From 111bf7f0c5f8c833465175eaa3cb93cd5a10a93b Mon Sep 17 00:00:00 2001
From: xzwgit <57032909+xzwgit@users.noreply.github.com>
Date: Tue, 20 Oct 2026 02:14:20 +0800
Subject: [PATCH] [Bugfix][Worker] Keep prefill-bearing batches out of FULL
 (decode-only) cudagraphs

A batch that contains prefill tokens can be dispatched into a varlen FULL cudagraph captured on a decode-shaped batch (one token per request) when the descriptor carries no max_query_len constraint: the padded replay re-shapes the prefill queries and corrupts the logits. Gate FULL candidates on allow_full=False whenever the batch has prefill.

Fixes the output corruption reported in #62417.
---
 vllm/v1/worker/gpu/cudagraph_utils.py | 3 +++
 vllm/v1/worker/gpu/dp_utils.py        | 2 ++
 vllm/v1/worker/gpu/model_runner.py    | 3 +++
 3 files changed, 8 insertions(+)

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 2bcad34e0560..a957da5ee12b 100644
--- a/vllm/v1/worker/gpu/dp_utils.py
+++ b/vllm/v1/worker/gpu/dp_utils.py
@@ -222,6 +222,7 @@ def dispatch_cg_and_sync_dp(
     allow_ubatching: bool = False,
     uniform_decode: 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.
 
@@ -283,6 +284,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 2d30fbb612cf..9bbfa867c49b 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -1764,12 +1764,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
@@ -1808,6 +1810,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,
         )
 
         if batch_desc.num_tokens == 0: