diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93b..9e59c60a94 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -77,6 +77,12 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( ) +# Upper bound for the token-keyed CUDA graph capture range chosen in +# _set_cudagraph_sizes. Capture cost grows with the number and size of captured +# graphs, so the token budget is only honoured up to this many tokens. +_TOKEN_KEYED_CUDAGRAPH_CAPTURE_LIMIT = 1024 + + class OptimizationLevel(IntEnum): """Optimization level enum.""" @@ -1690,6 +1696,32 @@ class VllmConfig: max_cudagraph_capture_size = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) + # The bound above is derived from the sequence count, which is + # the right bound for a step that is purely decode: one query + # position per running sequence. Under chunked prefill a step + # also carries prompt tokens, so its batch size is bounded by + # the token budget instead, and the two can differ by an order + # of magnitude -- max_num_seqs 32 with max_num_batched_tokens + # 8192 stops capture at 64 while real mixed prefill/decode + # steps run into the hundreds of tokens. Every such step then + # exceeds the largest captured graph and falls back to the + # eager path (see the note on this method: "If batch size > + # largest cudagraph_capture_sizes, cudagraph will not be + # used"), paying full Python dispatch on every layer. + # + # Take the token budget into account as well, bounded so that + # capture stays cheap in time and memory. An explicit + # max_cudagraph_capture_size still wins: this only fills in the + # default, and max_num_seqs is not read differently anywhere + # else -- 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, + _TOKEN_KEYED_CUDAGRAPH_CAPTURE_LIMIT, + ), + ) max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size)