Fix paged SWA retraction resume accounting (#33794)

Co-authored-by: zhisbug <1654062+zhisbug@users.noreply.github.com>
2026-08-06 16:32:54 -07:00
parent e0af47b03e
commit af7c62e337
2 changed files with 53 additions and 3 deletions
@@ -30,6 +30,51 @@ class FakeReceiver:


 class TestDecodeQueueCleanup(CustomTestCase):
+    def test_paged_swa_retraction_resume_uses_physical_page_budget(self):
+        page_size = 128
+        fill_len = 574
+        physical_tokens_per_req = 5 * page_size
+        physical_available = 18 * page_size
+
+        reqs = [
+            SimpleNamespace(
+                rid=f"req-{i}",
+                origin_input_ids=[0] * fill_len,
+                output_ids=[],
+                is_retracted=True,
+                load_kv_cache=MagicMock(),
+            )
+            for i in range(4)
+        ]
+
+        queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
+        queue.retracted_queue = reqs.copy()
+        queue.num_reserved_decode_tokens = 0
+        queue.req_to_token_pool = SimpleNamespace(available_size=lambda: len(reqs))
+        queue.token_to_kv_pool_allocator = SimpleNamespace(page_size=page_size)
+        queue.scheduler = SimpleNamespace(
+            sliding_window_size=2047,
+            server_args=SimpleNamespace(disable_radix_cache=True),
+        )
+        queue._uses_swa_tail_prealloc = MagicMock(return_value=True)
+        queue._swa_aware_allocatable_token_budgets = MagicMock(
+            return_value=(physical_available, physical_available)
+        )
+
+        def pre_alloc(_req):
+            nonlocal physical_available
+            self.assertGreaterEqual(physical_available, physical_tokens_per_req)
+            physical_available -= physical_tokens_per_req
+
+        queue._pre_alloc = MagicMock(side_effect=pre_alloc)
+
+        resumed = queue.resume_retracted_reqs()
+
+        self.assertEqual(resumed, reqs[:3])
+        self.assertEqual(queue.retracted_queue, reqs[3:])
+        self.assertEqual(physical_available, 3 * page_size)
+        self.assertEqual(queue._pre_alloc.call_count, 3)
+
    def test_prealloc_abort_clears_receiver_before_removing_request(self):
        receiver = FakeReceiver()
        req = SimpleNamespace(