fix large batch performance.

PatchouliTaisa · PatchouliTaisa · commit 1fbf2969a1e4 · 2025-12-02T15:49:16.000+08:00
Signed-off-by: PatchouliTaisa &lt;patchychen@tencent.com&gt;
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
@@ -749,17 +749,8 @@ def _update_after_schedule(
         # 3. If some tokens (e.g. spec tokens) are rejected later, the number of
         #    computed tokens will be adjusted in update_from_output.
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
-        spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens
         for req_id, num_scheduled_token in num_scheduled_tokens.items():
             request = self.requests[req_id]
-            # DEBUG LOG: Track num_computed_tokens update in scheduler
-            spec_tokens = spec_decode_tokens.get(req_id, [])
-            logger.info(f"[DEBUG-SCHED] _update_after_schedule: "
-                       f"req_id={req_id}, "
-                       f"num_computed_tokens_before={request.num_computed_tokens}, "
-                       f"num_scheduled_token={num_scheduled_token}, "
-                       f"spec_decode_tokens={spec_tokens}, "
-                       f"num_computed_tokens_after={request.num_computed_tokens + num_scheduled_token}")
             request.num_computed_tokens += num_scheduled_token
 
             # NOTE: _free_encoder_inputs relies on num_computed_tokens, which
@@ -1005,6 +996,7 @@ def update_from_output(
         pooler_outputs = model_runner_output.pooler_output
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
+        is_empty_draft_tokens = model_runner_output.is_empty_draft_tokens
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
@@ -1047,23 +1039,16 @@ def update_from_output(
                 sampled_token_ids[req_index] if sampled_token_ids else []
             )
 
+            req_is_empty_draft_tokens = (
+                is_empty_draft_tokens[req_index] if is_empty_draft_tokens else False
+            )
             scheduled_spec_token_ids = (
                 scheduler_output.scheduled_spec_decode_tokens.get(req_id)
             )
             if scheduled_spec_token_ids:
                 num_draft_tokens = len(scheduled_spec_token_ids)
                 num_accepted = len(generated_token_ids) - 1
                 num_rejected = num_draft_tokens - num_accepted
-                # DEBUG LOG: Track scheduler adjustment
-                logger.info(f"[DEBUG-SCHED] Adjusting in update_from_output: "
-                           f"req_id={req_id}, "
-                           f"scheduled_spec_token_ids={scheduled_spec_token_ids}, "
-                           f"num_draft_tokens={num_draft_tokens}, "
-                           f"generated_token_ids_len={len(generated_token_ids)}, "
-                           f"num_accepted={num_accepted}, "
-                           f"num_rejected={num_rejected}, "
-                           f"num_computed_tokens_before={request.num_computed_tokens}, "
-                           f"num_computed_tokens_after={request.num_computed_tokens - num_rejected if request.num_computed_tokens > 0 else request.num_computed_tokens}")
                 # num_computed_tokens represents the number of tokens
                 # processed in the current step, considering scheduled
                 # tokens and rejections. If some tokens are rejected,
@@ -1090,6 +1075,10 @@ def update_from_output(
             # Check for stop and update request status.
             # logger.info(f"In Scheduler::_update_request_with_output inside loop")
             # from fpdb import ForkedPdb; ForkedPdb().set_trace()
+
+            if req_is_empty_draft_tokens:
+                request.spec_token_ids = []
+
             if new_token_ids:
                 new_token_ids, stopped = self._update_request_with_output(
                     request, new_token_ids
diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py
@@ -181,6 +181,9 @@ class ModelRunnerOutput:
     # req_id -> num_nans_in_logits
     num_nans_in_logits: dict[str, int] | None = None
 
+    # [num_reqs]
+    is_empty_draft_tokens: list[bool] | None = None
+
 
 # ModelRunnerOutput wrapper for async scheduling.
 class AsyncModelRunnerOutput(ABC):
diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py
@@ -28,7 +28,7 @@
 @support_torch_compile(
     dynamic_arg_dims={
         "num_tokens_no_spec": 0,
-        "token_ids_gpu": [0, 1],
+        "token_ids_gpu": 0,
         "combined_mask": 0,
     }
 )
@@ -196,7 +196,7 @@ def _find_first_and_extract_all_n_parallel(
         results = torch.where(
             has_any_match.unsqueeze(1),
             extracted_sequences,
-            torch.full_like(extracted_sequences, 0),   # TODO:(patchy): Use -1 instead of 0.
+            torch.full_like(extracted_sequences, 0),
         )
 
         return results
@@ -248,7 +248,7 @@ def forward(
         mask = combined_mask.unsqueeze(1).expand(-1, self.k)
         draft_tokens = torch.where(mask, results, draft_tokens)
 
-        is_empty_draft_tokens = (draft_tokens == 0).all(dim=1)   # TODO:(patchy): Use -1 instead of 0.
+        is_empty_draft_tokens = (draft_tokens == 0).all(dim=1)
 
         return draft_tokens, is_empty_draft_tokens
 
@@ -296,6 +296,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device, runner=None):
         self.kernel.to(device)
         self.kernel.eval()
         max_batch_size = vllm_config.scheduler_config.max_num_seqs
+
+        # TODO(patchy): Remove this buffer, use
+        # token_ids_gpu_tensor in gpu_model_runner.py instead.
         self.backup_next_token_ids = CpuGpuBuffer(
             max_batch_size,
             dtype=torch.int32,
@@ -309,7 +312,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device, runner=None):
     def _dummy_run(self):
         token_ids, num_tokens, sampled_flags, valid_mask = self._generate_dummy_data(
             batch_size=self.max_num_seqs,
-            max_seq_len=min(self.max_model_len, 1024),
+            max_seq_len=self.max_model_len,
             vocab_size=self.vocab_size,
             pattern_len=self.k,
             repetition_rate=0.5,
@@ -354,33 +357,18 @@ def _generate_dummy_data(
             valid_mask: [batch_size] bool tensor
         """
         # Generate random token IDs
-        token_ids = torch.randint(
-            0, vocab_size, (batch_size, max_seq_len), dtype=torch.int32, device=device
+        token_ids = torch.zeros(
+            batch_size,
+            max_seq_len,
+            dtype=torch.int32,
+            device=device,
         )
 
         # Generate random sequence lengths
-        min_len = max(pattern_len * 2 + 3, max_seq_len // 2)
         num_tokens = torch.randint(
-            min_len, max_seq_len, (batch_size,), dtype=torch.int32, device=device
+            pattern_len, max_seq_len, (batch_size,), dtype=torch.int32, device=device
         )
 
-        # Inject n-gram repetitions using the tail pattern of each sequence
-        for i in range(batch_size):
-            seq_len = num_tokens[i].item()
-            if seq_len > pattern_len * 2:
-                # Pattern is the last pattern_len tokens of the valid sequence
-                src_pos = seq_len - pattern_len
-                num_reps = int(seq_len * repetition_rate / pattern_len)
-                for _ in range(num_reps):
-                    # Place the copied tail pattern somewhere before the tail
-                    tgt_pos = torch.randint(0, seq_len - pattern_len, (1,)).item()
-                    if tgt_pos == src_pos:
-                        continue
-
-                    token_ids[i, tgt_pos : tgt_pos + pattern_len] = token_ids[
-                        i, src_pos : src_pos + pattern_len
-                    ].clone()
-
         # All sequences have sampled tokens and are valid
         sampled_flags = torch.ones(batch_size, dtype=torch.bool, device=device)
         valid_mask = torch.ones(batch_size, dtype=torch.bool, device=device)
@@ -401,10 +389,7 @@ def propose(
 
         with set_forward_context(None, self.vllm_config):
             combined_mask = (
-                sampled_flags
-                & valid_mask
-                & (num_tokens_no_spec < self.max_model_len)
-                & (num_tokens_no_spec >= self.min_n)
+                sampled_flags & valid_mask & (num_tokens_no_spec >= self.min_n)
             )
 
             draft_tokens, is_empty_draft_tokens = self.kernel(
@@ -415,36 +400,6 @@ def propose(
 
             return draft_tokens, is_empty_draft_tokens
 
-    def prepare_next_token_ids_cpu(
-        self,
-        sampled_token_ids: list[np.ndarray],
-        requests: dict[str, CachedRequestState],
-        gpu_input_batch: InputBatch,
-        num_scheduled_tokens: dict[str, int],
-    ) -> torch.Tensor:
-        """
-        This function is used to prepare the inputs for speculative decoding.
-        It calculates the next token ids for each request based on the sampled
-        token ids from the CPU. If a request has no sampled token ids (e.g.,
-        during the initial decoding steps), it falls back to using the request
-        state to get the next token id.
-        """
-        req_ids = gpu_input_batch.req_ids
-        next_token_ids: list[int] = []
-        for i, token_ids in enumerate(sampled_token_ids):
-            if token_ids.shape[0] > 0:
-                # Common case.
-                next_token_id = token_ids[-1]
-            else:
-                # Partial prefill (rare case).
-                # Get the next token id from the request state.
-                req_id = req_ids[i]
-                req_state = requests[req_id]
-                seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]
-                next_token_id = req_state.get_token_id(seq_len)
-            next_token_ids.append(next_token_id)
-        return torch.tensor(next_token_ids, dtype=torch.int32, device=self.device)
-
     def prepare_next_token_ids_padded(
         self,
         common_attn_metadata: CommonAttentionMetadata,
@@ -463,8 +418,6 @@ def prepare_next_token_ids_padded(
         This function must use device functions to operate on the inputs, and
         should not introduce any blocking CPU-GPU synchronization.
         """
-        # TODO(Ben): Combine this into a custom fused kernel
-        # Precompute get_token_id for when there is no valid next token
         num_reqs = gpu_input_batch.num_reqs
         # Batch convert seq_lens to avoid multiple .item() calls
         seq_lens_list = common_attn_metadata.seq_lens_cpu[:num_reqs].tolist()
diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py
@@ -114,9 +114,6 @@ def __init__(
             pin_memory=False,
         )
         self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
-        self.token_ids_gpu_tensor = torch.zeros(
-            max_num_reqs, max_model_len, dtype=torch.int32, device=device
-        )
         self.is_token_ids_tensor = torch.zeros(
             (max_num_reqs, max_model_len), device="cpu", dtype=bool, pin_memory=False
         )
@@ -127,9 +124,6 @@ def __init__(
         self.req_prompt_embeds: dict[int, torch.Tensor] = {}
         self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)
         self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
-        self.num_tokens_no_spec_gpu = torch.zeros(
-            max_num_reqs, dtype=torch.int32, device=device
-        )
         self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
         self.num_computed_tokens_cpu_tensor = torch.zeros(
             (max_num_reqs,),
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
