| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 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 |
| 12 | import types |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | 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, |
| 24 | AiterMLAImpl, |
| 25 | AiterMLAMetadata, |
| 26 | ) |
| 27 | |
| 28 | NUM_HEADS = 12 |
| 29 | KV_LORA_RANK = 512 |
| 30 | QK_ROPE_HEAD_DIM = 64 |
| 31 | QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM |
| 32 | CONTEXT_LEN = 4096 |
| 33 | SCALE = 1.0 / math.sqrt(QK_HEAD_DIM) |
| 34 | |
| 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 | |
| 52 | |
| 53 | def _on_gfx950() -> bool: |
| 54 | if not (current_platform.is_rocm() and is_aiter_found()): |
| 55 | return False |
| 56 | from vllm.platforms.rocm import on_gfx950 |
| 57 | |
| 58 | return on_gfx950() |
| 59 | |
| 60 | |
| 61 | def _expected_tile_pad(q: torch.Tensor, num_heads: int, m: int = 16) -> torch.Tensor: |
| 62 | # get_mla_padded_q tiles the heads and slices to m, i.e. head i of the |
| 63 | # padded tensor is head (i % num_heads) of the input. |
| 64 | idx = [i % num_heads for i in range(m)] |
| 65 | return q[:, idx, :] |
| 66 | |
| 67 | |
| 68 | def _make_h12_decode_metadata(device: torch.device) -> AiterMLAMetadata: |
| 69 | page_indices = torch.arange(CONTEXT_LEN, dtype=torch.int32, device=device) |
| 70 | decode = AiterMLADecodeMetadata( |
| 71 | block_table=page_indices.view(1, -1), |
| 72 | seq_lens=torch.tensor([CONTEXT_LEN], dtype=torch.int32, device=device), |
| 73 | dcp_tot_seq_lens=None, |
| 74 | paged_kv_indptr=torch.tensor( |
| 75 | [0, CONTEXT_LEN], dtype=torch.int32, device=device |
| 76 | ), |
| 77 | paged_kv_indices=page_indices, |
| 78 | paged_kv_last_page_len=torch.ones(1, dtype=torch.int32, device=device), |
| 79 | qo_indptr=torch.tensor([0, 1], dtype=torch.int32, device=device), |
| 80 | attn_out_dtype=torch.bfloat16, |
| 81 | max_qo_len=1, |
| 82 | ) |
| 83 | return AiterMLAMetadata( |
| 84 | num_reqs=1, |
| 85 | max_query_len=1, |
| 86 | max_seq_len=CONTEXT_LEN, |
| 87 | num_actual_tokens=1, |
| 88 | query_start_loc=decode.qo_indptr, |
| 89 | slot_mapping=torch.tensor([CONTEXT_LEN - 1], dtype=torch.int64, device=device), |
| 90 | num_decodes=1, |
| 91 | num_decode_tokens=1, |
| 92 | num_prefills=0, |
| 93 | decode=decode, |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | def test_h12_query_is_tile_padded_to_h16(): |
| 98 | q = torch.arange(2 * 12 * 4, dtype=torch.bfloat16).view(2, 12, 4) |
| 99 | |
| 100 | padded_q = AiterMLAHelper.get_mla_padded_q(12, q) |
| 101 | |
| 102 | assert padded_q.shape == (2, 16, 4) |
| 103 | assert padded_q.is_contiguous() |
| 104 | # The real heads are untouched... |
| 105 | torch.testing.assert_close(padded_q[:, :12], q) |
| 106 | # ...and the 4 padding heads are the tiled wrap-around (heads 0..3), not |
| 107 | # zeros: MLA attention is per-head independent so duplicate query heads are |
| 108 | # harmless and get sliced back off the output. |
| 109 | torch.testing.assert_close(padded_q[:, 12:], q[:, :4]) |
| 110 | |
| 111 | |
| 112 | def test_h6_tp16_query_is_padded_to_h16(): |
| 113 | # TP16 puts 6 heads/rank. The old append-only padding produced cat(6, 6) = |
| 114 | # 12 heads and broke the asm kernel; tile-and-slice must reach exactly 16. |
| 115 | q = torch.arange(2 * 6 * 4, dtype=torch.bfloat16).view(2, 6, 4) |
| 116 | |
| 117 | padded_q = AiterMLAHelper.get_mla_padded_q(6, q) |
| 118 | |
| 119 | assert padded_q.shape == (2, 16, 4) |
| 120 | assert padded_q.is_contiguous() |
| 121 | torch.testing.assert_close(padded_q, _expected_tile_pad(q, 6)) |
| 122 | |
| 123 | |
| 124 | def test_h12_output_discards_padding_heads(): |
| 125 | o = torch.arange(2 * 16 * 4, dtype=torch.bfloat16).view(2, 16, 4) |
| 126 | |
| 127 | unpadded_o = AiterMLAHelper.get_mla_unpadded_o(12, o) |
| 128 | |
| 129 | assert unpadded_o.shape == (2, 12, 4) |
| 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 | |
| 202 | padded_q = AiterMLAHelper.get_mla_padded_q(8, q) |
| 203 | unpadded_o = AiterMLAHelper.get_mla_unpadded_o(8, padded_q) |
| 204 | |
| 205 | # Divisor counts still use repeat_interleave / strided unpad, unchanged. |
| 206 | torch.testing.assert_close(padded_q, q.repeat_interleave(2, dim=1)) |
| 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) |
| 225 | |
| 226 | padded_q = AiterMLAHelper.get_mla_padded_q(num_heads, q) |
| 227 | unpadded_o = AiterMLAHelper.get_mla_unpadded_o(num_heads, padded_q) |
| 228 | |
| 229 | assert padded_q.shape == (2, 16, 4) |
| 230 | assert padded_q.is_contiguous() |
| 231 | # The real heads survive the pad -> unpad round trip exactly. |
| 232 | torch.testing.assert_close(unpadded_o, q) |
| 233 | |
| 234 | |
| 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(): |
| 242 | for n in range(1, 129): |
| 243 | assert AiterMLAHelper.is_valid_num_heads(n) |
| 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") |
| 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") |
| 262 | |
| 263 | |
| 264 | def test_divisor_gluon_selection_follows_arch(): |
| 265 | # Divisor head counts keep Gluon only where the kernel exists (gfx950). |
| 266 | # On gfx942 / non-ROCm they must route to the asm persistent decode. |
| 267 | on_gfx950 = _on_gfx950() |
| 268 | assert AiterMLAHelper.use_gluon_decode(8, 1, "auto") is on_gfx950 |
| 269 | assert AiterMLAHelper.use_gluon_decode(4, 1, "auto") is on_gfx950 |
| 270 | assert AiterMLAHelper.use_gluon_decode(1, 1, "auto") is on_gfx950 |
| 271 | |
| 272 | |
| 273 | def test_asm_padding_env_default_is_auto(monkeypatch): |
| 274 | monkeypatch.delenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", raising=False) |
| 275 | import vllm.envs as envs |
| 276 | |
| 277 | assert envs.VLLM_ROCM_AITER_MLA_ASM_PADDING == "auto" |
| 278 | |
| 279 | |
| 280 | def test_asm_padding_env_force_asm_disables_gluon(monkeypatch): |
| 281 | monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "asm") |
| 282 | # Forcing the asm path: no small-head count uses Gluon on any arch. |
| 283 | for num_heads in (1, 2, 4, 8, 12): |
| 284 | assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, "auto") |
| 285 | |
| 286 | |
| 287 | def test_asm_padding_env_force_gluon_follows_arch(monkeypatch): |
| 288 | monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "gluon") |
| 289 | on_gfx950 = _on_gfx950() |
| 290 | # Forcing Gluon: any 1..15 single-token decode uses it where a build exists |
| 291 | # (gfx950), including non-divisor counts like 12; gfx942/non-ROCm still |
| 292 | # falls back to the asm path. |
| 293 | for num_heads in (1, 2, 4, 8, 12): |
| 294 | assert AiterMLAHelper.use_gluon_decode(num_heads, 1, "auto") is on_gfx950 |
| 295 | |
| 296 | |
| 297 | def test_asm_padding_env_auto_matches_arch_gate(monkeypatch): |
| 298 | monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "auto") |
| 299 | on_gfx950 = _on_gfx950() |
| 300 | # auto: divisor counts keep Gluon on gfx950, non-divisor counts take asm. |
| 301 | assert AiterMLAHelper.use_gluon_decode(8, 1, "auto") is on_gfx950 |
| 302 | assert not AiterMLAHelper.use_gluon_decode(12, 1, "auto") |
| 303 | |
| 304 | |
| 305 | @pytest.mark.skipif( |
| 306 | not _rocm_aiter_available(), |
| 307 | reason="12-head AITER MLA asm persistent decode needs ROCm + AITER", |
| 308 | ) |
| 309 | @torch.inference_mode() |
| 310 | def test_h12_aiter_mla_decode_matches_reference(): |
| 311 | """The 12-head AITER decode output must match an attention reference. |
| 312 | |
| 313 | 12 is a non-divisor of 16 so use_gluon_decode is False on every arch; the |
| 314 | query is tile-padded to 16 heads and served by the asm persistent decode |
| 315 | (the CDNA3 path this PR enables, and unchanged on CDNA4). |
| 316 | """ |
| 317 | device = torch.device("cuda:0") |
| 318 | torch.manual_seed(42) |
| 319 | |
| 320 | q = torch.randn(1, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16, device=device) |
| 321 | kv_cache = torch.randn( |
| 322 | CONTEXT_LEN, 1, QK_HEAD_DIM, dtype=torch.bfloat16, device=device |
| 323 | ) |
| 324 | |
| 325 | impl = object.__new__(AiterMLAImpl) |
| 326 | impl.num_heads = NUM_HEADS |
| 327 | impl.dcp_world_size = 1 |
| 328 | impl.kv_lora_rank = KV_LORA_RANK |
| 329 | impl.qk_rope_head_dim = QK_ROPE_HEAD_DIM |
| 330 | impl.scale = SCALE |
| 331 | impl.kv_cache_dtype = "auto" |
| 332 | |
| 333 | one = torch.ones(1, dtype=torch.float32, device=device) |
| 334 | layer = types.SimpleNamespace(_q_scale=one, _k_scale=one) |
| 335 | metadata = _make_h12_decode_metadata(device) |
| 336 | |
| 337 | out, _ = impl.forward_mqa(q, kv_cache, metadata, layer) |
| 338 | |
| 339 | assert out.shape[1] == NUM_HEADS # padding heads sliced back off |
| 340 | |
| 341 | key = kv_cache[:, 0].float().unsqueeze(0) |
| 342 | value = key[..., :KV_LORA_RANK] |
| 343 | out_ref = F.scaled_dot_product_attention( |
| 344 | q[0].float().unsqueeze(1), |
| 345 | key, |
| 346 | value, |
| 347 | scale=SCALE, |
| 348 | enable_gqa=True, |
| 349 | ).squeeze(1) |
| 350 | |
| 351 | assert torch.isfinite(out).all() |
| 352 | torch.testing.assert_close( |
| 353 | out.float(), |
| 354 | out_ref.unsqueeze(0), |
| 355 | atol=1e-2, |
| 356 | rtol=1e-2, |
| 357 | ) |