From 027f314050cc89a0e4d770b00aa901f23cdc3b8d Mon Sep 17 00:00:00 2001 From: cctry Date: Mon, 2 Feb 2026 14:38:15 -0800 Subject: [PATCH] [Fix] data race in req_to_token pool (#17850) --- python/sglang/srt/disaggregation/decode.py | 37 ++++---- .../decode_kvcache_offload_manager.py | 2 +- python/sglang/srt/disaggregation/prefill.py | 7 -- python/sglang/srt/managers/scheduler.py | 5 -- .../sglang/srt/managers/scheduler_pp_mixin.py | 2 +- python/sglang/srt/mem_cache/chunk_cache.py | 1 - python/sglang/srt/mem_cache/common.py | 44 ++++++---- .../sglang/srt/mem_cache/mamba_radix_cache.py | 19 ++--- python/sglang/srt/mem_cache/memory_pool.py | 85 ++++++++++--------- python/sglang/srt/mem_cache/radix_cache.py | 2 - .../sglang/srt/mem_cache/radix_cache_cpp.py | 1 - .../sglang/srt/mem_cache/swa_radix_cache.py | 2 - .../radix_cache/test_mamba_unittest.py | 13 +-- 13 files changed, 107 insertions(+), 113 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index d61e8dd8a..a0130cab5 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -25,7 +25,7 @@ import time 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 @@ class DecodeReqToTokenPool: 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 + 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, free_index: Union[int, List[int]]): - if isinstance(free_index, (int,)): - self.free_slots.append(free_index) - else: - self.free_slots.extend(free_index) + 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 @@ class DecodePreallocQueue: 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 diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index bf1d22cdf..1cfe615a7 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -191,7 +191,7 @@ class DecodeKVCacheOffloadManager: # 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): diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 35762a744..aa86731fa 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -632,13 +632,6 @@ class SchedulerDisaggregationPrefillMixin: ) 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(): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 29feebf21..b7eb45c4a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1792,11 +1792,6 @@ class Scheduler( 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() diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index e793350ad..17c0fb8c3 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -644,7 +644,7 @@ class SchedulerPPMixin: 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: " diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index a4377989b..f8e053285 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -68,7 +68,6 @@ class ChunkCache(BasePrefixCache): 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): diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index f77a654ef..7c7e82988 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 3ade398ba..ad4ff29e8 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -499,13 +499,6 @@ class MambaRadixCache(BasePrefixCache): 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: @@ -513,7 +506,7 @@ class MambaRadixCache(BasePrefixCache): 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 @@ class MambaRadixCache(BasePrefixCache): 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) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 5363abcf1..a5ad78b89 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -133,7 +133,6 @@ class ReqToTokenPool: device: str, enable_memory_saver: bool, ): - memory_saver_adapter = TorchMemorySaverAdapter.create( enable=enable_memory_saver ) @@ -145,7 +144,6 @@ class ReqToTokenPool: 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 @@ class ReqToTokenPool: 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:] + 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] - 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) + 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 @@ class HybridReqToTokenPool(ReqToTokenPool): # 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 @@ class HybridReqToTokenPool(ReqToTokenPool): 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") diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 4a7bb7229..4ff789431 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -451,7 +451,6 @@ class RadixCache(BasePrefixCache): 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) return token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] @@ -485,7 +484,6 @@ class RadixCache(BasePrefixCache): self.token_to_kv_pool_allocator.free(kv_indices[len(keys) :]) # Remove req slot release the cache lock - self.req_to_token_pool.free(req.req_pool_idx) self.dec_lock_ref(req.last_node) def cache_unfinished_req(self, req: Req, chunked=False): diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index c2b46f7b6..6ed293d1b 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -198,7 +198,6 @@ class RadixCacheCpp(BasePrefixCache): # Remove req slot release the cache lock self.dec_lock_ref(req.last_node) - self.req_to_token_pool.free(req.req_pool_idx) def cache_unfinished_req(self, req: Req, chunked=False): """Cache request when it is unfinished.""" diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 4b07b841f..b992b4727 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -461,7 +461,6 @@ class SWARadixCache(BasePrefixCache): 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) return token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] @@ -512,7 +511,6 @@ class SWARadixCache(BasePrefixCache): self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:]) # Remove req slot release the cache lock - self.req_to_token_pool.free(req.req_pool_idx) self.dec_lock_ref(req.last_node, req.swa_uuid_for_lock) def cache_unfinished_req(self, req: Req, chunked=False) -> None: diff --git a/test/registered/radix_cache/test_mamba_unittest.py b/test/registered/radix_cache/test_mamba_unittest.py index 4cc231095..2d7c7d24a 100644 --- a/test/registered/radix_cache/test_mamba_unittest.py +++ b/test/registered/radix_cache/test_mamba_unittest.py @@ -116,24 +116,25 @@ class TestMamba(unittest.TestCase): ) # alloc req - req_index = req_to_token_pool.alloc(1, [req]) + req_to_token_pool.alloc([req]) assert req_to_token_pool.available_size() == max_num_reqs - 1 assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size - 1 # free req - req_to_token_pool.free(req_index) + req_to_token_pool.free_mamba_cache(req) + req_to_token_pool.free(req) assert req_to_token_pool.available_size() == max_num_reqs assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size # alloc req without free mamba cache req.mamba_pool_idx = None - req_index = req_to_token_pool.alloc(1, [req]) - req_to_token_pool.free(req_index, free_mamba_cache=False) + req_to_token_pool.alloc([req]) + req_to_token_pool.free(req) assert req_to_token_pool.available_size() == max_num_reqs assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size - 1 # alloc again - req_index = req_to_token_pool.alloc(1, [req]) + req_to_token_pool.alloc([req]) assert req_to_token_pool.available_size() == max_num_reqs - 1 assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size - 1 @@ -225,7 +226,7 @@ class TestMamba(unittest.TestCase): origin_input_ids=[], sampling_params=sampling_params, ) - req_to_token_pool.alloc(1, reqs=[req]) + req_to_token_pool.alloc([req]) return req mamba_pool = req_to_token_pool.mamba_pool