[WIP][AMD][AgentX] Kimi-K3 vLLM 1P1D DSpark disagg on MI355X

.github/workflows/benchmark-multinode-tmpl.yml

@@ -209,6 +209,25 @@ jobs:
             done
           fi
 
+      - name: Clear root-owned benchmark leftovers (pre-checkout)
+        run: |
+          set -uo pipefail
+          case "${GITHUB_WORKSPACE:-}" in
+            */actions-runner/_work/*) ;;
+            *) echo "Refusing to clean unexpected workspace '${GITHUB_WORKSPACE:-}'"; exit 0 ;;
+          esac
+          for d in benchmark_logs benchmark_artifacts; do
+            target="$GITHUB_WORKSPACE/$d"
+            [ -e "$target" ] || continue
+            if rm -rf "$target" 2>/dev/null; then
+              echo "removed $target"
+            elif sudo rm -rf "$target"; then
+              echo "removed $target (sudo)"
+            else
+              echo "WARNING: could not remove $target" >&2
+            fi
+          done
+
       - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
         with:
           token: ${{ secrets.REPO_PAT }}
@@ -307,16 +326,19 @@ jobs:
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
         with:
           name: multinode_server_logs_${{ env.RESULT_FILENAME }}
-          # multinode launchers package server logs into this tarball.
-          path: multinode_server_logs.tar.gz
+          path: |
+            multinode_server_logs.tar.gz
+            benchmark_artifacts/slurm_job-*.out
+            benchmark_artifacts/slurm_job-*.err
           if-no-files-found: ignore
 
       - name: Upload agentic aggregated result
-        if: ${{ !inputs.eval-only && inputs.scenario-type == 'agentic-coding' }}
+        if: ${{ always() && !inputs.eval-only && inputs.scenario-type == 'agentic-coding' }}
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
         with:
           name: bmk_agentic_${{ env.RESULT_FILENAME }}
           path: ${{ env.RESULT_FILENAME }}_conc*.json
+          if-no-files-found: ignore
 
       - name: Upload agentic raw results
         if: ${{ always() && inputs.scenario-type == 'agentic-coding' }}

.github/workflows/e2e-tests.yml

@@ -242,6 +242,7 @@ jobs:
             conc: ${{ matrix.config.conc[0] }}
             kv-offloading: ${{ matrix.config.kv-offloading }}
             kv-offload-backend: ${{ matrix.config.kv-offload-backend }}
+            total-cpu-dram-gb: ${{ matrix.config['total-cpu-dram-gb'] }}
             duration: ${{ inputs.duration-override != '' && inputs.duration-override || matrix.config.duration }}
             run-eval: false
             scenario-type: agentic-coding

benchmarks/multi_node/agentic/kimik3_fp4_mi355x_vllm-disagg.sh

@@ -0,0 +1,166 @@
+#!/usr/bin/env bash
+
+# Agentic trace-replay recipe for a disaggregated vLLM server on MI355X
+# (Kimi-K3 MXFP4, DSpark draft, 1P1D and 2P1D TP8).
+#
+# CI sibling of agentic/dsv4_fp4_mi355x_sglang-disagg.sh: driven entirely by
+# environment variables and submits a SLURM job via amd_utils/submit.sh. Two
+# things differ from that recipe because the engine differs:
+#
+#   * the KV offload tier is vLLM's own SimpleCPUOffloadConnector, not HiCache,
+#     so none of the HICACHE_*/MC_* tunables apply. server_vllm.sh composes it
+#     with MoRIIO under MultiConnector on the prefill side (see
+#     build_kv_transfer_configs there).
+#   * the MoRI queue sizing is NOT set here. env.sh pins Kimi-K3 defaults
+#     (MORI_IO_SQ_BACKOFF_TIMEOUT_US=500000, QP_MAX_SEND_WR=8192, CQE=16384,
+#     SGE=2, TC_DISABLE=0) and job.slurm -e's them into the engine container;
+#     setting them here as well would create a second source of truth.
+#
+# The serve body (TP, cudagraph capture, fp8 KV, DSpark speculative-config, 1M
+# context) lives in amd_utils/models_vllm.yaml under Kimi-K3.
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "$SCRIPT_DIR/../../benchmark_lib.sh"
+
+# ISL/OSL are deliberately absent from this list. Agentic matrix entries carry no
+# isl/osl (generate_sweep_configs.py builds them from the trace corpus instead),
+# and they reach submit.sh only as positional filler that the agentic path never
+# reads -- BENCH_INPUT_LEN/BENCH_OUTPUT_LEN are consumed by bench.sh, not by
+# trace_replay.sh. Requiring them would abort the run over two unused numbers.
+check_env_vars \
+    CONC_LIST \
+    IMAGE \
+    SPEC_DECODING \
+    MODEL_PATH \
+    PREFILL_NUM_WORKERS \
+    PREFILL_TP \
+    PREFILL_EP \
+    PREFILL_DP_ATTN \
+    DECODE_NUM_WORKERS \
+    DECODE_TP \
+    DECODE_EP \
+    DECODE_DP_ATTN \
+    PREFILL_NODES \
+    DECODE_NODES \
+    RANDOM_RANGE_RATIO \
+    DURATION \
+    KV_OFFLOADING \
+    IS_AGENTIC \
+    FRAMEWORK
+
+if [[ -n "$SLURM_JOB_ID" ]]; then
+  echo "JOB $SLURM_JOB_ID running on $SLURMD_NODENAME"
+fi
+
+set -x
+
+# Use upstreamed multi_node scripts (no external clone needed)
+cd "$GITHUB_WORKSPACE/benchmarks/multi_node/amd_utils" || exit 1
+
+export ISL="${ISL:-0}"
+export OSL="${OSL:-0}"
+
+# ── SLURM / image ──
+# 12 h rather than the fixed-seq-len 8 h: a 3600 s replay sits behind ~2.8 TB of
+# weight loading on both roles plus cudagraph capture, and benchmark-multinode-tmpl.yml
+# already allows 780 min of wall clock for agentic-coding jobs.
+export TIME_LIMIT="${TIME_LIMIT:-12:00:00}"
+export MODEL_PATH=$MODEL_PATH
+export MODEL_NAME=$MODEL_NAME
+export CONTAINER_IMAGE=$IMAGE
+
+# ── MoRIIO write + wait_all ──
+# job.slurm derives a node-local image from CONTAINER_IMAGE with the exact
+# ROCm/mori commit SGLang MI35x pins (#341 / IOEngine.wait_all). The live server
+# only verifies the baked capability; it never installs a wheel at startup.
+export MORI_WAITALL_BUILD="${MORI_WAITALL_BUILD:-1}"
+export MORI_WAITALL_COMMIT="${MORI_WAITALL_COMMIT:-f7e6ac6863c53821bc7afb91a578cc6ce38fcad0}"
+export MORIIO_READ_MODE="${MORIIO_READ_MODE:-false}"
+
+# job.slurm's fallback pin is routinely garbage-collected from Docker Hub
+# (only ~16 nightlies are retained), so pin the tag this arm was validated
+# against. Keep it equal to the router version recorded in amd-master.yaml.
+export VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly-20260809-d2ba586}"
+
+# ── Identity / result naming ──
+export MODEL_PREFIX="${MODEL_PREFIX:-kimik3}"
+export PRECISION="${PRECISION:-fp4}"
+export RESULT_FILENAME="${RESULT_FILENAME:-${RUNNER_NAME:-kimik3-fp4-disagg-agentic}}"
+
+# ── Agentic benchmark params ──
+export DURATION="${DURATION:-3600}"
+# K3's native context. Kept equal to the --max-model-len in models_vllm.yaml so
+# the client's --max-context-length matches what the servers actually serve.
+export MAX_MODEL_LEN="${MAX_MODEL_LEN:-1048576}"
+
+# ── KV cache offloading ──
+# KV_OFFLOADING=none | dram (from YAML). The only backend this arm implements is
+# vLLM's built-in SimpleCPUOffloadConnector ("vllm-simple"); the capacity itself
+# comes from the matrix as TOTAL_CPU_DRAM_GB, sized from the PREFILL worker's
+# per-node GPU footprint because only prefill offloads to CPU DRAM today.
+#
+# check_env_vars above already rejects KV_OFFLOADING=dram without a backend and
+# without a positive TOTAL_CPU_DRAM_GB (benchmark_lib.sh), so the only thing left
+# to assert here is that the backend is one this arm can actually serve.
+export KV_OFFLOADING="${KV_OFFLOADING:-none}"
+if [[ "$KV_OFFLOADING" != "none" ]]; then
+  export KV_OFFLOAD_BACKEND="${KV_OFFLOAD_BACKEND:-vllm-simple}"
+  if [[ "$KV_OFFLOAD_BACKEND" != "vllm-simple" ]]; then
+    echo "ERROR: KV_OFFLOAD_BACKEND=$KV_OFFLOAD_BACKEND unsupported on this arm (vllm-simple only)" >&2
+    exit 1
+  fi
+fi
+
+# ── Server metrics ──
+# vLLM exposes Prometheus /metrics on every engine port unconditionally (no
+# --enable-metrics equivalent), so the aiperf server-metrics scrape always has
+# somewhere to point. server_vllm.sh builds the per-worker URL list.
+export ENABLE_METRICS="${ENABLE_METRICS:-1}"
+
+# ── MTP ──
+# Labels the result JSON only: the draft is configured by --speculative-config in
+# models_vllm.yaml (k=2 DSpark), and server_vllm.sh never reads this.
+export DECODE_MTP_SIZE="${DECODE_MTP_SIZE:-2}"
+
+# Derive EP/DP enable flags from the topology inputs.
+if [[ "${PREFILL_EP:-1}" -eq 1 ]]; then
+export PREFILL_ENABLE_EP=false
+else
+export PREFILL_ENABLE_EP=true
+fi
+
+if [[ "$PREFILL_DP_ATTN" == "true" ]]; then
+export PREFILL_ENABLE_DP=true
+else
+export PREFILL_ENABLE_DP=false
+fi
+
+if [[ "${DECODE_EP:-1}" -eq 1 ]]; then
+export DECODE_ENABLE_EP=false
+else
+export DECODE_ENABLE_EP=true
+fi
+
+if [[ "$DECODE_DP_ATTN" == "true" ]]; then
+export DECODE_ENABLE_DP=true
+else
+export DECODE_ENABLE_DP=false
+fi
+
+# Launch the job. CONC_LIST is space-delimited in YAML; submit.sh wants 'x'.
+JOB_ID=$(bash ./submit.sh $PREFILL_NODES \
+    $PREFILL_NUM_WORKERS \
+    $DECODE_NODES \
+    $DECODE_NUM_WORKERS \
+    $ISL $OSL "${CONC_LIST// /x}" inf \
+    ${PREFILL_ENABLE_EP} ${PREFILL_ENABLE_DP} \
+    ${DECODE_ENABLE_EP} ${DECODE_ENABLE_DP} \
+    ${PREFILL_TP} ${DECODE_TP} \
+    ${RANDOM_RANGE_RATIO})
+
+if [[ $? -ne 0 ]]; then
+    echo "Failed to submit job" >&2
+    exit 1
+fi
+
+echo "$JOB_ID"

