[None][perf] Integrate Mooncake store for MinimaxM3

docker/common/install_mooncake.sh

@@ -50,3 +50,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

examples/disaggregated/slurm/benchmark/disaggr_torch.slurm

@@ -138,6 +138,26 @@ else
     echo "TensorRT-LLM environment variables saved to ${full_logdir}/env_vars.json"
 fi
 
+# The Mooncake store bindings, when a worker config asks for the connector.
+# Images built from this repo bake them in (docker/common/install_mooncake.sh),
+# so this only confirms they are there: a job that discovers the gap later
+# fails once per rank, deep in engine startup.
+mooncake_enabled=false
+if grep -qs "mooncake-store" "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"; then
+    mooncake_enabled=true
+    # Both halves of the wheel are checked because they fail at different
+    # times: the connector needs mooncake.store in every context rank, and
+    # 'trtllm-serve mooncake_master' needs the binary on PATH.
+    if ! srun --container-name=${container_name} \
+        --container-mounts=${container_mount} --no-container-mount-home \
+        --mpi=pmix --overlap -N 1 -n 1 \
+        bash -c 'python3 -c "import mooncake.store" && command -v mooncake_master' \
+        &> ${full_logdir}/2_check_mooncake.log; then
+        cleanup_on_failure "A worker config requests the mooncake-store connector, but this image has neither the mooncake.store bindings nor the mooncake_master binary. Build the image from this repo so that docker/common/install_mooncake.sh runs. Check ${full_logdir}/2_check_mooncake.log for details"
+    fi
+    echo "Mooncake store bindings found in the image"
+fi
+
 # Get node lists and replace the placeholder with the actual node names
 echo "SLURM_NODELIST: ${SLURM_NODELIST}"
 all_nodes=($(scontrol show hostname $SLURM_NODELIST | sort))
@@ -154,6 +174,20 @@ client_cmds_base_file=${full_logdir}/client_cmds_base.sh
 client_cmds_file=${full_logdir}/client_cmds.sh
 replace_placeholder "${client_cmds_base_file}" "${all_nodes_str}" "${client_cmds_file}"
 
+# The pool is described in the worker configs and provisioned by trtllm-serve
+# during its own bringup: the context server starts the master, renders the
+# client config and publishes the master's address, and the generation servers
+# read that address to lend the pool their memory. Nothing here starts, waits
+# for or configures any of it.
+#
+# The one value a config written before submission cannot know is this job's log
+# directory, which the master's address is published into, so a __LOG_DIR__
+# placeholder in either worker config is filled in here.
+if [ "${mooncake_enabled}" = "true" ]; then
+    sed -i "s|__LOG_DIR__|${full_logdir}|g" \
+        "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"
+fi
+
 # Per-worker hostfile / gpu_map files for srun --distribution=arbitrary.
 # submit.py emits *_base.txt with <nodeN_placeholder>; rewrite them here.
 for base_file in "${full_logdir}"/hostfile_*_base.txt "${full_logdir}"/gpu_map_*_base.txt; do
@@ -176,6 +210,20 @@ cat ${start_server_cmds_file} | while read cmd; do
 done
 echo "Server is ready!"
 
+# A connector that failed to open its store handle does not stop the worker from
+# serving, it just never hits, so surface the startup lines here rather than
+# leaving them to be discovered after the benchmark. The registration line
+# carries the bytes/page figure the pool sizing depends on.
+if [ "${mooncake_enabled}" = "true" ]; then
+    echo "Mooncake store startup lines from the context workers:"
+    if ! grep -h "mooncake-store" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null; then
+        echo "  WARNING: no mooncake-store lines found. The connector may not have" \
+             "loaded; check ${full_logdir}/3_output_CTX_*.log and" \
+             "${full_logdir}/mooncake_master.log. The benchmark will still run," \
+             "but without the store."
+    fi
+fi
+
 # Start client commands
 echo "Starting client commands from ${client_cmds_file}..."
 while read -r cmd <&3; do
@@ -186,6 +234,78 @@ while read -r cmd <&3; do
     fi
 done 3< "${client_cmds_file}"
 
