minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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." ), "req_pool_indices is full! There is a bug in memory estimation."
fill_len = self._pre_alloc_fill_len(req) 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 req.kv_committed_len = fill_len
if prefix_len > 0: if prefix_len > 0:
@@ -1525,7 +1525,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
extend_num_tokens=fill_len, extend_num_tokens=fill_len,
swa_tail_len=self._swa_tail_len(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: else:
kv_loc = self.token_to_kv_pool_allocator.alloc_extend( kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
prefix_lens=torch.tensor( 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: if req is None or req_to_token_pool is None:
return 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 req_pool_idx = req.req_pool_idx
if kv_len <= 0 or req_pool_idx is None: if kv_len <= 0 or req_pool_idx is None:
return return
python/sglang/srt/managers/hisparse_coordinator.py
+4 -4
View File
@@ -268,7 +268,7 @@ class HiSparseCoordinator:
""" """
self.alloc_device_buffer(req) 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: if host_len <= self.device_buffer_size:
# Short sequences (seq_len <= device_buffer_size): the kernel fast path # Short sequences (seq_len <= device_buffer_size): the kernel fast path
# returns device_buffer_locs directly without any host loading, so we # 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: def _preload_to_device_buffer(self, req: Req) -> None:
"""Preload all tokens from host pool into the device buffer.""" """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] host_indices = self.req_to_host_pool[req.req_pool_idx, :n]
device_locs = self.req_to_device_buffer[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 allocated_len = req.extend_range.end
alloc_size = self.padded_buffer_size alloc_size = self.padded_buffer_size
else: else:
allocated_len = req.kv_allocated_len allocated_len = req.kv.kv_allocated_len
page_size = self.mem_pool_device.page_size page_size = self.mem_pool_device.page_size
# Allocate only enough for current tokens (page-aligned). # Allocate only enough for current tokens (page-aligned).
# When prefill already fills device_buffer_size, include the reserved page. # 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 # we just freed via free_hisparse_indices(all_hi). If left set, the
# subsequent release_kv_cache -> allocator.free -> free_hisparse path # subsequent release_kv_cache -> allocator.free -> free_hisparse path
# re-frees them (double-free into the page allocator's free list). # 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 # release memory -- only free actually-allocated buffer indices
current_cap = int(self.req_device_buffer_size[req.req_pool_idx]) 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 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): class Req(ReqDllmMixin):
"""The input and output status of a request.""" """The input and output status of a request."""
@@ -737,19 +748,13 @@ class Req(ReqDllmMixin):
# For req-level memory management # For req-level memory management
self.kv_committed_len = 0 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_committed_freed = False
self.kv_overallocated_freed = False self.kv_overallocated_freed = False
# for cross-encoder model # for cross-encoder model
self.token_type_ids = token_type_ids 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. # 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. # 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 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. # e.g., speculative decoding may allocate more KV cache than actually used.
assert ( assert (
not self.kv_overallocated_freed 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 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): def update_spec_correct_drafts_histogram(self, num_correct_drafts: int):
"""Update the speculative decoding acceptance histogram. """Update the speculative decoding acceptance histogram.
@@ -1518,11 +1523,11 @@ class Req(ReqDllmMixin):
self.mamba_cow_src_index = None self.mamba_cow_src_index = None
self.mamba_needs_clear = False self.mamba_needs_clear = False
self.already_computed = 0 self.already_computed = 0
self.kv_allocated_len = 0 self.kv.kv_allocated_len = 0
self.kv_committed_len = 0 self.kv_committed_len = 0
self.kv_committed_freed = False self.kv_committed_freed = False
self.kv_overallocated_freed = False self.kv_overallocated_freed = False
self.swa_evicted_seqlen = 0 self.kv.swa_evicted_seqlen = 0
self.extend_batch_idx = 0 self.extend_batch_idx = 0
self.decode_batch_idx = 0 self.decode_batch_idx = 0
@@ -2190,7 +2195,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# update req-level memory management fields # update req-level memory management fields
req.kv_committed_len = seq_len 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 input_embeds are available, store them
if req.input_embeds is not None: if req.input_embeds is not None:
@@ -2558,8 +2563,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
reserve = get_alloc_reserve_per_decode() reserve = get_alloc_reserve_per_decode()
total = 0 total = 0
for r in requests: for r in requests:
x = max(0, r.kv_committed_len + reserve - r.kv_allocated_len) x = max(0, r.kv_committed_len + reserve - r.kv.kv_allocated_len)
cur = r.kv_allocated_len cur = r.kv.kv_allocated_len
nxt = cur + x nxt = cur + x
total += ceil_align(nxt, page_size) - ceil_align(cur, page_size) total += ceil_align(nxt, page_size) - ceil_align(cur, page_size)
return total return total
@@ -2790,7 +2795,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
for req in self.reqs: for req in self.reqs:
req.decode_batch_idx += 1 req.decode_batch_idx += 1
req.kv_committed_len += 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 # New-tensor avoids racing model_worker_batch refs queued for
# overlap forward. # 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: if req.kv_committed_freed or req.req_pool_idx is None:
continue continue
allocated_len = req.kv_allocated_len allocated_len = req.kv.kv_allocated_len
if self.page_size > 1: if self.page_size > 1:
allocated_len = ceil_align(allocated_len, self.page_size) allocated_len = ceil_align(allocated_len, self.page_size)
assert req.cache_protected_len % self.page_size == 0 assert req.cache_protected_len % self.page_size == 0
@@ -258,7 +258,7 @@ class SchedulerInvariantChecker:
full_uncached += allocated_len - req.cache_protected_len full_uncached += allocated_len - req.cache_protected_len
if self.is_hybrid_swa: if self.is_hybrid_swa:
swa_uncached += allocated_len - max( 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 return full_uncached, swa_uncached
@@ -321,7 +321,7 @@ class SchedulerInvariantChecker:
f"req {req.rid}", f"req {req.rid}",
req.req_pool_idx, req.req_pool_idx,
req.kv_committed_len, req.kv_committed_len,
req.kv_allocated_len, req.kv.kv_allocated_len,
) )
sess = getattr(self.tree_cache, "slots", None) sess = getattr(self.tree_cache, "slots", None)
if sess: if sess:
@@ -332,7 +332,7 @@ class SchedulerInvariantChecker:
f"slot {sid[:8]}", f"slot {sid[:8]}",
slot.req_pool_idx, slot.req_pool_idx,
slot.kv_committed_len, slot.kv_committed_len,
slot.kv_allocated_len, slot.kv.kv_allocated_len,
) )
active = [ 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 req.req_pool_idx, :kv_committed_len
] ]
evict_floor = req.swa_evict_floor evict_floor = req.swa_evict_floor
evicted_seqlen = req.swa_evicted_seqlen evicted_seqlen = req.kv.swa_evicted_seqlen
if evicted_seqlen > evict_floor: if evicted_seqlen > evict_floor:
parts = [] parts = []
if evict_floor > 0: 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)) evict_floor = max(req.cache_protected_len, getattr(req, "swa_evict_floor", 0))
if page_size > 1 and evict_floor > req.cache_protected_len: if page_size > 1 and evict_floor > req.cache_protected_len:
evict_floor = -(-evict_floor // page_size) * page_size 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: if is_chunk_cache:
# Chunk cache builds no radix tree, so no tombstone-leaf concern; evict # 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. # No extra page margin is needed.
evict_threshold = pre_len - max(sliding_window_size, page_size) evict_threshold = pre_len - max(sliding_window_size, page_size)
new_swa_evicted_seqlen = max( new_swa_evicted_seqlen = max(
req.swa_evicted_seqlen, req.kv.swa_evicted_seqlen,
evict_threshold, evict_threshold,
) )
if page_size > 1: if page_size > 1:
new_swa_evicted_seqlen = (new_swa_evicted_seqlen // page_size) * page_size 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[ 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) token_to_kv_pool_allocator.free_swa(free_slots)
maybe_evict_dsv4_state_on_swa( maybe_evict_dsv4_state_on_swa(
token_to_kv_pool_allocator, req_to_token_pool, req, new_swa_evicted_seqlen 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): 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: if spec_algo is None and not global_server_args.strip_thinking_cache:
assert ( assert (
start_p == end_p 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: if page_size > 1:
start_p = ceil_align(start_p, page_size) 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 old_prefix_len = req.cache_protected_len
swa_evict_floor = req.swa_evict_floor 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: if self.page_size > 1 and swa_evict_floor > 0:
swa_evict_floor = -(-swa_evict_floor // self.page_size) * self.page_size 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, key=radix_key,
value=values, value=values,
prev_prefix_len=old_prefix_len, prev_prefix_len=old_prefix_len,
swa_evicted_seqlen=req.swa_evicted_seqlen, swa_evicted_seqlen=req.kv.swa_evicted_seqlen,
) )
) )
else: 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]: ) -> Optional[int]:
# Unfinished requests can already have an SWA-evicted prefix; preserve # Unfinished requests can already have an SWA-evicted prefix; preserve
# that boundary so insertion creates a tombstone instead of live SWA KV. # 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 return None
def free_out_of_window_slots( def free_out_of_window_slots(
@@ -604,7 +604,7 @@ class SWAComponent(TreeComponent):
req_to_token_pool=self.cache.req_to_token_pool, req_to_token_pool=self.cache.req_to_token_pool,
token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator, 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 ---- # ---- HiCache Hooks ----
python/sglang/srt/session/streaming_session.py
+31 -25
View File
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import copy
import logging import logging
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, Optional 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 from sglang.srt.utils.common import ceil_align
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_batch import Req, ReqKvInfo
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -36,6 +37,12 @@ class _VirtualNode:
pass pass
def _new_kv() -> ReqKvInfo:
from sglang.srt.managers.schedule_batch import ReqKvInfo
return ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
@dataclass @dataclass
class SessionSlot: class SessionSlot:
"""Holds KV state between streaming session turns.""" """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) # KV pool state (None means no KV is currently held by this slot)
req_pool_idx: Optional[int] = None req_pool_idx: Optional[int] = None
kv_committed_len: int = 0 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) # First req's radix tree node (for dec_lock_ref on session close)
last_node: Any = None last_node: Any = None
cache_protected_len: int = 0 cache_protected_len: int = 0
swa_uuid_for_lock: Optional[str] = None swa_uuid_for_lock: Optional[str] = None
# SWA state
swa_evicted_seqlen: int = 0
# Mamba states # Mamba states
mamba_pool_idx: Any = None mamba_pool_idx: Any = None
mamba_ping_pong_track_buffer: 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.""" """Save KV state from a finishing request into this slot."""
self.req_pool_idx = req.req_pool_idx self.req_pool_idx = req.req_pool_idx
self.kv_committed_len = req.kv_committed_len self.kv_committed_len = req.kv_committed_len
self.kv_allocated_len = req.kv_allocated_len self.kv = copy.copy(req.kv)
self.swa_evicted_seqlen = req.swa_evicted_seqlen
if is_first: if is_first:
self.last_node = req.last_node self.last_node = req.last_node
@@ -104,8 +107,7 @@ class SessionSlot:
"""Restore KV state from this slot into an incoming request.""" """Restore KV state from this slot into an incoming request."""
req.req_pool_idx = self.req_pool_idx req.req_pool_idx = self.req_pool_idx
req.kv_committed_len = self.kv_committed_len req.kv_committed_len = self.kv_committed_len
req.kv_allocated_len = self.kv_allocated_len req.kv = copy.copy(self.kv)
req.swa_evicted_seqlen = self.swa_evicted_seqlen
req.swa_uuid_for_lock = self.swa_uuid_for_lock req.swa_uuid_for_lock = self.swa_uuid_for_lock
req.mamba_pool_idx = self.mamba_pool_idx req.mamba_pool_idx = self.mamba_pool_idx
@@ -305,7 +307,7 @@ class StreamingSession(BasePrefixCache):
# the mamba pool; otherwise the abort orphans them. # the mamba pool; otherwise the abort orphans them.
slot = SessionSlot( slot = SessionSlot(
req_pool_idx=req.req_pool_idx, req_pool_idx=req.req_pool_idx,
kv_allocated_len=req.kv_allocated_len, kv=copy.copy(req.kv),
last_node=req.last_node, last_node=req.last_node,
cache_protected_len=req.cache_protected_len, cache_protected_len=req.cache_protected_len,
swa_uuid_for_lock=req.swa_uuid_for_lock, swa_uuid_for_lock=req.swa_uuid_for_lock,
@@ -317,7 +319,9 @@ class StreamingSession(BasePrefixCache):
# the abort fall-through doesn't double-free. # the abort fall-through doesn't double-free.
req.mamba_pool_idx = None req.mamba_pool_idx = None
req.mamba_ping_pong_track_buffer = 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) self.release_session(session_id)
req.req_pool_idx = None req.req_pool_idx = None
req.session.abort_req() req.session.abort_req()
@@ -339,7 +343,7 @@ class StreamingSession(BasePrefixCache):
# req clock (under overlap + honest committed the clock lags the in-flight # req clock (under overlap + honest committed the clock lags the in-flight
# verify by ~1, which would short-change inheritance). Clamp to allocated # verify by ~1, which would short-change inheritance). Clamp to allocated
# to keep committed <= allocated for prepare_for_decode. # 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. # Update req_nodes to this successfully finished request.
req.session.finish_req(req) req.session.finish_req(req)
@@ -411,7 +415,9 @@ class StreamingSession(BasePrefixCache):
protected_len = slot.cache_protected_len protected_len = slot.cache_protected_len
lock_node = slot.last_node lock_node = slot.last_node
tokens_freed = ( 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( logger.info(
"Session KV released: %s (%d tokens freed)", session_id, tokens_freed "Session KV released: %s (%d tokens freed)", session_id, tokens_freed
@@ -428,7 +434,7 @@ class StreamingSession(BasePrefixCache):
if slot.is_holding_kv: if slot.is_holding_kv:
start = protected_len start = protected_len
end = slot.kv_allocated_len end = slot.kv.kv_allocated_len
if start < end: if start < end:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
slot.req_pool_idx, start:end 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 active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs
) )
if slot.is_holding_kv and not in_batch: 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 total += allocated - slot.cache_protected_len
return total return total
@@ -470,9 +476,9 @@ class StreamingSession(BasePrefixCache):
active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs
) )
if slot.is_holding_kv and not in_batch: 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( total += allocated - max(
slot.cache_protected_len, slot.swa_evicted_seqlen slot.cache_protected_len, slot.kv.swa_evicted_seqlen
) )
return total return total
@@ -527,13 +533,13 @@ class StreamingSession(BasePrefixCache):
decoding pushes allocated above committed, or when retract retry's decoding pushes allocated above committed, or when retract retry's
logit-reserve pulls prefix_len below committed. logit-reserve pulls prefix_len below committed.
""" """
self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv_allocated_len) self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv.kv_allocated_len)
slot.kv_allocated_len = prefix_len slot.kv.kv_allocated_len = prefix_len
slot.kv_committed_len = min(slot.kv_committed_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) slot.kv.swa_evicted_seqlen = min(slot.kv.swa_evicted_seqlen, prefix_len)
req.kv_allocated_len = prefix_len req.kv.kv_allocated_len = prefix_len
req.kv_committed_len = min(req.kv_committed_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: def _trim_overshoot(self, req: Req, finished_len: int) -> None:
"""Trim slot KV to finished_len boundary. Spec v2 may overshoot """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. be released to avoid token/KV mismatch.
""" """
target = len(req.origin_input_ids) + finished_len target = len(req.origin_input_ids) + finished_len
self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len) self._free_kv_aligned(req.req_pool_idx, target, req.kv.kv_allocated_len)
req.kv_allocated_len = min(req.kv_allocated_len, target) req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target)
req.kv_committed_len = min(req.kv_committed_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] req.output_ids = req.output_ids[:finished_len]
def _free_kv_aligned(self, pool_idx: int, target: int, end: int) -> None: 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): for i, req in enumerate(batch.reqs):
committed_len = int(req.kv_committed_len) committed_len = int(req.kv_committed_len)
# Read the allocation watermark from the req object like EAGLE. # 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) reserved_len = max(cur_alloc_len, committed_len + 2 * block_size)
top_k = int(req.sampling_params.top_k) 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 # This request-side high-water mark is what release_kv_cache() uses to
# reclaim any DFLASH over-allocation if the request finishes later. # reclaim any DFLASH over-allocation if the request finishes later.
for i, req in enumerate(batch.reqs): 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. # Seed committed; overlap's resolve overwrites it with the published value.
batch.seq_lens_cpu = batch_seq_lens_cpu_t 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 nxt_kv_lens = [0] * bs
num_needed_tokens = 0 num_needed_tokens = 0
for i, r in enumerate(batch.reqs): 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. # max(cur, ...) clamps so adaptive downswitch cannot make nxt < cur.
# kv_committed_len is honest (bonus committed in resolve, not here), # 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. # 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 cur_kv_lens[i] = cur
nxt_kv_lens[i] = nxt nxt_kv_lens[i] = nxt
num_needed_tokens += nxt - cur num_needed_tokens += nxt - cur
r.kv_allocated_len = nxt r.kv.kv_allocated_len = nxt
r.decode_batch_idx += 1 r.decode_batch_idx += 1
cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu") 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: if req is None or req.req_pool_idx is None:
return 0 return 0
page_size = self.context.scheduler.page_size 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 @property
def lock_refs(self) -> int: 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 import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock from unittest.mock import MagicMock
import torch import torch
@@ -34,7 +35,7 @@ def _make_mock_req(
req.rid = rid req.rid = rid
req.req_pool_idx = req_pool_idx req.req_pool_idx = req_pool_idx
req.kv_committed_len = kv_committed_len 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_committed_freed = False
req.kv_overallocated_freed = False req.kv_overallocated_freed = False
req.prefix_indices = list(range(prefix_indices_len)) req.prefix_indices = list(range(prefix_indices_len))
@@ -47,7 +48,7 @@ def _make_mock_req(
def pop_overallocated(): def pop_overallocated():
assert not req.kv_overallocated_freed assert not req.kv_overallocated_freed
req.kv_overallocated_freed = True 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_committed_kv_cache = pop_committed
req.pop_overallocated_kv_cache = pop_overallocated 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, fill_ids=origin_input_ids + output_ids,
seqlen=len(origin_input_ids) + len(output_ids), seqlen=len(origin_input_ids) + len(output_ids),
req_pool_idx=None, req_pool_idx=None,
kv_allocated_len=0, kv=SimpleNamespace(kv_allocated_len=0),
kv_committed_len=0, kv_committed_len=0,
finished_reason=None, finished_reason=None,
hisparse_staging=False, hisparse_staging=False,
@@ -218,7 +218,7 @@ class TestHiSparseUnit(unittest.TestCase):
) )
self.assertIsNotNone(kv_loc, "KV alloc failed") self.assertIsNotNone(kv_loc, "KV alloc failed")
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc) 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.kv_committed_len = fill_len
req.full_untruncated_fill_ids = array("q", range(fill_len)) req.full_untruncated_fill_ids = array("q", range(fill_len))
req.extend_range = Range(0, fill_len) req.extend_range = Range(0, fill_len)
@@ -578,7 +578,7 @@ class TestHiSparseUnit(unittest.TestCase):
seq_len = fill_len + 1 seq_len = fill_len + 1
self.req_to_token_pool.write((req.req_pool_idx, fill_len), out_loc) 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 req.kv_committed_len = seq_len
self.coordinator.map_last_loc_to_buffer( 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.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.kv_committed_len, fill_len)
self.assertEqual(req.extend_range.length, 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) self.assertEqual(allocated_host_indices.numel(), rounded_len)
kv_loc = self.req_to_token_pool.req_to_token[ 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() ].clone()
self._cleanup_req(req, kv_loc, logical_only=True) self._cleanup_req(req, kv_loc, logical_only=True)
self._assert_sizes_restored(initial, "pd_decode_prealloc_hisparse") 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.rid = rid
self.req_pool_idx = rpi self.req_pool_idx = rpi
self.kv_committed_len = committed self.kv_committed_len = committed
self.kv_allocated_len = allocated self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0)
class _FakeSlot: class _FakeSlot:
def __init__(self, rpi, committed, allocated): def __init__(self, rpi, committed, allocated):
self.req_pool_idx = rpi self.req_pool_idx = rpi
self.kv_committed_len = committed 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 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 import unittest
from array import array from array import array
from types import SimpleNamespace
from unittest.mock import MagicMock from unittest.mock import MagicMock
import torch import torch
@@ -81,7 +82,7 @@ class MockReq:
self.prefix_indices = torch.empty(0, dtype=torch.int64) self.prefix_indices = torch.empty(0, dtype=torch.int64)
self.priority = 0 self.priority = 0
self.kv_committed_len = len(fill_ids) 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 self.kv_committed_freed = False
def get_fill_ids(self): def get_fill_ids(self):
@@ -92,7 +93,7 @@ class MockReq:
return self.kv_committed_len return self.kv_committed_len
def pop_overallocated_kv_cache(self): 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): 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: class _FakeReq:
req_pool_idx = 0 req_pool_idx = 0
swa_evict_floor = 3 swa_evict_floor = 3
swa_evicted_seqlen = 6 kv = SimpleNamespace(swa_evicted_seqlen=6)
def pop_committed_kv_cache(self): def pop_committed_kv_cache(self):
return 8 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.req_pool_idx = req_pool_idx
self.kv_committed_len = committed 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_committed_freed = False
self.kv_overallocated_freed = False self.kv_overallocated_freed = False
self.origin_input_ids = list(range(committed)) self.origin_input_ids = list(range(committed))
self.output_ids = [] self.output_ids = []
self.extra_key = None self.extra_key = None
self.swa_evicted_seqlen = 0
self.last_node = None self.last_node = None
self.cache_protected_len = 0 self.cache_protected_len = 0
self.swa_uuid_for_lock = None self.swa_uuid_for_lock = None
@@ -86,7 +88,7 @@ class _FakeReq:
assert not self.kv_overallocated_freed assert not self.kv_overallocated_freed
self.pop_overallocated_calls += 1 self.pop_overallocated_calls += 1
self.kv_overallocated_freed = True 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(): 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( tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0, req_pool_idx=0,
kv_committed_len=48, kv_committed_len=48,
kv_allocated_len=48, kv=SimpleNamespace(kv_allocated_len=48, swa_evicted_seqlen=0),
cache_protected_len=16, cache_protected_len=16,
) )
@@ -132,7 +134,7 @@ def test_preabort_detaches_session_and_preserves_slot():
slot = tree_cache.slots["session-a"] slot = tree_cache.slots["session-a"]
assert slot.req_pool_idx == 0 assert slot.req_pool_idx == 0
assert slot.kv_committed_len == 48 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 assert len(result.device_indices) == 0
@@ -179,7 +181,7 @@ def test_nth_mid_abort_nukes_session_slot():
tree_cache.slots["session-a"] = SessionSlot( tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0, req_pool_idx=0,
kv_committed_len=50, kv_committed_len=50,
kv_allocated_len=50, kv=SimpleNamespace(kv_allocated_len=50, swa_evicted_seqlen=0),
last_node=None, last_node=None,
cache_protected_len=0, 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 = _FakeReq("session-a", req_pool_idx=0, committed=40, allocated=44)
req.origin_input_ids = list(range(26)) req.origin_input_ids = list(range(26))
req.output_ids = list(range(14)) 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) tree_cache._trim_overshoot(req, finished_len=12)
target = 38 target = 38
assert req.kv_committed_len == target assert req.kv_committed_len == target
assert req.kv_allocated_len == target assert req.kv.kv_allocated_len == target
assert req.swa_evicted_seqlen == target assert req.kv.swa_evicted_seqlen == target
assert len(req.output_ids) == 12 assert len(req.output_ids) == 12
# Tail [38, 44) freed by _free_kv_aligned. # Tail [38, 44) freed by _free_kv_aligned.
assert len(allocator.freed) == 1 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, origin_input_ids=token_ids,
output_ids=[], output_ids=[],
cache_protected_len=cache_protected_len, cache_protected_len=cache_protected_len,
swa_evicted_seqlen=0, kv=SimpleNamespace(
swa_evicted_seqlen=0,
),
extra_key=None, extra_key=None,
last_node=tree.root_node, last_node=tree.root_node,
swa_uuid_for_lock=None, swa_uuid_for_lock=None,
@@ -162,7 +164,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
insert_len = seq_len // page_size * page_size insert_len = seq_len // page_size * page_size
self.assertLess( self.assertLess(
req.swa_evicted_seqlen, req.kv.swa_evicted_seqlen,
insert_len, insert_len,
f"page={page_size}, win={window}, seq={seq_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) ScheduleBatch._evict_swa(batch, req, seq_len - 1)
insert_len = seq_len // page_size * page_size 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.cache_finished_req(req, is_insert=True)
tree.sanity_check() tree.sanity_check()
@@ -209,9 +211,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
batch = _make_batch(tree, allocator, pool) batch = _make_batch(tree, allocator, pool)
ScheduleBatch._evict_swa(batch, req, seq_len - 1) 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( 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) tree.cache_finished_req(req, is_insert=True)
@@ -235,7 +237,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
batch = _make_batch(tree, allocator, pool) batch = _make_batch(tree, allocator, pool)
ScheduleBatch._evict_swa(batch, req, seq_len - 1) 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 -- # -- 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. # pre_len=15: 15-2-8=5, floor to 8 -> 0. Eviction stays within matched.
ScheduleBatch._evict_swa(batch, req2, first_len - 1) 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_ swa_evictable_before = tree.swa_evictable_size_
tree.cache_finished_req(req2, is_insert=True) tree.cache_finished_req(req2, is_insert=True)
@@ -295,12 +297,12 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1) ScheduleBatch._evict_swa(batch, req, seq_len - 1)
insert_len = seq_len // page_size * page_size insert_len = seq_len // page_size * page_size
self.assertGreater(req.swa_evicted_seqlen, 0, "Should have some eviction") self.assertGreater(req.kv.swa_evicted_seqlen, 0, "Should have some eviction")
self.assertLess(req.swa_evicted_seqlen, insert_len, "Should be partial") self.assertLess(req.kv.swa_evicted_seqlen, insert_len, "Should be partial")
tree.cache_finished_req(req, is_insert=True) 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.assertEqual(tree.swa_evictable_size_, swa_evictable_before + non_tombstone)
self.assertGreater(tree.full_evictable_size_, 0) self.assertGreater(tree.full_evictable_size_, 0)
tree.sanity_check() tree.sanity_check()
@@ -333,7 +335,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
allocator.free_swa(pool.req_to_token[0, :old_evicted]) allocator.free_swa(pool.req_to_token[0, :old_evicted])
req = _make_req(0, list(range(seq_len)), 0, tree) 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_ swa_evictable_before = tree.swa_evictable_size_
tree.cache_finished_req(req, is_insert=True) tree.cache_finished_req(req, is_insert=True)
@@ -363,7 +365,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1) ScheduleBatch._evict_swa(batch, req, seq_len - 1)
insert_len = seq_len // page_size * page_size 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.cache_finished_req(req, is_insert=True)
tree.sanity_check() tree.sanity_check()
@@ -386,7 +388,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1) ScheduleBatch._evict_swa(batch, req, seq_len - 1)
self.assertEqual( 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) 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 import unittest
from array import array from array import array
from types import SimpleNamespace
import torch import torch
@@ -35,6 +36,7 @@ class _DummyReq:
def __init__(self): def __init__(self):
self._kv_committed_len = 0 self._kv_committed_len = 0
self.swa_prefix_lock_released = False self.swa_prefix_lock_released = False
self.kv = SimpleNamespace(swa_evicted_seqlen=0)
def pop_committed_kv_cache(self): def pop_committed_kv_cache(self):
return self._kv_committed_len return self._kv_committed_len
@@ -575,7 +577,7 @@ class TestSWA(unittest.TestCase):
req.extra_key = None req.extra_key = None
req.last_node = tree.root_node req.last_node = tree.root_node
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
req.swa_evicted_seqlen = 0 req.kv.swa_evicted_seqlen = 0
req.cache_protected_len = 1 req.cache_protected_len = 1
# Intentionally mismatch to ensure code does not use len(prefix_indices). # Intentionally mismatch to ensure code does not use len(prefix_indices).
req.prefix_indices = torch.tensor([7, 8, 9, 10, 11], device=tree.device) 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.extra_key = None
req2.last_node = tree.root_node req2.last_node = tree.root_node
req2.swa_uuid_for_lock = None req2.swa_uuid_for_lock = None
req2.swa_evicted_seqlen = 0 req2.kv.swa_evicted_seqlen = 0
req2.cache_protected_len = 1 req2.cache_protected_len = 1
req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device) 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) kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv_committed_len = kv_len 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.last_node = cache.root_node
req.cache_protected_len = 0 req.cache_protected_len = 0
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
@@ -998,7 +998,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0 req.cache_protected_len = 0
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
req.extra_key = None req.extra_key = None
req.swa_evicted_seqlen = evicted_len req.kv.swa_evicted_seqlen = evicted_len
cache.cache_unfinished_req(req) cache.cache_unfinished_req(req)
@@ -1220,7 +1220,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0 req.cache_protected_len = 0
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
req.extra_key = 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() full_available_before_insert = allocator.full_attn_allocator.available_size()
@@ -1761,7 +1761,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0 req.cache_protected_len = 0
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
req.extra_key = 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() 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) cushion = max(self.cfg.sliding_window_size, self.cfg.page_size)
expected_evicted = (pre_len - 1) - cushion expected_evicted = (pre_len - 1) - cushion
self.assertEqual( self.assertEqual(
req.swa_evicted_seqlen, req.kv.swa_evicted_seqlen,
expected_evicted, expected_evicted,
f"swa_evicted_seqlen should advance to (pre_len-1) - cushion = " 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() swa_avail_after = allocator.swa_attn_allocator.available_size()
@@ -1813,13 +1813,13 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0 req.cache_protected_len = 0
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
req.extra_key = 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): with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
cache.cache_unfinished_req(req) cache.cache_unfinished_req(req)
self.assertEqual( self.assertEqual(
req.swa_evicted_seqlen, req.kv.swa_evicted_seqlen,
0, 0,
"Nothing should be evicted when prefill fits inside the cushion", "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, ): 1,
# streaming session slot save/restore and tail trimming # streaming session slot save/restore and tail trimming
(_SS, "SessionSlot.save_from_req", "kv_committed_len"): 1, (_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_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_committed_len"): 2,
(_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2, (_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2,
(_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1, (_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1,