diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c..c74e940 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1690,6 +1690,24 @@ class VllmConfig: max_cudagraph_capture_size = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) + # The bound above counts one query position per running + # sequence, which is the batch size of a pure-decode step. + # With chunked prefill a step also carries prompt tokens, so + # its width is set by the token budget rather than the + # sequence count. When the two disagree -- max_num_seqs 32 + # against max_num_batched_tokens 8192 -- every mixed + # prefill/decode step is wider than the largest captured + # graph and takes the eager path instead, paying Python + # dispatch per layer on the steps that dominate TTFT. + # Widen the default to cover those steps, capped so capture + # stays bounded in time and memory. An explicit + # max_cudagraph_capture_size is untouched: this only fills in + # the default, and the scheduler still admits exactly + # max_num_seqs sequences. + max_cudagraph_capture_size = max( + max_cudagraph_capture_size, + min(self.scheduler_config.max_num_batched_tokens, 768), + ) max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size)