From 182a054dbbdad92c0cc2f75ab8c091dd3548c499 Mon Sep 17 00:00:00 2001 From: Alex Bilichenko Date: Sun, 31 May 2026 18:34:41 -0700 Subject: [PATCH 1/4] feat: split temperature for reasoning vs answer phase Allow a user to set one temperature for the reasoning (thinking) phase of a model's output and a separate temperature for the answer (content) phase. Key design decisions: - No sampler changes: blend temperatures in _make_sampling_metadata() so the sampler sees the correct per-request temperature for the current phase - No ThinkingBudgetStateHolder changes: think mask is computed inline by scanning prompt+output tokens for think-start/think-end markers - Zero overhead when split is inactive: has_reasoning_temp_split flag gates the entire path - Requests with temperature != reasoning_temperature are added to both greedy_reqs and random_reqs to force the mixed sampling path, which correctly handles per-request phase-dependent temperature via torch.where(temp < EPS, greedy_sampled, random_sampled) Usage (OpenAI-compatible API): vllm_xargs: { "reasoning_temperature": 0.7 } Co-authored-by: Claude --- .../openai/chat_completion/protocol.py | 8 + .../entrypoints/openai/completion/protocol.py | 6 + vllm/entrypoints/openai/responses/protocol.py | 5 + vllm/sampling_params.py | 36 +++- vllm/v1/sample/metadata.py | 4 + vllm/v1/worker/gpu_input_batch.py | 159 +++++++++++++++++- 6 files changed, 214 insertions(+), 4 deletions(-) diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 73ecb3f35a1c..55c1d30fac06 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -612,12 +612,20 @@ def to_sampling_params( if self.kv_transfer_params: # Pass in kv_transfer_params via extra_args extra_args["kv_transfer_params"] = self.kv_transfer_params + + # Extract known extension parameters from extra_args before passing + # rest downstream. Pop them so they don't end up in extra_args. + reasoning_temperature: float | None = extra_args.pop( + "reasoning_temperature", None + ) + return SamplingParams.from_optional( n=self.n, presence_penalty=self.presence_penalty, frequency_penalty=self.frequency_penalty, repetition_penalty=repetition_penalty, temperature=temperature, + reasoning_temperature=reasoning_temperature, top_p=top_p, top_k=top_k, min_p=min_p, diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index cb793a415633..75e191afa7c6 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -312,12 +312,18 @@ def to_sampling_params( if self.kv_transfer_params: # Pass in kv_transfer_params via extra_args extra_args["kv_transfer_params"] = self.kv_transfer_params + + reasoning_temperature: float | None = extra_args.pop( + "reasoning_temperature", None + ) + return SamplingParams.from_optional( n=self.n, presence_penalty=self.presence_penalty, frequency_penalty=self.frequency_penalty, repetition_penalty=repetition_penalty, temperature=temperature, + reasoning_temperature=reasoning_temperature, top_p=top_p, top_k=top_k, min_p=min_p, diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 30a920663651..cb75228e639c 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -401,8 +401,13 @@ def to_sampling_params( if self.kv_transfer_params: extra_args["kv_transfer_params"] = self.kv_transfer_params + reasoning_temperature: float | None = extra_args.pop( + "reasoning_temperature", None + ) + return SamplingParams.from_optional( temperature=temperature, + reasoning_temperature=reasoning_temperature, top_p=top_p, top_k=top_k, max_tokens=max_tokens, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 6e0be9dbff57..d3c981f60dab 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -236,6 +236,13 @@ class SamplingParams( """Controls the randomness of the sampling. Lower values make the model more deterministic, while higher values make the model more random. Zero means greedy sampling.""" + reasoning_temperature: float = 1.0 + """Temperature used during the reasoning/thinking phase. + When the model is generating reasoning tokens (before the answer begins), + this temperature is used instead of the regular ``temperature``. + Set to 1.0 (the default) to use the same temperature for both phases; + set to 0.0 for greedy reasoning while using a stochastic answer, or + higher for stochastic reasoning with a greedy answer.""" top_p: float = 1.0 """Controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to 1 to consider all tokens.""" @@ -358,6 +365,7 @@ def from_optional( frequency_penalty: float | None = 0.0, repetition_penalty: float | None = 1.0, temperature: float | None = 1.0, + reasoning_temperature: float | None = None, top_p: float | None = 1.0, top_k: int = 0, min_p: float = 0.0, @@ -399,6 +407,9 @@ def from_optional( if repetition_penalty is None else repetition_penalty, temperature=1.0 if temperature is None else temperature, + reasoning_temperature=1.0 + if reasoning_temperature is None + else reasoning_temperature, top_p=1.0 if top_p is None else top_p, top_k=top_k, min_p=min_p, @@ -436,6 +447,19 @@ def __post_init__(self) -> None: ) self.temperature = max(self.temperature, _MAX_TEMP) + if 0 < self.reasoning_temperature < _MAX_TEMP: + logger.warning( + "reasoning_temperature %s is less than %s, which may cause " + "numerical errors nan or inf in tensors. We have maxed it " + "out to %s.", + self.reasoning_temperature, + _MAX_TEMP, + _MAX_TEMP, + ) + self.reasoning_temperature = max( + self.reasoning_temperature, _MAX_TEMP + ) + if self.seed == -1: self.seed = None @@ -467,8 +491,8 @@ def __post_init__(self) -> None: self._verify_args() - if self.temperature < _SAMPLING_EPS: - # Zero temperature means greedy sampling. + if self.temperature < _SAMPLING_EPS and self.reasoning_temperature < _SAMPLING_EPS: + # Both temperatures are effectively zero -> fully greedy sampling. self.top_p = 1.0 self.top_k = 0 self.min_p = 0.0 @@ -514,6 +538,13 @@ def _verify_args(self) -> None: parameter="temperature", value=self.temperature, ) + if self.reasoning_temperature < 0.0: + raise VLLMValidationError( + f"reasoning_temperature must be non-negative, got " + f"{self.reasoning_temperature}.", + parameter="reasoning_temperature", + value=self.reasoning_temperature, + ) if not 0.0 < self.top_p <= 1.0: raise VLLMValidationError( f"top_p must be in (0, 1], got {self.top_p}.", @@ -968,6 +999,7 @@ def __repr__(self) -> str: f"frequency_penalty={self.frequency_penalty}, " f"repetition_penalty={self.repetition_penalty}, " f"temperature={self.temperature}, " + f"reasoning_temperature={self.reasoning_temperature}, " f"top_p={self.top_p}, " f"top_k={self.top_k}, " f"min_p={self.min_p}, " diff --git a/vllm/v1/sample/metadata.py b/vllm/v1/sample/metadata.py index fa4ceac8e71e..ccd051b65d18 100644 --- a/vllm/v1/sample/metadata.py +++ b/vllm/v1/sample/metadata.py @@ -53,3 +53,7 @@ class SamplingMetadata: # When non-None, use ``holder.has_tracked_requests()`` to see if this batch applies # thinking-token-budget logits (holder may exist with an empty tracking set). thinking_budget_state_holder: ThinkingBudgetStateHolder | None = None + + # Temperature split for reasoning vs answer phase. + reasoning_temperature: torch.Tensor | None = None + think_mask: torch.Tensor | None = None diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 89d69c0bde64..362814c30450 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -188,6 +188,33 @@ def __init__( (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory ) self.temperature_cpu = self.temperature_cpu_tensor.numpy() + + # Reasoning-phase temperature split. + self.reasoning_temperature = torch.empty( + (max_num_reqs,), dtype=torch.float32, device=device + ) + self.reasoning_temperature_cpu_tensor = torch.empty( + (max_num_reqs,), + dtype=torch.float32, + device="cpu", + pin_memory=pin_memory, + ) + self.reasoning_temperature_cpu = ( + self.reasoning_temperature_cpu_tensor.numpy() + ) + self.has_reasoning_temp_split = False + # Think-boundary token IDs, populated from reasoning_config. + self.think_start_token_ids: list[int] = [] + self.think_end_token_ids: list[int] = [] + if reasoning_config is not None: + if reasoning_config.reasoning_start_token_ids: + self.think_start_token_ids = ( + reasoning_config.reasoning_start_token_ids + ) + if reasoning_config.reasoning_end_token_ids: + self.think_end_token_ids = ( + reasoning_config.reasoning_end_token_ids + ) self.greedy_reqs: set[str] = set() self.random_reqs: set[str] = set() @@ -378,7 +405,27 @@ def add_request( self.block_table.add_row(request.block_ids, req_index) if sampling_params := request.sampling_params: - if sampling_params.sampling_type == SamplingType.GREEDY: + has_reasoning_split = ( + sampling_params.reasoning_temperature + != sampling_params.temperature + ) + if has_reasoning_split: + # When the temperature split is active, the request must go + # into both greedy_reqs AND random_reqs so that the sampler + # takes the mixed path. This allows phase-dependent sampling: + # - thinking phase -> reasoning_temperature (may be >0) + # - answer phase -> temperature (may be 0.0, greedy) + # The mixed path ensures 0.0 temps don't cause NaN (clamped + # to 1.0 in apply_temperature), and torch.where() selects + # greedy vs random per-request based on the blended temp. + # Store the ORIGINAL temperature in temperature_cpu and the + # thinking-phase temp in reasoning_temperature_cpu. The + # blend in _make_sampling_metadata() chooses between them. + self.temperature_cpu[req_index] = sampling_params.temperature + self.greedy_reqs.add(req_id) + self.random_reqs.add(req_id) + self.has_reasoning_temp_split = True + elif sampling_params.sampling_type == SamplingType.GREEDY: # Should avoid division by zero later when apply_temperature. self.temperature_cpu[req_index] = 0.0 self.greedy_reqs.add(req_id) @@ -386,6 +433,10 @@ def add_request( self.temperature_cpu[req_index] = sampling_params.temperature self.random_reqs.add(req_id) + self.reasoning_temperature_cpu[req_index] = ( + sampling_params.reasoning_temperature + ) + self.top_p_cpu[req_index] = sampling_params.top_p if sampling_params.top_p < 1: self.top_p_reqs.add(req_id) @@ -644,6 +695,10 @@ def swap_states(self, i1: int, i2: int) -> None: self.temperature_cpu[i2], self.temperature_cpu[i1], ) + self.reasoning_temperature_cpu[i1], self.reasoning_temperature_cpu[i2] = ( + self.reasoning_temperature_cpu[i2], + self.reasoning_temperature_cpu[i1], + ) self.top_p_cpu[i1], self.top_p_cpu[i2] = self.top_p_cpu[i2], self.top_p_cpu[i1] self.top_k_cpu[i1], self.top_k_cpu[i2] = self.top_k_cpu[i2], self.top_k_cpu[i1] self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = ( @@ -771,7 +826,12 @@ def condense(self) -> None: (last_req_index, empty_index, MoveDirectionality.UNIDIRECTIONAL) ) - self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index] + self.temperature_cpu[empty_index] = self.temperature_cpu[ + last_req_index + ] + self.reasoning_temperature_cpu[empty_index] = ( + self.reasoning_temperature_cpu[last_req_index] + ) self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index] self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index] self.frequency_penalties_cpu[empty_index] = self.frequency_penalties_cpu[ @@ -828,6 +888,73 @@ def refresh_metadata(self): if batch_update: self.sampling_metadata = self._make_sampling_metadata() + @staticmethod + def _last_token_sequence_index( + token_ids: list[int], seq: list[int] + ) -> int: + """Find the last occurrence of *seq* in *token_ids*, or -1.""" + if not seq or not token_ids: + return -1 + for i in range(len(token_ids) - len(seq), -1, -1): + if token_ids[i : i + len(seq)] == seq: + return i + return -1 + + def _compute_think_mask(self, num_reqs: int) -> torch.Tensor | None: + """Build a boolean mask: ``True`` = request is in thinking mode. + + Scans each request's combined token sequence (prompt + output) for + think-start and think-end markers. Returns ``None`` when the split + is not active or there are no think-boundary markers configured. + """ + if not self.has_reasoning_temp_split: + return None + if not self.think_start_token_ids or not self.think_end_token_ids: + return None + + mask = torch.zeros(num_reqs, dtype=torch.bool, device=self.device) + start_len = len(self.think_start_token_ids) + end_len = len(self.think_end_token_ids) + + for i in range(num_reqs): + # Combine prompt tokens (numpy slice) with output tokens. + prompt_len = int(self.num_prompt_tokens[i]) + combined: list[int] = [] + if prompt_len > 0: + combined.extend( + self.token_ids_cpu[i, :prompt_len].tolist() + ) + if i < len(self.req_output_token_ids): + out = self.req_output_token_ids[i] + if out is not None: + combined.extend(out) + + if len(combined) < max(start_len, end_len): + continue + + # Scan from the end for the latest occurrence of each marker. + last_start = -1 + for j in range(len(combined) - start_len, -1, -1): + if ( + combined[j : j + start_len] + == self.think_start_token_ids + ): + last_start = j + break + last_end = -1 + for j in range(len(combined) - end_len, -1, -1): + if ( + combined[j : j + end_len] + == self.think_end_token_ids + ): + last_end = j + break + + if last_start >= 0 or last_end >= 0: + mask[i] = last_start > last_end + + return mask + def _make_sampling_metadata(self) -> SamplingMetadata: num_reqs = self.num_reqs if not self.all_greedy: @@ -885,6 +1012,7 @@ def _make_sampling_metadata(self) -> SamplingMetadata: or bool(self.bad_words_token_ids) or self.logitsprocs_need_output_token_ids or thinking_budget_tracks_reqs + or self.has_reasoning_temp_split ) output_token_ids = ( cast(list[list[int]], self.req_output_token_ids) @@ -911,10 +1039,37 @@ def _make_sampling_metadata(self) -> SamplingMetadata: req_index = self.req_id_to_index[req_id] logprob_token_ids_by_index[req_index] = token_ids + # --- Temperature split for reasoning vs answer phase --- + think_mask: torch.Tensor | None = None + reasoning_temperature_tensor: torch.Tensor | None = None + + if self.has_reasoning_temp_split and temperature is not None: + # Copy the reasoning-phase temperature to the GPU. + reasoning_temperature_tensor = copy_slice( + self.reasoning_temperature_cpu_tensor, + self.reasoning_temperature, + num_reqs, + ) + # Determine which requests are currently in thinking mode. + think_mask = self._compute_think_mask(num_reqs) + + if think_mask is not None: + # Blend: for requests in thinking mode, use + # reasoning_temperature; for answer mode, use the original + # temperature (already in ``temperature``). + blended = torch.where( + think_mask, reasoning_temperature_tensor, temperature + ) + # Write the blended result back into the GPU tensor. + self.temperature[:num_reqs].copy_(blended) + temperature = self.temperature[:num_reqs] + return SamplingMetadata( temperature=temperature, all_greedy=self.all_greedy, all_random=self.all_random, + reasoning_temperature=reasoning_temperature_tensor, + think_mask=think_mask, top_p=None if self.no_top_p else self.top_p[:num_reqs], top_k=None if self.no_top_k else self.top_k[:num_reqs], generators=self.generators, From fcb82e1d7543f704cea00dfed0b8df79f16ed138 Mon Sep 17 00:00:00 2001 From: Alex Bilichenko Date: Mon, 1 Jun 2026 00:12:03 -0700 Subject: [PATCH 2/4] Apply reasoning_temperature per decode step; fix unset-default gate The original implementation never applied reasoning_temperature at runtime: the answer/reasoning blend lived in InputBatch._make_sampling_metadata(), which only runs on batch composition changes, so the think-mask was computed once at request admission (output empty -> all-False) and frozen at the answer temperature for the whole generation. Verified on gemma-4-26B via per-token logprob analysis: reasoning tokens tracked the answer temperature regardless of reasoning_temperature. Changes: - Move the blend to a per-step InputBatch.update_reasoning_temperature(), called every decode step from _update_states after refresh_metadata(). It resets temperature from the answer-phase CPU base and applies where(in_think, reasoning_temperature, temperature) in place, so the effective temperature follows the model across the think/answer boundary. Gated on has_reasoning_temp_split (zero overhead otherwise). - Track the per-row think mask in update reasoning state, seeded from the prompt and refined from generated output each step. - Fix the activation gate: reasoning_temperature now defaults to None ("mirror temperature", no split) instead of 1.0, which previously turned the split on for almost every request (temperature != 1.0) and forced the whole batch onto the mixed greedy+random sampling path. - has_reasoning_temp_split is now derived from a per-request set, so it clears when split requests leave the batch (was a sticky bool). - Plumb reasoning_temperature through chat/completion/responses to_sampling_params() via vllm_xargs without mutating the request in place. - Remove dead SamplingMetadata.reasoning_temperature / think_mask fields. - Add tests/v1/sample/test_reasoning_temperature.py. Verified after the change: reasoning off-argmax rate moves 2.6% (rt=0) -> 8.4% (rt=1) while the answer phase stays near the greedy floor. AI assistance (Claude) was used for this change. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Alex Bilichenko --- tests/v1/sample/test_reasoning_temperature.py | 87 +++++++ .../openai/chat_completion/protocol.py | 4 +- .../entrypoints/openai/completion/protocol.py | 4 +- vllm/entrypoints/openai/responses/protocol.py | 4 +- vllm/sampling_params.py | 39 +-- vllm/v1/sample/metadata.py | 4 - vllm/v1/worker/gpu_input_batch.py | 242 +++++++++++------- vllm/v1/worker/gpu_model_runner.py | 4 + 8 files changed, 270 insertions(+), 118 deletions(-) create mode 100644 tests/v1/sample/test_reasoning_temperature.py diff --git a/tests/v1/sample/test_reasoning_temperature.py b/tests/v1/sample/test_reasoning_temperature.py new file mode 100644 index 000000000000..493bde9ea7bd --- /dev/null +++ b/tests/v1/sample/test_reasoning_temperature.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the reasoning/answer temperature split on SamplingParams. + +These lock in the contract the InputBatch split gate relies on: an unset +``reasoning_temperature`` must mean "mirror temperature" (no split, no +overhead), while a set, differing value opts the request into phase-dependent +sampling. +""" + +import pytest + +from vllm.sampling_params import _MAX_TEMP, SamplingParams, SamplingType + + +def test_default_is_unset(): + """An unspecified reasoning_temperature stays None (the no-split sentinel), + so an ordinary greedy request is not pushed onto the split path.""" + params = SamplingParams(temperature=0.0) + assert params.reasoning_temperature is None + # Plain greedy request remains greedy and gets the greedy top_p/k reset. + assert params.sampling_type == SamplingType.GREEDY + assert params.top_p == 1.0 + assert params.top_k == 0 + assert params.min_p == 0.0 + + +def test_from_optional_passes_none_through(): + """from_optional must not coerce an unset value to a concrete float; + coercing to 1.0 was what made the split fire for nearly every request.""" + params = SamplingParams.from_optional(temperature=0.7) + assert params.reasoning_temperature is None + + +@pytest.mark.parametrize("reasoning_temperature", [0.0, 0.7, 1.5]) +def test_split_is_set_when_specified(reasoning_temperature): + params = SamplingParams.from_optional( + temperature=0.0, reasoning_temperature=reasoning_temperature + ) + assert params.reasoning_temperature == reasoning_temperature + + +def test_greedy_reset_skipped_when_reasoning_is_random(): + """temperature=0 with a stochastic reasoning_temperature must NOT collapse + top_p/top_k/min_p, since the reasoning phase still samples randomly.""" + params = SamplingParams( + temperature=0.0, + reasoning_temperature=0.7, + top_p=0.5, + top_k=20, + min_p=0.1, + ) + assert params.top_p == 0.5 + assert params.top_k == 20 + assert params.min_p == 0.1 + + +def test_greedy_reset_applies_when_both_greedy(): + """temperature=0 and reasoning_temperature=0 is fully greedy and should + reset the nucleus parameters just like a plain greedy request.""" + params = SamplingParams( + temperature=0.0, + reasoning_temperature=0.0, + top_p=0.5, + top_k=20, + min_p=0.1, + ) + assert params.top_p == 1.0 + assert params.top_k == 0 + assert params.min_p == 0.0 + + +def test_negative_reasoning_temperature_rejected(): + with pytest.raises(ValueError): + SamplingParams(reasoning_temperature=-0.1) + + +def test_tiny_positive_reasoning_temperature_clamped(): + params = SamplingParams(reasoning_temperature=_MAX_TEMP / 2) + assert params.reasoning_temperature == _MAX_TEMP + + +def test_unset_reasoning_temperature_not_clamped(): + """None must survive __post_init__ untouched (no clamp, no validation + error).""" + params = SamplingParams(temperature=0.0) + assert params.reasoning_temperature is None diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 55c1d30fac06..6d6b7b72ff46 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -608,7 +608,9 @@ def to_sampling_params( else replace(self.structured_outputs, **structured_outputs_kwargs) ) - extra_args: dict[str, Any] = self.vllm_xargs if self.vllm_xargs else {} + # Copy so popping known extension params below does not mutate the + # request's ``vllm_xargs`` in place. + extra_args: dict[str, Any] = dict(self.vllm_xargs) if self.vllm_xargs else {} if self.kv_transfer_params: # Pass in kv_transfer_params via extra_args extra_args["kv_transfer_params"] = self.kv_transfer_params diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 75e191afa7c6..ce3505efa682 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -308,7 +308,9 @@ def to_sampling_params( else replace(self.structured_outputs, **structured_outputs_kwargs) ) - extra_args: dict[str, Any] = self.vllm_xargs if self.vllm_xargs else {} + # Copy so popping known extension params below does not mutate the + # request's ``vllm_xargs`` in place. + extra_args: dict[str, Any] = dict(self.vllm_xargs) if self.vllm_xargs else {} if self.kv_transfer_params: # Pass in kv_transfer_params via extra_args extra_args["kv_transfer_params"] = self.kv_transfer_params diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index cb75228e639c..791a4fb68f40 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -397,7 +397,9 @@ def to_sampling_params( if isinstance(stop, str): stop = [stop] - extra_args: dict[str, Any] = self.vllm_xargs if self.vllm_xargs else {} + # Copy so popping known extension params below does not mutate the + # request's ``vllm_xargs`` in place. + extra_args: dict[str, Any] = dict(self.vllm_xargs) if self.vllm_xargs else {} if self.kv_transfer_params: extra_args["kv_transfer_params"] = self.kv_transfer_params diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index d3c981f60dab..3158a0aa499e 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -236,13 +236,14 @@ class SamplingParams( """Controls the randomness of the sampling. Lower values make the model more deterministic, while higher values make the model more random. Zero means greedy sampling.""" - reasoning_temperature: float = 1.0 + reasoning_temperature: float | None = None """Temperature used during the reasoning/thinking phase. - When the model is generating reasoning tokens (before the answer begins), - this temperature is used instead of the regular ``temperature``. - Set to 1.0 (the default) to use the same temperature for both phases; - set to 0.0 for greedy reasoning while using a stochastic answer, or - higher for stochastic reasoning with a greedy answer.""" + When set and different from ``temperature``, reasoning/thinking tokens are + sampled at ``reasoning_temperature`` while the answer is sampled at + ``temperature``. For example ``temperature=0.0, reasoning_temperature=0.7`` + gives a greedy answer with stochastic reasoning. Leave unset (``None``, + the default) to use ``temperature`` for both phases — no split, no + overhead.""" top_p: float = 1.0 """Controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to 1 to consider all tokens.""" @@ -407,9 +408,7 @@ def from_optional( if repetition_penalty is None else repetition_penalty, temperature=1.0 if temperature is None else temperature, - reasoning_temperature=1.0 - if reasoning_temperature is None - else reasoning_temperature, + reasoning_temperature=reasoning_temperature, top_p=1.0 if top_p is None else top_p, top_k=top_k, min_p=min_p, @@ -447,7 +446,9 @@ def __post_init__(self) -> None: ) self.temperature = max(self.temperature, _MAX_TEMP) - if 0 < self.reasoning_temperature < _MAX_TEMP: + if self.reasoning_temperature is not None and ( + 0 < self.reasoning_temperature < _MAX_TEMP + ): logger.warning( "reasoning_temperature %s is less than %s, which may cause " "numerical errors nan or inf in tensors. We have maxed it " @@ -456,9 +457,7 @@ def __post_init__(self) -> None: _MAX_TEMP, _MAX_TEMP, ) - self.reasoning_temperature = max( - self.reasoning_temperature, _MAX_TEMP - ) + self.reasoning_temperature = max(self.reasoning_temperature, _MAX_TEMP) if self.seed == -1: self.seed = None @@ -491,8 +490,16 @@ def __post_init__(self) -> None: self._verify_args() - if self.temperature < _SAMPLING_EPS and self.reasoning_temperature < _SAMPLING_EPS: - # Both temperatures are effectively zero -> fully greedy sampling. + # ``reasoning_temperature is None`` means the reasoning phase mirrors + # ``temperature``, so the request is fully greedy iff ``temperature`` + # is. A set, non-greedy reasoning_temperature keeps top_p/top_k/min_p + # alive for the (random) reasoning phase even when the answer is greedy. + reasoning_temp_greedy = ( + self.reasoning_temperature is None + or self.reasoning_temperature < _SAMPLING_EPS + ) + if self.temperature < _SAMPLING_EPS and reasoning_temp_greedy: + # Both phases are effectively zero -> fully greedy sampling. self.top_p = 1.0 self.top_k = 0 self.min_p = 0.0 @@ -538,7 +545,7 @@ def _verify_args(self) -> None: parameter="temperature", value=self.temperature, ) - if self.reasoning_temperature < 0.0: + if self.reasoning_temperature is not None and self.reasoning_temperature < 0.0: raise VLLMValidationError( f"reasoning_temperature must be non-negative, got " f"{self.reasoning_temperature}.", diff --git a/vllm/v1/sample/metadata.py b/vllm/v1/sample/metadata.py index ccd051b65d18..fa4ceac8e71e 100644 --- a/vllm/v1/sample/metadata.py +++ b/vllm/v1/sample/metadata.py @@ -53,7 +53,3 @@ class SamplingMetadata: # When non-None, use ``holder.has_tracked_requests()`` to see if this batch applies # thinking-token-budget logits (holder may exist with an empty tracking set). thinking_budget_state_holder: ThinkingBudgetStateHolder | None = None - - # Temperature split for reasoning vs answer phase. - reasoning_temperature: torch.Tensor | None = None - think_mask: torch.Tensor | None = None diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 362814c30450..6b581f243d04 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -199,10 +199,24 @@ def __init__( device="cpu", pin_memory=pin_memory, ) - self.reasoning_temperature_cpu = ( - self.reasoning_temperature_cpu_tensor.numpy() + self.reasoning_temperature_cpu = self.reasoning_temperature_cpu_tensor.numpy() + # Requests whose reasoning_temperature differs from temperature. The + # batch-level split path is active iff this set is non-empty (see the + # ``has_reasoning_temp_split`` property), so it clears automatically + # once all such requests leave the batch. + self.reasoning_temp_split_reqs: set[str] = set() + # Think state seeded from each row's prompt (fallback when the output + # contains no marker yet). + self.prompt_in_think_cpu = np.zeros(max_num_reqs, dtype=bool) + # Live per-step think mask, refreshed every decode step in + # ``update_reasoning_temperature`` and copied to the GPU for the blend. + self.think_in_think_cpu_tensor = torch.zeros( + (max_num_reqs,), dtype=torch.bool, device="cpu", pin_memory=pin_memory + ) + self.think_in_think_cpu = self.think_in_think_cpu_tensor.numpy() + self.think_in_think_gpu = torch.zeros( + (max_num_reqs,), dtype=torch.bool, device=device ) - self.has_reasoning_temp_split = False # Think-boundary token IDs, populated from reasoning_config. self.think_start_token_ids: list[int] = [] self.think_end_token_ids: list[int] = [] @@ -405,9 +419,14 @@ def add_request( self.block_table.add_row(request.block_ids, req_index) if sampling_params := request.sampling_params: + # ``reasoning_temperature is None`` means "mirror temperature", so + # the split is only active when it is set AND differs. Comparing + # the raw default (None) here is what keeps ordinary requests off + # the mixed sampling path. + reasoning_temperature = sampling_params.reasoning_temperature has_reasoning_split = ( - sampling_params.reasoning_temperature - != sampling_params.temperature + reasoning_temperature is not None + and reasoning_temperature != sampling_params.temperature ) if has_reasoning_split: # When the temperature split is active, the request must go @@ -419,12 +438,12 @@ def add_request( # to 1.0 in apply_temperature), and torch.where() selects # greedy vs random per-request based on the blended temp. # Store the ORIGINAL temperature in temperature_cpu and the - # thinking-phase temp in reasoning_temperature_cpu. The - # blend in _make_sampling_metadata() chooses between them. + # thinking-phase temp in reasoning_temperature_cpu; + # update_reasoning_temperature() blends them each step. self.temperature_cpu[req_index] = sampling_params.temperature self.greedy_reqs.add(req_id) self.random_reqs.add(req_id) - self.has_reasoning_temp_split = True + self.reasoning_temp_split_reqs.add(req_id) elif sampling_params.sampling_type == SamplingType.GREEDY: # Should avoid division by zero later when apply_temperature. self.temperature_cpu[req_index] = 0.0 @@ -433,9 +452,21 @@ def add_request( self.temperature_cpu[req_index] = sampling_params.temperature self.random_reqs.add(req_id) - self.reasoning_temperature_cpu[req_index] = ( - sampling_params.reasoning_temperature - ) + if has_reasoning_split: + self.reasoning_temperature_cpu[req_index] = reasoning_temperature + # Seed the live think state from the prompt; the per-step update + # refines it from generated output. + in_think = self._prompt_in_think(request.prompt_token_ids) + self.prompt_in_think_cpu[req_index] = in_think + self.think_in_think_cpu[req_index] = in_think + else: + # Mirror temperature so the per-step blend is a no-op for rows + # that did not request a split. + self.reasoning_temperature_cpu[req_index] = self.temperature_cpu[ + req_index + ] + self.prompt_in_think_cpu[req_index] = False + self.think_in_think_cpu[req_index] = False self.top_p_cpu[req_index] = sampling_params.top_p if sampling_params.top_p < 1: @@ -595,6 +626,7 @@ def remove_request(self, req_id: str) -> int | None: self.greedy_reqs.discard(req_id) self.random_reqs.discard(req_id) + self.reasoning_temp_split_reqs.discard(req_id) self.top_p_reqs.discard(req_id) self.top_k_reqs.discard(req_id) self.frequency_penalties_reqs.discard(req_id) @@ -699,6 +731,14 @@ def swap_states(self, i1: int, i2: int) -> None: self.reasoning_temperature_cpu[i2], self.reasoning_temperature_cpu[i1], ) + self.prompt_in_think_cpu[i1], self.prompt_in_think_cpu[i2] = ( + self.prompt_in_think_cpu[i2], + self.prompt_in_think_cpu[i1], + ) + self.think_in_think_cpu[i1], self.think_in_think_cpu[i2] = ( + self.think_in_think_cpu[i2], + self.think_in_think_cpu[i1], + ) self.top_p_cpu[i1], self.top_p_cpu[i2] = self.top_p_cpu[i2], self.top_p_cpu[i1] self.top_k_cpu[i1], self.top_k_cpu[i2] = self.top_k_cpu[i2], self.top_k_cpu[i1] self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = ( @@ -826,12 +866,16 @@ def condense(self) -> None: (last_req_index, empty_index, MoveDirectionality.UNIDIRECTIONAL) ) - self.temperature_cpu[empty_index] = self.temperature_cpu[ - last_req_index - ] + self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index] self.reasoning_temperature_cpu[empty_index] = ( self.reasoning_temperature_cpu[last_req_index] ) + self.prompt_in_think_cpu[empty_index] = self.prompt_in_think_cpu[ + last_req_index + ] + self.think_in_think_cpu[empty_index] = self.think_in_think_cpu[ + last_req_index + ] self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index] self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index] self.frequency_penalties_cpu[empty_index] = self.frequency_penalties_cpu[ @@ -889,71 +933,95 @@ def refresh_metadata(self): self.sampling_metadata = self._make_sampling_metadata() @staticmethod - def _last_token_sequence_index( - token_ids: list[int], seq: list[int] - ) -> int: - """Find the last occurrence of *seq* in *token_ids*, or -1.""" - if not seq or not token_ids: + def _last_token_sequence_index(token_ids: list[int], seq: list[int]) -> int: + """Find the start index of the last occurrence of *seq* in + *token_ids*, or -1 if absent.""" + if not seq or len(token_ids) < len(seq): return -1 for i in range(len(token_ids) - len(seq), -1, -1): if token_ids[i : i + len(seq)] == seq: return i return -1 - def _compute_think_mask(self, num_reqs: int) -> torch.Tensor | None: - """Build a boolean mask: ``True`` = request is in thinking mode. - - Scans each request's combined token sequence (prompt + output) for - think-start and think-end markers. Returns ``None`` when the split - is not active or there are no think-boundary markers configured. - """ - if not self.has_reasoning_temp_split: - return None + def _prompt_in_think(self, prompt_token_ids: list[int] | None) -> bool: + """Whether *prompt_token_ids* ends inside a thinking section.""" + if not prompt_token_ids: + return False if not self.think_start_token_ids or not self.think_end_token_ids: - return None - - mask = torch.zeros(num_reqs, dtype=torch.bool, device=self.device) - start_len = len(self.think_start_token_ids) - end_len = len(self.think_end_token_ids) - - for i in range(num_reqs): - # Combine prompt tokens (numpy slice) with output tokens. - prompt_len = int(self.num_prompt_tokens[i]) - combined: list[int] = [] - if prompt_len > 0: - combined.extend( - self.token_ids_cpu[i, :prompt_len].tolist() - ) - if i < len(self.req_output_token_ids): - out = self.req_output_token_ids[i] - if out is not None: - combined.extend(out) - - if len(combined) < max(start_len, end_len): + return False + last_start = self._last_token_sequence_index( + prompt_token_ids, self.think_start_token_ids + ) + last_end = self._last_token_sequence_index( + prompt_token_ids, self.think_end_token_ids + ) + return last_start > last_end + + def _update_think_state(self, num_reqs: int) -> None: + """Refresh the per-row think mask in ``think_in_think_cpu``. + + For each split request, find the last think-start and think-end markers + in the generated output: the row is "in think" when a start marker is + the most recent of the two. If neither marker is present yet, fall back + to the prompt's think state. + + NOTE: this rescans the full output each step (O(output_len) per split + row). A scanned-prefix incremental version is unsafe under async + scheduling, where the in-flight token slot holds a placeholder (-1) + that is overwritten with the real token id a step later — advancing a + cursor over that slot skips the real marker. A correct incremental + version would need to track the last *committed* output length; left as + a follow-up optimization. + """ + start = self.think_start_token_ids + end = self.think_end_token_ids + for req_id in self.reasoning_temp_split_reqs: + i = self.req_id_to_index.get(req_id) + if i is None or i >= num_reqs: continue + out = self.req_output_token_ids[i] + if not out: + self.think_in_think_cpu[i] = bool(self.prompt_in_think_cpu[i]) + continue + last_start = self._last_token_sequence_index(out, start) + last_end = self._last_token_sequence_index(out, end) + if last_start >= 0 or last_end >= 0: + self.think_in_think_cpu[i] = last_start > last_end + else: + self.think_in_think_cpu[i] = bool(self.prompt_in_think_cpu[i]) - # Scan from the end for the latest occurrence of each marker. - last_start = -1 - for j in range(len(combined) - start_len, -1, -1): - if ( - combined[j : j + start_len] - == self.think_start_token_ids - ): - last_start = j - break - last_end = -1 - for j in range(len(combined) - end_len, -1, -1): - if ( - combined[j : j + end_len] - == self.think_end_token_ids - ): - last_end = j - break + def update_reasoning_temperature(self) -> None: + """Apply the reasoning/answer temperature split for the current step. - if last_start >= 0 or last_end >= 0: - mask[i] = last_start > last_end + Called every decode step (unlike ``_make_sampling_metadata``, which + only runs on batch changes). Recomputes the live ``temperature`` tensor + the sampler reads as ``where(in_think, reasoning_temperature, + temperature)``, so the effective temperature follows the model across + the think/answer boundary mid-generation. No-op (zero overhead) when no + request in the batch requested a split. + """ + if not self.has_reasoning_temp_split: + return + if not self.think_start_token_ids or not self.think_end_token_ids: + return + num_reqs = self.num_reqs + if num_reqs == 0 or self.all_greedy: + return - return mask + self._update_think_state(num_reqs) + # Reset to the answer-phase base, then blend in reasoning temps where + # the row is currently thinking. Starting from the CPU base each step + # keeps the blend idempotent (no drift from repeated in-place writes). + answer = copy_slice(self.temperature_cpu_tensor, self.temperature, num_reqs) + reasoning = copy_slice( + self.reasoning_temperature_cpu_tensor, + self.reasoning_temperature, + num_reqs, + ) + mask = copy_slice( + self.think_in_think_cpu_tensor, self.think_in_think_gpu, num_reqs + ) + torch.where(mask, reasoning, answer, out=self.temperature[:num_reqs]) def _make_sampling_metadata(self) -> SamplingMetadata: num_reqs = self.num_reqs @@ -1039,37 +1107,17 @@ def _make_sampling_metadata(self) -> SamplingMetadata: req_index = self.req_id_to_index[req_id] logprob_token_ids_by_index[req_index] = token_ids - # --- Temperature split for reasoning vs answer phase --- - think_mask: torch.Tensor | None = None - reasoning_temperature_tensor: torch.Tensor | None = None - - if self.has_reasoning_temp_split and temperature is not None: - # Copy the reasoning-phase temperature to the GPU. - reasoning_temperature_tensor = copy_slice( - self.reasoning_temperature_cpu_tensor, - self.reasoning_temperature, - num_reqs, - ) - # Determine which requests are currently in thinking mode. - think_mask = self._compute_think_mask(num_reqs) - - if think_mask is not None: - # Blend: for requests in thinking mode, use - # reasoning_temperature; for answer mode, use the original - # temperature (already in ``temperature``). - blended = torch.where( - think_mask, reasoning_temperature_tensor, temperature - ) - # Write the blended result back into the GPU tensor. - self.temperature[:num_reqs].copy_(blended) - temperature = self.temperature[:num_reqs] - + # The reasoning/answer temperature split is applied per decode step in + # ``update_reasoning_temperature`` (called after ``refresh_metadata``), + # not here: this metadata is only rebuilt on batch changes, so blending + # here would freeze the phase at admission time. ``temperature`` is the + # answer-phase base; the per-step update overwrites the same tensor + # (``self.temperature[:num_reqs]``, which this view aliases) in place + # before the sampler reads it. return SamplingMetadata( temperature=temperature, all_greedy=self.all_greedy, all_random=self.all_random, - reasoning_temperature=reasoning_temperature_tensor, - think_mask=think_mask, top_p=None if self.no_top_p else self.top_p[:num_reqs], top_k=None if self.no_top_k else self.top_k[:num_reqs], generators=self.generators, @@ -1241,6 +1289,10 @@ def update_async_spec_token_ids(self, draft_token_ids: list[list[int]]) -> None: def num_reqs(self) -> int: return len(self.req_id_to_index) + @property + def has_reasoning_temp_split(self) -> bool: + return len(self.reasoning_temp_split_reqs) > 0 + @property def all_greedy(self) -> bool: return len(self.random_reqs) == 0 diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index f82d2224a413..e59d94094cc5 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1440,6 +1440,10 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None self._may_reorder_batch(scheduler_output) # Refresh batch metadata with any pending updates. self.input_batch.refresh_metadata() + # Apply the reasoning/answer temperature split for this step (no-op + # unless a request set reasoning_temperature). Must run every step, + # after refresh_metadata, since the phase changes mid-generation. + self.input_batch.update_reasoning_temperature() # Incrementally update ngram_gpu tensors after batch is stable if is_ngram_gpu: From fe4189de1754b8260add05b5715d443e0c41a317 Mon Sep 17 00:00:00 2001 From: Alex Bilichenko Date: Mon, 1 Jun 2026 01:47:28 -0700 Subject: [PATCH 3/4] Fix reasoning temperature edge cases Co-authored-by: OpenAI Codex --- tests/v1/sample/test_reasoning_temperature.py | 30 +++++++++++++++++++ vllm/sampling_params.py | 27 ++++++++++++++++- vllm/v1/worker/gpu_model_runner.py | 5 +--- 3 files changed, 57 insertions(+), 5 deletions(-) diff --git a/tests/v1/sample/test_reasoning_temperature.py b/tests/v1/sample/test_reasoning_temperature.py index 493bde9ea7bd..405b86ddd34e 100644 --- a/tests/v1/sample/test_reasoning_temperature.py +++ b/tests/v1/sample/test_reasoning_temperature.py @@ -8,6 +8,8 @@ sampling. """ +from unittest.mock import Mock + import pytest from vllm.sampling_params import _MAX_TEMP, SamplingParams, SamplingType @@ -55,6 +57,34 @@ def test_greedy_reset_skipped_when_reasoning_is_random(): assert params.min_p == 0.1 +def test_sampling_type_uses_reasoning_temperature(): + params = SamplingParams( + temperature=0.0, + reasoning_temperature=0.7, + seed=123, + ) + assert params.sampling_type == SamplingType.RANDOM_SEED + + +def test_spec_decode_disables_reasoning_temperature_split(): + params = SamplingParams( + temperature=0.0, + reasoning_temperature=0.7, + top_p=0.5, + top_k=20, + min_p=0.1, + ) + assert params.sampling_type == SamplingType.RANDOM + + params._validate_spec_decode(Mock()) + + assert params.reasoning_temperature is None + assert params.sampling_type == SamplingType.GREEDY + assert params.top_p == 1.0 + assert params.top_k == 0 + assert params.min_p == 0.0 + + def test_greedy_reset_applies_when_both_greedy(): """temperature=0 and reasoning_temperature=0 is fully greedy and should reset the nucleus parameters just like a plain greedy request.""" diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 3158a0aa499e..5e89cabe5c67 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -698,7 +698,15 @@ def update_from_tokenizer(self, tokenizer: TokenizerLike) -> None: @cached_property def sampling_type(self) -> SamplingType: - if self.temperature < _SAMPLING_EPS: + reasoning_temperature = ( + self.temperature + if self.reasoning_temperature is None + else self.reasoning_temperature + ) + if ( + self.temperature < _SAMPLING_EPS + and reasoning_temperature < _SAMPLING_EPS + ): return SamplingType.GREEDY if self.seed is not None: return SamplingType.RANDOM_SEED @@ -856,6 +864,23 @@ def _validate_spec_decode( if speculative_config is None: return + if ( + self.reasoning_temperature is not None + and self.reasoning_temperature != self.temperature + ): + logger.warning_once( + "reasoning_temperature is not supported with speculative " + "decoding yet. Ignoring reasoning_temperature and using " + "temperature for all generated tokens." + ) + self.reasoning_temperature = None + self.__dict__.pop("sampling_type", None) + if self.temperature < _SAMPLING_EPS: + self.top_p = 1.0 + self.top_k = 0 + self.min_p = 0.0 + self._verify_greedy_sampling() + # Some sampling parameters are not yet compatible with spec decoding. if self.min_p > _SAMPLING_EPS or self.logit_bias: raise ValueError( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e59d94094cc5..d84428c3ae85 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1440,10 +1440,6 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None self._may_reorder_batch(scheduler_output) # Refresh batch metadata with any pending updates. self.input_batch.refresh_metadata() - # Apply the reasoning/answer temperature split for this step (no-op - # unless a request set reasoning_temperature). Must run every step, - # after refresh_metadata, since the phase changes mid-generation. - self.input_batch.update_reasoning_temperature() # Incrementally update ngram_gpu tensors after batch is stable if is_ngram_gpu: @@ -3500,6 +3496,7 @@ def _sample( # Update output token ids with tokens sampled in last step # if async scheduling and required by current sampling params. self.input_batch.update_async_output_token_ids() + self.input_batch.update_reasoning_temperature() if spec_decode_metadata is None: return self.sampler( logits=logits, From bdf361013d72f2dca5a3e1a2776f64c749619d82 Mon Sep 17 00:00:00 2001 From: Alex Bilichenko Date: Sun, 26 Jul 2026 20:06:53 +0000 Subject: [PATCH 4/4] fix: use PIN_MEMORY for reasoning temperature buffers Upstream replaced the local `pin_memory` argument in `InputBatch.__init__` with the module-level `PIN_MEMORY` constant from `vllm.utils.torch_utils`. When this branch is replayed onto current main, the two reasoning-temperature buffers still referenced the removed local and the engine died at startup: NameError: name 'pin_memory' is not defined vllm/v1/worker/gpu_input_batch.py:215 (InputBatch.__init__) Only these two buffers are ours; every other use in the file is upstream's and was already converted. Point them at PIN_MEMORY so the branch survives replay. This restores a fixup that previously lived only as a local commit on pook's integration branch and was never pushed here, so rebuilding the branch from scratch silently lost it. --- vllm/v1/worker/gpu_input_batch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 6b581f243d04..f1b414b1ecbf 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -197,7 +197,7 @@ def __init__( (max_num_reqs,), dtype=torch.float32, device="cpu", - pin_memory=pin_memory, + pin_memory=PIN_MEMORY, ) self.reasoning_temperature_cpu = self.reasoning_temperature_cpu_tensor.numpy() # Requests whose reasoning_temperature differs from temperature. The @@ -211,7 +211,7 @@ def __init__( # Live per-step think mask, refreshed every decode step in # ``update_reasoning_temperature`` and copied to the GPU for the blend. self.think_in_think_cpu_tensor = torch.zeros( - (max_num_reqs,), dtype=torch.bool, device="cpu", pin_memory=pin_memory + (max_num_reqs,), dtype=torch.bool, device="cpu", pin_memory=PIN_MEMORY ) self.think_in_think_cpu = self.think_in_think_cpu_tensor.numpy() self.think_in_think_gpu = torch.zeros(