Implementation of FP8 Sol-Attn on Wan2.1

docs/en/attention.md

@@ -71,6 +71,9 @@ class SparseAttentionConfig:
     sol_tau: float = 1.0                    # Sol-Attn routing threshold
     sol_threshold_type: str = "diag"        # "diag" or "exact"
     sol_kv_splits: int | str = "auto"       # "auto", 1, 2, or 4
+    sol_fp8: bool = False                    # FP8 Q/K/V Sol-Attn on SM90
+    sol_fp8_layer_start: int = 0             # First layer using FP8 Sol-Attn
+    sol_fp8_layer_end: int | None = None     # Exclusive end; None means all remaining layers
 ```
 
 ## Calling Flow
@@ -197,10 +200,20 @@ config = AttentionConfig.sol_attention()
 pipe_config.dit_config.attention_config = config
 ```
 
-Sol-Attn is used only for contiguous, noncausal BF16 self-attention with equal Q/K/V
-shapes and head dimension 128. Unsupported calls, dense warmup layers or timesteps,
-and kernel runtime failures fall back to the existing dense attention path. Ring/USP
-also remains dense because its online merge requires log-sum-exp output.
+Sol-Attn is used for contiguous, noncausal self-attention with equal Q/K/V shapes
+and head dimension 128. BF16 is supported by the architecture-specific kernels;
+SM90 additionally supports E4M3 Q/K/V with FP32 accumulation. Unsupported calls,
+dense warmup layers or timesteps, and kernel runtime failures fall back to the
+existing dense attention path. Ring/USP remains dense because its online merge
+requires log-sum-exp output.
+
+`sol_fp8_layer_start` and `sol_fp8_layer_end` restrict E4M3 Q/K/V to a half-open
+transformer-layer range. Sparse layers outside that range continue to use BF16
+Sol-Attn. This controls accumulated FP8 routing error in diffusion models.
+Setting `dense_timesteps=0`, `dense_layers=0`, and a negative `tau` forces all
+KV blocks onto the exact route. The Wan optimized example exposes this as
+`--attention fp8-dense`; `--attention fp8-sol` enables centroid routing with the
+same FP8 Q/K/V and QK/PV kernel.
 
 ### QwenImagePipeline / ZImagePipeline
 

docs/zh/attention.md

@@ -71,6 +71,9 @@ class SparseAttentionConfig:
     sol_tau: float = 1.0                    # Sol-Attn 路由阈值
     sol_threshold_type: str = "diag"        # "diag" 或 "exact"
     sol_kv_splits: int | str = "auto"       # "auto"、1、2 或 4
+    sol_fp8: bool = False                    # SM90 FP8 Q/K/V Sol-Attn
+    sol_fp8_layer_start: int = 0             # 启用 FP8 Sol-Attn 的首层
+    sol_fp8_layer_end: int | None = None     # 结束层（不包含）；None 表示其余所有层
 ```
 
 ## 调用流程
@@ -197,10 +200,19 @@ config = AttentionConfig.sol_attention()
 pipe_config.dit_config.attention_config = config
 ```
 
-Sol-Attn 仅用于连续、非因果、BF16、Q/K/V 形状相同且 head dimension 为 128 的
-self-attention。其他调用、dense 预热层/时间步以及内核运行失败都会回退到现有密集路径。
+Sol-Attn 用于连续、非因果、Q/K/V 形状相同且 head dimension 为 128 的
+self-attention。各架构内核支持 BF16，SM90 还支持使用 FP32 累加的 E4M3 Q/K/V。
+其他调用、dense 预热层/时间步以及内核运行失败都会回退到现有密集路径。
 Ring/USP 需要 LSE 做在线合并，因此仍使用支持 LSE 的密集后端。
 
