[Reasoning] Optimize TPOT for thinking budget when used with speculative decoding (#46662)

Signed-off-by: rishitdholakia13 <rishit+github@cohere.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

vllm/v1/sample/rejection_sampler.py

@@ -292,11 +292,8 @@ class RejectionSampler(nn.Module):
         any_penalties_or_bad_words = (
             sampling_metadata.bad_words_token_ids or has_penalties
         )
-        holder = sampling_metadata.thinking_budget_state_holder
-        needs_thinking = holder is not None and holder.has_tracked_requests()
-
         output_token_ids = sampling_metadata.output_token_ids
-        if any_penalties_or_bad_words or needs_thinking:
+        if any_penalties_or_bad_words:
             output_token_ids = self._combine_outputs_with_spec_tokens(
                 output_token_ids,
                 sampling_metadata.spec_token_ids,
@@ -305,9 +302,7 @@ class RejectionSampler(nn.Module):
         # Calculate indices of target logits.
         repeat_indices: torch.Tensor | None = None
         need_repeat_indices = (
-            sampling_metadata.allowed_token_ids_mask is not None
-            or has_penalties
-            or needs_thinking
+            sampling_metadata.allowed_token_ids_mask is not None or has_penalties
         )
         if need_repeat_indices:
             num_requests = len(metadata.num_draft_tokens)
@@ -337,6 +332,7 @@ class RejectionSampler(nn.Module):
                 logits = processor.apply_with_spec_decode(
                     logits, metadata.num_draft_tokens
                 )
+        holder = sampling_metadata.thinking_budget_state_holder
         if holder is not None and holder.has_tracked_requests():
             logits = holder.apply_to_logits(
                 logits,

vllm/v1/sample/sampler.py

@@ -378,13 +378,8 @@ class Sampler(nn.Module):
         any_penalties_or_bad_words = (
             bool(bad_words_token_ids) or not sampling_metadata.no_penalties
         )
-        holder = sampling_metadata.thinking_budget_state_holder
-        needs_thinking_combine = holder is not None and holder.has_tracked_requests()
-
         output_token_ids = sampling_metadata.output_token_ids
-        if predict_bonus_token and (
-            any_penalties_or_bad_words or needs_thinking_combine
-        ):
+        if predict_bonus_token and any_penalties_or_bad_words:
             # Combine base outputs with spec tokens when speculative decoding
             # is enabled.
             output_token_ids = self._combine_outputs_with_spec_tokens(
@@ -406,9 +401,11 @@ class Sampler(nn.Module):
 
         # Apply penalties (e.g., freq_penalties).
         logits = self.apply_penalties(logits, sampling_metadata, output_token_ids)
+        holder = sampling_metadata.thinking_budget_state_holder
         if holder is not None and holder.has_tracked_requests():
+            # Committed outputs only; spec drafts live in ``spec_token_ids``.
             holder.update_state(
-                output_token_ids,
+                sampling_metadata.output_token_ids,
                 sampling_metadata.spec_token_ids,
                 repeat_indices=None,
             )

vllm/v1/sample/thinking_budget_state.py

@@ -145,12 +145,6 @@ class ThinkingBudgetStateHolder:
                 state["spec_token_ids"] = []
             state["in_spec_mode"] = self.in_spec_mode
             state["force_index"] = []
-            if len(state["output_tok_ids"]) > 0:
-                spec_len = len(state["spec_token_ids"])
-                # Only strip draft suffix when there are spec tokens; ``[:-0]`` would
-                # clear the whole list (Python treats stop index 0 as "up to empty").
-                if spec_len > 0 and len(state["output_tok_ids"]) >= spec_len:
-                    state["output_tok_ids"] = state["output_tok_ids"][:-spec_len]
             self._update_think_state(state)
 
     def apply_to_logits(