[ROCm] Pad non-aligned AITER MLA heads (#51647)
Signed-off-by: Liuyinfeng01 <yinfeliu@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
yinfengLiu committed
Aug 19, 2026 at 05:37 UTC
5d8a4cf9761ec890ae9f11502fe0f4a13e9981a5
4 files changed
+200
-58
tests/kernels/attention/test_rocm_aiter_mla_head_padding.py
+107
-18
@@ -1,13 +1,11 @@
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.
3
+"""Head padding and kernel selection for the ROCm AITER MLA backend.
4
+
5
+The asm persistent decode requires a 16-aligned head count, so unaligned
6
+counts through 128 are tile-padded to the next multiple of 16 and sliced back
7
+off the output. Small divisor counts (1/2/4/8) preserve their existing
8
+repeat-interleave path and may keep the Gluon kernel on gfx950.
9
"""
10
11
import math
@@ -19,6 +17,7 @@ import torch.nn.functional as F
17
18
from vllm._aiter_ops import is_aiter_found
19
from vllm.platforms import current_platform
20
+from vllm.v1.attention.backends.mla import rocm_aiter_mla
21
from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
22
AiterMLADecodeMetadata,
23
AiterMLAHelper,
@@ -33,12 +32,20 @@ QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM
32
CONTEXT_LEN = 4096
33
SCALE = 1.0 / math.sqrt(QK_HEAD_DIM)
34
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.
35
+# Small non-divisor counts use tile-and-slice; divisors of 16 keep
36
+# repeat_interleave. Both pad to exactly 16 and round-trip.
37
NON_DIVISOR_HEADS = [3, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15]
38
DIVISOR_HEADS = [1, 2, 4, 8]
39
40
41
+@pytest.fixture(autouse=True)
42
+def _disable_native_h24(monkeypatch):
43
+ """Exercise the padding fallback unless a test explicitly enables H24."""
44
+ monkeypatch.setattr(
45
+ rocm_aiter_mla, "_aiter_mla_native_h24_supported", lambda: False
46
+ )
47
+
48
+
49
def _rocm_aiter_available() -> bool:
50
return current_platform.is_rocm() and is_aiter_found() and torch.cuda.is_available()
51
@@ -123,6 +130,72 @@ def test_h12_output_discards_padding_heads():
130
torch.testing.assert_close(unpadded_o, o[:, :12])
131
132
133
+def test_h24_query_is_tile_padded_to_h32():
134
+ q = torch.arange(2 * 24 * 4, dtype=torch.float32).view(2, 24, 4)
135
+
136
+ padded_q = AiterMLAHelper.get_mla_padded_q(24, q)
137
+
138
+ assert padded_q.shape == (2, 32, 4)
139
+ assert padded_q.is_contiguous()
140
+ torch.testing.assert_close(padded_q[:, :24], q)
141
+ torch.testing.assert_close(padded_q[:, 24:], q[:, :8])
142
+
143
+
144
+def test_h24_output_discards_h32_padding_heads():
145
+ o = torch.arange(2 * 32 * 4, dtype=torch.float32).view(2, 32, 4)
146
+
147
+ unpadded_o = AiterMLAHelper.get_mla_unpadded_o(24, o)
148
+
149
+ assert unpadded_o.shape == (2, 24, 4)
150
+ torch.testing.assert_close(unpadded_o, o[:, :24])
151
+
152
+
153
+def test_h24_reducer_without_metadata_still_pads_to_h32(monkeypatch):
154
+ monkeypatch.setattr(
155
+ rocm_aiter_mla, "_aiter_mla_native_h24_reducer_supported", lambda: True
156
+ )
157
+ monkeypatch.setattr(
158
+ rocm_aiter_mla, "_aiter_mla_native_h24_metadata_supported", lambda: False
159
+ )
160
+ monkeypatch.setattr(
161
+ rocm_aiter_mla,
162
+ "_aiter_mla_native_h24_supported",
163
+ lambda: (
164
+ rocm_aiter_mla._aiter_mla_native_h24_reducer_supported()
165
+ and rocm_aiter_mla._aiter_mla_native_h24_metadata_supported()
166
+ ),
167
+ )
168
+ q = torch.arange(2 * 24 * 4, dtype=torch.float32).view(2, 24, 4)
169
+
170
+ assert AiterMLAHelper.get_actual_mla_num_heads(24) == 32
171
+ padded_q = AiterMLAHelper.get_mla_padded_q(24, q)
172
+ assert padded_q.shape == (2, 32, 4)
173
+ torch.testing.assert_close(padded_q[:, :24], q)
174
+ torch.testing.assert_close(padded_q[:, 24:], q[:, :8])
175
+
176
+
177
+def test_native_h24_requires_reducer_and_metadata(monkeypatch):
178
+ monkeypatch.setattr(
179
+ rocm_aiter_mla, "_aiter_mla_native_h24_reducer_supported", lambda: True
180
+ )
181
+ monkeypatch.setattr(
182
+ rocm_aiter_mla, "_aiter_mla_native_h24_metadata_supported", lambda: True
183
+ )
184
+ monkeypatch.setattr(
185
+ rocm_aiter_mla,
186
+ "_aiter_mla_native_h24_supported",
187
+ lambda: (
188
+ rocm_aiter_mla._aiter_mla_native_h24_reducer_supported()
189
+ and rocm_aiter_mla._aiter_mla_native_h24_metadata_supported()
190
+ ),
191
+ )
192
+ q = torch.arange(2 * 24 * 4, dtype=torch.float32).view(2, 24, 4)
193
+
194
+ assert AiterMLAHelper.get_actual_mla_num_heads(24) == 24
195
+ assert AiterMLAHelper.get_mla_padded_q(24, q) is q
196
+ assert AiterMLAHelper.get_mla_unpadded_o(24, q) is q
197
+
198
+
199
def test_existing_divisor_head_mapping_is_unchanged():
200
q = torch.arange(2 * 8 * 4, dtype=torch.bfloat16).view(2, 8, 4)
201
@@ -134,6 +207,18 @@ def test_existing_divisor_head_mapping_is_unchanged():
207
torch.testing.assert_close(unpadded_o, q)
208
209
210
+@pytest.mark.parametrize("num_heads", [17, 24, 31])
211
+def test_unaligned_head_counts_round_trip_through_h32(num_heads: int):
212
+ q = torch.arange(2 * num_heads * 4, dtype=torch.float32).view(2, num_heads, 4)
213
+
214
+ padded_q = AiterMLAHelper.get_mla_padded_q(num_heads, q)
215
+ unpadded_o = AiterMLAHelper.get_mla_unpadded_o(num_heads, padded_q)
216
+
217
+ assert padded_q.shape == (2, 32, 4)
218
+ assert padded_q.is_contiguous()
219
+ torch.testing.assert_close(unpadded_o, q)
220
+
221
+
222
@pytest.mark.parametrize("num_heads", NON_DIVISOR_HEADS + DIVISOR_HEADS)
223
def test_all_small_head_counts_pad_to_16_and_round_trip(num_heads: int):
224
q = torch.arange(2 * num_heads * 4, dtype=torch.float32).view(2, num_heads, 4)
@@ -147,26 +232,30 @@ def test_all_small_head_counts_pad_to_16_and_round_trip(num_heads: int):
232
torch.testing.assert_close(unpadded_o, q)
233
234
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
235
+def test_aligned_h32_is_zero_copy():
236
+ q = torch.arange(2 * 32 * 4, dtype=torch.float32).view(2, 32, 4)
237
+ assert AiterMLAHelper.get_mla_padded_q(32, q) is q
238
+ assert AiterMLAHelper.get_mla_unpadded_o(32, q) is q
239
240
241
def test_is_valid_num_heads():
157
- for n in range(1, 16):
242
+ for n in range(1, 129):
243
assert AiterMLAHelper.is_valid_num_heads(n)
159
- assert AiterMLAHelper.is_valid_num_heads(16)
160
- assert AiterMLAHelper.is_valid_num_heads(32)
244
+ assert AiterMLAHelper.is_valid_num_heads(24)
245
+ assert AiterMLAHelper.is_valid_num_heads(127)
246
+ # Aligned counts remain valid above the range where padding is supported.
247
+ assert AiterMLAHelper.is_valid_num_heads(144)
248
assert not AiterMLAHelper.is_valid_num_heads(0)
249
+ assert not AiterMLAHelper.is_valid_num_heads(129)
250
251
252
def test_nondivisor_and_multitoken_never_use_gluon():
253
# Non-divisor decode always takes the asm path (12 heads/rank at TP8).
254
assert not AiterMLAHelper.use_gluon_decode(12, 1, "auto")
255
assert not AiterMLAHelper.use_gluon_decode(6, 1, "auto")
168
- # >=16 heads never pad, never Gluon.
256
+ # >=16 heads never use Gluon, including unaligned counts padded for asm.
257
assert not AiterMLAHelper.use_gluon_decode(16, 1, "auto")
258
+ assert not AiterMLAHelper.use_gluon_decode(24, 1, "auto")
259
# Multi-token (verify / qlen>1) is never the single-token Gluon decode.
260
assert not AiterMLAHelper.use_gluon_decode(8, 4, "auto")
261
assert not AiterMLAHelper.use_gluon_decode(12, 4, "auto")
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+9
-6
@@ -15,6 +15,7 @@ if not current_platform.is_rocm():
15
16
from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402
17
from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( # noqa: E402
18
+ AiterMLAHelper,
19
AiterMLAMetadataBuilder,
20
)
21
@@ -83,7 +84,7 @@ def _builder(
84
_uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len),
85
_use_persistent_metadata=False,
86
kernel_block_size=kernel_block_size,
86
- _num_attention_heads=16,
87
+ _num_attention_heads=AiterMLAHelper.get_actual_mla_num_heads(num_heads),
88
_mla_work_meta_data=torch.empty(1, dtype=torch.int32),
89
_mla_work_info_set=torch.empty(1, dtype=torch.int32),
90
_mla_work_indptr=torch.empty(1, dtype=torch.int32),
@@ -109,7 +110,7 @@ def test_backend_declares_uniform_batch_support():
110
)
111
112
112
-@pytest.mark.parametrize("num_heads", [8, 16, 32, 64, 128])
113
+@pytest.mark.parametrize("num_heads", [8, 16, 24, 32, 64, 128])
114
@pytest.mark.parametrize(
115
"spec_method, parallel_drafting",
116
[
@@ -127,8 +128,8 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(
128
):
129
"""Aiter init sizes the metadata for every query length decode can be handed.
130
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.
131
+ Sweeping num_heads asserts metadata is sized for the padded decode shape,
132
+ covering Kimi-K3 TP4's 24 -> 32 head path and native fp8 nhead=32 folding.
133
"""
134
135
dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16")
@@ -207,7 +208,7 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(
208
{
209
"max_batch_size": config.scheduler_config.max_num_seqs,
210
"max_qo_len": builder.reorder_batch_threshold,
210
- "num_attention_heads": max(16, num_heads),
211
+ "num_attention_heads": AiterMLAHelper.get_actual_mla_num_heads(num_heads),
212
"q_dtype": dtypes.fp8,
213
"kv_dtype": dtypes.fp8,
214
"is_sparse": False,
@@ -381,6 +382,7 @@ def test_decode_expands_kernel_block_page_indices(monkeypatch):
382
(1, 1, 16, "auto", True), # non-MTP decode
383
(4, 2, 16, "auto", True), # MTP deployment, in-range step
384
(4, 4, 16, "auto", True), # MTP deployment, full-qlen verification step
385
+ (1, 1, 24, "auto", True), # unaligned H24 pads to H32 persistent decode
386
(2, 4, 16, "auto", False), # step demand exceeds provisioned K -> fallback
387
(1, 1, 8, "auto", False), # divisor head count -> Gluon decode owns qlen==1
388
(4, 4, 8, "auto", False), # small head count -> Gluon flatten owns qlen>1
@@ -410,7 +412,8 @@ def test_persistent_metadata_gate(
412
K = _mtp_decode_qlen sizes the metadata buffers at init; a decode step gets
413
the pre-built schedule only when its qlen fits those buffers, otherwise it
414
falls back to the kernel computing its own. qlen==1 (non-MTP) must stay
413
- in-range -- dropping it is the regression this guards.
415
+ in-range -- dropping it is the regression this guards. This includes
416
+ unaligned H24, whose padded H32 decode uses persistent metadata.
417
418
Only the Gluon paths ignore the schedule, so the gate follows the routing
419
predicates rather than the raw head count. Reading `num_heads >= 16` instead
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+79
-31
@@ -3,6 +3,7 @@
3
4
import functools
5
from dataclasses import dataclass
6
+from pathlib import Path
7
from typing import ClassVar, Final
8
9
import torch
@@ -77,6 +78,48 @@ def _fp8_mla_prefill_supported() -> bool:
78
return True
79
80
81
+@functools.lru_cache(maxsize=1)
82
+def _aiter_mla_native_h24_reducer_supported() -> bool:
83
+ """Whether AITER's JIT reducer supports the native H24/512 shape."""
84
+ try:
85
+ from aiter.jit.core import AITER_CSRC_DIR
86
+
87
+ reduce_source = Path(AITER_CSRC_DIR) / "kernels" / "mla" / "reduce.cu"
88
+ source = "".join(reduce_source.read_text(encoding="utf-8").split())
89
+ except (ImportError, OSError):
90
+ return False
91
+ return "MLA_REDUCE_CASE_EF(NUM_HEAD,24,HEAD_DIM,512," in source
92
+
93
+
94
+@functools.lru_cache(maxsize=1)
95
+def _aiter_mla_native_h24_metadata_supported() -> bool:
96
+ """Whether AITER's fast MLA metadata planner accepts native H24.
97
+
98
+ The reducer and metadata planner have independent shape dispatch. Checking
99
+ only the reducer can route H24 into a planner that rejects it before the
100
+ attention kernel launches. Until AITER exposes a capability API, inspect
101
+ the shipped JIT source for an explicit native-H24 planner branch.
102
+ """
103
+ try:
104
+ from aiter.jit.core import AITER_CSRC_DIR
105
+
106
+ metadata_source = (
107
+ Path(AITER_CSRC_DIR) / "kernels" / "mla" / "metadata" / "v1_2_device.cuh"
108
+ )
109
+ source = "".join(metadata_source.read_text(encoding="utf-8").split())
110
+ except (ImportError, OSError):
111
+ return False
112
+ return "num_heads==24" in source
113
+
114
+
115
+def _aiter_mla_native_h24_supported() -> bool:
116
+ """Whether the complete AITER decode path supports native H24."""
117
+ return (
118
+ _aiter_mla_native_h24_reducer_supported()
119
+ and _aiter_mla_native_h24_metadata_supported()
120
+ )
121
+
122
+
123
@functools.lru_cache(maxsize=1)
124
def _gluon_mla_decode_supported() -> bool:
125
"""The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build.
@@ -308,10 +351,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
351
352
from aiter import dtypes, get_mla_metadata_info_v1
353
311
- # For num_attention_heads < 16 (e.g. kimi-k2.5 head=8 with TP8),
312
- # make sure get_mla_metadata_info_v1 / get_mla_metadata_v1 are consistent
313
- # with the actual tensor shape passed to mla_decode_fwd.
314
- self._num_attention_heads = max(16, self.num_heads)
354
+ # Keep metadata sizing consistent with the padded tensor shape passed
355
+ # to mla_decode_fwd.
356
+ self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads(
357
+ self.num_heads
358
+ )
359
kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
360
if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
361
kv_cache_dtype_str = "fp8"
@@ -839,14 +883,16 @@ def _expand_page_indices_kernel(
883
884
class AiterMLAHelper:
885
"""
842
- AITER MLA persistent (asm) decode requires num_heads >= 16. Head counts
843
- < 16 are padded up to exactly 16: divisors of 16 by repeat_interleave,
844
- other counts (e.g. 12 heads/rank at TP8, 6 at TP16) by tiling the query
845
- heads and slicing to 16. Non-divisor padded decodes take the asm path;
846
- divisors and max_qo_len > 1 small-head verify still use Gluon.
886
+ AITER MLA persistent (asm) decode requires a multiple of 16 heads. Unaligned
887
+ head counts through 128 are padded to the next multiple of 16 by tiling the
888
+ query heads and slicing to the padded size. Native H24 AITER builds bypass
889
+ that padding. Small divisors of 16 retain the existing repeat_interleave and
890
+ strided-unpad behavior. Native and aligned counts pass through without
891
+ copies.
892
"""
893
894
_AITER_MIN_MLA_HEADS: Final = 16
895
+ _AITER_MAX_PADDED_MLA_HEADS: Final = 128
896
# Largest qlen the padded gqa=16 asm decode has a bf16 persistent kernel
897
# for. Above it only the non-persistent qseqlen=8 entry exists, and the
898
# fold that reaches a persistent one is gfx950-only.
@@ -856,48 +902,50 @@ class AiterMLAHelper:
902
@staticmethod
903
def check_num_heads_validity(num_heads: int):
904
assert AiterMLAHelper.is_valid_num_heads(num_heads), (
859
- "ROCM AITER MLA requires 1-15 heads (padded to 16 for asm "
860
- "persistent decode; exact divisors of 16 may keep Gluon) or a "
861
- f"multiple of 16 heads, but got {num_heads}.\n"
905
+ "ROCM AITER MLA requires a positive multiple of 16 heads, or an "
906
+ "unaligned head count up to 128 (padded to the next multiple of "
907
+ f"16), but got {num_heads}.\n"
908
f"Try adjusting tensor_parallel_size value."
909
)
910
911
@staticmethod
912
def is_valid_num_heads(num_heads: int) -> bool:
867
- return num_heads > 0 and (
868
- num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS
869
- or num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
913
+ return (
914
+ num_heads > 0
915
+ and num_heads not in AiterMLAHelper._AITER_UNSUPPORTED_HEADS
916
+ and (
917
+ num_heads <= AiterMLAHelper._AITER_MAX_PADDED_MLA_HEADS
918
+ or num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
919
+ )
920
)
921
922
@staticmethod
923
def get_actual_mla_num_heads(num_heads: int) -> int:
874
- return max(num_heads, AiterMLAHelper._AITER_MIN_MLA_HEADS)
924
+ if num_heads == 24 and _aiter_mla_native_h24_supported():
925
+ return num_heads
926
+ m = AiterMLAHelper._AITER_MIN_MLA_HEADS
927
+ return -(-num_heads // m) * m
928
929
@staticmethod
930
def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:
878
- m = AiterMLAHelper._AITER_MIN_MLA_HEADS
879
- if num_heads >= m:
931
+ m = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
932
+ if num_heads == m:
933
return q
934
if m % num_heads == 0:
935
return q.repeat_interleave(m // num_heads, dim=1)
883
- # Non-divisor head counts (e.g. 12 heads/rank at TP8, 6 at TP16) cannot
884
- # be padded by repeat_interleave. Tile the query heads and slice to
885
- # exactly m; this reaches m for any 0 < num_heads < m (unlike a single
886
- # append, which under-pads when num_heads < m - num_heads). MLA
887
- # attention is independent per query head over the shared KV, so the
888
- # padding heads cannot affect heads [0:num_heads]; they are sliced back
889
- # off in get_mla_unpadded_o.
936
+ # Non-divisor head counts cannot be padded by repeat_interleave. Tile
937
+ # the query heads and slice to exactly m. MLA attention is independent
938
+ # per query head over the shared KV, so padding heads cannot affect
939
+ # heads [0:num_heads]; they are sliced back off the output.
940
reps = -(-m // num_heads) # ceil(m / num_heads)
891
- # Slicing a tiled tensor down to m yields a non-contiguous view whenever
892
- # reps * num_heads > m (the common case: TP8 12->24->16, TP16 6->18->16).
893
- # The asm persistent decode reads q as a packed [tokens, m, head_dim]
894
- # buffer, so materialize a contiguous copy. No-op when already contiguous.
941
+ # Slicing a tiled tensor yields a non-contiguous view. The asm decode
942
+ # reads q as packed [tokens, m, head_dim], so materialize it.
943
return q.repeat(1, reps, 1)[:, :m, :].contiguous()
944
945
@staticmethod
946
def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor:
899
- m = AiterMLAHelper._AITER_MIN_MLA_HEADS
900
- if num_heads >= m:
947
+ m = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
948
+ if num_heads == m:
949
return o
950
if m % num_heads == 0:
951
return o[:, :: m // num_heads, :]
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+5
-3
@@ -422,9 +422,11 @@ class ROCMAiterMLASparseMetadataBuilder(
422
# so the buffers are large enough for any decode shape we might see.
423
from aiter import dtypes, get_mla_metadata_info_v1
424
425
- # Aiter sparse MLA also requires num_heads >= 16 (will be padded by
426
- # AiterMLAHelper.get_mla_padded_q in forward).
427
- self._num_attention_heads = max(16, self.num_heads)
425
+ # Keep metadata sizing consistent with the padded tensor shape passed
426
+ # to the sparse decode kernel.
427
+ self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads(
428
+ self.num_heads
429
+ )
430
431
q_dtype = self.model_dtype
432
kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")