kimik3-fp4-b300-vllm-agentic: day-zero Kimi-K3 B300 agentic recipe (#2371)

Author: Oseltamivir

benchmarks/single_node/agentic/kimik3_fp4_b300_vllm.sh

@@ -0,0 +1,170 @@
+#!/usr/bin/env bash
+set -eo pipefail
+set -x
+
+# Agentic trace replay benchmark for Kimi-K3 (MXFP4) on B300 using vLLM.
+#
+# Day-zero single-node recipe. Serve flags follow the Kimi-K3 production recipe
+# already exercised by the DSpark AL collector
+# (benchmarks/single_node/speedbench/kimik3_fp4_b300_vllm.sh) and the upstream
+# recipes.vllm.ai/moonshotai/Kimi-K3 guidance, adapted to the agentic scenario.
+#
+# Required env vars:
+#   MODEL, TP, CONC, KV_OFFLOADING, TOTAL_CPU_DRAM_GB, RESULT_DIR, DURATION
+#
+# TP8 is the only single-node layout. The MXFP4 checkpoint is ~1.5 TB on disk;
+# at TP4 that is ~375 GB of weights per GPU against B300's 288 GB of HBM, so
+# only the full 8-GPU shard (~188 GB/GPU) fits. Do not add TP4/TP2 arms.
+#
+# Weights are pre-staged node-local: `Kimi-K3` is in the b300-nv launcher's
+# STAGED_MODELS, so MODEL_PATH resolves to the read-only /scratch/models/Kimi-K3
+# mount and no download happens on the runner. The download block below only
+# covers stand-alone runs.
+#
+# KV_OFFLOADING: `none` (GPU-resident) or `dram` with
+# KV_OFFLOAD_BACKEND=vllm-simple (SimpleCPUOffloadConnector).
+#
+# Note on the DRAM arm: Kimi-K3 is a KDA/MLA hybrid — its linear-attention (KDA)
+# layers carry a recurrent state rather than paged KV blocks, and
+# SimpleCPUOffloadConnector offloads uniform paged blocks with no
+# hybrid-geometry handling. The arm is expected to offload only the MLA layers'
+# paged blocks; watch the bring-up sweep for KV-geometry errors at server init.
+
+source "$(dirname "$0")/../../benchmark_lib.sh"
+
+check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION
+
+# The 2.8T MXFP4 checkpoint only fits across all 8 B300s (see header).
+if [ "$TP" -ne 8 ]; then
+    echo "Error: Kimi-K3 on B300 requires TP=8 (a ~1.5 TB MXFP4 checkpoint does not fit at TP<8), got TP='$TP'" >&2
+    exit 1
+fi
+
+if [[ -n "${EP_SIZE:-}" && "${EP_SIZE}" -gt 1 ]]; then
+    echo "Error: this recipe ships the pure-TP8 profile; EP_SIZE='$EP_SIZE' is not wired yet" >&2
+    exit 1
+fi
+
+if [[ -n "${SLURM_JOB_ID:-}" ]]; then
+    echo "JOB $SLURM_JOB_ID running on ${SLURMD_NODENAME:-unknown}"
+fi
+
+# `hf download` creates the target dir if missing and is itself idempotent.
+# When MODEL_PATH is unset (stand-alone runs), fall back to the HF_HUB_CACHE.
+# Either way, MODEL_PATH is what the server is launched with.
+if [[ -n "${MODEL_PATH:-}" ]]; then
+    if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then
+        hf download "$MODEL" --local-dir "$MODEL_PATH"
+    fi
+else
+    hf download "$MODEL"
+    export MODEL_PATH="$MODEL"
+fi
+nvidia-smi
+
+# ---- Resolve traces and install deps ----------------------------------------
+resolve_trace_source
+install_agentic_deps
+
+# ---- Kimi-K3 production serving environment ---------------------------------
+export NCCL_DMABUF_ENABLE=0
+export VLLM_ALLREDUCE_USE_FLASHINFER=1
+export VLLM_USE_RUST_FRONTEND=1
+# Loading ~1.5 TB of MXFP4 shards off the staged mount takes well past the
+# default readiness window even with fastsafetensors.
+export VLLM_ENGINE_READY_TIMEOUT_S=3600
+export PYTHONNOUSERSITE=1
+# AIPerf pins one pooled keep-alive connection per agentic session and reuses it
+# across turns, while the Rust frontend's default VLLM_HTTP_TIMEOUT_KEEP_ALIVE is
+# 5s. An inter-turn idle gap longer than that lets the client reuse a socket at
+# the moment the server closes it -> aiohttp ServerDisconnectedError -> AIPerf
+# treats it as a terminal warmup failure and aborts the whole job. This killed
+# the dram c4 arm ~15 min into run 30324907690 with a perfectly healthy server
+# (it kept serving after the client gave up). Outlast the client pool so the
+# race cannot occur. Same fix as glm5.2_fp4_b300_sglang.sh's
+# SGLANG_TIMEOUT_KEEP_ALIVE=900.
+export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=900
+# Agentic warmup dispatches large prompts at once; allow up to 15 minutes of TCP
+# progress before AIPerf declares a connection dead.
+export AIPERF_HTTP_TCP_USER_TIMEOUT=900000
+
+# ---- Server config ----------------------------------------------------------
+SERVER_LOG="$RESULT_DIR/server.log"
+mkdir -p "$RESULT_DIR"
+
+# ---- KV offloading ----------------------------------------------------------
+# The generated TOTAL_CPU_DRAM_GB budget is the aggregate host-DRAM pool for the
+# node; SimpleCPUOffloadConnector is sized per rank. At dram-utilization 0.63 on
+# cluster:b300-nv this resolves to ~220 GiB per rank across the 8 TP ranks.
+OFFLOAD_ARGS=()
+case "${KV_OFFLOAD_BACKEND:-}" in
+    "")
+        require_agentic_kv_offload_none
+        ;;
+    vllm-simple)
+        require_agentic_kv_offload_backend vllm-simple
+        CPU_BYTES_PER_RANK=$(( TOTAL_CPU_DRAM_GB * 1000 * 1000 * 1000 / TP ))
+        # Identical prefixes must hash to identical block keys run-to-run.
+        export PYTHONHASHSEED=42
+        # lazy_offload must be a JSON boolean, not a quoted string: the
+        # connector does bool(extra_config.get("lazy_offload", False)), and
+        # bool("false") is True in Python — a quoted "false" would silently
+        # turn lazy offload ON. Eager offload keeps block-hash behaviour
+        # aligned with the other B300 vllm-simple arms.
+        OFFLOAD_ARGS=(
+            --kv-transfer-config
+            "{\"kv_connector\":\"SimpleCPUOffloadConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"cpu_bytes_to_use_per_rank\":${CPU_BYTES_PER_RANK},\"lazy_offload\":false}}"
+        )
+        ;;
+    *)
+        echo "Error: unsupported KV_OFFLOAD_BACKEND='$KV_OFFLOAD_BACKEND' (expected empty or vllm-simple)" >&2
+        exit 1
+        ;;
+esac
+
+# Agentic fan-out: keep the scheduler headroom convention shared by the other
+# agentic recipes. Capture decode graphs only up to that batch size — a 93-layer
+# 2.8T model makes capturing vLLM's full 2048-wide ladder prohibitively slow.
+MAX_NUM_SEQS=$((2 * CONC))
+
+echo "Starting vllm server..."
+
+{ set +x; } 2>/dev/null
+VLLM_CMD=(
+    vllm serve "$MODEL_PATH" --served-model-name "$MODEL"
+    --host 0.0.0.0
+    --port "$PORT"
+    --tensor-parallel-size "$TP"
+    --gpu-memory-utilization 0.90
+    --max-num-seqs "$MAX_NUM_SEQS"
+    # Agentic replays run at the model's native context limit.
+    --max-model-len 1048576
+    --trust-remote-code
+    --load-format fastsafetensors
+    --moe-backend auto
+    --enable-prefix-caching
+    --kv-cache-dtype fp8
+    --reasoning-parser kimi_k3
+    --tool-call-parser kimi_k3
+    --enable-auto-tool-choice
+    # FP8 KV cache requires the prefill query quantization flag; MLA prefill
+    # runs on FlashInfer per the production recipe.
+    --attention-config '{"mla_prefill_backend":"FLASHINFER","use_prefill_query_quantization":true}'
+    --max-cudagraph-capture-size "$MAX_NUM_SEQS"
+    --disable-uvicorn-access-log
+    "${OFFLOAD_ARGS[@]}"
+)
+printf '%q ' "${VLLM_CMD[@]}" | tee "$RESULT_DIR/vllm_command.txt"
+printf '\n' | tee -a "$RESULT_DIR/vllm_command.txt"
+"${VLLM_CMD[@]}" > "$SERVER_LOG" 2>&1 &
+SERVER_PID=$!
+echo "Server PID: $SERVER_PID"
+
+wait_for_server_ready --port "$PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID"
+
+if [ "${EVAL_ONLY}" = "true" ]; then
+    run_eval --port "$PORT"
+else
+    build_replay_cmd "$RESULT_DIR"
+    run_agentic_replay_and_write_outputs "$RESULT_DIR"
+fi

