From cb39f232279545ea48d2d4600d9c064167c06e3c Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Sun, 13 Sep 2026 00:46:19 -0700 Subject: [PATCH 1/2] [HiCache] Rework the buffer-mode storage prefetch pipeline and retry bookkeeping Buffer-mode (host memory as transport staging, no L2 tier) storage prefetch used to decide a staged hit's fate inside the prefill adder, after admission had already charged host-hit tokens that a shrunk device prefix could no longer splice. This reworks the pipeline so the L3 hit is planned before selection and materialized only once, and makes the retry path bounded. Scheduler / prefill adder - `PrefillAdder.add_one_req` now selects a prefill shape without allocating (`_select_prefill_admission`), materializes host hits under the prefix lock via `init_load_back`, and commits the admission afterwards (`_commit_prefill_admission`). `init_load_back` may return None to ask for another admission attempt; over-delivered FULL tokens shrink the planned extend range instead of re-selecting, and under-delivery re-selects with no host hit. - The SWA budget gate lives in `_swa_admission_gate`, keeping the ring-slot exact-fit comparison for the unified-KV SWA ring. - Storage prefetch retries move into `StoragePrefetchRetries` (`mem_cache/storage_prefetch.py`): paced miss polls, immediate re-issues for known hits whose anchor moved, and a per-request attempt cap (`--hicache-storage-prefetch-retry-max-attempts`, default 8; the poll interval default becomes 8 passes). Requests past the cap are admitted with whatever the device holds. - HiCache event draining is factored into `_process_hicache_events` and called from the PD-prefill and PP scheduler loops, which had stopped draining prefetch/backup events. Buffer-mode pipeline - Staged holds are anchored by a FULL-only pin (`match_full_device_prefix`, `inc_full_pin`, `dec_full_pin` on the tree core, Python and Rust) so the device prefix under a parked fetch cannot be evicted while the SWA segment lock stays untouched. - Aux components size their prefetch staging in `prepare_prefetch` (`PreparePrefetchResult.staging_tokens`) and allocate it once the hit is known (`alloc_prefetch_staging`), instead of pre-allocating per query. The Rust components follow the same contract: PREFETCH builds take `staging_tokens` and carry no host buffer. - Admission-time device-capacity deferrals are bounded (`max_staged_admission_defers = 32`); past the cap the hold is dropped and the request recomputes. - Parked hits keep their turn ahead of newer hits, stale prefetch acks are ignored once a request re-issued its query, misses are accounted from the published hit count, and EAGLE bigram keys keep their boundary token when the prefetch span is trimmed or staged. - `swa_transient_size` reports SWA slots owned by staged holds so the scheduler invariant checker can account for them. Tests cover the new adder flow, the retry bookkeeping, the staged-hold lifecycle, the scheduler event draining, and the Rust tree-core pin. --- python/sglang/srt/arg_groups/fields/memory.py | 23 +- python/sglang/srt/disaggregation/prefill.py | 1 + .../npu/dsv4/c128_sidecar_component.py | 1 + python/sglang/srt/managers/schedule_batch.py | 9 +- python/sglang/srt/managers/schedule_policy.py | 398 +++++---- python/sglang/srt/managers/scheduler.py | 146 ++-- .../scheduler_components/invariant_checker.py | 2 +- .../sglang/srt/managers/scheduler_pp_mixin.py | 1 + .../sglang/srt/mem_cache/base_prefix_cache.py | 13 +- .../srt/mem_cache/buffer_mode/pipeline.py | 550 +++++++------ python/sglang/srt/mem_cache/hiradix_cache.py | 6 +- .../hybrid_cache/hybrid_cache_controller.py | 47 +- .../srt/mem_cache/rust_tree_core/adapter.py | 37 + .../sglang/srt/mem_cache/storage_prefetch.py | 85 ++ .../unified_cache/components/base.py | 16 +- .../unified_cache/components/full.py | 1 + .../unified_cache/components/mamba.py | 22 +- .../mem_cache/unified_cache/components/swa.py | 42 +- .../unified_cache/storage_attachment.py | 12 + .../unified_cache/unified_tree_core.py | 144 ++++ .../unified_tree_core_interface.py | 16 + .../srt/mem_cache/unified_radix_cache.py | 373 ++++++--- .../srt/observability/metrics_collector.py | 20 + .../sglang/srt/session/streaming_session.py | 10 + rust/sglang-radix-tree/src/components/full.rs | 1 + .../sglang-radix-tree/src/components/mamba.rs | 9 +- rust/sglang-radix-tree/src/components/mod.rs | 2 + rust/sglang-radix-tree/src/components/swa.rs | 15 +- rust/sglang-radix-tree/src/python_bindings.rs | 50 +- .../src/tests/components/full.rs | 5 +- .../src/tests/components/mamba.rs | 25 +- .../src/tests/components/swa.rs | 55 +- .../src/tests/unified_tree_core.rs | 20 + .../src/unified_tree_core.rs | 72 ++ .../unit/managers/test_prefill_adder.py | 185 +++++ .../test_scheduler_chunked_req_gate.py | 5 + .../managers/test_scheduler_hicache_events.py | 139 ++++ .../mem_cache/test_buffer_mode_sidecar.py | 15 +- ...test_hicache_staged_write_back_dispatch.py | 50 +- .../test_storage_prefetch_lifecycle.py | 498 +++++++++++ .../test_unified_radix_cache_unittest.py | 771 +++++++++++++++--- .../observability/test_stat_loggers_di.py | 5 + 42 files changed, 3035 insertions(+), 862 deletions(-) create mode 100644 python/sglang/srt/mem_cache/storage_prefetch.py create mode 100644 test/registered/unit/managers/test_scheduler_hicache_events.py create mode 100644 test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py diff --git a/python/sglang/srt/arg_groups/fields/memory.py b/python/sglang/srt/arg_groups/fields/memory.py index 76e414d2fe56..d00108b31eaf 100644 --- a/python/sglang/srt/arg_groups/fields/memory.py +++ b/python/sglang/srt/arg_groups/fields/memory.py @@ -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 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 512e46bdbd63..56604a281a55 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -592,6 +592,7 @@ def get_next_disagg_prefill_batch_to_run( 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 diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py b/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py index 9fdb0684d7fc..4d1dba298bfb 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py @@ -402,6 +402,7 @@ def build_hicache_transfers( 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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 9f251fd30a8f..7811e4477a8f 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -149,6 +149,7 @@ 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 @@ def __init__( 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 diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index b2bc89939076..e9e89f0be1bd 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -33,6 +33,7 @@ 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 @@ def _swa_req_never_fits( >= 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 @@ def add_one_req( 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), + # 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): + 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, ) - # 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 - - 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 - - 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) + if isinstance(admission, AddReqResult): + return admission - # 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 @@ def add_one_req( 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, ) ) + 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]) - 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 - ) + req.kv.cache_protected_len = len(req.prefix_indices) - 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 + # Successful materialization has no remaining admission gates. + self._commit_prefill_admission(req, admission, mamba_gap_reserve) - if self.dllm_config is not None: - if self.rem_dllm_tokens <= 0: - return AddReqResult.OTHER + # This verdict controls the next candidate, not the committed request. + return self.budget_state() - assert truncation_align_size is None, ( - "truncation_align_size is not supported for dllm prefill" - ) + 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 - if ( - tile_stop := self._check_prefill_tile_budget(input_tokens) - ) is not None: - return tile_stop + 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._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 + # 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 - # 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) - - 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: + 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 - ) - - 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 trunc_len <= 0: - return AddReqResult.OTHER - - 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 - ) + 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 - self.can_run_list.append(req) - self.new_chunked_req = req + if (verdict := self._check_prefill_tile_budget(tile_tokens)) is not None: + return verdict - 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 _PrefillAdmission(prefix_len, extend_len, max_new_tokens, is_chunked) - 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: """ diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2e3d1bcb268f..7f902853f00b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3074,7 +3074,7 @@ def handle_batch_generate_request( 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 @@ def _prefetch_kvcache(self, req: Req): ): 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 @@ def _prefetch_kvcache(self, req: Req): 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 @@ def _prefetch_kvcache(self, req: Req): 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 + 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, + ) + else: logger.debug( - "HiCache storage prefetch retry req=%s attempt=%d", + "HiCache storage prefetch re-issue req=%s attempt=%d", req.rid, req.storage_prefetch_retry_attempts, ) - self._prefetch_kvcache(req) + 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 @@ def _build_hisparse_decode_batch(self, reqs): # 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 @@ def _get_new_batch_prefill_raw( 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 @@ def _get_new_batch_prefill_raw( 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 @@ def _get_new_batch_prefill_raw( 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 @@ def _get_new_batch_prefill_raw( 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: diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index a7340ca20346..91d1df0e65b2 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -280,7 +280,7 @@ def _get_total_uncached_sizes( 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() diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 0d28ee302cef..48a34207cb4e 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -248,6 +248,7 @@ def event_loop_pp_disagg_prefill(self: Scheduler): 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 diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index cb1ed70325f1..0bafbf987fba 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -32,7 +32,9 @@ 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 @@ def full_protected_size(self): 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 @@ def pretty_print(self): 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() diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index 7ef87d2bab1e..6c42aa83cff4 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py @@ -50,6 +50,7 @@ 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 @@ ) 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 @@ def __init__( 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 @@ def is_idle(self) -> bool: 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 @@ def _free_staging_now( # ---- 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 @@ def try_lock_anchor(self, request: CacheRequestHandle) -> str: ): 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 @@ def release_anchor_lock(self, request: CacheRequestHandle) -> None: lock = self.anchor_locks.pop(request, None) if lock is None: return - self._cache.dec_lock_ref(lock.node_id, lock.lock_params) + self._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 @@ def pop_prefix_ctx(self, request: CacheRequestHandle) -> None: 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 @@ def stage_completed_prefetch( 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 @@ def stage_completed_prefetch( 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 @@ def stage_completed_prefetch( 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 @@ def stage_completed_prefetch( 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 _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 + ) -> Optional[tuple[torch.Tensor, NodeId]]: + """Materialize a selected prefill under the caller's prefix lock. - 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 @@ def init_load_back(self, params: InitLoadBackParams) -> tuple[torch.Tensor, Node 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 - - 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) + 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 _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 @@ def _drop(reason: Optional[str]) -> tuple[torch.Tensor, NodeId]: 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 @@ def try_finish_load_back(self, ack_id: int) -> bool: # 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 @@ def release_staged_hold( 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 diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 386f4de5cd0d..9c2e95b56806 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1738,11 +1738,6 @@ def pop_prefetch_loaded_tokens(self, handle: CacheRequestHandle) -> int: """ 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 @@ def prefetch_from_storage( 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( diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index b653cd3a735f..17b37747dee8 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -68,6 +68,7 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def prefetch( 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 @@ def prefetch( 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 @@ def _storage_hit_query(self, operation) -> tuple[list[str], int]: ) 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 @@ def _page_transfer_sidecar( 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 @@ def _page_transfer_sidecar( ) 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 @@ def _resolve_sidecar_kv_derived_pool_transfers(self, operation): 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: diff --git a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py index 9383e491608f..954b56791b7e 100644 --- a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py +++ b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py @@ -405,6 +405,26 @@ def node_by_id(self, node_id: NodeId) -> UnifiedTreeNode: 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 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: ) 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 @@ def build_hicache_transfers( 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 @@ def build_hicache_transfers( # TODO: Forward token ids when Rust Mamba prefetch consumes them. None, prefetch_tokens, + staging_tokens, last_hash, ) if transfers is None: diff --git a/python/sglang/srt/mem_cache/storage_prefetch.py b/python/sglang/srt/mem_cache/storage_prefetch.py new file mode 100644 index 000000000000..c1de15a05828 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage_prefetch.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache/components/base.py b/python/sglang/srt/mem_cache/unified_cache/components/base.py index 642d02f124fe..95ba7922504b 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/base.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/base.py @@ -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 @@ def prepare_prefetch( *, 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 @@ def build_hicache_transfers( 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. diff --git a/python/sglang/srt/mem_cache/unified_cache/components/full.py b/python/sglang/srt/mem_cache/unified_cache/components/full.py index 640bceccbd9c..1d4b255fad31 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/full.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/full.py @@ -350,6 +350,7 @@ def build_hicache_transfers( 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 diff --git a/python/sglang/srt/mem_cache/unified_cache/components/mamba.py b/python/sglang/srt/mem_cache/unified_cache/components/mamba.py index 70d7683de098..3c18cee301fc 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/mamba.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/mamba.py @@ -693,14 +693,15 @@ def prepare_prefetch( *, 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 @@ def build_hicache_transfers( 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 @@ def build_hicache_transfers( ] 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, ) diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa.py b/python/sglang/srt/mem_cache/unified_cache/components/swa.py index 31f92dd19e92..0780e0179d49 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa.py @@ -1007,28 +1007,23 @@ def prepare_prefetch( 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 @@ def build_hicache_transfers( 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 @@ def build_hicache_transfers( ] 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, ) diff --git a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py index 60c5a52b0155..82ee8bd21e01 100644 --- a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py +++ b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py @@ -114,6 +114,7 @@ def attach( ) 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 @@ def apply_runtime_config( 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 @@ def _release_pending_storage_ops(self) -> None: 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() diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index e452e3513e74..2b279540ca94 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -709,6 +709,22 @@ def dec_lock_ref( # 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 @@ def _update_best_if_valid(node): 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 @@ def build_hicache_transfers( 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 @@ def build_hicache_transfers( host_indices=host_indices, token_ids=token_ids, prefetch_tokens=prefetch_tokens, + staging_tokens=staging_tokens, last_hash=last_hash, ) @@ -2475,6 +2519,106 @@ def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None: 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: diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index 4fabcd2dc255..681567984b9e 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -391,6 +391,21 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: """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 @@ def build_hicache_transfers( 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.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index b060958bbb83..d417d7d95b1e 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -45,6 +45,7 @@ ) 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 @@ def __init__( 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 @@ def _reset_full(self) -> None: # 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 @@ def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None 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 @@ def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None ) = 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 @@ def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None 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 @@ def _build_sidecar_transfers( 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 @@ def prefetch_from_storage( 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 @@ def prefetch_from_storage( 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 @@ def prefetch_from_storage( 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 @@ def prefetch_from_storage( 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 @@ def prefetch_from_storage( 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 @@ def prefetch_from_storage( 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 @@ def prefetch_from_storage( 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 @@ def prefetch_from_storage( 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 @@ def _can_terminate_prefetch(self, operation: PrefetchOperation) -> bool: 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 @@ def _can_terminate_prefetch(self, operation: PrefetchOperation) -> bool: ) 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 @@ def check_prefetch_progress(self, request: CacheRequestHandle) -> bool: 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 @@ def check_prefetch_progress(self, request: CacheRequestHandle) -> bool: 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 @@ def _check_hybrid_prefetch_result( # 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 @@ def discard_storage_prefetch_accounting(self, request: CacheRequestHandle) -> No 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 @@ def pop_prefetch_loaded_span( self, request: CacheRequestHandle ) -> tuple[int, Optional[int]]: """Pop the loaded L3 token count and its absolute prefix start.""" - self._storage_prefetch_missed_rids.discard(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 @@ def _prefetch_occupied_span(self, prefetch_key, host_indices) -> int: 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 @@ def _try_alloc_storage_hit(operation) -> bool: # 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 @@ def _try_alloc_storage_hit(operation) -> bool: # 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 @@ def _try_alloc_storage_hit(operation) -> bool: 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 @@ def _drain_and_alloc_storage_hit(): 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 @@ def _drain_and_alloc_storage_hit(): 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 @@ def _drain_and_alloc_storage_hit(): 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 @@ def loading_check(self, finish_count: Optional[int] = None) -> None: 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 @@ def mamba_evictable_size(self) -> int: 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() diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 8f7d580e42e2..a92adbaad1ca 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -1944,6 +1944,7 @@ def __init__( "below_threshold", "host_capacity", "device_capacity", + "device_covered", "storage_transfer", "shrunk", "dropped", @@ -1952,6 +1953,17 @@ def __init__( **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 @@ def log_storage_prefetch_unfulfilled_tokens( **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) diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index c1a52a2ffbb5..96ec95f76468 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -24,6 +24,8 @@ 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 @@ def pretty_print(self): 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]]: diff --git a/rust/sglang-radix-tree/src/components/full.rs b/rust/sglang-radix-tree/src/components/full.rs index cd2ecfc41064..20f567b2c5da 100644 --- a/rust/sglang-radix-tree/src/components/full.rs +++ b/rust/sglang-radix-tree/src/components/full.rs @@ -391,6 +391,7 @@ impl TreeComponent for FullComponent { _host_indices: Option, _token_ids: Option<&[i64]>, _prefetch_tokens: usize, + _staging_tokens: usize, _last_hash: Option<&str>, ) -> Result>, TreeCoreRuntimeError> { Ok(match phase { diff --git a/rust/sglang-radix-tree/src/components/mamba.rs b/rust/sglang-radix-tree/src/components/mamba.rs index f2eddda7ece4..7d3e6cb50daa 100644 --- a/rust/sglang-radix-tree/src/components/mamba.rs +++ b/rust/sglang-radix-tree/src/components/mamba.rs @@ -494,6 +494,7 @@ impl TreeComponent for MambaComponent { host_indices: Option, _token_ids: Option<&[i64]>, _prefetch_tokens: usize, + staging_tokens: usize, _last_hash: Option<&str>, ) -> Result>, TreeCoreRuntimeError> { Ok(match phase { @@ -560,11 +561,13 @@ impl TreeComponent for MambaComponent { }]) } CacheTransferPhase::Prefetch => { - let host_indices = - host_indices.expect("Mamba PREFETCH build requires host indices"); + if staging_tokens == 0 { + return Ok(None); + } + // Staging is allocated once the hit is known; the placeholder + // key carries the single trailing page this pool loads. Some(vec![PoolTransfer { name: PoolName::Mamba, - host_indices: Some(host_indices), keys: Some(vec!["__placeholder__".to_string()]), hit_policy: PoolHitPolicy::TrailingPages, ..Default::default() diff --git a/rust/sglang-radix-tree/src/components/mod.rs b/rust/sglang-radix-tree/src/components/mod.rs index cf7154fe79f2..2337903c45ae 100644 --- a/rust/sglang-radix-tree/src/components/mod.rs +++ b/rust/sglang-radix-tree/src/components/mod.rs @@ -372,6 +372,7 @@ pub trait TreeComponent { host_indices: Option, token_ids: Option<&[i64]>, prefetch_tokens: usize, + staging_tokens: usize, last_hash: Option<&str>, ) -> Result>, TreeCoreRuntimeError> { // Python reference — base.py::TreeComponent.build_hicache_transfers: @@ -384,6 +385,7 @@ pub trait 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]]: // """Build transfer descriptors for this component in the given phase. diff --git a/rust/sglang-radix-tree/src/components/swa.rs b/rust/sglang-radix-tree/src/components/swa.rs index 4789adbe3c97..49bb05f94908 100644 --- a/rust/sglang-radix-tree/src/components/swa.rs +++ b/rust/sglang-radix-tree/src/components/swa.rs @@ -925,9 +925,10 @@ impl TreeComponent for SwaComponent { node_id: NodeIdx_, phase: CacheTransferPhase, _mamba_pool_idx: Option, - host_indices: Option, + _host_indices: Option, _token_ids: Option<&[i64]>, _prefetch_tokens: usize, + staging_tokens: usize, _last_hash: Option<&str>, ) -> Result>, TreeCoreRuntimeError> { // unified_kv keeps SWA as a device-only ring. @@ -1024,12 +1025,16 @@ impl TreeComponent for SwaComponent { }]) } CacheTransferPhase::Prefetch => { - let host_indices = host_indices.expect("SWA PREFETCH build requires host indices"); - let sw_pages = host_indices.numel() / tree_core.page_size; + // Staging is allocated once the hit is known; the placeholders + // carry the planned page count and the trailing hashes fill in + // at commit. + let num_pages = staging_tokens / tree_core.page_size; + if num_pages == 0 { + return Ok(None); + } Some(vec![PoolTransfer { name: PoolName::Swa, - host_indices: Some(host_indices), - keys: Some(vec!["__placeholder__".to_string(); sw_pages]), + keys: Some(vec!["__placeholder__".to_string(); num_pages]), hit_policy: PoolHitPolicy::TrailingPages, ..Default::default() }]) diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index c19793ec682a..a4e4e8524b53 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -999,6 +999,19 @@ impl TreeCoreBinding { MatchResultBinding::from_match_result(py, result) } + /// Read-only FULL-device match, independent of auxiliary components. + fn match_full_device_prefix( + &self, + py: Python<'_>, + params: &MatchParamsBinding, + ) -> (usize, NodeId, usize) { + let key = K::key_from(Cow::Borrowed(¶ms.key)); + let key = key.as_ref(); + let namespace = + KeyNamespaceRef::new(params.extra_key.as_deref(), params.cache_salt.as_deref()); + py.allow_threads(|| self.core().match_full_device_prefix(key, namespace)) + } + /// The empty match result anchored at the root. fn empty_match_result(&self, py: Python<'_>) -> PyResult { let result = py.allow_threads(|| self.core().empty_match_result()); @@ -1113,6 +1126,18 @@ impl TreeCoreBinding { Ok(IncLockRefResultBinding::from_result(result)) } + /// Pin only the FULL device values on a node's root path. + fn inc_full_pin(&self, py: Python<'_>, node_id: NodeId) -> PyResult<()> { + py.allow_threads(|| self.core().inc_full_pin(node_id)) + .map_err(node_access_error) + } + + /// Release a FULL-only root-path pin. + fn dec_full_pin(&self, py: Python<'_>, node_id: NodeId) -> PyResult<()> { + py.allow_threads(|| self.core().dec_full_pin(node_id)) + .map_err(node_access_error) + } + /// Decrease the reference count on a node's component locks. fn dec_lock_ref( &self, @@ -1486,6 +1511,7 @@ impl TreeCoreBinding { host_indices: Option, token_ids: Option>, prefetch_tokens: usize, + staging_tokens: usize, last_hash: Option, ) -> PyResult>>> { let component_type = parse_component_type(component_type)?; @@ -1500,6 +1526,7 @@ impl TreeCoreBinding { host_indices, token_ids.as_deref(), prefetch_tokens, + staging_tokens, last_hash.as_deref(), ) }) @@ -2423,6 +2450,15 @@ macro_rules! tree_core_binding { self.inner.match_prefix(py, params) } + /// Read-only FULL-device match, independent of auxiliary components. + fn match_full_device_prefix( + &self, + py: Python<'_>, + params: &MatchParamsBinding, + ) -> (usize, NodeId, usize) { + self.inner.match_full_device_prefix(py, params) + } + /// The empty match result anchored at the root. fn empty_match_result(&self, py: Python<'_>) -> PyResult { self.inner.empty_match_result(py) @@ -2472,6 +2508,16 @@ macro_rules! tree_core_binding { self.inner.inc_lock_ref(py, node_id, skip_lock_components) } + /// Pin only the FULL device values on a node's root path. + fn inc_full_pin(&self, py: Python<'_>, node_id: NodeId) -> PyResult<()> { + self.inner.inc_full_pin(py, node_id) + } + + /// Release a FULL-only root-path pin. + fn dec_full_pin(&self, py: Python<'_>, node_id: NodeId) -> PyResult<()> { + self.inner.dec_full_pin(py, node_id) + } + /// Decrease the reference count on a node's component locks. The /// receipt is required: a release must replay its acquire's evidence. #[pyo3(signature = (node_id, params, skip_swa = false))] @@ -2718,7 +2764,7 @@ macro_rules! tree_core_binding { /// Route a build_hicache_transfers call to the component for the given type. #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (component_type, node_id, phase, host_indices = None, token_ids = None, prefetch_tokens = 0, last_hash = None))] + #[pyo3(signature = (component_type, node_id, phase, host_indices = None, token_ids = None, prefetch_tokens = 0, staging_tokens = 0, last_hash = None))] fn build_hicache_transfers( &self, py: Python<'_>, @@ -2728,6 +2774,7 @@ macro_rules! tree_core_binding { host_indices: Option, token_ids: Option>, prefetch_tokens: usize, + staging_tokens: usize, last_hash: Option, ) -> PyResult>>> { self.inner.build_hicache_transfers( @@ -2738,6 +2785,7 @@ macro_rules! tree_core_binding { host_indices, token_ids, prefetch_tokens, + staging_tokens, last_hash, ) } diff --git a/rust/sglang-radix-tree/src/tests/components/full.rs b/rust/sglang-radix-tree/src/tests/components/full.rs index c466fd1981e1..956b97808ebb 100644 --- a/rust/sglang-radix-tree/src/tests/components/full.rs +++ b/rust/sglang-radix-tree/src/tests/components/full.rs @@ -2830,7 +2830,7 @@ fn build_hicache_transfers_returns_none_for_non_load_back_phases() { .build_hicache_transfers( &tc, a, phase, /* mamba_pool_idx = */ None, /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, - /* last_hash = */ None, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap(); assert!(transfers.is_none()); @@ -2850,6 +2850,7 @@ fn load_back_build_collects_the_evicted_suffix_ancestors_first() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap() @@ -2884,6 +2885,7 @@ fn load_back_build_returns_an_empty_cpu_transfer_for_a_device_backed_node() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap() @@ -2917,6 +2919,7 @@ fn load_back_build_panics_on_an_evicted_unbacked_node() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ); } diff --git a/rust/sglang-radix-tree/src/tests/components/mamba.rs b/rust/sglang-radix-tree/src/tests/components/mamba.rs index 1a1821e6be14..f223e7784a83 100644 --- a/rust/sglang-radix-tree/src/tests/components/mamba.rs +++ b/rust/sglang-radix-tree/src/tests/components/mamba.rs @@ -1071,6 +1071,7 @@ fn backup_host_build_carries_the_device_slot() { None, None, 0, + 0, None, ) .expect("live test node") @@ -1102,6 +1103,7 @@ fn backup_host_build_carries_the_device_slot() { None, None, 0, + 0, None, ) .expect("live test node") @@ -1122,6 +1124,7 @@ fn load_back_build_restores_the_host_only_node() { None, None, 0, + 0, None, ) .expect("live test node") @@ -1151,6 +1154,7 @@ fn load_back_build_skips_device_backed_and_bare_nodes() { None, None, 0, + 0, None, ) .expect("live test node") @@ -1173,6 +1177,7 @@ fn load_back_build_adds_the_per_request_cow_transfer() { None, None, 0, + 0, None, ) .unwrap() @@ -1388,6 +1393,7 @@ fn backup_storage_build_keys_the_trailing_hash() { None, None, 0, + 0, None, ) .expect("live test node") @@ -1403,6 +1409,7 @@ fn backup_storage_build_keys_the_trailing_hash() { None, None, 0, + 0, None, ) .expect("live test node") @@ -1417,6 +1424,7 @@ fn backup_storage_build_keys_the_trailing_hash() { None, None, 0, + 0, None, ) .expect("live test node") @@ -1434,23 +1442,28 @@ fn backup_storage_build_keys_the_trailing_hash() { } #[test] -fn prefetch_build_wraps_the_host_buffer_with_a_placeholder_key() { +fn prefetch_build_carries_a_placeholder_key_for_the_planned_slot() { let tc = mamba_core(/* page_size = */ 1); - let transfers = tc - .build_hicache_transfers( + let root_id = tc.arena.node(tc.arena.root()).id; + let build = |staging_tokens: usize| { + tc.build_hicache_transfers( MAMBA, - tc.arena.node(tc.arena.root()).id, + root_id, CacheTransferPhase::Prefetch, - Some(Tensor::from_slice(&[30i64])), + None, None, 0, + staging_tokens, None, ) .expect("live test node") - .unwrap(); + }; + let transfers = build(1).unwrap(); assert_eq!(transfers.len(), 1); assert_eq!(transfers[0].keys, Some(vec!["__placeholder__".to_string()])); assert_eq!(transfers[0].hit_policy, PoolHitPolicy::TrailingPages); + assert!(transfers[0].host_indices.is_none()); + assert!(build(0).is_none()); } #[test] diff --git a/rust/sglang-radix-tree/src/tests/components/swa.rs b/rust/sglang-radix-tree/src/tests/components/swa.rs index eab37178fc41..bfac6bc2f5e9 100644 --- a/rust/sglang-radix-tree/src/tests/components/swa.rs +++ b/rust/sglang-radix-tree/src/tests/components/swa.rs @@ -3595,6 +3595,7 @@ fn backup_storage_transfers_carry_trailing_page_keys() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap() @@ -3625,6 +3626,7 @@ fn backup_storage_is_none_without_host_value_or_hashes() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) }; @@ -3646,7 +3648,7 @@ fn build_transfers_are_gated_off_until_the_swa_host_pool_is_wired() { .build_hicache_transfers( &tc, a, phase, /* mamba_pool_idx = */ None, /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, - /* last_hash = */ None, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap(); assert!(transfers.is_none()); @@ -3662,6 +3664,7 @@ fn build_transfers_are_gated_off_until_the_swa_host_pool_is_wired() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap(); @@ -3684,6 +3687,7 @@ fn backup_host_build_wraps_the_device_value_as_int64() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap() @@ -3715,6 +3719,7 @@ fn backup_host_build_returns_none_for_a_tombstone() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap(); @@ -3765,6 +3770,7 @@ fn load_back_build_collects_host_only_nodes_within_the_window() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap() @@ -3805,6 +3811,7 @@ fn load_back_build_stops_at_the_window_boundary() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap() @@ -3838,6 +3845,7 @@ fn load_back_build_returns_none_when_the_window_is_on_device() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap(); @@ -3857,6 +3865,7 @@ fn load_back_build_rejects_a_bare_window_node() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ), Err(TreeCoreRuntimeError::SwaLoadBackMissingValue { node_id }) @@ -3880,6 +3889,7 @@ fn fallible_load_back_boundaries_reject_a_bare_window_node() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ), Err(TreeCoreRuntimeError::SwaLoadBackMissingValue { node_id: missing }) @@ -4112,21 +4122,26 @@ fn commit_hicache_transfers_routes_to_the_component() { } #[test] -fn prefetch_build_wraps_the_host_buffer_with_placeholder_keys() { +fn prefetch_build_sizes_the_placeholder_keys_from_the_staging_tokens() { let tc = swa_core(/* window = */ 4, /* page_size = */ 1); - let transfers = swa_component(4) - .build_hicache_transfers( - &tc, - tc.arena.root(), - CacheTransferPhase::Prefetch, - /* mamba_pool_idx = */ None, - /* host_indices = */ Some(Tensor::from_slice(&[30i64, 31])), - /* token_ids = */ None, - /* prefetch_tokens = */ 0, - /* last_hash = */ None, - ) - .unwrap() - .unwrap(); + let build = |staging_tokens: usize| { + swa_component(4) + .build_hicache_transfers( + &tc, + tc.arena.root(), + CacheTransferPhase::Prefetch, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + staging_tokens, + /* last_hash = */ None, + ) + .unwrap() + }; + // Staging is allocated once the hit is known: the build carries only the + // planned page count, never a host buffer. + let transfers = build(2).unwrap(); assert_eq!(transfers.len(), 1); assert_eq!( transfers[0].keys, @@ -4136,13 +4151,8 @@ fn prefetch_build_wraps_the_host_buffer_with_placeholder_keys() { ]) ); assert_eq!(transfers[0].hit_policy, PoolHitPolicy::TrailingPages); - assert!( - transfers[0] - .host_indices - .as_ref() - .unwrap() - .equal(&Tensor::from_slice(&[30i64, 31])) - ); + assert!(transfers[0].host_indices.is_none()); + assert!(build(0).is_none()); } #[test] @@ -5192,6 +5202,7 @@ fn backup_transfers( /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap() diff --git a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs index bd0765ed6998..5e81b9d0a839 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -1295,6 +1295,25 @@ fn match_prefix_splits_on_a_partial_match() { ); } +#[test] +fn match_full_device_prefix_is_read_only_and_accounts_the_pinned_node() { + let mut tc = core(); + let (a, _b) = matched_chain(&mut tc); + + let (matched_len, node_id, pinned_len) = + tc.match_full_device_prefix(&vec![1, 9], KeyNamespaceRef::new(None, None)); + + assert_eq!(matched_len, 1); + assert_eq!(node_id, tc.arena.node(a).id); + assert_eq!(pinned_len, 2); + assert_eq!(tc.arena.node(a).key, vec![1, 2]); + + tc.inc_full_pin(node_id).unwrap(); + assert_eq!(tc.arena.node(a).device_lock_ref(FULL), 1); + tc.dec_full_pin(node_id).unwrap(); + assert_eq!(tc.arena.node(a).device_lock_ref(FULL), 0); +} + #[test] fn match_prefix_stops_at_a_dead_node() { // An evicted, unbackuped child ends the traversal before it. @@ -4346,6 +4365,7 @@ fn fallible_node_boundaries_reject_stale_handles() { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ), Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id })) diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 9600aab958ca..ca42e67613bf 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -850,6 +850,34 @@ impl UnifiedTreeCore { Ok(result) } + /// Pin only the FULL device values on a node's root path. + pub fn inc_full_pin(&mut self, node_id: NodeId) -> Result<(), NodeAccessError> { + let node_idx = self.arena.resolve(node_id)?; + let full = self.component_by_type_(FULL); + full.acquire_component_lock( + self, + node_idx, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + self.update_evictable_leaf_sets_(node_idx); + Ok(()) + } + + /// Release a FULL-only root-path pin. + pub fn dec_full_pin(&mut self, node_id: NodeId) -> Result<(), NodeAccessError> { + let node_idx = self.arena.resolve(node_id)?; + let full = self.component_by_type_(FULL); + full.release_component_lock( + self, + node_idx, + &DecLockRefParams::default(), + /* lock_host = */ false, + ); + self.update_evictable_leaf_sets_(node_idx); + Ok(()) + } + /// A receipt releases only the node its acquire returned; a mispaired /// node would silently release (or steal) another holder's segment. fn assert_receipt_anchor_(&self, node_idx: NodeIdx_, params: &DecLockRefParams) { @@ -1035,6 +1063,44 @@ impl UnifiedTreeCore { ) } + /// Read-only FULL-device match, independent of auxiliary components. + /// Returns the request match and the complete root-path length pinned by + /// the deepest node; they differ when the key ends inside that node. + pub fn match_full_device_prefix( + &self, + key: &K, + namespace: KeyNamespaceRef<'_>, + ) -> (usize, NodeId, usize) { + let aligned_key_len = key.atom_len() / self.page_size * self.page_size; + let mut node_id = self.arena.root(); + let mut offset = 0; + let mut pinned_len = 0; + while offset < aligned_key_len { + let Some(child_id) = self.arena.child_on_page_in_namespace( + node_id, + namespace, + key.page_at(offset, self.page_size), + ) else { + break; + }; + let child = self.arena.node(child_id); + if !child.has_device_value(FULL) { + break; + } + let prefix_len = key.match_len(offset, &child.key, self.page_size); + if prefix_len == 0 { + break; + } + offset += prefix_len; + pinned_len += child.device_value_len(FULL); + node_id = child_id; + if prefix_len < child.key.atom_len() { + break; + } + } + (offset, self.arena.node(node_id).id, pinned_len) + } + /// Walk the tree for `key`; returns matched value chunks, the best match, /// the best device-resident match, its device value length, and any split action. pub fn match_prefix_helper_( @@ -3178,6 +3244,7 @@ impl UnifiedTreeCore { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap(); @@ -3217,6 +3284,7 @@ impl UnifiedTreeCore { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, ) .unwrap(); @@ -3244,6 +3312,7 @@ impl UnifiedTreeCore { host_indices: Option, token_ids: Option<&[i64]>, prefetch_tokens: usize, + staging_tokens: usize, last_hash: Option<&str>, ) -> Result>, TreeCoreRuntimeError> { let node_id = self.arena.resolve(node_id)?; @@ -3256,6 +3325,7 @@ impl UnifiedTreeCore { host_indices, token_ids, prefetch_tokens, + staging_tokens, last_hash, ) } @@ -3281,6 +3351,7 @@ impl UnifiedTreeCore { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, )? .unwrap(); @@ -3299,6 +3370,7 @@ impl UnifiedTreeCore { /* host_indices = */ None, /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* staging_tokens = */ 0, /* last_hash = */ None, )?; if let Some(transfers) = transfers diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 77bd0a6e3ff7..d8b76b84c866 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -2,6 +2,8 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch +import torch + import sglang.srt.managers.schedule_policy as schedule_policy from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_policy import ( @@ -57,6 +59,7 @@ def create_tree_cache( tree_cache.disable = False tree_cache.inc_lock_ref.return_value = IncLockRefResult() tree_cache.dec_lock_ref.return_value = DecLockRefResult() + tree_cache.buffer_pipeline = None return tree_cache def create_token_allocator( @@ -102,6 +105,7 @@ def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time) req.retracted_stain = False 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 @@ -158,6 +162,24 @@ def test_storage_prefetch_fulfillment_resolves_at_admission(self): req.cache_request_handle, fulfilled_tokens=0, reason="device_capacity" ) + self.mock_tree_cache.finish_storage_prefetch_admission.reset_mock() + req.host_hit_length = 4 + req.host_loaded_length = 4 + req.storage_hit_length = 8 + req.storage_hit_start = 4 + req.materialized_host_hit_len.return_value = 4 + req.fulfilled_storage_hit_len.return_value = 4 + req.needs_host_load_back.return_value = True + adder._account_prefill_cache_admission(req, prefix_len=8) + self.mock_tree_cache.finish_storage_prefetch_admission.assert_called_once_with( + req.cache_request_handle, + fulfilled_tokens=4, + reason="shrunk", + ) + self.assertEqual(adder.log_device_hit_tokens, 8) + self.assertEqual(adder.log_host_hit_tokens, 0) + self.assertEqual(adder.log_storage_hit_tokens, 12) + def test_retracted_storage_prefetch_accounting_is_omitted(self): adder = self.create_adder(self.create_running_batch()) req = self.create_mock_req( @@ -668,6 +690,169 @@ def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self): adder.add_one_req(req, has_chunked_req=False, truncation_align_size=None) self.assertIn(req, adder.can_run_list) + def test_load_back_delivery_mismatch_reselects_the_prefill_shape(self): + # Two incidents: a load that delivers nothing left the SWA gate sized + # for the tail and the allocator OOMed; a cache-mode load that also + # surfaces FULL device tokens behind a host-only SWA window tripped a + # strict promised==loaded check and crashed the scheduler. + WINDOW, PAGE = 128, 8 + SPAN, HOST_HIT = 1024, 1016 + self.mock_token_allocator.swa_available_size.return_value = 400 + self.mock_token_allocator.full_available_size.return_value = 100_000 + self.mock_token_allocator.available_size.return_value = 100_000 + self.mock_tree_cache.sliding_window_size = WINDOW + self.mock_tree_cache.is_tree_cache.return_value = False + + def run(delivered: int, remaining_after_load: int = 100_000): + self.mock_token_allocator.full_available_size.return_value = 100_000 + self.mock_token_allocator.swa_available_size.return_value = 400 + adder = self.create_adder(self.create_running_batch(), page_size=PAGE) + adder.is_hybrid_swa = True + req = self.create_mock_req("dropped-fetch", priority=0, max_new_tokens=8) + req.prefix_indices = torch.empty(0, dtype=torch.int64) + req.full_untruncated_fill_ids = list(range(SPAN)) + req.host_hit_length = HOST_HIT + req.swa_host_hit_length = WINDOW + req.needs_host_load_back.return_value = True + req.last_node = MagicMock() + req.best_match_node = MagicMock() + req.kv = SimpleNamespace(cache_protected_len=0) + + def set_extend_range(start, end): + req.extend_range = Range(start, end) + + req.set_extend_range = MagicMock(side_effect=set_extend_range) + req.sampling_params = SimpleNamespace(max_new_tokens=8, ignore_eos=False) + + def load_back(params): + self.mock_token_allocator.full_available_size.return_value = ( + remaining_after_load + ) + if remaining_after_load == 0: + self.mock_token_allocator.swa_available_size.return_value = 0 + return torch.arange(delivered, dtype=torch.int64), req.last_node + + self.mock_tree_cache.init_load_back.side_effect = load_back + verdict = adder.add_one_req( + req, has_chunked_req=False, truncation_align_size=None + ) + return verdict, list(adder.can_run_list), req + + # Promise kept: only the 8-token tail is prefilled, which fits. + _, admitted, _ = run(HOST_HIT) + self.assertEqual(len(admitted), 1) + # Nothing delivered: the whole span is prefilled and no longer fits, so + # admission must decline rather than OOM the pool. + verdict, admitted, _ = run(0) + self.assertIs(verdict, AddReqResult.NO_TOKEN) + self.assertEqual(admitted, []) + # Over-delivery: admitted with the loaded prefix, not the promise. + # The loaded prefix is now pinned and no longer part of the evictable + # budget. A successful load must not run admission gates again. + _, admitted, req = run(HOST_HIT + 4, remaining_after_load=0) + self.assertEqual(len(admitted), 1) + self.assertEqual(len(req.prefix_indices), HOST_HIT + 4) + self.assertEqual(req.kv.cache_protected_len, HOST_HIT + 4) + req.set_extend_range.assert_called_once_with(HOST_HIT + 4, SPAN) + # A partial FULL load stays fatal. + with self.assertRaisesRegex(RuntimeError, "promised"): + run(HOST_HIT // 2) + + def _create_host_hit_req(self, *, prefix_len=0, host_hit=8192, tail=1024): + req = self._create_delayer_req(prefix_len + host_hit + tail) + req.prefix_indices = torch.arange(prefix_len) + req.host_hit_length = host_hit + req.needs_host_load_back.return_value = True + req.best_match_node = req.last_node + req.kv = SimpleNamespace(cache_protected_len=prefix_len) + return req + + def test_successful_load_back_commits_the_selected_shape_once(self): + cases = ( + ("full", 0, 24, None, None, 8, 8), + ("full_unaligned", 0, 24, None, None, 7, 8), + ("retracted_unaligned", 0, 24, None, None, 7, 8), + ("chunk", 0, 24, 4, None, 4, 0), + ("aux_only", 24, 0, None, None, 8, 8), + ("overdelivery_full", 0, 24, None, None, 8, 8), + ("overdelivery_chunk", 0, 24, 4, None, 4, 0), + ("overdelivery_chunk_end", 0, 24, 8, None, 8, 8), + ( + "dllm", + 0, + 24, + None, + SimpleNamespace(block_size=4, max_running_requests=2), + 4, + 0, + ), + ( + "overdelivery_dllm", + 0, + 24, + None, + SimpleNamespace(block_size=4, max_running_requests=2), + 4, + 0, + ), + ) + for name, prefix_len, host_hit, chunk, dllm, extend, decode in cases: + with self.subTest(mode=name): + self.mock_tree_cache.reset_mock() + adder = self._create_delayer_adder( + available_tokens=100_000, + delayer=None, + page_size=2, + rem_chunk_tokens=chunk, + dllm_config=dllm, + ) + req = self._create_host_hit_req( + prefix_len=prefix_len, + host_hit=host_hit, + tail=extend if chunk is None and dllm is None else 8, + ) + req.retracted_stain = name == "retracted_unaligned" + if name.startswith("overdelivery"): + req.host_hit_length -= 4 + old_node, restored_node = req.last_node, object() + if name == "aux_only": + req.swa_host_hit_length = 8 + + def load_back(params): + self.assertIs(params.req, req) + tile_gate.assert_called_once() + tile_gate.return_value = AddReqResult.OTHER + return torch.arange(host_hit), restored_node + + self.mock_tree_cache.init_load_back.side_effect = load_back + with patch.object( + adder, "_check_prefill_tile_budget", return_value=None + ) as tile_gate: + adder.add_one_req(req, False, None) + tile_gate.assert_called_once() + self.mock_tree_cache.init_load_back.assert_called_once() + self.assertEqual(adder.can_run_list, [req]) + req.set_extend_range.assert_called_once_with(24, 24 + extend) + self.mock_tree_cache.inc_lock_ref.assert_any_call(restored_node) + self.assertIs( + self.mock_tree_cache.dec_lock_ref.call_args.args[0], old_node + ) + self.assertEqual(adder.log_hit_tokens, 24) + self.assertEqual(adder.log_input_tokens, extend) + self.assertEqual( + adder.reprocessed_log_input_tokens, + extend if req.retracted_stain else 0, + ) + self.assertEqual( + adder.rem_total_token_offset, + adder.ceil_paged_tokens(extend) + decode + 2, + ) + self.assertEqual( + adder.new_chunked_req is req, + name in ("chunk", "overdelivery_chunk"), + ) + self.mock_tree_cache.init_load_back.side_effect = None + def test_swa_new_tokens_clamps_remaining_not_total(self): # Remaining decode headroom must be min(max_new - generated, CLIP) # (subtract-then-clip). The reversed order (clip-then-subtract) zeroes diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index d13655d90cd0..94ac81dd538b 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -81,6 +81,10 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler: s.dllm_manager = None s.enable_hisparse = False s.enable_fpm = False + # Exercise the unconditional scheduler-loop HiCache event-drain point. + s.enable_hierarchical_cache = True + s.enable_hicache_storage = False + s.enable_unified_cache_external_linker = False s.last_batch = None s.require_mlp_sync = False s.spec_algorithm = MagicMock() @@ -104,6 +108,7 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler: side_effect=lambda batch, **_: batch ) s.update_running_batch = MagicMock(side_effect=lambda batch: batch) + tree_cache.check_hicache_events = MagicMock() s.tree_cache = tree_cache s.chunked_req = chunked_req s._pending_chunked_abort_req = None diff --git a/test/registered/unit/managers/test_scheduler_hicache_events.py b/test/registered/unit/managers/test_scheduler_hicache_events.py new file mode 100644 index 000000000000..79b764182ef0 --- /dev/null +++ b/test/registered/unit/managers/test_scheduler_hicache_events.py @@ -0,0 +1,139 @@ +"""HiCache progress must not depend on whether a prefill batch is admitted.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, call, patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.schedule_batch import NextBatchPlan +from sglang.srt.managers.scheduler import Scheduler + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestSchedulerHiCacheEvents(unittest.TestCase): + def setUp(self): + self.calls = Mock() + self.scheduler = s = Scheduler.__new__(Scheduler) + s.scheduler_stage_metrics = None + s.enable_hierarchical_cache = True + s.enable_unified_cache_external_linker = False + s.enable_hicache_storage = True + s.tree_cache = SimpleNamespace(check_hicache_events=self.calls.drain) + s._process_storage_prefetch_retries = self.calls.retry + s.process_pending_chunked_abort = Mock() + s.process_prefill_chunk = Mock() + s.dp_attn_adapter = Mock() + s.dp_attn_adapter.maybe_prepare_mlp_sync_batch.return_value = None + self.running_batch = SimpleNamespace( + batch_is_full=False, is_prefill_only=False, is_empty=lambda: True + ) + + def test_feature_gates(self): + for hierarchical, flexkv, linker, storage in ( + (False, False, False, False), + (True, False, False, False), + (False, True, False, False), + (False, False, True, False), + (True, False, False, True), + ): + with ( + self.subTest( + hierarchical=hierarchical, + flexkv=flexkv, + linker=linker, + storage=storage, + ), + patch( + "sglang.srt.managers.scheduler.get_memory", + return_value=SimpleNamespace(enable_flexkv=flexkv), + ), + ): + self.calls.reset_mock() + s = self.scheduler + s.enable_hierarchical_cache = hierarchical + s.enable_unified_cache_external_linker = linker + s.enable_hicache_storage = storage + s._process_hicache_events() + expected = [call.drain()] if hierarchical or flexkv or linker else [] + if storage: + expected.append(call.retry()) + self.assertEqual(self.calls.mock_calls, expected) + + def test_pd_prefill_drains_before_admission_even_with_empty_queue(self): + s = self.scheduler + s.resolve_waiting_queue_bootstrap = Mock() + s.get_new_batch_prefill = self.calls.admit + s.get_new_batch_prefill.return_value = NextBatchPlan( + batch_to_run=None, running_batch=self.running_batch + ) + for waiting_queue in ([], [SimpleNamespace(rid="pending_l3")]): + with self.subTest(waiting=bool(waiting_queue)): + s.waiting_queue = waiting_queue + self.calls.reset_mock() + for _ in range(2): + plan = s.get_next_disagg_prefill_batch_to_run( + running_batch=self.running_batch, last_batch=None + ) + self.assertIsNone(plan.batch_to_run) + self.assertEqual( + self.calls.mock_calls, + [call.drain(), call.retry(), call.admit(self.running_batch)] * 2, + ) + + def test_unified_drains_when_prefill_is_deferred(self): + s = self.scheduler + s.enable_fpm = False + s._abort_on_waiting_timeout = Mock() + s._abort_on_running_timeout = Mock() + s.dllm_config = None + s.chunked_req = None + s.enable_hisparse = False + s.require_mlp_sync = False + s._should_defer_prefill = self.calls.defer + s._should_defer_prefill.return_value = True + s.get_new_batch_prefill = Mock() + s.dp_attn_adapter.maybe_convert_decode_to_extend.return_value = None + s._arm_prefill_decode_interval = Mock() + s.ngram_embedding_manager = Mock() + s.ngram_embedding_manager.prepare_for_forward.return_value = None + + plan = s.get_next_batch_to_run(self.running_batch, None) + + self.assertIsNone(plan.batch_to_run) + s.get_new_batch_prefill.assert_not_called() + self.assertEqual( + self.calls.mock_calls, [call.drain(), call.retry(), call.defer()] + ) + + def test_pp_prefill_drains_before_admission(self): + s = self.scheduler + s.init_pp_loop_state = Mock() + s.pp_loop_size = 1 + s.ps = SimpleNamespace(pp_size=2) + s.pp_group = SimpleNamespace(is_last_rank=True) + s.running_mbs = [self.running_batch] + s.last_mbs = [None] + s.ingest_requests = Mock(return_value=[]) + s._pp_pd_get_bootstrapped_ids = Mock(return_value=[]) + s._pp_pd_get_prefill_transferred_ids = Mock(return_value=[]) + s._pp_commit_comm_work = Mock() + s.get_new_batch_prefill = self.calls.admit + # Stop the infinite loop at admission; no PP transport or GPU is needed. + s.get_new_batch_prefill.side_effect = StopIteration + + with self.assertRaises(StopIteration): + s.event_loop_pp_disagg_prefill() + + self.assertEqual( + self.calls.mock_calls, + [call.drain(), call.retry(), call.admit(self.running_batch)], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py index 622afdfe7dba..a62f10c1ea40 100644 --- a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py +++ b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py @@ -26,6 +26,7 @@ from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( BufferBackupSnapshot, ) +from sglang.srt.mem_cache.unified_radix_cache import _OngoingPrefetch from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=8, suite="base-a-test-cpu") @@ -244,13 +245,13 @@ def test_completed_prefetch_keeps_dsv4_full_and_swa_sidecars_for_h2d(self): cache.page_size = 2 cache.cache_controller.prefetch_tokens_occupied = len(host_indices) cache.ongoing_prefetch = { - req_id: ( - 0, - RadixKey(array("q", [1, 2, 3, 4])), - host_indices, - operation, - None, - {ComponentType.SWA: [swa]}, + req_id: _OngoingPrefetch( + anchor_node_id=0, + prefetch_key=RadixKey(array("q", [1, 2, 3, 4])), + host_indices=host_indices, + operation=operation, + anchor_lock_params=None, + comp_xfers={ComponentType.SWA: [swa]}, ) } cache.prefetch_loaded_tokens_by_reqid = {} diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index 7cbb73234748..7591102e382a 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -1,6 +1,7 @@ """Unit tests for HiCache staged write-back host-pool dispatch.""" import unittest +from array import array from contextlib import contextmanager from types import SimpleNamespace from unittest import mock @@ -264,15 +265,22 @@ def test_hybrid_load_forwards_merged_pool_transfers(self): controller._num_tokens_by_pool.assert_called_once_with(merged_op) self.assertEqual(controller.ack_load_queue[0].node_ids, [7, 7]) - def test_short_staged_swa_tail_resolves_device_covered_head(self): + def _short_swa_tail_pipeline(self, swa_page_size: int) -> BufferModePipeline: + """Pipeline holding one staged span [2, 8) whose 4-slot trailing SWA + window outruns the splice left by a device prefix of 6.""" handle = CacheRequestHandle("r", 0) pipeline = BufferModePipeline.__new__(BufferModePipeline) pipeline._cache = mock.Mock() + pipeline._cache.cache_controller.mem_pool_host.entry_map = { + PoolName.SWA: SimpleNamespace( + host_pool=SimpleNamespace(page_size=swa_page_size) + ) + } pipeline.release_staged_hold = mock.Mock(return_value=True) pipeline.staged_prefetches = { handle: SimpleNamespace( request=handle, - key_tokens=list(range(8)), + key_tokens=array("q", range(8)), extra_key=None, cache_salt=None, matched_len=2, @@ -289,14 +297,42 @@ def test_short_staged_swa_tail_resolves_device_covered_head(self): operation_id=1, ) } + return pipeline - self.assertEqual( - pipeline.plan_staged_splice(handle, device_prefix_len=6), (0, 0) + def test_short_staged_swa_tail_keeps_complete_window(self): + """FULL-prefix growth trims only FULL; SWA keeps its complete window.""" + handle = CacheRequestHandle("r", 0) + pipeline = self._short_swa_tail_pipeline(swa_page_size=2) + pipeline._cache.tree_core.is_eagle = False + pipeline._cache.tree_core.match_full_device_prefix.return_value = (6, 1, 6) + pipeline._cache.tree_core.collect_full_device_indices.return_value = _indices( + 0, 6 + ) + req = SimpleNamespace( + rid="r", + cache_request_handle=handle, + prefix_indices=_indices(0, 0), + kv=SimpleNamespace(cache_protected_len=0), + ) + self.assertTrue(pipeline.prepare_staged_prefetch(req)) + self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (2, 4)) + pipeline.release_staged_hold.assert_not_called() + + pipeline = self._short_swa_tail_pipeline(swa_page_size=4) + pipeline._cache.tree_core.is_eagle = False + pipeline._cache.tree_core.match_full_device_prefix.return_value = (6, 1, 6) + pipeline._cache.tree_core.collect_full_device_indices.return_value = _indices( + 0, 6 ) - pipeline._cache._resolve_storage_prefetch_tokens.assert_called_once_with( - handle, 4 + req = SimpleNamespace( + rid="r", + cache_request_handle=handle, + prefix_indices=_indices(0, 0), + kv=SimpleNamespace(cache_protected_len=0), ) - pipeline.release_staged_hold.assert_called_once_with(handle, reason="shrunk") + self.assertTrue(pipeline.prepare_staged_prefetch(req)) + self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (2, 4)) + pipeline.release_staged_hold.assert_not_called() def test_l2_transfer_maps_global_layers(self): host_pool = mock.Mock() diff --git a/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py b/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py new file mode 100644 index 000000000000..eb58879c97bf --- /dev/null +++ b/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py @@ -0,0 +1,498 @@ +"""Staged L3 prefetch lifecycle through the buffer pipeline; no GPU kernels.""" + +import tempfile +import unittest +from array import array +from collections import defaultdict, deque +from datetime import timedelta +from queue import Queue +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from sglang.srt.mem_cache.base_prefix_cache import ( + CacheRequestHandle, + InitLoadBackParams, +) +from sglang.srt.mem_cache.buffer_mode.pipeline import ( + BufferModePipeline, + _StagedPrefetch, +) +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, +) +from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries +from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedRadixCache, + _OngoingPrefetch, +) +from sglang.srt.mem_cache.utils import get_hash_str +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=20, suite="base-a-test-cpu") + +_REQ = CacheRequestHandle("r", 0) + + +def _staged_fixture(full_match=2): + cache = UnifiedRadixCache.__new__(UnifiedRadixCache) + cache.tree_core = SimpleNamespace( + page_size=2, + is_eagle=False, + enable_storage=True, + prefetch_anchor_info=lambda node: (None, None), + match_full_device_prefix=Mock(return_value=(full_match, 1, full_match)), + collect_full_device_indices=Mock(return_value=torch.arange(8)), + inc_full_pin=Mock(), + dec_full_pin=Mock(), + empty_match_result=SimpleNamespace( + last_device_node=0, device_indices=torch.arange(0) + ), + ) + cache.host_memory_mode = "buffer_only" + cache.linker = None + cache.storage_prefetch_retries = StoragePrefetchRetries() + cache.prefetch_loaded_tokens_by_reqid = {_REQ: 6} + cache.prefetch_loaded_storage_start_by_reqid = {_REQ: 2} + cache._storage_prefetch_hit_remaining_by_reqid = {} + cache.enable_storage_metrics = False + cache.storage_metrics_collector = None + cache.ongoing_prefetch = {} + cache._prefetch_outcome_stats = defaultdict(int) + cache.tree_components = [] + cache.prefetch_threshold = 2 + cache._build_sidecar_transfers = Mock(return_value=[]) + cache.supports_swa = lambda: True + cache.evict_for_alloc = Mock() + cache.token_to_kv_pool_allocator = SimpleNamespace( + full_available_size=Mock(return_value=100) + ) + cc = HybridCacheController.__new__(HybridCacheController) + cc.page_size = 2 + cc.get_hash_str = get_hash_str + cc.prefetch_queue = Queue() + cc.prefetch_tokens_occupied = 6 + cc.prefetch_rate_limited = lambda: False + cc.load = Mock(return_value=None) + cc.storage_backend = Mock() + cc.storage_backend.batch_exists.return_value = 0 + cc.mem_pool_host = SimpleNamespace( + free=Mock(), + entry_map={ + PoolName.SWA: SimpleNamespace(host_pool=SimpleNamespace(free=Mock())) + }, + ) + cache.cache_controller = cc + pipeline = BufferModePipeline.__new__(BufferModePipeline) + pipeline._cache = cache + pipeline.reset() + cache.buffer_pipeline = pipeline + pipeline.staged_prefetches[_REQ] = _StagedPrefetch( + request=_REQ, + key_tokens=array("q", range(8)), + extra_key=None, + cache_salt=None, + matched_len=2, + num_tokens=6, + occupied_tokens=6, + host_indices=torch.arange(6), + aux_xfers=[PoolTransfer(name=PoolName.SWA, host_indices=torch.arange(4))], + hash_values=["a", "b", "c"], + operation_id=1, + ) + req = SimpleNamespace( + rid="r", + cache_request_handle=_REQ, + prefix_indices=torch.arange(2), + last_node=1, + kv=SimpleNamespace(cache_protected_len=2), + extra_key=None, + cache_salt=None, + host_hit_length=0, + swa_host_hit_length=0, + host_hit_is_storage=False, + host_loaded_length=0, + storage_prefetch_last_match_len=4, + storage_prefetch_retry_attempts=0, + ) + return cache, pipeline, req + + +def _hit_drain_fixture(): + """A buffer-mode cache whose hit drain and outcome accounting are real.""" + cache = UnifiedRadixCache.__new__(UnifiedRadixCache) + cache.host_memory_mode = "buffer_only" + cache.prefetch_threshold = 2 + cache.enable_storage_metrics = False + cache.storage_metrics_collector = None + cache.storage_prefetch_retries = StoragePrefetchRetries() + cache._prefetch_outcome_stats = defaultdict(int) + cache._storage_prefetch_hit_remaining_by_reqid = {} + cache._record_storage_prefetch_hit = Mock() + cache.revoke_pending_prefetch = Mock() + cache.buffer_pipeline = SimpleNamespace(pending_hit_allocs=deque()) + cache.cache_controller = SimpleNamespace( + prefetch_hit_queue=Queue(), + ack_prefetch_queue=Queue(), + ack_backup_queue=Queue(), + host_mem_release_queue=Queue(), + extra_host_mem_release_queues={}, + ) + cache.ongoing_prefetch = {} + return cache + + +def _terminated_query(cache, rid, hit_tokens): + handle = CacheRequestHandle(rid, 0) + operation = SimpleNamespace( + request_id=rid, + handle=handle, + storage_hit_count=hit_tokens, + stats_requested_tokens=8, + is_terminated=lambda: True, + ) + cache.ongoing_prefetch[handle] = _OngoingPrefetch( + 0, RadixKey(array("q", range(8))), None, operation, None, {} + ) + cache.cache_controller.prefetch_hit_queue.put(operation) + + +def _two_rank_retry_trace(rank, rendezvous): + torch.distributed.init_process_group( + "gloo", + init_method=f"file://{rendezvous}", + rank=rank, + world_size=2, + timeout=timedelta(seconds=30), + ) + try: + cache, pipeline, req = _staged_fixture() + held = pipeline.staged_prefetches.pop(req.cache_request_handle) + cache.cache_controller.prefetch_tokens_occupied = 0 + waiting = [SimpleNamespace(rid="head", storage_prefetch_retry_attempts=0), req] + published = False + issued = [] + for step in range(7): + # Native completion arrives one pass earlier on rank 0. Only the + # agreed completion may publish staging to the admission path. + ready = torch.tensor([int(step >= rank + 1)]) + torch.distributed.all_reduce(ready, op=torch.distributed.ReduceOp.MIN) + if ready.item() and not published: + pipeline.staged_prefetches[req.cache_request_handle] = held + cache.cache_controller.prefetch_tokens_occupied = 6 + published = True + for retry_req, hit_end in cache.storage_prefetch_retries.pop_ready( + waiting, 2, 8 + ): + issued.append((step, retry_req.rid, hit_end)) + if pipeline.has_staged(req.cache_request_handle): + full_match = 2 if step < 3 else 0 + cache.tree_core.match_full_device_prefix.return_value = ( + full_match, + 1, + full_match, + ) + req.prefix_indices = torch.arange(full_match) + if cache.buffer_pipeline.prepare_staged_prefetch(req): + assert ( + cache.init_load_back( + InitLoadBackParams(None, req.host_hit_length, req=req) + ) + is None + ) + snapshot = ( + list(issued), + pipeline.has_staged(req.cache_request_handle), + cache.cache_controller.prefetch_tokens_occupied, + ) + snapshots = [None, None] + torch.distributed.all_gather_object(snapshots, snapshot) + assert snapshots[0] == snapshots[1], (step, snapshots) + assert issued == [(4, "r", 8)], issued + assert cache.cache_controller.load.call_count == 1 + assert cache.cache_controller.prefetch_tokens_occupied == 0 + finally: + torch.distributed.destroy_process_group() + + +class TestStagedPrefetchLifecycle(unittest.TestCase): + def test_trim_and_stage_preserve_raw_token_boundaries(self): + for bigram in (False, True): + for trims in ((2,), (2, 2), (8,), (2, 6)): + with self.subTest(bigram=bigram, trims=trims): + cache, pipeline, req = _staged_fixture() + pipeline.staged_prefetches.clear() + cache.tree_core.is_eagle = bigram + tokens = array("q", range(10 + int(bigram))) + cache.prefetch_from_storage( + req.cache_request_handle, + 0, + tokens[2:], + matched_prefix_tokens=tokens[:2], + storage_hit_end=10, + ) + info = cache.ongoing_prefetch[req.cache_request_handle] + operation = info.operation + self.assertEqual(len(info.prefetch_key), 8) + self.assertTrue(operation.assume_stored) + operation.hash_value = ["h0", "h1", "h2", "h3"] + operation.storage_hit_count = 8 + matched_len, hit_tokens = 2, 8 + for trim in trims: + matched_len += trim + info, hit_tokens, aux_tokens = ( + cache._trim_buffer_prefetch_full_head( + req.cache_request_handle, + info, + operation, + matched_len, + hit_tokens, + ) + ) + self.assertEqual(aux_tokens, 8) + self.assertEqual( + pipeline._prefetch_prefix_ctx[req.cache_request_handle][0], + list(tokens[:matched_len]), + ) + self.assertEqual( + list(info.prefetch_key.raw_token_ids()), + list(tokens[matched_len:]), + ) + # A capacity retry must retain the endpoint even when + # there is no FULL suffix and only SWA remains to load. + cache.tree_core.match_full_device_prefix.return_value = ( + matched_len, + 1, + matched_len, + ) + pipeline.anchor_lock_cap_tokens = 100 + pipeline.try_lock_anchor(req.cache_request_handle, hit_tokens) + anchor_key = ( + cache.tree_core.match_full_device_prefix.call_args.args[0] + ) + self.assertEqual(anchor_key.raw_token_ids(), tokens) + pipeline.release_anchor_lock(req.cache_request_handle) + cache.tree_core.match_full_device_prefix.reset_mock() + swa = PoolTransfer(name=PoolName.SWA, host_indices=torch.arange(4)) + operation.pool_transfers = [swa] + operation.host_indices = torch.arange(hit_tokens) + cache.ongoing_prefetch[req.cache_request_handle] = info._replace( + host_indices=operation.host_indices, comp_xfers={"swa": [swa]} + ) + cache.cache_controller.prefetch_tokens_occupied = hit_tokens + cache.storage_existence_cache = Mock() + pipeline.stage_completed_prefetch( + req.cache_request_handle, hit_tokens, operation.hash_value + ) + held = pipeline.staged_prefetches[req.cache_request_handle] + self.assertEqual(held.key_tokens, tokens) + self.assertEqual(held.matched_len, matched_len) + self.assertEqual(held.num_tokens, hit_tokens) + self.assertEqual( + len(RadixKey(held.key_tokens, is_bigram=bigram)), 10 + ) + # A joint match counts bigrams, not their extra raw boundary token. + req.prefix_indices = torch.arange(10) + req.kv.cache_protected_len = 10 + self.assertTrue(pipeline.prepare_staged_prefetch(req)) + self.assertFalse(pipeline.has_staged(req.cache_request_handle)) + cache.tree_core.match_full_device_prefix.assert_not_called() + + def test_two_rank_completion_capacity_and_anchor_loss(self): + with tempfile.TemporaryDirectory(prefix="prefetch-rank-test-") as directory: + torch.multiprocessing.spawn( + _two_rank_retry_trace, args=(f"{directory}/store",), nprocs=2, join=True + ) + + def test_next_pass_uses_fresh_joint_match(self): + cache, pipeline, req = _staged_fixture(full_match=8) + self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req)) + self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (0, 4)) + self.assertEqual(req.kv.cache_protected_len, 8) + plan = req.staged_prefetch_plan + self.assertIs( + plan.key.token_ids, + pipeline.staged_prefetches[req.cache_request_handle].key_tokens, + ) + cache.tree_core.match_full_device_prefix.assert_called_once() + # A twin finished first: the next pass's joint match runs past the + # staged span and is kept as is (a shrink would strand its recompute). + req.prefix_indices = torch.arange(12) + req.last_node = 9 + req.kv.cache_protected_len = 12 + self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req)) + self.assertIsNone(req.staged_prefetch_plan) + self.assertEqual(len(req.prefix_indices), 12) + self.assertEqual((req.last_node, req.kv.cache_protected_len), (9, 12)) + self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (0, 0)) + self.assertFalse(pipeline.has_staged(req.cache_request_handle)) + self.assertEqual(cache.cache_controller.prefetch_tokens_occupied, 0) + cache.cache_controller.mem_pool_host.free.assert_called_once() + cache.tree_core.match_full_device_prefix.assert_called_once() + + def test_capacity_retry_keeps_buffers_without_a_storage_retry(self): + for available in (0, 100): + with self.subTest(full_available=available): + cache, pipeline, req = _staged_fixture() + cache.cache_controller.load.return_value = None + cache.token_to_kv_pool_allocator.full_available_size.return_value = ( + available + ) + held = pipeline.staged_prefetches[req.cache_request_handle] + pipeline.anchor_lock_cap_tokens = 8 + pipeline._prefetch_prefix_ctx[req.cache_request_handle] = ( + [0, 1], + None, + None, + ) + self.assertEqual( + pipeline.try_lock_anchor(req.cache_request_handle, 0), ("locked", 2) + ) + anchor = pipeline.anchor_locks[req.cache_request_handle] + cache.tree_core.match_full_device_prefix.reset_mock() + for attempt in range(1, 4): + req.prefix_indices = torch.arange(2) + self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req)) + self.assertIsNone( + cache.init_load_back( + InitLoadBackParams( + best_match_node=None, + host_hit_length=req.host_hit_length, + req=req, + ) + ) + ) + self.assertIs( + pipeline.staged_prefetches[req.cache_request_handle], held + ) + self.assertIs( + pipeline.anchor_locks[req.cache_request_handle], anchor + ) + self.assertEqual(pipeline.anchor_locked_tokens_, 2) + self.assertEqual( + cache.storage_prefetch_retries.pop_ready([req], 0, 8), [] + ) + self.assertEqual( + cache.tree_core.match_full_device_prefix.call_count, attempt + ) + self.assertEqual( + cache.tree_core.collect_full_device_indices.call_count, attempt + ) + cache.cache_controller.mem_pool_host.free.assert_not_called() + self.assertEqual(cache.cache_controller.prefetch_tokens_occupied, 6) + cache.tree_core.dec_full_pin.assert_not_called() + pipeline.release_staged_hold(req.cache_request_handle) + cache.tree_core.dec_full_pin.assert_called_once_with(anchor.node_id) + self.assertEqual(pipeline.anchor_locked_tokens_, 0) + + def test_next_pass_replans_growth_and_refetches_anchor_loss_once(self): + cache, pipeline, req = _staged_fixture() + self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req)) + self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (6, 4)) + + # Tree changes occur while queued, before the next preparation pass. + cache.tree_core.match_full_device_prefix.return_value = (6, 1, 6) + req.prefix_indices = torch.arange(2) + self.assertTrue(cache.buffer_pipeline.prepare_staged_prefetch(req)) + self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (2, 4)) + self.assertTrue(pipeline.has_staged(req.cache_request_handle)) + self.assertEqual(cache.storage_prefetch_retries.pop_ready([req], 0, 8), []) + + cache.tree_core.match_full_device_prefix.return_value = (0, 0, 0) + req.prefix_indices = torch.arange(0) + self.assertFalse(cache.buffer_pipeline.prepare_staged_prefetch(req)) + self.assertFalse(pipeline.has_staged(req.cache_request_handle)) + self.assertEqual( + cache.storage_prefetch_retries.pop_ready([req], 0, 8), [(req, 8)] + ) + self.assertEqual(cache.storage_prefetch_retries.pop_ready([req], 0, 8), []) + self.assertEqual(cache.cache_controller.prefetch_tokens_occupied, 0) + cache.cache_controller.load.assert_not_called() + + def test_retry_budget_bounds_reissues_and_paces_capacity_misses(self): + """Past --hicache-storage-prefetch-retry-max-attempts a request stops + re-issuing; a rate-limited cache-mode query is paced, not re-issued.""" + retries = StoragePrefetchRetries() + head = SimpleNamespace(rid="head", storage_prefetch_retry_attempts=0) + req = SimpleNamespace(rid="r", storage_prefetch_retry_attempts=8) + retries.refetch(req.rid, 8) + self.assertEqual(retries.pop_ready([head, req], 0, 8), []) + req.storage_prefetch_retry_attempts = 7 + retries.refetch(req.rid, 8) + self.assertEqual(retries.pop_ready([head, req], 0, 8), [(req, 8)]) + # A paced retry yields to the queue head; an immediate one is re-issued. + retries.poll_miss(head.rid) + retries.refetch(req.rid, 8) + self.assertEqual(retries.pop_ready([head, req], 2, 8), [(req, 8)]) + + cache, _, req = _staged_fixture() + cache.host_memory_mode = "cache" + cache.buffer_pipeline = None + cache.cache_controller.prefetch_rate_limited = lambda: True + tokens = array("q", range(10)) + cache.prefetch_from_storage( + req.cache_request_handle, + 0, + tokens[2:], + matched_prefix_tokens=tokens[:2], + storage_hit_end=10, + ) + self.assertEqual(cache.ongoing_prefetch, {}) + retries = cache.storage_prefetch_retries + self.assertEqual(retries.pop_ready([head, req], 2, 8), []) + self.assertEqual(retries.pop_ready([head, req], 2, 8), []) + self.assertEqual(retries.pop_ready([head, req], 2, 8), [(req, 10)]) + + def test_staged_hold_drops_after_bounded_admission_deferrals(self): + """A hold that cannot be materialized after max_staged_admission_defers + passes is released, and the request re-plans without a new L3 query.""" + cache, pipeline, req = _staged_fixture() + cache.cache_controller.load.return_value = None + pipeline.max_staged_admission_defers = 3 + params = lambda: InitLoadBackParams(None, req.host_hit_length, req=req) + for attempt in range(1, 4): + req.prefix_indices = torch.arange(2) + self.assertTrue(pipeline.prepare_staged_prefetch(req)) + self.assertIsNone(cache.init_load_back(params())) + self.assertEqual(pipeline.has_staged(req.cache_request_handle), attempt < 3) + self.assertEqual(cache.cache_controller.prefetch_tokens_occupied, 0) + self.assertEqual(cache.storage_prefetch_retries.pop_ready([req], 0, 8), []) + self.assertTrue(pipeline.prepare_staged_prefetch(req)) + self.assertEqual((req.staged_prefetch_plan, req.storage_hit_length), (None, 0)) + + def test_controller_terminated_query_counts_as_an_l3_miss(self): + """A query the controller terminated (store miss or short hit) must feed + the L3-miss counters, or a store that lost pages reads as zero misses.""" + cache = _hit_drain_fixture() + _terminated_query(cache, "miss", hit_tokens=0) + _terminated_query(cache, "short", hit_tokens=2) + cache._drain_storage_control_queues_impl( + n_storage_hit=2, + n_ack_prefetch=0, + n_backup=0, + n_release=0, + extra_release_counts={}, + log_metrics=False, + ) + stats = cache._prefetch_outcome_stats + self.assertEqual( + ( + stats["revoked_full_miss"], + stats["revoked_insufficient"], + stats["l3_miss_tokens"], + ), + (1, 1, 14), + ) + self.assertEqual(cache.revoke_pending_prefetch.call_count, 2) + head = SimpleNamespace(rid="head", storage_prefetch_retry_attempts=0) + req = SimpleNamespace(rid="miss", storage_prefetch_retry_attempts=0) + retries = cache.storage_prefetch_retries + self.assertEqual(retries.pop_ready([head, req], 1, 8), []) + self.assertEqual(retries.pop_ready([head, req], 1, 8), [(req, None)]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index e6cc0fbfbcf0..1845436a15a6 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -26,6 +26,7 @@ ) from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req, ReqKvInfo +from sglang.srt.managers.schedule_policy import PrefillAdder from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( @@ -54,6 +55,7 @@ ReqToTokenPool, ) 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, @@ -888,6 +890,7 @@ def prefetch( last_hash=None, prefix_keys=None, extra_pools=None, + assume_stored=False, ): self.prefetch_args = ( request_id, @@ -984,7 +987,9 @@ def test_buffer_anchor_rematch_preserves_bigram_boundary(self): pipeline._cache = cache lock_ref = _device_lock_ref(cache, match.last_device_node, ComponentType.FULL) - self.assertEqual(pipeline.try_lock_anchor(req_id), "locked") + self.assertEqual( + pipeline.try_lock_anchor(req_id, remaining_full_tokens=1), ("locked", 4) + ) self.assertEqual(pipeline.anchor_locks[req_id].node_id, match.last_device_node) self.assertEqual( _device_lock_ref(cache, match.last_device_node, ComponentType.FULL), @@ -3258,11 +3263,14 @@ def _consume_staged_prefetch( extra_key=None, cache_salt=None, last_node=None, + on_dispatched=None, + return_req: bool = False, ): """Simulate the PrefillAdder consuming a staged prefetch at admission: init_load_back (buffer dispatch: device alloc + queued H2D), the batch start_loading flush, then pump until the ack commit lands. Returns - the spliced device indices (empty on degrade).""" + the spliced device indices (empty on degrade). ``on_dispatched`` runs + between the two, where the H2D is still in flight.""" from sglang.srt.mem_cache.base_prefix_cache import InitLoadBackParams f = cache.buffer_pipeline.staged_prefetches[req_id] @@ -3286,11 +3294,54 @@ def _consume_staged_prefetch( device=cache.tree_core.empty_match_result.device_indices.device, ) req.last_node = cache.root_node_handle() if last_node is None else last_node - new_indices, _last_node = cache.init_load_back( - InitLoadBackParams( - best_match_node=None, host_hit_length=f.num_tokens, req=req + joint = cache.match_prefix( + MatchPrefixParams( + key=RadixKey( + array("q", f.key_tokens), extra_key=extra_key, cache_salt=cache_salt + ) ) ) + req.prefix_indices = joint.device_indices + joint_covers_span = len(joint.device_indices) >= len(f.key_tokens) + adder = PrefillAdder.__new__(PrefillAdder) + adder.tree_cache = cache + with ( + mock.patch.object( + cache.tree_core, + "match_full_device_prefix", + wraps=cache.tree_core.match_full_device_prefix, + ) as full_match, + mock.patch.object( + cache.tree_core, + "collect_full_device_indices", + wraps=cache.tree_core.collect_full_device_indices, + ) as collect_indices, + ): + self.assertTrue( + cache.buffer_pipeline.prepare_staged_prefetch(req), + "production admission FULL rematch unexpectedly deferred", + ) + with adder._lock_node(req.last_node): + loaded = cache.init_load_back( + InitLoadBackParams( + best_match_node=None, + host_hit_length=req.host_hit_length, + req=req, + ) + ) + if joint_covers_span: + # A joint match past the staged span is kept without a rematch. + full_match.assert_not_called() + else: + full_match.assert_called_once() + collect_indices.assert_called_once() + new_indices = ( + cache.tree_core.empty_match_result.device_indices + if loaded is None + else loaded[0] + ) + if on_dispatched is not None: + on_dispatched() # Batch formation flushes the queued load into the batch's producer. cache.ready_to_load_host_cache() self._pump_hicache_until( @@ -3299,7 +3350,7 @@ def _consume_staged_prefetch( "staged-prefetch consumption did not commit", timeout=timeout, ) - return new_indices + return (new_indices, req) if return_req else new_indices def _all_page_hashes(self, cache, node): hashes = [] @@ -3952,6 +4003,7 @@ def _avail(): dtype=torch.int64, device=cons.tree_core.empty_match_result.device_indices.device, ) + self.assertTrue(cons.buffer_pipeline.prepare_staged_prefetch(req)) spliced, _last = cons.init_load_back( InitLoadBackParams( best_match_node=None, host_hit_length=held.num_tokens, req=req @@ -4110,10 +4162,7 @@ def test_buffer_only_cache_salt_uses_the_request_namespace(self): extra_key=extra_key, cache_salt=cache_salt, ) - self.assertIn(anchored_req, cons2.buffer_pipeline.anchor_locks) - self.assertEqual( - _device_lock_ref(cons2, anchor, ComponentType.FULL), lock_ref + 1 - ) + self.assertNotIn(anchored_req, cons2.buffer_pipeline.anchor_locks) self._pump_hicache_until( cons2, lambda: ( @@ -4122,6 +4171,10 @@ def test_buffer_only_cache_salt_uses_the_request_namespace(self): ), "salted mid-tree prefetch did not stage", ) + self.assertIn(anchored_req, cons2.buffer_pipeline.anchor_locks) + self.assertEqual( + _device_lock_ref(cons2, anchor, ComponentType.FULL), lock_ref + 1 + ) self._consume_staged_prefetch( cons2, anchored_req, @@ -4169,8 +4222,16 @@ def test_buffer_only_storage_prefetch_miss_marker_and_retry(self): ) self.assertFalse(cons.buffer_pipeline.has_staged(req_id)) self.assertEqual(stats["revoked_full_miss"], 1) - self.assertTrue(cons.pop_storage_prefetch_miss(req_id)) - self.assertFalse(cons.pop_storage_prefetch_miss(req_id)) # served once + waiting = [ + SimpleNamespace(rid="head", storage_prefetch_retry_attempts=0), + SimpleNamespace(rid=req_id.rid, storage_prefetch_retry_attempts=0), + ] + self.assertEqual(cons.storage_prefetch_retries.pop_ready(waiting, 1, 8), []) + self.assertEqual( + cons.storage_prefetch_retries.pop_ready(waiting, 1, 8), + [(waiting[1], None)], + ) + self.assertEqual(cons.storage_prefetch_retries.pop_ready(waiting, 1, 8), []) # Producer commits the span; the re-issued check (paced retry) hits # and stages what the first, too-early query could not see. @@ -4186,7 +4247,7 @@ def test_buffer_only_storage_prefetch_miss_marker_and_retry(self): ), "retried prefetch did not stage", ) - self.assertFalse(cons.pop_storage_prefetch_miss(req_id)) + self.assertNotIn(req_id.rid, cons.storage_prefetch_retries._pending) self.assertEqual(cons.pop_prefetch_loaded_tokens(req_id), len(seq)) # Unserved markers must not leak: abort cleanup ... @@ -4204,7 +4265,7 @@ def test_buffer_only_storage_prefetch_miss_marker_and_retry(self): "aborted-rid miss did not resolve", ) cons.finish(aborted_rid, CacheRequestOutcome.ABORT) - self.assertFalse(cons.pop_storage_prefetch_miss(aborted_rid)) + self.assertNotIn(aborted_rid.rid, cons.storage_prefetch_retries._pending) # A fully-device-matched (empty-suffix) decline also arms the retry: # the device match can evict while the request waits in the queue. @@ -4215,9 +4276,7 @@ def test_buffer_only_storage_prefetch_miss_marker_and_retry(self): None, None, ) - self.assertTrue( - cons.pop_storage_prefetch_miss(CacheRequestHandle("fully-matched", 0)) - ) + self.assertIn("fully-matched", cons.storage_prefetch_retries._pending) cons.sanity_check() def test_buffer_only_anchor_lock_cap_clamped_by_context_headroom(self): @@ -4298,7 +4357,14 @@ def test_buffer_load_back_swa_window_charged_at_admission(self): self.assertIsNotNone(cons_alloc.swa_attn_allocator.alloc(hold // ps * ps)) # The adder's SWA gate for this request (_swa_budget_for_req). - surfaced_swa_hit = cons.staged_prefetch_swa_tokens(req_id) + req = SimpleNamespace( + rid=req_id.rid, + cache_request_handle=req_id, + prefix_indices=cons.tree_core.empty_match_result.device_indices, + kv=SimpleNamespace(cache_protected_len=0), + ) + self.assertTrue(cons.buffer_pipeline.prepare_staged_prefetch(req)) + surfaced_swa_hit = req.swa_host_hit_length reserved = ( max(extend_need - window, 0) + min(extend_need + max_new, window) @@ -4320,6 +4386,7 @@ def test_buffer_load_back_swa_window_charged_at_admission(self): dtype=torch.int64, device=cons.tree_core.empty_match_result.device_indices.device, ) + self.assertTrue(cons.buffer_pipeline.prepare_staged_prefetch(req)) spliced, last_node = cons.init_load_back( InitLoadBackParams( best_match_node=None, host_hit_length=held.num_tokens, req=req @@ -4420,46 +4487,84 @@ def test_buffer_only_load_back_drops_on_sibling_published_span(self): k, v = self._snapshot_full_kv(cons_alloc, m.device_indices) self.assertTrue(torch.equal(k, sib_kv[0])) self.assertTrue(torch.equal(v, sib_kv[1])) - cons.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens.assert_not_called() + cons.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens.assert_called_once_with( + len(seq), "device_covered" + ) cons.sanity_check() - def test_buffer_only_load_back_drops_on_full_overlap_masked_by_swa_tombstone( - self, - ): - """Queued-UAF regression: live FULL under an SWA tombstone is invisible - to the unified match but still dedup-freed by insert; only the - full_kv_hit_length pre-check can drop the hold.""" + def test_buffer_only_load_back_uses_full_behind_swa_tombstone(self): + """FULL-only rematch keeps resident FULL and loads the complete SWA window.""" self._skip_unsupported_hicache_test() + self._skip_swa_window_repair_on_rust() if not self.cfg.has_swa: self.skipTest("masked overlap requires an SWA component") storage_dir = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) - seq = self._buffer_swa_seq() self._produce_buffer_l3(storage_dir, seq) + with self.subTest(window="all-tombstone"): + self._assert_masked_full_is_reused(storage_dir, seq, live_tail=0) + sw_pages = ( + self.cfg.sliding_window_size + self.cfg.page_size - 1 + ) // self.cfg.page_size + if sw_pages >= 2: + with self.subTest(window="partly-resident"): + self._assert_masked_full_is_reused( + storage_dir, seq, live_tail=self.cfg.page_size + ) + + def _assert_masked_full_is_reused(self, storage_dir, seq, live_tail: int): cons, cons_alloc, cons_rtp = build_fixture(self.cfg) self._init_buffer_hicache(cons, storage_dir) + cons.enable_storage_metrics = True + cons.storage_metrics_collector = mock.Mock() - # Masked state: nodes born with live FULL under SWA tombstones - # (sibling insert whose SWA ring had slid past the span). + # Masked state: nodes born with live FULL under SWA tombstones (a + # sibling insert whose SWA ring had slid past the span). live_tail is + # how far short of the span end the ring stopped. + live_from = len(seq) - live_tail value = self._alloc(cons_alloc, len(seq)) cons.insert( InsertParams( key=RadixKey(array("q", seq)), value=value[: len(seq)], - swa_evicted_seqlen=len(seq), + swa_evicted_seqlen=live_from, ) ) masked = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) self.assertEqual(len(masked.device_indices), 0, "unified match not masked") self.assertEqual(masked.full_kv_hit_length, len(seq), "live FULL not resident") + masked_full = value[: len(seq)].clone() + # Model the real SWA-eviction path: tombstoned positions have their + # allocator mapping cleared, while the live trailing portion remains. + cons_alloc.free_swa(masked_full[:live_from]) + self._fill_full_kv(cons_alloc, masked_full, marker=7) + masked_kv = self._snapshot_full_kv(cons_alloc, masked_full) + tail_full = masked_full[live_from:] + tail_swa = cons_alloc.full_to_swa_index_mapping[tail_full].clone() + self.assertTrue(bool((tail_swa > 0).all()), "tail SWA not left resident") avail0 = self._host_avail_sizes(cons) - req_id = CacheRequestHandle("masked-overlap", 0) + req_id = CacheRequestHandle(f"masked-overlap-{live_tail}", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) + # A parked SWA-only retry has no FULL suffix, but must re-pin its prefix. + with mock.patch.object(cons.swa_kv_pool_host, "alloc", return_value=None): + self._pump_hicache_until( + cons, + lambda: bool(cons.buffer_pipeline.pending_hit_allocs), + "SWA-only hit was not parked on the staging shortfall", + ) + self.assertEqual( + cons.ongoing_prefetch[req_id].operation.storage_hit_count, 0 + ) + self.assertEqual( + cons.buffer_pipeline._prefetch_prefix_ctx[req_id][0], list(seq) + ) + self.assertEqual(cons.buffer_pipeline.anchor_locks, {}) + self.assertEqual(cons.buffer_pipeline.anchor_locked_tokens_, 0) self._pump_hicache_until( cons, lambda: ( @@ -4468,21 +4573,210 @@ def test_buffer_only_load_back_drops_on_full_overlap_masked_by_swa_tombstone( ), "prefetch did not stage", ) + self.assertIn(req_id, cons.buffer_pipeline.anchor_locks) cons.pop_prefetch_loaded_tokens(req_id) - dev_avail0 = cons.token_to_kv_pool_allocator.available_size() + staged = cons.buffer_pipeline.staged_prefetches[req_id] + self.assertEqual(staged.matched_len, len(seq)) + self.assertEqual(staged.num_tokens, 0, "resident FULL must not be re-fetched") + window_tokens = sum( + len(t.host_indices) + for t in staged.aux_xfers + if t.name == PoolName.SWA and t.host_indices is not None + ) + self.assertGreater(window_tokens, live_tail, "window not partly tombstoned") + full_avail0 = cons_alloc.full_available_size() + mapping_before = cons_alloc.full_to_swa_index_mapping[masked_full].clone() + dispatched = {} + + def capture_swa_mapping(): + load = next(iter(cons.buffer_pipeline.ongoing_buffer_load_back.values())) + swa_xfer = next(t for t in load.aux_xfers if t.name == PoolName.SWA) + dispatched["device"] = swa_xfer.device_indices.clone() + full_window = masked_full[-len(swa_xfer.device_indices) :] + old_window = mapping_before[-len(swa_xfer.device_indices) :] + expected = torch.where(old_window > 0, old_window, swa_xfer.device_indices) + self.assertTrue( + torch.equal( + cons_alloc.full_to_swa_index_mapping[full_window], + expected, + ), + "admitted request did not fill the missing SWA mappings", + ) - # Every unified-length guard passes at 0 == 0; only the Full-only - # pre-check drops. - spliced = self._consume_staged_prefetch(cons, req_id, prefix_len=0) + spliced, req = self._consume_staged_prefetch( + cons, + req_id, + prefix_len=0, + on_dispatched=capture_swa_mapping, + return_req=True, + ) self.assertEqual(int(spliced.numel()), 0) - - self.assertEqual(cons.token_to_kv_pool_allocator.available_size(), dev_avail0) + self.assertEqual(len(req.prefix_indices), len(seq), "FULL rematch was lost") + self.assertEqual(len(dispatched["device"]), window_tokens) + self.assertEqual(cons_alloc.full_available_size(), full_avail0) + k, v = self._snapshot_full_kv(cons_alloc, masked_full) + self.assertTrue(torch.equal(k, masked_kv[0])) + self.assertTrue(torch.equal(v, masked_kv[1])) + self.assertTrue( + torch.equal(cons_alloc.full_to_swa_index_mapping[tail_full], tail_swa), + "live tail changed while dropping the staged overlap", + ) self.assertEqual(cons.buffer_pipeline.ongoing_buffer_load_back, {}) + self.assertEqual(cons.buffer_pipeline.anchor_locks, {}) self.assertEqual(cons.cache_controller.prefetch_tokens_occupied, 0) self.assertEqual(self._host_avail_sizes(cons), avail0) - # The masked FULL is still intact (nothing dedup-freed it). after = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) self.assertEqual(after.full_kv_hit_length, len(seq)) + self.assertEqual(len(after.device_indices), len(seq)) + cons.storage_metrics_collector.log_storage_prefetch_unfulfilled_tokens.assert_called_once_with( + len(seq), "device_covered" + ) + cons.sanity_check() + + def test_buffer_only_load_back_reuses_partial_masked_full(self): + """FULL-only rematch reuses a resident head and fetches only its tail.""" + self._skip_unsupported_hicache_test() + self._skip_swa_window_repair_on_rust() + if not self.cfg.has_swa: + self.skipTest("masked overlap requires an SWA component") + page_size = self.cfg.page_size + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + + seq = self._buffer_swa_seq() + _, expected_kv = self._produce_buffer_l3(storage_dir, seq, marker=3) + split_at = len(seq) - page_size + + cons, cons_alloc, cons_rtp = build_fixture(self.cfg) + self._init_buffer_hicache(cons, storage_dir) + + # Shares seq[:split_at] then diverges, so the FULL run ends mid-node. + sibling = seq[:split_at] + self._make_seq(9000, 1) + value = self._alloc(cons_alloc, len(sibling)) + cons.insert( + InsertParams( + key=RadixKey(array("q", sibling)), + value=value[: len(sibling)], + swa_evicted_seqlen=len(sibling), + ) + ) + live = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertEqual(len(live.device_indices), 0, "unified match not masked") + self.assertEqual(live.full_kv_hit_length, split_at, "FULL run not partial") + head_full = value[:split_at].clone() + cons_alloc.free_swa(head_full) + self._fill_full_kv(cons_alloc, head_full, marker=7) + head_kv = self._snapshot_full_kv(cons_alloc, head_full) + + avail0 = self._host_avail_sizes(cons) + req_id = CacheRequestHandle("masked-partial", 0) + cons.prefetch_from_storage( + req_id, cons.root_node_handle(), array("q", seq), None, None + ) + self._pump_hicache_until( + cons, + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), + "prefetch did not stage", + ) + cons.pop_prefetch_loaded_tokens(req_id) + staged = cons.buffer_pipeline.staged_prefetches[req_id] + self.assertEqual(staged.matched_len, split_at) + self.assertEqual(staged.num_tokens, page_size) + full_avail0 = cons_alloc.full_available_size() + + spliced, req = self._consume_staged_prefetch( + cons, req_id, prefix_len=0, return_req=True + ) + self.assertEqual(len(req.prefix_indices), split_at) + self.assertEqual(int(spliced.numel()), page_size) + self.assertEqual(cons_alloc.full_available_size(), full_avail0 - page_size) + k, v = self._snapshot_full_kv(cons_alloc, head_full) + self.assertTrue(torch.equal(k, head_kv[0])) + self.assertTrue(torch.equal(v, head_kv[1])) + tail_k, tail_v = self._snapshot_full_kv(cons_alloc, spliced) + self.assertTrue(torch.equal(tail_k, expected_kv[0][split_at:])) + self.assertTrue(torch.equal(tail_v, expected_kv[1][split_at:])) + + self.assertEqual(cons.buffer_pipeline.ongoing_buffer_load_back, {}) + self.assertEqual(cons.cache_controller.prefetch_tokens_occupied, 0) + self.assertEqual(self._host_avail_sizes(cons), avail0) + full_len, full_node, _ = cons.tree_core.match_full_device_prefix( + RadixKey(array("q", seq)) + ) + self.assertEqual(full_len, len(seq)) + full_indices = cons.tree_core.collect_full_device_indices( + full_node, cons.root_node_handle() + ) + self.assertTrue(torch.equal(full_indices[:split_at], head_full)) + self.assertTrue(torch.equal(full_indices[split_at:], spliced)) + cons.sanity_check() + + def test_buffer_only_masked_head_is_not_evicted_for_tail_load(self): + """Loading past a masked FULL head preserves its existing ownership.""" + self._skip_unsupported_hicache_test() + if not self.cfg.has_swa: + self.skipTest("masked overlap requires an SWA component") + page_size = self.cfg.page_size + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + + seq = self._buffer_swa_seq() + self._produce_buffer_l3(storage_dir, seq) + split_at = len(seq) - page_size + + cons, cons_alloc, cons_rtp = build_fixture(self.cfg) + self._init_buffer_hicache(cons, storage_dir) + + # The tree holds the head alone, SWA-tombstoned. It is a leaf, so it is + # the victim the tail's evict-before-alloc would pick. + value = self._alloc(cons_alloc, split_at) + cons.insert( + InsertParams( + key=RadixKey(array("q", seq[:split_at])), + value=value[:split_at], + swa_evicted_seqlen=split_at, + ) + ) + + req_id = CacheRequestHandle("adopt-under-evict", 0) + cons.prefetch_from_storage( + req_id, cons.root_node_handle(), array("q", seq), None, None + ) + self._pump_hicache_until( + cons, + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), + "prefetch did not stage", + ) + cons.pop_prefetch_loaded_tokens(req_id) + + # Admission's temporary request lock must protect the head on its own. + cons.buffer_pipeline.release_anchor_lock(req_id) + self.assertEqual(cons.buffer_pipeline.anchor_locks, {}) + # Starve FULL only; SWA stays free so the window still loads. + ballast = cons_alloc.full_attn_allocator.alloc(cons_alloc.full_available_size()) + self.assertIsNotNone(ballast) + self.addCleanup(cons_alloc.full_attn_allocator.free, ballast) + + def full_accounted(): + return ( + cons_alloc.full_available_size() + + cons.full_evictable_size() + + cons.full_protected_size() + ) + + accounted0 = full_accounted() + self._consume_staged_prefetch(cons, req_id, prefix_len=0) + + self.assertEqual(accounted0, full_accounted(), "a FULL page is free and owned") + # The resident FULL head was reused rather than evicted for the tail. + live = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertEqual(live.full_kv_hit_length, split_at) cons.sanity_check() def test_buffer_only_load_back_fail_stops_on_post_check_overlap(self): @@ -4534,6 +4828,7 @@ def adversarial_load(*args, **kwargs): device=cons.tree_core.empty_match_result.device_indices.device, ) req.last_node = cons.root_node_handle() + self.assertTrue(cons.buffer_pipeline.prepare_staged_prefetch(req)) with mock.patch.object(cons.cache_controller, "load", adversarial_load): with self.assertRaisesRegex(RuntimeError, "ownership violation"): cons.init_load_back( @@ -4549,8 +4844,12 @@ def test_buffer_only_load_back_trims_head_published_by_sibling(self): wasted). Consumption must instead splice the tail beyond the live prefix: sibling head slots stay untouched (add-only insert), the tail carries the producer's bytes, and the ack frees the entire - bounce including the trimmed head.""" + bounce including the trimmed head. + + A window-sized fetch still passes the global threshold, but after + sibling growth its FULL splice is shorter than its SWA transfer.""" self._skip_unsupported_hicache_test() + self._skip_swa_window_repair_on_rust() # Buffer-mode plan/commit logic is layout-independent, and each # hicache fixture retains its pools for the whole file run. Pin to # one config so the matrix does not exhaust a small CI GPU. @@ -4559,7 +4858,17 @@ def test_buffer_only_load_back_trims_head_published_by_sibling(self): storage_dir = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) - seq = self._buffer_swa_seq() + with self.subTest(span="over-window"): + self._assert_sibling_head_is_trimmed( + storage_dir, self._buffer_swa_seq(), window_sized=False + ) + with self.subTest(span="exact-window"): + seq = self._make_seq( + 5000, self.cfg.sliding_window_size // self.cfg.page_size + ) + self._assert_sibling_head_is_trimmed(storage_dir, seq, window_sized=True) + + def _assert_sibling_head_is_trimmed(self, storage_dir, seq, window_sized: bool): _, (expected_k, expected_v) = self._produce_buffer_l3( storage_dir, seq, marker=9 ) @@ -4568,7 +4877,7 @@ def test_buffer_only_load_back_trims_head_published_by_sibling(self): self._init_buffer_hicache(cons, storage_dir) avail0 = self._host_avail_sizes(cons) - req_id = CacheRequestHandle("growth-trim", 0) + req_id = CacheRequestHandle(f"growth-trim-{len(seq)}", 0) cons.prefetch_from_storage( req_id, cons.root_node_handle(), array("q", seq), None, None ) @@ -4590,10 +4899,28 @@ def test_buffer_only_load_back_trims_head_published_by_sibling(self): self._fill_full_kv(cons_alloc, sib.device_indices, marker=3) head_k, head_v = self._snapshot_full_kv(cons_alloc, sib.device_indices) + if window_sized: + # FULL shrank below a window; the complete SWA staging must survive. + staged = cons.buffer_pipeline.staged_prefetches[req_id] + aux_max = max( + ( + t.host_indices.numel() + for t in staged.aux_xfers + if t.host_indices is not None + ), + default=0, + ) + self.assertGreater(aux_max, len(seq) - len(head)) + # The surfaced host hit is the splice-able tail, not the full span. - kv_tokens, swa_tokens = cons.plan_staged_splice(req_id, len(head)) - self.assertEqual(kv_tokens, len(seq) - len(head)) - self.assertEqual(swa_tokens, cons.staged_prefetch_swa_tokens(req_id)) + req = SimpleNamespace( + rid=req_id.rid, + cache_request_handle=req_id, + prefix_indices=sib.device_indices, + kv=SimpleNamespace(cache_protected_len=len(sib.device_indices)), + ) + self.assertTrue(cons.buffer_pipeline.prepare_staged_prefetch(req)) + self.assertEqual(req.host_hit_length, len(seq) - len(head)) self.assertTrue(cons.buffer_pipeline.has_staged(req_id)) spliced = self._consume_staged_prefetch( @@ -4652,13 +4979,22 @@ def test_buffer_only_plan_frees_covered_hold(self): cons.pop_prefetch_loaded_tokens(req_id) self._insert(cons, cons_alloc, cons_rtp, seq) - self.assertEqual(cons.plan_staged_splice(req_id, len(seq)), (0, 0)) + live = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + req = mock.Mock( + rid=req_id.rid, + cache_request_handle=req_id, + prefix_indices=live.device_indices, + last_node=live.last_device_node, + ) + self.assertTrue(cons.buffer_pipeline.prepare_staged_prefetch(req)) + self.assertEqual((req.host_hit_length, req.swa_host_hit_length), (0, 0)) self.assertFalse(cons.buffer_pipeline.has_staged(req_id)) self.assertEqual(cons.buffer_pipeline.anchor_locks, {}) self.assertEqual(cons.cache_controller.prefetch_tokens_occupied, 0) self.assertEqual(self._host_avail_sizes(cons), avail0) # Idempotent once freed. - self.assertEqual(cons.plan_staged_splice(req_id, len(seq)), (0, 0)) + self.assertTrue(cons.buffer_pipeline.prepare_staged_prefetch(req)) + self.assertIsNone(req.staged_prefetch_plan) cons.sanity_check() def test_buffer_only_hit_commit_cancels_device_covered_fetch(self): @@ -4705,7 +5041,7 @@ def test_buffer_only_hit_commit_cancels_device_covered_fetch(self): self.assertFalse(cons.buffer_pipeline.has_staged(req_id)) self.assertTrue(cons.check_prefetch_progress(req_id)) self.assertEqual(cons.cache_controller.prefetch_tokens_occupied, 0) - self.assertFalse(cons.pop_storage_prefetch_miss(req_id)) + self.assertNotIn(req_id.rid, cons.storage_prefetch_retries._pending) # Aux staging released during the drain lands on queues sized before # it; a second drain flushes them. cons.drain_storage_control_queues() @@ -4713,11 +5049,7 @@ def test_buffer_only_hit_commit_cancels_device_covered_fetch(self): cons.sanity_check() def test_buffer_only_swa_window_semantics(self): - """SWA window handling across the three partial-window cases: - root-anchored sub-window sequence (the sequence IS its window), - mid-tree sub-window continuation (head = device ring state), and a - storage hit shorter than the requested window (shrunk, tail - released). Each was a zero-L3-reuse regression on Llama-4-Scout.""" + """Root and mid-tree requests and hits share a complete-window minimum.""" self._skip_unsupported_hicache_test() if not self.cfg.has_swa: self.skipTest("requires an SWA component") @@ -4728,11 +5060,16 @@ def test_buffer_only_swa_window_semantics(self): storage_dir = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) - # 1. Root-anchored sequence shorter than the window. - seq = self._make_seq(1, (window // self.cfg.page_size) - 1) + # 1. A root-anchored sub-window request is declined before querying. + seq = self._make_seq(1, sw_pages - 1) self._produce_buffer_l3(storage_dir, seq, marker=5) cons, _, _ = build_fixture(self.cfg) self._init_buffer_hicache(cons, storage_dir) + self.assertEqual(cons.prefetch_threshold, sw_pages * self.cfg.page_size) + self.assertEqual( + cons.cache_controller.prefetch_threshold, cons.prefetch_threshold + ) + avail = self._host_avail_sizes(cons) cons.prefetch_from_storage( CacheRequestHandle("short-req", 0), cons.root_node_handle(), @@ -4740,56 +5077,67 @@ def test_buffer_only_swa_window_semantics(self): None, None, ) - self._run_prefetch_to_completion(cons, CacheRequestHandle("short-req", 0)) - cons.drain_storage_control_queues() - mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - self.assertEqual(len(mc.device_indices), len(seq)) - self.assertIsNotNone( - _device_value(cons, mc.last_device_node, ComponentType.SWA) - ) + self.assertNotIn(CacheRequestHandle("short-req", 0), cons.ongoing_prefetch) + self.assertEqual(cons._prefetch_outcome_stats["declined_too_short"], 1) + self.assertEqual(cons._prefetch_outcome_stats["issued"], 0) + self.assertEqual(self._host_avail_sizes(cons), avail) cons.sanity_check() - # 2. Mid-tree continuation shorter than the window: the staged - # prefetch must carry an SWA transfer (not a KV-only degrade). + # 2. The same threshold applies mid-tree; an exact window is accepted. if sw_pages >= 2: seq_a = self._make_seq(1, max(2, sw_pages)) - seq_ab = seq_a + self._make_seq(900, sw_pages - 1) + seq_ab = seq_a + self._make_seq(900, sw_pages) self._produce_buffer_l3(storage_dir, seq_ab, marker=14) cons2, cons2_alloc, cons2_rtp = build_fixture(self.cfg) self._init_buffer_hicache(cons2, storage_dir) avail2 = self._host_avail_sizes(cons2) self._insert(cons2, cons2_alloc, cons2_rtp, seq_a) m = cons2.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a)))) + last_hash = cons2.get_last_hash_value(m.last_device_node) cons2.prefetch_from_storage( CacheRequestHandle("subwin-req", 0), m.last_device_node, + array("q", seq_ab[len(seq_a) : -self.cfg.page_size]), + last_hash, + None, + matched_prefix_tokens=list(seq_a), + ) + self.assertNotIn( + CacheRequestHandle("subwin-req", 0), cons2.ongoing_prefetch + ) + self.assertEqual(cons2._prefetch_outcome_stats["declined_too_short"], 1) + self.assertIn("subwin-req", cons2.storage_prefetch_retries._pending) + + cons2.prefetch_from_storage( + CacheRequestHandle("window-req", 0), + m.last_device_node, array("q", seq_ab[len(seq_a) :]), - cons2.get_last_hash_value(m.last_device_node), + last_hash, None, matched_prefix_tokens=list(seq_a), ) self._pump_hicache_until( cons2, lambda: ( - cons2.check_prefetch_progress(CacheRequestHandle("subwin-req", 0)) + cons2.check_prefetch_progress(CacheRequestHandle("window-req", 0)) and cons2.buffer_pipeline.has_staged( - CacheRequestHandle("subwin-req", 0) + CacheRequestHandle("window-req", 0) ) ), - "sub-window prefetch did not stage", + "whole-window mid-tree prefetch did not stage", ) self.assertTrue( any( t.name == PoolName.SWA for t in cons2.buffer_pipeline.staged_prefetches[ - CacheRequestHandle("subwin-req", 0) + CacheRequestHandle("window-req", 0) ].aux_xfers ), - "sub-window fetch degraded to KV-only", + "mid-tree fetch degraded to KV-only", ) spliced = self._consume_staged_prefetch( cons2, - CacheRequestHandle("subwin-req", 0), + CacheRequestHandle("window-req", 0), prefix_indices=m.device_indices, ) self.assertEqual(int(spliced.numel()), len(seq_ab) - len(seq_a)) @@ -4804,10 +5152,9 @@ def test_buffer_only_swa_window_semantics(self): self.assertEqual(self._host_avail_sizes(cons2), avail2) cons2.sanity_check() - # 3. Hit one page short of the requested window: the shrunk window - # is kept (its own trailing window) and the buffer tail released. - full = self._buffer_swa_seq() - stored = full[: window - self.cfg.page_size] + # 3. An accepted request with a sub-window L3 hit is also declined. + full = self._make_seq(6000, sw_pages + 1) + stored = full[: (sw_pages - 1) * self.cfg.page_size] self._produce_buffer_l3(storage_dir, stored, marker=6) cons3, _, _ = build_fixture(self.cfg) self._init_buffer_hicache(cons3, storage_dir) @@ -4819,16 +5166,25 @@ def test_buffer_only_swa_window_semantics(self): None, None, ) + operation = cons3.ongoing_prefetch[ + CacheRequestHandle("partial-req", 0) + ].operation self._run_prefetch_to_completion(cons3, CacheRequestHandle("partial-req", 0)) cons3.drain_storage_control_queues() + self.assertEqual(operation.storage_hit_count, len(stored)) self.assertEqual( len( cons3.match_prefix( MatchPrefixParams(key=RadixKey(array("q", stored))) ).device_indices ), - len(stored), - "partial hit lost its SWA window", + 0, + "sub-window hit should not be loaded", + ) + self.assertEqual(cons3._prefetch_outcome_stats["issued"], 1) + self.assertEqual(cons3._prefetch_outcome_stats["revoked_insufficient"], 1) + self.assertFalse( + cons3.buffer_pipeline.has_staged(CacheRequestHandle("partial-req", 0)) ) self.assertEqual(self._host_avail_sizes(cons3), avail3) cons3.sanity_check() @@ -5078,6 +5434,13 @@ def _skip_unsupported_hicache_test(self): self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks") return False + def _skip_swa_window_repair_on_rust(self): + # Buffer-mode consumption repairs SWA tombstones under the loaded + # window through swa_tombstone_ranges/attach_swa_window, which the + # Rust tree core does not implement yet. + if _selected_tree_core_test_backend() == "rust": + self.skipTest("buffer-mode SWA window repair is Python-core only") + def _simulate_backup(self, cache, node): """Simulate D->H backup over the whole root->node path (parent-first).""" for ancestor in self._path_chain(cache, node): @@ -5946,23 +6309,40 @@ def test_prepare_prefetch_swa(self): cache, _, _ = self._build_hicache_fixture() sw = cache.sliding_window_size swa = cache.components[ComponentType.SWA] - # zero-length prefetch -> does not participate, no alloc - prep = swa.prepare_prefetch(cache.root_node_handle(), prefetch_tokens=0) - self.assertFalse(prep.alloc_failed) - self.assertIsNone(prep.host_indices) + avail = cache.swa_kv_pool_host.available_size() + # zero-length prefetch -> does not participate + self.assertEqual( + swa.prepare_prefetch( + cache.root_node_handle(), prefetch_tokens=0 + ).staging_tokens, + 0, + ) # below a full window at the ROOT anchor -> the whole sequence is its # own trailing window (sub-window prompts stay reusable via storage) - prep = swa.prepare_prefetch(cache.root_node_handle(), prefetch_tokens=sw - 1) - self.assertEqual(int(prep.host_indices.numel()), sw - 1) - # a full window available -> participates, allocs one window of host pages - prep = swa.prepare_prefetch(cache.root_node_handle(), prefetch_tokens=sw) - self.assertEqual(int(prep.host_indices.numel()), sw) - # a non-participating component never allocs - prep = cache.components[ComponentType.FULL].prepare_prefetch( - cache.root_node_handle(), prefetch_tokens=sw - ) - self.assertFalse(prep.alloc_failed) - self.assertIsNone(prep.host_indices) + self.assertEqual( + swa.prepare_prefetch( + cache.root_node_handle(), prefetch_tokens=sw - 1 + ).staging_tokens, + sw - 1, + ) + # a full window available -> one window of host pages + self.assertEqual( + swa.prepare_prefetch( + cache.root_node_handle(), prefetch_tokens=sw + ).staging_tokens, + sw, + ) + # a non-participating component never stages + self.assertEqual( + cache.components[ComponentType.FULL] + .prepare_prefetch(cache.root_node_handle(), prefetch_tokens=sw) + .staging_tokens, + 0, + ) + # Sizing holds no host memory: the window is allocated at hit time, so + # a query in flight (hit or miss) never occupies the SWA host pool. + self.assertEqual(cache.swa_kv_pool_host.available_size(), avail) + self.assertEqual(int(swa.alloc_prefetch_staging(sw).numel()), sw) def test_prepare_prefetch_swa_pool_exhausted(self): if not self.cfg.has_swa or self.cfg.has_mamba or self.cfg.page_size != 1: @@ -5970,16 +6350,12 @@ def test_prepare_prefetch_swa_pool_exhausted(self): cache, _, _ = self._build_hicache_fixture() sw = cache.sliding_window_size swa = cache.components[ComponentType.SWA] - # pool can't satisfy even after evict -> participates but aborts (no buffer) + # pool can't satisfy even after evict -> no staging, the hit parks with ( mock.patch.object(cache.swa_kv_pool_host, "alloc", return_value=None), mock.patch.object(cache, "evict_host", autospec=True) as evict_host, ): - self.assertTrue( - swa.prepare_prefetch( - cache.root_node_handle(), prefetch_tokens=sw - ).alloc_failed - ) + self.assertIsNone(swa.alloc_prefetch_staging(sw)) # the retry must evict the SWA host pool, not the default (FULL) one evict_host.assert_called_once_with(sw, ComponentType.SWA) @@ -5987,11 +6363,66 @@ def test_prepare_prefetch_mamba(self): if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: self.skipTest("requires page_size=1 Full+Mamba") cache, _, _ = self._build_hicache_fixture() - # mamba always participates and allocs exactly one page - prep = cache.components[ComponentType.MAMBA].prepare_prefetch( - cache.root_node_handle(), prefetch_tokens=0 + # mamba always participates with exactly one state slot + mamba = cache.components[ComponentType.MAMBA] + self.assertEqual( + mamba.prepare_prefetch( + cache.root_node_handle(), prefetch_tokens=0 + ).staging_tokens, + 1, ) - self.assertEqual(int(prep.host_indices.numel()), 1) + self.assertEqual(int(mamba.alloc_prefetch_staging(1).numel()), 1) + + def test_buffer_only_hit_parks_on_swa_staging_shortfall(self): + """A hit whose SWA window cannot be staged parks whole and launches + once the pool frees. It is never forfeited, and it never launches with + the KV staged and the window not; a parked hit holds no host memory.""" + if not self.cfg.has_swa or self.cfg.has_mamba: + self.skipTest("SWA-only fixture required") + if self._skip_unsupported_hicache_test(): + return + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + seq = self._buffer_swa_seq() + self._produce_buffer_l3(storage_dir, seq) + cons, _, _ = build_fixture(self.cfg) + self._init_buffer_hicache(cons, storage_dir) + pool = cons.swa_kv_pool_host + avail = pool.available_size() + real_alloc = pool.alloc + # Both attempts of the first hit-time alloc fail (the second follows + # the host evict); the next drain's attempt succeeds. + shortfalls = [None, None] + + def alloc(num_tokens): + if shortfalls: + return shortfalls.pop() + return real_alloc(num_tokens) + + with mock.patch.object(pool, "alloc", side_effect=alloc): + cons.prefetch_from_storage( + CacheRequestHandle("r", 0), + cons.root_node_handle(), + array("q", seq), + None, + None, + ) + # The query holds no SWA staging. + self.assertEqual(pool.available_size(), avail) + self._pump_hicache_until( + cons, + lambda: bool(cons.buffer_pipeline.pending_hit_allocs), + "hit was not parked on the SWA staging shortfall", + ) + self.assertIn(CacheRequestHandle("r", 0), cons.ongoing_prefetch) + self.assertEqual(cons._prefetch_outcome_stats["declined_rate_limited"], 1) + # Parked, before the next drain retries it: nothing is held. + self.assertEqual(pool.available_size(), avail) + self._run_prefetch_to_completion(cons, CacheRequestHandle("r", 0)) + cons.drain_storage_control_queues() + m = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertEqual(len(m.device_indices), len(seq)) + cons.sanity_check() def test_hicache_swa_host_best_match_keeps_device_anchor(self): if not self.cfg.has_swa or self.cfg.has_mamba or self.cfg.page_size != 1: @@ -6704,6 +7135,64 @@ def test_load_back_success_publishes_fresh_mamba_slot(self): self._finish_pending_loads(cache) self._release_ongoing_load_back_locks(cache) + def test_admission_keeps_mamba_h2d_destination_on_full_overdelivery(self): + if ( + not self.cfg.has_mamba + or self.cfg.has_swa + or self.cfg.page_size != 1 + or self.cfg.enable_mamba_extra_buffer + or self.cfg.enable_int8_mamba_checkpoint + ): + self.skipTest("requires page_size=1 Full+Mamba without extra buffers") + cache, allocator, pool = build_fixture(replace(self.cfg, kv_size=20)) + self._init_hicache(cache, write_policy="write_back") + tokens = list(range(12)) + self._insert(cache, allocator, pool, tokens) + leaf = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + self._fill_mamba_state( + pool, _device_value(cache, leaf, ComponentType.MAMBA), marker=11 + ) + self._backup_node(cache, leaf) + self._set_aux_host_tombstone(cache, leaf, ComponentType.MAMBA) + self.assertEqual(cache.evict_host(12, ComponentType.FULL), 12) + + req = Req( + rid="mamba-overdelivery", + origin_input_text="", + origin_input_ids=array("q", range(16)), + sampling_params=SamplingParams(temperature=0, max_new_tokens=1), + ) + req.init_next_round_input(cache) + self.assertEqual((len(req.prefix_indices), req.host_hit_length), (0, 0)) + self.assertEqual(req.mamba_host_hit_length, 1) + adder = PrefillAdder( + page_size=1, + tree_cache=cache, + token_to_kv_pool_allocator=allocator, + running_batch=None, + new_token_ratio=1.0, + rem_input_tokens=100, + rem_chunk_tokens=None, + ) + adder.add_one_req(req, False, None) + + # The load pins 12 FULL slots, leaving 8 free. Rechecking the original + # 18-token demand would reject and scheduler cleanup would free the CoW + # destination even though the queued H2D still targets it. + self.assertEqual(adder.can_run_list, [req]) + self.assertEqual(len(req.prefix_indices), 12) + self.assertEqual(req.kv.cache_protected_len, 12) + slot = req.kv.mamba_pool_idx.unsqueeze(0) + other_slots = pool.mamba_allocator.alloc(pool.mamba_allocator.available_size()) + self.assertNotIn(slot.item(), other_slots.tolist()) + self._finish_pending_loads(cache) + temporal, _ = self._snapshot_mamba_state(pool, slot) + self.assertTrue(torch.all(temporal == 11)) + pool.mamba_allocator.free(other_slots) + cache.dec_lock_ref(req.last_node, req.lock_receipt) + def test_load_back_success_copies_mamba_state_into_request_slot(self): if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: self.skipTest("requires page_size=1 Full+Mamba") @@ -9105,7 +9594,7 @@ def test_prefetch_commit_applies_host_insert_actions_before_transfers(self): operation.request_id = "req" operation.completed_tokens = 8 cache.ongoing_prefetch = { - operation.handle: ( + operation.handle: _OngoingPrefetch( 7, list(range(8)), list(range(100, 108)), @@ -9596,9 +10085,9 @@ def test_enabling_storage_backfills_the_tree(self): class TestAnchorLockOutcomePolicy(CustomTestCase): """try_lock_anchor finds the anchor by re-matching the live tree (no carried node id to go stale): prefix intact -> lock the live node; - prefix shrunk -> anchor_lost so the caller cancels the storage IO - instead of gambling the read; cap_skip over budget (checked before the - match walk).""" + prefix shrunk -> anchor_lost so the caller re-plans; over the whole cap + -> cap_skip (launch unlocked, since no drain can ever fit it), merely + over budget -> cap_busy (park).""" _REQ = CacheRequestHandle("req-1", 0) _PREFIX = list(range(100, 100 + 8)) @@ -9615,27 +10104,29 @@ def _make_pipeline(self, cache, cap_tokens=10_000): pipeline._cache = cache return pipeline - def _make_cache(self, live_match_len): - from types import SimpleNamespace - + def _make_cache(self, live_match_len, pinned_len=None): cache = mock.MagicMock() cache.tree_core.is_eagle = False - cache.match_prefix.return_value = SimpleNamespace( - device_indices=list(range(live_match_len)), last_device_node=99 + cache.tree_core.match_full_device_prefix.return_value = ( + live_match_len, + 99, + live_match_len if pinned_len is None else pinned_len, ) return cache def test_intact_prefix_locks_live_node(self): cache = self._make_cache(live_match_len=len(self._PREFIX)) pipeline = self._make_pipeline(cache) - self.assertEqual(pipeline.try_lock_anchor(self._REQ), "locked") + self.assertEqual(pipeline.try_lock_anchor(self._REQ, 0)[0], "locked") self.assertEqual(pipeline.anchor_locks[self._REQ].node_id, 99) self.assertEqual(pipeline.anchor_locked_tokens_, len(self._PREFIX)) + cache.tree_core.inc_full_pin.assert_called_once_with(99) + cache.inc_lock_ref.assert_not_called() def test_shrunk_prefix_reports_anchor_lost(self): cache = self._make_cache(live_match_len=len(self._PREFIX) - 2) pipeline = self._make_pipeline(cache) - self.assertEqual(pipeline.try_lock_anchor(self._REQ), "anchor_lost") + self.assertEqual(pipeline.try_lock_anchor(self._REQ, 0)[0], "anchor_lost") self.assertEqual(pipeline.anchor_locks, {}) self.assertEqual(pipeline.anchor_locked_tokens_, 0) @@ -9643,7 +10134,7 @@ def test_storage_cleanup_releases_buffer_prefetch_anchor(self): cache = self._make_cache(live_match_len=len(self._PREFIX)) pipeline = self._make_pipeline(cache) cache.buffer_pipeline = pipeline - self.assertEqual(pipeline.try_lock_anchor(self._REQ), "locked") + self.assertEqual(pipeline.try_lock_anchor(self._REQ, 0)[0], "locked") host_indices = torch.arange(4) cache.ongoing_prefetch = { self._REQ: _OngoingPrefetch( @@ -9670,15 +10161,19 @@ def test_storage_cleanup_releases_buffer_prefetch_anchor(self): self.assertEqual(pipeline.anchor_locks, {}) self.assertEqual(pipeline.anchor_locked_tokens_, 0) self.assertNotIn(self._REQ, pipeline._prefetch_prefix_ctx) - cache.dec_lock_ref.assert_called_once_with( - 99, cache.inc_lock_ref.return_value.to_dec_params.return_value - ) + cache.tree_core.dec_full_pin.assert_called_once_with(99) + cache.dec_lock_ref.assert_not_called() cache.dec_host_lock_ref.assert_not_called() self.assertEqual(controller.prefetch_tokens_occupied, 8) def test_positive_hit_with_lost_anchor_is_reported_as_shrunk(self): cache = UnifiedRadixCache.__new__(UnifiedRadixCache) - cache._storage_prefetch_missed_rids = set() + cache.storage_prefetch_retries = StoragePrefetchRetries() + cache.ongoing_prefetch = { + self._REQ: SimpleNamespace( + operation=SimpleNamespace(storage_start=4, storage_hit_count=8) + ) + } cache._finish_storage_prefetch = mock.Mock() cache.revoke_pending_prefetch = mock.Mock() @@ -9687,30 +10182,48 @@ def test_positive_hit_with_lost_anchor_is_reported_as_shrunk(self): cache._finish_storage_prefetch.assert_called_once_with( self._REQ, fulfilled_tokens=0, reason="shrunk" ) - self.assertIn(self._REQ, cache._storage_prefetch_missed_rids) + # The eviction widened the span, so the request replans over it -- + # skipping the query, the prior hit having proved it stored, and the + # poll interval, which would only let admission run without it. + req = SimpleNamespace(rid=self._REQ.rid, storage_prefetch_retry_attempts=0) + self.assertEqual( + cache.storage_prefetch_retries.pop_ready([req], 0, 8), [(req, 12)] + ) cache.revoke_pending_prefetch.assert_called_once_with(self._REQ) - def test_over_cap_reports_cap_skip_before_matching(self): + def test_over_cap_reports_cap_skip_after_full_rematch(self): cache = self._make_cache(live_match_len=len(self._PREFIX)) pipeline = self._make_pipeline(cache, cap_tokens=len(self._PREFIX) - 1) - self.assertEqual(pipeline.try_lock_anchor(self._REQ), "cap_skip") + self.assertEqual(pipeline.try_lock_anchor(self._REQ, 0)[0], "cap_skip") self.assertEqual(pipeline.anchor_locks, {}) - cache.match_prefix.assert_not_called() + cache.tree_core.match_full_device_prefix.assert_called_once() def test_root_anchor_reports_no_anchor(self): cache = self._make_cache(live_match_len=0) + cache.ongoing_prefetch[self._REQ].prefetch_key = RadixKey(array("q", range(8))) pipeline = self._make_pipeline(cache) pipeline._prefetch_prefix_ctx[self._REQ] = ([], None, None) - self.assertEqual(pipeline.try_lock_anchor(self._REQ), "no_anchor") - cache.match_prefix.assert_not_called() + self.assertEqual(pipeline.try_lock_anchor(self._REQ, 8), ("no_anchor", 0)) + cache.tree_core.match_full_device_prefix.assert_called_once() - def test_already_locked_is_idempotent(self): - cache = self._make_cache(live_match_len=len(self._PREFIX)) + def test_anchor_lock_requires_release_before_retry(self): + cache = self._make_cache(live_match_len=8, pinned_len=12) pipeline = self._make_pipeline(cache) - self.assertEqual(pipeline.try_lock_anchor(self._REQ), "locked") - self.assertEqual(pipeline.try_lock_anchor(self._REQ), "locked") - self.assertEqual(pipeline.anchor_locked_tokens_, len(self._PREFIX)) - cache.match_prefix.assert_called_once() + self.assertEqual(pipeline.try_lock_anchor(self._REQ, 0), ("locked", 8)) + with self.assertRaisesRegex(AssertionError, "prefetch anchor already locked"): + pipeline.try_lock_anchor(self._REQ, 0) + self.assertEqual(pipeline.anchor_locked_tokens_, 12) + cache.tree_core.match_full_device_prefix.assert_called_once() + cache.tree_core.inc_full_pin.assert_called_once_with(99) + + pipeline.release_anchor_lock(self._REQ) + self.assertEqual(pipeline.anchor_locked_tokens_, 0) + cache.tree_core.dec_full_pin.assert_called_once_with(99) + cache.tree_core.match_full_device_prefix.return_value = (8, 100, 8) + self.assertEqual(pipeline.try_lock_anchor(self._REQ, 0), ("locked", 8)) + self.assertEqual(pipeline.anchor_locks[self._REQ].node_id, 100) + self.assertEqual(pipeline.anchor_locked_tokens_, 8) + self.assertEqual(cache.tree_core.match_full_device_prefix.call_count, 2) @unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA") diff --git a/test/registered/unit/observability/test_stat_loggers_di.py b/test/registered/unit/observability/test_stat_loggers_di.py index d7ac7258fcaf..667134ae0eb6 100644 --- a/test/registered/unit/observability/test_stat_loggers_di.py +++ b/test/registered/unit/observability/test_stat_loggers_di.py @@ -214,6 +214,7 @@ def test_storage_prefetch_lifecycle_metrics(self): collector.log_storage_prefetch_hit_tokens(21) collector.log_storage_prefetch_unfulfilled_tokens(4, "storage_transfer") + collector.log_storage_prefetch_deferred_tokens(7, "device_capacity") self.assertEqual( collector.storage_prefetch_hit_tokens_total.increments, [(labels, 21)] @@ -222,6 +223,10 @@ def test_storage_prefetch_lifecycle_metrics(self): collector.storage_prefetch_unfulfilled_tokens_total.increments, [({**labels, "reason": "storage_transfer"}, 4)], ) + self.assertEqual( + collector.storage_prefetch_deferred_tokens_total.increments, + [({**labels, "reason": "device_capacity"}, 7)], + ) if __name__ == "__main__": From 9aa0364d22de73e319ff9dec4ac9ac48030ae7d8 Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Mon, 14 Sep 2026 15:18:56 -0700 Subject: [PATCH 2/2] [HiCache] Align the Rust SWA prefetch test with staging and drop the per-fixture atexit pin - `test_swa_prefetch_commit_end_to_end` still built the SWA PREFETCH transfer from pre-allocated host indices; the build now carries the planned staging as placeholder keys (`staging_tokens`) and attaches the host buffer once the hit is known, and returns None when the pool takes no part. Assert the new contract on both sides. - `UnifiedRadixCache.init_hicache` registers `cache.shutdown` with atexit, which keeps every HiCache fixture's device and host pools alive until process exit. The shared cache suite builds hundreds of them, so the 1-GPU unit shard grew to the 32 GiB card limit and OOMed. Unregister the hook per fixture via addCleanup; peak device memory for the suite drops from ~33 GiB to ~7 GiB. --- .../mem_cache/test_rust_tree_core_integration.py | 15 ++++++++++++--- .../test_unified_radix_cache_unittest.py | 11 +++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py index 3a1d5f0edbcf..6e710ed370f0 100644 --- a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py +++ b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py @@ -1833,17 +1833,26 @@ def test_swa_prefetch_commit_end_to_end(): core.has_swa_host_pool = True anchor = core.match_prefix(MatchPrefixParams(key=_key([99]))).best_match_node - # The build wraps the host buffer with placeholder keys, trailing-pages policy. + # Without planned staging the SWA pool takes no part in the fetch. + assert ( + core.build_hicache_transfers( + ComponentType.SWA, anchor, CacheTransferPhase.PREFETCH + ) + is None + ) + + # The build carries the planned staging as placeholder keys, trailing-pages + # policy; the host buffer is attached once the hit is known. (xfer,) = core.build_hicache_transfers( ComponentType.SWA, anchor, CacheTransferPhase.PREFETCH, - host_indices=torch.tensor([30, 31], dtype=torch.int64), + staging_tokens=2, ) assert xfer.name == PoolName.SWA assert xfer.keys == ["__placeholder__", "__placeholder__"] assert xfer.hit_policy == PoolHitPolicy.TRAILING_PAGES - assert xfer.host_indices.tolist() == [30, 31] + assert xfer.host_indices is None # The prefetched suffix lands as one host node; its SWA host is a tombstone. insert_result = core.insert_host( diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 1845436a15a6..b2194679debf 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -1,5 +1,6 @@ """Unit tests for UnifiedRadixCache""" +import atexit import json import shutil import sys @@ -230,6 +231,13 @@ def _recede_session_coverage(self, session_id, leaf, fallback) -> None: pass +def _drop_hicache_atexit_pin(cache): + """init_hicache registers ``cache.shutdown`` with atexit, which pins the + fixture's device and host pools until process exit; hundreds of HiCache + fixtures in this suite would otherwise accumulate on the GPU.""" + atexit.unregister(cache.shutdown) + + class TestUnifiedRadixComponentRegistryOverride(CustomTestCase): def test_component_registry_override_is_instance_local(self): params = CacheInitParams( @@ -1064,6 +1072,7 @@ def _init_hicache(self, cache, *, write_policy: str = "write_through"): ) set_global_server_args_for_scheduler(server_args) cache.init_hicache(server_args, cache.cache_init_params) + self.addCleanup(_drop_hicache_atexit_pin, cache) self.addCleanup(cache.release_host_resources) cache.write_through_threshold = 1 << 30 cache.load_back_threshold = 0 @@ -5519,6 +5528,7 @@ def _init_hicache( server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size) set_global_server_args_for_scheduler(server_args) cache.init_hicache(server_args, cache.cache_init_params) + self.addCleanup(_drop_hicache_atexit_pin, cache) self.addCleanup(cache.release_host_resources) cache.write_through_threshold = 1 << 30 cache.load_back_threshold = 0 @@ -9691,6 +9701,7 @@ def kv_host_pool_wrapper(*args, **kwargs): server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size) set_global_server_args_for_scheduler(server_args) cache.init_hicache(server_args, cache.cache_init_params) + self.addCleanup(_drop_hicache_atexit_pin, cache) cache.write_through_threshold = 1 << 30 cache.load_back_threshold = 0