A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 15:15 UTC
[Bugfix] Fix inverted condition causing thinking_token_budget to be silently ignored (#41674)
Signed-off-by: Keyi Li <likey6688@gmail.com> Co-authored-by: Keyi Li <likey6688@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Keyi Li committed
May 14, 2026 at 21:48 UTC
bf610c2f56764e1b30bc6065f4ceace3d6e59036
2 files changed
+73
-1
tests/v1/logits_processors/test_correctness.py
+72
@@ -1122,3 +1122,75 @@ def test_thinking_budget_holder_empty_end_tokens_disables_row():
1122
)
1123
h.update_state([[THINK_START_TOKEN_ID, 1]], None, None)
1124
assert h._state[0]["thinking_token_budget"] == -1
1125
+
1126
+
1127
+def test_thinking_budget_enforced_without_penalties():
1128
+ """Regression test for gpu_input_batch.py bug.
1129
+
1130
+ When thinking_budget_tracks_reqs=True and no penalties/bad_words are set,
1131
+ the old code computed needs_output_token_ids=False (inverted condition:
1132
+ ``or not thinking_budget_tracks_reqs``), causing update_state to receive
1133
+ an empty list and skip _update_think_state for every request.
1134
+
1135
+ Fix: changed ``or not thinking_budget_tracks_reqs`` to
1136
+ ``or thinking_budget_tracks_reqs`` so that output_token_ids is populated
1137
+ whenever the thinking budget state holder has tracked requests.
1138
+
1139
+ This test verifies that update_state correctly calls _update_think_state
1140
+ (setting in_end=True) when given the real output_token_ids, and that
1141
+ passing an empty list (the pre-fix behavior) prevents budget enforcement.
1142
+ """
1143
+ vc = VllmConfig()
1144
+ vc.reasoning_config = MockReasoningConfig()
1145
+ budget = 3 # allow 3 thinking tokens
1146
+
1147
+ h = ThinkingBudgetStateHolder(
1148
+ vc.reasoning_config,
1149
+ vc.scheduler_config.max_num_seqs,
1150
+ 0,
1151
+ torch.device("cpu"),
1152
+ False,
1153
+ )
1154
+ output_token_ids: list[int] = []
1155
+ h.sync_batch(
1156
+ BatchUpdate(
1157
+ batch_size=1,
1158
+ removed=(),
1159
+ added=[
1160
+ (
1161
+ 0,
1162
+ SamplingParams(thinking_token_budget=budget),
1163
+ None,
1164
+ output_token_ids,
1165
+ )
1166
+ ],
1167
+ moved=(),
1168
+ )
1169
+ )
1170
+ assert h.has_tracked_requests()
1171
+
1172
+ # Simulate the buggy behavior: update_state receives empty list.
1173
+ # _update_think_state is skipped → in_end stays False → no budget enforcement.
1174
+ h.update_state([], None, None)
1175
+ assert not h._state[0].get("in_end", False), (
1176
+ "With empty output_token_ids, in_end should stay False (budget not yet tracked)"
1177
+ )
1178
+
1179
+ # Simulate the correct behavior: output_token_ids is the live list.
1180
+ # Step 1: think-start token appears.
1181
+ output_token_ids.append(THINK_START_TOKEN_ID)
1182
+ h.update_state([output_token_ids], None, None)
1183
+ assert not h._state[0].get("in_end", False), (
1184
+ "Still within budget after 0 think tokens"
1185
+ )
1186
+
1187
+ # Steps 2–4: 3 thinking tokens (hits the budget exactly).
1188
+ for tok in [1, 2, 3]:
1189
+ output_token_ids.append(tok)
1190
+ h.update_state([output_token_ids], None, None)
1191
+
1192
+ # After exactly `budget` thinking tokens the holder should force end token.
1193
+ assert h._state[0].get("in_end", False), (
1194
+ "Budget exceeded: in_end should be True so that apply_to_logits "
1195
+ "forces the end token"
1196
+ )
vllm/v1/worker/gpu_input_batch.py
+1
-1
@@ -884,7 +884,7 @@ class InputBatch:
884
not self.no_penalties
885
or bool(self.bad_words_token_ids)
886
or self.logitsprocs_need_output_token_ids
887
- or not thinking_budget_tracks_reqs
887
+ or thinking_budget_tracks_reqs
888
)
889
output_token_ids = (
890
cast(list[list[int]], self.req_output_token_ids)