[Tabrizian]

Description

The Mooncake store connector (#19235, #19171) shares context KV through a Mooncake pool, but only for full-attention layouts, and it refuses speculative decoding. A model that mixes sliding-window (SWA) and full-attention layers, served with a hidden-state drafter, can't use it. On long multi-turn prompts, it also blocks the executor while a reused prefix loads.

This PR makes the connector work for such models:

• Sliding-window layouts. Each window group saves its pages once per prompt, in the prompt's last chunk. It saves only the blocks the window still reads at the prompt end (its sinks plus its tail), reaching back swa_endpoint_rewind_tokens.
  • A lookup takes the longest full-attention hit. It then walks back over the rewind to the latest prompt end whose window is complete in the pool.
  • Optionally (TRTLLM_MOONCAKE_STORE_WINDOW_FALLBACK, on by default), it probes further back for an earlier prompt end with a complete window.
  • Loads cover only the blocks the runtime resumes at. kv_cache_layout now carries each group's num_sink_tokens.
  • validate_layout accepts window groups. Registration refuses a window layout without prefix-aware scheduling, which window loads need.
  • validate_llm_args refuses kv_cache_config.enable_swa_scratch_reuse and a nonzero speculative_config.context_recompute_tail. With either, the connector could save or load a window page it never verified.
• Asynchronous prefix loads through the V2 prefix reservations (reserve_prefix / release_prefix_reservation).
  • A request whose prefix is in the pool is parked while loader threads move its pages. Each thread has its own staging pool and stream.
  • The context instance keeps prefilling other requests in the meantime.
  • TRTLLM_MOONCAKE_STORE_LOAD_THREADS sets the thread count (default 2). TRTLLM_MOONCAKE_STORE_SYNC_LOAD=1 restores synchronous loads.
  • The node memory budget counts each thread's staging pool.
• Speculative decoding. Under a drafter, only prompt blocks are saved. They are filled before the first draft, so rejected drafts never rewind a page the connector addresses. The executor's connector-with-speculation gate now admits mooncake-store.
• Save-path fixes.
  • A context request's saves are capped at the end of the current chunk. Pages the cache allocates past it, a drafter's extra KV tokens, are never published with stale contents.
  • Window saves cover only what the request still holds after the last chunk.
• Long-prompt re-save (TRTLLM_MOONCAKE_STORE_RESAVE_MIN_TOKENS, off by default). A prompt that computed at least that many tokens lists its full-attention blocks again at its last chunk. That refreshes the chain's head in the pool and rewrites blocks the pool evicted during a long prefill. Otherwise one evicted head block makes the whole stored chain unreachable, and the next turn re-prefills all of it.
• Cheaper lookups.
  • The block-hash chain packs a chunk's tokens with one struct.pack and slices it per block. Keys are built column-wise per namespace.
  • Long existence probes split into concurrent chunks (TRTLLM_MOONCAKE_STORE_PROBE_CHUNK, default 4096 keys).
  • Hashes and keys are byte-identical to before (tested).
• Window-first save order (TRTLLM_MOONCAKE_STORE_WINDOW_FIRST, on by default). Window pages are queued ahead of a prompt's full-attention blocks, so a stored chain's window lands first.
• Statistics. The leader logs lookup outcomes every 30 s at INFO: full, rewound, recovered, window miss and full-attention miss, plus probe time. Each rank logs saver and loader throughput and queueing.

The changes are in one PR because they are interdependent: SWA lookups need the reservation path to load asynchronously, and the save fixes and re-save are all on the SWA save path.

Impact

• Full-attention-only layouts keep their behaviour. Their lookups and saves change only in speed (hashing, chunked probes) and in the save cap at chunk boundaries, which is a correctness fix.
• New behaviour is behind the reservation path, which the V2 manager already uses when a connector implements it and scheduler_config.enable_prefix_aware_scheduling is on.
• Measured on a disaggregated deployment (one context, one generation instance) replaying long multi-turn conversations of 100k-170k tokens, with an SWA + full-attention model and a hidden-state drafter:
  • The context's prefix reuse from the pool stays at ~94-96 % of prompt tokens.
  • Long-prompt re-save, together with master eviction settings (high watermark 0.97, eviction ratio 0.01), raised request throughput ~40 % over the same configuration without them.
  • The cheaper lookup path added ~4.5 % throughput.
• Known limitation: with a hidden-state drafter, a reused request's drafter state starts from the computed tail. Draft acceptance, and with it the verify steps' numerics, can therefore differ from local reuse. KV equality between store-served and locally reused prefixes was checked with drafting off.
• Not yet exercised: the unified target/draft KV cache from [None][feat] Unified target/draft KV cache for DFlash (port of #20178, #20486 and #20709) #20806. When draft layers share the V2 manager, their layer groups appear in the connector's layout and are saved and loaded like target groups. The measurements above predate it.

Test Coverage

• tests/unittest/_torch/executor/test_mooncake_store_connector.py: window saves, rewind and fallback lookups, window-first ordering, reservation loads (async and sync), loader threads, chunked probes, the chunk-end save cap, re-save, hash-chain and key equivalence with the per-token encoding, and the speculative-decoding prompt-only saves.
• Also run unchanged: test_mooncake_store_{cli,common,settings,real_pool}.py, test_kv_cache_layout.py, test_kv_connector_reservations.py and test_kv_connector_v2_prefix.py.
• The eight suites above: 260 passed, none skipped, on a GPU node in a container built from this branch's base, with the changed connector files mounted over it. The run predates [None][feat] Unified target/draft KV cache for DFlash (port of #20178, #20486 and #20709) #20806.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

[Tabrizian]

/bot run

[tensorrt-cicd]

PR_Github #81591 [ run ] triggered by Bot. Commit: 07b49ac Link to invocation

[tensorrt-cicd]

PR_Github #81591 [ run ] completed with state SUCCESS. Commit: 07b49ac
/LLM/main/L0_MergeRequest_PR pipeline #67257 completed with status: 'SUCCESS'

CI Report

Link to invocation