+`sol_fp8_layer_start` 和 `sol_fp8_layer_end` 用半开区间限制使用 E4M3 Q/K/V
+的 transformer 层，区间外的稀疏层继续使用 BF16 Sol-Attn，以控制扩散模型中
+逐层累积的 FP8 路由误差。
+
+设置 `dense_timesteps=0`、`dense_layers=0` 和负数 `tau` 会强制所有 KV block
+走 exact 路径。Wan 优化示例将其暴露为 `--attention fp8-dense`；
+`--attention fp8-sol` 使用相同的 FP8 Q/K/V 与 QK/PV kernel 并启用质心路由。
+
 ### QwenImagePipeline / ZImagePipeline
 
 仅支持密集注意力（图像生成没有时序维度）：

examples/wan_video/README.md

@@ -89,7 +89,6 @@ python examples/wan_video/wan21_1_3b_text_to_video_h100.py --resolution 480p --a
 **Features:**
 - Video Frame Interpolation (VFI) with RIFE model for 30fps output
 - CFG parallel when cfg_scale > 1
-
 #### wan21_1_3b_text_to_video_hf.py
 
 T2V with HuggingFace format loading.
@@ -173,6 +172,132 @@ pipe_config.dit_config.attention_config = AttentionConfig.sol_attention()
 
 Sol-Attn is built into TeleFuser. Eligible BF16 self-attention calls use the sparse kernel; unsupported calls
 automatically use the existing dense fallback. The defaults follow the official Wan2.1 profile: Morton3D token ordering, dense layer 0, and 10 dense warm-up steps for the standard 50-step schedule.
