From 32e632dfeb0860d11332754e641035ddd755f89f Mon Sep 17 00:00:00 2001 From: rishitdholakia13 <123388671+rishitdholakia13@users.noreply.github.com> Date: Tue, 14 Jul 2026 14:55:40 -0400 Subject: [PATCH] [Reasoning] Optimize TPOT for thinking budget when used with speculative decoding (#46662) Signed-off-by: rishitdholakia13 Co-authored-by: Cursor --- vllm/v1/sample/rejection_sampler.py | 10 +++------- vllm/v1/sample/sampler.py | 11 ++++------- vllm/v1/sample/thinking_budget_state.py | 6 ------ 3 files changed, 7 insertions(+), 20 deletions(-) diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 1324191be7..6930921383 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/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, diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index bb20432a08..ac8512f22c 100644 --- a/vllm/v1/sample/sampler.py +++ b/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, ) diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py index f19f612f48..efac011177 100644 --- a/vllm/v1/sample/thinking_budget_state.py +++ b/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(