Feat: Attention Smoothing to Improve Generation Quality

benchmarks/fp8_sol_attention_quality/README.md

@@ -0,0 +1,114 @@
+# MiniMax-H3 FP8 Sol Attention Quality
+
+Date: 2026-09-08
+
+## Scope
+
+- Model: `MiniMaxAI/MiniMax-H3`, FL2VA partition
+- Hardware: one NVIDIA H100 80 GB; no sequence or tensor parallelism
+- Request: T2VA, 1344x768, 4 seconds, 107 decoded frames, 50 denoising steps, seed 0
+- Prompt: `Steam rises from the ramen while the family talks in the background.`
+- Sol policy: 10 dense steps, 2 dense layers, `tau=1.0`, exact threshold
+- FP8 profile: tf-kernel W8A8 Linear plus post-RoPE E4M3 Q/K/V
+- Timing: cold measured requests in clean processes with no warm-up. FP8 values are means of two runs; the BF16
+  reference is one run.
+
+The two FP8 profiles differ only at the QKV quantization boundary. The final profile subtracts per-head sequence
+means from K and V, adds the V mean back to the attention output, and corrects the residual V mean after E4M3
+rounding. These are attention-equivalent transforms, not Linear SmoothQuant.
+
+## End-to-end results
+
+| Profile | Denoise mean / range (s) | Throughput (step/s) | Peak allocated (GiB) |
+|---|---:|---:|---:|
+| BF16 Linear + FA4 | 211.442 | 0.23647 | 64.67 |
+| FP8 Sol, unsmoothed | 148.446 / 148.372-148.520 | **0.33682** | 37.11 |
+| FP8 Sol, fused KV smoothing + V bias correction | 151.651 / 151.328-151.974 | 0.32971 | **37.11** |
+
+The final profile is 39.4% faster and uses 42.6% less peak allocated memory than the matched BF16 reference. Against
+unsmoothed FP8 Sol, smoothing adds 2.16% mean denoising time and reduces mean throughput by 2.11%; peak allocated
+memory is unchanged. The duplicate FP8 runs produced identical frame and audio SHA256 hashes within each profile.
+
+## Attention-boundary error
+
+Real post-QK-norm, post-RoPE Q/K/V were captured from the first active Sol layer. The live tensor shape was
+`(1, 32626, 56, 128)` within a 32640-token padded request. The measurement used heads 0-3 over the complete live K/V
+context. Dense-attention output error used 64 evenly spaced query positions and FP32 math against the original BF16
+Q/K/V. This isolates the quantization boundary; it is not a full-generation metric.
+
+| Tensor boundary | Unsmoothed MSE | Smoothed MSE | MSE reduction |
+|---|---:|---:|---:|
+| K at the quantizer input distribution | 9.380e-4 | 7.349e-4 | **21.65%** |
+| V at the quantizer input distribution | 1.647e-2 | 1.611e-2 | 2.16% |
+| Reconstructed V | 1.647e-2 | 1.611e-2 | 2.17% |
+| Per-head/channel reconstructed V mean bias | 1.044e-6 | 2.561e-14 | **>99.99999%** |
+| Dense attention output | 7.034e-4 | 6.459e-4 | **8.18%** |
+
+For the dense attention output, smoothing also improved cosine similarity from 0.999403 to 0.999452, relative L2
+error from 0.03455 to 0.03311, and SQNR from 29.23 to 29.60 dB. Q is intentionally unchanged by smoothing and had
+identical quantization output in both profiles.
+
+## Generated media comparison
+
+The following metrics compare decoded outputs with the same-seed BF16 trajectory. They measure numerical trajectory
+similarity, not absolute perceptual quality: a small attention perturbation can select a different valid diffusion
+trajectory.
+
+| Profile | Frame cosine | PSNR (dB) | SSIM mean / min |
+|---|---:|---:|---:|
+| FP8 Sol, unsmoothed | 0.87488 | 14.695 | 0.5464 / 0.5151 |
+| FP8 Sol, KV smoothing + V correction | **0.87729** | **14.800** | **0.5659 / 0.5382** |
+
+SSIM was evaluated on every fourth frame (27 synchronized frames). All-frame cosine and PSNR use all 107 decoded
+uint8 frames.
+
+| Profile | Waveform cosine | Waveform MSE | SI-SDR (dB) | Spectral convergence | Log-spectral distance (dB) |
+|---|---:|---:|---:|---:|---:|
+| FP8 Sol, unsmoothed | **0.55217** | **9.842e-5** | **-3.579** | **0.6301** | **12.484** |
+| FP8 Sol, KV smoothing + V correction | 0.53450 | 1.026e-4 | -3.980 | 0.6388 | 12.522 |
+
+Both FP8 videos are coherent and preserve the prompt's ramen-dining composition without corrupted frames. Smoothing
+improves every reported video trajectory metric in this seed, while the unsmoothed output is closer to BF16 on the
+reported audio trajectory metrics. This single-seed media comparison therefore does not establish a universal audio
+quality ranking; the attention-boundary metrics directly measure the error targeted by the implementation.
+
+## Fusion results
+
+The optimized path combines K and V sequence statistics into one Triton reduction and merges the BF16 dense prefix
+with the corrected sparse suffix in one output pass. At the current H3 live shape `(1, 32626, 56, 128)`, H100
+100-repetition means are:
+
+| Boundary operation | Before fusion (ms) | After fusion (ms) | Change |
+|---|---:|---:|---:|
+| Exact K/V sequence statistics | 1.0097 | 0.5319 | -47.3% |
+| Smoothed QKV quantization and correction statistics | 2.5795 | 2.1054 | -18.4% |
+| Output correction and dense-prefix merge | 0.8776 | 0.5540 | -36.9% |
+| Combined quantization/correction and output boundary | 3.4571 | 2.6594 | **-23.1%** |
+
+The GPU tests verify the fused K/V statistics, quantized tensors, correction, and fused prefix merge are bitwise
+identical to their unfused exact implementations. Fusion adds no large partial-statistics buffer.
+
+Several faster-looking approximations were rejected. A two-warp reduction changed long-sequence FP32 summation,
+sampling K/V centers damaged Sol routing quality, and a block-partial V-bias fusion was slower than directly scanning
+the compact FP8 V tensor. The remaining end-to-end cost is the exact global-statistics work; it was not removed by
+sacrificing the target error reduction.
+
+## Reproduction
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m tools.validation.benchmark_minimax_h3_fp8_sol_sp \
+  --gpu-num 1 --profile baseline --duration 4 --steps 50 --no-warmup \
+  --output outputs/h3_bf16.mp4
+
+CUDA_VISIBLE_DEVICES=0 python -m tools.validation.benchmark_minimax_h3_fp8_sol_sp \
+  --gpu-num 1 --profile optimized --duration 4 --steps 50 --no-warmup \
+  --sol-fp8-smoothing none --no-sol-fp8-v-bias-correction \
+  --output outputs/h3_fp8_sol_unsmoothed.mp4
+
+CUDA_VISIBLE_DEVICES=0 python -m tools.validation.benchmark_minimax_h3_fp8_sol_sp \
+  --gpu-num 1 --profile optimized --duration 4 --steps 50 --no-warmup \
+  --output outputs/h3_fp8_sol_kv_bias.mp4
+```
+
+The benchmark saves synchronized MP4, `.frames.npy`, `.audio.npy`, and metrics JSON artifacts. MP4 and NumPy
+artifacts remain local and are excluded from Git.

docs/en/blog/fp8_sol_attention.md

@@ -224,6 +224,23 @@ MiniMax-H3 has a packed multimodal sequence, so it needs two additional protecti
 prefix is registered as an exact KV sink, and prefix queries are recomputed with BF16 dense attention. The first ten
 steps and first two DiT layers also use matched packed FlashAttention-4. Token-refiner attention remains dense.
 
+MiniMax-H3 FP8 Sol also defaults to attention-specific sequence-mean smoothing. This is not Linear SmoothQuant and
+does not migrate scale between a Linear activation and its weights. For one head, TeleFuser applies
+
+$$K' = K - \operatorname{mean}_{T}(K), \qquad V' = V - \operatorname{mean}_{T}(V).$$
+
+The K shift adds the same scalar to every logit in a query row, so softmax is unchanged. Since every softmax row sums
+to one, the original result is recovered by adding `mean_T(V)` to the attention output. TeleFuser computes both means
+in FP32, fuses their subtraction into SM90 E4M3 preparation, and corrects the residual mean bias measured after V is
+rounded to E4M3. `sol_fp8_smoothing=none|k|kv` controls the two equivalent transforms, while
+`sol_fp8_v_bias_correction` controls the post-quantization V correction for ablation.
+The exact K/V statistics share one fused reduction, and a second Triton kernel merges the BF16 dense prefix with the
+corrected sparse suffix in one output pass. At the current H3 live shape `(1, 32626, 56, 128)`, this reduces the
+measured smoothing boundary by 23.1% while remaining bitwise identical to the unfused implementation. Two matched
+50-step runs show a 2.11% throughput cost versus unsmoothed FP8 Sol, with unchanged 37.11 GiB peak allocated memory.
+The matched single-H100 quality and performance ablation is recorded in
+[`benchmarks/fp8_sol_attention_quality`](https://github.com/Tele-AI/TeleFuser/tree/main/benchmarks/fp8_sol_attention_quality#readme).
+
 Unsupported shapes, dtypes, devices, or runtime kernel failures retain the public attention fallback. FP8 operands
 are dequantized before the BF16 fallback. Pure Ulysses sequence parallelism is supported: its all-to-all first
 produces full-sequence, local-head Q/K/V, then each rank computes FP8 scales and runs Sol independently. Ring and
@@ -328,13 +345,16 @@ python examples/wan_video/wan21_1_3b_text_to_video_optimized_h100.py \
   128. BF16 Sol has broader architecture fallbacks, but the performance result does not transfer to them.
 - MiniMax-H3 online `tf-kernel` FP8 Linear quantization is currently single-GPU only. Its TP/FSDP loading contract
   remains BF16.
-- QKV quantization and centroid preprocessing are separate kernels. Further fusion may reduce launch and memory-traffic overhead, but would increase specialization and register pressure.
+- Exact sequence statistics still precede QKV quantization. Fusing the K/V reduction and output merge reduces the
+  boundary overhead, but the remaining global reduction leaves about a 1% measured throughput difference versus
+  unsmoothed FP8 Sol. Sampled statistics were rejected because they changed Sol routing quality.
 - CuTe compilation is shape- and dtype-specific. Cold-start latency includes compilation; persistent services should
   evaluate warm steady state separately.
 - The best FP8 layer range is model- and checkpoint-dependent. An all-layer setting should not be treated as the
   default quality/performance point.
-- Peak allocated memory is a CUDA allocator metric, not total process or device memory. The experiments report one
-  run per configuration and do not establish variance bounds.
+- Peak allocated memory is a CUDA allocator metric, not total process or device memory. The matched unsmoothed and
+  final fused profiles report two runs; the other ablation points report one run, so these results do not establish
+  broad variance bounds.
 
 ## Related Work
 

docs/zh/blog/fp8_sol_attention.md

@@ -217,6 +217,22 @@ MiniMax-H3 使用 packed multimodal sequence，因此增加了两项保护：完
 KV sink，prefix query 使用 BF16 dense attention 重新计算。前十个 step、前两个 DiT layer 使用匹配的
 packed FlashAttention-4，token refiner 也始终保持 dense。
 
+MiniMax-H3 FP8 Sol 还默认使用 attention sequence-mean smoothing。它不是 Linear SmoothQuant，不会在
+Linear activation 与权重之间迁移 scale。对每个 attention head，TeleFuser 执行
+
+$$K' = K - \operatorname{mean}_{T}(K), \qquad V' = V - \operatorname{mean}_{T}(V).$$
+
+K 平移对同一 query row 的每个 logit 增加相同常数，因此 softmax 不变；softmax 每行权重和为 1，所以在
+attention output 加回 `mean_T(V)` 即恢复原始结果。TeleFuser 使用 FP32 统计均值，把中心化融合进 SM90
+E4M3 preparation，并补偿 V 舍入到 E4M3 后实测到的残余均值偏差。`sol_fp8_smoothing=none|k|kv` 用于
+分别消融两个等价变换，`sol_fp8_v_bias_correction` 控制量化后 V 偏差补偿。
+exact K/V 统计共享一个 fused reduction，另一个 Triton kernel 在一次写出中合并 BF16 dense prefix 与完成
+correction 的 sparse suffix。在当前 H3 live shape `(1, 32626, 56, 128)` 上，这将 smoothing boundary 耗时降低
+23.1%，并与融合前实现保持 bitwise 一致。两轮匹配的 50-step 实验中，相比 unsmoothed FP8 Sol 的吞吐开销
+为 2.11%，peak allocated 显存均为 37.11 GiB。
+匹配的单卡 H100 画质与性能消融记录在
+[`benchmarks/fp8_sol_attention_quality`](https://github.com/Tele-AI/TeleFuser/tree/main/benchmarks/fp8_sol_attention_quality#readme)。
+
 不支持的 shape、dtype、device 或 kernel runtime failure 会保留公共 attention fallback。FP8 operand 会先
 反量化再进入 BF16 fallback。纯 Ulysses sequence parallel 已支持：all-to-all 先得到完整 sequence、局部
 head 的 Q/K/V，每个 rank 再独立计算 FP8 scale 并运行 Sol。Ring 以及 Ulysses-ring 组合仍走 dense，因为
@@ -315,12 +331,12 @@ python examples/wan_video/wan21_1_3b_text_to_video_optimized_h100.py \
 - 已验证的 FP8 attention mainloop 面向 SM90、noncausal self-attention、相同 Q/K/V shape 和 128 维 head。
   BF16 Sol 有更广的 architecture fallback，但本文性能数据不能直接迁移到这些路径。
 - MiniMax-H3 在线 `tf-kernel` FP8 Linear 目前只支持单 GPU；其 TP/FSDP loading contract 仍为 BF16。
-- QKV quantization 与 centroid preprocessing 是独立 kernel。进一步融合可能减少 launch 与访存开销，
-  但也会提高 specialization 数量和 register pressure。
+- exact sequence statistics 仍必须先于 QKV quantization。融合 K/V reduction 与 output merge 后，剩余的
+  全局 reduction 使吞吐相比 unsmoothed FP8 Sol 仍约低 1%；采样统计会改变 Sol routing 质量，因此未采用。
 - CuTe 编译与 shape、dtype 绑定。冷启动结果包含编译，常驻服务还应单独评估 warm steady state。
 - 最佳 FP8 layer range 依赖模型与 checkpoint，不能把全层 FP8 当作默认质量/性能点。
-- Peak allocated 是 CUDA allocator 指标，不是进程或整张 GPU 的总显存。每个配置只有一次测量，尚未给出
-  方差范围。
+- Peak allocated 是 CUDA allocator 指标，不是进程或整张 GPU 的总显存。匹配的 unsmoothed 与最终 fused
+  profile 各测了两轮，其他消融点各测一轮，因此这些结果尚不能建立完整的方差范围。
 
 ## 相关工作
 

examples/minimax_h3/README.md

@@ -419,7 +419,11 @@ The same FL2VA example exposes dense/Sol and BF16/FP8 as independent switches. `
 tf-kernel W8A8 Linear GEMMs to the transformer blocks. `--attn-impl SOL_ATTN` enables the MiniMax-H3 Sol policy:
 the first 10 denoising steps and first 2 DiT layers remain dense, the full condition prefix is an exact KV sink, and
 prefix queries are recomputed with BF16 dense attention. Adding `--sol-fp8` quantizes post-RoPE Q/K/V in active sparse
