[WIP] Add Kimi Agentic with LMCache NIXL

.github/workflows/run-sweep.yml

@@ -556,7 +556,7 @@ jobs:
             precision: ${{ matrix.config.precision }}
             router: ${{ matrix.config.router && toJson(matrix.config.router) || '' }}
             kv-p2p-transfer: ${{ matrix.config['kv-p2p-transfer'] || '' }}
-            conc-list: '[${{ matrix.config.conc }}]'
+            conc-list: ${{ toJson(matrix.config.conc) }}
             spec-decoding: ${{ matrix.config.spec-decoding }}
             disagg: ${{ matrix.config.disagg }}
             prefill-hardware: ${{ matrix.config.prefill.hardware }}
@@ -577,7 +577,7 @@ jobs:
             decode-ep: ${{ matrix.config.decode.ep }}
             decode-dp-attn: ${{ matrix.config.decode.dp-attn }}
             decode-additional-settings: ${{ toJson(matrix.config.decode.additional-settings) }}
-            conc: ${{ matrix.config.conc }}
+            conc: ${{ matrix.config.conc[0] }}
             kv-offloading: ${{ matrix.config.kv-offloading }}
             kv-offload-backend: ${{ matrix.config['kv-offload-backend'].name }}
             kv-offload-backend-metadata: ${{ matrix.config['kv-offload-backend'] && toJson(matrix.config['kv-offload-backend']) || '' }}

benchmarks/benchmark_lib.sh

@@ -1346,6 +1346,18 @@ ensure_hf_cli() {
 }
 
 resolve_trace_source() {
+    if [[ -n "${AIPERF_TRACE_SOURCE_FLAG:-}" ]]; then
+        TRACE_SOURCE_FLAG="$AIPERF_TRACE_SOURCE_FLAG"
+        echo "Loading traces via AIPERF_TRACE_SOURCE_FLAG: $TRACE_SOURCE_FLAG"
+        return 0
+    fi
+
+    if [[ -n "${AIPERF_WEKA_TRACE_DIR:-}" ]]; then
+        TRACE_SOURCE_FLAG="--custom-dataset-type weka_trace --input-file $AIPERF_WEKA_TRACE_DIR"
+        echo "Loading traces via local Weka trace dir: $AIPERF_WEKA_TRACE_DIR"
+        return 0
+    fi
+
     # Per-recipe override: set WEKA_LOADER_OVERRIDE to one of the aiperf
     # public-dataset loader names allowed by the inferencex-agentx-mvp
     # scenario. Used by recipes whose servers have non-default context
@@ -1527,8 +1539,9 @@ build_replay_cmd() {
     # Default --num-dataset-entries is 100; the with-subagents Weka corpus
     # has 393. Cap at 393 so all unique traces are loaded (the loader treats
     # this as a ``min(cap, available)`` ceiling, not a target — see
-    # semianalysis_cc_traces_weka.py).
-    REPLAY_CMD+=" --num-dataset-entries 393"
+    # semianalysis_cc_traces_weka.py). Smoke tests can override this for a
+    # small local fixture without changing production recipes.
+    REPLAY_CMD+=" --num-dataset-entries ${AIPERF_NUM_DATASET_ENTRIES:-393}"
     # 1-second timeslices on the server-metrics scrape so the post-run
     # plotter has per-window time series (KV usage, cache hit rate,
     # throughput, etc.). Matches kv-cache-tester's poll_interval=1.0

benchmarks/multi_node/amd_utils/job.slurm

@@ -83,6 +83,9 @@ DECODE_MTP_SIZE=${DECODE_MTP_SIZE:-0}
 ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}"
 ROUTER_PORT="${ROUTER_PORT:-30000}"
 PROXY_PING_PORT="${PROXY_PING_PORT:-36367}"
+VLLM_ROUTER_POLICY="${VLLM_ROUTER_POLICY:-consistent_hash}"
+VLLM_ROUTER_PREFILL_POLICY="${VLLM_ROUTER_PREFILL_POLICY:-consistent_hash}"
+VLLM_ROUTER_DECODE_POLICY="${VLLM_ROUTER_DECODE_POLICY:-round_robin}"
 
 # =============================================================================
 # Model Path Resolution
@@ -113,6 +116,20 @@ if [[ "$ENGINE" == "vllm-disagg" ]]; then
     resolve_hf_cache_path() {
         local base_path=$1
         if [[ -d "${base_path}/snapshots" ]]; then
+            if [[ -f "${base_path}/refs/main" ]]; then
+                local ref_snapshot
+                ref_snapshot="$(cat "${base_path}/refs/main" 2>/dev/null | tr -d '[:space:]')"
+                if [[ -n "$ref_snapshot" && -d "${base_path}/snapshots/${ref_snapshot}" ]]; then
+                    echo "${base_path}/snapshots/${ref_snapshot}"
+                    return 0
+                fi
+            fi
+            local processor_snapshot
+            processor_snapshot=$(find "${base_path}/snapshots" -mindepth 2 -maxdepth 2 -name preprocessor_config.json -print -quit 2>/dev/null | xargs -r dirname)
+            if [[ -n "$processor_snapshot" ]]; then
+                echo "$processor_snapshot"
+                return 0
+            fi
             local snapshot=$(ls -1 "${base_path}/snapshots" 2>/dev/null | head -1)
             if [[ -n "$snapshot" ]]; then
                 echo "${base_path}/snapshots/${snapshot}"
@@ -124,20 +141,26 @@ if [[ "$ENGINE" == "vllm-disagg" ]]; then
     }
 
     MODEL_PATH=""
-    SEARCH_PATHS=(
-        "${MODEL_DIR}/${DISK_DIR_NAME}"
-        "${MODEL_DIR}/${MODEL_NAME}"
-        "/nfsdata/hf_hub_cache-0/${DISK_DIR_NAME}"
-        "/nfsdata/hf_hub_cache-0/${MODEL_NAME}"
-    )
-
-    for search_path in "${SEARCH_PATHS[@]}"; do
-        if [[ -d "$search_path" ]]; then
-            RESOLVED=$(resolve_hf_cache_path "$search_path")
-            MODEL_PATH="$RESOLVED"
-            echo "Found MODEL_PATH: $MODEL_PATH"
-            break
-        fi
+    MODEL_MOUNT_DIR=""
+    SEARCH_ROOTS=("$MODEL_DIR")
+    if [[ "$(hostname -s)" == mia1* ]]; then
+        SEARCH_ROOTS+=("/it-share/models" "/it-share/hf_cache" "/it-share/data")
+    fi
+    SEARCH_ROOTS+=("/nfsdata/hf_hub_cache-0" "/nfsdata/hf_hub_cache" "/nfsdata")
+
+    SEARCH_PATHS=()
+    for search_root in "${SEARCH_ROOTS[@]}"; do
+        SEARCH_PATHS+=("${search_root}/${DISK_DIR_NAME}" "${search_root}/${MODEL_NAME}")
+        for search_path in "${search_root}/${DISK_DIR_NAME}" "${search_root}/${MODEL_NAME}"; do
+            if [[ -d "$search_path" ]]; then
+                RESOLVED=$(resolve_hf_cache_path "$search_path")
+                MODEL_PATH="$RESOLVED"
+                MODEL_MOUNT_DIR="$search_root"
+                echo "Found MODEL_PATH: $MODEL_PATH"
+                echo "Found MODEL_MOUNT_DIR: $MODEL_MOUNT_DIR"
+                break 2
+            fi
+        done
     done
 
     if [[ -z "$MODEL_PATH" ]]; then
@@ -247,18 +270,21 @@ echo "NNODES:      ${NNODES}"
 echo "REPO DIR:    ${DI_REPO_DIR}"
 echo "USER:        ${USER_NAME}"
 
+
 # Reduce log spam
 export TQDM_MININTERVAL=20
 
 # Translate the host-resolved MODEL_PATH to the Docker mount namespace
-DOCKER_MODEL_PATH="${MODEL_PATH/#$MODEL_DIR//models}"
+MODEL_MOUNT_DIR="${MODEL_MOUNT_DIR:-$MODEL_DIR}"
+DOCKER_MODEL_PATH="${MODEL_PATH/#$MODEL_MOUNT_DIR//models}"
 
 export DI_REPO_DIR=$DI_REPO_DIR
 export WS_PATH=$WS_PATH
 export NNODES=$NNODES
 export NODE0_ADDR=$NODE0_ADDR
 export MODEL_PATH=$MODEL_PATH
 export MODEL_DIR=$MODEL_DIR
+export MODEL_MOUNT_DIR=$MODEL_MOUNT_DIR
 export xP=$xP
 export yD=$yD
 export MODEL_NAME=$MODEL_NAME
@@ -288,13 +314,39 @@ export RESULT_FILENAME="${RESULT_FILENAME:-}"
 export SPEC_DECODING="${SPEC_DECODING:-}"
 export IS_MULTINODE="${IS_MULTINODE:-false}"
 
+# Agentic / custom vLLM-disagg connector knobs (threaded from submit.sh)
+export IS_AGENTIC="${IS_AGENTIC:-0}"
+export DURATION="${DURATION:-1800}"
+export MODEL="${MODEL:-}"
+export ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}"
+export ROUTER_PORT="${ROUTER_PORT:-30000}"
+export VLLM_ROUTER_POLICY="${VLLM_ROUTER_POLICY:-consistent_hash}"
+export VLLM_ROUTER_PREFILL_POLICY="${VLLM_ROUTER_PREFILL_POLICY:-consistent_hash}"
+export VLLM_ROUTER_DECODE_POLICY="${VLLM_ROUTER_DECODE_POLICY:-round_robin}"
+export ENABLE_PREFIX_CACHING="${ENABLE_PREFIX_CACHING:-}"
+export MAX_MODEL_LEN="${MAX_MODEL_LEN:-}"
+export WEKA_LOADER_OVERRIDE="${WEKA_LOADER_OVERRIDE:-}"
+export AIPERF_TRACE_SOURCE_FLAG="${AIPERF_TRACE_SOURCE_FLAG:-}"
+export AIPERF_WEKA_TRACE_DIR="${AIPERF_WEKA_TRACE_DIR:-}"
+export AIPERF_NUM_DATASET_ENTRIES="${AIPERF_NUM_DATASET_ENTRIES:-}"
+export VLLM_BIND_IP="${VLLM_BIND_IP:-}"
+export SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-}"
+export PREFILL_KV_CONNECTOR="${PREFILL_KV_CONNECTOR:-moriio}"
+export DECODE_KV_CONNECTOR="${DECODE_KV_CONNECTOR:-moriio}"
+export MORI_RDMA_DEVICES="${MORI_RDMA_DEVICES:-}"
+export MORI_IO_TC="${MORI_IO_TC:-}"
+export MORI_RDMA_SL="${MORI_RDMA_SL:-}"
+export MORI_IO_SL="${MORI_IO_SL:-}"
+export UCX_IB_GID_INDEX="${UCX_IB_GID_INDEX:-}"
+export DOCKER_PULL_POLICY="${DOCKER_PULL_POLICY:-missing}"
+
 SANITIZED_USER=$(echo "$USER_NAME" | tr -c 'a-zA-Z0-9_.-' '_')
 export DOCKER_CONT_NAME="container_${ENGINE}_${SANITIZED_USER}_${MODEL_NAME}_${SLURM_JOB_ID}"
 
 # vLLM external router container.
 # NOTE: vllm/vllm-router only retains ~16 recent nightlies on Docker Hub; older
 # dated tags are garbage-collected (manifest unknown)
-VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly-20260629-e667ebb}"
+VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}"
 ROUTER_CONT_NAME="router_vllm_${SANITIZED_USER}_${SLURM_JOB_ID}"
 export RUN_FILE_FULL="$WS_PATH/${RUN_FILE}"
 
@@ -371,6 +423,37 @@ DOCKER_ENV_COMMON=(
     -e DECODE_ENABLE_DP=\$DECODE_ENABLE_DP
     -e DECODE_MTP_SIZE=\$DECODE_MTP_SIZE
     -e IS_MULTINODE=\$IS_MULTINODE
+    -e MODEL=\$MODEL
+    -e ROUTER_TYPE=\$ROUTER_TYPE
+    -e ROUTER_PORT=\$ROUTER_PORT
+    -e VLLM_ROUTER_POLICY=\$VLLM_ROUTER_POLICY
+    -e VLLM_ROUTER_PREFILL_POLICY=\$VLLM_ROUTER_PREFILL_POLICY
+    -e VLLM_ROUTER_DECODE_POLICY=\$VLLM_ROUTER_DECODE_POLICY
+    -e ENABLE_PREFIX_CACHING=\$ENABLE_PREFIX_CACHING
+    -e MAX_MODEL_LEN=\$MAX_MODEL_LEN
+    -e MAX_NUM_SEQS=\$MAX_NUM_SEQS
+    -e WEKA_LOADER_OVERRIDE=\$WEKA_LOADER_OVERRIDE
+    -e AIPERF_TRACE_SOURCE_FLAG=\$AIPERF_TRACE_SOURCE_FLAG
+    -e AIPERF_WEKA_TRACE_DIR=\$AIPERF_WEKA_TRACE_DIR
+    -e AIPERF_NUM_DATASET_ENTRIES=\$AIPERF_NUM_DATASET_ENTRIES
+    -e VLLM_BIND_IP=\$VLLM_BIND_IP
+    -e SERVED_MODEL_NAME=\$SERVED_MODEL_NAME
+    -e PREFILL_KV_CONNECTOR=\$PREFILL_KV_CONNECTOR
+    -e DECODE_KV_CONNECTOR=\$DECODE_KV_CONNECTOR
+    -e MORI_RDMA_DEVICES=\$MORI_RDMA_DEVICES
+    -e MORI_IO_TC=\$MORI_IO_TC
+    -e MORI_RDMA_SL=\$MORI_RDMA_SL
+    -e MORI_IO_SL=\$MORI_IO_SL
+    -e UCX_IB_GID_INDEX=\$UCX_IB_GID_INDEX
+    -e LMCACHE_HOST=\${LMCACHE_HOST:-}
+    -e LMCACHE_PORT=\${LMCACHE_PORT:-}
+    -e LMCACHE_HTTP_PORT=\${LMCACHE_HTTP_PORT:-}
+    -e LMCACHE_L1_SIZE_GB=\${LMCACHE_L1_SIZE_GB:-}
+    -e LMCACHE_L1_INIT_SIZE_GB=\${LMCACHE_L1_INIT_SIZE_GB:-}
+    -e LMCACHE_L1_READ_TTL_SECONDS=\${LMCACHE_L1_READ_TTL_SECONDS:-}
+    -e LMCACHE_CHUNK_SIZE=\${LMCACHE_CHUNK_SIZE:-}
+    -e LMCACHE_MAX_WORKERS=\${LMCACHE_MAX_WORKERS:-}
+    -e LMCACHE_MP_MQ_TIMEOUT=\${LMCACHE_MP_MQ_TIMEOUT:-}
 )
 
 # Engine-specific env vars
