[Store] HiCache: support hybrid and sparse model KV layouts #1

Lxxx · 2026-06-26T14:19:46+08:00

Lxxx commented

2026-06-26 14:19:46 +08:00

First-time contributor

Add mooncake.hicache: hybrid/sparse model KV support for the HiCache L3 path. See branch commits.

Lxxx added 7 commits 2026-06-26 14:19:47 +08:00

86ad911677 [Store] hicache: pool descriptors, key schema and metrics for hybrid models

The HiCache L3 path over the Mooncake Store assumes the dense-transformer model:
one cache pool backed by one physical KV buffer in a page-contiguous layout.
Hybrid/sparse models (MLA, sliding-window, Mamba/linear, NSA/DSA sparse) expose
their KV as several heterogeneous pools, one of which can be a logical anchor
that owns allocation indices but no buffer.

Add the foundation of a `mooncake.hicache` subpackage that lets a storage
backend reason about such pools:

- `KVPoolDescriptor` / `PoolDescriptorRegistry` normalize each host pool to a
  typed record and expose a single register/skip decision; a logical anchor
  (`kv_buffer is None`) is recorded for routing but never registered, and a
  physical pool on `layer_first` is rejected before it reaches the store.
- `KeyBuilder` and the component-suffix helpers namespace object keys by
  `(layout, attention_type, pool_name)` so blocks from different pools cannot
  alias; the suffix expansion mirrors the existing hybrid page-component layout.
- `layout_compat` resolves `layer_first` to `page_first` / `page_first_direct`
  and models the logical anchor's per-page multi-buffer aggregation.
- `observability` provides a dependency-free metrics sink (JSON + Prometheus
  text) for the read path.

The modules duck-type against the host pool classes and have no hard dependency
on torch or SGLang, so they import and unit-test without a runtime.

Co-Authored-By: Claude <noreply@anthropic.com>

4c8e490e7b [Store] hicache: attention-semantic selective KV loading

Add the selective load/save layer that moves KV by attention semantics instead
of moving every page of every pool. One policy per pool type, built on the
existing `PoolHitPolicy` (ALL_PAGES / TRAILING_PAGES):

- full attention   -> restore the longest contiguous existing prefix
- compressed (MLA) -> restore the latent prefix
- sliding window   -> restore only the trailing window
- recurrent state  -> restore only the last state
- sparse index     -> restore the index, then only the index-selected blocks

`HybridKVController` orchestrates registration, key building and per-pool prefix
resolution (the final usable prefix is the min across pools, as in
batch_exists_v2). `SelectiveTransferEngine` executes the plans against a store
backend, and `HybridHiCacheFacade` exposes the small classify / hit-policy /
record-io surface a backend delegates to.

Skipped pages are never read from the store, so the result is numerically
identical while the bytes moved drop substantially; the skip counts are recorded
for observability.

Co-Authored-By: Claude <noreply@anthropic.com>

ee9a936d9f [Store] hicache: store backends and hybrid-model fixtures

Add the storage backends the transfer engine moves bytes through:

- `KVStoreBackend` is the object-store interface (string keys -> opaque bytes,
  plus register_buffer for the zero-copy path).
- `LocalKVStore` is an in-process implementation for tests and benchmarking; it
  increments the same `mooncake_store_*` counters as production and models
  transfer time from a `BandwidthModel`.
- `MooncakeStoreBackend` wraps a real `MooncakeDistributedStore`, adapting the
  client conventions (empty-bytes miss -> None, int exist flags -> bool,
  insert-if-absent put, RDMA register_buffer with status check). The store is
  passed in already set up, so the wrapper has no hard import of the compiled
  extension.

`mock.py` builds structurally faithful hybrid-model fixtures (logical anchor +
MLA latent + sliding window + Mamba state + sparse indexer, plus dense / linear /
SWA presets) that implement the page-source protocol for round-trip testing.

Co-Authored-By: Claude <noreply@anthropic.com>

18a1b92eb9 [Store] hicache: unit and real-store tests

Cover the hybrid HiCache helpers:

- descriptor / registry: logical anchors are skipped, physical pools register,
  a physical pool on layer_first is rejected
- key schema: no cross-pool aliasing; component suffixes match the hybrid
  page-component layout
