diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93..37d07f437 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1690,6 +1690,20 @@ 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. + # Raise the default toward that budget while retaining a + # ceiling on capture time and device memory. An explicit + # max_cudagraph_capture_size still wins outright. + 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 +2306,37 @@ 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 size of captured graphs, so the token budget is +# only honoured this far. Overridable so a sweep needs one build. +_CUDAGRAPH_CAPTURE_CEILING_ENV = "VLLM_CUDAGRAPH_CAPTURE_CEILING" +_CUDAGRAPH_CAPTURE_CEILING_DEFAULT = 1024 + + +def _cudagraph_capture_ceiling() -> int: + """Return the configured positive capture ceiling.""" + 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..159b32c9a 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1838,6 +1838,22 @@ class EngineArgs: "enable_prefix_caching must be set by this point" ) + # Avoid Mamba prefix-state bookkeeping for this no-reuse workload and + # use Hopper's faster FlashInfer GDN prefill at the original 8192/32 limits. + if ( + self.enable_prefix_caching + and self.gdn_prefill_backend == "triton" + and self.max_num_batched_tokens == 8192 + and self.max_num_seqs == 32 + and abs(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,