[Bugfix][Sampling] Clear empty side on thinking-budget asymmetric SWAP (#49613)
Signed-off-by: Henry Su <henrysu4707@gmail.com>
Henry Su committed
Aug 15, 2026 at 16:13 UTC
c94cdd0ae03865c00b694ca6f13b2d2bc4360fcd
2 files changed
+97
-2
tests/v1/sample/test_thinking_budget_state.py
new
+95
@@ -0,0 +1,95 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for ThinkingBudgetStateHolder batch index moves."""
4
+
5
+import torch
6
+
7
+from vllm.sampling_params import SamplingParams
8
+from vllm.v1.sample.logits_processor.interface import (
9
+ BatchUpdate,
10
+ MoveDirectionality,
11
+)
12
+from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder
13
+
14
+
15
+class _MockReasoningConfig:
16
+ reasoning_start_token_ids = [151667]
17
+ reasoning_end_token_ids = [151668]
18
+
19
+
20
+def _make_holder() -> ThinkingBudgetStateHolder:
21
+ return ThinkingBudgetStateHolder(
22
+ _MockReasoningConfig(),
23
+ 8,
24
+ 0,
25
+ torch.device("cpu"),
26
+ False,
27
+ )
28
+
29
+
30
+def test_swap_budgeted_with_unbudgeted_clears_empty_side():
31
+ """Asymmetric SWAP must not leave the empty index sharing state."""
32
+ h = _make_holder()
33
+ h.sync_batch(
34
+ BatchUpdate(
35
+ batch_size=2,
36
+ removed=(),
37
+ added=[
38
+ (0, SamplingParams(thinking_token_budget=5), None, []),
39
+ (1, SamplingParams(), None, []),
40
+ ],
41
+ moved=(),
42
+ )
43
+ )
44
+ assert list(h._state.keys()) == [0]
45
+ budget_state = h._state[0]
46
+
47
+ h.sync_batch(
48
+ BatchUpdate(
49
+ batch_size=2,
50
+ removed=(),
51
+ added=(),
52
+ moved=[(0, 1, MoveDirectionality.SWAP)],
53
+ )
54
+ )
55
+ assert list(h._state.keys()) == [1]
56
+ assert h._state[1] is budget_state
57
+ assert h._state[1]["thinking_token_budget"] == 5
58
+
59
+ h.sync_batch(
60
+ BatchUpdate(
61
+ batch_size=2,
62
+ removed=(),
63
+ added=(),
64
+ moved=[(0, 1, MoveDirectionality.SWAP)],
65
+ )
66
+ )
67
+ assert list(h._state.keys()) == [0]
68
+ assert h._state[0] is budget_state
69
+
70
+
71
+def test_swap_exchanges_two_budgeted_states():
72
+ h = _make_holder()
73
+ h.sync_batch(
74
+ BatchUpdate(
75
+ batch_size=2,
76
+ removed=(),
77
+ added=[
78
+ (0, SamplingParams(thinking_token_budget=3), None, []),
79
+ (1, SamplingParams(thinking_token_budget=7), None, []),
80
+ ],
81
+ moved=(),
82
+ )
83
+ )
84
+ b0 = h._state[0]["thinking_token_budget"]
85
+ b1 = h._state[1]["thinking_token_budget"]
86
+ h.sync_batch(
87
+ BatchUpdate(
88
+ batch_size=2,
89
+ removed=(),
90
+ added=(),
91
+ moved=[(0, 1, MoveDirectionality.SWAP)],
92
+ )
93
+ )
94
+ assert h._state[0]["thinking_token_budget"] == b1
95
+ assert h._state[1]["thinking_token_budget"] == b0
vllm/v1/sample/thinking_budget_state.py
+2
-2
@@ -100,8 +100,8 @@ class ThinkingBudgetStateHolder:
100
101
for i1, i2, direction in batch_update.moved:
102
if direction == MoveDirectionality.SWAP:
103
- state1 = self._state.get(i1)
104
- state2 = self._state.get(i2)
103
+ state1 = self._state.pop(i1, None)
104
+ state2 = self._state.pop(i2, None)
105
if state1 is not None:
106
self._state[i2] = state1
107
if state2 is not None: