[ROCm][Perf] Extend QK-norm/RoPE/KV-cache fusion to MRoPE (#50212)

Signed-off-by: Vorapol Assavasangthong <Vorapol.Assavasangthong@amd.com>
Co-authored-by: Santosh Hiremath <Santosh.Hiremath@amd.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>

tests/compile/fusions_e2e/common.py

@@ -128,8 +128,11 @@ FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {
     ),
     "rms_quant_fusion": re.compile(r"rms_quant_fusion.py:\d+] Replaced (\d+) patterns"),
     "act_quant_fusion": re.compile(r"act_quant_fusion.py:\d+] Replaced (\d+) patterns"),
+    # The combined pass can consume these sites before the standalone pass.
     "norm_rope_fusion": re.compile(
-        r"qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on (\d+) sites"
+        r"(?:qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on|"
+        r"qk_norm_rope_kvcache_fusion.py:\d+] "
+        r"QK-Norm\+RoPE(?:/MRoPE)?\+KVCache fusion: replaced) (\d+)"
     ),
     "attn_quant_fusion": re.compile(
         r"attn_quant_fusion.py:\d+] Fused quant onto (\d+) attention nodes"

tests/compile/fusions_e2e/conftest.py

@@ -182,7 +182,11 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
             # TODO: Remove log counting in unit tests
             # once all matchers implement VllmFusionPatternMatcherPass
             n_expected = tp_size * num_ranges_activated
-            if match_name not in ("attn_quant_fusion", "act_quant_fusion"):
+            if match_name not in (
+                "attn_quant_fusion",
+                "act_quant_fusion",
+                "norm_rope_fusion",
+            ):
                 assert len(log_matches) == n_expected, (
                     f"Could not find {n_expected} {match_name} "
                     f"(found {len(log_matches)}) in:\n {log_holder.text}"
@@ -190,7 +194,16 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
 
             expected_matches = getattr(matches, match_name)
 
-            if match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:
+            if match_name == "norm_rope_fusion":
+                # Opaque LayerName lets the combined pass consume short-range
+                # sites before the standalone pass; count either owner.
+                assert len(log_matches) >= n_expected
+                assert all(m in (0, expected_matches) for m in log_matches)
+                assert sum(log_matches) == expected_matches * n_expected, (
+                    f"{match_name} expected {expected_matches * n_expected} "
+                    f"sites across both passes, found: {log_matches}"
+                )
+            elif match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:
                 # AR+rms+quant takes precedence over rms+quant if activated.
                 # That means we get full matching where ar+rms+quant was not
                 # activated, and less where it was (only the smallest range).

tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py

@@ -8,6 +8,7 @@ import pytest
 import torch
 from packaging.version import Version
 
+import vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion as fusion_module
 import vllm.config
 from tests.compile.backend import TestBackend
 from tests.v1.attention.utils import (
@@ -15,8 +16,12 @@ from tests.v1.attention.utils import (
     create_common_attn_metadata,
     dense_kv_cache_views,
 )
-from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
-from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
+from vllm._aiter_ops import (
+    IS_AITER_FOUND,
+    is_aiter_found_and_supported,
+    rocm_aiter_ops,
+)
+from vllm.compilation.passes.fusion.matcher_utils import MROPE_OP, ROTARY_OP
 from vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion import (
     QkNormRopeKvCacheFusionPass,
 )
@@ -37,8 +42,12 @@ from vllm.config import (
 from vllm.forward_context import get_forward_context, set_forward_context
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import _USE_LAYERNAME, _encode_layer_name
 from vllm.v1.attention.backend import (
     AttentionBackend,
     CommonAttentionMetadata,
@@ -94,6 +103,7 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
         device: torch.device,
         rotary_dim: int | None = None,
         prefix: str = "model.layers.0.self_attn.attn",
+        expected_query_quant_group_shape: GroupShape | None = None,
     ):
         super().__init__()
         self.num_heads = num_heads
@@ -133,6 +143,9 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
             attn_backend=attn_backend.get_class(),
         )
         self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()
+        if expected_query_quant_group_shape is not None:
+            assert self.attn.query_quant is not None
+            assert self.attn.query_quant.group_shape == expected_query_quant_group_shape
         assert not self.attn_backend.forward_includes_kv_cache_update, (
             f"Attention backend {self.attn_backend} does not support "
             "fuse_qk_norm_rope_kvcache."
@@ -227,18 +240,15 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
             self.kv_cache_dtype != self.dtype
             and self.attn.impl.supports_quant_query_input
         ):
-            q_fp8 = torch.empty_like(q, dtype=FP8_DTYPE)
-            torch.ops.vllm.rocm_aiter_per_tensor_quant(
-                q_fp8, q, self.attn._q_scale, False
-            )
-            q = q_fp8
+            assert self.attn.query_quant is not None
+            q, _ = self.attn.query_quant(q, self.attn._q_scale)
 
         # Final views + KV cache update
         q = q.view(-1, self.num_heads, self.head_size)
         k = k.view(-1, self.num_kv_heads, self.head_size)
         v = v.view(-1, self.num_kv_heads, self.head_size)
         kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
-            k, v, self.layer_name
+            k, v, _encode_layer_name(self.layer_name)
         )
         return q, k, v, kv_cache_dummy_dep
 
@@ -260,6 +270,64 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
         return [torch.ops.vllm.fused_qk_norm_rope_and_unified_kv_cache_update.default]
 
 
+class QKNormMRoPEKVCacheTestModel(QKNormRoPEKVCacheTestModel):
+    def __init__(
+        self,
+        *,
+        mrope_section: tuple[int, int, int],
+        mrope_interleaved: bool,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.rotary_emb = MRotaryEmbedding(
+            self.head_size,
+            rotary_dim=self.rotary_dim,
+            max_position_embeddings=4096,
+            base=10000,
+            is_neox_style=self.is_neox,
+            dtype=self.dtype,
+            mrope_section=list(mrope_section),
+            mrope_interleaved=mrope_interleaved,
+        )
+        self.enable_rope_custom_op = self.rotary_emb.enabled()
+
+    def forward(
+        self, qkv: torch.Tensor, positions: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        q, k, v, kv_cache_dummy_dep = super().forward(qkv, positions)
+        output = torch.empty(
+            q.shape[0],
+            self.num_heads * self.head_size,
+            device=q.device,
+            dtype=self.dtype,
+        )
+        torch.ops.vllm.unified_attention_with_output(
+            q,
+            k,
+            v,
+            output,
+            _encode_layer_name(self.layer_name),
+            kv_cache_dummy_dep=kv_cache_dummy_dep,
+        )
+        # Production decoder attention consumes K from the cache rather than
+        # its direct K argument. Returning `output` keeps that production user
+        # in the graph without exposing K to an unsupported consumer.
+        return q, output, v, kv_cache_dummy_dep
+
+    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
+        return [MROPE_OP, torch.ops.vllm.unified_kv_cache_update.default]
+
+    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
+        return [torch.ops.vllm.fused_qk_norm_mrope_and_unified_kv_cache_update.default]
+
+
+class QKNormMRoPEObservableKTestModel(QKNormMRoPEKVCacheTestModel):
+    def forward(
+        self, qkv: torch.Tensor, positions: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        return QKNormRoPEKVCacheTestModel.forward(self, qkv, positions)
+
+
 def _run_qk_norm_rope_kvcache_fusion_test(
     *,
     attn_backend: AttentionBackendEnum,
@@ -278,18 +346,33 @@ def _run_qk_norm_rope_kvcache_fusion_test(
     rms_norm_eps: float,
     custom_op: str,
     monkeypatch: pytest.MonkeyPatch,
+    mrope_section: tuple[int, int, int] | None = None,
+    mrope_interleaved: bool = False,
+    observe_mrope_k: bool = False,
+    expect_fusion: bool = True,
+    use_quark_scalar_query_scale: bool = False,
 ) -> None:
     device = os.environ.get("VLLM_TEST_CUDA_DEVICE", "cuda")
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
     torch.manual_seed(0)
 
+    # Set AITER toggles before VllmConfig initialization enables env caching.
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1")
+    rocm_aiter_ops.refresh_env_variables()
+
     vllm_config = VllmConfig(
         model_config=ModelConfig(dtype=dtype),
         cache_config=CacheConfig(
             block_size=block_size,
             cache_dtype=kv_cache_dtype,
         ),
+        quant_config=(
+            QuarkConfig(quant_config={"exclude": []})
+            if use_quark_scalar_query_scale
+            else None
+        ),
         compilation_config=CompilationConfig(
             mode=CompilationMode.VLLM_COMPILE,
             custom_ops=[custom_op],
@@ -302,6 +385,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
 
     with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
         m.setenv("VLLM_ROCM_USE_AITER", "1")
+        m.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1")
         m.setenv(
             "VLLM_ROCM_USE_AITER_TRITON_ROPE",
             "1" if enable_aiter_triton_rope else "0",
@@ -309,18 +393,49 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         m.setenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", use_shuffle_kv_layout)
         rocm_aiter_ops.refresh_env_variables()
 
-        model = QKNormRoPEKVCacheTestModel(
-            vllm_config=vllm_config,
-            attn_backend=attn_backend,
-            num_heads=num_heads,
-            num_kv_heads=num_kv_heads,
-            head_size=head_size,
-            rotary_dim=rotary_dim,
-            is_neox=is_neox,
-            rms_norm_eps=rms_norm_eps,
-            dtype=dtype,
-            device=torch.get_default_device(),
-        )
+        expected_query_quant_group_shape = None
+        if use_quark_scalar_query_scale:
+            expected_query_quant_group_shape = (
+                GroupShape(-1, num_heads * head_size)
+                if num_kv_heads == 1
+                else GroupShape.PER_TENSOR
+            )
+
+        model_kwargs = {
+            "vllm_config": vllm_config,
+            "attn_backend": attn_backend,
+            "num_heads": num_heads,
+            "num_kv_heads": num_kv_heads,
+            "head_size": head_size,
+            "rotary_dim": rotary_dim,
+            "is_neox": is_neox,
+            "rms_norm_eps": rms_norm_eps,
+            "dtype": dtype,
+            "device": torch.get_default_device(),
+            "expected_query_quant_group_shape": expected_query_quant_group_shape,
+        }
+        if mrope_section is None:
+            model = QKNormRoPEKVCacheTestModel(**model_kwargs)
+        else:
+            alternate_section = (8, 8, 8)
+            m.setattr(
+                fusion_module,
+                "_discover_mrope_configs",
+                lambda _: (
+                    (mrope_section, mrope_interleaved),
+                    (alternate_section, not mrope_interleaved),
+                ),
+            )
+            model_cls = (
+                QKNormMRoPEObservableKTestModel
+                if observe_mrope_k
+                else QKNormMRoPEKVCacheTestModel
+            )
+            model = model_cls(
+                **model_kwargs,
+                mrope_section=mrope_section,
+                mrope_interleaved=mrope_interleaved,
+            )
 
         fusion_pass = QkNormRopeKvCacheFusionPass(vllm_config)
         passes = [
@@ -337,7 +452,12 @@ def _run_qk_norm_rope_kvcache_fusion_test(
             num_heads * head_size + 2 * num_kv_heads * head_size,
             dtype=dtype,
         )
-        pos = torch.arange(num_tokens, dtype=torch.long)
+        if mrope_section is None:
+            pos = torch.arange(num_tokens, dtype=torch.long)
+        else:
+            # Real Qwen-VL position IDs can be a strided [3, N] view.
+            pos = torch.randint(0, 4096, (3, 2 * num_tokens), dtype=torch.long)[:, ::2]
+            assert not pos.is_contiguous()
 
         qkv_unfused = qkv.clone()
         pos_unfused = pos.clone()
@@ -356,7 +476,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
 
         # Run fused (compiled) forward
         torch._dynamo.mark_dynamic(qkv, 0)
-        torch._dynamo.mark_dynamic(pos, 0)
+        torch._dynamo.mark_dynamic(pos, 0 if mrope_section is None else 1)
         with set_forward_context(None, vllm_config):
             model_fused = torch.compile(model, backend=backend)
             forward_context = get_forward_context()
@@ -369,14 +489,33 @@ def _run_qk_norm_rope_kvcache_fusion_test(
             kv_cache_fused = attn_layer.kv_cache
         del dummy
 
-        assert fusion_pass.matched_count == 1
+        assert fusion_pass.matched_count == int(expect_fusion)
 
-        backend.check_before_ops(model.ops_in_model_before())
-        backend.check_after_ops(model.ops_in_model_after())
+        if expect_fusion:
+            backend.check_before_ops(model.ops_in_model_before())
+            backend.check_after_ops(model.ops_in_model_after())
+        else:
+            for op in model.ops_in_model_before():
+                assert backend.op_count(op, before=True) > 0
+                assert backend.op_count(op) > 0
+
+        if use_quark_scalar_query_scale:
+            expected_quant_op = (
+                torch.ops.vllm.rocm_aiter_per_tensor_quant.default
+                if num_kv_heads != 1
+                else torch.ops._C.static_scaled_fp8_quant.default
+            )
+            assert backend.op_count(expected_quant_op, before=True) > 0
+            # Query quantization remains separate after the prologue fusion.
+            assert backend.op_count(expected_quant_op) > 0
 
         # Sweep-backed (18.2k pts, PR #42749): native-rope ref worst 7.7e-3 -> 1e-2;
         # AITER-triton-rope ref is itself approximate (plateau 1.28e-2) -> 2e-2.
-        ATOL, RTOL = (2e-2, 2e-2) if enable_aiter_triton_rope else (1e-2, 1e-2)
+        ATOL, RTOL = (
+            (2e-2, 2e-2)
+            if enable_aiter_triton_rope or mrope_section is not None
+            else (1e-2, 1e-2)
+        )
         is_fp8_cache = model.kv_cache_dtype != dtype
 
         if q_fused.dtype == FP8_DTYPE:
@@ -387,7 +526,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         else:
             torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
 
-        if not is_fp8_cache:
+        if not expect_fusion or (not is_fp8_cache and mrope_section is None):
             # The AITER PTS kernel populates k_out only for non-FP8 caches.
             # With FP8, the kernel writes quantized K directly to the cache
             # and may leave k_out uninitialised.  In production this is fine
@@ -397,15 +536,16 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         # Should be bit exact since no processing had been done on v for both paths
         torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0)
 
-        # fp8 vs triton-rope ref requires loosening tolerance to 1.25e-1.
-        if is_fp8_cache and enable_aiter_triton_rope:
+        # Fused and unfused arithmetic can straddle an FP8 rounding boundary.
+        # Allow one E4M3 quantization step while still comparing every block.
+        if is_fp8_cache:
             cache_atol = cache_rtol = 1.25e-1
         else:
             cache_atol, cache_rtol = ATOL, RTOL
 
         torch.testing.assert_close(
-            kv_cache_unfused[0].float(),
-            kv_cache_fused[0].float(),
+            kv_cache_unfused.float(),
+            kv_cache_fused.float(),
             atol=cache_atol,
             rtol=cache_rtol,
         )
@@ -494,3 +634,225 @@ def test_qk_norm_rope_kvcache_fusion(
         custom_op=custom_op,
         monkeypatch=monkeypatch,
     )
+
+
+@pytest.mark.parametrize(
+    "mrope_section",
+    [
+        pytest.param((16, 24, 24), id="full_rotary"),
+        pytest.param((4, 6, 6), id="partial_rotary"),
+    ],
+)
+@pytest.mark.parametrize("mrope_interleaved", [False, True])
+@pytest.mark.parametrize("is_neox", [False, True])
+@pytest.mark.parametrize(
+    "kv_layout",
+    [
+        pytest.param(KVCacheLayout.LBHNC, id="head_major"),
+        pytest.param(KVCacheLayout.LBNHC, id="token_major"),
+        pytest.param(KVCacheLayout.LHBNC, id="block_major"),
+    ],
+)
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+@pytest.mark.skipif(
+    not IS_AITER_FOUND,
+    reason="Requires AITER",
+)
+def test_qk_norm_mrope_kvcache_fusion(
+    mrope_section: tuple[int, int, int],
+    mrope_interleaved: bool,
+    is_neox: bool,
+    kv_layout: KVCacheLayout,
+    kv_cache_dtype: str,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    _run_qk_norm_rope_kvcache_fusion_test(
+        attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        enable_aiter_triton_rope=False,
+        num_tokens=5,
+        num_heads=8,
+        num_kv_heads=2,
+        head_size=128,
+        rotary_dim=2 * sum(mrope_section),
+        block_size=16,
+        is_neox=is_neox,
+        use_shuffle_kv_layout="0",
+        kv_layout=kv_layout,
+        dtype=torch.bfloat16,
+        kv_cache_dtype=kv_cache_dtype,
+        rms_norm_eps=1e-6,
+        custom_op="+rotary_embedding",
+        monkeypatch=monkeypatch,
+        mrope_section=mrope_section,
+        mrope_interleaved=mrope_interleaved,
+    )
+
+
+@pytest.mark.parametrize(
+    ("num_heads", "num_kv_heads"),
+    [(32, 2), (16, 1), (8, 1)],
+    ids=["tp2", "tp4", "tp8"],
+)
+def test_qk_norm_mrope_kvcache_fusion_quark_scalar_query_scale(
+    num_heads: int,
+    num_kv_heads: int,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    _run_qk_norm_rope_kvcache_fusion_test(
+        attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        enable_aiter_triton_rope=False,
+        num_tokens=5,
+        num_heads=num_heads,
+        num_kv_heads=num_kv_heads,
+        head_size=128,
+        rotary_dim=128,
+        block_size=16,
+        is_neox=True,
+        use_shuffle_kv_layout="0",
+        kv_layout=KVCacheLayout.LBHNC,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="fp8",
+        rms_norm_eps=1e-6,
+        custom_op="+rotary_embedding",
+        monkeypatch=monkeypatch,
+        mrope_section=(24, 20, 20),
+        mrope_interleaved=True,
+        use_quark_scalar_query_scale=True,
+    )
+
+
+def test_mrope_fusion_rejects_observable_k(monkeypatch: pytest.MonkeyPatch):
+    _run_qk_norm_rope_kvcache_fusion_test(
+        attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        enable_aiter_triton_rope=False,
+        num_tokens=5,
+        num_heads=8,
+        num_kv_heads=2,
+        head_size=128,
+        rotary_dim=128,
+        block_size=16,
+        is_neox=True,
+        use_shuffle_kv_layout="0",
+        kv_layout=KVCacheLayout.LBHNC,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="auto",
+        rms_norm_eps=1e-6,
+        custom_op="+rotary_embedding",
+        monkeypatch=monkeypatch,
+        mrope_section=(16, 24, 24),
+        observe_mrope_k=True,
+        expect_fusion=False,
+    )
+
+
+def test_mrope_fusion_rejects_unaligned_rotary_dim(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    _run_qk_norm_rope_kvcache_fusion_test(
+        attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        enable_aiter_triton_rope=False,
+        num_tokens=5,
+        num_heads=8,
+        num_kv_heads=2,
+        head_size=128,
+        rotary_dim=6,
+        block_size=16,
+        is_neox=True,
+        use_shuffle_kv_layout="0",
+        kv_layout=KVCacheLayout.LBHNC,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="auto",
+        rms_norm_eps=1e-6,
+        custom_op="+rotary_embedding",
+        monkeypatch=monkeypatch,
+        mrope_section=(1, 1, 1),
+        expect_fusion=False,
+    )
+
+
+@pytest.mark.skipif(not _USE_LAYERNAME, reason="requires opaque LayerName")
+def test_opaque_layer_name_patterns_register_once(monkeypatch: pytest.MonkeyPatch):
+    device = os.environ.get("VLLM_TEST_CUDA_DEVICE", "cuda")
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(dtype=torch.bfloat16),
+        cache_config=CacheConfig(block_size=16, cache_dtype="fp8_e4m3"),
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=["+rotary_embedding"],
+            pass_config=PassConfig(fuse_qk_norm_rope_kvcache=True),
+        ),
+    )
+
+    with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
+        m.setenv("VLLM_ROCM_USE_AITER", "1")
+        rocm_aiter_ops.refresh_env_variables()
+        model_kwargs = {
+            "vllm_config": vllm_config,
+            "attn_backend": AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+            "num_heads": 8,
+            "num_kv_heads": 2,
+            "head_size": 128,
+            "is_neox": True,
+            "rms_norm_eps": 1e-6,
+            "dtype": torch.bfloat16,
+            "device": torch.device(device),
+        }
+        models = [
+            QKNormRoPEKVCacheTestModel(
+                **model_kwargs,
+                prefix=f"model.layers.{index}.self_attn.attn",
+            )
+            for index in range(2)
+        ]
+
+        # Opaque layer names are wildcards in the pattern graph. Registering
+        # once per physical layer raises "Duplicate pattern" on torch >= 2.11.
+        fusion_pass = QkNormRopeKvCacheFusionPass(vllm_config)
+
+    assert models
+    assert fusion_pass.patterns
+
+
+def _call_fused_mrope_impl(q_out: torch.Tensor) -> None:
+    device = q_out.device
+    fusion_module.fused_qk_norm_mrope_and_unified_kv_cache_update_impl(
+        q_out=q_out,
+        qkv=torch.zeros(2, 16, device=device),
+        positions=torch.arange(2, device=device),
+        q_weight=torch.ones(4, device=device),
+        k_weight=torch.ones(4, device=device),
+        rms_norm_eps=1e-6,
+        cos_sin_cache=torch.zeros(16, 4, device=device),
+        is_neox=True,
+        mrope_section_t=1,
+        mrope_section_h=1,
+        mrope_section_w=0,
+        is_interleaved=False,
+        rotary_dim=4,
+        layer_name="model.layers.0.self_attn.attn",
+    )
+
+
+def test_mrope_profiling_pass_initializes_query(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        fusion_module,
+        "get_attention_context",
+        lambda _: (None, None, None, None),
+    )
+    q_out = torch.full((2, 8), torch.nan, device="cpu")
+
+    _call_fused_mrope_impl(q_out)
+
+    torch.testing.assert_close(q_out, torch.zeros_like(q_out))
+
+
+def test_mrope_attention_context_failure_propagates(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    def fail_context(_):
+        raise RuntimeError("missing attention context")
+
+    monkeypatch.setattr(fusion_module, "get_attention_context", fail_context)
+
+    with pytest.raises(RuntimeError, match="missing attention context"):
+        _call_fused_mrope_impl(torch.empty(2, 8, device="cpu"))

vllm/_aiter_ops.py

@@ -3362,6 +3362,115 @@ class rocm_aiter_ops:
             rotary_dim=kernel_rotary_dim,
         )
 
+    @staticmethod
+    def fused_qk_norm_mrope_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,
+        head_dim: int,
+        is_neox: bool,
+        mrope_section: list[int],
+        is_interleaved: 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,
+        block_size: int,
+        x: int,
+        rotary_dim: int = 0,
+    ) -> None:
+        from aiter.ops.fused_qk_norm_mrope_cache_quant import (
+            fused_qk_norm_mrope_3d_cache_pts_quant_shuffle,
+        )
+
+        fused_qk_norm_mrope_3d_cache_pts_quant_shuffle(
+            qkv,
+            q_weight,
+            k_weight,
+            cos_sin_cache,
+            positions,
+            qkv.size(0),
+            num_heads_q,
+            num_heads_k,
+            num_heads_k,
+            head_dim,
+            is_neox,
+            mrope_section,
+            is_interleaved,
+            rms_norm_eps,
+            q_out,
+            k_cache,
+            v_cache,
+            slot_mapping,
+            k_scale,
+            v_scale,
+            None,
+            None,
+            False,
+            False,
+            block_size,
+            x,
+            rotary_dim,
+        )
+
+    @staticmethod
+    def do_qk_norm_mrope_kvcache_update(
+        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,
+        head_dim: int,
+        is_neox: bool,
+        mrope_section: list[int],
+        is_interleaved: bool,
+        rms_norm_eps: float,
+        q_out: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        k_scale: torch.Tensor,
+        v_scale: torch.Tensor,
+        kv_cache_dtype: str,
+        rotary_dim: int = 0,
+    ) -> None:
+        if kv_cache_dtype.startswith("fp8"):
+            key_cache = key_cache.view(current_platform.fp8_dtype())
+            value_cache = value_cache.view(current_platform.fp8_dtype())
+
+        rocm_aiter_ops.fused_qk_norm_mrope_and_cache(
+            qkv=qkv,
+            q_weight=q_weight,
+            k_weight=k_weight,
+            cos_sin_cache=cos_sin_cache,
+            positions=positions,
+            num_heads_q=num_heads_q,
+            num_heads_k=num_heads_k,
+            head_dim=head_dim,
+            is_neox=is_neox,
+            mrope_section=mrope_section,
+            is_interleaved=is_interleaved,
+            rms_norm_eps=rms_norm_eps,
+            q_out=q_out,
+            k_cache=key_cache,
+            v_cache=value_cache,
+            slot_mapping=slot_mapping,
+            k_scale=k_scale,
+            v_scale=v_scale,
+            block_size=key_cache.shape[1],
+            x=16 // key_cache.element_size(),
+            rotary_dim=rotary_dim,
+        )
+
     @staticmethod
     def triton_rope_and_cache(
         query: torch.Tensor,

vllm/compilation/passes/fusion/matcher_utils.py

@@ -31,6 +31,7 @@ from vllm.platforms import current_platform
 
 ROTARY_OP = torch.ops._C.rotary_embedding.default
 FLASHINFER_ROTARY_OP = torch.ops.vllm.flashinfer_rotary_embedding.default
+MROPE_OP = torch.ops.vllm.mrope.default
 
 QUANT_OPS: dict[QuantKey, OpOverload] = {
     kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default,  # noqa: E501
@@ -162,6 +163,74 @@ class MatcherRotaryEmbedding(MatcherCustomOp):
         return result
 
 
+class MatcherMRotaryEmbedding(MatcherCustomOp):
+    def __init__(
+        self,
+        is_neox: bool,
+        head_size: int,
+        rotary_dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+        mrope_section: tuple[int, int, int],
+        mrope_interleaved: bool,
+        enabled: bool | None = None,
+    ) -> None:
+        if enabled is None:
+            enabled = RotaryEmbedding.enabled()
+
+        super().__init__(enabled)
+        self.is_neox = is_neox
+        self.head_size = head_size
+        self.rotary_dim = rotary_dim
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.q_size = self.num_heads * self.head_size
+        self.kv_size = self.num_kv_heads * self.head_size
+        self.mrope_section = mrope_section
+        self.mrope_interleaved = mrope_interleaved
+
+    def inputs(self) -> list[torch.Tensor]:
+        positions = self.empty_int64(3, 5)
+        query = self.empty(5, self.q_size)
+        key = self.empty(5, self.kv_size)
+        cos_sin_cache = self.empty(4096, self.rotary_dim)
+        return [positions, query, key, cos_sin_cache]
+
+    def forward_custom(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        result = auto_functionalized(
+            MROPE_OP,
+            positions=positions,
+            query=query,
+            key=key,
+            cos_sin_cache=cos_sin_cache,
+            head_size=self.head_size,
+            rotary_dim=self.rotary_dim,
+            mrope_section_t=self.mrope_section[0],
+            mrope_section_h=self.mrope_section[1],
+            mrope_section_w=self.mrope_section[2],
+            mrope_interleaved=self.mrope_interleaved,
+            is_neox_style=self.is_neox,
+        )
+        return result[1], result[2]
+
+    def forward_native(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError(
+            "MRoPE fusion requires the rotary_embedding custom op"
+        )
+
+
 class MatcherRMSNormGated(MatcherCustomOp):
     """Matches RMSNormGated with norm_before_gate=True and group_size=None."""
 

vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import inspect
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import ParamSpec
 
 import torch
@@ -19,12 +19,19 @@ from vllm.model_executor.layers.attention.attention import (
     Attention,
     get_attention_context,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.platforms import current_platform
-from vllm.utils.torch_utils import direct_register_custom_op
+from vllm.utils.torch_utils import (
+    _USE_LAYERNAME,
+    LayerNameType,
+    _encode_layer_name,
+    _resolve_layer_name,
+    direct_register_custom_op,
+)
 
 from ..inductor_pass import enable_fake_mode
 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
-from .matcher_utils import MatcherRotaryEmbedding
+from .matcher_utils import MROPE_OP, MatcherMRotaryEmbedding, MatcherRotaryEmbedding
 from .rms_quant_fusion import empty_bf16, empty_fp32, empty_i64
 
 logger = init_logger(__name__)
@@ -51,8 +58,9 @@ def fused_qk_norm_rope_and_unified_kv_cache_update_impl(
     rms_norm_eps: float,
     cos_sin_cache: torch.Tensor,
     is_neox: bool,
-    layer_name: str = "",
+    layer_name: LayerNameType,
 ) -> torch.Tensor:
+    layer_name = _resolve_layer_name(layer_name)
     _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
     if layer_slot_mapping is not None:
         attn_layer.impl.do_qk_norm_rope_kvcache_update(
@@ -87,7 +95,7 @@ def fused_qk_norm_rope_and_unified_kv_cache_update_fake(
     rms_norm_eps: float,
     cos_sin_cache: torch.Tensor,
     is_neox: bool,
-    layer_name: str = "",
+    layer_name: LayerNameType,
 ) -> torch.Tensor:
     return torch.empty(0, device=qkv.device, dtype=qkv.dtype)
 
@@ -100,6 +108,76 @@ direct_register_custom_op(
 )
 
 
+def fused_qk_norm_mrope_and_unified_kv_cache_update_impl(
+    q_out: torch.Tensor,
+    qkv: torch.Tensor,
+    positions: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rms_norm_eps: float,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    rotary_dim: int,
+    layer_name: LayerNameType,
+) -> torch.Tensor:
+    layer_name = _resolve_layer_name(layer_name)
+    _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
+    if layer_slot_mapping is not None:
+        if positions.ndim == 1:
+            positions = positions.unsqueeze(0).expand(3, -1).contiguous()
+        attn_layer.impl.do_qk_norm_mrope_kvcache_update(
+            attn_layer,
+            qkv,
+            q_out,
+            positions,
+            q_weight,
+            k_weight,
+            rms_norm_eps,
+            cos_sin_cache,
+            is_neox,
+            (mrope_section_t, mrope_section_h, mrope_section_w),
+            is_interleaved,
+            rotary_dim,
+            kv_cache,
+            layer_slot_mapping,
+        )
+    else:
+        q_out.zero_()
+
+    return torch.empty(0, device=qkv.device, dtype=qkv.dtype)
+
+
+def fused_qk_norm_mrope_and_unified_kv_cache_update_fake(
+    q_out: torch.Tensor,
+    qkv: torch.Tensor,
+    positions: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rms_norm_eps: float,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    rotary_dim: int,
+    layer_name: LayerNameType,
+) -> torch.Tensor:
+    return torch.empty(0, device=qkv.device, dtype=qkv.dtype)
+
+
+direct_register_custom_op(
+    op_name="fused_qk_norm_mrope_and_unified_kv_cache_update",
+    op_func=fused_qk_norm_mrope_and_unified_kv_cache_update_impl,
+    mutates_args=["q_out"],
+    fake_impl=fused_qk_norm_mrope_and_unified_kv_cache_update_fake,
+)
+
+
 # ---------------------------------------------------------------------------
 # Pattern: QK-norm + RoPE + unified_kv_cache_update
 # ---------------------------------------------------------------------------
@@ -142,19 +220,27 @@ class QkNormRopeKvCachePattern:
         self.eps = eps
         self.is_neox = is_neox
         self.quant_query = quant_query
+        self.encoded_layer_name = _encode_layer_name(self.layer_name)
+        self.query_quant_group_shape = (
+            layer.query_quant.group_shape
+            if layer.query_quant is not None
+            else GroupShape.PER_TENSOR
+        )
 
         self.q_size = self.num_heads * self.head_size
         self.k_size = self.num_kv_heads * self.head_size
         self.v_size = self.num_kv_heads * self.head_size_v
 
-        self.rope_matcher = MatcherRotaryEmbedding(
-            is_neox=is_neox,
-            head_size=self.head_size,
-            num_heads=self.num_heads,
-            num_kv_heads=self.num_kv_heads,
+        self.rope_matcher: MatcherRotaryEmbedding | MatcherMRotaryEmbedding = (
+            MatcherRotaryEmbedding(
+                is_neox=is_neox,
+                head_size=self.head_size,
+                num_heads=self.num_heads,
+                num_kv_heads=self.num_kv_heads,
+            )
         )
 
-    def get_inputs(self) -> list[torch.Tensor]:
+    def get_inputs(self) -> list:
         T = 5
         L = 4096
         qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
@@ -166,8 +252,13 @@ class QkNormRopeKvCachePattern:
         if self.quant_query:
             q_scale = empty_fp32(1)
             inputs += [q_scale]
+        if _USE_LAYERNAME:
+            inputs.append(self.encoded_layer_name)
         return inputs
 
+    def _get_layer_name(self, layer_name: LayerNameType | None) -> LayerNameType:
+        return self.encoded_layer_name if layer_name is None else layer_name
+
     def pattern_non_fp8_quant_query(
         self,
         qkv: torch.Tensor,
@@ -175,6 +266,7 @@ class QkNormRopeKvCachePattern:
         q_weight: torch.Tensor,
         k_weight: torch.Tensor,
         cos_sin_cache: torch.Tensor,
+        layer_name: LayerNameType | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
         q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size)
@@ -190,7 +282,9 @@ class QkNormRopeKvCachePattern:
         q_rope = q_rope.view(-1, self.num_heads, self.head_size)
         k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size)
         v = v.view(-1, self.num_kv_heads, self.head_size_v)
-        dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name)
+        dummy = torch.ops.vllm.unified_kv_cache_update(
+            k_rope, v, self._get_layer_name(layer_name)
+        )
         return dummy, q_rope, k_rope, v
 
     def replacement_non_fp8_quant_query(
@@ -200,6 +294,7 @@ class QkNormRopeKvCachePattern:
         q_weight: torch.Tensor,
         k_weight: torch.Tensor,
         cos_sin_cache: torch.Tensor,
+        layer_name: LayerNameType | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         q_out = torch.empty(
             qkv.shape[0],
@@ -228,10 +323,49 @@ class QkNormRopeKvCachePattern:
             rms_norm_eps=self.eps,
             cos_sin_cache=cos_sin_cache,
             is_neox=self.is_neox,
-            layer_name=self.layer_name,
+            layer_name=self._get_layer_name(layer_name),
         )
         return results[0], results[1], results[2], v
 
+    def _quantize_query(
+        self,
+        query: torch.Tensor,
+        q_scale: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if self.query_quant_group_shape.is_per_tensor():
+            q_out = torch.empty_like(query, dtype=current_platform.fp8_dtype())
+            q_quant = auto_functionalized(
+                torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
+                out=q_out,
+                x=query,
+                scale=q_scale,
+                is_dynamic=False,
+            )
+            # The AITER op's schema marks scale mutable, so preserve its
+            # functionalized write-back as an output of the matched region.
+            return q_quant[1], q_quant[2]
+
+        # QuantFP8.forward_hip falls back to the generic static op for the
+        # historical per-head shape (-1, block_size). This also covers the
+        # degenerate one-local-KV-head case where a scalar scale is represented
+        # as (-1, q_size), rather than changing Attention's classification.
+        q_out = torch.empty(
+            query.shape,
+            device=query.device,
+            dtype=current_platform.fp8_dtype(),
+        )
+        q_quant = auto_functionalized(
+            torch.ops._C.static_scaled_fp8_quant.default,
+            result=q_out,
+            input=query,
+            scale=q_scale,
+            group_shape=[
+                self.query_quant_group_shape.row,
+                self.query_quant_group_shape.col,
+            ],
+        )
+        return q_quant[1], None
+
     def pattern_fp8_quant_query(
         self,
         qkv: torch.Tensor,
@@ -240,7 +374,8 @@ class QkNormRopeKvCachePattern:
         k_weight: torch.Tensor,
         cos_sin_cache: torch.Tensor,
         q_scale: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        layer_name: LayerNameType | None = None,
+    ) -> tuple[torch.Tensor, ...]:
         q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
         q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size)
         q_normed = vllm.ir.ops.rms_norm(q_by_head, q_weight, self.eps)
@@ -251,24 +386,15 @@ class QkNormRopeKvCachePattern:
         k_flat = k_normed.view(-1, self.k_size)
 
         q_rope, k_rope = self.rope_matcher(positions, q_flat, k_flat, cos_sin_cache)
-        # Match the quant-query op Attention.forward inserts (fp8 KV + UNIFIED).
-        # Explicit auto_functionalized (out=[1]) keeps the quant node in the pattern.
-        q_out = torch.empty_like(q_rope, dtype=current_platform.fp8_dtype())
-        q_quant = auto_functionalized(
-            torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
-            out=q_out,
-            x=q_rope,
-            scale=q_scale,
-            is_dynamic=False,
-        )
-        # `scale` is mutable: its copy_ write-back to _q_scale bumps the mutation
-        # region, so keep q flat (a reshape lands past the barrier and won't match).
-        q_rope_fp8 = q_quant[1]
-        q_scale_out = q_quant[2]
+        q_rope_fp8, q_scale_out = self._quantize_query(q_rope, q_scale)
 
         k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size)
         v = v.view(-1, self.num_kv_heads, self.head_size_v)
-        dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name)
+        dummy = torch.ops.vllm.unified_kv_cache_update(
+            k_rope, v, self._get_layer_name(layer_name)
+        )
+        if q_scale_out is None:
+            return dummy, q_rope_fp8, k_rope, v
         return dummy, q_rope_fp8, k_rope, v, q_scale_out
 
     def replacement_fp8_quant_query(
@@ -279,7 +405,8 @@ class QkNormRopeKvCachePattern:
         k_weight: torch.Tensor,
         cos_sin_cache: torch.Tensor,
         q_scale: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        layer_name: LayerNameType | None = None,
+    ) -> tuple[torch.Tensor, ...]:
         q_out = torch.empty(
             qkv.shape[0],
             self.num_heads,
@@ -307,22 +434,14 @@ class QkNormRopeKvCachePattern:
             rms_norm_eps=self.eps,
             cos_sin_cache=cos_sin_cache,
             is_neox=self.is_neox,
-            layer_name=self.layer_name,
+            layer_name=self._get_layer_name(layer_name),
         )
-        # Re-apply the quant on the kernel's bf16 q_out; fused op does not quant q.
-        # Same explicit auto_functionalized form as the pattern: [1] = quantized
-        # q, [2] = scale (returned so the buffer-writeback use is preserved).
+        # Re-apply the same quant form on the kernel's bf16 q_out; the fused
+        # operation does not quantize Q.
         q_fp8_flat = results[1].view(-1, self.q_size)
-        q_fp8_out = torch.empty_like(q_fp8_flat, dtype=current_platform.fp8_dtype())
-        q_requant = auto_functionalized(
-            torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
-            out=q_fp8_out,
-            x=q_fp8_flat,
-            scale=q_scale,
-            is_dynamic=False,
-        )
-        q_fp8 = q_requant[1]  # flat to mirror the pattern (see note above)
-        q_scale_out = q_requant[2]
+        q_fp8, q_scale_out = self._quantize_query(q_fp8_flat, q_scale)
+        if q_scale_out is None:
+            return results[0], q_fp8, results[2], v
         return results[0], q_fp8, results[2], v, q_scale_out
 
     @staticmethod
@@ -345,6 +464,9 @@ class QkNormRopeKvCachePattern:
 
         view_to_reshape(gm)
 
+    def _extra_check(self, match: pm.Match) -> bool:
+        return True
+
     def _register(self, pattern, replacement, pm_pass) -> None:
         trace_fn = QkNormRopeKvCachePattern.wrap_trace_fn(
             pm.fwd_only,
@@ -370,40 +492,338 @@ class QkNormRopeKvCachePattern:
             inputs,
             trace_fn,
             pm_pass,
+            extra_check=self._extra_check,
             search_fn_pattern=search_fn_pattern,
         )
 
     def register(self, pm_pass: PatternMatcherPass) -> None:
         # make_fx counts `self` in bound-method code params; wrap as plain fns.
-        # Distinct names per branch so mypy doesn't see one name, two signatures.
         if self.quant_query:
+            if _USE_LAYERNAME:
 
-            def pattern_q(qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale):
-                return self.pattern_fp8_quant_query(
+                def pattern_q_with_layer(
+                    qkv,
+                    positions,
+                    q_weight,
+                    k_weight,
+                    cos_sin_cache,
+                    q_scale,
+                    layer_name,
+                ):
+                    return self.pattern_fp8_quant_query(
+                        qkv,
+                        positions,
+                        q_weight,
+                        k_weight,
+                        cos_sin_cache,
+                        q_scale,
+                        layer_name,
+                    )
+
+                def replacement_q_with_layer(
+                    qkv,
+                    positions,
+                    q_weight,
+                    k_weight,
+                    cos_sin_cache,
+                    q_scale,
+                    layer_name,
+                ):
+                    return self.replacement_fp8_quant_query(
+                        qkv,
+                        positions,
+                        q_weight,
+                        k_weight,
+                        cos_sin_cache,
+                        q_scale,
+                        layer_name,
+                    )
+
+                self._register(pattern_q_with_layer, replacement_q_with_layer, pm_pass)
+            else:
+
+                def pattern_q_without_layer(
                     qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
-                )
+                ):
+                    return self.pattern_fp8_quant_query(
+                        qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
+                    )
 
-            def replacement_q(
-                qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
-            ):
-                return self.replacement_fp8_quant_query(
+                def replacement_q_without_layer(
                     qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
-                )
+                ):
+                    return self.replacement_fp8_quant_query(
+                        qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
+                    )
 
-            self._register(pattern_q, replacement_q, pm_pass)
+                self._register(
+                    pattern_q_without_layer, replacement_q_without_layer, pm_pass
+                )
         else:
+            if _USE_LAYERNAME:
 
-            def pattern_noq(qkv, positions, q_weight, k_weight, cos_sin_cache):
-                return self.pattern_non_fp8_quant_query(
+                def pattern_noq_with_layer(
+                    qkv,
+                    positions,
+                    q_weight,
+                    k_weight,
+                    cos_sin_cache,
+                    layer_name,
+                ):
+                    return self.pattern_non_fp8_quant_query(
+                        qkv,
+                        positions,
+                        q_weight,
+                        k_weight,
+                        cos_sin_cache,
+                        layer_name,
+                    )
+
+                def replacement_noq_with_layer(
+                    qkv,
+                    positions,
+                    q_weight,
+                    k_weight,
+                    cos_sin_cache,
+                    layer_name,
+                ):
+                    return self.replacement_non_fp8_quant_query(
+                        qkv,
+                        positions,
+                        q_weight,
+                        k_weight,
+                        cos_sin_cache,
+                        layer_name,
+                    )
+
+                self._register(
+                    pattern_noq_with_layer, replacement_noq_with_layer, pm_pass
+                )
+            else:
+
+                def pattern_noq_without_layer(
                     qkv, positions, q_weight, k_weight, cos_sin_cache
+                ):
+                    return self.pattern_non_fp8_quant_query(
+                        qkv, positions, q_weight, k_weight, cos_sin_cache
+                    )
+
+                def replacement_noq_without_layer(
+                    qkv, positions, q_weight, k_weight, cos_sin_cache
+                ):
+                    return self.replacement_non_fp8_quant_query(
+                        qkv, positions, q_weight, k_weight, cos_sin_cache
+                    )
+
+                self._register(
+                    pattern_noq_without_layer,
+                    replacement_noq_without_layer,
+                    pm_pass,
                 )
 
-            def replacement_noq(qkv, positions, q_weight, k_weight, cos_sin_cache):
-                return self.replacement_non_fp8_quant_query(
-                    qkv, positions, q_weight, k_weight, cos_sin_cache
-                )
 
-            self._register(pattern_noq, replacement_noq, pm_pass)
+class QkNormMRopeKvCachePattern(QkNormRopeKvCachePattern):
+    FUSED_OP = torch.ops.vllm.fused_qk_norm_mrope_and_unified_kv_cache_update.default
+
+    def __init__(
+        self,
+        layer: Attention,
+        eps: float,
+        is_neox: bool,
+        quant_query: bool,
+        rotary_dim: int,
+        mrope_section: tuple[int, int, int],
+        is_interleaved: bool,
+    ) -> None:
+        super().__init__(
+            layer=layer,
+            eps=eps,
+            is_neox=is_neox,
+            quant_query=quant_query,
+        )
+        self.rotary_dim = rotary_dim
+        self.mrope_section = mrope_section
+        self.is_interleaved = is_interleaved
+        self.rope_matcher = MatcherMRotaryEmbedding(
+            is_neox=is_neox,
+            head_size=self.head_size,
+            rotary_dim=rotary_dim,
+            num_heads=self.num_heads,
+            num_kv_heads=self.num_kv_heads,
+            mrope_section=mrope_section,
+            mrope_interleaved=is_interleaved,
+        )
+
+    def get_inputs(self) -> list[torch.Tensor]:
+        T = 5
+        L = 4096
+        qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
+        positions = empty_i64(3, T)
+        q_weight = empty_bf16(1, self.head_size)
+        k_weight = empty_bf16(1, self.head_size)
+        cos_sin_cache = empty_bf16(L, self.rotary_dim)
+        inputs = [qkv, positions, q_weight, k_weight, cos_sin_cache]
+        if self.quant_query:
+            inputs.append(empty_fp32(1))
+        if _USE_LAYERNAME:
+            inputs.append(self.encoded_layer_name)
+        return inputs
+
+    def _unfused_kv_views(self, qkv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # The supported backend reads K/V from the cache after the fused write.
+        # Keep view-compatible placeholders for the attention custom-op ABI.
+        _, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
+        k = k.view(qkv.shape[0], self.num_kv_heads, self.head_size)
+        v = v.view(qkv.shape[0], self.num_kv_heads, self.head_size_v)
+        return k, v
+
+    def _fused_results(
+        self,
+        qkv: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        layer_name: LayerNameType | None,
+    ) -> tuple[torch.Tensor, ...]:
+        q_out = torch.empty(
+            qkv.shape[0],
+            self.num_heads,
+            self.head_size,
+            device=qkv.device,
+            dtype=qkv.dtype,
+        )
+        return auto_functionalized(
+            self.FUSED_OP,
+            q_out=q_out,
+            qkv=qkv,
+            positions=positions,
+            q_weight=q_weight,
+            k_weight=k_weight,
+            rms_norm_eps=self.eps,
+            cos_sin_cache=cos_sin_cache,
+            is_neox=self.is_neox,
+            mrope_section_t=self.mrope_section[0],
+            mrope_section_h=self.mrope_section[1],
+            mrope_section_w=self.mrope_section[2],
+            is_interleaved=self.is_interleaved,
+            rotary_dim=self.rotary_dim,
+            layer_name=self._get_layer_name(layer_name),
+        )
+
+    def replacement_non_fp8_quant_query(
+        self,
+        qkv: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        layer_name: LayerNameType | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        k, v = self._unfused_kv_views(qkv)
+        results = self._fused_results(
+            qkv, positions, q_weight, k_weight, cos_sin_cache, layer_name
+        )
+        return results[0], results[1], k, v
+
+    def replacement_fp8_quant_query(
+        self,
+        qkv: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        q_scale: torch.Tensor,
+        layer_name: LayerNameType | None = None,
+    ) -> tuple[torch.Tensor, ...]:
+        k, v = self._unfused_kv_views(qkv)
+        results = self._fused_results(
+            qkv, positions, q_weight, k_weight, cos_sin_cache, layer_name
+        )
+        q_flat = results[1].view(-1, self.q_size)
+        q_fp8, q_scale_out = self._quantize_query(q_flat, q_scale)
+        if q_scale_out is None:
+            return results[0], q_fp8, k, v
+        return results[0], q_fp8, k, v, q_scale_out
+
+    def _extra_check(self, match: pm.Match) -> bool:
+        expected = {
+            "head_size": self.head_size,
+            "rotary_dim": self.rotary_dim,
+            "mrope_section_t": self.mrope_section[0],
+            "mrope_section_h": self.mrope_section[1],
+            "mrope_section_w": self.mrope_section[2],
+            "mrope_interleaved": self.is_interleaved,
+            "is_neox_style": self.is_neox,
+        }
+        mrope_node = next(
+            (node for node in match.nodes if node.args and node.args[0] == MROPE_OP),
+            None,
+        )
+        if mrope_node is None or any(
+            mrope_node.kwargs.get(name) != value for name, value in expected.items()
+        ):
+            return False
+
+        # The fused AITER MRoPE entrypoint writes rotated K directly to the
+        # cache and does not expose a BF16 K result. Returning the raw K view is
+        # valid only for decoder unified attention, whose backend consumes K
+        # from that cache. Reject graphs that observe K in any other way.
+        output_nodes = match.output_nodes()
+        if len(output_nodes) < 3 or output_nodes[2] is None:
+            return False
+        matched_nodes = set(match.nodes)
+        attention_op = torch.ops.vllm.unified_attention_with_output.default
+
+        def is_unified_attention_user(user: fx.Node) -> bool:
+            return user.op == "call_function" and (
+                user.target == attention_op
+                or (user.args and user.args[0] == attention_op)
+            )
+
+        return all(
+            user in matched_nodes or is_unified_attention_user(user)
+            for user in output_nodes[2].users
+        )
+
+
+def _discover_mrope_configs(
+    config: VllmConfig,
+) -> tuple[tuple[tuple[int, int, int], bool], ...]:
+    model_config = config.model_config
+    hf_configs = (
+        getattr(model_config, "hf_text_config", None),
+        getattr(model_config, "hf_config", None),
+    )
+    configs: set[tuple[tuple[int, int, int], bool]] = set()
+    for hf_config in hf_configs:
+        if hf_config is None:
+            continue
+        rope_parameters = getattr(hf_config, "rope_parameters", None)
+        if rope_parameters is None:
+            rope_parameters = getattr(hf_config, "rope_scaling", None)
+        if not isinstance(rope_parameters, Mapping):
+            continue
+        section = rope_parameters.get("mrope_section")
+        if not isinstance(section, (list, tuple)) or len(section) != 3:
+            continue
+        try:
+            parsed_section = (
+                int(section[0]),
+                int(section[1]),
+                int(section[2]),
+            )
+        except (TypeError, ValueError):
+            continue
+        if any(value <= 0 for value in parsed_section):
+            continue
+        configs.add(
+            (
+                parsed_section,
+                bool(rope_parameters.get("mrope_interleaved", False)),
+            )
+        )
+    return tuple(sorted(configs))
 
 
 # ---------------------------------------------------------------------------
@@ -412,7 +832,7 @@ class QkNormRopeKvCachePattern:
 
 
 class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):
-    """Fuse QK-norm + RoPE + KV cache update into a single AITER HIP kernel.
+    """Fuse QK-norm + RoPE/MRoPE + KV cache update into an AITER HIP kernel.
 
     Supersedes both QKNormRoPEFusionPass and RopeKVCacheFusionPass for
     attention layers that support the combined operation, eliminating two
@@ -429,22 +849,33 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):
 
         cc = config.compilation_config
         self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num
+        self.mrope_configs = (
+            _discover_mrope_configs(config)
+            if cc.is_custom_op_enabled("rotary_embedding")
+            else ()
+        )
 
         dtype = config.model_config.dtype
         if dtype not in (torch.bfloat16, torch.float16):
             logger.warning_once(
-                "QK Norm+RoPE+KVCache fusion not enabled: unsupported dtype %s", dtype
+                "QK Norm+RoPE/MRoPE+KVCache fusion not enabled: unsupported dtype %s",
+                dtype,
             )
             return
 
         attn_layers = get_layers_from_vllm_config(config, Attention)
 
         for _, layer in attn_layers.items():
-            if not layer.impl.fused_qk_norm_rope_kvcache_supported():
+            supports_rope = layer.impl.fused_qk_norm_rope_kvcache_supported()
+            supports_mrope = (
+                bool(self.mrope_configs)
+                and layer.impl.fused_qk_norm_mrope_kvcache_supported()
+            )
+            if not supports_rope and not supports_mrope:
                 continue
             if layer.head_size not in SUPPORTED_FUSED_QK_NORM_ROPE_KVCACHE_HEAD_DIMS:
                 logger.warning_once(
-                    "QK Norm+RoPE+KVCache fusion not enabled for a layer: "
+                    "QK Norm+RoPE/MRoPE+KVCache fusion not enabled for a layer: "
                     "head_size=%d is not supported by the "
                     "fused_qk_norm_rope_cache_pts_quant_shuffle kernel "
                     "(supported: %s). Falling back to the unfused path.",
@@ -455,22 +886,71 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):
             if layer.head_size_v != layer.head_size:
                 # The fused kernel uses a single head_dim for q/k/v.
                 logger.warning_once(
-                    "QK Norm+RoPE+KVCache fusion not enabled for a layer: "
+                    "QK Norm+RoPE/MRoPE+KVCache fusion not enabled for a layer: "
                     "head_size_v=%d differs from head_size=%d, which the fused "
                     "kernel does not support. Falling back to the unfused path.",
                     layer.head_size_v,
                     layer.head_size,
                 )
                 continue
-            for epsilon in [1e-5, 1e-6]:
-                for neox in [True, False]:
-                    for quant_q in [False, True]:
-                        QkNormRopeKvCachePattern(
-                            layer=layer,
-                            eps=epsilon,
-                            is_neox=neox,
-                            quant_query=quant_q,
-                        ).register(self.patterns)
+            if supports_rope:
+                for epsilon in [1e-5, 1e-6]:
+                    for neox in [True, False]:
+                        for quant_q in [False, True]:
+                            QkNormRopeKvCachePattern(
+                                layer=layer,
+                                eps=epsilon,
+                                is_neox=neox,
+                                quant_query=quant_q,
+                            ).register(self.patterns)
+
+            if supports_mrope:
+                for section, is_interleaved in self.mrope_configs:
+                    rotary_dim = 2 * sum(section)
+                    if rotary_dim > layer.head_size:
+                        logger.warning_once(
+                            "QK Norm+MRoPE+KVCache fusion not enabled for a "
+                            "layer: rotary_dim=%d exceeds head_size=%d.",
+                            rotary_dim,
+                            layer.head_size,
+                        )
+                        continue
+                    vector_width = layer.head_size // 32
+                    for epsilon in [1e-5, 1e-6]:
+                        for neox in [True, False]:
+                            # Each warp lane processes `vector_width` adjacent
+                            # values. NeoX additionally exchanges vectors
+                            # across the half-rotary boundary.
+                            rotary_alignment = vector_width * (2 if neox else 1)
+                            if rotary_dim % rotary_alignment != 0:
+                                logger.warning_once(
+                                    "QK Norm+MRoPE+KVCache fusion not enabled "
+                                    "for a layer/style: rotary_dim=%d must be "
+                                    "divisible by %d for head_size=%d and "
+                                    "is_neox=%s.",
+                                    rotary_dim,
+                                    rotary_alignment,
+                                    layer.head_size,
+                                    neox,
+                                )
+                                continue
+                            for quant_q in [False, True]:
+                                QkNormMRopeKvCachePattern(
+                                    layer=layer,
+                                    eps=epsilon,
+                                    is_neox=neox,
+                                    quant_query=quant_q,
+                                    rotary_dim=rotary_dim,
+                                    mrope_section=section,
+                                    is_interleaved=is_interleaved,
+                                ).register(self.patterns)
+
+            # Opaque LayerName is a pattern wildcard on torch >= 2.11, so
+            # homogeneous attention layers produce duplicate search patterns.
+            # As in the other LayerName-aware fusion passes, one supported
+            # layer registers all shape/style variants for the model.
+            if _USE_LAYERNAME:
+                break
 
         self.dump_patterns(config, self.patterns)
 
@@ -478,8 +958,8 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):
     def __call__(self, graph: fx.Graph) -> None:
         self.matched_count = self.patterns.apply(graph)
         logger.info(
-            "QK-Norm+RoPE+KVCache fusion: replaced %s pattern(s) "
-            "with AITER fused_qk_norm_rope_cache_pts_quant_shuffle",
+            "QK-Norm+RoPE/MRoPE+KVCache fusion: replaced %s pattern(s) "
+            "with an AITER fused attention-prologue kernel",
             self.matched_count,
         )
 
@@ -487,4 +967,10 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):
         return compile_range.end <= self.max_token_num
 
     def uuid(self) -> str:
-        return VllmInductorPass.hash_source(self, QkNormRopeKvCachePattern)
+        return VllmInductorPass.hash_source(
+            self,
+            QkNormRopeKvCachePattern,
+            QkNormMRopeKvCachePattern,
+            _discover_mrope_configs,
+            repr(self.mrope_configs),
+        )

vllm/config/compilation.py

@@ -154,15 +154,15 @@ class PassConfig:
     fuse_rope_kvcache: bool = None  # type: ignore[assignment]
     """Fuse the QK rope + KV cache ops."""
     fuse_qk_norm_rope_kvcache: bool = Field(default=None)  # type: ignore[assignment]
-    """Fuse QK RMSNorm + RoPE + KV cache update into a single AITER HIP
+    """Fuse QK RMSNorm + RoPE/MRoPE + KV cache update into an AITER HIP
     kernel. Supersedes both enable_qk_norm_rope_fusion and fuse_rope_kvcache
-    for layers that support it. Auto-enabled at O1+ on ROCm for models
-    with QK-norm (e.g. Qwen3-MoE)."""
+    for layers that support it. Auto-enabled at O2+ on ROCm for models
+    with QK-norm (e.g. Qwen3-MoE and Qwen3-VL-class architectures)."""
 
     rope_kvcache_fusion_max_token_num: int = 256
     """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode.
     Larger batch sizes e.g. during prefill will use the unfused kernels.
-    Also applies to the fused QK-Norm+RoPE+KVCache pass.
+    Also applies to the fused QK-Norm+RoPE/MRoPE+KVCache pass.
     """
 
     fi_allreduce_fusion_max_size_mb: float | None = None

vllm/config/vllm.py

@@ -242,7 +242,7 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool:
 
 
 def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool:
-    """Enable fused QK-norm + RoPE + KV cache update on ROCm with AITER."""
+    """Enable fused QK-norm + RoPE/MRoPE + KV cache update with AITER."""
     from vllm._aiter_ops import rocm_aiter_ops
 
     if not rocm_aiter_ops.is_enabled():

vllm/model_executor/layers/rotary_embedding/mrope.py

@@ -7,6 +7,7 @@ import torch
 
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import direct_register_custom_op
 
 from .base import RotaryEmbeddingBase
 from .yarn_scaling_rope import YaRNScalingRotaryEmbedding, yarn_get_mscale
@@ -240,6 +241,62 @@ def apply_interleaved_rope(x: torch.Tensor, mrope_section: list[int]) -> torch.T
     return torch.where(is_width, x[2], result)
 
 
+def _mrope_apply(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    head_size: int,
+    rotary_dim: int,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    mrope_interleaved: bool,
+    is_neox_style: bool,
+) -> None:
+    cos_sin = cos_sin_cache[positions]
+    cos, sin = cos_sin.chunk(2, dim=-1)
+    query_out, key_out = triton_mrope(
+        query,
+        key,
+        cos,
+        sin,
+        [mrope_section_t, mrope_section_h, mrope_section_w],
+        head_size,
+        rotary_dim,
+        mrope_interleaved,
+        is_neox_style,
+    )
+    if query_out.data_ptr() != query.data_ptr():
+        query.copy_(query_out.view_as(query))
+    if key_out.data_ptr() != key.data_ptr():
+        key.copy_(key_out.view_as(key))
+
+
+def _mrope_apply_fake(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    head_size: int,
+    rotary_dim: int,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    mrope_interleaved: bool,
+    is_neox_style: bool,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="mrope",
+    op_func=_mrope_apply,
+    mutates_args=["query", "key"],
+    fake_impl=_mrope_apply_fake,
+)
+
+
 class MRotaryEmbedding(RotaryEmbeddingBase):
     """Rotary Embedding with Multimodal Sections."""
 
@@ -426,6 +483,35 @@ class MRotaryEmbedding(RotaryEmbeddingBase):
         key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
         return query, key
 
+    def forward_hip(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor | None = None,
+        offsets: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if positions.ndim == 2:
+            assert key is not None
+            assert self.mrope_section
+            query_shape = query.shape
+            key_shape = key.shape
+            cos_sin_cache = self._match_cos_sin_cache_dtype(query)
+            torch.ops.vllm.mrope(
+                positions,
+                query,
+                key,
+                cos_sin_cache,
+                self.head_size,
+                self.rotary_dim,
+                self.mrope_section[0],
+                self.mrope_section[1],
+                self.mrope_section[2],
+                self.mrope_interleaved,
+                self.is_neox_style,
+            )
+            return query.reshape(query_shape), key.reshape(key_shape)
+        return super().forward_hip(positions, query, key)
+
     def forward_cpu(
         self,
         positions: torch.Tensor,

vllm/v1/attention/backend.py

@@ -926,6 +926,10 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
         """
         return False
 
+    def fused_qk_norm_mrope_kvcache_supported(self):
+        """Whether this implementation supports fused QKNorm+MRoPE+KVCache."""
+        return False
+
     def fused_rope_kvcache_supported(self):
         """Does this attention implementation support RoPE+KVCache fusion.
         This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops
@@ -955,6 +959,26 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
         """
         raise NotImplementedError
 
+    def do_qk_norm_mrope_kvcache_update(
+        self,
+        layer: AttentionLayer,
+        qkv: torch.Tensor,
+        q_out: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        rms_norm_eps: float,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        mrope_section: tuple[int, int, int],
+        is_interleaved: bool,
+        rotary_dim: int,
+        kv_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+    ):
+        """Apply QK-norm and MRoPE, then write K/V to the cache."""
+        raise NotImplementedError
+
     def do_rope_and_kv_cache_update(
         self,
         layer: AttentionLayer,

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

@@ -7,7 +7,7 @@ from typing import ClassVar
 import torch
 
 from vllm import _custom_ops as ops
-from vllm._aiter_ops import rocm_aiter_ops
+from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops
 from vllm.config import get_current_vllm_config_or_none
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
@@ -359,6 +359,9 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
     def fused_qk_norm_rope_kvcache_supported(self):
         return rocm_aiter_ops.is_enabled()
 
+    def fused_qk_norm_mrope_kvcache_supported(self) -> bool:
+        return IS_AITER_FOUND
+
     def do_qk_norm_rope_kvcache_update(
         self,
         layer: AttentionLayer,
@@ -397,6 +400,47 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
             use_shuffle_layout=False,
         )
 
+    def do_qk_norm_mrope_kvcache_update(
+        self,
+        layer: AttentionLayer,
+        qkv: torch.Tensor,
+        q_out: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        rms_norm_eps: float,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        mrope_section: tuple[int, int, int],
+        is_interleaved: bool,
+        rotary_dim: int,
+        kv_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+    ):
+        key_cache, value_cache = self._split_kv_cache(kv_cache)
+        rocm_aiter_ops.do_qk_norm_mrope_kvcache_update(
+            qkv=qkv,
+            q_weight=q_weight,
+            k_weight=k_weight,
+            cos_sin_cache=cos_sin_cache,
+            positions=positions,
+            num_heads_q=self.num_heads,
+            num_heads_k=self.num_kv_heads,
+            head_dim=self.head_size,
+            is_neox=is_neox,
+            mrope_section=list(mrope_section),
+            is_interleaved=is_interleaved,
+            rms_norm_eps=rms_norm_eps,
+            q_out=q_out,
+            key_cache=key_cache,
+            value_cache=value_cache,
+            slot_mapping=layer_slot_mapping,
+            k_scale=layer._k_scale_cpu,
+            v_scale=layer._v_scale_cpu,
+            kv_cache_dtype=self.kv_cache_dtype,
+            rotary_dim=rotary_dim,
+        )
+
     def do_rope_and_kv_cache_update(
         self,
         layer: AttentionLayer,