diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93..5421594e3 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1690,6 +1690,32 @@ class VllmConfig: max_cudagraph_capture_size = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) + # The bound above sizes a step that is purely decode: one + # query position per running sequence, times the speculative + # width. Under chunked prefill a step also carries prompt + # tokens, so its size is bounded by the token budget instead, + # and the two differ by more than an order of magnitude -- + # max_num_seqs 32 with max_num_batched_tokens 8192 stops + # capture in the low hundreds while a mixed step runs to + # thousands of tokens. Per this method's own contract, "if + # batch size > largest cudagraph_capture_sizes, cudagraph will + # not be used", so every mixed step then pays full Python + # dispatch per layer. + # + # Raise the default toward the token budget. The ceiling is + # what keeps capture affordable: graphs cost startup time, + # which is not on the measured path, but they also cost device + # memory, which competes with the KV cache and is the one + # resource that can fail a start outright. An explicit + # max_cudagraph_capture_size still wins outright; this only + # fills in the default. + max_cudagraph_capture_size = max( + max_cudagraph_capture_size, + min( + self.scheduler_config.max_num_batched_tokens, + _cudagraph_capture_ceiling(), + ), + ) max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size) @@ -2292,3 +2318,38 @@ def get_layers_from_vllm_config( for layer_name in layer_names if isinstance(layer := forward_context.get(layer_name), layer_type) } + + +# Ceiling on the token-keyed CUDA graph capture range chosen in +# _set_cudagraph_sizes when the config pins no explicit size. Capture cost +# grows with the number and the size of the captured graphs, so the token +# budget is only honoured this far. Overridable so a sweep needs one build +# rather than one build per ceiling. +_CUDAGRAPH_CAPTURE_CEILING_ENV = "VLLM_MIXED_STEP_CAPTURE_CEILING" +_CUDAGRAPH_CAPTURE_CEILING_DEFAULT = 1536 + + +def _cudagraph_capture_ceiling() -> int: + """Capture ceiling, from the environment when set and usable.""" + raw = os.environ.get(_CUDAGRAPH_CAPTURE_CEILING_ENV) + if raw is None: + return _CUDAGRAPH_CAPTURE_CEILING_DEFAULT + try: + ceiling = int(raw) + except ValueError: + logger.warning( + "%s=%r is not an integer; using %d", + _CUDAGRAPH_CAPTURE_CEILING_ENV, + raw, + _CUDAGRAPH_CAPTURE_CEILING_DEFAULT, + ) + return _CUDAGRAPH_CAPTURE_CEILING_DEFAULT + if ceiling < 1: + logger.warning( + "%s=%d is not positive; using %d", + _CUDAGRAPH_CAPTURE_CEILING_ENV, + ceiling, + _CUDAGRAPH_CAPTURE_CEILING_DEFAULT, + ) + return _CUDAGRAPH_CAPTURE_CEILING_DEFAULT + return ceiling