[Fix] data race in req_to_token pool (sgl-project#17850)

Author: cctry

python/sglang/srt/disaggregation/decode.py

@@ -25,7 +25,7 @@
 from collections import deque
 from dataclasses import dataclass
 from http import HTTPStatus
-from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Type
 
 import torch
 from torch.distributed import ProcessGroup
@@ -116,19 +116,31 @@ def write(self, indices, values):
     def available_size(self):
         return len(self.free_slots)
 
-    def alloc(self, need_size: int) -> List[int]:
+    def alloc(self, reqs: List["Req"]) -> Optional[List[int]]:
+        chunked = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None]
+        assert (
+            len(chunked) <= 1
+        ), "only one chunked request may reuse req_pool_idx in a batch"
+        assert all(
+            reqs[i].is_chunked > 0 or reqs[i].kv_committed_len > 0 for i in chunked
+        ), "request has req_pool_idx but is not chunked"
+
+        need_size = len(reqs) - len(chunked)
         if need_size > len(self.free_slots):
             return None
-
         select_index = self.free_slots[:need_size]
         self.free_slots = self.free_slots[need_size:]
-        return select_index
-
-    def free(self, free_index: Union[int, List[int]]):
-        if isinstance(free_index, (int,)):
-            self.free_slots.append(free_index)
-        else:
-            self.free_slots.extend(free_index)
+        offset = 0
+        for r in reqs:
+            if r.req_pool_idx is None:
+                r.req_pool_idx = select_index[offset]
+                offset += 1
+        return [r.req_pool_idx for r in reqs]
+
+    def free(self, req: "Req"):
+        assert req.req_pool_idx is not None, "request must have req_pool_idx"
+        self.free_slots.append(req.req_pool_idx)
+        req.req_pool_idx = None
 
     def clear(self):
         self.free_slots = list(range(self.size + self.pre_alloc_size))
@@ -652,17 +664,12 @@ def _allocatable_tokens(
 
     def _pre_alloc(self, req: Req) -> torch.Tensor:
         """Pre-allocate the memory for req_to_token and token_kv_pool"""
-        if isinstance(self.req_to_token_pool, HybridMambaDecodeReqToTokenPool):
-            req_pool_indices = self.req_to_token_pool.alloc(1, [req])
-        else:
-            req_pool_indices = self.req_to_token_pool.alloc(1)
+        req_pool_indices = self.req_to_token_pool.alloc([req])
 
         assert (
             req_pool_indices is not None
         ), "req_pool_indices is full! There is a bug in memory estimation."
 
-        req.req_pool_idx = req_pool_indices[0]
-
         # Alloc all tokens for the prebuilt req (except for the reserved input token for decoding)
         fill_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
         req.kv_allocated_len = fill_len

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -191,7 +191,7 @@ def _release_finished_req(self, req: Req, prefill_offloaded_len: int):
 
         # Free the incremental part of the request
         self.token_to_kv_pool_allocator.free(kv_indices)
-        self.req_to_token_pool.free(req.req_pool_idx)
+        self.req_to_token_pool.free(req)
         self.tree_cache.protected_size_ -= len(req.prefix_indices)
 
     def _check_backup_progress(self, finish_count):

python/sglang/srt/disaggregation/prefill.py

@@ -632,13 +632,6 @@ def process_prefill_chunk(self: Scheduler) -> None:
                 )
             else:
                 self.send_kv_chunk(self.chunked_req)
-            # chunked request keeps its rid but will get a new req_pool_idx
-            if self.tp_worker.model_runner.mambaish_config is not None:
-                self.req_to_token_pool.free(
-                    self.chunked_req.req_pool_idx, free_mamba_cache=False
-                )
-            else:
-                self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
             self.running_batch.batch_is_full = False
 
         if self.last_batch and self.last_batch.forward_mode.is_extend():

python/sglang/srt/managers/scheduler.py

@@ -1792,11 +1792,6 @@ def handle_batch_embedding_request(
 
     def stash_chunked_request(self, req: Req):
         self.tree_cache.cache_unfinished_req(req, chunked=True)
-        # Chunked request keeps its rid but will get a new req_pool_idx
-        if self.tp_worker.model_runner.mambaish_config is not None:
-            self.req_to_token_pool.free(req.req_pool_idx, free_mamba_cache=False)
-        else:
-            self.req_to_token_pool.free(req.req_pool_idx)
 
     def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
         self._abort_on_queued_timeout()

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -644,7 +644,7 @@ def profile_and_init_predictor(self: Scheduler):
                         req.req_pool_idx, : len(req.fill_ids)
                     ]
                     self.token_to_kv_pool_allocator.free(kv_indices)
-                    self.req_to_token_pool.free(req.req_pool_idx)
+                    self.req_to_token_pool.free(req)
 
             logger.info(
                 f"[PP Dynamic Chunk] [PP0] Profiled {len(seq_lens)} samples: "

python/sglang/srt/mem_cache/chunk_cache.py

@@ -68,7 +68,6 @@ def cache_finished_req(self, req: Req, is_insert: bool = True):
         kv_indices = self.req_to_token_pool.req_to_token[
             req.req_pool_idx, :kv_committed_len
         ]
-        self.req_to_token_pool.free(req.req_pool_idx)
         self.token_to_kv_pool_allocator.free(kv_indices)
 
     def cache_unfinished_req(self, req: Req, chunked=False):

python/sglang/srt/mem_cache/common.py

@@ -296,11 +296,11 @@ def alloc_paged_token_slots_extend(
 
 def alloc_req_slots(
     req_to_token_pool: ReqToTokenPool,
-    num_reqs: int,
-    reqs: list[Req] | None,
+    reqs: list[Req],
     tree_cache: BasePrefixCache | None,
 ) -> list[int]:
     """Allocate request slots from the pool."""
+    num_reqs = len(reqs)
     if isinstance(req_to_token_pool, HybridReqToTokenPool):
         mamba_available_size = req_to_token_pool.mamba_pool.available_size()
         factor = (
@@ -313,9 +313,7 @@ def alloc_req_slots(
             if tree_cache is not None and tree_cache.supports_mamba():
                 mamba_num = max(0, mamba_state_needed - mamba_available_size)
                 tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
-        req_pool_indices = req_to_token_pool.alloc(num_reqs, reqs)
-    else:
-        req_pool_indices = req_to_token_pool.alloc(num_reqs)
+    req_pool_indices = req_to_token_pool.alloc(reqs)
 
     if req_pool_indices is None:
         raise RuntimeError(
@@ -341,7 +339,6 @@ def alloc_for_extend(
     # free out-of-window swa tokens
     batch.maybe_evict_swa()
 
-    bs = len(batch.reqs)
     prefix_tensors = [r.prefix_indices for r in batch.reqs]
 
     # Create tensors for allocation
@@ -352,7 +349,7 @@ def alloc_for_extend(
 
     # Allocate req slots
     req_pool_indices = alloc_req_slots(
-        batch.req_to_token_pool, bs, batch.reqs, batch.tree_cache
+        batch.req_to_token_pool, batch.reqs, batch.tree_cache
     )
     req_pool_indices_cpu = torch.tensor(req_pool_indices, dtype=torch.int64)
     req_pool_indices_device = req_pool_indices_cpu.to(batch.device, non_blocking=True)
@@ -466,15 +463,21 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
 
 
 def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):
-    tree_cache.cache_finished_req(req, is_insert=is_insert)
-
     # MambaRadixCache may alloc mamba state before alloc KV cache
     if req.req_pool_idx is None:
         assert (
             tree_cache.supports_mamba()
-        ), "Only MambaRadixCache can handle abort with prefix cache hit before alloc"
+        ), "Only MambaRadixCache allow freeing before alloc"
+        # TODO (csy, hanming): clean up this early allocation logic
+        if req.mamba_pool_idx is not None:
+            tree_cache.req_to_token_pool.mamba_pool.free(
+                req.mamba_pool_idx.unsqueeze(-1)
+            )
+            req.mamba_pool_idx = None
         return
 
+    tree_cache.cache_finished_req(req, is_insert=is_insert)
+
     start_p, end_p = req.pop_overallocated_kv_cache()
 
     global_server_args = get_global_server_args()
@@ -489,13 +492,20 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
     if page_size > 1:
         start_p = ceil_align(start_p, page_size)
 
-    if start_p >= end_p:
-        return
-
-    indices_to_free = tree_cache.req_to_token_pool.req_to_token[req.req_pool_idx][
-        start_p:end_p
-    ]
-    tree_cache.token_to_kv_pool_allocator.free(indices_to_free)
+    if start_p < end_p:
+        indices_to_free = tree_cache.req_to_token_pool.req_to_token[req.req_pool_idx][
+            start_p:end_p
+        ]
+        tree_cache.token_to_kv_pool_allocator.free(indices_to_free)
+    # If the prefix cache doesn't manage mamba states, we must free them here.
+    if isinstance(tree_cache.req_to_token_pool, HybridReqToTokenPool) and (
+        not tree_cache.supports_mamba()
+    ):
+        assert (
+            req.mamba_pool_idx is not None
+        ), "mamba state is freed while the tree cache does not manage mamba states"
+        tree_cache.req_to_token_pool.free_mamba_cache(req)
+    tree_cache.req_to_token_pool.free(req)
 
 
 def available_and_evictable_str(tree_cache) -> str:

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -499,21 +499,14 @@ def insert(self, params: InsertParams) -> InsertResult:
 
     def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
         """Cache request when it finishes."""
-        # for abort with prefix cache hit and before alloc is called
-        if req.req_pool_idx is None:
-            if req.mamba_pool_idx is not None:
-                self.req_to_token_pool.mamba_pool.free(req.mamba_pool_idx.unsqueeze(-1))
-                req.mamba_pool_idx = None
-            return
-
         kv_committed_len = req.pop_committed_kv_cache()
 
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[
                 req.req_pool_idx, :kv_committed_len
             ]
             self.token_to_kv_pool_allocator.free(kv_indices)
-            self.req_to_token_pool.free(req.req_pool_idx)
+            self.req_to_token_pool.free_mamba_cache(req)
             return
 
         token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
@@ -588,11 +581,11 @@ def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
 
         free_mamba_cache = True if self.enable_mamba_extra_buffer else mamba_exist
 
-        self.req_to_token_pool.free(
-            req.req_pool_idx,
-            free_mamba_cache=free_mamba_cache,
-            mamba_ping_pong_track_buffer_to_keep=mamba_ping_pong_track_buffer_to_keep,
-        )
+        if free_mamba_cache:
+            self.req_to_token_pool.free_mamba_cache(
+                req,
+                mamba_ping_pong_track_buffer_to_keep=mamba_ping_pong_track_buffer_to_keep,
+            )
 
         self.dec_lock_ref(req.last_node)
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -133,7 +133,6 @@ def __init__(
         device: str,
         enable_memory_saver: bool,
     ):
-
         memory_saver_adapter = TorchMemorySaverAdapter.create(
             enable=enable_memory_saver
         )
@@ -145,7 +144,6 @@ def __init__(
             self.req_to_token = torch.zeros(
                 (size, max_context_len), dtype=torch.int32, device=device
             )
-
         self.free_slots = list(range(size))
 
     def write(self, indices, values):
@@ -154,20 +152,32 @@ def write(self, indices, values):
     def available_size(self):
         return len(self.free_slots)
 
-    def alloc(self, need_size: int) -> List[int]:
+    def alloc(self, reqs: list[Req]) -> Optional[List[int]]:
+        chunked = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None]
+        if not any(r.is_dllm() for r in reqs):
+            assert (
+                len(chunked) <= 1
+            ), "only one chunked request may reuse req_pool_idx in a batch"
+        assert all(
+            reqs[i].is_chunked > 0 or reqs[i].kv_committed_len > 0 for i in chunked
+        ), "request has req_pool_idx but is not chunked"
+
+        need_size = len(reqs) - len(chunked)
         if need_size > len(self.free_slots):
             return None
-
         select_index = self.free_slots[:need_size]
         self.free_slots = self.free_slots[need_size:]
-
-        return select_index
-
-    def free(self, free_index: Union[int, List[int]]):
-        if isinstance(free_index, (int,)):
-            self.free_slots.append(free_index)
-        else:
-            self.free_slots.extend(free_index)
+        offset = 0
+        for r in reqs:
+            if r.req_pool_idx is None:
+                r.req_pool_idx = select_index[offset]
+                offset += 1
+        return [r.req_pool_idx for r in reqs]
+
+    def free(self, req: Req):
+        assert req.req_pool_idx is not None, "request must have req_pool_idx"
+        self.free_slots.append(req.req_pool_idx)
+        req.req_pool_idx = None
 
     def clear(self):
         self.free_slots = list(range(self.size))
@@ -488,10 +498,9 @@ def _init_mamba_pool(
 
     # For chunk prefill req, we do not need to allocate mamba cache,
     # We could use allocated mamba cache instead.
-    def alloc(self, need_size: int, reqs: Optional[List["Req"]]) -> Optional[List[int]]:
-        assert reqs is not None
-        select_index = super().alloc(need_size)
-        if select_index == None:
+    def alloc(self, reqs: List["Req"]) -> Optional[List[int]]:
+        select_index = super().alloc(reqs)
+        if select_index is None:
             return None
 
         mamba_index = []
@@ -556,37 +565,29 @@ def get_mamba_ping_pong_other_idx(self, mamba_next_track_idx: int) -> int:
         else:
             return mamba_next_track_idx
 
-    # For chunk prefill, we can not free mamba cache, we need use it in the future
-    def free(
-        self,
-        free_index: Union[int, List[int]],
-        free_mamba_cache: bool = True,
-        mamba_ping_pong_track_buffer_to_keep: Optional[int] = None,
+    def free_mamba_cache(
+        self, req: "Req", mamba_ping_pong_track_buffer_to_keep: Optional[int] = None
     ):
-        if isinstance(free_index, (int,)):
-            free_index = [free_index]
-        super().free(free_index)
-        if free_mamba_cache:
-            mamba_index = self.req_index_to_mamba_index_mapping[free_index]
-            self.mamba_pool.free(mamba_index)
+        mamba_index = req.mamba_pool_idx
+        assert mamba_index is not None, "double free? mamba_index is None"
+        self.mamba_pool.free(mamba_index.unsqueeze(0))
+        req.mamba_pool_idx = None
 
-            if self.enable_mamba_extra_buffer:
+        if self.enable_mamba_extra_buffer:
+            mamba_ping_pong_track_buffer_to_free = (
+                self.req_index_to_mamba_ping_pong_track_buffer_mapping[req.req_pool_idx]
+            )
+            if mamba_ping_pong_track_buffer_to_keep is not None:
+                assert mamba_ping_pong_track_buffer_to_keep in [
+                    0,
+                    1,
+                ], f"mamba_ping_pong_track_buffer_to_keep must be 0 or 1, {mamba_ping_pong_track_buffer_to_keep=}"
+                idx_to_free = list(range(self.mamba_ping_pong_track_buffer_size))
+                idx_to_free.remove(mamba_ping_pong_track_buffer_to_keep)
                 mamba_ping_pong_track_buffer_to_free = (
-                    self.req_index_to_mamba_ping_pong_track_buffer_mapping[
-                        free_index
-                    ].squeeze(0)
+                    mamba_ping_pong_track_buffer_to_free[idx_to_free]
                 )
-                if mamba_ping_pong_track_buffer_to_keep is not None:
-                    assert mamba_ping_pong_track_buffer_to_keep in [
-                        0,
-                        1,
-                    ], f"mamba_ping_pong_track_buffer_to_keep must be 0 or 1, {mamba_ping_pong_track_buffer_to_keep=}"
-                    idx_to_free = list(range(self.mamba_ping_pong_track_buffer_size))
-                    idx_to_free.remove(mamba_ping_pong_track_buffer_to_keep)
-                    mamba_ping_pong_track_buffer_to_free = (
-                        mamba_ping_pong_track_buffer_to_free[idx_to_free]
-                    )
-                self.mamba_pool.free(mamba_ping_pong_track_buffer_to_free)
+            self.mamba_pool.free(mamba_ping_pong_track_buffer_to_free)
 
     def clear(self):
         logger.info("Reset HybridReqToTokenPool")