Fix DSpark warmup without sparse index buffer (#50693)
Signed-off-by: xijiade.aihemaiti <3146335281@qq.com> Signed-off-by: xijiaat <52128022+xijiaat@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>
xijiaat committed
Aug 11, 2026 at 03:12 UTC
d40c3e3c00ab904928ab51c9c62b7f514c39647d
1 file changed
+5
-2
vllm/models/deepseek_v4/nvidia/flashmla.py
+5
-2
@@ -96,8 +96,11 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
96
// self.compress_ratio
97
)
98
M = N + self.window_size + self.max_num_batched_tokens
99
- assert self.topk_indices_buffer is not None
100
- top_k = 0 if swa_only else self.topk_indices_buffer.shape[-1]
99
+ if swa_only:
100
+ top_k = 0
101
+ else:
102
+ assert self.topk_indices_buffer is not None
103
+ top_k = self.topk_indices_buffer.shape[-1]
104
combined_topk = round_up(top_k + self.window_size, 128)
105
current_workspace_manager().get_simultaneous(
106
((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),