From c94cdd0ae03865c00b694ca6f13b2d2bc4360fcd Mon Sep 17 00:00:00 2001
From: Henry Su <henrysu4707@gmail.com>
Date: Sat, 15 Aug 2026 16:13:02 -0500
Subject: [PATCH] [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 | 95 +++++++++++++++++++
 vllm/v1/sample/thinking_budget_state.py       |  4 +-
 2 files changed, 97 insertions(+), 2 deletions(-)
 create mode 100644 tests/v1/sample/test_thinking_budget_state.py

diff --git a/tests/v1/sample/test_thinking_budget_state.py b/tests/v1/sample/test_thinking_budget_state.py
new file mode 100644
index 0000000000..943734e8ee
--- /dev/null
+++ b/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
diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py
index efac011177..43880bf656 100644
--- a/vllm/v1/sample/thinking_budget_state.py
+++ b/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: