fix: thinking_token_budget has no effect when async scheduling is enabled

aphrodite/v1/sample/sampler.py

@@ -123,6 +123,21 @@ def forward(
         for processor in sampling_metadata.logitsprocs.argmax_invariant:
             logits = processor.apply(logits)
 
+        # ThinkingBudgetStateHolder.apply_to_logits() forces the reasoning
+        # end tokens into the logits once thinking_token_budget is exceeded.
+        # It was previously wired only into rejection_sampler.py (the
+        # speculative-decoding sampler path), so requests without
+        # speculative decoding never had budget overrun forced onto the
+        # logits here and generation ran past the budget every time. Mirror
+        # the same call for the normal (non-spec-decode) path.
+        thinking_budget_state_holder = sampling_metadata.thinking_budget_state_holder
+        if thinking_budget_state_holder is not None and thinking_budget_state_holder.has_tracked_requests():
+            logits = thinking_budget_state_holder.apply_to_logits(
+                logits,
+                predict_bonus_token=predict_bonus_token,
+                spec_token_ids=sampling_metadata.spec_token_ids,
+            )
+
         # Sample the next token.
         sampled, processed_logprobs = self.sample(logits, sampling_metadata)
         if processed_logprobs is not None:

aphrodite/v1/worker/gpu_model_runner.py

@@ -3349,6 +3349,19 @@ 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()
+        # ThinkingBudgetStateHolder.update_state() drives the per-request
+        # think/end state machine forward; it was previously never called
+        # from the decode loop (only sync_batch() ran, on batch add/remove/
+        # move), so budget overrun was never detected past the first step
+        # and thinking_token_budget had no effect once generation was under
+        # way. Call it here every step, before sampling, using this step's
+        # freshly-updated token lists.
+        thinking_budget_state_holder = sampling_metadata.thinking_budget_state_holder
+        if thinking_budget_state_holder is not None and thinking_budget_state_holder.has_tracked_requests():
+            thinking_budget_state_holder.update_state(
+                sampling_metadata.output_token_ids,
+                sampling_metadata.spec_token_ids,
+            )
         if spec_decode_metadata is None:
             return self.sampler(
                 logits=logits,

tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py

@@ -142,11 +142,36 @@ def server_qwen35_fp8_mtp_tp2():
         yield remote_server
 
 
+@pytest.fixture(scope="module")
+def server_async_scheduling():
+    """Same as ``server``, but leaves async scheduling at its default
+    (enabled), instead of passing --no-async-scheduling. Async scheduling
+    is the default for compatible executors, so this covers the actual
+    default configuration rather than only the opted-out one.
+    """
+    args = [
+        "--reasoning-parser",
+        "qwen3",
+        "--reasoning-config",
+        '{"reasoning_start_str": "<think>", "reasoning_end_str": "</think>"}',
+        "--max-model-len",
+        "2048",
+        "--enforce-eager",
+        "--gpu-memory-utilization",
+        "0.4",
+    ]
+    # thinking_token_budget is not yet supported by the V2 model runner.
+    env_dict = {"APHRODITE_USE_V2_MODEL_RUNNER": "0"}
+    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
+        yield remote_server
+
+
 @pytest_asyncio.fixture
-async def client(request, server, server_with_auto_reasoning_config):
+async def client(request, server, server_with_auto_reasoning_config, server_async_scheduling):
     server_map = {
         "default": server,
         "auto_config": server_with_auto_reasoning_config,
+        "async_scheduling": server_async_scheduling,
     }
     target_server = server_map[request.param]
     async with target_server.get_async_client() as async_client:
@@ -179,10 +204,18 @@ async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI):
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True)
+@pytest.mark.parametrize("client", ["default", "auto_config", "async_scheduling"], indirect=True)
 async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI):
     """Test that thinking_token_budget limits the number of reasoning tokens.
 
+    The "async_scheduling" param covers the actual default configuration
+    (async scheduling enabled): ThinkingBudgetStateHolder.update_state()
+    previously ran only from sync_batch() (add/remove/move bookkeeping),
+    never once per decode step, so a request's think state was never
+    advanced past the first step and the budget was silently ignored for
+    the rest of generation whenever async scheduling was on. "default" and
+    "auto_config" pass --no-async-scheduling and would not have caught this.
+
     Counts reasoning decode tokens by id, which is robust to how tokens are
     grouped into streamed chunks (a single chunk can carry several tokens under
     async scheduling / stream_interval > 1). Counting chunks under-counts.