-layers and dispatches the SM90 CuTe FP8 Sol mainloop.
+layers and dispatches the SM90 CuTe FP8 Sol mainloop. The MiniMax-H3 quality profile also centers K and V over the live
+sequence before E4M3 conversion, adds the V mean back after attention, and corrects the residual V mean bias caused by
+E4M3 rounding. These are attention-equivalent transforms rather than Linear SmoothQuant. K/V statistics use one exact
+fused reduction, and dense-prefix replacement is fused with the sparse-output correction; the optimized output is
+bitwise identical to the original exact smoothing path.
 
 ~~~bash
 # BF16 dense
@@ -435,22 +439,23 @@ python -m examples.minimax_h3.minimax_h3_fl2va_h100 \
 
 # FP8 Linear + FP8 Sol attention
 python -m examples.minimax_h3.minimax_h3_fl2va_h100 \
-  --gpu-num 4 --mode t2va --quantization fp8 --attn-impl SOL_ATTN --sol-fp8 \
+  --gpu-num 1 --mode t2va --quantization fp8 --attn-impl SOL_ATTN --sol-fp8 \
   --output outputs/h3_fp8_sol.mp4
 ~~~
 
 Use `--sol-dense-steps`, `--sol-dense-layers`, `--sol-tau`, `--sol-threshold-type`,
