mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

[Integration] TensorRT-LLM Store-backed KV cache connector #6

Open
feiyangyoufeng wants to merge 11 commits from feiyangyoufeng/Mooncake:feature/tensorrt-llm-kv-connector into main
pull from: feiyangyoufeng/Mooncake:feature/tensorrt-llm-kv-connector
merge into: mooncake-track:main
Conversation 0 Commits 11 Files Changed 38 +5167 -2
feiyangyoufeng commented 2026-07-08 18:43:03 +08:00
First-time contributor
Copy Link

Summary

This PR adds a Store-backed KV cache connector for TensorRT-LLM that gives it
cross-request / cross-instance / cross-CTX–GEN global KVCache reuse through the
Mooncake Distributed Store.

TensorRT-LLM already ships the Mooncake Transfer Engine as a C++
cache_transmission backend for point-to-point KV movement between a prefill
(CTX) worker and a decode (GEN) worker. That KV is ephemeral — when a request
finishes, its KV is freed, so the next request with the same system prompt or
conversation prefix re-runs prefill from scratch. This PR adds the missing layer:
a distributed KV pool with cache-hit scheduling, implemented through TensorRT-LLM's
official KvCacheConnectorScheduler / KvCacheConnectorWorker API.

It is config-only on the TensorRT-LLM side (no fork) and ships in the
mooncake wheel as mooncake.tensorrt_llm:

from tensorrt_llm import LLM
from mooncake.tensorrt_llm import make_kv_connector_config

llm = LLM(model=..., backend="pytorch",
          kv_connector_config=make_kv_connector_config())

What's included

  • mooncake-wheel/mooncake/tensorrt_llm/ — the connector package (scheduler,
    worker, Store client, key schema, paged-KV layout codec, heat-aware replication,
    metrics, a config layer, and a TensorRT-LLM-free simulator).
  • mooncake-wheel/tests/tensorrt_llm/ — unittest test cases (hermetic by
    default; real-Store tests gated on MC_BACKEND=mooncake).
  • benchmarks/tensorrt_llm_kv_connector/ — self-check, demo, benchmark scripts,
    CTX/GEN/Store config templates, and a one-click deploy script.
  • docs/source/getting_started/examples/tensorrt-llm-integration.md + README and
    toctree entries.

Design

  • Keyspace isolation. A fingerprint over model, engine, tokenizer, KV
    dtype/layout and TP/PP/EP prefixes every object key, so incompatible
    configurations occupy disjoint keyspaces and can never be reused across.
  • Chained prefix hashing, keyed with the per-request cache_salt, so a key
    match implies the entire prefix matches and tenants with different salts never
    share cache.
  • Self-describing blocks carry a layout fingerprint validated on read; a
    mismatch is rejected rather than served. KV bytes move verbatim, so float16
    / bfloat16 / fp8 round-trip bit-exactly.
  • Block-completion manifest with cross-rank quorum. Each rank writes a READY
    marker for the blocks it durably saved; the aggregate manifest the scheduler
    probes is published only once every rank's marker is present, so a partially
    written or single-rank block can never be advertised as a hit.
  • Fail-closed loads. Because the scheduler reports matched tokens up front
    (the runtime then skips prefill), a promised block that is missing or
    layout-rejected at load time raises rather than leaving a stale device slot.
  • Layer-granularity pipelining (opt-in). With MC_LAYER_GRANULARITY=1 the
    connector stores KV per layer and uses the real wait_for_layer_load /
    save_kv_layer hooks to overlap Store I/O with attention compute, instead of
    loading the whole block before the forward pass.
  • Heat-aware replication picks ReplicateConfig.replica_num from a decayed
    popularity estimate, giving hot shared prefixes more copies for read fan-out
    and fault tolerance.

The package depends only on numpy at import time; the Mooncake store client and
the TensorRT-LLM base classes are optional, with a process-shared file backend and
a faithful compat shim so the connector stays importable and testable without a
GPU, a TensorRT-LLM install, or a running Store.