+# Collect the store's traffic into one file. The per-event lines live at DEBUG
+# in the worker logs (module _torch), so this is only populated when a config
+# asks for that verbosity. The counts are what show the store did work rather
+# than merely started.
+if [ "${mooncake_enabled}" = "true" ]; then
+    mooncake_summary="${full_logdir}/9_mooncake_summary.log"
+    {
+        # The address file is retracted when the master stops, so fall back to
+        # the context server's log, which still names the master the run used.
+        echo "master: $(tr -d '[:space:]' < "${full_logdir}/master.addr" 2>/dev/null \
+            || grep -hoE "master at [0-9.]+:[0-9]+" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null \
+               | head -n 1 | awk '{print $3}' || echo unknown)"
+        echo
+        echo "== startup =="
+        grep -h "mooncake-store.*\(ready\|registered layout\)" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null || echo "(none)"
+        echo
+        echo "== event counts =="
+        for pattern in "matched" "loaded" "failed to load" "failed to save" \
+                       "lookup failed" "could not reserve connector prefix"; do
+            count=$(grep -h "mooncake-store.*${pattern}" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null | wc -l || true)
+            echo "${pattern}: ${count}"
+        done
+        echo
+        # Where the blocks physically went. The master names a segment for
+        # every allocation and a segment is one client process's donated
+        # memory, so grouping by segment host shows how much of the pool lives
+        # on a decode node rather than on the prefill node that computed it.
+        # Without a donor this section shows only the prefill node.
+        echo "== block placement by segment host =="
+        echo "(lending hosts: $(grep -hoE "GiB of [0-9.]+ is now part of the pool" \
+            "${full_logdir}"/3_output_GEN_*.log 2>/dev/null \
+            | awk '{print $3}' | sort -u | paste -sd, - || echo none))"
+        grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" \
+            "${full_logdir}/mooncake_master.log" 2>/dev/null \
+            | awk '{
+                    sub(/size=/, "", $2); sub(/segment=/, "", $3);
+                    split($3, parts, ":"); host = parts[1]; port = parts[2];
+                    allocs[host]++; bytes[host] += $2;
+                    if (!((host, port) in seen)) {
+                        seen[host, port] = 1;
+                        segs[host] = segs[host] " " port;
+                    }
+                    total_allocs++; total_bytes += $2;
+                }
+                END {
+                    if (total_allocs == 0) { print "(no allocations)"; exit }
+                    for (h in allocs) {
+                        printf "%-16s pages=%-7d %8.2f GiB  %5.1f%% of pool contents  segments:%s\n",
+                            h, allocs[h], bytes[h] / 1073741824,
+                            100 * bytes[h] / total_bytes, segs[h];
+                    }
+                    printf "%-16s pages=%-7d %8.2f GiB\n", "TOTAL", total_allocs, total_bytes / 1073741824;
+                }' || echo "(could not parse master log)"
+        echo
+        # Lending memory happens inside the generation servers, so their own
+        # logs are where a segment that never mounted shows up.
+        echo "== lent segments =="
+        grep -h "mooncake-store: .*\(lending memory\|part of the pool\|withdrew\)" \
+            "${full_logdir}"/3_output_GEN_*.log 2>/dev/null || echo "(none)"
+        echo
+        echo "== pool bringup (context server) =="
+        grep -h "mooncake-store:" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null \
+            | head -n 30 || echo "(none)"
+        echo
+        # The master's own glog output, not anything TensorRT-LLM writes.
+        echo "== master =="
+        tail -n 50 "${full_logdir}/mooncake_master.log" 2>/dev/null || echo "(no master log)"
+    } > "${mooncake_summary}" 2>&1
+    echo "Mooncake store summary written to ${mooncake_summary}"
+    cat "${mooncake_summary}"
+fi
+
 echo "Job completed successfully, total runtime: $SECONDS seconds"
 
 # try to kill the server and workers

examples/disaggregated/slurm/benchmark/start_worker.sh

@@ -54,6 +54,47 @@ fi
 
 echo "config_file: ${config_file}"
 