benchmarks/multi_node/amd_utils/apply_k3_moriio_patches.sh

@@ -0,0 +1,67 @@
+#!/usr/bin/env bash
+# Apply #51052 Kimi-K3 MoRIIO hybrid transfer into the engine container.
+#
+# This is deliberately a unified diff rather than a full-file overlay. An
+# overlay replaces whole modules, so it also pins every unrelated symbol in
+# those files to whichever vLLM it was captured from: a cb810-era snapshot
+# silently downgraded vllm/v1/core/kv_cache_utils.py and the newer engine died
+# with "cannot import name 'update_kv_cache_capacity'". The diff only touches
+# what #51052 changes, and fails loudly when it no longer fits the image.
+set -euo pipefail
+
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PATCH_FILE="${K3_MORIIO_PATCH:-$HERE/patches/k3_moriio_51052.patch}"
+ROOT="$(python3 -c 'import importlib.util as u, os; print(os.path.dirname(os.path.dirname(u.find_spec("vllm").origin)))')"
+
+if [[ -z "$ROOT" || ! -d "$ROOT/vllm" ]]; then
+    echo "[k3-moriio] ERROR: could not resolve vLLM root (ROOT='$ROOT')" >&2
+    exit 1
+fi
+
+MORIIO_DIR="$ROOT/vllm/distributed/kv_transfer/kv_connector/v1/moriio"
+if grep -RqsE '_draft_only_layers|as_attn_mamba' "$MORIIO_DIR" \
+   && python3 -c 'from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import MoRIIOConnector' 2>/dev/null; then
+    echo "[k3-moriio] #51052 already applied (markers + import ok)"
+    exit 0
+fi
+
+if [[ ! -f "$PATCH_FILE" ]]; then
+    echo "[k3-moriio] ERROR: patch not found at $PATCH_FILE" >&2
+    exit 1
+fi
+
+if grep -q $'\r' "$PATCH_FILE"; then
+    tmp=$(mktemp)
+    tr -d '\r' < "$PATCH_FILE" > "$tmp"
+    PATCH_FILE="$tmp"
+    trap 'rm -f "$tmp"' EXIT
+fi
+
+if (cd "$ROOT" && git apply -p1 "$PATCH_FILE" 2>/dev/null); then
+    echo "[k3-moriio] applied #51052 with git apply"
+elif patch -p1 -d "$ROOT" --forward --no-backup-if-mismatch < "$PATCH_FILE"; then
+    echo "[k3-moriio] applied #51052 with patch"
+else
+    echo "[k3-moriio] ERROR: #51052 does not apply to this image." >&2
+    echo "[k3-moriio] Regenerate patches/k3_moriio_51052.patch against it; do not" >&2
+    echo "[k3-moriio] substitute a full-file overlay, which downgrades unrelated modules." >&2
+    find "$ROOT" -name '*.rej' -newermt '-5 minutes' -print >&2 2>/dev/null || true
+    exit 1
+fi
+
+python3 - <<'PY'
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_common as c
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
+    MoRIIOConnector,
+)
+from vllm.v1.core import kv_cache_utils
+
+assert hasattr(c, "as_attn_mamba"), "as_attn_mamba missing after patch"
+assert MoRIIOConnector is not None
+# The engine imports this from kv_cache_utils; a stale replacement of that
+# module is exactly the failure mode this script now refuses to create.
+assert hasattr(kv_cache_utils, "update_kv_cache_capacity"), (
+    "kv_cache_utils lost update_kv_cache_capacity"
+)
+print("[k3-moriio] #51052 applied; IMPORT_OK")
+PY

