[Bugfix][Sampling] Clear empty side on thinking-budget asymmetric SWAP (#49613)

Signed-off-by: Henry Su <henrysu4707@gmail.com>

tests/v1/sample/test_thinking_budget_state.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for ThinkingBudgetStateHolder batch index moves."""
+
+import torch
+
+from vllm.sampling_params import SamplingParams
+from vllm.v1.sample.logits_processor.interface import (
+    BatchUpdate,
+    MoveDirectionality,
+)
+from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder
+
+
+class _MockReasoningConfig:
+    reasoning_start_token_ids = [151667]
+    reasoning_end_token_ids = [151668]
+
+
+def _make_holder() -> ThinkingBudgetStateHolder:
+    return ThinkingBudgetStateHolder(
+        _MockReasoningConfig(),
+        8,
+        0,
+        torch.device("cpu"),
+        False,
+    )
+
+
+def test_swap_budgeted_with_unbudgeted_clears_empty_side():
+    """Asymmetric SWAP must not leave the empty index sharing state."""
+    h = _make_holder()
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=2,
+            removed=(),
+            added=[
+                (0, SamplingParams(thinking_token_budget=5), None, []),
+                (1, SamplingParams(), None, []),
+            ],
+            moved=(),
+        )
+    )
+    assert list(h._state.keys()) == [0]
+    budget_state = h._state[0]
+
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=2,
+            removed=(),
+            added=(),
+            moved=[(0, 1, MoveDirectionality.SWAP)],
+        )
+    )
+    assert list(h._state.keys()) == [1]
+    assert h._state[1] is budget_state
+    assert h._state[1]["thinking_token_budget"] == 5
+
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=2,
+            removed=(),
+            added=(),
+            moved=[(0, 1, MoveDirectionality.SWAP)],
+        )
+    )
+    assert list(h._state.keys()) == [0]
+    assert h._state[0] is budget_state
+
+
+def test_swap_exchanges_two_budgeted_states():
+    h = _make_holder()
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=2,
+            removed=(),
+            added=[
+                (0, SamplingParams(thinking_token_budget=3), None, []),
+                (1, SamplingParams(thinking_token_budget=7), None, []),
+            ],
+            moved=(),
+        )
+    )
+    b0 = h._state[0]["thinking_token_budget"]
+    b1 = h._state[1]["thinking_token_budget"]
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=2,
+            removed=(),
+            added=(),
+            moved=[(0, 1, MoveDirectionality.SWAP)],
+        )
+    )
+    assert h._state[0]["thinking_token_budget"] == b1
+    assert h._state[1]["thinking_token_budget"] == b0

vllm/v1/sample/thinking_budget_state.py

@@ -100,8 +100,8 @@ class ThinkingBudgetStateHolder:
 
         for i1, i2, direction in batch_update.moved:
             if direction == MoveDirectionality.SWAP:
-                state1 = self._state.get(i1)
-                state2 = self._state.get(i2)
+                state1 = self._state.pop(i1, None)
+                state2 = self._state.pop(i2, None)
                 if state1 is not None:
                     self._state[i2] = state1
                 if state2 is not None: