Repository navigation
Conversation
…fix loads and long-prompt re-save Let the mooncake-store KV connector serve models that mix sliding-window and full-attention layers, with a hidden-state drafter: - Save a window group's pages once per prompt, at its last chunk (the window's sinks and tail at the prompt end). Look up the longest full-attention hit, walk back up to swa_endpoint_rewind_tokens for a prompt end whose window is stored, and optionally fall back to an earlier one. kv_cache_layout carries each group's num_sink_tokens. - Load stored prefixes through V2 prefix reservations on loader threads, each with its own staging pool and stream, while the request waits outside the batch (TRTLLM_MOONCAKE_STORE_LOAD_THREADS, _SYNC_LOAD). - Under speculative decoding save only prompt blocks, and admit mooncake-store through the executor's connector/speculation gate. - Cap a context request's saves at its chunk's end, and save only the window pages the request still holds. - Optionally re-list a long prompt's full-attention blocks at its last chunk, refreshing the chain's head in the pool (TRTLLM_MOONCAKE_STORE_RESAVE_MIN_TOKENS). - Build lookup hashes and keys with one struct.pack per chunk and per-namespace key columns (byte-identical), and split long existence probes into concurrent chunks. - Queue window pages ahead of full-attention blocks, log lookup, save and load statistics, refuse enable_swa_scratch_reuse and a nonzero context_recompute_tail, and document the connector's window support and environment variables. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Member
Author
|
/bot run |
Collaborator
|
PR_Github #81591 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #81591 [ run ] completed with state |
Collaborator
|
Hey @Tabrizian, my agent had some feedback on your PR. I think 1) and maybe 2) might be important to fix: Hi Iman — we read through #21270 (Mooncake store: SWA layouts, async prefix loads, long-prompt re-save) while comparing it with our own MSO Mooncake branch. Nice work, and the 94–96% pool reuse you reported is far better than what we got. A few findings you may want to look at. All of these come from reading the code, not from GPU runs.
|
Collaborator
|
And another race condition: PR #21270 (Mooncake store: sliding-window layouts, async prefix loads, re-save): window-page save raceStatus: draft for the user to send to the PR author. Nothing has been posted.
The raceThe save thread reads live device pages once the pass's event fires (
In each case the store gets bytes that are not this prompt's KV, under this prompt's valid key. A later hit Why page digests cannot see it
Repro (GB300, MSO 816B TEP4 + DFlash K=7, aggregated, greedy, 16 prompts of ~10-14k tokens, U2 harness)Cold pass: compute and save. Restart, then the hit pass loads from the store. Each hit is compared with local
Fix (the patch)
Separate candidates (real issues in the PR head, NOT in this patch)
Details are in Test scope note
|
pcastonguay
reviewed
Nov 9, 2026
pcastonguay
left a comment
pcastonguay
left a comment
Collaborator
There was a problem hiding this comment.
Suggested fix for the window-page save race described in my comment above. The patch is a single commit against this PR's head (07b49ac), and git apply --check is clean.
What it does
- In
wait_for_save, the window pages a pass saves are copied withcudaMemcpyAsyncon the execution stream into a pinned, store-registered host ring (SnapshotRing, newsnapshot.py). This happens before the event is recorded. The save thread then puts those pages from their ring slots and frees the slots. - Full-attention pages keep the live path, since the request holds them until its saves land.
- Volume is about 2 windows per prompt end (~130 pages on our model), not every pass. Default ring: 2048 slots per rank (
TRTLLM_MOONCAKE_STORE_WINDOW_SNAPSHOT_SLOTS). - If a pass overflows the ring, its window pages are dropped and lookups fall back to an earlier prompt end. It does not stall.
- Kill switch:
TRTLLM_MOONCAKE_STORE_WINDOW_SNAPSHOT=0. - Adds 2 CPU tests.
Verified
- PR head + patch: the connector CPU suite passes 72/72 (your 70 + the 2 new).
- On our GB300 MuseSparkOpen port (SWA + full attention + DFlash drafter), a reuse-heavy 16-prompt check of store hits against local block reuse:
- before the fix: 9/16, 13/16 (sync loads), 15/16; it never reproduces;
- after the fix: 16/16 bitwise in 3 of 3 runs, with page digests clean.
pr21270_window_snapshot.patch
From 29ae074cfc7c2a41a5d455a2a844a2a67fb6711f Mon Sep 17 00:00:00 2001
From: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com>
Date: Sun, 8 Nov 2026 20:41:39 -0500
Subject: [PATCH] [None][fix] Mooncake store: snapshot window pages before the
save thread reads them
The save thread reads live device pages after the pass's event, but the commit
that closes a prompt and every decode step after it release sliding-window
blocks as the window slides (see _append_window_pages), so a window page can be
recycled and rewritten before the thread copies it, and the store then serves
another request's KV for that window.
wait_for_save now copies each saved window page into a pinned, store-registered
host slot on the execution stream ahead of the pass's event; the save thread
puts those pages from their slots and releases them. Full-attention pages keep
the live path. A pass whose window pages do not fit the ring drops them (the
lookup falls back to an earlier prompt end). TRTLLM_MOONCAKE_STORE_WINDOW_SNAPSHOT=0
turns it off; TRTLLM_MOONCAKE_STORE_WINDOW_SNAPSHOT_SLOTS sizes the ring (default 2048).
Signed-off-by: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com>
---
.../connectors/mooncake_store/snapshot.py | 98 +++++++++++++++
.../connectors/mooncake_store/worker.py | 113 ++++++++++++++++--
.../executor/test_mooncake_store_connector.py | 46 +++++++
3 files changed, 250 insertions(+), 7 deletions(-)
create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/snapshot.py
diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/snapshot.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/snapshot.py
new file mode 100644
index 0000000000..1b29c262d3
--- /dev/null
+++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/snapshot.py
@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Save snapshots: copy a pass's saved window pages out before anything can reuse their slots.
+
+`wait_for_save` copies a pass's sliding-window pages into slots of a pinned,
+store-registered host ring with `cudaMemcpyAsync` on the execution stream,
+records the pass's event after the copies, and returns. The save thread waits
+for that event and puts straight from the slots.
+
+Why this closes the race without blocking: every later writer of a page slot is
+ordered behind the execution stream's earlier work. The next forward pass runs
+on that stream, and the KV cache manager moves pages on it. The commit that
+unlocks stale sliding-window pages happens on the host after `wait_for_save`
+returns, so the copies are queued before any reuse of those slots. The executor
+pays the host time to issue the copies (one `cudaMemcpyAsync` per page region)
+and the copy engine pays bytes / D2H bandwidth on the stream.
+"""
+
+import threading
+from typing import Callable, List, Optional, Sequence, Tuple
+
+from .staging import HostStagingPool
+
+
+class SnapshotRing:
+ """Slots of a `HostStagingPool` handed out per page and returned once put."""
+
+ def __init__(self, pool: HostStagingPool) -> None:
+ self._pool = pool
+ self._free: List[int] = list(range(pool.num_slots))
+ self._cond = threading.Condition()
+ #: Passes that had to wait for slots, and passes too large to fit.
+ self.waits = 0
+ self.oversized = 0
+
+ @property
+ def num_slots(self) -> int:
+ return self._pool.num_slots
+
+ @property
+ def free_slots(self) -> int:
+ with self._cond:
+ return len(self._free)
+
+ def try_reserve(self, count: int) -> Optional[List[int]]:
+ """`count` slots, or None when that many are not free right now."""
+ with self._cond:
+ if count > len(self._free):
+ return None
+ taken, self._free = self._free[:count], self._free[count:]
+ return taken
+
+ def reserve(self, count: int, alive: Callable[[], bool], timeout: float = 1.0) -> Optional[List[int]]:
+ """`count` slots, waiting for the save thread to return some.
+
+ Returns None when `count` can never fit or `alive()` turns false while
+ waiting, so the caller falls back to another way of protecting pages.
+ """
+ if count > self._pool.num_slots:
+ self.oversized += 1
+ return None
+ with self._cond:
+ waited = False
+ while count > len(self._free):
+ if not alive():
+ return None
+ waited = True
+ self._cond.wait(timeout)
+ if waited:
+ self.waits += 1
+ taken, self._free = self._free[:count], self._free[count:]
+ return taken
+
+ def release(self, slots: Sequence[int]) -> None:
+ with self._cond:
+ self._free.extend(slots)
+ self._cond.notify_all()
+
+ def gather(
+ self, slot: int, addresses: Sequence[int], sizes: Sequence[int], stream: int
+ ) -> Tuple[int, int]:
+ """Queue the copy of one page into `slot`; returns (slot address, bytes)."""
+ return self._pool.gather(slot, addresses, sizes, stream)
+
+ def close(self) -> None:
+ self._pool.close()
diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py
index 8ea9dd3105..3013e11bf8 100644
--- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py
+++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py
@@ -62,6 +62,7 @@ from .addressing import PageAddressing
from .config import CONFIG_PATH_ENV, MooncakeStoreConnectorConfig, local_address, pool_config
from .keys import KeyNamespace
from .ledger import record_segment
+from .snapshot import SnapshotRing
from .metadata import MooncakeStoreMetadata, PrefixLoadTransfers, RequestTransfers
from .staging import (
MAX_STAGING_BUFFER_BYTES,
@@ -103,6 +104,19 @@ DEFAULT_LOAD_THREADS = 2
#: request time instead of the batch.
SYNC_LOAD_ENV = "TRTLLM_MOONCAKE_STORE_SYNC_LOAD"
+#: Window-page save snapshots: `wait_for_save` copies every sliding-window page a pass saves
+#: into a pinned, store-registered host slot with `cudaMemcpyAsync` on the execution stream, ahead of the pass's
+#: event, and the save thread puts from that slot. Without it the save thread reads the live device page after the
+#: event, but the commit that closes a prompt (and every decode step after it) releases window blocks as the window
+#: slides, so a page can be recycled and rewritten before the thread reads it: the store then holds bytes that are
+#: not this prompt's KV, and a later load faithfully returns those wrong bytes. Full-attention
+#: pages stay held by the request until its saves land, so only window groups need it, which keeps the copy to the
+#: ~2 windows a prompt end saves. `TRTLLM_MOONCAKE_STORE_WINDOW_SNAPSHOT=0` turns it off; the slot count is
+#: `TRTLLM_MOONCAKE_STORE_WINDOW_SNAPSHOT_SLOTS` (default 2048 pages per rank).
+WINDOW_SNAPSHOT_ENV = "TRTLLM_MOONCAKE_STORE_WINDOW_SNAPSHOT"
+WINDOW_SNAPSHOT_SLOTS_ENV = "TRTLLM_MOONCAKE_STORE_WINDOW_SNAPSHOT_SLOTS"
+DEFAULT_WINDOW_SNAPSHOT_SLOTS = 2048
+
#: Whether a lookup that finds no window within the rewind looks further back for an
#: earlier prompt end whose window is stored, instead of reporting a miss (on by default; 0 turns it off).
WINDOW_FALLBACK_ENV = "TRTLLM_MOONCAKE_STORE_WINDOW_FALLBACK"
@@ -381,6 +395,9 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker):
# The cache's own life cycle per sliding-window group, which decides
# the blocks such a group still reads at a given history length.
self._life_cycles: Dict[int, AttnLifeCycle] = {}
+ #: Pinned slots window pages are copied into on the execution stream before a save reads them.
+ self._window_snapshot: Optional[SnapshotRing] = None
+ self._window_snapshot_dropped = 0
# Checked before the segment is mounted, since the kernel answers an
# unaffordable total by killing the process rather than by failing the
@@ -544,6 +561,7 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker):
self._life_cycles[group.layer_group_id] = AttnLifeCycle.make(
group.window_size, group.num_sink_tokens, layout.tokens_per_block
)
+ self._open_window_snapshot(addressing)
if self._life_cycles and self._config.role.loads and not self._prefix_aware_scheduling:
raise ValueError(
"The mooncake-store connector serves sliding-window layer groups only through prefix "
@@ -603,6 +621,58 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker):
f"mooncake-store worker rank {self._rank} registered layout: {addressing.describe()}"
)
+ def _open_window_snapshot(self, addressing: PageAddressing) -> None:
+ """Pinned slots for window-page save snapshots (see `WINDOW_SNAPSHOT_ENV`)."""
+ if not (self._config.role.saves and self._life_cycles and torch.cuda.is_available()):
+ return
+ if os.environ.get(WINDOW_SNAPSHOT_ENV, "1") == "0":
+ logger.warning(f"mooncake-store rank {self._rank}: {WINDOW_SNAPSHOT_ENV}=0, window pages are saved live")
+ return
+ slot_bytes = max(addressing.bytes_per_page(group) for group in self._life_cycles)
+ num_slots = int(os.environ.get(WINDOW_SNAPSHOT_SLOTS_ENV, DEFAULT_WINDOW_SNAPSHOT_SLOTS))
+ self._window_snapshot = SnapshotRing(
+ HostStagingPool(slot_bytes=slot_bytes, num_slots=num_slots, store=self._store, label="window-snapshot")
+ )
+
+ def _snapshot_window_pages(self, saves: List[RequestTransfers], stream) -> Tuple[List[RequestTransfers], Dict]:
+ """Copy this pass's window pages to pinned slots on `stream`; returns (saves, overrides).
+
+ overrides maps (layer group, page index) to (slot, slot address, bytes). A pass whose window pages
+ do not fit the ring drops them from the save (a later lookup falls back to an earlier prompt end)
+ rather than saving pages that may be recycled before the save thread reads them.
+ """
+ ring = self._window_snapshot
+ windowed = set(self._life_cycles)
+ wanted = []
+ for entry in saves:
+ for page in entry.pages:
+ if page.layer_group_id in windowed:
+ wanted.append((page.layer_group_id, page.page_index))
+ wanted = list(dict.fromkeys(wanted))
+ if not wanted:
+ return saves, {}
+ slots = ring.try_reserve(len(wanted))
+ if slots is None:
+ slots = ring.reserve(len(wanted), alive=lambda: self._save_thread is not None and self._save_thread.is_alive())
+ if slots is None:
+ self._window_snapshot_dropped += len(wanted)
+ logger.warning(
+ f"mooncake-store rank {self._rank}: {len(wanted)} window pages do not fit the "
+ f"{ring.num_slots}-slot window snapshot; not saving them ({self._window_snapshot_dropped} so far)"
+ )
+ kept = [
+ RequestTransfers(entry.request_id, [p for p in entry.pages if p.layer_group_id not in windowed])
+ for entry in saves
+ ]
+ return kept, {}
+ handle = _stream_handle(stream)
+ overrides = {}
+ for slot, (group, page_index) in zip(slots, wanted):
+ addresses, sizes = self._addressing.buffers(group, page_index)
+ slot_address, total = ring.gather(slot, addresses, sizes, handle)
+ overrides[(group, page_index)] = (slot, slot_address, total)
+ return saves, overrides
+
def _open_staging(self, addressing: PageAddressing) -> None:
"""Allocate and register the pinned slots pages will pass through.
@@ -1172,13 +1242,18 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker):
# The pages are written by kernels still queued on this stream, so the
# event is the handoff: the thread reads GPU memory only after the
# pass retires, without the executor loop waiting for it.
+ saves, overrides = list(metadata.saves), {}
+ if self._window_snapshot is not None:
+ # Queued on this stream ahead of the event, so the copies read the pages before any later
+ # forward, commit or recycle can rewrite them.
+ saves, overrides = self._snapshot_window_pages(saves, stream)
event = torch.cuda.Event()
event.record(stream)
with self._save_lock:
- for transfers in metadata.saves:
+ for transfers in saves:
self._outstanding_saves[transfers.request_id] += 1
- self._save_queue.put((event, list(metadata.saves), time.monotonic()))
+ self._save_queue.put((event, saves, time.monotonic(), overrides))
def get_finished(
self, finished_gen_req_ids: List[int], started_loading_req_ids: List[int]
@@ -1240,12 +1315,13 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker):
item = self._save_queue.get()
if item is None:
return
- event, transfers, queued_at = item
+ event, transfers, queued_at, *rest = item
+ overrides = rest[0] if rest else {}
began = time.monotonic()
written = present = written_bytes = 0
try:
event.synchronize()
- written, present, written_bytes = self._put(transfers)
+ written, present, written_bytes = self._put(transfers, overrides)
except Exception as exc:
# Broad on purpose: this is the thread boundary. The batch is
# dropped and the request retired by the `finally` below, the
@@ -1255,6 +1331,8 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker):
f"{type(exc).__name__}: {exc}\n{traceback.format_exc()}"
)
finally:
+ if overrides:
+ self._window_snapshot.release([slot for slot, _, _ in overrides.values()])
with self._save_lock:
for entry in transfers:
remaining = self._outstanding_saves.get(entry.request_id, 0) - 1
@@ -1270,18 +1348,38 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker):
f"mooncake-store rank {self._rank} saves: {report}, {self._save_queue.qsize()} queued"
)
- def _put(self, transfers: Sequence[RequestTransfers]) -> Tuple[int, int, int]:
+ def _put(self, transfers: Sequence[RequestTransfers], overrides: Optional[Dict] = None) -> Tuple[int, int, int]:
"""Write the pages the pool does not hold yet.
+ `overrides` (from `_snapshot_window_pages`) are window pages already copied to pinned slots: they
+ are put from those slots, the rest from the device pages through the save staging pool.
+
Returns:
Pages written, pages already stored, and the bytes written.
"""
- keys, addresses, sizes, _ = self._resolve(transfers)
+ overrides = overrides or {}
+ live, snap = [], []
+ for entry in transfers:
+ live.append(RequestTransfers(entry.request_id, [p for p in entry.pages if (p.layer_group_id, p.page_index) not in overrides]))
+ snap.extend((entry.request_id, p) for p in entry.pages if (p.layer_group_id, p.page_index) in overrides)
written = present_count = written_bytes = 0
+ if snap:
+ keys, _, _, _ = self._resolve([RequestTransfers(rid, [p]) for rid, p in snap])
+ addresses = [[overrides[(p.layer_group_id, p.page_index)][1]] for _, p in snap]
+ sizes = [[overrides[(p.layer_group_id, p.page_index)][2]] for _, p in snap]
+ w, pc, wb = self._put_resolved(keys, addresses, sizes, staged=True)
+ written, present_count, written_bytes = w, pc, wb
+ keys, addresses, sizes, _ = self._resolve(live)
if not keys:
return written, present_count, written_bytes
+ w, pc, wb = self._put_resolved(keys, addresses, sizes, staged=False)
+ return written + w, present_count + pc, written_bytes + wb
+
+ def _put_resolved(self, keys, addresses, sizes, staged: bool) -> Tuple[int, int, int]:
+ """`_put`'s transfer loop over resolved pages; `staged` pages already sit in registered host slots."""
+ written = present_count = written_bytes = 0
- staging = self._save_staging
+ staging = None if staged else self._save_staging
handle = _stream_handle(self._save_stream) if staging is not None else 0
for batch in zip(
@@ -1407,6 +1505,7 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker):
self._load_staging = None
self._async_load_stagings = []
self._save_staging = None
+ self._window_snapshot = None
self._save_stream = None
global _LOCAL_WORKER
if _LOCAL_WORKER is self:
diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py
index 4af7e53982..04543fba11 100644
--- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py
+++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py
@@ -1651,3 +1651,49 @@ def test_worker_refuses_a_window_layout_without_prefix_reservations(store_config
fake_store.workers.append(worker)
with pytest.raises(ValueError, match="enable_prefix_aware_scheduling"):
worker.register_kv_cache_layout(make_window_layout())
+
+
+def test_put_ships_window_pages_from_their_snapshot_slots(store_config, fake_store):
+ """mcgb2: a window page snapshotted by `wait_for_save` is put from its pinned slot, the rest live.
+
+ The store must receive exactly the slot address (the bytes copied on the execution stream before the
+ page could be recycled), never the device page the save thread would read later.
+ """
+ with make_worker(fake_store, layout=make_layout(num_groups=2)) as worker:
+ worker._save_staging = None
+ window_hash, full_hash = b"\x0a" * 16, b"\x0b" * 16
+ overrides = {(1, 3): (0, 0xABC000, 4096)}
+ written, present, _ = worker._put(
+ [RequestTransfers(1, [PageTransfer(full_hash, 0, 2), PageTransfer(window_hash, 1, 3)])],
+ overrides,
+ )
+ assert (written, present) == (2, 0)
+ by_key = {}
+ for keys, addresses, sizes in fake_store.put_calls:
+ for key, address, size in zip(keys, addresses, sizes):
+ by_key[key] = (address, size)
+ assert by_key[worker._namespaces[1].key(window_hash)] == ([0xABC000], [4096])
+ full_addresses, full_sizes = PageAddressing(make_layout(num_groups=2)).buffers(0, 2)
+ assert by_key[worker._namespaces[0].key(full_hash)] == (list(full_addresses), list(full_sizes))
+
+
+def test_drain_releases_window_snapshot_slots(store_config, fake_store):
+ """The save thread returns a pass's snapshot slots once its puts finish, so the ring cannot leak."""
+ released = []
+ with make_worker(fake_store, layout=make_layout(num_groups=2)) as worker:
+ worker._save_staging = None
+ worker._window_snapshot = SimpleNamespace(release=lambda slots: released.extend(slots))
+ with worker._save_lock:
+ worker._outstanding_saves[7] += 1
+ worker._save_queue.put(
+ (
+ SimpleNamespace(synchronize=lambda: None),
+ [RequestTransfers(7, [PageTransfer(b"\x0c" * 16, 1, 1)])],
+ time.monotonic(),
+ {(1, 1): (5, 0xDEF000, 4096)},
+ )
+ )
+ deadline = time.monotonic() + 5
+ while not released and time.monotonic() < deadline:
+ time.sleep(0.01)
+ assert released == [5]
--
2.50.1 (Apple Git-155)
|
|
||
| Only the blocks the window still reads at the prompt's end. The commit | ||
| that closes the last chunk releases every block below that window in | ||
| the same iteration, before the save thread copies anything, so a page |
Collaborator
There was a problem hiding this comment.
This release-before-copy also bites the pages inside the window. The prompt-closing commit, and every decode step after it, releases window blocks as the window slides. So a window page this pass queued for saving can be recycled and rewritten before the background thread copies it.
The save and load digests then both see the same overwritten bytes, so the digests can't catch it. The symptom is store hits that diverge from local block reuse, differently on every run. Suggested fix: snapshot the saved window pages in wait_for_save; the patch is in the review summary.
| try: | ||
| event.synchronize() | ||
| self._put(transfers) | ||
| written, present, written_bytes = self._put(transfers) |
Collaborator
There was a problem hiding this comment.
event.synchronize() orders this read after the kernels of the pass that recorded the event. It does not order it before later passes, which can already have reused released window blocks (see the comment on _append_window_pages). Live reads are fine for full-attention pages, because the request holds them until its saves land. Window pages need a copy taken on the execution stream before the event; the patch in the summary does that.
Tabrizian
marked this pull request as ready for review
November 10, 2026 17:37
Tabrizian
requested review from
Shixiaowei02,
YihuiLu512,
arysef,
lori-ren and
pcastonguay
November 10, 2026 17:37
Tabrizian
marked this pull request as draft
November 10, 2026 17:38
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.
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:
swa_endpoint_rewind_tokens.TRTLLM_MOONCAKE_STORE_WINDOW_FALLBACK, on by default), it probes further back for an earlier prompt end with a complete window.kv_cache_layoutnow carries each group'snum_sink_tokens.validate_layoutaccepts window groups. Registration refuses a window layout without prefix-aware scheduling, which window loads need.validate_llm_argsrefuseskv_cache_config.enable_swa_scratch_reuseand a nonzerospeculative_config.context_recompute_tail. With either, the connector could save or load a window page it never verified.reserve_prefix/release_prefix_reservation).TRTLLM_MOONCAKE_STORE_LOAD_THREADSsets the thread count (default 2).TRTLLM_MOONCAKE_STORE_SYNC_LOAD=1restores synchronous loads.mooncake-store.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.struct.packand slices it per block. Keys are built column-wise per namespace.TRTLLM_MOONCAKE_STORE_PROBE_CHUNK, default 4096 keys).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.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
scheduler_config.enable_prefix_aware_schedulingis on.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.test_mooncake_store_{cli,common,settings,real_pool}.py,test_kv_cache_layout.py,test_kv_connector_reservations.pyandtest_kv_connector_v2_prefix.py.PR Checklist