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, 03:45 UTC
enable persistent mla for sparse mla backend (#41990)
Signed-off-by: ganyi <ygan@amd.com> Signed-off-by: Douglas Lehr <Doug.Lehr@amd.com> Co-authored-by: ganyi <ygan@amd.com>
Douglas Lehr committed
May 7, 2026 at 22:10 UTC
baf068d8be09b22db636ae8c6a2082033e2b6911
2 files changed
+120
-5
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
-3
@@ -420,9 +420,6 @@ class AiterMLAHelper:
420
421
@staticmethod
422
def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:
423
- assert num_heads not in AiterMLAHelper._AITER_UNSUPPORTED_HEADS, (
424
- f"unsupported head_num: {num_heads}"
425
- )
423
return (
424
q
425
if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+120
-2
@@ -332,6 +332,15 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata):
332
block_size: int = 1
333
topk_tokens: int = 2048
334
335
+ # Persistent MLA metadata (only populated when persistent mode is enabled,
336
+ # i.e. when the aiter sparse decode kernel supports work-stealing splits).
337
+ work_meta_data: torch.Tensor | None = None
338
+ work_indptr: torch.Tensor | None = None
339
+ work_info_set: torch.Tensor | None = None
340
+ reduce_indptr: torch.Tensor | None = None
341
+ reduce_final_map: torch.Tensor | None = None
342
+ reduce_partial_map: torch.Tensor | None = None
343
+
344
345
@dataclass
346
class ROCMAiterMLASparseMetadataBuilder(
@@ -389,6 +398,66 @@ class ROCMAiterMLASparseMetadataBuilder(
398
[max_num_batched_tokens + 1], dtype=torch.int32, device=device
399
)
400
401
+ # ----- Persistent MLA metadata buffers -----
402
+ # The aiter sparse decode kernel supports a "persistent" path that
403
+ # uses precomputed work-splitting metadata for better load balancing
404
+ # across CUs. Mirrors the approach used in rocm_aiter_mla.py.
405
+ #
406
+ # In the sparse case each query token is its own "batch" entry in the
407
+ # qo_indptr (qo_indptr = [0, 1, 2, ..., num_tokens]) and max_qo_len=1.
408
+ # We pad get_mla_metadata_info_v1's batch_size to max_num_batched_tokens
409
+ # so the buffers are large enough for any decode shape we might see.
410
+ from aiter import dtypes, get_mla_metadata_info_v1
411
+
412
+ # Aiter sparse MLA also requires num_heads >= 16 (will be padded by
413
+ # AiterMLAHelper.get_mla_padded_q in forward).
414
+ self._num_attention_heads = max(16, self.num_heads)
415
+
416
+ q_dtype = self.model_dtype
417
+ kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
418
+ if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
419
+ kv_cache_dtype_str = "fp8"
420
+ else:
421
+ kv_cache_dtype_str = "bf16"
422
+ kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16)
423
+
424
+ (
425
+ (work_meta_data_size, work_meta_data_type),
426
+ (work_indptr_size, work_indptr_type),
427
+ (work_info_set_size, work_info_set_type),
428
+ (reduce_indptr_size, reduce_indptr_type),
429
+ (reduce_final_map_size, reduce_final_map_type),
430
+ (reduce_partial_map_size, reduce_partial_map_type),
431
+ ) = get_mla_metadata_info_v1(
432
+ max_num_batched_tokens,
433
+ 1,
434
+ self._num_attention_heads,
435
+ q_dtype,
436
+ kv_dtype,
437
+ is_sparse=True,
438
+ fast_mode=True,
439
+ )
440
+ self._mla_work_meta_data = torch.empty(
441
+ work_meta_data_size, dtype=work_meta_data_type, device=device
442
+ )
443
+ self._mla_work_indptr = torch.empty(
444
+ work_indptr_size, dtype=work_indptr_type, device=device
445
+ )
446
+ self._mla_work_info_set = torch.empty(
447
+ work_info_set_size, dtype=work_info_set_type, device=device
448
+ )
449
+ self._mla_reduce_indptr = torch.empty(
450
+ reduce_indptr_size, dtype=reduce_indptr_type, device=device
451
+ )
452
+ self._mla_reduce_final_map = torch.empty(
453
+ reduce_final_map_size, dtype=reduce_final_map_type, device=device
454
+ )
455
+ self._mla_reduce_partial_map = torch.empty(
456
+ reduce_partial_map_size,
457
+ dtype=reduce_partial_map_type,
458
+ device=device,
459
+ )
460
+
461
def build(
462
self,
463
common_prefix_len: int,
@@ -431,6 +500,33 @@ class ROCMAiterMLASparseMetadataBuilder(
500
paged_kv_indptr = self.paged_kv_indptr[: num_tokens + 1]
501
paged_kv_indices = self.paged_kv_indices[: num_tokens * self.topk_tokens]
502
503
+ # ----- Compute persistent MLA metadata -----
504
+ # The aiter sparse decode kernel uses qseqlen=1 (each query token is
505
+ # treated as its own batch entry), so persistent metadata can always
506
+ # be precomputed here. The kernel switches to the persistent
507
+ # work-stealing path automatically when work_meta_data is non-None.
508
+ from aiter import get_mla_metadata_v1
509
+
510
+ get_mla_metadata_v1(
511
+ qo_indptr,
512
+ paged_kv_indptr,
513
+ paged_kv_last_page_len,
514
+ self._num_attention_heads,
515
+ 1,
516
+ True,
517
+ self._mla_work_meta_data,
518
+ self._mla_work_info_set,
519
+ self._mla_work_indptr,
520
+ self._mla_reduce_indptr,
521
+ self._mla_reduce_final_map,
522
+ self._mla_reduce_partial_map,
523
+ page_size=1,
524
+ kv_granularity=16,
525
+ max_seqlen_qo=1,
526
+ uni_seqlen_qo=1,
527
+ fast_mode=True,
528
+ )
529
+
530
metadata = ROCMAiterMLASparseMetadata(
531
num_reqs=common_attn_metadata.num_reqs,
532
max_query_len=common_attn_metadata.max_query_len,
@@ -447,6 +543,12 @@ class ROCMAiterMLASparseMetadataBuilder(
543
paged_kv_last_page_len=paged_kv_last_page_len,
544
paged_kv_indices=paged_kv_indices,
545
paged_kv_indptr=paged_kv_indptr,
546
+ work_meta_data=self._mla_work_meta_data,
547
+ work_indptr=self._mla_work_indptr,
548
+ work_info_set=self._mla_work_info_set,
549
+ reduce_indptr=self._mla_reduce_indptr,
550
+ reduce_final_map=self._mla_reduce_final_map,
551
+ reduce_partial_map=self._mla_reduce_partial_map,
552
)
553
return metadata
554
@@ -525,6 +627,23 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
627
device=q.device,
628
)
629
630
+ # Build kwargs and forward the persistent MLA metadata when it has
631
+ # been computed. The aiter mla_decode_fwd switches to its
632
+ # work-stealing persistent kernel path when work_meta_data is given.
633
+ mla_kwargs: dict = dict(
634
+ q_scale=layer._q_scale,
635
+ kv_scale=layer._k_scale,
636
+ )
637
+ if attn_metadata.work_meta_data is not None:
638
+ mla_kwargs.update(
639
+ work_meta_data=attn_metadata.work_meta_data,
640
+ work_indptr=attn_metadata.work_indptr,
641
+ work_info_set=attn_metadata.work_info_set,
642
+ reduce_indptr=attn_metadata.reduce_indptr,
643
+ reduce_final_map=attn_metadata.reduce_final_map,
644
+ reduce_partial_map=attn_metadata.reduce_partial_map,
645
+ )
646
+
647
rocm_aiter_ops.mla_decode_fwd(
648
q,
649
kv_c_and_k_pe_cache,
@@ -535,8 +654,7 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
654
attn_metadata.paged_kv_indptr,
655
attn_metadata.paged_kv_indices,
656
attn_metadata.paged_kv_last_page_len,
538
- q_scale=layer._q_scale,
539
- kv_scale=layer._k_scale,
657
+ **mla_kwargs,
658
)
659
660
return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output)