iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][MLA] [K3] Fix fp8 KV cache decode on the AITER MLA backend (#51011)

Browse Source 
Signed-off-by: fanxingran <xingran.fan@amd.com>
This commit is contained in:
fanxingran
2026-08-10 11:42:10 -04:00
committed by GitHub
parent 8fea1d306f
commit 37fbf52084
4 changed files with 366 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_rocm_aiter_mla_head_padding.py
+13 -12
View File
@@ -163,22 +163,22 @@ def test_is_valid_num_heads():
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)
assert not AiterMLAHelper.use_gluon_decode(12, 1, "auto")
assert not AiterMLAHelper.use_gluon_decode(6, 1, "auto")
# >=16 heads never pad, never Gluon.
assert not AiterMLAHelper.use_gluon_decode(16, 1)
assert not AiterMLAHelper.use_gluon_decode(16, 1, "auto")
# 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)
assert not AiterMLAHelper.use_gluon_decode(8, 4, "auto")
assert not AiterMLAHelper.use_gluon_decode(12, 4, "auto")
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
assert AiterMLAHelper.use_gluon_decode(8, 1, "auto") is on_gfx950
assert AiterMLAHelper.use_gluon_decode(4, 1, "auto") is on_gfx950
assert AiterMLAHelper.use_gluon_decode(1, 1, "auto") is on_gfx950
def test_asm_padding_env_default_is_auto(monkeypatch):
@@ -192,7 +192,7 @@ 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)
assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, "auto")
def test_asm_padding_env_force_gluon_follows_arch(monkeypatch):
@@ -202,15 +202,15 @@ def test_asm_padding_env_force_gluon_follows_arch(monkeypatch):
# (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
assert AiterMLAHelper.use_gluon_decode(num_heads, 1, "auto") 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)
assert AiterMLAHelper.use_gluon_decode(8, 1, "auto") is on_gfx950
assert not AiterMLAHelper.use_gluon_decode(12, 1, "auto")
@pytest.mark.skipif(
@@ -237,6 +237,7 @@ def test_h12_aiter_mla_decode_matches_reference():
impl.num_heads = NUM_HEADS
impl.kv_lora_rank = KV_LORA_RANK
impl.scale = SCALE
impl.kv_cache_dtype = "auto"
one = torch.ones(1, dtype=torch.float32, device=device)
layer = types.SimpleNamespace(_q_scale=one, _k_scale=one)
tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py
+154
View File
@@ -0,0 +1,154 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Decode-kernel routing for the ROCm AITER MLA backend.
Gluon exposes a single fp8-KV regime, bh16bn128. It is a bf16-query kernel that
upcasts the cache in registers with a hardcoded scale of 1.0, and it asserts
batch_size == 1, so it cannot serve a decode batch. Every fp8 shape therefore
has to land on the asm kernels, which ship real fp8 variants for gqa=16. These
tests pin that down at the predicates, since the failure it prevents is either a
batch assertion or -- worse -- a silently wrong result.
"""
import pytest
import torch
from vllm.platforms import current_platform
if not current_platform.is_rocm():
pytest.skip("ROCm AITER MLA tests", allow_module_level=True)
from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402
from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper # noqa: E402
FP8_DTYPES = ["fp8", "fp8_e4m3", "fp8_e5m2"]
UNQUANTIZED_DTYPES = ["auto", "bfloat16"]
@pytest.fixture
def gluon_available(monkeypatch):
"""Pin the arch gate and the mode knob so only the dtype rules are in play.
Gluon ships a gfx950 build only, and VLLM_ROCM_AITER_MLA_ASM_PADDING can
force the asm path on any arch. Without pinning both, the expectations that
bf16 *keeps* Gluon would pass on a gfx950 host and fail on gfx942 for
reasons that have nothing to do with what these tests cover.
"""
monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True)
monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto")
@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES)
@pytest.mark.parametrize("num_heads", [1, 2, 4, 5, 6, 8, 12, 16, 32, 128])
@pytest.mark.parametrize("max_qo_len", [1, 2, 4, 5, 8, 15])
def test_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len):
"""No fp8 shape reaches either Gluon entry point.
The head count is deliberately swept across divisors of 16 as well as
non-divisors: the divisor case (e.g. 8 heads at TP8) is the one that stays
on Gluon for bf16 and so is the one an fp8 guard has to override.
"""
assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype)
assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES)
@pytest.mark.parametrize("num_heads", [1, 2, 4, 8, 12])
@pytest.mark.parametrize("mode", ["auto", "gluon", "asm"])
def test_fp8_never_routes_to_gluon_under_any_mode(
monkeypatch, kv_cache_dtype, num_heads, mode
):
"""VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force an fp8 cache onto Gluon.
The dtype guard deliberately precedes the mode knob: honouring an explicit
"gluon" request under fp8 would hand Gluon the batch it asserts against, so
it is overridden rather than obeyed. Pinned here because the override is a
correctness decision, not a preference.
"""
monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True)
monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: mode)
assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype)
assert not AiterMLAHelper.use_gluon_verify(num_heads, 8, kv_cache_dtype)
@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES)
@pytest.mark.parametrize("num_heads", [16, 17, 32, 64, 128])
@pytest.mark.parametrize("max_qo_len", [1, 4, 8])
def test_large_head_counts_never_use_gluon(kv_cache_dtype, num_heads, max_qo_len):
""">= 16 heads has always been asm-only; the dtype guard must not change it."""
assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype)
assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
@pytest.mark.parametrize("num_heads", [1, 2, 4, 8])
def test_unquantized_divisor_heads_keep_gluon_decode(
gluon_available, kv_cache_dtype, num_heads
):
"""Divisor head counts keep the Gluon single-token decode for bf16.
Gluon does not scale with KV length the way the asm persistent decode does,
so this is the faster path where it is usable.
"""
assert AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype)
@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
@pytest.mark.parametrize("num_heads", [3, 5, 6, 7, 9, 12, 15])
def test_unquantized_non_divisor_heads_use_asm_decode(kv_cache_dtype, num_heads):
"""Non-divisor head counts are padded to 16 and take the asm decode."""
assert not AiterMLAHelper.use_gluon_decode(num_heads, 1, kv_cache_dtype)
@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
@pytest.mark.parametrize("num_heads", [5, 8, 12])
@pytest.mark.parametrize("max_qo_len", [2, 4, 8, 15])
def test_unquantized_small_head_verify_keeps_gluon(
gluon_available, kv_cache_dtype, num_heads, max_qo_len
):
"""bf16 has no gqa<16, qseqlen>1 asm kernel, so verify still flattens onto Gluon.
Unlike the decode predicate, this one does not care whether the head count
divides 16 -- the flatten reshapes to qseqlen=1 either way.
"""
assert AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
# The verify flatten is a separate entry point from the single-token decode.
assert not AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype)
@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES)
@pytest.mark.parametrize("num_heads", [5, 8, 12, 16, 32])
def test_decode_and_verify_are_disjoint(kv_cache_dtype, num_heads):
"""The two Gluon entry points partition on qlen and never both fire."""
for max_qo_len in (1, 2, 8):
assert not (
AiterMLAHelper.use_gluon_decode(num_heads, max_qo_len, kv_cache_dtype)
and AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
)
@pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15])
def test_padded_query_is_contiguous(num_heads):
"""asm_mla.cu:805 requires Q.is_contiguous().
Padding a non-divisor head count to 16 tiles the query and slices it back
down, which yields a non-contiguous view whenever more than one tile is
needed (12 heads -> repeat to 24 -> slice to 16). The asm kernel rejects
that outright, so the padding has to materialize the result.
"""
q = torch.randn(4, num_heads, 576)
padded = AiterMLAHelper.get_mla_padded_q(num_heads, q)
assert padded.shape[1] == max(16, num_heads)
assert padded.is_contiguous()
@pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15, 16, 32])
def test_pad_unpad_round_trip_preserves_head_order(num_heads):
"""Unpadding recovers each original head from the padded output, in order."""
q = torch.arange(num_heads, dtype=torch.float32).view(1, num_heads, 1)
padded = AiterMLAHelper.get_mla_padded_q(num_heads, q)
unpadded = AiterMLAHelper.get_mla_unpadded_o(num_heads, padded)
assert unpadded.shape == q.shape
torch.testing.assert_close(unpadded, q)
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+135 -16
View File
@@ -64,10 +64,12 @@ def _builder(
kernel_block_size: int = 1,
max_decode_rows: int = 32,
num_heads: int = 16,
kv_cache_dtype: str = "auto",
):
return SimpleNamespace(
device=torch.device("cpu"),
num_heads=num_heads,
_kv_cache_dtype_str=kv_cache_dtype,
paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32),
paged_kv_indices=torch.empty(1024, dtype=torch.int32),
paged_kv_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32),
@@ -108,8 +110,22 @@ def test_backend_declares_uniform_batch_support():
@pytest.mark.parametrize("num_heads", [8, 16, 32, 64, 128])
def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads):
"""Aiter init passes real MTP qlen/dtypes to rocm_aiter_mla.py:349-355.
@pytest.mark.parametrize(
"spec_method, parallel_drafting",
[
("deepseek_mtp", False),
# A drafter that is not one of the historically recognized MTP methods,
# and a parallel one, so the threshold is 1 + 2 * num_spec rather than
# 1 + num_spec. Sizing the metadata off a method name instead leaves
# these at qlen=1 while the router still admits the full range.
("dspark", True),
("eagle", False),
],
)
def test_mtp_builder_init_sizes_native_fp8_metadata(
monkeypatch, num_heads, spec_method, parallel_drafting
):
"""Aiter init sizes the metadata for every query length decode can be handed.
Sweeping num_heads asserts the max(16, num_heads) clamp is what sizes the
metadata, covering the fp8 nhead=32 (TP4) fold path.
@@ -143,6 +159,13 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads):
def init_common_builder(self, *args, **kwargs):
self.num_heads = num_heads
# Mirror what _init_reorder_batch_threshold would have left behind: the
# metadata is sized from the routing threshold, so a stub that skips it
# would size for qlen=1 and hide the very mismatch this test covers.
spec = config.speculative_config
self.reorder_batch_threshold = (
1 + (2 if spec.parallel_drafting else 1) * spec.num_speculative_tokens
)
monkeypatch.setitem(
sys.modules,
@@ -161,8 +184,9 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads):
config = SimpleNamespace(
speculative_config=SimpleNamespace(
method="deepseek_mtp",
method=spec_method,
num_speculative_tokens=3,
parallel_drafting=parallel_drafting,
),
parallel_config=SimpleNamespace(tensor_parallel_size=8),
model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16),
@@ -182,7 +206,7 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads):
assert info_calls == [
{
"max_batch_size": config.scheduler_config.max_num_seqs,
"max_qo_len": config.speculative_config.num_speculative_tokens + 1,
"max_qo_len": builder.reorder_batch_threshold,
"num_attention_heads": max(16, num_heads),
"q_dtype": dtypes.fp8,
"kv_dtype": dtypes.fp8,
@@ -352,26 +376,46 @@ def test_decode_expands_kernel_block_page_indices(monkeypatch):
@pytest.mark.parametrize(
"mtp_decode_qlen, qo_len, num_heads, expect_persistent",
"mtp_decode_qlen, qo_len, num_heads, kv_cache_dtype, expect_persistent",
[
(1, 1, 16, True), # non-MTP decode
(4, 2, 16, True), # MTP deployment, in-range step
(4, 4, 16, True), # MTP deployment, full-qlen verification step
(2, 4, 16, False), # step demand exceeds provisioned K -> fallback
(1, 1, 8, False), # small head count -> Gluon decode owns qlen==1
(4, 4, 8, False), # small head count -> Gluon flatten owns qlen>1
(1, 1, 16, "auto", True), # non-MTP decode
(4, 2, 16, "auto", True), # MTP deployment, in-range step
(4, 4, 16, "auto", True), # MTP deployment, full-qlen verification step
(2, 4, 16, "auto", False), # step demand exceeds provisioned K -> fallback
(1, 1, 8, "auto", False), # divisor head count -> Gluon decode owns qlen==1
(4, 4, 8, "auto", False), # small head count -> Gluon flatten owns qlen>1
# A non-divisor head count is padded to 16 and runs the asm decode, so it
# needs the schedule even though num_heads < 16. Its verify still
# flattens onto Gluon though -- divisibility only steers the qlen==1
# decode, and bf16 has no small-head multi-token asm kernel either way.
(1, 1, 12, "auto", True),
(8, 8, 12, "auto", False),
# An fp8 cache never reaches Gluon, at any head count or qlen, so it
# always needs the schedule. (1, 1, 8) is the case the head-count gate
# got wrong: divisor head count, so it read as Gluon-owned.
(1, 1, 8, "fp8", True),
(1, 1, 12, "fp8", True),
(1, 1, 16, "fp8", True),
# DSpark verify shape. The fp8 fold path rejects non-persistent outright
# once qlen > 4, so these must not fall through.
(8, 8, 8, "fp8", True),
(8, 8, 12, "fp8", True),
],
)
def test_persistent_metadata_gate(
monkeypatch, mtp_decode_qlen, qo_len, num_heads, expect_persistent
monkeypatch, mtp_decode_qlen, qo_len, num_heads, kv_cache_dtype, expect_persistent
):
"""Persistent metadata is passed iff num_heads >= 16 and 1 <= max_qo_len <= K.
"""Persistent metadata is passed iff the asm decode runs and 1 <= qlen <= K.
K = _mtp_decode_qlen sizes the metadata buffers at init; a decode step gets
the pre-built schedule only when its qlen fits those buffers, otherwise it
falls back to the kernel computing its own. qlen==1 (non-MTP) must stay
in-range -- dropping it is the regression this guards. Fewer than 16 heads
is served by the Gluon decode paths, which never read this schedule.
in-range -- dropping it is the regression this guards.
Only the Gluon paths ignore the schedule, so the gate follows the routing
predicates rather than the raw head count. Reading `num_heads >= 16` instead
denies the schedule to two sets of shapes that do run the asm kernels: a
non-divisor head count padded up to 16, and any fp8 cache below 16 heads.
"""
get_mla_metadata_v1 = mock.MagicMock()
monkeypatch.setitem(
@@ -382,6 +426,12 @@ def test_persistent_metadata_gate(
monkeypatch.setattr(
rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
)
# The gate follows the routing now, so the small-head bf16 rows above only
# hold where a Gluon build exists. Pin the arch they describe: on gfx942
# those shapes run the asm decode and do get the schedule, which is a
# different assertion, made arch-free in the fp8 routing test.
monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True)
monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto")
# Uniform, non-CUDA-graph batch: every request has exactly qo_len tokens, so
# num_decode_tokens == sum(qo_len) and no dummy-row padding kicks in.
@@ -390,7 +440,11 @@ def test_persistent_metadata_gate(
0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32
)
metadata = AiterMLAMetadataBuilder._build_decode(
_builder(mtp_decode_qlen=mtp_decode_qlen, num_heads=num_heads),
_builder(
mtp_decode_qlen=mtp_decode_qlen,
num_heads=num_heads,
kv_cache_dtype=kv_cache_dtype,
),
block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
seq_lens_device=torch.tensor([8, 8], dtype=torch.int32),
max_seq_len=8,
@@ -406,3 +460,68 @@ def test_persistent_metadata_gate(
if expect_persistent:
assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == qo_len
assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == qo_len
@pytest.mark.parametrize(
"qo_len, num_heads, kv_cache_dtype, expect_persistent",
[
# A padded rank keeps the schedule wherever a persistent kernel exists:
# bf16 up to qseqlen 4, fp8 at every qlen (its fold requires it).
(1, 12, "auto", True),
(4, 12, "auto", True),
(1, 12, "fp8", True),
(8, 12, "fp8", True),
# Past qseqlen 4 bf16 has only the non-persistent entry, and the fold
# that would reach a persistent one is gfx950-only, so the padded rank
# must fall back rather than ask for a kernel that is not built.
(8, 12, "auto", False),
# A native 16-head rank is unchanged: it took the schedule before this
# backend keyed the gate off the routing, and still does.
(8, 16, "auto", True),
],
)
def test_persistent_metadata_gate_without_gluon_build(
monkeypatch, qo_len, num_heads, kv_cache_dtype, expect_persistent
):
"""The gate on an arch with no Gluon build, i.e. gfx942.
Both Gluon predicates read False there, so a small-head shape reaches the
padded asm decode instead of being flattened. The schedule then has to
follow what the asm dispatch actually ships: bf16 has no gqa=16 persistent
kernel above qseqlen 4, and the q-row fold that reaches one on gfx950 is
guarded on that arch, so asking for the schedule there selects a kernel
that does not exist.
"""
get_mla_metadata_v1 = mock.MagicMock()
monkeypatch.setitem(
sys.modules,
"aiter",
SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
)
monkeypatch.setattr(
rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
)
monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: False)
monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto")
num_reqs = 2
query_start_loc = torch.arange(
0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32
)
metadata = AiterMLAMetadataBuilder._build_decode(
_builder(
mtp_decode_qlen=qo_len,
num_heads=num_heads,
kv_cache_dtype=kv_cache_dtype,
),
block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
seq_lens_device=torch.tensor([8, 8], dtype=torch.int32),
max_seq_len=8,
query_start_loc_cpu=query_start_loc,
query_start_loc_device=query_start_loc,
num_decode_tokens=num_reqs * qo_len,
dcp_tot_seq_lens_device=None,
)
assert metadata.has_persistent_metadata is expect_persistent
assert get_mla_metadata_v1.called is expect_persistent
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+64 -35
View File
@@ -26,7 +26,7 @@ from vllm.v1.attention.backend import (
CommonAttentionMetadata,
MultipleOf,
)
from vllm.v1.kv_cache_interface import AttentionSpec
from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
logger = init_logger(__name__)
@@ -269,17 +269,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
self.compilation_config = vllm_config.compilation_config
self.decode_attn_out_dtype = vllm_config.model_config.dtype
# MTP/deepseek_mtp verification runs decode with qlen = num_spec + 1;
# any other config (including no spec) stays at single-token decode.
speculative_config = vllm_config.speculative_config
if (
speculative_config is not None
and speculative_config.method in ("mtp", "deepseek_mtp")
and speculative_config.num_speculative_tokens is not None
):
self._mtp_decode_qlen = int(speculative_config.num_speculative_tokens) + 1
else:
self._mtp_decode_qlen = 1
# reorder_batch_threshold is the largest query length decode can be
# handed, and already accounts for the drafting scheme. A method-name
# whitelist sizes unlisted drafters for qlen=1, which closes the
# persistent gate below and makes aiter raise a KeyError mid-run.
self._mtp_decode_qlen = self.reorder_batch_threshold or 1
# Store the kernel block size from the spec. When kernel_block_size=1
# (no spec-dec), behavior is identical to the original. When > 1
@@ -327,6 +321,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
torch.float16: dtypes.fp16,
torch.bfloat16: dtypes.bf16,
}[kv_cache_spec.dtype]
# _build_decode needs the cache dtype to pick the decode kernel; keep
# the normalized string instead of dropping it at the end of __init__.
self._kv_cache_dtype_str = kv_cache_dtype_str
# MLAAttention quantizes decode Q to FP8 before calling this backend
# whenever the KV cache is FP8 and supports_quant_query_input is true.
q_dtype = (
@@ -624,7 +621,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
]
)
use_gluon_decode = AiterMLAHelper.use_gluon_decode(
self.num_heads, int(max_qo_len)
self.num_heads, int(max_qo_len), self._kv_cache_dtype_str
)
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
@@ -693,13 +690,26 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
else:
qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]
# Pass persistent metadata for every uniform decode we sized buffers for
# (normal qlen==1 through MTP verification qlen==K): the fp8 nhead=32 fold
# path breaks without it. qlen>K falls back to kernel-internal metadata.
# Small-head (<16) decode takes the Gluon paths and never consumes it.
# Only the asm decode consumes the schedule, so gate on the routing
# rather than on num_heads >= 16, which denies it to a padded rank
# running the same asm kernels. The two predicates are disjoint --
# decode is qlen==1, verify is qlen>1 -- and cover both Gluon entries.
has_persistent_metadata = False
use_persistent_metadata = (
self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
not AiterMLAHelper.use_gluon_decode(
self.num_heads, max_qo_len, self._kv_cache_dtype_str
)
and not AiterMLAHelper.use_gluon_verify(
self.num_heads, max_qo_len, self._kv_cache_dtype_str
)
# A padded rank has no bf16 persistent kernel past qlen 4 where the
# gfx950 fold is absent; the non-persistent entry covers it. fp8
# keeps the schedule -- its fold rejects non-persistent outright.
and (
self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN
or is_quantized_kv_cache(self._kv_cache_dtype_str)
)
and max_qo_len >= 1
and max_qo_len <= self._mtp_decode_qlen
)
@@ -833,6 +843,10 @@ class AiterMLAHelper:
"""
_AITER_MIN_MLA_HEADS: Final = 16
# Largest qlen the padded gqa=16 asm decode has a bf16 persistent kernel
# for. Above it only the non-persistent qseqlen=8 entry exists, and the
# fold that reaches a persistent one is gfx950-only.
_ASM_PADDED_MAX_PS_QLEN: Final = 4
_AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = ()
@staticmethod
@@ -888,17 +902,19 @@ class AiterMLAHelper:
return o[:, :num_heads, :]
@staticmethod
def use_gluon_decode(num_heads: int, max_qo_len: int) -> bool:
# 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.
def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
# Small-head (<16) single-token decode takes either the Gluon kernel or
# the padded asm persistent decode, selected by
# VLLM_ROCM_AITER_MLA_ASM_PADDING and the arch (Gluon is gfx950 only).
m = AiterMLAHelper._AITER_MIN_MLA_HEADS
if num_heads >= m or max_qo_len != 1:
return False
# Gluon's only fp8-KV regime, bh16bn128, is a bf16-query kernel with a
# hardcoded scale that asserts batch_size == 1, so it cannot serve a
# decode batch. Checked before the mode knob: an explicit "gluon"
# request under fp8 would assert immediately.
if is_quantized_kv_cache(kv_cache_dtype):
return False
mode = _aiter_mla_small_head_mode()
if mode == "asm":
return False
@@ -907,6 +923,23 @@ class AiterMLAHelper:
return gluon_supported
return m % num_heads == 0 and gluon_supported
@staticmethod
def use_gluon_verify(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
"""Whether a small-head multi-token verify is flattened onto Gluon.
bf16 has no gqa<16, qseqlen>1 asm kernel, so the verify is flattened
into per-token Gluon decodes. fp8 has one via the q-row fold and must
not come here: the flatten hands Gluon the batch size its fp8 regime
asserts against. A predicate rather than inline in forward_mqa so the
builder sees the same answer the impl acts on.
"""
if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= 1:
return False
if is_quantized_kv_cache(kv_cache_dtype):
return False
# Same arch and mode gating as use_gluon_decode.
return _aiter_mla_small_head_mode() != "asm" and _gluon_mla_decode_supported()
class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
def __init__(
@@ -1191,15 +1224,11 @@ 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()
# Arch, mode and dtype gating all live in use_gluon_verify, so that the
# builder -- which has to know whether the asm decode will run -- sees
# the same answer as this branch.
if AiterMLAHelper.use_gluon_verify(
self.num_heads, int(decode.max_qo_len), self.kv_cache_dtype
):
qlen = int(decode.max_qo_len)
if type(q) is tuple: