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

Unify memory management across (overlap, non-overlap) x (page>=1) x (spec, non-spec, spec v2) x (retract, finished) (#12224)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-11-11 02:56:22 +08:00
committed by GitHub
parent 838bcb0d93
commit 665416f6dd
24 changed files with 193 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+14 -29
View File
@@ -51,6 +51,7 @@ from sglang.srt.managers.schedule_batch import FINISH_ABORT, RequestStage, Sched
from sglang.srt.managers.utils import GenerationBatchResult
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.common import release_kv_cache
from sglang.srt.mem_cache.memory_pool import (
HybridLinearKVPool,
HybridReqToTokenPool,
@@ -618,37 +619,21 @@ class DecodePreallocQueue:
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
req.kv_committed_len = fill_len
if self.token_to_kv_pool_allocator.page_size == 1:
kv_loc = self.token_to_kv_pool_allocator.alloc(
len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
)
kv_loc = self.token_to_kv_pool_allocator.alloc(fill_len)
else:
num_tokens = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
device = self.token_to_kv_pool_allocator.device
kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
prefix_lens=torch.tensor(
[0],
dtype=torch.int64,
device=self.token_to_kv_pool_allocator.device,
),
prefix_lens_cpu=torch.tensor(
[0],
dtype=torch.int64,
),
seq_lens=torch.tensor(
[num_tokens],
dtype=torch.int64,
device=self.token_to_kv_pool_allocator.device,
),
seq_lens_cpu=torch.tensor(
[num_tokens],
dtype=torch.int64,
),
last_loc=torch.tensor(
[-1],
dtype=torch.int64,
device=self.token_to_kv_pool_allocator.device,
),
extend_num_tokens=num_tokens,
prefix_lens=torch.tensor([0], dtype=torch.int64, device=device),
prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
last_loc=torch.tensor([-1], dtype=torch.int64, device=device),
extend_num_tokens=fill_len,
)
assert (
@@ -763,7 +748,7 @@ class DecodeTransferQueue:
[decode_req.req], decode_req.req.return_logprob
)
# release pre-allocated kv cache, but don't insert into the tree since it's failed
self.tree_cache.cache_finished_req(decode_req.req, is_insert=False)
release_kv_cache(decode_req.req, self.tree_cache, is_insert=False)
indices_to_remove.add(i)
if self.scheduler.enable_metrics:
self.scheduler.metrics_collector.increment_transfer_failed_reqs()
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+10 -4
View File
@@ -1,6 +1,9 @@
from __future__ import annotations
import logging
import threading
import time
from typing import TYPE_CHECKING
import torch
@@ -18,6 +21,9 @@ from sglang.srt.mem_cache.memory_pool_host import (
)
from sglang.srt.server_args import ServerArgs
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
logger = logging.getLogger(__name__)
@@ -177,14 +183,14 @@ class DecodeKVCacheOffloadManager:
)
finish_count -= 1
def _release_finished_req(self, req, prefill_offloaded_len):
def _release_finished_req(self, req: Req, prefill_offloaded_len: int):
# FIXME: not sure which length to use here: kv_allocated_len or kv_committed_len
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx,
: len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0),
req.req_pool_idx, prefill_offloaded_len : req.kv_allocated_len
]
# Free the incremental part of the request
self.token_to_kv_pool_allocator.free(kv_indices[prefill_offloaded_len:])
self.token_to_kv_pool_allocator.free(kv_indices)
self.req_to_token_pool.free(req.req_pool_idx)
def _check_backup_progress(self, finish_count):
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
+2 -1
View File
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.srt.disaggregation.utils import prepare_abort
from sglang.srt.mem_cache.common import release_kv_cache
from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
@@ -119,7 +120,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
# Grammar accept_token can raise ValueError if the token is not in the grammar.
# This can happen if the grammar is not set correctly or the token is invalid.
error_message = f"Grammar accept_token failed for req {req.rid} with token {req.output_ids[-1]}: {e}"
self.tree_cache.cache_finished_req(req)
release_kv_cache(req, self.tree_cache)
prepare_abort(
req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR
)
python/sglang/srt/disaggregation/prefill.py
+4 -3
View File
@@ -48,6 +48,7 @@ from sglang.srt.managers.schedule_batch import (
RequestStage,
ScheduleBatch,
)
from sglang.srt.mem_cache.common import release_kv_cache
from sglang.srt.mem_cache.memory_pool import (
HybridLinearKVPool,
NSATokenToKVPool,
@@ -468,7 +469,7 @@ class SchedulerDisaggregationPrefillMixin:
# Grammar accept_token can raise ValueError if the token is not in the grammar.
# This can happen if the grammar is not set correctly or the token is invalid.
error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
self.tree_cache.cache_finished_req(req)
release_kv_cache(req, self.tree_cache)
prepare_abort(
req,
error_message,
@@ -534,7 +535,7 @@ class SchedulerDisaggregationPrefillMixin:
if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]:
undone_reqs.append(req)
elif poll == KVPoll.Success: # transfer done
self.tree_cache.cache_finished_req(req) # unlock the tree
release_kv_cache(req, self.tree_cache) # unlock the tree
req.finished_reason = FINISH_LENGTH(length=0)
# FIXME: clean up req's data in transfer engine
if hasattr(req.disagg_kv_sender, "clear"):
@@ -547,7 +548,7 @@ class SchedulerDisaggregationPrefillMixin:
except Exception as e:
error_message += f" with exception {e}"
logger.warning(error_message)
self.tree_cache.cache_finished_req(req) # unlock the tree
release_kv_cache(req, self.tree_cache) # unlock the tree
prepare_abort(
req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR
)
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+3 -3
View File
@@ -8,7 +8,7 @@ from einops import rearrange
from sglang.srt.custom_op import CustomOp
from sglang.srt.layers.layernorm import LayerNorm
from sglang.srt.utils import add_prefix, align, is_cuda, is_hip, is_npu
from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu
if is_cuda():
try:
@@ -114,7 +114,7 @@ class Indexer(CustomOp):
self.fuse_wk_and_weights_proj = fuse_wk_and_weights_proj
if is_cuda():
self.sm_count = deep_gemm.get_num_sms()
self.half_device_sm_count = align(self.sm_count // 2, 8)
self.half_device_sm_count = ceil_align(self.sm_count // 2, 8)
self.wq_b = ReplicatedLinear(
self.q_lora_rank,
@@ -511,7 +511,7 @@ class Indexer(CustomOp):
end_pos = seq_len
topk_indices = index_score.topk(min(topk, end_pos), dim=-1)[1].squeeze(0)
pad_len = align(topk_indices.shape[-1], 2048) - topk_indices.shape[-1]
pad_len = ceil_align(topk_indices.shape[-1], 2048) - topk_indices.shape[-1]
topk_indices = torch.nn.functional.pad(
topk_indices, (0, pad_len), "constant", -1
)
python/sglang/srt/layers/quantization/fp8_kernel.py
+3 -3
View File
@@ -25,7 +25,7 @@ import triton.language as tl
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.utils import (
align,
ceil_align,
direct_register_custom_op,
get_bool_env_var,
get_device_core_count,
@@ -442,8 +442,8 @@ def create_per_token_group_quant_fp8_output_scale(
assert column_major_scales and scale_tma_aligned
*x_batch, x_q_mn, x_q_k = x_shape
x_s_mn, x_s_k = x_q_mn, x_q_k // 128
aligned_mn = align(x_s_mn, 4)
aligned_k = align(x_s_k, 4)
aligned_mn = ceil_align(x_s_mn, 4)
aligned_k = ceil_align(x_s_k, 4)
# TODO(FIXME): Fix cuda kernel and recover here to empty.
return torch.empty(
(*x_batch, aligned_k // 4, aligned_mn),
python/sglang/srt/layers/quantization/fp8_utils.py
+2 -2
View File
@@ -27,7 +27,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
w8a8_block_fp8_matmul_triton,
)
from sglang.srt.utils import (
align,
ceil_align,
get_bool_env_var,
get_cuda_version,
get_device_capability,
@@ -495,7 +495,7 @@ def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
x_padded = torch.zeros(
(align(m, 128), align(n, 128)), dtype=x.dtype, device=x.device
(ceil_align(m, 128), ceil_align(n, 128)), dtype=x.dtype, device=x.device
)
x_padded[:m, :n] = x
x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
python/sglang/srt/managers/schedule_batch.py
+50 -1
View File
@@ -66,6 +66,7 @@ from sglang.srt.mem_cache.common import (
alloc_for_decode,
alloc_for_extend,
evict_from_tree_cache,
release_kv_cache,
)
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
@@ -479,6 +480,12 @@ class Req:
self.session_id = session_id
self.input_embeds = input_embeds
# For req-level memory management
self.kv_committed_len = 0
self.kv_allocated_len = 0
self.kv_committed_freed = False
self.kv_overallocated_freed = False
# for corss-endoder model
self.token_type_ids = token_type_ids
@@ -693,6 +700,35 @@ class Req:
return self.output_ids[: self.finished_len]
return self.output_ids
def pop_committed_kv_cache(self) -> int:
"""Return the length of committed KV cache and mark them as freed."""
# NOTE: This function is called exactly once after the request is finished.
global_server_args = get_global_server_args()
topk = global_server_args.speculative_eagle_topk
enable_kv_committed_len = topk is None or topk == 1
if enable_kv_committed_len:
assert (
not self.kv_committed_freed
), f"Committed KV cache already freed ({self.kv_committed_len=})"
self.kv_committed_freed = True
return self.kv_committed_len
else:
return len(self.origin_input_ids) + max(len(self.output_ids) - 1, 0)
def pop_overallocated_kv_cache(self) -> Tuple[int, int]:
"""Return the range of over-allocated KV cache and mark them as freed."""
# NOTE: This function is called when there is over-allocation of KV cache.
# Over-allocation: we allocate more KV cache than the committed length.
# e.g., speculative decoding may allocate more KV cache than actually used.
assert (
not self.kv_overallocated_freed
), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv_allocated_len=}"
self.kv_overallocated_freed = True
return self.kv_committed_len, self.kv_allocated_len
def add_latency(self, stage: RequestStage):
if self.metrics_collector is None:
return
@@ -918,6 +954,10 @@ class Req:
self.is_chunked = 0
self.mamba_pool_idx = None
self.already_computed = 0
self.kv_allocated_len = 0
self.kv_committed_len = 0
self.kv_committed_freed = False
self.kv_overallocated_freed = False
def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
token_indices = req_to_token_pool.req_to_token[
@@ -1262,6 +1302,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
req.req_pool_idx = req_pool_indices[i]
assert seq_len - pre_len == req.extend_input_len
# update req-level memory management fields
req.kv_committed_len = seq_len
req.kv_allocated_len = seq_len
# If input_embeds are available, store them
if req.input_embeds is not None:
# If req.input_embeds is already a list, append its content directly
@@ -1536,7 +1580,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.req_to_token_pool, self.token_to_kv_pool_allocator
)
# TODO (csy): for preempted requests, we may want to insert into the tree
self.tree_cache.cache_finished_req(req, is_insert=False)
release_kv_cache(req, self.tree_cache, is_insert=False)
# NOTE(lsyin): we should use the newly evictable memory instantly.
num_tokens = remaing_req_count * envs.SGLANG_RETRACT_DECODE_STEPS.get()
evict_from_tree_cache(self.tree_cache, num_tokens)
@@ -1614,6 +1658,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Allocate memory
self.out_cache_loc = alloc_for_decode(self, token_per_req=1)
# Update req-level memory management fields
for req in self.reqs:
req.kv_committed_len += 1
req.kv_allocated_len += 1
# Update seq_lens after allocation
if self.enable_overlap:
# Do not use in-place operations in the overlap mode
python/sglang/srt/managers/scheduler.py
+3 -2
View File
@@ -149,6 +149,7 @@ from sglang.srt.managers.scheduler_update_weights_mixin import (
from sglang.srt.managers.session_controller import Session
from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length
from sglang.srt.mem_cache.chunk_cache import ChunkCache, SWAChunkCache
from sglang.srt.mem_cache.common import release_kv_cache
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
from sglang.srt.mem_cache.radix_cache import RadixCache
@@ -2410,11 +2411,11 @@ class Scheduler(
self.send_to_tokenizer.send_output(AbortReq(rid=req.rid), req)
# For disaggregation decode mode, the request in the waiting queue has KV cache allocated.
if self.disaggregation_mode == DisaggregationMode.DECODE:
self.tree_cache.cache_finished_req(req)
release_kv_cache(req, self.tree_cache)
# For mamba radix cache
if req.mamba_pool_idx is not None:
self.tree_cache.cache_finished_req(req, is_insert=False)
release_kv_cache(req, self.tree_cache, is_insert=False)
logger.debug(f"Abort queued request. {req.rid=}")
# Delete the requests in the grammar queue
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+15 -73
View File
@@ -20,8 +20,8 @@ from sglang.srt.managers.schedule_batch import (
RequestStage,
ScheduleBatch,
)
from sglang.srt.mem_cache.common import release_kv_cache
from sglang.srt.tracing.trace import trace_slice, trace_slice_batch, trace_slice_end
from sglang.srt.utils.common import ceil_div
if TYPE_CHECKING:
from sglang.srt.managers.scheduler import (
@@ -55,7 +55,7 @@ class SchedulerOutputProcessorMixin:
req.rid,
thread_finish_flag=True,
)
self.tree_cache.cache_finished_req(req)
release_kv_cache(req, self.tree_cache)
# Note: Logprobs should be handled on the prefill engine.
trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
@@ -102,26 +102,8 @@ class SchedulerOutputProcessorMixin:
logprob_pt = 0
for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
if self.enable_overlap and req.is_retracted and len(req.output_ids) > 0:
req_idx = batch.req_pool_indices[i]
seq_len = len(req.origin_input_ids) + len(req.output_ids)
pos = batch.req_to_token_pool.req_to_token[req_idx][
seq_len - 1 : seq_len
]
self.token_to_kv_pool_allocator.free(pos)
continue
if (
self.is_mixed_chunk
and self.enable_overlap
and (req.finished() or req.is_retracted)
):
# Free the one delayed token for the mixed decode batch
j = len(batch.out_cache_loc) - len(batch.reqs) + i
self.token_to_kv_pool_allocator.free(batch.out_cache_loc[j : j + 1])
continue
if req.is_retracted:
if req.finished() or req.is_retracted:
# decode req in mixed batch or retracted req
continue
if req.is_chunked <= 0:
@@ -130,7 +112,7 @@ class SchedulerOutputProcessorMixin:
req.check_finished()
if req.finished():
self.tree_cache.cache_finished_req(req)
release_kv_cache(req, self.tree_cache)
req.time_stats.completion_time = time.perf_counter()
elif not batch.decoding_reqs or req not in batch.decoding_reqs:
# This updates radix so others can match
@@ -259,7 +241,7 @@ class SchedulerOutputProcessorMixin:
req.check_finished()
if req.finished():
self.tree_cache.cache_finished_req(req)
release_kv_cache(req, self.tree_cache)
else:
self.tree_cache.cache_unfinished_req(req)
else:
@@ -327,42 +309,17 @@ class SchedulerOutputProcessorMixin:
self.token_to_kv_pool_allocator.free_group_begin()
# NOTE: in any case, we should check finish here
# if finished, also clean up committed kv cache and over-allocated kv cache here
# Check finish condition
# NOTE: the length of reqs and next_token_ids don't match if it is spec decoding.
# We should ignore using next_token_ids for spec decoding cases.
for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
req: Req
if self.enable_overlap and (req.finished() or req.is_retracted):
indices_to_free = None
if batch.spec_algorithm.is_eagle():
from sglang.srt.speculative.eagle_info import EagleDraftInput
end_p = allocate_lens_list[i]
start_p = end_p - EagleDraftInput.ALLOC_LEN_PER_DECODE
if self.page_size > 1:
start_p = ceil_div(start_p, self.page_size) * self.page_size
indices_to_free = self.req_to_token_pool.req_to_token[
req.req_pool_idx
][start_p:end_p]
else:
if self.page_size == 1:
# Free the one extra delayed token
indices_to_free = batch.out_cache_loc[i : i + 1]
else:
if (
len(req.origin_input_ids) + len(req.output_ids) - 1
) % self.page_size == 0:
# Only free when the extra token is in a new page
indices_to_free = batch.out_cache_loc[i : i + 1]
if indices_to_free is not None:
self.token_to_kv_pool_allocator.free(indices_to_free)
continue
if req.is_retracted:
# NOTE: This (req.finished() or req.is_retracted) should only happen when overlap scheduling is enabled.
# (currently not, e.g. Eagle V1 still check finish during forward)
# And all the over-allocated tokens will be freed in `release_kv_cache`.
continue
new_accepted_len = 1
@@ -376,27 +333,12 @@ class SchedulerOutputProcessorMixin:
req.check_finished(new_accepted_len)
if req.finished():
if batch.is_v2_eagle and self.cur_batch.forward_mode.is_extend():
# FIXME(lsyin): fix the messy logic here
# 1) when not overlap (v2 impl), we free the extra tokens in the req
# 2) overlap eagle and the current batch is prefill. This seq will not run extra iteration.
start_p = batch.seq_lens_cpu[i] + accept_lens_list[i]
end_p = allocate_lens_list[i]
if self.page_size > 1:
start_p = ceil_div(start_p, self.page_size) * self.page_size
indices_to_free = self.req_to_token_pool.req_to_token[
req.req_pool_idx
][start_p:end_p]
self.token_to_kv_pool_allocator.free(indices_to_free)
if self.server_args.disaggregation_decode_enable_offload_kvcache:
# Asynchronously offload KV cache; cache_finished_req will be called after Device->Host transfer completes
# Asynchronously offload KV cache; release_kv_cache will be called after Device->Host transfer completes
if not self.decode_offload_manager.offload_kv_cache(req):
self.tree_cache.cache_finished_req(req)
release_kv_cache(req, self.tree_cache)
else:
self.tree_cache.cache_finished_req(req)
release_kv_cache(req, self.tree_cache)
req.time_stats.completion_time = time.perf_counter()
python/sglang/srt/mem_cache/chunk_cache.py
+3 -3
View File
@@ -52,10 +52,10 @@ class ChunkCache(BasePrefixCache):
)
def cache_finished_req(self, req: Req, is_insert: bool = True):
kv_committed_len = req.pop_committed_kv_cache()
# For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx,
# For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids
: len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0),
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)
python/sglang/srt/mem_cache/common.py
+26
View File
@@ -14,6 +14,7 @@ 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
from sglang.srt.utils.common import ceil_align
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
@@ -462,6 +463,31 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
return out_cache_loc
def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):
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()
page_size = global_server_args.page_size
spec_algo = global_server_args.speculative_algorithm
if spec_algo is None:
assert (
start_p == end_p
), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv_allocated_len=}"
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)
def available_and_evictable_str(tree_cache) -> str:
token_to_kv_pool_allocator = tree_cache.token_to_kv_pool_allocator
if isinstance(token_to_kv_pool_allocator, SWATokenToKVPoolAllocator):
python/sglang/srt/mem_cache/mamba_radix_cache.py
+6 -6
View File
@@ -430,21 +430,21 @@ class MambaRadixCache(BasePrefixCache):
value = torch.tensor([x for x in key.token_ids], dtype=torch.int64)
return self._insert_helper(self.root_node, key, value, mamba_value)
def cache_finished_req(self, req: Req, is_insert=True) -> None:
def cache_finished_req(self, req: Req, is_insert: bool = True):
"""Cache request when it finishes."""
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,
: len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0),
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
cache_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
token_ids = (req.origin_input_ids + req.output_ids)[:cache_len]
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids)
req.req_pool_idx, :kv_committed_len
]
page_aligned_len = len(kv_indices)
python/sglang/srt/mem_cache/radix_cache.py
+5 -5
View File
@@ -341,21 +341,21 @@ class RadixCache(BasePrefixCache):
def cache_finished_req(self, req: Req, is_insert: bool = True):
"""Cache request when it finishes."""
all_token_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
committed_kv_len = req.pop_committed_kv_cache()
if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :all_token_len
req.req_pool_idx, :committed_kv_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)[:all_token_len]
token_ids = (req.origin_input_ids + req.output_ids)[:committed_kv_len]
# For EAGLE radix cache, we will convert the key to bigram key, e.g. [1,2,3,4] -> [(1,2), (2,3), (3,4)], the length will -1. ((len([(1,2), (2,3), (3,4)]) = len([1,2,3,4]) - 1))
# So for the corresponding kv length should also -1. Then we get the actual_kv_len, and use it to do later calculation and slicing.
actual_kv_len = all_token_len - 1 if self.is_eagle else all_token_len
actual_kv_len = committed_kv_len - 1 if self.is_eagle else committed_kv_len
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :all_token_len
req.req_pool_idx, :committed_kv_len
]
if self.page_size != 1:
python/sglang/srt/mem_cache/radix_cache_cpp.py
+7 -6
View File
@@ -165,15 +165,16 @@ class RadixCacheCpp(BasePrefixCache):
def cache_finished_req(self, req: Req, is_insert: bool = True):
"""Cache request when it finishes."""
assert req.req_pool_idx is not None
all_token_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
token_ids = (req.origin_input_ids + req.output_ids)[:all_token_len]
overall_len = len(token_ids) # prefill + decode
kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, :overall_len]
kv_committed_len = req.pop_committed_kv_cache()
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len
]
# NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned
# it will automatically align them, but length of them should be equal
old_prefix_len = len(req.prefix_indices) // self.page_size * self.page_size
page_aligned_overall_len = overall_len // self.page_size * self.page_size
page_aligned_overall_len = kv_committed_len // self.page_size * self.page_size
if is_insert:
new_prefix_len = self._insert(
@@ -190,7 +191,7 @@ class RadixCacheCpp(BasePrefixCache):
)
# need to free the unaligned part, since it cannot be inserted into the radix tree
if page_aligned_overall_len < overall_len:
if page_aligned_overall_len < kv_committed_len:
# NOTE: sglang PagedAllocator support unaligned free (which will automatically align it)
self.token_to_kv_pool.free(kv_indices[page_aligned_overall_len:])
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
+4 -3
View File
@@ -219,16 +219,17 @@ class LMCRadixCache(RadixCache):
return base_res
def cache_finished_req(self, req: "Req", is_insert: bool = True) -> None: # type: ignore[override]
def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: # type: ignore[override]
"""On request completion, insert device KV into radix and store to LMCache."""
super().cache_finished_req(req, is_insert=is_insert)
if not is_insert:
return
token_ids = (req.origin_input_ids + req.output_ids)[:-1]
kv_committed_len = req.pop_committed_kv_cache()
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids)
req.req_pool_idx, :kv_committed_len
]
_, new_last_node, _, _ = self.match_prefix(RadixKey(token_ids, req.extra_key))
python/sglang/srt/mem_cache/swa_radix_cache.py
+5 -5
View File
@@ -441,21 +441,21 @@ class SWARadixCache(BasePrefixCache):
def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
"""Cache request when it finishes."""
all_token_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
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, :all_token_len
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)[:all_token_len]
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
# For EAGLE radix cache, we will convert the key to bigram key, e.g. [1,2,3,4] -> [(1,2), (2,3), (3,4)], the length will -1. ((len([(1,2), (2,3), (3,4)]) = len([1,2,3,4]) - 1))
# So for the corresponding kv length should also -1. Then we get the actual_kv_len, and use it to do later calculation and slicing.
actual_kv_len = all_token_len - 1 if self.is_eagle else all_token_len
actual_kv_len = kv_committed_len - 1 if self.is_eagle else kv_committed_len
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :all_token_len
req.req_pool_idx, :kv_committed_len
]
if self.page_size != 1:
python/sglang/srt/speculative/eagle_info.py
+8
View File
@@ -116,6 +116,8 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
len(batch.input_ids),
)
end_offset = batch.seq_lens + self.draft_token_num
for req in batch.reqs:
req.kv_allocated_len += 1
else:
prefix_lens = batch.seq_lens
prefix_lens_cpu = batch.seq_lens_cpu
@@ -415,6 +417,9 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
if page_size == 1:
# TODO: boolean array index leads to a device sync. Remove it.
token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
for i, req in enumerate(batch.reqs):
req.kv_committed_len += accept_length_list[i] + 1
req.kv_allocated_len = req.kv_committed_len
else:
if self.topk == 1:
# Only evict full empty page. Do not evict partial empty page
@@ -426,6 +431,9 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
next_power_of_2(self.draft_token_num),
)
token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
for i, req in enumerate(batch.reqs):
req.kv_committed_len += accept_length_list[i] + 1
req.kv_allocated_len = req.kv_committed_len
else:
# Shift the accepted tokens to the beginning.
# Only evict the last part
python/sglang/srt/speculative/eagle_info_v2.py
+4
View File
@@ -129,6 +129,10 @@ class EagleDraftInputV2Mixin:
batch.seq_lens_cpu = batch.seq_lens.cpu()
batch.seq_lens_sum = batch.seq_lens_cpu.sum().item()
for i, req in enumerate(batch.reqs):
req.kv_committed_len = batch.seq_lens_cpu[i].item()
req.kv_allocated_len = req.kv_committed_len + self.ALLOC_LEN_PER_DECODE
def prepare_for_v2_draft(
self: EagleDraftInput,
req_to_token_pool: ReqToTokenPool,
python/sglang/srt/speculative/eagle_worker.py
+2
View File
@@ -364,6 +364,8 @@ class EAGLEWorker(TpModelWorker):
# [ topk 0 ] [ topk 1 ]
# [iter=0, iter=1, iter=2] [iter=0, iter=1, iter=2]
if self.page_size == 1:
for req in batch.reqs:
req.kv_allocated_len += self.speculative_num_steps * self.topk
out_cache_loc, token_to_kv_pool_state_backup = alloc_token_slots(
batch.tree_cache,
num_seqs * self.speculative_num_steps * self.topk,
python/sglang/srt/speculative/ngram_info.py
+10 -2
View File
@@ -195,7 +195,9 @@ class NgramVerifyInput(SpecInput):
logits_output.hidden_states = logits_output.hidden_states[self.accept_index]
self.verified_id = self.predict[self.accept_index]
def _free_cache(self, batch: ScheduleBatch, page_size: int):
def _free_cache(
self, batch: ScheduleBatch, page_size: int, accept_length_cpu: torch.Tensor
):
bs = batch.batch_size()
# Free the KV cache for unaccepted tokens
if page_size == 1:
@@ -250,6 +252,11 @@ class NgramVerifyInput(SpecInput):
)
batch.out_cache_loc = tgt_cache_loc
accept_length_list = accept_length_cpu.tolist()
for i, req in enumerate(batch.reqs):
req.kv_committed_len += accept_length_list[i] + 1
req.kv_allocated_len = req.kv_committed_len
assign_req_to_token_pool[(bs,)](
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
@@ -416,11 +423,12 @@ class NgramVerifyInput(SpecInput):
# self._sampling_verify(batch, logits_output, sampling_info)
self._fill_requests(batch, logits_output)
self._free_cache(batch, page_size)
accept_length_cpu = self.accept_length.cpu()
num_accepted_tokens = accept_length_cpu.sum().item()
self._free_cache(batch, page_size, accept_length_cpu)
batch.seq_lens.add_(self.accept_length + 1)
batch.seq_lens_cpu.add_(accept_length_cpu + 1)
python/sglang/srt/utils/common.py
+1 -1
View File
@@ -2981,7 +2981,7 @@ def configure_gc_logger():
# COPIED FROM DeepGEMM
def align(x: int, y: int) -> int:
def ceil_align(x: int, y: int) -> int:
return ceil_div(x, y) * y