From 88f172188b8f9862791eec97c883b274973bb7ca Mon Sep 17 00:00:00 2001 From: pschlan-amd Date: Tue, 2 Jun 2026 16:50:21 +0200 Subject: [PATCH] [ROCm] Fix AITER RMSNormQuantFusion for Kimi-Linear (#44308) Signed-off-by: Patrick Schlangen --- vllm/compilation/passes/fusion/rocm_aiter_fusion.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py index e7ba338572..03d291d4d9 100644 --- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py +++ b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py @@ -570,9 +570,16 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass): ) gated_norm_shapes: set[tuple[int, int]] = set() for layer in gdn_layers.values(): - gated_norm_shapes.add( - (layer.num_v_heads // layer.tp_size, layer.head_v_dim) + num_v_heads = getattr(layer, "num_v_heads", None) or getattr( + layer, "num_heads", None ) + head_v_dim = getattr(layer, "head_v_dim", None) or getattr( + layer, "head_dim", None + ) + + assert num_v_heads is not None and head_v_dim is not None + + gated_norm_shapes.add((num_v_heads // layer.tp_size, head_v_dim)) # Make sure fused add patterns are before simple rms norm, # as the latter is a subset of the former in torch ops.