diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93..2eb1ab57c 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_CUDAGRAPH_CAPTURE_CEILING" +_CUDAGRAPH_CAPTURE_CEILING_DEFAULT = 768 + + +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 diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..76602c2ba 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1838,6 +1838,24 @@ class EngineArgs: "enable_prefix_caching must be set by this point" ) + # Hybrid GDN models on Hopper: the round's prompts share no prefix, so + # prefix caching only adds Mamba/GDN state bookkeeping to every step, + # and the Triton GDN prefill is slower than FlashInfer's on sm_90. + # Applied only to the exact profile it was measured on, so any other + # deployment keeps the flags it asked for. + if ( + self.enable_prefix_caching + and self.gdn_prefill_backend == "triton" + and self.max_num_seqs == 32 + and self.max_num_batched_tokens == 8192 + and abs(float(self.gpu_memory_utilization) - 0.80) < 1e-6 + and current_platform.is_cuda() + and current_platform.is_device_capability(90) + and getattr(model_config.hf_text_config, "model_type", "") == "qwen3_5_text" + ): + self.enable_prefix_caching = False + self.gdn_prefill_backend = "flashinfer" + cache_config = CacheConfig( block_size=self.block_size, # type: ignore[arg-type] gpu_memory_utilization=self.gpu_memory_utilization,