[TRTLLM-11316][feat] Add kvbm disagg tests and support kvbm disagg slurm task

examples/disaggregated/slurm/benchmark/config.yaml

@@ -47,6 +47,23 @@ environment:
   worker_env_var: "TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes NCCL_GRAPH_MIXING_SUPPORT=0"
   server_env_var: "TRTLLM_SERVER_DISABLE_GC=1"
 
+# Service Discovery Configuration
+# Workers register with the disagg server via service discovery instead of
+# static server lists.  Two backends are supported:
+#   "http"  – the disagg server itself is the registry (no extra services)
+#   "etcd"  – an etcd server is started on the disagg node for registration
+service_discovery:
+  backend: "http"   # Options: "http", "etcd"
+  etcd_port: 2379   # Port for etcd (only used when backend is "etcd")
+  cluster_name: ""  # Optional namespace to isolate clusters in etcd
+
+# KVBM Configuration (Dynamo KV Block Manager)
+# When enabled, injects kv_connector_config into worker configs
+# and installs the kvbm pip package.
+kvbm:
+  enabled: false   # Set to true to enable kvbm
+  package: "kvbm"  # pip package name or path (e.g. "kvbm" or "/path/to/kvbm.whl")
+
 # Profiling Configuration
 profiling:
   nsys_on: false  # Set to true to enable profiling

examples/disaggregated/slurm/benchmark/disaggr_torch.slurm

@@ -17,6 +17,8 @@ while [[ $# -gt 0 ]]; do
         --build-wheel) build_wheel="$2"; shift 2 ;;
         --cuda-architectures) cuda_architectures="$2"; shift 2 ;;
         --trtllm-wheel-path) trtllm_wheel_path="$2"; shift 2 ;;
+        --kvbm-enabled) kvbm_enabled="$2"; shift 2 ;;
+        --kvbm-package) kvbm_package="$2"; shift 2 ;;
         *)
             echo "Unknown argument: $1"
             exit 1
@@ -39,6 +41,8 @@ echo "  container_image: ${container_image}"
 echo "  build_wheel: ${build_wheel}"
 echo "  cuda_architectures: ${cuda_architectures}"
 echo "  trtllm_wheel_path: ${trtllm_wheel_path}"
+echo "  kvbm_enabled: ${kvbm_enabled}"
+echo "  kvbm_package: ${kvbm_package}"
 
 # Set TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1 for gen_only_no_context mode
 if [ "${benchmark_mode}" = "gen_only_no_context" ]; then
@@ -138,6 +142,19 @@ else
     echo "TensorRT-LLM environment variables saved to ${full_logdir}/env_vars.json"
 fi
 
+# Install kvbm if enabled
+if [ "${kvbm_enabled}" = "true" ]; then
+    echo "Installing kvbm package: ${kvbm_package}..."
+    if ! srun --container-name=${container_name} \
+        --container-mounts=${container_mount} --no-container-mount-home \
+        --mpi=pmix --overlap -N $SLURM_NNODES --ntasks-per-node=1 \
+        bash -c "pip install ${kvbm_package}" \
+        &> ${full_logdir}/2.1_kvbm_install.log; then
+        cleanup_on_failure "kvbm installation failed. Check ${full_logdir}/2.1_kvbm_install.log for details"
+    fi
+    echo "kvbm installation completed successfully"
+fi
+
 # Get node lists and replace the placeholder with the actual node names
 echo "SLURM_NODELIST: ${SLURM_NODELIST}"
 all_nodes=($(scontrol show hostname $SLURM_NODELIST | sort))

examples/disaggregated/slurm/benchmark/disaggr_torch_dwdp.slurm

@@ -17,6 +17,8 @@ while [[ $# -gt 0 ]]; do
         --build-wheel) build_wheel="$2"; shift 2 ;;
         --cuda-architectures) cuda_architectures="$2"; shift 2 ;;
         --trtllm-wheel-path) trtllm_wheel_path="$2"; shift 2 ;;
+        --kvbm-enabled) kvbm_enabled="$2"; shift 2 ;;
+        --kvbm-package) kvbm_package="$2"; shift 2 ;;
         *)
             echo "Unknown argument: $1"
             exit 1

examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh

