[Bugfix][Spec Decode] Size DFlash query buffers for cudagraph-padded batches (#50065)
Signed-off-by: siddhant-bharti <sbharti@together.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Siddhant Bharti committed
Jul 28, 2026 at 12:34 UTC
0b6aa3c47ce69a1f3c8a19cafe3b9dc2871d1f6b
1 file changed
+7
-3
vllm/v1/spec_decode/dflash.py
+7
-3
@@ -38,8 +38,12 @@ class DFlashProposer(SpecDecodeBaseProposer):
38
39
# Only next_token_ids and mask tokens are query tokens, all other context is K/V
40
self.max_query_tokens = self.max_batch_size * (1 + self.num_speculative_tokens)
41
+ self.max_padded_query_tokens = max(
42
+ self.max_query_tokens,
43
+ vllm_config.compilation_config.max_cudagraph_capture_size or 0,
44
+ )
45
# Positions covers both context states + query states
42
- self.max_positions = self.max_num_tokens + self.max_query_tokens
46
+ self.max_positions = self.max_num_tokens + self.max_padded_query_tokens
47
48
# Separate context buffers to keep query buffer addresses stable for CUDA graphs
49
self._context_slot_mapping_buffer = torch.zeros(
@@ -48,7 +52,7 @@ class DFlashProposer(SpecDecodeBaseProposer):
52
device=device,
53
)
54
self._slot_mapping_buffer = torch.zeros(
51
- self.max_query_tokens,
55
+ self.max_padded_query_tokens,
56
dtype=torch.int64,
57
device=device,
58
)
@@ -58,7 +62,7 @@ class DFlashProposer(SpecDecodeBaseProposer):
62
device=device,
63
)
64
self.positions = torch.zeros(
61
- self.max_query_tokens,
65
+ self.max_padded_query_tokens,
66
dtype=torch.int64,
67
device=device,
68
)