From 4c3c17d43bbdb6ec082914a482052bb7ca21784b Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Sat, 4 Jul 2026 01:21:43 -0500 Subject: [PATCH] [ROCm] Disable persistent sparse-MLA kernel for chunked-prefill continuations (#47567) Signed-off-by: Rohan Potdar Co-authored-by: Claude Opus 4.8 (1M context) --- .../backends/mla/rocm_aiter_mla_sparse.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 9982b7aacb..55bba1a7ee 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -536,7 +536,17 @@ class ROCMAiterMLASparseMetadataBuilder( self._num_attention_heads, clamped_seq_lens.tobytes(), ) - if metadata_key != self._prev_metadata_key: + # The persistent MLA kernel is numerically wrong for multi-token prefill + # batches; errors compound across chunked prefill and break long-context + # decode (vllm#47042). Use it only for decode and single-chunk prefills, + # not chunked-prefill continuations (>1 query token, seq_len > query_len). + step_query_lens = seg_lengths + total_seq_lens = common_attn_metadata.seq_lens_cpu[:num_reqs].numpy() + is_chunked_continuation = (step_query_lens > 1) & ( + total_seq_lens > step_query_lens + ) + use_persistent = not is_chunked_continuation.any() + if use_persistent and metadata_key != self._prev_metadata_key: from aiter import get_mla_metadata_v1 get_mla_metadata_v1( @@ -576,7 +586,7 @@ class ROCMAiterMLASparseMetadataBuilder( paged_kv_last_page_len=paged_kv_last_page_len, paged_kv_indices=paged_kv_indices, paged_kv_indptr=paged_kv_indptr, - work_meta_data=self._mla_work_meta_data, + work_meta_data=self._mla_work_meta_data if use_persistent else None, work_indptr=self._mla_work_indptr, work_info_set=self._mla_work_info_set, reduce_indptr=self._mla_reduce_indptr,