@@ -1,14 +1,13 @@
 #!/bin/bash
 
-# aiperf-based benchmark script for disaggregated serving
+# AIPerf-based benchmark script for disaggregated serving.
 # Args: model_name dataset_file multi_round num_gen_servers concurrency_list streaming log_path hostname port ucx_warmup_requests
 
-set -e
-set -u
+set -euo pipefail
 trap 'echo "Error occurred at line $LINENO"; exit 1' ERR
 
 if [ "$#" -lt 10 ]; then
-    echo "Error: Missing required arguments, got $# arguments, args: $@"
+    echo "Error: Missing required arguments, got $# arguments, args: $*"
     echo "Usage: $0 model_name dataset_file multi_round num_gen_servers concurrency_list streaming log_path hostname port ucx_warmup_requests"
     exit 1
 fi
@@ -30,71 +29,218 @@ if [[ ${SLURM_PROCID} != "0" ]]; then
     exit 0
 fi
 
-# Always install/upgrade aiperf to ensure we have the version with trust_remote_code fix
-# (container may have an older version with parallel_decode.py that lacks trust_remote_code)
-echo "Installing aiperf..."
-pip install --force-reinstall --no-deps 'aiperf @ git+https://github.com/ai-dynamo/aiperf.git@ac3d91652e5e024bfb4ac38d48603423aad666bc'
+strip_outer_quotes() {
+    local value="$1"
+    if [[ "${value}" == \'*\' ]] && [[ "${value}" == *\' ]]; then
+        value="${value:1:${#value}-2}"
+    fi
+    if [[ "${value}" == \"*\" ]] && [[ "${value}" == *\" ]]; then
+        value="${value:1:${#value}-2}"
+    fi
+    printf "%s" "${value}"
+}
+
+for env_name in HF_HOME XDG_CACHE_HOME PIP_CACHE_DIR AIPERF_SOURCE_DIR AIPERF_CUSTOM_DATASET_TYPE AIPERF_BENCHMARK_DURATION_SECONDS AIPERF_MAX_CONTEXT_LENGTH AIPERF_TOKENIZER AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT AIPERF_DATASET_CONFIGURATION_TIMEOUT AIPERF_TRACE_SPLIT_DIR AIPERF_ARTIFACT_BASE_DIR TMPDIR HF_MODULES_CACHE; do
+    if [ -n "${!env_name:-}" ]; then
+        printf -v "${env_name}" "%s" "$(strip_outer_quotes "${!env_name}")"
+        export "${env_name}"
+    fi
+done
+
+export HF_HOME=${HF_HOME:-/lustre/fsw/coreai_comparch_trtllm/lizhiz/.cache/huggingface}
+export XDG_CACHE_HOME=${XDG_CACHE_HOME:-${log_path}/cache/xdg}
+export PIP_CACHE_DIR=${PIP_CACHE_DIR:-${log_path}/cache/pip}
+export TMPDIR=${TMPDIR:-${log_path}/cache/tmp}
+export HF_MODULES_CACHE=${HF_MODULES_CACHE:-${log_path}/cache/hf_modules}
+export AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT=${AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT:-900}
+export AIPERF_DATASET_CONFIGURATION_TIMEOUT=${AIPERF_DATASET_CONFIGURATION_TIMEOUT:-900}
+export HF_HUB_TRUST_REMOTE_CODE=1
+trace_split_root="${AIPERF_TRACE_SPLIT_DIR:-${log_path}/weka_trace_json}"
+artifact_base_dir="${AIPERF_ARTIFACT_BASE_DIR:-${log_path}}"
+mkdir -p "${HF_HOME}" "${XDG_CACHE_HOME}" "${PIP_CACHE_DIR}" "${TMPDIR}" "${HF_MODULES_CACHE}" "${trace_split_root}" "${artifact_base_dir}"
+
+install_aiperf() {
+    local source_dir="${AIPERF_SOURCE_DIR:-/lustre/fsw/coreai_comparch_trtllm/lizhiz/offload-sa-k25/aiperf}"
+
+    if [ -d "${source_dir}/src/aiperf" ]; then
+        echo "Installing AIPerf from local source: ${source_dir}"
+        python -m pip install --force-reinstall --no-deps -e "${source_dir}"
+    else
+        echo "Local AIPerf source not found at ${source_dir}; installing Cam's AgentX branch"
+        python -m pip install --force-reinstall --no-deps \
+            'aiperf @ git+https://github.com/cquil11/aiperf.git@cjq/agentx-v0.3-subagents'
+    fi
+}
+
+prepare_dataset() {
+    dataset_type="${AIPERF_CUSTOM_DATASET_TYPE:-}"
+    if [ -z "${dataset_type}" ]; then
+        case "${dataset_file}" in
+            *.jsonl) dataset_type="weka_trace" ;;
+            *) dataset_type="mooncake_trace" ;;
+        esac
+    fi
+
+    dataset_input="${dataset_file}"
+    if [ "${dataset_type}" = "weka_trace" ] && [ -f "${dataset_file}" ] && [[ "${dataset_file}" == *.jsonl ]]; then
+        dataset_input="${trace_split_root}"
+        if [ ! -f "${dataset_input}/.complete" ]; then
+            echo "Splitting Weka JSONL ${dataset_file} into ${dataset_input}"
+            mkdir -p "${dataset_input}"
+            python - "${dataset_file}" "${dataset_input}" <<'PY'
+import json
+import pathlib
+import re
+import sys
+
+src = pathlib.Path(sys.argv[1])
+out = pathlib.Path(sys.argv[2])
+complete = out / ".complete"
+
+if complete.exists():
+    print(f"Weka JSON split already complete: {out}")
+    raise SystemExit(0)
+
+existing = list(out.glob("*.json"))
+if existing:
+    raise SystemExit(
+        f"{out} contains JSON files but .complete is absent; use a fresh log directory"
+    )
+
+count = 0
+with src.open("r", encoding="utf-8") as fh:
+    for idx, line in enumerate(fh):
+        line = line.strip()
+        if not line:
+            continue
+        obj = json.loads(line)
+        trace_id = str(obj.get("id") or obj.get("trace_id") or f"trace_{idx:06d}")
+        safe = re.sub(r"[^A-Za-z0-9._-]+", "_", trace_id).strip("._")
+        safe = (safe or f"trace_{idx:06d}")[:180]
+        with (out / f"{idx:06d}_{safe}.json").open("w", encoding="utf-8") as out_f:
+            json.dump(obj, out_f, separators=(",", ":"))
+        count += 1
+
+complete.write_text(f"{count}\n", encoding="utf-8")
+print(f"Wrote {count} Weka trace JSON files to {out}")
+PY
+        fi
+    fi
+}
+
+install_aiperf
+python -m pip install "crick~=0.0.8"
+prepare_dataset
 
 # warmup requests for ucx connections
 if [ "${ucx_warmup_requests}" -gt 0 ]; then
     echo "warming up ucx connections with small requests... ${ucx_warmup_requests}"
     python -m tensorrt_llm.serve.scripts.benchmark_serving \
-        --model ${model_name} \
+        --model "${model_name}" \
         --dataset-name random \
         --random-ids \
         --random-input-len 100 \
         --random-output-len 10 \
-        --num-prompts ${ucx_warmup_requests} \
-        --host ${hostname} \
-        --port ${port} \
+        --num-prompts "${ucx_warmup_requests}" \
+        --host "${hostname}" \
+        --port "${port}" \
         --ignore-eos \
         --trust-remote-code \
         --non-streaming
     echo "UCX warmup done"
 fi
 
-# Trust remote code globally for custom tokenizers in parallel workers
-export HF_HUB_TRUST_REMOTE_CODE=1
+streaming_args=()
+case "${streaming}" in
+    true|True|TRUE|1|yes|Yes|YES) streaming_args=(--streaming) ;;
+    *) streaming_args=() ;;
+esac
 
