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

[HiCache] Rework the buffer-mode storage prefetch pipeline and retry bookkeeping (#39283)

Browse Source
This commit is contained in:
Zhiqiang Xie
2026-09-15 10:48:49 -07:00
committed by GitHub
parent 03ea13a545
commit 7f5dd19256
43 changed files with 3062 additions and 869 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/fields/memory.py
+16 -7
View File
@@ -177,17 +177,26 @@ class Memory(msgspec.Struct):
int,
Arg(
help=(
"Scheduling passes a queued request waits after a storage "
"prefetch miss before the availability check is retried "
"(under load the first check can run before the needed "
"backup commits). 0 disables retries."
"Scheduling passes a queued request waits before its storage "
"availability check is re-issued, when the prefetch found "
"nothing and a backup may still be committing (under load the "
"first check can run before it does). A re-issue that waits on "
"staging or a moved match instead goes out on the next pass. "
"Only passes that reach prefill scheduling count. 0 disables "
"miss retries; known-hit deferrals are always re-issued."
),
),
] = 0
] = 8
hicache_storage_prefetch_retry_max_attempts: A[
int,
"Maximum storage prefetch retries per request when --hicache-storage-prefetch-retry-poll-interval is set.",
] = 4
Arg(
help=(
"Storage availability re-issues a queued request may make, paced "
"miss polls and immediate re-issues alike; past the cap it is "
"admitted with whatever the device holds. 0 disables re-issues."
),
),
] = 8
# -------------------------------------------------------------------------
# Unified Radix Cache
python/sglang/srt/disaggregation/prefill.py
+1
View File
@@ -592,6 +592,7 @@ class SchedulerDisaggregationPrefillMixin:
last_batch: Optional[ScheduleBatch],
) -> NextBatchPlan:
self.process_pending_chunked_abort()
self._process_hicache_events()
# HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it
# Otherwise, it hangs under high concurrency
python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py
+1
View File
@@ -402,6 +402,7 @@ class C128SidecarComponent(TreeComponent):
host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0,
staging_tokens: int = 0,
last_hash: Optional[str] = None,
) -> Optional[list[PoolTransfer]]:
ct = self.component_type
python/sglang/srt/managers/schedule_batch.py
+5 -4
View File
@@ -149,6 +149,7 @@ if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
from sglang.srt.managers.scheduler_components.metrics_reporter import PrefillStats
from sglang.srt.mem_cache.storage_prefetch import StagedPrefetchPlan
from sglang.srt.session.session_controller import Session
from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm
@@ -1127,14 +1128,14 @@ class Req(ReqDllmMixin):
self.host_loaded_length = 0
# Buffer-mode host memory is transport staging, not an L2 cache tier.
self.host_hit_is_storage = False
# Storage prefetch retry state while queued
# (see Scheduler._retry_missed_storage_prefetches).
self.storage_prefetch_retry_pending = False
self.storage_prefetch_retry_wait_polls = 0
self.storage_prefetch_retry_attempts = 0
self.staged_prefetch_plan: Optional[StagedPrefetchPlan] = None
# Receipt of the tree lock held on last_node (anchor, SWA boundary,
# skipped components); every release replays it unchanged.
self.lock_receipt: DecLockRefParams = DecLockRefParams()
# Device/host prefix used to plan the latest L3 lookup. Admission uses
# it to detect newly exposed storage demand after queue-time eviction.
self.storage_prefetch_last_match_len: Optional[int] = None
# Whether the prefill-time SWA tree lock has been released early
self.swa_prefix_lock_released: bool = False
# Logical-page KV sharding: rotation base of the chain this request
python/sglang/srt/managers/schedule_policy.py
+192 -212
View File
@@ -33,6 +33,7 @@ import os
import random
from collections import Counter
from contextlib import contextmanager
from dataclasses import dataclass
from enum import Enum, auto
from functools import lru_cache
from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
@@ -548,6 +549,14 @@ class AddReqResult(Enum):
OTHER = auto() # Other reasons to stop adding requests
@dataclass(frozen=True, slots=True)
class _PrefillAdmission:
prefix_len: int
extend_len: int
max_new_tokens: int
is_chunked: bool
class PrefillAdder:
def __init__(
self,
@@ -869,6 +878,40 @@ class PrefillAdder:
>= capacity
)
def _swa_admission_gate(
self,
req: Req,
extend_input_len: int,
swa_host_hit_length: int,
chunk_tokens_limit: Optional[int],
) -> tuple[Optional[AddReqResult], Optional[int]]:
"""SWA-pool gate: a non-None verdict rejects; otherwise the returned chunk
limit stands, tightened to the pool cap when never-fits fires."""
max_new_tokens = self._swa_new_tokens(req)
swa_needed = self._swa_budget_for_req(
extend_input_len, max_new_tokens, swa_host_hit_length=swa_host_hit_length
)
# Ring-slot capacity is exact, so needing exactly what is left still
# fits; the legacy SWA-token path keeps its conservative `>=`.
fits = (
swa_needed <= self.rem_swa_tokens
if self._swa_req_ring
else swa_needed < self.rem_swa_tokens
)
if fits:
return None, chunk_tokens_limit
if not self._swa_req_never_fits(
extend_input_len, max_new_tokens, swa_host_hit_length
):
return AddReqResult.NO_TOKEN, chunk_tokens_limit
swa_cap = self._swa_chunk_cap(max_new_tokens, swa_host_hit_length)
if self.rem_chunk_tokens is None or swa_cap <= 0:
return AddReqResult.NO_TOKEN, chunk_tokens_limit
current = (
self.rem_chunk_tokens if chunk_tokens_limit is None else chunk_tokens_limit
)
return None, min(current, swa_cap)
def _mamba_gap_budget_for_req(self, req: Req) -> int:
"""Shared-gap reservation (full-token-equivalents) for a request's new
mamba state. Charged only on the SHARED Mamba pool (`_mamba_slot_cost > 0`)
@@ -1326,87 +1369,23 @@ class PrefillAdder:
mamba_gap_reserve = self._mamba_gap_budget_for_req(req)
total_tokens += mamba_gap_reserve
# adjusting the input_tokens based on host_hit_length and page_size
real_input_tokens = cand_extend_input_len - req.host_hit_length
real_input_tokens = self.ceil_paged_tokens(real_input_tokens)
prefix_len = len(req.prefix_indices)
if total_tokens >= self.rem_total_tokens:
return AddReqResult.NO_TOKEN
chunk_tokens_limit = self.rem_chunk_tokens
if self.is_hybrid_swa:
# host-hit prefix is loaded back, not re-prefilled, so the SWA peak is
# driven only by the freshly-prefilled tail (the loaded window is
# charged separately via swa_host_hit_length).
swa_needed = self._swa_budget_for_req(
real_input_tokens,
self._swa_new_tokens(req),
swa_host_hit_length=req.swa_host_hit_length,
)
# Ring-slot capacity is exact, so needing exactly what is left still
# fits; the legacy SWA-token path keeps its conservative `>=`.
if (
swa_needed > self.rem_swa_tokens
if self._swa_req_ring
else swa_needed >= self.rem_swa_tokens
):
if not self._swa_req_never_fits(
real_input_tokens,
self._swa_new_tokens(req),
req.swa_host_hit_length,
):
return AddReqResult.NO_TOKEN
swa_cap = self._swa_chunk_cap(
self._swa_new_tokens(req), req.swa_host_hit_length
)
if self.rem_chunk_tokens is None or swa_cap <= 0:
return AddReqResult.NO_TOKEN
chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap)
if (
self.rem_chunk_tokens is None
and len(self.can_run_list) != 0
and real_input_tokens >= self.rem_input_tokens
):
# If without chunked prefill:
# - if the can_run_list is not empty, we satisfy the constraint of (max_prefill_tokens)
# - if the can_run_list is empty, always accept the first prefill request
return AddReqResult.OTHER
# The temporary pin excludes this prefix from the evictable budget.
# Selection itself neither allocates slots nor materializes host hits.
with self._lock_node(req.last_node):
# self.rem_total_tokens may decrease after the lock acquisition
if total_tokens >= self.rem_total_tokens:
return AddReqResult.NO_TOKEN
admission = self._select_prefill_admission(
req,
total_tokens=total_tokens,
host_hit_length=req.host_hit_length,
swa_host_hit_length=req.swa_host_hit_length,
truncation_align_size=truncation_align_size,
)
if isinstance(admission, AddReqResult):
return admission
if self.is_hybrid_swa:
# self.rem_swa_tokens may decrease after the lock acquisition
swa_needed = self._swa_budget_for_req(
real_input_tokens,
self._swa_new_tokens(req),
swa_host_hit_length=req.swa_host_hit_length,
)
if (
swa_needed > self.rem_swa_tokens
if self._swa_req_ring
else swa_needed >= self.rem_swa_tokens
):
if not self._swa_req_never_fits(
real_input_tokens,
self._swa_new_tokens(req),
req.swa_host_hit_length,
):
return AddReqResult.NO_TOKEN
swa_cap = self._swa_chunk_cap(
self._swa_new_tokens(req), req.swa_host_hit_length
)
if self.rem_chunk_tokens is None or swa_cap <= 0:
return AddReqResult.NO_TOKEN
chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap)
# Negotiate only after every KV-budget gate (a NO_TOKEN rank must
# report not-prefillable via finalize()) and before init_load_back
# (a delay verdict must not start KV load-back).
# A rejected candidate must not report prefillable or queue H2D.
if (self.prefill_delayer_single_pass is not None) and (
not self.prefill_delayer_single_pass.negotiate_should_allow_prefill(
local_prefillable=True,
@@ -1419,166 +1398,167 @@ class PrefillAdder:
return AddReqResult.OTHER
if req.needs_host_load_back():
new_indices, req.last_node = self.tree_cache.init_load_back(
promised_host_hit = req.host_hit_length
loaded = self.tree_cache.init_load_back(
InitLoadBackParams(
best_match_node=req.best_match_node,
host_hit_length=req.host_hit_length,
req=req,
)
)
req.host_loaded_length = len(new_indices)
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
prefix_len = len(req.prefix_indices)
req.kv.cache_protected_len = prefix_len
raw_input_tokens = len(req.full_untruncated_fill_ids) - len(
req.prefix_indices
)
input_tokens = self.ceil_paged_tokens(raw_input_tokens)
# Whether the request fits whole. Against the raw length under
# exact-chunk-fill, so a request whose ceiled length would spill is
# not needlessly split into a second chunk.
chunk_fit_tokens = (
raw_input_tokens if self.exact_chunk_fill else input_tokens
)
if (
self.rem_chunk_tokens is None
and len(self.can_run_list) != 0
and input_tokens >= self.rem_input_tokens
):
# If without chunked prefill:
# - if the can_run_list is not empty, we satisfy the constraint of (max_prefill_tokens)
# - if the can_run_list is empty, always accept the first prefill request
return AddReqResult.OTHER
if self.dllm_config is not None:
if self.rem_dllm_tokens <= 0:
if loaded is None:
return AddReqResult.OTHER
new_indices, req.last_node = loaded
req.host_loaded_length = len(new_indices)
if 0 < req.host_loaded_length < promised_host_hit:
raise RuntimeError(
"HiCache load-back must commit all promised FULL tokens or none: "
f"req={req.rid} promised={promised_host_hit} "
f"loaded={req.host_loaded_length}"
)
if req.host_loaded_length > promised_host_hit:
# A load can expose resident FULL behind host-only aux state; its
# H2D is queued, so keep the approved budget and shrink the work.
prefix_len = len(req.prefix_indices) + req.host_loaded_length
extend_len = admission.extend_len
if self.dllm_config is None:
extend_len = min(
extend_len, len(req.full_untruncated_fill_ids) - prefix_len
)
is_chunked = admission.is_chunked and (
prefix_len + extend_len < len(req.full_untruncated_fill_ids)
)
max_new_tokens = admission.max_new_tokens
if admission.is_chunked and not is_chunked:
max_new_tokens = min(
req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS
)
admission = _PrefillAdmission(
prefix_len, extend_len, max_new_tokens, is_chunked
)
elif req.host_loaded_length < promised_host_hit:
# No FULL was loaded; recomputation may no longer fit.
admission = self._select_prefill_admission(
req,
total_tokens=total_tokens,
host_hit_length=0,
swa_host_hit_length=0,
truncation_align_size=truncation_align_size,
)
if isinstance(admission, AddReqResult):
return admission
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
req.kv.cache_protected_len = len(req.prefix_indices)
assert truncation_align_size is None, (
"truncation_align_size is not supported for dllm prefill"
)
# Successful materialization has no remaining admission gates.
self._commit_prefill_admission(req, admission, mamba_gap_reserve)
if (
tile_stop := self._check_prefill_tile_budget(input_tokens)
) is not None:
return tile_stop
# This verdict controls the next candidate, not the committed request.
return self.budget_state()
self._add_dllm_req(req, prefix_len)
self._req_inc_lock_ref(req)
elif chunk_tokens_limit is None or chunk_fit_tokens <= chunk_tokens_limit:
if (
tile_stop := self._check_prefill_tile_budget(input_tokens)
) is not None:
return tile_stop
def _select_prefill_admission(
self,
req: Req,
*,
total_tokens: int,
host_hit_length: int,
swa_host_hit_length: int,
truncation_align_size: Optional[int],
) -> _PrefillAdmission | AddReqResult:
"""Select a prefill shape without allocating or publishing cached KV."""
if total_tokens >= self.rem_total_tokens:
return AddReqResult.NO_TOKEN
# Non-chunked prefill — the whole sequence is committed this iter.
req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
self.can_run_list.append(req)
prefix_len = len(req.prefix_indices) + host_hit_length
extend_len = len(req.full_untruncated_fill_ids) - prefix_len
input_tokens = self.ceil_paged_tokens(extend_len)
# Whether the request fits whole. Against the raw length under
# exact-chunk-fill, so a request whose ceiled length would spill is
# not needlessly split into a second chunk.
chunk_fit_tokens = extend_len if self.exact_chunk_fill else input_tokens
chunk_tokens_limit = self.rem_chunk_tokens
if self.is_hybrid_swa:
verdict, chunk_tokens_limit = self._swa_admission_gate(
req, input_tokens, swa_host_hit_length, chunk_tokens_limit
)
if verdict is not None:
return verdict
self._req_inc_lock_ref(req)
self._update_prefill_budget(
prefix_len,
req.extend_range.length,
min(
req.sampling_params.max_new_tokens,
CLIP_MAX_NEW_TOKENS,
),
req.retracted_stain,
mamba_gap_reserve=mamba_gap_reserve,
compute_charge=raw_input_tokens if self.exact_chunk_fill else None,
)
self._account_prefill_cache_admission(req, prefix_len)
elif self.exact_chunk_fill:
# Without chunking, allow the first request even above the input cap.
if (
self.rem_chunk_tokens is None
and self.can_run_list
and input_tokens >= self.rem_input_tokens
):
return AddReqResult.OTHER
is_chunked = False
max_new_tokens = min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS)
tile_tokens = input_tokens
if self.dllm_config is not None:
assert truncation_align_size is None, (
"truncation_align_size is not supported for dllm prefill"
)
extend_len = (
min(self.rem_dllm_tokens, self.dllm_block_size)
// self.page_size
* self.page_size
)
if extend_len <= 0:
return AddReqResult.OTHER
max_new_tokens = 0
elif chunk_tokens_limit is not None and chunk_fit_tokens > chunk_tokens_limit:
if self.exact_chunk_fill:
# Take the remainder verbatim so the batch hits exactly
# chunked_prefill_size. `chunk_fit_tokens > chunk_tokens_limit`
# here, so this never runs past the end of the prompt. Uses the
# limit rather than rem_chunk_tokens so an SWA-capped chunk stays
# capped.
trunc_len = chunk_tokens_limit
if trunc_len <= 0:
return AddReqResult.OTHER
extend_len = chunk_tokens_limit
if truncation_align_size is not None:
if trunc_len < truncation_align_size:
return AddReqResult.OTHER
trunc_len = truncation_align_size * (
trunc_len // truncation_align_size
extend_len = (
extend_len // truncation_align_size * truncation_align_size
)
if (
tile_stop := self._check_prefill_tile_budget(trunc_len)
) is not None:
return tile_stop
req.set_extend_range(
len(req.prefix_indices), len(req.prefix_indices) + trunc_len
)
self.can_run_list.append(req)
self.new_chunked_req = req
self._req_inc_lock_ref(req)
self._update_prefill_budget(
prefix_len,
trunc_len,
0,
req.retracted_stain,
mamba_gap_reserve=mamba_gap_reserve,
compute_charge=trunc_len,
)
self._account_prefill_cache_admission(req, prefix_len)
else:
# Make sure at least one page is available
trunc_len = chunk_tokens_limit // self.page_size * self.page_size
if trunc_len <= 0:
return AddReqResult.OTHER
# When truncation align size is set, we want to assert that the prefill prefix length is multiple of truncation align size
# A typical use case is when deterministic inference is enabled with flashinfer attention backend,
# we need the prefill prefix length to be multiple of attention split size
extend_len = chunk_tokens_limit // self.page_size * self.page_size
if truncation_align_size is not None:
if trunc_len < truncation_align_size:
return AddReqResult.OTHER
else:
trunc_len = truncation_align_size * (
trunc_len // truncation_align_size
)
extend_len = (
extend_len // truncation_align_size * truncation_align_size
)
end = (prefix_len + extend_len) // self.page_size * self.page_size
extend_len = end - prefix_len
if extend_len <= 0:
return AddReqResult.OTHER
is_chunked = True
max_new_tokens = 0
tile_tokens = extend_len
now_input_len = trunc_len + len(req.prefix_indices)
now_input_len = now_input_len // self.page_size * self.page_size
trunc_len = now_input_len - len(req.prefix_indices)
if (verdict := self._check_prefill_tile_budget(tile_tokens)) is not None:
return verdict
if trunc_len <= 0:
return AddReqResult.OTHER
return _PrefillAdmission(prefix_len, extend_len, max_new_tokens, is_chunked)
if (
tile_stop := self._check_prefill_tile_budget(trunc_len)
) is not None:
return tile_stop
# Chunked prefill
req.set_extend_range(
len(req.prefix_indices), len(req.prefix_indices) + trunc_len
)
self.can_run_list.append(req)
self.new_chunked_req = req
self._req_inc_lock_ref(req)
self._update_prefill_budget(
prefix_len,
trunc_len,
0,
req.retracted_stain,
mamba_gap_reserve=mamba_gap_reserve,
)
self._account_prefill_cache_admission(req, prefix_len)
return self.budget_state()
def _commit_prefill_admission(
self, req: Req, admission: _PrefillAdmission, mamba_gap_reserve: int
) -> None:
assert len(req.prefix_indices) == admission.prefix_len
req.set_extend_range(
admission.prefix_len, admission.prefix_len + admission.extend_len
)
self._req_inc_lock_ref(req)
self.can_run_list.append(req)
if admission.is_chunked:
self.new_chunked_req = req
self._update_prefill_budget(
admission.prefix_len,
admission.extend_len,
admission.max_new_tokens,
req.retracted_stain,
mamba_gap_reserve=mamba_gap_reserve,
# Compute budgets are billed forward-pass tokens under exact-chunk-fill.
compute_charge=admission.extend_len if self.exact_chunk_fill else None,
)
self._account_prefill_cache_admission(req, admission.prefix_len)
def preempt_to_schedule(self, req: Req) -> bool:
"""
python/sglang/srt/managers/scheduler.py
+90 -58
View File
@@ -3074,7 +3074,7 @@ class Scheduler(
for tokenized_req in recv_req:
self.handle_generate_request(tokenized_req)
def _prefetch_kvcache(self, req: Req):
def _prefetch_kvcache(self, req: Req, storage_hit_end: Optional[int] = None):
if self.enable_hicache_storage:
req.init_next_round_input(self.tree_cache, cow_mamba=False)
tree_cache = self.tree_cache
@@ -3093,6 +3093,9 @@ class Scheduler(
):
last_host_node = req.last_node
matched_len = len(req.prefix_indices) + req.host_hit_length
req.storage_prefetch_last_match_len = matched_len
if (
tree_cache.is_backuped(last_host_node)
or tree_cache.is_root(last_host_node)
@@ -3101,7 +3104,6 @@ class Scheduler(
and tree_cache.get_last_hash_value(last_host_node) is not None
)
):
matched_len = len(req.prefix_indices) + req.host_hit_length
match_end = req._compute_max_prefix_len(
len(req.full_untruncated_fill_ids)
)
@@ -3120,41 +3122,78 @@ class Scheduler(
matched_prefix_tokens=req.full_untruncated_fill_ids[:matched_len],
extra_key=req.extra_key,
cache_salt=req.cache_salt,
storage_hit_end=storage_hit_end,
)
def _retry_missed_storage_prefetches(self):
"""Re-issue the availability check for queued requests whose prefetch
missed. Pacing counts scheduling passes so TP ranks re-issue on the
same pass; a sweep (the admission loop stops at the first
unschedulable request) covers the whole queue."""
interval = get_memory().hicache_storage_prefetch_retry_poll_interval
if interval <= 0 or not self.waiting_queue:
def _process_storage_prefetch_retries(self):
"""Issue due L3 attempts in the current waiting-queue order."""
retries = self.tree_cache.storage_prefetch_retries
if retries is None:
return
memory = get_memory()
for req, storage_hit_end in retries.pop_ready(
self.waiting_queue,
memory.hicache_storage_prefetch_retry_poll_interval,
memory.hicache_storage_prefetch_retry_max_attempts,
):
self._retry_storage_prefetch(req, storage_hit_end)
def _retry_storage_prefetch(
self, req: Req, storage_hit_end: Optional[int] = None
) -> None:
req.storage_prefetch_retry_attempts += 1
max_attempts = get_memory().hicache_storage_prefetch_retry_max_attempts
for req in self.waiting_queue:
if self.tree_cache.pop_storage_prefetch_miss(req.cache_request_handle):
req.storage_prefetch_retry_pending = True
req.storage_prefetch_retry_wait_polls = 0
if (
not req.storage_prefetch_retry_pending
or req.storage_prefetch_retry_attempts >= max_attempts
):
continue
req.storage_prefetch_retry_wait_polls += 1
if req.storage_prefetch_retry_wait_polls <= interval:
continue
req.storage_prefetch_retry_pending = False
req.storage_prefetch_retry_attempts += 1
logger.debug(
"HiCache storage prefetch retry req=%s attempt=%d",
if req.storage_prefetch_retry_attempts >= max_attempts:
logger.warning(
"HiCache storage prefetch reissue cap reached req=%s attempts=%d; "
"the request is admitted without further L3 lookups",
req.rid,
req.storage_prefetch_retry_attempts,
)
self._prefetch_kvcache(req)
else:
logger.debug(
"HiCache storage prefetch re-issue req=%s attempt=%d",
req.rid,
req.storage_prefetch_retry_attempts,
)
self._prefetch_kvcache(req, storage_hit_end)
def _prefetch_after_device_hit_loss(self, req: Req) -> bool:
"""Re-query an L3 range newly exposed by queue-time device eviction."""
previous_match_len = req.storage_prefetch_last_match_len
buffer_pipeline = self.tree_cache.buffer_pipeline
if not previous_match_len or (
buffer_pipeline is not None
and buffer_pipeline.has_staged(req.cache_request_handle)
):
return False
current_match_len = len(req.prefix_indices) + req.host_hit_length
if current_match_len >= previous_match_len:
return False
if (
req.storage_prefetch_retry_attempts
>= get_memory().hicache_storage_prefetch_retry_max_attempts
):
# Past the re-issue cap the shorter live match is admitted as is.
req.storage_prefetch_last_match_len = current_match_len
return False
logger.warning(
"HiCache device prefix shrank before admission req=%s "
"lookup_match=%d current_match=%d; reissuing storage lookup",
req.rid,
previous_match_len,
current_match_len,
)
self._retry_storage_prefetch(req)
return True
def _add_request_to_queue(self, req: Req, is_retracted: bool = False):
if not self._set_or_validate_priority(req):
return
if is_retracted:
req.storage_prefetch_retry_attempts = 0
req.storage_prefetch_last_match_len = None
req.staged_prefetch_plan = None
if self.disaggregation_mode == DisaggregationMode.NULL:
if self._abort_on_queued_limit(req):
return
@@ -3504,11 +3543,24 @@ class Scheduler(
# todo hisparse, maybe other info to contain for the new batch
return batch
def _process_hicache_events(self) -> None:
# The HiCache drain is TP-wide consensus; run it before rank-local
# decisions (_should_defer_prefill) or ranks enter different collectives.
if (
self.enable_hierarchical_cache
or get_memory().enable_flexkv
or self.enable_unified_cache_external_linker
):
self.tree_cache.check_hicache_events()
if self.enable_hicache_storage:
self._process_storage_prefetch_retries()
@scheduler_stage_method(SCHEDULER_STAGE_GET_NEXT_BATCH)
def get_next_batch_to_run(
self, running_batch: ScheduleBatch, last_batch: Optional[ScheduleBatch]
) -> NextBatchPlan:
self.process_pending_chunked_abort()
self._process_hicache_events()
if self.enable_fpm:
self._fpm_batch_t0 = time.monotonic()
@@ -3710,15 +3762,6 @@ class Scheduler(
for req in ready_grammar_requests:
self._add_request_to_queue(req)
if (
self.enable_hierarchical_cache
or get_memory().enable_flexkv
or self.enable_unified_cache_external_linker
):
self.tree_cache.check_hicache_events()
if self.enable_hicache_storage:
self._retry_missed_storage_prefetches()
if self.enable_priority_preemption or self.is_hybrid_swa:
# Reset batch_is_full to try preemption with a prefill adder.
running_batch.batch_is_full = False
@@ -3823,6 +3866,7 @@ class Scheduler(
mamba_allocator = getattr(self.req_to_token_pool, "mamba_allocator", None)
if mamba_allocator is not None:
mamba_allocator.alloc_group_begin(len(self.waiting_queue))
buffer_pipeline = self.tree_cache.buffer_pipeline
# Get requests from the waiting queue to a new prefill batch
for req in self.waiting_queue:
if self.enable_lora and not self.can_schedule_lora_req(req, running_loras):
@@ -3870,32 +3914,16 @@ class Scheduler(
req.host_hit_is_storage = False
req.init_next_round_input(self.tree_cache)
if self.enable_hicache_storage and (
self._prefetch_after_device_hit_loss(req)
):
continue
if (
self.enable_hicache_storage
and get_memory().hicache_host_memory_mode == "buffer_only"
and buffer_pipeline is not None
and not buffer_pipeline.prepare_staged_prefetch(req)
):
# Buffer mode: surface a staged prefetch as the request's host
# hit (consumed through init_load_back) plus its SWA window,
# which consumption allocates and the request lock pins —
# uncharged, the batch alloc can OOM. Planned against the same
# live prefix admission uses, so only the splice-able span
# tail is charged and unusable holds are freed. Set AFTER
# init_next_round_input (which recomputes host_hit). Mamba
# (fenced in init_hicache) will need the same charge via
# mamba_host_hit_length.
held_tokens, held_swa_tokens = self.tree_cache.plan_staged_splice(
req.cache_request_handle, len(req.prefix_indices)
)
if held_tokens > 0:
req.host_hit_length = held_tokens
req.swa_host_hit_length = held_swa_tokens
req.storage_hit_length = held_tokens
req.storage_hit_start = len(req.prefix_indices)
req.host_hit_is_storage = True
elif not (req.host_hit_is_storage and req.host_loaded_length > 0):
req.storage_hit_length = 0
req.storage_hit_start = None
req.host_hit_is_storage = False
continue
res = adder.add_one_req(
req,
has_chunked_req=(self.chunked_req is not None),
@@ -3943,6 +3971,10 @@ class Scheduler(
return None, running_batch
can_run_set = set(can_run_list)
retries = self.tree_cache.storage_prefetch_retries
if self.enable_hicache_storage and retries is not None:
for req in can_run_list:
retries.cancel(req.rid)
self.waiting_queue = [x for x in self.waiting_queue if x not in can_run_set]
if adder.preempt_list:
for req in adder.preempt_list:
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+1 -1
View File
@@ -280,7 +280,7 @@ class SchedulerInvariantChecker:
batches.append(running_batch)
full_uncached = 0
swa_uncached = 0
swa_uncached = self.tree_cache.swa_transient_size()
counted: set[int] = set()
reqs = [req for batch in batches for req in batch.reqs]
chunked_req = self.get_chunked_req()
python/sglang/srt/managers/scheduler_pp_mixin.py
+1
View File
@@ -248,6 +248,7 @@ class SchedulerPPMixin:
self.process_prefill_chunk(
last_batch=self.last_batch, running_batch=self.running_batch
)
self._process_hicache_events()
prefill_plan = self.get_new_batch_prefill(self.running_batch)
batch = prefill_plan.batch_to_run
self.running_batch = prefill_plan.running_batch
python/sglang/srt/mem_cache/base_prefix_cache.py
+11 -2
View File
@@ -32,7 +32,9 @@ from sglang.srt.runtime_context import get_observability
if TYPE_CHECKING:
from sglang.srt.managers.cache_controller import HiCacheController
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries
from sglang.srt.mem_cache.unified_cache.cache_action import (
CacheAction,
ComponentAction,
@@ -332,6 +334,8 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
None # metrics collector for the cache
)
cache_controller: Optional[HiCacheController] = None
buffer_pipeline: Optional[BufferModePipeline] = None
storage_prefetch_retries: Optional[StoragePrefetchRetries] = None
# Set by caches that publish KV placement events; None means they don't.
kv_events: Optional[KVCacheEventRecorder] = None
@@ -489,6 +493,10 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
def swa_protected_size(self):
return 0
def swa_transient_size(self):
"""Allocated SWA tokens owned outside the request and tree views."""
return 0
def total_size(self):
raise NotImplementedError()
@@ -498,9 +506,10 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
def init_load_back(
self,
params: InitLoadBackParams,
) -> Tuple[torch.Tensor, Any]:
) -> Optional[Tuple[torch.Tensor, Any]]:
"""
Preparing KV cache loading from host to device.
Prepare host-to-device loading. None means retry admission; an empty
tensor can be a successful auxiliary-only load or a recompute fallback.
"""
raise NotImplementedError()
python/sglang/srt/mem_cache/buffer_mode/pipeline.py
+317 -231
View File
@@ -50,6 +50,7 @@ from sglang.srt.mem_cache.hicache_storage import (
SidecarPoolSpec,
)
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.storage_prefetch import StagedPrefetchPlan
from sglang.srt.mem_cache.unified_cache.cache_action import RebuildFullToSWAMapping
from sglang.srt.mem_cache.unified_cache.components import (
CacheTransferPhase,
@@ -62,6 +63,7 @@ from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
)
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.pool_host import HostPoolGroup
from sglang.srt.mem_cache.unified_cache.components import SWAComponent
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
@@ -101,7 +103,7 @@ class _StagedPrefetch(msgspec.Struct):
"""
request: CacheRequestHandle
key_tokens: list[int]
key_tokens: array
extra_key: Optional[str]
cache_salt: Optional[str]
matched_len: int
@@ -115,7 +117,8 @@ class _StagedPrefetch(msgspec.Struct):
class _OngoingBufferLoadBack(msgspec.Struct):
"""A buffer-mode load-back awaiting its H2D ack: the span is already
tree-resident; only the host bounce remains to free.
tree-resident. The host bounce and any redundant auxiliary device slots
remain owned here until the copy completes.
"""
request: CacheRequestHandle
@@ -124,13 +127,13 @@ class _OngoingBufferLoadBack(msgspec.Struct):
aux_xfers: list[PoolTransfer]
host_indices: torch.Tensor
hash_values: list[str]
aux_device_releases: list[tuple[PoolName, torch.Tensor]]
class _AnchorLock(msgspec.Struct):
"""Pins a staged prefetch's device anchor from IO commit to consumption."""
"""Pins a staged prefetch's FULL device anchor until consumption."""
node_id: NodeId
lock_params: DecLockRefParams
tokens: int
@@ -152,21 +155,6 @@ def _untrack_content_refs(refs: dict[str, int], hash_values: list[str]) -> None:
refs[h] = n
def staged_splice_tokens(f: _StagedPrefetch, device_prefix_len: int) -> int:
"""Tokens a staged prefetch can still splice beyond the live device
prefix; 0 = unusable hold (prefix shrunk below the span, span fully
device-resident, or the trim would cut into a staged aux trailing
window — aux pools splice whole or not at all)."""
span_end = f.matched_len + f.num_tokens
if device_prefix_len < f.matched_len or device_prefix_len >= span_end:
return 0
splice_tokens = span_end - device_prefix_len
for t in f.aux_xfers:
if t.host_indices is not None and t.host_indices.numel() > splice_tokens:
return 0
return splice_tokens
def validate_buffer_only_stack(
sidecar_pool_specs: list[SidecarPoolSpec],
host_pool_group: HostPoolGroup,
@@ -260,18 +248,32 @@ class BufferModePipeline:
full_pool.size - max_context_len,
),
)
logger.info(
"BufferModePipeline anchor_lock_cap_tokens=%d",
self.anchor_lock_cap_tokens,
)
if self.anchor_lock_cap_tokens == 0:
logger.warning(
"BufferModePipeline anchor_lock_cap_tokens=0 (pool=%d, "
"max_context_len=%d): every prefetch launches with its splice "
"base unpinned. Shrink --context-length or grow the KV pool.",
full_pool.size,
max_context_len,
)
else:
logger.info(
"BufferModePipeline anchor_lock_cap_tokens=%d",
self.anchor_lock_cap_tokens,
)
self.reset()
# Arbitrary: a deferral means the admission budget and the allocator disagree
# about free slots, which waiting on decode rarely fixes.
max_staged_admission_defers: int = 32
def reset(self) -> None:
# Load pipeline: hits awaiting a staging grant (park-and-retry),
# enqueue-time prefix context, completed prefetches staged until
# prefill admission, and load-backs in flight (keyed by synthetic
# negative ack id).
self.pending_hit_allocs: deque = deque()
self._staged_admission_defers: dict[CacheRequestHandle, int] = {}
self._prefetch_prefix_ctx: dict[
CacheRequestHandle, tuple[list[int], Optional[str], Optional[str]]
] = {}
@@ -310,6 +312,15 @@ class BufferModePipeline:
or self.ongoing_backup
)
def swa_transient_size(self) -> int:
"""SWA destinations kept alive only until an in-flight H2D completes."""
return sum(
len(device_indices)
for load_back in self.ongoing_buffer_load_back.values()
for pool_name, device_indices in load_back.aux_device_releases
if pool_name == PoolName.SWA
)
# ---- backup pipeline (device -> staging -> storage) ----
def _backup_parent_covered(self, state: BufferBackupState) -> bool:
@@ -744,21 +755,48 @@ class BufferModePipeline:
# ---- load back pipeline (storage -> staging -> device) ----
def try_lock_anchor(self, request: CacheRequestHandle) -> str:
def try_lock_anchor(
self, request: CacheRequestHandle, remaining_full_tokens: int
) -> tuple[str, int]:
"""Pin the staged prefetch's device anchor so eviction cannot
invalidate the splice, finding it by re-matching the live tree
(carried node ids go stale via splits and eviction; the walk is
O(prefix path)). Returns "locked", "no_anchor" (nothing to pin),
"cap_skip" (over cap; launches unlocked), or "anchor_lost" (splice
base gone — the caller cancels the storage IO)."""
if request in self.anchor_locks:
return "locked"
prefix_ctx = self._prefetch_prefix_ctx.get(request)
if not prefix_ctx or not prefix_ctx[0]:
return "no_anchor" # root anchor: nothing to pin
prefix_tokens, extra_key, cache_salt = prefix_ctx
O(request hit span)). Returns "locked", "no_anchor" (nothing to pin),
"cap_skip" (bigger than the whole cap; launches unlocked), "cap_busy"
(fits, but the budget is taken; the caller parks), or "anchor_lost"
(splice base gone -- the caller re-plans)."""
assert request not in self.anchor_locks, (
f"prefetch anchor already locked: {request.rid}"
)
prefix_tokens, extra_key, cache_salt = self._prefetch_prefix_ctx[request]
matched_len = len(prefix_tokens)
if self.anchor_locked_tokens_ + matched_len > self.anchor_lock_cap_tokens:
assert matched_len + remaining_full_tokens > 0, (
f"empty prefetch span: {request.rid}"
)
full_key_tokens = array("q", prefix_tokens)
if remaining_full_tokens or self._cache.tree_core.is_eagle:
info = self._cache.ongoing_prefetch[request]
raw_len = remaining_full_tokens + int(info.prefetch_key.is_bigram)
full_key_tokens.extend(info.prefetch_key.token_ids[:raw_len])
cache = self._cache
matched_full, anchor_node, anchor_tokens = (
cache.tree_core.match_full_device_prefix(
RadixKey(
full_key_tokens,
extra_key=extra_key,
is_bigram=cache.tree_core.is_eagle,
cache_salt=cache_salt,
)
)
)
if matched_full < matched_len:
return "anchor_lost", matched_full
if anchor_tokens == 0:
return "no_anchor", matched_full
if self.anchor_locked_tokens_ + anchor_tokens > self.anchor_lock_cap_tokens:
# Parking for a pin no drain can ever satisfy deadlocks the
# request, so only a pin that still fits the cap is worth a wait.
over_cap = anchor_tokens > self.anchor_lock_cap_tokens
self._anchor_lock_cap_skips += 1
if (
self._anchor_lock_cap_skips <= 3
@@ -766,42 +804,21 @@ class BufferModePipeline:
):
logger.warning(
"HiCache anchor-lock cap reached (skip %d): locked=%d "
"want=%d cap=%d; launching unlocked.",
"want=%d cap=%d; %s.",
self._anchor_lock_cap_skips,
self.anchor_locked_tokens_,
matched_len,
self.anchor_lock_cap_tokens,
)
return "cap_skip"
cache = self._cache
anchor_tokens = array("q", prefix_tokens)
if cache.tree_core.is_eagle:
# The suffix owns the boundary token shared with the last matched
# bigram, so include it when rebuilding the anchor key.
info = cache.ongoing_prefetch.get(request)
if info is None or not info.prefetch_key.token_ids:
return "anchor_lost"
anchor_tokens.append(info.prefetch_key.token_ids[0])
match = cache.match_prefix(
MatchPrefixParams(
key=RadixKey(
anchor_tokens,
extra_key=extra_key,
is_bigram=cache.tree_core.is_eagle,
cache_salt=cache_salt,
self.anchor_lock_cap_tokens,
"launching unlocked" if over_cap else "parking",
)
)
)
if len(match.device_indices) < matched_len:
return "anchor_lost"
lock_params = cache.inc_lock_ref(match.last_device_node).to_dec_params()
return ("cap_skip" if over_cap else "cap_busy"), matched_full
cache.tree_core.inc_full_pin(anchor_node)
self.anchor_locks[request] = _AnchorLock(
node_id=match.last_device_node,
lock_params=lock_params,
tokens=matched_len,
node_id=anchor_node,
tokens=anchor_tokens,
)
self.anchor_locked_tokens_ += matched_len
return "locked"
self.anchor_locked_tokens_ += anchor_tokens
return "locked", matched_full
def release_anchor_lock(self, request: CacheRequestHandle) -> None:
"""Drop a staged prefetch's anchor lock (idempotent; called at every
@@ -809,7 +826,7 @@ class BufferModePipeline:
lock = self.anchor_locks.pop(request, None)
if lock is None:
return
self._cache.dec_lock_ref(lock.node_id, lock.lock_params)
self._cache.tree_core.dec_full_pin(lock.node_id)
self.anchor_locked_tokens_ -= lock.tokens
assert self.anchor_locked_tokens_ >= 0, (
f"anchor-lock accounting corrupted: locked={self.anchor_locked_tokens_} "
@@ -861,6 +878,85 @@ class BufferModePipeline:
def has_staged(self, request: CacheRequestHandle) -> bool:
return request in self.staged_prefetches
def prepare_staged_prefetch(self, req: Req) -> bool:
"""Rebuild the admission plan from this pass's joint FULL/SWA match."""
req.staged_prefetch_plan = None
f = self.staged_prefetches.get(req.cache_request_handle)
if f is None:
if not (req.host_hit_is_storage and req.host_loaded_length > 0):
self._clear_storage_hit(req)
return True
if len(req.prefix_indices) >= f.matched_len + f.num_tokens:
# The joint match already covers the staged span; a shorter FULL-only
# prefix would strand the slots recomputed below cache_protected_len.
self._resolve_device_covered(req, f)
return True
key = RadixKey(
f.key_tokens,
extra_key=f.extra_key,
is_bigram=self._cache.tree_core.is_eagle,
cache_salt=f.cache_salt,
)
matched_len, node_id, _ = self._cache.tree_core.match_full_device_prefix(key)
if matched_len < f.matched_len:
logger.warning(
"HiCache staged prefetch deferred req=%s reason=shrunk "
"matched=%d now=%d tokens=%d",
req.rid,
f.matched_len,
matched_len,
f.num_tokens,
)
self._refetch_staged(f)
return False
root_id = self._cache.tree_core.empty_match_result.last_device_node
full_indices = self._cache.tree_core.collect_full_device_indices(
node_id, root_id
)[:matched_len]
assert len(full_indices) == matched_len
req.prefix_indices = full_indices
req.last_node = node_id
req.kv.cache_protected_len = matched_len
full_tokens = max(0, f.matched_len + f.num_tokens - matched_len)
swa_tokens = sum(
len(t.host_indices)
for t in f.aux_xfers
if t.name == PoolName.SWA and t.host_indices is not None
)
if full_tokens == 0 and swa_tokens == 0:
self._resolve_device_covered(req, f)
return True
req.host_hit_length = full_tokens
req.swa_host_hit_length = swa_tokens
req.storage_hit_length = full_tokens
req.storage_hit_start = matched_len if full_tokens else None
req.host_hit_is_storage = True
req.staged_prefetch_plan = StagedPrefetchPlan(
f.operation_id, key, matched_len, full_tokens, swa_tokens
)
return True
def _resolve_device_covered(self, req: Req, f: _StagedPrefetch) -> None:
req.host_hit_length = 0
req.swa_host_hit_length = 0
self._clear_storage_hit(req)
self._cache._resolve_storage_prefetch_tokens(
req.cache_request_handle, f.num_tokens, reason="device_covered"
)
self.release_staged_hold(req.cache_request_handle, reason=None)
@staticmethod
def _clear_storage_hit(req: Req) -> None:
req.storage_hit_length = 0
req.storage_hit_start = None
req.host_hit_is_storage = False
def _refetch_staged(self, f: _StagedPrefetch) -> None:
self.release_staged_hold(f.request, reason="shrunk")
self._cache.storage_prefetch_retries.refetch(
f.request.rid, f.matched_len + f.num_tokens
)
@staticmethod
def _occupied_span(host_indices) -> int:
"""Occupancy units a buffer-mode prefetch holds: granted at
@@ -877,14 +973,11 @@ class BufferModePipeline:
it as host_hit_length and the adder consumes it via init_load_back.
Always returns True (ready is a stable, revisited state)."""
cache = self._cache
(
_anchor,
prefetch_key,
host_indices,
operation,
_lock_params,
comp_xfers,
) = cache.ongoing_prefetch.pop(request)
info = cache.ongoing_prefetch.pop(request)
prefetch_key = info.prefetch_key
host_indices = info.host_indices
operation = info.operation
comp_xfers = info.comp_xfers
cc = cache.cache_controller
prefix_ctx = self._prefetch_prefix_ctx.pop(request, None)
prefix_tokens = prefix_ctx[0] if prefix_ctx is not None else None
@@ -898,7 +991,10 @@ class BufferModePipeline:
if transfer.indices_from_pool is not None
)
if num_tokens == 0 or prefix_tokens is None:
has_aux = any(
t.host_indices is not None and t.host_indices.numel() > 0 for t in aux_xfers
)
if (num_tokens == 0 and not has_aux) or prefix_tokens is None:
# Nothing usable fetched: recompute.
cache.discard_storage_prefetch_accounting(request)
self.release_anchor_lock(request)
@@ -921,7 +1017,13 @@ class BufferModePipeline:
self.staged_prefetches[request] = _StagedPrefetch(
request=request,
key_tokens=prefix_tokens + list(prefetch_key[:num_tokens].token_ids),
key_tokens=array(
"q",
prefix_tokens
+ list(
prefetch_key.token_ids[: num_tokens + int(prefetch_key.is_bigram)]
),
),
extra_key=prefetch_key.extra_key,
cache_salt=prefetch_key.cache_salt,
matched_len=len(prefix_tokens),
@@ -936,62 +1038,19 @@ class BufferModePipeline:
cache.prefetch_loaded_storage_start_by_reqid[request] = operation.storage_start
return True
def plan_staged_splice(
self, request: CacheRequestHandle, device_prefix_len: int
) -> tuple[int, int]:
"""(kv, swa) host-hit tokens consumption will splice given the
request's live device prefix, so admission charges no phantom
tokens. Frees a hold that can no longer splice: surfaced as 0 but
kept, it would leak — the adder only consumes surfaced host hits."""
f = self.staged_prefetches.get(request)
if f is None:
return 0, 0
splice_tokens = staged_splice_tokens(f, device_prefix_len)
if splice_tokens == 0:
covered_tokens = self._resolve_staged_device_coverage(f, device_prefix_len)
logger.info(
"HiCache staged prefetch released req=%s matched=%d "
"device_prefix=%d tokens=%d",
request.rid,
f.matched_len,
device_prefix_len,
f.num_tokens,
)
reason = None if covered_tokens == f.num_tokens else "shrunk"
self.release_staged_hold(request, reason=reason)
return 0, 0
return splice_tokens, self.staged_prefetch_swa_tokens(request)
def init_load_back(
self, params: InitLoadBackParams
) -> Optional[tuple[torch.Tensor, NodeId]]:
"""Materialize a selected prefill under the caller's prefix lock.
def _resolve_staged_device_coverage(
self, f: _StagedPrefetch, device_prefix_len: int
) -> int:
covered_tokens = min(max(device_prefix_len - f.matched_len, 0), f.num_tokens)
self._cache._resolve_storage_prefetch_tokens(f.request, covered_tokens)
return covered_tokens
def staged_prefetch_swa_tokens(self, request: CacheRequestHandle) -> int:
"""SWA device tokens consuming this staged prefetch will allocate (the
staged trailing window); surfaced as the request's swa_host_hit_length
so the adder's SWA gate charges the admission-time alloc."""
f = self.staged_prefetches.get(request)
if f is None:
return 0
return sum(
len(t.host_indices)
for t in f.aux_xfers
if t.name == PoolName.SWA and t.host_indices is not None
)
def init_load_back(self, params: InitLoadBackParams) -> tuple[torch.Tensor, NodeId]:
"""Consume the staged prefetch at prefill admission: device alloc,
layer-gated H2D, and a plain insert so downstream sees ordinary tree
state. The splice base is the request's live device prefix — growth
trims to the span tail beyond it; unusable holds drop and the
request recomputes.
The caller has finished selecting its prefill shape and must acquire
the request lock after success, without further admission gates. The
prefix lock protects allocation-time eviction. None retains staging
and its anchor for the next admission attempt.
Ownership contract: cc.load queues the H2D before insert adjudicates
ownership, so the live pre-checks below must prove the insert can
only ADD nodes — a dedup would free slots the in-flight copy still
ownership, so the prepared boundary must ensure the insert can only
ADD nodes — a dedup would free slots the in-flight copy still
targets (queued use-after-free)."""
cache = self._cache
req = params.req
@@ -999,96 +1058,58 @@ class BufferModePipeline:
request = req.cache_request_handle
empty = cache.tree_core.empty_match_result.device_indices
unchanged = (empty, req.last_node)
f = self.staged_prefetches.pop(request, None)
f = self.staged_prefetches.get(request)
if f is None:
self.release_anchor_lock(request)
return unchanged
cc = cache.cache_controller
plan = req.staged_prefetch_plan
assert plan is not None, f"staged prefetch was not planned for {req.rid}"
assert f.operation_id == plan.operation_id
assert (f.extra_key, f.cache_salt) == (req.extra_key, req.cache_salt)
assert (req.host_hit_length, req.swa_host_hit_length) == (
plan.full_tokens,
plan.swa_tokens,
), f"staged load-back budget changed for {req.rid}"
def _drop(reason: Optional[str]) -> tuple[torch.Tensor, NodeId]:
cache._finish_storage_prefetch(request, fulfilled_tokens=0, reason=reason)
self.release_anchor_lock(request)
self._free_staging_now(f.host_indices, f.aux_xfers)
cc.prefetch_tokens_occupied -= f.occupied_tokens
# Nothing spliced: keep the surfaced host-hit fields truthful.
req.host_hit_length = 0
req.swa_host_hit_length = 0
req.storage_hit_length = 0
req.storage_hit_start = None
req.host_hit_is_storage = False
return unchanged
# A hold staged under a different namespace than the consuming request
# must never splice (wrong-namespace publish = duplicate slot
# ownership); unreachable while the prefetch key is request-derived.
if f.extra_key != req.extra_key or f.cache_salt != req.cache_salt:
logger.error(
"HiCache staged prefetch dropped req=%s reason=namespace "
"staged=%s req=%s",
req.rid,
(f.extra_key, f.cache_salt),
(req.extra_key, req.cache_salt),
)
return _drop("dropped")
splice_base = len(req.prefix_indices)
splice_tokens = staged_splice_tokens(f, splice_base)
if splice_tokens == 0:
covered_tokens = self._resolve_staged_device_coverage(f, splice_base)
def _defer_for_capacity(pool: str) -> None:
defers = self._staged_admission_defers.get(request, 0) + 1
self._staged_admission_defers[request] = defers
cache._log_storage_prefetch_deferred(f.num_tokens, "device_capacity")
if defers < self.max_staged_admission_defers:
logger.warning(
"HiCache staged prefetch deferred at admission req=%s "
"reason=device_capacity pool=%s tokens=%d defers=%d",
req.rid,
pool,
f.num_tokens,
defers,
)
return
# Still unmaterializable: drop the hold so the admission loop stops
# breaking on this request, which recomputes on its next pass.
logger.warning(
"HiCache staged prefetch dropped req=%s matched=%d now=%d "
"tokens_wasted=%d locked=%s",
"HiCache staged prefetch dropped after %d device_capacity "
"deferrals req=%s pool=%s tokens=%d",
defers,
req.rid,
f.matched_len,
splice_base,
pool,
f.num_tokens,
request in self.anchor_locks,
)
reason = None if covered_tokens == f.num_tokens else "shrunk"
return _drop(reason)
self.release_staged_hold(request, reason="device_capacity")
req.staged_prefetch_plan = None
splice_base = plan.device_prefix_len
assert len(req.prefix_indices) == splice_base
trim_tokens = splice_base - f.matched_len
assert trim_tokens % cache.page_size == 0, (
f"staged splice trim not page-aligned req={req.rid}: "
f"matched={f.matched_len} splice_base={splice_base}"
)
cache._resolve_storage_prefetch_tokens(request, trim_tokens)
key = RadixKey(
array("q", f.key_tokens),
extra_key=f.extra_key,
is_bigram=cache.tree_core.is_eagle,
cache_salt=f.cache_salt,
).page_aligned(cache.page_size)
key = plan.key
span_end = f.matched_len + f.num_tokens
# Live ownership pre-check at the splice base: the unified length
# detects a stale request view (req matched before a later publish),
# full_kv_hit_length detects FULL overlap the insert would dedup-free
# (an SWA tombstone can mask live FULL from the unified match alone).
live = cache.match_prefix(MatchPrefixParams(key=key))
if (
len(live.device_indices) != splice_base
or live.full_kv_hit_length != splice_base
):
logger.warning(
"HiCache staged prefetch dropped req=%s reason=overlap "
"splice_base=%d live_unified=%d live_full=%d tokens_wasted=%d "
"locked=%s",
req.rid,
splice_base,
len(live.device_indices),
live.full_kv_hit_length,
f.num_tokens,
request in self.anchor_locks,
)
available_end = min(
span_end,
len(live.device_indices),
live.full_kv_hit_length,
)
available_overlap = max(0, available_end - splice_base)
cache._resolve_storage_prefetch_tokens(request, available_overlap)
return _drop(None if available_overlap == splice_tokens else "shrunk")
load_tokens = plan.full_tokens
# Evict-before-alloc (mirrors _load_back_transfers): the budget gate
# counts evictable pages, but cc.load draws from free slots only.
@@ -1096,87 +1117,147 @@ class BufferModePipeline:
avail = cache.token_to_kv_pool_allocator.full_available_size()
else:
avail = cache.token_to_kv_pool_allocator.available_size()
if avail < splice_tokens:
needed = splice_tokens - avail
if avail < load_tokens:
needed = load_tokens - avail
cache.evict_for_alloc(EvictParams(num_tokens=needed))
if cache.supports_swa():
avail = cache.token_to_kv_pool_allocator.full_available_size()
else:
avail = cache.token_to_kv_pool_allocator.available_size()
if avail < splice_tokens:
# Genuinely no room (locked pages): recompute.
return _drop("device_capacity")
if avail < load_tokens:
return _defer_for_capacity("full")
load_back_id = -(f.operation_id) - 1
# The full trailing-window aux transfer is independent of the shorter
# FULL suffix and may remain nonempty for an aux-only load.
load_xfers = list(f.aux_xfers)
staged_swa = next(
(
len(t.host_indices)
for t in load_xfers
if t.name == PoolName.SWA and t.host_indices is not None
),
0,
)
device_indices = cc.load(
host_indices=f.host_indices[trim_tokens:],
node_id=load_back_id,
extra_pools=f.aux_xfers or None,
extra_pools=load_xfers or None,
)
if device_indices is None:
# Transient allocator shortfall despite the evict: recompute
# (init_load_back's degrade contract).
return _drop("device_capacity")
# load() allocates all pools atomically before queueing H2D, so the
# staged host buffers remain reusable after either pool is short.
return _defer_for_capacity("full_or_aux")
del self.staged_prefetches[request]
self._staged_admission_defers.pop(request, None)
req.staged_prefetch_plan = None
cache._resolve_storage_prefetch_tokens(
request, trim_tokens, reason="device_covered"
)
swa_dev = next(
(
t.device_indices
for t in f.aux_xfers
for t in load_xfers
if t.name == PoolName.SWA
and t.device_indices is not None
and t.device_indices.numel() > 0
),
None,
)
aux_device_releases: list[tuple[PoolName, torch.Tensor]] = []
if swa_dev is not None:
# Register the trailing window's FULL->SWA translation NOW: the
# admitted request's attention reads the window through this
# mapping during the layer-gated forward.
cache._apply_cache_action(
RebuildFullToSWAMapping([device_indices[-len(swa_dev) :]], [swa_dev])
# Register the window's FULL->SWA translation now (attention reads
# through it). Keep SWA slots another request may still hold; their
# redundant H2D destinations are reclaimed at the transfer ack.
full_window = torch.cat([req.prefix_indices, device_indices])[
-len(swa_dev) :
]
allocator = cache.token_to_kv_pool_allocator
old_swa = allocator.full_to_swa_index_mapping[full_window.to(torch.int64)]
missing = old_swa <= 0
window_start = span_end - len(swa_dev)
repair_end = min(splice_base, span_end)
tree_missing = torch.zeros_like(missing)
tail_start = max(splice_base, window_start)
tree_missing[tail_start - window_start :] = True
repair_ranges = []
if window_start < repair_end:
repair_ranges = cache.tree_core.swa_tombstone_ranges(
key, window_start, repair_end
)
for repair_start, repair_end_ in repair_ranges:
repair_slice = slice(
repair_start - window_start, repair_end_ - window_start
)
tree_missing[repair_slice] = True
assert torch.equal(missing, tree_missing), (
"SWA tree and allocator residency disagree for restored window "
f"[{window_start}, {span_end})"
)
for repair_start, repair_end_ in repair_ranges:
repair_slice = slice(
repair_start - window_start, repair_end_ - window_start
)
for action in cache.tree_core.attach_swa_window(
key,
repair_start,
repair_end_,
swa_dev[repair_slice],
):
cache._apply_cache_action(action)
if bool(missing.any()):
cache._apply_cache_action(
RebuildFullToSWAMapping(
[full_window[missing]],
[swa_dev[missing]],
)
)
if bool((~missing).any()):
aux_device_releases.append((PoolName.SWA, swa_dev[~missing]))
# Publish via a plain insert under the admission lock choreography;
# the caller's request lock then pins the span (load_back pattern).
# prev_prefix_len covers the already-device-resident head.
insert_result = cache.insert(
InsertParams(
key=key,
value=torch.cat([req.prefix_indices, device_indices]),
prev_prefix_len=splice_base,
swa_evicted_seqlen=(
max(0, span_end - len(swa_dev)) if swa_dev is not None else 0
),
swa_evicted_seqlen=(span_end - staged_swa) if staged_swa else 0,
)
)
self.ongoing_buffer_load_back[load_back_id] = _OngoingBufferLoadBack(
request=f.request,
num_tokens=splice_tokens,
num_tokens=load_tokens,
occupied_tokens=f.occupied_tokens,
aux_xfers=f.aux_xfers,
# The full staged bounce (not the trimmed H2D source): the ack
# frees it whole, trimmed head included.
host_indices=f.host_indices,
hash_values=f.hash_values,
aux_device_releases=aux_device_releases,
)
m = cache.match_prefix(MatchPrefixParams(key=key))
match = cache.match_prefix(MatchPrefixParams(key=key))
self.release_anchor_lock(request)
canonical = m.device_indices[splice_base:span_end]
if len(m.device_indices) < span_end or not torch.equal(
canonical, device_indices
):
canonical = match.device_indices[splice_base:span_end]
owned = len(match.device_indices) >= span_end and torch.equal(
match.device_indices[splice_base:span_end], device_indices
)
if not owned:
# Fail-stop: the insert freed or replaced slots the in-flight H2D
# still targets; continuing risks silent KV corruption.
raise RuntimeError(
f"HiCache buffer load-back ownership violation req={f.request.rid}: "
"HiCache buffer load-back ownership violation "
f"req={f.request.rid}: "
f"insert prefix_len={insert_result.prefix_len} "
f"expected={splice_base}, adopted={len(m.device_indices)} "
f"span_end={span_end}, canonical_matches_incoming="
f"{len(m.device_indices) >= span_end and torch.equal(canonical, device_indices)}; "
f"expected={splice_base}, matched={len(match.device_indices)} "
f"span_end={span_end} splice_base={splice_base}; "
f"in-flight H2D targets freed slots"
)
# Canonical ownership: return the post-insert tree slice, never the
# raw cc.load allocation (torch.equal here; the tree slice is truth).
return canonical, m.last_device_node
return canonical, match.last_device_node
def try_finish_load_back(self, ack_id: int) -> bool:
"""Fill ack: free the host bounce and return True when the ack id is
@@ -1191,6 +1272,10 @@ class BufferModePipeline:
# The H2D consumed the bounce buffers; free them outright.
self._free_staging_now(f.host_indices, f.aux_xfers)
for pool_name, device_indices in f.aux_device_releases:
entry = cc.mem_pool_host.entry_map[pool_name]
free_fn = entry.device_free_fn or entry.device_pool.free
free_fn(device_indices)
cc.prefetch_tokens_occupied -= f.occupied_tokens
logger.info(
@@ -1213,6 +1298,7 @@ class BufferModePipeline:
and for holds that can no longer splice. Returns True when a hold
existed."""
self.release_anchor_lock(request)
self._staged_admission_defers.pop(request, None)
staged = self.staged_prefetches.pop(request, None)
if staged is None:
return False
python/sglang/srt/mem_cache/hiradix_cache.py
+1 -5
View File
@@ -1738,11 +1738,6 @@ class HiRadixCache(RadixCache):
"""
return self.prefetch_loaded_tokens_by_reqid.pop(handle.rid, 0)
def pop_storage_prefetch_miss(self, handle: CacheRequestHandle) -> bool:
"""Storage prefetch miss markers are not tracked on the dense path;
the scheduler's paced availability-check retry is inert here."""
return False
def match_prefix(self, params: MatchPrefixParams):
if self.disable:
return self._empty_match_result
@@ -1787,6 +1782,7 @@ class HiRadixCache(RadixCache):
matched_prefix_tokens: Optional[List[int]] = None,
extra_key: Optional[str] = None,
cache_salt: Optional[str] = None,
storage_hit_end: Optional[int] = None,
):
req_id = handle.rid
prefetch_key = RadixKey(
python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py
+43 -4
View File
@@ -68,6 +68,7 @@ class PrefetchOperation(StorageOperation):
last_hash: Optional[str] = None,
prefix_keys: Optional[List[str]] = None,
pool_transfers: Optional[list[PoolTransfer]] = None,
assume_stored: bool = False,
):
self.handle = handle
self.request_id = handle.rid
@@ -75,6 +76,9 @@ class PrefetchOperation(StorageOperation):
self._terminated_flag = False
self.storage_hit_count = 0
self.start_time = time.monotonic()
# Take the whole span as present instead of querying for it; the read
# itself is fail-soft, so a wrong guess shortens the fetch.
self.assume_stored = assume_stored
super().__init__(
None,
token_ids,
@@ -83,6 +87,13 @@ class PrefetchOperation(StorageOperation):
pool_transfers=pool_transfers,
)
self.pool_transfers_done = not bool(pool_transfers)
# The Python transfer worker leaves the unfinished tail to the ACK drain;
# a controller that releases it itself must set this False.
self.ack_releases_incomplete_host_indices = True
# Buffer mode may trim already-device-resident FULL pages after the
# query. Trailing sidecars still use the untrimmed hit endpoint.
self.sidecar_hash_values: Optional[list[str]] = None
self.sidecar_hit_pages = 0
def mark_terminate(self):
with self._lock:
@@ -550,6 +561,7 @@ class HybridCacheController(BaseHiCacheController):
last_hash: Optional[str] = None,
prefix_keys: Optional[List[str]] = None,
extra_pools: Optional[list[PoolTransfer]] = None,
assume_stored: bool = False,
) -> PrefetchOperation:
operation = PrefetchOperation(
handle,
@@ -557,6 +569,7 @@ class HybridCacheController(BaseHiCacheController):
last_hash,
prefix_keys=prefix_keys,
pool_transfers=extra_pools,
assume_stored=assume_stored,
)
self.prefetch_queue.put(operation)
return operation
@@ -585,6 +598,13 @@ class HybridCacheController(BaseHiCacheController):
)
operation.all_hash_values = hash_value
if operation.assume_stored:
# A prior hit on a suffix of this span proved it stored, and writes
# are prefix-covered, so re-querying only adds a round trip.
kv_hit_pages = len(hash_value)
operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages)
return hash_value, kv_hit_pages * self.page_size
extra_info = HiCacheStorageExtraInfo(
prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None
)
@@ -661,9 +681,13 @@ class HybridCacheController(BaseHiCacheController):
for transfer in operation.pool_transfers
if transfer.indices_from_pool != PoolName.KV
]
self._sync_trailing_keys(
transfers_nonkv, operation.hash_value, kv_completed_pages
sidecar_hashes = operation.sidecar_hash_values or operation.hash_value
sidecar_hit_pages = (
operation.sidecar_hit_pages
if operation.sidecar_hash_values is not None
else kv_completed_pages
)
self._sync_trailing_keys(transfers_nonkv, sidecar_hashes, sidecar_hit_pages)
self._resolve_sidecar_nonkv_derived_pool_transfers(operation)
results = self.storage_backend.batch_get_v2(transfers_nonkv)
pool_hits = count_pool_hits(results)
@@ -679,6 +703,22 @@ class HybridCacheController(BaseHiCacheController):
)
return
def trim_prefetch_full_head(self, operation, trim_tokens: int) -> None:
"""Drop a device-resident FULL head after the availability query;
trailing sidecars keep the untrimmed hit endpoint."""
if trim_tokens <= 0:
return
assert trim_tokens % self.page_size == 0
trim_pages = trim_tokens // self.page_size
original_hashes = list(operation.hash_value)
assert trim_pages <= len(original_hashes)
if operation.sidecar_hash_values is None:
operation.sidecar_hash_values = original_hashes
operation.sidecar_hit_pages = len(original_hashes)
operation.hash_value = original_hashes[trim_pages:]
operation.storage_hit_count -= trim_tokens
operation.storage_start += trim_tokens
def _page_backup(self, operation):
# MLA KV is replicated across TP ranks and should still be written only
# by TP0. Rank-sharded sidecars still need every TP rank.
@@ -761,8 +801,7 @@ class HybridCacheController(BaseHiCacheController):
for transfer in operation.pool_transfers:
if transfer.indices_from_pool == PoolName.KV:
transfer.host_indices = operation.host_indices
if transfer.keys is None:
transfer.keys = operation.hash_value
transfer.keys = operation.hash_value
def _resolve_sidecar_nonkv_derived_pool_transfers(self, operation):
for transfer in operation.pool_transfers:
python/sglang/srt/mem_cache/rust_tree_core/adapter.py
+37
View File
@@ -405,6 +405,26 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
def root_node(self) -> UnifiedTreeNode:
raise NotImplementedError("root_node: not yet ported to the Rust tree core")
def swa_tombstone_ranges(
self, key: RadixKey, start: int, end: int
) -> list[tuple[int, int]]:
raise NotImplementedError(
"swa_tombstone_ranges: buffer-mode SWA window repair is not yet "
"ported to the Rust tree core"
)
def attach_swa_window(
self,
key: RadixKey,
window_start: int,
window_end: int,
swa_values: torch.Tensor,
) -> list:
raise NotImplementedError(
"attach_swa_window: buffer-mode SWA window repair is not yet "
"ported to the Rust tree core"
)
def inc_lock_ref(
self,
node_id: NodeId,
@@ -556,6 +576,21 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
)
return _match_result_from_binding(result)
def match_full_device_prefix(self, key: RadixKey) -> tuple[int, NodeId, int]:
return self._binding.match_full_device_prefix(
self._bindings.MatchParamsBinding(
key=_radix_key_buffer(key),
extra_key=key.extra_key,
cache_salt=key.cache_salt,
)
)
def inc_full_pin(self, node_id: NodeId) -> None:
self._binding.inc_full_pin(node_id)
def dec_full_pin(self, node_id: NodeId) -> None:
self._binding.dec_full_pin(node_id)
@property
def empty_match_result(self) -> MatchResult:
return self._empty_match_result
@@ -747,6 +782,7 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0,
staging_tokens: int = 0,
last_hash: Optional[str] = None,
) -> Optional[list[PoolTransfer]]:
transfers = self._binding.build_hicache_transfers(
@@ -757,6 +793,7 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
# TODO: Forward token ids when Rust Mamba prefetch consumes them.
None,
prefetch_tokens,
staging_tokens,
last_hash,
)
if transfers is None:
python/sglang/srt/mem_cache/storage_prefetch.py
+85
View File
@@ -0,0 +1,85 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.radix_cache import RadixKey
@dataclass
class _StoragePrefetchRetry:
immediate: bool
storage_hit_end: Optional[int] = None
due_step: Optional[int] = None
class StoragePrefetchRetries:
"""New L3 attempts only (parked IO and retained staging retry at their owner);
step deadlines keep TP ranks in lockstep, retries follow queue order, and a
request past its re-issue budget is admitted with what the device holds."""
def __init__(self):
self._pending: dict[str, _StoragePrefetchRetry] = {}
self._step = 0
def poll_miss(self, req_id: str, storage_hit_end: Optional[int] = None) -> None:
self._pending[req_id] = _StoragePrefetchRetry(False, storage_hit_end)
def refetch(self, req_id: str, storage_hit_end: Optional[int] = None) -> None:
self._pending[req_id] = _StoragePrefetchRetry(True, storage_hit_end)
def cancel(self, req_id: str) -> None:
self._pending.pop(req_id, None)
def clear(self) -> None:
self._pending.clear()
def pop_ready(
self, waiting_queue: list[Req], interval: int, max_attempts: int
) -> list[tuple[Req, Optional[int]]]:
self._step += 1
if not waiting_queue or not self._pending:
return []
# A speculative miss must never delay the queue head.
head_id = waiting_queue[0].rid
head_retry = self._pending.get(head_id)
if head_retry is not None and not head_retry.immediate:
self.cancel(head_id)
if not any(
retry.due_step is None or retry.due_step <= self._step
for retry in self._pending.values()
):
return []
ready = []
for req in waiting_queue:
retry = self._pending.get(req.rid)
if retry is None:
continue
if req.storage_prefetch_retry_attempts >= max_attempts:
self.cancel(req.rid)
continue
if not retry.immediate:
if interval <= 0:
self.cancel(req.rid)
continue
if retry.due_step is None:
retry.due_step = self._step + interval
if retry.due_step > self._step:
continue
self.cancel(req.rid)
ready.append((req, retry.storage_hit_end))
return ready
@dataclass(frozen=True)
class StagedPrefetchPlan:
"""One admission pass, computed before promoting the joint FULL/SWA match."""
operation_id: int
key: RadixKey
device_prefix_len: int
full_tokens: int
swa_tokens: int
python/sglang/srt/mem_cache/unified_cache/components/base.py
+10 -6
View File
@@ -73,12 +73,11 @@ class PrepareLoadBackResult:
@dataclasses.dataclass(frozen=True)
class PreparePrefetchResult:
"""Outcome of prepare_prefetch; default = nothing to prepare."""
"""Outcome of prepare_prefetch; default = the component takes no part."""
# Host pool exhausted; the caller aborts the prefetch.
alloc_failed: bool = False
# The component's pre-allocated host buffer (None = skip the build).
host_indices: Optional[torch.Tensor] = None
# Host staging the fetch needs from this component, in the component
# pool's units. Allocated at hit time, next to the KV staging; 0 = none.
staging_tokens: int = 0
class CacheTransferPhase(str, Enum):
@@ -666,9 +665,13 @@ class TreeComponent(ABC):
*,
prefetch_tokens: int = 0,
) -> PreparePrefetchResult:
"""Cache-level host pre-allocation before a prefetch builds its transfers."""
"""Size the host staging a prefetch from node_id needs from this component."""
return PreparePrefetchResult()
def alloc_prefetch_staging(self, num_tokens: int) -> Optional[torch.Tensor]:
"""Allocate prefetch staging sized by prepare_prefetch, once the hit is known."""
return None
def build_hicache_transfers(
self,
node: UnifiedTreeNode,
@@ -678,6 +681,7 @@ class TreeComponent(ABC):
host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0,
staging_tokens: int = 0,
last_hash: Optional[str] = None,
) -> Optional[list[PoolTransfer]]:
"""Build transfer descriptors for this component in the given phase.
python/sglang/srt/mem_cache/unified_cache/components/full.py
+1
View File
@@ -350,6 +350,7 @@ class FullComponent(TreeComponent):
host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0,
staging_tokens: int = 0,
last_hash: Optional[str] = None,
) -> Optional[list[PoolTransfer]]:
ct = self.component_type
python/sglang/srt/mem_cache/unified_cache/components/mamba.py
+13 -9
View File
@@ -693,14 +693,15 @@ class MambaComponent(TreeComponent):
*,
prefetch_tokens: int = 0,
) -> PreparePrefetchResult:
host_indices = self.cache.host_pool_group.alloc(
1,
pool=PoolName.MAMBA,
reclaim=lambda size: self.cache.evict_host(size, ComponentType.MAMBA),
)
# One state slot per fetch, allocated once the hit is known.
return PreparePrefetchResult(staging_tokens=1)
def alloc_prefetch_staging(self, num_tokens: int) -> Optional[torch.Tensor]:
host_indices = self._mamba_pool_host.alloc(num_tokens)
if host_indices is None:
return PreparePrefetchResult(alloc_failed=True)
return PreparePrefetchResult(host_indices=host_indices)
self.cache.evict_host(num_tokens, ComponentType.MAMBA)
host_indices = self._mamba_pool_host.alloc(num_tokens)
return host_indices
def build_hicache_transfers(
self,
@@ -711,6 +712,7 @@ class MambaComponent(TreeComponent):
host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0,
staging_tokens: int = 0,
last_hash: Optional[str] = None,
) -> Optional[list[PoolTransfer]]:
ct = self.component_type
@@ -770,11 +772,13 @@ class MambaComponent(TreeComponent):
]
if phase == CacheTransferPhase.PREFETCH:
assert host_indices is not None
if staging_tokens == 0:
return None
# Staging is allocated once the hit is known; the placeholder key
# carries the single trailing page this pool loads.
return [
PoolTransfer(
name=PoolName.MAMBA,
host_indices=host_indices,
keys=["__placeholder__"],
hit_policy=PoolHitPolicy.TRAILING_PAGES,
)
python/sglang/srt/mem_cache/unified_cache/components/swa.py
+19 -23
View File
@@ -1007,28 +1007,23 @@ class SWAComponent(TreeComponent):
prefetch_pages = prefetch_tokens // self.cache.page_size
if prefetch_pages >= sw_pages:
num_pages = sw_pages
elif prefetch_pages <= 0:
return PreparePrefetchResult()
elif (
self.tree_core.is_root(node_id) or self.tree_core.is_host_memory_buffer_only
):
# Sub-window fetch: at root the sequence IS its window; mid-tree
# (buffer mode) the window head is the device prefix's own ring
# state, so only the suffix needs fetching.
elif prefetch_pages > 0 and self.tree_core.is_root(node_id):
# At root the sequence is shorter than the window, so its whole
# SWA is the window -- complete, not a partial fetch.
num_pages = prefetch_pages
else:
# Cache-mode graft: a mid-tree window head is not
# device-guaranteed, require a full window.
# Mid-tree short span: the window head would have to come from the
# device prefix, which the match validator does not promise.
return PreparePrefetchResult()
num_tokens = num_pages * self.cache.page_size
host_indices = self.cache.host_pool_group.alloc(
num_tokens,
pool=PoolName.SWA,
reclaim=lambda size: self.cache.evict_host(size, ComponentType.SWA),
)
return PreparePrefetchResult(staging_tokens=num_pages * self.cache.page_size)
def alloc_prefetch_staging(self, num_tokens: int) -> Optional[torch.Tensor]:
assert self._swa_kv_pool_host is not None
host_indices = self._swa_kv_pool_host.alloc(num_tokens)
if host_indices is None:
return PreparePrefetchResult(alloc_failed=True)
return PreparePrefetchResult(host_indices=host_indices)
self.cache.evict_host(num_tokens, ComponentType.SWA)
host_indices = self._swa_kv_pool_host.alloc(num_tokens)
return host_indices
def build_hicache_transfers(
self,
@@ -1039,6 +1034,7 @@ class SWAComponent(TreeComponent):
host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0,
staging_tokens: int = 0,
last_hash: Optional[str] = None,
) -> Optional[list[PoolTransfer]]:
ct = self.component_type
@@ -1116,14 +1112,14 @@ class SWAComponent(TreeComponent):
]
if phase == CacheTransferPhase.PREFETCH:
assert host_indices is not None
# Keys are unknowable at build time; placeholders carry the
# count, _sync_trailing_keys fills the real trailing hashes.
num_pages = host_indices.numel() // self.tree_core.page_size
# Staging is allocated once the hit is known; the placeholders carry
# the planned page count and _sync_trailing_keys fills the real hashes.
num_pages = staging_tokens // self.tree_core.page_size
if num_pages == 0:
return None
return [
PoolTransfer(
name=PoolName.SWA,
host_indices=host_indices,
keys=["__placeholder__"] * num_pages,
hit_policy=PoolHitPolicy.TRAILING_PAGES,
)
python/sglang/srt/mem_cache/unified_cache/storage_attachment.py
+12
View File
@@ -114,6 +114,7 @@ class StorageAttachment:
)
try:
prefetch_threshold = self.resolve_prefetch_threshold(prefetch_threshold)
controller.attach_storage_backend(
storage_backend=storage_backend,
prefetch_threshold=prefetch_threshold,
@@ -253,6 +254,16 @@ class StorageAttachment:
else:
cache.storage_metrics_collector = None
def resolve_prefetch_threshold(self, configured: int) -> int:
"""Use the same complete-window minimum for every buffer-mode anchor."""
cache = self._cache
window = cache.sliding_window_size
if not window or cache.host_memory_mode != "buffer_only":
return configured
page_size = cache.page_size
window_tokens = ((window + page_size - 1) // page_size) * page_size
return max(configured, window_tokens)
def _resolve_metrics_collector(
self,
storage_backend: Optional[str],
@@ -409,3 +420,4 @@ class StorageAttachment:
cache.discard_storage_prefetch_accounting(handle)
cache.prefetch_loaded_tokens_by_reqid.clear()
cache.prefetch_loaded_storage_start_by_reqid.clear()
cache.storage_prefetch_retries.clear()
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+144
View File
@@ -709,6 +709,22 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
# TODO: delta is not aggregated from components; no caller uses it yet.
return DecLockRefResult()
def inc_full_pin(self, node_id: NodeId) -> None:
"""Pin only the FULL device slots on the node's root path; the SWA segment
lock is left alone since its receipt does not survive a window re-attach."""
node = self.node_by_id(node_id)
self.components_by_type[BASE_COMPONENT_TYPE].acquire_component_lock(
node=node, result=IncLockRefResult()
)
self._update_evictable_leaf_sets(node)
def dec_full_pin(self, node_id: NodeId) -> None:
node = self.node_by_id(node_id)
self.components_by_type[BASE_COMPONENT_TYPE].release_component_lock(
node=node, params=None
)
self._update_evictable_leaf_sets(node)
def dec_swa_lock_only(
self,
node_id: NodeId,
@@ -872,6 +888,32 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
action,
)
def match_full_device_prefix(self, key: RadixKey) -> tuple[int, NodeId, int]:
"""Read-only FULL-device match, independent of auxiliary components; the
third result counts the whole deepest node (a partial match pins it all)."""
key, _ = key.maybe_to_bigram_view(self.is_eagle)
key = key.page_aligned(self.page_size)
node = self.root_node
matched_len = 0
pinned_len = 0
while len(key) > 0:
child = node.children.get(key.child_key(self.page_size))
if child is None:
break
value = child.component_data[BASE_COMPONENT_TYPE].value
if value is None:
break
prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len == 0:
break
matched_len += prefix_len
node = child
pinned_len += len(child.key)
if prefix_len < len(child.key):
break
key = key[prefix_len:]
return matched_len, node.id, pinned_len
def _match_post_processor(
self,
params: MatchPrefixParams,
@@ -2191,6 +2233,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0,
staging_tokens: int = 0,
last_hash: Optional[str] = None,
) -> Optional[list[PoolTransfer]]:
"""Route a build_hicache_transfers call to the component for the given type."""
@@ -2200,6 +2243,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
host_indices=host_indices,
token_ids=token_ids,
prefetch_tokens=prefetch_tokens,
staging_tokens=staging_tokens,
last_hash=last_hash,
)
@@ -2475,6 +2519,106 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self._update_duplicate_tracking(node)
self.kv_events.record_store(node, medium=StorageMedium.CPU)
def _walk_span(self, key: RadixKey, end: int):
"""Yield nodes and their matched spans along ``key`` through ``end``."""
node = self.root_node
pos = 0
remaining = key
while pos < end and len(remaining) > 0:
child = node.children.get(remaining.child_key(self.page_size))
if child is None or (child.evicted and not child.backuped):
return
prefix_len = child.key.match(remaining, page_size=self.page_size)
if prefix_len == 0:
return
yield child, pos, prefix_len
if prefix_len < len(child.key):
return
node = child
pos += prefix_len
remaining = remaining[prefix_len:]
def swa_tombstone_ranges(
self, key: RadixKey, start: int, end: int
) -> list[tuple[int, int]]:
"""Return the maximal SWA-tombstoned ranges within ``[start, end)``."""
ranges: list[tuple[int, int]] = []
for child, pos, prefix_len in self._walk_span(key, end):
seg_end = pos + prefix_len
if seg_end <= start:
continue
if child.component_data[ComponentType.SWA].value is not None:
continue
lo, hi = max(start, pos), min(end, seg_end)
if ranges and ranges[-1][1] == lo:
ranges[-1] = (ranges[-1][0], hi)
else:
ranges.append((lo, hi))
if hi >= end:
break
return ranges
def attach_swa_window(
self,
key: RadixKey,
window_start: int,
window_end: int,
swa_values: torch.Tensor,
) -> list[CacheAction | ComponentAction]:
"""Attach a loaded SWA slice to an existing tombstoned tree span."""
assert len(swa_values) == window_end - window_start, (
f"attach_swa_window size mismatch: got {len(swa_values)} "
f"for [{window_start}, {window_end})"
)
actions: list[CacheAction | ComponentAction] = []
covered = window_start
for child, pos, prefix_len in list(self._walk_span(key, window_end)):
seg_end = pos + prefix_len
if seg_end <= window_start:
continue
seg_start = max(pos, window_start)
assign_end = min(seg_end, window_end)
assert child.component_data[ComponentType.SWA].value is None, (
f"attach_swa_window over live SWA at [{seg_start}, {assign_end}) "
f"of [{window_start}, {window_end})"
)
self._attach_swa_segment(
child,
pos,
seg_start,
assign_end,
swa_values[seg_start - window_start : assign_end - window_start],
actions,
)
covered = assign_end
if covered >= window_end:
break
assert covered == window_end, (
f"attach_swa_window covered {covered} of [{window_start}, {window_end})"
)
return actions
def _attach_swa_segment(
self,
node: UnifiedTreeNode,
node_start: int,
seg_start: int,
seg_end: int,
values: torch.Tensor,
actions: list[CacheAction | ComponentAction],
) -> None:
target = node
if seg_start > node_start:
_, action = self._split_node(target.key, target, seg_start - node_start)
if action is not None:
actions.append(action)
if seg_start + len(target.key) > seg_end:
fragment, action = self._split_node(target.key, target, seg_end - seg_start)
if action is not None:
actions.append(action)
target = fragment
self.set_component_device_value(target.id, ComponentType.SWA, values.clone())
def set_component_device_value(
self, node_id: NodeId, component_type: ComponentType, value: torch.Tensor
) -> None:
python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py
+16
View File
@@ -391,6 +391,21 @@ class UnifiedTreeCoreInterface(ABC):
"""Match a key against the tree; returns device indices + boundary NodeIds."""
...
@abstractmethod
def match_full_device_prefix(self, key: RadixKey) -> tuple[int, NodeId, int]:
"""Return (matched tokens, deepest node, FULL tokens pinned by it)."""
...
@abstractmethod
def inc_full_pin(self, node_id: NodeId) -> None:
"""Pin only FULL device values on the node's root path."""
...
@abstractmethod
def dec_full_pin(self, node_id: NodeId) -> None:
"""Release a pin acquired by inc_full_pin."""
...
def supports_fast_match_prefix(self) -> bool:
"""Whether matching every waiting request is cheap enough for scheduling."""
return False
@@ -498,6 +513,7 @@ class UnifiedTreeCoreInterface(ABC):
host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0,
staging_tokens: int = 0,
last_hash: Optional[str] = None,
) -> Optional[list[PoolTransfer]]:
"""Build a component's HiCache transfers for the given node and phase."""
python/sglang/srt/mem_cache/unified_radix_cache.py
+251 -122
View File
@@ -45,6 +45,7 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
)
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.mem_cache.unified_cache.cache_action import (
BackupKV,
@@ -281,6 +282,7 @@ class UnifiedRadixCache(BasePrefixCache):
self._prefetch_outcome_stats: dict[str, float] = {
"attempts": 0,
"issued": 0,
"issued_assumed_stored": 0,
"declined_too_short": 0,
"declined_rate_limited": 0,
"declined_anchor_lost": 0,
@@ -389,12 +391,8 @@ class UnifiedRadixCache(BasePrefixCache):
# Rank-agreed L3-hit tokens not yet resolved as usable or unfulfilled.
# Cache-mode entries survive L3->L2 until H2D succeeds or admission
# fails; buffer-mode entries survive staging until the H2D ack.
self._storage_prefetch_hit_remaining_by_reqid: dict[
CacheRequestHandle, int
] = {}
# Attempts whose storage prefetch resolved without a usable result;
# popped by the scheduler to pace availability-check retries.
self._storage_prefetch_missed_rids: set[CacheRequestHandle] = set()
self._storage_prefetch_hit_remaining_by_reqid: dict[str, int] = {}
self.storage_prefetch_retries = StoragePrefetchRetries()
self.ongoing_backup: dict[int, tuple[NodeId, DecLockRefParams]] = {}
if self.buffer_pipeline is not None:
self.buffer_pipeline.reset()
@@ -431,6 +429,7 @@ class UnifiedRadixCache(BasePrefixCache):
self.load_cache_event = threading.Event()
self.sidecar_pool_specs.clear()
self.extra_metric_labels = get_observability().extra_metric_labels
self._storage_attachment = StorageAttachment(self)
# Parse storage config once, share with assembler and tree
storage_backend = get_memory().hicache_storage_backend
@@ -449,6 +448,11 @@ class UnifiedRadixCache(BasePrefixCache):
) = HybridCacheController.parse_storage_backend_extra_config(
get_memory().hicache_storage_backend_extra_config
)
storage_prefetch_threshold = (
self._storage_attachment.resolve_prefetch_threshold(
storage_prefetch_threshold
)
)
attach_hybrid_pool_to_unified_cache(
self,
@@ -515,7 +519,6 @@ class UnifiedRadixCache(BasePrefixCache):
self.prefetch_stop_policy = get_memory().hicache_storage_prefetch_policy
# Runtime attach/detach of the L3 backend (startup, admin API, atexit).
self._storage_attachment = StorageAttachment(self)
atexit.register(self.shutdown)
if storage_backend is not None:
@@ -1795,11 +1798,10 @@ class UnifiedRadixCache(BasePrefixCache):
if phase == CacheTransferPhase.BACKUP_HOST
else indices_source.host_indices
)
defer_kv_sidecar = (
phase == CacheTransferPhase.PREFETCH
and spec.indices_from_pool == PoolName.KV
)
if (indices is None or len(indices) == 0) and not defer_kv_sidecar:
# Prefetch staging is allocated at hit time, so a prefetch sidecar
# resolves its indices from its source then, whatever the source.
deferred = phase == CacheTransferPhase.PREFETCH
if (indices is None or len(indices) == 0) and not deferred:
continue
transfers.append(
PoolTransfer(
@@ -1913,11 +1915,33 @@ class UnifiedRadixCache(BasePrefixCache):
matched_prefix_tokens: Optional[list[int]] = None,
extra_key: Optional[str] = None,
cache_salt: Optional[str] = None,
storage_hit_end: Optional[int] = None,
) -> None:
if not self.enable_storage or self.cache_controller is None:
return
req_id = request.rid
self.storage_prefetch_retries.cancel(req_id)
buffer_mode = self.host_memory_mode == "buffer_only"
if request in self.ongoing_prefetch or (
self.buffer_pipeline is not None
and self.buffer_pipeline.has_staged(request)
):
return
# Prefix-covered writes justify extending a known hit to the left,
# never beyond its confirmed right endpoint (including the SWA window).
storage_start = len(matched_prefix_tokens or [])
boundary_tokens = int(self.tree_core.is_eagle)
assume_stored = (
storage_hit_end is not None
and storage_start
< storage_hit_end
<= storage_start + len(new_input_tokens) - boundary_tokens
)
if assume_stored:
new_input_tokens = new_input_tokens[
: storage_hit_end - storage_start + boundary_tokens
]
# Key the span by the request's namespace, not the anchor's (a root
# anchor has none): a span published under the wrong namespace gets
# re-owned by the request's own insert (double free).
@@ -1936,6 +1960,9 @@ class UnifiedRadixCache(BasePrefixCache):
is_bigram=self.tree_core.is_eagle,
cache_salt=cache_salt,
).page_aligned(self.page_size)
assume_stored = (
assume_stored and storage_start + len(prefetch_key) == storage_hit_end
)
prefetch_length = len(prefetch_key)
stats = self._prefetch_outcome_stats
if prefetch_length > 0:
@@ -1944,18 +1971,13 @@ class UnifiedRadixCache(BasePrefixCache):
if prefetch_length > 0:
stats["declined_too_short"] += 1
# A too-short/fully-matched suffix can become a full recompute if
# the device match evicts while queued; arm the paced retry.
self._storage_prefetch_missed_rids.add(request)
# the device match evicts while queued; arm the retry.
self.storage_prefetch_retries.poll_miss(req_id)
return
if not buffer_mode and self.cache_controller.prefetch_rate_limited():
stats["declined_rate_limited"] += 1
self._storage_prefetch_missed_rids.add(request)
return
if request in self.ongoing_prefetch or (
buffer_mode and self.buffer_pipeline.has_staged(request)
):
# A fetch (or an unconsumed hold) already exists for this attempt;
# overwriting would leak its staging slots.
# Paced: the limiter clears as transfers finish, not on the next pass.
self.storage_prefetch_retries.poll_miss(req_id, storage_hit_end)
return
# Buffer mode holds no tree state during the fetch: buffers are
@@ -1966,18 +1988,15 @@ class UnifiedRadixCache(BasePrefixCache):
else self.inc_host_lock_ref(last_host_node_id).to_dec_params()
)
comp_xfers: dict[ComponentType, list[PoolTransfer]] = {}
alloc_failed = False
for ct in self.tree_components:
if ct == BASE_COMPONENT_TYPE:
continue
# Pre-allocate the component's prefetch host buffer so the build stays pure.
# Size the component's staging now; it is allocated at hit time,
# next to the KV staging, so the query holds no host memory.
prep = self.components[ct].prepare_prefetch(
last_host_node_id, prefetch_tokens=len(prefetch_key)
)
if prep.alloc_failed:
alloc_failed = True
break
if prep.host_indices is None:
if prep.staging_tokens == 0:
continue
transfers = self.tree_core.build_hicache_transfers(
ct,
@@ -1985,8 +2004,8 @@ class UnifiedRadixCache(BasePrefixCache):
CacheTransferPhase.PREFETCH,
token_ids=prefetch_key.token_ids,
prefetch_tokens=len(prefetch_key),
staging_tokens=prep.staging_tokens,
last_hash=last_hash,
host_indices=prep.host_indices,
)
if transfers:
comp_xfers[ct] = transfers
@@ -1994,25 +2013,6 @@ class UnifiedRadixCache(BasePrefixCache):
sidecar_xfers = self._build_sidecar_transfers(
CacheTransferPhase.PREFETCH, kv_xfer, comp_xfers
)
if alloc_failed:
# The whole storage fetch is forfeited over one aux staging
# alloc (e.g. a single SWA window) — count it, or write-burst
# starvation of the aux pool reads as generic hit-rate loss.
if (
self.enable_storage_metrics
and self.storage_metrics_collector is not None
):
self.storage_metrics_collector.log_prefetch_aux_alloc_failed_tokens(
len(prefetch_key)
)
self.cache_controller.append_host_mem_release(
extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
)
if anchor_lock_params is not None:
self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
# Forfeited over transient staging pressure; retryable.
self._storage_prefetch_missed_rids.add(request)
return
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
aux_xfers.extend(sidecar_xfers)
@@ -2022,12 +2022,15 @@ class UnifiedRadixCache(BasePrefixCache):
last_hash,
prefix_keys,
extra_pools=aux_xfers or None,
assume_stored=assume_stored,
)
stats["issued"] += 1
if assume_stored:
stats["issued_assumed_stored"] += 1
# Snapshot the requested span for L3 miss-token accounting at the
# rank-synchronized query outcome.
operation.stats_requested_tokens = prefetch_length
operation.storage_start = len(matched_prefix_tokens or [])
operation.storage_start = storage_start
self.ongoing_prefetch[request] = _OngoingPrefetch(
last_host_node_id,
prefetch_key,
@@ -2037,16 +2040,14 @@ class UnifiedRadixCache(BasePrefixCache):
comp_xfers,
)
if buffer_mode:
# The query reads no tree state; pinning across it would hold the
# anchor cap for a storage round trip, so the pin waits for IO commit.
self.buffer_pipeline.set_prefix_ctx(
request,
matched_prefix_tokens,
extra_key=extra_key,
cache_salt=cache_salt,
)
# Pin the just-matched anchor now: deferred to IO commit it is
# often already deleted under churn. The IO-commit call remains
# as the second chance that decides the fetch's fate.
self.buffer_pipeline.try_lock_anchor(request)
else:
# Cache mode reserves the requested span up front; buffer mode
# grants occupancy later at hit-alloc time, sized to the hit.
@@ -2064,14 +2065,12 @@ class UnifiedRadixCache(BasePrefixCache):
if self.prefetch_stop_policy == "best_effort":
return True
if self.prefetch_stop_policy == "wait_complete":
# Completion is committed by the ACK drain (sync-thread MIN-reduced
# progress, rank-min queue drain); a vote here only adds collectives.
return False
elif self.prefetch_stop_policy == "timeout":
# Wall-clock time may differ among ranks, all-reduce is needed to ensure
# all ranks reach the same final result. Otherwise PP/TP ranks will diverge.
#
# For TP, if any rank reaches the timeout, the final result is timeout.
#
# For PP, PP0 makes the decision and other ranks follow PP0's decision.
if self.prefetch_stop_policy == "timeout":
# Wall clocks differ across ranks: PP0 decides, TP takes MAX (any
# rank timed out) and _all_reduce broadcasts the verdict along PP.
should_terminate = False
if self.pp_rank == 0:
should_terminate = self._prefetch_timeout_check_linear_func(operation)
@@ -2080,8 +2079,7 @@ class UnifiedRadixCache(BasePrefixCache):
)
self._all_reduce(should_terminate_tensor, torch.distributed.ReduceOp.MAX)
return should_terminate_tensor.item() == 1
else:
return True
return True
def has_ongoing_prefetch(self, handle: CacheRequestHandle) -> bool:
return handle in self.ongoing_prefetch
@@ -2091,7 +2089,7 @@ class UnifiedRadixCache(BasePrefixCache):
if request not in self.ongoing_prefetch:
return True
_, _, _, operation, _, _ = self.ongoing_prefetch[request]
operation = self.ongoing_prefetch[request].operation
# Determine whether or not we should terminate this prefetch request.
should_terminate = operation.is_terminated() or self._can_terminate_prefetch(
@@ -2103,7 +2101,7 @@ class UnifiedRadixCache(BasePrefixCache):
self.cache_controller.terminate_prefetch(operation)
if operation.host_indices is None:
self._storage_prefetch_missed_rids.add(request)
self.storage_prefetch_retries.poll_miss(request.rid)
self.revoke_pending_prefetch(request)
else:
self._handle_prefetch_result(operation)
@@ -2283,17 +2281,24 @@ class UnifiedRadixCache(BasePrefixCache):
# Drop the KV beliefs from the first page any pool failed to serve;
# the next insert then re-writes that span through one FULL check,
# restoring the missing aux pages.
keep_pages = completed_tokens // self.page_size
invalidation_hashes = (
operation.sidecar_hash_values
if operation.sidecar_hash_values is not None
else hash_value
)
trimmed_pages = len(invalidation_hashes) - len(hash_value)
keep_pages = trimmed_pages + completed_tokens // self.page_size
for transfer, count in zip(pool_transfers, pool_hit_pages):
if transfer.keys is None:
keep_pages = 0
elif count < len(transfer.keys):
# Aux transfers key the chain's trailing pages.
keep_pages = min(
keep_pages, max(0, len(hash_value) - len(transfer.keys))
keep_pages,
max(0, len(invalidation_hashes) - len(transfer.keys)),
)
self.storage_existence_cache.invalidate_beyond(
PoolName.KV, hash_value, keep_pages=keep_pages
PoolName.KV, invalidation_hashes, keep_pages=keep_pages
)
# The controller's prefetch IO thread already releases the untransferred
# tail (host_indices[completed_tokens:])
@@ -2411,14 +2416,19 @@ class UnifiedRadixCache(BasePrefixCache):
self._storage_prefetch_hit_remaining_by_reqid.pop(request, None)
def _handle_storage_prefetch_anchor_loss(self, request: CacheRequestHandle) -> None:
operation = self.ongoing_prefetch[request].operation
storage_hit_end = operation.storage_start + operation.storage_hit_count
self._finish_storage_prefetch(request, fulfilled_tokens=0, reason="shrunk")
# The span is still L3-resident; retry from the shorter live match.
self._storage_prefetch_missed_rids.add(request)
self.revoke_pending_prefetch(request)
self.storage_prefetch_retries.refetch(request.rid, storage_hit_end)
def _log_storage_prefetch_deferred(self, num_tokens: int, reason: str) -> None:
if self.enable_storage_metrics and self.storage_metrics_collector is not None:
self.storage_metrics_collector.log_storage_prefetch_deferred_tokens(
num_tokens, reason
)
def pop_prefetch_loaded_tokens(self, request: CacheRequestHandle) -> int:
# The request is being scheduled; a still-unserved miss marker is moot.
self._storage_prefetch_missed_rids.discard(request)
self.prefetch_loaded_storage_start_by_reqid.pop(request, None)
return self.prefetch_loaded_tokens_by_reqid.pop(request, 0)
@@ -2426,43 +2436,19 @@ class UnifiedRadixCache(BasePrefixCache):
self, request: CacheRequestHandle
) -> tuple[int, Optional[int]]:
"""Pop the loaded L3 token count and its absolute prefix start."""
self._storage_prefetch_missed_rids.discard(request)
return (
self.prefetch_loaded_tokens_by_reqid.pop(request, 0),
self.prefetch_loaded_storage_start_by_reqid.pop(request, None),
)
def pop_storage_prefetch_miss(self, request: CacheRequestHandle) -> bool:
"""True once per resolved storage-prefetch miss for a live request;
the scheduler uses it to arm the paced availability-check retry."""
if request in self._storage_prefetch_missed_rids:
self._storage_prefetch_missed_rids.discard(request)
return True
return False
def plan_staged_splice(
self, request: CacheRequestHandle, device_prefix_len: int
) -> tuple[int, int]:
"""(kv, swa) host-hit tokens a staged buffer-mode prefetch will splice
given the request's live device prefix; frees unusable holds."""
if self.buffer_pipeline is None:
return 0, 0
return self.buffer_pipeline.plan_staged_splice(request, device_prefix_len)
def staged_prefetch_swa_tokens(self, request: CacheRequestHandle) -> int:
"""SWA device tokens consuming a staged buffer-mode prefetch will
allocate; surfaced as the request's swa_host_hit_length."""
if self.buffer_pipeline is None:
return 0
return self.buffer_pipeline.staged_prefetch_swa_tokens(request)
@rank_consensus(same_params=True)
def release_aborted_request(self, request: CacheRequestHandle) -> None:
rid = request.rid
if self.linker is not None:
self.linker.release_request(request.rid)
self.prefetch_loaded_tokens_by_reqid.pop(request, None)
self.prefetch_loaded_storage_start_by_reqid.pop(request, None)
self._storage_prefetch_missed_rids.discard(request)
self.linker.release_request(rid)
self.prefetch_loaded_tokens_by_reqid.pop(rid, None)
self.prefetch_loaded_storage_start_by_reqid.pop(rid, None)
self.storage_prefetch_retries.cancel(rid)
if (
self.buffer_pipeline is not None
and self.buffer_pipeline.release_staged_hold(request)
@@ -2545,6 +2531,90 @@ class UnifiedRadixCache(BasePrefixCache):
return len(host_indices) if host_indices is not None else 0
return len(prefetch_key)
def _alloc_prefetch_aux_staging(
self, info: _OngoingPrefetch, hit_tokens: int
) -> bool:
"""Bind the aux staging (SWA window, mamba state) next to the KV staging;
all pools or none, so a fetch never launches half-staged."""
cc = self.cache_controller
hit_pages = hit_tokens // self.page_size
taken: list[PoolTransfer] = []
for ct, transfers in info.comp_xfers.items():
for transfer in transfers:
if transfer.host_indices is not None or transfer.indices_from_pool:
continue
entry = cc.mem_pool_host.entry_map.get(transfer.name)
pool_page_size = (
entry.host_pool.page_size if entry is not None else self.page_size
)
num_tokens = min(len(transfer.keys or ()), hit_pages) * pool_page_size
if num_tokens == 0:
continue
host_indices = self.components[ct].alloc_prefetch_staging(num_tokens)
if host_indices is None:
if (
self.enable_storage_metrics
and self.storage_metrics_collector is not None
):
self.storage_metrics_collector.log_prefetch_aux_alloc_failed_tokens(
num_tokens
)
cc.append_host_mem_release(extra_pools=taken)
for t in taken:
t.host_indices = None
return False
transfer.host_indices = host_indices
taken.append(transfer)
return True
def _trim_buffer_prefetch_full_head(
self,
request: CacheRequestHandle,
info: _OngoingPrefetch,
operation,
full_match_len: int,
hit_tokens: int,
) -> tuple[_OngoingPrefetch, int, int]:
"""Drop FULL pages the hit-time rematch found on device; trailing aux
transfers keep the query's endpoint so SWA still fetches its window."""
# A capacity-deferred hit comes back already trimmed; size its sidecars
# from the pre-trim boundary or the SWA staging comes up short.
original_hit_tokens = (
operation.sidecar_hit_pages * self.page_size
if operation.sidecar_hash_values is not None
else hit_tokens
)
trim_tokens = min(hit_tokens, max(0, full_match_len - operation.storage_start))
if trim_tokens == 0:
return info, hit_tokens, original_hit_tokens
assert trim_tokens % self.page_size == 0
old_key = info.prefetch_key
prefix_ctx = self.buffer_pipeline._prefetch_prefix_ctx[request]
prefix_tokens, extra_key, cache_salt = prefix_ctx
# Prefix context excludes the boundary token; the suffix owns it even
# when all FULL pages are trimmed and only sidecars remain to fetch.
new_prefix_tokens = prefix_tokens + list(old_key.token_ids[:trim_tokens])
self.buffer_pipeline._prefetch_prefix_ctx[request] = (
new_prefix_tokens,
extra_key,
cache_salt,
)
info = info._replace(
prefetch_key=RadixKey(
old_key.raw_token_ids()[trim_tokens:],
extra_key=old_key.extra_key,
is_bigram=old_key.is_bigram,
cache_salt=old_key.cache_salt,
)
)
self.ongoing_prefetch[request] = info
self.cache_controller.trim_prefetch_full_head(operation, trim_tokens)
self._resolve_storage_prefetch_tokens(
request, trim_tokens, reason="device_covered"
)
return info, hit_tokens - trim_tokens, original_hit_tokens
def revoke_pending_prefetch(self, request: CacheRequestHandle) -> None:
info = self.ongoing_prefetch.pop(request, None)
self._finish_storage_prefetch(request, fulfilled_tokens=0, reason="dropped")
@@ -2630,13 +2700,19 @@ class UnifiedRadixCache(BasePrefixCache):
# ongoing_prefetch, so wait_complete keeps gating admission.
return False
if buffer_mode:
# IO commit: pin before the bounce alloc so a cancel is a
# plain revoke and a parked op keeps its pin; a fetch whose
# splice base is gone is not worth its storage read.
if self.buffer_pipeline.try_lock_anchor(request) == "anchor_lost":
# IO commit pins the splice base; a parked op held no pin, so a
# moved or vanished anchor is detected (and recovered) here.
anchor, full_match_len = self.buffer_pipeline.try_lock_anchor(
request, operation.storage_hit_count
)
if anchor == "anchor_lost":
self._prefetch_outcome_stats["declined_anchor_lost"] += 1
self._handle_storage_prefetch_anchor_loss(request)
return True
if anchor == "cap_busy":
# Other prefetches hold the budget; ours fits once they
# drain, so park rather than launch this one unprotected.
return False
if self.buffer_pipeline.staged_span_covered(
request, operation.storage_hit_count
):
@@ -2644,10 +2720,15 @@ class UnifiedRadixCache(BasePrefixCache):
# splice, so skip the storage read.
self._prefetch_outcome_stats["declined_device_covered"] += 1
self._finish_storage_prefetch(
request, fulfilled_tokens=0, reason=None
request, fulfilled_tokens=0, reason="device_covered"
)
self.revoke_pending_prefetch(request)
return True
info, hit_tokens, aux_hit_tokens = self._trim_buffer_prefetch_full_head(
request, info, operation, full_match_len, hit_tokens
)
else:
aux_hit_tokens = hit_tokens
alloc_len = hit_tokens
host_indices = cc.mem_pool_host.alloc(alloc_len)
if host_indices is None:
@@ -2665,11 +2746,43 @@ class UnifiedRadixCache(BasePrefixCache):
host_indices = cc.mem_pool_host.alloc(alloc_len)
if host_indices is None:
if buffer_mode:
# Parked ops hold no pin: release and re-take at the next
# attempt, which is also how a moved anchor gets noticed.
self.buffer_pipeline.release_anchor_lock(request)
self._log_storage_prefetch_deferred(
max(alloc_len, aux_hit_tokens), "host_capacity"
)
return False
self._finish_storage_prefetch(
request, fulfilled_tokens=0, reason="host_capacity"
)
self.revoke_pending_prefetch(request)
self.storage_prefetch_retries.poll_miss(
request.rid, operation.storage_start + operation.storage_hit_count
)
self._log_storage_prefetch_deferred(
max(alloc_len, aux_hit_tokens), "host_capacity"
)
return True
if not self._alloc_prefetch_aux_staging(info, aux_hit_tokens):
# Same outcome as a KV shortfall: nothing stays staged.
cc.append_host_mem_release(host_indices=host_indices)
if buffer_mode:
self.buffer_pipeline.release_anchor_lock(request)
self._log_storage_prefetch_deferred(
max(alloc_len, aux_hit_tokens), "host_capacity"
)
return False
self._finish_storage_prefetch(
request, fulfilled_tokens=0, reason="host_capacity"
)
self.revoke_pending_prefetch(request)
self.storage_prefetch_retries.poll_miss(
request.rid, operation.storage_start + operation.storage_hit_count
)
self._log_storage_prefetch_deferred(
max(alloc_len, aux_hit_tokens), "host_capacity"
)
return True
self._resolve_storage_prefetch_tokens(
@@ -2718,7 +2831,10 @@ class UnifiedRadixCache(BasePrefixCache):
else None
),
)
self._storage_prefetch_missed_rids.add(request)
# The published (rank-agreed) hit count is the L3 verdict;
# without this, store misses never reach l3_miss_tokens.
self._account_prefetch_outcome(operation, revoked=True)
self.storage_prefetch_retries.poll_miss(request.rid)
self.revoke_pending_prefetch(request)
continue
if hit_tokens < self.prefetch_threshold:
@@ -2727,12 +2843,17 @@ class UnifiedRadixCache(BasePrefixCache):
self._finish_storage_prefetch(
request, fulfilled_tokens=0, reason="below_threshold"
)
self._storage_prefetch_missed_rids.add(request)
self.storage_prefetch_retries.poll_miss(request.rid)
self.revoke_pending_prefetch(request)
continue
self._invalidate_absent_from_hit_query(operation)
self._account_prefetch_outcome(operation, revoked=False)
if not _try_alloc_storage_hit(operation):
# A parked hit keeps its turn: a newer hit must not take the
# staging or anchor budget the parked head is waiting for.
parked_ahead = buffer_mode and bool(
self.buffer_pipeline.pending_hit_allocs
)
if parked_ahead or not _try_alloc_storage_hit(operation):
# Counted once at first parking, not per retry tick.
self._prefetch_outcome_stats["declined_rate_limited"] += 1
self.buffer_pipeline.pending_hit_allocs.append(operation)
@@ -2740,29 +2861,32 @@ class UnifiedRadixCache(BasePrefixCache):
def _drain_ack_prefetch():
for ack in _drain_queue(cc.ack_prefetch_queue, n_ack_prefetch):
operation = ack.operation
info = self.ongoing_prefetch.get(operation.handle)
is_current = info is not None and info.operation is operation
if ack.completed_tokens is not None:
if operation.handle in self.ongoing_prefetch:
if is_current:
assert operation.completed_tokens <= ack.completed_tokens
operation.completed_tokens = ack.completed_tokens
if ack.pool_hits is not None:
if operation.handle in self.ongoing_prefetch:
if is_current:
operation.pool_storage_result.update_extra_pool_hit_pages(
ack.pool_hits
)
operation.pool_transfers_done = True
if ack.completed_req:
if operation.handle in self.ongoing_prefetch:
# check_prefetch_progress() is not called for this attempt yet.
if is_current:
# check_prefetch_progress() is not called for this rid yet.
# Let us insert the prefetch result into the radix tree.
self._handle_prefetch_result(operation)
cc.append_host_mem_release(
operation.host_indices[operation.completed_tokens :],
(
operation.pool_transfers
if not operation.pool_transfers_done
else None
),
)
if operation.ack_releases_incomplete_host_indices:
cc.append_host_mem_release(
operation.host_indices[operation.completed_tokens :],
(
operation.pool_transfers
if not operation.pool_transfers_done
else None
),
)
def _drain_backup():
drained = 0
@@ -3109,10 +3233,10 @@ class UnifiedRadixCache(BasePrefixCache):
def init_load_back(
self,
params: InitLoadBackParams,
) -> tuple[torch.Tensor, NodeId]:
) -> Optional[tuple[torch.Tensor, NodeId]]:
"""Prepare KV cache loading from host to device.
Returns (device_indices, last_node). Buffer mode dispatches to the
staged-prefetch consumption (BufferModePipeline.init_load_back)."""
Returns (device_indices, last_node), or None when buffer-mode
admission must retry without committing a load."""
if self.buffer_pipeline is not None:
return self.buffer_pipeline.init_load_back(params)
best_match_node_id = params.best_match_node
@@ -3343,6 +3467,11 @@ class UnifiedRadixCache(BasePrefixCache):
def swa_protected_size(self) -> int:
return self.tree_core.swa_protected_size()
def swa_transient_size(self) -> int:
if self.buffer_pipeline is None:
return 0
return self.buffer_pipeline.swa_transient_size()
def mamba_protected_size(self) -> int:
return self.tree_core.mamba_protected_size()
python/sglang/srt/observability/metrics_collector.py
+20
View File
@@ -1944,6 +1944,7 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
"below_threshold",
"host_capacity",
"device_capacity",
"device_covered",
"storage_transfer",
"shrunk",
"dropped",
@@ -1952,6 +1953,17 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
**self.labels, reason=reason
)
self.storage_prefetch_deferred_tokens_total = Counter(
name="sglang:storage_prefetch_deferred_tokens_total",
documentation="Storage-prefetch token-attempts deferred for later "
"reuse, by transient capacity reason.",
labelnames=list(labels.keys()) + ["reason"],
)
for reason in ("host_capacity", "device_capacity"):
self.storage_prefetch_deferred_tokens_total.labels(
**self.labels, reason=reason
)
self.backup_dropped_tokens_total = Counter(
name="sglang:hicache_backup_dropped_tokens_total",
documentation="Buffer-mode backup tokens that never reached L3 "
@@ -2041,6 +2053,14 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
**self.labels, reason=reason
).inc(num_tokens)
def log_storage_prefetch_deferred_tokens(
self, num_tokens: int, reason: str
) -> None:
if num_tokens > 0:
self.storage_prefetch_deferred_tokens_total.labels(
**self.labels, reason=reason
).inc(num_tokens)
def log_backup_dropped_tokens(self, dropped_tokens: int):
if dropped_tokens > 0:
self.backup_dropped_tokens_total.labels(**self.labels).inc(dropped_tokens)
python/sglang/srt/session/streaming_session.py
+10
View File
@@ -24,6 +24,8 @@ from sglang.srt.utils.common import ceil_align, is_npu
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline
from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries
logger = logging.getLogger(__name__)
@@ -572,6 +574,14 @@ class StreamingSession(BasePrefixCache):
def init_load_back(self, params: InitLoadBackParams):
return self.inner.init_load_back(params)
@property
def buffer_pipeline(self) -> Optional[BufferModePipeline]:
return self.inner.buffer_pipeline
@property
def storage_prefetch_retries(self) -> Optional[StoragePrefetchRetries]:
return self.inner.storage_prefetch_retries
def pop_prefetch_loaded_span(
self, handle: CacheRequestHandle
) -> tuple[int, Optional[int]]: