[AMD][DI][CI] 6/N Add Kimi-K2.6 MXFP4 wide-EP16 2P1D nightly recipes (#33333)

This commit is contained in:
Zhaoyi Li
2026-08-03 18:51:55 -07:00
committed by GitHub
parent 03f44c978a
commit a84e70eb1e
4 changed files with 633 additions and 63 deletions
+363 -57
View File
@@ -101,11 +101,15 @@ python3 -c 'import yaml' 2>/dev/null || pip install pyyaml -q 2>/dev/null \
# Emit KEY=value lines and eval them (robust single-level command substitution;
# avoids a nested read<<EOF/$(<<PY) heredoc that misparses on some shells).
RECIPE_VARS="$(python3 - "$CONFIG_FILE" <<'PY'
import sys, yaml
import sys, yaml, shlex
r = yaml.safe_load(open(sys.argv[1]))
rt = r["runtime"]; b = r["backend"]["sglang_config"]; bn = r["bench"]
res = r.get("resources", {})
def emit(k, v): print(f"{k}={v}")
# Quoted emit for values that may contain spaces (per-role wide-EP flag/env
# strings). eval treats the RHS as a single shell word. Existing emits stay
# space-free so the EP<=8 recipes are byte-identical.
def emitq(k, v): print(f"{k}={shlex.quote(str(v))}")
emit("IMAGE", rt["image"])
# Attention backend: single (`attention_backend`) or split
# (`prefill_attention_backend`/`decode_attention_backend`). Empty when absent so
@@ -114,6 +118,20 @@ emit("ATTN", rt.get("attention_backend", ""))
emit("PATTN", rt.get("prefill_attention_backend", ""))
emit("DATTN", rt.get("decode_attention_backend", ""))
emit("IB", rt["ib_devices"])
# Wide-EP (EP > GPUs/node) options; empty for the single-node EP<=8 recipes so
# the flags/env below are dropped and their argv stays byte-identical.
# moe_a2a_backend -> --moe-a2a-backend <x> --deepep-mode normal (MoE all-to-all)
# dist_socket_ifname -> NCCL_/GLOO_SOCKET_IFNAME for cross-node torch-dist init
emit("A2A", rt.get("moe_a2a_backend", ""))
emit("DIST_SOCK", rt.get("dist_socket_ifname", ""))
# KV transfer backend for the P->D handoff. Defaults to mori (the pre-wide
# hardcoded value) so EP<=8 recipes stay byte-identical; wide-EP spur recipes
# set mooncake, which is the validated cross-node KV path there.
emit("XFER", rt.get("kv_transfer_backend", "mori"))
# SGLANG_USE_ROCM700A toggles the ROCm-7.0.0-alpha codepath. Default 1 (the
# pre-wide hardcoded value) keeps EP<=8 recipes byte-identical; the validated
# wide-EP run on the rocm720 0715 image needs 0, set via runtime.rocm700a.
emit("ROCM700A", rt.get("rocm700a", 1))
emit("PPORT", rt["prefill_port"])
emit("DPORT", rt["decode_port"])
emit("PBOOT", rt["prefill_bootstrap_port"])
@@ -133,6 +151,29 @@ emit("PTP", b["prefill"]["tensor-parallel-size"])
emit("DTP", b["decode"]["tensor-parallel-size"])
emit("PEP", b["prefill"].get("expert-parallel-size", 1))
emit("PDP", b["prefill"].get("data-parallel-size", 1))
# Decode-side EP/DP: default to the prefill values so a recipe that omits them
# (every EP<=8 recipe today, where both roles are identical) is byte-identical.
# Oren's wide-EP recipes set decode EP/DP=16 while prefill stays EP8.
emit("DEP", b["decode"].get("expert-parallel-size", b["prefill"].get("expert-parallel-size", 1)))
emit("DDP", b["decode"].get("data-parallel-size", b["prefill"].get("data-parallel-size", 1)))
# Wide-EP per-role overrides (optional `runtime.wide_ep` block). Absent for
# EP<=8 recipes -> all defaults collapse to the existing single-value knobs, so
# the generated prefill/decode argv is unchanged. Present only in the asymmetric
# narrow-prefill/wide-decode (Oren) EP16 recipes.
we = rt.get("wide_ep", {}) or {}
emit("KVDTYPE", we.get("kv_cache_dtype", ""))
emit("PMEMFRAC", we.get("prefill_mem_fraction_static", rt["mem_fraction_static"]))
emit("DMEMFRAC", we.get("decode_mem_fraction_static", rt["mem_fraction_static"]))
emit("PCHUNK", we.get("prefill_chunked_prefill_size", rt["chunked_prefill_size"]))
emit("PMAXREQ", we.get("prefill_max_running_requests", rt["max_running_requests"]))
emit("DMAXREQ", we.get("decode_max_running_requests", rt["max_running_requests"]))
emitq("PEXTRA", we.get("prefill_extra_flags", ""))
emitq("DEXTRA", we.get("decode_extra_flags", ""))
emitq("WECOMMON", we.get("common_extra_flags", ""))
def render_env(d):
return " ".join(f"-e {k}={v}" for k, v in (d or {}).items())
emitq("PENV", render_env(we.get("prefill_extra_env")))
emitq("DENV", render_env(we.get("decode_extra_env")))
m = r.get("mtp", {}) or {}
emit("MTP_ENABLED", 1 if m.get("enabled") else 0)
emit("MTP_ALGO", m.get("algorithm", "EAGLE"))
@@ -165,7 +206,27 @@ eval "$RECIPE_VARS"
if [[ -n "${IMAGE_OVERRIDE:-}" ]]; then
IMAGE="$IMAGE_OVERRIDE"
fi
echo "recipe: image=$IMAGE attn=${ATTN:-$PATTN/$DATTN} ib=$IB ptp=$PTP dtp=$DTP concs=$CONCS isl=$ISL osl=$OSL"
# Nodes per engine: an engine whose TP exceeds one node's GPU count spans
# ceil(TP/GPUS_PER_NODE) nodes and needs torch-dist multi-node init. EP<=8
# recipes give 1 (single node) so all downstream multi-node logic no-ops.
GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
PN_PER=$(( (PTP + GPUS_PER_NODE - 1) / GPUS_PER_NODE ))
DN_PER=$(( (DTP + GPUS_PER_NODE - 1) / GPUS_PER_NODE ))
# torch-dist rendezvous port for a multi-node engine. Default 29500 (torch's
# conventional MASTER_PORT); avoids :5000, which a node-local daemon holds on
# some clusters (spur). Overridable per environment.
DIST_PORT="${DIST_PORT:-29500}"
# A role may have multiple single-node engines (PW>1, each PN_PER=1: the router
# fans out over them) OR one multi-node wide engine (PN_PER>1, workers==1). What
# is NOT wired is *multiple copies of a multi-node engine* (a wide engine with
# workers>1), because the drive split assumes one contiguous node block per wide
# engine and the router only knows one endpoint per wide engine.
if (( (PN_PER > 1 && PW > 1) || (DN_PER > 1 && DW > 1) )); then
echo "ERROR: a wide engine (nodes/engine>1) cannot have >1 worker of that role (got PW=$PW PN_PER=$PN_PER DW=$DW DN_PER=$DN_PER)" >&2
exit 1
fi
echo "recipe: image=$IMAGE attn=${ATTN:-$PATTN/$DATTN} ib=$IB ptp=$PTP dtp=$DTP pn_per=$PN_PER dn_per=$DN_PER a2a=${A2A:-<none>} concs=$CONCS isl=$ISL osl=$OSL"
# ---------------------------------------------------------------------------
# Shared NFS scratch (visible to login node + compute nodes). Raw bench output
@@ -219,7 +280,7 @@ fi
DSV4_ENV=(
-e SGLANG_DEFAULT_THINKING=1 -e SGLANG_DSV4_REASONING_EFFORT=max
-e SGLANG_OPT_DEEPGEMM_HC_PRENORM=false -e SGLANG_USE_AITER=1
-e SGLANG_USE_ROCM700A=1 -e SGLANG_OPT_USE_FUSED_COMPRESS=true
-e SGLANG_USE_ROCM700A=$ROCM700A -e SGLANG_OPT_USE_FUSED_COMPRESS=true
-e SGLANG_OPT_USE_FUSED_COMPRESS_TRITON=true
-e SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton
-e SGLANG_OPT_FP8_WO_A_GEMM=false -e SGLANG_OPT_USE_JIT_INDEXER_METADATA=false
@@ -234,6 +295,40 @@ DSV4_ENV_STR="${DSV4_ENV[*]}"
# DSV4 env must not leak into it; the DSV4 recipes keep the string above.
[[ "$HAS_MODEL" == "1" ]] && DSV4_ENV_STR=""
MORI_ENV="-e MORI_DISABLE_AUTO_XGMI=1 -e NCCL_IB_HCA=ionic -e NCCL_IB_GID_INDEX=1 -e NCCL_CROSS_NIC=1"
# Wide-EP (engine spans >1 node) adds mori all-to-all MoE tuning + the cross-node
# torch-dist socket NIC. Gated on nodes-per-engine>1 so EP<=8 recipes are untouched.
if (( PN_PER > 1 || DN_PER > 1 )); then
# TC=104 (SL=3) = the ionic lossless RoCE queue (DSCP26/pri3); TC=96 is the
# lossy pri0 queue (~1% BW) and wedges cross-node a2a under load. bf16
# dispatch/combine matches the validated wide-EP run (job 13196).
# The base MORI_ENV sets NCCL_IB_HCA=ionic (a spur-ism); on this fabric the
# IB device names are the recipe's $IB (rdma0..7), and a wide engine's
# cross-node TP/attention collectives ride NCCL, so point NCCL at the real
# HCAs. Docker last-wins => this overrides the base value for wide recipes.
MORI_ENV="$MORI_ENV \
-e NCCL_IB_HCA=$IB \
-e MORI_IB_GID_INDEX=1 \
-e SGLANG_MORI_DISPATCH_DTYPE=bf16 -e SGLANG_MORI_COMBINE_DTYPE=bf16 \
-e SGLANG_MORI_QP_PER_TRANSFER=4 -e SGLANG_MORI_NUM_WORKERS=4 \
-e MORI_IO_SQ_BACKOFF_TIMEOUT_US=50000 -e MORI_IO_QP_MAX_SEND_WR=16384 \
-e MORI_IO_QP_MAX_CQE=32768 -e MORI_IO_QP_MAX_SGE=4 \
-e MORI_SHMEM_MODE=ISOLATION -e MORI_EP_LAUNCH_CONFIG_MODE=AUTO -e MORI_APP_LOG_LEVEL=INFO \
-e MORI_RDMA_SL=3 -e MORI_RDMA_TC=104 -e MORI_IO_SL=3 -e MORI_IO_TC=104 -e MORI_IO_TC_DISABLE=0 \
-e SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD=4096 \
-e SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=3600 -e SGLANG_DISAGGREGATION_WAITING_TIMEOUT=3600 \
-e SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS=32 -e SGLANG_EAGER_INPUT_NO_COPY=true \
-e MORI_BOOTSTRAP_TIMEOUT=300"
# The overlap plan stream is a wide-EP perf knob, but it makes
# EAGLEWorkerV2.verify call attn_backend.update_verify_buffers_to_fill_after_draft,
# which the DSV4 MLA backend does not implement (NotImplementedError). Only
# enable it for non-MTP wide-EP; MTP wide-EP uses the plain verify path (as EP8 does).
if [[ "$MTP_ENABLED" != "1" ]]; then
MORI_ENV="$MORI_ENV -e SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1"
fi
if [[ -n "$DIST_SOCK" ]]; then
MORI_ENV="$MORI_ENV -e GLOO_SOCKET_IFNAME=$DIST_SOCK -e NCCL_SOCKET_IFNAME=$DIST_SOCK -e MORI_SOCKET_IFNAME=$DIST_SOCK"
fi
fi
# Model-specific docker `-e` env + sglang server args from the recipe's optional
# `model:` block, written as bash arrays to model_flags.sh (sourced by
@@ -263,8 +358,8 @@ with open(sys.argv[2], "w") as f:
f.write(f"MODEL_SERVER_ARGS=({q(server_args)})\n")
PY
# Resolve host ionic userspace mounts on each compute node. This keeps server
# and benchmark containers compatible with the host RDMA ABI.
# Resolve host ionic userspace mounts on each compute node. This keeps server,
# benchmark, and wide-EP containers compatible with the host RDMA ABI.
cat > "$WORKDIR/ionic_mounts.sh" <<'IONIC_EOF'
IONIC_MOUNTS=()
_ionic_provider="/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so"
@@ -292,23 +387,47 @@ fi
IONIC_EOF
# Optional topology / speculative-decode flags driven by the recipe. Base recipes
# (EP1/DP1, no mtp) leave EXTRA_FLAGS empty, preserving prior behavior exactly.
EXTRA_FLAGS=""
(( PDP > 1 )) && EXTRA_FLAGS="$EXTRA_FLAGS --enable-dp-attention --dp-size $PDP"
(( PEP > 1 )) && EXTRA_FLAGS="$EXTRA_FLAGS --ep-size $PEP"
[[ -n "$MAXTOK" ]] && EXTRA_FLAGS="$EXTRA_FLAGS --max-total-tokens $MAXTOK"
# (EP1/DP1, no mtp) leave the extra strings empty, preserving prior behavior.
#
# EP/DP are now PER ROLE: the DP-attention + ep-size flags come from PDP/PEP for
# prefill and DDP/DEP for decode. For every EP<=8 recipe DDP==PDP and DEP==PEP
# (decode inherits prefill), so both role strings equal the old single EXTRA_FLAGS
# and the generated argv is byte-identical. Oren's wide-EP recipes set decode
# EP/DP=16 while prefill stays EP8.
PREFILL_DPEP=""
(( PDP > 1 )) && PREFILL_DPEP="$PREFILL_DPEP --enable-dp-attention --dp-size $PDP"
(( PEP > 1 )) && PREFILL_DPEP="$PREFILL_DPEP --ep-size $PEP"
DECODE_DPEP=""
(( DDP > 1 )) && DECODE_DPEP="$DECODE_DPEP --enable-dp-attention --dp-size $DDP"
(( DEP > 1 )) && DECODE_DPEP="$DECODE_DPEP --ep-size $DEP"
# Flags shared by both roles (a2a backend, mtp). --max-total-tokens stays here for
# non-wide recipes; wide recipes carry a per-role prefill_max_total via wide_ep.
EXTRA_COMMON=""
[[ -n "$A2A" ]] && EXTRA_COMMON="$EXTRA_COMMON --moe-a2a-backend $A2A --deepep-mode normal"
[[ -n "$MAXTOK" ]] && EXTRA_COMMON="$EXTRA_COMMON --max-total-tokens $MAXTOK"
if [[ "$MTP_ENABLED" == "1" ]]; then
EXTRA_FLAGS="$EXTRA_FLAGS --speculative-algorithm $MTP_ALGO \
EXTRA_COMMON="$EXTRA_COMMON --speculative-algorithm $MTP_ALGO \
--speculative-num-steps $MTP_STEPS --speculative-eagle-topk $MTP_TOPK \
--speculative-num-draft-tokens $MTP_DRAFT"
# EAGLE3 (and other draft-model algos) need an external draft checkpoint;
# built-in EAGLE (DSV4) omits draft_model_path and this stays unset.
if [[ -n "$MTP_DRAFT_PATH" ]]; then
DRAFT_RESOLVED="$(resolve_snapshot "$MTP_DRAFT_PATH")" || exit 1
EXTRA_FLAGS="$EXTRA_FLAGS --speculative-draft-model-path $DRAFT_RESOLVED"
EXTRA_COMMON="$EXTRA_COMMON --speculative-draft-model-path $DRAFT_RESOLVED"
fi
fi
echo "extra flags: ${EXTRA_FLAGS:-<none>} (pep=$PEP pdp=$PDP mtp=$MTP_ENABLED algo=$MTP_ALGO)"
# Prefix a leading space only when the arg is non-empty (keeps EP<=8 argv byte-
# identical: the wide-only strings are empty and contribute nothing).
sp() { [[ -n "$1" ]] && printf ' %s' "$1"; return 0; }
# --kv-cache-dtype is emitted only when the recipe sets wide_ep.kv_cache_dtype;
# the pre-wide DSV4 path had no such flag, so EP<=8 recipes omit it.
KV_FLAG=""
[[ -n "$KVDTYPE" ]] && KV_FLAG=" --kv-cache-dtype $KVDTYPE"
# Assemble the per-role tail: role DP/EP + shared + wide common + role-specific
# wide extras. All wide pieces (WECOMMON/PEXTRA/DEXTRA) are empty for EP<=8.
PREFILL_TAIL="$PREFILL_DPEP$EXTRA_COMMON$(sp "$WECOMMON")$(sp "$PEXTRA")"
DECODE_TAIL="$DECODE_DPEP$EXTRA_COMMON$(sp "$WECOMMON")$(sp "$DEXTRA")"
echo "prefill tail:${PREFILL_TAIL:-<none>} | decode tail:${DECODE_TAIL:-<none>} (pep=$PEP pdp=$PDP dep=$DEP ddp=$DDP mtp=$MTP_ENABLED)"
if [[ "$HAS_MODEL" == "1" ]]; then
# Generic path (e.g. Kimi): attention + swa from the recipe, model parsers /
@@ -321,25 +440,46 @@ if [[ "$HAS_MODEL" == "1" ]]; then
[[ -n "$DATTN" ]] && ATTN_FLAGS="$ATTN_FLAGS --decode-attention-backend $DATTN"
SWA_FLAG=""
[[ -n "$SWA" ]] && SWA_FLAG=" --swa-full-tokens-ratio $SWA"
COMMON_FLAGS="--trust-remote-code --tp $PTP --disable-radix-cache \
$ATTN_FLAGS --max-running-requests $MAXREQ --page-size $PAGE \
--mem-fraction-static $MEMFRAC$SWA_FLAG \
PREFILL_COMMON_FLAGS="--trust-remote-code --tp $PTP --disable-radix-cache \
$ATTN_FLAGS --max-running-requests $PMAXREQ --page-size $PAGE \
--mem-fraction-static $PMEMFRAC$SWA_FLAG \
--chunked-prefill-size $PCHUNK \
--disaggregation-transfer-backend $XFER --disaggregation-ib-device $IB$KV_FLAG$PREFILL_TAIL"
DECODE_COMMON_FLAGS="--trust-remote-code --tp $DTP --disable-radix-cache \
$ATTN_FLAGS --max-running-requests $DMAXREQ --page-size $PAGE \
--mem-fraction-static $DMEMFRAC$SWA_FLAG \
--chunked-prefill-size $CHUNK \
--disaggregation-transfer-backend mori --disaggregation-ib-device $IB$EXTRA_FLAGS"
--disaggregation-transfer-backend $XFER --disaggregation-ib-device $IB$KV_FLAG$DECODE_TAIL"
else
# DSV4 path: byte-identical to the pre-Kimi launcher.
COMMON_FLAGS="--trust-remote-code --tp $PTP --disable-radix-cache \
--attention-backend $ATTN --max-running-requests $MAXREQ --page-size $PAGE \
--mem-fraction-static $MEMFRAC --swa-full-tokens-ratio $SWA \
# DSV4 path: for EP<=8 recipes (PTP==DTP, no wide_ep) both role strings equal
# the pre-Kimi launcher's COMMON_FLAGS exactly.
PREFILL_COMMON_FLAGS="--trust-remote-code --tp $PTP --disable-radix-cache \
--attention-backend $ATTN --max-running-requests $PMAXREQ --page-size $PAGE \
--mem-fraction-static $PMEMFRAC --swa-full-tokens-ratio $SWA \
--chunked-prefill-size $PCHUNK --disable-shared-experts-fusion \
--tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 \
--disaggregation-transfer-backend $XFER --disaggregation-ib-device $IB$KV_FLAG$PREFILL_TAIL"
DECODE_COMMON_FLAGS="--trust-remote-code --tp $DTP --disable-radix-cache \
--attention-backend $ATTN --max-running-requests $DMAXREQ --page-size $PAGE \
--mem-fraction-static $DMEMFRAC --swa-full-tokens-ratio $SWA \
--chunked-prefill-size $CHUNK --disable-shared-experts-fusion \
--tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 \
--disaggregation-transfer-backend mori --disaggregation-ib-device $IB$EXTRA_FLAGS"
--disaggregation-transfer-backend $XFER --disaggregation-ib-device $IB$KV_FLAG$DECODE_TAIL"
fi
DOCKER_COMMON="--rm --network host --ipc host --shm-size 32g --privileged \
--security-opt seccomp=unconfined \
--device /dev/kfd --device /dev/dri --device /dev/infiniband \
-v /it-share:/it-share:ro -v $HOME:/host_home $CHECKOUT_DOCKER_ARGS"
# Optional extra docker args (e.g. bind-mounting a locally-built lib for
# validation). Empty by default so the docker argv is byte-identical otherwise.
[[ -n "${EXTRA_DOCKER_ARGS:-}" ]] && DOCKER_COMMON="$DOCKER_COMMON ${EXTRA_DOCKER_ARGS}"
# Per-role wide-EP docker env (MORI dispatch-token tuning etc.). Empty for EP<=8
# recipes; carries its own leading space so an empty value leaves the docker argv
# byte-identical (no stray double space).
PENV_ARG=""; [[ -n "$PENV" ]] && PENV_ARG=" $PENV"
DENV_ARG=""; [[ -n "$DENV" ]] && DENV_ARG=" $DENV"
# ---------------------------------------------------------------------------
# Write per-role scripts that srun dispatches to each compute node.
@@ -490,6 +630,80 @@ print(f"[checkout-router] Router={Router}")
PY
EOF
if (( PN_PER > 1 || DN_PER > 1 )); then
# Wide-EP path: an engine spans >1 node. Entry scripts add torch-dist
# rendezvous args from NODE_RANK/NNODES/DIST_ADDR; launch scripts take the
# per-node rank as $1..$3 and forward it into the container. Only reached when
# a recipe sets TP>GPUS_PER_NODE, so EP<=8 recipes never take this branch.
cat > "$WORKDIR/prefill_entry.sh" <<EOF
#!/bin/bash
set -euo pipefail
CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
source "\$CIDIR/model_flags.sh"
bash "\$CIDIR/install_checkout_sglang.sh"
if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
fi
DIST_ARGS=""
if [[ "\${NNODES:-1}" != "1" ]]; then
DIST_ARGS="--nnodes \$NNODES --node-rank \$NODE_RANK --dist-init-addr \$DIST_ADDR:\${DIST_PORT:-29500}"
fi
exec python3 -m sglang.launch_server \
--model-path $MODEL_PATH --host 0.0.0.0 --port $PPORT \
$PREFILL_COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \$DIST_ARGS \
--disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT
EOF
cat > "$WORKDIR/decode_entry.sh" <<EOF
#!/bin/bash
set -euo pipefail
CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
source "\$CIDIR/model_flags.sh"
bash "\$CIDIR/install_checkout_sglang.sh"
if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
fi
DIST_ARGS=""
if [[ "\${NNODES:-1}" != "1" ]]; then
DIST_ARGS="--nnodes \$NNODES --node-rank \$NODE_RANK --dist-init-addr \$DIST_ADDR:\${DIST_PORT:-29500}"
fi
exec python3 -m sglang.launch_server \
--model-path $MODEL_PATH --host 0.0.0.0 --port $DPORT \
$DECODE_COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \$DIST_ARGS \
--disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT
EOF
cat > "$WORKDIR/prefill.sh" <<EOF
#!/bin/bash
source "$WORKDIR/model_flags.sh"
source "$WORKDIR/ionic_mounts.sh"
NODE_RANK="\${1:-0}"; NNODES="\${2:-1}"; DIST_ADDR="\${3:-}"
docker rm -f mi355x_prefill 2>/dev/null || true
docker run $DOCKER_COMMON --name mi355x_prefill \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
-e NODE_RANK="\$NODE_RANK" -e NNODES="\$NNODES" -e DIST_ADDR="\$DIST_ADDR" -e DIST_PORT=${DIST_PORT} \
"\${IONIC_MOUNTS[@]}" \
$MORI_ENV$PENV_ARG $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
$IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/prefill_entry.sh
EOF
cat > "$WORKDIR/decode.sh" <<EOF
#!/bin/bash
source "$WORKDIR/model_flags.sh"
source "$WORKDIR/ionic_mounts.sh"
NODE_RANK="\${1:-0}"; NNODES="\${2:-1}"; DIST_ADDR="\${3:-}"
docker rm -f mi355x_decode 2>/dev/null || true
docker run $DOCKER_COMMON --name mi355x_decode \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
-e NODE_RANK="\$NODE_RANK" -e NNODES="\$NNODES" -e DIST_ADDR="\$DIST_ADDR" -e DIST_PORT=${DIST_PORT} \
"\${IONIC_MOUNTS[@]}" \
$MORI_ENV$DENV_ARG $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
$IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/decode_entry.sh
EOF
else
# Single-node-per-engine path (all EP<=8 recipes): byte-identical to the
# pre-wide launcher. No NODE_RANK/NNODES/DIST_* plumbing is emitted.
cat > "$WORKDIR/prefill_entry.sh" <<EOF
#!/bin/bash
set -euo pipefail
@@ -501,7 +715,7 @@ if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
fi
exec python3 -m sglang.launch_server \
--model-path $MODEL_PATH --host 0.0.0.0 --port $PPORT \
$COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
$PREFILL_COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
--disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT
EOF
@@ -516,7 +730,7 @@ if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
fi
exec python3 -m sglang.launch_server \
--model-path $MODEL_PATH --host 0.0.0.0 --port $DPORT \
$COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
$DECODE_COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
--disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT
EOF
@@ -526,7 +740,7 @@ source "$WORKDIR/model_flags.sh"
source "$WORKDIR/ionic_mounts.sh"
docker rm -f mi355x_prefill 2>/dev/null || true
docker run $DOCKER_COMMON "\${IONIC_MOUNTS[@]}" --name mi355x_prefill \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV$PENV_ARG $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
$IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/prefill_entry.sh
EOF
@@ -536,9 +750,10 @@ source "$WORKDIR/model_flags.sh"
source "$WORKDIR/ionic_mounts.sh"
docker rm -f mi355x_decode 2>/dev/null || true
docker run $DOCKER_COMMON "\${IONIC_MOUNTS[@]}" --name mi355x_decode \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV$DENV_ARG $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
$IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/decode_entry.sh
EOF
fi
# Probe payload + validator (separate files to avoid quoting inside the
# bench.sh `bash -lc '...'` block). One real request exercises the full
@@ -556,15 +771,39 @@ if not (t and t.strip()):
print("[probe] ok:", t[:80].replace("\n", " "))
PY
# Prefill health-wait + router --prefill args. For PW=1 this is a single endpoint
# (byte-identical to the pre-fan-out launcher); for PW>1 it iterates the comma-
# separated engine node0 IPs in PCSV ($3), health-waits each, and passes one
# --prefill per engine so the router fans requests across all prefill engines.
if (( PW > 1 )); then
PREFILL_WAIT_ROUTER=" IFS=',' read -ra PIPS <<< \"\${PCSV:-\$PIP}\"
PREFILL_ARGS=\"\"
for pip in \"\${PIPS[@]}\"; do
echo \"[wait] prefill \$pip\"; for i in \$(seq 1 600); do curl -sf http://\$pip:$PPORT/health >/dev/null && break; sleep 5; done
PREFILL_ARGS=\"\$PREFILL_ARGS --prefill http://\$pip:$PPORT $PBOOT\"
done
echo \"[wait] decode\"; for i in \$(seq 1 600); do curl -sf http://\$DIP:$DPORT/health >/dev/null && break; sleep 5; done
python3 -m sglang_router.launch_router \\
--pd-disaggregation \\
\$PREFILL_ARGS \\
--decode http://\$DIP:$DPORT \\
--host 0.0.0.0 --port $LBPORT \\
--disable-circuit-breaker &"
else
PREFILL_WAIT_ROUTER=" echo \"[wait] prefill\"; for i in \$(seq 1 600); do curl -sf http://\$PIP:$PPORT/health >/dev/null && break; sleep 5; done
echo \"[wait] decode\"; for i in \$(seq 1 600); do curl -sf http://\$DIP:$DPORT/health >/dev/null && break; sleep 5; done
python3 -m sglang_router.launch_router --pd-disaggregation --prefill http://\$PIP:$PPORT $PBOOT --decode http://\$DIP:$DPORT --host 0.0.0.0 --port $LBPORT --disable-circuit-breaker &"
fi
# Bench script runs on the prefill node; \$PIP/\$DIP injected at srun time.
cat > "$WORKDIR/bench.sh" <<EOF
#!/bin/bash
set -e
PIP=\$1; DIP=\$2
PIP=\$1; DIP=\$2; PCSV=\${3:-\$PIP}
source "$WORKDIR/ionic_mounts.sh"
docker rm -f mi355x_bench 2>/dev/null || true
docker run $DOCKER_COMMON "\${IONIC_MOUNTS[@]}" --name mi355x_bench \
-e PIP=\$PIP -e DIP=\$DIP \
-e PIP=\$PIP -e DIP=\$DIP -e PCSV=\$PCSV \
$IMAGE bash -lc '
CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
bash \$CIDIR/install_checkout_sglang.sh
@@ -574,14 +813,7 @@ docker run $DOCKER_COMMON "\${IONIC_MOUNTS[@]}" --name mi355x_bench \
export PYTHONPATH=/sgl-workspace/sglang/python:\${PYTHONPATH:-}
fi
bash \$CIDIR/install_checkout_router.sh
echo "[wait] prefill"; for i in \$(seq 1 600); do curl -sf http://\$PIP:$PPORT/health >/dev/null && break; sleep 5; done
echo "[wait] decode"; for i in \$(seq 1 600); do curl -sf http://\$DIP:$DPORT/health >/dev/null && break; sleep 5; done
python3 -m sglang_router.launch_router \
--pd-disaggregation \
--prefill http://\$PIP:$PPORT $PBOOT \
--decode http://\$DIP:$DPORT \
--host 0.0.0.0 --port $LBPORT \
--disable-circuit-breaker &
$PREFILL_WAIT_ROUTER
for i in \$(seq 1 30); do curl -sf http://127.0.0.1:$LBPORT/health >/dev/null && break; sleep 2; done
echo "[probe] PD end-to-end check via LB"
curl -sf -X POST http://127.0.0.1:$LBPORT/generate \
@@ -632,18 +864,63 @@ chmod +x "$WORKDIR"/*.sh
cat > "$WORKDIR/drive.sh" <<'DRIVE'
#!/bin/bash
set -x
WORKDIR="$1"; PW="${2:-1}"; DW="${3:-1}"
WORKDIR="$1"; PW="${2:-1}"; DW="${3:-1}"; PN_PER="${4:-1}"; DN_PER="${5:-1}"; DIST_NIC="${6:-}"
mapfile -t NODES < <(scontrol show hostnames "$SLURM_JOB_NODELIST")
PNODES=("${NODES[@]:0:PW}")
DNODES=("${NODES[@]:PW:DW}")
# Resolve a node's IP. When a cross-node dist NIC is named (wide engines set
# $DIST_NIC to the recipe's dist_socket_ifname), read the address off that
# interface on the node itself -- the multi-node dist-init-addr must be a
# LOCALLY-BINDABLE IP (torch-dist + tokenizer ZMQ bind to it), and a node's
# forward DNS can be stale/point at a non-local mgmt alias (observed: a decode
# root whose hostname resolved to an unpingable IP, ZMQ bind => "Cannot assign
# requested address"). Fall back to Slurm's NodeAddr, then DNS.
resolve_ip() {
local n="$1" ip="" addr="" field
if [[ -n "$DIST_NIC" ]]; then
ip=$(srun --overlap -N1 --nodelist="$n" ip -4 -o addr show "$DIST_NIC" 2>/dev/null \
| awk '{print $4}' | cut -d/ -f1 | head -1)
fi
if [[ -z "$ip" ]]; then
for field in $(scontrol show node "$n" -o); do
case "$field" in
NodeAddr=*) addr="${field#NodeAddr=}"; break ;;
esac
done
[[ -n "$addr" ]] || addr="$n"
read -r ip _ < <(getent ahostsv4 "$addr")
fi
printf '%s\n' "$ip"
}
# SLURM canonicalizes (sorts) SLURM_JOB_NODELIST, so the requested --nodelist
# order is already lost here. To keep a "slow-root" node out of any engine's
# rank0 slot -- e.g. mia1-p01-g20, whose MORI EP bootstrap reaches the connect
# stage ~16s late and loses the hardcoded 10s worker-connect race whenever it is
# rank0 -- push every node named in SLURM_DIST_TAIL to the END of the list. It
# then lands in a trailing decode-worker position, where its late ranks connect
# into the decode root's 30s accept window (harmless). Comma-separated; unmatched
# names ignored; empty (every EP<=8 recipe) => no-op, order unchanged.
if [[ -n "${SLURM_DIST_TAIL:-}" ]]; then
IFS=',' read -ra _TAIL <<< "$SLURM_DIST_TAIL"
_HEAD=(); _TL=()
for n in "${NODES[@]}"; do
_skip=0; for t in "${_TAIL[@]}"; do [[ "$n" == "$t" ]] && { _skip=1; break; }; done
if (( _skip )); then _TL+=("$n"); else _HEAD+=("$n"); fi
done
NODES=("${_HEAD[@]}" "${_TL[@]}")
echo "[drive] SLURM_DIST_TAIL=$SLURM_DIST_TAIL -> node order: ${NODES[*]}"
fi
# Each engine may span PN_PER/DN_PER nodes (ceil(TP/GPUs-per-node)); PN_PER=1 for
# EP<=8 so this reduces to the original one-node-per-worker split.
PN_TOTAL=$((PW * PN_PER)); DN_TOTAL=$((DW * DN_PER))
PNODES=("${NODES[@]:0:PN_TOTAL}")
DNODES=("${NODES[@]:PN_TOTAL:DN_TOTAL}")
PNODE="${PNODES[0]}"; DNODE="${DNODES[0]}"
PIP=$(getent ahostsv4 "$PNODE" | head -1 | awk '{print $1}')
DIP=$(getent ahostsv4 "$DNODE" | head -1 | awk '{print $1}')
PIP=$(resolve_ip "$PNODE")
DIP=$(resolve_ip "$DNODE")
echo "[drive] prefill nodes: ${PNODES[*]} ; decode nodes: ${DNODES[*]}"
echo "[drive] bench targets prefill=$PNODE($PIP) decode=$DNODE($DIP)"
if (( PW > 1 || DW > 1 )); then
echo "[drive] NOTE: router + bench use the first prefill and first decode only;"
echo "[drive] multi-prefill/multi-decode fan-out is not wired yet (LB work)."
if (( DW > 1 )); then
echo "[drive] NOTE: router + bench use the first decode engine only;"
echo "[drive] multi-decode fan-out is not wired yet (LB work)."
fi
# Each server's srun runs here on the login node and returns exactly when its
# compute-node container exits. Wrap it so the return code lands in a marker
@@ -652,19 +929,37 @@ fi
# died and with what code. (A hung-but-alive server is NOT caught here; that is
# bounded by bench.sh's health-wait timeout.)
rm -f "$WORKDIR"/server_exit_* "$WORKDIR/bench_exit"
for n in "${PNODES[@]}"; do
( srun --overlap -N1 --nodelist="$n" bash "$WORKDIR/prefill.sh" > "$WORKDIR/prefill_$n.log" 2>&1
# Launch PW prefill engines; each spans PN_PER nodes as its own torch-dist group
# (engine node0 = dist-init addr; node_rank = position WITHIN the engine, so an
# engine's ranks are 0..PN_PER-1, not a global index). PN_PER=1 => NNODES=1 in
# the entry script => dist args dropped => byte-identical single-node launch.
# Collect each engine's node0 IP into PCSV for the router's prefill fan-out.
PCSV=""
for ((k=0; k<PW; k++)); do
e0="${PNODES[k*PN_PER]}"
eip=$(resolve_ip "$e0")
PCSV="${PCSV:+$PCSV,}$eip"
for ((j=0; j<PN_PER; j++)); do
n="${PNODES[k*PN_PER + j]}"
( srun --overlap -N1 --nodelist="$n" bash "$WORKDIR/prefill.sh" "$j" "$PN_PER" "$eip" > "$WORKDIR/prefill_$n.log" 2>&1
echo "prefill@$n rc=$?" > "$WORKDIR/server_exit_prefill_$n" ) &
done
for n in "${DNODES[@]}"; do
( srun --overlap -N1 --nodelist="$n" bash "$WORKDIR/decode.sh" > "$WORKDIR/decode_$n.log" 2>&1
done
for ((k=0; k<DW; k++)); do
e0="${DNODES[k*DN_PER]}"
eip=$(resolve_ip "$e0")
for ((j=0; j<DN_PER; j++)); do
n="${DNODES[k*DN_PER + j]}"
( srun --overlap -N1 --nodelist="$n" bash "$WORKDIR/decode.sh" "$j" "$DN_PER" "$eip" > "$WORKDIR/decode_$n.log" 2>&1
echo "decode@$n rc=$?" > "$WORKDIR/server_exit_decode_$n" ) &
done
done
echo "[drive] prefill engine endpoints (fan-out): $PCSV"
sleep 5
# Bench in the background with its own marker, so the wait loop is purely file
# based: finish when bench writes its marker, abort if any server marker shows up
# first (a server died before the sweep completed).
( srun --overlap -N1 --nodelist="$PNODE" bash "$WORKDIR/bench.sh" "$PIP" "$DIP" > "$WORKDIR/bench.log" 2>&1
( srun --overlap -N1 --nodelist="$PNODE" bash "$WORKDIR/bench.sh" "$PIP" "$DIP" "$PCSV" > "$WORKDIR/bench.log" 2>&1
echo $? > "$WORKDIR/bench_exit" ) &
BENCH_BG=$!
# Stream bench output live and poll the markers with xtrace OFF, so the console
@@ -704,14 +999,19 @@ NODELIST_ARG=()
EXCLUSIVE_ARG=()
[[ "${SLURM_EXCLUSIVE:-1}" == "1" ]] && EXCLUSIVE_ARG=(--exclusive)
# Keep the scheduler off known-bad nodes (e.g. a host whose ionic RDMA driver
# ABI mismatches the container, where MORI reports "no active RDMA device" and
# the disagg server dies on init). Comma-separated node list.
# Optional comma-separated nodes to keep the scheduler off.
EXCLUDE_ARG=()
[[ -n "${SLURM_EXCLUDE:-}" ]] && EXCLUDE_ARG=(--exclude="$SLURM_EXCLUDE")
# One node per prefill/decode worker (TP == GPUs/node). 1P1D -> 2 nodes.
TOTAL_NODES=$((PW + DW))
# Nodes = sum over engines of nodes-per-engine. EP<=8 (PN_PER=DN_PER=1) gives the
# original PW+DW (1P1D -> 2 nodes); wide EP16 1P1D gives 2+2 = 4 nodes.
TOTAL_NODES=$(( PW * PN_PER + DW * DN_PER ))
# Keep g20 out of a wide engine's root position, where its slower MORI
# bootstrap can miss the worker-connect timeout.
if [[ "$MATRIX_CONFIG_NAME" == *-2p1d-ep16* ]]; then
export SLURM_DIST_TAIL="${SLURM_DIST_TAIL:-mia1-p01-g20}"
fi
# Name the allocation <RUNNER_NAME>-<GITHUB_RUN_ID>-<config> so the workflow's
# cleanup steps can scancel precisely instead of a blanket `squeue --me` that
@@ -722,7 +1022,7 @@ JOB_NAME="mi355x-ci-${RUNNER_NAME:-norunner}-${GITHUB_RUN_ID:-0}-${MATRIX_CONFIG
set +e
salloc -p "$SLURM_PARTITION" -N"$TOTAL_NODES" "${NODELIST_ARG[@]}" "${EXCLUDE_ARG[@]}" "${EXCLUSIVE_ARG[@]}" \
--job-name "$JOB_NAME" -t "$TIME_LIMIT" \
bash "$WORKDIR/drive.sh" "$WORKDIR" "$PW" "$DW"
bash "$WORKDIR/drive.sh" "$WORKDIR" "$PW" "$DW" "$PN_PER" "$DN_PER" "$DIST_SOCK"
SALLOC_RC=$?
set -e
@@ -769,12 +1069,18 @@ fi
# median_e2e_latency_ms -> median_e2el_ms (E2E latency; /1000 -> s)
# (none; injected here) -> model_id (served model, from $MODEL_PATH)
# ---------------------------------------------------------------------------
TOTAL_GPUS=$((PTP + DTP))
# GPU counts are per-ROLE totals across all engines: PW prefill engines of PTP
# GPUs each, DW decode engines of DTP each. For EP<=8 (PW=DW=1) this is PTP/DTP
# exactly, so the filename fields are unchanged. process_result.py reads the
# _ctx_/_gen_ fields as PREFILL_GPUS/DECODE_GPUS for its per-GPU throughput math,
# so they must be the role totals (Oren EP16: ctx=2*8=16, gen=1*16=16, gpus=32).
PREFILL_GPUS_TOTAL=$((PW * PTP)); DECODE_GPUS_TOTAL=$((DW * DTP))
TOTAL_GPUS=$((PREFILL_GPUS_TOTAL + DECODE_GPUS_TOTAL))
PROCESSED=0
for C in ${CONCS//,/ }; do
RAW="$WORKDIR/raw_conc${C}.json"
[[ -f "$RAW" ]] || { echo "WARN: missing $RAW"; continue; }
DEST="$GITHUB_WORKSPACE/${RESULT_FILENAME}_${MATRIX_CONFIG_NAME}_conc${C}_gpus_${TOTAL_GPUS}_ctx_${PTP}_gen_${DTP}.json"
DEST="$GITHUB_WORKSPACE/${RESULT_FILENAME}_${MATRIX_CONFIG_NAME}_conc${C}_gpus_${TOTAL_GPUS}_ctx_${PREFILL_GPUS_TOTAL}_gen_${DECODE_GPUS_TOTAL}.json"
MODEL_ID="$MODEL_PATH" python3 - "$RAW" "$DEST" "$C" <<'PY'
import json, os, sys
raw_path, dest, conc = sys.argv[1], sys.argv[2], int(sys.argv[3])
+33
View File
@@ -356,3 +356,36 @@ kimik26-fp8-mi355x-mtp-sglang:
search-space:
- conc-list: [1, 8, 16, 32, 64, 128, 256]
config_file: scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-mtp.yaml
# Kimi-K2.6 MXFP4 wide-EP16 2P1D: aiter MoE path, needs only the wide-EP launcher, not #32048.
kimik26-mxfp4-mi355x-ep16-sglang:
model: amd/Kimi-K2.6-MXFP4
model-prefix: kimik26
model_path: /it-share/model_coverage/models--amd--Kimi-K2.6-MXFP4
runner: mi355x
precision: mxfp4
framework: sglang
multinode: true
disagg: true
seq-len-configs:
- isl: 1024
osl: 1024
search-space:
- conc-list: [1, 8, 16, 32, 64, 128, 256]
config_file: scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/2p1d-ep16-mxfp4.yaml
kimik26-mxfp4-mi355x-ep16-mtp-sglang:
model: amd/Kimi-K2.6-MXFP4
model-prefix: kimik26
model_path: /it-share/model_coverage/models--amd--Kimi-K2.6-MXFP4
runner: mi355x
precision: mxfp4
framework: sglang
multinode: true
disagg: true
seq-len-configs:
- isl: 1024
osl: 1024
search-space:
- conc-list: [1, 8, 16, 32, 64, 128, 256]
config_file: scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/2p1d-ep16-mtp-mxfp4.yaml
@@ -0,0 +1,122 @@
# MI355X Kimi-K2.6 (MXFP4 experts, FP8 KV) 4-node 2P1D disaggregation recipe (MTP) — narrow-prefill EP8
# + wide-decode EP16 (mirrors the DSV4-Pro Oren config: wide EP only helps decode).
#
# Two prefill engines (EP8, one node each; the router fans requests across both) +
# one decode engine (EP16) spanning 2 nodes. Still one logical P/D pair per role
# group, 4 nodes total. nodes-per-engine = ceil(TP/8): prefill 8->1, decode 16->2,
# so the launcher emits cross-node --nnodes/--node-rank/--dist-init-addr for the
# decode engine only. Prefill EP8 keeps MoE all-to-all INTRA-node (XGMI); decode
# gets wide EP16 across nodes. KV (prefill TP8 -> decode TP16) is carried over mori.
#
# Kimi-specific bits vs the DSV4-Pro EP16 recipe: split attention backends
# (aiter prefill / triton decode), the Kimi model env + parsers, and the Kimi
# model path. Everything else (2P1D topology, mori a2a + KV, dist init) is shared.
#
# Consumed by:
# * scripts/ci/slurm/process_result.py reads `resources` and
# `backend.sglang_config` (TP/EP/DP + worker counts) for the summary table.
# * scripts/ci/slurm/launch_mi355x.sh reads `runtime`, `bench`, `model`, `mtp`.
resources:
prefill_workers: 2
decode_workers: 1
backend:
sglang_config:
prefill:
tensor-parallel-size: 8
expert-parallel-size: 8
data-parallel-size: 8
decode:
tensor-parallel-size: 16
expert-parallel-size: 16
data-parallel-size: 16
# Model-specific docker env + sglang server args (written verbatim via
# model_flags.sh). Each server arg + value is a SEPARATE list item.
model:
env:
SGLANG_USE_AITER: 1
SGLANG_ROCM_FUSED_DECODE_MLA: 0
server_args:
- --model-loader-extra-config
- '{"enable_multithread_load": true}'
- --reasoning-parser
- kimi_k2
- --tool-call-parser
- kimi_k2
runtime:
image: lmsysorg/sglang-rocm:v0.5.15.post1-rocm720-mi35x-20260715
# Kimi uses split attention backends (aiter prefill / triton decode), not a
# single --attention-backend.
prefill_attention_backend: aiter
decode_attention_backend: triton
# RoCE HCAs (8/node) for mori MoE all-to-all AND the P->D KV transfer.
ib_devices: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
# Wide-EP MoE all-to-all backend (cross-node expert dispatch/combine).
moe_a2a_backend: mori
# KV P->D transfer backend (mori for both a2a and KV on this cluster).
kv_transfer_backend: mori
# Cross-node torch-distributed NIC for the wide decode engine's dist init.
dist_socket_ifname: eno0
# rocm720 0715 image needs the ROCm-7.0.0-alpha path OFF (validated).
rocm700a: 0
prefill_port: 30025
decode_port: 30026
prefill_bootstrap_port: 8998
decode_bootstrap_port: 9001
lb_port: 8000
# Base defaults; the wide_ep block overrides mem-fraction / max-req per role.
mem_fraction_static: 0.90
page_size: 256
max_running_requests: 256
chunked_prefill_size: 131072
swa_full_tokens_ratio: 0.1
# Per-role wide-EP tuning (starting point cloned from the validated DSV4-Pro
# EP16 run; may need Kimi-specific retuning). Prefill and decode differ.
wide_ep:
kv_cache_dtype: fp8_e4m3
prefill_mem_fraction_static: 0.8
decode_mem_fraction_static: 0.85
prefill_chunked_prefill_size: 131072
prefill_max_running_requests: 1024
decode_max_running_requests: 1024
common_extra_flags: "--moe-dense-tp-size 1 --enable-dp-lm-head --decode-log-interval 100 --watchdog-timeout 3600 --load-balance-method round_robin"
prefill_extra_flags: "--context-length 9217 --max-total-tokens 262144"
decode_extra_flags: "--disable-cuda-graph --prefill-round-robin-balance"
prefill_extra_env:
MORI_MAX_DISPATCH_TOKENS_PREFILL: 8192
MORI_MAX_DISPATCH_TOKENS_DECODE: 256
SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK: 16384
decode_extra_env:
MORI_MAX_DISPATCH_TOKENS_DECODE: 64
MORI_MOE_MAX_INPUT_TOKENS_DECODE: 332
SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK: 128
# EAGLE3 speculative decoding with the external Kimi-K2.6 draft model (same
# draft + hyperparams as the validated EP8 1p1d-mtp.yaml). Decode stays eager
# (--disable-cuda-graph); MXFP4 experts run the aiter MoE path (no int4 route-1
# dense-compaction), so cuda-graph could be revisited once the leg is green.
mtp:
enabled: true
algorithm: EAGLE3
num_steps: 3
eagle_topk: 1
num_draft_tokens: 4
draft_model_path: /it-share/model_coverage/models--lightseekorg--kimi-k2.6-eagle3.1-mla
bench:
# bench_serving --max-concurrency sweep; one result JSON per concurrency.
concurrencies: [1, 8, 16, 32, 64, 128, 256]
num_prompts_factor: 4 # num-prompts = concurrency * factor
random_range_ratio: 1.0
# Correctness gate run through the PD path before the perf sweep (full GSM8K,
# 8-shot, accuracy > 0.92). Mirrors the single-node Kimi-K2.6 eval threshold.
accuracy:
enabled: true
num_shots: 8
num_questions: 1319 # full GSM8K test set
threshold: 0.92
@@ -0,0 +1,109 @@
# MI355X Kimi-K2.6 (MXFP4 experts, FP8 KV) 4-node 2P1D disaggregation recipe — narrow-prefill EP8
# + wide-decode EP16 (mirrors the DSV4-Pro Oren config: wide EP only helps decode).
#
# Two prefill engines (EP8, one node each; the router fans requests across both) +
# one decode engine (EP16) spanning 2 nodes. Still one logical P/D pair per role
# group, 4 nodes total. nodes-per-engine = ceil(TP/8): prefill 8->1, decode 16->2,
# so the launcher emits cross-node --nnodes/--node-rank/--dist-init-addr for the
# decode engine only. Prefill EP8 keeps MoE all-to-all INTRA-node (XGMI); decode
# gets wide EP16 across nodes. KV (prefill TP8 -> decode TP16) is carried over mori.
#
# Kimi-specific bits vs the DSV4-Pro EP16 recipe: split attention backends
# (aiter prefill / triton decode), the Kimi model env + parsers, and the Kimi
# model path. Everything else (2P1D topology, mori a2a + KV, dist init) is shared.
#
# Consumed by:
# * scripts/ci/slurm/process_result.py reads `resources` and
# `backend.sglang_config` (TP/EP/DP + worker counts) for the summary table.
# * scripts/ci/slurm/launch_mi355x.sh reads `runtime`, `bench`, `model`, `mtp`.
resources:
prefill_workers: 2
decode_workers: 1
backend:
sglang_config:
prefill:
tensor-parallel-size: 8
expert-parallel-size: 8
data-parallel-size: 8
decode:
tensor-parallel-size: 16
expert-parallel-size: 16
data-parallel-size: 16
# Model-specific docker env + sglang server args (written verbatim via
# model_flags.sh). Each server arg + value is a SEPARATE list item.
model:
env:
SGLANG_USE_AITER: 1
SGLANG_ROCM_FUSED_DECODE_MLA: 0
server_args:
- --model-loader-extra-config
- '{"enable_multithread_load": true}'
- --reasoning-parser
- kimi_k2
- --tool-call-parser
- kimi_k2
runtime:
image: lmsysorg/sglang-rocm:v0.5.15.post1-rocm720-mi35x-20260715
# Kimi uses split attention backends (aiter prefill / triton decode), not a
# single --attention-backend.
prefill_attention_backend: aiter
decode_attention_backend: triton
# RoCE HCAs (8/node) for mori MoE all-to-all AND the P->D KV transfer.
ib_devices: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
# Wide-EP MoE all-to-all backend (cross-node expert dispatch/combine).
moe_a2a_backend: mori
# KV P->D transfer backend (mori for both a2a and KV on this cluster).
kv_transfer_backend: mori
# Cross-node torch-distributed NIC for the wide decode engine's dist init.
dist_socket_ifname: eno0
# rocm720 0715 image needs the ROCm-7.0.0-alpha path OFF (validated).
rocm700a: 0
prefill_port: 30025
decode_port: 30026
prefill_bootstrap_port: 8998
decode_bootstrap_port: 9001
lb_port: 8000
# Base defaults; the wide_ep block overrides mem-fraction / max-req per role.
mem_fraction_static: 0.90
page_size: 256
max_running_requests: 256
chunked_prefill_size: 131072
swa_full_tokens_ratio: 0.1
# Per-role wide-EP tuning (starting point cloned from the validated DSV4-Pro
# EP16 run; may need Kimi-specific retuning). Prefill and decode differ.
wide_ep:
kv_cache_dtype: fp8_e4m3
prefill_mem_fraction_static: 0.8
decode_mem_fraction_static: 0.85
prefill_chunked_prefill_size: 131072
prefill_max_running_requests: 1024
decode_max_running_requests: 1024
common_extra_flags: "--moe-dense-tp-size 1 --enable-dp-lm-head --decode-log-interval 100 --watchdog-timeout 3600 --load-balance-method round_robin"
prefill_extra_flags: "--context-length 9217 --max-total-tokens 262144"
decode_extra_flags: "--disable-cuda-graph --prefill-round-robin-balance"
prefill_extra_env:
MORI_MAX_DISPATCH_TOKENS_PREFILL: 8192
MORI_MAX_DISPATCH_TOKENS_DECODE: 256
SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK: 16384
decode_extra_env:
MORI_MAX_DISPATCH_TOKENS_DECODE: 64
MORI_MOE_MAX_INPUT_TOKENS_DECODE: 332
SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK: 128
bench:
# bench_serving --max-concurrency sweep; one result JSON per concurrency.
concurrencies: [1, 8, 16, 32, 64, 128, 256]
num_prompts_factor: 4 # num-prompts = concurrency * factor
random_range_ratio: 1.0
# Correctness gate run through the PD path before the perf sweep (full GSM8K,
# 8-shot, accuracy > 0.92). Mirrors the single-node Kimi-K2.6 eval threshold.
accuracy:
enabled: true
num_shots: 8
num_questions: 1319 # full GSM8K test set
threshold: 0.92