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

[Bugfix][ROCm] AITER MLA: size MTP verification decode metadata for real qlen/dtype (#45227)

Browse Source 
Signed-off-by: Edwin Lim <edwin.lim@mangoboost.io>
Signed-off-by: Chaemin Lim <chaemin.lim@mangoboost.io>
Signed-off-by: QinPR <1905873179@qq.com>
Co-authored-by: Edwin Lim <edwin.lim@mangoboost.io>
Co-authored-by: limeward <32970461+edwinlim0919@users.noreply.github.com>
Co-authored-by: QinPR <1905873179@qq.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
chaeminlim-mb
2026-07-30 20:50:34 -05:00
committed by GitHub
co-authored by Edwin Lim limeward QinPR Claude Opus 4.8
parent 4f1da84eb5
commit f1899b2ffd
2 changed files with 554 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+408
View File
@@ -0,0 +1,408 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys
from types import SimpleNamespace
from unittest import mock
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 ( # noqa: E402
AiterMLAMetadataBuilder,
)
class _NoOpTritonKernel:
def __getitem__(self, grid):
self.grid = grid
return self
def __call__(self, *args, **kwargs):
pass
class _ExpandPageIndicesKernel:
def __getitem__(self, grid):
self.grid = grid
return self
def __call__(
self,
page_indices,
block_table_tensor,
stride,
paged_kv_indptr,
seq_lens_for_kernel,
*,
KERNEL_BLOCK_SIZE,
BLOCK_SIZE,
):
self.kernel_block_size = KERNEL_BLOCK_SIZE
for req_idx in range(self.grid[0]):
out_start = int(paged_kv_indptr[req_idx].item())
seq_len = int(seq_lens_for_kernel[req_idx].item())
for token_idx in range(seq_len):
block_id = int(
block_table_tensor[req_idx, token_idx // KERNEL_BLOCK_SIZE].item()
)
page_indices[out_start + token_idx] = (
block_id * KERNEL_BLOCK_SIZE + token_idx % KERNEL_BLOCK_SIZE
)
def _builder(
*,
mtp_decode_qlen: int,
has_full_cudagraphs: bool = False,
kernel_block_size: int = 1,
max_decode_rows: int = 32,
num_heads: int = 16,
):
return SimpleNamespace(
device=torch.device("cpu"),
num_heads=num_heads,
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),
qo_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32),
compilation_config=SimpleNamespace(
cudagraph_mode=SimpleNamespace(
has_full_cudagraphs=lambda: has_full_cudagraphs
)
),
_mtp_decode_qlen=mtp_decode_qlen,
_uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len),
_use_persistent_metadata=False,
kernel_block_size=kernel_block_size,
_num_attention_heads=16,
_mla_work_meta_data=torch.empty(1, dtype=torch.int32),
_mla_work_info_set=torch.empty(1, dtype=torch.int32),
_mla_work_indptr=torch.empty(1, dtype=torch.int32),
_mla_reduce_indptr=torch.empty(1, dtype=torch.int32),
_mla_reduce_final_map=torch.empty(1, dtype=torch.int32),
_mla_reduce_partial_map=torch.empty(1, dtype=torch.int32),
_mla_q_dtype=torch.bfloat16,
_mla_kv_dtype=torch.bfloat16,
decode_attn_out_dtype=torch.bfloat16,
)
def test_backend_declares_uniform_batch_support():
# UNIFORM/UNIFORM_BATCH is unconditional: MTP yields uniform qlen>1 and
# non-MTP yields qlen==1, both uniform batches.
assert (
AiterMLAMetadataBuilder.query_len_support
== rocm_aiter_mla.QueryLenSupport.UNIFORM
)
assert (
AiterMLAMetadataBuilder._cudagraph_support
== rocm_aiter_mla.AttentionCGSupport.UNIFORM_BATCH
)
@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.
Sweeping num_heads asserts the max(16, num_heads) clamp is what sizes the
metadata, covering the fp8 nhead=32 (TP4) fold path.
"""
dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16")
info_calls = []
def get_mla_metadata_info_v1(
max_batch_size,
max_qo_len,
num_attention_heads,
q_dtype,
kv_dtype,
*,
is_sparse,
fast_mode,
):
info_calls.append(
{
"max_batch_size": max_batch_size,
"max_qo_len": max_qo_len,
"num_attention_heads": num_attention_heads,
"q_dtype": q_dtype,
"kv_dtype": kv_dtype,
"is_sparse": is_sparse,
"fast_mode": fast_mode,
}
)
return tuple((1, torch.int32) for _ in range(6))
def init_common_builder(self, *args, **kwargs):
self.num_heads = num_heads
monkeypatch.setitem(
sys.modules,
"aiter",
SimpleNamespace(
dtypes=dtypes,
get_mla_metadata_info_v1=get_mla_metadata_info_v1,
),
)
monkeypatch.setattr(
rocm_aiter_mla.MLACommonMetadataBuilder,
"__init__",
init_common_builder,
)
monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False)
config = SimpleNamespace(
speculative_config=SimpleNamespace(
method="deepseek_mtp",
num_speculative_tokens=3,
),
parallel_config=SimpleNamespace(tensor_parallel_size=8),
model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16),
scheduler_config=SimpleNamespace(max_num_seqs=2),
cache_config=SimpleNamespace(cache_dtype="fp8_e4m3"),
compilation_config=SimpleNamespace(
cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False)
),
)
builder = AiterMLAMetadataBuilder(
kv_cache_spec=SimpleNamespace(block_size=1, dtype=torch.bfloat16),
layer_names=["layer.0"],
vllm_config=config,
device=torch.device("cpu"),
)
assert info_calls == [
{
"max_batch_size": config.scheduler_config.max_num_seqs,
"max_qo_len": config.speculative_config.num_speculative_tokens + 1,
"num_attention_heads": max(16, num_heads),
"q_dtype": dtypes.fp8,
"kv_dtype": dtypes.fp8,
"is_sparse": False,
"fast_mode": True,
}
]
assert builder._mla_q_dtype == dtypes.fp8
assert builder._mla_kv_dtype == dtypes.fp8
def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch):
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()
)
metadata = AiterMLAMetadataBuilder._build_decode(
_builder(mtp_decode_qlen=4),
block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
seq_lens_device=torch.tensor([7, 5], dtype=torch.int32),
max_seq_len=7,
query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32),
query_start_loc_device=torch.tensor([0, 4, 8], dtype=torch.int32),
num_decode_tokens=8,
dcp_tot_seq_lens_device=None,
)
assert metadata.max_qo_len == 4
assert torch.equal(metadata.seq_lens, torch.tensor([7, 5], dtype=torch.int32))
assert torch.equal(metadata.qo_indptr, torch.tensor([0, 4, 8], dtype=torch.int32))
assert metadata.has_persistent_metadata
assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == 4
assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == 4
def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr(
monkeypatch,
):
"""Full-CG zero-qo rows follow rocm_aiter_mla.py:608-657,717-759."""
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()
)
mtp_qlen = 4
seq_lens = torch.tensor([7, 0], dtype=torch.int32)
qo_lens = torch.tensor([mtp_qlen, 0], dtype=torch.int32)
expected_seq_lens = torch.where(qo_lens > 0, seq_lens, mtp_qlen)
expected_paged_kv_indptr = torch.cat(
[
torch.zeros(1, dtype=torch.int32),
expected_seq_lens.cumsum(dim=0, dtype=torch.int32),
]
)
expected_qo_indptr = torch.arange(
0,
(seq_lens.numel() + 1) * mtp_qlen,
step=mtp_qlen,
dtype=torch.int32,
)
builder = _builder(
mtp_decode_qlen=mtp_qlen,
has_full_cudagraphs=True,
max_decode_rows=4,
)
metadata = AiterMLAMetadataBuilder._build_decode(
builder,
block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
seq_lens_device=seq_lens,
max_seq_len=int(seq_lens.max().item()),
query_start_loc_cpu=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32),
query_start_loc_device=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32),
num_decode_tokens=seq_lens.numel() * mtp_qlen,
dcp_tot_seq_lens_device=None,
)
assert metadata.max_qo_len == mtp_qlen
assert torch.equal(metadata.seq_lens, expected_seq_lens)
assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr)
assert torch.equal(metadata.qo_indptr, expected_qo_indptr)
assert torch.all(
builder.paged_kv_indptr[expected_paged_kv_indptr.numel() :]
== expected_paged_kv_indptr[-1]
)
assert torch.all(
builder.qo_indptr[expected_qo_indptr.numel() :] == expected_qo_indptr[-1]
)
assert metadata.has_persistent_metadata
assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == mtp_qlen
assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == mtp_qlen
def test_decode_expands_kernel_block_page_indices(monkeypatch):
"""kernel_block_size>1 expands b -> b*K+offset at rocm_aiter_mla.py:696-704."""
expand_kernel = _ExpandPageIndicesKernel()
monkeypatch.setattr(rocm_aiter_mla, "_expand_page_indices_kernel", expand_kernel)
# qlen==1 now takes the persistent-metadata path, which imports
# get_mla_metadata_v1 from aiter; mock it so this test needs no real kernel.
monkeypatch.setitem(
sys.modules,
"aiter",
SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()),
)
kernel_block_size = 2
block_table = torch.tensor(
[
[10, 11, 99],
[20, 21, 22],
],
dtype=torch.int32,
)
seq_lens = torch.tensor([3, 5], dtype=torch.int32)
expected_paged_kv_indptr = torch.cat(
[
torch.zeros(1, dtype=torch.int32),
seq_lens.cumsum(dim=0, dtype=torch.int32),
]
)
expected_indices = torch.tensor(
[
int(block_table[req_idx, token_idx // kernel_block_size].item())
* kernel_block_size
+ token_idx % kernel_block_size
for req_idx, seq_len in enumerate(seq_lens.tolist())
for token_idx in range(seq_len)
],
dtype=torch.int32,
)
metadata = AiterMLAMetadataBuilder._build_decode(
_builder(
mtp_decode_qlen=1,
kernel_block_size=kernel_block_size,
),
block_table_tensor=block_table,
seq_lens_device=seq_lens,
max_seq_len=int(seq_lens.max().item()),
query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.int32),
query_start_loc_device=torch.tensor([0, 1, 2], dtype=torch.int32),
num_decode_tokens=seq_lens.numel(),
dcp_tot_seq_lens_device=None,
)
assert metadata.max_qo_len == 1
assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr)
assert torch.equal(
metadata.paged_kv_indices[: expected_indices.numel()],
expected_indices,
)
assert expand_kernel.grid == (seq_lens.numel(),)
assert expand_kernel.kernel_block_size == kernel_block_size
@pytest.mark.parametrize(
"mtp_decode_qlen, qo_len, num_heads, 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
],
)
def test_persistent_metadata_gate(
monkeypatch, mtp_decode_qlen, qo_len, num_heads, expect_persistent
):
"""Persistent metadata is passed iff num_heads >= 16 and 1 <= max_qo_len <= 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.
"""
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()
)
# 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.
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=mtp_decode_qlen, num_heads=num_heads),
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.max_qo_len == qo_len
assert metadata.has_persistent_metadata is expect_persistent
assert get_mla_metadata_v1.called is expect_persistent
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
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+146 -35
View File
@@ -132,7 +132,7 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
min_kv_seq_len: int = 1
# Small-head decode uses Gluon (avoids padding to 16).
use_gluon_decode: bool = False
# Whether persistent MLA metadata was computed (only for qseqlen=1)
# Whether persistent MLA metadata was computed
has_persistent_metadata: bool = False
@@ -171,6 +171,46 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
@staticmethod
def _uniform_padded_mtp_qo_len(
qo_len: torch.Tensor,
max_qo_len: int,
num_decode_tokens: int,
) -> int:
num_reqs = qo_len.numel()
if num_reqs == 0 or num_decode_tokens <= 0:
return 0
# Full-CG pads q to a captured token count while leaving
# query_start_loc flat for dummy requests. Only synthesize dummy rows
# when every padded request maps to the same qlen and the q buffer has
# exactly that many rows.
if num_decode_tokens <= int(qo_len.sum().item()):
return 0
if num_decode_tokens % num_reqs != 0:
return 0
uniform_qo_len = num_decode_tokens // num_reqs
if uniform_qo_len <= 1:
return 0
positive_qo_len = qo_len[qo_len > 0]
if positive_qo_len.numel() == qo_len.numel():
return 0
if positive_qo_len.numel() > 0:
if max_qo_len != uniform_qo_len:
return 0
if not torch.all(positive_qo_len == uniform_qo_len):
return 0
zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten()
if zero_positions.numel() > 0:
first_zero = int(zero_positions[0].item())
if torch.any(qo_len[first_zero:] > 0):
return 0
return uniform_qo_len
def __init__(
self,
kv_cache_spec: AttentionSpec,
@@ -185,6 +225,18 @@ 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
# Store the kernel block size from the spec. When kernel_block_size=1
# (no spec-dec), behavior is identical to the original. When > 1
# (e.g. 16 with Eagle3), we expand block-level indices into per-token
@@ -222,14 +274,20 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
# make sure get_mla_metadata_info_v1 / get_mla_metadata_v1 are consistent
# with the actual tensor shape passed to mla_decode_fwd.
self._num_attention_heads = max(16, self.num_heads)
q_dtype = self.decode_attn_out_dtype
kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
kv_cache_dtype_str = "fp8"
q_dtype = dtypes.fp8
kv_dtype = dtypes.fp8
else:
kv_cache_dtype_str = "bf16"
kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16)
kv_dtype = {
torch.float16: dtypes.fp16,
torch.bfloat16: dtypes.bf16,
}[kv_cache_spec.dtype]
# 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 = (
dtypes.fp8 if kv_cache_dtype_str == "fp8" else self.decode_attn_out_dtype
)
# Persist for get_mla_metadata_v1 (decode build): omitting these causes
# wrong split/reduce metadata for the gfx950 fp8 nhead=32 fold path.
self._mla_q_dtype = q_dtype
@@ -243,7 +301,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
(reduce_partial_map_size, reduce_partial_map_type),
) = get_mla_metadata_info_v1(
max_num_reqs,
1,
self._mtp_decode_qlen,
self._num_attention_heads,
q_dtype,
kv_dtype,
@@ -488,20 +546,39 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
) -> AiterMLADecodeMetadata:
device = self.device
num_reqs = seq_lens_device.size(0)
qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
max_qo_len = qo_len.max().item()
padded_mtp_qo_len = self._uniform_padded_mtp_qo_len(
qo_len, max_qo_len, num_decode_tokens
)
if padded_mtp_qo_len > 0:
max_qo_len = padded_mtp_qo_len
pad_uniform_mtp = padded_mtp_qo_len > 0
seq_lens_for_kernel = seq_lens_device
num_kernel_reqs = num_reqs
if pad_uniform_mtp:
qo_lens_device = (
query_start_loc_device[1 : num_reqs + 1]
- query_start_loc_device[:num_reqs]
).to(torch.int32)
seq_lens_for_kernel = torch.where(
qo_lens_device > 0,
seq_lens_for_kernel,
seq_lens_for_kernel.new_full((), max_qo_len),
)
# The aiter kernel always operates with page_size=1 (the wrapper
# flattens kv_buffer). last_page_len is always 1.
paged_kv_last_page_len = self.paged_kv_last_page_len[:num_reqs]
paged_kv_last_page_len = self.paged_kv_last_page_len[:num_kernel_reqs]
# indptr: cumsum of seq_lens (one page per token in the flat view)
paged_kv_indptr = torch.cat(
[
torch.zeros(1, dtype=seq_lens_device.dtype, device=device),
seq_lens_device.cumsum(dim=0, dtype=torch.int32),
torch.zeros(1, dtype=torch.int32, device=device),
seq_lens_for_kernel.cumsum(dim=0, dtype=torch.int32),
]
)
qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
max_qo_len = qo_len.max().item()
use_gluon_decode = AiterMLAHelper.use_gluon_decode(
self.num_heads, int(max_qo_len)
)
@@ -519,44 +596,75 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
block_table_tensor,
block_table_tensor.stride(0),
paged_kv_indptr,
seq_lens_device,
seq_lens_for_kernel,
KERNEL_BLOCK_SIZE=self.kernel_block_size,
BLOCK_SIZE=1024,
)
paged_kv_indices = self.paged_kv_indices
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
self.paged_kv_indptr[: 1 + num_reqs].copy_(
self.paged_kv_indptr[: 1 + num_kernel_reqs].copy_(
paged_kv_indptr, non_blocking=True
)
self.paged_kv_indptr[1 + num_reqs :].fill_(paged_kv_indptr[-1])
paged_kv_indptr = self.paged_kv_indptr[: 1 + num_reqs]
self.paged_kv_indptr[1 + num_kernel_reqs :].fill_(paged_kv_indptr[-1])
paged_kv_indptr = self.paged_kv_indptr[: 1 + num_kernel_reqs]
# paged_kv_last_page_len already uses the pre-initialized buffer slice
# (set above), so no copy needed - buffer is always 1s.
self.qo_indptr[: 1 + num_reqs].copy_(
query_start_loc_device, non_blocking=True
if pad_uniform_mtp:
qo_indptr_src = torch.arange(
0,
(num_kernel_reqs + 1) * max_qo_len,
step=max_qo_len,
dtype=torch.int32,
device=device,
)
else:
qo_indptr_src = query_start_loc_device[: 1 + num_kernel_reqs]
self.qo_indptr[: 1 + num_kernel_reqs].copy_(
qo_indptr_src, non_blocking=True
)
self.qo_indptr[1 + num_reqs :] = query_start_loc_device[-1]
qo_indptr = self.qo_indptr[: 1 + num_reqs]
self.qo_indptr[1 + num_kernel_reqs :] = qo_indptr_src[-1]
qo_indptr = self.qo_indptr[: 1 + num_kernel_reqs]
else:
qo_indptr = torch.arange(
0, num_reqs + 1, step=1, dtype=torch.int32, device=device
)
if max_qo_len == 1:
qo_indptr = torch.arange(
0,
num_kernel_reqs + 1,
step=1,
dtype=torch.int32,
device=device,
)
else:
if pad_uniform_mtp:
qo_indptr = torch.arange(
0,
(num_kernel_reqs + 1) * max_qo_len,
step=max_qo_len,
dtype=torch.int32,
device=device,
)
else:
qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]
# The aiter MLA ASM kernel only supports qseqlen=1 (single-token
# decode). With speculative decoding, the verification step has
# qseqlen > 1 (e.g. 8 for spec7). get_mla_metadata_v1 calls
# get_heuristic_kernel_mla which fails for qseqlen > 1.
# We track whether persistent metadata was successfully computed
# so forward_mqa can skip passing it (falling back to the kernel
# computing its own metadata internally, like v0.18.0).
# 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.
has_persistent_metadata = False
if max_qo_len == 1 and not use_gluon_decode:
use_persistent_metadata = (
self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
and max_qo_len >= 1
and max_qo_len <= self._mtp_decode_qlen
)
if use_persistent_metadata:
from aiter import get_mla_metadata_v1
uni_qo_len = (
max_qo_len if pad_uniform_mtp or torch.all(qo_len == max_qo_len) else -1
)
get_mla_metadata_v1(
qo_indptr,
paged_kv_indptr,
@@ -573,7 +681,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
page_size=1,
kv_granularity=16,
max_seqlen_qo=max_qo_len,
uni_seqlen_qo=max_qo_len,
uni_seqlen_qo=uni_qo_len,
fast_mode=True,
dtype_q=self._mla_q_dtype,
dtype_kv=self._mla_kv_dtype,
@@ -582,7 +690,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
attn_metadata = AiterMLADecodeMetadata(
block_table=block_table_tensor,
seq_lens=seq_lens_device,
seq_lens=seq_lens_for_kernel,
paged_kv_indptr=paged_kv_indptr,
paged_kv_indices=paged_kv_indices,
paged_kv_last_page_len=paged_kv_last_page_len,
@@ -1076,13 +1184,16 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
dtype=attn_metadata.decode.attn_out_dtype,
device=q.device,
)
if decode.max_qo_len > 1 and not decode.has_persistent_metadata:
# MTP verification can call the AITER MLA decode kernel with
# qlen > 1. If that path is running without persistent metadata,
# zero-fill so unwritten lanes cannot leak into logits.
o.zero_()
kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
# Build kwargs for mla_decode_fwd. Pass persistent metadata only
# when it was successfully computed (qseqlen=1 decode steps).
# For multi-token verification steps (spec-dec), the kernel falls
# back to computing metadata internally.
# when it was successfully computed.
mla_kwargs = dict(
q_scale=layer._q_scale,
kv_scale=layer._k_scale,