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 14, 05:30 UTC
[Bugfix] Fix CUDA graph selection bug in FlashInfer at high concurrency (#26499)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Oct 9, 2025 at 17:12 UTC
6e783bc54b03f2bea2fc7996b565006652a3726c
1 file changed
+9
-2
vllm/v1/attention/backends/flashinfer.py
+9
-2
@@ -296,6 +296,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
296
)
297
max_num_reqs = vllm_config.scheduler_config.max_num_seqs
298
max_num_pages = max_num_reqs * max_num_pages_per_req
299
+ speculative_config = vllm_config.speculative_config
300
+ num_spec_tokens = (
301
+ speculative_config.num_speculative_tokens
302
+ if speculative_config is not None
303
+ else 0
304
+ )
305
self.enable_cuda_graph = (
306
self.compilation_config.cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
307
)
@@ -306,7 +312,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
312
int, BatchDecodeWithPagedKVCacheWrapper
313
] = {}
314
self._decode_cudagraph_max_bs = min(
309
- max_num_reqs, self.compilation_config.max_capture_size
315
+ (1 + num_spec_tokens) * max_num_reqs,
316
+ self.compilation_config.max_capture_size,
317
)
318
319
self.num_qo_heads = self.model_config.get_num_attention_heads(
@@ -679,7 +686,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
686
use_cudagraph = (
687
self.enable_cuda_graph
688
and pure_decode
682
- and num_decodes <= self._decode_cudagraph_max_bs
689
+ and num_decode_tokens <= self._decode_cudagraph_max_bs
690
)
691
if use_cudagraph:
692
num_input_tokens = self.vllm_config.pad_for_cudagraph(