-`--sol-fp8-layer-start`, and `--sol-fp8-layer-end` to override the policy for controlled ablations. The defaults
-match the released H100 MiniMax-H3 Sol profile.
+`--sol-fp8-layer-start`, and `--sol-fp8-layer-end` to override the policy for controlled ablations. FP8 Sol enables
+K+V smoothing and V bias correction by default, so the standard example needs no additional quality flags. The
+dedicated validation benchmark retains explicit overrides for comparing against an unsmoothed profile.
 
-For a warmed four-GPU comparison that saves synchronized MP4s, decoded arrays, throughput, and sampled device-memory
-peaks, run the two profiles below. Four GPUs use Ulysses2 x TP2 for both profiles.
+For a warmed matched comparison that saves synchronized MP4s, decoded arrays, throughput, and sampled device-memory
+peaks, run the two single-GPU profiles below. Pass `--gpu-num 4` to retain the earlier Ulysses2 x TP2 benchmark mode.
 
 ~~~bash
 python -m tools.validation.benchmark_minimax_h3_fp8_sol_sp \
-  --profile baseline --output benchmarks/minimax_h3_fp8_sol/baseline_sp2_tp2_bf16_flash.mp4
+  --gpu-num 1 --profile baseline --output benchmarks/minimax_h3_fp8_sol/baseline_bf16_flash.mp4
 python -m tools.validation.benchmark_minimax_h3_fp8_sol_sp \
