From ed49ee0767e1c07857f80cb5f7f85b1dea0af8e0 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 15 Sep 2026 14:39:19 -0700 Subject: [PATCH 01/24] [None][feat] Mooncake store part 1: pool, CLI, and V2 scheduler preemption Splits out the part of the Mooncake store integration that does not depend on KV connector support in KVCacheManagerV2, so it can be reviewed and merged without waiting on that work. The store side is complete: the pool master and its lifecycle, segment donation from nodes that run no connector, the JSON config, block hashing and key namespacing, and the pinned host slots pages pass through where GPUDirect RDMA is unavailable. `trtllm-serve` provisions the pool during bringup, and `mooncake_master` / `mooncake_donor` cover the parts of a pool that cannot belong to a server. The connector that moves KV pages in and out of the pool needs the KV cache layout description, and follows separately. Also adds preemption to the V2 scheduler, which is what a full pool falls back to when there is no cache tier below GPU to suspend into: suspended pages stay HELD and unevictable there, so suspension frees nothing. A victim gives its pages up and re-prefills. Alongside it, a deadlock detector fails loudly when consecutive scheduling passes can neither schedule nor reclaim anything, instead of spinning at full speed while looking healthy. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/common/install_mooncake.sh | 57 ++ .../attribution/scan/metadata/mooncake.yml | 5 +- tensorrt_llm/_torch/pyexecutor/_util.py | 2 + .../connectors/mooncake_store/__init__.py | 67 ++ .../connectors/mooncake_store/config.py | 268 +++++++ .../connectors/mooncake_store/donor.py | 171 +++++ .../connectors/mooncake_store/keys.py | 134 ++++ .../connectors/mooncake_store/master.py | 596 ++++++++++++++++ .../connectors/mooncake_store/metadata.py | 62 ++ .../connectors/mooncake_store/staging.py | 303 ++++++++ .../connectors/mooncake_store/validation.py | 60 ++ .../_torch/pyexecutor/connectors/registry.py | 27 + .../kv_cache/kv_cache_manager_v2.py | 39 ++ .../_torch/pyexecutor/py_executor_creator.py | 14 + .../pyexecutor/scheduler/scheduler_v2.py | 241 +++++-- tensorrt_llm/commands/mooncake.py | 294 ++++++++ tensorrt_llm/commands/serve.py | 144 ++-- tensorrt_llm/grpc/smg/server.py | 7 +- tensorrt_llm/llmapi/llm_args.py | 189 ++++- .../usage/llm_args_golden_manifest.json | 3 +- .../integration/test_lists/test-db/l0_a10.yml | 3 + .../kv_cache/test_kv_cache_v2_scheduler.py | 235 +++++++ .../executor/test_mooncake_store_common.py | 551 +++++++++++++++ .../executor/test_mooncake_store_donor.py | 201 ++++++ .../executor/test_mooncake_store_master.py | 654 ++++++++++++++++++ .../api_stability/references/llm.yaml | 4 + 26 files changed, 4228 insertions(+), 103 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py create mode 100644 tensorrt_llm/commands/mooncake.py create mode 100644 tests/unittest/_torch/executor/test_mooncake_store_common.py create mode 100644 tests/unittest/_torch/executor/test_mooncake_store_donor.py create mode 100644 tests/unittest/_torch/executor/test_mooncake_store_master.py diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index 0935c7f31a2f..a145b3756f62 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -53,3 +53,60 @@ cd ../.. rm -rf Mooncake echo "export LD_LIBRARY_PATH=${MOONCAKE_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" + +# 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` 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". +# +# - 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. +# +# 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 +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 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, 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}" + +# 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/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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f830d3c23191..93e662cb58ba 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3920,6 +3920,8 @@ def create_py_executor_instance( enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, # A disaggregated generation worker must not replay context locally. enable_recompute_pause=not is_disagg, + 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/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py new file mode 100644 index 000000000000..ee60e8d07b67 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -0,0 +1,67 @@ +# 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. +"""A Mooncake distributed store to back a KV cache connector. + +The store is a shared CPU memory pool addressed by content, 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, so a context server can write pages +here and still hand off over NIXL. + +The pool is described in `KvCacheConnectorConfig.mooncake_store`, which lets +`trtllm-serve` provision it during bringup 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, so `donor.py` lends +a node's memory to the pool without giving it a connector. Both need the +Mooncake Python bindings (`pip install mooncake-transfer-engine`). + +`keys.py` and `staging.py` hold what the store side shares with the connector +that moves pages in and out of the pool: how a block of tokens becomes a store +key, and how pages reach the fabric on hosts without GPUDirect RDMA. +""" + +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, +) + +__all__ = [ + "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", + "MooncakeStoreConnectorConfig", + "StoreRole", + "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 new file mode 100644 index 000000000000..577bd22676dd --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -0,0 +1,268 @@ +# 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__ = [ + "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" +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 +#: 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"} + +_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) + + +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 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, 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: + 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.""" + + master_server_address: str + metadata_server: str = DEFAULT_METADATA_SERVER + 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 + #: Pass pages through a pinned host buffer instead of registering the KV + #: 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. 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 + + 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") + 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": + """Read the topology from a vLLM-compatible Mooncake JSON config.""" + with open(path) as handle: + raw = json.load(handle) + return MooncakeStoreConnectorConfig( + 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( + 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)), + 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 + def from_env() -> "MooncakeStoreConnectorConfig": + """Load the JSON config, then apply the TensorRT-LLM env overrides.""" + 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), " + "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() + + 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 + 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: + """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/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py new file mode 100644 index 000000000000..4fe167a5bdaf --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -0,0 +1,171 @@ +# 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 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, 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 +import time +from typing import Any, Iterator, Optional + +from tensorrt_llm.logger import logger + +from .config import DEFAULT_METADATA_SERVER, 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", + "maybe_donate_segment", +] + +#: A donor never transfers, but `setup` rejects a zero-sized transfer buffer. +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 = DEFAULT_METADATA_SERVER, + 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 how the master + and the engines reading from it identify the capacity. + + 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 + 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() + 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( + 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} " + f"local_buffer_size={local_buffer_size} bytes" + ) + + store = MooncakeDistributedStore() + started = time.monotonic() + status = store.setup( + host, + metadata_server, + segment_size, + local_buffer_size, + protocol, + device_name, + master_server_address, + ) + elapsed = time.monotonic() - started + if status != 0: + raise RuntimeError( + 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: {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 + finally: + # The segment stays mounted while anything references the handle. + del store + 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, 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 + + # 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 " + "it; resolving the master now" + ) + master_address = resolve_master_address(donation.master_server_address, master_timeout()) + wait_for_master(master_address) + with donate_segment( + master_server_address=master_address, + segment_size=parse_size(donation.segment_size), + protocol=donation.protocol, + 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/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py new file mode 100644 index 000000000000..56609877843d --- /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, 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 +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/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py new file mode 100644 index 000000000000..f2f3a21a3c91 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -0,0 +1,596 @@ +# 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. + +`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 +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, List, Optional, Sequence, Tuple + +from tensorrt_llm.logger import logger + +from ..registry import uses_connector +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" +DEFAULT_MASTER_BINARY = "mooncake_master" +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 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 +#: rather than the address itself. +ADDRESS_FILE_SCHEME = "file://" +#: 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 + + +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: + """The address this host is known by inside the pool. + + 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()) + 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 resolve_master_address(address: str, timeout: float) -> str: + """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 + + path = address[len(ADDRESS_FILE_SCHEME) :] + started = time.monotonic() + deadline = started + timeout + announced = started + logger.info(f"mooncake-store: reading the master's address from {path}") + while True: + try: + published = open(path).read().strip() + except FileNotFoundError: + published = "" + if published: + logger.info(f"mooncake-store: {path} names the master at {published}") + return published + now = time.monotonic() + if now - announced >= 5.0: + announced = now + # 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 " + "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 " + 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, + timeout: float, + process: Optional[subprocess.Popen] = None, + 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 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 as it happens, since silence here is + indistinguishable from a hang elsewhere in bringup. + """ + 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} 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 time.monotonic() - started + except OSError as exc: + last_error = exc + 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." + 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) + + +#: 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 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] = {} + 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 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 + 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. + + 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. + """ + 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 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, + "master_server_address": master_address, + "protocol": pool.protocol, + "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", + "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 + # 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 + + +@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 " + "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." + ) + + host = local_address() + log_path = os.path.join(run_dir, MASTER_LOG_NAME) + + # 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 = [ + 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 + ) + 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: + 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} 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 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 never sees a partial address. + 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, 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 +) -> 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. 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 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) + try: + with _published_address(master.address, _address_files(run_dir, address_file)): + yield master + finally: + master.stop() + logger.info(f"mooncake-store: master at {master.address} stopped") + + +@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) + 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 + with contextlib.ExitStack() as stack: + try: + if pool.launch_master: + master = _launch_master(pool, run_dir) + master_address = master.address + # 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) + ) + ) + else: + 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) + # 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 + logger.info( + f"mooncake-store: {CONFIG_PATH_ENV}={config_path} " + f"({json.dumps(config, sort_keys=True)})" + ) + # 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; " + "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 +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, since such a deployment is told about its pool + through `MOONCAKE_CONFIG_PATH` instead. + """ + 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/_torch/pyexecutor/connectors/mooncake_store/metadata.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py new file mode 100644 index 000000000000..f1eb2db7ee14 --- /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/staging.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py new file mode 100644 index 000000000000..966897b178c4 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py @@ -0,0 +1,303 @@ +# 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, 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 + +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 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 + + +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 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: " + 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, since 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. + """ + + 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 + + # 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 + ) + 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/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py new file mode 100644 index 000000000000..4df451e42b5a --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -0,0 +1,60 @@ +# 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 tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +__all__ = ["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." + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py index 9a00cdadd7fd..dcf5e24f3642 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", @@ -35,4 +40,26 @@ "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", + }, } + + +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/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 371d7ed85327..e4e619c367af 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3775,6 +3775,45 @@ def resume_request(self, req: LlmRequest) -> bool: return False return self._resume_and_restore(req.py_request_id, kv_cache) + # ---- preemption ---- + # + # 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. + + @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 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, 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. + + Returns whether the pages were released. Callers must check: pages + still being read out of, such as by a connector save in flight, are + not released until that completes. + """ + self._release_preempted(req) + return True + + def _release_preempted(self, req: LlmRequest) -> None: + self.free_resources(req) + req.py_num_connector_matched_tokens = 0 + # ---- prepare_resources ---- def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 6204a1cd88bf..1747142b76d8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -47,6 +47,7 @@ resolve_cache_transceiver_config, uses_fp4_mla_attention, uses_vswa_kv_cache_layout) from .connectors.kv_cache_connector import KvCacheConnectorManager +from .connectors.registry import uses_connector from .dwdp import DwdpManager, get_global_dwdp_manager from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder from .hang_diagnostics import monitor_executor_initialization @@ -396,6 +397,19 @@ def _create_py_executor_impl( kv_cache_config.enable_block_reuse = False kv_cache_config.enable_partial_reuse = False + # 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( + "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 + # The tokenizer is stripped from MPI kwargs in proxy.py to avoid pickle # failures with trust_remote_code models. Reload it from the checkpoint # when guided decoding needs it. diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 2929fe386200..76877fc0db6b 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 @@ -168,8 +168,13 @@ def __init__( cross_kv_cache_manager: object | None = None, # KVCacheManagerV2 for enc-dec cross-attn enable_prefix_aware_scheduling: bool = True, enable_recompute_pause: 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, the original plus tokens generated so far, to this. + 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 ) @@ -443,6 +448,19 @@ def _schedule_loop(self, active_requests, inflight_request_ids): req_it += 1 + # 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: + 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. # @@ -468,6 +486,8 @@ def _schedule_loop(self, active_requests, inflight_request_ids): and not self._has_context_chunk_budget(budget) ): continue + if req.py_request_id in preempted_ids: + continue # Probe context requests before peft_pages_needed and before # _try_schedule_context so that a deferral costs nothing: KV pages # are allocated inline, so a skip decided after prepare_context @@ -499,7 +519,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: @@ -521,50 +543,20 @@ def _schedule_loop(self, active_requests, inflight_request_ids): if first_new_block is not None: contributed_blocks.add(first_new_block) - # 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 secondary 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 - and not self._has_pending_connector_load(r) - ) - if ( - num_gen_candidates > 0 - and not evicted - and not recompute_paused - and not inflight_request_ids - and not any(self._has_pending_connector_load(r) for r in active_requests) - ): - # A connector rejects every tier below GPU at bring-up - # (`PyExecutor._reject_non_gpu_cache_tiers`), so offering those - # two settings there is advice the user cannot act on. - if getattr(self.kv_cache_manager, "kv_connector_manager", None) is not None: - remedy = ( - "A KV connector is attached, which requires a GPU-only cache, " - "so no secondary tier can be configured. Increase " - "kv_cache_config.max_tokens or kv_cache_config." - "free_gpu_memory_fraction, or lower max_num_tokens to hand " - "memory back to the KV pool." - ) - else: - remedy = ( - "Configure kv_cache_config.host_cache_size or " - "kv_cache_config.disk_cache_size, or increase " - "kv_cache_config.max_tokens." - ) - raise RuntimeError( - f"V2 scheduler deadlock: {num_gen_candidates} generation " - f"request(s) active but none could be scheduled or " - f"evicted or recompute-paused. KV cache pool is likely exhausted with no " - f"secondary cache tier for suspend/resume offload. {remedy}" - ) + 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 + or recompute_paused + ), + ) return ( scheduled_encoder, @@ -734,7 +726,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). @@ -747,9 +742,9 @@ def _try_schedule_context( """ first_chunk = req.is_first_context_chunk if self.chunking_enabled: - result = self._try_schedule_context_chunked(req, budget) + result = self._try_schedule_context_chunked(req, budget, preempt_for_pages) else: - result = self._try_schedule_context_full(req, budget) + result = self._try_schedule_context_full(req, budget, preempt_for_pages) if first_chunk and result[0] is not ScheduleAction.SCHEDULED: # Failed admission must not retain prefix-reuse holds. Suspension @@ -766,7 +761,10 @@ def _try_schedule_context( return result 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. @@ -805,6 +803,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._try_allocate_context(req, context_tokens + draft_len): + # 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 cross_action = self._try_schedule_cross_context(req) @@ -825,7 +828,10 @@ def _has_context_chunk_budget(self, budget: BudgetTracker) -> bool: ) 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. @@ -883,6 +889,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 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( @@ -910,6 +919,8 @@ def _try_schedule_context_chunked( # V2 resizes KV cache directly in the scheduler, so include # draft tokens for last chunk. if not self._try_allocate_context(req, resize_tokens): + # Out of pages, as in _try_schedule_context_full. + preempt_for_pages(req) return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) @@ -1375,6 +1386,134 @@ 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. + + 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. + """ + if self.kv_cache_manager.has_cache_tier_below_gpu: + 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 + + self.kv_cache_manager.preempt_request(victim) + logger.debug( + f"[V2Scheduler] Preempting request {victim.py_request_id} " + f"(state={victim.state.name})" + ) + self._clear_request_runtime_state(victim) + if self.draft_kv_cache_manager is not None: + self.draft_kv_cache_manager.free_resources(victim) + # Rewrites the prompt to include what was generated and resets + # state to CONTEXT_INIT, so the request re-enters as an ordinary + # prefill. + victim.pause(self.max_input_len) + evicted.append(victim) + preempted_ids.add(victim.py_request_id) + return True + + return False + + # 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( + 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, 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 + 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 + and not self._has_pending_connector_load(r) + ) + 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 + + # A connector load in flight releases its pages when it lands, so a + # pass that reclaims nothing while one is outstanding is not a stall. + if any(self._has_pending_connector_load(r) for r in active_requests): + self._stalled_schedules = 0 + return + + self._stalled_schedules += 1 + if self._stalled_schedules < self._DEADLOCK_STALL_ITERS: + return + + # A connector rejects every tier below GPU at bring-up + # (`PyExecutor._reject_non_gpu_cache_tiers`), so offering host_cache_size + # there is advice the user cannot act on. + if getattr(self.kv_cache_manager, "kv_connector_manager", None) is not None: + remedy = ( + "A KV connector is attached, which requires a GPU-only cache, " + "so no secondary tier can be configured. Increase " + "kv_cache_config.max_tokens or " + "kv_cache_config.free_gpu_memory_fraction, or lower " + "max_num_tokens to hand memory back to the KV pool." + ) + else: + remedy = ( + "Configure kv_cache_config.host_cache_size, increase " + "kv_cache_config.max_tokens, or lower max_batch_size." + ) + 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. {remedy}" + ) + def _free_kv_caches(self, req: LlmRequest) -> None: self.kv_cache_manager.free_resources(req) if self.draft_kv_cache_manager is not None: diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py new file mode 100644 index 000000000000..37d3cf0a6846 --- /dev/null +++ b/tensorrt_llm/commands/mooncake.py @@ -0,0 +1,294 @@ +# 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: it describes the pool in +`kv_connector_config.mooncake_store` and `trtllm-serve` provisions it during +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 +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 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() + + 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.", +) +@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 + `mooncake_store.launch_master` instead. + """ + # 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 + + 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: + 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: + # 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 {master.log_path}" + ) + stopping.wait(1.0) + now = time.monotonic() + # Distinguishes a dead master from a dead 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") +@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, else P2PHANDSHAKE.", +) +@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_METADATA_SERVER, + ) + + 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) + resolved = resolve_master_address(master, master_timeout()) + wait_for_master(resolved) + + 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") or DEFAULT_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" + ) + + # 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: + stopping.wait() + continue + 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" + ) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 2bf9ae8f9a5a..7958203ce554 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -32,8 +32,8 @@ from importlib.util import find_spec from pathlib import Path from types import FrameType -from typing import (TYPE_CHECKING, Any, Dict, NamedTuple, NoReturn, Optional, - Sequence, Set) +from typing import (TYPE_CHECKING, Any, Dict, Iterator, NamedTuple, NoReturn, + Optional, Sequence, Set) import click import torch @@ -47,6 +47,7 @@ from tensorrt_llm._utils import mpi_rank, set_prometheus_multiproc_dir from tensorrt_llm.commands import _telemetry as _command_telemetry 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 @@ -58,7 +59,9 @@ 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, + 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 from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory, @@ -620,6 +623,49 @@ def _terminate_attached_frontends(children: list) -> None: child.kill() +@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. + + 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 pool and lend a second segment. + """ + if not owns_engine: + yield + return + + 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): + # 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 + + 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( host: str, port: int, @@ -701,48 +747,50 @@ def launch_server( # until uvicorn takes it over, so no one can steal the port in between. _publish_bound_address(report_addr, host, port) - if backend == 'pytorch': - llm_args.pop("build_config", None) - llm = PyTorchLLM(**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, - internal_disagg_auth_key=internal_disagg_auth_key) - _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_cache_pool( + llm_args, owns_engine=not multi_frontend.is_attached_frontend): + if backend == 'pytorch': + llm_args.pop("build_config", None) + llm = PyTorchLLM(**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, + internal_disagg_auth_key=internal_disagg_auth_key) + _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_mm_encoder_server( @@ -2745,7 +2793,11 @@ 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 + # deployments where a pool outlives or spans them. + "mooncake_master": mooncake_master, + "mooncake_donor": mooncake_donor, }) if __name__ == "__main__": diff --git a/tensorrt_llm/grpc/smg/server.py b/tensorrt_llm/grpc/smg/server.py index 8e891ac6556b..a7b3bfe55d7a 100644 --- a/tensorrt_llm/grpc/smg/server.py +++ b/tensorrt_llm/grpc/smg/server.py @@ -126,7 +126,12 @@ def signal_handler() -> None: logger.info("LLM engine stopped") logger.info("Shutdown complete") - uvloop.run(serve_grpc_async()) + # Imported here rather than at module scope: tensorrt_llm.commands.serve + # reaches into this module to launch the server. + from tensorrt_llm.commands.serve import _provision_kv_cache_pool + + with _provision_kv_cache_pool(llm_args): + uvloop.run(serve_grpc_async()) __all__ = ["launch_smg_server"] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d596c1d38a4e..2e3d4569e7c1 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2303,6 +2303,166 @@ 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, 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. 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. 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.") + 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_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, + 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.") + 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": + 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.") + 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, 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( + ..., + 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. @@ -2317,7 +2477,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= @@ -2332,11 +2493,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: @@ -2354,6 +2520,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 @@ -6024,6 +6196,17 @@ def is_partial_model_loading(self) -> bool: 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/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 257d1ba70c11..07e67f75ce52 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -781,7 +781,8 @@ "allowed_values": [ "lmcache", "lmcache-mp", - "kvbm" + "kvbm", + "mooncake-store" ], "capture_policy": "allowlist|none", "kind": "categorical", diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index b35205380b5f..eb5fce771dfe 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -41,6 +41,9 @@ l0_a10: - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_kv_cache_layout.py - unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py + - unittest/_torch/executor/test_mooncake_store_common.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/executor/test_resource_manager.py - unittest/_torch/executor/test_seq_slot_sizing.py diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 7a941a960796..09e05c0c00a0 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -57,6 +57,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 req.py_multimodal_data = None @@ -172,6 +173,7 @@ def make_kv_cache_manager( block_reuse_policy=BlockReusePolicy.PER_REQUEST, first_new_block_fn=None, is_vswa=False, + has_cache_tier_below_gpu=True, ): mgr = Mock() mgr.kv_connector_manager = None @@ -216,6 +218,9 @@ def suspend_request(req): mgr.suspend_request.side_effect = suspend_request 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 + mgr.preempt_request.side_effect = lambda req: True return mgr @@ -239,6 +244,7 @@ def make_scheduler( cross_kv_cache_manager: Mock | None = None, enable_prefix_aware_scheduling: bool = True, enable_recompute_pause: 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 @@ -256,6 +262,8 @@ def make_scheduler( kwargs["draft_kv_cache_manager"] = draft_kv_cache_manager 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, @@ -1285,6 +1293,233 @@ 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. + + 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): + 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_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 the request is + # simply skipped. + assert ids(out.context_requests) == [99] + + 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. + """ + + 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, can_evict=True) + 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 # =========================================================================== diff --git a/tests/unittest/_torch/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py new file mode 100644 index 000000000000..8692a594d011 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -0,0 +1,551 @@ +# 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 pieces every deployment shares. + +Covers how a block of tokens becomes a store key, how the JSON config is read, +which deployments are refused at startup, and the pinned host slots pages pass +through where GPUDirect RDMA is unavailable. + +Runs without a Mooncake installation: the store handle is replaced by an +in-process fake that records what it was handed. +""" + +import json +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import staging as staging_module +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 ( + MooncakeStoreMetadata, + PageTransfer, + RequestTransfers, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.staging import ( + HostStagingPool, + describe_batch_for_get, + plan_slot_geometry, + stage_batch_for_put, + unstage_batch_after_get, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.validation import validate_llm_args +from tensorrt_llm._torch.pyexecutor.connectors.registry import uses_connector +from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig + +TOKENS_PER_BLOCK = 4 + +#: Device addresses the staging tests gather from and scatter back to. Never +#: dereferenced: the copies are recorded rather than issued. +PAGE_ADDRESSES = [0xA000, 0xB000] +PAGE_SIZES = [64, 128] +PAGE_BYTES = sum(PAGE_SIZES) + + +# ---- fixtures and fakes ---- + + +class FakeStore: + """Records buffer registrations, nothing more. + + Staging only ever asks the store to register its host buffer; the transfer + calls belong to the connector. + """ + + def __init__(self, register_status=0): + self.registered = [] + self._register_status = register_status + + def register_buffer(self, address, size): + self.registered.append((address, size)) + return self._register_status + + +@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 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))), + ) + return copies + + +def make_pool(*, slot_bytes=PAGE_BYTES, num_slots=4, store=None): + return HostStagingPool( + slot_bytes=slot_bytes, + num_slots=num_slots, + store=store if store is not None else FakeStore(), + label="save", + ) + + +# ---- 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 + + +# ---- 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): + 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) + 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 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) + 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, 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( + 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" + + +@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"})) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + config = MooncakeStoreConnectorConfig.from_env() + assert config.resolve_model_key("/models/MiniMax-M3/") == "MiniMax-M3" + + +@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() + + +# ---- 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) + + +# ---- connector identification ---- +# +# 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. + + +@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(): + config = KvCacheConnectorConfig(connector="mooncake-store") + with pytest.raises(ValueError, match="Unknown connector preset"): + uses_connector(config, "mooncake-stroe") + + +# ---- host staging ---- + + +@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_staging_pool_registers_its_whole_buffer_once(): + """One registration covering every slot, since the store addresses bytes.""" + store = FakeStore() + pool = make_pool(slot_bytes=256, num_slots=4, store=store) + + assert pool.slot_bytes == 256 + assert pool.num_slots == 4 + assert store.registered == [(pool.slot_address(0), 256 * 4)] + + +def test_staging_pool_refuses_to_start_when_registration_fails(): + # Host memory that cannot be registered is the one failure staging exists to + # rule out, so it has to be loud rather than fall back to the pools. + with pytest.raises(RuntimeError, match="register_buffer failed"): + make_pool(store=FakeStore(register_status=-1)) + + +def test_staging_pool_slots_are_contiguous_and_bounds_checked(): + pool = make_pool(slot_bytes=256, num_slots=4) + base = pool.slot_address(0) + + assert [pool.slot_address(i) for i in range(4)] == [base + 256 * i for i in range(4)] + with pytest.raises(IndexError): + pool.slot_address(4) + with pytest.raises(IndexError): + pool.slot_address(-1) + + +def test_staging_pool_gather_concatenates_regions_in_region_order(staged_copies): + """The slot layout is what the zero-copy path would have written.""" + pool = make_pool(num_slots=4) + slot = pool.slot_address(1) + + address, total = pool.gather(1, PAGE_ADDRESSES, PAGE_SIZES, stream=0) + + assert (address, total) == (slot, PAGE_BYTES) + assert staged_copies == [ + (slot, PAGE_ADDRESSES[0], PAGE_SIZES[0]), + (slot + PAGE_SIZES[0], PAGE_ADDRESSES[1], PAGE_SIZES[1]), + ] + + +def test_staging_pool_scatter_is_the_inverse_of_gather(staged_copies): + pool = make_pool(num_slots=4) + slot = pool.slot_address(1) + + pool.scatter(1, PAGE_ADDRESSES, PAGE_SIZES, stream=0) + + # Same split in the same order, with source and destination swapped. + assert staged_copies == [ + (PAGE_ADDRESSES[0], slot, PAGE_SIZES[0]), + (PAGE_ADDRESSES[1], slot + PAGE_SIZES[0], PAGE_SIZES[1]), + ] + + +@pytest.mark.parametrize("operation", ["gather", "scatter", "reserve"]) +def test_staging_pool_rejects_a_page_wider_than_a_slot(staged_copies, operation): + # The pool is sized from the layout's largest page, so an overflow means the + # layout changed after registration rather than a bad argument. + pool = make_pool(slot_bytes=PAGE_BYTES - 1) + + with pytest.raises(ValueError, match="does not fit"): + if operation == "reserve": + pool.reserve(PAGE_BYTES) + else: + getattr(pool, operation)(0, PAGE_ADDRESSES, PAGE_SIZES, stream=0) + + +def test_stage_batch_for_put_hands_the_store_one_buffer_per_page(staged_copies): + pool = make_pool(num_slots=4) + addresses, sizes = stage_batch_for_put( + pool, [PAGE_ADDRESSES, PAGE_ADDRESSES], [PAGE_SIZES, PAGE_SIZES], stream=0 + ) + + # Two regions per page collapse into the one slot the store reads. + assert addresses == [[pool.slot_address(0)], [pool.slot_address(1)]] + assert sizes == [[PAGE_BYTES], [PAGE_BYTES]] + assert len(staged_copies) == 4 + + +def test_describe_batch_for_get_points_at_slots_without_copying(staged_copies): + """Nothing to gather in this direction: the slots are the destination.""" + pool = make_pool(num_slots=4) + + addresses, sizes = describe_batch_for_get(pool, [PAGE_SIZES, PAGE_SIZES]) + + assert addresses == [[pool.slot_address(0)], [pool.slot_address(1)]] + assert sizes == [[PAGE_BYTES], [PAGE_BYTES]] + assert staged_copies == [] + + +@pytest.mark.parametrize("helper", ["put", "get"]) +def test_batch_helpers_reject_a_batch_wider_than_the_pool(staged_copies, helper): + pool = make_pool(num_slots=1) + batch_sizes = [PAGE_SIZES, PAGE_SIZES] + + with pytest.raises(ValueError, match="exceeds 1 staging slots"): + if helper == "put": + stage_batch_for_put(pool, [PAGE_ADDRESSES, PAGE_ADDRESSES], batch_sizes, stream=0) + else: + describe_batch_for_get(pool, batch_sizes) + + +def test_unstage_batch_after_get_scatters_every_page_by_default(staged_copies): + pool = make_pool(num_slots=4) + pages = [[0xA000, 0xB000], [0xC000, 0xD000]] + + unstage_batch_after_get(pool, pages, [PAGE_SIZES, PAGE_SIZES], stream=0) + + assert [copy[0] for copy in staged_copies] == [0xA000, 0xB000, 0xC000, 0xD000] + + +def test_unstage_batch_after_get_leaves_the_pages_not_asked_for_alone(staged_copies): + """A page whose read failed keeps its device bytes rather than slot garbage.""" + pool = make_pool(num_slots=4) + pages = [[0xA000, 0xB000], [0xC000, 0xD000]] + + unstage_batch_after_get(pool, pages, [PAGE_SIZES, PAGE_SIZES], stream=0, only=[1]) + + assert [copy[0] for copy in staged_copies] == [0xC000, 0xD000] + + +# ---- metadata ---- + + +def test_metadata_is_falsy_until_there_is_work(): + """The worker skips the iteration entirely on an empty work list.""" + assert not MooncakeStoreMetadata() + assert MooncakeStoreMetadata(loads=[RequestTransfers(request_id=1)]) + assert MooncakeStoreMetadata(saves=[RequestTransfers(request_id=1)]) + + +def test_request_transfers_do_not_share_a_page_list(): + first = RequestTransfers(request_id=1) + second = RequestTransfers(request_id=2) + + first.pages.append(PageTransfer(block_hash=b"\x01" * 16, layer_group_id=0, page_index=3)) + + assert second.pages == [] 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..cda334bc385a --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -0,0 +1,201 @@ +# 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 under test is what +the donor asks Mooncake for and how long it holds it. +""" + +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, + maybe_donate_segment, +) +from tensorrt_llm.llmapi.llm_args import MooncakeDonationConfig + +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 + + +@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", + ) + 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, + segment_size, + local_buffer_size, + protocol, + device_name, + master, + ) = fake_bindings.instances[0].setup_args + + assert registered_host == expected_host + assert metadata_server == "P2PHANDSHAKE" + assert segment_size == expected_size + assert protocol == "rdma" + assert device_name == expected_device + assert master == "10.0.0.1:50051" + 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") + + +@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_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_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 new file mode 100644 index 000000000000..618bcfbb8512 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -0,0 +1,654 @@ +# 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 +import threading +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, + resolve_master_address, +) +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.pid = 4242 + 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. 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, 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: + stdout.write(log_text.encode()) + stdout.flush() + 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 ---- + + +@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(): + """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), + ) + # 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), + ) + + +# ---- 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_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") + written = master_module._client_config(pool, "host:50051") + assert "cache_prefix" not in written + assert "staging_buffer_bytes" not in written + + +@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_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_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, so an unprobeable one passes through.""" + 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): + """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)) + 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 + # 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 + + 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 + # 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" + + +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 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() + + +# ---- the entry point servers call ---- + + +@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.""" + 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 + + +# ---- 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): + """An existing file is not the same as a published address.""" + 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 to 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): + """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" + 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 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" + 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 + + +# ---- a master a server launched, made findable ---- + + +def test_a_launched_master_publishes_where_its_run_left_its_logs(fake_master, tmp_path): + """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): + """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" + 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 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()}" + + with pytest.raises(TimeoutError, match=address): + master_module.wait_for_master(address) + + +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 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") + 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): + """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")) == "" + + +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" diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 8728888a2868..9ecbbc61e23b 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -291,6 +291,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 dbbdfff43ce004851c595c201338617e4472a690 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Fri, 18 Sep 2026 19:37:16 -0700 Subject: [PATCH 02/24] [None][chore] Address comments from CodeRabbit HostStagingPool now keeps the store handle and has an idempotent close() that unregisters the staging buffer before it is freed. The registration is against an address range, so a buffer released while the store still holds one leaves the fabric pointing at reusable memory. A failed unregistration keeps the buffer alive rather than freeing it. The mooncake_store package now exports MooncakeStoreConnectorScheduler and MooncakeStoreConnectorWorker. Both refuse construction with a message naming what this build can do with a pool, so the registry preset fails as a configuration this build cannot serve rather than as an AttributeError from the loader, after the provisioned master is already up. The OpenEngine gRPC path wraps its launch in _provision_kv_cache_pool, as launch_server and launch_smg_server already do. Without it a deployment that set mooncake_store or mooncake_donation and served over that path got no master, no client config, and no donated segment. MooncakeStoreConfig.global_segment_size and local_buffer_size opt out of telemetry. They were the only capturable fields left in that model and were missing from the golden manifest, so the committed manifest was stale. They size and address one site's pool; the connector field already records that the store is on. mooncake_donor takes its transfer buffer from a --local_buffer_size flag instead of a local_buffer_size_donor config key that nothing writes. The shared local_buffer_size is sized for an engine worker, and a donor never transfers through its buffer. Tests: real preemption against a full GPU-only pool, exercising _release_preempted rather than a stub; staging close() and its failure path; the PREFIX and MODEL_KEY environment overrides that feed KeyNamespace; the registered preset resolving to classes that report themselves unavailable. The store_config fixture also clears TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST so it no longer depends on the developer or CI environment. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/__init__.py | 4 ++ .../connectors/mooncake_store/connector.py | 54 +++++++++++++++ .../connectors/mooncake_store/staging.py | 21 ++++++ .../_torch/pyexecutor/connectors/registry.py | 1 + tensorrt_llm/commands/mooncake.py | 14 +++- tensorrt_llm/commands/serve.py | 14 ++-- tensorrt_llm/llmapi/llm_args.py | 6 ++ .../kv_cache/test_kvcm2_integration.py | 68 +++++++++++++++++++ .../executor/test_mooncake_store_common.py | 55 ++++++++++++++- 9 files changed, 230 insertions(+), 7 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index ee60e8d07b67..be28d0590594 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -34,9 +34,11 @@ `keys.py` and `staging.py` hold what the store side shares with the connector that moves pages in and out of the pool: how a block of tokens becomes a store key, and how pages reach the fabric on hosts without GPUDirect RDMA. +`connector.py` holds the connector classes, which are placeholders. """ from .config import MooncakeStoreConnectorConfig, StoreRole, parse_size +from .connector import MooncakeStoreConnectorScheduler, MooncakeStoreConnectorWorker from .donor import DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, maybe_donate_segment from .master import ( local_address, @@ -52,6 +54,8 @@ __all__ = [ "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", "MooncakeStoreConnectorConfig", + "MooncakeStoreConnectorScheduler", + "MooncakeStoreConnectorWorker", "StoreRole", "donate_segment", "local_address", diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py new file mode 100644 index 000000000000..2e5dc2e66cc0 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py @@ -0,0 +1,54 @@ +# 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. +"""Placeholders for the connector that moves KV pages in and out of the pool. + +The store side stands on its own: a pool can be provisioned, lent memory, and +shared with an engine elsewhere without a connector here. Moving pages needs the +KV cache layout description, so neither class is implemented here. + +`CONNECTOR_REGISTRY` names both and `py_executor_creator` resolves them by name, +so they exist to turn `connector: mooncake-store` into a clear refusal instead +of an `AttributeError` raised once the pool is already provisioned. + +Neither subclasses `KvCacheConnectorWorker` or `KvCacheConnectorScheduler`, +whose `__init_subclass__` checks a method set there is nothing here to satisfy. +""" + +__all__ = ["MooncakeStoreConnectorScheduler", "MooncakeStoreConnectorWorker"] + +_UNAVAILABLE = ( + "The mooncake-store KV cache connector is not available in this build, so an " + "engine here cannot read or write the pool. The pool itself is: " + "'trtllm-serve mooncake_master' owns one, 'trtllm-serve mooncake_donor' and " + "the mooncake_donation setting lend it host memory, and an engine that does " + "have the connector can share it. Drop kv_connector_config.connector to " + "serve without one." +) + + +class _Unavailable: + """Reports what this build can and cannot do with a Mooncake pool.""" + + def __init__(self, llm_args): + del llm_args + raise NotImplementedError(_UNAVAILABLE) + + +class MooncakeStoreConnectorWorker(_Unavailable): + """Moves this rank's KV pages between its cache and the pool.""" + + +class MooncakeStoreConnectorScheduler(_Unavailable): + """Decides which pages a request loads from the pool and saves to it.""" diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py index 966897b178c4..98e3b61a0a62 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py @@ -127,6 +127,7 @@ def __init__( ): self._slot_bytes = int(slot_bytes) self._num_slots = int(num_slots) + self._store = store self._label = label # Page-locking is a correctness requirement here rather than a @@ -152,6 +153,26 @@ def __init__( f"(pinned={pin})" ) + def close(self) -> None: + """Hand the registration back before the buffer is freed. + + The store registers an address range rather than the tensor, so a buffer + freed while it still holds one leaves the fabric able to reach memory the + allocator has since handed out again. Called once the connector's pending + transfers have drained. Idempotent, and a failed unregistration keeps the + buffer rather than freeing memory the store may still reach. + """ + if self._buffer is None: + return + status = self._store.unregister_buffer(self._base) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.unregister_buffer failed with status " + f"{status} for the {self._label} host staging buffer at " + f"{self._base:#x}. The buffer is kept alive rather than freed." + ) + self._buffer = None + @property def num_slots(self) -> int: """Pages this pool can hold at once.""" diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py index dcf5e24f3642..c0a40bf480a6 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py @@ -40,6 +40,7 @@ "connector_scheduler_class": "DynamoKVBMConnectorLeader", "connector_worker_class": "DynamoKVBMConnectorWorker", }, + # Both classes currently refuse construction; see the module's docstring. "mooncake-store": { "connector_module": "tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", "connector_scheduler_class": "MooncakeStoreConnectorScheduler", diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index 37d3cf0a6846..de779107f4b3 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -199,6 +199,15 @@ def mooncake_master( default=None, help="Mooncake metadata service. Defaults to --config's, else P2PHANDSHAKE.", ) +@click.option( + "--local_buffer_size", + type=str, + default=None, + help="Mooncake transfer buffer for this process. Deliberately " + "separate from a config's local_buffer_size, which is sized for " + "an engine worker: a donor never transfers, and only needs one " + "because setup rejects a zero-sized buffer. Defaults to 64MiB.", +) @click.option( "--ready_file", type=str, @@ -221,6 +230,7 @@ def mooncake_donor( protocol: Optional[str], device_name: Optional[str], metadata_server: Optional[str], + local_buffer_size: Optional[str], ready_file: Optional[str], heartbeat_seconds: int, ): @@ -266,8 +276,8 @@ def mooncake_donor( 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") or DEFAULT_METADATA_SERVER), - local_buffer_size=parse_size( - raw.get("local_buffer_size_donor", DEFAULT_DONOR_LOCAL_BUFFER_SIZE) + local_buffer_size=( + parse_size(local_buffer_size) if local_buffer_size else DEFAULT_DONOR_LOCAL_BUFFER_SIZE ), ) as host: if ready_file: diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 7958203ce554..37d136d3223a 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1602,10 +1602,16 @@ def _serve_llm(): "https://buf.build/gen/python " "\"tensorrt_llm[openengine]\"`.") from error - launch_grpc_server(host, - port, - llm_args, - served_model_name=served_model_name) + # launch_smg_server provisions from inside itself; OpenEngine's + # server is an optional package this repo does not own, so the + # pool is brought up around it here instead. Either way the + # context has to outlive engine construction, which happens + # inside the launch call. + with _provision_kv_cache_pool(llm_args): + launch_grpc_server(host, + port, + llm_args, + served_model_name=served_model_name) else: # Default: launch OpenAI HTTP server launch_server( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 2e3d4569e7c1..f44032cd1d3f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2315,6 +2315,10 @@ class MooncakeStoreConfig(StrictBaseModel): Setting this makes `trtllm-serve` render the Mooncake client config and export `MOONCAKE_CONFIG_PATH` itself. An inherited `MOONCAKE_CONFIG_PATH` still wins, so an externally managed pool stays reachable. + + Fields opt out of telemetry because they size and address one site's pool + rather than saying which features are in use; `kv_connector_config.connector` + already records that the store is on. """ master_server_address: Optional[str] = Field( None, @@ -2365,11 +2369,13 @@ class MooncakeStoreConfig(StrictBaseModel): "with protocol 'tcp'.") global_segment_size: Union[int, str] = Field( "16GiB", + telemetry=False, 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", + telemetry=False, description="Per-process Mooncake transfer buffer, not pool capacity.") transfer_batch_size: int = Field(64, telemetry=False, diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index fbb26fedafd0..1f3b77e7b170 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -1397,6 +1397,7 @@ class _ContextRequest: return_perf_metrics: bool = False context_current_position: int = 0 py_connector_served_position: int = 0 + py_num_connector_matched_tokens: int = 0 prepopulated_prompt: tuple[int, int] | None = None multimodal_hashes: None = None multimodal_positions: None = None @@ -1595,6 +1596,73 @@ def _run_context( _update_context_resources(manager, batch) +def _try_run_context(manager: KVCacheManagerV2, request: _ContextRequest) -> bool: + """Run a context request all the way through, reporting whether it fit.""" + batch = _prepare_context_resources(manager, request) + if not manager.prepare_context(request): + return False + request.context_remaining_length = request.prompt_len - request.context_current_position + if not manager.resize_context(request, num_tokens=request.context_remaining_length): + return False + request.context_current_position = request.prompt_len + request.context_remaining_length = 0 + _update_context_resources(manager, batch) + return True + + +def test_preempt_request_gives_a_full_pool_back_its_pages( + manager: KVCacheManagerV2, +) -> None: + """Preemption is the only way out of a full GPU-only pool. + + With GPU as the last cache level a suspended page stays HELD, which the + eviction controller refuses to move, so suspension frees nothing. This + covers the release itself: the pages returning, and the connector's + matched-token count dropping. + """ + victim = _ContextRequest(1, list(range(MAX_SEQ_LEN)), MAX_SEQ_LEN, "conv-1") + victim.py_num_connector_matched_tokens = TOKENS_PER_BLOCK + started: list[_ContextRequest] = [victim] + + try: + assert _try_run_context(manager, victim) + assert manager.is_request_active(victim.py_request_id) + + # The pool's capacity in requests follows from the fixture's layout and + # windows, so find it rather than hard-code it. Distinct tokens per + # request keep block reuse from hiding the pressure. + blocked = None + for request_id in range(2, 12): + candidate = _ContextRequest( + request_id, + [request_id * 1000 + i for i in range(MAX_SEQ_LEN)], + MAX_SEQ_LEN, + f"conv-{request_id}", + ) + started.append(candidate) + if not _try_run_context(manager, candidate): + blocked = candidate + break + assert blocked is not None, "the pool never filled, so there is nothing to preempt for" + + # Its own partial allocation goes back first, as the scheduler's retry + # does. What is missing is the victim's pages. + manager.free_resources(blocked) + + assert manager.preempt_request(victim) is True + assert not manager.is_request_active(victim.py_request_id) + assert victim.py_request_id not in manager.kv_cache_map + # A replay recomputes from its own reuse match, so a stale count would + # skip a prefix the pool no longer holds for it. + assert victim.py_num_connector_matched_tokens == 0 + + blocked.context_current_position = 0 + assert _try_run_context(manager, blocked) + finally: + for request in started: + _free_if_active(manager, request) + + def test_per_conversation_policy_delays_commit_until_last_context_chunk( manager: KVCacheManagerV2, ) -> None: diff --git a/tests/unittest/_torch/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py index 8692a594d011..9e02826d32e6 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_common.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -22,6 +22,7 @@ in-process fake that records what it was handed. """ +import importlib import json from types import SimpleNamespace @@ -71,14 +72,20 @@ class FakeStore: calls belong to the connector. """ - def __init__(self, register_status=0): + def __init__(self, register_status=0, unregister_status=0): self.registered = [] + self.unregistered = [] self._register_status = register_status + self._unregister_status = unregister_status def register_buffer(self, address, size): self.registered.append((address, size)) return self._register_status + def unregister_buffer(self, address): + self.unregistered.append(address) + return self._unregister_status + @pytest.fixture def store_config(tmp_path, monkeypatch): @@ -100,6 +107,7 @@ def store_config(tmp_path, monkeypatch): 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) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST", raising=False) return path @@ -317,6 +325,20 @@ def test_config_staging_env_override(store_config, monkeypatch, value, expected) assert MooncakeStoreConnectorConfig.from_env().stage_through_host is expected +def test_config_namespace_env_overrides(store_config, monkeypatch): + # Both feed KeyNamespace, which decides whether two engines share cache. + # Getting either wrong silently loses every hit, or lets one engine read + # another's pages. + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_PREFIX", "tenant-a") + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_MODEL_KEY", "llama-3.1-8b@rev7") + + config = MooncakeStoreConnectorConfig.from_env() + + assert config.cache_prefix == "tenant-a" + # Overrides the JSON's model_key, not just the basename default. + assert config.resolve_model_key("/models/test-model") == "llama-3.1-8b@rev7" + + 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"): @@ -384,6 +406,16 @@ def test_uses_connector_rejects_an_unknown_preset(): uses_connector(config, "mooncake-stroe") +def test_the_registered_preset_resolves_to_refusing_classes(): + """py_executor_creator resolves both by name from the package.""" + config = KvCacheConnectorConfig(connector="mooncake-store") + module = importlib.import_module(config.connector_module) + + for class_name in (config.connector_scheduler_class, config.connector_worker_class): + with pytest.raises(NotImplementedError, match="not available in this build"): + getattr(module, class_name)(llm_args=None) + + # ---- host staging ---- @@ -427,6 +459,27 @@ def test_staging_pool_refuses_to_start_when_registration_fails(): make_pool(store=FakeStore(register_status=-1)) +def test_staging_pool_close_unregisters_before_the_buffer_goes(): + store = FakeStore() + pool = make_pool(slot_bytes=256, num_slots=4, store=store) + base = pool.slot_address(0) + + pool.close() + pool.close() + + assert store.unregistered == [base] + + +def test_staging_pool_close_keeps_the_buffer_when_unregistration_fails(): + store = FakeStore(unregister_status=-1) + pool = make_pool(slot_bytes=256, num_slots=4, store=store) + + with pytest.raises(RuntimeError, match="unregister_buffer failed"): + pool.close() + + assert pool._buffer is not None + + def test_staging_pool_slots_are_contiguous_and_bounds_checked(): pool = make_pool(slot_bytes=256, num_slots=4) base = pool.slot_address(0) From 502cda79292e277a9226a32a7f3189aa548aca04 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Fri, 18 Sep 2026 19:37:41 -0700 Subject: [PATCH 03/24] [None][fix] Address comments from Thor Preemption pauses the victim in place, which resets mSeqSlot without telling SeqSlotManager: slot_mapping still holds the slot while py_seq_slot is None, so the victim's next schedule reaches SeqSlotManager.prepare_resources, calls add_slot, and trips its request-id assertion. pause() also rewrites mPromptLen and mMaxNewTokens without resyncing py_prompt_len, py_max_new_tokens, py_seq_slot, draft tokens or py_decoding_iter, and _prepare_tp_inputs feeds py_prompt_len straight into the attention metadata. Both follow from the scheduler owning a teardown it only did half of. _try_preempt_for_pages now hands the victim to the executor on recompute_paused, the channel the generation side already uses, so _terminate_recompute_paused_requests frees the rest of the request's resources through the resource manager, sequence slot included, and reset_for_recompute rewrites the prompt and resyncs the Python-side mirrors of it. The KV release stays inline, since the blocked request has to allocate in this same pass. Two consequences. Victim selection uses _is_recompute_pause_candidate, so it also skips GENERATION_TO_COMPLETE and generation requests whose multimodal replay data was released, both of which a re-prefill needs. And preemption is gated on enable_recompute_pause: a disaggregated generation worker received its context KV rather than computing it, so it cannot replay a prefill at all, and a full pool there is a deadlock the detector should report rather than something to thrash against. max_input_len has no reader left, so it is gone from the scheduler along with the _util.py plumbing that fed it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 2 -- .../pyexecutor/scheduler/scheduler_v2.py | 30 +++++++++--------- .../kv_cache/test_kv_cache_v2_scheduler.py | 31 +++++++++++++------ 3 files changed, 38 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 93e662cb58ba..f830d3c23191 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3920,8 +3920,6 @@ def create_py_executor_instance( enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, # A disaggregated generation worker must not replay context locally. enable_recompute_pause=not is_disagg, - 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/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 76877fc0db6b..a622b58336f8 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -168,12 +168,8 @@ def __init__( cross_kv_cache_manager: object | None = None, # KVCacheManagerV2 for enc-dec cross-attn enable_prefix_aware_scheduling: bool = True, enable_recompute_pause: 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, the original plus tokens generated so far, to this. - 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 @@ -458,7 +454,7 @@ def preempt_for_pages(req: LlmRequest) -> bool: 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 + requests_list, protected, inflight_request_ids, recompute_paused, preempted_ids ) # --- Phase 2: schedule deferred context / encoder requests --- @@ -1391,7 +1387,7 @@ def _try_preempt_for_pages( requests_list: RequestList, protected_ids: set[int], inflight_request_ids: set[int], - evicted: RequestList, + recompute_paused: RequestList, preempted_ids: set[int], ) -> bool: """Release one started request's KV cache so another can allocate. @@ -1400,10 +1396,22 @@ def _try_preempt_for_pages( `KVCacheManagerV2.preempt_request`. With a cache tier below GPU, suspension is cheaper and keeps the pages, so that path is left alone. + The victim leaves on `recompute_paused`, the same channel the generation + side uses, because a re-prefill needs more teardown than the KV cache: + the executor frees the request's remaining resources, its sequence slot + included, and `reset_for_recompute` rewrites the prompt and resyncs the + Python-side mirrors of it. Pausing the request here instead would leave + the slot owned by SeqSlotManager while `py_seq_slot` is None, which + asserts on the next schedule. + Returns True when pages became available in this iteration. """ if self.kv_cache_manager.has_cache_tier_below_gpu: return False + # A disaggregated generation worker received its context KV rather + # than computing it, so it cannot replay a prefill at all. + if not self.enable_recompute_pause: + return False # Newest first, so the requests closest to completing keep their # pages and the pool drains instead of thrashing. @@ -1411,9 +1419,7 @@ def _try_preempt_for_pages( 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): + if not self._is_recompute_pause_candidate(victim, inflight_request_ids): continue if not self.kv_cache_manager.is_request_active(victim.py_request_id): continue @@ -1426,11 +1432,7 @@ def _try_preempt_for_pages( self._clear_request_runtime_state(victim) if self.draft_kv_cache_manager is not None: self.draft_kv_cache_manager.free_resources(victim) - # Rewrites the prompt to include what was generated and resets - # state to CONTEXT_INIT, so the request re-enters as an ordinary - # prefill. - victim.pause(self.max_input_len) - evicted.append(victim) + recompute_paused.append(victim) preempted_ids.add(victim.py_request_id) return True diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 09e05c0c00a0..bc1e10bf1dc2 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -244,7 +244,6 @@ def make_scheduler( cross_kv_cache_manager: Mock | None = None, enable_prefix_aware_scheduling: bool = True, enable_recompute_pause: 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 @@ -262,8 +261,6 @@ def make_scheduler( kwargs["draft_kv_cache_manager"] = draft_kv_cache_manager 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, @@ -1323,26 +1320,42 @@ def test_out_of_pages_preempts_started_request(self): out = sched.schedule_request(reqs, set()) mgr.preempt_request.assert_called_once_with(victim) - assert ids(out.paused_requests) == [99] + assert ids(out.recompute_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): + def test_released_victim_is_left_for_the_executor_to_reset(self): + """The rest of the teardown belongs to the recompute-pause path.""" 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) + sched = make_scheduler(mgr, max_num_tokens=1000) 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()) + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) - victim.pause.assert_called_once_with(4096) + victim.pause.assert_not_called() + victim.reset_for_recompute.assert_not_called() + assert ids(out.recompute_paused_requests) == [99] assert victim.py_batch_idx is None + def test_disagg_generation_worker_never_preempts(self): + """It received its context KV, so it cannot replay a prefill.""" + 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, enable_recompute_pause=False) + 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_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( @@ -1441,7 +1454,7 @@ def test_chunked_context_out_of_pages_preempts(self): 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] + assert ids(out.recompute_paused_requests) == [99] # =========================================================================== From a6b5124f729090a2994cd8d84f9ba10903026151 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Sat, 19 Sep 2026 17:25:14 -0700 Subject: [PATCH 04/24] [None][chore] Address follow-up comments from CodeRabbit A deferral behind an in-flight prefix contributor leaves every scheduled list empty, while the duplicate it deferred stays in pending_ctx and still counts as a context candidate, so the deadlock detector reads a working engine as stalled. The contributor running is the progress, so the deferral now says so directly. Counting a non-empty inflight_request_ids as progress instead would blunt the detector under overlap scheduling, where something is almost always in flight. mooncake_donor treated an empty --local_buffer_size as an omitted one and silently took the 64 MiB default. Only None means omitted now, so an empty or malformed value reaches parse_size, and both size options report a bad value as a usage error rather than a traceback. Tests: preemption releasing the draft pool alongside the target; the whole preemption handoff, where the scheduler preempts, the executor frees and resets the victim, and a later pass admits the request that was blocked; and a deferral behind an in-flight contributor repeated past the stall threshold without raising. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../pyexecutor/scheduler/scheduler_v2.py | 7 ++ tensorrt_llm/commands/mooncake.py | 20 ++++- .../kv_cache/test_kv_cache_v2_scheduler.py | 73 +++++++++++++++++++ 3 files changed, 96 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index a622b58336f8..f3aba65eb123 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -470,6 +470,11 @@ def preempt_for_pages(req: LlmRequest) -> bool: contributed_blocks = self._collect_contributed_blocks( requests_list, pending_ctx, inflight_request_ids ) + # A deferral behind an in-flight contributor leaves nothing on any of + # the scheduled lists, so the deadlock detector below would read the + # iteration as a stall even though the contributor is running. Deferring + # is itself the progress in that case. + deferred_behind_contributor = False for req in pending_ctx: if budget.requests_full: @@ -504,6 +509,7 @@ def preempt_for_pages(req: LlmRequest) -> bool: f"Deferring context request {req.py_request_id}: its first new " "block is already contributed by a request that runs this iteration" ) + deferred_behind_contributor = True continue peft_pages = budget.peft_pages_needed(req) if peft_pages is None: @@ -551,6 +557,7 @@ def preempt_for_pages(req: LlmRequest) -> bool: or disagg_candidates or evicted or recompute_paused + or deferred_behind_contributor ), ) diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index de779107f4b3..56d237001560 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -265,7 +265,21 @@ def mooncake_donor( f"naming one (or set {CONFIG_PATH_ENV})." ) - donating = parse_size(segment_size) + def size_option(name: str, value: str) -> int: + """Parse a size option, reporting a bad one as a usage error.""" + try: + return parse_size(value) + except ValueError as exc: + raise click.UsageError(f"{name}: {exc}") from exc + + donating = size_option("--segment_size", segment_size) + # None means the option was left off. An empty string was passed, so it goes + # to parse_size and is rejected rather than silently taking the default. + buffer_size = ( + DEFAULT_DONOR_LOCAL_BUFFER_SIZE + if local_buffer_size is None + else size_option("--local_buffer_size", local_buffer_size) + ) resolved = resolve_master_address(master, master_timeout()) wait_for_master(resolved) @@ -276,9 +290,7 @@ def mooncake_donor( 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") or DEFAULT_METADATA_SERVER), - local_buffer_size=( - parse_size(local_buffer_size) if local_buffer_size else DEFAULT_DONOR_LOCAL_BUFFER_SIZE - ), + local_buffer_size=buffer_size, ) as host: if ready_file: with open(ready_file, "w") as handle: diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index bc1e10bf1dc2..7aa85db79a60 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -1300,6 +1300,11 @@ def _out_of_pages_for(request_id): return lambda req, n: req.py_request_id != request_id +#: What the executor passes to reset_for_recompute. Its choice, not the +#: scheduler's, so the exact value does not matter here. +UNBOUNDED_MAX_INPUT_LEN = 0x7FFFFFFF + + class TestContextPreemption: """Releasing a started request's pages when suspension cannot help. @@ -1343,6 +1348,56 @@ def test_released_victim_is_left_for_the_executor_to_reset(self): assert ids(out.recompute_paused_requests) == [99] assert victim.py_batch_idx is None + def test_preemption_releases_the_draft_pool_too(self): + """The draft pool mirrors the target, so a half-released victim leaks.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + draft_mgr = make_kv_cache_manager() + sched = make_scheduler(mgr, max_num_tokens=1000, draft_kv_cache_manager=draft_mgr) + 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) + draft_mgr.free_resources.assert_called_once_with(victim) + assert ids(out.recompute_paused_requests) == [99] + + def test_blocked_request_is_admitted_after_the_executor_recomputes(self): + """The whole handoff, not just the scheduler's half of it. + + A preemption is only worth anything if the blocked request gets in on a + later pass, and that depends on the executor freeing the victim and + resetting it for recompute in between. + """ + out_of_pages = {0} + + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: req.py_request_id not in out_of_pages, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + blocked = make_ctx_request(0, 100) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + first = sched.schedule_request([blocked, victim], set()) + assert ids(first.context_requests) == [] + assert ids(first.recompute_paused_requests) == [99] + + # Stand in for the executor: free the victim's resources, reset it for + # recompute, and let the pages it gave up satisfy the blocked request. + for req in first.recompute_paused_requests: + mgr.free_resources(req) + req.reset_for_recompute(UNBOUNDED_MAX_INPUT_LEN) + req.is_first_context_chunk = True + out_of_pages.clear() + + second = sched.schedule_request([blocked, victim], set()) + + assert 0 in ids(second.context_requests) + victim.reset_for_recompute.assert_called_once() + def test_disagg_generation_worker_never_preempts(self): """It received its context KV, so it cannot replay a prefill.""" mgr = make_kv_cache_manager( @@ -3647,6 +3702,24 @@ def test_deferral_never_empties_the_batch_alone(self): assert ids(out.context_requests) == [] assert mgr.resize_context.call_count == 2 + def test_deferral_behind_an_inflight_contributor_is_not_a_stall(self): + """Deferring is progress when the contributor is the one in flight. + + Nothing reaches a scheduled list on such a pass, but the duplicate + stays in pending_ctx and still counts as a candidate, so without an + explicit signal the detector reads a working engine as hung. + """ + mgr = self._keyed_manager({5: b"blockA", 1: b"blockA"}) + sched = make_scheduler(mgr, max_num_tokens=10000) + sched._DEADLOCK_STALL_ITERS = 3 + contributor = make_ctx_request(5, 100, is_first_context_chunk=False) + duplicate = make_ctx_request(1, 500) + reqs = [duplicate, contributor] + + for _ in range(sched._DEADLOCK_STALL_ITERS + 1): + out = sched.schedule_request(reqs, {5}) + assert ids(out.context_requests) == [] + def test_registered_contributor_cannot_be_evicted(self): """Eviction and deferral cannot collide. From 6edc03970e56c702e94904ca88e84e77050bd079 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Sat, 19 Sep 2026 17:25:45 -0700 Subject: [PATCH 05/24] [None][fix] Address comments from zhaoyuanh-nvidia The key namespace defaulted to the model path's basename, so two engines sharing a pool that load different checkpoints under the same directory name, org-a/model and org-b/model or two revisions mounted alike, agreed on a namespace while disagreeing on what the pages mean. Each would read the other's KV as its own, and nothing about it would look wrong. No derivation of the identity is both cheap and safe, so the key is required rather than defaulted. resolve_model_key raises when it is unset, which covers a hand-written Mooncake JSON reached through MOONCAKE_CONFIG_PATH, and KvCacheConnectorConfig rejects a described pool without one, which fails at startup rather than at the first transfer. mooncake_store.model_key carries it for a server that provisions its own pool, and the rendered client config passes it to the ranks that server spawns. TRTLLM_MOONCAKE_STORE_MODEL_KEY still overrides per process. The pool object keeps the field optional because 'trtllm-serve mooncake_master' builds one to run a master with no engine behind it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/config.py | 26 ++++++++++++++----- .../connectors/mooncake_store/master.py | 2 ++ tensorrt_llm/llmapi/llm_args.py | 19 ++++++++++++++ .../executor/test_mooncake_store_common.py | 9 ++++--- .../executor/test_mooncake_store_master.py | 22 +++++++++++++--- 5 files changed, 66 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 577bd22676dd..b73caf63182e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -157,10 +157,10 @@ class MooncakeStoreConnectorConfig: 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. + #: Identity the keys are namespaced by. Two engines share cache only when + #: they agree on this, and two that disagree about what it names read each + #: other's pages, so it is required rather than defaulted. See + #: :meth:`resolve_model_key`. 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. @@ -262,7 +262,21 @@ def with_env_overrides(self) -> "MooncakeStoreConnectorConfig": 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.""" + """The model identity to namespace keys by. + + Deliberately has no default. Deriving one from the model path would + make two checkpoints that happen to share a directory name, such as + `org-a/model` and `org-b/model` or two revisions mounted alike, agree + on a namespace while disagreeing on what the pages mean, and each would + read the other's KV as its own. + """ if self.model_key: return self.model_key - return os.path.basename(str(model).rstrip("/")) or str(model) + raise ValueError( + f"The mooncake-store connector needs a model key to namespace its " + f"pool keys by, and there is no safe default: set " + f"kv_connector_config.mooncake_store.model_key, the model_key field " + f"of the Mooncake JSON config, or ${MODEL_KEY_ENV}. Give it a value " + f"that separates this checkpoint from any other an engine sharing " + f"the pool might load, rather than one derived from {model!r}." + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index f2f3a21a3c91..93f1f4843182 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -340,6 +340,8 @@ def _client_config( } if pool.cache_prefix is not None: config["cache_prefix"] = pool.cache_prefix + if pool.model_key is not None: + config["model_key"] = pool.model_key # 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: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f44032cd1d3f..cc9b3247b1ea 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2384,6 +2384,15 @@ class MooncakeStoreConfig(StrictBaseModel): None, description="Key namespace for the pool. Bump it after any change to " "page layout or contents. Defaults to 'trtllm'.") + model_key: Optional[str] = Field( + None, + telemetry=False, + description="Identity the keys are namespaced by. Required whenever a " + "pool is described here, because engines sharing a pool read each " + "other's pages exactly when they agree on this. Give checkpoints that " + "differ in weights, revision, quantization, or anything else that " + "changes the KV distinct values. TRTLLM_MOONCAKE_STORE_MODEL_KEY " + "overrides it for a single process.") stage_through_host: bool = Field( False, telemetry=False, @@ -2532,6 +2541,16 @@ def _resolve_preset(self) -> "KvCacheConnectorConfig": "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.") + if (self.mooncake_store is not None + and self.mooncake_store.model_key is None): + raise ValueError( + "mooncake_store.model_key is required. Engines sharing a pool " + "read each other's pages exactly when they agree on this, so " + "there is no safe default: deriving one from the model path " + "would give org-a/model and org-b/model, or two revisions " + "mounted under one directory name, the same namespace for " + "different weights. Set it to a value that separates this " + "checkpoint from any other an engine on this pool may load.") return self diff --git a/tests/unittest/_torch/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py index 9e02826d32e6..8bcb00d51fb9 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_common.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -309,12 +309,15 @@ def test_config_metadata_server_falls_back_to_the_handshake(tmp_path, monkeypatc assert MooncakeStoreConnectorConfig.from_env().metadata_server == "P2PHANDSHAKE" -def test_config_model_key_defaults_to_basename(store_config, tmp_path, monkeypatch): +def test_config_requires_an_explicit_model_key(store_config, tmp_path, monkeypatch): + # A default derived from the path would let org-a/model and org-b/model + # agree on a namespace while disagreeing on what the pages mean. 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" + with pytest.raises(ValueError, match="needs a model key"): + config.resolve_model_key("/models/MiniMax-M3/") @pytest.mark.parametrize( @@ -335,7 +338,7 @@ def test_config_namespace_env_overrides(store_config, monkeypatch): config = MooncakeStoreConnectorConfig.from_env() assert config.cache_prefix == "tenant-a" - # Overrides the JSON's model_key, not just the basename default. + # Overrides the JSON's model_key rather than only supplying a missing one. assert config.resolve_model_key("/models/test-model") == "llama-3.1-8b@rev7" diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py index 618bcfbb8512..7b96ad99f9f8 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_master.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -185,17 +185,26 @@ def test_pool_is_rejected_unless_the_connector_is_mooncake_store(): with pytest.raises(ValueError, match="mooncake_store describes a Mooncake pool"): KvCacheConnectorConfig( connector="lmcache", - mooncake_store=MooncakeStoreConfig(launch_master=True), + mooncake_store=MooncakeStoreConfig(launch_master=True, model_key="m"), ) # 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), + mooncake_store=MooncakeStoreConfig(launch_master=True, model_key="m"), ) +def test_a_described_pool_needs_a_model_key(): + """Two checkpoints that agree on the namespace read each other's pages.""" + with pytest.raises(ValueError, match="mooncake_store.model_key is required"): + KvCacheConnectorConfig( + connector="mooncake-store", + mooncake_store=MooncakeStoreConfig(launch_master=True), + ) + + # ---- the rendered client config ---- @@ -208,6 +217,7 @@ def test_client_config_is_what_the_connector_reads_back(tmp_path): global_segment_size="64GiB", local_buffer_size="4GiB", cache_prefix="trtllm-m3", + model_key="minimax-m3@rev7", stage_through_host=True, transfer_batch_size=32, ) @@ -222,6 +232,9 @@ def test_client_config_is_what_the_connector_reads_back(tmp_path): assert parsed.global_segment_size == 64 * 1024**3 assert parsed.local_buffer_size == 4 * 1024**3 assert parsed.cache_prefix == "trtllm-m3" + # Reaches the ranks the server spawns, which otherwise have no model key + # and would refuse to name a page. + assert parsed.resolve_model_key("/models/ignored") == "minimax-m3@rev7" assert parsed.stage_through_host is True assert parsed.transfer_batch_size == 32 @@ -231,6 +244,7 @@ def test_client_config_omits_the_fields_the_pool_left_unset(): pool = MooncakeStoreConfig(master_server_address="host:50051") written = master_module._client_config(pool, "host:50051") assert "cache_prefix" not in written + assert "model_key" not in written assert "staging_buffer_bytes" not in written @@ -424,7 +438,9 @@ def test_provisioning_is_a_no_op_unless_a_pool_is_described(config): def test_a_described_pool_is_provisioned(running_master): config = KvCacheConnectorConfig( connector="mooncake-store", - mooncake_store=MooncakeStoreConfig(master_server_address=running_master), + mooncake_store=MooncakeStoreConfig( + master_server_address=running_master, model_key="test-model" + ), ) with maybe_provision_pool(config): written = json.loads(open(os.environ[CONFIG_PATH_ENV]).read()) From 1cd820b94d806c502f7805953ab486623f778a5b Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:39:45 -0700 Subject: [PATCH 06/24] [None][fix] Offer --telemetry/--no-telemetry on the mooncake commands mooncake_master and mooncake_donor are registered in the telemetry-aware trtllm-serve group but rejected the opt-out the group documents, so 'trtllm-serve mooncake_master --no-telemetry' failed with 'No such option'. Both now take the standard flag and apply it. Addresses review comment from Mgluhovskoi. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/commands/mooncake.py | 34 +++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index 56d237001560..33cf50755df0 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -31,7 +31,33 @@ import click +import tensorrt_llm.usage as usage from tensorrt_llm.logger import logger +from tensorrt_llm.usage.config import UsageContext + +#: Both commands sit in the telemetry-aware `trtllm-serve` group, so they have +#: to offer the same opt-out the group documents. +_telemetry_option = click.option( + "--telemetry/--no-telemetry", + default=True, + help="Enable or disable anonymous usage telemetry collection.", +) + + +def _apply_cli_telemetry(telemetry: bool) -> None: + """Honor --no-telemetry for a command that reads no config of its own. + + The group already applies the flag it finds in argv, so this matters when + the command is reached through its callback rather than through the CLI. + """ + if telemetry: + return + usage.apply_usage_session_config( + {"disabled": True}, + default_usage_context=UsageContext.CLI_SERVE.value, + component="server", + lifecycle_phase="config_validation", + ) def _until_signalled() -> threading.Event: @@ -98,6 +124,7 @@ def stop(signum, _frame): show_default=True, help="Interval between liveness lines. 0 disables them.", ) +@_telemetry_option def mooncake_master( rpc_port: int, metrics_port: int, @@ -105,12 +132,15 @@ def mooncake_master( address_file: Optional[str], run_dir: Optional[str], heartbeat_seconds: int, + telemetry: bool, ): """Run a mooncake_master for as long as this command runs. A single server with a pool of its own should set `mooncake_store.launch_master` instead. """ + _apply_cli_telemetry(telemetry) + # 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 @@ -223,6 +253,7 @@ def mooncake_master( show_default=True, help="Interval between liveness lines. 0 disables them.", ) +@_telemetry_option def mooncake_donor( master_server_address: Optional[str], segment_size: str, @@ -233,12 +264,15 @@ def mooncake_donor( local_buffer_size: Optional[str], ready_file: Optional[str], heartbeat_seconds: int, + telemetry: bool, ): """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. """ + _apply_cli_telemetry(telemetry) + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, From 2189f028d787924c98eb7fc697565700a5b05d79 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:45:05 -0700 Subject: [PATCH 07/24] [None][fix] Hand a signalled mooncake command to the telemetry boundary The SIGINT/SIGTERM handler set a threading.Event and let the command return normally, so the boundary saw a clean pre-model exit and suppressed the terminal report. Event.set() also takes a lock from inside a synchronous handler. Both commands now install _telemetry.raise_signal_exit, which unwinds the master's and the donor's context managers, so cleanup still runs and the signal reaches the outer report. The log line moves out of the handler to after the release. Adds test_mooncake_store_cli.py covering both signals on both commands: exactly one terminal event, classified as a signal, with the held resource released. It also covers the new opt-out flag. Addresses review comment from Mgluhovskoi. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/commands/mooncake.py | 84 +++++--- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../executor/test_mooncake_store_cli.py | 201 ++++++++++++++++++ 3 files changed, 254 insertions(+), 32 deletions(-) create mode 100644 tests/unittest/_torch/executor/test_mooncake_store_cli.py diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index 33cf50755df0..a0fea3ac2e39 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -21,20 +21,25 @@ from nodes that run no connector. """ +import contextlib import json import os import signal import tempfile -import threading import time from typing import Optional import click import tensorrt_llm.usage as usage +from tensorrt_llm.commands import _telemetry as _command_telemetry from tensorrt_llm.logger import logger from tensorrt_llm.usage.config import UsageContext +#: How long a donor with heartbeats turned off sleeps between wakeups. Only +#: the signal that ends the command interrupts it, so the value is arbitrary. +_IDLE_POLL_SECONDS = 60.0 + #: Both commands sit in the telemetry-aware `trtllm-serve` group, so they have #: to offer the same opt-out the group documents. _telemetry_option = click.option( @@ -60,22 +65,32 @@ def _apply_cli_telemetry(telemetry: bool) -> None: ) -def _until_signalled() -> threading.Event: - """An event that SIGINT and SIGTERM set. +@contextlib.contextmanager +def _signal_handoff(): + """Turn SIGINT and SIGTERM into the exit the telemetry boundary expects. 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() - def stop(signum, _frame): - logger.info(f"mooncake-store: signal {signum} received, shutting down") - stopping.set() + `raise_signal_exit` unwinds those context managers and carries the signal + number out to `trtllm-serve`'s boundary, which is what classifies the exit + as a signal. A handler that only set a flag and let the command return + normally would be reported as a clean pre-model exit instead, and taking + `threading.Event`'s lock from inside a synchronous handler is unsafe on + its own. + Wrap this around the resource so the log line below follows the release. + """ for received in (signal.SIGINT, signal.SIGTERM): - signal.signal(received, stop) - return stopping + signal.signal(received, _command_telemetry.raise_signal_exit) + try: + yield + except _command_telemetry.SignalExit as stopping: + # Logged here rather than in the handler, which must not take the + # logging lock. + logger.info(f"mooncake-store: signal {stopping.signal_number} received, shut down") + raise @click.command("mooncake_master") @@ -158,8 +173,7 @@ def mooncake_master( or tempfile.mkdtemp(prefix="trtllm-mooncake-master-") ) - stopping = _until_signalled() - with running_master(pool, run_dir, address_file=address_file) as master: + with _signal_handoff(), 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 " @@ -167,14 +181,15 @@ def mooncake_master( ) started = time.monotonic() announced = started - while not stopping.is_set(): + # Left for a signal to end, which raises through the sleep below. + while True: if (code := master.process.poll()) is not None: # 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 {master.log_path}" ) - stopping.wait(1.0) + time.sleep(1.0) now = time.monotonic() # Distinguishes a dead master from a dead fabric. if heartbeat_seconds > 0 and now - announced >= heartbeat_seconds: @@ -317,15 +332,19 @@ def size_option(name: str, value: str) -> int: resolved = resolve_master_address(master, master_timeout()) wait_for_master(resolved) - 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") or DEFAULT_METADATA_SERVER), - local_buffer_size=buffer_size, - ) as host: + with ( + _signal_handoff(), + 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") or DEFAULT_METADATA_SERVER + ), + local_buffer_size=buffer_size, + ) as host, + ): if ready_file: with open(ready_file, "w") as handle: handle.write(f"{host} {donating}\n") @@ -336,15 +355,16 @@ def size_option(name: str, value: str) -> int: ) # Idle by design: a put or get here would make this node a traffic - # client, which is what donation exists to avoid. + # client, which is what donation exists to avoid. Left for a signal to + # end, which raises through the sleep below. started = time.monotonic() - while not stopping.is_set(): + while True: if heartbeat_seconds <= 0: - stopping.wait() + time.sleep(_IDLE_POLL_SECONDS) continue - 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" - ) + time.sleep(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" + ) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index eb5fce771dfe..d41067842c82 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/kv_cache/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_kv_cache_layout.py - unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py + - unittest/_torch/executor/test_mooncake_store_cli.py - unittest/_torch/executor/test_mooncake_store_common.py - unittest/_torch/executor/test_mooncake_store_donor.py - unittest/_torch/executor/test_mooncake_store_master.py diff --git a/tests/unittest/_torch/executor/test_mooncake_store_cli.py b/tests/unittest/_torch/executor/test_mooncake_store_cli.py new file mode 100644 index 000000000000..e517d3b8a300 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_cli.py @@ -0,0 +1,201 @@ +# 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 two Mooncake pool commands as `trtllm-serve` runs them. + +Runs without a Mooncake installation and without a GPU. The resource each +command holds is a context manager that records its own release, and the wait +the command would otherwise spend idling is where the signal under test is +delivered to this process. +""" + +import contextlib +import os +import signal +import time +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors import mooncake_store +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import CONFIG_PATH_ENV +from tensorrt_llm.commands import _telemetry +from tensorrt_llm.commands import mooncake as mooncake_commands +from tensorrt_llm.usage.config import UsageContext + +HANDLED_SIGNALS = [signal.SIGINT, signal.SIGTERM] +COMMANDS = ["mooncake_master", "mooncake_donor"] + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + """No ambient pool config: these tests drive the commands from their flags.""" + monkeypatch.delenv(CONFIG_PATH_ENV, raising=False) + + +@pytest.fixture(autouse=True) +def restore_signal_handlers(): + """Both commands install handlers process-wide, so put pytest's back.""" + installed = {number: signal.getsignal(number) for number in HANDLED_SIGNALS} + yield + for number, handler in installed.items(): + signal.signal(number, handler) + + +@pytest.fixture(autouse=True) +def terminal_outcomes(monkeypatch): + """Collect what the shared boundary reports, without a usage session. + + Autouse because a command reaching the real session would apply this + process's opt-out for the rest of the run. + """ + outcomes = [] + monkeypatch.setattr( + _telemetry.usage, "report_exit", lambda outcome, **_kwargs: outcomes.append(outcome) + ) + monkeypatch.setattr(_telemetry.usage, "apply_usage_session_config", lambda *a, **k: True) + monkeypatch.setattr(_telemetry.usage, "set_lifecycle_phase", lambda *a, **k: None) + monkeypatch.setattr(_telemetry.usage, "get_observed_signal", lambda: 0) + monkeypatch.setattr(_telemetry.usage, "record_observed_signal", lambda *a, **k: None) + return outcomes + + +class Resource: + """Stands in for the master process or the mounted segment. + + Records its release, so a test can tell a signal that unwound the command + from one that ended it with the `finally` never reached. + """ + + def __init__(self, held): + self.held = held + self.released = False + + @contextlib.contextmanager + def holding(self, *_args, **_kwargs): + try: + yield self.held + finally: + self.released = True + + +class CommandUnderTest: + """One command, its stubbed resource, and the arguments that reach it.""" + + def __init__(self, resource: Resource, args: list[str]): + self.resource = resource + self.args = args + + def run(self, *extra_args: str) -> None: + group = _telemetry.TelemetryGroup( + name="trtllm-serve", + telemetry_usage_context=UsageContext.CLI_SERVE, + telemetry_component="server", + commands={ + "mooncake_master": mooncake_commands.mooncake_master, + "mooncake_donor": mooncake_commands.mooncake_donor, + }, + ) + group.main( + args=[*self.args, *extra_args], + prog_name="trtllm-serve", + standalone_mode=False, + ) + + +@pytest.fixture +def master(monkeypatch, tmp_path) -> CommandUnderTest: + resource = Resource( + SimpleNamespace( + address="127.0.0.1:50051", + log_path=str(tmp_path / "master.log"), + # Alive, so the command idles rather than reporting a dead master. + process=SimpleNamespace(poll=lambda: None), + ) + ) + monkeypatch.setattr(mooncake_store, "running_master", resource.holding) + return CommandUnderTest(resource, ["mooncake_master", "--run_dir", str(tmp_path)]) + + +@pytest.fixture +def donor(monkeypatch) -> CommandUnderTest: + resource = Resource("10.0.0.1:12345") + monkeypatch.setattr(mooncake_store, "donate_segment", resource.holding) + monkeypatch.setattr(mooncake_store, "resolve_master_address", lambda address, _timeout: address) + monkeypatch.setattr(mooncake_store, "wait_for_master", lambda _address: None) + return CommandUnderTest( + resource, + ["mooncake_donor", "--master_server_address", "127.0.0.1:50051", "--segment_size", "1GiB"], + ) + + +@pytest.fixture +def command(request, master, donor) -> CommandUnderTest: + return {"mooncake_master": master, "mooncake_donor": donor}[request.param] + + +def signal_on_idle(monkeypatch, number: int) -> None: + """Deliver `number` to this process the first time a command idles. + + The handler runs at the next bytecode boundary in the main thread, so the + loop here only has to give the interpreter one. It bounds the wait rather + than blocking, so a handler that never fires fails the test instead of + hanging it. + """ + + def sleep(_seconds): + os.kill(os.getpid(), number) + for _ in range(100): + time.sleep(0.01) + + monkeypatch.setattr( + mooncake_commands, "time", SimpleNamespace(sleep=sleep, monotonic=time.monotonic) + ) + + +@pytest.mark.parametrize("number", HANDLED_SIGNALS) +@pytest.mark.parametrize("command", COMMANDS, indirect=True) +def test_a_signal_is_reported_once_and_still_releases_what_was_held( + monkeypatch, terminal_outcomes, command, number +): + """The signal reaches the boundary, and the resource is still given up. + + A handler that only set a flag would let the command return normally, + which the boundary reads as a clean exit before any model was loaded. + """ + signal_on_idle(monkeypatch, number) + + with pytest.raises(_telemetry.SignalExit) as stopped: + command.run() + + assert stopped.value.signal_number == number + assert command.resource.released + + assert len(terminal_outcomes) == 1 + outcome = terminal_outcomes[0] + assert outcome.termination_kind == "signal" + assert outcome.signal_number == number + assert outcome.exit_code == 128 + number + + +@pytest.mark.parametrize("flag", ["--telemetry", "--no-telemetry"]) +@pytest.mark.parametrize("command", COMMANDS, indirect=True) +def test_the_opt_out_the_group_documents_is_accepted(monkeypatch, command, flag): + """Without the option Click rejects the flag as a usage error instead.""" + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + command.run(flag) + + assert command.resource.released From e23d334c29a19f7715daa3fc81d531f50ed04d45 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:50:07 -0700 Subject: [PATCH 08/24] [None][fix] Do not call a pending KV transfer a V2 scheduler deadlock Counting waiting context requests made the detector reachable on a context server, where asynchronous sends hold KV pages after prefill. Those requests sit past the schedulable states, so they never reach the scheduled lists and made_progress does not see them, while the new CONTEXT_INIT requests they block keep failing to allocate. The detector would then raise well before the transfer's own timeout. A pass with any request in DISAGG_CONTEXT_TRANS_IN_PROGRESS, DISAGG_CONTEXT_COMPLETE or DISAGG_GENERATION_TRANS_IN_PROGRESS now resets the stall count: the pages are coming back, and a send that never lands belongs to the transfer layer to report. Tests cover all three holder states and that the reprieve ends with the transfer. Addresses review comment from nv-xtf. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../pyexecutor/scheduler/scheduler_v2.py | 23 +++++++++ .../kv_cache/test_kv_cache_v2_scheduler.py | 51 +++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index f3aba65eb123..90fb2bac9434 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -1451,6 +1451,19 @@ def _try_preempt_for_pages( # PEFT budget, IndexMapper slots) clear long before. _DEADLOCK_STALL_ITERS = 1000 + # States in which an in-flight KV transfer still owns pages it is about to + # release: a context-only request sending its cache once prefill finished, + # one whose send landed and which the executor has yet to reap, and a + # generation request receiving a cache. None of them are schedulable, so + # the capacity they are about to return is invisible to `made_progress`. + _TRANSFER_HOLDING_STATE_VALUES = frozenset( + { + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS.value, + LlmRequestState.DISAGG_CONTEXT_COMPLETE.value, + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS.value, + } + ) + def _detect_deadlock( self, active_requests: RequestList, @@ -1495,6 +1508,16 @@ def _detect_deadlock( self._stalled_schedules = 0 return + # Waiting on a transfer is not a deadlock: the pages it holds come + # back when it lands. This matters on a context server, where a pool + # filled by asynchronous sends blocks every new CONTEXT_INIT request + # while nothing the scheduler can see makes progress. A send that + # never lands is the transfer layer's timeout to report, not a + # scheduling failure, so the count resets rather than merely pausing. + if any(req.state_value in self._TRANSFER_HOLDING_STATE_VALUES for req in active_requests): + self._stalled_schedules = 0 + return + self._stalled_schedules += 1 if self._stalled_schedules < self._DEADLOCK_STALL_ITERS: return diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 7aa85db79a60..4a8215eb6e64 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -38,6 +38,8 @@ CONTEXT_INIT = LlmRequestState.CONTEXT_INIT.value # 10 GEN_IN_PROGRESS = LlmRequestState.GENERATION_IN_PROGRESS.value # 13 GEN_TO_COMPLETE = LlmRequestState.GENERATION_TO_COMPLETE.value # 14 +DISAGG_CTX_TRANS_IN_PROGRESS = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS.value # 21 +DISAGG_CTX_COMPLETE = LlmRequestState.DISAGG_CONTEXT_COMPLETE.value # 22 # --------------------------------------------------------------------------- @@ -1587,6 +1589,55 @@ def test_all_candidates_inflight_never_raises(self): for _ in range(5): sched.schedule_request(reqs, {0}) + @pytest.mark.parametrize( + "holder_state", + [DISAGG_CTX_TRANS_IN_PROGRESS, DISAGG_CTX_COMPLETE, DISAGG_GEN_TRANS_IN_PROGRESS], + ) + def test_a_pending_transfer_holding_pages_is_not_a_deadlock(self, holder_state): + """A context server's pool can be full of sends that have not landed. + + Those requests are past the schedulable states, so they never reach + the scheduled lists, and the new context requests they block cannot + allocate. Without this the detector would call that a deadlock well + before the transfer's own timeout has anything to say. + """ + 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), + make_filtered_request(1, state_value=holder_state), + ] + + for _ in range(sched._DEADLOCK_STALL_ITERS + 2): + sched.schedule_request(reqs, set()) + + assert sched._stalled_schedules == 0 + + def test_a_deadlock_behind_a_finished_transfer_is_still_reported(self): + """The reprieve lasts only as long as the transfer does.""" + 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 + blocked = make_ctx_request(0, 100, is_first_context_chunk=False) + sending = make_filtered_request(1, state_value=DISAGG_CTX_TRANS_IN_PROGRESS) + + for _ in range(5): + sched.schedule_request([blocked, sending], set()) + + # The send landed and the executor reaped it, but the pool is still + # full, so the stall is now the scheduler's to report. + for _ in range(sched._DEADLOCK_STALL_ITERS - 1): + sched.schedule_request([blocked], set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request([blocked], set()) + # =========================================================================== # PEFT / LoRA From e11bdc7ded569439a1095f46b2420275ff56edb1 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:52:05 -0700 Subject: [PATCH 09/24] [None][fix] Reserve preempted KV pages for the request that freed them A successful preemption only skipped the triggering request, so a context request later in the same pass could take the pages while the one that paid a re-prefill for them waited until the next iteration. Repeated, that leaves the triggering request blocked indefinitely while the scheduler still looks like it is making progress. The preemption now ends the phase 2 loop, which holds the capacity for one iteration until the triggering request re-enters prepare_context. A preemption that finds no victim still lets the pass continue, since there is nothing to reserve. Tests document the ordering, including that a committed generation request is unaffected. Addresses review comment from nv-xtf. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../pyexecutor/scheduler/scheduler_v2.py | 17 +++++- .../kv_cache/test_kv_cache_v2_scheduler.py | 55 +++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 90fb2bac9434..e087d2530115 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -450,6 +450,17 @@ def _schedule_loop(self, active_requests, inflight_request_ids): preempted_ids: set[int] = set() def preempt_for_pages(req: LlmRequest) -> bool: + """Free pages for `req` by giving up one started request. + + A success ends the phase 2 loop rather than only skipping `req`. + The pages are reserved for the request that paid a re-prefill for + them: letting a later, smaller context request take them instead + would leave `req` blocked and force it to preempt again on the + next pass, which can repeat indefinitely while the scheduler + still looks like it is making progress. The cost of holding them + is one iteration of admission, a couple of milliseconds, against + a whole prefill replayed for nothing. + """ 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) @@ -810,7 +821,8 @@ def _try_schedule_context_full( # 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) + if preempt_for_pages(req): + return ScheduleAction.STOP, 0, False return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) @@ -923,7 +935,8 @@ def _try_schedule_context_chunked( # draft tokens for last chunk. if not self._try_allocate_context(req, resize_tokens): # Out of pages, as in _try_schedule_context_full. - preempt_for_pages(req) + if preempt_for_pages(req): + return ScheduleAction.STOP, 0, False return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 4a8215eb6e64..7d461ed07c1c 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -1426,6 +1426,61 @@ def test_preempted_victim_not_scheduled_in_the_same_pass(self): assert ids(out.context_requests) == [] + def test_freed_pages_are_reserved_for_the_request_that_preempted(self): + """No later context request may spend them first. + + Otherwise the request that paid a re-prefill for the pages stays + blocked and preempts again next pass, which can repeat without ever + admitting it while the scheduler still looks like it is progressing. + """ + 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) + blocked = make_ctx_request(0, 100) + # Would fit in what the preemption releases, and is only behind + # `blocked` by arrival order. + later = make_ctx_request(1, 10) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([blocked, later, victim], set()) + + assert ids(out.recompute_paused_requests) == [99] + assert ids(out.context_requests) == [] + + def test_a_failed_preemption_still_lets_the_pass_continue(self): + """With nothing to give up there is no capacity to reserve.""" + 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) + # Only first chunks, so there is no preemption victim among them. + reqs = [make_ctx_request(0, 100), make_ctx_request(1, 10)] + + out = sched.schedule_request(reqs, set()) + + mgr.preempt_request.assert_not_called() + assert ids(out.context_requests) == [1] + + def test_generation_scheduled_before_a_preemption_is_unaffected(self): + """Phase 1 has already committed, so ending phase 2 costs it nothing.""" + 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) + blocked = make_ctx_request(1, 100) + later = make_ctx_request(2, 10) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_gen_request(0), blocked, later, victim], set()) + + assert ids(out.generation_requests) == [0] + assert ids(out.recompute_paused_requests) == [99] + 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), From ad5a85e7dd34b948e9329b4ce4b73946ae2a9674 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:55:42 -0700 Subject: [PATCH 10/24] [None][chore] Tighten comments on the review follow-ups Drop wording that argued against how the code used to behave, so each comment stands on its own, and remove the repetition between the transfer-state list and its use, between the preemption policy and the test that covers it, and between the two idle loops. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../pyexecutor/scheduler/scheduler_v2.py | 30 ++++++++----------- tensorrt_llm/commands/mooncake.py | 13 +++----- .../kv_cache/test_kv_cache_v2_scheduler.py | 16 ++++------ .../executor/test_mooncake_store_cli.py | 6 ++-- 4 files changed, 24 insertions(+), 41 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index e087d2530115..17cd4c8da41f 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -452,14 +452,11 @@ def _schedule_loop(self, active_requests, inflight_request_ids): def preempt_for_pages(req: LlmRequest) -> bool: """Free pages for `req` by giving up one started request. - A success ends the phase 2 loop rather than only skipping `req`. - The pages are reserved for the request that paid a re-prefill for - them: letting a later, smaller context request take them instead - would leave `req` blocked and force it to preempt again on the - next pass, which can repeat indefinitely while the scheduler - still looks like it is making progress. The cost of holding them - is one iteration of admission, a couple of milliseconds, against - a whole prefill replayed for nothing. + A success ends the phase 2 loop, reserving the pages for the + request that paid a re-prefill for them. Letting a later context + request take them instead would leave `req` to preempt again on + the next pass, repeating without ever admitting it. The cost is + one iteration of admission. """ protected = {r.py_request_id for r in scheduled_gen} protected.update(r.py_request_id for r in scheduled_ctx) @@ -1465,10 +1462,9 @@ def _try_preempt_for_pages( _DEADLOCK_STALL_ITERS = 1000 # States in which an in-flight KV transfer still owns pages it is about to - # release: a context-only request sending its cache once prefill finished, - # one whose send landed and which the executor has yet to reap, and a - # generation request receiving a cache. None of them are schedulable, so - # the capacity they are about to return is invisible to `made_progress`. + # release: a context-only request sending its cache after prefill, one + # whose send landed and which the executor has yet to reap, and a + # generation request receiving a cache. _TRANSFER_HOLDING_STATE_VALUES = frozenset( { LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS.value, @@ -1521,12 +1517,10 @@ def _detect_deadlock( self._stalled_schedules = 0 return - # Waiting on a transfer is not a deadlock: the pages it holds come - # back when it lands. This matters on a context server, where a pool - # filled by asynchronous sends blocks every new CONTEXT_INIT request - # while nothing the scheduler can see makes progress. A send that - # never lands is the transfer layer's timeout to report, not a - # scheduling failure, so the count resets rather than merely pausing. + # Waiting on a transfer is not a deadlock: the pages come back when it + # lands. None of those states are schedulable, so a context server + # whose pool is full of pending sends shows no progress here at all. A + # send that never lands is the transfer layer's timeout to report. if any(req.state_value in self._TRANSFER_HOLDING_STATE_VALUES for req in active_requests): self._stalled_schedules = 0 return diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index a0fea3ac2e39..5aa6c7d05bf6 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -74,13 +74,10 @@ def _signal_handoff(): the master unreaped or the pool advertising memory that has gone. `raise_signal_exit` unwinds those context managers and carries the signal - number out to `trtllm-serve`'s boundary, which is what classifies the exit - as a signal. A handler that only set a flag and let the command return - normally would be reported as a clean pre-model exit instead, and taking - `threading.Event`'s lock from inside a synchronous handler is unsafe on - its own. + number out to `trtllm-serve`, which is what reports the exit as a signal + rather than as a clean one. - Wrap this around the resource so the log line below follows the release. + Wrap this around the resource so the log below follows the release. """ for received in (signal.SIGINT, signal.SIGTERM): signal.signal(received, _command_telemetry.raise_signal_exit) @@ -181,7 +178,6 @@ def mooncake_master( ) started = time.monotonic() announced = started - # Left for a signal to end, which raises through the sleep below. while True: if (code := master.process.poll()) is not None: # The pool is gone once the master dies, and every client is @@ -355,8 +351,7 @@ def size_option(name: str, value: str) -> int: ) # Idle by design: a put or get here would make this node a traffic - # client, which is what donation exists to avoid. Left for a signal to - # end, which raises through the sleep below. + # client, which is what donation exists to avoid. started = time.monotonic() while True: if heartbeat_seconds <= 0: diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 7d461ed07c1c..970cbbcc5a2e 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -1427,20 +1427,15 @@ def test_preempted_victim_not_scheduled_in_the_same_pass(self): assert ids(out.context_requests) == [] def test_freed_pages_are_reserved_for_the_request_that_preempted(self): - """No later context request may spend them first. - - Otherwise the request that paid a re-prefill for the pages stays - blocked and preempts again next pass, which can repeat without ever - admitting it while the scheduler still looks like it is progressing. - """ + """No later context request may spend them first.""" 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) blocked = make_ctx_request(0, 100) - # Would fit in what the preemption releases, and is only behind - # `blocked` by arrival order. + # Fits in what the preemption releases, and is only behind `blocked` + # in arrival order. later = make_ctx_request(1, 10) victim = make_ctx_request(99, 100, is_first_context_chunk=False) @@ -1652,9 +1647,8 @@ def test_a_pending_transfer_holding_pages_is_not_a_deadlock(self, holder_state): """A context server's pool can be full of sends that have not landed. Those requests are past the schedulable states, so they never reach - the scheduled lists, and the new context requests they block cannot - allocate. Without this the detector would call that a deadlock well - before the transfer's own timeout has anything to say. + the scheduled lists, and the context requests they block cannot + allocate. The transfer's own timeout covers a send that never lands. """ mgr = make_kv_cache_manager( resize_context_fn=lambda req, n: False, diff --git a/tests/unittest/_torch/executor/test_mooncake_store_cli.py b/tests/unittest/_torch/executor/test_mooncake_store_cli.py index e517d3b8a300..3f486fb1375f 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_cli.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_cli.py @@ -150,7 +150,7 @@ def signal_on_idle(monkeypatch, number: int) -> None: The handler runs at the next bytecode boundary in the main thread, so the loop here only has to give the interpreter one. It bounds the wait rather - than blocking, so a handler that never fires fails the test instead of + than blocking, so a handler that never fires fails the test rather than hanging it. """ @@ -171,8 +171,8 @@ def test_a_signal_is_reported_once_and_still_releases_what_was_held( ): """The signal reaches the boundary, and the resource is still given up. - A handler that only set a flag would let the command return normally, - which the boundary reads as a clean exit before any model was loaded. + A command that returned normally instead would be reported as a clean + exit before any model was loaded. """ signal_on_idle(monkeypatch, number) From 4929e6ec6fff7f6a002755b4f586f5056b4f8420 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 16:11:28 -0700 Subject: [PATCH 11/24] [None][chore] Shorten the mooncake CMake package cleanup The CMake-generated mooncake package lands in the directory mooncake-integration/CMakeLists.txt selects, so repeating that one expression removes it directly instead of scanning sys.path for candidates. Drops the install step from 57 added lines to 33. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/common/install_mooncake.sh | 60 ++++++++++--------------------- 1 file changed, 18 insertions(+), 42 deletions(-) diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index a145b3756f62..bc61ae8d9034 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -57,56 +57,32 @@ echo "export LD_LIBRARY_PATH=${MOONCAKE_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> # 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. +# wheel below. # -# 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". -# -# - 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. -# -# 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 -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 +# `make install` also emits a `mooncake` package that omits +# libmooncake_store.so, so importing mooncake.store from it fails. It has to go +# before the wheel is installed: CMake writes +# store.cpython-312-x86_64-linux-gnu.so where the wheel writes store.so, and +# importlib prefers the interpreter-tagged suffix, so the broken extension +# would win even after pip reports success. The directory is the one +# mooncake-integration/CMakeLists.txt chose, which this repeats. +MOONCAKE_CMAKE_PACKAGE="$(python3 -c "import sys; print([s for s in sys.path if 'packages' in s][0])")/mooncake" +echo "removing CMake-generated mooncake package: ${MOONCAKE_CMAKE_PACKAGE}" +rm -rf "${MOONCAKE_CMAKE_PACKAGE}" -# 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, 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. +# 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, 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}" # Fail the build rather than ship an image whose import is broken. python3 - <<'PY' -from mooncake.store import MooncakeDistributedStore import mooncake.store -MooncakeDistributedStore() +mooncake.store.MooncakeDistributedStore() print(f"mooncake.store OK: {mooncake.store.__file__}") PY From 83793db3cfcee5b8ade3a2de4807fda07afcfcca Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:16:16 -0700 Subject: [PATCH 12/24] [None][chore] Defer the Mooncake LlmArgs surface to the connector MR Part 1 landed `MooncakeStoreConfig`, `MooncakeDonationConfig`, and `TorchLlmArgs.mooncake_donation` ahead of the connector that reads them. Ten of those fields were only ever serialized into the Mooncake client config, whose sole reader is `MooncakeStoreConnectorConfig.from_file`, and nothing called it: the connector classes were placeholders that refused construction. `mooncake_donation` was read only by `trtllm-serve`, so `LLM(mooncake_donation=...)` validated and then did nothing. Holding the API back until the connector lands also keeps this MR off the files that the KV-connector V2 work is changing underneath it. Everything the fields configured is still reachable. The two commands this MR already added cover the master and the donor, and `mooncake_master` gains `--config` so it renders the client config into its run directory, where `provisioned_config_path` already looks for it. A deployment describes its pool in the Mooncake JSON and the `TRTLLM_MOONCAKE_STORE_*` variables, which is the schema the connector reads and is shared with vLLM. `provision_pool` and `running_master` now take a `PoolSpec` dataclass instead of duck-typing the Pydantic model, so the library no longer depends on `llm_args` at all and the user-facing spelling can be settled alongside the connector. Removed with the API they served: the `mooncake-store` registry preset and `uses_connector`, the partial-reuse gate keyed off it, the placeholder connector classes, the startup gates in `validation.py`, and the `trtllm-serve` provisioning wrapper. The V2 scheduler preemption work is unchanged. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/__init__.py | 30 +-- .../connectors/mooncake_store/connector.py | 54 ----- .../connectors/mooncake_store/donor.py | 54 +---- .../connectors/mooncake_store/master.py | 185 +++++++++++---- .../connectors/mooncake_store/validation.py | 60 ----- .../_torch/pyexecutor/connectors/registry.py | 28 --- .../_torch/pyexecutor/py_executor_creator.py | 14 -- tensorrt_llm/commands/mooncake.py | 51 +++-- tensorrt_llm/commands/serve.py | 151 ++++-------- tensorrt_llm/grpc/smg/server.py | 7 +- tensorrt_llm/llmapi/llm_args.py | 214 +----------------- .../usage/llm_args_golden_manifest.json | 3 +- .../executor/test_mooncake_store_cli.py | 121 +++++++++- .../executor/test_mooncake_store_common.py | 92 +------- .../executor/test_mooncake_store_donor.py | 96 ++------ .../executor/test_mooncake_store_master.py | 149 ++++++------ .../api_stability/references/llm.yaml | 4 - 17 files changed, 453 insertions(+), 860 deletions(-) delete mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py delete mode 100644 tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index be28d0590594..9c122862d0f5 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -24,48 +24,48 @@ pool addressed by content. The two compose, so a context server can write pages here and still hand off over NIXL. -The pool is described in `KvCacheConnectorConfig.mooncake_store`, which lets -`trtllm-serve` provision it during bringup 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, so `donor.py` lends -a node's memory to the pool without giving it a connector. Both need the -Mooncake Python bindings (`pip install mooncake-transfer-engine`). +`master.py` brings a pool up: it resolves or launches the `mooncake_master` +and renders the client config the workers read. Capacity comes only from +processes that open a store handle, which in a disaggregated deployment is the +context servers alone, so `donor.py` lends a node's memory to the pool without +giving it a connector. `trtllm-serve mooncake_master` and `mooncake_donor` +expose both. Both need the Mooncake Python bindings (`pip install +mooncake-transfer-engine`). `keys.py` and `staging.py` hold what the store side shares with the connector that moves pages in and out of the pool: how a block of tokens becomes a store -key, and how pages reach the fabric on hosts without GPUDirect RDMA. -`connector.py` holds the connector classes, which are placeholders. +key, and how pages reach the fabric on hosts without GPUDirect RDMA. The +connector itself, and the `LlmArgs` surface that selects it, land with the KV +cache manager V2 support it depends on. """ from .config import MooncakeStoreConnectorConfig, StoreRole, parse_size -from .connector import MooncakeStoreConnectorScheduler, MooncakeStoreConnectorWorker -from .donor import DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, maybe_donate_segment +from .donor import DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment from .master import ( + PoolSpec, local_address, master_timeout, - maybe_provision_pool, provision_pool, resolve_device_name, resolve_master_address, running_master, wait_for_master, + write_client_config, ) __all__ = [ "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", "MooncakeStoreConnectorConfig", - "MooncakeStoreConnectorScheduler", - "MooncakeStoreConnectorWorker", + "PoolSpec", "StoreRole", "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", + "write_client_config", ] diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py deleted file mode 100644 index 2e5dc2e66cc0..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py +++ /dev/null @@ -1,54 +0,0 @@ -# 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. -"""Placeholders for the connector that moves KV pages in and out of the pool. - -The store side stands on its own: a pool can be provisioned, lent memory, and -shared with an engine elsewhere without a connector here. Moving pages needs the -KV cache layout description, so neither class is implemented here. - -`CONNECTOR_REGISTRY` names both and `py_executor_creator` resolves them by name, -so they exist to turn `connector: mooncake-store` into a clear refusal instead -of an `AttributeError` raised once the pool is already provisioned. - -Neither subclasses `KvCacheConnectorWorker` or `KvCacheConnectorScheduler`, -whose `__init_subclass__` checks a method set there is nothing here to satisfy. -""" - -__all__ = ["MooncakeStoreConnectorScheduler", "MooncakeStoreConnectorWorker"] - -_UNAVAILABLE = ( - "The mooncake-store KV cache connector is not available in this build, so an " - "engine here cannot read or write the pool. The pool itself is: " - "'trtllm-serve mooncake_master' owns one, 'trtllm-serve mooncake_donor' and " - "the mooncake_donation setting lend it host memory, and an engine that does " - "have the connector can share it. Drop kv_connector_config.connector to " - "serve without one." -) - - -class _Unavailable: - """Reports what this build can and cannot do with a Mooncake pool.""" - - def __init__(self, llm_args): - del llm_args - raise NotImplementedError(_UNAVAILABLE) - - -class MooncakeStoreConnectorWorker(_Unavailable): - """Moves this rank's KV pages between its cache and the pool.""" - - -class MooncakeStoreConnectorScheduler(_Unavailable): - """Decides which pages a request loads from the pool and saves to it.""" diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py index 4fe167a5bdaf..1373b82f08ee 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -20,10 +20,10 @@ 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, 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. +Running `trtllm-serve mooncake_donor` on a generation node puts that node's +memory into the same 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 @@ -35,23 +35,16 @@ import contextlib import time -from typing import Any, Iterator, Optional +from typing import Iterator, Optional from tensorrt_llm.logger import logger -from .config import DEFAULT_METADATA_SERVER, parse_size -from .master import ( - local_address, - master_timeout, - resolve_device_name, - resolve_master_address, - wait_for_master, -) +from .config import DEFAULT_METADATA_SERVER +from .master import local_address __all__ = [ "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", "donate_segment", - "maybe_donate_segment", ] #: A donor never transfers, but `setup` rejects a zero-sized transfer buffer. @@ -136,36 +129,3 @@ def donate_segment( 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, 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 - - # 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 " - "it; resolving the master now" - ) - master_address = resolve_master_address(donation.master_server_address, master_timeout()) - wait_for_master(master_address) - with donate_segment( - master_server_address=master_address, - segment_size=parse_size(donation.segment_size), - protocol=donation.protocol, - 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 93f1f4843182..b5b96a7301b8 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -12,22 +12,22 @@ # 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. +"""Bring a Mooncake store's pool up and tell its workers where to find it. 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. +`provision_pool` produces both. It resolves the master, either launching one or +checking that the configured one answers, renders the client config, and exports +`MOONCAKE_CONFIG_PATH`. 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`. +A master launched by the process that provisions lives and dies with it, so it +suits one engine talking to its own pool. Several engines sharing a pool, or a +pool meant to survive a restart, want `trtllm-serve mooncake_master`, which +holds one for as long as that command runs and publishes its address for the +rest of the deployment to read back. """ import contextlib @@ -38,23 +38,23 @@ import subprocess # nosec B404 import tempfile import time -from dataclasses import dataclass -from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple +from dataclasses import dataclass, fields +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union from tensorrt_llm.logger import logger -from ..registry import uses_connector -from .config import CLIENT_CONFIG_NAME, CONFIG_PATH_ENV, RUN_DIR_ENV +from .config import CLIENT_CONFIG_NAME, CONFIG_PATH_ENV, DEFAULT_METADATA_SERVER, RUN_DIR_ENV __all__ = [ + "PoolSpec", "local_address", - "maybe_provision_pool", "master_timeout", "provision_pool", "resolve_device_name", "resolve_master_address", "running_master", "wait_for_master", + "write_client_config", ] #: Override the binary that `launch_master` runs. @@ -75,6 +75,69 @@ LOG_TAIL_LINES = 20 +@dataclass(frozen=True) +class PoolSpec: + """A pool to provision: whose master owns it, and what its workers are told. + + Deliberately internal. The settings below end up in the Mooncake client + config, whose schema is vLLM's and is the contract the connector reads; + this is only the shape `provision_pool` takes them in. Keeping it out of + `LlmArgs` leaves the user-facing spelling to be settled alongside the + connector that gives these fields meaning. + + Exactly one of `master_server_address` and `launch_master` says where the + master is, which `provision_pool` enforces rather than this class, since a + caller may fill the two in from different places. + """ + + #: `host:port`, or `file://` naming a file that holds one. The file + #: is how to reach a master whose host a scheduler chose. + master_server_address: Optional[str] = None + #: Start a `mooncake_master` and use it, instead of joining one. + launch_master: bool = False + #: Extra path to publish a launched master's address to, beyond the copy + #: always written to the run directory. + master_address_file: Optional[str] = None + master_port: int = 50051 + master_metrics_port: int = 9004 + master_eviction_ratio: float = 0.05 + metadata_server: str = DEFAULT_METADATA_SERVER + #: `rdma` or `tcp`. TCP is for bring-up only. + protocol: str = "rdma" + #: RDMA device to transfer over, from `ibv_devinfo`. Empty auto-detects. + device_name: str = "" + #: Host memory each worker process contributes. Pool capacity is this times + #: the number of processes that open a store handle. + global_segment_size: Union[int, str] = "16GiB" + #: Per-process Mooncake transfer buffer, not pool capacity. + local_buffer_size: Union[int, str] = "1GiB" + transfer_batch_size: int = 64 + #: Left unset so the rendered config omits the key and the connector's own + #: default applies, rather than carrying a second copy of it here. + cache_prefix: Optional[str] = None + #: Identity the keys are namespaced by. Engines sharing a pool read each + #: other's pages exactly when they agree on this, so there is no safe + #: default; see `MooncakeStoreConnectorConfig.resolve_model_key`. + model_key: Optional[str] = None + #: Copy pages through a pinned host buffer instead of registering the KV + #: pools, for hosts without GPUDirect RDMA. + stage_through_host: bool = False + staging_buffer_bytes: Optional[Union[int, str]] = None + + @staticmethod + def from_json(raw: Dict[str, Any], **overrides: Any) -> "PoolSpec": + """Build a spec from a Mooncake client config, plus what a caller adds. + + Lets a deployment describe its pool in the same JSON the connector + reads, rather than in a second format, and lets the CLI layer on the + parts that are the command's rather than the pool's. + """ + known = {f.name for f in fields(PoolSpec)} + settings = {key: value for key, value in raw.items() if key in known} + settings.update({key: value for key, value in overrides.items() if value is not None}) + return PoolSpec(**settings) + + 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: @@ -319,7 +382,7 @@ def wait_for_master(master_address: str, timeout: Optional[float] = None) -> Opt def _client_config( - pool: Any, master_address: str, device_name: Optional[str] = None + pool: PoolSpec, master_address: str, device_name: Optional[str] = None ) -> Dict[str, Any]: """Render the Mooncake client config for a pool. @@ -349,6 +412,27 @@ def _client_config( return config +def write_client_config(pool: PoolSpec, master_address: str, run_dir: str) -> str: + """Render `pool`'s client config into `run_dir` and return its path. + + Workers reach it either through `MOONCAKE_CONFIG_PATH`, which + `provision_pool` exports, or by looking in `$TRTLLM_MOONCAKE_RUN_DIR` when + an external launcher started them too early to inherit it; see + `provisioned_config_path`. Writing it next to the master's address file is + what lets `trtllm-serve mooncake_master --run_dir` describe a whole pool + without any engine having to. + """ + os.makedirs(run_dir, exist_ok=True) + 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) + logger.info(f"mooncake-store: wrote {config_path} ({json.dumps(config, sort_keys=True)})") + return config_path + + @dataclass class LaunchedMaster: """A `mooncake_master` owned by this process.""" @@ -368,7 +452,7 @@ def stop(self, timeout: float = 10.0) -> None: self.process.wait() -def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: +def _launch_master(pool: PoolSpec, 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) @@ -466,7 +550,7 @@ def _address_files(run_dir: str, extra: Optional[str] = None) -> List[str]: @contextlib.contextmanager def running_master( - pool: Any, run_dir: str, address_file: Optional[str] = None + pool: PoolSpec, run_dir: str, address_file: Optional[str] = None ) -> Iterator[LaunchedMaster]: """Run a master whose lifetime is this process's rather than an engine's. @@ -489,24 +573,25 @@ def running_master( @contextlib.contextmanager -def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optional[str]]: +def provision_pool(pool: PoolSpec, 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`. + pool: The pool to bring up. 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. """ + _check_exactly_one_master(pool) + 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." + "the pool it names is used as is and nothing is provisioned here." ) yield None return @@ -544,29 +629,21 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona 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) + config_path = write_client_config(pool, master_address, run_dir) # 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 - logger.info( - f"mooncake-store: {CONFIG_PATH_ENV}={config_path} " - f"({json.dumps(config, sort_keys=True)})" - ) + logger.info(f"mooncake-store: {CONFIG_PATH_ENV}={config_path}") # 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 " + "mooncake-store: each rank that opens a handle will 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" + "total capacity is that times the number of such ranks, plus " + "whatever any donor lends" ) yield config_path finally: @@ -579,20 +656,28 @@ def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optiona 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. +def _check_exactly_one_master(pool: PoolSpec) -> None: + """Reject a pool that names no master, or two. - 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. + Checked here rather than on `PoolSpec` because a caller may fill the two + fields in from different places, such as a command line and a config file, + and is entitled to an incomplete spec until it has read both. """ - 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 + if pool.launch_master and pool.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 pool.launch_master and not pool.master_server_address: + raise ValueError( + "mooncake-store: needs a master. Set master_server_address to join " + "an existing pool, or launch_master to start one." + ) + if pool.master_address_file and not pool.launch_master: + raise ValueError( + "mooncake-store: master_address_file publishes the address of a " + "master started here, so it needs launch_master. To read an " + "address a master elsewhere published, set master_server_address " + "to file://." + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py deleted file mode 100644 index 4df451e42b5a..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 tensorrt_llm.llmapi.llm_args import TorchLlmArgs - -__all__ = ["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." - ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py index c0a40bf480a6..9a00cdadd7fd 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py @@ -19,11 +19,6 @@ 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", @@ -40,27 +35,4 @@ "connector_scheduler_class": "DynamoKVBMConnectorLeader", "connector_worker_class": "DynamoKVBMConnectorWorker", }, - # Both classes currently refuse construction; see the module's docstring. - "mooncake-store": { - "connector_module": "tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", - "connector_scheduler_class": "MooncakeStoreConnectorScheduler", - "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 1747142b76d8..6204a1cd88bf 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -47,7 +47,6 @@ resolve_cache_transceiver_config, uses_fp4_mla_attention, uses_vswa_kv_cache_layout) from .connectors.kv_cache_connector import KvCacheConnectorManager -from .connectors.registry import uses_connector from .dwdp import DwdpManager, get_global_dwdp_manager from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder from .hang_diagnostics import monitor_executor_initialization @@ -397,19 +396,6 @@ def _create_py_executor_impl( kv_cache_config.enable_block_reuse = False kv_cache_config.enable_partial_reuse = False - # 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( - "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 - # The tokenizer is stripped from MPI kwargs in proxy.py to avoid pickle # failures with trust_remote_code models. Reload it from the checkpoint # when guided decoding needs it. diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index 5aa6c7d05bf6..2a147276056e 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -12,13 +12,15 @@ # 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. +"""The two pieces of a Mooncake pool that are not any one engine's to run. -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. 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. +`mooncake_master` owns a pool for as long as the command runs, and with +`--config` also writes the client config its workers read into `--run_dir`. +Pointing every server at that directory with `$TRTLLM_MOONCAKE_RUN_DIR`, or at +the file with `$MOONCAKE_CONFIG_PATH`, is how a pool is described today. + +`mooncake_donor` lends a node's host memory to a pool it does not otherwise +use, which is how capacity comes from nodes whose engines have no connector. """ import contextlib @@ -129,6 +131,16 @@ def _signal_handoff(): help="Where to keep the master's log. Defaults to " "$TRTLLM_MOONCAKE_RUN_DIR, else a temporary directory.", ) +@click.option( + "--config", + type=str, + default=None, + help="Mooncake JSON config describing the pool. When given, a copy " + "naming this master is written to --run_dir as the client config " + "every worker reads, so no external script has to render one. " + "Workers find it by setting $TRTLLM_MOONCAKE_RUN_DIR to that " + "directory, or $MOONCAKE_CONFIG_PATH to the file.", +) @click.option( "--heartbeat_seconds", type=int, @@ -143,22 +155,31 @@ def mooncake_master( eviction_ratio: float, address_file: Optional[str], run_dir: Optional[str], + config: Optional[str], heartbeat_seconds: int, telemetry: bool, ): - """Run a mooncake_master for as long as this command runs. - - A single server with a pool of its own should set - `mooncake_store.launch_master` instead. - """ + """Run a mooncake_master for as long as this command runs.""" _apply_cli_telemetry(telemetry) # 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 + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( + PoolSpec, + running_master, + write_client_config, + ) + + raw = {} + if config: + with open(config) as handle: + raw = json.load(handle) + # This command is the master, so whichever one the config names is not + # the one being described here. + raw.pop("master_server_address", None) - pool = MooncakeStoreConfig( + pool = PoolSpec.from_json( + raw, launch_master=True, master_port=rpc_port, master_metrics_port=metrics_port, @@ -171,6 +192,8 @@ def mooncake_master( ) with _signal_handoff(), running_master(pool, run_dir, address_file=address_file) as master: + if config: + write_client_config(pool, master.address, run_dir) logger.info( f"mooncake-store: this master owns the pool until this command " f"stops; address {master.address}, log {master.log_path}, metrics " diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 37d136d3223a..f94e7ffc5688 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -32,8 +32,8 @@ from importlib.util import find_spec from pathlib import Path from types import FrameType -from typing import (TYPE_CHECKING, Any, Dict, Iterator, NamedTuple, NoReturn, - Optional, Sequence, Set) +from typing import (TYPE_CHECKING, Any, Dict, NamedTuple, NoReturn, Optional, + Sequence, Set) import click import torch @@ -59,9 +59,7 @@ parse_disagg_config_file, parse_metadata_server_config_file, validate_config_bool) -from tensorrt_llm.llmapi.llm_args import (KvCacheConnectorConfig, - MooncakeDonationConfig, - MultimodalConfig, TorchLlmArgs) +from tensorrt_llm.llmapi.llm_args import 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, @@ -623,49 +621,6 @@ def _terminate_attached_frontends(children: list) -> None: child.kill() -@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. - - 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 pool and lend a second segment. - """ - if not owns_engine: - yield - return - - 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): - # 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 - - 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( host: str, port: int, @@ -747,50 +702,48 @@ def launch_server( # until uvicorn takes it over, so no one can steal the port in between. _publish_bound_address(report_addr, host, port) - with _provision_kv_cache_pool( - llm_args, owns_engine=not multi_frontend.is_attached_frontend): - if backend == 'pytorch': - llm_args.pop("build_config", None) - llm = PyTorchLLM(**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, - internal_disagg_auth_key=internal_disagg_auth_key) - _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) + if backend == 'pytorch': + llm_args.pop("build_config", None) + llm = PyTorchLLM(**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, + internal_disagg_auth_key=internal_disagg_auth_key) + _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_mm_encoder_server( @@ -1602,16 +1555,10 @@ def _serve_llm(): "https://buf.build/gen/python " "\"tensorrt_llm[openengine]\"`.") from error - # launch_smg_server provisions from inside itself; OpenEngine's - # server is an optional package this repo does not own, so the - # pool is brought up around it here instead. Either way the - # context has to outlive engine construction, which happens - # inside the launch call. - with _provision_kv_cache_pool(llm_args): - launch_grpc_server(host, - port, - llm_args, - served_model_name=served_model_name) + launch_grpc_server(host, + port, + llm_args, + served_model_name=served_model_name) else: # Default: launch OpenAI HTTP server launch_server( diff --git a/tensorrt_llm/grpc/smg/server.py b/tensorrt_llm/grpc/smg/server.py index a7b3bfe55d7a..8e891ac6556b 100644 --- a/tensorrt_llm/grpc/smg/server.py +++ b/tensorrt_llm/grpc/smg/server.py @@ -126,12 +126,7 @@ def signal_handler() -> None: logger.info("LLM engine stopped") logger.info("Shutdown complete") - # Imported here rather than at module scope: tensorrt_llm.commands.serve - # reaches into this module to launch the server. - from tensorrt_llm.commands.serve import _provision_kv_cache_pool - - with _provision_kv_cache_pool(llm_args): - uvloop.run(serve_grpc_async()) + uvloop.run(serve_grpc_async()) __all__ = ["launch_smg_server"] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index cc9b3247b1ea..d596c1d38a4e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2303,181 +2303,6 @@ 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, 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. An inherited `MOONCAKE_CONFIG_PATH` - still wins, so an externally managed pool stays reachable. - - Fields opt out of telemetry because they size and address one site's pool - rather than saying which features are in use; `kv_connector_config.connector` - already records that the store is on. - """ - 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. 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.") - 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_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, - 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", - telemetry=False, - 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", - telemetry=False, - 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'.") - model_key: Optional[str] = Field( - None, - telemetry=False, - description="Identity the keys are namespaced by. Required whenever a " - "pool is described here, because engines sharing a pool read each " - "other's pages exactly when they agree on this. Give checkpoints that " - "differ in weights, revision, quantization, or anything else that " - "changes the KV distinct values. TRTLLM_MOONCAKE_STORE_MODEL_KEY " - "overrides it for a single process.") - 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.") - 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": - 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.") - 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, 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( - ..., - 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. @@ -2492,8 +2317,7 @@ 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', - 'mooncake-store')) + telemetry=TelemetryField.categorical('lmcache', 'lmcache-mp', 'kvbm')) connector_module: Optional[str] = Field( None, description= @@ -2508,16 +2332,11 @@ 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, uses_connector) + from tensorrt_llm._torch.pyexecutor.connectors.registry import \ + CONNECTOR_REGISTRY if self.connector is not None: preset = CONNECTOR_REGISTRY.get(self.connector) if preset is None: @@ -2535,22 +2354,6 @@ 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.") - if (self.mooncake_store is not None - and self.mooncake_store.model_key is None): - raise ValueError( - "mooncake_store.model_key is required. Engines sharing a pool " - "read each other's pages exactly when they agree on this, so " - "there is no safe default: deriving one from the model path " - "would give org-a/model and org-b/model, or two revisions " - "mounted under one directory name, the same namespace for " - "different weights. Set it to a value that separates this " - "checkpoint from any other an engine on this pool may load.") return self @@ -6221,17 +6024,6 @@ def is_partial_model_loading(self) -> bool: 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/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 07e67f75ce52..257d1ba70c11 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -781,8 +781,7 @@ "allowed_values": [ "lmcache", "lmcache-mp", - "kvbm", - "mooncake-store" + "kvbm" ], "capture_policy": "allowlist|none", "kind": "categorical", diff --git a/tests/unittest/_torch/executor/test_mooncake_store_cli.py b/tests/unittest/_torch/executor/test_mooncake_store_cli.py index 3f486fb1375f..0233637255c5 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_cli.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_cli.py @@ -21,11 +21,13 @@ """ import contextlib +import json import os import signal import time from types import SimpleNamespace +import click import pytest from tensorrt_llm._torch.pyexecutor.connectors import mooncake_store @@ -75,15 +77,20 @@ class Resource: """Stands in for the master process or the mounted segment. Records its release, so a test can tell a signal that unwound the command - from one that ended it with the `finally` never reached. + from one that ended it with the `finally` never reached, and what it was + asked for, so a test can check what the command resolved before asking. """ def __init__(self, held): self.held = held self.released = False + #: `None` until the resource is taken, which distinguishes a command + #: that gave up beforehand from one that took it. + self.taken_with = None @contextlib.contextmanager - def holding(self, *_args, **_kwargs): + def holding(self, *args, **_kwargs): + self.taken_with = args try: yield self.held finally: @@ -199,3 +206,113 @@ def test_the_opt_out_the_group_documents_is_accepted(monkeypatch, command, flag) command.run(flag) assert command.resource.released + + +# ---- what the donor resolves before it offers anything ---- + + +@pytest.fixture +def segment(monkeypatch) -> Resource: + """Stub only the segment, so the command's address handling still runs.""" + resource = Resource("10.0.0.1:12345") + monkeypatch.setattr(mooncake_store, "donate_segment", resource.holding) + return resource + + +def donor_running(resource: Resource, *args: str) -> CommandUnderTest: + return CommandUnderTest(resource, ["mooncake_donor", *args]) + + +def test_a_donor_resolves_a_published_master_address_before_joining(monkeypatch, segment, tmp_path): + """What lets a generation node 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") + probed = [] + monkeypatch.setattr(mooncake_store, "wait_for_master", probed.append) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running( + segment, + "--master_server_address", + f"file://{address_file}", + "--segment_size", + "2GiB", + ).run() + + # The resolved address is both what was probed and what Mooncake is given: + # Mooncake cannot dial a file:// URL. + assert probed == ["10.0.0.9:50051"] + assert segment.taken_with[0] == "10.0.0.9:50051" + # A size string reaching Mooncake unparsed would be a segment of nothing. + assert segment.taken_with[1] == 2 * 1024**3 + + +def test_a_donor_reports_an_unreachable_master_before_offering_the_segment(monkeypatch, segment): + """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(mooncake_store, "resolve_master_address", lambda address, _t: address) + monkeypatch.setattr(mooncake_store, "wait_for_master", refuse) + + with pytest.raises(TimeoutError, match="10.0.0.1:50051"): + donor_running(segment, "--master_server_address", "10.0.0.1:50051").run() + + assert segment.taken_with is None + + +@pytest.mark.parametrize("size", ["", "16 GB!"]) +def test_a_donor_rejects_a_size_it_cannot_parse(monkeypatch, segment, size): + """Caught as a usage error rather than as a segment of some other size.""" + monkeypatch.setattr(mooncake_store, "resolve_master_address", lambda address, _t: address) + monkeypatch.setattr(mooncake_store, "wait_for_master", lambda _address: None) + + with pytest.raises(click.UsageError, match="--segment_size"): + donor_running( + segment, "--master_server_address", "10.0.0.1:50051", "--segment_size", size + ).run() + + assert segment.taken_with is None + + +# ---- describing the pool from the master's own run directory ---- + + +def test_the_master_writes_the_client_config_it_was_given(monkeypatch, master, tmp_path): + """Lets workers be pointed at the run directory instead of a hand-written file.""" + config = tmp_path / "pool.json" + config.write_text( + json.dumps({"protocol": "tcp", "global_segment_size": "8GiB", "model_key": "m"}) + ) + written = {} + monkeypatch.setattr( + mooncake_store, + "write_client_config", + lambda pool, address, run_dir: written.update(pool=pool, address=address), + ) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + master.run("--config", str(config)) + + # Written only once the master answers, and naming that master rather than + # whichever one the config was copied from. + assert written["address"] == "127.0.0.1:50051" + assert written["pool"].protocol == "tcp" + assert written["pool"].global_segment_size == "8GiB" + assert written["pool"].launch_master is True + assert written["pool"].master_server_address is None + + +def test_the_master_writes_no_client_config_unless_asked(monkeypatch, master): + """A deployment that renders its own must not have it overwritten.""" + calls = [] + monkeypatch.setattr(mooncake_store, "write_client_config", lambda *args: calls.append(args)) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + master.run() + + assert calls == [] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py index 8bcb00d51fb9..0347ca026eb7 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_common.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -15,16 +15,14 @@ """Unit tests for the Mooncake store pieces every deployment shares. Covers how a block of tokens becomes a store key, how the JSON config is read, -which deployments are refused at startup, and the pinned host slots pages pass -through where GPUDirect RDMA is unavailable. +and the pinned host slots pages pass through where GPUDirect RDMA is +unavailable. Runs without a Mooncake installation: the store handle is replaced by an in-process fake that records what it was handed. """ -import importlib import json -from types import SimpleNamespace import pytest @@ -49,9 +47,6 @@ stage_batch_for_put, unstage_batch_after_get, ) -from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.validation import validate_llm_args -from tensorrt_llm._torch.pyexecutor.connectors.registry import uses_connector -from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig TOKENS_PER_BLOCK = 4 @@ -111,17 +106,6 @@ def store_config(tmp_path, monkeypatch): 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 staged_copies(monkeypatch): """Record the copies staging would issue, instead of running them.""" @@ -348,75 +332,9 @@ def test_config_rejects_a_non_boolean_staging_env(store_config, monkeypatch): MooncakeStoreConnectorConfig.from_env() -# ---- 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) - - -# ---- connector identification ---- -# -# 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. - - -@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(): - config = KvCacheConnectorConfig(connector="mooncake-store") - with pytest.raises(ValueError, match="Unknown connector preset"): - uses_connector(config, "mooncake-stroe") - - -def test_the_registered_preset_resolves_to_refusing_classes(): - """py_executor_creator resolves both by name from the package.""" - config = KvCacheConnectorConfig(connector="mooncake-store") - module = importlib.import_module(config.connector_module) - - for class_name in (config.connector_scheduler_class, config.connector_worker_class): - with pytest.raises(NotImplementedError, match="not available in this build"): - getattr(module, class_name)(llm_args=None) +# The connector itself, the registry preset that selects it, and the startup +# gates on parallelism and index-V caching land with the KV cache manager V2 +# support they depend on. # ---- host staging ---- diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py index cda334bc385a..0fa98a8da1f1 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_donor.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -28,9 +28,7 @@ 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 @@ -76,40 +74,19 @@ def setup(self, *args): return Refusing -@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", - ) - 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", - ) - ) +def test_a_donor_registers_the_segment_it_was_asked_for(fake_bindings): + """Mooncake's `setup` is positional, so all seven arguments are pinned.""" + donation = donate_segment( + "10.0.0.1:50051", + 32 * GIB, + protocol="rdma", + device_name="mlx5_0", + metadata_server="P2PHANDSHAKE", + hostname="10.0.0.5", + ) with donation as host: - assert host == expected_host + assert host == "10.0.0.5" ( registered_host, metadata_server, @@ -120,11 +97,11 @@ def test_a_donor_registers_the_segment_it_was_asked_for( master, ) = fake_bindings.instances[0].setup_args - assert registered_host == expected_host + assert registered_host == "10.0.0.5" assert metadata_server == "P2PHANDSHAKE" - assert segment_size == expected_size + assert segment_size == 32 * GIB assert protocol == "rdma" - assert device_name == expected_device + assert device_name == "mlx5_0" assert master == "10.0.0.1:50051" assert local_buffer_size == DEFAULT_DONOR_LOCAL_BUFFER_SIZE @@ -156,46 +133,5 @@ def test_missing_bindings_are_reported_as_the_separate_component_they_are(monkey 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_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_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 == [] +# The address a donor joins is resolved and probed by the `mooncake_donor` +# command before it gets here; see test_mooncake_store_cli.py. diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py index 7b96ad99f9f8..ffe5a6b05722 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_master.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -34,13 +34,13 @@ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( CONFIG_PATH_ENV, MooncakeStoreConnectorConfig, + provisioned_config_path, ) from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.master import ( - maybe_provision_pool, + PoolSpec, provision_pool, resolve_master_address, ) -from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig, MooncakeStoreConfig def free_port() -> int: @@ -175,34 +175,12 @@ def running_master(): ], ids=["two_masters", "no_master", "publishing_without_launching"], ) -def test_pool_needs_exactly_one_master(kwargs, message): +def test_provisioning_needs_exactly_one_master(kwargs, message): + """Rejected before anything is launched, dialed, or written.""" with pytest.raises(ValueError, match=message): - MooncakeStoreConfig(**kwargs) - - -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, model_key="m"), - ) - # 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, model_key="m"), - ) - - -def test_a_described_pool_needs_a_model_key(): - """Two checkpoints that agree on the namespace read each other's pages.""" - with pytest.raises(ValueError, match="mooncake_store.model_key is required"): - KvCacheConnectorConfig( - connector="mooncake-store", - mooncake_store=MooncakeStoreConfig(launch_master=True), - ) + with provision_pool(PoolSpec(**kwargs)): + pytest.fail("provisioning should not have yielded") + assert CONFIG_PATH_ENV not in os.environ # ---- the rendered client config ---- @@ -210,7 +188,7 @@ def test_a_described_pool_needs_a_model_key(): 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( + pool = PoolSpec( master_server_address="10.0.0.1:50051", protocol="rdma", device_name="mlx5_0", @@ -241,13 +219,49 @@ def test_client_config_is_what_the_connector_reads_back(tmp_path): 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") + pool = PoolSpec(master_server_address="host:50051") written = master_module._client_config(pool, "host:50051") assert "cache_prefix" not in written assert "model_key" not in written assert "staging_buffer_bytes" not in written +def test_a_written_client_config_is_where_the_ranks_look_for_it(tmp_path, monkeypatch): + """A rank an external launcher started never inherits MOONCAKE_CONFIG_PATH. + + Writing the config into the run directory is what lets + 'trtllm-serve mooncake_master --run_dir' describe a pool to servers it + did not spawn. + """ + pool = PoolSpec(master_server_address="10.0.0.1:50051", model_key="m") + + written = master_module.write_client_config(pool, "10.0.0.1:50051", str(tmp_path)) + + monkeypatch.setenv(master_module.RUN_DIR_ENV, str(tmp_path)) + assert provisioned_config_path() == written + assert MooncakeStoreConnectorConfig.from_file(written).model_key == "m" + + +def test_a_spec_takes_the_pool_settings_from_the_json_the_connector_reads(): + """One schema for the pool, rather than a second spelling of it.""" + pool = PoolSpec.from_json( + { + "protocol": "tcp", + "global_segment_size": "8GiB", + "model_key": "m", + # Not a PoolSpec field; the client config carries settings that + # only the connector reads, and they must not break the parse. + "role": "both", + }, + launch_master=True, + ) + + assert pool.protocol == "tcp" + assert pool.global_segment_size == "8GiB" + assert pool.model_key == "m" + assert pool.launch_master is True + + @pytest.mark.parametrize( "address, expected", [ @@ -265,7 +279,7 @@ def test_master_addresses_are_split_or_declined(address, expected): def test_provisioning_points_the_workers_at_a_running_master(running_master): - pool = MooncakeStoreConfig(master_server_address=running_master) + pool = PoolSpec(master_server_address=running_master) with provision_pool(pool) as config_path: # The workers are spawned inside this window and are told about the @@ -280,7 +294,7 @@ def test_provisioning_points_the_workers_at_a_running_master(running_master): 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( + pool = PoolSpec( master_server_address=running_master, stage_through_host=True, staging_buffer_bytes="4GiB", @@ -294,7 +308,7 @@ def test_a_staging_buffer_can_be_sized_where_staging_is_turned_on(running_master 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()}") + pool = PoolSpec(master_server_address=f"127.0.0.1:{free_port()}") with pytest.raises(TimeoutError, match="did not accept connections"): with provision_pool(pool): @@ -304,7 +318,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, so an unprobeable one passes through.""" - pool = MooncakeStoreConfig(master_server_address="unix:///var/run/mooncake") + pool = PoolSpec(master_server_address="unix:///var/run/mooncake") with provision_pool(pool) as config_path: written = json.loads(open(config_path).read()) @@ -316,7 +330,7 @@ def test_an_inherited_config_path_wins(monkeypatch, tmp_path): harness_config = tmp_path / "harness.json" harness_config.write_text("{}") monkeypatch.setenv(CONFIG_PATH_ENV, str(harness_config)) - pool = MooncakeStoreConfig(launch_master=True) + pool = PoolSpec(launch_master=True) with provision_pool(pool) as config_path: assert config_path is None @@ -331,7 +345,7 @@ def test_an_inherited_config_path_wins(monkeypatch, tmp_path): 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) + pool = PoolSpec(launch_master=True, master_port=port) with provision_pool(pool) as config_path: written = json.loads(open(config_path).read()) @@ -349,7 +363,7 @@ def test_a_launched_master_is_named_in_the_config_and_stopped_on_exit(fake_maste 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( + pool = PoolSpec( launch_master=True, master_port=port, master_metrics_port=free_port(), @@ -370,7 +384,7 @@ def test_a_launched_master_gets_the_flags_and_logging_it_needs(fake_master): 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()) + pool = PoolSpec(launch_master=True, master_port=free_port()) with pytest.raises(RuntimeError, match="exited with code 3"): with provision_pool(pool): @@ -381,7 +395,7 @@ def test_a_master_that_dies_during_startup_says_so(fake_master): 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()) + pool = PoolSpec(launch_master=True, master_port=free_port()) with pytest.raises(TimeoutError, match="did not accept connections"): with provision_pool(pool): @@ -395,7 +409,7 @@ def test_a_missing_master_binary_names_the_alternatives(monkeypatch): "shutil", SimpleNamespace(which=lambda _name: None, rmtree=shutil.rmtree), ) - pool = MooncakeStoreConfig(launch_master=True) + pool = PoolSpec(launch_master=True) with pytest.raises(FileNotFoundError, match="master_server_address"): with provision_pool(pool): @@ -406,7 +420,7 @@ 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) + pool = PoolSpec(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) @@ -417,37 +431,6 @@ def test_a_run_dir_keeps_the_master_log_and_the_config(fake_master, tmp_path): assert (run_dir / master_module.CLIENT_CONFIG_NAME).exists() -# ---- the entry point servers call ---- - - -@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.""" - 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, model_key="test-model" - ), - ) - 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 - - # ---- reaching a master whose host nobody knew in advance ---- @@ -492,7 +475,7 @@ def test_a_standalone_master_publishes_an_address_that_can_be_dialed(fake_master port = free_port() fake_master.arm(listen_on=port) address_file = tmp_path / "master.addr" - pool = MooncakeStoreConfig(launch_master=True, master_port=port) + pool = PoolSpec(launch_master=True, master_port=port) with master_module.running_master( pool, str(tmp_path / "run"), address_file=str(address_file) @@ -509,7 +492,7 @@ def test_a_stopped_master_leaves_no_address_behind(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) + pool = PoolSpec(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() @@ -523,7 +506,7 @@ def test_a_standalone_master_keeps_its_log(fake_master, tmp_path): port = free_port() fake_master.arm(listen_on=port) run_dir = tmp_path / "run" - pool = MooncakeStoreConfig(launch_master=True, master_port=port) + pool = PoolSpec(launch_master=True, master_port=port) with master_module.running_master(pool, str(run_dir)): pass @@ -536,8 +519,8 @@ def test_provisioning_joins_a_master_it_was_never_given_the_address_of(fake_mast 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}") + standalone = PoolSpec(launch_master=True, master_port=port) + worker = PoolSpec(master_server_address=f"file://{address_file}") with master_module.running_master( standalone, str(tmp_path / "run"), address_file=str(address_file) @@ -557,7 +540,7 @@ def test_a_launched_master_publishes_where_its_run_left_its_logs(fake_master, tm port = free_port() fake_master.arm(listen_on=port) run_dir = tmp_path / "run" - pool = MooncakeStoreConfig(launch_master=True, master_port=port) + pool = PoolSpec(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() @@ -571,9 +554,7 @@ def test_a_launched_master_can_be_published_where_the_donors_look(fake_master, t 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) - ) + pool = PoolSpec(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}") @@ -613,7 +594,7 @@ def test_a_master_that_died_starting_is_reported_with_its_last_words(fake_master """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") - pool = MooncakeStoreConfig(launch_master=True, master_port=free_port()) + pool = PoolSpec(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)): @@ -664,7 +645,7 @@ def test_the_detected_device_is_what_the_workers_are_told(fake_master, tmp_path, 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") + pool = PoolSpec(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" diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 9ecbbc61e23b..8728888a2868 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -291,10 +291,6 @@ 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 93bd853dc0c9b8e0f62ac5f28e61651d30ed6e0c Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 23 Sep 2026 04:31:25 +0000 Subject: [PATCH 13/24] Address comments from nv-xtf Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 5 ++ .../pyexecutor/scheduler/scheduler_v2.py | 27 ++++++- .../kv_cache/test_kv_cache_v2_scheduler.py | 71 +++++++++++++++++++ 3 files changed, 102 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3fc930fa8a48..01eaf0d10e1f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -745,6 +745,11 @@ def __init__( self.resource_manager, should_store_blocks=self.enable_disagg_partial_reuse_store) + # Wire the transfer manager into the V2 scheduler's deadlock detector. + if hasattr(self.scheduler, "set_async_transfer_manager"): + self.scheduler.set_async_transfer_manager( + self.async_transfer_manager) + # Router is built after async_transfer_manager so KVCacheAwareADPRouter # can receive the transfer-manager reference at construction time. self.adp_router: ADPRouter = ADPRouter.create( diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 17cd4c8da41f..a4c1acd50dec 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -256,6 +256,20 @@ def __init__( os.environ.get("TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN", "0") == "1" ) + # Registered by PyExecutor; see set_async_transfer_manager. + self._async_transfer_manager = None + + def set_async_transfer_manager(self, mgr) -> None: + """Register the AsyncTransferManager the deadlock detector consults. + + A finished disaggregated context sender leaves active_requests once its + response is emitted, but the transfer manager still owns it and its + pinned KV pages until the send lands. Without this reference the + detector cannot see those pages and may report a false deadlock on a + context server whose pool is full of in-flight sends. + """ + self._async_transfer_manager = mgr + @property def scheduling_state_range( self, @@ -1521,7 +1535,18 @@ def _detect_deadlock( # lands. None of those states are schedulable, so a context server # whose pool is full of pending sends shows no progress here at all. A # send that never lands is the transfer layer's timeout to report. - if any(req.state_value in self._TRANSFER_HOLDING_STATE_VALUES for req in active_requests): + # + # The active_requests scan covers senders still on the list and the + # generation-side receiver. The transfer manager covers finished + # senders that have already left active_requests (see + # set_async_transfer_manager). + transfer_holding = any( + req.state_value in self._TRANSFER_HOLDING_STATE_VALUES for req in active_requests + ) or ( + self._async_transfer_manager is not None + and self._async_transfer_manager.has_any_inflight_requests() + ) + if transfer_holding: self._stalled_schedules = 0 return diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 970cbbcc5a2e..cc663b218a4d 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -1687,6 +1687,77 @@ def test_a_deadlock_behind_a_finished_transfer_is_still_reported(self): with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): sched.schedule_request([blocked], set()) + def test_inflight_send_that_left_active_list_is_not_a_deadlock(self): + """A sender off the active list but still in flight is not a deadlock. + + The sender never appears in active_requests, so only the transfer + manager can tell the detector its pinned pages are coming back. + """ + 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 + transfer_mgr = Mock() + transfer_mgr.has_any_inflight_requests.return_value = True + sched.set_async_transfer_manager(transfer_mgr) + # Only the blocked request is on the active list; the sender has left. + blocked = make_ctx_request(0, 100, is_first_context_chunk=False) + + for _ in range(sched._DEADLOCK_STALL_ITERS + 2): + sched.schedule_request([blocked], set()) + + assert sched._stalled_schedules == 0 + + def test_blocked_request_schedules_after_send_completes(self): + """Once the off-list send lands, its pages free the blocked request.""" + pages_held = [True] + + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: not pages_held[0], + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + transfer_mgr = Mock() + transfer_mgr.has_any_inflight_requests.return_value = True + sched.set_async_transfer_manager(transfer_mgr) + blocked = make_ctx_request(0, 100, is_first_context_chunk=False) + + for _ in range(sched._DEADLOCK_STALL_ITERS + 2): + out = sched.schedule_request([blocked], set()) + assert len(out.context_requests) == 0 + assert sched._stalled_schedules == 0 + + # The send lands: the transfer manager drops it and its pages return. + transfer_mgr.has_any_inflight_requests.return_value = False + pages_held[0] = False + out = sched.schedule_request([blocked], set()) + assert ids(out.context_requests) == [0] + + def test_finished_send_still_deadlocks_with_transfer_manager_wired(self): + """The transfer-manager reprieve ends the moment no send is in flight. + + With the manager wired but reporting nothing in flight, a still-full + pool must deadlock exactly as it does without a manager attached. + """ + 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 + transfer_mgr = Mock() + transfer_mgr.has_any_inflight_requests.return_value = False + sched.set_async_transfer_manager(transfer_mgr) + blocked = make_ctx_request(0, 100, is_first_context_chunk=False) + + for _ in range(sched._DEADLOCK_STALL_ITERS - 1): + sched.schedule_request([blocked], set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request([blocked], set()) + # =========================================================================== # PEFT / LoRA From fdfb2f3f0baf29fe61a53654770a1afe7d8390a5 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 23 Sep 2026 04:51:26 +0000 Subject: [PATCH 14/24] Address comments from Coderabbit Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/commands/mooncake.py | 2 +- .../executor/test_mooncake_store_cli.py | 166 +++++++++++++++++- 2 files changed, 166 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index 2a147276056e..4ef6fda20ca6 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -356,7 +356,7 @@ def size_option(name: str, value: str) -> int: donate_segment( resolved, donating, - protocol=protocol or raw.get("protocol", "rdma"), + protocol=protocol or raw.get("protocol") or "rdma", device_name=device_name or raw.get("device_name", "") or "", metadata_server=( metadata_server or raw.get("metadata_server") or DEFAULT_METADATA_SERVER diff --git a/tests/unittest/_torch/executor/test_mooncake_store_cli.py b/tests/unittest/_torch/executor/test_mooncake_store_cli.py index 0233637255c5..dfeadfa7b9c2 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_cli.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_cli.py @@ -31,7 +31,10 @@ import pytest from tensorrt_llm._torch.pyexecutor.connectors import mooncake_store -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, +) from tensorrt_llm.commands import _telemetry from tensorrt_llm.commands import mooncake as mooncake_commands from tensorrt_llm.usage.config import UsageContext @@ -277,6 +280,167 @@ def test_a_donor_rejects_a_size_it_cannot_parse(monkeypatch, segment, size): assert segment.taken_with is None +# ---- how the donor resolves its flags against a config ---- + + +class RecordingDonation: + """Records the arguments the donor resolves before donate_segment is called. + + The command idles inside the donation context, so a test breaks the loop + with a signal once the resolved arguments and the ready file are in hand. + """ + + def __init__(self, host: str = "10.0.0.1:12345"): + self.host = host + self.resolved: dict = {} + + @contextlib.contextmanager + def holding(self, master, segment, **kwargs): + self.resolved = dict(master=master, segment=segment, **kwargs) + yield self.host + + +@pytest.fixture +def donation(monkeypatch) -> RecordingDonation: + """Drive the donor to the point of donation, capturing what it resolved.""" + recorder = RecordingDonation() + monkeypatch.setattr(mooncake_store, "donate_segment", recorder.holding) + monkeypatch.setattr(mooncake_store, "resolve_master_address", lambda address, _t: address) + monkeypatch.setattr(mooncake_store, "wait_for_master", lambda _address: None) + return recorder + + +def _config_file(tmp_path, **entries) -> str: + path = tmp_path / "pool.json" + path.write_text(json.dumps(entries)) + return str(path) + + +def test_the_donor_uses_its_documented_defaults_without_a_config(monkeypatch, donation): + """With only the required master, transport and buffer defaults apply.""" + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running(donation, "--master_server_address", "10.0.0.1:50051").run() + + assert donation.resolved["protocol"] == "rdma" + assert donation.resolved["device_name"] == "" + assert donation.resolved["metadata_server"] == DEFAULT_METADATA_SERVER + assert donation.resolved["local_buffer_size"] == mooncake_store.DEFAULT_DONOR_LOCAL_BUFFER_SIZE + + +def test_the_donor_resolves_every_option_from_the_cli(monkeypatch, donation): + """Flags alone describe the donation, with no config on hand.""" + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running( + donation, + "--master_server_address", + "10.0.0.1:50051", + "--segment_size", + "2GiB", + "--protocol", + "tcp", + "--device_name", + "mlx5_0", + "--metadata_server", + "http://cli:8080/metadata", + ).run() + + assert donation.resolved["master"] == "10.0.0.1:50051" + assert donation.resolved["segment"] == 2 * 1024**3 + assert donation.resolved["protocol"] == "tcp" + assert donation.resolved["device_name"] == "mlx5_0" + assert donation.resolved["metadata_server"] == "http://cli:8080/metadata" + + +def test_the_donor_falls_back_to_the_config_for_unset_options(monkeypatch, donation, tmp_path): + """A config names the master and the transport when no flag overrides it.""" + config = _config_file( + tmp_path, + master_server_address="10.0.0.9:50051", + protocol="tcp", + device_name="mlx5_1", + metadata_server="http://config:8080/metadata", + ) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running(donation, "--config", config, "--segment_size", "1GiB").run() + + assert donation.resolved["master"] == "10.0.0.9:50051" + assert donation.resolved["protocol"] == "tcp" + assert donation.resolved["device_name"] == "mlx5_1" + assert donation.resolved["metadata_server"] == "http://config:8080/metadata" + + +def test_the_donor_prefers_the_cli_over_the_config(monkeypatch, donation, tmp_path): + """A flag wins over the same setting in the config.""" + config = _config_file( + tmp_path, + master_server_address="10.0.0.9:50051", + protocol="tcp", + device_name="mlx5_1", + metadata_server="http://config:8080/metadata", + ) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running( + donation, + "--config", + config, + "--master_server_address", + "10.0.0.1:50051", + "--segment_size", + "1GiB", + "--protocol", + "rdma", + "--device_name", + "mlx5_0", + "--metadata_server", + "http://cli:8080/metadata", + ).run() + + assert donation.resolved["master"] == "10.0.0.1:50051" + assert donation.resolved["protocol"] == "rdma" + assert donation.resolved["device_name"] == "mlx5_0" + assert donation.resolved["metadata_server"] == "http://cli:8080/metadata" + + +def test_the_donor_defaults_the_protocol_to_rdma_when_the_config_value_is_null( + monkeypatch, donation, tmp_path +): + """A null protocol in the config must not reach setup unnormalized.""" + config = _config_file(tmp_path, master_server_address="10.0.0.9:50051", protocol=None) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running(donation, "--config", config, "--segment_size", "1GiB").run() + + assert donation.resolved["protocol"] == "rdma" + + +def test_the_donor_announces_the_mounted_segment_in_its_ready_file(monkeypatch, donation, tmp_path): + """A launcher waits on this file before letting prefill fill the pool.""" + ready_file = tmp_path / "ready" + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running( + donation, + "--master_server_address", + "10.0.0.1:50051", + "--segment_size", + "2GiB", + "--ready_file", + str(ready_file), + ).run() + + assert ready_file.read_text() == f"{donation.host} {2 * 1024**3}\n" + + # ---- describing the pool from the master's own run directory ---- From 36c55d09368d6f62e7e01e0716aa0b263a277446 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 23 Sep 2026 06:29:46 +0000 Subject: [PATCH 15/24] Logging for scheduler deadlock Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../pyexecutor/scheduler/scheduler_v2.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index a4c1acd50dec..c11210f6b5ed 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -15,6 +15,7 @@ import enum import os +from collections import Counter from typing import Callable, Optional from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy @@ -1504,6 +1505,11 @@ def _detect_deadlock( generation requests at all. """ if made_progress: + if self._stalled_schedules: + logger.debug( + f"[V2Scheduler] Stall cleared after {self._stalled_schedules} pass(es); " + "scheduling resumed." + ) self._stalled_schedules = 0 return @@ -1522,6 +1528,11 @@ def _detect_deadlock( ) if num_gen_candidates == 0 and num_ctx_candidates == 0: # Legitimately idle: nothing to schedule. + if self._stalled_schedules: + logger.debug( + f"[V2Scheduler] Stall cleared after {self._stalled_schedules} pass(es); " + "no candidates remain." + ) self._stalled_schedules = 0 return @@ -1547,11 +1558,30 @@ def _detect_deadlock( and self._async_transfer_manager.has_any_inflight_requests() ) if transfer_holding: + if self._stalled_schedules: + logger.debug( + f"[V2Scheduler] {num_gen_candidates} generation and " + f"{num_ctx_candidates} context request(s) cannot allocate, but KV " + "transfers in flight will release pages; stall cleared." + ) self._stalled_schedules = 0 return self._stalled_schedules += 1 if self._stalled_schedules < self._DEADLOCK_STALL_ITERS: + # Warn as the stall builds so the raise is not a surprise. A quarter + # of the threshold keeps this to a handful of lines beforehand. + if self._stalled_schedules == 1: + logger.debug( + "[V2Scheduler] Scheduling stall started: " + f"{self._summarize_stall(active_requests, pending_ctx, inflight_request_ids, preempted_ids)}" + ) + elif self._stalled_schedules % max(1, self._DEADLOCK_STALL_ITERS // 4) == 0: + logger.warning( + f"[V2Scheduler] Stalled {self._stalled_schedules}/" + f"{self._DEADLOCK_STALL_ITERS} consecutive passes: " + f"{self._summarize_stall(active_requests, pending_ctx, inflight_request_ids, preempted_ids)}" + ) return # A connector rejects every tier below GPU at bring-up @@ -1570,6 +1600,13 @@ def _detect_deadlock( "Configure kv_cache_config.host_cache_size, increase " "kv_cache_config.max_tokens, or lower max_batch_size." ) + # Logged before the raise so the diagnostic survives even if the + # exception is caught or truncated upstream. + logger.warning( + f"[V2Scheduler] Declaring deadlock after {self._stalled_schedules} passes: " + f"{self._summarize_stall(active_requests, pending_ctx, inflight_request_ids, preempted_ids)} " + f"{remedy}" + ) raise RuntimeError( f"V2 scheduler deadlock: {num_gen_candidates} generation and " f"{num_ctx_candidates} context request(s) active but none could " @@ -1578,6 +1615,45 @@ def _detect_deadlock( f"pool is likely exhausted. {remedy}" ) + def _summarize_stall( + self, + active_requests: RequestList, + pending_ctx: RequestList, + inflight_request_ids: set[int], + preempted_ids: set[int], + ) -> str: + """One-line breakdown of a stalled pass for the deadlock logs. + + Re-derives the blocked candidates so the per-pass detector stays cheap: + this runs only on the rare logging branches. + """ + gen = [ + r + 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 + ] + ctx = [ + r + for r in pending_ctx + if r.py_request_id not in preempted_ids and r.request_id not in inflight_request_ids + ] + states = Counter(LlmRequestState(r.state_value).name for r in gen + ctx) + sample_ids = [r.py_request_id for r in (gen + ctx)[:5]] + has_connector = getattr(self.kv_cache_manager, "kv_connector_manager", None) is not None + transfers_in_flight = ( + self._async_transfer_manager is not None + and self._async_transfer_manager.has_any_inflight_requests() + ) + return ( + f"{len(gen)} generation and {len(ctx)} context request(s) blocked; " + f"states={dict(states)}; sample_ids={sample_ids}; " + f"kv_connector={has_connector}; " + f"cache_tier_below_gpu={self.kv_cache_manager.has_cache_tier_below_gpu}; " + f"transfers_in_flight={transfers_in_flight}" + ) + def _free_kv_caches(self, req: LlmRequest) -> None: self.kv_cache_manager.free_resources(req) if self.draft_kv_cache_manager is not None: From 3f81b48334df5be7d85cd645b4018829fc0109eb Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 23 Sep 2026 13:46:18 -0700 Subject: [PATCH 16/24] [None][chore] Share the mooncake client pin between the wheel and container Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/Dockerfile.multi | 3 ++- docker/common/install_mooncake.sh | 19 +++++++++++-------- requirements-mooncake.txt | 19 +++++++++++++++++++ setup.py | 5 +++++ .../connectors/mooncake_store/__init__.py | 4 ++-- .../connectors/mooncake_store/donor.py | 6 +++--- 6 files changed, 42 insertions(+), 14 deletions(-) create mode 100644 requirements-mooncake.txt diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index f61052c605a3..68340b65c81e 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -133,6 +133,7 @@ COPY --from=triton /opt/tritonserver/caches /opt/tritonserver/caches # Install Triton deps, Mooncake, and CI tooling packages RUN --mount=type=bind,source=docker/common,target=/opt/docker/common \ + --mount=type=bind,source=requirements-mooncake.txt,target=/opt/requirements-mooncake.txt \ --mount=type=secret,id=github_clone_token \ . /opt/docker/common/github_auth.sh && \ GITHUB_MIRROR=${GITHUB_MIRROR} bash /opt/docker/common/install_triton.sh && \ @@ -151,7 +152,7 @@ COPY scripts scripts COPY tensorrt_llm tensorrt_llm COPY triton_kernels triton_kernels COPY 3rdparty 3rdparty -COPY setup.py requirements.txt requirements-dev.txt requirements-grpc-smg.txt requirements-openengine.txt constraints.txt LICENSE README.md ./ +COPY setup.py requirements.txt requirements-dev.txt requirements-grpc-smg.txt requirements-openengine.txt requirements-mooncake.txt constraints.txt LICENSE README.md ./ ENV CCACHE_DIR=/root/.cache/ccache # Build the TRT-LLM wheel diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index bc61ae8d9034..f47af3189cf2 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -70,14 +70,17 @@ MOONCAKE_CMAKE_PACKAGE="$(python3 -c "import sys; print([s for s in sys.path if echo "removing CMake-generated mooncake package: ${MOONCAKE_CMAKE_PACKAGE}" rm -rf "${MOONCAKE_CMAKE_PACKAGE}" -# 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, 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}" +# The Mooncake Python client (MooncakeDistributedStore) is installed from +# requirements-mooncake.txt, the pin shared with the standalone wheel's +# `tensorrt-llm[mooncake]` extra. That file explains why the pin is independent +# of MOONCAKE_VERSION above and why the CUDA 13 wheel is required. +# The Docker build bind-mounts the file at /opt; a run from a checkout finds it +# relative to this script instead. +MOONCAKE_REQUIREMENTS="${MOONCAKE_REQUIREMENTS:-/opt/requirements-mooncake.txt}" +if [ ! -f "${MOONCAKE_REQUIREMENTS}" ]; then + MOONCAKE_REQUIREMENTS="$(dirname "${BASH_SOURCE[0]}")/../../requirements-mooncake.txt" +fi +pip3 install --no-cache-dir -r "${MOONCAKE_REQUIREMENTS}" # Fail the build rather than ship an image whose import is broken. python3 - <<'PY' diff --git a/requirements-mooncake.txt b/requirements-mooncake.txt new file mode 100644 index 000000000000..5daa279bade0 --- /dev/null +++ b/requirements-mooncake.txt @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Pins the Mooncake Python client (MooncakeDistributedStore, imported as +# `mooncake.store`) used by the mooncake-store KV cache connector and the +# `trtllm-serve mooncake_master` / `mooncake_donor` commands. +# +# This is the single source of truth for the client version: the standalone +# wheel exposes it through the `tensorrt-llm[mooncake]` extra (see setup.py) and +# the container installs this same file from docker/common/install_mooncake.sh, +# so both agree on one pin. The C++ transfer engine that the cache transceiver +# links against is a separate component built from source in that script. +# +# The plain `mooncake-transfer-engine` wheel is built against CUDA 12, while +# TensorRT-LLM ships CUDA 13 (see cuda-python>=13 in requirements.txt), so its +# extensions cannot resolve libcudart.so.12. `mooncake-transfer-engine-cuda13` +# is the same project built for CUDA 13; it is versioned independently, with +# releases starting at 0.3.9, so its version does not track the source build. +mooncake-transfer-engine-cuda13==0.3.13 diff --git a/setup.py b/setup.py index b14b2588cb07..8d122d9db0d6 100644 --- a/setup.py +++ b/setup.py @@ -156,6 +156,10 @@ def has_ext_modules(self): # may carry gateway-specific options (such as an --extra-index-url) without # affecting the default dependency graph. grpc_smg_deps, _ = parse_requirements(Path("requirements-grpc-smg.txt")) +# The mooncake-store KV cache connector is opt-in: its Python client is heavy +# and CUDA-ABI specific, so the default installation must not carry it. The pin +# lives in requirements-mooncake.txt, the same file the container build reads. +mooncake_deps, _ = parse_requirements(Path("requirements-mooncake.txt")) constraints_file = Path("constraints.txt") if constraints_file.exists(): constraints, _ = parse_requirements(constraints_file) @@ -706,6 +710,7 @@ def get_build_state_options(): "openengine": openengine_deps, "mx": mx_deps, "grpc-smg": grpc_smg_deps, + "mooncake": mooncake_deps, }, zip_safe=True, install_requires=required_deps, diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index 9c122862d0f5..0f5c88b86911 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -29,8 +29,8 @@ processes that open a store handle, which in a disaggregated deployment is the context servers alone, so `donor.py` lends a node's memory to the pool without giving it a connector. `trtllm-serve mooncake_master` and `mooncake_donor` -expose both. Both need the Mooncake Python bindings (`pip install -mooncake-transfer-engine`). +expose both. Both need the Mooncake Python bindings, an optional dependency +installed with `pip install tensorrt-llm[mooncake]`. `keys.py` and `staging.py` hold what the store side shares with the connector that moves pages in and out of the pool: how a block of tokens becomes a store diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py index 1373b82f08ee..542936bb11fe 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -74,9 +74,9 @@ def donate_segment( 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}" + "Donating memory needs the Mooncake Python bindings, an optional " + "dependency: install them with `pip install tensorrt-llm[mooncake]`. " + f"The C++ transfer engine is a separate component: {exc}" ) from exc host = hostname or local_address() From 0284774d6f0681dc2d93b7d616f9afc37766d50a Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:14:25 -0700 Subject: [PATCH 17/24] [None][fix] Drop the Mooncake import check from the container build The check imported mooncake.store, whose extension links against libcuda.so.1. That library comes from the host driver and is injected by the NVIDIA container runtime, so it is absent on every image build host and the check could never pass. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/common/install_mooncake.sh | 8 -------- 1 file changed, 8 deletions(-) diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index f47af3189cf2..d61b9bfa2678 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -81,11 +81,3 @@ if [ ! -f "${MOONCAKE_REQUIREMENTS}" ]; then MOONCAKE_REQUIREMENTS="$(dirname "${BASH_SOURCE[0]}")/../../requirements-mooncake.txt" fi pip3 install --no-cache-dir -r "${MOONCAKE_REQUIREMENTS}" - -# Fail the build rather than ship an image whose import is broken. -python3 - <<'PY' -import mooncake.store - -mooncake.store.MooncakeDistributedStore() -print(f"mooncake.store OK: {mooncake.store.__file__}") -PY From ba8aca4cd3fcd203a02a83a916de5ef4af848edf Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Mon, 28 Sep 2026 17:03:58 -0700 Subject: [PATCH 18/24] [None][fix] Install the Mooncake client in the release image install_mooncake.sh runs in tritondevel, and release starts FROM DEVEL_IMAGE. That is tritondevel for the internal trtllm image, so it inherits Mooncake there, but ngc-release names the NGC devel image, which is built from the devel stage. The published NGC container therefore had neither the mooncake.store bindings the donor needs nor the mooncake_master binary that mooncake_master launches, so both commands failed in the one image most deployments start from. The release stage installs the wheel's mooncake extra, which is the same requirements-mooncake.txt pin the container build already uses, so the two cannot drift. Where DEVEL_IMAGE already carries it the install is a no-op. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/Dockerfile.multi | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 68340b65c81e..6a53c0405c09 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -182,7 +182,12 @@ RUN --mount=type=cache,target=/root/.cache/pip --mount=type=bind,from=wheel,sour fi && \ TRTLLM_WHEEL=$(find /tmp/wheel -maxdepth 1 -name 'tensorrt_llm*.whl' -print -quit) && \ test -n "${TRTLLM_WHEEL}" && \ - pip install "${TRTLLM_WHEEL}[mx]" + # The mooncake extra carries the Python client and the mooncake_master + # binary that the mooncake_master and mooncake_donor commands need. It is + # installed here rather than inherited, because release starts from + # DEVEL_IMAGE and the NGC image names the devel stage, not the tritondevel + # one that builds Mooncake. Where it is inherited this is already satisfied. + pip install "${TRTLLM_WHEEL}[mx,mooncake]" RUN --mount=type=bind,source=README.md,target=/mnt/ctx/README.md \ --mount=type=bind,source=docs,target=/mnt/ctx/docs \ From 929f0b8b098c620a7c97e6a2a157fe8f68b31cca Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 29 Sep 2026 05:43:16 +0000 Subject: [PATCH 19/24] [None][infra] Update CI image tags for the Mooncake client install Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- jenkins/current_image_tags.properties | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties index b615eb04d28a..4e0819926677 100644 --- a/jenkins/current_image_tags.properties +++ b/jenkins/current_image_tags.properties @@ -13,8 +13,8 @@ # images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead. IMAGE_NAME=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm -LLM_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.08-py3-x86_64-ubuntu24.04-skip-tritondevel-202609210610-19345 -LLM_SBSA_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.08-py3-sbsa-ubuntu24.04-skip-tritondevel-202609210610-19345 -LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py310-skip-tritondevel-202609210610-19345 -LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py312-skip-tritondevel-202609210610-19345 -LLM_SBSA_WHEEL_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-skip-tritondevel-202609210610-19345 +LLM_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.08-py3-x86_64-ubuntu24.04-skip-tritondevel-202609290532-19235 +LLM_SBSA_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.08-py3-sbsa-ubuntu24.04-skip-tritondevel-202609290532-19235 +LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py310-skip-tritondevel-202609290532-19235 +LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py312-skip-tritondevel-202609290532-19235 +LLM_SBSA_WHEEL_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-skip-tritondevel-202609290532-19235 From 27e6b0b6e59da3e6a2f320c533965c728690f389 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 29 Sep 2026 12:03:28 -0700 Subject: [PATCH 20/24] [None][chore] Keep the Mooncake client out of the container images Nothing in a container build or in CI needs the Mooncake Python client. Nothing compiles or links against it, `mooncake.store` is imported inside the one function that uses it, and the four store tests stub the module and fake `shutil.which` rather than importing either. Installing it in tritondevel only added 98MB to every CI image and turned each pin bump into an image rebuild and a current_image_tags.properties update, which then has to be sequenced against the other image PRs in flight. So the client is now reached only through `tensorrt-llm[mooncake]`, the way grpc-smg and openengine already are, and no shipped container carries the pin. install_mooncake.sh, current_image_tags.properties and the attribution metadata go back to what they were: the C++ transfer engine the cache transceiver links against is unaffected and still builds from source there, and attribution covers container-shipped components only. The one build file left is the wheel stage's COPY, which setup.py needs to read the pin. launch_master pointed at install_mooncake.sh for its missing binary, which no longer installs it, so it names the extra instead. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/Dockerfile.multi | 8 +----- docker/common/install_mooncake.sh | 28 ------------------- .../source/installation/installation-guide.md | 16 +++++++++++ jenkins/current_image_tags.properties | 10 +++---- requirements-mooncake.txt | 19 ++++--------- .../attribution/scan/metadata/mooncake.yml | 5 +--- setup.py | 3 +- .../connectors/mooncake_store/master.py | 5 ++-- 8 files changed, 32 insertions(+), 62 deletions(-) diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 6a53c0405c09..d2c4b0f9c515 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -133,7 +133,6 @@ COPY --from=triton /opt/tritonserver/caches /opt/tritonserver/caches # Install Triton deps, Mooncake, and CI tooling packages RUN --mount=type=bind,source=docker/common,target=/opt/docker/common \ - --mount=type=bind,source=requirements-mooncake.txt,target=/opt/requirements-mooncake.txt \ --mount=type=secret,id=github_clone_token \ . /opt/docker/common/github_auth.sh && \ GITHUB_MIRROR=${GITHUB_MIRROR} bash /opt/docker/common/install_triton.sh && \ @@ -182,12 +181,7 @@ RUN --mount=type=cache,target=/root/.cache/pip --mount=type=bind,from=wheel,sour fi && \ TRTLLM_WHEEL=$(find /tmp/wheel -maxdepth 1 -name 'tensorrt_llm*.whl' -print -quit) && \ test -n "${TRTLLM_WHEEL}" && \ - # The mooncake extra carries the Python client and the mooncake_master - # binary that the mooncake_master and mooncake_donor commands need. It is - # installed here rather than inherited, because release starts from - # DEVEL_IMAGE and the NGC image names the devel stage, not the tritondevel - # one that builds Mooncake. Where it is inherited this is already satisfied. - pip install "${TRTLLM_WHEEL}[mx,mooncake]" + pip install "${TRTLLM_WHEEL}[mx]" RUN --mount=type=bind,source=README.md,target=/mnt/ctx/README.md \ --mount=type=bind,source=docs,target=/mnt/ctx/docs \ diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index d61b9bfa2678..0935c7f31a2f 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -53,31 +53,3 @@ cd ../.. rm -rf Mooncake echo "export LD_LIBRARY_PATH=${MOONCAKE_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" - -# 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 below. -# -# `make install` also emits a `mooncake` package that omits -# libmooncake_store.so, so importing mooncake.store from it fails. It has to go -# before the wheel is installed: CMake writes -# store.cpython-312-x86_64-linux-gnu.so where the wheel writes store.so, and -# importlib prefers the interpreter-tagged suffix, so the broken extension -# would win even after pip reports success. The directory is the one -# mooncake-integration/CMakeLists.txt chose, which this repeats. -MOONCAKE_CMAKE_PACKAGE="$(python3 -c "import sys; print([s for s in sys.path if 'packages' in s][0])")/mooncake" -echo "removing CMake-generated mooncake package: ${MOONCAKE_CMAKE_PACKAGE}" -rm -rf "${MOONCAKE_CMAKE_PACKAGE}" - -# The Mooncake Python client (MooncakeDistributedStore) is installed from -# requirements-mooncake.txt, the pin shared with the standalone wheel's -# `tensorrt-llm[mooncake]` extra. That file explains why the pin is independent -# of MOONCAKE_VERSION above and why the CUDA 13 wheel is required. -# The Docker build bind-mounts the file at /opt; a run from a checkout finds it -# relative to this script instead. -MOONCAKE_REQUIREMENTS="${MOONCAKE_REQUIREMENTS:-/opt/requirements-mooncake.txt}" -if [ ! -f "${MOONCAKE_REQUIREMENTS}" ]; then - MOONCAKE_REQUIREMENTS="$(dirname "${BASH_SOURCE[0]}")/../../requirements-mooncake.txt" -fi -pip3 install --no-cache-dir -r "${MOONCAKE_REQUIREMENTS}" diff --git a/docs/source/installation/installation-guide.md b/docs/source/installation/installation-guide.md index 0e4bea8511d3..d14bf383217d 100644 --- a/docs/source/installation/installation-guide.md +++ b/docs/source/installation/installation-guide.md @@ -37,6 +37,10 @@ Sanity check the installation by running the following inside the container: python3 -c "import tensorrt_llm" ``` +The container carries the default dependency set. Features behind a `pip` extra, such as the Mooncake Python +client, are not pre-installed and are added the same way as in a standalone installation; see +[Install optional extras](#install-optional-extras). + (linux)= ## Option 2: Install on Linux via `pip` @@ -89,6 +93,18 @@ pip3 uninstall nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base \ pip3 install --ignore-installed pip setuptools wheel && pip3 install tensorrt_llm ``` +#### Install optional extras + +A few features are reached through `pip` extras rather than the default dependency set, so an installation +that does not use them does not carry their dependencies either. + +The `mooncake` extra installs the Mooncake Python client, which the mooncake-store KV cache connector and the +`trtllm-serve mooncake_master` and `trtllm-serve mooncake_donor` commands need: + +```bash +pip3 install tensorrt_llm[mooncake] +``` + #### Install a nightly release Nightly wheels are published to a dedicated package index. Query that index to list compatible versions, and use diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties index 4e0819926677..b615eb04d28a 100644 --- a/jenkins/current_image_tags.properties +++ b/jenkins/current_image_tags.properties @@ -13,8 +13,8 @@ # images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead. IMAGE_NAME=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm -LLM_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.08-py3-x86_64-ubuntu24.04-skip-tritondevel-202609290532-19235 -LLM_SBSA_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.08-py3-sbsa-ubuntu24.04-skip-tritondevel-202609290532-19235 -LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py310-skip-tritondevel-202609290532-19235 -LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py312-skip-tritondevel-202609290532-19235 -LLM_SBSA_WHEEL_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-skip-tritondevel-202609290532-19235 +LLM_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.08-py3-x86_64-ubuntu24.04-skip-tritondevel-202609210610-19345 +LLM_SBSA_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.08-py3-sbsa-ubuntu24.04-skip-tritondevel-202609210610-19345 +LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py310-skip-tritondevel-202609210610-19345 +LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-rocky8-x86_64-rocky8-py312-skip-tritondevel-202609210610-19345 +LLM_SBSA_WHEEL_DOCKER_IMAGE=artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:cuda-13.4.1-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-skip-tritondevel-202609210610-19345 diff --git a/requirements-mooncake.txt b/requirements-mooncake.txt index 5daa279bade0..cecbb743ce47 100644 --- a/requirements-mooncake.txt +++ b/requirements-mooncake.txt @@ -2,18 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 # Pins the Mooncake Python client (MooncakeDistributedStore, imported as -# `mooncake.store`) used by the mooncake-store KV cache connector and the -# `trtllm-serve mooncake_master` / `mooncake_donor` commands. -# -# This is the single source of truth for the client version: the standalone -# wheel exposes it through the `tensorrt-llm[mooncake]` extra (see setup.py) and -# the container installs this same file from docker/common/install_mooncake.sh, -# so both agree on one pin. The C++ transfer engine that the cache transceiver -# links against is a separate component built from source in that script. -# -# The plain `mooncake-transfer-engine` wheel is built against CUDA 12, while -# TensorRT-LLM ships CUDA 13 (see cuda-python>=13 in requirements.txt), so its -# extensions cannot resolve libcudart.so.12. `mooncake-transfer-engine-cuda13` -# is the same project built for CUDA 13; it is versioned independently, with -# releases starting at 0.3.9, so its version does not track the source build. +# `mooncake.store`) behind the `tensorrt-llm[mooncake]` extra, used by the +# mooncake-store KV cache connector and the `trtllm-serve mooncake_master` / +# `mooncake_donor` commands. The `-cuda13` variant is the one built for the +# CUDA version TensorRT-LLM ships; it is versioned independently of the C++ +# transfer engine that docker/common/install_mooncake.sh builds from source. mooncake-transfer-engine-cuda13==0.3.13 diff --git a/scripts/attribution/scan/metadata/mooncake.yml b/scripts/attribution/scan/metadata/mooncake.yml index b1e6dab2c141..c8d51e0f81b8 100644 --- a/scripts/attribution/scan/metadata/mooncake.yml +++ b/scripts/attribution/scan/metadata/mooncake.yml @@ -1,8 +1,5 @@ name: mooncake -description: Mooncake transfer engine and distributed store for distributed KV cache +description: Mooncake transfer engine for distributed KV cache source: container directory_matches: - /usr/local/Mooncake -- mooncake -basename_matches: -- mooncake_transfer_engine diff --git a/setup.py b/setup.py index 8d122d9db0d6..0b101071cbd6 100644 --- a/setup.py +++ b/setup.py @@ -157,8 +157,7 @@ def has_ext_modules(self): # affecting the default dependency graph. grpc_smg_deps, _ = parse_requirements(Path("requirements-grpc-smg.txt")) # The mooncake-store KV cache connector is opt-in: its Python client is heavy -# and CUDA-ABI specific, so the default installation must not carry it. The pin -# lives in requirements-mooncake.txt, the same file the container build reads. +# and CUDA-ABI specific, so no container and no default installation carries it. mooncake_deps, _ = parse_requirements(Path("requirements-mooncake.txt")) constraints_file = Path("constraints.txt") if constraints_file.exists(): diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index b5b96a7301b8..555cdcc06a6f 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -459,8 +459,9 @@ def _launch_master(pool: PoolSpec, run_dir: str) -> LaunchedMaster: 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 " - "docker/common/install_mooncake.sh installs. Point " + "Mooncake master. It comes from the Mooncake Python client, an " + "optional dependency: install it with " + "`pip install tensorrt-llm[mooncake]`. Otherwise point " f"{MASTER_BINARY_ENV} at the binary, or drop launch_master and " "set master_server_address to a master you run yourself." ) From a5142cae1ebcffb81cf78b549c3dfbfe6e4cfe34 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 29 Sep 2026 12:03:53 -0700 Subject: [PATCH 21/24] [None][fix] Withdraw the donor's ready file along with the segment --ready_file is what a launcher waits on before letting prefill write into the pool, but the donor wrote it and left it there. The file outlived the segment it described, so a launcher reusing the path saw the gate already cleared and could start prefill against capacity this donor had stopped contributing. The signal-exit test read the file after the command returned, so it passed either way. The write moves into a context manager entered after `as host`, which puts the file's lifetime inside the segment's on both ends. A launcher that consumes the file by removing it is not an error on the way out. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tensorrt_llm/commands/mooncake.py | 45 +++++++++++++++---- .../executor/test_mooncake_store_cli.py | 30 ++++++++++--- 2 files changed, 61 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index 4ef6fda20ca6..47b3b8e32fcd 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -92,6 +92,39 @@ def _signal_handoff(): raise +@contextlib.contextmanager +def _announce_ready(ready_file: Optional[str], host: str, segment_size: int): + """Hold a readiness file for exactly as long as the segment is mounted. + + Launchers gate prefill on this file existing, so it has to go when the + capacity does. Left behind, it would clear that gate for a pool the donor + no longer contributes to, and prefill would start writing into a pool + short of the memory the launcher waited for. + """ + if ready_file is None: + yield + return + + with open(ready_file, "w") as handle: + handle.write(f"{host} {segment_size}\n") + logger.info( + f"mooncake-store: announced this segment in {ready_file}, so a " + "launcher waiting on the pool's capacity can proceed" + ) + try: + yield + finally: + # A launcher that consumes the file by removing it has already had the + # signal, so its absence here is a success and not worth failing the + # shutdown over. + with contextlib.suppress(OSError): + os.unlink(ready_file) + logger.info( + f"mooncake-store: withdrew the announcement in {ready_file}; a " + "launcher waiting on this pool's capacity will wait again" + ) + + @click.command("mooncake_master") @click.option( "--rpc_port", @@ -363,16 +396,10 @@ def size_option(name: str, value: str) -> int: ), local_buffer_size=buffer_size, ) as host, + # After `as host`, so the announcement names the segment that mounted + # and is withdrawn before the segment is. + _announce_ready(ready_file, host, donating), ): - 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 traffic # client, which is what donation exists to avoid. started = time.monotonic() diff --git a/tests/unittest/_torch/executor/test_mooncake_store_cli.py b/tests/unittest/_torch/executor/test_mooncake_store_cli.py index dfeadfa7b9c2..1f9e6a0f5c59 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_cli.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_cli.py @@ -155,16 +155,22 @@ def command(request, master, donor) -> CommandUnderTest: return {"mooncake_master": master, "mooncake_donor": donor}[request.param] -def signal_on_idle(monkeypatch, number: int) -> None: +def signal_on_idle(monkeypatch, number: int, on_idle=None) -> None: """Deliver `number` to this process the first time a command idles. The handler runs at the next bytecode boundary in the main thread, so the loop here only has to give the interpreter one. It bounds the wait rather than blocking, so a handler that never fires fails the test rather than hanging it. + + `on_idle` runs before the signal, which is the only point at which a + command's resources are all held. Anything a command releases on its way + out has to be observed from there rather than after `run` returns. """ def sleep(_seconds): + if on_idle is not None: + on_idle() os.kill(os.getpid(), number) for _ in range(100): time.sleep(0.01) @@ -422,10 +428,23 @@ def test_the_donor_defaults_the_protocol_to_rdma_when_the_config_value_is_null( assert donation.resolved["protocol"] == "rdma" -def test_the_donor_announces_the_mounted_segment_in_its_ready_file(monkeypatch, donation, tmp_path): - """A launcher waits on this file before letting prefill fill the pool.""" +def test_the_donors_ready_file_lasts_exactly_as_long_as_the_segment( + monkeypatch, donation, tmp_path +): + """A launcher waits on this file before letting prefill fill the pool. + + So the file has to appear once the segment is mounted and go when it is + withdrawn. Outliving the segment would clear that gate for capacity this + donor no longer contributes, and prefill would start writing into a pool + short of the memory the launcher waited for. + """ ready_file = tmp_path / "ready" - signal_on_idle(monkeypatch, signal.SIGTERM) + announced = [] + signal_on_idle( + monkeypatch, + signal.SIGTERM, + on_idle=lambda: announced.append(ready_file.read_text()), + ) with pytest.raises(_telemetry.SignalExit): donor_running( @@ -438,7 +457,8 @@ def test_the_donor_announces_the_mounted_segment_in_its_ready_file(monkeypatch, str(ready_file), ).run() - assert ready_file.read_text() == f"{donation.host} {2 * 1024**3}\n" + assert announced == [f"{donation.host} {2 * 1024**3}\n"] + assert not ready_file.exists() # ---- describing the pool from the master's own run directory ---- From 45bb57d997d69b79762fbc5e015c781ce36d3033 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 29 Sep 2026 12:53:39 -0700 Subject: [PATCH 22/24] [None][fix] Match the install the donor's import error actually names The error stopped naming `pip install mooncake-transfer-engine` when it started naming the `tensorrt-llm[mooncake]` extra, but the test kept matching the old package, so it could not pass. The brackets are escaped because `pytest.raises(match=...)` searches a regex, where `[mooncake]` would be a character class. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tests/unittest/_torch/executor/test_mooncake_store_donor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py index 0fa98a8da1f1..ea51408b622f 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_donor.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -128,7 +128,7 @@ def test_missing_bindings_are_reported_as_the_separate_component_they_are(monkey monkeypatch.setitem(sys.modules, "mooncake", None) monkeypatch.setitem(sys.modules, "mooncake.store", None) - with pytest.raises(ImportError, match="mooncake-transfer-engine"): + with pytest.raises(ImportError, match=r"tensorrt-llm\[mooncake\]"): with donate_segment("10.0.0.1:50051", GIB): pytest.fail("donation should not have yielded") From 4df81102dbb6de764c05e5c5a667764fd5efdf34 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 30 Sep 2026 15:33:24 -0700 Subject: [PATCH 23/24] [None][chore] Install the Mooncake client with the default dependency set Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- docker/Dockerfile.multi | 2 +- docs/source/installation/installation-guide.md | 16 ---------------- requirements-mooncake.txt | 10 ---------- requirements.txt | 1 + setup.py | 4 ---- .../connectors/mooncake_store/__init__.py | 4 ++-- .../connectors/mooncake_store/donor.py | 7 ++++--- .../connectors/mooncake_store/master.py | 10 +++++----- .../_torch/executor/test_mooncake_store_donor.py | 2 +- 9 files changed, 14 insertions(+), 42 deletions(-) delete mode 100644 requirements-mooncake.txt diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index d2c4b0f9c515..f61052c605a3 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -151,7 +151,7 @@ COPY scripts scripts COPY tensorrt_llm tensorrt_llm COPY triton_kernels triton_kernels COPY 3rdparty 3rdparty -COPY setup.py requirements.txt requirements-dev.txt requirements-grpc-smg.txt requirements-openengine.txt requirements-mooncake.txt constraints.txt LICENSE README.md ./ +COPY setup.py requirements.txt requirements-dev.txt requirements-grpc-smg.txt requirements-openengine.txt constraints.txt LICENSE README.md ./ ENV CCACHE_DIR=/root/.cache/ccache # Build the TRT-LLM wheel diff --git a/docs/source/installation/installation-guide.md b/docs/source/installation/installation-guide.md index d14bf383217d..0e4bea8511d3 100644 --- a/docs/source/installation/installation-guide.md +++ b/docs/source/installation/installation-guide.md @@ -37,10 +37,6 @@ Sanity check the installation by running the following inside the container: python3 -c "import tensorrt_llm" ``` -The container carries the default dependency set. Features behind a `pip` extra, such as the Mooncake Python -client, are not pre-installed and are added the same way as in a standalone installation; see -[Install optional extras](#install-optional-extras). - (linux)= ## Option 2: Install on Linux via `pip` @@ -93,18 +89,6 @@ pip3 uninstall nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base \ pip3 install --ignore-installed pip setuptools wheel && pip3 install tensorrt_llm ``` -#### Install optional extras - -A few features are reached through `pip` extras rather than the default dependency set, so an installation -that does not use them does not carry their dependencies either. - -The `mooncake` extra installs the Mooncake Python client, which the mooncake-store KV cache connector and the -`trtllm-serve mooncake_master` and `trtllm-serve mooncake_donor` commands need: - -```bash -pip3 install tensorrt_llm[mooncake] -``` - #### Install a nightly release Nightly wheels are published to a dedicated package index. Query that index to list compatible versions, and use diff --git a/requirements-mooncake.txt b/requirements-mooncake.txt deleted file mode 100644 index cecbb743ce47..000000000000 --- a/requirements-mooncake.txt +++ /dev/null @@ -1,10 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Pins the Mooncake Python client (MooncakeDistributedStore, imported as -# `mooncake.store`) behind the `tensorrt-llm[mooncake]` extra, used by the -# mooncake-store KV cache connector and the `trtllm-serve mooncake_master` / -# `mooncake_donor` commands. The `-cuda13` variant is the one built for the -# CUDA version TensorRT-LLM ships; it is versioned independently of the C++ -# transfer engine that docker/common/install_mooncake.sh builds from source. -mooncake-transfer-engine-cuda13==0.3.13 diff --git a/requirements.txt b/requirements.txt index d3945937aef6..308cfe9b386b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -116,3 +116,4 @@ cache-dit>=1.3.5 librosa msgpack uvloop>=0.19.0 +mooncake-transfer-engine-cuda13==0.3.13 diff --git a/setup.py b/setup.py index 0b101071cbd6..b14b2588cb07 100644 --- a/setup.py +++ b/setup.py @@ -156,9 +156,6 @@ def has_ext_modules(self): # may carry gateway-specific options (such as an --extra-index-url) without # affecting the default dependency graph. grpc_smg_deps, _ = parse_requirements(Path("requirements-grpc-smg.txt")) -# The mooncake-store KV cache connector is opt-in: its Python client is heavy -# and CUDA-ABI specific, so no container and no default installation carries it. -mooncake_deps, _ = parse_requirements(Path("requirements-mooncake.txt")) constraints_file = Path("constraints.txt") if constraints_file.exists(): constraints, _ = parse_requirements(constraints_file) @@ -709,7 +706,6 @@ def get_build_state_options(): "openengine": openengine_deps, "mx": mx_deps, "grpc-smg": grpc_smg_deps, - "mooncake": mooncake_deps, }, zip_safe=True, install_requires=required_deps, diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index 0f5c88b86911..e5219cbb3d05 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -29,8 +29,8 @@ processes that open a store handle, which in a disaggregated deployment is the context servers alone, so `donor.py` lends a node's memory to the pool without giving it a connector. `trtllm-serve mooncake_master` and `mooncake_donor` -expose both. Both need the Mooncake Python bindings, an optional dependency -installed with `pip install tensorrt-llm[mooncake]`. +expose both. Both need the Mooncake Python bindings, which `tensorrt-llm` pulls +in as `mooncake-transfer-engine-cuda13`. `keys.py` and `staging.py` hold what the store side shares with the connector that moves pages in and out of the pool: how a block of tokens becomes a store diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py index 542936bb11fe..071390ca286d 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -74,9 +74,10 @@ def donate_segment( from mooncake.store import MooncakeDistributedStore except ImportError as exc: raise ImportError( - "Donating memory needs the Mooncake Python bindings, an optional " - "dependency: install them with `pip install tensorrt-llm[mooncake]`. " - f"The C++ transfer engine is a separate component: {exc}" + "Donating memory needs the Mooncake Python bindings, which " + "`tensorrt-llm` pulls in as `mooncake-transfer-engine-cuda13`, so " + "reinstall that package if this environment dropped it. The C++ " + f"transfer engine is a separate component: {exc}" ) from exc host = hostname or local_address() diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index 555cdcc06a6f..f2de46a10feb 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -459,11 +459,11 @@ def _launch_master(pool: PoolSpec, run_dir: str) -> LaunchedMaster: if resolved is None: raise FileNotFoundError( f"{binary!r} is not on PATH, so launch_master cannot start a " - "Mooncake master. It comes from the Mooncake Python client, an " - "optional dependency: install it with " - "`pip install tensorrt-llm[mooncake]`. Otherwise point " - f"{MASTER_BINARY_ENV} at the binary, or drop launch_master and " - "set master_server_address to a master you run yourself." + "Mooncake master. It comes from the Mooncake Python client, which " + "`tensorrt-llm` pulls in as `mooncake-transfer-engine-cuda13`, so " + "reinstall that package if this environment dropped it. Otherwise " + f"point {MASTER_BINARY_ENV} at the binary, or drop launch_master " + "and set master_server_address to a master you run yourself." ) host = local_address() diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py index ea51408b622f..36cbc3aded57 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_donor.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -128,7 +128,7 @@ def test_missing_bindings_are_reported_as_the_separate_component_they_are(monkey monkeypatch.setitem(sys.modules, "mooncake", None) monkeypatch.setitem(sys.modules, "mooncake.store", None) - with pytest.raises(ImportError, match=r"tensorrt-llm\[mooncake\]"): + with pytest.raises(ImportError, match=r"mooncake-transfer-engine-cuda13"): with donate_segment("10.0.0.1:50051", GIB): pytest.fail("donation should not have yielded") From 6b235123e1065f7c902bb8fa7cc6ee6a41ca2ad9 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:23:03 -0700 Subject: [PATCH 24/24] [None][fix] Let each Mooncake client discover its own RDMA device Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../connectors/mooncake_store/__init__.py | 2 - .../connectors/mooncake_store/master.py | 87 ++----------------- .../executor/test_mooncake_store_master.py | 55 ++++-------- 3 files changed, 23 insertions(+), 121 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py index e5219cbb3d05..1bf4602e2088 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -46,7 +46,6 @@ local_address, master_timeout, provision_pool, - resolve_device_name, resolve_master_address, running_master, wait_for_master, @@ -63,7 +62,6 @@ "master_timeout", "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/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index f2de46a10feb..e94b0bc37350 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -50,7 +50,6 @@ "local_address", "master_timeout", "provision_pool", - "resolve_device_name", "resolve_master_address", "running_master", "wait_for_master", @@ -104,7 +103,9 @@ class PoolSpec: metadata_server: str = DEFAULT_METADATA_SERVER #: `rdma` or `tcp`. TCP is for bring-up only. protocol: str = "rdma" - #: RDMA device to transfer over, from `ibv_devinfo`. Empty auto-detects. + #: RDMA device to transfer over, from `ibv_devinfo`. One config serves a + #: whole pool, so naming a device here names it for every node; left empty + #: each client discovers its own, which is what differing HCAs need. device_name: str = "" #: Host memory each worker process contributes. Pool capacity is this times #: the number of processes that open a store handle. @@ -285,78 +286,6 @@ def _wait_until_accepting( time.sleep(0.5) -#: 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 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] = {} - 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 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 - 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. @@ -381,9 +310,7 @@ def wait_for_master(master_address: str, timeout: Optional[float] = None) -> Opt return elapsed -def _client_config( - pool: PoolSpec, master_address: str, device_name: Optional[str] = None -) -> Dict[str, Any]: +def _client_config(pool: PoolSpec, 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 @@ -394,7 +321,7 @@ def _client_config( "metadata_server": pool.metadata_server, "master_server_address": master_address, "protocol": pool.protocol, - "device_name": pool.device_name if device_name is None else device_name, + "device_name": pool.device_name, "global_segment_size": pool.global_segment_size, "local_buffer_size": pool.local_buffer_size, "role": "both", @@ -424,9 +351,7 @@ def write_client_config(pool: PoolSpec, master_address: str, run_dir: str) -> st """ os.makedirs(run_dir, exist_ok=True) config_path = os.path.join(run_dir, CLIENT_CONFIG_NAME) - config = _client_config( - pool, master_address, resolve_device_name(pool.protocol, pool.device_name) - ) + config = _client_config(pool, master_address) with open(config_path, "w") as handle: json.dump(config, handle, indent=2) logger.info(f"mooncake-store: wrote {config_path} ({json.dumps(config, sort_keys=True)})") diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py index ffe5a6b05722..9eb27056a037 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_master.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -601,51 +601,30 @@ def test_a_master_that_died_starting_is_reported_with_its_last_words(fake_master pytest.fail("provisioning should not have yielded") -# ---- choosing the fabric without naming it in a config ---- +# ---- which fabric a client transfers over ---- -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_the_pool_config_leaves_the_fabric_to_each_client(fake_master, tmp_path): + """One config describes a whole pool, so it must not name one node's HCAs. + Mooncake discovers devices locally on each client when `device_name` is + empty. Resolving a device while writing the config would instead freeze + whatever the provisioning host happens to have into a file every node + reads, which is wrong as soon as their HCAs differ. + """ + port = free_port() + fake_master.arm(listen_on=port) + pool = PoolSpec(launch_master=True, master_port=port, protocol="rdma") -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")) == "" + with provision_pool(pool, run_dir=str(tmp_path / "run")) as config_path: + assert json.loads(open(config_path).read())["device_name"] == "" -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)) +def test_a_named_device_is_what_the_workers_are_told(fake_master, tmp_path): + """Naming one is still how a deployment overrides that discovery.""" port = free_port() fake_master.arm(listen_on=port) - pool = PoolSpec(launch_master=True, master_port=port, protocol="rdma") + pool = PoolSpec(launch_master=True, master_port=port, protocol="rdma", device_name="mlx5_7") with provision_pool(pool, run_dir=str(tmp_path / "run")) as config_path: - assert json.loads(open(config_path).read())["device_name"] == "mlx5_0" + assert json.loads(open(config_path).read())["device_name"] == "mlx5_7"