diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 00722ccc24..7490cbb27b 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -18,6 +18,7 @@ from vllm.config import ( VllmConfig, set_current_vllm_config, ) +from vllm.config.reasoning import ReasoningConfig from vllm.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, @@ -255,6 +256,31 @@ def test_select_common_block_size_uses_largest_shared_int(): assert selected_size == 64 +def test_reasoning_config_without_custom_logitsprocs_does_not_need_output_token_ids( + dist_init, +): + vllm_config = get_vllm_config() + assert vllm_config.model_config.logits_processors is None + reasoning_config = ReasoningConfig( + reasoning_start_str="", reasoning_end_str="" + ) + reasoning_config._reasoning_start_token_ids = [1] + reasoning_config._reasoning_end_token_ids = [2] + vllm_config.reasoning_config = reasoning_config + + with set_current_vllm_config(vllm_config): + model_config = vllm_config.model_config + num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config) + head_size = model_config.get_head_size() + vllm_config.compilation_config.static_forward_context["layer.0"] = Attention( + num_heads, head_size, 0.1 + ) + runner = GPUModelRunner(vllm_config, torch.device("cpu")) + + assert runner.input_batch.thinking_budget_state_holder is not None + assert runner.input_batch.logitsprocs_need_output_token_ids is False + + @pytest.mark.skip_global_cleanup @pytest.mark.parametrize( ("world_size", "is_last_rank", "expected_calls"), diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e9b23f1c61..199470aaf8 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -702,11 +702,9 @@ class GPUModelRunner( custom_logitsprocs, ), # We currently don't know whether a particular custom logits processor - # uses output token ids so we set this conservatively. - # ThinkingTokenBudgetLogitsProcessor also needs output token ids to - # correctly track think start/end token sequences in async scheduling. - logitsprocs_need_output_token_ids=bool(custom_logitsprocs) - or self.vllm_config.reasoning_config is not None, + # uses output token ids so we set this conservatively. Thinking-budget + # tracking is requested dynamically when a budgeted request is in the batch. + logitsprocs_need_output_token_ids=bool(custom_logitsprocs), is_pooling_model=self.is_pooling_model, cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, reasoning_config=self.vllm_config.reasoning_config,