diff --git a/tests/compile/fusions_e2e/common.py b/tests/compile/fusions_e2e/common.py index 77c3ba5112..1f92ff5a19 100644 --- a/tests/compile/fusions_e2e/common.py +++ b/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" diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index f3df3c358e..a2cbe0135b 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/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). diff --git a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py index 6965346a58..ee557235fd 100644 --- a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py +++ b/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")) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index e5ce96cf24..0a28c06ee7 100644 --- a/vllm/_aiter_ops.py +++ b/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, diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 7fd44dd005..6d959c3067 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/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.""" diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py index 80eeb648f4..f418e8a7ef 100644 --- a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py +++ b/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), + ) diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 8ccc7332bd..eefcf1f49e 100644 --- a/vllm/config/compilation.py +++ b/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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index e1c80be67b..35b214f980 100644 --- a/vllm/config/vllm.py +++ b/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(): diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 185bfd650c..42013926d7 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/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, diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index afafa3c925..a403c9ae39 100644 --- a/vllm/v1/attention/backend.py +++ b/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, diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 9e3267c526..711a8f8755 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/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,