[ROCm][MLA] Use asm decode for non-divisor small head counts (#50578)
Signed-off-by: vanshbhatia-amd <210711135+vanshbhatia-amd@users.noreply.github.com> Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Signed-off-by: Liuyinfeng01 <yinfeliu@amd.com> Co-authored-by: vanshbhatia-amd <210711135+vanshbhatia-amd@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Liuyinfeng01 <yinfeliu@amd.com>
vanshbhatia-amd committed
Aug 6, 2026 at 14:56 UTC
d8eabdbfbe93ecc8a8d5cb8a55c5067a443a8796
4 files changed
+393
-25
tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py
+12
-8
@@ -32,11 +32,9 @@ def _on_rocm_with_aiter() -> bool:
32
return current_platform.is_rocm() and is_aiter_found()
33
34
35
-# Unlike the fp8 persistent-decode metadata this file's neighbour guards, the
36
-# verify flatten is not gfx950-only: it is selected by head count alone and is
37
-# reached on any ROCm AITER MLA deployment with fewer than 16 heads per rank.
38
-# The Gluon kernel is replaced by a spy here, so nothing below needs it to be
39
-# runnable either.
35
+# The production verify flatten is gfx950-only. This test replaces the Gluon
36
+# kernel with a spy and forces the architecture feature probe so the metadata
37
+# regression remains testable on every ROCm AITER CI runner.
38
pytestmark = pytest.mark.skipif(
39
not _on_rocm_with_aiter(),
40
reason="ROCM_AITER_MLA verify flatten requires ROCm and AITER",
@@ -186,9 +184,15 @@ def _run_verify_block():
184
185
# The Gluon kernel only reads the metadata this test is about, so a spy in
186
# its place keeps the assertions independent of the AITER build.
189
- with patch(
190
- "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon",
191
- lambda: spy,
187
+ with (
188
+ patch(
189
+ "vllm.v1.attention.backends.mla.rocm_aiter_mla._gluon_mla_decode_supported",
190
+ lambda: True,
191
+ ),
192
+ patch(
193
+ "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon",
194
+ lambda: spy,
195
+ ),
196
):
197
impl.forward_mqa((q_nope, q_pe), kv_cache, metadata, layer=None)
198
tests/kernels/attention/test_rocm_aiter_mla_head_padding.py
new
+265
@@ -0,0 +1,265 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Head-count padding + kernel selection for the ROCm AITER MLA backend.
4
+
5
+Kimi-K3 at TP8 puts 12 heads/rank and at TP16 puts 6 heads/rank on the AITER
6
+MLA decode. The asm persistent decode requires exactly 16 heads, so small
7
+head counts are tile-padded to 16 and the padding heads are sliced back off
8
+the output. Divisor counts (1/2/4/8) may keep the Gluon kernel, but only on
9
+gfx950 where that kernel has a build; every small head count on gfx942 (which
10
+has no Gluon build) is routed to the asm persistent decode instead.
11
+"""
12
+
13
+import math
14
+import types
15
+
16
+import pytest
17
+import torch
18
+import torch.nn.functional as F
19
+
20
+from vllm._aiter_ops import is_aiter_found
21
+from vllm.platforms import current_platform
22
+from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
23
+ AiterMLADecodeMetadata,
24
+ AiterMLAHelper,
25
+ AiterMLAImpl,
26
+ AiterMLAMetadata,
27
+)
28
+
29
+NUM_HEADS = 12
30
+KV_LORA_RANK = 512
31
+QK_ROPE_HEAD_DIM = 64
32
+QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM
33
+CONTEXT_LEN = 4096
34
+SCALE = 1.0 / math.sqrt(QK_HEAD_DIM)
35
+
36
+# Non-divisor counts must go through the tile-and-slice path; divisor counts
37
+# (of 16) keep repeat_interleave. Both must pad to exactly 16 and round-trip.
38
+NON_DIVISOR_HEADS = [3, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15]
39
+DIVISOR_HEADS = [1, 2, 4, 8]
40
+
41
+
42
+def _rocm_aiter_available() -> bool:
43
+ return current_platform.is_rocm() and is_aiter_found() and torch.cuda.is_available()
44
+
45
+
46
+def _on_gfx950() -> bool:
47
+ if not (current_platform.is_rocm() and is_aiter_found()):
48
+ return False
49
+ from vllm.platforms.rocm import on_gfx950
50
+
51
+ return on_gfx950()
52
+
53
+
54
+def _expected_tile_pad(q: torch.Tensor, num_heads: int, m: int = 16) -> torch.Tensor:
55
+ # get_mla_padded_q tiles the heads and slices to m, i.e. head i of the
56
+ # padded tensor is head (i % num_heads) of the input.
57
+ idx = [i % num_heads for i in range(m)]
58
+ return q[:, idx, :]
59
+
60
+
61
+def _make_h12_decode_metadata(device: torch.device) -> AiterMLAMetadata:
62
+ page_indices = torch.arange(CONTEXT_LEN, dtype=torch.int32, device=device)
63
+ decode = AiterMLADecodeMetadata(
64
+ block_table=page_indices.view(1, -1),
65
+ seq_lens=torch.tensor([CONTEXT_LEN], dtype=torch.int32, device=device),
66
+ dcp_tot_seq_lens=None,
67
+ paged_kv_indptr=torch.tensor(
68
+ [0, CONTEXT_LEN], dtype=torch.int32, device=device
69
+ ),
70
+ paged_kv_indices=page_indices,
71
+ paged_kv_last_page_len=torch.ones(1, dtype=torch.int32, device=device),
72
+ qo_indptr=torch.tensor([0, 1], dtype=torch.int32, device=device),
73
+ attn_out_dtype=torch.bfloat16,
74
+ max_qo_len=1,
75
+ )
76
+ return AiterMLAMetadata(
77
+ num_reqs=1,
78
+ max_query_len=1,
79
+ max_seq_len=CONTEXT_LEN,
80
+ num_actual_tokens=1,
81
+ query_start_loc=decode.qo_indptr,
82
+ slot_mapping=torch.tensor([CONTEXT_LEN - 1], dtype=torch.int64, device=device),
83
+ num_decodes=1,
84
+ num_decode_tokens=1,
85
+ num_prefills=0,
86
+ decode=decode,
87
+ )
88
+
89
+
90
+def test_h12_query_is_tile_padded_to_h16():
91
+ q = torch.arange(2 * 12 * 4, dtype=torch.bfloat16).view(2, 12, 4)
92
+
93
+ padded_q = AiterMLAHelper.get_mla_padded_q(12, q)
94
+
95
+ assert padded_q.shape == (2, 16, 4)
96
+ assert padded_q.is_contiguous()
97
+ # The real heads are untouched...
98
+ torch.testing.assert_close(padded_q[:, :12], q)
99
+ # ...and the 4 padding heads are the tiled wrap-around (heads 0..3), not
100
+ # zeros: MLA attention is per-head independent so duplicate query heads are
101
+ # harmless and get sliced back off the output.
102
+ torch.testing.assert_close(padded_q[:, 12:], q[:, :4])
103
+
104
+
105
+def test_h6_tp16_query_is_padded_to_h16():
106
+ # TP16 puts 6 heads/rank. The old append-only padding produced cat(6, 6) =
107
+ # 12 heads and broke the asm kernel; tile-and-slice must reach exactly 16.
108
+ q = torch.arange(2 * 6 * 4, dtype=torch.bfloat16).view(2, 6, 4)
109
+
110
+ padded_q = AiterMLAHelper.get_mla_padded_q(6, q)
111
+
112
+ assert padded_q.shape == (2, 16, 4)
113
+ assert padded_q.is_contiguous()
114
+ torch.testing.assert_close(padded_q, _expected_tile_pad(q, 6))
115
+
116
+
117
+def test_h12_output_discards_padding_heads():
118
+ o = torch.arange(2 * 16 * 4, dtype=torch.bfloat16).view(2, 16, 4)
119
+
120
+ unpadded_o = AiterMLAHelper.get_mla_unpadded_o(12, o)
121
+
122
+ assert unpadded_o.shape == (2, 12, 4)
123
+ torch.testing.assert_close(unpadded_o, o[:, :12])
124
+
125
+
126
+def test_existing_divisor_head_mapping_is_unchanged():
127
+ q = torch.arange(2 * 8 * 4, dtype=torch.bfloat16).view(2, 8, 4)
128
+
129
+ padded_q = AiterMLAHelper.get_mla_padded_q(8, q)
130
+ unpadded_o = AiterMLAHelper.get_mla_unpadded_o(8, padded_q)
131
+
132
+ # Divisor counts still use repeat_interleave / strided unpad, unchanged.
133
+ torch.testing.assert_close(padded_q, q.repeat_interleave(2, dim=1))
134
+ torch.testing.assert_close(unpadded_o, q)
135
+
136
+
137
+@pytest.mark.parametrize("num_heads", NON_DIVISOR_HEADS + DIVISOR_HEADS)
138
+def test_all_small_head_counts_pad_to_16_and_round_trip(num_heads: int):
139
+ q = torch.arange(2 * num_heads * 4, dtype=torch.float32).view(2, num_heads, 4)
140
+
141
+ padded_q = AiterMLAHelper.get_mla_padded_q(num_heads, q)
142
+ unpadded_o = AiterMLAHelper.get_mla_unpadded_o(num_heads, padded_q)
143
+
144
+ assert padded_q.shape == (2, 16, 4)
145
+ assert padded_q.is_contiguous()
146
+ # The real heads survive the pad -> unpad round trip exactly.
147
+ torch.testing.assert_close(unpadded_o, q)
148
+
149
+
150
+def test_num_heads_ge_16_is_passthrough():
151
+ q = torch.arange(2 * 16 * 4, dtype=torch.float32).view(2, 16, 4)
152
+ assert AiterMLAHelper.get_mla_padded_q(16, q) is q
153
+ assert AiterMLAHelper.get_mla_unpadded_o(16, q) is q
154
+
155
+
156
+def test_is_valid_num_heads():
157
+ for n in range(1, 16):
158
+ assert AiterMLAHelper.is_valid_num_heads(n)
159
+ assert AiterMLAHelper.is_valid_num_heads(16)
160
+ assert AiterMLAHelper.is_valid_num_heads(32)
161
+ assert not AiterMLAHelper.is_valid_num_heads(0)
162
+
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)
168
+ # >=16 heads never pad, never Gluon.
169
+ assert not AiterMLAHelper.use_gluon_decode(16, 1)
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)
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
182
+
183
+
184
+def test_asm_padding_env_default_is_auto(monkeypatch):
185
+ monkeypatch.delenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", raising=False)
186
+ import vllm.envs as envs
187
+
188
+ assert envs.VLLM_ROCM_AITER_MLA_ASM_PADDING == "auto"
189
+
190
+
191
+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)
196
+
197
+
198
+def test_asm_padding_env_force_gluon_follows_arch(monkeypatch):
199
+ monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "gluon")
200
+ on_gfx950 = _on_gfx950()
201
+ # Forcing Gluon: any 1..15 single-token decode uses it where a build exists
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
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)
214
+
215
+
216
+@pytest.mark.skipif(
217
+ not _rocm_aiter_available(),
218
+ reason="12-head AITER MLA asm persistent decode needs ROCm + AITER",
219
+)
220
+@torch.inference_mode()
221
+def test_h12_aiter_mla_decode_matches_reference():
222
+ """The 12-head AITER decode output must match an attention reference.
223
+
224
+ 12 is a non-divisor of 16 so use_gluon_decode is False on every arch; the
225
+ query is tile-padded to 16 heads and served by the asm persistent decode
226
+ (the CDNA3 path this PR enables, and unchanged on CDNA4).
227
+ """
228
+ device = torch.device("cuda:0")
229
+ torch.manual_seed(42)
230
+
231
+ q = torch.randn(1, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16, device=device)
232
+ kv_cache = torch.randn(
233
+ CONTEXT_LEN, 1, QK_HEAD_DIM, dtype=torch.bfloat16, device=device
234
+ )
235
+
236
+ impl = object.__new__(AiterMLAImpl)
237
+ impl.num_heads = NUM_HEADS
238
+ impl.kv_lora_rank = KV_LORA_RANK
239
+ impl.scale = SCALE
240
+
241
+ one = torch.ones(1, dtype=torch.float32, device=device)
242
+ layer = types.SimpleNamespace(_q_scale=one, _k_scale=one)
243
+ metadata = _make_h12_decode_metadata(device)
244
+
245
+ out, _ = impl.forward_mqa(q, kv_cache, metadata, layer)
246
+
247
+ assert out.shape[1] == NUM_HEADS # padding heads sliced back off
248
+
249
+ key = kv_cache[:, 0].float().unsqueeze(0)
250
+ value = key[..., :KV_LORA_RANK]
251
+ out_ref = F.scaled_dot_product_attention(
252
+ q[0].float().unsqueeze(1),
253
+ key,
254
+ value,
255
+ scale=SCALE,
256
+ enable_gqa=True,
257
+ ).squeeze(1)
258
+
259
+ assert torch.isfinite(out).all()
260
+ torch.testing.assert_close(
261
+ out.float(),
262
+ out_ref.unsqueeze(0),
263
+ atol=1e-2,
264
+ rtol=1e-2,
265
+ )
vllm/envs.py
+15
@@ -133,6 +133,7 @@ if TYPE_CHECKING:
133
VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False
134
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
135
VLLM_ROCM_USE_AITER_MLA: bool = True
136
+ VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto"
137
VLLM_ROCM_USE_AITER_MHA: bool = True
138
VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False
139
VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False
@@ -1237,6 +1238,20 @@ environment_variables: dict[str, Callable[[], Any]] = {
1238
"VLLM_ROCM_USE_AITER_MLA": lambda: (
1239
os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1")
1240
),
1241
+ # Small-head (<16) AITER MLA decode kernel selection. Small head counts
1242
+ # (e.g. Kimi-K3: 12 heads/rank at TP8, 6 at TP16) can decode either through
1243
+ # the Gluon small-head kernel or through the padded persistent-scheduling
1244
+ # (PS) ASM kernel. "auto" (default) keeps Gluon for head counts that divide
1245
+ # 16 where a Gluon build exists (gfx950/CDNA4) and otherwise uses the padded
1246
+ # PS ASM decode; "gluon" forces the Gluon path wherever a build exists;
1247
+ # "asm" forces the padded PS ASM decode. On gfx942/CDNA3 there is no Gluon
1248
+ # build, so the ASM path is always used regardless of this setting.
1249
+ "VLLM_ROCM_AITER_MLA_ASM_PADDING": env_with_choices(
1250
+ "VLLM_ROCM_AITER_MLA_ASM_PADDING",
1251
+ "auto",
1252
+ ["auto", "gluon", "asm"],
1253
+ case_sensitive=False,
1254
+ ),
1255
# Whether to use aiter mha ops.
1256
# By default is enabled.
1257
"VLLM_ROCM_USE_AITER_MHA": lambda: (
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+101
-17
@@ -77,6 +77,50 @@ def _fp8_mla_prefill_supported() -> bool:
77
return True
78
79
80
+@functools.lru_cache(maxsize=1)
81
+def _gluon_mla_decode_supported() -> bool:
82
+ """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build.
83
+
84
+ Its tiling needs ~160 KiB of LDS, which exceeds CDNA3's 64 KiB, so on
85
+ gfx942 there is no kernel to fall through to and selecting it asserts
86
+ (``mla_gluon requires gfx950``). Restrict Gluon decode to gfx950; other
87
+ archs use the asm persistent decode, which ``get_mla_padded_q`` makes
88
+ correct for any 1..15 heads.
89
+ """
90
+ try:
91
+ from vllm.platforms.rocm import on_gfx950
92
+ except Exception: # noqa: BLE001
93
+ return False
94
+ return on_gfx950()
95
+
96
+
97
+def _aiter_mla_small_head_mode() -> str:
98
+ """Small-head (<16) MLA decode kernel selection.
99
+
100
+ Controlled by ``VLLM_ROCM_AITER_MLA_ASM_PADDING``:
101
+
102
+ - ``"auto"`` (default): let the arch decide -- divisor head counts keep the
103
+ Gluon decode where a build exists (gfx950), everything else (non-divisor
104
+ counts and all counts on gfx942) uses the padded persistent-scheduling
105
+ ASM decode.
106
+ - ``"gluon"``: prefer the Gluon path wherever a build exists.
107
+ - ``"asm"``: force the padded persistent-scheduling ASM decode.
108
+
109
+ On gfx942 (no Gluon build) the ASM path is always used regardless of this
110
+ setting; ``"gluon"`` there falls back to ASM with a one-time warning.
111
+ """
112
+ import vllm.envs as envs
113
+
114
+ mode = (envs.VLLM_ROCM_AITER_MLA_ASM_PADDING or "auto").lower()
115
+ if mode == "gluon" and not _gluon_mla_decode_supported():
116
+ logger.warning_once(
117
+ "VLLM_ROCM_AITER_MLA_ASM_PADDING=gluon requested, but this device "
118
+ "has no Gluon MLA decode build (Gluon requires gfx950); using the "
119
+ "padded persistent-scheduling ASM decode instead."
120
+ )
121
+ return mode
122
+
123
+
124
class AiterMLABackend(MLACommonBackend):
125
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
126
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
@@ -781,8 +825,11 @@ def _expand_page_indices_kernel(
825
826
class AiterMLAHelper:
827
"""
784
- AITER MLA implementation requires num_heads >= 16. If num_heads < 16 and
785
- 16 % num_heads == 0, we can pad q to 16 heads; otherwise AITER has to fail.
828
+ AITER MLA persistent (asm) decode requires num_heads >= 16. Head counts
829
+ < 16 are padded up to exactly 16: divisors of 16 by repeat_interleave,
830
+ other counts (e.g. 12 heads/rank at TP8, 6 at TP16) by tiling the query
831
+ heads and slicing to 16. Non-divisor padded decodes take the asm path;
832
+ divisors and max_qo_len > 1 small-head verify still use Gluon.
833
"""
834
835
_AITER_MIN_MLA_HEADS: Final = 16
@@ -791,8 +838,9 @@ class AiterMLAHelper:
838
@staticmethod
839
def check_num_heads_validity(num_heads: int):
840
assert AiterMLAHelper.is_valid_num_heads(num_heads), (
794
- "ROCM AITER MLA requires 1-15 heads for Gluon decode or a multiple "
795
- f"of 16 heads for persistent decode, but got {num_heads}.\n"
841
+ "ROCM AITER MLA requires 1-15 heads (padded to 16 for asm "
842
+ "persistent decode; exact divisors of 16 may keep Gluon) or a "
843
+ f"multiple of 16 heads, but got {num_heads}.\n"
844
f"Try adjusting tensor_parallel_size value."
845
)
846
@@ -809,25 +857,55 @@ class AiterMLAHelper:
857
858
@staticmethod
859
def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:
812
- return (
813
- q
814
- if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
815
- else q.repeat_interleave(
816
- AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, dim=1
817
- )
818
- )
860
+ m = AiterMLAHelper._AITER_MIN_MLA_HEADS
861
+ if num_heads >= m:
862
+ return q
863
+ if m % num_heads == 0:
864
+ return q.repeat_interleave(m // num_heads, dim=1)
865
+ # Non-divisor head counts (e.g. 12 heads/rank at TP8, 6 at TP16) cannot
866
+ # be padded by repeat_interleave. Tile the query heads and slice to
867
+ # exactly m; this reaches m for any 0 < num_heads < m (unlike a single
868
+ # append, which under-pads when num_heads < m - num_heads). MLA
869
+ # attention is independent per query head over the shared KV, so the
870
+ # padding heads cannot affect heads [0:num_heads]; they are sliced back
871
+ # off in get_mla_unpadded_o.
872
+ reps = -(-m // num_heads) # ceil(m / num_heads)
873
+ # Slicing a tiled tensor down to m yields a non-contiguous view whenever
874
+ # reps * num_heads > m (the common case: TP8 12->24->16, TP16 6->18->16).
875
+ # The asm persistent decode reads q as a packed [tokens, m, head_dim]
876
+ # buffer, so materialize a contiguous copy. No-op when already contiguous.
877
+ return q.repeat(1, reps, 1)[:, :m, :].contiguous()
878
879
@staticmethod
880
def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor:
822
- return (
823
- o
824
- if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
825
- else o[:, :: AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, :]
826
- )
881
+ m = AiterMLAHelper._AITER_MIN_MLA_HEADS
882
+ if num_heads >= m:
883
+ return o
884
+ if m % num_heads == 0:
885
+ return o[:, :: m // num_heads, :]
886
+ # Undo the tile-padding from get_mla_padded_q: the real heads are the
887
+ # first num_heads.
888
+ return o[:, :num_heads, :]
889
890
@staticmethod
891
def use_gluon_decode(num_heads: int, max_qo_len: int) -> bool:
830
- return num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS and max_qo_len == 1
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.
899
+ m = AiterMLAHelper._AITER_MIN_MLA_HEADS
900
+ if num_heads >= m or max_qo_len != 1:
901
+ return False
902
+ mode = _aiter_mla_small_head_mode()
903
+ if mode == "asm":
904
+ return False
905
+ gluon_supported = _gluon_mla_decode_supported()
906
+ if mode == "gluon":
907
+ return gluon_supported
908
+ return m % num_heads == 0 and gluon_supported
909
910
911
class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
@@ -1113,9 +1191,15 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1191
# target is checking draft tokens, so position t must not see t+1 --
1192
# and attention rows are independent, so giving row t the KV range
1193
# [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()
1203
):
1204
qlen = int(decode.max_qo_len)
1205
if type(q) is tuple: