minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Introduce req.kv container for coupled owned kv field lifecycle (#29427)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-15 14:40:38 +08:00
committed by GitHub
parent 201ddeaba1
commit d8d76c4d12
24 changed files with 132 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+2 -2
View File
@@ -1420,7 +1420,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
), "req_pool_indices is full! There is a bug in memory estimation."
fill_len = self._pre_alloc_fill_len(req)
req.kv_allocated_len = fill_len
req.kv.kv_allocated_len = fill_len
req.kv_committed_len = fill_len
if prefix_len > 0:
@@ -1525,7 +1525,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
extend_num_tokens=fill_len,
swa_tail_len=self._swa_tail_len(fill_len),
)
req.swa_evicted_seqlen = fill_len - self._swa_tail_len(fill_len)
req.kv.swa_evicted_seqlen = fill_len - self._swa_tail_len(fill_len)
else:
kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
prefix_lens=torch.tensor(
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py
+1 -1
View File
@@ -663,7 +663,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
if req is None or req_to_token_pool is None:
return
kv_len = max(req.kv_committed_len, req.kv_allocated_len)
kv_len = max(req.kv_committed_len, req.kv.kv_allocated_len)
req_pool_idx = req.req_pool_idx
if kv_len <= 0 or req_pool_idx is None:
return
python/sglang/srt/managers/hisparse_coordinator.py
+4 -4
View File
@@ -268,7 +268,7 @@ class HiSparseCoordinator:
"""
self.alloc_device_buffer(req)
host_len = self.host_token_len(req.kv_allocated_len)
host_len = self.host_token_len(req.kv.kv_allocated_len)
if host_len <= self.device_buffer_size:
# Short sequences (seq_len <= device_buffer_size): the kernel fast path
# returns device_buffer_locs directly without any host loading, so we
@@ -293,7 +293,7 @@ class HiSparseCoordinator:
def _preload_to_device_buffer(self, req: Req) -> None:
"""Preload all tokens from host pool into the device buffer."""
n = self.host_token_len(req.kv_allocated_len)
n = self.host_token_len(req.kv.kv_allocated_len)
host_indices = self.req_to_host_pool[req.req_pool_idx, :n]
device_locs = self.req_to_device_buffer[req.req_pool_idx, :n]
@@ -311,7 +311,7 @@ class HiSparseCoordinator:
allocated_len = req.extend_range.end
alloc_size = self.padded_buffer_size
else:
allocated_len = req.kv_allocated_len
allocated_len = req.kv.kv_allocated_len
page_size = self.mem_pool_device.page_size
# Allocate only enough for current tokens (page-aligned).
# When prefill already fills device_buffer_size, include the reserved page.
@@ -766,7 +766,7 @@ class HiSparseCoordinator:
# we just freed via free_hisparse_indices(all_hi). If left set, the
# subsequent release_kv_cache -> allocator.free -> free_hisparse path
# re-frees them (double-free into the page allocator's free list).
allocated_len = req.kv_allocated_len
allocated_len = req.kv.kv_allocated_len
# release memory -- only free actually-allocated buffer indices
current_cap = int(self.req_device_buffer_size[req.req_pool_idx])
python/sglang/srt/managers/schedule_batch.py
+20 -15
View File
@@ -661,6 +661,17 @@ class ReqLogprob:
output_token_ids_logprobs_idx: Optional[list] = None
@dataclasses.dataclass(slots=True, kw_only=True)
class ReqKvInfo:
kv_allocated_len: int
# The length of KV that have been removed in swa cache.
# SWA KV cache eviction behavior differs by cache type:
# - Radix cache: KV in range [cache_protected_len, swa_evicted_seqlen) is freed manually in
# `ScheduleBatch.maybe_evict_swa`; KV in range [0, cache_protected_len) is freed during radix cache eviction.
# - Chunk cache: KV in range [0, swa_evicted_seqlen) is freed manually in `ScheduleBatch.maybe_evict_swa`.
swa_evicted_seqlen: int
class Req(ReqDllmMixin):
"""The input and output status of a request."""
@@ -737,19 +748,13 @@ class Req(ReqDllmMixin):
# For req-level memory management
self.kv_committed_len = 0
self.kv_allocated_len = 0
self.kv: ReqKvInfo = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
self.kv_committed_freed = False
self.kv_overallocated_freed = False
# for cross-encoder model
self.token_type_ids = token_type_ids
# The length of KV that have been removed in swa cache.
# SWA KV cache eviction behavior differs by cache type:
# - Radix cache: KV in range [cache_protected_len, swa_evicted_seqlen) is freed manually in
# `ScheduleBatch.maybe_evict_swa`; KV in range [0, cache_protected_len) is freed during radix cache eviction.
# - Chunk cache: KV in range [0, swa_evicted_seqlen) is freed manually in `ScheduleBatch.maybe_evict_swa`.
self.swa_evicted_seqlen = 0
# Tokens in [0, swa_evict_floor) are protected from SWA window eviction.
# This is used by prefill-aware SWA models such as Unlimited-OCR to keep prompt/image KV visible during decode.
self.swa_evict_floor: int = 0
@@ -1094,9 +1099,9 @@ class Req(ReqDllmMixin):
# 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=}"
), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv.kv_allocated_len=}"
self.kv_overallocated_freed = True
return self._cache_commit_len(), self.kv_allocated_len
return self._cache_commit_len(), self.kv.kv_allocated_len
def update_spec_correct_drafts_histogram(self, num_correct_drafts: int):
"""Update the speculative decoding acceptance histogram.
@@ -1518,11 +1523,11 @@ class Req(ReqDllmMixin):
self.mamba_cow_src_index = None
self.mamba_needs_clear = False
self.already_computed = 0
self.kv_allocated_len = 0
self.kv.kv_allocated_len = 0
self.kv_committed_len = 0
self.kv_committed_freed = False
self.kv_overallocated_freed = False
self.swa_evicted_seqlen = 0
self.kv.swa_evicted_seqlen = 0
self.extend_batch_idx = 0
self.decode_batch_idx = 0
@@ -2190,7 +2195,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# update req-level memory management fields
req.kv_committed_len = seq_len
req.kv_allocated_len = seq_len
req.kv.kv_allocated_len = seq_len
# If input_embeds are available, store them
if req.input_embeds is not None:
@@ -2558,8 +2563,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
reserve = get_alloc_reserve_per_decode()
total = 0
for r in requests:
x = max(0, r.kv_committed_len + reserve - r.kv_allocated_len)
cur = r.kv_allocated_len
x = max(0, r.kv_committed_len + reserve - r.kv.kv_allocated_len)
cur = r.kv.kv_allocated_len
nxt = cur + x
total += ceil_align(nxt, page_size) - ceil_align(cur, page_size)
return total
@@ -2790,7 +2795,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
for req in self.reqs:
req.decode_batch_idx += 1
req.kv_committed_len += 1
req.kv_allocated_len += 1
req.kv.kv_allocated_len += 1
# New-tensor avoids racing model_worker_batch refs queued for
# overlap forward.
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+4 -4
View File
@@ -250,7 +250,7 @@ class SchedulerInvariantChecker:
if req.kv_committed_freed or req.req_pool_idx is None:
continue
allocated_len = req.kv_allocated_len
allocated_len = req.kv.kv_allocated_len
if self.page_size > 1:
allocated_len = ceil_align(allocated_len, self.page_size)
assert req.cache_protected_len % self.page_size == 0
@@ -258,7 +258,7 @@ class SchedulerInvariantChecker:
full_uncached += allocated_len - req.cache_protected_len
if self.is_hybrid_swa:
swa_uncached += allocated_len - max(
req.cache_protected_len, req.swa_evicted_seqlen
req.cache_protected_len, req.kv.swa_evicted_seqlen
)
return full_uncached, swa_uncached
@@ -321,7 +321,7 @@ class SchedulerInvariantChecker:
f"req {req.rid}",
req.req_pool_idx,
req.kv_committed_len,
req.kv_allocated_len,
req.kv.kv_allocated_len,
)
sess = getattr(self.tree_cache, "slots", None)
if sess:
@@ -332,7 +332,7 @@ class SchedulerInvariantChecker:
f"slot {sid[:8]}",
slot.req_pool_idx,
slot.kv_committed_len,
slot.kv_allocated_len,
slot.kv.kv_allocated_len,
)
active = [
python/sglang/srt/mem_cache/chunk_cache.py
+1 -1
View File
@@ -156,7 +156,7 @@ class PureSWAChunkCache(SWAChunkCache):
req.req_pool_idx, :kv_committed_len
]
evict_floor = req.swa_evict_floor
evicted_seqlen = req.swa_evicted_seqlen
evicted_seqlen = req.kv.swa_evicted_seqlen
if evicted_seqlen > evict_floor:
parts = []
if evict_floor > 0:
python/sglang/srt/mem_cache/common.py
+6 -6
View File
@@ -77,7 +77,7 @@ def free_swa_out_of_window_slots(
evict_floor = max(req.cache_protected_len, getattr(req, "swa_evict_floor", 0))
if page_size > 1 and evict_floor > req.cache_protected_len:
evict_floor = -(-evict_floor // page_size) * page_size
req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, evict_floor)
req.kv.swa_evicted_seqlen = max(req.kv.swa_evicted_seqlen, evict_floor)
if is_chunk_cache:
# Chunk cache builds no radix tree, so no tombstone-leaf concern; evict
@@ -90,22 +90,22 @@ def free_swa_out_of_window_slots(
# No extra page margin is needed.
evict_threshold = pre_len - max(sliding_window_size, page_size)
new_swa_evicted_seqlen = max(
req.swa_evicted_seqlen,
req.kv.swa_evicted_seqlen,
evict_threshold,
)
if page_size > 1:
new_swa_evicted_seqlen = (new_swa_evicted_seqlen // page_size) * page_size
if new_swa_evicted_seqlen > req.swa_evicted_seqlen:
if new_swa_evicted_seqlen > req.kv.swa_evicted_seqlen:
free_slots = req_to_token_pool.req_to_token[
req.req_pool_idx, req.swa_evicted_seqlen : new_swa_evicted_seqlen
req.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen
]
token_to_kv_pool_allocator.free_swa(free_slots)
maybe_evict_dsv4_state_on_swa(
token_to_kv_pool_allocator, req_to_token_pool, req, new_swa_evicted_seqlen
)
req.swa_evicted_seqlen = new_swa_evicted_seqlen
req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen
def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
@@ -661,7 +661,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
if spec_algo is None and not global_server_args.strip_thinking_cache:
assert (
start_p == end_p
), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv_allocated_len=}"
), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv.kv_allocated_len=}"
if page_size > 1:
start_p = ceil_align(start_p, page_size)
python/sglang/srt/mem_cache/pure_swa_radix_cache.py
+1 -1
View File
@@ -93,7 +93,7 @@ class PureSWARadixCache(RadixCache):
old_prefix_len = req.cache_protected_len
swa_evict_floor = req.swa_evict_floor
swa_evicted_seqlen = req.swa_evicted_seqlen
swa_evicted_seqlen = req.kv.swa_evicted_seqlen
if self.page_size > 1 and swa_evict_floor > 0:
swa_evict_floor = -(-swa_evict_floor // self.page_size) * self.page_size
python/sglang/srt/mem_cache/swa_radix_cache.py
+1 -1
View File
@@ -486,7 +486,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
key=radix_key,
value=values,
prev_prefix_len=old_prefix_len,
swa_evicted_seqlen=req.swa_evicted_seqlen,
swa_evicted_seqlen=req.kv.swa_evicted_seqlen,
)
)
else:
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+2 -2
View File
@@ -589,7 +589,7 @@ class SWAComponent(TreeComponent):
) -> Optional[int]:
# Unfinished requests can already have an SWA-evicted prefix; preserve
# that boundary so insertion creates a tombstone instead of live SWA KV.
insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen
insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
return None
def free_out_of_window_slots(
@@ -604,7 +604,7 @@ class SWAComponent(TreeComponent):
req_to_token_pool=self.cache.req_to_token_pool,
token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator,
)
insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen
insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
# ---- HiCache Hooks ----
python/sglang/srt/session/streaming_session.py
+31 -25
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import copy
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, Optional
@@ -20,7 +21,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
from sglang.srt.utils.common import ceil_align
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_batch import Req, ReqKvInfo
logger = logging.getLogger(__name__)
@@ -36,6 +37,12 @@ class _VirtualNode:
pass
def _new_kv() -> ReqKvInfo:
from sglang.srt.managers.schedule_batch import ReqKvInfo
return ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
@dataclass
class SessionSlot:
"""Holds KV state between streaming session turns."""
@@ -45,16 +52,13 @@ class SessionSlot:
# KV pool state (None means no KV is currently held by this slot)
req_pool_idx: Optional[int] = None
kv_committed_len: int = 0
kv_allocated_len: int = 0
kv: ReqKvInfo = field(default_factory=_new_kv)
# First req's radix tree node (for dec_lock_ref on session close)
last_node: Any = None
cache_protected_len: int = 0
swa_uuid_for_lock: Optional[str] = None
# SWA state
swa_evicted_seqlen: int = 0
# Mamba states
mamba_pool_idx: Any = None
mamba_ping_pong_track_buffer: Any = None
@@ -71,8 +75,7 @@ class SessionSlot:
"""Save KV state from a finishing request into this slot."""
self.req_pool_idx = req.req_pool_idx
self.kv_committed_len = req.kv_committed_len
self.kv_allocated_len = req.kv_allocated_len
self.swa_evicted_seqlen = req.swa_evicted_seqlen
self.kv = copy.copy(req.kv)
if is_first:
self.last_node = req.last_node
@@ -104,8 +107,7 @@ class SessionSlot:
"""Restore KV state from this slot into an incoming request."""
req.req_pool_idx = self.req_pool_idx
req.kv_committed_len = self.kv_committed_len
req.kv_allocated_len = self.kv_allocated_len
req.swa_evicted_seqlen = self.swa_evicted_seqlen
req.kv = copy.copy(self.kv)
req.swa_uuid_for_lock = self.swa_uuid_for_lock
req.mamba_pool_idx = self.mamba_pool_idx
@@ -305,7 +307,7 @@ class StreamingSession(BasePrefixCache):
# the mamba pool; otherwise the abort orphans them.
slot = SessionSlot(
req_pool_idx=req.req_pool_idx,
kv_allocated_len=req.kv_allocated_len,
kv=copy.copy(req.kv),
last_node=req.last_node,
cache_protected_len=req.cache_protected_len,
swa_uuid_for_lock=req.swa_uuid_for_lock,
@@ -317,7 +319,9 @@ class StreamingSession(BasePrefixCache):
# the abort fall-through doesn't double-free.
req.mamba_pool_idx = None
req.mamba_ping_pong_track_buffer = None
slot.kv_allocated_len = max(slot.kv_allocated_len, req.kv_allocated_len)
slot.kv.kv_allocated_len = max(
slot.kv.kv_allocated_len, req.kv.kv_allocated_len
)
self.release_session(session_id)
req.req_pool_idx = None
req.session.abort_req()
@@ -339,7 +343,7 @@ class StreamingSession(BasePrefixCache):
# req clock (under overlap + honest committed the clock lags the in-flight
# verify by ~1, which would short-change inheritance). Clamp to allocated
# to keep committed <= allocated for prepare_for_decode.
slot.kv_committed_len = min(target, slot.kv_allocated_len)
slot.kv_committed_len = min(target, slot.kv.kv_allocated_len)
# Update req_nodes to this successfully finished request.
req.session.finish_req(req)
@@ -411,7 +415,9 @@ class StreamingSession(BasePrefixCache):
protected_len = slot.cache_protected_len
lock_node = slot.last_node
tokens_freed = (
max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0
max(0, slot.kv.kv_allocated_len - protected_len)
if slot.is_holding_kv
else 0
)
logger.info(
"Session KV released: %s (%d tokens freed)", session_id, tokens_freed
@@ -428,7 +434,7 @@ class StreamingSession(BasePrefixCache):
if slot.is_holding_kv:
start = protected_len
end = slot.kv_allocated_len
end = slot.kv.kv_allocated_len
if start < end:
kv_indices = self.req_to_token_pool.req_to_token[
slot.req_pool_idx, start:end
@@ -454,7 +460,7 @@ class StreamingSession(BasePrefixCache):
active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs
)
if slot.is_holding_kv and not in_batch:
allocated = ceil_align(slot.kv_allocated_len, self.page_size)
allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
total += allocated - slot.cache_protected_len
return total
@@ -470,9 +476,9 @@ class StreamingSession(BasePrefixCache):
active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs
)
if slot.is_holding_kv and not in_batch:
allocated = ceil_align(slot.kv_allocated_len, self.page_size)
allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
total += allocated - max(
slot.cache_protected_len, slot.swa_evicted_seqlen
slot.cache_protected_len, slot.kv.swa_evicted_seqlen
)
return total
@@ -527,13 +533,13 @@ class StreamingSession(BasePrefixCache):
decoding pushes allocated above committed, or when retract retry's
logit-reserve pulls prefix_len below committed.
"""
self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv_allocated_len)
slot.kv_allocated_len = prefix_len
self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv.kv_allocated_len)
slot.kv.kv_allocated_len = prefix_len
slot.kv_committed_len = min(slot.kv_committed_len, prefix_len)
slot.swa_evicted_seqlen = min(slot.swa_evicted_seqlen, prefix_len)
req.kv_allocated_len = prefix_len
slot.kv.swa_evicted_seqlen = min(slot.kv.swa_evicted_seqlen, prefix_len)
req.kv.kv_allocated_len = prefix_len
req.kv_committed_len = min(req.kv_committed_len, prefix_len)
req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, prefix_len)
req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, prefix_len)
def _trim_overshoot(self, req: Req, finished_len: int) -> None:
"""Trim slot KV to finished_len boundary. Spec v2 may overshoot
@@ -542,10 +548,10 @@ class StreamingSession(BasePrefixCache):
be released to avoid token/KV mismatch.
"""
target = len(req.origin_input_ids) + finished_len
self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len)
req.kv_allocated_len = min(req.kv_allocated_len, target)
self._free_kv_aligned(req.req_pool_idx, target, req.kv.kv_allocated_len)
req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target)
req.kv_committed_len = min(req.kv_committed_len, target)
req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, target)
req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target)
req.output_ids = req.output_ids[:finished_len]
def _free_kv_aligned(self, pool_idx: int, target: int, end: int) -> None:
python/sglang/srt/speculative/dflash_info_v2.py
+4 -2
View File
@@ -155,7 +155,7 @@ class DFlashDraftInputV2(SpecInput):
for i, req in enumerate(batch.reqs):
committed_len = int(req.kv_committed_len)
# Read the allocation watermark from the req object like EAGLE.
cur_alloc_len = int(req.kv_allocated_len)
cur_alloc_len = int(req.kv.kv_allocated_len)
reserved_len = max(cur_alloc_len, committed_len + 2 * block_size)
top_k = int(req.sampling_params.top_k)
@@ -233,7 +233,9 @@ class DFlashDraftInputV2(SpecInput):
# This request-side high-water mark is what release_kv_cache() uses to
# reclaim any DFLASH over-allocation if the request finishes later.
for i, req in enumerate(batch.reqs):
req.kv_allocated_len = max(req.kv_allocated_len, int(nxt_kv_lens_cpu_t[i]))
req.kv.kv_allocated_len = max(
req.kv.kv_allocated_len, int(nxt_kv_lens_cpu_t[i])
)
# Seed committed; overlap's resolve overwrites it with the published value.
batch.seq_lens_cpu = batch_seq_lens_cpu_t
python/sglang/srt/speculative/eagle_utils.py
+2 -2
View File
@@ -833,7 +833,7 @@ def eagle_prepare_for_decode(batch: ScheduleBatch):
nxt_kv_lens = [0] * bs
num_needed_tokens = 0
for i, r in enumerate(batch.reqs):
cur = r.kv_allocated_len
cur = r.kv.kv_allocated_len
# max(cur, ...) clamps so adaptive downswitch cannot make nxt < cur.
# kv_committed_len is honest (bonus committed in resolve, not here),
# so it lags batch.seq_lens by ~1 verify in overlap; 2*alloc absorbs.
@@ -841,7 +841,7 @@ def eagle_prepare_for_decode(batch: ScheduleBatch):
cur_kv_lens[i] = cur
nxt_kv_lens[i] = nxt
num_needed_tokens += nxt - cur
r.kv_allocated_len = nxt
r.kv.kv_allocated_len = nxt
r.decode_batch_idx += 1
cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu")
python/sglang/test/scripted_runtime/req_handle.py
+1 -1
View File
@@ -45,7 +45,7 @@ class ScriptedReqHandle:
if req is None or req.req_pool_idx is None:
return 0
page_size = self.context.scheduler.page_size
return (req.kv_allocated_len + page_size - 1) // page_size
return (req.kv.kv_allocated_len + page_size - 1) // page_size
@property
def lock_refs(self) -> int:
test/registered/disaggregation/test_specv2_kvcache_offloading.py
+3 -2
View File
@@ -8,6 +8,7 @@ Requires: torch, sglang (run in an environment with sglang installed)
"""
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
import torch
@@ -34,7 +35,7 @@ def _make_mock_req(
req.rid = rid
req.req_pool_idx = req_pool_idx
req.kv_committed_len = kv_committed_len
req.kv_allocated_len = kv_allocated_len
req.kv = SimpleNamespace(kv_allocated_len=kv_allocated_len)
req.kv_committed_freed = False
req.kv_overallocated_freed = False
req.prefix_indices = list(range(prefix_indices_len))
@@ -47,7 +48,7 @@ def _make_mock_req(
def pop_overallocated():
assert not req.kv_overallocated_freed
req.kv_overallocated_freed = True
return req.kv_committed_len, req.kv_allocated_len
return req.kv_committed_len, req.kv.kv_allocated_len
req.pop_committed_kv_cache = pop_committed
req.pop_overallocated_kv_cache = pop_overallocated
test/registered/unit/managers/test_hisparse_unit.py
+5 -5
View File
@@ -50,7 +50,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
fill_ids=origin_input_ids + output_ids,
seqlen=len(origin_input_ids) + len(output_ids),
req_pool_idx=None,
kv_allocated_len=0,
kv=SimpleNamespace(kv_allocated_len=0),
kv_committed_len=0,
finished_reason=None,
hisparse_staging=False,
@@ -218,7 +218,7 @@ class TestHiSparseUnit(unittest.TestCase):
)
self.assertIsNotNone(kv_loc, "KV alloc failed")
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
req.kv_allocated_len = fill_len
req.kv.kv_allocated_len = fill_len
req.kv_committed_len = fill_len
req.full_untruncated_fill_ids = array("q", range(fill_len))
req.extend_range = Range(0, fill_len)
@@ -578,7 +578,7 @@ class TestHiSparseUnit(unittest.TestCase):
seq_len = fill_len + 1
self.req_to_token_pool.write((req.req_pool_idx, fill_len), out_loc)
req.kv_allocated_len = seq_len
req.kv.kv_allocated_len = seq_len
req.kv_committed_len = seq_len
self.coordinator.map_last_loc_to_buffer(
@@ -769,7 +769,7 @@ class TestHiSparseUnit(unittest.TestCase):
self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len],
)
)
self.assertEqual(req.kv_allocated_len, fill_len)
self.assertEqual(req.kv.kv_allocated_len, fill_len)
self.assertEqual(req.kv_committed_len, fill_len)
self.assertEqual(req.extend_range.length, fill_len)
@@ -786,7 +786,7 @@ class TestHiSparseUnit(unittest.TestCase):
self.assertEqual(allocated_host_indices.numel(), rounded_len)
kv_loc = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.kv_allocated_len
req.req_pool_idx, : req.kv.kv_allocated_len
].clone()
self._cleanup_req(req, kv_loc, logical_only=True)
self._assert_sizes_restored(initial, "pd_decode_prealloc_hisparse")
test/registered/unit/managers/test_kv_page_invariants.py
+2 -2
View File
@@ -48,14 +48,14 @@ class _FakeReq:
self.rid = rid
self.req_pool_idx = rpi
self.kv_committed_len = committed
self.kv_allocated_len = allocated
self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0)
class _FakeSlot:
def __init__(self, rpi, committed, allocated):
self.req_pool_idx = rpi
self.kv_committed_len = committed
self.kv_allocated_len = allocated
self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0)
self.is_holding_kv = True
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+3 -2
View File
@@ -27,6 +27,7 @@ register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
import unittest
from array import array
from types import SimpleNamespace
from unittest.mock import MagicMock
import torch
@@ -81,7 +82,7 @@ class MockReq:
self.prefix_indices = torch.empty(0, dtype=torch.int64)
self.priority = 0
self.kv_committed_len = len(fill_ids)
self.kv_allocated_len = len(fill_ids)
self.kv = SimpleNamespace(kv_allocated_len=len(fill_ids))
self.kv_committed_freed = False
def get_fill_ids(self):
@@ -92,7 +93,7 @@ class MockReq:
return self.kv_committed_len
def pop_overallocated_kv_cache(self):
return (self.kv_committed_len, self.kv_allocated_len)
return (self.kv_committed_len, self.kv.kv_allocated_len)
def _make_req(fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None):
test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py
+1 -1
View File
@@ -23,7 +23,7 @@ class _FakeAllocator:
class _FakeReq:
req_pool_idx = 0
swa_evict_floor = 3
swa_evicted_seqlen = 6
kv = SimpleNamespace(swa_evicted_seqlen=6)
def pop_committed_kv_cache(self):
return 8
test/registered/unit/mem_cache/test_streaming_session_unit.py
+11 -9
View File
@@ -57,13 +57,15 @@ class _FakeReq:
)
self.req_pool_idx = req_pool_idx
self.kv_committed_len = committed
self.kv_allocated_len = allocated
self.kv = SimpleNamespace(
kv_allocated_len=allocated,
swa_evicted_seqlen=0,
)
self.kv_committed_freed = False
self.kv_overallocated_freed = False
self.origin_input_ids = list(range(committed))
self.output_ids = []
self.extra_key = None
self.swa_evicted_seqlen = 0
self.last_node = None
self.cache_protected_len = 0
self.swa_uuid_for_lock = None
@@ -86,7 +88,7 @@ class _FakeReq:
assert not self.kv_overallocated_freed
self.pop_overallocated_calls += 1
self.kv_overallocated_freed = True
return self.kv_committed_len, self.kv_allocated_len
return self.kv_committed_len, self.kv.kv_allocated_len
def test_preabort_detaches_session_and_preserves_slot():
@@ -112,7 +114,7 @@ def test_preabort_detaches_session_and_preserves_slot():
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=48,
kv_allocated_len=48,
kv=SimpleNamespace(kv_allocated_len=48, swa_evicted_seqlen=0),
cache_protected_len=16,
)
@@ -132,7 +134,7 @@ def test_preabort_detaches_session_and_preserves_slot():
slot = tree_cache.slots["session-a"]
assert slot.req_pool_idx == 0
assert slot.kv_committed_len == 48
assert slot.kv_allocated_len == 48
assert slot.kv.kv_allocated_len == 48
assert len(result.device_indices) == 0
@@ -179,7 +181,7 @@ def test_nth_mid_abort_nukes_session_slot():
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=50,
kv_allocated_len=50,
kv=SimpleNamespace(kv_allocated_len=50, swa_evicted_seqlen=0),
last_node=None,
cache_protected_len=0,
)
@@ -230,14 +232,14 @@ def test_trim_overshoot_postcondition():
req = _FakeReq("session-a", req_pool_idx=0, committed=40, allocated=44)
req.origin_input_ids = list(range(26))
req.output_ids = list(range(14))
req.swa_evicted_seqlen = 42
req.kv.swa_evicted_seqlen = 42
tree_cache._trim_overshoot(req, finished_len=12)
target = 38
assert req.kv_committed_len == target
assert req.kv_allocated_len == target
assert req.swa_evicted_seqlen == target
assert req.kv.kv_allocated_len == target
assert req.kv.swa_evicted_seqlen == target
assert len(req.output_ids) == 12
# Tail [38, 44) freed by _free_kv_aligned.
assert len(allocator.freed) == 1
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
+15 -13
View File
@@ -106,7 +106,9 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
origin_input_ids=token_ids,
output_ids=[],
cache_protected_len=cache_protected_len,
swa_evicted_seqlen=0,
kv=SimpleNamespace(
swa_evicted_seqlen=0,
),
extra_key=None,
last_node=tree.root_node,
swa_uuid_for_lock=None,
@@ -162,7 +164,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
insert_len = seq_len // page_size * page_size
self.assertLess(
req.swa_evicted_seqlen,
req.kv.swa_evicted_seqlen,
insert_len,
f"page={page_size}, win={window}, seq={seq_len}",
)
@@ -187,7 +189,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
insert_len = seq_len // page_size * page_size
self.assertLess(req.swa_evicted_seqlen, insert_len)
self.assertLess(req.kv.swa_evicted_seqlen, insert_len)
tree.cache_finished_req(req, is_insert=True)
tree.sanity_check()
@@ -209,9 +211,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
batch = _make_batch(tree, allocator, pool)
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
self.assertLess(req.swa_evicted_seqlen, seq_len)
self.assertLess(req.kv.swa_evicted_seqlen, seq_len)
self.assertEqual(
req.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
)
tree.cache_finished_req(req, is_insert=True)
@@ -235,7 +237,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
batch = _make_batch(tree, allocator, pool)
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
self.assertEqual(req.swa_evicted_seqlen, 0)
self.assertEqual(req.kv.swa_evicted_seqlen, 0)
# -- Insert case 1: swa_evicted <= total_prefix_length --
@@ -264,7 +266,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
# pre_len=15: 15-2-8=5, floor to 8 -> 0. Eviction stays within matched.
ScheduleBatch._evict_swa(batch, req2, first_len - 1)
self.assertLessEqual(req2.swa_evicted_seqlen, first_len)
self.assertLessEqual(req2.kv.swa_evicted_seqlen, first_len)
swa_evictable_before = tree.swa_evictable_size_
tree.cache_finished_req(req2, is_insert=True)
@@ -295,12 +297,12 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
insert_len = seq_len // page_size * page_size
self.assertGreater(req.swa_evicted_seqlen, 0, "Should have some eviction")
self.assertLess(req.swa_evicted_seqlen, insert_len, "Should be partial")
self.assertGreater(req.kv.swa_evicted_seqlen, 0, "Should have some eviction")
self.assertLess(req.kv.swa_evicted_seqlen, insert_len, "Should be partial")
tree.cache_finished_req(req, is_insert=True)
non_tombstone = insert_len - req.swa_evicted_seqlen
non_tombstone = insert_len - req.kv.swa_evicted_seqlen
self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + non_tombstone)
self.assertGreater(tree.full_evictable_size_, 0)
tree.sanity_check()
@@ -333,7 +335,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
allocator.free_swa(pool.req_to_token[0, :old_evicted])
req = _make_req(0, list(range(seq_len)), 0, tree)
req.swa_evicted_seqlen = old_evicted
req.kv.swa_evicted_seqlen = old_evicted
swa_evictable_before = tree.swa_evictable_size_
tree.cache_finished_req(req, is_insert=True)
@@ -363,7 +365,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
insert_len = seq_len // page_size * page_size
self.assertLess(req.swa_evicted_seqlen, insert_len, f"turn {turn}")
self.assertLess(req.kv.swa_evicted_seqlen, insert_len, f"turn {turn}")
tree.cache_finished_req(req, is_insert=True)
tree.sanity_check()
@@ -386,7 +388,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
self.assertEqual(
req.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
)
tree.cache_finished_req(req, is_insert=True)
test/registered/unit/mem_cache/test_swa_unittest.py
+4 -2
View File
@@ -1,5 +1,6 @@
import unittest
from array import array
from types import SimpleNamespace
import torch
@@ -35,6 +36,7 @@ class _DummyReq:
def __init__(self):
self._kv_committed_len = 0
self.swa_prefix_lock_released = False
self.kv = SimpleNamespace(swa_evicted_seqlen=0)
def pop_committed_kv_cache(self):
return self._kv_committed_len
@@ -575,7 +577,7 @@ class TestSWA(unittest.TestCase):
req.extra_key = None
req.last_node = tree.root_node
req.swa_uuid_for_lock = None
req.swa_evicted_seqlen = 0
req.kv.swa_evicted_seqlen = 0
req.cache_protected_len = 1
# Intentionally mismatch to ensure code does not use len(prefix_indices).
req.prefix_indices = torch.tensor([7, 8, 9, 10, 11], device=tree.device)
@@ -610,7 +612,7 @@ class TestSWA(unittest.TestCase):
req2.extra_key = None
req2.last_node = tree.root_node
req2.swa_uuid_for_lock = None
req2.swa_evicted_seqlen = 0
req2.kv.swa_evicted_seqlen = 0
req2.cache_protected_len = 1
req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device)
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+8 -8
View File
@@ -881,7 +881,7 @@ class UnifiedRadixCacheSuite:
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv_committed_len = kv_len
req.kv_allocated_len = kv_len
req.kv.kv_allocated_len = kv_len
req.last_node = cache.root_node
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
@@ -998,7 +998,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.swa_evicted_seqlen = evicted_len
req.kv.swa_evicted_seqlen = evicted_len
cache.cache_unfinished_req(req)
@@ -1220,7 +1220,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.swa_evicted_seqlen = 0
req.kv.swa_evicted_seqlen = 0
full_available_before_insert = allocator.full_attn_allocator.available_size()
@@ -1761,7 +1761,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.swa_evicted_seqlen = 0
req.kv.swa_evicted_seqlen = 0
swa_avail_before = allocator.swa_attn_allocator.available_size()
@@ -1771,10 +1771,10 @@ class UnifiedRadixCacheSuite:
cushion = max(self.cfg.sliding_window_size, self.cfg.page_size)
expected_evicted = (pre_len - 1) - cushion
self.assertEqual(
req.swa_evicted_seqlen,
req.kv.swa_evicted_seqlen,
expected_evicted,
f"swa_evicted_seqlen should advance to (pre_len-1) - cushion = "
f"{expected_evicted}, got {req.swa_evicted_seqlen}",
f"{expected_evicted}, got {req.kv.swa_evicted_seqlen}",
)
swa_avail_after = allocator.swa_attn_allocator.available_size()
@@ -1813,13 +1813,13 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.swa_evicted_seqlen = 0
req.kv.swa_evicted_seqlen = 0
with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
cache.cache_unfinished_req(req)
self.assertEqual(
req.swa_evicted_seqlen,
req.kv.swa_evicted_seqlen,
0,
"Nothing should be evicted when prefill fits inside the cushion",
)
test/registered/unit/spec/test_decode_bookkeeping_ownership.py
-2
View File
@@ -80,9 +80,7 @@ _OWNER_SITES = {
): 1,
# streaming session slot save/restore and tail trimming
(_SS, "SessionSlot.save_from_req", "kv_committed_len"): 1,
(_SS, "SessionSlot.save_from_req", "kv_allocated_len"): 1,
(_SS, "SessionSlot.restore_to_req", "kv_committed_len"): 1,
(_SS, "SessionSlot.restore_to_req", "kv_allocated_len"): 1,
(_SS, "StreamingSession._free_tail", "kv_committed_len"): 2,
(_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2,
(_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1,