[gcanlin]

Summary

This PR fuses the MiniMax H3 modulation operations into dedicated Triton kernels using FP32 accumulation. It removes redundant gather, elementwise, RMSNorm, and intermediate BF16 operations from each transformer block.

Performance

Benchmarked on MiniMax H3 Ref2VA with USP=4, eager execution, a 5-second output, and 49 denoising iterations. Both configurations were warmed up before profiling and end-to-end measurement.

End-to-End Performance

Metric	Baseline	FP32 Fused Modulation	Improvement
Engine E2E latency	136.961 s	134.707 s	2.253 s / 1.65%
Diffusion engine execution	134.695 s	132.199 s	2.496 s / 1.85%
Reported step latency	2.739 s	2.694 s	45.1 ms / 1.65%

Single DiT Forward Profile

Metric	Baseline	FP32 Fused Modulation	Improvement
GPU wall time	2,575.64 ms	2,534.49 ms	41.16 ms / 1.60%
CUDA kernel count	1,969	1,234	735 / 37.33%
Total CUDA kernel duration	2,571.83 ms	2,531.96 ms	39.87 ms / 1.55%

The fusion removes 294 gather kernels, 490 BF16 elementwise kernels, 98 in-place add kernels, and 98 standalone RMSNorm kernels per DiT forward, replacing them with 147 fused modulation kernels.

Accuracy

The implementation uses FP32 accumulation and casts the final results back to the destination tensor dtype. It is therefore not bit-exact with the baseline because intermediate BF16 rounding is removed.

Metric	Result
PSNR	39.55 dB
SSIM	0.9750

Before:

minimax_h3_modulation_off_measured_50step.mp4

After:

minimax_h3_modulation_on_fp32_measured_50step.mp4

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[gcanlin]

cc @hsliuustc0106 @david6666666 @lishunyang12

[vllm-omni-review-bot]

This PR appears to be related to model: MinimaxH3.

Model owners: @david6666666

@gcanlin, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[Sy0307]

LGTM

[hsliuustc0106]

All four wrappers gate the Triton path on x.is_cpu only, so NPU (and XPU/MUSA) tensors fall through to a raw Triton launch. MiniMax H3 is served on Ascend — this same transformer file opts NPU into ascend_laser_attention and the packed npu_attn_varlen path, and these blocks currently run norm1/norm2 through RMSNorm.forward_npu → torch_npu.npu_rms_norm. After this change an NPU forward reaches _rms_norm_indexed_scale_shift_kernel with torch_npu tensors on the first block and fails at the launch. The npu-ci lane doesn't run H3, so CI can't catch this.

The in-repo precedent is _fused_cuda_supported (vllm_omni/diffusion/layers/fused_qk_norm_rope.py:133): fuse only when HAS_TRITON and current_platform.is_cuda() and x.is_cuda (+ dtype check), otherwise run the eager reference. Since the eager bodies already exist here, I'd extend the gate the same way; on the non-CUDA path it's also worth routing the norm through the RMSNorm module so NPU keeps npu_rms_norm instead of the bf16 eager fallback.

[hsliuustc0106]

No test covers the new module. A CPU test in the style of tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py would pin the fallback math almost for free. Worth stating in the test (or the module docstring) that the CPU fallback intentionally keeps the old two-step numerics — norm rounded to input dtype, then affine — while the Triton kernel single-rounds in FP32, so CPU asserts validate the fallback and the PR-level PSNR is the GPU-parity evidence.

[hsliuustc0106]

Minor: the rms-norm kernels write outputs via row * hidden_size and all loads/stores assume unit column stride, so the wrappers only accept row-major contiguous inputs. Today's callers pass fresh contiguous tensors, but a strided input would silently read/write the wrong elements. An x.is_contiguous() assert (or a stated requirement in the module docstring) would pin the assumption.

[hsliuustc0106]

amd-ci is red on head (build 10670) while the CUDA/Intel/NPU lanes pass; the logs are behind auth on my side. The AMD lane runs the CPU-marked diffusion tests on a ROCm host, and the CPU fallbacks here should cover those, so I can't correlate the failure from the diff — could you check whether it's related or infra?