fix: thinking_token_budget has no effect when async scheduling is enabled

[Nakanokensetsu]

Summary

ThinkingBudgetStateHolder tracks each request's <think>/</think> state and forces the reasoning end tokens onto the logits once thinking_token_budget is exceeded, via two methods. Both had a wiring gap:

• update_state() (advances the per-request think/end state machine from the latest sampled tokens) was only ever called from sync_batch() (batch add/remove/move bookkeeping), never once per decode step. A request's think state was never advanced past the moment it entered the batch, so budget overrun was never detected once generation was under way.
• apply_to_logits() (forces the reasoning end token(s) into the logits) 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, thinking_token_budget silently had no effect for any request on the normal (non-speculative-decoding) sampling path once async scheduling is enabled — which is the default for compatible executors (AphroditeConfig.__post_init__, scheduler_config.async_scheduling is None -> True). Reasoning ran unbounded regardless of the configured budget, up to max_tokens.

The existing tests in test_thinking_token_budget.py didn't 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.

Reproduction

Manual repro against Qwen/Qwen3-0.6B with async scheduling left at its default (enabled), thinking_token_budget=5, max_tokens=100:

before: reasoning_token_count=None, total_decode_tokens=100
        (budget ignored entirely, ran to max_tokens without ever closing </think>)
after:  reasoning_token_count=5, total_decode_tokens=17
        (budget respected exactly, natural completion)

Fix

• Call update_state() every decode step from GPUModelRunner._sample().
• Call apply_to_logits() from Sampler.forward(), mirroring the existing rejection_sampler.py call.

Both paths are now covered regardless of speculative decoding or async scheduling state.

Test plan

• Existing test_thinking_token_budget.py suite (6 non-MTP cases, default/auto_config params) passes unchanged
• Added 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
• Manual before/after repro documented above

🤖 Generated with Claude Code