[None][feat] Mooncake store part 1: pool, CLI, and V2 scheduler preemption

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

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

tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py

@@ -0,0 +1,71 @@
+# 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.
+`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,
+    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",
+    "MooncakeStoreConnectorScheduler",
+    "MooncakeStoreConnectorWorker",
+    "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",
+]