[TRTLLM-14815][feat] Enable disaggregated serving for Kimi K3

examples/disaggregated/slurm/benchmark/run_benchmark.sh

@@ -5,6 +5,22 @@ set -e
 set -u
 trap 'echo "Error occurred at line $LINENO"; exit 1' ERR
 
+# Container runtimes (pyxis/enroot) reset image-defined variables like PATH
+# at container start, so values passed via srun --export are lost for them.
+# Allow the launcher config to prepend entries from inside the container.
+# The client_cmds srun passes these via --export with single-quoted values
+# (submit.py convert_envs_to_str), and srun keeps the quotes literal — strip them.
+TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND:-}"
+TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}"
+TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND:-}"
+TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}"
+if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then
+    export PATH="${TRTLLM_PATH_PREPEND}:${PATH}"
+fi
+if [ -n "${TRTLLM_PYTHONPATH_PREPEND:-}" ]; then
+    export PYTHONPATH="${TRTLLM_PYTHONPATH_PREPEND}${PYTHONPATH:+:${PYTHONPATH}}"
+fi
+
 # Add parameter validation
 if [ "$#" -lt 10 ]; then
     echo "Error: Missing required arguments, got $# arguments, args: $@"

examples/disaggregated/slurm/benchmark/start_server.sh

@@ -5,4 +5,26 @@ set -x
 
 config_file=$1
 
-trtllm-serve disaggregated -c ${config_file} -t 7200 -r 7200
+# Container runtimes (pyxis/enroot) reset image-defined variables like PATH
+# at container start, so values passed via srun --export are lost for them.
+# Allow the launcher config to prepend entries from inside the container.
+# srun --export keeps any quotes in the exported values literal; strip them.
+TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND:-}"
+TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}"
+TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND:-}"
+TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}"
+if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then
+    export PATH="${TRTLLM_PATH_PREPEND}:${PATH}"
+fi
+if [ -n "${TRTLLM_PYTHONPATH_PREPEND:-}" ]; then
+    export PYTHONPATH="${TRTLLM_PYTHONPATH_PREPEND}${PYTHONPATH:+:${PYTHONPATH}}"
+fi
+
+# In-place (.pth-style) TRT-LLM installs may lack the trtllm-serve console
+# script; fall back to the module entry point in that case.
+trtllm_serve_cmd="trtllm-serve"
+if ! command -v trtllm-serve >/dev/null 2>&1; then
+    trtllm_serve_cmd="python3 -m tensorrt_llm.commands.serve"
+fi
+
+${trtllm_serve_cmd} disaggregated -c ${config_file} -t 7200 -r 7200

examples/disaggregated/slurm/benchmark/start_worker.sh

@@ -38,8 +38,30 @@ else
     export CUDA_VISIBLE_DEVICES=${cuda_devices}
 fi
 
-# Clear UCX_TLS for specific clusters
-unset UCX_TLS
+# Container runtimes (pyxis/enroot) reset image-defined variables like PATH
+# at container start, so values passed via srun --export are lost for them.
+# Allow the launcher config to prepend entries from inside the container.
+# srun --export keeps any quotes in the exported values literal; strip them.
+TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND:-}"
+TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}"
+TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND:-}"
+TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}"
+if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then
+    export PATH="${TRTLLM_PATH_PREPEND}:${PATH}"
+fi
+if [ -n "${TRTLLM_PYTHONPATH_PREPEND:-}" ]; then
+    export PYTHONPATH="${TRTLLM_PYTHONPATH_PREPEND}${PYTHONPATH:+:${PYTHONPATH}}"
+fi
+
+# Clear UCX_TLS for specific clusters. Some clusters instead need an
+# explicit transport list (e.g. NVL72 nodes whose verbs transports cannot
+# initialize): set TRTLLM_WORKER_UCX_TLS in worker_env_var to re-pin
+# UCX_TLS here, after the container-provided value is cleared.
+if [ -n "${TRTLLM_WORKER_UCX_TLS:-}" ]; then
+    export UCX_TLS="${TRTLLM_WORKER_UCX_TLS}"
+else
+    unset UCX_TLS
+fi
 
 echo "SLURM_PROCID: ${SLURM_PROCID}, hostname: $(hostname), instance_id: ${instance_id}"
 echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}"
@@ -63,7 +85,14 @@ else
     nsys_prefix="nsys profile -o ${nsys_file} -f true -t cuda,nvtx,python-gil -c cudaProfilerApi --cuda-graph-trace node --capture-range-end=stop --gpu-metrics-devices=none"
 fi
 
+# In-place (.pth-style) TRT-LLM installs may lack the trtllm-serve console
+# script; fall back to the module entry point in that case.
+trtllm_serve_cmd="trtllm-serve"
+if ! command -v trtllm-serve >/dev/null 2>&1; then
+    trtllm_serve_cmd="python3 -m tensorrt_llm.commands.serve"
+fi
+
 ${nsys_prefix} trtllm-llmapi-launch ${numa_bind_cmd} \
-    trtllm-serve ${model_path} \
+    ${trtllm_serve_cmd} ${model_path} \
         --host $(hostname) --port ${port} \
         --config ${config_file}

examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml

@@ -0,0 +1,83 @@
+# Kimi K3 KDA-payload cache-transceiver micro-benchmark.
+# Run from the repository root:
+#   python3 examples/disaggregated/slurm/cache_transceiver_test/submit.py \
+#       -c examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml
+#
+# Synthetic KV geometry sized so ONE request's per-rank transfer equals the
+# exact per-request KDA state payload of Kimi K3: 454,459,392 bytes
+# (69 layers x (conv [3*96*128, 4] bf16 + delta [96,128,128] fp32)).
+# Per-rank bytes/token = 69 layers * kvFactor(2) * (24/4 kv heads) * 128 * 2B
+#                      = 211,968 B -> 2144 tokens == 454,459,392 B exactly.
+# request_lengths 512 / 1024 also match MLA-latent payloads at ~4k / ~8k ISL
+# (bf16) for the KDA-vs-MLA comparison.
+#
+# Fill in the <...> placeholders for your cluster before submitting.
+
+slurm:
+  partition: "<partition>"
+  account: "<account>"
+  job_time: "00:30:00"
+  job_name: "ctt_kda_payload"
+  extra_args: "--gpus-per-node=4"
+
+hardware:
+  gpus_per_node: 4
+
+environment:
+  container_image: "<container_image>"
+  # Mount the TRT-LLM checkout into the container.
+  container_mount: "<container_mount>"  # Format: path1:path1,path2:path2
+  # TRT-LLM comes from the checkout's venv (PATH below) + checkout code
+  # (PYTHONPATH below); skip pip install.
+  trtllm_repo: ""
+  # Output directory (results.json, logs/, csv/); point at any writable
+  # scratch location.
+  work_dir: "<full_path_to_work_dir>"
+  trtllm_wheel_path: ""
+  build_wheel: false
+  cuda_architectures: ""
+
+test_matrix:
+  combinations:
+    # NIXL/PYTHON dropped for now: V2 _exchange_rank_info mpi_allgather can
+    # hang under this harness (srun --mpi=pmix).
+    - {backend: "NIXL", runtime: "CPP"}
+    - {backend: "UCX", runtime: "CPP"}
+  cache_manager_versions: ["V1"]
+  # 512 tok = 108.5 MB/rank (~ MLA bf16 @4k ISL), 1024 = 217.1 MB (~ MLA @8k),
+  # 2144 = 454.46 MB = exact K3 per-request KDA payload.
+  request_lengths: [512, 1024, 2144]
+  num_requests_per_length: 8
+  warmup_requests: 2
+
+kv_cache:
+  num_layers: 69
+  num_kv_heads: 24       # /tp4 -> 6 heads per rank
+  head_dim: 128
+  tokens_per_block: 32
+  dtype: "HALF"
+  max_tokens_in_buffer: 4096
+
+parallel:
+  ctx_tp: 4
+  ctx_pp: 1
+  gen_tp: 4
+  gen_pp: 1
+
+ucx_env_sweep:
+  # Single sweep. PATH/PYTHONPATH are exported inside the container prelude so
+  # python3 resolves to the checkout's venv (system-site-packages=true,
+  # so container torch/mpi4py are visible) and the checkout on PYTHONPATH
+  # wins. Keep your cluster's SLURM bin dir on PATH so host-side srun works.
+  # The UCX_TLS pin (no verbs) is for clusters where verbs transports cannot
+  # initialize on the compute nodes; drop the pin where verbs works.
+  - name: "venv_no_verbs"
+    env:
+      UCX_TLS: "tcp,self,sm,cuda_copy,cuda_ipc"
+      PATH: "<trtllm_venv_bin_dir>:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:<slurm_bin_dir>"
+      PYTHONPATH: "<trtllm_repo>"
+
+run:
+  timeout_per_cell_s: 120
+  max_sweep_s: 900
+  capture_proto_info: true

examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py

@@ -287,6 +287,35 @@ class _TransferError(Exception):
     pass
 
 
+def _wait_ctx_complete(xcvr, rid, runtime):
+    """Block until this ctx request's send finishes (or errors).
+
+    * C++ transceiver: check_context_transfer_status(None) is a true block-all;
+      a single call suffices (see the comment at the call site).
+    * PYTHON (V2) transceiver: even under block_all, each TxSession wait is
+      bounded by kv_transfer_sender_future_timeout_ms (default 1000 ms). A
+      slower peer handshake makes the call log "TxSession ... timed out" and
+      return with the request still DISAGG_CONTEXT_TRANS_IN_PROGRESS -- in the
+      real executor that is benign (the session stays open and is re-polled
+      every iteration), but returning here would free and refill the KV blocks
+      mid-flight, so the receiver reads the NEXT request's pattern (verify FAIL
+      on every request except the last). So poll until
+      this rid lands in the completed/failed lists. Collectively safe: every
+      ctx rank loops on the same rid and the per-call consensus makes all
+      ranks observe completion on the same iteration. The per-cell
+      signal.alarm and the hang detector bound the loop.
+    """
+    if runtime != "PYTHON":
+        xcvr.check_context_transfer_status(None)
+        return
+    while True:
+        completed, failed = xcvr.check_context_transfer_status(None)
+        if rid in failed:
+            raise _TransferError(f"ctx transfer failed for rid={rid}")
+        if rid in completed:
+            return
+
+
 def _wait_gen_complete(xcvr, req, runtime):
     """Block until this gen request's receive finishes (or errors).
 
@@ -372,7 +401,9 @@ def run_one_request(
         # transfer is still in progress. NIXL/UCX cold-start connection setup can
         # exceed that, so the harness would free the request mid-transfer, leaving
         # the gen side hung and the ctx sender thread asserting on a freed session.
-        xcvr.check_context_transfer_status(None)
+        # The PYTHON runtime additionally needs a poll loop on top of block_all;
+        # see _wait_ctx_complete.
+        _wait_ctx_complete(xcvr, rid, runtime)
         state = req.state
         tensorrt_llm.logger.info(f"[ctx r{rank}] rid={rid}: transfer DONE (send), state={state}")
         free_sequence(kvm, req, kv_handle, use_v2)
@@ -648,6 +679,13 @@ def arm_watchdog(combination_idx, reqlen_idx, what):
             backend=backend,
             transceiver_runtime=(None if runtime == "CPP" else "PYTHON"),
             max_tokens_in_buffer=cfg["kv_cache"]["max_tokens_in_buffer"],
+            # PYTHON (V2) only; 0 keeps bounce off. With bounce on, the KV data
+            # rides a fabric-VMM staging buffer (CU_MEM_HANDLE_TYPE_FABRIC), which
+            # is what lets UCX pick cuda_ipc across NVL72 nodes -- direct
+            # pool-to-pool transfers from non-fabric allocations fall back to
+            # much slower host-staged tcp, so enable bounce for cross-node
+            # transfers inside an NVLink domain.
+            kv_cache_bounce_size_mb=int(cfg["kv_cache"].get("bounce_size_mb", 0)),
         )
 
         # Build the cache manager + transceiver ONCE per case (the manager is

examples/kimi_k3/README.md

@@ -193,4 +193,3 @@ default cache manager.
   TRTLLM-14904.
 - FP8 KV cache (`kv_cache_config.dtype: fp8`) is not yet supported.
 - Speculative decoding: suffix-automaton speculation is supported for aggregated serving (`speculative_config: {decoding_type: SA}` in the extra LLM API options). Combining speculation with disaggregated serving is not yet supported.
-- Disaggregated serving is not yet supported.