| 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 | Gluon is still the route for supported non-DCP decode shapes. DCP multi-token |
| 13 | verification uses segmented MLA instead, so it depends on neither Gluon's |
| 14 | query-head nor its long-context pipeline limits. |
| 15 | """ |
| 16 | |
| 17 | import pytest |
| 18 | import torch |
| 19 | |
| 20 | from vllm.platforms import current_platform |
| 21 | |
| 22 | if not current_platform.is_rocm(): |
| 23 | pytest.skip("ROCm AITER MLA tests", allow_module_level=True) |
| 24 | |
| 25 | from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402 |
| 26 | from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper # noqa: E402 |
| 27 | |
| 28 | FP8_DTYPES = ["fp8", "fp8_e4m3", "fp8_e5m2"] |
| 29 | UNQUANTIZED_DTYPES = ["auto", "bfloat16"] |
| 30 | |
| 31 | |
| 32 | @pytest.fixture |
| 33 | def gluon_available(monkeypatch): |
| 34 | """Pin the arch gate and the mode knob so only the dtype rules are in play. |
| 35 | |
| 36 | Gluon ships a gfx950 build only, and VLLM_ROCM_AITER_MLA_ASM_PADDING can |
| 37 | force the asm path on any arch. Without pinning both, the expectations that |
| 38 | bf16 *keeps* Gluon would pass on a gfx950 host and fail on gfx942 for |
| 39 | reasons that have nothing to do with what these tests cover. |
| 40 | """ |
| 41 | monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True) |
| 42 | monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto") |
| 43 | |
| 44 | |
| 45 | @pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES) |
| 46 | @pytest.mark.parametrize("num_heads", [1, 2, 4, 5, 6, 8, 12, 16, 32, 128]) |
| 47 | @pytest.mark.parametrize("max_qo_len", [1, 2, 4, 5, 8, 15]) |
| 48 | def test_non_dcp_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len): |
| 49 | """No non-DCP fp8 shape reaches either Gluon entry point. |
| 50 | |
| 51 | The head count is deliberately swept across divisors of 16 as well as |
| 52 | non-divisors: the divisor case (e.g. 8 heads at TP8) is the one that stays |
| 53 | on Gluon for bf16 and so is the one an fp8 guard has to override. |
| 54 | """ |
| 55 | assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype) |
| 56 | assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype) |
| 57 | |
| 58 | |
| 59 | @pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES) |
| 60 | @pytest.mark.parametrize("num_heads", [1, 2, 4, 8, 12]) |
| 61 | @pytest.mark.parametrize("mode", ["auto", "gluon", "asm"]) |
| 62 | def test_fp8_never_routes_to_gluon_under_any_mode( |
| 63 | monkeypatch, kv_cache_dtype, num_heads, mode |
| 64 | ): |
| 65 | """VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force a non-DCP fp8 cache onto Gluon. |
| 66 | |
| 67 | The dtype guard deliberately precedes the mode knob: honouring an explicit |
| 68 | "gluon" request under fp8 would hand Gluon the batch it asserts against, so |
| 69 | it is overridden rather than obeyed. Pinned here because the override is a |
| 70 | correctness decision, not a preference. |
| 71 | """ |
| 72 | monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True) |
| 73 | monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: mode) |
| 74 | assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype) |
| 75 | assert not AiterMLAHelper.use_gluon_verify(num_heads, 8, kv_cache_dtype) |
| 76 | |
| 77 | |
| 78 | @pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES) |
| 79 | @pytest.mark.parametrize("num_heads", [16, 17, 32, 64, 128]) |
| 80 | @pytest.mark.parametrize("max_qo_len", [1, 4, 8]) |
| 81 | def test_large_head_counts_never_use_gluon(kv_cache_dtype, num_heads, max_qo_len): |
| 82 | """>= 16 heads has always been asm-only; the dtype guard must not change it.""" |
| 83 | assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype) |
| 84 | assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype) |
| 85 | |
| 86 | |
| 87 | @pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES) |
| 88 | @pytest.mark.parametrize("num_heads", [8, 16, 32, 96]) |
| 89 | @pytest.mark.parametrize("max_qo_len", [2, 3, 8]) |
| 90 | def test_dcp_multitoken_verify_never_uses_gluon( |
| 91 | gluon_available, kv_cache_dtype, num_heads, max_qo_len |
| 92 | ): |
| 93 | assert not AiterMLAHelper.use_gluon_verify( |
| 94 | num_heads, max_qo_len, kv_cache_dtype, dcp_world_size=8 |
| 95 | ) |
| 96 | |
| 97 | |
| 98 | def test_segmented_dcp_verify_does_not_depend_on_gluon(monkeypatch): |
| 99 | monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: False) |
| 100 | monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True) |
| 101 | |
| 102 | assert rocm_aiter_mla._segmented_dcp_verify_supported(8, 1) |
| 103 | # Round-robin interleaving other than 1 is not served by this route. |
| 104 | assert not rocm_aiter_mla._segmented_dcp_verify_supported(8, 4) |
| 105 | |
| 106 | |
| 107 | @pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES) |
| 108 | @pytest.mark.parametrize("num_heads", [1, 2, 4, 8]) |
| 109 | def test_unquantized_divisor_heads_keep_gluon_decode( |
| 110 | gluon_available, kv_cache_dtype, num_heads |
| 111 | ): |
| 112 | """Divisor head counts keep the Gluon single-token decode for bf16. |
| 113 | |
| 114 | Gluon does not scale with KV length the way the asm persistent decode does, |
| 115 | so this is the faster path where it is usable. |
| 116 | """ |
| 117 | assert AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype) |
| 118 | |
| 119 | |
| 120 | @pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES) |
| 121 | @pytest.mark.parametrize("num_heads", [3, 5, 6, 7, 9, 12, 15]) |
| 122 | def test_unquantized_non_divisor_heads_use_asm_decode(kv_cache_dtype, num_heads): |
| 123 | """Non-divisor head counts are padded to 16 and take the asm decode.""" |
| 124 | assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype) |
| 125 | |
| 126 | |
| 127 | @pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES) |
| 128 | @pytest.mark.parametrize("num_heads", [5, 8, 12]) |
| 129 | @pytest.mark.parametrize("max_qo_len", [2, 4, 8, 15]) |
| 130 | def test_unquantized_small_head_verify_keeps_gluon( |
| 131 | gluon_available, kv_cache_dtype, num_heads, max_qo_len |
| 132 | ): |
| 133 | """bf16 has no gqa<16, qseqlen>1 asm kernel, so verify still flattens onto Gluon. |
| 134 | |
| 135 | Unlike the decode predicate, this one does not care whether the head count |
| 136 | divides 16 -- the flatten reshapes to qseqlen=1 either way. |
| 137 | """ |
| 138 | assert AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype) |
| 139 | # The verify flatten is a separate entry point from the single-token decode. |
| 140 | assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype) |
| 141 | |
| 142 | |
| 143 | @pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES) |
| 144 | @pytest.mark.parametrize("num_heads", [5, 8, 12, 16, 32]) |
| 145 | def test_decode_and_verify_are_disjoint(kv_cache_dtype, num_heads): |
| 146 | """The two Gluon entry points partition on qlen and never both fire.""" |
| 147 | for max_qo_len in (1, 2, 8): |
| 148 | assert not ( |
| 149 | AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype) |
| 150 | and AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype) |
| 151 | ) |
| 152 | |
| 153 | |
| 154 | @pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES) |
| 155 | def test_oversized_kv_cache_refuses_gluon(gluon_available, kv_cache_dtype): |
| 156 | """Past 2 GiB per layer Gluon drops its KV bounds mask and reads out of range. |
| 157 | |
| 158 | ``None`` is the profiling case, before the cache is sized: nothing to |
| 159 | refuse yet, so the usual route stands. |
| 160 | """ |
| 161 | bound = rocm_aiter_mla._GLUON_MAX_KV_CACHE_BYTES |
| 162 | for kv_cache_bytes in (None, bound): |
| 163 | assert AiterMLAHelper.use_gluon_decode(8, 1, kv_cache_dtype, kv_cache_bytes) |
| 164 | assert AiterMLAHelper.use_gluon_verify( |
| 165 | 8, 8, kv_cache_dtype, kv_cache_bytes=kv_cache_bytes |
| 166 | ) |
| 167 | assert not AiterMLAHelper.use_gluon_decode(8, 1, kv_cache_dtype, bound + 1) |
| 168 | assert not AiterMLAHelper.use_gluon_verify( |
| 169 | 8, 8, kv_cache_dtype, kv_cache_bytes=bound + 1 |
| 170 | ) |
| 171 | |
| 172 | |
| 173 | @pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15]) |
| 174 | def test_padded_query_is_contiguous(num_heads): |
| 175 | """asm_mla.cu:805 requires Q.is_contiguous(). |
| 176 | |
| 177 | Padding a non-divisor head count to 16 tiles the query and slices it back |
| 178 | down, which yields a non-contiguous view whenever more than one tile is |
| 179 | needed (12 heads -> repeat to 24 -> slice to 16). The asm kernel rejects |
| 180 | that outright, so the padding has to materialize the result. |
| 181 | """ |
| 182 | q = torch.randn(4, num_heads, 576) |
| 183 | padded = AiterMLAHelper.get_mla_padded_q(num_heads, q) |
| 184 | |
| 185 | assert padded.shape[1] == max(16, num_heads) |
| 186 | assert padded.is_contiguous() |
| 187 | |
| 188 | |
| 189 | @pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15, 16, 32]) |
| 190 | def test_pad_unpad_round_trip_preserves_head_order(num_heads): |
| 191 | """Unpadding recovers each original head from the padded output, in order.""" |
| 192 | q = torch.arange(num_heads, dtype=torch.float32).view(1, num_heads, 1) |
| 193 | padded = AiterMLAHelper.get_mla_padded_q(num_heads, q) |
| 194 | unpadded = AiterMLAHelper.get_mla_unpadded_o(num_heads, padded) |
| 195 | |
| 196 | assert unpadded.shape == q.shape |
| 197 | torch.testing.assert_close(unpadded, q) |
| 198 | |
| 199 | |
| 200 | def test_a_non_causal_block_never_routes_to_gluon(gluon_available): |
| 201 | """A small-head block that normally uses Gluon must use ASM when non-causal.""" |
| 202 | assert not AiterMLAHelper.use_gluon_verify(12, 8, "auto", causal=False) |