[None][perf] Integrate Mooncake store for MinimaxM3 - #18676
Merged
brb-nv merged 24 commits intoSep 5, 2026
Conversation
1 task
…iner The CMake install in install_mooncake.sh only exposes the C++ transfer engine, which is what the cache transceiver links against. The mooncake-store KV cache connector needs MooncakeDistributedStore instead, and that class only exists in the Python bindings, which the source build does not produce. Install the wheel pinned to the same upstream version already built from source, so the store client and the transfer engine cannot drift apart, and widen the attribution entry to cover the wheel's payload alongside the source install. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Regular block reuse never leaves the instance that computed a prefix, so a context server recomputes prefixes its neighbours already have. This connector publishes KV pages into a Mooncake store -- a shared CPU pool addressed by content -- so any engine can replay them, and composes with the existing point-to-point cache transceiver rather than replacing it. Built on KVCacheManagerV2's register_kv_cache_layout, since a V2 page is a set of strided byte ranges per layer group rather than one pool tensor, which is also the shape Mooncake's multi-buffer batch APIs take. V2 exposes no block hashes to a connector, so identity is a blake2b chain over (parent, salt, block tokens) computed leader-side; the key namespace additionally pins the model, shard, layer group and page geometry so any mismatch reads as a miss instead of as garbage. Loads run synchronously in start_load_kv and fail loudly: the runtime has already counted those tokens as computed, so a partial load is a wrong answer. Saves are handed to a background thread behind a CUDA event, and the leader reports such requests as saving asynchronously so their pages stay pinned until the writes land. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Adds the startup gates, the 'mooncake-store' registry preset so the connector can be selected by name, and unit tests covering the pieces that decide whether a cache hit is correct. Every gate rejects a configuration whose failure mode is a wrong answer rather than a slow one. Context parallelism gives a rank a slice of the sequence instead of whole blocks, so one key would name different bytes per rank. Sliding-window attention makes a page's validity depend on where the window sits, which is a property of the reader rather than of the tokens. MiniMax-M3's index-V cache lives outside the paged pools, so a replayed prefix would pair stored index-K with stale index-V -- the same restriction disaggregated serving already applies. Pipeline parallelism is refused as untested rather than unsound. Beam search, attention DP, non-GPU cache tiers and Mamba caches are already rejected for all connectors in py_executor. Tests run without Mooncake or a GPU: the store is an in-process fake and the layout is synthesized from integers, which is all the addressing arithmetic needs. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Describe what the store buys over local block reuse, how pages are keyed and which configurations are refused, so an operator can tell the store apart from the similarly named transfer engine used for prefill/decode handoff. Ship a trtllm-serve config as a starting point. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Follows the telemetry allowlist gaining the new connector preset. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…odes Pool capacity comes only from processes that open a store handle, and only the context workers configure the KV connector, so every byte of the pool was prefill-node memory. That made the store a prefill-DRAM-caches-prefill-GPU tier largely duplicating TensorRT-LLM's native host offload, rather than the cross-node pool it is meant to be. Add a capacity-only donor process per generation node: it contributes host memory and then idles, never issuing a put or get. Prefill-written KV can then live on decode-side DRAM while the generation engine stays connector-free, so it keeps its cache transceiver for the prefill-to-decode handoff. A donor is a separate process rather than a new StoreRole because the roles describe traffic -- producer writes, consumer reads, both does both -- and none of them means "contribute memory only". Also bring the master and client config up inside the benchmark job so a run needs no manual setup, resolve MOONCAKE_CONFIG_PATH from the log directory whose path is not known when the worker environment is built, and report block placement grouped by segment host. That last figure is what distinguishes a spanning pool from a prefill-only one: a single host means the donors are absent or not being allocated into. Verified on MiniMax-M3 at TP=2 for both prefill and decode: the pool grew from 32GiB on one host to 64GiB across two, and of 4.58GiB written by prefill, 1.45GiB (32%) landed on the decode node, with no load or save failures and no change in throughput or latency. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Registering the GPU KV pools directly with Mooncake requires GPUDirect RDMA, which is unavailable on GB300 nodes without nvidia_peermem: ibv_reg_mr fails with EFAULT and the TCP transport segfaults in its memcpy worker pool. Add an opt-in stage_through_host mode that gathers pages into pinned host slots and registers those instead, so only host memory is ever exposed to the transport. Bind the save thread to the rank's device, captured on the executor thread at layout registration. Torch's current device is thread-local and a new thread starts at 0, so the thread was creating its stream on device 0 while the KV pages lived on the rank's device -- every staged copy failed with cudaErrorInvalidValue on every rank except 0. State copy directions explicitly rather than inferring them from pointers, and report the operands and current device when a copy fails. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… is used The store is addressed by whole blocks. The connector receives the device match as num_computed_tokens and offers only blocks beyond it, but it can resume only from a block boundary, so a match ending mid-block makes it decline the lookup and the store is never consulted. enable_partial_reuse is exactly what puts the match off a boundary, so it trades part of one block of device reuse for every stored block of the remaining prefix. The default is true, which made the pathological combination the one a user gets by saying nothing. On MiniMax-M3 it declined 97.2% of lookups and held actual prompt cache read at 35% against a 96% ceiling, so a 1.6 TB pool measured as though it were absent; forcing it off reached 93.5% and 2.18x the output token throughput. Coerce rather than reject, since a wrong answer is not at stake and refusing to start over a default no one chose is worse than fixing it and saying so. This sits in py_executor_creator beside the FORCE_DETERMINISTIC coercion because the KV cache manager reads the flag when it builds its block pools, which happens well before the connector is constructed -- so the connector cannot police this from its own startup gates. uses_connector compares the resolved module rather than the connector name, so a config that spells out connector_module instead of using the preset is still recognized. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
mooncake_usage.md is the entry point someone inheriting this work should read first: install, the configuration that works, how to tell from the logs whether the pool is being used, what it measured, and what bites. It stays short by pointing at the reference doc, the SLURM runbook and the working configs rather than restating them. It leads on partial reuse because that single flag decided whether the feature did anything at all, and on which reuse metric counts store hits, because the Prometheus counters exclude them and reading the wrong one makes a working store look inert. Also corrects the runbook, which predicted an unaligned local match would be rare with 128-token blocks. It was the common case: 97.2% of lookups. The prediction is worth replacing rather than deleting, since the arithmetic returns if tokens_per_block changes. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…spill to The V2 scheduler's only way to reclaim GPU pages was to suspend a request, which unpins its pages so the eviction controller can migrate them one cache level down. With GPU as the last level a suspended page stays HELD, which is not evictable there, so suspension frees nothing. A KV connector run is exactly that configuration: the manager drops the automatic host tier because tier migration would reassign GPU slots and invalidate the device addresses the connector registered. A prefill server then fills its pool, admits nothing further, and spins at full speed scheduling nothing -- looking healthy to the hang detector and to /health while burning the rest of its wall clock. Give the scheduler a second reclaim action for that case. Preemption closes the victim's KVCache instead of parking it, which returns its committed blocks to the radix tree as reusable prefix and leaves the pages DROPPABLE, evictable at every level. The data is not discarded: it stays resident and locally matchable until something else actually needs the space, and with a connector attached the blocks already written to the store come back through the ordinary prefix load. Recompute is the always-correct fallback. Pages are not released while the connector still has saves reading out of them, which would otherwise let a later request overwrite the bytes mid-transfer and publish them under a valid hash. The victim goes through the same request_finished/get_finished handshake a finished request uses, and the executor resets it to context state once every rank reports the saves retired. Configurations with a tier below GPU keep the existing suspend-based path untouched. Also replace the deadlock detector's single-iteration check with a consecutive-stall counter that counts context candidates as well. The old check only looked at generation requests, which a disaggregated prefill server does not have, and raising on one stalled iteration would misfire on the transient deferrals the scheduler makes for multimodal chunk alignment and PEFT budget. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…ingup The mooncake-store connector needs a reachable mooncake_master and a MOONCAKE_CONFIG_PATH naming it, both of which only the SLURM benchmark harness knew how to produce. A plain trtllm-serve therefore could not use the connector without borrowing that harness. Describe the pool in kv_connector_config.mooncake_store instead and the server provisions it itself: resolve the master, render the client config, export MOONCAKE_CONFIG_PATH before the ranks that open store handles are spawned, and tear down what it started on exit. launch_master: true starts a master for a single engine; master_server_address joins one with its own lifetime, which is what sharing a pool or surviving a restart requires. An inherited MOONCAKE_CONFIG_PATH still wins and logs that it did, so the harness path is unchanged. An external master is probed during startup rather than left to fail inside store.setup on every rank after the model has loaded. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
MooncakeStoreConfig could turn stage_through_host on but not size the buffer it stages through, leaving it at the connector's 512MiB default. A buffer that cannot hold transfer_batch_size pages reduces the batch rather than failing, so the ceiling that setting implies was reachable only by writing MOONCAKE_CONFIG_PATH by hand -- which is the thing describing the pool in the config is meant to replace. The field is left out of the rendered config when unset so the connector's default stays the one definition of it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
A server that owns its pool provisions it from its own config. The pools it cannot own still needed a script of someone else's: one shared by several engines, or outliving a restart, needs a master that is not any of them, and a pool whose capacity should include nodes that run no connector needs those nodes to hold segments. Both were scripts under mooncake_disagg, so those deployments were assembled from a benchmark harness rather than from what TensorRT-LLM ships. 'trtllm-serve mooncake_master' runs a master for as long as it runs and publishes where it landed; 'trtllm-serve mooncake_donor' lends a node's memory while leaving that engine connector-free. Donation stays out of StoreRole on purpose -- the roles describe an engine's traffic, and contributing memory is capacity, so making it a role would start a generation server reading or writing the store to get its DRAM in. master_server_address also accepts file://<path>, which is what makes the master reachable without anyone writing its address down: its host is whatever the scheduler chose, so a config settled beforehand cannot name it, and the wait to read it is also the wait for it to exist. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The benchmark harness implemented the master's lifetime -- glog settings so it logs at all, a readiness probe, an address file -- and located mooncake_segment_donor.py through environment.trtllm_repo. All of that now has a shipped equivalent, so the script calls 'trtllm-serve mooncake_master' and 'trtllm-serve mooncake_donor' instead and the donor script is deleted rather than left as a second implementation of the same setup() call. The install step becomes a fallback: docker/common/install_mooncake.sh bakes the wheel into images built from this repo, so probe for the bindings and the binary first and only reach for trtllm_repo when an older image has neither. Two consequences worth knowing when reading a log directory: the address file now holds host:port rather than a bare host, and the master's glog moved to <log_dir>/mooncake_master.log, leaving 2_mooncake_master.log to the launching command -- so the block placement and eviction sections read the former. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
A pool assembled by a launch script is a pool only that script can assemble. The parts a server could not own -- a master with its own lifetime, memory lent by a node whose engine uses no connector -- were commands the harness ran, which left the shape of a deployment split between the configs and a 1000-line SLURM script. Both become config. A launched master publishes its address, so a server lending memory can find the pool a context server owns without anyone writing an address down, and mooncake_donation contributes host memory from a server that never reads or writes the store. The harness now installs the bindings and substitutes its log directory; nothing else about the pool is its business. The run directory turns out to be load-bearing for more than logs. Provisioning reaches the ranks the LLM constructor spawns by exporting MOONCAKE_CONFIG_PATH, but under trtllm-llmapi-launch each rank is its own task and was already running, so those ranks now read the rendered config back from the run directory. Without that, every rank but the leader of a multi-GPU server failed during bringup. Bringup narrates itself throughout, because a pool that came up wrong is otherwise visible only as a low hit rate hours later: what was resolved from where, the segment in bytes as well as GiB, the capacity arithmetic, and the tail of the master's own log when it dies during startup. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Tighten the connector prose and error messages, and drop unit tests that only asserted a default value, a non-None return, or message wording. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
mooncake_disagg/ and mooncake_usage.md were runbooks and configs for local experiments; docs/source/features/kv-cache-connector.md covers the connector for users. The two places that pointed at the scratch install script now point at docker/common/install_mooncake.sh, and the SLURM harness checks the image for the bindings instead of installing them per job. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
watch_job.sh polled a log directory for local experiment monitoring; nothing in the harness invokes it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
TRTLLM_STALL_REPORT_SEC came out of debugging the no-evict deadlock and is unrelated to the mooncake-store connector. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…d cases Seven groups covered the same code path with trivially different inputs: connector recognition, provisioning no-ops, pool master validation, omitted client-config fields, JSON size parsing, the rank's device, and the two donor entry paths. Each is now one parametrized test. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
force-pushed
the
user/brb/m3-mooncake-store
branch
from
September 4, 2026 17:12
12522bc to
652948c
Compare
…n the M3 stage The mooncake-store connector tests and the KVCacheV2Scheduler mock tests are CPU-only, so they fit the single-GPU M3 pre-merge stage. The scheduler tests were not listed anywhere, and this branch adds cases to them, so list the file alongside the three new mooncake-store test modules. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
added
the
api-compatible
brb-nv
marked this pull request as ready for review
September 4, 2026 17:29
brb-nv
requested review from
pcicotti,
peihu-nv and
zheyuf
and removed request for
a team,
FrankD412,
Mgluhovskoi,
QiJune,
StanleySun639,
YihuiLu512,
arysef,
brnguyen2,
chienchunhung,
chuangz0,
crazydemo,
liji-nv,
lori-ren,
lowsfer,
yiqingy0 and
yuanjingx87
September 4, 2026 17:30
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #71584 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #71584 [ run ] completed with state
|
pcicotti
approved these changes
Sep 4, 2026
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
force-pushed
the
user/brb/m3-mooncake-store
branch
from
September 5, 2026 00:09
8005175 to
6b11d7f
Compare
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #71633 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #71633 [ run ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Built on top of #17974.
This MR integrates Mooncake into TRTLLM for KVCMv2. When Mooncake is being used, native offloading is turned off.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.