Repository navigation
[None][feat] Close the Mooncake store feature gap: async loads, attention DP, speculative prefill and KV connector support in KVCacheManagerV2 - #21195
Draft
SimengLiu-nv wants to merge 20 commits into
Conversation
…ock multiple (reuse_match_alignment) A one-model EAGLE3 or MTP draft sets reuse_match_backoff, so the radix tree trims every local match by the draft lookahead and a match of K whole blocks ends a few tokens short of a block boundary. A KV connector that stores and offers whole blocks is queried from that endpoint and has nothing to offer for any request with a local hit. Add KVCacheManagerConfig::reuseMatchAlignment. After the backoff, BlockRadixTree::match trims the final match down to a multiple of the alignment, and pruneMatch re-aligns before validating pages and again whenever SSM snapshot truncation moves the endpoint off the alignment, so the probe, the claim and the first-new-block probe all trim once in the same tree walk. The attention-only diagnostic pass is aligned too; the content-divergence depth stays raw. The default of 1 keeps token-granular reuse, and validate() refuses an alignment that is not a positive divisor of tokens_per_block. The nanobind constructor gains a trailing reuse_match_alignment=1 kwarg and a read-write field, and the runtime field spec lists it so that dataclasses.replace() rebuilds of the config keep it. The struct layout changes, so the bindings and the kvCacheManagerV2 gtests must be rebuilt together with the library. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…by KVCacheManagerV2 A connector that stores and offers KV by whole blocks is queried from the local prefix match. Under a one-model draft that match is backed off by the draft's prompt lookahead and ends a few tokens short of a block boundary, so such a connector has nothing to offer for any request with a local hit. KvCacheConnectorWorker gains two opt-in declarations next to supports_attention_dp, both False by default and honoured only when exactly True: - requires_block_aligned_reuse: when a backoff is in effect and partial reuse is off, KVCacheManagerV2 passes its ledger tokens_per_block as the core reuse_match_alignment, so every local match is aligned down to a block boundary and the connector serves the trimmed block with the rest of the prefix. A paired draft pool aligns too, keeping the joint claim on a boundary. - drops_reuse_match_backoff: the target manager drops the one-model draft backoff instead and uses whole local matches. Draft KV at the tail of a reused block may then come from another continuation, which affects draft acceptance only. A draft manager never drops its own backoff, and the raw draft_prompt_lookahead stays for the sharing lender. The manager reads both through one helper that ignores a capacity-only connector and tolerates partially built managers and connector doubles without a worker. Nothing declares them yet, and speculative decoding is still refused with a connector, so production behaviour is unchanged. The connector guide documents the declarations in the worker interface, and the shutdown and capacity_only worker bullets that had drifted under the attention data parallelism section move back there. Signed-off-by: Simeng Liu <simengl@nvidia.com>
KvCacheConnectorManager and its scheduler-output manager skipped placeholder requests only when attention DP was on, by testing `enable_attention_dp and req.is_dummy_request`. A dummy request carries no prompt a connector could load or save, whatever the parallel mode, so the check now goes through one module helper, `_is_dummy_request`, used by the scheduler output, prefix reservation, query, commit, request_finished and update_state_after_alloc paths. The helper reads both `is_dummy_request` (bound by the C++ LlmRequest, so readable on the base class C++ callers hand over) and the Python-only `is_attention_dp_dummy`, and compares each against True. The explicit comparison keeps a MagicMock request, which answers every attribute truthily, from being mistaken for a dummy and silently skipped. Under tensor parallelism this is defensive rather than a behaviour change: C++ addSequence already skips dummies, KVCacheManagerV2 does not ask the connector for them, and padding dummies exist only under attention DP. Every TP rank sees the same flags, so the early returns before the leader broadcast stay collective-consistent. The output manager's `enable_attention_dp` keyword is now dead and is removed together with the one call site that passed it. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…d failed (reports_failed_loads) A page a connector 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 today the only safe answer is for the connector to raise, which stops the executor for one request's cache miss. This adds an opt-in protocol that hands such a request back for a restart instead. A worker declares `reports_failed_loads = True` and implements `take_failed_load_requests()` (polled after every `start_load_kv`), optionally `take_failed_async_load_requests()` (polled after every `get_finished`), and `drop_bound_saves(request_ids)`. The manager masks the declaration for a capacity-only role, refuses a declaring worker that leaves its hooks at the base defaults, and refuses the combination with prefix reservations, whose completion protocol has no failure outcome. The executor refuses it on KVCacheManagerV1, which cannot restart a request. For each reported request the executor drops its bound saves, then calls the new `KVCacheManagerV2.drop_context_allocation`, which is `free_resources` followed by a context rewind, and releases any speculative-decoding slot. The cache is closed without committing, so the never-filled range does not reach the radix tree, and the scheduler readmits the request and asks the connector afresh. `_kv_connector_prepare_batch` now parks, pads an attention-DP batch the connector emptied (`_pad_attention_dp_batch_emptied_by_connector`), then binds the plan and starts loads (`_kv_connector_start_loads`), which is where synchronous failures are taken. Asynchronous failures are taken right after the terminations in `_kv_connector_terminate_requests`, so a parked load whose cancellation was deferred is not restarted. The attention-DP recovery loop drops restarted parked requests from the requests it waits for, so a restarted request (CONTEXT_INIT with no cache) is never put back into the batch, and a loop whose only parked request restarted stops waiting and retries padding. A rebuilt plan goes through `_kv_connector_start_loads` too. A restart is owner-local, so it applies on a single rank or an attention-DP owner; a tensor-parallel report is fatal. Connector exceptions, deadline misses and misreported IDs remain fail-stop through `_kv_connector_fail`, with no new collective or vote field. A misreported ID is a synchronous ID outside the batch, an asynchronous ID whose load has not been reported finished, or one an earlier `get_finished` already returned as finished: a finished load stays in `finished_async_loading_requests` after the scheduler admits the request, so only the `get_finished` call the poll follows can still name a request that is parked. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…l MTP/EAGLE3 on disaggregated context servers KVCacheManagerV2 refuses speculative decoding next to a KV connector because a rejected draft token shrinks a request's page list underneath the connector, which only hears about appended pages. A context-only request on a disaggregated context server never rejects a draft token, so its page list only grows, and one-model MTP and EAGLE3 keep their draft KV in the target pool the connector has registered. This lets such a server run speculative decoding next to a connector that declares it can. Both connector halves gain a `supports_speculative_prefill` class attribute, modelled on `supports_attention_dp`. The scheduler's declaration promises that its keys cover the draft's prompt lookahead past each block and that it saves only computed blocks; the worker's promises whole-page transfers of every layer group, draft layers included. `KvCacheConnectorManager` takes a keyword-only `speculative_prefill_only` and refuses it, after the attention-DP and failed-load checks, unless both present halves declare support (`declares_speculative_prefill`; a TP follower has no scheduler). The executor creator computes the mode once the connector classes are resolved: one-model MTP (MTP-Eagle or vanilla) or EAGLE3 with a positive draft length, a cache transceiver backend, and both declarations. It is never set for an explicit `use_kv_cache_manager_v2=False`, so V1 keeps constructing exactly as before. The speculative-decoding refusal moves after class resolution, keeps its wording, and is lifted only in this mode. The executor accepts the mode on KVCacheManagerV2 only (an "auto" that resolves to V1 is refused), lifts its own speculative refusal for it, and refuses a separate draft KV cache, which the connector never sees. A request that is not context-only is refused at admission by the module-level `_validate_connector_admission`, the first check in `_validate_request`, so only that request fails, for example a /health_generate probe sent straight to the context server. Being a module-level helper read through getattr, it needs no stub in the minimal executors unit tests build. As a backstop, the manager refuses such a request again in `query_num_new_matched_tokens` and `reserve_prefix` before the leader broadcast, so every TP rank raises together, and in `build_scheduler_output` before any output state is recorded. Dummy requests and a capacity-only connector are never refused, which also lets a capacity-role generation server whose classes declare the mode run one-model speculation. No connector declares the mode yet, so behaviour is unchanged until one opts in. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…egated servers without an eviction tier With no eviction tier, a full-attention disaggregated server admitted requests against their prompt-sized allocations only. Incoming transfers and new prompts could then take pages that admitted requests still needed: a decode request ran out of room for its output, and a connector prefill, whose prefix is restored only after the final batch is chosen, found the pages for that restore taken by other requests' chunks. KVCacheManagerV2 can now keep a logical completion bound, in ledger blocks, per admitted request and admit a new one only while every bound still fits each GPU pool's free and evictable slots. No page is allocated early. Disaggregated generation init reserves through the output limit plus the maximum generation lead, and frees a refused request's cache instead of suspending it. A connector first chunk reserves through the end of its prompt when the connector has no source reservations; the V2 scheduler consults the reservation after every budget, chunking and multimodal-alignment refusal and defers the request instead of preempting an admitted one. The reservation is opt-in. A connector asks for it with a new KvCacheConnectorWorker declaration, requires_completion_reservation, which the manager reads from the attached connector's worker like requires_block_aligned_reuse, so every rank sees it. Only an explicit True counts, and a capacity-only connector's is ignored. The mooncake_store worker declares it: registering its pages drops the automatic host tier, and its prefix restores are chosen only after the final batch. A server opts in without such a connector through a new prototype KvCacheConfig field, enable_disagg_completion_reservation, False by default. It lets a disaggregated generation server with no connector and host_cache_size: 0 reserve its decode headroom; the connector prefill reservation still needs a connector that may serve the prompt. A target manager that cannot apply the field logs a warning rather than ignore it silently: without a connector the automatic host tier stays unless host_cache_size is 0. The field is Python-only, so the C++ KvCacheConfig conversion leaves it out, and usage telemetry captures it like the other KvCacheConfig booleans. A server with neither opt-in keeps its existing admission; that includes host_cache_size: 0 and page-registering connectors such as LMCache and KVBM, whose automatic host tier the manager drops. Beyond either opt-in, the reservation applies only to a disaggregated target manager with no tier below GPU, beam width 1, a single pipeline stage and full attention on every layer. Pipeline parallelism is excluded because each rank decides alone and a follower refusing a generation-init request that rank 0 admitted would not be reconciled. A bound is released by free_resources, by revert_allocate_context and a failed first-chunk resize while the chunk is unconsumed, and on shutdown. The receive-ordering tests bind the real prepare_disagg_gen_init onto a spec'd Mock, which would otherwise answer the new class-level state with truthy child mocks, so their fixture sets both reservation maps explicitly. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…booking and uncharged async parks The KVCacheManagerV2 scheduler sizes and charges a first context chunk before the connector is asked for its prefix, which only happens in prepare_resources once the batch is final. A chunk the connector then loads asynchronously is parked before the forward pass, so its charge is lost budget, and an offer that moves the chunk window has to grow its pages after every later admission of the pass took theirs; when that grow fails, only the whole blocks the chunk's pages hold are loaded and the rest of the offer is recomputed. kv_connector_config.admission_query (KvCacheConnectorAdmissionConfig, prototype) lets the scheduler ask a connector that declares KvCacheConnectorScheduler.supports_admission_query once a first chunk is sized and allocated. The answer is stored on the request as an epoch-checked ConnectorAdmissionOffer and used by prepare_resources, so each allocation is still queried and reported once. KVCacheManagerV2.cover_admission_offer grows the allocation to the shifted window right away. When the pool cannot cover it, the chunk is skipped and the existing first-chunk release returns its pages, up to max_offer_deferrals passes or until the request has just preempted another one for pages; deferrals count as scheduler progress. With max_uncharged_parks > 0, up to that many chunks per pass that will park are admitted with a 0-token charge; they still take a request slot. The manager reports the capability from the scheduler class and turns it off under attention DP. The executor ignores the config with a warning for a capacity-only role, prefix-reservation connectors and connectors without the capability, and refuses it with the overlap scheduler or without the KVCacheV2Scheduler. A request admitted without a charge that does not park fails the executor through _kv_connector_fail, because other parked loads may already own destination pages. The cover is a no-op under SWA scratch reuse, whose pre-grown replay would otherwise read the grow as a skipped forward. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…PARK_STATS) Tuning admission-time connector queries on a disaggregated context server needs to see, per iteration, how much of the context budget went to requests the connector parked, how many first chunks were passed over for budget, how connector offers were placed, and where the executor thread spent its time. TLLM_CTX_PARK_STATS=1 appends a flat ", ctx_park = key=value ..." suffix, with no braces so existing parsers keep matching, to every print_iter_log line of the non-overlap executor loop. ParkStats records the wall and executor-thread CPU time of the sched/prep/meta/start/fwd/upd/term phases, the parked requests and tokens, the requests admitted without a token charge, the context tokens left for the forward pass, and the scheduler's budget-limited first chunks (BudgetTracker.num_ctx_budget_skips, including FORCE_CHUNK first chunks that wait for a fresh budget). KVCacheManagerV2 counts each connector query and offer, how the offer was placed, the admission-time grows, and the host time in the query, the placement and the page report, enabled by the executor and drained on every rank each iteration. TLLM_GIL_PROBE_US adds the wake-up overshoot of a sleeping probe thread as a GIL contention proxy; the executor stops it on shutdown. Both switches are debug-only environment variables, like the TLLM_PROFILE_* ones next to them in profiling.py, and add no configuration field. 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>
…n server may hold; fail a stream before its first chunk with an HTTP error A generation leg takes its prompt KV cache from the context server's one-shot send session named by disagg_request_id. When that transfer fails, the generation server raises in its stream generator after the 200 headers, so the proxy saw a ClientPayloadError with no lines yielded and re-posted the leg under a freshly minted disagg_request_id. No context server holds KV cache for that id, so the replay waited forever and the client request never completed. The transient-TCP budget made this wider: a server disconnect or connection reset can arrive after the generation server already read the request. The client now uses the generation-leg contract of open PR NVIDIA#20938, with its text kept verbatim so a later merge of that PR is a no-op. A generation leg keeps its id on retry, and it is re-sent only when it cannot have reached the server: the connection was refused (ClientConnectorError) or reset while the request was being written (ConnectionResetError). Any failure after the response headers, and any other error, fails the request. Context requests are unchanged. The error-status check in _post_with_retry now runs before the stream branch as well, so a streaming leg answered with a 4xx/5xx JSON body raises a ClientResponseError carrying the worker's status instead of tripping the "Response is not streaming" assertion. On the server, _wrap_entry_point pulls a relayed stream's first chunk before it sends the headers. A chat or completion stream that fails before its first chunk now reaches _handle_exception and becomes an HTTP error (and an Anthropic error envelope on /v1/messages) instead of a 200 whose body is cut off. The returned relay is already started, so closing it before it is iterated still closes the upstream stream and releases the router reservation. Responses streams keep their 200 plus an `error` event, because guard_responses_stream stays inside the first-chunk pull. total_responses now counts a stream only once its first chunk arrived. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…generation server Without in-flight cancellation, the generation-side KV-transfer timeout only flags the request in check_transfer_timeouts. The cancel and the error response happen in the response pass, which runs only after a forward step. A generation server with nothing else scheduled therefore kept a request whose KV never arrived until the client gave up, even after logging "Observed timeout on generation request". DisaggTransferCoordinator.poll_gen_transfers, which runs at every loop head, now cancels and fails the receives that check_transfer_timeouts flagged whenever in-flight cancellation is inactive. This covers any transceiver without that capability, including KvCacheTransceiverV2 and TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 on an unsupporting transceiver. A cancel that finishes at once fails the request through fail_timed_out, inline or buffered for handle_timeouts_synced under multi-rank ADP. A cancel the transceiver refuses is asked only once: the Python transceiver has then marked the session cancelled and the next reap_gen_receives fails the request through the transfer-error path, while the C++ transceiver cannot cancel a receive already in flight and would only repeat its "Cannot cancel request" warning at the idle loop rate. Requests the user cancelled are left to the cancel path. No collective is added. The in-flight cancellation path is unchanged. fail_timed_out now moves a request still in DISAGG_GENERATION_TRANS_IN_PROGRESS to DISAGG_TRANS_ERROR before failing it. cancel_request does not change the state, so PyExecutor._handle_errors parked the failure for a transfer that no longer existed, and _sweep_transfer_held_failures then dropped it because the request had already left active_requests: no error response was sent and the KV was never freed. This also fixes the same loss on the existing response-pass timeout path. Signed-off-by: Simeng Liu <simengl@nvidia.com>
…stimated uncached prefill The default attention-DP router balances ranks by prompt length. When a KV cache connector restores prefixes from a pool shared by all attention-DP ranks, most of a later turn's prompt is restored rather than computed, so equal-length prompts can bring very different prefill work and one rank can end up working through a large prefill while its partners pad in lockstep. WorkAwareADPRouter estimates a new request's prefill as its prompt growth since the previous request with the same conversation_id (the whole prompt for a first turn, a request without an id, or a shrinking prompt; at least one token). A rank's load is the summed estimate of its requests still in the context phase, and new requests go, largest first, to the least-loaded rank below max_num_active_requests. Every request, including explicitly placed ones, is estimated once in routing order and charged to its rank; n>1 child ids inherit the parent's estimate, and the estimate is published as py_uncached_prefill_tokens through the existing hook while the V2 scheduler's long-prefill demotion consumes it. The conversation and per-request tables are LRU-bounded and change only inside route_requests, which every rank runs on identical inputs; create_rank_state only reads them, so placements stay consistent across ranks. The NVIDIA#19369 connector fields in the gathered rank state are untouched. The router is selected by the new prototype field attention_dp_config.enable_work_aware_routing, replacing the fork's TRTLLM_ADP_WORK_AWARE_ROUTING environment variable so the choice is validated and identical on every rank. ADPRouter.create gains a kv_connector_manager argument, which PyExecutor passes through, and builds the router only with a connector that moves KV; otherwise it falls back to the default router with a warning, staying silent for the KV-cache estimation executor, which is built without its connector. The field cannot be combined with enable_kv_cache_aware_routing or kv_cache_routing_conversation_affinity. Signed-off-by: Simeng Liu <simengl@nvidia.com>
This branch has not been deployed
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.
Summary
This PR brings the Mooncake KV-cache-store work from the fork branch
M3-mooncake-maintomain. That branch carries MiniMax-M3-specific changes, so it is not rebased; each feature is re-implemented againstmaininstead:connectors/mooncake_storepackage. Everything it already has stays: pool manifest, capacity role, ledger, GPUDirect, LoRA and multimodal scoping, node budget.The PR builds on #19369 (attention DP for V2 KV cache connectors) and keeps its behaviour: the status vote, the transfer deadlines and the
_kv_connector_failfail-stop.The generic KV connector, KVCacheManagerV2 and scheduler support (group A) and the work-aware router (commit 20) are opt-in. With all new knobs at their defaults and no connector declaring the new class attributes, they leave
main's behaviour unchanged. The in-treemooncake_storeconnector declares those attributes, so existingmooncake_storedeployments change behaviour by default. The disaggregated-serving fixes in commits 18 and 19 have no knob and apply to every disaggregated deployment. Both are listed under Behaviour changes.The PR has 20 commits in three groups. Group A comes first, B depends on A, and C is independent and can be split out.
A. Generic KV connector, KVCacheManagerV2 and scheduler support (commits 1-8)
reuse_match_alignmentin the V2 radix-tree match, with its nanobind binding and runtime config field.requires_block_aligned_reuseanddrops_reuse_match_backoff. KVCacheManagerV2 reads them when a one-model draft would back local reuse matches off a block boundary.reports_failed_loads).supports_speculative_prefill.requires_completion_reservation(mooncake_store does) orkv_cache_config.enable_disagg_completion_reservationis set (for generation servers without a connector and withhost_cache_size: 0).main's admission. A set knob that cannot apply logs a warning.kv_connector_config.admission_query.TLLM_CTX_PARK_STATSandTLLM_GIL_PROBE_US. This commit is debug-only and can be dropped.B.
mooncake_store(commits 9-17)local_reuse_backoff(off/align) for whole-block local reuse under one-model drafts.C. Disaggregated serving (commits 18-20, independent)
User-visible knobs
kv_cache_config.enable_disagg_completion_reservationkv_connector_config.admission_query.{max_uncharged_parks, max_offer_deferrals}disable_overlap_scheduler=Trueonly; ignored with a warning under attention DP, for a capacity-only role, with prefix reservations, or when the connector scheduler does not declaresupports_admission_querykv_connector_config.mooncake_store.local_reuse_backoffoffofforalign; only matters with one-model draftskv_connector_config.mooncake_store.soft_pin_ttl_ms--max_kv_soft_pin_ttlkv_connector_config.mooncake_store.group_blockskv_connector_config.mooncake_store.stats_intervalkv_connector_config.mooncake_store.staging_buffer_bytes1GiBkv_connector_config.mooncake_store.{async_load, async_load_workers}attention_dp_config.enable_work_aware_routingsupports_speculative_prefill,reports_failed_loads(plustake_failed_load_requests,take_failed_async_load_requests,drop_bound_saves),requires_block_aligned_reuse,drops_reuse_match_backoff,requires_completion_reservation,supports_admission_queryKVCacheManagerConfig.reuse_match_alignment(C++ and binding)The fork's environment variables became config fields:
TRTLLM_MOONCAKE_STORE_SOFT_PIN_TTL_MS,TRTLLM_MOONCAKE_STORE_GROUP_BLOCKS,TRTLLM_MOONCAKE_HOLE_STATS,TLLM_M3_LOAD_PHASE_TIMING,TRTLLM_MOONCAKE_STORE_LOCAL_REUSE_BACKOFF,TLLM_CTX_UNCHARGED_ASYNC_PARKS,TLLM_CTX_PARK_OFFER_DEFERRALSandTRTLLM_ADP_WORK_AWARE_ROUTING.Telemetry: four new fields are captured and need telemetry/privacy CODEOWNER review:
kv_cache_config.enable_disagg_completion_reservation,kv_connector_config.admission_query.max_uncharged_parks,kv_connector_config.admission_query.max_offer_deferralsandattention_dp_config.enable_work_aware_routing. Every newmooncake_storefield setstelemetry=False. The golden manifest and both serve schemas are regenerated in the commit that adds each field.Behaviour changes
supports_speculative_prefill. A capacity-role generation server may also run one-model drafts.mooncake_storerole transfers KV (its host tier is already dropped) now reserves completion headroom when every layer uses full attention, beam width is 1 and there is a single pipeline stage. Windowed or recurrent models keepmain's admission.max_tokenslowers that server's concurrency.role: capacityservers are unaffected.mooncake_storepublishes only blocks whose KV has been computed. On a single rank or under attention DP, a prefix page missing from the pool restarts that request instead of failing the server.kv_transfer_timeout_ms(60 s by default). Every timeout path now moves the request toDISAGG_TRANS_ERROR.Intentionally not included
mooncake_store_m3package andmooncake-store-m3preset,trtllm-serve mooncake_master_m3,mooncake_donorandlaunch_master. The pool manifest and therole: capacityserver onmaincover them.TLLM_CTX_APPEND_RESERVE,_ctx_append_reserve*and_ctx_pending_init. Admission parks are re-expressed againstmain's scheduler_v2 without it.alignwithasync_load(d169724). Under that combination a one-block offer parks; this is a follow-up.Test coverage
Both trees ran the same test selections on GB300 (aarch64, 4 GPUs per node):
mainat 0746a23 and this branch. Each ran in the CI dev image built from its own tree, with Slurm and MPI variables cleared as injenkins/scripts/slurm_run.sh. The full_torch/executor,disaggregatedand e2e rows also applied the UCX fixup fromjenkins/scripts/slurm_env_setup.sh.main0746a23pytest -m cpu_only _torch/executor_torch/test_connector*.py,_torch/disaggregationdisaggregated, serving routes, forward-error survivalllmapi/test_llm_args.py, serve config schemas,usage,api_stability,kv_cache_manager_v2_testsscripts/generate_trtllm_serve_schemas.py --check,scripts/generate_llm_args_golden_manifest.py --checkkv_cache_manager_v2_tests/test_kv_cache_manager_v2.pykvCacheManagerV2*Testgtest binarieskvCacheManagerV2ConfigTest7 to 8 tests)_torch/executordisaggregated, with the cluster's injected UCX settingsmaindisaggregated, with theoci-jhbUCX rule fromtests/test_common/cluster_env.pyllmapi/test_llm_api_connector_mooncake.pyEvery failure listed also fails on
main:tests/unittest/_torch/test_connector_adp.py. Newermainskips that file under https://nvbugs/7003528.test_pytorch_model_engine.py::DecoderForwardTestCase::test_mrope_*.disaggregatedfailures under the injected settings are this cluster's UCX environment. UCX pickstcp/rdma_vf_rail0and binds an IPv6 address the container cannot assign, so NIXL cannot create its UCX backend.tests/test_common/cluster_env.pydocuments the same issue.oci-jhbUCX rule, almost everything passes.test_cache_transceiver_precheck_e2e.py::test_precheck_declared_architecture_without_checkpointand fourtest_py_cache_transceiver_mp.py::test_v2_transceiver_mpcases.test_v2_transceiver_mp[ctx_first_sync-v2_mp_tp1_pp2_to_tp1_pp1], is the casemain's run hung on until the 45-minute limit.test_py_cache_transceiver_mp.pyran separately: 580 passed on each tree, outcomes identical.Revisions: the CPU rows ran at e8eea61828, which differs from the head only in the e2e test file. The rows for touched GPU files,
kv_cache_manager_v2, and the C++ gtests ran at 181aeecd18. Since then only test stubs, the warning guard on the estimation manager and its test, and the e2e test changed, and the full_torch/executoranddisaggregatedGPU rows and the e2e row reran at the head, 2cebc83. There is no C++ change after the gtest run.Also:
pre-commitis clean on every changed file.main18c972e without conflicts. On the four files both sides change,main's schema changes are only the version string, andrufffinds no undefined or unused names in the merged files.Coordination with open upstream PRs
reuse_match_alignment, and a speculative-decoding gatetests/unittest/_torch/executor; test paths will need a rebaseDesign notes (internal cookbook): /home/simengl/shared/trt_repos/trtllm-cookbook/projects/minimax-m3-mooncake/design/2026-10-27-mooncake-gap-port/ and /home/simengl/shared/trt_repos/trtllm-cookbook/projects/minimax-m3-mooncake/analysis/upstream-gap/2026-10-27-main-vs-M3-mooncake-main.md
🤖 Generated with Claude Code