minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD][DI][CI] 3/N Add Kimi K2.6 FP8 MI355X 1P1D nightly recipes (#29855)

Browse Source 
Co-authored-by: bingxche <bingxche@amd.com>
This commit is contained in:
Zhaoyi Li
2026-07-04 21:05:12 -07:00
committed by GitHub
co-authored by bingxche
parent a37bc2456d
commit 67361ff91b
4 changed files with 303 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/ci/slurm/launch_mi355x.sh
+107 -22
View File
@@ -59,19 +59,26 @@ if [[ -z "$MODEL_PATH" ]]; then
fi
# Resolve a HuggingFace cache dir (models--org--name) to its live snapshot dir.
# Lets nightly-configs point at the shared cache without hardcoding a snapshot
# hash; if MODEL_PATH is already a concrete snapshot (or plain dir), use as-is.
if [[ -f "$MODEL_PATH/refs/main" && -d "$MODEL_PATH/snapshots" ]]; then
SNAP_HASH="$(cat "$MODEL_PATH/refs/main")"
RESOLVED="$MODEL_PATH/snapshots/$SNAP_HASH"
if [[ -d "$RESOLVED" ]]; then
echo "resolved snapshot: $MODEL_PATH -> $RESOLVED"
MODEL_PATH="$RESOLVED"
else
echo "ERROR: refs/main=$SNAP_HASH but $RESOLVED missing" >&2
exit 1
# Lets nightly-configs / recipes point at the shared cache without hardcoding a
# snapshot hash; a concrete snapshot dir (or plain dir) is returned unchanged.
# Used for both MODEL_PATH and an optional speculative draft model path.
resolve_snapshot() {
local p="$1"
if [[ -f "$p/refs/main" && -d "$p/snapshots" ]]; then
local hash resolved
hash="$(cat "$p/refs/main")"
resolved="$p/snapshots/$hash"
if [[ -n "$hash" && -d "$resolved" ]]; then
echo "resolved snapshot: $p -> $resolved" >&2
echo "$resolved"
return 0
fi
echo "ERROR: refs/main=$hash but $resolved missing" >&2
return 1
fi
fi
echo "$p"
}
MODEL_PATH="$(resolve_snapshot "$MODEL_PATH")" || exit 1
# ---------------------------------------------------------------------------
# Parse the recipe (runtime + bench + topology) into shell vars.
@@ -89,7 +96,12 @@ rt = r["runtime"]; b = r["backend"]["sglang_config"]; bn = r["bench"]
res = r.get("resources", {})
def emit(k, v): print(f"{k}={v}")
emit("IMAGE", rt["image"])
emit("ATTN", rt["attention_backend"])
# Attention backend: single (`attention_backend`) or split
# (`prefill_attention_backend`/`decode_attention_backend`). Empty when absent so
# the flag is dropped for a model that omits it.
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"])
emit("PPORT", rt["prefill_port"])
emit("DPORT", rt["decode_port"])
@@ -100,16 +112,23 @@ emit("MEMFRAC", rt["mem_fraction_static"])
emit("PAGE", rt["page_size"])
emit("MAXREQ", rt["max_running_requests"])
emit("CHUNK", rt["chunked_prefill_size"])
emit("SWA", rt["swa_full_tokens_ratio"])
# swa is DSV4-specific; emit empty when a model omits it so the flag is dropped.
emit("SWA", rt.get("swa_full_tokens_ratio", ""))
# 1 when the recipe carries a `model:` block (env + server_args written to
# model_flags.sh); 0 for the DSV4 recipes, which keep the hardcoded DSV4 path.
emit("HAS_MODEL", 1 if r.get("model") else 0)
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))
m = r.get("mtp", {}) or {}
emit("MTP_ENABLED", 1 if m.get("enabled") else 0)
emit("MTP_ALGO", m.get("algorithm", "EAGLE"))
emit("MTP_STEPS", m.get("num_steps", 3))
emit("MTP_TOPK", m.get("eagle_topk", 1))
emit("MTP_DRAFT", m.get("num_draft_tokens", 4))
# External draft checkpoint (EAGLE3 etc.); empty for DSV4's built-in EAGLE head.
emit("MTP_DRAFT_PATH", m.get("draft_model_path", ""))
# Worker counts double as node counts here: one server per node (TP == GPUs/node).
# 1P1D today; bumping these reserves 2P2D / 1P3D / 3P1D. Multi-node-per-worker
# (TP > GPUs/node, needs --dist-init-addr/--nnodes/--node-rank) is out of scope.
@@ -134,7 +153,7 @@ eval "$RECIPE_VARS"
if [[ -n "${IMAGE_OVERRIDE:-}" ]]; then
IMAGE="$IMAGE_OVERRIDE"
fi
echo "recipe: image=$IMAGE attn=$ATTN ib=$IB ptp=$PTP dtp=$DTP concs=$CONCS isl=$ISL osl=$OSL"
echo "recipe: image=$IMAGE attn=${ATTN:-$PATTN/$DATTN} ib=$IB ptp=$PTP dtp=$DTP concs=$CONCS isl=$ISL osl=$OSL"
# ---------------------------------------------------------------------------
# Shared NFS scratch (visible to login node + compute nodes). Raw bench output
@@ -182,26 +201,82 @@ DSV4_ENV=(
-e AITER_BF16_FP8_MOE_BOUND=0 -e SGLANG_DSV4_FP4_EXPERTS=$FP4_EXPERTS
)
DSV4_ENV_STR="${DSV4_ENV[*]}"
# A recipe carrying a `model:` block supplies its OWN docker env (below), so the
# 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"
# 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
# prefill.sh/decode.sh). DSV4 recipes have no `model:` block -> empty arrays, so
# their generated docker argv is unchanged. Each server arg + its value MUST be a
# separate YAML list item so shlex.quote keeps "--foo" and "bar" as two tokens.
python3 - "$CONFIG_FILE" "$WORKDIR/model_flags.sh" <<'PY'
import shlex, sys, yaml
r = yaml.safe_load(open(sys.argv[1]))
model = r.get("model", {}) or {}
env = model.get("env", {}) or {}
server_args = model.get("server_args", []) or []
# YAML true/false parse to Python bool; render lowercase so env values stay
# byte-identical to shell (`=false`, not `=False`) -- SGLang parsing is
# case-sensitive for some of these.
def fmt(v):
if isinstance(v, bool):
return "true" if v else "false"
return str(v)
env_args = []
for k, v in env.items():
env_args += ["-e", f"{k}={fmt(v)}"]
def q(items):
return " ".join(shlex.quote(fmt(x)) for x in items)
with open(sys.argv[2], "w") as f:
f.write(f"MODEL_ENV_ARGS=({q(env_args)})\n")
f.write(f"MODEL_SERVER_ARGS=({q(server_args)})\n")
PY
# 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"
if [[ "$MTP_ENABLED" == "1" ]]; then
EXTRA_FLAGS="$EXTRA_FLAGS --speculative-algorithm EAGLE \
EXTRA_FLAGS="$EXTRA_FLAGS --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"
fi
fi
echo "extra flags: ${EXTRA_FLAGS:-<none>} (pep=$PEP pdp=$PDP mtp=$MTP_ENABLED)"
echo "extra flags: ${EXTRA_FLAGS:-<none>} (pep=$PEP pdp=$PDP mtp=$MTP_ENABLED algo=$MTP_ALGO)"
COMMON_FLAGS="--trust-remote-code --tp $PTP --disable-radix-cache \
if [[ "$HAS_MODEL" == "1" ]]; then
# Generic path (e.g. Kimi): attention + swa from the recipe, model parsers /
# quirks ride MODEL_SERVER_ARGS. Single `--attention-backend` when the recipe
# sets `attention_backend`; split `--prefill-/--decode-attention-backend` when
# it sets the per-role keys. swa dropped when the recipe omits it.
ATTN_FLAGS=""
[[ -n "$ATTN" ]] && ATTN_FLAGS="$ATTN_FLAGS --attention-backend $ATTN"
[[ -n "$PATTN" ]] && ATTN_FLAGS="$ATTN_FLAGS --prefill-attention-backend $PATTN"
[[ -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 \
--chunked-prefill-size $CHUNK \
--disaggregation-transfer-backend mori --disaggregation-ib-device $IB$EXTRA_FLAGS"
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 \
--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"
fi
DOCKER_COMMON="--rm --network host --ipc host --shm-size 32g --privileged \
--security-opt seccomp=unconfined \
@@ -211,24 +286,34 @@ DOCKER_COMMON="--rm --network host --ipc host --shm-size 32g --privileged \
# ---------------------------------------------------------------------------
# Write per-role scripts that srun dispatches to each compute node.
# ---------------------------------------------------------------------------
# These are UNQUOTED `<<EOF` heredocs, so $MORI_ENV/$DSV4_ENV_STR/$COMMON_FLAGS/
# $IMAGE/$MODEL_PATH expand now (at generation). model_flags.sh is sourced at
# runtime for the model's env/server arrays, so the `${MODEL_ENV_ARGS[@]}` /
# `${MODEL_SERVER_ARGS[@]}` refs are backslash-escaped to survive into the script
# and expand after `source`. For DSV4 those arrays are empty and $DSV4_ENV_STR is
# set, so the resulting docker argv is byte-identical to the pre-Kimi launcher.
cat > "$WORKDIR/prefill.sh" <<EOF
#!/bin/bash
source "$WORKDIR/model_flags.sh"
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 $MORI_ENV $DSV4_ENV_STR \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
$IMAGE python3 -m sglang.launch_server \
--model-path $MODEL_PATH --host 0.0.0.0 --port $PPORT \
$COMMON_FLAGS --disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT
$COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
--disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT
EOF
cat > "$WORKDIR/decode.sh" <<EOF
#!/bin/bash
source "$WORKDIR/model_flags.sh"
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 $MORI_ENV $DSV4_ENV_STR \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
$IMAGE python3 -m sglang.launch_server \
--model-path $MODEL_PATH --host 0.0.0.0 --port $DPORT \
$COMMON_FLAGS --disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT
$COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
--disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT
EOF
# Probe payload + validator (separate files to avoid quoting inside the
scripts/ci/slurm/nightly-configs.yaml
+36
View File
@@ -320,3 +320,39 @@ dsv4pro-fp4-mi355x-dp8ep8-mtp-sglang:
search-space:
- conc-list: [1, 8, 16, 32, 64, 128, 256]
config_file: scripts/ci/slurm/recipes/mi355x-fp4/dsv4pro/1k1k/1p1d-dp8ep8-mtp.yaml
# Kimi-K2.6 (FP8) 2-node 1P1D. Demonstrates the launcher is model-agnostic: all
# Kimi-specific config lives in the recipe's `model:` block + split attention
# runtime, with nothing hardcoded in launch_mi355x.sh. Base + EAGLE3 MTP (the
# MTP leg uses an external draft checkpoint via mtp.draft_model_path).
kimik26-fp8-mi355x-sglang:
model: moonshotai/Kimi-K2.6
model-prefix: kimik26
model_path: /it-share/model_coverage/models--moonshotai--Kimi-K2.6
runner: mi355x
precision: fp8
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/1p1d.yaml
kimik26-fp8-mi355x-mtp-sglang:
model: moonshotai/Kimi-K2.6
model-prefix: kimik26
model_path: /it-share/model_coverage/models--moonshotai--Kimi-K2.6
runner: mi355x
precision: fp8
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/1p1d-mtp.yaml
scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-mtp.yaml
+85
View File
@@ -0,0 +1,85 @@
# MI355X Kimi-K2.6 (FP8) 2-node 1P1D disaggregation recipe (base + EAGLE3 MTP).
#
# Self-contained (no inheritance): same as 1p1d.yaml plus the `mtp:` block. Uses
# EAGLE3 speculative decoding with an EXTERNAL draft checkpoint (unlike DSV4's
# built-in EAGLE NextN head): the launcher resolves mtp.draft_model_path through
# the HF-cache snapshot logic and appends --speculative-draft-model-path. The
# draft dir must live under /it-share (the container's :ro mount). MTP is applied
# to both prefill and decode. Mirrors the Kimi-K2.6 serving cookbook.
#
# 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: 1
decode_workers: 1
backend:
sglang_config:
prefill:
tensor-parallel-size: 8
expert-parallel-size: 1
data-parallel-size: 1
decode:
tensor-parallel-size: 8
expert-parallel-size: 1
data-parallel-size: 1
# 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}'
- --watchdog-timeout
- 1200
- --reasoning-parser
- kimi_k2
- --tool-call-parser
- kimi_k2
# EAGLE3 speculative decoding with an external draft checkpoint.
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
runtime:
image: lmsysorg/sglang-rocm:v0.5.13.post1-rocm720-mi35x-20260623
# Kimi uses split attention backends (aiter prefill / triton decode), not a
# single --attention-backend.
prefill_attention_backend: aiter
decode_attention_backend: triton
# RoCE HCAs MORI uses for cross-node KV transfer.
ib_devices: rdma0,rdma1,rdma2,rdma3
prefill_port: 30025
decode_port: 30026
prefill_bootstrap_port: 8998
decode_bootstrap_port: 9001
lb_port: 8000
mem_fraction_static: 0.90
page_size: 256
max_running_requests: 256
chunked_prefill_size: 8192
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. Mirrors the
# registered single-node Kimi-K2.6 eval (full GSM8K, 8-shot, accuracy > 0.92).
accuracy:
enabled: true
num_shots: 8
num_questions: 1319 # full GSM8K test set
threshold: 0.92
scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d.yaml
+75
View File
@@ -0,0 +1,75 @@
# MI355X Kimi-K2.6 (FP8) 2-node 1P1D disaggregation recipe (base).
#
# All Kimi-specific config lives in this recipe's `model:` block (docker env +
# sglang server args) and `runtime` (split attention backends); nothing about
# Kimi is hardcoded in launch_mi355x.sh. Mirrors the single-node registered test
# test/registered/amd/accuracy/mi35x/test_kimi_k26_eval_mi35x.py (TP8, split
# attention backends, multithread loader, GSM8K > 0.92).
#
# 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: 1
decode_workers: 1
backend:
sglang_config:
prefill:
tensor-parallel-size: 8
expert-parallel-size: 1
data-parallel-size: 1
decode:
tensor-parallel-size: 8
expert-parallel-size: 1
data-parallel-size: 1
# 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}'
- --watchdog-timeout
- 1200
- --reasoning-parser
- kimi_k2
- --tool-call-parser
- kimi_k2
runtime:
image: lmsysorg/sglang-rocm:v0.5.13.post1-rocm720-mi35x-20260623
# Kimi uses split attention backends (aiter prefill / triton decode), not a
# single --attention-backend.
prefill_attention_backend: aiter
decode_attention_backend: triton
# RoCE HCAs MORI uses for cross-node KV transfer.
ib_devices: rdma0,rdma1,rdma2,rdma3
prefill_port: 30025
decode_port: 30026
prefill_bootstrap_port: 8998
decode_bootstrap_port: 9001
lb_port: 8000
mem_fraction_static: 0.90
page_size: 256
max_running_requests: 256
chunked_prefill_size: 8192
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. Mirrors the
# registered single-node Kimi-K2.6 eval (full GSM8K, 8-shot, accuracy > 0.92).
accuracy:
enabled: true
num_shots: 8
num_questions: 1319 # full GSM8K test set
threshold: 0.92