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 4, 22:30 UTC
[ROCm][Feature] Enable AITER MLA attention backend to work with Eagle3 speculative decoding on ROCm (#39616)
Signed-off-by: larryli2-amd <larryli2@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
larryli2-amd committed
Apr 20, 2026 at 22:44 UTC
7243e02aa1c6e01ab1d2a50aa7e27ebc9b2c6719
2 files changed
+130
-58
docs/design/attention_backends.md
+1
-1
@@ -215,7 +215,7 @@ configuration.
215
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
216
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
217
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
218
-| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
218
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
219
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
220
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
221
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+129
-57
@@ -44,7 +44,11 @@ class AiterMLABackend(MLACommonBackend):
44
45
@staticmethod
46
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
47
- return [1]
47
+ # The aiter MLA decode kernel always operates with page_size=1
48
+ # internally (the wrapper flattens kv_buffer via .view(-1, 1, 1, H)).
49
+ # We support any kernel_block_size by expanding block-level indices
50
+ # into per-token flat indices in the metadata builder.
51
+ return [MultipleOf(1)]
52
53
@staticmethod
54
def get_name() -> str:
@@ -74,6 +78,8 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
78
attn_out_dtype: torch.dtype = torch.bfloat16
79
# The max query output length: int
80
max_qo_len: int | None = None
81
+ # Whether persistent MLA metadata was computed (only for qseqlen=1)
82
+ has_persistent_metadata: bool = False
83
84
85
@dataclass
@@ -105,7 +111,16 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
111
112
self.compilation_config = vllm_config.compilation_config
113
self.decode_attn_out_dtype = vllm_config.model_config.dtype
108
- # kernel block size is always 1.
114
+
115
+ # Store the kernel block size from the spec. When kernel_block_size=1
116
+ # (no spec-dec), behavior is identical to the original. When > 1
117
+ # (e.g. 16 with Eagle3), we expand block-level indices into per-token
118
+ # flat indices since the aiter kernel always uses page_size=1 internally.
119
+ self.kernel_block_size = kv_cache_spec.block_size
120
+
121
+ # In the flat view (.view(-1,1,1,H)), each token is its own page,
122
+ # so max_num_pages_per_req = max_model_len regardless of
123
+ # kernel_block_size.
124
max_num_pages_per_req = vllm_config.model_config.max_model_len
125
max_num_reqs = vllm_config.scheduler_config.max_num_seqs
126
max_num_pages = max_num_reqs * max_num_pages_per_req
@@ -115,8 +130,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
130
# so we can only use the persistent buffer if a cudagraph is actually
131
# being used.
132
118
- # paged_kv_last_page_len is always 1s (kernel block size is always 1),
119
- # so we create it once and reuse slices in both eager and cudagraph modes.
133
+ # paged_kv_last_page_len is always 1s (the aiter kernel always sees
134
+ # page_size=1 after .view(-1,1,1,H) flattening), so we create it
135
+ # once and reuse slices in both eager and cudagraph modes.
136
self.paged_kv_last_page_len = torch.ones(
137
max_num_reqs, dtype=torch.int32, device=device
138
)
@@ -196,14 +212,14 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
212
num_decode_tokens: int,
213
dcp_tot_seq_lens_device: torch.Tensor | None,
214
) -> AiterMLADecodeMetadata:
199
- # kernel block size is always 1, although the kv block size is not 1.
215
device = self.device
216
num_reqs = seq_lens_device.size(0)
217
203
- # kernel block size is always 1, so each page has exactly 1 token.
204
- # last_page_len is always 1 - just slice the pre-initialized buffer.
218
+ # The aiter kernel always operates with page_size=1 (the wrapper
219
+ # flattens kv_buffer). last_page_len is always 1.
220
paged_kv_last_page_len = self.paged_kv_last_page_len[:num_reqs]
221
222
+ # indptr: cumsum of seq_lens (one page per token in the flat view)
223
paged_kv_indptr = torch.cat(
224
[
225
torch.zeros(1, dtype=seq_lens_device.dtype, device=device),
@@ -215,11 +231,19 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
231
232
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
233
self.paged_kv_indices.fill_(-1)
218
- _copy_page_indices_kernel[(num_reqs,)](
234
+
235
+ # Expand block_table entries into per-token flat indices.
236
+ # When kernel_block_size=1, this degrades to a direct copy (identical
237
+ # to the original _copy_page_indices_kernel).
238
+ # When kernel_block_size=K>1, block_table entry b covering K tokens
239
+ # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
240
+ _expand_page_indices_kernel[(num_reqs,)](
241
self.paged_kv_indices,
242
block_table_tensor,
243
block_table_tensor.stride(0),
244
paged_kv_indptr,
245
+ seq_lens_device,
246
+ KERNEL_BLOCK_SIZE=self.kernel_block_size,
247
BLOCK_SIZE=1024,
248
)
249
paged_kv_indices = self.paged_kv_indices
@@ -245,27 +269,37 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
269
0, num_reqs + 1, step=1, dtype=torch.int32, device=device
270
)
271
248
- from aiter import get_mla_metadata_v1
249
-
250
- get_mla_metadata_v1(
251
- qo_indptr,
252
- paged_kv_indptr,
253
- paged_kv_last_page_len,
254
- self._num_attention_heads,
255
- 1,
256
- True,
257
- self._mla_work_meta_data,
258
- self._mla_work_info_set,
259
- self._mla_work_indptr,
260
- self._mla_reduce_indptr,
261
- self._mla_reduce_final_map,
262
- self._mla_reduce_partial_map,
263
- page_size=1,
264
- kv_granularity=16,
265
- max_seqlen_qo=max_qo_len,
266
- uni_seqlen_qo=max_qo_len,
267
- fast_mode=True,
268
- )
272
+ # The aiter MLA ASM kernel only supports qseqlen=1 (single-token
273
+ # decode). With speculative decoding, the verification step has
274
+ # qseqlen > 1 (e.g. 8 for spec7). get_mla_metadata_v1 calls
275
+ # get_heuristic_kernel_mla which fails for qseqlen > 1.
276
+ # We track whether persistent metadata was successfully computed
277
+ # so forward_mqa can skip passing it (falling back to the kernel
278
+ # computing its own metadata internally, like v0.18.0).
279
+ has_persistent_metadata = False
280
+ if max_qo_len == 1:
281
+ from aiter import get_mla_metadata_v1
282
+
283
+ get_mla_metadata_v1(
284
+ qo_indptr,
285
+ paged_kv_indptr,
286
+ paged_kv_last_page_len,
287
+ self._num_attention_heads,
288
+ 1,
289
+ True,
290
+ self._mla_work_meta_data,
291
+ self._mla_work_info_set,
292
+ self._mla_work_indptr,
293
+ self._mla_reduce_indptr,
294
+ self._mla_reduce_final_map,
295
+ self._mla_reduce_partial_map,
296
+ page_size=1,
297
+ kv_granularity=16,
298
+ max_seqlen_qo=max_qo_len,
299
+ uni_seqlen_qo=max_qo_len,
300
+ fast_mode=True,
301
+ )
302
+ has_persistent_metadata = True
303
304
attn_metadata = AiterMLADecodeMetadata(
305
block_table=block_table_tensor,
@@ -277,6 +311,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
311
dcp_tot_seq_lens=dcp_tot_seq_lens_device,
312
max_qo_len=max_qo_len,
313
attn_out_dtype=self.decode_attn_out_dtype,
314
+ has_persistent_metadata=has_persistent_metadata,
315
)
316
317
return attn_metadata
@@ -290,41 +325,67 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
325
attn_metadata = super().build(
326
common_prefix_len, common_attn_metadata, fast_build
327
)
293
- attn_metadata.work_meta_data = self._mla_work_meta_data
294
- attn_metadata.work_indptr = self._mla_work_indptr
295
- attn_metadata.work_info_set = self._mla_work_info_set
296
- attn_metadata.reduce_indptr = self._mla_reduce_indptr
297
- attn_metadata.reduce_final_map = self._mla_reduce_final_map
298
- attn_metadata.reduce_partial_map = self._mla_reduce_partial_map
328
+ if (
329
+ attn_metadata.decode is not None
330
+ and attn_metadata.decode.has_persistent_metadata
331
+ ):
332
+ attn_metadata.work_meta_data = self._mla_work_meta_data
333
+ attn_metadata.work_indptr = self._mla_work_indptr
334
+ attn_metadata.work_info_set = self._mla_work_info_set
335
+ attn_metadata.reduce_indptr = self._mla_reduce_indptr
336
+ attn_metadata.reduce_final_map = self._mla_reduce_final_map
337
+ attn_metadata.reduce_partial_map = self._mla_reduce_partial_map
338
return attn_metadata
339
340
341
@triton.jit
303
-def _copy_page_indices_kernel(
342
+def _expand_page_indices_kernel(
343
page_indices,
344
block_table,
345
block_table_stride,
307
- cu_num_blocks,
346
+ cu_num_tokens,
347
+ seq_lens,
348
+ KERNEL_BLOCK_SIZE: tl.constexpr,
349
BLOCK_SIZE: tl.constexpr,
350
):
310
- """Copy block table rows into a flat page_indices buffer using indptr.
311
- Avoids blocking boolean mask indexing (tensor[mask]) which has
312
- data-dependent output size and forces sync.
313
- This is the same kernel as introduced in backends/flashinfer.py.
351
+ """Expand block table entries into per-token flat page indices.
352
+
353
+ The aiter MLA kernel always operates with page_size=1 internally
354
+ (kv_buffer is flattened via .view(-1, 1, 1, H)). This kernel converts
355
+ block-level indices from the block table into individual token positions
356
+ in the flattened KV buffer.
357
+
358
+ When KERNEL_BLOCK_SIZE=1: block_idx=t, offset=0, flat=block_id
359
+ (equivalent to a direct copy -- no regression from the original kernel).
360
+
361
+ When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens)
362
+ is expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
363
"""
364
req_idx = tl.program_id(0)
365
row_ptr = block_table + req_idx * block_table_stride
317
- start_idx = tl.load(cu_num_blocks + req_idx)
318
- end_idx = tl.load(cu_num_blocks + req_idx + 1)
319
- num_blocks = end_idx - start_idx
366
+ start_idx = tl.load(cu_num_tokens + req_idx)
367
+ num_tokens = tl.load(seq_lens + req_idx)
368
369
offset = tl.arange(0, BLOCK_SIZE)
322
- for i in tl.range(0, num_blocks, BLOCK_SIZE):
323
- block_ids = tl.load(row_ptr + i + offset, mask=i + offset < num_blocks)
370
+ for i in tl.range(0, num_tokens, BLOCK_SIZE):
371
+ token_offsets = i + offset
372
+ mask = token_offsets < num_tokens
373
+
374
+ # Which block in the block table does this token belong to?
375
+ block_idx = token_offsets // KERNEL_BLOCK_SIZE
376
+ # Offset within that block
377
+ offset_in_block = token_offsets % KERNEL_BLOCK_SIZE
378
+
379
+ # Load the block ID from the block table
380
+ block_ids = tl.load(row_ptr + block_idx, mask=mask)
381
+
382
+ # Compute flat index in the flattened kv_buffer
383
+ flat_indices = block_ids * KERNEL_BLOCK_SIZE + offset_in_block
384
+
385
tl.store(
325
- page_indices + start_idx + i + offset,
326
- block_ids,
327
- mask=i + offset < num_blocks,
386
+ page_indices + start_idx + token_offsets,
387
+ flat_indices,
388
+ mask=mask,
389
)
390
391
@@ -426,6 +487,24 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
487
488
kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
489
490
+ # Build kwargs for mla_decode_fwd. Pass persistent metadata only
491
+ # when it was successfully computed (qseqlen=1 decode steps).
492
+ # For multi-token verification steps (spec-dec), the kernel falls
493
+ # back to computing metadata internally.
494
+ mla_kwargs = dict(
495
+ q_scale=layer._q_scale,
496
+ kv_scale=layer._k_scale,
497
+ )
498
+ if attn_metadata.work_meta_data is not None:
499
+ mla_kwargs.update(
500
+ work_meta_data=attn_metadata.work_meta_data,
501
+ work_indptr=attn_metadata.work_indptr,
502
+ work_info_set=attn_metadata.work_info_set,
503
+ reduce_indptr=attn_metadata.reduce_indptr,
504
+ reduce_final_map=attn_metadata.reduce_final_map,
505
+ reduce_partial_map=attn_metadata.reduce_partial_map,
506
+ )
507
+
508
rocm_aiter_ops.mla_decode_fwd(
509
q,
510
kv_buffer,
@@ -436,14 +515,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
515
attn_metadata.decode.paged_kv_indptr,
516
attn_metadata.decode.paged_kv_indices,
517
attn_metadata.decode.paged_kv_last_page_len,
439
- q_scale=layer._q_scale,
440
- kv_scale=layer._k_scale,
441
- work_meta_data=attn_metadata.work_meta_data,
442
- work_indptr=attn_metadata.work_indptr,
443
- work_info_set=attn_metadata.work_info_set,
444
- reduce_indptr=attn_metadata.reduce_indptr,
445
- reduce_final_map=attn_metadata.reduce_final_map,
446
- reduce_partial_map=attn_metadata.reduce_partial_map,
518
+ **mla_kwargs,
519
)
520
521
if self._needs_head_repeat: