diff --git a/benchmark/hicache/bench_buffer_mode.py b/benchmark/hicache/bench_buffer_mode.py new file mode 100644 index 000000000..5f60807a8 --- /dev/null +++ b/benchmark/hicache/bench_buffer_mode.py @@ -0,0 +1,404 @@ +"""Benchmark driver for HiCache buffer_only mode. + +Runs a matrix of (model, host-memory-mode) server configs through the two +workloads buffer mode targets, and reports hit rates, latency, and the +buffer-mode pipeline counters side by side: + +- multiturn: growing-history conversations (delegates to bench_multiturn.py, + offline random tokens, round barrier, fixed seed); +- longctx: long shared prefix + divergent continuations, measured warm + (device tier) and again after /flush_cache (through-storage tier). + +Example (dense, cache-vs-buffer): + python benchmark/hicache/bench_buffer_mode.py \ + --model /path/to/model --modes cache,buffer \ + --workloads multiturn,longctx --buffer-size-gb 2 --cache-ratio 2 + +SWA / Mamba hybrids run the same way (add --tp for large hybrids); the +unified radix tree is selected automatically for them. +""" + +import argparse +import json +import os +import shutil +import subprocess +import sys +import tempfile +import time + +import requests + +BENCH_DIR = os.path.dirname(os.path.abspath(__file__)) + +BUFFER_METRICS = [ + "sglang:hicache_existence_cache_skipped_pages_total", + "sglang:hicache_backup_dropped_tokens_total", + "sglang:hicache_pending_write_queue_depth", + "sglang:hicache_host_used_tokens", +] + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=str, required=True) + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--port", type=int, default=31212) + parser.add_argument( + "--launch-module", + type=str, + default="sglang.launch_server", + help="server entry module (e.g. sglang_meta.launch_server for " + "Meta-internal model families)", + ) + parser.add_argument( + "--modes", + type=str, + default="cache,buffer", + help="comma list of deployments to compare: cache (alias cache_wt), " + "cache_wb (write_back), buffer", + ) + parser.add_argument( + "--workloads", type=str, default="multiturn,longctx", help="comma list" + ) + parser.add_argument("--buffer-size-gb", type=int, default=2) + parser.add_argument("--cache-ratio", type=float, default=2.0) + parser.add_argument( + "--host-ratio", + type=float, + default=0.0, + help="if > 0, size EVERY mode's host pool as ratio x device pool " + "(overrides --buffer-size-gb / --cache-ratio) for apples-to-apples runs", + ) + parser.add_argument( + "--max-total-tokens", + type=int, + default=0, + help="cap the device KV pool so the working set overflows into the hierarchy", + ) + parser.add_argument("--mem-fraction-static", type=float, default=0.5) + parser.add_argument("--page-size", type=int, default=64) + parser.add_argument("--attention-backend", type=str, default="") + parser.add_argument( + "--storage-backend", + type=str, + default="file", + help="L3 backend: file, or meta_cache_store (spawns a local-mode store " + "under the run's storage dir; requires --launch-module " + "sglang_meta.launch_server)", + ) + parser.add_argument( + "--sub-question-input-length", + type=int, + default=0, + help="Input tokens per FOLLOW-UP turn (0 = same as --request-length). " + "Set with a large --request-length for long-context multiturn: " + "100K first prompt + small follow-ups.", + ) + parser.add_argument( + "--prefetch-policy", + type=str, + default="wait_complete", + choices=["best_effort", "wait_complete", "timeout"], + help="hicache storage prefetch stop policy (server default is " + "timeout; benches historically used wait_complete)", + ) + parser.add_argument( + "--storage-extra-config", + type=str, + default="", + help="JSON merged over the backend's default extra-config " + '(e.g. \'{"capacity_gb": "120"}\')', + ) + parser.add_argument( + "--extra-server-args", + type=str, + default="", + help="space-separated extra args appended to the server command", + ) + # multiturn shape + parser.add_argument("--num-clients", type=int, default=12) + parser.add_argument("--num-rounds", type=int, default=4) + parser.add_argument("--request-length", type=int, default=1024) + parser.add_argument("--output-length", type=int, default=64) + parser.add_argument("--request-rate", type=int, default=4) + # longctx shape + parser.add_argument("--prefix-tokens", type=int, default=6144) + parser.add_argument("--num-prefixes", type=int, default=4) + parser.add_argument("--continuations", type=int, default=4) + parser.add_argument("--out", type=str, default="bench_buffer_mode_results.json") + return parser.parse_args() + + +def scrape_metrics(base_url): + try: + text = requests.get(f"{base_url}/metrics", timeout=30).text + except Exception: + return {} + out = {} + for name in BUFFER_METRICS: + total, found = 0.0, False + for line in text.splitlines(): + if line.startswith(name + "{") or line.startswith(name + " "): + total += float(line.rsplit(" ", 1)[1]) + found = True + if found: + out[name.split(":")[1]] = total + return out + + +class Server: + def __init__(self, args, mode, storage_dir): + self.args = args + self.mode = mode + self.storage_dir = storage_dir + self.base_url = f"http://127.0.0.1:{args.port}" + self.proc = None + + def launch(self): + a = self.args + extra_config = {"prefetch_threshold": 64} + if a.storage_backend == "meta_cache_store": + # Local-mode store: in-process server on rank 0, data under the + # run's storage dir (same isolation as the file backend). + extra_config.update( + { + "cluster": "mks", + "local_mode": "true", + "data_root": os.path.join(self.storage_dir, "mcs"), + "capacity_gb": "120", + } + ) + if a.storage_extra_config: + extra_config.update(json.loads(a.storage_extra_config)) + cmd = [ + sys.executable, + "-m", + a.launch_module, + "--model-path", + a.model, + "--port", + str(a.port), + "--tp-size", + str(a.tp), + "--mem-fraction-static", + str(a.mem_fraction_static), + "--page-size", + str(a.page_size), + "--enable-hierarchical-cache", + "--enable-cache-report", + "--enable-metrics", + "--hicache-write-policy", + "write_through", + "--hicache-storage-prefetch-policy", + a.prefetch_policy, + "--hicache-storage-backend-extra-config", + json.dumps(extra_config), + ] + if self.mode == "buffer": + cmd += [ + "--hicache-host-memory-mode", + "buffer_only", + "--hicache-storage-backend", + a.storage_backend, + ] + if a.host_ratio > 0: + cmd += ["--hicache-ratio", str(a.host_ratio)] + else: + cmd += ["--hicache-size", str(a.buffer_size_gb)] + else: + write_policy = "write_back" if self.mode == "cache_wb" else "write_through" + cmd += [ + "--hicache-ratio", + str(a.host_ratio if a.host_ratio > 0 else a.cache_ratio), + "--hicache-storage-backend", + a.storage_backend, + ] + # Override the default write policy appended below. + cmd = [c if c != "write_through" else write_policy for c in cmd] + if a.max_total_tokens > 0: + cmd += ["--max-total-tokens", str(a.max_total_tokens)] + if a.attention_backend: + cmd += ["--attention-backend", a.attention_backend] + if a.extra_server_args: + cmd += a.extra_server_args.split() + env = { + **os.environ, + "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": self.storage_dir, + } + log_stem = os.path.splitext(os.path.basename(self.args.out))[0] + self.log = open(f"/tmp/{log_stem}_{self.mode}_server.log", "w") + self.proc = subprocess.Popen(cmd, stdout=self.log, stderr=self.log, env=env) + deadline = time.time() + 2400 + while time.time() < deadline: + if self.proc.poll() is not None: + raise RuntimeError(f"server died during launch; see {self.log.name}") + try: + if ( + requests.get(f"{self.base_url}/health", timeout=5).status_code + == 200 + ): + return + except Exception: + pass + time.sleep(2) + raise RuntimeError("server did not become healthy in time") + + def flush(self): + requests.post( + f"{self.base_url}/flush_cache", params={"timeout": 60}, timeout=90 + ).raise_for_status() + + def kill(self): + if self.proc is not None: + from sglang.srt.utils import kill_process_tree + + kill_process_tree(self.proc.pid) + self.proc = None + + +def run_multiturn(server, args): + log_file = tempfile.mktemp(suffix=".jsonl") + cmd = [ + sys.executable, + os.path.join(BENCH_DIR, "bench_multiturn.py"), + "--model-path", + args.model, + "--port", + str(args.port), + "--disable-auto-run", + "--disable-random-sample", + "--enable-round-barrier", + "--num-clients", + str(args.num_clients), + "--max-parallel", + str(args.num_clients), + "--num-rounds", + str(args.num_rounds), + "--request-length", + str(args.request_length), + "--output-length", + str(args.output_length), + "--request-rate", + str(args.request_rate), + "--seed", + "1", + "--log-file", + log_file, + ] + if args.sub_question_input_length > 0: + cmd += [ + "--sub-question-input-length", + str(args.sub_question_input_length), + ] + subprocess.run(cmd, check=True, timeout=7200) + with open(log_file) as f: + records = [json.loads(line) for line in f if line.strip()] + record = records[-1] + result = { + "hit_rate": record["summary"]["cache_hit_rate"], + "avg_ttft_s": record["summary"]["average_ttft"], + "p90_ttft_s": record["summary"].get("p90_ttft", 0), + } + for round_key, round_data in record.get("round", {}).items(): + result[f"{round_key}_hit"] = round(round_data["cache_hit_rate"], 4) + result[f"{round_key}_ttft_s"] = round(round_data["average_ttft"], 3) + return result + + +def _gen_token_prompt(tokenizer, n_tokens, seed): + import random as _random + + rng = _random.Random(seed) + vocab_size = min(tokenizer.vocab_size - 1000, 32000) + ids = [rng.randrange(1000, vocab_size) for _ in range(n_tokens * 2)] + text = tokenizer.decode(ids) + ids = tokenizer.encode(text)[:n_tokens] + return tokenizer.decode(ids) + + +def run_longctx(server, args): + """Long shared prefix + divergent continuations, warm then post-flush.""" + from sglang.benchmark.utils import get_tokenizer + + tokenizer = get_tokenizer(args.model) + prefixes = [ + _gen_token_prompt(tokenizer, args.prefix_tokens, seed=900 + g) + for g in range(args.num_prefixes) + ] + tails = [ + _gen_token_prompt(tokenizer, 128, seed=9900 + c) + for c in range(args.continuations) + ] + + def one_pass(tag): + latencies, cached, prompt_tokens = [], 0, 0 + for prefix in prefixes: + for tail in tails: + start = time.perf_counter() + res = requests.post( + f"{server.base_url}/generate", + json={ + "text": prefix + tail, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": 8, + "ignore_eos": True, + }, + }, + timeout=600, + ).json() + latencies.append(time.perf_counter() - start) + meta = res.get("meta_info", {}) + cached += int(meta.get("cached_tokens", 0)) + prompt_tokens += int(meta.get("prompt_tokens", 1)) + return { + f"{tag}_hit_rate": round(cached / max(prompt_tokens, 1), 4), + f"{tag}_avg_latency_s": round(sum(latencies) / len(latencies), 3), + f"{tag}_max_latency_s": round(max(latencies), 3), + } + + result = one_pass("warm") + # Give write-backs a moment to settle before dropping device state. + time.sleep(5) + server.flush() + result.update(one_pass("replay")) + return result + + +def main(): + args = parse_args() + results = {} + for mode in args.modes.split(","): + mode = mode.strip() + storage_dir = tempfile.mkdtemp(prefix=f"bench_buffer_{mode}_") + server = Server(args, mode, storage_dir) + print(f"\n=== launching {mode} server ===", flush=True) + try: + server.launch() + mode_result = {} + for workload in args.workloads.split(","): + workload = workload.strip() + print(f"--- {mode}: running {workload} ---", flush=True) + if workload == "multiturn": + mode_result["multiturn"] = run_multiturn(server, args) + elif workload == "longctx": + mode_result["longctx"] = run_longctx(server, args) + server.flush() + mode_result["hicache_metrics"] = scrape_metrics(server.base_url) + results[mode] = mode_result + finally: + server.kill() + shutil.rmtree(storage_dir, ignore_errors=True) + + print("\n===== results =====") + print(json.dumps(results, indent=2)) + with open(args.out, "w") as f: + json.dump( + {"model": args.model, "args": vars(args), "results": results}, f, indent=2 + ) + print(f"saved to {args.out}") + + +if __name__ == "__main__": + main() diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 6a1089254..31ee7c5ac 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -17,7 +17,7 @@ import logging import threading import time from queue import Empty, Queue -from typing import TYPE_CHECKING, List, NamedTuple, Optional +from typing import TYPE_CHECKING, Callable, List, NamedTuple, Optional import torch @@ -203,6 +203,13 @@ class StorageOperation: self.completed_tokens = 0 self.hash_value = hash_value if hash_value is not None else [] self.prefix_keys = prefix_keys + # Full queried page-hash chain, set by _storage_hit_query before + # hash_value is truncated to the hit boundary; the tail is the + # absence signal that invalidates buffer-mode existence beliefs. + self.all_hash_values: Optional[List[str]] = None + # Prefetch-outcome accounting, set at enqueue by the tree cache. + self.stats_requested_tokens = 0 + self.stats_total_tokens = 0 self.id = StorageOperation.counter StorageOperation.counter += 1 @@ -211,6 +218,15 @@ class StorageOperation: return self.id < other.id +# Buffer-mode staging budgets. Prefetch staging is latency-critical +# (wait_complete gates TTFT), so loads may fill the pool up to this fraction +# before new prefetches are declined. +HICACHE_LOAD_POOL_USAGE_FRACTION = 0.9 +# Write-staging floor: writes are deferrable, so the flush gate grows the +# write window dynamically into whatever load staging is not using. +HICACHE_WRITE_STAGING_POOL_FRACTION = 0.2 + + class PrefetchOperation(StorageOperation): def __init__( self, @@ -262,8 +278,10 @@ class HiCacheController: model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, enable_storage_metrics: bool = False, + host_memory_mode: str = "cache", ): self.tp_group = tp_group + self.host_memory_mode = host_memory_mode self.attn_cp_group = attn_cp_group self.attn_tp_group = attn_tp_group self.pp_group = pp_group @@ -283,6 +301,9 @@ class HiCacheController: self.storage_backend = None self.storage_backend_type = None self.enable_storage_metrics = enable_storage_metrics + # Buffer mode: wired by the tree cache after attach; the load rate + # limiter subtracts write staging from actual pool usage. + self.host_write_staged_tokens_fn: Optional[Callable[[], int]] = None # Draft KV pool support (best-effort piggyback on target L2/L3 ops). self.has_draft = False @@ -501,8 +522,16 @@ class HiCacheController: self.enable_storage = True # todo: threshold policy for prefetching self.prefetch_threshold = max(prefetch_threshold, self.page_size) - # Budget speculative prefetch at half the host pool, leaving the rest for the write-back staging path. - self.prefetch_capacity_limit = int(0.5 * self.mem_pool_host.size) + if self.host_memory_mode == "buffer_only": + # The whole pool is transient staging; loads may fill it up + # to this fraction, and the tree's write flush gate yields + # to live fetch demand (the write fraction is a floor). + self.prefetch_capacity_limit = int( + HICACHE_LOAD_POOL_USAGE_FRACTION * self.mem_pool_host.size + ) + else: + # Budget speculative prefetch at half the host pool, leaving the rest for the write-back staging path. + self.prefetch_capacity_limit = int(0.5 * self.mem_pool_host.size) # tracking the number of tokens locked in prefetching, updated by the main scheduler thread self.prefetch_tokens_occupied = 0 @@ -1050,6 +1079,16 @@ class HiCacheController: """ Rate limit the prefetching operations to avoid overwhelming the storage backend. """ + if self.host_memory_mode == "buffer_only": + # Gate on real pool usage: buffer mode allocates hit-sized, so + # prefetch_tokens_occupied's requested spans overstate it. Pool + # state mutates only at scheduler-thread lockstep points, so this + # stays TP-deterministic. Write staging is the write budget's + # usage; charging it here would park hits behind its storage drain. + used = self.mem_pool_host.size - self.mem_pool_host.available_size() + if self.host_write_staged_tokens_fn is not None: + used -= self.host_write_staged_tokens_fn() + return max(0, used) >= self.prefetch_capacity_limit # cancel prefetch if too much memory is occupied if self.prefetch_tokens_occupied >= self.prefetch_capacity_limit: return True @@ -1066,6 +1105,7 @@ class HiCacheController: page_hashes = self.get_hash_str( tokens_to_fetch, last_hash, page_size=self.page_size ) + operation.all_hash_values = page_hashes for start in range(0, len(page_hashes), STORAGE_BATCH_SIZE): batch_hashes = page_hashes[start : start + STORAGE_BATCH_SIZE] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b60db6484..ca3fb39df 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2698,8 +2698,28 @@ class Scheduler( if self.enable_hicache_storage: req.init_next_round_input(self.tree_cache, cow_mamba=False) tree_cache = self.tree_cache - if tree_cache.is_backuped(req.last_host_node) or tree_cache.is_root( - req.last_host_node + buffer_mode = self.server_args.hicache_host_memory_mode == "buffer_only" + last_host_node = req.last_host_node + # Buffer mode host-backups nothing, so match_prefix anchors at + # root; re-anchor on the deepest device node. The anchor is only + # read for hash/extra-key context here (never locked), so a device + # node serves. Cache mode keeps the is_backuped gate below: its + # write-through prefix is contiguous from root, so an unbacked + # anchor means a guaranteed storage miss. + if ( + buffer_mode + and tree_cache.is_root(last_host_node) + and not tree_cache.is_root(req.last_node) + ): + last_host_node = req.last_node + + if ( + tree_cache.is_backuped(last_host_node) + or tree_cache.is_root(last_host_node) + or ( + buffer_mode + 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( @@ -2707,16 +2727,17 @@ class Scheduler( ) new_input_tokens = req.full_untruncated_fill_ids[matched_len:match_end] prefix_keys = ( - tree_cache.get_prefix_hash_values(req.last_host_node) + tree_cache.get_prefix_hash_values(last_host_node) if tree_cache.hicache_storage_pass_prefix_keys else None ) tree_cache.prefetch_from_storage( req.rid, - req.last_host_node, + last_host_node, new_input_tokens, - tree_cache.get_last_hash_value(req.last_host_node), + tree_cache.get_last_hash_value(last_host_node), prefix_keys, + matched_prefix_tokens=req.full_untruncated_fill_ids[:matched_len], ) def _add_request_to_queue(self, req: Req, is_retracted: bool = False): @@ -3329,6 +3350,23 @@ class Scheduler( req.storage_hit_length = loaded_tokens req.init_next_round_input(self.tree_cache) + if ( + self.enable_hicache_storage + and self.server_args.hicache_host_memory_mode == "buffer_only" + ): + # 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. 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 = self.tree_cache.staged_prefetch_tokens(req.rid) + if held_tokens > 0: + req.host_hit_length = held_tokens + req.swa_host_hit_length = ( + self.tree_cache.staged_prefetch_swa_tokens(req.rid) + ) res = adder.add_one_req( req, has_chunked_req=(self.chunked_req is not None), @@ -4146,6 +4184,11 @@ class Scheduler( if tc.enable_storage: idle &= len(tc.ongoing_prefetch) == 0 idle &= len(tc.ongoing_backup) == 0 + if self.server_args.hicache_host_memory_mode == "buffer_only": + # Queued writes, staged prefetches, and in-flight + # storage writes still hold host staging + # (buffer-mode unified tree only). + idle &= tc.buffer_pipeline.is_idle() return idle diff --git a/python/sglang/srt/mem_cache/buffer_mode/__init__.py b/python/sglang/srt/mem_cache/buffer_mode/__init__.py new file mode 100644 index 000000000..a9f578666 --- /dev/null +++ b/python/sglang/srt/mem_cache/buffer_mode/__init__.py @@ -0,0 +1,9 @@ +"""Buffer-only HiCache host memory mode (--hicache-host-memory-mode buffer_only). + +Host RAM is a transient staging buffer between the GPU and the L3 storage +backend, never an L2 cache tier: writes stage device KV through op-owned +host bounces into storage and free them at the storage ack; reads fetch +storage hits into op-owned bounces and publish them into the device tree at +prefill admission. ``UnifiedRadixCache`` composes ``BufferModePipeline`` for +the two transfer pipelines and dispatches to it at the mode branches. +""" diff --git a/python/sglang/srt/mem_cache/buffer_mode/buffer_page_cache.py b/python/sglang/srt/mem_cache/buffer_mode/buffer_page_cache.py new file mode 100644 index 000000000..883e6c71b --- /dev/null +++ b/python/sglang/srt/mem_cache/buffer_mode/buffer_page_cache.py @@ -0,0 +1,396 @@ +"""Refcounted content-addressed page cache over the buffer-mode host pool. + +NOT WIRED YET: staged spans register as ``(pool, page_hash) -> (slots, +refcount)`` so prefetches can be served zero-copy from local staging +(write-around / promote-on-read retention, zero-ref LRU reclaim); keys are +the content-chained page hashes, so entries survive node deletion, splits, +and recompute. Counterpart of ``StorageExistenceCache`` (beliefs about +STORAGE, dedupes writes); this tracks LOCAL HOST RAM and dedupes loads. + +TP determinism: replicas must stay identical across attention ranks — the +cache feeds scheduler-visible structure, so divergence is a collective +hang, not a soft miss. Preconditions when wiring: (1) mutate only on the +scheduler thread at lockstep points; (2) rank-reduce any fold anchored by a +per-rank storage outcome (hit count, revoke) before it picks a mutation; +(3) controller queues stay FIFO and single-threaded so MIN-count drains +process the same prefix on every rank. +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Callable, Optional, Sequence + +import torch + +from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, + PoolName, + PoolTransfer, +) + + +class _PageRef: + """One cached page: host slot span, reader refcount, and whether to + retain at refs==0 (write-around: write staging frees at its storage ack + unless a read promoted it). ``first_slot`` mirrors ``slots[0]`` as a + plain int because per-page tensor-scalar reads are too slow in + ``release``.""" + + __slots__ = ("slots", "first_slot", "refs", "retain") + + def __init__( + self, slots: torch.Tensor, first_slot: int, refs: int = 1, retain: bool = True + ): + self.slots = slots + self.first_slot = first_slot + self.refs = refs + self.retain = retain + + +class BufferPageCache: + def __init__(self) -> None: + # (pool, page_hash) -> _PageRef; slots stay allocated in the host + # pool for as long as the entry exists. + self._entries: dict[tuple[str, str], _PageRef] = {} + # Per-pool zero-ref LRU (head = coldest): reclaim victims. + self._zero_ref: dict[str, OrderedDict[str, None]] = {} + # Per-pool slot tokens held by the cache (refed + zero-ref). + self._held_tokens: dict[str, int] = {} + # Per-pool slot tokens at refs=0 (reclaimable under pressure). + self._zero_ref_tokens: dict[str, int] = {} + + def __len__(self) -> int: + return len(self._entries) + + def num_zero_ref_pages(self) -> int: + return sum(len(lru) for lru in self._zero_ref.values()) + + def held_tokens(self, pool: str) -> int: + return self._held_tokens.get(pool, 0) + + def zero_ref_tokens(self, pool: str) -> int: + """Slot tokens reclaimable right now (zero-ref cached pages). + Occupancy/rate-limit gates must treat these as free-able, not used: + a pool full of zero-ref cache is one reclaim away from empty.""" + return self._zero_ref_tokens.get(pool, 0) + + def register( + self, + pool: str, + hashes: Sequence[str], + host_indices: torch.Tensor, + page_size: int, + retain: bool = True, + ) -> int: + """Cache a staged span, one entry per page, refs=1 (the staging op); + returns the number of pages newly cached. ``retain=False`` = + write-around (freed at last ref unless a read hit promotes it); a + duplicate hash keeps the existing entry and the newcomer's slots + stay op-owned for a raw free at release.""" + assert len(host_indices) == len(hashes) * page_size + registered = 0 + entries = self._entries + # One batched read of the page-boundary slot ids (see _PageRef). + first_slots = host_indices[::page_size].tolist() + for i, page_hash in enumerate(hashes): + key = (pool, page_hash) + existing = entries.get(key) + if existing is not None: + existing.retain = existing.retain or retain + continue + entries[key] = _PageRef( + host_indices[i * page_size : (i + 1) * page_size], + first_slots[i], + retain=retain, + ) + registered += 1 + if registered: + self._held_tokens[pool] = ( + self._held_tokens.get(pool, 0) + registered * page_size + ) + return registered + + def contains(self, pool: str, page_hash: str) -> bool: + """Non-mutating presence probe (no LRU touch).""" + return (pool, page_hash) in self._entries + + def peek_run_len(self, pool: str, hashes: Sequence[str]) -> int: + """Length of the leading run of cached pages. Non-mutating.""" + entries = self._entries + run = 0 + for page_hash in hashes: + if (pool, page_hash) not in entries: + break + run += 1 + return run + + def acquire(self, pool: str, hashes: Sequence[str]) -> Optional[torch.Tensor]: + """refs++ on every page and return the gathered slot tensor (pages + expanded to token slots, in page order). All-or-nothing: returns + None without mutating if any page is missing.""" + entries = self._entries + refs = [] + for page_hash in hashes: + entry = entries.get((pool, page_hash)) + if entry is None: + return None + refs.append(entry) + zero_ref = self._zero_ref.get(pool) + for page_hash, entry in zip(hashes, refs): + if entry.refs == 0 and zero_ref is not None: + if zero_ref.pop(page_hash, None) is not None: + self._zero_ref_tokens[pool] -= len(entry.slots) + entry.refs += 1 + # Read demand proven: promote write-around pages to retained. + entry.retain = True + return torch.cat([entry.slots for entry in refs]) + + def release( + self, + pool: str, + hashes: Sequence[str], + host_indices: torch.Tensor, + page_size: int, + ) -> Optional[torch.Tensor]: + """Drop one ref per page of a span: at refs==0 retained pages move + to the zero-ref LRU tail while write-around pages return their + slots for an immediate free. Duplicate-staging slots (canonical + entry lives elsewhere) are returned for a raw free without touching + the canonical refcount.""" + assert len(host_indices) == len(hashes) * page_size + leftover: list[torch.Tensor] = [] + entries = self._entries + # One batched read of the page-boundary slot ids (see _PageRef). + first_slots = host_indices[::page_size].tolist() + for i, page_hash in enumerate(hashes): + entry = entries.get((pool, page_hash)) + if entry is None or entry.first_slot != first_slots[i]: + leftover.append(host_indices[i * page_size : (i + 1) * page_size]) + continue + assert entry.refs > 0, "release without a matching acquire/register" + entry.refs -= 1 + if entry.refs == 0: + if entry.retain: + self._zero_ref.setdefault(pool, OrderedDict())[page_hash] = None + self._zero_ref_tokens[pool] = self._zero_ref_tokens.get( + pool, 0 + ) + len(entry.slots) + else: + del entries[(pool, page_hash)] + self._held_tokens[pool] -= len(entry.slots) + leftover.append(entry.slots) + if not leftover: + return None + return torch.cat(leftover) + + def reclaim( + self, + pool: str, + need_tokens: int, + free: Callable[[torch.Tensor], int], + ) -> int: + """Pop zero-ref LRU heads, free their slots back to the host pool, + and drop the entries. Called under allocation pressure only. Returns + the number of slot tokens freed (may undershoot when everything + left is refed).""" + zero_ref = self._zero_ref.get(pool) + if not zero_ref or need_tokens <= 0: + return 0 + freed = 0 + batch: list[torch.Tensor] = [] + while zero_ref and freed < need_tokens: + page_hash, _ = zero_ref.popitem(last=False) + entry = self._entries.pop((pool, page_hash)) + batch.append(entry.slots) + freed += len(entry.slots) + if batch: + free(torch.cat(batch)) + self._held_tokens[pool] -= freed + self._zero_ref_tokens[pool] -= freed + return freed + + +class BufferPageCacheOps: + """Pool-facing operations over a :class:`BufferPageCache`: span/hold + registration and release keyed the way the storage write keys them, + pressure reclaim, and the SWA-folded continuation fold. The caller owns + the collectives — rank-reduce any fold anchored by a per-rank storage + outcome before acting on it (see the module docstring).""" + + def __init__( + self, + page_cache: BufferPageCache, + mem_pool_host, + sw_window_pages_fn: Callable[[], int], + ): + # Rebound by the owner when the structure is recreated (reset). + self.page_cache = page_cache + self._mem_pool_host = mem_pool_host + # SWA window in KV pages when SWA stages through a host pool + # (0 = KV-only: no trailing window in the fold). + self._sw_window_pages_fn = sw_window_pages_fn + + def aux_window_keys( + self, hash_values: list[str], transfer: PoolTransfer + ) -> Optional[list[str]]: + """Trailing KV page hashes keying an aux transfer's staged window + (one key per aux-pool page), recomputed from the rank-synced span + hashes so registration and release always agree across ranks.""" + if transfer.host_indices is None or transfer.host_indices.numel() == 0: + return None + if transfer.indices_from_pool is not None: + return None # sidecar rides another pool's slots; nothing to key + entry = self._mem_pool_host.entry_map.get(transfer.name) + if entry is None: + return None + pool_page_size = entry.host_pool.page_size + num_keys = len(transfer.host_indices) // pool_page_size + if num_keys == 0 or num_keys > len(hash_values): + return None + return hash_values[-num_keys:] + + def register_span( + self, + pool: PoolName, + hashes: list[str], + host_indices: torch.Tensor, + retain: bool = True, + ) -> None: + """Cache a page-aligned staged span (refs=1 for the staging op).""" + if not hashes: + return + entry = self._mem_pool_host.entry_map.get(pool) + if entry is None: + return + self.page_cache.register( + pool, hashes, host_indices, entry.host_pool.page_size, retain=retain + ) + + def release_span( + self, + pool: PoolName, + hashes: list[str], + host_indices: torch.Tensor, + ) -> None: + """Drop the staging op's ref on a span; zero-ref pages stay cached + (servable) until pressure reclaims them. Op-owned duplicate slots + (their hash was cached elsewhere) are freed raw, as before.""" + if host_indices is None or host_indices.numel() == 0: + return + entry = self._mem_pool_host.entry_map.get(pool) + if entry is None: + return + if not hashes: + entry.host_pool.free(host_indices) + return + leftover = self.page_cache.release( + pool, hashes, host_indices, entry.host_pool.page_size + ) + if leftover is not None and leftover.numel() > 0: + entry.host_pool.free(leftover) + + def register_hold( + self, + hash_values: list[str], + host_indices: torch.Tensor, + aux_xfers: list[PoolTransfer], + retain: bool = True, + ) -> None: + """Register a staged KV span plus its aux windows (SWA/Mamba states + keyed by their trailing KV page hashes, same keying the storage + write uses). ``retain=False`` = write-around: servable only while + the staging op pins the slots, freed at the last release unless a + read hit promotes it.""" + self.register_span(PoolName.KV, hash_values, host_indices, retain=retain) + for transfer in aux_xfers: + keys = self.aux_window_keys(hash_values, transfer) + if keys is not None: + self.register_span( + transfer.name, keys, transfer.host_indices, retain=retain + ) + + def release_hold( + self, + hash_values: list[str], + host_indices: torch.Tensor, + aux_xfers: list[PoolTransfer], + ) -> None: + """Mirror of register_hold for every hold retirement path + (storage-ack, fill H2D-ack, staged drop, abort).""" + self.release_span(PoolName.KV, hash_values, host_indices) + for transfer in aux_xfers: + if transfer.indices_from_pool is not None: + continue + keys = self.aux_window_keys(hash_values, transfer) + self.release_span(transfer.name, keys or [], transfer.host_indices) + + def reclaim(self, pool: PoolName, num_tokens: int) -> int: + """Free just enough zero-ref cached pages for an allocation of + num_tokens to succeed. Scheduler-thread only (lockstep pressure + points: staging-hit alloc, prepare_prefetch, cc.write).""" + entry = self._mem_pool_host.entry_map.get(pool) + if entry is None: + return 0 + shortfall = num_tokens - entry.host_pool.available_size() + if shortfall <= 0: + return 0 + return self.page_cache.reclaim(pool, shortfall, entry.host_pool.free) + + def continuation_run(self, chain: list[str], start_pages: int) -> int: + """Longest cached run continuing the span at page ``start_pages`` + (0 = leading run), folded for SWA: the joint span's trailing window + must be fully cache-servable, mirroring batch_exists_v2's + trailing_pages fold. Non-mutating and rank-deterministic.""" + page_cache = self.page_cache + kv_run = page_cache.peek_run_len(PoolName.KV, chain[start_pages:]) + if kv_run == 0: + return 0 + sw_pages = self._sw_window_pages_fn() + if sw_pages == 0: + return kv_run + for cont in range(kv_run, 0, -1): + joint = start_pages + cont + window = min(sw_pages, joint) + if cont < window: + # Window straddles into the head; only possible for + # anchored runs, and shrinking cont cannot fix it. + break + if all( + page_cache.contains(PoolName.SWA, chain[i]) + for i in range(joint - window, joint) + ): + return cont + return 0 + + def acquire_span( + self, chain: list[str], start_pages: int, cont_pages: int + ) -> Optional[tuple[torch.Tensor, list[PoolTransfer]]]: + """Acquire a folded continuation run: its KV pages plus the JOINT + span's trailing SWA window (refs++ on every page). Returns + (kv_slots, aux_xfers) or None (with no refs held) if a page + vanished since the fold — defensive; fold and acquire run in the + same lockstep step.""" + page_cache = self.page_cache + cont_hashes = list(chain[start_pages : start_pages + cont_pages]) + kv_slots = page_cache.acquire(PoolName.KV, cont_hashes) + if kv_slots is None: + return None + aux_xfers: list[PoolTransfer] = [] + sw_pages = self._sw_window_pages_fn() + if sw_pages > 0: + joint = start_pages + cont_pages + window_hashes = list(chain[joint - min(sw_pages, joint) : joint]) + swa_slots = page_cache.acquire(PoolName.SWA, window_hashes) + if swa_slots is None: + self.release_span(PoolName.KV, cont_hashes, kv_slots) + return None + aux_xfers.append( + PoolTransfer( + name=PoolName.SWA, + host_indices=swa_slots, + keys=window_hashes, + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ) + return kv_slots, aux_xfers diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py new file mode 100644 index 000000000..8e71eb911 --- /dev/null +++ b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py @@ -0,0 +1,816 @@ +"""Buffer-only mode transfer pipelines for the unified radix cache. + +``BufferModePipeline`` owns all buffer-mode state and the two pipelines that +move KV through the transient host staging buffer: + +- backup (write path): admission-gated FIFO intents, head-of-line D2H + staging launches, storage writes at the D2H ack, staging freed at the + storage ack; +- load back (read path): completed storage fetches parked as op-owned host + bounces, consumed at prefill admission via a device alloc + layer-gated + H2D + plain tree insert, bounce freed at the H2D ack. + +The pipeline is an intimate collaborator of ``UnifiedRadixCache``: it is +constructed by ``init_hicache`` only when ``--hicache-host-memory-mode +buffer_only`` is active, and it drives tree/controller operations (insert, +match, evict, lock refs, cache actions) through the owning cache. All +buffer-mode-only state lives here; the cache dispatches to this object at +its mode branches. + +TP-lockstep contract: every mutation runs on the scheduler thread at +rank-synchronized points (insert walks, rank-MIN-reduced drains, ack +drains), so per-rank state never diverges. There is no runtime +verification; a violation surfaces as an unexplained collective hang. +""" + +from __future__ import annotations + +import logging +from array import array +from collections import deque +from typing import TYPE_CHECKING, Optional + +import msgspec +import torch + +from sglang.srt.managers.cache_controller import HICACHE_WRITE_STAGING_POOL_FRACTION +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + EvictParams, + InitLoadBackParams, + InsertParams, + MatchPrefixParams, +) +from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer +from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.unified_cache.cache_action import RebuildFullToSWAMapping +from sglang.srt.mem_cache.unified_cache.components import ( + BASE_COMPONENT_TYPE, + ComponentType, +) +from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( + NodeId, + UnifiedTreeNode, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.unified_cache.components import SWAComponent + from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache + +logger = logging.getLogger(__name__) + + +class _UnifiedBackupIntent(msgspec.Struct): + """Buffer-mode backup intent, unpinned while queued. + + Snapshots node identity at enqueue time: a split rewrites the node's + key/hash in place while these copies stay intact, so + ``node.hash_value != hash_values`` doubles as split detection and a None + FULL device value as eviction detection (``_backup_intent_stale``). + """ + + node: UnifiedTreeNode + node_id: int + hash_values: list[str] + key: RadixKey + prefix_keys: Optional[list[str]] = None + + +class _UnifiedBufferBackupEntry(msgspec.Struct): + """A buffer-mode backup after its D2H launch: intent + staging slots. + + ``host_indices`` are FULL-pool staging slots; ``aux_xfers`` carry the + staged aux-pool slots (e.g. the SWA window). All are freed at the + storage-write ack — host memory is never retained as a cache tier. + """ + + intent: _UnifiedBackupIntent + host_indices: torch.Tensor + aux_xfers: list[PoolTransfer] + lock_params: DecLockRefParams + + +class _StagedPrefetch(msgspec.Struct): + """A completed buffer-mode fetch parked until prefill admission: only + the op-owned host bounce exists (no device state, nothing in the tree). + """ + + req_id: str + key_tokens: list[int] + extra_key: Optional[str] + matched_len: int + num_tokens: int + occupied_tokens: int + host_indices: torch.Tensor + aux_xfers: list[PoolTransfer] + hash_values: list[str] + operation_id: int + + +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. + """ + + req_id: str + num_tokens: int + occupied_tokens: int + aux_xfers: list[PoolTransfer] + host_indices: torch.Tensor + hash_values: list[str] + + +def _track_content_refs(refs: dict[str, int], hash_values: list[str]) -> None: + """Add one content ref per page hash (at D2H launch). Refcounted, + not a flag: several launched entries can carry the same content + (duplicate staging of republished spans).""" + for h in hash_values: + refs[h] = refs.get(h, 0) + 1 + + +def _untrack_content_refs(refs: dict[str, int], hash_values: list[str]) -> None: + """Drop one content ref per page hash (at storage-ack).""" + for h in hash_values: + n = refs.get(h, 0) - 1 + if n <= 0: + refs.pop(h, None) + else: + refs[h] = n + + +def validate_buffer_only_stack( + sidecar_pool_specs: list, swa_component: Optional[SWAComponent] +) -> None: + """Post-assembly buffer-mode fences. + + Sidecar pools (DSv4 compressed regions) and unified_kv SWA (device-only + ring, never offloaded) have no per-pool staging path yet. + """ + if sidecar_pool_specs: + raise ValueError( + "--hicache-host-memory-mode buffer_only does not support " + "sidecar storage pools (DeepSeek-V4 compressed regions)." + ) + swa = swa_component + if swa is not None and swa._swa_kv_pool_host is None: + # Only reachable on SWA models with the unified_kv layout (SWA as + # a device-only ring): without a host pool the window can neither + # stage for writes nor fetch for load-backs. + raise ValueError( + "--hicache-host-memory-mode buffer_only on SWA models " + "requires an SWA host staging pool; the unified_kv layout " + "keeps SWA as a device-only ring." + ) + if swa is not None and swa._swa_kv_pool_host is not None: + # Below two windows the pool cannot hold a staging write AND the + # loads-priority reserve (_aux_loads_margin floors at one + # window), so every window-carrying intent would be dropped as + # oversize and SWA storage coverage would silently be zero. + window_tokens = swa.full_window_pages * swa._swa_kv_pool_host.page_size + if swa._swa_kv_pool_host.size < 2 * window_tokens: + raise ValueError( + "--hicache-host-memory-mode buffer_only requires an SWA " + f"host pool of at least two trailing windows " + f"({2 * window_tokens} tokens; got " + f"{swa._swa_kv_pool_host.size}): one staging a write " + "while one stays reserved for prefetch window allocs." + ) + + +class BufferModePipeline: + """All buffer-mode state plus the backup and load-back pipelines. + + Constructed by ``UnifiedRadixCache.init_hicache`` when host memory mode + is ``buffer_only``; ``cache.buffer_pipeline is None`` elsewhere, which + the cache's mode branches use as the dispatch test. + """ + + def __init__( + self, + cache: UnifiedRadixCache, + swa_window_pages: int, + write_backlog_cap: int, + ): + self._cache = cache + # SWA window size in KV pages when the SWA component stages through + # a host pool (0 = KV-only: no trailing window staged). Static after + # pool assembly. + self._swa_window_pages = swa_window_pages + # Metadata-only pending-write backlog cap; beyond it new intents + # are dropped at admission (re-trigger on a later hit). + self.write_backlog_cap = write_backlog_cap + self.reset() + + 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._prefetch_prefix_ctx: dict[str, list[int]] = {} + self.staged_prefetches: dict[str, _StagedPrefetch] = {} + self.ongoing_buffer_load_back: dict[int, _OngoingBufferLoadBack] = {} + # Backup pipeline: FIFO intents awaiting a D2H slot, node ids + # anywhere in flight (dedupes re-triggers), and a content refcount + # of every page hash between D2H launch and storage-ack — admission + # skips content covered by beliefs + launched writes. + self.pending_write_queue: deque[_UnifiedBackupIntent] = deque() + self.inflight_backup_node_ids: set[int] = set() + self.inflight_backup_hashes: dict[str, int] = {} + # Backups between D2H launch and D2H ack (keyed by node id), then + # between storage-write launch and storage ack (keyed by operation + # id). Mirrors the cache-mode ongoing_write_through/ongoing_backup + # stages, with buffer entries. + self.ongoing_write_through: dict[int, _UnifiedBufferBackupEntry] = {} + self.ongoing_backup: dict[int, _UnifiedBufferBackupEntry] = {} + self.write_staged_tokens_ = 0 + self.write_backlog_tokens_ = 0 + self._backlog_cap_hits = 0 + + def is_idle(self) -> bool: + """No queued writes, staged prefetches, or storage writes in flight + (all of which hold host staging or would re-trigger IO).""" + return not ( + self.pending_write_queue or self.staged_prefetches or self.ongoing_backup + ) + + # ---- backup pipeline (device -> staging -> storage) ---- + + def _backup_parent_covered(self, node: UnifiedTreeNode) -> bool: + """Only admit a node whose parent is stored/in-flight: writing above + a dropped parent creates a permanent longest-prefix hole.""" + parent = node.parent + if ( + parent is self._cache.root_node + or parent.id in self.inflight_backup_node_ids + ): + return True + last_hash = parent.get_last_hash_value() + return last_hash is not None and self._cache.storage_existence_cache.contains( + PoolName.KV, last_hash + ) + + def _log_backup_dropped(self, num_tokens: int) -> None: + cache = self._cache + if cache.enable_storage_metrics and cache.storage_metrics_collector is not None: + cache.storage_metrics_collector.log_backup_dropped_tokens(num_tokens) + + def enqueue_backup_intent(self, node: UnifiedTreeNode) -> None: + """Snapshot a backup intent and commit it to the write queue. + Admission gates: belief skip, parent-cover, backlog cap, oversize. + Drops are silent; the node re-triggers on a later hit.""" + if not self._cache.enable_storage or not node.hash_value: + return + if node.id in self.inflight_backup_node_ids: + return + # Admission cover: beliefs plus content past its D2H launch. The + # launched cover keeps republished content (fill inserts under new + # node ids) from re-writing while the original write drains. + if self._cache.storage_existence_cache.covers_all( + PoolName.KV, node.hash_value, extra_cover=self.inflight_backup_hashes + ): + return + intent_tokens = len(node.hash_value) * self._cache.page_size + if self.write_backlog_tokens_ >= self.write_backlog_cap: + # The cap sits at 2x the intrinsic live-backlog ceiling (see + # init_hicache), so reaching it means leaked accounting or a + # broken stale sweep — a bug, not load. + self._backlog_cap_hits += 1 + if self._backlog_cap_hits <= 3 or self._backlog_cap_hits % 1000 == 0: + logger.error( + "HiCache write backlog cap hit (occurrence %d): " + "backlog=%d cap=%d queue=%d. Live backlog is bounded " + "by the device pool span, so this indicates a " + "stale-sweep or accounting leak.", + self._backlog_cap_hits, + self.write_backlog_tokens_, + self.write_backlog_cap, + len(self.pending_write_queue), + ) + self._log_backup_dropped(intent_tokens) + return + # A span larger than any pool's whole staging capacity can never + # stage; admitting it would wedge the head-of-line queue forever. + if not self._backup_parent_covered(node) or self._backup_oversize( + node, intent_tokens + ): + self._log_backup_dropped(intent_tokens) + return + + prefix_keys = ( + node.get_prefix_hash_values(node.parent) + if self._cache.hicache_storage_pass_prefix_keys + else None + ) + intent = _UnifiedBackupIntent( + node=node, + node_id=node.id, + hash_values=list(node.hash_value), + key=node.key, + prefix_keys=prefix_keys, + ) + self.pending_write_queue.append(intent) + self.inflight_backup_node_ids.add(node.id) + self.write_backlog_tokens_ += intent_tokens + + def _build_aux_staging_transfers( + self, node: UnifiedTreeNode + ) -> Optional[list[PoolTransfer]]: + """Keys-only aux transfers mirroring what BACKUP_STORAGE would write; + sizes the per-pool oversize gate (beliefs do not consult these).""" + transfers: list[PoolTransfer] = [] + if ComponentType.SWA in self._cache.components: + cd = node.component_data[ComponentType.SWA] + if cd.value is not None: + num_pages = len(cd.value) // self._cache.page_size + if num_pages > 0: + transfers.append( + PoolTransfer( + name=PoolName.SWA, + keys=node.hash_value[-num_pages:], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ) + return transfers or None + + def _backup_oversize( + self, + node: UnifiedTreeNode, + intent_tokens: int, + aux_xfers: Optional[list[PoolTransfer]] = None, + ) -> bool: + """True if any pool's staging need exceeds that pool's write-usable + capacity (total for KV, total minus the loads-priority margin for aux + pools — matching ``_aux_budget_blocked``'s admission ceiling): such an + intent could never stage and would wedge the FIFO head.""" + cc = self._cache.cache_controller + if intent_tokens > cc.mem_pool_host.size: + return True + if aux_xfers is None: + aux_xfers = self._build_aux_staging_transfers(node) + for t in aux_xfers or (): + entry = cc.mem_pool_host.entry_map.get(t.name) + if entry is not None and ( + len(t.keys) * entry.host_pool.page_size + > entry.host_pool.size - self._aux_loads_margin(entry.host_pool) + ): + return True + return False + + def _aux_loads_margin(self, host_pool) -> int: + """Aux-pool tokens reserved for loads: at least one trailing window + (prepare_prefetch allocates its window here and a failed alloc + forfeits the whole prefetch), plus a 10% burst absorber mirroring + live_cap.""" + return max( + self._swa_window_pages * host_pool.page_size, + host_pool.size // 10, + ) + + def _backup_intent_stale(self, intent: _UnifiedBackupIntent) -> bool: + # Arena-lookup failure = deleted, hash mismatch vs the enqueue-time + # snapshot = split, a None FULL device value = evicted. Stale + # intents drop silently; the node re-triggers on a later hit. + node = intent.node + try: + self._cache.tree_core.node_by_id(intent.node_id) + except KeyError: + return True + return ( + node.component_data[BASE_COMPONENT_TYPE].value is None + or node.hash_value != intent.hash_values + ) + + def _sweep_stale_backup_intents(self) -> None: + """Cancel stale intents anywhere in the queue, not just at the head: + a dead intent would otherwise inflate the backlog accounting and + hold FIFO position ahead of live segments.""" + if not self.pending_write_queue: + return + page_size = self._cache.page_size + survivors: deque[_UnifiedBackupIntent] = deque() + for intent in self.pending_write_queue: + if self._backup_intent_stale(intent): + self.inflight_backup_node_ids.discard(intent.node_id) + self.write_backlog_tokens_ -= len(intent.hash_values) * page_size + continue + survivors.append(intent) + self.pending_write_queue = survivors + + def flush_pending_writes(self) -> None: + """Launch D2H transfers for admitted intents, head-of-line: device + locks and staging slots are taken only here, when capacity allows.""" + if not self.pending_write_queue: + return + cc = self._cache.cache_controller + self._sweep_stale_backup_intents() + # Loads have priority (writes are deferrable): the write window is + # the pool minus prefetch occupancy minus a 10% margin, floored at + # the configured fraction. + pool_tokens = cc.mem_pool_host.size + live_cap = max( + int(HICACHE_WRITE_STAGING_POOL_FRACTION * pool_tokens), + pool_tokens - cc.prefetch_tokens_occupied - pool_tokens // 10, + ) + while self.pending_write_queue: + intent = self.pending_write_queue[0] + intent_tokens = len(intent.hash_values) * self._cache.page_size + if not self._backup_parent_covered(intent.node) or self._backup_oversize( + intent.node, intent_tokens + ): + # Unwritable intent (dropped parent or unstageable size): + # cascade the drop down the chain rather than creating a + # permanent storage hole / stalling the head-of-line queue. + self.pending_write_queue.popleft() + self.inflight_backup_node_ids.discard(intent.node_id) + self.write_backlog_tokens_ -= intent_tokens + self._log_backup_dropped(intent_tokens) + continue + if self.write_staged_tokens_ >= live_cap: + # Yield to live fetch demand; retry next round. + break + if self._aux_budget_blocked(intent): + # An aux pool lacks staging headroom: yield at the gate + # instead of failing the alloc inside cc.write; acks free + # aux staging, retry next round. + break + if not self._launch_backup_intent(intent): + # Pool full of in-flight staging and nothing reclaimable + # (the tree never holds host values in buffer mode): + # defer, head-of-line; pending acks will free slots. + break + self.pending_write_queue.popleft() + + def _launch_backup_intent(self, intent: _UnifiedBackupIntent) -> bool: + """Launch one admitted intent's D2H (staging alloc + device lock + + async copy); the caller removes it from pending_write_queue. Returns + False when staging cannot be allocated. From a successful launch the + intent always reaches its storage-ack, so its content joins the + LAUNCHED cover consulted by admission.""" + cache = self._cache + cc = cache.cache_controller + node = intent.node + # Build aux transfers from the node's CURRENT state: a SWA span + # tombstoned since admission backs up FULL-only, as in cache mode. + device_value, comp_xfers = cache.tree_core.build_backup_spec(node.id) + aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] + host_indices = cc.write( + device_value, + node_id=node.id, + extra_pools=aux_xfers or None, + ) + if host_indices is None: + return False + _track_content_refs(self.inflight_backup_hashes, intent.hash_values) + # NOTE: no commit_backup — the node must never appear + # host-resident in buffer mode; staging slots live in the entry. + lock_params = cache.inc_lock_ref(node.id).to_dec_params() + self.ongoing_write_through[node.id] = _UnifiedBufferBackupEntry( + intent=intent, + host_indices=host_indices, + aux_xfers=aux_xfers, + lock_params=lock_params, + ) + self.write_staged_tokens_ += len(host_indices) + self.write_backlog_tokens_ -= len(intent.hash_values) * cache.page_size + return True + + def _aux_budget_blocked(self, intent: _UnifiedBackupIntent) -> bool: + """True when an aux pool cannot stage this intent right now (free + minus the loads-priority margin falls short of the need): defer at + the gate instead of failing the alloc inside cc.write and blocking + pure-KV intents behind an unallocatable head. The margin enforces + loads-have-priority on aux pools the way live_cap does on the KV + pool; avail already reflects prefetch-held slots, so no occupancy + subtraction here.""" + aux = self._build_aux_staging_transfers(intent.node) + if not aux: + return False + cc = self._cache.cache_controller + for t in aux: + entry = cc.mem_pool_host.entry_map.get(t.name) + if entry is None: + continue + need = len(t.keys) * entry.host_pool.page_size + headroom = entry.host_pool.available_size() - self._aux_loads_margin( + entry.host_pool + ) + if need > headroom: + return True + return False + + def _aux_window_keys( + self, hash_values: list[str], transfer: PoolTransfer + ) -> Optional[list[str]]: + """Trailing KV page hashes keying an aux transfer's staged window + (one key per aux-pool page).""" + if transfer.host_indices is None or transfer.host_indices.numel() == 0: + return None + if transfer.indices_from_pool is not None: + return None # sidecar rides another pool's slots; nothing to key + entry = self._cache.cache_controller.mem_pool_host.entry_map.get(transfer.name) + if entry is None: + return None + num_keys = len(transfer.host_indices) // entry.host_pool.page_size + if num_keys == 0 or num_keys > len(hash_values): + return None + return hash_values[-num_keys:] + + def finish_backup_ack(self, ack_id: int) -> None: + """D2H confirmed: drop the device lock and enqueue the storage write + (which reads from the staging copy, so device eviction may proceed).""" + entry = self.ongoing_write_through.pop(ack_id) + intent = entry.intent + self._cache.dec_lock_ref(intent.node_id, entry.lock_params) + + # Every aux pool writes a trailing snapshot keyed by the last KV page + # hashes it covers: the SWA window spans page_size-sized pages, the + # Mamba state is a single slot (host pool page_size 1 -> one key). + storage_xfers: list[PoolTransfer] = [] + for staged in entry.aux_xfers: + keys = self._aux_window_keys(intent.hash_values, staged) + if keys is None: + continue + storage_xfers.append( + PoolTransfer( + name=staged.name, + host_indices=staged.host_indices, + keys=keys, + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ) + operation_id = self._cache.cache_controller.write_storage( + entry.host_indices, + intent.key.token_ids, + intent.hash_values, + intent.prefix_keys, + extra_pools=storage_xfers or None, + ) + self.ongoing_backup[operation_id] = entry + + def finish_storage_write_ack(self, operation_id: int) -> None: + """Storage write acked (rank-synced drain): free the entry's staging + outright. Existence entries are added unconditionally + (completed_tokens can diverge across ranks under backend failure) to + keep admission decisions TP-deterministic. No-op for operations this + pipeline does not own (e.g. acks for already-reset state).""" + entry = self.ongoing_backup.pop(operation_id, None) + if entry is None: + return + intent = entry.intent + self._cache.storage_existence_cache.add(PoolName.KV, intent.hash_values) + self._free_staging_now(entry.host_indices, entry.aux_xfers) + self.write_staged_tokens_ -= len(entry.host_indices) + self.inflight_backup_node_ids.discard(entry.intent.node_id) + _untrack_content_refs(self.inflight_backup_hashes, intent.hash_values) + + def _free_staging_now( + self, host_indices: torch.Tensor, aux_xfers: list[PoolTransfer] + ) -> None: + """Synchronously free a staging span (KV + aux pools) on the + scheduler thread; buffer-mode acks/drops all run here, so frees + land before the tick's next gate reads pool availability.""" + cc = self._cache.cache_controller + if host_indices is not None and host_indices.numel() > 0: + cc.mem_pool_host.free(host_indices) + for t in aux_xfers or (): + if ( + t.host_indices is None + or t.host_indices.numel() == 0 + or t.indices_from_pool is not None + ): + continue + entry = cc.mem_pool_host.entry_map.get(t.name) + if entry is not None: + entry.host_pool.free(t.host_indices) + + # ---- load back pipeline (storage -> staging -> device) ---- + + def set_prefix_ctx(self, req_id: str, matched_prefix_tokens) -> None: + """Record the device-matched prefix at prefetch enqueue; consumed at + staging commit to build the full-span tree key.""" + self._prefetch_prefix_ctx[req_id] = list(matched_prefix_tokens or []) + + def pop_prefix_ctx(self, req_id: str) -> None: + self._prefetch_prefix_ctx.pop(req_id, None) + + def has_staged(self, req_id: str) -> bool: + return req_id in self.staged_prefetches + + @staticmethod + def _occupied_span(host_indices) -> int: + """Occupancy units a buffer-mode prefetch holds: granted at + hit-alloc, sized to the allocation (0 while still querying).""" + return len(host_indices) if host_indices is not None else 0 + + def stage_completed_prefetch( + self, + req_id: str, + num_tokens: int, + hash_value: list[str], + ) -> bool: + """Park the completed fetch as a held bounce; the scheduler surfaces + 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(req_id) + cc = cache.cache_controller + prefix_tokens = self._prefetch_prefix_ctx.pop(req_id, None) + aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] + + if num_tokens == 0 or prefix_tokens is None: + # Nothing usable fetched: recompute. + cc.append_host_mem_release( + host_indices[:num_tokens], extra_pools=aux_xfers or None + ) + cc.prefetch_tokens_occupied -= self._occupied_span(host_indices) + cache.prefetch_loaded_tokens_by_reqid[req_id] = 0 + return True + + staged_pages = num_tokens // cache.page_size + staged_hashes = hash_value[:staged_pages] + staged_kv = host_indices[:num_tokens] + # Feed existence beliefs from the storage-fetched pages: the fetch + # itself is the evidence, so feeding is sound even if this staged + # prefetch is later dropped unconsumed. + cache.storage_existence_cache.add(PoolName.KV, list(staged_hashes)) + occupied_tokens = self._occupied_span(host_indices) + + self.staged_prefetches[req_id] = _StagedPrefetch( + req_id=req_id, + key_tokens=prefix_tokens + list(prefetch_key[:num_tokens].token_ids), + extra_key=prefetch_key.extra_key, + matched_len=len(prefix_tokens), + num_tokens=num_tokens, + occupied_tokens=occupied_tokens, + host_indices=staged_kv, + aux_xfers=aux_xfers, + hash_values=staged_hashes, + operation_id=operation.id, + ) + cache.prefetch_loaded_tokens_by_reqid[req_id] = num_tokens + return True + + def staged_prefetch_tokens(self, req_id: str) -> int: + """Tokens a staged prefetch would splice (0 = no hold); surfaced by the + scheduler as the request's host_hit_length.""" + f = self.staged_prefetches.get(req_id) + return f.num_tokens if f is not None else 0 + + def staged_prefetch_swa_tokens(self, req_id: str) -> int: + """SWA device tokens consuming this staged prefetch will allocate (the + staged trailing window); surfaced as the request's swa_host_hit_length + so the adder's SWA gate charges the admission-time alloc.""" + f = self.staged_prefetches.get(req_id) + if f is None: + return 0 + return sum( + len(t.host_indices) + for t in f.aux_xfers + if t.name == PoolName.SWA and t.host_indices is not None + ) + + def init_load_back(self, params: InitLoadBackParams) -> tuple[torch.Tensor, NodeId]: + """Buffer-mode branch of init_load_back: consume the staged prefetch + at prefill admission — device alloc (evict-before-alloc), layer-gated + H2D, and a plain insert so downstream sees ordinary tree state. + Misaligned or alloc-failed holds drop; the request recomputes.""" + cache = self._cache + req = params.req + assert req is not None + empty = cache.tree_core.empty_match_result.device_indices + unchanged = (empty, req.last_node) + f = self.staged_prefetches.pop(req.rid, None) + if f is None: + return unchanged + cc = cache.cache_controller + + def _drop() -> tuple[torch.Tensor, NodeId]: + self._free_staging_now(f.host_indices, f.aux_xfers) + cc.prefetch_tokens_occupied -= f.occupied_tokens + return unchanged + + # Splice-validity: the span only fits if the device prefix still + # ends exactly at the enqueue-time matched_len. + if len(req.prefix_indices) != f.matched_len: + # Prefix moved while held (leaf eviction or sibling extension): + # drop and recompute. + return _drop() + + # Evict-before-alloc (mirrors _load_back_transfers): the budget gate + # counts evictable pages, but cc.load draws from free slots only. + 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 < f.num_tokens: + needed = f.num_tokens - avail + evicted = cache.evict(EvictParams(num_tokens=needed)) + if evicted.num_tokens_evicted < needed: + # Genuinely no room (locked pages): recompute. + return _drop() + + load_back_id = -(f.operation_id) - 1 + device_indices = cc.load( + host_indices=f.host_indices, + node_id=load_back_id, + extra_pools=f.aux_xfers or None, + ) + if device_indices is None: + # Transient allocator shortfall despite the evict: recompute + # (init_load_back's degrade contract). + return _drop() + + swa_dev = next( + ( + t.device_indices + for t in f.aux_xfers + if t.name == PoolName.SWA + and t.device_indices is not None + and t.device_indices.numel() > 0 + ), + None, + ) + 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]) + ) + + # Publish via a plain insert under the admission lock choreography; + # the caller's request lock then pins the span (load_back pattern). + key = RadixKey( + array("q", f.key_tokens), + extra_key=f.extra_key, + is_bigram=cache.tree_core.is_eagle, + ).page_aligned(cache.page_size) + span_end = f.matched_len + f.num_tokens + cache.insert( + InsertParams( + key=key, + value=torch.cat([req.prefix_indices, device_indices]), + prev_prefix_len=f.matched_len, + swa_evicted_seqlen=( + max(0, span_end - len(swa_dev)) if swa_dev is not None else 0 + ), + ) + ) + self.ongoing_buffer_load_back[load_back_id] = _OngoingBufferLoadBack( + req_id=f.req_id, + num_tokens=f.num_tokens, + occupied_tokens=f.occupied_tokens, + aux_xfers=f.aux_xfers, + host_indices=f.host_indices, + hash_values=f.hash_values, + ) + m = cache.match_prefix(MatchPrefixParams(key=key)) + if len(m.device_indices) < span_end: + # The insert walk did not adopt the full span (should not happen + # for a locked prefix); the slots are tree-owned/evictable — do + # not splice, the request recomputes. + return unchanged + return device_indices, m.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 + a buffer-mode load-back. The span was published at admission; the + ack never touches the tree (existence beliefs were fed from the + storage-fetched pages at staging commit).""" + f = self.ongoing_buffer_load_back.pop(ack_id, None) + if f is None: + return False + cache = self._cache + cc = cache.cache_controller + + # The H2D consumed the bounce buffers; free them outright. + self._free_staging_now(f.host_indices, f.aux_xfers) + + cc.prefetch_tokens_occupied -= f.occupied_tokens + logger.info( + "HiCache prefetch fill committed req=%s filled=%d occupied=%d", + f.req_id, + f.num_tokens, + cc.prefetch_tokens_occupied, + ) + if cache.enable_storage_metrics and cache.storage_metrics_collector is not None: + cache.storage_metrics_collector.log_prefetched_tokens(f.num_tokens) + return True + + def release_aborted_staged(self, rid: str) -> bool: + """Free an aborted request's staged prefetch (nothing device-side + exists yet — only the bounce). Returns True when a hold existed.""" + staged = self.staged_prefetches.pop(rid, None) + if staged is None: + return False + self._free_staging_now(staged.host_indices, staged.aux_xfers) + self._cache.cache_controller.prefetch_tokens_occupied -= staged.occupied_tokens + return True diff --git a/python/sglang/srt/mem_cache/buffer_mode/storage_existence_cache.py b/python/sglang/srt/mem_cache/buffer_mode/storage_existence_cache.py new file mode 100644 index 000000000..dcb3ff3c5 --- /dev/null +++ b/python/sglang/srt/mem_cache/buffer_mode/storage_existence_cache.py @@ -0,0 +1,89 @@ +"""Local existence cache for HiCache buffer_only mode. + +In buffer mode host memory holds no persistent copy, so without a local +existence signal every re-insert of a hot prefix re-stages and re-writes to +L3 storage. This cache is that signal: a bounded LRU of (pool, page-hash) +entries *believed* present in storage. + +Semantics are advisory, not authoritative: + +- A hit skips the redundant D2H + storage write. +- A stale positive (backend evicted the data) costs skipped write-backs until + a prefetch hit-query shortfall invalidates the entries; the next insert + then writes the data back. Never a correctness issue — at worst one cold + recompute, the same as any cache miss. +- A miss (entry LRU-evicted or never seen) just costs one redundant write + (idempotent: storage keys are content-addressed). + +Keys are the content-chained page hashes already computed at insert time, so +lookups never hash anything and entries survive node deletion, splits, and +recompute (same tokens => same chain). Page hashes are chained and the write +path is prefix-contiguous (parent-cover gate), so a node's own page set is +the only thing a caller needs to check. + +TP determinism: replicas stay identical because every mutation happens on the +scheduler thread at lockstep points with cross-rank-reduced inputs +(storage-ack drain, prefetch-hit drain, fill commit). Do not touch it from +anywhere else. +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Container, Iterable, Sequence + +# ~131K entries; at ~150-250 B/entry this is <= ~30 MB and covers roughly +# 8M KV tokens at page size 64 (aux-pool entries included). Coverage per MB +# scales with page size — small-page configs simply remember fewer tokens. +HICACHE_EXISTENCE_CACHE_MAX_ENTRIES = 128 * 1024 + + +class StorageExistenceCache: + def __init__(self, max_entries: int = HICACHE_EXISTENCE_CACHE_MAX_ENTRIES): + self.max_entries = max_entries + self._entries: OrderedDict[tuple[str, str], None] = OrderedDict() + + def __len__(self) -> int: + return len(self._entries) + + def add(self, pool: str, hashes: Iterable[str]) -> None: + entries = self._entries + for h in hashes: + entries[(pool, h)] = None + entries.move_to_end((pool, h)) + while len(entries) > self.max_entries: + entries.popitem(last=False) + + def contains(self, pool: str, page_hash: str) -> bool: + entries = self._entries + if (pool, page_hash) not in entries: + return False + entries.move_to_end((pool, page_hash)) + return True + + def contains_all(self, pool: str, hashes: Iterable[str]) -> bool: + return all(self.contains(pool, h) for h in hashes) + + def covers_all( + self, + pool: str, + hashes: Iterable[str], + extra_cover: Container[str] = frozenset(), + ) -> bool: + """True when every page is believed stored or sits in + ``extra_cover`` (e.g. content past its D2H launch, which always + reaches its storage-ack). LRU-touches the believed entries.""" + return all(self.contains(pool, h) or h in extra_cover for h in hashes) + + def invalidate_beyond( + self, pool: str, hashes: Sequence[str], keep_pages: int + ) -> None: + """Ground-truth heal from a prefetch hit query: discard beliefs + beyond the leading ``keep_pages`` of a hash chain (the folded + usable cut). The next insert re-writes the discarded span, closing + stale positives and aux holes at the cut.""" + for h in hashes[keep_pages:]: + self._entries.pop((pool, h), None) + + def clear(self) -> None: + self._entries.clear() diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b86485657..dc65a7828 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1775,6 +1775,8 @@ class HiRadixCache(RadixCache): new_input_tokens: List[int], last_hash: Optional[str] = None, prefix_keys: Optional[List[str]] = None, + # Scheduler-call parity with UnifiedRadixCache; unused in cache mode. + matched_prefix_tokens: Optional[List[int]] = None, ): prefetch_key = RadixKey( new_input_tokens, 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 2c47e326d..4d8a5e683 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 @@ -112,6 +112,7 @@ class HybridCacheController(BaseHiCacheController): storage_backend_extra_config: Optional[dict] = None, transfer_layer_num: Optional[int] = None, enable_storage_metrics: bool = False, + host_memory_mode: str = "cache", ): startup_storage_backend = storage_backend self.extra_host_mem_release_queues: dict[PoolName, Queue[torch.Tensor]] = {} @@ -131,6 +132,7 @@ class HybridCacheController(BaseHiCacheController): model_name=model_name, storage_backend_extra_config=storage_backend_extra_config, enable_storage_metrics=enable_storage_metrics, + host_memory_mode=host_memory_mode, ) # Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)), # not just the full attention layers reported by full_kv_pool. @@ -587,6 +589,7 @@ class HybridCacheController(BaseHiCacheController): hash_value = self.get_hash_str( operation.token_ids, operation.last_hash, page_size=self.page_size ) + operation.all_hash_values = hash_value extra_info = HiCacheStorageExtraInfo( prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None @@ -780,6 +783,28 @@ class HybridCacheController(BaseHiCacheController): continue trailing_n = len(transfer.keys) if transfer.keys else 1 transfer.keys = all_hashes[max(0, kv_hit_pages - trailing_n) : kv_hit_pages] + if transfer.host_indices is None: + continue + entry = self.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 + ) + needed = len(transfer.keys) * pool_page_size + if transfer.host_indices.numel() > needed: + # The hit undershot the pre-allocated window buffer. Backends + # fetch keys zipped against the buffer head, so shrink the + # transfer to match and release the tail now — otherwise the + # length mismatch makes batch_get_v2 fetch nothing and the + # whole window is silently lost downstream. + self.append_host_mem_release( + extra_pools=[ + PoolTransfer( + name=transfer.name, + host_indices=transfer.host_indices[needed:], + ) + ] + ) + transfer.host_indices = transfer.host_indices[:needed] def _resolve_pool_transfers_allocation( self, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index d071d7534..d1e55258a 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -307,6 +307,7 @@ def build_kv_only_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + host_memory_mode=server_args.hicache_host_memory_mode, ) if params.mtp_draft_device_pools: cache_controller.set_mtp_draft_pools(params.mtp_draft_device_pools) @@ -377,6 +378,7 @@ def build_hybrid_swa_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + host_memory_mode=server_args.hicache_host_memory_mode, ) if mtp_swa_device_pools: cache_controller.set_mtp_draft_pools(mtp_swa_device_pools) @@ -664,6 +666,7 @@ def build_deepseek_v4_hicache_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + host_memory_mode=server_args.hicache_host_memory_mode, ) if mtp_swa_device_buffers: cache_controller.set_mtp_draft_pools(mtp_swa_device_buffers) @@ -759,6 +762,7 @@ def build_hybrid_mamba_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + host_memory_mode=server_args.hicache_host_memory_mode, ) if mtp_draft_device_pools: cache_controller.set_mtp_draft_pools(mtp_draft_device_pools) @@ -874,6 +878,7 @@ def build_hybrid_mamba_swa_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + host_memory_mode=server_args.hicache_host_memory_mode, ) return host_pool_group, cache_controller @@ -951,6 +956,7 @@ def build_anchor_sidecar_stack( storage_backend_extra_config=storage_backend_extra_config, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, + host_memory_mode=server_args.hicache_host_memory_mode, ) if mtp_draft_device_pools: cache_controller.set_mtp_draft_pools(mtp_draft_device_pools) diff --git a/python/sglang/srt/mem_cache/pool_host/mamba.py b/python/sglang/srt/mem_cache/pool_host/mamba.py index 22d1f8b17..8560ec8a5 100644 --- a/python/sglang/srt/mem_cache/pool_host/mamba.py +++ b/python/sglang/srt/mem_cache/pool_host/mamba.py @@ -129,7 +129,7 @@ class MambaPoolHost(HostKVCache): for conv_state in device_pool.mamba_cache.conv ] - self.init_kv_buffer() + self.kv_buffer = self.init_kv_buffer() self._init_write_back_staging_buffers() self.lock = threading.RLock() self.clear() @@ -201,6 +201,12 @@ class MambaPoolHost(HostKVCache): allocator=self.allocator, ) ) + # destroy() unregisters via kv_buffer; without this list the pinned + # registrations leak past the buffers' mmap. 0-element buffers + # (conv-only models' temporal state) were never registered. + return [ + buf for buf in (self.temporal_buffer, *self.conv_buffer) if buf.numel() > 0 + ] def _init_write_back_staging_buffers(self): self.temporal_staging_buffer = None diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index 8aaec2641..de3f94c6e 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -237,6 +237,16 @@ def create_tree_cache(ctx: TreeCacheBuildContext) -> BasePrefixCache: cache = default_radix_cache_factory(ctx) source = "default" + if ( + ctx.server_args.enable_hierarchical_cache + and ctx.server_args.hicache_host_memory_mode == "buffer_only" + and type(cache).__name__ != "UnifiedRadixCache" + ): + raise ValueError( + "--hicache-host-memory-mode buffer_only is only implemented for " + f"the unified radix tree; this model selected {type(cache).__name__}." + ) + if ctx.server_args.enable_session_radix_cache and not getattr( cache, "enable_session_radix_cache", False ): diff --git a/python/sglang/srt/mem_cache/storage/backend_factory.py b/python/sglang/srt/mem_cache/storage/backend_factory.py index 0fe83fdaf..d24a54190 100644 --- a/python/sglang/srt/mem_cache/storage/backend_factory.py +++ b/python/sglang/srt/mem_cache/storage/backend_factory.py @@ -158,7 +158,7 @@ class StorageBackendFactory: mem_pool_host: Any, ) -> HiCacheStorage: """Create built-in backend with original initialization logic.""" - if backend_name == "file": + if backend_name in ("file", "sim"): return backend_class(storage_config) elif backend_name == "nixl": return backend_class(storage_config) @@ -198,6 +198,10 @@ StorageBackendFactory.register_backend( "file", "sglang.srt.mem_cache.hicache_storage", "HiCacheFile" ) +StorageBackendFactory.register_backend( + "sim", "sglang.srt.mem_cache.storage.sim_storage", "SimHiCacheStorage" +) + StorageBackendFactory.register_backend( "nixl", "sglang.srt.mem_cache.storage.nixl.hicache_nixl", diff --git a/python/sglang/srt/mem_cache/storage/sim_storage.py b/python/sglang/srt/mem_cache/storage/sim_storage.py new file mode 100644 index 000000000..8c4aaf47d --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/sim_storage.py @@ -0,0 +1,254 @@ +"""Deterministic no-IO storage simulator for HiCache benchmarking. + +Stores KEYS ONLY (served KV is garbage — benchmark-only, for ignore_eos +workloads where nothing reads the generated text) and sleeps +``bytes / bandwidth + op latency`` on the calling backup/prefetch thread, +so pipeline dynamics are preserved while the medium is exactly +reproducible. Bandwidth is per scheduler rank (each rank ships its own +shard). extra_config knobs: ``sim_write_gbps`` (default 5.0; <=0 = +infinite), ``sim_read_gbps`` (default = write), ``sim_op_latency_us`` +(default 100, applied to exists queries too). +""" + +from __future__ import annotations + +import logging +import threading +import time +from typing import Any, List, Optional + +import torch + +from sglang.srt.mem_cache.hicache_storage import ( + HiCacheStorage, + HiCacheStorageConfig, + HiCacheStorageExtraInfo, + PoolHitPolicy, + PoolName, + PoolTransfer, + PoolTransferResult, +) + +logger = logging.getLogger(__name__) + + +class SimHiCacheStorage(HiCacheStorage): + def __init__(self, storage_config: HiCacheStorageConfig): + extra = storage_config.extra_config or {} + self.write_gbps = float(extra.get("sim_write_gbps", 5.0)) + self.read_gbps = float(extra.get("sim_read_gbps", self.write_gbps)) + self.op_latency_s = float(extra.get("sim_op_latency_us", 100.0)) * 1e-6 + # Scoped key -> True. Keys only; there are no bytes to store. + self._keys: set[str] = set() + self._lock = threading.Lock() + logger.info( + "SimHiCacheStorage: write=%.2f GB/s read=%.2f GB/s latency=%.0fus " + "(per rank; <=0 GB/s = infinite)", + self.write_gbps, + self.read_gbps, + self.op_latency_s * 1e6, + ) + + # ---- timing model ---- + + def _sleep_io(self, num_bytes: int, gbps: float) -> None: + delay = self.op_latency_s + if gbps > 0: + delay += num_bytes / (gbps * 1e9) + if delay > 0: + time.sleep(delay) + + def _pool_bytes(self, name: PoolName, num_slots: int) -> int: + return num_slots * self.registered_pools[name].size_per_token + + @staticmethod + def _scoped(name: PoolName, key: str) -> str: + return key if name == PoolName.KV else f"{key}.{name}" + + # ---- single-key surface (generic controller paths) ---- + + def get( + self, + key: str, + target_location: Optional[Any] = None, + target_sizes: Optional[Any] = None, + ) -> torch.Tensor | None: + with self._lock: + present = key in self._keys + return target_location if present else None + + def set( + self, + key: str, + value: Optional[Any] = None, + target_location: Optional[Any] = None, + target_sizes: Optional[Any] = None, + ) -> bool: + with self._lock: + self._keys.add(key) + return True + + def exists(self, key: str) -> bool: + with self._lock: + return key in self._keys + + # ---- batch v0/v1 (generic page funcs) ---- + + def batch_get( + self, + keys: List[str], + target_locations: Optional[Any] = None, + target_sizes: Optional[Any] = None, + ) -> List[torch.Tensor | None]: + locations = target_locations or [None] * len(keys) + with self._lock: + present = [k in self._keys for k in keys] + num_bytes = sum( + loc.numel() * loc.element_size() + for loc, hit in zip(locations, present) + if hit and loc is not None + ) + self._sleep_io(num_bytes, self.read_gbps) + return [loc if hit else None for loc, hit in zip(locations, present)] + + def batch_set( + self, + keys: List[str], + values: Optional[Any] = None, + target_locations: Optional[Any] = None, + target_sizes: Optional[Any] = None, + ) -> bool: + num_bytes = sum(v.numel() * v.element_size() for v in values or ()) + self._sleep_io(num_bytes, self.write_gbps) + with self._lock: + self._keys.update(keys) + return True + + def batch_get_v1( + self, + keys: List[str], + host_indices: torch.Tensor, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> List[bool]: + with self._lock: + present = [k in self._keys for k in keys] + hit_slots = (len(host_indices) // max(1, len(keys))) * sum(present) + self._sleep_io(hit_slots * self.mem_pool_host.size_per_token, self.read_gbps) + return present + + def batch_set_v1( + self, + keys: List[str], + host_indices: torch.Tensor, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> List[bool]: + self._sleep_io( + len(host_indices) * self.mem_pool_host.size_per_token, self.write_gbps + ) + with self._lock: + self._keys.update(keys) + return [True] * len(keys) + + # ---- batch v2 (hybrid multi-pool paths) ---- + + def batch_exists( + self, + keys: List[str], + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> int: + self._sleep_io(0, self.read_gbps) + with self._lock: + for i, key in enumerate(keys): + if key not in self._keys: + return i + return len(keys) + + def batch_exists_v2( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> PoolTransferResult: + """Same fold semantics as HiCacheFile.batch_exists_v2, over the + in-memory key set.""" + self._sleep_io(0, self.read_gbps) + with self._lock: + snapshot = self._keys.copy() + + kv_pages = next( + (i for i in range(len(keys)) if keys[i] not in snapshot), len(keys) + ) + hit_count: dict[str, int] = {PoolName.KV: kv_pages} if kv_pages else {} + final_pages = kv_pages + + for transfer in pool_transfers or []: + if final_pages == 0: + break + name = transfer.name + if transfer.hit_policy == PoolHitPolicy.ALL_PAGES: + boundary = next( + ( + i + for i in range(kv_pages) + if self._scoped(name, keys[i]) not in snapshot + ), + kv_pages, + ) + else: # trailing_pages + trailing = max(1, len(transfer.keys) if transfer.keys else 1) + boundary = 0 + for prefix_len in range(kv_pages, 0, -1): + if all( + self._scoped(name, keys[i]) in snapshot + for i in range(max(0, prefix_len - trailing), prefix_len) + ): + boundary = prefix_len + break + if boundary: + hit_count[name] = boundary + final_pages = min(final_pages, boundary) + + return PoolTransferResult(final_pages, hit_count) + + def _batch_v2( + self, transfers: List[PoolTransfer], gbps: float, record: bool + ) -> dict[str, List[bool]]: + results: dict[str, List[bool]] = {} + num_bytes = 0 + for t in transfers: + t_keys = t.keys or [] + if t.host_indices is not None: + num_bytes += self._pool_bytes(t.name, len(t.host_indices)) + scoped = [self._scoped(t.name, k) for k in t_keys] + with self._lock: + if record: + self._keys.update(scoped) + results[t.name] = [True] * len(t_keys) + else: + results[t.name] = [k in self._keys for k in scoped] + self._sleep_io(num_bytes, gbps) + return results + + def batch_get_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> dict[str, List[bool]]: + return self._batch_v2(transfers, self.read_gbps, record=False) + + def batch_set_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> dict[str, List[bool]]: + return self._batch_v2(transfers, self.write_gbps, record=True) + + # ---- misc ---- + + def clear(self) -> bool: + with self._lock: + self._keys.clear() + return True + + def get_stats(self): + return None diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index 1ff56aaf1..ac49edef7 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -71,6 +71,9 @@ class SWAComponent(TreeComponent): super().__init__(cache, params) self._session_leaf_covered_len: dict[str, dict[UnifiedTreeNode, int]] = {} self.sliding_window_size = params.sliding_window_size + self.full_window_pages = ( + self.sliding_window_size + params.page_size - 1 + ) // params.page_size # HiCache state: set to host SWA pool when HiCache enabled self._swa_kv_pool_host = None @@ -770,12 +773,27 @@ class SWAComponent(TreeComponent): # unified_kv keeps SWA as a device-only ring -- nothing to prefetch into. if self._swa_kv_pool_host is None: return PreparePrefetchResult() - sw_pages = ( - self.cache.sliding_window_size + self.cache.page_size - 1 - ) // self.cache.page_size - if sw_pages == 0 or prefetch_tokens // self.cache.page_size < sw_pages: + sw_pages = self.full_window_pages + if sw_pages == 0: return PreparePrefetchResult() - num_tokens = sw_pages * self.cache.page_size + 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.cache.host_memory_mode == "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. + num_pages = prefetch_pages + else: + # Cache-mode graft: a mid-tree window head is not + # device-guaranteed, require a full window. + return PreparePrefetchResult() + num_tokens = num_pages * self.cache.page_size host_indices = self._swa_kv_pool_host.alloc(num_tokens) if host_indices is None: self.cache.evict_host(num_tokens, ComponentType.SWA) @@ -869,12 +887,14 @@ class SWAComponent(TreeComponent): if phase == CacheTransferPhase.PREFETCH: assert host_indices is not None - sw_pages = host_indices.numel() // self.tree_core.page_size + # 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 return [ PoolTransfer( name=PoolName.SWA, host_indices=host_indices, - keys=["__placeholder__"] * sw_pages, + keys=["__placeholder__"] * num_pages, hit_policy=PoolHitPolicy.TRAILING_PAGES, ) ] @@ -997,6 +1017,13 @@ class SWAComponent(TreeComponent): and insert_result.inserted_host_node is not None else None ) + if anchor is not self.tree_core.root_node: + # Cache-mode graft commit only (buffer fills never reach here): + # a hit-shrunk window mid-tree is missing its head — drop it. + # Root anchors are complete windows of their own. + if window_require_pages < self.full_window_pages: + self._release_swa_host(host_indices, cache_actions) + return if ( target is None or window_require_pages == 0 diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 2f0dfb517..ba99c9385 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -25,6 +25,13 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.buffer_mode.pipeline import ( + BufferModePipeline, + validate_buffer_only_stack, +) +from sglang.srt.mem_cache.buffer_mode.storage_existence_cache import ( + StorageExistenceCache, +) from sglang.srt.mem_cache.common import RetractionBackup from sglang.srt.mem_cache.hicache_storage import ( PoolHitPolicy, @@ -71,6 +78,7 @@ from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401 ) from sglang.srt.observability.metrics_collector import ( STAT_LOGGER_ROLE_STORAGE, + StorageMetrics, StorageMetricsCollector, resolve_collector_class, ) @@ -227,6 +235,30 @@ class UnifiedRadixCache(BasePrefixCache): self.prefetch_timeout_base = 1.0 self.prefetch_timeout_per_page = 0.25 self.hicache_storage_pass_prefix_keys = False + # Buffer-only host memory mode (host RAM as transient GPU↔storage + # staging, not an L2 tier); resolved in init_hicache, which also + # constructs the pipeline collaborator (None = cache mode). + self.host_memory_mode = "cache" + self.buffer_pipeline: Optional[BufferModePipeline] = None + # Write-side dedupe: beliefs about what storage already holds, so + # re-inserts of hot prefixes skip the redundant backup. + self.storage_existence_cache = StorageExistenceCache() + # Cumulative prefetch-outcome counters, exported through the + # log_storage_metrics flow. + self._prefetch_outcome_stats: dict[str, float] = { + "attempts": 0, + "issued": 0, + "declined_too_short": 0, + "declined_rate_limited": 0, + "revoked_insufficient": 0, + "revoked_full_miss": 0, + "l3_demand_requests": 0, + "l3_miss_tokens": 0, + "l1l2_miss_tokens": 0, + "l3_demand_total_tokens": 0, + "l3_sum_rate_all": 0.0, + "l3_sum_rate_main_weighted": 0.0, + } self.reset() logger.info( @@ -316,6 +348,8 @@ class UnifiedRadixCache(BasePrefixCache): self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {} self.ongoing_prefetch: dict[str, _OngoingPrefetch] = {} self.ongoing_backup: dict[int, tuple[NodeId, DecLockRefParams]] = {} + if self.buffer_pipeline is not None: + self.buffer_pipeline.reset() if self.cache_controller is not None: self.cache_controller.reset() @@ -326,6 +360,20 @@ class UnifiedRadixCache(BasePrefixCache): def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None: """Initialize HiCache infrastructure.""" + self.host_memory_mode = server_args.hicache_host_memory_mode + if self.host_memory_mode == "buffer_only": + # FULL and FULL+SWA only: Mamba has no state-handoff channel on + # the admission-time load-back read path and is not layer-gated. + # Lifting the fence also needs the admission charge: a staged + # state slot is request-pinned at consumption and must ride + # req.mamba_host_hit_length the way the SWA window does. + supported = {ComponentType.FULL, ComponentType.SWA} + if not set(self.tree_components) <= supported: + raise ValueError( + "--hicache-host-memory-mode buffer_only supports only " + "FULL/SWA unified trees; got components " + f"{sorted(ct.name for ct in self.tree_components)}." + ) from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( attach_hybrid_pool_to_unified_cache, ) @@ -368,6 +416,28 @@ class UnifiedRadixCache(BasePrefixCache): swa = self.components[ComponentType.SWA] self.tree_core.has_swa_host_pool = swa._swa_kv_pool_host is not None + if self.host_memory_mode == "buffer_only": + swa = self.components.get(ComponentType.SWA) + validate_buffer_only_stack( + sidecar_pool_specs=self.sidecar_pool_specs, swa_component=swa + ) + self.buffer_pipeline = BufferModePipeline( + cache=self, + swa_window_pages=( + swa.full_window_pages + if swa is not None and self.tree_core.has_swa_host_pool + else 0 + ), + # Leak backstop only: live queued tokens are intrinsically + # bounded by the FULL device pool (one intent per node, stale + # intents swept per tick), so a cap that binds on live + # content would drop-newest and punch storage holes. + write_backlog_cap=2 * self.token_to_kv_pool_allocator.size_full, + ) + self.cache_controller.host_write_staged_tokens_fn = ( + lambda: self.buffer_pipeline.write_staged_tokens_ + ) + # State initialization self.write_through_threshold = ( 1 if server_args.hicache_write_policy == "write_through" else 2 @@ -541,6 +611,9 @@ class UnifiedRadixCache(BasePrefixCache): request_by_type: dict[ComponentType, int], tracker: dict[ComponentType, int], ) -> None: + # Buffer mode: eviction always wins over queued backup intents — a + # destroyed victim's intent is stale-swept and the content rewrites + # after its recompute. for ct in self.tree_components: request_cnt = request_by_type[ct] # Skip eviction walk if request is already met @@ -926,6 +999,10 @@ class UnifiedRadixCache(BasePrefixCache): self, num_tokens: int, component_type: ComponentType = BASE_COMPONENT_TYPE ) -> int: """Evict host resources for a specific component to free host pool space.""" + if self.host_memory_mode == "buffer_only": + # The tree never holds host values in buffer mode, and staging + # is operation-owned (freed at each ack): nothing is evictable. + return 0 result = self.tree_core.drive_host_eviction(component_type, num_tokens) self._free_values(result.device_frees, result.host_frees) return result.tracker.get(component_type, 0) @@ -1170,6 +1247,16 @@ class UnifiedRadixCache(BasePrefixCache): self, action: BackupKV, write_back: bool = False ) -> int: """Run a backup action top-down, stopping at the first failed backup.""" + if self.buffer_pipeline is not None: + # Buffer mode bypasses the host-backup contiguity below: nothing + # is ever host-backuped here. Contiguity comes from end-to-end + # FIFO ordering instead (BackupKV chains are parent-before-child + # and every pipeline stage drains in order). + for node_id in action.node_ids: + self.buffer_pipeline.enqueue_backup_intent( + self.tree_core.node_by_id(node_id) + ) + return 0 written = 0 for node_id in action.node_ids: device_value, comp_xfers = self.tree_core.build_backup_spec(node_id) @@ -1249,6 +1336,10 @@ class UnifiedRadixCache(BasePrefixCache): ) def _finish_write_through_ack(self, ack_id: int) -> None: + if self.buffer_pipeline is not None: + self.buffer_pipeline.finish_backup_ack(ack_id) + return + lock_node_id, lock_params, publish_node_ids = self.ongoing_write_through.pop( ack_id ) @@ -1468,10 +1559,12 @@ class UnifiedRadixCache(BasePrefixCache): new_input_tokens: list[int], last_hash: Optional[str] = None, prefix_keys: Optional[list[str]] = None, + matched_prefix_tokens: Optional[list[int]] = None, ) -> None: if not self.enable_storage or self.cache_controller is None: return + buffer_mode = self.host_memory_mode == "buffer_only" extra_key, cache_salt = self.tree_core.prefetch_anchor_info(last_host_node_id) prefetch_key = RadixKey( new_input_tokens, @@ -1480,13 +1573,30 @@ class UnifiedRadixCache(BasePrefixCache): cache_salt=cache_salt, ).page_aligned(self.page_size) prefetch_length = len(prefetch_key) - if ( - prefetch_length < self.prefetch_threshold - or self.cache_controller.prefetch_rate_limited() + stats = self._prefetch_outcome_stats + if prefetch_length > 0: + stats["attempts"] += 1 + if prefetch_length < self.prefetch_threshold: + if prefetch_length > 0: + stats["declined_too_short"] += 1 + return + if not buffer_mode and self.cache_controller.prefetch_rate_limited(): + stats["declined_rate_limited"] += 1 + return + if req_id in self.ongoing_prefetch or ( + buffer_mode and self.buffer_pipeline.has_staged(req_id) ): + # A fetch (or an unconsumed hold) already exists for this rid; + # overwriting would leak its staging slots. return - anchor_lock_params = self.inc_host_lock_ref(last_host_node_id).to_dec_params() + # Buffer mode holds no tree state during the fetch: buffers are + # operation-owned, so the anchor needs no pin. + anchor_lock_params = ( + None + if buffer_mode + 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: @@ -1517,10 +1627,21 @@ class UnifiedRadixCache(BasePrefixCache): 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], ) - self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) + if anchor_lock_params is not None: + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) return aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] @@ -1532,6 +1653,13 @@ class UnifiedRadixCache(BasePrefixCache): prefix_keys, extra_pools=aux_xfers or None, ) + stats["issued"] += 1 + # Snapshots for the L3 miss accounting at the query outcome (the + # hit/revoke drains): requested span and total prompt length. + operation.stats_requested_tokens = prefetch_length + operation.stats_total_tokens = prefetch_length + len( + matched_prefix_tokens or [] + ) self.ongoing_prefetch[req_id] = _OngoingPrefetch( last_host_node_id, prefetch_key, @@ -1540,7 +1668,12 @@ class UnifiedRadixCache(BasePrefixCache): anchor_lock_params, comp_xfers, ) - self.cache_controller.prefetch_tokens_occupied += len(prefetch_key) + if buffer_mode: + self.buffer_pipeline.set_prefix_ctx(req_id, matched_prefix_tokens) + else: + # Cache mode reserves the requested span up front; buffer mode + # grants occupancy later at hit-alloc time, sized to the hit. + self.cache_controller.prefetch_tokens_occupied += len(prefetch_key) def _prefetch_timeout_check_linear_func(self, operation: PrefetchOperation) -> bool: return ( @@ -1622,6 +1755,16 @@ class UnifiedRadixCache(BasePrefixCache): # Hybrid all-or-nothing check failed; result already discarded. return True + if self.buffer_pipeline is not None: + # No graft: release the rank-local tail beyond the synced usable + # length, then park the bounce for admission-time consumption. + self.cache_controller.append_host_mem_release( + host_indices[min_completed_tokens:completed_tokens] + ) + return self.buffer_pipeline.stage_completed_prefetch( + req_id, min_completed_tokens, hash_value + ) + fetched_key = prefetch_key[:min_completed_tokens] insert_result = self.tree_core.insert_host( last_host_node_id, @@ -1751,9 +1894,14 @@ class UnifiedRadixCache(BasePrefixCache): host_indices=host_indices[:completed_tokens], extra_pools=pool_transfers, ) - self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) + if anchor_lock_params is not None: + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) del self.ongoing_prefetch[req_id] - self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) + if self.buffer_pipeline is not None: + self.buffer_pipeline.pop_prefix_ctx(req_id) + self.cache_controller.prefetch_tokens_occupied -= ( + self._prefetch_occupied_span(prefetch_key, host_indices) + ) self.prefetch_loaded_tokens_by_reqid[req_id] = 0 logger.warning( "HiCache hybrid prefetch discarded req=%s completed=%d requested=%d", @@ -1773,8 +1921,27 @@ class UnifiedRadixCache(BasePrefixCache): def pop_prefetch_loaded_tokens(self, req_id: str) -> int: return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0) + def staged_prefetch_tokens(self, req_id: str) -> int: + """Tokens a staged buffer-mode prefetch would splice (0 = no hold); + surfaced by the scheduler as the request's host_hit_length.""" + if self.buffer_pipeline is None: + return 0 + return self.buffer_pipeline.staged_prefetch_tokens(req_id) + + def staged_prefetch_swa_tokens(self, req_id: str) -> int: + """SWA device tokens consuming a staged buffer-mode prefetch will + allocate; surfaced as the request's swa_host_hit_length.""" + if self.buffer_pipeline is None: + return 0 + return self.buffer_pipeline.staged_prefetch_swa_tokens(req_id) + def release_aborted_request(self, rid: str) -> None: self.prefetch_loaded_tokens_by_reqid.pop(rid, None) + if ( + self.buffer_pipeline is not None + and self.buffer_pipeline.release_aborted_staged(rid) + ): + return if rid not in self.ongoing_prefetch: return @@ -1793,13 +1960,73 @@ class UnifiedRadixCache(BasePrefixCache): completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) self._barrier_attn_groups() - self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) + if anchor_lock_params is not None: + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) del self.ongoing_prefetch[rid] + if self.buffer_pipeline is not None: + self.buffer_pipeline.pop_prefix_ctx(rid) self.cache_controller.append_host_mem_release( host_indices=host_indices[:completed_tokens], extra_pools=[x for xfers in comp_xfers.values() for x in xfers], ) - self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) + # Buffer mode granted occupancy at hit-alloc, sized to the bounce; + # cache mode reserved the requested span at enqueue. + self.cache_controller.prefetch_tokens_occupied -= self._prefetch_occupied_span( + prefetch_key, host_indices + ) + + def _invalidate_absent_from_hit_query(self, operation) -> None: + """Drop KV beliefs beyond the folded usable cut (rank-synced): the + next insert then re-writes the node (all pools), healing stale + positives and aux holes at the cut through one FULL check.""" + if self.host_memory_mode != "buffer_only": + return + chain = operation.all_hash_values + if chain is None: + return + self.storage_existence_cache.invalidate_beyond( + PoolName.KV, chain, keep_pages=operation.storage_hit_count // self.page_size + ) + + def _account_prefetch_outcome(self, operation, revoked: bool) -> None: + """Feed the cumulative prefetch-outcome counters at the (rank-synced) + query outcome: T = prompt tokens, L = requested, m = L3-miss.""" + requested = operation.stats_requested_tokens + if requested <= 0: + return + stats = self._prefetch_outcome_stats + hit = max(0, min(operation.storage_hit_count, requested)) + if revoked: + if hit > 0: + stats["revoked_insufficient"] += 1 + else: + stats["revoked_full_miss"] += 1 + miss = requested - hit + total = max(operation.stats_total_tokens, requested, 1) + stats["l3_demand_requests"] += 1 + stats["l1l2_miss_tokens"] += requested + stats["l3_miss_tokens"] += miss + stats["l3_demand_total_tokens"] += total + stats["l3_sum_rate_all"] += miss / total + stats["l3_sum_rate_main_weighted"] += (miss / requested) * total + + def prefetch_outcome_stats_snapshot(self) -> dict: + """Cumulative counters + instantaneous occupancy, in the schema + log_prefetch_stats consumers expect.""" + cc = self.cache_controller + cap = max(cc.prefetch_capacity_limit, 1) + return { + **self._prefetch_outcome_stats, + "occupancy_ratio": cc.prefetch_tokens_occupied / cap, + } + + def _prefetch_occupied_span(self, prefetch_key, host_indices) -> int: + """Occupancy units held by a prefetch: cache mode reserves the + requested span at enqueue; buffer mode grants at hit-alloc, sized + to the allocation (0 while still querying / parked).""" + if self.host_memory_mode == "buffer_only": + return len(host_indices) if host_indices is not None else 0 + return len(prefetch_key) def _revoke_pending_prefetch(self, req_id: str) -> None: info = self.ongoing_prefetch.pop(req_id, None) @@ -1809,17 +2036,27 @@ class UnifiedRadixCache(BasePrefixCache): last_host_node_id, prefetch_key, _host_indices, - _operation, + operation, anchor_lock_params, comp_xfers, ) = info + self._invalidate_absent_from_hit_query(operation) + if self.buffer_pipeline is not None: + self.buffer_pipeline.pop_prefix_ctx(req_id) cc = self.cache_controller cc.append_host_mem_release( extra_pools=[x for xfers in comp_xfers.values() for x in xfers] ) - self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) + if anchor_lock_params is not None: + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) + # Every revoke path runs before the bounce alloc, so buffer mode + # holds no occupancy here; post-alloc aborts go through + # release_aborted_request instead. + assert _host_indices is None or self.host_memory_mode != "buffer_only" cc.prefetch_tokens_occupied = max( - 0, cc.prefetch_tokens_occupied - len(prefetch_key) + 0, + cc.prefetch_tokens_occupied + - self._prefetch_occupied_span(prefetch_key, _host_indices), ) def _drain_storage_control_queues_impl( @@ -1842,56 +2079,100 @@ class UnifiedRadixCache(BasePrefixCache): drained += 1 yield item + buffer_mode = self.host_memory_mode == "buffer_only" + + def _try_alloc_storage_hit(operation) -> bool: + """Allocate the hit-sized bounce and launch the transfer. + Returns False when staging pressure defers the allocation + (buffer mode parks and retries; cache mode revokes).""" + req_id = operation.request_id + info = self.ongoing_prefetch.get(req_id) + if info is None: + return True # aborted/cleaned; nothing to retry + if operation.is_terminated(): + self._revoke_pending_prefetch(req_id) + return True + + if buffer_mode and cc.prefetch_rate_limited(): + # Pool is load-saturated: hold the KNOWN hit until staged + # prefetches ahead of us are consumed. The op stays in + # ongoing_prefetch, so wait_complete keeps gating admission. + return False + alloc_len = operation.storage_hit_count + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + self.evict_host(alloc_len) + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None and not buffer_mode: + # Memory-pressure fallback: a shorter page-aligned prefix. + # (Cache mode only — buffer mode parks for the full hit.) + available_size = cc.mem_pool_host.available_size() + alloc_len = min( + operation.storage_hit_count, + available_size - (available_size % self.page_size), + ) + if alloc_len >= self.prefetch_threshold: + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + if buffer_mode: + return False + self._revoke_pending_prefetch(req_id) + return True + + operation.storage_hit_count = alloc_len + operation.hash_value = operation.hash_value[: alloc_len // self.page_size] + operation.host_indices = host_indices + self.ongoing_prefetch[req_id] = info._replace(host_indices=host_indices) + if buffer_mode: + cc.prefetch_tokens_occupied += alloc_len + cc.prefetch_buffer.put(operation) + return True + def _drain_and_alloc_storage_hit(): + # Parked hits first (FIFO fairness with retries; buffer only). + if buffer_mode: + parked = self.buffer_pipeline.pending_hit_allocs + while parked: + if not _try_alloc_storage_hit(parked[0]): + break + parked.popleft() for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit): req_id = operation.request_id info = self.ongoing_prefetch.get(req_id) if info is None: - # request already aborted/cleaned up, skip + # Request already aborted/cleaned up; still flush the + # query's absent-hash feedback. + self._invalidate_absent_from_hit_query(operation) continue if operation.is_terminated(): - # request was aborted while the storage query was in flight + # Aborted while the storage query was in flight. self._revoke_pending_prefetch(req_id) continue if operation.storage_hit_count < self.prefetch_threshold: - # not to prefetch if not enough benefits + # Below-threshold hit: classify + feed the L3 miss + # accounting, then revoke (not enough benefit). + self._account_prefetch_outcome(operation, revoked=True) self._revoke_pending_prefetch(req_id) continue - - alloc_len = operation.storage_hit_count - host_indices = cc.mem_pool_host.alloc(alloc_len) - if host_indices is None: - self.evict_host(alloc_len) - host_indices = cc.mem_pool_host.alloc(alloc_len) - if host_indices is None: - # Memory-pressure fallback: a shorter page-aligned prefix. - available_size = cc.mem_pool_host.available_size() - alloc_len = min( - operation.storage_hit_count, - available_size - (available_size % self.page_size), - ) - if alloc_len >= self.prefetch_threshold: - host_indices = cc.mem_pool_host.alloc(alloc_len) - if host_indices is None: - self._revoke_pending_prefetch(req_id) - continue - - operation.storage_hit_count = alloc_len - operation.hash_value = operation.hash_value[ - : alloc_len // self.page_size - ] - operation.host_indices = host_indices - self.ongoing_prefetch[req_id] = info._replace(host_indices=host_indices) - cc.prefetch_buffer.put(operation) + self._invalidate_absent_from_hit_query(operation) + self._account_prefetch_outcome(operation, revoked=False) + if 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) def _drain_backup(): drained = 0 for operation in _drain_queue(cc.ack_backup_queue, n_backup): drained += 1 - entry = self.ongoing_backup.pop(operation.id, None) - if entry is not None: - node_id, lock_params = entry - self.dec_host_lock_ref(node_id, lock_params) + if buffer_mode: + # Storage write acked: free the staging. + self.buffer_pipeline.finish_storage_write_ack(operation.id) + else: + entry = self.ongoing_backup.pop(operation.id, None) + if entry is not None: + node_id, lock_params = entry + self.dec_host_lock_ref(node_id, lock_params) if ( log_metrics and self.enable_storage_metrics @@ -2054,6 +2335,9 @@ class UnifiedRadixCache(BasePrefixCache): logger.error("Failed to clear hierarchical cache storage backend: %s", e) return False if ok: + # L3 is empty now: every storage-presence belief is stale, and a + # retained positive would skip that page's backup forever. + self.storage_existence_cache.clear() logger.info("Hierarchical cache storage backend cleared successfully!") return ok @@ -2205,6 +2489,11 @@ class UnifiedRadixCache(BasePrefixCache): ack = cc.ack_load_queue.pop(0) ack.finish_event.synchronize() for ack_id in ack.node_ids: + if ( + self.buffer_pipeline is not None + and self.buffer_pipeline.try_finish_load_back(ack_id) + ): + continue node, lock_params, host_lock_params = self.ongoing_load_back.pop(ack_id) self.dec_lock_ref(node, lock_params) self.dec_host_lock_ref(node, host_lock_params) @@ -2233,7 +2522,10 @@ class UnifiedRadixCache(BasePrefixCache): params: InitLoadBackParams, ) -> tuple[torch.Tensor, NodeId]: """Prepare KV cache loading from host to device. - Returns (device_indices, last_node) tuple.""" + Returns (device_indices, last_node). Buffer mode dispatches to the + staged-prefetch consumption (BufferModePipeline.init_load_back).""" + if self.buffer_pipeline is not None: + return self.buffer_pipeline.init_load_back(params) best_match_node_id = params.best_match_node mem_quota = params.mem_quota req = params.req @@ -2306,10 +2598,16 @@ class UnifiedRadixCache(BasePrefixCache): extra_release_counts=extra_release_counts, log_metrics=True, ) + if self.buffer_pipeline is not None: + self.buffer_pipeline.flush_pending_writes() if self.enable_storage_metrics and self.storage_metrics_collector is not None: - self.storage_metrics_collector.log_storage_metrics( - self.cache_controller.storage_backend.get_stats() - ) + storage_metrics = self.cache_controller.storage_backend.get_stats() + if storage_metrics is None: + # Backends without native stats (e.g. file) still carry the + # controller-side prefetch outcome counters. + storage_metrics = StorageMetrics() + storage_metrics.prefetch_stats = self.prefetch_outcome_stats_snapshot() + self.storage_metrics_collector.log_storage_metrics(storage_metrics) def ready_to_load_host_cache(self) -> int: """Notify the cache controller to start the KV cache loading.""" @@ -2464,9 +2762,15 @@ class UnifiedRadixCache(BasePrefixCache): # Pass ongoing ops as lightweight (id, node_id) pairs so the tree core # can resolve + validate them without reaching into Controller state. - ongoing_write_through = [ - (nid, wt.node_id) for nid, wt in self.ongoing_write_through.items() - ] + if self.buffer_pipeline is not None: + ongoing_write_through = [ + (nid, entry.intent.node_id) + for nid, entry in self.buffer_pipeline.ongoing_write_through.items() + ] + else: + ongoing_write_through = [ + (nid, wt.node_id) for nid, wt in self.ongoing_write_through.items() + ] ongoing_load_back = [ (nid, lb.node_id) for nid, lb in self.ongoing_load_back.items() ] diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index adc15d97a..69175915f 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -1852,9 +1852,11 @@ class StorageMetricsCollector(_StatLoggerDIMixin): labels: Dict[str, str], ): from prometheus_client import Counter as _PromCounter + from prometheus_client import Gauge as _PromGauge from prometheus_client import Histogram as _PromHistogram Counter = self._counter_cls or _PromCounter + Gauge = self._gauge_cls or _PromGauge Histogram = self._histogram_cls or _PromHistogram self.labels = labels @@ -1871,6 +1873,21 @@ class StorageMetricsCollector(_StatLoggerDIMixin): labelnames=labels.keys(), ) + self.backup_dropped_tokens_total = Counter( + name="sglang:hicache_backup_dropped_tokens_total", + documentation="Buffer-mode backup tokens dropped by write-path rate " + "limiting (backlog cap or dropped-parent cascade).", + labelnames=labels.keys(), + ) + + self.prefetch_aux_alloc_failed_tokens_total = Counter( + name="sglang:hicache_prefetch_aux_alloc_failed_tokens_total", + documentation="Prefetch tokens abandoned because an aux pool " + "(e.g. the SWA trailing window) could not allocate host staging " + "— the whole storage fetch is forfeited, not just the aux part.", + labelnames=labels.keys(), + ) + bucket_io = [ 1, 5, @@ -1925,6 +1942,16 @@ class StorageMetricsCollector(_StatLoggerDIMixin): if backuped_tokens > 0: self.backuped_tokens_total.labels(**self.labels).inc(backuped_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) + + def log_prefetch_aux_alloc_failed_tokens(self, num_tokens: int): + if num_tokens > 0: + self.prefetch_aux_alloc_failed_tokens_total.labels(**self.labels).inc( + num_tokens + ) + def _log_histogram(self, histogram, data: Union[int, float]): histogram.labels(**self.labels).observe(data) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2502efe37..75d3543fe 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2668,14 +2668,22 @@ class ServerArgs: enable_hierarchical_cache: A[bool, "Enable hierarchical cache", NS("memory")] = ( False ) + hicache_host_memory_mode: A[ + str, + Arg( + help="Whether host memory is a persistent HiCache tier (cache) or a transient staging buffer between GPU and the storage backend (buffer_only). buffer_only requires --hicache-storage-backend.", + choices=["cache", "buffer_only"], + ), + NS("memory"), + ] = "cache" hicache_ratio: A[ Optional[float], - "The ratio of the size of host KV cache memory pool to the size of device pool. Defaults to 2.0, or 1.0 for host-pool decode retraction.", + "The ratio of the size of host KV cache memory pool to the size of device pool. Defaults to 2.0 in cache mode, 1.2 in buffer_only mode, or 1.0 for host-pool decode retraction.", NS("memory"), ] = None hicache_size: A[ int, - "The size of host KV cache memory pool in gigabytes, which will override the hicache_ratio if set.", + "The size of host KV cache memory pool in gigabytes. Overrides --hicache-ratio in either host memory mode.", NS("memory"), ] = 0 hicache_write_policy: A[ @@ -2714,6 +2722,7 @@ class ServerArgs: help="The storage backend for hierarchical KV cache. Built-in backends: file, mooncake, hf3fs, nixl, aibrix. For dynamic backend, use --hicache-storage-backend-extra-config to specify: backend_name (custom name), module_path (Python module path), class_name (backend class name).", choices=[ "file", + "sim", "mooncake", "hf3fs", "nixl", @@ -7393,8 +7402,20 @@ class ServerArgs: ) def _handle_hicache_ratio_default(self): + """Default the host/device ratio per host memory mode. + + Runs before the dummy-model boundary: direct HostKVCache consumers + (unit fixtures, dummy-model launches) must never see a None ratio. + buffer_only stages in flight rather than retaining, so it needs only + enough to cover the write backlog plus parked prefetches. + + A decode server keeps the ratio unset here: kv_cache_builder resolves + it against the retraction-backup backend (1.0 for host_pool, else 2.0). + """ if self.hicache_ratio is None and self.disaggregation_mode != "decode": - self.hicache_ratio = 2.0 + self.hicache_ratio = ( + 1.2 if self.hicache_host_memory_mode == "buffer_only" else 2.0 + ) def _handle_hicache(self): """Normalize hicache-related knobs into a valid runtime configuration. @@ -7414,6 +7435,8 @@ class ServerArgs: ): return + self._validate_hicache_host_memory_mode() + # Step 1: Initial layout-io compatibility normalization. self._resolve_layout_io_compatibility() @@ -7423,6 +7446,51 @@ class ServerArgs: # Step 3: DCP compatibility for the L2 (device<->host) path. self._resolve_hicache_dcp_compatibility() + def _validate_hicache_host_memory_mode(self): + if self.hicache_host_memory_mode not in ("cache", "buffer_only"): + raise ValueError( + "hicache_host_memory_mode must be 'cache' or 'buffer_only', " + f"got {self.hicache_host_memory_mode!r}" + ) + + # Both modes are defaulted upstream (a decode server resolves the + # ratio later, in kv_cache_builder), so this fires only if that + # defaulting regresses -- never build an unsized host pool. + if ( + self.hicache_size <= 0 + and self.hicache_ratio is None + and self.disaggregation_mode != "decode" + ): + raise ValueError( + f"--hicache-host-memory-mode {self.hicache_host_memory_mode} " + "requires a host pool size: pass --hicache-size or " + "--hicache-ratio." + ) + + if self.hicache_host_memory_mode == "cache": + return + + if self.hicache_storage_backend is None: + raise ValueError( + "--hicache-host-memory-mode buffer_only requires a storage backend " + "(--hicache-storage-backend): host memory is only a staging buffer " + "and all cached data lives in storage." + ) + if self.hicache_write_policy == "write_back": + raise ValueError( + "--hicache-host-memory-mode buffer_only does not support " + "--hicache-write-policy write_back; use write_through or " + "write_through_selective." + ) + if self.disaggregation_mode == "decode": + raise ValueError( + "--hicache-host-memory-mode buffer_only is not supported on " + "decode instances: the decode-side prefetch and offload paths " + "bypass the buffer-mode pipeline, fetching without its prefix " + "context and never consuming its staged holds. Prefill " + "instances share the standard scheduler path and are supported." + ) + def _resolve_hicache_dcp_compatibility(self): if self.dcp_size <= 1 or not self.enable_hierarchical_cache: return 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 932eae344..5155fc87b 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 @@ -92,8 +92,8 @@ from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=16, stage="base-b", runner_config="1-gpu-small") -register_amd_ci(est_time=16, suite="stage-b-test-1-gpu-small-amd") +register_cuda_ci(est_time=50, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=50, suite="stage-b-test-1-gpu-small-amd") @dataclass(frozen=True) @@ -595,39 +595,20 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): return leaf def _init_hicache(self, cache, *, write_policy: str = "write_through"): - import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler - - # Wrap the host-pool factory (not MHATokenToKVPoolHost directly) - # because the assembler picks between MHATokenToKVPoolHost and - # AsymmetricMHATokenToKVPoolHost via get_mha_host_pool_cls(device_pool). - orig_get_mha_host_pool_cls = assembler.get_mha_host_pool_cls - - def get_mha_host_pool_cls_wrapper(device_pool): - host_pool_cls = orig_get_mha_host_pool_cls(device_pool) - - def kv_host_pool_wrapper(*args, **kwargs): - kwargs["pin_memory"] = False - return host_pool_cls(*args, **kwargs) - - return kv_host_pool_wrapper - - patcher = mock.patch.object( - assembler, - "get_mha_host_pool_cls", - side_effect=get_mha_host_pool_cls_wrapper, - ) - patcher.start() - self.addCleanup(patcher.stop) - + # Production config: kernel IO backend + page_first layout with + # PINNED host pools (kernel GPU DMA requires them). Pools track their + # cudaHostRegister'd pointers and unregister on destroy()/GC, so the + # many fixtures sharing this process cannot collide on recycled + # address ranges (rc=712). server_args = ServerArgs( model_path="dummy", page_size=self.cfg.page_size, - hicache_io_backend="direct", - hicache_mem_layout="page_first_direct", + hicache_io_backend="kernel", hicache_write_policy=write_policy, ) set_global_server_args_for_scheduler(server_args) cache.init_hicache(server_args, cache.cache_init_params) + self.addCleanup(cache.release_host_resources) cache.write_through_threshold = 1 << 30 cache.load_back_threshold = 0 @@ -2468,28 +2449,87 @@ class UnifiedRadixCacheSuite: for n in self._path_chain(cache, node): cache.write_backup_storage(n.id) + def _ongoing_l3_backups(self, cache): + """Storage writes in flight (buffer mode tracks them on the pipeline).""" + if cache.buffer_pipeline is not None: + return cache.buffer_pipeline.ongoing_backup + return cache.ongoing_backup + def _flush_l3_backups(self, cache, timeout: float = 10.0): """Wait for backup threads to finish, then drain acks (release locks).""" deadline = time.time() + timeout - while cache.ongoing_backup and time.time() < deadline: + while self._ongoing_l3_backups(cache) and time.time() < deadline: cache.drain_storage_control_queues() - if cache.ongoing_backup: + if self._ongoing_l3_backups(cache): time.sleep(0.01) cache.drain_storage_control_queues() - self.assertFalse(cache.ongoing_backup, "L3 backups did not complete in time") + self.assertFalse( + self._ongoing_l3_backups(cache), "L3 backups did not complete in time" + ) def _run_prefetch_to_completion(self, cache, req_id, timeout: float = 10.0): deadline = time.time() + timeout while time.time() < deadline: # Host memory is reserved (and IO started) by the scheduler-thread - # drain once the L3 hit count is known, so pump it like the real - # scheduler loop does (check_hicache_events before progress checks). - cache.drain_storage_control_queues() + # drain once the L3 hit count is known, and the buffer-mode fill + # commits at its H2D ack in loading_check — pump the full event + # round like the real scheduler loop does. + cache.check_hicache_events() if cache.check_prefetch_progress(req_id): + # Buffer mode parks a completed fetch as a staged prefetch; + # consume it like the PrefillAdder would so callers see the + # span tree-resident. + if ( + cache.buffer_pipeline is not None + and cache.buffer_pipeline.has_staged(req_id) + ): + self._consume_staged_prefetch(cache, req_id, timeout=timeout) return time.sleep(0.01) self.fail(f"prefetch {req_id} did not complete in time") + def _consume_staged_prefetch( + self, cache, req_id, prefix_len=None, prefix_indices=None, timeout: float = 10.0 + ): + """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).""" + from sglang.srt.mem_cache.base_prefix_cache import InitLoadBackParams + + f = cache.buffer_pipeline.staged_prefetches[req_id] + if prefix_len is None: + prefix_len = f.matched_len + req = mock.Mock() + req.rid = req_id + if prefix_indices is not None: + # Spliceable mid-anchor consumption publishes value=cat(prefix, + # fill) — the real device prefix is required (zeros would insert + # bogus slots into the tree). + assert len(prefix_indices) == prefix_len + req.prefix_indices = prefix_indices + else: + req.prefix_indices = torch.zeros( + prefix_len, + dtype=torch.int64, + device=cache.tree_core.empty_match_result.device_indices.device, + ) + req.last_node = cache.root_node.id + new_indices, _last_node = cache.init_load_back( + InitLoadBackParams( + best_match_node=None, host_hit_length=f.num_tokens, req=req + ) + ) + # Batch formation flushes the queued load into the batch's producer. + cache.ready_to_load_host_cache() + self._pump_hicache_until( + cache, + lambda: not cache.buffer_pipeline.ongoing_buffer_load_back, + "staged-prefetch consumption did not commit", + timeout=timeout, + ) + return new_indices + def _all_page_hashes(self, cache, node): hashes = [] for n in self._path_chain(cache, node): @@ -2606,6 +2646,494 @@ class UnifiedRadixCacheSuite: self.assertTrue(torch.equal(loaded_v, expected_v)) cons.sanity_check() + # ================================================================ + # Buffer-only host memory mode (host = transient staging, L3 = cache) + # ================================================================ + + def _init_buffer_hicache( + self, + cache, + storage_dir, + prefetch_policy: str = "wait_complete", + storage_extra: Optional[dict] = None, + ): + if self.cfg.has_mamba: + self.skipTest( + "buffer_only is FULL/SWA-only (no Mamba state-handoff channel " + "on the admission-time load-back read path)" + ) + self._init_hicache( + cache, + storage_backend="file", + storage_dir=storage_dir, + prefetch_threshold=1, + host_memory_mode="buffer_only", + prefetch_policy=prefetch_policy, + storage_extra=storage_extra, + ) + + def _pump_hicache_until(self, cache, cond, msg, timeout: float = 10.0): + deadline = time.time() + timeout + while time.time() < deadline: + cache.check_hicache_events() + if cond(): + return + time.sleep(0.01) + self.fail(msg) + + def _host_avail_sizes(self, cache): + group = cache.cache_controller.mem_pool_host + return {entry.name: entry.host_pool.available_size() for entry in group.entries} + + def _storage_exists_count(self, cache, page_hashes, pool_transfers=None): + """Ground-truth longest-prefix existence count from the backend, + folded across pools like the prefetch hit query.""" + from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo + + backend = cache.cache_controller.storage_backend + extra_info = HiCacheStorageExtraInfo(prefix_keys=None) + if pool_transfers: + result = backend.batch_exists_v2(page_hashes, pool_transfers, extra_info) + return min(result.kv_hit_pages, len(page_hashes)) + return backend.batch_exists(page_hashes, extra_info) + + def _buffer_backup_and_wait(self, cache, node): + # Parent-first over the whole path, mirroring the production trigger + # (_inc_hit_count fires per matched node on the insert walk): the SWA + # component may have split the leaf at the window boundary. + pipeline = cache.buffer_pipeline + for n in self._path_chain(cache, node): + pipeline.enqueue_backup_intent(n) + self.assertIn(n.id, pipeline.inflight_backup_node_ids) + self._pump_hicache_until( + cache, + lambda: not pipeline.inflight_backup_node_ids + and not pipeline.ongoing_backup, + "buffer backup pipeline did not drain", + ) + + def _produce_buffer_l3(self, storage_dir, seq, marker=None): + """Producer tree in buffer mode: insert seq and push it to L3.""" + prod, prod_alloc, prod_rtp = build_fixture(self.cfg) + self._init_buffer_hicache(prod, storage_dir) + self._insert(prod, prod_alloc, prod_rtp, seq) + leaf = prod.resolve_node_handle( + prod.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq))) + ).last_device_node + ) + expected = None + if marker is not None: + m = prod.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self._fill_full_kv(prod_alloc, m.device_indices, marker=marker) + expected = self._snapshot_full_kv(prod_alloc, m.device_indices) + self._buffer_backup_and_wait(prod, leaf) + return leaf, expected + + def _buffer_swa_seq(self, min_pages=4): + """Sequence long enough for SWA prefetch (one full window + 1).""" + num_pages = min_pages + if self.cfg.has_swa: + sw_pages = ( + self.cfg.sliding_window_size + self.cfg.page_size - 1 + ) // self.cfg.page_size + num_pages = max(num_pages, sw_pages + 1) + return self._make_seq(1, num_pages) + + def test_buffer_only_write_path_roundtrip(self): + """Write path end to end: admission -> D2H staging -> storage write + -> free. Staging and locks fully released, pages stored under every + pool namespace, beliefs registered, re-hits skipped.""" + self._skip_unsupported_hicache_test() + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + self._init_buffer_hicache(cache, storage_dir) + avail0 = self._host_avail_sizes(cache) + + seq_a = self._make_seq(1, 2) + seq_ab = seq_a + self._make_seq(500, 2) + self._insert(cache, allocator, req_to_token_pool, seq_a) + self._insert(cache, allocator, req_to_token_pool, seq_ab) + leaf = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq_ab))) + ).last_device_node + ) + chain = self._path_chain(cache, leaf) + + self._buffer_backup_and_wait(cache, leaf) + self.assertFalse(leaf.backuped) + self.assertEqual(leaf.component_data[ComponentType.FULL].lock_ref, 0) + self.assertEqual(self._host_avail_sizes(cache), avail0) + self.assertEqual(cache.buffer_pipeline.write_backlog_tokens_, 0) + page_hashes = self._all_page_hashes(cache, leaf) + self.assertEqual( + self._storage_exists_count( + cache, + page_hashes, + cache.buffer_pipeline._build_aux_staging_transfers(leaf), + ), + len(page_hashes), + ) + for n in chain: + self.assertTrue( + cache.storage_existence_cache.contains_all(PoolName.KV, n.hash_value) + ) + # Re-hit absorbed by the (FULL-focused) belief skip. + cache.buffer_pipeline.enqueue_backup_intent(leaf) + self.assertNotIn(leaf.id, cache.buffer_pipeline.inflight_backup_node_ids) + cache.sanity_check() + + def test_buffer_only_read_path_roundtrip(self): + """Read path end to end: prefetch -> staged (host bounce only, + nothing device-side, unmatchable, stable readiness, counters fed) -> + admission-time load-back against a saturated-but-evictable pool + (evict-before-alloc) publishing pre-ack -> ack frees the bounce. + Data bytes match the producer's; no CPU-tier KV events anywhere; + declines feed the outcome counters.""" + self._skip_unsupported_hicache_test() + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + + seq = self._buffer_swa_seq() + _, (expected_k, expected_v) = self._produce_buffer_l3( + storage_dir, seq, marker=7 + ) + + cons, cons_alloc, cons_rtp = build_fixture( + self.cfg, enable_kv_cache_events=True + ) + self._init_buffer_hicache(cons, storage_dir) + cons.take_events() + avail0 = self._host_avail_sizes(cons) + dev_avail0 = cons.token_to_kv_pool_allocator.available_size() + stats = cons._prefetch_outcome_stats + + req_id = "buffer-read-roundtrip" + cons.prefetch_from_storage( + req_id, cons.root_node.id, array("q", seq), None, None + ) + self.assertEqual((stats["attempts"], stats["issued"]), (1, 1)) + self._pump_hicache_until( + cons, + lambda: cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id), + "prefetch did not stage", + ) + # Staged: bounce occupies host staging; nothing device-side; span + # unmatchable; readiness stable; hit accounting reported once. + self.assertNotEqual(self._host_avail_sizes(cons), avail0) + self.assertEqual(cons.token_to_kv_pool_allocator.available_size(), dev_avail0) + self.assertEqual( + len( + cons.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq))) + ).device_indices + ), + 0, + ) + self.assertTrue(cons.check_prefetch_progress(req_id)) + self.assertEqual(cons.pop_prefetch_loaded_tokens(req_id), len(seq)) + self.assertEqual(stats["l3_demand_requests"], 1) + self.assertEqual(stats["l3_miss_tokens"], 0) + + # Saturate the device pool with unrelated evictable spans: the + # load-back must evict, not degrade to recompute (run-2 regression). + def _avail(): + if cons.supports_swa(): + return cons.token_to_kv_pool_allocator.full_available_size() + return cons.token_to_kv_pool_allocator.available_size() + + filler_base = 90000 + while _avail() >= self.cfg.page_size: + pages = max(1, min(2048, _avail()) // self.cfg.page_size) + self._insert(cons, cons_alloc, cons_rtp, self._make_seq(filler_base, pages)) + filler_base += 1000 + self.assertLess(_avail(), len(seq), "pool not saturated") + + # Consume WITHOUT pumping the ack: published (matchable, same slot + # ids) before any ack lands. + from sglang.srt.mem_cache.base_prefix_cache import InitLoadBackParams + + held = cons.buffer_pipeline.staged_prefetches[req_id] + req = mock.Mock() + req.rid = req_id + req.last_node = cons.root_node.id + req.prefix_indices = torch.zeros( + held.matched_len, + dtype=torch.int64, + device=cons.tree_core.empty_match_result.device_indices.device, + ) + spliced, _last = cons.init_load_back( + InitLoadBackParams( + best_match_node=None, host_hit_length=held.num_tokens, req=req + ) + ) + self.assertEqual(int(spliced.numel()), len(seq)) + cons.ready_to_load_host_cache() + m = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertTrue(torch.equal(m.device_indices, spliced)) + + # Ack: bounce freed, beliefs fed, tree holds no host values, and the + # loaded KV bytes equal the producer's. + self._pump_hicache_until( + cons, + lambda: not cons.buffer_pipeline.ongoing_buffer_load_back + and self._host_avail_sizes(cons) == avail0, + "load-back ack did not free the bounce", + ) + mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertEqual(mc.host_hit_length, 0) + self.assertEqual(len(mc.device_indices), len(seq)) + leaf = cons.resolve_node_handle(mc.last_device_node) + for cur in self._path_chain(cons, leaf): + for cd in cur.component_data: + self.assertIsNone(cd.host_value) + self.assertTrue( + cons.storage_existence_cache.contains_all( + PoolName.KV, self._all_page_hashes(cons, leaf) + ) + ) + loaded_k, loaded_v = self._snapshot_full_kv(cons_alloc, mc.device_indices) + self.assertTrue(torch.equal(loaded_k, expected_k)) + self.assertTrue(torch.equal(loaded_v, expected_v)) + self.assertEqual(cons.cache_controller.prefetch_tokens_occupied, 0) + cpu_events = [ + e + for e in cons.take_events() + if isinstance(e, (BlockStored, BlockRemoved)) + and e.medium == StorageMedium.CPU + ] + self.assertEqual(cpu_events, []) + + self.assertIn("occupancy_ratio", cons.prefetch_outcome_stats_snapshot()) + cons.sanity_check() + + def test_buffer_load_back_swa_window_charged_at_admission(self): + """Admission contract: a request the SWA budget gate accepts must be + allocatable at batch time (_swa_reserved_tokens: "an admitted request + cannot OOM"). Regression: buffer mode surfaced a staged prefetch as + host_hit_length only, so the gate never charged the SWA window that + consumption (init_load_back -> cc.load) allocates and the request + lock pins; with the rest of the SWA pool batch-held, the batch alloc + fell short by up to one window and raised the fail-loud prefill OOM + (prod scheduler crash, 2026-08-17).""" + self._skip_unsupported_hicache_test() + if not self.cfg.has_swa: + self.skipTest("SWA-specific admission accounting") + from sglang.srt.mem_cache.allocation import ( + alloc_paged_token_slots_extend, + alloc_token_slots, + ) + from sglang.srt.mem_cache.base_prefix_cache import InitLoadBackParams + + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + ps = self.cfg.page_size + window = self.cfg.sliding_window_size + seq = self._buffer_swa_seq() # one full window + 1 page + self._produce_buffer_l3(storage_dir, seq) + + cons, cons_alloc, _ = build_fixture(self.cfg) + self._init_buffer_hicache(cons, storage_dir) + req_id = "buffer-swa-admission-oom" + cons.prefetch_from_storage( + req_id, cons.root_node.id, 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", + ) + + # Batch-held SWA (chunk allocs, decode windows) is neither free nor + # evictable: leave one token less than window + extend_need, enough + # for an un-charged gate to accept. + extend_need = 2 * ps + 1 + max_new = 8 + self.assertIsNotNone( + cons_alloc.swa_attn_allocator.alloc( + cons_alloc.swa_available_size() - (window + extend_need - 1) + ) + ) + + # The adder's SWA gate for this request (_swa_budget_for_req). + surfaced_swa_hit = cons.staged_prefetch_swa_tokens(req_id) + reserved = ( + max(extend_need - window, 0) + + min(extend_need + max_new, window) + + ps + + (surfaced_swa_hit + ps - 1) // ps * ps + ) + budget = cons_alloc.swa_available_size() + cons.swa_evictable_size() + + # Consume at admission (init_load_back + request lock). + held = cons.buffer_pipeline.staged_prefetches[req_id] + req = mock.Mock() + req.rid = req_id + req.last_node = cons.root_node.id + req.prefix_indices = torch.zeros( + held.matched_len, + dtype=torch.int64, + device=cons.tree_core.empty_match_result.device_indices.device, + ) + spliced, last_node = cons.init_load_back( + InitLoadBackParams( + best_match_node=None, host_hit_length=held.num_tokens, req=req + ) + ) + self.assertEqual(int(spliced.numel()), len(seq), "load-back degraded") + cons.ready_to_load_host_cache() + cons.inc_lock_ref(last_node) # _req_inc_lock_ref + + self.assertEqual(cons.swa_evictable_size(), 0) # window is protected + # FULL stays roomy: only SWA can fail below. + self.assertGreaterEqual(cons_alloc.full_available_size(), extend_need + ps) + + if reserved <= budget: + try: + if ps == 1: + alloc_token_slots(cons, extend_need) + else: # paged batch path — the production crash site + prefix_len = int(spliced.numel()) + prefix_t = torch.tensor( + [prefix_len], dtype=torch.int64, device=spliced.device + ) + seq_t = torch.tensor( + [prefix_len + extend_need], + dtype=torch.int64, + device=spliced.device, + ) + alloc_paged_token_slots_extend( + cons, + prefix_t, + prefix_t.cpu(), + seq_t, + seq_t.cpu(), + spliced[-1:], + extend_need, + ) + except RuntimeError as e: + self.fail( + f"gate admitted (reserved={reserved} <= budget={budget}) " + f"but the batch alloc OOMed: {e}" + ) + else: + # Rejection must come from the surfaced window charge. + self.assertGreaterEqual(surfaced_swa_hit, window) + + 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.""" + self._skip_unsupported_hicache_test() + if not self.cfg.has_swa: + self.skipTest("requires an SWA component") + window = self.cfg.sliding_window_size + if window <= self.cfg.page_size: + self.skipTest("window fits in one page") + sw_pages = (window + self.cfg.page_size - 1) // self.cfg.page_size + 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) + self._produce_buffer_l3(storage_dir, seq, marker=5) + cons, _, _ = build_fixture(self.cfg) + self._init_buffer_hicache(cons, storage_dir) + cons.prefetch_from_storage( + "short-req", cons.root_node.id, array("q", seq), None, None + ) + self._run_prefetch_to_completion(cons, "short-req") + 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( + cons.resolve_node_handle(mc.last_device_node) + .component_data[ComponentType.SWA] + .value + ) + cons.sanity_check() + + # 2. Mid-tree continuation shorter than the window: the staged + # prefetch must carry an SWA transfer (not a KV-only degrade). + 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) + 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)))) + cons2.prefetch_from_storage( + "subwin-req", + m.last_device_node, + array("q", seq_ab[len(seq_a) :]), + cons2.get_last_hash_value(m.last_device_node), + None, + matched_prefix_tokens=list(seq_a), + ) + self._pump_hicache_until( + cons2, + lambda: cons2.check_prefetch_progress("subwin-req") + and cons2.buffer_pipeline.has_staged("subwin-req"), + "sub-window prefetch did not stage", + ) + self.assertTrue( + any( + t.name == PoolName.SWA + for t in cons2.buffer_pipeline.staged_prefetches[ + "subwin-req" + ].aux_xfers + ), + "sub-window fetch degraded to KV-only", + ) + spliced = self._consume_staged_prefetch( + cons2, "subwin-req", prefix_indices=m.device_indices + ) + self.assertEqual(int(spliced.numel()), len(seq_ab) - len(seq_a)) + self.assertEqual( + len( + cons2.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq_ab))) + ).device_indices + ), + len(seq_ab), + ) + 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] + self._produce_buffer_l3(storage_dir, stored, marker=6) + cons3, _, _ = build_fixture(self.cfg) + self._init_buffer_hicache(cons3, storage_dir) + avail3 = self._host_avail_sizes(cons3) + cons3.prefetch_from_storage( + "partial-req", cons3.root_node.id, array("q", full), None, None + ) + self._run_prefetch_to_completion(cons3, "partial-req") + cons3.drain_storage_control_queues() + self.assertEqual( + len( + cons3.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", stored))) + ).device_indices + ), + len(stored), + "partial hit lost its SWA window", + ) + self.assertEqual(self._host_avail_sizes(cons3), avail3) + cons3.sanity_check() + # ---------- TP consistency for SWA prefetch (all-or-nothing) ---------- def _patch_tp_all_reduce(self, cache, drop_swa: bool): @@ -2898,44 +3426,9 @@ class UnifiedRadixCacheSuite: storage_dir: Optional[str] = None, prefetch_threshold: Optional[int] = None, prefetch_policy: str = "wait_complete", + host_memory_mode: str = "cache", + storage_extra: Optional[dict] = None, ): - import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler - - # See _init_hicache: wrap the factory rather than MHATokenToKVPoolHost - # directly so the pin_memory=False override applies to both - # MHATokenToKVPoolHost and AsymmetricMHATokenToKVPoolHost. - orig_get_mha_host_pool_cls = assembler.get_mha_host_pool_cls - orig_mamba_host_pool = assembler.MambaPoolHost - - def get_mha_host_pool_cls_wrapper(device_pool): - host_pool_cls = orig_get_mha_host_pool_cls(device_pool) - - def kv_host_pool_wrapper(*args, **kwargs): - kwargs["pin_memory"] = False - return host_pool_cls(*args, **kwargs) - - return kv_host_pool_wrapper - - def mamba_host_pool_wrapper(*args, **kwargs): - kwargs["pin_memory"] = False - return orig_mamba_host_pool(*args, **kwargs) - - patchers = [ - mock.patch.object( - assembler, - "get_mha_host_pool_cls", - side_effect=get_mha_host_pool_cls_wrapper, - ), - mock.patch.object( - assembler, - "MambaPoolHost", - side_effect=mamba_host_pool_wrapper, - ), - ] - for patcher in patchers: - patcher.start() - self.addCleanup(patcher.stop) - storage_extra_config = None if storage_backend == "file": from sglang.srt.runtime_context import get_parallel @@ -2957,22 +3450,26 @@ class UnifiedRadixCacheSuite: extra = {} if prefetch_threshold is not None: extra["prefetch_threshold"] = prefetch_threshold + if storage_extra: + extra.update(storage_extra) storage_extra_config = json.dumps(extra) if extra else None + # See _init_hicache: production kernel IO backend, pinned pools. server_args = ServerArgs( model_path="dummy", page_size=self.cfg.page_size, - hicache_io_backend="direct", - hicache_mem_layout="page_first_direct", + hicache_io_backend="kernel", hicache_write_policy=write_policy, hicache_storage_backend=storage_backend, hicache_storage_backend_extra_config=storage_extra_config, hicache_storage_prefetch_policy=prefetch_policy, + hicache_host_memory_mode=host_memory_mode, ) # See build_fixture for why _mamba_cache_chunk_size is preset. 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(cache.release_host_resources) cache.write_through_threshold = 1 << 30 cache.load_back_threshold = 0 if storage_backend is not None: @@ -3020,8 +3517,8 @@ class UnifiedRadixCacheSuite: if node is not cache.root_node: self._backup_node(cache, node) - def _load_back_node(self, cache, node): - loaded = cache.load_back(node.id) + def _load_back_node(self, cache, node, req=None): + loaded = cache.load_back(node.id, req=req) self.assertTrue(loaded) producer_id = cache.ready_to_load_host_cache() self.assertNotEqual(producer_id, -1) @@ -3729,10 +4226,14 @@ class UnifiedRadixCacheSuite: cache, _, _ = self._build_hicache_fixture() sw = cache.sliding_window_size swa = cache.components[ComponentType.SWA] - # below a full window -> does not participate, no alloc - prep = swa.prepare_prefetch(cache.root_node.id, prefetch_tokens=sw - 1) + # zero-length prefetch -> does not participate, no alloc + prep = swa.prepare_prefetch(cache.root_node.id, prefetch_tokens=0) self.assertFalse(prep.alloc_failed) self.assertIsNone(prep.host_indices) + # 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.id, 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.id, prefetch_tokens=sw) self.assertEqual(int(prep.host_indices.numel()), sw) @@ -6293,6 +6794,7 @@ class TestPrefetchCommitOrdering(CustomTestCase): cache = mock.MagicMock() cache.page_size = 1 cache.enable_storage_metrics = False + cache.buffer_pipeline = None # cache-mode commit path walk_action = object() insert_result = mock.MagicMock() insert_result.cache_actions = [walk_action]