-echo "Hostname: ${hostname}, Port: ${port}"
-echo "Starting aiperf benchmark..."
+run_weka_trace_benchmark() {
+    local concurrency=$1
+    local artifact_dir="${artifact_base_dir}/concurrency_${concurrency}"
+    local benchmark_duration="${AIPERF_BENCHMARK_DURATION_SECONDS:-1800}"
+    local tokenizer="${AIPERF_TOKENIZER:-${model_name}}"
+    local max_context_length="${AIPERF_MAX_CONTEXT_LENGTH:-167000}"
 
-concurrency_list=$(echo "${concurrency_list}" | tr ',' ' ')
-for concurrency in ${concurrency_list}; do
-    concurrency=$((concurrency))
-    request_count=$((concurrency * multi_round))
-    # benchmark_duration: 20min per round
-    benchmark_duration=$((multi_round * 1200))
-    echo "Benchmarking with concurrency ${concurrency} ... ${request_count} requests, duration ${benchmark_duration}s"
-    mkdir -p ${log_path}/concurrency_${concurrency}
+    echo "Benchmarking Weka trace with concurrency ${concurrency}, duration ${benchmark_duration}s"
+    mkdir -p "${artifact_dir}"
+
+    aiperf profile \
+        --scenario inferencex-agentx-mvp \
+        --model "${model_name}" \
+        --tokenizer "${tokenizer}" \
+        --tokenizer-trust-remote-code \
+        --max-context-length "${max_context_length}" \
+        --endpoint-type chat \
+        --url "http://${hostname}:${port}" \
+        "${streaming_args[@]}" \
+        --ui simple \
+        --input-file "${dataset_input}" \
+        --artifact-dir "${artifact_dir}" \
+        --concurrency "${concurrency}" \
+        --concurrency-ramp-duration "${AIPERF_CONCURRENCY_RAMP_DURATION_SECONDS:-60}" \
+        --custom-dataset-type weka_trace \
+        --benchmark-duration "${benchmark_duration}" \
+        --benchmark-grace-period "${AIPERF_BENCHMARK_GRACE_PERIOD_SECONDS:-60}" \
+        --workers-max "${AIPERF_WORKERS_MAX:-200}" \
+        --request-timeout-seconds "${AIPERF_REQUEST_TIMEOUT_SECONDS:-1200}" \
+        --profile-export-level records \
+        --extra-inputs ignore_eos:true \
+        --record-processors "${AIPERF_RECORD_PROCESSORS:-8}" \
+        --use-server-token-count
+}
+
+run_trace_benchmark() {
+    local concurrency=$1
+    local request_count=$((concurrency * multi_round))
+    local benchmark_duration=$((multi_round * 1200))
+    local artifact_dir="${artifact_base_dir}/concurrency_${concurrency}"
+
+    echo "Benchmarking ${dataset_type} with concurrency ${concurrency} ... ${request_count} requests, duration ${benchmark_duration}s"
+    mkdir -p "${artifact_dir}"
 
     aiperf profile \
-        -m ${model_name} \
-        --tokenizer ${model_name} \
+        --model "${model_name}" \
+        --tokenizer "${model_name}" \
         --tokenizer-trust-remote-code \
-        --url http://${hostname}:${port} \
-        --streaming \
+        --url "http://${hostname}:${port}" \
+        "${streaming_args[@]}" \
         --ui simple \
-        --input-file ${dataset_file} \
-        --artifact-dir ${log_path}/concurrency_${concurrency} \
-        --concurrency ${concurrency} \
+        --input-file "${dataset_input}" \
+        --artifact-dir "${artifact_dir}" \
+        --concurrency "${concurrency}" \
         --concurrency-ramp-duration 60 \
-        --custom-dataset-type mooncake_trace \
-        --benchmark-duration ${benchmark_duration} \
+        --custom-dataset-type "${dataset_type}" \
+        --benchmark-duration "${benchmark_duration}" \
         --benchmark-grace-period 60 \
         --workers-max 200 \
         --request-timeout-seconds 1200 \
         --profile-export-level records \
         --extra-inputs ignore_eos:true \
-        --request-count ${request_count} \
+        --request-count "${request_count}" \
         --record-processors 8
+}
+
+echo "Hostname: ${hostname}, Port: ${port}"
+echo "AIPerf trace split dir: ${trace_split_root}"
+echo "AIPerf artifact base dir: ${artifact_base_dir}"
+echo "Starting aiperf benchmark with dataset_type=${dataset_type}, dataset_input=${dataset_input}"
 
+concurrency_list=$(echo "${concurrency_list}" | tr ',' ' ')
+for concurrency in ${concurrency_list}; do
+    concurrency=$((concurrency))
+    if [ "${dataset_type}" = "weka_trace" ]; then
+        run_weka_trace_benchmark "${concurrency}"
+    else
+        run_trace_benchmark "${concurrency}"
+    fi
     echo "Benchmark with concurrency ${concurrency} done"
 done
 
 # Fetch perf metrics from disagg server
 echo "Fetching perf metrics from http://${hostname}:${port}/perf_metrics ..."
-curl -s "http://${hostname}:${port}/perf_metrics" > ${log_path}/perf_metrics.json 2>&1 || true
+curl -s "http://${hostname}:${port}/perf_metrics" > "${log_path}/perf_metrics.json" 2>&1 || true
 if [ -s "${log_path}/perf_metrics.json" ]; then
     echo "Perf metrics saved to ${log_path}/perf_metrics.json"
 else

examples/disaggregated/slurm/benchmark/start_worker.sh

@@ -11,31 +11,31 @@ numa_bind=${5}
 log_dir=${6}
 enable_nsys=${7}
 config_file=${8}
-# CUDA_VISIBLE_DEVICES selection:
-#   - Default packing (no gpu_map file): each node is dedicated to one
-#     worker, so SLURM_LOCALID maps directly to the physical GPU id.
-#   - Compact packing (gpu_map file emitted by submit.py): two workers may
-#     share a node and would both see LOCALID=0, so look up the per-worker
-#     gpu_map "<rank> <host> <local_gpu_id>" by SLURM_PROCID. srun
-#     --distribution=arbitrary assigns PROCID in hostfile order, so it
-#     indexes directly into the map.
-gpu_map_file="${log_dir}/gpu_map_${role}_${instance_id}.txt"
-if [ -f "${gpu_map_file}" ]; then
-    gpu_id=$(awk -v p="${SLURM_PROCID}" '$1==p {print $3; exit}' "${gpu_map_file}")
-    if [ -z "${gpu_id}" ]; then
-        echo "ERROR: no GPU mapping for SLURM_PROCID=${SLURM_PROCID} in ${gpu_map_file}" >&2
-        exit 1
-    fi
-    export CUDA_VISIBLE_DEVICES=${gpu_id}
+server_role=${9}
+disagg_cluster_uri=${10}
+
+# Do not set CUDA_VISIBLE_DEVICES here. Let Slurm/enroot expose the devices
+# assigned to this step; TRT-LLM relies on that visibility for peer checks.
+
+# Only clear UCX_TLS if the benchmark config didn't explicitly set one.
+if [ -z "${UCX_TLS:-}" ]; then
+    unset UCX_TLS
 else
