diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c..704cf43 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -77,6 +77,21 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( ) +# Ceiling for the chunked-prefill extension of the CUDA graph capture range in +# _set_cudagraph_sizes. The prompt side of a mixed step is what carries it past the +# decode-only bound, so this holds one whole prompt plus a full decode batch, over +# the whole prompt-length range rather than its short half. Beyond it a step is +# several prompts at once, which is rare enough not to pay for. +_CHUNKED_PREFILL_CAPTURE_LIMIT = 2816 + +# Capture sizes step by 16 above 256. Padding is bounded by the stride, so its cost +# relative to the batch falls as the batch grows, while every size costs capture +# time. Keep the fine stride over the whole range the decode-side bound can reach +# and widen it only across the extension, where 64 tokens is under 6% of the batch. +_COARSE_CAPTURE_SIZE_THRESHOLD = 1024 +_COARSE_CAPTURE_SIZE_STRIDE = 64 + + class OptimizationLevel(IntEnum): """Optimization level enum.""" @@ -1690,6 +1705,30 @@ 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 size of a step that is purely decode. Under chunked + # prefill a step also carries prompt tokens and is bounded by the + # token budget instead, and the two differ by an order of + # magnitude: max_num_seqs 32 stops capture at 64 while a step that + # admits one 2k-token prompt runs to ~2k. Every such step exceeds + # the largest captured graph and takes the eager path (see this + # method's note: "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, capped so capture + # stays affordable in time and memory. Only the default is filled + # in here; an explicit max_cudagraph_capture_size still wins, and + # max_num_seqs keeps its meaning everywhere else -- the scheduler + # still admits exactly max_num_seqs sequences. + if self.scheduler_config.enable_chunked_prefill: + max_cudagraph_capture_size = max( + max_cudagraph_capture_size, + min( + self.scheduler_config.max_num_batched_tokens, + _CHUNKED_PREFILL_CAPTURE_LIMIT, + ), + ) max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size) @@ -1729,8 +1768,32 @@ class VllmConfig: if max_cudagraph_capture_size >= 256: # Step size 16 for larger batch sizes cudagraph_capture_sizes += list( - range(256, max_cudagraph_capture_size + 1, 16) + range( + 256, + min( + max_cudagraph_capture_size, + _COARSE_CAPTURE_SIZE_THRESHOLD, + ) + + 1, + 16, + ) + ) + if max_cudagraph_capture_size > _COARSE_CAPTURE_SIZE_THRESHOLD: + # Wider step size across the extension, where padding is a + # small fraction of the batch, so the added range costs few + # extra graphs. Sizes at or below the threshold are unchanged, + # so no batch pads further than it did before. + cudagraph_capture_sizes += list( + range( + _COARSE_CAPTURE_SIZE_THRESHOLD + + _COARSE_CAPTURE_SIZE_STRIDE, + max_cudagraph_capture_size + 1, + _COARSE_CAPTURE_SIZE_STRIDE, + ) ) + # The stride need not divide the range; keep the ceiling itself + # capturable so the largest step still gets a graph. + cudagraph_capture_sizes.append(max_cudagraph_capture_size) # ensure max_num_tokens is captured if within max capture size if ( max_num_tokens <= max_cudagraph_capture_size