benchmarks/multi_node/amd_utils/ensure_mori_wait_all.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+# Verify that job.slurm selected the source-built mori #341 image. Serving must
+# never mutate site-packages or silently fall back to the Python polling path.
+set -euo pipefail
+
+python3 - <<'PY'
+import mori
+from mori.io import IOEngine, StatusCode
+
+version = getattr(mori, "__version__", "?")
+wait_all = hasattr(IOEngine, "wait_all")
+print("[k3-mori-waitall] mori=", version, "at", mori.__file__)
+print("[k3-mori-waitall] IOEngine.wait_all=", wait_all)
+if not wait_all:
+    raise RuntimeError(
+        "IOEngine.wait_all missing: job.slurm must prepare the pinned "
+        "ROCm/mori#341 derived image before serving"
+    )
+assert StatusCode is not None
+PY

benchmarks/multi_node/amd_utils/env.sh

@@ -50,11 +50,17 @@ export NCCL_IB_HCA=${NCCL_IB_HCA:-$IBDEVICES}
 # =============================================================================
 # Shared by the vLLM MoRIIOConnector and the SGLang/MoRI KV-transfer path.
 
-export MORI_IO_SQ_BACKOFF_TIMEOUT_US=50000
-export MORI_IO_QP_MAX_SEND_WR=16384
-export MORI_IO_QP_MAX_CQE=32768
-export MORI_IO_QP_MAX_SGE=2
-export MORI_IO_TC_DISABLE=0
+if [[ "${MODEL_NAME:-}" == "Kimi-K3" ]]; then
+    export MORI_IO_SQ_BACKOFF_TIMEOUT_US="${MORI_IO_SQ_BACKOFF_TIMEOUT_US:-500000}"
+    export MORI_IO_QP_MAX_SEND_WR="${MORI_IO_QP_MAX_SEND_WR:-8192}"
+    export MORI_IO_QP_MAX_CQE="${MORI_IO_QP_MAX_CQE:-16384}"
+else
+    export MORI_IO_SQ_BACKOFF_TIMEOUT_US="${MORI_IO_SQ_BACKOFF_TIMEOUT_US:-50000}"
+    export MORI_IO_QP_MAX_SEND_WR="${MORI_IO_QP_MAX_SEND_WR:-16384}"
+    export MORI_IO_QP_MAX_CQE="${MORI_IO_QP_MAX_CQE:-32768}"
+fi
+export MORI_IO_QP_MAX_SGE="${MORI_IO_QP_MAX_SGE:-2}"
+export MORI_IO_TC_DISABLE="${MORI_IO_TC_DISABLE:-0}"
 
 # QoS/DSCP configuration
 # Priority order: 1) Set by runner, 2) Detect via nicctl, 3) Detect from hostname
@@ -180,6 +186,15 @@ $1 == "DSCP" && $2 == ":" && $NF == p {
     set +x
     echo "[INFO] IBDEVICES=$IBDEVICES  UCX_NET_DEVICES=$UCX_NET_DEVICES  NCCL_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME  UCX_IB_GID_INDEX=$UCX_IB_GID_INDEX  UCX_IB_TRAFFIC_CLASS=${UCX_IB_TRAFFIC_CLASS:-unset}"
 
+    if [[ "$MODEL_NAME" == "Kimi-K3" ]]; then
+        export MORI_IO_SQ_BACKOFF_TIMEOUT_US="${MORI_IO_SQ_BACKOFF_TIMEOUT_US:-500000}"
+        export MORI_IO_QP_MAX_SEND_WR="${MORI_IO_QP_MAX_SEND_WR:-8192}"
+        export MORI_IO_QP_MAX_CQE="${MORI_IO_QP_MAX_CQE:-16384}"
+        export MORI_IO_QP_MAX_SGE="${MORI_IO_QP_MAX_SGE:-2}"
+        export MORI_IO_TC_DISABLE="${MORI_IO_TC_DISABLE:-0}"
+        echo "[INFO] Kimi-K3 MoRI IO: SQ_BACKOFF=${MORI_IO_SQ_BACKOFF_TIMEOUT_US}us SEND_WR=${MORI_IO_QP_MAX_SEND_WR} CQE=${MORI_IO_QP_MAX_CQE} SGE=${MORI_IO_QP_MAX_SGE} TC_DISABLE=${MORI_IO_TC_DISABLE}"
+    fi
+
 else
     # =========================================================================
     # SGLang-specific environment