[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>
rishitdholakia13 committed
Jul 14, 2026 at 14:55 UTC
32e632dfeb0860d11332754e641035ddd755f89f
3 files changed
+7
-20
vllm/v1/sample/rejection_sampler.py
+3
-7
@@ -292,11 +292,8 @@ class RejectionSampler(nn.Module):
292
any_penalties_or_bad_words = (
293
sampling_metadata.bad_words_token_ids or has_penalties
294
)
295
- holder = sampling_metadata.thinking_budget_state_holder
296
- needs_thinking = holder is not None and holder.has_tracked_requests()
297
-
295
output_token_ids = sampling_metadata.output_token_ids
299
- if any_penalties_or_bad_words or needs_thinking:
296
+ if any_penalties_or_bad_words:
297
output_token_ids = self._combine_outputs_with_spec_tokens(
298
output_token_ids,
299
sampling_metadata.spec_token_ids,
@@ -305,9 +302,7 @@ class RejectionSampler(nn.Module):
302
# Calculate indices of target logits.
303
repeat_indices: torch.Tensor | None = None
304
need_repeat_indices = (
308
- sampling_metadata.allowed_token_ids_mask is not None
309
- or has_penalties
310
- or needs_thinking
305
+ sampling_metadata.allowed_token_ids_mask is not None or has_penalties
306
)
307
if need_repeat_indices:
308
num_requests = len(metadata.num_draft_tokens)
@@ -337,6 +332,7 @@ class RejectionSampler(nn.Module):
332
logits = processor.apply_with_spec_decode(
333
logits, metadata.num_draft_tokens
334
)
335
+ holder = sampling_metadata.thinking_budget_state_holder
336
if holder is not None and holder.has_tracked_requests():
337
logits = holder.apply_to_logits(
338
logits,
vllm/v1/sample/sampler.py
+4
-7
@@ -378,13 +378,8 @@ class Sampler(nn.Module):
378
any_penalties_or_bad_words = (
379
bool(bad_words_token_ids) or not sampling_metadata.no_penalties
380
)
381
- holder = sampling_metadata.thinking_budget_state_holder
382
- needs_thinking_combine = holder is not None and holder.has_tracked_requests()
383
-
381
output_token_ids = sampling_metadata.output_token_ids
385
- if predict_bonus_token and (
386
- any_penalties_or_bad_words or needs_thinking_combine
387
- ):
382
+ if predict_bonus_token and any_penalties_or_bad_words:
383
# Combine base outputs with spec tokens when speculative decoding
384
# is enabled.
385
output_token_ids = self._combine_outputs_with_spec_tokens(
@@ -406,9 +401,11 @@ class Sampler(nn.Module):
401
402
# Apply penalties (e.g., freq_penalties).
403
logits = self.apply_penalties(logits, sampling_metadata, output_token_ids)
404
+ holder = sampling_metadata.thinking_budget_state_holder
405
if holder is not None and holder.has_tracked_requests():
406
+ # Committed outputs only; spec drafts live in ``spec_token_ids``.
407
holder.update_state(
411
- output_token_ids,
408
+ sampling_metadata.output_token_ids,
409
sampling_metadata.spec_token_ids,
410
repeat_indices=None,
411
)
vllm/v1/sample/thinking_budget_state.py
-6
@@ -145,12 +145,6 @@ class ThinkingBudgetStateHolder:
145
state["spec_token_ids"] = []
146
state["in_spec_mode"] = self.in_spec_mode
147
state["force_index"] = []
148
- if len(state["output_tok_ids"]) > 0:
149
- spec_len = len(state["spec_token_ids"])
150
- # Only strip draft suffix when there are spec tokens; ``[:-0]`` would
151
- # clear the whole list (Python treats stop index 0 as "up to empty").
152
- if spec_len > 0 and len(state["output_tok_ids"]) >= spec_len:
153
- state["output_tok_ids"] = state["output_tok_ids"][:-spec_len]
148
self._update_think_state(state)
149
150
def apply_to_logits(