[Graph][Fusion] Add MatmulAllReduceAddRMSNorm graph fusion for npugraph_ex.

[ForBetterCodeNine]

What this PR does / why we need it?

This PR builds upon PR #5011 and aims to further enhance the npu_graph_ex_passes module. Based on prior work, we have added graph optimization support for the add_rms_quant fused operator in scenarios where a bias term is present—ensuring the fusion pattern is correctly registered and matched into the computation graph.

This time, we performed the operator fusion of MatmulAllReduceAddRMSNorm and added corresponding ST test cases for regression monitoring.

Does this PR introduce any user-facing change?

No

How was this patch tested?

• vLLM version: v0.13.0
• vLLM main: vllm-project/vllm@2c24bc6

[gemini-code-assist]

Code Review

This pull request introduces a new GraphEX fusion pass for Matmul + AllReduce + AddRMSNorm and an end-to-end test for the fused operator. While the initiative is good, the implementation contains several critical issues that must be addressed. The test file has a bug in tensor shape initialization and an unused parameter. The fusion pass file has multiple critical bugs, including a typo, incorrect argument passing, a missing base class inheritance, and a potential correctness issue in the replacement logic. Please review the detailed comments for fixes.

[gemini-code-assist]

The shape of tensor x2 is incorrect. For torch.nn.functional.linear(a, b) where a is x1 with shape [m, k], the weight b (x2 here) should have shape [n, k]. The current code initializes x2 with shape [m, k]. This will cause a shape mismatch error during the npu_add_rms_norm operation, as the residual tensor has shape [m, n], while the output of linear will have an incorrect shape.

Suggested change

	x2 = torch.normal(0, 0.1, [m, k], dtype=DTYPE).npu(rank)
	x2 = torch.normal(0, 0.1, [n, k], dtype=DTYPE).npu(rank)

[gemini-code-assist]

The is_allgather_add_out parameter is hardcoded to False. The pattern being replaced produces a full tensor for the second output (add_out), as it's derived from an all-reduced tensor. If is_allgather_add_out=False causes the fused operator to return a sharded add_out, this would be a correctness issue, as subsequent layers expect a full tensor. The corresponding test for this operator uses is_allgather_add_out=True. It should likely be True here as well to ensure the replacement is correct for middle layers.

Suggested change

	False,
	True,

[gemini-code-assist]

There is a typo torch.rann. It should be torch.randn. This will cause an AttributeError when get_inputs is called.

Suggested change

	residual = torch.rann(batch_size, seq_len, hidden_size, device="npu")
	residual = torch.randn(batch_size, seq_len, hidden_size, device="npu")

[gemini-code-assist]

The example_inputs argument to torchair.register_replacement expects a callable, but it's being passed the result of self.get_inputs(). This will cause an error during pattern registration. It should be self.get_inputs to pass the method itself.

Suggested change

	example_inputs=self.get_inputs(),
	example_inputs=self.get_inputs,

[gemini-code-assist]

The class GraphEXMatmulAllReduceAddRMSNormPass calls super().__init__ but does not inherit from any base class, which will cause a TypeError. Based on its usage in GraphFusionPassManager, it should inherit from VllmInductorPass. You'll also need to add from vllm.compilation.vllm_inductor_pass import VllmInductorPass at the top of the file.

Suggested change

	class GraphEXMatmulAllReduceAddRMSNormPass:
	def __init__(self, vllm_config: VllmConfig):
	super().__init__(vllm_config)
	from vllm.compilation.vllm_inductor_pass import VllmInductorPass
	class GraphEXMatmulAllReduceAddRMSNormPass(VllmInductorPass):
	def __init__(self, vllm_config: VllmConfig):
	super().__init__(vllm_config)

[gemini-code-assist]

The batch_size parameter is defined but never used within the worker function. It should be removed from the function signature to improve code clarity. The call to mp.spawn on line 134 and the args tuple on line 133 should also be updated accordingly.

Suggested change

	def worker(rank, ep_world_size, batch_size, m, k, n):
	def worker(rank, ep_world_size, m, k, n):

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[wxsIcey]

Is this test related to graph_ex?

[wxsIcey]

This change is caused by the introduction of custom operators. As a result, AddRmsNorm in the graph is replaced with AddRmsNormBias.