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

[mem_cache] Move req_pool_idx into ReqKvInfo (#37094)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-30 14:46:21 -07:00
committed by GitHub
parent 8a87079dbb
commit 007ef5e23a
68 changed files with 372 additions and 357 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/beam_search/coordinator.py
+3 -3
View File
@@ -308,7 +308,7 @@ class BeamCoordinator(msgspec.Struct, kw_only=True):
if final: if final:
# Prefill-terminated (max_new_tokens == 1): no spawn, but the relay # Prefill-terminated (max_new_tokens == 1): no spawn, but the relay
# slot still needs a token so an overshoot step has a valid input. # slot still needs a token so an overshoot step has a valid input.
self._stash_next_tokens([req.req_pool_idx], next_tokens[:1]) self._stash_next_tokens([req.kv.req_pool_idx], next_tokens[:1])
return return
self._spawn_member_rows(group, req) self._spawn_member_rows(group, req)
req.output_ids.append(0) # length placeholder; DAG owns history req.output_ids.append(0) # length placeholder; DAG owns history
@@ -339,13 +339,13 @@ class BeamCoordinator(msgspec.Struct, kw_only=True):
alias_members_prompt_kv( alias_members_prompt_kv(
self.req_to_token_pool.req_to_token, self.req_to_token_pool.req_to_token,
member_rows, member_rows,
leader.req_pool_idx, leader.kv.req_pool_idx,
group.prompt_len, group.prompt_len,
) )
group.member_rows = member_rows group.member_rows = member_rows
group.member_rows_cpu = torch.tensor(rows, dtype=torch.int64) group.member_rows_cpu = torch.tensor(rows, dtype=torch.int64)
leader_row = torch.tensor( leader_row = torch.tensor(
[leader.req_pool_idx], dtype=torch.int64, device=device [leader.kv.req_pool_idx], dtype=torch.int64, device=device
) )
group.all_rows = torch.cat([leader_row, member_rows]) group.all_rows = torch.cat([leader_row, member_rows])
python/sglang/srt/disaggregation/common/staging_handler.py
+1 -1
View File
@@ -411,7 +411,7 @@ class DecodeStagingHandler:
staging_view = self.staging_allocator.buffer.buffer[staging_offset:] staging_view = self.staging_allocator.buffer.buffer[staging_offset:]
req_pool_idx = decode_req.req.req_pool_idx req_pool_idx = decode_req.req.kv.req_pool_idx
# page_start is suffix-relative (pages after the decode-side cached # page_start is suffix-relative (pages after the decode-side cached
# prefix); req_to_token rows are absolute. # prefix); req_to_token rows are absolute.
prefix_tokens = decode_req.req.kv.cache_protected_len prefix_tokens = decode_req.req.kv.cache_protected_len
python/sglang/srt/disaggregation/decode.py
+20 -20
View File
@@ -188,7 +188,7 @@ class DecodeReqToTokenPool:
def alloc(self, reqs: List[Req]) -> Optional[List[int]]: def alloc(self, reqs: List[Req]) -> Optional[List[int]]:
# Indices of reqs that already have a req_pool_idx and will reuse # Indices of reqs that already have a req_pool_idx and will reuse
# their existing slot (e.g. chunked prefill continuing across chunks). # their existing slot (e.g. chunked prefill continuing across chunks).
reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None] reusing = [i for i, r in enumerate(reqs) if r.kv.req_pool_idx is not None]
assert ( assert (
len(reusing) <= 1 len(reusing) <= 1
), "only one chunked request may reuse req_pool_idx in a batch" ), "only one chunked request may reuse req_pool_idx in a batch"
@@ -204,16 +204,16 @@ class DecodeReqToTokenPool:
self.free_slots = self.free_slots[need_size:] self.free_slots = self.free_slots[need_size:]
offset = 0 offset = 0
for r in reqs: for r in reqs:
if r.req_pool_idx is None: if r.kv.req_pool_idx is None:
r.req_pool_idx = select_index[offset] r.kv.req_pool_idx = select_index[offset]
self.req_generation[r.req_pool_idx] += 1 self.req_generation[r.kv.req_pool_idx] += 1
offset += 1 offset += 1
return [r.req_pool_idx for r in reqs] return [r.kv.req_pool_idx for r in reqs]
def free(self, req: Req): def free(self, req: Req):
assert req.req_pool_idx is not None, "request must have req_pool_idx" assert req.kv.req_pool_idx is not None, "request must have req_pool_idx"
self.free_slots.append(req.req_pool_idx) self.free_slots.append(req.kv.req_pool_idx)
req.req_pool_idx = None req.kv.req_pool_idx = None
def clear(self): def clear(self):
self.free_slots = list(range(1, self._alloc_size)) self.free_slots = list(range(1, self._alloc_size))
@@ -1376,7 +1376,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
else: else:
# Only send delta indices (beyond prefix) to prefill. # Only send delta indices (beyond prefix) to prefill.
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx decode_req.req.kv.req_pool_idx
][total_prefix_len:origin_input_len] ][total_prefix_len:origin_input_len]
kv_indices = ( kv_indices = (
self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer( self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer(
@@ -1390,7 +1390,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
return [ return [
self.req_to_token_pool.translate_mamba_indices( self.req_to_token_pool.translate_mamba_indices(
self.req_to_token_pool.req_index_to_mamba_index_mapping[ self.req_to_token_pool.req_index_to_mamba_index_mapping[
decode_req.req.req_pool_idx decode_req.req.kv.req_pool_idx
] ]
) )
.cpu() .cpu()
@@ -1402,7 +1402,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
window_start = max(total_prefix_len, seq_len - window_size) window_start = max(total_prefix_len, seq_len - window_size)
window_start = page_align_floor(window_start, page_size) window_start = page_align_floor(window_start, page_size)
window_kv_indices_full = self.req_to_token_pool.req_to_token[ window_kv_indices_full = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx, window_start:seq_len decode_req.req.kv.req_pool_idx, window_start:seq_len
] ]
window_kv_indices_swa = ( window_kv_indices_swa = (
self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
@@ -1413,7 +1413,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
def _full_kv_pages_payload(): def _full_kv_pages_payload():
kv_indices_full = self.req_to_token_pool.req_to_token[ kv_indices_full = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx, :seq_len decode_req.req.kv.req_pool_idx, :seq_len
] ]
# Indexer lives on device pool; always use device page_size # Indexer lives on device pool; always use device page_size
device_page_size = self.token_to_kv_pool.page_size device_page_size = self.token_to_kv_pool.page_size
@@ -1426,7 +1426,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
window_size = self.token_to_kv_pool.unified_swa_window window_size = self.token_to_kv_pool.unified_swa_window
window_start = max(0, seq_len - window_size) window_start = max(0, seq_len - window_size)
positions = np.arange(window_start, seq_len, dtype=np.int64) positions = np.arange(window_start, seq_len, dtype=np.int64)
state_slot = int(decode_req.req.req_pool_idx) state_slot = int(decode_req.req.kv.req_pool_idx)
ring_rows = state_slot * ring_stride + (positions % ring_stride) ring_rows = state_slot * ring_stride + (positions % ring_stride)
return ring_rows.astype(np.int32) return ring_rows.astype(np.int32)
@@ -1434,7 +1434,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
online = is_dsv4_c128_online_enabled() online = is_dsv4_c128_online_enabled()
ring_size = 1 if online else self.token_to_kv_pool.get_ring_size(128) ring_size = 1 if online else self.token_to_kv_pool.get_ring_size(128)
return get_dsv4_c128_state_indices( return get_dsv4_c128_state_indices(
int(decode_req.req.req_pool_idx), int(decode_req.req.kv.req_pool_idx),
seq_len, seq_len,
online=online, online=online,
ring_size=ring_size, ring_size=ring_size,
@@ -1446,7 +1446,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
self.token_to_kv_pool, "clear_c128_req_state", None self.token_to_kv_pool, "clear_c128_req_state", None
) )
if clear_c128_state is not None: if clear_c128_state is not None:
clear_c128_state(int(decode_req.req.req_pool_idx)) clear_c128_state(int(decode_req.req.kv.req_pool_idx))
payloads = { payloads = {
StateType.MAMBA: _mamba_payload, StateType.MAMBA: _mamba_payload,
StateType.SWA: _swa_payload, StateType.SWA: _swa_payload,
@@ -1465,7 +1465,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
payloads.update( payloads.update(
dsv4_state_payloads( dsv4_state_payloads(
self.req_to_token_pool, self.req_to_token_pool,
decode_req.req.req_pool_idx, decode_req.req.kv.req_pool_idx,
seq_len, seq_len,
self.token_to_kv_pool_allocator.page_size, self.token_to_kv_pool_allocator.page_size,
prefix_len=total_prefix_len, prefix_len=total_prefix_len,
@@ -1494,7 +1494,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
# the C4 sparse physical-slot mapping; carry their logical page IDs # the C4 sparse physical-slot mapping; carry their logical page IDs
# alongside the independently allocated C4 host page IDs. # alongside the independently allocated C4 host page IDs.
full_kv_indices = self.req_to_token_pool.req_to_token[ full_kv_indices = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx, decode_req.req.kv.req_pool_idx,
prefix_len:origin_input_len, prefix_len:origin_input_len,
] ]
device_page_indices = kv_to_page_indices( device_page_indices = kv_to_page_indices(
@@ -1787,7 +1787,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
if prefix_len > 0: if prefix_len > 0:
self.req_to_token_pool.write( self.req_to_token_pool.write(
(req.req_pool_idx, slice(0, prefix_len)), prefix_indices (req.kv.req_pool_idx, slice(0, prefix_len)), prefix_indices
) )
# TODO(retraction): when retraction is implemented with radix cache # TODO(retraction): when retraction is implemented with radix cache
@@ -1842,7 +1842,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
host_indices = coordinator.mem_pool_host.alloc_paged_token_slots( host_indices = coordinator.mem_pool_host.alloc_paged_token_slots(
coordinator.req_to_host_pool, coordinator.req_to_host_pool,
coordinator.req_to_host_pool_allocated_len, coordinator.req_to_host_pool_allocated_len,
req.req_pool_idx, req.kv.req_pool_idx,
0, 0,
coordinator.host_token_len(fill_len), coordinator.host_token_len(fill_len),
) )
@@ -1872,7 +1872,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
self.req_to_token_pool.write( self.req_to_token_pool.write(
( (
req.req_pool_idx, req.kv.req_pool_idx,
slice(total_prefix_len, total_prefix_len + len(kv_loc)), slice(total_prefix_len, total_prefix_len + len(kv_loc)),
), ),
kv_loc, kv_loc,
python/sglang/srt/disaggregation/decode_hicache_mixin.py
+1 -1
View File
@@ -306,7 +306,7 @@ class DecodeHiCacheTransferMixin:
self.tree_cache.req_to_token_pool.write( self.tree_cache.req_to_token_pool.write(
( (
decode_req.req.req_pool_idx, decode_req.req.kv.req_pool_idx,
slice(prefix_match.l1_prefix_len, prefix_match.decode_prefix_len), slice(prefix_match.l1_prefix_len, prefix_match.decode_prefix_len),
), ),
decode_req.hicache_restored_kv_indices, decode_req.hicache_restored_kv_indices,
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+7 -7
View File
@@ -123,10 +123,10 @@ class DecodeKVCacheOffloadManager:
if self.cache_controller is None or self.decode_host_mem_pool is None: if self.cache_controller is None or self.decode_host_mem_pool is None:
return False return False
if req.req_pool_idx == -1 or len(req.output_ids) == 0: if req.kv.req_pool_idx == -1 or len(req.output_ids) == 0:
return False return False
token_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx] token_indices = self.req_to_token_pool.req_to_token[req.kv.req_pool_idx]
if token_indices.dim() == 0 or token_indices.numel() == 0: if token_indices.dim() == 0 or token_indices.numel() == 0:
return False return False
@@ -251,7 +251,7 @@ class DecodeKVCacheOffloadManager:
# so a previously-released request must be skipped here to avoid # so a previously-released request must be skipped here to avoid
# non-idempotent side effects (e.g. tree_cache.protected_size_ # non-idempotent side effects (e.g. tree_cache.protected_size_
# double-decrement, host pool double-free). # double-decrement, host pool double-free).
if req.req_pool_idx is None or req.req_pool_idx == -1: if req.kv.req_pool_idx is None or req.kv.req_pool_idx == -1:
return return
kv_committed_len = req.effective_kv_committed_len() kv_committed_len = req.effective_kv_committed_len()
@@ -264,13 +264,13 @@ class DecodeKVCacheOffloadManager:
state = self.offloaded_state.get(req) state = self.offloaded_state.get(req)
if state is not None and state.prefill_len > 0: if state is not None and state.prefill_len > 0:
prefill_indices = self.req_to_token_pool.req_to_token[ prefill_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : state.prefill_len req.kv.req_pool_idx, : state.prefill_len
] ]
self.token_to_kv_pool_allocator.free(prefill_indices) self.token_to_kv_pool_allocator.free(prefill_indices)
start = start_offset start = start_offset
end = kv_committed_len end = kv_committed_len
# Free the incremental part of the request (DSA-aware) # Free the incremental part of the request (DSA-aware)
kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, start:end] kv_indices = self.req_to_token_pool.req_to_token[req.kv.req_pool_idx, start:end]
self.token_to_kv_pool_allocator.free(kv_indices) self.token_to_kv_pool_allocator.free(kv_indices)
# Free over-allocated KV cache slots (e.g. from speculative decoding v2). # Free over-allocated KV cache slots (e.g. from speculative decoding v2).
@@ -280,7 +280,7 @@ class DecodeKVCacheOffloadManager:
start_p = ceil_align(start_p, self.page_size) start_p = ceil_align(start_p, self.page_size)
if start_p < end_p: if start_p < end_p:
overalloc_indices = self.req_to_token_pool.req_to_token[ overalloc_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, start_p:end_p req.kv.req_pool_idx, start_p:end_p
] ]
self.token_to_kv_pool_allocator.free(overalloc_indices) self.token_to_kv_pool_allocator.free(overalloc_indices)
@@ -329,7 +329,7 @@ class DecodeKVCacheOffloadManager:
"""Free any remaining tail KV that was not offloaded due to non-aligned length.""" """Free any remaining tail KV that was not offloaded due to non-aligned length."""
# ReqToTokenPool.free sets req_pool_idx to None on release, so # ReqToTokenPool.free sets req_pool_idx to None on release, so
# guard against both sentinels here. # guard against both sentinels here.
if req.req_pool_idx is None or req.req_pool_idx == -1: if req.kv.req_pool_idx is None or req.kv.req_pool_idx == -1:
return return
state = self.offloaded_state.get(req) state = self.offloaded_state.get(req)
if state is None: if state is None:
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
+2 -2
View File
@@ -43,10 +43,10 @@ class ScheduleBatchDisaggregationDecodeMixin:
# Fill the tensor in one pass # Fill the tensor in one pass
offset = 0 offset = 0
for i, req in enumerate(reqs): for i, req in enumerate(reqs):
req_pool_indices.append(req.req_pool_idx) req_pool_indices.append(req.kv.req_pool_idx)
pre_len = len(req.prefix_indices) pre_len = len(req.prefix_indices)
chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][ chunk = self.req_to_token_pool.req_to_token[req.kv.req_pool_idx][
pre_len : pre_len + req.extend_range.length pre_len : pre_len + req.extend_range.length
] ]
assert ( assert (
python/sglang/srt/disaggregation/prefill.py
+8 -8
View File
@@ -1036,7 +1036,7 @@ class SchedulerDisaggregationPrefillMixin:
else: else:
logger.warning(error_message) logger.warning(error_message)
req.time_stats.trace_ctx.abort(abort_info={"reason": error_message}) req.time_stats.trace_ctx.abort(abort_info={"reason": error_message})
if req.is_holding_kv or req.mamba_pool_idx is not None: if req.kv.is_held or req.mamba_pool_idx is not None:
release_kv_cache(req, self.tree_cache) release_kv_cache(req, self.tree_cache)
maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator) maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator)
req.pending_bootstrap = False req.pending_bootstrap = False
@@ -1215,7 +1215,7 @@ class SchedulerDisaggregationPrefillMixin:
return [ return [
self.req_to_token_pool.translate_mamba_indices( self.req_to_token_pool.translate_mamba_indices(
self.req_to_token_pool.req_index_to_mamba_index_mapping[ self.req_to_token_pool.req_index_to_mamba_index_mapping[
req.req_pool_idx req.kv.req_pool_idx
] ]
) )
.cpu() .cpu()
@@ -1227,7 +1227,7 @@ class SchedulerDisaggregationPrefillMixin:
window_start = max(req.disagg_decode_prefix_len, seq_len - window_size) window_start = max(req.disagg_decode_prefix_len, seq_len - window_size)
window_start = (window_start // page_size) * page_size window_start = (window_start // page_size) * page_size
window_kv_indices_full = self.req_to_token_pool.req_to_token[ window_kv_indices_full = self.req_to_token_pool.req_to_token[
req.req_pool_idx, window_start:seq_len req.kv.req_pool_idx, window_start:seq_len
] ]
window_kv_indices_swa = ( window_kv_indices_swa = (
self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
@@ -1238,7 +1238,7 @@ class SchedulerDisaggregationPrefillMixin:
def _full_kv_pages_payload(): def _full_kv_pages_payload():
kv_indices_full = self.req_to_token_pool.req_to_token[ kv_indices_full = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :seq_len req.kv.req_pool_idx, :seq_len
] ]
return kv_to_page_indices(kv_indices_full, page_size) return kv_to_page_indices(kv_indices_full, page_size)
@@ -1251,7 +1251,7 @@ class SchedulerDisaggregationPrefillMixin:
window_size = _pool.unified_swa_window window_size = _pool.unified_swa_window
window_start = max(0, seq_len - window_size) window_start = max(0, seq_len - window_size)
positions = np.arange(window_start, seq_len, dtype=np.int64) positions = np.arange(window_start, seq_len, dtype=np.int64)
state_slot = int(req.req_pool_idx) state_slot = int(req.kv.req_pool_idx)
ring_rows = state_slot * ring_stride + (positions % ring_stride) ring_rows = state_slot * ring_stride + (positions % ring_stride)
return ring_rows.astype(np.int32) return ring_rows.astype(np.int32)
@@ -1265,7 +1265,7 @@ class SchedulerDisaggregationPrefillMixin:
) )
) )
return get_dsv4_c128_state_indices( return get_dsv4_c128_state_indices(
int(req.req_pool_idx), int(req.kv.req_pool_idx),
c128_seq_len, c128_seq_len,
online=online, online=online,
ring_size=ring_size, ring_size=ring_size,
@@ -1295,7 +1295,7 @@ class SchedulerDisaggregationPrefillMixin:
payloads.update( payloads.update(
dsv4_state_payloads( dsv4_state_payloads(
self.req_to_token_pool, self.req_to_token_pool,
req.req_pool_idx, req.kv.req_pool_idx,
seq_len, seq_len,
page_size, page_size,
prefix_len=req.disagg_decode_prefix_len, prefix_len=req.disagg_decode_prefix_len,
@@ -1321,7 +1321,7 @@ class SchedulerDisaggregationPrefillMixin:
for seg_start, seg_end in segments: for seg_start, seg_end in segments:
is_final_segment = seg_end == end_idx is_final_segment = seg_end == end_idx
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, seg_start:seg_end req.kv.req_pool_idx, seg_start:seg_end
] ]
# Unified memory: req_to_token holds VIRTUAL ids; the transfer needs # Unified memory: req_to_token holds VIRTUAL ids; the transfer needs
# physical ones. Per segment, since each is its own gather. # physical ones. Per segment, since each is its own gather.
python/sglang/srt/hardware_backend/mlx/tp_worker.py
+1 -1
View File
@@ -502,7 +502,7 @@ class MlxTpModelWorker(TpModelWorker):
full_token_ids=full_token_ids, full_token_ids=full_token_ids,
prefix_slot_ids=prefix_slot_ids, prefix_slot_ids=prefix_slot_ids,
new_slot_ids=req_new_slots, new_slot_ids=req_new_slots,
req_pool_idx=req.req_pool_idx, req_pool_idx=req.kv.req_pool_idx,
req=req, req=req,
needs_logits=self._chunk_needs_logits(req), needs_logits=self._chunk_needs_logits(req),
logit_edit_row=edit_rows[req.rid] if edit_rows else None, logit_edit_row=edit_rows[req.rid] if edit_rows else None,
python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py
+1 -1
View File
@@ -259,7 +259,7 @@ class C128SidecarComponent(TreeComponent):
cache_len = logical_len // group_tokens * group_tokens cache_len = logical_len // group_tokens * group_tokens
num_pages = cache_len // group_tokens num_pages = cache_len // group_tokens
insert_params.c128_value = self.cache.req_to_token_pool.req_to_c128_sidecar[ insert_params.c128_value = self.cache.req_to_token_pool.req_to_c128_sidecar[
int(req.req_pool_idx), :num_pages int(req.kv.req_pool_idx), :num_pages
].clone() ].clone()
return cache_len + 1 if self.tree_core.is_eagle and cache_len > 0 else cache_len return cache_len + 1 if self.tree_core.is_eagle and cache_len > 0 else cache_len
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py
+1 -1
View File
@@ -466,7 +466,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.kv_committed_len, req.kv.kv_allocated_len) kv_len = max(req.kv.kv_committed_len, req.kv.kv_allocated_len)
req_pool_idx = req.req_pool_idx req_pool_idx = req.kv.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/hardware_backend/npu/dsv4/dsv4_common_hooks.py
+2 -2
View File
@@ -104,7 +104,7 @@ def dsv4_prealloc_kwargs(allocator, req, fill_len, req_to_token_pool, *, device)
return {} return {}
return dict( return dict(
req_pool_indices=torch.tensor( req_pool_indices=torch.tensor(
[req.req_pool_idx], dtype=torch.int64, device=device [req.kv.req_pool_idx], dtype=torch.int64, device=device
), ),
req_to_token_pool=req_to_token_pool, req_to_token_pool=req_to_token_pool,
) )
@@ -130,7 +130,7 @@ def write_dsv4_prealloc_tables(
prealloc path (no ScheduleBatch); no-op without bundle / DSV4 tables.""" prealloc path (no ScheduleBatch); no-op without bundle / DSV4 tables."""
if bundle is None or not hasattr(req_to_token_pool, "write_c128"): if bundle is None or not hasattr(req_to_token_pool, "write_c128"):
return return
rp = torch.tensor([req.req_pool_idx]) rp = torch.tensor([req.kv.req_pool_idx])
pl = torch.tensor([prefix_len]) pl = torch.tensor([prefix_len])
sl = torch.tensor([fill_len]) sl = torch.tensor([fill_len])
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py
+3 -3
View File
@@ -83,15 +83,15 @@ class DSV4ReqToTokenTablesMixin:
Prefix matching can happen before a request slot is allocated, so the Prefix matching can happen before a request slot is allocated, so the
page ids are temporarily carried by ``Req`` and installed by ``alloc``. page ids are temporarily carried by ``Req`` and installed by ``alloc``.
""" """
if req.req_pool_idx is None: if req.kv.req_pool_idx is None:
req.c128_prefix_page_ids = page_ids req.c128_prefix_page_ids = page_ids
return return
self._dsv4_allocator.replace_req_c128_prefix( self._dsv4_allocator.replace_req_c128_prefix(
int(req.req_pool_idx), page_ids, self int(req.kv.req_pool_idx), page_ids, self
) )
def alloc(self, reqs): def alloc(self, reqs):
fresh = [req.req_pool_idx is None for req in reqs] fresh = [req.kv.req_pool_idx is None for req in reqs]
indices = super().alloc(reqs) indices = super().alloc(reqs)
if indices is None: if indices is None:
return None return None
python/sglang/srt/managers/hisparse_coordinator.py
+31 -31
View File
@@ -349,7 +349,7 @@ class HiSparseCoordinator:
req.hisparse_staging = True req.hisparse_staging = True
full_kv_indices = self.req_to_token_pool.req_to_token[ full_kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.extend_range.end req.kv.req_pool_idx, : req.extend_range.end
].to(dtype=torch.int64, copy=True) ].to(dtype=torch.int64, copy=True)
device_indices = ( device_indices = (
self.mem_pool_device.translate_loc_from_full_to_hisparse_device( self.mem_pool_device.translate_loc_from_full_to_hisparse_device(
@@ -361,7 +361,7 @@ class HiSparseCoordinator:
host_indices = self.mem_pool_host.alloc_paged_token_slots( host_indices = self.mem_pool_host.alloc_paged_token_slots(
self.req_to_host_pool, self.req_to_host_pool,
self.req_to_host_pool_allocated_len, self.req_to_host_pool_allocated_len,
req.req_pool_idx, req.kv.req_pool_idx,
0, 0,
prefill_len, prefill_len,
) )
@@ -411,11 +411,11 @@ class HiSparseCoordinator:
# Long sequence: reset device_buffer_tokens to -1 so the kernel # Long sequence: reset device_buffer_tokens to -1 so the kernel
# sees all slots as empty -> every top-k lookup is a miss -> host load. # sees all slots as empty -> every top-k lookup is a miss -> host load.
self.req_device_buffer_tokens[ self.req_device_buffer_tokens[
:, req.req_pool_idx, : self.device_buffer_size :, req.kv.req_pool_idx, : self.device_buffer_size
] = -1 ] = -1
req.hisparse_staging = False req.hisparse_staging = False
self._skip_first_backup[req.req_pool_idx] = True self._skip_first_backup[req.kv.req_pool_idx] = True
logger.debug("HiSparse: admitting request %s directly", req.rid) logger.debug("HiSparse: admitting request %s directly", req.rid)
def host_token_len(self, kv_allocated_len: int) -> int: def host_token_len(self, kv_allocated_len: int) -> int:
@@ -426,8 +426,8 @@ 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.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.kv.req_pool_idx, :n]
device_locs = self.req_to_device_buffer[req.req_pool_idx, :n] device_locs = self.req_to_device_buffer[req.kv.req_pool_idx, :n]
for layer_id in range(self.mem_pool_device.layer_num): for layer_id in range(self.mem_pool_device.layer_num):
self.mem_pool_host.load_to_device_per_layer( self.mem_pool_host.load_to_device_per_layer(
@@ -456,7 +456,7 @@ class HiSparseCoordinator:
compressed_logical_indices = ( compressed_logical_indices = (
self.mem_pool_device.translate_loc_from_full_to_compressed( self.mem_pool_device.translate_loc_from_full_to_compressed(
self.req_to_token_pool.req_to_token[req.req_pool_idx, :allocated_len] self.req_to_token_pool.req_to_token[req.kv.req_pool_idx, :allocated_len]
) )
) )
compressed_len = len(compressed_logical_indices) compressed_len = len(compressed_logical_indices)
@@ -475,13 +475,13 @@ class HiSparseCoordinator:
raise RuntimeError("HiSparse alloc_device_buffer returned None") raise RuntimeError("HiSparse alloc_device_buffer returned None")
buffer_indices = buffer_indices.to(torch.int32) buffer_indices = buffer_indices.to(torch.int32)
self.req_to_device_buffer[req.req_pool_idx, :alloc_size] = buffer_indices self.req_to_device_buffer[req.kv.req_pool_idx, :alloc_size] = buffer_indices
self.req_device_buffer_size[req.req_pool_idx] = alloc_size self.req_device_buffer_size[req.kv.req_pool_idx] = alloc_size
self.req_device_buffer_tokens[ self.req_device_buffer_tokens[
:, req.req_pool_idx, : self.device_buffer_size :, req.kv.req_pool_idx, : self.device_buffer_size
] = self._device_buffer_arange_i32 ] = self._device_buffer_arange_i32
self.req_device_buffer_token_locs[:, req.req_pool_idx, :alloc_size] = ( self.req_device_buffer_token_locs[:, req.kv.req_pool_idx, :alloc_size] = (
buffer_indices[:alloc_size] buffer_indices[:alloc_size]
) )
@@ -584,7 +584,7 @@ class HiSparseCoordinator:
_, _, req = self.ack_staging_queue.pop(0) _, _, req = self.ack_staging_queue.pop(0)
# prepare device buffer and update req # prepare device buffer and update req
self.alloc_device_buffer(req) self.alloc_device_buffer(req)
self._skip_first_backup[req.req_pool_idx] = True self._skip_first_backup[req.kv.req_pool_idx] = True
req.hisparse_staging = False req.hisparse_staging = False
finish_count -= 1 finish_count -= 1
ready_reqs.append(req) ready_reqs.append(req)
@@ -863,21 +863,21 @@ class HiSparseCoordinator:
prefill_len = req.extend_range.end prefill_len = req.extend_range.end
allocated_locs = self.req_to_token_pool.req_to_token[ allocated_locs = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :prefill_len req.kv.req_pool_idx, :prefill_len
] ]
self.token_to_kv_pool_allocator.free_hisparse(allocated_locs) self.token_to_kv_pool_allocator.free_hisparse(allocated_locs)
# Free host memory that was allocated during admit_request_into_staging # Free host memory that was allocated during admit_request_into_staging
host_indices = self.mem_pool_host.allocated_host_indices( host_indices = self.mem_pool_host.allocated_host_indices(
self.req_to_host_pool, self.req_to_host_pool,
req.req_pool_idx, req.kv.req_pool_idx,
self.req_to_host_pool_allocated_len[req.req_pool_idx], self.req_to_host_pool_allocated_len[req.kv.req_pool_idx],
) )
if host_indices.numel() > 0: if host_indices.numel() > 0:
self.mem_pool_host.free(host_indices) self.mem_pool_host.free(host_indices)
self.req_to_host_pool[req.req_pool_idx, :] = -1 self.req_to_host_pool[req.kv.req_pool_idx, :] = -1
self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0 self.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = 0
self._skip_first_backup[req.req_pool_idx] = False self._skip_first_backup[req.kv.req_pool_idx] = False
req.hisparse_staging = False req.hisparse_staging = False
def retract_req(self, req: Req) -> None: def retract_req(self, req: Req) -> None:
@@ -901,15 +901,15 @@ class HiSparseCoordinator:
allocated_len = req.kv.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.kv.req_pool_idx])
if current_cap > 0: if current_cap > 0:
side_buf_hi = self.req_to_device_buffer[req.req_pool_idx, :current_cap] side_buf_hi = self.req_to_device_buffer[req.kv.req_pool_idx, :current_cap]
all_hi = torch.unique(side_buf_hi[side_buf_hi > 0]) all_hi = torch.unique(side_buf_hi[side_buf_hi > 0])
if all_hi.numel() > 0: if all_hi.numel() > 0:
self.token_to_kv_pool_allocator.free_hisparse_indices(all_hi) self.token_to_kv_pool_allocator.free_hisparse_indices(all_hi)
allocated_locs = self.req_to_token_pool.req_to_token[ allocated_locs = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :allocated_len req.kv.req_pool_idx, :allocated_len
] ]
compressed_locs = self.mem_pool_device.translate_loc_from_full_to_compressed( compressed_locs = self.mem_pool_device.translate_loc_from_full_to_compressed(
allocated_locs allocated_locs
@@ -918,21 +918,21 @@ class HiSparseCoordinator:
host_indices = self.mem_pool_host.allocated_host_indices( host_indices = self.mem_pool_host.allocated_host_indices(
self.req_to_host_pool, self.req_to_host_pool,
req.req_pool_idx, req.kv.req_pool_idx,
self.req_to_host_pool_allocated_len[req.req_pool_idx], self.req_to_host_pool_allocated_len[req.kv.req_pool_idx],
) )
if host_indices.numel() > 0: if host_indices.numel() > 0:
self.mem_pool_host.free(host_indices) self.mem_pool_host.free(host_indices)
# clear req info # clear req info
self.req_device_buffer_tokens[:, req.req_pool_idx, :] = -1 self.req_device_buffer_tokens[:, req.kv.req_pool_idx, :] = -1
self.req_device_buffer_token_locs[:, req.req_pool_idx, :] = -1 self.req_device_buffer_token_locs[:, req.kv.req_pool_idx, :] = -1
self.req_to_device_buffer[req.req_pool_idx, :] = 0 self.req_to_device_buffer[req.kv.req_pool_idx, :] = 0
self.req_device_buffer_size[req.req_pool_idx] = 0 self.req_device_buffer_size[req.kv.req_pool_idx] = 0
self.req_to_host_pool[req.req_pool_idx, :] = -1 self.req_to_host_pool[req.kv.req_pool_idx, :] = -1
self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0 self.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = 0
self.lru_slots[:, req.req_pool_idx, :].copy_(self._lru_init) self.lru_slots[:, req.kv.req_pool_idx, :].copy_(self._lru_init)
self._skip_first_backup[req.req_pool_idx] = False self._skip_first_backup[req.kv.req_pool_idx] = False
def _run_swap_in_kernel( def _run_swap_in_kernel(
self, self,
python/sglang/srt/managers/schedule_batch.py
+10 -10
View File
@@ -816,7 +816,8 @@ class ReqLogprob:
@dataclasses.dataclass(slots=True, kw_only=True) @dataclasses.dataclass(slots=True, kw_only=True)
class ReqKvInfo: class ReqKvInfo:
# Device KV a request holds outside the prefix cache. Always present on the Req; # Device KV a request holds outside the prefix cache. Always present on the Req;
# whether any KV is held is `req.req_pool_idx is not None` (Req.is_holding_kv). # whether any KV is held is `is_held` (a row is registered).
req_pool_idx: Optional[int] = None # req_to_token row, the register for the slots
# The request's own KV is [cache_protected_len, kv_allocated_len). # The request's own KV is [cache_protected_len, kv_allocated_len).
cache_protected_len: int = 0 # tree cache owns [0, here) (matched or inserted) cache_protected_len: int = 0 # tree cache owns [0, here) (matched or inserted)
@@ -835,6 +836,10 @@ class ReqKvInfo:
lo = ceil_align(lo, page_size) lo = ceil_align(lo, page_size)
return lo return lo
@property
def is_held(self) -> bool:
return self.req_pool_idx is not None
@property @property
def is_released(self) -> bool: def is_released(self) -> bool:
return self.kv_allocated_len == 0 and self.swa_evicted_seqlen == 0 return self.kv_allocated_len == 0 and self.swa_evicted_seqlen == 0
@@ -970,7 +975,6 @@ class Req(ReqDllmMixin):
self.routing_key = routing_key self.routing_key = routing_key
# Memory pool info # Memory pool info
self.req_pool_idx: Optional[int] = None
self.mamba_pool_idx: Optional[torch.Tensor] = None # shape (1) self.mamba_pool_idx: Optional[torch.Tensor] = None # shape (1)
self.mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None # shape (2) self.mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None # shape (2)
self.mamba_next_track_idx: Optional[int] = None # 0 or 1 self.mamba_next_track_idx: Optional[int] = None # 0 or 1
@@ -1278,10 +1282,6 @@ class Req(ReqDllmMixin):
or self.mamba_host_hit_length > 0 or self.mamba_host_hit_length > 0
) )
@property
def is_holding_kv(self) -> bool:
return self.req_pool_idx is not None
def effective_kv_committed_len(self) -> int: def effective_kv_committed_len(self) -> int:
# Report only the prompt prefix so thinking + answer fall into the # Report only the prompt prefix so thinking + answer fall into the
# overallocated range and are reclaimed by release_kv_cache. #22373. # overallocated range and are reclaimed by release_kv_cache. #22373.
@@ -1748,7 +1748,7 @@ 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
assert not self.is_holding_kv, "expect it is already released" assert not self.kv.is_held, "expect it is already released"
self.kv.kv_committed_len = 0 self.kv.kv_committed_len = 0
self.extend_batch_idx = 0 self.extend_batch_idx = 0
self.decode_batch_idx = 0 self.decode_batch_idx = 0
@@ -1773,7 +1773,7 @@ class Req(ReqDllmMixin):
def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
token_indices = req_to_token_pool.req_to_token[ token_indices = req_to_token_pool.req_to_token[
self.req_pool_idx, : self.seqlen - 1 self.kv.req_pool_idx, : self.seqlen - 1
] ]
# Copies over both the kv cache and mamba state if available # Copies over both the kv cache and mamba state if available
mamba_pool = self._mamba_pool_needing_backup( mamba_pool = self._mamba_pool_needing_backup(
@@ -1793,7 +1793,7 @@ class Req(ReqDllmMixin):
def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
assert self.retraction_backup is not None assert self.retraction_backup is not None
token_indices = req_to_token_pool.req_to_token[ token_indices = req_to_token_pool.req_to_token[
self.req_pool_idx, : self.seqlen - 1 self.kv.req_pool_idx, : self.seqlen - 1
] ]
# Loads both the kv cache and mamba state if exists # Loads both the kv cache and mamba state if exists
mamba_cpu = self.retraction_backup.mamba_cpu mamba_cpu = self.retraction_backup.mamba_cpu
@@ -3492,7 +3492,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# seqlen progress is monotonic per KV handle. # seqlen progress is monotonic per KV handle.
if ( if (
req.decode_batch_idx >= 1 req.decode_batch_idx >= 1
and req.is_holding_kv and req.kv.is_held
and req.seqlen - 1 - sliding_window_size and req.seqlen - 1 - sliding_window_size
>= req.kv.swa_evicted_seqlen + eviction_interval >= req.kv.swa_evicted_seqlen + eviction_interval
): ):
python/sglang/srt/managers/scheduler.py
+3 -3
View File
@@ -3124,7 +3124,7 @@ class Scheduler(
if self.chunked_req is not req: if self.chunked_req is not req:
# Already past chunked prefill; the running-batch abort path handles # Already past chunked prefill; the running-batch abort path handles
# it. Drop the marker once the request is actually gone. # it. Drop the marker once the request is actually gone.
if req.finished() or req.req_pool_idx is None: if req.finished() or req.kv.req_pool_idx is None:
self._pending_chunked_abort_req = None self._pending_chunked_abort_req = None
return return
@@ -3161,7 +3161,7 @@ class Scheduler(
spec_algorithm=self.spec_algorithm, spec_algorithm=self.spec_algorithm,
) )
req_pool_indices = [r.req_pool_idx for r in reqs] req_pool_indices = [r.kv.req_pool_idx for r in reqs]
batch.req_pool_indices = torch.tensor( batch.req_pool_indices = torch.tensor(
req_pool_indices, dtype=torch.int64, device=device req_pool_indices, dtype=torch.int64, device=device
) )
@@ -4868,7 +4868,7 @@ class Scheduler(
_make_abort_req(req), req _make_abort_req(req), req
) )
if ( if (
req.req_pool_idx is not None req.kv.req_pool_idx is not None
or getattr(req, "mamba_pool_idx", None) is not None or getattr(req, "mamba_pool_idx", None) is not None
): ):
release_kv_cache(req, self.tree_cache, is_insert=False) release_kv_cache(req, self.tree_cache, is_insert=False)
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+2 -2
View File
@@ -136,7 +136,7 @@ class SchedulerBatchResultProcessor:
start_len = req.routed_experts_start_len start_len = req.routed_experts_start_len
seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop) seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop)
req.routed_experts = capturer.get_topk( req.routed_experts = capturer.get_topk(
req_pool_idx=req.req_pool_idx, req_pool_idx=req.kv.req_pool_idx,
seqlen=seqlen, seqlen=seqlen,
req_to_token_pool=self.req_to_token_pool, req_to_token_pool=self.req_to_token_pool,
start_len=start_len, start_len=start_len,
@@ -166,7 +166,7 @@ class SchedulerBatchResultProcessor:
return return
seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop) seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop)
req.indexer_topk = capturer.get_topk( req.indexer_topk = capturer.get_topk(
req_pool_idx=req.req_pool_idx, req_pool_idx=req.kv.req_pool_idx,
seqlen=seqlen, seqlen=seqlen,
req_to_token_pool=self.req_to_token_pool, req_to_token_pool=self.req_to_token_pool,
) )
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+5 -5
View File
@@ -252,7 +252,7 @@ class SchedulerInvariantChecker:
swa_uncached = 0 swa_uncached = 0
for batch in batches: for batch in batches:
for req in batch.reqs: for req in batch.reqs:
if not req.is_holding_kv: if not req.kv.is_held:
continue continue
allocated_len = req.kv.kv_allocated_len allocated_len = req.kv.kv_allocated_len
@@ -324,23 +324,23 @@ class SchedulerInvariantChecker:
batch = self.get_last_batch() batch = self.get_last_batch()
if batch is not None: if batch is not None:
for req in batch.reqs: for req in batch.reqs:
if not req.is_holding_kv: if not req.kv.is_held:
continue continue
_add_owner( _add_owner(
req, req,
f"req {req.rid}", f"req {req.rid}",
req.req_pool_idx, req.kv.req_pool_idx,
req.kv.kv_committed_len, req.kv.kv_committed_len,
req.kv.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:
for sid, slot in sess.items(): for sid, slot in sess.items():
if getattr(slot, "is_holding_kv", False): if slot.kv.is_held:
_add_owner( _add_owner(
slot, slot,
f"slot {sid[:8]}", f"slot {sid[:8]}",
slot.req_pool_idx, slot.kv.req_pool_idx,
slot.kv.kv_committed_len, slot.kv.kv_committed_len,
slot.kv.kv_allocated_len, slot.kv.kv_allocated_len,
) )
python/sglang/srt/managers/scheduler_components/pool_stats_observer.py
+2 -2
View File
@@ -172,8 +172,8 @@ class SchedulerPoolStatsObserver:
if batch is None or batch.is_empty(): if batch is None or batch.is_empty():
continue continue
for req in batch.reqs: for req in batch.reqs:
if req.req_pool_idx is not None: if req.kv.req_pool_idx is not None:
idxs.add(req.req_pool_idx) idxs.add(req.kv.req_pool_idx)
return idxs return idxs
def session_held_tokens(self) -> int: def session_held_tokens(self) -> int:
python/sglang/srt/managers/scheduler_pp_mixin.py
+2 -2
View File
@@ -723,9 +723,9 @@ class SchedulerPPMixin:
latencies.append(latency_ms) latencies.append(latency_ms)
# Release KV and Mamba cache # Release KV and Mamba cache
if req.is_holding_kv: if req.kv.is_held:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.extend_range.end req.kv.req_pool_idx, : req.extend_range.end
] ]
self.token_to_kv_pool_allocator.free(kv_indices) self.token_to_kv_pool_allocator.free(kv_indices)
if req.mamba_pool_idx is not None: if req.mamba_pool_idx is not None:
python/sglang/srt/mem_cache/allocation.py
+1 -1
View File
@@ -297,7 +297,7 @@ def alloc_for_extend(
reuse_kv = None reuse_kv = None
if batch.is_dllm(): if batch.is_dllm():
reuse_kv = [ reuse_kv = [
r.req_pool_idx is not None and bool(r.dllm_incomplete_ids) r.kv.req_pool_idx is not None and bool(r.dllm_incomplete_ids)
for r in batch.reqs for r in batch.reqs
] ]
python/sglang/srt/mem_cache/chunk_cache.py
+3 -3
View File
@@ -82,13 +82,13 @@ class ChunkCache(BasePrefixCache):
# For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids # For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids
# The protected prefix is not this req's to free. # The protected prefix is not this req's to free.
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle req.kv.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle
] ]
self.token_to_kv_pool_allocator.free(kv_indices) self.token_to_kv_pool_allocator.free(kv_indices)
def cache_unfinished_req(self, req: Req, chunked=False): def cache_unfinished_req(self, req: Req, chunked=False):
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.extend_range.end req.kv.req_pool_idx, : req.extend_range.end
] ]
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
@@ -158,7 +158,7 @@ class PureSWAChunkCache(SWAChunkCache):
): ):
kv_committed_len = kv_len_to_handle kv_committed_len = kv_len_to_handle
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len req.kv.req_pool_idx, :kv_committed_len
] ]
# The cache_protected_len prefix is not this req's to free. # The cache_protected_len prefix is not this req's to free.
protected_len = req.kv.cache_protected_len protected_len = req.kv.cache_protected_len
python/sglang/srt/mem_cache/common.py
+9 -9
View File
@@ -62,7 +62,7 @@ def free_swa_out_of_window_slots(
is_chunk_cache: bool = False, is_chunk_cache: bool = False,
retain_floor: int | None = None, retain_floor: int | None = None,
) -> None: ) -> None:
if not req.is_holding_kv: if not req.kv.is_held:
return return
# For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window # For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window
@@ -99,7 +99,7 @@ def free_swa_out_of_window_slots(
if new_swa_evicted_seqlen > req.kv.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.kv.swa_evicted_seqlen : new_swa_evicted_seqlen req.kv.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)
req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen
@@ -199,9 +199,9 @@ def retraction_discard(req: Req, tree_cache: BasePrefixCache, backend: str) -> N
def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True): def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):
assert (not req.is_holding_kv) == req.kv.is_released assert (not req.kv.is_held) == req.kv.is_released
# MambaRadixCache may alloc mamba state before alloc KV cache # MambaRadixCache may alloc mamba state before alloc KV cache
if not req.is_holding_kv: if not req.kv.is_held:
assert ( assert (
tree_cache.supports_mamba() tree_cache.supports_mamba()
), "Only MambaRadixCache allow freeing before alloc" ), "Only MambaRadixCache allow freeing before alloc"
@@ -222,8 +222,8 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
# StreamingSession.cache_finished_req handles speculative tail trim # StreamingSession.cache_finished_req handles speculative tail trim
# internally, then sets req_pool_idx = None. # internally, then sets req_pool_idx = None.
assert (not req.is_holding_kv) == req.kv.is_released assert (not req.kv.is_held) == req.kv.is_released
if not req.is_holding_kv: if not req.kv.is_held:
return return
start_p, end_p = effective_kv_committed_len, req.kv.kv_allocated_len start_p, end_p = effective_kv_committed_len, req.kv.kv_allocated_len
@@ -261,9 +261,9 @@ def _release_overallocated_kv_indices(
start_p = ceil_align(start_p, page_size) start_p = ceil_align(start_p, page_size)
if start_p < end_p: if start_p < end_p:
indices_to_free = tree_cache.req_to_token_pool.req_to_token[req.req_pool_idx][ indices_to_free = tree_cache.req_to_token_pool.req_to_token[
start_p:end_p req.kv.req_pool_idx
] ][start_p:end_p]
# start_p is aligned to the allocator's physical page size above, so it # start_p is aligned to the allocator's physical page size above, so it
# never shares a page with cache_finished_req's tail free in this group. # never shares a page with cache_finished_req's tail free in this group.
allocator.free_segment(indices_to_free, start_pos=start_p) allocator.free_segment(indices_to_free, start_pos=start_p)
python/sglang/srt/mem_cache/mamba_radix_cache.py
+5 -5
View File
@@ -548,7 +548,7 @@ class MambaRadixCache(BasePrefixCache):
"""Cache request when it finishes.""" """Cache request when it finishes."""
if self.disable: if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_len_to_handle req.kv.req_pool_idx, :kv_len_to_handle
] ]
self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0) self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0)
self.req_to_token_pool.free_mamba_cache(req) self.req_to_token_pool.free_mamba_cache(req)
@@ -556,7 +556,7 @@ class MambaRadixCache(BasePrefixCache):
token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_len_to_handle req.kv.req_pool_idx, :kv_len_to_handle
] ]
if is_insert: if is_insert:
@@ -676,7 +676,7 @@ class MambaRadixCache(BasePrefixCache):
def _skip_cache_unfinished_req(req: Req) -> None: def _skip_cache_unfinished_req(req: Req) -> None:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.extend_range.end req.kv.req_pool_idx, : req.extend_range.end
] ]
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
@@ -693,7 +693,7 @@ class MambaRadixCache(BasePrefixCache):
return _skip_cache_unfinished_req(req) return _skip_cache_unfinished_req(req)
kv_indices_orig = self.req_to_token_pool.req_to_token[ kv_indices_orig = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids) req.kv.req_pool_idx, : len(token_ids)
] ]
# kv_indices is the kv indices to be cached # kv_indices is the kv indices to be cached
kv_indices = kv_indices_orig[:cache_len] kv_indices = kv_indices_orig[:cache_len]
@@ -786,7 +786,7 @@ class MambaRadixCache(BasePrefixCache):
), f"{new_prefix_len=}, {len(new_indices)=}" ), f"{new_prefix_len=}, {len(new_indices)=}"
self.req_to_token_pool.write( self.req_to_token_pool.write(
(req.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))),
new_indices[req.kv.cache_protected_len :], new_indices[req.kv.cache_protected_len :],
) )
python/sglang/srt/mem_cache/memory_pool.py
+11 -9
View File
@@ -292,7 +292,7 @@ class ReqToTokenPool:
def alloc(self, reqs: list[Req]) -> Optional[List[int]]: def alloc(self, reqs: list[Req]) -> Optional[List[int]]:
# Indices of reqs that already have a req_pool_idx and will reuse # Indices of reqs that already have a req_pool_idx and will reuse
# their existing slot (e.g. chunked prefill continuing across chunks). # their existing slot (e.g. chunked prefill continuing across chunks).
reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None] reusing = [i for i, r in enumerate(reqs) if r.kv.req_pool_idx is not None]
# NOTE: this check is relaxed temporarily # NOTE: this check is relaxed temporarily
# https://github.com/sgl-project/sglang/pull/20476 # https://github.com/sgl-project/sglang/pull/20476
# if not any(r.is_dllm() for r in reqs): # if not any(r.is_dllm() for r in reqs):
@@ -309,10 +309,10 @@ class ReqToTokenPool:
return None return None
offset = 0 offset = 0
for r in reqs: for r in reqs:
if r.req_pool_idx is None: if r.kv.req_pool_idx is None:
r.req_pool_idx = select_index[offset] r.kv.req_pool_idx = select_index[offset]
offset += 1 offset += 1
return [r.req_pool_idx for r in reqs] return [r.kv.req_pool_idx for r in reqs]
def alloc_rows(self, need_size: int) -> Optional[List[int]]: def alloc_rows(self, need_size: int) -> Optional[List[int]]:
"""Take need_size rows and bump their generation, with no Req bound to """Take need_size rows and bump their generation, with no Req bound to
@@ -338,9 +338,9 @@ class ReqToTokenPool:
self.free_slots.extend(indices) self.free_slots.extend(indices)
def free(self, req: Req): def free(self, req: Req):
assert req.req_pool_idx is not None, "request must have req_pool_idx" assert req.kv.req_pool_idx is not None, "request must have req_pool_idx"
self.free_rows([req.req_pool_idx]) self.free_rows([req.kv.req_pool_idx])
req.req_pool_idx = None req.kv.req_pool_idx = None
def clear(self): def clear(self):
self.free_slots = list(range(1, self._alloc_size)) self.free_slots = list(range(1, self._alloc_size))
@@ -1493,7 +1493,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
set_mamba_track_indices_from_reqs reads correct slot indices. set_mamba_track_indices_from_reqs reads correct slot indices.
""" """
req.mamba_ping_pong_track_buffer[idx] = value req.mamba_ping_pong_track_buffer[idx] = value
self.req_index_to_mamba_ping_pong_track_buffer_mapping[req.req_pool_idx] = ( self.req_index_to_mamba_ping_pong_track_buffer_mapping[req.kv.req_pool_idx] = (
req.mamba_ping_pong_track_buffer req.mamba_ping_pong_track_buffer
) )
@@ -1530,7 +1530,9 @@ class HybridReqToTokenPool(ReqToTokenPool):
if self.enable_mamba_extra_buffer: if self.enable_mamba_extra_buffer:
mamba_ping_pong_track_buffer_to_free = ( mamba_ping_pong_track_buffer_to_free = (
self.req_index_to_mamba_ping_pong_track_buffer_mapping[req.req_pool_idx] self.req_index_to_mamba_ping_pong_track_buffer_mapping[
req.kv.req_pool_idx
]
) )
if mamba_ping_pong_track_buffer_to_keep is not None: if mamba_ping_pong_track_buffer_to_keep is not None:
assert mamba_ping_pong_track_buffer_to_keep in [ assert mamba_ping_pong_track_buffer_to_keep in [
python/sglang/srt/mem_cache/pure_swa_radix_cache.py
+2 -2
View File
@@ -78,14 +78,14 @@ class PureSWARadixCache(RadixCache):
kv_committed_len = kv_len_to_handle kv_committed_len = kv_len_to_handle
if self.disable: if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len req.kv.req_pool_idx, :kv_committed_len
] ]
self.token_to_kv_pool_allocator.free(kv_indices) self.token_to_kv_pool_allocator.free(kv_indices)
return return
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len req.kv.req_pool_idx, :kv_committed_len
] ]
radix_key = RadixKey( radix_key = RadixKey(
python/sglang/srt/mem_cache/radix_cache.py
+4 -4
View File
@@ -467,7 +467,7 @@ class RadixCache(BasePrefixCache):
if self.disable: if self.disable:
# The protected prefix is not this req's to free. # The protected prefix is not this req's to free.
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle req.kv.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle
] ]
self.token_to_kv_pool_allocator.free_segment( self.token_to_kv_pool_allocator.free_segment(
kv_indices, start_pos=req.kv.cache_protected_len kv_indices, start_pos=req.kv.cache_protected_len
@@ -476,7 +476,7 @@ class RadixCache(BasePrefixCache):
token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids) req.kv.req_pool_idx, : len(token_ids)
] ]
radix_key = RadixKey( radix_key = RadixKey(
@@ -520,7 +520,7 @@ class RadixCache(BasePrefixCache):
token_ids = req.get_fill_ids() token_ids = req.get_fill_ids()
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids) req.kv.req_pool_idx, : len(token_ids)
] ]
radix_key = RadixKey( radix_key = RadixKey(
@@ -558,7 +558,7 @@ class RadixCache(BasePrefixCache):
), f"{len(new_indices)=}, {len(radix_key)=}" ), f"{len(new_indices)=}, {len(radix_key)=}"
self.req_to_token_pool.write( self.req_to_token_pool.write(
(req.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))),
new_indices[req.kv.cache_protected_len :], new_indices[req.kv.cache_protected_len :],
) )
python/sglang/srt/mem_cache/radix_cache_cpp.py
+5 -5
View File
@@ -185,10 +185,10 @@ class RadixCacheCpp(BasePrefixCache):
): ):
"""Cache request when it finishes.""" """Cache request when it finishes."""
self._reject_cache_salt(req.cache_salt) self._reject_cache_salt(req.cache_salt)
assert req.req_pool_idx is not None assert req.kv.req_pool_idx is not None
token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_len_to_handle req.kv.req_pool_idx, :kv_len_to_handle
].to(dtype=torch.int64, copy=True) ].to(dtype=torch.int64, copy=True)
# NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned
@@ -223,11 +223,11 @@ class RadixCacheCpp(BasePrefixCache):
def cache_unfinished_req(self, req: Req, chunked=False): def cache_unfinished_req(self, req: Req, chunked=False):
"""Cache request when it is unfinished.""" """Cache request when it is unfinished."""
self._reject_cache_salt(req.cache_salt) self._reject_cache_salt(req.cache_salt)
assert req.req_pool_idx is not None assert req.kv.req_pool_idx is not None
token_ids = req.get_fill_ids() token_ids = req.get_fill_ids()
prefill_len = len(token_ids) # prefill only (maybe chunked) prefill_len = len(token_ids) # prefill only (maybe chunked)
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :prefill_len req.kv.req_pool_idx, :prefill_len
].to(dtype=torch.int64, copy=True) ].to(dtype=torch.int64, copy=True)
# NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned
@@ -254,7 +254,7 @@ class RadixCacheCpp(BasePrefixCache):
) )
reused_indices = new_indices[old_prefix_len:new_prefix_len] reused_indices = new_indices[old_prefix_len:new_prefix_len]
self.req_to_token_pool.req_to_token[ self.req_to_token_pool.req_to_token[
req.req_pool_idx, old_prefix_len:new_prefix_len req.kv.req_pool_idx, old_prefix_len:new_prefix_len
] = reused_indices ] = reused_indices
if req.last_node != new_last_node: if req.last_node != new_last_node:
python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py
+4 -4
View File
@@ -142,18 +142,18 @@ class SparseCoordinator:
Registers the request in the state tracker to enable sparse attention processing. Registers the request in the state tracker to enable sparse attention processing.
""" """
if req.req_pool_idx is not None: if req.kv.req_pool_idx is not None:
self.states.register(req.req_pool_idx, len(req.origin_input_ids)) self.states.register(req.kv.req_pool_idx, len(req.origin_input_ids))
def on_request_end(self, req: "Req") -> None: def on_request_end(self, req: "Req") -> None:
""" """
Handle request end event. Called when a request is completed or aborted. Handle request end event. Called when a request is completed or aborted.
Cleans up request-specific state and releases resources. Cleans up request-specific state and releases resources.
""" """
if req.req_pool_idx is None: if req.kv.req_pool_idx is None:
return return
self.states.clear(req.req_pool_idx) self.states.clear(req.kv.req_pool_idx)
# TODO: Implement request end handling # TODO: Implement request end handling
# - Release host indices if any were allocated for offloading # - Release host indices if any were allocated for offloading
python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py
+1 -1
View File
@@ -409,7 +409,7 @@ class FlexKVRadixCache(RadixCache):
if not token_ids: if not token_ids:
return return
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len req.kv.req_pool_idx, :kv_committed_len
] ]
# Anchor on the new last_device_node so FlexKV's lock matches # Anchor on the new last_device_node so FlexKV's lock matches
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
+1 -1
View File
@@ -463,7 +463,7 @@ class LMCRadixCache(RadixCache):
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len req.kv.req_pool_idx, :kv_committed_len
] ]
# Use super() to avoid a redundant LOOKUP — we only need new_last_node from radix. # Use super() to avoid a redundant LOOKUP — we only need new_last_node from radix.
python/sglang/srt/mem_cache/swa_radix_cache.py
+5 -5
View File
@@ -465,14 +465,14 @@ class SWARadixCache(BasePrefixCache):
"""Cache request when it finishes.""" """Cache request when it finishes."""
if self.disable: if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_len_to_handle req.kv.req_pool_idx, :kv_len_to_handle
] ]
self.token_to_kv_pool_allocator.free(kv_indices) self.token_to_kv_pool_allocator.free(kv_indices)
return return
token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_len_to_handle req.kv.req_pool_idx, :kv_len_to_handle
] ]
radix_key = RadixKey( radix_key = RadixKey(
@@ -516,7 +516,7 @@ class SWARadixCache(BasePrefixCache):
"""Cache request when it is unfinished.""" """Cache request when it is unfinished."""
if self.disable: if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.extend_range.end req.kv.req_pool_idx, : req.extend_range.end
] ]
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
@@ -525,7 +525,7 @@ class SWARadixCache(BasePrefixCache):
token_ids = req.get_fill_ids() token_ids = req.get_fill_ids()
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids) req.kv.req_pool_idx, : len(token_ids)
] ]
radix_key = RadixKey( radix_key = RadixKey(
@@ -561,7 +561,7 @@ class SWARadixCache(BasePrefixCache):
assert old_prefix_len <= len(new_indices), f"{old_prefix_len=}, {new_indices=}" assert old_prefix_len <= len(new_indices), f"{old_prefix_len=}, {new_indices=}"
assert new_prefix_len <= len(new_indices), f"{new_prefix_len=}, {new_indices=}" assert new_prefix_len <= len(new_indices), f"{new_prefix_len=}, {new_indices=}"
self.req_to_token_pool.write( self.req_to_token_pool.write(
(req.req_pool_idx, slice(old_prefix_len, len(new_indices))), (req.kv.req_pool_idx, slice(old_prefix_len, len(new_indices))),
new_indices[old_prefix_len:], new_indices[old_prefix_len:],
) )
python/sglang/srt/mem_cache/unified_radix_cache.py
+7 -7
View File
@@ -825,7 +825,7 @@ class UnifiedRadixCache(BasePrefixCache):
if self.disable: if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_len_to_handle req.kv.req_pool_idx, :kv_len_to_handle
] ]
self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0) self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0)
for comp in self._components_tuple: for comp in self._components_tuple:
@@ -834,7 +834,7 @@ class UnifiedRadixCache(BasePrefixCache):
token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_len_to_handle req.kv.req_pool_idx, :kv_len_to_handle
] ]
result = None result = None
@@ -918,13 +918,13 @@ class UnifiedRadixCache(BasePrefixCache):
if self.disable: if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids) req.kv.req_pool_idx, : len(token_ids)
] ]
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
return return
kv_indices_orig = self.req_to_token_pool.req_to_token[ kv_indices_orig = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids) req.kv.req_pool_idx, : len(token_ids)
] ]
# components prepare insert data + return effective cache_len # components prepare insert data + return effective cache_len
@@ -991,7 +991,7 @@ class UnifiedRadixCache(BasePrefixCache):
new_indices new_indices
), f"{new_prefix_len=}, {len(new_indices)=}" ), f"{new_prefix_len=}, {len(new_indices)=}"
self.req_to_token_pool.write( self.req_to_token_pool.write(
(req.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))),
new_indices[req.kv.cache_protected_len :], new_indices[req.kv.cache_protected_len :],
) )
@@ -1160,7 +1160,7 @@ class UnifiedRadixCache(BasePrefixCache):
) -> tuple[torch.Tensor, list[PoolTransfer]]: ) -> tuple[torch.Tensor, list[PoolTransfer]]:
num_tokens = req.seqlen - 1 num_tokens = req.seqlen - 1
full_indices = self.req_to_token_pool.req_to_token[ full_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :num_tokens req.kv.req_pool_idx, :num_tokens
].to(torch.int64) ].to(torch.int64)
full_indices = self._pad_retraction_indices(full_indices, self.page_size) full_indices = self._pad_retraction_indices(full_indices, self.page_size)
@@ -1171,7 +1171,7 @@ class UnifiedRadixCache(BasePrefixCache):
window_start = max(0, num_tokens - self.sliding_window_size) window_start = max(0, num_tokens - self.sliding_window_size)
window_start = window_start // self.page_size * self.page_size window_start = window_start // self.page_size * self.page_size
window_indices = self.req_to_token_pool.req_to_token[ window_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, window_start:num_tokens req.kv.req_pool_idx, window_start:num_tokens
].to(torch.int64) ].to(torch.int64)
swa_indices = kv_cache.translate_loc_from_full_to_swa(window_indices) swa_indices = kv_cache.translate_loc_from_full_to_swa(window_indices)
assert bool( assert bool(
python/sglang/srt/session/streaming_session.py
+21 -30
View File
@@ -45,7 +45,6 @@ class SessionSlot:
virtual_node: _VirtualNode = field(default_factory=_VirtualNode) virtual_node: _VirtualNode = field(default_factory=_VirtualNode)
# KV pool state # KV pool state
req_pool_idx: Optional[int] = None
kv: ReqKvInfo = field(default_factory=ReqKvInfo) kv: ReqKvInfo = field(default_factory=ReqKvInfo)
# 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)
@@ -63,15 +62,8 @@ class SessionSlot:
mamba_last_track_seqlen: Any = None mamba_last_track_seqlen: Any = None
mamba_branching_seqlen: Any = None mamba_branching_seqlen: Any = None
@property
def is_holding_kv(self) -> bool:
"""Whether this slot currently holds KV pool resources."""
return self.req_pool_idx is not None
def save_from_req(self, req: Req, is_first: bool): def save_from_req(self, req: Req, is_first: bool):
"""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
if is_first: if is_first:
self.last_node = req.last_node self.last_node = req.last_node
self.swa_uuid_for_lock = req.swa_uuid_for_lock self.swa_uuid_for_lock = req.swa_uuid_for_lock
@@ -93,7 +85,6 @@ class SessionSlot:
# Ownership moved to the slot; clear the req's references so a later # Ownership moved to the slot; clear the req's references so a later
# alloc/retract path cannot mistake slot-owned mamba state for its own. # alloc/retract path cannot mistake slot-owned mamba state for its own.
req.req_pool_idx = None
req.kv = ReqKvInfo() req.kv = ReqKvInfo()
req.mamba_pool_idx = None req.mamba_pool_idx = None
req.mamba_ping_pong_track_buffer = None req.mamba_ping_pong_track_buffer = None
@@ -104,7 +95,6 @@ class SessionSlot:
def restore_to_req(self, req: Req): def restore_to_req(self, req: Req):
"""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.kv = copy.copy(self.kv) req.kv = copy.copy(self.kv)
req.swa_uuid_for_lock = self.swa_uuid_for_lock req.swa_uuid_for_lock = self.swa_uuid_for_lock
req.skip_lock_node_ids = self.skip_lock_node_ids req.skip_lock_node_ids = self.skip_lock_node_ids
@@ -189,7 +179,7 @@ class StreamingSession(BasePrefixCache):
return session_id in self.slots return session_id in self.slots
def any_holding_kv(self) -> bool: def any_holding_kv(self) -> bool:
return any(s.is_holding_kv for s in self.slots.values()) return any(s.kv.is_held for s in self.slots.values())
# -- Try-handle entries for composition (see class docstring) -- # -- Try-handle entries for composition (see class docstring) --
@@ -217,7 +207,7 @@ class StreamingSession(BasePrefixCache):
if not _is_streaming(req): if not _is_streaming(req):
return None return None
slot = self.slots.get(req.session.session_id) slot = self.slots.get(req.session.session_id)
if slot is None or not slot.is_holding_kv: if slot is None or not slot.kv.is_held:
return None return None
if req.to_finish is not None: if req.to_finish is not None:
req.session.abort_req() req.session.abort_req()
@@ -290,7 +280,7 @@ class StreamingSession(BasePrefixCache):
self._free_tail(slot, req, prefix_len) self._free_tail(slot, req, prefix_len)
device_indices = self.req_to_token_pool.req_to_token[ device_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :prefix_len req.kv.req_pool_idx, :prefix_len
].to(dtype=torch.int64) ].to(dtype=torch.int64)
return MatchResult( return MatchResult(
@@ -330,7 +320,6 @@ class StreamingSession(BasePrefixCache):
# return the (possibly extra_buffer ping-pong) slots to # return the (possibly extra_buffer ping-pong) slots to
# 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,
kv=copy.copy(req.kv), kv=copy.copy(req.kv),
last_node=req.last_node, last_node=req.last_node,
swa_uuid_for_lock=req.swa_uuid_for_lock, swa_uuid_for_lock=req.swa_uuid_for_lock,
@@ -347,7 +336,6 @@ class StreamingSession(BasePrefixCache):
slot.kv.kv_allocated_len, req.kv.kv_allocated_len 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.kv = ReqKvInfo() req.kv = ReqKvInfo()
req.session.abort_req() req.session.abort_req()
return True return True
@@ -386,7 +374,7 @@ class StreamingSession(BasePrefixCache):
return False return False
if chunked: if chunked:
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.extend_range.end req.kv.req_pool_idx, : req.extend_range.end
] ]
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
return True return True
@@ -441,9 +429,7 @@ class StreamingSession(BasePrefixCache):
protected_len = slot.kv.cache_protected_len protected_len = slot.kv.cache_protected_len
lock_node = slot.last_node lock_node = slot.last_node
tokens_freed = ( tokens_freed = (
max(0, slot.kv.kv_allocated_len - protected_len) max(0, slot.kv.kv_allocated_len - protected_len) if slot.kv.is_held else 0
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
@@ -458,12 +444,12 @@ class StreamingSession(BasePrefixCache):
), ),
) )
if slot.is_holding_kv: if slot.kv.is_held:
start = protected_len start = protected_len
end = slot.kv.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.kv.req_pool_idx, start:end
] ]
self.token_to_kv_pool_allocator.free(kv_indices) self.token_to_kv_pool_allocator.free(kv_indices)
self.req_to_token_pool.free(slot) self.req_to_token_pool.free(slot)
@@ -483,9 +469,10 @@ class StreamingSession(BasePrefixCache):
total = 0 total = 0
for slot in self.slots.values(): for slot in self.slots.values():
in_batch = ( in_batch = (
active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs active_pool_idxs is not None
and slot.kv.req_pool_idx in active_pool_idxs
) )
if slot.is_holding_kv and not in_batch: if slot.kv.is_held and not in_batch:
allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size) allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
total += allocated - slot.kv.cache_protected_len total += allocated - slot.kv.cache_protected_len
return total return total
@@ -499,9 +486,10 @@ class StreamingSession(BasePrefixCache):
total = 0 total = 0
for slot in self.slots.values(): for slot in self.slots.values():
in_batch = ( in_batch = (
active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs active_pool_idxs is not None
and slot.kv.req_pool_idx in active_pool_idxs
) )
if slot.is_holding_kv and not in_batch: if slot.kv.is_held and not in_batch:
allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size) allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
total += allocated - max( total += allocated - max(
slot.kv.cache_protected_len, slot.kv.swa_evicted_seqlen slot.kv.cache_protected_len, slot.kv.swa_evicted_seqlen
@@ -513,9 +501,9 @@ class StreamingSession(BasePrefixCache):
def _owned(s): def _owned(s):
in_batch = ( in_batch = (
active_pool_idxs is not None and s.req_pool_idx in active_pool_idxs active_pool_idxs is not None and s.kv.req_pool_idx in active_pool_idxs
) )
return s.is_holding_kv and not in_batch return s.kv.is_held and not in_batch
return sum(_owned(s) for s in self.slots.values()) return sum(_owned(s) for s in self.slots.values())
@@ -529,7 +517,8 @@ class StreamingSession(BasePrefixCache):
total = 0 total = 0
for slot in self.slots.values(): for slot in self.slots.values():
in_batch = ( in_batch = (
active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs active_pool_idxs is not None
and slot.kv.req_pool_idx in active_pool_idxs
) )
if in_batch: if in_batch:
continue continue
@@ -559,7 +548,9 @@ 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.kv_allocated_len) self._free_kv_aligned(
slot.kv.req_pool_idx, prefix_len, slot.kv.kv_allocated_len
)
slot.kv.kv_allocated_len = prefix_len slot.kv.kv_allocated_len = prefix_len
slot.kv.kv_committed_len = min(slot.kv.kv_committed_len, prefix_len) slot.kv.kv_committed_len = min(slot.kv.kv_committed_len, prefix_len)
slot.kv.swa_evicted_seqlen = min(slot.kv.swa_evicted_seqlen, prefix_len) slot.kv.swa_evicted_seqlen = min(slot.kv.swa_evicted_seqlen, prefix_len)
@@ -574,7 +565,7 @@ 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.kv_allocated_len) self._free_kv_aligned(req.kv.req_pool_idx, target, req.kv.kv_allocated_len)
req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target) req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target)
req.kv.kv_committed_len = min(req.kv.kv_committed_len, target) req.kv.kv_committed_len = min(req.kv.kv_committed_len, target)
req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target) req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target)
python/sglang/test/scripted_runtime/req_handle.py
+1 -1
View File
@@ -42,7 +42,7 @@ class ScriptedReqHandle:
@property @property
def kv_pages(self) -> int: def kv_pages(self) -> int:
req = self.req req = self.req
if req is None or not req.is_holding_kv: if req is None or not req.kv.is_held:
return 0 return 0
page_size = self.context.scheduler.page_size page_size = self.context.scheduler.page_size
return (req.kv.kv_allocated_len + page_size - 1) // page_size return (req.kv.kv_allocated_len + page_size - 1) // page_size
test/manual/chunked_prefill/test_scripted_abort.py
+14 -14
View File
@@ -23,7 +23,7 @@ def _drain_until_released(t: ScriptedContext, *handles: ScriptedReqHandle):
if all( if all(
h.kv_pages == 0 h.kv_pages == 0
and h.lock_refs == 0 and h.lock_refs == 0
and (h.req is None or h.req.req_pool_idx is None) and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles for h in handles
): ):
return return
@@ -57,7 +57,7 @@ class TestAbortBasic(ScriptedTestCase):
r.kv_pages == 0 r.kv_pages == 0
), f"abort must release KV; r.kv_pages={r.kv_pages} after abort" ), f"abort must release KV; r.kv_pages={r.kv_pages} after abort"
assert ( assert (
r.req is None or r.req.req_pool_idx is None r.req is None or r.req.kv.req_pool_idx is None
), f"abort must release row; r.req={r.req} after abort" ), f"abort must release row; r.req={r.req} after abort"
assert ( assert (
r.lock_refs == 0 r.lock_refs == 0
@@ -78,7 +78,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
def test_abort_at_chunk_mid(self): def test_abort_at_chunk_mid(self):
self.server.execute_script(self._script_abort_at_chunk_mid) self.server.execute_script(self._script_abort_at_chunk_mid)
@@ -124,7 +124,7 @@ class TestAbortBasic(ScriptedTestCase):
t.abort(r) t.abort(r)
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
assert r.lock_refs == 0 assert r.lock_refs == 0
def test_abort_at_admission_step(self): def test_abort_at_admission_step(self):
@@ -137,7 +137,7 @@ class TestAbortBasic(ScriptedTestCase):
t.abort(r) t.abort(r)
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
def test_abort_then_start_same_step_new_rid(self): def test_abort_then_start_same_step_new_rid(self):
self.server.execute_script(self._script_abort_then_start_same_step_new_rid) self.server.execute_script(self._script_abort_then_start_same_step_new_rid)
@@ -191,7 +191,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from _drain_until_released(t, *reqs) yield from _drain_until_released(t, *reqs)
for r in reqs: for r in reqs:
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
def test_abort_unknown_rid_noop(self): def test_abort_unknown_rid_noop(self):
self.server.execute_script(self._script_abort_unknown_rid_noop) self.server.execute_script(self._script_abort_unknown_rid_noop)
@@ -264,7 +264,7 @@ class TestAbortBasic(ScriptedTestCase):
t.abort(r) t.abort(r)
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
assert r.lock_refs == 0 assert r.lock_refs == 0
def test_double_abort_idempotent(self): def test_double_abort_idempotent(self):
@@ -351,7 +351,7 @@ class TestAbortBasic(ScriptedTestCase):
f"aborted req revived and ran another chunk; " f"aborted req revived and ran another chunk; "
f"chunks_done went {chunks_after_abort} -> {r.chunks_done}" f"chunks_done went {chunks_after_abort} -> {r.chunks_done}"
) )
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
def test_abort_mid_chunk_no_extra_radix_node(self): def test_abort_mid_chunk_no_extra_radix_node(self):
self.server.execute_script(self._script_abort_mid_chunk_no_extra_radix_node) self.server.execute_script(self._script_abort_mid_chunk_no_extra_radix_node)
@@ -369,7 +369,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
chunks_after_release = r.chunks_done chunks_after_release = r.chunks_done
for _ in range(4): for _ in range(4):
yield yield
@@ -406,7 +406,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from run_until_finished(r2) yield from run_until_finished(r2)
assert r2.finished, "resubmit under same rid must complete independently" assert r2.finished, "resubmit under same rid must complete independently"
assert r1.kv_pages == 0, "aborted r1 must release KV before resubmit" assert r1.kv_pages == 0, "aborted r1 must release KV before resubmit"
assert r1.req is None or r1.req.req_pool_idx is None assert r1.req is None or r1.req.kv.req_pool_idx is None
assert r1.lock_refs == 0 assert r1.lock_refs == 0
def test_abort_during_gap_inflight_middle_chunks_positive(self): def test_abort_during_gap_inflight_middle_chunks_positive(self):
@@ -429,7 +429,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
assert not r.is_chunking, "aborted gap req must not re-enter chunking" assert not r.is_chunking, "aborted gap req must not re-enter chunking"
yield yield
@@ -511,7 +511,7 @@ class TestAbortBasic(ScriptedTestCase):
assert r1.kv_pages == 0, ( assert r1.kv_pages == 0, (
f"force_retract + abort same yield must release KV; got " f"{r1.kv_pages}" f"force_retract + abort same yield must release KV; got " f"{r1.kv_pages}"
) )
assert r1.req is None or r1.req.req_pool_idx is None, ( assert r1.req is None or r1.req.kv.req_pool_idx is None, (
f"force_retract + abort same yield must release row; got " f"{r1.req}" f"force_retract + abort same yield must release row; got " f"{r1.req}"
) )
assert r1.lock_refs == 0, ( assert r1.lock_refs == 0, (
@@ -540,7 +540,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from run_until(r2, lambda h: h.is_chunking) yield from run_until(r2, lambda h: h.is_chunking)
assert r1.kv_pages == 0 assert r1.kv_pages == 0
assert r1.req is None or r1.req.req_pool_idx is None assert r1.req is None or r1.req.kv.req_pool_idx is None
assert r1.lock_refs == 0 assert r1.lock_refs == 0
yield from run_until_finished(r2) yield from run_until_finished(r2)
assert r2.finished, "baton handoff must let r2 complete" assert r2.finished, "baton handoff must let r2 complete"
@@ -570,7 +570,7 @@ class TestAbortPP(ScriptedTestCase):
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
assert r.lock_refs == 0 assert r.lock_refs == 0
assert r.finished assert r.finished
test/manual/chunked_prefill/test_scripted_kv_pressure.py
+2 -2
View File
@@ -409,13 +409,13 @@ class TestKVPressureSmallPool(ScriptedTestCase):
if ( if (
r_chunk.kv_pages == 0 r_chunk.kv_pages == 0
and r_chunk.lock_refs == 0 and r_chunk.lock_refs == 0
and (r_chunk.req is None or r_chunk.req.req_pool_idx is None) and (r_chunk.req is None or r_chunk.req.kv.req_pool_idx is None)
): ):
break break
yield yield
assert r_chunk.kv_pages == 0, f"kv_pages={r_chunk.kv_pages}" assert r_chunk.kv_pages == 0, f"kv_pages={r_chunk.kv_pages}"
assert r_chunk.lock_refs == 0, f"lock_refs={r_chunk.lock_refs}" assert r_chunk.lock_refs == 0, f"lock_refs={r_chunk.lock_refs}"
assert r_chunk.req is None or r_chunk.req.req_pool_idx is None assert r_chunk.req is None or r_chunk.req.kv.req_pool_idx is None
t.abort(ballast) t.abort(ballast)
for _ in range(200): for _ in range(200):
test/manual/chunked_prefill/test_scripted_lifecycle.py
+20 -12
View File
@@ -18,7 +18,7 @@ def _drain_until_released(t, *handles):
if all( if all(
h.kv_pages == 0 h.kv_pages == 0
and h.lock_refs == 0 and h.lock_refs == 0
and (h.req is None or h.req.req_pool_idx is None) and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles for h in handles
): ):
return return
@@ -223,7 +223,7 @@ class TestLifecycleBasic(ScriptedTestCase):
r = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) r = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True)
yield from run_until_finished(r) yield from run_until_finished(r)
assert r.finished assert r.finished
assert r.req.req_pool_idx is None assert r.req.kv.req_pool_idx is None
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.lock_refs == 0 assert r.lock_refs == 0
@@ -235,7 +235,7 @@ class TestLifecycleBasic(ScriptedTestCase):
r1 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) r1 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True)
yield from run_until_finished(r1) yield from run_until_finished(r1)
yield from _drain_until_released(t, r1) yield from _drain_until_released(t, r1)
assert r1.req.req_pool_idx is None and r1.kv_pages == 0 and r1.lock_refs == 0 assert r1.req.kv.req_pool_idx is None and r1.kv_pages == 0 and r1.lock_refs == 0
r1_output_len = len(r1.req.output_ids) r1_output_len = len(r1.req.output_ids)
r2 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) r2 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True)
@@ -243,7 +243,7 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from _drain_until_released(t, r2) yield from _drain_until_released(t, r2)
assert r1.finished and r2.finished assert r1.finished and r2.finished
assert r1_output_len == 2 and len(r2.req.output_ids) == 2 assert r1_output_len == 2 and len(r2.req.output_ids) == 2
assert r2.req.req_pool_idx is None and r2.kv_pages == 0 and r2.lock_refs == 0 assert r2.req.kv.req_pool_idx is None and r2.kv_pages == 0 and r2.lock_refs == 0
def test_five_seq_clean(self): def test_five_seq_clean(self):
self.server.execute_script(self._script_five_seq_clean) self.server.execute_script(self._script_five_seq_clean)
@@ -256,7 +256,7 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r) yield from run_until_finished(r)
assert r.finished assert r.finished
assert len(r.req.output_ids) == 2 assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None assert r.req.kv.req_pool_idx is None
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.lock_refs == 0 assert r.lock_refs == 0
reqs.append(r) reqs.append(r)
@@ -300,7 +300,9 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r) yield from run_until_finished(r)
assert r.finished assert r.finished
assert len(r.req.output_ids) == 2 assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 assert (
r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
)
if prompt == VERY_LONG_PROMPT_LEN: if prompt == VERY_LONG_PROMPT_LEN:
assert r.chunks_done == 8 assert r.chunks_done == 8
else: else:
@@ -319,7 +321,7 @@ class TestLifecycleBasic(ScriptedTestCase):
assert r.finished assert r.finished
assert len(r.req.output_ids) == 1 assert len(r.req.output_ids) == 1
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
assert r.kv_pages == 0 and r.lock_refs == 0 assert r.kv_pages == 0 and r.lock_refs == 0
if L > DEFAULT_CHUNK_SIZE: if L > DEFAULT_CHUNK_SIZE:
assert ( assert (
@@ -341,7 +343,7 @@ class TestLifecycleBasic(ScriptedTestCase):
assert r.finished assert r.finished
assert len(r.req.output_ids) == 1 assert len(r.req.output_ids) == 1
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
assert r.kv_pages == 0 and r.lock_refs == 0 assert r.kv_pages == 0 and r.lock_refs == 0
if L > DEFAULT_CHUNK_SIZE: if L > DEFAULT_CHUNK_SIZE:
assert ( assert (
@@ -360,7 +362,9 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r) yield from run_until_finished(r)
assert r.finished assert r.finished
assert len(r.req.output_ids) == 2 assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 assert (
r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
)
for _ in range(20): for _ in range(20):
yield yield
@@ -377,7 +381,9 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r) yield from run_until_finished(r)
assert r.finished assert r.finished
assert len(r.req.output_ids) == 2 assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 assert (
r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
)
if L == VERY_LONG_PROMPT_LEN: if L == VERY_LONG_PROMPT_LEN:
assert r.chunks_done == 8 assert r.chunks_done == 8
else: else:
@@ -394,7 +400,9 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r) yield from run_until_finished(r)
assert r.finished assert r.finished
assert len(r.req.output_ids) == 2 assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 assert (
r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
)
for _ in range(5): for _ in range(5):
yield yield
t.flush_cache() t.flush_cache()
@@ -432,7 +440,7 @@ class TestLifecycleBasic(ScriptedTestCase):
) )
assert r.finished or _error_message(r) is not None assert r.finished or _error_message(r) is not None
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
assert r.lock_refs == 0 assert r.lock_refs == 0
test/manual/chunked_prefill/test_scripted_lora.py
+3 -1
View File
@@ -222,7 +222,9 @@ class TestLoRAAdapterEviction(ScriptedTestCase):
t.abort(r_a) t.abort(r_a)
for _ in range(12): for _ in range(12):
if r_a.kv_pages == 0 and (r_a.req is None or r_a.req.req_pool_idx is None): if r_a.kv_pages == 0 and (
r_a.req is None or r_a.req.kv.req_pool_idx is None
):
break break
yield yield
test/manual/chunked_prefill/test_scripted_multi_req.py
+1 -1
View File
@@ -173,7 +173,7 @@ class TestMultiReqBasic(ScriptedTestCase):
for _ in range(5): for _ in range(5):
yield yield
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
def test_rid_reuse_after_finish(self): def test_rid_reuse_after_finish(self):
self.server.execute_script(self._script_rid_reuse_after_finish) self.server.execute_script(self._script_rid_reuse_after_finish)
test/manual/chunked_prefill/test_scripted_pp.py
+1 -1
View File
@@ -31,7 +31,7 @@ def _expected_chunks(prompt_len: int, chunk_size: int) -> int:
def _drain_until_released(t, *handles): def _drain_until_released(t, *handles):
for _ in range(16): for _ in range(16):
if all( if all(
h.kv_pages == 0 and (h.req is None or h.req.req_pool_idx is None) h.kv_pages == 0 and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles for h in handles
): ):
return return
test/manual/chunked_prefill/test_scripted_priority.py
+2 -2
View File
@@ -168,13 +168,13 @@ class TestPriorityBasic(ScriptedTestCase):
if ( if (
r.kv_pages == 0 r.kv_pages == 0
and r.lock_refs == 0 and r.lock_refs == 0
and (r.req is None or r.req.req_pool_idx is None) and (r.req is None or r.req.kv.req_pool_idx is None)
): ):
break break
yield yield
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.lock_refs == 0 assert r.lock_refs == 0
assert r.req is None or r.req.req_pool_idx is None assert r.req is None or r.req.kv.req_pool_idx is None
t.continue_generation() t.continue_generation()
yield yield
assert r.kv_pages == 0 and r.lock_refs == 0 assert r.kv_pages == 0 and r.lock_refs == 0
test/manual/chunked_prefill/test_scripted_regression.py
+9 -9
View File
@@ -19,7 +19,7 @@ def _drain_until_released(t, *handles):
if all( if all(
h.kv_pages == 0 h.kv_pages == 0
and h.lock_refs == 0 and h.lock_refs == 0
and (h.req is None or h.req.req_pool_idx is None) and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles for h in handles
): ):
return return
@@ -41,7 +41,7 @@ class TestRegressionBasic(ScriptedTestCase):
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req.req_pool_idx is None assert r.req.kv.req_pool_idx is None
assert r.lock_refs == 0 assert r.lock_refs == 0
assert not r.is_chunking assert not r.is_chunking
assert r.req.inflight_middle_chunks == 0 assert r.req.inflight_middle_chunks == 0
@@ -58,7 +58,7 @@ class TestRegressionBasic(ScriptedTestCase):
yield yield
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.req.req_pool_idx is None assert r.req.kv.req_pool_idx is None
assert r.lock_refs == 0 assert r.lock_refs == 0
assert not r.is_chunking assert not r.is_chunking
@@ -240,7 +240,7 @@ class TestRegressionBasic(ScriptedTestCase):
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
assert r.req.req_pool_idx is not None, "row must be held mid-chunk" assert r.req.kv.req_pool_idx is not None, "row must be held mid-chunk"
assert r.kv_pages > 0, "committed KV must be held mid-chunk" assert r.kv_pages > 0, "committed KV must be held mid-chunk"
assert r.lock_refs >= 1, "radix lock_ref must be held mid-chunk" assert r.lock_refs >= 1, "radix lock_ref must be held mid-chunk"
@@ -248,8 +248,8 @@ class TestRegressionBasic(ScriptedTestCase):
yield from _drain_until_released(t, r) yield from _drain_until_released(t, r)
assert ( assert (
r.req.req_pool_idx is None r.req.kv.req_pool_idx is None
), f"96d4749094: abort must release row; got row_idx={r.req.req_pool_idx!r}" ), f"96d4749094: abort must release row; got row_idx={r.req.kv.req_pool_idx!r}"
assert ( assert (
r.kv_pages == 0 r.kv_pages == 0
), f"96d4749094: abort must release KV; got kv_pages={r.kv_pages}" ), f"96d4749094: abort must release KV; got kv_pages={r.kv_pages}"
@@ -269,14 +269,14 @@ class TestRegressionBasic(ScriptedTestCase):
def _script_pause_retract_releases_waiting_chunked_resume(t: ScriptedContext): def _script_pause_retract_releases_waiting_chunked_resume(t: ScriptedContext):
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
assert r.req.req_pool_idx is not None and r.kv_pages > 0 and r.lock_refs >= 1 assert r.req.kv.req_pool_idx is not None and r.kv_pages > 0 and r.lock_refs >= 1
t.pause_generation(mode="retract") t.pause_generation(mode="retract")
yield yield
assert r.req.req_pool_idx is None, ( assert r.req.kv.req_pool_idx is None, (
f"f38e69f87d: pause(retract) must release waiting " f"f38e69f87d: pause(retract) must release waiting "
f"chunked-resume row; got row_idx={r.req.req_pool_idx!r}" f"chunked-resume row; got row_idx={r.req.kv.req_pool_idx!r}"
) )
assert r.kv_pages == 0 assert r.kv_pages == 0
assert r.lock_refs == 0 assert r.lock_refs == 0
test/manual/scheduler/test_scripted_pp_abort.py
+1 -1
View File
@@ -14,7 +14,7 @@ def _drain_until_released(t: ScriptedContext, *handles: ScriptedReqHandle):
if all( if all(
h.kv_pages == 0 h.kv_pages == 0
and h.lock_refs == 0 and h.lock_refs == 0
and (h.req is None or h.req.req_pool_idx is None) and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles for h in handles
): ):
return return
test/registered/unit/disaggregation/test_disaggregation_wire.py
+1 -1
View File
@@ -102,7 +102,7 @@ class TestDisaggregationWire(unittest.TestCase):
def test_prebuilt_skips_unused_prompt_tensor(self): def test_prebuilt_skips_unused_prompt_tensor(self):
req = SimpleNamespace( req = SimpleNamespace(
req_pool_idx=0, kv=SimpleNamespace(req_pool_idx=0),
prefix_indices=[0, 1], prefix_indices=[0, 1],
extend_range=SimpleNamespace(length=3), extend_range=SimpleNamespace(length=3),
origin_input_ids=[0, 1, 2, 3, 4], origin_input_ids=[0, 1, 2, 3, 4],
test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py
+3 -2
View File
@@ -33,9 +33,10 @@ def _make_mock_req(
"""Create a mock Req with the KV cache state needed for testing.""" """Create a mock Req with the KV cache state needed for testing."""
req = MagicMock() req = MagicMock()
req.rid = rid req.rid = rid
req.req_pool_idx = req_pool_idx
req.kv = ReqKvInfo( req.kv = ReqKvInfo(
kv_committed_len=kv_committed_len, kv_allocated_len=kv_allocated_len req_pool_idx=req_pool_idx,
kv_committed_len=kv_committed_len,
kv_allocated_len=kv_allocated_len,
) )
req.prefix_indices = list(range(prefix_indices_len)) req.prefix_indices = list(range(prefix_indices_len))
req.effective_kv_committed_len = lambda: req.kv.kv_committed_len req.effective_kv_committed_len = lambda: req.kv.kv_committed_len
test/registered/unit/hardware_backend/mlx/test_attention_patching.py
+6 -6
View File
@@ -719,7 +719,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
full_token_ids=[11, 12, 13], full_token_ids=[11, 12, 13],
prefix_slot_ids=[2, 3], prefix_slot_ids=[2, 3],
new_slot_ids=[4], new_slot_ids=[4],
req_pool_idx=req.req_pool_idx, req_pool_idx=req.kv.req_pool_idx,
) )
runner.eval_pending(pending) runner.eval_pending(pending)
runner.prefill_finalize(pending) runner.prefill_finalize(pending)
@@ -775,7 +775,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
full_token_ids=token_ids, full_token_ids=token_ids,
prefix_slot_ids=[], prefix_slot_ids=[],
new_slot_ids=list(range(1, 71)), new_slot_ids=list(range(1, 71)),
req_pool_idx=req.req_pool_idx, req_pool_idx=req.kv.req_pool_idx,
req=req, req=req,
) )
runner.eval_pending(pending) runner.eval_pending(pending)
@@ -837,7 +837,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
full_token_ids=token_ids, full_token_ids=token_ids,
prefix_slot_ids=list(range(1, 65)), prefix_slot_ids=list(range(1, 65)),
new_slot_ids=list(range(65, 258)), new_slot_ids=list(range(65, 258)),
req_pool_idx=req.req_pool_idx, req_pool_idx=req.kv.req_pool_idx,
req=req, req=req,
) )
runner.eval_pending(pending) runner.eval_pending(pending)
@@ -917,7 +917,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
req = FakeRequest() req = FakeRequest()
req_indices = pool.alloc([req]) req_indices = pool.alloc([req])
auxiliary_state_idx = pool.get_auxiliary_state_indices(req.req_pool_idx) auxiliary_state_idx = pool.get_auxiliary_state_indices(req.kv.req_pool_idx)
pool.free(req) pool.free(req)
# Which free slot a fresh alloc gets is not semantically meaningful # Which free slot a fresh alloc gets is not semantically meaningful
@@ -925,7 +925,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
self.assertEqual(len(req_indices), 1) self.assertEqual(len(req_indices), 1)
self.assertIn(req_indices[0], range(1, pool.size + 1)) self.assertIn(req_indices[0], range(1, pool.size + 1))
self.assertIsNotNone(auxiliary_state_idx) self.assertIsNotNone(auxiliary_state_idx)
self.assertIsNone(req.req_pool_idx) self.assertIsNone(req.kv.req_pool_idx)
self.assertIsNotNone(req.mamba_pool_idx) self.assertIsNotNone(req.mamba_pool_idx)
self.assertIs(pool.mamba_allocator, pool.mamba_pool) self.assertIs(pool.mamba_allocator, pool.mamba_pool)
self.assertEqual(pool.auxiliary_state_pool.available_size(), 3) self.assertEqual(pool.auxiliary_state_pool.available_size(), 3)
@@ -1518,7 +1518,7 @@ if _HAS_MLX:
class FakeRequest: class FakeRequest:
def __init__(self): def __init__(self):
self.req_pool_idx = None self.kv = SimpleNamespace(req_pool_idx=None)
self.mamba_pool_idx = None self.mamba_pool_idx = None
self.inflight_middle_chunks = 0 self.inflight_middle_chunks = 0
test/registered/unit/hardware_backend/mlx/test_max_running_requests.py
+1 -1
View File
@@ -119,8 +119,8 @@ def _hybrid_stub_for_initialize(
def _fake_req(): def _fake_req():
return SimpleNamespace( return SimpleNamespace(
req_pool_idx=None,
inflight_middle_chunks=0, inflight_middle_chunks=0,
kv=SimpleNamespace(req_pool_idx=None),
mamba_pool_idx=None, mamba_pool_idx=None,
mamba_ping_pong_track_buffer=None, mamba_ping_pong_track_buffer=None,
) )
test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py
+1 -1
View File
@@ -169,7 +169,7 @@ class _FakeReq:
self.rid = rid self.rid = rid
self.prefix_indices = torch.empty(0, dtype=torch.long) self.prefix_indices = torch.empty(0, dtype=torch.long)
self.fill_ids = [0] self.fill_ids = [0]
self.req_pool_idx = req_pool_idx self.kv = SimpleNamespace(req_pool_idx=req_pool_idx)
# Mirrors Req's chunk-finality contract read by # Mirrors Req's chunk-finality contract read by
# MlxTpModelWorker._chunk_needs_logits: extend_range=None means # MlxTpModelWorker._chunk_needs_logits: extend_range=None means
# "not truncated" (final chunk / plain prefill). # "not truncated" (final chunk / plain prefill).
test/registered/unit/layers/test_minicpm_sparse_cache.py
+4 -4
View File
@@ -7,6 +7,7 @@ import torch
from sglang.srt.layers.attention.minicpm.cache import ( from sglang.srt.layers.attention.minicpm.cache import (
attach_compressed_cache, attach_compressed_cache,
) )
from sglang.srt.managers.schedule_batch import ReqKvInfo
from sglang.srt.managers.scheduler_components.invariant_checker import ( from sglang.srt.managers.scheduler_components.invariant_checker import (
SchedulerInvariantChecker, SchedulerInvariantChecker,
) )
@@ -66,8 +67,8 @@ def make_pool_and_req(capacity: int = 64):
enable_memory_saver=False, enable_memory_saver=False,
) )
req = SimpleNamespace( req = SimpleNamespace(
req_pool_idx=None,
inflight_middle_chunks=0, inflight_middle_chunks=0,
kv=SimpleNamespace(req_pool_idx=None),
) )
req_pool_idx = pool.alloc([req])[0] req_pool_idx = pool.alloc([req])[0]
return pool, req, req_pool_idx, allocator return pool, req, req_pool_idx, allocator
@@ -224,8 +225,7 @@ def test_streaming_session_release_frees_compressed_slots():
) )
) )
session.slots["session-a"] = SessionSlot( session.slots["session-a"] = SessionSlot(
req_pool_idx=req_pool_idx, kv=ReqKvInfo(req_pool_idx=req_pool_idx, kv_allocated_len=16),
kv=SimpleNamespace(kv_allocated_len=16, cache_protected_len=0),
) )
session.release_session("session-a") session.release_session("session-a")
@@ -293,7 +293,7 @@ def test_partial_failure_rolls_back_and_free_releases_every_slot():
assert len(cache.free_slots) == 0 assert len(cache.free_slots) == 0
pool.free(req) pool.free(req)
assert req.req_pool_idx is None assert req.kv.req_pool_idx is None
assert allocator.available_size() == 12 assert allocator.available_size() == 12
assert len(cache.free_slots) == 8 assert len(cache.free_slots) == 8
test/registered/unit/managers/test_hisparse_unit.py
+26 -25
View File
@@ -49,8 +49,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
output_ids=output_ids, output_ids=output_ids,
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, kv=SimpleNamespace(req_pool_idx=None, kv_allocated_len=0, kv_committed_len=0),
kv=SimpleNamespace(kv_allocated_len=0, kv_committed_len=0),
finished_reason=None, finished_reason=None,
hisparse_staging=False, hisparse_staging=False,
staging=False, staging=False,
@@ -190,11 +189,11 @@ class TestHiSparseUnit(unittest.TestCase):
"""Allocate a req_pool_idx for the request.""" """Allocate a req_pool_idx for the request."""
indices = self.req_to_token_pool.alloc([req]) indices = self.req_to_token_pool.alloc([req])
self.assertIsNotNone(indices, "Failed to allocate req pool slot") self.assertIsNotNone(indices, "Failed to allocate req pool slot")
return req.req_pool_idx return req.kv.req_pool_idx
def _free_req_slot(self, req): def _free_req_slot(self, req):
"""Free the req_pool_idx.""" """Free the req_pool_idx."""
if req.req_pool_idx is not None: if req.kv.req_pool_idx is not None:
self.req_to_token_pool.free(req) self.req_to_token_pool.free(req)
def _alloc_kv(self, req, fill_len, *, logical_only=False): def _alloc_kv(self, req, fill_len, *, logical_only=False):
@@ -216,7 +215,9 @@ class TestHiSparseUnit(unittest.TestCase):
extend_num_tokens=fill_len, extend_num_tokens=fill_len,
) )
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.kv.req_pool_idx, slice(0, len(kv_loc))), kv_loc
)
req.kv.kv_allocated_len = fill_len req.kv.kv_allocated_len = fill_len
req.kv.kv_committed_len = fill_len req.kv.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))
@@ -254,8 +255,8 @@ class TestHiSparseUnit(unittest.TestCase):
host_indices = host_pool.alloc(fill_len) host_indices = host_pool.alloc(fill_len)
self.assertIsNotNone(host_indices, "Host alloc failed") self.assertIsNotNone(host_indices, "Host alloc failed")
host_indices = host_indices.to(device="cuda") host_indices = host_indices.to(device="cuda")
self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len] = host_indices self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :fill_len] = host_indices
self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx] = fill_len self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = fill_len
for lid in range(LAYER_NUM): for lid in range(LAYER_NUM):
for i in range(fill_len): for i in range(fill_len):
host_pool.kv_buffer[lid][host_indices[i]] = self._kv_pattern(lid, i) host_pool.kv_buffer[lid][host_indices[i]] = self._kv_pattern(lid, i)
@@ -291,7 +292,7 @@ class TestHiSparseUnit(unittest.TestCase):
def _make_batch_tensors(self, reqs, fill_lens): def _make_batch_tensors(self, reqs, fill_lens):
"""Build (req_pool_indices [int64], seq_lens [int32]) on cuda.""" """Build (req_pool_indices [int64], seq_lens [int32]) on cuda."""
rpi = torch.tensor( rpi = torch.tensor(
[r.req_pool_idx for r in reqs], dtype=torch.int64, device="cuda" [r.kv.req_pool_idx for r in reqs], dtype=torch.int64, device="cuda"
) )
sls = torch.tensor(fill_lens, dtype=torch.int32, device="cuda") sls = torch.tensor(fill_lens, dtype=torch.int32, device="cuda")
return rpi, sls return rpi, sls
@@ -566,7 +567,7 @@ class TestHiSparseUnit(unittest.TestCase):
kv_loc = self._alloc_kv(req, fill_len) kv_loc = self._alloc_kv(req, fill_len)
self.coordinator.alloc_device_buffer(req) self.coordinator.alloc_device_buffer(req)
self.coordinator._skip_first_backup[req.req_pool_idx] = True self.coordinator._skip_first_backup[req.kv.req_pool_idx] = True
out_loc = self.allocator.alloc(1) out_loc = self.allocator.alloc(1)
self.assertIsNotNone(out_loc) self.assertIsNotNone(out_loc)
@@ -576,7 +577,7 @@ class TestHiSparseUnit(unittest.TestCase):
self.assertTrue(torch.all(stale_loc > 0), "Temporary mapping should exist") self.assertTrue(torch.all(stale_loc > 0), "Temporary mapping should exist")
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.kv.req_pool_idx, fill_len), out_loc)
req.kv.kv_allocated_len = seq_len req.kv.kv_allocated_len = seq_len
req.kv.kv_committed_len = seq_len req.kv.kv_committed_len = seq_len
@@ -584,10 +585,10 @@ class TestHiSparseUnit(unittest.TestCase):
seq_lens=torch.tensor([seq_len], dtype=torch.int64, device=device), seq_lens=torch.tensor([seq_len], dtype=torch.int64, device=device),
out_cache_loc=out_loc, out_cache_loc=out_loc,
req_pool_indices=torch.tensor( req_pool_indices=torch.tensor(
[req.req_pool_idx], dtype=torch.int64, device=device [req.kv.req_pool_idx], dtype=torch.int64, device=device
), ),
seq_lens_cpu=torch.tensor([seq_len], dtype=torch.int64), seq_lens_cpu=torch.tensor([seq_len], dtype=torch.int64),
req_pool_indices_cpu=torch.tensor([req.req_pool_idx], dtype=torch.int64), req_pool_indices_cpu=torch.tensor([req.kv.req_pool_idx], dtype=torch.int64),
) )
remapped_loc = self.allocator.full_to_hisparse_device_index_mapping[out_loc] remapped_loc = self.allocator.full_to_hisparse_device_index_mapping[out_loc]
@@ -626,7 +627,7 @@ class TestHiSparseUnit(unittest.TestCase):
ready = self.coordinator.collect_ready_reqs() ready = self.coordinator.collect_ready_reqs()
self.assertEqual(len(ready), 1) self.assertEqual(len(ready), 1)
self.assertFalse(req.hisparse_staging) self.assertFalse(req.hisparse_staging)
self.assertTrue(self.coordinator._skip_first_backup[req.req_pool_idx]) self.assertTrue(self.coordinator._skip_first_backup[req.kv.req_pool_idx])
tokens = self._build_topk_tokens(fill_len) tokens = self._build_topk_tokens(fill_len)
batch = tokens.unsqueeze(0) batch = tokens.unsqueeze(0)
@@ -662,11 +663,11 @@ class TestHiSparseUnit(unittest.TestCase):
ready = self.coordinator.collect_ready_reqs() ready = self.coordinator.collect_ready_reqs()
self.assertEqual(ready, [req]) self.assertEqual(ready, [req])
host_row = self.coordinator.req_to_host_pool[req.req_pool_idx, :rounded_len] host_row = self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :rounded_len]
self.assertTrue(torch.all(host_row >= 0)) self.assertTrue(torch.all(host_row >= 0))
self.assertEqual(torch.unique(host_row).numel(), rounded_len) self.assertEqual(torch.unique(host_row).numel(), rounded_len)
self.assertEqual( self.assertEqual(
int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]),
rounded_len, rounded_len,
) )
@@ -674,7 +675,7 @@ class TestHiSparseUnit(unittest.TestCase):
next_host_index = self.coordinator.mem_pool_host.alloc_paged_token_slots( next_host_index = self.coordinator.mem_pool_host.alloc_paged_token_slots(
self.coordinator.req_to_host_pool, self.coordinator.req_to_host_pool,
self.coordinator.req_to_host_pool_allocated_len, self.coordinator.req_to_host_pool_allocated_len,
req.req_pool_idx, req.kv.req_pool_idx,
fill_len, fill_len,
1, 1,
) )
@@ -691,8 +692,8 @@ class TestHiSparseUnit(unittest.TestCase):
expected_total = rounded_len + expected_new_pages * self.page_size expected_total = rounded_len + expected_new_pages * self.page_size
allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices( allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices(
self.coordinator.req_to_host_pool, self.coordinator.req_to_host_pool,
req.req_pool_idx, req.kv.req_pool_idx,
int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]),
) )
self.assertEqual(allocated_host_indices.numel(), expected_total) self.assertEqual(allocated_host_indices.numel(), expected_total)
@@ -714,9 +715,9 @@ class TestHiSparseUnit(unittest.TestCase):
self.coordinator.admit_request_direct(req) self.coordinator.admit_request_direct(req)
self.assertFalse(req.staging) self.assertFalse(req.staging)
self.assertTrue(self.coordinator._skip_first_backup[req.req_pool_idx]) self.assertTrue(self.coordinator._skip_first_backup[req.kv.req_pool_idx])
buf_tokens = self.coordinator.req_device_buffer_tokens[ buf_tokens = self.coordinator.req_device_buffer_tokens[
:, req.req_pool_idx, :DEVICE_BUFFER_SIZE :, req.kv.req_pool_idx, :DEVICE_BUFFER_SIZE
] ]
self.assertTrue(torch.all(buf_tokens == -1)) self.assertTrue(torch.all(buf_tokens == -1))
@@ -765,7 +766,7 @@ class TestHiSparseUnit(unittest.TestCase):
self.assertTrue( self.assertTrue(
torch.equal( torch.equal(
host_indices, host_indices,
self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len], self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :fill_len],
) )
) )
self.assertEqual(req.kv.kv_allocated_len, fill_len) self.assertEqual(req.kv.kv_allocated_len, fill_len)
@@ -774,18 +775,18 @@ class TestHiSparseUnit(unittest.TestCase):
rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size
self.assertEqual( self.assertEqual(
int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]),
rounded_len, rounded_len,
) )
allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices( allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices(
self.coordinator.req_to_host_pool, self.coordinator.req_to_host_pool,
req.req_pool_idx, req.kv.req_pool_idx,
int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]),
) )
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.kv_allocated_len req.kv.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
+19 -20
View File
@@ -5,6 +5,7 @@ from types import SimpleNamespace
import torch import torch
from sglang.srt.managers.schedule_batch import ReqKvInfo
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase
@@ -43,23 +44,15 @@ def _make_checker(page_size=_PAGE_SIZE, row_width=4096, num_reqs=8, free_pages=N
return _FakeChecker(), rtt, tc, alloc return _FakeChecker(), rtt, tc, alloc
class _FakeReq: class _FakeOwner:
def __init__(self, rid, rpi, committed, allocated): # A req or a session slot; the checker reads only `kv` (and `rid` for reqs).
def __init__(self, rpi, committed, allocated, rid=None):
self.rid = rid self.rid = rid
self.req_pool_idx = rpi self.kv = ReqKvInfo(
self.kv = SimpleNamespace( req_pool_idx=rpi,
kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0 kv_committed_len=committed,
kv_allocated_len=allocated,
) )
self.is_holding_kv = True
class _FakeSlot:
def __init__(self, rpi, committed, allocated):
self.req_pool_idx = rpi
self.kv = SimpleNamespace(
kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0
)
self.is_holding_kv = True
class TestKVPageInvariants(CustomTestCase): class TestKVPageInvariants(CustomTestCase):
@@ -70,21 +63,23 @@ class TestKVPageInvariants(CustomTestCase):
rtt[0, :256] = torch.arange(_PAGE_SIZE) # req 0 owns page 0 rtt[0, :256] = torch.arange(_PAGE_SIZE) # req 0 owns page 0
rtt[1, :256] = torch.arange(_PAGE_SIZE, 2 * _PAGE_SIZE) # req 1 owns page 1 rtt[1, :256] = torch.arange(_PAGE_SIZE, 2 * _PAGE_SIZE) # req 1 owns page 1
chk.get_last_batch = lambda: SimpleNamespace( chk.get_last_batch = lambda: SimpleNamespace(
reqs=[_FakeReq("a", 0, 256, 256), _FakeReq("b", 1, 200, 256)] reqs=[_FakeOwner(0, 256, 256, rid="a"), _FakeOwner(1, 200, 256, rid="b")]
) )
chk._check_kv_page_invariants() chk._check_kv_page_invariants()
self.assertEqual(chk.count_memory_leak_warnings, 0) self.assertEqual(chk.count_memory_leak_warnings, 0)
def test_committed_gt_allocated_raises(self): def test_committed_gt_allocated_raises(self):
chk, rtt, tc, alloc = _make_checker() chk, rtt, tc, alloc = _make_checker()
chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 145, 144)]) chk.get_last_batch = lambda: SimpleNamespace(
reqs=[_FakeOwner(0, 145, 144, rid="a")]
)
with self.assertRaises(AssertionError): with self.assertRaises(AssertionError):
chk._check_kv_page_invariants() chk._check_kv_page_invariants()
def test_slot_committed_gt_allocated_raises(self): def test_slot_committed_gt_allocated_raises(self):
chk, rtt, tc, alloc = _make_checker() chk, rtt, tc, alloc = _make_checker()
chk.get_last_batch = lambda: None chk.get_last_batch = lambda: None
tc.slots = {"s1": _FakeSlot(0, 145, 144)} tc.slots = {"s1": _FakeOwner(0, 145, 144)}
with self.assertRaises(AssertionError): with self.assertRaises(AssertionError):
chk._check_kv_page_invariants() chk._check_kv_page_invariants()
@@ -94,14 +89,18 @@ class TestKVPageInvariants(CustomTestCase):
rtt[0, :3] = torch.tensor( rtt[0, :3] = torch.tensor(
[5 * _PAGE_SIZE, 5 * _PAGE_SIZE + 1, 5 * _PAGE_SIZE + 2] [5 * _PAGE_SIZE, 5 * _PAGE_SIZE + 1, 5 * _PAGE_SIZE + 2]
) )
chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 3, 3)]) chk.get_last_batch = lambda: SimpleNamespace(
reqs=[_FakeOwner(0, 3, 3, rid="a")]
)
with self.assertRaises(ValueError): with self.assertRaises(ValueError):
chk._check_kv_page_invariants() chk._check_kv_page_invariants()
def test_free_pool_duplicate_raises(self): def test_free_pool_duplicate_raises(self):
chk, rtt, tc, alloc = _make_checker(free_pages=torch.tensor([3, 3, 4])) chk, rtt, tc, alloc = _make_checker(free_pages=torch.tensor([3, 3, 4]))
rtt[0, :1] = torch.tensor([10 * _PAGE_SIZE]) # owner page 10, not in free rtt[0, :1] = torch.tensor([10 * _PAGE_SIZE]) # owner page 10, not in free
chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 1, 1)]) chk.get_last_batch = lambda: SimpleNamespace(
reqs=[_FakeOwner(0, 1, 1, rid="a")]
)
with self.assertRaises(ValueError): with self.assertRaises(ValueError):
chk._check_kv_page_invariants() chk._check_kv_page_invariants()
test/registered/unit/managers/test_priority_scheduling_disaggregation.py
+1 -2
View File
@@ -93,11 +93,10 @@ class TestDecodePreallocQueuePriority(unittest.TestCase):
priority=priority, priority=priority,
origin_input_ids=[1, 2, 3], origin_input_ids=[1, 2, 3],
output_ids=[], output_ids=[],
req_pool_idx=int(priority) % 8,
finished_reason=FINISH_ABORT("failed") if failed else None, finished_reason=FINISH_ABORT("failed") if failed else None,
return_logprob=False, return_logprob=False,
sampling_params=SimpleNamespace(max_new_tokens=8), sampling_params=SimpleNamespace(max_new_tokens=8),
kv=SimpleNamespace(cache_protected_len=0), kv=SimpleNamespace(req_pool_idx=int(priority) % 8, cache_protected_len=0),
time_stats=MagicMock(), time_stats=MagicMock(),
) )
return SimpleNamespace( return SimpleNamespace(
test/registered/unit/managers/test_schedule_batch_req_pool_indices.py
+1 -1
View File
@@ -18,7 +18,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _make_req(req_pool_idx, origin_input_ids, output_ids): def _make_req(req_pool_idx, origin_input_ids, output_ids):
return types.SimpleNamespace( return types.SimpleNamespace(
req_pool_idx=req_pool_idx, kv=types.SimpleNamespace(req_pool_idx=req_pool_idx),
origin_input_ids=origin_input_ids, origin_input_ids=origin_input_ids,
output_ids=output_ids, output_ids=output_ids,
return_logprob=False, return_logprob=False,
test/registered/unit/managers/test_scheduler_chunked_req_gate.py
+1 -2
View File
@@ -34,11 +34,10 @@ def _make_req(
req.output_ids = array("q") req.output_ids = array("q")
req.full_untruncated_fill_ids = array("q", fill_ids) req.full_untruncated_fill_ids = array("q", fill_ids)
req.prefix_indices = prefix_indices req.prefix_indices = prefix_indices
req.req_pool_idx = req_pool_idx
req.extend_range = Range(fill_len - extend_input_len, fill_len) req.extend_range = Range(fill_len - extend_input_len, fill_len)
req.inflight_middle_chunks = 0 req.inflight_middle_chunks = 0
req.host_hit_length = 0 req.host_hit_length = 0
req.kv = ReqKvInfo() req.kv = ReqKvInfo(req_pool_idx=req_pool_idx)
req.skip_radix_cache_insert = False req.skip_radix_cache_insert = False
req.last_node = None req.last_node = None
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+1 -1
View File
@@ -75,13 +75,13 @@ class MockReq:
"q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids "q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids
) )
self.output_ids = array("q", [fill_ids[-1]] if len(fill_ids) > 1 else []) self.output_ids = array("q", [fill_ids[-1]] if len(fill_ids) > 1 else [])
self.req_pool_idx = req_pool_idx
self.last_node = last_node self.last_node = last_node
self.extra_key = None self.extra_key = None
self.cache_salt = None self.cache_salt = None
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 = SimpleNamespace( self.kv = SimpleNamespace(
req_pool_idx=req_pool_idx,
kv_committed_len=len(fill_ids), kv_committed_len=len(fill_ids),
kv_allocated_len=len(fill_ids), kv_allocated_len=len(fill_ids),
cache_protected_len=cache_protected_len, cache_protected_len=cache_protected_len,
test/registered/unit/mem_cache/test_decode_retraction_backup.py
+5 -3
View File
@@ -129,12 +129,14 @@ class TestDecodeRetractionBackup(unittest.TestCase):
) )
def _admit_req(self, env, num_tokens: int): def _admit_req(self, env, num_tokens: int):
req = SimpleNamespace(rid="request", req_pool_idx=None, seqlen=num_tokens + 1) req = SimpleNamespace(
rid="request", kv=SimpleNamespace(req_pool_idx=None), seqlen=num_tokens + 1
)
self.assertIsNotNone(env.req_to_token_pool.alloc([req])) self.assertIsNotNone(env.req_to_token_pool.alloc([req]))
source_indices = env.allocator.alloc(num_tokens) source_indices = env.allocator.alloc(num_tokens)
self.assertIsNotNone(source_indices) self.assertIsNotNone(source_indices)
env.req_to_token_pool.write( env.req_to_token_pool.write(
(req.req_pool_idx, slice(0, num_tokens)), source_indices (req.kv.req_pool_idx, slice(0, num_tokens)), source_indices
) )
return req, source_indices return req, source_indices
@@ -200,7 +202,7 @@ class TestDecodeRetractionBackup(unittest.TestCase):
self.assertIsNotNone(destination_indices) self.assertIsNotNone(destination_indices)
self.assertFalse(torch.equal(source_indices, destination_indices)) self.assertFalse(torch.equal(source_indices, destination_indices))
req_to_token_pool.write( req_to_token_pool.write(
(req.req_pool_idx, slice(0, self.num_tokens)), destination_indices (req.kv.req_pool_idx, slice(0, self.num_tokens)), destination_indices
) )
cache.retraction_restore(req, backup) cache.retraction_restore(req, backup)
test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py
+5 -5
View File
@@ -61,10 +61,10 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False):
return SimpleNamespace( return SimpleNamespace(
rid=rid, rid=rid,
prefix_indices=torch.tensor(prefix, dtype=torch.int32), prefix_indices=torch.tensor(prefix, dtype=torch.int32),
req_pool_idx=req_pool_idx,
dllm_incomplete_ids=array("q", range(block_size)) if reuse else array("q"), dllm_incomplete_ids=array("q", range(block_size)) if reuse else array("q"),
inflight_middle_chunks=1 if req_pool_idx is not None else 0, inflight_middle_chunks=1 if req_pool_idx is not None else 0,
kv=SimpleNamespace( kv=SimpleNamespace(
req_pool_idx=req_pool_idx,
kv_committed_len=len(prefix) if req_pool_idx is not None else 0, kv_committed_len=len(prefix) if req_pool_idx is not None else 0,
kv_allocated_len=( kv_allocated_len=(
len(prefix) + block_size if req_pool_idx is not None else 0 len(prefix) + block_size if req_pool_idx is not None else 0
@@ -75,8 +75,8 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False):
def _remove_allocated_req_slots(pool, *reqs): def _remove_allocated_req_slots(pool, *reqs):
for req in reqs: for req in reqs:
if req.req_pool_idx in pool.free_slots: if req.kv.req_pool_idx in pool.free_slots:
pool.free_slots.remove(req.req_pool_idx) pool.free_slots.remove(req.kv.req_pool_idx)
def _make_batch(pool, allocator, reqs, extend_lens): def _make_batch(pool, allocator, reqs, extend_lens):
@@ -105,8 +105,8 @@ def _make_batch(pool, allocator, reqs, extend_lens):
def _seed_retained_block(pool, req, values): def _seed_retained_block(pool, req, values):
prefix_len = len(req.prefix_indices) prefix_len = len(req.prefix_indices)
pool.req_to_token[req.req_pool_idx, :prefix_len] = req.prefix_indices pool.req_to_token[req.kv.req_pool_idx, :prefix_len] = req.prefix_indices
pool.req_to_token[req.req_pool_idx, prefix_len : prefix_len + len(values)] = ( pool.req_to_token[req.kv.req_pool_idx, prefix_len : prefix_len + len(values)] = (
torch.tensor(values, dtype=torch.int32) torch.tensor(values, dtype=torch.int32)
) )
test/registered/unit/mem_cache/test_hisparse_allocator.py
+2 -2
View File
@@ -106,7 +106,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
def alloc(self, reqs): def alloc(self, reqs):
for item in reqs: for item in reqs:
item.req_pool_idx = 0 item.kv.req_pool_idx = 0
return torch.tensor([0], dtype=torch.int64) return torch.tensor([0], dtype=torch.int64)
def write(self, indices, values): def write(self, indices, values):
@@ -160,7 +160,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
regular_host_alloc.assert_called_once_with( regular_host_alloc.assert_called_once_with(
coordinator.req_to_host_pool, coordinator.req_to_host_pool,
coordinator.req_to_host_pool_allocated_len, coordinator.req_to_host_pool_allocated_len,
req.req_pool_idx, req.kv.req_pool_idx,
0, 0,
len(host_indices), len(host_indices),
) )
test/registered/unit/mem_cache/test_paged_free_segment.py
+1 -1
View File
@@ -131,7 +131,7 @@ class TestFreeSegment(unittest.TestCase):
token_to_kv_pool_allocator=alloc, token_to_kv_pool_allocator=alloc,
req_to_token_pool=SimpleNamespace(req_to_token=row.unsqueeze(0)), req_to_token_pool=SimpleNamespace(req_to_token=row.unsqueeze(0)),
) )
req = SimpleNamespace(req_pool_idx=0) req = SimpleNamespace(kv=SimpleNamespace(req_pool_idx=0))
before = len(alloc.free_pages) before = len(alloc.free_pages)
alloc.free_group_begin() alloc.free_group_begin()
test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py
+4 -3
View File
@@ -21,11 +21,12 @@ class _FakeAllocator:
class _FakeReq: class _FakeReq:
req_pool_idx = 0
def __init__(self): def __init__(self):
self.kv = SimpleNamespace( self.kv = SimpleNamespace(
swa_evicted_seqlen=6, swa_evict_floor=3, cache_protected_len=0 req_pool_idx=0,
swa_evicted_seqlen=6,
swa_evict_floor=3,
cache_protected_len=0,
) )
def pop_committed_kv_cache(self): def pop_committed_kv_cache(self):
test/registered/unit/mem_cache/test_radix_cache_unit.py
+1 -2
View File
@@ -518,8 +518,7 @@ class TestRadixCache(unittest.TestCase):
) )
cache.req_to_token_pool = ReqToTokenPool(request_indices.clone()) cache.req_to_token_pool = ReqToTokenPool(request_indices.clone())
req = unittest.mock.Mock( req = unittest.mock.Mock(
req_pool_idx=0, kv=SimpleNamespace(req_pool_idx=0, cache_protected_len=0),
kv=SimpleNamespace(cache_protected_len=0),
extra_key=None, extra_key=None,
cache_salt=None, cache_salt=None,
priority=0, priority=0,
test/registered/unit/mem_cache/test_retraction_mamba_backup.py
+2 -2
View File
@@ -2,7 +2,7 @@ import unittest
import torch import torch
from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_batch import Req, ReqKvInfo
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
@@ -39,7 +39,7 @@ class _Allocator:
def _req_and_pool(): def _req_and_pool():
req = object.__new__(Req) req = object.__new__(Req)
req.req_pool_idx = 0 req.kv = ReqKvInfo(req_pool_idx=0)
req.origin_input_ids = [1, 2] req.origin_input_ids = [1, 2]
req.output_ids = [3] req.output_ids = [3]
req.mamba_pool_idx = torch.tensor(1) req.mamba_pool_idx = torch.tensor(1)
test/registered/unit/mem_cache/test_streaming_session_unit.py
+14 -14
View File
@@ -2,7 +2,7 @@ from types import SimpleNamespace
import torch import torch
from sglang.srt.managers.schedule_batch import FINISH_ABORT from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo
from sglang.srt.mem_cache.base_prefix_cache import MatchResult from sglang.srt.mem_cache.base_prefix_cache import MatchResult
from sglang.srt.session.streaming_session import SessionSlot, StreamingSession from sglang.srt.session.streaming_session import SessionSlot, StreamingSession
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
@@ -24,8 +24,8 @@ class _FakeReqToTokenPool:
self.free_slots = [] self.free_slots = []
def free(self, req): def free(self, req):
self.free_slots.append(req.req_pool_idx) self.free_slots.append(req.kv.req_pool_idx)
req.req_pool_idx = None req.kv.req_pool_idx = None
class _FakeInnerCache: class _FakeInnerCache:
@@ -67,8 +67,8 @@ class _FakeReq:
abort_req=lambda: None, abort_req=lambda: None,
_inflight=False, _inflight=False,
) )
self.req_pool_idx = req_pool_idx self.kv = ReqKvInfo(
self.kv = SimpleNamespace( req_pool_idx=req_pool_idx,
kv_committed_len=committed, kv_committed_len=committed,
kv_allocated_len=allocated, kv_allocated_len=allocated,
swa_evicted_seqlen=0, swa_evicted_seqlen=0,
@@ -112,8 +112,8 @@ def test_preabort_detaches_session_and_preserves_slot():
) )
tree_cache = StreamingSession(inner) tree_cache = StreamingSession(inner)
tree_cache.slots["session-a"] = SessionSlot( tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0, kv=ReqKvInfo(
kv=SimpleNamespace( req_pool_idx=0,
kv_committed_len=48, kv_committed_len=48,
kv_allocated_len=48, kv_allocated_len=48,
swa_evicted_seqlen=0, swa_evicted_seqlen=0,
@@ -135,7 +135,7 @@ def test_preabort_detaches_session_and_preserves_slot():
assert req.session is None assert req.session is None
# Slot untouched. # Slot untouched.
slot = tree_cache.slots["session-a"] slot = tree_cache.slots["session-a"]
assert slot.req_pool_idx == 0 assert slot.kv.req_pool_idx == 0
assert slot.kv.kv_committed_len == 48 assert slot.kv.kv_committed_len == 48
assert slot.kv.kv_allocated_len == 48 assert slot.kv.kv_allocated_len == 48
assert len(result.device_indices) == 0 assert len(result.device_indices) == 0
@@ -160,7 +160,7 @@ def test_first_mid_abort_nukes_ephemeral_slot():
# Slot must NOT be created. # Slot must NOT be created.
assert "session-a" not in tree_cache.slots assert "session-a" not in tree_cache.slots
# Transient pool slot freed. # Transient pool slot freed.
assert req.req_pool_idx is None assert req.kv.req_pool_idx is None
assert req_to_token_pool.free_slots == [0] assert req_to_token_pool.free_slots == [0]
assert len(allocator.freed) == 1 assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(20)) assert allocator.freed[0].tolist() == list(range(20))
@@ -179,8 +179,8 @@ def test_nth_mid_abort_nukes_session_slot():
# Session already has a slot from a previous turn. # Session already has a slot from a previous turn.
tree_cache.slots["session-a"] = SessionSlot( tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0, kv=ReqKvInfo(
kv=SimpleNamespace( req_pool_idx=0,
kv_committed_len=50, kv_committed_len=50,
kv_allocated_len=50, kv_allocated_len=50,
swa_evicted_seqlen=0, swa_evicted_seqlen=0,
@@ -202,7 +202,7 @@ def test_nth_mid_abort_nukes_session_slot():
assert allocator.freed[0].tolist() == list(range(65)) assert allocator.freed[0].tolist() == list(range(65))
# Pool slot returned. # Pool slot returned.
assert req_to_token_pool.free_slots == [0] assert req_to_token_pool.free_slots == [0]
assert req.req_pool_idx is None assert req.kv.req_pool_idx is None
def test_release_session_threads_mamba_skip_ids(): def test_release_session_threads_mamba_skip_ids():
@@ -220,8 +220,8 @@ def test_release_session_threads_mamba_skip_ids():
lock_node = SimpleNamespace(id=42) lock_node = SimpleNamespace(id=42)
tree_cache.slots["session-a"] = SessionSlot( tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0, kv=ReqKvInfo(
kv=SimpleNamespace( req_pool_idx=0,
kv_committed_len=50, kv_committed_len=50,
kv_allocated_len=50, kv_allocated_len=50,
swa_evicted_seqlen=0, swa_evicted_seqlen=0,
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
+3 -3
View File
@@ -103,11 +103,11 @@ def _build_swa_tree(page_size, sliding_window_size, kv_size=1024, kv_size_swa=51
def _make_req(req_pool_idx, token_ids, cache_protected_len, tree): def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
"""Mock Req with fields needed by _evict_swa and cache_finished_req.""" """Mock Req with fields needed by _evict_swa and cache_finished_req."""
req = SimpleNamespace( req = SimpleNamespace(
req_pool_idx=req_pool_idx,
is_holding_kv=True,
origin_input_ids=token_ids, origin_input_ids=token_ids,
output_ids=[], output_ids=[],
kv=ReqKvInfo(cache_protected_len=cache_protected_len), kv=ReqKvInfo(
req_pool_idx=req_pool_idx, cache_protected_len=cache_protected_len
),
extra_key=None, extra_key=None,
cache_salt=None, cache_salt=None,
last_node=tree.root_node, last_node=tree.root_node,
test/registered/unit/mem_cache/test_swa_unittest.py
+5 -5
View File
@@ -669,13 +669,13 @@ class TestSWA(unittest.TestCase):
# Case 1: is_insert=True should pass bigram key and use cache_protected_len. # Case 1: is_insert=True should pass bigram key and use cache_protected_len.
req = _DummyReq() req = _DummyReq()
req.req_pool_idx = 0 req.kv.req_pool_idx = 0
req.origin_input_ids = array("q", [1, 2, 3, 4, 5, 6]) req.origin_input_ids = array("q", [1, 2, 3, 4, 5, 6])
req.output_ids = array("q") req.output_ids = array("q")
req._kv_committed_len = len(req.origin_input_ids) req._kv_committed_len = len(req.origin_input_ids)
kv_indices = allocator.alloc(req._kv_committed_len) kv_indices = allocator.alloc(req._kv_committed_len)
req_to_token_pool.write( req_to_token_pool.write(
(req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices (req.kv.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices
) )
req.extra_key = None req.extra_key = None
req.cache_salt = None req.cache_salt = None
@@ -707,13 +707,13 @@ class TestSWA(unittest.TestCase):
# Case 2: is_insert=False should free [cache_protected_len:page_aligned_len] # Case 2: is_insert=False should free [cache_protected_len:page_aligned_len]
# even when len(prefix_indices) is intentionally larger. # even when len(prefix_indices) is intentionally larger.
req2 = _DummyReq() req2 = _DummyReq()
req2.req_pool_idx = 1 req2.kv.req_pool_idx = 1
req2.origin_input_ids = array("q", [11, 12, 13, 14, 15, 16]) req2.origin_input_ids = array("q", [11, 12, 13, 14, 15, 16])
req2.output_ids = array("q") req2.output_ids = array("q")
req2._kv_committed_len = len(req2.origin_input_ids) req2._kv_committed_len = len(req2.origin_input_ids)
kv_indices2 = allocator.alloc(req2._kv_committed_len) kv_indices2 = allocator.alloc(req2._kv_committed_len)
req_to_token_pool.write( req_to_token_pool.write(
(req2.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2 (req2.kv.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2
) )
req2.extra_key = None req2.extra_key = None
req2.cache_salt = None req2.cache_salt = None
@@ -911,7 +911,7 @@ class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase):
token_ids = array("q", range(1, num_tokens + 1)) token_ids = array("q", range(1, num_tokens + 1))
req = _DummyReq() req = _DummyReq()
req.req_pool_idx = 0 req.kv.req_pool_idx = 0
req.origin_input_ids = token_ids req.origin_input_ids = token_ids
req.output_ids = array("q") req.output_ids = array("q")
req.get_fill_ids = lambda: token_ids req.get_fill_ids = lambda: token_ids
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+1 -1
View File
@@ -646,7 +646,7 @@ def bench_cache_finished(
req.kv.kv_committed_len = len(seq) req.kv.kv_committed_len = len(seq)
if hasattr(lr, "swa_uuid_for_lock"): if hasattr(lr, "swa_uuid_for_lock"):
req.swa_uuid_for_lock = lr.swa_uuid_for_lock req.swa_uuid_for_lock = lr.swa_uuid_for_lock
env.rtp.req_to_token[req.req_pool_idx, : len(kv_indices)] = kv_indices env.rtp.req_to_token[req.kv.req_pool_idx, : len(kv_indices)] = kv_indices
req_items.append(req) req_items.append(req)
if not req_items: if not req_items:
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+16 -12
View File
@@ -1235,7 +1235,7 @@ class UnifiedRadixCacheSuite:
req.output_ids = array("q", output_ids) req.output_ids = array("q", output_ids)
kv_len = len(input_ids) + len(output_ids) kv_len = len(input_ids) + len(output_ids)
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.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
@@ -1275,7 +1275,7 @@ class UnifiedRadixCacheSuite:
) )
kv_len = req.extend_range.end kv_len = req.extend_range.end
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.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.kv.kv_allocated_len = kv_len req.kv.kv_allocated_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
@@ -1297,7 +1297,7 @@ class UnifiedRadixCacheSuite:
start_p = ((start_p + ps - 1) // ps) * ps start_p = ((start_p + ps - 1) // ps) * ps
if start_p < end_p: if start_p < end_p:
allocator.free( allocator.free(
req_to_token_pool.req_to_token[req.req_pool_idx][start_p:end_p] req_to_token_pool.req_to_token[req.kv.req_pool_idx][start_p:end_p]
) )
prompt_aligned = (len(prompt_ids) // ps) * ps prompt_aligned = (len(prompt_ids) // ps) * ps
@@ -1320,7 +1320,7 @@ class UnifiedRadixCacheSuite:
req.output_ids = array("q") req.output_ids = array("q")
kv_len = len(tokens) kv_len = len(tokens)
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.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
@@ -1355,7 +1355,7 @@ class UnifiedRadixCacheSuite:
) )
kv_len = len(tokens) kv_len = len(tokens)
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.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
@@ -1392,7 +1392,9 @@ class UnifiedRadixCacheSuite:
req.full_untruncated_fill_ids = array("q", tokens) req.full_untruncated_fill_ids = array("q", tokens)
req.set_extend_range(0, len(req.full_untruncated_fill_ids)) req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc(allocator, len(tokens)) kv_indices = self._alloc(allocator, len(tokens))
req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices) req_to_token_pool.write(
(req.kv.req_pool_idx, slice(0, len(tokens))), kv_indices
)
req.kv.kv_committed_len = len(tokens) req.kv.kv_committed_len = len(tokens)
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
@@ -1494,7 +1496,7 @@ class UnifiedRadixCacheSuite:
req.output_ids = array("q") req.output_ids = array("q")
kv_len = len(input_ids) kv_len = len(input_ids)
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.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
@@ -1619,7 +1621,7 @@ class UnifiedRadixCacheSuite:
req.set_extend_range(0, len(req.full_untruncated_fill_ids)) req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_len = len(tokens) kv_len = len(tokens)
fresh_value = self._alloc(allocator, kv_len) fresh_value = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), fresh_value) req_to_token_pool.write((req.kv.req_pool_idx, slice(0, kv_len)), fresh_value)
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
@@ -2206,7 +2208,7 @@ class UnifiedRadixCacheSuite:
req.full_untruncated_fill_ids = array("q", tokens) req.full_untruncated_fill_ids = array("q", tokens)
req.set_extend_range(0, len(req.full_untruncated_fill_ids)) req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc(allocator, pre_len) kv_indices = self._alloc(allocator, pre_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) req_to_token_pool.write((req.kv.req_pool_idx, slice(0, pre_len)), kv_indices)
req.kv.kv_committed_len = pre_len req.kv.kv_committed_len = pre_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
@@ -2295,7 +2297,7 @@ class UnifiedRadixCacheSuite:
req.full_untruncated_fill_ids = array("q", tokens) req.full_untruncated_fill_ids = array("q", tokens)
req.set_extend_range(0, len(req.full_untruncated_fill_ids)) req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc(allocator, pre_len) kv_indices = self._alloc(allocator, pre_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices) req_to_token_pool.write((req.kv.req_pool_idx, slice(0, pre_len)), kv_indices)
req.kv.kv_committed_len = pre_len req.kv.kv_committed_len = pre_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
@@ -6790,7 +6792,9 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
req = self._make_req(req_to_token_pool, tokens) req = self._make_req(req_to_token_pool, tokens)
kv_indices = allocator.alloc(len(tokens)) kv_indices = allocator.alloc(len(tokens))
self.assertIsNotNone(kv_indices) self.assertIsNotNone(kv_indices)
req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices) req_to_token_pool.write(
(req.kv.req_pool_idx, slice(0, len(tokens))), kv_indices
)
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
cache.cache_finished_req( cache.cache_finished_req(
@@ -8104,7 +8108,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase):
req.full_untruncated_fill_ids = array("q", tokens) req.full_untruncated_fill_ids = array("q", tokens)
req.set_extend_range(0, len(req.full_untruncated_fill_ids)) req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc_paged(allocator, seq_len) kv_indices = self._alloc_paged(allocator, seq_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, seq_len)), kv_indices) req_to_token_pool.write((req.kv.req_pool_idx, slice(0, seq_len)), kv_indices)
req.kv.kv_committed_len = seq_len req.kv.kv_committed_len = seq_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
test/registered/xpu/test_lmcache_radix_cache.py
+10 -3
View File
@@ -66,13 +66,16 @@ def _make_req(rid, req_pool_idx, token_ids, tree):
SimpleNamespace pattern in test_swa_eviction_boundary.py).""" SimpleNamespace pattern in test_swa_eviction_boundary.py)."""
req = SimpleNamespace( req = SimpleNamespace(
rid=rid, rid=rid,
req_pool_idx=req_pool_idx,
origin_input_ids=token_ids, origin_input_ids=token_ids,
output_ids=[], output_ids=[],
extra_key=None, extra_key=None,
cache_salt=None, cache_salt=None,
last_node=tree.root_node, last_node=tree.root_node,
kv=SimpleNamespace(cache_protected_len=0, kv_committed_len=len(token_ids)), kv=SimpleNamespace(
req_pool_idx=req_pool_idx,
cache_protected_len=0,
kv_committed_len=len(token_ids),
),
priority=0, priority=0,
kv_committed_freed=False, kv_committed_freed=False,
) )
@@ -162,7 +165,11 @@ class TestLMCRadixCacheXPU(unittest.TestCase):
# commit it as a finished request (inserts into radix + stores to # commit it as a finished request (inserts into radix + stores to
# LMCache on tree.store_stream). # LMCache on tree.store_stream).
req_pool_idx = req_to_token_pool.alloc( req_pool_idx = req_to_token_pool.alloc(
[SimpleNamespace(req_pool_idx=None, inflight_middle_chunks=0)] [
SimpleNamespace(
inflight_middle_chunks=0, kv=SimpleNamespace(req_pool_idx=None)
)
]
)[0] )[0]
kv_slots = allocator.alloc(self.INPUT_LEN) kv_slots = allocator.alloc(self.INPUT_LEN)
self.assertIsNotNone(kv_slots) self.assertIsNotNone(kv_slots)