diff --git a/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py b/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py index 7401203629..81c622db00 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py @@ -32,11 +32,9 @@ def _on_rocm_with_aiter() -> bool: return current_platform.is_rocm() and is_aiter_found() -# Unlike the fp8 persistent-decode metadata this file's neighbour guards, the -# verify flatten is not gfx950-only: it is selected by head count alone and is -# reached on any ROCm AITER MLA deployment with fewer than 16 heads per rank. -# The Gluon kernel is replaced by a spy here, so nothing below needs it to be -# runnable either. +# The production verify flatten is gfx950-only. This test replaces the Gluon +# kernel with a spy and forces the architecture feature probe so the metadata +# regression remains testable on every ROCm AITER CI runner. pytestmark = pytest.mark.skipif( not _on_rocm_with_aiter(), reason="ROCM_AITER_MLA verify flatten requires ROCm and AITER", @@ -186,9 +184,15 @@ def _run_verify_block(): # The Gluon kernel only reads the metadata this test is about, so a spy in # its place keeps the assertions independent of the AITER build. - with patch( - "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon", - lambda: spy, + with ( + patch( + "vllm.v1.attention.backends.mla.rocm_aiter_mla._gluon_mla_decode_supported", + lambda: True, + ), + patch( + "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon", + lambda: spy, + ), ): impl.forward_mqa((q_nope, q_pe), kv_cache, metadata, layer=None) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py b/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py new file mode 100644 index 0000000000..82ddce3164 --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_mla_head_padding.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Head-count padding + kernel selection for the ROCm AITER MLA backend. + +Kimi-K3 at TP8 puts 12 heads/rank and at TP16 puts 6 heads/rank on the AITER +MLA decode. The asm persistent decode requires exactly 16 heads, so small +head counts are tile-padded to 16 and the padding heads are sliced back off +the output. Divisor counts (1/2/4/8) may keep the Gluon kernel, but only on +gfx950 where that kernel has a build; every small head count on gfx942 (which +has no Gluon build) is routed to the asm persistent decode instead. +""" + +import math +import types + +import pytest +import torch +import torch.nn.functional as F + +from vllm._aiter_ops import is_aiter_found +from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( + AiterMLADecodeMetadata, + AiterMLAHelper, + AiterMLAImpl, + AiterMLAMetadata, +) + +NUM_HEADS = 12 +KV_LORA_RANK = 512 +QK_ROPE_HEAD_DIM = 64 +QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM +CONTEXT_LEN = 4096 +SCALE = 1.0 / math.sqrt(QK_HEAD_DIM) + +# Non-divisor counts must go through the tile-and-slice path; divisor counts +# (of 16) keep repeat_interleave. Both must pad to exactly 16 and round-trip. +NON_DIVISOR_HEADS = [3, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15] +DIVISOR_HEADS = [1, 2, 4, 8] + + +def _rocm_aiter_available() -> bool: + return current_platform.is_rocm() and is_aiter_found() and torch.cuda.is_available() + + +def _on_gfx950() -> bool: + if not (current_platform.is_rocm() and is_aiter_found()): + return False + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() + + +def _expected_tile_pad(q: torch.Tensor, num_heads: int, m: int = 16) -> torch.Tensor: + # get_mla_padded_q tiles the heads and slices to m, i.e. head i of the + # padded tensor is head (i % num_heads) of the input. + idx = [i % num_heads for i in range(m)] + return q[:, idx, :] + + +def _make_h12_decode_metadata(device: torch.device) -> AiterMLAMetadata: + page_indices = torch.arange(CONTEXT_LEN, dtype=torch.int32, device=device) + decode = AiterMLADecodeMetadata( + block_table=page_indices.view(1, -1), + seq_lens=torch.tensor([CONTEXT_LEN], dtype=torch.int32, device=device), + dcp_tot_seq_lens=None, + paged_kv_indptr=torch.tensor( + [0, CONTEXT_LEN], dtype=torch.int32, device=device + ), + paged_kv_indices=page_indices, + paged_kv_last_page_len=torch.ones(1, dtype=torch.int32, device=device), + qo_indptr=torch.tensor([0, 1], dtype=torch.int32, device=device), + attn_out_dtype=torch.bfloat16, + max_qo_len=1, + ) + return AiterMLAMetadata( + num_reqs=1, + max_query_len=1, + max_seq_len=CONTEXT_LEN, + num_actual_tokens=1, + query_start_loc=decode.qo_indptr, + slot_mapping=torch.tensor([CONTEXT_LEN - 1], dtype=torch.int64, device=device), + num_decodes=1, + num_decode_tokens=1, + num_prefills=0, + decode=decode, + ) + + +def test_h12_query_is_tile_padded_to_h16(): + q = torch.arange(2 * 12 * 4, dtype=torch.bfloat16).view(2, 12, 4) + + padded_q = AiterMLAHelper.get_mla_padded_q(12, q) + + assert padded_q.shape == (2, 16, 4) + assert padded_q.is_contiguous() + # The real heads are untouched... + torch.testing.assert_close(padded_q[:, :12], q) + # ...and the 4 padding heads are the tiled wrap-around (heads 0..3), not + # zeros: MLA attention is per-head independent so duplicate query heads are + # harmless and get sliced back off the output. + torch.testing.assert_close(padded_q[:, 12:], q[:, :4]) + + +def test_h6_tp16_query_is_padded_to_h16(): + # TP16 puts 6 heads/rank. The old append-only padding produced cat(6, 6) = + # 12 heads and broke the asm kernel; tile-and-slice must reach exactly 16. + q = torch.arange(2 * 6 * 4, dtype=torch.bfloat16).view(2, 6, 4) + + padded_q = AiterMLAHelper.get_mla_padded_q(6, q) + + assert padded_q.shape == (2, 16, 4) + assert padded_q.is_contiguous() + torch.testing.assert_close(padded_q, _expected_tile_pad(q, 6)) + + +def test_h12_output_discards_padding_heads(): + o = torch.arange(2 * 16 * 4, dtype=torch.bfloat16).view(2, 16, 4) + + unpadded_o = AiterMLAHelper.get_mla_unpadded_o(12, o) + + assert unpadded_o.shape == (2, 12, 4) + torch.testing.assert_close(unpadded_o, o[:, :12]) + + +def test_existing_divisor_head_mapping_is_unchanged(): + q = torch.arange(2 * 8 * 4, dtype=torch.bfloat16).view(2, 8, 4) + + padded_q = AiterMLAHelper.get_mla_padded_q(8, q) + unpadded_o = AiterMLAHelper.get_mla_unpadded_o(8, padded_q) + + # Divisor counts still use repeat_interleave / strided unpad, unchanged. + torch.testing.assert_close(padded_q, q.repeat_interleave(2, dim=1)) + torch.testing.assert_close(unpadded_o, q) + + +@pytest.mark.parametrize("num_heads", NON_DIVISOR_HEADS + DIVISOR_HEADS) +def test_all_small_head_counts_pad_to_16_and_round_trip(num_heads: int): + q = torch.arange(2 * num_heads * 4, dtype=torch.float32).view(2, num_heads, 4) + + padded_q = AiterMLAHelper.get_mla_padded_q(num_heads, q) + unpadded_o = AiterMLAHelper.get_mla_unpadded_o(num_heads, padded_q) + + assert padded_q.shape == (2, 16, 4) + assert padded_q.is_contiguous() + # The real heads survive the pad -> unpad round trip exactly. + torch.testing.assert_close(unpadded_o, q) + + +def test_num_heads_ge_16_is_passthrough(): + q = torch.arange(2 * 16 * 4, dtype=torch.float32).view(2, 16, 4) + assert AiterMLAHelper.get_mla_padded_q(16, q) is q + assert AiterMLAHelper.get_mla_unpadded_o(16, q) is q + + +def test_is_valid_num_heads(): + for n in range(1, 16): + assert AiterMLAHelper.is_valid_num_heads(n) + assert AiterMLAHelper.is_valid_num_heads(16) + assert AiterMLAHelper.is_valid_num_heads(32) + assert not AiterMLAHelper.is_valid_num_heads(0) + + +def test_nondivisor_and_multitoken_never_use_gluon(): + # Non-divisor decode always takes the asm path (12 heads/rank at TP8). + assert not AiterMLAHelper.use_gluon_decode(12, 1) + assert not AiterMLAHelper.use_gluon_decode(6, 1) + # >=16 heads never pad, never Gluon. + assert not AiterMLAHelper.use_gluon_decode(16, 1) + # Multi-token (verify / qlen>1) is never the single-token Gluon decode. + assert not AiterMLAHelper.use_gluon_decode(8, 4) + assert not AiterMLAHelper.use_gluon_decode(12, 4) + + +def test_divisor_gluon_selection_follows_arch(): + # Divisor head counts keep Gluon only where the kernel exists (gfx950). + # On gfx942 / non-ROCm they must route to the asm persistent decode. + on_gfx950 = _on_gfx950() + assert AiterMLAHelper.use_gluon_decode(8, 1) is on_gfx950 + assert AiterMLAHelper.use_gluon_decode(4, 1) is on_gfx950 + assert AiterMLAHelper.use_gluon_decode(1, 1) is on_gfx950 + + +def test_asm_padding_env_default_is_auto(monkeypatch): + monkeypatch.delenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", raising=False) + import vllm.envs as envs + + assert envs.VLLM_ROCM_AITER_MLA_ASM_PADDING == "auto" + + +def test_asm_padding_env_force_asm_disables_gluon(monkeypatch): + monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "asm") + # Forcing the asm path: no small-head count uses Gluon on any arch. + for num_heads in (1, 2, 4, 8, 12): + assert not AiterMLAHelper.use_gluon_decode(num_heads, 1) + + +def test_asm_padding_env_force_gluon_follows_arch(monkeypatch): + monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "gluon") + on_gfx950 = _on_gfx950() + # Forcing Gluon: any 1..15 single-token decode uses it where a build exists + # (gfx950), including non-divisor counts like 12; gfx942/non-ROCm still + # falls back to the asm path. + for num_heads in (1, 2, 4, 8, 12): + assert AiterMLAHelper.use_gluon_decode(num_heads, 1) is on_gfx950 + + +def test_asm_padding_env_auto_matches_arch_gate(monkeypatch): + monkeypatch.setenv("VLLM_ROCM_AITER_MLA_ASM_PADDING", "auto") + on_gfx950 = _on_gfx950() + # auto: divisor counts keep Gluon on gfx950, non-divisor counts take asm. + assert AiterMLAHelper.use_gluon_decode(8, 1) is on_gfx950 + assert not AiterMLAHelper.use_gluon_decode(12, 1) + + +@pytest.mark.skipif( + not _rocm_aiter_available(), + reason="12-head AITER MLA asm persistent decode needs ROCm + AITER", +) +@torch.inference_mode() +def test_h12_aiter_mla_decode_matches_reference(): + """The 12-head AITER decode output must match an attention reference. + + 12 is a non-divisor of 16 so use_gluon_decode is False on every arch; the + query is tile-padded to 16 heads and served by the asm persistent decode + (the CDNA3 path this PR enables, and unchanged on CDNA4). + """ + device = torch.device("cuda:0") + torch.manual_seed(42) + + q = torch.randn(1, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16, device=device) + kv_cache = torch.randn( + CONTEXT_LEN, 1, QK_HEAD_DIM, dtype=torch.bfloat16, device=device + ) + + impl = object.__new__(AiterMLAImpl) + impl.num_heads = NUM_HEADS + impl.kv_lora_rank = KV_LORA_RANK + impl.scale = SCALE + + one = torch.ones(1, dtype=torch.float32, device=device) + layer = types.SimpleNamespace(_q_scale=one, _k_scale=one) + metadata = _make_h12_decode_metadata(device) + + out, _ = impl.forward_mqa(q, kv_cache, metadata, layer) + + assert out.shape[1] == NUM_HEADS # padding heads sliced back off + + key = kv_cache[:, 0].float().unsqueeze(0) + value = key[..., :KV_LORA_RANK] + out_ref = F.scaled_dot_product_attention( + q[0].float().unsqueeze(1), + key, + value, + scale=SCALE, + enable_gqa=True, + ).squeeze(1) + + assert torch.isfinite(out).all() + torch.testing.assert_close( + out.float(), + out_ref.unsqueeze(0), + atol=1e-2, + rtol=1e-2, + ) diff --git a/vllm/envs.py b/vllm/envs.py index 98014dc498..3f23fb6152 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -133,6 +133,7 @@ if TYPE_CHECKING: VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True + VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" VLLM_ROCM_USE_AITER_MHA: bool = True VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False @@ -1237,6 +1238,20 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ROCM_USE_AITER_MLA": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1") ), + # Small-head (<16) AITER MLA decode kernel selection. Small head counts + # (e.g. Kimi-K3: 12 heads/rank at TP8, 6 at TP16) can decode either through + # the Gluon small-head kernel or through the padded persistent-scheduling + # (PS) ASM kernel. "auto" (default) keeps Gluon for head counts that divide + # 16 where a Gluon build exists (gfx950/CDNA4) and otherwise uses the padded + # PS ASM decode; "gluon" forces the Gluon path wherever a build exists; + # "asm" forces the padded PS ASM decode. On gfx942/CDNA3 there is no Gluon + # build, so the ASM path is always used regardless of this setting. + "VLLM_ROCM_AITER_MLA_ASM_PADDING": env_with_choices( + "VLLM_ROCM_AITER_MLA_ASM_PADDING", + "auto", + ["auto", "gluon", "asm"], + case_sensitive=False, + ), # Whether to use aiter mha ops. # By default is enabled. "VLLM_ROCM_USE_AITER_MHA": lambda: ( diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 43584998a0..f10b9cff36 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -77,6 +77,50 @@ def _fp8_mla_prefill_supported() -> bool: return True +@functools.lru_cache(maxsize=1) +def _gluon_mla_decode_supported() -> bool: + """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build. + + Its tiling needs ~160 KiB of LDS, which exceeds CDNA3's 64 KiB, so on + gfx942 there is no kernel to fall through to and selecting it asserts + (``mla_gluon requires gfx950``). Restrict Gluon decode to gfx950; other + archs use the asm persistent decode, which ``get_mla_padded_q`` makes + correct for any 1..15 heads. + """ + try: + from vllm.platforms.rocm import on_gfx950 + except Exception: # noqa: BLE001 + return False + return on_gfx950() + + +def _aiter_mla_small_head_mode() -> str: + """Small-head (<16) MLA decode kernel selection. + + Controlled by ``VLLM_ROCM_AITER_MLA_ASM_PADDING``: + + - ``"auto"`` (default): let the arch decide -- divisor head counts keep the + Gluon decode where a build exists (gfx950), everything else (non-divisor + counts and all counts on gfx942) uses the padded persistent-scheduling + ASM decode. + - ``"gluon"``: prefer the Gluon path wherever a build exists. + - ``"asm"``: force the padded persistent-scheduling ASM decode. + + On gfx942 (no Gluon build) the ASM path is always used regardless of this + setting; ``"gluon"`` there falls back to ASM with a one-time warning. + """ + import vllm.envs as envs + + mode = (envs.VLLM_ROCM_AITER_MLA_ASM_PADDING or "auto").lower() + if mode == "gluon" and not _gluon_mla_decode_supported(): + logger.warning_once( + "VLLM_ROCM_AITER_MLA_ASM_PADDING=gluon requested, but this device " + "has no Gluon MLA decode build (Gluon requires gfx950); using the " + "padded persistent-scheduling ASM decode instead." + ) + return mode + + class AiterMLABackend(MLACommonBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ @@ -781,8 +825,11 @@ def _expand_page_indices_kernel( class AiterMLAHelper: """ - AITER MLA implementation requires num_heads >= 16. If num_heads < 16 and - 16 % num_heads == 0, we can pad q to 16 heads; otherwise AITER has to fail. + AITER MLA persistent (asm) decode requires num_heads >= 16. Head counts + < 16 are padded up to exactly 16: divisors of 16 by repeat_interleave, + other counts (e.g. 12 heads/rank at TP8, 6 at TP16) by tiling the query + heads and slicing to 16. Non-divisor padded decodes take the asm path; + divisors and max_qo_len > 1 small-head verify still use Gluon. """ _AITER_MIN_MLA_HEADS: Final = 16 @@ -791,8 +838,9 @@ class AiterMLAHelper: @staticmethod def check_num_heads_validity(num_heads: int): assert AiterMLAHelper.is_valid_num_heads(num_heads), ( - "ROCM AITER MLA requires 1-15 heads for Gluon decode or a multiple " - f"of 16 heads for persistent decode, but got {num_heads}.\n" + "ROCM AITER MLA requires 1-15 heads (padded to 16 for asm " + "persistent decode; exact divisors of 16 may keep Gluon) or a " + f"multiple of 16 heads, but got {num_heads}.\n" f"Try adjusting tensor_parallel_size value." ) @@ -809,25 +857,55 @@ class AiterMLAHelper: @staticmethod def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor: - return ( - q - if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS - else q.repeat_interleave( - AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, dim=1 - ) - ) + m = AiterMLAHelper._AITER_MIN_MLA_HEADS + if num_heads >= m: + return q + if m % num_heads == 0: + return q.repeat_interleave(m // num_heads, dim=1) + # Non-divisor head counts (e.g. 12 heads/rank at TP8, 6 at TP16) cannot + # be padded by repeat_interleave. Tile the query heads and slice to + # exactly m; this reaches m for any 0 < num_heads < m (unlike a single + # append, which under-pads when num_heads < m - num_heads). MLA + # attention is independent per query head over the shared KV, so the + # padding heads cannot affect heads [0:num_heads]; they are sliced back + # off in get_mla_unpadded_o. + reps = -(-m // num_heads) # ceil(m / num_heads) + # Slicing a tiled tensor down to m yields a non-contiguous view whenever + # reps * num_heads > m (the common case: TP8 12->24->16, TP16 6->18->16). + # The asm persistent decode reads q as a packed [tokens, m, head_dim] + # buffer, so materialize a contiguous copy. No-op when already contiguous. + return q.repeat(1, reps, 1)[:, :m, :].contiguous() @staticmethod def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor: - return ( - o - if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS - else o[:, :: AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, :] - ) + m = AiterMLAHelper._AITER_MIN_MLA_HEADS + if num_heads >= m: + return o + if m % num_heads == 0: + return o[:, :: m // num_heads, :] + # Undo the tile-padding from get_mla_padded_q: the real heads are the + # first num_heads. + return o[:, :num_heads, :] @staticmethod def use_gluon_decode(num_heads: int, max_qo_len: int) -> bool: - return num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS and max_qo_len == 1 + # Small-head (<16) single-token decode can use either the Gluon kernel + # or the padded asm persistent decode, selected by + # VLLM_ROCM_AITER_MLA_ASM_PADDING (see _aiter_mla_small_head_mode) and + # the arch: Gluon only has a gfx950 build. In "auto" (default) mode + # divisor counts keep Gluon on gfx950 and everything else -- non-divisor + # counts (e.g. 12 heads/rank at TP8) and all counts on gfx942 -- takes + # the asm path, which get_mla_padded_q pads to exactly 16. + m = AiterMLAHelper._AITER_MIN_MLA_HEADS + if num_heads >= m or max_qo_len != 1: + return False + mode = _aiter_mla_small_head_mode() + if mode == "asm": + return False + gluon_supported = _gluon_mla_decode_supported() + if mode == "gluon": + return gluon_supported + return m % num_heads == 0 and gluon_supported class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): @@ -1113,9 +1191,15 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): # target is checking draft tokens, so position t must not see t+1 -- # and attention rows are independent, so giving row t the KV range # [0, context + t] is exactly causal multi-token attention. + # Gluon only has a gfx950 build, so gate on the arch (mirrors + # use_gluon_decode). VLLM_ROCM_AITER_MLA_ASM_PADDING=asm also forces the + # asm path here. Otherwise this falls through to the asm decode, which + # pads to 16 heads and handles qlen>1 verify directly. if ( self.num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS and int(decode.max_qo_len) > 1 + and _aiter_mla_small_head_mode() != "asm" + and _gluon_mla_decode_supported() ): qlen = int(decode.max_qo_len) if type(q) is tuple: