A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 10:00 UTC
[ROCm][MTP] Support MTP for AITER MLA backend (#28624)
Signed-off-by: ganyi <ygan@amd.com>
Pleaplusone committed
Dec 16, 2025 at 22:10 UTC
9dbbc59b151163b7c08a2fde362b6aaedd9bc343
1 file changed
+15
-9
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+15
-9
@@ -15,6 +15,7 @@ from vllm.v1.attention.backends.mla.common import (
15
MLACommonImpl,
16
MLACommonMetadata,
17
MLACommonMetadataBuilder,
18
+ QueryLenSupport,
19
)
20
from vllm.v1.attention.backends.utils import AttentionCGSupport
21
from vllm.v1.kv_cache_interface import AttentionSpec
@@ -51,6 +52,8 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
52
qo_indptr: torch.Tensor | None = None
53
# The dtype of MLA out tensor
54
attn_out_dtype: torch.dtype = torch.bfloat16
55
+ # The max query output length: int
56
+ max_qo_len: int | None = None
57
58
59
class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):
@@ -60,9 +63,8 @@ class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):
63
class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
64
# TODO(luka, lucas): audit this as part of:
65
# https://github.com/vllm-project/vllm/issues/22945
63
- _cudagraph_support: ClassVar[AttentionCGSupport] = (
64
- AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
65
- )
66
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
67
+ query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
68
69
def __init__(
70
self,
@@ -97,8 +99,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
99
max_num_reqs, dtype=torch.int32, device=device
100
)
101
100
- self.qo_indptr = torch.arange(
101
- 0, max_num_reqs + 1, dtype=torch.int32, device=device
102
+ self.qo_indptr = torch.zeros(
103
+ max_num_reqs + 1, dtype=torch.int32, device=device
104
)
105
106
def _build_decode(
@@ -128,6 +130,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
130
seq_lens_device.cumsum(dim=0, dtype=torch.int32),
131
]
132
)
133
+ qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
134
+ max_qo_len = qo_len.max().item()
135
136
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
137
num_actual_pages = paged_kv_indices.size(0)
@@ -150,6 +154,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
154
self.paged_kv_last_page_len[num_reqs:].fill_(1)
155
paged_kv_last_page_len = self.paged_kv_last_page_len[:num_reqs]
156
157
+ self.qo_indptr[: 1 + num_reqs].copy_(
158
+ query_start_loc_device, non_blocking=True
159
+ )
160
+ self.qo_indptr[1 + num_reqs :] = query_start_loc_device[-1]
161
qo_indptr = self.qo_indptr[: 1 + num_reqs]
162
163
else:
@@ -165,6 +173,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
173
paged_kv_last_page_len=paged_kv_last_page_len,
174
qo_indptr=qo_indptr,
175
dcp_tot_seq_lens=dcp_tot_seq_lens_device,
176
+ max_qo_len=max_qo_len,
177
attn_out_dtype=self.decode_attn_out_dtype,
178
)
179
@@ -255,16 +264,13 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
264
265
kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
266
258
- # max_seqlen_qo must be 1 except for MTP
259
- # TODO: Find the best value for MTP
260
- max_seqlen_qo = 1
267
rocm_aiter_ops.mla_decode_fwd(
268
q,
269
kv_buffer,
270
o,
271
self.scale,
272
attn_metadata.decode.qo_indptr,
267
- max_seqlen_qo,
273
+ attn_metadata.decode.max_qo_len,
274
attn_metadata.decode.paged_kv_indptr,
275
attn_metadata.decode.paged_kv_indices,
276
attn_metadata.decode.paged_kv_last_page_len,