configs/nvidia-master.yaml

@@ -1578,6 +1578,37 @@ kimik2.5-fp4-b300-vllm-agentic-mtp:
       - { tp: 4, ep: 1, spec-decoding: mtp, dcp-size: 4, kv-offloading: none, conc-list: [32, 64, 80, 96, 112, 128] }
       - { tp: 4, ep: 1, spec-decoding: mtp, dcp-size: 4, kv-offloading: dram, kv-offload-backend: { name: native }, conc-list: [64, 80, 96, 112, 128, 144, 160] }
 
+kimik3-fp4-b300-vllm-agentic:
+  # Day-zero Kimi-K3 recipe. `vllm/vllm-openai:kimi-k3` is the pre-release
+  # Kimi-K3 build (vLLM 0.1.dev19262+gb6bbf29dd, pushed 2026-07-27); K3 ships as
+  # the out-of-tree `vllm.models.kimi_k3` plugin package (DSpark MLA + KDA
+  # kernels), so a generic vllm-openai release cannot serve this checkpoint.
+  image: vllm/vllm-openai:kimi-k3
+  model: moonshotai/Kimi-K3
+  model-prefix: kimik3
+  runner: cluster:b300-nv
+  precision: fp4
+  framework: vllm
+  multinode: false
+  scenarios:
+    # Agentic-coding only: no fixed-seq-len (1k1k / 8k1k) arms for this recipe.
+    agentic-coding:
+      # 0.63 resolves to ~220 GiB of host DRAM per rank across the 8 TP ranks on
+      # cluster:b300-nv (total-cpu-dram-gb 1889), which is what the
+      # SimpleCPUOffloadConnector pool is sized to in the script.
+    - dram-utilization: 0.63
+      search-space:
+      # TP8 is the only single-node layout that fits: the MXFP4 checkpoint is
+      # ~1.5 TB (~188 GB/GPU across 8 B300s), and TP4 would need ~375 GB/GPU
+      # against 288 GB of HBM. The conc ceiling is set by the ~70 GB/GPU left
+      # for MLA KV after weights at gpu-memory-utilization 0.90.
+      # Both arms share one ladder so GPU-resident and DRAM-offload are directly
+      # comparable at equal concurrency (same shape as the kimik2.5 B300 sister).
+      # TP8 GPU-resident
+      - { tp: 8, ep: 1, kv-offloading: none, conc-list: [1, 2, 4, 8, 16, 24] }
+      # TP8 SimpleCPUOffload (host DRAM)
+      - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [1, 2, 4, 8, 16, 24] }
+
 dsr1-fp8-b200-trt:
   image: nvcr.io#nvidia/tensorrt-llm/release:1.3.0rc14
   model: deepseek-ai/DeepSeek-R1-0528

