From a84e70eb1e66c9cf548f3c7d48256708ba39429e Mon Sep 17 00:00:00 2001 From: Zhaoyi Li <36555117+Lzy17@users.noreply.github.com> Date: Mon, 3 Aug 2026 20:51:55 -0500 Subject: [PATCH] [AMD][DI][CI] 6/N Add Kimi-K2.6 MXFP4 wide-EP16 2P1D nightly recipes (#33333) --- scripts/ci/slurm/launch_mi355x.sh | 432 +++++++++++++++--- scripts/ci/slurm/nightly-configs.yaml | 33 ++ .../kimik26/1k1k/2p1d-ep16-mtp-mxfp4.yaml | 122 +++++ .../kimik26/1k1k/2p1d-ep16-mxfp4.yaml | 109 +++++ 4 files changed, 633 insertions(+), 63 deletions(-) create mode 100644 scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/2p1d-ep16-mtp-mxfp4.yaml create mode 100644 scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/2p1d-ep16-mxfp4.yaml diff --git a/scripts/ci/slurm/launch_mi355x.sh b/scripts/ci/slurm/launch_mi355x.sh index 494cf9073..83cb5e903 100755 --- a/scripts/ci/slurm/launch_mi355x.sh +++ b/scripts/ci/slurm/launch_mi355x.sh @@ -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< 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 --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:-} 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:-} (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:-} | decode tail:${DECODE_TAIL:-} (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,7 +630,81 @@ print(f"[checkout-router] Router={Router}") PY EOF -cat > "$WORKDIR/prefill_entry.sh" < 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" < "$WORKDIR/decode_entry.sh" < "$WORKDIR/prefill.sh" </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" </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" < "$WORKDIR/decode_entry.sh" < "$WORKDIR/decode_entry.sh" < "$WORKDIR/prefill.sh" < "$WORKDIR/prefill.sh" </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 -cat > "$WORKDIR/decode.sh" < "$WORKDIR/decode.sh" </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" </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 - echo "prefill@$n rc=$?" > "$WORKDIR/server_exit_prefill_$n" ) & +# 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 "$WORKDIR/prefill_$n.log" 2>&1 + echo "prefill@$n rc=$?" > "$WORKDIR/server_exit_prefill_$n" ) & + done done -for n in "${DNODES[@]}"; do - ( srun --overlap -N1 --nodelist="$n" bash "$WORKDIR/decode.sh" > "$WORKDIR/decode_$n.log" 2>&1 - echo "decode@$n rc=$?" > "$WORKDIR/server_exit_decode_$n" ) & +for ((k=0; k "$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 -- 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]) diff --git a/scripts/ci/slurm/nightly-configs.yaml b/scripts/ci/slurm/nightly-configs.yaml index 80b5f4b6c..071e86348 100644 --- a/scripts/ci/slurm/nightly-configs.yaml +++ b/scripts/ci/slurm/nightly-configs.yaml @@ -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 diff --git a/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/2p1d-ep16-mtp-mxfp4.yaml b/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/2p1d-ep16-mtp-mxfp4.yaml new file mode 100644 index 000000000..96d19d371 --- /dev/null +++ b/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 diff --git a/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/2p1d-ep16-mxfp4.yaml b/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/2p1d-ep16-mxfp4.yaml new file mode 100644 index 000000000..d3f73abe0 --- /dev/null +++ b/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/2p1d-ep16-mxfp4.yaml @@ -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