-    export CUDA_VISIBLE_DEVICES=${SLURM_LOCALID}
+    echo "Using UCX_TLS: ${UCX_TLS}"
 fi
 
-# Clear UCX_TLS for specific clusters
-unset UCX_TLS
+# Resolve KVBM leader ZMQ hostname to IPv4 so ZMQ can bind (leader) and
+# connect (workers) using the same address across nodes.
+if [ -n "${DYN_KVBM_LEADER_ZMQ_HOST:-}" ]; then
+    resolved_ip=$(getent ahostsv4 "${DYN_KVBM_LEADER_ZMQ_HOST}" | awk '{print $1}' | head -1)
+    if [ -n "${resolved_ip}" ]; then
+        export DYN_KVBM_LEADER_ZMQ_HOST="${resolved_ip}"
+        echo "Resolved DYN_KVBM_LEADER_ZMQ_HOST to ${resolved_ip}"
+    fi
+fi
 
 echo "SLURM_PROCID: ${SLURM_PROCID}, hostname: $(hostname), instance_id: ${instance_id}"
-echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}"
+echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-<unset>}"
 
 if [ "${numa_bind}" = "true" ]; then
     numa_bind_cmd="numactl -m 0,1"
@@ -59,4 +59,6 @@ fi
 ${nsys_prefix} trtllm-llmapi-launch ${numa_bind_cmd} \
     trtllm-serve ${model_path} \
         --host $(hostname) --port ${port} \
-        --config ${config_file}
+        --config ${config_file} \
+        --server_role ${server_role} \
+        --disagg_cluster_uri ${disagg_cluster_uri}