Validation

  • ruff check, ruff format, and codespell pass on all new files.
  • Hermetic test suite: 45 passed, subtests, 3 skipped (real-Store tests).
  • Against a real Mooncake Store (master + HTTP metadata, tcp): the
    self-check passes (cold→warm full hit, bit-exact reuse, salt isolation), the
    layer-granularity path round-trips bit-exactly with its manifest published in
    the real store, and an agentic trace of 144 requests warms from 0% to ~88%
    block hit ratio     with every request verified bit-exact.
  • Benchmark matrix (agentic workload): 0.87 block hit ratio, ~6.5× modeled
    TTFT     at a conservative 3000 tok/s prefill, 0 errors, numerical
    consistency OK across backends. A freshly restarted instance reuses the warmed
    Store immediately (1.0 hit ratio, 0 errors), demonstrating cross-instance
    resilience.

Notes

  • The block hit ratio is a property of the connector and workload, so it is
    identical across the memory / file / mooncake backends; only the measured Store
    bandwidth differs.
  • Prefill compute time in the benchmark is modelled with a single
    --prefill-tok-per-s knob; hit ratio, bytes moved, bandwidth and numerical
    consistency are measured against the real Store.
## Summary This PR adds a **Store-backed KV cache connector for TensorRT-LLM** that gives it cross-request / cross-instance / cross-CTX–GEN **global KVCache reuse** through the Mooncake Distributed Store. TensorRT-LLM already ships the Mooncake **Transfer Engine** as a C++ `cache_transmission` backend for point-to-point KV movement between a prefill (CTX) worker and a decode (GEN) worker. That KV is ephemeral — when a request finishes, its KV is freed, so the next request with the same system prompt or conversation prefix re-runs prefill from scratch. This PR adds the missing layer: a distributed KV pool with cache-hit scheduling, implemented through TensorRT-LLM's official `KvCacheConnectorScheduler` / `KvCacheConnectorWorker` API. It is **config-only on the TensorRT-LLM side** (no fork) and ships in the `mooncake` wheel as `mooncake.tensorrt_llm`: ```python from tensorrt_llm import LLM from mooncake.tensorrt_llm import make_kv_connector_config llm = LLM(model=..., backend="pytorch", kv_connector_config=make_kv_connector_config()) ``` ## What's included - `mooncake-wheel/mooncake/tensorrt_llm/` — the connector package (scheduler, worker, Store client, key schema, paged-KV layout codec, heat-aware replication, metrics, a config layer, and a TensorRT-LLM-free simulator). - `mooncake-wheel/tests/tensorrt_llm/` — `unittest` test cases (hermetic by default; real-Store tests gated on `MC_BACKEND=mooncake`). - `benchmarks/tensorrt_llm_kv_connector/` — self-check, demo, benchmark scripts, CTX/GEN/Store config templates, and a one-click deploy script. - `docs/source/getting_started/examples/tensorrt-llm-integration.md` + README and toctree entries. ## Design - **Keyspace isolation.** A fingerprint over model, engine, tokenizer, KV dtype/layout and TP/PP/EP prefixes every object key, so incompatible configurations occupy disjoint keyspaces and can never be reused across. - **Chained prefix hashing**, keyed with the per-request `cache_salt`, so a key match implies the entire prefix matches and tenants with different salts never share cache. - **Self-describing blocks** carry a layout fingerprint validated on read; a mismatch is rejected rather than served. KV bytes move verbatim, so `float16` / `bfloat16` / `fp8` round-trip bit-exactly. - **Block-completion manifest with cross-rank quorum.** Each rank writes a READY marker for the blocks it durably saved; the aggregate manifest the scheduler probes is published only once every rank's marker is present, so a partially written or single-rank block can never be advertised as a hit. - **Fail-closed loads.** Because the scheduler reports matched tokens up front (the runtime then skips prefill), a promised block that is missing or layout-rejected at load time raises rather than leaving a stale device slot. - **Layer-granularity pipelining (opt-in).** With `MC_LAYER_GRANULARITY=1` the connector stores KV per layer and uses the real `wait_for_layer_load` / `save_kv_layer` hooks to overlap Store I/O with attention compute, instead of loading the whole block before the forward pass. - **Heat-aware replication** picks `ReplicateConfig.replica_num` from a decayed popularity estimate, giving hot shared prefixes more copies for read fan-out and fault tolerance. The package depends only on `numpy` at import time; the Mooncake store client and the TensorRT-LLM base classes are optional, with a process-shared file backend and a faithful compat shim so the connector stays importable and testable without a GPU, a TensorRT-LLM install, or a running Store. ## Validation - `ruff check`, `ruff format`, and `codespell` pass on all new files. - Hermetic test suite: **45 passed, subtests, 3 skipped** (real-Store tests). - Against a **real Mooncake Store** (master + HTTP metadata, `tcp`): the self-check passes (cold→warm full hit, bit-exact reuse, salt isolation), the layer-granularity path round-trips bit-exactly with its manifest published in the real store, and an agentic trace of 144 requests warms from **0% to ~88% block hit ratio** with every request verified bit-exact. - Benchmark matrix (agentic workload): **0.87 block hit ratio**, ~**6.5× modeled TTFT** at a conservative 3000 tok/s prefill, **0 errors**, numerical consistency OK across backends. A freshly restarted instance reuses the warmed Store immediately (**1.0 hit ratio, 0 errors**), demonstrating cross-instance resilience. ## Notes - The block hit ratio is a property of the connector and workload, so it is identical across the memory / file / mooncake backends; only the measured Store bandwidth differs. - Prefill compute time in the benchmark is modelled with a single `--prefill-tok-per-s` knob; hit ratio, bytes moved, bandwidth and numerical consistency are measured against the real Store.
feiyangyoufeng added 11 commits 2026-07-08 18:43:05 +08:00
9519a831c5 [Integration] Add TensorRT-LLM Store-backed KV cache connector 
Add `mooncake.tensorrt_llm`, a KV cache connector that gives TensorRT-LLM
global KVCache reuse on top of the Mooncake Distributed Store. It implements
TensorRT-LLM's `KvCacheConnectorScheduler` / `KvCacheConnectorWorker` interfaces
and is selected purely through `KvCacheConnectorConfig`, so no fork of the host
runtime is required.

The scheduler probes the Store for the longest matching block prefix in
`get_num_new_matched_tokens` and emits per-block load/save metadata in
`build_connector_meta`; the worker performs the batched Store gets/puts and
bit-exact device copies in `start_load_kv` / `wait_for_save`.

Design points:

* Keyspace isolation: a fingerprint over model, engine, tokenizer, KV
  dtype/layout and TP/PP/EP prefixes every object key, so incompatible
  configurations occupy disjoint keyspaces and cannot be reused across.
* Chained prefix hashing keyed with the per-request cache salt, so a key match
  implies the whole prefix matches and different tenants never share cache.
* Self-describing blocks with a layout fingerprint; a mismatch on read is
  treated as a miss rather than a corrupt reuse. KV bytes move verbatim, so
  float16 / bfloat16 / fp8 round-trip bit-exactly.
* Heat-aware replication picks `ReplicateConfig.replica_num` from a decayed
  popularity estimate, giving hot shared prefixes more copies.

The package depends only on numpy at import time; the Mooncake store client and
TensorRT-LLM base classes are optional, with a process-shared file backend and a
faithful compat shim so the connector stays importable and testable without a
GPU, a TensorRT-LLM install, or a running Store.
489061dddc [Integration] Add tests for the TensorRT-LLM KV cache connector 
Cover the connector end to end with `unittest` test cases that run under both
pytest and `python -m unittest`:

* keying: chained prefix hashing, salt isolation, schema fingerprint sensitivity.
* layout: bit-exact block codec across dtypes, layout-mismatch refusal, the
  cross-TP head remap.
* scheduler: cold/warm/partial prefix matching, salt isolation, the block
  boundary guard, and the load_kv_async==False invariant.
