[ROCm] Disable persistent sparse-MLA kernel for chunked-prefill continuations (#47567)
Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rohan Potdar committed
Jul 4, 2026 at 01:21 UTC
4c3c17d43bbdb6ec082914a482052bb7ca21784b
1 file changed
+12
-2
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+12
-2
@@ -536,7 +536,17 @@ class ROCMAiterMLASparseMetadataBuilder(
536
self._num_attention_heads,
537
clamped_seq_lens.tobytes(),
538
)
539
- if metadata_key != self._prev_metadata_key:
539
+ # The persistent MLA kernel is numerically wrong for multi-token prefill
540
+ # batches; errors compound across chunked prefill and break long-context
541
+ # decode (vllm#47042). Use it only for decode and single-chunk prefills,
542
+ # not chunked-prefill continuations (>1 query token, seq_len > query_len).
543
+ step_query_lens = seg_lengths
544
+ total_seq_lens = common_attn_metadata.seq_lens_cpu[:num_reqs].numpy()
545
+ is_chunked_continuation = (step_query_lens > 1) & (
546
+ total_seq_lens > step_query_lens
547
+ )
548
+ use_persistent = not is_chunked_continuation.any()
549
+ if use_persistent and metadata_key != self._prev_metadata_key:
550
from aiter import get_mla_metadata_v1
551
552
get_mla_metadata_v1(
@@ -576,7 +586,7 @@ class ROCMAiterMLASparseMetadataBuilder(
586
paged_kv_last_page_len=paged_kv_last_page_len,
587
paged_kv_indices=paged_kv_indices,
588
paged_kv_indptr=paged_kv_indptr,
579
- work_meta_data=self._mla_work_meta_data,
589
+ work_meta_data=self._mla_work_meta_data if use_persistent else None,
590
work_indptr=self._mla_work_indptr,
591
work_info_set=self._mla_work_info_set,
592
reduce_indptr=self._mla_reduce_indptr,