Repository navigation
[None][feat] Mooncake safe shared-pool reuse and asynchronous loading - #10
Closed
SimengLiu-nv wants to merge 12 commits into
Closed
SimengLiu-nv wants to merge 12 commits into
SimengLiu-nv wants to merge 12 commits into
Conversation
Signed-off-by: Simeng Liu <simengl@nvidia.com>
BlockHashChain.extend encoded every token with int.to_bytes and joined the pieces per block, which dominated the lookup of a long prompt. It now encodes all newly completed blocks with a single np.asarray(dtype="<i8").tobytes() call and digests consecutive 8 * tokens_per_block byte windows of that buffer, returning early when no new full block exists. '<i8' is the same signed 8-byte little-endian encoding int.to_bytes produced, so every stored key is byte-identical and pages already in a pool stay addressable. A token outside int64 still raises OverflowError, now before any block is appended. Tests pin the encoding against a frozen copy of the per-token loop for an unscoped, a cache_salt and a multimodal scope (including -1 and both int64 extremes, fed incrementally), pin golden digests captured from the unmodified code, and check out-of-range tokens. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…save once it is queued The scheduler adapter published every block that had a hash and a page. The whole prompt is hashed at lookup, and pages can run ahead of computation: the post-allocation token-budget trim shrinks a context chunk whose pages were already allocated, and a served prefix grows capacity by the extra KV tokens a one-model draft needs. Such a block was saved with KV that had not been computed, and since the save mark moved past it and _put never overwrites an existing key, a later prefix hit loaded it. Saves are now bounded by the blocks this pass completes: computed_position, plus this iteration's honoured restore (which the runtime leaves out of computed_position), plus num_scheduled_tokens. The offer is trimmed to what the runtime honoured. update_state_after_alloc settles the outstanding offer to request.py_num_connector_matched_tokens, which KVCacheManagerV2 reaches through the base by-layer-group fold, and raises when that count is not a whole-block part of the offer. _loads_for now raises when an honoured block has no destination page instead of silently dropping the tail, because the runtime already counts those tokens as computed. A save counts as stored only once the worker queued it. _saves_for returns the block ordinal its saves end at without moving the save mark; build_connector_meta records it as unconfirmed, and the next build promotes it only for requests the worker's wait_for_save reported through the new take_submitted_saves. Saves bound for a pass that never reached wait_for_save are therefore emitted again rather than lost. The report is taken on every build, even with nothing unconfirmed, so a report about a request reset since cannot confirm that request's next pass if that pass never runs; request_reset also drops the request's unconfirmed entry. The scheduler test helpers gain FakeWorker.take_submitted_saves, submit_saves, settle_offer, saved_pages and request_data's computed_position / num_scheduled_tokens keywords, and make_request carries a settable py_num_connector_matched_tokens. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…art on a single rank A page the pool offered at lookup can be gone by the time it is loaded: evicted from a full pool, or left behind by a put that never completed. The runtime has already counted those positions as computed, so the worker used to fail the server on any missing page. The worker now declares reports_failed_loads. On a rank that holds whole pages and its own batch (a single-rank server, or an attention-DP owner once this connector supports attention DP) a request that lost a page is recorded for take_failed_load_requests instead, and the executor restarts only that request without the offered prefix. Under tensor parallelism the ranks hold shards of the same page and cannot restart a request alone, so there the load still raises with the existing "failed to load ... already reported as computed" error. Loads go through a new _load_pages(keys, addresses, sizes, owners, staging, handle, batch_size, *, recover), which later work reuses for asynchronous loads. It skips the remaining pages of a request that already lost one, treats a result list of the wrong length as all pages failed, scatters only the staged slots that landed and syncs the stream per batch, so no write into a restarted request's pages is in flight when it is reported. It returns the failed requests, the failed page count, the first missing key and the bytes that landed. A failed request's saves leave the bound metadata at once through drop_bound_saves, which edits the saves list wait_for_save reads in place and is idempotent with the manager's forget_request; otherwise the pass would publish pages the request never computed under its valid keys. The fake-store tests cover the single-rank report and its TP raise, the skip of a failed request's later pages, the save drop for a mixed batch and for a batch the failed request was all of, and staged loads that scatter only the landed pages. TP variants set _recover_failed_loads on the built worker rather than changing the world size, which would change every key namespace. Signed-off-by: Simeng Liu <simengl@nvidia.com>
The attention-DP framework creates a scheduler adapter and a worker on every
owner and asks a connector to opt in, but mooncake-store declared nothing,
so the executor creator refused it under enable_attention_dp, including for
a capacity-only server. Its keys also assumed tensor parallelism: each rank
keyed pages w{N}r{rank} and a lookup required every rank's shard, so under
attention DP no owner could ever find a prefix, because no other owner
writes the same page under a different rank.
Both connector halves now declare supports_attention_dp. The worker keys
pages by attention shard rather than by global rank: 0 of 1 under attention
DP, where every owner holds all heads, and the global rank and world size
under tensor parallelism. All owners therefore share one key space (w1r0),
a prefix one owner saved is a hit for every other owner and for a
single-rank server with the same model key and page layout, and a lookup
asks one key per layer group per block. Tensor-parallel keys are
byte-identical to before. The global rank still names the process in the
segment ledger, thread names and logs. Request state, transfers, completion
and the store handle were already owner-local.
With the failed-load restart in place, a page lost between lookup and load
restarts only the affected request on its owner instead of fail-stopping
every owner; the docs and the worker docstring now say so without a hedge.
Docstrings that described a single rank-0 leader now describe the scheduler
adapter, which runs on rank 0 under tensor parallelism and on every owner
under attention DP. The connector docs gain an attention-DP subsection for
Mooncake, drop attention DP from the unsupported list, replace the line
that said no Mooncake backend exists, and state that aggregated serving
under attention DP is not supported yet: when the attention-DP gate skips a
prepared pass, a generation page already reported to the connector can be
reassigned without the connector hearing of it. Disaggregated context
servers and roles that save nothing are unaffected.
Tests: the fake-store suite covers owners sharing pages across rank pairs,
tensor-parallel lookups still needing every shard, pinned key strings for
a single rank and for rank 2 of 4, the class declarations against
validate_attention_dp, the restart formula on an owner versus a TP rank,
and a capacity-only worker under attention DP. The real-pool suite saves a
page from one owner and loads it byte-identical on another while a TP rank
misses it. A new TP2 attention-DP end-to-end test, parametrized over the
overlap scheduler and listed in l0_dgx_h100, checks that both owners join
the pool, that rank 0 serves prompt tokens from it, and that every replayed
copy agrees with the cold output.
Signed-off-by: Simeng Liu <simengl@nvidia.com>
…l drafts (local_reuse_backoff) A one-model draft (MTP or EAGLE3) makes KVCacheManagerV2 back every local reuse match off by its prompt lookahead, so a local match of whole blocks ends one or more tokens short of a block boundary. mooncake-store offers only from a block boundary, so after any local hit it would offer nothing. The worker now fills in the two whole-block reuse declarations the cache manager reads once it is built after the connector. A role that reads the pool (both, consumer) always declares requires_block_aligned_reuse, and declares drops_reuse_match_backoff according to the new kv_connector_config.mooncake_store.local_reuse_backoff: - "off", the default, drops the backoff and takes local matches whole. The draft KV at the tail of a reused block may belong to another continuation, which can lower draft acceptance but never the output. - "align" keeps the backoff and has the manager align local matches down to a block, so the pool serves the trimmed block with the rest of the prefix at the cost of that block's transfer. Producer and capacity roles never offer, so they declare neither and keep the manager's own reuse. Every rank resolves the same role and arguments, so the declarations agree across the ranks of a server. The mode is an engine setting read from llm_args through the new config.local_reuse_backoff helper on every rank. It is not rendered into the client config or restated from an inherited one, so a deployment that inherits MOONCAKE_CONFIG_PATH without a mooncake_store block gets "off". The field is telemetry=False. The setting takes effect once the connector may run next to one-model speculative decoding. The single-block synchronous offer that the source branch made under "align" with asynchronous loads is not ported; this connector loads synchronously. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…text servers The speculative-prefill-only mode lets one-model MTP or EAGLE3 run next to a KV connector on a disaggregated context server, but only for a connector whose two halves declare supports_speculative_prefill. mooncake-store now declares it on both and makes the promises behind the declaration true. A one-model draft writes the KV at a position from the prompt tokens after it, so a stored block's draft-layer KV depends on tokens past the block. BlockHashChain takes a keyword-only prompt_lookahead: with a lookahead of L, block b is hashed over tokens[b*T : (b+1)*T + L] at a one-block stride over the same int64 payload, so a block exists only once its lookahead does, and a non-zero L re-digests the scope seed with b"prompt-lookahead" and L. With L = 0 the seed and every key are byte-identical to before, so pools already populated stay addressable, while a speculative server never shares pages with a server without a draft or with a different lookahead. The scheduler adapter reads the lookahead with draft_prompt_lookahead (one for MTP-Eagle and EAGLE3, max_draft_len for vanilla MTP), imported inside its constructor so importing the connector package does not pull in every drafter. It raises when max_draft_len is still unresolved, since vanilla MTP's lookahead is its draft length and lookahead-0 keys would silently name draft KV that reads ahead. Offers are already bounded by the chain, so a block whose lookahead runs past the prompt is neither offered nor saved. Saves were already capped at computed blocks, which this mode needs because one-model speculation allocates KV slots past every context chunk. The worker moves whole pages of every layer group, draft layers included, and its per-layer hooks are no-ops. Before the store is opened it runs the new validate_speculative_config, an allow-list of one-model MTP and EAGLE3 whose draft layers share the target KV cache: any other drafter, or a separate draft KV cache, would leave missing or stale draft KV behind a restored prefix and is refused. A capacity role addresses no page and skips the check, so a disaggregated generation server may run one-model MTP or EAGLE3, with or without attention DP. Tests pin key bytes built by hand from the framed scope and the lookahead digest for lookahead 0, 1 and 3 with salt and multimodal scopes, and cover lookahead keys, the unresolved draft length, chunk-ahead capacity pages, the drafter allow-list and the creator's acceptance of the real classes. An end-to-end test runs Qwen3-8B with EAGLE3 as a cold and a warm disaggregated context server against a live pool, with a generation server pulling each request, and requires reuse, matching first and draft tokens, and the refusal of a request that is not context-only; it is listed in l0_h100. Signed-off-by: Simeng Liu <simengl@nvidia.com>
A full Mooncake pool evicts on the master, and two of its behaviours cost
hits that the connector can avoid with per-put options.
Mooncake 0.3.13's first eviction pass makes up a shortfall from the
shard-ordered rest of the pool when lookups re-lease part of its oldest
frontier mid-eviction. That evicts blocks from the middle of prefixes still
in use, and a prefix is replayed only up to its first missing block, so
every stored block behind such a hole is recomputed. Soft-pinned objects are
evicted only by the second, age-ordered pass. The new
mooncake_store.soft_pin_ttl_ms soft-pins every page a server writes for that
many milliseconds.
A hit on a block needs every page of it: all layer groups and, under tensor
parallelism, every rank's shard. The master evicts pages independently, so
it can keep a shard that holds space but can never hit again. With the new
mooncake_store.group_blocks, every page of a block is put into one eviction
group, named by the new KeyNamespace.group as
<namespace>/<model_key>/w<world>/<block hash>, which leaves out the shard,
the layer group, the tokens per block and the page size but keeps the
attention world size.
Both fields are off by default, telemetry=False, and apply only to a role
that writes. They are rendered into the client config, read back by
MooncakeStoreConnectorConfig.from_file, and restated from an inherited
MOONCAKE_CONFIG_PATH; a config without a TTL restates the block to off. The
JSON keys match the source branch's, so an existing client config keeps
working, while its TRTLLM_MOONCAKE_STORE_SOFT_PIN_TTL_MS and
TRTLLM_MOONCAKE_STORE_GROUP_BLOCKS environment overrides are not read.
group_blocks must be a JSON boolean, since bool("false") is True.
The worker builds a Mooncake ReplicateConfig through _replicate_config
(the source branch's _put_config, renamed so it no longer shadows the
worker's _put_config attribute). A soft pin alone shares one config across
every put; grouped puts build one per call with a group per key that
survived the existence filter. The constructor builds a probe config before
the segment is mounted, so bindings older than 0.3.13 fail at startup with
an ImportError naming SoftPinAction and the ReplicateConfig fields. With
neither option the put keeps its three-argument form. A put whose every page
is refused under a soft pin logs once that a TTL above the master's
--max_kv_soft_pin_ttl (24 h by default) makes the master refuse every put.
Signed-off-by: Simeng Liu <simengl@nvidia.com>
…tats_interval) Tuning a shared pool needs two numbers the connector did not report: how often a prefix lookup stops at a block the pool evicted from the middle of a chain it otherwise still holds, and where the time of a prefix load goes. The source branch read them behind three environment variables and a fixed interval; here one setting, mooncake_store.stats_interval (seconds, gt=0, finite, telemetry=False), turns on both. It is rendered into the client config, read back by MooncakeStoreConnectorConfig.from_file and restated from an inherited MOONCAKE_CONFIG_PATH like the eviction settings; a config without it restates the block to off. TRTLLM_MOONCAKE_HOLE_STATS, TLLM_M3_LOAD_PHASE_TIMING and TLLM_MOONCAKE_LOAD_PHASE_TIMING are not read. Only a role that reads the pool counts anything, since only it looks prefixes up and loads them. count_prefix_hit then folds its existing batch_is_exist answer into _PrefixHoleStats: lookups, looked-up and hit blocks, holes, the missing blocks inside them, and the stored tokens behind a hole that are recomputed anyway. The hit itself is unchanged and no store call is added. A "prefix-hole stats" line is logged at info level at most once per interval. The synchronous load times its phases with _timed_phase, which sums wall and thread CPU seconds per phase and reads no clock when stats are off: resolve, the allocation fence, and, inside _load_pages through its new phases argument, describe, get, unstage and sync. Each start_load_kv pass that moved pages is folded into _LoadStats with its pages, store calls, landed bytes and transfer time, and get_finished logs a "load stats" line at most once per interval. _LoadStats records and drains under a lock, so loads that later run on background threads can record into the same counters. Both lines are flushed when the worker shuts down. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…host staging and admission-query support A prefix the pool serves is loaded inside the executor iteration that schedules the request, so every other request in the batch waits on the transfer. With the new mooncake_store.async_load, a role that reads the pool offers every hit asynchronously instead: the runtime parks the request outside the batch and async_load_workers background threads per rank load its pages while other requests keep computing. The scheduler adapter computes a parked request's transfer when its allocation is reported. update_state_after_alloc now delegates to the new update_state_after_alloc_by_layer_group, which settles the honoured count first, as before, and for an asynchronous offer builds the transfer from the request's full page lists, refusing a missing or invalid destination page. An offer honoured with no tokens still queues an empty transfer, because the runtime parks the request regardless. The transfer travels once, in MooncakeStoreMetadata.async_loads of the same iteration's metadata; the blocks it covers are never written back, and request_reset drops one not yet handed out. The adapter also declares supports_admission_query: a query only refills the request's offer state and transfers are built only once the allocation is reported, so a repeated admission-time answer replaces the earlier one. Every rank queues its own pages when it binds the metadata, so under tensor parallelism each rank loads its own shard. Binding never raises, since it runs outside the attention-DP failure vote; what cannot be queued (async_load off, no load thread, no allocation fence) is kept for the next start_load_kv or get_finished to raise. A load thread waits for the KV cache manager's allocation fence on its own thread, fetches through the shared _load_pages, and get_finished reports a request only after the runtime announced it and its pages landed or one was found missing. A missing page hands the request back through take_failed_async_load_requests on a single rank or an attention-DP owner, and fails the rank under tensor parallelism. Any exception on a load thread is kept and raised by every later start_load_kv and get_finished until shutdown. start_load_kv now consumes the bound synchronous loads, so a second call on the same metadata starts nothing. The pinned host memory per direction under stage_through_host becomes the configurable staging_buffer_bytes (default 1GiB, the previous fixed ceiling), which replaces staging.MAX_STAGING_BUFFER_BYTES and is what the node memory check counts. With async_load each load thread gets its own staging pool and CUDA stream from an even share of the load direction's budget, and no synchronous load pool is pinned, so the check still counts one budget per direction; a pool whose share is below one page is warned about. Shutdown signals the save and load threads together, joins them against one deadline, skips loads still queued, and unregisters every staging pool before closing the store, keeping any pool whose unregistration failed. The three settings are telemetry-free MooncakeStoreConfig fields, rendered into the client config, read back by MooncakeStoreConnectorConfig.from_file (async_load must be a JSON boolean) and restated from an inherited MOONCAKE_CONFIG_PATH. Asynchronous loads record into the stats_interval load line, which gains failed_pages and, with async_load, wait_ms and the load threads' queue. The integration list runs a new end-to-end test comparing an asynchronously loaded prefix with a synchronous run. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…ling policy Signed-off-by: Simeng Liu <simengl@nvidia.com>
Signed-off-by: Simeng Liu <simengl@nvidia.com>
This was referenced Nov 4, 2026
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
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.
Superseded by NVIDIA#21656 targeting
NVIDIA/TensorRT-LLM:main, per the requested upstream destination. The source branch is retained.Mooncake must publish computed KV only, preserve cache identity across supported layouts, and finish or discard transfers before pages can be reused.
This PR implements computed-only publication, failed-load recovery, attention-DP shared keys, speculative/local reuse, completion-reservation activation, eviction controls, diagnostics, asynchronous loads, and bounded host staging. It rejects aggregated attention-DP producer/both roles at startup; consumer/capacity roles and supported disaggregated deployments remain available. Async loading stays off by default. Admission-query activation belongs to the scheduling PR.
Tests cover the startup gate and use supported ADP deployments. Upgrade documentation requires draining old writers and resetting affected pool entries because preserved key bytes cannot identify previously published invalid KV.
Extracted from NVIDIA/TensorRT-LLM#21195 at
4918c76ee1958bf0b6d01840588f742a7d0e743b, rooted at fork maincbef70dbc9879798b4106c694928d59976cf0ac1. Original feature commits are retained where possible, with split-specific activation and review fixes in separate commits.Part 2 of 5: C8–C16 plus C6's Mooncake activation, excluding C16's admission-query activation. Targets
codex/pr21195-runtimefor a clean incremental diff.Validation: Scoped pre-commit checks passed. Changed Python files parse, JSON files parse, and
git diff --checkpasses. These are static checks, not native/GPU runtime evidence. Full TensorRT-LLM imports and schema regeneration are blocked in this environment by missingtensorrt_llm.bindings; generated snapshots were mechanically replayed, and the combined config models/schemas/telemetry manifest match the original head. The aggregate PR's earlier results do not validate these split heads. A source-isolated check passed all 24 ADP role/transceiver combinations using the actual enum and guard. Native, real-pool, DMA lifetime, multi-rank, speculative, and GPU smoke tests remain pending.Five-PR split: runtime → Mooncake → scheduling. Independent: HTTP streams and idle receive timeouts. The source PR remains the umbrella.