vLLM V1 LogitsProcessor that enforces a per-request reasoning (thinking) token budget.

Models like Nemotron-Nano use <think>/</think> delimiters (same convention as DeepSeek-R1). This processor monitors generated tokens, counts those inside the thinking block, and forces end tokens when the budget is reached.

Per-request parameters (via SamplingParams.extra_args):

thinking_budget (int): Maximum number of thinking tokens allowed before forcing the end sequence. Required to activate the processor for a given request.

thinking_budget_grace_period (int, optional): Number of tokens before the budget at which \n and end-token logits start being boosted. Defaults to 10 % of thinking_budget.

think_start_tokens (str, optional): Text that marks the beginning of a thinking block. The processor tokenizes this string at request time. Defaults to "<think>".

think_end_tokens (str, optional): Text to force when the budget is reached. The processor tokenizes this string at request time. Defaults to "\n</think>". Can be set to a custom closing such as "Reached thinking limit.\n</think>".

Usage — offline with vllm.LLM::

llm = LLM(model=model, logits_processors=[ReasoningBudgetLogitsProcessor], …) params = SamplingParams( temperature=0.6, max_tokens=256, extra_args={“thinking_budget”: 64}, ) outputs = llm.generate(prompts, params)

Usage — online with vllm serve::

vllm serve <model>
—logits-processors ’[”…/reasoning_budget_logits_processor:ReasoningBudgetLogitsProcessor”]‘

then per-request via the OpenAI client:

extra_body={“vllm_xargs”: {“thinking_budget”: 64}}

Module Contents

Classes

Name	Description
`ReasoningBudgetLogitsProcessor`	Enforce a per-request thinking-token budget for reasoning models.
`RequestState`	Mutable per-request tracking state.

Data

_DEFAULT_GRACE_RATIO

_GRACE_LOGIT_BOOST

logger

API

class nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor(
    vllm_config: 'VllmConfig',
    device: torch.device,
    is_pin_memory: bool
)

Bases: LogitsProcessor

Enforce a per-request thinking-token budget for reasoning models.

The processor tracks thinking-start / thinking-end boundaries in each request’s output. Once a request’s thinking token count enters the grace window, \n and end-token logits are boosted. At the hard limit all logits except the next forced end token are set to -inf.

neg_inf

newline_ids

list[int] = self._encode(self.tokenizer, '\n')

req_states

dict[int, RequestState] = {}

think_end_detect_ids

list[int] = self._encode(self.tokenizer, '</think>')

think_end_force_ids

list[int] = self._encode(self.tokenizer, '\n</think>')

think_start_ids

list[int] = self._encode(self.tokenizer, '<think>')

tokenizer

= self._load_tokenizer(vllm_config)

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor._apply_grace_boost(
    logits: torch.Tensor,
    batch_idx: int,
    state: nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.RequestState
) -> None

staticmethod

Additively boost newline and end-token logits.

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor._device_tensor(
    data: list,
    dtype: torch.dtype
) -> torch.Tensor

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor._encode(
    tokenizer,
    text: str
) -> list[int]

staticmethod

Encode text without special tokens.

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor._force_end_token(
    logits: torch.Tensor,
    batch_idx: int,
    state: nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.RequestState
) -> None

Set all logits to -inf except the next token in the end sequence. Advances forcing_end_idx each call.

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor._load_tokenizer(
    vllm_config: 'VllmConfig'
)

staticmethod

Obtain a tokenizer from the vLLM config.

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor._new_state(
    params: vllm.SamplingParams,
    prompt_tok_ids: list[int] | None,
    output_tok_ids: list[int]
) -> nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.RequestState | None

Called by process_dict_updates for each newly added request.

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor._prompt_ends_with(
    prompt_tok_ids: list[int] | None,
    pattern: list[int],
    skip_ids: list[int]
) -> bool

staticmethod

Return True if prompt_tok_ids ends with pattern, ignoring any trailing tokens whose ID is in skip_ids (e.g. newlines).

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor._scan_tokens(
    state: nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.RequestState,
    from_idx: int
) -> None

staticmethod

Update state by scanning output_tok_ids[from_idx:].

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor.apply(
    logits: torch.Tensor
) -> torch.Tensor

Modify the logits in place to force thinking-budget exits when the limit is reached.

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor.is_argmax_invariant() -> bool

Return whether this processor preserves argmax behavior (it does not).

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor.update_state(
    batch_update: vllm.v1.sample.logits_processor.BatchUpdate | None
) -> None

Sync per-request thinking state with the batch update from vLLM.

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.ReasoningBudgetLogitsProcessor.validate_params(
    sampling_params: vllm.SamplingParams
)

classmethod

Validate thinking-budget-related extra_args on the provided sampling params.

class nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.RequestState(
    max_thinking_tokens: int,
    grace_period: int,
    detect_start_ids: list[int],
    detect_end_ids: list[int],
    force_end_ids: list[int],
    thinking_token_count: int = 0,
    inside_thinking: bool = False,
    stopped_thinking: bool = False,
    forcing_end_idx: int = -1,
    output_tok_ids: list[int] = list(),
    prev_output_length: int = 0
)

Dataclass

Mutable per-request tracking state.

detect_end_ids

list[int]

detect_start_ids

list[int]

force_end_ids

list[int]

forcing_end_idx

int = -1

grace_period

int

inside_thinking

bool = False

max_thinking_tokens

int

output_tok_ids

list[int] = field(default_factory=list)

prev_output_length

int = 0

stopped_thinking

bool = False

thinking_token_count

int = 0

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor._DEFAULT_GRACE_RATIO = 0.1

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor._GRACE_LOGIT_BOOST = 5.0

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor.logger = logging.getLogger(__name__)