[Model][Hardware][AMD]: Part 1/2 -> Enable e2e QK Norm + RoPE + KV Cache runtime fusion for Qwen3-30B-A3B on ROCM_AITER_FA, and ROCM_AITER_UNIFIED_ATTN

vllm/_aiter_ops.py

@@ -2128,6 +2128,63 @@ def triton_fp4_gemm_dynamic_quant(
         gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)
         return y
 
+    @staticmethod
+    def fused_qk_norm_rope_and_cache(
+        qkv: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        positions: torch.Tensor,
+        num_heads_q: int,
+        num_heads_k: int,
+        num_heads_v: int,
+        head_dim: int,
+        is_neox: bool,
+        rms_norm_eps: float,
+        q_out: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        k_scale: torch.Tensor,
+        v_scale: torch.Tensor,
+        k_out: torch.Tensor | None,
+        v_out: torch.Tensor | None,
+        return_kv: bool,
+        use_shuffle_layout: bool,
+        block_size: int,
+        x: int,
+    ):
+        from aiter.ops.fused_qk_norm_rope_cache_quant import (
+            fused_qk_norm_rope_cache_pts_quant_shuffle,
+        )
+
+        fused_qk_norm_rope_cache_pts_quant_shuffle(
+            qkv,
+            q_weight,
+            k_weight,
+            cos_sin_cache,
+            positions,
+            qkv.size(0),
+            num_heads_q,
+            num_heads_k,
+            num_heads_v,
+            head_dim,
+            is_neox,
+            rms_norm_eps,
+            q_out,
+            k_cache,
+            v_cache,
+            slot_mapping,
+            k_scale,
+            v_scale,
+            k_out,
+            v_out,
+            return_kv,
+            use_shuffle_layout,
+            block_size,
+            x,
+        )
+
     @staticmethod
     def triton_rope_and_cache(
         query: torch.Tensor,

vllm/compilation/passes/fusion/qk_norm_rope_fusion.py

@@ -25,6 +25,11 @@
 
 FUSED_QK_ROPE_OP = torch.ops._C.fused_qk_norm_rope.default
 
+# Head dimensions supported by csrc/fused_qknorm_rope_kernel.cu's
+# launchFusedQKNormRope and launchFusedQKNormRopeNTokenHeads dispatchers.
+# Keep in sync with the switch statements in that file.
+SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS: tuple[int, ...] = (64, 128, 256)
+
 P = ParamSpec("P")
 
 
@@ -67,6 +72,9 @@ def __init__(
         self.eps = eps
         self.is_neox = is_neox
         self.rope_flashinfer = rope_flashinfer
+        # match_rocm_aiter is auto-detected by MatcherRotaryEmbedding via
+        # rocm_aiter_ops.is_triton_rotary_embed_enabled(), so no need to
+        # enumerate AITER variants here.
         self.rope_matcher = MatcherRotaryEmbedding(
             is_neox=is_neox,
             head_size=self.head_dim,
@@ -186,7 +194,12 @@ def replacement(
 
 
 class QKNormRoPEFusionPass(VllmPatternMatcherPass):
-    """Fuse Q/K RMSNorm + RoPE into fused_qk_norm_rope when the custom op exists."""
+    """Fuse Q/K RMSNorm + RoPE into fused_qk_norm_rope when the custom op exists.
+
+    Registers patterns for both standard vLLM ops and ROCm AITER ops
+    (when AITER is enabled), so the fusion fires regardless of which
+    RMSNorm/RoPE implementation the graph uses.
+    """
 
     @enable_fake_mode
     def __init__(self, config: VllmConfig) -> None:
@@ -202,7 +215,6 @@ def __init__(self, config: VllmConfig) -> None:
             )
             return
 
-        # use one attn layer to get meta (such as head_dim) for QkNormRopePattern
         attn_layers: dict[str, Attention] = get_layers_from_vllm_config(
             config, Attention
         )
@@ -213,6 +225,20 @@ def __init__(self, config: VllmConfig) -> None:
             return
         layer = next(iter(attn_layers.values()))
 
+        if layer.head_size not in SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS:
+            logger.warning_once(
+                "QK Norm+RoPE fusion not enabled: layer head_size=%d is not "
+                "supported by fused_qk_norm_rope kernel (supported: %s). "
+                "Falling back to unfused QK norm + RoPE path.",
+                layer.head_size,
+                SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS,
+            )
+            return
+
+        # RMS norm variants are no longer iterated: after the vLLM IR migration (#33825)
+        # AITER rope variants are also not iterated: `MatcherRotaryEmbedding`
+        # auto-detects via `rocm_aiter_ops.is_triton_rotary_embed_enabled()`
+        # and selects the right rotary op.
         for epsilon in [1e-5, 1e-6]:
             for neox in [True, False]:
                 if RotaryEmbedding.enabled():