Skip to content

Bound forced-token row index to the owning request in ThinkingBudgetStateHolder #34

Description

@randomvariable

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    area/ideaOptimization idea candidate for evaluation

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions