From 37fbf5208413486e5553c1b287af64f6af6afdcf Mon Sep 17 00:00:00 2001 From: fanxingran Date: Mon, 10 Aug 2026 23:42:10 +0800 Subject: [PATCH] [ROCm][MLA] [K3] Fix fp8 KV cache decode on the AITER MLA backend (#51011) Signed-off-by: fanxingran --- .../test_rocm_aiter_mla_head_padding.py | 25 +-- .../test_rocm_aiter_mla_fp8_decode_routing.py | 154 ++++++++++++++++++ .../test_rocm_aiter_mla_mtp_split.py | 151 +++++++++++++++-- .../attention/backends/mla/rocm_aiter_mla.py | 99 +++++++---- 4 files changed, 366 insertions(+), 63 deletions(-) create mode 100644 tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py diff --git a/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py b/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py index 82ddce3164..744f1852ba 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py @@ -163,22 +163,22 @@ def test_is_valid_num_heads(): def test_nondivisor_and_multitoken_never_use_gluon(): # Non-divisor decode always takes the asm path (12 heads/rank at TP8). - assert not AiterMLAHelper.use_gluon_decode(12, 1) - assert not AiterMLAHelper.use_gluon_decode(6, 1) + assert not AiterMLAHelper.use_gluon_decode(12, 1, "auto") + assert not AiterMLAHelper.use_gluon_decode(6, 1, "auto") # >=16 heads never pad, never Gluon. - assert not AiterMLAHelper.use_gluon_decode(16, 1) + assert not AiterMLAHelper.use_gluon_decode(16, 1, "auto") # Multi-token (verify / qlen>1) is never the single-token Gluon decode. - assert not AiterMLAHelper.use_gluon_decode(8, 4) - assert not AiterMLAHelper.use_gluon_decode(12, 4) + assert not AiterMLAHelper.use_gluon_decode(8, 4, "auto") + assert not AiterMLAHelper.use_gluon_decode(12, 4, "auto") def test_divisor_gluon_selection_follows_arch(): # Divisor head counts keep Gluon only where the kernel exists (gfx950). # On gfx942 / non-ROCm they must route to the asm persistent decode. on_gfx950 = _on_gfx950() - assert AiterMLAHelper.use_gluon_decode(8, 1) is on_gfx950 - assert AiterMLAHelper.use_gluon_decode(4, 1) is on_gfx950 - assert AiterMLAHelper.use_gluon_decode(1, 1) is on_gfx950 + assert AiterMLAHelper.use_gluon_decode(8, 1, "auto") is on_gfx950 + assert AiterMLAHelper.use_gluon_decode(4, 1, "auto") is on_gfx950 + assert AiterMLAHelper.use_gluon_decode(1, 1, "auto") is on_gfx950 def test_asm_padding_env_default_is_auto(monkeypatch): @@ -192,7 +192,7 @@ def test_asm_padding_env_force_asm_disables_gluon(monkeypatch): monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "asm") # Forcing the asm path: no small-head count uses Gluon on any arch. for num_heads in (1, 2, 4, 8, 12): - assert not AiterMLAHelper.use_gluon_decode(num_heads, 1) + assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, "auto") def test_asm_padding_env_force_gluon_follows_arch(monkeypatch): @@ -202,15 +202,15 @@ def test_asm_padding_env_force_gluon_follows_arch(monkeypatch): # (gfx950), including non-divisor counts like 12; gfx942/non-ROCm still # falls back to the asm path. for num_heads in (1, 2, 4, 8, 12): - assert AiterMLAHelper.use_gluon_decode(num_heads, 1) is on_gfx950 + assert AiterMLAHelper.use_gluon_decode(num_heads, 1, "auto") is on_gfx950 def test_asm_padding_env_auto_matches_arch_gate(monkeypatch): monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "auto") on_gfx950 = _on_gfx950() # auto: divisor counts keep Gluon on gfx950, non-divisor counts take asm. - assert AiterMLAHelper.use_gluon_decode(8, 1) is on_gfx950 - assert not AiterMLAHelper.use_gluon_decode(12, 1) + assert AiterMLAHelper.use_gluon_decode(8, 1, "auto") is on_gfx950 + assert not AiterMLAHelper.use_gluon_decode(12, 1, "auto") @pytest.mark.skipif( @@ -237,6 +237,7 @@ def test_h12_aiter_mla_decode_matches_reference(): impl.num_heads = NUM_HEADS impl.kv_lora_rank = KV_LORA_RANK impl.scale = SCALE + impl.kv_cache_dtype = "auto" one = torch.ones(1, dtype=torch.float32, device=device) layer = types.SimpleNamespace(_q_scale=one, _k_scale=one) diff --git a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py new file mode 100644 index 0000000000..31e16c6250 --- /dev/null +++ b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Decode-kernel routing for the ROCm AITER MLA backend. + +Gluon exposes a single fp8-KV regime, bh16bn128. It is a bf16-query kernel that +upcasts the cache in registers with a hardcoded scale of 1.0, and it asserts +batch_size == 1, so it cannot serve a decode batch. Every fp8 shape therefore +has to land on the asm kernels, which ship real fp8 variants for gqa=16. These +tests pin that down at the predicates, since the failure it prevents is either a +batch assertion or -- worse -- a silently wrong result. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm AITER MLA tests", allow_module_level=True) + +from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402 +from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper # noqa: E402 + +FP8_DTYPES = ["fp8", "fp8_e4m3", "fp8_e5m2"] +UNQUANTIZED_DTYPES = ["auto", "bfloat16"] + + +@pytest.fixture +def gluon_available(monkeypatch): + """Pin the arch gate and the mode knob so only the dtype rules are in play. + + Gluon ships a gfx950 build only, and VLLM_ROCM_AITER_MLA_ASM_PADDING can + force the asm path on any arch. Without pinning both, the expectations that + bf16 *keeps* Gluon would pass on a gfx950 host and fail on gfx942 for + reasons that have nothing to do with what these tests cover. + """ + monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True) + monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto") + + +@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES) +@pytest.mark.parametrize("num_heads", [1, 2, 4, 5, 6, 8, 12, 16, 32, 128]) +@pytest.mark.parametrize("max_qo_len", [1, 2, 4, 5, 8, 15]) +def test_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len): + """No fp8 shape reaches either Gluon entry point. + + The head count is deliberately swept across divisors of 16 as well as + non-divisors: the divisor case (e.g. 8 heads at TP8) is the one that stays + on Gluon for bf16 and so is the one an fp8 guard has to override. + """ + assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype) + assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype) + + +@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES) +@pytest.mark.parametrize("num_heads", [1, 2, 4, 8, 12]) +@pytest.mark.parametrize("mode", ["auto", "gluon", "asm"]) +def test_fp8_never_routes_to_gluon_under_any_mode( + monkeypatch, kv_cache_dtype, num_heads, mode +): + """VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force an fp8 cache onto Gluon. + + The dtype guard deliberately precedes the mode knob: honouring an explicit + "gluon" request under fp8 would hand Gluon the batch it asserts against, so + it is overridden rather than obeyed. Pinned here because the override is a + correctness decision, not a preference. + """ + monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True) + monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: mode) + assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype) + assert not AiterMLAHelper.use_gluon_verify(num_heads, 8, kv_cache_dtype) + + +@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES) +@pytest.mark.parametrize("num_heads", [16, 17, 32, 64, 128]) +@pytest.mark.parametrize("max_qo_len", [1, 4, 8]) +def test_large_head_counts_never_use_gluon(kv_cache_dtype, num_heads, max_qo_len): + """>= 16 heads has always been asm-only; the dtype guard must not change it.""" + assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype) + assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype) + + +@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES) +@pytest.mark.parametrize("num_heads", [1, 2, 4, 8]) +def test_unquantized_divisor_heads_keep_gluon_decode( + gluon_available, kv_cache_dtype, num_heads +): + """Divisor head counts keep the Gluon single-token decode for bf16. + + Gluon does not scale with KV length the way the asm persistent decode does, + so this is the faster path where it is usable. + """ + assert AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype) + + +@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES) +@pytest.mark.parametrize("num_heads", [3, 5, 6, 7, 9, 12, 15]) +def test_unquantized_non_divisor_heads_use_asm_decode(kv_cache_dtype, num_heads): + """Non-divisor head counts are padded to 16 and take the asm decode.""" + assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype) + + +@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES) +@pytest.mark.parametrize("num_heads", [5, 8, 12]) +@pytest.mark.parametrize("max_qo_len", [2, 4, 8, 15]) +def test_unquantized_small_head_verify_keeps_gluon( + gluon_available, kv_cache_dtype, num_heads, max_qo_len +): + """bf16 has no gqa<16, qseqlen>1 asm kernel, so verify still flattens onto Gluon. + + Unlike the decode predicate, this one does not care whether the head count + divides 16 -- the flatten reshapes to qseqlen=1 either way. + """ + assert AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype) + # The verify flatten is a separate entry point from the single-token decode. + assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype) + + +@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES) +@pytest.mark.parametrize("num_heads", [5, 8, 12, 16, 32]) +def test_decode_and_verify_are_disjoint(kv_cache_dtype, num_heads): + """The two Gluon entry points partition on qlen and never both fire.""" + for max_qo_len in (1, 2, 8): + assert not ( + AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype) + and AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype) + ) + + +@pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15]) +def test_padded_query_is_contiguous(num_heads): + """asm_mla.cu:805 requires Q.is_contiguous(). + + Padding a non-divisor head count to 16 tiles the query and slices it back + down, which yields a non-contiguous view whenever more than one tile is + needed (12 heads -> repeat to 24 -> slice to 16). The asm kernel rejects + that outright, so the padding has to materialize the result. + """ + q = torch.randn(4, num_heads, 576) + padded = AiterMLAHelper.get_mla_padded_q(num_heads, q) + + assert padded.shape[1] == max(16, num_heads) + assert padded.is_contiguous() + + +@pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15, 16, 32]) +def test_pad_unpad_round_trip_preserves_head_order(num_heads): + """Unpadding recovers each original head from the padded output, in order.""" + q = torch.arange(num_heads, dtype=torch.float32).view(1, num_heads, 1) + padded = AiterMLAHelper.get_mla_padded_q(num_heads, q) + unpadded = AiterMLAHelper.get_mla_unpadded_o(num_heads, padded) + + assert unpadded.shape == q.shape + torch.testing.assert_close(unpadded, q) diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index 844bee0f71..c5eb921a22 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -64,10 +64,12 @@ def _builder( kernel_block_size: int = 1, max_decode_rows: int = 32, num_heads: int = 16, + kv_cache_dtype: str = "auto", ): return SimpleNamespace( device=torch.device("cpu"), num_heads=num_heads, + _kv_cache_dtype_str=kv_cache_dtype, paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32), paged_kv_indices=torch.empty(1024, dtype=torch.int32), paged_kv_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32), @@ -108,8 +110,22 @@ def test_backend_declares_uniform_batch_support(): @pytest.mark.parametrize("num_heads", [8, 16, 32, 64, 128]) -def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads): - """Aiter init passes real MTP qlen/dtypes to rocm_aiter_mla.py:349-355. +@pytest.mark.parametrize( + "spec_method, parallel_drafting", + [ + ("deepseek_mtp", False), + # A drafter that is not one of the historically recognized MTP methods, + # and a parallel one, so the threshold is 1 + 2 * num_spec rather than + # 1 + num_spec. Sizing the metadata off a method name instead leaves + # these at qlen=1 while the router still admits the full range. + ("dspark", True), + ("eagle", False), + ], +) +def test_mtp_builder_init_sizes_native_fp8_metadata( + monkeypatch, num_heads, spec_method, parallel_drafting +): + """Aiter init sizes the metadata for every query length decode can be handed. Sweeping num_heads asserts the max(16, num_heads) clamp is what sizes the metadata, covering the fp8 nhead=32 (TP4) fold path. @@ -143,6 +159,13 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads): def init_common_builder(self, *args, **kwargs): self.num_heads = num_heads + # Mirror what _init_reorder_batch_threshold would have left behind: the + # metadata is sized from the routing threshold, so a stub that skips it + # would size for qlen=1 and hide the very mismatch this test covers. + spec = config.speculative_config + self.reorder_batch_threshold = ( + 1 + (2 if spec.parallel_drafting else 1) * spec.num_speculative_tokens + ) monkeypatch.setitem( sys.modules, @@ -161,8 +184,9 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads): config = SimpleNamespace( speculative_config=SimpleNamespace( - method="deepseek_mtp", + method=spec_method, num_speculative_tokens=3, + parallel_drafting=parallel_drafting, ), parallel_config=SimpleNamespace(tensor_parallel_size=8), model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16), @@ -182,7 +206,7 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads): assert info_calls == [ { "max_batch_size": config.scheduler_config.max_num_seqs, - "max_qo_len": config.speculative_config.num_speculative_tokens + 1, + "max_qo_len": builder.reorder_batch_threshold, "num_attention_heads": max(16, num_heads), "q_dtype": dtypes.fp8, "kv_dtype": dtypes.fp8, @@ -352,26 +376,46 @@ def test_decode_expands_kernel_block_page_indices(monkeypatch): @pytest.mark.parametrize( - "mtp_decode_qlen, qo_len, num_heads, expect_persistent", + "mtp_decode_qlen, qo_len, num_heads, kv_cache_dtype, expect_persistent", [ - (1, 1, 16, True), # non-MTP decode - (4, 2, 16, True), # MTP deployment, in-range step - (4, 4, 16, True), # MTP deployment, full-qlen verification step - (2, 4, 16, False), # step demand exceeds provisioned K -> fallback - (1, 1, 8, False), # small head count -> Gluon decode owns qlen==1 - (4, 4, 8, False), # small head count -> Gluon flatten owns qlen>1 + (1, 1, 16, "auto", True), # non-MTP decode + (4, 2, 16, "auto", True), # MTP deployment, in-range step + (4, 4, 16, "auto", True), # MTP deployment, full-qlen verification step + (2, 4, 16, "auto", False), # step demand exceeds provisioned K -> fallback + (1, 1, 8, "auto", False), # divisor head count -> Gluon decode owns qlen==1 + (4, 4, 8, "auto", False), # small head count -> Gluon flatten owns qlen>1 + # A non-divisor head count is padded to 16 and runs the asm decode, so it + # needs the schedule even though num_heads < 16. Its verify still + # flattens onto Gluon though -- divisibility only steers the qlen==1 + # decode, and bf16 has no small-head multi-token asm kernel either way. + (1, 1, 12, "auto", True), + (8, 8, 12, "auto", False), + # An fp8 cache never reaches Gluon, at any head count or qlen, so it + # always needs the schedule. (1, 1, 8) is the case the head-count gate + # got wrong: divisor head count, so it read as Gluon-owned. + (1, 1, 8, "fp8", True), + (1, 1, 12, "fp8", True), + (1, 1, 16, "fp8", True), + # DSpark verify shape. The fp8 fold path rejects non-persistent outright + # once qlen > 4, so these must not fall through. + (8, 8, 8, "fp8", True), + (8, 8, 12, "fp8", True), ], ) def test_persistent_metadata_gate( - monkeypatch, mtp_decode_qlen, qo_len, num_heads, expect_persistent + monkeypatch, mtp_decode_qlen, qo_len, num_heads, kv_cache_dtype, expect_persistent ): - """Persistent metadata is passed iff num_heads >= 16 and 1 <= max_qo_len <= K. + """Persistent metadata is passed iff the asm decode runs and 1 <= qlen <= K. K = _mtp_decode_qlen sizes the metadata buffers at init; a decode step gets the pre-built schedule only when its qlen fits those buffers, otherwise it falls back to the kernel computing its own. qlen==1 (non-MTP) must stay - in-range -- dropping it is the regression this guards. Fewer than 16 heads - is served by the Gluon decode paths, which never read this schedule. + in-range -- dropping it is the regression this guards. + + Only the Gluon paths ignore the schedule, so the gate follows the routing + predicates rather than the raw head count. Reading `num_heads >= 16` instead + denies the schedule to two sets of shapes that do run the asm kernels: a + non-divisor head count padded up to 16, and any fp8 cache below 16 heads. """ get_mla_metadata_v1 = mock.MagicMock() monkeypatch.setitem( @@ -382,6 +426,12 @@ def test_persistent_metadata_gate( monkeypatch.setattr( rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() ) + # The gate follows the routing now, so the small-head bf16 rows above only + # hold where a Gluon build exists. Pin the arch they describe: on gfx942 + # those shapes run the asm decode and do get the schedule, which is a + # different assertion, made arch-free in the fp8 routing test. + monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True) + monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto") # Uniform, non-CUDA-graph batch: every request has exactly qo_len tokens, so # num_decode_tokens == sum(qo_len) and no dummy-row padding kicks in. @@ -390,7 +440,11 @@ def test_persistent_metadata_gate( 0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32 ) metadata = AiterMLAMetadataBuilder._build_decode( - _builder(mtp_decode_qlen=mtp_decode_qlen, num_heads=num_heads), + _builder( + mtp_decode_qlen=mtp_decode_qlen, + num_heads=num_heads, + kv_cache_dtype=kv_cache_dtype, + ), block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), seq_lens_device=torch.tensor([8, 8], dtype=torch.int32), max_seq_len=8, @@ -406,3 +460,68 @@ def test_persistent_metadata_gate( if expect_persistent: assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == qo_len assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == qo_len + + +@pytest.mark.parametrize( + "qo_len, num_heads, kv_cache_dtype, expect_persistent", + [ + # A padded rank keeps the schedule wherever a persistent kernel exists: + # bf16 up to qseqlen 4, fp8 at every qlen (its fold requires it). + (1, 12, "auto", True), + (4, 12, "auto", True), + (1, 12, "fp8", True), + (8, 12, "fp8", True), + # Past qseqlen 4 bf16 has only the non-persistent entry, and the fold + # that would reach a persistent one is gfx950-only, so the padded rank + # must fall back rather than ask for a kernel that is not built. + (8, 12, "auto", False), + # A native 16-head rank is unchanged: it took the schedule before this + # backend keyed the gate off the routing, and still does. + (8, 16, "auto", True), + ], +) +def test_persistent_metadata_gate_without_gluon_build( + monkeypatch, qo_len, num_heads, kv_cache_dtype, expect_persistent +): + """The gate on an arch with no Gluon build, i.e. gfx942. + + Both Gluon predicates read False there, so a small-head shape reaches the + padded asm decode instead of being flattened. The schedule then has to + follow what the asm dispatch actually ships: bf16 has no gqa=16 persistent + kernel above qseqlen 4, and the q-row fold that reaches one on gfx950 is + guarded on that arch, so asking for the schedule there selects a kernel + that does not exist. + """ + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: False) + monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto") + + num_reqs = 2 + query_start_loc = torch.arange( + 0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32 + ) + metadata = AiterMLAMetadataBuilder._build_decode( + _builder( + mtp_decode_qlen=qo_len, + num_heads=num_heads, + kv_cache_dtype=kv_cache_dtype, + ), + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=torch.tensor([8, 8], dtype=torch.int32), + max_seq_len=8, + query_start_loc_cpu=query_start_loc, + query_start_loc_device=query_start_loc, + num_decode_tokens=num_reqs * qo_len, + dcp_tot_seq_lens_device=None, + ) + + assert metadata.has_persistent_metadata is expect_persistent + assert get_mla_metadata_v1.called is expect_persistent diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index f10b9cff36..ccbdf72890 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -26,7 +26,7 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, MultipleOf, ) -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache logger = init_logger(__name__) @@ -269,17 +269,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): self.compilation_config = vllm_config.compilation_config self.decode_attn_out_dtype = vllm_config.model_config.dtype - # MTP/deepseek_mtp verification runs decode with qlen = num_spec + 1; - # any other config (including no spec) stays at single-token decode. - speculative_config = vllm_config.speculative_config - if ( - speculative_config is not None - and speculative_config.method in ("mtp", "deepseek_mtp") - and speculative_config.num_speculative_tokens is not None - ): - self._mtp_decode_qlen = int(speculative_config.num_speculative_tokens) + 1 - else: - self._mtp_decode_qlen = 1 + # reorder_batch_threshold is the largest query length decode can be + # handed, and already accounts for the drafting scheme. A method-name + # whitelist sizes unlisted drafters for qlen=1, which closes the + # persistent gate below and makes aiter raise a KeyError mid-run. + self._mtp_decode_qlen = self.reorder_batch_threshold or 1 # Store the kernel block size from the spec. When kernel_block_size=1 # (no spec-dec), behavior is identical to the original. When > 1 @@ -327,6 +321,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): torch.float16: dtypes.fp16, torch.bfloat16: dtypes.bf16, }[kv_cache_spec.dtype] + # _build_decode needs the cache dtype to pick the decode kernel; keep + # the normalized string instead of dropping it at the end of __init__. + self._kv_cache_dtype_str = kv_cache_dtype_str # MLAAttention quantizes decode Q to FP8 before calling this backend # whenever the KV cache is FP8 and supports_quant_query_input is true. q_dtype = ( @@ -624,7 +621,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): ] ) use_gluon_decode = AiterMLAHelper.use_gluon_decode( - self.num_heads, int(max_qo_len) + self.num_heads, int(max_qo_len), self._kv_cache_dtype_str ) if self.compilation_config.cudagraph_mode.has_full_cudagraphs(): @@ -693,13 +690,26 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): else: qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs] - # Pass persistent metadata for every uniform decode we sized buffers for - # (normal qlen==1 through MTP verification qlen==K): the fp8 nhead=32 fold - # path breaks without it. qlen>K falls back to kernel-internal metadata. - # Small-head (<16) decode takes the Gluon paths and never consumes it. + # Only the asm decode consumes the schedule, so gate on the routing + # rather than on num_heads >= 16, which denies it to a padded rank + # running the same asm kernels. The two predicates are disjoint -- + # decode is qlen==1, verify is qlen>1 -- and cover both Gluon entries. has_persistent_metadata = False use_persistent_metadata = ( - self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS + not AiterMLAHelper.use_gluon_decode( + self.num_heads, max_qo_len, self._kv_cache_dtype_str + ) + and not AiterMLAHelper.use_gluon_verify( + self.num_heads, max_qo_len, self._kv_cache_dtype_str + ) + # A padded rank has no bf16 persistent kernel past qlen 4 where the + # gfx950 fold is absent; the non-persistent entry covers it. fp8 + # keeps the schedule -- its fold rejects non-persistent outright. + and ( + self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS + or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN + or is_quantized_kv_cache(self._kv_cache_dtype_str) + ) and max_qo_len >= 1 and max_qo_len <= self._mtp_decode_qlen ) @@ -833,6 +843,10 @@ class AiterMLAHelper: """ _AITER_MIN_MLA_HEADS: Final = 16 + # Largest qlen the padded gqa=16 asm decode has a bf16 persistent kernel + # for. Above it only the non-persistent qseqlen=8 entry exists, and the + # fold that reaches a persistent one is gfx950-only. + _ASM_PADDED_MAX_PS_QLEN: Final = 4 _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = () @staticmethod @@ -888,17 +902,19 @@ class AiterMLAHelper: return o[:, :num_heads, :] @staticmethod - def use_gluon_decode(num_heads: int, max_qo_len: int) -> bool: - # Small-head (<16) single-token decode can use either the Gluon kernel - # or the padded asm persistent decode, selected by - # VLLM_ROCM_AITER_MLA_ASM_PADDING (see _aiter_mla_small_head_mode) and - # the arch: Gluon only has a gfx950 build. In "auto" (default) mode - # divisor counts keep Gluon on gfx950 and everything else -- non-divisor - # counts (e.g. 12 heads/rank at TP8) and all counts on gfx942 -- takes - # the asm path, which get_mla_padded_q pads to exactly 16. + def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool: + # Small-head (<16) single-token decode takes either the Gluon kernel or + # the padded asm persistent decode, selected by + # VLLM_ROCM_AITER_MLA_ASM_PADDING and the arch (Gluon is gfx950 only). m = AiterMLAHelper._AITER_MIN_MLA_HEADS if num_heads >= m or max_qo_len != 1: return False + # Gluon's only fp8-KV regime, bh16bn128, is a bf16-query kernel with a + # hardcoded scale that asserts batch_size == 1, so it cannot serve a + # decode batch. Checked before the mode knob: an explicit "gluon" + # request under fp8 would assert immediately. + if is_quantized_kv_cache(kv_cache_dtype): + return False mode = _aiter_mla_small_head_mode() if mode == "asm": return False @@ -907,6 +923,23 @@ class AiterMLAHelper: return gluon_supported return m % num_heads == 0 and gluon_supported + @staticmethod + def use_gluon_verify(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool: + """Whether a small-head multi-token verify is flattened onto Gluon. + + bf16 has no gqa<16, qseqlen>1 asm kernel, so the verify is flattened + into per-token Gluon decodes. fp8 has one via the q-row fold and must + not come here: the flatten hands Gluon the batch size its fp8 regime + asserts against. A predicate rather than inline in forward_mqa so the + builder sees the same answer the impl acts on. + """ + if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= 1: + return False + if is_quantized_kv_cache(kv_cache_dtype): + return False + # Same arch and mode gating as use_gluon_decode. + return _aiter_mla_small_head_mode() != "asm" and _gluon_mla_decode_supported() + class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): def __init__( @@ -1191,15 +1224,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): # target is checking draft tokens, so position t must not see t+1 -- # and attention rows are independent, so giving row t the KV range # [0, context + t] is exactly causal multi-token attention. - # Gluon only has a gfx950 build, so gate on the arch (mirrors - # use_gluon_decode). VLLM_ROCM_AITER_MLA_ASM_PADDING=asm also forces the - # asm path here. Otherwise this falls through to the asm decode, which - # pads to 16 heads and handles qlen>1 verify directly. - if ( - self.num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS - and int(decode.max_qo_len) > 1 - and _aiter_mla_small_head_mode() != "asm" - and _gluon_mla_decode_supported() + # Arch, mode and dtype gating all live in use_gluon_verify, so that the + # builder -- which has to know whether the asm decode will run -- sees + # the same answer as this branch. + if AiterMLAHelper.use_gluon_verify( + self.num_heads, int(decode.max_qo_len), self.kv_cache_dtype ): qlen = int(decode.max_qo_len) if type(q) is tuple: