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 5, 06:00 UTC
[Bug] Fix rocm sparse attn indexer issue (#39225)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Wentao Ye committed
Apr 13, 2026 at 10:53 UTC
ccd0d1d9067a0bf24330a87044dca272e4a5228c
1 file changed
+5
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+5
@@ -532,6 +532,11 @@ def rocm_aiter_sparse_attn_indexer(
532
has_prefill = layer_attn_metadata.num_prefills > 0
533
num_decode_tokens = layer_attn_metadata.num_decode_tokens
534
535
+ # during speculative decoding, k may be padded to the CUDA graph batch
536
+ # size while slot_mapping only covers actual tokens.
537
+ num_tokens = slot_mapping.shape[0]
538
+ k = k[:num_tokens]
539
+
540
ops.indexer_k_quant_and_cache(
541
k,
542
kv_cache,