[vLLM IR][RMSNorm] Port Mixer2RMSNormGated to vLLM IR Ops

[wxsIcey]

Purpose

Register Mixer2RMSNormGated as a vllm IR op and rewrite Mixer2RMSNormGated.forward_native to dispatch correctly across all tensor-parallel configurations.

The implementation handles four cases:

Case	Condition	Issue	Solution
1	n_groups=1, tp_size>1	Variance must be computed across all ranks (one global norm group, each rank holds only a slice)	AllReduce local sum-of-squares → compute global variance
2	n_groups=1, tp_size=1	No TP, local data is complete	Use IR op directly
3	n_groups>1, n_groups % tp_size != 0	Group boundaries straddle rank boundaries (a rank may hold half a group), local norm is incorrect	AllGather full tensor → normalize locally → slice back to local rank
4	n_groups>1, n_groups % tp_size == 0	Each rank holds an integer number of complete groups, variance can be computed independently	Use IR op directly

Cases 2 and 4 require no collective communication and are handled by the IR op. Cases 1 and 3 require cross-rank communication that cannot be fused into a single kernel, so they are handled with explicit AllReduce / AllGather before calling into local computation.

Because forward_native now covers all cases (including the optimized IR op paths for cases 2 and 4), forward_cuda is fully redundant and can be removed.

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces the mixer2_rms_norm_gated IR operator and its Triton implementation, refactoring the Mamba Mixer2 layer to utilize this new operator. The changes include updating kernel configurations, platform-specific priorities, and the underlying Triton kernel for gated layer normalization. A potential TypeError was identified in the native implementation of mixer2_rms_norm_gated when the weight is None, and a code suggestion was provided to handle this case.

[wxsIcey]

This pr has the same issue as #38798: when using wrap_triton rather than custom op, garbled characters are output. However, I found that setting enforce_eager=True resulted in normal output. It seems this issue is related to torch.compile().

I hope to get some help. cc@zou3519

[ProExpertProg]

@wxsIcey I think you're missing run_functional_passes=False in the lowering pass I think - for some reason that flag removes the triton kernel from the replacement if set to True. See

vllm/vllm/compilation/passes/ir/lowering_pass.py

Lines 95 to 97 in 06de5e1

	match.replace_by_example(
	ir_op_impl.impl_fn, bound_args.args, run_functional_passes=False
	)

[wxsIcey]

@wxsIcey I think you're missing run_functional_passes=False in the lowering pass I think - for some reason that flag removes the triton kernel from the replacement if set to True. See

vllm/vllm/compilation/passes/ir/lowering_pass.py

Lines 95 to 97 in 06de5e1

	match.replace_by_example(
	ir_op_impl.impl_fn, bound_args.args, run_functional_passes=False
	)

It seems to have no effect.

[tomeras91]

Left a nit comment
Also - do you plan to post benchmark results before/after this change? I understand we don't expect any perf diff (IR ops still go through torch.compile), but would like to verify that since this PR changes the code path significantly..

[tomeras91]

nit: docstring should have vllm.ir.ops.mixer2_rms_norm_gated instead of vllm.ir.ops.rms_norm_gated.
(Or change the op name to rms_norm_gated)

[wxsIcey]

Thanks for your review. I will change it.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @wxsIcey.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[wxsIcey]

Left a nit comment Also - do you plan to post benchmark results before/after this change? I understand we don't expect any perf diff (IR ops still go through torch.compile), but would like to verify that since this PR changes the code path significantly..

Thanks for the review. This PR is currently on low priority. I will add benchmark tests once it's ready.

[chaojun-zhang]

@wxsIcey I added this IR on XPU, please review wxsIcey#12

[wxsIcey]

@wxsIcey I added this IR on XPU, please review wxsIcey#12

Thanks for your work, I will merge it.

[chaojun-zhang]

@wxsIcey I added this IR on XPU, please review wxsIcey#12

Thanks for your work, I will merge it.

@wxsIcey, I found that the Triton provider fails to pass IR lowering. I attempted to fix this by registering the Triton kernel as a torch custom op. Please take a look wxsIcey#14, thanks

[wxsIcey]

@wxsIcey I added this IR on XPU, please review wxsIcey#12

Thanks for your work, I will merge it.

@wxsIcey, I found that the Triton provider fails to pass IR lowering. I attempted to fix this by registering the Triton kernel as a torch custom op. Please take a look wxsIcey#14, thanks

This is a known issue. make_fx does not handle the triton operator correctly. You can see the discussion in #38798. We need to figure out why wrap_trion cannot be used.