[Bugfix] Take padded path for ragged decode batches in sparse_attn_indexer (#52500)
Signed-off-by: pavelzak <pavel.zakharov@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Pavel Zakharov committed
Sep 19, 2026 at 13:14 UTC
133b71e0beec3c7bcaae5b0839a03e7e88d20bc7
1 file changed
+15
-2
vllm/model_executor/layers/sparse_attn_indexer.py
+15
-2
@@ -642,6 +642,19 @@ def sparse_attn_indexer(
642
assert decode_metadata is not None
643
kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
644
decode_lens = decode_metadata.decode_lens
645
+ # requires_padding can be computed False for ragged warmup/mixed
646
+ # batches (seen on SM120 TP=2: 8 tokens over 6 seqs -> uniform
647
+ # reshape below would crash). The token-count/batch-size divisibility
648
+ # check is a necessary-but-not-sufficient proxy for "uniform batch"
649
+ # (e.g. decode_lens=[1, 2, 3] divides evenly by 3 but isn't uniform);
650
+ # it catches the specific case seen in practice without requiring a
651
+ # device sync to inspect decode_lens directly, which would break
652
+ # CUDA-graph-capture safety on this path. Shared by both the pack
653
+ # branch below and the matching unpack branch further down — they
654
+ # must agree on which path was taken.
655
+ needs_padded_path = decode_metadata.requires_padding or (
656
+ num_decode_tokens % decode_lens.shape[0] != 0
657
+ )
658
if num_decode_tokens == 0:
659
padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:])
660
padded_q_scale = (
@@ -649,7 +662,7 @@ def sparse_attn_indexer(
662
if q_scale is not None
663
else None
664
)
652
- elif decode_metadata.requires_padding:
665
+ elif needs_padded_path:
666
# pad in edge case where we have short chunked prefill length <
667
# decode_threshold since we unstrictly split
668
# prefill and decode by decode_threshold
@@ -773,7 +786,7 @@ def sparse_attn_indexer(
786
cp_kv_cache_interleave_size,
787
)
788
776
- if decode_metadata.requires_padding:
789
+ if needs_padded_path:
790
# if padded, we need to unpack
791
# the topk indices removing padded tokens
792
topk_indices = unpack_seq_triton(