[gcanlin]

Summary

This PR adds a local SiluAndMul layer under vllm_omni.diffusion.layers and uses it in the MiniMax H3 MLP. The implementation preserves vLLM-Omni platform dispatch through CustomOp and current_omni_platform, while the CUDA path invokes the fused torch.ops._C.silu_and_mul kernel. This avoids depending on the vLLM SiluAndMul class and prevents the diffusion configuration from falling back to the native two-kernel implementation.

Performance

Benchmarked on MiniMax H3 Ref2VA with USP=4, eager execution, a 5-second output, and 49 denoising iterations. The existing main-branch result was reused as the baseline. The fused implementation was warmed up separately before both the 2-step profile and the 50-step end-to-end measurement.

End-to-End Performance

Metric	Baseline	Local Fused SiluAndMul	Improvement
Engine E2E latency	136.961 s	135.104 s	1.857 s / 1.36%
Diffusion engine execution	134.695 s	132.838 s	1.857 s / 1.38%
Reported step latency	2.739 s	2.702 s	37.1 ms / 1.35%

Single DiT Forward Profile

Metric	Baseline	Local Fused SiluAndMul	Improvement
GPU wall time	2,575.64 ms	2,546.81 ms	28.83 ms / 1.12%
CUDA kernel count	1,969	1,920	49 / 2.49%
Total CUDA kernel duration	2,571.83 ms	2,543.15 ms	28.68 ms / 1.12%
MLP activation	57.57 ms / 98 kernels	16.26 ms / 49 kernels	41.32 ms / 49 kernels

The profile confirms that every MiniMax H3 transformer layer dispatches to one vllm::act_and_mul_kernel instead of separate native SiLU and multiplication kernels. NCCL SendRecv time was 18.85 ms higher in the optimized profile due to run-to-run variation, partially masking the activation improvement in the overall DiT forward time.

Accuracy

The optimized 50-step output is byte-for-byte identical to the baseline. Both files have the SHA-256 digest bc383673714d4252467549335f01f92e52f26a4376a2f41aefd4e540c3c2021b.

[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.

[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.

[hsliuustc0106]

@fan2956 does it apply to npu as well?

[hsliuustc0106]

Verified from the code: as written, NPU keeps the native two-kernel path (forward_npu → forward_native), so Ascend behavior is unchanged — but a fused op already exists in-repo (torch_npu.npu_swiglu; see my comment on forward_npu) if you want the fusion to apply there too.

[hsliuustc0106]

NPU is left on the two-kernel native path even though a fused operator with this exact packed contract already exists in-repo: torch_npu.npu_swiglu(gate_up, dim=-1) (wrapped at vllm_omni/platforms/npu/models/minimax_h3.py:57, covered by test_minimax_h3_qwen3vl_swiglu.py at 2e-2 tolerance). RMSNorm.forward_npu (vllm_omni/diffusion/layers/norm.py:125) sets the lazy-import precedent:

def forward_npu(self, x):
    import torch_npu
    return torch_npu.npu_swiglu(x, dim=-1)

With that the DiT MLP gets the same one-kernel activation on Ascend and the PR goal ("avoid the native two-kernel fallback") holds on NPU too.

[hsliuustc0106]

Nit: self.op is resolved eagerly when the platform is cuda/rocm/musa, but the musa path never uses it (forward_musa → forward_cuda → x.device.type != "cuda" → native), and construction would crash on a musa build that doesn't register _C.silu_and_mul. Resolving the op lazily inside forward_cuda (or dropping is_musa() from the condition) avoids both.

[hsliuustc0106]

The byte-identical accuracy result surprises me: torch.ops._C.silu_and_mul computes silu in FP32 and rounds once to BF16, while the baseline F.silu(gate) * up materializes silu as BF16 before the multiply (double rounding), so I'd expect ULP-level divergence to propagate through 49 steps into the output bytes — the modulation PR (#6281) measures PSNR 39.55 dB from exactly this class of rounding change. Can you confirm how the comparison was produced (fused op active in both runs? same checkpoint/seed as the baseline)?

[hsliuustc0106]

Both amd-ci (build 10690) and npu-ci (build 5498) are red on head while the CUDA/Intel lanes pass; logs are behind auth on my side. The npu lane doesn't run MiniMax H3 and the NPU path here is a pure fallback to the old eager math, so I can't correlate from the diff — could you check whether these are related or infra?