From c58953d90a56c03576b0a002581342a74bd01a23 Mon Sep 17 00:00:00 2001 From: Dmitrii Sergeev Date: Tue, 11 Aug 2026 20:26:05 +0200 Subject: [PATCH] O(1) slot allocation in ReqToTokenPool.alloc() (#32208) Co-authored-by: Zhiqiang Xie --- python/sglang/srt/mem_cache/memory_pool.py | 10 ++++++++-- .../mlx/test_attention_patching.py | 5 ++++- .../unit/mem_cache/test_dllm_fdfo_kv_reuse.py | 16 +++++++++++++--- 3 files changed, 25 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 5f3728ce8..55b1f7269 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -306,8 +306,14 @@ class ReqToTokenPool: need_size = len(reqs) - len(reusing) if need_size > len(self.free_slots): return None - select_index = self.free_slots[:need_size] - self.free_slots = self.free_slots[need_size:] + if need_size > 0: + # Pop from the tail: O(need_size), unlike a prefix pop which is + # O(len(free_slots)). + select_index = self.free_slots[-need_size:] + del self.free_slots[-need_size:] + else: + # Handled separately: free_slots[-0:] is the entire list, not []. + select_index = [] offset = 0 for r in reqs: if r.req_pool_idx is None: diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index f62f5b414..fb9d1e5ab 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -909,7 +909,10 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): auxiliary_state_idx = pool.get_auxiliary_state_indices(req.req_pool_idx) pool.free(req) - self.assertEqual(req_indices, [1]) + # Which free slot a fresh alloc gets is not semantically meaningful + # (see ReqToTokenPool.alloc); only pin that it's a real, valid slot. + self.assertEqual(len(req_indices), 1) + self.assertIn(req_indices[0], range(1, pool.size + 1)) self.assertIsNotNone(auxiliary_state_idx) self.assertIsNone(req.req_pool_idx) self.assertIsNotNone(req.mamba_pool_idx) diff --git a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py index 787ed1827..79222d6e9 100644 --- a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py +++ b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py @@ -136,10 +136,17 @@ class TestDllmFdfoKvReuse(unittest.TestCase): out, _, req_pool_indices_cpu = alloc_for_extend(batch) self.assertEqual(allocator.alloc_calls, [4]) - self.assertEqual(req_pool_indices_cpu.tolist(), [1, 2]) + # Allocation order is not semantically meaningful (ReqToTokenPool.alloc + # picks whichever free slot is cheapest to pop), so only pin the + # reused row's index and that the fresh row got a different, real slot. + self.assertEqual(req_pool_indices_cpu[0].item(), 1) + fresh_idx = req_pool_indices_cpu[1].item() + self.assertNotEqual(fresh_idx, 1) self.assertEqual(out.tolist(), [100, 101, 102, 103, 200, 201, 202, 203]) self.assertEqual(self.pool.req_to_token[1, 4:8].tolist(), [100, 101, 102, 103]) - self.assertEqual(self.pool.req_to_token[2, 4:8].tolist(), [200, 201, 202, 203]) + self.assertEqual( + self.pool.req_to_token[fresh_idx, 4:8].tolist(), [200, 201, 202, 203] + ) self.assertEqual(reused.kv.kv_allocated_len, 8) self.assertEqual(fresh.kv.kv_allocated_len, 8) @@ -174,7 +181,10 @@ class TestDllmFdfoKvReuse(unittest.TestCase): batch = _make_batch(self.pool, allocator, [reused, fresh], [4, 4]) out, _, req_pool_indices_cpu = alloc_for_extend(batch) - self.assertEqual(req_pool_indices_cpu.tolist(), [1, 2]) + # See test_alloc_for_extend_mixed_reuse_allocates_only_fresh_and_writes_rows: + # allocation order is not semantically meaningful. + self.assertEqual(req_pool_indices_cpu[0].item(), 1) + self.assertNotEqual(req_pool_indices_cpu[1].item(), 1) self.assertEqual(out.tolist(), [100, 101, 102, 103, 500, 501, 502, 503]) self.assertEqual( allocator.extend_calls,