Problem
ThinkingBudgetStateHolder._apply_forcing_to_logits computes the row to force as:
mask_idx = self.cu_num_tokens[seq_idx] + force_idx
and then bounds it only by self._mask_capacity and logits.shape[0]. Neither
bound relates to the request that owns seq_idx.
In speculative decoding a request occupies max(1, len(spec_token_ids)) rows of
the target-logits tensor. If force_idx ever exceeds that span, mask_idx
lands inside the next request's rows and the end-of-thinking token is forced
into a different sequence's logits — silently, with no error. If the overshoot
happens on the last request in the batch it is instead clipped away and the
force is dropped.
Why this matters now
Two separate bugs of exactly this shape were found and fixed during the
reasoning_answer_reserve work (#28):
- the budgeted same-step revisit branch recomputed
force_index from the
thinking budget alone and could select [spec_len], which is one past the
request's last target row;
- the same branch could overwrite a marker-continuation position with row 0.
Both fixes work by keeping force_index in range. Nothing prevents the next
change from reintroducing the same class of fault, and the failure is invisible:
wrong tokens in someone else's sequence rather than an exception.
Proposal
Bound mask_idx to the owning request's row span rather than to the tensor:
row_span = 1 if predict_bonus_token else max(1, len(state["spec_token_ids"]))
limit = self.cu_num_tokens[seq_idx] + row_span
and treat an out-of-span force_idx as a bug — assert, or log once and skip,
rather than writing.
This touches the shared thinking_token_budget path, so it wants its own
change with its own regression coverage rather than riding along with a feature.
Acceptance
- A forced row outside the owning request's span is rejected rather than written.
- Regression test with at least two requests of differing draft lengths that
fails if a cross-request write is reintroduced.
- Existing
thinking_token_budget and reasoning_answer_reserve behaviour
unchanged.
Duplicate check
Problem
ThinkingBudgetStateHolder._apply_forcing_to_logitscomputes the row to force as:and then bounds it only by
self._mask_capacityandlogits.shape[0]. Neitherbound relates to the request that owns
seq_idx.In speculative decoding a request occupies
max(1, len(spec_token_ids))rows ofthe target-logits tensor. If
force_idxever exceeds that span,mask_idxlands inside the next request's rows and the end-of-thinking token is forced
into a different sequence's logits — silently, with no error. If the overshoot
happens on the last request in the batch it is instead clipped away and the
force is dropped.
Why this matters now
Two separate bugs of exactly this shape were found and fixed during the
reasoning_answer_reservework (#28):force_indexfrom thethinking budget alone and could select
[spec_len], which is one past therequest's last target row;
Both fixes work by keeping
force_indexin range. Nothing prevents the nextchange from reintroducing the same class of fault, and the failure is invisible:
wrong tokens in someone else's sequence rather than an exception.
Proposal
Bound
mask_idxto the owning request's row span rather than to the tensor:and treat an out-of-span
force_idxas a bug — assert, or log once and skip,rather than writing.
This touches the shared
thinking_token_budgetpath, so it wants its ownchange with its own regression coverage rather than riding along with a feature.
Acceptance
fails if a cross-request write is reintroduced.
thinking_token_budgetandreasoning_answer_reservebehaviourunchanged.
Duplicate check
thinking_token_budgetPRs([Model Runner V2] Support thinking_token_budget for mrv2 vllm-project/vllm#44510, [Feat] Support thinking_token_budget in Model Runner V2 vllm-project/vllm#46727, [Bugfix][Sampling] Clear empty side on thinking-budget asymmetric SWAP vllm-project/vllm#49613, [Bugfix][Reasoning] Properly detect reasoning end when using thinking_token_budget vllm-project/vllm#43210) are MRv2 support and unrelated bugfixes.