diff --git a/benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py b/benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py new file mode 100644 index 000000000..c5947d3e1 --- /dev/null +++ b/benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py @@ -0,0 +1,444 @@ +"""Benchmark CP shared-KV prefetch collective layouts. + +This benchmark isolates the current CP shared-KV prefetch merge pattern: + + dense_all_reduce: + each rank materializes its owner-lane pages into a full dense prefix + buffer, zeros non-owned pages, then all-reduces the full dense buffer. + +Against two all-gather candidates: + + packed_all_gather_scatter: + each rank gathers only owner-lane pages, all-gathers owner-packed pages, + then scatters them back into the old dense slot layout. + + packed_all_gather_rank_major: + each rank gathers only owner-lane pages and keeps the all-gather output in + rank-major owner-packed order. This models the intended design where + dense_locs/page_inverse are remapped instead of scattering back. + +Run on the target CUDA machine, for example: + + torchrun --nproc_per_node=8 \\ + benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py \\ + --payload both --tokens 4096 8192 16384 32768 40320 65536 + +Do not use this as an end-to-end throughput benchmark. It is a microbenchmark +for materialize + collective + optional scatter cost. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import statistics +import time +from dataclasses import asdict, dataclass +from typing import Callable, Iterable + +import torch +import torch.distributed as dist + + +@dataclass(frozen=True) +class CaseConfig: + payload: str + prefix_tokens: int + page_size: int + kv_dim: int + index_page_bytes: int + dtype: str + warmup: int + repeat: int + + +@dataclass(frozen=True) +class BenchResult: + payload: str + prefix_tokens: int + prefix_pages: int + path: str + full_mib: float + packed_mib_per_rank: float + gpu_ms_avg: float + gpu_ms_p50: float + gpu_ms_min: float + gpu_ms_max_rank_avg: float + cpu_ms_avg: float + cpu_ms_p50: float + cpu_ms_max_rank_avg: float + + +def _rank_env() -> tuple[int, int, int]: + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + return rank, world_size, local_rank + + +def _dtype(name: str) -> torch.dtype: + if name == "bf16": + return torch.bfloat16 + if name == "fp16": + return torch.float16 + if name == "fp32": + return torch.float32 + raise ValueError(f"unsupported dtype: {name}") + + +def _payload_shape_and_dtype(cfg: CaseConfig) -> tuple[tuple[int, ...], torch.dtype]: + if cfg.payload == "mla": + return (cfg.page_size, cfg.kv_dim), _dtype(cfg.dtype) + if cfg.payload == "index": + return (cfg.index_page_bytes,), torch.uint8 + raise ValueError(f"unsupported payload: {cfg.payload}") + + +def _owner_counts(prefix_pages: int, world_size: int) -> list[int]: + return [len(range(rank, prefix_pages, world_size)) for rank in range(world_size)] + + +def _make_source_pages( + *, + max_owned_pages: int, + page_shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + rank: int, +) -> torch.Tensor: + # Page 0 is kept as a sentinel to match CP shared-KV physical page usage. + shape = (max_owned_pages + 1, *page_shape) + if dtype is torch.uint8: + src = torch.empty(shape, dtype=dtype, device=device) + src.random_(0, 127) + else: + src = torch.empty(shape, dtype=dtype, device=device) + src.normal_(mean=float(rank + 1), std=0.01) + src[0].zero_() + return src + + +def _copy_owned_to_dense( + *, + dense_pages: torch.Tensor, + src_pages: torch.Tensor, + owned_slots: torch.Tensor, + owned_count: int, +) -> None: + if owned_count <= 0: + return + dense_pages[owned_slots].copy_(src_pages[1 : owned_count + 1]) + + +def _copy_owned_to_packed( + *, + local_packed: torch.Tensor, + src_pages: torch.Tensor, + owned_count: int, +) -> None: + local_packed.zero_() + if owned_count > 0: + local_packed[:owned_count].copy_(src_pages[1 : owned_count + 1]) + + +def _scatter_rank_major_to_dense( + *, + dense_pages: torch.Tensor, + gathered: torch.Tensor, + prefix_pages: int, + owner_counts: list[int], + world_size: int, + device: torch.device, +) -> None: + dense_pages.zero_() + for owner_rank, count in enumerate(owner_counts): + if count <= 0: + continue + slots = torch.arange( + owner_rank, + prefix_pages, + world_size, + device=device, + dtype=torch.long, + ) + dense_pages[slots].copy_(gathered[owner_rank, :count]) + + +def _time_path( + fn: Callable[[], None], + *, + warmup: int, + repeat: int, + stream: torch.cuda.Stream, + group: dist.ProcessGroup, + device: torch.device, +) -> tuple[list[float], list[float]]: + for _ in range(warmup): + fn() + torch.cuda.synchronize(device) + dist.barrier(group=group) + + gpu_ms: list[float] = [] + cpu_ms: list[float] = [] + for _ in range(repeat): + dist.barrier(group=group) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize(device) + t0 = time.perf_counter() + with torch.cuda.stream(stream): + start.record(stream) + fn() + end.record(stream) + end.synchronize() + t1 = time.perf_counter() + gpu_ms.append(float(start.elapsed_time(end))) + cpu_ms.append((t1 - t0) * 1000.0) + dist.barrier(group=group) + return gpu_ms, cpu_ms + + +def _summarize_across_ranks( + *, + local_gpu_ms: list[float], + local_cpu_ms: list[float], + group: dist.ProcessGroup, + device: torch.device, +) -> tuple[dict[str, float], dict[str, float]]: + local = torch.tensor( + [ + statistics.mean(local_gpu_ms), + statistics.median(local_gpu_ms), + min(local_gpu_ms), + statistics.mean(local_cpu_ms), + statistics.median(local_cpu_ms), + ], + dtype=torch.float64, + device=device, + ) + world_size = dist.get_world_size(group=group) + gathered = torch.empty((world_size, local.numel()), dtype=local.dtype, device=device) + dist.all_gather_into_tensor(gathered, local, group=group) + gathered_cpu = gathered.cpu() + gpu = { + "avg": float(gathered_cpu[:, 0].mean().item()), + "p50": float(gathered_cpu[:, 1].mean().item()), + "min": float(gathered_cpu[:, 2].mean().item()), + "max_rank_avg": float(gathered_cpu[:, 0].max().item()), + } + cpu = { + "avg": float(gathered_cpu[:, 3].mean().item()), + "p50": float(gathered_cpu[:, 4].mean().item()), + "max_rank_avg": float(gathered_cpu[:, 3].max().item()), + } + return gpu, cpu + + +def _bench_case( + cfg: CaseConfig, + *, + group: dist.ProcessGroup, + device: torch.device, + rank: int, + world_size: int, +) -> list[BenchResult]: + prefix_pages = math.ceil(cfg.prefix_tokens / cfg.page_size) + owner_counts = _owner_counts(prefix_pages, world_size) + owned_count = owner_counts[rank] + max_owned_pages = max(owner_counts) + page_shape, dtype = _payload_shape_and_dtype(cfg) + src_pages = _make_source_pages( + max_owned_pages=max_owned_pages, + page_shape=page_shape, + dtype=dtype, + device=device, + rank=rank, + ) + + dense_pages = torch.empty((prefix_pages, *page_shape), dtype=dtype, device=device) + local_packed = torch.empty( + (max_owned_pages, *page_shape), dtype=dtype, device=device + ) + gathered = torch.empty( + (world_size, max_owned_pages, *page_shape), dtype=dtype, device=device + ) + owned_slots = torch.arange( + rank, + prefix_pages, + world_size, + device=device, + dtype=torch.long, + ) + stream = torch.cuda.Stream(device=device) + full_mib = dense_pages.numel() * dense_pages.element_size() / (1024 * 1024) + packed_mib = local_packed.numel() * local_packed.element_size() / (1024 * 1024) + + def dense_all_reduce() -> None: + dense_pages.zero_() + _copy_owned_to_dense( + dense_pages=dense_pages, + src_pages=src_pages, + owned_slots=owned_slots, + owned_count=owned_count, + ) + dist.all_reduce(dense_pages, op=dist.ReduceOp.SUM, group=group) + + def packed_all_gather_scatter() -> None: + _copy_owned_to_packed( + local_packed=local_packed, + src_pages=src_pages, + owned_count=owned_count, + ) + dist.all_gather_into_tensor(gathered, local_packed, group=group) + _scatter_rank_major_to_dense( + dense_pages=dense_pages, + gathered=gathered, + prefix_pages=prefix_pages, + owner_counts=owner_counts, + world_size=world_size, + device=device, + ) + + def packed_all_gather_rank_major() -> None: + _copy_owned_to_packed( + local_packed=local_packed, + src_pages=src_pages, + owned_count=owned_count, + ) + dist.all_gather_into_tensor(gathered, local_packed, group=group) + + paths: list[tuple[str, Callable[[], None]]] = [ + ("dense_all_reduce", dense_all_reduce), + ("packed_all_gather_scatter", packed_all_gather_scatter), + ("packed_all_gather_rank_major", packed_all_gather_rank_major), + ] + results: list[BenchResult] = [] + for name, fn in paths: + gpu_ms, cpu_ms = _time_path( + fn, + warmup=cfg.warmup, + repeat=cfg.repeat, + stream=stream, + group=group, + device=device, + ) + gpu, cpu = _summarize_across_ranks( + local_gpu_ms=gpu_ms, + local_cpu_ms=cpu_ms, + group=group, + device=device, + ) + results.append( + BenchResult( + payload=cfg.payload, + prefix_tokens=cfg.prefix_tokens, + prefix_pages=prefix_pages, + path=name, + full_mib=full_mib, + packed_mib_per_rank=packed_mib, + gpu_ms_avg=gpu["avg"], + gpu_ms_p50=gpu["p50"], + gpu_ms_min=gpu["min"], + gpu_ms_max_rank_avg=gpu["max_rank_avg"], + cpu_ms_avg=cpu["avg"], + cpu_ms_p50=cpu["p50"], + cpu_ms_max_rank_avg=cpu["max_rank_avg"], + ) + ) + return results + + +def _print_result(result: BenchResult) -> None: + print(json.dumps(asdict(result), sort_keys=True), flush=True) + print( + f"{result.payload:5s} tokens={result.prefix_tokens:6d} " + f"pages={result.prefix_pages:5d} {result.path:28s} " + f"gpu_avg={result.gpu_ms_avg:8.3f}ms " + f"gpu_max_rank={result.gpu_ms_max_rank_avg:8.3f}ms " + f"cpu_avg={result.cpu_ms_avg:8.3f}ms " + f"full={result.full_mib:8.1f}MiB " + f"packed/rank={result.packed_mib_per_rank:8.1f}MiB", + flush=True, + ) + + +def _payloads(value: str) -> Iterable[str]: + if value == "both": + return ("mla", "index") + return (value,) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark CP shared-KV prefetch all-reduce vs all-gather layouts." + ) + parser.add_argument("--payload", choices=("mla", "index", "both"), default="both") + parser.add_argument( + "--tokens", + type=int, + nargs="+", + default=[4096, 8192, 16384, 32768, 40320, 65536], + help="Prefix token counts to benchmark.", + ) + parser.add_argument("--page-size", type=int, default=64) + parser.add_argument("--kv-dim", type=int, default=576) + parser.add_argument( + "--index-page-bytes", + type=int, + default=8448, + help="Bytes per index page. GLM5 NSA observed default is 8448.", + ) + parser.add_argument("--dtype", choices=("bf16", "fp16", "fp32"), default="bf16") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--repeat", type=int, default=20) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + rank, world_size, local_rank = _rank_env() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + dist.init_process_group("nccl", init_method="env://") + group = dist.group.WORLD + + if rank == 0: + print( + "CP shared-KV prefetch collective benchmark " + f"world_size={world_size} tokens={args.tokens} payload={args.payload}", + flush=True, + ) + + try: + for payload in _payloads(args.payload): + for tokens in args.tokens: + cfg = CaseConfig( + payload=payload, + prefix_tokens=int(tokens), + page_size=args.page_size, + kv_dim=args.kv_dim, + index_page_bytes=args.index_page_bytes, + dtype=args.dtype, + warmup=args.warmup, + repeat=args.repeat, + ) + for result in _bench_case( + cfg, + group=group, + device=device, + rank=rank, + world_size=world_size, + ): + if rank == 0: + _print_result(result) + finally: + dist.barrier(group=group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_no_collective_capacity_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_no_collective_capacity_plan.md index b933d7f3f..56c0d5a74 100644 --- a/docs/advanced_features/nsa_prefill_cp_hicache_no_collective_capacity_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_hicache_no_collective_capacity_plan.md @@ -185,12 +185,13 @@ waits for D2H write of layer `i`. - Planner skips device-valid nodes, pending backup nodes, host-protected nodes, pinned nodes, malformed metadata, and non-host-backed leaves. - Current eviction behavior is unchanged. -- **Online debug option added:** `SGLANG_CP_HICACHE_CAPACITY_DEBUG=1`. - - When enabled, write admission computes the deterministic pre-reserve plan and - compares its eviction-needed decision with the current collective - `reserve_slots_max` result. - - Logs use `[HiCache-capacity-debug] write_admission_compare ...`. - - This is observer-only and does not remove or add collectives. +- **Write admission is now owner-lane authoritative.** + - The old observer/debug scalar planner (`current_compatible_need` plus + `SGLANG_CP_HICACHE_CAPACITY_DEBUG`) has been removed from the hot path. + - Write admission uses the per-owner target/draft deficit vector directly, + evicts deterministic host victims only for lanes with deficit, and fails + fast if a subsequent local reservation failure contradicts the vector view. + - No CP capacity all-reduce is used for write admission. ### P1: Observer-only ledger diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md index 4682f75f3..176018fc6 100644 --- a/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md @@ -226,17 +226,12 @@ all_reduce per node victim in a tight host-eviction loop all_reduce every scheduler tick when no completion prefix advanced ``` -Current correctness note: CP host reservation now synchronizes -`required_host_slots` with `all_reduce(MAX)` before the host-eviction retry. -This forces every rank into the same reserve/evict/retry branch and avoids -collective mismatches when one rank is host-full and another rank reserves -successfully. It is intentionally a coarse slow-path collective, not a -per-layer collective, but it can become a performance cost when host pressure -is frequent because every reservation failure pays at least one rank-wide MAX -sync and retry failures pay a second one. Treat this as a correctness guard to -be amortized later with batched reservation epochs, deterministic host -watermarks, or less frequent proactive host eviction; do not move it into the -per-layer data path. +Current correctness note: CP host reservation no longer uses the old +`required_host_slots` scalar `all_reduce(MAX)`. The owner-lane vector computed +from `CpHiCacheNodeMetadata.page_owners` is now the authoritative admission +view for target and draft host capacity. A reservation failure after this +vector predicts no deficit is treated as an invariant violation rather than +falling back to another collective. ## Per-Layer Backup Data Plane @@ -466,10 +461,11 @@ paths: 2. **Pending split behavior.** The chosen first pass is defer/requeue the request that would split a node with pending backup. Do not split the in-flight backup op, including for future `bs > 1`. -3. **Ack batching threshold.** Current `writing_check()` can all-reduce on - every progress poll. The first per-layer implementation should keep one - final logical ack per node and check it at final visibility time, not per - layer; batching threshold for final commit remains a later performance pass. +3. **Ack batching threshold.** `writing_check()` no longer enters the TP + all-reduce while a write is ongoing but no final ack exists. The remaining + `MIN` is the final host-visibility barrier for ack entries that do exist: + it should not run per layer, and removing it requires a different global + visibility protocol rather than a local-only commit. 4. **Failure policy under reserve mismatch.** Capacity pressure should skip backup; malformed target/draft metadata should fail fast. If a local rank fails reservation after deterministic host eviction while another succeeds, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d29199e73..3a7f612fa 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -215,7 +215,6 @@ class Envs: SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1) SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False) SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False) - SGLANG_CP_HICACHE_CAPACITY_DEBUG = EnvBool(False) SGLANG_DISABLE_TAI_BIGRAM = EnvBool(False) SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False) SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index 0474b9888..e2c05bc83 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import time from dataclasses import dataclass from typing import Any, Optional @@ -107,6 +108,10 @@ def _debug_owned_pages_count( try: if logical_pages.numel() == 0: return 0 + if logical_pages.is_cuda: + # Temporary CPU timing logs must not add a CUDA synchronization via + # Tensor.item(); skip exact debug counts for GPU metadata. + return -1 return int(layout.owned_pages_mask(logical_pages).sum().item()) except Exception: logger.exception("Failed to count CP shared KV prefetch owned pages.") @@ -124,6 +129,145 @@ def _debug_handle_keys( return tuple(sorted(handles.keys())) +def _cpu_timing_start() -> float: + if not cp_shared_kv_mla_prefetch_log_enabled(): + return 0.0 + return time.perf_counter() + + +def _cpu_timing_ms(start: float) -> float: + if start <= 0.0: + return -1.0 + return (time.perf_counter() - start) * 1000.0 + + +@dataclass +class _PrefetchCpuTiming: + start_count: int = 0 + start_total_ms: float = 0.0 + start_max_ms: float = 0.0 + get_total_ms: float = 0.0 + materialize_total_ms: float = 0.0 + reduce_enqueue_total_ms: float = 0.0 + consume_count: int = 0 + consume_total_ms: float = 0.0 + consume_wait_total_ms: float = 0.0 + consume_suffix_total_ms: float = 0.0 + consume_remap_total_ms: float = 0.0 + + def record_start( + self, + *, + total_ms: float, + get_ms: float, + materialize_ms: float, + reduce_enqueue_ms: float, + ) -> None: + if total_ms < 0.0: + return + self.start_count += 1 + self.start_total_ms += total_ms + self.start_max_ms = max(self.start_max_ms, total_ms) + self.get_total_ms += max(get_ms, 0.0) + self.materialize_total_ms += max(materialize_ms, 0.0) + self.reduce_enqueue_total_ms += max(reduce_enqueue_ms, 0.0) + + def record_consume( + self, + *, + total_ms: float, + wait_ms: float, + suffix_ms: float, + remap_ms: float, + ) -> None: + if total_ms < 0.0: + return + self.consume_count += 1 + self.consume_total_ms += total_ms + self.consume_wait_total_ms += max(wait_ms, 0.0) + self.consume_suffix_total_ms += max(suffix_ms, 0.0) + self.consume_remap_total_ms += max(remap_ms, 0.0) + + +def _log_prefetch_cpu_start( + *, + log_fn: Any, + layout: CpSharedKVLayout, + timing: _PrefetchCpuTiming, + path: str, + layer_id: int, + total_ms: float, + get_ms: float, + materialize_ms: float, + reduce_enqueue_ms: float, + prefix_pages: int, + total_slots: int, + dense_units: int, +) -> None: + if not cp_shared_kv_mla_prefetch_log_enabled() or total_ms < 0.0: + return + if cp_shared_kv_mla_prefetch_should_log_layer(layer_id): + log_fn( + "cpu_timing path=%s stage=start layer=%s total_ms=%.3f " + "get_ms=%.3f materialize_ms=%.3f reduce_enqueue_ms=%.3f " + "prefix_pages=%s total_slots=%s dense_units=%s", + path, + layer_id, + total_ms, + get_ms, + materialize_ms, + reduce_enqueue_ms, + prefix_pages, + total_slots, + dense_units, + ) + if layout.cp_rank == 0 and timing.start_count > 0 and timing.start_count % 16 == 0: + starts = timing.start_count + log_fn( + "cpu_summary path=%s starts=%s avg_start_ms=%.3f max_start_ms=%.3f " + "avg_get_ms=%.3f avg_materialize_ms=%.3f avg_reduce_enqueue_ms=%.3f", + path, + starts, + timing.start_total_ms / starts, + timing.start_max_ms, + timing.get_total_ms / starts, + timing.materialize_total_ms / starts, + timing.reduce_enqueue_total_ms / starts, + ) + + +def _log_prefetch_cpu_consume( + *, + log_fn: Any, + timing: _PrefetchCpuTiming, + path: str, + layer_id: int, + total_ms: float, + wait_ms: float, + suffix_ms: float, + remap_ms: float, + prefix_pages: int, + suffix_slots: int, +) -> None: + if not cp_shared_kv_mla_prefetch_log_enabled() or total_ms < 0.0: + return + if cp_shared_kv_mla_prefetch_should_log_layer(layer_id): + log_fn( + "cpu_timing path=%s stage=consume layer=%s total_ms=%.3f " + "wait_event_ms=%.3f suffix_ms=%.3f remap_ms=%.3f " + "prefix_pages=%s suffix_slots=%s consume_count=%s", + path, + layer_id, + total_ms, + wait_ms, + suffix_ms, + remap_ms, + prefix_pages, + suffix_slots, + timing.consume_count, + ) + + @dataclass class CpSharedKVMlaPrefetchHandle: layer_id: int @@ -174,6 +318,7 @@ class CpSharedKVMlaPrefetcher: self.handles: dict[int, CpSharedKVMlaPrefetchHandle] = {} self.pending_attention_handle: Optional[CpSharedKVMlaPrefetchHandle] = None self.disabled = False + self._cpu_timing = _PrefetchCpuTiming() @classmethod def maybe_create( @@ -280,6 +425,8 @@ class CpSharedKVMlaPrefetcher: return None prefetch_stream = stream if stream is not None else torch.cuda.Stream() + create_cpu = _cpu_timing_start() + get_cpu = _cpu_timing_start() try: first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0)) kv_cache = _prefetch_pool_get_key_buffer( @@ -288,6 +435,8 @@ class CpSharedKVMlaPrefetcher: stream=prefetch_stream, path="mla", ) + get_ms = _cpu_timing_ms(get_cpu) + remap_cpu = _cpu_timing_start() remap = get_or_build_shared_token_kv_slot_remap( forward_batch, kv_cache=kv_cache, @@ -295,6 +444,7 @@ class CpSharedKVMlaPrefetcher: layout=layout, page_size=page_size, ) + remap_ms = _cpu_timing_ms(remap_cpu) except Exception: logger.exception("Failed to initialize CP shared KV MLA prefetcher.") return None @@ -316,6 +466,20 @@ class CpSharedKVMlaPrefetcher: remap.dense_num_pages, page_size, ) + create_total_ms = _cpu_timing_ms(create_cpu) + _prefetch_log( + "cpu_timing path=mla stage=create cp_rank=%s cp_size=%s " + "total_ms=%.3f get_ms=%.3f remap_ms=%.3f prefix_pages=%s " + "total_slots=%s dense_pages=%s", + layout.cp_rank, + layout.cp_size, + create_total_ms, + get_ms, + remap_ms, + prefix_pages, + int(remap.slot_logical_pages.numel()), + remap.dense_num_pages, + ) return cls( layout=layout, @@ -386,9 +550,13 @@ class CpSharedKVMlaPrefetcher: ) return None + consume_cpu = _cpu_timing_start() + wait_cpu = _cpu_timing_start() torch.cuda.current_stream().wait_event(handle.event) + wait_ms = _cpu_timing_ms(wait_cpu) dense_kv_cache = handle.dense_kv_cache suffix_slots = self.total_slots - self.prefix_pages + suffix_ms = 0.0 if self.prefix_pages < self.total_slots: self._log_layer( @@ -400,6 +568,7 @@ class CpSharedKVMlaPrefetcher: self.total_slots, suffix_slots, ) + suffix_cpu = _cpu_timing_start() materialize_local_token_kv_page_slots_into( kv_cache=kv_cache, dense_kv_cache=dense_kv_cache, @@ -430,6 +599,7 @@ class CpSharedKVMlaPrefetcher: suffix_rows.start, suffix_rows.stop, ) + suffix_ms = _cpu_timing_ms(suffix_cpu) self._log_layer( layer_id, @@ -440,6 +610,7 @@ class CpSharedKVMlaPrefetcher: int(dense_kv_cache.shape[0]), ) + remap_cpu = _cpu_timing_start() logical_locs = filter_locs_mappable_to_physical_pool( logical_locs=logical_locs, layout=self.layout, @@ -450,6 +621,26 @@ class CpSharedKVMlaPrefetcher: page_inverse=self.page_inverse, page_size=self.page_size, ) + remap_ms = _cpu_timing_ms(remap_cpu) + total_ms = _cpu_timing_ms(consume_cpu) + self._cpu_timing.record_consume( + total_ms=total_ms, + wait_ms=wait_ms, + suffix_ms=suffix_ms, + remap_ms=remap_ms, + ) + _log_prefetch_cpu_consume( + log_fn=self._log, + timing=self._cpu_timing, + path="mla", + layer_id=layer_id, + total_ms=total_ms, + wait_ms=wait_ms, + suffix_ms=suffix_ms, + remap_ms=remap_ms, + prefix_pages=self.prefix_pages, + suffix_slots=suffix_slots, + ) return dense_kv_cache, dense_locs def start_next_layer_prefix( @@ -496,6 +687,8 @@ class CpSharedKVMlaPrefetcher: ) return + start_cpu = _cpu_timing_start() + get_cpu = _cpu_timing_start() try: kv_cache = _prefetch_pool_get_key_buffer( token_to_kv_pool=token_to_kv_pool, @@ -503,6 +696,7 @@ class CpSharedKVMlaPrefetcher: stream=self.stream, path="mla", ) + get_ms = _cpu_timing_ms(get_cpu) except Exception: logger.exception( "Failed to get next-layer KV cache for CP shared KV MLA prefetch." @@ -517,33 +711,36 @@ class CpSharedKVMlaPrefetcher: try: current_stream = torch.cuda.current_stream() - self.stream.wait_stream(current_stream) prefix_rows = slot_range_to_token_slice( self.page_size, 0, self.prefix_pages, ) + materialize_cpu = _cpu_timing_start() + dense_kv_cache = kv_cache.new_zeros( + (self.dense_num_pages * self.page_size, *kv_cache.shape[1:]) + ) + self._log_next_layer( + next_layer_id, + "start_prefix_begin next_layer=%s start_slot=0 end_slot=%s " + "dense_rows=%s", + next_layer_id, + self.prefix_pages, + int(dense_kv_cache.shape[0]), + ) + materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + page_size=self.page_size, + start_slot=0, + end_slot=self.prefix_pages, + ) + materialize_ms = _cpu_timing_ms(materialize_cpu) + reduce_cpu = _cpu_timing_start() + self.stream.wait_stream(current_stream) with torch.cuda.stream(self.stream): - dense_kv_cache = kv_cache.new_zeros( - (self.dense_num_pages * self.page_size, *kv_cache.shape[1:]) - ) - self._log_next_layer( - next_layer_id, - "start_prefix_begin next_layer=%s start_slot=0 end_slot=%s " - "dense_rows=%s", - next_layer_id, - self.prefix_pages, - int(dense_kv_cache.shape[0]), - ) - materialize_local_token_kv_page_slots_into( - kv_cache=kv_cache, - dense_kv_cache=dense_kv_cache, - slot_logical_pages=self.slot_logical_pages, - layout=self.layout, - page_size=self.page_size, - start_slot=0, - end_slot=self.prefix_pages, - ) event = _all_reduce_materialized_buffer_async( dense_kv_cache[prefix_rows], cp_size=self.layout.cp_size, @@ -553,6 +750,7 @@ class CpSharedKVMlaPrefetcher: nvtx_cp_rank=self.layout.cp_rank, nvtx_rows=(prefix_rows.start, prefix_rows.stop), ) + reduce_enqueue_ms = _cpu_timing_ms(reduce_cpu) if event is None: self.disabled = True logger.warning( @@ -571,6 +769,27 @@ class CpSharedKVMlaPrefetcher: next_layer_id, ) return + total_ms = _cpu_timing_ms(start_cpu) + self._cpu_timing.record_start( + total_ms=total_ms, + get_ms=get_ms, + materialize_ms=materialize_ms, + reduce_enqueue_ms=reduce_enqueue_ms, + ) + _log_prefetch_cpu_start( + log_fn=self._log, + layout=self.layout, + timing=self._cpu_timing, + path="mla", + layer_id=next_layer_id, + total_ms=total_ms, + get_ms=get_ms, + materialize_ms=materialize_ms, + reduce_enqueue_ms=reduce_enqueue_ms, + prefix_pages=self.prefix_pages, + total_slots=self.total_slots, + dense_units=int(dense_kv_cache.shape[0]), + ) except Exception: logger.exception("Failed to start CP shared KV MLA prefix prefetch.") self.disabled = True @@ -717,6 +936,7 @@ class CpSharedKVIndexPrefetcher: self.handles: dict[int, CpSharedKVIndexPrefetchHandle] = {} self.pending_attention_handle: Optional[CpSharedKVIndexPrefetchHandle] = None self.disabled = False + self._cpu_timing = _PrefetchCpuTiming() @classmethod def maybe_create( @@ -864,6 +1084,8 @@ class CpSharedKVIndexPrefetcher: return None prefetch_stream = stream if stream is not None else torch.cuda.Stream() + create_cpu = _cpu_timing_start() + get_cpu = _cpu_timing_start() try: first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0)) page_buffer = _prefetch_pool_get_index_buffer( @@ -871,12 +1093,15 @@ class CpSharedKVIndexPrefetcher: layer_id=first_layer_id, stream=prefetch_stream, ) + get_ms = _cpu_timing_ms(get_cpu) + remap_cpu = _cpu_timing_start() remap = get_or_build_shared_paged_buffer_slot_remap( forward_batch, page_buffer=page_buffer, logical_pages=real_page_table, layout=layout, ) + remap_ms = _cpu_timing_ms(remap_cpu) except Exception as exc: _index_prefetch_fallback_log( "init_exception", @@ -903,6 +1128,20 @@ class CpSharedKVIndexPrefetcher: remap.dense_num_pages, page_size, ) + create_total_ms = _cpu_timing_ms(create_cpu) + _prefetch_log( + "cpu_timing path=index stage=create cp_rank=%s cp_size=%s " + "total_ms=%.3f get_ms=%.3f remap_ms=%.3f prefix_pages=%s " + "total_slots=%s dense_pages=%s", + layout.cp_rank, + layout.cp_size, + create_total_ms, + get_ms, + remap_ms, + prefix_pages, + int(remap.slot_logical_pages.numel()), + remap.dense_num_pages, + ) return cls( layout=layout, @@ -972,9 +1211,13 @@ class CpSharedKVIndexPrefetcher: ) return None + consume_cpu = _cpu_timing_start() + wait_cpu = _cpu_timing_start() torch.cuda.current_stream().wait_event(handle.event) + wait_ms = _cpu_timing_ms(wait_cpu) dense_page_buffer = handle.dense_page_buffer suffix_slots = self.total_slots - self.prefix_pages + suffix_ms = 0.0 if self.prefix_pages < self.total_slots: self._log_layer( @@ -986,6 +1229,7 @@ class CpSharedKVIndexPrefetcher: self.total_slots, suffix_slots, ) + suffix_cpu = _cpu_timing_start() materialize_local_paged_buffer_page_slots_into( page_buffer=page_buffer, dense_page_buffer=dense_page_buffer, @@ -1014,6 +1258,7 @@ class CpSharedKVIndexPrefetcher: suffix_rows.start, suffix_rows.stop, ) + suffix_ms = _cpu_timing_ms(suffix_cpu) self._log_layer( layer_id, @@ -1024,6 +1269,7 @@ class CpSharedKVIndexPrefetcher: int(dense_page_buffer.shape[0]), ) + remap_cpu = _cpu_timing_start() logical_pages = filter_pages_mappable_to_physical_pool( logical_pages=logical_pages, layout=self.layout, @@ -1033,6 +1279,26 @@ class CpSharedKVIndexPrefetcher: logical_pages, page_inverse=self.page_inverse, ) + remap_ms = _cpu_timing_ms(remap_cpu) + total_ms = _cpu_timing_ms(consume_cpu) + self._cpu_timing.record_consume( + total_ms=total_ms, + wait_ms=wait_ms, + suffix_ms=suffix_ms, + remap_ms=remap_ms, + ) + _log_prefetch_cpu_consume( + log_fn=self._log, + timing=self._cpu_timing, + path="index", + layer_id=layer_id, + total_ms=total_ms, + wait_ms=wait_ms, + suffix_ms=suffix_ms, + remap_ms=remap_ms, + prefix_pages=self.prefix_pages, + suffix_slots=suffix_slots, + ) return dense_page_buffer, dense_pages def start_next_layer_prefix( @@ -1079,12 +1345,15 @@ class CpSharedKVIndexPrefetcher: ) return + start_cpu = _cpu_timing_start() + get_cpu = _cpu_timing_start() try: page_buffer = _prefetch_pool_get_index_buffer( token_to_kv_pool=token_to_kv_pool, layer_id=next_layer_id, stream=self.stream, ) + get_ms = _cpu_timing_ms(get_cpu) except Exception: logger.exception( "Failed to get next-layer index buffer for CP shared KV index prefetch." @@ -1099,28 +1368,31 @@ class CpSharedKVIndexPrefetcher: try: current_stream = torch.cuda.current_stream() - self.stream.wait_stream(current_stream) prefix_rows = slot_range_to_page_slice(0, self.prefix_pages) + materialize_cpu = _cpu_timing_start() + dense_page_buffer = page_buffer.new_zeros( + (self.dense_num_pages, *page_buffer.shape[1:]) + ) + self._log_next_layer( + next_layer_id, + "index_start_prefix_begin next_layer=%s start_slot=0 " + "end_slot=%s dense_pages=%s", + next_layer_id, + self.prefix_pages, + int(dense_page_buffer.shape[0]), + ) + materialize_local_paged_buffer_page_slots_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + start_slot=0, + end_slot=self.prefix_pages, + ) + materialize_ms = _cpu_timing_ms(materialize_cpu) + reduce_cpu = _cpu_timing_start() + self.stream.wait_stream(current_stream) with torch.cuda.stream(self.stream): - dense_page_buffer = page_buffer.new_zeros( - (self.dense_num_pages, *page_buffer.shape[1:]) - ) - self._log_next_layer( - next_layer_id, - "index_start_prefix_begin next_layer=%s start_slot=0 " - "end_slot=%s dense_pages=%s", - next_layer_id, - self.prefix_pages, - int(dense_page_buffer.shape[0]), - ) - materialize_local_paged_buffer_page_slots_into( - page_buffer=page_buffer, - dense_page_buffer=dense_page_buffer, - slot_logical_pages=self.slot_logical_pages, - layout=self.layout, - start_slot=0, - end_slot=self.prefix_pages, - ) event = _all_reduce_materialized_buffer_async( dense_page_buffer[prefix_rows], cp_size=self.layout.cp_size, @@ -1130,6 +1402,7 @@ class CpSharedKVIndexPrefetcher: nvtx_cp_rank=self.layout.cp_rank, nvtx_rows=(prefix_rows.start, prefix_rows.stop), ) + reduce_enqueue_ms = _cpu_timing_ms(reduce_cpu) if event is None: self.disabled = True _index_prefetch_fallback_log( @@ -1148,6 +1421,27 @@ class CpSharedKVIndexPrefetcher: next_layer_id, ) return + total_ms = _cpu_timing_ms(start_cpu) + self._cpu_timing.record_start( + total_ms=total_ms, + get_ms=get_ms, + materialize_ms=materialize_ms, + reduce_enqueue_ms=reduce_enqueue_ms, + ) + _log_prefetch_cpu_start( + log_fn=self._log, + layout=self.layout, + timing=self._cpu_timing, + path="index", + layer_id=next_layer_id, + total_ms=total_ms, + get_ms=get_ms, + materialize_ms=materialize_ms, + reduce_enqueue_ms=reduce_enqueue_ms, + prefix_pages=self.prefix_pages, + total_slots=self.total_slots, + dense_units=int(dense_page_buffer.shape[0]), + ) except Exception: logger.exception("Failed to start CP shared KV index prefix prefetch.") self.disabled = True diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 418301926..3c9bd7831 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -303,6 +303,17 @@ class CpLoadBackEvictionPlan: remaining_deficit_by_owner: Tuple[int, ...] +@dataclass(frozen=True) +class CpWriteAdmission: + node_id: int + phase: str + required_by_owner: Tuple[int, ...] + target_available_by_owner: Tuple[int, ...] + draft_available_by_owner: Tuple[int, ...] + deficit_by_owner: Tuple[int, ...] + eviction_plan: CpHiCacheEvictionPlan + + class HiCachePendingBackupSplit(Exception): def __init__(self, node: TreeNode): self.node = node @@ -523,10 +534,6 @@ class HiRadixCache(RadixCache): # calls are on latency-sensitive scheduler paths; keep logging coarse # enough to avoid per-request spam while still exposing hot collectives. self._cp_hicache_collective_stats = {} - self._cp_hicache_capacity_debug = ( - envs.SGLANG_CP_HICACHE_CAPACITY_DEBUG.get() - ) - self._cp_hicache_capacity_debug_stats = {} # track per-request tokens loaded from storage (L3 hits) # key: request_id, value: number of tokens actually loaded from storage self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {} @@ -949,7 +956,7 @@ class HiRadixCache(RadixCache): host_hit_len=plan.host_hit_len, ) - def _cp_device_leaf_is_load_back_victim(self, node: TreeNode) -> bool: + def _cp_device_node_is_load_back_victim_base(self, node: TreeNode) -> bool: if node == getattr(self, "root_node", None): return False if getattr(node, "lock_ref", 0) > 0: @@ -969,6 +976,21 @@ class HiRadixCache(RadixCache): return False return True + def _cp_device_node_is_load_back_victim_after_plan( + self, node: TreeNode, planned_evicted_nodes: set + ) -> bool: + if not self._cp_device_node_is_load_back_victim_base(node): + return False + for child in getattr(node, "children", {}).values(): + if child in planned_evicted_nodes: + continue + if not getattr(child, "evicted", False): + return False + return True + + def _cp_device_leaf_is_load_back_victim(self, node: TreeNode) -> bool: + return self._cp_device_node_is_load_back_victim_after_plan(node, set()) + def _cp_load_back_node_owner_page_counts( self, node: TreeNode, cp_size: int ) -> Tuple[int, ...]: @@ -987,6 +1009,32 @@ class HiRadixCache(RadixCache): owners = torch.remainder(logical_pages - 1, cp_size) return tuple(int((owners == owner).sum().item()) for owner in range(cp_size)) + def _cp_load_back_ancestor_unlock_contribution( + self, + node: TreeNode, + deficits: List[int], + planned_evicted_nodes: set, + cp_size: int, + ) -> int: + ancestor = getattr(node, "parent", None) + while ancestor is not None and ancestor != getattr(self, "root_node", None): + if ancestor in planned_evicted_nodes: + return 0 + if getattr(ancestor, "value", None) is None: + ancestor = getattr(ancestor, "parent", None) + continue + if not self._cp_device_node_is_load_back_victim_base(ancestor): + return 0 + counts = self._cp_load_back_node_owner_page_counts(ancestor, cp_size) + contribution = sum( + min(int(count), int(deficit)) + for count, deficit in zip(counts, deficits) + ) + if contribution > 0: + return int(contribution) + ancestor = getattr(ancestor, "parent", None) + return 0 + def _plan_cp_load_back_owner_lane_evictions( self, plan: CpLoadBackPlan ) -> CpLoadBackEvictionPlan: @@ -994,29 +1042,39 @@ class HiRadixCache(RadixCache): cp_size = len(deficits) planned_freed = [0 for _ in range(cp_size)] victims: List[TreeNode] = [] - used_node_ids = set() + planned_evicted_nodes = set() + candidate_nodes = set(getattr(self, "evictable_leaves", set())) while any(v > 0 for v in deficits): best_node = None best_counts = None best_score = None - for node in list(getattr(self, "evictable_leaves", set())): - node_id = getattr(node, "id", None) - if node_id in used_node_ids: + for node in list(candidate_nodes): + if node in planned_evicted_nodes: continue - if not self._cp_device_leaf_is_load_back_victim(node): + if not self._cp_device_node_is_load_back_victim_after_plan( + node, planned_evicted_nodes + ): continue counts = self._cp_load_back_node_owner_page_counts(node, cp_size) contribution = sum( min(int(count), int(deficit)) for count, deficit in zip(counts, deficits) ) + unlock_contribution = 0 if contribution <= 0: + unlock_contribution = ( + self._cp_load_back_ancestor_unlock_contribution( + node, deficits, planned_evicted_nodes, cp_size + ) + ) + if contribution <= 0 and unlock_contribution <= 0: continue score = ( -int(contribution), + -int(unlock_contribution), self.eviction_strategy.get_priority(node), - int(node_id or 0), + int(getattr(node, "id", 0) or 0), ) if best_score is None or score < best_score: best_score = score @@ -1027,10 +1085,23 @@ class HiRadixCache(RadixCache): break victims.append(best_node) - used_node_ids.add(getattr(best_node, "id", None)) + planned_evicted_nodes.add(best_node) + candidate_nodes.discard(best_node) for owner, count in enumerate(best_counts): planned_freed[owner] += int(count) deficits[owner] = max(0, deficits[owner] - int(count)) + ancestor = getattr(best_node, "parent", None) + while ancestor is not None and ancestor != getattr(self, "root_node", None): + if ancestor in planned_evicted_nodes: + break + if self._cp_device_node_is_load_back_victim_after_plan( + ancestor, planned_evicted_nodes + ): + candidate_nodes.add(ancestor) + break + if getattr(ancestor, "value", None) is not None: + break + ancestor = getattr(ancestor, "parent", None) return CpLoadBackEvictionPlan( victims=tuple(victims), @@ -1127,15 +1198,9 @@ class HiRadixCache(RadixCache): ) return refreshed - def _cp_capacity_debug_enabled(self) -> bool: - return bool( - getattr(self, "_cp_hicache_capacity_debug", False) - or envs.SGLANG_CP_HICACHE_CAPACITY_DEBUG.get() - ) - - def _cp_build_write_admission_plan( + def _cp_build_write_admission( self, device_indices: torch.Tensor, *, node_id: int, phase: str - ) -> dict: + ) -> CpWriteAdmission: required = self._cp_required_host_tokens_by_rank(device_indices) snapshot = self._cp_host_capacity_snapshot() target_available = self._cp_host_available_tokens_by_rank(snapshot) @@ -1147,90 +1212,15 @@ class HiRadixCache(RadixCache): max(0, req - avail) for req, avail in zip(required, draft_available) ) deficit = tuple(max(a, b) for a, b in zip(target_deficit, draft_deficit)) - current_compatible_need = tuple( - req if missing > 0 else 0 for req, missing in zip(required, deficit) - ) eviction_plan = self._plan_cp_host_evictions(deficit) - return { - "node_id": node_id, - "phase": phase, - "required": required, - "target_available": target_available, - "draft_available": draft_available, - "deficit": deficit, - "current_compatible_need": current_compatible_need, - "eviction_plan": eviction_plan, - } - - def _cp_debug_build_write_admission_plan( - self, device_indices: torch.Tensor, *, node_id: int, phase: str - ) -> Optional[dict]: - if not self._cp_capacity_debug_enabled(): - return None - try: - return self._cp_build_write_admission_plan( - device_indices, node_id=node_id, phase=phase - ) - except Exception as exc: - logger.warning( - "[HiCache-capacity-debug] failed to build write admission plan: node_id=%d phase=%s error=%s", - node_id, - phase, - exc, - ) - return None - - def _cp_debug_compare_write_admission( - self, - plan: Optional[dict], - *, - result, - planned_required_slots: int, - ) -> None: - if plan is None or not self._cp_capacity_debug_enabled(): - return - stats = getattr(self, "_cp_hicache_capacity_debug_stats", None) - if stats is None: - stats = self._cp_hicache_capacity_debug_stats = {} - tag = "write_admission" - entry = stats.setdefault(tag, {"count": 0, "mismatch": 0}) - entry["count"] += 1 - - planned_required_slots = max(plan["current_compatible_need"], default=0) - planned_needs_eviction = planned_required_slots > 0 - reservation_failed = isinstance(result, HiCacheWriteFailure) - mismatch = ( - reservation_failed - and not planned_needs_eviction - and all(v <= 0 for v in plan["eviction_plan"].remaining_deficit) - ) - if mismatch: - entry["mismatch"] += 1 - should_log = mismatch or entry["count"] == 1 or entry["count"] % 128 == 0 - if not should_log: - return - - eviction_plan = plan["eviction_plan"] - log_fn = logger.warning if mismatch else logger.info - log_fn( - "[HiCache-capacity-debug] write_admission_compare node_id=%d phase=%s " - "count=%d mismatch=%d planned_required=%d local_failed=%s " - "planned_current_required=%d required=%s target_avail=%s draft_avail=%s " - "deficit=%s planned_freed=%s remaining_deficit=%s victims=%s", - plan["node_id"], - plan["phase"], - entry["count"], - entry["mismatch"], - int(planned_required_slots), - isinstance(result, HiCacheWriteFailure), - int(planned_required_slots), - plan["required"], - plan["target_available"], - plan["draft_available"], - plan["deficit"], - eviction_plan.planned_freed, - eviction_plan.remaining_deficit, - [getattr(node, "id", None) for node in eviction_plan.victims], + return CpWriteAdmission( + node_id=node_id, + phase=phase, + required_by_owner=required, + target_available_by_owner=target_available, + draft_available_by_owner=draft_available, + deficit_by_owner=deficit, + eviction_plan=eviction_plan, ) def shutdown(self): @@ -1753,10 +1743,10 @@ class HiRadixCache(RadixCache): self.dec_node_lock_ref(node) return node - def _evict_cp_host_plan( - self, plan: dict, *, node_id: int, phase: str + def _evict_cp_host_for_write_admission( + self, admission: CpWriteAdmission, *, node_id: int, phase: str ) -> bool: - eviction_plan = plan["eviction_plan"] + eviction_plan = admission.eviction_plan if any(v > 0 for v in eviction_plan.remaining_deficit): logger.warning( "[CP_HICACHE_FALLBACK][cp_host_reservation_plan_insufficient] " @@ -1765,10 +1755,10 @@ class HiRadixCache(RadixCache): "planned_freed=%s remaining_deficit=%s victims=%s", node_id, phase, - plan["required"], - plan["target_available"], - plan["draft_available"], - plan["deficit"], + admission.required_by_owner, + admission.target_available_by_owner, + admission.draft_available_by_owner, + admission.deficit_by_owner, eviction_plan.planned_freed, eviction_plan.remaining_deficit, [getattr(node, "id", None) for node in eviction_plan.victims], @@ -1810,59 +1800,55 @@ class HiRadixCache(RadixCache): def _reserve_write_cp_indices_no_collective( self, device_indices: torch.Tensor, node_id: int ): - plan = self._cp_build_write_admission_plan( + admission = self._cp_build_write_admission( device_indices, node_id=node_id, phase="initial" ) - planned_required_slots = max(plan["current_compatible_need"], default=0) - if planned_required_slots > 0: - if not self._evict_cp_host_plan(plan, node_id=node_id, phase="initial"): - return HiCacheWriteFailure(required_host_slots=planned_required_slots) + if any(v > 0 for v in admission.deficit_by_owner): + if not self._evict_cp_host_for_write_admission( + admission, node_id=node_id, phase="initial" + ): + return HiCacheWriteFailure( + required_host_slots=max(admission.eviction_plan.remaining_deficit) + ) result = self.cache_controller.reserve_write_cp( device_indices=device_indices, node_id=node_id, ) - debug_plan = plan if self._cp_capacity_debug_enabled() else None - self._cp_debug_compare_write_admission( - debug_plan, - result=result, - planned_required_slots=planned_required_slots, - ) if not isinstance(result, HiCacheWriteFailure): return result - retry_plan = self._cp_build_write_admission_plan( + retry_admission = self._cp_build_write_admission( device_indices, node_id=node_id, phase="retry_after_local_failure" ) - retry_required_slots = max( - retry_plan["current_compatible_need"], default=int(result.required_host_slots) - ) - if retry_required_slots <= 0: + if not any(v > 0 for v in retry_admission.deficit_by_owner) and not any( + v > 0 for v in retry_admission.eviction_plan.remaining_deficit + ): raise RuntimeError( "CP HiCache host reservation failed although deterministic " - "capacity planner predicted no eviction need: " + "owner-lane admission predicted no host deficit: " f"node_id={node_id} required_host_slots={result.required_host_slots}" ) - if not self._evict_cp_host_plan( - retry_plan, node_id=node_id, phase="retry_after_local_failure" + if not self._evict_cp_host_for_write_admission( + retry_admission, node_id=node_id, phase="retry_after_local_failure" ): - return HiCacheWriteFailure(required_host_slots=retry_required_slots) + return HiCacheWriteFailure( + required_host_slots=max( + retry_admission.eviction_plan.remaining_deficit, + default=int(result.required_host_slots), + ) + ) logger.info( "[HiCache-write] write_backup CP retry after deterministic host eviction: " - "node_id=%d needed_slots=%d", + "node_id=%d deficit_by_owner=%s", node_id, - retry_required_slots, + retry_admission.deficit_by_owner, ) result = self.cache_controller.reserve_write_cp( device_indices=device_indices, node_id=node_id, ) - self._cp_debug_compare_write_admission( - retry_plan if self._cp_capacity_debug_enabled() else None, - result=result, - planned_required_slots=retry_required_slots, - ) if not isinstance(result, HiCacheWriteFailure): return result @@ -2258,6 +2244,16 @@ class HiRadixCache(RadixCache): return ack_queue_len = len(self.cache_controller.ack_write_queue) + # With per-layer CP backup, ongoing_write_through is populated before + # the final write ack is appended. Polling a TP all-reduce while the + # ack queue is empty is pure scheduler overhead: there is no candidate + # radix-state transition to make yet. The ack itself is appended at a + # deterministic layer-end / post-forward point across CP ranks; once it + # exists, the MIN below still gates host visibility on all ranks having + # completed the corresponding local transfer. + if ack_queue_len == 0: + return + finish_count = 0 for _, finish_event, ack_list in self.cache_controller.ack_write_queue: if not finish_event.query(): @@ -2728,6 +2724,10 @@ class HiRadixCache(RadixCache): result = self.inc_lock_ref(ancester_node) delta = result.delta + if len(nodes_to_load) == 0: + self.dec_lock_ref(ancester_node) + return None + # load it all or not at all. The scalar length remains only a # coarse upper bound; final CP admission is the owner-lane vector # check below. diff --git a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py index b3aa49ca8..cf9c625ca 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py @@ -272,6 +272,47 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): self.assertEqual(target.value.tolist(), loaded.tolist()) self.assertIn(target.id, cache.ongoing_load_back) + def test_load_back_evicts_blocking_leaf_to_unlock_parent_owner_lane(self): + allocator = _make_allocator() + # Only owner lane 1 is initially available. The target needs owner lane + # 0. The current leaf itself belongs to owner 1, but evicting it makes + # its parent evictable and the parent frees owner 0. + allocator.free_pages = torch.tensor([2], dtype=torch.int64) + cache = _make_cache(allocator) + + parent = _make_node( + 40, + 600, + [0], + value=torch.tensor([4, 5, 6, 7], dtype=torch.int64), + priority=0, + ) + _attach_child(cache, cache.root_node, parent) + child = _make_node( + 41, + 700, + [1], + value=torch.tensor([8, 9, 10, 11], dtype=torch.int64), + priority=0, + ) + _attach_child(cache, parent, child) + cache.evictable_leaves.add(child) + cache.evictable_size_ = len(parent.key) + len(child.key) + + target = _make_node(42, 800, [0], value=None, priority=10) + _attach_child(cache, cache.root_node, target) + + loaded = cache.load_back(target, mem_quota=100) + + self.assertIsNotNone(loaded) + self.assertEqual(cache.cache_controller.load_calls, 1) + self.assertEqual( + [indices.tolist() for indices in cache.cache_controller.evicted_device_indices], + [[8, 9, 10, 11], [4, 5, 6, 7]], + ) + self.assertEqual(target.value.tolist(), loaded.tolist()) + self.assertIn(target.id, cache.ongoing_load_back) + def test_load_back_failure_leaves_node_unassigned_and_unlocked(self): allocator = _make_allocator() allocator.free_pages = torch.tensor([1, 2], dtype=torch.int64) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index d3592c2fa..478eb0565 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -64,6 +64,7 @@ except (ImportError, RuntimeError): "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()", "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor", "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor", + "moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> Tensor[]", ): try: _sgl_kernel_lib.define(_schema) @@ -83,6 +84,7 @@ for _schema in ( "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()", "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor", "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor", + "moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> Tensor[]", ): try: _sgl_kernel_lib.define(_schema) @@ -629,6 +631,20 @@ class FakeTokenAllocator: def available_size(self): return 0 + def compute_owner_lane_stats(self, page_owners): + if len(page_owners) == 0: + return [], [], [] + cp_size = max(int(owner) for owner in page_owners) + 1 + required = [0] * cp_size + for owner in page_owners: + required[int(owner)] += 1 + available = list(required) + deficits = [0] * cp_size + return required, available, deficits + + def allocator_state_str(self): + return "FakeTokenAllocator" + class TestHiRadixCacheCPBackup(CustomTestCase): def test_session_aware_cache_forwards_cp_hicache_prepare(self): @@ -1243,7 +1259,7 @@ class TestHiRadixCacheCPBackup(CustomTestCase): with self.assertRaisesRegex( RuntimeError, - "planner predicted no eviction need", + "owner-lane admission predicted no host deficit", ): cache.write_backup(node) @@ -1868,6 +1884,8 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): cache._uses_cp_hicache = True cache.root_node = TreeNode() cache.device = "cpu" + cache.page_size = 1 + cache.token_to_kv_pool_allocator = FakeTokenAllocator() cache.load_back_threshold = 1 cache.evictable_size_ = 0 cache.metrics_collector = None @@ -1911,6 +1929,8 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): cache._uses_cp_hicache = True cache.root_node = TreeNode() cache.root_node.value = torch.empty((0,), dtype=torch.int64) + cache.page_size = 1 + cache.token_to_kv_pool_allocator = FakeTokenAllocator() cache.load_back_threshold = 5 cache.evictable_size_ = 0 cache.metrics_collector = None