From 58e9d1dc411123a6966c9257f6cd159d77c34b49 Mon Sep 17 00:00:00 2001 From: Nakanokensetsu <281529662+Nakanokensetsu@users.noreply.github.com> Date: Thu, 13 Aug 2026 16:02:55 +0900 Subject: [PATCH 1/2] fix: thinking_token_budget has no effect when async scheduling is enabled ThinkingBudgetStateHolder tracks each request's / state and forces the reasoning end tokens onto the logits once thinking_token_budget is exceeded, via two methods: - update_state(): advances the per-request think/end state machine using the latest sampled output tokens. This was only ever called from sync_batch() (batch add/remove/move bookkeeping), never once per decode step, so a request's think state was never advanced past the moment it entered the batch. Budget overrun was therefore never detected once generation was under way. - apply_to_logits(): forces the reasoning end token(s) into the logits. This was only wired into rejection_sampler.py, the speculative-decoding sampler path. Sampler.forward(), used whenever speculative decoding is not active, never called it, so even a correctly-tracked budget overrun was never forced onto the logits. Together these mean thinking_token_budget silently had no effect for any request using the normal (non-speculative-decoding) sampling path once async scheduling was enabled -- which it is by default for compatible executors (AphroditeConfig.__post_init__, scheduler_config.async_scheduling is None -> True). Reasoning would run unbounded regardless of the budget value, up to max_tokens. The existing tests in test_thinking_token_budget.py did not catch this because their server fixtures explicitly pass --no-async-scheduling. Under that config, _make_sampling_metadata() happens to call update_state() from a different, sufficient path, masking the bug. Manual reproduction against Qwen/Qwen3-0.6B with async scheduling left at its default (enabled) confirmed the failure and the fix: thinking_token_budget=5, max_tokens=100 before: reasoning_token_count=None, total_decode_tokens=100 (budget ignored entirely, ran to max_tokens without closing ) after: reasoning_token_count=5, total_decode_tokens=17 (budget respected exactly, natural completion) Fix: call update_state() every decode step from GPUModelRunner._sample(), and call apply_to_logits() from Sampler.forward() mirroring the existing rejection_sampler.py call, so both paths are covered regardless of speculative decoding or async scheduling state. Also adds a new "async_scheduling" server fixture/param to test_thinking_token_budget_limits_reasoning that leaves async scheduling at its default instead of disabling it, so this regression is covered going forward. --- aphrodite/v1/sample/sampler.py | 18 +++++++++ aphrodite/v1/worker/gpu_model_runner.py | 16 ++++++++ .../test_thinking_token_budget.py | 39 ++++++++++++++++++- 3 files changed, 71 insertions(+), 2 deletions(-) diff --git a/aphrodite/v1/sample/sampler.py b/aphrodite/v1/sample/sampler.py index 4bbfb94e3d..358f2198a1 100644 --- a/aphrodite/v1/sample/sampler.py +++ b/aphrodite/v1/sample/sampler.py @@ -123,6 +123,24 @@ 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: diff --git a/aphrodite/v1/worker/gpu_model_runner.py b/aphrodite/v1/worker/gpu_model_runner.py index b8a76c7781..bd89b7c64a 100644 --- a/aphrodite/v1/worker/gpu_model_runner.py +++ b/aphrodite/v1/worker/gpu_model_runner.py @@ -3349,6 +3349,22 @@ 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, diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index 21fe0860a3..ae45b03ccf 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py +++ b/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": "", "reasoning_end_str": ""}', + "--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,20 @@ 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. From eb0a49f7bb597dc50256bf632a6e59c05463400c Mon Sep 17 00:00:00 2001 From: Nakanokensetsu <281529662+Nakanokensetsu@users.noreply.github.com> Date: Fri, 14 Aug 2026 07:22:09 +0900 Subject: [PATCH 2/2] style: apply ruff-format (single-line boolean conditions) CI's ruff-format hook wanted these three multi-line `if (...)` / `@pytest.mark.parametrize(...)` blocks collapsed to single lines. No behavior change. --- aphrodite/v1/sample/sampler.py | 5 +---- aphrodite/v1/worker/gpu_model_runner.py | 5 +---- .../openai/chat_completion/test_thinking_token_budget.py | 4 +--- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/aphrodite/v1/sample/sampler.py b/aphrodite/v1/sample/sampler.py index 358f2198a1..743018a623 100644 --- a/aphrodite/v1/sample/sampler.py +++ b/aphrodite/v1/sample/sampler.py @@ -131,10 +131,7 @@ def forward( # 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() - ): + 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, diff --git a/aphrodite/v1/worker/gpu_model_runner.py b/aphrodite/v1/worker/gpu_model_runner.py index bd89b7c64a..34c7d218fc 100644 --- a/aphrodite/v1/worker/gpu_model_runner.py +++ b/aphrodite/v1/worker/gpu_model_runner.py @@ -3357,10 +3357,7 @@ def _sample( # 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() - ): + 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, diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index ae45b03ccf..37d62c5269 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py @@ -204,9 +204,7 @@ async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI): @pytest.mark.asyncio -@pytest.mark.parametrize( - "client", ["default", "auto_config", "async_scheduling"], 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.