-  --profile optimized --output benchmarks/minimax_h3_fp8_sol/optimized_sp2_tp2_fp8_sol_exact.mp4
+  --gpu-num 1 --profile optimized --output benchmarks/minimax_h3_fp8_sol/optimized_fp8_sol_exact.mp4
 ~~~
 
 For matched BF16/TorchAO-FP8/tf-kernel-FP8/NF4 profiling, use the validation benchmark. It writes the synchronized MP4 plus a JSON report

examples/minimax_h3/common.py

@@ -182,6 +182,8 @@ def load_minimax_h3_pipeline(
     sol_threshold_type: str = "exact",
     sol_fp8_layer_start: int = 0,
     sol_fp8_layer_end: int | None = None,
+    sol_fp8_smoothing: str = "kv",
+    sol_fp8_v_bias_correction: bool = True,
     feature_cache_config: FeatureCacheConfig | None = None,
     adaln_cache_path: str | Path | None = None,
     online_adaln_cache: bool = False,
@@ -278,6 +280,8 @@ def load_minimax_h3_pipeline(
             sol_fp8_layer_end=sol_fp8_layer_end,
             attention_chunks=attention_chunks,
             ulysses_sequence_mode=ulysses_sequence_mode,
+            sol_fp8_smoothing=sol_fp8_smoothing,
+            sol_fp8_v_bias_correction=sol_fp8_v_bias_correction,
         )
         if attn_impl == AttnImplType.SOL_ATTN
         else AttentionConfig.dense_attention(

examples/minimax_h3/minimax_h3_fl2va_h100.py

@@ -99,6 +99,8 @@ def get_pipeline(
     sol_threshold_type: str = "exact",
     sol_fp8_layer_start: int = 0,
     sol_fp8_layer_end: int | None = None,
+    sol_fp8_smoothing: str = "kv",
+    sol_fp8_v_bias_correction: bool = True,
     enable_feature_cache: bool = False,
     feature_cache_model_type: str = PPL_CONFIG["feature_cache_model_type"],
     feature_cache_n_derivatives: int = PPL_CONFIG["feature_cache_n_derivatives"],
@@ -129,6 +131,8 @@ def get_pipeline(
         sol_threshold_type=sol_threshold_type,
         sol_fp8_layer_start=sol_fp8_layer_start,
         sol_fp8_layer_end=sol_fp8_layer_end,
+        sol_fp8_smoothing=sol_fp8_smoothing,
+        sol_fp8_v_bias_correction=sol_fp8_v_bias_correction,
         feature_cache_config=FeatureCacheConfig(
             enabled=enable_feature_cache,
             model_type=feature_cache_model_type,