From 07b49ac22e10146d127ac562829bf5b00e899f2d Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Thu, 29 Oct 2026 13:05:38 -0400 Subject: [PATCH] [None][feat] Mooncake store: sliding-window layouts, asynchronous prefix 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> --- docs/source/features/kv-cache-connector.md | 43 +- .../pyexecutor/connectors/kv_cache_layout.py | 10 + .../connectors/mooncake_store/keys.py | 29 +- .../connectors/mooncake_store/metadata.py | 19 +- .../connectors/mooncake_store/scheduler.py | 373 +++++++- .../connectors/mooncake_store/validation.py | 48 +- .../connectors/mooncake_store/worker.py | 793 ++++++++++++++++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 9 +- .../_torch/pyexecutor/py_executor_creator.py | 7 +- .../executor/test_mooncake_store_connector.py | 756 ++++++++++++++++- 10 files changed, 1977 insertions(+), 110 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 2d93e25c9194..0e3091fc66b0 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -193,13 +193,14 @@ following configuration limits at bring-up. | Configuration | Mechanism | |---|---| -| Speculative decoding | Rejected draft tokens shrink a request's page list, and the freed slot goes to whichever request allocates next. The connector is never told the tail block moved. | +| Speculative decoding | Rejected draft tokens shrink a request's page list, and the freed slot goes to whichever request allocates next. The connector is never told the tail block moved. `mooncake-store` is exempt: under a drafter it saves only prompt blocks, which are filled before the first draft. | | A capacity scheduler policy other than `GUARANTEED_NO_EVICT` with KV cache manager V1 | V1 does not reset the connector block delta when a request is destroyed and replayed. V2 resets that state and permits replay; an active connector load retains its allocation until completion. | | A host or disk cache tier | Tier eviction reassigns the GPU slot. See [KV cache tiers are GPU-only under a connector](#kv-cache-tiers-are-gpu-only-under-a-connector). | The exact set the runtime refuses depends on your cache configuration; the bring-up error is authoritative. Speculative decoding is unsupported with a connector wherever it is not refused — -the same page-list shrink happens there. +the same page-list shrink happens there — except with `mooncake-store`, which saves only prompt +blocks while drafts are on. ##### `RequestData` fields that may not be populated @@ -487,12 +488,38 @@ settings. A `capacity` server keeps both as configured. Checked at startup, before any request is admitted, because each one's failure mode is KV replayed without all of the state it was computed with — a wrong answer rather than a slow one. -Beyond beam search, attention data parallelism, speculative decoding and Mamba/hybrid models, which -are rejected for [every connector](#a-page-slot-must-not-be-reassigned-underneath-the-connector), -this one also rejects KV cache manager V1, context parallelism, pipeline parallelism, -sliding-window attention, and `sparse_attention_config` with an index-V cache unless -`sparse_disable_index_value=True`; it bypasses individual requests that carry a LoRA adapter or -multimodal content without hashes, leaving the rest of the deployment served. +Beyond beam search, attention data parallelism and Mamba/hybrid models, which are rejected for +[every connector](#a-page-slot-must-not-be-reassigned-underneath-the-connector), this one also +rejects KV cache manager V1, context parallelism, pipeline parallelism, and +`sparse_attention_config` with an index-V cache unless `sparse_disable_index_value=True`; it +bypasses individual requests that carry a LoRA adapter or multimodal content without hashes, +leaving the rest of the deployment served. Speculative decoding is accepted: only prompt blocks +are saved while drafts are on. + +With sliding-window layers it also rejects `kv_cache_config.enable_swa_scratch_reuse`, a nonzero +`speculative_config.context_recompute_tail`, and `scheduler_config.enable_prefix_aware_scheduling=False`. +Each would let it save or load a window page it never verified. + +### Sliding-window layers + +A sliding-window group keeps only the blocks its window still reads, so the connector saves its +pages once per prompt, in the prompt's last chunk: the window's sinks and tail at the prompt end. +That is where the next turn of a conversation resumes. A lookup takes the longest full-attention +hit, then walks back up to `block_reuse_config.swa_endpoint_rewind_tokens` for a prompt end whose +window is complete in the pool. With none found, it looks further back for the latest earlier end +that has one. A stored prefix is loaded through a prefix reservation: the request waits outside +the batch while loader threads move its pages, and the server keeps computing the others. + +### Environment variables + +| Variable | Default | Effect | +|---|---|---| +| `TRTLLM_MOONCAKE_STORE_LOAD_THREADS` | `2` | Loader threads per loading rank. Each pins its own staging pool, which the node memory check counts. | +| `TRTLLM_MOONCAKE_STORE_SYNC_LOAD` | `0` | `1` loads reserved prefixes synchronously in the executor's step and starts no loader threads. | +| `TRTLLM_MOONCAKE_STORE_PROBE_CHUNK` | `4096` | Keys per existence-probe call; longer probes go out as concurrent chunks. `0` sends one call. | +| `TRTLLM_MOONCAKE_STORE_WINDOW_FALLBACK` | `1` | `0` reports a miss when no window is stored within the rewind, instead of looking further back. | +| `TRTLLM_MOONCAKE_STORE_WINDOW_FIRST` | `1` | `0` queues a prompt's full-attention blocks for saving before its window pages. | +| `TRTLLM_MOONCAKE_STORE_RESAVE_MIN_TOKENS` | `0` | A prompt that computed at least this many tokens lists its full-attention blocks again at its last chunk. This refreshes the chain's head in the pool and rewrites blocks evicted during a long prefill. `0` turns it off. | ### How pages reach the pool diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py index b0749a168390..a81a486b5439 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py @@ -165,6 +165,10 @@ class KvCacheLayerGroupLayout: #: Attention window for this group, or None for full attention. window_size: Optional[int] regions: Tuple[KvCacheRegion, ...] + #: Leading tokens a sliding window keeps however far it has moved, or None. + #: With ``window_size`` it fixes which blocks the group still reads: the + #: same ``AttnLifeCycle`` the cache masks stale pages with. + num_sink_tokens: Optional[int] = None @property def bytes_per_page(self) -> int: @@ -254,6 +258,11 @@ def _window_size(init_config, local_layer_id: int) -> Optional[int]: return None if window is None else int(window) +def _num_sink_tokens(init_config, local_layer_id: int) -> Optional[int]: + sinks = getattr(init_config.layers[local_layer_id], "num_sink_tokens", None) + return None if sinks is None else int(sinks) + + def build_kv_cache_layout_v2(manager: "KVCacheManagerV2") -> KvCacheLayout: """Describe the cache's page regions for a KV connector. @@ -318,6 +327,7 @@ def build_kv_cache_layout_v2(manager: "KVCacheManagerV2") -> KvCacheLayout: layer_ids=tuple(global_by_local[lid] for lid in local_layer_ids), window_size=_window_size(init_config, local_layer_ids[0]), regions=tuple(regions), + num_sink_tokens=_num_sink_tokens(init_config, local_layer_ids[0]), ) ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py index 89db8dc3f467..cc0cfc72d196 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -25,7 +25,9 @@ bytes a page is. A change to any of those reads as a miss rather than as garbage. """ +import functools import hashlib +import struct from dataclasses import dataclass, field from typing import List, Optional, Sequence, Tuple @@ -123,14 +125,21 @@ def extend(self, tokens: Sequence[int]) -> Sequence[bytes]: f"token list shrank from {len(self._hashes)} to {num_full_blocks} " "full blocks; a hash chain belongs to exactly one request" ) - for ordinal in range(len(self._hashes), num_full_blocks): - start = ordinal * self._tokens_per_block - block = tokens[start : start + self._tokens_per_block] - parent = self._hashes[-1] if self._hashes else self._seed - # Fixed-width little-endian token ids: a delimiter-free encoding - # would let two different token sequences serialize identically. - payload = b"".join(int(token).to_bytes(8, "little", signed=True) for token in block) - self._hashes.append(_digest(parent, payload)) + first = len(self._hashes) + if first == num_full_blocks: + return self._hashes + # Fixed-width little-endian token ids: a delimiter-free encoding would + # let two different token sequences serialize identically. The new + # blocks are packed in one call and sliced per block, which gives the + # same payloads as packing each token, at a fraction of the cost on the + # scheduler's path. + start, end = first * self._tokens_per_block, num_full_blocks * self._tokens_per_block + packed = struct.pack(f"<{end - start}q", *tokens[start:end]) + step = 8 * self._tokens_per_block + parent = self._hashes[-1] if self._hashes else self._seed + for offset in range(0, len(packed), step): + parent = _digest(parent, packed[offset : offset + step]) + self._hashes.append(parent) return self._hashes @@ -149,9 +158,9 @@ class KeyNamespace: tokens_per_block: int bytes_per_page: int - @property + @functools.cached_property def prefix(self) -> str: - """The literal string every key in this namespace starts with.""" + """The literal string every key in this namespace starts with (computed once: a lookup builds thousands).""" return ( f"{self.namespace}/{self.model_key}" f"/w{self.world_size}r{self.rank}" diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py index f1eb2db7ee14..e2eac6671f84 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py @@ -23,7 +23,7 @@ from dataclasses import dataclass, field from typing import List -__all__ = ["MooncakeStoreMetadata", "PageTransfer", "RequestTransfers"] +__all__ = ["MooncakeStoreMetadata", "PageTransfer", "PrefixLoadTransfers", "RequestTransfers"] @dataclass @@ -50,13 +50,28 @@ class RequestTransfers: pages: List[PageTransfer] = field(default_factory=list) +@dataclass +class PrefixLoadTransfers: + """The pages of one confirmed prefix reservation. + + Reported back through `get_finished_prefix_loads` once they are on the + device. An asynchronous one is moved off the executor thread while its + request waits outside the batch. + """ + + reservation_id: int + is_async: bool + transfers: RequestTransfers + + @dataclass class MooncakeStoreMetadata: """Loads to perform before the next forward pass, saves to start after it.""" loads: List[RequestTransfers] = field(default_factory=list) saves: List[RequestTransfers] = field(default_factory=list) + prefix_loads: List[PrefixLoadTransfers] = field(default_factory=list) def __bool__(self) -> bool: """Whether there is any work at all this iteration.""" - return bool(self.loads or self.saves) + return bool(self.loads or self.saves or self.prefix_loads) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py index 0bda85e66232..f6829ced24d4 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -22,8 +22,17 @@ * the page slot index per block ordinal, accumulated across iterations. The manager reports only newly allocated indices each step, but a block is allocated before it is full and is savable only once it is full. + +A sliding-window layer group keeps only the blocks its window still reads, so +its pages are saved once per prompt, in the prompt's last chunk: the blocks the +window still reads where the prompt ends. That is what the next turn of a +conversation resumes from, so a lookup resumes only at a prompt end whose +window is stored, searching back up to `swa_endpoint_rewind_tokens` from the +longest full-attention hit. Full-attention groups are saved block by block. """ +import os +from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple from tensorrt_llm.bindings.internal.batch_manager import LlmRequest @@ -31,15 +40,25 @@ from tensorrt_llm.logger import logger from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX -from ..kv_cache_connector import KvCacheConnectorScheduler, RequestData, SchedulerOutput +from ..kv_cache_connector import KvCacheConnectorScheduler, PrefixLoad, RequestData, SchedulerOutput from .config import MooncakeStoreConnectorConfig from .keys import BlockHashChain, ReuseScope -from .metadata import MooncakeStoreMetadata, PageTransfer, RequestTransfers +from .metadata import MooncakeStoreMetadata, PageTransfer, PrefixLoadTransfers, RequestTransfers from .validation import validate_llm_args -from .worker import MooncakeStoreConnectorWorker, resolve_local_worker +from .worker import ( + RESAVE_MIN_TOKENS_ENV, + SYNC_LOAD_ENV, + MooncakeStoreConnectorWorker, + resolve_local_worker, +) __all__ = ["MooncakeStoreConnectorScheduler"] +#: Whether a prompt's window pages go to the save thread ahead of its full-attention blocks (on by default; 0 queues +#: the full-attention blocks first). A lookup that lands between the two then finds the previous prompt end intact +#: instead of full-attention blocks with no window. +WINDOW_FIRST_ENV = "TRTLLM_MOONCAKE_STORE_WINDOW_FIRST" + def _multimodal_item_digests(request: LlmRequest) -> Optional[Tuple[bytes, ...]]: """Content digest per multimodal item of *request*, in prompt order. @@ -94,6 +113,17 @@ def _reuse_scope(request: LlmRequest) -> Optional[ReuseScope]: return ReuseScope(cache_salt=request.cache_salt, multimodal_digests=digests) +@dataclass +class _Reservation: + """What `reserve_prefix` promised, in tokens, and the windows it verified.""" + + start: int + end: int + #: Per sliding-window group, the first block from which the store held every + #: page up to the promised end. + window_present_from: Dict[int, int] = field(default_factory=dict) + + class _RequestState: """Per-request bookkeeping that has to outlive a single iteration.""" @@ -105,10 +135,28 @@ class _RequestState: "load_first_block", "load_blocks", "emitted_saves", + "request", + "window_present_from", + "saved_window", + "prompt_blocks", + "reservations", + "resume_block", + "resaved", ) - def __init__(self, chain: BlockHashChain): + def __init__(self, chain: BlockHashChain, request: LlmRequest, prompt_blocks: int): self.chain = chain + #: The request itself, read for where this iteration's chunk starts and + #: whether it finishes the prompt. + self.request = request + #: Full blocks of the prompt, the only blocks saved under speculative + #: decoding. + self.prompt_blocks = prompt_blocks + #: Per sliding-window group, the first block of the offer's window the + #: store was found to hold. + self.window_present_from: Dict[int, int] = {} + #: Whether the window pages the prompt ends with were handed over. + self.saved_window = False #: The request's tokens, accumulated from the per-step deltas. self.tokens: List[int] = [] #: Page slot index per block ordinal, per layer group. @@ -119,6 +167,13 @@ def __init__(self, chain: BlockHashChain): self.load_first_block = 0 self.load_blocks = 0 self.emitted_saves = False + #: Open prefix reservations by id, until released. + self.reservations: Dict[int, _Reservation] = {} + #: The block the request's computation started at (the prefix the local cache + #: and the store gave it), read at its first scheduled chunk. + self.resume_block: Optional[int] = None + #: Whether its whole prompt was listed again at the last chunk (long prefills only). + self.resaved = False class MooncakeStoreConnectorScheduler(KvCacheConnectorScheduler): @@ -130,12 +185,32 @@ def __init__(self, llm_args: TorchLlmArgs): validate_llm_args(llm_args) self._config = MooncakeStoreConnectorConfig.resolve(llm_args) self._tokens_per_block = int(llm_args.kv_cache_config.tokens_per_block) + # How far back from the longest full-attention hit a lookup searches for + # a prompt end whose window is stored. + reuse_config = getattr(llm_args.kv_cache_config, "block_reuse_config", None) + self._window_rewind_tokens = int( + getattr(reuse_config, "swa_endpoint_rewind_tokens", 0) or 0 + ) + # Rejected draft tokens shrink a request's page list during generation, + # and the freed slot goes to whichever request allocates next, while the + # page indices here only ever grow. Prompt blocks are allocated and + # filled before the first draft, so saving only those never names a + # slot another request owns. A disaggregated context server, which never + # generates, loses nothing by it. + spec_config = getattr(llm_args, "speculative_config", None) + self._prompt_blocks_only = bool( + spec_config is not None and (getattr(spec_config, "max_draft_len", 0) or 0) > 0 + ) + self._async_loads = os.environ.get(SYNC_LOAD_ENV, "0") != "1" + self._window_first = os.environ.get(WINDOW_FIRST_ENV, "1") != "0" + self._resave_min_tokens = int(os.environ.get(RESAVE_MIN_TOKENS_ENV, "0")) self._requests: Dict[int, _RequestState] = {} self._worker: Optional[MooncakeStoreConnectorWorker] = None logger.warning( f"mooncake-store leader ready (role={self._config.role.value}, " - f"tokens_per_block={self._tokens_per_block})" + f"tokens_per_block={self._tokens_per_block}, " + f"reserved prefixes load {'asynchronously' if self._async_loads else 'synchronously'})" ) def wait_for_initialization(self): @@ -156,31 +231,92 @@ def get_num_new_matched_tokens( Returns: Tokens the store can supply, and `False` for a synchronous load. """ + found = self._lookup(request, num_computed_tokens) + if found is None: + return 0, False + state, first_block, hit_blocks, window_present_from = found + state.load_first_block = first_block + state.load_blocks = hit_blocks + state.window_present_from = window_present_from + return hit_blocks * self._tokens_per_block, False + + def reserve_prefix( + self, request: LlmRequest, num_computed_tokens: int, reservation_id: int + ) -> Tuple[int, bool]: + """Promise the longest stored prefix beyond the local match, loaded later. + + The runtime confirms all or part of the promise when it admits the + request, and the pages move only then, from `SchedulerOutput.prefix_loads`. + Nothing is pinned in the pool: a page evicted in between fails its load, + which is fatal under tensor parallelism, as on the synchronous path. + + Returns: + Tokens promised from `num_computed_tokens`, and whether they load + asynchronously, with the request parked outside the batch meanwhile. + """ + found = self._lookup(request, num_computed_tokens) + if found is None: + return 0, False + state, first_block, hit_blocks, window_present_from = found + start = first_block * self._tokens_per_block + end = start + hit_blocks * self._tokens_per_block + state.reservations[reservation_id] = _Reservation(start, end, window_present_from) + return end - start, self._async_loads + + def release_prefix_reservation( + self, request: LlmRequest, reservation_id: int, start: int, end: int + ) -> None: + """Forget the released part of a reservation. + + The pool holds no pin for it, so there is nothing to hand back: the + record only serves to check a confirmed load against the windows its + lookup verified. + """ + state = self._requests.get(request.request_id) + reservation = None if state is None else state.reservations.get(reservation_id) + if reservation is None: + return + if start <= reservation.start and end >= reservation.end: + del state.reservations[reservation_id] + elif end >= reservation.end: + reservation.end = min(reservation.end, start) + elif start <= reservation.start: + reservation.start = max(reservation.start, end) + + def _lookup( + self, request: LlmRequest, num_computed_tokens: int + ) -> Optional[Tuple[_RequestState, int, int, Dict[int, int]]]: + """The stored prefix after the local match: state, first block, blocks, windows. + + None when there is nothing to offer, including for a request whose pages + no key can name; such a request gets no state and so no saves either. + """ if self._config.capacity_only: # Neither a lookup nor a save will read the hash chain this would # otherwise build, which costs a SHA-256 over the whole prompt of # every request. - return 0, False + return None scope = _reuse_scope(request) if scope is None: # Leaving it without a `_RequestState` suppresses its saves too, # since `build_connector_meta` skips requests it has no state for. - return 0, False + return None tokens = request.get_tokens(0) state = self._state_for(request, tokens, scope) state.load_first_block = 0 state.load_blocks = 0 + state.window_present_from = {} if not self._config.role.loads: - return 0, False + return None # A partial local match means the boundary block is half computed on # device. Overwriting it with a stored page would discard tokens the # runtime already counted, so only whole-block offers are made. if num_computed_tokens % self._tokens_per_block: - return 0, False + return None first_block = num_computed_tokens // self._tokens_per_block # Stop one token short of the prompt: the runtime still has to run a @@ -189,20 +325,26 @@ def get_num_new_matched_tokens( last_block = (len(tokens) - 1) // self._tokens_per_block candidates = state.chain.hashes[first_block:last_block] if not candidates: - return 0, False - - hit_blocks = self._require_worker().count_prefix_hit(candidates) + return None + + window_present_from: Dict[int, int] = {} + if self._windowed_groups(): + hit = self._require_worker().find_prefix_hit( + candidates, first_block, self._window_rewind_tokens + ) + hit_blocks = hit.blocks + window_present_from = dict(hit.window_present_from) + else: + hit_blocks = self._require_worker().count_prefix_hit(candidates) if hit_blocks == 0: - return 0, False + return None - state.load_first_block = first_block - state.load_blocks = hit_blocks logger.debug( f"mooncake-store matched {hit_blocks} blocks " f"({hit_blocks * self._tokens_per_block} tokens) " f"for request {request.request_id}" ) - return hit_blocks * self._tokens_per_block, False + return state, first_block, hit_blocks, window_present_from def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): """No-op: page indices are read from the scheduler output instead. @@ -212,11 +354,20 @@ def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): `RequestData.new_block_ids_by_layer_group`. """ + def update_state_after_alloc_by_layer_group( + self, request: LlmRequest, block_ids_by_layer_group: List[List[int]] + ): + """No-op, as `update_state_after_alloc`: the scheduler output carries these.""" + # ---- work lists ---- def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeStoreMetadata: """Turn this iteration's scheduled requests into load and save lists.""" metadata = MooncakeStoreMetadata() + # Confirmed reservations, each listed once; an asynchronous one's + # request is parked outside the batch until its pages are reported. + for load in scheduler_output.prefix_loads: + metadata.prefix_loads.append(self._prefix_load_transfers(load)) for request_data in (*scheduler_output.new_requests, *scheduler_output.cached_requests): state = self._requests.get(request_data.request_id) if state is None: @@ -227,6 +378,11 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeSto state.tokens.extend(request_data.new_tokens) state.chain.extend(state.tokens) self._record_pages(state, request_data) + if state.resume_block is None: + # The first chunk starts where the local cache and the store left off. Unlike + # `computed_position`, this counts the store's tokens too. + position = int(getattr(state.request, "context_current_position", 0) or 0) + state.resume_block = position // self._tokens_per_block loads = self._loads_for(state, request_data) if loads.pages: @@ -264,6 +420,15 @@ def request_finished(self, request: LlmRequest, cache_block_ids: List[int]) -> b state = self._requests.pop(request.request_id, None) return bool(state is not None and state.emitted_saves) + def request_finished_by_layer_group( + self, request: LlmRequest, cache_block_ids_by_layer_group: List[List[int]] + ) -> bool: + """`request_finished` for a cache of several layer groups, a sliding window's. + + The answer does not depend on the page indices, so it is the same. + """ + return self.request_finished(request, []) + def request_reset(self, request: LlmRequest) -> None: """Forget a request whose pages were released for it to run again. @@ -279,12 +444,64 @@ def _require_worker(self) -> MooncakeStoreConnectorWorker: self._worker = resolve_local_worker() return self._worker + def _windowed_groups(self) -> Tuple[int, ...]: + """Sliding-window layer groups of the registered layout.""" + return tuple(getattr(self._require_worker(), "windowed_group_ids", ())) + + def _prefix_load_transfers(self, load: PrefixLoad) -> PrefixLoadTransfers: + """The pages a confirmed reservation fills, into the pages the runtime allocated. + + A sliding-window group has pages only for the blocks its window reads at + the load's end, and each of those has to be one the lookup found + stored, or the replay would read a page nothing wrote. + """ + state = self._requests.get(load.request_id) + reservation = None if state is None else state.reservations.get(load.reservation_id) + if reservation is None: + raise RuntimeError( + f"mooncake-store: prefix load {load.reservation_id} of request " + f"{load.request_id} has no reservation behind it" + ) + tokens_per_block = self._tokens_per_block + first, end = load.start // tokens_per_block, load.end // tokens_per_block + windowed = self._windowed_groups() + worker = self._require_worker() + transfers = RequestTransfers(load.request_id) + for layer_group_id, indices in enumerate(load.block_ids_by_layer_group): + is_window = layer_group_id in windowed + # Sinks sit below the window's tail and every lookup requires them. + sinks = worker.sink_blocks(layer_group_id) if is_window else 0 + present_from = reservation.window_present_from.get(layer_group_id, end) + for block in range(first, end): + page_index = indices[block] if block < len(indices) else BAD_PAGE_INDEX + if page_index == BAD_PAGE_INDEX and is_window: + # The window has passed this block, so nothing reads it. + continue + if page_index == BAD_PAGE_INDEX or (is_window and sinks <= block < present_from): + raise RuntimeError( + f"mooncake-store: request {load.request_id} loads blocks " + f"[{first}, {end}) and layer group {layer_group_id} would read " + f"block {block}, which the store was not found to hold. Loading " + "the rest would leave that page unwritten." + ) + transfers.pages.append( + PageTransfer(state.chain.hashes[block], layer_group_id, page_index) + ) + # The store already has what it supplies, so none of it is written back. + state.saved_upto = max(state.saved_upto, end) + return PrefixLoadTransfers(load.reservation_id, load.is_async, transfers) + def _state_for( self, request: LlmRequest, tokens: List[int], scope: ReuseScope ) -> _RequestState: state = self._requests.get(request.request_id) if state is None: - state = _RequestState(BlockHashChain(self._tokens_per_block, scope=scope)) + # Looked up before the request is scheduled, so `tokens` is its prompt. + state = _RequestState( + BlockHashChain(self._tokens_per_block, scope=scope), + request, + len(tokens) // self._tokens_per_block, + ) self._requests[request.request_id] = state # The lookup runs before the request is scheduled, so the chain has to # be ready before any scheduler output exists. @@ -309,6 +526,11 @@ def _addressable_blocks(self, state: _RequestState) -> int: def _loads_for(self, state: _RequestState, request_data: RequestData) -> RequestTransfers: transfers = RequestTransfers(request_data.request_id) + if state.load_blocks: + windowed = self._windowed_groups() + if windowed: + self._append_window_loads(state, transfers, windowed) + return transfers limit = self._addressable_blocks(state) for offset in range(state.load_blocks): block = state.load_first_block + offset @@ -322,17 +544,52 @@ def _loads_for(self, state: _RequestState, request_data: RequestData) -> Request def _saves_for(self, state: _RequestState, request_data: RequestData) -> RequestTransfers: transfers = RequestTransfers(request_data.request_id) + windowed = self._windowed_groups() limit = self._addressable_blocks(state) - for block in range(state.saved_upto, limit): - self._append_pages(state, transfers, block) + if self._prompt_blocks_only: + limit = min(limit, state.prompt_blocks) + request = state.request + chunk_size = getattr(request, "context_chunk_size", None) + if getattr(request, "is_context_init_state", False) and chunk_size is not None: + # The cache allocates a little past the chunk (a drafter's extra KV tokens), so a page can exist for a + # block only the next chunk computes. Save just the blocks this forward completes. + completed = int(request.context_current_position) + int(chunk_size) + limit = min(limit, completed // self._tokens_per_block) + save_window = bool(windowed) and not state.saved_window and self._finishes_prompt(state) + if save_window and self._window_first: + self._append_window_pages(state, transfers, windowed) + first = state.saved_upto + if self._resave_min_tokens and not state.resaved and self._finishes_prompt(state): + computed = len(state.tokens) - (state.resume_block or 0) * self._tokens_per_block + if computed >= self._resave_min_tokens: + # A long prefill wrote its first blocks minutes ago and never touched them again, so the pool can + # evict the head of the chain before the conversation's next turn arrives. Listing the whole prompt + # once more makes `_put` refresh what the pool still holds and rewrite what it evicted, from pages + # this request still holds. + first = 0 + state.resaved = True + for block in range(first, limit): + self._append_pages(state, transfers, block, skip=windowed) state.saved_upto = max(state.saved_upto, limit) + if save_window: + if not self._window_first: + self._append_window_pages(state, transfers, windowed) + state.saved_window = True return transfers - def _append_pages(self, state: _RequestState, transfers: RequestTransfers, block: int) -> None: - """Add one block's page from every layer group, or none of them.""" + def _append_pages( + self, + state: _RequestState, + transfers: RequestTransfers, + block: int, + skip: Tuple[int, ...] = (), + ) -> None: + """Add one block's page from every layer group but `skip`, or none of them.""" block_hash = state.chain.hashes[block] pages: List[PageTransfer] = [] for layer_group_id, indices in state.pages.items(): + if layer_group_id in skip: + continue page_index = indices[block] if page_index == BAD_PAGE_INDEX: # A sliding window dropped the block in this group, and a @@ -340,3 +597,77 @@ def _append_pages(self, state: _RequestState, transfers: RequestTransfers, block return pages.append(PageTransfer(block_hash, layer_group_id, page_index)) transfers.pages.extend(pages) + + @staticmethod + def _finishes_prompt(state: _RequestState) -> bool: + """Whether this iteration computes the last chunk of the request's prompt.""" + request = state.request + return bool(request.is_context_init_state and request.is_last_context_chunk) + + def _append_window_pages( + self, state: _RequestState, transfers: RequestTransfers, windowed: Tuple[int, ...] + ) -> None: + """Add the window pages a later turn resumes from: the ones this prompt ends with. + + 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 + further back (the rewind margin) may already hold another request's KV. + A lookup therefore resumes exactly at a stored prompt end, or through + the fallback at an earlier one. + """ + worker = self._require_worker() + tokens_per_block = self._tokens_per_block + chunk_start = int(state.request.context_current_position) + prompt_end = min(len(state.tokens) // tokens_per_block, len(state.chain.hashes)) + earliest_resume = prompt_end * tokens_per_block + for layer_group_id in windowed: + indices = state.pages.get(layer_group_id, []) + sinks = worker.sink_blocks(layer_group_id) + _, needed_from = worker.stale_blocks(layer_group_id, earliest_resume) + _, held_from = worker.stale_blocks(layer_group_id, chunk_start) + first = max(needed_from, held_from, sinks) + for block in (*range(min(sinks, prompt_end)), *range(first, prompt_end)): + page_index = indices[block] if block < len(indices) else BAD_PAGE_INDEX + if page_index == BAD_PAGE_INDEX: + continue + transfers.pages.append( + PageTransfer(state.chain.hashes[block], layer_group_id, page_index) + ) + + def _append_window_loads( + self, state: _RequestState, transfers: RequestTransfers, windowed: Tuple[int, ...] + ) -> None: + """Loads for a layout with sliding windows, up to where the runtime resumes. + + The runtime may honour less of the offer than was made. It resumes at + this chunk's start and gives a window group pages only for the blocks + it reads from there, so those are what is loaded, with every honoured + full-attention block. A window block the lookup did not find stored + means the replay would read a page nothing wrote, which is refused. + """ + worker = self._require_worker() + resume = int(state.request.context_current_position) // self._tokens_per_block + end = min(state.load_first_block + state.load_blocks, resume) + end = min(end, self._addressable_blocks(state)) + for layer_group_id, indices in state.pages.items(): + is_window = layer_group_id in windowed + # Sinks sit below the window's tail and every lookup requires them. + sinks = worker.sink_blocks(layer_group_id) if is_window else 0 + present_from = state.window_present_from.get(layer_group_id, end) + for block in range(state.load_first_block, end): + page_index = indices[block] + if page_index == BAD_PAGE_INDEX and is_window: + # The window has passed this block, so nothing reads it. + continue + if page_index == BAD_PAGE_INDEX or (is_window and sinks <= block < present_from): + raise RuntimeError( + f"mooncake-store: request {transfers.request_id} resumes at " + f"block {resume} of the {state.load_first_block + state.load_blocks} " + f"offered, and layer group {layer_group_id} would read block " + f"{block}, which the store was not found to hold. Loading the " + "rest would leave that page unwritten." + ) + transfers.pages.append( + PageTransfer(state.chain.hashes[block], layer_group_id, page_index) + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py index 917e8ca2c3de..336a960218bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -62,6 +62,27 @@ def validate_llm_args(llm_args: TorchLlmArgs) -> None: "is untested. Run with tensor parallelism only." ) + if getattr(llm_args.kv_cache_config, "enable_swa_scratch_reuse", False): + raise NotImplementedError( + "The mooncake-store connector does not support " + "kv_cache_config.enable_swa_scratch_reuse. The connector saves a " + "sliding window's pages at the end of a prompt from the page indices " + "the cache reported when it allocated them, and scratch reuse lets a " + "prefill write a window block into a page it does not keep." + ) + + spec_config = getattr(llm_args, "speculative_config", None) + recompute_tail = ( + getattr(spec_config, "context_recompute_tail", 0) if spec_config is not None else 0 + ) + if recompute_tail: + raise NotImplementedError( + "The mooncake-store connector does not support " + "speculative_config.context_recompute_tail. The runtime trims a reserved prefix to leave that tail for " + "the drafter after the lookup verified the sliding window at the full end, so the window the trimmed " + "load needs was never checked against the store. Leave context_recompute_tail at 0." + ) + # Only configs with an index-V cache declare the setting. sparse_config = llm_args.sparse_attention_config if sparse_config is not None and not getattr(sparse_config, "sparse_disable_index_value", True): @@ -103,6 +124,7 @@ def validate_node_budget( *, ranks_on_node: Optional[int] = None, available_bytes: Optional[int] = None, + load_threads: int = 1, ) -> None: """Reject a segment size this node cannot afford once every rank asks. @@ -118,6 +140,8 @@ def validate_node_budget( Defaults to the local MPI communicator's size. available_bytes: Host memory to measure against. Defaults to the kernel's own figure. + load_threads: Asynchronous loader threads per loading rank, each + pinning a staging pool of its own. """ from tensorrt_llm._utils import local_mpi_size @@ -131,9 +155,10 @@ def validate_node_budget( segment_claim = ranks_on_node * max(0, config.global_segment_size) # Staging is pinned for the process's lifetime and comes out of the same - # DRAM, so it belongs in the sum even when no segment is mounted. One pool - # per direction, and a capacity-only rank opens neither. - directions = int(config.role.loads) + int(config.role.saves) + # DRAM, so it belongs in the sum even when no segment is mounted. Loads get + # a pool for the executor and one per asynchronous loader thread, saves + # one, and a capacity-only rank opens none. + directions = (1 + load_threads) * int(config.role.loads) + int(config.role.saves) staging_claim = ranks_on_node * directions * MAX_STAGING_BUFFER_BYTES claimed = segment_claim + staging_claim @@ -187,17 +212,14 @@ def validate_node_budget( def validate_layout(layout: "KvCacheLayout") -> None: - """Reject KV cache geometries this connector cannot key correctly.""" - windowed = [group.layer_group_id for group in layout.groups if group.window_size is not None] - if windowed: - raise NotImplementedError( - "The mooncake-store connector does not support sliding-window " - f"attention (layer groups {windowed} declare a window size). A page's " - "validity then depends on where the window sits, which is a property " - "of the request that read it rather than of the tokens it holds, so " - "content-addressed reuse across instances is not sound." - ) + """Reject KV cache geometries this connector cannot key correctly. + Sliding-window groups are accepted. A page's bytes are fixed by the tokens + up to and including its block, which its key already chains, so a stored + page is never wrong for a request that shares the prefix. What a window + changes is which pages exist: the scheduler saves a window's pages only at + the end of each prompt, and a lookup resumes only where those are present. + """ if not layout.groups: raise ValueError( "The KV cache layout describes no layer groups, so there is nothing " diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index d27167d0936b..8ea9dd3105b7 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -19,8 +19,10 @@ knows how a page is addressed and how a key is spelled, which is why the leader, colocated with rank 0's worker, asks it to run prefix lookups. -Loads are synchronous: the runtime has already told the scheduler those tokens -are computed, so a failed load is a wrong answer rather than a slow one. +Loads must land before the request computes: the runtime has already counted +those tokens, so a failed load is a wrong answer rather than a slow one. A +confirmed prefix reservation loads on a loader thread while its request waits +outside the batch (`SYNC_LOAD_ENV` loads it in the executor's step instead). Saves are asynchronous and gated on a CUDA event, since the pages are complete only once the forward pass that wrote them has retired and blocking the @@ -28,18 +30,22 @@ scheduler reports such a request as saving asynchronously, which keeps its pages pinned until `get_finished` says the writes landed. A save that fails is logged and dropped: a write that did not land costs a future cache miss and -nothing else, so the request retires either way. Only the save thread failing -to start is fatal, because that would strand every request behind pages -nothing will release. +nothing else, so the request retires either way. Only a save or loader thread +failing to start is fatal, because that would strand every request behind pages +nothing will release or loads nothing will run. A capacity-only worker opens its handle and stops there, with no layout, no -buffer registration and no save thread, so a node can lend host memory to the +buffer registration and no save or loader threads, so a node can lend host memory to the pool without an HCA that can pin GPU pages. """ +import os import threading +import time import traceback from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field from queue import Queue from typing import Dict, List, Optional, Sequence, Set, Tuple @@ -48,6 +54,7 @@ from tensorrt_llm._utils import mpi_rank, mpi_world_size from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.logger import logger +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttnLifeCycle from ..kv_cache_connector import KvCacheConnectorWorker from ..kv_cache_layout import KvCacheLayout @@ -55,7 +62,7 @@ from .config import CONFIG_PATH_ENV, MooncakeStoreConnectorConfig, local_address, pool_config from .keys import KeyNamespace from .ledger import record_segment -from .metadata import MooncakeStoreMetadata, RequestTransfers +from .metadata import MooncakeStoreMetadata, PrefixLoadTransfers, RequestTransfers from .staging import ( MAX_STAGING_BUFFER_BYTES, HostStagingPool, @@ -67,7 +74,7 @@ from .staging import sync_stream as _sync_stream from .validation import validate_layout, validate_llm_args, validate_node_budget -__all__ = ["MooncakeStoreConnectorWorker", "resolve_local_worker"] +__all__ = ["MooncakeStoreConnectorWorker", "PrefixHit", "resolve_local_worker"] #: Set by the worker's constructor so the leader, which the executor builds in #: the same process on rank 0, can reach the store handle without a second @@ -76,11 +83,50 @@ _LOCAL_WORKER: Optional["MooncakeStoreConnectorWorker"] = None _LOCAL_WORKER_READY = threading.Event() -#: Seconds `shutdown` gives the save thread to drain its queue and return. +#: Seconds `shutdown` gives the save thread, and each loader thread, to drain its queue and return. #: Sized for a backlog of pages over a congested fabric rather than a healthy #: one, since what follows the wait is only safe once the thread has stopped. SAVE_DRAIN_TIMEOUT = 30.0 +#: Seconds between the loader, save and lookup statistics reports. +LOAD_STATS_INTERVAL = 30.0 + +#: Asynchronous prefix-load threads per loading rank, each with a staging pool +#: and stream of its own. +LOAD_THREADS_ENV = "TRTLLM_MOONCAKE_STORE_LOAD_THREADS" +DEFAULT_LOAD_THREADS = 2 + + +#: Set to 1 for reserved prefixes to load synchronously, in the executor's step that computes the rest of the prompt, +#: instead of on loader threads; read by the scheduler, and here to start no loader threads. Asynchronous loads (the +#: default) park the request until its pages are on the device while the others keep running, so a load costs that +#: request time instead of the batch. +SYNC_LOAD_ENV = "TRTLLM_MOONCAKE_STORE_SYNC_LOAD" + +#: 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" + +#: Keys per existence-probe call, 0 for one call. A lookup over a long prefix probes one key per block per rank, which +#: the master answers on one RPC thread while the scheduler waits; split, the calls run concurrently on the master's +#: thread pool. +PROBE_CHUNK_ENV = "TRTLLM_MOONCAKE_STORE_PROBE_CHUNK" +DEFAULT_PROBE_CHUNK = 4096 +#: Concurrent probe calls per lookup. +PROBE_THREADS = 4 + +#: Prompts that computed at least this many tokens list every full-attention block again at their last chunk, so the +#: save path refreshes the chain's head and rewrites what the pool evicted during the prefill. 0 turns it off. Read by +#: the scheduler. +RESAVE_MIN_TOKENS_ENV = "TRTLLM_MOONCAKE_STORE_RESAVE_MIN_TOKENS" + + +def _load_thread_count() -> int: + count = int(os.environ.get(LOAD_THREADS_ENV, DEFAULT_LOAD_THREADS)) + if count < 1: + raise ValueError(f"{LOAD_THREADS_ENV} must be at least 1, got {count}") + return count + def resolve_local_worker(timeout: float = 60.0) -> "MooncakeStoreConnectorWorker": """The worker living in this process, once it has been constructed. @@ -144,6 +190,26 @@ def _open_store(config: MooncakeStoreConnectorConfig): return store, hostname +def _interleaved_keys( + namespaces: Sequence[KeyNamespace], block_hashes: Sequence[bytes] +) -> List[str]: + """Every namespace's key of every block, block by block in namespace order. + + Built per namespace from its fixed prefix, then interleaved: a lookup over a long prompt builds one key per block + per rank on the scheduler's path, and calling `KeyNamespace.key` per key costs several times as much. + """ + if not namespaces or not block_hashes: + return [] + hexes = [block_hash.hex() for block_hash in block_hashes] + columns = [] + for namespace in namespaces: + prefix = namespace.prefix + "/" + columns.append([prefix + digest for digest in hexes]) + if len(columns) == 1: + return columns[0] + return [key for row in zip(*columns) for key in row] + + def _batched(items: Sequence, size: int): for start in range(0, len(items), size): yield items[start : start + size] @@ -160,6 +226,135 @@ def _stream_handle(stream: Optional[torch.cuda.Stream]) -> int: return int(stream.cuda_stream) +class _LoadStats: + """The loader thread's running totals, reported once per `LOAD_STATS_INTERVAL`.""" + + def __init__(self): + self._reset(time.monotonic()) + + def _reset(self, now: float) -> None: + self.since = now + self.loads = self.pages = self.loaded_bytes = 0 + self.busy = self.waited = self.max_wait = 0.0 + + def add(self, pages: int, loaded_bytes: int, wait: float, busy: float) -> Optional[str]: + """Count one load; past the interval, return the report and start a new one.""" + self.loads += 1 + self.pages += pages + self.loaded_bytes += loaded_bytes + self.busy += busy + self.waited += wait + self.max_wait = max(self.max_wait, wait) + now = time.monotonic() + elapsed = now - self.since + if elapsed < LOAD_STATS_INTERVAL: + return None + report = ( + f"{self.loads} loads, {self.pages} pages, {self.loaded_bytes / (1 << 30):.2f} GiB " + f"in {elapsed:.0f}s ({self.busy:.1f}s busy, " + f"{self.loaded_bytes / (1 << 30) / max(self.busy, 1e-9):.1f} GiB/s while busy); " + f"queue wait mean {self.waited / self.loads:.3f}s max {self.max_wait:.3f}s" + ) + self._reset(now) + return report + + +class _SaveStats: + """The save thread's running totals, reported once per `LOAD_STATS_INTERVAL`.""" + + def __init__(self): + self._reset(time.monotonic()) + + def _reset(self, now: float) -> None: + self.since = now + self.batches = self.written = self.present = self.written_bytes = 0 + self.busy = self.waited = self.max_wait = 0.0 + + def add( + self, written: int, present: int, written_bytes: int, wait: float, busy: float + ) -> Optional[str]: + """Count one queued save; past the interval, return the report and start a new one.""" + self.batches += 1 + self.written += written + self.present += present + self.written_bytes += written_bytes + self.busy += busy + self.waited += wait + self.max_wait = max(self.max_wait, wait) + now = time.monotonic() + elapsed = now - self.since + if elapsed < LOAD_STATS_INTERVAL: + return None + report = ( + f"{self.batches} batches, {self.written} pages written ({self.written_bytes / (1 << 30):.2f} GiB), " + f"{self.present} already stored, in {elapsed:.0f}s ({self.busy:.1f}s busy); " + f"queue wait mean {self.waited / self.batches:.3f}s max {self.max_wait:.3f}s" + ) + self._reset(now) + return report + + +class _LookupStats: + """Lookup outcomes on the leader, reported once per `LOAD_STATS_INTERVAL`. + + full: the window at the full-attention hit's end is stored. rewound: one within the rewind is. recovered: an + earlier end found by the fallback. window_miss: full-attention blocks but no usable window. fa_miss: not even + the first full-attention block. + """ + + OUTCOMES = ("full", "rewound", "recovered", "window_miss", "fa_miss") + + def __init__(self): + self._reset(time.monotonic()) + + def _reset(self, now: float) -> None: + self.since = now + self.counts = dict.fromkeys(self.OUTCOMES, 0) + self.hit_blocks = self.lost_blocks = 0 + self.probe_seconds = 0.0 + self.max_probe = 0.0 + + def add( + self, outcome: str, hit_blocks: int, lost_blocks: int, probe_seconds: float + ) -> Optional[str]: + self.counts[outcome] += 1 + self.hit_blocks += hit_blocks + self.lost_blocks += lost_blocks + self.probe_seconds += probe_seconds + self.max_probe = max(self.max_probe, probe_seconds) + now = time.monotonic() + elapsed = now - self.since + if elapsed < LOAD_STATS_INTERVAL: + return None + total = sum(self.counts.values()) + report = ( + f"{total} lookups (" + + ", ".join(f"{name} {count}" for name, count in self.counts.items()) + + f"), {self.hit_blocks} blocks hit, {self.lost_blocks} blocks behind the full-attention hit lost, " + f"probe mean {self.probe_seconds / max(1, total) * 1e3:.1f}ms max {self.max_probe * 1e3:.1f}ms " + f"in {elapsed:.0f}s" + ) + self._reset(now) + return report + + +@dataclass(frozen=True) +class PrefixHit: + """How far a request can resume from the store, and what backs it.""" + + #: Leading candidate blocks the store can replay. + blocks: int + #: Per sliding-window group, the first block from which the store holds + #: every page up to the hit's end. A replay the runtime cuts short reads its + #: window from further back, which is sound only down to this block. + window_present_from: Dict[int, int] = field(default_factory=dict) + + +def _live_blocks(sinks: int, stale_end: int, lo: int, hi: int) -> List[int]: + """Block ordinals in ``[lo, hi)`` a window still reads: its sinks and its tail.""" + return [*range(lo, min(sinks, hi)), *range(max(lo, sinks, stale_end), hi)] + + class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): """Moves KV pages between this rank's GPU cache and the Mooncake pool.""" @@ -183,11 +378,26 @@ def __init__(self, llm_args: TorchLlmArgs): # The same namespaces for every rank. A prefix is only reusable when all # shards of it are present, so a lookup has to ask about all of them. self._peer_namespaces: Dict[int, Tuple[KeyNamespace, ...]] = {} + # 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] = {} # Checked before the segment is mounted, since the kernel answers an # unaffordable total by killing the process rather than by failing the # allocation. - validate_node_budget(self._config) + # No loader threads, and no staging pools for them, when reserved prefixes load synchronously. + self._num_load_threads = ( + _load_thread_count() + if self._config.role.loads and os.environ.get(SYNC_LOAD_ENV, "0") != "1" + else 0 + ) + # A sliding-window layout is served through prefix reservations only; see register_kv_cache_layout. + self._prefix_aware_scheduling = bool( + getattr( + getattr(llm_args, "scheduler_config", None), "enable_prefix_aware_scheduling", True + ) + ) + validate_node_budget(self._config, load_threads=self._num_load_threads) self._store, self._segment_host = _open_store(self._config) record_segment( @@ -224,6 +434,29 @@ def __init__(self, llm_args: TorchLlmArgs): # `register_kv_cache_layout`. A failure inside the loop is dropped # there instead. self._save_startup_error: Optional[BaseException] = None + # Asynchronous prefix loads run on threads of their own, each with its + # own staging slots and stream, while their requests wait outside the + # batch. The threads share one queue. + self._load_queue: "Queue[Optional[Tuple[int, RequestTransfers, Optional[torch.cuda.Event], float]]]" = Queue() + self._load_threads: List[threading.Thread] = [] + self._load_startup_error: Optional[BaseException] = None + self._async_load_stagings: List[HostStagingPool] = [] + self._window_fallback = os.environ.get(WINDOW_FALLBACK_ENV, "1") != "0" + self._lookup_stats = _LookupStats() + self._probe_chunk = int(os.environ.get(PROBE_CHUNK_ENV, DEFAULT_PROBE_CHUNK)) + # Lookups run on the leader's scheduler, so only a loading role probes. + self._probe_pool: Optional[ThreadPoolExecutor] = ( + ThreadPoolExecutor( + max_workers=PROBE_THREADS, thread_name_prefix=f"mooncake-store-probe-{self._rank}" + ) + if self._config.role.loads and self._probe_chunk > 0 + else None + ) + self._load_lock = threading.Lock() + # Reservations whose pages landed, until `get_finished_prefix_loads`. + self._finished_prefix_loads: List[int] = [] + # The loader thread's first failure, raised on the executor thread. + self._load_error: Optional[BaseException] = None global _LOCAL_WORKER _LOCAL_WORKER = self @@ -306,6 +539,19 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: self._namespace(rank, layer_group_id, bytes_per_page) for rank in range(self._world_size) ) + for group in layout.groups: + if group.window_size is not None: + self._life_cycles[group.layer_group_id] = AttnLifeCycle.make( + group.window_size, group.num_sink_tokens, layout.tokens_per_block + ) + 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 " + "reservations, which scheduler_config.enable_prefix_aware_scheduling=False turns off. Without " + "them the runtime can honour a stored prefix only up to a point whose window the lookup never " + "verified (its fallback when the full offer does not fit), and the load would then read a " + "window page nothing wrote. Leave enable_prefix_aware_scheduling on." + ) if self._config.role.saves: self._save_thread = threading.Thread( @@ -326,6 +572,33 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: "could never write to the pool." ) from startup_error + if self._config.role.loads: + started = [] + for index in range(self._num_load_threads): + # Set once the thread has finished its device and stream setup, + # whether or not that succeeded. + event = threading.Event() + thread = threading.Thread( + target=self._drain_loads, + args=(index, event), + name=f"mooncake-store-load-{self._rank}-{index}", + daemon=True, + ) + thread.start() + self._load_threads.append(thread) + started.append(event) + # As for the save thread: a loader that cannot bind this rank's + # device would leave every parked request waiting for good. + for event in started: + event.wait() + with self._load_lock: + startup_error, self._load_startup_error = self._load_startup_error, None + if startup_error is not None: + raise RuntimeError( + f"mooncake-store: the load thread for rank {self._rank} " + "could not start, so asynchronous prefix loads would never land." + ) from startup_error + logger.warning( f"mooncake-store worker rank {self._rank} registered layout: {addressing.describe()}" ) @@ -352,6 +625,17 @@ def _open_staging(self, addressing: PageAddressing) -> None: store=self._store, label="load", ) + # Each loader thread's own slots, so none waits on another's or the + # executor's. + self._async_load_stagings = [ + HostStagingPool( + slot_bytes=slot_bytes, + num_slots=num_slots, + store=self._store, + label=f"async-load-{index}", + ) + for index in range(self._num_load_threads) + ] if self._config.role.saves: self._save_staging = HostStagingPool( slot_bytes=slot_bytes, @@ -392,6 +676,35 @@ def is_registered(self) -> bool: """Whether a KV cache layout has been registered yet.""" return self._addressing is not None + @property + def windowed_group_ids(self) -> Tuple[int, ...]: + """Layer groups with a sliding window, in layout order. + + Such a group keeps only the tail of a sequence, so the scheduler saves + its pages where prompts end and a lookup has to land on one of those. + """ + return tuple(self._life_cycles) + + def sink_blocks(self, layer_group_id: int) -> int: + """Leading blocks a sliding-window group keeps however far it moves.""" + life_cycle = self._life_cycles.get(layer_group_id) + return 0 if life_cycle is None else int(life_cycle.num_sink_blocks) + + def stale_blocks(self, layer_group_id: int, history_length: int) -> Tuple[int, int]: + """Block ordinals ``[start, end)`` a group no longer reads at ``history_length``. + + Answered by the cache's own life cycle, as + `KVCacheManagerV2._stale_block_range` is, so the blocks a lookup asks + for are the ones the runtime gives pages to. Empty for full attention. + """ + life_cycle = self._life_cycles.get(layer_group_id) + if life_cycle is None: + return 0, 0 + stale = life_cycle.get_stale_range(history_length, self._addressing.tokens_per_block) + # Read `beg`/`end` rather than unpacking: only the pure-Python + # `HalfOpenRange` is a tuple subclass, the bound one is not. + return int(stale.beg), int(stale.end) + def count_prefix_hit(self, block_hashes: Sequence[bytes]) -> int: """How many leading blocks of `block_hashes` are fully present. @@ -407,53 +720,394 @@ def count_prefix_hit(self, block_hashes: Sequence[bytes]) -> int: """ if not block_hashes or self._addressing is None: return 0 + return self._count_leading(block_hashes, tuple(self._peer_namespaces)) + + def find_prefix_hit( + self, block_hashes: Sequence[bytes], first_block: int, rewind_tokens: int = 0 + ) -> PrefixHit: + """The longest stored prefix a request can resume from, windows included. + + Full-attention groups need every block of the prefix. A sliding-window + group needs only the blocks it still reads where the prefix ends, and + the store holds those only near the ends of earlier prompts, so the + search starts at the longest full-attention hit and walks back up to + `rewind_tokens`, taking the first prompt end whose window is complete. + With no window this is `count_prefix_hit`. + + When no end within the rewind has its window, the full-attention blocks + can still reach past the last stored window (a save that has not landed + yet, or a window evicted on its own). Rather than give up the whole + prefix, the search then probes the window pages further back and + resumes at the latest earlier end whose window is complete. + + Args: + block_hashes: Candidate hashes in block order. + first_block: Block ordinal of `block_hashes[0]`. Everything before it + is on device already, window included. + rewind_tokens: How far back from the full-attention hit to look for + a prompt end whose window is stored. + + Returns: + The usable prefix in blocks, and for each sliding-window group the + first block from which the store holds every page up to its end. + """ + started = time.monotonic() + hit, outcome, fa_end = self._find_prefix_hit(block_hashes, first_block, rewind_tokens) + if self._life_cycles and block_hashes and self._addressing is not None: + report = self._lookup_stats.add( + outcome, + hit.blocks, + max(0, fa_end - first_block - hit.blocks), + time.monotonic() - started, + ) + if report: + logger.info(f"mooncake-store rank {self._rank} lookups: {report}") + return hit + + def _find_prefix_hit( + self, block_hashes: Sequence[bytes], first_block: int, rewind_tokens: int + ) -> Tuple[PrefixHit, str, int]: + """`find_prefix_hit`, which of `_LookupStats.OUTCOMES` it came to, and where the full-attention hit ends.""" + if not block_hashes or self._addressing is None: + return PrefixHit(0), "fa_miss", first_block + full_groups = tuple(g for g in self._peer_namespaces if g not in self._life_cycles) + hit = self._count_leading(block_hashes, full_groups) if full_groups else len(block_hashes) + if not self._life_cycles or hit == 0: + return PrefixHit(hit), ("full" if hit else "fa_miss"), first_block + hit + + tokens_per_block = self._addressing.tokens_per_block + end = first_block + hit + # The search for a stored prompt end goes back at most `rewind_tokens`; + # past that, only the fallback looks. + lowest_end = max(first_block + 1, end - -(-rewind_tokens // tokens_per_block)) + + # Every block some candidate end reads, probed in one call: the sinks, + # and the tail from where the lowest end's window starts. + probed: Dict[int, List[int]] = {} + for layer_group_id in self._life_cycles: + _, stale_end = self.stale_blocks(layer_group_id, lowest_end * tokens_per_block) + probed[layer_group_id] = _live_blocks( + self.sink_blocks(layer_group_id), stale_end, first_block, end + ) + stored = self._probe_window_blocks(block_hashes, first_block, probed) + if stored is None: + return PrefixHit(0), "window_miss", end + + for candidate in range(end, lowest_end - 1, -1): + if self._window_complete(stored, first_block, candidate): + outcome = "full" if candidate == end else "rewound" + hit_at = PrefixHit( + candidate - first_block, self._present_from(stored, first_block, candidate) + ) + return hit_at, outcome, end + + if not self._window_fallback or lowest_end <= first_block + 1: + return PrefixHit(0), "window_miss", end + + # The fallback: probe each window group's blocks below what the first + # call covered, then take the latest earlier end with a complete window. + below: Dict[int, List[int]] = {} + for layer_group_id, blocks in probed.items(): + low = max(first_block, self.sink_blocks(layer_group_id)) + first_probed = next((b for b in blocks if b >= low), end) + below[layer_group_id] = list(range(low, first_probed)) + more = self._probe_window_blocks(block_hashes, first_block, below) + if more is None: + return PrefixHit(0), "window_miss", end + for layer_group_id, blocks in more.items(): + stored[layer_group_id] |= blocks + # Missing-block prefix sums per group over [first_block, end), so each + # candidate's window is checked in constant time. + missing: Dict[int, List[int]] = {} + for layer_group_id, blocks in stored.items(): + sums = [0] + for block in range(first_block, end): + sums.append(sums[-1] + (block not in blocks)) + missing[layer_group_id] = sums + for candidate in range(lowest_end - 1, first_block, -1): + history_length = candidate * tokens_per_block + complete = True + for layer_group_id, sums in missing.items(): + sinks = self.sink_blocks(layer_group_id) + tail_from = min( + candidate, + max(first_block, sinks, self.stale_blocks(layer_group_id, history_length)[1]), + ) + sink_hi = max(first_block, min(sinks, candidate)) + if ( + sums[sink_hi - first_block] + or sums[candidate - first_block] - sums[tail_from - first_block] + ): + complete = False + break + if complete: + return ( + PrefixHit( + candidate - first_block, self._present_from(stored, first_block, candidate) + ), + "recovered", + end, + ) + return PrefixHit(0), "window_miss", end + def _probe_window_blocks( + self, block_hashes: Sequence[bytes], first_block: int, blocks_by_group: Dict[int, List[int]] + ) -> Optional[Dict[int, Set[int]]]: + """Which of the given window-group blocks every rank holds, or None if the probe failed.""" keys: List[str] = [] - for block_hash in block_hashes: - for namespaces in self._peer_namespaces.values(): - keys.extend(namespace.key(block_hash) for namespace in namespaces) - keys_per_block = len(keys) // len(block_hashes) + for layer_group_id, blocks in blocks_by_group.items(): + keys.extend( + _interleaved_keys( + self._peer_namespaces[layer_group_id], + [block_hashes[block - first_block] for block in blocks], + ) + ) + present = self._probe(keys) if keys else [] + if present is None: + return None + stored: Dict[int, Set[int]] = {} + index = 0 + for layer_group_id, blocks in blocks_by_group.items(): + stored[layer_group_id] = set() + for block in blocks: + statuses = present[index : index + self._world_size] + index += self._world_size + if all(status == 1 for status in statuses): + stored[layer_group_id].add(block) + return stored + + def _window_complete( + self, stored: Dict[int, Set[int]], first_block: int, candidate: int + ) -> bool: + """Whether every window group holds the blocks it reads at `candidate`.""" + history_length = candidate * self._addressing.tokens_per_block + return all( + set( + _live_blocks( + self.sink_blocks(layer_group_id), + self.stale_blocks(layer_group_id, history_length)[1], + first_block, + candidate, + ) + ) + <= stored[layer_group_id] + for layer_group_id in self._life_cycles + ) + + @staticmethod + def _present_from( + stored: Dict[int, Set[int]], first_block: int, candidate: int + ) -> Dict[int, int]: + """Per window group, the first block of the stored run that ends at `candidate`.""" + present_from = {} + for layer_group_id, blocks in stored.items(): + block = candidate + while block - 1 >= first_block and block - 1 in blocks: + block -= 1 + present_from[layer_group_id] = block + return present_from + + def _count_leading(self, block_hashes: Sequence[bytes], group_ids: Sequence[int]) -> int: + """Leading blocks whose pages every rank holds in each of `group_ids`.""" + namespaces = [ + namespace + for layer_group_id in group_ids + for namespace in self._peer_namespaces[layer_group_id] + ] + keys = _interleaved_keys(namespaces, block_hashes) + keys_per_block = len(namespaces) + + present = self._probe(keys) + if present is None: + return 0 + + # Mooncake reports 1 for present, 0 for absent and a negative value for + # a failed probe. Anything but a definite 1 is treated as a miss, and + # the hit ends at the block holding the first one. + first_miss = next( + (index for index, status in enumerate(present) if status != 1), len(present) + ) + return first_miss // keys_per_block + + def _probe(self, keys: List[str]) -> Optional[List[int]]: + """`batch_is_exist` over `keys`, or None when the probe itself failed. + Long key lists go out as concurrent chunks of `PROBE_CHUNK_ENV` keys, answered in order. + """ + chunk = self._probe_chunk + if self._probe_pool is None or chunk <= 0 or len(keys) <= chunk: + chunks = [keys] + else: + chunks = [keys[start : start + chunk] for start in range(0, len(keys), chunk)] try: - present = self._store.batch_is_exist(keys) + if len(chunks) == 1: + answers = [self._store.batch_is_exist(chunks[0])] + else: + futures = [ + self._probe_pool.submit(self._store.batch_is_exist, part) for part in chunks + ] + answers = [future.result() for future in futures] except Exception as exc: logger.warning( f"mooncake-store lookup failed; treating as a miss: " f"{type(exc).__name__}: {exc}\n{traceback.format_exc()}" ) - return 0 - - if len(present) != len(keys): - logger.warning( - f"mooncake-store batch_is_exist returned {len(present)} results for " - f"{len(keys)} keys; treating as a miss" - ) - return 0 + return None - hit_blocks = 0 - for index in range(len(block_hashes)): - window = present[index * keys_per_block : (index + 1) * keys_per_block] - # Mooncake reports 1 for present, 0 for absent and a negative value - # for a failed probe. Anything but a definite 1 is treated as a miss. - if not all(status == 1 for status in window): - break - hit_blocks += 1 - return hit_blocks + present: List[int] = [] + for part, answer in zip(chunks, answers): + if len(answer) != len(part): + logger.warning( + f"mooncake-store batch_is_exist returned {len(answer)} results for " + f"{len(part)} keys; treating as a miss" + ) + return None + present.extend(answer) + return present # ---- load path ---- def start_load_kv(self, stream: torch.cuda.Stream): - """Pull every scheduled page into its GPU slot before the forward pass.""" + """Pull this pass's pages into their GPU slots, or hand them to the loader thread. + + Synchronous loads land before the forward pass. An asynchronous prefix + load's request is parked outside the batch, so its pages go to the + loader thread and are reported through `get_finished_prefix_loads`. + """ metadata: Optional[MooncakeStoreMetadata] = self.get_connector_meta() - if metadata is None or not metadata.loads: + if metadata is None: return + self._raise_load_error() - keys, addresses, sizes, total_pages = self._resolve(metadata.loads) - if not keys: + now: List[RequestTransfers] = list(metadata.loads) + finished: List[int] = [] + for prefix_load in getattr(metadata, "prefix_loads", ()): + if prefix_load.is_async: + self._queue_async_load(prefix_load, stream) + else: + now.append(prefix_load.transfers) + finished.append(prefix_load.reservation_id) + if now: + staging = self._load_staging + self._load_pages(now, staging, _stream_handle(stream) if staging is not None else 0) + if finished: + with self._load_lock: + self._finished_prefix_loads.extend(finished) + + def get_finished_prefix_loads(self) -> List[int]: + """Reservations whose pages are on this rank's device since the last call.""" + self._raise_load_error() + with self._load_lock: + finished, self._finished_prefix_loads = self._finished_prefix_loads, [] + return finished + + def _queue_async_load(self, prefix_load: PrefixLoadTransfers, stream) -> None: + """Hand a parked request's pages to the loader thread. + + The pages were just allocated, and their previous owner's last forward + may still be queued on the executor's stream, so the loader waits for + an event recorded there before it writes them. + """ + if not prefix_load.transfers.pages: + with self._load_lock: + self._finished_prefix_loads.append(prefix_load.reservation_id) return + ready = None + if torch.cuda.is_available(): + ready = torch.cuda.Event() + ready.record(stream) + self._load_queue.put( + (prefix_load.reservation_id, prefix_load.transfers, ready, time.monotonic()) + ) - staging = self._load_staging - handle = _stream_handle(stream) if staging is not None else 0 + def _raise_load_error(self) -> None: + """Surface a loader-thread failure on the executor thread. + + A parked request whose pages did not all land can neither run nor be + reported done, and under tensor parallelism the ranks cannot agree to + restart it, so the failure is fatal, as on the synchronous path. + """ + with self._load_lock: + error = self._load_error + if error is not None: + raise RuntimeError( + f"mooncake-store: an asynchronous prefix load failed on rank {self._rank}" + ) from error + + def _drain_loads(self, index: int, started: threading.Event) -> None: + """A loader thread: one parked request at a time, on its own stream and staging.""" + try: + if self._device_index is not None: + # A new thread starts on device 0; see `_drain_saves`. + torch.cuda.set_device(self._device_index) + staging = self._async_load_stagings[index] if self._async_load_stagings else None + stream = ( + torch.cuda.Stream() if staging is not None and torch.cuda.is_available() else None + ) + handle = _stream_handle(stream) + except Exception as exc: + logger.error( + f"mooncake-store load thread failed to start on rank {self._rank}: " + f"{type(exc).__name__}: {exc}" + ) + with self._load_lock: + if self._load_startup_error is None: + self._load_startup_error = exc + return + finally: + started.set() + stats = _LoadStats() + while True: + item = self._load_queue.get() + if item is None: + return + reservation_id, transfers, ready, queued_at = item + began = time.monotonic() + try: + if ready is not None: + if stream is not None: + # The store fills host slots and may run ahead; only + # the scatters onto the pages wait for their release. + stream.wait_event(ready) + else: + ready.synchronize() + loaded_bytes = self._load_pages([transfers], staging, handle) + except Exception as exc: + # Thread boundary: the executor thread raises it on its next + # call, rather than leaving the request parked forever. + logger.error( + f"mooncake-store rank {self._rank} failed the asynchronous load " + f"of reservation {reservation_id}: {type(exc).__name__}: {exc}" + ) + with self._load_lock: + if self._load_error is None: + self._load_error = exc + continue + with self._load_lock: + self._finished_prefix_loads.append(reservation_id) + report = stats.add( + len(transfers.pages), loaded_bytes, began - queued_at, time.monotonic() - began + ) + if report: + logger.info( + f"mooncake-store rank {self._rank} loader {index}: {report}, " + f"{self._load_queue.qsize()} queued" + ) + + def _load_pages( + self, transfers: Sequence[RequestTransfers], staging: Optional[HostStagingPool], handle: int + ) -> int: + """Fetch pages batch by batch and land them on the device before returning. + + Shared by the executor thread (synchronous loads, on its stream) and the + loader thread (on its own stream and staging pool). + + Returns: + The bytes loaded. + """ + keys, addresses, sizes, total_pages = self._resolve(transfers) + if not keys: + return 0 for batch in zip( _batched(keys, self._batch_size), @@ -491,6 +1145,7 @@ def start_load_kv(self, stream: torch.cuda.Stream): _sync_stream(handle) logger.debug(f"mooncake-store rank {self._rank} loaded {total_pages} pages") + return sum(sum(page_sizes) for page_sizes in sizes) def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream): """No-op: loads complete in `start_load_kv`. @@ -511,7 +1166,7 @@ def save_kv_layer(self, layer_idx: int, stream: torch.cuda.Stream): def wait_for_save(self, stream: torch.cuda.Stream): """Hand this pass's saves to the background thread, gated on an event.""" metadata: Optional[MooncakeStoreMetadata] = self.get_connector_meta() - if metadata is None or not metadata.saves or not self._config.role.saves: + if metadata is None or not self._config.role.saves or not metadata.saves: return # The pages are written by kernels still queued on this stream, so the @@ -523,7 +1178,7 @@ def wait_for_save(self, stream: torch.cuda.Stream): with self._save_lock: for transfers in metadata.saves: self._outstanding_saves[transfers.request_id] += 1 - self._save_queue.put((event, list(metadata.saves))) + self._save_queue.put((event, list(metadata.saves), time.monotonic())) def get_finished( self, finished_gen_req_ids: List[int], started_loading_req_ids: List[int] @@ -580,14 +1235,17 @@ def _drain_saves(self) -> None: return finally: self._save_started.set() + stats = _SaveStats() while True: item = self._save_queue.get() if item is None: return - event, transfers = item + event, transfers, queued_at = item + began = time.monotonic() + written = present = written_bytes = 0 try: event.synchronize() - self._put(transfers) + written, present, written_bytes = self._put(transfers) 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 @@ -604,11 +1262,24 @@ def _drain_saves(self) -> None: self._outstanding_saves.pop(entry.request_id, None) else: self._outstanding_saves[entry.request_id] = remaining + report = stats.add( + written, present, written_bytes, began - queued_at, time.monotonic() - began + ) + if report: + logger.info( + f"mooncake-store rank {self._rank} saves: {report}, {self._save_queue.qsize()} queued" + ) + + def _put(self, transfers: Sequence[RequestTransfers]) -> Tuple[int, int, int]: + """Write the pages the pool does not hold yet. - def _put(self, transfers: Sequence[RequestTransfers]) -> None: + Returns: + Pages written, pages already stored, and the bytes written. + """ keys, addresses, sizes, _ = self._resolve(transfers) + written = present_count = written_bytes = 0 if not keys: - return + return written, present_count, written_bytes staging = self._save_staging handle = _stream_handle(self._save_stream) if staging is not None else 0 @@ -628,6 +1299,7 @@ def _put(self, transfers: Sequence[RequestTransfers]) -> None: for index, status in enumerate(present) if status != 1 # absent, or a failed probe we retry as a write ] + present_count += len(batch_keys) - len(pending) if not pending: continue source_addresses = [batch_addresses[i] for i in pending] @@ -646,6 +1318,12 @@ def _put(self, transfers: Sequence[RequestTransfers]) -> None: source_sizes, ) failures = sum(1 for result in results if not isinstance(result, int) or result < 0) + written += len(pending) - failures + written_bytes += sum( + sum(source_sizes[i]) + for i, result in enumerate(results) + if isinstance(result, int) and result >= 0 + ) if failures: # A dropped write only costs a future cache miss, which is not # worth failing a request that already answered. @@ -653,6 +1331,7 @@ def _put(self, transfers: Sequence[RequestTransfers]) -> None: f"mooncake-store rank {self._rank} failed to save {failures} of " f"{len(pending)} pages" ) + return written, present_count, written_bytes # ---- shared ---- @@ -684,10 +1363,16 @@ def _resolve( return keys, addresses, sizes, pages def shutdown(self) -> None: - """Stop the save thread, then release what it was reading. Idempotent.""" - thread = self._save_thread - if thread is not None: - self._save_queue.put(None) + """Stop the save and load threads, then release what they used. Idempotent.""" + threads = ( + [("save thread", self._save_thread, self._save_queue)] if self._save_thread else [] + ) + threads += [("load thread", thread, self._load_queue) for thread in self._load_threads] + # One sentinel per thread: a loader stops at the first one it takes, + # after the loads queued ahead of it. + for _, _, queue in threads: + queue.put(None) + for label, thread, _ in threads: thread.join(timeout=SAVE_DRAIN_TIMEOUT) if thread.is_alive(): # The thread reads the KV pools through the store handle and @@ -696,13 +1381,18 @@ def shutdown(self) -> None: # end: the process is going down anyway, the thread is a daemon # and does not hold it open, and a later call retries the join. logger.error( - f"mooncake-store rank {self._rank}: the save thread did not stop " - f"within {SAVE_DRAIN_TIMEOUT:g}s, so the store handle and its " + f"mooncake-store rank {self._rank}: the {label} " + f"did not stop within {SAVE_DRAIN_TIMEOUT:g}s, so the store handle and its " "staging buffers are left in place rather than freed under a " "transfer still reading them." ) return - self._save_thread = None + self._save_thread = None + self._load_threads = [] + if self._probe_pool is not None: + # Lookups run on the scheduler thread, which is done by now. + self._probe_pool.shutdown(wait=True) + self._probe_pool = None store, self._store = self._store, None if store is not None: try: @@ -715,6 +1405,7 @@ def shutdown(self) -> None: # Released only after the store is closed, since it holds registrations # against this memory. self._load_staging = None + self._async_load_stagings = [] self._save_staging = None self._save_stream = None global _LOCAL_WORKER diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9711f72f6f11..f4128a53af79 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1171,7 +1171,14 @@ def _maybe_init_kv_connector_manager(self): is_kv_cache_manager_v2 = isinstance(self.kv_cache_manager, KVCacheManagerV2) - if is_kv_cache_manager_v2 and self.max_draft_len > 0: + from .connectors.registry import uses_connector + + # mooncake-store saves only prompt blocks under speculative + # decoding, which are filled before the first draft. + if (is_kv_cache_manager_v2 and self.max_draft_len > 0 + and not uses_connector( + getattr(self.llm_args, "kv_connector_config", None), + "mooncake-store")): raise NotImplementedError( "KV Cache Connector is not supported with speculative " "decoding. Rejected draft tokens " diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 095424b7f37a..5ba8385c00e7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -1019,9 +1019,12 @@ def allocation_scope(current_stage: ExecutorMemoryType): # Rejected draft tokens shrink a request's page list, and the freed slot # goes to whichever request allocates next. The connector is only told # about pages appended since the last report, so it would keep - # addressing a slot another request now owns. + # addressing a slot another request now owns. mooncake-store saves only + # prompt blocks under speculative decoding, which are filled before the + # first draft, so it never does. if (spec_config is not None and spec_config.max_draft_len > 0 - and v2_selection is True): + and v2_selection is True + and not uses_connector(kv_connector_config, "mooncake-store")): raise NotImplementedError( "KV connector is not supported with speculative decoding. " "Disable speculative decoding to run a connector.") diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index d7bae75b56af..4af7e53982d6 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -35,6 +35,7 @@ import torch from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( + PrefixLoad, RequestData, SchedulerOutput, ) @@ -44,19 +45,29 @@ KvCacheLayout, KvCacheRegion, ) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import scheduler as scheduler_module from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import staging as staging_module from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import worker as worker_module from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.addressing import PageAddressing from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.keys import BlockHashChain from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.metadata import ( + MooncakeStoreMetadata, PageTransfer, + PrefixLoadTransfers, RequestTransfers, ) from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.scheduler import ( MooncakeStoreConnectorScheduler, ) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.validation import ( + NODE_BUDGET_RESERVE_BYTES, + validate_layout, + validate_llm_args, + validate_node_budget, +) from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.worker import ( MooncakeStoreConnectorWorker, + PrefixHit, ) from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX @@ -188,7 +199,7 @@ def fake_store(monkeypatch): monkeypatch.setattr(worker_module, "_open_store", lambda _config: (store, "10.0.0.1")) # The node budget is a property of the machine the test happens to run on, # which is not what any of these tests are about. - monkeypatch.setattr(worker_module, "validate_node_budget", lambda _config: None) + monkeypatch.setattr(worker_module, "validate_node_budget", lambda _config, **_kwargs: None) yield store for worker in store.workers: worker.shutdown() @@ -368,6 +379,7 @@ def refuse(_keys): ( SimpleNamespace(synchronize=lambda: None), [RequestTransfers(42, [PageTransfer(b"\x01" * 16, 0, 1)])], + time.monotonic(), ) ) @@ -441,7 +453,9 @@ def fake_cuda(monkeypatch): monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "current_device", lambda: 3) monkeypatch.setattr(torch.cuda, "set_device", lambda index: recorded.append(index)) - monkeypatch.setattr(torch.cuda, "Stream", lambda *_args, **_kwargs: SimpleNamespace()) + monkeypatch.setattr( + torch.cuda, "Stream", lambda *_args, **_kwargs: SimpleNamespace(cuda_stream=0) + ) monkeypatch.setattr(worker_module, "HostStagingPool", StubStagingPool) return recorded @@ -561,6 +575,30 @@ def refuse(_index): pytest.fail("registration should not have completed") +def test_a_load_thread_that_cannot_start_fails_registration( + store_config, fake_store, fake_cuda, monkeypatch +): + """A parked request waits on the load thread, so one that died must stop bringup.""" + calls = [] + + def refuse_after_the_save_thread(index): + calls.append(index) + # The save thread binds the device first and registration waits for it. + if len(calls) > 1: + raise RuntimeError("cudaSetDevice failed") + + monkeypatch.setattr(torch.cuda, "set_device", refuse_after_the_save_thread) + + worker = MooncakeStoreConnectorWorker(make_llm_args()) + fake_store.workers.append(worker) + try: + with pytest.raises(RuntimeError, match="load thread"): + worker.register_kv_cache_layout(make_layout()) + finally: + # The save thread did start, so it has to be stopped. + worker.shutdown() + + def test_shutdown_leaves_the_store_open_under_a_save_still_reading( store_config, fake_store, staged_copies, monkeypatch ): @@ -589,6 +627,7 @@ def blocking_put(*args, **kwargs): ( SimpleNamespace(synchronize=lambda: None), [RequestTransfers(1, [PageTransfer(b"\x05" * 16, 0, 0)])], + time.monotonic(), ) ) assert reading.wait(30.0), "the save thread never reached the store" @@ -899,3 +938,716 @@ def test_scheduler_bypasses_multimodal_requests_it_cannot_identify(store_config) assert metadata.saves == [] # Nothing was handed to a save thread, so nothing pins its pages either. assert scheduler.request_finished(request, [7, 8, 9]) is False + + +# ---- sliding windows ---- + +#: The window of the sliding group in these tests: two blocks. +WINDOW = 2 * TOKENS_PER_BLOCK + + +def make_window_layout(*, num_slots=16): + """Group 0 is full attention, group 1 a sliding window of `WINDOW` tokens.""" + full = make_layout(num_slots=num_slots).groups[0] + windowed = KvCacheLayerGroupLayout( + layer_group_id=1, + layer_ids=(1,), + window_size=WINDOW, + regions=( + KvCacheRegion( + base=0x100000, + size=128, + stride=128, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=1, role="key"),), + ), + ), + ) + return KvCacheLayout(tokens_per_block=TOKENS_PER_BLOCK, groups=(full, windowed)) + + +def store_window(fake_store, worker, hashes, end_block): + """Store the window pages a prefix of `end_block` blocks ends with.""" + _, stale_end = worker.stale_blocks(1, end_block * TOKENS_PER_BLOCK) + for block in range(stale_end, end_block): + fake_store.objects.add(worker._namespaces[1].key(hashes[block])) + return stale_end + + +def test_validate_layout_accepts_a_sliding_window(): + validate_layout(make_window_layout()) + + +def reference_chain(tokens, tokens_per_block): + """The chain under the per-token `int.to_bytes` encoding, which the packed fast path must match byte for byte.""" + import hashlib + + seed = BlockHashChain(tokens_per_block)._seed + hashes = [] + for ordinal in range(len(tokens) // tokens_per_block): + block = tokens[ordinal * tokens_per_block : (ordinal + 1) * tokens_per_block] + hasher = hashlib.blake2b(digest_size=16) + hasher.update(hashes[-1] if hashes else seed) + hasher.update(b"".join(int(t).to_bytes(8, "little", signed=True) for t in block)) + hashes.append(hasher.digest()) + return hashes + + +def test_hash_chain_matches_the_per_token_encoding_across_extends(): + import random + + rng = random.Random(7) + tokens = [rng.randrange(-(2**40), 2**40) for _ in range(40 * TOKENS_PER_BLOCK + 5)] + chain = BlockHashChain(TOKENS_PER_BLOCK) + for end in (3, TOKENS_PER_BLOCK, 7 * TOKENS_PER_BLOCK + 1, len(tokens)): + chain.extend(tokens[:end]) + assert list(chain.hashes) == reference_chain(tokens[:end], TOKENS_PER_BLOCK) + + +def test_interleaved_keys_match_namespace_keys(store_config, fake_store): + with make_worker(fake_store, layout=make_window_layout()) as worker: + namespaces = ( + list(worker._namespaces.values()) * 3 + ) # several namespaces, as a multi-rank peer list + hashes = [bytes([index]) * 16 for index in range(9)] + expected = [namespace.key(block_hash) for block_hash in hashes for namespace in namespaces] + assert worker_module._interleaved_keys(namespaces, hashes) == expected + assert worker_module._interleaved_keys(namespaces[:1], hashes) == [ + namespaces[0].key(h) for h in hashes + ] + assert worker_module._interleaved_keys(namespaces, []) == [] + + +def test_worker_window_hit_needs_the_window_at_the_prefix_end(store_config, fake_store): + with make_worker(fake_store, layout=make_window_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(6)] + for block_hash in hashes: + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + + # Every full-attention page but no window anywhere: no end is servable. + assert worker.find_prefix_hit(hashes, 0) == PrefixHit(0) + fake_store.exist_calls.clear() + + window_from = store_window(fake_store, worker, hashes, 6) + assert 0 < window_from < 6 + assert worker.find_prefix_hit(hashes, 0) == PrefixHit(6, {1: window_from}) + # The window group's pages before the window are never asked for. + probed = {key for call in fake_store.exist_calls for key in call} + assert worker._namespaces[1].key(hashes[0]) not in probed + + +def test_worker_window_hit_walks_back_over_the_rewind_only(store_config, fake_store): + with make_worker(fake_store, layout=make_window_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(6)] + for block_hash in hashes: + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + # An earlier prompt ended at block 5, one block short of the full hit. + store_window(fake_store, worker, hashes, 5) + + worker._window_fallback = False + assert worker.find_prefix_hit(hashes, 0, rewind_tokens=0).blocks == 0 + assert worker.find_prefix_hit(hashes, 0, rewind_tokens=TOKENS_PER_BLOCK).blocks == 5 + + +def test_worker_window_fallback_resumes_at_an_earlier_prompt_end(store_config, fake_store): + """Full-attention pages that outrun every stored window cost one prompt's delta, not the prefix.""" + with make_worker(fake_store, layout=make_window_layout(num_slots=32)) as worker: + hashes = [bytes([index]) * 16 for index in range(12)] + for block_hash in hashes: + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + # Earlier prompts ended at blocks 4 and 7; the latest one, at 12, has no window yet. + store_window(fake_store, worker, hashes, 4) + window_from = store_window(fake_store, worker, hashes, 7) + + assert worker.find_prefix_hit(hashes, 0, rewind_tokens=TOKENS_PER_BLOCK) == PrefixHit( + 7, {1: window_from} + ) + assert worker._lookup_stats.counts["recovered"] == 1 + + worker._window_fallback = False + assert worker.find_prefix_hit(hashes, 0, rewind_tokens=TOKENS_PER_BLOCK).blocks == 0 + assert worker._lookup_stats.counts["window_miss"] == 1 + + +def test_worker_window_fallback_needs_a_complete_window(store_config, fake_store): + with make_worker(fake_store, layout=make_window_layout(num_slots=32)) as worker: + hashes = [bytes([index]) * 16 for index in range(12)] + for block_hash in hashes: + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + window_from = store_window(fake_store, worker, hashes, 4) + # The window at block 8 lost its last page, so only the end at 4 is complete. + store_window(fake_store, worker, hashes, 8) + fake_store.objects.discard(worker._namespaces[1].key(hashes[7])) + + assert worker.find_prefix_hit(hashes, 0) == PrefixHit(4, {1: window_from}) + + +def test_worker_window_fallback_respects_what_is_on_device(store_config, fake_store): + with make_worker(fake_store, layout=make_window_layout(num_slots=32)) as worker: + hashes = [bytes([index]) * 16 for index in range(12)] + for block_hash in hashes[3:]: + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + store_window(fake_store, worker, hashes, 8) + # Resuming at 8 from a device prefix of 3 blocks: the store serves 5. + assert worker.find_prefix_hit(hashes[3:], 3).blocks == 5 + # No stored end above the device prefix: nothing to serve. + assert worker.find_prefix_hit(hashes[9:], 9).blocks == 0 + + +def test_worker_window_hit_skips_what_is_on_device(store_config, fake_store): + with make_worker(fake_store, layout=make_window_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(6)] + for block_hash in hashes[4:]: + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + # Blocks 0..3 are on device, window included, so the store only needs + # the part of the end's window from block 4 on. + _, stale_end = worker.stale_blocks(1, 6 * TOKENS_PER_BLOCK) + for block in range(max(4, stale_end), 6): + fake_store.objects.add(worker._namespaces[1].key(hashes[block])) + assert worker.find_prefix_hit(hashes[4:], 4).blocks == 2 + + +class FakeWindowWorker: + """A lookup service over group 0 (full) and group 1 (a `WINDOW` window). + + The window reads the blocks overlapping its last `WINDOW` tokens; the real + life cycle may keep a little more, which these tests do not depend on. + """ + + windowed_group_ids = (1,) + + def __init__(self, hit: PrefixHit): + self.hit = hit + self.queries = [] + + def find_prefix_hit(self, block_hashes, first_block, rewind_tokens=0): + self.queries.append((list(block_hashes), first_block, rewind_tokens)) + return PrefixHit(min(self.hit.blocks, len(block_hashes)), self.hit.window_present_from) + + def sink_blocks(self, layer_group_id): + return 0 + + def stale_blocks(self, layer_group_id, history_length): + return 0, max(0, (history_length - WINDOW) // TOKENS_PER_BLOCK) + + +def make_window_scheduler(store_config, hit=PrefixHit(0), rewind_tokens=0): + llm_args = make_llm_args() + llm_args.kv_cache_config.block_reuse_config = SimpleNamespace( + swa_endpoint_rewind_tokens=rewind_tokens + ) + scheduler = MooncakeStoreConnectorScheduler(llm_args) + scheduler._worker = FakeWindowWorker(hit) + return scheduler + + +def make_context_request(request_id, tokens, *, position=0, last_chunk=True): + request = make_request(request_id, tokens) + request.context_current_position = position + request.is_context_init_state = True + request.is_last_context_chunk = last_chunk + return request + + +def window_data(request_id, tokens, full_pages, window_pages, computed_position=0): + data = request_data(request_id, tokens, full_pages) + data.computed_position = computed_position + data.new_block_ids_by_layer_group = [list(full_pages), list(window_pages)] + return data + + +def pages_of(transfers, layer_group_id): + return [page.page_index for page in transfers.pages if page.layer_group_id == layer_group_id] + + +def test_scheduler_saves_the_window_once_at_the_prompts_end(store_config): + scheduler = make_window_scheduler(store_config) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_context_request(1, tokens, last_chunk=False) + scheduler.get_num_new_matched_tokens(request, 0) + + # The first report carries the whole prompt and the first chunk's pages. + first = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[window_data(1, tokens, [10, 11, 12], [20, 21, 22])]) + ) + # A chunk short of the prompt's end saves full attention only. + assert pages_of(first.saves[0], 0) == [10, 11, 12] + assert pages_of(first.saves[0], 1) == [] + + request.context_current_position = 12 + request.is_last_context_chunk = True + last = scheduler.build_connector_meta( + SchedulerOutput(cached_requests=[window_data(1, [], [13, 14, 15], [23, 24, 25], 12)]) + ) + # The window the prompt ends with: the blocks overlapping its last WINDOW + # tokens, blocks 4 and 5. + assert pages_of(last.saves[0], 0) == [13, 14, 15] + assert pages_of(last.saves[0], 1) == [24, 25] + assert scheduler.request_finished(request, []) is True + + +def test_scheduler_saves_the_window_ahead_of_the_full_attention_blocks(store_config, monkeypatch): + for window_first, first_group in (("1", 1), ("0", 0)): + monkeypatch.setenv(scheduler_module.WINDOW_FIRST_ENV, window_first) + scheduler = make_window_scheduler(store_config) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_context_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + metadata = scheduler.build_connector_meta( + SchedulerOutput( + new_requests=[ + window_data(1, tokens, [10, 11, 12, 13, 14, 15], [20, 21, 22, 23, 24, 25]) + ] + ) + ) + (saves,) = metadata.saves + assert saves.pages[0].layer_group_id == first_group + assert pages_of(saves, 0) == [10, 11, 12, 13, 14, 15] + assert pages_of(saves, 1) == [24, 25] + + +def test_scheduler_window_save_stops_at_what_the_request_still_holds(store_config): + """The last chunk's commit releases the rewind margin before the save thread runs.""" + scheduler = make_window_scheduler(store_config, rewind_tokens=TOKENS_PER_BLOCK) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_context_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + metadata = scheduler.build_connector_meta( + SchedulerOutput( + new_requests=[ + window_data(1, tokens, [10, 11, 12, 13, 14, 15], [20, 21, 22, 23, 24, 25]) + ] + ) + ) + # Only the window the prompt ends with, despite one block of rewind. + assert pages_of(metadata.saves[0], 1) == [24, 25] + + +def test_scheduler_window_save_skips_blocks_the_window_has_passed(store_config): + scheduler = make_window_scheduler(store_config, rewind_tokens=4 * TOKENS_PER_BLOCK) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_context_request(1, tokens, position=4 * TOKENS_PER_BLOCK) + scheduler.get_num_new_matched_tokens(request, 4 * TOKENS_PER_BLOCK) + + # The last chunk starts at block 4: blocks 0 and 1 have left the window and + # may hold another request's KV by now, even with a rewind that reaches them. + metadata = scheduler.build_connector_meta( + SchedulerOutput( + new_requests=[ + window_data(1, tokens, [10, 11, 12, 13, 14, 15], [20, 21, 22, 23, 24, 25], 16) + ] + ) + ) + assert pages_of(metadata.saves[0], 1) == [24, 25] + + +def chunked_prompt(scheduler, tokens): + """A cold prompt of 8 blocks in two chunks of 4; returns the two iterations' metadata.""" + request = make_context_request(1, tokens, last_chunk=False) + scheduler.get_num_new_matched_tokens(request, 0) + first = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[window_data(1, tokens, [10, 11, 12, 13], [20, 21, 22, 23])]) + ) + request.context_current_position = 4 * TOKENS_PER_BLOCK + request.is_last_context_chunk = True + last = scheduler.build_connector_meta( + SchedulerOutput( + cached_requests=[ + window_data(1, [], [14, 15, 16, 17], [24, 25, 26, 27], 4 * TOKENS_PER_BLOCK) + ] + ) + ) + return first, last + + +def test_scheduler_resaves_a_long_prompt_at_its_last_chunk(store_config, monkeypatch): + monkeypatch.setenv(scheduler_module.RESAVE_MIN_TOKENS_ENV, str(6 * TOKENS_PER_BLOCK)) + first, last = chunked_prompt( + make_window_scheduler(store_config), list(range(8 * TOKENS_PER_BLOCK)) + ) + assert pages_of(first.saves[0], 0) == [10, 11, 12, 13] + # The whole prompt again, so the save path refreshes the head the first chunk wrote. + assert pages_of(last.saves[0], 0) == [10, 11, 12, 13, 14, 15, 16, 17] + assert pages_of(last.saves[0], 1) == [26, 27] + + +def test_scheduler_does_not_resave_short_prompts_or_by_default(store_config, monkeypatch): + tokens = list(range(8 * TOKENS_PER_BLOCK)) + _, last = chunked_prompt(make_window_scheduler(store_config), tokens) + assert pages_of(last.saves[0], 0) == [14, 15, 16, 17] + monkeypatch.setenv(scheduler_module.RESAVE_MIN_TOKENS_ENV, str(16 * TOKENS_PER_BLOCK)) + _, last = chunked_prompt(make_window_scheduler(store_config), tokens) + assert pages_of(last.saves[0], 0) == [14, 15, 16, 17] + + +def test_scheduler_loads_window_pages_only_where_the_runtime_resumes(store_config): + scheduler = make_window_scheduler(store_config, hit=PrefixHit(4, {1: 2})) + tokens = list(range(6 * TOKENS_PER_BLOCK + 1)) + request = make_context_request(1, tokens) + assert scheduler.get_num_new_matched_tokens(request, 0) == (4 * TOKENS_PER_BLOCK, False) + + # Honoured in full: the cache gives the window pages for blocks 2 and 3 only. + request.context_current_position = 4 * TOKENS_PER_BLOCK + metadata = scheduler.build_connector_meta( + SchedulerOutput( + new_requests=[ + window_data( + 1, + tokens, + [10, 11, 12, 13, 14, 15, 16], + [BAD_PAGE_INDEX, BAD_PAGE_INDEX, 22, 23, 24, 25, 26], + ) + ] + ) + ) + assert pages_of(metadata.loads[0], 0) == [10, 11, 12, 13] + assert pages_of(metadata.loads[0], 1) == [22, 23] + + +def test_scheduler_refuses_a_short_honour_whose_window_is_not_stored(store_config): + scheduler = make_window_scheduler(store_config, hit=PrefixHit(4, {1: 2})) + tokens = list(range(6 * TOKENS_PER_BLOCK + 1)) + request = make_context_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + # The runtime resumes at block 2, whose window reads blocks 0 and 1: pages + # the store was never found to hold, since prompts do not end there. + request.context_current_position = 2 * TOKENS_PER_BLOCK + data = window_data(1, tokens, [10, 11, 12, 13, 14, 15, 16], [20, 21, 22, 23, 24, 25, 26]) + with pytest.raises(RuntimeError, match="not found to hold"): + scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + + +def test_scheduler_passes_the_rewind_and_local_prefix_to_the_lookup(store_config): + scheduler = make_window_scheduler(store_config, rewind_tokens=8) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + scheduler.get_num_new_matched_tokens(make_context_request(1, tokens), 2 * TOKENS_PER_BLOCK) + _, first_block, rewind_tokens = scheduler._worker.queries[0] + assert (first_block, rewind_tokens) == (2, 8) + + +def test_scheduler_saves_only_prompt_blocks_under_speculative_decoding(store_config): + """Rejected drafts rewind a request's page list, so a generated block's page may change hands.""" + llm_args = make_llm_args() + llm_args.speculative_config = SimpleNamespace(max_draft_len=3) + scheduler = MooncakeStoreConnectorScheduler(llm_args) + scheduler._worker = FakeWorker(0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + first = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])]) + ) + assert [page.page_index for page in first.saves[0].pages] == [4, 5] + + more_tokens = list(range(2 * TOKENS_PER_BLOCK, 3 * TOKENS_PER_BLOCK)) + second = scheduler.build_connector_meta( + SchedulerOutput(cached_requests=[request_data(1, more_tokens, [6])]) + ) + assert second.saves == [] + + +def test_scheduler_implements_the_per_layer_group_callbacks(store_config): + """A sliding-window cache has several layer groups, and the flat callbacks are never called for it.""" + scheduler = make_window_scheduler(store_config) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_context_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + scheduler.update_state_after_alloc_by_layer_group(request, [[4, 5], [6, 7]]) + scheduler.build_connector_meta( + SchedulerOutput(new_requests=[window_data(1, tokens, [4, 5], [6, 7])]) + ) + assert scheduler.request_finished_by_layer_group(request, [[4, 5], [6, 7]]) is True + assert scheduler.request_finished_by_layer_group(request, [[4, 5], [6, 7]]) is False + + +# ---- prefix reservations ---- + + +def wait_for(probe, timeout=5.0): + """Poll `probe` until it returns something truthy or `timeout` passes.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + value = probe() + if value: + return value + time.sleep(0.01) + return probe() + + +def confirmed_load(reservation_id, request_id, start, end, pages_by_group, tokens, is_async=True): + return PrefixLoad( + reservation_id, + request_id, + start, + end, + is_async, + [list(p) for p in pages_by_group], + list(tokens), + None, + ) + + +def test_scheduler_reserve_prefix_promises_an_asynchronous_load(store_config): + scheduler = make_window_scheduler(store_config, hit=PrefixHit(4, {1: 2})) + request = make_context_request(1, list(range(6 * TOKENS_PER_BLOCK + 1))) + assert scheduler.reserve_prefix(request, 0, 11) == (4 * TOKENS_PER_BLOCK, True) + + +def test_scheduler_reserve_prefix_loads_synchronously_on_request(store_config, monkeypatch): + monkeypatch.setenv(scheduler_module.SYNC_LOAD_ENV, "1") + scheduler = make_window_scheduler(store_config, hit=PrefixHit(4, {1: 2})) + request = make_context_request(1, list(range(6 * TOKENS_PER_BLOCK + 1))) + assert scheduler.reserve_prefix(request, 0, 11) == (4 * TOKENS_PER_BLOCK, False) + + +def test_scheduler_reserve_prefix_offers_nothing_without_a_hit(store_config): + scheduler = make_window_scheduler(store_config, hit=PrefixHit(0)) + request = make_context_request(1, list(range(6 * TOKENS_PER_BLOCK + 1))) + assert scheduler.reserve_prefix(request, 0, 11) == (0, False) + + +def test_scheduler_turns_a_confirmed_reservation_into_page_transfers(store_config): + scheduler = make_window_scheduler(store_config, hit=PrefixHit(4, {1: 2})) + tokens = list(range(6 * TOKENS_PER_BLOCK + 1)) + request = make_context_request(1, tokens) + scheduler.reserve_prefix(request, 0, 11) + + full = [10, 11, 12, 13, 14, 15, 16] + window = [BAD_PAGE_INDEX, BAD_PAGE_INDEX, 22, 23, 24, 25, 26] + load = confirmed_load(11, 1, 0, 4 * TOKENS_PER_BLOCK, [full, window], tokens) + metadata = scheduler.build_connector_meta(SchedulerOutput(prefix_loads=[load])) + (prefix_load,) = metadata.prefix_loads + assert (prefix_load.reservation_id, prefix_load.is_async) == (11, True) + assert pages_of(prefix_load.transfers, 0) == [10, 11, 12, 13] + # The window reads blocks 2 and 3 at the load's end; 0 and 1 have no page. + assert pages_of(prefix_load.transfers, 1) == [22, 23] + + # Once the request runs, nothing is loaded again and the store's blocks are + # not written back. + request.context_current_position = 4 * TOKENS_PER_BLOCK + later = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[window_data(1, tokens, full, window, 4 * TOKENS_PER_BLOCK)]) + ) + assert later.loads == [] and later.prefix_loads == [] + assert pages_of(later.saves[0], 0) == [14, 15] + assert pages_of(later.saves[0], 1) == [24, 25] + + +def test_scheduler_refuses_a_confirmed_load_whose_window_was_not_found(store_config): + scheduler = make_window_scheduler(store_config, hit=PrefixHit(4, {1: 2})) + tokens = list(range(6 * TOKENS_PER_BLOCK + 1)) + request = make_context_request(1, tokens) + scheduler.reserve_prefix(request, 0, 11) + + # Confirmed only up to block 2, whose window reads blocks 0 and 1: pages the + # lookup never found, since no prompt ended there. + load = confirmed_load( + 11, 1, 0, 2 * TOKENS_PER_BLOCK, [[10, 11, 12, 13], [20, 21, 22, 23]], tokens + ) + with pytest.raises(RuntimeError, match="not found to hold"): + scheduler.build_connector_meta(SchedulerOutput(prefix_loads=[load])) + + +def test_scheduler_release_trims_then_forgets_a_reservation(store_config): + scheduler = make_window_scheduler(store_config, hit=PrefixHit(4, {1: 2})) + request = make_context_request(1, list(range(6 * TOKENS_PER_BLOCK + 1))) + scheduler.reserve_prefix(request, 0, 11) + reservations = scheduler._requests[1].reservations + + scheduler.release_prefix_reservation(request, 11, 2 * TOKENS_PER_BLOCK, 4 * TOKENS_PER_BLOCK) + assert (reservations[11].start, reservations[11].end) == (0, 2 * TOKENS_PER_BLOCK) + scheduler.release_prefix_reservation(request, 11, 0, 2 * TOKENS_PER_BLOCK) + assert 11 not in reservations + + +def test_worker_loads_a_parked_requests_pages_off_the_executor_thread( + store_config, fake_store, staged_copies +): + layout = make_layout() + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x05" * 16 + key = worker._namespaces[0].key(block_hash) + fake_store.objects.add(key) + transfers = RequestTransfers(7, [PageTransfer(block_hash, 0, 3)]) + worker.bind_connector_meta( + MooncakeStoreMetadata(prefix_loads=[PrefixLoadTransfers(21, True, transfers)]) + ) + worker.start_load_kv(None) + + assert wait_for(worker.get_finished_prefix_loads) == [21] + (keys, addresses, _sizes) = fake_store.get_calls[0] + assert keys == [key] + # Through a loader's own slots, never the executor's. + loader_slots = {pool.slot_address(0) for pool in worker._async_load_stagings} + assert addresses[0][0] in loader_slots + assert addresses[0][0] != worker._load_staging.slot_address(0) + assert len(staged_copies) == 1 + + +def test_worker_starts_one_loader_per_configured_thread( + store_config, fake_store, staged_copies, monkeypatch +): + monkeypatch.setenv(worker_module.LOAD_THREADS_ENV, "3") + with make_staged_worker(fake_store, store_config, layout=make_layout()) as worker: + assert len(worker._load_threads) == 3 + # Each with slots of its own. + assert len({id(pool) for pool in worker._async_load_stagings}) == 3 + assert all(thread.is_alive() for thread in worker._load_threads) + assert worker._load_threads == [] + + +def test_worker_refuses_a_loader_count_below_one(store_config, fake_store, monkeypatch): + monkeypatch.setenv(worker_module.LOAD_THREADS_ENV, "0") + with pytest.raises(ValueError, match="at least 1"): + MooncakeStoreConnectorWorker(make_llm_args()) + + +def test_node_budget_counts_a_staging_pool_per_loader_thread(): + gib = 1 << 30 + config = SimpleNamespace(global_segment_size=0, role=SimpleNamespace(loads=True, saves=True)) + # The executor's pool, two loaders' and the save pool: 4 GiB of pinned staging. + validate_node_budget( + config, ranks_on_node=1, available_bytes=NODE_BUDGET_RESERVE_BYTES + 4 * gib, load_threads=2 + ) + with pytest.raises(ValueError, match="pin 4.0 GiB of staging"): + validate_node_budget( + config, + ranks_on_node=1, + available_bytes=NODE_BUDGET_RESERVE_BYTES + 4 * gib - 1, + load_threads=2, + ) + + +def test_worker_reports_a_synchronous_prefix_load_at_once(store_config, fake_store, staged_copies): + layout = make_layout() + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x06" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + transfers = RequestTransfers(7, [PageTransfer(block_hash, 0, 3)]) + worker.bind_connector_meta( + MooncakeStoreMetadata(prefix_loads=[PrefixLoadTransfers(22, False, transfers)]) + ) + worker.start_load_kv(None) + + assert worker.get_finished_prefix_loads() == [22] + assert worker.get_finished_prefix_loads() == [] + (_keys, addresses, _sizes) = fake_store.get_calls[0] + assert addresses == [[worker._load_staging.slot_address(0)]] + + +def test_worker_raises_a_failed_asynchronous_load_on_the_executor_thread( + store_config, fake_store, staged_copies +): + layout = make_layout() + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x08" * 16 + key = worker._namespaces[0].key(block_hash) + fake_store.objects.add(key) + fake_store.fail_gets_for.add(key) + transfers = RequestTransfers(7, [PageTransfer(block_hash, 0, 3)]) + worker.bind_connector_meta( + MooncakeStoreMetadata(prefix_loads=[PrefixLoadTransfers(23, True, transfers)]) + ) + worker.start_load_kv(None) + + def raised(): + try: + worker.get_finished_prefix_loads() + except RuntimeError: + return True + return False + + # The parked request can neither run nor be reported done, so the + # failure surfaces on the executor thread instead of hanging it. + assert wait_for(raised) + # A failed read is never scattered over the destination page. + assert staged_copies == [] + + +def test_worker_splits_a_long_probe_into_concurrent_chunks(store_config, fake_store, monkeypatch): + monkeypatch.setenv(worker_module.PROBE_CHUNK_ENV, "3") + with make_worker(fake_store, layout=make_layout()) as worker: + keys = [f"k{index}" for index in range(8)] + fake_store.objects.update(keys[:5]) + fake_store.exist_calls.clear() + + assert worker._probe(keys) == [1, 1, 1, 1, 1, 0, 0, 0] + # Three calls of at most three keys, answered in key order. + assert sorted(len(call) for call in fake_store.exist_calls) == [2, 3, 3] + assert sorted(key for call in fake_store.exist_calls for key in call) == sorted(keys) + + +def test_worker_probe_chunking_can_be_turned_off(store_config, fake_store, monkeypatch): + monkeypatch.setenv(worker_module.PROBE_CHUNK_ENV, "0") + with make_worker(fake_store, layout=make_layout()) as worker: + keys = [f"k{index}" for index in range(8)] + fake_store.exist_calls.clear() + assert worker._probe(keys) == [0] * 8 + assert [len(call) for call in fake_store.exist_calls] == [8] + + +def test_worker_reports_a_failed_chunk_as_a_miss(store_config, fake_store, monkeypatch): + monkeypatch.setenv(worker_module.PROBE_CHUNK_ENV, "3") + with make_worker(fake_store, layout=make_layout()) as worker: + real = fake_store.batch_is_exist + + def flaky(keys): + if "k4" in keys: + raise RuntimeError("master unreachable") + return real(keys) + + monkeypatch.setattr(fake_store, "batch_is_exist", flaky) + assert worker._probe([f"k{index}" for index in range(8)]) is None + + +def test_scheduler_saves_only_the_blocks_a_chunk_completes(store_config): + """A drafter's extra KV tokens give the next chunk's first block a page before it is computed.""" + scheduler = make_window_scheduler(store_config) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_context_request(1, tokens, last_chunk=False) + request.context_chunk_size = 2 * TOKENS_PER_BLOCK + scheduler.get_num_new_matched_tokens(request, 0) + + first = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[window_data(1, tokens, [10, 11, 12], [20, 21, 22])]) + ) + # Blocks 0 and 1 are computed; block 2 only has its page. + assert pages_of(first.saves[0], 0) == [10, 11] + + request.context_current_position = 2 * TOKENS_PER_BLOCK + request.context_chunk_size = 4 * TOKENS_PER_BLOCK + request.is_last_context_chunk = True + last = scheduler.build_connector_meta( + SchedulerOutput(cached_requests=[window_data(1, [], [13, 14, 15], [23, 24, 25], 8)]) + ) + assert pages_of(last.saves[0], 0) == [12, 13, 14, 15] + + +def test_validate_refuses_a_spec_recompute_tail(): + llm_args = make_llm_args() + llm_args.speculative_config = SimpleNamespace(max_draft_len=7, context_recompute_tail=2048) + with pytest.raises(NotImplementedError, match="context_recompute_tail"): + validate_llm_args(llm_args) + llm_args.speculative_config.context_recompute_tail = 0 + validate_llm_args(llm_args) + + +def test_worker_starts_no_loader_threads_for_synchronous_loads( + store_config, fake_store, staged_copies, monkeypatch +): + monkeypatch.setenv(worker_module.SYNC_LOAD_ENV, "1") + with make_staged_worker(fake_store, store_config, layout=make_layout()) as worker: + assert worker._load_threads == [] + assert worker._async_load_stagings == [] + + +def test_worker_refuses_a_window_layout_without_prefix_reservations(store_config, fake_store): + llm_args = make_llm_args() + llm_args.scheduler_config = SimpleNamespace(enable_prefix_aware_scheduling=False) + worker = MooncakeStoreConnectorWorker(llm_args) + fake_store.workers.append(worker) + with pytest.raises(ValueError, match="enable_prefix_aware_scheduling"): + worker.register_kv_cache_layout(make_window_layout())