[ROCm][Bugfix] Key sparse-MLA persistent metadata on per-request context lengths (#47766)
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 7, 2026 at 19:22 UTC
4aceabf8c1a40f8d576ba9b83e4b9a8854eda2d5
1 file changed
+12
-15
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+12
-15
@@ -522,31 +522,28 @@ class ROCMAiterMLASparseMetadataBuilder(
522
# treated as its own batch entry), so persistent metadata can always
523
# be precomputed here. The kernel switches to the persistent
524
# work-stealing path automatically when work_meta_data is non-None.
525
- # The output is a deterministic function of (num_tokens, max_query_len,
526
- # num_heads, min(seq_lens, topk_tokens)); fingerprint those CPU-side
527
- # and skip the launch when nothing changed.
525
+ # The output is a deterministic function of the per-request query and
526
+ # context lengths (both clamped to topk_tokens, past which per-token KV
527
+ # length saturates) and num_heads; fingerprint those CPU-side and skip
528
+ # the launch when nothing changed.
529
num_reqs = common_attn_metadata.num_reqs
530
clamped_seq_lens = np.minimum(
531
common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(),
532
self.topk_tokens,
533
)
534
+ clamped_context_lens = np.minimum(
535
+ common_attn_metadata.seq_lens_cpu[:num_reqs].numpy() - seg_lengths,
536
+ self.topk_tokens,
537
+ )
538
metadata_key = (
539
num_tokens,
540
int(common_attn_metadata.max_query_len),
541
self._num_attention_heads,
542
clamped_seq_lens.tobytes(),
543
+ clamped_context_lens.tobytes(),
544
+ seg_lengths.tobytes(),
545
)
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:
546
+ if metadata_key != self._prev_metadata_key:
547
from aiter import get_mla_metadata_v1
548
549
get_mla_metadata_v1(
@@ -586,7 +583,7 @@ class ROCMAiterMLASparseMetadataBuilder(
583
paged_kv_last_page_len=paged_kv_last_page_len,
584
paged_kv_indices=paged_kv_indices,
585
paged_kv_indptr=paged_kv_indptr,
589
- work_meta_data=self._mla_work_meta_data if use_persistent else None,
586
+ work_meta_data=self._mla_work_meta_data,
587
work_indptr=self._mla_work_indptr,
588
work_info_set=self._mla_work_info_set,
589
reduce_indptr=self._mla_reduce_indptr,