Fix mamba radix cache eviction logic in alloc_req_slots (#11616)
Signed-off-by: rogeryoungh <rogeryoungh@foxmail.com>
This commit is contained in:
@@ -10,6 +10,7 @@ import triton.language as tl
|
||||
from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
|
||||
from sglang.srt.mem_cache.chunk_cache import ChunkCache, SWAChunkCache
|
||||
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
|
||||
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.utils import support_triton
|
||||
@@ -292,9 +293,15 @@ def alloc_req_slots(
|
||||
req_to_token_pool: ReqToTokenPool,
|
||||
num_reqs: int,
|
||||
reqs: list[Req] | None,
|
||||
tree_cache: BasePrefixCache | None,
|
||||
) -> list[int]:
|
||||
"""Allocate request slots from the pool."""
|
||||
if isinstance(req_to_token_pool, HybridReqToTokenPool):
|
||||
mamba_available_size = req_to_token_pool.mamba_pool.available_size()
|
||||
if mamba_available_size < num_reqs:
|
||||
if tree_cache is not None and isinstance(tree_cache, MambaRadixCache):
|
||||
mamba_num = max(0, num_reqs - mamba_available_size)
|
||||
tree_cache.evict_mamba(mamba_num)
|
||||
req_pool_indices = req_to_token_pool.alloc(num_reqs, reqs)
|
||||
else:
|
||||
req_pool_indices = req_to_token_pool.alloc(num_reqs)
|
||||
@@ -337,7 +344,9 @@ def alloc_for_extend(
|
||||
extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True)
|
||||
|
||||
# Allocate req slots
|
||||
req_pool_indices = alloc_req_slots(batch.req_to_token_pool, bs, batch.reqs)
|
||||
req_pool_indices = alloc_req_slots(
|
||||
batch.req_to_token_pool, bs, 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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user