[HiCache] Buffer-only mode for HiCache host memory layer (#34798)
This commit is contained in:
@@ -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()
|
||||||
@@ -17,7 +17,7 @@ import logging
|
|||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
from queue import Empty, Queue
|
from queue import Empty, Queue
|
||||||
from typing import TYPE_CHECKING, List, NamedTuple, Optional
|
from typing import TYPE_CHECKING, Callable, List, NamedTuple, Optional
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
@@ -203,6 +203,13 @@ class StorageOperation:
|
|||||||
self.completed_tokens = 0
|
self.completed_tokens = 0
|
||||||
self.hash_value = hash_value if hash_value is not None else []
|
self.hash_value = hash_value if hash_value is not None else []
|
||||||
self.prefix_keys = prefix_keys
|
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
|
self.id = StorageOperation.counter
|
||||||
StorageOperation.counter += 1
|
StorageOperation.counter += 1
|
||||||
@@ -211,6 +218,15 @@ class StorageOperation:
|
|||||||
return self.id < other.id
|
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):
|
class PrefetchOperation(StorageOperation):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -262,8 +278,10 @@ class HiCacheController:
|
|||||||
model_name: Optional[str] = None,
|
model_name: Optional[str] = None,
|
||||||
storage_backend_extra_config: Optional[dict] = None,
|
storage_backend_extra_config: Optional[dict] = None,
|
||||||
enable_storage_metrics: bool = False,
|
enable_storage_metrics: bool = False,
|
||||||
|
host_memory_mode: str = "cache",
|
||||||
):
|
):
|
||||||
self.tp_group = tp_group
|
self.tp_group = tp_group
|
||||||
|
self.host_memory_mode = host_memory_mode
|
||||||
self.attn_cp_group = attn_cp_group
|
self.attn_cp_group = attn_cp_group
|
||||||
self.attn_tp_group = attn_tp_group
|
self.attn_tp_group = attn_tp_group
|
||||||
self.pp_group = pp_group
|
self.pp_group = pp_group
|
||||||
@@ -283,6 +301,9 @@ class HiCacheController:
|
|||||||
self.storage_backend = None
|
self.storage_backend = None
|
||||||
self.storage_backend_type = None
|
self.storage_backend_type = None
|
||||||
self.enable_storage_metrics = enable_storage_metrics
|
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).
|
# Draft KV pool support (best-effort piggyback on target L2/L3 ops).
|
||||||
self.has_draft = False
|
self.has_draft = False
|
||||||
@@ -501,8 +522,16 @@ class HiCacheController:
|
|||||||
self.enable_storage = True
|
self.enable_storage = True
|
||||||
# todo: threshold policy for prefetching
|
# todo: threshold policy for prefetching
|
||||||
self.prefetch_threshold = max(prefetch_threshold, self.page_size)
|
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.
|
if self.host_memory_mode == "buffer_only":
|
||||||
self.prefetch_capacity_limit = int(0.5 * self.mem_pool_host.size)
|
# 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
|
# tracking the number of tokens locked in prefetching, updated by the main scheduler thread
|
||||||
self.prefetch_tokens_occupied = 0
|
self.prefetch_tokens_occupied = 0
|
||||||
|
|
||||||
@@ -1050,6 +1079,16 @@ class HiCacheController:
|
|||||||
"""
|
"""
|
||||||
Rate limit the prefetching operations to avoid overwhelming the storage backend.
|
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
|
# cancel prefetch if too much memory is occupied
|
||||||
if self.prefetch_tokens_occupied >= self.prefetch_capacity_limit:
|
if self.prefetch_tokens_occupied >= self.prefetch_capacity_limit:
|
||||||
return True
|
return True
|
||||||
@@ -1066,6 +1105,7 @@ class HiCacheController:
|
|||||||
page_hashes = self.get_hash_str(
|
page_hashes = self.get_hash_str(
|
||||||
tokens_to_fetch, last_hash, page_size=self.page_size
|
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):
|
for start in range(0, len(page_hashes), STORAGE_BATCH_SIZE):
|
||||||
batch_hashes = page_hashes[start : start + STORAGE_BATCH_SIZE]
|
batch_hashes = page_hashes[start : start + STORAGE_BATCH_SIZE]
|
||||||
|
|||||||
@@ -2698,8 +2698,28 @@ class Scheduler(
|
|||||||
if self.enable_hicache_storage:
|
if self.enable_hicache_storage:
|
||||||
req.init_next_round_input(self.tree_cache, cow_mamba=False)
|
req.init_next_round_input(self.tree_cache, cow_mamba=False)
|
||||||
tree_cache = self.tree_cache
|
tree_cache = self.tree_cache
|
||||||
if tree_cache.is_backuped(req.last_host_node) or tree_cache.is_root(
|
buffer_mode = self.server_args.hicache_host_memory_mode == "buffer_only"
|
||||||
req.last_host_node
|
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
|
matched_len = len(req.prefix_indices) + req.host_hit_length
|
||||||
match_end = req._compute_max_prefix_len(
|
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]
|
new_input_tokens = req.full_untruncated_fill_ids[matched_len:match_end]
|
||||||
prefix_keys = (
|
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
|
if tree_cache.hicache_storage_pass_prefix_keys
|
||||||
else None
|
else None
|
||||||
)
|
)
|
||||||
tree_cache.prefetch_from_storage(
|
tree_cache.prefetch_from_storage(
|
||||||
req.rid,
|
req.rid,
|
||||||
req.last_host_node,
|
last_host_node,
|
||||||
new_input_tokens,
|
new_input_tokens,
|
||||||
tree_cache.get_last_hash_value(req.last_host_node),
|
tree_cache.get_last_hash_value(last_host_node),
|
||||||
prefix_keys,
|
prefix_keys,
|
||||||
|
matched_prefix_tokens=req.full_untruncated_fill_ids[:matched_len],
|
||||||
)
|
)
|
||||||
|
|
||||||
def _add_request_to_queue(self, req: Req, is_retracted: bool = False):
|
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.storage_hit_length = loaded_tokens
|
||||||
|
|
||||||
req.init_next_round_input(self.tree_cache)
|
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(
|
res = adder.add_one_req(
|
||||||
req,
|
req,
|
||||||
has_chunked_req=(self.chunked_req is not None),
|
has_chunked_req=(self.chunked_req is not None),
|
||||||
@@ -4146,6 +4184,11 @@ class Scheduler(
|
|||||||
if tc.enable_storage:
|
if tc.enable_storage:
|
||||||
idle &= len(tc.ongoing_prefetch) == 0
|
idle &= len(tc.ongoing_prefetch) == 0
|
||||||
idle &= len(tc.ongoing_backup) == 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
|
return idle
|
||||||
|
|
||||||
|
|||||||
@@ -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.
|
||||||
|
"""
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -1775,6 +1775,8 @@ class HiRadixCache(RadixCache):
|
|||||||
new_input_tokens: List[int],
|
new_input_tokens: List[int],
|
||||||
last_hash: Optional[str] = None,
|
last_hash: Optional[str] = None,
|
||||||
prefix_keys: Optional[List[str]] = None,
|
prefix_keys: Optional[List[str]] = None,
|
||||||
|
# Scheduler-call parity with UnifiedRadixCache; unused in cache mode.
|
||||||
|
matched_prefix_tokens: Optional[List[int]] = None,
|
||||||
):
|
):
|
||||||
prefetch_key = RadixKey(
|
prefetch_key = RadixKey(
|
||||||
new_input_tokens,
|
new_input_tokens,
|
||||||
|
|||||||
@@ -112,6 +112,7 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
storage_backend_extra_config: Optional[dict] = None,
|
storage_backend_extra_config: Optional[dict] = None,
|
||||||
transfer_layer_num: Optional[int] = None,
|
transfer_layer_num: Optional[int] = None,
|
||||||
enable_storage_metrics: bool = False,
|
enable_storage_metrics: bool = False,
|
||||||
|
host_memory_mode: str = "cache",
|
||||||
):
|
):
|
||||||
startup_storage_backend = storage_backend
|
startup_storage_backend = storage_backend
|
||||||
self.extra_host_mem_release_queues: dict[PoolName, Queue[torch.Tensor]] = {}
|
self.extra_host_mem_release_queues: dict[PoolName, Queue[torch.Tensor]] = {}
|
||||||
@@ -131,6 +132,7 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
model_name=model_name,
|
model_name=model_name,
|
||||||
storage_backend_extra_config=storage_backend_extra_config,
|
storage_backend_extra_config=storage_backend_extra_config,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
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)),
|
# 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.
|
# not just the full attention layers reported by full_kv_pool.
|
||||||
@@ -587,6 +589,7 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
hash_value = self.get_hash_str(
|
hash_value = self.get_hash_str(
|
||||||
operation.token_ids, operation.last_hash, page_size=self.page_size
|
operation.token_ids, operation.last_hash, page_size=self.page_size
|
||||||
)
|
)
|
||||||
|
operation.all_hash_values = hash_value
|
||||||
|
|
||||||
extra_info = HiCacheStorageExtraInfo(
|
extra_info = HiCacheStorageExtraInfo(
|
||||||
prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None
|
prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None
|
||||||
@@ -780,6 +783,28 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
continue
|
continue
|
||||||
trailing_n = len(transfer.keys) if transfer.keys else 1
|
trailing_n = len(transfer.keys) if transfer.keys else 1
|
||||||
transfer.keys = all_hashes[max(0, kv_hit_pages - trailing_n) : kv_hit_pages]
|
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(
|
def _resolve_pool_transfers_allocation(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -307,6 +307,7 @@ def build_kv_only_stack(
|
|||||||
storage_backend_extra_config=storage_backend_extra_config,
|
storage_backend_extra_config=storage_backend_extra_config,
|
||||||
transfer_layer_num=transfer_layer_num,
|
transfer_layer_num=transfer_layer_num,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
enable_storage_metrics=enable_storage_metrics,
|
||||||
|
host_memory_mode=server_args.hicache_host_memory_mode,
|
||||||
)
|
)
|
||||||
if params.mtp_draft_device_pools:
|
if params.mtp_draft_device_pools:
|
||||||
cache_controller.set_mtp_draft_pools(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,
|
storage_backend_extra_config=storage_backend_extra_config,
|
||||||
transfer_layer_num=transfer_layer_num,
|
transfer_layer_num=transfer_layer_num,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
enable_storage_metrics=enable_storage_metrics,
|
||||||
|
host_memory_mode=server_args.hicache_host_memory_mode,
|
||||||
)
|
)
|
||||||
if mtp_swa_device_pools:
|
if mtp_swa_device_pools:
|
||||||
cache_controller.set_mtp_draft_pools(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,
|
storage_backend_extra_config=storage_backend_extra_config,
|
||||||
transfer_layer_num=transfer_layer_num,
|
transfer_layer_num=transfer_layer_num,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
enable_storage_metrics=enable_storage_metrics,
|
||||||
|
host_memory_mode=server_args.hicache_host_memory_mode,
|
||||||
)
|
)
|
||||||
if mtp_swa_device_buffers:
|
if mtp_swa_device_buffers:
|
||||||
cache_controller.set_mtp_draft_pools(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,
|
storage_backend_extra_config=storage_backend_extra_config,
|
||||||
transfer_layer_num=transfer_layer_num,
|
transfer_layer_num=transfer_layer_num,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
enable_storage_metrics=enable_storage_metrics,
|
||||||
|
host_memory_mode=server_args.hicache_host_memory_mode,
|
||||||
)
|
)
|
||||||
if mtp_draft_device_pools:
|
if mtp_draft_device_pools:
|
||||||
cache_controller.set_mtp_draft_pools(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,
|
storage_backend_extra_config=storage_backend_extra_config,
|
||||||
transfer_layer_num=transfer_layer_num,
|
transfer_layer_num=transfer_layer_num,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
enable_storage_metrics=enable_storage_metrics,
|
||||||
|
host_memory_mode=server_args.hicache_host_memory_mode,
|
||||||
)
|
)
|
||||||
return host_pool_group, cache_controller
|
return host_pool_group, cache_controller
|
||||||
|
|
||||||
@@ -951,6 +956,7 @@ def build_anchor_sidecar_stack(
|
|||||||
storage_backend_extra_config=storage_backend_extra_config,
|
storage_backend_extra_config=storage_backend_extra_config,
|
||||||
transfer_layer_num=transfer_layer_num,
|
transfer_layer_num=transfer_layer_num,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
enable_storage_metrics=enable_storage_metrics,
|
||||||
|
host_memory_mode=server_args.hicache_host_memory_mode,
|
||||||
)
|
)
|
||||||
if mtp_draft_device_pools:
|
if mtp_draft_device_pools:
|
||||||
cache_controller.set_mtp_draft_pools(mtp_draft_device_pools)
|
cache_controller.set_mtp_draft_pools(mtp_draft_device_pools)
|
||||||
|
|||||||
@@ -129,7 +129,7 @@ class MambaPoolHost(HostKVCache):
|
|||||||
for conv_state in device_pool.mamba_cache.conv
|
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._init_write_back_staging_buffers()
|
||||||
self.lock = threading.RLock()
|
self.lock = threading.RLock()
|
||||||
self.clear()
|
self.clear()
|
||||||
@@ -201,6 +201,12 @@ class MambaPoolHost(HostKVCache):
|
|||||||
allocator=self.allocator,
|
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):
|
def _init_write_back_staging_buffers(self):
|
||||||
self.temporal_staging_buffer = None
|
self.temporal_staging_buffer = None
|
||||||
|
|||||||
@@ -237,6 +237,16 @@ def create_tree_cache(ctx: TreeCacheBuildContext) -> BasePrefixCache:
|
|||||||
cache = default_radix_cache_factory(ctx)
|
cache = default_radix_cache_factory(ctx)
|
||||||
source = "default"
|
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(
|
if ctx.server_args.enable_session_radix_cache and not getattr(
|
||||||
cache, "enable_session_radix_cache", False
|
cache, "enable_session_radix_cache", False
|
||||||
):
|
):
|
||||||
|
|||||||
@@ -158,7 +158,7 @@ class StorageBackendFactory:
|
|||||||
mem_pool_host: Any,
|
mem_pool_host: Any,
|
||||||
) -> HiCacheStorage:
|
) -> HiCacheStorage:
|
||||||
"""Create built-in backend with original initialization logic."""
|
"""Create built-in backend with original initialization logic."""
|
||||||
if backend_name == "file":
|
if backend_name in ("file", "sim"):
|
||||||
return backend_class(storage_config)
|
return backend_class(storage_config)
|
||||||
elif backend_name == "nixl":
|
elif backend_name == "nixl":
|
||||||
return backend_class(storage_config)
|
return backend_class(storage_config)
|
||||||
@@ -198,6 +198,10 @@ StorageBackendFactory.register_backend(
|
|||||||
"file", "sglang.srt.mem_cache.hicache_storage", "HiCacheFile"
|
"file", "sglang.srt.mem_cache.hicache_storage", "HiCacheFile"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
StorageBackendFactory.register_backend(
|
||||||
|
"sim", "sglang.srt.mem_cache.storage.sim_storage", "SimHiCacheStorage"
|
||||||
|
)
|
||||||
|
|
||||||
StorageBackendFactory.register_backend(
|
StorageBackendFactory.register_backend(
|
||||||
"nixl",
|
"nixl",
|
||||||
"sglang.srt.mem_cache.storage.nixl.hicache_nixl",
|
"sglang.srt.mem_cache.storage.nixl.hicache_nixl",
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -71,6 +71,9 @@ class SWAComponent(TreeComponent):
|
|||||||
super().__init__(cache, params)
|
super().__init__(cache, params)
|
||||||
self._session_leaf_covered_len: dict[str, dict[UnifiedTreeNode, int]] = {}
|
self._session_leaf_covered_len: dict[str, dict[UnifiedTreeNode, int]] = {}
|
||||||
self.sliding_window_size = params.sliding_window_size
|
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
|
# HiCache state: set to host SWA pool when HiCache enabled
|
||||||
self._swa_kv_pool_host = None
|
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.
|
# unified_kv keeps SWA as a device-only ring -- nothing to prefetch into.
|
||||||
if self._swa_kv_pool_host is None:
|
if self._swa_kv_pool_host is None:
|
||||||
return PreparePrefetchResult()
|
return PreparePrefetchResult()
|
||||||
sw_pages = (
|
sw_pages = self.full_window_pages
|
||||||
self.cache.sliding_window_size + self.cache.page_size - 1
|
if sw_pages == 0:
|
||||||
) // self.cache.page_size
|
|
||||||
if sw_pages == 0 or prefetch_tokens // self.cache.page_size < sw_pages:
|
|
||||||
return PreparePrefetchResult()
|
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)
|
host_indices = self._swa_kv_pool_host.alloc(num_tokens)
|
||||||
if host_indices is None:
|
if host_indices is None:
|
||||||
self.cache.evict_host(num_tokens, ComponentType.SWA)
|
self.cache.evict_host(num_tokens, ComponentType.SWA)
|
||||||
@@ -869,12 +887,14 @@ class SWAComponent(TreeComponent):
|
|||||||
|
|
||||||
if phase == CacheTransferPhase.PREFETCH:
|
if phase == CacheTransferPhase.PREFETCH:
|
||||||
assert host_indices is not None
|
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 [
|
return [
|
||||||
PoolTransfer(
|
PoolTransfer(
|
||||||
name=PoolName.SWA,
|
name=PoolName.SWA,
|
||||||
host_indices=host_indices,
|
host_indices=host_indices,
|
||||||
keys=["__placeholder__"] * sw_pages,
|
keys=["__placeholder__"] * num_pages,
|
||||||
hit_policy=PoolHitPolicy.TRAILING_PAGES,
|
hit_policy=PoolHitPolicy.TRAILING_PAGES,
|
||||||
)
|
)
|
||||||
]
|
]
|
||||||
@@ -997,6 +1017,13 @@ class SWAComponent(TreeComponent):
|
|||||||
and insert_result.inserted_host_node is not None
|
and insert_result.inserted_host_node is not None
|
||||||
else 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 (
|
if (
|
||||||
target is None
|
target is None
|
||||||
or window_require_pages == 0
|
or window_require_pages == 0
|
||||||
|
|||||||
@@ -25,6 +25,13 @@ from sglang.srt.mem_cache.base_prefix_cache import (
|
|||||||
MatchPrefixParams,
|
MatchPrefixParams,
|
||||||
MatchResult,
|
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.common import RetractionBackup
|
||||||
from sglang.srt.mem_cache.hicache_storage import (
|
from sglang.srt.mem_cache.hicache_storage import (
|
||||||
PoolHitPolicy,
|
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 (
|
from sglang.srt.observability.metrics_collector import (
|
||||||
STAT_LOGGER_ROLE_STORAGE,
|
STAT_LOGGER_ROLE_STORAGE,
|
||||||
|
StorageMetrics,
|
||||||
StorageMetricsCollector,
|
StorageMetricsCollector,
|
||||||
resolve_collector_class,
|
resolve_collector_class,
|
||||||
)
|
)
|
||||||
@@ -227,6 +235,30 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
self.prefetch_timeout_base = 1.0
|
self.prefetch_timeout_base = 1.0
|
||||||
self.prefetch_timeout_per_page = 0.25
|
self.prefetch_timeout_per_page = 0.25
|
||||||
self.hicache_storage_pass_prefix_keys = False
|
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()
|
self.reset()
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -316,6 +348,8 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
|
self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
|
||||||
self.ongoing_prefetch: dict[str, _OngoingPrefetch] = {}
|
self.ongoing_prefetch: dict[str, _OngoingPrefetch] = {}
|
||||||
self.ongoing_backup: dict[int, tuple[NodeId, DecLockRefParams]] = {}
|
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:
|
if self.cache_controller is not None:
|
||||||
self.cache_controller.reset()
|
self.cache_controller.reset()
|
||||||
@@ -326,6 +360,20 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
|
|
||||||
def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None:
|
def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None:
|
||||||
"""Initialize HiCache infrastructure."""
|
"""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 (
|
from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
|
||||||
attach_hybrid_pool_to_unified_cache,
|
attach_hybrid_pool_to_unified_cache,
|
||||||
)
|
)
|
||||||
@@ -368,6 +416,28 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
swa = self.components[ComponentType.SWA]
|
swa = self.components[ComponentType.SWA]
|
||||||
self.tree_core.has_swa_host_pool = swa._swa_kv_pool_host is not None
|
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
|
# State initialization
|
||||||
self.write_through_threshold = (
|
self.write_through_threshold = (
|
||||||
1 if server_args.hicache_write_policy == "write_through" else 2
|
1 if server_args.hicache_write_policy == "write_through" else 2
|
||||||
@@ -541,6 +611,9 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
request_by_type: dict[ComponentType, int],
|
request_by_type: dict[ComponentType, int],
|
||||||
tracker: dict[ComponentType, int],
|
tracker: dict[ComponentType, int],
|
||||||
) -> None:
|
) -> 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:
|
for ct in self.tree_components:
|
||||||
request_cnt = request_by_type[ct]
|
request_cnt = request_by_type[ct]
|
||||||
# Skip eviction walk if request is already met
|
# 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
|
self, num_tokens: int, component_type: ComponentType = BASE_COMPONENT_TYPE
|
||||||
) -> int:
|
) -> int:
|
||||||
"""Evict host resources for a specific component to free host pool space."""
|
"""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)
|
result = self.tree_core.drive_host_eviction(component_type, num_tokens)
|
||||||
self._free_values(result.device_frees, result.host_frees)
|
self._free_values(result.device_frees, result.host_frees)
|
||||||
return result.tracker.get(component_type, 0)
|
return result.tracker.get(component_type, 0)
|
||||||
@@ -1170,6 +1247,16 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
self, action: BackupKV, write_back: bool = False
|
self, action: BackupKV, write_back: bool = False
|
||||||
) -> int:
|
) -> int:
|
||||||
"""Run a backup action top-down, stopping at the first failed backup."""
|
"""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
|
written = 0
|
||||||
for node_id in action.node_ids:
|
for node_id in action.node_ids:
|
||||||
device_value, comp_xfers = self.tree_core.build_backup_spec(node_id)
|
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:
|
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(
|
lock_node_id, lock_params, publish_node_ids = self.ongoing_write_through.pop(
|
||||||
ack_id
|
ack_id
|
||||||
)
|
)
|
||||||
@@ -1468,10 +1559,12 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
new_input_tokens: list[int],
|
new_input_tokens: list[int],
|
||||||
last_hash: Optional[str] = None,
|
last_hash: Optional[str] = None,
|
||||||
prefix_keys: Optional[list[str]] = None,
|
prefix_keys: Optional[list[str]] = None,
|
||||||
|
matched_prefix_tokens: Optional[list[int]] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
if not self.enable_storage or self.cache_controller is None:
|
if not self.enable_storage or self.cache_controller is None:
|
||||||
return
|
return
|
||||||
|
|
||||||
|
buffer_mode = self.host_memory_mode == "buffer_only"
|
||||||
extra_key, cache_salt = self.tree_core.prefetch_anchor_info(last_host_node_id)
|
extra_key, cache_salt = self.tree_core.prefetch_anchor_info(last_host_node_id)
|
||||||
prefetch_key = RadixKey(
|
prefetch_key = RadixKey(
|
||||||
new_input_tokens,
|
new_input_tokens,
|
||||||
@@ -1480,13 +1573,30 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
cache_salt=cache_salt,
|
cache_salt=cache_salt,
|
||||||
).page_aligned(self.page_size)
|
).page_aligned(self.page_size)
|
||||||
prefetch_length = len(prefetch_key)
|
prefetch_length = len(prefetch_key)
|
||||||
if (
|
stats = self._prefetch_outcome_stats
|
||||||
prefetch_length < self.prefetch_threshold
|
if prefetch_length > 0:
|
||||||
or self.cache_controller.prefetch_rate_limited()
|
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
|
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]] = {}
|
comp_xfers: dict[ComponentType, list[PoolTransfer]] = {}
|
||||||
alloc_failed = False
|
alloc_failed = False
|
||||||
for ct in self.tree_components:
|
for ct in self.tree_components:
|
||||||
@@ -1517,10 +1627,21 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
CacheTransferPhase.PREFETCH, kv_xfer, comp_xfers
|
CacheTransferPhase.PREFETCH, kv_xfer, comp_xfers
|
||||||
)
|
)
|
||||||
if alloc_failed:
|
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(
|
self.cache_controller.append_host_mem_release(
|
||||||
extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
|
extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
|
||||||
)
|
)
|
||||||
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
|
return
|
||||||
|
|
||||||
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
|
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
|
||||||
@@ -1532,6 +1653,13 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
prefix_keys,
|
prefix_keys,
|
||||||
extra_pools=aux_xfers or None,
|
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(
|
self.ongoing_prefetch[req_id] = _OngoingPrefetch(
|
||||||
last_host_node_id,
|
last_host_node_id,
|
||||||
prefetch_key,
|
prefetch_key,
|
||||||
@@ -1540,7 +1668,12 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
anchor_lock_params,
|
anchor_lock_params,
|
||||||
comp_xfers,
|
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:
|
def _prefetch_timeout_check_linear_func(self, operation: PrefetchOperation) -> bool:
|
||||||
return (
|
return (
|
||||||
@@ -1622,6 +1755,16 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
# Hybrid all-or-nothing check failed; result already discarded.
|
# Hybrid all-or-nothing check failed; result already discarded.
|
||||||
return True
|
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]
|
fetched_key = prefetch_key[:min_completed_tokens]
|
||||||
insert_result = self.tree_core.insert_host(
|
insert_result = self.tree_core.insert_host(
|
||||||
last_host_node_id,
|
last_host_node_id,
|
||||||
@@ -1751,9 +1894,14 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
host_indices=host_indices[:completed_tokens],
|
host_indices=host_indices[:completed_tokens],
|
||||||
extra_pools=pool_transfers,
|
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]
|
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
|
self.prefetch_loaded_tokens_by_reqid[req_id] = 0
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"HiCache hybrid prefetch discarded req=%s completed=%d requested=%d",
|
"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:
|
def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
|
||||||
return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
|
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:
|
def release_aborted_request(self, rid: str) -> None:
|
||||||
self.prefetch_loaded_tokens_by_reqid.pop(rid, 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:
|
if rid not in self.ongoing_prefetch:
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -1793,13 +1960,73 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
|
|
||||||
completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
|
completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
|
||||||
self._barrier_attn_groups()
|
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]
|
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(
|
self.cache_controller.append_host_mem_release(
|
||||||
host_indices=host_indices[:completed_tokens],
|
host_indices=host_indices[:completed_tokens],
|
||||||
extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
|
extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
|
||||||
)
|
)
|
||||||
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:
|
def _revoke_pending_prefetch(self, req_id: str) -> None:
|
||||||
info = self.ongoing_prefetch.pop(req_id, None)
|
info = self.ongoing_prefetch.pop(req_id, None)
|
||||||
@@ -1809,17 +2036,27 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
last_host_node_id,
|
last_host_node_id,
|
||||||
prefetch_key,
|
prefetch_key,
|
||||||
_host_indices,
|
_host_indices,
|
||||||
_operation,
|
operation,
|
||||||
anchor_lock_params,
|
anchor_lock_params,
|
||||||
comp_xfers,
|
comp_xfers,
|
||||||
) = info
|
) = 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 = self.cache_controller
|
||||||
cc.append_host_mem_release(
|
cc.append_host_mem_release(
|
||||||
extra_pools=[x for xfers in comp_xfers.values() for x in xfers]
|
extra_pools=[x for xfers in comp_xfers.values() for x in xfers]
|
||||||
)
|
)
|
||||||
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(
|
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(
|
def _drain_storage_control_queues_impl(
|
||||||
@@ -1842,56 +2079,100 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
drained += 1
|
drained += 1
|
||||||
yield item
|
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():
|
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):
|
for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit):
|
||||||
req_id = operation.request_id
|
req_id = operation.request_id
|
||||||
info = self.ongoing_prefetch.get(req_id)
|
info = self.ongoing_prefetch.get(req_id)
|
||||||
if info is None:
|
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
|
continue
|
||||||
if operation.is_terminated():
|
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)
|
self._revoke_pending_prefetch(req_id)
|
||||||
continue
|
continue
|
||||||
if operation.storage_hit_count < self.prefetch_threshold:
|
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)
|
self._revoke_pending_prefetch(req_id)
|
||||||
continue
|
continue
|
||||||
|
self._invalidate_absent_from_hit_query(operation)
|
||||||
alloc_len = operation.storage_hit_count
|
self._account_prefetch_outcome(operation, revoked=False)
|
||||||
host_indices = cc.mem_pool_host.alloc(alloc_len)
|
if not _try_alloc_storage_hit(operation):
|
||||||
if host_indices is None:
|
# Counted once at first parking, not per retry tick.
|
||||||
self.evict_host(alloc_len)
|
self._prefetch_outcome_stats["declined_rate_limited"] += 1
|
||||||
host_indices = cc.mem_pool_host.alloc(alloc_len)
|
self.buffer_pipeline.pending_hit_allocs.append(operation)
|
||||||
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)
|
|
||||||
|
|
||||||
def _drain_backup():
|
def _drain_backup():
|
||||||
drained = 0
|
drained = 0
|
||||||
for operation in _drain_queue(cc.ack_backup_queue, n_backup):
|
for operation in _drain_queue(cc.ack_backup_queue, n_backup):
|
||||||
drained += 1
|
drained += 1
|
||||||
entry = self.ongoing_backup.pop(operation.id, None)
|
if buffer_mode:
|
||||||
if entry is not None:
|
# Storage write acked: free the staging.
|
||||||
node_id, lock_params = entry
|
self.buffer_pipeline.finish_storage_write_ack(operation.id)
|
||||||
self.dec_host_lock_ref(node_id, lock_params)
|
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 (
|
if (
|
||||||
log_metrics
|
log_metrics
|
||||||
and self.enable_storage_metrics
|
and self.enable_storage_metrics
|
||||||
@@ -2054,6 +2335,9 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
logger.error("Failed to clear hierarchical cache storage backend: %s", e)
|
logger.error("Failed to clear hierarchical cache storage backend: %s", e)
|
||||||
return False
|
return False
|
||||||
if ok:
|
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!")
|
logger.info("Hierarchical cache storage backend cleared successfully!")
|
||||||
return ok
|
return ok
|
||||||
|
|
||||||
@@ -2205,6 +2489,11 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
ack = cc.ack_load_queue.pop(0)
|
ack = cc.ack_load_queue.pop(0)
|
||||||
ack.finish_event.synchronize()
|
ack.finish_event.synchronize()
|
||||||
for ack_id in ack.node_ids:
|
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)
|
node, lock_params, host_lock_params = self.ongoing_load_back.pop(ack_id)
|
||||||
self.dec_lock_ref(node, lock_params)
|
self.dec_lock_ref(node, lock_params)
|
||||||
self.dec_host_lock_ref(node, host_lock_params)
|
self.dec_host_lock_ref(node, host_lock_params)
|
||||||
@@ -2233,7 +2522,10 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
params: InitLoadBackParams,
|
params: InitLoadBackParams,
|
||||||
) -> tuple[torch.Tensor, NodeId]:
|
) -> tuple[torch.Tensor, NodeId]:
|
||||||
"""Prepare KV cache loading from host to device.
|
"""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
|
best_match_node_id = params.best_match_node
|
||||||
mem_quota = params.mem_quota
|
mem_quota = params.mem_quota
|
||||||
req = params.req
|
req = params.req
|
||||||
@@ -2306,10 +2598,16 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
extra_release_counts=extra_release_counts,
|
extra_release_counts=extra_release_counts,
|
||||||
log_metrics=True,
|
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:
|
if self.enable_storage_metrics and self.storage_metrics_collector is not None:
|
||||||
self.storage_metrics_collector.log_storage_metrics(
|
storage_metrics = self.cache_controller.storage_backend.get_stats()
|
||||||
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:
|
def ready_to_load_host_cache(self) -> int:
|
||||||
"""Notify the cache controller to start the KV cache loading."""
|
"""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
|
# Pass ongoing ops as lightweight (id, node_id) pairs so the tree core
|
||||||
# can resolve + validate them without reaching into Controller state.
|
# can resolve + validate them without reaching into Controller state.
|
||||||
ongoing_write_through = [
|
if self.buffer_pipeline is not None:
|
||||||
(nid, wt.node_id) for nid, wt in self.ongoing_write_through.items()
|
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 = [
|
ongoing_load_back = [
|
||||||
(nid, lb.node_id) for nid, lb in self.ongoing_load_back.items()
|
(nid, lb.node_id) for nid, lb in self.ongoing_load_back.items()
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -1852,9 +1852,11 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
|
|||||||
labels: Dict[str, str],
|
labels: Dict[str, str],
|
||||||
):
|
):
|
||||||
from prometheus_client import Counter as _PromCounter
|
from prometheus_client import Counter as _PromCounter
|
||||||
|
from prometheus_client import Gauge as _PromGauge
|
||||||
from prometheus_client import Histogram as _PromHistogram
|
from prometheus_client import Histogram as _PromHistogram
|
||||||
|
|
||||||
Counter = self._counter_cls or _PromCounter
|
Counter = self._counter_cls or _PromCounter
|
||||||
|
Gauge = self._gauge_cls or _PromGauge
|
||||||
Histogram = self._histogram_cls or _PromHistogram
|
Histogram = self._histogram_cls or _PromHistogram
|
||||||
|
|
||||||
self.labels = labels
|
self.labels = labels
|
||||||
@@ -1871,6 +1873,21 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
|
|||||||
labelnames=labels.keys(),
|
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 = [
|
bucket_io = [
|
||||||
1,
|
1,
|
||||||
5,
|
5,
|
||||||
@@ -1925,6 +1942,16 @@ class StorageMetricsCollector(_StatLoggerDIMixin):
|
|||||||
if backuped_tokens > 0:
|
if backuped_tokens > 0:
|
||||||
self.backuped_tokens_total.labels(**self.labels).inc(backuped_tokens)
|
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]):
|
def _log_histogram(self, histogram, data: Union[int, float]):
|
||||||
histogram.labels(**self.labels).observe(data)
|
histogram.labels(**self.labels).observe(data)
|
||||||
|
|
||||||
|
|||||||
@@ -2668,14 +2668,22 @@ class ServerArgs:
|
|||||||
enable_hierarchical_cache: A[bool, "Enable hierarchical cache", NS("memory")] = (
|
enable_hierarchical_cache: A[bool, "Enable hierarchical cache", NS("memory")] = (
|
||||||
False
|
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[
|
hicache_ratio: A[
|
||||||
Optional[float],
|
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"),
|
NS("memory"),
|
||||||
] = None
|
] = None
|
||||||
hicache_size: A[
|
hicache_size: A[
|
||||||
int,
|
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"),
|
NS("memory"),
|
||||||
] = 0
|
] = 0
|
||||||
hicache_write_policy: A[
|
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).",
|
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=[
|
choices=[
|
||||||
"file",
|
"file",
|
||||||
|
"sim",
|
||||||
"mooncake",
|
"mooncake",
|
||||||
"hf3fs",
|
"hf3fs",
|
||||||
"nixl",
|
"nixl",
|
||||||
@@ -7393,8 +7402,20 @@ class ServerArgs:
|
|||||||
)
|
)
|
||||||
|
|
||||||
def _handle_hicache_ratio_default(self):
|
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":
|
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):
|
def _handle_hicache(self):
|
||||||
"""Normalize hicache-related knobs into a valid runtime configuration.
|
"""Normalize hicache-related knobs into a valid runtime configuration.
|
||||||
@@ -7414,6 +7435,8 @@ class ServerArgs:
|
|||||||
):
|
):
|
||||||
return
|
return
|
||||||
|
|
||||||
|
self._validate_hicache_host_memory_mode()
|
||||||
|
|
||||||
# Step 1: Initial layout-io compatibility normalization.
|
# Step 1: Initial layout-io compatibility normalization.
|
||||||
self._resolve_layout_io_compatibility()
|
self._resolve_layout_io_compatibility()
|
||||||
|
|
||||||
@@ -7423,6 +7446,51 @@ class ServerArgs:
|
|||||||
# Step 3: DCP compatibility for the L2 (device<->host) path.
|
# Step 3: DCP compatibility for the L2 (device<->host) path.
|
||||||
self._resolve_hicache_dcp_compatibility()
|
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):
|
def _resolve_hicache_dcp_compatibility(self):
|
||||||
if self.dcp_size <= 1 or not self.enable_hierarchical_cache:
|
if self.dcp_size <= 1 or not self.enable_hierarchical_cache:
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -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.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||||
from sglang.test.test_utils import CustomTestCase
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
register_cuda_ci(est_time=16, stage="base-b", runner_config="1-gpu-small")
|
register_cuda_ci(est_time=50, stage="base-b", runner_config="1-gpu-small")
|
||||||
register_amd_ci(est_time=16, suite="stage-b-test-1-gpu-small-amd")
|
register_amd_ci(est_time=50, suite="stage-b-test-1-gpu-small-amd")
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
@@ -595,39 +595,20 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase):
|
|||||||
return leaf
|
return leaf
|
||||||
|
|
||||||
def _init_hicache(self, cache, *, write_policy: str = "write_through"):
|
def _init_hicache(self, cache, *, write_policy: str = "write_through"):
|
||||||
import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler
|
# Production config: kernel IO backend + page_first layout with
|
||||||
|
# PINNED host pools (kernel GPU DMA requires them). Pools track their
|
||||||
# Wrap the host-pool factory (not MHATokenToKVPoolHost directly)
|
# cudaHostRegister'd pointers and unregister on destroy()/GC, so the
|
||||||
# because the assembler picks between MHATokenToKVPoolHost and
|
# many fixtures sharing this process cannot collide on recycled
|
||||||
# AsymmetricMHATokenToKVPoolHost via get_mha_host_pool_cls(device_pool).
|
# address ranges (rc=712).
|
||||||
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)
|
|
||||||
|
|
||||||
server_args = ServerArgs(
|
server_args = ServerArgs(
|
||||||
model_path="dummy",
|
model_path="dummy",
|
||||||
page_size=self.cfg.page_size,
|
page_size=self.cfg.page_size,
|
||||||
hicache_io_backend="direct",
|
hicache_io_backend="kernel",
|
||||||
hicache_mem_layout="page_first_direct",
|
|
||||||
hicache_write_policy=write_policy,
|
hicache_write_policy=write_policy,
|
||||||
)
|
)
|
||||||
set_global_server_args_for_scheduler(server_args)
|
set_global_server_args_for_scheduler(server_args)
|
||||||
cache.init_hicache(server_args, cache.cache_init_params)
|
cache.init_hicache(server_args, cache.cache_init_params)
|
||||||
|
self.addCleanup(cache.release_host_resources)
|
||||||
cache.write_through_threshold = 1 << 30
|
cache.write_through_threshold = 1 << 30
|
||||||
cache.load_back_threshold = 0
|
cache.load_back_threshold = 0
|
||||||
|
|
||||||
@@ -2468,28 +2449,87 @@ class UnifiedRadixCacheSuite:
|
|||||||
for n in self._path_chain(cache, node):
|
for n in self._path_chain(cache, node):
|
||||||
cache.write_backup_storage(n.id)
|
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):
|
def _flush_l3_backups(self, cache, timeout: float = 10.0):
|
||||||
"""Wait for backup threads to finish, then drain acks (release locks)."""
|
"""Wait for backup threads to finish, then drain acks (release locks)."""
|
||||||
deadline = time.time() + timeout
|
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()
|
cache.drain_storage_control_queues()
|
||||||
if cache.ongoing_backup:
|
if self._ongoing_l3_backups(cache):
|
||||||
time.sleep(0.01)
|
time.sleep(0.01)
|
||||||
cache.drain_storage_control_queues()
|
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):
|
def _run_prefetch_to_completion(self, cache, req_id, timeout: float = 10.0):
|
||||||
deadline = time.time() + timeout
|
deadline = time.time() + timeout
|
||||||
while time.time() < deadline:
|
while time.time() < deadline:
|
||||||
# Host memory is reserved (and IO started) by the scheduler-thread
|
# 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
|
# drain once the L3 hit count is known, and the buffer-mode fill
|
||||||
# scheduler loop does (check_hicache_events before progress checks).
|
# commits at its H2D ack in loading_check — pump the full event
|
||||||
cache.drain_storage_control_queues()
|
# round like the real scheduler loop does.
|
||||||
|
cache.check_hicache_events()
|
||||||
if cache.check_prefetch_progress(req_id):
|
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
|
return
|
||||||
time.sleep(0.01)
|
time.sleep(0.01)
|
||||||
self.fail(f"prefetch {req_id} did not complete in time")
|
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):
|
def _all_page_hashes(self, cache, node):
|
||||||
hashes = []
|
hashes = []
|
||||||
for n in self._path_chain(cache, node):
|
for n in self._path_chain(cache, node):
|
||||||
@@ -2606,6 +2646,494 @@ class UnifiedRadixCacheSuite:
|
|||||||
self.assertTrue(torch.equal(loaded_v, expected_v))
|
self.assertTrue(torch.equal(loaded_v, expected_v))
|
||||||
cons.sanity_check()
|
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) ----------
|
# ---------- TP consistency for SWA prefetch (all-or-nothing) ----------
|
||||||
|
|
||||||
def _patch_tp_all_reduce(self, cache, drop_swa: bool):
|
def _patch_tp_all_reduce(self, cache, drop_swa: bool):
|
||||||
@@ -2898,44 +3426,9 @@ class UnifiedRadixCacheSuite:
|
|||||||
storage_dir: Optional[str] = None,
|
storage_dir: Optional[str] = None,
|
||||||
prefetch_threshold: Optional[int] = None,
|
prefetch_threshold: Optional[int] = None,
|
||||||
prefetch_policy: str = "wait_complete",
|
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
|
storage_extra_config = None
|
||||||
if storage_backend == "file":
|
if storage_backend == "file":
|
||||||
from sglang.srt.runtime_context import get_parallel
|
from sglang.srt.runtime_context import get_parallel
|
||||||
@@ -2957,22 +3450,26 @@ class UnifiedRadixCacheSuite:
|
|||||||
extra = {}
|
extra = {}
|
||||||
if prefetch_threshold is not None:
|
if prefetch_threshold is not None:
|
||||||
extra["prefetch_threshold"] = prefetch_threshold
|
extra["prefetch_threshold"] = prefetch_threshold
|
||||||
|
if storage_extra:
|
||||||
|
extra.update(storage_extra)
|
||||||
storage_extra_config = json.dumps(extra) if extra else None
|
storage_extra_config = json.dumps(extra) if extra else None
|
||||||
|
|
||||||
|
# See _init_hicache: production kernel IO backend, pinned pools.
|
||||||
server_args = ServerArgs(
|
server_args = ServerArgs(
|
||||||
model_path="dummy",
|
model_path="dummy",
|
||||||
page_size=self.cfg.page_size,
|
page_size=self.cfg.page_size,
|
||||||
hicache_io_backend="direct",
|
hicache_io_backend="kernel",
|
||||||
hicache_mem_layout="page_first_direct",
|
|
||||||
hicache_write_policy=write_policy,
|
hicache_write_policy=write_policy,
|
||||||
hicache_storage_backend=storage_backend,
|
hicache_storage_backend=storage_backend,
|
||||||
hicache_storage_backend_extra_config=storage_extra_config,
|
hicache_storage_backend_extra_config=storage_extra_config,
|
||||||
hicache_storage_prefetch_policy=prefetch_policy,
|
hicache_storage_prefetch_policy=prefetch_policy,
|
||||||
|
hicache_host_memory_mode=host_memory_mode,
|
||||||
)
|
)
|
||||||
# See build_fixture for why _mamba_cache_chunk_size is preset.
|
# 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)
|
server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size)
|
||||||
set_global_server_args_for_scheduler(server_args)
|
set_global_server_args_for_scheduler(server_args)
|
||||||
cache.init_hicache(server_args, cache.cache_init_params)
|
cache.init_hicache(server_args, cache.cache_init_params)
|
||||||
|
self.addCleanup(cache.release_host_resources)
|
||||||
cache.write_through_threshold = 1 << 30
|
cache.write_through_threshold = 1 << 30
|
||||||
cache.load_back_threshold = 0
|
cache.load_back_threshold = 0
|
||||||
if storage_backend is not None:
|
if storage_backend is not None:
|
||||||
@@ -3020,8 +3517,8 @@ class UnifiedRadixCacheSuite:
|
|||||||
if node is not cache.root_node:
|
if node is not cache.root_node:
|
||||||
self._backup_node(cache, node)
|
self._backup_node(cache, node)
|
||||||
|
|
||||||
def _load_back_node(self, cache, node):
|
def _load_back_node(self, cache, node, req=None):
|
||||||
loaded = cache.load_back(node.id)
|
loaded = cache.load_back(node.id, req=req)
|
||||||
self.assertTrue(loaded)
|
self.assertTrue(loaded)
|
||||||
producer_id = cache.ready_to_load_host_cache()
|
producer_id = cache.ready_to_load_host_cache()
|
||||||
self.assertNotEqual(producer_id, -1)
|
self.assertNotEqual(producer_id, -1)
|
||||||
@@ -3729,10 +4226,14 @@ class UnifiedRadixCacheSuite:
|
|||||||
cache, _, _ = self._build_hicache_fixture()
|
cache, _, _ = self._build_hicache_fixture()
|
||||||
sw = cache.sliding_window_size
|
sw = cache.sliding_window_size
|
||||||
swa = cache.components[ComponentType.SWA]
|
swa = cache.components[ComponentType.SWA]
|
||||||
# below a full window -> does not participate, no alloc
|
# zero-length prefetch -> does not participate, no alloc
|
||||||
prep = swa.prepare_prefetch(cache.root_node.id, prefetch_tokens=sw - 1)
|
prep = swa.prepare_prefetch(cache.root_node.id, prefetch_tokens=0)
|
||||||
self.assertFalse(prep.alloc_failed)
|
self.assertFalse(prep.alloc_failed)
|
||||||
self.assertIsNone(prep.host_indices)
|
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
|
# a full window available -> participates, allocs one window of host pages
|
||||||
prep = swa.prepare_prefetch(cache.root_node.id, prefetch_tokens=sw)
|
prep = swa.prepare_prefetch(cache.root_node.id, prefetch_tokens=sw)
|
||||||
self.assertEqual(int(prep.host_indices.numel()), sw)
|
self.assertEqual(int(prep.host_indices.numel()), sw)
|
||||||
@@ -6293,6 +6794,7 @@ class TestPrefetchCommitOrdering(CustomTestCase):
|
|||||||
cache = mock.MagicMock()
|
cache = mock.MagicMock()
|
||||||
cache.page_size = 1
|
cache.page_size = 1
|
||||||
cache.enable_storage_metrics = False
|
cache.enable_storage_metrics = False
|
||||||
|
cache.buffer_pipeline = None # cache-mode commit path
|
||||||
walk_action = object()
|
walk_action = object()
|
||||||
insert_result = mock.MagicMock()
|
insert_result = mock.MagicMock()
|
||||||
insert_result.cache_actions = [walk_action]
|
insert_result.cache_actions = [walk_action]
|
||||||
|
|||||||
Reference in New Issue
Block a user