# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Per-batch thinking token budget state; applied after penalties at sample time.""" from typing import TYPE_CHECKING, Any import torch from vllm.platforms import current_platform from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.sample.logits_processor.interface import ( BatchUpdate, MoveDirectionality, ) if TYPE_CHECKING: from vllm.config.reasoning import ReasoningConfig def maybe_create_thinking_budget_state_holder( reasoning_config: "ReasoningConfig | None", max_num_seqs: int, num_spec_tokens: int, device: torch.device, is_pin_memory: bool, ) -> "ThinkingBudgetStateHolder | None": if reasoning_config is None: return None return ThinkingBudgetStateHolder( reasoning_config, max_num_seqs, num_spec_tokens, device, is_pin_memory ) class ThinkingBudgetStateHolder: """Tracks thinking sections and forces end tokens when budget is exceeded.""" think_start_token_ids: list[int] think_end_token_ids: list[int] def __init__( self, reasoning_config: "ReasoningConfig | None", max_num_seqs: int, num_spec_tokens: int, device: torch.device, is_pin_memory: bool, ): _ = is_pin_memory # API parity with logits processors max_num_reqs = max_num_seqs self.in_spec_mode = num_spec_tokens > 0 self.num_spec_tokens = num_spec_tokens # No separate enable flag: a non-``None`` ``reasoning_config`` is the switch. self.is_enabled = reasoning_config is not None if reasoning_config is None: self.think_start_token_ids = [] self.think_end_token_ids = [] else: rs = reasoning_config.reasoning_start_token_ids re = reasoning_config.reasoning_end_token_ids self.think_start_token_ids = rs if rs else [] self.think_end_token_ids = re if re else [] self.device = device self._state: dict[int, dict[str, Any]] = {} self.cu_num_tokens: dict[int, int] = {} if self.num_spec_tokens > 0: self._mask_capacity = max_num_reqs * (self.num_spec_tokens + 1) else: self._mask_capacity = max_num_reqs def has_tracked_requests(self) -> bool: """True when ``sync_batch`` has state for a ``thinking_token_budget`` row. Used to decide whether sampling needs output-token rows and spec combining; distinct from merely having a holder instance (reasoning may be on with no budgeted requests in this batch). """ return bool(self._state) def sync_batch(self, batch_update: BatchUpdate | None) -> None: """Add/remove/move per-request state only (no _update_think_state).""" if not self.is_enabled or not batch_update: return for index in batch_update.removed: self._state.pop(index, None) for index, params, prompt_tok_ids, output_tok_ids in batch_update.added: thinking_token_budget = params.thinking_token_budget if thinking_token_budget is not None: self._state[index] = self._init_state_entry( prompt_tok_ids, thinking_token_budget ) self._state[index]["output_tok_ids"] = output_tok_ids self._state[index]["spec_token_ids"] = [] else: self._state.pop(index, None) for i1, i2, direction in batch_update.moved: if direction == MoveDirectionality.SWAP: state1 = self._state.get(i1) state2 = self._state.get(i2) if state1 is not None: self._state[i2] = state1 if state2 is not None: self._state[i1] = state2 else: state = self._state.pop(i1, None) if state is not None: self._state[i2] = state def update_state( self, output_token_ids: list[list[int]], spec_token_ids: list[list[int]] | None, repeat_indices: torch.Tensor | None = None, ) -> None: """Refresh output/spec from sampling rows and recompute think state.""" if not self.is_enabled or not self._state: return spec_lists = spec_token_ids or [] last_row_for_req: dict[int, int] | None = None if repeat_indices is not None: last_row_for_req = {} rpt = repeat_indices.cpu().tolist() for batch_row, req_i in enumerate(rpt): last_row_for_req[req_i] = batch_row for seq_idx, state in list(self._state.items()): if last_row_for_req is not None: output_row: int | None = last_row_for_req.get(seq_idx) if output_row is None or output_row >= len(output_token_ids): continue state["output_tok_ids"] = output_token_ids[output_row] elif seq_idx >= len(output_token_ids): continue else: state["output_tok_ids"] = output_token_ids[seq_idx] if seq_idx < len(spec_lists): state["spec_token_ids"] = list(spec_lists[seq_idx]) else: 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( self, logits: torch.Tensor, predict_bonus_token: bool, spec_token_ids: list[list[int]] | None, ) -> torch.Tensor: """Mask and bump logits for forced end-of-thinking tokens.""" if not self.is_enabled or not self._state: return logits spec_lists = spec_token_ids or [] return self._apply_forcing_to_logits(logits, predict_bonus_token, spec_lists) @staticmethod def _find_last_sequence_index(target_list: list[int], token_ids: list[int]) -> int: if not token_ids: return -1 for i in range(len(target_list) - len(token_ids), -1, -1): if target_list[i : i + len(token_ids)] == token_ids: return i return -1 def _init_state_entry( self, prompt_tok_ids: list[int] | None, thinking_token_budget: int ) -> dict[str, Any]: if prompt_tok_ids is None: last_start = -1 last_end = -1 in_think = False think_count = 0 start_thinking = -1 countdown = thinking_token_budget continue_thinking = False in_end = False else: start_thinking = -1 countdown = thinking_token_budget continue_thinking = False in_end = False last_start = self._find_last_sequence_index( prompt_tok_ids, self.think_start_token_ids ) last_end = self._find_last_sequence_index( prompt_tok_ids, self.think_end_token_ids ) in_think = last_start > last_end # load metrics such as think count, start thinking # if request is in thinking mode, already if in_think: think_count = len(prompt_tok_ids) - ( last_start + len(self.think_start_token_ids) ) start_thinking = len(prompt_tok_ids) - think_count - 1 countdown -= think_count continue_thinking = True # check if the token is exhausted within prompt token_exhausted = thinking_token_budget - think_count in_end = token_exhausted <= 0 else: think_count = 0 return { "in_think": in_think, "in_end": in_end, "check_count_down": countdown, "think_count": think_count, "end_count": 0, "prompt_tok_ids": prompt_tok_ids, "output_tok_ids": [], "thinking_token_budget": thinking_token_budget, "prev_output_length": 0, "spec_token_ids": [], "force_index": [], "start_thinking": start_thinking, "end_thinking": -1, "in_spec_mode": False, "bonus_token_forced": False, "continue_thinking": continue_thinking, "scan_offset": 0, } def _update_think_state(self, state: dict[str, Any]) -> None: if state.get("thinking_token_budget", -1) == -1: return if len(self.think_end_token_ids) == 0: state["thinking_token_budget"] = -1 state["in_end"] = False state["force_index"] = [] return if state["start_thinking"] == -1: scan_offset = state.get("scan_offset", 0) output_slice = state.get("output_tok_ids", [])[scan_offset:] start_thinking = self._find_last_sequence_index( output_slice, self.think_start_token_ids ) if start_thinking >= 0: start_thinking += scan_offset state["start_thinking"] = start_thinking if state["end_thinking"] == -1: scan_offset = state.get("scan_offset", 0) output_slice = state.get("output_tok_ids", [])[scan_offset:] end_thinking = self._find_last_sequence_index( output_slice, self.think_end_token_ids ) if end_thinking >= 0: end_thinking += scan_offset state["end_thinking"] = end_thinking if ( not state.get("in_end", False) and state["start_thinking"] >= 0 and state["end_thinking"] >= 0 and state["end_thinking"] > state["start_thinking"] and not state.get("continue_thinking", False) ): state["in_think"] = False state["think_count"] = 0 state["continue_thinking"] = False state["start_thinking"] = -1 state["end_thinking"] = -1 state["scan_offset"] = len(state.get("output_tok_ids", [])) state["check_count_down"] = state["thinking_token_budget"] return if state["start_thinking"] == -1: return if state["continue_thinking"]: sampled_tokens_from_previous_step = len( state.get("output_tok_ids", []) ) - state.get("prev_output_length", 0) else: if state["prev_output_length"] == 0: sampled_tokens_from_previous_step = len( state.get("output_tok_ids", []) ) - len(self.think_start_token_ids) else: sampled_tokens_from_previous_step = ( len(state.get("output_tok_ids", [])) - state["prev_output_length"] ) current_step_countdown = ( state["check_count_down"] - sampled_tokens_from_previous_step ) predicted_countdown = current_step_countdown - len(state["spec_token_ids"]) - 1 # We only proceed further if we have counted down the thinking budget # to 0 or less and when we are in the "in think" mode. # Exception: when continue_thinking=True and a natural is # detected (end_thinking != -1), fall through to handle the exit — # even if the budget hasn't expired yet. For continue_thinking=False, # the early natural-end detection block above already handles it. natural_end_with_continue = ( state.get("continue_thinking", False) and state["end_thinking"] != -1 ) if ( not state.get("in_end", False) and predicted_countdown >= 0 and state["start_thinking"] > -1 and not natural_end_with_continue ): state["check_count_down"] = current_step_countdown state["prev_output_length"] = len(state.get("output_tok_ids", [])) return output = state.get("output_tok_ids", []) if not output: # When in_end was set at init (budget=0, prompt already in think), # we must force the first generated token to be the end token; # otherwise apply() sees in_end=True but force_index=[] and # allows an extra thinking token. if state.get("in_end", False): state["force_index"] = [0] return # Track previous output length for incremental processing prev_length = state.get("prev_output_length", 0) current_length = len(output) if current_length <= prev_length: if state.get("in_end", False): remaining_budget = state["thinking_token_budget"] - state["think_count"] spec_len = len(state["spec_token_ids"]) if spec_len > 0: if 0 < remaining_budget < spec_len: state["force_index"] = [remaining_budget] elif remaining_budget <= 0: state["force_index"] = [0] else: state["force_index"] = [spec_len] else: state["force_index"] = [0] return state["prev_output_length"] = current_length start_len = len(self.think_start_token_ids) absolute_start_pos = state["start_thinking"] if state["continue_thinking"] and state["end_thinking"] > -1: absolute_end_pos = state["end_thinking"] + len( state.get("prompt_tok_ids") or [] ) else: absolute_end_pos = state["end_thinking"] # Update state based on recent sequences # This is the case where we are in end mode, but the rejection sampler # rejected a token before the end token, # so we need to go back to think mode and wait for the next end token # eg with 999: [2,4,5,999] -> [3,-1,-1,-1] if state["in_end"] and state["end_count"] == 0: new_tokens = output[prev_length:] stopping_thinking = ( self.think_end_token_ids[state["end_count"]] in new_tokens ) if not stopping_thinking: state["in_think"] = True state["in_end"] = False state["end_count"] = 0 state["bonus_token_forced"] = False if not state["in_end"]: if absolute_start_pos >= 0 and absolute_end_pos >= 0: # Case: ......... - entering think mode if absolute_start_pos > absolute_end_pos: new_think_count = current_length - (absolute_start_pos + start_len) state["in_think"] = True state["think_count"] = new_think_count else: # Case: ......... - exiting think mode state["in_think"] = False state["think_count"] = 0 state["continue_thinking"] = False state["start_thinking"] = -1 state["end_thinking"] = -1 state["scan_offset"] = len(state.get("output_tok_ids", [])) elif absolute_start_pos >= 0 and not state["continue_thinking"]: # Found think start - entering think mode new_think_count = current_length - (absolute_start_pos + start_len) state["in_think"] = True state["think_count"] = new_think_count elif absolute_end_pos >= 0: # Found think end - exiting think mode state["in_think"] = False state["think_count"] = 0 state["continue_thinking"] = False state["start_thinking"] = -1 state["end_thinking"] = -1 state["scan_offset"] = len(state.get("output_tok_ids", [])) elif state["in_think"]: # Continue thinking mode, increment count by new tokens prompt_tok_ids = state.get("prompt_tok_ids") or [] think_tokens_in_prompt = len(prompt_tok_ids) - ( absolute_start_pos + start_len ) state["think_count"] = ( len(state["output_tok_ids"]) + think_tokens_in_prompt ) if state["in_think"]: remaining_budget = max( 0, state["thinking_token_budget"] - state["think_count"] ) state["check_count_down"] = remaining_budget else: state["check_count_down"] = state["thinking_token_budget"] total_thinking_tokens = ( state["think_count"] + len(state["spec_token_ids"]) + 1 ) # Check if need to transition to end mode # If we have more thinking tokens than the budget, # we need to transition to end mode if ( state["in_think"] and total_thinking_tokens > state["thinking_token_budget"] ): # Calculate force_index: position within spec_token_ids where # forcing starts. If we're already over budget without spec # tokens, force from position 0. Force from the position # where budget is exceeded. state["in_think"] = False state["in_end"] = True state["end_count"] = 0 state["check_count_down"] = state["thinking_token_budget"] remaining_budget = state["thinking_token_budget"] - state["think_count"] spec_len = len(state["spec_token_ids"]) if 0 < remaining_budget < spec_len: state["force_index"] = [remaining_budget] elif remaining_budget <= 0: state["force_index"] = [0] else: # remaining_budget >= spec_len: all spec tokens are within # budget; force the bonus token position state["force_index"] = [len(state["spec_token_ids"])] else: state["force_index"] = [] if len(state["spec_token_ids"]) > 0: for i, token_id in enumerate(state["spec_token_ids"]): if state["end_count"] + 1 < len(self.think_end_token_ids): if token_id == self.think_end_token_ids[state["end_count"] + 1]: state["end_count"] += 1 else: state["end_count"] += 1 state["force_index"] = [i] break else: state["end_count"] += 1 if len(state["force_index"]) == 0: state["end_count"] += 1 state["force_index"] = [len(state["spec_token_ids"])] else: state["end_count"] += 1 state["force_index"] = [0] if state["end_count"] >= len(self.think_end_token_ids): state.update( { "in_end": False, "end_count": 0, "check_count_down": state["thinking_token_budget"], "start_thinking": -1, "end_thinking": -1, "think_count": 0, "continue_thinking": False, "scan_offset": len(state.get("output_tok_ids", [])), } ) def _apply_forcing_to_logits( self, logits: torch.Tensor, predict_bonus_token: bool, spec_token_ids_for_layout: list[list[int]], ) -> torch.Tensor: cumulative_total = 0 self.cu_num_tokens.clear() n_layout = len(spec_token_ids_for_layout) if self._state: n_layout = max(n_layout, max(self._state.keys()) + 1) for index in range(n_layout): self.cu_num_tokens[index] = cumulative_total spec_tokens = ( spec_token_ids_for_layout[index] if index < len(spec_token_ids_for_layout) else [] ) if self.in_spec_mode: cumulative_total += len(spec_tokens) if not predict_bonus_token else 1 else: cumulative_total += 1 # Build the active index / forced-token lists entirely on CPU so we # avoid per-iteration scalar sync writes to GPU tensors. active_indices_cpu: list[int] = [] force_tokens_cpu: list[int] = [] for seq_idx in sorted(self._state.keys()): if seq_idx not in self.cu_num_tokens: continue state = self._state[seq_idx] if state.get("in_end", False): # logits processor in spec mode are called twice # once for bonus token logits and # second time for the target logits # in case the force index is bonus token index # we change the force index to 0 if predict_bonus_token: if state.get("force_index") and state["force_index"][0] < len( state["spec_token_ids"] ): continue else: state["force_index"] = [0] # continue enforcing the end thinking tokens if state["end_count"] > 0: state["bonus_token_forced"] = False if state and not state["bonus_token_forced"]: force_index = state.get("force_index", []) if len(force_index) == 0: continue end_count = state.get("end_count", 0) for force_idx in force_index: if end_count < len(self.think_end_token_ids): mask_idx = self.cu_num_tokens[seq_idx] + force_idx if ( mask_idx < self._mask_capacity and mask_idx < logits.shape[0] ): active_indices_cpu.append(mask_idx) force_tokens_cpu.append( self.think_end_token_ids[end_count] ) if predict_bonus_token: if state["end_count"] > 0: state["bonus_token_forced"] = False state["force_index"] = [] else: state["bonus_token_forced"] = True if active_indices_cpu: device = logits.device if current_platform.is_rocm() and logits.is_contiguous(): # Flattened index_fill avoids ROCm faults seen with 2-D # advanced-indexing writes on the thinking-budget path. vocab_size = logits.shape[1] flat_indices_cpu = [ row * vocab_size + token for row, token in zip(active_indices_cpu, force_tokens_cpu) ] flat_indices = async_tensor_h2d( flat_indices_cpu, dtype=torch.long, device=device ) logits.view(-1).index_fill_(0, flat_indices, 1e9) elif current_platform.is_rocm(): fill = logits.new_tensor(1e9) for row, token in zip(active_indices_cpu, force_tokens_cpu): logits[row, token] = fill else: active_indices = async_tensor_h2d( active_indices_cpu, dtype=torch.long, device=device ) force_tokens = async_tensor_h2d( force_tokens_cpu, dtype=torch.long, device=device ) # Avoid CPU->GPU sync. fill = logits.new_full((len(active_indices_cpu),), 1e9) logits.index_put_((active_indices, force_tokens), fill) return logits