+#### wan21_1_3b_text_to_video_optimized_h100.py
+
+Provides one entry point for independently enabling attention and quantization
+optimizations. The defaults are `--attention dense --quantization none`, which
+run the BF16 baseline.
+
+Attention choices:
+
+- `dense`: BF16 PyTorch SDPA
+- `sol`: BF16 Sol-Attn with dense warm-up and fallback calls
+- `fp8-dense`: E4M3 Q/K/V with FP8 QK/PV WGMMA; all KV blocks are exact
+- `fp8-sol`: the same FP8 kernel with Sol routing enabled
+
+For both FP8 modes, `--fp8-layer-start` and `--fp8-layer-end` select the
+half-open transformer-layer range that uses FP8 Q/K/V. Other layers use the
+corresponding BF16 dense or Sol path. Restricting FP8 Q/K/V to middle layers
+avoids accumulating small quantization changes across the full denoiser.
+
+Quantization choices:
+
+- `none`: BF16 DiT
+- `tf-kernel-fp8`: TeleFuser dynamic W8A8 FP8 GEMM
+
+`--fp8-linear-scope all` quantizes every transformer-block Linear layer.
+`--fp8-linear-scope ffn` keeps self/cross-attention projections in BF16 and
+quantizes the 60 FFN Linear layers. The default `auto` selects `all`; generated
+video validation shows that all-Linear FP8 preserves quality. Attention Q/K/V
+are more sensitive, so their default FP8 layer range is 10-19 for this 30-layer
+Wan2.1 model.
+
+Only DiT transformer-block Linear layers are quantized; the VAE and text encoder
+remain BF16. Select the two optimization axes independently:
+
+```bash
+# Dense + BF16 baseline
+python examples/wan_video/wan21_1_3b_text_to_video_optimized_h100.py \
+    --model-root /path/to/Wan2.1-T2V-1.3B \
+    --attention dense --quantization none
+
+# Sol-Attn + BF16
+python examples/wan_video/wan21_1_3b_text_to_video_optimized_h100.py \
+    --model-root /path/to/Wan2.1-T2V-1.3B \
+    --attention sol --quantization none
+
+# FP8 Dense: exact FP8 attention + FP8 Linear
+python examples/wan_video/wan21_1_3b_text_to_video_optimized_h100.py \
+    --model-root /path/to/Wan2.1-T2V-1.3B \
+    --attention fp8-dense \
+    --quantization tf-kernel-fp8 \
+    --fp8-layer-start 10 \
+    --fp8-layer-end 20
+
+# FP8 Sol: routed FP8 attention + FP8 Linear
+python examples/wan_video/wan21_1_3b_text_to_video_optimized_h100.py \
+    --model-root /path/to/Wan2.1-T2V-1.3B \
+    --attention fp8-sol \
+    --quantization tf-kernel-fp8 \
+    --fp8-layer-start 10 \
+    --fp8-layer-end 20 \
+    --dense-timesteps 10 \
+    --dense-layers 1 \
+    --tau 1.0 \
+    --threshold-type diag \
+    --kv-splits auto
+```
+
+In this example, FP8 means the E4M3 attention implementation rather than a
+BF16-attention run with only its Linear layers quantized. Post-RoPE Q/K/V are
+quantized and QK/PV run through the CuTe SM90 WGMMA mainloop. Q/K use one scale
+per 64-token block, V uses per-channel scales and a K-major layout, and FP32
+accumulators are used throughout. `fp8-dense` forces every routed KV block onto
+the exact path, while `fp8-sol` permits centroid approximation. `auto` selects
+two KV splits for long FP8 sequences, which is faster at Wan's sequence length
+without changing the FP32 accumulation contract.
+Partial tiles are physically padded while the original sequence length remains
+masked in the kernel. FP8 split execution restores the represented N64 route
+length before PV, matching the BF16 summed-centroid contract. The accompanying
+FP8 Linear GEMMs use the tf-kernel backend. Self-attention Q/K/V projections
+share one dynamic activation quantization instead of quantizing the same input
+three times. With a partial FP8 layer range, FP8 Dense sends unquantized layers
+to SDPA and FP8 Sol sends unquantized sparse layers to Triton, avoiding a second
+CuTe specialization in the cold-start path.
+The final log reports generation time, frames per second, and peak allocated and
+reserved CUDA memory.
+
+##### H100 benchmark
+
+This clean-process cold-start benchmark runs each configuration in a separate process
+with no other GPU processes on one H100 80GB. It uses the official Wan2.1 T2V-1.3B example prompt, `832x480`,
+81 frames, 50 UniPC steps, CFG 5.0, sigma shift 5.0, and seed 42. Generation timing
+starts after pipeline loading, so it includes first-execution kernel/JIT costs but
+excludes model loading. Peak memory is `torch.cuda.max_memory_allocated()` over the
+same generation interval.
+
+| Quantization | Attention | Throughput (frames/s) | Peak allocated (GiB) |
+| --- | --- | ---: | ---: |
+| BF16 | Dense | 0.8491 | 16.147 |
+| BF16 | Sol-Attn | 1.1090 | 17.023 |
+| FP8 | Dense (Q/K/V layers 10-19, exact) | 0.8739 | 15.730 |
+| FP8 | Sol-Attn (Q/K/V layers 10-19) | 1.1565 | 15.730 |
+
+Both FP8 rows quantize all 300 transformer-block Linear layers and use the same
+E4M3 attention layer range. FP8 Dense therefore measures this implementation's
+exact QK/PV path, not BF16 SDPA with only Linear quantization. Against the
+corresponding BF16 output, FP8 Dense measures 22.0257 dB PSNR / 0.828783 SSIM,
+and FP8 Sol measures 20.8502 dB PSNR / 0.792656 SSIM.
+
+The benchmark prompt is:
+
+> Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely
+> on a spotlighted stage.
+
+Example command (change `--attention` and `--quantization` for each ablation):
+
+```bash
+python examples/wan_video/wan21_1_3b_text_to_video_optimized_h100.py \
+    --model-root /path/to/Wan2.1-T2V-1.3B \
+    --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+    --attention fp8-sol \
+    --quantization tf-kernel-fp8 \
+    --fp8-linear-scope all \
+    --fp8-layer-start 10 --fp8-layer-end 20 \
+    --width 832 --height 480 \
+    --num-frames 81 --num-inference-steps 50 \
+    --sample-solver unipc --cfg-scale 5.0 --sigma-shift 5.0 --seed 42
+```
 
 #### wan21_1_3b_text_to_video_cache_calibrate.py