[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>

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)
 

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,
+    )

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: (

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: