perf(minimax-h3): optimize Ulysses attention and DiT kernels

[L17807593]

TeleFuser MiniMax-H3: Five Lossless Optimizations

1. Summary

Under a four-GPU setup with five prompt/seed pairs, the five lossless optimizations reduce the average wall time from 78.546 s to 75.766 s, achieving a 3.539% average speedup. The generated outputs remain bit-exact with the baseline.

Configuration	Avg. wall time	Speedup vs. baseline	Avg. denoise time	DiT calls
Baseline	78.546 s	0.000%	75.737 s	49
Five lossless optimizations	75.766 s	3.539%	72.934 s	49

2. Optimizations

2.1 Overlap Ulysses communication with FlashAttention 4

Added AttentionConfig.attention_chunks. The generic default remains 1 to preserve existing behavior; the H100 example uses 2.

MiniMax-H3 attention splits local heads into two chunks. Each chunk enters FA4 immediately after the fused scatter, while its gather is submitted asynchronously, allowing computation for the next chunk to overlap with communication for the previous chunk.

New interfaces include:

• ulysses_scatter_qkv_qknorm_rope_chunk_async
• ulysses_gather_heads_chunk_async

The new path is enabled only when Ulysses, FA4, CUDA BF16, and the supported tensor layout are all available. Otherwise, execution falls back to the original attention path.

2.2 Fuse Q/K RMSNorm, RoPE, and Ulysses Pack

Added the Triton kernel pack_qkv_qknorm_rope_destination_major, which performs the following in a single kernel:

1. Q/K RMSNorm;
2. BF16-compatible rounding;
3. RoPE;
4. Q/K/V destination-major Ulysses packing.

This reduces kernel launches, intermediate layout conversions, and temporary tensors. Guards on dtype, device, weights, RoPE cache, head topology, and packed shape ensure that unsupported cases fall back to the original implementation.

2.3 Add fixed-valid and valid-only attention modes

FA4 packed attention now supports:

• fixed_valid
• pad_fixed_valid_output

For MiniMax-H3, valid video tokens are at the front of the packed tensor, while the tail contains alignment padding or auxiliary short sequences. The fixed-valid path runs FA4 only on the valid sequence with a fixed shape. The valid-only mode does not restore the invalid tail after attention, reducing attention and subsequent communication work.

The generic defaults remain:

fixed_valid=False
ulysses_sequence_mode="padded"

Therefore, existing call sites retain their original behavior.

2.4 Fuse valid-only gather with zero-tail merge

Added the Triton kernel merge_ulysses_head_chunk. It writes each received All-to-All head chunk directly into the final destination tensor and clears the invalid tail in the same kernel, avoiding separate relayout, copy, and zero-tail operations.

2.5 Fuse RMSNorm with indexed AdaLN modulation

Added the Triton kernel indexed_rmsnorm_scale_shift_bf16, which fuses the following sequence in DiT blocks and the final layer:

RMSNorm -> indexed scale/shift modulation

The implementation preserves the original BF16 store/load rounding boundaries. CUDA unit tests verify element-wise bit-exact equivalence with the previous path. The fused kernel is enabled only for eager CUDA BF16, contiguous tensors, and valid shapes; other cases fall back to the original _modulate(norm(hidden), ...) implementation.

3. Per-case performance

Case	Baseline (s)	Five lossless optimizations (s)	Speedup
P0/S0	78.426	75.745	3.414%
P0/S42	78.420	75.893	3.224%
P0/S1234	78.801	75.740	3.884%
P1/S0	78.612	75.706	3.697%
P1/S42	78.473	75.745	3.476%

4. Quality verification

All quality comparisons use the baseline generated with the same prompt and seed, comparing the final video and audio MP4 outputs.

Configuration	Exact MP4 SHA match	Video PSNR	Video SSIM	Audio cosine	Audio SNR	Audio RMSE
Five lossless optimizations	5/5	inf	1.000000	1.000000	inf	0

Additional hash check:

• baseline == lossless: 5/5.

The five optimizations are therefore bit-exact with the baseline at the final video and audio file level, while providing a 3.539% average wall-time speedup.