+# The mooncake-store pool is described in the worker config and provisioned by
+# trtllm-serve during bringup. Anchoring its run directory here keeps the
+# master's log, the rendered client config and the published address in the
+# job's log directory rather than in a temporary directory that shutdown
+# removes, and it is how the ranks srun started, which never inherited the
+# leader's environment, find that client config. An inherited
+# MOONCAKE_CONFIG_PATH still wins, so an externally managed pool stays reachable.
+export TRTLLM_MOONCAKE_RUN_DIR="${log_dir}"
+
+# The generation servers wait for a master the context server starts. Both are
+# launched together and the master comes up before its model loads, but the wait
+# spans container start on another node, so it is given far more than the 60s
+# default. Too short fails the job; too long costs nothing when the master is
+# already there.
+export TRTLLM_MOONCAKE_MASTER_TIMEOUT="${TRTLLM_MOONCAKE_MASTER_TIMEOUT:-900}"
+
+# MiniMax-M3's MSA sparse attention JIT-compiles its FMHA kernels on first use,
+# from inside the attention forward pass. One TP rank runs ninja while the
+# others block on a file lock, so an uncached variant stalls the whole executor
+# loop for ~8s, or ~70s when an iteration needs several. The cache defaults to
+# ~/.cache, which is thrown away because the container is started with
+# --no-container-mount-home, so every job would pay the compiles again during
+# serving. Anchoring it next to this script puts it on the mounted filesystem
+# at a path identical across jobs, so only the first run compiles.
+if [ -z "${MINFER_FMHA_CACHE_DIR:-}" ]; then
+    export MINFER_FMHA_CACHE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/.cache/minfer/fmha_sm100"
+    mkdir -p "${MINFER_FMHA_CACHE_DIR}"
+    echo "MINFER_FMHA_CACHE_DIR: ${MINFER_FMHA_CACHE_DIR}"
+fi
+
+# Per-transfer KV timings (size, queue/transfer latency, throughput) as CSV next
+# to the worker logs. These separate slow prefill from a slow prefill-to-decode
+# handoff, which the aggregate benchmark numbers cannot. An explicit setting
+# wins, and KV_TRANSFER_PERF_LOG=false turns it off.
+if [ "${KV_TRANSFER_PERF_LOG:-true}" = "true" ] \
+    && [ -z "${TLLM_KV_TRANSFER_PERF_LOG_FILE:-}" ]; then
+    export TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1
+    export TLLM_KV_TRANSFER_PERF_LOG_FILE="${log_dir}/kv_transfer_perf"
+    echo "TLLM_KV_TRANSFER_PERF_LOG_FILE: ${TLLM_KV_TRANSFER_PERF_LOG_FILE}"
+fi
+
 nsys_prefix=""
 if [ "${enable_nsys}" != "true" ]; then
     echo "nsys is not enabled, start normal flow"

examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml

@@ -0,0 +1,49 @@
+# Extra LLM API options for trtllm-serve with the Mooncake store KV connector.
+#
+# Offloads KV pages to a Mooncake distributed store, a shared CPU memory pool
+# addressed by content, so a prefix computed by one engine can be replayed by
+# another. This is a different component from the Mooncake transfer engine used
+# by the C++ cache transceiver for disaggregated prefill/decode handoff; the two
+# compose rather than conflict.
+#
+# Prerequisites:
+#   - Mooncake Python bindings: pip install mooncake-transfer-engine
+#     (present in the release container; the C++ source build does not
+#     provide MooncakeDistributedStore)
+#   - A running Mooncake master, and a metadata server unless using
+#     P2PHANDSHAKE. See https://kvcache-ai.github.io/Mooncake/
+#   - MOONCAKE_CONFIG_PATH pointing at a Mooncake JSON config, for example:
+#       {
+#         "metadata_server": "http://127.0.0.1:8080/metadata",
+#         "master_server_address": "127.0.0.1:50051",
+#         "protocol": "rdma",
+#         "device_name": "mlx5_0",
+#         "global_segment_size": "32GiB",
+#         "local_buffer_size": "1GiB"
+#       }
+#
+# Optional environment overrides:
+#   TRTLLM_MOONCAKE_STORE_ROLE       producer | consumer | both   (default both)
+#   TRTLLM_MOONCAKE_STORE_PREFIX     key prefix, to isolate deployments sharing
+#                                    one pool                    (default trtllm)
+#   TRTLLM_MOONCAKE_STORE_MODEL_KEY  identity keys are namespaced by
+#                                    (default: model directory basename)
+#
+# Example:
+#   export MOONCAKE_CONFIG_PATH=/path/to/mooncake.json
+#   trtllm-serve <hf_model> --backend pytorch --host 0.0.0.0 --port 8000 \
+#     --extra_llm_api_options /path/to/this/file
+
+kv_cache_config:
+  # The connector describes its pools through register_kv_cache_layout, which
+  # only KVCacheManagerV2 implements.
+  use_kv_cache_manager_v2: true
+  # Local reuse still runs first; the store serves whatever the device missed.
+  enable_block_reuse: true
+  # GPU-only tiers are required: a page evicted to host or disk has its GPU slot
+  # reassigned, which would invalidate the addresses registered with the store.
+  host_cache_size: 0
+  disk_cache_size: 0
+
+kv_connector_config:
+  connector: mooncake-store

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/_util.py

@@ -2714,6 +2714,8 @@ def create_py_executor_instance(
             cross_kv_cache_manager=cross_kv_cache_manager,
             no_schedule_until_state=no_schedule_until_state,
             enable_prefix_aware_scheduling=enable_prefix_aware_scheduling,
+            max_input_len=max_seq_len
+            if max_seq_len is not None else 0x7fffffff,
         )
     elif (scheduler_config is not None
           and scheduler_config.use_python_scheduler):