[Core][Engine] only materialize tokens when thinking budget is in req (#47053)
Signed-off-by: walterbm <walter.beller.morales@gmail.com>
Walter Beller-Morales committed
Jul 7, 2026 at 23:02 UTC
dd127d82ed29c40b7daf6e751add49ff371b1d9d
2 files changed
+29
-5
tests/v1/worker/test_gpu_model_runner.py
+26
@@ -18,6 +18,7 @@ from vllm.config import (
18
VllmConfig,
19
set_current_vllm_config,
20
)
21
+from vllm.config.reasoning import ReasoningConfig
22
from vllm.distributed.parallel_state import (
23
init_distributed_environment,
24
initialize_model_parallel,
@@ -255,6 +256,31 @@ def test_select_common_block_size_uses_largest_shared_int():
256
assert selected_size == 64
257
258
259
+def test_reasoning_config_without_custom_logitsprocs_does_not_need_output_token_ids(
260
+ dist_init,
261
+):
262
+ vllm_config = get_vllm_config()
263
+ assert vllm_config.model_config.logits_processors is None
264
+ reasoning_config = ReasoningConfig(
265
+ reasoning_start_str="<think>", reasoning_end_str="</think>"
266
+ )
267
+ reasoning_config._reasoning_start_token_ids = [1]
268
+ reasoning_config._reasoning_end_token_ids = [2]
269
+ vllm_config.reasoning_config = reasoning_config
270
+
271
+ with set_current_vllm_config(vllm_config):
272
+ model_config = vllm_config.model_config
273
+ num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)
274
+ head_size = model_config.get_head_size()
275
+ vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(
276
+ num_heads, head_size, 0.1
277
+ )
278
+ runner = GPUModelRunner(vllm_config, torch.device("cpu"))
279
+
280
+ assert runner.input_batch.thinking_budget_state_holder is not None
281
+ assert runner.input_batch.logitsprocs_need_output_token_ids is False
282
+
283
+
284
@pytest.mark.skip_global_cleanup
285
@pytest.mark.parametrize(
286
("world_size", "is_last_rank", "expected_calls"),
vllm/v1/worker/gpu_model_runner.py
+3
-5
@@ -702,11 +702,9 @@ class GPUModelRunner(
702
custom_logitsprocs,
703
),
704
# We currently don't know whether a particular custom logits processor
705
- # uses output token ids so we set this conservatively.
706
- # ThinkingTokenBudgetLogitsProcessor also needs output token ids to
707
- # correctly track think start/end token sequences in async scheduling.
708
- logitsprocs_need_output_token_ids=bool(custom_logitsprocs)
709
- or self.vllm_config.reasoning_config is not None,
705
+ # uses output token ids so we set this conservatively. Thinking-budget
706
+ # tracking is requested dynamically when a budgeted request is in the batch.
707
+ logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
708
is_pooling_model=self.is_pooling_model,
709
cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
710
reasoning_config=self.vllm_config.reasoning_config,