chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
cctry
2026-02-02 14:38:15 -08:00
committed by GitHub
parent cbf1500390
commit 027f314050
13 changed files with 107 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
python/sglang/srt/disaggregation/decode.py
View File

@@ -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

2
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
View File

@@ -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):

7
python/sglang/srt/disaggregation/prefill.py
View File

@@ -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():

5
python/sglang/srt/managers/scheduler.py
View File

@@ -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()

2
python/sglang/srt/managers/scheduler_pp_mixin.py
View File

@@ -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: "

1
python/sglang/srt/mem_cache/chunk_cache.py
View File

@@ -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):

44
python/sglang/srt/mem_cache/common.py
View File

@@ -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:

19
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -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)

85
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -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")

2
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -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):

1
python/sglang/srt/mem_cache/radix_cache_cpp.py
View File

@@ -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."""

2
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -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: