From dad8c074e76de83d8e39242aec49d7499181fa83 Mon Sep 17 00:00:00 2001 From: cctry Date: Mon, 14 Sep 2026 15:03:03 -0700 Subject: [PATCH] Scope prefetch cache state to the request attempt (#39318) --- .../disaggregation/decode_hicache_mixin.py | 15 +- python/sglang/srt/disaggregation/prefill.py | 11 +- python/sglang/srt/managers/schedule_batch.py | 8 + python/sglang/srt/managers/schedule_policy.py | 6 +- python/sglang/srt/managers/scheduler.py | 32 +-- .../sglang/srt/mem_cache/base_prefix_cache.py | 33 ++- .../srt/mem_cache/buffer_mode/pipeline.py | 124 ++++---- python/sglang/srt/mem_cache/hiradix_cache.py | 33 ++- .../hybrid_cache/hybrid_cache_controller.py | 10 +- .../storage/flexkv/flexkv_connector.py | 103 +++---- .../storage/flexkv/flexkv_radix_cache.py | 65 ++--- .../unified_cache/storage_attachment.py | 22 +- .../srt/mem_cache/unified_radix_cache.py | 265 +++++++++--------- .../sglang/srt/session/streaming_session.py | 22 +- .../test_decode_hicache_tree_core.py | 6 +- .../test_prefill_abort_result_cleanup.py | 13 +- .../unit/managers/test_prefill_adder.py | 8 +- .../unit/managers/test_scheduler_timeouts.py | 11 +- .../mem_cache/test_buffer_mode_sidecar.py | 5 +- ...test_hicache_staged_write_back_dispatch.py | 16 +- .../test_unified_radix_cache_unittest.py | 120 +++++--- 21 files changed, 541 insertions(+), 387 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index 5f0dac9aa..acfddcecd 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/python/sglang/srt/disaggregation/decode_hicache_mixin.py @@ -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( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 285401064..512e46bdb 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -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 " diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 199dac90c..9f251fd30 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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.""" diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 59a7fe2e6..b2bc89939 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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, ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f6bf9b53a..e9cba392f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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() diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 9695d1048..77e152903 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -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: """ diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index c5f5e0561..7ef87d2ba 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py @@ -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 diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index d1fa1bdd9..386f4de5c 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 69e3fa84b..b653cd3a7 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -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, diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py index 7c4211179..0f9f204c7 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py @@ -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 diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py index 7eb53ddf7..deba6a699 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py index 2c02d29b2..60c5a52b0 100644 --- a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py +++ b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py @@ -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() diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index de31d9022..420fa8cc1 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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( diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 3d513a9b8..c1a52a2ff 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -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() diff --git a/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py b/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py index a40389402..adc255ed6 100644 --- a/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py +++ b/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py @@ -10,6 +10,7 @@ from sglang.srt.disaggregation.decode_hicache_mixin import ( DecodeHiCachePreallocMixin, DecodePrefixMatch, ) +from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -26,6 +27,7 @@ class TestDecodeHiCacheTreeCore(CustomTestCase): tree_cache = SimpleNamespace( hicache_storage_pass_prefix_keys=True, ongoing_prefetch=ongoing_prefetch, + has_ongoing_prefetch=ongoing_prefetch.__contains__, is_backuped=Mock(return_value=True), is_root=Mock(return_value=False), get_last_hash_value=Mock(return_value="h2"), @@ -39,6 +41,7 @@ class TestDecodeHiCacheTreeCore(CustomTestCase): ) req = SimpleNamespace( rid="req-0", + cache_request_handle=CacheRequestHandle("req-0", 0), origin_input_ids=[0, 1, 2, 3, 4, 5, 6, 7], extra_key="model", cache_salt="tenant-a", @@ -68,7 +71,7 @@ class TestDecodeHiCacheTreeCore(CustomTestCase): self.assertTrue(prefix_match.prefetch_registered) tree_cache.prefetch_from_storage.assert_called_once_with( - "req-0", + req.cache_request_handle, 22, [4, 5], "h2", @@ -88,6 +91,7 @@ class TestDecodeHiCacheTreeCore(CustomTestCase): harness = SimpleNamespace(tree_cache=tree_cache) req = SimpleNamespace( rid="req-0", + cache_request_handle=CacheRequestHandle("req-0", 0), origin_input_ids=[0, 1, 2, 3, 4, 5], extra_key=None, cache_salt=None, diff --git a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py index af280b2bb..4286a29f5 100644 --- a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py +++ b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py @@ -11,6 +11,10 @@ from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, ) from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.mem_cache.base_prefix_cache import ( + CacheRequestHandle, + CacheRequestOutcome, +) from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci @@ -20,6 +24,7 @@ register_cpu_ci(est_time=1, suite="base-a-test-cpu") class _Req: def __init__(self, *, inflight_middle_chunks: int, allocated: bool = True): self.rid = "aborted-prefill" + self.cache_request_handle = CacheRequestHandle(self.rid, 0) self.inflight_middle_chunks = inflight_middle_chunks self.kv = ReqKvInfo( req_pool_idx=1 if allocated else None, @@ -105,7 +110,9 @@ def test_aborted_final_result_releases_hybrid_cache( maybe_cache_unfinished_req.assert_not_called() req.disagg_kv_sender.abort.assert_called_once_with() scheduler.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(7) - scheduler.tree_cache.release_aborted_request.assert_called_once_with(req.rid) + scheduler.tree_cache.finish.assert_called_once_with( + req.cache_request_handle, CacheRequestOutcome.ABORT + ) scheduler.output_streamer.stream_output.assert_called_once_with([req], False) scheduler.send_kv_chunk.assert_not_called() assert req.output_ids == [] @@ -254,7 +261,9 @@ def test_sampling_mask_abort_preserves_error_and_releases_once( release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False) req.disagg_kv_sender.abort.assert_called_once_with() scheduler.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(7) - scheduler.tree_cache.release_aborted_request.assert_called_once_with(req.rid) + scheduler.tree_cache.finish.assert_called_once_with( + req.cache_request_handle, CacheRequestOutcome.ABORT + ) scheduler.output_streamer.stream_output.assert_called_once_with([req], False) scheduler.send_kv_chunk.assert_not_called() diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 2fbe8e1ee..77bd0a6e3 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -10,6 +10,7 @@ from sglang.srt.managers.schedule_policy import ( estimate_prefill_extend_tile_metrics, ) from sglang.srt.mem_cache.base_prefix_cache import ( + CacheRequestHandle, DecLockRefResult, IncLockRefResult, ) @@ -92,6 +93,7 @@ class TestPrefillAdder(CustomTestCase): def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time=0): req = MagicMock(spec=Req) req.rid = str(rid) + req.cache_request_handle = CacheRequestHandle(req.rid, 0) req.priority = priority req.prefix_indices = [] req.full_untruncated_fill_ids = [] @@ -139,7 +141,7 @@ class TestPrefillAdder(CustomTestCase): adder._account_prefill_cache_admission(req, prefix_len=12) self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with( - "storage-hit", + req.cache_request_handle, fulfilled_tokens=8, reason=None, ) @@ -153,7 +155,7 @@ class TestPrefillAdder(CustomTestCase): req.fulfilled_storage_hit_len.return_value = 0 adder._account_prefill_cache_admission(req, prefix_len=0) self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with( - "storage-hit", fulfilled_tokens=0, reason="device_capacity" + req.cache_request_handle, fulfilled_tokens=0, reason="device_capacity" ) def test_retracted_storage_prefetch_accounting_is_omitted(self): @@ -166,7 +168,7 @@ class TestPrefillAdder(CustomTestCase): adder._account_prefill_cache_admission(req, prefix_len=8) self.mock_tree_cache.discard_storage_prefetch_accounting.assert_called_once_with( - "retracted-storage-hit" + req.cache_request_handle ) self.mock_tree_cache.finish_storage_prefetch_admission.assert_not_called() diff --git a/test/registered/unit/managers/test_scheduler_timeouts.py b/test/registered/unit/managers/test_scheduler_timeouts.py index ca67a8c8d..518b5a679 100644 --- a/test/registered/unit/managers/test_scheduler_timeouts.py +++ b/test/registered/unit/managers/test_scheduler_timeouts.py @@ -17,6 +17,10 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel maybe_stub_sgl_kernel() from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.mem_cache.base_prefix_cache import ( + CacheRequestHandle, + CacheRequestOutcome, +) register_cpu_ci(est_time=12, suite="base-a-test-cpu") @@ -24,6 +28,7 @@ register_cpu_ci(est_time=12, suite="base-a-test-cpu") class _FakeReq: def __init__(self, rid, wait_entry=0.0, forward_entry=0.0, is_finished=False): self.rid = rid + self.cache_request_handle = CacheRequestHandle(rid, 0) self.to_finish = None self.beam_group = None self._finished = is_finished @@ -83,11 +88,13 @@ class TestQueuedLimitAbort(CustomTestCase): s.enable_priority_scheduling = True s.schedule_low_priority_values_first = False s.enable_hierarchical_cache = True - s.tree_cache = MagicMock(spec=["release_aborted_request"]) + s.tree_cache = MagicMock(spec=["finish"]) self.assertFalse(s._abort_on_queued_limit(incoming)) - s.tree_cache.release_aborted_request.assert_called_once_with("candidate") + s.tree_cache.finish.assert_called_once_with( + candidate.cache_request_handle, CacheRequestOutcome.ABORT + ) self.assertEqual(s.waiting_queue, []) diff --git a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py index 3a9fcae76..622afdfe7 100644 --- a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py +++ b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock import torch +from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle from sglang.srt.mem_cache.buffer_mode.pipeline import ( BufferModePipeline, _UnifiedBackupIntent, @@ -237,7 +238,7 @@ class TestBufferModeSidecar(unittest.TestCase): storage_start=0, ) host_indices = torch.arange(4, dtype=torch.int64) - req_id = "sidecar-prefetch" + req_id = CacheRequestHandle("sidecar-prefetch", 0) cache = MagicMock() cache.page_size = 2 @@ -263,7 +264,7 @@ class TestBufferModeSidecar(unittest.TestCase): self.assertTrue( pipeline.stage_completed_prefetch( - req_id=req_id, + request=req_id, num_tokens=len(host_indices), hash_value=["page-0", "page-1"], ) diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index 76bcea13a..7cbb73234 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -9,6 +9,7 @@ import torch from sglang.srt.managers.cache_controller import CacheOperation, HiCacheController from sglang.srt.mem_cache import l2_transfer as transfer_module +from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline from sglang.srt.mem_cache.hicache_storage import ( PoolHitPolicy, @@ -264,12 +265,13 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase): self.assertEqual(controller.ack_load_queue[0].node_ids, [7, 7]) def test_short_staged_swa_tail_resolves_device_covered_head(self): + handle = CacheRequestHandle("r", 0) pipeline = BufferModePipeline.__new__(BufferModePipeline) pipeline._cache = mock.Mock() pipeline.release_staged_hold = mock.Mock(return_value=True) pipeline.staged_prefetches = { - "r": SimpleNamespace( - req_id="r", + handle: SimpleNamespace( + request=handle, key_tokens=list(range(8)), extra_key=None, cache_salt=None, @@ -288,9 +290,13 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase): ) } - self.assertEqual(pipeline.plan_staged_splice("r", device_prefix_len=6), (0, 0)) - pipeline._cache._resolve_storage_prefetch_tokens.assert_called_once_with("r", 4) - pipeline.release_staged_hold.assert_called_once_with("r", reason="shrunk") + self.assertEqual( + pipeline.plan_staged_splice(handle, device_prefix_len=6), (0, 0) + ) + pipeline._cache._resolve_storage_prefetch_tokens.assert_called_once_with( + handle, 4 + ) + pipeline.release_staged_hold.assert_called_once_with(handle, reason="shrunk") def test_l2_transfer_maps_global_layers(self): host_pool = mock.Mock() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index ef79f571e..719324844 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -30,6 +30,8 @@ from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req, ReqKvInfo from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( + CacheRequestHandle, + CacheRequestOutcome, DecLockRefParams, EvictParams, InitLoadBackParams, @@ -899,7 +901,9 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): controller = FakeCacheController() cache.cache_controller = controller - cache.prefetch_from_storage("req", cache.root_node_handle(), tokens) + cache.prefetch_from_storage( + CacheRequestHandle("req", 0), cache.root_node_handle(), tokens + ) _, storage_key, _, _, _ = controller.prefetch_args self.assertIsInstance(storage_key, RadixKey) @@ -954,7 +958,7 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): ) self.assertEqual(len(match.device_indices), len(prefix_tokens) - 1) - req_id = "bigram-anchor" + req_id = CacheRequestHandle("bigram-anchor", 0) prefetch_key = RadixKey( array("q", [prefix_tokens[-1], 6, 7, 8, 9]), extra_key=extra_key, @@ -3193,7 +3197,8 @@ class UnifiedRadixCacheSuite: if prefix_len is None: prefix_len = f.matched_len req = mock.Mock() - req.rid = req_id + req.rid = req_id.rid + req.cache_request_handle = req_id req.extra_key = extra_key req.cache_salt = cache_salt if prefix_indices is not None: @@ -3315,7 +3320,7 @@ class UnifiedRadixCacheSuite: storage_dir=storage_dir, prefetch_threshold=1, ) - req_id = "l3-prefetch-req" + req_id = CacheRequestHandle("l3-prefetch-req", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -3387,7 +3392,7 @@ class UnifiedRadixCacheSuite: storage_dir=storage_dir, prefetch_threshold=1, ) - req_id = "abort-req" + req_id = CacheRequestHandle("abort-req", 0) cc = cons.cache_controller kv_pool_available_size_before = cc.mem_pool_host.get_pool( @@ -3449,7 +3454,7 @@ class UnifiedRadixCacheSuite: self.assertFalse(op.pool_transfers_done) self.assertEqual(cc.host_mem_release_queue.qsize(), 0) self.assertEqual(swa_release_q.qsize(), 0) - cons.release_aborted_request(req_id) + cons.finish(req_id, CacheRequestOutcome.ABORT) self.assertEqual(cc.host_mem_release_queue.qsize(), 0) self.assertEqual(swa_release_q.qsize(), 0) @@ -3542,7 +3547,7 @@ class UnifiedRadixCacheSuite: storage_dir=storage_dir, prefetch_threshold=1, ) - req_id = "abort-req" + req_id = CacheRequestHandle("abort-req", 0) cc = cons.cache_controller kv_pool_available_size_before = cc.mem_pool_host.get_pool( @@ -3594,7 +3599,7 @@ class UnifiedRadixCacheSuite: self.assertEqual(swa_release_q.qsize(), 0) # --- Act: abort without committing the prefetch. --- - cons.release_aborted_request(req_id) + cons.finish(req_id, CacheRequestOutcome.ABORT) self.assertTrue(op.pool_transfers_done) self.assertGreater(swa_release_q.qsize(), 0) @@ -3815,7 +3820,7 @@ class UnifiedRadixCacheSuite: dev_avail0 = cons.token_to_kv_pool_allocator.available_size() stats = cons._prefetch_outcome_stats - req_id = "buffer-read-roundtrip" + req_id = CacheRequestHandle("buffer-read-roundtrip", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -3865,7 +3870,8 @@ class UnifiedRadixCacheSuite: held = cons.buffer_pipeline.staged_prefetches[req_id] req = mock.Mock() - req.rid = req_id + req.rid = req_id.rid + req.cache_request_handle = req_id req.extra_key = None req.cache_salt = None req.last_node = cons.root_node_handle() @@ -3954,7 +3960,7 @@ class UnifiedRadixCacheSuite: ), len(seq), ) - root_req = "salted-root-prefetch" + root_req = CacheRequestHandle("salted-root-prefetch", 0) cons.prefetch_from_storage( root_req, cons.root_node_handle(), @@ -4021,7 +4027,7 @@ class UnifiedRadixCacheSuite: ) anchor = prefix_match.last_device_node lock_ref = _device_lock_ref(cons2, anchor, ComponentType.FULL) - anchored_req = "salted-mid-tree-prefetch" + anchored_req = CacheRequestHandle("salted-mid-tree-prefetch", 0) cons2.prefetch_from_storage( anchored_req, anchor, @@ -4080,7 +4086,7 @@ class UnifiedRadixCacheSuite: stats = cons._prefetch_outcome_stats # Query BEFORE any producer wrote the span: full miss -> revoked. - req_id = "early-query-miss" + req_id = CacheRequestHandle("early-query-miss", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -4112,7 +4118,7 @@ class UnifiedRadixCacheSuite: self.assertEqual(cons.pop_prefetch_loaded_tokens(req_id), len(seq)) # Unserved markers must not leak: abort cleanup ... - aborted_rid = "aborted-miss" + aborted_rid = CacheRequestHandle("aborted-miss", 0) cons.prefetch_from_storage( aborted_rid, cons.root_node_handle(), @@ -4125,15 +4131,21 @@ class UnifiedRadixCacheSuite: lambda: cons.check_prefetch_progress(aborted_rid), "aborted-rid miss did not resolve", ) - cons.release_aborted_request(aborted_rid) + cons.finish(aborted_rid, CacheRequestOutcome.ABORT) self.assertFalse(cons.pop_storage_prefetch_miss(aborted_rid)) # A fully-device-matched (empty-suffix) decline also arms the retry: # the device match can evict while the request waits in the queue. cons.prefetch_from_storage( - "fully-matched", cons.root_node_handle(), array("q", []), None, None + CacheRequestHandle("fully-matched", 0), + cons.root_node_handle(), + array("q", []), + None, + None, + ) + self.assertTrue( + cons.pop_storage_prefetch_miss(CacheRequestHandle("fully-matched", 0)) ) - self.assertTrue(cons.pop_storage_prefetch_miss("fully-matched")) cons.sanity_check() def test_buffer_only_anchor_lock_cap_clamped_by_context_headroom(self): @@ -4192,7 +4204,7 @@ class UnifiedRadixCacheSuite: cons, cons_alloc, _ = build_fixture(self.cfg) self._init_buffer_hicache(cons, storage_dir) - req_id = "buffer-swa-admission-oom" + req_id = CacheRequestHandle("buffer-swa-admission-oom", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -4226,7 +4238,8 @@ class UnifiedRadixCacheSuite: # Consume at admission (init_load_back + request lock). held = cons.buffer_pipeline.staged_prefetches[req_id] req = mock.Mock() - req.rid = req_id + req.rid = req_id.rid + req.cache_request_handle = req_id req.extra_key = None req.cache_salt = None req.last_node = cons.root_node_handle() @@ -4297,7 +4310,7 @@ class UnifiedRadixCacheSuite: cons.storage_metrics_collector = mock.Mock() avail0 = self._host_avail_sizes(cons) - req_id = "sibling-publish" + req_id = CacheRequestHandle("sibling-publish", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -4371,7 +4384,7 @@ class UnifiedRadixCacheSuite: self.assertEqual(masked.full_kv_hit_length, len(seq), "live FULL not resident") avail0 = self._host_avail_sizes(cons) - req_id = "masked-overlap" + req_id = CacheRequestHandle("masked-overlap", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -4414,7 +4427,7 @@ class UnifiedRadixCacheSuite: cons, cons_alloc, cons_rtp = build_fixture(self.cfg) self._init_buffer_hicache(cons, storage_dir) - req_id = "post-check-overlap" + req_id = CacheRequestHandle("post-check-overlap", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -4439,7 +4452,8 @@ class UnifiedRadixCacheSuite: f = cons.buffer_pipeline.staged_prefetches[req_id] req = mock.Mock() - req.rid = req_id + req.rid = req_id.rid + req.cache_request_handle = req_id req.extra_key = None req.cache_salt = None req.prefix_indices = torch.zeros( @@ -4482,7 +4496,7 @@ class UnifiedRadixCacheSuite: self._init_buffer_hicache(cons, storage_dir) avail0 = self._host_avail_sizes(cons) - req_id = "growth-trim" + req_id = CacheRequestHandle("growth-trim", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -4551,7 +4565,7 @@ class UnifiedRadixCacheSuite: self._init_buffer_hicache(cons, storage_dir) avail0 = self._host_avail_sizes(cons) - req_id = "covered-hold" + req_id = CacheRequestHandle("covered-hold", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -4597,7 +4611,7 @@ class UnifiedRadixCacheSuite: avail0 = self._host_avail_sizes(cons) stats = cons._prefetch_outcome_stats - req_id = "covered-at-commit" + req_id = CacheRequestHandle("covered-at-commit", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -4648,9 +4662,13 @@ class UnifiedRadixCacheSuite: cons, _, _ = build_fixture(self.cfg) self._init_buffer_hicache(cons, storage_dir) cons.prefetch_from_storage( - "short-req", cons.root_node_handle(), array("q", seq), None, None + CacheRequestHandle("short-req", 0), + cons.root_node_handle(), + array("q", seq), + None, + None, ) - self._run_prefetch_to_completion(cons, "short-req") + self._run_prefetch_to_completion(cons, CacheRequestHandle("short-req", 0)) cons.drain_storage_control_queues() mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) self.assertEqual(len(mc.device_indices), len(seq)) @@ -4671,7 +4689,7 @@ class UnifiedRadixCacheSuite: self._insert(cons2, cons2_alloc, cons2_rtp, seq_a) m = cons2.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a)))) cons2.prefetch_from_storage( - "subwin-req", + CacheRequestHandle("subwin-req", 0), m.last_device_node, array("q", seq_ab[len(seq_a) :]), cons2.get_last_hash_value(m.last_device_node), @@ -4681,8 +4699,10 @@ class UnifiedRadixCacheSuite: self._pump_hicache_until( cons2, lambda: ( - cons2.check_prefetch_progress("subwin-req") - and cons2.buffer_pipeline.has_staged("subwin-req") + cons2.check_prefetch_progress(CacheRequestHandle("subwin-req", 0)) + and cons2.buffer_pipeline.has_staged( + CacheRequestHandle("subwin-req", 0) + ) ), "sub-window prefetch did not stage", ) @@ -4690,13 +4710,15 @@ class UnifiedRadixCacheSuite: any( t.name == PoolName.SWA for t in cons2.buffer_pipeline.staged_prefetches[ - "subwin-req" + CacheRequestHandle("subwin-req", 0) ].aux_xfers ), "sub-window fetch degraded to KV-only", ) spliced = self._consume_staged_prefetch( - cons2, "subwin-req", prefix_indices=m.device_indices + cons2, + CacheRequestHandle("subwin-req", 0), + prefix_indices=m.device_indices, ) self.assertEqual(int(spliced.numel()), len(seq_ab) - len(seq_a)) self.assertEqual( @@ -4719,9 +4741,13 @@ class UnifiedRadixCacheSuite: self._init_buffer_hicache(cons3, storage_dir) avail3 = self._host_avail_sizes(cons3) cons3.prefetch_from_storage( - "partial-req", cons3.root_node_handle(), array("q", full), None, None + CacheRequestHandle("partial-req", 0), + cons3.root_node_handle(), + array("q", full), + None, + None, ) - self._run_prefetch_to_completion(cons3, "partial-req") + self._run_prefetch_to_completion(cons3, CacheRequestHandle("partial-req", 0)) cons3.drain_storage_control_queues() self.assertEqual( len( @@ -4817,7 +4843,7 @@ class UnifiedRadixCacheSuite: # Baseline (single rank) must actually adopt SWA, else the TP assertions # below would be vacuous -> skip. base = self._l3_consumer(storage_dir) - self._consume_prefetch(base, seq, "base") + self._consume_prefetch(base, seq, CacheRequestHandle("base", 0)) if not self._swa_host_on_path(base, seq): self.skipTest("fixture does not exercise SWA L3 prefetch") return storage_dir, seq @@ -4833,7 +4859,7 @@ class UnifiedRadixCacheSuite: cons = self._l3_consumer(storage_dir) cons.tp_world_size = 2 self._patch_tp_prefetch_sync(cons, drop_swa=True) - self._consume_prefetch(cons, seq, "drop") + self._consume_prefetch(cons, seq, CacheRequestHandle("drop", 0)) m = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) self.assertEqual(m.host_hit_length, 0) @@ -4853,7 +4879,7 @@ class UnifiedRadixCacheSuite: cons = self._l3_consumer(storage_dir) cons.tp_world_size = 2 self._patch_tp_prefetch_sync(cons, drop_swa=False) # peer == local - self._consume_prefetch(cons, seq, "keep") + self._consume_prefetch(cons, seq, CacheRequestHandle("keep", 0)) m = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) self.assertEqual(m.host_hit_length, len(seq)) @@ -4875,7 +4901,7 @@ class UnifiedRadixCacheSuite: cons.tp_world_size = 2 self._patch_tp_prefetch_sync(cons, drop_swa=True) avail_before = cons.swa_kv_pool_host.available_size() - self._consume_prefetch(cons, seq, "drop") + self._consume_prefetch(cons, seq, CacheRequestHandle("drop", 0)) self.assertEqual( cons.match_prefix( @@ -9003,10 +9029,11 @@ class TestPrefetchCommitOrdering(CustomTestCase): insert_result.host_insert_dropped = False cache.tree_core.insert_host.return_value = insert_result operation = mock.MagicMock() + operation.handle = CacheRequestHandle("req", 0) operation.request_id = "req" operation.completed_tokens = 8 cache.ongoing_prefetch = { - operation.request_id: ( + operation.handle: ( 7, list(range(8)), list(range(100, 108)), @@ -9041,7 +9068,11 @@ class TestPrefetchCommitOrdering(CustomTestCase): cache._handle_prefetch_result = _handle_prefetch_result - self.assertTrue(UnifiedRadixCache.check_prefetch_progress(cache, "req")) + self.assertTrue( + UnifiedRadixCache.check_prefetch_progress( + cache, CacheRequestHandle("req", 0) + ) + ) self.assertEqual([c[0] for c in order.mock_calls], ["apply", "commit", "apply"]) self.assertEqual(applied[0], [walk_action]) @@ -9182,8 +9213,9 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase): }, ) anchor_lock_params = cache.inc_host_lock_ref(parent_id).to_dec_params() - req_id = "drop-all-resources" - operation.request_id = req_id + req_id = CacheRequestHandle("drop-all-resources", 0) + operation.handle = req_id + operation.request_id = req_id.rid cache.ongoing_prefetch[req_id] = _OngoingPrefetch( parent_id, prefetch_key, @@ -9496,7 +9528,7 @@ class TestAnchorLockOutcomePolicy(CustomTestCase): instead of gambling the read; cap_skip over budget (checked before the match walk).""" - _REQ = "req-1" + _REQ = CacheRequestHandle("req-1", 0) _PREFIX = list(range(100, 100 + 8)) def _make_pipeline(self, cache, cap_tokens=10_000):