From 0b6aa3c47ce69a1f3c8a19cafe3b9dc2871d1f6b Mon Sep 17 00:00:00 2001 From: Siddhant Bharti <42143349+siddhant-bharti@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:34:22 -0700 Subject: [PATCH] [Bugfix][Spec Decode] Size DFlash query buffers for cudagraph-padded batches (#50065) Signed-off-by: siddhant-bharti Co-authored-by: Claude Fable 5 --- vllm/v1/spec_decode/dflash.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index 5d1c062921..626dd36dea 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -38,8 +38,12 @@ class DFlashProposer(SpecDecodeBaseProposer): # Only next_token_ids and mask tokens are query tokens, all other context is K/V self.max_query_tokens = self.max_batch_size * (1 + self.num_speculative_tokens) + self.max_padded_query_tokens = max( + self.max_query_tokens, + vllm_config.compilation_config.max_cudagraph_capture_size or 0, + ) # Positions covers both context states + query states - self.max_positions = self.max_num_tokens + self.max_query_tokens + self.max_positions = self.max_num_tokens + self.max_padded_query_tokens # Separate context buffers to keep query buffer addresses stable for CUDA graphs self._context_slot_mapping_buffer = torch.zeros( @@ -48,7 +52,7 @@ class DFlashProposer(SpecDecodeBaseProposer): device=device, ) self._slot_mapping_buffer = torch.zeros( - self.max_query_tokens, + self.max_padded_query_tokens, dtype=torch.int64, device=device, ) @@ -58,7 +62,7 @@ class DFlashProposer(SpecDecodeBaseProposer): device=device, ) self.positions = torch.zeros( - self.max_query_tokens, + self.max_padded_query_tokens, dtype=torch.int64, device=device, )