nemo_voice_agent.vllm.v1.sample.logits_processor

Submodules

nemo_voice_agent.vllm.v1.sample.logits_processor.reasoning_budget_logits_processor

Package Contents

Classes

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

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.