A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 04:45 UTC
[ROCm] Fix AITER RMSNormQuantFusion for Kimi-Linear (#44308)
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
pschlan-amd committed
Jun 2, 2026 at 16:50 UTC
88f172188b8f9862791eec97c883b274973bb7ca
1 file changed
+9
-2
vllm/compilation/passes/fusion/rocm_aiter_fusion.py
+9
-2
@@ -570,9 +570,16 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
570
)
571
gated_norm_shapes: set[tuple[int, int]] = set()
572
for layer in gdn_layers.values():
573
- gated_norm_shapes.add(
574
- (layer.num_v_heads // layer.tp_size, layer.head_v_dim)
573
+ num_v_heads = getattr(layer, "num_v_heads", None) or getattr(
574
+ layer, "num_heads", None
575
)
576
+ head_v_dim = getattr(layer, "head_v_dim", None) or getattr(
577
+ layer, "head_dim", None
578
+ )
579
+
580
+ assert num_v_heads is not None and head_v_dim is not None
581
+
582
+ gated_norm_shapes.add((num_v_heads // layer.tp_size, head_v_dim))
583
584
# Make sure fused add patterns are before simple rms norm,
585
# as the latter is a subset of the former in torch ops.