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

Scope prefetch cache state to the request attempt (#39318)

Browse Source
This commit is contained in:
cctry
2026-09-14 15:03:03 -07:00
committed by GitHub
parent d72e59508b
commit dad8c074e7
21 changed files with 541 additions and 387 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_hicache_mixin.py
+9 -6
View File
@@ -12,6 +12,7 @@ import torch
from sglang.srt.disaggregation.base import KVPoll
from sglang.srt.managers.schedule_policy import match_prefix_for_req
from sglang.srt.mem_cache.base_prefix_cache import (
CacheRequestOutcome,
InitLoadBackParams,
)
@@ -129,7 +130,7 @@ class DecodeHiCachePreallocMixin:
else None
)
self.tree_cache.prefetch_from_storage(
req.rid,
req.cache_request_handle,
prefix_match.last_host_node,
suffix,
last_hash,
@@ -137,8 +138,8 @@ class DecodeHiCachePreallocMixin:
extra_key=req.extra_key,
cache_salt=req.cache_salt,
)
prefix_match.prefetch_registered = (
req.rid in self.tree_cache.ongoing_prefetch
prefix_match.prefetch_registered = self.tree_cache.has_ongoing_prefetch(
req.cache_request_handle
)
except Exception as e:
logger.warning(
@@ -186,7 +187,9 @@ class DecodeHiCacheTransferMixin:
decode_req.prefix_match is not None
and decode_req.prefix_match.prefetch_registered
):
self.tree_cache.release_aborted_request(decode_req.req.rid)
self.tree_cache.finish(
decode_req.req.cache_request_handle, CacheRequestOutcome.ABORT
)
if decode_req.hicache_restored_node is not None:
self.tree_cache.dec_lock_ref(
decode_req.hicache_restored_node,
@@ -207,9 +210,9 @@ class DecodeHiCacheTransferMixin:
# Wait for L3 -> L2 prefetch to drain (skip when no L3 hit).
if pm.l3_storage_hit_length > 0:
if not self.tree_cache.check_prefetch_progress(dr.req.rid):
if not self.tree_cache.check_prefetch_progress(dr.req.cache_request_handle):
return False
self.tree_cache.pop_prefetch_loaded_tokens(dr.req.rid)
self.tree_cache.pop_prefetch_loaded_tokens(dr.req.cache_request_handle)
# Re-match: req.last_node / prefix_indices updated to current device state.
rematch = match_prefix_for_req(
python/sglang/srt/disaggregation/prefill.py
+9 -2
View File
@@ -66,6 +66,7 @@ from sglang.srt.managers.schedule_batch import (
Req,
ScheduleBatch,
)
from sglang.srt.mem_cache.base_prefix_cache import CacheRequestOutcome
from sglang.srt.mem_cache.common import (
kv_to_page_indices,
kv_to_page_num,
@@ -999,6 +1000,9 @@ class SchedulerDisaggregationPrefillMixin:
if not isinstance(req.finished_reason, FINISH_ABORT):
req.finished_reason = FINISH_LENGTH(length=0)
release_kv_cache(req, self.tree_cache) # unlock the tree
self.tree_cache.finish(
req.cache_request_handle, CacheRequestOutcome.SUCCESS
)
# FIXME: clean up req's data in transfer engine
req.disagg_kv_sender.clear()
done_reqs.append(req)
@@ -1070,6 +1074,7 @@ class SchedulerDisaggregationPrefillMixin:
logger.warning(error_message)
req.time_stats.trace_ctx.abort(abort_info={"reason": error_message})
release_kv_cache(req, self.tree_cache) # unlock the tree
self._release_aborted_request(req)
if not isinstance(req.finished_reason, FINISH_ABORT):
prepare_abort(
req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR
@@ -1111,7 +1116,7 @@ class SchedulerDisaggregationPrefillMixin:
maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator)
req.pending_bootstrap = False
if self.enable_hicache_storage:
self.tree_cache.release_aborted_request(req.rid)
self.tree_cache.finish(req.cache_request_handle, CacheRequestOutcome.ABORT)
if req.kv.holds_kv or req.kv.holds_mamba:
release_kv_cache(req, self.tree_cache, is_insert=False)
return True
@@ -1143,7 +1148,7 @@ class SchedulerDisaggregationPrefillMixin:
if self.metrics_reporter.enable_metrics:
self.metrics_collector.increment_bootstrap_failed_reqs()
if self.enable_hicache_storage:
self.tree_cache.release_aborted_request(req.rid)
self.tree_cache.finish(req.cache_request_handle, CacheRequestOutcome.ABORT)
def handle_pending_bootstrap(self: Scheduler, req: Req, poll: KVPoll) -> bool:
"""Return True when bootstrap is finalized and KV transfer can proceed."""
@@ -1466,6 +1471,7 @@ class SchedulerDisaggregationPrefillMixin:
"""Release KV cache and requeue an optimistic prefill request."""
max_attempts = get_disagg().optimistic_prefill_attempts
maybe_cache_unfinished_req(req, self.tree_cache)
self._release_aborted_request(req)
release_kv_cache(req, self.tree_cache)
req.reset_for_retract()
req.output_ids = array("q")
@@ -1478,6 +1484,7 @@ class SchedulerDisaggregationPrefillMixin:
req.output_dsa_topk_indices = None
req.pending_bootstrap = True
req.time_stats.reset_prefill_retry_time()
req.advance_cache_request_handle()
if req.prefill_attempt_count >= max_attempts:
logger.info(
f"Req {req.rid} exhausted optimistic prefill attempts "
python/sglang/srt/managers/schedule_batch.py
+8
View File
@@ -104,6 +104,7 @@ from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
CacheRequestHandle,
DecLockRefParams,
MatchPrefixParams,
zero_match_result,
@@ -979,6 +980,7 @@ class Req(ReqDllmMixin):
):
# Input and output info
self.rid = rid
self.cache_request_handle = CacheRequestHandle(rid=rid, attempt_id=0)
self.origin_input_ids = origin_input_ids
self.origin_input_ids_unpadded = (
origin_input_ids_unpadded
@@ -1335,6 +1337,12 @@ class Req(ReqDllmMixin):
# Snapshot of the scheduler prefill-token counter taken at waiting_queue entry; used by HRRN aging.
self.arrival_processed_tokens: int = 0
def advance_cache_request_handle(self) -> None:
self.cache_request_handle = dataclasses.replace(
self.cache_request_handle,
attempt_id=self.cache_request_handle.attempt_id + 1,
)
@property
def seqlen(self) -> int:
"""Get the current sequence length of the request."""
python/sglang/srt/managers/schedule_policy.py
+4 -2
View File
@@ -980,7 +980,9 @@ class PrefillAdder:
if req.retracted_stain:
# Retraction attribution is intentionally omitted for now; discard
# its lifecycle state so a later abort cannot report it as a drop.
self.tree_cache.discard_storage_prefetch_accounting(req.rid)
self.tree_cache.discard_storage_prefetch_accounting(
req.cache_request_handle
)
return
if prefix_len > 0:
@@ -1005,7 +1007,7 @@ class PrefillAdder:
else "shrunk"
)
self.tree_cache.finish_storage_prefetch_admission(
req.rid,
req.cache_request_handle,
fulfilled_tokens=fulfilled_storage_hit,
reason=reason,
)
python/sglang/srt/managers/scheduler.py
+15 -17
View File
@@ -279,6 +279,7 @@ from sglang.srt.managers.utils import (
validate_input_length,
)
from sglang.srt.mem_cache import kv_cache_builder
from sglang.srt.mem_cache.base_prefix_cache import CacheRequestOutcome
from sglang.srt.mem_cache.common import (
maybe_cache_unfinished_req,
release_kv_cache,
@@ -3111,7 +3112,7 @@ class Scheduler(
else None
)
tree_cache.prefetch_from_storage(
req.rid,
req.cache_request_handle,
last_host_node,
new_input_tokens,
tree_cache.get_last_hash_value(last_host_node),
@@ -3131,7 +3132,7 @@ class Scheduler(
return
max_attempts = get_memory().hicache_storage_prefetch_retry_max_attempts
for req in self.waiting_queue:
if self.tree_cache.pop_storage_prefetch_miss(req.rid):
if self.tree_cache.pop_storage_prefetch_miss(req.cache_request_handle):
req.storage_prefetch_retry_pending = True
req.storage_prefetch_retry_wait_polls = 0
if (
@@ -3208,14 +3209,9 @@ class Scheduler(
return False
return True
def _release_aborted_request(self, rid: str) -> None:
def _release_aborted_request(self, req: Req) -> None:
"""Drop the cache-side state an aborted request left behind."""
if (
self.enable_hierarchical_cache
or self.enable_hicache_storage
or self.enable_unified_cache_external_linker
):
self.tree_cache.release_aborted_request(rid)
self.tree_cache.finish(req.cache_request_handle, CacheRequestOutcome.ABORT)
def _abort_on_queued_limit(self, recv_req: Req) -> bool:
"""Abort an incoming or existing request if the waiting queue is full. Returns True if the incoming request is aborted."""
@@ -3243,7 +3239,7 @@ class Scheduler(
direction * recv_req.priority < direction * candidate_req.priority
)
if abort_existing_req:
self._release_aborted_request(candidate_req.rid)
self._release_aborted_request(candidate_req)
self.waiting_queue.pop(idx)
self.beam_coordinator.retire_group(candidate_req)
req_to_abort = candidate_req
@@ -3457,7 +3453,7 @@ class Scheduler(
req, self.req_to_metadata_buffer_idx_allocator
)
req.pending_bootstrap = False
self._release_aborted_request(req.rid)
self._release_aborted_request(req)
release_kv_cache(req, self.tree_cache, is_insert=False)
self.chunked_req = None
@@ -3855,14 +3851,16 @@ class Scheduler(
break
if self.enable_hicache_storage:
prefetch_done = self.tree_cache.check_prefetch_progress(req.rid)
prefetch_done = self.tree_cache.check_prefetch_progress(
req.cache_request_handle
)
if not prefetch_done:
# skip staging requests that are ongoing prefetch
continue
# Pop the L3-loaded span. Unified cache exposes its absolute
# start so cache-mode L2/L3 attribution survives L3-tail eviction.
loaded_tokens, loaded_start = self.tree_cache.pop_prefetch_loaded_span(
req.rid
req.cache_request_handle
)
if loaded_tokens > 0:
req.storage_hit_length = loaded_tokens
@@ -3886,7 +3884,7 @@ class Scheduler(
# (fenced in init_hicache) will need the same charge via
# mamba_host_hit_length.
held_tokens, held_swa_tokens = self.tree_cache.plan_staged_splice(
req.rid, len(req.prefix_indices)
req.cache_request_handle, len(req.prefix_indices)
)
if held_tokens > 0:
req.host_hit_length = held_tokens
@@ -5203,7 +5201,7 @@ class Scheduler(
# This only works for requests that have not started anything.
# We still need to send something back to TokenizerManager to clean up the state.
req = self.waiting_queue.pop(i)
self._release_aborted_request(req.rid)
self._release_aborted_request(req)
self.beam_coordinator.retire_group(req)
# Without the initiator's reason the tokenizer falls back to a
# generic abort message.
@@ -5239,7 +5237,7 @@ class Scheduler(
for req in self.dllm_manager.pop_aborted_reqs(
recv_req.abort_all, recv_req.rid
):
self._release_aborted_request(req.rid)
self._release_aborted_request(req)
self.ipc_channels.send_to_tokenizer.send_output(
_make_abort_req(req), req
)
@@ -5259,7 +5257,7 @@ class Scheduler(
for req in self.disagg_prefill_bootstrap_queue.queue:
if recv_req.abort_all or req.rid.startswith(recv_req.rid):
logger.debug(f"Abort bootstrap queue request. {req.rid=}")
self._release_aborted_request(req.rid)
self._release_aborted_request(req)
if hasattr(req.disagg_kv_sender, "abort"):
req.disagg_kv_sender.abort()
python/sglang/srt/mem_cache/base_prefix_cache.py
+29 -4
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import dataclasses
import time
from abc import ABC, abstractmethod
from enum import Enum, auto
from typing import (
TYPE_CHECKING,
Any,
@@ -38,6 +39,17 @@ if TYPE_CHECKING:
)
@dataclasses.dataclass(frozen=True)
class CacheRequestHandle:
rid: str
attempt_id: int
class CacheRequestOutcome(Enum):
SUCCESS = auto()
ABORT = auto()
@runtime_checkable
class PrefixCacheTrait(Protocol):
req_to_token_pool: ReqToTokenPool
@@ -346,6 +358,14 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
tens of seconds (see HostKVCache.destroy). Idempotent.
"""
def release_aborted_request(self, handle: CacheRequestHandle) -> None:
"""Release attempt state; caches without prefetch state have nothing to drop."""
def finish(self, handle: CacheRequestHandle, outcome: CacheRequestOutcome) -> None:
"""Finish an attempt without cancelling successful asynchronous cache work."""
if outcome != CacheRequestOutcome.SUCCESS:
self.release_aborted_request(handle)
@abstractmethod
def reset(self):
pass
@@ -484,19 +504,24 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
raise NotImplementedError()
def finish_storage_prefetch_admission(
self, req_id: str, fulfilled_tokens: int, reason: Optional[str]
self,
handle: CacheRequestHandle,
fulfilled_tokens: int,
reason: Optional[str],
) -> None:
"""Resolve storage-hit accounting once a request is admitted.
Non-storage caches have no lifecycle state to resolve.
"""
def discard_storage_prefetch_accounting(self, req_id: str) -> None:
def discard_storage_prefetch_accounting(self, handle: CacheRequestHandle) -> None:
"""Forget storage-hit lifecycle state without emitting a result."""
def pop_prefetch_loaded_span(self, req_id: str) -> tuple[int, Optional[int]]:
def pop_prefetch_loaded_span(
self, handle: CacheRequestHandle
) -> tuple[int, Optional[int]]:
"""Pop L3-loaded tokens and their absolute prefix start, if known."""
return self.pop_prefetch_loaded_tokens(req_id), None
return self.pop_prefetch_loaded_tokens(handle), None
def ready_to_load_host_cache(self) -> Any:
"""
python/sglang/srt/mem_cache/buffer_mode/pipeline.py
+65 -59
View File
@@ -36,6 +36,7 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.managers.cache_controller import HICACHE_WRITE_STAGING_POOL_FRACTION
from sglang.srt.mem_cache.base_prefix_cache import (
CacheRequestHandle,
DecLockRefParams,
EvictParams,
InitLoadBackParams,
@@ -99,7 +100,7 @@ class _StagedPrefetch(msgspec.Struct):
the op-owned host bounce exists (no device state, nothing in the tree).
"""
req_id: str
request: CacheRequestHandle
key_tokens: list[int]
extra_key: Optional[str]
cache_salt: Optional[str]
@@ -117,7 +118,7 @@ class _OngoingBufferLoadBack(msgspec.Struct):
tree-resident; only the host bounce remains to free.
"""
req_id: str
request: CacheRequestHandle
num_tokens: int
occupied_tokens: int
aux_xfers: list[PoolTransfer]
@@ -272,9 +273,9 @@ class BufferModePipeline:
# negative ack id).
self.pending_hit_allocs: deque = deque()
self._prefetch_prefix_ctx: dict[
str, tuple[list[int], Optional[str], Optional[str]]
CacheRequestHandle, tuple[list[int], Optional[str], Optional[str]]
] = {}
self.staged_prefetches: dict[str, _StagedPrefetch] = {}
self.staged_prefetches: dict[CacheRequestHandle, _StagedPrefetch] = {}
self.ongoing_buffer_load_back: dict[int, _OngoingBufferLoadBack] = {}
# Backup pipeline: FIFO intents awaiting a D2H slot, node ids
# anywhere in flight (dedupes re-triggers), and a content refcount
@@ -292,8 +293,8 @@ class BufferModePipeline:
self.write_staged_tokens_ = 0
self.write_backlog_tokens_ = 0
self._backlog_cap_hits = 0
# rid-keyed anchor locks; released idempotently at every exit.
self.anchor_locks: dict[str, _AnchorLock] = {}
# Attempt-keyed locks stop stale completions from unlocking retries.
self.anchor_locks: dict[CacheRequestHandle, _AnchorLock] = {}
self.anchor_locked_tokens_ = 0
self._anchor_lock_cap_skips = 0
@@ -743,16 +744,16 @@ class BufferModePipeline:
# ---- load back pipeline (storage -> staging -> device) ----
def try_lock_anchor(self, req_id: str) -> str:
def try_lock_anchor(self, request: CacheRequestHandle) -> str:
"""Pin the staged prefetch's device anchor so eviction cannot
invalidate the splice, finding it by re-matching the live tree
(carried node ids go stale via splits and eviction; the walk is
O(prefix path)). Returns "locked", "no_anchor" (nothing to pin),
"cap_skip" (over cap; launches unlocked), or "anchor_lost" (splice
base gone — the caller cancels the storage IO)."""
if req_id in self.anchor_locks:
if request in self.anchor_locks:
return "locked"
prefix_ctx = self._prefetch_prefix_ctx.get(req_id)
prefix_ctx = self._prefetch_prefix_ctx.get(request)
if not prefix_ctx or not prefix_ctx[0]:
return "no_anchor" # root anchor: nothing to pin
prefix_tokens, extra_key, cache_salt = prefix_ctx
@@ -777,7 +778,7 @@ class BufferModePipeline:
if cache.tree_core.is_eagle:
# The suffix owns the boundary token shared with the last matched
# bigram, so include it when rebuilding the anchor key.
info = cache.ongoing_prefetch.get(req_id)
info = cache.ongoing_prefetch.get(request)
if info is None or not info.prefetch_key.token_ids:
return "anchor_lost"
anchor_tokens.append(info.prefetch_key.token_ids[0])
@@ -794,7 +795,7 @@ class BufferModePipeline:
if len(match.device_indices) < matched_len:
return "anchor_lost"
lock_params = cache.inc_lock_ref(match.last_device_node).to_dec_params()
self.anchor_locks[req_id] = _AnchorLock(
self.anchor_locks[request] = _AnchorLock(
node_id=match.last_device_node,
lock_params=lock_params,
tokens=matched_len,
@@ -802,28 +803,30 @@ class BufferModePipeline:
self.anchor_locked_tokens_ += matched_len
return "locked"
def release_anchor_lock(self, req_id: str) -> None:
def release_anchor_lock(self, request: CacheRequestHandle) -> None:
"""Drop a staged prefetch's anchor lock (idempotent; called at every
consume/drop/abort exit)."""
lock = self.anchor_locks.pop(req_id, None)
lock = self.anchor_locks.pop(request, None)
if lock is None:
return
self._cache.dec_lock_ref(lock.node_id, lock.lock_params)
self.anchor_locked_tokens_ -= lock.tokens
assert self.anchor_locked_tokens_ >= 0, (
f"anchor-lock accounting corrupted: locked={self.anchor_locked_tokens_} "
f"after releasing {req_id}"
f"after releasing {request.rid}"
)
def staged_span_covered(self, req_id: str, span_tokens: int) -> bool:
def staged_span_covered(
self, request: CacheRequestHandle, span_tokens: int
) -> bool:
"""True when the live device tree already covers the fetch's whole
would-be span (prefix + the storage-hit tokens): nothing would be
left to splice at consumption, so the IO-commit caller cancels
before the bounce alloc and the storage read."""
info = self._cache.ongoing_prefetch.get(req_id)
info = self._cache.ongoing_prefetch.get(request)
if info is None or span_tokens <= 0:
return False
prefix_tokens, _, _ = self._prefetch_prefix_ctx[req_id]
prefix_tokens, _, _ = self._prefetch_prefix_ctx[request]
span_key = info.prefetch_key
full_tokens = array("q", prefix_tokens)
full_tokens.extend(span_key[:span_tokens].token_ids)
@@ -838,7 +841,7 @@ class BufferModePipeline:
def set_prefix_ctx(
self,
req_id: str,
request: CacheRequestHandle,
matched_prefix_tokens,
extra_key: Optional[str] = None,
cache_salt: Optional[str] = None,
@@ -846,17 +849,17 @@ class BufferModePipeline:
"""Record the device-matched prefix (and its tree-key namespace) at
prefetch enqueue; consumed at staging commit to build the full-span
tree key, and by try_lock_anchor to re-match a stale anchor."""
self._prefetch_prefix_ctx[req_id] = (
self._prefetch_prefix_ctx[request] = (
list(matched_prefix_tokens or []),
extra_key,
cache_salt,
)
def pop_prefix_ctx(self, req_id: str) -> None:
self._prefetch_prefix_ctx.pop(req_id, None)
def pop_prefix_ctx(self, request: CacheRequestHandle) -> None:
self._prefetch_prefix_ctx.pop(request, None)
def has_staged(self, req_id: str) -> bool:
return req_id in self.staged_prefetches
def has_staged(self, request: CacheRequestHandle) -> bool:
return request in self.staged_prefetches
@staticmethod
def _occupied_span(host_indices) -> int:
@@ -866,7 +869,7 @@ class BufferModePipeline:
def stage_completed_prefetch(
self,
req_id: str,
request: CacheRequestHandle,
num_tokens: int,
hash_value: list[str],
) -> bool:
@@ -881,9 +884,9 @@ class BufferModePipeline:
operation,
_lock_params,
comp_xfers,
) = cache.ongoing_prefetch.pop(req_id)
) = cache.ongoing_prefetch.pop(request)
cc = cache.cache_controller
prefix_ctx = self._prefetch_prefix_ctx.pop(req_id, None)
prefix_ctx = self._prefetch_prefix_ctx.pop(request, None)
prefix_tokens = prefix_ctx[0] if prefix_ctx is not None else None
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
# Component transfers are already present in comp_xfers. Preserve the
@@ -897,14 +900,14 @@ class BufferModePipeline:
if num_tokens == 0 or prefix_tokens is None:
# Nothing usable fetched: recompute.
cache.discard_storage_prefetch_accounting(req_id)
self.release_anchor_lock(req_id)
cache.discard_storage_prefetch_accounting(request)
self.release_anchor_lock(request)
cc.append_host_mem_release(
host_indices[:num_tokens], extra_pools=aux_xfers or None
)
cc.prefetch_tokens_occupied -= self._occupied_span(host_indices)
cache.prefetch_loaded_tokens_by_reqid[req_id] = 0
cache.prefetch_loaded_storage_start_by_reqid.pop(req_id, None)
cache.prefetch_loaded_tokens_by_reqid[request] = 0
cache.prefetch_loaded_storage_start_by_reqid.pop(request, None)
return True
staged_pages = num_tokens // cache.page_size
@@ -916,8 +919,8 @@ class BufferModePipeline:
cache.storage_existence_cache.add(PoolName.KV, list(staged_hashes))
occupied_tokens = self._occupied_span(host_indices)
self.staged_prefetches[req_id] = _StagedPrefetch(
req_id=req_id,
self.staged_prefetches[request] = _StagedPrefetch(
request=request,
key_tokens=prefix_tokens + list(prefetch_key[:num_tokens].token_ids),
extra_key=prefetch_key.extra_key,
cache_salt=prefetch_key.cache_salt,
@@ -929,18 +932,18 @@ class BufferModePipeline:
hash_values=staged_hashes,
operation_id=operation.id,
)
cache.prefetch_loaded_tokens_by_reqid[req_id] = num_tokens
cache.prefetch_loaded_storage_start_by_reqid[req_id] = operation.storage_start
cache.prefetch_loaded_tokens_by_reqid[request] = num_tokens
cache.prefetch_loaded_storage_start_by_reqid[request] = operation.storage_start
return True
def plan_staged_splice(
self, req_id: str, device_prefix_len: int
self, request: CacheRequestHandle, device_prefix_len: int
) -> tuple[int, int]:
"""(kv, swa) host-hit tokens consumption will splice given the
request's live device prefix, so admission charges no phantom
tokens. Frees a hold that can no longer splice: surfaced as 0 but
kept, it would leak — the adder only consumes surfaced host hits."""
f = self.staged_prefetches.get(req_id)
f = self.staged_prefetches.get(request)
if f is None:
return 0, 0
splice_tokens = staged_splice_tokens(f, device_prefix_len)
@@ -949,28 +952,28 @@ class BufferModePipeline:
logger.info(
"HiCache staged prefetch released req=%s matched=%d "
"device_prefix=%d tokens=%d",
req_id,
request.rid,
f.matched_len,
device_prefix_len,
f.num_tokens,
)
reason = None if covered_tokens == f.num_tokens else "shrunk"
self.release_staged_hold(req_id, reason=reason)
self.release_staged_hold(request, reason=reason)
return 0, 0
return splice_tokens, self.staged_prefetch_swa_tokens(req_id)
return splice_tokens, self.staged_prefetch_swa_tokens(request)
def _resolve_staged_device_coverage(
self, f: _StagedPrefetch, device_prefix_len: int
) -> int:
covered_tokens = min(max(device_prefix_len - f.matched_len, 0), f.num_tokens)
self._cache._resolve_storage_prefetch_tokens(f.req_id, covered_tokens)
self._cache._resolve_storage_prefetch_tokens(f.request, covered_tokens)
return covered_tokens
def staged_prefetch_swa_tokens(self, req_id: str) -> int:
def staged_prefetch_swa_tokens(self, request: CacheRequestHandle) -> int:
"""SWA device tokens consuming this staged prefetch will allocate (the
staged trailing window); surfaced as the request's swa_host_hit_length
so the adder's SWA gate charges the admission-time alloc."""
f = self.staged_prefetches.get(req_id)
f = self.staged_prefetches.get(request)
if f is None:
return 0
return sum(
@@ -993,17 +996,18 @@ class BufferModePipeline:
cache = self._cache
req = params.req
assert req is not None
request = req.cache_request_handle
empty = cache.tree_core.empty_match_result.device_indices
unchanged = (empty, req.last_node)
f = self.staged_prefetches.pop(req.rid, None)
f = self.staged_prefetches.pop(request, None)
if f is None:
self.release_anchor_lock(req.rid)
self.release_anchor_lock(request)
return unchanged
cc = cache.cache_controller
def _drop(reason: Optional[str]) -> tuple[torch.Tensor, NodeId]:
cache._finish_storage_prefetch(req.rid, fulfilled_tokens=0, reason=reason)
self.release_anchor_lock(req.rid)
cache._finish_storage_prefetch(request, fulfilled_tokens=0, reason=reason)
self.release_anchor_lock(request)
self._free_staging_now(f.host_indices, f.aux_xfers)
cc.prefetch_tokens_occupied -= f.occupied_tokens
# Nothing spliced: keep the surfaced host-hit fields truthful.
@@ -1038,7 +1042,7 @@ class BufferModePipeline:
f.matched_len,
splice_base,
f.num_tokens,
req.rid in self.anchor_locks,
request in self.anchor_locks,
)
reason = None if covered_tokens == f.num_tokens else "shrunk"
return _drop(reason)
@@ -1047,7 +1051,7 @@ class BufferModePipeline:
f"staged splice trim not page-aligned req={req.rid}: "
f"matched={f.matched_len} splice_base={splice_base}"
)
cache._resolve_storage_prefetch_tokens(req.rid, trim_tokens)
cache._resolve_storage_prefetch_tokens(request, trim_tokens)
key = RadixKey(
array("q", f.key_tokens),
@@ -1075,7 +1079,7 @@ class BufferModePipeline:
len(live.device_indices),
live.full_kv_hit_length,
f.num_tokens,
req.rid in self.anchor_locks,
request in self.anchor_locks,
)
available_end = min(
span_end,
@@ -1083,7 +1087,7 @@ class BufferModePipeline:
live.full_kv_hit_length,
)
available_overlap = max(0, available_end - splice_base)
cache._resolve_storage_prefetch_tokens(req.rid, available_overlap)
cache._resolve_storage_prefetch_tokens(request, available_overlap)
return _drop(None if available_overlap == splice_tokens else "shrunk")
# Evict-before-alloc (mirrors _load_back_transfers): the budget gate
@@ -1145,7 +1149,7 @@ class BufferModePipeline:
)
)
self.ongoing_buffer_load_back[load_back_id] = _OngoingBufferLoadBack(
req_id=f.req_id,
request=f.request,
num_tokens=splice_tokens,
occupied_tokens=f.occupied_tokens,
aux_xfers=f.aux_xfers,
@@ -1155,7 +1159,7 @@ class BufferModePipeline:
hash_values=f.hash_values,
)
m = cache.match_prefix(MatchPrefixParams(key=key))
self.release_anchor_lock(req.rid)
self.release_anchor_lock(request)
canonical = m.device_indices[splice_base:span_end]
if len(m.device_indices) < span_end or not torch.equal(
canonical, device_indices
@@ -1163,7 +1167,7 @@ class BufferModePipeline:
# Fail-stop: the insert freed or replaced slots the in-flight H2D
# still targets; continuing risks silent KV corruption.
raise RuntimeError(
f"HiCache buffer load-back ownership violation req={f.req_id}: "
f"HiCache buffer load-back ownership violation req={f.request.rid}: "
f"insert prefix_len={insert_result.prefix_len} "
f"expected={splice_base}, adopted={len(m.device_indices)} "
f"span_end={span_end}, canonical_matches_incoming="
@@ -1191,26 +1195,28 @@ class BufferModePipeline:
cc.prefetch_tokens_occupied -= f.occupied_tokens
logger.info(
"HiCache prefetch fill committed req=%s filled=%d occupied=%d locked=%d",
f.req_id,
f.request.rid,
f.num_tokens,
cc.prefetch_tokens_occupied,
self.anchor_locked_tokens_,
)
cache._finish_storage_prefetch(
f.req_id, fulfilled_tokens=f.num_tokens, reason=None
f.request, fulfilled_tokens=f.num_tokens, reason=None
)
return True
def release_staged_hold(self, rid: str, reason: Optional[str] = None) -> bool:
def release_staged_hold(
self, request: CacheRequestHandle, reason: Optional[str] = None
) -> bool:
"""Free a staged hold outright — anchor pin, host bounce (KV + aux),
occupancy grant; nothing device-side exists yet. Called for aborts
and for holds that can no longer splice. Returns True when a hold
existed."""
self.release_anchor_lock(rid)
staged = self.staged_prefetches.pop(rid, None)
self.release_anchor_lock(request)
staged = self.staged_prefetches.pop(request, None)
if staged is None:
return False
self._cache._finish_storage_prefetch(rid, fulfilled_tokens=0, reason=reason)
self._cache._finish_storage_prefetch(request, fulfilled_tokens=0, reason=reason)
self._free_staging_now(staged.host_indices, staged.aux_xfers)
self._cache.cache_controller.prefetch_tokens_occupied -= staged.occupied_tokens
return True
python/sglang/srt/mem_cache/hiradix_cache.py
+25 -8
View File
@@ -16,6 +16,7 @@ from sglang.srt.disaggregation.kv_events import StorageMedium
from sglang.srt.distributed.communication_tags import P2PTag
from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation
from sglang.srt.mem_cache.base_prefix_cache import (
CacheRequestHandle,
DecLockRefParams,
DecLockRefResult,
EvictParams,
@@ -1511,8 +1512,13 @@ class HiRadixCache(RadixCache):
extra_kwargs = {}
if prefetch_op_cls is HybridPrefetchOperation:
extra_kwargs["pool_transfers"] = self._get_extra_pools().get("extra_pools")
request = (
CacheRequestHandle("__storage_hit_query__", 0)
if prefetch_op_cls is HybridPrefetchOperation
else "__storage_hit_query__"
)
operation = prefetch_op_cls(
"__storage_hit_query__",
request,
prefetch_key,
last_hash,
prefix_keys,
@@ -1626,7 +1632,11 @@ class HiRadixCache(RadixCache):
0, cc.prefetch_tokens_occupied - len(prefetch_key)
)
def check_prefetch_progress(self, req_id: str) -> bool:
def has_ongoing_prefetch(self, handle: CacheRequestHandle) -> bool:
return handle.rid in self.ongoing_prefetch
def check_prefetch_progress(self, handle: CacheRequestHandle) -> bool:
req_id = handle.rid
if req_id not in self.ongoing_prefetch:
# there is no ongoing prefetch for this request or it has been revoked
return True
@@ -1720,15 +1730,15 @@ class HiRadixCache(RadixCache):
usable_pages = min(usable_pages, *pool_hit_pages)
return usable_pages * self.page_size
def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
def pop_prefetch_loaded_tokens(self, handle: CacheRequestHandle) -> int:
"""
Pop and return the number of tokens loaded from storage for a request.
Returns 0 if no prefetch was done or was revoked.
This should be called after check_prefetch_progress() returns True.
"""
return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
return self.prefetch_loaded_tokens_by_reqid.pop(handle.rid, 0)
def pop_storage_prefetch_miss(self, req_id: str) -> bool:
def pop_storage_prefetch_miss(self, handle: CacheRequestHandle) -> bool:
"""Storage prefetch miss markers are not tracked on the dense path;
the scheduler's paced availability-check retry is inert here."""
return False
@@ -1768,7 +1778,7 @@ class HiRadixCache(RadixCache):
def prefetch_from_storage(
self,
req_id: str,
handle: CacheRequestHandle,
last_host_node: TreeNode,
new_input_tokens: List[int],
last_hash: Optional[str] = None,
@@ -1778,6 +1788,7 @@ class HiRadixCache(RadixCache):
extra_key: Optional[str] = None,
cache_salt: Optional[str] = None,
):
req_id = handle.rid
prefetch_key = RadixKey(
new_input_tokens,
extra_key=extra_key,
@@ -1798,8 +1809,13 @@ class HiRadixCache(RadixCache):
# NOTE: host_indices is no longer pre-allocated here. It is allocated
# lazily in _drain_and_alloc_storage_hit() once the L3 storage hit count is known,
# so we only reserve host memory for pages that actually hit.
request = (
handle
if isinstance(self.cache_controller, HybridCacheController)
else req_id
)
operation = self.cache_controller.prefetch(
req_id,
request,
prefetch_key,
last_hash,
prefix_keys,
@@ -1999,7 +2015,8 @@ class HiRadixCache(RadixCache):
self._inc_hit_count(new_node, chunked)
return InsertResult(prefix_len=total_prefix_length)
def release_aborted_request(self, rid: str):
def release_aborted_request(self, handle: CacheRequestHandle):
rid = handle.rid
# Clean up storage hit tracking for aborted request
self.prefetch_loaded_tokens_by_reqid.pop(rid, None)
python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py
+6 -4
View File
@@ -24,6 +24,7 @@ from sglang.srt.managers.cache_controller import (
from sglang.srt.managers.cache_controller import (
StorageOperation as BaseStorageOperation,
)
from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle
from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorageExtraInfo,
PoolHitPolicy,
@@ -62,13 +63,14 @@ class StorageOperation(BaseStorageOperation):
class PrefetchOperation(StorageOperation):
def __init__(
self,
request_id: str,
handle: CacheRequestHandle,
token_ids: List[int],
last_hash: Optional[str] = None,
prefix_keys: Optional[List[str]] = None,
pool_transfers: Optional[list[PoolTransfer]] = None,
):
self.request_id = request_id
self.handle = handle
self.request_id = handle.rid
self._lock = threading.Lock()
self._terminated_flag = False
self.storage_hit_count = 0
@@ -543,14 +545,14 @@ class HybridCacheController(BaseHiCacheController):
def prefetch(
self,
request_id: str,
handle: CacheRequestHandle,
new_input_tokens: List[int],
last_hash: Optional[str] = None,
prefix_keys: Optional[List[str]] = None,
extra_pools: Optional[list[PoolTransfer]] = None,
) -> PrefetchOperation:
operation = PrefetchOperation(
request_id,
handle,
new_input_tokens,
last_hash,
prefix_keys=prefix_keys,
python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py
+55 -48
View File
@@ -35,6 +35,7 @@ from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import torch
from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle
from sglang.srt.mem_cache.storage.flexkv.flexkv_comm import (
CMD_LAYERWISE,
CMD_PUT_META,
@@ -219,14 +220,20 @@ class FlexKVConnector:
# 10. Per-rank in-flight tracking.
# Loads
self._pending_lookups: Dict[str, int] = {} # rid -> fkv_task_id
self._inflight_loads: Dict[int, int] = {} # producer_id -> rid hashlike
self._pending_lookups: Dict[
CacheRequestHandle, int
] = {} # handle -> fkv_task_id
self._inflight_loads: Dict[int, int] = {} # producer_id -> task id
self._completed_layerwise: List[int] = []
self._launched_load_tids: List[int] = [] # leader-only, for periodic drain
# Stores
self._inflight_stores: Dict[str, int] = {} # rid -> fkv_task_id
self._inflight_stores: Dict[
CacheRequestHandle, int
] = {} # handle -> fkv_task_id
# Prefetches
self._ongoing_prefetches: Dict[str, int] = {} # rid -> fkv_task_id
self._ongoing_prefetches: Dict[
CacheRequestHandle, int
] = {} # handle -> fkv_task_id
self._prefetch_enabled = bool(
self.cache_config.enable_ssd
or self.cache_config.enable_remote
@@ -248,7 +255,7 @@ class FlexKVConnector:
self,
token_ids: List[int],
token_mask: torch.Tensor,
rid: Optional[str] = None,
handle: Optional[CacheRequestHandle] = None,
) -> Tuple[int, int]:
"""Page-aligned prefix lookup against FlexKV.
@@ -256,8 +263,8 @@ class FlexKVConnector:
token_ids: full token id sequence we'd like to check.
token_mask: 1-D bool tensor or array, True for "this token is
*not* already on GPU and is a candidate for load-back".
rid: if set and hit > 0, the held FlexKV task id is stashed
under this key so a later ``retrieve_kv(rid, slots)`` call
handle: if set and hit > 0, the held FlexKV task id is stashed
under this key so a later ``retrieve_kv(handle, slots)`` call
can resolve it. If not set, the held task is cancelled when
hit > 0 and the caller didn't ask to track it.
@@ -304,29 +311,29 @@ class FlexKVConnector:
hit_length = aligned
# Decide what to do with the held task. Three cases:
# 1. hit_length > 0 and rid given → stash for retrieve_kv later.
# 2. hit_length > 0 and rid is None → cancel; caller can't use it.
# 1. hit_length > 0 and handle given → stash for retrieve_kv later.
# 2. hit_length > 0 and handle is None → cancel; caller can't use it.
# 3. hit_length == 0 → no work to do; FlexKV already marked the
# empty graph COMPLETED inside get_match, cancel would warn.
if hit_length > 0 and rid is not None and fkv_task_id >= 0:
self._pending_lookups[rid] = fkv_task_id
if hit_length > 0 and handle is not None and fkv_task_id >= 0:
self._pending_lookups[handle] = fkv_task_id
elif hit_length > 0 and fkv_task_id >= 0 and self._sync_ctx.is_sync_leader:
assert self.kv_manager is not None
self.kv_manager.cancel([fkv_task_id])
return fkv_task_id, hit_length
def release_pending(self, rid: str) -> None:
"""Cancel the task held by an earlier ``lookup_kv(rid=...)`` that
def release_pending(self, handle: CacheRequestHandle) -> None:
"""Cancel the task held by an earlier ``lookup_kv(handle=...)`` that
won't be followed by a ``retrieve_kv`` (e.g. allocation failed)."""
fkv_task_id = self._pending_lookups.pop(rid, -1)
fkv_task_id = self._pending_lookups.pop(handle, -1)
if fkv_task_id >= 0 and self._sync_ctx.is_sync_leader:
assert self.kv_manager is not None
self.kv_manager.cancel([fkv_task_id])
def retrieve_kv(
self,
rid: str,
handle: CacheRequestHandle,
slot_mapping: torch.Tensor,
) -> int:
"""Synchronous load: ``launch`` + ``wait``.
@@ -335,7 +342,7 @@ class FlexKVConnector:
responsible for having allocated ``slot_mapping`` of length
equal to ``hit_length`` from a prior ``lookup_kv``.
"""
fkv_task_id = self._pending_lookups.pop(rid, -1)
fkv_task_id = self._pending_lookups.pop(handle, -1)
if fkv_task_id < 0:
return 0
@@ -371,7 +378,7 @@ class FlexKVConnector:
def start_load_kv_layerwise(
self,
rid: str,
handle: CacheRequestHandle,
slot_mapping: torch.Tensor,
) -> Tuple[int, int]:
"""Layerwise load. Fires ``launch(layerwise_transfer=True)`` and
@@ -382,7 +389,7 @@ class FlexKVConnector:
"start_load_kv_layerwise called but layerwise transfer is "
"disabled. Set FLEXKV_ENABLE_LAYERWISE_TRANSFER=1."
)
fkv_task_id = self._pending_lookups.pop(rid, -1)
fkv_task_id = self._pending_lookups.pop(handle, -1)
if fkv_task_id < 0:
return 0, -1
@@ -450,7 +457,7 @@ class FlexKVConnector:
def store_kv(
self,
rid: str,
handle: CacheRequestHandle,
token_ids: List[int],
kv_indices: torch.Tensor,
) -> int:
@@ -507,7 +514,7 @@ class FlexKVConnector:
as_batch=False,
layerwise_transfer=False,
)
self._inflight_stores[rid] = fkv_task_id
self._inflight_stores[handle] = fkv_task_id
return fkv_task_id
return -1
@@ -530,28 +537,28 @@ class FlexKVConnector:
filtered = kv_indices[unmatched_mask]
slot_mapping_cpu = self._to_cpu_int64(filtered)
self._send_slot_mapping_to_remote(fkv_task_id, slot_mapping_cpu)
self._inflight_stores[rid] = fkv_task_id
self._inflight_stores[handle] = fkv_task_id
return fkv_task_id
def check_completed_stores(self) -> List[str]:
"""Return rids whose stores have completed since the last call."""
completed_rids: List[str] = []
def check_completed_stores(self) -> List[CacheRequestHandle]:
"""Return request handles whose stores have completed since the last call."""
completed_handles: List[CacheRequestHandle] = []
completed_dict: Dict[int, Any] = {}
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
if self._inflight_stores:
fk_to_rid = {v: k for k, v in self._inflight_stores.items()}
fk_to_handle = {v: k for k, v in self._inflight_stores.items()}
try:
completed_dict = self.kv_manager.try_wait(
task_ids=list(fk_to_rid.keys())
task_ids=list(fk_to_handle.keys())
)
except Exception as exc: # noqa: BLE001
logger.debug("[FlexKV] check_completed_stores: %s", exc)
completed_dict = {}
for fk_tid in completed_dict:
rid = fk_to_rid[fk_tid]
completed_rids.append(rid)
self._inflight_stores.pop(rid, None)
handle = fk_to_handle[fk_tid]
completed_handles.append(handle)
self._inflight_stores.pop(handle, None)
if self._sync_ctx.is_pp_sender:
self._sync_ctx.scatter_pp(
@@ -569,20 +576,20 @@ class FlexKVConnector:
)
fk_ids = payload.get("completed_fk_ids", [])
if fk_ids and self._inflight_stores:
fk_to_rid = {v: k for k, v in self._inflight_stores.items()}
fk_to_handle = {v: k for k, v in self._inflight_stores.items()}
for fk_tid in fk_ids:
if fk_tid in fk_to_rid:
rid = fk_to_rid[fk_tid]
completed_rids.append(rid)
self._inflight_stores.pop(rid, None)
if fk_tid in fk_to_handle:
handle = fk_to_handle[fk_tid]
completed_handles.append(handle)
self._inflight_stores.pop(handle, None)
if self._sync_ctx.needs_sync:
completed_rids = self._sync_ctx.scatter(completed_rids)
return completed_rids
completed_handles = self._sync_ctx.scatter(completed_handles)
return completed_handles
def wait_store(self, rid: str, timeout: float = 30.0) -> bool:
"""Block until a single store task identified by ``rid`` finishes."""
fkv_task_id = self._inflight_stores.pop(rid, -1)
def wait_store(self, handle: CacheRequestHandle, timeout: float = 30.0) -> bool:
"""Block until a single store task identified by ``handle`` finishes."""
fkv_task_id = self._inflight_stores.pop(handle, -1)
if fkv_task_id < 0:
return True
if not self._sync_ctx.is_sync_leader or self.kv_manager is None:
@@ -600,8 +607,8 @@ class FlexKVConnector:
# Public API — prefetch
# ------------------------------------------------------------------
def prefetch_async(self, rid: str, token_ids: List[int]) -> int:
if not self._prefetch_enabled or not rid:
def prefetch_async(self, handle: CacheRequestHandle, token_ids: List[int]) -> int:
if not self._prefetch_enabled:
return -1
task_id = -1
if self._sync_ctx.is_sync_leader and self.kv_manager is not None:
@@ -616,13 +623,13 @@ class FlexKVConnector:
payload = self._sync_ctx.scatter({"task_id": task_id})
task_id = payload["task_id"]
if task_id >= 0:
self._ongoing_prefetches[rid] = task_id
self._ongoing_prefetches[handle] = task_id
return task_id
def check_prefetch_progress(self, rid: str) -> bool:
def check_prefetch_progress(self, handle: CacheRequestHandle) -> bool:
if not self._prefetch_enabled:
return True
task_id = self._ongoing_prefetches.get(rid, -1)
task_id = self._ongoing_prefetches.get(handle, -1)
if task_id < 0:
return True
done = False
@@ -637,14 +644,14 @@ class FlexKVConnector:
payload = self._sync_ctx.scatter({"done": done})
done = payload["done"]
if done:
self._ongoing_prefetches.pop(rid, None)
self._ongoing_prefetches.pop(handle, None)
return done
def cancel_prefetch(self, rid: str) -> None:
self._pending_lookups.pop(rid, None)
def cancel_prefetch(self, handle: CacheRequestHandle) -> None:
self._pending_lookups.pop(handle, None)
# FlexKV doesn't currently support prefetch cancellation, but
# we still drop our tracking entry.
self._ongoing_prefetches.pop(rid, None)
self._ongoing_prefetches.pop(handle, None)
# ------------------------------------------------------------------
# Layerwise transfer hooks
python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py
+33 -32
View File
@@ -34,6 +34,7 @@ from typing import TYPE_CHECKING, Optional, Tuple
import torch
from sglang.srt.mem_cache.base_prefix_cache import (
CacheRequestHandle,
EvictParams,
EvictResult,
InitLoadBackParams,
@@ -123,11 +124,11 @@ class FlexKVRadixCache(RadixCache):
# Two-phase MP load: stash marker between ``match_prefix`` and
# ``init_load_back``.
self._load_markers: dict[str, _LoadBackMarker] = {}
self._load_markers: dict[CacheRequestHandle, _LoadBackMarker] = {}
# ``store_kv`` is async — we keep a lock on the source node
# until FlexKV signals completion, draining in ``evict`` /
# ``check_hicache_events``.
self._inflight_store_nodes: dict[str, TreeNode] = {}
self._inflight_store_nodes: dict[CacheRequestHandle, TreeNode] = {}
self._node_lock = threading.Lock()
# ------------------------------------------------------------------
@@ -205,7 +206,7 @@ class FlexKVRadixCache(RadixCache):
token_mask[device_len:] = True
fkv_task_id, hit = self.flexkv_connector.lookup_kv(
token_ids=token_ids, token_mask=token_mask, rid=req.rid
token_ids=token_ids, token_mask=token_mask, handle=req.cache_request_handle
)
if hit <= 0:
return base_res
@@ -215,7 +216,7 @@ class FlexKVRadixCache(RadixCache):
token_ids_snap = token_ids[:]
else:
token_ids_snap = token_ids
self._load_markers[req.rid] = _LoadBackMarker(
self._load_markers[req.cache_request_handle] = _LoadBackMarker(
key=RadixKey(
token_ids_snap,
key.extra_key,
@@ -249,10 +250,10 @@ class FlexKVRadixCache(RadixCache):
# Quick LOOKUP first to discover how many slots we'd need.
token_mask = torch.zeros(len(token_ids), dtype=torch.bool)
token_mask[device_len:] = True
# No rid here — IP mode self-pops; pass a synthetic stable key.
synthetic_rid = f"_ip_{id(key)}"
# No handle here — IP mode self-pops; pass a synthetic stable key.
synthetic_handle = CacheRequestHandle(f"_ip_{id(key)}", 0)
_, hit = self.flexkv_connector.lookup_kv(
token_ids=token_ids, token_mask=token_mask, rid=synthetic_rid
token_ids=token_ids, token_mask=token_mask, handle=synthetic_handle
)
if hit <= 0:
return base_res
@@ -263,7 +264,7 @@ class FlexKVRadixCache(RadixCache):
uncached_len=hit,
last_node=last_node,
load_fn=lambda slot_mapping: self.flexkv_connector.start_load_kv_layerwise(
synthetic_rid, slot_mapping
synthetic_handle, slot_mapping
)[0],
)
if result is None:
@@ -288,12 +289,12 @@ class FlexKVRadixCache(RadixCache):
load; inserts the resulting TreeNode."""
req = params.req
last_node: TreeNode = params.best_match_node
marker = self._load_markers.pop(req.rid, None)
marker = self._load_markers.pop(req.cache_request_handle, None)
if marker is None:
# ``match_prefix`` decided there was no work to do, but the
# scheduler still called us. Release any held task and
# return an empty load.
self.flexkv_connector.release_pending(req.rid)
self.flexkv_connector.release_pending(req.cache_request_handle)
return (
torch.empty((0,), dtype=torch.int64, device=self.device),
last_node,
@@ -305,7 +306,7 @@ class FlexKVRadixCache(RadixCache):
uncached_len=params.host_hit_length,
last_node=last_node,
load_fn=lambda slot_mapping: self.flexkv_connector.retrieve_kv(
req.rid, slot_mapping
req.cache_request_handle, slot_mapping
),
)
if result is None:
@@ -313,7 +314,7 @@ class FlexKVRadixCache(RadixCache):
# already cancels/cleans up on failure paths; release_pending
# is idempotent for the case where allocation failed before
# we even popped the held task.
self.flexkv_connector.release_pending(req.rid)
self.flexkv_connector.release_pending(req.cache_request_handle)
return (
torch.empty((0,), dtype=torch.int64, device=self.device),
last_node,
@@ -392,7 +393,7 @@ class FlexKVRadixCache(RadixCache):
req, is_insert=is_insert, kv_len_to_handle=kv_len_to_handle
)
if not is_insert:
self._load_markers.pop(req.rid, None)
self._load_markers.pop(req.cache_request_handle, None)
return
# Compute the committed prefix mirroring LMCRadixCache's logic.
@@ -431,7 +432,7 @@ class FlexKVRadixCache(RadixCache):
try:
with torch.cuda.stream(self.store_stream):
fkv_task_id = self.flexkv_connector.store_kv(
rid=req.rid,
handle=req.cache_request_handle,
token_ids=list(token_ids),
kv_indices=kv_indices,
)
@@ -446,7 +447,7 @@ class FlexKVRadixCache(RadixCache):
return
with self._node_lock:
self._inflight_store_nodes[req.rid] = new_last_node
self._inflight_store_nodes[req.cache_request_handle] = new_last_node
# ------------------------------------------------------------------
# evict + completion draining
@@ -473,12 +474,12 @@ class FlexKVRadixCache(RadixCache):
self.flexkv_connector.drain_launched_loads()
def _drain_completed_stores(self) -> None:
completed_rids = self.flexkv_connector.check_completed_stores()
if not completed_rids:
completed_handles = self.flexkv_connector.check_completed_stores()
if not completed_handles:
return
with self._node_lock:
for rid in completed_rids:
node = self._inflight_store_nodes.pop(rid, None)
for handle in completed_handles:
node = self._inflight_store_nodes.pop(handle, None)
if node is not None:
self.dec_lock_ref(node)
@@ -486,33 +487,33 @@ class FlexKVRadixCache(RadixCache):
# Optional pass-throughs used by the scheduler
# ------------------------------------------------------------------
def release_aborted_request(self, rid: str) -> None:
def release_aborted_request(self, handle: CacheRequestHandle) -> None:
"""Clean up tracking for an aborted request without invoking FlexKV."""
self._load_markers.pop(rid, None)
self._load_markers.pop(handle, None)
with self._node_lock:
node = self._inflight_store_nodes.pop(rid, None)
node = self._inflight_store_nodes.pop(handle, None)
if node is not None:
self.dec_lock_ref(node)
self.flexkv_connector.release_pending(rid)
self.flexkv_connector.cancel_prefetch(rid)
self.flexkv_connector.release_pending(handle)
self.flexkv_connector.cancel_prefetch(handle)
def prefetch_from_storage(
self, rid: str, last_host_node: TreeNode, token_ids
self, handle: CacheRequestHandle, last_host_node: TreeNode, token_ids
) -> None:
"""Kick off an opportunistic prefetch (SSD/Remote → CPU)."""
try:
self.flexkv_connector.prefetch_async(rid, list(token_ids))
self.flexkv_connector.prefetch_async(handle, list(token_ids))
except Exception as exc: # noqa: BLE001
logger.debug("[FlexKV] prefetch_from_storage: %s", exc)
def check_prefetch_progress(self, rid: str) -> bool:
return self.flexkv_connector.check_prefetch_progress(rid)
def check_prefetch_progress(self, handle: CacheRequestHandle) -> bool:
return self.flexkv_connector.check_prefetch_progress(handle)
def terminate_prefetch(self, rid: str) -> None:
self.flexkv_connector.cancel_prefetch(rid)
def terminate_prefetch(self, handle: CacheRequestHandle) -> None:
self.flexkv_connector.cancel_prefetch(handle)
def pop_prefetch_loaded_tokens(self, rid: str) -> int:
# FlexKV doesn't expose per-rid prefetched token counts yet.
def pop_prefetch_loaded_tokens(self, handle: CacheRequestHandle) -> int:
# FlexKV doesn't expose per-handle prefetched token counts yet.
return 0
@property
python/sglang/srt/mem_cache/unified_cache/storage_attachment.py
+11 -11
View File
@@ -364,23 +364,23 @@ class StorageAttachment:
cache = self._cache
controller = cache.cache_controller
for req_id in list(cache.ongoing_prefetch):
info = cache.ongoing_prefetch[req_id]
for handle in list(cache.ongoing_prefetch):
info = cache.ongoing_prefetch[handle]
try:
cache.discard_storage_prefetch_accounting(req_id)
cache.discard_storage_prefetch_accounting(handle)
if info.host_indices is None:
# Host pages were never allocated for this operation.
cache.revoke_pending_prefetch(req_id)
cache.revoke_pending_prefetch(handle)
continue
completed_tokens, _ = controller.terminate_prefetch(info.operation)
del cache.ongoing_prefetch[req_id]
del cache.ongoing_prefetch[handle]
if info.anchor_lock_params is not None:
cache.dec_host_lock_ref(
info.anchor_node_id, info.anchor_lock_params
)
if cache.buffer_pipeline is not None:
cache.buffer_pipeline.pop_prefix_ctx(req_id)
cache.buffer_pipeline.release_anchor_lock(req_id)
cache.buffer_pipeline.pop_prefix_ctx(handle)
cache.buffer_pipeline.release_anchor_lock(handle)
controller.append_host_mem_release(
host_indices=info.host_indices[:completed_tokens],
extra_pools=[
@@ -395,8 +395,8 @@ class StorageAttachment:
),
)
except Exception:
logger.exception("Failed to release pending prefetch %s", req_id)
cache.ongoing_prefetch.pop(req_id, None)
logger.exception("Failed to release pending prefetch %s", handle.rid)
cache.ongoing_prefetch.pop(handle, None)
for ack_id in list(cache.ongoing_backup):
node_id, lock_params = cache.ongoing_backup.pop(ack_id)
@@ -405,7 +405,7 @@ class StorageAttachment:
except Exception:
logger.exception("Failed to release host lock for backup op %s", ack_id)
for req_id in list(cache._storage_prefetch_hit_remaining_by_reqid):
cache.discard_storage_prefetch_accounting(req_id)
for handle in list(cache._storage_prefetch_hit_remaining_by_reqid):
cache.discard_storage_prefetch_accounting(handle)
cache.prefetch_loaded_tokens_by_reqid.clear()
cache.prefetch_loaded_storage_start_by_reqid.clear()
python/sglang/srt/mem_cache/unified_radix_cache.py
+136 -129
View File
@@ -18,6 +18,7 @@ from sglang.srt.mem_cache.allocator.page_interleave import (
)
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
CacheRequestHandle,
DecLockRefParams,
DecLockRefResult,
EvictParams,
@@ -40,6 +41,7 @@ from sglang.srt.mem_cache.common import RetractionBackup
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer, SidecarPoolSpec
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
HybridCacheController,
PrefetchOperation,
)
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
from sglang.srt.mem_cache.radix_cache import RadixKey
@@ -381,16 +383,18 @@ class UnifiedRadixCache(BasePrefixCache):
self.ongoing_write_through: dict[int, _OngoingWriteThrough] = {}
self.ongoing_load_back: dict[int, _OngoingLoadBack] = {}
self.enable_storage = False
self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
self.prefetch_loaded_storage_start_by_reqid: dict[str, int] = {}
self.ongoing_prefetch: dict[str, _OngoingPrefetch] = {}
self.prefetch_loaded_tokens_by_reqid: dict[CacheRequestHandle, int] = {}
self.prefetch_loaded_storage_start_by_reqid: dict[CacheRequestHandle, int] = {}
self.ongoing_prefetch: dict[CacheRequestHandle, _OngoingPrefetch] = {}
# Rank-agreed L3-hit tokens not yet resolved as usable or unfulfilled.
# Cache-mode entries survive L3->L2 until H2D succeeds or admission
# fails; buffer-mode entries survive staging until the H2D ack.
self._storage_prefetch_hit_remaining_by_reqid: dict[str, int] = {}
# Rids whose storage prefetch resolved without a usable result;
self._storage_prefetch_hit_remaining_by_reqid: dict[
CacheRequestHandle, int
] = {}
# Attempts whose storage prefetch resolved without a usable result;
# popped by the scheduler to pace availability-check retries.
self._storage_prefetch_missed_rids: set[str] = set()
self._storage_prefetch_missed_rids: set[CacheRequestHandle] = set()
self.ongoing_backup: dict[int, tuple[NodeId, DecLockRefParams]] = {}
if self.buffer_pipeline is not None:
self.buffer_pipeline.reset()
@@ -1881,12 +1885,8 @@ class UnifiedRadixCache(BasePrefixCache):
if len(prefetch_key) < self.prefetch_threshold:
return 0
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
PrefetchOperation,
)
operation = PrefetchOperation(
"__storage_hit_query__",
CacheRequestHandle("__storage_hit_query__", 0),
prefetch_key,
last_hash,
prefix_keys,
@@ -1903,7 +1903,7 @@ class UnifiedRadixCache(BasePrefixCache):
@rank_consensus(same_params=["req_id", "len(new_input_tokens)"])
def prefetch_from_storage(
self,
req_id: str,
request: CacheRequestHandle,
last_host_node_id: NodeId,
new_input_tokens: list[int],
last_hash: Optional[str] = None,
@@ -1943,16 +1943,16 @@ class UnifiedRadixCache(BasePrefixCache):
stats["declined_too_short"] += 1
# A too-short/fully-matched suffix can become a full recompute if
# the device match evicts while queued; arm the paced retry.
self._storage_prefetch_missed_rids.add(req_id)
self._storage_prefetch_missed_rids.add(request)
return
if not buffer_mode and self.cache_controller.prefetch_rate_limited():
stats["declined_rate_limited"] += 1
self._storage_prefetch_missed_rids.add(req_id)
self._storage_prefetch_missed_rids.add(request)
return
if req_id in self.ongoing_prefetch or (
buffer_mode and self.buffer_pipeline.has_staged(req_id)
if request in self.ongoing_prefetch or (
buffer_mode and self.buffer_pipeline.has_staged(request)
):
# A fetch (or an unconsumed hold) already exists for this rid;
# A fetch (or an unconsumed hold) already exists for this attempt;
# overwriting would leak its staging slots.
return
@@ -2009,13 +2009,13 @@ class UnifiedRadixCache(BasePrefixCache):
if anchor_lock_params is not None:
self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
# Forfeited over transient staging pressure; retryable.
self._storage_prefetch_missed_rids.add(req_id)
self._storage_prefetch_missed_rids.add(request)
return
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
aux_xfers.extend(sidecar_xfers)
operation = self.cache_controller.prefetch(
req_id,
request,
prefetch_key,
last_hash,
prefix_keys,
@@ -2026,7 +2026,7 @@ class UnifiedRadixCache(BasePrefixCache):
# rank-synchronized query outcome.
operation.stats_requested_tokens = prefetch_length
operation.storage_start = len(matched_prefix_tokens or [])
self.ongoing_prefetch[req_id] = _OngoingPrefetch(
self.ongoing_prefetch[request] = _OngoingPrefetch(
last_host_node_id,
prefetch_key,
None,
@@ -2036,7 +2036,7 @@ class UnifiedRadixCache(BasePrefixCache):
)
if buffer_mode:
self.buffer_pipeline.set_prefix_ctx(
req_id,
request,
matched_prefix_tokens,
extra_key=extra_key,
cache_salt=cache_salt,
@@ -2044,7 +2044,7 @@ class UnifiedRadixCache(BasePrefixCache):
# Pin the just-matched anchor now: deferred to IO commit it is
# often already deleted under churn. The IO-commit call remains
# as the second chance that decides the fetch's fate.
self.buffer_pipeline.try_lock_anchor(req_id)
self.buffer_pipeline.try_lock_anchor(request)
else:
# Cache mode reserves the requested span up front; buffer mode
# grants occupancy later at hit-alloc time, sized to the hit.
@@ -2081,12 +2081,15 @@ class UnifiedRadixCache(BasePrefixCache):
else:
return True
def has_ongoing_prefetch(self, handle: CacheRequestHandle) -> bool:
return handle in self.ongoing_prefetch
@rank_consensus(same_params=True, same_results=True)
def check_prefetch_progress(self, req_id: str) -> bool:
if req_id not in self.ongoing_prefetch:
def check_prefetch_progress(self, request: CacheRequestHandle) -> bool:
if request not in self.ongoing_prefetch:
return True
_, _, _, operation, _, _ = self.ongoing_prefetch[req_id]
_, _, _, operation, _, _ = self.ongoing_prefetch[request]
# Determine whether or not we should terminate this prefetch request.
should_terminate = operation.is_terminated() or self._can_terminate_prefetch(
@@ -2098,8 +2101,8 @@ class UnifiedRadixCache(BasePrefixCache):
self.cache_controller.terminate_prefetch(operation)
if operation.host_indices is None:
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
self._storage_prefetch_missed_rids.add(request)
self.revoke_pending_prefetch(request)
else:
self._handle_prefetch_result(operation)
return True
@@ -2112,7 +2115,7 @@ class UnifiedRadixCache(BasePrefixCache):
# That is, when this function returns the host memory referenced must be inserted
# into the radix tree or released to pool.
req_id = operation.request_id
request = operation.handle
completed_tokens = operation.completed_tokens
hash_value = operation.hash_value
@@ -2123,12 +2126,12 @@ class UnifiedRadixCache(BasePrefixCache):
_,
anchor_lock_params,
comp_xfers,
) = self.ongoing_prefetch[req_id]
) = self.ongoing_prefetch[request]
# All PP/TP ranks will get the same `min_completed_tokens`, because `completed_tokens`
# and `pool_hits` in their operations are same. No need to sync cross-rank here.
if not self._check_hybrid_prefetch_result(
req_id,
request,
operation,
completed_tokens,
hash_value,
@@ -2143,7 +2146,7 @@ class UnifiedRadixCache(BasePrefixCache):
allocated_tokens = len(host_indices)
if completed_tokens < allocated_tokens:
self._resolve_storage_prefetch_tokens(
req_id,
request,
allocated_tokens - completed_tokens,
reason="storage_transfer",
)
@@ -2158,7 +2161,7 @@ class UnifiedRadixCache(BasePrefixCache):
# No graft: release the rank-local tail beyond the synced usable
# length, then park the bounce for admission-time consumption.
return self.buffer_pipeline.stage_completed_prefetch(
req_id, completed_tokens, hash_value
request, completed_tokens, hash_value
)
fetched_key = prefetch_key[:completed_tokens]
@@ -2173,8 +2176,8 @@ class UnifiedRadixCache(BasePrefixCache):
self._apply_cache_actions(insert_result.cache_actions)
if insert_result.host_insert_dropped:
self._resolve_storage_prefetch_tokens(req_id, insert_result.prefix_len)
self._finish_storage_prefetch(req_id, fulfilled_tokens=0, reason="dropped")
self._resolve_storage_prefetch_tokens(request, insert_result.prefix_len)
self._finish_storage_prefetch(request, fulfilled_tokens=0, reason="dropped")
self.cache_controller.append_host_mem_release(
host_indices=host_indices[:completed_tokens],
extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
@@ -2201,23 +2204,23 @@ class UnifiedRadixCache(BasePrefixCache):
# Cache mode has only completed L3 -> L2 here. Keep the usable
# storage span unresolved until admission proves that L2 -> L1
# load-back actually materialized it for this request.
self._resolve_storage_prefetch_tokens(req_id, insert_result.prefix_len)
self._resolve_storage_prefetch_tokens(request, insert_result.prefix_len)
self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
del self.ongoing_prefetch[req_id]
del self.ongoing_prefetch[request]
self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)
self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage
self.prefetch_loaded_tokens_by_reqid[request] = loaded_from_storage
if loaded_from_storage > 0:
self.prefetch_loaded_storage_start_by_reqid[req_id] = (
self.prefetch_loaded_storage_start_by_reqid[request] = (
operation.storage_start + insert_result.prefix_len
)
else:
self.prefetch_loaded_storage_start_by_reqid.pop(req_id, None)
self.prefetch_loaded_storage_start_by_reqid.pop(request, None)
logger.info(
"HiCache prefetch %s req=%s completed=%d matched=%d loaded=%d occupied=%d",
"dropped" if insert_result.host_insert_dropped else "success",
req_id,
request.rid,
completed_tokens,
insert_result.prefix_len,
loaded_from_storage,
@@ -2227,7 +2230,7 @@ class UnifiedRadixCache(BasePrefixCache):
def _check_hybrid_prefetch_result(
self,
req_id: str,
request: CacheRequestHandle,
operation: PrefetchOperation,
completed_tokens: int,
hash_value: list[str],
@@ -2297,23 +2300,23 @@ class UnifiedRadixCache(BasePrefixCache):
extra_pools=pool_transfers if operation.pool_transfers_done else None,
)
self._finish_storage_prefetch(
req_id, fulfilled_tokens=0, reason="storage_transfer"
request, fulfilled_tokens=0, reason="storage_transfer"
)
if anchor_lock_params is not None:
self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
if self.buffer_pipeline is not None:
self.buffer_pipeline.pop_prefix_ctx(req_id)
self.buffer_pipeline.release_anchor_lock(req_id)
del self.ongoing_prefetch[req_id]
self.buffer_pipeline.pop_prefix_ctx(request)
self.buffer_pipeline.release_anchor_lock(request)
del self.ongoing_prefetch[request]
self.cache_controller.prefetch_tokens_occupied -= (
self._prefetch_occupied_span(prefetch_key, host_indices)
)
self.prefetch_loaded_tokens_by_reqid[req_id] = 0
self.prefetch_loaded_storage_start_by_reqid.pop(req_id, None)
self.prefetch_loaded_tokens_by_reqid[request] = 0
self.prefetch_loaded_storage_start_by_reqid.pop(request, None)
logger.warning(
"HiCache hybrid prefetch discarded req=%s completed=%d requested=%d "
"kv_beliefs_kept_pages=%d",
req_id,
request.rid,
completed_tokens,
expected_tokens,
keep_pages,
@@ -2321,7 +2324,9 @@ class UnifiedRadixCache(BasePrefixCache):
return False
return True
def _record_storage_prefetch_hit(self, req_id: str, num_tokens: int) -> None:
def _record_storage_prefetch_hit(
self, request: CacheRequestHandle, num_tokens: int
) -> None:
"""Start accounting one rank-agreed positive L3 query result."""
if (
num_tokens <= 0
@@ -2329,23 +2334,23 @@ class UnifiedRadixCache(BasePrefixCache):
or self.storage_metrics_collector is None
):
return
if req_id in self._storage_prefetch_hit_remaining_by_reqid:
if request in self._storage_prefetch_hit_remaining_by_reqid:
logger.warning(
"Replacing unresolved storage-hit accounting req=%s old=%d new=%d",
req_id,
self._storage_prefetch_hit_remaining_by_reqid[req_id],
request.rid,
self._storage_prefetch_hit_remaining_by_reqid[request],
num_tokens,
)
self.discard_storage_prefetch_accounting(req_id)
self._storage_prefetch_hit_remaining_by_reqid[req_id] = num_tokens
self.discard_storage_prefetch_accounting(request)
self._storage_prefetch_hit_remaining_by_reqid[request] = num_tokens
self.storage_metrics_collector.log_storage_prefetch_hit_tokens(num_tokens)
def _resolve_storage_prefetch_tokens(
self, req_id: str, num_tokens: int, reason: Optional[str] = None
self, request: CacheRequestHandle, num_tokens: int, reason: Optional[str] = None
) -> None:
if num_tokens <= 0:
return
remaining = self._storage_prefetch_hit_remaining_by_reqid.get(req_id)
remaining = self._storage_prefetch_hit_remaining_by_reqid.get(request)
if remaining is None:
return
dropped = min(num_tokens, remaining)
@@ -2353,7 +2358,7 @@ class UnifiedRadixCache(BasePrefixCache):
logger.warning(
"Storage-prefetch accounting exceeded remaining "
"tokens req=%s requested=%d remaining=%d reason=%s",
req_id,
request.rid,
num_tokens,
remaining,
reason,
@@ -2364,21 +2369,21 @@ class UnifiedRadixCache(BasePrefixCache):
)
remaining -= dropped
if remaining:
self._storage_prefetch_hit_remaining_by_reqid[req_id] = remaining
self._storage_prefetch_hit_remaining_by_reqid[request] = remaining
else:
self._storage_prefetch_hit_remaining_by_reqid.pop(req_id, None)
self._storage_prefetch_hit_remaining_by_reqid.pop(request, None)
def _finish_storage_prefetch(
self, req_id: str, fulfilled_tokens: int, reason: Optional[str]
self, request: CacheRequestHandle, fulfilled_tokens: int, reason: Optional[str]
) -> None:
remaining = self._storage_prefetch_hit_remaining_by_reqid.pop(req_id, None)
remaining = self._storage_prefetch_hit_remaining_by_reqid.pop(request, None)
if remaining is None:
return
if fulfilled_tokens > remaining:
logger.warning(
"Storage-prefetch fulfilled accounting exceeded remaining "
"tokens req=%s fulfilled=%d remaining=%d",
req_id,
request.rid,
fulfilled_tokens,
remaining,
)
@@ -2389,7 +2394,7 @@ class UnifiedRadixCache(BasePrefixCache):
)
def finish_storage_prefetch_admission(
self, req_id: str, fulfilled_tokens: int, reason: Optional[str]
self, request: CacheRequestHandle, fulfilled_tokens: int, reason: Optional[str]
) -> None:
"""Resolve a cache-mode L3 hit after request admission.
@@ -2397,70 +2402,72 @@ class UnifiedRadixCache(BasePrefixCache):
host allocation is transport staging rather than a resident L2 hit.
"""
if self.host_memory_mode == "cache":
self._finish_storage_prefetch(req_id, fulfilled_tokens, reason)
self._finish_storage_prefetch(request, fulfilled_tokens, reason)
def discard_storage_prefetch_accounting(self, req_id: str) -> None:
def discard_storage_prefetch_accounting(self, request: CacheRequestHandle) -> None:
"""Drop lifecycle state for cases intentionally excluded from metrics."""
self._storage_prefetch_hit_remaining_by_reqid.pop(req_id, None)
self._storage_prefetch_hit_remaining_by_reqid.pop(request, None)
def _handle_storage_prefetch_anchor_loss(self, req_id: str) -> None:
self._finish_storage_prefetch(req_id, fulfilled_tokens=0, reason="shrunk")
def _handle_storage_prefetch_anchor_loss(self, request: CacheRequestHandle) -> None:
self._finish_storage_prefetch(request, fulfilled_tokens=0, reason="shrunk")
# The span is still L3-resident; retry from the shorter live match.
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
self._storage_prefetch_missed_rids.add(request)
self.revoke_pending_prefetch(request)
def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
def pop_prefetch_loaded_tokens(self, request: CacheRequestHandle) -> int:
# The request is being scheduled; a still-unserved miss marker is moot.
self._storage_prefetch_missed_rids.discard(req_id)
self.prefetch_loaded_storage_start_by_reqid.pop(req_id, None)
return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
self._storage_prefetch_missed_rids.discard(request)
self.prefetch_loaded_storage_start_by_reqid.pop(request, None)
return self.prefetch_loaded_tokens_by_reqid.pop(request, 0)
def pop_prefetch_loaded_span(self, req_id: str) -> tuple[int, Optional[int]]:
def pop_prefetch_loaded_span(
self, request: CacheRequestHandle
) -> tuple[int, Optional[int]]:
"""Pop the loaded L3 token count and its absolute prefix start."""
self._storage_prefetch_missed_rids.discard(req_id)
self._storage_prefetch_missed_rids.discard(request)
return (
self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0),
self.prefetch_loaded_storage_start_by_reqid.pop(req_id, None),
self.prefetch_loaded_tokens_by_reqid.pop(request, 0),
self.prefetch_loaded_storage_start_by_reqid.pop(request, None),
)
def pop_storage_prefetch_miss(self, req_id: str) -> bool:
def pop_storage_prefetch_miss(self, request: CacheRequestHandle) -> bool:
"""True once per resolved storage-prefetch miss for a live request;
the scheduler uses it to arm the paced availability-check retry."""
if req_id in self._storage_prefetch_missed_rids:
self._storage_prefetch_missed_rids.discard(req_id)
if request in self._storage_prefetch_missed_rids:
self._storage_prefetch_missed_rids.discard(request)
return True
return False
def plan_staged_splice(
self, req_id: str, device_prefix_len: int
self, request: CacheRequestHandle, device_prefix_len: int
) -> tuple[int, int]:
"""(kv, swa) host-hit tokens a staged buffer-mode prefetch will splice
given the request's live device prefix; frees unusable holds."""
if self.buffer_pipeline is None:
return 0, 0
return self.buffer_pipeline.plan_staged_splice(req_id, device_prefix_len)
return self.buffer_pipeline.plan_staged_splice(request, device_prefix_len)
def staged_prefetch_swa_tokens(self, req_id: str) -> int:
def staged_prefetch_swa_tokens(self, request: CacheRequestHandle) -> int:
"""SWA device tokens consuming a staged buffer-mode prefetch will
allocate; surfaced as the request's swa_host_hit_length."""
if self.buffer_pipeline is None:
return 0
return self.buffer_pipeline.staged_prefetch_swa_tokens(req_id)
return self.buffer_pipeline.staged_prefetch_swa_tokens(request)
@rank_consensus(same_params=True)
def release_aborted_request(self, rid: str) -> None:
def release_aborted_request(self, request: CacheRequestHandle) -> None:
if self.linker is not None:
self.linker.release_request(rid)
self.prefetch_loaded_tokens_by_reqid.pop(rid, None)
self.prefetch_loaded_storage_start_by_reqid.pop(rid, None)
self._storage_prefetch_missed_rids.discard(rid)
self.linker.release_request(request.rid)
self.prefetch_loaded_tokens_by_reqid.pop(request, None)
self.prefetch_loaded_storage_start_by_reqid.pop(request, None)
self._storage_prefetch_missed_rids.discard(request)
if (
self.buffer_pipeline is not None
and self.buffer_pipeline.release_staged_hold(rid)
and self.buffer_pipeline.release_staged_hold(request)
):
return
self.discard_storage_prefetch_accounting(rid)
if rid not in self.ongoing_prefetch:
self.discard_storage_prefetch_accounting(request)
if request not in self.ongoing_prefetch:
return
(
@@ -2470,19 +2477,19 @@ class UnifiedRadixCache(BasePrefixCache):
operation,
anchor_lock_params,
comp_xfers,
) = self.ongoing_prefetch[rid]
) = self.ongoing_prefetch[request]
if operation.host_indices is None:
self.cache_controller.terminate_prefetch(operation)
self.revoke_pending_prefetch(rid)
self.revoke_pending_prefetch(request)
return
completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
if anchor_lock_params is not None:
self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
del self.ongoing_prefetch[rid]
del self.ongoing_prefetch[request]
if self.buffer_pipeline is not None:
self.buffer_pipeline.pop_prefix_ctx(rid)
self.buffer_pipeline.release_anchor_lock(rid)
self.buffer_pipeline.pop_prefix_ctx(request)
self.buffer_pipeline.release_anchor_lock(request)
pool_transfers = [x for xfers in comp_xfers.values() for x in xfers]
self.cache_controller.append_host_mem_release(
host_indices=host_indices[:completed_tokens],
@@ -2536,9 +2543,9 @@ class UnifiedRadixCache(BasePrefixCache):
return len(host_indices) if host_indices is not None else 0
return len(prefetch_key)
def revoke_pending_prefetch(self, req_id: str) -> None:
info = self.ongoing_prefetch.pop(req_id, None)
self._finish_storage_prefetch(req_id, fulfilled_tokens=0, reason="dropped")
def revoke_pending_prefetch(self, request: CacheRequestHandle) -> None:
info = self.ongoing_prefetch.pop(request, None)
self._finish_storage_prefetch(request, fulfilled_tokens=0, reason="dropped")
if info is None:
return
(
@@ -2551,8 +2558,8 @@ class UnifiedRadixCache(BasePrefixCache):
) = info
self._invalidate_absent_from_hit_query(operation)
if self.buffer_pipeline is not None:
self.buffer_pipeline.pop_prefix_ctx(req_id)
self.buffer_pipeline.release_anchor_lock(req_id)
self.buffer_pipeline.pop_prefix_ctx(request)
self.buffer_pipeline.release_anchor_lock(request)
cc = self.cache_controller
cc.append_host_mem_release(
extra_pools=[x for xfers in comp_xfers.values() for x in xfers]
@@ -2606,13 +2613,13 @@ class UnifiedRadixCache(BasePrefixCache):
"""Allocate the hit-sized bounce and launch the transfer.
Returns False when staging pressure defers the allocation
(buffer mode parks and retries; cache mode revokes)."""
req_id = operation.request_id
info = self.ongoing_prefetch.get(req_id)
request = operation.handle
info = self.ongoing_prefetch.get(request)
hit_tokens = operation.storage_hit_count
if info is None:
return True # aborted/cleaned; nothing to retry
if operation.is_terminated():
self.revoke_pending_prefetch(req_id)
self.revoke_pending_prefetch(request)
return True
if buffer_mode and cc.prefetch_rate_limited():
@@ -2624,20 +2631,20 @@ class UnifiedRadixCache(BasePrefixCache):
# IO commit: pin before the bounce alloc so a cancel is a
# plain revoke and a parked op keeps its pin; a fetch whose
# splice base is gone is not worth its storage read.
if self.buffer_pipeline.try_lock_anchor(req_id) == "anchor_lost":
if self.buffer_pipeline.try_lock_anchor(request) == "anchor_lost":
self._prefetch_outcome_stats["declined_anchor_lost"] += 1
self._handle_storage_prefetch_anchor_loss(req_id)
self._handle_storage_prefetch_anchor_loss(request)
return True
if self.buffer_pipeline.staged_span_covered(
req_id, operation.storage_hit_count
request, operation.storage_hit_count
):
# Live tree already covers the span: nothing left to
# splice, so skip the storage read.
self._prefetch_outcome_stats["declined_device_covered"] += 1
self._finish_storage_prefetch(
req_id, fulfilled_tokens=0, reason=None
request, fulfilled_tokens=0, reason=None
)
self.revoke_pending_prefetch(req_id)
self.revoke_pending_prefetch(request)
return True
alloc_len = hit_tokens
host_indices = cc.mem_pool_host.alloc(alloc_len)
@@ -2658,18 +2665,18 @@ class UnifiedRadixCache(BasePrefixCache):
if buffer_mode:
return False
self._finish_storage_prefetch(
req_id, fulfilled_tokens=0, reason="host_capacity"
request, fulfilled_tokens=0, reason="host_capacity"
)
self.revoke_pending_prefetch(req_id)
self.revoke_pending_prefetch(request)
return True
self._resolve_storage_prefetch_tokens(
req_id, hit_tokens - alloc_len, reason="host_capacity"
request, hit_tokens - alloc_len, reason="host_capacity"
)
operation.storage_hit_count = alloc_len
operation.hash_value = operation.hash_value[: alloc_len // self.page_size]
operation.host_indices = host_indices
self.ongoing_prefetch[req_id] = info._replace(host_indices=host_indices)
self.ongoing_prefetch[request] = info._replace(host_indices=host_indices)
if buffer_mode:
cc.prefetch_tokens_occupied += alloc_len
cc.prefetch_buffer.put(operation)
@@ -2684,24 +2691,24 @@ class UnifiedRadixCache(BasePrefixCache):
break
parked.popleft()
for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit):
req_id = operation.request_id
request = operation.handle
hit_tokens = operation.storage_hit_count
info = self.ongoing_prefetch.get(req_id)
info = self.ongoing_prefetch.get(request)
if info is None:
# Request already aborted/cleaned up; still flush the
# query's absent-hash feedback.
self._invalidate_absent_from_hit_query(operation)
if hit_tokens > 0:
self.discard_storage_prefetch_accounting(req_id)
self.discard_storage_prefetch_accounting(request)
continue
if hit_tokens > 0:
self._record_storage_prefetch_hit(req_id, hit_tokens)
self._record_storage_prefetch_hit(request, hit_tokens)
if operation.is_terminated():
# Controller-side miss termination (retryable) or an abort
# race (abort cleanup discards the marker).
if hit_tokens > 0:
self._finish_storage_prefetch(
req_id,
request,
fulfilled_tokens=0,
reason=(
"below_threshold"
@@ -2709,17 +2716,17 @@ class UnifiedRadixCache(BasePrefixCache):
else None
),
)
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
self._storage_prefetch_missed_rids.add(request)
self.revoke_pending_prefetch(request)
continue
if hit_tokens < self.prefetch_threshold:
# Below-threshold hits are not worth the transfer.
self._account_prefetch_outcome(operation, revoked=True)
self._finish_storage_prefetch(
req_id, fulfilled_tokens=0, reason="below_threshold"
request, fulfilled_tokens=0, reason="below_threshold"
)
self._storage_prefetch_missed_rids.add(req_id)
self.revoke_pending_prefetch(req_id)
self._storage_prefetch_missed_rids.add(request)
self.revoke_pending_prefetch(request)
continue
self._invalidate_absent_from_hit_query(operation)
self._account_prefetch_outcome(operation, revoked=False)
@@ -2732,18 +2739,18 @@ class UnifiedRadixCache(BasePrefixCache):
for ack in _drain_queue(cc.ack_prefetch_queue, n_ack_prefetch):
operation = ack.operation
if ack.completed_tokens is not None:
if operation.request_id in self.ongoing_prefetch:
if operation.handle in self.ongoing_prefetch:
assert operation.completed_tokens <= ack.completed_tokens
operation.completed_tokens = ack.completed_tokens
if ack.pool_hits is not None:
if operation.request_id in self.ongoing_prefetch:
if operation.handle in self.ongoing_prefetch:
operation.pool_storage_result.update_extra_pool_hit_pages(
ack.pool_hits
)
operation.pool_transfers_done = True
if ack.completed_req:
if operation.request_id in self.ongoing_prefetch:
# check_prefetch_progress() is not called for this rid yet.
if operation.handle in self.ongoing_prefetch:
# check_prefetch_progress() is not called for this attempt yet.
# Let us insert the prefetch result into the radix tree.
self._handle_prefetch_result(operation)
cc.append_host_mem_release(
python/sglang/srt/session/streaming_session.py
+16 -6
View File
@@ -9,6 +9,8 @@ import torch
from sglang.srt.managers.schedule_batch import ReqKvInfo
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
CacheRequestHandle,
CacheRequestOutcome,
DecLockRefParams,
DecLockRefResult,
EvictParams,
@@ -360,6 +362,12 @@ class StreamingSession(BasePrefixCache):
return
self.inner.cache_unfinished_req(req, **kwargs)
def finish(self, handle: CacheRequestHandle, outcome: CacheRequestOutcome) -> None:
self.inner.finish(handle, outcome)
def release_aborted_request(self, handle: CacheRequestHandle) -> None:
self.inner.release_aborted_request(handle)
def evict(self, params: EvictParams) -> EvictResult:
return self.inner.evict(params)
@@ -564,16 +572,18 @@ class StreamingSession(BasePrefixCache):
def init_load_back(self, params: InitLoadBackParams):
return self.inner.init_load_back(params)
def pop_prefetch_loaded_span(self, req_id: str) -> tuple[int, Optional[int]]:
return self.inner.pop_prefetch_loaded_span(req_id)
def pop_prefetch_loaded_span(
self, handle: CacheRequestHandle
) -> tuple[int, Optional[int]]:
return self.inner.pop_prefetch_loaded_span(handle)
def finish_storage_prefetch_admission(
self, req_id: str, fulfilled_tokens: int, reason: Optional[str]
self, handle: CacheRequestHandle, fulfilled_tokens: int, reason: Optional[str]
) -> None:
self.inner.finish_storage_prefetch_admission(req_id, fulfilled_tokens, reason)
self.inner.finish_storage_prefetch_admission(handle, fulfilled_tokens, reason)
def discard_storage_prefetch_accounting(self, req_id: str) -> None:
self.inner.discard_storage_prefetch_accounting(req_id)
def discard_storage_prefetch_accounting(self, handle: CacheRequestHandle) -> None:
self.inner.discard_storage_prefetch_accounting(handle)
def ready_to_load_host_cache(self):
return self.inner.ready_to_load_host_cache()