[ROCm][MLA] [K3] Fix fp8 KV cache decode on the AITER MLA backend (#51011)
Signed-off-by: fanxingran <xingran.fan@amd.com>
fanxingran committed
Aug 10, 2026 at 23:42 UTC
37fbf5208413486e5553c1b287af64f6af6afdcf
4 files changed
+366
-63
tests/kernels/attention/test_rocm_aiter_mla_head_padding.py
+13
-12
@@ -163,22 +163,22 @@ def test_is_valid_num_heads():
163
164
def test_nondivisor_and_multitoken_never_use_gluon():
165
# Non-divisor decode always takes the asm path (12 heads/rank at TP8).
166
- assert not AiterMLAHelper.use_gluon_decode(12, 1)
167
- assert not AiterMLAHelper.use_gluon_decode(6, 1)
166
+ assert not AiterMLAHelper.use_gluon_decode(12, 1, "auto")
167
+ assert not AiterMLAHelper.use_gluon_decode(6, 1, "auto")
168
# >=16 heads never pad, never Gluon.
169
- assert not AiterMLAHelper.use_gluon_decode(16, 1)
169
+ assert not AiterMLAHelper.use_gluon_decode(16, 1, "auto")
170
# Multi-token (verify / qlen>1) is never the single-token Gluon decode.
171
- assert not AiterMLAHelper.use_gluon_decode(8, 4)
172
- assert not AiterMLAHelper.use_gluon_decode(12, 4)
171
+ assert not AiterMLAHelper.use_gluon_decode(8, 4, "auto")
172
+ assert not AiterMLAHelper.use_gluon_decode(12, 4, "auto")
173
174
175
def test_divisor_gluon_selection_follows_arch():
176
# Divisor head counts keep Gluon only where the kernel exists (gfx950).
177
# On gfx942 / non-ROCm they must route to the asm persistent decode.
178
on_gfx950 = _on_gfx950()
179
- assert AiterMLAHelper.use_gluon_decode(8, 1) is on_gfx950
180
- assert AiterMLAHelper.use_gluon_decode(4, 1) is on_gfx950
181
- assert AiterMLAHelper.use_gluon_decode(1, 1) is on_gfx950
179
+ assert AiterMLAHelper.use_gluon_decode(8, 1, "auto") is on_gfx950
180
+ assert AiterMLAHelper.use_gluon_decode(4, 1, "auto") is on_gfx950
181
+ assert AiterMLAHelper.use_gluon_decode(1, 1, "auto") is on_gfx950
182
183
184
def test_asm_padding_env_default_is_auto(monkeypatch):
@@ -192,7 +192,7 @@ def test_asm_padding_env_force_asm_disables_gluon(monkeypatch):
192
monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "asm")
193
# Forcing the asm path: no small-head count uses Gluon on any arch.
194
for num_heads in (1, 2, 4, 8, 12):
195
- assert not AiterMLAHelper.use_gluon_decode(num_heads, 1)
195
+ assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, "auto")
196
197
198
def test_asm_padding_env_force_gluon_follows_arch(monkeypatch):
@@ -202,15 +202,15 @@ def test_asm_padding_env_force_gluon_follows_arch(monkeypatch):
202
# (gfx950), including non-divisor counts like 12; gfx942/non-ROCm still
203
# falls back to the asm path.
204
for num_heads in (1, 2, 4, 8, 12):
205
- assert AiterMLAHelper.use_gluon_decode(num_heads, 1) is on_gfx950
205
+ assert AiterMLAHelper.use_gluon_decode(num_heads, 1, "auto") is on_gfx950
206
207
208
def test_asm_padding_env_auto_matches_arch_gate(monkeypatch):
209
monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "auto")
210
on_gfx950 = _on_gfx950()
211
# auto: divisor counts keep Gluon on gfx950, non-divisor counts take asm.
212
- assert AiterMLAHelper.use_gluon_decode(8, 1) is on_gfx950
213
- assert not AiterMLAHelper.use_gluon_decode(12, 1)
212
+ assert AiterMLAHelper.use_gluon_decode(8, 1, "auto") is on_gfx950
213
+ assert not AiterMLAHelper.use_gluon_decode(12, 1, "auto")
214
215
216
@pytest.mark.skipif(
@@ -237,6 +237,7 @@ def test_h12_aiter_mla_decode_matches_reference():
237
impl.num_heads = NUM_HEADS
238
impl.kv_lora_rank = KV_LORA_RANK
239
impl.scale = SCALE
240
+ impl.kv_cache_dtype = "auto"
241
242
one = torch.ones(1, dtype=torch.float32, device=device)
243
layer = types.SimpleNamespace(_q_scale=one, _k_scale=one)
tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py
new
+154
@@ -0,0 +1,154 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Decode-kernel routing for the ROCm AITER MLA backend.
4
+
5
+Gluon exposes a single fp8-KV regime, bh16bn128. It is a bf16-query kernel that
6
+upcasts the cache in registers with a hardcoded scale of 1.0, and it asserts
7
+batch_size == 1, so it cannot serve a decode batch. Every fp8 shape therefore
8
+has to land on the asm kernels, which ship real fp8 variants for gqa=16. These
9
+tests pin that down at the predicates, since the failure it prevents is either a
10
+batch assertion or -- worse -- a silently wrong result.
11
+"""
12
+
13
+import pytest
14
+import torch
15
+
16
+from vllm.platforms import current_platform
17
+
18
+if not current_platform.is_rocm():
19
+ pytest.skip("ROCm AITER MLA tests", allow_module_level=True)
20
+
21
+from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402
22
+from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper # noqa: E402
23
+
24
+FP8_DTYPES = ["fp8", "fp8_e4m3", "fp8_e5m2"]
25
+UNQUANTIZED_DTYPES = ["auto", "bfloat16"]
26
+
27
+
28
+@pytest.fixture
29
+def gluon_available(monkeypatch):
30
+ """Pin the arch gate and the mode knob so only the dtype rules are in play.
31
+
32
+ Gluon ships a gfx950 build only, and VLLM_ROCM_AITER_MLA_ASM_PADDING can
33
+ force the asm path on any arch. Without pinning both, the expectations that
34
+ bf16 *keeps* Gluon would pass on a gfx950 host and fail on gfx942 for
35
+ reasons that have nothing to do with what these tests cover.
36
+ """
37
+ monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True)
38
+ monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto")
39
+
40
+
41
+@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES)
42
+@pytest.mark.parametrize("num_heads", [1, 2, 4, 5, 6, 8, 12, 16, 32, 128])
43
+@pytest.mark.parametrize("max_qo_len", [1, 2, 4, 5, 8, 15])
44
+def test_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len):
45
+ """No fp8 shape reaches either Gluon entry point.
46
+
47
+ The head count is deliberately swept across divisors of 16 as well as
48
+ non-divisors: the divisor case (e.g. 8 heads at TP8) is the one that stays
49
+ on Gluon for bf16 and so is the one an fp8 guard has to override.
50
+ """
51
+ assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype)
52
+ assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
53
+
54
+
55
+@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES)
56
+@pytest.mark.parametrize("num_heads", [1, 2, 4, 8, 12])
57
+@pytest.mark.parametrize("mode", ["auto", "gluon", "asm"])
58
+def test_fp8_never_routes_to_gluon_under_any_mode(
59
+ monkeypatch, kv_cache_dtype, num_heads, mode
60
+):
61
+ """VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force an fp8 cache onto Gluon.
62
+
63
+ The dtype guard deliberately precedes the mode knob: honouring an explicit
64
+ "gluon" request under fp8 would hand Gluon the batch it asserts against, so
65
+ it is overridden rather than obeyed. Pinned here because the override is a
66
+ correctness decision, not a preference.
67
+ """
68
+ monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True)
69
+ monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: mode)
70
+ assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype)
71
+ assert not AiterMLAHelper.use_gluon_verify(num_heads, 8, kv_cache_dtype)
72
+
73
+
74
+@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES)
75
+@pytest.mark.parametrize("num_heads", [16, 17, 32, 64, 128])
76
+@pytest.mark.parametrize("max_qo_len", [1, 4, 8])
77
+def test_large_head_counts_never_use_gluon(kv_cache_dtype, num_heads, max_qo_len):
78
+ """>= 16 heads has always been asm-only; the dtype guard must not change it."""
79
+ assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype)
80
+ assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
81
+
82
+
83
+@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
84
+@pytest.mark.parametrize("num_heads", [1, 2, 4, 8])
85
+def test_unquantized_divisor_heads_keep_gluon_decode(
86
+ gluon_available, kv_cache_dtype, num_heads
87
+):
88
+ """Divisor head counts keep the Gluon single-token decode for bf16.
89
+
90
+ Gluon does not scale with KV length the way the asm persistent decode does,
91
+ so this is the faster path where it is usable.
92
+ """
93
+ assert AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype)
94
+
95
+
96
+@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
97
+@pytest.mark.parametrize("num_heads", [3, 5, 6, 7, 9, 12, 15])
98
+def test_unquantized_non_divisor_heads_use_asm_decode(kv_cache_dtype, num_heads):
99
+ """Non-divisor head counts are padded to 16 and take the asm decode."""
100
+ assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype)
101
+
102
+
103
+@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
104
+@pytest.mark.parametrize("num_heads", [5, 8, 12])
105
+@pytest.mark.parametrize("max_qo_len", [2, 4, 8, 15])
106
+def test_unquantized_small_head_verify_keeps_gluon(
107
+ gluon_available, kv_cache_dtype, num_heads, max_qo_len
108
+):
109
+ """bf16 has no gqa<16, qseqlen>1 asm kernel, so verify still flattens onto Gluon.
110
+
111
+ Unlike the decode predicate, this one does not care whether the head count
112
+ divides 16 -- the flatten reshapes to qseqlen=1 either way.
113
+ """
114
+ assert AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
115
+ # The verify flatten is a separate entry point from the single-token decode.
116
+ assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype)
117
+
118
+
119
+@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES)
120
+@pytest.mark.parametrize("num_heads", [5, 8, 12, 16, 32])
121
+def test_decode_and_verify_are_disjoint(kv_cache_dtype, num_heads):
122
+ """The two Gluon entry points partition on qlen and never both fire."""
123
+ for max_qo_len in (1, 2, 8):
124
+ assert not (
125
+ AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype)
126
+ and AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
127
+ )
128
+
129
+
130
+@pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15])
131
+def test_padded_query_is_contiguous(num_heads):
132
+ """asm_mla.cu:805 requires Q.is_contiguous().
133
+
134
+ Padding a non-divisor head count to 16 tiles the query and slices it back
135
+ down, which yields a non-contiguous view whenever more than one tile is
136
+ needed (12 heads -> repeat to 24 -> slice to 16). The asm kernel rejects
137
+ that outright, so the padding has to materialize the result.
138
+ """
139
+ q = torch.randn(4, num_heads, 576)
140
+ padded = AiterMLAHelper.get_mla_padded_q(num_heads, q)
141
+
142
+ assert padded.shape[1] == max(16, num_heads)
143
+ assert padded.is_contiguous()
144
+
145
+
146
+@pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15, 16, 32])
147
+def test_pad_unpad_round_trip_preserves_head_order(num_heads):
148
+ """Unpadding recovers each original head from the padded output, in order."""
149
+ q = torch.arange(num_heads, dtype=torch.float32).view(1, num_heads, 1)
150
+ padded = AiterMLAHelper.get_mla_padded_q(num_heads, q)
151
+ unpadded = AiterMLAHelper.get_mla_unpadded_o(num_heads, padded)
152
+
153
+ assert unpadded.shape == q.shape
154
+ torch.testing.assert_close(unpadded, q)
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+135
-16
@@ -64,10 +64,12 @@ def _builder(
64
kernel_block_size: int = 1,
65
max_decode_rows: int = 32,
66
num_heads: int = 16,
67
+ kv_cache_dtype: str = "auto",
68
):
69
return SimpleNamespace(
70
device=torch.device("cpu"),
71
num_heads=num_heads,
72
+ _kv_cache_dtype_str=kv_cache_dtype,
73
paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32),
74
paged_kv_indices=torch.empty(1024, dtype=torch.int32),
75
paged_kv_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32),
@@ -108,8 +110,22 @@ def test_backend_declares_uniform_batch_support():
110
111
112
@pytest.mark.parametrize("num_heads", [8, 16, 32, 64, 128])
111
-def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads):
112
- """Aiter init passes real MTP qlen/dtypes to rocm_aiter_mla.py:349-355.
113
+@pytest.mark.parametrize(
114
+ "spec_method, parallel_drafting",
115
+ [
116
+ ("deepseek_mtp", False),
117
+ # A drafter that is not one of the historically recognized MTP methods,
118
+ # and a parallel one, so the threshold is 1 + 2 * num_spec rather than
119
+ # 1 + num_spec. Sizing the metadata off a method name instead leaves
120
+ # these at qlen=1 while the router still admits the full range.
121
+ ("dspark", True),
122
+ ("eagle", False),
123
+ ],
124
+)
125
+def test_mtp_builder_init_sizes_native_fp8_metadata(
126
+ monkeypatch, num_heads, spec_method, parallel_drafting
127
+):
128
+ """Aiter init sizes the metadata for every query length decode can be handed.
129
130
Sweeping num_heads asserts the max(16, num_heads) clamp is what sizes the
131
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):
159
160
def init_common_builder(self, *args, **kwargs):
161
self.num_heads = num_heads
162
+ # Mirror what _init_reorder_batch_threshold would have left behind: the
163
+ # metadata is sized from the routing threshold, so a stub that skips it
164
+ # would size for qlen=1 and hide the very mismatch this test covers.
165
+ spec = config.speculative_config
166
+ self.reorder_batch_threshold = (
167
+ 1 + (2 if spec.parallel_drafting else 1) * spec.num_speculative_tokens
168
+ )
169
170
monkeypatch.setitem(
171
sys.modules,
@@ -161,8 +184,9 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads):
184
185
config = SimpleNamespace(
186
speculative_config=SimpleNamespace(
164
- method="deepseek_mtp",
187
+ method=spec_method,
188
num_speculative_tokens=3,
189
+ parallel_drafting=parallel_drafting,
190
),
191
parallel_config=SimpleNamespace(tensor_parallel_size=8),
192
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):
206
assert info_calls == [
207
{
208
"max_batch_size": config.scheduler_config.max_num_seqs,
185
- "max_qo_len": config.speculative_config.num_speculative_tokens + 1,
209
+ "max_qo_len": builder.reorder_batch_threshold,
210
"num_attention_heads": max(16, num_heads),
211
"q_dtype": dtypes.fp8,
212
"kv_dtype": dtypes.fp8,
@@ -352,26 +376,46 @@ def test_decode_expands_kernel_block_page_indices(monkeypatch):
376
377
378
@pytest.mark.parametrize(
355
- "mtp_decode_qlen, qo_len, num_heads, expect_persistent",
379
+ "mtp_decode_qlen, qo_len, num_heads, kv_cache_dtype, expect_persistent",
380
[
357
- (1, 1, 16, True), # non-MTP decode
358
- (4, 2, 16, True), # MTP deployment, in-range step
359
- (4, 4, 16, True), # MTP deployment, full-qlen verification step
360
- (2, 4, 16, False), # step demand exceeds provisioned K -> fallback
361
- (1, 1, 8, False), # small head count -> Gluon decode owns qlen==1
362
- (4, 4, 8, False), # small head count -> Gluon flatten owns qlen>1
381
+ (1, 1, 16, "auto", True), # non-MTP decode
382
+ (4, 2, 16, "auto", True), # MTP deployment, in-range step
383
+ (4, 4, 16, "auto", True), # MTP deployment, full-qlen verification step
384
+ (2, 4, 16, "auto", False), # step demand exceeds provisioned K -> fallback
385
+ (1, 1, 8, "auto", False), # divisor head count -> Gluon decode owns qlen==1
386
+ (4, 4, 8, "auto", False), # small head count -> Gluon flatten owns qlen>1
387
+ # A non-divisor head count is padded to 16 and runs the asm decode, so it
388
+ # needs the schedule even though num_heads < 16. Its verify still
389
+ # flattens onto Gluon though -- divisibility only steers the qlen==1
390
+ # decode, and bf16 has no small-head multi-token asm kernel either way.
391
+ (1, 1, 12, "auto", True),
392
+ (8, 8, 12, "auto", False),
393
+ # An fp8 cache never reaches Gluon, at any head count or qlen, so it
394
+ # always needs the schedule. (1, 1, 8) is the case the head-count gate
395
+ # got wrong: divisor head count, so it read as Gluon-owned.
396
+ (1, 1, 8, "fp8", True),
397
+ (1, 1, 12, "fp8", True),
398
+ (1, 1, 16, "fp8", True),
399
+ # DSpark verify shape. The fp8 fold path rejects non-persistent outright
400
+ # once qlen > 4, so these must not fall through.
401
+ (8, 8, 8, "fp8", True),
402
+ (8, 8, 12, "fp8", True),
403
],
404
)
405
def test_persistent_metadata_gate(
366
- monkeypatch, mtp_decode_qlen, qo_len, num_heads, expect_persistent
406
+ monkeypatch, mtp_decode_qlen, qo_len, num_heads, kv_cache_dtype, expect_persistent
407
):
368
- """Persistent metadata is passed iff num_heads >= 16 and 1 <= max_qo_len <= K.
408
+ """Persistent metadata is passed iff the asm decode runs and 1 <= qlen <= K.
409
410
K = _mtp_decode_qlen sizes the metadata buffers at init; a decode step gets
411
the pre-built schedule only when its qlen fits those buffers, otherwise it
412
falls back to the kernel computing its own. qlen==1 (non-MTP) must stay
373
- in-range -- dropping it is the regression this guards. Fewer than 16 heads
374
- is served by the Gluon decode paths, which never read this schedule.
413
+ in-range -- dropping it is the regression this guards.
414
+
415
+ Only the Gluon paths ignore the schedule, so the gate follows the routing
416
+ predicates rather than the raw head count. Reading `num_heads >= 16` instead
417
+ denies the schedule to two sets of shapes that do run the asm kernels: a
418
+ non-divisor head count padded up to 16, and any fp8 cache below 16 heads.
419
"""
420
get_mla_metadata_v1 = mock.MagicMock()
421
monkeypatch.setitem(
@@ -382,6 +426,12 @@ def test_persistent_metadata_gate(
426
monkeypatch.setattr(
427
rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
428
)
429
+ # The gate follows the routing now, so the small-head bf16 rows above only
430
+ # hold where a Gluon build exists. Pin the arch they describe: on gfx942
431
+ # those shapes run the asm decode and do get the schedule, which is a
432
+ # different assertion, made arch-free in the fp8 routing test.
433
+ monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True)
434
+ monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto")
435
436
# Uniform, non-CUDA-graph batch: every request has exactly qo_len tokens, so
437
# num_decode_tokens == sum(qo_len) and no dummy-row padding kicks in.
@@ -390,7 +440,11 @@ def test_persistent_metadata_gate(
440
0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32
441
)
442
metadata = AiterMLAMetadataBuilder._build_decode(
393
- _builder(mtp_decode_qlen=mtp_decode_qlen, num_heads=num_heads),
443
+ _builder(
444
+ mtp_decode_qlen=mtp_decode_qlen,
445
+ num_heads=num_heads,
446
+ kv_cache_dtype=kv_cache_dtype,
447
+ ),
448
block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
449
seq_lens_device=torch.tensor([8, 8], dtype=torch.int32),
450
max_seq_len=8,
@@ -406,3 +460,68 @@ def test_persistent_metadata_gate(
460
if expect_persistent:
461
assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == qo_len
462
assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == qo_len
463
+
464
+
465
+@pytest.mark.parametrize(
466
+ "qo_len, num_heads, kv_cache_dtype, expect_persistent",
467
+ [
468
+ # A padded rank keeps the schedule wherever a persistent kernel exists:
469
+ # bf16 up to qseqlen 4, fp8 at every qlen (its fold requires it).
470
+ (1, 12, "auto", True),
471
+ (4, 12, "auto", True),
472
+ (1, 12, "fp8", True),
473
+ (8, 12, "fp8", True),
474
+ # Past qseqlen 4 bf16 has only the non-persistent entry, and the fold
475
+ # that would reach a persistent one is gfx950-only, so the padded rank
476
+ # must fall back rather than ask for a kernel that is not built.
477
+ (8, 12, "auto", False),
478
+ # A native 16-head rank is unchanged: it took the schedule before this
479
+ # backend keyed the gate off the routing, and still does.
480
+ (8, 16, "auto", True),
481
+ ],
482
+)
483
+def test_persistent_metadata_gate_without_gluon_build(
484
+ monkeypatch, qo_len, num_heads, kv_cache_dtype, expect_persistent
485
+):
486
+ """The gate on an arch with no Gluon build, i.e. gfx942.
487
+
488
+ Both Gluon predicates read False there, so a small-head shape reaches the
489
+ padded asm decode instead of being flattened. The schedule then has to
490
+ follow what the asm dispatch actually ships: bf16 has no gqa=16 persistent
491
+ kernel above qseqlen 4, and the q-row fold that reaches one on gfx950 is
492
+ guarded on that arch, so asking for the schedule there selects a kernel
493
+ that does not exist.
494
+ """
495
+ get_mla_metadata_v1 = mock.MagicMock()
496
+ monkeypatch.setitem(
497
+ sys.modules,
498
+ "aiter",
499
+ SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
500
+ )
501
+ monkeypatch.setattr(
502
+ rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
503
+ )
504
+ monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: False)
505
+ monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto")
506
+
507
+ num_reqs = 2
508
+ query_start_loc = torch.arange(
509
+ 0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32
510
+ )
511
+ metadata = AiterMLAMetadataBuilder._build_decode(
512
+ _builder(
513
+ mtp_decode_qlen=qo_len,
514
+ num_heads=num_heads,
515
+ kv_cache_dtype=kv_cache_dtype,
516
+ ),
517
+ block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
518
+ seq_lens_device=torch.tensor([8, 8], dtype=torch.int32),
519
+ max_seq_len=8,
520
+ query_start_loc_cpu=query_start_loc,
521
+ query_start_loc_device=query_start_loc,
522
+ num_decode_tokens=num_reqs * qo_len,
523
+ dcp_tot_seq_lens_device=None,
524
+ )
525
+
526
+ assert metadata.has_persistent_metadata is expect_persistent
527
+ assert get_mla_metadata_v1.called is expect_persistent
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+64
-35
@@ -26,7 +26,7 @@ from vllm.v1.attention.backend import (
26
CommonAttentionMetadata,
27
MultipleOf,
28
)
29
-from vllm.v1.kv_cache_interface import AttentionSpec
29
+from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
30
31
logger = init_logger(__name__)
32
@@ -269,17 +269,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
269
self.compilation_config = vllm_config.compilation_config
270
self.decode_attn_out_dtype = vllm_config.model_config.dtype
271
272
- # MTP/deepseek_mtp verification runs decode with qlen = num_spec + 1;
273
- # any other config (including no spec) stays at single-token decode.
274
- speculative_config = vllm_config.speculative_config
275
- if (
276
- speculative_config is not None
277
- and speculative_config.method in ("mtp", "deepseek_mtp")
278
- and speculative_config.num_speculative_tokens is not None
279
- ):
280
- self._mtp_decode_qlen = int(speculative_config.num_speculative_tokens) + 1
281
- else:
282
- self._mtp_decode_qlen = 1
272
+ # reorder_batch_threshold is the largest query length decode can be
273
+ # handed, and already accounts for the drafting scheme. A method-name
274
+ # whitelist sizes unlisted drafters for qlen=1, which closes the
275
+ # persistent gate below and makes aiter raise a KeyError mid-run.
276
+ self._mtp_decode_qlen = self.reorder_batch_threshold or 1
277
278
# Store the kernel block size from the spec. When kernel_block_size=1
279
# (no spec-dec), behavior is identical to the original. When > 1
@@ -327,6 +321,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
321
torch.float16: dtypes.fp16,
322
torch.bfloat16: dtypes.bf16,
323
}[kv_cache_spec.dtype]
324
+ # _build_decode needs the cache dtype to pick the decode kernel; keep
325
+ # the normalized string instead of dropping it at the end of __init__.
326
+ self._kv_cache_dtype_str = kv_cache_dtype_str
327
# MLAAttention quantizes decode Q to FP8 before calling this backend
328
# whenever the KV cache is FP8 and supports_quant_query_input is true.
329
q_dtype = (
@@ -624,7 +621,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
621
]
622
)
623
use_gluon_decode = AiterMLAHelper.use_gluon_decode(
627
- self.num_heads, int(max_qo_len)
624
+ self.num_heads, int(max_qo_len), self._kv_cache_dtype_str
625
)
626
627
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
@@ -693,13 +690,26 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
690
else:
691
qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]
692
696
- # Pass persistent metadata for every uniform decode we sized buffers for
697
- # (normal qlen==1 through MTP verification qlen==K): the fp8 nhead=32 fold
698
- # path breaks without it. qlen>K falls back to kernel-internal metadata.
699
- # Small-head (<16) decode takes the Gluon paths and never consumes it.
693
+ # Only the asm decode consumes the schedule, so gate on the routing
694
+ # rather than on num_heads >= 16, which denies it to a padded rank
695
+ # running the same asm kernels. The two predicates are disjoint --
696
+ # decode is qlen==1, verify is qlen>1 -- and cover both Gluon entries.
697
has_persistent_metadata = False
698
use_persistent_metadata = (
702
- self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
699
+ not AiterMLAHelper.use_gluon_decode(
700
+ self.num_heads, max_qo_len, self._kv_cache_dtype_str
701
+ )
702
+ and not AiterMLAHelper.use_gluon_verify(
703
+ self.num_heads, max_qo_len, self._kv_cache_dtype_str
704
+ )
705
+ # A padded rank has no bf16 persistent kernel past qlen 4 where the
706
+ # gfx950 fold is absent; the non-persistent entry covers it. fp8
707
+ # keeps the schedule -- its fold rejects non-persistent outright.
708
+ and (
709
+ self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
710
+ or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN
711
+ or is_quantized_kv_cache(self._kv_cache_dtype_str)
712
+ )
713
and max_qo_len >= 1
714
and max_qo_len <= self._mtp_decode_qlen
715
)
@@ -833,6 +843,10 @@ class AiterMLAHelper:
843
"""
844
845
_AITER_MIN_MLA_HEADS: Final = 16
846
+ # Largest qlen the padded gqa=16 asm decode has a bf16 persistent kernel
847
+ # for. Above it only the non-persistent qseqlen=8 entry exists, and the
848
+ # fold that reaches a persistent one is gfx950-only.
849
+ _ASM_PADDED_MAX_PS_QLEN: Final = 4
850
_AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = ()
851
852
@staticmethod
@@ -888,17 +902,19 @@ class AiterMLAHelper:
902
return o[:, :num_heads, :]
903
904
@staticmethod
891
- def use_gluon_decode(num_heads: int, max_qo_len: int) -> bool:
892
- # Small-head (<16) single-token decode can use either the Gluon kernel
893
- # or the padded asm persistent decode, selected by
894
- # VLLM_ROCM_AITER_MLA_ASM_PADDING (see _aiter_mla_small_head_mode) and
895
- # the arch: Gluon only has a gfx950 build. In "auto" (default) mode
896
- # divisor counts keep Gluon on gfx950 and everything else -- non-divisor
897
- # counts (e.g. 12 heads/rank at TP8) and all counts on gfx942 -- takes
898
- # the asm path, which get_mla_padded_q pads to exactly 16.
905
+ def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
906
+ # Small-head (<16) single-token decode takes either the Gluon kernel or
907
+ # the padded asm persistent decode, selected by
908
+ # VLLM_ROCM_AITER_MLA_ASM_PADDING and the arch (Gluon is gfx950 only).
909
m = AiterMLAHelper._AITER_MIN_MLA_HEADS
910
if num_heads >= m or max_qo_len != 1:
911
return False
912
+ # Gluon's only fp8-KV regime, bh16bn128, is a bf16-query kernel with a
913
+ # hardcoded scale that asserts batch_size == 1, so it cannot serve a
914
+ # decode batch. Checked before the mode knob: an explicit "gluon"
915
+ # request under fp8 would assert immediately.
916
+ if is_quantized_kv_cache(kv_cache_dtype):
917
+ return False
918
mode = _aiter_mla_small_head_mode()
919
if mode == "asm":
920
return False
@@ -907,6 +923,23 @@ class AiterMLAHelper:
923
return gluon_supported
924
return m % num_heads == 0 and gluon_supported
925
926
+ @staticmethod
927
+ def use_gluon_verify(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
928
+ """Whether a small-head multi-token verify is flattened onto Gluon.
929
+
930
+ bf16 has no gqa<16, qseqlen>1 asm kernel, so the verify is flattened
931
+ into per-token Gluon decodes. fp8 has one via the q-row fold and must
932
+ not come here: the flatten hands Gluon the batch size its fp8 regime
933
+ asserts against. A predicate rather than inline in forward_mqa so the
934
+ builder sees the same answer the impl acts on.
935
+ """
936
+ if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= 1:
937
+ return False
938
+ if is_quantized_kv_cache(kv_cache_dtype):
939
+ return False
940
+ # Same arch and mode gating as use_gluon_decode.
941
+ return _aiter_mla_small_head_mode() != "asm" and _gluon_mla_decode_supported()
942
+
943
944
class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
945
def __init__(
@@ -1191,15 +1224,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1224
# target is checking draft tokens, so position t must not see t+1 --
1225
# and attention rows are independent, so giving row t the KV range
1226
# [0, context + t] is exactly causal multi-token attention.
1194
- # Gluon only has a gfx950 build, so gate on the arch (mirrors
1195
- # use_gluon_decode). VLLM_ROCM_AITER_MLA_ASM_PADDING=asm also forces the
1196
- # asm path here. Otherwise this falls through to the asm decode, which
1197
- # pads to 16 heads and handles qlen>1 verify directly.
1198
- if (
1199
- self.num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS
1200
- and int(decode.max_qo_len) > 1
1201
- and _aiter_mla_small_head_mode() != "asm"
1202
- and _gluon_mla_decode_supported()
1227
+ # Arch, mode and dtype gating all live in use_gluon_verify, so that the
1228
+ # builder -- which has to know whether the asm decode will run -- sees
1229
+ # the same answer as this branch.
1230
+ if AiterMLAHelper.use_gluon_verify(
1231
+ self.num_heads, int(decode.max_qo_len), self.kv_cache_dtype
1232
):
1233
qlen = int(decode.max_qo_len)
1234
if type(q) is tuple: