O(1) slot allocation in ReqToTokenPool.alloc() (#32208)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

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)

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,