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:
# Prefill-terminated (max_new_tokens == 1): no spawn, but the relay
# 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
self._spawn_member_rows(group, req)
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(
self.req_to_token_pool.req_to_token,
member_rows,
leader.req_pool_idx,
leader.kv.req_pool_idx,
group.prompt_len,
)
group.member_rows = member_rows
group.member_rows_cpu = torch.tensor(rows, dtype=torch.int64)
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])
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:]
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
# prefix); req_to_token rows are absolute.
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]]:
# Indices of reqs that already have a req_pool_idx and will reuse
# 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 (
len(reusing) <= 1
), "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:]
offset = 0
for r in reqs:
if r.req_pool_idx is None:
r.req_pool_idx = select_index[offset]
self.req_generation[r.req_pool_idx] += 1
if r.kv.req_pool_idx is None:
r.kv.req_pool_idx = select_index[offset]
self.req_generation[r.kv.req_pool_idx] += 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):
assert req.req_pool_idx is not None, "request must have req_pool_idx"
self.free_slots.append(req.req_pool_idx)
req.req_pool_idx = None
assert req.kv.req_pool_idx is not None, "request must have req_pool_idx"
self.free_slots.append(req.kv.req_pool_idx)
req.kv.req_pool_idx = None
def clear(self):
self.free_slots = list(range(1, self._alloc_size))
@@ -1376,7 +1376,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
else:
# Only send delta indices (beyond prefix) to prefill.
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]
kv_indices = (
self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer(
@@ -1390,7 +1390,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
return [
self.req_to_token_pool.translate_mamba_indices(
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()
@@ -1402,7 +1402,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
window_start = max(total_prefix_len, seq_len - window_size)
window_start = page_align_floor(window_start, page_size)
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 = (
self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
@@ -1413,7 +1413,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
def _full_kv_pages_payload():
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
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_start = max(0, seq_len - window_size)
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)
return ring_rows.astype(np.int32)
@@ -1434,7 +1434,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
online = is_dsv4_c128_online_enabled()
ring_size = 1 if online else self.token_to_kv_pool.get_ring_size(128)
return get_dsv4_c128_state_indices(
int(decode_req.req.req_pool_idx),
int(decode_req.req.kv.req_pool_idx),
seq_len,
online=online,
ring_size=ring_size,
@@ -1446,7 +1446,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
self.token_to_kv_pool, "clear_c128_req_state", 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 = {
StateType.MAMBA: _mamba_payload,
StateType.SWA: _swa_payload,
@@ -1465,7 +1465,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
payloads.update(
dsv4_state_payloads(
self.req_to_token_pool,
decode_req.req.req_pool_idx,
decode_req.req.kv.req_pool_idx,
seq_len,
self.token_to_kv_pool_allocator.page_size,
prefix_len=total_prefix_len,
@@ -1494,7 +1494,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
# the C4 sparse physical-slot mapping; carry their logical page IDs
# alongside the independently allocated C4 host page IDs.
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,
]
device_page_indices = kv_to_page_indices(
@@ -1787,7 +1787,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
if prefix_len > 0:
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
@@ -1842,7 +1842,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
host_indices = coordinator.mem_pool_host.alloc_paged_token_slots(
coordinator.req_to_host_pool,
coordinator.req_to_host_pool_allocated_len,
req.req_pool_idx,
req.kv.req_pool_idx,
0,
coordinator.host_token_len(fill_len),
)
@@ -1872,7 +1872,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
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)),
),
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(
(
decode_req.req.req_pool_idx,
decode_req.req.kv.req_pool_idx,
slice(prefix_match.l1_prefix_len, prefix_match.decode_prefix_len),
),
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:
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
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:
return False
@@ -251,7 +251,7 @@ class DecodeKVCacheOffloadManager:
# so a previously-released request must be skipped here to avoid
# non-idempotent side effects (e.g. tree_cache.protected_size_
# 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
kv_committed_len = req.effective_kv_committed_len()
@@ -264,13 +264,13 @@ class DecodeKVCacheOffloadManager:
state = self.offloaded_state.get(req)
if state is not None and state.prefill_len > 0:
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)
start = start_offset
end = kv_committed_len
# 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)
# 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)
if start_p < end_p:
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)
@@ -329,7 +329,7 @@ class DecodeKVCacheOffloadManager:
"""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
# 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
state = self.offloaded_state.get(req)
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
offset = 0
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)
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
]
assert (
python/sglang/srt/disaggregation/prefill.py
+8 -8
View File
@@ -1036,7 +1036,7 @@ class SchedulerDisaggregationPrefillMixin:
else:
logger.warning(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)
maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator)
req.pending_bootstrap = False
@@ -1215,7 +1215,7 @@ class SchedulerDisaggregationPrefillMixin:
return [
self.req_to_token_pool.translate_mamba_indices(
self.req_to_token_pool.req_index_to_mamba_index_mapping[
req.req_pool_idx
req.kv.req_pool_idx
]
)
.cpu()
@@ -1227,7 +1227,7 @@ class SchedulerDisaggregationPrefillMixin:
window_start = max(req.disagg_decode_prefix_len, seq_len - window_size)
window_start = (window_start // page_size) * page_size
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 = (
self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
@@ -1238,7 +1238,7 @@ class SchedulerDisaggregationPrefillMixin:
def _full_kv_pages_payload():
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)
@@ -1251,7 +1251,7 @@ class SchedulerDisaggregationPrefillMixin:
window_size = _pool.unified_swa_window
window_start = max(0, seq_len - window_size)
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)
return ring_rows.astype(np.int32)
@@ -1265,7 +1265,7 @@ class SchedulerDisaggregationPrefillMixin:
)
)
return get_dsv4_c128_state_indices(
int(req.req_pool_idx),
int(req.kv.req_pool_idx),
c128_seq_len,
online=online,
ring_size=ring_size,
@@ -1295,7 +1295,7 @@ class SchedulerDisaggregationPrefillMixin:
payloads.update(
dsv4_state_payloads(
self.req_to_token_pool,
req.req_pool_idx,
req.kv.req_pool_idx,
seq_len,
page_size,
prefix_len=req.disagg_decode_prefix_len,
@@ -1321,7 +1321,7 @@ class SchedulerDisaggregationPrefillMixin:
for seg_start, seg_end in segments:
is_final_segment = seg_end == end_idx
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
# 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,
prefix_slot_ids=prefix_slot_ids,
new_slot_ids=req_new_slots,
req_pool_idx=req.req_pool_idx,
req_pool_idx=req.kv.req_pool_idx,
req=req,
needs_logits=self._chunk_needs_logits(req),
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
num_pages = cache_len // group_tokens
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()
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:
return
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:
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 dict(
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,
)
@@ -130,7 +130,7 @@ def write_dsv4_prealloc_tables(
prealloc path (no ScheduleBatch); no-op without bundle / DSV4 tables."""
if bundle is None or not hasattr(req_to_token_pool, "write_c128"):
return
rp = torch.tensor([req.req_pool_idx])
rp = torch.tensor([req.kv.req_pool_idx])
pl = torch.tensor([prefix_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
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
return
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):
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)
if indices is None:
return None
python/sglang/srt/managers/hisparse_coordinator.py
+31 -31
View File
@@ -349,7 +349,7 @@ class HiSparseCoordinator:
req.hisparse_staging = True
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)
device_indices = (
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(
self.req_to_host_pool,
self.req_to_host_pool_allocated_len,
req.req_pool_idx,
req.kv.req_pool_idx,
0,
prefill_len,
)
@@ -411,11 +411,11 @@ class HiSparseCoordinator:
# 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.
self.req_device_buffer_tokens[
:, req.req_pool_idx, : self.device_buffer_size
:, req.kv.req_pool_idx, : self.device_buffer_size
] = -1
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)
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:
"""Preload all tokens from host pool into the device buffer."""
n = self.host_token_len(req.kv.kv_allocated_len)
host_indices = self.req_to_host_pool[req.req_pool_idx, :n]
device_locs = self.req_to_device_buffer[req.req_pool_idx, :n]
host_indices = self.req_to_host_pool[req.kv.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):
self.mem_pool_host.load_to_device_per_layer(
@@ -456,7 +456,7 @@ class HiSparseCoordinator:
compressed_logical_indices = (
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)
@@ -475,13 +475,13 @@ class HiSparseCoordinator:
raise RuntimeError("HiSparse alloc_device_buffer returned None")
buffer_indices = buffer_indices.to(torch.int32)
self.req_to_device_buffer[req.req_pool_idx, :alloc_size] = buffer_indices
self.req_device_buffer_size[req.req_pool_idx] = alloc_size
self.req_to_device_buffer[req.kv.req_pool_idx, :alloc_size] = buffer_indices
self.req_device_buffer_size[req.kv.req_pool_idx] = alloc_size
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.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]
)
@@ -584,7 +584,7 @@ class HiSparseCoordinator:
_, _, req = self.ack_staging_queue.pop(0)
# prepare device buffer and update 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
finish_count -= 1
ready_reqs.append(req)
@@ -863,21 +863,21 @@ class HiSparseCoordinator:
prefill_len = req.extend_range.end
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)
# Free host memory that was allocated during admit_request_into_staging
host_indices = self.mem_pool_host.allocated_host_indices(
self.req_to_host_pool,
req.req_pool_idx,
self.req_to_host_pool_allocated_len[req.req_pool_idx],
req.kv.req_pool_idx,
self.req_to_host_pool_allocated_len[req.kv.req_pool_idx],
)
if host_indices.numel() > 0:
self.mem_pool_host.free(host_indices)
self.req_to_host_pool[req.req_pool_idx, :] = -1
self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0
self._skip_first_backup[req.req_pool_idx] = False
self.req_to_host_pool[req.kv.req_pool_idx, :] = -1
self.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = 0
self._skip_first_backup[req.kv.req_pool_idx] = False
req.hisparse_staging = False
def retract_req(self, req: Req) -> None:
@@ -901,15 +901,15 @@ class HiSparseCoordinator:
allocated_len = req.kv.kv_allocated_len
# release memory -- only free actually-allocated buffer indices
current_cap = int(self.req_device_buffer_size[req.req_pool_idx])
current_cap = int(self.req_device_buffer_size[req.kv.req_pool_idx])
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])
if all_hi.numel() > 0:
self.token_to_kv_pool_allocator.free_hisparse_indices(all_hi)
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(
allocated_locs
@@ -918,21 +918,21 @@ class HiSparseCoordinator:
host_indices = self.mem_pool_host.allocated_host_indices(
self.req_to_host_pool,
req.req_pool_idx,
self.req_to_host_pool_allocated_len[req.req_pool_idx],
req.kv.req_pool_idx,
self.req_to_host_pool_allocated_len[req.kv.req_pool_idx],
)
if host_indices.numel() > 0:
self.mem_pool_host.free(host_indices)
# clear req info
self.req_device_buffer_tokens[:, req.req_pool_idx, :] = -1
self.req_device_buffer_token_locs[:, req.req_pool_idx, :] = -1
self.req_to_device_buffer[req.req_pool_idx, :] = 0
self.req_device_buffer_size[req.req_pool_idx] = 0
self.req_to_host_pool[req.req_pool_idx, :] = -1
self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0
self.lru_slots[:, req.req_pool_idx, :].copy_(self._lru_init)
self._skip_first_backup[req.req_pool_idx] = False
self.req_device_buffer_tokens[:, req.kv.req_pool_idx, :] = -1
self.req_device_buffer_token_locs[:, req.kv.req_pool_idx, :] = -1
self.req_to_device_buffer[req.kv.req_pool_idx, :] = 0
self.req_device_buffer_size[req.kv.req_pool_idx] = 0
self.req_to_host_pool[req.kv.req_pool_idx, :] = -1
self.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = 0
self.lru_slots[:, req.kv.req_pool_idx, :].copy_(self._lru_init)
self._skip_first_backup[req.kv.req_pool_idx] = False
def _run_swap_in_kernel(
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)
class ReqKvInfo:
# 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).
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)
return lo
@property
def is_held(self) -> bool:
return self.req_pool_idx is not None
@property
def is_released(self) -> bool:
return self.kv_allocated_len == 0 and self.swa_evicted_seqlen == 0
@@ -970,7 +975,6 @@ class Req(ReqDllmMixin):
self.routing_key = routing_key
# Memory pool info
self.req_pool_idx: Optional[int] = None
self.mamba_pool_idx: Optional[torch.Tensor] = None # shape (1)
self.mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None # shape (2)
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
)
@property
def is_holding_kv(self) -> bool:
return self.req_pool_idx is not None
def effective_kv_committed_len(self) -> int:
# Report only the prompt prefix so thinking + answer fall into the
# 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_needs_clear = False
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.extend_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):
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
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):
assert self.retraction_backup is not None
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
mamba_cpu = self.retraction_backup.mamba_cpu
@@ -3492,7 +3492,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# seqlen progress is monotonic per KV handle.
if (
req.decode_batch_idx >= 1
and req.is_holding_kv
and req.kv.is_held
and req.seqlen - 1 - sliding_window_size
>= 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:
# Already past chunked prefill; the running-batch abort path handles
# 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
return
@@ -3161,7 +3161,7 @@ class Scheduler(
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(
req_pool_indices, dtype=torch.int64, device=device
)
@@ -4868,7 +4868,7 @@ class Scheduler(
_make_abort_req(req), req
)
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
):
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
seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop)
req.routed_experts = capturer.get_topk(
req_pool_idx=req.req_pool_idx,
req_pool_idx=req.kv.req_pool_idx,
seqlen=seqlen,
req_to_token_pool=self.req_to_token_pool,
start_len=start_len,
@@ -166,7 +166,7 @@ class SchedulerBatchResultProcessor:
return
seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop)
req.indexer_topk = capturer.get_topk(
req_pool_idx=req.req_pool_idx,
req_pool_idx=req.kv.req_pool_idx,
seqlen=seqlen,
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
for batch in batches:
for req in batch.reqs:
if not req.is_holding_kv:
if not req.kv.is_held:
continue
allocated_len = req.kv.kv_allocated_len
@@ -324,23 +324,23 @@ class SchedulerInvariantChecker:
batch = self.get_last_batch()
if batch is not None:
for req in batch.reqs:
if not req.is_holding_kv:
if not req.kv.is_held:
continue
_add_owner(
req,
f"req {req.rid}",
req.req_pool_idx,
req.kv.req_pool_idx,
req.kv.kv_committed_len,
req.kv.kv_allocated_len,
)
sess = getattr(self.tree_cache, "slots", None)
if sess:
for sid, slot in sess.items():
if getattr(slot, "is_holding_kv", False):
if slot.kv.is_held:
_add_owner(
slot,
f"slot {sid[:8]}",
slot.req_pool_idx,
slot.kv.req_pool_idx,
slot.kv.kv_committed_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():
continue
for req in batch.reqs:
if req.req_pool_idx is not None:
idxs.add(req.req_pool_idx)
if req.kv.req_pool_idx is not None:
idxs.add(req.kv.req_pool_idx)
return idxs
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)
# 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[
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)
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
if batch.is_dllm():
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
]
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
# The protected prefix is not this req's to free.
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)
def cache_unfinished_req(self, req: Req, chunked=False):
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 = kv_indices.to(dtype=torch.int64, copy=True)
@@ -158,7 +158,7 @@ class PureSWAChunkCache(SWAChunkCache):
):
kv_committed_len = kv_len_to_handle
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.
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,
retain_floor: int | None = None,
) -> None:
if not req.is_holding_kv:
if not req.kv.is_held:
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
@@ -99,7 +99,7 @@ def free_swa_out_of_window_slots(
if new_swa_evicted_seqlen > req.kv.swa_evicted_seqlen:
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)
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):
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
if not req.is_holding_kv:
if not req.kv.is_held:
assert (
tree_cache.supports_mamba()
), "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
# internally, then sets req_pool_idx = None.
assert (not req.is_holding_kv) == req.kv.is_released
if not req.is_holding_kv:
assert (not req.kv.is_held) == req.kv.is_released
if not req.kv.is_held:
return
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)
if start_p < end_p:
indices_to_free = tree_cache.req_to_token_pool.req_to_token[req.req_pool_idx][
start_p:end_p
]
indices_to_free = tree_cache.req_to_token_pool.req_to_token[
req.kv.req_pool_idx
][start_p:end_p]
# 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.
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."""
if self.disable:
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.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]
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:
@@ -676,7 +676,7 @@ class MambaRadixCache(BasePrefixCache):
def _skip_cache_unfinished_req(req: Req) -> None:
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
@@ -693,7 +693,7 @@ class MambaRadixCache(BasePrefixCache):
return _skip_cache_unfinished_req(req)
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 = kv_indices_orig[:cache_len]
@@ -786,7 +786,7 @@ class MambaRadixCache(BasePrefixCache):
), f"{new_prefix_len=}, {len(new_indices)=}"
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 :],
)
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]]:
# Indices of reqs that already have a req_pool_idx and will reuse
# 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
# https://github.com/sgl-project/sglang/pull/20476
# if not any(r.is_dllm() for r in reqs):
@@ -309,10 +309,10 @@ class ReqToTokenPool:
return None
offset = 0
for r in reqs:
if r.req_pool_idx is None:
r.req_pool_idx = select_index[offset]
if r.kv.req_pool_idx is None:
r.kv.req_pool_idx = select_index[offset]
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]]:
"""Take need_size rows and bump their generation, with no Req bound to
@@ -338,9 +338,9 @@ class ReqToTokenPool:
self.free_slots.extend(indices)
def free(self, req: Req):
assert req.req_pool_idx is not None, "request must have req_pool_idx"
self.free_rows([req.req_pool_idx])
req.req_pool_idx = None
assert req.kv.req_pool_idx is not None, "request must have req_pool_idx"
self.free_rows([req.kv.req_pool_idx])
req.kv.req_pool_idx = None
def clear(self):
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.
"""
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
)
@@ -1530,7 +1530,9 @@ class HybridReqToTokenPool(ReqToTokenPool):
if self.enable_mamba_extra_buffer:
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:
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
if self.disable:
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)
return
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len
req.kv.req_pool_idx, :kv_committed_len
]
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:
# The protected prefix is not this req's to free.
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(
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]
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(
@@ -520,7 +520,7 @@ class RadixCache(BasePrefixCache):
token_ids = req.get_fill_ids()
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(
@@ -558,7 +558,7 @@ class RadixCache(BasePrefixCache):
), f"{len(new_indices)=}, {len(radix_key)=}"
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 :],
)
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."""
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]
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)
# 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):
"""Cache request when it is unfinished."""
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()
prefill_len = len(token_ids) # prefill only (maybe chunked)
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)
# 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]
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
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.
"""
if req.req_pool_idx is not None:
self.states.register(req.req_pool_idx, len(req.origin_input_ids))
if req.kv.req_pool_idx is not None:
self.states.register(req.kv.req_pool_idx, len(req.origin_input_ids))
def on_request_end(self, req: "Req") -> None:
"""
Handle request end event. Called when a request is completed or aborted.
Cleans up request-specific state and releases resources.
"""
if req.req_pool_idx is None:
if req.kv.req_pool_idx is None:
return
self.states.clear(req.req_pool_idx)
self.states.clear(req.kv.req_pool_idx)
# TODO: Implement request end handling
# - 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:
return
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
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]
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.
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."""
if self.disable:
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)
return
token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
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(
@@ -516,7 +516,7 @@ class SWARadixCache(BasePrefixCache):
"""Cache request when it is unfinished."""
if self.disable:
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
@@ -525,7 +525,7 @@ class SWARadixCache(BasePrefixCache):
token_ids = req.get_fill_ids()
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(
@@ -561,7 +561,7 @@ class SWARadixCache(BasePrefixCache):
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=}"
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:],
)
python/sglang/srt/mem_cache/unified_radix_cache.py
+7 -7
View File
@@ -825,7 +825,7 @@ class UnifiedRadixCache(BasePrefixCache):
if self.disable:
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)
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]
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
@@ -918,13 +918,13 @@ class UnifiedRadixCache(BasePrefixCache):
if self.disable:
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)
return
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
@@ -991,7 +991,7 @@ class UnifiedRadixCache(BasePrefixCache):
new_indices
), f"{new_prefix_len=}, {len(new_indices)=}"
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 :],
)
@@ -1160,7 +1160,7 @@ class UnifiedRadixCache(BasePrefixCache):
) -> tuple[torch.Tensor, list[PoolTransfer]]:
num_tokens = req.seqlen - 1
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)
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 = window_start // self.page_size * self.page_size
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)
swa_indices = kv_cache.translate_loc_from_full_to_swa(window_indices)
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)
# KV pool state
req_pool_idx: Optional[int] = None
kv: ReqKvInfo = field(default_factory=ReqKvInfo)
# 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_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):
"""Save KV state from a finishing request into this slot."""
self.req_pool_idx = req.req_pool_idx
if is_first:
self.last_node = req.last_node
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
# alloc/retract path cannot mistake slot-owned mamba state for its own.
req.req_pool_idx = None
req.kv = ReqKvInfo()
req.mamba_pool_idx = None
req.mamba_ping_pong_track_buffer = None
@@ -104,7 +95,6 @@ class SessionSlot:
def restore_to_req(self, req: Req):
"""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.swa_uuid_for_lock = self.swa_uuid_for_lock
req.skip_lock_node_ids = self.skip_lock_node_ids
@@ -189,7 +179,7 @@ class StreamingSession(BasePrefixCache):
return session_id in self.slots
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) --
@@ -217,7 +207,7 @@ class StreamingSession(BasePrefixCache):
if not _is_streaming(req):
return None
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
if req.to_finish is not None:
req.session.abort_req()
@@ -290,7 +280,7 @@ class StreamingSession(BasePrefixCache):
self._free_tail(slot, req, prefix_len)
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)
return MatchResult(
@@ -330,7 +320,6 @@ class StreamingSession(BasePrefixCache):
# return the (possibly extra_buffer ping-pong) slots to
# the mamba pool; otherwise the abort orphans them.
slot = SessionSlot(
req_pool_idx=req.req_pool_idx,
kv=copy.copy(req.kv),
last_node=req.last_node,
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
)
self.release_session(session_id)
req.req_pool_idx = None
req.kv = ReqKvInfo()
req.session.abort_req()
return True
@@ -386,7 +374,7 @@ class StreamingSession(BasePrefixCache):
return False
if chunked:
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)
return True
@@ -441,9 +429,7 @@ class StreamingSession(BasePrefixCache):
protected_len = slot.kv.cache_protected_len
lock_node = slot.last_node
tokens_freed = (
max(0, slot.kv.kv_allocated_len - protected_len)
if slot.is_holding_kv
else 0
max(0, slot.kv.kv_allocated_len - protected_len) if slot.kv.is_held else 0
)
logger.info(
"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
end = slot.kv.kv_allocated_len
if start < end:
kv_indices = self.req_to_token_pool.req_to_token[
slot.req_pool_idx, start:end
slot.kv.req_pool_idx, start:end
]
self.token_to_kv_pool_allocator.free(kv_indices)
self.req_to_token_pool.free(slot)
@@ -483,9 +469,10 @@ class StreamingSession(BasePrefixCache):
total = 0
for slot in self.slots.values():
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)
total += allocated - slot.kv.cache_protected_len
return total
@@ -499,9 +486,10 @@ class StreamingSession(BasePrefixCache):
total = 0
for slot in self.slots.values():
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)
total += allocated - max(
slot.kv.cache_protected_len, slot.kv.swa_evicted_seqlen
@@ -513,9 +501,9 @@ class StreamingSession(BasePrefixCache):
def _owned(s):
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())
@@ -529,7 +517,8 @@ class StreamingSession(BasePrefixCache):
total = 0
for slot in self.slots.values():
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:
continue
@@ -559,7 +548,9 @@ class StreamingSession(BasePrefixCache):
decoding pushes allocated above committed, or when retract retry's
logit-reserve pulls prefix_len below committed.
"""
self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv.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_committed_len = min(slot.kv.kv_committed_len, 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.
"""
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_committed_len = min(req.kv.kv_committed_len, 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
def kv_pages(self) -> int:
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
page_size = self.context.scheduler.page_size
return (req.kv.kv_allocated_len + page_size - 1) // page_size