perf-changelog.yaml

@@ -5145,3 +5145,17 @@
     - "Add a one-node 4P1D TP1-prefill/TEP4-decode concurrency-4096 point on the refreshed image while preserving the legacy 4P2D DEP2-prefill/TEP4-decode point unchanged"
   pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2310
 
+
+- config-keys:
+    - kimik3-fp4-b300-vllm-agentic
+  description:
+    - "Add day-zero single-node agentic-coding recipe for Kimi-K3 (MXFP4, 2.8T MoE, 896 experts, KDA/MLA hybrid, 1M native context) on B300 with vLLM. Agentic-coding scenario only -- no fixed-seq-len (1k1k / 8k1k) arms."
+    - "Image vllm/vllm-openai:kimi-k3 (vLLM 0.1.dev19262+gb6bbf29dd, pushed 2026-07-27): K3 ships as the out-of-tree vllm.models.kimi_k3 plugin package (DSpark MLA + KDA kernels, registering KimiK3ForConditionalGeneration / KimiLinearForCausalLM / KimiK3MTPModel), so a generic vllm-openai release cannot serve this checkpoint."
+    - "Weights are pre-staged node-local at /scratch/models/Kimi-K3 (96/96 shards, 1.5 TB, verified on all 18 b300 nodes); Kimi-K3 is already in the b300-nv launcher STAGED_MODELS list, so MODEL_PATH resolves to the read-only staged mount and the runner performs no download."
+    - "Search space is two TP8 arms at conc 1-24: GPU-resident (kv-offloading none) and host-DRAM offload (kv-offloading dram, SimpleCPUOffloadConnector via kv-offload-backend vllm-simple). Both share one ladder so the arms are directly comparable at equal concurrency. TP8 is the only single-node layout that fits: the MXFP4 checkpoint is ~1.5 TB (~188 GB/GPU across 8 B300s), whereas TP4 would need ~375 GB/GPU against 288 GB of HBM."
+    - "DRAM arm: dram-utilization 0.63 resolves to total-cpu-dram-gb 1889 on cluster:b300-nv, which the script divides across the 8 TP ranks into cpu_bytes_to_use_per_rank (~220 GiB/rank), with eager offload (lazy_offload false). Note K3 is a KDA/MLA hybrid -- the KDA layers hold a recurrent state rather than paged KV blocks, so only the MLA layers' paged blocks are expected to offload; the bring-up sweep is the check for KV-geometry errors at server init."
+    - "Serve flags follow the K3 production recipe already exercised by the DSpark AL collector: --load-format fastsafetensors, --moe-backend auto, --enable-prefix-caching, --kv-cache-dtype fp8 with --attention-config mla_prefill_backend=FLASHINFER + use_prefill_query_quantization, kimi_k3 reasoning/tool-call parsers, NCCL_DMABUF_ENABLE=0, VLLM_ALLREDUCE_USE_FLASHINFER=1, VLLM_USE_RUST_FRONTEND=1."
+    - "Sets VLLM_HTTP_TIMEOUT_KEEP_ALIVE=900 and AIPERF_HTTP_TCP_USER_TIMEOUT=900000. AIPerf pins one pooled keep-alive connection per agentic session and reuses it across turns, while the Rust frontend defaults to a 5s keep-alive; an inter-turn idle gap past that lets the client reuse a socket exactly as the server closes it, producing an aiohttp ServerDisconnectedError that AIPerf escalates to a terminal warmup failure. This aborted the dram c4 arm ~15 min into run 30324907690 against a healthy server. Same mitigation as glm5.2_fp4_b300_sglang.sh's SGLANG_TIMEOUT_KEEP_ALIVE=900."
+    - "Bring-up validated in run 30326393603: all 12 configs green, zero ServerDisconnectedError after the keep-alive fix. GPU KV resolves to 42.23 GiB / 3,249,215 tokens, i.e. ~3.1 max-length requests, against MAX_NUM_SEQS = 2*CONC."
+    - "Measured behaviour: below conc 8 the GPU-resident and DRAM arms are within run-to-run noise (1-5%). At conc 16 and 24 the GPU-resident arm thrashes -- prefix cache hit rate 2.7%, TTFT p50 86s and 191s, 49.6 and 54.9 output tok/s -- because the working set exceeds GPU KV and prefixes are recomputed. The DRAM arm holds TTFT p50 0.85s and 6.2s for 245.0 and 260.6 output tok/s (4-5x), with the CPU tier serving a 62% external prefix cache hit rate at conc 24. The high-conc GPU-resident points are retained deliberately as the honest baseline that makes the offload gain legible."
+  pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2371