Feat: Attention Smoothing to Improve Generation Quality

[Uxito-Ada]

Summary

Follow up on #35 by improving MiniMax-H3 FP8 Sol-Attn quality with attention-equivalent K/V sequence-mean smoothing.

On the H100 MiniMax-H3, the final profile reduces attention-output MSE by 8.18% versus unsmoothed FP8 Sol, remains 39.4% faster than BF16, and uses 42.6% less peak allocated GPU memory.

Changes

• Centralize K in FP32 before E4M3 quantization
• Centralize V, restore its sequence mean after attention
• Fuse exact K/V statistics into one Triton reduction
• Merge BF16 dense-prefix replacement and sparse-output correction in one output pass
• Enable KV smoothing and V-bias correction by default in AttentionConfig.sol_attention; the standard MiniMax-H3 example no longer needs dedicated quality flags, while the validation benchmark keeps explicit ablation overrides.

This is attention smoothing, not SmoothQuant, where no scale is migrated between Linear activations and weights.

MiniMax-H3: Unsmoothed vs Smoothed FP8 Sol

Workload: 1 x H100 80GB, no TP/SP, FL2VA T2VA 1344x768, 4 seconds / 107 decoded frames at 24 fps, 50 denoising steps, seed 0, exact Sol threshold, 10 dense steps, 2 dense layers, cold processes without warm-up, and MP4 serialization excluded from generation timing. Each FP8 profile was measured twice.

Prompt:

Steam rises from the ramen while the family talks in the background.

Performance overview

• The final smoothed profile is 39.4% faster and uses 42.6% less peak allocated memory than BF16.
• Against unsmoothed FP8 Sol, exact smoothing adds 2.16% denoising time and costs 2.11% throughput, with unchanged peak allocated memory, while improving the generation quality as the following.

Generated Output Comparison

BF16 Linear + FlashAttention 4

bf16_fa4.mp4

FP8 Linear + FP8 Sol, no smoothing

fp8_sol_unsmoothed.mp4

FP8 Linear + FP8 Sol, KV smoothing + V correction

fp8_sol_smoothed.mp4

Quantitative analysis, against the same-seed BF16 output:

Profile	Frame cosine	PSNR	SSIM mean / min	Audio cosine	SI-SDR	Log-spectral distance
FP8 Sol, no smoothing	0.87488	14.695 dB	0.5464 / 0.5151	0.55217	-3.579 dB	12.484 dB
FP8 Sol, KV smoothing + V correction	0.87729	14.800 dB	0.5659 / 0.5382	0.53450	-3.980 dB	12.522 dB

Smoothing improves every reported video metric.

Smooth Attention Kernel Accuracy

Measurement setup:

• Real post-QK-norm, post-RoPE Q/K/V were captured from the first active Sol layer at live shape (1, 32626, 56, 128)
• Covers heads 0-3 over the complete 32,626-token K/V context
• Attention-output error uses 64 evenly spaced query positions and FP32 dense-attention math

Boundary	Unsmoothed MSE	Smoothed MSE	Benefit
quanted K	9.380e-4	7.349e-4	21.65%
dequanted V	1.647e-2	1.611e-2	2.17%
dequanted V's mean	1.044e-6	2.561e-14	>99.99999%
attention output (against dense)	7.034e-4	6.459e-4	8.18%

Dense-attention output cosine improves 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.

Testing

ruff check <PR-scoped files>
ruff format --check <PR-scoped files>
git diff --check
# All checks passed

python -m pytest -q -m "not gpu" \
  tests/unit/models/test_minimax_h3_dit.py \
  tests/unit/ops/test_fp8_attention_smoothing.py \
  tests/unit/ops/test_sol_attention.py \
  tests/unit/pipelines/minimax_h3/test_examples.py
# 76 passed, 8 deselected

python -m pytest -q -m gpu \
  tests/unit/models/test_minimax_h3_dit.py \
  tests/unit/ops/test_fp8_attention_smoothing.py \
  tests/unit/ops/test_sol_attention.py \
  tests/unit/pipelines/minimax_h3/test_examples.py
# 8 passed, 76 deselected on one H100

python scripts/docs/prepare_cookbook.py
# Passed

GitHub CI also passes lint, build, server tests, and unit tests on Python 3.10, 3.11, and 3.12.

Architecture Support

• Native FP8 Sol and fused smoothing kernels: NVIDIA Hopper SM90.
• Single-GPU MiniMax-H3: validated in this PR.
• Pure Ulysses and MiniMax-H3 TP + Ulysses: compatibility retained from Enable Sequence Parallelism on FP8 Sol-Attn #35.
• Fused Ulysses + FlashAttention 4 path: retained for dense attention.
• Ring and Ulysses-Ring: dense fallback retained.
• Wan FP8 Sol behavior: unchanged.

[lzx1413]

It's better to add some comparison videos.

[lzx1413]

KV smoothing can be enabled by default, eliminating the need for command-line arguments and simplifying example configurations.