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 10, 01:30 UTC
[Bugfix] [DeepSeek-V3.2] fix sparse_attn_indexer padding (#32175)
Signed-off-by: Kebe <mail@kebe7jun.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Kebe committed
Jan 16, 2026 at 12:21 UTC
5de6dd0662da663aa57127e82499d2c2d604c3d0
1 file changed
+9
-2
vllm/model_executor/models/deepseek_v2.py
+9
-2
@@ -717,13 +717,20 @@ def sparse_attn_indexer(
717
# decode_threshold since we unstrictly split
718
# prefill and decode by decode_threshold
719
# (currently set to 1 + speculative tokens)
720
+
721
+ # [num_decode_tokens, n_head, head_dim] -> [bs, 1+next_n, n_head, head_dim]
722
padded_q_fp8_decode_tokens = pack_seq_triton(
723
q_fp8[:num_decode_tokens], decode_lens
724
)
725
+ # [num_decode_tokens, n_head] -> [bs, 1+next_n, n_head]
726
+ padded_weights = pack_seq_triton(weights[:num_decode_tokens], decode_lens)
727
+ # [bs, 1+next_n, n_head] -> [bs * next_n, n_head]
728
+ padded_weights = padded_weights.flatten(0, 1)
729
else:
730
padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
731
decode_lens.shape[0], -1, *q_fp8.shape[1:]
732
)
733
+ padded_weights = weights
734
# TODO: move and optimize below logic with triton kernels
735
batch_size = padded_q_fp8_decode_tokens.shape[0]
736
next_n = padded_q_fp8_decode_tokens.shape[1]
@@ -739,14 +746,14 @@ def sparse_attn_indexer(
746
logits = fp8_paged_mqa_logits_func(
747
padded_q_fp8_decode_tokens,
748
kv_cache,
742
- weights[:num_padded_tokens],
749
+ padded_weights[:num_padded_tokens],
750
decode_metadata.seq_lens,
751
decode_metadata.block_table,
752
decode_metadata.schedule_metadata,
753
max_model_len=max_model_len,
754
)
755
num_rows = logits.shape[0]
749
- topk_indices = topk_indices_buffer[:num_decode_tokens, :topk_tokens]
756
+ topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
757
758
torch.ops._C.top_k_per_row_decode(
759
logits,