fix: preserve HiCache load-back allocator headroom

2026-05-08 00:18:33 +08:00
parent 1f074f434e
commit 95bacb8862
10 changed files with 172 additions and 37 deletions
@@ -442,6 +442,40 @@ class TestPrefillAdder(CustomTestCase):
        self.assertEqual(adder2.rem_chunk_tokens, 0)  # 3 - 3 = 0
        self.assertEqual(result3, AddReqResult.OTHER)

+    def test_host_load_back_passes_mem_quota(self):
+        running_batch = self.create_running_batch()
+        self.mock_token_allocator.available_size.return_value = 512
+        self.mock_tree_cache.init_load_back.return_value = (
+            __import__("torch").tensor([1, 2, 3, 4], dtype=__import__("torch").int64),
+            "loaded_node",
+        )
+        adder = self.create_adder(
+            running_batch,
+            page_size=64,
+            rem_input_tokens=4096,
+            rem_total_tokens=4096,
+        )
+        req = self.create_mock_req("req", priority=0, max_new_tokens=16)
+        req.extend_input_len = 256
+        req.host_hit_length = 128
+        req.prefix_indices = __import__("torch").empty(
+            (0,), dtype=__import__("torch").int64
+        )
+        req.last_node = object()
+        req.last_host_node = object()
+        req.fill_ids = list(range(256))
+        req.cache_protected_len = 0
+        req.set_extend_input_len = lambda value: setattr(req, "extend_input_len", value)
+        req.sampling_params.ignore_eos = False
+
+        result = adder.add_one_req(
+            req, has_chunked_req=False, truncation_align_size=None
+        )
+
+        self.assertNotEqual(result, AddReqResult.NO_TOKEN)
+        params = self.mock_tree_cache.init_load_back.call_args.args[0]
+        self.assertEqual(params.mem_quota, 256)
+

 if __name__ == "__main__":
    unittest.main()