@@ -387,6 +470,13 @@ if [[ "$ENGINE" == "vllm-disagg" ]]; then
         -e UCX_LOG_LEVEL=warn
         -e HSA_ENABLE_SDMA=1
         -e PROXY_STREAM_IDLE_TIMEOUT=\${PROXY_STREAM_IDLE_TIMEOUT:-300}
+        -e VLLM_MORIIO_CONNECTOR_READ_MODE=\${VLLM_MORIIO_CONNECTOR_READ_MODE:-1}
+        -e IBDEVICES=\${IBDEVICES:-}
+        -e MORI_RDMA_TC=\${MORI_RDMA_TC:-}
+        -e MORI_IO_TC=\${MORI_IO_TC:-}
+        -e MORI_RDMA_SL=\${MORI_RDMA_SL:-}
+        -e MORI_IO_SL=\${MORI_IO_SL:-}
+        -e NCCL_IB_DISABLE=\${NCCL_IB_DISABLE:-}
         -e PYTHONPYCACHEPREFIX=/tmp/pycache
     )
 elif [[ "$ENGINE" == "atom-disagg" ]]; then
@@ -428,6 +518,36 @@ echo \"Rank \$SLURM_PROCID on \$(hostname)\"
 eval \"\$DOCKER_CMD_DETECT\"
 echo \"[docker-detect] rank \$SLURM_PROCID: DOCKER_CMD=\$DOCKER_CMD\"
 
