diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 13ff66fac8d9..53580ecaa84b 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -528,12 +528,15 @@ def dispatch( num_active_loras: int, max_query_len: int | None = None, num_ubatches: int = 1, + allow_full: bool = True, ) -> BatchExecutionDescriptor: """Find matching cudagraph descriptor from priority-ordered candidates.""" effective_loras = self._resolve_effective_loras(num_active_loras) key = (num_tokens, effective_loras) if self._graphs_captured and num_tokens > 0 and key in self._candidates: for desc in self._candidates[key]: + if not allow_full and desc.cg_mode == CUDAGraphMode.FULL: + continue if _is_compatible( desc, num_reqs, diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index fb9f4c02d3e4..b0d2bc81c68e 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -223,6 +223,7 @@ def dispatch_cg_and_sync_dp( uniform_decode: bool = False, allow_skip_drafts: bool = False, dp_sync_state: DPSyncState | None = None, + allow_full: bool = True, ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: """Pick a cudagraph descriptor for this batch, agreeing it across DP ranks. @@ -287,6 +288,7 @@ def dispatch_cg_and_sync_dp( dp_sync_state.uniform_token_count if dp_sync_state else uniform_token_count, num_active_loras=num_active_loras, max_query_len=max_query_len, + allow_full=allow_full, ) if dp_size == 1: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 31b2b6ad5132..487a2f70d31c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1774,12 +1774,14 @@ def execute_model( num_reqs = len(scheduler_output.num_scheduled_tokens) num_toks = scheduler_output.total_num_scheduled_tokens max_query_len: int | None = max(scheduler_output.num_scheduled_tokens.values()) + allow_full = True batch_req_state, uniform_tok_count = self.gather_batch_req_state( scheduler_output, dummy_run ) if batch_req_state is not None: num_toks = batch_req_state.num_tokens if batch_req_state.has_prefill: + allow_full = False # Varlen decode graphs replay decodes only, and their bound # alone would admit a short prefill. max_query_len = None @@ -1818,6 +1820,7 @@ def execute_model( self.ubatch_runner is not None and not skip_attn_for_dummy_run ), uniform_decode=uniform_tok_count == self.decode_query_len, + allow_full=allow_full, allow_skip_drafts=scheduler_output.skip_drafts, )