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

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

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