* worker: bit-exact encode/decode (incl. bfloat16), negative shape/dtype
  refusal, a full put->get block round-trip, and safe miss on an absent block.
* replicate, metrics, and an end-to-end simulator check of hit ratio and
  numerical consistency.

The hermetic tests use the in-memory backend and need no GPU or Store.
`test_real_mooncake.py` exercises a real master and is skipped unless
`MC_BACKEND=mooncake` is set.
b8d8a235e0 [Integration] Add benchmarks and deployment scripts for the TRT-LLM connector 
Add a reproducible benchmark and deployment harness under
`benchmarks/tensorrt_llm_kv_connector`:

* `selfcheck.py` reports the environment and runs a CTX->Store->GEN round-trip
  that asserts bit-exact reuse.
* `agentic_trace.py` records the cache warm-up curve over a multi-turn trace;
  `run_matrix.py` sweeps backends (memory / file / mooncake) plus a
  cross-instance restart scenario; `plot.py` renders the figures.
* `deploy/run_ctx_gen_store.sh` brings up a local Mooncake master + metadata
  server, prints the client environment, and launches CTX/GEN workers.
* `configs/` holds CTX / GEN / Store templates wired to the connector preset.

Block hit ratio and per-request numerical consistency are measured against a
real Store; prefill compute time is modelled with a single throughput knob, so
the reported TTFT speedup follows transparently from the measured hit ratio.
92f0323aec [Doc] Document the TensorRT-LLM KV cache connector 
Add an integration guide describing the architecture, the one-line
`KvCacheConnectorConfig` setup, the Store/keyspace environment variables and the
correctness model, wire it into the docs toctree, and add a TensorRT-LLM
Integration section plus an update entry to the top-level README.
aa75b26a9a [Integration] Add layer-granularity pipelining and block-completion manifest 
Add an opt-in per-layer load/save path (`MC_LAYER_GRANULARITY=1`) that uses the
connector's `wait_for_layer_load` / `save_kv_layer` hooks to overlap Store I/O
with attention compute: on load, a layer's attention can begin as soon as that
layer's KV is resident while later layers are still in flight; on save, each
layer is persisted the moment its forward work is enqueued. The default block
granularity (all layers in one object) is unchanged.

Introduce a per-block completion manifest: rank 0 publishes a marker only after
a block's shards (and, in layer mode, all layers) are durable, and the scheduler
keys its existence probe on the manifest instead of a single rank-0 data shard.
A reported hit therefore means the whole block is present across ranks, so a
missing nonzero-rank shard can no longer produce a false hit.

Drop the half-wired asynchronous-save flag: `request_finished` now always reports
synchronous saving, matching `get_finished` (which returns nothing pending),
removing a path that could otherwise defer block reclamation indefinitely.

Add tests for the layer-granularity round-trip, manifest timing (published only
after the final layer; never by a nonzero rank), and a keyspace-isolation matrix
asserting every fingerprint field independently yields disjoint keys.
ab2de27942 [Integration] Add one-command showcase and label modeled TTFT explicitly 
Add a `showcase` subcommand to the deploy script that brings up a Store, runs the
correctness gate, the benchmark matrix and the warm-up trace against it, renders
the figures, and tears the Store down on exit -- a single reproducible command.

Rename the benchmark TTFT outputs to mark them as modeled (`ttft_speedup_modeled`,
table column "TTFT x modeled"): hit ratio, bytes moved, bandwidth and numerical
consistency are measured against a real Store, while prefill compute time is
modeled from one throughput knob, and the field names now make that explicit.

Document the load/save granularity options and the `MC_LAYER_GRANULARITY` switch
in the integration guide and benchmark README.
55ce4b8a0b [Integration] Harden the block-completion manifest into an all-ranks commit 
Make the manifest a real cross-rank commit so a reported hit means the block is
durable on every rank, not just rank 0:

* Each rank writes a per-rank READY marker for the blocks it actually saved.
* Rank 0 publishes the aggregate MANIFEST only for blocks whose READY markers are
  present on all `world_size` ranks; a fast rank 0 can no longer expose a hit
  before another rank's shard exists.

Commit manifests only for blocks whose shard write succeeded (track the saved
ops rather than committing from the full metadata), and surface a non-zero
`put_batch` status as an error.

Guard the layer-granularity path: it requires a known layer count, so
`MC_LAYER_GRANULARITY` with an unknown count now falls back to block granularity
with a warning instead of publishing a manifest for a never-written block. Also
reset the per-pass layer-load scratch at the start of every load so a request
with no loads cannot drain a previous request's buffers.

Add tests for the multi-rank quorum (a missing rank shard blocks the manifest),
the layer-count fallback, and the per-request scratch reset.
3d11549b4a [Integration] Add a modeled block-vs-layer KV-load timeline 
Add a small, explicitly-modeled illustration of the per-layer pipeline benefit:
block granularity waits for the whole-block load before the forward pass, while
layer granularity overlaps each layer's load with the previous layers' compute.
The script reports the modeled time-to-first-token for both schedules and a
figure, clearly labelled as modeled (hit ratio, bytes and bandwidth remain the
measured quantities elsewhere).
b785207328 [Integration] Gate manifest commit on durable writes and make quorum live 
Three correctness fixes to the commit protocol:

* Publish a block's manifest only when its data writes actually succeeded. The
  block path skips the manifest on a non-zero `put_batch` status; the layer path
  tracks per-block failure across all layers (encode or put) and excludes any
  block that failed on any layer, so a final-layer success can no longer commit a
  block whose earlier layer write failed.
* Remove the cross-rank liveness gap: any rank that observes READY markers from
  all ranks publishes the aggregate manifest (idempotent), instead of only rank 0
  probing once -- so whichever rank finishes last guarantees publication without
  a coordinator re-running.
* Synchronize the stream before encoding KV in the layer save path, so a layer
  persisted right after its work is enqueued captures finalised bytes rather than
  partial device state; the overlap with later layers' compute is preserved.

Also surface non-zero store status on marker writes as an error. Add tests for
the durable-write gating (a failed put publishes no manifest) and the multi-rank
quorum semantics.
19fbea39ef [Integration] Fail closed when a promised KV block is missing at load 
The scheduler reports matched tokens before the forward pass, so TensorRT-LLM
skips prefill for those positions. If the promised KV is then missing (evicted or
lease-expired) or fails layout validation at load time, continuing would leave a
stale device slot and silently corrupt the output. Both the block and layer load
paths now raise `KvConnectorError` in that case instead of treating it as a miss,
matching the runtime's "a counted hit must be delivered" contract.

Document the admission-to-load eviction window (mitigated by the master's KV
lease TTL) and the token+salt identity model in the integration guide's
correctness section.
Auto Label PRs / triage (pull_request) Failing after 2m20s 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
99e7304373 test(store): normalize allocator test headings
Some checks failed
Auto Label PRs / triage (pull_request) Failing after 2m20s
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
This pull request can be merged automatically.
You are not authorized to merge this pull request.
You can also view command line instructions.

Step 1:

From your project repository, check out a new branch and test the changes.
git checkout -b feiyangyoufeng-feature/tensorrt-llm-kv-connector main
git pull feature/tensorrt-llm-kv-connector

Step 2:

Merge the changes and update on Gitea.
git checkout main
git merge --no-ff feiyangyoufeng-feature/tensorrt-llm-kv-connector
git push origin main
Sign in to join this conversation.
Reviewers
No reviewers
Labels
Clear labels
No items
No Label
Milestone
Clear milestone
No items
No Milestone
Projects
Clear projects
No project
Assignees
Clear assignees
No Assignees
1 Participants
Notifications
Due Date
The due date is invalid or out of range. Please use the format 'yyyy-mm-dd'.

No due date set.

Dependencies

No dependencies set.

Reference: mooncake-track/Mooncake#6
No description provided.
Delete Branch "feiyangyoufeng/Mooncake:feature/tensorrt-llm-kv-connector"

Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?