[ROCm][DSv3.2][Perf] Cap sparse MLA decode KV-splits with a work-per-split heuristic (#46832)
Signed-off-by: Frida Andersson <fanderss@amd.com>
frida-andersson committed
Jul 18, 2026 at 18:39 UTC
e94243893dd30256f58644ad4ecf779be757dff8
1 file changed
+21
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+21
@@ -369,6 +369,8 @@ class ROCMAiterMLASparseMetadataBuilder(
369
self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
370
self.mla_dims = get_mla_dims(self.model_config)
371
self.topk_tokens = vllm_config.model_config.hf_config.index_topk
372
+ # Bounds the KV-split heuristic (see `_sparse_decode_max_split`).
373
+ self._num_compute_units = current_platform.num_compute_units()
374
self.max_model_len_tensor = torch.tensor(
375
[self.model_config.max_model_len], device=device, dtype=torch.int32
376
)
@@ -464,6 +466,21 @@ class ROCMAiterMLASparseMetadataBuilder(
466
self._prev_indices_extent: int = 0
467
self._prev_metadata_key: tuple | None = None
468
469
+ def _sparse_decode_max_split(self, max_seq_len: int) -> int:
470
+ """Cap ``max_split_per_batch`` for the aiter sparse-MLA decode reduce.
471
+
472
+ The reduce only covers the selected tokens per row (``<= topk_tokens``),
473
+ so aiter's default (``-1`` => split across every CU) over-fragments it.
474
+ Mirror ``triton_mla.py``: aim for a minimum work per split, round to a
475
+ power of two, and cap by the CU count. Numerics are unchanged.
476
+ """
477
+ effective_len = min(max_seq_len, self.topk_tokens)
478
+ min_work_per_split = 128
479
+ ideal_splits = triton.next_power_of_2(
480
+ max(1, effective_len // min_work_per_split)
481
+ )
482
+ return min(ideal_splits, self._num_compute_units)
483
+
484
def build(
485
self,
486
common_prefix_len: int,
@@ -547,6 +564,9 @@ class ROCMAiterMLASparseMetadataBuilder(
564
if metadata_key != self._prev_metadata_key:
565
from aiter import get_mla_metadata_v1
566
567
+ max_split_per_batch = self._sparse_decode_max_split(
568
+ int(common_attn_metadata.max_seq_len)
569
+ )
570
get_mla_metadata_v1(
571
qo_indptr,
572
paged_kv_indptr,
@@ -565,6 +585,7 @@ class ROCMAiterMLASparseMetadataBuilder(
585
max_seqlen_qo=1,
586
uni_seqlen_qo=1,
587
fast_mode=True,
588
+ max_split_per_batch=max_split_per_batch,
589
)
590
# The persistent metadata buffers are read by graph replay. Order
591
# the async metadata write before the graph-captured decode kernel.