[rbrugaro-amd]

Summary

Add a torch.inductor post-grad pattern-matcher pass that fuses the paired q_a_layernorm and kv_a_layernorm RMS norm operations in DeepSeek-V3 / Kimi-K2 MLA attention into a single fused_qk_rmsnorm HIP kernel call via AITer.

In the unfused graph, each MLA layer runs two separate rms_norm calls on the q and kv compressed latents:

qkv_lora = gemm(hidden)
q_c, kv_lora = split(qkv_lora, [1536, 576])
kv_c, k_pe   = split(kv_lora,  [512,  64])
q_c  = rms_norm(q_c,  q_weight,  eps)   # norm #1
kv_c = rms_norm(kv_c, kv_weight, eps)   # norm #2

The pass uses torch._inductor.pattern_matcher.register_replacement to match the connected subgraph — both splits and both norms — and rewrites it to:

q_c, kv_lora = split(qkv_lora, [1536, 576])
kv_c, k_pe   = split(kv_lora,  [512,  64])
q_normed, kv_normed = fused_mla_dual_rms_norm(
    q_c, q_weight, kv_c, kv_weight, eps1, eps2)

This reduces kernel launch overhead from 2 launches to 1 per MLA layer (61 layers in Kimi-K2).

Changes

• vllm/_aiter_ops.py — register fused_mla_dual_rms_norm custom op wrapping AITer's fused_qk_rmsnorm HIP kernel
• vllm/compilation/passes/fusion/rocm_aiter_fusion.py — MLADualRMSNormPattern + MLADualRMSNormFusionPass using register_replacement pattern matching
• vllm/config/compilation.py / vllm/config/vllm.py — fuse_mla_dual_rms_norm PassConfig flag (ROCm-only, auto-enabled when AITer is available at O1+)
• vllm/compilation/passes/pass_manager.py — register the pass
• tests/compile/passes/test_fuse_mla_dual_rms_norm.py — unit test verifying pattern match, op replacement, and numerical correctness

Enable

Enabled by default when AITer is available (optimization level ≥ O1). Can also be toggled explicitly:

--compilation-config='{"pass_config": {"fuse_mla_dual_rms_norm": true}}'

AITer dependency

Requires AITer PR #2442 (fused_qk_rmsnorm HIP kernel, commit 149a8f613).

Accuracy

Benchmark	Metric	Score	Threshold	Status
GSM8K (5-shot, 250 samples)	exact_match (flexible-extract)	0.94	0.90	✅ PASS
GSM8K (5-shot, 250 samples)	exact_match (strict-match)	0.94	0.90	✅ PASS

Model: Kimi-K2-Thinking-MXFP4, TP=4, MI355X, fp8 KV cache

Trace showing fused kernel

Two separate RMSNorm CK kernels _ZN7ck_tile6kentryILi1ENS_12Rmsnorm2dFwdINS_27Rmsnorm2dFwdPipelineOnePassINS_27Rmsnorm2dFwdPipelineProblemIDF replaced with single AITER kernel ZN5aiter23fused_qk_rmsnorm_kernel

Performance

Kimi-K2-Thinking-MXFP4, TP=4, MI355X

Config	Concurrency	Baseline (tok/s/GPU)	QK RMSNorm Fusion (tok/s/GPU)	Improvement
1k/1k	4	190	194	1.02x
1k/1k	8	321	326	1.01x
1k/1k	16	521	533	1.02x
1k/1k	32	783	801	1.02x
1k/1k	64	1173	1177	1.00x
8k/1k	4	799	819	1.03x
8k/1k	8	1348	1367	1.01x
8k/1k	16	2051	2074	1.01x
8k/1k	32	2917	2975	1.02x
8k/1k	64	4038	4026	1.00x
			Geomean uplift	1.02x

Test plan

• Unit test: pattern fires, ops replaced, numerical correctness (test_fuse_mla_dual_rms_norm)
• Pass fires on all TP workers: MLADualRMSNormFusionPass: fused 1 q/kv norm pair(s)
• FX graph dumps confirm before/after rewrite
• GSM8K accuracy check passes (0.94 >= 0.90 threshold)
• Throughput benchmark (see table above)

Related

• AITer fused_qk_rmsnorm HIP kernel: ROCm/aiter#2442

[gemini-code-assist]

Code Review

This pull request introduces the MLADualRMSNormFusionPass to fuse paired RMS norms in MLA attention for DeepSeek-V3 and Kimi-K2 models, utilizing the aiter fused kernel on ROCm platforms. The changes include the implementation of the FX graph pass, custom operator registration, and integration into the pass manager and configuration system. Feedback indicates a potential topological ordering issue in the graph rewrite logic that could lead to compilation failures if input nodes are not correctly hoisted.

[gemini-code-assist]

The current hoisting logic only ensures that kv_split and kv_c are placed after kv_lora. However, q_c is also an input to the fused node (inserted after kv_c). Since q_c and kv_lora are siblings (both users of q_split), their relative order in the FX graph's linked list is not guaranteed. If q_c appears after kv_c in the linked list, the fused node will be inserted before its input q_c, violating topological order and causing compilation failures. You should also hoist q_c to ensure it precedes the fused operation.

        kv_split = kv_c.args[0]
        kv_lora = kv_split.args[0]
        q_split = q_c.args[0]

        # Hoist nodes to ensure topological order:
        # q_split -> q_c -> kv_lora -> kv_split -> kv_c
        q_split.append(q_c)
        q_c.append(kv_lora)
        kv_lora.append(kv_split)
        kv_split.append(kv_c)

[rbrugaro-amd]

Thanks for the input @Rohan138 I've reimplemented with pattern matching instead and moved it to the existing files as you suggested.

[ProExpertProg]

Looks good overall, just a few nits!

[ProExpertProg]

Nit:

Suggested change

	result = torch.ops.vllm.fused_mla_dual_rms_norm(
	q_c, q_weight, kv_c, kv_weight, eps, eps,
	)
	return result[0], result[1], k_pe
	q_normed, kv_normed = torch.ops.vllm.fused_mla_dual_rms_norm(
	q_c, q_weight, kv_c, kv_weight, eps, eps,
	)
	return q_normed, kv_normed, k_pe

[ProExpertProg]

Let's use the new VllmFusionPatternMatcherPass please - will make it cleaner!

[ProExpertProg]

This seems enough for now

[rbrugaro-amd]

thx for reviewing @ProExpertProg! I've addressed the comments, please take a look.

[mergify]

Hi @rbrugaro-amd, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @rbrugaro-amd, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10