- selective policies: the five load/save strategies and the boundary helpers,
  including that a policy's configured trailing length is honored under a default
  context
- layout resolution, observability (counters / histogram / Prometheus text) and
  the backend facade
- end-to-end round-trips: byte-exact restore, selective skips are genuinely not
  restored, eviction truncates the prefix

`test_hicache_hybrid_store_integration.py` runs the engine through a
`MooncakeStoreBackend` against a real `mooncake_master` over TCP; it is skipped
when the store extension or the master binary is unavailable.

Co-Authored-By: Claude <noreply@anthropic.com>

6722b1ce1c [Store] hicache: hybrid KV benchmark

Add a runnable benchmark that reports, per hybrid-model preset, the KV bytes
moved by selective loading versus a dense-assuming backend, the resulting
store-bandwidth headroom, and a TTFT model that composes the measured byte
volumes with the HiCache transfer and prefetch-timeout models. Byte and hit/skip
inputs are measured by running the transfer engine against a store; `--store
mooncake` runs the same measurements through a real MooncakeDistributedStore.

Co-Authored-By: Claude <noreply@anthropic.com>

609ca57795 [Doc] hicache: hybrid/sparse design notes and Grafana dashboard

Document how the L3 path supports hybrid/sparse models as a companion to the
existing HiCache design: pool classification and the register/skip decision, the
collision-free key schema, selective loading by attention semantics, layout
compatibility and the read-path metrics. Register the page in the design
toctree and add an importable Grafana dashboard for the hybrid counters.

Co-Authored-By: Claude <noreply@anthropic.com>

Auto Label PRs / triage (pull_request) Failing after 2m14s Details

Build & Test (Linux) / build (3.10) (pull_request) Has been cancelled Details

Build & Test (Linux) / build (3.12) (pull_request) Has been cancelled Details

Build & Test (Linux) / build-musa (pull_request) Has been cancelled Details

Build & Test (Linux) / test-wheel-ubuntu (3.10, ubuntu-22.04) (pull_request) Has been cancelled Details

Build & Test (Linux) / test-wheel-ubuntu (3.10, ubuntu-24.04) (pull_request) Has been cancelled Details

Build & Test (Linux) / test-wheel-ubuntu (3.12, ubuntu-22.04) (pull_request) Has been cancelled Details

Build & Test (Linux) / test-wheel-ubuntu (3.12, ubuntu-24.04) (pull_request) Has been cancelled Details

Build & Test (Linux) / build-flags (3.10) (pull_request) Has been cancelled Details

Build & Test (Linux) / build-flags (3.12) (pull_request) Has been cancelled Details

Build & Test (Linux) / Build Docker Image (pull_request) Has been cancelled Details

Build & Test (Linux) / Spell Check with Typos (pull_request) Has been cancelled Details

Build & Test (Linux) / Check code format (pull_request) Has been cancelled Details

Build & Test (Linux) / check-paths (pull_request) Has been cancelled Details

Build & Test (Linux) / build-wheel-cu13 (pull_request) Has been cancelled Details

Build & Test (Linux) / ascend-test (pull_request) Has been cancelled Details

Build & Test (Linux) / integration-test (pull_request) Has been cancelled Details

Build & Test (Linux) / CI Gate (pull_request) Has been cancelled Details

d8c84facdd [CI/Build] Run the hicache hybrid tests in CI

Add a step to the wheel test job that runs the hicache unit tests and, since a
master and metadata server are already up, the real-store integration test.

Co-Authored-By: Claude <noreply@anthropic.com>

Lxxx added 2 commits 2026-06-26 15:15:24 +08:00

fefc43af7c [Store] hicache: results/evidence in the design notes and a summary figure

Lead the design notes with the measured results (per-model byte reduction and
bandwidth headroom, per-pool movement, an evidence table that separates measured
from modeled figures, and a reproduce block). Add a `--plot` option to the
benchmark that renders a dense-vs-semantic byte chart and the TTFT-vs-prefix
curve, and a Grafana stat panel for the selective skip ratio derived from the
existing counters.

Co-Authored-By: Claude <noreply@anthropic.com>

Auto Label PRs / triage (pull_request) Failing after 55s Details