From 01fd5ce6ff020f3c5d09442b1c517484eb43ed4d Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:43:09 -0700 Subject: [PATCH 01/24] [None][chore] Install the Mooncake Python store bindings in the container The CMake install in install_mooncake.sh only exposes the C++ transfer engine, which is what the cache transceiver links against. The mooncake-store KV cache connector needs MooncakeDistributedStore instead, and that class only exists in the Python bindings, which the source build does not produce. Install the wheel pinned to the same upstream version already built from source, so the store client and the transfer engine cannot drift apart, and widen the attribution entry to cover the wheel's payload alongside the source install. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/common/install_mooncake.sh | 8 ++++++++ scripts/attribution/scan/metadata/mooncake.yml | 5 ++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index badd5f0eb6f5..d648be2f60a6 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -50,3 +50,11 @@ cd ../.. rm -rf Mooncake echo "export LD_LIBRARY_PATH=${MOONCAKE_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" + +# The source build above only produces the C++ transfer engine, which is what +# the cache transceiver links against. MooncakeDistributedStore -- the shared +# CPU pool behind the mooncake-store KV cache connector -- is only reachable +# through the Python bindings, and those are not part of the CMake install. Take +# them from the wheel at the same upstream version so the store client and the +# transfer engine cannot drift apart. +pip3 install --no-cache-dir "mooncake-transfer-engine==${MOONCAKE_VERSION#v}" diff --git a/scripts/attribution/scan/metadata/mooncake.yml b/scripts/attribution/scan/metadata/mooncake.yml index c8d51e0f81b8..b1e6dab2c141 100644 --- a/scripts/attribution/scan/metadata/mooncake.yml +++ b/scripts/attribution/scan/metadata/mooncake.yml @@ -1,5 +1,8 @@ name: mooncake -description: Mooncake transfer engine for distributed KV cache +description: Mooncake transfer engine and distributed store for distributed KV cache source: container directory_matches: - /usr/local/Mooncake +- mooncake +basename_matches: +- mooncake_transfer_engine From 92199c7e2215fe874c5aa8bb8ea482436454706f Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:57:49 -0700 Subject: [PATCH 02/24] [None][feat] Add a Mooncake distributed store KV cache connector Regular block reuse never leaves the instance that computed a prefix, so a context server recomputes prefixes its neighbours already have. This connector publishes KV pages into a Mooncake store -- a shared CPU pool addressed by content -- so any engine can replay them, and composes with the existing point-to-point cache transceiver rather than replacing it. Built on KVCacheManagerV2's register_kv_cache_layout, since a V2 page is a set of strided byte ranges per layer group rather than one pool tensor, which is also the shape Mooncake's multi-buffer batch APIs take. V2 exposes no block hashes to a connector, so identity is a blake2b chain over (parent, salt, block tokens) computed leader-side; the key namespace additionally pins the model, shard, layer group and page geometry so any mismatch reads as a miss instead of as garbage. Loads run synchronously in start_load_kv and fail loudly: the runtime has already counted those tokens as computed, so a partial load is a wrong answer. Saves are handed to a background thread behind a CUDA event, and the leader reports such requests as saving asynchronously so their pages stay pinned until the writes land. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/__init__.py | 47 ++ .../connectors/mooncake_store/addressing.py | 152 +++++ .../connectors/mooncake_store/config.py | 202 +++++++ .../connectors/mooncake_store/keys.py | 134 +++++ .../connectors/mooncake_store/metadata.py | 62 ++ .../connectors/mooncake_store/scheduler.py | 292 ++++++++++ .../connectors/mooncake_store/worker.py | 529 ++++++++++++++++++ 7 files changed, 1418 insertions(+) create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py new file mode 100644 index 000000000000..939f846bf8dd --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""KV cache connector backed by a Mooncake distributed store. + +Offloads KV pages to a shared CPU memory pool so a prefix computed by one engine +can be replayed by another, which regular block reuse cannot do because it never +leaves the instance that computed it. + +This is a different component from the Mooncake transfer engine that the C++ +cache transceiver uses for disaggregated prefill/decode handoff: that moves KV +point to point between two known peers, while this one publishes pages into a +pool addressed by content. The two compose -- a context server can write pages +here and still hand off over NIXL. + +Requires ``KVCacheManagerV2``, which is the manager that can describe its pools +to a connector (``register_kv_cache_layout``), and the Mooncake Python bindings +(``pip install mooncake-transfer-engine``). + +Enable it with:: + + kv_connector_config = KvCacheConnectorConfig(connector="mooncake-store") + +with ``MOONCAKE_CONFIG_PATH`` pointing at a Mooncake JSON config. +""" + +from .config import MooncakeStoreConnectorConfig, StoreRole +from .scheduler import MooncakeStoreConnectorScheduler +from .worker import MooncakeStoreConnectorWorker + +__all__ = [ + "MooncakeStoreConnectorConfig", + "MooncakeStoreConnectorScheduler", + "MooncakeStoreConnectorWorker", + "StoreRole", +] diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py new file mode 100644 index 000000000000..089dfbcc99b7 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Turning a ``KvCacheLayout`` into addresses Mooncake can transfer. + +Mooncake's batch APIs take, per key, a list of ``(address, size)`` buffers. That +is exactly the shape of a V2 page: a layer group's regions each contribute one +byte range at ``base + stride * page_index``, and the concatenation of those +ranges in region order is the page's payload. + +Region order is therefore load-bearing -- it is the value's serialization -- and +``build_kv_cache_layout_v2`` derives it from the allocator's own aggregation, so +it is stable for a given model and parallel layout. ``bytes_per_page`` goes into +the key namespace to keep a geometry change from being read as a valid page. +""" + +from typing import Dict, Iterable, List, Sequence, Tuple + +from ..kv_cache_layout import KvCacheLayout, KvCacheRegion + +__all__ = ["PageAddressing", "merge_intervals"] + + +def merge_intervals(intervals: Iterable[Tuple[int, int]]) -> List[Tuple[int, int]]: + """Collapse ``(start, end)`` byte ranges into a minimal disjoint cover. + + Registration is per range and a range may not be registered twice, but + several regions routinely live inside one pool allocation: sliding-window + layer groups share it, and a non-uniform slot (MiniMax-M3's index-K sitting + beside K/V) splits one pool into several regions. Merging first means the + caller does not have to know which case it is in. + """ + ordered = sorted((int(start), int(end)) for start, end in intervals if end > start) + merged: List[Tuple[int, int]] = [] + for start, end in ordered: + if merged and start <= merged[-1][1]: + previous_start, previous_end = merged[-1] + merged[-1] = (previous_start, max(previous_end, end)) + else: + merged.append((start, end)) + return merged + + +class PageAddressing: + """Resolves ``(layer group, page index)`` to the byte ranges of that page.""" + + def __init__(self, layout: KvCacheLayout): + self._layout = layout + self._regions: Dict[int, Tuple[KvCacheRegion, ...]] = {} + self._bytes_per_page: Dict[int, int] = {} + self._num_slots: Dict[int, int] = {} + for group in layout.groups: + if not group.regions: + raise ValueError( + f"layer group {group.layer_group_id} has no KV regions; there " + "is nothing for the connector to transfer" + ) + self._regions[group.layer_group_id] = group.regions + self._bytes_per_page[group.layer_group_id] = group.bytes_per_page + # Every region of a group is drawn from the same pool group, so they + # share a slot count; disagreement would mean the page index space is + # not the single space the layout documents. + slot_counts = {region.num_slots for region in group.regions} + if len(slot_counts) != 1: + raise ValueError( + f"layer group {group.layer_group_id} mixes slot counts " + f"{sorted(slot_counts)}; page indices would be ambiguous" + ) + self._num_slots[group.layer_group_id] = slot_counts.pop() + + @property + def layout(self) -> KvCacheLayout: + """The layout this addressing was built from.""" + return self._layout + + @property + def layer_group_ids(self) -> Tuple[int, ...]: + """Layer group ids covered, in layout order.""" + return tuple(group.layer_group_id for group in self._layout.groups) + + @property + def tokens_per_block(self) -> int: + """Tokens held by one page.""" + return self._layout.tokens_per_block + + def bytes_per_page(self, layer_group_id: int) -> int: + """Total payload size of one page of ``layer_group_id``.""" + return self._bytes_per_page[layer_group_id] + + def num_slots(self, layer_group_id: int) -> int: + """Number of page slots addressable in ``layer_group_id``.""" + return self._num_slots[layer_group_id] + + def buffers(self, layer_group_id: int, page_index: int) -> Tuple[List[int], List[int]]: + """Addresses and sizes of one page, in the order they concatenate. + + Args: + layer_group_id: Layer group the page index is scoped to. + page_index: Page slot index within that group. + + Returns: + Parallel lists of device addresses and byte counts. + """ + regions = self._regions[layer_group_id] + num_slots = self._num_slots[layer_group_id] + if not 0 <= page_index < num_slots: + raise IndexError( + f"page index {page_index} out of range [0, {num_slots}) for layer " + f"group {layer_group_id}" + ) + addresses = [region.base + region.stride * page_index for region in regions] + sizes = [region.size for region in regions] + return addresses, sizes + + def registration_ranges(self) -> List[Tuple[int, int]]: + """Byte ranges to hand to ``register_buffer``, deduplicated and merged. + + A region's slots are strided rather than packed, so the range covering it + is the whole span from the first slot to the end of the last. Registering + the span is what makes every slot's address valid for RDMA, and merging + keeps a shared pool from being registered once per region. + """ + spans: List[Tuple[int, int]] = [] + for regions in self._regions.values(): + for region in regions: + span_end = region.base + region.stride * (region.num_slots - 1) + region.size + spans.append((region.base, span_end)) + return merge_intervals(spans) + + def describe(self) -> str: + """A one-line summary for startup logs.""" + parts: Sequence[str] = [ + f"lg{group.layer_group_id}(" + f"layers={len(group.layer_ids)}, " + f"regions={len(group.regions)}, " + f"bytes/page={group.bytes_per_page}, " + f"slots={self._num_slots[group.layer_group_id]}, " + f"window={group.window_size})" + for group in self._layout.groups + ] + return f"tokens_per_block={self.tokens_per_block}, " + ", ".join(parts) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py new file mode 100644 index 000000000000..0e69d7270e72 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Configuration for the Mooncake store KV cache connector. + +Topology settings are read from the JSON file named by ``MOONCAKE_CONFIG_PATH``, +the same file and environment variable the vLLM Mooncake store connector uses, +so one deployment can point both engines at the same pool. + +``KvCacheConnectorConfig`` carries no free-form dictionary, so the two settings +that are TensorRT-LLM's rather than Mooncake's -- the read/write role and the +key prefix -- are also taken from the environment. +""" + +import json +import os +import re +from dataclasses import dataclass +from enum import Enum +from typing import Any, Optional + +__all__ = [ + "CONFIG_PATH_ENV", + "MooncakeStoreConnectorConfig", + "ROLE_ENV", + "StoreRole", +] + +CONFIG_PATH_ENV = "MOONCAKE_CONFIG_PATH" +ROLE_ENV = "TRTLLM_MOONCAKE_STORE_ROLE" +CACHE_PREFIX_ENV = "TRTLLM_MOONCAKE_STORE_PREFIX" +MODEL_KEY_ENV = "TRTLLM_MOONCAKE_STORE_MODEL_KEY" + +DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200 +DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 +DEFAULT_CACHE_PREFIX = "trtllm" + +_SIZE_UNITS = { + "": 1, + "b": 1, + "k": 1000, + "kb": 1000, + "m": 1000**2, + "mb": 1000**2, + "g": 1000**3, + "gb": 1000**3, + "t": 1000**4, + "tb": 1000**4, + "kib": 1024, + "mib": 1024**2, + "gib": 1024**3, + "tib": 1024**4, +} +_SIZE_RE = re.compile(r"^\s*([0-9]+(?:\.[0-9]+)?)\s*([a-zA-Z]*)\s*$") + + +class StoreRole(Enum): + """Which directions of traffic this engine is allowed to drive. + + A disaggregated deployment typically runs context servers as ``both`` and + leaves generation servers unconfigured: generated tokens are rarely a reused + prefix, so writing them costs bandwidth for no hit rate. + """ + + PRODUCER = "producer" + CONSUMER = "consumer" + BOTH = "both" + + @property + def loads(self) -> bool: + """Whether this role reads previously stored KV back onto the GPU.""" + return self is not StoreRole.PRODUCER + + @property + def saves(self) -> bool: + """Whether this role writes newly computed KV into the store.""" + return self is not StoreRole.CONSUMER + + +def _parse_size(value: Any) -> int: + """Accept either a byte count or a suffixed string such as ``"4GiB"``.""" + if isinstance(value, bool): + raise ValueError(f"expected a size, got {value!r}") + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + match = _SIZE_RE.match(str(value)) + if match is None: + raise ValueError(f"cannot parse size {value!r}") + magnitude, unit = match.groups() + scale = _SIZE_UNITS.get(unit.lower()) + if scale is None: + raise ValueError(f"unknown size unit {unit!r} in {value!r}") + return int(float(magnitude) * scale) + + +@dataclass(frozen=True) +class MooncakeStoreConnectorConfig: + """Everything needed to open a store handle and name keys in it.""" + + metadata_server: str + master_server_address: str + protocol: str = "rdma" + device_name: str = "" + global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE + local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE + local_hostname: Optional[str] = None + tenant_id: Optional[str] = None + role: StoreRole = StoreRole.BOTH + cache_prefix: str = DEFAULT_CACHE_PREFIX + #: Identity the keys are namespaced by. Two engines only share cache when + #: they agree on this, so it defaults to the model directory's basename + #: rather than its full path: the same checkpoint is routinely mounted + #: somewhere else on another host, which is exactly the case sharing is for. + model_key: Optional[str] = None + #: How many page keys go into one store call. Bounds the size of a single + #: RPC without bounding how much a request may transfer. + transfer_batch_size: int = 64 + + def __post_init__(self) -> None: + """Reject settings that would fail later, inside a transfer.""" + if not self.master_server_address: + raise ValueError("master_server_address is required") + if self.local_buffer_size <= 0: + raise ValueError("local_buffer_size must be > 0") + if self.global_segment_size < 0: + raise ValueError("global_segment_size must be >= 0") + if self.transfer_batch_size <= 0: + raise ValueError("transfer_batch_size must be > 0") + + @staticmethod + def from_file(path: str) -> "MooncakeStoreConnectorConfig": + """Read the topology from a vLLM-compatible Mooncake JSON config.""" + with open(path) as handle: + raw = json.load(handle) + return MooncakeStoreConnectorConfig( + metadata_server=raw.get("metadata_server", ""), + master_server_address=raw.get("master_server_address", ""), + protocol=raw.get("protocol", "rdma"), + device_name=raw.get("device_name", ""), + global_segment_size=_parse_size( + raw.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE) + ), + local_buffer_size=_parse_size(raw.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)), + local_hostname=raw.get("local_hostname") or None, + tenant_id=raw.get("tenant_id") or None, + role=StoreRole(str(raw.get("role", StoreRole.BOTH.value)).strip().lower()), + cache_prefix=str(raw.get("cache_prefix", DEFAULT_CACHE_PREFIX)), + model_key=raw.get("model_key") or None, + transfer_batch_size=int(raw.get("transfer_batch_size", 64)), + ) + + @staticmethod + def from_env() -> "MooncakeStoreConnectorConfig": + """Load the JSON config, then apply the TensorRT-LLM env overrides.""" + path = os.getenv(CONFIG_PATH_ENV) + if not path: + raise ValueError( + f"The mooncake-store connector needs {CONFIG_PATH_ENV} set to a " + "Mooncake JSON config (metadata_server, master_server_address, " + "protocol, device_name, global_segment_size, local_buffer_size)." + ) + config = MooncakeStoreConnectorConfig.from_file(path) + return config.with_env_overrides() + + def with_env_overrides(self) -> "MooncakeStoreConnectorConfig": + """Apply ``TRTLLM_MOONCAKE_STORE_*`` on top of the file's settings.""" + import dataclasses + + updates: dict[str, Any] = {} + role = os.getenv(ROLE_ENV) + if role: + try: + updates["role"] = StoreRole(role.strip().lower()) + except ValueError as exc: + known = ", ".join(member.value for member in StoreRole) + raise ValueError(f"{ROLE_ENV}={role!r} is not one of: {known}") from exc + prefix = os.getenv(CACHE_PREFIX_ENV) + if prefix: + updates["cache_prefix"] = prefix + model_key = os.getenv(MODEL_KEY_ENV) + if model_key: + updates["model_key"] = model_key + return dataclasses.replace(self, **updates) if updates else self + + def resolve_model_key(self, model: Any) -> str: + """The model identity to namespace keys by, given the configured model.""" + if self.model_key: + return self.model_key + return os.path.basename(str(model).rstrip("/")) or str(model) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py new file mode 100644 index 000000000000..fe5ea6965e79 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Block identity and store key naming for the Mooncake store connector. + +``KVCacheManagerV2`` exposes no block hashes to a connector -- ``RequestData`` +reports them empty -- so content identity is derived here instead. The chain is +the standard one: a block's hash covers its own tokens *and* every token before +it, so a key can only be reused by a request whose prefix is byte-identical. + +A key is ``/``. The namespace pins down everything that +would make the stored bytes mean something different: the model, the shard that +produced them, the layer group inside that shard, the tokens each page holds and +how many bytes a page is. Anything that changes those reads as a cache miss +rather than as garbage. +""" + +import hashlib +from dataclasses import dataclass +from typing import List, Optional, Sequence + +__all__ = [ + "BlockHashChain", + "KeyNamespace", + "HASH_DIGEST_BYTES", +] + +#: 128 bits. Collisions decide whether one request reads another's KV, so the +#: digest is sized to make that negligible over any realistic cache lifetime, +#: while staying half the width of a full blake2b digest in every key. +HASH_DIGEST_BYTES = 16 + + +def _digest(*parts: bytes) -> bytes: + hasher = hashlib.blake2b(digest_size=HASH_DIGEST_BYTES) + for part in parts: + hasher.update(part) + return hasher.digest() + + +class BlockHashChain: + """Rolling hashes of a request's full blocks, one entry per block ordinal. + + Extended in place as a request's token list grows, so generation steps cost + one digest per newly completed block rather than a rehash of the prompt. + """ + + def __init__(self, tokens_per_block: int, cache_salt: Optional[str] = None): + if tokens_per_block <= 0: + raise ValueError(f"tokens_per_block must be > 0, got {tokens_per_block}") + self._tokens_per_block = int(tokens_per_block) + # The salt seeds the chain rather than being mixed into every block, so + # a request carrying a different salt diverges from the first block on. + salt_bytes = b"" if cache_salt is None else str(cache_salt).encode() + self._seed = _digest(b"salt", salt_bytes) + self._hashes: List[bytes] = [] + + @property + def tokens_per_block(self) -> int: + """Tokens covered by each entry in the chain.""" + return self._tokens_per_block + + @property + def hashes(self) -> Sequence[bytes]: + """Hashes computed so far, indexed by block ordinal.""" + return self._hashes + + def extend(self, tokens: Sequence[int]) -> Sequence[bytes]: + """Grow the chain to cover every full block of ``tokens``. + + Args: + tokens: The request's complete token list, prompt first. Must be an + extension of what was passed previously; a request's tokens only + ever grow, so a shorter list means the caller mixed up requests. + + Returns: + The full chain, indexed by block ordinal. + """ + num_full_blocks = len(tokens) // self._tokens_per_block + if num_full_blocks < len(self._hashes): + raise ValueError( + 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)) + return self._hashes + + +@dataclass(frozen=True) +class KeyNamespace: + """The part of a store key that is fixed for one shard and layer group.""" + + cache_prefix: str + model_key: str + #: Global rank of the shard whose KV these bytes are, and the world size it + #: was produced under. Both are needed: rank 3 of 8 holds different heads + #: than rank 3 of 4. + rank: int + world_size: int + layer_group_id: int + tokens_per_block: int + bytes_per_page: int + + @property + def prefix(self) -> str: + """The literal string every key in this namespace starts with.""" + return ( + f"{self.cache_prefix}/{self.model_key}" + f"/w{self.world_size}r{self.rank}" + f"/lg{self.layer_group_id}" + f"/t{self.tokens_per_block}b{self.bytes_per_page}" + ) + + def key(self, block_hash: bytes) -> str: + """The store key holding one page of this namespace.""" + return f"{self.prefix}/{block_hash.hex()}" diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py new file mode 100644 index 000000000000..6246877805f6 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The per-iteration work list the scheduler hands the workers. + +Instances are broadcast from rank 0 to every worker, so these carry only plain +data: a page's identity (its block hash) and where that page currently lives on +this rank (a layer group and a page slot index). Deliberately no store keys -- +each worker prefixes its own rank namespace, so one broadcast serves all shards. +""" + +from dataclasses import dataclass, field +from typing import List + +__all__ = ["MooncakeStoreMetadata", "PageTransfer", "RequestTransfers"] + + +@dataclass +class PageTransfer: + """One page of one layer group, to move in either direction.""" + + #: Content identity from ``BlockHashChain``; names the key, not the location. + block_hash: bytes + layer_group_id: int + #: Page slot index within ``layer_group_id``, as reported by + #: ``RequestData.new_block_ids_by_layer_group``. + page_index: int + + +@dataclass +class RequestTransfers: + """Pages belonging to one request, kept together for save bookkeeping. + + The worker owes ``get_finished`` an answer per request, so a save's owner has + to survive the trip from scheduler to worker. + """ + + request_id: int + pages: List[PageTransfer] = field(default_factory=list) + + +@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) + + def __bool__(self) -> bool: + """Whether there is any work at all this iteration.""" + return bool(self.loads or self.saves) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py new file mode 100644 index 000000000000..05d1f6e3551c --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Leader side of the Mooncake store KV cache connector. + +Runs only on rank 0. It decides what to load and what to save; the workers do +the moving. Two pieces of bookkeeping make that possible, and both exist because +``KVCacheManagerV2`` reports ``RequestData.block_hashes`` empty: + +* a hash chain per request, so a block has a content identity at all; +* 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 only savable once it is full, so the index + has to be remembered from the step that reported it. +""" + +from typing import Dict, List, Optional, Tuple + +from tensorrt_llm.bindings.internal.batch_manager import LlmRequest +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +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 .config import MooncakeStoreConnectorConfig +from .keys import BlockHashChain +from .metadata import MooncakeStoreMetadata, PageTransfer, RequestTransfers +from .worker import MooncakeStoreConnectorWorker, resolve_local_worker + +__all__ = ["MooncakeStoreConnectorScheduler"] + + +class _RequestState: + """Per-request bookkeeping that has to outlive a single iteration.""" + + __slots__ = ( + "chain", + "tokens", + "pages", + "saved_upto", + "load_first_block", + "load_blocks", + "emitted_saves", + ) + + def __init__(self, chain: BlockHashChain): + self.chain = chain + #: The request's tokens, accumulated from the per-step deltas. + self.tokens: List[int] = [] + #: Page slot index per block ordinal, per layer group. + self.pages: Dict[int, List[int]] = {} + #: First block ordinal not yet considered for saving. + self.saved_upto = 0 + #: The offer made by ``get_num_new_matched_tokens``, in block ordinals. + self.load_first_block = 0 + self.load_blocks = 0 + self.emitted_saves = False + + +class MooncakeStoreConnectorScheduler(KvCacheConnectorScheduler): + """Chooses which pages the Mooncake pool serves and which it receives.""" + + def __init__(self, llm_args: TorchLlmArgs): + super().__init__(llm_args) + + self._config = MooncakeStoreConnectorConfig.from_env() + self._tokens_per_block = int(llm_args.kv_cache_config.tokens_per_block) + self._requests: Dict[int, _RequestState] = {} + self._worker: Optional[MooncakeStoreConnectorWorker] = None + + logger.info( + "mooncake-store leader ready (role=%s, tokens_per_block=%d)", + self._config.role.value, + self._tokens_per_block, + ) + + def wait_for_initialization(self): + """Bind to the process-local worker, which owns the store handle. + + Called after the executor has built both halves and registered the KV + cache layout, which is what the worker needs before it can name a key. + """ + self._worker = resolve_local_worker() + + # ---- lookup ---- + + def get_num_new_matched_tokens( + self, request: LlmRequest, num_computed_tokens: int + ) -> Tuple[int, bool]: + """Offer the longest stored prefix beyond what the device already has. + + Args: + request: The request being scheduled. + num_computed_tokens: Tokens already matched in the local KV cache. + + Returns: + Tokens the store can supply, and ``False`` for a synchronous load. + """ + tokens = request.get_tokens(0) + state = self._state_for(request, tokens) + state.load_first_block = 0 + state.load_blocks = 0 + + if not self._config.role.loads: + return 0, False + + # 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 + + first_block = num_computed_tokens // self._tokens_per_block + # Stop one token short of the prompt: the runtime still has to run a + # forward pass for this request, and it cannot do that with nothing left + # to compute. + 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) + if hit_blocks == 0: + return 0, False + + state.load_first_block = first_block + state.load_blocks = hit_blocks + logger.debug( + "mooncake-store matched %d blocks (%d tokens) for request %d", + hit_blocks, + hit_blocks * self._tokens_per_block, + request.request_id, + ) + return hit_blocks * self._tokens_per_block, False + + def cancel_load(self, request: LlmRequest, start: int, end: int): + """Drop offered blocks whose tokens the runtime will not consume. + + Loads here are synchronous and nothing has been transferred yet, so this + is exact: the offer is truncated before ``build_connector_meta`` turns it + into work. + """ + state = self._requests.get(request.request_id) + if state is None or state.load_blocks == 0: + return + kept = 0 + for offset in range(state.load_blocks): + block = state.load_first_block + offset + block_start = block * self._tokens_per_block + if block_start + self._tokens_per_block > start and block_start < end: + break + kept += 1 + state.load_blocks = kept + + def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): + """No-op: page indices are read from the scheduler output instead. + + The flat ``block_ids`` here are a single space, but a V2 page index is + scoped to a layer group. ``RequestData.new_block_ids_by_layer_group`` is + the form that stays correct for every model, so that is the only source + this connector uses. + """ + + # ---- work lists ---- + + def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeStoreMetadata: + """Turn this iteration's scheduled requests into load and save lists.""" + metadata = MooncakeStoreMetadata() + for request_data in (*scheduler_output.new_requests, *scheduler_output.cached_requests): + state = self._requests.get(request_data.request_id) + if state is None: + # Only requests that went through get_num_new_matched_tokens have + # a hash chain. Generation-only requests never do, and the + # connector manager refuses them outright. + continue + + state.tokens.extend(request_data.new_tokens) + state.chain.extend(state.tokens) + self._record_pages(state, request_data) + + loads = self._loads_for(state, request_data) + if loads.pages: + metadata.loads.append(loads) + + # Whatever the store just supplied, and whatever the local cache + # matched, is not ours to write back: the store already has the + # former, and the latter was never allocated during this run. + state.saved_upto = max(state.saved_upto, state.load_first_block + state.load_blocks) + # An offer is consumed once. The load is issued in exactly the + # iteration the runtime allocated pages to hold it. + state.load_blocks = 0 + + if self._config.role.saves: + saves = self._saves_for(state, request_data) + if saves.pages: + state.emitted_saves = True + metadata.saves.append(saves) + return metadata + + def request_finished(self, request: LlmRequest, cache_block_ids: List[int]) -> bool: + """Report whether pages must stay pinned for in-flight saves. + + Returns: + True when this request handed any page to the background save + thread. Its pages are the source of those RDMA reads, so freeing + them now would let a later request overwrite bytes mid-transfer. + """ + state = self._requests.pop(request.request_id, None) + return bool(state is not None and state.emitted_saves) + + # ---- internals ---- + + def _require_worker(self) -> MooncakeStoreConnectorWorker: + if self._worker is None: + self._worker = resolve_local_worker() + return self._worker + + def _state_for(self, request: LlmRequest, tokens: List[int]) -> _RequestState: + state = self._requests.get(request.request_id) + if state is None: + state = _RequestState( + BlockHashChain(self._tokens_per_block, cache_salt=request.cache_salt) + ) + self._requests[request.request_id] = state + # Hashing the prompt here rather than waiting for the first scheduler + # output is the whole point: the lookup happens before the request is + # scheduled, so the chain has to be ready before any metadata exists. + state.chain.extend(tokens) + return state + + def _record_pages(self, state: _RequestState, request_data: RequestData) -> None: + """Append this step's newly allocated page indices, by block ordinal.""" + by_group = request_data.new_block_ids_by_layer_group + if not by_group: + # Under a single layer group the manager also mirrors that group's + # indices into the flat ``new_block_ids``, but it does not say which + # group they belong to, so there is nothing safe to record from it. + return + for layer_group_id, indices in by_group.items(): + state.pages.setdefault(layer_group_id, []).extend(int(index) for index in indices) + + def _addressable_blocks(self, state: _RequestState) -> int: + """Block ordinals that are both hashed and backed by a page everywhere.""" + if not state.pages: + return 0 + return min(len(state.chain.hashes), min(len(indices) for indices in state.pages.values())) + + def _loads_for(self, state: _RequestState, request_data: RequestData) -> RequestTransfers: + transfers = RequestTransfers(request_data.request_id) + limit = self._addressable_blocks(state) + for offset in range(state.load_blocks): + block = state.load_first_block + offset + if block >= limit: + # The runtime allocated fewer pages than it accepted tokens for. + # It reports the shortfall through cancel_load; until then the + # unaddressable tail is simply not loaded. + break + self._append_pages(state, transfers, block) + return transfers + + def _saves_for(self, state: _RequestState, request_data: RequestData) -> RequestTransfers: + transfers = RequestTransfers(request_data.request_id) + limit = self._addressable_blocks(state) + for block in range(state.saved_upto, limit): + self._append_pages(state, transfers, block) + state.saved_upto = max(state.saved_upto, limit) + 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.""" + block_hash = state.chain.hashes[block] + pages: List[PageTransfer] = [] + for layer_group_id, indices in state.pages.items(): + page_index = indices[block] + if page_index == BAD_PAGE_INDEX: + # The block has no page in this group -- a sliding window has + # already dropped it. A partial page is not a usable cache entry, + # so the whole block is skipped. + return + pages.append(PageTransfer(block_hash, layer_group_id, page_index)) + transfers.pages.extend(pages) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py new file mode 100644 index 000000000000..6a21f9a8178d --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -0,0 +1,529 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Worker side of the Mooncake store KV cache connector. + +One worker per rank owns a ``MooncakeDistributedStore`` handle and moves pages +between that pool and its own GPU KV cache. It is also the only place that knows +how a page is addressed and how a key is spelled, which is why the leader -- +colocated with rank 0's worker in the same process -- asks it to run prefix +lookups instead of rebuilding that knowledge. + +Loads are synchronous: the runtime has already told the scheduler those tokens +are computed, so the bytes must be in place before the forward pass reads them, +and a failed load is a wrong answer rather than a slow one. + +Saves are asynchronous and gated on a CUDA event. The pages are only complete +once the forward pass that wrote them has retired, and blocking the executor +loop on an RDMA write is exactly the cost the store is supposed to avoid. The +scheduler reports such a request as saving asynchronously, which keeps its pages +pinned until ``get_finished`` says the writes landed. +""" + +import threading +from collections import defaultdict +from queue import Queue +from typing import Dict, List, Optional, Sequence, Set, Tuple + +import torch + +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 ..kv_cache_connector import KvCacheConnectorWorker +from ..kv_cache_layout import KvCacheLayout +from .addressing import PageAddressing +from .config import CONFIG_PATH_ENV, MooncakeStoreConnectorConfig +from .keys import KeyNamespace +from .metadata import MooncakeStoreMetadata, RequestTransfers + +__all__ = ["MooncakeStoreConnectorWorker", "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 +#: connection or an out-of-band channel. See ``py_executor_creator``, which +#: constructs scheduler and worker concurrently for exactly this kind of +#: mutual dependency. +_LOCAL_WORKER: Optional["MooncakeStoreConnectorWorker"] = None +_LOCAL_WORKER_READY = threading.Event() + + +def resolve_local_worker(timeout: float = 60.0) -> "MooncakeStoreConnectorWorker": + """The worker living in this process, once it has been constructed. + + Args: + timeout: Seconds to wait. Construction is concurrent with the leader's, + so a short wait is expected; exceeding it means the worker failed. + + Returns: + The process-local worker. + """ + if not _LOCAL_WORKER_READY.wait(timeout): + raise RuntimeError( + "The mooncake-store leader could not find a worker in its process. " + "The leader only runs on rank 0, where the executor also builds a " + "worker, so this means worker construction failed." + ) + assert _LOCAL_WORKER is not None + return _LOCAL_WORKER + + +def _open_store(config: MooncakeStoreConnectorConfig): + """Connect to the Mooncake master and return a live store handle.""" + try: + from mooncake.store import MooncakeDistributedStore + except ImportError as exc: + raise ImportError( + "The mooncake-store connector needs the Mooncake Python bindings " + "(`pip install mooncake-transfer-engine`). The C++ transfer engine " + "built into the container is a different component and does not " + "provide MooncakeDistributedStore." + ) from exc + + store = MooncakeDistributedStore() + hostname = config.local_hostname or _default_hostname() + setup_kwargs = {} + if config.tenant_id: + setup_kwargs["tenant_id"] = config.tenant_id + status = store.setup( + hostname, + config.metadata_server, + config.global_segment_size, + config.local_buffer_size, + config.protocol, + config.device_name, + config.master_server_address, + **setup_kwargs, + ) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.setup failed with status {status} " + f"(master={config.master_server_address!r}, " + f"metadata={config.metadata_server!r}, protocol={config.protocol!r}). " + f"Check the config named by {CONFIG_PATH_ENV}." + ) + return store + + +def _default_hostname() -> str: + import socket + + return socket.gethostbyname(socket.gethostname()) + + +def _batched(items: Sequence, size: int): + for start in range(0, len(items), size): + yield items[start : start + size] + + +class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): + """Moves KV pages between this rank's GPU cache and the Mooncake pool.""" + + def __init__(self, llm_args: TorchLlmArgs): + super().__init__(llm_args) + + self._config = MooncakeStoreConnectorConfig.from_env() + self._rank = mpi_rank() + self._world_size = mpi_world_size() + self._model_key = self._config.resolve_model_key(llm_args.model) + + self._addressing: Optional[PageAddressing] = None + # Namespaces for this rank, used for both directions of transfer. + self._namespaces: Dict[int, KeyNamespace] = {} + # 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, ...]] = {} + + self._store = _open_store(self._config) + + self._save_queue: "Queue[Optional[Tuple[torch.cuda.Event, List[RequestTransfers]]]]" = ( + Queue() + ) + self._save_thread: Optional[threading.Thread] = None + self._save_lock = threading.Lock() + # Save submissions still in flight, per request. + self._outstanding_saves: Dict[int, int] = defaultdict(int) + # Requests the runtime has told us are done producing KV. Their pages + # stay pinned until we report them back through ``get_finished``. + self._closed_requests: Set[int] = set() + self._save_error: Optional[BaseException] = None + + global _LOCAL_WORKER + _LOCAL_WORKER = self + _LOCAL_WORKER_READY.set() + + logger.info( + "mooncake-store worker rank %d/%d ready (role=%s, model_key=%s, master=%s)", + self._rank, + self._world_size, + self._config.role.value, + self._model_key, + self._config.master_server_address, + ) + + # ---- registration ---- + + def register_kv_caches(self, kv_cache_tensor: torch.Tensor): + """Reject the V1 single-pool registration. + + Raises: + NotImplementedError: Always. Identity here is a hash chain the + connector computes itself, keyed per layer group, and the V1 + manager supplies real block hashes over a single flat block + space instead. Running the V2 addressing against V1 block ids + would silently mislabel pages, so V1 is refused rather than + approximated. + """ + raise NotImplementedError( + "The mooncake-store connector requires KVCacheManagerV2. Set " + "kv_cache_config.use_kv_cache_manager_v2=True." + ) + + def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: + """Register the KV pools with Mooncake and start the save thread.""" + if self._addressing is not None: + raise RuntimeError("KV cache layout already registered") + + addressing = PageAddressing(layout) + for start, end in addressing.registration_ranges(): + status = self._store.register_buffer(start, end - start) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.register_buffer failed with status " + f"{status} for [{start:#x}, {end:#x}). Without registration " + "the store cannot read or write these pages." + ) + + self._addressing = addressing + for layer_group_id in addressing.layer_group_ids: + bytes_per_page = addressing.bytes_per_page(layer_group_id) + self._namespaces[layer_group_id] = self._namespace( + self._rank, layer_group_id, bytes_per_page + ) + self._peer_namespaces[layer_group_id] = tuple( + self._namespace(rank, layer_group_id, bytes_per_page) + for rank in range(self._world_size) + ) + + if self._config.role.saves: + self._save_thread = threading.Thread( + target=self._drain_saves, + name=f"mooncake-store-save-{self._rank}", + daemon=True, + ) + self._save_thread.start() + + logger.info( + "mooncake-store worker rank %d registered layout: %s", + self._rank, + addressing.describe(), + ) + + def _namespace(self, rank: int, layer_group_id: int, bytes_per_page: int) -> KeyNamespace: + return KeyNamespace( + cache_prefix=self._config.cache_prefix, + model_key=self._model_key, + rank=rank, + world_size=self._world_size, + layer_group_id=layer_group_id, + tokens_per_block=self._addressing.tokens_per_block, + bytes_per_page=bytes_per_page, + ) + + # ---- leader-facing lookup ---- + + @property + def config(self) -> MooncakeStoreConnectorConfig: + """The resolved connector configuration.""" + return self._config + + @property + def is_registered(self) -> bool: + """Whether a KV cache layout has been registered yet.""" + return self._addressing is not None + + def count_prefix_hit(self, block_hashes: Sequence[bytes]) -> int: + """How many leading blocks of ``block_hashes`` are fully present. + + A block counts only when every layer group and every rank has its page, + because a prefix is replayed as a whole. The scan stops at the first + incomplete block: the runtime consumes a prefix, so a later hit is not + usable on its own. + + Args: + block_hashes: Candidate hashes in block order. + + Returns: + Length of the usable prefix, in blocks. + """ + if not block_hashes or self._addressing is None: + return 0 + + 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) + + try: + present = self._store.batch_is_exist(keys) + except Exception: + logger.warning("mooncake-store lookup failed; treating as a miss", exc_info=True) + return 0 + + if len(present) != len(keys): + logger.warning( + "mooncake-store batch_is_exist returned %d results for %d keys; treating as a miss", + len(present), + len(keys), + ) + return 0 + + 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 + + # ---- load path ---- + + def start_load_kv(self, stream: torch.cuda.Stream): + """Pull every scheduled page into its GPU slot before the forward pass.""" + metadata: Optional[MooncakeStoreMetadata] = self.get_connector_meta() + if metadata is None or not metadata.loads: + return + self._reraise_save_error() + + keys, addresses, sizes, total_pages = self._resolve(metadata.loads) + if not keys: + return + + for batch in zip( + _batched(keys, self._config.transfer_batch_size), + _batched(addresses, self._config.transfer_batch_size), + _batched(sizes, self._config.transfer_batch_size), + ): + batch_keys, batch_addresses, batch_sizes = batch + results = self._store.batch_get_into_multi_buffers( + list(batch_keys), list(batch_addresses), list(batch_sizes) + ) + failed = [ + key + for key, result in zip(batch_keys, results) + if not isinstance(result, int) or result < 0 + ] + if failed or len(results) != len(batch_keys): + # The runtime already counted these tokens as computed, so a + # partial load leaves the forward pass reading uninitialized KV + # and silently producing wrong tokens. Fail loudly instead. + raise RuntimeError( + f"mooncake-store failed to load {len(failed) or len(batch_keys)} of " + f"{len(batch_keys)} pages; the affected KV slots were already " + f"reported as computed. First failure: {failed[:1]}" + ) + + logger.debug("mooncake-store rank %d loaded %d pages", self._rank, total_pages) + + def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream): + """No-op: loads complete in ``start_load_kv``. + + Transfers are whole pages, so a page's bytes for every layer in a group + land in one store call rather than layer by layer. There is nothing left + outstanding by the time the first layer runs. + """ + + def save_kv_layer(self, layer_idx: int, stream: torch.cuda.Stream): + """No-op: saves are submitted once per pass in ``wait_for_save``. + + A page is only complete when every layer of its group has written its + slice, so there is no correct per-layer submission point. + """ + + # ---- save path ---- + + 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: + return + self._reraise_save_error() + + # The pages are written by kernels still queued on this stream. The event + # is the handoff: the thread reads GPU memory only after the pass retires, + # and the executor loop is not blocked waiting for that. + event = torch.cuda.Event() + event.record(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))) + + def get_finished( + self, finished_gen_req_ids: List[int], started_loading_req_ids: List[int] + ) -> Tuple[List[int], List[int]]: + """Report which requests' saves have landed. + + Args: + finished_gen_req_ids: Requests that will produce no further KV. + started_loading_req_ids: Requests loading asynchronously. Always + empty here, since ``get_num_new_matched_tokens`` only ever + offers synchronous loads; echoed back so the runtime does not + wait on something that already happened. + + Returns: + Requests that have finished saving, and requests that have finished + loading. + """ + self._reraise_save_error() + with self._save_lock: + self._closed_requests.update(finished_gen_req_ids) + finished_saving = [ + request_id + for request_id in self._closed_requests + if self._outstanding_saves.get(request_id, 0) == 0 + ] + for request_id in finished_saving: + self._closed_requests.discard(request_id) + self._outstanding_saves.pop(request_id, None) + return finished_saving, list(started_loading_req_ids) + + def _drain_saves(self) -> None: + torch.cuda.set_device(torch.cuda.current_device()) + while True: + item = self._save_queue.get() + if item is None: + return + event, transfers = item + try: + event.synchronize() + self._put(transfers) + except Exception as exc: + # Broad on purpose: this is the thread boundary. Anything that + # escapes here would be lost, so it is stashed and re-raised on + # the executor thread at the next connector call. + logger.error("mooncake-store save failed on rank %d: %s", self._rank, exc) + with self._save_lock: + if self._save_error is None: + self._save_error = exc + finally: + with self._save_lock: + for entry in transfers: + remaining = self._outstanding_saves.get(entry.request_id, 0) - 1 + if remaining <= 0: + self._outstanding_saves.pop(entry.request_id, None) + else: + self._outstanding_saves[entry.request_id] = remaining + + def _put(self, transfers: Sequence[RequestTransfers]) -> None: + keys, addresses, sizes, _ = self._resolve(transfers) + if not keys: + return + + for batch in zip( + _batched(keys, self._config.transfer_batch_size), + _batched(addresses, self._config.transfer_batch_size), + _batched(sizes, self._config.transfer_batch_size), + ): + batch_keys, batch_addresses, batch_sizes = batch + # Skip pages another rank or another instance already wrote. The + # scheduler cannot know this: it holds no store handle, and the + # answer changes between the time it builds metadata and now. + present = self._store.batch_is_exist(list(batch_keys)) + pending = [ + index + for index, status in enumerate(present) + if status != 1 # absent, or a failed probe we retry as a write + ] + if not pending: + continue + results = self._store.batch_put_from_multi_buffers( + [batch_keys[i] for i in pending], + [batch_addresses[i] for i in pending], + [batch_sizes[i] for i in pending], + ) + failures = sum(1 for result in results if not isinstance(result, int) or result < 0) + if failures: + # A dropped write only costs a future cache miss, so it is worth + # a warning rather than failing a request that already answered. + logger.warning( + "mooncake-store rank %d failed to save %d of %d pages", + self._rank, + failures, + len(pending), + ) + + # ---- shared ---- + + def _resolve( + self, transfers: Sequence[RequestTransfers] + ) -> Tuple[List[str], List[List[int]], List[List[int]], int]: + """Expand per-request page transfers into parallel store call arguments.""" + if self._addressing is None: + raise RuntimeError("KV cache layout has not been registered") + keys: List[str] = [] + addresses: List[List[int]] = [] + sizes: List[List[int]] = [] + pages = 0 + for entry in transfers: + for page in entry.pages: + namespace = self._namespaces.get(page.layer_group_id) + if namespace is None: + raise KeyError( + f"layer group {page.layer_group_id} is not in the registered " + "layout; the scheduler and worker disagree about the model" + ) + page_addresses, page_sizes = self._addressing.buffers( + page.layer_group_id, page.page_index + ) + keys.append(namespace.key(page.block_hash)) + addresses.append(page_addresses) + sizes.append(page_sizes) + pages += 1 + return keys, addresses, sizes, pages + + def _reraise_save_error(self) -> None: + with self._save_lock: + error = self._save_error + self._save_error = None + if error is not None: + raise RuntimeError("mooncake-store background save failed") from error + + def shutdown(self) -> None: + """Stop the save thread and release the store handle. Idempotent.""" + thread, self._save_thread = self._save_thread, None + if thread is not None: + self._save_queue.put(None) + thread.join(timeout=30.0) + store, self._store = self._store, None + if store is not None: + try: + store.close() + except Exception: + logger.warning("mooncake-store close failed", exc_info=True) + global _LOCAL_WORKER + if _LOCAL_WORKER is self: + _LOCAL_WORKER = None + _LOCAL_WORKER_READY.clear() + + def __del__(self): + try: + self.shutdown() + except Exception: # noqa: S110 - interpreter teardown, nothing left to report to + pass From b49a4b60bc302edeedbd812284a1e600b77cf131 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:13:07 -0700 Subject: [PATCH 03/24] [None][feat] Gate, register and test the mooncake-store connector Adds the startup gates, the 'mooncake-store' registry preset so the connector can be selected by name, and unit tests covering the pieces that decide whether a cache hit is correct. Every gate rejects a configuration whose failure mode is a wrong answer rather than a slow one. Context parallelism gives a rank a slice of the sequence instead of whole blocks, so one key would name different bytes per rank. Sliding-window attention makes a page's validity depend on where the window sits, which is a property of the reader rather than of the tokens. MiniMax-M3's index-V cache lives outside the paged pools, so a replayed prefix would pair stored index-K with stale index-V -- the same restriction disaggregated serving already applies. Pipeline parallelism is refused as untested rather than unsound. Beam search, attention DP, non-GPU cache tiers and Mamba caches are already rejected for all connectors in py_executor. Tests run without Mooncake or a GPU: the store is an in-process fake and the layout is synthesized from integers, which is all the addressing arithmetic needs. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/scheduler.py | 2 + .../connectors/mooncake_store/validation.py | 84 +++ .../connectors/mooncake_store/worker.py | 3 + .../_torch/pyexecutor/connectors/registry.py | 5 + tensorrt_llm/llmapi/llm_args.py | 3 +- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../executor/test_mooncake_store_connector.py | 714 ++++++++++++++++++ 7 files changed, 811 insertions(+), 1 deletion(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py create mode 100644 tests/unittest/_torch/executor/test_mooncake_store_connector.py diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py index 05d1f6e3551c..1096acc826d9 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -36,6 +36,7 @@ from .config import MooncakeStoreConnectorConfig from .keys import BlockHashChain from .metadata import MooncakeStoreMetadata, PageTransfer, RequestTransfers +from .validation import validate_llm_args from .worker import MooncakeStoreConnectorWorker, resolve_local_worker __all__ = ["MooncakeStoreConnectorScheduler"] @@ -74,6 +75,7 @@ class MooncakeStoreConnectorScheduler(KvCacheConnectorScheduler): def __init__(self, llm_args: TorchLlmArgs): super().__init__(llm_args) + validate_llm_args(llm_args) self._config = MooncakeStoreConnectorConfig.from_env() self._tokens_per_block = int(llm_args.kv_cache_config.tokens_per_block) self._requests: Dict[int, _RequestState] = {} diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py new file mode 100644 index 000000000000..8eb9a633b6ca --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Startup gates for the Mooncake store connector. + +Every rejection here is a configuration whose failure mode is a wrong answer +rather than a slow one: KV that gets replayed without all of the state it was +computed with. Beam search, attention data parallelism, host and disk cache +tiers, and Mamba caches are rejected for all connectors in ``py_executor``, so +they are not repeated. + +Checks run at construction, before any request is admitted, so a bad deployment +fails at startup instead of after the first cache hit. +""" + +from typing import TYPE_CHECKING + +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +if TYPE_CHECKING: + from ..kv_cache_layout import KvCacheLayout + +__all__ = ["validate_layout", "validate_llm_args"] + + +def validate_llm_args(llm_args: TorchLlmArgs) -> None: + """Reject parallel and model configurations this connector cannot serve.""" + if getattr(llm_args, "context_parallel_size", 1) > 1: + raise NotImplementedError( + "The mooncake-store connector does not support context parallelism. " + "A stored page is keyed by the tokens it holds, but under context " + "parallelism a rank holds a slice of the sequence rather than whole " + "blocks of it, so the same key would name different bytes on " + "different ranks." + ) + + if getattr(llm_args, "pipeline_parallel_size", 1) > 1: + raise NotImplementedError( + "The mooncake-store connector does not support pipeline parallelism. " + "Keys are namespaced per rank, so each stage would store only its own " + "layers and a prefix hit would require every stage to agree; that path " + "is untested. Run with tensor parallelism only." + ) + + sparse_config = getattr(llm_args, "sparse_attention_config", None) + if sparse_config is not None and not getattr(sparse_config, "sparse_disable_index_value", True): + raise NotImplementedError( + "The mooncake-store connector requires " + "sparse_attention_config.sparse_disable_index_value=True. The index-V " + "cache is a plain tensor outside the KV cache manager's paged pools, " + "so it is neither described to the connector nor transferred; a " + "replayed prefix would carry index-K from the store alongside stale " + "index-V. This is the same restriction disaggregated serving applies." + ) + + +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." + ) + + if not layout.groups: + raise ValueError( + "The KV cache layout describes no layer groups, so there is nothing " + "for the mooncake-store connector to transfer." + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 6a21f9a8178d..9dd65ed7d3fc 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -48,6 +48,7 @@ from .config import CONFIG_PATH_ENV, MooncakeStoreConnectorConfig from .keys import KeyNamespace from .metadata import MooncakeStoreMetadata, RequestTransfers +from .validation import validate_layout, validate_llm_args __all__ = ["MooncakeStoreConnectorWorker", "resolve_local_worker"] @@ -134,6 +135,7 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): def __init__(self, llm_args: TorchLlmArgs): super().__init__(llm_args) + validate_llm_args(llm_args) self._config = MooncakeStoreConnectorConfig.from_env() self._rank = mpi_rank() self._world_size = mpi_world_size() @@ -196,6 +198,7 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: if self._addressing is not None: raise RuntimeError("KV cache layout already registered") + validate_layout(layout) addressing = PageAddressing(layout) for start, end in addressing.registration_ranges(): status = self._store.register_buffer(start, end - start) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py index 9a00cdadd7fd..70dca022b763 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py @@ -35,4 +35,9 @@ "connector_scheduler_class": "DynamoKVBMConnectorLeader", "connector_worker_class": "DynamoKVBMConnectorWorker", }, + "mooncake-store": { + "connector_module": "tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + "connector_scheduler_class": "MooncakeStoreConnectorScheduler", + "connector_worker_class": "MooncakeStoreConnectorWorker", + }, } diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 7b5d014d3994..76693db6c02c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1959,7 +1959,8 @@ class KvCacheConnectorConfig(StrictBaseModel): description="Named connector preset (e.g. 'lmcache'). " "When set, connector_module/scheduler_class/worker_class are " "auto-populated from the preset registry.", - telemetry=TelemetryField.categorical('lmcache', 'lmcache-mp', 'kvbm')) + telemetry=TelemetryField.categorical('lmcache', 'lmcache-mp', 'kvbm', + 'mooncake-store')) connector_module: Optional[str] = Field( None, description= diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 083895fe6d59..db8d31642d2c 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -41,6 +41,7 @@ l0_a10: - unittest/_torch/executor/test_kv_cache_compression_manager.py - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_kv_cache_layout.py + - unittest/_torch/executor/test_mooncake_store_connector.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py - unittest/_torch/modules/dwdp/test_dwdp_manager.py diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py new file mode 100644 index 000000000000..dd507da4efd4 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -0,0 +1,714 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the Mooncake store KV cache connector. + +Runs without a Mooncake installation and without a GPU: the store handle is +replaced by an in-process fake, and the KV cache layout is synthesized from +plain integers, which is all the addressing arithmetic needs. +""" + +import json +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( + RequestData, + SchedulerOutput, +) +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import ( + KvCacheBufferRef, + KvCacheLayerGroupLayout, + KvCacheLayout, + KvCacheRegion, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import worker as worker_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.addressing import ( + PageAddressing, + merge_intervals, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + MooncakeStoreConnectorConfig, + StoreRole, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.keys import ( + BlockHashChain, + KeyNamespace, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.metadata import ( + PageTransfer, + RequestTransfers, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.scheduler import ( + MooncakeStoreConnectorScheduler, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.validation import ( + validate_layout, + validate_llm_args, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.worker import ( + MooncakeStoreConnectorWorker, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX + +TOKENS_PER_BLOCK = 4 + + +# ---- fixtures and fakes ---- + + +class FakeStore: + """Records calls and remembers which keys exist, nothing more.""" + + def __init__(self): + self.objects = set() + self.registered = [] + self.put_calls = [] + self.get_calls = [] + self.exist_calls = [] + self.closed = False + self.fail_gets_for = set() + #: Workers built against this store, torn down by the fixture. + self.workers = [] + + def register_buffer(self, address, size): + self.registered.append((address, size)) + return 0 + + def batch_is_exist(self, keys): + self.exist_calls.append(list(keys)) + return [1 if key in self.objects else 0 for key in keys] + + def batch_put_from_multi_buffers(self, keys, addresses, sizes, *_args, **_kwargs): + self.put_calls.append((list(keys), [list(a) for a in addresses], [list(s) for s in sizes])) + self.objects.update(keys) + return [sum(size) for size in sizes] + + def batch_get_into_multi_buffers(self, keys, addresses, sizes): + self.get_calls.append((list(keys), [list(a) for a in addresses], [list(s) for s in sizes])) + return [ + -1 if (key in self.fail_gets_for or key not in self.objects) else sum(size) + for key, size in zip(keys, sizes) + ] + + def close(self): + self.closed = True + + +def make_layout(*, num_groups=1, regions_per_group=1, num_slots=8, window_size=None): + """A layout whose regions are laid out back to back in a fake address space.""" + groups = [] + base = 0x1000 + for group_id in range(num_groups): + regions = [] + for region_id in range(regions_per_group): + size = 64 * (region_id + 1) + stride = size + regions.append( + KvCacheRegion( + base=base, + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=group_id, role="key"),), + ) + ) + base += stride * num_slots + groups.append( + KvCacheLayerGroupLayout( + layer_group_id=group_id, + layer_ids=(group_id,), + window_size=window_size, + regions=tuple(regions), + ) + ) + return KvCacheLayout(tokens_per_block=TOKENS_PER_BLOCK, groups=tuple(groups)) + + +@pytest.fixture +def store_config(tmp_path, monkeypatch): + path = tmp_path / "mooncake.json" + path.write_text( + json.dumps( + { + "metadata_server": "http://127.0.0.1:8080/metadata", + "master_server_address": "127.0.0.1:50051", + "protocol": "tcp", + "device_name": "", + "global_segment_size": "1GiB", + "local_buffer_size": "256MiB", + "model_key": "test-model", + } + ) + ) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_ROLE", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_PREFIX", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_MODEL_KEY", raising=False) + return path + + +def make_llm_args(): + return SimpleNamespace( + model="/models/test-model", + kv_cache_config=SimpleNamespace(tokens_per_block=TOKENS_PER_BLOCK), + tensor_parallel_size=1, + pipeline_parallel_size=1, + context_parallel_size=1, + sparse_attention_config=None, + ) + + +@pytest.fixture +def fake_store(monkeypatch): + """Replace the store handle, and tear down any worker a test builds.""" + store = FakeStore() + monkeypatch.setattr(worker_module, "_open_store", lambda _config: store) + yield store + for worker in store.workers: + worker.shutdown() + worker_module._LOCAL_WORKER = None + worker_module._LOCAL_WORKER_READY.clear() + + +def make_worker(fake_store, *, layout=None): + worker = MooncakeStoreConnectorWorker(make_llm_args()) + fake_store.workers.append(worker) + if layout is not None: + worker.register_kv_cache_layout(layout) + return worker + + +def make_request(request_id, tokens, cache_salt=None): + return SimpleNamespace( + request_id=request_id, + cache_salt=cache_salt, + get_tokens=lambda _beam=0, _tokens=tuple(tokens): list(_tokens), + ) + + +# ---- keys ---- + + +def test_hash_chain_is_deterministic_and_prefix_sensitive(): + tokens = list(range(3 * TOKENS_PER_BLOCK)) + first = list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + second = list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + assert first == second + + # Changing a token in block 0 must change every hash after it, which is what + # makes a key safe to share: a hit implies the whole prefix matched. + altered = list(tokens) + altered[0] += 1 + changed = list(BlockHashChain(TOKENS_PER_BLOCK).extend(altered)) + assert all(a != b for a, b in zip(first, changed)) + + +def test_hash_chain_ignores_partial_trailing_block(): + full = list(range(2 * TOKENS_PER_BLOCK)) + chain = BlockHashChain(TOKENS_PER_BLOCK) + assert len(chain.extend(full)) == 2 + assert len(chain.extend(full + [99])) == 2 + + +def test_hash_chain_extends_incrementally(): + tokens = list(range(4 * TOKENS_PER_BLOCK)) + incremental = BlockHashChain(TOKENS_PER_BLOCK) + for end in range(0, len(tokens) + 1, TOKENS_PER_BLOCK): + incremental.extend(tokens[:end]) + assert list(incremental.hashes) == list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + + +def test_hash_chain_separates_cache_salts(): + tokens = list(range(TOKENS_PER_BLOCK)) + unsalted = BlockHashChain(TOKENS_PER_BLOCK).extend(tokens) + salted = BlockHashChain(TOKENS_PER_BLOCK, cache_salt="tenant-a").extend(tokens) + other = BlockHashChain(TOKENS_PER_BLOCK, cache_salt="tenant-b").extend(tokens) + assert unsalted[0] != salted[0] != other[0] + assert salted[0] != other[0] + + +def test_hash_chain_rejects_shrinking_token_list(): + chain = BlockHashChain(TOKENS_PER_BLOCK) + chain.extend(list(range(2 * TOKENS_PER_BLOCK))) + with pytest.raises(ValueError, match="shrank"): + chain.extend(list(range(TOKENS_PER_BLOCK))) + + +def test_key_namespace_separates_every_dimension(): + base = dict( + cache_prefix="trtllm", + model_key="m", + rank=0, + world_size=2, + layer_group_id=0, + tokens_per_block=32, + bytes_per_page=1024, + ) + block_hash = b"\x01" * 16 + reference = KeyNamespace(**base).key(block_hash) + for field, value in [ + ("cache_prefix", "other"), + ("model_key", "n"), + ("rank", 1), + ("world_size", 4), + ("layer_group_id", 1), + ("tokens_per_block", 64), + ("bytes_per_page", 2048), + ]: + assert KeyNamespace(**{**base, field: value}).key(block_hash) != reference + + +# ---- addressing ---- + + +@pytest.mark.parametrize( + "intervals,expected", + [ + ([], []), + ([(0, 10)], [(0, 10)]), + ([(0, 10), (10, 20)], [(0, 20)]), + ([(0, 10), (5, 20)], [(0, 20)]), + ([(0, 10), (20, 30)], [(0, 10), (20, 30)]), + ([(20, 30), (0, 10)], [(0, 10), (20, 30)]), + ([(0, 100), (10, 20)], [(0, 100)]), + ([(0, 0), (5, 10)], [(5, 10)]), + ], +) +def test_merge_intervals(intervals, expected): + assert merge_intervals(intervals) == expected + + +def test_page_addressing_resolves_every_region_of_a_page(): + layout = make_layout(regions_per_group=3, num_slots=4) + addressing = PageAddressing(layout) + regions = layout.groups[0].regions + + addresses, sizes = addressing.buffers(0, 2) + assert sizes == [region.size for region in regions] + assert addresses == [region.base + region.stride * 2 for region in regions] + assert addressing.bytes_per_page(0) == sum(region.size for region in regions) + + +def test_page_addressing_rejects_out_of_range_page(): + addressing = PageAddressing(make_layout(num_slots=4)) + with pytest.raises(IndexError): + addressing.buffers(0, 4) + with pytest.raises(IndexError): + addressing.buffers(0, -1) + + +def test_page_addressing_registration_covers_every_slot_once(): + layout = make_layout(num_groups=2, regions_per_group=2, num_slots=4) + ranges = PageAddressing(layout).registration_ranges() + + # Regions were laid out back to back, so the whole span merges into one. + all_regions = [region for group in layout.groups for region in group.regions] + lowest = min(region.base for region in all_regions) + highest = max( + region.base + region.stride * (region.num_slots - 1) + region.size for region in all_regions + ) + assert ranges == [(lowest, highest)] + + +def test_page_addressing_rejects_mixed_slot_counts(): + region_a = KvCacheRegion(base=0, size=8, stride=8, num_slots=4, buffers=()) + region_b = KvCacheRegion(base=64, size=8, stride=8, num_slots=8, buffers=()) + layout = KvCacheLayout( + tokens_per_block=TOKENS_PER_BLOCK, + groups=( + KvCacheLayerGroupLayout( + layer_group_id=0, + layer_ids=(0,), + window_size=None, + regions=(region_a, region_b), + ), + ), + ) + with pytest.raises(ValueError, match="slot counts"): + PageAddressing(layout) + + +# ---- config ---- + + +def test_config_reads_sizes_with_units(store_config): + config = MooncakeStoreConnectorConfig.from_env() + assert config.global_segment_size == 1024**3 + assert config.local_buffer_size == 256 * 1024**2 + assert config.role is StoreRole.BOTH + assert config.resolve_model_key("/models/ignored") == "test-model" + + +def test_config_role_comes_from_environment(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "producer") + config = MooncakeStoreConnectorConfig.from_env() + assert config.role is StoreRole.PRODUCER + assert config.role.saves and not config.role.loads + + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "consumer") + config = MooncakeStoreConnectorConfig.from_env() + assert config.role.loads and not config.role.saves + + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "nonsense") + with pytest.raises(ValueError, match="TRTLLM_MOONCAKE_STORE_ROLE"): + MooncakeStoreConnectorConfig.from_env() + + +def test_config_requires_the_env_var(monkeypatch): + monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + with pytest.raises(ValueError, match="MOONCAKE_CONFIG_PATH"): + MooncakeStoreConnectorConfig.from_env() + + +def test_config_model_key_defaults_to_basename(store_config, tmp_path, monkeypatch): + path = tmp_path / "no_model_key.json" + path.write_text(json.dumps({"master_server_address": "127.0.0.1:50051"})) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + config = MooncakeStoreConnectorConfig.from_env() + assert config.resolve_model_key("/models/MiniMax-M3/") == "MiniMax-M3" + + +# ---- validation ---- + + +@pytest.mark.parametrize( + "field,value,match", + [ + ("context_parallel_size", 2, "context parallelism"), + ("pipeline_parallel_size", 2, "pipeline parallelism"), + ], +) +def test_validate_llm_args_rejects_unsupported_parallelism(field, value, match): + args = make_llm_args() + setattr(args, field, value) + with pytest.raises(NotImplementedError, match=match): + validate_llm_args(args) + + +def test_validate_llm_args_rejects_m3_index_value_cache(): + args = make_llm_args() + args.sparse_attention_config = SimpleNamespace(sparse_disable_index_value=False) + with pytest.raises(NotImplementedError, match="sparse_disable_index_value"): + validate_llm_args(args) + + args.sparse_attention_config = SimpleNamespace(sparse_disable_index_value=True) + validate_llm_args(args) + + +def test_validate_layout_rejects_sliding_window(): + with pytest.raises(NotImplementedError, match="sliding-window"): + validate_layout(make_layout(window_size=1024)) + validate_layout(make_layout()) + + +# ---- worker ---- + + +def test_worker_registers_every_pool_range(store_config, fake_store): + layout = make_layout(num_groups=2, regions_per_group=2) + worker = make_worker(fake_store, layout=layout) + assert fake_store.registered == [ + (start, end - start) for start, end in PageAddressing(layout).registration_ranges() + ] + assert worker.is_registered + + +def test_worker_rejects_v1_pool_registration(store_config, fake_store): + worker = make_worker(fake_store) + with pytest.raises(NotImplementedError, match="KVCacheManagerV2"): + worker.register_kv_caches(None) + + +def test_worker_prefix_hit_needs_every_layer_group(store_config, fake_store): + layout = make_layout(num_groups=2) + worker = make_worker(fake_store, layout=layout) + hashes = [bytes([index]) * 16 for index in range(3)] + + assert worker.count_prefix_hit(hashes) == 0 + + # Populate blocks 0 and 1 completely, and block 2 only partially. + for block in range(2): + for group_id in range(2): + fake_store.objects.add(worker._namespaces[group_id].key(hashes[block])) + fake_store.objects.add(worker._namespaces[0].key(hashes[2])) + + assert worker.count_prefix_hit(hashes) == 2 + + +def test_worker_prefix_hit_stops_at_the_first_gap(store_config, fake_store): + worker = make_worker(fake_store, layout=make_layout()) + hashes = [bytes([index]) * 16 for index in range(3)] + # Block 1 missing: block 2 is unusable even though it is present, because a + # prefix is replayed contiguously. + fake_store.objects.add(worker._namespaces[0].key(hashes[0])) + fake_store.objects.add(worker._namespaces[0].key(hashes[2])) + assert worker.count_prefix_hit(hashes) == 1 + + +def test_worker_load_raises_when_a_page_is_missing(store_config, fake_store): + worker = make_worker(fake_store, layout=make_layout()) + transfers = RequestTransfers(7, [PageTransfer(b"\x00" * 16, 0, 1)]) + worker.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) + with pytest.raises(RuntimeError, match="already"): + worker.start_load_kv(None) + + +def test_worker_load_addresses_the_requested_page(store_config, fake_store): + layout = make_layout(regions_per_group=2) + worker = make_worker(fake_store, layout=layout) + block_hash = b"\x00" * 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(SimpleNamespace(loads=[transfers], saves=[])) + worker.start_load_kv(None) + + (keys, addresses, sizes) = fake_store.get_calls[0] + expected_addresses, expected_sizes = PageAddressing(layout).buffers(0, 3) + assert keys == [key] + assert addresses == [expected_addresses] + assert sizes == [expected_sizes] + + +def test_worker_save_skips_pages_already_in_the_store(store_config, fake_store): + worker = make_worker(fake_store, layout=make_layout()) + hashes = [bytes([index]) * 16 for index in range(2)] + fake_store.objects.add(worker._namespaces[0].key(hashes[0])) + + worker._put( + [ + RequestTransfers( + 1, + [PageTransfer(hashes[0], 0, 0), PageTransfer(hashes[1], 0, 1)], + ) + ] + ) + assert len(fake_store.put_calls) == 1 + assert fake_store.put_calls[0][0] == [worker._namespaces[0].key(hashes[1])] + + +def test_worker_reports_a_request_finished_once_its_saves_drain(store_config, fake_store): + worker = make_worker(fake_store, layout=make_layout()) + + # One submission outstanding: the request is closed but must not be released. + worker._outstanding_saves[42] = 1 + assert worker.get_finished([42], []) == ([], []) + + worker._outstanding_saves.pop(42) + assert worker.get_finished([], []) == ([42], []) + # Reported once only. + assert worker.get_finished([], []) == ([], []) + + +def test_worker_reports_a_request_with_no_saves_immediately(store_config, fake_store): + worker = make_worker(fake_store, layout=make_layout()) + assert worker.get_finished([9], [5]) == ([9], [5]) + + +def test_worker_shutdown_closes_the_store(store_config, fake_store): + worker = make_worker(fake_store, layout=make_layout()) + worker.shutdown() + assert fake_store.closed + worker.shutdown() + + +# ---- scheduler ---- + + +class FakeWorker: + """Stands in for the process-local worker's lookup service.""" + + def __init__(self, hit_blocks=0): + self.hit_blocks = hit_blocks + self.queries = [] + + def count_prefix_hit(self, block_hashes): + self.queries.append(list(block_hashes)) + return min(self.hit_blocks, len(block_hashes)) + + +def make_scheduler(store_config, hit_blocks=0): + scheduler = MooncakeStoreConnectorScheduler(make_llm_args()) + scheduler._worker = FakeWorker(hit_blocks) + return scheduler + + +def request_data(request_id, new_tokens, page_indices, layer_group_id=0): + return RequestData( + request_id=request_id, + new_tokens=list(new_tokens), + new_block_ids=list(page_indices), + computed_position=0, + num_scheduled_tokens=len(new_tokens), + new_block_ids_by_layer_group={layer_group_id: list(page_indices)}, + ) + + +def test_scheduler_offers_the_stored_prefix(store_config): + scheduler = make_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (2 * TOKENS_PER_BLOCK, False) + + +def test_scheduler_never_offers_the_whole_prompt(store_config): + scheduler = make_scheduler(store_config, hit_blocks=99) + # Exactly three full blocks: the last one is withheld so the runtime still + # has a token to run a forward pass on. + request = make_request(1, list(range(3 * TOKENS_PER_BLOCK))) + matched, _ = scheduler.get_num_new_matched_tokens(request, 0) + assert matched == 2 * TOKENS_PER_BLOCK + + +def test_scheduler_declines_partial_local_matches(store_config): + scheduler = make_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, TOKENS_PER_BLOCK + 1) == (0, False) + + +def test_scheduler_offers_nothing_as_a_producer(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "producer") + scheduler = make_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False) + assert scheduler._worker.queries == [] + + +def test_scheduler_skips_local_prefix_when_looking_up(store_config): + scheduler = make_scheduler(store_config, hit_blocks=1) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 2 * TOKENS_PER_BLOCK) + # Blocks 0 and 1 are on device already; candidates start at block 2 and stop + # short of the final block. + full_chain = list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + assert scheduler._worker.queries[0] == full_chain[2:5] + + +def test_scheduler_builds_loads_for_the_offered_blocks(store_config): + scheduler = make_scheduler(store_config, hit_blocks=2) + tokens = list(range(5 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + output = SchedulerOutput(new_requests=[request_data(1, tokens, [10, 11, 12, 13, 14])]) + metadata = scheduler.build_connector_meta(output) + + assert [page.page_index for page in metadata.loads[0].pages] == [10, 11] + # Blocks 0 and 1 came from the store, so only blocks 2..4 are written back. + assert [page.page_index for page in metadata.saves[0].pages] == [12, 13, 14] + + +def test_scheduler_does_not_resave_blocks_across_iterations(store_config): + scheduler = make_scheduler(store_config, hit_blocks=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] + + # A generation step completes one more block; only that block is saved. + 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 [page.page_index for page in second.saves[0].pages] == [6] + + +def test_scheduler_waits_for_a_block_to_fill_before_saving(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(TOKENS_PER_BLOCK + 1)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + metadata = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])]) + ) + # Page 5 holds a single token, so only the full block is offered up. + assert [page.page_index for page in metadata.saves[0].pages] == [4] + + +def test_scheduler_saves_nothing_as_a_consumer(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "consumer") + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + metadata = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])]) + ) + assert metadata.saves == [] + + +def test_scheduler_skips_blocks_without_a_page_in_every_group(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + data = request_data(1, tokens, [4, 5]) + data.new_block_ids_by_layer_group[1] = [7, BAD_PAGE_INDEX] + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + + # Block 1 has no page in group 1, so neither of its halves is stored; block 0 + # contributes one page per group. + assert [(page.layer_group_id, page.page_index) for page in metadata.saves[0].pages] == [ + (0, 4), + (1, 7), + ] + + +def test_scheduler_cancel_load_truncates_the_offer(store_config): + scheduler = make_scheduler(store_config, hit_blocks=3) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + # The runtime will not consume anything from block 1 onwards. + scheduler.cancel_load(request, TOKENS_PER_BLOCK, 6 * TOKENS_PER_BLOCK) + metadata = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, list(range(10, 16)))]) + ) + assert [page.page_index for page in metadata.loads[0].pages] == [10] + + +def test_scheduler_request_finished_pins_pages_only_when_saving(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + scheduler.build_connector_meta(SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])])) + assert scheduler.request_finished(request, [4, 5]) is True + # State is dropped with the request, so a second call reports nothing pending. + assert scheduler.request_finished(request, [4, 5]) is False + + +def test_scheduler_request_finished_is_false_without_saves(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + request = make_request(1, list(range(TOKENS_PER_BLOCK - 1))) + scheduler.get_num_new_matched_tokens(request, 0) + assert scheduler.request_finished(request, []) is False + + +def test_scheduler_isolates_requests_by_cache_salt(store_config): + scheduler = make_scheduler(store_config, hit_blocks=1) + tokens = list(range(3 * TOKENS_PER_BLOCK)) + scheduler.get_num_new_matched_tokens(make_request(1, tokens, cache_salt="a"), 0) + scheduler.get_num_new_matched_tokens(make_request(2, tokens, cache_salt="b"), 0) + assert scheduler._worker.queries[0] != scheduler._worker.queries[1] From ead8bd090336c752a486581d9b95d323f0b74f0f Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:16:36 -0700 Subject: [PATCH 04/24] [None][doc] Document the mooncake-store connector Describe what the store buys over local block reuse, how pages are keyed and which configurations are refused, so an operator can tell the store apart from the similarly named transfer engine used for prefill/decode handoff. Ship a trtllm-serve config as a starting point. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docs/source/features/kv-cache-connector.md | 85 +++++++++++++++++++ ...trtllm_mooncake_store_connector_extra.yaml | 49 +++++++++++ 2 files changed, 134 insertions(+) create mode 100644 examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 89ae82287b88..c4f82b71e116 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -64,6 +64,11 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **Description**: Called at initialization. Provides the worker with the GPU KV cache tensors. * **Arguments**: `kv_cache_tensor` is the underlying storage tensor for the KV cache. +* **`register_kv_cache_layout(self, layout: KvCacheLayout)`** + * **Description**: Called at initialization **instead of** `register_kv_caches` when the KV cache manager is `KVCacheManagerV2`, whose memory cannot be expressed as one tensor: there is one slot address space per pool and one page-index space per layer group. The default implementation raises, so a connector that does not implement it can only run on V1. + * **Arguments**: `layout` describes the byte ranges that repeat per page slot. Each `KvCacheLayerGroupLayout` carries a tuple of `KvCacheRegion`s, and the bytes for page slot `i` of a region live at `region.base + region.stride * i` for `region.size` bytes -- or equivalently at `region.as_tensor()[i]`. Page indices arriving in `RequestData.new_block_ids_by_layer_group` are scoped to a layer group and index that group's regions. + * **Why regions rather than a tensor**: because the ranges are described rather than implied, the same structure covers MLA (a pool simply has no `value` buffer), sliding-window and hybrid models (one layer group per window size), and non-uniform slots such as MiniMax-M3's index-K buffer sitting beside K/V, without any of them being a special case. + * **`start_load_kv(self, stream: torch.cuda.Stream)`** * **Description**: Initiates the loading of KV blocks from the external source into the GPU memory. * **Arguments**: `stream` is the CUDA stream where the forward pass is executed in. @@ -81,6 +86,86 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **Description**: Polled by the runtime to check the status of asynchronous operations. * **Returns**: Two lists of request IDs: those that have finished saving, and those that have finished loading. +## Built-in Connectors + +Named presets can be selected without naming a module or class: + +```python +from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig + +kv_connector_config = KvCacheConnectorConfig(connector="mooncake-store") +``` + +The available presets are `lmcache`, `lmcache-mp`, `kvbm` and `mooncake-store`. The first three are external packages; `mooncake-store` ships with TensorRT-LLM and is described below. + +### Mooncake distributed store (`mooncake-store`) + +Publishes KV pages into a [Mooncake](https://github.com/kvcache-ai/Mooncake) store -- a shared CPU memory pool addressed by content -- so a prefix computed by one engine can be replayed by another. Regular block reuse cannot do this, because it never leaves the instance that computed the prefix. + +This is a **different component** from the Mooncake transfer engine that the C++ cache transceiver uses for disaggregated prefill/decode handoff. That moves KV point to point between two known peers; this publishes pages into a pool that any peer can read. The two compose: a context server can write pages into the store and still hand off to a generation server over NIXL. + +#### Requirements + +* `KVCacheManagerV2` (`kv_cache_config.use_kv_cache_manager_v2: true`), since that is the manager that can describe its pools through `register_kv_cache_layout`. +* The Mooncake Python bindings: `pip install mooncake-transfer-engine`. These are installed in the release container; the source build of the C++ transfer engine does not provide them. +* A running Mooncake master (and metadata server, unless using `P2PHANDSHAKE`). See the [Mooncake documentation](https://kvcache-ai.github.io/Mooncake/). +* GPU-only KV cache tiers: set `kv_cache_config.host_cache_size: 0` and `disk_cache_size: 0`. A page evicted to another tier has its GPU slot reassigned, which would invalidate the addresses registered with the store. + +#### Configuration + +Topology comes from a JSON file named by `MOONCAKE_CONFIG_PATH`, using the same schema as the vLLM Mooncake store connector so one deployment can point both engines at the same pool: + +```json +{ + "metadata_server": "http://127.0.0.1:8080/metadata", + "master_server_address": "127.0.0.1:50051", + "protocol": "rdma", + "device_name": "mlx5_0", + "global_segment_size": "32GiB", + "local_buffer_size": "1GiB" +} +``` + +Two further settings are TensorRT-LLM's rather than Mooncake's, and are read from the environment because `KvCacheConnectorConfig` carries no free-form dictionary: + +| Variable | Default | Meaning | +|---|---|---| +| `TRTLLM_MOONCAKE_STORE_ROLE` | `both` | `producer` writes only, `consumer` reads only, `both` does both. | +| `TRTLLM_MOONCAKE_STORE_PREFIX` | `trtllm` | Leading component of every key, for isolating deployments that share a pool. | +| `TRTLLM_MOONCAKE_STORE_MODEL_KEY` | model directory basename | Identity keys are namespaced by. Two engines share cache only when they agree on it, so the default is the basename rather than the full path -- the same checkpoint is routinely mounted elsewhere on another host, which is exactly what sharing is for. | + +In a disaggregated deployment, run context servers as `both` and leave generation servers unconfigured. Generated tokens are rarely a reused prefix, so writing them costs bandwidth for no hit rate. + +#### How it keys pages + +`KVCacheManagerV2` reports `RequestData.block_hashes` empty, so the connector derives block identity itself: a blake2b chain where each block's hash covers its own tokens *and* every token before it, seeded by the request's `cache_salt`. A key is `//wr/lg/tb/`. The namespace pins down everything that would make the stored bytes mean something different, so a mismatched shard count, layer group or page geometry reads as a cache miss rather than as garbage. + +The value for one key is the concatenation of that layer group's regions for one page slot, handed to Mooncake's multi-buffer batch APIs as a list of `(address, size)` pairs. + +#### Transfer behavior + +* **Loads are synchronous**, performed in `start_load_kv` before the forward pass. A failed load raises: the runtime has already counted those tokens as computed, so a partial load is a wrong answer rather than a slow one. +* **Saves are asynchronous**, handed to a background thread behind a CUDA event recorded on the forward stream. The pages are only complete once the pass that wrote them retires, and blocking the executor loop on an RDMA write is the cost the store exists to avoid. The leader reports such requests as saving asynchronously, so their pages stay pinned until `get_finished` confirms the writes landed. A dropped save is logged rather than raised -- it only costs a future cache miss. +* Pages the store already holds are skipped, so several ranks or instances converging on the same prefix write it once. + +#### Unsupported configurations + +These are rejected at startup, before any request is admitted: + +| Configuration | Reason | +|---|---| +| Context parallelism | A rank holds a slice of the sequence rather than whole blocks of it, so one key would name different bytes on different ranks. | +| Sliding-window attention / VSWA | A page's validity depends on where the window sits, which is a property of the request that read it rather than of the tokens it holds. | +| MiniMax-M3 with `sparse_disable_index_value: false` | The index-V cache is a plain tensor outside the paged pools, so a replayed prefix would pair stored index-K with stale index-V. Disaggregated serving applies the same restriction. | +| Pipeline parallelism | Untested rather than unsound. Use tensor parallelism. | +| `KVCacheManagerV1` | Identity here is a per-layer-group hash chain; V1 supplies real block hashes over a single flat block space. | + +Beam search, attention data parallelism, non-GPU cache tiers and Mamba caches are rejected for all connectors by the executor. + +#### Example + +`examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml` is a starting point for `trtllm-serve`. + ## Example Implementation The file `examples/llm-api/llm_kv_cache_connector.py` provides a reference implementation of a **Persistent KV Cache**. diff --git a/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml new file mode 100644 index 000000000000..85e8530de354 --- /dev/null +++ b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml @@ -0,0 +1,49 @@ +# Extra LLM API options for trtllm-serve with the Mooncake store KV connector. +# +# Offloads KV pages to a Mooncake distributed store -- a shared CPU memory pool +# addressed by content -- so a prefix computed by one engine can be replayed by +# another. This is a different component from the Mooncake transfer engine used +# by the C++ cache transceiver for disaggregated prefill/decode handoff; the two +# compose rather than conflict. +# +# Prerequisites: +# - Mooncake Python bindings: pip install mooncake-transfer-engine +# (present in the release container; the C++ source build does not +# provide MooncakeDistributedStore) +# - A running Mooncake master, and a metadata server unless using +# P2PHANDSHAKE. See https://kvcache-ai.github.io/Mooncake/ +# - MOONCAKE_CONFIG_PATH pointing at a Mooncake JSON config, for example: +# { +# "metadata_server": "http://127.0.0.1:8080/metadata", +# "master_server_address": "127.0.0.1:50051", +# "protocol": "rdma", +# "device_name": "mlx5_0", +# "global_segment_size": "32GiB", +# "local_buffer_size": "1GiB" +# } +# +# Optional environment overrides: +# TRTLLM_MOONCAKE_STORE_ROLE producer | consumer | both (default both) +# TRTLLM_MOONCAKE_STORE_PREFIX key prefix, to isolate deployments sharing +# one pool (default trtllm) +# TRTLLM_MOONCAKE_STORE_MODEL_KEY identity keys are namespaced by +# (default: model directory basename) +# +# Example: +# export MOONCAKE_CONFIG_PATH=/path/to/mooncake.json +# trtllm-serve --backend pytorch --host 0.0.0.0 --port 8000 \ +# --extra_llm_api_options /path/to/this/file + +kv_cache_config: + # The connector describes its pools through register_kv_cache_layout, which + # only KVCacheManagerV2 implements. + use_kv_cache_manager_v2: true + # Local reuse still runs first; the store serves whatever the device missed. + enable_block_reuse: true + # GPU-only tiers are required: a page evicted to host or disk has its GPU slot + # reassigned, which would invalidate the addresses registered with the store. + host_cache_size: 0 + disk_cache_size: 0 + +kv_connector_config: + connector: mooncake-store From 8a32553241786973eebb748758a7069679e66f16 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:19:35 -0700 Subject: [PATCH 05/24] [None][chore] Add mooncake-store to the LLM args golden manifest Follows the telemetry allowlist gaining the new connector preset. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/usage/llm_args_golden_manifest.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index cc9b5a7ff7d6..1345c2af1e2e 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -798,7 +798,8 @@ "allowed_values": [ "lmcache", "lmcache-mp", - "kvbm" + "kvbm", + "mooncake-store" ], "annotation": "Optional[str]", "converter": "allowlist", From 8bf20594edd022641a38e86359cf52602e3fe761 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 27 Aug 2026 18:14:14 +0000 Subject: [PATCH 06/24] save changes to install process Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/common/install_mooncake.sh | 72 +- .../slurm/benchmark/disaggr_torch.slurm | 25 + mooncake_disagg/README.md | 784 ++++++++++++++++++ mooncake_disagg/ctx_config.yaml | 30 + mooncake_disagg/disagg_config.yaml | 18 + mooncake_disagg/gen_config.yaml | 22 + mooncake_disagg/install_mooncake_runtime.sh | 135 +++ mooncake_disagg/m3_agg_mooncake.yaml | 89 ++ mooncake_disagg/m3_ctx_mooncake.yaml | 68 ++ mooncake_disagg/m3_disagg_config.yaml | 20 + mooncake_disagg/m3_gen_mooncake.yaml | 68 ++ mooncake_disagg/mooncake.json | 11 + mooncake_disagg/mooncake_api_surface_test.py | 116 +++ mooncake_disagg/mooncake_smoke_test.py | 67 ++ .../usage/llm_args_golden_manifest.json | 7 + .../executor/test_mooncake_store_connector.py | 156 ++-- 16 files changed, 1609 insertions(+), 79 deletions(-) create mode 100644 mooncake_disagg/README.md create mode 100644 mooncake_disagg/ctx_config.yaml create mode 100644 mooncake_disagg/disagg_config.yaml create mode 100644 mooncake_disagg/gen_config.yaml create mode 100755 mooncake_disagg/install_mooncake_runtime.sh create mode 100644 mooncake_disagg/m3_agg_mooncake.yaml create mode 100644 mooncake_disagg/m3_ctx_mooncake.yaml create mode 100644 mooncake_disagg/m3_disagg_config.yaml create mode 100644 mooncake_disagg/m3_gen_mooncake.yaml create mode 100644 mooncake_disagg/mooncake.json create mode 100644 mooncake_disagg/mooncake_api_surface_test.py create mode 100644 mooncake_disagg/mooncake_smoke_test.py diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index d648be2f60a6..ccd5e04de1d9 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -51,10 +51,68 @@ rm -rf Mooncake echo "export LD_LIBRARY_PATH=${MOONCAKE_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" -# The source build above only produces the C++ transfer engine, which is what -# the cache transceiver links against. MooncakeDistributedStore -- the shared -# CPU pool behind the mooncake-store KV cache connector -- is only reachable -# through the Python bindings, and those are not part of the CMake install. Take -# them from the wheel at the same upstream version so the store client and the -# transfer engine cannot drift apart. -pip3 install --no-cache-dir "mooncake-transfer-engine==${MOONCAKE_VERSION#v}" +# The source build above is only useful for the C++ transfer engine, which is +# what the cache transceiver links against. MooncakeDistributedStore -- the +# shared CPU pool behind the mooncake-store KV cache connector -- comes from the +# Python wheel instead, for two reasons. +# +# First, `make install` does emit a `mooncake` Python package, but an unusable +# one: it omits libmooncake_store.so, so importing mooncake.store raises +# ImportError. It must be deleted, and deleting it is not optional in either of +# the two places it can land. +# +# mooncake-integration/CMakeLists.txt chooses its install directory with +# python3 -c "import sys; print([s for s in sys.path if 'packages' in s][0])" +# i.e. the first sys.path entry whose name merely contains "packages". +# +# - With nvidia-cutlass-dsl installed (the normal case here: the devel stage +# removes it, then constraints.txt reinstalls it), that first match is +# nvidia_cutlass_dsl/dsl_packages, because nvidia_cutlass_dsl_packages.pth +# does sys.path.insert(0) on it. The broken package then outranks +# dist-packages on every interpreter start, so no amount of pip installing +# can fix the import. CUTLASS DSL does not reference `mooncake` at all, so +# removing it is safe. +# - Without it, the match is dist-packages itself, and the broken package +# collides with the wheel. That is the more insidious case: CMake writes +# store.cpython-312-x86_64-linux-gnu.so while the wheel writes store.so, and +# importlib prefers the interpreter-tagged suffix, so the broken extension +# still wins even after pip reports success. +# +# Remove the package outright wherever it landed, before pip installs the real +# one. Nothing legitimate owns a `mooncake` package at this point, and removing +# rather than trying to identify individual leftovers keeps this correct in the +# dist-packages case, where pip would overwrite __init__.py and leave no marker +# to key on. +python3 - <<'PY' +import os +import shutil +import sys +import sysconfig + +paths = sysconfig.get_paths() +for entry in list(sys.path) + [paths["purelib"], paths["platlib"]]: + if not entry: + continue + package = os.path.join(entry, "mooncake") + if os.path.isdir(package): + print(f"removing CMake-generated mooncake package: {package}") + shutil.rmtree(package, ignore_errors=True) +PY + +# Second, the `mooncake-transfer-engine` wheel is built against CUDA 12 and +# these images ship CUDA 13 only, so its extensions cannot resolve +# libcudart.so.12. `mooncake-transfer-engine-cuda13` is the same project built +# for CUDA 13. It is versioned independently and its releases start at 0.3.9, so +# it cannot track MOONCAKE_VERSION above; the store client only has to agree +# with the mooncake_master it connects to, and the wheel supplies both. +MOONCAKE_WHEEL_VERSION="0.3.13" +pip3 install --no-cache-dir "mooncake-transfer-engine-cuda13==${MOONCAKE_WHEEL_VERSION}" + +# Fail the build rather than ship an image whose import is broken. +python3 - <<'PY' +from mooncake.store import MooncakeDistributedStore +import mooncake.store + +MooncakeDistributedStore() +print(f"mooncake.store OK: {mooncake.store.__file__}") +PY diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm index d73430ba5820..4472b07b2ea8 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm @@ -138,6 +138,31 @@ else echo "TensorRT-LLM environment variables saved to ${full_logdir}/env_vars.json" fi +# Install the Mooncake Python store bindings, but only when a worker config asks +# for the mooncake-store KV connector. The bindings shipped in the container +# images are unusable (see mooncake_disagg/README.md section 2), and the fix has +# to be reapplied per job: --container-name gives each node a container that +# lives for the job, so anything installed here survives to the worker sruns but +# not into the next job. +if grep -qs "mooncake-store" "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"; then + mooncake_install_script="" + if [ -n "${trtllm_repo:-}" ]; then + mooncake_install_script="${trtllm_repo}/mooncake_disagg/install_mooncake_runtime.sh" + fi + if [ -z "${mooncake_install_script}" ] || [ ! -f "${mooncake_install_script}" ]; then + cleanup_on_failure "A worker config requests the mooncake-store connector, but mooncake_disagg/install_mooncake_runtime.sh was not found under trtllm_repo='${trtllm_repo:-}'. Set environment.trtllm_repo to a checkout that contains it, or bake the bindings into the container image." + fi + echo "Installing Mooncake store bindings on all nodes..." + if ! srun --container-name=${container_name} \ + --container-mounts=${container_mount} --no-container-mount-home \ + --mpi=pmix --overlap -N $SLURM_NNODES --ntasks-per-node=1 \ + bash -c "MOONCAKE_WHEEL='${MOONCAKE_WHEEL:-}' bash ${mooncake_install_script}" \ + &> ${full_logdir}/2_install_mooncake.log; then + cleanup_on_failure "Mooncake store bindings installation failed. Check ${full_logdir}/2_install_mooncake.log for details" + fi + echo "Mooncake store bindings installation completed successfully" +fi + # Get node lists and replace the placeholder with the actual node names echo "SLURM_NODELIST: ${SLURM_NODELIST}" all_nodes=($(scontrol show hostname $SLURM_NODELIST | sort)) diff --git a/mooncake_disagg/README.md b/mooncake_disagg/README.md new file mode 100644 index 000000000000..f34b3a042ec1 --- /dev/null +++ b/mooncake_disagg/README.md @@ -0,0 +1,784 @@ +# Validating the mooncake-store KV connector on MiniMax-M3 + +This is a runbook for testing the `mooncake-store` KV cache connector +(commits `f3c092187e`..`b8d3f43c43`) on MiniMax-M3 under load, using the SLURM +disaggregated benchmark harness in +`examples/disaggregated/slurm/benchmark/`. + +The unit tests in `tests/unittest/_torch/executor/test_mooncake_store_connector.py` +cover the pieces that decide whether a cache hit is *correct* -- key namespacing, +hash chaining, page addressing, the startup gates. They deliberately do not run +a store, a model, or two engines. What is untested is everything that decides +whether the feature is *worth having*: whether a real prefix survives the round +trip, whether one M3 instance can replay a prefix another computed, and whether +the synchronous load path costs less than the prefill it avoids. + +## 1. The claim under test + +Local block reuse never leaves the instance that computed the prefix. The store +publishes KV pages into a shared, content-addressed CPU pool so any engine can +replay them. So there are exactly three things the store can do that local reuse +cannot, and each gets its own experiment in §7: + +1. **Cross-instance reuse.** A request routed to context instance B replays a + prefix computed on instance A. +2. **Survival across restarts.** Pages outlive the process that wrote them. +3. **Pool capacity beyond one host.** The pool is the sum of every worker's + segment rather than one node's host memory. + +### Why M3 makes this a hard test rather than an easy one + +The residency measurements in `../m3-kv-residency-measurement-README.md` (taken +on this same model and workload) found that M3 production traffic already serves +**97.0% of prompt tokens from local cache**, with eviction responsible for under +0.7% of misses. On a *single* instance there is almost no headroom for the store +to recover -- over 99% of misses are prefixes never cached anywhere, which no +store can serve either. + +That is not an argument against the feature; it is an argument about where to +look. Set expectations accordingly: + +- Do not expect a single-instance hit-rate improvement. Expect roughly zero. +- The store's value on M3 is concentrated in the cross-instance and + post-restart cases, where local reuse scores zero by construction. +- The connector's loads are **synchronous** (`start_load_kv`, before the forward + pass), so every loaded byte is fully exposed to TTFT. The host-offload tier it + replaces achieved 38-43% overlap at 45-51 GiB/s. At M3's context lengths a + loaded prefix is gigabytes, so a store hit is a win only when it displaces + real prefill, and a *needless* store hit is pure added latency. +- The design already rules out the worst version of that: the leader is handed + `num_computed_tokens` (the local match) and offers only blocks *beyond* it, so + the store cannot re-fetch something the GPU already holds. It cannot regress a + local hit; it can only add latency on a genuine local miss that it then fails + to make cheaper. That is what experiment 4 measures. + +## 2. Prerequisites: is Mooncake actually installed? + +**Short answer: probably not the part this connector needs.** Check before you +burn an allocation. + +Two different Mooncake components exist, and TensorRT-LLM's history treats them +differently: + +| Component | What uses it | How it gets installed | Since | +|---|---|---|---| +| C++ transfer engine (`/usr/local/Mooncake`) | the C++ cache transceiver's Mooncake backend | CMake source build in `docker/common/install_mooncake.sh` | PR #8447, Nov 2025 | +| Python bindings (`mooncake.store.MooncakeDistributedStore`) | **this connector** | pip wheel, added to the same script | commit `d36dae435e`, **this branch** | + +So Mooncake has been in the container images for months, but only usefully as +the C++ library. Three consequences: + +- Any image built before commit `d36dae435e` lacks a working set of bindings. + The image pinned in `jenkins/current_image_tags.properties` is tagged + `202607211045` (2026-07-21), which predates that commit, so **the currently + pinned CI image does not have them**. +- The CMake install *does* drop a `mooncake` Python package into the image, but + it is both broken and actively harmful: it shadows the working one. This is the + single biggest time sink in this section; see "Fix" below. +- The wheel is also the only usable source of the `mooncake_master` and + `mooncake_http_metadata_server` entry points. The CMake build's + `/usr/local/Mooncake/bin/mooncake_master` does exist and does run, but it is + the 0.3.7 build and it is not what ends up on `PATH` once the wheel is + installed. + +Also note `install_mooncake.sh` runs only in the `tritondevel` stage of +`docker/Dockerfile.multi`, and is skipped entirely on Rocky8. The CI image and +the internal `trtllm_build` release image descend from `tritondevel`, so they +get it; the NGC `release` image descends from the plain `devel` stage, so it +does not. + +### Verify + +Inside the container you will actually run: + +```bash +python3 -c "from mooncake.store import MooncakeDistributedStore; print('store bindings OK')" +command -v mooncake_master || ls /usr/local/Mooncake/bin +``` + +### Fix + +Run `mooncake_disagg/install_mooncake_runtime.sh` inside the container. It takes +under ten seconds on a warm pip cache, it is idempotent, and it verifies itself, +so it is safe in a SLURM prolog on every node. + +```bash +bash mooncake_disagg/install_mooncake_runtime.sh +``` + +A bare `pip3 install mooncake-transfer-engine` is *not* enough, and its failure +mode is what makes this step so confusing: pip reports success and the import +still fails. Two independent problems, both of which the script handles. + +**1. A broken `mooncake` package that pip cannot displace.** The CMake build in +`install_mooncake.sh` emits its own `mooncake` Python package, omitting +`libmooncake_store.so`, so it cannot load. `mooncake-integration/CMakeLists.txt` +chooses where to put it with: + +```cmake +COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print([s for s in sys.path if 'packages' in s][0])" +``` + +-- the *first* `sys.path` entry whose name merely contains `"packages"`. That +gives two different failures depending on what else is installed, and both +produce the same confusing symptom: an `ImportError` **after a `pip install` +that reported success**. + +- **With `nvidia-cutlass-dsl` present** (the normal case: the `devel` stage + uninstalls it at `Dockerfile.multi:58`, then `constraints.txt` pulls it back + in) the first match is `nvidia_cutlass_dsl/dsl_packages`, because + `nvidia_cutlass_dsl_packages.pth` does `sys.path.insert(0, ...)` on it. It + therefore outranks `dist-packages` on every interpreter start and shadows the + wheel permanently. CUTLASS DSL never references `mooncake`, so deleting it + breaks nothing. +- **Without it**, the match is `dist-packages` itself and the broken package + *collides* with the wheel. This one is nastier: CMake writes + `store.cpython-312-.so` while the wheel writes `store.so`, and + `importlib.machinery.EXTENSION_SUFFIXES` puts the interpreter-tagged suffix + first, so the broken extension still wins. pip also overwrites `__init__.py`, + erasing the `# Auto-generated by CMake` marker, so afterwards there is no + reliable way to tell leftover files from wheel files. + +Because of that second case, `install_mooncake_runtime.sh` removes any +`mooncake` package directory outright and reinstalls, rather than trying to +identify individual bad files. Both cases are covered. + +**2. The default wheel is built for CUDA 12.** `mooncake-transfer-engine` links +against `libcudart.so.12`; containers from `pytorch-26.05` on ship CUDA 13 only. +Every extension and the `mooncake_master` binary then fail to load: + +``` +ImportError: libcudart.so.12: cannot open shared object file +``` + +Use **`mooncake-transfer-engine-cuda13`**, the same project built for CUDA 13, +which needs no shim. Its releases start at 0.3.9, so it cannot match the +`MOONCAKE_VERSION` pin (`0.3.7.post2`) in `install_mooncake.sh`. That drift is +safe here: `/usr/local/Mooncake` backs the *cache transceiver's* Mooncake +backend, a different feature that these configs do not use +(`cache_transceiver_config.backend: "NIXL"`). The connector only ever talks to +the wheel, and the wheel also supplies the `mooncake_master` that lands on +`PATH` ahead of the CMake one, so client and master stay matched. Revisit this +only if you set the transceiver backend to `MOONCAKE`. + +If you would rather match `install_mooncake.sh` exactly, the script keeps that +path working and applies the `libcudart.so.12` shim for you: + +```bash +MOONCAKE_WHEEL="mooncake-transfer-engine==0.3.7.post2" \ + bash mooncake_disagg/install_mooncake_runtime.sh +``` + +Both wheel choices were validated against the full set of store methods the +connector calls -- see "Validating the install" below. + +### How often does the script need to run? + +It depends on whether the container filesystem persists, because the script +writes into `dist-packages` inside the container, not into your checkout. + +| Situation | How often | +|---|---| +| Long-lived container you `docker exec` into | **Once.** It survives until the container is deleted; `docker restart` keeps it. | +| SLURM via `disaggr_torch.slurm` | **Once per job, per node** -- and the harness now does it for you, see below. | +| Image built from this branch | **Never.** `install_mooncake.sh` now does it at build time and fails the build if the import does not work. | + +For the SLURM case this is already wired up: `disaggr_torch.slurm` now runs the +script on every node, right after its `pip install -e .[devel]` step, and gates +it on whether a worker config actually asks for the connector: + +```bash +if grep -qs "mooncake-store" "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"; then +``` + +So arms A and B of the run matrix pay nothing, arm C installs automatically, and +there is no new config key to remember. It resolves the script as +`${trtllm_repo}/mooncake_disagg/install_mooncake_runtime.sh` and fails the job +with an explicit message if `environment.trtllm_repo` is unset or does not +contain it -- which is the case if you benchmark from +`environment.trtllm_wheel_path` instead, so use an image with the bindings baked +in for that path. Output lands in `/2_install_mooncake.log`. Set +`MOONCAKE_WHEEL` in the submitting environment to override the wheel; it is +forwarded to every node. + +Note that `--container-name` gives each node a container that lives for the whole +job, so the install survives from that step through to the worker `srun`s. It +does not survive into the *next* job, which is why this runs per job rather than +once. Do not try to persist it via `~/.local` unless home is genuinely shared and +mounted (`disaggr_torch.slurm` passes `--no-container-mount-home` to most of its +`srun` calls). + +Baking an image is the only option that removes the step entirely. Given the +script takes about eight seconds from cold, that is a convenience decision rather +than a necessity. + +### Will the shadow package come back? + +The root cause is upstream in Mooncake's `CMakeLists.txt` and is **not** fixed; +both scripts clean up after it. Practically: + +- **Images built from this branch:** no. The cleanup runs in the same script, + immediately after `make install` and before the wheel install, and the build + now fails if `import mooncake.store` does not work. +- **Any pre-existing image**, including the one pinned in + `jenkins/current_image_tags.properties` (`202607211045`): the broken package is + baked in, so the runtime script is required. +- **Inside a running container:** only if something re-runs Mooncake's CMake + install. Reinstalling `nvidia-cutlass-dsl` does not recreate it -- that package + has never shipped a `mooncake` directory; it only supplies the `.pth` that made + CMake choose the wrong destination. +- **If Mooncake is ever upgraded** to a version that fixes its install path, or + the `.pth` ordering changes, the cleanup becomes a no-op rather than a hazard. + +### Validating the install + +Two scripts in this directory, in increasing order of strictness. Both need a +running master and `MOONCAKE_CONFIG_PATH`, exactly like a real worker: + +```bash +mooncake_master --rpc_port=50051 --metrics_port=9004 & + +export MOONCAKE_CONFIG_PATH=$PWD/mooncake.json # TCP config; edit the master address +python3 mooncake_disagg/mooncake_smoke_test.py # setup + put/get round trip +python3 mooncake_disagg/mooncake_api_surface_test.py # needs a GPU +``` + +`mooncake_smoke_test.py` proves the bindings load and `store.setup()` succeeds +with the same argument list `worker.py` passes. `mooncake_api_surface_test.py` is +the one that matters when changing wheel versions: the connector's hot path never +uses `put`/`get`, it uses `register_buffer` plus the zero-copy +`batch_put_from_multi_buffers` / `batch_get_into_multi_buffers` / `batch_is_exist` +calls against registered GPU pages. Those take `list[list[int]]` -- one buffer +list per key, because `PageAddressing.page_buffers` returns one address per +layer-group region -- and that is the signature most likely to drift. + +Then the unit tests, which need no store and no GPU: + +```bash +pytest tests/unittest/_torch/executor/test_mooncake_store_connector.py +``` + +## 3. Topology + +``` + mooncake_master (1 CPU core, its own job) + ^ ^ + register/put/get| | + ┌─────────────────────┴──┐ ┌──┴──────────────────────┐ + │ CTX instance 0 TP=4 │ │ CTX instance 1 TP=4 │ store: role=both + └────────────┬───────────┘ └───────────┬─────────────┘ + │ NIXL KV handoff │ + └──────────┬────────────────┘ + v + ┌──────────────────────────┐ + │ GEN instance TP=4 │ no connector at all + └──────────────────────────┘ + ^ + round-robin│ + ┌──────────┴───────────┐ + │ trtllm-serve disagg │ <- benchmark_serving client + └──────────────────────┘ +``` + +12 GPUs = 3 nodes at 4 GPUs/node. The generation worker deliberately has **no** +`kv_connector_config`: generated tokens are rarely a reused prefix, and that +absence is the only way to express "off" (`StoreRole` has no off value). It also +lets the generation worker keep its host cache tier and `MAX_UTILIZATION` +scheduler, both of which the connector would forbid. + +## 4. Step 1 -- run the Mooncake master + +`master_server_address` is mandatory, so a master must exist and be reachable +from every worker. The benchmark harness has no hook for launching a side +process, and the workers read `MOONCAKE_CONFIG_PATH` at startup, so the master's +address has to be known *before* the benchmark job is submitted. Run it as its +own long-lived job: + +```bash +# mooncake_master.sbatch +#!/bin/bash +#SBATCH --job-name=mooncake-master +#SBATCH --nodes=1 +#SBATCH --time=08:00:00 +#SBATCH --output=%x-%j.out + +srun --container-image=$CONTAINER_IMAGE \ + --container-mounts=$WORK_DIR:$WORK_DIR \ + bash -lc ' + hostname -i | awk "{print \$1}" > '"$WORK_DIR"'/master.addr + exec mooncake_master \ + --rpc_port=50051 \ + --metrics_port=9004 \ + --eviction_ratio=0.05 + ' +``` + +Flag names above were read out of the shipped `mooncake_master` binary. Run +`mooncake_master --help` inside the container to confirm defaults and to see the +rest (`--rpc_address`, `--rpc_thread_num`, `--default_kv_lease_ttl`, +`--eviction_high_watermark_ratio`, `--enable_http_metadata_server`, +`--cluster_id`, `--root_fs_dir`). + +Keeping the master in a separate job is what makes experiment 3 (§7) possible: +the pool outlives the engines, so a second benchmark job finds a warm store. + +Then write the client config, substituting the address the master job just +recorded. The schema is vLLM's, so one pool can serve both engines: + +```bash +MASTER_IP=$(cat $WORK_DIR/master.addr) +cat > $WORK_DIR/mooncake.json <:/metadata` with + `mooncake_http_metadata_server` (or the master's own + `--enable_http_metadata_server`) only if you need the shared-metadata + behaviour; confirm the port from `--help` rather than assuming. +- `device_name`: pick from `ibv_devinfo` on a compute node. For first bring-up + only, `"protocol": "tcp"` with `"device_name": ""` removes RDMA from the + variable list -- that is what `mooncake.json` in this directory currently + does. Do not draw performance conclusions from a TCP run. +- `global_segment_size` is contributed **per worker process**, so the pool is + `global_segment_size x (ctx instances x TP)` = 8 segments here. +- Sizing: after startup, each worker logs its page geometry (§8). Pool bytes for + a corpus of `T` unique prefix tokens is + `T / tokens_per_block x Σ_layer_groups bytes_per_page x world_size`. + As an anchor, the residency work measured M3 at ~22 KiB/token aggregate across + TP=4 (fp8 KV plus a per-rank-replicated index-K), so ~21 GiB per million + unique prefix tokens. Confirm against your own log line rather than trusting + that number. +- `role`/`cache_prefix` can also be overridden per process by + `TRTLLM_MOONCAKE_STORE_ROLE` and `TRTLLM_MOONCAKE_STORE_PREFIX`. Bump the + prefix whenever you change anything that should not be shared with an earlier + run's pages. + +## 5. Step 2 -- the harness config + +Copy `examples/disaggregated/slurm/benchmark/config.yaml` and replace the +`worker_config` section with M3's. `submit.py` serializes `worker_config.ctx` +and `worker_config.gen` straight to `ctx_config.yaml`/`gen_config.yaml` with +`yaml.dump`, so any LLM-API key passes through untouched -- including +`kv_connector_config`. + +The context worker below is `m3_ctx_mooncake.yaml` from this directory; the +generation worker is `m3_gen_mooncake.yaml`. Every deviation from the production +M3 config is marked and explained there, and those comments are the reason to +read those two files rather than treating this block as self-explanatory. + +```yaml +# m3_store_2ctx.yaml +slurm: + script_file: "disaggr_torch.slurm" + partition: "" + account: "" + job_time: "04:00:00" + job_name: "m3-mooncake-store" + extra_args: "" + set_segment: true + numa_bind: true # GB200/GB300 NVL72 + +benchmark: + mode: "e2e" + use_nv_sa_benchmark: false + multi_round: 8 # num_prompts = concurrency x multi_round + streaming: true + concurrency_list: "8" + input_length: 131072 # log-dir naming only; the dataset is authoritative + output_length: 1024 + dataset_file: "/m3_shared_prefix.jsonl" + +hardware: + gpus_per_node: 4 + num_ctx_servers: 2 # >= 2 is the whole point; see experiment 2 + num_gen_servers: 1 + +environment: + container_mount: "" + container_image: "" + model_path: "" + trtllm_repo: "" + build_wheel: false + trtllm_wheel_path: "" + work_dir: "" + worker_env_var: "TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes NCCL_GRAPH_MIXING_SUPPORT=0" + # Only the context workers open a store handle. + ctx_worker_env_var: "MOONCAKE_CONFIG_PATH=/mooncake.json TRTLLM_MOONCAKE_STORE_ROLE=both TRTLLM_MOONCAKE_STORE_PREFIX=trtllm-m3-run1" + server_env_var: "TRTLLM_SERVER_DISABLE_GC=1" + +profiling: + nsys_on: false + ctx_profile_range: "10-30" + gen_profile_range: "200-250" + +accuracy: + enable_accuracy_test: false + tasks: {} + +worker_config: + ctx: + # ---- contents of m3_ctx_mooncake.yaml, plus parallelism ---- + tensor_parallel_size: 4 + moe_expert_parallel_size: 4 + pipeline_parallel_size: 1 # gated: connector refuses PP > 1 + context_parallel_size: 1 # gated: connector refuses CP > 1 + enable_attention_dp: false # required: dummy DP-balancing requests reach the hooks + max_seq_len: 1048576 + max_num_tokens: 16384 + max_batch_size: 20 + sparse_attention_config: + algorithm: minimax_m3 + implementation: msa + indexer_kv_dtype: fp8 + sparse_disable_index_value: true # gated: index-V lives outside the paged pools + fuse_qkv_index_projection: true + kv_cache_config: + free_gpu_memory_fraction: 0.94 + enable_block_reuse: true + tokens_per_block: 128 + use_kv_cache_manager_v2: true # required: only V2 can describe its pools + dtype: fp8 + event_buffer_max_size: 0 + host_cache_size: 0 # gated: must be explicit 0, not omitted + disk_cache_size: 0 + scheduler_config: + capacity_scheduler_policy: GUARANTEED_NO_EVICT # gated + cache_transceiver_config: + backend: "NIXL" + transceiver_runtime: "PYTHON" # M3 is always-V2; C++ transceiver is refused + enable_chunked_prefill: true + enable_autotuner: true + trust_remote_code: true + reasoning_parser: minimax_m3 + stream_interval: 20 + print_iter_log: true + num_postprocess_workers: 8 + # Required to see any reuse number at all -- see section 8. All three + # default to false, and without them /metrics returns an empty list. + enable_iter_perf_stats: true + enable_iter_req_stats: true + return_perf_metrics: true + kv_connector_config: + connector: mooncake-store # <-- the only line experiment 1 removes + + gen: + # ---- contents of m3_gen_mooncake.yaml; no connector, so no gates ---- + tensor_parallel_size: 4 + moe_expert_parallel_size: 4 + enable_attention_dp: false + max_seq_len: 1048576 + max_num_tokens: 16384 + max_batch_size: 20 + sparse_attention_config: + algorithm: minimax_m3 + implementation: msa + indexer_kv_dtype: fp8 + sparse_disable_index_value: true # must match ctx: it changes the model + fuse_qkv_index_projection: true + kv_cache_config: + free_gpu_memory_fraction: 0.94 + enable_block_reuse: true + block_reuse_policy: per_conversation + tokens_per_block: 128 + use_kv_cache_manager_v2: true + dtype: fp8 + event_buffer_max_size: 0 + host_cache_size: 388554555392 # kept: no connector here + scheduler_config: + capacity_scheduler_policy: MAX_UTILIZATION # kept: no connector here + cache_transceiver_config: + backend: "NIXL" + transceiver_runtime: "PYTHON" # must match ctx + enable_chunked_prefill: true + enable_autotuner: true + trust_remote_code: true + reasoning_parser: minimax_m3 + stream_interval: 20 + print_iter_log: true + num_postprocess_workers: 8 + enable_iter_perf_stats: true + enable_iter_req_stats: true + return_perf_metrics: true +``` + +Eagle3 is left off. It is not gated, but `MiniMaxM3KVCacheManagerV2` sets +`supports_shared_draft_layers`, so draft layers join the unified V2 cache and +therefore the registered layout -- extra page geometry the store must key +correctly, on a path with no coverage. Turn it on only after a clean run +without it. + +Submit with: + +```bash +cd examples/disaggregated/slurm/benchmark +python3 submit.py -c /m3_store_2ctx.yaml --dry-run # inspect first +python3 submit.py -c /m3_store_2ctx.yaml +``` + +## 6. Step 3 -- the workload + +`run_benchmark.sh` invokes `benchmark_serving` with +`--dataset-name trtllm_custom --dataset-path `, so you supply a +JSONL file. `CustomDataset` reads `input.messages[1].content` as the prompt, +`input.max_tokens` as the output length, and skips re-tokenization when +`input.num_tokens` is present. It shuffles the file on load, which is what +spreads repeated prefixes apart in time. + +The workload must have **repeated prefixes across requests**, because that is +the only structure a content-addressed store can exploit. `P` distinct prefixes +each repeated `R` times, with a unique suffix per request so no two requests are +identical: + +```python +# gen_shared_prefix_dataset.py +import json, random +from transformers import AutoTokenizer + +MODEL = "" +NUM_PREFIXES = 8 # P distinct shared prefixes +REPEATS = 8 # R requests per prefix -> P*R = 64 total +PREFIX_TOKENS = 131072 # must be >> tokens_per_block (128) to be worth storing +SUFFIX_TOKENS = 512 +OUTPUT_TOKENS = 1024 +OUT = "m3_shared_prefix.jsonl" + +tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) +rng = random.Random(1234) +vocab = tok.vocab_size + +def text_of(n_tokens, seed): + r = random.Random(seed) + ids = [r.randrange(1000, vocab - 1000) for _ in range(int(n_tokens * 1.3))] + text = tok.decode(ids, skip_special_tokens=True) + # Re-encode and trim: decode/encode is not a round trip, so measure. + ids = tok.encode(text, add_special_tokens=False)[:n_tokens] + return tok.decode(ids, skip_special_tokens=True), len(ids) + +prefixes = [text_of(PREFIX_TOKENS, 100 + i) for i in range(NUM_PREFIXES)] + +with open(OUT, "w") as f: + for i, (prefix, plen) in enumerate(prefixes): + for r in range(REPEATS): + suffix, slen = text_of(SUFFIX_TOKENS, 900000 + i * 1000 + r) + f.write(json.dumps({"input": { + "messages": [{"role": "system", "content": ""}, + {"role": "user", "content": prefix + suffix}], + "max_tokens": OUTPUT_TOKENS, + "num_tokens": plen + slen, + }}) + "\n") +``` + +Size the file against what the client will actually request. +`run_benchmark.sh` computes +`num_prompts = (concurrency x num_gen_servers) x multi_round`, so the §5 config +(`concurrency_list: "8"`, `multi_round: 8`, one generation server) asks for 64 +prompts -- which is why `P x R` above is 64. Ask for more than the file holds +and the extra is not sampled; write more than you ask for and the tail of your +repeat structure never runs. + +Random token text is deliberate: it defeats any accidental prefix sharing +between "distinct" prefixes, so the hit rate you measure is the one you +designed. If you would rather test genuine production traffic, substitute a +real multi-turn trace -- but keep an eye on whether it actually contains +repeated prefixes, since without them the store has nothing to do and a flat +result means nothing. + +## 7. Step 4 -- the run matrix + +Three arms, and the middle one is the one people skip: + +| Arm | Config | Purpose | +|---|---|---| +| **A** production reference | today's M3 config: host tier on, `MAX_UTILIZATION`, no connector | where you are today | +| **B** gated baseline | arm A's config edited to satisfy every connector gate (`host_cache_size: 0`, `GUARANTEED_NO_EVICT`, ...), still no connector | isolates what the gates cost | +| **C** store | arm B plus `kv_connector_config` | isolates what the store adds | + +Comparing C against A alone conflates two independent changes: the store's +benefit and the loss of a 362 GiB host cache tier plus a scheduler policy +change. **The store's efficacy is C vs B.** A vs B tells you the entry price, +and A vs C tells you whether the whole package is deployable. All three are +worth knowing and they answer different questions. + +Then, within that: + +**Experiment 1 -- does it work at all (`num_ctx_servers: 1`).** +Arm C, one context instance, small `PREFIX_TOKENS` (say 4096) and a short run. +You are looking for a clean startup, the registration log line, non-zero store +traffic, no load failures, and coherent output text. Do this before spending an +allocation on anything larger. Expect no throughput change; local reuse already +serves this case. + +**Experiment 2 -- cross-instance reuse (`num_ctx_servers: 2`).** +The router defaults to round-robin, so consecutive requests alternate between +context instances and roughly half of each prefix's repeats land on the instance +that did not compute it. Those are the requests local reuse must recompute from +scratch and the store can serve. Compare arm C against arm B on: +- TTFT p50/p99 (the store's whole thesis is prefill avoided) +- `reused_blocks_per_request` distribution (§8) +- output tokens/s/GPU + +This is the primary result. A store that does not win here does not work. + +**Experiment 3 -- survival across restarts.** +Run experiment 2's arm C twice, same `TRTLLM_MOONCAKE_STORE_PREFIX`, with the +master job left running between them. The second job starts with an empty local +cache but a warm pool. First-round TTFT should fall toward the warm steady-state +value. Local reuse scores zero here by construction, so any improvement is +attributable to the store alone -- which makes this the cleanest signal in the +whole matrix, and the cheapest to run. + +**Experiment 4 -- the cost when there is nothing to gain.** +Arm C against arm B on a workload with *no* repeated prefixes (unique prompts). +This measures pure overhead: lookups, key hashing on the leader, background +saves competing for host bandwidth. Ideally indistinguishable from arm B. This +is the arm that catches a feature that helps its benchmark and hurts the fleet. + +## 8. Step 5 -- reading the results + +### Did the connector even load? + +Every context worker logs at INFO on startup: + +``` +mooncake-store leader ready (role=both, tokens_per_block=128) +mooncake-store worker rank 0/4 ready (role=both, model_key=MiniMax-M3-NVFP4, master=10.0.0.5:50051) +mooncake-store worker rank 0 registered layout: tokens_per_block=128, lg0(layers=N, regions=..., bytes/page=..., slots=..., window=None) +``` + +```bash +grep -h "mooncake-store" /3_output_CTX_*.log | head -40 +``` + +The registration line is the one to keep: `bytes/page` per layer group is what +your pool sizing in §4 depends on, and `window=None` is the confirmation that no +sliding-window group is present (one would have aborted startup). + +### Is the store actually moving pages? + +Hit and transfer counts are at DEBUG. The connector logs under module `_torch`, +so: + +``` +TLLM_LOG_LEVEL_BY_MODULE="debug:_torch" +``` + +added to `environment.ctx_worker_env_var`. This is verbose -- it enables DEBUG +for all of `_torch` -- so use it for experiment 1 and for diagnosis, not for the +runs you intend to quote numbers from. The lines worth counting: + +``` +mooncake-store matched N blocks (M tokens) for request R # leader, a hit +mooncake-store rank K loaded P pages # worker, a load +``` + +The store's own counters are the alternative that costs nothing at runtime: +`mooncake_master --metrics_port=9004` exposes pool-level statistics over HTTP. +Scrape it before and after a run and diff. + +### Which reuse number means what + +This distinction matters and is easy to get backwards: + +| Signal | Where | Includes store hits? | +|---|---|---| +| `reused_blocks_per_request`, `kv_cache_hit_rate_per_request` | per-request iteration stats | **Yes.** `_reserve_connector_prefix` calls `set_prepopulated_prompt_len` with the connector-served position, and these derive from `mPrepopulatedPromptLen`. | +| `kv_cache_iter_reused_blocks`, `kv_cache_iter_reuse_rate` | `GET /prometheus/metrics` | **No.** These come from the local V2 reuse tree's committed stats. | + +So **store hits ≈ per-request reuse − local-tree reuse**. Confirm that +relationship on experiment 3, where the local tree starts empty and the +difference is unambiguous, before relying on it elsewhere. + +Getting at either one requires the three flags added to the worker configs in +§5, all of which default to false: + +- `enable_iter_perf_stats: true` -- without it `get_latest_iteration_stats` + short-circuits and `GET /metrics` returns `[]`. +- `enable_iter_req_stats: true` -- needed for the *per-request* half of the + table above. +- `return_perf_metrics: true` -- mounts `/prometheus/metrics`. `GET /metrics` + (plain JSON iteration stats) is routed unconditionally but still needs + `enable_iter_perf_stats`. + +`print_iter_log: true` is worth keeping on, but note it prints iteration timing +and KV *utilization* only -- no reuse counters. Do not go looking for hit rates +there. + +Per-request client-side results land in `/concurrency_/result.json` +with TTFT/TPOT/ITL/E2EL percentiles, which is where the headline numbers for +§7's arms come from. + +### Failure signatures + +| Log line | Meaning | +|---|---| +| `mooncake-store failed to load N of M pages` (raises) | **Stop.** The runtime had already counted those tokens as computed, so this is the tripwire against silently wrong answers. Do not treat as flaky. | +| `mooncake-store background save failed` | A save thread exception, re-raised on the executor thread. | +| `mooncake-store rank K failed to save N of M pages` (warning) | Dropped write. Costs a future miss, not correctness. A trickle is tolerable; a flood means the pool is full or the master is overloaded. | +| `mooncake-store lookup failed; treating as a miss` (warning) | Probe failed. Degrades to no-store behavior. | +| `could not reserve connector prefix up to N, falling back to the local match` (debug) | Out of GPU pages. The store offered more than the engine could hold -- expected under pressure, but frequent occurrences mean the offer is outrunning capacity. | + +### Sanity check that is not a performance number + +Run a handful of prompts through arms B and C with temperature 0 and compare the +text. A store that returns the wrong bytes shows up as degraded output long +before it shows up as an error. `accuracy.enable_accuracy_test: true` with gsm8k +gives a coarser version of the same check. + +## 9. Things that will bite + +- **`host_cache_size: 0` must be written explicitly.** Left at its default of + `None`, V2 still provisions a host tier, and the gate rejects it. Falsy is not + the same as absent here. +- **The gates change the config out from under you.** `GUARANTEED_NO_EVICT` + instead of `MAX_UTILIZATION`, no host tier, no attention DP. That is why + arm B exists. +- **`sparse_disable_index_value: true` changes the model, not just the cache.** + Hold it fixed across every arm, generation workers included, or you are + comparing two different models. +- **Key namespace pins world size and rank.** Change TP and every stored page + becomes unreachable -- a miss, not an error. Same for `tokens_per_block`, the + layer group set, and `bytes_per_page`. +- **`model_key` defaults to the checkpoint directory's basename.** Two hosts + mounting the same checkpoint at different paths still share cache, which is + intended; two *different* checkpoints in identically-named directories also + share it, which is not. Set `TRTLLM_MOONCAKE_STORE_MODEL_KEY` explicitly for + anything long-lived. +- **Stale pages across code changes.** The key namespace does not include a + build hash. After changing anything about page layout or contents, bump + `TRTLLM_MOONCAKE_STORE_PREFIX` or restart the master. +- **UCX warmup requests hit the store too.** `run_benchmark.sh` sends + `2 x ctx_instances x gen_instances` 100-token requests before the real run. + Harmless, but they are in the counters. +- **`enable_chunked_prefill` interacts with the offer.** The connector offers + only whole blocks and only when the local match is block-aligned; a partial + local match disables the store for that request entirely. With + `tokens_per_block: 128` this is rare, but it explains occasional zero-offer + requests. +- **`block_reuse_policy: per_conversation` is off on the context worker** in + these configs. It is not gated, but the connector derives its own + `cache_salt`-seeded hash chain and the interaction is untested. Restore it + only after the store is proven, and treat it as its own experiment. + +## 10. What this does not test + +Worth stating so the results are not oversold: single-node only insofar as the +master is one process (no HA master, no `--root_fs_dir` persistence); no +pipeline or context parallelism (both refused); no VSWA or sliding-window model +(refused); no Eagle3; no shared pool between TensorRT-LLM and vLLM, though the +config schema is deliberately compatible with it. Load bandwidth under +contention from many simultaneous large prefixes is exercised only incidentally +by concurrency, not measured directly -- if experiment 2 shows a TTFT +regression at high concurrency despite hits, that is the first thing to profile. diff --git a/mooncake_disagg/ctx_config.yaml b/mooncake_disagg/ctx_config.yaml new file mode 100644 index 000000000000..2cfa4359fc9e --- /dev/null +++ b/mooncake_disagg/ctx_config.yaml @@ -0,0 +1,30 @@ +# Context (prefill) worker: reads AND writes the Mooncake store. +# +# export MOONCAKE_CONFIG_PATH=/abs/path/to/mooncake.json +# CUDA_VISIBLE_DEVICES=0 trtllm-serve \ +# --host localhost --port 8001 --server_role CONTEXT \ +# --config ./ctx_config.yaml + +kv_cache_config: + # The connector describes its pools through register_kv_cache_layout, + # which only KVCacheManagerV2 implements. + use_kv_cache_manager_v2: true + # Local reuse still runs first; the store serves whatever the device missed. + enable_block_reuse: true + # GPU-only tiers are required: a page evicted to host or disk has its GPU + # slot reassigned, which would invalidate the addresses registered with + # the store. + host_cache_size: 0 + disk_cache_size: 0 + free_gpu_memory_fraction: 0.2 + +cache_transceiver_config: + # Requiring KVCacheManagerV2 for the connector forces the Python + # transceiver: the C++ one is bound to the V1 BaseKVCacheManager and + # raises on a V2 manager. NIXL is the only backend the Python + # transceiver supports. + backend: "NIXL" + transceiver_runtime: "PYTHON" + +kv_connector_config: + connector: mooncake-store diff --git a/mooncake_disagg/disagg_config.yaml b/mooncake_disagg/disagg_config.yaml new file mode 100644 index 000000000000..ca1f35a7072d --- /dev/null +++ b/mooncake_disagg/disagg_config.yaml @@ -0,0 +1,18 @@ +# Router. Only needs to know where the workers are; each worker's own LLM +# args come from its --config file (ctx_config.yaml / gen_config.yaml). +# +# trtllm-serve disaggregated -c ./disagg_config.yaml + +hostname: localhost +port: 8000 +backend: "pytorch" + +context_servers: + num_instances: 1 + urls: + - "localhost:8001" + +generation_servers: + num_instances: 1 + urls: + - "localhost:8002" diff --git a/mooncake_disagg/gen_config.yaml b/mooncake_disagg/gen_config.yaml new file mode 100644 index 000000000000..4b8ed0b35262 --- /dev/null +++ b/mooncake_disagg/gen_config.yaml @@ -0,0 +1,22 @@ +# Generation (decode) worker: does not touch the Mooncake store at all. +# +# There is deliberately no kv_connector_config here. "Decode none" is the +# absence of a connector, not a StoreRole -- StoreRole only has +# producer / consumer / both. Generated tokens are rarely a reused prefix, +# so writing them would cost bandwidth for no hit rate. +# +# CUDA_VISIBLE_DEVICES=1 trtllm-serve \ +# --host localhost --port 8002 --server_role GENERATION \ +# --config ./gen_config.yaml + +kv_cache_config: + # Only to match the context side's transceiver, which must be the Python + # one there. Nothing here is required by the store: this worker never + # opens a store handle. + use_kv_cache_manager_v2: true + enable_block_reuse: true + free_gpu_memory_fraction: 0.8 + +cache_transceiver_config: + backend: "NIXL" + transceiver_runtime: "PYTHON" diff --git a/mooncake_disagg/install_mooncake_runtime.sh b/mooncake_disagg/install_mooncake_runtime.sh new file mode 100755 index 000000000000..27c2a658885c --- /dev/null +++ b/mooncake_disagg/install_mooncake_runtime.sh @@ -0,0 +1,135 @@ +#!/bin/bash +# Make the Mooncake Python store bindings importable inside a TensorRT-LLM +# container, so the mooncake-store KV connector can start. +# +# Two things break a plain `pip install mooncake-transfer-engine` in the +# containers built by docker/common/install_mooncake.sh: +# +# 1. The CMake source build in that script emits its own, unusable `mooncake` +# Python package (it omits libmooncake_store.so). mooncake-integration's +# CMakeLists picks the install directory with +# python3 -c "import sys; print([s for s in sys.path if 'packages' in s][0])" +# -- the first sys.path entry whose name merely contains "packages". With +# nvidia-cutlass-dsl installed that is nvidia_cutlass_dsl/dsl_packages, +# which nvidia_cutlass_dsl_packages.pth sys.path.insert(0)s, so it shadows +# anything pip installs. Without it, the package lands in dist-packages and +# collides with the wheel: CMake writes store.cpython-312-.so, the +# wheel writes store.so, and importlib prefers the interpreter-tagged +# suffix, so the broken extension still wins. Either way the symptom is +# `ImportError: libmooncake_store.so` *after* pip reports success. +# Because pip overwrites __init__.py in the collision case, leftovers are +# not reliably identifiable after the fact -- so this script removes the +# package directory outright and reinstalls, rather than trying to tell +# good files from bad. +# +# 2. The `mooncake-transfer-engine` wheel is linked against libcudart.so.12, +# while containers from pytorch-26.05 on ship CUDA 13 only. +# `mooncake-transfer-engine-cuda13` is the same project built for CUDA 13 +# and needs no shim, so it is the default here. Its releases start at 0.3.9, +# so it cannot match the 0.3.7.post2 pin in install_mooncake.sh -- see the +# note below on why that is safe. +# +# Version drift against /usr/local/Mooncake: that CMake-built C++ library backs +# the *cache transceiver's* Mooncake backend, a different feature. The connector +# only ever talks to the wheel, and the wheel also supplies the mooncake_master +# that lands on PATH, so client and master stay matched. Revisit only if you set +# cache_transceiver_config.backend to MOONCAKE (these configs use NIXL). +# +# Set MOONCAKE_WHEEL to override, e.g. +# MOONCAKE_WHEEL="mooncake-transfer-engine==0.3.7.post2" +# to match install_mooncake.sh exactly; the libcudart.so.12 shim is then applied +# automatically. +# +# Idempotent, and cheap on re-runs: if the install is already correct it exits +# without contacting the network, so it is safe in a SLURM prolog on every node. + +set -euo pipefail + +MOONCAKE_WHEEL="${MOONCAKE_WHEEL:-mooncake-transfer-engine-cuda13==0.3.13}" +WHEEL_NAME="${MOONCAKE_WHEEL%%[=<>]*}" +SITE_PACKAGES="$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" + +echo ">> target wheel: ${MOONCAKE_WHEEL}" + +# Fast path: already correct, so do not touch the network. +if pip3 show "${WHEEL_NAME}" >/dev/null 2>&1 && + python3 -c 'from mooncake.store import MooncakeDistributedStore; MooncakeDistributedStore()' >/dev/null 2>&1; then + echo ">> already installed and importable; nothing to do" + python3 -c 'import mooncake.store; print(" resolved extension:", mooncake.store.__file__)' + exit 0 +fi + +# Purge every `mooncake` package directory on the search path, whatever wrote it. +# Distinguishing CMake leftovers from wheel files is unreliable once pip has +# overwritten __init__.py, so remove and reinstall instead. +python3 - <<'PY' +import os +import shutil +import sys +import sysconfig + +paths = sysconfig.get_paths() +for entry in list(sys.path) + [paths["purelib"], paths["platlib"]]: + if not entry: + continue + package = os.path.join(entry, "mooncake") + if os.path.isdir(package): + print(f">> removing existing mooncake package: {package}") + shutil.rmtree(package, ignore_errors=True) +PY + +# The two distributions install the same `mooncake` package, so leaving both +# registered produces a half-overwritten directory. +for distribution in mooncake-transfer-engine mooncake-transfer-engine-cuda13; do + if pip3 show "${distribution}" >/dev/null 2>&1; then + echo ">> unregistering ${distribution}" + pip3 uninstall -y -q "${distribution}" >/dev/null 2>&1 || true + fi +done + +pip3 install --no-cache-dir "${MOONCAKE_WHEEL}" + +# Only the CUDA 12 wheel needs the runtime shim. Drop it into the wheel's own +# RPATH directory so the Python extensions and the mooncake_* binaries all find +# it without LD_LIBRARY_PATH being set in each process. +if ldd "${SITE_PACKAGES}"/mooncake/store*.so 2>/dev/null | grep -q "libcudart.so.12 => not found"; then + echo ">> wheel needs libcudart.so.12, which this container lacks; installing it" + pip3 install --no-cache-dir nvidia-cuda-runtime-cu12 + CUDART12="${SITE_PACKAGES}/nvidia/cuda_runtime/lib/libcudart.so.12" + [[ -f "${CUDART12}" ]] || { echo "libcudart.so.12 not found after install" >&2; exit 1; } + mkdir -p "${SITE_PACKAGES}/mooncake_transfer_engine.libs" + ln -sf "${CUDART12}" "${SITE_PACKAGES}/mooncake_transfer_engine.libs/libcudart.so.12" + echo ">> linked libcudart.so.12 into the wheel's RPATH directory" +fi + +# Verify, because every failure above stays silent until a worker starts. +echo ">> verifying" +python3 - <<'PY' +import mooncake.store +from mooncake.store import MooncakeDistributedStore + +MooncakeDistributedStore() +print(" mooncake.store imports and instantiates: OK") +print(f" resolved extension: {mooncake.store.__file__}") +PY + +# Every extension module and the master binary must have a resolvable link line. +# ldd the real ELF files, not /usr/local/bin/mooncake_master, which is a Python +# console script. +unresolved=0 +for elf in "${SITE_PACKAGES}"/mooncake/*.so "${SITE_PACKAGES}"/mooncake/mooncake_master; do + [[ -e "${elf}" ]] || continue + if missing="$(ldd "${elf}" 2>&1 | grep 'not found')"; then + echo " $(basename "${elf}"): unresolved -> ${missing}" >&2 + unresolved=1 + fi +done +[[ "${unresolved}" -eq 0 ]] || { echo "unresolved shared libraries; see above" >&2; exit 1; } +echo " all mooncake ELF link lines resolve: OK" + +for entry in mooncake_master mooncake_http_metadata_server; do + path="$(command -v "${entry}")" || { echo "${entry} is not on PATH" >&2; exit 1; } + echo " ${entry}: OK (${path})" +done + +echo ">> done" diff --git a/mooncake_disagg/m3_agg_mooncake.yaml b/mooncake_disagg/m3_agg_mooncake.yaml new file mode 100644 index 000000000000..a2c0251f0f0f --- /dev/null +++ b/mooncake_disagg/m3_agg_mooncake.yaml @@ -0,0 +1,89 @@ +# MiniMax-M3-NVFP4, aggregated, with the mooncake-store KV connector. +# +# Derived from the production M3 serving config. Every deviation from it is +# marked CONNECTOR: with the gate that forces it. Startup gates live in +# py_executor_creator.py (~line 830), py_executor._maybe_init_kv_connector_manager, +# and connectors/mooncake_store/validation.py. + +max_seq_len: 1048576 +max_num_tokens: 16384 +max_batch_size: 20 + +cuda_graph_config: + enable_padding: true + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20] + +torch_compile_config: + enable_fullgraph: true + enable_inductor: false + enable_piecewise_cuda_graph: true + capture_num_tokens: [1, 16, 128, 256, 512, 768, 1024, 1280, 1536, 1792, 2048] + enable_userbuffers: true + max_num_streams: 3 + +moe_config: + backend: TRTLLM + use_low_precision_moe_combine: true + +sparse_attention_config: + algorithm: minimax_m3 + implementation: msa + indexer_kv_dtype: fp8 + # Required by the connector: index-V is a plain tensor outside the paged + # pools, so it is never described or transferred. A replayed prefix would + # pair stored index-K with stale index-V. + sparse_disable_index_value: true + fuse_qkv_index_projection: true + +kv_cache_config: + free_gpu_memory_fraction: 0.94 + enable_block_reuse: true + tokens_per_block: 128 + use_kv_cache_manager_v2: true + dtype: fp8 + event_buffer_max_size: 0 + + # CONNECTOR: was host_cache_size: 388554555392. + # _reject_non_gpu_cache_tiers rejects every tier below GPU, because a + # registered region is only a valid device address while its page stays + # pinned to GPU -- eviction reassigns the slot. Both must be an explicit + # 0: V2 provisions a host tier when the field is left at its default of + # None, which is falsy but still yields a tier. + host_cache_size: 0 + disk_cache_size: 0 + + # CONNECTOR: was block_reuse_policy: per_conversation. + # Not gated, but the connector derives its own blake2b hash chain seeded + # by cache_salt, so reuse-policy interactions are untested. Restore this + # after the store is proven. + # block_reuse_policy: per_conversation + +# CONNECTOR: was MAX_UTILIZATION. py_executor_creator raises +# "KV connector is only supported with guaranteed no evict scheduler policy." +scheduler_config: + capacity_scheduler_policy: GUARANTEED_NO_EVICT + +# CONNECTOR: Eagle3 disabled for the first validation run. +# Not gated, but MiniMaxM3KVCacheManagerV2 sets supports_shared_draft_layers, +# so draft layers join the unified V2 cache and therefore the registered +# layout. That adds page geometry the store must key correctly, on a path +# with no coverage. Re-enable once cold/warm passes without it. +# speculative_config: +# decoding_type: Eagle3 +# max_draft_len: 3 +# speculative_model: /path/to/eagle3/draft + +enable_chunked_prefill: true +enable_autotuner: true +trust_remote_code: true +reasoning_parser: minimax_m3 +stream_interval: 20 +print_iter_log: true +num_postprocess_workers: 8 + +# Required: dummy requests inserted for cross-DP balancing flow through the +# connector hooks and are indistinguishable from real requests. +enable_attention_dp: false + +kv_connector_config: + connector: mooncake-store diff --git a/mooncake_disagg/m3_ctx_mooncake.yaml b/mooncake_disagg/m3_ctx_mooncake.yaml new file mode 100644 index 000000000000..f937d24e3c21 --- /dev/null +++ b/mooncake_disagg/m3_ctx_mooncake.yaml @@ -0,0 +1,68 @@ +# MiniMax-M3-NVFP4 CONTEXT (prefill) worker: reads AND writes the store. +# +# Same as m3_agg_mooncake.yaml plus the transceiver. Every CONNECTOR: note +# there applies here too. + +max_seq_len: 1048576 +max_num_tokens: 16384 +max_batch_size: 20 + +cuda_graph_config: + enable_padding: true + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20] + +torch_compile_config: + enable_fullgraph: true + enable_inductor: false + enable_piecewise_cuda_graph: true + capture_num_tokens: [1, 16, 128, 256, 512, 768, 1024, 1280, 1536, 1792, 2048] + enable_userbuffers: true + max_num_streams: 3 + +moe_config: + backend: TRTLLM + use_low_precision_moe_combine: true + +sparse_attention_config: + algorithm: minimax_m3 + implementation: msa + indexer_kv_dtype: fp8 + sparse_disable_index_value: true + fuse_qkv_index_projection: true + +kv_cache_config: + free_gpu_memory_fraction: 0.94 + enable_block_reuse: true + tokens_per_block: 128 + use_kv_cache_manager_v2: true + dtype: fp8 + event_buffer_max_size: 0 + # CONNECTOR: GPU-only tiers. Mooncake takes over the CPU-offload role. + host_cache_size: 0 + disk_cache_size: 0 + +# MANDATORY, and not only because of the connector. M3 sets +# sparse_attention_config, so get_kv_cache_manager_cls routes to +# MiniMaxM3KVCacheManagerV2 unconditionally -- use_kv_cache_manager_v2 is not +# even consulted on that branch. KVCacheManagerV2 cannot drive the C++ +# transceiver, and M3 does not override get_preferred_transceiver_runtime, so +# 'auto' would resolve to C++ and be rejected. +cache_transceiver_config: + backend: "NIXL" + transceiver_runtime: "PYTHON" + +# CONNECTOR: was MAX_UTILIZATION; connectors require guaranteed-no-evict. +scheduler_config: + capacity_scheduler_policy: GUARANTEED_NO_EVICT + +enable_chunked_prefill: true +enable_autotuner: true +trust_remote_code: true +reasoning_parser: minimax_m3 +stream_interval: 20 +print_iter_log: true +num_postprocess_workers: 8 +enable_attention_dp: false + +kv_connector_config: + connector: mooncake-store diff --git a/mooncake_disagg/m3_disagg_config.yaml b/mooncake_disagg/m3_disagg_config.yaml new file mode 100644 index 000000000000..4e51aa45e154 --- /dev/null +++ b/mooncake_disagg/m3_disagg_config.yaml @@ -0,0 +1,20 @@ +# Router for the M3 mooncake-store validation. Worker LLM args come from each +# worker's own --config file, not from here. +# +# 1 ctx + 1 gen at TP=4 fits 8 GPUs. To demonstrate cross-instance reuse +# without restarting anything, raise context_servers.num_instances to 2 and +# add a second URL (needs 12 GPUs at TP=4). + +hostname: localhost +port: 8000 +backend: "pytorch" + +context_servers: + num_instances: 1 + urls: + - "localhost:8001" + +generation_servers: + num_instances: 1 + urls: + - "localhost:8002" diff --git a/mooncake_disagg/m3_gen_mooncake.yaml b/mooncake_disagg/m3_gen_mooncake.yaml new file mode 100644 index 000000000000..c6e84cf6d0c1 --- /dev/null +++ b/mooncake_disagg/m3_gen_mooncake.yaml @@ -0,0 +1,68 @@ +# MiniMax-M3-NVFP4 GENERATION (decode) worker: does not touch the store. +# +# There is deliberately no kv_connector_config. That absence is the whole of +# "decode-none" -- StoreRole has only producer / consumer / both, so there is +# no role value meaning "off". +# +# Because no connector runs here, three of the context worker's constraints +# do NOT apply, and this file keeps the production values instead. + +max_seq_len: 1048576 +max_num_tokens: 16384 +max_batch_size: 20 + +cuda_graph_config: + enable_padding: true + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20] + +torch_compile_config: + enable_fullgraph: true + enable_inductor: false + enable_piecewise_cuda_graph: true + capture_num_tokens: [1, 16, 128, 256, 512, 768, 1024, 1280, 1536, 1792, 2048] + enable_userbuffers: true + max_num_streams: 3 + +moe_config: + backend: TRTLLM + use_low_precision_moe_combine: true + +sparse_attention_config: + algorithm: minimax_m3 + implementation: msa + indexer_kv_dtype: fp8 + sparse_disable_index_value: true + fuse_qkv_index_projection: true + +kv_cache_config: + free_gpu_memory_fraction: 0.94 + enable_block_reuse: true + block_reuse_policy: per_conversation + tokens_per_block: 128 + use_kv_cache_manager_v2: true + dtype: fp8 + event_buffer_max_size: 0 + # Kept: no connector here, so _reject_non_gpu_cache_tiers never runs and + # decode keeps its host tier. Drop to 0 first if the Python transceiver + # misbehaves, since that pairing is the less-travelled path. + host_cache_size: 388554555392 + +# Must match the context worker: both ends of the handoff run the same +# transceiver, and M3's always-V2 manager rules out the C++ one. +cache_transceiver_config: + backend: "NIXL" + transceiver_runtime: "PYTHON" + +# Kept: the guaranteed-no-evict requirement is a connector gate, and no +# connector runs on this worker. +scheduler_config: + capacity_scheduler_policy: MAX_UTILIZATION + +enable_chunked_prefill: true +enable_autotuner: true +trust_remote_code: true +reasoning_parser: minimax_m3 +stream_interval: 20 +print_iter_log: true +num_postprocess_workers: 8 +enable_attention_dp: false diff --git a/mooncake_disagg/mooncake.json b/mooncake_disagg/mooncake.json new file mode 100644 index 000000000000..2a17f7d2d760 --- /dev/null +++ b/mooncake_disagg/mooncake.json @@ -0,0 +1,11 @@ +{ + "metadata_server": "P2PHANDSHAKE", + "master_server_address": "127.0.0.1:50051", + "protocol": "tcp", + "device_name": "", + "global_segment_size": "16GiB", + "local_buffer_size": "1GiB", + "role": "both", + "cache_prefix": "trtllm", + "transfer_batch_size": 64 +} diff --git a/mooncake_disagg/mooncake_api_surface_test.py b/mooncake_disagg/mooncake_api_surface_test.py new file mode 100644 index 000000000000..a359b7ea5903 --- /dev/null +++ b/mooncake_disagg/mooncake_api_surface_test.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Exercise every MooncakeDistributedStore method the connector calls. + +``mooncake_smoke_test.py`` only proves the install loads and can round-trip a +byte string. The connector's hot path never uses ``put``/``get``: it registers +the KV pools and then moves pages with the ``batch_*_multi_buffers`` zero-copy +calls. Those are the calls whose signatures could drift between wheel versions, +so this checks them against real registered GPU memory, in the same order +``worker.py`` uses them. + +Needs a running mooncake_master and MOONCAKE_CONFIG_PATH, same as the connector. +""" + +import json +import os +import socket +import sys + +import torch + +CONFIG_PATH = os.environ.get("MOONCAKE_CONFIG_PATH") +if not CONFIG_PATH: + sys.exit("Set MOONCAKE_CONFIG_PATH to the Mooncake JSON config first.") + +with open(CONFIG_PATH) as handle: + cfg = json.load(handle) + + +def parse_size(value): + if isinstance(value, int): + return value + units = {"KiB": 1 << 10, "MiB": 1 << 20, "GiB": 1 << 30, "TiB": 1 << 40} + for suffix, scale in units.items(): + if value.endswith(suffix): + return int(float(value[: -len(suffix)]) * scale) + return int(value) + + +import mooncake # noqa: E402 +from mooncake.store import MooncakeDistributedStore # noqa: E402 + +print(f"mooncake package: {mooncake.__path__[0]}") + +store = MooncakeDistributedStore() +hostname = cfg.get("local_hostname") or socket.gethostbyname(socket.gethostname()) +status = store.setup( + hostname, + cfg["metadata_server"], + parse_size(cfg.get("global_segment_size", "1GiB")), + parse_size(cfg.get("local_buffer_size", "256MiB")), + cfg.get("protocol", "tcp"), + cfg.get("device_name", ""), + cfg["master_server_address"], +) +assert status == 0, f"setup failed with status {status}" +print("setup: OK") + +# Stand in for a KV pool. PageAddressing.page_buffers returns one address per +# layer-group region, so a page is scattered across REGIONS buffers rather than +# contiguous -- which is why the batch calls take list[list[int]]. Model two +# strided regions so the scatter-gather path is actually exercised. +PAGES = 8 +REGIONS = 2 +REGION_BYTES = 128 * 1024 +STRIDE = REGION_BYTES # slots within a region are strided, as in the real layout +pool = torch.empty(REGIONS * PAGES * STRIDE, dtype=torch.uint8, device="cuda") +region_bases = [pool.data_ptr() + r * PAGES * STRIDE for r in range(REGIONS)] + +status = store.register_buffer(pool.data_ptr(), pool.numel()) +assert status == 0, f"register_buffer failed with status {status}" +print(f"register_buffer: OK ({pool.numel()} bytes of GPU memory at {pool.data_ptr():#x})") + +prefix = cfg.get("cache_prefix", "trtllm") +keys = [f"{prefix}/api-surface/page{i}" for i in range(PAGES)] +addresses = [[base + i * STRIDE for base in region_bases] for i in range(PAGES)] +sizes = [[REGION_BYTES] * REGIONS for _ in range(PAGES)] + +# Distinct content per (page, region), so a mixed-up address or size cannot pass. +view = pool.view(REGIONS, PAGES, STRIDE) +for r in range(REGIONS): + for i in range(PAGES): + view[r, i, :REGION_BYTES] = (i * 31 + r * 97 + 7) % 256 +expected = pool.clone() + +present = store.batch_is_exist(keys) +assert len(present) == PAGES, f"batch_is_exist returned {len(present)} of {PAGES}" +assert all(status != 1 for status in present), f"keys already present: {present}" +print(f"batch_is_exist (absent): OK {list(present)}") + +results = store.batch_put_from_multi_buffers(keys, addresses, sizes) +assert len(results) == PAGES, f"batch_put returned {len(results)} of {PAGES}" +bad = [(k, r) for k, r in zip(keys, results) if not isinstance(r, int) or r < 0] +assert not bad, f"batch_put_from_multi_buffers failed: {bad}" +print(f"batch_put_from_multi_buffers: OK {list(results)}") + +present = store.batch_is_exist(keys) +assert all(status == 1 for status in present), f"keys missing after put: {present}" +print("batch_is_exist (present): OK") + +pool.zero_() +results = store.batch_get_into_multi_buffers(keys, addresses, sizes) +assert len(results) == PAGES, f"batch_get returned {len(results)} of {PAGES}" +bad = [(k, r) for k, r in zip(keys, results) if not isinstance(r, int) or r < 0] +assert not bad, f"batch_get_into_multi_buffers failed: {bad}" +print(f"batch_get_into_multi_buffers: OK {list(results)}") + +torch.cuda.synchronize() +assert torch.equal(pool, expected), "page contents differ after the round trip" +print("GPU page contents byte-for-byte identical: OK") + +for key in keys: + store.remove(key) +store.close() +print("remove + close: OK") + +print("\nPASS: the full connector API surface works on this install.") diff --git a/mooncake_disagg/mooncake_smoke_test.py b/mooncake_disagg/mooncake_smoke_test.py new file mode 100644 index 000000000000..214ac0c46313 --- /dev/null +++ b/mooncake_disagg/mooncake_smoke_test.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Prove that a Mooncake install can actually serve the mooncake-store connector. + +Mirrors the ``store.setup(...)`` call in +``tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py`` and then +does one round trip, so a pass here means the connector's own startup path will +work. Reads the same ``MOONCAKE_CONFIG_PATH`` file the connector reads. +""" + +import json +import os +import socket +import sys + +CONFIG_PATH = os.environ.get("MOONCAKE_CONFIG_PATH") +if not CONFIG_PATH: + sys.exit("Set MOONCAKE_CONFIG_PATH to the Mooncake JSON config first.") + +with open(CONFIG_PATH) as handle: + cfg = json.load(handle) + + +def parse_size(value): + if isinstance(value, int): + return value + units = {"KiB": 1 << 10, "MiB": 1 << 20, "GiB": 1 << 30, "TiB": 1 << 40} + for suffix, scale in units.items(): + if value.endswith(suffix): + return int(float(value[: -len(suffix)]) * scale) + return int(value) + + +from mooncake.store import MooncakeDistributedStore # noqa: E402 + +print("import mooncake.store: OK") + +store = MooncakeDistributedStore() +hostname = cfg.get("local_hostname") or socket.gethostbyname(socket.gethostname()) +status = store.setup( + hostname, + cfg["metadata_server"], + parse_size(cfg.get("global_segment_size", "1GiB")), + parse_size(cfg.get("local_buffer_size", "256MiB")), + cfg.get("protocol", "tcp"), + cfg.get("device_name", ""), + cfg["master_server_address"], +) +if status != 0: + sys.exit(f"store.setup failed with status {status}") +print(f"store.setup: OK (host={hostname}, master={cfg['master_server_address']})") + +key = f"{cfg.get('cache_prefix', 'trtllm')}/smoke-test" +payload = bytes(range(256)) * 4096 # 1 MiB, non-trivial content + +assert store.put(key, payload) == 0, "put failed" +print(f"put {len(payload)} bytes: OK") + +assert store.is_exist(key) == 1, "is_exist did not report the key" +print("is_exist: OK") + +got = store.get(key) +assert got == payload, f"round trip mismatch: got {len(got)} bytes" +print("get + byte-for-byte compare: OK") + +store.remove(key) +print("remove: OK") +print("\nPASS: this install can back the mooncake-store connector.") diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 1345c2af1e2e..18b784e62495 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1409,6 +1409,13 @@ "kind": "value", "path": "sparse_attention_config.enable_heuristic_topk" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.fuse_qkv_index_projection" + }, { "allowed_values": [ "triton", diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index dd507da4efd4..292259c209ab 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -19,6 +19,7 @@ plain integers, which is all the addressing arithmetic needs. """ +import contextlib import json from types import SimpleNamespace @@ -80,7 +81,8 @@ def __init__(self): self.exist_calls = [] self.closed = False self.fail_gets_for = set() - #: Workers built against this store, torn down by the fixture. + #: Workers built against this store. ``make_worker`` shuts each one + #: down; the fixture repeats it as a backstop for early failures. self.workers = [] def register_buffer(self, address, size): @@ -183,12 +185,22 @@ def fake_store(monkeypatch): worker_module._LOCAL_WORKER_READY.clear() +@contextlib.contextmanager def make_worker(fake_store, *, layout=None): + """Build a worker and shut it down before the test call phase ends. + + Registering a layout starts the background save thread, and + pytest-threadleak snapshots threads around the call phase only, so + fixture teardown would run too late to keep it quiet. + """ worker = MooncakeStoreConnectorWorker(make_llm_args()) fake_store.workers.append(worker) if layout is not None: worker.register_kv_cache_layout(layout) - return worker + try: + yield worker + finally: + worker.shutdown() def make_request(request_id, tokens, cache_salt=None): @@ -419,111 +431,111 @@ def test_validate_layout_rejects_sliding_window(): def test_worker_registers_every_pool_range(store_config, fake_store): layout = make_layout(num_groups=2, regions_per_group=2) - worker = make_worker(fake_store, layout=layout) - assert fake_store.registered == [ - (start, end - start) for start, end in PageAddressing(layout).registration_ranges() - ] - assert worker.is_registered + with make_worker(fake_store, layout=layout) as worker: + assert fake_store.registered == [ + (start, end - start) for start, end in PageAddressing(layout).registration_ranges() + ] + assert worker.is_registered def test_worker_rejects_v1_pool_registration(store_config, fake_store): - worker = make_worker(fake_store) - with pytest.raises(NotImplementedError, match="KVCacheManagerV2"): - worker.register_kv_caches(None) + with make_worker(fake_store) as worker: + with pytest.raises(NotImplementedError, match="KVCacheManagerV2"): + worker.register_kv_caches(None) def test_worker_prefix_hit_needs_every_layer_group(store_config, fake_store): layout = make_layout(num_groups=2) - worker = make_worker(fake_store, layout=layout) - hashes = [bytes([index]) * 16 for index in range(3)] + with make_worker(fake_store, layout=layout) as worker: + hashes = [bytes([index]) * 16 for index in range(3)] - assert worker.count_prefix_hit(hashes) == 0 + assert worker.count_prefix_hit(hashes) == 0 - # Populate blocks 0 and 1 completely, and block 2 only partially. - for block in range(2): - for group_id in range(2): - fake_store.objects.add(worker._namespaces[group_id].key(hashes[block])) - fake_store.objects.add(worker._namespaces[0].key(hashes[2])) + # Populate blocks 0 and 1 completely, and block 2 only partially. + for block in range(2): + for group_id in range(2): + fake_store.objects.add(worker._namespaces[group_id].key(hashes[block])) + fake_store.objects.add(worker._namespaces[0].key(hashes[2])) - assert worker.count_prefix_hit(hashes) == 2 + assert worker.count_prefix_hit(hashes) == 2 def test_worker_prefix_hit_stops_at_the_first_gap(store_config, fake_store): - worker = make_worker(fake_store, layout=make_layout()) - hashes = [bytes([index]) * 16 for index in range(3)] - # Block 1 missing: block 2 is unusable even though it is present, because a - # prefix is replayed contiguously. - fake_store.objects.add(worker._namespaces[0].key(hashes[0])) - fake_store.objects.add(worker._namespaces[0].key(hashes[2])) - assert worker.count_prefix_hit(hashes) == 1 + with make_worker(fake_store, layout=make_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(3)] + # Block 1 missing: block 2 is unusable even though it is present, because a + # prefix is replayed contiguously. + fake_store.objects.add(worker._namespaces[0].key(hashes[0])) + fake_store.objects.add(worker._namespaces[0].key(hashes[2])) + assert worker.count_prefix_hit(hashes) == 1 def test_worker_load_raises_when_a_page_is_missing(store_config, fake_store): - worker = make_worker(fake_store, layout=make_layout()) - transfers = RequestTransfers(7, [PageTransfer(b"\x00" * 16, 0, 1)]) - worker.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) - with pytest.raises(RuntimeError, match="already"): - worker.start_load_kv(None) + with make_worker(fake_store, layout=make_layout()) as worker: + transfers = RequestTransfers(7, [PageTransfer(b"\x00" * 16, 0, 1)]) + worker.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) + with pytest.raises(RuntimeError, match="already"): + worker.start_load_kv(None) def test_worker_load_addresses_the_requested_page(store_config, fake_store): layout = make_layout(regions_per_group=2) - worker = make_worker(fake_store, layout=layout) - block_hash = b"\x00" * 16 - key = worker._namespaces[0].key(block_hash) - fake_store.objects.add(key) + with make_worker(fake_store, layout=layout) as worker: + block_hash = b"\x00" * 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(SimpleNamespace(loads=[transfers], saves=[])) - worker.start_load_kv(None) + transfers = RequestTransfers(7, [PageTransfer(block_hash, 0, 3)]) + worker.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) + worker.start_load_kv(None) - (keys, addresses, sizes) = fake_store.get_calls[0] - expected_addresses, expected_sizes = PageAddressing(layout).buffers(0, 3) - assert keys == [key] - assert addresses == [expected_addresses] - assert sizes == [expected_sizes] + (keys, addresses, sizes) = fake_store.get_calls[0] + expected_addresses, expected_sizes = PageAddressing(layout).buffers(0, 3) + assert keys == [key] + assert addresses == [expected_addresses] + assert sizes == [expected_sizes] def test_worker_save_skips_pages_already_in_the_store(store_config, fake_store): - worker = make_worker(fake_store, layout=make_layout()) - hashes = [bytes([index]) * 16 for index in range(2)] - fake_store.objects.add(worker._namespaces[0].key(hashes[0])) - - worker._put( - [ - RequestTransfers( - 1, - [PageTransfer(hashes[0], 0, 0), PageTransfer(hashes[1], 0, 1)], - ) - ] - ) - assert len(fake_store.put_calls) == 1 - assert fake_store.put_calls[0][0] == [worker._namespaces[0].key(hashes[1])] + with make_worker(fake_store, layout=make_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(2)] + fake_store.objects.add(worker._namespaces[0].key(hashes[0])) + + worker._put( + [ + RequestTransfers( + 1, + [PageTransfer(hashes[0], 0, 0), PageTransfer(hashes[1], 0, 1)], + ) + ] + ) + assert len(fake_store.put_calls) == 1 + assert fake_store.put_calls[0][0] == [worker._namespaces[0].key(hashes[1])] def test_worker_reports_a_request_finished_once_its_saves_drain(store_config, fake_store): - worker = make_worker(fake_store, layout=make_layout()) + with make_worker(fake_store, layout=make_layout()) as worker: + # One submission outstanding: the request is closed but must not be released. + worker._outstanding_saves[42] = 1 + assert worker.get_finished([42], []) == ([], []) - # One submission outstanding: the request is closed but must not be released. - worker._outstanding_saves[42] = 1 - assert worker.get_finished([42], []) == ([], []) - - worker._outstanding_saves.pop(42) - assert worker.get_finished([], []) == ([42], []) - # Reported once only. - assert worker.get_finished([], []) == ([], []) + worker._outstanding_saves.pop(42) + assert worker.get_finished([], []) == ([42], []) + # Reported once only. + assert worker.get_finished([], []) == ([], []) def test_worker_reports_a_request_with_no_saves_immediately(store_config, fake_store): - worker = make_worker(fake_store, layout=make_layout()) - assert worker.get_finished([9], [5]) == ([9], [5]) + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker.get_finished([9], [5]) == ([9], [5]) def test_worker_shutdown_closes_the_store(store_config, fake_store): - worker = make_worker(fake_store, layout=make_layout()) - worker.shutdown() - assert fake_store.closed - worker.shutdown() + with make_worker(fake_store, layout=make_layout()) as worker: + worker.shutdown() + assert fake_store.closed + # Idempotent: a second call must not raise or reopen anything. + worker.shutdown() # ---- scheduler ---- From 32625df74ab7140cc01859ded67179f832498dff Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 27 Aug 2026 14:53:08 -0700 Subject: [PATCH 07/24] [None][feat] Span the Mooncake store pool across prefill and decode nodes Pool capacity comes only from processes that open a store handle, and only the context workers configure the KV connector, so every byte of the pool was prefill-node memory. That made the store a prefill-DRAM-caches-prefill-GPU tier largely duplicating TensorRT-LLM's native host offload, rather than the cross-node pool it is meant to be. Add a capacity-only donor process per generation node: it contributes host memory and then idles, never issuing a put or get. Prefill-written KV can then live on decode-side DRAM while the generation engine stays connector-free, so it keeps its cache transceiver for the prefill-to-decode handoff. A donor is a separate process rather than a new StoreRole because the roles describe traffic -- producer writes, consumer reads, both does both -- and none of them means "contribute memory only". Also bring the master and client config up inside the benchmark job so a run needs no manual setup, resolve MOONCAKE_CONFIG_PATH from the log directory whose path is not known when the worker environment is built, and report block placement grouped by segment host. That last figure is what distinguishes a spanning pool from a prefill-only one: a single host means the donors are absent or not being allocated into. Verified on MiniMax-M3 at TP=2 for both prefill and decode: the pool grew from 32GiB on one host to 64GiB across two, and of 4.58GiB written by prefill, 1.45GiB (32%) landed on the decode node, with no load or save failures and no change in throughput or latency. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../slurm/benchmark/disaggr_torch.slurm | 258 +++++++++++++++++- .../slurm/benchmark/start_worker.sh | 10 + .../slurm/benchmark/watch_job.sh | 105 +++++++ mooncake_disagg/README.md | 111 ++++++-- mooncake_disagg/mooncake_segment_donor.py | 223 +++++++++++++++ 5 files changed, 686 insertions(+), 21 deletions(-) create mode 100644 examples/disaggregated/slurm/benchmark/watch_job.sh create mode 100644 mooncake_disagg/mooncake_segment_donor.py diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm index 4472b07b2ea8..d72f756832e8 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm @@ -101,7 +101,7 @@ elif [ -d "${trtllm_repo}" ]; then if [ "${build_wheel}" = "true" ]; then echo "Building TensorRT-LLM wheel on one node..." - build_command="python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --use_ccache --clean" + build_command="python3 ./scripts/build_wheel.py --use_ccache" if [ -n "${cuda_architectures:-}" ]; then build_command="${build_command} --cuda_architectures \"${cuda_architectures}\"" fi @@ -144,14 +144,21 @@ fi # to be reapplied per job: --container-name gives each node a container that # lives for the job, so anything installed here survives to the worker sruns but # not into the next job. +mooncake_enabled=false if grep -qs "mooncake-store" "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"; then + mooncake_enabled=true mooncake_install_script="" + mooncake_donor_script="" if [ -n "${trtllm_repo:-}" ]; then mooncake_install_script="${trtllm_repo}/mooncake_disagg/install_mooncake_runtime.sh" + mooncake_donor_script="${trtllm_repo}/mooncake_disagg/mooncake_segment_donor.py" fi if [ -z "${mooncake_install_script}" ] || [ ! -f "${mooncake_install_script}" ]; then cleanup_on_failure "A worker config requests the mooncake-store connector, but mooncake_disagg/install_mooncake_runtime.sh was not found under trtllm_repo='${trtllm_repo:-}'. Set environment.trtllm_repo to a checkout that contains it, or bake the bindings into the container image." fi + if [ "${MOONCAKE_DONOR_SEGMENT_SIZE:-32GiB}" != "0" ] && [ ! -f "${mooncake_donor_script}" ]; then + cleanup_on_failure "mooncake_disagg/mooncake_segment_donor.py was not found under trtllm_repo='${trtllm_repo:-}'. It contributes the generation nodes' host memory to the pool; without it the pool is prefill-node memory only. Set MOONCAKE_DONOR_SEGMENT_SIZE=0 to accept that and skip the donors." + fi echo "Installing Mooncake store bindings on all nodes..." if ! srun --container-name=${container_name} \ --container-mounts=${container_mount} --no-container-mount-home \ @@ -179,6 +186,178 @@ client_cmds_base_file=${full_logdir}/client_cmds_base.sh client_cmds_file=${full_logdir}/client_cmds.sh replace_placeholder "${client_cmds_base_file}" "${all_nodes_str}" "${client_cmds_file}" +# Bring up the Mooncake store's pool: a master process plus the client config +# that names it. Both are per job and land in the log directory, so nothing has +# to be prepared by hand and concurrent jobs do not share state. +# +# The master lives inside this job, which means the pool dies with it. That is +# the right default for the single-job experiments, but it makes the +# survival-across-restarts case impossible to test, so set +# MOONCAKE_MASTER_ADDRESS in the submitting environment to reuse a master +# running as its own longer-lived job (mooncake_disagg/README.md section 4); +# this block then only writes the client config. +if [ "${mooncake_enabled}" = "true" ]; then + mooncake_master_port=50051 + mooncake_master_metrics_port=9004 + mooncake_master_addr="${MOONCAKE_MASTER_ADDRESS:-}" + + if [ -z "${mooncake_master_addr}" ]; then + mooncake_master_node="${all_nodes[0]}" + mooncake_addr_file="${full_logdir}/mooncake_master.addr" + rm -f "${mooncake_addr_file}" + echo "Starting mooncake_master on ${mooncake_master_node}..." + # --overlap because every GPU on this node is already claimed by a + # worker; the master is a CPU-only process sharing the node with them. + # + # mooncake_master logs through glog, which writes to files under /tmp + # inside the container unless told otherwise -- so without + # GLOG_logtostderr the log below would be empty. GLOG_v=1 adds the + # per-RPC lines that show workers registering segments and putting and + # getting keys, which is the only view of the pool's side of the + # conversation short of scraping the metrics port. + srun -l --container-name=${container_name} \ + --container-mounts=${container_mount} --no-container-mount-home \ + --mpi=pmix --overlap --nodelist=${mooncake_master_node} -N 1 -n 1 \ + bash -c "export GLOG_logtostderr=1 GLOG_v=\"${MOONCAKE_MASTER_GLOG_V:-1}\"; \ + addr=\$(hostname -I 2>/dev/null | awk '{print \$1}'); \ + [ -n \"\${addr}\" ] || addr=\$(hostname -f); \ + echo \"\${addr}\" > ${mooncake_addr_file}; \ + echo \"mooncake_master starting on \$(hostname) at \${addr}:${mooncake_master_port} (metrics ${mooncake_master_metrics_port})\"; \ + exec mooncake_master --rpc_port=${mooncake_master_port} --metrics_port=${mooncake_master_metrics_port} --eviction_ratio=0.05" \ + &> ${full_logdir}/2_mooncake_master.log & + + for _ in $(seq 1 60); do + if [ -s "${mooncake_addr_file}" ]; then + break + fi + sleep 1 + done + if [ ! -s "${mooncake_addr_file}" ]; then + cleanup_on_failure "mooncake_master did not report its address within 60s. Check ${full_logdir}/2_mooncake_master.log for details" + fi + mooncake_master_addr="$(tr -d '[:space:]' < ${mooncake_addr_file}):${mooncake_master_port}" + + # A worker that opens its store handle before the master accepts + # connections fails at startup, so wait for the port rather than the + # process. /dev/tcp keeps this dependency-free. + mooncake_master_ready=false + mooncake_wait_start=${SECONDS} + for _ in $(seq 1 60); do + if (exec 3<>/dev/tcp/${mooncake_master_addr%:*}/${mooncake_master_port}) 2>/dev/null; then + mooncake_master_ready=true + break + fi + sleep 1 + done + if [ "${mooncake_master_ready}" != "true" ]; then + cleanup_on_failure "mooncake_master at ${mooncake_master_addr} is not accepting connections. Check ${full_logdir}/2_mooncake_master.log for details" + fi + echo "mooncake_master ready at ${mooncake_master_addr} on node ${mooncake_master_node} after $((SECONDS - mooncake_wait_start))s of waiting" + echo "mooncake_master metrics: http://${mooncake_master_addr%:*}:${mooncake_master_metrics_port}" + echo "mooncake_master log: ${full_logdir}/2_mooncake_master.log" + else + echo "Using externally managed mooncake_master at ${mooncake_master_addr}" + fi + + # The schema is vLLM's, so one pool can serve both engines. Defaults are + # the first-bring-up ones: TCP removes RDMA from the variable list, at the + # cost of any performance conclusion. Set MOONCAKE_PROTOCOL=rdma with a + # MOONCAKE_DEVICE_NAME from ibv_devinfo for a run worth quoting. + # global_segment_size is contributed per worker process, so the pool is + # this value times (ctx instances x world size). + cat > "${full_logdir}/mooncake.json" < ${full_logdir}/2_mooncake_donor_${donor_node}.log & + done + + # Wait for the segments before any worker starts, so that the first + # blocks prefill writes can already be placed on a decode node. A donor + # that never mounts is a hard failure: the run would otherwise quietly + # fall back to the prefill-only pool this is meant to replace. + for donor_ready_file in "${mooncake_donor_ready_files[@]}"; do + for _ in $(seq 1 120); do + if [ -s "${donor_ready_file}" ]; then + break + fi + sleep 1 + done + if [ ! -s "${donor_ready_file}" ]; then + donor_node="$(basename "${donor_ready_file}" .ready)" + donor_node="${donor_node#mooncake_donor_}" + cleanup_on_failure "The memory donor on ${donor_node} did not mount its segment within 120s. Check ${full_logdir}/2_mooncake_donor_${donor_node}.log; if the node is short on memory, lower MOONCAKE_DONOR_SEGMENT_SIZE or the generation worker's kv_cache_config.host_cache_size" + fi + done + echo "Pool capacity: ${MOONCAKE_GLOBAL_SEGMENT_SIZE:-16GiB} per context worker process" \ + "+ ${mooncake_donor_size} per generation node (${#mooncake_donor_nodes[@]} donor(s):" \ + "${mooncake_donor_nodes[*]})" + fi +fi + # Per-worker hostfile / gpu_map files for srun --distribution=arbitrary. # submit.py emits *_base.txt with ; rewrite them here. for base_file in "${full_logdir}"/hostfile_*_base.txt "${full_logdir}"/gpu_map_*_base.txt; do @@ -201,6 +380,20 @@ cat ${start_server_cmds_file} | while read cmd; do done echo "Server is ready!" +# A connector that failed to open its store handle does not stop the worker from +# serving, it just silently never hits, so surface the startup lines here rather +# than leaving them to be discovered after the benchmark. The registration line +# carries the bytes/page figure the pool sizing depends on. +if [ "${mooncake_enabled}" = "true" ]; then + echo "Mooncake store startup lines from the context workers:" + if ! grep -h "mooncake-store" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null; then + echo " WARNING: no mooncake-store lines found. The connector may not have" \ + "loaded; check ${full_logdir}/3_output_CTX_*.log and" \ + "${full_logdir}/2_mooncake_master.log. The benchmark will still run," \ + "but without the store." + fi +fi + # Start client commands echo "Starting client commands from ${client_cmds_file}..." while read -r cmd <&3; do @@ -211,6 +404,69 @@ while read -r cmd <&3; do fi done 3< "${client_cmds_file}" +# Collect the store's traffic into one file. The per-event lines live at DEBUG in +# the worker logs (module _torch), so this is only populated when a config asks +# for that verbosity; the counts are what distinguish "the store ran" from "the +# store ran and did something". +if [ "${mooncake_enabled}" = "true" ]; then + mooncake_summary="${full_logdir}/9_mooncake_summary.log" + { + echo "master: ${mooncake_master_addr}" + echo + echo "== startup ==" + grep -h "mooncake-store.*\(ready\|registered layout\)" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null || echo "(none)" + echo + echo "== event counts ==" + for pattern in "matched" "loaded" "failed to load" "failed to save" \ + "lookup failed" "could not reserve connector prefix"; do + count=$(grep -h "mooncake-store.*${pattern}" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null | wc -l || true) + echo "${pattern}: ${count}" + done + echo + # Where the blocks physically went. The master names the segment for + # every allocation, and a segment is one client process's donated + # memory, so grouping by segment host answers the question the donors + # exist for: how much of the pool's contents lives on a decode node + # rather than on the prefill node that computed it. Without a donor this + # section shows a single host, which is the prefill node. + echo "== block placement by segment host ==" + echo "(donor hosts: $(cat "${full_logdir}"/mooncake_donor_*.ready 2>/dev/null | awk '{print $1}' | paste -sd, - || echo none))" + grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" \ + "${full_logdir}/2_mooncake_master.log" 2>/dev/null \ + | awk '{ + sub(/size=/, "", $2); sub(/segment=/, "", $3); + split($3, parts, ":"); host = parts[1]; port = parts[2]; + allocs[host]++; bytes[host] += $2; + if (!((host, port) in seen)) { + seen[host, port] = 1; + segs[host] = segs[host] " " port; + } + total_allocs++; total_bytes += $2; + } + END { + if (total_allocs == 0) { print "(no allocations)"; exit } + for (h in allocs) { + printf "%-16s pages=%-7d %8.2f GiB %5.1f%% of pool contents segments:%s\n", + h, allocs[h], bytes[h] / 1073741824, + 100 * bytes[h] / total_bytes, segs[h]; + } + printf "%-16s pages=%-7d %8.2f GiB\n", "TOTAL", total_allocs, total_bytes / 1073741824; + }' || echo "(could not parse master log)" + echo + echo "== donors ==" + for donor_log in "${full_logdir}"/2_mooncake_donor_*.log; do + [ -f "${donor_log}" ] || continue + echo "--- $(basename "${donor_log}") ---" + grep -h "donating\|mounted\|failed\|exited" "${donor_log}" 2>/dev/null || tail -n 5 "${donor_log}" + done + echo + echo "== master ==" + tail -n 50 "${full_logdir}/2_mooncake_master.log" 2>/dev/null || echo "(no master log)" + } > "${mooncake_summary}" 2>&1 + echo "Mooncake store summary written to ${mooncake_summary}" + cat "${mooncake_summary}" +fi + echo "Job completed successfully, total runtime: $SECONDS seconds" # try to kill the server and workers diff --git a/examples/disaggregated/slurm/benchmark/start_worker.sh b/examples/disaggregated/slurm/benchmark/start_worker.sh index 0a5b5897b773..fa8eca6425ae 100644 --- a/examples/disaggregated/slurm/benchmark/start_worker.sh +++ b/examples/disaggregated/slurm/benchmark/start_worker.sh @@ -54,6 +54,16 @@ fi echo "config_file: ${config_file}" +# The mooncake-store KV connector reads its pool topology from +# MOONCAKE_CONFIG_PATH. disaggr_torch.slurm generates one per job in the log +# directory, whose path is not known when submit.py builds the worker +# environment; an explicit setting still wins, so pointing at an externally +# managed pool remains possible. +if [ -z "${MOONCAKE_CONFIG_PATH:-}" ] && [ -f "${log_dir}/mooncake.json" ]; then + export MOONCAKE_CONFIG_PATH="${log_dir}/mooncake.json" + echo "MOONCAKE_CONFIG_PATH: ${MOONCAKE_CONFIG_PATH}" +fi + nsys_prefix="" if [ "${enable_nsys}" != "true" ]; then echo "nsys is not enabled, start normal flow" diff --git a/examples/disaggregated/slurm/benchmark/watch_job.sh b/examples/disaggregated/slurm/benchmark/watch_job.sh new file mode 100644 index 000000000000..613f3714e68d --- /dev/null +++ b/examples/disaggregated/slurm/benchmark/watch_job.sh @@ -0,0 +1,105 @@ +#!/bin/bash +# Poll a disaggregated benchmark log directory and print one EVENT line per +# state change. Intended for tracking a running job without tailing megabytes +# of worker log. +# +# bash watch_job.sh [poll_seconds] [max_minutes] +set -uo pipefail + +log_dir="${1:?usage: watch_job.sh [poll_seconds] [max_minutes]}" +poll="${2:-20}" +max_minutes="${3:-60}" +deadline=$(( SECONDS + max_minutes * 60 )) + +declare -A seen + +announce() { + key="$1"; shift + if [ -z "${seen[$key]:-}" ]; then + seen[$key]=1 + echo "EVENT: $*" + fi +} + +count_matches() { + grep -h "$1" "${log_dir}"/3_output_CTX_*.log 2>/dev/null | wc -l || true +} + +# How much of the pool's contents lives on each node. A segment is one client +# process's donated memory, so a single host here means the pool is prefill-only +# and the memory donors are either absent or not being allocated into. +placement() { + grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" \ + "${log_dir}/2_mooncake_master.log" 2>/dev/null \ + | awk '{sub(/size=/,"",$2); sub(/segment=/,"",$3); split($3,p,":"); + n[p[1]]++; b[p[1]]+=$2; t+=$2} + END {if (t == 0) {print "(none yet)"; exit} + for (h in n) printf "%s:%d pages/%.2fGiB/%.0f%% ", h, n[h], b[h]/1073741824, 100*b[h]/t}' +} + +while [ ${SECONDS} -lt ${deadline} ]; do + for ready in "${log_dir}"/mooncake_donor_*.ready; do + [ -s "${ready}" ] || continue + node="$(basename "${ready}" .ready)"; node="${node#mooncake_donor_}" + announce "donor_${node}" "memory donor on ${node} mounted its segment:" \ + "$(cat "${ready}")" + done + + for role in CTX GEN; do + f="${log_dir}/3_output_${role}_0.log" + [ -f "$f" ] || continue + grep -qs "Server started at\|Application startup complete\|Uvicorn running" "$f" \ + && announce "${role}_up" "${role} worker serving" + done + + if grep -qs "registered layout" "${log_dir}"/3_output_CTX_*.log; then + announce "registered" "connector registered KV layout:" \ + "$(grep -h "registered layout" "${log_dir}"/3_output_CTX_*.log 2>/dev/null | head -n 1 | cut -c1-400)" + fi + + matched=$(count_matches "mooncake-store matched") + loaded=$(count_matches "mooncake-store rank") + [ "${matched:-0}" -gt 0 ] && announce "first_match" "first store hit; matched lines=${matched}" + [ "${loaded:-0}" -gt 0 ] && announce "first_load" "first store load; loaded lines=${loaded}" + + # The point of the donors: report as soon as a second host appears, since + # that is the first moment prefill-written KV is provably on decode DRAM. + hosts=$(grep -o "segment=[0-9.]*:" "${log_dir}/2_mooncake_master.log" 2>/dev/null \ + | sort -u | wc -l || true) + [ "${hosts:-0}" -ge 2 ] && announce "multi_host" \ + "pool spans ${hosts} hosts; placement: $(placement)" + + for pattern in "failed to load" "failed to save" "lookup failed" "background save failed"; do + if grep -qs "mooncake-store.*${pattern}" "${log_dir}"/3_output_CTX_*.log; then + announce "fail_${pattern// /_}" "PROBLEM: mooncake-store ${pattern}" + fi + done + + [ -f "${log_dir}/6_bench.log" ] && announce "bench_started" "benchmark client started" + ls "${log_dir}"/concurrency_*/result.json >/dev/null 2>&1 \ + && announce "result" "result.json written" + + if ls "${log_dir}"/8_done_*.txt >/dev/null 2>&1; then + echo "EVENT: job finished (8_done marker present)" + echo "FINAL: matched=${matched:-0} loaded=${loaded:-0}" + echo "FINAL placement: $(placement)" + exit 0 + fi + + # A dead batch script leaves the tree untouched; report it rather than + # polling a corpse until the deadline. + if grep -qs "Job completed successfully" "${log_dir}"/slurm-*.out; then + echo "EVENT: batch script reported completion" + exit 0 + fi + if grep -qs "^Error: " "${log_dir}"/slurm-*.out; then + echo "EVENT: PROBLEM: batch script hit cleanup_on_failure" + grep -h "^Error: " "${log_dir}"/slurm-*.out | tail -n 3 + exit 1 + fi + + sleep "${poll}" +done + +echo "EVENT: watcher deadline reached after ${max_minutes} minutes" +echo "FINAL: matched=$(count_matches 'mooncake-store matched') loaded=$(count_matches 'mooncake-store rank')" diff --git a/mooncake_disagg/README.md b/mooncake_disagg/README.md index f34b3a042ec1..7ed4c20c12ee 100644 --- a/mooncake_disagg/README.md +++ b/mooncake_disagg/README.md @@ -262,19 +262,19 @@ pytest tests/unittest/_torch/executor/test_mooncake_store_connector.py ``` mooncake_master (1 CPU core, its own job) - ^ ^ - register/put/get| | - ┌─────────────────────┴──┐ ┌──┴──────────────────────┐ - │ CTX instance 0 TP=4 │ │ CTX instance 1 TP=4 │ store: role=both - └────────────┬───────────┘ └───────────┬─────────────┘ - │ NIXL KV handoff │ - └──────────┬────────────────┘ - v - ┌──────────────────────────┐ - │ GEN instance TP=4 │ no connector at all - └──────────────────────────┘ - ^ - round-robin│ + ^ ^ ^ + register/put/get| | |mount segment only + ┌─────────────────────┴──┐ ┌──┴──────────────────────┐ │ + │ CTX instance 0 TP=4 │ │ CTX instance 1 TP=4 │ │ store: role=both + └────────────┬───────────┘ └───────────┬─────────────┘ │ + │ NIXL KV handoff │ │ + └──────────┬────────────────┘ │ + v │ + ┌──────────────────────────┐ ┌──────────────┴────────────┐ + │ GEN instance TP=4 │ │ segment donor (same node) │ + └──────────────────────────┘ └───────────────────────────┘ + ^ no connector, no put/get, + round-robin│ contributes host memory ┌──────────┴───────────┐ │ trtllm-serve disagg │ <- benchmark_serving client └──────────────────────┘ @@ -286,13 +286,58 @@ absence is the only way to express "off" (`StoreRole` has no off value). It also lets the generation worker keep its host cache tier and `MAX_UTILIZATION` scheduler, both of which the connector would forbid. +### Why the generation node still needs a donor process + +Pool capacity comes only from processes that open a store handle: `setup` +registers `global_segment_size` bytes of the calling process's host memory, and +the master then places blocks in it. Since only the context workers configure +the connector, only they contribute memory — so by default every byte of the +pool is prefill-node DRAM, and the store is a prefill-DRAM-caches-prefill-GPU +tier that largely duplicates TensorRT-LLM's native host offload. Confirm this on +any run by grouping the master's `allocation_succeeded ... segment=:` +lines by host: a single host means a prefill-only pool. + +`mooncake_segment_donor.py` closes that gap. One donor per generation node opens +a handle, contributes memory, and then idles forever without a single put or get, +so the pool spans both sides while the generation engine stays connector-free +and keeps its cache transceiver for the KV handoff. A donor is deliberately not +a `StoreRole`: the roles describe traffic (`producer` writes, `consumer` reads, +`both`), and none of them means "contribute memory only", so capacity and +traffic have to be separate processes. + +`disaggr_torch.slurm` starts the donors automatically, reading the generation +nodes off the generated worker commands and waiting for each segment to mount +before any worker starts — so the first blocks prefill writes can already land +on a decode node. Tune with `MOONCAKE_DONOR_SEGMENT_SIZE` (default `32GiB`, set +`0` to keep the pool prefill-only) and `MOONCAKE_DONOR_NODES` to override node +selection. + +The donated memory is charged to the donor process and competes with the +generation worker's own `kv_cache_config.host_cache_size` on that node, so size +the two together. The worker logs its own share as `KV cache manager v2 host +cache quota set to N GiB`, **per rank**, against the `available host memory` it +reports on the same line. + ## 4. Step 1 -- run the Mooncake master `master_server_address` is mandatory, so a master must exist and be reachable -from every worker. The benchmark harness has no hook for launching a side -process, and the workers read `MOONCAKE_CONFIG_PATH` at startup, so the master's -address has to be known *before* the benchmark job is submitted. Run it as its -own long-lived job: +from every worker. + +**For a single-job experiment you can skip this section.** +`disaggr_torch.slurm` now starts a `mooncake_master` on the first node of the +allocation, waits for its port to accept connections, and writes +`/mooncake.json` naming it; `start_worker.sh` then resolves +`MOONCAKE_CONFIG_PATH` from the log directory, which is why the harness config +does not set it. Defaults are the bring-up ones (TCP, 16GiB per worker); +`MOONCAKE_PROTOCOL`, `MOONCAKE_DEVICE_NAME`, `MOONCAKE_GLOBAL_SEGMENT_SIZE` and +`MOONCAKE_LOCAL_BUFFER_SIZE` in the submitting environment override them, and +the master's own log lands in `/2_mooncake_master.log`. + +That master dies with the job, so read on if you need a pool that outlives one +allocation -- which experiment 3 does, by construction. Run it as its own +long-lived job and export `MOONCAKE_MASTER_ADDRESS=:50051` before +`submit.py`; the harness then skips launching one and only writes the client +config pointing at yours. ```bash # mooncake_master.sbatch @@ -323,7 +368,8 @@ Keeping the master in a separate job is what makes experiment 3 (§7) possible: the pool outlives the engines, so a second benchmark job finds a warm store. Then write the client config, substituting the address the master job just -recorded. The schema is vLLM's, so one pool can serve both engines: +recorded -- or let `disaggr_torch.slurm` generate it, as above. The schema is +vLLM's, so one pool can serve both engines: ```bash MASTER_IP=$(cat $WORK_DIR/master.addr) @@ -415,8 +461,10 @@ environment: trtllm_wheel_path: "" work_dir: "" worker_env_var: "TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes NCCL_GRAPH_MIXING_SUPPORT=0" - # Only the context workers open a store handle. - ctx_worker_env_var: "MOONCAKE_CONFIG_PATH=/mooncake.json TRTLLM_MOONCAKE_STORE_ROLE=both TRTLLM_MOONCAKE_STORE_PREFIX=trtllm-m3-run1" + # Only the context workers open a store handle. MOONCAKE_CONFIG_PATH is + # deliberately absent: the harness generates the file per job in the log + # directory, whose path is not known when submit.py builds this environment. + ctx_worker_env_var: "TRTLLM_MOONCAKE_STORE_ROLE=both TRTLLM_MOONCAKE_STORE_PREFIX=trtllm-m3-run1" server_env_var: "TRTLLM_SERVER_DISABLE_GC=1" profiling: @@ -688,6 +736,29 @@ The store's own counters are the alternative that costs nothing at runtime: `mooncake_master --metrics_port=9004` exposes pool-level statistics over HTTP. Scrape it before and after a run and diff. +### Where did the pages land? + +Page counts alone do not say whether the pool is doing anything the native host +offload tier could not. For that, group the master's allocations by segment host +— a segment is one client process's donated memory, so the host tells you which +node the block physically lives on: + +```bash +grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" 2_mooncake_master.log \ + | awk '{sub(/size=/,"",$2); sub(/segment=/,"",$3); split($3,p,":"); + n[p[1]]++; b[p[1]]+=$2} + END {for (h in n) printf "%-16s pages=%-7d %.2f GiB\n", h, n[h], b[h]/1073741824}' +``` + +One host means a prefill-only pool (see §3). Two or more, with the generation +node among them, means blocks written by prefill are living on decode-side DRAM +and being read back from there. `disaggr_torch.slurm` writes this breakdown into +`9_mooncake_summary.log` at the end of every run, alongside the donor hosts, so +it needs running by hand only when diagnosing a partial run. + +Requires `GLOG_v=1` on the master, which `disaggr_torch.slurm` sets; raise it +with `MOONCAKE_MASTER_GLOG_V`. + ### Which reuse number means what This distinction matters and is easy to get backwards: diff --git a/mooncake_disagg/mooncake_segment_donor.py b/mooncake_disagg/mooncake_segment_donor.py new file mode 100644 index 000000000000..f7161c271cfc --- /dev/null +++ b/mooncake_disagg/mooncake_segment_donor.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Donate this node's host memory to a Mooncake store pool, without reading or +writing it. + +Pool capacity comes only from processes that open a store handle: ``setup`` +registers ``global_segment_size`` bytes of the calling process's host memory and +the master then places blocks in it. In a disaggregated deployment only the +context servers configure the KV connector, so only they call ``setup``, and the +pool is entirely prefill-node memory -- which makes the store a +prefill-DRAM-caches-prefill-GPU tier, overlapping what TensorRT-LLM's native +host offload already does. + +Running this alongside a generation server puts that node's memory into the same +pool. Prefill then writes blocks that land on decode-side DRAM, and reads them +back, while the generation engine itself stays free of any connector: it neither +reads nor writes the store, so it keeps its single cache transceiver for the +prefill-to-decode KV handoff. + +A donor is deliberately not a ``StoreRole``. The roles describe an engine's +traffic (``producer`` writes, ``consumer`` reads, ``both``), and none of them +means "contribute memory only" -- attaching a connector to the generation server +to get its DRAM into the pool would also start it reading or writing. Capacity +and traffic are separate concerns, so donation is a separate process. + +The donated memory is charged to this process, so it competes with the +generation server's own ``kv_cache_config.host_cache_size`` on the same node. +Size the two together. +""" + +import argparse +import json +import os +import re +import signal +import sys +import threading +import time + +_SIZE_UNITS = { + "": 1, + "b": 1, + "k": 1000, + "kb": 1000, + "m": 1000**2, + "mb": 1000**2, + "g": 1000**3, + "gb": 1000**3, + "t": 1000**4, + "tb": 1000**4, + "kib": 1024, + "mib": 1024**2, + "gib": 1024**3, + "tib": 1024**4, +} +_SIZE_RE = re.compile(r"^\s*([0-9]+(?:\.[0-9]+)?)\s*([a-zA-Z]*)\s*$") + +# The donor never transfers, so its staging buffer is dead weight; setup still +# rejects a zero one. +DEFAULT_LOCAL_BUFFER_SIZE = "64MiB" + + +def parse_size(value) -> int: + """Accept either a byte count or a suffixed string such as ``"32GiB"``. + + Mirrors the connector's parser so a size means the same thing in + ``mooncake.json`` and on this script's command line. + """ + if isinstance(value, bool): + raise ValueError(f"expected a size, got {value!r}") + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + match = _SIZE_RE.match(str(value)) + if match is None: + raise ValueError(f"cannot parse size {value!r}") + magnitude, unit = match.groups() + scale = _SIZE_UNITS.get(unit.lower()) + if scale is None: + raise ValueError(f"unknown size unit {unit!r} in {value!r}") + return int(float(magnitude) * scale) + + +def log(message: str) -> None: + print(f"[donor {time.strftime('%H:%M:%S')}] {message}", flush=True) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", + default=os.getenv("MOONCAKE_CONFIG_PATH"), + help="Mooncake JSON config naming the pool to join. Defaults to " + "$MOONCAKE_CONFIG_PATH.", + ) + parser.add_argument( + "--segment-size", + default="32GiB", + help="Host memory to contribute, e.g. 32GiB. Overrides the config's " + "global_segment_size, which is sized for an engine worker rather " + "than a node donating spare memory.", + ) + parser.add_argument( + "--ready-file", + default=None, + help="File to create once the segment is mounted, for launchers that " + "must not start writing to the pool before it has this capacity.", + ) + parser.add_argument( + "--heartbeat-seconds", + type=int, + default=300, + help="Interval between liveness lines. 0 disables them.", + ) + args = parser.parse_args() + + if not args.config: + parser.error("--config is required when MOONCAKE_CONFIG_PATH is unset") + + with open(args.config) as handle: + raw = json.load(handle) + + master = raw.get("master_server_address", "") + if not master: + parser.error(f"{args.config} has no master_server_address") + + segment_size = parse_size(args.segment_size) + local_buffer_size = parse_size( + raw.get("local_buffer_size_donor", DEFAULT_LOCAL_BUFFER_SIZE) + ) + protocol = raw.get("protocol", "rdma") + device_name = raw.get("device_name", "") or "" + metadata_server = raw.get("metadata_server", "") + + try: + from mooncake.store import MooncakeDistributedStore + except ImportError as exc: + log( + "the Mooncake Python bindings are missing " + "(`pip install mooncake-transfer-engine`); the C++ transfer engine " + f"in the container is a different component: {exc}" + ) + return 1 + + import socket + + hostname = socket.gethostbyname(socket.gethostname()) + + log( + f"joining pool at {master} as a capacity-only client: " + f"host={hostname} protocol={protocol} device={device_name or '(none)'} " + f"donating={segment_size / 1024 ** 3:.1f}GiB" + ) + + # Held for the process's lifetime: dropping the handle unmounts the segment + # and the master starts reporting the blocks living in it as lost. + store = MooncakeDistributedStore() + status = store.setup( + hostname, + metadata_server, + segment_size, + local_buffer_size, + protocol, + device_name, + master, + ) + if status != 0: + log( + f"setup failed with status {status}. The master must already be " + f"accepting connections at {master}, and protocol={protocol!r} " + "must be usable from this node." + ) + return 1 + + log(f"segment mounted; {segment_size / 1024 ** 3:.1f}GiB now available to the pool") + + if args.ready_file: + with open(args.ready_file, "w") as handle: + handle.write(f"{hostname} {segment_size}\n") + + stop = threading.Event() + + def handle_signal(signum, _frame): + log(f"received signal {signum}; unmounting segment") + stop.set() + + signal.signal(signal.SIGTERM, handle_signal) + signal.signal(signal.SIGINT, handle_signal) + + # Idle by design. Any put or get here would make this node a store client in + # the traffic sense, which is what keeping the generation engine + # connector-free is meant to avoid. + heartbeat = args.heartbeat_seconds + started = time.monotonic() + while not stop.is_set(): + if heartbeat > 0: + if stop.wait(heartbeat): + break + log(f"alive, donating for {(time.monotonic() - started) / 60:.0f}m") + else: + stop.wait() + + del store + log("exited") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From bd0acb40a51f3f3dd10ec1fc8425602df6d3def4 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:58:01 -0700 Subject: [PATCH 08/24] [None][fix] Stage mooncake-store KV pages through pinned host memory Registering the GPU KV pools directly with Mooncake requires GPUDirect RDMA, which is unavailable on GB300 nodes without nvidia_peermem: ibv_reg_mr fails with EFAULT and the TCP transport segfaults in its memcpy worker pool. Add an opt-in stage_through_host mode that gathers pages into pinned host slots and registers those instead, so only host memory is ever exposed to the transport. Bind the save thread to the rank's device, captured on the executor thread at layout registration. Torch's current device is thread-local and a new thread starts at 0, so the thread was creating its stream on device 0 while the KV pages lived on the rank's device -- every staged copy failed with cudaErrorInvalidValue on every rank except 0. State copy directions explicitly rather than inferring them from pointers, and report the operands and current device when a copy fails. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/__init__.py | 9 + .../connectors/mooncake_store/config.py | 33 ++ .../connectors/mooncake_store/scheduler.py | 12 +- .../connectors/mooncake_store/staging.py | 310 ++++++++++++++++++ .../connectors/mooncake_store/worker.py | 211 +++++++++--- .../executor/test_mooncake_store_connector.py | 244 ++++++++++++++ 6 files changed, 772 insertions(+), 47 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index 939f846bf8dd..1664e461421e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -33,6 +33,15 @@ kv_connector_config = KvCacheConnectorConfig(connector="mooncake-store") with ``MOONCAKE_CONFIG_PATH`` pointing at a Mooncake JSON config. + +By default the KV pools themselves are registered with Mooncake, so the store +reads and writes device memory and no copy is added. That needs the HCA to be +able to pin GPU pages -- GPUDirect RDMA, through ``nvidia_peermem`` or dma-buf. +Where it is missing, registration fails on every pool range and the connector +cannot start; ``"stage_through_host": true`` in the JSON config (or +``TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST=1``) then routes pages through a +pinned host buffer instead, so only host memory is ever registered. The stored +bytes are the same either way, so the two modes can share a pool. """ from .config import MooncakeStoreConnectorConfig, StoreRole diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 0e69d7270e72..948f47e7fa60 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -34,6 +34,7 @@ "CONFIG_PATH_ENV", "MooncakeStoreConnectorConfig", "ROLE_ENV", + "STAGE_THROUGH_HOST_ENV", "StoreRole", ] @@ -41,10 +42,15 @@ ROLE_ENV = "TRTLLM_MOONCAKE_STORE_ROLE" CACHE_PREFIX_ENV = "TRTLLM_MOONCAKE_STORE_PREFIX" MODEL_KEY_ENV = "TRTLLM_MOONCAKE_STORE_MODEL_KEY" +STAGE_THROUGH_HOST_ENV = "TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST" DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200 DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 DEFAULT_CACHE_PREFIX = "trtllm" +DEFAULT_STAGING_BUFFER_SIZE = 536870912 + +_TRUE = {"1", "true", "yes", "on"} +_FALSE = {"0", "false", "no", "off"} _SIZE_UNITS = { "": 1, @@ -128,6 +134,15 @@ class MooncakeStoreConnectorConfig: #: How many page keys go into one store call. Bounds the size of a single #: RPC without bounding how much a request may transfer. transfer_batch_size: int = 64 + #: Pass pages through a pinned host buffer instead of registering the KV + #: pools with Mooncake. Costs a copy in each direction and buys independence + #: from GPUDirect RDMA, without which registering device memory fails + #: outright. Leave off wherever the pool can reach GPU memory. + stage_through_host: bool = False + #: Ceiling on the pinned allocation per direction when staging. The pool is + #: sized from the layout's largest page, so this caps how many pages may be + #: in flight rather than how large one may be. + staging_buffer_bytes: int = DEFAULT_STAGING_BUFFER_SIZE def __post_init__(self) -> None: """Reject settings that would fail later, inside a transfer.""" @@ -139,6 +154,8 @@ def __post_init__(self) -> None: raise ValueError("global_segment_size must be >= 0") if self.transfer_batch_size <= 0: raise ValueError("transfer_batch_size must be > 0") + if self.stage_through_host and self.staging_buffer_bytes <= 0: + raise ValueError("staging_buffer_bytes must be > 0 when staging is on") @staticmethod def from_file(path: str) -> "MooncakeStoreConnectorConfig": @@ -160,6 +177,10 @@ def from_file(path: str) -> "MooncakeStoreConnectorConfig": cache_prefix=str(raw.get("cache_prefix", DEFAULT_CACHE_PREFIX)), model_key=raw.get("model_key") or None, transfer_batch_size=int(raw.get("transfer_batch_size", 64)), + stage_through_host=bool(raw.get("stage_through_host", False)), + staging_buffer_bytes=_parse_size( + raw.get("staging_buffer_bytes", DEFAULT_STAGING_BUFFER_SIZE) + ), ) @staticmethod @@ -193,6 +214,18 @@ def with_env_overrides(self) -> "MooncakeStoreConnectorConfig": model_key = os.getenv(MODEL_KEY_ENV) if model_key: updates["model_key"] = model_key + staging = os.getenv(STAGE_THROUGH_HOST_ENV) + if staging: + normalized = staging.strip().lower() + if normalized in _TRUE: + updates["stage_through_host"] = True + elif normalized in _FALSE: + updates["stage_through_host"] = False + else: + known = ", ".join(sorted(_TRUE | _FALSE)) + raise ValueError( + f"{STAGE_THROUGH_HOST_ENV}={staging!r} is not a boolean; use one of: {known}" + ) return dataclasses.replace(self, **updates) if updates else self def resolve_model_key(self, model: Any) -> str: diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py index 1096acc826d9..bb24c01be475 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -82,9 +82,8 @@ def __init__(self, llm_args: TorchLlmArgs): self._worker: Optional[MooncakeStoreConnectorWorker] = None logger.info( - "mooncake-store leader ready (role=%s, tokens_per_block=%d)", - self._config.role.value, - self._tokens_per_block, + f"mooncake-store leader ready (role={self._config.role.value}, " + f"tokens_per_block={self._tokens_per_block})" ) def wait_for_initialization(self): @@ -139,10 +138,9 @@ def get_num_new_matched_tokens( state.load_first_block = first_block state.load_blocks = hit_blocks logger.debug( - "mooncake-store matched %d blocks (%d tokens) for request %d", - hit_blocks, - hit_blocks * self._tokens_per_block, - request.request_id, + 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 diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py new file mode 100644 index 000000000000..484fb79320cd --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Pinned host slots that stand in for GPU pages when the pool cannot reach them. + +The connector's default path registers the KV pools themselves with Mooncake, so +the store reads and writes device memory directly. That needs the HCA to be able +to pin GPU pages -- GPUDirect RDMA, via ``nvidia_peermem`` or dma-buf. Where that +is unavailable, ``ibv_reg_mr`` fails on every pool range and the connector cannot +start at all. + +Staging trades a copy for that dependency. Mooncake is given a pinned host buffer +instead of the pools, and each page passes through a slot in it: gathered from its +device regions before a write, scattered back to them after a read. The store then +only ever registers host memory, which needs no GPUDirect. + +A slot holds the page's regions concatenated in region order, which is precisely +the payload the zero-copy path would have produced from the same regions. The +stored bytes are therefore identical either way, so a pool written by one path is +readable by the other -- including by another engine sharing the pool. + +Copies go through ``cudaMemcpyAsync`` rather than the batched Triton kernel in +``disaggregation/native/bounce/gather_scatter.py``. That kernel is the better tool +for device-to-device gather, but here one side is host memory: the copy engines +move it over the host link by DMA, whereas a kernel would do it with scattered +stores from the SMs. +""" + +from typing import List, Optional, Sequence, Tuple + +import torch + +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart + +from tensorrt_llm._utils import CUASSERT +from tensorrt_llm.logger import logger + +__all__ = ["HostStagingPool", "plan_slot_geometry", "sync_stream"] + +#: Stated rather than inferred from the pointers: the direction is known at each +#: call site, and saying so keeps a copy from being misread if a host pointer is +#: ever outside the unified address space. +_DEVICE_TO_HOST = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost +_HOST_TO_DEVICE = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice + + +def _memcpy_async(dst: int, src: int, size: int, kind, stream: int) -> None: + """One asynchronous copy between a device page and a host slot.""" + status = cudart.cudaMemcpyAsync(int(dst), int(src), int(size), kind, stream)[0] + if status == cudart.cudaError_t.cudaSuccess: + return + # Raised here rather than through CUASSERT so the operands are in the + # message. A bare cudaErrorInvalidValue from a copy says nothing about which + # of the three plausible causes it was. + device = torch.cuda.current_device() if torch.cuda.is_available() else None + raise RuntimeError( + f"cudaMemcpyAsync failed with {status} staging a KV page: " + f"dst={int(dst):#x} src={int(src):#x} size={size} " + f"stream={int(stream):#x} current_device={device}. An invalid value here " + "is usually a stream created on a different device than the pages, which " + "happens when a thread issues the copy without inheriting the rank's " + "device -- torch's current device is thread-local." + ) + + +def sync_stream(stream: int) -> None: + """Wait for a stream's copies to finish, given its raw handle.""" + CUASSERT(cudart.cudaStreamSynchronize(stream)) + + +def plan_slot_geometry( + max_bytes_per_page: int, + transfer_batch_size: int, + budget_bytes: int, +) -> Tuple[int, int]: + """Choose how many pages may be staged at once, and how wide a slot is. + + A slot has to hold the largest page any layer group produces, so the page size + is a floor on the allocation: a budget below one page is raised to one rather + than refused, since the alternative is not starting. + + Args: + max_bytes_per_page: Largest page payload across layer groups. + transfer_batch_size: Pages the connector puts in one store call. There is + no point staging more than that. + budget_bytes: Ceiling on this pool's pinned allocation. + + Returns: + Slot width in bytes, and the number of slots. + """ + if max_bytes_per_page <= 0: + raise ValueError(f"max_bytes_per_page must be > 0, got {max_bytes_per_page}") + if transfer_batch_size <= 0: + raise ValueError(f"transfer_batch_size must be > 0, got {transfer_batch_size}") + + affordable = budget_bytes // max_bytes_per_page + num_slots = max(1, min(transfer_batch_size, affordable)) + return max_bytes_per_page, num_slots + + +class HostStagingPool: + """A registered pinned buffer, sliced into per-page slots. + + One pool serves one direction. Loads run on the executor thread and saves on + the connector's background thread, so sharing slots between them would need a + lock on the transfer path for no benefit -- the two pools are independent. + """ + + def __init__( + self, + *, + slot_bytes: int, + num_slots: int, + store, + label: str, + ): + self._slot_bytes = int(slot_bytes) + self._num_slots = int(num_slots) + self._label = label + + # Pinned unconditionally, unlike the ``prefer_pinned`` heuristic used for + # transfer buffers elsewhere: this memory is handed to the store to + # register, so page-locking it is a correctness property of the + # registration rather than a copy-speed preference. + pin = torch.cuda.is_available() + self._buffer = torch.empty( + self._slot_bytes * self._num_slots, dtype=torch.uint8, pin_memory=pin + ) + self._base = int(self._buffer.data_ptr()) + + status = store.register_buffer(self._base, self._buffer.numel()) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.register_buffer failed with status " + f"{status} for the {label} host staging buffer at " + f"[{self._base:#x}, {self._base + self._buffer.numel():#x}). Host " + f"memory registration failing points at the pool or the fabric " + f"rather than at GPUDirect, which is what staging avoids." + ) + logger.info( + f"mooncake-store {label} staging: {self._num_slots} slots x " + f"{self._slot_bytes} B = {self._buffer.numel() / 1024**2:.1f} MiB pinned " + f"(pinned={pin})" + ) + + @property + def num_slots(self) -> int: + """Pages this pool can hold at once.""" + return self._num_slots + + @property + def slot_bytes(self) -> int: + """Capacity of one slot.""" + return self._slot_bytes + + def slot_address(self, index: int) -> int: + """Address of slot ``index``.""" + if not 0 <= index < self._num_slots: + raise IndexError(f"slot {index} out of range [0, {self._num_slots})") + return self._base + index * self._slot_bytes + + def _check_fits(self, total: int) -> None: + if total > self._slot_bytes: + raise ValueError( + f"a {total} B page does not fit the {self._slot_bytes} B " + f"{self._label} staging slot; the pool was sized from the layout's " + "largest page, so this means the layout changed after registration" + ) + + def gather( + self, + index: int, + addresses: Sequence[int], + sizes: Sequence[int], + stream: int, + ) -> Tuple[int, int]: + """Copy one page's device regions into slot ``index``, concatenated. + + Args: + index: Slot to fill. + addresses: Device addresses of the page's regions, in region order. + sizes: Byte counts matching ``addresses``. + stream: CUDA stream handle the copies are issued on. + + Returns: + The slot's address and the total bytes written, ready to hand to the + store as a single buffer. + """ + total = sum(sizes) + self._check_fits(total) + destination = self.slot_address(index) + offset = 0 + for address, size in zip(addresses, sizes, strict=True): + _memcpy_async(destination + offset, address, size, _DEVICE_TO_HOST, stream) + offset += size + return destination, total + + def scatter( + self, + index: int, + addresses: Sequence[int], + sizes: Sequence[int], + stream: int, + ) -> None: + """Copy slot ``index`` back out to one page's device regions. + + The inverse of :meth:`gather`, walking the regions in the same order so + the split matches the concatenation the slot holds. + """ + self._check_fits(sum(sizes)) + source = self.slot_address(index) + offset = 0 + for address, size in zip(addresses, sizes, strict=True): + _memcpy_async(address, source + offset, size, _HOST_TO_DEVICE, stream) + offset += size + + def reserve(self, total: int) -> None: + """Assert a page of ``total`` bytes is stageable, without copying.""" + self._check_fits(total) + + +def stage_batch_for_put( + pool: HostStagingPool, + addresses: Sequence[Sequence[int]], + sizes: Sequence[Sequence[int]], + stream: int, +) -> Tuple[List[List[int]], List[List[int]]]: + """Gather a batch of device pages into slots and describe them for the store. + + Args: + pool: Slots to stage through. The batch must not exceed its slot count. + addresses: Per-page device region addresses. + sizes: Per-page device region sizes. + stream: Stream the copies are issued on. The caller must synchronize it + before the store reads the slots. + + Returns: + Per-page address and size lists, each a single staged buffer. + """ + if len(addresses) > pool.num_slots: + raise ValueError( + f"batch of {len(addresses)} pages exceeds {pool.num_slots} staging slots" + ) + staged_addresses: List[List[int]] = [] + staged_sizes: List[List[int]] = [] + for index, (page_addresses, page_sizes) in enumerate(zip(addresses, sizes, strict=True)): + slot, total = pool.gather(index, page_addresses, page_sizes, stream) + staged_addresses.append([slot]) + staged_sizes.append([total]) + return staged_addresses, staged_sizes + + +def describe_batch_for_get( + pool: HostStagingPool, + sizes: Sequence[Sequence[int]], +) -> Tuple[List[List[int]], List[List[int]]]: + """Describe slots for the store to read a batch into, before scattering. + + Unlike the put direction there is nothing to copy first: the slots are the + destination, and :func:`unstage_batch_after_get` moves the bytes on once the + store has filled them. + """ + if len(sizes) > pool.num_slots: + raise ValueError(f"batch of {len(sizes)} pages exceeds {pool.num_slots} staging slots") + staged_addresses: List[List[int]] = [] + staged_sizes: List[List[int]] = [] + for index, page_sizes in enumerate(sizes): + total = sum(page_sizes) + pool.reserve(total) + staged_addresses.append([pool.slot_address(index)]) + staged_sizes.append([total]) + return staged_addresses, staged_sizes + + +def unstage_batch_after_get( + pool: HostStagingPool, + addresses: Sequence[Sequence[int]], + sizes: Sequence[Sequence[int]], + stream: int, + only: Optional[Sequence[int]] = None, +) -> None: + """Scatter filled slots back to their device pages. + + Args: + pool: Slots the store just wrote into. + addresses: Per-page device region addresses. + sizes: Per-page device region sizes. + stream: Stream the copies are issued on. The caller must synchronize it + before the pages are read. + only: Slot indices to scatter. Defaults to all of them; a caller that + knows some reads failed passes the rest so a failed page is not + written over its device slot with whatever the slot held. + """ + indices = range(len(addresses)) if only is None else only + for index in indices: + pool.scatter(index, addresses[index], sizes[index], stream) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 9dd65ed7d3fc..344c041e565f 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -32,6 +32,7 @@ """ import threading +import traceback from collections import defaultdict from queue import Queue from typing import Dict, List, Optional, Sequence, Set, Tuple @@ -48,6 +49,14 @@ from .config import CONFIG_PATH_ENV, MooncakeStoreConnectorConfig from .keys import KeyNamespace from .metadata import MooncakeStoreMetadata, RequestTransfers +from .staging import ( + HostStagingPool, + describe_batch_for_get, + plan_slot_geometry, + stage_batch_for_put, + sync_stream as _sync_stream, + unstage_batch_after_get, +) from .validation import validate_layout, validate_llm_args __all__ = ["MooncakeStoreConnectorWorker", "resolve_local_worker"] @@ -129,6 +138,17 @@ def _batched(items: Sequence, size: int): yield items[start : start + size] +def _stream_handle(stream) -> int: + """The raw CUDA stream handle behind a torch stream, or a handle as given. + + ``None`` maps to 0, the default stream, which is what the runtime passes when + it has no stream of its own to offer. + """ + if stream is None: + return 0 + return int(getattr(stream, "cuda_stream", stream)) + + class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): """Moves KV pages between this rank's GPU cache and the Mooncake pool.""" @@ -155,6 +175,16 @@ def __init__(self, llm_args: TorchLlmArgs): ) self._save_thread: Optional[threading.Thread] = None self._save_lock = threading.Lock() + # Host staging, when the pool cannot register device memory. One pool per + # direction so the executor thread and the save thread never share slots. + self._load_staging: Optional[HostStagingPool] = None + self._save_staging: Optional[HostStagingPool] = None + self._save_stream: Optional[torch.cuda.Stream] = None + # This rank's device, captured on the executor thread. The save thread + # cannot ask for it itself; see _drain_saves. + self._device_index: Optional[int] = None + # Pages per store call. Staging narrows this to the slots it can afford. + self._batch_size = self._config.transfer_batch_size # Save submissions still in flight, per request. self._outstanding_saves: Dict[int, int] = defaultdict(int) # Requests the runtime has told us are done producing KV. Their pages @@ -167,12 +197,9 @@ def __init__(self, llm_args: TorchLlmArgs): _LOCAL_WORKER_READY.set() logger.info( - "mooncake-store worker rank %d/%d ready (role=%s, model_key=%s, master=%s)", - self._rank, - self._world_size, - self._config.role.value, - self._model_key, - self._config.master_server_address, + f"mooncake-store worker rank {self._rank}/{self._world_size} ready " + f"(role={self._config.role.value}, model_key={self._model_key}, " + f"master={self._config.master_server_address})" ) # ---- registration ---- @@ -200,14 +227,25 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: validate_layout(layout) addressing = PageAddressing(layout) - for start, end in addressing.registration_ranges(): - status = self._store.register_buffer(start, end - start) - if status != 0: - raise RuntimeError( - f"MooncakeDistributedStore.register_buffer failed with status " - f"{status} for [{start:#x}, {end:#x}). Without registration " - "the store cannot read or write these pages." - ) + # Read here, on the executor thread, because it is thread-local and the + # save thread would otherwise see device 0 rather than this rank's. + if torch.cuda.is_available(): + self._device_index = torch.cuda.current_device() + if self._config.stage_through_host: + self._open_staging(addressing) + else: + for start, end in addressing.registration_ranges(): + status = self._store.register_buffer(start, end - start) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.register_buffer failed with status " + f"{status} for [{start:#x}, {end:#x}). Without registration " + "the store cannot read or write these pages. Registering " + "device memory needs GPUDirect RDMA (nvidia_peermem or " + "dma-buf); where that is unavailable, set " + "stage_through_host to pass pages through pinned host " + "memory instead." + ) self._addressing = addressing for layer_group_id in addressing.layer_group_ids: @@ -229,11 +267,51 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: self._save_thread.start() logger.info( - "mooncake-store worker rank %d registered layout: %s", - self._rank, - addressing.describe(), + f"mooncake-store worker rank {self._rank} registered layout: " + f"{addressing.describe()}" ) + def _open_staging(self, addressing: PageAddressing) -> None: + """Allocate and register the pinned slots pages will pass through. + + Only the directions this role drives get a pool, since each one costs its + own pinned allocation and a consumer never gathers, nor a producer + scatter. The GPU pools are deliberately left unregistered: reaching them + is what staging exists to avoid needing. + """ + max_bytes_per_page = max( + addressing.bytes_per_page(layer_group_id) + for layer_group_id in addressing.layer_group_ids + ) + slot_bytes, num_slots = plan_slot_geometry( + max_bytes_per_page, + self._config.transfer_batch_size, + self._config.staging_buffer_bytes, + ) + if self._config.role.loads: + self._load_staging = HostStagingPool( + slot_bytes=slot_bytes, + num_slots=num_slots, + store=self._store, + label="load", + ) + if self._config.role.saves: + self._save_staging = HostStagingPool( + slot_bytes=slot_bytes, + num_slots=num_slots, + store=self._store, + label="save", + ) + self._batch_size = min(self._config.transfer_batch_size, num_slots) + if self._batch_size < self._config.transfer_batch_size: + logger.warning( + f"mooncake-store rank {self._rank} reduced its transfer batch from " + f"{self._config.transfer_batch_size} to {self._batch_size} pages: " + f"staging {max_bytes_per_page} B pages within " + f"{self._config.staging_buffer_bytes} B does not fit more. Raise " + f"staging_buffer_bytes to restore the configured batch size." + ) + def _namespace(self, rank: int, layer_group_id: int, bytes_per_page: int) -> KeyNamespace: return KeyNamespace( cache_prefix=self._config.cache_prefix, @@ -282,15 +360,17 @@ def count_prefix_hit(self, block_hashes: Sequence[bytes]) -> int: try: present = self._store.batch_is_exist(keys) - except Exception: - logger.warning("mooncake-store lookup failed; treating as a miss", exc_info=True) + 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( - "mooncake-store batch_is_exist returned %d results for %d keys; treating as a miss", - len(present), - len(keys), + f"mooncake-store batch_is_exist returned {len(present)} results for " + f"{len(keys)} keys; treating as a miss" ) return 0 @@ -317,14 +397,21 @@ def start_load_kv(self, stream: torch.cuda.Stream): if not keys: return + staging = self._load_staging + handle = _stream_handle(stream) if staging is not None else 0 + for batch in zip( - _batched(keys, self._config.transfer_batch_size), - _batched(addresses, self._config.transfer_batch_size), - _batched(sizes, self._config.transfer_batch_size), + _batched(keys, self._batch_size), + _batched(addresses, self._batch_size), + _batched(sizes, self._batch_size), ): batch_keys, batch_addresses, batch_sizes = batch + if staging is None: + target_addresses, target_sizes = list(batch_addresses), list(batch_sizes) + else: + target_addresses, target_sizes = describe_batch_for_get(staging, batch_sizes) results = self._store.batch_get_into_multi_buffers( - list(batch_keys), list(batch_addresses), list(batch_sizes) + list(batch_keys), target_addresses, target_sizes ) failed = [ key @@ -340,8 +427,16 @@ def start_load_kv(self, stream: torch.cuda.Stream): f"{len(batch_keys)} pages; the affected KV slots were already " f"reported as computed. First failure: {failed[:1]}" ) + if staging is not None: + # Only reached when every page in the batch landed, so no slot + # holding a failed read is copied over a device page. + unstage_batch_after_get(staging, batch_addresses, batch_sizes, handle) + # The slots are reused by the next batch and the forward pass + # reads these pages, so the scatter has to be complete before + # either happens. + _sync_stream(handle) - logger.debug("mooncake-store rank %d loaded %d pages", self._rank, total_pages) + logger.debug(f"mooncake-store rank {self._rank} loaded {total_pages} pages") def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream): """No-op: loads complete in ``start_load_kv``. @@ -408,7 +503,20 @@ def get_finished( return finished_saving, list(started_loading_req_ids) def _drain_saves(self) -> None: - torch.cuda.set_device(torch.cuda.current_device()) + # Torch's current device is thread-local and a new thread starts at 0, so + # this has to be the device captured on the executor thread rather than + # whatever this one defaults to. Getting it wrong only shows up once the + # thread issues CUDA work of its own: the stream would belong to device 0 + # while the KV pointers belong to the rank's device, and the copy fails + # with cudaErrorInvalidValue on every rank except 0. + if self._device_index is not None: + torch.cuda.set_device(self._device_index) + if self._save_staging is not None and torch.cuda.is_available(): + # Owned by this thread so the gather never queues behind the + # executor's work, and created after set_device so it lands on the + # rank's device. Without a device there is nothing to order and the + # default stream handle stands in, which is what unit tests exercise. + self._save_stream = torch.cuda.Stream() while True: item = self._save_queue.get() if item is None: @@ -421,7 +529,10 @@ def _drain_saves(self) -> None: # Broad on purpose: this is the thread boundary. Anything that # escapes here would be lost, so it is stashed and re-raised on # the executor thread at the next connector call. - logger.error("mooncake-store save failed on rank %d: %s", self._rank, exc) + logger.error( + f"mooncake-store save failed on rank {self._rank}: " + f"{type(exc).__name__}: {exc}" + ) with self._save_lock: if self._save_error is None: self._save_error = exc @@ -439,10 +550,13 @@ def _put(self, transfers: Sequence[RequestTransfers]) -> None: if not keys: return + staging = self._save_staging + handle = _stream_handle(self._save_stream) if staging is not None else 0 + for batch in zip( - _batched(keys, self._config.transfer_batch_size), - _batched(addresses, self._config.transfer_batch_size), - _batched(sizes, self._config.transfer_batch_size), + _batched(keys, self._batch_size), + _batched(addresses, self._batch_size), + _batched(sizes, self._batch_size), ): batch_keys, batch_addresses, batch_sizes = batch # Skip pages another rank or another instance already wrote. The @@ -456,20 +570,29 @@ def _put(self, transfers: Sequence[RequestTransfers]) -> None: ] if not pending: continue + source_addresses = [batch_addresses[i] for i in pending] + source_sizes = [batch_sizes[i] for i in pending] + if staging is not None: + # Gathering after the existence filter means a page that is + # already in the pool costs no copy. + source_addresses, source_sizes = stage_batch_for_put( + staging, source_addresses, source_sizes, handle + ) + # The store reads the slots on this thread, so they have to be + # filled first. + _sync_stream(handle) results = self._store.batch_put_from_multi_buffers( [batch_keys[i] for i in pending], - [batch_addresses[i] for i in pending], - [batch_sizes[i] for i in pending], + source_addresses, + source_sizes, ) failures = sum(1 for result in results if not isinstance(result, int) or result < 0) if failures: # A dropped write only costs a future cache miss, so it is worth # a warning rather than failing a request that already answered. logger.warning( - "mooncake-store rank %d failed to save %d of %d pages", - self._rank, - failures, - len(pending), + f"mooncake-store rank {self._rank} failed to save {failures} of " + f"{len(pending)} pages" ) # ---- shared ---- @@ -518,8 +641,16 @@ def shutdown(self) -> None: if store is not None: try: store.close() - except Exception: - logger.warning("mooncake-store close failed", exc_info=True) + except Exception as exc: + logger.warning( + f"mooncake-store close failed: {type(exc).__name__}: {exc}\n" + f"{traceback.format_exc()}" + ) + # Released only after the store is closed: it holds registrations against + # this memory, and the save thread was already joined above. + self._load_staging = None + self._save_staging = None + self._save_stream = None global _LOCAL_WORKER if _LOCAL_WORKER is self: _LOCAL_WORKER = None diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 292259c209ab..6b8725907e74 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -21,9 +21,11 @@ import contextlib import json +import time from types import SimpleNamespace import pytest +import torch from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( RequestData, @@ -52,9 +54,11 @@ PageTransfer, RequestTransfers, ) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import staging as staging_module from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.scheduler import ( MooncakeStoreConnectorScheduler, ) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.staging import plan_slot_geometry from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.validation import ( validate_layout, validate_llm_args, @@ -538,6 +542,246 @@ def test_worker_shutdown_closes_the_store(store_config, fake_store): worker.shutdown() +# ---- host staging ---- + + +@contextlib.contextmanager +def make_staged_worker(fake_store, store_config, *, layout, budget=None): + """A worker configured to pass pages through pinned host slots.""" + raw = json.loads(store_config.read_text()) + raw["stage_through_host"] = True + if budget is not None: + raw["staging_buffer_bytes"] = budget + store_config.write_text(json.dumps(raw)) + with make_worker(fake_store, layout=layout) as worker: + yield worker + + +@pytest.fixture +def staged_copies(monkeypatch): + """Record the copies staging would issue, instead of running them.""" + copies = [] + monkeypatch.setattr( + staging_module, + "_memcpy_async", + lambda dst, src, size, kind, stream: copies.append((int(dst), int(src), int(size))), + ) + # Imported into the worker by name, so the worker's binding is the one that + # has to be replaced. + monkeypatch.setattr(worker_module, "_sync_stream", lambda _stream: None) + return copies + + +@pytest.fixture +def fake_cuda(monkeypatch): + """Present a CUDA device on a host that has none, recording set_device calls. + + Only safe for paths that do not allocate or launch; it exists to exercise the + device bookkeeping around the save thread. + """ + recorded = [] + 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)) + return recorded + + +@pytest.mark.parametrize( + "page_bytes,batch,budget,expected_slots", + [ + (1024, 64, 1 << 20, 64), # budget is ample: the full batch stages + (1024, 64, 8 * 1024, 8), # budget binds before the batch does + (1024, 8, 1 << 20, 8), # batch binds before the budget does + (1024, 64, 1024, 1), # exactly one page fits + (1024, 64, 1, 1), # below one page, raised to one rather than refused + ], +) +def test_plan_slot_geometry(page_bytes, batch, budget, expected_slots): + slot_bytes, num_slots = plan_slot_geometry(page_bytes, batch, budget) + # A slot always holds a whole page; the budget bounds the count, not the width. + assert slot_bytes == page_bytes + assert num_slots == expected_slots + + +@pytest.mark.parametrize("bad", [(0, 8, 1024), (-1, 8, 1024), (1024, 0, 1024)]) +def test_plan_slot_geometry_rejects_degenerate_inputs(bad): + with pytest.raises(ValueError): + plan_slot_geometry(*bad) + + +def test_config_reads_staging_from_the_json(store_config): + raw = json.loads(store_config.read_text()) + raw["stage_through_host"] = True + raw["staging_buffer_bytes"] = "256MiB" + store_config.write_text(json.dumps(raw)) + + config = MooncakeStoreConnectorConfig.from_env() + assert config.stage_through_host is True + assert config.staging_buffer_bytes == 256 * 1024**2 + + +def test_config_defaults_to_zero_copy(store_config): + assert MooncakeStoreConnectorConfig.from_env().stage_through_host is False + + +@pytest.mark.parametrize( + "value,expected", [("1", True), ("true", True), ("on", True), ("0", False), ("off", False)] +) +def test_config_staging_env_override(store_config, monkeypatch, value, expected): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST", value) + assert MooncakeStoreConnectorConfig.from_env().stage_through_host is expected + + +def test_config_rejects_a_non_boolean_staging_env(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST", "sometimes") + with pytest.raises(ValueError, match="not a boolean"): + MooncakeStoreConnectorConfig.from_env() + + +def test_staging_registers_host_buffers_and_never_the_pools( + store_config, fake_store, staged_copies +): + layout = make_layout(regions_per_group=2) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + # Registering the pools is the step that needs GPUDirect, so staging must + # not do it at all -- that is the whole point of the mode. + pool_ranges = PageAddressing(layout).registration_ranges() + registered_starts = {address for address, _size in fake_store.registered} + assert registered_starts.isdisjoint({start for start, _end in pool_ranges}) + + # One pinned buffer per direction, since the default role is ``both``. + assert len(fake_store.registered) == 2 + assert registered_starts == { + worker._load_staging.slot_address(0), + worker._save_staging.slot_address(0), + } + + +def test_staging_put_hands_the_store_one_host_buffer_per_page( + store_config, fake_store, staged_copies +): + layout = make_layout(regions_per_group=3) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x07" * 16 + worker._put([RequestTransfers(1, [PageTransfer(block_hash, 0, 2)])]) + + device_addresses, device_sizes = PageAddressing(layout).buffers(0, 2) + slot = worker._save_staging.slot_address(0) + (keys, addresses, sizes) = fake_store.put_calls[0] + + assert keys == [worker._namespaces[0].key(block_hash)] + # The store sees one contiguous host buffer, and its length is the sum of + # the device regions. That equality is what keeps a staged write + # byte-identical to a zero-copy one, so either path can read the other's + # pages. + assert addresses == [[slot]] + assert sizes == [[sum(device_sizes)]] + + # Every region was gathered, in order, into its place in the slot. + expected = [] + offset = 0 + for address, size in zip(device_addresses, device_sizes): + expected.append((slot + offset, address, size)) + offset += size + assert staged_copies == expected + + +def test_staging_get_scatters_back_to_the_device_regions( + store_config, fake_store, staged_copies +): + layout = make_layout(regions_per_group=3) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x03" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + + worker.bind_connector_meta( + SimpleNamespace(loads=[RequestTransfers(7, [PageTransfer(block_hash, 0, 5)])], saves=[]) + ) + worker.start_load_kv(None) + + device_addresses, device_sizes = PageAddressing(layout).buffers(0, 5) + slot = worker._load_staging.slot_address(0) + (_keys, addresses, sizes) = fake_store.get_calls[0] + assert addresses == [[slot]] + assert sizes == [[sum(device_sizes)]] + + # The scatter is the mirror of the gather: same split, opposite direction. + expected = [] + offset = 0 + for address, size in zip(device_addresses, device_sizes): + expected.append((address, slot + offset, size)) + offset += size + assert staged_copies == expected + + +def test_staging_does_not_scatter_a_failed_load(store_config, fake_store, staged_copies): + layout = make_layout() + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x09" * 16 + key = worker._namespaces[0].key(block_hash) + fake_store.objects.add(key) + fake_store.fail_gets_for.add(key) + + worker.bind_connector_meta( + SimpleNamespace(loads=[RequestTransfers(7, [PageTransfer(block_hash, 0, 1)])], saves=[]) + ) + with pytest.raises(RuntimeError, match="failed to load"): + worker.start_load_kv(None) + + # A failed read leaves the slot holding whatever it held before. Copying + # that onto the page would put unrelated bytes where the runtime already + # promised computed KV. + assert staged_copies == [] + + +def test_worker_captures_the_ranks_device_at_registration(store_config, fake_store, fake_cuda): + with make_worker(fake_store, layout=make_layout()) as worker: + # Read on the executor thread, where it is correct. torch's current + # device is thread-local, so the save thread cannot ask for it itself. + assert worker._device_index == 3 + + +def test_save_thread_adopts_the_ranks_device_not_the_thread_default( + store_config, fake_store, fake_cuda +): + """Regression: the save thread must not run on torch's default device. + + It issues staging copies against pointers owned by the rank's device. A + stream created on device 0 instead fails every copy with + cudaErrorInvalidValue, and only on ranks other than 0 -- which is exactly how + this escaped into a run. + """ + with make_worker(fake_store, layout=make_layout()): + deadline = time.monotonic() + 5.0 + while 3 not in fake_cuda and time.monotonic() < deadline: + time.sleep(0.01) + assert 3 in fake_cuda, f"save thread set devices {fake_cuda}, expected the rank's 3" + # Never the thread-local default, which is what the bug did. + assert 0 not in fake_cuda + + +def test_staging_narrows_the_batch_to_the_budget(store_config, fake_store, staged_copies): + layout = make_layout(regions_per_group=2, num_slots=8) + page_bytes = PageAddressing(layout).bytes_per_page(0) + with make_staged_worker( + fake_store, store_config, layout=layout, budget=2 * page_bytes + ) as worker: + assert worker._save_staging.num_slots == 2 + assert worker._batch_size == 2 + + hashes = [bytes([index]) * 16 for index in range(5)] + worker._put( + [ + RequestTransfers( + 1, [PageTransfer(block_hash, 0, index) for index, block_hash in enumerate(hashes)] + ) + ] + ) + # Five pages through two slots: three calls, and no call wider than the + # slot count, which is the constraint staging adds. + assert [len(keys) for keys, _a, _s in fake_store.put_calls] == [2, 2, 1] + + # ---- scheduler ---- From 41d312e3865d5ffa0197766e38c8ef5865b73d80 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:36:14 -0700 Subject: [PATCH 09/24] [None][fix] Force partial reuse off when the mooncake-store connector is used The store is addressed by whole blocks. The connector receives the device match as num_computed_tokens and offers only blocks beyond it, but it can resume only from a block boundary, so a match ending mid-block makes it decline the lookup and the store is never consulted. enable_partial_reuse is exactly what puts the match off a boundary, so it trades part of one block of device reuse for every stored block of the remaining prefix. The default is true, which made the pathological combination the one a user gets by saying nothing. On MiniMax-M3 it declined 97.2% of lookups and held actual prompt cache read at 35% against a 96% ceiling, so a 1.6 TB pool measured as though it were absent; forcing it off reached 93.5% and 2.18x the output token throughput. Coerce rather than reject, since a wrong answer is not at stake and refusing to start over a default no one chose is worse than fixing it and saying so. This sits in py_executor_creator beside the FORCE_DETERMINISTIC coercion because the KV cache manager reads the flag when it builds its block pools, which happens well before the connector is constructed -- so the connector cannot police this from its own startup gates. uses_connector compares the resolved module rather than the connector name, so a config that spells out connector_module instead of using the preset is still recognized. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docs/source/features/kv-cache-connector.md | 6 ++++ .../_torch/pyexecutor/connectors/registry.py | 22 ++++++++++++ .../_torch/pyexecutor/py_executor_creator.py | 14 ++++++++ .../executor/test_mooncake_store_connector.py | 34 +++++++++++++++++++ 4 files changed, 76 insertions(+) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index c4f82b71e116..1338d0e32359 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -136,6 +136,12 @@ Two further settings are TensorRT-LLM's rather than Mooncake's, and are read fro In a disaggregated deployment, run context servers as `both` and leave generation servers unconfigured. Generated tokens are rarely a reused prefix, so writing them costs bandwidth for no hit rate. +#### Partial block reuse is forced off + +`kv_cache_config.enable_partial_reuse` is set to `false` when this connector is configured, with a warning, whether or not it was requested explicitly. It defaults to `true`, so most deployments will see that warning. + +The store is addressed by whole blocks. The connector is handed the device match as `num_computed_tokens` and offers only blocks beyond it, but it can only resume from a block boundary -- so when the device match ends mid-block, it declines the lookup and the store is not consulted at all. Partial reuse is precisely what puts the match off a boundary, which means it trades part of one block of device reuse for every stored block of the remaining prefix. Measured on MiniMax-M3, leaving it enabled declined 97.2% of lookups and left actual prompt cache read at 35% against a 96% ceiling; forcing it off raised that to 94% and roughly doubled throughput. + #### How it keys pages `KVCacheManagerV2` reports `RequestData.block_hashes` empty, so the connector derives block identity itself: a blake2b chain where each block's hash covers its own tokens *and* every token before it, seeded by the request's `cache_salt`. A key is `//wr/lg/tb/`. The namespace pins down everything that would make the stored bytes mean something different, so a mismatched shard count, layer group or page geometry reads as a cache miss rather than as garbage. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py index 70dca022b763..ed7a4098c43b 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py @@ -19,6 +19,11 @@ it is resolved at runtime via importlib in py_executor_creator.py. """ +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig + CONNECTOR_REGISTRY: dict[str, dict[str, str]] = { "lmcache": { "connector_module": "lmcache.integration.tensorrt_llm.tensorrt_adapter", @@ -41,3 +46,20 @@ "connector_worker_class": "MooncakeStoreConnectorWorker", }, } + + +def uses_connector(kv_connector_config: Optional["KvCacheConnectorConfig"], name: str) -> bool: + """Report whether a connector config resolves to the named preset. + + Compares the resolved module rather than the ``connector`` field, so a + config that names the module explicitly instead of using the preset is + still recognized. Accepts ``None`` to save every caller a null check. + """ + if kv_connector_config is None: + return False + preset = CONNECTOR_REGISTRY.get(name) + if preset is None: + raise ValueError( + f"Unknown connector preset: {name!r}. Known presets: {list(CONNECTOR_REGISTRY)}" + ) + return kv_connector_config.connector_module == preset["connector_module"] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 67a48544c9b6..2f9cd43c7d65 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -42,6 +42,7 @@ validate_feature_combination) from .config_utils import is_hybrid_linear, is_minimax_m3 from .connectors.kv_cache_connector import KvCacheConnectorManager +from .connectors.registry import uses_connector from .dwdp import DwdpManager from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder from .model_engine import PyTorchModelEngine @@ -382,6 +383,19 @@ def create_py_executor( kv_cache_config.enable_block_reuse = False kv_cache_config.enable_partial_reuse = False + # Must happen before the KV cache manager is built below, since the manager + # reads enable_partial_reuse to construct its block pools. + if (kv_cache_config.enable_partial_reuse + and uses_connector(kv_connector_config, "mooncake-store")): + logger.warning( + "Disabling partial reuse: it is not usable with the mooncake-store " + "connector. The store is addressed by whole blocks, so a partial " + "device match leaves the matched length off a block boundary and " + "the connector declines the lookup rather than resume a block from " + "the middle. Partial reuse therefore trades part of one block for " + "every stored block of the remaining prefix.") + kv_cache_config.enable_partial_reuse = False + decoding_config = llm_args.decoding_config # The tokenizer is stripped from MPI kwargs in proxy.py to avoid pickle diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 6b8725907e74..e28692c22f46 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -66,6 +66,8 @@ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.worker import ( MooncakeStoreConnectorWorker, ) +from tensorrt_llm._torch.pyexecutor.connectors.registry import uses_connector +from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX TOKENS_PER_BLOCK = 4 @@ -430,6 +432,38 @@ def test_validate_layout_rejects_sliding_window(): validate_layout(make_layout()) +# ---- connector identification ---- +# +# py_executor_creator turns partial reuse off for this connector, and finds it +# through uses_connector. Missing the config would silently cost the reuse the +# store exists to provide, so the recognition itself is worth pinning down. + + +def test_uses_connector_recognizes_the_preset(): + config = KvCacheConnectorConfig(connector="mooncake-store") + assert uses_connector(config, "mooncake-store") + + +def test_uses_connector_recognizes_a_hand_written_module(): + config = KvCacheConnectorConfig( + connector_module="tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + connector_scheduler_class="MooncakeStoreConnectorScheduler", + connector_worker_class="MooncakeStoreConnectorWorker", + ) + assert uses_connector(config, "mooncake-store") + + +def test_uses_connector_separates_connectors_and_tolerates_none(): + assert not uses_connector(KvCacheConnectorConfig(connector="kvbm"), "mooncake-store") + assert not uses_connector(None, "mooncake-store") + + +def test_uses_connector_rejects_an_unknown_preset(): + config = KvCacheConnectorConfig(connector="mooncake-store") + with pytest.raises(ValueError, match="Unknown connector preset"): + uses_connector(config, "mooncake-stroe") + + # ---- worker ---- From 6056b19c2ec0926fda57080b91cf72faa5237ed6 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:45:01 -0700 Subject: [PATCH 10/24] [None][doc] Add a handover guide for the mooncake-store connector mooncake_usage.md is the entry point someone inheriting this work should read first: install, the configuration that works, how to tell from the logs whether the pool is being used, what it measured, and what bites. It stays short by pointing at the reference doc, the SLURM runbook and the working configs rather than restating them. It leads on partial reuse because that single flag decided whether the feature did anything at all, and on which reuse metric counts store hits, because the Prometheus counters exclude them and reading the wrong one makes a working store look inert. Also corrects the runbook, which predicted an unaligned local match would be rare with 128-token blocks. It was the common case: 97.2% of lookups. The prediction is worth replacing rather than deleting, since the arithmetic returns if tokens_per_block changes. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- mooncake_disagg/README.md | 14 ++- mooncake_usage.md | 195 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 204 insertions(+), 5 deletions(-) create mode 100644 mooncake_usage.md diff --git a/mooncake_disagg/README.md b/mooncake_disagg/README.md index 7ed4c20c12ee..514f8e49c9fc 100644 --- a/mooncake_disagg/README.md +++ b/mooncake_disagg/README.md @@ -833,11 +833,15 @@ gives a coarser version of the same check. - **UCX warmup requests hit the store too.** `run_benchmark.sh` sends `2 x ctx_instances x gen_instances` 100-token requests before the real run. Harmless, but they are in the counters. -- **`enable_chunked_prefill` interacts with the offer.** The connector offers - only whole blocks and only when the local match is block-aligned; a partial - local match disables the store for that request entirely. With - `tokens_per_block: 128` this is rare, but it explains occasional zero-offer - requests. +- **A partial local match disables the store for that request entirely,** and + this is the dominant failure mode rather than the rare one it was predicted to + be here. The connector offers only whole blocks and only when the local match + is block-aligned, and `enable_partial_reuse` (default `true`) is exactly what + puts the match off a boundary: measured on M3, it declined **97.2% of + lookups**, so a 1.6 TB pool measured as if it were absent. Turning it off took + actual prompt cache read from 35% to 94%. `py_executor_creator` now forces it + off for this connector, so the hazard is gone, but the arithmetic is worth + knowing before changing `tokens_per_block`. See `../mooncake_usage.md` §3. - **`block_reuse_policy: per_conversation` is off on the context worker** in these configs. It is not gated, but the connector derives its own `cache_salt`-seeded hash chain and the interaction is untested. Restore it diff --git a/mooncake_usage.md b/mooncake_usage.md new file mode 100644 index 000000000000..8d455c5bb1ac --- /dev/null +++ b/mooncake_usage.md @@ -0,0 +1,195 @@ +# Using the mooncake-store KV connector + +The `mooncake-store` connector publishes KV cache pages into a shared, +content-addressed pool in host DRAM, so a prefix computed by one engine can be +replayed by another. It is a KV cache *connector*, unrelated to the Mooncake +*transfer engine* that the cache transceiver can use for prefill/decode handoff +— different component, different config, and they are usually not both in play. + +Use it when local block reuse leaves reuse on the table: several context +instances that see the same prefixes, prefixes that should outlive a restart, or +a working set larger than one node's host memory. It replaces TensorRT-LLM's +native host offload tier rather than layering on top of it (`host_cache_size` +must be `0`). + +This page is the entry point. Depth lives elsewhere: + +| For | Read | +|---|---| +| API surface, gates, keying | `docs/source/features/kv-cache-connector.md` § *Mooncake distributed store* | +| Full SLURM runbook, install debugging, experiment matrix | `mooncake_disagg/README.md` | +| Working configs | `mooncake_disagg/m3_ctx_mooncake.yaml`, `m3_gen_mooncake.yaml` | +| Correctness tests (no GPU, no store) | `tests/unittest/_torch/executor/test_mooncake_store_connector.py` | + +## 1. Install + +The connector needs the Mooncake **Python** bindings +(`mooncake.store.MooncakeDistributedStore`). Containers have shipped the C++ +library for months without them, and a CMake-installed `mooncake` package +shadows the working wheel, so a bare `pip install` reports success and the +import still fails. Inside the container: + +```bash +bash mooncake_disagg/install_mooncake_runtime.sh # idempotent, ~8s warm +python3 -c "from mooncake.store import MooncakeDistributedStore; print('ok')" +``` + +`disaggr_torch.slurm` runs this per node automatically when a worker config +mentions `mooncake-store`. Images built from this branch have it baked in. +`mooncake_disagg/README.md` §2 explains why it is this awkward. + +## 2. Configure + +A `mooncake_master` process must be reachable, and every worker needs +`MOONCAKE_CONFIG_PATH` pointing at a JSON client config naming it. The SLURM +harness starts the master and writes that JSON per job; outside SLURM, see +§4 of the runbook. + +Put the connector on the **context** workers only: + +```yaml +kv_connector_config: + connector: mooncake-store +kv_cache_config: + use_kv_cache_manager_v2: true # required: only V2 describes its pools + enable_block_reuse: true + host_cache_size: 0 # required, and must be explicit, not omitted + disk_cache_size: 0 + tokens_per_block: 128 +scheduler_config: + capacity_scheduler_policy: GUARANTEED_NO_EVICT # required +enable_attention_dp: false # required +``` + +Generation workers deliberately get no `kv_connector_config`: generated tokens +are rarely a reused prefix, and leaving it off is the only way to say "off" +(`StoreRole` has no off value). It also lets them keep their host cache tier and +`MAX_UTILIZATION` scheduler, both of which the connector forbids. + +Per-process environment, on the workers that open a handle: + +| Variable | Purpose | +|---|---| +| `TRTLLM_MOONCAKE_STORE_ROLE` | `producer` / `consumer` / `both` | +| `TRTLLM_MOONCAKE_STORE_PREFIX` | Cache namespace. Bump it after any change to page layout or contents. | +| `TRTLLM_MOONCAKE_STORE_MODEL_KEY` | Defaults to the checkpoint directory's basename — set it explicitly for anything long-lived. | + +Pool capacity comes only from processes that open a store handle, so a +prefill-only connector gives a prefill-only pool. `mooncake_segment_donor.py` +contributes host memory from the generation nodes without any traffic; +`disaggr_torch.slurm` starts one per generation node +(`MOONCAKE_DONOR_SEGMENT_SIZE`, default `32GiB`). + +## 3. Partial reuse must be off — now enforced + +This is the one setting that decides whether the feature works at all. + +The store is addressed by whole blocks. The connector is handed +`num_computed_tokens`, the device match, and offers only blocks beyond it — but +it can only continue from a block boundary, so when the device match lands +mid-block it declines the lookup entirely. `enable_partial_reuse=true` is +precisely what makes the device match land mid-block, so it trades part of one +block of device reuse for *every* stored block of the remaining prefix. + +On MiniMax-M3 that guard declined **97.2% of lookups**. The pool was never +asked, and a 1.6 TB pool measured as if it were not there. + +`py_executor_creator` now forces `enable_partial_reuse=false` whenever this +connector is configured, and says so: + +``` +Disabling partial reuse: it is not usable with the mooncake-store connector... +``` + +Nothing to set; the warning fires even from the default (`true`), and is the +confirmation that the coercion ran. The field is otherwise untouched, so +configs that already set `false` are unaffected. + +## 4. Verify a run + +Startup, at INFO, on every context worker: + +```bash +grep -h "mooncake-store" /3_output_CTX_*.log | head -40 +``` + +`registered layout: ... bytes/page=...` is the line to keep — pool sizing +depends on it, and `window=None` confirms no sliding-window group (one would +have aborted startup). + +Then check that the pool spans the hosts you expect. +`disaggr_torch.slurm` writes the per-segment breakdown to +`/9_mooncake_summary.log`; a single host means a prefill-only pool. +Pool occupancy and eviction come from `/2_mooncake_master.log`. + +**Which reuse number counts store hits:** per-request stats +(`reused_blocks_per_request`, `kv_cache_hit_rate_per_request`) **do**; +`/prometheus/metrics` iteration counters (`kv_cache_iter_reused_blocks`) **do +not** — those come from the local reuse tree. So store hits ≈ per-request reuse +− local-tree reuse. All of this needs `enable_iter_perf_stats`, +`enable_iter_req_stats` and `return_perf_metrics`, which all default to false. + +## 5. What it measured + +MiniMax-M3-NVFP4 on GB300, 1 context server (TP=2) + 2 generation servers +(TP=4), connector on context only, ~1.6 TB pool (160 GiB per context rank plus +640 GiB donated per generation node), real conversation trace. + +"Baseline" is the same configuration with partial reuse left at its default. + +| Run | Theoretical hit | Actual hit | Output tok/s | +|---|---|---|---| +| c50 baseline | 96.29% | 35.19% | 319.43 | +| **c50 with partial reuse off** | 96.64% | **93.53%** | **697.29** (2.18x) | +| c70 baseline | 96.01% | 38.44% | 397.43 | +| **c70 with partial reuse off** | 96.51% | **86.46%** | **643.57** (1.62x) | + +A 61-point gap between the reuse the workload allowed and the reuse the system +achieved closed to 3 points. For comparison, native host offload on the same +workload reached 35.59% actual hit at 318.14 tok/s — it wrote 1.83 TB to host +and read 32.5 GB back, behaving as a write-only tier. + +Where the reuse comes from, in steady state (attribution counters, c50): +**~95% of all reuse is served by the pool and ~5% by the device cache.** The +residual ~3–5% of misses are prefixes never written by anyone, which no store +can serve. Blocks stranded behind a contiguity gap measured exactly zero, as did +unattributed blocks. + +**Peak is at c50, not higher.** By c70 the pool runs 85–90% full with active +eviction, hit rate falls to 86% and throughput with it. Concurrency headroom is +a function of pool size; size the pool for the working set rather than assuming +the c50 result scales. + +## 6. Things that will bite + +- **`host_cache_size: 0` must be written explicitly.** Left at its `None` + default, V2 still provisions a host tier and startup is rejected. Falsy is not + the same as absent. +- **The key namespace pins world size, rank, `tokens_per_block`, layer groups + and `bytes_per_page`.** Change tensor parallelism and every stored page + becomes unreachable — a miss, not an error. +- **No build hash in the key.** After changing page layout or contents, bump + `TRTLLM_MOONCAKE_STORE_PREFIX` or restart the master. +- **Loads are synchronous** (`start_load_kv`, before the forward pass), so every + loaded byte is exposed to TTFT. A store hit wins only when it displaces real + prefill. +- **`mooncake-store failed to load N of M pages` is not flaky.** The runtime had + already counted those tokens as computed; it is the tripwire against a wrong + answer. Stop and investigate. +- **Rejected outright at startup:** pipeline or context parallelism, + sliding-window attention, attention DP, beam search, host/disk cache tiers, + `MAX_UTILIZATION`, and M3's index-V cache unless + `sparse_disable_index_value=true`. `mooncake_disagg/README.md` §9 has the rest. + +## 7. Diagnostics + +The `user/brb/m3-mooncake-store-instrumentation` branch carries attribution +counters that sort every reusable block of every prompt into served-from-device, +served-from-pool, stranded behind a contiguity gap, never written, evicted, or +unattributed, with cumulative and per-window reporting. That is what produced +§5's split. They are kept off this branch because they cost a store probe on +lookups the connector would otherwise decline. + +Enable with `MOONCAKE_DEBUG_COVERAGE=1`, which must be set via +`environment.ctx_worker_env_var` — the SLURM harness consumes `MOONCAKE_*` +itself to build the pool config and does not forward it to workers. From abe576cb7d3194c7aeb6ba024f931b655f28f08d Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:22:48 -0700 Subject: [PATCH 11/24] [None][fix] Preempt context requests when the KV pool has nothing to spill to The V2 scheduler's only way to reclaim GPU pages was to suspend a request, which unpins its pages so the eviction controller can migrate them one cache level down. With GPU as the last level a suspended page stays HELD, which is not evictable there, so suspension frees nothing. A KV connector run is exactly that configuration: the manager drops the automatic host tier because tier migration would reassign GPU slots and invalidate the device addresses the connector registered. A prefill server then fills its pool, admits nothing further, and spins at full speed scheduling nothing -- looking healthy to the hang detector and to /health while burning the rest of its wall clock. Give the scheduler a second reclaim action for that case. Preemption closes the victim's KVCache instead of parking it, which returns its committed blocks to the radix tree as reusable prefix and leaves the pages DROPPABLE, evictable at every level. The data is not discarded: it stays resident and locally matchable until something else actually needs the space, and with a connector attached the blocks already written to the store come back through the ordinary prefix load. Recompute is the always-correct fallback. Pages are not released while the connector still has saves reading out of them, which would otherwise let a later request overwrite the bytes mid-transfer and publish them under a valid hash. The victim goes through the same request_finished/get_finished handshake a finished request uses, and the executor resets it to context state once every rank reports the saves retired. Configurations with a tier below GPU keep the existing suspend-based path untouched. Also replace the deadlock detector's single-iteration check with a consecutive-stall counter that counts context candidates as well. The old check only looked at generation requests, which a disaggregated prefill server does not have, and raising on one stalled iteration would misfire on the transient deferrals the scheduler makes for multimodal chunk alignment and PEFT budget. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 2 + .../_torch/pyexecutor/kv_cache_manager_v2.py | 105 ++++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 22 ++ .../pyexecutor/scheduler/scheduler_v2.py | 219 +++++++++++--- .../executor/test_kv_cache_v2_scheduler.py | 276 ++++++++++++++++++ 5 files changed, 584 insertions(+), 40 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 92843c42f3aa..6bb66c0a53e2 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2714,6 +2714,8 @@ def create_py_executor_instance( cross_kv_cache_manager=cross_kv_cache_manager, no_schedule_until_state=no_schedule_until_state, enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, + max_input_len=max_seq_len + if max_seq_len is not None else 0x7fffffff, ) elif (scheduler_config is not None and scheduler_config.use_python_scheduler): diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 944036b6155b..2f219fe16cb0 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -980,13 +980,14 @@ def append_to_kv_heads_per_layer( cache_tiers: List[CacheTierConfig] = [GpuCacheTierConfig(quota=int(quota))] if kv_connector_manager is not None and kv_cache_config.host_cache_size is None: # A KV connector registers device addresses for its pages, and a - # page evicted to another tier has its GPU slot reassigned. The - # automatic host tier below exists only to give the MAX_UTILIZATION - # scheduler's suspend/resume somewhere to spill to, and a connector - # run cannot use that policy (py_executor_creator requires - # GUARANTEED_NO_EVICT), so skip it rather than silently migrating - # pages out from under the connector. An explicitly configured - # host_cache_size is left alone and rejected loudly at bring-up. + # page evicted to another tier has its GPU slot reassigned, so the + # automatic host tier below would silently migrate pages out from + # under the connector. An explicitly configured host_cache_size is + # left alone and rejected loudly at bring-up. + # + # That leaves the scheduler without a tier to spill to, where + # suspension frees nothing. What reclaims pages instead is + # preemption -- see KVCacheManagerV2.preempt_request. host_quota = 0 logger.info( "KV cache manager v2 host tier disabled: a KV connector is attached " @@ -1173,6 +1174,9 @@ def append_to_kv_heads_per_layer( ) self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) self._early_freed_index_requests: set[int] = set() + # Requests whose pages a connector is still reading from, so the + # release half of `preempt_request` has to wait. See that method. + self._pending_preemption: Dict[int, LlmRequest] = {} self._prepare_page_table_tensor(index_mapper_capacity) self._log_kv_cache_pool_lifecycle_mapping() @@ -2614,6 +2618,89 @@ def resume_request(self, req: LlmRequest) -> bool: return False return self._resume_and_restore(req.py_request_id, kv_cache) + # ---- preemption ---- + # + # Suspension is the cheap way to free pages, but it only unpins them: the + # eviction controller then migrates them one cache level down. With GPU as + # the last level a suspended page stays `HELD`, which + # `CacheLevelManager.is_evictable` refuses to evict, so suspension frees + # nothing and the scheduler has no way out of a full pool. + # + # Preemption is the fallback for that case. It gives the pages up instead + # of parking them, which costs a re-prefill but always works. + + @property + def has_cache_tier_below_gpu(self) -> bool: + """True when a suspended page has somewhere to be evicted to.""" + return len(self.impl.cache_tier_list) > 1 + + def has_pending_preemption(self) -> bool: + """True while a deferred preemption is still waiting on a connector.""" + return bool(self._pending_preemption) + + def preempt_request(self, req: LlmRequest) -> bool: + """Give up *req*'s KV cache so its pages can be reclaimed. + + Unlike :meth:`suspend_request` this does not keep the pages. Closing + the request's ``_KVCache`` returns its committed blocks to the radix + tree as reusable prefix and leaves their pages ``DROPPABLE``, which is + evictable at every level -- including the last, where ``HELD`` is not. + So the data is not thrown away: it stays resident and locally matchable + until something else actually needs the space. + + The request is reset to context state by the caller and re-prefills + whatever it can no longer match. With a connector attached the blocks + it already wrote to the store come back through the ordinary prefix + load, so the reload is usually cheap, and recompute is the + always-correct fallback when the store no longer has them. + + Returns True when the pages were released. When the connector still has + saves in flight the release is deferred and this returns False: those + saves read directly out of these pages, so freeing them now would let a + later request overwrite the bytes mid-transfer and publish them to the + store under a valid hash. Callers must not count on the pages until + :meth:`try_complete_preemption` has run for this request. + """ + if self.kv_connector_manager is None: + self._release_preempted(req) + return True + + # Same handshake the finish path uses, for the same reason. The + # request lands in DISAGG_CONTEXT_TRANS_IN_PROGRESS, out of the + # schedulable range, and its `_KVCache` keeps holding the pages until + # every rank reports the save retired through `get_finished`. + if self.kv_connector_manager.request_finished( + req, self.get_connector_page_indices(req) + ): + self._pending_preemption[req.py_request_id] = req + return False + + self._release_preempted(req) + return True + + def try_complete_preemption(self, req: LlmRequest) -> bool: + """Release pages for a request whose deferred preemption just cleared. + + Returns False when *req* was not awaiting preemption, which is how the + caller tells a preempted request apart from an ordinary finished one + in the connector's ``get_finished`` output. + """ + if self._pending_preemption.pop(req.py_request_id, None) is None: + return False + self._release_preempted(req) + return True + + def _release_preempted(self, req: LlmRequest) -> None: + self.free_resources(req) + # Ask the connector again on re-admission rather than reusing the + # memoised offer: the store has strictly more of this prefix now than + # it did when the request was first admitted. + req.py_connector_prefix_start = None + req.py_connector_prefix_end = None + req.py_connector_load_async = False + req.py_connector_delivered = False + req.py_num_connector_matched_tokens = 0 + # ---- prepare_resources ---- @nvtx_range("prepare_resources_kv_cache_manager_v2") @@ -3430,6 +3517,10 @@ def release_index_slot(self, request_id: int) -> None: self._early_freed_index_requests.add(request_id) def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + # A request awaiting preemption can still be cancelled or fail while + # its saves drain. Dropping the entry here keeps a dead request from + # blocking every later preemption via has_pending_preemption(). + self._pending_preemption.pop(request.py_request_id, None) self._release_undelivered_connector_prefix(request) if self.conversation_manager is not None: self.conversation_manager.finish_request(request) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 534035e33ac2..40581baaa0df 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3750,8 +3750,30 @@ def _kv_connector_terminate_requests(self): if self.kv_connector_manager: reqs_to_terminate = self.kv_connector_manager.get_finished() for req in reqs_to_terminate: + if self._resume_preempted_request(req): + continue self._end_transfer_and_maybe_terminate(req) + def _resume_preempted_request(self, request: LlmRequest) -> bool: + """Complete a preemption whose connector saves have now retired. + + The scheduler preempts a request by handing it to the connector the + same way a finished request is handed over, so that its pages stay put + until every rank reports the in-flight saves done. Both kinds come + back through ``get_finished``; only the KV cache manager knows which + is which. + + Returns True when *request* was preempted rather than finished, in + which case its pages are now released and it is back in context state + awaiting a re-prefill. + """ + if not self._is_kv_manager_v2: + return False + if not self.kv_cache_manager.try_complete_preemption(request): + return False + request.pause(self.max_input_len) + return True + def _kv_connector_wait_for_save(self): if self.kv_connector_manager is not None: self.kv_connector_manager.worker.wait_for_save( diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index dd9b67132834..305a1627c248 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -15,7 +15,7 @@ import enum import os -from typing import Optional +from typing import Callable, Optional from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy from tensorrt_llm.logger import logger @@ -154,8 +154,13 @@ def __init__( draft_kv_cache_manager: object | None = None, # KVCacheManagerV2 for MTP draft layers cross_kv_cache_manager: object | None = None, # KVCacheManagerV2 for enc-dec cross-attn enable_prefix_aware_scheduling: bool = True, + max_input_len: int = 0x7FFFFFFF, ) -> None: self.max_num_tokens = max_num_tokens + # Only read when preempting: LlmRequest.pause clamps the rewritten + # prompt (original prompt plus tokens generated so far) to it. + self.max_input_len = max_input_len + self._stalled_schedules = 0 self.max_num_requests = ( scheduler_capacity if scheduler_capacity is not None else max_batch_size ) @@ -398,11 +403,28 @@ def _schedule_loop(self, active_requests, inflight_request_ids): req_it += 1 + # Requests whose pages were given up during this pass. They must not + # be scheduled afterwards: a victim that is itself a started context + # request is already sitting in pending_ctx, and re-admitting it here + # would spend the pages the preemption just released on the very + # request that released them. + preempted_ids: set[int] = set() + + def preempt_for_pages(req: LlmRequest) -> bool: + protected = {r.py_request_id for r in scheduled_gen} + protected.update(r.py_request_id for r in scheduled_ctx) + protected.add(req.py_request_id) + return self._try_preempt_for_pages( + requests_list, protected, inflight_request_ids, evicted, preempted_ids + ) + # --- Phase 2: schedule deferred context / encoder requests --- # Generation PEFT pages are now fully committed in the budget. for req in pending_ctx: if budget.requests_full: break + if req.py_request_id in preempted_ids: + continue peft_pages = budget.peft_pages_needed(req) if peft_pages is None: continue @@ -413,7 +435,9 @@ def _schedule_loop(self, active_requests, inflight_request_ids): scheduled_encoder.append(req) budget.commit(req, tokens, peft_pages) else: - action, tokens, chunking_flag = self._try_schedule_context(req, budget) + action, tokens, chunking_flag = self._try_schedule_context( + req, budget, preempt_for_pages + ) if action is ScheduleAction.STOP: break if action is ScheduleAction.SKIP: @@ -422,28 +446,19 @@ def _schedule_loop(self, active_requests, inflight_request_ids): scheduled_ctx.append(req) budget.commit(req, tokens, peft_pages) - # Deadlock detection: if generation requests exist but none were - # scheduled and none were evicted, no forward pass will run and no - # KV cache pages will ever be freed — the scheduler will spin - # forever. This typically happens when the KV cache pool is - # exhausted and no host cache tier is available for suspend/resume. - if not scheduled_gen and not scheduled_ctx: - num_gen_candidates = sum( - 1 - for r in active_requests - if r.is_generation_in_progress_state - and not r.is_generation_to_complete_state - and r.request_id not in inflight_request_ids - ) - if num_gen_candidates > 0 and not evicted: - raise RuntimeError( - f"V2 scheduler deadlock: {num_gen_candidates} generation " - f"request(s) active but none could be scheduled or " - f"evicted. KV cache pool is likely exhausted with no " - f"host cache tier for suspend/resume offload. " - f"Configure kv_cache_config.host_cache_size or increase " - f"kv_cache_config.max_tokens." - ) + self._detect_deadlock( + active_requests, + inflight_request_ids, + pending_ctx, + preempted_ids, + made_progress=bool( + scheduled_gen + or scheduled_ctx + or scheduled_encoder + or disagg_candidates + or evicted + ), + ) return ( scheduled_encoder, @@ -505,7 +520,10 @@ def _try_schedule_disagg_gen_init( return ScheduleAction.SCHEDULED, 0 def _try_schedule_context( - self, req: LlmRequest, budget: BudgetTracker + self, + req: LlmRequest, + budget: BudgetTracker, + preempt_for_pages: Callable[[LlmRequest], bool], ) -> tuple[ScheduleAction, int, bool]: """Try to schedule a context request (chunked or non-chunked). @@ -521,11 +539,14 @@ def _try_schedule_context( # connector being asked or told twice is the per-request query state # (see KVCacheManagerV2._connector_prefix_position). if self.chunking_enabled: - return self._try_schedule_context_chunked(req, budget) - return self._try_schedule_context_full(req, budget) + return self._try_schedule_context_chunked(req, budget, preempt_for_pages) + return self._try_schedule_context_full(req, budget, preempt_for_pages) def _try_schedule_context_full( - self, req: LlmRequest, budget: BudgetTracker + self, + req: LlmRequest, + budget: BudgetTracker, + preempt_for_pages: Callable[[LlmRequest], bool], ) -> tuple[ScheduleAction, int, bool]: """Try to schedule a non-chunked context request. @@ -564,6 +585,11 @@ def _try_schedule_context_full( # V2 resizes KV cache directly in the scheduler (no separate # prepareResources for main cache), so include draft tokens. if not self.kv_cache_manager.resize_context(req, context_tokens + draft_len): + # Out of pages. Give one started request up so this one can + # proceed, and retry next iteration rather than now: a failed + # resize leaves a first chunk suspended, so the retry has to go + # back through prepare_context to resume it. + preempt_for_pages(req) return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) @@ -574,7 +600,10 @@ def _try_schedule_context_full( return ScheduleAction.SCHEDULED, req_tokens, False def _try_schedule_context_chunked( - self, req: LlmRequest, budget: BudgetTracker + self, + req: LlmRequest, + budget: BudgetTracker, + preempt_for_pages: Callable[[LlmRequest], bool], ) -> tuple[ScheduleAction, int, bool]: """FCFS interleaved chunking for a single context request. @@ -635,10 +664,10 @@ def _try_schedule_context_chunked( chunk_size = (chunk_size // self.chunk_unit_size) * self.chunk_unit_size if chunk_size <= 0: - # TODO: consider suspending first-chunk KVCache to release - # GPU pages. Currently we skip without suspend to avoid - # pathological suspend/resume cycles. suspend_request is - # only called from eviction (_try_evict_for_gen). + # Out of token budget, not out of pages, so releasing this + # request's pages would not help: next iteration gets a fresh + # budget. Deliberately not suspended, to avoid pathological + # suspend/resume cycles. return ScheduleAction.SKIP, 0, False chunk_size = self._align_chunk_to_mm_block( @@ -666,6 +695,8 @@ def _try_schedule_context_chunked( # V2 resizes KV cache directly in the scheduler, so include # draft tokens for last chunk. if not self.kv_cache_manager.resize_context(req, resize_tokens): + # Out of pages — see the same call in _try_schedule_context_full. + preempt_for_pages(req) return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) @@ -1010,6 +1041,128 @@ def _suspend_request(self, req: LlmRequest) -> None: def _clear_request_runtime_state(self, req: LlmRequest) -> None: req.py_batch_idx = None + def _try_preempt_for_pages( + self, + requests_list: RequestList, + protected_ids: set[int], + inflight_request_ids: set[int], + evicted: RequestList, + preempted_ids: set[int], + ) -> bool: + """Release one started request's KV cache so another can allocate. + + This is the fallback for a pool that suspension cannot drain -- see + ``KVCacheManagerV2.preempt_request``. With a cache tier below GPU, + suspension is cheaper and keeps the pages, so leave that path alone. + + Returns True when pages became available in this iteration. A + connector defers the release until its in-flight saves retire, in + which case this returns False and the caller should give up for now: + the pages arrive a few iterations later. + """ + if self.kv_cache_manager.has_cache_tier_below_gpu: + return False + + if self.kv_cache_manager.has_pending_preemption(): + # One victim at a time. A full pool would otherwise preempt the + # whole batch while the first release is still draining, and + # every one of those requests would have to re-prefill. + return False + + # Newest first, so the requests closest to completing keep their + # pages and the pool drains instead of thrashing. + for i in range(len(requests_list) - 1, -1, -1): + victim = requests_list[i] + if victim.py_request_id in protected_ids: + continue + if victim.request_id in inflight_request_ids: + continue + if not self._is_started_request(victim): + continue + if not self.kv_cache_manager.is_request_active(victim.py_request_id): + continue + + released = self.kv_cache_manager.preempt_request(victim) + logger.debug( + f"[V2Scheduler] Preempting request {victim.py_request_id} " + f"(state={victim.state.name}), pages " + f"{'released' if released else 'pending connector saves'}" + ) + self._clear_request_runtime_state(victim) + if self.draft_kv_cache_manager is not None: + self.draft_kv_cache_manager.free_resources(victim) + if released: + # Rewrites the prompt to include whatever was generated and + # resets state to CONTEXT_INIT with the chunk position at 0, + # so the request re-enters as an ordinary prefill next + # iteration. Deferred releases are paused by the executor + # once the connector reports the saves retired. + victim.pause(self.max_input_len) + evicted.append(victim) + preempted_ids.add(victim.py_request_id) + return released + + return False + + # Consecutive scheduling passes that reclaimed nothing before the + # scheduler calls it a deadlock. A stalled pass costs ~2ms, so this trips + # in seconds, while the transient one-iteration deferrals (multimodal + # chunk alignment, PEFT budget, IndexMapper slots) clear long before it. + _DEADLOCK_STALL_ITERS = 1000 + + def _detect_deadlock( + self, + active_requests: RequestList, + inflight_request_ids: set[int], + pending_ctx: RequestList, + preempted_ids: set[int], + made_progress: bool, + ) -> None: + """Fail loudly when no request can be scheduled or reclaimed. + + Without this the executor spins at full speed while scheduling + nothing: the loop looks healthy to the hang detector and to + ``/health``, and the job burns its wall clock. Context candidates + count alongside generation ones because a disaggregated prefill + server has no generation requests at all, and counting only those + left it spinning silently. + """ + if made_progress: + self._stalled_schedules = 0 + return + + num_gen_candidates = sum( + 1 + for r in active_requests + if r.is_generation_in_progress_state + and not r.is_generation_to_complete_state + and r.request_id not in inflight_request_ids + ) + num_ctx_candidates = sum( + 1 + for r in pending_ctx + if r.py_request_id not in preempted_ids + and r.request_id not in inflight_request_ids + ) + if num_gen_candidates == 0 and num_ctx_candidates == 0: + # Legitimately idle: nothing to schedule. + self._stalled_schedules = 0 + return + + self._stalled_schedules += 1 + if self._stalled_schedules < self._DEADLOCK_STALL_ITERS: + return + + raise RuntimeError( + f"V2 scheduler deadlock: {num_gen_candidates} generation and " + f"{num_ctx_candidates} context request(s) active but none could " + f"be scheduled, suspended or preempted in " + f"{self._stalled_schedules} consecutive attempts. The KV cache " + f"pool is likely exhausted. Configure " + f"kv_cache_config.host_cache_size, increase " + f"kv_cache_config.max_tokens, or lower max_batch_size." + ) + def _is_evictable(self, req: LlmRequest) -> bool: """A started request whose KV cache is still active on GPU. diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index 0a7c67e9e702..a77f517c0e2b 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -160,6 +160,9 @@ def make_kv_cache_manager( resize_context_fn=None, prepare_disagg_gen_init_fn=None, try_allocate_generation_fn=None, + has_cache_tier_below_gpu=True, + preempt_request_fn=None, + has_pending_preemption=False, ): mgr = Mock() mgr.tokens_per_block = tokens_per_block @@ -170,6 +173,11 @@ def make_kv_cache_manager( mgr.try_allocate_generation.side_effect = try_allocate_generation_fn or (lambda req: True) mgr.suspend_request.return_value = None mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active + # Preemption is the fallback for a pool with nothing under GPU to spill + # to, so the default here (a host tier exists) leaves it switched off. + mgr.has_cache_tier_below_gpu = has_cache_tier_below_gpu + mgr.has_pending_preemption.return_value = has_pending_preemption + mgr.preempt_request.side_effect = preempt_request_fn or (lambda req: True) return mgr @@ -191,6 +199,7 @@ def make_scheduler( no_schedule_after_state: LlmRequestState | None = None, cross_kv_cache_manager: Mock | None = None, enable_prefix_aware_scheduling: bool = True, + max_input_len: int | None = None, ) -> object: """Create KVCacheV2Scheduler, patching isinstance check for mock mgr.""" from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler @@ -206,6 +215,8 @@ def make_scheduler( kwargs["no_schedule_after_state"] = no_schedule_after_state if cross_kv_cache_manager is not None: kwargs["cross_kv_cache_manager"] = cross_kv_cache_manager + if max_input_len is not None: + kwargs["max_input_len"] = max_input_len return KVCacheV2Scheduler( max_batch_size=max_batch_size, max_num_tokens=max_num_tokens, @@ -827,6 +838,271 @@ def test_self_eviction_no_started_in_range(self): assert len(out.context_requests) == 0 +# =========================================================================== +# Preemption (context side, no cache tier below GPU) +# =========================================================================== + + +def _out_of_pages_for(request_id): + """resize_context that only fails for *request_id*.""" + return lambda req, n: req.py_request_id != request_id + + +class TestContextPreemption: + """Releasing a started request's pages when suspension cannot help. + + Suspension only unpins pages so the eviction controller can migrate them + one level down; with GPU as the last level a suspended page stays HELD and + unevictable, so it frees nothing. These tests cover the fallback that + gives the pages up instead. + """ + + def test_out_of_pages_preempts_started_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + reqs = [make_ctx_request(0, 100), victim] + + out = sched.schedule_request(reqs, set()) + + mgr.preempt_request.assert_called_once_with(victim) + assert ids(out.paused_requests) == [99] + # Deferred to the next iteration: a failed resize leaves the first + # chunk suspended, so the retry has to go back through + # prepare_context. + assert len(out.context_requests) == 0 + + def test_released_victim_is_reset_to_context_state(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000, max_input_len=4096) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + victim.py_batch_idx = 7 + + sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + victim.pause.assert_called_once_with(4096) + assert victim.py_batch_idx is None + + def test_deferred_release_leaves_pause_to_the_executor(self): + """A connector still reading these pages defers the release. + + Freeing them now would let a later request overwrite bytes + mid-transfer, so the executor pauses the request only once the + connector reports the saves retired. + """ + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + preempt_request_fn=lambda req: False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + mgr.preempt_request.assert_called_once_with(victim) + victim.pause.assert_not_called() + assert ids(out.paused_requests) == [99] + + def test_preempted_victim_not_scheduled_in_the_same_pass(self): + """Re-admitting the victim would spend the pages it just released.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + assert ids(out.context_requests) == [] + + def test_skipped_when_a_cache_tier_exists_below_gpu(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=True, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + mgr.preempt_request.assert_not_called() + # Suspension is cheaper and keeps the pages, so that path is left + # exactly as it was: the request is simply skipped. + assert ids(out.context_requests) == [99] + + def test_one_victim_at_a_time_while_a_release_is_draining(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + has_pending_preemption=True, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + mgr.preempt_request.assert_not_called() + + def test_never_preempts_a_scheduled_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(1), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + # gen0 is scheduled in phase 1; ctx1 then runs out of pages and must + # not take the pages out from under it. + reqs = [make_gen_request(0), make_ctx_request(1, 100)] + + out = sched.schedule_request(reqs, set()) + + assert ids(out.generation_requests) == [0] + mgr.preempt_request.assert_not_called() + + def test_never_preempts_itself(self): + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + req = make_ctx_request(0, 100, is_first_context_chunk=False) + + sched.schedule_request([req], set()) + + mgr.preempt_request.assert_not_called() + + def test_never_preempts_an_inflight_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + sched.schedule_request([make_ctx_request(0, 100), victim], {99}) + + mgr.preempt_request.assert_not_called() + + def test_never_preempts_a_first_chunk_or_suspended_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + # First chunk: holds no pages worth taking. + first_chunk = make_ctx_request(98, 100, is_first_context_chunk=True) + # Already suspended: preempting it frees nothing extra. + suspended = make_ctx_request(99, 100, is_first_context_chunk=False) + mgr.kv_cache_map[suspended.py_request_id].is_active = False + + sched.schedule_request([make_ctx_request(0, 100), first_chunk, suspended], set()) + + mgr.preempt_request.assert_not_called() + + def test_chunked_context_out_of_pages_preempts(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler( + mgr, + max_num_tokens=1000, + ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, 64), + ) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 500), victim], set()) + + mgr.preempt_request.assert_called_once_with(victim) + assert ids(out.paused_requests) == [99] + + +# =========================================================================== +# Deadlock detection +# =========================================================================== + + +class TestDeadlockDetection: + """The scheduler must fail loudly rather than spin scheduling nothing. + + A stalled pass costs a couple of milliseconds, so an undetected stall + burns a job's whole wall clock while the loop still looks healthy to the + hang detector and to /health. + """ + + def test_raises_after_repeated_stalls_with_context_candidates(self): + """A prefill-only worker has no generation requests to count.""" + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + reqs = [make_ctx_request(0, 100, is_first_context_chunk=False)] + + for _ in range(2): + sched.schedule_request(reqs, set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request(reqs, set()) + + def test_raises_after_repeated_stalls_with_generation_candidates(self): + mgr = make_kv_cache_manager(try_allocate_generation_fn=lambda req: False) + sched = make_scheduler(mgr, max_num_tokens=100) + sched._DEADLOCK_STALL_ITERS = 3 + # Self-eviction suspends it on the first pass, which counts as + # progress; afterwards it is inactive and nothing can be reclaimed. + reqs = [make_gen_request(0)] + + for _ in range(3): + sched.schedule_request(reqs, set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request(reqs, set()) + + def test_transient_stall_does_not_raise(self): + """One bad iteration is normal; the counter has to reset.""" + fail = [True] + + def resize_fn(req, n): + return not fail[0] + + mgr = make_kv_cache_manager( + resize_context_fn=resize_fn, has_cache_tier_below_gpu=False + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + reqs = [make_ctx_request(0, 100, is_first_context_chunk=False)] + + for _ in range(10): + sched.schedule_request(reqs, set()) + fail[0] = not fail[0] + + def test_idle_scheduler_never_raises(self): + mgr = make_kv_cache_manager() + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 2 + + for _ in range(5): + out = sched.schedule_request([], set()) + assert len(out.context_requests) == 0 + + def test_all_candidates_inflight_never_raises(self): + """Requests in the PP pipeline are progressing, just not here.""" + mgr = make_kv_cache_manager(resize_context_fn=lambda req, n: False) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 2 + reqs = [make_ctx_request(0, 100)] + + for _ in range(5): + sched.schedule_request(reqs, {0}) + + # =========================================================================== # PEFT / LoRA # =========================================================================== From 8549d5b79a607342d30204fae165ef4e24526e86 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:09:20 -0700 Subject: [PATCH 12/24] [None][feat] provision the Mooncake store pool during trtllm-serve bringup The mooncake-store connector needs a reachable mooncake_master and a MOONCAKE_CONFIG_PATH naming it, both of which only the SLURM benchmark harness knew how to produce. A plain trtllm-serve therefore could not use the connector without borrowing that harness. Describe the pool in kv_connector_config.mooncake_store instead and the server provisions it itself: resolve the master, render the client config, export MOONCAKE_CONFIG_PATH before the ranks that open store handles are spawned, and tear down what it started on exit. launch_master: true starts a master for a single engine; master_server_address joins one with its own lifetime, which is what sharing a pool or surviving a restart requires. An inherited MOONCAKE_CONFIG_PATH still wins and logs that it did, so the harness path is unchanged. An external master is probed during startup rather than left to fail inside store.setup on every rank after the model has loaded. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docs/source/features/kv-cache-connector.md | 25 +- mooncake_disagg/README.md | 9 + mooncake_disagg/m3_agg_mooncake.yaml | 20 + mooncake_usage.md | 34 +- .../connectors/mooncake_store/__init__.py | 8 + .../connectors/mooncake_store/master.py | 315 ++++++++++++++ tensorrt_llm/commands/serve.py | 136 +++--- tensorrt_llm/llmapi/llm_args.py | 106 ++++- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../executor/test_mooncake_store_master.py | 404 ++++++++++++++++++ 10 files changed, 999 insertions(+), 59 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py create mode 100644 tests/unittest/_torch/executor/test_mooncake_store_master.py diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 1338d0e32359..a00a68aa5a92 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -108,12 +108,29 @@ This is a **different component** from the Mooncake transfer engine that the C++ * `KVCacheManagerV2` (`kv_cache_config.use_kv_cache_manager_v2: true`), since that is the manager that can describe its pools through `register_kv_cache_layout`. * The Mooncake Python bindings: `pip install mooncake-transfer-engine`. These are installed in the release container; the source build of the C++ transfer engine does not provide them. -* A running Mooncake master (and metadata server, unless using `P2PHANDSHAKE`). See the [Mooncake documentation](https://kvcache-ai.github.io/Mooncake/). +* A reachable Mooncake master (and metadata server, unless using `P2PHANDSHAKE`). See the [Mooncake documentation](https://kvcache-ai.github.io/Mooncake/). `trtllm-serve` can start one for a single engine; see below. * GPU-only KV cache tiers: set `kv_cache_config.host_cache_size: 0` and `disk_cache_size: 0`. A page evicted to another tier has its GPU slot reassigned, which would invalidate the addresses registered with the store. #### Configuration -Topology comes from a JSON file named by `MOONCAKE_CONFIG_PATH`, using the same schema as the vLLM Mooncake store connector so one deployment can point both engines at the same pool: +Describe the pool in `kv_connector_config.mooncake_store` and `trtllm-serve` provisions it during bringup: it resolves the master, renders the client config, and exports `MOONCAKE_CONFIG_PATH` before the ranks that open store handles are spawned. + +```yaml +kv_connector_config: + connector: mooncake-store + mooncake_store: + master_server_address: 10.0.0.1:50051 # a master with its own lifetime + protocol: rdma + device_name: mlx5_0 + global_segment_size: 32GiB + local_buffer_size: 1GiB +``` + +Replacing `master_server_address` with `launch_master: true` makes the server start a `mooncake_master` itself and use it, so a single-instance deployment needs nothing prepared outside `trtllm-serve`. **That master lives and dies with the server**, which makes it wrong for anything else: several engines that should share one pool would each get their own, and a pool meant to survive a restart cannot be owned by the thing restarting. Those deployments run a master with its own lifetime and name it in `master_server_address`. + +`TRTLLM_MOONCAKE_MASTER_BINARY` overrides the binary a launched master runs, and `TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) how long startup waits for any master to accept connections -- reaching a master that is not there otherwise fails inside every rank after the model has loaded. Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated client config and the master's log, which are otherwise in a temporary directory removed at shutdown. + +Topology can equally come from a JSON file named by `MOONCAKE_CONFIG_PATH`, using the same schema as the vLLM Mooncake store connector so one deployment can point both engines at the same pool: ```json { @@ -126,7 +143,9 @@ Topology comes from a JSON file named by `MOONCAKE_CONFIG_PATH`, using the same } ``` -Two further settings are TensorRT-LLM's rather than Mooncake's, and are read from the environment because `KvCacheConnectorConfig` carries no free-form dictionary: +An inherited `MOONCAKE_CONFIG_PATH` wins over `mooncake_store` and is logged as doing so, so an orchestrator that already provisions the pool -- as the SLURM benchmark harness does -- keeps working unchanged. + +Three further settings are TensorRT-LLM's rather than Mooncake's, and stay in the environment because they are per process rather than per pool: | Variable | Default | Meaning | |---|---|---| diff --git a/mooncake_disagg/README.md b/mooncake_disagg/README.md index 514f8e49c9fc..9ee1795ffce6 100644 --- a/mooncake_disagg/README.md +++ b/mooncake_disagg/README.md @@ -323,6 +323,15 @@ reports on the same line. `master_server_address` is mandatory, so a master must exist and be reachable from every worker. +**Outside SLURM you can skip this section too.** A worker config carrying +`kv_connector_config.mooncake_store` makes `trtllm-serve` provision the pool +during its own bringup — `launch_master: true` starts a master for that server +alone, `master_server_address` joins one that already exists — and write the +client config itself. That covers aggregated and single-instance runs, which is +what `m3_agg_mooncake.yaml` now does; `mooncake_usage.md` §2 has the table. The +rest of this section is about the master the experiments below need, which +outlives any one server and therefore cannot be owned by one. + **For a single-job experiment you can skip this section.** `disaggr_torch.slurm` now starts a `mooncake_master` on the first node of the allocation, waits for its port to accept connections, and writes diff --git a/mooncake_disagg/m3_agg_mooncake.yaml b/mooncake_disagg/m3_agg_mooncake.yaml index a2c0251f0f0f..42f73c018d56 100644 --- a/mooncake_disagg/m3_agg_mooncake.yaml +++ b/mooncake_disagg/m3_agg_mooncake.yaml @@ -87,3 +87,23 @@ enable_attention_dp: false kv_connector_config: connector: mooncake-store + + # Aggregated means one engine, which is the only shape a server-owned + # master is right for: it dies with the server, so nothing else can be + # sharing the pool and nothing can expect it to survive a restart. Point + # master_server_address at a master of its own for either of those. + # + # An inherited MOONCAKE_CONFIG_PATH wins over this block, so the SLURM + # harness keeps naming its own pool. + mooncake_store: + launch_master: true + # TCP removes RDMA from the variable list at the cost of any + # performance conclusion. Set protocol: rdma with a device_name from + # ibv_devinfo for a run worth quoting. + protocol: tcp + device_name: "" + # Contributed per worker process, so the pool is this times the world + # size. Sized for the working set, not for what the node can spare: + # by ~85-90% full the master evicts and the hit rate follows it down. + global_segment_size: 160GiB + local_buffer_size: 4GiB diff --git a/mooncake_usage.md b/mooncake_usage.md index 8d455c5bb1ac..2932ef105087 100644 --- a/mooncake_usage.md +++ b/mooncake_usage.md @@ -41,15 +41,39 @@ mentions `mooncake-store`. Images built from this branch have it baked in. ## 2. Configure A `mooncake_master` process must be reachable, and every worker needs -`MOONCAKE_CONFIG_PATH` pointing at a JSON client config naming it. The SLURM -harness starts the master and writes that JSON per job; outside SLURM, see -§4 of the runbook. +`MOONCAKE_CONFIG_PATH` pointing at a JSON client config naming it. + +Three ways to get there, in increasing order of how much you have to arrange: + +| Deployment | Master | +|---|---| +| One `trtllm-serve`, own pool | `mooncake_store: {launch_master: true}` — the server starts it | +| Several engines, or a pool that outlives them | `mooncake_store: {master_server_address: host:50051}` — a master you run | +| SLURM benchmark harness | Nothing: `disaggr_torch.slurm` starts the master and writes the JSON per job | + +The first two make `trtllm-serve` render the client config and export +`MOONCAKE_CONFIG_PATH` itself. An inherited `MOONCAKE_CONFIG_PATH` wins over +both and says so in the log, which is why the harness path is unaffected. +`mooncake_disagg/README.md` §4 covers running a master as its own SLURM job. + +A launched master dies with the server, so use it only for a single engine: +two context servers that each launch one get two disjoint pools, and the +survival-across-restart case is impossible by construction. + +Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated JSON and the master's log, +which otherwise sit in a temporary directory that shutdown removes. +`TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) bounds the wait for the port; +that wait is also what turns an unreachable external master from a failure in +every rank after the model loads into one line before it starts. Put the connector on the **context** workers only: ```yaml kv_connector_config: connector: mooncake-store + mooncake_store: # omit when an orchestrator sets + launch_master: true # MOONCAKE_CONFIG_PATH for you + protocol: tcp # rdma with a device_name for real numbers kv_cache_config: use_kv_cache_manager_v2: true # required: only V2 describes its pools enable_block_reuse: true @@ -120,7 +144,9 @@ have aborted startup). Then check that the pool spans the hosts you expect. `disaggr_torch.slurm` writes the per-segment breakdown to `/9_mooncake_summary.log`; a single host means a prefill-only pool. -Pool occupancy and eviction come from `/2_mooncake_master.log`. +Pool occupancy and eviction come from `/2_mooncake_master.log`, or +from `$TRTLLM_MOONCAKE_RUN_DIR/mooncake_master.log` when `trtllm-serve` +launched the master — the startup line reports the path either way. **Which reuse number counts store hits:** per-request stats (`reused_blocks_per_request`, `kv_cache_hit_rate_per_request`) **do**; diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index 1664e461421e..a07b7a6aace0 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -34,6 +34,11 @@ with ``MOONCAKE_CONFIG_PATH`` pointing at a Mooncake JSON config. +Describing the pool in ``KvCacheConnectorConfig.mooncake_store`` instead lets +``trtllm-serve`` provision it during bringup -- resolving or launching the +master and writing that JSON itself -- so no external script has to. See +``master.py``. + By default the KV pools themselves are registered with Mooncake, so the store reads and writes device memory and no copy is added. That needs the HCA to be able to pin GPU pages -- GPUDirect RDMA, through ``nvidia_peermem`` or dma-buf. @@ -45,6 +50,7 @@ """ from .config import MooncakeStoreConnectorConfig, StoreRole +from .master import maybe_provision_pool, provision_pool from .scheduler import MooncakeStoreConnectorScheduler from .worker import MooncakeStoreConnectorWorker @@ -53,4 +59,6 @@ "MooncakeStoreConnectorScheduler", "MooncakeStoreConnectorWorker", "StoreRole", + "maybe_provision_pool", + "provision_pool", ] diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py new file mode 100644 index 000000000000..16f91d6c32ee --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Bring the Mooncake store's pool up as part of a server's own startup. + +The connector needs two things that are not the engine's to produce: a +reachable ``mooncake_master``, and a JSON client config named by +``MOONCAKE_CONFIG_PATH`` that points every worker at it. Both were the SLURM +harness's job, which left a single ``trtllm-serve`` unable to use the connector +without borrowing that harness. + +``provision_pool`` does the same work inside the serving process. It resolves +the master -- launching one here, or checking that the configured one answers +-- renders the client config, and exports ``MOONCAKE_CONFIG_PATH``, which +reaches the ranks because the LLM constructor spawns them from this process. +Everything it started is torn down when the context exits. + +A master launched here lives and dies with the server, so it is only right for +one engine talking to its own pool. Several engines sharing a pool, or a pool +meant to survive a restart, need a master with its own lifetime, named by +``master_server_address``. +""" + +import contextlib +import json +import os +import shutil +import socket +import subprocess # nosec B404 +import tempfile +import time +from dataclasses import dataclass +from typing import Any, Dict, Iterator, Optional, Tuple + +from tensorrt_llm.logger import logger + +from ..registry import uses_connector +from .config import CONFIG_PATH_ENV + +__all__ = ["maybe_provision_pool", "provision_pool"] + +#: Override the binary that ``launch_master`` runs. +MASTER_BINARY_ENV = "TRTLLM_MOONCAKE_MASTER_BINARY" +#: How long to wait for a master to accept connections, in seconds. +MASTER_TIMEOUT_ENV = "TRTLLM_MOONCAKE_MASTER_TIMEOUT" +#: Where to keep the generated client config and the master's log. Set it to +#: keep them after shutdown; otherwise they live in a temporary directory. +RUN_DIR_ENV = "TRTLLM_MOONCAKE_RUN_DIR" + +DEFAULT_MASTER_BINARY = "mooncake_master" +DEFAULT_MASTER_TIMEOUT = 60.0 +CLIENT_CONFIG_NAME = "mooncake.json" +MASTER_LOG_NAME = "mooncake_master.log" + + +def _local_address() -> str: + """The address this host is known by inside the pool. + + Deliberately the same derivation the connector worker uses for its own + hostname, so the master and the segments registering with it agree on + which host they are on. + """ + try: + return socket.gethostbyname(socket.gethostname()) + except OSError: + return "127.0.0.1" + + +def _master_timeout() -> float: + raw = os.getenv(MASTER_TIMEOUT_ENV) + if not raw: + return DEFAULT_MASTER_TIMEOUT + try: + timeout = float(raw) + except ValueError as exc: + raise ValueError(f"{MASTER_TIMEOUT_ENV}={raw!r} is not a number") from exc + if timeout <= 0: + raise ValueError(f"{MASTER_TIMEOUT_ENV}={raw!r} must be > 0") + return timeout + + +def _split_address(address: str) -> Optional[Tuple[str, int]]: + """Split ``host:port``, or return ``None`` if it is not in that form.""" + host, separator, port = address.rpartition(":") + if not separator or not port.isdigit(): + return None + return host.strip("[]"), int(port) + + +def _wait_until_accepting( + host: str, + port: int, + timeout: float, + process: Optional[subprocess.Popen] = None, + hint: str = "", +) -> None: + """Block until the master accepts connections. + + A worker that opens its store handle before the master is listening fails + outright, so the port -- not the presence of a process -- is what the + ordering has to wait on. When the master is ours, its exit is checked first + each pass, so a master that died is reported as that rather than as a + timeout. + """ + deadline = time.monotonic() + timeout + while True: + if process is not None and (code := process.poll()) is not None: + raise RuntimeError(f"mooncake_master exited with code {code} during startup.{hint}") + try: + with socket.create_connection((host, port), timeout=1.0): + return + except OSError as exc: + last_error = exc + if time.monotonic() >= deadline: + raise TimeoutError( + f"The Mooncake master at {host}:{port} did not accept " + f"connections within {timeout:g}s ({last_error}). Raise " + f"{MASTER_TIMEOUT_ENV} if it is only slow to start.{hint}" + ) + time.sleep(0.5) + + +def _client_config(pool: Any, master_address: str) -> Dict[str, Any]: + """Render the Mooncake client config for a pool. + + The schema is vLLM's, so one pool can serve both engines. ``role`` is + written as ``both`` because the file describes the pool; which directions + of traffic a given process drives is its own + ``TRTLLM_MOONCAKE_STORE_ROLE``. + """ + config: Dict[str, Any] = { + "metadata_server": pool.metadata_server, + "master_server_address": master_address, + "protocol": pool.protocol, + "device_name": pool.device_name, + "global_segment_size": pool.global_segment_size, + "local_buffer_size": pool.local_buffer_size, + "role": "both", + "transfer_batch_size": pool.transfer_batch_size, + "stage_through_host": pool.stage_through_host, + } + if pool.cache_prefix is not None: + config["cache_prefix"] = pool.cache_prefix + return config + + +@dataclass +class LaunchedMaster: + """A ``mooncake_master`` owned by this process.""" + + process: subprocess.Popen + address: str + log_path: str + + def stop(self, timeout: float = 10.0) -> None: + if self.process.poll() is not None: + return + self.process.terminate() + try: + self.process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.wait() + + +def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: + """Start a master on this host and wait for it to answer.""" + binary = os.getenv(MASTER_BINARY_ENV) or DEFAULT_MASTER_BINARY + resolved = shutil.which(binary) + if resolved is None: + raise FileNotFoundError( + f"{binary!r} is not on PATH, so launch_master cannot start a " + "Mooncake master. It ships with the Mooncake runtime, which " + "mooncake_disagg/install_mooncake_runtime.sh installs. Point " + f"{MASTER_BINARY_ENV} at the binary, or drop launch_master and " + "set master_server_address to a master you run yourself." + ) + + host = _local_address() + log_path = os.path.join(run_dir, MASTER_LOG_NAME) + hint = f" See {log_path}." + + # mooncake_master logs through glog, which writes files under /tmp unless + # told otherwise, so without GLOG_logtostderr the log below stays empty. + # GLOG_v=1 adds the per-RPC lines showing segments registering and keys + # moving, which is the only view of the pool's side of the conversation + # short of scraping the metrics port. + env = dict(os.environ, GLOG_logtostderr="1") + env.setdefault("GLOG_v", "1") + command = [ + resolved, + f"--rpc_port={pool.master_port}", + f"--metrics_port={pool.master_metrics_port}", + f"--eviction_ratio={pool.master_eviction_ratio}", + ] + + logger.info(f"mooncake-store: starting {' '.join(command)} on {host}") + with open(log_path, "wb") as log_file: + process = subprocess.Popen( # nosec B603 + command, env=env, stdout=log_file, stderr=subprocess.STDOUT + ) + master = LaunchedMaster( + process=process, address=f"{host}:{pool.master_port}", log_path=log_path + ) + try: + _wait_until_accepting(host, pool.master_port, _master_timeout(), process=process, hint=hint) + except BaseException: + master.stop() + raise + + logger.info( + f"mooncake-store: master ready at {master.address} " + f"(metrics http://{host}:{pool.master_metrics_port}, log {log_path})" + ) + return master + + +@contextlib.contextmanager +def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optional[str]]: + """Make ``pool`` reachable and name it in this process's environment. + + Yields the path of the client config written, or ``None`` when an inherited + ``MOONCAKE_CONFIG_PATH`` was left in charge. + + Args: + pool: A ``MooncakeStoreConfig``. + run_dir: Where to write the client config and the master's log. + Defaults to ``TRTLLM_MOONCAKE_RUN_DIR``, else a temporary directory + that is removed on exit. + """ + inherited = os.getenv(CONFIG_PATH_ENV) + if inherited: + logger.info( + f"mooncake-store: {CONFIG_PATH_ENV}={inherited} is already set, so " + "kv_connector_config.mooncake_store is ignored and the pool it " + "names is used as is." + ) + yield None + return + + keep_run_dir = bool(run_dir or os.getenv(RUN_DIR_ENV)) + run_dir = run_dir or os.getenv(RUN_DIR_ENV) or tempfile.mkdtemp(prefix="trtllm-mooncake-") + os.makedirs(run_dir, exist_ok=True) + + master: Optional[LaunchedMaster] = None + exported = False + try: + if pool.launch_master: + master = _launch_master(pool, run_dir) + master_address = master.address + else: + master_address = pool.master_server_address + # Reaching a master that is not there fails inside store.setup on + # every rank, after the model has been loaded. Spend a socket now. + if (endpoint := _split_address(master_address)) is None: + logger.warning( + f"mooncake-store: cannot parse master_server_address=" + f"{master_address!r} as host:port, so its reachability is " + "left for the workers to discover." + ) + else: + _wait_until_accepting(*endpoint, _master_timeout()) + logger.info(f"mooncake-store: using the master at {master_address}") + + config_path = os.path.join(run_dir, CLIENT_CONFIG_NAME) + config = _client_config(pool, master_address) + with open(config_path, "w") as handle: + json.dump(config, handle, indent=2) + # The ranks that open store handles are spawned by the LLM constructor, + # inheriting this environment; that is the only reason exporting it + # here reaches them. + os.environ[CONFIG_PATH_ENV] = config_path + exported = True + logger.info( + f"mooncake-store: {CONFIG_PATH_ENV}={config_path} " + f"({json.dumps(config, sort_keys=True)})" + ) + yield config_path + finally: + if exported: + os.environ.pop(CONFIG_PATH_ENV, None) + if master is not None: + master.stop() + if not keep_run_dir: + shutil.rmtree(run_dir, ignore_errors=True) + + +@contextlib.contextmanager +def maybe_provision_pool(kv_connector_config: Any) -> Iterator[None]: + """Provision the pool if this deployment asked the server to. + + A no-op for every other connector, and for a ``mooncake-store`` config + that left ``mooncake_store`` unset: that deployment is told about its pool + through ``MOONCAKE_CONFIG_PATH``, which is how the SLURM harness drives it. + """ + if not uses_connector(kv_connector_config, "mooncake-store"): + yield + return + pool = kv_connector_config.mooncake_store + if pool is None: + yield + return + with provision_pool(pool): + yield diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 2dfe1e91ffe3..fdbf4a479a0a 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1,5 +1,6 @@ import asyncio import atexit +import contextlib import gc import importlib import inspect @@ -14,7 +15,8 @@ import time import uuid from pathlib import Path -from typing import Any, Dict, NamedTuple, Optional, Sequence, Set +from typing import (Any, ContextManager, Dict, NamedTuple, Optional, Sequence, + Set) import click import torch @@ -38,7 +40,8 @@ parse_disagg_config_file, parse_metadata_server_config_file, validate_config_bool) -from tensorrt_llm.llmapi.llm_args import MultimodalConfig, TorchLlmArgs +from tensorrt_llm.llmapi.llm_args import (KvCacheConnectorConfig, + MultimodalConfig, TorchLlmArgs) from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr, split_mpi_env from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory, @@ -514,6 +517,36 @@ def _terminate_attached_frontends(children: list) -> None: child.kill() +def _provision_kv_connector_pool(llm_args: dict, + owns_engine: bool = True) -> ContextManager: + """Bring up the shared cache a KV connector needs, for its lifetime. + + Connectors backed by a cluster-wide pool need it reachable before any rank + opens a handle to it, and the ranks are spawned by the LLM constructor. + Entering this around that construction is what makes the pool part of + `trtllm-serve` bringup rather than something a launch script has to + arrange; a deployment that arranges it anyway is detected and left alone. + + Only the process that owns the engine provisions anything: an attached + frontend re-execs this command line but shares the launcher's executor, + so it would otherwise stand up a second, private pool. + """ + if not owns_engine: + return contextlib.nullcontext() + + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import \ + maybe_provision_pool + + connector_config = llm_args.get("kv_connector_config") + if isinstance(connector_config, dict): + # A YAML config section arrives here unvalidated. Coerce it now, since + # the pool has to be described before the LLM constructor would do it, + # and hand the validated model on so it is not parsed twice. + connector_config = KvCacheConnectorConfig(**connector_config) + llm_args["kv_connector_config"] = connector_config + return maybe_provision_pool(connector_config) + + def launch_server( host: str, port: int, @@ -570,53 +603,55 @@ def launch_server( raise RuntimeError(f"Failed to bind socket to {host}:{port}: {e}. " f"Port holder(s): {holder}") - if backend == 'pytorch': - llm_args.pop("build_config", None) - llm = PyTorchLLM(**llm_args) - elif backend == '_autodeploy': - from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM - - # AutoDeploy does not support build_config - llm_args.pop("build_config", None) - llm = AutoDeployLLM(**llm_args) - else: - raise click.BadParameter( - f"{backend} is not a known backend, check help for available options.", - param_hint="backend") - - # The finally below is the cleanup boundary for the attached - # frontends: it must cover everything from their spawn through - # server construction, middleware registration, and runtime, or a - # failure in between leaks the child processes. - frontend_children = [] - try: - if multi_frontend.is_launcher: - frontend_children = _spawn_attached_frontends( - llm, multi_frontend.num_frontends) - - server = OpenAIServer( - generator=llm, - model=model, - tool_parser=tool_parser, - server_role=server_role, - metadata_server_cfg=metadata_server_cfg, - disagg_cluster_config=disagg_cluster_config, - multimodal_server_config=multimodal_server_config, - chat_template=chat_template, - allow_request_chat_template=allow_request_chat_template, - input_processor_workers=num_input_processor_workers, - media_load_workers=num_media_load_workers) - _apply_fastapi_middlewares(server.app, middleware) - - # Optionally disable GC (default: not disabled) - if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": - gc.disable() - - _signal_frontend_ready(multi_frontend) - uvloop.run(server(host, port, sockets=[s])) - finally: - if frontend_children: - _terminate_attached_frontends(frontend_children) + with _provision_kv_connector_pool( + llm_args, owns_engine=not multi_frontend.is_attached_frontend): + if backend == 'pytorch': + llm_args.pop("build_config", None) + llm = PyTorchLLM(**llm_args) + elif backend == '_autodeploy': + from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM + + # AutoDeploy does not support build_config + llm_args.pop("build_config", None) + llm = AutoDeployLLM(**llm_args) + else: + raise click.BadParameter( + f"{backend} is not a known backend, check help for available options.", + param_hint="backend") + + # The finally below is the cleanup boundary for the attached + # frontends: it must cover everything from their spawn through + # server construction, middleware registration, and runtime, or a + # failure in between leaks the child processes. + frontend_children = [] + try: + if multi_frontend.is_launcher: + frontend_children = _spawn_attached_frontends( + llm, multi_frontend.num_frontends) + + server = OpenAIServer( + generator=llm, + model=model, + tool_parser=tool_parser, + server_role=server_role, + metadata_server_cfg=metadata_server_cfg, + disagg_cluster_config=disagg_cluster_config, + multimodal_server_config=multimodal_server_config, + chat_template=chat_template, + allow_request_chat_template=allow_request_chat_template, + input_processor_workers=num_input_processor_workers, + media_load_workers=num_media_load_workers) + _apply_fastapi_middlewares(server.app, middleware) + + # Optionally disable GC (default: not disabled) + if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": + gc.disable() + + _signal_frontend_ready(multi_frontend) + uvloop.run(server(host, port, sockets=[s])) + finally: + if frontend_children: + _terminate_attached_frontends(frontend_children) def launch_grpc_server(host: str, @@ -739,7 +774,8 @@ def signal_handler(): logger.info("Shutdown complete") - uvloop.run(serve_grpc_async()) + with _provision_kv_connector_pool(llm_args): + uvloop.run(serve_grpc_async()) def launch_mm_encoder_server( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 76693db6c02c..5ae7fc40a182 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1945,6 +1945,97 @@ def num_capture_layers(self) -> int: return 0 +class MooncakeStoreConfig(StrictBaseModel): + """The Mooncake store pool the `mooncake-store` connector should join. + + Describes the *pool*: which master owns it, how workers reach it, and how + much memory each contributes. A worker's own relationship to the pool + (read/write role, key namespace, model identity) stays in the + ``TRTLLM_MOONCAKE_STORE_*`` environment variables, because it is per + process while this is per deployment. + + Setting this makes `trtllm-serve` render the Mooncake client config and + export ``MOONCAKE_CONFIG_PATH`` itself, so no external script has to. + An inherited ``MOONCAKE_CONFIG_PATH`` still wins, which is how the SLURM + harness keeps pointing workers at a pool it manages. + + Every field opts out of telemetry: they describe one site's pool -- its + ports, its fabric, how much memory it was given -- rather than which + features are in use, which ``kv_connector_config.connector`` already says. + """ + master_server_address: Optional[str] = Field( + None, + description="Address (host:port) of an already-running " + "mooncake_master. Mutually exclusive with launch_master.") + launch_master: bool = Field( + False, + telemetry=False, + description="Start a mooncake_master in this server's process group " + "and use it. The pool then dies with the server, so this is only " + "correct for a single engine: several engines sharing a pool, or a " + "pool that must outlive a restart, need master_server_address.") + master_port: int = Field( + 50051, + telemetry=False, + description="RPC port for a master started by launch_master.") + master_metrics_port: int = Field( + 9004, + telemetry=False, + description="Prometheus port for a master started by launch_master.") + master_eviction_ratio: float = Field( + 0.05, + telemetry=False, + description="Fraction of the pool a master started by launch_master " + "frees per eviction pass.") + metadata_server: str = Field( + "P2PHANDSHAKE", + description="Mooncake metadata service. P2PHANDSHAKE keeps a separate " + "metadata process out of the deployment.") + protocol: str = Field( + "rdma", + description="Transport for page traffic: 'rdma' or 'tcp'. " + "TCP is for bring-up only; it invalidates performance conclusions.") + device_name: str = Field( + "", + description="RDMA device to transfer over, from ibv_devinfo. Empty " + "with protocol 'tcp'.") + global_segment_size: Union[int, str] = Field( + "16GiB", + description="Host memory each worker process contributes to the pool. " + "Pool capacity is this times the number of processes that open a " + "store handle, so a prefill-only connector gives a prefill-only pool.") + local_buffer_size: Union[int, str] = Field( + "1GiB", + description="Per-process Mooncake transfer buffer, not pool capacity.") + transfer_batch_size: int = Field(64, + telemetry=False, + description="Page keys per store call.") + cache_prefix: Optional[str] = Field( + None, + description="Key namespace for the pool. Bump it after any change to " + "page layout or contents. Defaults to 'trtllm'.") + stage_through_host: bool = Field( + False, + telemetry=False, + description="Copy pages through a pinned host buffer instead of " + "registering the KV pools with Mooncake. Needed where the HCA cannot " + "pin GPU pages (no GPUDirect RDMA); costs a copy each way.") + + @model_validator(mode="after") + def _require_exactly_one_master(self) -> "MooncakeStoreConfig": + if self.launch_master and self.master_server_address: + raise ValueError( + "mooncake_store: set either launch_master or " + "master_server_address, not both. launch_master starts a " + "master here; master_server_address joins an existing pool.") + if not self.launch_master and not self.master_server_address: + raise ValueError( + "mooncake_store: needs a master. Set master_server_address to " + "join an existing pool, or launch_master: true to start one " + "for this server alone.") + return self + + class KvCacheConnectorConfig(StrictBaseModel): """Configuration for the KV Cache Connector. @@ -1975,11 +2066,16 @@ class KvCacheConnectorConfig(StrictBaseModel): description="URL for an external connector server " "(e.g. 'tcp://localhost:5555'). Connectors that run in " "multi-process mode use this to reach the cache server.") + mooncake_store: Optional[MooncakeStoreConfig] = Field( + None, + description="Pool topology for the 'mooncake-store' connector. When " + "set, trtllm-serve provisions the pool during bringup instead of " + "requiring MOONCAKE_CONFIG_PATH from an external script.") @model_validator(mode="after") def _resolve_preset(self) -> "KvCacheConnectorConfig": - from tensorrt_llm._torch.pyexecutor.connectors.registry import \ - CONNECTOR_REGISTRY + from tensorrt_llm._torch.pyexecutor.connectors.registry import ( + CONNECTOR_REGISTRY, uses_connector) if self.connector is not None: preset = CONNECTOR_REGISTRY.get(self.connector) if preset is None: @@ -1997,6 +2093,12 @@ def _resolve_preset(self) -> "KvCacheConnectorConfig": raise ValueError("connector_scheduler_class is required") if self.connector_worker_class is None: raise ValueError("connector_worker_class is required") + if self.mooncake_store is not None and not uses_connector( + self, "mooncake-store"): + raise ValueError( + "mooncake_store describes a Mooncake pool, but this config " + f"resolves to connector_module={self.connector_module!r}. " + "Set connector: mooncake-store, or drop mooncake_store.") return self diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index db8d31642d2c..b807f3e46f42 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -42,6 +42,7 @@ l0_a10: - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_kv_cache_layout.py - unittest/_torch/executor/test_mooncake_store_connector.py + - unittest/_torch/executor/test_mooncake_store_master.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py - unittest/_torch/modules/dwdp/test_dwdp_manager.py diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py new file mode 100644 index 000000000000..c1b0c1e244f1 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -0,0 +1,404 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for provisioning a Mooncake store pool during server bringup. + +Runs without a Mooncake installation and without a GPU. A master this process +launches is a fake standing in for ``Popen`` that opens the RPC port, which is +all the readiness handshake ever observes; a master someone else runs is a +plain socket. +""" + +import json +import os +import shutil +import socket +import subprocess +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import master as master_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + CONFIG_PATH_ENV, + MooncakeStoreConnectorConfig, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.master import ( + maybe_provision_pool, + provision_pool, +) +from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig, MooncakeStoreConfig + + +def free_port() -> int: + with socket.socket() as probe: + probe.bind(("", 0)) + return probe.getsockname()[1] + + +class FakeMasterProcess: + """The slice of ``Popen`` that launching a master actually drives. + + ``listen_on`` makes it answer on that port, which is what a real master + does last and what the readiness wait keys off. ``exit_code`` makes it a + master that failed to start. + """ + + def __init__(self, command, env, listen_on=None, exit_code=None): + self.command = command + self.env = env + self.terminated = False + self.killed = False + self._exit_code = exit_code + self._listener = None + if listen_on is not None: + self._listener = socket.socket() + self._listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self._listener.bind(("", listen_on)) + self._listener.listen(8) + + def poll(self): + return self._exit_code + + def terminate(self): + self.terminated = True + self._exit_code = -15 + if self._listener is not None: + self._listener.close() + self._listener = None + + def wait(self, timeout=None): + return self._exit_code + + def kill(self): + self.killed = True + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + """No ambient pool: these tests are about what provisioning does itself.""" + for name in ( + CONFIG_PATH_ENV, + master_module.MASTER_BINARY_ENV, + master_module.MASTER_TIMEOUT_ENV, + master_module.RUN_DIR_ENV, + ): + monkeypatch.delenv(name, raising=False) + # Nothing here is slow to start, so a wait that runs long is a failure + # rather than something that needs more time. + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "10") + + +@pytest.fixture +def fake_master(monkeypatch): + """Replace the master binary and its process with in-process fakes. + + Returns a callable that arms the fake and, once provisioning has run, the + launched instance is available as ``.process`` for inspection. + """ + + class Launcher: + def __init__(self): + self.process = None + + def arm(self, listen_on=None, exit_code=None): + + def popen(command, env=None, **_kwargs): + self.process = FakeMasterProcess( + command, env, listen_on=listen_on, exit_code=exit_code + ) + return self.process + + # Swap the modules as this module sees them rather than patching + # attributes on the shared stdlib ones. + monkeypatch.setattr( + master_module, + "shutil", + SimpleNamespace(which=lambda name: f"/opt/bin/{name}", rmtree=shutil.rmtree), + ) + monkeypatch.setattr( + master_module, + "subprocess", + SimpleNamespace( + Popen=popen, + STDOUT=subprocess.STDOUT, + TimeoutExpired=subprocess.TimeoutExpired, + ), + ) + + return Launcher() + + +@pytest.fixture +def running_master(): + """A socket standing in for a master someone else is running.""" + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(8) + try: + yield f"127.0.0.1:{listener.getsockname()[1]}" + finally: + listener.close() + + +# ---- configuration ---- + + +def test_pool_needs_exactly_one_master(): + with pytest.raises(ValueError, match="not both"): + MooncakeStoreConfig(launch_master=True, master_server_address="host:50051") + with pytest.raises(ValueError, match="needs a master"): + MooncakeStoreConfig() + + +def test_pool_is_rejected_on_another_connector(): + with pytest.raises(ValueError, match="mooncake_store describes a Mooncake pool"): + KvCacheConnectorConfig( + connector="lmcache", + mooncake_store=MooncakeStoreConfig(launch_master=True), + ) + + +def test_pool_is_accepted_on_the_module_spelled_out(): + """A config naming the module instead of the preset is still the connector.""" + config = KvCacheConnectorConfig( + connector_module="tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + connector_scheduler_class="MooncakeStoreConnectorScheduler", + connector_worker_class="MooncakeStoreConnectorWorker", + mooncake_store=MooncakeStoreConfig(launch_master=True), + ) + assert config.mooncake_store.launch_master + + +# ---- the rendered client config ---- + + +def test_client_config_is_what_the_connector_reads_back(tmp_path): + """The generated JSON has to survive the connector's own parser.""" + pool = MooncakeStoreConfig( + master_server_address="10.0.0.1:50051", + protocol="rdma", + device_name="mlx5_0", + global_segment_size="64GiB", + local_buffer_size="4GiB", + cache_prefix="trtllm-m3", + stage_through_host=True, + transfer_batch_size=32, + ) + path = tmp_path / "mooncake.json" + path.write_text(json.dumps(master_module._client_config(pool, "10.0.0.1:50051"))) + + parsed = MooncakeStoreConnectorConfig.from_file(str(path)) + assert parsed.master_server_address == "10.0.0.1:50051" + assert parsed.metadata_server == "P2PHANDSHAKE" + assert parsed.protocol == "rdma" + assert parsed.device_name == "mlx5_0" + assert parsed.global_segment_size == 64 * 1024**3 + assert parsed.local_buffer_size == 4 * 1024**3 + assert parsed.cache_prefix == "trtllm-m3" + assert parsed.stage_through_host is True + assert parsed.transfer_batch_size == 32 + + +def test_client_config_leaves_an_unset_prefix_to_the_connector(): + pool = MooncakeStoreConfig(master_server_address="host:50051") + assert "cache_prefix" not in master_module._client_config(pool, "host:50051") + + +@pytest.mark.parametrize( + "address, expected", + [ + ("host:50051", ("host", 50051)), + ("[::1]:50051", ("::1", 50051)), + ("unix:///var/run/mooncake", None), + ("host", None), + ], +) +def test_master_addresses_are_split_or_declined(address, expected): + assert master_module._split_address(address) == expected + + +# ---- provisioning against a master someone else runs ---- + + +def test_provisioning_points_the_workers_at_a_running_master(running_master): + pool = MooncakeStoreConfig(master_server_address=running_master) + + with provision_pool(pool) as config_path: + # The workers are spawned inside this window and are told about the + # pool through the environment, so both have to hold while it is open. + assert os.environ[CONFIG_PATH_ENV] == config_path + written = json.loads(open(config_path).read()) + assert written["master_server_address"] == running_master + + assert CONFIG_PATH_ENV not in os.environ + assert not os.path.exists(config_path) + + +def test_provisioning_fails_before_the_model_loads_if_the_master_is_absent(monkeypatch): + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") + pool = MooncakeStoreConfig(master_server_address=f"127.0.0.1:{free_port()}") + + with pytest.raises(TimeoutError, match="did not accept connections"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + assert CONFIG_PATH_ENV not in os.environ + + +def test_an_unparseable_master_address_is_left_to_the_workers(): + """Not every address is host:port; that is the worker's problem, not ours.""" + pool = MooncakeStoreConfig(master_server_address="unix:///var/run/mooncake") + + with provision_pool(pool) as config_path: + written = json.loads(open(config_path).read()) + assert written["master_server_address"] == "unix:///var/run/mooncake" + + +def test_an_inherited_config_path_wins(monkeypatch, tmp_path): + """The SLURM harness names the pool this way; provisioning must defer.""" + harness_config = tmp_path / "harness.json" + harness_config.write_text("{}") + monkeypatch.setenv(CONFIG_PATH_ENV, str(harness_config)) + pool = MooncakeStoreConfig(launch_master=True) + + with provision_pool(pool) as config_path: + assert config_path is None + assert os.environ[CONFIG_PATH_ENV] == str(harness_config) + + assert os.environ[CONFIG_PATH_ENV] == str(harness_config) + + +# ---- provisioning with a master of our own ---- + + +def test_a_launched_master_is_named_in_the_config_and_stopped_on_exit(fake_master): + port = free_port() + fake_master.arm(listen_on=port) + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with provision_pool(pool) as config_path: + written = json.loads(open(config_path).read()) + host, _, named_port = written["master_server_address"].rpartition(":") + assert int(named_port) == port + # Whatever the config names has to be dialable: it is all a worker on + # another host is given. + with socket.create_connection((host, port), timeout=5): + pass + + assert fake_master.process.terminated + assert not fake_master.process.killed + + +def test_a_launched_master_gets_the_flags_and_logging_it_needs(fake_master): + port = free_port() + fake_master.arm(listen_on=port) + pool = MooncakeStoreConfig( + launch_master=True, + master_port=port, + master_metrics_port=free_port(), + master_eviction_ratio=0.1, + ) + + with provision_pool(pool): + command = fake_master.process.command + assert command[0].endswith("mooncake_master") + assert f"--rpc_port={port}" in command + assert f"--metrics_port={pool.master_metrics_port}" in command + assert "--eviction_ratio=0.1" in command + # glog writes to files under /tmp unless told otherwise, which would + # leave the master's log -- the only view of the pool's own side of + # the conversation -- empty. + assert fake_master.process.env["GLOG_logtostderr"] == "1" + assert fake_master.process.env["GLOG_v"] == "1" + + +def test_a_master_that_dies_during_startup_says_so(fake_master): + fake_master.arm(exit_code=3) + pool = MooncakeStoreConfig(launch_master=True, master_port=free_port()) + + with pytest.raises(RuntimeError, match="exited with code 3"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + assert CONFIG_PATH_ENV not in os.environ + + +def test_a_master_that_never_listens_times_out(monkeypatch, fake_master): + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") + fake_master.arm() + pool = MooncakeStoreConfig(launch_master=True, master_port=free_port()) + + with pytest.raises(TimeoutError, match="did not accept connections"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + assert fake_master.process.terminated + + +def test_a_missing_master_binary_names_the_alternatives(monkeypatch): + monkeypatch.setattr( + master_module, + "shutil", + SimpleNamespace(which=lambda _name: None, rmtree=shutil.rmtree), + ) + pool = MooncakeStoreConfig(launch_master=True) + + with pytest.raises(FileNotFoundError, match="master_server_address"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + + +def test_a_run_dir_keeps_the_master_log_and_the_config(fake_master, tmp_path): + port = free_port() + fake_master.arm(listen_on=port) + run_dir = tmp_path / "pool" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with provision_pool(pool, run_dir=str(run_dir)) as config_path: + assert config_path == str(run_dir / master_module.CLIENT_CONFIG_NAME) + + # An explicit run directory has to outlive the run that filled it: the + # master's log is where pool occupancy and eviction are read from. + assert (run_dir / master_module.MASTER_LOG_NAME).exists() + assert (run_dir / master_module.CLIENT_CONFIG_NAME).exists() + + +# ---- the entry point servers call ---- + + +def test_other_connectors_are_left_alone(): + config = KvCacheConnectorConfig(connector="lmcache") + with maybe_provision_pool(config): + assert CONFIG_PATH_ENV not in os.environ + + +def test_no_connector_at_all_is_left_alone(): + with maybe_provision_pool(None): + assert CONFIG_PATH_ENV not in os.environ + + +def test_a_pool_left_undescribed_stays_the_environment_contract(): + """Without ``mooncake_store``, MOONCAKE_CONFIG_PATH is still the only input.""" + config = KvCacheConnectorConfig(connector="mooncake-store") + with maybe_provision_pool(config): + assert CONFIG_PATH_ENV not in os.environ + + +def test_a_described_pool_is_provisioned(running_master): + config = KvCacheConnectorConfig( + connector="mooncake-store", + mooncake_store=MooncakeStoreConfig(master_server_address=running_master), + ) + with maybe_provision_pool(config): + written = json.loads(open(os.environ[CONFIG_PATH_ENV]).read()) + assert written["master_server_address"] == running_master + assert CONFIG_PATH_ENV not in os.environ From d16989ddd002d0db65eed6e3f01539ca59035b32 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:27:35 -0700 Subject: [PATCH 13/24] [None][fix] Let the config size the Mooncake staging buffer MooncakeStoreConfig could turn stage_through_host on but not size the buffer it stages through, leaving it at the connector's 512MiB default. A buffer that cannot hold transfer_batch_size pages reduces the batch rather than failing, so the ceiling that setting implies was reachable only by writing MOONCAKE_CONFIG_PATH by hand -- which is the thing describing the pool in the config is meant to replace. The field is left out of the rendered config when unset so the connector's default stays the one definition of it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/master.py | 4 ++++ tensorrt_llm/llmapi/llm_args.py | 8 +++++++ .../executor/test_mooncake_store_master.py | 21 +++++++++++++++++++ 3 files changed, 33 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index 16f91d6c32ee..2ebaecb8e73b 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -152,6 +152,10 @@ def _client_config(pool: Any, master_address: str) -> Dict[str, Any]: } if pool.cache_prefix is not None: config["cache_prefix"] = pool.cache_prefix + # Left out when unset so the connector's own default applies, rather than + # restating it here for the two to drift apart. + if pool.staging_buffer_bytes is not None: + config["staging_buffer_bytes"] = pool.staging_buffer_bytes return config diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5ae7fc40a182..b639a2a718f7 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2020,6 +2020,14 @@ class MooncakeStoreConfig(StrictBaseModel): description="Copy pages through a pinned host buffer instead of " "registering the KV pools with Mooncake. Needed where the HCA cannot " "pin GPU pages (no GPUDirect RDMA); costs a copy each way.") + staging_buffer_bytes: Optional[Union[int, str]] = Field( + None, + telemetry=False, + description="Size of the buffer stage_through_host copies through, " + "per process. Pages move transfer_batch_size at a time, and a buffer " + "that cannot hold that many reduces the batch instead of failing, so " + "undersizing it costs throughput quietly. Defaults to the connector's " + "own 512MiB.") @model_validator(mode="after") def _require_exactly_one_master(self) -> "MooncakeStoreConfig": diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py index c1b0c1e244f1..d3aef17774ce 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_master.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -246,6 +246,27 @@ def test_provisioning_points_the_workers_at_a_running_master(running_master): assert not os.path.exists(config_path) +def test_a_staging_buffer_can_be_sized_where_staging_is_turned_on(running_master): + """Undersizing it silently shrinks the transfer batch, so it must be settable.""" + pool = MooncakeStoreConfig( + master_server_address=running_master, + stage_through_host=True, + staging_buffer_bytes="4GiB", + ) + + with provision_pool(pool) as config_path: + written = json.loads(open(config_path).read()) + assert written["stage_through_host"] is True + assert written["staging_buffer_bytes"] == "4GiB" + + +def test_an_unsized_staging_buffer_leaves_the_connector_its_default(running_master): + pool = MooncakeStoreConfig(master_server_address=running_master) + + with provision_pool(pool) as config_path: + assert "staging_buffer_bytes" not in json.loads(open(config_path).read()) + + def test_provisioning_fails_before_the_model_loads_if_the_master_is_absent(monkeypatch): monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") pool = MooncakeStoreConfig(master_server_address=f"127.0.0.1:{free_port()}") From a3404ea34bd065b6bfdf57407f66e3af28493c40 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:55:10 -0700 Subject: [PATCH 14/24] [None][feat] Ship the Mooncake pool parts a server cannot own A server that owns its pool provisions it from its own config. The pools it cannot own still needed a script of someone else's: one shared by several engines, or outliving a restart, needs a master that is not any of them, and a pool whose capacity should include nodes that run no connector needs those nodes to hold segments. Both were scripts under mooncake_disagg, so those deployments were assembled from a benchmark harness rather than from what TensorRT-LLM ships. 'trtllm-serve mooncake_master' runs a master for as long as it runs and publishes where it landed; 'trtllm-serve mooncake_donor' lends a node's memory while leaving that engine connector-free. Donation stays out of StoreRole on purpose -- the roles describe an engine's traffic, and contributing memory is capacity, so making it a role would start a generation server reading or writing the store to get its DRAM in. master_server_address also accepts file://, which is what makes the master reachable without anyone writing its address down: its host is whatever the scheduler chose, so a config settled beforehand cannot name it, and the wait to read it is also the wait for it to exist. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/__init__.py | 18 +- .../connectors/mooncake_store/config.py | 8 +- .../connectors/mooncake_store/donor.py | 121 ++++++++++ .../connectors/mooncake_store/master.py | 97 +++++++- tensorrt_llm/commands/mooncake.py | 222 ++++++++++++++++++ tensorrt_llm/commands/serve.py | 8 +- tensorrt_llm/llmapi/llm_args.py | 7 +- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../executor/test_mooncake_store_donor.py | 133 +++++++++++ .../executor/test_mooncake_store_master.py | 105 +++++++++ 10 files changed, 704 insertions(+), 16 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py create mode 100644 tensorrt_llm/commands/mooncake.py create mode 100644 tests/unittest/_torch/executor/test_mooncake_store_donor.py diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index a07b7a6aace0..0a1b7ec5ff58 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -39,6 +39,11 @@ master and writing that JSON itself -- so no external script has to. See ``master.py``. +Capacity, separately, comes only from processes that open a store handle, which +in a disaggregated deployment is the context servers alone. ``donor.py`` lends +a node's memory to the pool without giving it a connector, so the generation +nodes can hold cache they never read. + By default the KV pools themselves are registered with Mooncake, so the store reads and writes device memory and no copy is added. That needs the HCA to be able to pin GPU pages -- GPUDirect RDMA, through ``nvidia_peermem`` or dma-buf. @@ -49,16 +54,25 @@ bytes are the same either way, so the two modes can share a pool. """ -from .config import MooncakeStoreConnectorConfig, StoreRole -from .master import maybe_provision_pool, provision_pool +from .config import MooncakeStoreConnectorConfig, StoreRole, parse_size +from .donor import DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment +from .master import (local_address, master_timeout, maybe_provision_pool, + provision_pool, resolve_master_address, running_master) from .scheduler import MooncakeStoreConnectorScheduler from .worker import MooncakeStoreConnectorWorker __all__ = [ + "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", "MooncakeStoreConnectorConfig", "MooncakeStoreConnectorScheduler", "MooncakeStoreConnectorWorker", "StoreRole", + "donate_segment", + "local_address", + "master_timeout", "maybe_provision_pool", + "parse_size", "provision_pool", + "resolve_master_address", + "running_master", ] diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 948f47e7fa60..09ca8cc0bf3e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -94,7 +94,7 @@ def saves(self) -> bool: return self is not StoreRole.CONSUMER -def _parse_size(value: Any) -> int: +def parse_size(value: Any) -> int: """Accept either a byte count or a suffixed string such as ``"4GiB"``.""" if isinstance(value, bool): raise ValueError(f"expected a size, got {value!r}") @@ -167,10 +167,10 @@ def from_file(path: str) -> "MooncakeStoreConnectorConfig": master_server_address=raw.get("master_server_address", ""), protocol=raw.get("protocol", "rdma"), device_name=raw.get("device_name", ""), - global_segment_size=_parse_size( + global_segment_size=parse_size( raw.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE) ), - local_buffer_size=_parse_size(raw.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)), + local_buffer_size=parse_size(raw.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)), local_hostname=raw.get("local_hostname") or None, tenant_id=raw.get("tenant_id") or None, role=StoreRole(str(raw.get("role", StoreRole.BOTH.value)).strip().lower()), @@ -178,7 +178,7 @@ def from_file(path: str) -> "MooncakeStoreConnectorConfig": model_key=raw.get("model_key") or None, transfer_batch_size=int(raw.get("transfer_batch_size", 64)), stage_through_host=bool(raw.get("stage_through_host", False)), - staging_buffer_bytes=_parse_size( + staging_buffer_bytes=parse_size( raw.get("staging_buffer_bytes", DEFAULT_STAGING_BUFFER_SIZE) ), ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py new file mode 100644 index 000000000000..dbb105086cfd --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Put a node's host memory into a Mooncake pool without reading or writing it. + +Pool capacity comes only from processes that open a store handle: ``setup`` +registers ``global_segment_size`` bytes of the caller's host memory and the +master then places blocks in it. In a disaggregated deployment only the context +servers configure the connector, so only they call ``setup``, and the pool is +entirely prefill-node memory -- which makes the store a +prefill-DRAM-caches-prefill-GPU tier, overlapping what TensorRT-LLM's own host +offload already does. + +Donating alongside a generation server puts that node's memory into the same +pool. Prefill then writes blocks that land on decode-side DRAM and reads them +back, while the generation engine stays free of any connector: it neither reads +nor writes the store, so it keeps its single cache transceiver for the +prefill-to-decode handoff. + +Donation is deliberately not a ``StoreRole``. The roles describe an engine's +traffic -- ``producer`` writes, ``consumer`` reads, ``both`` does both -- and +none of them means "contribute memory only", so attaching a connector to a +generation server to get its DRAM into the pool would also start it reading or +writing. Capacity and traffic are separate concerns, which is why this holds a +handle of its own rather than being a setting on the connector. + +The memory is charged to the donating process, so it competes with anything +else on the node -- a generation server's ``kv_cache_config.host_cache_size`` +above all. Size the two together. +""" + +import contextlib +from typing import Iterator, Optional + +from tensorrt_llm.logger import logger + +from .master import local_address + +__all__ = ["DEFAULT_DONOR_LOCAL_BUFFER_SIZE", "donate_segment"] + +#: A donor never transfers, so its transfer buffer is dead weight; ``setup`` +#: still rejects a zero one. +DEFAULT_DONOR_LOCAL_BUFFER_SIZE = 64 * 1024**2 + + +@contextlib.contextmanager +def donate_segment( + master_server_address: str, + segment_size: int, + protocol: str = "rdma", + device_name: str = "", + metadata_server: str = "", + local_buffer_size: int = DEFAULT_DONOR_LOCAL_BUFFER_SIZE, + hostname: Optional[str] = None, +) -> Iterator[str]: + """Hold ``segment_size`` bytes of this node's memory in the pool. + + Yields the host the segment is registered under, which is what the master + and the engines reading from it identify the capacity by. + + The handle is held for the duration: dropping it unmounts the segment, and + the master starts reporting the blocks that lived in it as lost. So the + caller must stay inside this context for as long as the capacity is meant + to exist, which for a donor is its whole run. + """ + try: + from mooncake.store import MooncakeDistributedStore + except ImportError as exc: + raise ImportError( + "Donating memory needs the Mooncake Python bindings " + "(`pip install mooncake-transfer-engine`). The C++ transfer engine " + f"in the container is a different component: {exc}" + ) from exc + + host = hostname or local_address() + logger.info( + f"mooncake-store: joining the pool at {master_server_address} as " + f"capacity only: host={host} protocol={protocol} " + f"device={device_name or '(none)'} " + f"donating={segment_size / 1024 ** 3:.1f}GiB" + ) + + store = MooncakeDistributedStore() + status = store.setup( + host, + metadata_server, + segment_size, + local_buffer_size, + protocol, + device_name, + master_server_address, + ) + if status != 0: + raise RuntimeError( + f"Mooncake store.setup failed with status {status}. The master at " + f"{master_server_address} must already be accepting connections, " + f"and protocol={protocol!r} must be usable from {host}." + ) + + logger.info( + f"mooncake-store: {segment_size / 1024 ** 3:.1f}GiB from {host} is now " + "part of the pool" + ) + try: + yield host + finally: + # Explicit because the segment stays mounted for as long as anything + # references the handle, and "as long as this context" is the contract. + del store + logger.info(f"mooncake-store: withdrew the segment donated from {host}") diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index 2ebaecb8e73b..fe244d818d02 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -48,7 +48,14 @@ from ..registry import uses_connector from .config import CONFIG_PATH_ENV -__all__ = ["maybe_provision_pool", "provision_pool"] +__all__ = [ + "local_address", + "maybe_provision_pool", + "master_timeout", + "provision_pool", + "resolve_master_address", + "running_master", +] #: Override the binary that ``launch_master`` runs. MASTER_BINARY_ENV = "TRTLLM_MOONCAKE_MASTER_BINARY" @@ -62,9 +69,12 @@ DEFAULT_MASTER_TIMEOUT = 60.0 CLIENT_CONFIG_NAME = "mooncake.json" MASTER_LOG_NAME = "mooncake_master.log" +#: Prefix that makes ``master_server_address`` name a file holding the address +#: rather than the address itself. +ADDRESS_FILE_SCHEME = "file://" -def _local_address() -> str: +def local_address() -> str: """The address this host is known by inside the pool. Deliberately the same derivation the connector worker uses for its own @@ -77,7 +87,7 @@ def _local_address() -> str: return "127.0.0.1" -def _master_timeout() -> float: +def master_timeout() -> float: raw = os.getenv(MASTER_TIMEOUT_ENV) if not raw: return DEFAULT_MASTER_TIMEOUT @@ -98,6 +108,42 @@ def _split_address(address: str) -> Optional[Tuple[str, int]]: return host.strip("[]"), int(port) +def resolve_master_address(address: str, timeout: float) -> str: + """Read a ``file://`` address through, and pass anything else along. + + A master with its own lifetime is on whichever host its scheduler gave + it, which is not known when the worker configs are written. Naming the + file it publishes to instead keeps the address out of the config and out + of a launch script: ``trtllm-serve mooncake_master --address-file`` writes + it, every worker's ``master_server_address`` names the same path, and the + wait here is also the wait for the master to exist at all. + """ + if not address.startswith(ADDRESS_FILE_SCHEME): + return address + + path = address[len(ADDRESS_FILE_SCHEME):] + deadline = time.monotonic() + timeout + while True: + # Written whole by the master command, so a non-empty file is a + # complete address rather than a prefix of one. + try: + published = open(path).read().strip() + except FileNotFoundError: + published = "" + if published: + logger.info(f"mooncake-store: {path} names the master at {published}") + return published + if time.monotonic() >= deadline: + raise TimeoutError( + f"No Mooncake master address appeared in {path} within " + f"{timeout:g}s. Start one with 'trtllm-serve mooncake_master " + f"--address-file {path}', or name a reachable host:port in " + f"master_server_address. Raise {MASTER_TIMEOUT_ENV} if the " + "master is only slow to start." + ) + time.sleep(0.5) + + def _wait_until_accepting( host: str, port: int, @@ -191,7 +237,7 @@ def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: "set master_server_address to a master you run yourself." ) - host = _local_address() + host = local_address() log_path = os.path.join(run_dir, MASTER_LOG_NAME) hint = f" See {log_path}." @@ -218,7 +264,7 @@ def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: process=process, address=f"{host}:{pool.master_port}", log_path=log_path ) try: - _wait_until_accepting(host, pool.master_port, _master_timeout(), process=process, hint=hint) + _wait_until_accepting(host, pool.master_port, master_timeout(), process=process, hint=hint) except BaseException: master.stop() raise @@ -230,6 +276,41 @@ def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: return master +@contextlib.contextmanager +def running_master( + pool: Any, run_dir: str, address_file: Optional[str] = None +) -> Iterator[LaunchedMaster]: + """Run a master whose lifetime is this process's rather than an engine's. + + ``provision_pool`` covers the server that owns its pool. Everything else -- + several engines on one pool, a pool that has to survive a restart -- needs + the master somewhere that is not any of them, which is what this is for. + + ``address_file`` receives ``host:port`` once the master answers, so the + workers can name the file instead of an address nobody knows until the + scheduler has placed this process. + """ + os.makedirs(run_dir, exist_ok=True) + master = _launch_master(pool, run_dir) + try: + if address_file: + # Renamed into place so a reader sees either nothing or the whole + # address. A half-written one would be dialed as if it were real. + staging = f"{address_file}.partial" + with open(staging, "w") as handle: + handle.write(f"{master.address}\n") + os.replace(staging, address_file) + logger.info(f"mooncake-store: published {master.address} to {address_file}") + yield master + finally: + if address_file: + # The address outliving the master would send the next run's + # workers to a port with nothing behind it. + with contextlib.suppress(OSError): + os.remove(address_file) + master.stop() + + @contextlib.contextmanager def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optional[str]]: """Make ``pool`` reachable and name it in this process's environment. @@ -264,7 +345,9 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona master = _launch_master(pool, run_dir) master_address = master.address else: - master_address = pool.master_server_address + master_address = resolve_master_address( + pool.master_server_address, master_timeout() + ) # Reaching a master that is not there fails inside store.setup on # every rank, after the model has been loaded. Spend a socket now. if (endpoint := _split_address(master_address)) is None: @@ -274,7 +357,7 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona "left for the workers to discover." ) else: - _wait_until_accepting(*endpoint, _master_timeout()) + _wait_until_accepting(*endpoint, master_timeout()) logger.info(f"mooncake-store: using the master at {master_address}") config_path = os.path.join(run_dir, CLIENT_CONFIG_NAME) diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py new file mode 100644 index 000000000000..814832535c7f --- /dev/null +++ b/tensorrt_llm/commands/mooncake.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The two pieces of a Mooncake pool that outlive any one engine. + +A server that owns its pool needs neither of these: it describes the pool in +`kv_connector_config.mooncake_store` and `trtllm-serve` provisions it during +bringup. They exist for the pools it cannot own -- one shared by several +engines, one that has to survive a restart, one whose capacity has to come from +nodes that run no connector -- so that those deployments are still assembled +from things TensorRT-LLM ships rather than from a launch script of one's own. +""" + +import json +import os +import signal +import tempfile +import threading +import time +from typing import Optional + +import click + +from tensorrt_llm.logger import logger + + +def _until_signalled() -> threading.Event: + """An event that SIGINT and SIGTERM set. + + Both commands hold a resource -- a child process, a mounted segment -- + whose release is in a ``finally``. Default SIGTERM handling would skip it, + leaving the master unreaped or the pool advertising memory that has gone. + """ + stopping = threading.Event() + + def stop(signum, _frame): + logger.info(f"mooncake-store: signal {signum} received, shutting down") + stopping.set() + + for received in (signal.SIGINT, signal.SIGTERM): + signal.signal(received, stop) + return stopping + + +@click.command("mooncake_master") +@click.option("--rpc_port", + type=int, + default=50051, + show_default=True, + help="Port the store clients reach the master on.") +@click.option("--metrics_port", + type=int, + default=9004, + show_default=True, + help="Prometheus port. Pool occupancy and eviction are read " + "from here or from the master's log.") +@click.option("--eviction_ratio", + type=float, + default=0.05, + show_default=True, + help="Fraction of the pool freed per eviction pass.") +@click.option("--address_file", + type=str, + default=None, + help="File to publish 'host:port' to once the master answers. " + "Workers name it as master_server_address: file://, which " + "is how they reach a master whose host the scheduler chose. " + "Removed on exit so a stale address is never dialed.") +@click.option("--run_dir", + type=str, + default=None, + help="Where to keep the master's log. Defaults to " + "$TRTLLM_MOONCAKE_RUN_DIR, else a temporary directory.") +def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, + address_file: Optional[str], run_dir: Optional[str]): + """Run a mooncake_master for as long as this command runs. + + For a pool that must not belong to any one engine: several servers sharing + it, or one that has to still be there after a server restarts. A single + server with a pool of its own should set `mooncake_store.launch_master` + instead and skip this entirely. + """ + # Imported here rather than at module scope so that reaching any other + # subcommand, or --help, does not pay for the connector package. + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import \ + running_master + from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig + + pool = MooncakeStoreConfig( + launch_master=True, + master_port=rpc_port, + master_metrics_port=metrics_port, + master_eviction_ratio=eviction_ratio, + ) + run_dir = run_dir or os.getenv( + "TRTLLM_MOONCAKE_RUN_DIR") or tempfile.mkdtemp( + prefix="trtllm-mooncake-master-") + + stopping = _until_signalled() + with running_master(pool, run_dir, address_file=address_file) as master: + while not stopping.is_set(): + if (code := master.process.poll()) is not None: + # Its own death is the interesting outcome: the pool is gone + # and every client is about to start failing. + raise click.ClickException( + f"mooncake_master exited with code {code}. See " + f"{master.log_path}") + stopping.wait(1.0) + + +@click.command("mooncake_donor") +@click.option("--master_server_address", + type=str, + default=None, + help="Master to join, as host:port or file:// naming a " + "file that holds one. Defaults to the master_server_address in " + "--config.") +@click.option("--segment_size", + type=str, + default="32GiB", + show_default=True, + help="Host memory to contribute from this node. Deliberately " + "separate from a config's global_segment_size, which is sized " + "for an engine worker rather than a node lending what it can " + "spare.") +@click.option("--config", + type=str, + default=None, + help="Mooncake JSON config describing the pool, for the " + "settings not given here. Defaults to $MOONCAKE_CONFIG_PATH.") +@click.option("--protocol", + type=str, + default=None, + help="Transport, 'rdma' or 'tcp'. Defaults to --config's, else " + "rdma.") +@click.option("--device_name", + type=str, + default=None, + help="RDMA device, from ibv_devinfo. Defaults to --config's.") +@click.option("--metadata_server", + type=str, + default=None, + help="Mooncake metadata service. Defaults to --config's.") +@click.option("--ready_file", + type=str, + default=None, + help="File to create once the segment is mounted, for launchers " + "that must not let prefill start writing before the pool has " + "this capacity.") +@click.option("--heartbeat_seconds", + type=int, + default=300, + show_default=True, + help="Interval between liveness lines. 0 disables them.") +def mooncake_donor(master_server_address: Optional[str], segment_size: str, + config: Optional[str], protocol: Optional[str], + device_name: Optional[str], metadata_server: Optional[str], + ready_file: Optional[str], heartbeat_seconds: int): + """Lend this node's host memory to a Mooncake pool, for as long as it runs. + + Pool capacity comes only from processes that open a store handle, and in a + disaggregated deployment only the context servers do -- so the pool is + prefill-node memory, caching prefill's own GPUs. Running this on the + generation nodes puts their memory in the same pool while leaving those + engines connector-free. + """ + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( + DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, master_timeout, + parse_size, resolve_master_address) + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import \ + CONFIG_PATH_ENV + + raw = {} + config = config or os.getenv(CONFIG_PATH_ENV) + if config: + with open(config) as handle: + raw = json.load(handle) + + master = master_server_address or raw.get("master_server_address", "") + if not master: + raise click.UsageError( + "No master to join. Pass --master_server_address, or a --config " + f"naming one (or set {CONFIG_PATH_ENV}).") + + donating = parse_size(segment_size) + stopping = _until_signalled() + with donate_segment( + resolve_master_address(master, master_timeout()), + donating, + protocol=protocol or raw.get("protocol", "rdma"), + device_name=device_name or raw.get("device_name", "") or "", + metadata_server=metadata_server or raw.get("metadata_server", ""), + local_buffer_size=parse_size( + raw.get("local_buffer_size_donor", + DEFAULT_DONOR_LOCAL_BUFFER_SIZE)), + ) as host: + if ready_file: + with open(ready_file, "w") as handle: + handle.write(f"{host} {donating}\n") + + # Idle by design. A put or get here would make this node a client in + # the traffic sense, which is the thing keeping the generation engine + # connector-free is meant to avoid. + started = time.monotonic() + while not stopping.is_set(): + if heartbeat_seconds <= 0: + stopping.wait() + continue + if not stopping.wait(heartbeat_seconds): + logger.info("mooncake-store: still donating after " + f"{(time.monotonic() - started) / 60:.0f}m") diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index fdbf4a479a0a..d63059a49d13 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -29,6 +29,7 @@ from tensorrt_llm import MultimodalEncoder from tensorrt_llm._utils import mpi_rank, set_prometheus_multiproc_dir from tensorrt_llm.commands._serve_stability import stability_option +from tensorrt_llm.commands.mooncake import mooncake_donor, mooncake_master from tensorrt_llm.commands.utils import (collect_explicit_cli_keys, get_is_diffusion_only_model) from tensorrt_llm.executor.utils import MAX_NUM_FRONTENDS, LlmLauncherEnvs @@ -2530,7 +2531,12 @@ def resolve_command(self, ctx, args): "disaggregated": disaggregated, "disaggregated_mpi_worker": disaggregated_mpi_worker, "mm_embedding_serve": serve_encoder, - "embeddings": serve_embedding + "embeddings": serve_embedding, + # The parts of a Mooncake pool that cannot belong to a server, for the + # deployments where a pool outlives or spans them. A server that owns + # its pool provisions it from its own config and needs neither. + "mooncake_master": mooncake_master, + "mooncake_donor": mooncake_donor, }) if __name__ == "__main__": diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index b639a2a718f7..fea3bffb3c67 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1965,8 +1965,11 @@ class MooncakeStoreConfig(StrictBaseModel): """ master_server_address: Optional[str] = Field( None, - description="Address (host:port) of an already-running " - "mooncake_master. Mutually exclusive with launch_master.") + description="Address of an already-running mooncake_master, as " + "host:port or file:// naming a file that holds one -- which is " + "how to reach a master whose host a scheduler chose, since " + "'trtllm-serve mooncake_master --address_file' publishes it there " + "once it answers. Mutually exclusive with launch_master.") launch_master: bool = Field( False, telemetry=False, diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index b807f3e46f42..66de76c0e3d4 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -42,6 +42,7 @@ l0_a10: - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_kv_cache_layout.py - unittest/_torch/executor/test_mooncake_store_connector.py + - unittest/_torch/executor/test_mooncake_store_donor.py - unittest/_torch/executor/test_mooncake_store_master.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py new file mode 100644 index 000000000000..83dcf2f180f4 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for lending a node's host memory to a Mooncake pool. + +Runs without a Mooncake installation and without a GPU: the store is a fake +recording what ``setup`` was called with, since the contract being tested is +what the donor asks Mooncake for and how long it holds it, not what Mooncake +then does. +""" + +import sys +from types import ModuleType + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import donor as donor_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.donor import ( + DEFAULT_DONOR_LOCAL_BUFFER_SIZE, + donate_segment, +) + +GIB = 1024**3 + + +class FakeStore: + """The slice of ``MooncakeDistributedStore`` a donor drives.""" + + instances = [] + + def __init__(self): + self.setup_args = None + self.status = 0 + FakeStore.instances.append(self) + + def setup(self, *args): + self.setup_args = args + return self.status + + +@pytest.fixture +def fake_bindings(monkeypatch): + """Stand in for ``mooncake.store``, which is not installed here.""" + FakeStore.instances = [] + package = ModuleType("mooncake") + store = ModuleType("mooncake.store") + store.MooncakeDistributedStore = FakeStore + package.store = store + monkeypatch.setitem(sys.modules, "mooncake", package) + monkeypatch.setitem(sys.modules, "mooncake.store", store) + return FakeStore + + +@pytest.fixture +def failing_bindings(fake_bindings): + """Bindings whose ``setup`` refuses, as an unreachable master would.""" + + class Refusing(fake_bindings): + + def setup(self, *args): + super().setup(*args) + return 7 + + sys.modules["mooncake.store"].MooncakeDistributedStore = Refusing + return Refusing + + +def test_a_donor_registers_the_segment_it_was_asked_for(fake_bindings): + with donate_segment( + "10.0.0.1:50051", + 32 * GIB, + protocol="rdma", + device_name="mlx5_0", + metadata_server="P2PHANDSHAKE", + hostname="10.0.0.5", + ) as host: + assert host == "10.0.0.5" + ( + registered_host, + metadata_server, + segment_size, + local_buffer_size, + protocol, + device_name, + master, + ) = fake_bindings.instances[0].setup_args + + assert registered_host == "10.0.0.5" + assert metadata_server == "P2PHANDSHAKE" + assert segment_size == 32 * GIB + assert protocol == "rdma" + assert device_name == "mlx5_0" + assert master == "10.0.0.1:50051" + # The donor never transfers, so its transfer buffer is dead weight -- but + # setup rejects a zero one, hence a token rather than nothing. + assert local_buffer_size == DEFAULT_DONOR_LOCAL_BUFFER_SIZE + + +def test_a_donor_that_cannot_join_says_which_master_it_could_not_reach(failing_bindings): + with pytest.raises(RuntimeError, match="status 7"): + with donate_segment("10.0.0.1:50051", GIB, hostname="10.0.0.5"): + pytest.fail("donation should not have yielded") + + +def test_a_donor_given_no_host_registers_under_the_pool_s_view_of_this_node( + fake_bindings, monkeypatch): + """The master and the segments registering with it must agree on the host.""" + monkeypatch.setattr(donor_module, "local_address", lambda: "10.1.2.3") + + with donate_segment("10.0.0.1:50051", GIB) as host: + assert host == "10.1.2.3" + assert fake_bindings.instances[0].setup_args[0] == "10.1.2.3" + + +def test_missing_bindings_are_reported_as_the_separate_component_they_are(monkeypatch): + """The container's C++ transfer engine is not these Python bindings.""" + monkeypatch.setitem(sys.modules, "mooncake", None) + monkeypatch.setitem(sys.modules, "mooncake.store", None) + + with pytest.raises(ImportError, match="mooncake-transfer-engine"): + with donate_segment("10.0.0.1:50051", GIB): + pytest.fail("donation should not have yielded") diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py index d3aef17774ce..f307f4cd448f 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_master.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -25,6 +25,7 @@ import shutil import socket import subprocess +import threading from types import SimpleNamespace import pytest @@ -37,6 +38,7 @@ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.master import ( maybe_provision_pool, provision_pool, + resolve_master_address, ) from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig, MooncakeStoreConfig @@ -423,3 +425,106 @@ def test_a_described_pool_is_provisioned(running_master): written = json.loads(open(os.environ[CONFIG_PATH_ENV]).read()) assert written["master_server_address"] == running_master assert CONFIG_PATH_ENV not in os.environ + + +# ---- reaching a master whose host nobody knew in advance ---- + + +@pytest.mark.parametrize("address", ["10.0.0.1:50051", "unix:///var/run/mooncake"]) +def test_an_address_that_is_not_a_file_passes_through(address): + assert resolve_master_address(address, timeout=1.0) == address + + +def test_a_published_address_is_read_from_the_file_that_names_it(tmp_path): + published = tmp_path / "master.addr" + published.write_text("10.0.0.7:50051\n") + + assert resolve_master_address(f"file://{published}", timeout=1.0) == "10.0.0.7:50051" + + +def test_an_address_not_published_yet_is_waited_for(tmp_path): + """Master and workers are started together; neither one orders the other.""" + published = tmp_path / "master.addr" + threading.Timer(0.5, published.write_text, ["10.0.0.9:50051\n"]).start() + + assert resolve_master_address(f"file://{published}", timeout=10.0) == "10.0.0.9:50051" + + +def test_an_empty_address_file_is_not_taken_for_an_address(tmp_path): + """It exists, which is not the same as holding somewhere to connect to.""" + published = tmp_path / "master.addr" + published.write_text("") + + with pytest.raises(TimeoutError, match="No Mooncake master address"): + resolve_master_address(f"file://{published}", timeout=1.0) + + +def test_an_unpublished_address_names_the_command_that_publishes_it(tmp_path): + with pytest.raises(TimeoutError, match="--address-file"): + resolve_master_address(f"file://{tmp_path / 'absent'}", timeout=1.0) + + +# ---- a master with a lifetime of its own ---- + + +def test_a_standalone_master_publishes_an_address_that_can_be_dialed(fake_master, tmp_path): + port = free_port() + fake_master.arm(listen_on=port) + address_file = tmp_path / "master.addr" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with master_module.running_master( + pool, str(tmp_path / "run"), address_file=str(address_file) + ) as master: + assert resolve_master_address(f"file://{address_file}", timeout=5.0) == master.address + host, _, named_port = master.address.rpartition(":") + assert int(named_port) == port + with socket.create_connection((host, port), timeout=5): + pass + + +def test_a_stopped_master_leaves_no_address_behind(fake_master, tmp_path): + """A stale address would send the next run's workers at a dead port.""" + port = free_port() + fake_master.arm(listen_on=port) + address_file = tmp_path / "master.addr" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with master_module.running_master( + pool, str(tmp_path / "run"), address_file=str(address_file) + ): + assert address_file.exists() + + assert not address_file.exists() + assert fake_master.process.terminated + + +def test_a_standalone_master_keeps_its_log(fake_master, tmp_path): + """Its whole point is outliving servers, so its history is worth more.""" + port = free_port() + fake_master.arm(listen_on=port) + run_dir = tmp_path / "run" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with master_module.running_master(pool, str(run_dir)): + pass + + assert (run_dir / master_module.MASTER_LOG_NAME).exists() + + +def test_provisioning_joins_a_master_it_was_never_given_the_address_of(fake_master, tmp_path): + """The point of the file: no config and no script names a host.""" + port = free_port() + fake_master.arm(listen_on=port) + address_file = tmp_path / "master.addr" + standalone = MooncakeStoreConfig(launch_master=True, master_port=port) + worker = MooncakeStoreConfig(master_server_address=f"file://{address_file}") + + with master_module.running_master( + standalone, str(tmp_path / "run"), address_file=str(address_file) + ) as master: + with provision_pool(worker) as config_path: + # Mooncake cannot dial a file:// URL, so what reaches the workers + # has to be the address it resolved to. + written = json.loads(open(config_path).read()) + assert written["master_server_address"] == master.address From 7199b90b632f38b80880e5bc0550fc035f3cca0b Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:05:49 -0700 Subject: [PATCH 15/24] [None][chore] Assemble the Mooncake pool from what TensorRT-LLM ships The benchmark harness implemented the master's lifetime -- glog settings so it logs at all, a readiness probe, an address file -- and located mooncake_segment_donor.py through environment.trtllm_repo. All of that now has a shipped equivalent, so the script calls 'trtllm-serve mooncake_master' and 'trtllm-serve mooncake_donor' instead and the donor script is deleted rather than left as a second implementation of the same setup() call. The install step becomes a fallback: docker/common/install_mooncake.sh bakes the wheel into images built from this repo, so probe for the bindings and the binary first and only reach for trtllm_repo when an older image has neither. Two consequences worth knowing when reading a log directory: the address file now holds host:port rather than a bare host, and the master's glog moved to /mooncake_master.log, leaving 2_mooncake_master.log to the launching command -- so the block placement and eviction sections read the former. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docs/source/features/kv-cache-connector.md | 34 ++- .../slurm/benchmark/disaggr_torch.slurm | 125 +++++----- mooncake_disagg/README.md | 71 ++++-- mooncake_disagg/mooncake_segment_donor.py | 223 ------------------ mooncake_usage.md | 30 ++- 5 files changed, 163 insertions(+), 320 deletions(-) delete mode 100644 mooncake_disagg/mooncake_segment_donor.py diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index a00a68aa5a92..9a1d7bdc311f 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -126,9 +126,39 @@ kv_connector_config: local_buffer_size: 1GiB ``` -Replacing `master_server_address` with `launch_master: true` makes the server start a `mooncake_master` itself and use it, so a single-instance deployment needs nothing prepared outside `trtllm-serve`. **That master lives and dies with the server**, which makes it wrong for anything else: several engines that should share one pool would each get their own, and a pool meant to survive a restart cannot be owned by the thing restarting. Those deployments run a master with its own lifetime and name it in `master_server_address`. +Replacing `master_server_address` with `launch_master: true` makes the server start a `mooncake_master` itself and use it, so a single-instance deployment needs nothing prepared outside `trtllm-serve`. **That master lives and dies with the server**, which makes it wrong for anything else: several engines that should share one pool would each get their own, and a pool meant to survive a restart cannot be owned by the thing restarting. -`TRTLLM_MOONCAKE_MASTER_BINARY` overrides the binary a launched master runs, and `TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) how long startup waits for any master to accept connections -- reaching a master that is not there otherwise fails inside every rank after the model has loaded. Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated client config and the master's log, which are otherwise in a temporary directory removed at shutdown. +Those deployments run the master as its own command instead: + +```bash +trtllm-serve mooncake_master --rpc_port 50051 --address_file /shared/master.addr +``` + +The pool then lasts as long as that command, independently of any engine. `--address_file` receives `host:port` once the master accepts connections, and `master_server_address` accepts `file://` as well as a literal address: + +```yaml +kv_connector_config: + connector: mooncake-store + mooncake_store: + master_server_address: file:///shared/master.addr +``` + +This is what makes a master reachable without anyone writing its address down. Under a scheduler its host is not known when the configs are written; publishing it to a file the configs already name closes that gap, and a server reading the file waits for it, so the master and the engines can be started in any order. The file is removed when the master stops, so a stale address is never dialed. + +`TRTLLM_MOONCAKE_MASTER_BINARY` overrides the binary a launched master runs, and `TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) how long startup waits for any master to accept connections or publish its address -- reaching a master that is not there otherwise fails inside every rank after the model has loaded. Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated client config and the master's log, which are otherwise in a temporary directory removed at shutdown. + +#### Pool capacity + +Capacity comes only from processes that open a store handle, and `global_segment_size` is what each contributes -- so the pool is that value times the number of such processes. In a disaggregated deployment the connector belongs on the context servers only, which makes every byte of the pool prefill-node memory: prefill's DRAM caching prefill's GPUs, largely duplicating what `kv_cache_config.host_cache_size` already does. + +To give the pool memory from nodes that run no connector, run a donor on them: + +```bash +trtllm-serve mooncake_donor --master_server_address file:///shared/master.addr \ + --segment_size 160GiB --protocol rdma --device_name mlx5_0 +``` + +A donor holds a segment and issues no reads or writes, so a generation node can hold pages that prefill wrote while its engine stays connector-free and keeps its cache transceiver for the prefill-to-decode handoff. Contributing memory is deliberately not a `TRTLLM_MOONCAKE_STORE_ROLE`: the roles describe an engine's traffic, and every one of them reads or writes, so expressing capacity as a role would start that engine using the store. The donated memory is charged to the donor process and competes with anything else on the node, `kv_cache_config.host_cache_size` above all, so size the two together. Topology can equally come from a JSON file named by `MOONCAKE_CONFIG_PATH`, using the same schema as the vLLM Mooncake store connector so one deployment can point both engines at the same pool: diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm index d72f756832e8..a7616df19fa1 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm @@ -138,36 +138,42 @@ else echo "TensorRT-LLM environment variables saved to ${full_logdir}/env_vars.json" fi -# Install the Mooncake Python store bindings, but only when a worker config asks -# for the mooncake-store KV connector. The bindings shipped in the container -# images are unusable (see mooncake_disagg/README.md section 2), and the fix has -# to be reapplied per job: --container-name gives each node a container that -# lives for the job, so anything installed here survives to the worker sruns but -# not into the next job. +# The Mooncake store bindings, when a worker config asks for the connector. +# Images built from this repo bake them in (docker/common/install_mooncake.sh), +# so the install below is only a fallback for images that predate it. Either +# way it is per job: --container-name gives each node a container that lives +# for the job, so anything installed here survives to the worker sruns but not +# into the next job. mooncake_enabled=false if grep -qs "mooncake-store" "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"; then mooncake_enabled=true - mooncake_install_script="" - mooncake_donor_script="" - if [ -n "${trtllm_repo:-}" ]; then - mooncake_install_script="${trtllm_repo}/mooncake_disagg/install_mooncake_runtime.sh" - mooncake_donor_script="${trtllm_repo}/mooncake_disagg/mooncake_segment_donor.py" - fi - if [ -z "${mooncake_install_script}" ] || [ ! -f "${mooncake_install_script}" ]; then - cleanup_on_failure "A worker config requests the mooncake-store connector, but mooncake_disagg/install_mooncake_runtime.sh was not found under trtllm_repo='${trtllm_repo:-}'. Set environment.trtllm_repo to a checkout that contains it, or bake the bindings into the container image." - fi - if [ "${MOONCAKE_DONOR_SEGMENT_SIZE:-32GiB}" != "0" ] && [ ! -f "${mooncake_donor_script}" ]; then - cleanup_on_failure "mooncake_disagg/mooncake_segment_donor.py was not found under trtllm_repo='${trtllm_repo:-}'. It contributes the generation nodes' host memory to the pool; without it the pool is prefill-node memory only. Set MOONCAKE_DONOR_SEGMENT_SIZE=0 to accept that and skip the donors." - fi - echo "Installing Mooncake store bindings on all nodes..." - if ! srun --container-name=${container_name} \ + # Both halves of the wheel are load-bearing and fail at different times: + # the connector needs mooncake.store in every context rank, and + # 'trtllm-serve mooncake_master' needs the binary on PATH. + if srun --container-name=${container_name} \ --container-mounts=${container_mount} --no-container-mount-home \ - --mpi=pmix --overlap -N $SLURM_NNODES --ntasks-per-node=1 \ - bash -c "MOONCAKE_WHEEL='${MOONCAKE_WHEEL:-}' bash ${mooncake_install_script}" \ + --mpi=pmix --overlap -N 1 -n 1 \ + bash -c 'python3 -c "import mooncake.store" && command -v mooncake_master' \ &> ${full_logdir}/2_install_mooncake.log; then - cleanup_on_failure "Mooncake store bindings installation failed. Check ${full_logdir}/2_install_mooncake.log for details" + echo "Mooncake store bindings already in the image; nothing to install" + else + mooncake_install_script="" + if [ -n "${trtllm_repo:-}" ]; then + mooncake_install_script="${trtllm_repo}/mooncake_disagg/install_mooncake_runtime.sh" + fi + if [ -z "${mooncake_install_script}" ] || [ ! -f "${mooncake_install_script}" ]; then + cleanup_on_failure "A worker config requests the mooncake-store connector, this image does not have the bindings, and mooncake_disagg/install_mooncake_runtime.sh was not found under trtllm_repo='${trtllm_repo:-}'. Use an image built from this repo, or set environment.trtllm_repo to a checkout that contains the script." + fi + echo "Installing Mooncake store bindings on all nodes..." + if ! srun --container-name=${container_name} \ + --container-mounts=${container_mount} --no-container-mount-home \ + --mpi=pmix --overlap -N $SLURM_NNODES --ntasks-per-node=1 \ + bash -c "MOONCAKE_WHEEL='${MOONCAKE_WHEEL:-}' bash ${mooncake_install_script}" \ + &>> ${full_logdir}/2_install_mooncake.log; then + cleanup_on_failure "Mooncake store bindings installation failed. Check ${full_logdir}/2_install_mooncake.log for details" + fi + echo "Mooncake store bindings installation completed successfully" fi - echo "Mooncake store bindings installation completed successfully" fi # Get node lists and replace the placeholder with the actual node names @@ -209,52 +215,37 @@ if [ "${mooncake_enabled}" = "true" ]; then # --overlap because every GPU on this node is already claimed by a # worker; the master is a CPU-only process sharing the node with them. # - # mooncake_master logs through glog, which writes to files under /tmp - # inside the container unless told otherwise -- so without - # GLOG_logtostderr the log below would be empty. GLOG_v=1 adds the - # per-RPC lines that show workers registering segments and putting and - # getting keys, which is the only view of the pool's side of the - # conversation short of scraping the metrics port. + # The command owns the glog settings the master needs to log at all, + # the wait for its port, and the address file, so none of that is + # repeated here. --run_dir keeps the master's own log next to the job's + # rather than in a temporary directory removed at shutdown. srun -l --container-name=${container_name} \ --container-mounts=${container_mount} --no-container-mount-home \ --mpi=pmix --overlap --nodelist=${mooncake_master_node} -N 1 -n 1 \ - bash -c "export GLOG_logtostderr=1 GLOG_v=\"${MOONCAKE_MASTER_GLOG_V:-1}\"; \ - addr=\$(hostname -I 2>/dev/null | awk '{print \$1}'); \ - [ -n \"\${addr}\" ] || addr=\$(hostname -f); \ - echo \"\${addr}\" > ${mooncake_addr_file}; \ - echo \"mooncake_master starting on \$(hostname) at \${addr}:${mooncake_master_port} (metrics ${mooncake_master_metrics_port})\"; \ - exec mooncake_master --rpc_port=${mooncake_master_port} --metrics_port=${mooncake_master_metrics_port} --eviction_ratio=0.05" \ + trtllm-serve mooncake_master \ + --rpc_port ${mooncake_master_port} \ + --metrics_port ${mooncake_master_metrics_port} \ + --address_file ${mooncake_addr_file} \ + --run_dir ${full_logdir} \ &> ${full_logdir}/2_mooncake_master.log & - for _ in $(seq 1 60); do + # The address is published only once the master accepts connections, + # so this is the readiness wait as well. A worker that opens its store + # handle before then fails outright. + for _ in $(seq 1 120); do if [ -s "${mooncake_addr_file}" ]; then break fi sleep 1 done if [ ! -s "${mooncake_addr_file}" ]; then - cleanup_on_failure "mooncake_master did not report its address within 60s. Check ${full_logdir}/2_mooncake_master.log for details" - fi - mooncake_master_addr="$(tr -d '[:space:]' < ${mooncake_addr_file}):${mooncake_master_port}" - - # A worker that opens its store handle before the master accepts - # connections fails at startup, so wait for the port rather than the - # process. /dev/tcp keeps this dependency-free. - mooncake_master_ready=false - mooncake_wait_start=${SECONDS} - for _ in $(seq 1 60); do - if (exec 3<>/dev/tcp/${mooncake_master_addr%:*}/${mooncake_master_port}) 2>/dev/null; then - mooncake_master_ready=true - break - fi - sleep 1 - done - if [ "${mooncake_master_ready}" != "true" ]; then - cleanup_on_failure "mooncake_master at ${mooncake_master_addr} is not accepting connections. Check ${full_logdir}/2_mooncake_master.log for details" + cleanup_on_failure "mooncake_master did not publish an address to ${mooncake_addr_file} within 120s. Check ${full_logdir}/2_mooncake_master.log for details" fi - echo "mooncake_master ready at ${mooncake_master_addr} on node ${mooncake_master_node} after $((SECONDS - mooncake_wait_start))s of waiting" + # Published as host:port, so the port is not appended here. + mooncake_master_addr="$(tr -d '[:space:]' < ${mooncake_addr_file})" + echo "mooncake_master ready at ${mooncake_master_addr} on node ${mooncake_master_node}" echo "mooncake_master metrics: http://${mooncake_master_addr%:*}:${mooncake_master_metrics_port}" - echo "mooncake_master log: ${full_logdir}/2_mooncake_master.log" + echo "mooncake_master log: ${full_logdir}/mooncake_master.log" else echo "Using externally managed mooncake_master at ${mooncake_master_addr}" fi @@ -292,7 +283,9 @@ EOF # # A donor is a separate process rather than a connector role because the # roles describe traffic, not capacity: producer/consumer/both all read or - # write, and none of them means "contribute memory only". + # write, and none of them means "contribute memory only". It is a + # subcommand of trtllm-serve rather than a script, so a deployment outside + # this harness can lend memory the same way. # # The donated memory is charged to the donor process and competes with the # generation worker's own kv_cache_config.host_cache_size on that node, so @@ -328,10 +321,10 @@ EOF srun -l --container-name=${container_name} \ --container-mounts=${container_mount} --no-container-mount-home \ --mpi=pmix --overlap --nodelist=${donor_node} -N 1 -n 1 \ - python3 ${mooncake_donor_script} \ + trtllm-serve mooncake_donor \ --config ${full_logdir}/mooncake.json \ - --segment-size ${mooncake_donor_size} \ - --ready-file ${donor_ready_file} \ + --segment_size ${mooncake_donor_size} \ + --ready_file ${donor_ready_file} \ &> ${full_logdir}/2_mooncake_donor_${donor_node}.log & done @@ -389,7 +382,7 @@ if [ "${mooncake_enabled}" = "true" ]; then if ! grep -h "mooncake-store" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null; then echo " WARNING: no mooncake-store lines found. The connector may not have" \ "loaded; check ${full_logdir}/3_output_CTX_*.log and" \ - "${full_logdir}/2_mooncake_master.log. The benchmark will still run," \ + "${full_logdir}/mooncake_master.log. The benchmark will still run," \ "but without the store." fi fi @@ -432,7 +425,7 @@ if [ "${mooncake_enabled}" = "true" ]; then echo "== block placement by segment host ==" echo "(donor hosts: $(cat "${full_logdir}"/mooncake_donor_*.ready 2>/dev/null | awk '{print $1}' | paste -sd, - || echo none))" grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" \ - "${full_logdir}/2_mooncake_master.log" 2>/dev/null \ + "${full_logdir}/mooncake_master.log" 2>/dev/null \ | awk '{ sub(/size=/, "", $2); sub(/segment=/, "", $3); split($3, parts, ":"); host = parts[1]; port = parts[2]; @@ -457,11 +450,13 @@ if [ "${mooncake_enabled}" = "true" ]; then for donor_log in "${full_logdir}"/2_mooncake_donor_*.log; do [ -f "${donor_log}" ] || continue echo "--- $(basename "${donor_log}") ---" - grep -h "donating\|mounted\|failed\|exited" "${donor_log}" 2>/dev/null || tail -n 5 "${donor_log}" + grep -h "donating\|part of the pool\|withdrew\|failed" "${donor_log}" 2>/dev/null || tail -n 5 "${donor_log}" done echo + # The master's own log, which is glog and separate from the launching + # command's output in 2_mooncake_master.log. echo "== master ==" - tail -n 50 "${full_logdir}/2_mooncake_master.log" 2>/dev/null || echo "(no master log)" + tail -n 50 "${full_logdir}/mooncake_master.log" 2>/dev/null || echo "(no master log)" } > "${mooncake_summary}" 2>&1 echo "Mooncake store summary written to ${mooncake_summary}" cat "${mooncake_summary}" diff --git a/mooncake_disagg/README.md b/mooncake_disagg/README.md index 9ee1795ffce6..84f1d4035453 100644 --- a/mooncake_disagg/README.md +++ b/mooncake_disagg/README.md @@ -297,13 +297,21 @@ tier that largely duplicates TensorRT-LLM's native host offload. Confirm this on any run by grouping the master's `allocation_succeeded ... segment=:` lines by host: a single host means a prefill-only pool. -`mooncake_segment_donor.py` closes that gap. One donor per generation node opens -a handle, contributes memory, and then idles forever without a single put or get, -so the pool spans both sides while the generation engine stays connector-free -and keeps its cache transceiver for the KV handoff. A donor is deliberately not -a `StoreRole`: the roles describe traffic (`producer` writes, `consumer` reads, -`both`), and none of them means "contribute memory only", so capacity and -traffic have to be separate processes. +`trtllm-serve mooncake_donor` closes that gap. One donor per generation node +opens a handle, contributes memory, and then idles forever without a single put +or get, so the pool spans both sides while the generation engine stays +connector-free and keeps its cache transceiver for the KV handoff. Donation is +deliberately not a `StoreRole`: the roles describe traffic (`producer` writes, +`consumer` reads, `both`), and none of them means "contribute memory only", so +capacity and traffic have to be separate processes. + +Being a subcommand rather than a script, it is also how a deployment outside +this harness lends memory: + +```bash +trtllm-serve mooncake_donor --master_server_address 10.0.0.1:50051 \ + --segment_size 160GiB --protocol rdma --device_name mlx5_0 +``` `disaggr_torch.slurm` starts the donors automatically, reading the generation nodes off the generated worker commands and waiting for each segment to mount @@ -339,8 +347,10 @@ allocation, waits for its port to accept connections, and writes `MOONCAKE_CONFIG_PATH` from the log directory, which is why the harness config does not set it. Defaults are the bring-up ones (TCP, 16GiB per worker); `MOONCAKE_PROTOCOL`, `MOONCAKE_DEVICE_NAME`, `MOONCAKE_GLOBAL_SEGMENT_SIZE` and -`MOONCAKE_LOCAL_BUFFER_SIZE` in the submitting environment override them, and -the master's own log lands in `/2_mooncake_master.log`. +`MOONCAKE_LOCAL_BUFFER_SIZE` in the submitting environment override them. The +master itself is `trtllm-serve mooncake_master`, so its own log lands in +`/mooncake_master.log` and the launching command's output in +`/2_mooncake_master.log`. That master dies with the job, so read on if you need a pool that outlives one allocation -- which experiment 3 does, by construction. Run it as its own @@ -358,27 +368,39 @@ config pointing at yours. srun --container-image=$CONTAINER_IMAGE \ --container-mounts=$WORK_DIR:$WORK_DIR \ - bash -lc ' - hostname -i | awk "{print \$1}" > '"$WORK_DIR"'/master.addr - exec mooncake_master \ - --rpc_port=50051 \ - --metrics_port=9004 \ - --eviction_ratio=0.05 - ' + trtllm-serve mooncake_master \ + --rpc_port 50051 \ + --metrics_port 9004 \ + --address_file $WORK_DIR/master.addr \ + --run_dir $WORK_DIR ``` -Flag names above were read out of the shipped `mooncake_master` binary. Run -`mooncake_master --help` inside the container to confirm defaults and to see the -rest (`--rpc_address`, `--rpc_thread_num`, `--default_kv_lease_ttl`, +The master runs for as long as the command does, and `--address_file` receives +`host:port` **once it accepts connections** — so waiting for that file is +waiting for readiness, and its absence after the job starts is a failure rather +than a slow start. It is removed on exit, so a stale address is never dialed. +`--run_dir` keeps the master's log at `$WORK_DIR/mooncake_master.log`, which is +where pool occupancy and eviction are read from. + +`TRTLLM_MOONCAKE_MASTER_BINARY` overrides the binary this runs, and +`TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) how long it waits for the port. +Run `mooncake_master --help` inside the container for the flags this does not +surface (`--rpc_address`, `--rpc_thread_num`, `--default_kv_lease_ttl`, `--eviction_high_watermark_ratio`, `--enable_http_metadata_server`, `--cluster_id`, `--root_fs_dir`). Keeping the master in a separate job is what makes experiment 3 (§7) possible: the pool outlives the engines, so a second benchmark job finds a warm store. -Then write the client config, substituting the address the master job just -recorded -- or let `disaggr_torch.slurm` generate it, as above. The schema is -vLLM's, so one pool can serve both engines: +Workers can now be pointed at it without anyone writing the address down: +`master_server_address: file://$WORK_DIR/master.addr` in a config's +`mooncake_store` block makes each server read it during bringup and wait if the +master job has not started yet. That is what makes a master whose host the +scheduler chose usable from a config settled beforehand. + +Failing that, write the client config, substituting the address the master job +just recorded -- or let `disaggr_torch.slurm` generate it, as above. The schema +is vLLM's, so one pool can serve both engines: ```bash MASTER_IP=$(cat $WORK_DIR/master.addr) @@ -765,8 +787,9 @@ and being read back from there. `disaggr_torch.slurm` writes this breakdown into `9_mooncake_summary.log` at the end of every run, alongside the donor hosts, so it needs running by hand only when diagnosing a partial run. -Requires `GLOG_v=1` on the master, which `disaggr_torch.slurm` sets; raise it -with `MOONCAKE_MASTER_GLOG_V`. +Requires `GLOG_v=1` on the master, which `trtllm-serve mooncake_master` sets +unless `GLOG_v` is already in its environment -- so raise it by exporting +`GLOG_v` to that command. ### Which reuse number means what diff --git a/mooncake_disagg/mooncake_segment_donor.py b/mooncake_disagg/mooncake_segment_donor.py deleted file mode 100644 index f7161c271cfc..000000000000 --- a/mooncake_disagg/mooncake_segment_donor.py +++ /dev/null @@ -1,223 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Donate this node's host memory to a Mooncake store pool, without reading or -writing it. - -Pool capacity comes only from processes that open a store handle: ``setup`` -registers ``global_segment_size`` bytes of the calling process's host memory and -the master then places blocks in it. In a disaggregated deployment only the -context servers configure the KV connector, so only they call ``setup``, and the -pool is entirely prefill-node memory -- which makes the store a -prefill-DRAM-caches-prefill-GPU tier, overlapping what TensorRT-LLM's native -host offload already does. - -Running this alongside a generation server puts that node's memory into the same -pool. Prefill then writes blocks that land on decode-side DRAM, and reads them -back, while the generation engine itself stays free of any connector: it neither -reads nor writes the store, so it keeps its single cache transceiver for the -prefill-to-decode KV handoff. - -A donor is deliberately not a ``StoreRole``. The roles describe an engine's -traffic (``producer`` writes, ``consumer`` reads, ``both``), and none of them -means "contribute memory only" -- attaching a connector to the generation server -to get its DRAM into the pool would also start it reading or writing. Capacity -and traffic are separate concerns, so donation is a separate process. - -The donated memory is charged to this process, so it competes with the -generation server's own ``kv_cache_config.host_cache_size`` on the same node. -Size the two together. -""" - -import argparse -import json -import os -import re -import signal -import sys -import threading -import time - -_SIZE_UNITS = { - "": 1, - "b": 1, - "k": 1000, - "kb": 1000, - "m": 1000**2, - "mb": 1000**2, - "g": 1000**3, - "gb": 1000**3, - "t": 1000**4, - "tb": 1000**4, - "kib": 1024, - "mib": 1024**2, - "gib": 1024**3, - "tib": 1024**4, -} -_SIZE_RE = re.compile(r"^\s*([0-9]+(?:\.[0-9]+)?)\s*([a-zA-Z]*)\s*$") - -# The donor never transfers, so its staging buffer is dead weight; setup still -# rejects a zero one. -DEFAULT_LOCAL_BUFFER_SIZE = "64MiB" - - -def parse_size(value) -> int: - """Accept either a byte count or a suffixed string such as ``"32GiB"``. - - Mirrors the connector's parser so a size means the same thing in - ``mooncake.json`` and on this script's command line. - """ - if isinstance(value, bool): - raise ValueError(f"expected a size, got {value!r}") - if isinstance(value, int): - return value - if isinstance(value, float): - return int(value) - match = _SIZE_RE.match(str(value)) - if match is None: - raise ValueError(f"cannot parse size {value!r}") - magnitude, unit = match.groups() - scale = _SIZE_UNITS.get(unit.lower()) - if scale is None: - raise ValueError(f"unknown size unit {unit!r} in {value!r}") - return int(float(magnitude) * scale) - - -def log(message: str) -> None: - print(f"[donor {time.strftime('%H:%M:%S')}] {message}", flush=True) - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--config", - default=os.getenv("MOONCAKE_CONFIG_PATH"), - help="Mooncake JSON config naming the pool to join. Defaults to " - "$MOONCAKE_CONFIG_PATH.", - ) - parser.add_argument( - "--segment-size", - default="32GiB", - help="Host memory to contribute, e.g. 32GiB. Overrides the config's " - "global_segment_size, which is sized for an engine worker rather " - "than a node donating spare memory.", - ) - parser.add_argument( - "--ready-file", - default=None, - help="File to create once the segment is mounted, for launchers that " - "must not start writing to the pool before it has this capacity.", - ) - parser.add_argument( - "--heartbeat-seconds", - type=int, - default=300, - help="Interval between liveness lines. 0 disables them.", - ) - args = parser.parse_args() - - if not args.config: - parser.error("--config is required when MOONCAKE_CONFIG_PATH is unset") - - with open(args.config) as handle: - raw = json.load(handle) - - master = raw.get("master_server_address", "") - if not master: - parser.error(f"{args.config} has no master_server_address") - - segment_size = parse_size(args.segment_size) - local_buffer_size = parse_size( - raw.get("local_buffer_size_donor", DEFAULT_LOCAL_BUFFER_SIZE) - ) - protocol = raw.get("protocol", "rdma") - device_name = raw.get("device_name", "") or "" - metadata_server = raw.get("metadata_server", "") - - try: - from mooncake.store import MooncakeDistributedStore - except ImportError as exc: - log( - "the Mooncake Python bindings are missing " - "(`pip install mooncake-transfer-engine`); the C++ transfer engine " - f"in the container is a different component: {exc}" - ) - return 1 - - import socket - - hostname = socket.gethostbyname(socket.gethostname()) - - log( - f"joining pool at {master} as a capacity-only client: " - f"host={hostname} protocol={protocol} device={device_name or '(none)'} " - f"donating={segment_size / 1024 ** 3:.1f}GiB" - ) - - # Held for the process's lifetime: dropping the handle unmounts the segment - # and the master starts reporting the blocks living in it as lost. - store = MooncakeDistributedStore() - status = store.setup( - hostname, - metadata_server, - segment_size, - local_buffer_size, - protocol, - device_name, - master, - ) - if status != 0: - log( - f"setup failed with status {status}. The master must already be " - f"accepting connections at {master}, and protocol={protocol!r} " - "must be usable from this node." - ) - return 1 - - log(f"segment mounted; {segment_size / 1024 ** 3:.1f}GiB now available to the pool") - - if args.ready_file: - with open(args.ready_file, "w") as handle: - handle.write(f"{hostname} {segment_size}\n") - - stop = threading.Event() - - def handle_signal(signum, _frame): - log(f"received signal {signum}; unmounting segment") - stop.set() - - signal.signal(signal.SIGTERM, handle_signal) - signal.signal(signal.SIGINT, handle_signal) - - # Idle by design. Any put or get here would make this node a store client in - # the traffic sense, which is what keeping the generation engine - # connector-free is meant to avoid. - heartbeat = args.heartbeat_seconds - started = time.monotonic() - while not stop.is_set(): - if heartbeat > 0: - if stop.wait(heartbeat): - break - log(f"alive, donating for {(time.monotonic() - started) / 60:.0f}m") - else: - stop.wait() - - del store - log("exited") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/mooncake_usage.md b/mooncake_usage.md index 2932ef105087..0b506fb269eb 100644 --- a/mooncake_usage.md +++ b/mooncake_usage.md @@ -48,7 +48,7 @@ Three ways to get there, in increasing order of how much you have to arrange: | Deployment | Master | |---|---| | One `trtllm-serve`, own pool | `mooncake_store: {launch_master: true}` — the server starts it | -| Several engines, or a pool that outlives them | `mooncake_store: {master_server_address: host:50051}` — a master you run | +| Several engines, or a pool that outlives them | `trtllm-serve mooncake_master --address_file P`, then `mooncake_store: {master_server_address: file://P}` | | SLURM benchmark harness | Nothing: `disaggr_torch.slurm` starts the master and writes the JSON per job | The first two make `trtllm-serve` render the client config and export @@ -58,7 +58,15 @@ both and says so in the log, which is why the harness path is unaffected. A launched master dies with the server, so use it only for a single engine: two context servers that each launch one get two disjoint pools, and the -survival-across-restart case is impossible by construction. +survival-across-restart case is impossible by construction. Those cases want +row two, where the master is its own command and nothing else's lifetime +bounds it. + +`master_server_address` takes a plain `host:port` or `file://`. The file +is what a scheduler-placed master needs: its host is not known when the configs +are written, `--address_file` publishes it once the master answers, and a +server reading it waits for the master to exist. Nobody has to write an address +down, and a stale one cannot be dialed because the file is removed on exit. Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated JSON and the master's log, which otherwise sit in a temporary directory that shutdown removes. @@ -99,10 +107,20 @@ Per-process environment, on the workers that open a handle: | `TRTLLM_MOONCAKE_STORE_MODEL_KEY` | Defaults to the checkpoint directory's basename — set it explicitly for anything long-lived. | Pool capacity comes only from processes that open a store handle, so a -prefill-only connector gives a prefill-only pool. `mooncake_segment_donor.py` -contributes host memory from the generation nodes without any traffic; -`disaggr_torch.slurm` starts one per generation node -(`MOONCAKE_DONOR_SEGMENT_SIZE`, default `32GiB`). +prefill-only connector gives a prefill-only pool — which caches prefill's GPUs +in prefill's own DRAM, largely duplicating the native host offload. +`trtllm-serve mooncake_donor` contributes host memory from a node without any +traffic, leaving that engine connector-free: + +```bash +trtllm-serve mooncake_donor --master_server_address file://$WORK_DIR/master.addr \ + --segment_size 160GiB --protocol rdma --device_name mlx5_0 +``` + +`disaggr_torch.slurm` starts one per generation node already +(`MOONCAKE_DONOR_SEGMENT_SIZE`, default `32GiB`, `0` to skip). Donated memory +is charged to the donor process, so it competes with that node's own +`kv_cache_config.host_cache_size` — size the two together. ## 3. Partial reuse must be off — now enforced From ca850792002c758c46fe6bda41791c6990705858 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:38:18 -0700 Subject: [PATCH 16/24] [None][feat] Configure the whole Mooncake pool, launch none of it A pool assembled by a launch script is a pool only that script can assemble. The parts a server could not own -- a master with its own lifetime, memory lent by a node whose engine uses no connector -- were commands the harness ran, which left the shape of a deployment split between the configs and a 1000-line SLURM script. Both become config. A launched master publishes its address, so a server lending memory can find the pool a context server owns without anyone writing an address down, and mooncake_donation contributes host memory from a server that never reads or writes the store. The harness now installs the bindings and substitutes its log directory; nothing else about the pool is its business. The run directory turns out to be load-bearing for more than logs. Provisioning reaches the ranks the LLM constructor spawns by exporting MOONCAKE_CONFIG_PATH, but under trtllm-llmapi-launch each rank is its own task and was already running, so those ranks now read the rendered config back from the run directory. Without that, every rank but the leader of a multi-GPU server failed during bringup. Bringup narrates itself throughout, because a pool that came up wrong is otherwise visible only as a low hit rate hours later: what was resolved from where, the segment in bytes as well as GiB, the capacity arithmetic, and the tail of the master's own log when it dies during startup. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docs/source/features/kv-cache-connector.md | 50 ++- .../slurm/benchmark/disaggr_torch.slurm | 194 ++-------- .../slurm/benchmark/start_worker.sh | 48 ++- mooncake_disagg/README.md | 74 ++-- mooncake_disagg/m3_ctx_mooncake.yaml | 28 ++ mooncake_disagg/m3_gen_mooncake.yaml | 16 + mooncake_usage.md | 68 +++- .../connectors/mooncake_store/__init__.py | 9 +- .../connectors/mooncake_store/config.py | 36 +- .../connectors/mooncake_store/donor.py | 91 ++++- .../connectors/mooncake_store/master.py | 366 ++++++++++++++---- tensorrt_llm/commands/mooncake.py | 40 +- tensorrt_llm/commands/serve.py | 53 ++- tensorrt_llm/llmapi/llm_args.py | 81 ++++ .../executor/test_mooncake_store_connector.py | 40 ++ .../executor/test_mooncake_store_donor.py | 81 ++++ .../executor/test_mooncake_store_master.py | 156 +++++++- .../api_stability/references/llm.yaml | 4 + 18 files changed, 1090 insertions(+), 345 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 9a1d7bdc311f..eff31a6addfe 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -128,7 +128,15 @@ kv_connector_config: Replacing `master_server_address` with `launch_master: true` makes the server start a `mooncake_master` itself and use it, so a single-instance deployment needs nothing prepared outside `trtllm-serve`. **That master lives and dies with the server**, which makes it wrong for anything else: several engines that should share one pool would each get their own, and a pool meant to survive a restart cannot be owned by the thing restarting. -Those deployments run the master as its own command instead: +A master started this way still publishes its address, to `master.addr` in the run directory and to `master_address_file` if one is named. That is what lets other processes find a pool this server owns -- the donors below, most of all -- and what makes a finished run's logs say which master it used: + +```yaml +mooncake_store: + launch_master: true + master_address_file: /shared/master.addr +``` + +The two cases above -- several engines, or surviving a restart -- run the master as its own command instead: ```bash trtllm-serve mooncake_master --rpc_port 50051 --address_file /shared/master.addr @@ -147,19 +155,53 @@ This is what makes a master reachable without anyone writing its address down. U `TRTLLM_MOONCAKE_MASTER_BINARY` overrides the binary a launched master runs, and `TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) how long startup waits for any master to accept connections or publish its address -- reaching a master that is not there otherwise fails inside every rank after the model has loaded. Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated client config and the master's log, which are otherwise in a temporary directory removed at shutdown. +#### Servers whose ranks the launcher starts + +Provisioning happens in the server process and reaches the ranks that open store handles by exporting `MOONCAKE_CONFIG_PATH` for them to inherit. That holds when the LLM constructor spawns them, and does not when the launcher starts one task per rank -- as `trtllm-llmapi-launch` under a scheduler does -- because those ranks were already running. + +Naming a shared run directory covers that case: the rendered config is read back from `$TRTLLM_MOONCAKE_RUN_DIR/mooncake.json` by any rank that inherited no path, so every rank of a multi-GPU server joins the pool its own leader provisioned. The directory has to be one they all see, which under a scheduler means the job's own, and it is where the master's log and published address already go: + +```bash +export TRTLLM_MOONCAKE_RUN_DIR=/shared/run/$SLURM_JOB_ID +srun trtllm-llmapi-launch trtllm-serve "$model" --config ctx.yaml +``` + +Without it, a rank that inherited nothing fails during bringup naming `MOONCAKE_CONFIG_PATH`, rather than serving without a store. + +#### Reading bringup in the log + +Everything the pool is assembled from is logged under the `mooncake-store:` prefix before the model loads, because a pool that came up wrong is otherwise visible only as a low hit rate hours later. In order: the run directory, the master's command line and pid, the address it published and where, the rendered client config in full, and the capacity each rank will contribute. A server lending memory logs the master it resolved, the segment in both GiB and bytes, and the transport -- a size string parsed wrong is otherwise invisible until the pool starts evicting far too eagerly. + +Both waits narrate themselves every five seconds, since waiting for a master in another job is normal and indistinguishable from a hang if it is silent. A master that dies during startup has the tail of its own log quoted in the failure, which is where the reason (a port in use, a bad flag) actually is. + #### Pool capacity Capacity comes only from processes that open a store handle, and `global_segment_size` is what each contributes -- so the pool is that value times the number of such processes. In a disaggregated deployment the connector belongs on the context servers only, which makes every byte of the pool prefill-node memory: prefill's DRAM caching prefill's GPUs, largely duplicating what `kv_cache_config.host_cache_size` already does. -To give the pool memory from nodes that run no connector, run a donor on them: +To give the pool memory from nodes whose engines run no connector, ask those servers to lend it: + +```yaml +# generation server -- no connector, memory only +mooncake_donation: + master_server_address: file:///shared/master.addr + segment_size: 320GiB + protocol: rdma + device_name: mlx5_1 +``` + +`trtllm-serve` then holds that segment for as long as the server runs, so a generation node holds pages prefill wrote while its own engine stays connector-free and keeps its cache transceiver for the prefill-to-decode handoff. The server is ready only once the segment is mounted, which makes its readiness the signal that the pool has this capacity. + +Lending memory is deliberately outside `kv_connector_config`, and not a `TRTLLM_MOONCAKE_STORE_ROLE` either. Both of those attach a connector, and a connector reads or writes -- `producer`, `consumer` and `both` all describe traffic, and none of them means "contribute memory only" -- so expressing capacity there would start this server using the store. Capacity and traffic are separate, and configured separately. + +Size is charged **per server process, not per rank**, unlike `global_segment_size`. Two servers on one node lend twice this. The memory is charged to the process and competes with everything else on the node, `kv_cache_config.host_cache_size` above all, so size the two together. + +A node that runs no server at all can still lend, as its own command: ```bash trtllm-serve mooncake_donor --master_server_address file:///shared/master.addr \ --segment_size 160GiB --protocol rdma --device_name mlx5_0 ``` -A donor holds a segment and issues no reads or writes, so a generation node can hold pages that prefill wrote while its engine stays connector-free and keeps its cache transceiver for the prefill-to-decode handoff. Contributing memory is deliberately not a `TRTLLM_MOONCAKE_STORE_ROLE`: the roles describe an engine's traffic, and every one of them reads or writes, so expressing capacity as a role would start that engine using the store. The donated memory is charged to the donor process and competes with anything else on the node, `kv_cache_config.host_cache_size` above all, so size the two together. - Topology can equally come from a JSON file named by `MOONCAKE_CONFIG_PATH`, using the same schema as the vLLM Mooncake store connector so one deployment can point both engines at the same pool: ```json diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm index a7616df19fa1..6283ae6c0828 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm @@ -192,163 +192,18 @@ client_cmds_base_file=${full_logdir}/client_cmds_base.sh client_cmds_file=${full_logdir}/client_cmds.sh replace_placeholder "${client_cmds_base_file}" "${all_nodes_str}" "${client_cmds_file}" -# Bring up the Mooncake store's pool: a master process plus the client config -# that names it. Both are per job and land in the log directory, so nothing has -# to be prepared by hand and concurrent jobs do not share state. +# The pool is described in the worker configs and provisioned by trtllm-serve +# during its own bringup: the context server starts the master, renders the +# client config and publishes the master's address, and the generation servers +# read that address to lend the pool their memory. Nothing here starts, waits +# for or configures any of it. # -# The master lives inside this job, which means the pool dies with it. That is -# the right default for the single-job experiments, but it makes the -# survival-across-restarts case impossible to test, so set -# MOONCAKE_MASTER_ADDRESS in the submitting environment to reuse a master -# running as its own longer-lived job (mooncake_disagg/README.md section 4); -# this block then only writes the client config. +# The one value a config written before submission cannot know is where this +# job's log directory is, and the master's address is published into it, so a +# __LOG_DIR__ placeholder in either worker config is filled in here. if [ "${mooncake_enabled}" = "true" ]; then - mooncake_master_port=50051 - mooncake_master_metrics_port=9004 - mooncake_master_addr="${MOONCAKE_MASTER_ADDRESS:-}" - - if [ -z "${mooncake_master_addr}" ]; then - mooncake_master_node="${all_nodes[0]}" - mooncake_addr_file="${full_logdir}/mooncake_master.addr" - rm -f "${mooncake_addr_file}" - echo "Starting mooncake_master on ${mooncake_master_node}..." - # --overlap because every GPU on this node is already claimed by a - # worker; the master is a CPU-only process sharing the node with them. - # - # The command owns the glog settings the master needs to log at all, - # the wait for its port, and the address file, so none of that is - # repeated here. --run_dir keeps the master's own log next to the job's - # rather than in a temporary directory removed at shutdown. - srun -l --container-name=${container_name} \ - --container-mounts=${container_mount} --no-container-mount-home \ - --mpi=pmix --overlap --nodelist=${mooncake_master_node} -N 1 -n 1 \ - trtllm-serve mooncake_master \ - --rpc_port ${mooncake_master_port} \ - --metrics_port ${mooncake_master_metrics_port} \ - --address_file ${mooncake_addr_file} \ - --run_dir ${full_logdir} \ - &> ${full_logdir}/2_mooncake_master.log & - - # The address is published only once the master accepts connections, - # so this is the readiness wait as well. A worker that opens its store - # handle before then fails outright. - for _ in $(seq 1 120); do - if [ -s "${mooncake_addr_file}" ]; then - break - fi - sleep 1 - done - if [ ! -s "${mooncake_addr_file}" ]; then - cleanup_on_failure "mooncake_master did not publish an address to ${mooncake_addr_file} within 120s. Check ${full_logdir}/2_mooncake_master.log for details" - fi - # Published as host:port, so the port is not appended here. - mooncake_master_addr="$(tr -d '[:space:]' < ${mooncake_addr_file})" - echo "mooncake_master ready at ${mooncake_master_addr} on node ${mooncake_master_node}" - echo "mooncake_master metrics: http://${mooncake_master_addr%:*}:${mooncake_master_metrics_port}" - echo "mooncake_master log: ${full_logdir}/mooncake_master.log" - else - echo "Using externally managed mooncake_master at ${mooncake_master_addr}" - fi - - # The schema is vLLM's, so one pool can serve both engines. Defaults are - # the first-bring-up ones: TCP removes RDMA from the variable list, at the - # cost of any performance conclusion. Set MOONCAKE_PROTOCOL=rdma with a - # MOONCAKE_DEVICE_NAME from ibv_devinfo for a run worth quoting. - # global_segment_size is contributed per worker process, so the pool is - # this value times (ctx instances x world size). - cat > "${full_logdir}/mooncake.json" < ${full_logdir}/2_mooncake_donor_${donor_node}.log & - done - - # Wait for the segments before any worker starts, so that the first - # blocks prefill writes can already be placed on a decode node. A donor - # that never mounts is a hard failure: the run would otherwise quietly - # fall back to the prefill-only pool this is meant to replace. - for donor_ready_file in "${mooncake_donor_ready_files[@]}"; do - for _ in $(seq 1 120); do - if [ -s "${donor_ready_file}" ]; then - break - fi - sleep 1 - done - if [ ! -s "${donor_ready_file}" ]; then - donor_node="$(basename "${donor_ready_file}" .ready)" - donor_node="${donor_node#mooncake_donor_}" - cleanup_on_failure "The memory donor on ${donor_node} did not mount its segment within 120s. Check ${full_logdir}/2_mooncake_donor_${donor_node}.log; if the node is short on memory, lower MOONCAKE_DONOR_SEGMENT_SIZE or the generation worker's kv_cache_config.host_cache_size" - fi - done - echo "Pool capacity: ${MOONCAKE_GLOBAL_SEGMENT_SIZE:-16GiB} per context worker process" \ - "+ ${mooncake_donor_size} per generation node (${#mooncake_donor_nodes[@]} donor(s):" \ - "${mooncake_donor_nodes[*]})" - fi + sed -i "s|__LOG_DIR__|${full_logdir}|g" \ + "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml" fi # Per-worker hostfile / gpu_map files for srun --distribution=arbitrary. @@ -404,7 +259,12 @@ done 3< "${client_cmds_file}" if [ "${mooncake_enabled}" = "true" ]; then mooncake_summary="${full_logdir}/9_mooncake_summary.log" { - echo "master: ${mooncake_master_addr}" + # The address file is retracted when the master stops, so a stale + # address is never dialed; the context server's log still says which + # master the run used. + echo "master: $(tr -d '[:space:]' < "${full_logdir}/master.addr" 2>/dev/null \ + || grep -hoE "master at [0-9.]+:[0-9]+" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null \ + | head -n 1 | awk '{print $3}' || echo unknown)" echo echo "== startup ==" grep -h "mooncake-store.*\(ready\|registered layout\)" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null || echo "(none)" @@ -423,7 +283,9 @@ if [ "${mooncake_enabled}" = "true" ]; then # rather than on the prefill node that computed it. Without a donor this # section shows a single host, which is the prefill node. echo "== block placement by segment host ==" - echo "(donor hosts: $(cat "${full_logdir}"/mooncake_donor_*.ready 2>/dev/null | awk '{print $1}' | paste -sd, - || echo none))" + echo "(lending hosts: $(grep -hoE "GiB of [0-9.]+ is now part of the pool" \ + "${full_logdir}"/3_output_GEN_*.log 2>/dev/null \ + | awk '{print $3}' | sort -u | paste -sd, - || echo none))" grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" \ "${full_logdir}/mooncake_master.log" 2>/dev/null \ | awk '{ @@ -446,15 +308,17 @@ if [ "${mooncake_enabled}" = "true" ]; then printf "%-16s pages=%-7d %8.2f GiB\n", "TOTAL", total_allocs, total_bytes / 1073741824; }' || echo "(could not parse master log)" echo - echo "== donors ==" - for donor_log in "${full_logdir}"/2_mooncake_donor_*.log; do - [ -f "${donor_log}" ] || continue - echo "--- $(basename "${donor_log}") ---" - grep -h "donating\|part of the pool\|withdrew\|failed" "${donor_log}" 2>/dev/null || tail -n 5 "${donor_log}" - done + # Lending memory happens inside the generation servers, so their own + # logs are where a segment that never mounted shows up. + echo "== lent segments ==" + grep -h "mooncake-store: .*\(lending memory\|part of the pool\|withdrew\)" \ + "${full_logdir}"/3_output_GEN_*.log 2>/dev/null || echo "(none)" + echo + echo "== pool bringup (context server) ==" + grep -h "mooncake-store:" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null \ + | head -n 30 || echo "(none)" echo - # The master's own log, which is glog and separate from the launching - # command's output in 2_mooncake_master.log. + # The master's own log, glog rather than anything TensorRT-LLM writes. echo "== master ==" tail -n 50 "${full_logdir}/mooncake_master.log" 2>/dev/null || echo "(no master log)" } > "${mooncake_summary}" 2>&1 diff --git a/examples/disaggregated/slurm/benchmark/start_worker.sh b/examples/disaggregated/slurm/benchmark/start_worker.sh index fa8eca6425ae..34aebdef7894 100644 --- a/examples/disaggregated/slurm/benchmark/start_worker.sh +++ b/examples/disaggregated/slurm/benchmark/start_worker.sh @@ -54,14 +54,46 @@ fi echo "config_file: ${config_file}" -# The mooncake-store KV connector reads its pool topology from -# MOONCAKE_CONFIG_PATH. disaggr_torch.slurm generates one per job in the log -# directory, whose path is not known when submit.py builds the worker -# environment; an explicit setting still wins, so pointing at an externally -# managed pool remains possible. -if [ -z "${MOONCAKE_CONFIG_PATH:-}" ] && [ -f "${log_dir}/mooncake.json" ]; then - export MOONCAKE_CONFIG_PATH="${log_dir}/mooncake.json" - echo "MOONCAKE_CONFIG_PATH: ${MOONCAKE_CONFIG_PATH}" +# The mooncake-store pool is described in the worker config and provisioned by +# trtllm-serve during bringup. Anchoring its run directory here is what puts the +# master's log, the client config it renders and the address it publishes in the +# job's log directory rather than in a temporary directory that shutdown +# removes -- and it is how the ranks the launcher started, which never inherited +# the leader's environment, find that client config. An inherited +# MOONCAKE_CONFIG_PATH still wins, so an externally managed pool stays reachable. +export TRTLLM_MOONCAKE_RUN_DIR="${log_dir}" + +# The generation servers wait for a master the context server starts. Both are +# launched together and the master comes up before its model loads, but the wait +# spans container start on another node, so it is given far more than the 60s +# default: too short fails the job, too long costs nothing when the master is +# there. +export TRTLLM_MOONCAKE_MASTER_TIMEOUT="${TRTLLM_MOONCAKE_MASTER_TIMEOUT:-900}" + +# MiniMax-M3's MSA sparse attention JIT-compiles its FMHA kernels on first use, +# from inside the attention forward pass: one TP rank runs ninja while the others +# block on a file lock, so an uncached variant stalls the whole executor loop for +# ~8s (and ~70s when an iteration needs several). The cache defaults to +# ~/.cache, which is thrown away here because the container is started with +# --no-container-mount-home, making every job pay the compiles again during +# serving. Anchor it next to this script instead: that path is on the mounted +# filesystem and identical across jobs, so only the first run compiles. +if [ -z "${MINFER_FMHA_CACHE_DIR:-}" ]; then + export MINFER_FMHA_CACHE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/.cache/minfer/fmha_sm100" + mkdir -p "${MINFER_FMHA_CACHE_DIR}" + echo "MINFER_FMHA_CACHE_DIR: ${MINFER_FMHA_CACHE_DIR}" +fi + +# Per-transfer KV timings (size, queue/transfer latency, throughput) as CSV next +# to the worker logs. This is what tells apart "prefill is slow" from "the +# prefill->decode handoff is slow", which the aggregate benchmark numbers +# cannot. Same rationale as above for defaulting the path here: an explicit +# setting wins, and it can be turned off with KV_TRANSFER_PERF_LOG=false. +if [ "${KV_TRANSFER_PERF_LOG:-true}" = "true" ] \ + && [ -z "${TLLM_KV_TRANSFER_PERF_LOG_FILE:-}" ]; then + export TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1 + export TLLM_KV_TRANSFER_PERF_LOG_FILE="${log_dir}/kv_transfer_perf" + echo "TLLM_KV_TRANSFER_PERF_LOG_FILE: ${TLLM_KV_TRANSFER_PERF_LOG_FILE}" fi nsys_prefix="" diff --git a/mooncake_disagg/README.md b/mooncake_disagg/README.md index 84f1d4035453..e6cba4cc7543 100644 --- a/mooncake_disagg/README.md +++ b/mooncake_disagg/README.md @@ -297,29 +297,39 @@ tier that largely duplicates TensorRT-LLM's native host offload. Confirm this on any run by grouping the master's `allocation_succeeded ... segment=:` lines by host: a single host means a prefill-only pool. -`trtllm-serve mooncake_donor` closes that gap. One donor per generation node -opens a handle, contributes memory, and then idles forever without a single put -or get, so the pool spans both sides while the generation engine stays -connector-free and keeps its cache transceiver for the KV handoff. Donation is -deliberately not a `StoreRole`: the roles describe traffic (`producer` writes, -`consumer` reads, `both`), and none of them means "contribute memory only", so -capacity and traffic have to be separate processes. +`mooncake_donation` on the generation worker closes that gap. The server opens +a handle, contributes memory, and then holds it without a single put or get for +as long as it runs, so the pool spans both sides while its engine stays +connector-free and keeps its cache transceiver for the KV handoff: -Being a subcommand rather than a script, it is also how a deployment outside -this harness lends memory: +```yaml +# gen worker config -- no kv_connector_config anywhere near it +mooncake_donation: + master_server_address: file:///$WORK_DIR/master.addr + segment_size: 640GiB + protocol: rdma +``` + +Donation is deliberately outside `kv_connector_config`, and not a `StoreRole` +either: the roles describe traffic (`producer` writes, `consumer` reads, +`both`), none of them means "contribute memory only", and configuring capacity +there would start this server using the store. The size is charged **per server +process, not per rank** — unlike `global_segment_size` — so two servers on one +node lend twice this. + +The server is ready only once its segment is mounted, which makes readiness the +signal that the pool has the capacity, and means the first blocks prefill writes +can already land on a decode node. Set `segment_size: 0`, or leave the section +out, to keep the pool prefill-only. + +A node running no server lends as its own command instead, which is also how a +machine with no GPUs contributes: ```bash trtllm-serve mooncake_donor --master_server_address 10.0.0.1:50051 \ --segment_size 160GiB --protocol rdma --device_name mlx5_0 ``` -`disaggr_torch.slurm` starts the donors automatically, reading the generation -nodes off the generated worker commands and waiting for each segment to mount -before any worker starts — so the first blocks prefill writes can already land -on a decode node. Tune with `MOONCAKE_DONOR_SEGMENT_SIZE` (default `32GiB`, set -`0` to keep the pool prefill-only) and `MOONCAKE_DONOR_NODES` to override node -selection. - The donated memory is charged to the donor process and competes with the generation worker's own `kv_cache_config.host_cache_size` on that node, so size the two together. The worker logs its own share as `KV cache manager v2 host @@ -340,17 +350,21 @@ what `m3_agg_mooncake.yaml` now does; `mooncake_usage.md` §2 has the table. The rest of this section is about the master the experiments below need, which outlives any one server and therefore cannot be owned by one. -**For a single-job experiment you can skip this section.** -`disaggr_torch.slurm` now starts a `mooncake_master` on the first node of the -allocation, waits for its port to accept connections, and writes -`/mooncake.json` naming it; `start_worker.sh` then resolves -`MOONCAKE_CONFIG_PATH` from the log directory, which is why the harness config -does not set it. Defaults are the bring-up ones (TCP, 16GiB per worker); -`MOONCAKE_PROTOCOL`, `MOONCAKE_DEVICE_NAME`, `MOONCAKE_GLOBAL_SEGMENT_SIZE` and -`MOONCAKE_LOCAL_BUFFER_SIZE` in the submitting environment override them. The -master itself is `trtllm-serve mooncake_master`, so its own log lands in -`/mooncake_master.log` and the launching command's output in -`/2_mooncake_master.log`. +**For a single-job experiment you can skip this section.** The harness starts +no master and writes no client config: the context worker's `launch_master: +true` does both, on the context node, and publishes the address the generation +workers' `mooncake_donation` reads. `disaggr_torch.slurm` contributes exactly +two things — it installs the bindings on every node, and it substitutes +`__LOG_DIR__` in the worker configs, since the run directory is the one value a +config written before submission cannot know. Everything else, the pool sizes +and the HCA included, is in the config; no `MOONCAKE_*` variable is read from +the submitting environment any more. + +The master's log lands in `/mooncake_master.log` and its address in +`/master.addr` while it runs, because `start_worker.sh` sets +`TRTLLM_MOONCAKE_RUN_DIR` to the log directory. That is also what lets the +context server's other ranks read the rendered `mooncake.json`: they are +separate srun tasks that never inherited the leader's environment. That master dies with the job, so read on if you need a pool that outlives one allocation -- which experiment 3 does, by construction. Run it as its own @@ -493,8 +507,8 @@ environment: work_dir: "" worker_env_var: "TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes NCCL_GRAPH_MIXING_SUPPORT=0" # Only the context workers open a store handle. MOONCAKE_CONFIG_PATH is - # deliberately absent: the harness generates the file per job in the log - # directory, whose path is not known when submit.py builds this environment. + # deliberately absent: the context server renders that file itself, into the + # log directory, and its own ranks read it back from there. ctx_worker_env_var: "TRTLLM_MOONCAKE_STORE_ROLE=both TRTLLM_MOONCAKE_STORE_PREFIX=trtllm-m3-run1" server_env_var: "TRTLLM_SERVER_DISABLE_GC=1" @@ -775,7 +789,7 @@ offload tier could not. For that, group the master's allocations by segment host node the block physically lives on: ```bash -grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" 2_mooncake_master.log \ +grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" mooncake_master.log \ | awk '{sub(/size=/,"",$2); sub(/segment=/,"",$3); split($3,p,":"); n[p[1]]++; b[p[1]]+=$2} END {for (h in n) printf "%-16s pages=%-7d %.2f GiB\n", h, n[h], b[h]/1073741824}' diff --git a/mooncake_disagg/m3_ctx_mooncake.yaml b/mooncake_disagg/m3_ctx_mooncake.yaml index f937d24e3c21..0d25d67bb4ac 100644 --- a/mooncake_disagg/m3_ctx_mooncake.yaml +++ b/mooncake_disagg/m3_ctx_mooncake.yaml @@ -66,3 +66,31 @@ enable_attention_dp: false kv_connector_config: connector: mooncake-store + # The pool, provisioned by this server during its own bringup: it starts + # the master, renders the client config and exports MOONCAKE_CONFIG_PATH + # before the ranks that open store handles exist. Nothing outside + # trtllm-serve prepares any of it. + # + # Only right because there is one context server. Two that each launch a + # master get two disjoint pools, which is what 'trtllm-serve + # mooncake_master' plus master_server_address: file:// is for. + mooncake_store: + launch_master: true + # Where the master's address is published, for the generation + # workers to lend the pool memory. One also goes to the run + # directory regardless; under the SLURM harness that is the job's + # log directory, and __LOG_DIR__ is substituted there. + master_address_file: __LOG_DIR__/master.addr + protocol: rdma + # device_name omitted on purpose: the fastest active HCAs on this + # node are detected, which keeps one config portable across node + # types. Name them to pin it. + # + # Charged per rank, so a TP=2 server contributes twice this. + global_segment_size: 160GiB + # Registering the KV pools with the HCA needs nvidia_peermem; without + # it registration fails on every range, and staging registers only + # host memory instead. 1GiB holds a full transfer_batch_size of + # pages -- less silently reduces the batch rather than failing. + stage_through_host: true + staging_buffer_bytes: 1GiB diff --git a/mooncake_disagg/m3_gen_mooncake.yaml b/mooncake_disagg/m3_gen_mooncake.yaml index c6e84cf6d0c1..bb0102d0f7d0 100644 --- a/mooncake_disagg/m3_gen_mooncake.yaml +++ b/mooncake_disagg/m3_gen_mooncake.yaml @@ -6,6 +6,22 @@ # # Because no connector runs here, three of the context worker's constraints # do NOT apply, and this file keeps the production values instead. +# +# It does lend the pool this node's host memory. Capacity comes only from +# processes that open a store handle, so without this every byte of the pool +# would be prefill-node DRAM caching prefill's own GPUs -- which is what the +# native host tier already does. Lending is not a connector: this engine still +# never reads or writes the store, and keeps its cache transceiver for the +# prefill-to-decode handoff. +mooncake_donation: + # The master the context server started and published. A server lending + # memory waits for it, so the two can start in any order. + master_server_address: file://__LOG_DIR__/master.addr + # Charged per server process, not per rank as global_segment_size is. Two + # servers on one node lend twice this, and it competes with this node's own + # kv_cache_config.host_cache_size -- size the two together. + segment_size: 640GiB + protocol: rdma max_seq_len: 1048576 max_num_tokens: 16384 diff --git a/mooncake_usage.md b/mooncake_usage.md index 0b506fb269eb..a12b55ddee40 100644 --- a/mooncake_usage.md +++ b/mooncake_usage.md @@ -47,14 +47,25 @@ Three ways to get there, in increasing order of how much you have to arrange: | Deployment | Master | |---|---| -| One `trtllm-serve`, own pool | `mooncake_store: {launch_master: true}` — the server starts it | +| One `trtllm-serve`, own pool — **including the SLURM harness** | `mooncake_store: {launch_master: true}` — the server starts it | | Several engines, or a pool that outlives them | `trtllm-serve mooncake_master --address_file P`, then `mooncake_store: {master_server_address: file://P}` | -| SLURM benchmark harness | Nothing: `disaggr_torch.slurm` starts the master and writes the JSON per job | +| An externally provisioned pool | Nothing in the config: an inherited `MOONCAKE_CONFIG_PATH` wins over `mooncake_store` and says so in the log | The first two make `trtllm-serve` render the client config and export -`MOONCAKE_CONFIG_PATH` itself. An inherited `MOONCAKE_CONFIG_PATH` wins over -both and says so in the log, which is why the harness path is unaffected. -`mooncake_disagg/README.md` §4 covers running a master as its own SLURM job. +`MOONCAKE_CONFIG_PATH` itself. Nothing outside `trtllm-serve` starts a master, +writes a JSON config or picks an HCA — the SLURM harness included, which is why +`disaggr_torch.slurm` now only installs the bindings and tells the configs which +directory the run is in. `mooncake_disagg/README.md` §4 covers running a master +as its own SLURM job, for the second row. + +One thing a scheduler-launched server does need: `TRTLLM_MOONCAKE_RUN_DIR` +pointing somewhere all its ranks can read. Provisioning happens in the server +process and reaches the ranks it spawns through the environment, but under +`trtllm-llmapi-launch` each rank is its own task and was already running, so +those ranks read the rendered config back from that directory instead. Without +it they fail during bringup naming `MOONCAKE_CONFIG_PATH`. `start_worker.sh` +sets it to the job's log directory, which is also where the master's log and +published address land. A launched master dies with the server, so use it only for a single engine: two context servers that each launch one get two disjoint pools, and the @@ -108,20 +119,38 @@ Per-process environment, on the workers that open a handle: Pool capacity comes only from processes that open a store handle, so a prefill-only connector gives a prefill-only pool — which caches prefill's GPUs -in prefill's own DRAM, largely duplicating the native host offload. -`trtllm-serve mooncake_donor` contributes host memory from a node without any -traffic, leaving that engine connector-free: +in prefill's own DRAM, largely duplicating the native host offload. Ask the +generation servers to lend their memory and the pool spans both sides while +their engines stay connector-free: + +```yaml +# generation worker — no connector, memory only +mooncake_donation: + master_server_address: file:///$WORK_DIR/master.addr + segment_size: 320GiB # per server process, not per rank + protocol: rdma + device_name: mlx5_1 +``` + +The context worker publishes the address this reads by adding +`master_address_file: $WORK_DIR/master.addr` next to its `launch_master: true`; +one is written to the run directory regardless. Startup order stops mattering, +because a server lending memory waits for the master rather than needing to +follow it. + +Note the granularity: `global_segment_size` is charged per *rank*, +`segment_size` per *server process*. Two generation servers on one node lend +`segment_size` each. It is charged to the process, so it competes with that +node's own `kv_cache_config.host_cache_size` — size the two together. + +A node running no server can lend as its own command, which is also how the +pool gets memory from a machine with no GPUs at all: ```bash trtllm-serve mooncake_donor --master_server_address file://$WORK_DIR/master.addr \ --segment_size 160GiB --protocol rdma --device_name mlx5_0 ``` -`disaggr_torch.slurm` starts one per generation node already -(`MOONCAKE_DONOR_SEGMENT_SIZE`, default `32GiB`, `0` to skip). Donated memory -is charged to the donor process, so it competes with that node's own -`kv_cache_config.host_cache_size` — size the two together. - ## 3. Partial reuse must be off — now enforced This is the one setting that decides whether the feature works at all. @@ -162,9 +191,9 @@ have aborted startup). Then check that the pool spans the hosts you expect. `disaggr_torch.slurm` writes the per-segment breakdown to `/9_mooncake_summary.log`; a single host means a prefill-only pool. -Pool occupancy and eviction come from `/2_mooncake_master.log`, or -from `$TRTLLM_MOONCAKE_RUN_DIR/mooncake_master.log` when `trtllm-serve` -launched the master — the startup line reports the path either way. +Pool occupancy and eviction come from the master's own log, +`$TRTLLM_MOONCAKE_RUN_DIR/mooncake_master.log` — under the harness that is +`/mooncake_master.log`. The startup line reports the path either way. **Which reuse number counts store hits:** per-request stats (`reused_blocks_per_request`, `kv_cache_hit_rate_per_request`) **do**; @@ -234,6 +263,7 @@ unattributed, with cumulative and per-window reporting. That is what produced §5's split. They are kept off this branch because they cost a store probe on lookups the connector would otherwise decline. -Enable with `MOONCAKE_DEBUG_COVERAGE=1`, which must be set via -`environment.ctx_worker_env_var` — the SLURM harness consumes `MOONCAKE_*` -itself to build the pool config and does not forward it to workers. +Enable with `MOONCAKE_DEBUG_COVERAGE=1`, set via +`environment.ctx_worker_env_var`: `slurm.extra_args` reaches the harness rather +than the worker processes, so a flag the connector reads has to go where the +worker environment is built. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index 0a1b7ec5ff58..cda431c99d61 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -55,9 +55,11 @@ """ from .config import MooncakeStoreConnectorConfig, StoreRole, parse_size -from .donor import DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment +from .donor import (DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, + maybe_donate_segment) from .master import (local_address, master_timeout, maybe_provision_pool, - provision_pool, resolve_master_address, running_master) + provision_pool, resolve_device_name, + resolve_master_address, running_master, wait_for_master) from .scheduler import MooncakeStoreConnectorScheduler from .worker import MooncakeStoreConnectorWorker @@ -70,9 +72,12 @@ "donate_segment", "local_address", "master_timeout", + "maybe_donate_segment", "maybe_provision_pool", "parse_size", "provision_pool", + "resolve_device_name", "resolve_master_address", "running_master", + "wait_for_master", ] diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 09ca8cc0bf3e..d6fb47628c0f 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -31,14 +31,22 @@ from typing import Any, Optional __all__ = [ + "CLIENT_CONFIG_NAME", "CONFIG_PATH_ENV", "MooncakeStoreConnectorConfig", "ROLE_ENV", + "RUN_DIR_ENV", "STAGE_THROUGH_HOST_ENV", "StoreRole", + "provisioned_config_path", ] CONFIG_PATH_ENV = "MOONCAKE_CONFIG_PATH" +#: Where a server keeps the client config it renders and the master's log. Set +#: it to keep them after shutdown; otherwise they live in a temporary directory. +RUN_DIR_ENV = "TRTLLM_MOONCAKE_RUN_DIR" +#: Name the rendered client config takes in the run directory. +CLIENT_CONFIG_NAME = "mooncake.json" ROLE_ENV = "TRTLLM_MOONCAKE_STORE_ROLE" CACHE_PREFIX_ENV = "TRTLLM_MOONCAKE_STORE_PREFIX" MODEL_KEY_ENV = "TRTLLM_MOONCAKE_STORE_MODEL_KEY" @@ -112,6 +120,27 @@ def parse_size(value: Any) -> int: return int(float(magnitude) * scale) +def provisioned_config_path() -> Optional[str]: + """The client config a server on this node rendered, if there is one. + + ``provision_pool`` writes one and exports ``MOONCAKE_CONFIG_PATH``, which + reaches the ranks the LLM constructor spawns, since they inherit that + environment. Ranks the launcher started instead -- one task per rank, which + is how a server spanning several GPUs is launched under a scheduler -- were + already running by then and never see it. Reading the config back from the + run directory is what lets those ranks join the pool their own leader + provisioned. + + Only possible when the deployment named that directory: it otherwise + defaults to a per-process temporary one, which no other rank could read. + """ + run_dir = os.getenv(RUN_DIR_ENV) + if not run_dir: + return None + path = os.path.join(run_dir, CLIENT_CONFIG_NAME) + return path if os.path.exists(path) else None + + @dataclass(frozen=True) class MooncakeStoreConnectorConfig: """Everything needed to open a store handle and name keys in it.""" @@ -186,12 +215,15 @@ def from_file(path: str) -> "MooncakeStoreConnectorConfig": @staticmethod def from_env() -> "MooncakeStoreConnectorConfig": """Load the JSON config, then apply the TensorRT-LLM env overrides.""" - path = os.getenv(CONFIG_PATH_ENV) + path = os.getenv(CONFIG_PATH_ENV) or provisioned_config_path() if not path: raise ValueError( f"The mooncake-store connector needs {CONFIG_PATH_ENV} set to a " "Mooncake JSON config (metadata_server, master_server_address, " - "protocol, device_name, global_segment_size, local_buffer_size)." + "protocol, device_name, global_segment_size, local_buffer_size), " + "or kv_connector_config.mooncake_store set so the server renders " + f"one -- into ${RUN_DIR_ENV} if this rank was started by the " + "launcher rather than spawned by the server." ) config = MooncakeStoreConnectorConfig.from_file(path) return config.with_env_overrides() diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py index dbb105086cfd..b27f6fbfddfe 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -41,13 +41,20 @@ """ import contextlib -from typing import Iterator, Optional +import time +from typing import Any, Iterator, Optional from tensorrt_llm.logger import logger -from .master import local_address +from .config import parse_size +from .master import (local_address, master_timeout, resolve_device_name, + resolve_master_address, wait_for_master) -__all__ = ["DEFAULT_DONOR_LOCAL_BUFFER_SIZE", "donate_segment"] +__all__ = [ + "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", + "donate_segment", + "maybe_donate_segment", +] #: A donor never transfers, so its transfer buffer is dead weight; ``setup`` #: still rejects a zero one. @@ -84,14 +91,22 @@ def donate_segment( ) from exc host = hostname or local_address() + donated = f"{segment_size / 1024 ** 3:.1f}GiB" + # Every argument is echoed, with the byte counts spelled out next to the + # human-readable form: a segment that is a thousandth of the intended size + # is a size string parsed wrong, and it otherwise shows up only as a pool + # that evicts far too eagerly, days later. logger.info( - f"mooncake-store: joining the pool at {master_server_address} as " - f"capacity only: host={host} protocol={protocol} " - f"device={device_name or '(none)'} " - f"donating={segment_size / 1024 ** 3:.1f}GiB" + f"mooncake-store: lending memory to the pool at {master_server_address} " + f"as capacity only, no reads or writes: host={host} " + f"segment_size={donated} ({segment_size} bytes) " + f"protocol={protocol} device={device_name or '(none)'} " + f"metadata_server={metadata_server or '(none)'} " + f"local_buffer_size={local_buffer_size} bytes" ) store = MooncakeDistributedStore() + started = time.monotonic() status = store.setup( host, metadata_server, @@ -101,16 +116,22 @@ def donate_segment( device_name, master_server_address, ) + elapsed = time.monotonic() - started if status != 0: raise RuntimeError( - f"Mooncake store.setup failed with status {status}. The master at " - f"{master_server_address} must already be accepting connections, " - f"and protocol={protocol!r} must be usable from {host}." + f"Mooncake store.setup failed with status {status} after " + f"{elapsed:.1f}s, so no memory was lent to the pool. The master at " + f"{master_server_address} must already be accepting connections; " + f"protocol={protocol!r} with device={device_name or '(none)'!r} " + f"must be usable from {host}; and this node must have " + f"{donated} of memory to spare, which it does not if its own " + "kv_cache_config.host_cache_size has already claimed it." ) logger.info( - f"mooncake-store: {segment_size / 1024 ** 3:.1f}GiB from {host} is now " - "part of the pool" + f"mooncake-store: {donated} of {host} is now part of the pool, " + f"registered in {elapsed:.1f}s; the master at {master_server_address} " + "can place blocks here from now on" ) try: yield host @@ -118,4 +139,48 @@ def donate_segment( # Explicit because the segment stays mounted for as long as anything # references the handle, and "as long as this context" is the contract. del store - logger.info(f"mooncake-store: withdrew the segment donated from {host}") + logger.info( + f"mooncake-store: withdrew the {donated} lent from {host}; the " + "master will report blocks that lived there as lost" + ) + + +@contextlib.contextmanager +def maybe_donate_segment(donation: Any) -> Iterator[Optional[str]]: + """Lend memory for this process's lifetime if the config asked to. + + Args: + donation: A ``MooncakeDonationConfig``, or ``None`` to do nothing -- + which is every deployment that does not lend memory, so callers + need no condition of their own. + + Yields the host the segment is registered under, or ``None``. + """ + if donation is None: + yield None + return + + # A generation server has no other reason to resolve a master, so the + # address it lends against is worth saying out loud before the wait: this + # is the one place bringup blocks on a component from a different job. + logger.info( + "mooncake-store: mooncake_donation is set, so this server lends host " + f"memory to the pool at {donation.master_server_address} without using " + "it; resolving the master now" + ) + master_address = resolve_master_address( + donation.master_server_address, master_timeout() + ) + # Checked before setup so an absent master reads as one, rather than as + # the status code setup returns for everything. + wait_for_master(master_address) + with donate_segment( + master_server_address=master_address, + segment_size=parse_size(donation.segment_size), + protocol=donation.protocol, + # Which HCAs this node has is the node's business, so a config that + # leaves it open stays usable on every node type in the deployment. + device_name=resolve_device_name(donation.protocol, donation.device_name), + metadata_server=donation.metadata_server, + ) as host: + yield host diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index fe244d818d02..442fe3135d47 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -41,37 +41,58 @@ import tempfile import time from dataclasses import dataclass -from typing import Any, Dict, Iterator, Optional, Tuple +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple from tensorrt_llm.logger import logger from ..registry import uses_connector -from .config import CONFIG_PATH_ENV +from .config import CLIENT_CONFIG_NAME, CONFIG_PATH_ENV, RUN_DIR_ENV __all__ = [ "local_address", "maybe_provision_pool", "master_timeout", "provision_pool", + "resolve_device_name", "resolve_master_address", "running_master", + "wait_for_master", ] #: Override the binary that ``launch_master`` runs. MASTER_BINARY_ENV = "TRTLLM_MOONCAKE_MASTER_BINARY" #: How long to wait for a master to accept connections, in seconds. MASTER_TIMEOUT_ENV = "TRTLLM_MOONCAKE_MASTER_TIMEOUT" -#: Where to keep the generated client config and the master's log. Set it to -#: keep them after shutdown; otherwise they live in a temporary directory. -RUN_DIR_ENV = "TRTLLM_MOONCAKE_RUN_DIR" - DEFAULT_MASTER_BINARY = "mooncake_master" DEFAULT_MASTER_TIMEOUT = 60.0 -CLIENT_CONFIG_NAME = "mooncake.json" MASTER_LOG_NAME = "mooncake_master.log" +#: Name a launched master's address is always published under in the run +#: directory, so "which master is this pool on" is answerable from the logs of +#: a run that named no address file. +MASTER_ADDRESS_NAME = "master.addr" #: Prefix that makes ``master_server_address`` name a file holding the address #: rather than the address itself. ADDRESS_FILE_SCHEME = "file://" +#: Lines of the master's log to quote when startup fails. Its last words are +#: usually the whole diagnosis -- a port in use, a bad flag -- and they are +#: otherwise in a file the reader has to be told exists. +LOG_TAIL_LINES = 20 + + +def _log_tail(path: str, lines: int = LOG_TAIL_LINES) -> str: + """The end of the master's log, ready to append to a failure message.""" + try: + with open(path, errors="replace") as handle: + tail = handle.read().splitlines()[-lines:] + except OSError as exc: + return f" Its log at {path} could not be read: {exc}." + if not tail: + return ( + f" Its log at {path} is empty, which usually means it failed " + "before glog opened; check that the binary runs at all." + ) + quoted = "\n ".join(tail) + return f" The last {len(tail)} lines of {path}:\n {quoted}" def local_address() -> str: @@ -122,7 +143,10 @@ def resolve_master_address(address: str, timeout: float) -> str: return address path = address[len(ADDRESS_FILE_SCHEME):] - deadline = time.monotonic() + timeout + started = time.monotonic() + deadline = started + timeout + announced = started + logger.info(f"mooncake-store: reading the master's address from {path}") while True: # Written whole by the master command, so a non-empty file is a # complete address rather than a prefix of one. @@ -133,7 +157,17 @@ def resolve_master_address(address: str, timeout: float) -> str: if published: logger.info(f"mooncake-store: {path} names the master at {published}") return published - if time.monotonic() >= deadline: + now = time.monotonic() + if now - announced >= 5.0: + announced = now + # Waiting for a master in another job is the normal case here, so + # this is progress rather than trouble -- but only if it is said. + logger.info( + f"mooncake-store: no master address in {path} yet " + f"({now - started:.0f}s of {timeout:g}s); waiting for the " + "master to start and publish it" + ) + if now >= deadline: raise TimeoutError( f"No Mooncake master address appeared in {path} within " f"{timeout:g}s. Start one with 'trtllm-serve mooncake_master " @@ -149,35 +183,157 @@ def _wait_until_accepting( port: int, timeout: float, process: Optional[subprocess.Popen] = None, - hint: str = "", -) -> None: - """Block until the master accepts connections. + log_path: Optional[str] = None, +) -> float: + """Block until the master accepts connections, and say how long it took. A worker that opens its store handle before the master is listening fails outright, so the port -- not the presence of a process -- is what the ordering has to wait on. When the master is ours, its exit is checked first each pass, so a master that died is reported as that rather than as a timeout. + + The wait is narrated while it happens: silence here is indistinguishable + from a hang somewhere else in bringup, and this is one of the two places a + Mooncake deployment stalls. """ - deadline = time.monotonic() + timeout + started = time.monotonic() + deadline = started + timeout + announced = started while True: if process is not None and (code := process.poll()) is not None: - raise RuntimeError(f"mooncake_master exited with code {code} during startup.{hint}") + raise RuntimeError( + f"mooncake_master exited with code {code} after " + f"{time.monotonic() - started:.1f}s, before it accepted " + f"connections on {host}:{port}." + f"{_log_tail(log_path) if log_path else ''}" + ) try: with socket.create_connection((host, port), timeout=1.0): - return + return time.monotonic() - started except OSError as exc: last_error = exc - if time.monotonic() >= deadline: + now = time.monotonic() + if now >= deadline: raise TimeoutError( f"The Mooncake master at {host}:{port} did not accept " f"connections within {timeout:g}s ({last_error}). Raise " - f"{MASTER_TIMEOUT_ENV} if it is only slow to start.{hint}" + f"{MASTER_TIMEOUT_ENV} if it is only slow to start." + f"{_log_tail(log_path) if log_path else ''}" + ) + if now - announced >= 5.0: + announced = now + logger.info( + f"mooncake-store: still waiting for the master at {host}:{port}" + f" ({now - started:.0f}s of {timeout:g}s, {last_error})" ) time.sleep(0.5) -def _client_config(pool: Any, master_address: str) -> Dict[str, Any]: +#: Where the InfiniBand devices of a host are described. +IB_SYSFS_ROOT = "/sys/class/infiniband" + + +def _highest_rate_ib_devices(sysfs_root: Optional[str] = None) -> List[str]: + """The active InfiniBand devices on the compute fabric, fastest first. + + A node's HCAs are not interchangeable. On GB300 six are exposed, of which + four run at 800Gb/s -- two per NUMA node, one per GPU -- while the rest + share a PCI device with an Ethernet port and are the storage or management + adapter. Taking every device at the highest rate picks the compute fabric + on any node type, where a hardcoded name would be wrong on the next one. + """ + sysfs_root = sysfs_root or IB_SYSFS_ROOT + rated: Dict[str, int] = {} + try: + devices = sorted(os.listdir(sysfs_root)) + except OSError: + return [] + for device in devices: + port = os.path.join(sysfs_root, device, "ports", "1") + + def attribute(name: str) -> str: + try: + with open(os.path.join(port, name)) as handle: + return handle.read().strip() + except OSError: + return "" + + if attribute("link_layer") != "InfiniBand": + continue + if "ACTIVE" not in attribute("state"): + continue + # "800 Gb/sec (4X XDR)" + rate = attribute("rate").split() + if not rate or not rate[0].isdigit(): + continue + rated[device] = int(rate[0]) + + if not rated: + return [] + fastest = max(rated.values()) + return [device for device, rate in sorted(rated.items()) if rate == fastest] + + +def resolve_device_name(protocol: str, + configured: str, + sysfs_root: Optional[str] = None) -> str: + """The RDMA devices to transfer over, detected if the config left it open. + + Which HCAs a node has is a property of the node, not of the deployment, so + requiring it in a config makes that config specific to one machine type. + Detecting it keeps ``protocol: rdma`` portable, and leaving ``device_name`` + set overrides this for a node where the choice has to be made by hand. + """ + if configured or protocol != "rdma": + return configured + detected = _highest_rate_ib_devices(sysfs_root) + if not detected: + logger.warning( + "mooncake-store: protocol is rdma but no active InfiniBand device " + f"was found under {sysfs_root or IB_SYSFS_ROOT}, so device_name is " + "left empty for Mooncake's own discovery. Set device_name to " + "choose explicitly." + ) + return "" + joined = ",".join(detected) + logger.info( + f"mooncake-store: transferring over the fastest active InfiniBand " + f"devices on this host: {joined}" + ) + return joined + + +def wait_for_master(master_address: str, timeout: Optional[float] = None) -> Optional[float]: + """Block until the master at ``master_address`` accepts connections. + + Every user of a pool it did not start wants this: reaching a master that + is not there otherwise fails deep inside ``store.setup``, in every rank, + after the model has loaded, as a status code. One socket beforehand turns + that into a line that names the address and the wait. + + Returns how long it took, or ``None`` if the address was not in + ``host:port`` form and could not be checked. + """ + timeout = master_timeout() if timeout is None else timeout + endpoint = _split_address(master_address) + if endpoint is None: + logger.warning( + f"mooncake-store: cannot parse master_server_address=" + f"{master_address!r} as host:port, so its reachability is left " + "for the workers to discover." + ) + return None + elapsed = _wait_until_accepting(*endpoint, timeout) + logger.info( + f"mooncake-store: the master at {master_address} answered in {elapsed:.1f}s" + ) + return elapsed + + +def _client_config(pool: Any, + master_address: str, + device_name: Optional[str] = None) -> Dict[str, Any]: """Render the Mooncake client config for a pool. The schema is vLLM's, so one pool can serve both engines. ``role`` is @@ -189,7 +345,7 @@ def _client_config(pool: Any, master_address: str) -> Dict[str, Any]: "metadata_server": pool.metadata_server, "master_server_address": master_address, "protocol": pool.protocol, - "device_name": pool.device_name, + "device_name": pool.device_name if device_name is None else device_name, "global_segment_size": pool.global_segment_size, "local_buffer_size": pool.local_buffer_size, "role": "both", @@ -239,7 +395,6 @@ def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: host = local_address() log_path = os.path.join(run_dir, MASTER_LOG_NAME) - hint = f" See {log_path}." # mooncake_master logs through glog, which writes files under /tmp unless # told otherwise, so without GLOG_logtostderr the log below stays empty. @@ -263,19 +418,67 @@ def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: master = LaunchedMaster( process=process, address=f"{host}:{pool.master_port}", log_path=log_path ) + logger.info( + f"mooncake-store: master pid={process.pid} logging to {log_path} " + f"(GLOG_v={env['GLOG_v']}); waiting for it to accept connections" + ) try: - _wait_until_accepting(host, pool.master_port, master_timeout(), process=process, hint=hint) + elapsed = _wait_until_accepting( + host, pool.master_port, master_timeout(), process=process, log_path=log_path + ) except BaseException: master.stop() raise logger.info( - f"mooncake-store: master ready at {master.address} " + f"mooncake-store: master ready at {master.address} after {elapsed:.1f}s " f"(metrics http://{host}:{pool.master_metrics_port}, log {log_path})" ) return master +@contextlib.contextmanager +def _published_address(address: str, paths: Sequence[str]) -> Iterator[None]: + """Write ``address`` to every path for the life of the context. + + Publishing is how anything else finds this master: a donor or a second + server names the path in ``master_server_address`` as ``file://`` + and reads it back. Retracting on the way out is as important as writing, + since an address that outlives its master sends the next run's workers to + a port with nothing behind it. + """ + for path in paths: + directory = os.path.dirname(path) + if directory: + os.makedirs(directory, exist_ok=True) + # Renamed into place so a reader sees either nothing or the whole + # address. A half-written one would be dialed as if it were real. + staging = f"{path}.partial" + with open(staging, "w") as handle: + handle.write(f"{address}\n") + os.replace(staging, path) + logger.info(f"mooncake-store: published master {address} to {path}") + try: + yield + finally: + for path in paths: + with contextlib.suppress(OSError): + os.remove(path) + logger.info(f"mooncake-store: withdrew the master address at {path}") + + +def _address_files(run_dir: str, extra: Optional[str] = None) -> List[str]: + """Where a master this process starts should publish its address. + + Always the run directory, so a reader who was told nothing can still find + out which master a run used, plus wherever the deployment asked for. + """ + paths = [os.path.join(run_dir, MASTER_ADDRESS_NAME)] + if extra and os.path.abspath(extra) not in {os.path.abspath(p) for p in paths}: + paths.append(extra) + return paths + + @contextlib.contextmanager def running_master( pool: Any, run_dir: str, address_file: Optional[str] = None @@ -288,27 +491,17 @@ def running_master( ``address_file`` receives ``host:port`` once the master answers, so the workers can name the file instead of an address nobody knows until the - scheduler has placed this process. + scheduler has placed this process. One is written to ``run_dir`` either + way. """ os.makedirs(run_dir, exist_ok=True) master = _launch_master(pool, run_dir) try: - if address_file: - # Renamed into place so a reader sees either nothing or the whole - # address. A half-written one would be dialed as if it were real. - staging = f"{address_file}.partial" - with open(staging, "w") as handle: - handle.write(f"{master.address}\n") - os.replace(staging, address_file) - logger.info(f"mooncake-store: published {master.address} to {address_file}") - yield master + with _published_address(master.address, _address_files(run_dir, address_file)): + yield master finally: - if address_file: - # The address outliving the master would send the next run's - # workers to a port with nothing behind it. - with contextlib.suppress(OSError): - os.remove(address_file) master.stop() + logger.info(f"mooncake-store: master at {master.address} stopped") @contextlib.contextmanager @@ -337,50 +530,71 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona keep_run_dir = bool(run_dir or os.getenv(RUN_DIR_ENV)) run_dir = run_dir or os.getenv(RUN_DIR_ENV) or tempfile.mkdtemp(prefix="trtllm-mooncake-") os.makedirs(run_dir, exist_ok=True) + if keep_run_dir: + logger.info(f"mooncake-store: provisioning the pool, run directory {run_dir}") + else: + logger.info( + f"mooncake-store: provisioning the pool in {run_dir}, which is " + f"removed at shutdown along with the master's log; set " + f"{RUN_DIR_ENV} to keep them" + ) master: Optional[LaunchedMaster] = None exported = False - try: - if pool.launch_master: - master = _launch_master(pool, run_dir) - master_address = master.address - else: - master_address = resolve_master_address( - pool.master_server_address, master_timeout() - ) - # Reaching a master that is not there fails inside store.setup on - # every rank, after the model has been loaded. Spend a socket now. - if (endpoint := _split_address(master_address)) is None: - logger.warning( - f"mooncake-store: cannot parse master_server_address=" - f"{master_address!r} as host:port, so its reachability is " - "left for the workers to discover." + with contextlib.ExitStack() as stack: + try: + if pool.launch_master: + master = _launch_master(pool, run_dir) + master_address = master.address + # Even a master that only this server uses publishes: it is + # how its donors reach it, and how the log of a finished run + # still says which pool it was. + stack.enter_context( + _published_address( + master_address, _address_files(run_dir, pool.master_address_file) + ) ) else: - _wait_until_accepting(*endpoint, master_timeout()) - logger.info(f"mooncake-store: using the master at {master_address}") - - config_path = os.path.join(run_dir, CLIENT_CONFIG_NAME) - config = _client_config(pool, master_address) - with open(config_path, "w") as handle: - json.dump(config, handle, indent=2) - # The ranks that open store handles are spawned by the LLM constructor, - # inheriting this environment; that is the only reason exporting it - # here reaches them. - os.environ[CONFIG_PATH_ENV] = config_path - exported = True - logger.info( - f"mooncake-store: {CONFIG_PATH_ENV}={config_path} " - f"({json.dumps(config, sort_keys=True)})" - ) - yield config_path - finally: - if exported: - os.environ.pop(CONFIG_PATH_ENV, None) - if master is not None: - master.stop() - if not keep_run_dir: - shutil.rmtree(run_dir, ignore_errors=True) + master_address = resolve_master_address( + pool.master_server_address, master_timeout() + ) + wait_for_master(master_address) + logger.info(f"mooncake-store: using the master at {master_address}") + + config_path = os.path.join(run_dir, CLIENT_CONFIG_NAME) + config = _client_config( + pool, master_address, + resolve_device_name(pool.protocol, pool.device_name)) + with open(config_path, "w") as handle: + json.dump(config, handle, indent=2) + # This reaches the ranks the LLM constructor spawns, which + # inherit it. A rank the launcher started instead was already + # running, and reads the config out of the run directory -- see + # provisioned_config_path. + os.environ[CONFIG_PATH_ENV] = config_path + exported = True + logger.info( + f"mooncake-store: {CONFIG_PATH_ENV}={config_path} " + f"({json.dumps(config, sort_keys=True)})" + ) + # Capacity is the pool's least obvious property and the one that + # explains a low hit rate, so state the arithmetic rather than + # leaving it to be done from global_segment_size later. + logger.info( + "mooncake-store: this server's ranks will each contribute " + f"global_segment_size={pool.global_segment_size} to the pool; " + "total capacity is that times the number of ranks that open a " + "handle, plus whatever any mooncake_donation adds" + ) + yield config_path + finally: + if exported: + os.environ.pop(CONFIG_PATH_ENV, None) + if master is not None: + master.stop() + logger.info(f"mooncake-store: master at {master.address} stopped") + if not keep_run_dir: + shutil.rmtree(run_dir, ignore_errors=True) @contextlib.contextmanager diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index 814832535c7f..8e1018fa4e0f 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -82,8 +82,14 @@ def stop(signum, _frame): default=None, help="Where to keep the master's log. Defaults to " "$TRTLLM_MOONCAKE_RUN_DIR, else a temporary directory.") +@click.option("--heartbeat_seconds", + type=int, + default=300, + show_default=True, + help="Interval between liveness lines. 0 disables them.") def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, - address_file: Optional[str], run_dir: Optional[str]): + address_file: Optional[str], run_dir: Optional[str], + heartbeat_seconds: int): """Run a mooncake_master for as long as this command runs. For a pool that must not belong to any one engine: several servers sharing @@ -109,6 +115,12 @@ def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, stopping = _until_signalled() with running_master(pool, run_dir, address_file=address_file) as master: + logger.info( + f"mooncake-store: this master owns the pool until this command " + f"stops; address {master.address}, log {master.log_path}, metrics " + f"http://{master.address.rsplit(':', 1)[0]}:{metrics_port}/metrics") + started = time.monotonic() + announced = started while not stopping.is_set(): if (code := master.process.poll()) is not None: # Its own death is the interesting outcome: the pool is gone @@ -117,6 +129,14 @@ def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, f"mooncake_master exited with code {code}. See " f"{master.log_path}") stopping.wait(1.0) + now = time.monotonic() + # Says the pool is still there, which is the question asked of + # this log when clients start failing: master or fabric? + if heartbeat_seconds > 0 and now - announced >= heartbeat_seconds: + announced = now + logger.info( + f"mooncake-store: master at {master.address} alive after " + f"{(now - started) / 60:.0f}m") @click.command("mooncake_donor") @@ -177,7 +197,7 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, """ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, master_timeout, - parse_size, resolve_master_address) + parse_size, resolve_master_address, wait_for_master) from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import \ CONFIG_PATH_ENV @@ -194,9 +214,14 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, f"naming one (or set {CONFIG_PATH_ENV}).") donating = parse_size(segment_size) + resolved = resolve_master_address(master, master_timeout()) + # Before setup, so "the master is not up yet" is reported as that and not + # as the status code setup returns for every kind of failure. + wait_for_master(resolved) + stopping = _until_signalled() with donate_segment( - resolve_master_address(master, master_timeout()), + resolved, donating, protocol=protocol or raw.get("protocol", "rdma"), device_name=device_name or raw.get("device_name", "") or "", @@ -208,6 +233,9 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, if ready_file: with open(ready_file, "w") as handle: handle.write(f"{host} {donating}\n") + logger.info(f"mooncake-store: announced this segment in " + f"{ready_file}, so a launcher waiting on the pool's " + "capacity can proceed") # Idle by design. A put or get here would make this node a client in # the traffic sense, which is the thing keeping the generation engine @@ -218,5 +246,7 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, stopping.wait() continue if not stopping.wait(heartbeat_seconds): - logger.info("mooncake-store: still donating after " - f"{(time.monotonic() - started) / 60:.0f}m") + logger.info( + f"mooncake-store: {host} still lending " + f"{donating / 1024 ** 3:.1f}GiB to the pool at {master} " + f"after {(time.monotonic() - started) / 60:.0f}m") diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index d63059a49d13..38bf2e742acb 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -15,8 +15,7 @@ import time import uuid from pathlib import Path -from typing import (Any, ContextManager, Dict, NamedTuple, Optional, Sequence, - Set) +from typing import Any, Dict, Iterator, NamedTuple, Optional, Sequence, Set import click import torch @@ -42,6 +41,7 @@ parse_metadata_server_config_file, validate_config_bool) from tensorrt_llm.llmapi.llm_args import (KvCacheConnectorConfig, + MooncakeDonationConfig, MultimodalConfig, TorchLlmArgs) from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr, split_mpi_env @@ -518,25 +518,33 @@ def _terminate_attached_frontends(children: list) -> None: child.kill() -def _provision_kv_connector_pool(llm_args: dict, - owns_engine: bool = True) -> ContextManager: - """Bring up the shared cache a KV connector needs, for its lifetime. +@contextlib.contextmanager +def _provision_kv_cache_pool(llm_args: dict, + owns_engine: bool = True) -> Iterator[None]: + """Bring up the shared cache this server needs or feeds, for its lifetime. - Connectors backed by a cluster-wide pool need it reachable before any rank - opens a handle to it, and the ranks are spawned by the LLM constructor. - Entering this around that construction is what makes the pool part of - `trtllm-serve` bringup rather than something a launch script has to - arrange; a deployment that arranges it anyway is detected and left alone. + Two things the engine cannot produce for itself. A connector backed by a + cluster-wide pool needs it reachable before any rank opens a handle, and + the ranks are spawned by the LLM constructor; entering this around that + construction is what makes the pool part of `trtllm-serve` bringup rather + than something a launch script has to arrange. A deployment that arranges + it anyway is detected and left alone. - Only the process that owns the engine provisions anything: an attached - frontend re-execs this command line but shares the launcher's executor, - so it would otherwise stand up a second, private pool. + Separately, a server may lend the pool host memory without using it, which + is how a pool spans nodes whose engines have no connector. That segment + has to be mounted before traffic arrives, and held for as long as the + pages placed in it are expected to be there -- so, this context. + + Only the process that owns the engine does either: an attached frontend + re-execs this command line but shares the launcher's executor, so it would + otherwise stand up a second, private pool and lend a second segment. """ if not owns_engine: - return contextlib.nullcontext() + yield + return - from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import \ - maybe_provision_pool + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( + maybe_donate_segment, maybe_provision_pool) connector_config = llm_args.get("kv_connector_config") if isinstance(connector_config, dict): @@ -545,7 +553,14 @@ def _provision_kv_connector_pool(llm_args: dict, # and hand the validated model on so it is not parsed twice. connector_config = KvCacheConnectorConfig(**connector_config) llm_args["kv_connector_config"] = connector_config - return maybe_provision_pool(connector_config) + + donation = llm_args.get("mooncake_donation") + if isinstance(donation, dict): + donation = MooncakeDonationConfig(**donation) + llm_args["mooncake_donation"] = donation + + with maybe_provision_pool(connector_config), maybe_donate_segment(donation): + yield def launch_server( @@ -604,7 +619,7 @@ def launch_server( raise RuntimeError(f"Failed to bind socket to {host}:{port}: {e}. " f"Port holder(s): {holder}") - with _provision_kv_connector_pool( + with _provision_kv_cache_pool( llm_args, owns_engine=not multi_frontend.is_attached_frontend): if backend == 'pytorch': llm_args.pop("build_config", None) @@ -775,7 +790,7 @@ def signal_handler(): logger.info("Shutdown complete") - with _provision_kv_connector_pool(llm_args): + with _provision_kv_cache_pool(llm_args): uvloop.run(serve_grpc_async()) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index fea3bffb3c67..7313bd98bc39 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1977,6 +1977,14 @@ class MooncakeStoreConfig(StrictBaseModel): "and use it. The pool then dies with the server, so this is only " "correct for a single engine: several engines sharing a pool, or a " "pool that must outlive a restart, need master_server_address.") + master_address_file: Optional[str] = Field( + None, + telemetry=False, + description="Where a master started by launch_master should publish " + "its host:port, for donors and other servers to read back as " + "file://. One is always written to the run directory; set this " + "to put a second copy somewhere the rest of the deployment already " + "names, such as a shared filesystem. Removed when the master stops.") master_port: int = Field( 50051, telemetry=False, @@ -2044,9 +2052,71 @@ def _require_exactly_one_master(self) -> "MooncakeStoreConfig": "mooncake_store: needs a master. Set master_server_address to " "join an existing pool, or launch_master: true to start one " "for this server alone.") + if self.master_address_file and not self.launch_master: + raise ValueError( + "mooncake_store: master_address_file publishes the address of " + "a master this server starts, so it needs launch_master: " + "true. To read an address a master elsewhere published, set " + "master_server_address: file://.") return self +class MooncakeDonationConfig(StrictBaseModel): + """Host memory this server lends to a Mooncake pool it does not use. + + Pool capacity comes only from processes that open a store handle, and in a + disaggregated deployment only the context servers configure the connector. + The pool is then entirely prefill-node memory: prefill's DRAM caching + prefill's GPUs, which largely duplicates what + ``kv_cache_config.host_cache_size`` already does. Setting this on the + generation servers puts their memory into the same pool, so prefill writes + blocks that land on decode-side DRAM, while the generation engine stays + free of any connector and keeps its cache transceiver for the handoff. + + Lending memory is deliberately separate from + ``kv_connector_config``. That config attaches a connector, and a connector + reads and writes; there is no setting on it that means "contribute memory + only", so expressing capacity there would start this server using the + store. Capacity and traffic are different things and are configured + separately. + + The memory is charged to this process and competes with everything else on + the node, ``kv_cache_config.host_cache_size`` above all, so size the two + together. + + Every field opts out of telemetry: they describe one site's pool and how + much of this node was given to it, not which features are in use. + """ + master_server_address: str = Field( + ..., + telemetry=False, + description="Master of the pool to lend memory to, as host:port or " + "file:// naming a file that holds one. The file is what a " + "context server's launch_master publishes, so the generation servers " + "can name a path instead of a host chosen by a scheduler, and they " + "wait for the master rather than having to start after it.") + segment_size: Union[int, str] = Field( + "32GiB", + telemetry=False, + description="Host memory this server contributes. Charged once per " + "server process, not per rank, so a node running several servers " + "contributes this much for each of them.") + protocol: str = Field( + "rdma", + telemetry=False, + description="Transport the pool's traffic reaches this memory over: " + "'rdma' or 'tcp'. Must match the pool's.") + device_name: str = Field( + "", + telemetry=False, + description="RDMA device to serve the segment over, from ibv_devinfo. " + "Empty with protocol 'tcp'.") + metadata_server: str = Field( + "P2PHANDSHAKE", + telemetry=False, + description="Mooncake metadata service. Must match the pool's.") + + class KvCacheConnectorConfig(StrictBaseModel): """Configuration for the KV Cache Connector. @@ -5333,6 +5403,17 @@ def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: status="prototype", ) + mooncake_donation: Optional[MooncakeDonationConfig] = Field( + default=None, + description="Host memory to lend to a Mooncake pool this server does " + "not otherwise use. Separate from kv_connector_config because it adds " + "capacity without attaching a connector, which is what lets a " + "generation server hold pages for a pool only prefill reads and " + "writes. Honored by trtllm-serve, which holds the segment for the " + "server's lifetime.", + status="prototype", + ) + mm_encoder_only: bool = Field( default=False, description= diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index e28692c22f46..0a768e9738b8 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -387,10 +387,50 @@ def test_config_role_comes_from_environment(store_config, monkeypatch): def test_config_requires_the_env_var(monkeypatch): monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_RUN_DIR", raising=False) with pytest.raises(ValueError, match="MOONCAKE_CONFIG_PATH"): MooncakeStoreConnectorConfig.from_env() +def test_config_falls_back_to_the_run_directory(tmp_path, monkeypatch): + # A rank the launcher started was already running when its leader + # provisioned the pool, so it never inherited the exported path and reads + # the rendered config out of the shared run directory instead. + monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + monkeypatch.setenv("TRTLLM_MOONCAKE_RUN_DIR", str(tmp_path)) + (tmp_path / "mooncake.json").write_text( + json.dumps({"master_server_address": "10.0.0.1:50051", "global_segment_size": "8GiB"}) + ) + + config = MooncakeStoreConnectorConfig.from_env() + + assert config.master_server_address == "10.0.0.1:50051" + assert config.global_segment_size == 8 * 1024**3 + + +def test_config_run_directory_without_a_rendered_config_still_asks(tmp_path, monkeypatch): + # An empty run directory means no leader provisioned anything, which is a + # missing pool rather than a default one. + monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + monkeypatch.setenv("TRTLLM_MOONCAKE_RUN_DIR", str(tmp_path)) + with pytest.raises(ValueError, match="MOONCAKE_CONFIG_PATH"): + MooncakeStoreConnectorConfig.from_env() + + +def test_config_env_var_wins_over_the_run_directory(tmp_path, monkeypatch): + # An externally managed pool stays reachable: the run directory is only + # consulted when nothing was passed in. + named = tmp_path / "external.json" + named.write_text(json.dumps({"master_server_address": "external:50051"})) + (tmp_path / "mooncake.json").write_text( + json.dumps({"master_server_address": "provisioned:50051"}) + ) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(named)) + monkeypatch.setenv("TRTLLM_MOONCAKE_RUN_DIR", str(tmp_path)) + + assert MooncakeStoreConnectorConfig.from_env().master_server_address == "external:50051" + + def test_config_model_key_defaults_to_basename(store_config, tmp_path, monkeypatch): path = tmp_path / "no_model_key.json" path.write_text(json.dumps({"master_server_address": "127.0.0.1:50051"})) diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py index 83dcf2f180f4..d8c2b040f1d2 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_donor.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -29,7 +29,9 @@ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.donor import ( DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, + maybe_donate_segment, ) +from tensorrt_llm.llmapi.llm_args import MooncakeDonationConfig GIB = 1024**3 @@ -131,3 +133,82 @@ def test_missing_bindings_are_reported_as_the_separate_component_they_are(monkey with pytest.raises(ImportError, match="mooncake-transfer-engine"): with donate_segment("10.0.0.1:50051", GIB): pytest.fail("donation should not have yielded") + + +@pytest.fixture +def reachable_master(monkeypatch): + """Skip the socket probe: these tests are about what donation asks for.""" + monkeypatch.setattr(donor_module, "wait_for_master", lambda address: 0.0) + + +def test_a_server_that_was_asked_to_lend_memory_does(fake_bindings, reachable_master, + monkeypatch): + """The config-driven path is what makes a generation server a donor.""" + monkeypatch.setattr(donor_module, "local_address", lambda: "10.1.2.3") + donation = MooncakeDonationConfig( + master_server_address="10.0.0.1:50051", + segment_size="320GiB", + protocol="rdma", + device_name="mlx5_1", + ) + + with maybe_donate_segment(donation) as host: + assert host == "10.1.2.3" + registered_host, metadata_server, segment_size, _, protocol, device, master = ( + fake_bindings.instances[0].setup_args) + + assert registered_host == "10.1.2.3" + assert metadata_server == "P2PHANDSHAKE" + # A size string reaching Mooncake unparsed would be a segment of nothing. + assert segment_size == 320 * GIB + assert protocol == "rdma" + assert device == "mlx5_1" + assert master == "10.0.0.1:50051" + + +def test_a_server_that_was_not_asked_lends_nothing(fake_bindings): + """Every deployment that does not lend memory takes this path.""" + with maybe_donate_segment(None) as host: + assert host is None + assert fake_bindings.instances == [] + + +def test_a_published_master_address_is_read_before_joining(fake_bindings, reachable_master, + tmp_path): + """What lets a generation server name a path instead of a scheduler's choice.""" + address_file = tmp_path / "master.addr" + address_file.write_text("10.0.0.9:50051\n") + donation = MooncakeDonationConfig( + master_server_address=f"file://{address_file}", + segment_size=GIB, + ) + + with maybe_donate_segment(donation): + assert fake_bindings.instances[0].setup_args[6] == "10.0.0.9:50051" + + +def test_the_segment_is_withdrawn_when_the_server_stops(fake_bindings, reachable_master): + """The handle is the segment: holding it for the server's life is the point.""" + donation = MooncakeDonationConfig(master_server_address="10.0.0.1:50051") + + with maybe_donate_segment(donation): + store = fake_bindings.instances[0] + # Nothing to assert on the fake beyond its existence -- the contract is + # that the reference is dropped, which is what unmounts the segment. + assert store.setup_args is not None + + +def test_an_unreachable_master_is_reported_before_the_segment_is_offered( + fake_bindings, monkeypatch): + """Otherwise this is a status code from setup, with no address in it.""" + + def refuse(address): + raise TimeoutError(f"The Mooncake master at {address} did not accept connections") + + monkeypatch.setattr(donor_module, "wait_for_master", refuse) + donation = MooncakeDonationConfig(master_server_address="10.0.0.1:50051") + + with pytest.raises(TimeoutError, match="10.0.0.1:50051"): + with maybe_donate_segment(donation): + pytest.fail("donation should not have yielded") + assert fake_bindings.instances == [] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py index f307f4cd448f..b50c40ed74a4 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_master.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -60,6 +60,7 @@ class FakeMasterProcess: def __init__(self, command, env, listen_on=None, exit_code=None): self.command = command self.env = env + self.pid = 4242 self.terminated = False self.killed = False self._exit_code = exit_code @@ -114,9 +115,14 @@ class Launcher: def __init__(self): self.process = None - def arm(self, listen_on=None, exit_code=None): + def arm(self, listen_on=None, exit_code=None, log_text=None): - def popen(command, env=None, **_kwargs): + def popen(command, env=None, stdout=None, **_kwargs): + # A real master writes its own log through glog, and what it + # says there is the diagnosis when it fails to start. + if log_text is not None and stdout is not None: + stdout.write(log_text.encode()) + stdout.flush() self.process = FakeMasterProcess( command, env, listen_on=listen_on, exit_code=exit_code ) @@ -164,6 +170,15 @@ def test_pool_needs_exactly_one_master(): MooncakeStoreConfig() +def test_publishing_an_address_needs_a_master_to_publish(): + """Reading a published address is master_server_address, not this.""" + with pytest.raises(ValueError, match="needs launch_master"): + MooncakeStoreConfig( + master_server_address="host:50051", + master_address_file="/shared/master.addr", + ) + + def test_pool_is_rejected_on_another_connector(): with pytest.raises(ValueError, match="mooncake_store describes a Mooncake pool"): KvCacheConnectorConfig( @@ -528,3 +543,140 @@ def test_provisioning_joins_a_master_it_was_never_given_the_address_of(fake_mast # has to be the address it resolved to. written = json.loads(open(config_path).read()) assert written["master_server_address"] == master.address + + +# ---- a master a server launched, made findable ---- + + +def test_a_launched_master_publishes_where_its_run_left_its_logs(fake_master, tmp_path): + """So a finished run's logs still say which pool it used.""" + port = free_port() + fake_master.arm(listen_on=port) + run_dir = tmp_path / "run" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with provision_pool(pool, run_dir=str(run_dir)): + address = (run_dir / master_module.MASTER_ADDRESS_NAME).read_text().strip() + assert address.endswith(f":{port}") + + assert not (run_dir / master_module.MASTER_ADDRESS_NAME).exists() + + +def test_a_launched_master_can_be_published_where_the_donors_look(fake_master, tmp_path): + """The whole reason a server launching a master can still have donors.""" + port = free_port() + fake_master.arm(listen_on=port) + shared = tmp_path / "shared" / "master.addr" + pool = MooncakeStoreConfig( + launch_master=True, master_port=port, master_address_file=str(shared) + ) + + with provision_pool(pool, run_dir=str(tmp_path / "run")): + assert resolve_master_address(f"file://{shared}", timeout=5.0).endswith(f":{port}") + + # Retracted, so the next run's donors wait for a live master rather than + # joining a pool that no longer exists. + assert not shared.exists() + + +def test_a_half_written_address_is_never_read(tmp_path): + """A reader sees the whole address or nothing, never a prefix of one.""" + target = tmp_path / "master.addr" + + with master_module._published_address("10.0.0.7:50051", [str(target)]): + assert not (tmp_path / "master.addr.partial").exists() + assert target.read_text().strip() == "10.0.0.7:50051" + + +# ---- saying why bringup is stuck ---- + + +def test_an_absent_master_is_named_rather_than_left_to_store_setup(monkeypatch): + """Otherwise this is a status code, in every rank, after the model loads.""" + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") + address = f"127.0.0.1:{free_port()}" + + with pytest.raises(TimeoutError, match=address): + master_module.wait_for_master(address) + + +def test_a_master_that_answers_is_reported_with_the_wait_it_cost(running_master): + assert master_module.wait_for_master(running_master) is not None + + +def test_an_address_of_a_shape_we_cannot_probe_is_not_fatal(): + """Mooncake may accept addresses this cannot dial; leave them to it.""" + assert master_module.wait_for_master("unix:///var/run/mooncake") is None + + +def test_a_master_that_died_starting_is_reported_with_its_last_words(fake_master, tmp_path): + """The reason is in its log, which nobody reads unless it is quoted.""" + run_dir = tmp_path / "run" + fake_master.arm( + exit_code=1, log_text="E0903 bind(50051) failed: Address already in use\n") + pool = MooncakeStoreConfig(launch_master=True, master_port=free_port()) + + with pytest.raises(RuntimeError, match="Address already in use"): + with provision_pool(pool, run_dir=str(run_dir)): + pytest.fail("provisioning should not have yielded") + + +# ---- choosing the fabric without naming it in a config ---- + + +def fake_hca(root, device, link_layer="InfiniBand", state="4: ACTIVE", rate="800 Gb/sec"): + port = root / device / "ports" / "1" + port.mkdir(parents=True) + (port / "link_layer").write_text(f"{link_layer}\n") + (port / "state").write_text(f"{state}\n") + (port / "rate").write_text(f"{rate}\n") + + +def test_the_compute_fabric_is_picked_over_the_management_adapter(tmp_path): + """A node's HCAs are not interchangeable: only some are the fast fabric.""" + fake_hca(tmp_path, "mlx5_0") + fake_hca(tmp_path, "mlx5_1") + fake_hca(tmp_path, "mlx5_2", rate="400 Gb/sec") + fake_hca(tmp_path, "mlx5_3", state="1: DOWN") + fake_hca(tmp_path, "mlx5_4", link_layer="Ethernet") + + assert master_module.resolve_device_name( + "rdma", "", sysfs_root=str(tmp_path)) == "mlx5_0,mlx5_1" + + +def test_a_named_device_is_not_second_guessed(tmp_path): + fake_hca(tmp_path, "mlx5_0") + assert master_module.resolve_device_name( + "rdma", "mlx5_7", sysfs_root=str(tmp_path)) == "mlx5_7" + + +def test_tcp_needs_no_device_and_looks_for_none(tmp_path): + assert master_module.resolve_device_name("tcp", "", sysfs_root=str(tmp_path)) == "" + + +def test_a_node_without_infiniband_is_left_to_mooncake_s_own_discovery(tmp_path): + """Better than failing: Mooncake may still find something usable.""" + assert master_module.resolve_device_name( + "rdma", "", sysfs_root=str(tmp_path / "absent")) == "" + + +def test_the_detected_device_is_what_the_workers_are_told(fake_master, tmp_path, + monkeypatch): + sysfs = tmp_path / "sysfs" + fake_hca(sysfs, "mlx5_0") + monkeypatch.setattr(master_module, "IB_SYSFS_ROOT", str(sysfs)) + port = free_port() + fake_master.arm(listen_on=port) + pool = MooncakeStoreConfig(launch_master=True, master_port=port, protocol="rdma") + + with provision_pool(pool, run_dir=str(tmp_path / "run")) as config_path: + assert json.loads(open(config_path).read())["device_name"] == "mlx5_0" + + +def test_an_empty_master_log_says_what_that_means(tmp_path): + """Empty means it failed before glog opened, which reads as no log at all.""" + empty = tmp_path / "mooncake_master.log" + empty.write_text("") + + assert "empty" in master_module._log_tail(str(empty)) + assert "could not be read" in master_module._log_tail(str(tmp_path / "absent.log")) diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 1f5e97d1b74f..f2aefd22f9c6 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -243,6 +243,10 @@ methods: annotation: Optional[tensorrt_llm.llmapi.llm_args.KvCacheConnectorConfig] default: null status: prototype + mooncake_donation: + annotation: Optional[tensorrt_llm.llmapi.llm_args.MooncakeDonationConfig] + default: null + status: prototype enable_lm_head_tp_in_adp: annotation: bool default: False From fb3cc71b898ddf4aa3a99f9b5938fbe5ad1cd330 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:01:40 -0700 Subject: [PATCH 17/24] [None][chore] Polish mooncake-store docs, config handling, and tests Tighten the connector prose and error messages, and drop unit tests that only asserted a default value, a non-None return, or message wording. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/common/install_mooncake.sh | 57 +++--- docs/source/developer-guide/overview.md | 21 ++ docs/source/features/kv-cache-connector.md | 38 ++-- .../slurm/benchmark/disaggr_torch.slurm | 48 +++-- .../slurm/benchmark/start_worker.sh | 33 ++- .../slurm/benchmark/watch_job.sh | 12 +- ...trtllm_mooncake_store_connector_extra.yaml | 4 +- mooncake_disagg/README.md | 168 +++++++-------- mooncake_disagg/gen_config.yaml | 4 +- mooncake_disagg/install_mooncake_runtime.sh | 61 ++---- mooncake_disagg/m3_agg_mooncake.yaml | 2 +- mooncake_disagg/m3_ctx_mooncake.yaml | 6 +- mooncake_disagg/m3_gen_mooncake.yaml | 6 +- mooncake_disagg/mooncake_api_surface_test.py | 15 +- mooncake_disagg/mooncake_smoke_test.py | 6 +- mooncake_usage.md | 74 +++---- .../connectors/mooncake_store/__init__.py | 36 ++-- .../connectors/mooncake_store/addressing.py | 22 +- .../connectors/mooncake_store/config.py | 38 ++-- .../connectors/mooncake_store/donor.py | 80 +++----- .../connectors/mooncake_store/keys.py | 8 +- .../connectors/mooncake_store/master.py | 192 ++++++++---------- .../connectors/mooncake_store/metadata.py | 8 +- .../connectors/mooncake_store/scheduler.py | 18 +- .../connectors/mooncake_store/staging.py | 71 +++---- .../connectors/mooncake_store/validation.py | 2 +- .../connectors/mooncake_store/worker.py | 66 +++--- .../_torch/pyexecutor/connectors/registry.py | 6 +- .../_torch/pyexecutor/hang_detector.py | 48 +++++ .../_torch/pyexecutor/kv_cache_manager_v2.py | 58 +++--- tensorrt_llm/_torch/pyexecutor/py_executor.py | 8 +- .../_torch/pyexecutor/py_executor_creator.py | 4 +- .../pyexecutor/scheduler/scheduler_v2.py | 68 +++---- tensorrt_llm/commands/mooncake.py | 48 ++--- tensorrt_llm/commands/serve.py | 37 ++-- tensorrt_llm/llmapi/llm_args.py | 58 +++--- .../executor/test_hang_detector_kill.py | 76 +++++++ .../executor/test_kv_cache_v2_scheduler.py | 19 +- .../executor/test_mooncake_store_connector.py | 62 +++--- .../executor/test_mooncake_store_donor.py | 26 +-- .../executor/test_mooncake_store_master.py | 80 +++----- 41 files changed, 838 insertions(+), 856 deletions(-) diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index ccd5e04de1d9..a05d08c1f7b7 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -51,38 +51,29 @@ rm -rf Mooncake echo "export LD_LIBRARY_PATH=${MOONCAKE_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" -# The source build above is only useful for the C++ transfer engine, which is -# what the cache transceiver links against. MooncakeDistributedStore -- the -# shared CPU pool behind the mooncake-store KV cache connector -- comes from the -# Python wheel instead, for two reasons. +# The source build above provides only the C++ transfer engine, which is what +# the cache transceiver links against. MooncakeDistributedStore, the shared CPU +# pool behind the mooncake-store KV cache connector, comes from the Python +# wheel instead, for two reasons. # -# First, `make install` does emit a `mooncake` Python package, but an unusable -# one: it omits libmooncake_store.so, so importing mooncake.store raises -# ImportError. It must be deleted, and deleting it is not optional in either of -# the two places it can land. +# First, `make install` emits a `mooncake` Python package that omits +# libmooncake_store.so, so importing mooncake.store raises ImportError. It has +# to be removed wherever it landed, and where that is depends on the +# environment: mooncake-integration/CMakeLists.txt picks its install directory +# as the first sys.path entry whose name merely contains "packages". # -# mooncake-integration/CMakeLists.txt chooses its install directory with -# python3 -c "import sys; print([s for s in sys.path if 'packages' in s][0])" -# i.e. the first sys.path entry whose name merely contains "packages". +# - With nvidia-cutlass-dsl installed, that is +# nvidia_cutlass_dsl/dsl_packages, which nvidia_cutlass_dsl_packages.pth +# puts at sys.path[0], so it shadows anything pip installs. CUTLASS DSL +# does not reference `mooncake`, so removing the package is safe. +# - Without it, the package lands in dist-packages and collides with the +# wheel: CMake writes store.cpython-312-x86_64-linux-gnu.so, the wheel +# writes store.so, and importlib prefers the interpreter-tagged suffix, so +# the broken extension wins even after pip reports success. # -# - With nvidia-cutlass-dsl installed (the normal case here: the devel stage -# removes it, then constraints.txt reinstalls it), that first match is -# nvidia_cutlass_dsl/dsl_packages, because nvidia_cutlass_dsl_packages.pth -# does sys.path.insert(0) on it. The broken package then outranks -# dist-packages on every interpreter start, so no amount of pip installing -# can fix the import. CUTLASS DSL does not reference `mooncake` at all, so -# removing it is safe. -# - Without it, the match is dist-packages itself, and the broken package -# collides with the wheel. That is the more insidious case: CMake writes -# store.cpython-312-x86_64-linux-gnu.so while the wheel writes store.so, and -# importlib prefers the interpreter-tagged suffix, so the broken extension -# still wins even after pip reports success. -# -# Remove the package outright wherever it landed, before pip installs the real -# one. Nothing legitimate owns a `mooncake` package at this point, and removing -# rather than trying to identify individual leftovers keeps this correct in the -# dist-packages case, where pip would overwrite __init__.py and leave no marker -# to key on. +# Remove the directory outright rather than trying to identify leftovers, since +# pip overwrites __init__.py in the collision case and leaves no marker to key +# on. python3 - <<'PY' import os import shutil @@ -99,12 +90,12 @@ for entry in list(sys.path) + [paths["purelib"], paths["platlib"]]: shutil.rmtree(package, ignore_errors=True) PY -# Second, the `mooncake-transfer-engine` wheel is built against CUDA 12 and +# Second, the `mooncake-transfer-engine` wheel is built against CUDA 12 while # these images ship CUDA 13 only, so its extensions cannot resolve # libcudart.so.12. `mooncake-transfer-engine-cuda13` is the same project built -# for CUDA 13. It is versioned independently and its releases start at 0.3.9, so -# it cannot track MOONCAKE_VERSION above; the store client only has to agree -# with the mooncake_master it connects to, and the wheel supplies both. +# for CUDA 13. It is versioned independently, with releases starting at 0.3.9, +# so it cannot track MOONCAKE_VERSION above. The store client only has to agree +# with the mooncake_master it connects to, and this wheel supplies both. MOONCAKE_WHEEL_VERSION="0.3.13" pip3 install --no-cache-dir "mooncake-transfer-engine-cuda13==${MOONCAKE_WHEEL_VERSION}" diff --git a/docs/source/developer-guide/overview.md b/docs/source/developer-guide/overview.md index d8fe31631612..bd8d8cd5b77c 100644 --- a/docs/source/developer-guide/overview.md +++ b/docs/source/developer-guide/overview.md @@ -122,3 +122,24 @@ export TLLM_LOG_LEVEL_BY_MODULE="debug:_torch,runtime;info:serve" ``` This example sets the global level to `warning` but enables `debug` output for `_torch` and `runtime` modules, and `info` for `serve`. Valid levels: `trace`, `debug`, `verbose`, `info`, `warning`, `error`, `internal_error`. + +### Diagnosing Slow Iterations + +`TRTLLM_STALL_REPORT_SEC` dumps every thread's stack to stderr whenever a single +executor iteration takes longer than the given number of seconds: + +```bash +# Any iteration slower than 5s dumps all thread stacks. +export TRTLLM_STALL_REPORT_SEC=5 +``` + +Set it somewhat above the normal iteration time, which the `host_step_time` field +of the per-iteration log line reports. Unlike the hang detector, this neither +kills the process nor stops the run, so it is safe to leave on for a whole +benchmark. It stays silent while iterations are under the threshold and does not +report time the loop spends idle waiting for requests. + +Background threads are included in the dump, which is usually the point: a slow +iteration that is not blocked in any obvious call is often waiting on a KV +transfer, connector or sampler thread that the main thread's stack does not +explain. diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index eff31a6addfe..6af39a1d79bf 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -29,7 +29,7 @@ These methods run on the leader process and drive the connector's behavior. * **`build_connector_meta(self, scheduler_output: SchedulerOutput) -> object`** * **Description**: The core orchestration method. Called during the scheduling phase. It examines the current requests and decides which blocks need to be loaded from or saved to the external store. - * **Arguments**: `scheduler_output` contains information about new requests, blocks allocated, current request states, and the cumulative `RequestData.block_hashes` chain. `block_hashes` is read directly from each KV cache block's stored hash, which the KV cache manager commits as soon as a block becomes full -- the value matches the hash that KV cache events will subsequently emit for the same block. The chain only covers beam 0; the executor rejects `kv_connector_config` at startup when `max_beam_width > 1`, so connectors may assume beam-width-1 inputs. + * **Arguments**: `scheduler_output` contains information about new requests, blocks allocated, current request states, and the cumulative `RequestData.block_hashes` chain. `block_hashes` is read directly from each KV cache block's stored hash, which the KV cache manager commits as soon as a block becomes full, so the value matches the hash that KV cache events will subsequently emit for the same block. The chain only covers beam 0; the executor rejects `kv_connector_config` at startup when `max_beam_width > 1`, so connectors may assume beam-width-1 inputs. * **Returns**: An arbitrary metadata object (picklable) that describes the tasks for the workers. This object is broadcasted to all workers. * **`get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> tuple[int, bool]`** @@ -47,14 +47,14 @@ These methods run on the leader process and drive the connector's behavior. * **`cancel_load(self, request: LlmRequest, start: int, end: int)`** * **Description**: Optional, with a no-op default. Tells the connector that the runtime will not consume KV it offered from `get_num_new_matched_tokens` for prompt tokens `[start, end)`, so any ownership taken for that range can be released. Offsets are absolute prompt positions, on the same scale as `num_computed_tokens`. - * **When it fires**: only on `KVCacheManagerV2`, which asks during a speculative scheduling pass and resolves the answer later. Two things can happen in between, and both are reported here: the runtime may fail to allocate pages to cover the offer, in which case the request falls back to computing the prefix locally; or the request may be cancelled, time out or fail before it ever reaches a batch, in which case the whole offer is released. A third case -- the local cache overtaking part of the offer because another request committed the same prefix -- is handled by the same callback but cannot arise today, since a request's local match is fixed when its cache is created and only its own completed forward passes extend it. + * **When it fires**: only on `KVCacheManagerV2`, which asks during a speculative scheduling pass and resolves the answer later. Two things can happen in between, and both are reported here: the runtime may fail to allocate pages to cover the offer, in which case the request falls back to computing the prefix locally; or the request may be cancelled, time out or fail before it ever reaches a batch, in which case the whole offer is released. A third case, the local cache overtaking part of the offer because another request committed the same prefix, is handled by the same callback but cannot arise today, since a request's local match is fixed when its cache is created and only its own completed forward passes extend it. * **Caveat**: best-effort. For a synchronous load nothing has been transferred yet, so cancelling is exact. For `is_async=True` the transfer necessarily started inside `get_num_new_matched_tokens`, so it may already be in flight. ##### Serving a prefix on `KVCacheManagerV2` -V1 answers `get_num_new_matched_tokens` from C++ while the block manager holds its radix-tree mutex, so the local match and the query are atomic and the answer is consumed immediately. V2 has no such mutex, and its scheduling pass is speculative -- a prepared request can still be dropped at the token budget, at resize, at multimodal alignment or at cross attention, and retried in a later iteration. +V1 answers `get_num_new_matched_tokens` from C++ while the block manager holds its radix-tree mutex, so the local match and the query are atomic and the answer is consumed immediately. V2 has no such mutex, and its scheduling pass is speculative: a prepared request can still be dropped at the token budget, at resize, at multimodal alignment or at cross attention, and retried in a later iteration. -The contract for connectors is unchanged, and in particular `get_num_new_matched_tokens` is still called **exactly once per request** on both managers -- a request that is asked and then deferred is not asked again when it comes back. What differs is that on V2 the runtime may resolve the answer in a later iteration than the one it asked in, and may by then be unable to honour part or all of it. That is what `cancel_load` reports. +`get_num_new_matched_tokens` is still called **exactly once per request** on both managers, so a request that is asked and then deferred is not asked again when it comes back. What differs is that on V2 the runtime may resolve the answer in a later iteration than the one it asked in, and may by then be unable to honour part or all of it. That is what `cancel_load` reports. #### 2. Worker Interface (`KvCacheConnectorWorker`) @@ -66,7 +66,7 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **`register_kv_cache_layout(self, layout: KvCacheLayout)`** * **Description**: Called at initialization **instead of** `register_kv_caches` when the KV cache manager is `KVCacheManagerV2`, whose memory cannot be expressed as one tensor: there is one slot address space per pool and one page-index space per layer group. The default implementation raises, so a connector that does not implement it can only run on V1. - * **Arguments**: `layout` describes the byte ranges that repeat per page slot. Each `KvCacheLayerGroupLayout` carries a tuple of `KvCacheRegion`s, and the bytes for page slot `i` of a region live at `region.base + region.stride * i` for `region.size` bytes -- or equivalently at `region.as_tensor()[i]`. Page indices arriving in `RequestData.new_block_ids_by_layer_group` are scoped to a layer group and index that group's regions. + * **Arguments**: `layout` describes the byte ranges that repeat per page slot. Each `KvCacheLayerGroupLayout` carries a tuple of `KvCacheRegion`s, and the bytes for page slot `i` of a region live at `region.base + region.stride * i` for `region.size` bytes, or equivalently at `region.as_tensor()[i]`. Page indices arriving in `RequestData.new_block_ids_by_layer_group` are scoped to a layer group and index that group's regions. * **Why regions rather than a tensor**: because the ranges are described rather than implied, the same structure covers MLA (a pool simply has no `value` buffer), sliding-window and hybrid models (one layer group per window size), and non-uniform slots such as MiniMax-M3's index-K buffer sitting beside K/V, without any of them being a special case. * **`start_load_kv(self, stream: torch.cuda.Stream)`** @@ -100,7 +100,7 @@ The available presets are `lmcache`, `lmcache-mp`, `kvbm` and `mooncake-store`. ### Mooncake distributed store (`mooncake-store`) -Publishes KV pages into a [Mooncake](https://github.com/kvcache-ai/Mooncake) store -- a shared CPU memory pool addressed by content -- so a prefix computed by one engine can be replayed by another. Regular block reuse cannot do this, because it never leaves the instance that computed the prefix. +Publishes KV pages into a [Mooncake](https://github.com/kvcache-ai/Mooncake) store, a shared CPU memory pool addressed by content, so a prefix computed by one engine can be replayed by another. Regular block reuse cannot do this, because it never leaves the instance that computed the prefix. This is a **different component** from the Mooncake transfer engine that the C++ cache transceiver uses for disaggregated prefill/decode handoff. That moves KV point to point between two known peers; this publishes pages into a pool that any peer can read. The two compose: a context server can write pages into the store and still hand off to a generation server over NIXL. @@ -128,7 +128,7 @@ kv_connector_config: Replacing `master_server_address` with `launch_master: true` makes the server start a `mooncake_master` itself and use it, so a single-instance deployment needs nothing prepared outside `trtllm-serve`. **That master lives and dies with the server**, which makes it wrong for anything else: several engines that should share one pool would each get their own, and a pool meant to survive a restart cannot be owned by the thing restarting. -A master started this way still publishes its address, to `master.addr` in the run directory and to `master_address_file` if one is named. That is what lets other processes find a pool this server owns -- the donors below, most of all -- and what makes a finished run's logs say which master it used: +A master started this way still publishes its address, to `master.addr` in the run directory and to `master_address_file` if one is named. That is how other processes, the donors below above all, find a pool this server owns, and how a finished run's logs say which master it used: ```yaml mooncake_store: @@ -136,7 +136,7 @@ mooncake_store: master_address_file: /shared/master.addr ``` -The two cases above -- several engines, or surviving a restart -- run the master as its own command instead: +The two cases above, several engines or surviving a restart, run the master as its own command instead: ```bash trtllm-serve mooncake_master --rpc_port 50051 --address_file /shared/master.addr @@ -153,11 +153,11 @@ kv_connector_config: This is what makes a master reachable without anyone writing its address down. Under a scheduler its host is not known when the configs are written; publishing it to a file the configs already name closes that gap, and a server reading the file waits for it, so the master and the engines can be started in any order. The file is removed when the master stops, so a stale address is never dialed. -`TRTLLM_MOONCAKE_MASTER_BINARY` overrides the binary a launched master runs, and `TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) how long startup waits for any master to accept connections or publish its address -- reaching a master that is not there otherwise fails inside every rank after the model has loaded. Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated client config and the master's log, which are otherwise in a temporary directory removed at shutdown. +`TRTLLM_MOONCAKE_MASTER_BINARY` overrides the binary a launched master runs, and `TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) sets how long startup waits for any master to accept connections or publish its address. Without that wait, a master that is not there yet fails inside every rank after the model has loaded. Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated client config and the master's log, which are otherwise in a temporary directory removed at shutdown. #### Servers whose ranks the launcher starts -Provisioning happens in the server process and reaches the ranks that open store handles by exporting `MOONCAKE_CONFIG_PATH` for them to inherit. That holds when the LLM constructor spawns them, and does not when the launcher starts one task per rank -- as `trtllm-llmapi-launch` under a scheduler does -- because those ranks were already running. +Provisioning happens in the server process and reaches the ranks that open store handles by exporting `MOONCAKE_CONFIG_PATH` for them to inherit. That holds when the LLM constructor spawns them. It does not when the launcher starts one task per rank, as `trtllm-llmapi-launch` under a scheduler does, because those ranks were already running. Naming a shared run directory covers that case: the rendered config is read back from `$TRTLLM_MOONCAKE_RUN_DIR/mooncake.json` by any rank that inherited no path, so every rank of a multi-GPU server joins the pool its own leader provisioned. The directory has to be one they all see, which under a scheduler means the job's own, and it is where the master's log and published address already go: @@ -170,18 +170,18 @@ Without it, a rank that inherited nothing fails during bringup naming `MOONCAKE_ #### Reading bringup in the log -Everything the pool is assembled from is logged under the `mooncake-store:` prefix before the model loads, because a pool that came up wrong is otherwise visible only as a low hit rate hours later. In order: the run directory, the master's command line and pid, the address it published and where, the rendered client config in full, and the capacity each rank will contribute. A server lending memory logs the master it resolved, the segment in both GiB and bytes, and the transport -- a size string parsed wrong is otherwise invisible until the pool starts evicting far too eagerly. +Everything the pool is assembled from is logged under the `mooncake-store:` prefix before the model loads, because a pool that came up wrong is otherwise visible only as a low hit rate hours later. In order: the run directory, the master's command line and pid, the address it published and where, the rendered client config in full, and the capacity each rank will contribute. A server lending memory logs the master it resolved, the segment in both GiB and bytes, and the transport, so that a size string parsed wrong is caught before the pool starts evicting far too eagerly. -Both waits narrate themselves every five seconds, since waiting for a master in another job is normal and indistinguishable from a hang if it is silent. A master that dies during startup has the tail of its own log quoted in the failure, which is where the reason (a port in use, a bad flag) actually is. +Both waits report progress every five seconds, since waiting for a master in another job is normal and indistinguishable from a hang if it is silent. A master that dies during startup has the tail of its own log quoted in the failure, which is where the reason, a port in use or a bad flag, actually is. #### Pool capacity -Capacity comes only from processes that open a store handle, and `global_segment_size` is what each contributes -- so the pool is that value times the number of such processes. In a disaggregated deployment the connector belongs on the context servers only, which makes every byte of the pool prefill-node memory: prefill's DRAM caching prefill's GPUs, largely duplicating what `kv_cache_config.host_cache_size` already does. +Capacity comes only from processes that open a store handle, and `global_segment_size` is what each contributes, so the pool is that value times the number of such processes. In a disaggregated deployment the connector belongs on the context servers only, which makes every byte of the pool prefill-node memory: prefill's DRAM caching prefill's GPUs, largely duplicating what `kv_cache_config.host_cache_size` already does. To give the pool memory from nodes whose engines run no connector, ask those servers to lend it: ```yaml -# generation server -- no connector, memory only +# generation server: no connector, memory only mooncake_donation: master_server_address: file:///shared/master.addr segment_size: 320GiB @@ -191,7 +191,7 @@ mooncake_donation: `trtllm-serve` then holds that segment for as long as the server runs, so a generation node holds pages prefill wrote while its own engine stays connector-free and keeps its cache transceiver for the prefill-to-decode handoff. The server is ready only once the segment is mounted, which makes its readiness the signal that the pool has this capacity. -Lending memory is deliberately outside `kv_connector_config`, and not a `TRTLLM_MOONCAKE_STORE_ROLE` either. Both of those attach a connector, and a connector reads or writes -- `producer`, `consumer` and `both` all describe traffic, and none of them means "contribute memory only" -- so expressing capacity there would start this server using the store. Capacity and traffic are separate, and configured separately. +Lending memory is deliberately outside `kv_connector_config`, and not a `TRTLLM_MOONCAKE_STORE_ROLE` either. Both of those attach a connector, and a connector reads or writes: `producer`, `consumer` and `both` all describe traffic, and none of them means "contribute memory only". Expressing capacity there would therefore start this server using the store. Capacity and traffic are separate, and configured separately. Size is charged **per server process, not per rank**, unlike `global_segment_size`. Two servers on one node lend twice this. The memory is charged to the process and competes with everything else on the node, `kv_cache_config.host_cache_size` above all, so size the two together. @@ -215,7 +215,7 @@ Topology can equally come from a JSON file named by `MOONCAKE_CONFIG_PATH`, usin } ``` -An inherited `MOONCAKE_CONFIG_PATH` wins over `mooncake_store` and is logged as doing so, so an orchestrator that already provisions the pool -- as the SLURM benchmark harness does -- keeps working unchanged. +An inherited `MOONCAKE_CONFIG_PATH` wins over `mooncake_store` and is logged as doing so, so an orchestrator that already provisions the pool, as the SLURM benchmark harness does, keeps working unchanged. Three further settings are TensorRT-LLM's rather than Mooncake's, and stay in the environment because they are per process rather than per pool: @@ -223,7 +223,7 @@ Three further settings are TensorRT-LLM's rather than Mooncake's, and stay in th |---|---|---| | `TRTLLM_MOONCAKE_STORE_ROLE` | `both` | `producer` writes only, `consumer` reads only, `both` does both. | | `TRTLLM_MOONCAKE_STORE_PREFIX` | `trtllm` | Leading component of every key, for isolating deployments that share a pool. | -| `TRTLLM_MOONCAKE_STORE_MODEL_KEY` | model directory basename | Identity keys are namespaced by. Two engines share cache only when they agree on it, so the default is the basename rather than the full path -- the same checkpoint is routinely mounted elsewhere on another host, which is exactly what sharing is for. | +| `TRTLLM_MOONCAKE_STORE_MODEL_KEY` | model directory basename | Identity keys are namespaced by. Two engines share cache only when they agree on it, so the default is the basename rather than the full path, since the same checkpoint is routinely mounted elsewhere on another host, which is exactly what sharing is for. | In a disaggregated deployment, run context servers as `both` and leave generation servers unconfigured. Generated tokens are rarely a reused prefix, so writing them costs bandwidth for no hit rate. @@ -231,7 +231,7 @@ In a disaggregated deployment, run context servers as `both` and leave generatio `kv_cache_config.enable_partial_reuse` is set to `false` when this connector is configured, with a warning, whether or not it was requested explicitly. It defaults to `true`, so most deployments will see that warning. -The store is addressed by whole blocks. The connector is handed the device match as `num_computed_tokens` and offers only blocks beyond it, but it can only resume from a block boundary -- so when the device match ends mid-block, it declines the lookup and the store is not consulted at all. Partial reuse is precisely what puts the match off a boundary, which means it trades part of one block of device reuse for every stored block of the remaining prefix. Measured on MiniMax-M3, leaving it enabled declined 97.2% of lookups and left actual prompt cache read at 35% against a 96% ceiling; forcing it off raised that to 94% and roughly doubled throughput. +The store is addressed by whole blocks. The connector is handed the device match as `num_computed_tokens` and offers only blocks beyond it, but it can resume only from a block boundary, so when the device match ends mid-block it declines the lookup and the store is not consulted at all. Partial reuse is precisely what puts the match off a boundary, so it trades part of one block of device reuse for every stored block of the remaining prefix. Measured on MiniMax-M3, leaving it enabled declined 97.2% of lookups and left actual prompt cache read at 35% against a 96% ceiling; forcing it off raised that to 94% and roughly doubled throughput. #### How it keys pages @@ -242,7 +242,7 @@ The value for one key is the concatenation of that layer group's regions for one #### Transfer behavior * **Loads are synchronous**, performed in `start_load_kv` before the forward pass. A failed load raises: the runtime has already counted those tokens as computed, so a partial load is a wrong answer rather than a slow one. -* **Saves are asynchronous**, handed to a background thread behind a CUDA event recorded on the forward stream. The pages are only complete once the pass that wrote them retires, and blocking the executor loop on an RDMA write is the cost the store exists to avoid. The leader reports such requests as saving asynchronously, so their pages stay pinned until `get_finished` confirms the writes landed. A dropped save is logged rather than raised -- it only costs a future cache miss. +* **Saves are asynchronous**, handed to a background thread behind a CUDA event recorded on the forward stream. The pages are only complete once the pass that wrote them retires, and blocking the executor loop on an RDMA write is the cost the store exists to avoid. The leader reports such requests as saving asynchronously, so their pages stay pinned until `get_finished` confirms the writes landed. A dropped save is logged rather than raised, since it only costs a future cache miss. * Pages the store already holds are skipped, so several ranks or instances converging on the same prefix write it once. #### Unsupported configurations diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm index 6283ae6c0828..e3acd8bda3be 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm @@ -101,7 +101,7 @@ elif [ -d "${trtllm_repo}" ]; then if [ "${build_wheel}" = "true" ]; then echo "Building TensorRT-LLM wheel on one node..." - build_command="python3 ./scripts/build_wheel.py --use_ccache" + build_command="python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --use_ccache --clean" if [ -n "${cuda_architectures:-}" ]; then build_command="${build_command} --cuda_architectures \"${cuda_architectures}\"" fi @@ -141,14 +141,14 @@ fi # The Mooncake store bindings, when a worker config asks for the connector. # Images built from this repo bake them in (docker/common/install_mooncake.sh), # so the install below is only a fallback for images that predate it. Either -# way it is per job: --container-name gives each node a container that lives -# for the job, so anything installed here survives to the worker sruns but not -# into the next job. +# way it is per job, since --container-name gives each node a container that +# lives for the job: anything installed here survives to the worker sruns but +# not into the next job. mooncake_enabled=false if grep -qs "mooncake-store" "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"; then mooncake_enabled=true - # Both halves of the wheel are load-bearing and fail at different times: - # the connector needs mooncake.store in every context rank, and + # Both halves of the wheel are checked because they fail at different + # times: the connector needs mooncake.store in every context rank, and # 'trtllm-serve mooncake_master' needs the binary on PATH. if srun --container-name=${container_name} \ --container-mounts=${container_mount} --no-container-mount-home \ @@ -198,9 +198,9 @@ replace_placeholder "${client_cmds_base_file}" "${all_nodes_str}" "${client_cmds # read that address to lend the pool their memory. Nothing here starts, waits # for or configures any of it. # -# The one value a config written before submission cannot know is where this -# job's log directory is, and the master's address is published into it, so a -# __LOG_DIR__ placeholder in either worker config is filled in here. +# The one value a config written before submission cannot know is this job's log +# directory, which the master's address is published into, so a __LOG_DIR__ +# placeholder in either worker config is filled in here. if [ "${mooncake_enabled}" = "true" ]; then sed -i "s|__LOG_DIR__|${full_logdir}|g" \ "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml" @@ -229,8 +229,8 @@ done echo "Server is ready!" # A connector that failed to open its store handle does not stop the worker from -# serving, it just silently never hits, so surface the startup lines here rather -# than leaving them to be discovered after the benchmark. The registration line +# serving, it just never hits, so surface the startup lines here rather than +# leaving them to be discovered after the benchmark. The registration line # carries the bytes/page figure the pool sizing depends on. if [ "${mooncake_enabled}" = "true" ]; then echo "Mooncake store startup lines from the context workers:" @@ -252,16 +252,15 @@ while read -r cmd <&3; do fi done 3< "${client_cmds_file}" -# Collect the store's traffic into one file. The per-event lines live at DEBUG in -# the worker logs (module _torch), so this is only populated when a config asks -# for that verbosity; the counts are what distinguish "the store ran" from "the -# store ran and did something". +# Collect the store's traffic into one file. The per-event lines live at DEBUG +# in the worker logs (module _torch), so this is only populated when a config +# asks for that verbosity. The counts are what show the store did work rather +# than merely started. if [ "${mooncake_enabled}" = "true" ]; then mooncake_summary="${full_logdir}/9_mooncake_summary.log" { - # The address file is retracted when the master stops, so a stale - # address is never dialed; the context server's log still says which - # master the run used. + # The address file is retracted when the master stops, so fall back to + # the context server's log, which still names the master the run used. echo "master: $(tr -d '[:space:]' < "${full_logdir}/master.addr" 2>/dev/null \ || grep -hoE "master at [0-9.]+:[0-9]+" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null \ | head -n 1 | awk '{print $3}' || echo unknown)" @@ -276,12 +275,11 @@ if [ "${mooncake_enabled}" = "true" ]; then echo "${pattern}: ${count}" done echo - # Where the blocks physically went. The master names the segment for - # every allocation, and a segment is one client process's donated - # memory, so grouping by segment host answers the question the donors - # exist for: how much of the pool's contents lives on a decode node - # rather than on the prefill node that computed it. Without a donor this - # section shows a single host, which is the prefill node. + # Where the blocks physically went. The master names a segment for + # every allocation and a segment is one client process's donated + # memory, so grouping by segment host shows how much of the pool lives + # on a decode node rather than on the prefill node that computed it. + # Without a donor this section shows only the prefill node. echo "== block placement by segment host ==" echo "(lending hosts: $(grep -hoE "GiB of [0-9.]+ is now part of the pool" \ "${full_logdir}"/3_output_GEN_*.log 2>/dev/null \ @@ -318,7 +316,7 @@ if [ "${mooncake_enabled}" = "true" ]; then grep -h "mooncake-store:" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null \ | head -n 30 || echo "(none)" echo - # The master's own log, glog rather than anything TensorRT-LLM writes. + # The master's own glog output, not anything TensorRT-LLM writes. echo "== master ==" tail -n 50 "${full_logdir}/mooncake_master.log" 2>/dev/null || echo "(no master log)" } > "${mooncake_summary}" 2>&1 diff --git a/examples/disaggregated/slurm/benchmark/start_worker.sh b/examples/disaggregated/slurm/benchmark/start_worker.sh index 34aebdef7894..cf93aa095678 100644 --- a/examples/disaggregated/slurm/benchmark/start_worker.sh +++ b/examples/disaggregated/slurm/benchmark/start_worker.sh @@ -55,29 +55,29 @@ fi echo "config_file: ${config_file}" # The mooncake-store pool is described in the worker config and provisioned by -# trtllm-serve during bringup. Anchoring its run directory here is what puts the -# master's log, the client config it renders and the address it publishes in the +# trtllm-serve during bringup. Anchoring its run directory here keeps the +# master's log, the rendered client config and the published address in the # job's log directory rather than in a temporary directory that shutdown -# removes -- and it is how the ranks the launcher started, which never inherited -# the leader's environment, find that client config. An inherited +# removes, and it is how the ranks srun started, which never inherited the +# leader's environment, find that client config. An inherited # MOONCAKE_CONFIG_PATH still wins, so an externally managed pool stays reachable. export TRTLLM_MOONCAKE_RUN_DIR="${log_dir}" # The generation servers wait for a master the context server starts. Both are # launched together and the master comes up before its model loads, but the wait # spans container start on another node, so it is given far more than the 60s -# default: too short fails the job, too long costs nothing when the master is -# there. +# default. Too short fails the job; too long costs nothing when the master is +# already there. export TRTLLM_MOONCAKE_MASTER_TIMEOUT="${TRTLLM_MOONCAKE_MASTER_TIMEOUT:-900}" # MiniMax-M3's MSA sparse attention JIT-compiles its FMHA kernels on first use, -# from inside the attention forward pass: one TP rank runs ninja while the others -# block on a file lock, so an uncached variant stalls the whole executor loop for -# ~8s (and ~70s when an iteration needs several). The cache defaults to -# ~/.cache, which is thrown away here because the container is started with -# --no-container-mount-home, making every job pay the compiles again during -# serving. Anchor it next to this script instead: that path is on the mounted -# filesystem and identical across jobs, so only the first run compiles. +# from inside the attention forward pass. One TP rank runs ninja while the +# others block on a file lock, so an uncached variant stalls the whole executor +# loop for ~8s, or ~70s when an iteration needs several. The cache defaults to +# ~/.cache, which is thrown away because the container is started with +# --no-container-mount-home, so every job would pay the compiles again during +# serving. Anchoring it next to this script puts it on the mounted filesystem +# at a path identical across jobs, so only the first run compiles. if [ -z "${MINFER_FMHA_CACHE_DIR:-}" ]; then export MINFER_FMHA_CACHE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/.cache/minfer/fmha_sm100" mkdir -p "${MINFER_FMHA_CACHE_DIR}" @@ -85,10 +85,9 @@ if [ -z "${MINFER_FMHA_CACHE_DIR:-}" ]; then fi # Per-transfer KV timings (size, queue/transfer latency, throughput) as CSV next -# to the worker logs. This is what tells apart "prefill is slow" from "the -# prefill->decode handoff is slow", which the aggregate benchmark numbers -# cannot. Same rationale as above for defaulting the path here: an explicit -# setting wins, and it can be turned off with KV_TRANSFER_PERF_LOG=false. +# to the worker logs. These separate slow prefill from a slow prefill-to-decode +# handoff, which the aggregate benchmark numbers cannot. An explicit setting +# wins, and KV_TRANSFER_PERF_LOG=false turns it off. if [ "${KV_TRANSFER_PERF_LOG:-true}" = "true" ] \ && [ -z "${TLLM_KV_TRANSFER_PERF_LOG_FILE:-}" ]; then export TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1 diff --git a/examples/disaggregated/slurm/benchmark/watch_job.sh b/examples/disaggregated/slurm/benchmark/watch_job.sh index 613f3714e68d..4b27a2322f14 100644 --- a/examples/disaggregated/slurm/benchmark/watch_job.sh +++ b/examples/disaggregated/slurm/benchmark/watch_job.sh @@ -26,8 +26,8 @@ count_matches() { } # How much of the pool's contents lives on each node. A segment is one client -# process's donated memory, so a single host here means the pool is prefill-only -# and the memory donors are either absent or not being allocated into. +# process's donated memory, so a single host here means the pool is +# prefill-only: the donors are absent or not being allocated into. placement() { grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" \ "${log_dir}/2_mooncake_master.log" 2>/dev/null \ @@ -62,8 +62,8 @@ while [ ${SECONDS} -lt ${deadline} ]; do [ "${matched:-0}" -gt 0 ] && announce "first_match" "first store hit; matched lines=${matched}" [ "${loaded:-0}" -gt 0 ] && announce "first_load" "first store load; loaded lines=${loaded}" - # The point of the donors: report as soon as a second host appears, since - # that is the first moment prefill-written KV is provably on decode DRAM. + # A second host is the first moment prefill-written KV is provably on + # decode DRAM, which is what the donors exist for. hosts=$(grep -o "segment=[0-9.]*:" "${log_dir}/2_mooncake_master.log" 2>/dev/null \ | sort -u | wc -l || true) [ "${hosts:-0}" -ge 2 ] && announce "multi_host" \ @@ -86,8 +86,8 @@ while [ ${SECONDS} -lt ${deadline} ]; do exit 0 fi - # A dead batch script leaves the tree untouched; report it rather than - # polling a corpse until the deadline. + # A dead batch script leaves the tree untouched, so report it rather than + # polling until the deadline. if grep -qs "Job completed successfully" "${log_dir}"/slurm-*.out; then echo "EVENT: batch script reported completion" exit 0 diff --git a/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml index 85e8530de354..350c77b65941 100644 --- a/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml +++ b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml @@ -1,7 +1,7 @@ # Extra LLM API options for trtllm-serve with the Mooncake store KV connector. # -# Offloads KV pages to a Mooncake distributed store -- a shared CPU memory pool -# addressed by content -- so a prefix computed by one engine can be replayed by +# Offloads KV pages to a Mooncake distributed store, a shared CPU memory pool +# addressed by content, so a prefix computed by one engine can be replayed by # another. This is a different component from the Mooncake transfer engine used # by the C++ cache transceiver for disaggregated prefill/decode handoff; the two # compose rather than conflict. diff --git a/mooncake_disagg/README.md b/mooncake_disagg/README.md index e6cba4cc7543..e160704212ae 100644 --- a/mooncake_disagg/README.md +++ b/mooncake_disagg/README.md @@ -6,7 +6,7 @@ disaggregated benchmark harness in `examples/disaggregated/slurm/benchmark/`. The unit tests in `tests/unittest/_torch/executor/test_mooncake_store_connector.py` -cover the pieces that decide whether a cache hit is *correct* -- key namespacing, +cover the pieces that decide whether a cache hit is *correct*: key namespacing, hash chaining, page addressing, the startup gates. They deliberately do not run a store, a model, or two engines. What is untested is everything that decides whether the feature is *worth having*: whether a real prefix survives the round @@ -32,8 +32,8 @@ The residency measurements in `../m3-kv-residency-measurement-README.md` (taken on this same model and workload) found that M3 production traffic already serves **97.0% of prompt tokens from local cache**, with eviction responsible for under 0.7% of misses. On a *single* instance there is almost no headroom for the store -to recover -- over 99% of misses are prefixes never cached anywhere, which no -store can serve either. +to recover, since over 99% of misses are prefixes never cached anywhere, which +no store can serve either. That is not an argument against the feature; it is an argument about where to look. Set expectations accordingly: @@ -63,7 +63,7 @@ differently: | Component | What uses it | How it gets installed | Since | |---|---|---|---| | C++ transfer engine (`/usr/local/Mooncake`) | the C++ cache transceiver's Mooncake backend | CMake source build in `docker/common/install_mooncake.sh` | PR #8447, Nov 2025 | -| Python bindings (`mooncake.store.MooncakeDistributedStore`) | **this connector** | pip wheel, added to the same script | commit `d36dae435e`, **this branch** | +| Python bindings (`mooncake.store.MooncakeDistributedStore`) | **this connector** | pip wheel, added to the same script | commit `d36dae435e` | So Mooncake has been in the container images for months, but only usefully as the C++ library. Three consequences: @@ -119,7 +119,7 @@ chooses where to put it with: COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print([s for s in sys.path if 'packages' in s][0])" ``` --- the *first* `sys.path` entry whose name merely contains `"packages"`. That +That is the *first* `sys.path` entry whose name merely contains `"packages"`. It gives two different failures depending on what else is installed, and both produce the same confusing symptom: an `ImportError` **after a `pip install` that reported success**. @@ -170,7 +170,7 @@ MOONCAKE_WHEEL="mooncake-transfer-engine==0.3.7.post2" \ ``` Both wheel choices were validated against the full set of store methods the -connector calls -- see "Validating the install" below. +connector calls; see "Validating the install" below. ### How often does the script need to run? @@ -180,10 +180,10 @@ writes into `dist-packages` inside the container, not into your checkout. | Situation | How often | |---|---| | Long-lived container you `docker exec` into | **Once.** It survives until the container is deleted; `docker restart` keeps it. | -| SLURM via `disaggr_torch.slurm` | **Once per job, per node** -- and the harness now does it for you, see below. | -| Image built from this branch | **Never.** `install_mooncake.sh` now does it at build time and fails the build if the import does not work. | +| SLURM via `disaggr_torch.slurm` | **Once per job, per node**, which the harness does for you, see below. | +| Image built from this repo | **Never.** `install_mooncake.sh` does it at build time and fails the build if the import does not work. | -For the SLURM case this is already wired up: `disaggr_torch.slurm` now runs the +For the SLURM case this is already wired up: `disaggr_torch.slurm` runs the script on every node, right after its `pip install -e .[devel]` step, and gates it on whether a worker config actually asks for the connector: @@ -195,7 +195,7 @@ So arms A and B of the run matrix pay nothing, arm C installs automatically, and there is no new config key to remember. It resolves the script as `${trtllm_repo}/mooncake_disagg/install_mooncake_runtime.sh` and fails the job with an explicit message if `environment.trtllm_repo` is unset or does not -contain it -- which is the case if you benchmark from +contain it, which is the case if you benchmark from `environment.trtllm_wheel_path` instead, so use an image with the bindings baked in for that path. Output lands in `/2_install_mooncake.log`. Set `MOONCAKE_WHEEL` in the submitting environment to override the wheel; it is @@ -217,16 +217,16 @@ than a necessity. The root cause is upstream in Mooncake's `CMakeLists.txt` and is **not** fixed; both scripts clean up after it. Practically: -- **Images built from this branch:** no. The cleanup runs in the same script, +- **Images built from this repo:** no. The cleanup runs in the same script, immediately after `make install` and before the wheel install, and the build - now fails if `import mooncake.store` does not work. + fails if `import mooncake.store` does not work. - **Any pre-existing image**, including the one pinned in `jenkins/current_image_tags.properties` (`202607211045`): the broken package is baked in, so the runtime script is required. - **Inside a running container:** only if something re-runs Mooncake's CMake - install. Reinstalling `nvidia-cutlass-dsl` does not recreate it -- that package - has never shipped a `mooncake` directory; it only supplies the `.pth` that made - CMake choose the wrong destination. + install. Reinstalling `nvidia-cutlass-dsl` does not recreate it: that package + has never shipped a `mooncake` directory, and only supplies the `.pth` that + made CMake choose the wrong destination. - **If Mooncake is ever upgraded** to a version that fixes its install path, or the `.pth` ordering changes, the cleanup becomes a no-op rather than a hazard. @@ -248,9 +248,9 @@ with the same argument list `worker.py` passes. `mooncake_api_surface_test.py` i the one that matters when changing wheel versions: the connector's hot path never uses `put`/`get`, it uses `register_buffer` plus the zero-copy `batch_put_from_multi_buffers` / `batch_get_into_multi_buffers` / `batch_is_exist` -calls against registered GPU pages. Those take `list[list[int]]` -- one buffer +calls against registered GPU pages. Those take `list[list[int]]`, one buffer list per key, because `PageAddressing.page_buffers` returns one address per -layer-group region -- and that is the signature most likely to drift. +layer-group region, and that is the signature most likely to drift. Then the unit tests, which need no store and no GPU: @@ -291,9 +291,9 @@ scheduler, both of which the connector would forbid. Pool capacity comes only from processes that open a store handle: `setup` registers `global_segment_size` bytes of the calling process's host memory, and the master then places blocks in it. Since only the context workers configure -the connector, only they contribute memory — so by default every byte of the -pool is prefill-node DRAM, and the store is a prefill-DRAM-caches-prefill-GPU -tier that largely duplicates TensorRT-LLM's native host offload. Confirm this on +the connector, only they contribute memory, so by default every byte of the pool +is prefill-node DRAM and the store is a prefill-DRAM-caches-prefill-GPU tier +that largely duplicates TensorRT-LLM's native host offload. Confirm this on any run by grouping the master's `allocation_succeeded ... segment=:` lines by host: a single host means a prefill-only pool. @@ -303,7 +303,7 @@ as long as it runs, so the pool spans both sides while its engine stays connector-free and keeps its cache transceiver for the KV handoff: ```yaml -# gen worker config -- no kv_connector_config anywhere near it +# gen worker config: no kv_connector_config anywhere near it mooncake_donation: master_server_address: file:///$WORK_DIR/master.addr segment_size: 640GiB @@ -314,7 +314,7 @@ Donation is deliberately outside `kv_connector_config`, and not a `StoreRole` either: the roles describe traffic (`producer` writes, `consumer` reads, `both`), none of them means "contribute memory only", and configuring capacity there would start this server using the store. The size is charged **per server -process, not per rank** — unlike `global_segment_size` — so two servers on one +process, not per rank**, unlike `global_segment_size`, so two servers on one node lend twice this. The server is ready only once its segment is mounted, which makes readiness the @@ -336,17 +336,17 @@ the two together. The worker logs its own share as `KV cache manager v2 host cache quota set to N GiB`, **per rank**, against the `available host memory` it reports on the same line. -## 4. Step 1 -- run the Mooncake master +## 4. Step 1: run the Mooncake master `master_server_address` is mandatory, so a master must exist and be reachable from every worker. **Outside SLURM you can skip this section too.** A worker config carrying `kv_connector_config.mooncake_store` makes `trtllm-serve` provision the pool -during its own bringup — `launch_master: true` starts a master for that server -alone, `master_server_address` joins one that already exists — and write the -client config itself. That covers aggregated and single-instance runs, which is -what `m3_agg_mooncake.yaml` now does; `mooncake_usage.md` §2 has the table. The +during its own bringup, with `launch_master: true` starting a master for that +server alone and `master_server_address` joining one that already exists, and +write the client config itself. That covers aggregated and single-instance runs, +as `m3_agg_mooncake.yaml` does; `mooncake_usage.md` §2 has the table. The rest of this section is about the master the experiments below need, which outlives any one server and therefore cannot be owned by one. @@ -354,11 +354,11 @@ outlives any one server and therefore cannot be owned by one. no master and writes no client config: the context worker's `launch_master: true` does both, on the context node, and publishes the address the generation workers' `mooncake_donation` reads. `disaggr_torch.slurm` contributes exactly -two things — it installs the bindings on every node, and it substitutes +two things: it installs the bindings on every node, and it substitutes `__LOG_DIR__` in the worker configs, since the run directory is the one value a config written before submission cannot know. Everything else, the pool sizes -and the HCA included, is in the config; no `MOONCAKE_*` variable is read from -the submitting environment any more. +and the HCA included, is in the config, and no `MOONCAKE_*` variable is read +from the submitting environment. The master's log lands in `/mooncake_master.log` and its address in `/master.addr` while it runs, because `start_worker.sh` sets @@ -367,7 +367,7 @@ context server's other ranks read the rendered `mooncake.json`: they are separate srun tasks that never inherited the leader's environment. That master dies with the job, so read on if you need a pool that outlives one -allocation -- which experiment 3 does, by construction. Run it as its own +allocation, which experiment 3 does by construction. Run it as its own long-lived job and export `MOONCAKE_MASTER_ADDRESS=:50051` before `submit.py`; the harness then skips launching one and only writes the client config pointing at yours. @@ -390,9 +390,9 @@ srun --container-image=$CONTAINER_IMAGE \ ``` The master runs for as long as the command does, and `--address_file` receives -`host:port` **once it accepts connections** — so waiting for that file is -waiting for readiness, and its absence after the job starts is a failure rather -than a slow start. It is removed on exit, so a stale address is never dialed. +`host:port` **once it accepts connections**, so waiting for that file is waiting +for readiness, and its absence after the job starts is a failure rather than a +slow start. It is removed on exit, so a stale address is never dialed. `--run_dir` keeps the master's log at `$WORK_DIR/mooncake_master.log`, which is where pool occupancy and eviction are read from. @@ -406,14 +406,14 @@ surface (`--rpc_address`, `--rpc_thread_num`, `--default_kv_lease_ttl`, Keeping the master in a separate job is what makes experiment 3 (§7) possible: the pool outlives the engines, so a second benchmark job finds a warm store. -Workers can now be pointed at it without anyone writing the address down: +Workers are pointed at it without anyone writing the address down: `master_server_address: file://$WORK_DIR/master.addr` in a config's `mooncake_store` block makes each server read it during bringup and wait if the master job has not started yet. That is what makes a master whose host the scheduler chose usable from a config settled beforehand. Failing that, write the client config, substituting the address the master job -just recorded -- or let `disaggr_torch.slurm` generate it, as above. The schema +just recorded, or let `disaggr_torch.slurm` generate it, as above. The schema is vLLM's, so one pool can serve both engines: ```bash @@ -441,8 +441,8 @@ EOF behaviour; confirm the port from `--help` rather than assuming. - `device_name`: pick from `ibv_devinfo` on a compute node. For first bring-up only, `"protocol": "tcp"` with `"device_name": ""` removes RDMA from the - variable list -- that is what `mooncake.json` in this directory currently - does. Do not draw performance conclusions from a TCP run. + variable list, which is what `mooncake.json` in this directory does. Do not + draw performance conclusions from a TCP run. - `global_segment_size` is contributed **per worker process**, so the pool is `global_segment_size x (ctx instances x TP)` = 8 segments here. - Sizing: after startup, each worker logs its page geometry (§8). Pool bytes for @@ -457,13 +457,13 @@ EOF prefix whenever you change anything that should not be shared with an earlier run's pages. -## 5. Step 2 -- the harness config +## 5. Step 2: the harness config Copy `examples/disaggregated/slurm/benchmark/config.yaml` and replace the `worker_config` section with M3's. `submit.py` serializes `worker_config.ctx` and `worker_config.gen` straight to `ctx_config.yaml`/`gen_config.yaml` with -`yaml.dump`, so any LLM-API key passes through untouched -- including -`kv_connector_config`. +`yaml.dump`, so any LLM-API key passes through untouched, `kv_connector_config` +included. The context worker below is `m3_ctx_mooncake.yaml` from this directory; the generation worker is `m3_gen_mooncake.yaml`. Every deviation from the production @@ -499,7 +499,7 @@ hardware: environment: container_mount: "" - container_image: "" + container_image: "" model_path: "" trtllm_repo: "" build_wheel: false @@ -559,7 +559,7 @@ worker_config: stream_interval: 20 print_iter_log: true num_postprocess_workers: 8 - # Required to see any reuse number at all -- see section 8. All three + # Required to see any reuse number at all; see section 8. All three # default to false, and without them /metrics returns an empty list. enable_iter_perf_stats: true enable_iter_req_stats: true @@ -609,8 +609,8 @@ worker_config: Eagle3 is left off. It is not gated, but `MiniMaxM3KVCacheManagerV2` sets `supports_shared_draft_layers`, so draft layers join the unified V2 cache and -therefore the registered layout -- extra page geometry the store must key -correctly, on a path with no coverage. Turn it on only after a clean run +therefore the registered layout. That is extra page geometry the store must key +correctly, on a path with no coverage, so turn it on only after a clean run without it. Submit with: @@ -621,7 +621,7 @@ python3 submit.py -c /m3_store_2ctx.yaml --dry-run # inspect first python3 submit.py -c /m3_store_2ctx.yaml ``` -## 6. Step 3 -- the workload +## 6. Step 3: the workload `run_benchmark.sh` invokes `benchmark_serving` with `--dataset-name trtllm_custom --dataset-path `, so you supply a @@ -678,18 +678,18 @@ Size the file against what the client will actually request. `run_benchmark.sh` computes `num_prompts = (concurrency x num_gen_servers) x multi_round`, so the §5 config (`concurrency_list: "8"`, `multi_round: 8`, one generation server) asks for 64 -prompts -- which is why `P x R` above is 64. Ask for more than the file holds +prompts, which is why `P x R` above is 64. Ask for more than the file holds and the extra is not sampled; write more than you ask for and the tail of your repeat structure never runs. Random token text is deliberate: it defeats any accidental prefix sharing between "distinct" prefixes, so the hit rate you measure is the one you designed. If you would rather test genuine production traffic, substitute a -real multi-turn trace -- but keep an eye on whether it actually contains -repeated prefixes, since without them the store has nothing to do and a flat -result means nothing. +real multi-turn trace, but keep an eye on whether it actually contains repeated +prefixes, since without them the store has nothing to do and a flat result means +nothing. -## 7. Step 4 -- the run matrix +## 7. Step 4: the run matrix Three arms, and the middle one is the one people skip: @@ -707,14 +707,14 @@ worth knowing and they answer different questions. Then, within that: -**Experiment 1 -- does it work at all (`num_ctx_servers: 1`).** +**Experiment 1: does it work at all (`num_ctx_servers: 1`).** Arm C, one context instance, small `PREFIX_TOKENS` (say 4096) and a short run. You are looking for a clean startup, the registration log line, non-zero store traffic, no load failures, and coherent output text. Do this before spending an allocation on anything larger. Expect no throughput change; local reuse already serves this case. -**Experiment 2 -- cross-instance reuse (`num_ctx_servers: 2`).** +**Experiment 2: cross-instance reuse (`num_ctx_servers: 2`).** The router defaults to round-robin, so consecutive requests alternate between context instances and roughly half of each prefix's repeats land on the instance that did not compute it. Those are the requests local reuse must recompute from @@ -725,21 +725,21 @@ scratch and the store can serve. Compare arm C against arm B on: This is the primary result. A store that does not win here does not work. -**Experiment 3 -- survival across restarts.** +**Experiment 3: survival across restarts.** Run experiment 2's arm C twice, same `TRTLLM_MOONCAKE_STORE_PREFIX`, with the master job left running between them. The second job starts with an empty local cache but a warm pool. First-round TTFT should fall toward the warm steady-state value. Local reuse scores zero here by construction, so any improvement is -attributable to the store alone -- which makes this the cleanest signal in the +attributable to the store alone, which makes this the cleanest signal in the whole matrix, and the cheapest to run. -**Experiment 4 -- the cost when there is nothing to gain.** +**Experiment 4: the cost when there is nothing to gain.** Arm C against arm B on a workload with *no* repeated prefixes (unique prompts). This measures pure overhead: lookups, key hashing on the leader, background saves competing for host bandwidth. Ideally indistinguishable from arm B. This is the arm that catches a feature that helps its benchmark and hurts the fleet. -## 8. Step 5 -- reading the results +## 8. Step 5: reading the results ### Did the connector even load? @@ -768,9 +768,9 @@ so: TLLM_LOG_LEVEL_BY_MODULE="debug:_torch" ``` -added to `environment.ctx_worker_env_var`. This is verbose -- it enables DEBUG -for all of `_torch` -- so use it for experiment 1 and for diagnosis, not for the -runs you intend to quote numbers from. The lines worth counting: +added to `environment.ctx_worker_env_var`. This is verbose, since it enables +DEBUG for all of `_torch`, so use it for experiment 1 and for diagnosis rather +than for the runs you intend to quote numbers from. The lines worth counting: ``` mooncake-store matched N blocks (M tokens) for request R # leader, a hit @@ -784,9 +784,9 @@ Scrape it before and after a run and diff. ### Where did the pages land? Page counts alone do not say whether the pool is doing anything the native host -offload tier could not. For that, group the master's allocations by segment host -— a segment is one client process's donated memory, so the host tells you which -node the block physically lives on: +offload tier could not. For that, group the master's allocations by segment +host. A segment is one client process's donated memory, so the host tells you +which node the block physically lives on: ```bash grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" mooncake_master.log \ @@ -802,8 +802,8 @@ and being read back from there. `disaggr_torch.slurm` writes this breakdown into it needs running by hand only when diagnosing a partial run. Requires `GLOG_v=1` on the master, which `trtllm-serve mooncake_master` sets -unless `GLOG_v` is already in its environment -- so raise it by exporting -`GLOG_v` to that command. +unless `GLOG_v` is already in its environment, so raise it by exporting `GLOG_v` +to that command. ### Which reuse number means what @@ -814,24 +814,24 @@ This distinction matters and is easy to get backwards: | `reused_blocks_per_request`, `kv_cache_hit_rate_per_request` | per-request iteration stats | **Yes.** `_reserve_connector_prefix` calls `set_prepopulated_prompt_len` with the connector-served position, and these derive from `mPrepopulatedPromptLen`. | | `kv_cache_iter_reused_blocks`, `kv_cache_iter_reuse_rate` | `GET /prometheus/metrics` | **No.** These come from the local V2 reuse tree's committed stats. | -So **store hits ≈ per-request reuse − local-tree reuse**. Confirm that +So **store hits are roughly per-request reuse minus local-tree reuse**. Confirm that relationship on experiment 3, where the local tree starts empty and the difference is unambiguous, before relying on it elsewhere. Getting at either one requires the three flags added to the worker configs in §5, all of which default to false: -- `enable_iter_perf_stats: true` -- without it `get_latest_iteration_stats` +- `enable_iter_perf_stats: true`, without which `get_latest_iteration_stats` short-circuits and `GET /metrics` returns `[]`. -- `enable_iter_req_stats: true` -- needed for the *per-request* half of the +- `enable_iter_req_stats: true`, needed for the *per-request* half of the table above. -- `return_perf_metrics: true` -- mounts `/prometheus/metrics`. `GET /metrics` +- `return_perf_metrics: true`, which mounts `/prometheus/metrics`. `GET /metrics` (plain JSON iteration stats) is routed unconditionally but still needs `enable_iter_perf_stats`. `print_iter_log: true` is worth keeping on, but note it prints iteration timing -and KV *utilization* only -- no reuse counters. Do not go looking for hit rates -there. +and KV *utilization* only, with no reuse counters. Do not go looking for hit +rates there. Per-request client-side results land in `/concurrency_/result.json` with TTFT/TPOT/ITL/E2EL percentiles, which is where the headline numbers for @@ -845,7 +845,7 @@ with TTFT/TPOT/ITL/E2EL percentiles, which is where the headline numbers for | `mooncake-store background save failed` | A save thread exception, re-raised on the executor thread. | | `mooncake-store rank K failed to save N of M pages` (warning) | Dropped write. Costs a future miss, not correctness. A trickle is tolerable; a flood means the pool is full or the master is overloaded. | | `mooncake-store lookup failed; treating as a miss` (warning) | Probe failed. Degrades to no-store behavior. | -| `could not reserve connector prefix up to N, falling back to the local match` (debug) | Out of GPU pages. The store offered more than the engine could hold -- expected under pressure, but frequent occurrences mean the offer is outrunning capacity. | +| `could not reserve connector prefix up to N, falling back to the local match` (debug) | Out of GPU pages. The store offered more than the engine could hold, which is expected under pressure, but frequent occurrences mean the offer is outrunning capacity. | ### Sanity check that is not a performance number @@ -866,7 +866,7 @@ gives a coarser version of the same check. Hold it fixed across every arm, generation workers included, or you are comparing two different models. - **Key namespace pins world size and rank.** Change TP and every stored page - becomes unreachable -- a miss, not an error. Same for `tokens_per_block`, the + becomes unreachable, as a miss rather than an error. Same for `tokens_per_block`, the layer group set, and `bytes_per_page`. - **`model_key` defaults to the checkpoint directory's basename.** Two hosts mounting the same checkpoint at different paths still share cache, which is @@ -879,15 +879,15 @@ gives a coarser version of the same check. - **UCX warmup requests hit the store too.** `run_benchmark.sh` sends `2 x ctx_instances x gen_instances` 100-token requests before the real run. Harmless, but they are in the counters. -- **A partial local match disables the store for that request entirely,** and - this is the dominant failure mode rather than the rare one it was predicted to - be here. The connector offers only whole blocks and only when the local match - is block-aligned, and `enable_partial_reuse` (default `true`) is exactly what - puts the match off a boundary: measured on M3, it declined **97.2% of - lookups**, so a 1.6 TB pool measured as if it were absent. Turning it off took - actual prompt cache read from 35% to 94%. `py_executor_creator` now forces it - off for this connector, so the hazard is gone, but the arithmetic is worth - knowing before changing `tokens_per_block`. See `../mooncake_usage.md` §3. +- **A partial local match disables the store for that request entirely,** and it + is the dominant failure mode rather than a rare one. The connector offers only + whole blocks and only when the local match is block-aligned, and + `enable_partial_reuse` (default `true`) is exactly what puts the match off a + boundary: measured on M3, it declined **97.2% of lookups**, so a 1.6 TB pool + measured as if it were absent. Turning it off took actual prompt cache read + from 35% to 94%. `py_executor_creator` forces it off for this connector, so + the hazard cannot be hit, but the arithmetic is worth knowing before changing + `tokens_per_block`. See `../mooncake_usage.md` §3. - **`block_reuse_policy: per_conversation` is off on the context worker** in these configs. It is not gated, but the connector derives its own `cache_salt`-seeded hash chain and the interaction is untested. Restore it @@ -901,5 +901,5 @@ pipeline or context parallelism (both refused); no VSWA or sliding-window model (refused); no Eagle3; no shared pool between TensorRT-LLM and vLLM, though the config schema is deliberately compatible with it. Load bandwidth under contention from many simultaneous large prefixes is exercised only incidentally -by concurrency, not measured directly -- if experiment 2 shows a TTFT -regression at high concurrency despite hits, that is the first thing to profile. +by concurrency, not measured directly. If experiment 2 shows a TTFT regression +at high concurrency despite hits, that is the first thing to profile. diff --git a/mooncake_disagg/gen_config.yaml b/mooncake_disagg/gen_config.yaml index 4b8ed0b35262..6851670e3a78 100644 --- a/mooncake_disagg/gen_config.yaml +++ b/mooncake_disagg/gen_config.yaml @@ -1,8 +1,8 @@ # Generation (decode) worker: does not touch the Mooncake store at all. # # There is deliberately no kv_connector_config here. "Decode none" is the -# absence of a connector, not a StoreRole -- StoreRole only has -# producer / consumer / both. Generated tokens are rarely a reused prefix, +# absence of a connector, not a StoreRole: StoreRole only has producer, +# consumer and both. Generated tokens are rarely a reused prefix, # so writing them would cost bandwidth for no hit rate. # # CUDA_VISIBLE_DEVICES=1 trtllm-serve \ diff --git a/mooncake_disagg/install_mooncake_runtime.sh b/mooncake_disagg/install_mooncake_runtime.sh index 27c2a658885c..8c669ca00f0d 100755 --- a/mooncake_disagg/install_mooncake_runtime.sh +++ b/mooncake_disagg/install_mooncake_runtime.sh @@ -1,46 +1,30 @@ #!/bin/bash # Make the Mooncake Python store bindings importable inside a TensorRT-LLM -# container, so the mooncake-store KV connector can start. +# container, so the mooncake-store KV connector can start. Images built by +# docker/common/install_mooncake.sh already have this; run it on images that +# predate that, or to change the wheel. # -# Two things break a plain `pip install mooncake-transfer-engine` in the -# containers built by docker/common/install_mooncake.sh: +# Two things break a plain `pip install mooncake-transfer-engine` in those +# containers, both explained in docker/common/install_mooncake.sh: the CMake +# source build leaves behind an unusable `mooncake` package that shadows or +# collides with the wheel, and the default wheel is linked against +# libcudart.so.12 while these images ship CUDA 13 only. Either way the symptom +# is `ImportError: libmooncake_store.so` after pip reports success. # -# 1. The CMake source build in that script emits its own, unusable `mooncake` -# Python package (it omits libmooncake_store.so). mooncake-integration's -# CMakeLists picks the install directory with -# python3 -c "import sys; print([s for s in sys.path if 'packages' in s][0])" -# -- the first sys.path entry whose name merely contains "packages". With -# nvidia-cutlass-dsl installed that is nvidia_cutlass_dsl/dsl_packages, -# which nvidia_cutlass_dsl_packages.pth sys.path.insert(0)s, so it shadows -# anything pip installs. Without it, the package lands in dist-packages and -# collides with the wheel: CMake writes store.cpython-312-.so, the -# wheel writes store.so, and importlib prefers the interpreter-tagged -# suffix, so the broken extension still wins. Either way the symptom is -# `ImportError: libmooncake_store.so` *after* pip reports success. -# Because pip overwrites __init__.py in the collision case, leftovers are -# not reliably identifiable after the fact -- so this script removes the -# package directory outright and reinstalls, rather than trying to tell -# good files from bad. +# `mooncake-transfer-engine-cuda13` needs no CUDA 12 shim, so it is the default +# here. Its releases start at 0.3.9 and so cannot match the pin in +# install_mooncake.sh, which is safe: that CMake-built C++ library backs the +# cache transceiver's Mooncake backend, a different feature, while the +# connector only ever talks to the wheel. The wheel also supplies the +# mooncake_master that lands on PATH, so client and master stay matched. +# Revisit only if cache_transceiver_config.backend is set to MOONCAKE. # -# 2. The `mooncake-transfer-engine` wheel is linked against libcudart.so.12, -# while containers from pytorch-26.05 on ship CUDA 13 only. -# `mooncake-transfer-engine-cuda13` is the same project built for CUDA 13 -# and needs no shim, so it is the default here. Its releases start at 0.3.9, -# so it cannot match the 0.3.7.post2 pin in install_mooncake.sh -- see the -# note below on why that is safe. -# -# Version drift against /usr/local/Mooncake: that CMake-built C++ library backs -# the *cache transceiver's* Mooncake backend, a different feature. The connector -# only ever talks to the wheel, and the wheel also supplies the mooncake_master -# that lands on PATH, so client and master stay matched. Revisit only if you set -# cache_transceiver_config.backend to MOONCAKE (these configs use NIXL). -# -# Set MOONCAKE_WHEEL to override, e.g. +# Set MOONCAKE_WHEEL to override, for example # MOONCAKE_WHEEL="mooncake-transfer-engine==0.3.7.post2" -# to match install_mooncake.sh exactly; the libcudart.so.12 shim is then applied -# automatically. +# to match install_mooncake.sh exactly; the libcudart.so.12 shim is then +# applied automatically. # -# Idempotent, and cheap on re-runs: if the install is already correct it exits +# Idempotent and cheap on re-runs: if the install is already correct it exits # without contacting the network, so it is safe in a SLURM prolog on every node. set -euo pipefail @@ -59,9 +43,8 @@ if pip3 show "${WHEEL_NAME}" >/dev/null 2>&1 && exit 0 fi -# Purge every `mooncake` package directory on the search path, whatever wrote it. -# Distinguishing CMake leftovers from wheel files is unreliable once pip has -# overwritten __init__.py, so remove and reinstall instead. +# Purge every `mooncake` package directory on the search path, whatever wrote +# it, since leftovers cannot be told apart from wheel files reliably. python3 - <<'PY' import os import shutil diff --git a/mooncake_disagg/m3_agg_mooncake.yaml b/mooncake_disagg/m3_agg_mooncake.yaml index 42f73c018d56..39610cc63899 100644 --- a/mooncake_disagg/m3_agg_mooncake.yaml +++ b/mooncake_disagg/m3_agg_mooncake.yaml @@ -46,7 +46,7 @@ kv_cache_config: # CONNECTOR: was host_cache_size: 388554555392. # _reject_non_gpu_cache_tiers rejects every tier below GPU, because a # registered region is only a valid device address while its page stays - # pinned to GPU -- eviction reassigns the slot. Both must be an explicit + # pinned to GPU, and eviction reassigns the slot. Both must be an explicit # 0: V2 provisions a host tier when the field is left at its default of # None, which is falsy but still yields a tier. host_cache_size: 0 diff --git a/mooncake_disagg/m3_ctx_mooncake.yaml b/mooncake_disagg/m3_ctx_mooncake.yaml index 0d25d67bb4ac..1a2c1b20ece5 100644 --- a/mooncake_disagg/m3_ctx_mooncake.yaml +++ b/mooncake_disagg/m3_ctx_mooncake.yaml @@ -43,7 +43,7 @@ kv_cache_config: # MANDATORY, and not only because of the connector. M3 sets # sparse_attention_config, so get_kv_cache_manager_cls routes to -# MiniMaxM3KVCacheManagerV2 unconditionally -- use_kv_cache_manager_v2 is not +# MiniMaxM3KVCacheManagerV2 unconditionally, and use_kv_cache_manager_v2 is not # even consulted on that branch. KVCacheManagerV2 cannot drive the C++ # transceiver, and M3 does not override get_preferred_transceiver_runtime, so # 'auto' would resolve to C++ and be rejected. @@ -90,7 +90,7 @@ kv_connector_config: global_segment_size: 160GiB # Registering the KV pools with the HCA needs nvidia_peermem; without # it registration fails on every range, and staging registers only - # host memory instead. 1GiB holds a full transfer_batch_size of - # pages -- less silently reduces the batch rather than failing. + # host memory instead. 1GiB holds a full transfer_batch_size of pages; + # less silently reduces the batch rather than failing. stage_through_host: true staging_buffer_bytes: 1GiB diff --git a/mooncake_disagg/m3_gen_mooncake.yaml b/mooncake_disagg/m3_gen_mooncake.yaml index bb0102d0f7d0..e666099674b6 100644 --- a/mooncake_disagg/m3_gen_mooncake.yaml +++ b/mooncake_disagg/m3_gen_mooncake.yaml @@ -1,7 +1,7 @@ # MiniMax-M3-NVFP4 GENERATION (decode) worker: does not touch the store. # # There is deliberately no kv_connector_config. That absence is the whole of -# "decode-none" -- StoreRole has only producer / consumer / both, so there is +# "decode-none": StoreRole has only producer, consumer and both, so there is # no role value meaning "off". # # Because no connector runs here, three of the context worker's constraints @@ -9,7 +9,7 @@ # # It does lend the pool this node's host memory. Capacity comes only from # processes that open a store handle, so without this every byte of the pool -# would be prefill-node DRAM caching prefill's own GPUs -- which is what the +# would be prefill-node DRAM caching prefill's own GPUs, which is what the # native host tier already does. Lending is not a connector: this engine still # never reads or writes the store, and keeps its cache transceiver for the # prefill-to-decode handoff. @@ -19,7 +19,7 @@ mooncake_donation: master_server_address: file://__LOG_DIR__/master.addr # Charged per server process, not per rank as global_segment_size is. Two # servers on one node lend twice this, and it competes with this node's own - # kv_cache_config.host_cache_size -- size the two together. + # kv_cache_config.host_cache_size, so size the two together. segment_size: 640GiB protocol: rdma diff --git a/mooncake_disagg/mooncake_api_surface_test.py b/mooncake_disagg/mooncake_api_surface_test.py index a359b7ea5903..4c0129f6243f 100644 --- a/mooncake_disagg/mooncake_api_surface_test.py +++ b/mooncake_disagg/mooncake_api_surface_test.py @@ -1,12 +1,11 @@ #!/usr/bin/env python3 """Exercise every MooncakeDistributedStore method the connector calls. -``mooncake_smoke_test.py`` only proves the install loads and can round-trip a -byte string. The connector's hot path never uses ``put``/``get``: it registers -the KV pools and then moves pages with the ``batch_*_multi_buffers`` zero-copy -calls. Those are the calls whose signatures could drift between wheel versions, -so this checks them against real registered GPU memory, in the same order -``worker.py`` uses them. +`mooncake_smoke_test.py` only proves the install loads and can round-trip a byte +string. The connector's hot path never uses `put` or `get`: it registers the KV +pools and then moves pages with the `batch_*_multi_buffers` zero-copy calls. +Those signatures could drift between wheel versions, so this checks them against +real registered GPU memory, in the same order `worker.py` uses them. Needs a running mooncake_master and MOONCAKE_CONFIG_PATH, same as the connector. """ @@ -57,8 +56,8 @@ def parse_size(value): # Stand in for a KV pool. PageAddressing.page_buffers returns one address per # layer-group region, so a page is scattered across REGIONS buffers rather than -# contiguous -- which is why the batch calls take list[list[int]]. Model two -# strided regions so the scatter-gather path is actually exercised. +# contiguous, which is why the batch calls take list[list[int]]. Two strided +# regions here so the scatter-gather path is actually exercised. PAGES = 8 REGIONS = 2 REGION_BYTES = 128 * 1024 diff --git a/mooncake_disagg/mooncake_smoke_test.py b/mooncake_disagg/mooncake_smoke_test.py index 214ac0c46313..db0b1dc8dee9 100644 --- a/mooncake_disagg/mooncake_smoke_test.py +++ b/mooncake_disagg/mooncake_smoke_test.py @@ -1,10 +1,10 @@ #!/usr/bin/env python3 """Prove that a Mooncake install can actually serve the mooncake-store connector. -Mirrors the ``store.setup(...)`` call in -``tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py`` and then +Mirrors the `store.setup` call in +`tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py` and then does one round trip, so a pass here means the connector's own startup path will -work. Reads the same ``MOONCAKE_CONFIG_PATH`` file the connector reads. +work. Reads the same `MOONCAKE_CONFIG_PATH` file the connector reads. """ import json diff --git a/mooncake_usage.md b/mooncake_usage.md index a12b55ddee40..f2a4e8d84286 100644 --- a/mooncake_usage.md +++ b/mooncake_usage.md @@ -3,8 +3,9 @@ The `mooncake-store` connector publishes KV cache pages into a shared, content-addressed pool in host DRAM, so a prefix computed by one engine can be replayed by another. It is a KV cache *connector*, unrelated to the Mooncake -*transfer engine* that the cache transceiver can use for prefill/decode handoff -— different component, different config, and they are usually not both in play. +*transfer engine* that the cache transceiver can use for prefill/decode handoff: +a different component with a different config, and the two are rarely both in +play. Use it when local block reuse leaves reuse on the table: several context instances that see the same prefixes, prefixes that should outlive a restart, or @@ -35,7 +36,7 @@ python3 -c "from mooncake.store import MooncakeDistributedStore; print('ok')" ``` `disaggr_torch.slurm` runs this per node automatically when a worker config -mentions `mooncake-store`. Images built from this branch have it baked in. +mentions `mooncake-store`, and images built from this repo have it baked in. `mooncake_disagg/README.md` §2 explains why it is this awkward. ## 2. Configure @@ -47,14 +48,14 @@ Three ways to get there, in increasing order of how much you have to arrange: | Deployment | Master | |---|---| -| One `trtllm-serve`, own pool — **including the SLURM harness** | `mooncake_store: {launch_master: true}` — the server starts it | +| One `trtllm-serve`, own pool, **including the SLURM harness** | `mooncake_store: {launch_master: true}`, so the server starts it | | Several engines, or a pool that outlives them | `trtllm-serve mooncake_master --address_file P`, then `mooncake_store: {master_server_address: file://P}` | | An externally provisioned pool | Nothing in the config: an inherited `MOONCAKE_CONFIG_PATH` wins over `mooncake_store` and says so in the log | The first two make `trtllm-serve` render the client config and export `MOONCAKE_CONFIG_PATH` itself. Nothing outside `trtllm-serve` starts a master, -writes a JSON config or picks an HCA — the SLURM harness included, which is why -`disaggr_torch.slurm` now only installs the bindings and tells the configs which +writes a JSON config or picks an HCA, the SLURM harness included: all +`disaggr_torch.slurm` does is install the bindings and tell the configs which directory the run is in. `mooncake_disagg/README.md` §4 covers running a master as its own SLURM job, for the second row. @@ -115,16 +116,16 @@ Per-process environment, on the workers that open a handle: |---|---| | `TRTLLM_MOONCAKE_STORE_ROLE` | `producer` / `consumer` / `both` | | `TRTLLM_MOONCAKE_STORE_PREFIX` | Cache namespace. Bump it after any change to page layout or contents. | -| `TRTLLM_MOONCAKE_STORE_MODEL_KEY` | Defaults to the checkpoint directory's basename — set it explicitly for anything long-lived. | +| `TRTLLM_MOONCAKE_STORE_MODEL_KEY` | Defaults to the checkpoint directory's basename. Set it explicitly for anything long-lived. | Pool capacity comes only from processes that open a store handle, so a -prefill-only connector gives a prefill-only pool — which caches prefill's GPUs -in prefill's own DRAM, largely duplicating the native host offload. Ask the +prefill-only connector gives a prefill-only pool, caching prefill's GPUs in +prefill's own DRAM and largely duplicating the native host offload. Ask the generation servers to lend their memory and the pool spans both sides while their engines stay connector-free: ```yaml -# generation worker — no connector, memory only +# generation worker: no connector, memory only mooncake_donation: master_server_address: file:///$WORK_DIR/master.addr segment_size: 320GiB # per server process, not per rank @@ -141,7 +142,7 @@ follow it. Note the granularity: `global_segment_size` is charged per *rank*, `segment_size` per *server process*. Two generation servers on one node lend `segment_size` each. It is charged to the process, so it competes with that -node's own `kv_cache_config.host_cache_size` — size the two together. +node's own `kv_cache_config.host_cache_size`. Size the two together. A node running no server can lend as its own command, which is also how the pool gets memory from a machine with no GPUs at all: @@ -151,30 +152,28 @@ trtllm-serve mooncake_donor --master_server_address file://$WORK_DIR/master.addr --segment_size 160GiB --protocol rdma --device_name mlx5_0 ``` -## 3. Partial reuse must be off — now enforced +## 3. Partial reuse is forced off This is the one setting that decides whether the feature works at all. -The store is addressed by whole blocks. The connector is handed -`num_computed_tokens`, the device match, and offers only blocks beyond it — but -it can only continue from a block boundary, so when the device match lands -mid-block it declines the lookup entirely. `enable_partial_reuse=true` is -precisely what makes the device match land mid-block, so it trades part of one -block of device reuse for *every* stored block of the remaining prefix. +The store is addressed by whole blocks. The connector is handed the device match +as `num_computed_tokens` and offers only blocks beyond it, but it can continue +only from a block boundary, so when the device match lands mid-block it declines +the lookup entirely. `enable_partial_reuse=true` is precisely what makes the +device match land mid-block, so it trades part of one block of device reuse for +*every* stored block of the remaining prefix. On MiniMax-M3 that declined 97.2% +of lookups, leaving a 1.6 TB pool measuring as if it were not there. -On MiniMax-M3 that guard declined **97.2% of lookups**. The pool was never -asked, and a 1.6 TB pool measured as if it were not there. - -`py_executor_creator` now forces `enable_partial_reuse=false` whenever this -connector is configured, and says so: +`py_executor_creator` therefore forces `enable_partial_reuse=false` whenever +this connector is configured, and says so: ``` Disabling partial reuse: it is not usable with the mooncake-store connector... ``` -Nothing to set; the warning fires even from the default (`true`), and is the -confirmation that the coercion ran. The field is otherwise untouched, so -configs that already set `false` are unaffected. +There is nothing to set. The warning fires even from the default of `true`, and +is the confirmation that the coercion ran. Configs that already set `false` are +unaffected. ## 4. Verify a run @@ -184,7 +183,7 @@ Startup, at INFO, on every context worker: grep -h "mooncake-store" /3_output_CTX_*.log | head -40 ``` -`registered layout: ... bytes/page=...` is the line to keep — pool sizing +`registered layout: ... bytes/page=...` is the line to keep, since pool sizing depends on it, and `window=None` confirms no sliding-window group (one would have aborted startup). @@ -192,15 +191,16 @@ Then check that the pool spans the hosts you expect. `disaggr_torch.slurm` writes the per-segment breakdown to `/9_mooncake_summary.log`; a single host means a prefill-only pool. Pool occupancy and eviction come from the master's own log, -`$TRTLLM_MOONCAKE_RUN_DIR/mooncake_master.log` — under the harness that is +`$TRTLLM_MOONCAKE_RUN_DIR/mooncake_master.log`, which under the harness is `/mooncake_master.log`. The startup line reports the path either way. **Which reuse number counts store hits:** per-request stats (`reused_blocks_per_request`, `kv_cache_hit_rate_per_request`) **do**; `/prometheus/metrics` iteration counters (`kv_cache_iter_reused_blocks`) **do -not** — those come from the local reuse tree. So store hits ≈ per-request reuse -− local-tree reuse. All of this needs `enable_iter_perf_stats`, -`enable_iter_req_stats` and `return_perf_metrics`, which all default to false. +not**, since those come from the local reuse tree. Store hits are therefore +roughly per-request reuse minus local-tree reuse. All of this needs +`enable_iter_perf_stats`, `enable_iter_req_stats` and `return_perf_metrics`, +which all default to false. ## 5. What it measured @@ -219,16 +219,16 @@ MiniMax-M3-NVFP4 on GB300, 1 context server (TP=2) + 2 generation servers A 61-point gap between the reuse the workload allowed and the reuse the system achieved closed to 3 points. For comparison, native host offload on the same -workload reached 35.59% actual hit at 318.14 tok/s — it wrote 1.83 TB to host -and read 32.5 GB back, behaving as a write-only tier. +workload reached 35.59% actual hit at 318.14 tok/s: it wrote 1.83 TB to host and +read 32.5 GB back, behaving as a write-only tier. Where the reuse comes from, in steady state (attribution counters, c50): **~95% of all reuse is served by the pool and ~5% by the device cache.** The -residual ~3–5% of misses are prefixes never written by anyone, which no store -can serve. Blocks stranded behind a contiguity gap measured exactly zero, as did +residual 3-5% of misses are prefixes never written by anyone, which no store can +serve. Blocks stranded behind a contiguity gap measured exactly zero, as did unattributed blocks. -**Peak is at c50, not higher.** By c70 the pool runs 85–90% full with active +**Peak is at c50, not higher.** By c70 the pool runs 85-90% full with active eviction, hit rate falls to 86% and throughput with it. Concurrency headroom is a function of pool size; size the pool for the working set rather than assuming the c50 result scales. @@ -240,7 +240,7 @@ the c50 result scales. the same as absent. - **The key namespace pins world size, rank, `tokens_per_block`, layer groups and `bytes_per_page`.** Change tensor parallelism and every stored page - becomes unreachable — a miss, not an error. + becomes unreachable, as a miss rather than an error. - **No build hash in the key.** After changing page layout or contents, bump `TRTLLM_MOONCAKE_STORE_PREFIX` or restart the master. - **Loads are synchronous** (`start_load_kv`, before the forward pass), so every diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index cda431c99d61..378dbda6f4d3 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -21,37 +21,29 @@ This is a different component from the Mooncake transfer engine that the C++ cache transceiver uses for disaggregated prefill/decode handoff: that moves KV point to point between two known peers, while this one publishes pages into a -pool addressed by content. The two compose -- a context server can write pages +pool addressed by content. The two compose, so a context server can write pages here and still hand off over NIXL. -Requires ``KVCacheManagerV2``, which is the manager that can describe its pools -to a connector (``register_kv_cache_layout``), and the Mooncake Python bindings -(``pip install mooncake-transfer-engine``). +Requires `KVCacheManagerV2`, the manager that can describe its pools to a +connector through `register_kv_cache_layout`, and the Mooncake Python bindings +(`pip install mooncake-transfer-engine`). Enable it with:: kv_connector_config = KvCacheConnectorConfig(connector="mooncake-store") -with ``MOONCAKE_CONFIG_PATH`` pointing at a Mooncake JSON config. +with `MOONCAKE_CONFIG_PATH` pointing at a Mooncake JSON config. Describing the +pool in `KvCacheConnectorConfig.mooncake_store` instead lets `trtllm-serve` +provision it during bringup, so no external script has to; see `master.py`. -Describing the pool in ``KvCacheConnectorConfig.mooncake_store`` instead lets -``trtllm-serve`` provision it during bringup -- resolving or launching the -master and writing that JSON itself -- so no external script has to. See -``master.py``. +Capacity comes only from processes that open a store handle, which in a +disaggregated deployment is the context servers alone. `donor.py` lends a +node's memory to the pool without giving it a connector. -Capacity, separately, comes only from processes that open a store handle, which -in a disaggregated deployment is the context servers alone. ``donor.py`` lends -a node's memory to the pool without giving it a connector, so the generation -nodes can hold cache they never read. - -By default the KV pools themselves are registered with Mooncake, so the store -reads and writes device memory and no copy is added. That needs the HCA to be -able to pin GPU pages -- GPUDirect RDMA, through ``nvidia_peermem`` or dma-buf. -Where it is missing, registration fails on every pool range and the connector -cannot start; ``"stage_through_host": true`` in the JSON config (or -``TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST=1``) then routes pages through a -pinned host buffer instead, so only host memory is ever registered. The stored -bytes are the same either way, so the two modes can share a pool. +By default the KV pools themselves are registered with Mooncake, which requires +GPUDirect RDMA. Where that is unavailable, `"stage_through_host": true` in the +JSON config, or `TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST=1`, routes pages +through a pinned host buffer instead; see `staging.py`. """ from .config import MooncakeStoreConnectorConfig, StoreRole, parse_size diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py index 089dfbcc99b7..712e34b3978e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py @@ -12,16 +12,16 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Turning a ``KvCacheLayout`` into addresses Mooncake can transfer. +"""Turning a `KvCacheLayout` into addresses Mooncake can transfer. -Mooncake's batch APIs take, per key, a list of ``(address, size)`` buffers. That +Mooncake's batch APIs take, per key, a list of `(address, size)` buffers. That is exactly the shape of a V2 page: a layer group's regions each contribute one -byte range at ``base + stride * page_index``, and the concatenation of those +byte range at `base + stride * page_index`, and the concatenation of those ranges in region order is the page's payload. -Region order is therefore load-bearing -- it is the value's serialization -- and -``build_kv_cache_layout_v2`` derives it from the allocator's own aggregation, so -it is stable for a given model and parallel layout. ``bytes_per_page`` goes into +Region order is therefore the value's serialization, and it is stable for a +given model and parallel layout because `build_kv_cache_layout_v2` derives it +from the allocator's own aggregation. `bytes_per_page` goes into the key namespace to keep a geometry change from being read as a valid page. """ @@ -33,7 +33,7 @@ def merge_intervals(intervals: Iterable[Tuple[int, int]]) -> List[Tuple[int, int]]: - """Collapse ``(start, end)`` byte ranges into a minimal disjoint cover. + """Collapse `(start, end)` byte ranges into a minimal disjoint cover. Registration is per range and a range may not be registered twice, but several regions routinely live inside one pool allocation: sliding-window @@ -53,7 +53,7 @@ def merge_intervals(intervals: Iterable[Tuple[int, int]]) -> List[Tuple[int, int class PageAddressing: - """Resolves ``(layer group, page index)`` to the byte ranges of that page.""" + """Resolves `(layer group, page index)` to the byte ranges of that page.""" def __init__(self, layout: KvCacheLayout): self._layout = layout @@ -95,11 +95,11 @@ def tokens_per_block(self) -> int: return self._layout.tokens_per_block def bytes_per_page(self, layer_group_id: int) -> int: - """Total payload size of one page of ``layer_group_id``.""" + """Total payload size of one page of `layer_group_id`.""" return self._bytes_per_page[layer_group_id] def num_slots(self, layer_group_id: int) -> int: - """Number of page slots addressable in ``layer_group_id``.""" + """Number of page slots addressable in `layer_group_id`.""" return self._num_slots[layer_group_id] def buffers(self, layer_group_id: int, page_index: int) -> Tuple[List[int], List[int]]: @@ -124,7 +124,7 @@ def buffers(self, layer_group_id: int, page_index: int) -> Tuple[List[int], List return addresses, sizes def registration_ranges(self) -> List[Tuple[int, int]]: - """Byte ranges to hand to ``register_buffer``, deduplicated and merged. + """Byte ranges to hand to `register_buffer`, deduplicated and merged. A region's slots are strided rather than packed, so the range covering it is the whole span from the first slot to the end of the last. Registering diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index d6fb47628c0f..008407106ff1 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -14,13 +14,13 @@ # limitations under the License. """Configuration for the Mooncake store KV cache connector. -Topology settings are read from the JSON file named by ``MOONCAKE_CONFIG_PATH``, +Topology settings are read from the JSON file named by `MOONCAKE_CONFIG_PATH`, the same file and environment variable the vLLM Mooncake store connector uses, so one deployment can point both engines at the same pool. -``KvCacheConnectorConfig`` carries no free-form dictionary, so the two settings -that are TensorRT-LLM's rather than Mooncake's -- the read/write role and the -key prefix -- are also taken from the environment. +`KvCacheConnectorConfig` carries no free-form dictionary, so the two settings +that are TensorRT-LLM's rather than Mooncake's, the read/write role and the key +prefix, are also taken from the environment. """ import json @@ -82,7 +82,7 @@ class StoreRole(Enum): """Which directions of traffic this engine is allowed to drive. - A disaggregated deployment typically runs context servers as ``both`` and + A disaggregated deployment typically runs context servers as `both` and leaves generation servers unconfigured: generated tokens are rarely a reused prefix, so writing them costs bandwidth for no hit rate. """ @@ -103,7 +103,7 @@ def saves(self) -> bool: def parse_size(value: Any) -> int: - """Accept either a byte count or a suffixed string such as ``"4GiB"``.""" + """Accept either a byte count or a suffixed string such as `"4GiB"`.""" if isinstance(value, bool): raise ValueError(f"expected a size, got {value!r}") if isinstance(value, int): @@ -123,16 +123,13 @@ def parse_size(value: Any) -> int: def provisioned_config_path() -> Optional[str]: """The client config a server on this node rendered, if there is one. - ``provision_pool`` writes one and exports ``MOONCAKE_CONFIG_PATH``, which - reaches the ranks the LLM constructor spawns, since they inherit that - environment. Ranks the launcher started instead -- one task per rank, which - is how a server spanning several GPUs is launched under a scheduler -- were - already running by then and never see it. Reading the config back from the - run directory is what lets those ranks join the pool their own leader - provisioned. + `provision_pool` writes one and exports `MOONCAKE_CONFIG_PATH`, which the + ranks the LLM constructor spawns inherit. Ranks an external launcher + started, one task per rank, were already running by then and never see it, + so they read the config back from the run directory instead. - Only possible when the deployment named that directory: it otherwise - defaults to a per-process temporary one, which no other rank could read. + Only possible when the deployment named that directory, since it otherwise + defaults to a per-process temporary one that no other rank could read. """ run_dir = os.getenv(RUN_DIR_ENV) if not run_dir: @@ -164,11 +161,10 @@ class MooncakeStoreConnectorConfig: #: RPC without bounding how much a request may transfer. transfer_batch_size: int = 64 #: Pass pages through a pinned host buffer instead of registering the KV - #: pools with Mooncake. Costs a copy in each direction and buys independence - #: from GPUDirect RDMA, without which registering device memory fails - #: outright. Leave off wherever the pool can reach GPU memory. + #: pools with Mooncake. Costs a copy each way, but works without GPUDirect + #: RDMA, which registering device memory requires. stage_through_host: bool = False - #: Ceiling on the pinned allocation per direction when staging. The pool is + #: Ceiling on the pinned allocation per direction when staging. Slots are #: sized from the layout's largest page, so this caps how many pages may be #: in flight rather than how large one may be. staging_buffer_bytes: int = DEFAULT_STAGING_BUFFER_SIZE @@ -222,14 +218,14 @@ def from_env() -> "MooncakeStoreConnectorConfig": "Mooncake JSON config (metadata_server, master_server_address, " "protocol, device_name, global_segment_size, local_buffer_size), " "or kv_connector_config.mooncake_store set so the server renders " - f"one -- into ${RUN_DIR_ENV} if this rank was started by the " + f"one, into ${RUN_DIR_ENV} if this rank was started by the " "launcher rather than spawned by the server." ) config = MooncakeStoreConnectorConfig.from_file(path) return config.with_env_overrides() def with_env_overrides(self) -> "MooncakeStoreConnectorConfig": - """Apply ``TRTLLM_MOONCAKE_STORE_*`` on top of the file's settings.""" + """Apply `TRTLLM_MOONCAKE_STORE_*` on top of the file's settings.""" import dataclasses updates: dict[str, Any] = {} diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py index b27f6fbfddfe..42a765de62bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -14,30 +14,23 @@ # limitations under the License. """Put a node's host memory into a Mooncake pool without reading or writing it. -Pool capacity comes only from processes that open a store handle: ``setup`` -registers ``global_segment_size`` bytes of the caller's host memory and the -master then places blocks in it. In a disaggregated deployment only the context -servers configure the connector, so only they call ``setup``, and the pool is -entirely prefill-node memory -- which makes the store a -prefill-DRAM-caches-prefill-GPU tier, overlapping what TensorRT-LLM's own host -offload already does. +Pool capacity comes only from processes that open a store handle: `setup` +registers `global_segment_size` bytes of the caller's host memory and the master +then places blocks in it. In a disaggregated deployment only the context servers +configure the connector, so the pool is entirely prefill-node memory, which +overlaps what TensorRT-LLM's own host offload already does. Donating alongside a generation server puts that node's memory into the same -pool. Prefill then writes blocks that land on decode-side DRAM and reads them -back, while the generation engine stays free of any connector: it neither reads -nor writes the store, so it keeps its single cache transceiver for the -prefill-to-decode handoff. - -Donation is deliberately not a ``StoreRole``. The roles describe an engine's -traffic -- ``producer`` writes, ``consumer`` reads, ``both`` does both -- and -none of them means "contribute memory only", so attaching a connector to a -generation server to get its DRAM into the pool would also start it reading or -writing. Capacity and traffic are separate concerns, which is why this holds a -handle of its own rather than being a setting on the connector. - -The memory is charged to the donating process, so it competes with anything -else on the node -- a generation server's ``kv_cache_config.host_cache_size`` -above all. Size the two together. +pool, so prefill writes blocks that land on decode-side DRAM. The generation +engine stays free of any connector and keeps its single cache transceiver for +the prefill-to-decode handoff. + +Donation is not a `StoreRole`. The roles describe an engine's traffic and none +of them means "contribute memory only", so capacity and traffic stay separate +concerns and a donor holds a store handle of its own. + +The memory is charged to the donating process, so size it together with that +node's `kv_cache_config.host_cache_size`. """ import contextlib @@ -56,8 +49,7 @@ "maybe_donate_segment", ] -#: A donor never transfers, so its transfer buffer is dead weight; ``setup`` -#: still rejects a zero one. +#: A donor never transfers, but `setup` rejects a zero-sized transfer buffer. DEFAULT_DONOR_LOCAL_BUFFER_SIZE = 64 * 1024**2 @@ -71,15 +63,14 @@ def donate_segment( local_buffer_size: int = DEFAULT_DONOR_LOCAL_BUFFER_SIZE, hostname: Optional[str] = None, ) -> Iterator[str]: - """Hold ``segment_size`` bytes of this node's memory in the pool. + """Hold `segment_size` bytes of this node's memory in the pool. - Yields the host the segment is registered under, which is what the master - and the engines reading from it identify the capacity by. + Yields the host the segment is registered under, which is how the master + and the engines reading from it identify the capacity. - The handle is held for the duration: dropping it unmounts the segment, and - the master starts reporting the blocks that lived in it as lost. So the - caller must stay inside this context for as long as the capacity is meant - to exist, which for a donor is its whole run. + Dropping the store handle unmounts the segment and the master starts + reporting the blocks that lived in it as lost, so the caller must stay + inside this context for as long as the capacity is meant to exist. """ try: from mooncake.store import MooncakeDistributedStore @@ -92,10 +83,8 @@ def donate_segment( host = hostname or local_address() donated = f"{segment_size / 1024 ** 3:.1f}GiB" - # Every argument is echoed, with the byte counts spelled out next to the - # human-readable form: a segment that is a thousandth of the intended size - # is a size string parsed wrong, and it otherwise shows up only as a pool - # that evicts far too eagerly, days later. + # Byte counts are spelled out next to the human-readable form. A misparsed + # size string otherwise surfaces only as a pool that evicts far too eagerly. logger.info( f"mooncake-store: lending memory to the pool at {master_server_address} " f"as capacity only, no reads or writes: host={host} " @@ -136,8 +125,7 @@ def donate_segment( try: yield host finally: - # Explicit because the segment stays mounted for as long as anything - # references the handle, and "as long as this context" is the contract. + # The segment stays mounted while anything references the handle. del store logger.info( f"mooncake-store: withdrew the {donated} lent from {host}; the " @@ -150,19 +138,17 @@ def maybe_donate_segment(donation: Any) -> Iterator[Optional[str]]: """Lend memory for this process's lifetime if the config asked to. Args: - donation: A ``MooncakeDonationConfig``, or ``None`` to do nothing -- - which is every deployment that does not lend memory, so callers - need no condition of their own. + donation: A `MooncakeDonationConfig`, or `None` to do nothing, so + callers need no condition of their own. - Yields the host the segment is registered under, or ``None``. + Yields the host the segment is registered under, or `None`. """ if donation is None: yield None return - # A generation server has no other reason to resolve a master, so the - # address it lends against is worth saying out loud before the wait: this - # is the one place bringup blocks on a component from a different job. + # Bringup blocks here on a master that may belong to a different job, so + # name the address before waiting on it. logger.info( "mooncake-store: mooncake_donation is set, so this server lends host " f"memory to the pool at {donation.master_server_address} without using " @@ -171,15 +157,13 @@ def maybe_donate_segment(donation: Any) -> Iterator[Optional[str]]: master_address = resolve_master_address( donation.master_server_address, master_timeout() ) - # Checked before setup so an absent master reads as one, rather than as - # the status code setup returns for everything. + # Checked before setup so an absent master is reported as such, rather than + # as the status code setup returns for every kind of failure. wait_for_master(master_address) with donate_segment( master_server_address=master_address, segment_size=parse_size(donation.segment_size), protocol=donation.protocol, - # Which HCAs this node has is the node's business, so a config that - # leaves it open stays usable on every node type in the deployment. device_name=resolve_device_name(donation.protocol, donation.device_name), metadata_server=donation.metadata_server, ) as host: diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py index fe5ea6965e79..56609877843d 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -14,12 +14,12 @@ # limitations under the License. """Block identity and store key naming for the Mooncake store connector. -``KVCacheManagerV2`` exposes no block hashes to a connector -- ``RequestData`` -reports them empty -- so content identity is derived here instead. The chain is +`KVCacheManagerV2` exposes no block hashes to a connector, since `RequestData` +reports them empty, so content identity is derived here instead. The chain is the standard one: a block's hash covers its own tokens *and* every token before it, so a key can only be reused by a request whose prefix is byte-identical. -A key is ``/``. The namespace pins down everything that +A key is `/`. The namespace pins down everything that would make the stored bytes mean something different: the model, the shard that produced them, the layer group inside that shard, the tokens each page holds and how many bytes a page is. Anything that changes those reads as a cache miss @@ -77,7 +77,7 @@ def hashes(self) -> Sequence[bytes]: return self._hashes def extend(self, tokens: Sequence[int]) -> Sequence[bytes]: - """Grow the chain to cover every full block of ``tokens``. + """Grow the chain to cover every full block of `tokens`. Args: tokens: The request's complete token list, prompt first. Must be an diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index 442fe3135d47..94c4d8d7a928 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -14,22 +14,20 @@ # limitations under the License. """Bring the Mooncake store's pool up as part of a server's own startup. -The connector needs two things that are not the engine's to produce: a -reachable ``mooncake_master``, and a JSON client config named by -``MOONCAKE_CONFIG_PATH`` that points every worker at it. Both were the SLURM -harness's job, which left a single ``trtllm-serve`` unable to use the connector -without borrowing that harness. - -``provision_pool`` does the same work inside the serving process. It resolves -the master -- launching one here, or checking that the configured one answers --- renders the client config, and exports ``MOONCAKE_CONFIG_PATH``, which -reaches the ranks because the LLM constructor spawns them from this process. -Everything it started is torn down when the context exits. - -A master launched here lives and dies with the server, so it is only right for -one engine talking to its own pool. Several engines sharing a pool, or a pool -meant to survive a restart, need a master with its own lifetime, named by -``master_server_address``. +The connector needs two things that are not the engine's to produce: a reachable +`mooncake_master`, and a JSON client config named by `MOONCAKE_CONFIG_PATH` that +points every worker at it. + +`provision_pool` does that work inside the serving process. It resolves the +master, either launching one here or checking that the configured one answers, +renders the client config, and exports `MOONCAKE_CONFIG_PATH`, which reaches the +ranks because the LLM constructor spawns them from this process. Everything it +started is torn down when the context exits. + +A master launched here lives and dies with the server, so it suits one engine +talking to its own pool. Several engines sharing a pool, or a pool meant to +survive a restart, need a master with its own lifetime named by +`master_server_address`. """ import contextlib @@ -59,7 +57,7 @@ "wait_for_master", ] -#: Override the binary that ``launch_master`` runs. +#: Override the binary that `launch_master` runs. MASTER_BINARY_ENV = "TRTLLM_MOONCAKE_MASTER_BINARY" #: How long to wait for a master to accept connections, in seconds. MASTER_TIMEOUT_ENV = "TRTLLM_MOONCAKE_MASTER_TIMEOUT" @@ -67,15 +65,13 @@ DEFAULT_MASTER_TIMEOUT = 60.0 MASTER_LOG_NAME = "mooncake_master.log" #: Name a launched master's address is always published under in the run -#: directory, so "which master is this pool on" is answerable from the logs of -#: a run that named no address file. +#: directory, so even a run that named no address file records its pool. MASTER_ADDRESS_NAME = "master.addr" -#: Prefix that makes ``master_server_address`` name a file holding the address +#: Prefix that makes `master_server_address` name a file holding the address #: rather than the address itself. ADDRESS_FILE_SCHEME = "file://" -#: Lines of the master's log to quote when startup fails. Its last words are -#: usually the whole diagnosis -- a port in use, a bad flag -- and they are -#: otherwise in a file the reader has to be told exists. +#: Lines of the master's log to quote when startup fails, since its last words +#: (a port in use, a bad flag) are usually the whole diagnosis. LOG_TAIL_LINES = 20 @@ -98,9 +94,8 @@ def _log_tail(path: str, lines: int = LOG_TAIL_LINES) -> str: def local_address() -> str: """The address this host is known by inside the pool. - Deliberately the same derivation the connector worker uses for its own - hostname, so the master and the segments registering with it agree on - which host they are on. + Uses the same derivation as the connector worker's own hostname, so the + master and the segments registering with it agree on which host they are on. """ try: return socket.gethostbyname(socket.gethostname()) @@ -122,7 +117,7 @@ def master_timeout() -> float: def _split_address(address: str) -> Optional[Tuple[str, int]]: - """Split ``host:port``, or return ``None`` if it is not in that form.""" + """Split `host:port`, or return `None` if it is not in that form.""" host, separator, port = address.rpartition(":") if not separator or not port.isdigit(): return None @@ -130,14 +125,14 @@ def _split_address(address: str) -> Optional[Tuple[str, int]]: def resolve_master_address(address: str, timeout: float) -> str: - """Read a ``file://`` address through, and pass anything else along. - - A master with its own lifetime is on whichever host its scheduler gave - it, which is not known when the worker configs are written. Naming the - file it publishes to instead keeps the address out of the config and out - of a launch script: ``trtllm-serve mooncake_master --address-file`` writes - it, every worker's ``master_server_address`` names the same path, and the - wait here is also the wait for the master to exist at all. + """Read a `file://` address through, and pass anything else along. + + A master with its own lifetime runs on whichever host its scheduler gave + it, which is not known when the worker configs are written. Naming the file + it publishes to keeps the address out of both the config and the launch + script: `trtllm-serve mooncake_master --address-file` writes it, every + worker's `master_server_address` names the same path, and the wait here + doubles as the wait for the master to exist at all. """ if not address.startswith(ADDRESS_FILE_SCHEME): return address @@ -148,8 +143,6 @@ def resolve_master_address(address: str, timeout: float) -> str: announced = started logger.info(f"mooncake-store: reading the master's address from {path}") while True: - # Written whole by the master command, so a non-empty file is a - # complete address rather than a prefix of one. try: published = open(path).read().strip() except FileNotFoundError: @@ -160,8 +153,8 @@ def resolve_master_address(address: str, timeout: float) -> str: now = time.monotonic() if now - announced >= 5.0: announced = now - # Waiting for a master in another job is the normal case here, so - # this is progress rather than trouble -- but only if it is said. + # Waiting on a master in another job is normal here, so say so + # rather than letting the wait look like a hang. logger.info( f"mooncake-store: no master address in {path} yet " f"({now - started:.0f}s of {timeout:g}s); waiting for the " @@ -188,14 +181,13 @@ def _wait_until_accepting( """Block until the master accepts connections, and say how long it took. A worker that opens its store handle before the master is listening fails - outright, so the port -- not the presence of a process -- is what the - ordering has to wait on. When the master is ours, its exit is checked first - each pass, so a master that died is reported as that rather than as a + outright, so the ordering has to wait on the port rather than on the + presence of a process. When the master is ours, its exit is checked first + each pass, so a master that died is reported as such rather than as a timeout. - The wait is narrated while it happens: silence here is indistinguishable - from a hang somewhere else in bringup, and this is one of the two places a - Mooncake deployment stalls. + The wait is narrated as it happens, since silence here is + indistinguishable from a hang elsewhere in bringup. """ started = time.monotonic() deadline = started + timeout @@ -238,10 +230,10 @@ def _highest_rate_ib_devices(sysfs_root: Optional[str] = None) -> List[str]: """The active InfiniBand devices on the compute fabric, fastest first. A node's HCAs are not interchangeable. On GB300 six are exposed, of which - four run at 800Gb/s -- two per NUMA node, one per GPU -- while the rest - share a PCI device with an Ethernet port and are the storage or management - adapter. Taking every device at the highest rate picks the compute fabric - on any node type, where a hardcoded name would be wrong on the next one. + four run at 800Gb/s (two per NUMA node, one per GPU) while the rest share a + PCI device with an Ethernet port and serve storage or management. Taking + every device at the highest rate picks the compute fabric on any node type, + where a hardcoded name would be wrong on the next one. """ sysfs_root = sysfs_root or IB_SYSFS_ROOT rated: Dict[str, int] = {} @@ -281,9 +273,9 @@ def resolve_device_name(protocol: str, """The RDMA devices to transfer over, detected if the config left it open. Which HCAs a node has is a property of the node, not of the deployment, so - requiring it in a config makes that config specific to one machine type. - Detecting it keeps ``protocol: rdma`` portable, and leaving ``device_name`` - set overrides this for a node where the choice has to be made by hand. + requiring it in a config would tie that config to one machine type. + Detecting it keeps `protocol: rdma` portable; setting `device_name` + overrides the detection. """ if configured or protocol != "rdma": return configured @@ -305,15 +297,14 @@ def resolve_device_name(protocol: str, def wait_for_master(master_address: str, timeout: Optional[float] = None) -> Optional[float]: - """Block until the master at ``master_address`` accepts connections. + """Block until the master at `master_address` accepts connections. - Every user of a pool it did not start wants this: reaching a master that - is not there otherwise fails deep inside ``store.setup``, in every rank, - after the model has loaded, as a status code. One socket beforehand turns - that into a line that names the address and the wait. + Reaching a master that is not there otherwise fails deep inside + `store.setup`, in every rank, after the model has loaded, as a bare status + code. One socket beforehand turns that into a line naming the address. - Returns how long it took, or ``None`` if the address was not in - ``host:port`` form and could not be checked. + Returns how long it took, or `None` if the address was not in `host:port` + form and could not be checked. """ timeout = master_timeout() if timeout is None else timeout endpoint = _split_address(master_address) @@ -336,10 +327,9 @@ def _client_config(pool: Any, device_name: Optional[str] = None) -> Dict[str, Any]: """Render the Mooncake client config for a pool. - The schema is vLLM's, so one pool can serve both engines. ``role`` is - written as ``both`` because the file describes the pool; which directions - of traffic a given process drives is its own - ``TRTLLM_MOONCAKE_STORE_ROLE``. + The schema is vLLM's, so one pool can serve both engines. `role` is written + as `both` because the file describes the pool; the directions of traffic a + given process drives come from its own `TRTLLM_MOONCAKE_STORE_ROLE`. """ config: Dict[str, Any] = { "metadata_server": pool.metadata_server, @@ -354,8 +344,8 @@ def _client_config(pool: Any, } if pool.cache_prefix is not None: config["cache_prefix"] = pool.cache_prefix - # Left out when unset so the connector's own default applies, rather than - # restating it here for the two to drift apart. + # Left out when unset so the connector's own default applies instead of a + # second copy of it here. if pool.staging_buffer_bytes is not None: config["staging_buffer_bytes"] = pool.staging_buffer_bytes return config @@ -363,7 +353,7 @@ def _client_config(pool: Any, @dataclass class LaunchedMaster: - """A ``mooncake_master`` owned by this process.""" + """A `mooncake_master` owned by this process.""" process: subprocess.Popen address: str @@ -396,11 +386,11 @@ def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: host = local_address() log_path = os.path.join(run_dir, MASTER_LOG_NAME) - # mooncake_master logs through glog, which writes files under /tmp unless - # told otherwise, so without GLOG_logtostderr the log below stays empty. - # GLOG_v=1 adds the per-RPC lines showing segments registering and keys - # moving, which is the only view of the pool's side of the conversation - # short of scraping the metrics port. + # glog writes to files under /tmp unless redirected, so without + # GLOG_logtostderr the log opened below stays empty. GLOG_v=1 adds the + # per-RPC lines showing segments registering and keys moving, which is the + # only view of the pool's own side of the conversation short of scraping + # the metrics port. env = dict(os.environ, GLOG_logtostderr="1") env.setdefault("GLOG_v", "1") command = [ @@ -439,20 +429,18 @@ def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: @contextlib.contextmanager def _published_address(address: str, paths: Sequence[str]) -> Iterator[None]: - """Write ``address`` to every path for the life of the context. + """Write `address` to every path for the life of the context. Publishing is how anything else finds this master: a donor or a second - server names the path in ``master_server_address`` as ``file://`` - and reads it back. Retracting on the way out is as important as writing, - since an address that outlives its master sends the next run's workers to - a port with nothing behind it. + server names the path as `file://` in `master_server_address`. + Retracting on the way out matters as much as writing, since an address that + outlives its master sends the next run's workers to a dead port. """ for path in paths: directory = os.path.dirname(path) if directory: os.makedirs(directory, exist_ok=True) - # Renamed into place so a reader sees either nothing or the whole - # address. A half-written one would be dialed as if it were real. + # Renamed into place so a reader never sees a partial address. staging = f"{path}.partial" with open(staging, "w") as handle: handle.write(f"{address}\n") @@ -470,8 +458,7 @@ def _published_address(address: str, paths: Sequence[str]) -> Iterator[None]: def _address_files(run_dir: str, extra: Optional[str] = None) -> List[str]: """Where a master this process starts should publish its address. - Always the run directory, so a reader who was told nothing can still find - out which master a run used, plus wherever the deployment asked for. + Always the run directory, plus wherever the deployment asked for. """ paths = [os.path.join(run_dir, MASTER_ADDRESS_NAME)] if extra and os.path.abspath(extra) not in {os.path.abspath(p) for p in paths}: @@ -485,14 +472,13 @@ def running_master( ) -> Iterator[LaunchedMaster]: """Run a master whose lifetime is this process's rather than an engine's. - ``provision_pool`` covers the server that owns its pool. Everything else -- - several engines on one pool, a pool that has to survive a restart -- needs - the master somewhere that is not any of them, which is what this is for. + `provision_pool` covers the server that owns its pool. Several engines on + one pool, or a pool that has to survive a restart, need the master + somewhere that is not any of them. - ``address_file`` receives ``host:port`` once the master answers, so the - workers can name the file instead of an address nobody knows until the - scheduler has placed this process. One is written to ``run_dir`` either - way. + `address_file` receives `host:port` once the master answers, so workers can + name the file instead of an address nobody knows until the scheduler has + placed this process. One is written to `run_dir` either way. """ os.makedirs(run_dir, exist_ok=True) master = _launch_master(pool, run_dir) @@ -506,15 +492,15 @@ def running_master( @contextlib.contextmanager def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optional[str]]: - """Make ``pool`` reachable and name it in this process's environment. + """Make `pool` reachable and name it in this process's environment. - Yields the path of the client config written, or ``None`` when an inherited - ``MOONCAKE_CONFIG_PATH`` was left in charge. + Yields the path of the client config written, or `None` when an inherited + `MOONCAKE_CONFIG_PATH` was left in charge. Args: - pool: A ``MooncakeStoreConfig``. + pool: A `MooncakeStoreConfig`. run_dir: Where to write the client config and the master's log. - Defaults to ``TRTLLM_MOONCAKE_RUN_DIR``, else a temporary directory + Defaults to `TRTLLM_MOONCAKE_RUN_DIR`, else a temporary directory that is removed on exit. """ inherited = os.getenv(CONFIG_PATH_ENV) @@ -546,9 +532,8 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona if pool.launch_master: master = _launch_master(pool, run_dir) master_address = master.address - # Even a master that only this server uses publishes: it is - # how its donors reach it, and how the log of a finished run - # still says which pool it was. + # Published even when only this server uses it, since that is + # how its donors reach it. stack.enter_context( _published_address( master_address, _address_files(run_dir, pool.master_address_file) @@ -567,9 +552,9 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona resolve_device_name(pool.protocol, pool.device_name)) with open(config_path, "w") as handle: json.dump(config, handle, indent=2) - # This reaches the ranks the LLM constructor spawns, which - # inherit it. A rank the launcher started instead was already - # running, and reads the config out of the run directory -- see + # Inherited by the ranks the LLM constructor spawns. Ranks an + # external launcher started were already running, so they read the + # config out of the run directory instead; see # provisioned_config_path. os.environ[CONFIG_PATH_ENV] = config_path exported = True @@ -577,9 +562,8 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona f"mooncake-store: {CONFIG_PATH_ENV}={config_path} " f"({json.dumps(config, sort_keys=True)})" ) - # Capacity is the pool's least obvious property and the one that - # explains a low hit rate, so state the arithmetic rather than - # leaving it to be done from global_segment_size later. + # Capacity is what explains a low hit rate, so state the + # arithmetic instead of leaving it to be derived later. logger.info( "mooncake-store: this server's ranks will each contribute " f"global_segment_size={pool.global_segment_size} to the pool; " @@ -601,9 +585,9 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona def maybe_provision_pool(kv_connector_config: Any) -> Iterator[None]: """Provision the pool if this deployment asked the server to. - A no-op for every other connector, and for a ``mooncake-store`` config - that left ``mooncake_store`` unset: that deployment is told about its pool - through ``MOONCAKE_CONFIG_PATH``, which is how the SLURM harness drives it. + A no-op for every other connector, and for a `mooncake-store` config that + left `mooncake_store` unset, since such a deployment is told about its pool + through `MOONCAKE_CONFIG_PATH` instead. """ if not uses_connector(kv_connector_config, "mooncake-store"): yield diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py index 6246877805f6..0e35227cf606 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py @@ -30,11 +30,11 @@ class PageTransfer: """One page of one layer group, to move in either direction.""" - #: Content identity from ``BlockHashChain``; names the key, not the location. + #: Content identity from `BlockHashChain`; names the key, not the location. block_hash: bytes layer_group_id: int - #: Page slot index within ``layer_group_id``, as reported by - #: ``RequestData.new_block_ids_by_layer_group``. + #: Page slot index within `layer_group_id`, as reported by + #: `RequestData.new_block_ids_by_layer_group`. page_index: int @@ -42,7 +42,7 @@ class PageTransfer: class RequestTransfers: """Pages belonging to one request, kept together for save bookkeeping. - The worker owes ``get_finished`` an answer per request, so a save's owner has + The worker owes `get_finished` an answer per request, so a save's owner has to survive the trip from scheduler to worker. """ diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py index bb24c01be475..9be55a5121df 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -16,7 +16,7 @@ Runs only on rank 0. It decides what to load and what to save; the workers do the moving. Two pieces of bookkeeping make that possible, and both exist because -``KVCacheManagerV2`` reports ``RequestData.block_hashes`` empty: +`KVCacheManagerV2` reports `RequestData.block_hashes` empty: * a hash chain per request, so a block has a content identity at all; * the page slot index per block ordinal, accumulated across iterations. The @@ -63,7 +63,7 @@ def __init__(self, chain: BlockHashChain): self.pages: Dict[int, List[int]] = {} #: First block ordinal not yet considered for saving. self.saved_upto = 0 - #: The offer made by ``get_num_new_matched_tokens``, in block ordinals. + #: The offer made by `get_num_new_matched_tokens`, in block ordinals. self.load_first_block = 0 self.load_blocks = 0 self.emitted_saves = False @@ -106,7 +106,7 @@ def get_num_new_matched_tokens( num_computed_tokens: Tokens already matched in the local KV cache. Returns: - Tokens the store can supply, and ``False`` for a synchronous load. + Tokens the store can supply, and `False` for a synchronous load. """ tokens = request.get_tokens(0) state = self._state_for(request, tokens) @@ -148,7 +148,7 @@ def cancel_load(self, request: LlmRequest, start: int, end: int): """Drop offered blocks whose tokens the runtime will not consume. Loads here are synchronous and nothing has been transferred yet, so this - is exact: the offer is truncated before ``build_connector_meta`` turns it + is exact: the offer is truncated before `build_connector_meta` turns it into work. """ state = self._requests.get(request.request_id) @@ -166,8 +166,8 @@ def cancel_load(self, request: LlmRequest, start: int, end: int): def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): """No-op: page indices are read from the scheduler output instead. - The flat ``block_ids`` here are a single space, but a V2 page index is - scoped to a layer group. ``RequestData.new_block_ids_by_layer_group`` is + The flat `block_ids` here are a single space, but a V2 page index is + scoped to a layer group. `RequestData.new_block_ids_by_layer_group` is the form that stays correct for every model, so that is the only source this connector uses. """ @@ -244,7 +244,7 @@ def _record_pages(self, state: _RequestState, request_data: RequestData) -> None by_group = request_data.new_block_ids_by_layer_group if not by_group: # Under a single layer group the manager also mirrors that group's - # indices into the flat ``new_block_ids``, but it does not say which + # indices into the flat `new_block_ids`, but it does not say which # group they belong to, so there is nothing safe to record from it. return for layer_group_id, indices in by_group.items(): @@ -284,8 +284,8 @@ def _append_pages(self, state: _RequestState, transfers: RequestTransfers, block for layer_group_id, indices in state.pages.items(): page_index = indices[block] if page_index == BAD_PAGE_INDEX: - # The block has no page in this group -- a sliding window has - # already dropped it. A partial page is not a usable cache entry, + # The block has no page in this group, because a sliding window + # dropped it. A partial page is not a usable cache entry, # so the whole block is skipped. return pages.append(PageTransfer(block_hash, layer_group_id, page_index)) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py index 484fb79320cd..1441367f67c5 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py @@ -16,25 +16,24 @@ The connector's default path registers the KV pools themselves with Mooncake, so the store reads and writes device memory directly. That needs the HCA to be able -to pin GPU pages -- GPUDirect RDMA, via ``nvidia_peermem`` or dma-buf. Where that -is unavailable, ``ibv_reg_mr`` fails on every pool range and the connector cannot -start at all. - -Staging trades a copy for that dependency. Mooncake is given a pinned host buffer -instead of the pools, and each page passes through a slot in it: gathered from its -device regions before a write, scattered back to them after a read. The store then -only ever registers host memory, which needs no GPUDirect. - -A slot holds the page's regions concatenated in region order, which is precisely -the payload the zero-copy path would have produced from the same regions. The -stored bytes are therefore identical either way, so a pool written by one path is -readable by the other -- including by another engine sharing the pool. - -Copies go through ``cudaMemcpyAsync`` rather than the batched Triton kernel in -``disaggregation/native/bounce/gather_scatter.py``. That kernel is the better tool -for device-to-device gather, but here one side is host memory: the copy engines -move it over the host link by DMA, whereas a kernel would do it with scattered -stores from the SMs. +to pin GPU pages, which means GPUDirect RDMA via `nvidia_peermem` or dma-buf. +Where that is unavailable, `ibv_reg_mr` fails on every pool range and the +connector cannot start. + +Staging trades a copy for that dependency. Mooncake is given a pinned host +buffer instead of the pools, and each page passes through a slot in it: gathered +from its device regions before a write, scattered back to them after a read. The +store then only ever registers host memory. + +A slot holds the page's regions concatenated in region order, which is exactly +the payload the zero-copy path produces from the same regions. The stored bytes +are therefore identical either way, so a pool written by one path is readable by +the other, including by another engine sharing the pool. + +Copies go through `cudaMemcpyAsync` rather than the batched Triton kernel in +`disaggregation/native/bounce/gather_scatter.py`. That kernel is the better tool +for device-to-device gather, but here one side is host memory, which the copy +engines move over the host link by DMA. """ from typing import List, Optional, Sequence, Tuple @@ -51,9 +50,8 @@ __all__ = ["HostStagingPool", "plan_slot_geometry", "sync_stream"] -#: Stated rather than inferred from the pointers: the direction is known at each -#: call site, and saying so keeps a copy from being misread if a host pointer is -#: ever outside the unified address space. +#: Stated explicitly rather than inferred from the pointers, which would be +#: wrong for a host pointer outside the unified address space. _DEVICE_TO_HOST = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost _HOST_TO_DEVICE = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice @@ -63,9 +61,8 @@ def _memcpy_async(dst: int, src: int, size: int, kind, stream: int) -> None: status = cudart.cudaMemcpyAsync(int(dst), int(src), int(size), kind, stream)[0] if status == cudart.cudaError_t.cudaSuccess: return - # Raised here rather than through CUASSERT so the operands are in the - # message. A bare cudaErrorInvalidValue from a copy says nothing about which - # of the three plausible causes it was. + # Raised here rather than through CUASSERT so the operands appear in the + # message; a bare cudaErrorInvalidValue names no cause. device = torch.cuda.current_device() if torch.cuda.is_available() else None raise RuntimeError( f"cudaMemcpyAsync failed with {status} staging a KV page: " @@ -73,7 +70,7 @@ def _memcpy_async(dst: int, src: int, size: int, kind, stream: int) -> None: f"stream={int(stream):#x} current_device={device}. An invalid value here " "is usually a stream created on a different device than the pages, which " "happens when a thread issues the copy without inheriting the rank's " - "device -- torch's current device is thread-local." + "device, since torch's current device is thread-local." ) @@ -115,9 +112,9 @@ def plan_slot_geometry( class HostStagingPool: """A registered pinned buffer, sliced into per-page slots. - One pool serves one direction. Loads run on the executor thread and saves on - the connector's background thread, so sharing slots between them would need a - lock on the transfer path for no benefit -- the two pools are independent. + One pool serves one direction. Loads run on the executor thread and saves + on the connector's background thread, so sharing slots between them would + need a lock on the transfer path for no benefit. """ def __init__( @@ -132,10 +129,8 @@ def __init__( self._num_slots = int(num_slots) self._label = label - # Pinned unconditionally, unlike the ``prefer_pinned`` heuristic used for - # transfer buffers elsewhere: this memory is handed to the store to - # register, so page-locking it is a correctness property of the - # registration rather than a copy-speed preference. + # Page-locking is a correctness requirement here rather than a + # copy-speed preference: this memory is handed to the store to register. pin = torch.cuda.is_available() self._buffer = torch.empty( self._slot_bytes * self._num_slots, dtype=torch.uint8, pin_memory=pin @@ -168,7 +163,7 @@ def slot_bytes(self) -> int: return self._slot_bytes def slot_address(self, index: int) -> int: - """Address of slot ``index``.""" + """Address of slot `index`.""" if not 0 <= index < self._num_slots: raise IndexError(f"slot {index} out of range [0, {self._num_slots})") return self._base + index * self._slot_bytes @@ -188,12 +183,12 @@ def gather( sizes: Sequence[int], stream: int, ) -> Tuple[int, int]: - """Copy one page's device regions into slot ``index``, concatenated. + """Copy one page's device regions into slot `index`, concatenated. Args: index: Slot to fill. addresses: Device addresses of the page's regions, in region order. - sizes: Byte counts matching ``addresses``. + sizes: Byte counts matching `addresses`. stream: CUDA stream handle the copies are issued on. Returns: @@ -216,7 +211,7 @@ def scatter( sizes: Sequence[int], stream: int, ) -> None: - """Copy slot ``index`` back out to one page's device regions. + """Copy slot `index` back out to one page's device regions. The inverse of :meth:`gather`, walking the regions in the same order so the split matches the concatenation the slot holds. @@ -229,7 +224,7 @@ def scatter( offset += size def reserve(self, total: int) -> None: - """Assert a page of ``total`` bytes is stageable, without copying.""" + """Assert a page of `total` bytes is stageable, without copying.""" self._check_fits(total) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py index 8eb9a633b6ca..2281685a0a5d 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -17,7 +17,7 @@ Every rejection here is a configuration whose failure mode is a wrong answer rather than a slow one: KV that gets replayed without all of the state it was computed with. Beam search, attention data parallelism, host and disk cache -tiers, and Mamba caches are rejected for all connectors in ``py_executor``, so +tiers, and Mamba caches are rejected for all connectors in `py_executor`, so they are not repeated. Checks run at construction, before any request is admitted, so a bad deployment diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 344c041e565f..9f7e02846f6f 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -14,10 +14,10 @@ # limitations under the License. """Worker side of the Mooncake store KV cache connector. -One worker per rank owns a ``MooncakeDistributedStore`` handle and moves pages +One worker per rank owns a `MooncakeDistributedStore` handle and moves pages between that pool and its own GPU KV cache. It is also the only place that knows -how a page is addressed and how a key is spelled, which is why the leader -- -colocated with rank 0's worker in the same process -- asks it to run prefix +how a page is addressed and how a key is spelled, which is why the leader, +colocated with rank 0's worker in the same process, asks it to run prefix lookups instead of rebuilding that knowledge. Loads are synchronous: the runtime has already told the scheduler those tokens @@ -28,7 +28,7 @@ once the forward pass that wrote them has retired, and blocking the executor loop on an RDMA write is exactly the cost the store is supposed to avoid. The scheduler reports such a request as saving asynchronously, which keeps its pages -pinned until ``get_finished`` says the writes landed. +pinned until `get_finished` says the writes landed. """ import threading @@ -63,7 +63,7 @@ #: 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 -#: connection or an out-of-band channel. See ``py_executor_creator``, which +#: connection or an out-of-band channel. See `py_executor_creator`, which #: constructs scheduler and worker concurrently for exactly this kind of #: mutual dependency. _LOCAL_WORKER: Optional["MooncakeStoreConnectorWorker"] = None @@ -141,7 +141,7 @@ def _batched(items: Sequence, size: int): def _stream_handle(stream) -> int: """The raw CUDA stream handle behind a torch stream, or a handle as given. - ``None`` maps to 0, the default stream, which is what the runtime passes when + `None` maps to 0, the default stream, which is what the runtime passes when it has no stream of its own to offer. """ if stream is None: @@ -175,20 +175,18 @@ def __init__(self, llm_args: TorchLlmArgs): ) self._save_thread: Optional[threading.Thread] = None self._save_lock = threading.Lock() - # Host staging, when the pool cannot register device memory. One pool per - # direction so the executor thread and the save thread never share slots. + # Host staging, when the pool cannot register device memory. self._load_staging: Optional[HostStagingPool] = None self._save_staging: Optional[HostStagingPool] = None self._save_stream: Optional[torch.cuda.Stream] = None - # This rank's device, captured on the executor thread. The save thread - # cannot ask for it itself; see _drain_saves. + # This rank's device, captured on the executor thread; see _drain_saves. self._device_index: Optional[int] = None # Pages per store call. Staging narrows this to the slots it can afford. self._batch_size = self._config.transfer_batch_size # Save submissions still in flight, per request. self._outstanding_saves: Dict[int, int] = defaultdict(int) # Requests the runtime has told us are done producing KV. Their pages - # stay pinned until we report them back through ``get_finished``. + # stay pinned until we report them back through `get_finished`. self._closed_requests: Set[int] = set() self._save_error: Optional[BaseException] = None @@ -227,8 +225,8 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: validate_layout(layout) addressing = PageAddressing(layout) - # Read here, on the executor thread, because it is thread-local and the - # save thread would otherwise see device 0 rather than this rank's. + # Torch's current device is thread-local, so read it here on the + # executor thread; the save thread would otherwise see device 0. if torch.cuda.is_available(): self._device_index = torch.cuda.current_device() if self._config.stage_through_host: @@ -274,10 +272,9 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: def _open_staging(self, addressing: PageAddressing) -> None: """Allocate and register the pinned slots pages will pass through. - Only the directions this role drives get a pool, since each one costs its - own pinned allocation and a consumer never gathers, nor a producer - scatter. The GPU pools are deliberately left unregistered: reaching them - is what staging exists to avoid needing. + Only the directions this role drives get a pool, since each one costs a + pinned allocation of its own. The GPU pools are left unregistered, + which is the point of the mode. """ max_bytes_per_page = max( addressing.bytes_per_page(layer_group_id) @@ -336,7 +333,7 @@ def is_registered(self) -> bool: return self._addressing is not None def count_prefix_hit(self, block_hashes: Sequence[bytes]) -> int: - """How many leading blocks of ``block_hashes`` are fully present. + """How many leading blocks of `block_hashes` are fully present. A block counts only when every layer group and every rank has its page, because a prefix is replayed as a whole. The scan stops at the first @@ -428,18 +425,17 @@ def start_load_kv(self, stream: torch.cuda.Stream): f"reported as computed. First failure: {failed[:1]}" ) if staging is not None: - # Only reached when every page in the batch landed, so no slot + # Only reached once every page in the batch landed, so no slot # holding a failed read is copied over a device page. unstage_batch_after_get(staging, batch_addresses, batch_sizes, handle) - # The slots are reused by the next batch and the forward pass - # reads these pages, so the scatter has to be complete before - # either happens. + # The next batch reuses the slots and the forward pass reads + # these pages, so the scatter has to complete before either. _sync_stream(handle) logger.debug(f"mooncake-store rank {self._rank} loaded {total_pages} pages") def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream): - """No-op: loads complete in ``start_load_kv``. + """No-op: loads complete in `start_load_kv`. Transfers are whole pages, so a page's bytes for every layer in a group land in one store call rather than layer by layer. There is nothing left @@ -447,7 +443,7 @@ def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream): """ def save_kv_layer(self, layer_idx: int, stream: torch.cuda.Stream): - """No-op: saves are submitted once per pass in ``wait_for_save``. + """No-op: saves are submitted once per pass in `wait_for_save`. A page is only complete when every layer of its group has written its slice, so there is no correct per-layer submission point. @@ -481,7 +477,7 @@ def get_finished( Args: finished_gen_req_ids: Requests that will produce no further KV. started_loading_req_ids: Requests loading asynchronously. Always - empty here, since ``get_num_new_matched_tokens`` only ever + empty here, since `get_num_new_matched_tokens` only ever offers synchronous loads; echoed back so the runtime does not wait on something that already happened. @@ -503,10 +499,8 @@ def get_finished( return finished_saving, list(started_loading_req_ids) def _drain_saves(self) -> None: - # Torch's current device is thread-local and a new thread starts at 0, so - # this has to be the device captured on the executor thread rather than - # whatever this one defaults to. Getting it wrong only shows up once the - # thread issues CUDA work of its own: the stream would belong to device 0 + # A new thread starts on device 0, so adopt the device captured on the + # executor thread. Otherwise a stream created below belongs to device 0 # while the KV pointers belong to the rank's device, and the copy fails # with cudaErrorInvalidValue on every rank except 0. if self._device_index is not None: @@ -514,8 +508,7 @@ def _drain_saves(self) -> None: if self._save_staging is not None and torch.cuda.is_available(): # Owned by this thread so the gather never queues behind the # executor's work, and created after set_device so it lands on the - # rank's device. Without a device there is nothing to order and the - # default stream handle stands in, which is what unit tests exercise. + # rank's device. self._save_stream = torch.cuda.Stream() while True: item = self._save_queue.get() @@ -573,13 +566,12 @@ def _put(self, transfers: Sequence[RequestTransfers]) -> None: source_addresses = [batch_addresses[i] for i in pending] source_sizes = [batch_sizes[i] for i in pending] if staging is not None: - # Gathering after the existence filter means a page that is - # already in the pool costs no copy. + # Gathered after the existence filter, so a page already in the + # pool costs no copy. source_addresses, source_sizes = stage_batch_for_put( staging, source_addresses, source_sizes, handle ) - # The store reads the slots on this thread, so they have to be - # filled first. + # The store reads the slots on this thread, so fill them first. _sync_stream(handle) results = self._store.batch_put_from_multi_buffers( [batch_keys[i] for i in pending], @@ -646,8 +638,8 @@ def shutdown(self) -> None: f"mooncake-store close failed: {type(exc).__name__}: {exc}\n" f"{traceback.format_exc()}" ) - # Released only after the store is closed: it holds registrations against - # this memory, and the save thread was already joined above. + # Released only after the store is closed, since it holds registrations + # against this memory. self._load_staging = None self._save_staging = None self._save_stream = None diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py index ed7a4098c43b..dcf5e24f3642 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py @@ -51,9 +51,9 @@ def uses_connector(kv_connector_config: Optional["KvCacheConnectorConfig"], name: str) -> bool: """Report whether a connector config resolves to the named preset. - Compares the resolved module rather than the ``connector`` field, so a - config that names the module explicitly instead of using the preset is - still recognized. Accepts ``None`` to save every caller a null check. + Compares the resolved module rather than the `connector` field, so a config + that names the module explicitly instead of using the preset is still + recognized. Accepts `None` to save every caller a null check. """ if kv_connector_config is None: return False diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index 2ae692ed5902..ff7200544ad8 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio +import faulthandler import os import signal import sys @@ -26,6 +27,13 @@ # 137 == 128 + SIGKILL(9): the exit code a shell reports for a SIGKILL'd process. _HARD_KILL_EXIT_CODE = 137 +#: Seconds an iteration may take before all thread stacks are dumped to stderr. +#: Unset or non-positive disables it. This diagnoses a *slow* loop, as opposed to +#: the hung loop `HangDetector` exists for, so it neither kills the process nor +#: stops the run. Set it somewhat above the normal iteration time and read the +#: dumps out of the worker log afterwards. +STALL_REPORT_ENV = "TRTLLM_STALL_REPORT_SEC" + def _best_effort_flush_streams() -> None: """Flush stdout/stderr without ever raising; diagnostics must not block hard kill.""" @@ -100,6 +108,41 @@ def __init__( self.lock = threading.Lock() self.active = False self._detected = False + self._stall_report_sec = self._read_stall_report_sec() + if self._stall_report_sec > 0: + logger.info( + f"Stall reporting enabled: dumping all thread stacks for any " + f"iteration exceeding {self._stall_report_sec}s " + f"({STALL_REPORT_ENV})." + ) + + @staticmethod + def _read_stall_report_sec() -> float: + raw = os.environ.get(STALL_REPORT_ENV, "") + if not raw.strip(): + return 0.0 + try: + return float(raw) + except ValueError: + logger.warning(f"Ignoring {STALL_REPORT_ENV}={raw!r}: not a number.") + return 0.0 + + def _arm_stall_report(self) -> None: + """Schedule a stack dump if this iteration runs long. + + `faulthandler`'s timer lives in a thread that does not take the GIL, so + unlike `print_all_stacks` it still fires when the loop is blocked inside + a native call. That is the case worth diagnosing, since a stall in pure + Python would already show up in a profile. + """ + if self._stall_report_sec <= 0: + return + faulthandler.dump_traceback_later(self._stall_report_sec, repeat=False, exit=False) + + def _cancel_stall_report(self) -> None: + if self._stall_report_sec <= 0: + return + faulthandler.cancel_dump_traceback_later() def start(self): """Enable hang detection.""" @@ -129,11 +172,16 @@ def detected(self): def checkpoint(self): """Reset hang detection timer.""" self.cancel_task() + self._arm_stall_report() if self.active: self.task = asyncio.run_coroutine_threadsafe(self._detect_hang(), self.loop) def cancel_task(self): """Cancel the hang detection task.""" + # Disarmed here rather than in checkpoint() so that pause(), used by the + # request-queue wait and the cross-rank broadcast probe, does not report + # a stall for time the loop is legitimately idle. + self._cancel_stall_report() if self.task is not None and not self.task.done(): self.task.cancel() self.task = None diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 2f219fe16cb0..fb7a90706cc7 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -986,8 +986,8 @@ def append_to_kv_heads_per_layer( # left alone and rejected loudly at bring-up. # # That leaves the scheduler without a tier to spill to, where - # suspension frees nothing. What reclaims pages instead is - # preemption -- see KVCacheManagerV2.preempt_request. + # suspension frees nothing, so pages are reclaimed by preemption + # instead. See KVCacheManagerV2.preempt_request. host_quota = 0 logger.info( "KV cache manager v2 host tier disabled: a KV connector is attached " @@ -1175,7 +1175,7 @@ def append_to_kv_heads_per_layer( self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) self._early_freed_index_requests: set[int] = set() # Requests whose pages a connector is still reading from, so the - # release half of `preempt_request` has to wait. See that method. + # release half of `preempt_request` has to wait. self._pending_preemption: Dict[int, LlmRequest] = {} self._prepare_page_table_tensor(index_mapper_capacity) @@ -2620,11 +2620,10 @@ def resume_request(self, req: LlmRequest) -> bool: # ---- preemption ---- # - # Suspension is the cheap way to free pages, but it only unpins them: the - # eviction controller then migrates them one cache level down. With GPU as - # the last level a suspended page stays `HELD`, which - # `CacheLevelManager.is_evictable` refuses to evict, so suspension frees - # nothing and the scheduler has no way out of a full pool. + # Suspension only unpins pages; the eviction controller then migrates them + # one cache level down. With GPU as the last level a suspended page stays + # `HELD`, which `CacheLevelManager.is_evictable` refuses to evict, so + # suspension frees nothing and the scheduler has no way out of a full pool. # # Preemption is the fallback for that case. It gives the pages up instead # of parking them, which costs a re-prefill but always works. @@ -2642,33 +2641,32 @@ def preempt_request(self, req: LlmRequest) -> bool: """Give up *req*'s KV cache so its pages can be reclaimed. Unlike :meth:`suspend_request` this does not keep the pages. Closing - the request's ``_KVCache`` returns its committed blocks to the radix - tree as reusable prefix and leaves their pages ``DROPPABLE``, which is - evictable at every level -- including the last, where ``HELD`` is not. - So the data is not thrown away: it stays resident and locally matchable - until something else actually needs the space. + the request's `_KVCache` returns its committed blocks to the radix tree + as reusable prefix and leaves their pages `DROPPABLE`, which is + evictable at every level, unlike `HELD`. The data is not thrown away: + it stays resident and locally matchable until something else needs the + space. The request is reset to context state by the caller and re-prefills - whatever it can no longer match. With a connector attached the blocks - it already wrote to the store come back through the ordinary prefix - load, so the reload is usually cheap, and recompute is the - always-correct fallback when the store no longer has them. + whatever it can no longer match. With a connector attached, blocks it + already wrote to the store come back through the ordinary prefix load, + and recompute is the fallback when the store no longer has them. Returns True when the pages were released. When the connector still has - saves in flight the release is deferred and this returns False: those - saves read directly out of these pages, so freeing them now would let a - later request overwrite the bytes mid-transfer and publish them to the - store under a valid hash. Callers must not count on the pages until + saves in flight the release is deferred and this returns False, because + those saves read directly out of these pages: freeing them now would + let a later request overwrite the bytes mid-transfer and publish them + under a valid hash. Callers must not count on the pages until :meth:`try_complete_preemption` has run for this request. """ if self.kv_connector_manager is None: self._release_preempted(req) return True - # Same handshake the finish path uses, for the same reason. The - # request lands in DISAGG_CONTEXT_TRANS_IN_PROGRESS, out of the - # schedulable range, and its `_KVCache` keeps holding the pages until - # every rank reports the save retired through `get_finished`. + # The same handshake the finish path uses: the request lands in + # DISAGG_CONTEXT_TRANS_IN_PROGRESS, out of the schedulable range, and + # its `_KVCache` keeps holding the pages until every rank reports the + # save retired through `get_finished`. if self.kv_connector_manager.request_finished( req, self.get_connector_page_indices(req) ): @@ -2682,8 +2680,8 @@ def try_complete_preemption(self, req: LlmRequest) -> bool: """Release pages for a request whose deferred preemption just cleared. Returns False when *req* was not awaiting preemption, which is how the - caller tells a preempted request apart from an ordinary finished one - in the connector's ``get_finished`` output. + caller tells a preempted request apart from an ordinary finished one in + the connector's `get_finished` output. """ if self._pending_preemption.pop(req.py_request_id, None) is None: return False @@ -2693,8 +2691,8 @@ def try_complete_preemption(self, req: LlmRequest) -> bool: def _release_preempted(self, req: LlmRequest) -> None: self.free_resources(req) # Ask the connector again on re-admission rather than reusing the - # memoised offer: the store has strictly more of this prefix now than - # it did when the request was first admitted. + # memoised offer, since the store now has more of this prefix than it + # did when the request was first admitted. req.py_connector_prefix_start = None req.py_connector_prefix_end = None req.py_connector_load_async = False @@ -3519,7 +3517,7 @@ def release_index_slot(self, request_id: int) -> None: def free_resources(self, request: LlmRequest, pin_on_release: bool = False): # A request awaiting preemption can still be cancelled or fail while # its saves drain. Dropping the entry here keeps a dead request from - # blocking every later preemption via has_pending_preemption(). + # blocking every later preemption through has_pending_preemption. self._pending_preemption.pop(request.py_request_id, None) self._release_undelivered_connector_prefix(request) if self.conversation_manager is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 40581baaa0df..0b65452ab01a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3758,10 +3758,10 @@ def _resume_preempted_request(self, request: LlmRequest) -> bool: """Complete a preemption whose connector saves have now retired. The scheduler preempts a request by handing it to the connector the - same way a finished request is handed over, so that its pages stay put - until every rank reports the in-flight saves done. Both kinds come - back through ``get_finished``; only the KV cache manager knows which - is which. + same way a finished request is handed over, so its pages stay put until + every rank reports the in-flight saves done. Both kinds come back + through `get_finished`, and only the KV cache manager knows which is + which. Returns True when *request* was preempted rather than finished, in which case its pages are now released and it is back in context state diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 2f9cd43c7d65..a3b15f56c734 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -383,8 +383,8 @@ def create_py_executor( kv_cache_config.enable_block_reuse = False kv_cache_config.enable_partial_reuse = False - # Must happen before the KV cache manager is built below, since the manager - # reads enable_partial_reuse to construct its block pools. + # Must happen before the KV cache manager is built, since the manager reads + # enable_partial_reuse to construct its block pools. if (kv_cache_config.enable_partial_reuse and uses_connector(kv_connector_config, "mooncake-store")): logger.warning( diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 305a1627c248..80861f8060fd 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -158,7 +158,7 @@ def __init__( ) -> None: self.max_num_tokens = max_num_tokens # Only read when preempting: LlmRequest.pause clamps the rewritten - # prompt (original prompt plus tokens generated so far) to it. + # prompt, the original plus tokens generated so far, to this. self.max_input_len = max_input_len self._stalled_schedules = 0 self.max_num_requests = ( @@ -403,11 +403,9 @@ def _schedule_loop(self, active_requests, inflight_request_ids): req_it += 1 - # Requests whose pages were given up during this pass. They must not - # be scheduled afterwards: a victim that is itself a started context - # request is already sitting in pending_ctx, and re-admitting it here - # would spend the pages the preemption just released on the very - # request that released them. + # Requests whose pages were given up during this pass. A victim that is + # itself a started context request still sits in pending_ctx, so + # re-admitting it would spend the pages its own preemption released. preempted_ids: set[int] = set() def preempt_for_pages(req: LlmRequest) -> bool: @@ -585,10 +583,10 @@ def _try_schedule_context_full( # V2 resizes KV cache directly in the scheduler (no separate # prepareResources for main cache), so include draft tokens. if not self.kv_cache_manager.resize_context(req, context_tokens + draft_len): - # Out of pages. Give one started request up so this one can - # proceed, and retry next iteration rather than now: a failed - # resize leaves a first chunk suspended, so the retry has to go - # back through prepare_context to resume it. + # Out of pages. Give up one started request so this one can + # proceed, and retry next iteration: a failed resize leaves a + # first chunk suspended, so the retry has to go back through + # prepare_context to resume it. preempt_for_pages(req) return ScheduleAction.SKIP, 0, False @@ -664,10 +662,9 @@ def _try_schedule_context_chunked( chunk_size = (chunk_size // self.chunk_unit_size) * self.chunk_unit_size if chunk_size <= 0: - # Out of token budget, not out of pages, so releasing this - # request's pages would not help: next iteration gets a fresh - # budget. Deliberately not suspended, to avoid pathological - # suspend/resume cycles. + # Out of token budget rather than out of pages, so releasing pages + # would not help; the next iteration gets a fresh budget. Not + # suspended either, to avoid pathological suspend/resume cycles. return ScheduleAction.SKIP, 0, False chunk_size = self._align_chunk_to_mm_block( @@ -695,7 +692,7 @@ def _try_schedule_context_chunked( # V2 resizes KV cache directly in the scheduler, so include # draft tokens for last chunk. if not self.kv_cache_manager.resize_context(req, resize_tokens): - # Out of pages — see the same call in _try_schedule_context_full. + # Out of pages, as in _try_schedule_context_full. preempt_for_pages(req) return ScheduleAction.SKIP, 0, False @@ -1051,22 +1048,21 @@ def _try_preempt_for_pages( ) -> bool: """Release one started request's KV cache so another can allocate. - This is the fallback for a pool that suspension cannot drain -- see - ``KVCacheManagerV2.preempt_request``. With a cache tier below GPU, - suspension is cheaper and keeps the pages, so leave that path alone. + The fallback for a pool that suspension cannot drain; see + `KVCacheManagerV2.preempt_request`. With a cache tier below GPU, + suspension is cheaper and keeps the pages, so that path is left alone. Returns True when pages became available in this iteration. A connector defers the release until its in-flight saves retire, in - which case this returns False and the caller should give up for now: - the pages arrive a few iterations later. + which case this returns False and the pages arrive a few iterations + later. """ if self.kv_cache_manager.has_cache_tier_below_gpu: return False if self.kv_cache_manager.has_pending_preemption(): - # One victim at a time. A full pool would otherwise preempt the - # whole batch while the first release is still draining, and - # every one of those requests would have to re-prefill. + # One victim at a time, or a full pool would preempt the whole + # batch while the first release is still draining. return False # Newest first, so the requests closest to completing keep their @@ -1092,11 +1088,10 @@ def _try_preempt_for_pages( if self.draft_kv_cache_manager is not None: self.draft_kv_cache_manager.free_resources(victim) if released: - # Rewrites the prompt to include whatever was generated and - # resets state to CONTEXT_INIT with the chunk position at 0, - # so the request re-enters as an ordinary prefill next - # iteration. Deferred releases are paused by the executor - # once the connector reports the saves retired. + # Rewrites the prompt to include what was generated and resets + # state to CONTEXT_INIT, so the request re-enters as an + # ordinary prefill. Deferred releases are paused by the + # executor once the connector reports the saves retired. victim.pause(self.max_input_len) evicted.append(victim) preempted_ids.add(victim.py_request_id) @@ -1104,10 +1099,10 @@ def _try_preempt_for_pages( return False - # Consecutive scheduling passes that reclaimed nothing before the - # scheduler calls it a deadlock. A stalled pass costs ~2ms, so this trips - # in seconds, while the transient one-iteration deferrals (multimodal - # chunk alignment, PEFT budget, IndexMapper slots) clear long before it. + # Consecutive scheduling passes that reclaimed nothing before this counts + # as a deadlock. A stalled pass costs ~2ms, so it trips within seconds, + # while transient one-iteration deferrals (multimodal chunk alignment, + # PEFT budget, IndexMapper slots) clear long before. _DEADLOCK_STALL_ITERS = 1000 def _detect_deadlock( @@ -1121,11 +1116,10 @@ def _detect_deadlock( """Fail loudly when no request can be scheduled or reclaimed. Without this the executor spins at full speed while scheduling - nothing: the loop looks healthy to the hang detector and to - ``/health``, and the job burns its wall clock. Context candidates - count alongside generation ones because a disaggregated prefill - server has no generation requests at all, and counting only those - left it spinning silently. + nothing, which looks healthy to the hang detector and to `/health` + while the job burns its wall clock. Context candidates count alongside + generation ones because a disaggregated prefill server has no + generation requests at all. """ if made_progress: self._stalled_schedules = 0 diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index 8e1018fa4e0f..f9c5423fa37b 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -14,12 +14,11 @@ # limitations under the License. """The two pieces of a Mooncake pool that outlive any one engine. -A server that owns its pool needs neither of these: it describes the pool in +A server that owns its pool needs neither: it describes the pool in `kv_connector_config.mooncake_store` and `trtllm-serve` provisions it during -bringup. They exist for the pools it cannot own -- one shared by several -engines, one that has to survive a restart, one whose capacity has to come from -nodes that run no connector -- so that those deployments are still assembled -from things TensorRT-LLM ships rather than from a launch script of one's own. +bringup. These commands exist for the pools it cannot own, such as one shared by +several engines, one that has to survive a restart, or one whose capacity comes +from nodes that run no connector. """ import json @@ -38,9 +37,9 @@ def _until_signalled() -> threading.Event: """An event that SIGINT and SIGTERM set. - Both commands hold a resource -- a child process, a mounted segment -- - whose release is in a ``finally``. Default SIGTERM handling would skip it, - leaving the master unreaped or the pool advertising memory that has gone. + Both commands hold a resource, a child process or a mounted segment, whose + release is in a `finally`. Default SIGTERM handling would skip it, leaving + the master unreaped or the pool advertising memory that has gone. """ stopping = threading.Event() @@ -92,13 +91,11 @@ def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, heartbeat_seconds: int): """Run a mooncake_master for as long as this command runs. - For a pool that must not belong to any one engine: several servers sharing - it, or one that has to still be there after a server restarts. A single - server with a pool of its own should set `mooncake_store.launch_master` - instead and skip this entirely. + A single server with a pool of its own should set + `mooncake_store.launch_master` instead. """ - # Imported here rather than at module scope so that reaching any other - # subcommand, or --help, does not pay for the connector package. + # Imported lazily so other subcommands and --help do not pay for the + # connector package. from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import \ running_master from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig @@ -123,15 +120,15 @@ def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, announced = started while not stopping.is_set(): if (code := master.process.poll()) is not None: - # Its own death is the interesting outcome: the pool is gone - # and every client is about to start failing. + # The pool is gone once the master dies, and every client is + # about to start failing. raise click.ClickException( f"mooncake_master exited with code {code}. See " f"{master.log_path}") stopping.wait(1.0) now = time.monotonic() - # Says the pool is still there, which is the question asked of - # this log when clients start failing: master or fabric? + # Distinguishes a dead master from a dead fabric once clients + # start failing. if heartbeat_seconds > 0 and now - announced >= heartbeat_seconds: announced = now logger.info( @@ -189,11 +186,8 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, ready_file: Optional[str], heartbeat_seconds: int): """Lend this node's host memory to a Mooncake pool, for as long as it runs. - Pool capacity comes only from processes that open a store handle, and in a - disaggregated deployment only the context servers do -- so the pool is - prefill-node memory, caching prefill's own GPUs. Running this on the - generation nodes puts their memory in the same pool while leaving those - engines connector-free. + Running this on the generation nodes puts their memory into the pool while + leaving those engines connector-free. """ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, master_timeout, @@ -215,8 +209,7 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, donating = parse_size(segment_size) resolved = resolve_master_address(master, master_timeout()) - # Before setup, so "the master is not up yet" is reported as that and not - # as the status code setup returns for every kind of failure. + # Before setup, so an absent master is reported as such. wait_for_master(resolved) stopping = _until_signalled() @@ -237,9 +230,8 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, f"{ready_file}, so a launcher waiting on the pool's " "capacity can proceed") - # Idle by design. A put or get here would make this node a client in - # the traffic sense, which is the thing keeping the generation engine - # connector-free is meant to avoid. + # Idle by design: a put or get here would make this node a traffic + # client, which is what donation exists to avoid. started = time.monotonic() while not stopping.is_set(): if heartbeat_seconds <= 0: diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 38bf2e742acb..d1f946ddcf9f 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -523,21 +523,19 @@ def _provision_kv_cache_pool(llm_args: dict, owns_engine: bool = True) -> Iterator[None]: """Bring up the shared cache this server needs or feeds, for its lifetime. - Two things the engine cannot produce for itself. A connector backed by a - cluster-wide pool needs it reachable before any rank opens a handle, and - the ranks are spawned by the LLM constructor; entering this around that - construction is what makes the pool part of `trtllm-serve` bringup rather - than something a launch script has to arrange. A deployment that arranges - it anyway is detected and left alone. - - Separately, a server may lend the pool host memory without using it, which - is how a pool spans nodes whose engines have no connector. That segment - has to be mounted before traffic arrives, and held for as long as the - pages placed in it are expected to be there -- so, this context. - - Only the process that owns the engine does either: an attached frontend + A connector backed by a cluster-wide pool needs that pool reachable before + any rank opens a handle, and the ranks are spawned by the LLM constructor, + so this wraps the construction. A deployment that provisions the pool + externally is detected and left alone. + + A server may also lend the pool host memory without using it, which is how + a pool spans nodes whose engines have no connector. That segment has to be + mounted before traffic arrives and held for as long as the pages placed in + it are expected to be there. + + Only the process that owns the engine does either. An attached frontend re-execs this command line but shares the launcher's executor, so it would - otherwise stand up a second, private pool and lend a second segment. + otherwise stand up a second pool and lend a second segment. """ if not owns_engine: yield @@ -548,9 +546,9 @@ def _provision_kv_cache_pool(llm_args: dict, connector_config = llm_args.get("kv_connector_config") if isinstance(connector_config, dict): - # A YAML config section arrives here unvalidated. Coerce it now, since - # the pool has to be described before the LLM constructor would do it, - # and hand the validated model on so it is not parsed twice. + # A YAML config section arrives unvalidated, and the pool has to be + # described before the LLM constructor would coerce it. Hand the + # validated model on so it is not parsed twice. connector_config = KvCacheConnectorConfig(**connector_config) llm_args["kv_connector_config"] = connector_config @@ -2547,9 +2545,8 @@ def resolve_command(self, ctx, args): "disaggregated_mpi_worker": disaggregated_mpi_worker, "mm_embedding_serve": serve_encoder, "embeddings": serve_embedding, - # The parts of a Mooncake pool that cannot belong to a server, for the - # deployments where a pool outlives or spans them. A server that owns - # its pool provisions it from its own config and needs neither. + # The parts of a Mooncake pool that cannot belong to a server, for + # deployments where a pool outlives or spans them. "mooncake_master": mooncake_master, "mooncake_donor": mooncake_donor, }) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 7313bd98bc39..6d22c1262b26 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1948,25 +1948,20 @@ def num_capture_layers(self) -> int: class MooncakeStoreConfig(StrictBaseModel): """The Mooncake store pool the `mooncake-store` connector should join. - Describes the *pool*: which master owns it, how workers reach it, and how - much memory each contributes. A worker's own relationship to the pool - (read/write role, key namespace, model identity) stays in the - ``TRTLLM_MOONCAKE_STORE_*`` environment variables, because it is per - process while this is per deployment. + Describes the pool: which master owns it, how workers reach it, and how + much memory each contributes. A worker's own relationship to the pool, such + as its read/write role and key namespace, stays in the + `TRTLLM_MOONCAKE_STORE_*` environment variables, since that is per process + while this is per deployment. Setting this makes `trtllm-serve` render the Mooncake client config and - export ``MOONCAKE_CONFIG_PATH`` itself, so no external script has to. - An inherited ``MOONCAKE_CONFIG_PATH`` still wins, which is how the SLURM - harness keeps pointing workers at a pool it manages. - - Every field opts out of telemetry: they describe one site's pool -- its - ports, its fabric, how much memory it was given -- rather than which - features are in use, which ``kv_connector_config.connector`` already says. + export `MOONCAKE_CONFIG_PATH` itself. An inherited `MOONCAKE_CONFIG_PATH` + still wins, so an externally managed pool stays reachable. """ master_server_address: Optional[str] = Field( None, description="Address of an already-running mooncake_master, as " - "host:port or file:// naming a file that holds one -- which is " + "host:port or file:// naming a file that holds one. The file is " "how to reach a master whose host a scheduler chose, since " "'trtllm-serve mooncake_master --address_file' publishes it there " "once it answers. Mutually exclusive with launch_master.") @@ -2064,28 +2059,21 @@ def _require_exactly_one_master(self) -> "MooncakeStoreConfig": class MooncakeDonationConfig(StrictBaseModel): """Host memory this server lends to a Mooncake pool it does not use. - Pool capacity comes only from processes that open a store handle, and in a - disaggregated deployment only the context servers configure the connector. - The pool is then entirely prefill-node memory: prefill's DRAM caching - prefill's GPUs, which largely duplicates what - ``kv_cache_config.host_cache_size`` already does. Setting this on the - generation servers puts their memory into the same pool, so prefill writes - blocks that land on decode-side DRAM, while the generation engine stays - free of any connector and keeps its cache transceiver for the handoff. - - Lending memory is deliberately separate from - ``kv_connector_config``. That config attaches a connector, and a connector - reads and writes; there is no setting on it that means "contribute memory - only", so expressing capacity there would start this server using the - store. Capacity and traffic are different things and are configured - separately. - - The memory is charged to this process and competes with everything else on - the node, ``kv_cache_config.host_cache_size`` above all, so size the two - together. - - Every field opts out of telemetry: they describe one site's pool and how - much of this node was given to it, not which features are in use. + Pool capacity comes only from processes that open a store handle, which in + a disaggregated deployment is the context servers alone. Setting this on + the generation servers puts their memory into the same pool, so prefill + writes blocks that land on decode-side DRAM while the generation engine + stays free of any connector and keeps its cache transceiver for the + handoff. + + Capacity is kept separate from `kv_connector_config` because attaching a + connector would also start this server reading and writing the store. + + The memory is charged to this process, so size it together with + `kv_cache_config.host_cache_size`. + + Fields opt out of telemetry because they describe one site's pool rather + than which features are in use. """ master_server_address: str = Field( ..., diff --git a/tests/unittest/_torch/executor/test_hang_detector_kill.py b/tests/unittest/_torch/executor/test_hang_detector_kill.py index b4f122bb8ad3..3dd00982c941 100644 --- a/tests/unittest/_torch/executor/test_hang_detector_kill.py +++ b/tests/unittest/_torch/executor/test_hang_detector_kill.py @@ -64,6 +64,82 @@ def test_pause_suppresses_detection(): assert hd.detected() is False +# Driver for the stall-reporting tests. faulthandler writes straight to fd 2, +# so the dump can only be observed from another process; markers go to stderr +# too, which is what makes their interleaving with the dump assertable. +_STALL_REPORT_SCRIPT = """ +import os, sys, time +from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector + +def mark(name): + print(f"MARK {name}", file=sys.stderr, flush=True) + +os.environ["TRTLLM_STALL_REPORT_SEC"] = "1" +hd = HangDetector(timeout=300, on_detected=lambda: None) +with hd: + mark("fast-begin") + hd.checkpoint() + time.sleep(0.2) + mark("fast-end") + + hd.checkpoint() + mark("paused-begin") + with hd.pause(): + time.sleep(2.0) + mark("paused-end") + + hd.checkpoint() + mark("slow-begin") + time.sleep(2.0) + mark("slow-end") + +# Same slow iteration, but with the knob unset: must stay silent. +os.environ.pop("TRTLLM_STALL_REPORT_SEC", None) +off = HangDetector(timeout=300, on_detected=lambda: None) +with off: + off.checkpoint() + mark("disabled-begin") + time.sleep(2.0) + mark("disabled-end") +""" + + +def test_stall_report_fires_only_for_slow_iterations(): + """The stack dump lands in the slow iteration, and nowhere else. + + Covers all four cases in one subprocess because each one would otherwise + pay a cold `import tensorrt_llm`. + """ + proc = subprocess.run( + [sys.executable, "-c", _STALL_REPORT_SCRIPT], + env={**os.environ, "TLLM_DISABLE_MPI": "1"}, + timeout=300, + capture_output=True, + ) + stderr = proc.stderr.decode(errors="replace") + assert proc.returncode == 0, f"driver failed: {stderr[-2000:]}" + + def index_of(marker): + position = stderr.find(f"MARK {marker}") + assert position != -1, f"missing marker {marker!r} in:\n{stderr[-2000:]}" + return position + + # faulthandler's dump is headed by "Timeout (H:MM:SS)!". + dumps = [ + position + for position in range(len(stderr)) + if stderr.startswith("Timeout (", position) + ] + assert len(dumps) == 1, ( + f"expected exactly one stack dump, got {len(dumps)}:\n{stderr[-2000:]}" + ) + # Only the slow iteration should have produced it: not the fast one, not + # the paused window (which is longer than the threshold), and not the run + # with the knob unset. + assert index_of("slow-begin") < dumps[0] < index_of("slow-end") + assert "MARK disabled-end" in stderr + + def test_propagate_hard_kill_self_sigkills_without_mpi(): """With MPI disabled, propagate_hard_kill self-SIGKILLs the process. diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index a77f517c0e2b..e3adca7ddeb3 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -173,8 +173,7 @@ def make_kv_cache_manager( mgr.try_allocate_generation.side_effect = try_allocate_generation_fn or (lambda req: True) mgr.suspend_request.return_value = None mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active - # Preemption is the fallback for a pool with nothing under GPU to spill - # to, so the default here (a host tier exists) leaves it switched off. + # The default here has a cache tier below GPU, which leaves preemption off. mgr.has_cache_tier_below_gpu = has_cache_tier_below_gpu mgr.has_pending_preemption.return_value = has_pending_preemption mgr.preempt_request.side_effect = preempt_request_fn or (lambda req: True) @@ -851,10 +850,9 @@ def _out_of_pages_for(request_id): class TestContextPreemption: """Releasing a started request's pages when suspension cannot help. - Suspension only unpins pages so the eviction controller can migrate them - one level down; with GPU as the last level a suspended page stays HELD and - unevictable, so it frees nothing. These tests cover the fallback that - gives the pages up instead. + With GPU as the last cache level a suspended page stays HELD and + unevictable, so suspension frees nothing. These tests cover the fallback + that gives the pages up instead. """ def test_out_of_pages_preempts_started_request(self): @@ -934,8 +932,8 @@ def test_skipped_when_a_cache_tier_exists_below_gpu(self): out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) mgr.preempt_request.assert_not_called() - # Suspension is cheaper and keeps the pages, so that path is left - # exactly as it was: the request is simply skipped. + # Suspension is cheaper and keeps the pages, so the request is + # simply skipped. assert ids(out.context_requests) == [99] def test_one_victim_at_a_time_while_a_release_is_draining(self): @@ -1033,8 +1031,7 @@ class TestDeadlockDetection: """The scheduler must fail loudly rather than spin scheduling nothing. A stalled pass costs a couple of milliseconds, so an undetected stall - burns a job's whole wall clock while the loop still looks healthy to the - hang detector and to /health. + burns a job's whole wall clock. """ def test_raises_after_repeated_stalls_with_context_candidates(self): @@ -1057,7 +1054,7 @@ def test_raises_after_repeated_stalls_with_generation_candidates(self): sched = make_scheduler(mgr, max_num_tokens=100) sched._DEADLOCK_STALL_ITERS = 3 # Self-eviction suspends it on the first pass, which counts as - # progress; afterwards it is inactive and nothing can be reclaimed. + # progress. Afterwards it is inactive and nothing can be reclaimed. reqs = [make_gen_request(0)] for _ in range(3): diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 0a768e9738b8..52426dc49b5c 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -87,8 +87,8 @@ def __init__(self): self.exist_calls = [] self.closed = False self.fail_gets_for = set() - #: Workers built against this store. ``make_worker`` shuts each one - #: down; the fixture repeats it as a backstop for early failures. + #: Workers built against this store. `make_worker` shuts each one down; + #: the fixture repeats it as a backstop for early failures. self.workers = [] def register_buffer(self, address, size): @@ -393,7 +393,7 @@ def test_config_requires_the_env_var(monkeypatch): def test_config_falls_back_to_the_run_directory(tmp_path, monkeypatch): - # A rank the launcher started was already running when its leader + # A rank an external launcher started was already running when its leader # provisioned the pool, so it never inherited the exported path and reads # the rendered config out of the shared run directory instead. monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) @@ -418,8 +418,8 @@ def test_config_run_directory_without_a_rendered_config_still_asks(tmp_path, mon def test_config_env_var_wins_over_the_run_directory(tmp_path, monkeypatch): - # An externally managed pool stays reachable: the run directory is only - # consulted when nothing was passed in. + # An externally managed pool stays reachable, since the run directory is + # only consulted when nothing was passed in. named = tmp_path / "external.json" named.write_text(json.dumps({"master_server_address": "external:50051"})) (tmp_path / "mooncake.json").write_text( @@ -474,9 +474,9 @@ def test_validate_layout_rejects_sliding_window(): # ---- connector identification ---- # -# py_executor_creator turns partial reuse off for this connector, and finds it -# through uses_connector. Missing the config would silently cost the reuse the -# store exists to provide, so the recognition itself is worth pinning down. +# py_executor_creator turns partial reuse off for this connector and finds it +# through uses_connector. Failing to recognize the config would silently cost +# the reuse the store exists to provide. def test_uses_connector_recognizes_the_preset(): @@ -541,8 +541,8 @@ def test_worker_prefix_hit_needs_every_layer_group(store_config, fake_store): def test_worker_prefix_hit_stops_at_the_first_gap(store_config, fake_store): with make_worker(fake_store, layout=make_layout()) as worker: hashes = [bytes([index]) * 16 for index in range(3)] - # Block 1 missing: block 2 is unusable even though it is present, because a - # prefix is replayed contiguously. + # With block 1 missing, block 2 is unusable even though it is present, + # because a prefix is replayed contiguously. fake_store.objects.add(worker._namespaces[0].key(hashes[0])) fake_store.objects.add(worker._namespaces[0].key(hashes[2])) assert worker.count_prefix_hit(hashes) == 1 @@ -640,8 +640,7 @@ def staged_copies(monkeypatch): "_memcpy_async", lambda dst, src, size, kind, stream: copies.append((int(dst), int(src), int(size))), ) - # Imported into the worker by name, so the worker's binding is the one that - # has to be replaced. + # Imported into the worker by name, so replace the worker's binding. monkeypatch.setattr(worker_module, "_sync_stream", lambda _stream: None) return copies @@ -650,8 +649,8 @@ def staged_copies(monkeypatch): def fake_cuda(monkeypatch): """Present a CUDA device on a host that has none, recording set_device calls. - Only safe for paths that do not allocate or launch; it exists to exercise the - device bookkeeping around the save thread. + Only safe for paths that do not allocate or launch, and exists to exercise + the device bookkeeping around the save thread. """ recorded = [] monkeypatch.setattr(torch.cuda, "is_available", lambda: True) @@ -672,7 +671,7 @@ def fake_cuda(monkeypatch): ) def test_plan_slot_geometry(page_bytes, batch, budget, expected_slots): slot_bytes, num_slots = plan_slot_geometry(page_bytes, batch, budget) - # A slot always holds a whole page; the budget bounds the count, not the width. + # A slot always holds a whole page: the budget bounds the count, not the width. assert slot_bytes == page_bytes assert num_slots == expected_slots @@ -694,10 +693,6 @@ def test_config_reads_staging_from_the_json(store_config): assert config.staging_buffer_bytes == 256 * 1024**2 -def test_config_defaults_to_zero_copy(store_config): - assert MooncakeStoreConnectorConfig.from_env().stage_through_host is False - - @pytest.mark.parametrize( "value,expected", [("1", True), ("true", True), ("on", True), ("0", False), ("off", False)] ) @@ -717,13 +712,13 @@ def test_staging_registers_host_buffers_and_never_the_pools( ): layout = make_layout(regions_per_group=2) with make_staged_worker(fake_store, store_config, layout=layout) as worker: - # Registering the pools is the step that needs GPUDirect, so staging must - # not do it at all -- that is the whole point of the mode. + # Registering the pools is the step that needs GPUDirect, so staging + # must not do it at all. pool_ranges = PageAddressing(layout).registration_ranges() registered_starts = {address for address, _size in fake_store.registered} assert registered_starts.isdisjoint({start for start, _end in pool_ranges}) - # One pinned buffer per direction, since the default role is ``both``. + # One pinned buffer per direction, since the default role is `both`. assert len(fake_store.registered) == 2 assert registered_starts == { worker._load_staging.slot_address(0), @@ -744,10 +739,9 @@ def test_staging_put_hands_the_store_one_host_buffer_per_page( (keys, addresses, sizes) = fake_store.put_calls[0] assert keys == [worker._namespaces[0].key(block_hash)] - # The store sees one contiguous host buffer, and its length is the sum of + # The store sees one contiguous host buffer whose length is the sum of # the device regions. That equality is what keeps a staged write - # byte-identical to a zero-copy one, so either path can read the other's - # pages. + # byte-identical to a zero-copy one. assert addresses == [[slot]] assert sizes == [[sum(device_sizes)]] @@ -802,35 +796,32 @@ def test_staging_does_not_scatter_a_failed_load(store_config, fake_store, staged with pytest.raises(RuntimeError, match="failed to load"): worker.start_load_kv(None) - # A failed read leaves the slot holding whatever it held before. Copying - # that onto the page would put unrelated bytes where the runtime already - # promised computed KV. + # A failed read leaves the slot holding whatever it held before, and + # copying that onto the page would put unrelated bytes where the + # runtime already promised computed KV. assert staged_copies == [] def test_worker_captures_the_ranks_device_at_registration(store_config, fake_store, fake_cuda): with make_worker(fake_store, layout=make_layout()) as worker: - # Read on the executor thread, where it is correct. torch's current - # device is thread-local, so the save thread cannot ask for it itself. assert worker._device_index == 3 def test_save_thread_adopts_the_ranks_device_not_the_thread_default( store_config, fake_store, fake_cuda ): - """Regression: the save thread must not run on torch's default device. + """The save thread must not run on torch's default device. It issues staging copies against pointers owned by the rank's device. A stream created on device 0 instead fails every copy with - cudaErrorInvalidValue, and only on ranks other than 0 -- which is exactly how - this escaped into a run. + cudaErrorInvalidValue, and only on ranks other than 0. """ with make_worker(fake_store, layout=make_layout()): deadline = time.monotonic() + 5.0 while 3 not in fake_cuda and time.monotonic() < deadline: time.sleep(0.01) assert 3 in fake_cuda, f"save thread set devices {fake_cuda}, expected the rank's 3" - # Never the thread-local default, which is what the bug did. + # Never the thread-local default. assert 0 not in fake_cuda @@ -851,8 +842,7 @@ def test_staging_narrows_the_batch_to_the_budget(store_config, fake_store, stage ) ] ) - # Five pages through two slots: three calls, and no call wider than the - # slot count, which is the constraint staging adds. + # Five pages through two slots, and no call wider than the slot count. assert [len(keys) for keys, _a, _s in fake_store.put_calls] == [2, 2, 1] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py index d8c2b040f1d2..841f42e5f5db 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_donor.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -14,10 +14,9 @@ # limitations under the License. """Unit tests for lending a node's host memory to a Mooncake pool. -Runs without a Mooncake installation and without a GPU: the store is a fake -recording what ``setup`` was called with, since the contract being tested is -what the donor asks Mooncake for and how long it holds it, not what Mooncake -then does. +Runs without a Mooncake installation and without a GPU. The store is a fake +recording what `setup` was called with, since the contract under test is what +the donor asks Mooncake for and how long it holds it. """ import sys @@ -37,7 +36,7 @@ class FakeStore: - """The slice of ``MooncakeDistributedStore`` a donor drives.""" + """The slice of `MooncakeDistributedStore` a donor drives.""" instances = [] @@ -53,7 +52,7 @@ def setup(self, *args): @pytest.fixture def fake_bindings(monkeypatch): - """Stand in for ``mooncake.store``, which is not installed here.""" + """Stand in for `mooncake.store`, which is not installed here.""" FakeStore.instances = [] package = ModuleType("mooncake") store = ModuleType("mooncake.store") @@ -66,7 +65,7 @@ def fake_bindings(monkeypatch): @pytest.fixture def failing_bindings(fake_bindings): - """Bindings whose ``setup`` refuses, as an unreachable master would.""" + """Bindings whose `setup` refuses, as an unreachable master would.""" class Refusing(fake_bindings): @@ -104,8 +103,6 @@ def test_a_donor_registers_the_segment_it_was_asked_for(fake_bindings): assert protocol == "rdma" assert device_name == "mlx5_0" assert master == "10.0.0.1:50051" - # The donor never transfers, so its transfer buffer is dead weight -- but - # setup rejects a zero one, hence a token rather than nothing. assert local_buffer_size == DEFAULT_DONOR_LOCAL_BUFFER_SIZE @@ -187,17 +184,6 @@ def test_a_published_master_address_is_read_before_joining(fake_bindings, reacha assert fake_bindings.instances[0].setup_args[6] == "10.0.0.9:50051" -def test_the_segment_is_withdrawn_when_the_server_stops(fake_bindings, reachable_master): - """The handle is the segment: holding it for the server's life is the point.""" - donation = MooncakeDonationConfig(master_server_address="10.0.0.1:50051") - - with maybe_donate_segment(donation): - store = fake_bindings.instances[0] - # Nothing to assert on the fake beyond its existence -- the contract is - # that the reference is dropped, which is what unmounts the segment. - assert store.setup_args is not None - - def test_an_unreachable_master_is_reported_before_the_segment_is_offered( fake_bindings, monkeypatch): """Otherwise this is a status code from setup, with no address in it.""" diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py index b50c40ed74a4..a9c8f7e3e66e 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_master.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -15,8 +15,8 @@ """Unit tests for provisioning a Mooncake store pool during server bringup. Runs without a Mooncake installation and without a GPU. A master this process -launches is a fake standing in for ``Popen`` that opens the RPC port, which is -all the readiness handshake ever observes; a master someone else runs is a +launches is a fake standing in for `Popen` that opens the RPC port, which is +all the readiness handshake ever observes. A master someone else runs is a plain socket. """ @@ -50,11 +50,11 @@ def free_port() -> int: class FakeMasterProcess: - """The slice of ``Popen`` that launching a master actually drives. + """The slice of `Popen` that launching a master actually drives. - ``listen_on`` makes it answer on that port, which is what a real master - does last and what the readiness wait keys off. ``exit_code`` makes it a - master that failed to start. + `listen_on` makes it answer on that port, which is what a real master does + last and what the readiness wait keys off. `exit_code` makes it a master + that failed to start. """ def __init__(self, command, env, listen_on=None, exit_code=None): @@ -107,8 +107,8 @@ def clean_env(monkeypatch): def fake_master(monkeypatch): """Replace the master binary and its process with in-process fakes. - Returns a callable that arms the fake and, once provisioning has run, the - launched instance is available as ``.process`` for inspection. + Returns a callable that arms the fake. Once provisioning has run, the + launched instance is available as `.process` for inspection. """ class Launcher: @@ -118,8 +118,7 @@ def __init__(self): def arm(self, listen_on=None, exit_code=None, log_text=None): def popen(command, env=None, stdout=None, **_kwargs): - # A real master writes its own log through glog, and what it - # says there is the diagnosis when it fails to start. + # A real master writes its log through this handle. if log_text is not None and stdout is not None: stdout.write(log_text.encode()) stdout.flush() @@ -171,7 +170,7 @@ def test_pool_needs_exactly_one_master(): def test_publishing_an_address_needs_a_master_to_publish(): - """Reading a published address is master_server_address, not this.""" + """This option only writes an address; reading one is master_server_address.""" with pytest.raises(ValueError, match="needs launch_master"): MooncakeStoreConfig( master_server_address="host:50051", @@ -179,23 +178,20 @@ def test_publishing_an_address_needs_a_master_to_publish(): ) -def test_pool_is_rejected_on_another_connector(): +def test_pool_is_rejected_unless_the_connector_is_mooncake_store(): + """The validator keys off the connector, however that was spelled.""" with pytest.raises(ValueError, match="mooncake_store describes a Mooncake pool"): KvCacheConnectorConfig( connector="lmcache", mooncake_store=MooncakeStoreConfig(launch_master=True), ) - - -def test_pool_is_accepted_on_the_module_spelled_out(): - """A config naming the module instead of the preset is still the connector.""" - config = KvCacheConnectorConfig( + # Naming the module rather than the preset selects the same connector. + KvCacheConnectorConfig( connector_module="tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", connector_scheduler_class="MooncakeStoreConnectorScheduler", connector_worker_class="MooncakeStoreConnectorWorker", mooncake_store=MooncakeStoreConfig(launch_master=True), ) - assert config.mooncake_store.launch_master # ---- the rendered client config ---- @@ -295,7 +291,7 @@ def test_provisioning_fails_before_the_model_loads_if_the_master_is_absent(monke def test_an_unparseable_master_address_is_left_to_the_workers(): - """Not every address is host:port; that is the worker's problem, not ours.""" + """Not every address is host:port, so an unprobeable one passes through.""" pool = MooncakeStoreConfig(master_server_address="unix:///var/run/mooncake") with provision_pool(pool) as config_path: @@ -304,7 +300,7 @@ def test_an_unparseable_master_address_is_left_to_the_workers(): def test_an_inherited_config_path_wins(monkeypatch, tmp_path): - """The SLURM harness names the pool this way; provisioning must defer.""" + """An externally managed pool names itself this way, so provisioning defers.""" harness_config = tmp_path / "harness.json" harness_config.write_text("{}") monkeypatch.setenv(CONFIG_PATH_ENV, str(harness_config)) @@ -329,8 +325,8 @@ def test_a_launched_master_is_named_in_the_config_and_stopped_on_exit(fake_maste written = json.loads(open(config_path).read()) host, _, named_port = written["master_server_address"].rpartition(":") assert int(named_port) == port - # Whatever the config names has to be dialable: it is all a worker on - # another host is given. + # The address in the config is all a worker on another host gets, so + # it has to be dialable. with socket.create_connection((host, port), timeout=5): pass @@ -354,9 +350,8 @@ def test_a_launched_master_gets_the_flags_and_logging_it_needs(fake_master): assert f"--rpc_port={port}" in command assert f"--metrics_port={pool.master_metrics_port}" in command assert "--eviction_ratio=0.1" in command - # glog writes to files under /tmp unless told otherwise, which would - # leave the master's log -- the only view of the pool's own side of - # the conversation -- empty. + # Without these the master logs to a file under /tmp and the log the + # run directory holds stays empty. assert fake_master.process.env["GLOG_logtostderr"] == "1" assert fake_master.process.env["GLOG_v"] == "1" @@ -404,7 +399,7 @@ def test_a_run_dir_keeps_the_master_log_and_the_config(fake_master, tmp_path): with provision_pool(pool, run_dir=str(run_dir)) as config_path: assert config_path == str(run_dir / master_module.CLIENT_CONFIG_NAME) - # An explicit run directory has to outlive the run that filled it: the + # An explicit run directory outlives the run that filled it, since the # master's log is where pool occupancy and eviction are read from. assert (run_dir / master_module.MASTER_LOG_NAME).exists() assert (run_dir / master_module.CLIENT_CONFIG_NAME).exists() @@ -425,7 +420,7 @@ def test_no_connector_at_all_is_left_alone(): def test_a_pool_left_undescribed_stays_the_environment_contract(): - """Without ``mooncake_store``, MOONCAKE_CONFIG_PATH is still the only input.""" + """Without `mooncake_store`, MOONCAKE_CONFIG_PATH is still the only input.""" config = KvCacheConnectorConfig(connector="mooncake-store") with maybe_provision_pool(config): assert CONFIG_PATH_ENV not in os.environ @@ -466,7 +461,7 @@ def test_an_address_not_published_yet_is_waited_for(tmp_path): def test_an_empty_address_file_is_not_taken_for_an_address(tmp_path): - """It exists, which is not the same as holding somewhere to connect to.""" + """An existing file is not the same as a published address.""" published = tmp_path / "master.addr" published.write_text("") @@ -499,7 +494,7 @@ def test_a_standalone_master_publishes_an_address_that_can_be_dialed(fake_master def test_a_stopped_master_leaves_no_address_behind(fake_master, tmp_path): - """A stale address would send the next run's workers at a dead port.""" + """A stale address would send the next run's workers to a dead port.""" port = free_port() fake_master.arm(listen_on=port) address_file = tmp_path / "master.addr" @@ -515,7 +510,7 @@ def test_a_stopped_master_leaves_no_address_behind(fake_master, tmp_path): def test_a_standalone_master_keeps_its_log(fake_master, tmp_path): - """Its whole point is outliving servers, so its history is worth more.""" + """A standalone master outlives the servers that used it, so its log is kept.""" port = free_port() fake_master.arm(listen_on=port) run_dir = tmp_path / "run" @@ -528,7 +523,7 @@ def test_a_standalone_master_keeps_its_log(fake_master, tmp_path): def test_provisioning_joins_a_master_it_was_never_given_the_address_of(fake_master, tmp_path): - """The point of the file: no config and no script names a host.""" + """The address file is how workers reach a master no config names a host for.""" port = free_port() fake_master.arm(listen_on=port) address_file = tmp_path / "master.addr" @@ -549,7 +544,7 @@ def test_provisioning_joins_a_master_it_was_never_given_the_address_of(fake_mast def test_a_launched_master_publishes_where_its_run_left_its_logs(fake_master, tmp_path): - """So a finished run's logs still say which pool it used.""" + """A finished run's logs still say which pool it used.""" port = free_port() fake_master.arm(listen_on=port) run_dir = tmp_path / "run" @@ -563,7 +558,7 @@ def test_a_launched_master_publishes_where_its_run_left_its_logs(fake_master, tm def test_a_launched_master_can_be_published_where_the_donors_look(fake_master, tmp_path): - """The whole reason a server launching a master can still have donors.""" + """This is what lets a server that launched its own master have donors.""" port = free_port() fake_master.arm(listen_on=port) shared = tmp_path / "shared" / "master.addr" @@ -592,7 +587,7 @@ def test_a_half_written_address_is_never_read(tmp_path): def test_an_absent_master_is_named_rather_than_left_to_store_setup(monkeypatch): - """Otherwise this is a status code, in every rank, after the model loads.""" + """Otherwise the failure is a bare status code in every rank, after loading.""" monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") address = f"127.0.0.1:{free_port()}" @@ -600,17 +595,13 @@ def test_an_absent_master_is_named_rather_than_left_to_store_setup(monkeypatch): master_module.wait_for_master(address) -def test_a_master_that_answers_is_reported_with_the_wait_it_cost(running_master): - assert master_module.wait_for_master(running_master) is not None - - def test_an_address_of_a_shape_we_cannot_probe_is_not_fatal(): """Mooncake may accept addresses this cannot dial; leave them to it.""" assert master_module.wait_for_master("unix:///var/run/mooncake") is None def test_a_master_that_died_starting_is_reported_with_its_last_words(fake_master, tmp_path): - """The reason is in its log, which nobody reads unless it is quoted.""" + """The reason is in the master's log, which is only read if the error quotes it.""" run_dir = tmp_path / "run" fake_master.arm( exit_code=1, log_text="E0903 bind(50051) failed: Address already in use\n") @@ -655,7 +646,7 @@ def test_tcp_needs_no_device_and_looks_for_none(tmp_path): def test_a_node_without_infiniband_is_left_to_mooncake_s_own_discovery(tmp_path): - """Better than failing: Mooncake may still find something usable.""" + """Falling back beats failing, since Mooncake may still find a usable device.""" assert master_module.resolve_device_name( "rdma", "", sysfs_root=str(tmp_path / "absent")) == "" @@ -671,12 +662,3 @@ def test_the_detected_device_is_what_the_workers_are_told(fake_master, tmp_path, with provision_pool(pool, run_dir=str(tmp_path / "run")) as config_path: assert json.loads(open(config_path).read())["device_name"] == "mlx5_0" - - -def test_an_empty_master_log_says_what_that_means(tmp_path): - """Empty means it failed before glog opened, which reads as no log at all.""" - empty = tmp_path / "mooncake_master.log" - empty.write_text("") - - assert "empty" in master_module._log_tail(str(empty)) - assert "could not be read" in master_module._log_tail(str(tmp_path / "absent.log")) From d7db7f234b872fb6a90d4fff86942b1c0ea7e486 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:05:54 -0700 Subject: [PATCH 18/24] [None][chore] Drop the local mooncake experiment scratch from the branch mooncake_disagg/ and mooncake_usage.md were runbooks and configs for local experiments; docs/source/features/kv-cache-connector.md covers the connector for users. The two places that pointed at the scratch install script now point at docker/common/install_mooncake.sh, and the SLURM harness checks the image for the bindings instead of installing them per job. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../slurm/benchmark/disaggr_torch.slurm | 30 +- mooncake_disagg/README.md | 905 ------------------ mooncake_disagg/ctx_config.yaml | 30 - mooncake_disagg/disagg_config.yaml | 18 - mooncake_disagg/gen_config.yaml | 22 - mooncake_disagg/install_mooncake_runtime.sh | 118 --- mooncake_disagg/m3_agg_mooncake.yaml | 109 --- mooncake_disagg/m3_ctx_mooncake.yaml | 96 -- mooncake_disagg/m3_disagg_config.yaml | 20 - mooncake_disagg/m3_gen_mooncake.yaml | 84 -- mooncake_disagg/mooncake.json | 11 - mooncake_disagg/mooncake_api_surface_test.py | 115 --- mooncake_disagg/mooncake_smoke_test.py | 67 -- mooncake_usage.md | 269 ------ .../connectors/mooncake_store/master.py | 2 +- 15 files changed, 7 insertions(+), 1889 deletions(-) delete mode 100644 mooncake_disagg/README.md delete mode 100644 mooncake_disagg/ctx_config.yaml delete mode 100644 mooncake_disagg/disagg_config.yaml delete mode 100644 mooncake_disagg/gen_config.yaml delete mode 100755 mooncake_disagg/install_mooncake_runtime.sh delete mode 100644 mooncake_disagg/m3_agg_mooncake.yaml delete mode 100644 mooncake_disagg/m3_ctx_mooncake.yaml delete mode 100644 mooncake_disagg/m3_disagg_config.yaml delete mode 100644 mooncake_disagg/m3_gen_mooncake.yaml delete mode 100644 mooncake_disagg/mooncake.json delete mode 100644 mooncake_disagg/mooncake_api_surface_test.py delete mode 100644 mooncake_disagg/mooncake_smoke_test.py delete mode 100644 mooncake_usage.md diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm index e3acd8bda3be..dcbd5588936c 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm @@ -140,40 +140,22 @@ fi # The Mooncake store bindings, when a worker config asks for the connector. # Images built from this repo bake them in (docker/common/install_mooncake.sh), -# so the install below is only a fallback for images that predate it. Either -# way it is per job, since --container-name gives each node a container that -# lives for the job: anything installed here survives to the worker sruns but -# not into the next job. +# so this only confirms they are there: a job that discovers the gap later +# fails once per rank, deep in engine startup. mooncake_enabled=false if grep -qs "mooncake-store" "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"; then mooncake_enabled=true # Both halves of the wheel are checked because they fail at different # times: the connector needs mooncake.store in every context rank, and # 'trtllm-serve mooncake_master' needs the binary on PATH. - if srun --container-name=${container_name} \ + if ! srun --container-name=${container_name} \ --container-mounts=${container_mount} --no-container-mount-home \ --mpi=pmix --overlap -N 1 -n 1 \ bash -c 'python3 -c "import mooncake.store" && command -v mooncake_master' \ - &> ${full_logdir}/2_install_mooncake.log; then - echo "Mooncake store bindings already in the image; nothing to install" - else - mooncake_install_script="" - if [ -n "${trtllm_repo:-}" ]; then - mooncake_install_script="${trtllm_repo}/mooncake_disagg/install_mooncake_runtime.sh" - fi - if [ -z "${mooncake_install_script}" ] || [ ! -f "${mooncake_install_script}" ]; then - cleanup_on_failure "A worker config requests the mooncake-store connector, this image does not have the bindings, and mooncake_disagg/install_mooncake_runtime.sh was not found under trtllm_repo='${trtllm_repo:-}'. Use an image built from this repo, or set environment.trtllm_repo to a checkout that contains the script." - fi - echo "Installing Mooncake store bindings on all nodes..." - if ! srun --container-name=${container_name} \ - --container-mounts=${container_mount} --no-container-mount-home \ - --mpi=pmix --overlap -N $SLURM_NNODES --ntasks-per-node=1 \ - bash -c "MOONCAKE_WHEEL='${MOONCAKE_WHEEL:-}' bash ${mooncake_install_script}" \ - &>> ${full_logdir}/2_install_mooncake.log; then - cleanup_on_failure "Mooncake store bindings installation failed. Check ${full_logdir}/2_install_mooncake.log for details" - fi - echo "Mooncake store bindings installation completed successfully" + &> ${full_logdir}/2_check_mooncake.log; then + cleanup_on_failure "A worker config requests the mooncake-store connector, but this image has neither the mooncake.store bindings nor the mooncake_master binary. Build the image from this repo so that docker/common/install_mooncake.sh runs. Check ${full_logdir}/2_check_mooncake.log for details" fi + echo "Mooncake store bindings found in the image" fi # Get node lists and replace the placeholder with the actual node names diff --git a/mooncake_disagg/README.md b/mooncake_disagg/README.md deleted file mode 100644 index e160704212ae..000000000000 --- a/mooncake_disagg/README.md +++ /dev/null @@ -1,905 +0,0 @@ -# Validating the mooncake-store KV connector on MiniMax-M3 - -This is a runbook for testing the `mooncake-store` KV cache connector -(commits `f3c092187e`..`b8d3f43c43`) on MiniMax-M3 under load, using the SLURM -disaggregated benchmark harness in -`examples/disaggregated/slurm/benchmark/`. - -The unit tests in `tests/unittest/_torch/executor/test_mooncake_store_connector.py` -cover the pieces that decide whether a cache hit is *correct*: key namespacing, -hash chaining, page addressing, the startup gates. They deliberately do not run -a store, a model, or two engines. What is untested is everything that decides -whether the feature is *worth having*: whether a real prefix survives the round -trip, whether one M3 instance can replay a prefix another computed, and whether -the synchronous load path costs less than the prefill it avoids. - -## 1. The claim under test - -Local block reuse never leaves the instance that computed the prefix. The store -publishes KV pages into a shared, content-addressed CPU pool so any engine can -replay them. So there are exactly three things the store can do that local reuse -cannot, and each gets its own experiment in §7: - -1. **Cross-instance reuse.** A request routed to context instance B replays a - prefix computed on instance A. -2. **Survival across restarts.** Pages outlive the process that wrote them. -3. **Pool capacity beyond one host.** The pool is the sum of every worker's - segment rather than one node's host memory. - -### Why M3 makes this a hard test rather than an easy one - -The residency measurements in `../m3-kv-residency-measurement-README.md` (taken -on this same model and workload) found that M3 production traffic already serves -**97.0% of prompt tokens from local cache**, with eviction responsible for under -0.7% of misses. On a *single* instance there is almost no headroom for the store -to recover, since over 99% of misses are prefixes never cached anywhere, which -no store can serve either. - -That is not an argument against the feature; it is an argument about where to -look. Set expectations accordingly: - -- Do not expect a single-instance hit-rate improvement. Expect roughly zero. -- The store's value on M3 is concentrated in the cross-instance and - post-restart cases, where local reuse scores zero by construction. -- The connector's loads are **synchronous** (`start_load_kv`, before the forward - pass), so every loaded byte is fully exposed to TTFT. The host-offload tier it - replaces achieved 38-43% overlap at 45-51 GiB/s. At M3's context lengths a - loaded prefix is gigabytes, so a store hit is a win only when it displaces - real prefill, and a *needless* store hit is pure added latency. -- The design already rules out the worst version of that: the leader is handed - `num_computed_tokens` (the local match) and offers only blocks *beyond* it, so - the store cannot re-fetch something the GPU already holds. It cannot regress a - local hit; it can only add latency on a genuine local miss that it then fails - to make cheaper. That is what experiment 4 measures. - -## 2. Prerequisites: is Mooncake actually installed? - -**Short answer: probably not the part this connector needs.** Check before you -burn an allocation. - -Two different Mooncake components exist, and TensorRT-LLM's history treats them -differently: - -| Component | What uses it | How it gets installed | Since | -|---|---|---|---| -| C++ transfer engine (`/usr/local/Mooncake`) | the C++ cache transceiver's Mooncake backend | CMake source build in `docker/common/install_mooncake.sh` | PR #8447, Nov 2025 | -| Python bindings (`mooncake.store.MooncakeDistributedStore`) | **this connector** | pip wheel, added to the same script | commit `d36dae435e` | - -So Mooncake has been in the container images for months, but only usefully as -the C++ library. Three consequences: - -- Any image built before commit `d36dae435e` lacks a working set of bindings. - The image pinned in `jenkins/current_image_tags.properties` is tagged - `202607211045` (2026-07-21), which predates that commit, so **the currently - pinned CI image does not have them**. -- The CMake install *does* drop a `mooncake` Python package into the image, but - it is both broken and actively harmful: it shadows the working one. This is the - single biggest time sink in this section; see "Fix" below. -- The wheel is also the only usable source of the `mooncake_master` and - `mooncake_http_metadata_server` entry points. The CMake build's - `/usr/local/Mooncake/bin/mooncake_master` does exist and does run, but it is - the 0.3.7 build and it is not what ends up on `PATH` once the wheel is - installed. - -Also note `install_mooncake.sh` runs only in the `tritondevel` stage of -`docker/Dockerfile.multi`, and is skipped entirely on Rocky8. The CI image and -the internal `trtllm_build` release image descend from `tritondevel`, so they -get it; the NGC `release` image descends from the plain `devel` stage, so it -does not. - -### Verify - -Inside the container you will actually run: - -```bash -python3 -c "from mooncake.store import MooncakeDistributedStore; print('store bindings OK')" -command -v mooncake_master || ls /usr/local/Mooncake/bin -``` - -### Fix - -Run `mooncake_disagg/install_mooncake_runtime.sh` inside the container. It takes -under ten seconds on a warm pip cache, it is idempotent, and it verifies itself, -so it is safe in a SLURM prolog on every node. - -```bash -bash mooncake_disagg/install_mooncake_runtime.sh -``` - -A bare `pip3 install mooncake-transfer-engine` is *not* enough, and its failure -mode is what makes this step so confusing: pip reports success and the import -still fails. Two independent problems, both of which the script handles. - -**1. A broken `mooncake` package that pip cannot displace.** The CMake build in -`install_mooncake.sh` emits its own `mooncake` Python package, omitting -`libmooncake_store.so`, so it cannot load. `mooncake-integration/CMakeLists.txt` -chooses where to put it with: - -```cmake -COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print([s for s in sys.path if 'packages' in s][0])" -``` - -That is the *first* `sys.path` entry whose name merely contains `"packages"`. It -gives two different failures depending on what else is installed, and both -produce the same confusing symptom: an `ImportError` **after a `pip install` -that reported success**. - -- **With `nvidia-cutlass-dsl` present** (the normal case: the `devel` stage - uninstalls it at `Dockerfile.multi:58`, then `constraints.txt` pulls it back - in) the first match is `nvidia_cutlass_dsl/dsl_packages`, because - `nvidia_cutlass_dsl_packages.pth` does `sys.path.insert(0, ...)` on it. It - therefore outranks `dist-packages` on every interpreter start and shadows the - wheel permanently. CUTLASS DSL never references `mooncake`, so deleting it - breaks nothing. -- **Without it**, the match is `dist-packages` itself and the broken package - *collides* with the wheel. This one is nastier: CMake writes - `store.cpython-312-.so` while the wheel writes `store.so`, and - `importlib.machinery.EXTENSION_SUFFIXES` puts the interpreter-tagged suffix - first, so the broken extension still wins. pip also overwrites `__init__.py`, - erasing the `# Auto-generated by CMake` marker, so afterwards there is no - reliable way to tell leftover files from wheel files. - -Because of that second case, `install_mooncake_runtime.sh` removes any -`mooncake` package directory outright and reinstalls, rather than trying to -identify individual bad files. Both cases are covered. - -**2. The default wheel is built for CUDA 12.** `mooncake-transfer-engine` links -against `libcudart.so.12`; containers from `pytorch-26.05` on ship CUDA 13 only. -Every extension and the `mooncake_master` binary then fail to load: - -``` -ImportError: libcudart.so.12: cannot open shared object file -``` - -Use **`mooncake-transfer-engine-cuda13`**, the same project built for CUDA 13, -which needs no shim. Its releases start at 0.3.9, so it cannot match the -`MOONCAKE_VERSION` pin (`0.3.7.post2`) in `install_mooncake.sh`. That drift is -safe here: `/usr/local/Mooncake` backs the *cache transceiver's* Mooncake -backend, a different feature that these configs do not use -(`cache_transceiver_config.backend: "NIXL"`). The connector only ever talks to -the wheel, and the wheel also supplies the `mooncake_master` that lands on -`PATH` ahead of the CMake one, so client and master stay matched. Revisit this -only if you set the transceiver backend to `MOONCAKE`. - -If you would rather match `install_mooncake.sh` exactly, the script keeps that -path working and applies the `libcudart.so.12` shim for you: - -```bash -MOONCAKE_WHEEL="mooncake-transfer-engine==0.3.7.post2" \ - bash mooncake_disagg/install_mooncake_runtime.sh -``` - -Both wheel choices were validated against the full set of store methods the -connector calls; see "Validating the install" below. - -### How often does the script need to run? - -It depends on whether the container filesystem persists, because the script -writes into `dist-packages` inside the container, not into your checkout. - -| Situation | How often | -|---|---| -| Long-lived container you `docker exec` into | **Once.** It survives until the container is deleted; `docker restart` keeps it. | -| SLURM via `disaggr_torch.slurm` | **Once per job, per node**, which the harness does for you, see below. | -| Image built from this repo | **Never.** `install_mooncake.sh` does it at build time and fails the build if the import does not work. | - -For the SLURM case this is already wired up: `disaggr_torch.slurm` runs the -script on every node, right after its `pip install -e .[devel]` step, and gates -it on whether a worker config actually asks for the connector: - -```bash -if grep -qs "mooncake-store" "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"; then -``` - -So arms A and B of the run matrix pay nothing, arm C installs automatically, and -there is no new config key to remember. It resolves the script as -`${trtllm_repo}/mooncake_disagg/install_mooncake_runtime.sh` and fails the job -with an explicit message if `environment.trtllm_repo` is unset or does not -contain it, which is the case if you benchmark from -`environment.trtllm_wheel_path` instead, so use an image with the bindings baked -in for that path. Output lands in `/2_install_mooncake.log`. Set -`MOONCAKE_WHEEL` in the submitting environment to override the wheel; it is -forwarded to every node. - -Note that `--container-name` gives each node a container that lives for the whole -job, so the install survives from that step through to the worker `srun`s. It -does not survive into the *next* job, which is why this runs per job rather than -once. Do not try to persist it via `~/.local` unless home is genuinely shared and -mounted (`disaggr_torch.slurm` passes `--no-container-mount-home` to most of its -`srun` calls). - -Baking an image is the only option that removes the step entirely. Given the -script takes about eight seconds from cold, that is a convenience decision rather -than a necessity. - -### Will the shadow package come back? - -The root cause is upstream in Mooncake's `CMakeLists.txt` and is **not** fixed; -both scripts clean up after it. Practically: - -- **Images built from this repo:** no. The cleanup runs in the same script, - immediately after `make install` and before the wheel install, and the build - fails if `import mooncake.store` does not work. -- **Any pre-existing image**, including the one pinned in - `jenkins/current_image_tags.properties` (`202607211045`): the broken package is - baked in, so the runtime script is required. -- **Inside a running container:** only if something re-runs Mooncake's CMake - install. Reinstalling `nvidia-cutlass-dsl` does not recreate it: that package - has never shipped a `mooncake` directory, and only supplies the `.pth` that - made CMake choose the wrong destination. -- **If Mooncake is ever upgraded** to a version that fixes its install path, or - the `.pth` ordering changes, the cleanup becomes a no-op rather than a hazard. - -### Validating the install - -Two scripts in this directory, in increasing order of strictness. Both need a -running master and `MOONCAKE_CONFIG_PATH`, exactly like a real worker: - -```bash -mooncake_master --rpc_port=50051 --metrics_port=9004 & - -export MOONCAKE_CONFIG_PATH=$PWD/mooncake.json # TCP config; edit the master address -python3 mooncake_disagg/mooncake_smoke_test.py # setup + put/get round trip -python3 mooncake_disagg/mooncake_api_surface_test.py # needs a GPU -``` - -`mooncake_smoke_test.py` proves the bindings load and `store.setup()` succeeds -with the same argument list `worker.py` passes. `mooncake_api_surface_test.py` is -the one that matters when changing wheel versions: the connector's hot path never -uses `put`/`get`, it uses `register_buffer` plus the zero-copy -`batch_put_from_multi_buffers` / `batch_get_into_multi_buffers` / `batch_is_exist` -calls against registered GPU pages. Those take `list[list[int]]`, one buffer -list per key, because `PageAddressing.page_buffers` returns one address per -layer-group region, and that is the signature most likely to drift. - -Then the unit tests, which need no store and no GPU: - -```bash -pytest tests/unittest/_torch/executor/test_mooncake_store_connector.py -``` - -## 3. Topology - -``` - mooncake_master (1 CPU core, its own job) - ^ ^ ^ - register/put/get| | |mount segment only - ┌─────────────────────┴──┐ ┌──┴──────────────────────┐ │ - │ CTX instance 0 TP=4 │ │ CTX instance 1 TP=4 │ │ store: role=both - └────────────┬───────────┘ └───────────┬─────────────┘ │ - │ NIXL KV handoff │ │ - └──────────┬────────────────┘ │ - v │ - ┌──────────────────────────┐ ┌──────────────┴────────────┐ - │ GEN instance TP=4 │ │ segment donor (same node) │ - └──────────────────────────┘ └───────────────────────────┘ - ^ no connector, no put/get, - round-robin│ contributes host memory - ┌──────────┴───────────┐ - │ trtllm-serve disagg │ <- benchmark_serving client - └──────────────────────┘ -``` - -12 GPUs = 3 nodes at 4 GPUs/node. The generation worker deliberately has **no** -`kv_connector_config`: generated tokens are rarely a reused prefix, and that -absence is the only way to express "off" (`StoreRole` has no off value). It also -lets the generation worker keep its host cache tier and `MAX_UTILIZATION` -scheduler, both of which the connector would forbid. - -### Why the generation node still needs a donor process - -Pool capacity comes only from processes that open a store handle: `setup` -registers `global_segment_size` bytes of the calling process's host memory, and -the master then places blocks in it. Since only the context workers configure -the connector, only they contribute memory, so by default every byte of the pool -is prefill-node DRAM and the store is a prefill-DRAM-caches-prefill-GPU tier -that largely duplicates TensorRT-LLM's native host offload. Confirm this on -any run by grouping the master's `allocation_succeeded ... segment=:` -lines by host: a single host means a prefill-only pool. - -`mooncake_donation` on the generation worker closes that gap. The server opens -a handle, contributes memory, and then holds it without a single put or get for -as long as it runs, so the pool spans both sides while its engine stays -connector-free and keeps its cache transceiver for the KV handoff: - -```yaml -# gen worker config: no kv_connector_config anywhere near it -mooncake_donation: - master_server_address: file:///$WORK_DIR/master.addr - segment_size: 640GiB - protocol: rdma -``` - -Donation is deliberately outside `kv_connector_config`, and not a `StoreRole` -either: the roles describe traffic (`producer` writes, `consumer` reads, -`both`), none of them means "contribute memory only", and configuring capacity -there would start this server using the store. The size is charged **per server -process, not per rank**, unlike `global_segment_size`, so two servers on one -node lend twice this. - -The server is ready only once its segment is mounted, which makes readiness the -signal that the pool has the capacity, and means the first blocks prefill writes -can already land on a decode node. Set `segment_size: 0`, or leave the section -out, to keep the pool prefill-only. - -A node running no server lends as its own command instead, which is also how a -machine with no GPUs contributes: - -```bash -trtllm-serve mooncake_donor --master_server_address 10.0.0.1:50051 \ - --segment_size 160GiB --protocol rdma --device_name mlx5_0 -``` - -The donated memory is charged to the donor process and competes with the -generation worker's own `kv_cache_config.host_cache_size` on that node, so size -the two together. The worker logs its own share as `KV cache manager v2 host -cache quota set to N GiB`, **per rank**, against the `available host memory` it -reports on the same line. - -## 4. Step 1: run the Mooncake master - -`master_server_address` is mandatory, so a master must exist and be reachable -from every worker. - -**Outside SLURM you can skip this section too.** A worker config carrying -`kv_connector_config.mooncake_store` makes `trtllm-serve` provision the pool -during its own bringup, with `launch_master: true` starting a master for that -server alone and `master_server_address` joining one that already exists, and -write the client config itself. That covers aggregated and single-instance runs, -as `m3_agg_mooncake.yaml` does; `mooncake_usage.md` §2 has the table. The -rest of this section is about the master the experiments below need, which -outlives any one server and therefore cannot be owned by one. - -**For a single-job experiment you can skip this section.** The harness starts -no master and writes no client config: the context worker's `launch_master: -true` does both, on the context node, and publishes the address the generation -workers' `mooncake_donation` reads. `disaggr_torch.slurm` contributes exactly -two things: it installs the bindings on every node, and it substitutes -`__LOG_DIR__` in the worker configs, since the run directory is the one value a -config written before submission cannot know. Everything else, the pool sizes -and the HCA included, is in the config, and no `MOONCAKE_*` variable is read -from the submitting environment. - -The master's log lands in `/mooncake_master.log` and its address in -`/master.addr` while it runs, because `start_worker.sh` sets -`TRTLLM_MOONCAKE_RUN_DIR` to the log directory. That is also what lets the -context server's other ranks read the rendered `mooncake.json`: they are -separate srun tasks that never inherited the leader's environment. - -That master dies with the job, so read on if you need a pool that outlives one -allocation, which experiment 3 does by construction. Run it as its own -long-lived job and export `MOONCAKE_MASTER_ADDRESS=:50051` before -`submit.py`; the harness then skips launching one and only writes the client -config pointing at yours. - -```bash -# mooncake_master.sbatch -#!/bin/bash -#SBATCH --job-name=mooncake-master -#SBATCH --nodes=1 -#SBATCH --time=08:00:00 -#SBATCH --output=%x-%j.out - -srun --container-image=$CONTAINER_IMAGE \ - --container-mounts=$WORK_DIR:$WORK_DIR \ - trtllm-serve mooncake_master \ - --rpc_port 50051 \ - --metrics_port 9004 \ - --address_file $WORK_DIR/master.addr \ - --run_dir $WORK_DIR -``` - -The master runs for as long as the command does, and `--address_file` receives -`host:port` **once it accepts connections**, so waiting for that file is waiting -for readiness, and its absence after the job starts is a failure rather than a -slow start. It is removed on exit, so a stale address is never dialed. -`--run_dir` keeps the master's log at `$WORK_DIR/mooncake_master.log`, which is -where pool occupancy and eviction are read from. - -`TRTLLM_MOONCAKE_MASTER_BINARY` overrides the binary this runs, and -`TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) how long it waits for the port. -Run `mooncake_master --help` inside the container for the flags this does not -surface (`--rpc_address`, `--rpc_thread_num`, `--default_kv_lease_ttl`, -`--eviction_high_watermark_ratio`, `--enable_http_metadata_server`, -`--cluster_id`, `--root_fs_dir`). - -Keeping the master in a separate job is what makes experiment 3 (§7) possible: -the pool outlives the engines, so a second benchmark job finds a warm store. - -Workers are pointed at it without anyone writing the address down: -`master_server_address: file://$WORK_DIR/master.addr` in a config's -`mooncake_store` block makes each server read it during bringup and wait if the -master job has not started yet. That is what makes a master whose host the -scheduler chose usable from a config settled beforehand. - -Failing that, write the client config, substituting the address the master job -just recorded, or let `disaggr_torch.slurm` generate it, as above. The schema -is vLLM's, so one pool can serve both engines: - -```bash -MASTER_IP=$(cat $WORK_DIR/master.addr) -cat > $WORK_DIR/mooncake.json <:/metadata` with - `mooncake_http_metadata_server` (or the master's own - `--enable_http_metadata_server`) only if you need the shared-metadata - behaviour; confirm the port from `--help` rather than assuming. -- `device_name`: pick from `ibv_devinfo` on a compute node. For first bring-up - only, `"protocol": "tcp"` with `"device_name": ""` removes RDMA from the - variable list, which is what `mooncake.json` in this directory does. Do not - draw performance conclusions from a TCP run. -- `global_segment_size` is contributed **per worker process**, so the pool is - `global_segment_size x (ctx instances x TP)` = 8 segments here. -- Sizing: after startup, each worker logs its page geometry (§8). Pool bytes for - a corpus of `T` unique prefix tokens is - `T / tokens_per_block x Σ_layer_groups bytes_per_page x world_size`. - As an anchor, the residency work measured M3 at ~22 KiB/token aggregate across - TP=4 (fp8 KV plus a per-rank-replicated index-K), so ~21 GiB per million - unique prefix tokens. Confirm against your own log line rather than trusting - that number. -- `role`/`cache_prefix` can also be overridden per process by - `TRTLLM_MOONCAKE_STORE_ROLE` and `TRTLLM_MOONCAKE_STORE_PREFIX`. Bump the - prefix whenever you change anything that should not be shared with an earlier - run's pages. - -## 5. Step 2: the harness config - -Copy `examples/disaggregated/slurm/benchmark/config.yaml` and replace the -`worker_config` section with M3's. `submit.py` serializes `worker_config.ctx` -and `worker_config.gen` straight to `ctx_config.yaml`/`gen_config.yaml` with -`yaml.dump`, so any LLM-API key passes through untouched, `kv_connector_config` -included. - -The context worker below is `m3_ctx_mooncake.yaml` from this directory; the -generation worker is `m3_gen_mooncake.yaml`. Every deviation from the production -M3 config is marked and explained there, and those comments are the reason to -read those two files rather than treating this block as self-explanatory. - -```yaml -# m3_store_2ctx.yaml -slurm: - script_file: "disaggr_torch.slurm" - partition: "" - account: "" - job_time: "04:00:00" - job_name: "m3-mooncake-store" - extra_args: "" - set_segment: true - numa_bind: true # GB200/GB300 NVL72 - -benchmark: - mode: "e2e" - use_nv_sa_benchmark: false - multi_round: 8 # num_prompts = concurrency x multi_round - streaming: true - concurrency_list: "8" - input_length: 131072 # log-dir naming only; the dataset is authoritative - output_length: 1024 - dataset_file: "/m3_shared_prefix.jsonl" - -hardware: - gpus_per_node: 4 - num_ctx_servers: 2 # >= 2 is the whole point; see experiment 2 - num_gen_servers: 1 - -environment: - container_mount: "" - container_image: "" - model_path: "" - trtllm_repo: "" - build_wheel: false - trtllm_wheel_path: "" - work_dir: "" - worker_env_var: "TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes NCCL_GRAPH_MIXING_SUPPORT=0" - # Only the context workers open a store handle. MOONCAKE_CONFIG_PATH is - # deliberately absent: the context server renders that file itself, into the - # log directory, and its own ranks read it back from there. - ctx_worker_env_var: "TRTLLM_MOONCAKE_STORE_ROLE=both TRTLLM_MOONCAKE_STORE_PREFIX=trtllm-m3-run1" - server_env_var: "TRTLLM_SERVER_DISABLE_GC=1" - -profiling: - nsys_on: false - ctx_profile_range: "10-30" - gen_profile_range: "200-250" - -accuracy: - enable_accuracy_test: false - tasks: {} - -worker_config: - ctx: - # ---- contents of m3_ctx_mooncake.yaml, plus parallelism ---- - tensor_parallel_size: 4 - moe_expert_parallel_size: 4 - pipeline_parallel_size: 1 # gated: connector refuses PP > 1 - context_parallel_size: 1 # gated: connector refuses CP > 1 - enable_attention_dp: false # required: dummy DP-balancing requests reach the hooks - max_seq_len: 1048576 - max_num_tokens: 16384 - max_batch_size: 20 - sparse_attention_config: - algorithm: minimax_m3 - implementation: msa - indexer_kv_dtype: fp8 - sparse_disable_index_value: true # gated: index-V lives outside the paged pools - fuse_qkv_index_projection: true - kv_cache_config: - free_gpu_memory_fraction: 0.94 - enable_block_reuse: true - tokens_per_block: 128 - use_kv_cache_manager_v2: true # required: only V2 can describe its pools - dtype: fp8 - event_buffer_max_size: 0 - host_cache_size: 0 # gated: must be explicit 0, not omitted - disk_cache_size: 0 - scheduler_config: - capacity_scheduler_policy: GUARANTEED_NO_EVICT # gated - cache_transceiver_config: - backend: "NIXL" - transceiver_runtime: "PYTHON" # M3 is always-V2; C++ transceiver is refused - enable_chunked_prefill: true - enable_autotuner: true - trust_remote_code: true - reasoning_parser: minimax_m3 - stream_interval: 20 - print_iter_log: true - num_postprocess_workers: 8 - # Required to see any reuse number at all; see section 8. All three - # default to false, and without them /metrics returns an empty list. - enable_iter_perf_stats: true - enable_iter_req_stats: true - return_perf_metrics: true - kv_connector_config: - connector: mooncake-store # <-- the only line experiment 1 removes - - gen: - # ---- contents of m3_gen_mooncake.yaml; no connector, so no gates ---- - tensor_parallel_size: 4 - moe_expert_parallel_size: 4 - enable_attention_dp: false - max_seq_len: 1048576 - max_num_tokens: 16384 - max_batch_size: 20 - sparse_attention_config: - algorithm: minimax_m3 - implementation: msa - indexer_kv_dtype: fp8 - sparse_disable_index_value: true # must match ctx: it changes the model - fuse_qkv_index_projection: true - kv_cache_config: - free_gpu_memory_fraction: 0.94 - enable_block_reuse: true - block_reuse_policy: per_conversation - tokens_per_block: 128 - use_kv_cache_manager_v2: true - dtype: fp8 - event_buffer_max_size: 0 - host_cache_size: 388554555392 # kept: no connector here - scheduler_config: - capacity_scheduler_policy: MAX_UTILIZATION # kept: no connector here - cache_transceiver_config: - backend: "NIXL" - transceiver_runtime: "PYTHON" # must match ctx - enable_chunked_prefill: true - enable_autotuner: true - trust_remote_code: true - reasoning_parser: minimax_m3 - stream_interval: 20 - print_iter_log: true - num_postprocess_workers: 8 - enable_iter_perf_stats: true - enable_iter_req_stats: true - return_perf_metrics: true -``` - -Eagle3 is left off. It is not gated, but `MiniMaxM3KVCacheManagerV2` sets -`supports_shared_draft_layers`, so draft layers join the unified V2 cache and -therefore the registered layout. That is extra page geometry the store must key -correctly, on a path with no coverage, so turn it on only after a clean run -without it. - -Submit with: - -```bash -cd examples/disaggregated/slurm/benchmark -python3 submit.py -c /m3_store_2ctx.yaml --dry-run # inspect first -python3 submit.py -c /m3_store_2ctx.yaml -``` - -## 6. Step 3: the workload - -`run_benchmark.sh` invokes `benchmark_serving` with -`--dataset-name trtllm_custom --dataset-path `, so you supply a -JSONL file. `CustomDataset` reads `input.messages[1].content` as the prompt, -`input.max_tokens` as the output length, and skips re-tokenization when -`input.num_tokens` is present. It shuffles the file on load, which is what -spreads repeated prefixes apart in time. - -The workload must have **repeated prefixes across requests**, because that is -the only structure a content-addressed store can exploit. `P` distinct prefixes -each repeated `R` times, with a unique suffix per request so no two requests are -identical: - -```python -# gen_shared_prefix_dataset.py -import json, random -from transformers import AutoTokenizer - -MODEL = "" -NUM_PREFIXES = 8 # P distinct shared prefixes -REPEATS = 8 # R requests per prefix -> P*R = 64 total -PREFIX_TOKENS = 131072 # must be >> tokens_per_block (128) to be worth storing -SUFFIX_TOKENS = 512 -OUTPUT_TOKENS = 1024 -OUT = "m3_shared_prefix.jsonl" - -tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) -rng = random.Random(1234) -vocab = tok.vocab_size - -def text_of(n_tokens, seed): - r = random.Random(seed) - ids = [r.randrange(1000, vocab - 1000) for _ in range(int(n_tokens * 1.3))] - text = tok.decode(ids, skip_special_tokens=True) - # Re-encode and trim: decode/encode is not a round trip, so measure. - ids = tok.encode(text, add_special_tokens=False)[:n_tokens] - return tok.decode(ids, skip_special_tokens=True), len(ids) - -prefixes = [text_of(PREFIX_TOKENS, 100 + i) for i in range(NUM_PREFIXES)] - -with open(OUT, "w") as f: - for i, (prefix, plen) in enumerate(prefixes): - for r in range(REPEATS): - suffix, slen = text_of(SUFFIX_TOKENS, 900000 + i * 1000 + r) - f.write(json.dumps({"input": { - "messages": [{"role": "system", "content": ""}, - {"role": "user", "content": prefix + suffix}], - "max_tokens": OUTPUT_TOKENS, - "num_tokens": plen + slen, - }}) + "\n") -``` - -Size the file against what the client will actually request. -`run_benchmark.sh` computes -`num_prompts = (concurrency x num_gen_servers) x multi_round`, so the §5 config -(`concurrency_list: "8"`, `multi_round: 8`, one generation server) asks for 64 -prompts, which is why `P x R` above is 64. Ask for more than the file holds -and the extra is not sampled; write more than you ask for and the tail of your -repeat structure never runs. - -Random token text is deliberate: it defeats any accidental prefix sharing -between "distinct" prefixes, so the hit rate you measure is the one you -designed. If you would rather test genuine production traffic, substitute a -real multi-turn trace, but keep an eye on whether it actually contains repeated -prefixes, since without them the store has nothing to do and a flat result means -nothing. - -## 7. Step 4: the run matrix - -Three arms, and the middle one is the one people skip: - -| Arm | Config | Purpose | -|---|---|---| -| **A** production reference | today's M3 config: host tier on, `MAX_UTILIZATION`, no connector | where you are today | -| **B** gated baseline | arm A's config edited to satisfy every connector gate (`host_cache_size: 0`, `GUARANTEED_NO_EVICT`, ...), still no connector | isolates what the gates cost | -| **C** store | arm B plus `kv_connector_config` | isolates what the store adds | - -Comparing C against A alone conflates two independent changes: the store's -benefit and the loss of a 362 GiB host cache tier plus a scheduler policy -change. **The store's efficacy is C vs B.** A vs B tells you the entry price, -and A vs C tells you whether the whole package is deployable. All three are -worth knowing and they answer different questions. - -Then, within that: - -**Experiment 1: does it work at all (`num_ctx_servers: 1`).** -Arm C, one context instance, small `PREFIX_TOKENS` (say 4096) and a short run. -You are looking for a clean startup, the registration log line, non-zero store -traffic, no load failures, and coherent output text. Do this before spending an -allocation on anything larger. Expect no throughput change; local reuse already -serves this case. - -**Experiment 2: cross-instance reuse (`num_ctx_servers: 2`).** -The router defaults to round-robin, so consecutive requests alternate between -context instances and roughly half of each prefix's repeats land on the instance -that did not compute it. Those are the requests local reuse must recompute from -scratch and the store can serve. Compare arm C against arm B on: -- TTFT p50/p99 (the store's whole thesis is prefill avoided) -- `reused_blocks_per_request` distribution (§8) -- output tokens/s/GPU - -This is the primary result. A store that does not win here does not work. - -**Experiment 3: survival across restarts.** -Run experiment 2's arm C twice, same `TRTLLM_MOONCAKE_STORE_PREFIX`, with the -master job left running between them. The second job starts with an empty local -cache but a warm pool. First-round TTFT should fall toward the warm steady-state -value. Local reuse scores zero here by construction, so any improvement is -attributable to the store alone, which makes this the cleanest signal in the -whole matrix, and the cheapest to run. - -**Experiment 4: the cost when there is nothing to gain.** -Arm C against arm B on a workload with *no* repeated prefixes (unique prompts). -This measures pure overhead: lookups, key hashing on the leader, background -saves competing for host bandwidth. Ideally indistinguishable from arm B. This -is the arm that catches a feature that helps its benchmark and hurts the fleet. - -## 8. Step 5: reading the results - -### Did the connector even load? - -Every context worker logs at INFO on startup: - -``` -mooncake-store leader ready (role=both, tokens_per_block=128) -mooncake-store worker rank 0/4 ready (role=both, model_key=MiniMax-M3-NVFP4, master=10.0.0.5:50051) -mooncake-store worker rank 0 registered layout: tokens_per_block=128, lg0(layers=N, regions=..., bytes/page=..., slots=..., window=None) -``` - -```bash -grep -h "mooncake-store" /3_output_CTX_*.log | head -40 -``` - -The registration line is the one to keep: `bytes/page` per layer group is what -your pool sizing in §4 depends on, and `window=None` is the confirmation that no -sliding-window group is present (one would have aborted startup). - -### Is the store actually moving pages? - -Hit and transfer counts are at DEBUG. The connector logs under module `_torch`, -so: - -``` -TLLM_LOG_LEVEL_BY_MODULE="debug:_torch" -``` - -added to `environment.ctx_worker_env_var`. This is verbose, since it enables -DEBUG for all of `_torch`, so use it for experiment 1 and for diagnosis rather -than for the runs you intend to quote numbers from. The lines worth counting: - -``` -mooncake-store matched N blocks (M tokens) for request R # leader, a hit -mooncake-store rank K loaded P pages # worker, a load -``` - -The store's own counters are the alternative that costs nothing at runtime: -`mooncake_master --metrics_port=9004` exposes pool-level statistics over HTTP. -Scrape it before and after a run and diff. - -### Where did the pages land? - -Page counts alone do not say whether the pool is doing anything the native host -offload tier could not. For that, group the master's allocations by segment -host. A segment is one client process's donated memory, so the host tells you -which node the block physically lives on: - -```bash -grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" mooncake_master.log \ - | awk '{sub(/size=/,"",$2); sub(/segment=/,"",$3); split($3,p,":"); - n[p[1]]++; b[p[1]]+=$2} - END {for (h in n) printf "%-16s pages=%-7d %.2f GiB\n", h, n[h], b[h]/1073741824}' -``` - -One host means a prefill-only pool (see §3). Two or more, with the generation -node among them, means blocks written by prefill are living on decode-side DRAM -and being read back from there. `disaggr_torch.slurm` writes this breakdown into -`9_mooncake_summary.log` at the end of every run, alongside the donor hosts, so -it needs running by hand only when diagnosing a partial run. - -Requires `GLOG_v=1` on the master, which `trtllm-serve mooncake_master` sets -unless `GLOG_v` is already in its environment, so raise it by exporting `GLOG_v` -to that command. - -### Which reuse number means what - -This distinction matters and is easy to get backwards: - -| Signal | Where | Includes store hits? | -|---|---|---| -| `reused_blocks_per_request`, `kv_cache_hit_rate_per_request` | per-request iteration stats | **Yes.** `_reserve_connector_prefix` calls `set_prepopulated_prompt_len` with the connector-served position, and these derive from `mPrepopulatedPromptLen`. | -| `kv_cache_iter_reused_blocks`, `kv_cache_iter_reuse_rate` | `GET /prometheus/metrics` | **No.** These come from the local V2 reuse tree's committed stats. | - -So **store hits are roughly per-request reuse minus local-tree reuse**. Confirm that -relationship on experiment 3, where the local tree starts empty and the -difference is unambiguous, before relying on it elsewhere. - -Getting at either one requires the three flags added to the worker configs in -§5, all of which default to false: - -- `enable_iter_perf_stats: true`, without which `get_latest_iteration_stats` - short-circuits and `GET /metrics` returns `[]`. -- `enable_iter_req_stats: true`, needed for the *per-request* half of the - table above. -- `return_perf_metrics: true`, which mounts `/prometheus/metrics`. `GET /metrics` - (plain JSON iteration stats) is routed unconditionally but still needs - `enable_iter_perf_stats`. - -`print_iter_log: true` is worth keeping on, but note it prints iteration timing -and KV *utilization* only, with no reuse counters. Do not go looking for hit -rates there. - -Per-request client-side results land in `/concurrency_/result.json` -with TTFT/TPOT/ITL/E2EL percentiles, which is where the headline numbers for -§7's arms come from. - -### Failure signatures - -| Log line | Meaning | -|---|---| -| `mooncake-store failed to load N of M pages` (raises) | **Stop.** The runtime had already counted those tokens as computed, so this is the tripwire against silently wrong answers. Do not treat as flaky. | -| `mooncake-store background save failed` | A save thread exception, re-raised on the executor thread. | -| `mooncake-store rank K failed to save N of M pages` (warning) | Dropped write. Costs a future miss, not correctness. A trickle is tolerable; a flood means the pool is full or the master is overloaded. | -| `mooncake-store lookup failed; treating as a miss` (warning) | Probe failed. Degrades to no-store behavior. | -| `could not reserve connector prefix up to N, falling back to the local match` (debug) | Out of GPU pages. The store offered more than the engine could hold, which is expected under pressure, but frequent occurrences mean the offer is outrunning capacity. | - -### Sanity check that is not a performance number - -Run a handful of prompts through arms B and C with temperature 0 and compare the -text. A store that returns the wrong bytes shows up as degraded output long -before it shows up as an error. `accuracy.enable_accuracy_test: true` with gsm8k -gives a coarser version of the same check. - -## 9. Things that will bite - -- **`host_cache_size: 0` must be written explicitly.** Left at its default of - `None`, V2 still provisions a host tier, and the gate rejects it. Falsy is not - the same as absent here. -- **The gates change the config out from under you.** `GUARANTEED_NO_EVICT` - instead of `MAX_UTILIZATION`, no host tier, no attention DP. That is why - arm B exists. -- **`sparse_disable_index_value: true` changes the model, not just the cache.** - Hold it fixed across every arm, generation workers included, or you are - comparing two different models. -- **Key namespace pins world size and rank.** Change TP and every stored page - becomes unreachable, as a miss rather than an error. Same for `tokens_per_block`, the - layer group set, and `bytes_per_page`. -- **`model_key` defaults to the checkpoint directory's basename.** Two hosts - mounting the same checkpoint at different paths still share cache, which is - intended; two *different* checkpoints in identically-named directories also - share it, which is not. Set `TRTLLM_MOONCAKE_STORE_MODEL_KEY` explicitly for - anything long-lived. -- **Stale pages across code changes.** The key namespace does not include a - build hash. After changing anything about page layout or contents, bump - `TRTLLM_MOONCAKE_STORE_PREFIX` or restart the master. -- **UCX warmup requests hit the store too.** `run_benchmark.sh` sends - `2 x ctx_instances x gen_instances` 100-token requests before the real run. - Harmless, but they are in the counters. -- **A partial local match disables the store for that request entirely,** and it - is the dominant failure mode rather than a rare one. The connector offers only - whole blocks and only when the local match is block-aligned, and - `enable_partial_reuse` (default `true`) is exactly what puts the match off a - boundary: measured on M3, it declined **97.2% of lookups**, so a 1.6 TB pool - measured as if it were absent. Turning it off took actual prompt cache read - from 35% to 94%. `py_executor_creator` forces it off for this connector, so - the hazard cannot be hit, but the arithmetic is worth knowing before changing - `tokens_per_block`. See `../mooncake_usage.md` §3. -- **`block_reuse_policy: per_conversation` is off on the context worker** in - these configs. It is not gated, but the connector derives its own - `cache_salt`-seeded hash chain and the interaction is untested. Restore it - only after the store is proven, and treat it as its own experiment. - -## 10. What this does not test - -Worth stating so the results are not oversold: single-node only insofar as the -master is one process (no HA master, no `--root_fs_dir` persistence); no -pipeline or context parallelism (both refused); no VSWA or sliding-window model -(refused); no Eagle3; no shared pool between TensorRT-LLM and vLLM, though the -config schema is deliberately compatible with it. Load bandwidth under -contention from many simultaneous large prefixes is exercised only incidentally -by concurrency, not measured directly. If experiment 2 shows a TTFT regression -at high concurrency despite hits, that is the first thing to profile. diff --git a/mooncake_disagg/ctx_config.yaml b/mooncake_disagg/ctx_config.yaml deleted file mode 100644 index 2cfa4359fc9e..000000000000 --- a/mooncake_disagg/ctx_config.yaml +++ /dev/null @@ -1,30 +0,0 @@ -# Context (prefill) worker: reads AND writes the Mooncake store. -# -# export MOONCAKE_CONFIG_PATH=/abs/path/to/mooncake.json -# CUDA_VISIBLE_DEVICES=0 trtllm-serve \ -# --host localhost --port 8001 --server_role CONTEXT \ -# --config ./ctx_config.yaml - -kv_cache_config: - # The connector describes its pools through register_kv_cache_layout, - # which only KVCacheManagerV2 implements. - use_kv_cache_manager_v2: true - # Local reuse still runs first; the store serves whatever the device missed. - enable_block_reuse: true - # GPU-only tiers are required: a page evicted to host or disk has its GPU - # slot reassigned, which would invalidate the addresses registered with - # the store. - host_cache_size: 0 - disk_cache_size: 0 - free_gpu_memory_fraction: 0.2 - -cache_transceiver_config: - # Requiring KVCacheManagerV2 for the connector forces the Python - # transceiver: the C++ one is bound to the V1 BaseKVCacheManager and - # raises on a V2 manager. NIXL is the only backend the Python - # transceiver supports. - backend: "NIXL" - transceiver_runtime: "PYTHON" - -kv_connector_config: - connector: mooncake-store diff --git a/mooncake_disagg/disagg_config.yaml b/mooncake_disagg/disagg_config.yaml deleted file mode 100644 index ca1f35a7072d..000000000000 --- a/mooncake_disagg/disagg_config.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# Router. Only needs to know where the workers are; each worker's own LLM -# args come from its --config file (ctx_config.yaml / gen_config.yaml). -# -# trtllm-serve disaggregated -c ./disagg_config.yaml - -hostname: localhost -port: 8000 -backend: "pytorch" - -context_servers: - num_instances: 1 - urls: - - "localhost:8001" - -generation_servers: - num_instances: 1 - urls: - - "localhost:8002" diff --git a/mooncake_disagg/gen_config.yaml b/mooncake_disagg/gen_config.yaml deleted file mode 100644 index 6851670e3a78..000000000000 --- a/mooncake_disagg/gen_config.yaml +++ /dev/null @@ -1,22 +0,0 @@ -# Generation (decode) worker: does not touch the Mooncake store at all. -# -# There is deliberately no kv_connector_config here. "Decode none" is the -# absence of a connector, not a StoreRole: StoreRole only has producer, -# consumer and both. Generated tokens are rarely a reused prefix, -# so writing them would cost bandwidth for no hit rate. -# -# CUDA_VISIBLE_DEVICES=1 trtllm-serve \ -# --host localhost --port 8002 --server_role GENERATION \ -# --config ./gen_config.yaml - -kv_cache_config: - # Only to match the context side's transceiver, which must be the Python - # one there. Nothing here is required by the store: this worker never - # opens a store handle. - use_kv_cache_manager_v2: true - enable_block_reuse: true - free_gpu_memory_fraction: 0.8 - -cache_transceiver_config: - backend: "NIXL" - transceiver_runtime: "PYTHON" diff --git a/mooncake_disagg/install_mooncake_runtime.sh b/mooncake_disagg/install_mooncake_runtime.sh deleted file mode 100755 index 8c669ca00f0d..000000000000 --- a/mooncake_disagg/install_mooncake_runtime.sh +++ /dev/null @@ -1,118 +0,0 @@ -#!/bin/bash -# Make the Mooncake Python store bindings importable inside a TensorRT-LLM -# container, so the mooncake-store KV connector can start. Images built by -# docker/common/install_mooncake.sh already have this; run it on images that -# predate that, or to change the wheel. -# -# Two things break a plain `pip install mooncake-transfer-engine` in those -# containers, both explained in docker/common/install_mooncake.sh: the CMake -# source build leaves behind an unusable `mooncake` package that shadows or -# collides with the wheel, and the default wheel is linked against -# libcudart.so.12 while these images ship CUDA 13 only. Either way the symptom -# is `ImportError: libmooncake_store.so` after pip reports success. -# -# `mooncake-transfer-engine-cuda13` needs no CUDA 12 shim, so it is the default -# here. Its releases start at 0.3.9 and so cannot match the pin in -# install_mooncake.sh, which is safe: that CMake-built C++ library backs the -# cache transceiver's Mooncake backend, a different feature, while the -# connector only ever talks to the wheel. The wheel also supplies the -# mooncake_master that lands on PATH, so client and master stay matched. -# Revisit only if cache_transceiver_config.backend is set to MOONCAKE. -# -# Set MOONCAKE_WHEEL to override, for example -# MOONCAKE_WHEEL="mooncake-transfer-engine==0.3.7.post2" -# to match install_mooncake.sh exactly; the libcudart.so.12 shim is then -# applied automatically. -# -# Idempotent and cheap on re-runs: if the install is already correct it exits -# without contacting the network, so it is safe in a SLURM prolog on every node. - -set -euo pipefail - -MOONCAKE_WHEEL="${MOONCAKE_WHEEL:-mooncake-transfer-engine-cuda13==0.3.13}" -WHEEL_NAME="${MOONCAKE_WHEEL%%[=<>]*}" -SITE_PACKAGES="$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" - -echo ">> target wheel: ${MOONCAKE_WHEEL}" - -# Fast path: already correct, so do not touch the network. -if pip3 show "${WHEEL_NAME}" >/dev/null 2>&1 && - python3 -c 'from mooncake.store import MooncakeDistributedStore; MooncakeDistributedStore()' >/dev/null 2>&1; then - echo ">> already installed and importable; nothing to do" - python3 -c 'import mooncake.store; print(" resolved extension:", mooncake.store.__file__)' - exit 0 -fi - -# Purge every `mooncake` package directory on the search path, whatever wrote -# it, since leftovers cannot be told apart from wheel files reliably. -python3 - <<'PY' -import os -import shutil -import sys -import sysconfig - -paths = sysconfig.get_paths() -for entry in list(sys.path) + [paths["purelib"], paths["platlib"]]: - if not entry: - continue - package = os.path.join(entry, "mooncake") - if os.path.isdir(package): - print(f">> removing existing mooncake package: {package}") - shutil.rmtree(package, ignore_errors=True) -PY - -# The two distributions install the same `mooncake` package, so leaving both -# registered produces a half-overwritten directory. -for distribution in mooncake-transfer-engine mooncake-transfer-engine-cuda13; do - if pip3 show "${distribution}" >/dev/null 2>&1; then - echo ">> unregistering ${distribution}" - pip3 uninstall -y -q "${distribution}" >/dev/null 2>&1 || true - fi -done - -pip3 install --no-cache-dir "${MOONCAKE_WHEEL}" - -# Only the CUDA 12 wheel needs the runtime shim. Drop it into the wheel's own -# RPATH directory so the Python extensions and the mooncake_* binaries all find -# it without LD_LIBRARY_PATH being set in each process. -if ldd "${SITE_PACKAGES}"/mooncake/store*.so 2>/dev/null | grep -q "libcudart.so.12 => not found"; then - echo ">> wheel needs libcudart.so.12, which this container lacks; installing it" - pip3 install --no-cache-dir nvidia-cuda-runtime-cu12 - CUDART12="${SITE_PACKAGES}/nvidia/cuda_runtime/lib/libcudart.so.12" - [[ -f "${CUDART12}" ]] || { echo "libcudart.so.12 not found after install" >&2; exit 1; } - mkdir -p "${SITE_PACKAGES}/mooncake_transfer_engine.libs" - ln -sf "${CUDART12}" "${SITE_PACKAGES}/mooncake_transfer_engine.libs/libcudart.so.12" - echo ">> linked libcudart.so.12 into the wheel's RPATH directory" -fi - -# Verify, because every failure above stays silent until a worker starts. -echo ">> verifying" -python3 - <<'PY' -import mooncake.store -from mooncake.store import MooncakeDistributedStore - -MooncakeDistributedStore() -print(" mooncake.store imports and instantiates: OK") -print(f" resolved extension: {mooncake.store.__file__}") -PY - -# Every extension module and the master binary must have a resolvable link line. -# ldd the real ELF files, not /usr/local/bin/mooncake_master, which is a Python -# console script. -unresolved=0 -for elf in "${SITE_PACKAGES}"/mooncake/*.so "${SITE_PACKAGES}"/mooncake/mooncake_master; do - [[ -e "${elf}" ]] || continue - if missing="$(ldd "${elf}" 2>&1 | grep 'not found')"; then - echo " $(basename "${elf}"): unresolved -> ${missing}" >&2 - unresolved=1 - fi -done -[[ "${unresolved}" -eq 0 ]] || { echo "unresolved shared libraries; see above" >&2; exit 1; } -echo " all mooncake ELF link lines resolve: OK" - -for entry in mooncake_master mooncake_http_metadata_server; do - path="$(command -v "${entry}")" || { echo "${entry} is not on PATH" >&2; exit 1; } - echo " ${entry}: OK (${path})" -done - -echo ">> done" diff --git a/mooncake_disagg/m3_agg_mooncake.yaml b/mooncake_disagg/m3_agg_mooncake.yaml deleted file mode 100644 index 39610cc63899..000000000000 --- a/mooncake_disagg/m3_agg_mooncake.yaml +++ /dev/null @@ -1,109 +0,0 @@ -# MiniMax-M3-NVFP4, aggregated, with the mooncake-store KV connector. -# -# Derived from the production M3 serving config. Every deviation from it is -# marked CONNECTOR: with the gate that forces it. Startup gates live in -# py_executor_creator.py (~line 830), py_executor._maybe_init_kv_connector_manager, -# and connectors/mooncake_store/validation.py. - -max_seq_len: 1048576 -max_num_tokens: 16384 -max_batch_size: 20 - -cuda_graph_config: - enable_padding: true - batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20] - -torch_compile_config: - enable_fullgraph: true - enable_inductor: false - enable_piecewise_cuda_graph: true - capture_num_tokens: [1, 16, 128, 256, 512, 768, 1024, 1280, 1536, 1792, 2048] - enable_userbuffers: true - max_num_streams: 3 - -moe_config: - backend: TRTLLM - use_low_precision_moe_combine: true - -sparse_attention_config: - algorithm: minimax_m3 - implementation: msa - indexer_kv_dtype: fp8 - # Required by the connector: index-V is a plain tensor outside the paged - # pools, so it is never described or transferred. A replayed prefix would - # pair stored index-K with stale index-V. - sparse_disable_index_value: true - fuse_qkv_index_projection: true - -kv_cache_config: - free_gpu_memory_fraction: 0.94 - enable_block_reuse: true - tokens_per_block: 128 - use_kv_cache_manager_v2: true - dtype: fp8 - event_buffer_max_size: 0 - - # CONNECTOR: was host_cache_size: 388554555392. - # _reject_non_gpu_cache_tiers rejects every tier below GPU, because a - # registered region is only a valid device address while its page stays - # pinned to GPU, and eviction reassigns the slot. Both must be an explicit - # 0: V2 provisions a host tier when the field is left at its default of - # None, which is falsy but still yields a tier. - host_cache_size: 0 - disk_cache_size: 0 - - # CONNECTOR: was block_reuse_policy: per_conversation. - # Not gated, but the connector derives its own blake2b hash chain seeded - # by cache_salt, so reuse-policy interactions are untested. Restore this - # after the store is proven. - # block_reuse_policy: per_conversation - -# CONNECTOR: was MAX_UTILIZATION. py_executor_creator raises -# "KV connector is only supported with guaranteed no evict scheduler policy." -scheduler_config: - capacity_scheduler_policy: GUARANTEED_NO_EVICT - -# CONNECTOR: Eagle3 disabled for the first validation run. -# Not gated, but MiniMaxM3KVCacheManagerV2 sets supports_shared_draft_layers, -# so draft layers join the unified V2 cache and therefore the registered -# layout. That adds page geometry the store must key correctly, on a path -# with no coverage. Re-enable once cold/warm passes without it. -# speculative_config: -# decoding_type: Eagle3 -# max_draft_len: 3 -# speculative_model: /path/to/eagle3/draft - -enable_chunked_prefill: true -enable_autotuner: true -trust_remote_code: true -reasoning_parser: minimax_m3 -stream_interval: 20 -print_iter_log: true -num_postprocess_workers: 8 - -# Required: dummy requests inserted for cross-DP balancing flow through the -# connector hooks and are indistinguishable from real requests. -enable_attention_dp: false - -kv_connector_config: - connector: mooncake-store - - # Aggregated means one engine, which is the only shape a server-owned - # master is right for: it dies with the server, so nothing else can be - # sharing the pool and nothing can expect it to survive a restart. Point - # master_server_address at a master of its own for either of those. - # - # An inherited MOONCAKE_CONFIG_PATH wins over this block, so the SLURM - # harness keeps naming its own pool. - mooncake_store: - launch_master: true - # TCP removes RDMA from the variable list at the cost of any - # performance conclusion. Set protocol: rdma with a device_name from - # ibv_devinfo for a run worth quoting. - protocol: tcp - device_name: "" - # Contributed per worker process, so the pool is this times the world - # size. Sized for the working set, not for what the node can spare: - # by ~85-90% full the master evicts and the hit rate follows it down. - global_segment_size: 160GiB - local_buffer_size: 4GiB diff --git a/mooncake_disagg/m3_ctx_mooncake.yaml b/mooncake_disagg/m3_ctx_mooncake.yaml deleted file mode 100644 index 1a2c1b20ece5..000000000000 --- a/mooncake_disagg/m3_ctx_mooncake.yaml +++ /dev/null @@ -1,96 +0,0 @@ -# MiniMax-M3-NVFP4 CONTEXT (prefill) worker: reads AND writes the store. -# -# Same as m3_agg_mooncake.yaml plus the transceiver. Every CONNECTOR: note -# there applies here too. - -max_seq_len: 1048576 -max_num_tokens: 16384 -max_batch_size: 20 - -cuda_graph_config: - enable_padding: true - batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20] - -torch_compile_config: - enable_fullgraph: true - enable_inductor: false - enable_piecewise_cuda_graph: true - capture_num_tokens: [1, 16, 128, 256, 512, 768, 1024, 1280, 1536, 1792, 2048] - enable_userbuffers: true - max_num_streams: 3 - -moe_config: - backend: TRTLLM - use_low_precision_moe_combine: true - -sparse_attention_config: - algorithm: minimax_m3 - implementation: msa - indexer_kv_dtype: fp8 - sparse_disable_index_value: true - fuse_qkv_index_projection: true - -kv_cache_config: - free_gpu_memory_fraction: 0.94 - enable_block_reuse: true - tokens_per_block: 128 - use_kv_cache_manager_v2: true - dtype: fp8 - event_buffer_max_size: 0 - # CONNECTOR: GPU-only tiers. Mooncake takes over the CPU-offload role. - host_cache_size: 0 - disk_cache_size: 0 - -# MANDATORY, and not only because of the connector. M3 sets -# sparse_attention_config, so get_kv_cache_manager_cls routes to -# MiniMaxM3KVCacheManagerV2 unconditionally, and use_kv_cache_manager_v2 is not -# even consulted on that branch. KVCacheManagerV2 cannot drive the C++ -# transceiver, and M3 does not override get_preferred_transceiver_runtime, so -# 'auto' would resolve to C++ and be rejected. -cache_transceiver_config: - backend: "NIXL" - transceiver_runtime: "PYTHON" - -# CONNECTOR: was MAX_UTILIZATION; connectors require guaranteed-no-evict. -scheduler_config: - capacity_scheduler_policy: GUARANTEED_NO_EVICT - -enable_chunked_prefill: true -enable_autotuner: true -trust_remote_code: true -reasoning_parser: minimax_m3 -stream_interval: 20 -print_iter_log: true -num_postprocess_workers: 8 -enable_attention_dp: false - -kv_connector_config: - connector: mooncake-store - # The pool, provisioned by this server during its own bringup: it starts - # the master, renders the client config and exports MOONCAKE_CONFIG_PATH - # before the ranks that open store handles exist. Nothing outside - # trtllm-serve prepares any of it. - # - # Only right because there is one context server. Two that each launch a - # master get two disjoint pools, which is what 'trtllm-serve - # mooncake_master' plus master_server_address: file:// is for. - mooncake_store: - launch_master: true - # Where the master's address is published, for the generation - # workers to lend the pool memory. One also goes to the run - # directory regardless; under the SLURM harness that is the job's - # log directory, and __LOG_DIR__ is substituted there. - master_address_file: __LOG_DIR__/master.addr - protocol: rdma - # device_name omitted on purpose: the fastest active HCAs on this - # node are detected, which keeps one config portable across node - # types. Name them to pin it. - # - # Charged per rank, so a TP=2 server contributes twice this. - global_segment_size: 160GiB - # Registering the KV pools with the HCA needs nvidia_peermem; without - # it registration fails on every range, and staging registers only - # host memory instead. 1GiB holds a full transfer_batch_size of pages; - # less silently reduces the batch rather than failing. - stage_through_host: true - staging_buffer_bytes: 1GiB diff --git a/mooncake_disagg/m3_disagg_config.yaml b/mooncake_disagg/m3_disagg_config.yaml deleted file mode 100644 index 4e51aa45e154..000000000000 --- a/mooncake_disagg/m3_disagg_config.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# Router for the M3 mooncake-store validation. Worker LLM args come from each -# worker's own --config file, not from here. -# -# 1 ctx + 1 gen at TP=4 fits 8 GPUs. To demonstrate cross-instance reuse -# without restarting anything, raise context_servers.num_instances to 2 and -# add a second URL (needs 12 GPUs at TP=4). - -hostname: localhost -port: 8000 -backend: "pytorch" - -context_servers: - num_instances: 1 - urls: - - "localhost:8001" - -generation_servers: - num_instances: 1 - urls: - - "localhost:8002" diff --git a/mooncake_disagg/m3_gen_mooncake.yaml b/mooncake_disagg/m3_gen_mooncake.yaml deleted file mode 100644 index e666099674b6..000000000000 --- a/mooncake_disagg/m3_gen_mooncake.yaml +++ /dev/null @@ -1,84 +0,0 @@ -# MiniMax-M3-NVFP4 GENERATION (decode) worker: does not touch the store. -# -# There is deliberately no kv_connector_config. That absence is the whole of -# "decode-none": StoreRole has only producer, consumer and both, so there is -# no role value meaning "off". -# -# Because no connector runs here, three of the context worker's constraints -# do NOT apply, and this file keeps the production values instead. -# -# It does lend the pool this node's host memory. Capacity comes only from -# processes that open a store handle, so without this every byte of the pool -# would be prefill-node DRAM caching prefill's own GPUs, which is what the -# native host tier already does. Lending is not a connector: this engine still -# never reads or writes the store, and keeps its cache transceiver for the -# prefill-to-decode handoff. -mooncake_donation: - # The master the context server started and published. A server lending - # memory waits for it, so the two can start in any order. - master_server_address: file://__LOG_DIR__/master.addr - # Charged per server process, not per rank as global_segment_size is. Two - # servers on one node lend twice this, and it competes with this node's own - # kv_cache_config.host_cache_size, so size the two together. - segment_size: 640GiB - protocol: rdma - -max_seq_len: 1048576 -max_num_tokens: 16384 -max_batch_size: 20 - -cuda_graph_config: - enable_padding: true - batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20] - -torch_compile_config: - enable_fullgraph: true - enable_inductor: false - enable_piecewise_cuda_graph: true - capture_num_tokens: [1, 16, 128, 256, 512, 768, 1024, 1280, 1536, 1792, 2048] - enable_userbuffers: true - max_num_streams: 3 - -moe_config: - backend: TRTLLM - use_low_precision_moe_combine: true - -sparse_attention_config: - algorithm: minimax_m3 - implementation: msa - indexer_kv_dtype: fp8 - sparse_disable_index_value: true - fuse_qkv_index_projection: true - -kv_cache_config: - free_gpu_memory_fraction: 0.94 - enable_block_reuse: true - block_reuse_policy: per_conversation - tokens_per_block: 128 - use_kv_cache_manager_v2: true - dtype: fp8 - event_buffer_max_size: 0 - # Kept: no connector here, so _reject_non_gpu_cache_tiers never runs and - # decode keeps its host tier. Drop to 0 first if the Python transceiver - # misbehaves, since that pairing is the less-travelled path. - host_cache_size: 388554555392 - -# Must match the context worker: both ends of the handoff run the same -# transceiver, and M3's always-V2 manager rules out the C++ one. -cache_transceiver_config: - backend: "NIXL" - transceiver_runtime: "PYTHON" - -# Kept: the guaranteed-no-evict requirement is a connector gate, and no -# connector runs on this worker. -scheduler_config: - capacity_scheduler_policy: MAX_UTILIZATION - -enable_chunked_prefill: true -enable_autotuner: true -trust_remote_code: true -reasoning_parser: minimax_m3 -stream_interval: 20 -print_iter_log: true -num_postprocess_workers: 8 -enable_attention_dp: false diff --git a/mooncake_disagg/mooncake.json b/mooncake_disagg/mooncake.json deleted file mode 100644 index 2a17f7d2d760..000000000000 --- a/mooncake_disagg/mooncake.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "metadata_server": "P2PHANDSHAKE", - "master_server_address": "127.0.0.1:50051", - "protocol": "tcp", - "device_name": "", - "global_segment_size": "16GiB", - "local_buffer_size": "1GiB", - "role": "both", - "cache_prefix": "trtllm", - "transfer_batch_size": 64 -} diff --git a/mooncake_disagg/mooncake_api_surface_test.py b/mooncake_disagg/mooncake_api_surface_test.py deleted file mode 100644 index 4c0129f6243f..000000000000 --- a/mooncake_disagg/mooncake_api_surface_test.py +++ /dev/null @@ -1,115 +0,0 @@ -#!/usr/bin/env python3 -"""Exercise every MooncakeDistributedStore method the connector calls. - -`mooncake_smoke_test.py` only proves the install loads and can round-trip a byte -string. The connector's hot path never uses `put` or `get`: it registers the KV -pools and then moves pages with the `batch_*_multi_buffers` zero-copy calls. -Those signatures could drift between wheel versions, so this checks them against -real registered GPU memory, in the same order `worker.py` uses them. - -Needs a running mooncake_master and MOONCAKE_CONFIG_PATH, same as the connector. -""" - -import json -import os -import socket -import sys - -import torch - -CONFIG_PATH = os.environ.get("MOONCAKE_CONFIG_PATH") -if not CONFIG_PATH: - sys.exit("Set MOONCAKE_CONFIG_PATH to the Mooncake JSON config first.") - -with open(CONFIG_PATH) as handle: - cfg = json.load(handle) - - -def parse_size(value): - if isinstance(value, int): - return value - units = {"KiB": 1 << 10, "MiB": 1 << 20, "GiB": 1 << 30, "TiB": 1 << 40} - for suffix, scale in units.items(): - if value.endswith(suffix): - return int(float(value[: -len(suffix)]) * scale) - return int(value) - - -import mooncake # noqa: E402 -from mooncake.store import MooncakeDistributedStore # noqa: E402 - -print(f"mooncake package: {mooncake.__path__[0]}") - -store = MooncakeDistributedStore() -hostname = cfg.get("local_hostname") or socket.gethostbyname(socket.gethostname()) -status = store.setup( - hostname, - cfg["metadata_server"], - parse_size(cfg.get("global_segment_size", "1GiB")), - parse_size(cfg.get("local_buffer_size", "256MiB")), - cfg.get("protocol", "tcp"), - cfg.get("device_name", ""), - cfg["master_server_address"], -) -assert status == 0, f"setup failed with status {status}" -print("setup: OK") - -# Stand in for a KV pool. PageAddressing.page_buffers returns one address per -# layer-group region, so a page is scattered across REGIONS buffers rather than -# contiguous, which is why the batch calls take list[list[int]]. Two strided -# regions here so the scatter-gather path is actually exercised. -PAGES = 8 -REGIONS = 2 -REGION_BYTES = 128 * 1024 -STRIDE = REGION_BYTES # slots within a region are strided, as in the real layout -pool = torch.empty(REGIONS * PAGES * STRIDE, dtype=torch.uint8, device="cuda") -region_bases = [pool.data_ptr() + r * PAGES * STRIDE for r in range(REGIONS)] - -status = store.register_buffer(pool.data_ptr(), pool.numel()) -assert status == 0, f"register_buffer failed with status {status}" -print(f"register_buffer: OK ({pool.numel()} bytes of GPU memory at {pool.data_ptr():#x})") - -prefix = cfg.get("cache_prefix", "trtllm") -keys = [f"{prefix}/api-surface/page{i}" for i in range(PAGES)] -addresses = [[base + i * STRIDE for base in region_bases] for i in range(PAGES)] -sizes = [[REGION_BYTES] * REGIONS for _ in range(PAGES)] - -# Distinct content per (page, region), so a mixed-up address or size cannot pass. -view = pool.view(REGIONS, PAGES, STRIDE) -for r in range(REGIONS): - for i in range(PAGES): - view[r, i, :REGION_BYTES] = (i * 31 + r * 97 + 7) % 256 -expected = pool.clone() - -present = store.batch_is_exist(keys) -assert len(present) == PAGES, f"batch_is_exist returned {len(present)} of {PAGES}" -assert all(status != 1 for status in present), f"keys already present: {present}" -print(f"batch_is_exist (absent): OK {list(present)}") - -results = store.batch_put_from_multi_buffers(keys, addresses, sizes) -assert len(results) == PAGES, f"batch_put returned {len(results)} of {PAGES}" -bad = [(k, r) for k, r in zip(keys, results) if not isinstance(r, int) or r < 0] -assert not bad, f"batch_put_from_multi_buffers failed: {bad}" -print(f"batch_put_from_multi_buffers: OK {list(results)}") - -present = store.batch_is_exist(keys) -assert all(status == 1 for status in present), f"keys missing after put: {present}" -print("batch_is_exist (present): OK") - -pool.zero_() -results = store.batch_get_into_multi_buffers(keys, addresses, sizes) -assert len(results) == PAGES, f"batch_get returned {len(results)} of {PAGES}" -bad = [(k, r) for k, r in zip(keys, results) if not isinstance(r, int) or r < 0] -assert not bad, f"batch_get_into_multi_buffers failed: {bad}" -print(f"batch_get_into_multi_buffers: OK {list(results)}") - -torch.cuda.synchronize() -assert torch.equal(pool, expected), "page contents differ after the round trip" -print("GPU page contents byte-for-byte identical: OK") - -for key in keys: - store.remove(key) -store.close() -print("remove + close: OK") - -print("\nPASS: the full connector API surface works on this install.") diff --git a/mooncake_disagg/mooncake_smoke_test.py b/mooncake_disagg/mooncake_smoke_test.py deleted file mode 100644 index db0b1dc8dee9..000000000000 --- a/mooncake_disagg/mooncake_smoke_test.py +++ /dev/null @@ -1,67 +0,0 @@ -#!/usr/bin/env python3 -"""Prove that a Mooncake install can actually serve the mooncake-store connector. - -Mirrors the `store.setup` call in -`tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py` and then -does one round trip, so a pass here means the connector's own startup path will -work. Reads the same `MOONCAKE_CONFIG_PATH` file the connector reads. -""" - -import json -import os -import socket -import sys - -CONFIG_PATH = os.environ.get("MOONCAKE_CONFIG_PATH") -if not CONFIG_PATH: - sys.exit("Set MOONCAKE_CONFIG_PATH to the Mooncake JSON config first.") - -with open(CONFIG_PATH) as handle: - cfg = json.load(handle) - - -def parse_size(value): - if isinstance(value, int): - return value - units = {"KiB": 1 << 10, "MiB": 1 << 20, "GiB": 1 << 30, "TiB": 1 << 40} - for suffix, scale in units.items(): - if value.endswith(suffix): - return int(float(value[: -len(suffix)]) * scale) - return int(value) - - -from mooncake.store import MooncakeDistributedStore # noqa: E402 - -print("import mooncake.store: OK") - -store = MooncakeDistributedStore() -hostname = cfg.get("local_hostname") or socket.gethostbyname(socket.gethostname()) -status = store.setup( - hostname, - cfg["metadata_server"], - parse_size(cfg.get("global_segment_size", "1GiB")), - parse_size(cfg.get("local_buffer_size", "256MiB")), - cfg.get("protocol", "tcp"), - cfg.get("device_name", ""), - cfg["master_server_address"], -) -if status != 0: - sys.exit(f"store.setup failed with status {status}") -print(f"store.setup: OK (host={hostname}, master={cfg['master_server_address']})") - -key = f"{cfg.get('cache_prefix', 'trtllm')}/smoke-test" -payload = bytes(range(256)) * 4096 # 1 MiB, non-trivial content - -assert store.put(key, payload) == 0, "put failed" -print(f"put {len(payload)} bytes: OK") - -assert store.is_exist(key) == 1, "is_exist did not report the key" -print("is_exist: OK") - -got = store.get(key) -assert got == payload, f"round trip mismatch: got {len(got)} bytes" -print("get + byte-for-byte compare: OK") - -store.remove(key) -print("remove: OK") -print("\nPASS: this install can back the mooncake-store connector.") diff --git a/mooncake_usage.md b/mooncake_usage.md deleted file mode 100644 index f2a4e8d84286..000000000000 --- a/mooncake_usage.md +++ /dev/null @@ -1,269 +0,0 @@ -# Using the mooncake-store KV connector - -The `mooncake-store` connector publishes KV cache pages into a shared, -content-addressed pool in host DRAM, so a prefix computed by one engine can be -replayed by another. It is a KV cache *connector*, unrelated to the Mooncake -*transfer engine* that the cache transceiver can use for prefill/decode handoff: -a different component with a different config, and the two are rarely both in -play. - -Use it when local block reuse leaves reuse on the table: several context -instances that see the same prefixes, prefixes that should outlive a restart, or -a working set larger than one node's host memory. It replaces TensorRT-LLM's -native host offload tier rather than layering on top of it (`host_cache_size` -must be `0`). - -This page is the entry point. Depth lives elsewhere: - -| For | Read | -|---|---| -| API surface, gates, keying | `docs/source/features/kv-cache-connector.md` § *Mooncake distributed store* | -| Full SLURM runbook, install debugging, experiment matrix | `mooncake_disagg/README.md` | -| Working configs | `mooncake_disagg/m3_ctx_mooncake.yaml`, `m3_gen_mooncake.yaml` | -| Correctness tests (no GPU, no store) | `tests/unittest/_torch/executor/test_mooncake_store_connector.py` | - -## 1. Install - -The connector needs the Mooncake **Python** bindings -(`mooncake.store.MooncakeDistributedStore`). Containers have shipped the C++ -library for months without them, and a CMake-installed `mooncake` package -shadows the working wheel, so a bare `pip install` reports success and the -import still fails. Inside the container: - -```bash -bash mooncake_disagg/install_mooncake_runtime.sh # idempotent, ~8s warm -python3 -c "from mooncake.store import MooncakeDistributedStore; print('ok')" -``` - -`disaggr_torch.slurm` runs this per node automatically when a worker config -mentions `mooncake-store`, and images built from this repo have it baked in. -`mooncake_disagg/README.md` §2 explains why it is this awkward. - -## 2. Configure - -A `mooncake_master` process must be reachable, and every worker needs -`MOONCAKE_CONFIG_PATH` pointing at a JSON client config naming it. - -Three ways to get there, in increasing order of how much you have to arrange: - -| Deployment | Master | -|---|---| -| One `trtllm-serve`, own pool, **including the SLURM harness** | `mooncake_store: {launch_master: true}`, so the server starts it | -| Several engines, or a pool that outlives them | `trtllm-serve mooncake_master --address_file P`, then `mooncake_store: {master_server_address: file://P}` | -| An externally provisioned pool | Nothing in the config: an inherited `MOONCAKE_CONFIG_PATH` wins over `mooncake_store` and says so in the log | - -The first two make `trtllm-serve` render the client config and export -`MOONCAKE_CONFIG_PATH` itself. Nothing outside `trtllm-serve` starts a master, -writes a JSON config or picks an HCA, the SLURM harness included: all -`disaggr_torch.slurm` does is install the bindings and tell the configs which -directory the run is in. `mooncake_disagg/README.md` §4 covers running a master -as its own SLURM job, for the second row. - -One thing a scheduler-launched server does need: `TRTLLM_MOONCAKE_RUN_DIR` -pointing somewhere all its ranks can read. Provisioning happens in the server -process and reaches the ranks it spawns through the environment, but under -`trtllm-llmapi-launch` each rank is its own task and was already running, so -those ranks read the rendered config back from that directory instead. Without -it they fail during bringup naming `MOONCAKE_CONFIG_PATH`. `start_worker.sh` -sets it to the job's log directory, which is also where the master's log and -published address land. - -A launched master dies with the server, so use it only for a single engine: -two context servers that each launch one get two disjoint pools, and the -survival-across-restart case is impossible by construction. Those cases want -row two, where the master is its own command and nothing else's lifetime -bounds it. - -`master_server_address` takes a plain `host:port` or `file://`. The file -is what a scheduler-placed master needs: its host is not known when the configs -are written, `--address_file` publishes it once the master answers, and a -server reading it waits for the master to exist. Nobody has to write an address -down, and a stale one cannot be dialed because the file is removed on exit. - -Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated JSON and the master's log, -which otherwise sit in a temporary directory that shutdown removes. -`TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) bounds the wait for the port; -that wait is also what turns an unreachable external master from a failure in -every rank after the model loads into one line before it starts. - -Put the connector on the **context** workers only: - -```yaml -kv_connector_config: - connector: mooncake-store - mooncake_store: # omit when an orchestrator sets - launch_master: true # MOONCAKE_CONFIG_PATH for you - protocol: tcp # rdma with a device_name for real numbers -kv_cache_config: - use_kv_cache_manager_v2: true # required: only V2 describes its pools - enable_block_reuse: true - host_cache_size: 0 # required, and must be explicit, not omitted - disk_cache_size: 0 - tokens_per_block: 128 -scheduler_config: - capacity_scheduler_policy: GUARANTEED_NO_EVICT # required -enable_attention_dp: false # required -``` - -Generation workers deliberately get no `kv_connector_config`: generated tokens -are rarely a reused prefix, and leaving it off is the only way to say "off" -(`StoreRole` has no off value). It also lets them keep their host cache tier and -`MAX_UTILIZATION` scheduler, both of which the connector forbids. - -Per-process environment, on the workers that open a handle: - -| Variable | Purpose | -|---|---| -| `TRTLLM_MOONCAKE_STORE_ROLE` | `producer` / `consumer` / `both` | -| `TRTLLM_MOONCAKE_STORE_PREFIX` | Cache namespace. Bump it after any change to page layout or contents. | -| `TRTLLM_MOONCAKE_STORE_MODEL_KEY` | Defaults to the checkpoint directory's basename. Set it explicitly for anything long-lived. | - -Pool capacity comes only from processes that open a store handle, so a -prefill-only connector gives a prefill-only pool, caching prefill's GPUs in -prefill's own DRAM and largely duplicating the native host offload. Ask the -generation servers to lend their memory and the pool spans both sides while -their engines stay connector-free: - -```yaml -# generation worker: no connector, memory only -mooncake_donation: - master_server_address: file:///$WORK_DIR/master.addr - segment_size: 320GiB # per server process, not per rank - protocol: rdma - device_name: mlx5_1 -``` - -The context worker publishes the address this reads by adding -`master_address_file: $WORK_DIR/master.addr` next to its `launch_master: true`; -one is written to the run directory regardless. Startup order stops mattering, -because a server lending memory waits for the master rather than needing to -follow it. - -Note the granularity: `global_segment_size` is charged per *rank*, -`segment_size` per *server process*. Two generation servers on one node lend -`segment_size` each. It is charged to the process, so it competes with that -node's own `kv_cache_config.host_cache_size`. Size the two together. - -A node running no server can lend as its own command, which is also how the -pool gets memory from a machine with no GPUs at all: - -```bash -trtllm-serve mooncake_donor --master_server_address file://$WORK_DIR/master.addr \ - --segment_size 160GiB --protocol rdma --device_name mlx5_0 -``` - -## 3. Partial reuse is forced off - -This is the one setting that decides whether the feature works at all. - -The store is addressed by whole blocks. The connector is handed the device match -as `num_computed_tokens` and offers only blocks beyond it, but it can continue -only from a block boundary, so when the device match lands mid-block it declines -the lookup entirely. `enable_partial_reuse=true` is precisely what makes the -device match land mid-block, so it trades part of one block of device reuse for -*every* stored block of the remaining prefix. On MiniMax-M3 that declined 97.2% -of lookups, leaving a 1.6 TB pool measuring as if it were not there. - -`py_executor_creator` therefore forces `enable_partial_reuse=false` whenever -this connector is configured, and says so: - -``` -Disabling partial reuse: it is not usable with the mooncake-store connector... -``` - -There is nothing to set. The warning fires even from the default of `true`, and -is the confirmation that the coercion ran. Configs that already set `false` are -unaffected. - -## 4. Verify a run - -Startup, at INFO, on every context worker: - -```bash -grep -h "mooncake-store" /3_output_CTX_*.log | head -40 -``` - -`registered layout: ... bytes/page=...` is the line to keep, since pool sizing -depends on it, and `window=None` confirms no sliding-window group (one would -have aborted startup). - -Then check that the pool spans the hosts you expect. -`disaggr_torch.slurm` writes the per-segment breakdown to -`/9_mooncake_summary.log`; a single host means a prefill-only pool. -Pool occupancy and eviction come from the master's own log, -`$TRTLLM_MOONCAKE_RUN_DIR/mooncake_master.log`, which under the harness is -`/mooncake_master.log`. The startup line reports the path either way. - -**Which reuse number counts store hits:** per-request stats -(`reused_blocks_per_request`, `kv_cache_hit_rate_per_request`) **do**; -`/prometheus/metrics` iteration counters (`kv_cache_iter_reused_blocks`) **do -not**, since those come from the local reuse tree. Store hits are therefore -roughly per-request reuse minus local-tree reuse. All of this needs -`enable_iter_perf_stats`, `enable_iter_req_stats` and `return_perf_metrics`, -which all default to false. - -## 5. What it measured - -MiniMax-M3-NVFP4 on GB300, 1 context server (TP=2) + 2 generation servers -(TP=4), connector on context only, ~1.6 TB pool (160 GiB per context rank plus -640 GiB donated per generation node), real conversation trace. - -"Baseline" is the same configuration with partial reuse left at its default. - -| Run | Theoretical hit | Actual hit | Output tok/s | -|---|---|---|---| -| c50 baseline | 96.29% | 35.19% | 319.43 | -| **c50 with partial reuse off** | 96.64% | **93.53%** | **697.29** (2.18x) | -| c70 baseline | 96.01% | 38.44% | 397.43 | -| **c70 with partial reuse off** | 96.51% | **86.46%** | **643.57** (1.62x) | - -A 61-point gap between the reuse the workload allowed and the reuse the system -achieved closed to 3 points. For comparison, native host offload on the same -workload reached 35.59% actual hit at 318.14 tok/s: it wrote 1.83 TB to host and -read 32.5 GB back, behaving as a write-only tier. - -Where the reuse comes from, in steady state (attribution counters, c50): -**~95% of all reuse is served by the pool and ~5% by the device cache.** The -residual 3-5% of misses are prefixes never written by anyone, which no store can -serve. Blocks stranded behind a contiguity gap measured exactly zero, as did -unattributed blocks. - -**Peak is at c50, not higher.** By c70 the pool runs 85-90% full with active -eviction, hit rate falls to 86% and throughput with it. Concurrency headroom is -a function of pool size; size the pool for the working set rather than assuming -the c50 result scales. - -## 6. Things that will bite - -- **`host_cache_size: 0` must be written explicitly.** Left at its `None` - default, V2 still provisions a host tier and startup is rejected. Falsy is not - the same as absent. -- **The key namespace pins world size, rank, `tokens_per_block`, layer groups - and `bytes_per_page`.** Change tensor parallelism and every stored page - becomes unreachable, as a miss rather than an error. -- **No build hash in the key.** After changing page layout or contents, bump - `TRTLLM_MOONCAKE_STORE_PREFIX` or restart the master. -- **Loads are synchronous** (`start_load_kv`, before the forward pass), so every - loaded byte is exposed to TTFT. A store hit wins only when it displaces real - prefill. -- **`mooncake-store failed to load N of M pages` is not flaky.** The runtime had - already counted those tokens as computed; it is the tripwire against a wrong - answer. Stop and investigate. -- **Rejected outright at startup:** pipeline or context parallelism, - sliding-window attention, attention DP, beam search, host/disk cache tiers, - `MAX_UTILIZATION`, and M3's index-V cache unless - `sparse_disable_index_value=true`. `mooncake_disagg/README.md` §9 has the rest. - -## 7. Diagnostics - -The `user/brb/m3-mooncake-store-instrumentation` branch carries attribution -counters that sort every reusable block of every prompt into served-from-device, -served-from-pool, stranded behind a contiguity gap, never written, evicted, or -unattributed, with cumulative and per-window reporting. That is what produced -§5's split. They are kept off this branch because they cost a store probe on -lookups the connector would otherwise decline. - -Enable with `MOONCAKE_DEBUG_COVERAGE=1`, set via -`environment.ctx_worker_env_var`: `slurm.extra_args` reaches the harness rather -than the worker processes, so a flag the connector reads has to go where the -worker environment is built. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index 94c4d8d7a928..d8f8d38d40b3 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -378,7 +378,7 @@ def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: raise FileNotFoundError( f"{binary!r} is not on PATH, so launch_master cannot start a " "Mooncake master. It ships with the Mooncake runtime, which " - "mooncake_disagg/install_mooncake_runtime.sh installs. Point " + "docker/common/install_mooncake.sh installs. Point " f"{MASTER_BINARY_ENV} at the binary, or drop launch_master and " "set master_server_address to a master you run yourself." ) From 2ba3bcaf32a16e1b69cc6ce066dd6fc53bfd6c58 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:07:33 -0700 Subject: [PATCH 19/24] [None][chore] Drop the benchmark job-watching helper watch_job.sh polled a log directory for local experiment monitoring; nothing in the harness invokes it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../slurm/benchmark/watch_job.sh | 105 ------------------ 1 file changed, 105 deletions(-) delete mode 100644 examples/disaggregated/slurm/benchmark/watch_job.sh diff --git a/examples/disaggregated/slurm/benchmark/watch_job.sh b/examples/disaggregated/slurm/benchmark/watch_job.sh deleted file mode 100644 index 4b27a2322f14..000000000000 --- a/examples/disaggregated/slurm/benchmark/watch_job.sh +++ /dev/null @@ -1,105 +0,0 @@ -#!/bin/bash -# Poll a disaggregated benchmark log directory and print one EVENT line per -# state change. Intended for tracking a running job without tailing megabytes -# of worker log. -# -# bash watch_job.sh [poll_seconds] [max_minutes] -set -uo pipefail - -log_dir="${1:?usage: watch_job.sh [poll_seconds] [max_minutes]}" -poll="${2:-20}" -max_minutes="${3:-60}" -deadline=$(( SECONDS + max_minutes * 60 )) - -declare -A seen - -announce() { - key="$1"; shift - if [ -z "${seen[$key]:-}" ]; then - seen[$key]=1 - echo "EVENT: $*" - fi -} - -count_matches() { - grep -h "$1" "${log_dir}"/3_output_CTX_*.log 2>/dev/null | wc -l || true -} - -# How much of the pool's contents lives on each node. A segment is one client -# process's donated memory, so a single host here means the pool is -# prefill-only: the donors are absent or not being allocated into. -placement() { - grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" \ - "${log_dir}/2_mooncake_master.log" 2>/dev/null \ - | awk '{sub(/size=/,"",$2); sub(/segment=/,"",$3); split($3,p,":"); - n[p[1]]++; b[p[1]]+=$2; t+=$2} - END {if (t == 0) {print "(none yet)"; exit} - for (h in n) printf "%s:%d pages/%.2fGiB/%.0f%% ", h, n[h], b[h]/1073741824, 100*b[h]/t}' -} - -while [ ${SECONDS} -lt ${deadline} ]; do - for ready in "${log_dir}"/mooncake_donor_*.ready; do - [ -s "${ready}" ] || continue - node="$(basename "${ready}" .ready)"; node="${node#mooncake_donor_}" - announce "donor_${node}" "memory donor on ${node} mounted its segment:" \ - "$(cat "${ready}")" - done - - for role in CTX GEN; do - f="${log_dir}/3_output_${role}_0.log" - [ -f "$f" ] || continue - grep -qs "Server started at\|Application startup complete\|Uvicorn running" "$f" \ - && announce "${role}_up" "${role} worker serving" - done - - if grep -qs "registered layout" "${log_dir}"/3_output_CTX_*.log; then - announce "registered" "connector registered KV layout:" \ - "$(grep -h "registered layout" "${log_dir}"/3_output_CTX_*.log 2>/dev/null | head -n 1 | cut -c1-400)" - fi - - matched=$(count_matches "mooncake-store matched") - loaded=$(count_matches "mooncake-store rank") - [ "${matched:-0}" -gt 0 ] && announce "first_match" "first store hit; matched lines=${matched}" - [ "${loaded:-0}" -gt 0 ] && announce "first_load" "first store load; loaded lines=${loaded}" - - # A second host is the first moment prefill-written KV is provably on - # decode DRAM, which is what the donors exist for. - hosts=$(grep -o "segment=[0-9.]*:" "${log_dir}/2_mooncake_master.log" 2>/dev/null \ - | sort -u | wc -l || true) - [ "${hosts:-0}" -ge 2 ] && announce "multi_host" \ - "pool spans ${hosts} hosts; placement: $(placement)" - - for pattern in "failed to load" "failed to save" "lookup failed" "background save failed"; do - if grep -qs "mooncake-store.*${pattern}" "${log_dir}"/3_output_CTX_*.log; then - announce "fail_${pattern// /_}" "PROBLEM: mooncake-store ${pattern}" - fi - done - - [ -f "${log_dir}/6_bench.log" ] && announce "bench_started" "benchmark client started" - ls "${log_dir}"/concurrency_*/result.json >/dev/null 2>&1 \ - && announce "result" "result.json written" - - if ls "${log_dir}"/8_done_*.txt >/dev/null 2>&1; then - echo "EVENT: job finished (8_done marker present)" - echo "FINAL: matched=${matched:-0} loaded=${loaded:-0}" - echo "FINAL placement: $(placement)" - exit 0 - fi - - # A dead batch script leaves the tree untouched, so report it rather than - # polling until the deadline. - if grep -qs "Job completed successfully" "${log_dir}"/slurm-*.out; then - echo "EVENT: batch script reported completion" - exit 0 - fi - if grep -qs "^Error: " "${log_dir}"/slurm-*.out; then - echo "EVENT: PROBLEM: batch script hit cleanup_on_failure" - grep -h "^Error: " "${log_dir}"/slurm-*.out | tail -n 3 - exit 1 - fi - - sleep "${poll}" -done - -echo "EVENT: watcher deadline reached after ${max_minutes} minutes" -echo "FINAL: matched=$(count_matches 'mooncake-store matched') loaded=$(count_matches 'mooncake-store rank')" From 5dc8ffbc1974176bbb3afb2bf71d74c2c45e1aeb Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:10:39 -0700 Subject: [PATCH 20/24] [None][chore] Drop the stall-report diagnostic from the branch TRTLLM_STALL_REPORT_SEC came out of debugging the no-evict deadlock and is unrelated to the mooncake-store connector. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docs/source/developer-guide/overview.md | 21 ----- .../_torch/pyexecutor/hang_detector.py | 48 ------------ .../executor/test_hang_detector_kill.py | 76 ------------------- 3 files changed, 145 deletions(-) diff --git a/docs/source/developer-guide/overview.md b/docs/source/developer-guide/overview.md index bd8d8cd5b77c..d8fe31631612 100644 --- a/docs/source/developer-guide/overview.md +++ b/docs/source/developer-guide/overview.md @@ -122,24 +122,3 @@ export TLLM_LOG_LEVEL_BY_MODULE="debug:_torch,runtime;info:serve" ``` This example sets the global level to `warning` but enables `debug` output for `_torch` and `runtime` modules, and `info` for `serve`. Valid levels: `trace`, `debug`, `verbose`, `info`, `warning`, `error`, `internal_error`. - -### Diagnosing Slow Iterations - -`TRTLLM_STALL_REPORT_SEC` dumps every thread's stack to stderr whenever a single -executor iteration takes longer than the given number of seconds: - -```bash -# Any iteration slower than 5s dumps all thread stacks. -export TRTLLM_STALL_REPORT_SEC=5 -``` - -Set it somewhat above the normal iteration time, which the `host_step_time` field -of the per-iteration log line reports. Unlike the hang detector, this neither -kills the process nor stops the run, so it is safe to leave on for a whole -benchmark. It stays silent while iterations are under the threshold and does not -report time the loop spends idle waiting for requests. - -Background threads are included in the dump, which is usually the point: a slow -iteration that is not blocked in any obvious call is often waiting on a KV -transfer, connector or sampler thread that the main thread's stack does not -explain. diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index ff7200544ad8..2ae692ed5902 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio -import faulthandler import os import signal import sys @@ -27,13 +26,6 @@ # 137 == 128 + SIGKILL(9): the exit code a shell reports for a SIGKILL'd process. _HARD_KILL_EXIT_CODE = 137 -#: Seconds an iteration may take before all thread stacks are dumped to stderr. -#: Unset or non-positive disables it. This diagnoses a *slow* loop, as opposed to -#: the hung loop `HangDetector` exists for, so it neither kills the process nor -#: stops the run. Set it somewhat above the normal iteration time and read the -#: dumps out of the worker log afterwards. -STALL_REPORT_ENV = "TRTLLM_STALL_REPORT_SEC" - def _best_effort_flush_streams() -> None: """Flush stdout/stderr without ever raising; diagnostics must not block hard kill.""" @@ -108,41 +100,6 @@ def __init__( self.lock = threading.Lock() self.active = False self._detected = False - self._stall_report_sec = self._read_stall_report_sec() - if self._stall_report_sec > 0: - logger.info( - f"Stall reporting enabled: dumping all thread stacks for any " - f"iteration exceeding {self._stall_report_sec}s " - f"({STALL_REPORT_ENV})." - ) - - @staticmethod - def _read_stall_report_sec() -> float: - raw = os.environ.get(STALL_REPORT_ENV, "") - if not raw.strip(): - return 0.0 - try: - return float(raw) - except ValueError: - logger.warning(f"Ignoring {STALL_REPORT_ENV}={raw!r}: not a number.") - return 0.0 - - def _arm_stall_report(self) -> None: - """Schedule a stack dump if this iteration runs long. - - `faulthandler`'s timer lives in a thread that does not take the GIL, so - unlike `print_all_stacks` it still fires when the loop is blocked inside - a native call. That is the case worth diagnosing, since a stall in pure - Python would already show up in a profile. - """ - if self._stall_report_sec <= 0: - return - faulthandler.dump_traceback_later(self._stall_report_sec, repeat=False, exit=False) - - def _cancel_stall_report(self) -> None: - if self._stall_report_sec <= 0: - return - faulthandler.cancel_dump_traceback_later() def start(self): """Enable hang detection.""" @@ -172,16 +129,11 @@ def detected(self): def checkpoint(self): """Reset hang detection timer.""" self.cancel_task() - self._arm_stall_report() if self.active: self.task = asyncio.run_coroutine_threadsafe(self._detect_hang(), self.loop) def cancel_task(self): """Cancel the hang detection task.""" - # Disarmed here rather than in checkpoint() so that pause(), used by the - # request-queue wait and the cross-rank broadcast probe, does not report - # a stall for time the loop is legitimately idle. - self._cancel_stall_report() if self.task is not None and not self.task.done(): self.task.cancel() self.task = None diff --git a/tests/unittest/_torch/executor/test_hang_detector_kill.py b/tests/unittest/_torch/executor/test_hang_detector_kill.py index 3dd00982c941..b4f122bb8ad3 100644 --- a/tests/unittest/_torch/executor/test_hang_detector_kill.py +++ b/tests/unittest/_torch/executor/test_hang_detector_kill.py @@ -64,82 +64,6 @@ def test_pause_suppresses_detection(): assert hd.detected() is False -# Driver for the stall-reporting tests. faulthandler writes straight to fd 2, -# so the dump can only be observed from another process; markers go to stderr -# too, which is what makes their interleaving with the dump assertable. -_STALL_REPORT_SCRIPT = """ -import os, sys, time -from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector - -def mark(name): - print(f"MARK {name}", file=sys.stderr, flush=True) - -os.environ["TRTLLM_STALL_REPORT_SEC"] = "1" -hd = HangDetector(timeout=300, on_detected=lambda: None) -with hd: - mark("fast-begin") - hd.checkpoint() - time.sleep(0.2) - mark("fast-end") - - hd.checkpoint() - mark("paused-begin") - with hd.pause(): - time.sleep(2.0) - mark("paused-end") - - hd.checkpoint() - mark("slow-begin") - time.sleep(2.0) - mark("slow-end") - -# Same slow iteration, but with the knob unset: must stay silent. -os.environ.pop("TRTLLM_STALL_REPORT_SEC", None) -off = HangDetector(timeout=300, on_detected=lambda: None) -with off: - off.checkpoint() - mark("disabled-begin") - time.sleep(2.0) - mark("disabled-end") -""" - - -def test_stall_report_fires_only_for_slow_iterations(): - """The stack dump lands in the slow iteration, and nowhere else. - - Covers all four cases in one subprocess because each one would otherwise - pay a cold `import tensorrt_llm`. - """ - proc = subprocess.run( - [sys.executable, "-c", _STALL_REPORT_SCRIPT], - env={**os.environ, "TLLM_DISABLE_MPI": "1"}, - timeout=300, - capture_output=True, - ) - stderr = proc.stderr.decode(errors="replace") - assert proc.returncode == 0, f"driver failed: {stderr[-2000:]}" - - def index_of(marker): - position = stderr.find(f"MARK {marker}") - assert position != -1, f"missing marker {marker!r} in:\n{stderr[-2000:]}" - return position - - # faulthandler's dump is headed by "Timeout (H:MM:SS)!". - dumps = [ - position - for position in range(len(stderr)) - if stderr.startswith("Timeout (", position) - ] - assert len(dumps) == 1, ( - f"expected exactly one stack dump, got {len(dumps)}:\n{stderr[-2000:]}" - ) - # Only the slow iteration should have produced it: not the fast one, not - # the paused window (which is longer than the threshold), and not the run - # with the knob unset. - assert index_of("slow-begin") < dumps[0] < index_of("slow-end") - assert "MARK disabled-end" in stderr - - def test_propagate_hard_kill_self_sigkills_without_mpi(): """With MPI disabled, propagate_hard_kill self-SIGKILLs the process. From 652948cb67b2f5bde550285b599e03221fb0ff08 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:19:31 -0700 Subject: [PATCH 21/24] [None][test] Collapse redundant mooncake-store tests into parametrized cases Seven groups covered the same code path with trivially different inputs: connector recognition, provisioning no-ops, pool master validation, omitted client-config fields, JSON size parsing, the rank's device, and the two donor entry paths. Each is now one parametrized test. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../executor/test_mooncake_store_connector.py | 71 +++++++++---------- .../executor/test_mooncake_store_donor.py | 61 ++++++++-------- .../executor/test_mooncake_store_master.py | 68 +++++++++--------- 3 files changed, 96 insertions(+), 104 deletions(-) diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 52426dc49b5c..4c8ee16ba6b7 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -362,12 +362,23 @@ def test_page_addressing_rejects_mixed_slot_counts(): # ---- config ---- -def test_config_reads_sizes_with_units(store_config): +def test_config_reads_sizes_and_staging_from_the_json(store_config): + """Sizes arrive as unit strings, and staging is off until the JSON asks.""" config = MooncakeStoreConnectorConfig.from_env() assert config.global_segment_size == 1024**3 assert config.local_buffer_size == 256 * 1024**2 assert config.role is StoreRole.BOTH assert config.resolve_model_key("/models/ignored") == "test-model" + assert config.stage_through_host is False + + raw = json.loads(store_config.read_text()) + raw["stage_through_host"] = True + raw["staging_buffer_bytes"] = "256MiB" + store_config.write_text(json.dumps(raw)) + + config = MooncakeStoreConnectorConfig.from_env() + assert config.stage_through_host is True + assert config.staging_buffer_bytes == 256 * 1024**2 def test_config_role_comes_from_environment(store_config, monkeypatch): @@ -479,23 +490,25 @@ def test_validate_layout_rejects_sliding_window(): # the reuse the store exists to provide. -def test_uses_connector_recognizes_the_preset(): - config = KvCacheConnectorConfig(connector="mooncake-store") - assert uses_connector(config, "mooncake-store") - - -def test_uses_connector_recognizes_a_hand_written_module(): - config = KvCacheConnectorConfig( - connector_module="tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", - connector_scheduler_class="MooncakeStoreConnectorScheduler", - connector_worker_class="MooncakeStoreConnectorWorker", - ) - assert uses_connector(config, "mooncake-store") - - -def test_uses_connector_separates_connectors_and_tolerates_none(): - assert not uses_connector(KvCacheConnectorConfig(connector="kvbm"), "mooncake-store") - assert not uses_connector(None, "mooncake-store") +@pytest.mark.parametrize( + "config, expected", + [ + (KvCacheConnectorConfig(connector="mooncake-store"), True), + ( + KvCacheConnectorConfig( + connector_module="tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + connector_scheduler_class="MooncakeStoreConnectorScheduler", + connector_worker_class="MooncakeStoreConnectorWorker", + ), + True, + ), + (KvCacheConnectorConfig(connector="kvbm"), False), + (None, False), + ], + ids=["preset", "hand_written_module", "another_connector", "no_connector"], +) +def test_uses_connector_recognizes_the_connector_however_it_is_spelled(config, expected): + assert uses_connector(config, "mooncake-store") is expected def test_uses_connector_rejects_an_unknown_preset(): @@ -682,17 +695,6 @@ def test_plan_slot_geometry_rejects_degenerate_inputs(bad): plan_slot_geometry(*bad) -def test_config_reads_staging_from_the_json(store_config): - raw = json.loads(store_config.read_text()) - raw["stage_through_host"] = True - raw["staging_buffer_bytes"] = "256MiB" - store_config.write_text(json.dumps(raw)) - - config = MooncakeStoreConnectorConfig.from_env() - assert config.stage_through_host is True - assert config.staging_buffer_bytes == 256 * 1024**2 - - @pytest.mark.parametrize( "value,expected", [("1", True), ("true", True), ("on", True), ("0", False), ("off", False)] ) @@ -802,12 +804,7 @@ def test_staging_does_not_scatter_a_failed_load(store_config, fake_store, staged assert staged_copies == [] -def test_worker_captures_the_ranks_device_at_registration(store_config, fake_store, fake_cuda): - with make_worker(fake_store, layout=make_layout()) as worker: - assert worker._device_index == 3 - - -def test_save_thread_adopts_the_ranks_device_not_the_thread_default( +def test_the_ranks_device_is_captured_and_adopted_by_the_save_thread( store_config, fake_store, fake_cuda ): """The save thread must not run on torch's default device. @@ -816,7 +813,9 @@ def test_save_thread_adopts_the_ranks_device_not_the_thread_default( stream created on device 0 instead fails every copy with cudaErrorInvalidValue, and only on ranks other than 0. """ - with make_worker(fake_store, layout=make_layout()): + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._device_index == 3 + deadline = time.monotonic() + 5.0 while 3 not in fake_cuda and time.monotonic() < deadline: time.sleep(0.01) diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py index 841f42e5f5db..e3a649878cc6 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_donor.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -77,16 +77,38 @@ def setup(self, *args): return Refusing -def test_a_donor_registers_the_segment_it_was_asked_for(fake_bindings): - with donate_segment( +@pytest.mark.parametrize("entry", ["direct", "config"]) +def test_a_donor_registers_the_segment_it_was_asked_for( + entry, fake_bindings, reachable_master, monkeypatch): + """Both entry paths must reach Mooncake with the same seven setup arguments. + + The config-driven path is what makes a generation server a donor, and it + derives the hostname and parses the size string on the way: a size string + reaching Mooncake unparsed would be a segment of nothing. + """ + if entry == "direct": + expected_host, expected_size, expected_device = "10.0.0.5", 32 * GIB, "mlx5_0" + donation = donate_segment( "10.0.0.1:50051", 32 * GIB, protocol="rdma", device_name="mlx5_0", metadata_server="P2PHANDSHAKE", hostname="10.0.0.5", - ) as host: - assert host == "10.0.0.5" + ) + else: + monkeypatch.setattr(donor_module, "local_address", lambda: "10.1.2.3") + expected_host, expected_size, expected_device = "10.1.2.3", 320 * GIB, "mlx5_1" + donation = maybe_donate_segment( + MooncakeDonationConfig( + master_server_address="10.0.0.1:50051", + segment_size="320GiB", + protocol="rdma", + device_name="mlx5_1", + )) + + with donation as host: + assert host == expected_host ( registered_host, metadata_server, @@ -97,11 +119,11 @@ def test_a_donor_registers_the_segment_it_was_asked_for(fake_bindings): master, ) = fake_bindings.instances[0].setup_args - assert registered_host == "10.0.0.5" + assert registered_host == expected_host assert metadata_server == "P2PHANDSHAKE" - assert segment_size == 32 * GIB + assert segment_size == expected_size assert protocol == "rdma" - assert device_name == "mlx5_0" + assert device_name == expected_device assert master == "10.0.0.1:50051" assert local_buffer_size == DEFAULT_DONOR_LOCAL_BUFFER_SIZE @@ -138,31 +160,6 @@ def reachable_master(monkeypatch): monkeypatch.setattr(donor_module, "wait_for_master", lambda address: 0.0) -def test_a_server_that_was_asked_to_lend_memory_does(fake_bindings, reachable_master, - monkeypatch): - """The config-driven path is what makes a generation server a donor.""" - monkeypatch.setattr(donor_module, "local_address", lambda: "10.1.2.3") - donation = MooncakeDonationConfig( - master_server_address="10.0.0.1:50051", - segment_size="320GiB", - protocol="rdma", - device_name="mlx5_1", - ) - - with maybe_donate_segment(donation) as host: - assert host == "10.1.2.3" - registered_host, metadata_server, segment_size, _, protocol, device, master = ( - fake_bindings.instances[0].setup_args) - - assert registered_host == "10.1.2.3" - assert metadata_server == "P2PHANDSHAKE" - # A size string reaching Mooncake unparsed would be a segment of nothing. - assert segment_size == 320 * GIB - assert protocol == "rdma" - assert device == "mlx5_1" - assert master == "10.0.0.1:50051" - - def test_a_server_that_was_not_asked_lends_nothing(fake_bindings): """Every deployment that does not lend memory takes this path.""" with maybe_donate_segment(None) as host: diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py index a9c8f7e3e66e..03c678e779a4 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_master.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -162,20 +162,23 @@ def running_master(): # ---- configuration ---- -def test_pool_needs_exactly_one_master(): - with pytest.raises(ValueError, match="not both"): - MooncakeStoreConfig(launch_master=True, master_server_address="host:50051") - with pytest.raises(ValueError, match="needs a master"): - MooncakeStoreConfig() - - -def test_publishing_an_address_needs_a_master_to_publish(): - """This option only writes an address; reading one is master_server_address.""" - with pytest.raises(ValueError, match="needs launch_master"): - MooncakeStoreConfig( - master_server_address="host:50051", - master_address_file="/shared/master.addr", - ) +@pytest.mark.parametrize( + "kwargs, message", + [ + (dict(launch_master=True, master_server_address="host:50051"), "not both"), + (dict(), "needs a master"), + # master_address_file only writes an address; reading one is + # master_server_address, so publishing without launching is incoherent. + ( + dict(master_server_address="host:50051", master_address_file="/shared/master.addr"), + "needs launch_master", + ), + ], + ids=["two_masters", "no_master", "publishing_without_launching"], +) +def test_pool_needs_exactly_one_master(kwargs, message): + with pytest.raises(ValueError, match=message): + MooncakeStoreConfig(**kwargs) def test_pool_is_rejected_unless_the_connector_is_mooncake_store(): @@ -224,9 +227,12 @@ def test_client_config_is_what_the_connector_reads_back(tmp_path): assert parsed.transfer_batch_size == 32 -def test_client_config_leaves_an_unset_prefix_to_the_connector(): +def test_client_config_omits_the_fields_the_pool_left_unset(): + """An absent key leaves the connector its own default; a null would not.""" pool = MooncakeStoreConfig(master_server_address="host:50051") - assert "cache_prefix" not in master_module._client_config(pool, "host:50051") + written = master_module._client_config(pool, "host:50051") + assert "cache_prefix" not in written + assert "staging_buffer_bytes" not in written @pytest.mark.parametrize( @@ -273,13 +279,6 @@ def test_a_staging_buffer_can_be_sized_where_staging_is_turned_on(running_master assert written["staging_buffer_bytes"] == "4GiB" -def test_an_unsized_staging_buffer_leaves_the_connector_its_default(running_master): - pool = MooncakeStoreConfig(master_server_address=running_master) - - with provision_pool(pool) as config_path: - assert "staging_buffer_bytes" not in json.loads(open(config_path).read()) - - def test_provisioning_fails_before_the_model_loads_if_the_master_is_absent(monkeypatch): monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") pool = MooncakeStoreConfig(master_server_address=f"127.0.0.1:{free_port()}") @@ -408,20 +407,17 @@ def test_a_run_dir_keeps_the_master_log_and_the_config(fake_master, tmp_path): # ---- the entry point servers call ---- -def test_other_connectors_are_left_alone(): - config = KvCacheConnectorConfig(connector="lmcache") - with maybe_provision_pool(config): - assert CONFIG_PATH_ENV not in os.environ - - -def test_no_connector_at_all_is_left_alone(): - with maybe_provision_pool(None): - assert CONFIG_PATH_ENV not in os.environ - - -def test_a_pool_left_undescribed_stays_the_environment_contract(): +@pytest.mark.parametrize( + "config", + [ + KvCacheConnectorConfig(connector="lmcache"), + None, + KvCacheConnectorConfig(connector="mooncake-store"), + ], + ids=["another_connector", "no_connector", "pool_left_undescribed"], +) +def test_provisioning_is_a_no_op_unless_a_pool_is_described(config): """Without `mooncake_store`, MOONCAKE_CONFIG_PATH is still the only input.""" - config = KvCacheConnectorConfig(connector="mooncake-store") with maybe_provision_pool(config): assert CONFIG_PATH_ENV not in os.environ From e6d39fbfda8bbcd560429e7f571ca31c23835ff7 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Fri, 4 Sep 2026 10:15:16 -0700 Subject: [PATCH 22/24] [None][test] Run the mooncake-store and KV-cache-v2 scheduler tests on the M3 stage The mooncake-store connector tests and the KVCacheV2Scheduler mock tests are CPU-only, so they fit the single-GPU M3 pre-merge stage. The scheduler tests were not listed anywhere, and this branch adds cases to them, so list the file alongside the three new mooncake-store test modules. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200_m3.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/integration/test_lists/test-db/l0_b200_m3.yml b/tests/integration/test_lists/test-db/l0_b200_m3.yml index 8e5c9eacbfde..07e1d2e95d21 100644 --- a/tests/integration/test_lists/test-db/l0_b200_m3.yml +++ b/tests/integration/test_lists/test-db/l0_b200_m3.yml @@ -18,6 +18,10 @@ l0_b200_m3: - unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py - unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py - unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py + - unittest/_torch/executor/test_kv_cache_v2_scheduler.py + - unittest/_torch/executor/test_mooncake_store_connector.py + - unittest/_torch/executor/test_mooncake_store_donor.py + - unittest/_torch/executor/test_mooncake_store_master.py - unittest/_torch/models/test_minimax_m3.py - unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py From 6b11d7f7231b3b4f94876c336ac42989fba0db24 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:25:14 +0000 Subject: [PATCH 23/24] formatting Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/__init__.py | 16 +- .../connectors/mooncake_store/donor.py | 15 +- .../connectors/mooncake_store/master.py | 20 +- .../connectors/mooncake_store/staging.py | 4 +- .../connectors/mooncake_store/worker.py | 8 +- .../_torch/pyexecutor/kv_cache_manager_v2.py | 4 +- .../pyexecutor/scheduler/scheduler_v2.py | 9 +- tensorrt_llm/commands/mooncake.py | 269 +++++++++++------- .../executor/test_kv_cache_v2_scheduler.py | 13 +- .../executor/test_mooncake_store_connector.py | 9 +- .../executor/test_mooncake_store_donor.py | 18 +- .../executor/test_mooncake_store_master.py | 22 +- 12 files changed, 226 insertions(+), 181 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index 378dbda6f4d3..bae898f8619a 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -47,11 +47,17 @@ """ from .config import MooncakeStoreConnectorConfig, StoreRole, parse_size -from .donor import (DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, - maybe_donate_segment) -from .master import (local_address, master_timeout, maybe_provision_pool, - provision_pool, resolve_device_name, - resolve_master_address, running_master, wait_for_master) +from .donor import DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, maybe_donate_segment +from .master import ( + local_address, + master_timeout, + maybe_provision_pool, + provision_pool, + resolve_device_name, + resolve_master_address, + running_master, + wait_for_master, +) from .scheduler import MooncakeStoreConnectorScheduler from .worker import MooncakeStoreConnectorWorker diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py index 42a765de62bb..15514d69a996 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -40,8 +40,13 @@ from tensorrt_llm.logger import logger from .config import parse_size -from .master import (local_address, master_timeout, resolve_device_name, - resolve_master_address, wait_for_master) +from .master import ( + local_address, + master_timeout, + resolve_device_name, + resolve_master_address, + wait_for_master, +) __all__ = [ "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", @@ -82,7 +87,7 @@ def donate_segment( ) from exc host = hostname or local_address() - donated = f"{segment_size / 1024 ** 3:.1f}GiB" + donated = f"{segment_size / 1024**3:.1f}GiB" # Byte counts are spelled out next to the human-readable form. A misparsed # size string otherwise surfaces only as a pool that evicts far too eagerly. logger.info( @@ -154,9 +159,7 @@ def maybe_donate_segment(donation: Any) -> Iterator[Optional[str]]: f"memory to the pool at {donation.master_server_address} without using " "it; resolving the master now" ) - master_address = resolve_master_address( - donation.master_server_address, master_timeout() - ) + master_address = resolve_master_address(donation.master_server_address, master_timeout()) # Checked before setup so an absent master is reported as such, rather than # as the status code setup returns for every kind of failure. wait_for_master(master_address) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index d8f8d38d40b3..f2f3a21a3c91 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -137,7 +137,7 @@ def resolve_master_address(address: str, timeout: float) -> str: if not address.startswith(ADDRESS_FILE_SCHEME): return address - path = address[len(ADDRESS_FILE_SCHEME):] + path = address[len(ADDRESS_FILE_SCHEME) :] started = time.monotonic() deadline = started + timeout announced = started @@ -267,9 +267,7 @@ def attribute(name: str) -> str: return [device for device, rate in sorted(rated.items()) if rate == fastest] -def resolve_device_name(protocol: str, - configured: str, - sysfs_root: Optional[str] = None) -> str: +def resolve_device_name(protocol: str, configured: str, sysfs_root: Optional[str] = None) -> str: """The RDMA devices to transfer over, detected if the config left it open. Which HCAs a node has is a property of the node, not of the deployment, so @@ -316,15 +314,13 @@ def wait_for_master(master_address: str, timeout: Optional[float] = None) -> Opt ) return None elapsed = _wait_until_accepting(*endpoint, timeout) - logger.info( - f"mooncake-store: the master at {master_address} answered in {elapsed:.1f}s" - ) + logger.info(f"mooncake-store: the master at {master_address} answered in {elapsed:.1f}s") return elapsed -def _client_config(pool: Any, - master_address: str, - device_name: Optional[str] = None) -> Dict[str, Any]: +def _client_config( + pool: Any, master_address: str, device_name: Optional[str] = None +) -> Dict[str, Any]: """Render the Mooncake client config for a pool. The schema is vLLM's, so one pool can serve both engines. `role` is written @@ -548,8 +544,8 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona config_path = os.path.join(run_dir, CLIENT_CONFIG_NAME) config = _client_config( - pool, master_address, - resolve_device_name(pool.protocol, pool.device_name)) + pool, master_address, resolve_device_name(pool.protocol, pool.device_name) + ) with open(config_path, "w") as handle: json.dump(config, handle, indent=2) # Inherited by the ranks the LLM constructor spawns. Ranks an diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py index 1441367f67c5..966897b178c4 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py @@ -247,9 +247,7 @@ def stage_batch_for_put( Per-page address and size lists, each a single staged buffer. """ if len(addresses) > pool.num_slots: - raise ValueError( - f"batch of {len(addresses)} pages exceeds {pool.num_slots} staging slots" - ) + raise ValueError(f"batch of {len(addresses)} pages exceeds {pool.num_slots} staging slots") staged_addresses: List[List[int]] = [] staged_sizes: List[List[int]] = [] for index, (page_addresses, page_sizes) in enumerate(zip(addresses, sizes, strict=True)): diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 9f7e02846f6f..6209d44c7e95 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -54,9 +54,9 @@ describe_batch_for_get, plan_slot_geometry, stage_batch_for_put, - sync_stream as _sync_stream, unstage_batch_after_get, ) +from .staging import sync_stream as _sync_stream from .validation import validate_layout, validate_llm_args __all__ = ["MooncakeStoreConnectorWorker", "resolve_local_worker"] @@ -265,8 +265,7 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: self._save_thread.start() logger.info( - f"mooncake-store worker rank {self._rank} registered layout: " - f"{addressing.describe()}" + f"mooncake-store worker rank {self._rank} registered layout: {addressing.describe()}" ) def _open_staging(self, addressing: PageAddressing) -> None: @@ -523,8 +522,7 @@ def _drain_saves(self) -> None: # escapes here would be lost, so it is stashed and re-raised on # the executor thread at the next connector call. logger.error( - f"mooncake-store save failed on rank {self._rank}: " - f"{type(exc).__name__}: {exc}" + f"mooncake-store save failed on rank {self._rank}: {type(exc).__name__}: {exc}" ) with self._save_lock: if self._save_error is None: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index fb7a90706cc7..2f711ed2ab17 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2667,9 +2667,7 @@ def preempt_request(self, req: LlmRequest) -> bool: # DISAGG_CONTEXT_TRANS_IN_PROGRESS, out of the schedulable range, and # its `_KVCache` keeps holding the pages until every rank reports the # save retired through `get_finished`. - if self.kv_connector_manager.request_finished( - req, self.get_connector_page_indices(req) - ): + if self.kv_connector_manager.request_finished(req, self.get_connector_page_indices(req)): self._pending_preemption[req.py_request_id] = req return False diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 80861f8060fd..23e922976e3e 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -450,11 +450,7 @@ def preempt_for_pages(req: LlmRequest) -> bool: pending_ctx, preempted_ids, made_progress=bool( - scheduled_gen - or scheduled_ctx - or scheduled_encoder - or disagg_candidates - or evicted + scheduled_gen or scheduled_ctx or scheduled_encoder or disagg_candidates or evicted ), ) @@ -1135,8 +1131,7 @@ def _detect_deadlock( num_ctx_candidates = sum( 1 for r in pending_ctx - if r.py_request_id not in preempted_ids - and r.request_id not in inflight_request_ids + if r.py_request_id not in preempted_ids and r.request_id not in inflight_request_ids ) if num_gen_candidates == 0 and num_ctx_candidates == 0: # Legitimately idle: nothing to schedule. diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index f9c5423fa37b..aeb06315c3ca 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -53,42 +53,59 @@ def stop(signum, _frame): @click.command("mooncake_master") -@click.option("--rpc_port", - type=int, - default=50051, - show_default=True, - help="Port the store clients reach the master on.") -@click.option("--metrics_port", - type=int, - default=9004, - show_default=True, - help="Prometheus port. Pool occupancy and eviction are read " - "from here or from the master's log.") -@click.option("--eviction_ratio", - type=float, - default=0.05, - show_default=True, - help="Fraction of the pool freed per eviction pass.") -@click.option("--address_file", - type=str, - default=None, - help="File to publish 'host:port' to once the master answers. " - "Workers name it as master_server_address: file://, which " - "is how they reach a master whose host the scheduler chose. " - "Removed on exit so a stale address is never dialed.") -@click.option("--run_dir", - type=str, - default=None, - help="Where to keep the master's log. Defaults to " - "$TRTLLM_MOONCAKE_RUN_DIR, else a temporary directory.") -@click.option("--heartbeat_seconds", - type=int, - default=300, - show_default=True, - help="Interval between liveness lines. 0 disables them.") -def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, - address_file: Optional[str], run_dir: Optional[str], - heartbeat_seconds: int): +@click.option( + "--rpc_port", + type=int, + default=50051, + show_default=True, + help="Port the store clients reach the master on.", +) +@click.option( + "--metrics_port", + type=int, + default=9004, + show_default=True, + help="Prometheus port. Pool occupancy and eviction are read " + "from here or from the master's log.", +) +@click.option( + "--eviction_ratio", + type=float, + default=0.05, + show_default=True, + help="Fraction of the pool freed per eviction pass.", +) +@click.option( + "--address_file", + type=str, + default=None, + help="File to publish 'host:port' to once the master answers. " + "Workers name it as master_server_address: file://, which " + "is how they reach a master whose host the scheduler chose. " + "Removed on exit so a stale address is never dialed.", +) +@click.option( + "--run_dir", + type=str, + default=None, + help="Where to keep the master's log. Defaults to " + "$TRTLLM_MOONCAKE_RUN_DIR, else a temporary directory.", +) +@click.option( + "--heartbeat_seconds", + type=int, + default=300, + show_default=True, + help="Interval between liveness lines. 0 disables them.", +) +def mooncake_master( + rpc_port: int, + metrics_port: int, + eviction_ratio: float, + address_file: Optional[str], + run_dir: Optional[str], + heartbeat_seconds: int, +): """Run a mooncake_master for as long as this command runs. A single server with a pool of its own should set @@ -96,8 +113,7 @@ def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, """ # Imported lazily so other subcommands and --help do not pay for the # connector package. - from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import \ - running_master + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import running_master from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig pool = MooncakeStoreConfig( @@ -106,16 +122,19 @@ def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, master_metrics_port=metrics_port, master_eviction_ratio=eviction_ratio, ) - run_dir = run_dir or os.getenv( - "TRTLLM_MOONCAKE_RUN_DIR") or tempfile.mkdtemp( - prefix="trtllm-mooncake-master-") + run_dir = ( + run_dir + or os.getenv("TRTLLM_MOONCAKE_RUN_DIR") + or tempfile.mkdtemp(prefix="trtllm-mooncake-master-") + ) stopping = _until_signalled() with running_master(pool, run_dir, address_file=address_file) as master: logger.info( f"mooncake-store: this master owns the pool until this command " f"stops; address {master.address}, log {master.log_path}, metrics " - f"http://{master.address.rsplit(':', 1)[0]}:{metrics_port}/metrics") + f"http://{master.address.rsplit(':', 1)[0]}:{metrics_port}/metrics" + ) started = time.monotonic() announced = started while not stopping.is_set(): @@ -123,8 +142,8 @@ def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, # The pool is gone once the master dies, and every client is # about to start failing. raise click.ClickException( - f"mooncake_master exited with code {code}. See " - f"{master.log_path}") + f"mooncake_master exited with code {code}. See {master.log_path}" + ) stopping.wait(1.0) now = time.monotonic() # Distinguishes a dead master from a dead fabric once clients @@ -133,67 +152,93 @@ def mooncake_master(rpc_port: int, metrics_port: int, eviction_ratio: float, announced = now logger.info( f"mooncake-store: master at {master.address} alive after " - f"{(now - started) / 60:.0f}m") + f"{(now - started) / 60:.0f}m" + ) @click.command("mooncake_donor") -@click.option("--master_server_address", - type=str, - default=None, - help="Master to join, as host:port or file:// naming a " - "file that holds one. Defaults to the master_server_address in " - "--config.") -@click.option("--segment_size", - type=str, - default="32GiB", - show_default=True, - help="Host memory to contribute from this node. Deliberately " - "separate from a config's global_segment_size, which is sized " - "for an engine worker rather than a node lending what it can " - "spare.") -@click.option("--config", - type=str, - default=None, - help="Mooncake JSON config describing the pool, for the " - "settings not given here. Defaults to $MOONCAKE_CONFIG_PATH.") -@click.option("--protocol", - type=str, - default=None, - help="Transport, 'rdma' or 'tcp'. Defaults to --config's, else " - "rdma.") -@click.option("--device_name", - type=str, - default=None, - help="RDMA device, from ibv_devinfo. Defaults to --config's.") -@click.option("--metadata_server", - type=str, - default=None, - help="Mooncake metadata service. Defaults to --config's.") -@click.option("--ready_file", - type=str, - default=None, - help="File to create once the segment is mounted, for launchers " - "that must not let prefill start writing before the pool has " - "this capacity.") -@click.option("--heartbeat_seconds", - type=int, - default=300, - show_default=True, - help="Interval between liveness lines. 0 disables them.") -def mooncake_donor(master_server_address: Optional[str], segment_size: str, - config: Optional[str], protocol: Optional[str], - device_name: Optional[str], metadata_server: Optional[str], - ready_file: Optional[str], heartbeat_seconds: int): +@click.option( + "--master_server_address", + type=str, + default=None, + help="Master to join, as host:port or file:// naming a " + "file that holds one. Defaults to the master_server_address in " + "--config.", +) +@click.option( + "--segment_size", + type=str, + default="32GiB", + show_default=True, + help="Host memory to contribute from this node. Deliberately " + "separate from a config's global_segment_size, which is sized " + "for an engine worker rather than a node lending what it can " + "spare.", +) +@click.option( + "--config", + type=str, + default=None, + help="Mooncake JSON config describing the pool, for the " + "settings not given here. Defaults to $MOONCAKE_CONFIG_PATH.", +) +@click.option( + "--protocol", + type=str, + default=None, + help="Transport, 'rdma' or 'tcp'. Defaults to --config's, else rdma.", +) +@click.option( + "--device_name", + type=str, + default=None, + help="RDMA device, from ibv_devinfo. Defaults to --config's.", +) +@click.option( + "--metadata_server", + type=str, + default=None, + help="Mooncake metadata service. Defaults to --config's.", +) +@click.option( + "--ready_file", + type=str, + default=None, + help="File to create once the segment is mounted, for launchers " + "that must not let prefill start writing before the pool has " + "this capacity.", +) +@click.option( + "--heartbeat_seconds", + type=int, + default=300, + show_default=True, + help="Interval between liveness lines. 0 disables them.", +) +def mooncake_donor( + master_server_address: Optional[str], + segment_size: str, + config: Optional[str], + protocol: Optional[str], + device_name: Optional[str], + metadata_server: Optional[str], + ready_file: Optional[str], + heartbeat_seconds: int, +): """Lend this node's host memory to a Mooncake pool, for as long as it runs. Running this on the generation nodes puts their memory into the pool while leaving those engines connector-free. """ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( - DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, master_timeout, - parse_size, resolve_master_address, wait_for_master) - from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import \ - CONFIG_PATH_ENV + DEFAULT_DONOR_LOCAL_BUFFER_SIZE, + donate_segment, + master_timeout, + parse_size, + resolve_master_address, + wait_for_master, + ) + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import CONFIG_PATH_ENV raw = {} config = config or os.getenv(CONFIG_PATH_ENV) @@ -205,7 +250,8 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, if not master: raise click.UsageError( "No master to join. Pass --master_server_address, or a --config " - f"naming one (or set {CONFIG_PATH_ENV}).") + f"naming one (or set {CONFIG_PATH_ENV})." + ) donating = parse_size(segment_size) resolved = resolve_master_address(master, master_timeout()) @@ -214,21 +260,23 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, stopping = _until_signalled() with donate_segment( - resolved, - donating, - protocol=protocol or raw.get("protocol", "rdma"), - device_name=device_name or raw.get("device_name", "") or "", - metadata_server=metadata_server or raw.get("metadata_server", ""), - local_buffer_size=parse_size( - raw.get("local_buffer_size_donor", - DEFAULT_DONOR_LOCAL_BUFFER_SIZE)), + resolved, + donating, + protocol=protocol or raw.get("protocol", "rdma"), + device_name=device_name or raw.get("device_name", "") or "", + metadata_server=metadata_server or raw.get("metadata_server", ""), + local_buffer_size=parse_size( + raw.get("local_buffer_size_donor", DEFAULT_DONOR_LOCAL_BUFFER_SIZE) + ), ) as host: if ready_file: with open(ready_file, "w") as handle: handle.write(f"{host} {donating}\n") - logger.info(f"mooncake-store: announced this segment in " - f"{ready_file}, so a launcher waiting on the pool's " - "capacity can proceed") + logger.info( + f"mooncake-store: announced this segment in " + f"{ready_file}, so a launcher waiting on the pool's " + "capacity can proceed" + ) # Idle by design: a put or get here would make this node a traffic # client, which is what donation exists to avoid. @@ -240,5 +288,6 @@ def mooncake_donor(master_server_address: Optional[str], segment_size: str, if not stopping.wait(heartbeat_seconds): logger.info( f"mooncake-store: {host} still lending " - f"{donating / 1024 ** 3:.1f}GiB to the pool at {master} " - f"after {(time.monotonic() - started) / 60:.0f}m") + f"{donating / 1024**3:.1f}GiB to the pool at {master} " + f"after {(time.monotonic() - started) / 60:.0f}m" + ) diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index e3adca7ddeb3..2c24d0765f3c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -53,6 +53,7 @@ def make_gen_request( req.lora_task_id = lora_task_id req.is_context_init_state = False req.is_generation_in_progress_state = True + req.is_generation_to_complete_state = False req.is_first_context_chunk = is_first_context_chunk req.py_encoder_output_ready_event = None return req @@ -171,7 +172,13 @@ def make_kv_cache_manager( mgr.resize_context.side_effect = resize_context_fn or (lambda req, n: True) mgr.prepare_disagg_gen_init.side_effect = prepare_disagg_gen_init_fn or (lambda req: True) mgr.try_allocate_generation.side_effect = try_allocate_generation_fn or (lambda req: True) - mgr.suspend_request.return_value = None + + def _suspend(req): + # Mirrors KVCacheManagerV2.suspend_request: the cache stops being + # active on GPU, so the request is no longer an eviction victim. + mgr.kv_cache_map[req.py_request_id].is_active = False + + mgr.suspend_request.side_effect = _suspend mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active # The default here has a cache tier below GPU, which leaves preemption off. mgr.has_cache_tier_below_gpu = has_cache_tier_below_gpu @@ -1069,9 +1076,7 @@ def test_transient_stall_does_not_raise(self): def resize_fn(req, n): return not fail[0] - mgr = make_kv_cache_manager( - resize_context_fn=resize_fn, has_cache_tier_below_gpu=False - ) + mgr = make_kv_cache_manager(resize_context_fn=resize_fn, has_cache_tier_below_gpu=False) sched = make_scheduler(mgr, max_num_tokens=1000) sched._DEADLOCK_STALL_ITERS = 3 reqs = [make_ctx_request(0, 100, is_first_context_chunk=False)] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 4c8ee16ba6b7..03cf093a1104 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -37,6 +37,7 @@ KvCacheLayout, KvCacheRegion, ) +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, @@ -54,7 +55,6 @@ PageTransfer, RequestTransfers, ) -from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import staging as staging_module from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.scheduler import ( MooncakeStoreConnectorScheduler, ) @@ -756,9 +756,7 @@ def test_staging_put_hands_the_store_one_host_buffer_per_page( assert staged_copies == expected -def test_staging_get_scatters_back_to_the_device_regions( - store_config, fake_store, staged_copies -): +def test_staging_get_scatters_back_to_the_device_regions(store_config, fake_store, staged_copies): layout = make_layout(regions_per_group=3) with make_staged_worker(fake_store, store_config, layout=layout) as worker: block_hash = b"\x03" * 16 @@ -837,7 +835,8 @@ def test_staging_narrows_the_batch_to_the_budget(store_config, fake_store, stage worker._put( [ RequestTransfers( - 1, [PageTransfer(block_hash, 0, index) for index, block_hash in enumerate(hashes)] + 1, + [PageTransfer(block_hash, 0, index) for index, block_hash in enumerate(hashes)], ) ] ) diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py index e3a649878cc6..cda334bc385a 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_donor.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -68,7 +68,6 @@ def failing_bindings(fake_bindings): """Bindings whose `setup` refuses, as an unreachable master would.""" class Refusing(fake_bindings): - def setup(self, *args): super().setup(*args) return 7 @@ -79,7 +78,8 @@ def setup(self, *args): @pytest.mark.parametrize("entry", ["direct", "config"]) def test_a_donor_registers_the_segment_it_was_asked_for( - entry, fake_bindings, reachable_master, monkeypatch): + entry, fake_bindings, reachable_master, monkeypatch +): """Both entry paths must reach Mooncake with the same seven setup arguments. The config-driven path is what makes a generation server a donor, and it @@ -105,7 +105,8 @@ def test_a_donor_registers_the_segment_it_was_asked_for( segment_size="320GiB", protocol="rdma", device_name="mlx5_1", - )) + ) + ) with donation as host: assert host == expected_host @@ -135,7 +136,8 @@ def test_a_donor_that_cannot_join_says_which_master_it_could_not_reach(failing_b def test_a_donor_given_no_host_registers_under_the_pool_s_view_of_this_node( - fake_bindings, monkeypatch): + fake_bindings, monkeypatch +): """The master and the segments registering with it must agree on the host.""" monkeypatch.setattr(donor_module, "local_address", lambda: "10.1.2.3") @@ -167,8 +169,9 @@ def test_a_server_that_was_not_asked_lends_nothing(fake_bindings): assert fake_bindings.instances == [] -def test_a_published_master_address_is_read_before_joining(fake_bindings, reachable_master, - tmp_path): +def test_a_published_master_address_is_read_before_joining( + fake_bindings, reachable_master, tmp_path +): """What lets a generation server name a path instead of a scheduler's choice.""" address_file = tmp_path / "master.addr" address_file.write_text("10.0.0.9:50051\n") @@ -182,7 +185,8 @@ def test_a_published_master_address_is_read_before_joining(fake_bindings, reacha def test_an_unreachable_master_is_reported_before_the_segment_is_offered( - fake_bindings, monkeypatch): + fake_bindings, monkeypatch +): """Otherwise this is a status code from setup, with no address in it.""" def refuse(address): diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py index 03c678e779a4..618bcfbb8512 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_master.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -116,7 +116,6 @@ def __init__(self): self.process = None def arm(self, listen_on=None, exit_code=None, log_text=None): - def popen(command, env=None, stdout=None, **_kwargs): # A real master writes its log through this handle. if log_text is not None and stdout is not None: @@ -496,9 +495,7 @@ def test_a_stopped_master_leaves_no_address_behind(fake_master, tmp_path): address_file = tmp_path / "master.addr" pool = MooncakeStoreConfig(launch_master=True, master_port=port) - with master_module.running_master( - pool, str(tmp_path / "run"), address_file=str(address_file) - ): + with master_module.running_master(pool, str(tmp_path / "run"), address_file=str(address_file)): assert address_file.exists() assert not address_file.exists() @@ -599,8 +596,7 @@ def test_an_address_of_a_shape_we_cannot_probe_is_not_fatal(): def test_a_master_that_died_starting_is_reported_with_its_last_words(fake_master, tmp_path): """The reason is in the master's log, which is only read if the error quotes it.""" run_dir = tmp_path / "run" - fake_master.arm( - exit_code=1, log_text="E0903 bind(50051) failed: Address already in use\n") + fake_master.arm(exit_code=1, log_text="E0903 bind(50051) failed: Address already in use\n") pool = MooncakeStoreConfig(launch_master=True, master_port=free_port()) with pytest.raises(RuntimeError, match="Address already in use"): @@ -627,14 +623,14 @@ def test_the_compute_fabric_is_picked_over_the_management_adapter(tmp_path): fake_hca(tmp_path, "mlx5_3", state="1: DOWN") fake_hca(tmp_path, "mlx5_4", link_layer="Ethernet") - assert master_module.resolve_device_name( - "rdma", "", sysfs_root=str(tmp_path)) == "mlx5_0,mlx5_1" + assert ( + master_module.resolve_device_name("rdma", "", sysfs_root=str(tmp_path)) == "mlx5_0,mlx5_1" + ) def test_a_named_device_is_not_second_guessed(tmp_path): fake_hca(tmp_path, "mlx5_0") - assert master_module.resolve_device_name( - "rdma", "mlx5_7", sysfs_root=str(tmp_path)) == "mlx5_7" + assert master_module.resolve_device_name("rdma", "mlx5_7", sysfs_root=str(tmp_path)) == "mlx5_7" def test_tcp_needs_no_device_and_looks_for_none(tmp_path): @@ -643,12 +639,10 @@ def test_tcp_needs_no_device_and_looks_for_none(tmp_path): def test_a_node_without_infiniband_is_left_to_mooncake_s_own_discovery(tmp_path): """Falling back beats failing, since Mooncake may still find a usable device.""" - assert master_module.resolve_device_name( - "rdma", "", sysfs_root=str(tmp_path / "absent")) == "" + assert master_module.resolve_device_name("rdma", "", sysfs_root=str(tmp_path / "absent")) == "" -def test_the_detected_device_is_what_the_workers_are_told(fake_master, tmp_path, - monkeypatch): +def test_the_detected_device_is_what_the_workers_are_told(fake_master, tmp_path, monkeypatch): sysfs = tmp_path / "sysfs" fake_hca(sysfs, "mlx5_0") monkeypatch.setattr(master_module, "IB_SYSFS_ROOT", str(sysfs)) From 8665cf36e945c9fd4e796eb12ed16397fdaa5ae7 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:15:50 +0000 Subject: [PATCH 24/24] address comment from Pietro Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docs/source/features/kv-cache-connector.md | 2 ++ .../pyexecutor/connectors/mooncake_store/config.py | 9 +++++++-- .../pyexecutor/connectors/mooncake_store/donor.py | 6 +++--- tensorrt_llm/commands/mooncake.py | 9 ++++++--- .../_torch/executor/test_mooncake_store_connector.py | 11 +++++++++++ 5 files changed, 29 insertions(+), 8 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 6af39a1d79bf..05c3f10af8d3 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -215,6 +215,8 @@ Topology can equally come from a JSON file named by `MOONCAKE_CONFIG_PATH`, usin } ``` +Only `master_server_address` is required. `metadata_server` may be left out, in which case it is `P2PHANDSHAKE`, Mooncake's peer-to-peer handshake, which is also what `mooncake_store` and `mooncake_donation` default to; the example above names a metadata service instead. + An inherited `MOONCAKE_CONFIG_PATH` wins over `mooncake_store` and is logged as doing so, so an orchestrator that already provisions the pool, as the SLURM benchmark harness does, keeps working unchanged. Three further settings are TensorRT-LLM's rather than Mooncake's, and stay in the environment because they are per process rather than per pool: diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 008407106ff1..577bd22676dd 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -56,6 +56,11 @@ DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 DEFAULT_CACHE_PREFIX = "trtllm" DEFAULT_STAGING_BUFFER_SIZE = 536870912 +#: Mooncake's own peer-to-peer handshake, which keeps a separate metadata +#: process out of the deployment. Nothing else is a sensible fallback: an empty +#: connstring is not one of the forms `store.setup` accepts, so a config that +#: leaves the field out means this rather than meaning no metadata service. +DEFAULT_METADATA_SERVER = "P2PHANDSHAKE" _TRUE = {"1", "true", "yes", "on"} _FALSE = {"0", "false", "no", "off"} @@ -142,8 +147,8 @@ def provisioned_config_path() -> Optional[str]: class MooncakeStoreConnectorConfig: """Everything needed to open a store handle and name keys in it.""" - metadata_server: str master_server_address: str + metadata_server: str = DEFAULT_METADATA_SERVER protocol: str = "rdma" device_name: str = "" global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE @@ -188,8 +193,8 @@ def from_file(path: str) -> "MooncakeStoreConnectorConfig": with open(path) as handle: raw = json.load(handle) return MooncakeStoreConnectorConfig( - metadata_server=raw.get("metadata_server", ""), master_server_address=raw.get("master_server_address", ""), + metadata_server=raw.get("metadata_server") or DEFAULT_METADATA_SERVER, protocol=raw.get("protocol", "rdma"), device_name=raw.get("device_name", ""), global_segment_size=parse_size( diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py index 15514d69a996..83c010797d9e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -39,7 +39,7 @@ from tensorrt_llm.logger import logger -from .config import parse_size +from .config import DEFAULT_METADATA_SERVER, parse_size from .master import ( local_address, master_timeout, @@ -64,7 +64,7 @@ def donate_segment( segment_size: int, protocol: str = "rdma", device_name: str = "", - metadata_server: str = "", + metadata_server: str = DEFAULT_METADATA_SERVER, local_buffer_size: int = DEFAULT_DONOR_LOCAL_BUFFER_SIZE, hostname: Optional[str] = None, ) -> Iterator[str]: @@ -95,7 +95,7 @@ def donate_segment( f"as capacity only, no reads or writes: host={host} " f"segment_size={donated} ({segment_size} bytes) " f"protocol={protocol} device={device_name or '(none)'} " - f"metadata_server={metadata_server or '(none)'} " + f"metadata_server={metadata_server} " f"local_buffer_size={local_buffer_size} bytes" ) diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index aeb06315c3ca..e4cfc6a67371 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -198,7 +198,7 @@ def mooncake_master( "--metadata_server", type=str, default=None, - help="Mooncake metadata service. Defaults to --config's.", + help="Mooncake metadata service. Defaults to --config's, else P2PHANDSHAKE.", ) @click.option( "--ready_file", @@ -238,7 +238,10 @@ def mooncake_donor( resolve_master_address, wait_for_master, ) - from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import CONFIG_PATH_ENV + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + CONFIG_PATH_ENV, + DEFAULT_METADATA_SERVER, + ) raw = {} config = config or os.getenv(CONFIG_PATH_ENV) @@ -264,7 +267,7 @@ def mooncake_donor( donating, protocol=protocol or raw.get("protocol", "rdma"), device_name=device_name or raw.get("device_name", "") or "", - metadata_server=metadata_server or raw.get("metadata_server", ""), + metadata_server=(metadata_server or raw.get("metadata_server") or DEFAULT_METADATA_SERVER), local_buffer_size=parse_size( raw.get("local_buffer_size_donor", DEFAULT_DONOR_LOCAL_BUFFER_SIZE) ), diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 03cf093a1104..20396ed352da 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -442,6 +442,17 @@ def test_config_env_var_wins_over_the_run_directory(tmp_path, monkeypatch): assert MooncakeStoreConnectorConfig.from_env().master_server_address == "external:50051" +@pytest.mark.parametrize("named", [{}, {"metadata_server": ""}], ids=["omitted", "empty"]) +def test_config_metadata_server_falls_back_to_the_handshake(tmp_path, monkeypatch, named): + # No metadata service means Mooncake's peer-to-peer handshake. An empty + # connstring is not one of the forms setup accepts, so leaving the field + # out of a hand-written config must not reach it. + path = tmp_path / "metadata.json" + path.write_text(json.dumps({"master_server_address": "127.0.0.1:50051", **named})) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + assert MooncakeStoreConnectorConfig.from_env().metadata_server == "P2PHANDSHAKE" + + def test_config_model_key_defaults_to_basename(store_config, tmp_path, monkeypatch): path = tmp_path / "no_model_key.json" path.write_text(json.dumps({"master_server_address": "127.0.0.1:50051"}))