+pull_docker_image_if_needed() {
+    local image=\"\$1\"
+    case \"${DOCKER_PULL_POLICY}\" in
+        always)
+            echo \"[docker-pull] pulling \$image on \$(hostname) (policy=always)\"
+            \$DOCKER_CMD pull \"\$image\"
+            ;;
+        missing|if-not-present)
+            if ! \$DOCKER_CMD image inspect \"\$image\" >/dev/null 2>&1; then
+                echo \"[docker-pull] pulling \$image on \$(hostname) (policy=missing)\"
+                \$DOCKER_CMD pull \"\$image\"
+            else
+                echo \"[docker-pull] using existing \$image on \$(hostname) (policy=missing)\"
+            fi
+            ;;
+        never|false|0)
+            echo \"[docker-pull] skipping pull for \$image on \$(hostname) (policy=${DOCKER_PULL_POLICY})\"
+            ;;
+        *)
+            echo \"[docker-pull] invalid DOCKER_PULL_POLICY=${DOCKER_PULL_POLICY}; expected always, missing, or never\" >&2
+            exit 1
+            ;;
+    esac
+}
+
+pull_docker_image_if_needed \"$DOCKER_IMAGE_NAME\"
+if [[ \"$ENGINE\" == \"vllm-disagg\" && \"$ROUTER_TYPE\" == \"vllm-router\" && \"\$SLURM_PROCID\" == \"0\" ]]; then
+    pull_docker_image_if_needed \"$VLLM_ROUTER_IMAGE\"
+fi
+
 # Enable out-of-tree RDMA library mounts for atom-disagg (mooncake requires host RDMA stack)
 RDMA_MOUNTS=()
 if [[ "$ENGINE" == "atom-disagg" ]]; then
@@ -507,27 +627,47 @@ fi
 fi  # end: if ENGINE == atom-disagg
 
 # Pre-clean (idempotent)
-\$DOCKER_CMD ps -aq --filter \"$CONT_FILTER\" | xargs -r \$DOCKER_CMD rm -f || true
-\$DOCKER_CMD ps -aq | xargs -r \$DOCKER_CMD stop || true
+mkdir -p \"${BENCHMARK_LOGS_DIR}/runtime_logs\"
+\$DOCKER_CMD rm -f \"$DOCKER_CONT_NAME\" 2>/dev/null || true
+if [[ \"$ENGINE\" == \"vllm-disagg\" && \"$ROUTER_TYPE\" == \"vllm-router\" && \"\$SLURM_PROCID\" == \"0\" ]]; then
+    \$DOCKER_CMD rm -f \"$ROUTER_CONT_NAME\" 2>/dev/null || true
+fi
 
 # Start vLLM external router container on node 0
 if [[ \"$ENGINE\" == \"vllm-disagg\" && \"$ROUTER_TYPE\" == \"vllm-router\" && \"\$SLURM_PROCID\" == \"0\" ]]; then
     \$DOCKER_CMD rm -f \"$ROUTER_CONT_NAME\" 2>/dev/null || true
+    ROUTER_KV_CONNECTOR=\"\${VLLM_ROUTER_KV_CONNECTOR:-moriio}\"
+    if [[ \"\${PREFILL_KV_CONNECTOR:-}\" == \"lmcache-nixl\" ]]; then
+        ROUTER_KV_CONNECTOR=\"\${VLLM_ROUTER_KV_CONNECTOR:-nixl}\"
+    fi
+    ROUTER_ENDPOINT_ARGS=\"--vllm-discovery-address 0.0.0.0:${PROXY_PING_PORT}\"
+    if [[ \"\${PREFILL_KV_CONNECTOR:-}\" == \"lmcache-nixl\" ]]; then
+        ROUTER_ENDPOINT_ARGS=\"\"
+        IFS=',' read -r -a _router_ips <<< \"${IPADDRS}\"
+        for ((i=0; i<${xP} && i<\${#_router_ips[@]}; i++)); do
+            ROUTER_ENDPOINT_ARGS+=\" --prefill http://\${_router_ips[\$i]}:${SERVER_PORT:-2584}\"
+        done
+        for ((i=${xP}; i<\${#_router_ips[@]}; i++)); do
+            ROUTER_ENDPOINT_ARGS+=\" --decode http://\${_router_ips[\$i]}:${SERVER_PORT:-2584}\"
+        done
+    fi
+    echo \"[ROUTER_CMD] image=$VLLM_ROUTER_IMAGE vllm-router --vllm-pd-disaggregation --kv-connector \${ROUTER_KV_CONNECTOR} \${ROUTER_ENDPOINT_ARGS} --worker-startup-timeout-secs \${VLLM_ROUTER_WORKER_STARTUP_TIMEOUT_SECS:-1800} --port ${ROUTER_PORT} --host 0.0.0.0 --policy ${VLLM_ROUTER_POLICY} --prefill-policy ${VLLM_ROUTER_PREFILL_POLICY} --decode-policy ${VLLM_ROUTER_DECODE_POLICY} --log-level info\"
     \$DOCKER_CMD run -d \
         --name \"$ROUTER_CONT_NAME\" \
         --network host \
         --ulimit nofile=1048576:1048576 \
-        -v /tmp:/run_logs \
+        -v ${BENCHMARK_LOGS_DIR}/runtime_logs:/run_logs \
         \"$VLLM_ROUTER_IMAGE\" \
         bash -lc \"mkdir -p /run_logs/slurm_job-${SLURM_JOB_ID} && exec vllm-router \
             --vllm-pd-disaggregation \
-            --kv-connector moriio \
-            --vllm-discovery-address 0.0.0.0:${PROXY_PING_PORT} \
+            --kv-connector \${ROUTER_KV_CONNECTOR} \
+            \${ROUTER_ENDPOINT_ARGS} \
+            --worker-startup-timeout-secs \${VLLM_ROUTER_WORKER_STARTUP_TIMEOUT_SECS:-1800} \
             --port ${ROUTER_PORT} \
             --host 0.0.0.0 \
-            --policy consistent_hash \
-            --prefill-policy consistent_hash \
-            --decode-policy consistent_hash \
+            --policy ${VLLM_ROUTER_POLICY} \
+            --prefill-policy ${VLLM_ROUTER_PREFILL_POLICY} \
+            --decode-policy ${VLLM_ROUTER_DECODE_POLICY} \
             --log-level info 2>&1 | tee /run_logs/slurm_job-${SLURM_JOB_ID}/vllm_router_\$(hostname).log \"
 fi
 
@@ -566,10 +706,10 @@ fi
     --privileged \
     -v /sys:/sys \
     $(command -v nicctl >/dev/null 2>&1 && echo "-v $(which nicctl):/usr/sbin/nicctl") \
-    -v ${MODEL_DIR}:/models \
+    -v ${MODEL_MOUNT_DIR}:/models \
     -v \$HOME/.ssh:/root/.ssh \
     --shm-size 128G \
-    -v /tmp:/run_logs \
+    -v ${BENCHMARK_LOGS_DIR}/runtime_logs:/run_logs \
     -v ${BENCHMARK_LOGS_DIR}:/benchmark_logs \
     -v ${DI_REPO_DIR}:${DOCKER_MOUNT_PATH} \
     ${EXTRA_DOCKER_MOUNTS:-} \