From 54c2e31945a35e47f3ecdd2f6248b372975b92d4 Mon Sep 17 00:00:00 2001 From: bingxche Date: Mon, 6 Jul 2026 08:36:21 -0500 Subject: [PATCH 1/8] ci: run MI355X disagg with checkout sglang --- .../workflows/nightly-amd-mi355x-disagg.yml | 2 + scripts/ci/slurm/launch_mi355x.sh | 144 ++++++++++++++++-- 2 files changed, 135 insertions(+), 11 deletions(-) diff --git a/.github/workflows/nightly-amd-mi355x-disagg.yml b/.github/workflows/nightly-amd-mi355x-disagg.yml index d16e140bfc1a..63bbcfa8504c 100644 --- a/.github/workflows/nightly-amd-mi355x-disagg.yml +++ b/.github/workflows/nightly-amd-mi355x-disagg.yml @@ -118,6 +118,8 @@ jobs: steps: - name: Checkout code uses: actions/checkout@v4 + with: + fetch-depth: 0 - name: Clean up prior Slurm jobs from this runner continue-on-error: true diff --git a/scripts/ci/slurm/launch_mi355x.sh b/scripts/ci/slurm/launch_mi355x.sh index 0635b71b442f..89cf86ed1a67 100755 --- a/scripts/ci/slurm/launch_mi355x.sh +++ b/scripts/ci/slurm/launch_mi355x.sh @@ -24,6 +24,11 @@ # SLURM_NODELIST - optional explicit node pin (else scheduler chooses) # SLURM_EXCLUDE - optional comma-separated nodes to keep the scheduler # off (e.g. hosts with a broken RDMA driver) +# SGLANG_USE_CHECKOUT_RUNTIME +# - default 1. Reinstall this workflow checkout's Python +# sglang package inside each runtime container before +# launching servers/bench. Set 0 to use the image's +# baked-in sglang package. # RUNNER_NAME - GitHub runner name (a built-in default env var) # GITHUB_RUN_ID - GitHub Actions run id (a built-in default env var) # The allocation is named @@ -52,6 +57,11 @@ set -x SLURM_PARTITION="${SLURM_PARTITION:-amd-sglang}" TIME_LIMIT="${TIME_LIMIT:-02:30:00}" MODEL_PATH="${MODEL_PATH:-${MODEL:-}}" +SGLANG_USE_CHECKOUT_RUNTIME="${SGLANG_USE_CHECKOUT_RUNTIME:-1}" +case "${SGLANG_USE_CHECKOUT_RUNTIME,,}" in + 0|false|no|off) SGLANG_USE_CHECKOUT_RUNTIME=0 ;; + *) SGLANG_USE_CHECKOUT_RUNTIME=1 ;; +esac if [[ -z "$MODEL_PATH" ]]; then echo "ERROR: set MODEL_PATH (local snapshot) or MODEL" >&2 @@ -163,6 +173,23 @@ WORKDIR="$HOME/.mi355x_ci/${MATRIX_CONFIG_NAME}" rm -rf "$WORKDIR"; mkdir -p "$WORKDIR" SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# Stage the workflow checkout on shared NFS so Slurm compute-node containers can +# reinstall the same code SHA the workflow checked out. The container gets a +# read-only mount and copies it to /tmp before mutating pyproject.toml. +CHECKOUT_DOCKER_ARGS="-e SGLANG_USE_CHECKOUT_RUNTIME=$SGLANG_USE_CHECKOUT_RUNTIME" +if [[ "$SGLANG_USE_CHECKOUT_RUNTIME" == "1" ]]; then + CHECKOUT_STAGE="$WORKDIR/checkout" + CHECKOUT_SHA="$(git -C "$GITHUB_WORKSPACE" rev-parse HEAD)" + echo "Staging checkout runtime: sha=$CHECKOUT_SHA -> $CHECKOUT_STAGE" + rm -rf "$CHECKOUT_STAGE" + mkdir -p "$CHECKOUT_STAGE" + tar --exclude='__pycache__' --exclude='*.pyc' \ + -C "$GITHUB_WORKSPACE" -cf - . | tar -C "$CHECKOUT_STAGE" -xf - + CHECKOUT_DOCKER_ARGS="$CHECKOUT_DOCKER_ARGS -e SGLANG_CHECKOUT_SHA=$CHECKOUT_SHA -v $CHECKOUT_STAGE:/sglang-checkout:ro" +else + echo "SGLANG_USE_CHECKOUT_RUNTIME=0; using sglang package baked into image." +fi + # Accuracy-gate helpers (written when enabled). Pre-stage the GSM8K test set on # shared NFS from the login node (which has internet) so the in-container eval # doesn't depend on compute-node connectivity; fall back to in-container @@ -281,7 +308,7 @@ 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" +-v /it-share:/it-share:ro -v $HOME:/host_home $CHECKOUT_DOCKER_ARGS" # --------------------------------------------------------------------------- # Write per-role scripts that srun dispatches to each compute node. @@ -292,16 +319,109 @@ DOCKER_COMMON="--rm --network host --ipc host --shm-size 32g --privileged \ # `${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/install_checkout_sglang.sh" <<'EOF' +#!/bin/bash +set -euo pipefail + +case "${SGLANG_USE_CHECKOUT_RUNTIME:-1}" in + 0|false|False|FALSE|no|No|NO|off|Off|OFF) + echo "[checkout-sglang] disabled; using image-baked sglang" + exit 0 + ;; +esac + +CHECKOUT_SRC="${CHECKOUT_SRC:-/sglang-checkout}" +RUNTIME_CHECKOUT="${RUNTIME_CHECKOUT:-/tmp/sglang-checkout-runtime}" + +if [[ ! -f "$CHECKOUT_SRC/python/sglang/version.py" ]]; then + echo "[checkout-sglang] ERROR: invalid checkout mount: $CHECKOUT_SRC" >&2 + exit 1 +fi + +echo "[checkout-sglang] reinstalling sglang from $CHECKOUT_SRC" +rm -rf "$RUNTIME_CHECKOUT" +mkdir -p "$RUNTIME_CHECKOUT" +tar --exclude='__pycache__' --exclude='*.pyc' \ + -C "$CHECKOUT_SRC" -cf - . | tar -C "$RUNTIME_CHECKOUT" -xf - + +git config --global --add safe.directory "$RUNTIME_CHECKOUT" || true + +# The ROCm pyproject variant is the one used by AMD CI. Mutate only the private +# /tmp copy so prefill/decode/bench never race on the read-only checkout mount. +rm -f "$RUNTIME_CHECKOUT/python/pyproject.toml" +cp "$RUNTIME_CHECKOUT/python/pyproject_other.toml" "$RUNTIME_CHECKOUT/python/pyproject.toml" +for f in README.md LICENSE; do + if [[ -f "$RUNTIME_CHECKOUT/$f" && ! -e "$RUNTIME_CHECKOUT/python/$f" ]]; then + cp "$RUNTIME_CHECKOUT/$f" "$RUNTIME_CHECKOUT/python/$f" + fi +done + +python3 -m pip uninstall -y sglang || true +python3 -m pip install --no-deps --no-build-isolation -e "$RUNTIME_CHECKOUT/python" + +export RUNTIME_CHECKOUT +export PYTHONPATH="$RUNTIME_CHECKOUT/python:${PYTHONPATH:-}" +python3 - <<'PY' +import importlib.metadata +import os +import subprocess +import sglang + +checkout = os.environ["RUNTIME_CHECKOUT"] +expected = os.path.realpath(os.path.join(checkout, "python", "sglang")) + os.sep +actual = os.path.realpath(os.path.dirname(sglang.__file__)) + os.sep +try: + sha = subprocess.check_output( + ["git", "-C", checkout, "rev-parse", "HEAD"], text=True + ).strip() +except Exception: + sha = os.environ.get("SGLANG_CHECKOUT_SHA", "unknown") + +print(f"[checkout-sglang] sha={sha}") +print(f"[checkout-sglang] sglang_file={sglang.__file__}") +print(f"[checkout-sglang] sglang_version={importlib.metadata.version('sglang')}") +if not actual.startswith(expected): + raise SystemExit(f"sglang did not import from checkout: {sglang.__file__}") +PY +EOF + +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 $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 "\${MODEL_SERVER_ARGS[@]}" \ - --disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT + $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 $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 "\${MODEL_SERVER_ARGS[@]}" \ - --disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT + $IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/decode_entry.sh EOF # Probe payload + validator (separate files to avoid quoting inside the @@ -341,7 +458,13 @@ docker rm -f mi355x_bench 2>/dev/null || true docker run $DOCKER_COMMON --name mi355x_bench \ -e PIP=\$PIP -e DIP=\$DIP \ $IMAGE bash -lc ' - export PYTHONPATH=/sgl-workspace/sglang/python:\$PYTHONPATH + CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME} + bash \$CIDIR/install_checkout_sglang.sh + if [ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]; then + export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-} + else + export PYTHONPATH=/sgl-workspace/sglang/python:\${PYTHONPATH:-} + fi 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 \ @@ -351,7 +474,6 @@ docker run $DOCKER_COMMON --name mi355x_bench \ --host 0.0.0.0 --port $LBPORT \ --disable-circuit-breaker & for i in \$(seq 1 30); do curl -sf http://127.0.0.1:$LBPORT/health >/dev/null && break; sleep 2; done - CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME} echo "[probe] PD end-to-end check via LB" curl -sf -X POST http://127.0.0.1:$LBPORT/generate \ -H "content-type: application/json" -d @\$CIDIR/probe.json > \$CIDIR/probe_out.json \ From 371ef5e9b703edeaa12be2f6bb1c06c399e3f2df Mon Sep 17 00:00:00 2001 From: bingxche Date: Mon, 6 Jul 2026 09:02:54 -0500 Subject: [PATCH 2/8] ci: clean MI355X disagg scratch after run --- .github/workflows/nightly-amd-mi355x-disagg.yml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/.github/workflows/nightly-amd-mi355x-disagg.yml b/.github/workflows/nightly-amd-mi355x-disagg.yml index 63bbcfa8504c..85423ad9cd00 100644 --- a/.github/workflows/nightly-amd-mi355x-disagg.yml +++ b/.github/workflows/nightly-amd-mi355x-disagg.yml @@ -219,6 +219,16 @@ jobs: scancel $ACTIVE_JOBS fi + - name: Clean up MI355X scratch + if: always() + continue-on-error: true + run: | + LOG_DIR="$HOME/.mi355x_ci/${MATRIX_CONFIG_NAME}" + if [ -d "$LOG_DIR" ]; then + echo "Removing MI355X scratch directory: $LOG_DIR" + rm -rf "$LOG_DIR" + fi + collect-results: needs: nightly-mi355x-benchmark if: github.repository == 'sgl-project/sglang' && always() From 1e0890e73933fd543089697dcea95a9d027e9fd3 Mon Sep 17 00:00:00 2001 From: Bingxu Chen Date: Tue, 7 Jul 2026 15:33:05 +0800 Subject: [PATCH 3/8] [AMD] Cap DSV4 Flash max_total_num_tokens (#30313) Co-authored-by: YC Yen-Ching Tseng --- scripts/ci/slurm/launch_mi355x.sh | 2 ++ .../recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml | 1 + .../ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8.yaml | 1 + .../ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-mtp.yaml | 1 + scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d.yaml | 1 + .../recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml | 1 + .../ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8.yaml | 1 + .../ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-mtp.yaml | 1 + scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d.yaml | 1 + 9 files changed, 10 insertions(+) diff --git a/scripts/ci/slurm/launch_mi355x.sh b/scripts/ci/slurm/launch_mi355x.sh index 89cf86ed1a67..4e5274e4d99b 100755 --- a/scripts/ci/slurm/launch_mi355x.sh +++ b/scripts/ci/slurm/launch_mi355x.sh @@ -121,6 +121,7 @@ emit("LBPORT", rt["lb_port"]) emit("MEMFRAC", rt["mem_fraction_static"]) emit("PAGE", rt["page_size"]) emit("MAXREQ", rt["max_running_requests"]) +emit("MAXTOK", rt.get("max_total_tokens", "")) emit("CHUNK", rt["chunked_prefill_size"]) # swa is DSV4-specific; emit empty when a model omits it so the flag is dropped. emit("SWA", rt.get("swa_full_tokens_ratio", "")) @@ -266,6 +267,7 @@ PY 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" if [[ "$MTP_ENABLED" == "1" ]]; then EXTRA_FLAGS="$EXTRA_FLAGS --speculative-algorithm $MTP_ALGO \ --speculative-num-steps $MTP_STEPS --speculative-eagle-topk $MTP_TOPK \ diff --git a/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml b/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml index da710c6aa9b9..5c52f17a1a4d 100644 --- a/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml +++ b/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml @@ -33,6 +33,7 @@ runtime: mem_fraction_static: 0.90 page_size: 256 max_running_requests: 256 + max_total_tokens: 8551168 chunked_prefill_size: 8192 swa_full_tokens_ratio: 0.1 diff --git a/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8.yaml b/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8.yaml index 01a12bfec289..02678af98dfe 100644 --- a/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8.yaml +++ b/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-dp8ep8.yaml @@ -33,6 +33,7 @@ runtime: mem_fraction_static: 0.90 page_size: 256 max_running_requests: 256 + max_total_tokens: 8551168 chunked_prefill_size: 8192 swa_full_tokens_ratio: 0.1 diff --git a/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-mtp.yaml b/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-mtp.yaml index ec7c8a93cdff..25c2ad8994fb 100644 --- a/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-mtp.yaml +++ b/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d-mtp.yaml @@ -33,6 +33,7 @@ runtime: mem_fraction_static: 0.90 page_size: 256 max_running_requests: 256 + max_total_tokens: 8551168 chunked_prefill_size: 8192 swa_full_tokens_ratio: 0.1 diff --git a/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d.yaml b/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d.yaml index 3116ce6fdc0d..7da8a1d6401a 100644 --- a/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d.yaml +++ b/scripts/ci/slurm/recipes/mi355x-fp4/dsv4flash/1k1k/1p1d.yaml @@ -35,6 +35,7 @@ runtime: mem_fraction_static: 0.90 page_size: 256 max_running_requests: 256 + max_total_tokens: 8551168 chunked_prefill_size: 8192 swa_full_tokens_ratio: 0.1 diff --git a/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml b/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml index 826b1782a8d1..7a3623e31131 100644 --- a/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml +++ b/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8-mtp.yaml @@ -33,6 +33,7 @@ runtime: mem_fraction_static: 0.90 page_size: 256 max_running_requests: 256 + max_total_tokens: 8551168 chunked_prefill_size: 8192 swa_full_tokens_ratio: 0.1 diff --git a/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8.yaml b/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8.yaml index 82b52e8df75c..f20abe3ef099 100644 --- a/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8.yaml +++ b/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-dp8ep8.yaml @@ -33,6 +33,7 @@ runtime: mem_fraction_static: 0.90 page_size: 256 max_running_requests: 256 + max_total_tokens: 8551168 chunked_prefill_size: 8192 swa_full_tokens_ratio: 0.1 diff --git a/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-mtp.yaml b/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-mtp.yaml index 5546a1fcb958..89901377067c 100644 --- a/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-mtp.yaml +++ b/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d-mtp.yaml @@ -33,6 +33,7 @@ runtime: mem_fraction_static: 0.90 page_size: 256 max_running_requests: 256 + max_total_tokens: 8551168 chunked_prefill_size: 8192 swa_full_tokens_ratio: 0.1 diff --git a/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d.yaml b/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d.yaml index 8539748aed8c..6e35d15270e7 100644 --- a/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d.yaml +++ b/scripts/ci/slurm/recipes/mi355x-fp8/dsv4flash/1k1k/1p1d.yaml @@ -33,6 +33,7 @@ runtime: mem_fraction_static: 0.90 page_size: 256 max_running_requests: 256 + max_total_tokens: 8551168 chunked_prefill_size: 8192 swa_full_tokens_ratio: 0.1 From 0aac9766b11e96584141d74c454d2361d6899ace Mon Sep 17 00:00:00 2001 From: bingxche Date: Tue, 7 Jul 2026 06:17:22 -0500 Subject: [PATCH 4/8] ci: run MI355X disagg router from checkout --- .../workflows/nightly-amd-mi355x-disagg.yml | 1 + scripts/ci/slurm/launch_mi355x.sh | 84 ++++++++++++++++++- 2 files changed, 81 insertions(+), 4 deletions(-) diff --git a/.github/workflows/nightly-amd-mi355x-disagg.yml b/.github/workflows/nightly-amd-mi355x-disagg.yml index 85423ad9cd00..5b3898ea76c7 100644 --- a/.github/workflows/nightly-amd-mi355x-disagg.yml +++ b/.github/workflows/nightly-amd-mi355x-disagg.yml @@ -120,6 +120,7 @@ jobs: uses: actions/checkout@v4 with: fetch-depth: 0 + persist-credentials: false - name: Clean up prior Slurm jobs from this runner continue-on-error: true diff --git a/scripts/ci/slurm/launch_mi355x.sh b/scripts/ci/slurm/launch_mi355x.sh index 4e5274e4d99b..5297f3852be0 100755 --- a/scripts/ci/slurm/launch_mi355x.sh +++ b/scripts/ci/slurm/launch_mi355x.sh @@ -26,9 +26,10 @@ # off (e.g. hosts with a broken RDMA driver) # SGLANG_USE_CHECKOUT_RUNTIME # - default 1. Reinstall this workflow checkout's Python -# sglang package inside each runtime container before -# launching servers/bench. Set 0 to use the image's -# baked-in sglang package. +# sglang package inside each runtime container, and the +# checkout sglang-router package inside the bench +# container, before launching servers/bench. Set 0 to +# use the image's baked-in packages. # RUNNER_NAME - GitHub runner name (a built-in default env var) # GITHUB_RUN_ID - GitHub Actions run id (a built-in default env var) # The allocation is named @@ -184,7 +185,7 @@ if [[ "$SGLANG_USE_CHECKOUT_RUNTIME" == "1" ]]; then echo "Staging checkout runtime: sha=$CHECKOUT_SHA -> $CHECKOUT_STAGE" rm -rf "$CHECKOUT_STAGE" mkdir -p "$CHECKOUT_STAGE" - tar --exclude='__pycache__' --exclude='*.pyc' \ + tar --exclude='__pycache__' --exclude='*.pyc' --exclude='.git/config' \ -C "$GITHUB_WORKSPACE" -cf - . | tar -C "$CHECKOUT_STAGE" -xf - CHECKOUT_DOCKER_ARGS="$CHECKOUT_DOCKER_ARGS -e SGLANG_CHECKOUT_SHA=$CHECKOUT_SHA -v $CHECKOUT_STAGE:/sglang-checkout:ro" else @@ -387,6 +388,80 @@ if not actual.startswith(expected): PY EOF +cat > "$WORKDIR/install_checkout_router.sh" <<'EOF' +#!/bin/bash +set -euo pipefail + +case "${SGLANG_USE_CHECKOUT_RUNTIME:-1}" in + 0|false|False|FALSE|no|No|NO|off|Off|OFF) + echo "[checkout-router] disabled; using image-baked sglang-router" + python3 - <<'PY' || true +import importlib.metadata +import sglang_router + +print(f"[checkout-router] sglang_router_file={sglang_router.__file__}") +print( + "[checkout-router] sglang_router_version=" + f"{importlib.metadata.version('sglang-router')}" +) +try: + import sglang_router.sglang_router_rs as rs + + print(f"[checkout-router] sglang_router_rs_file={rs.__file__}") +except Exception as exc: + print(f"[checkout-router] sglang_router_rs_import_error={exc}") +PY + exit 0 + ;; +esac + +RUNTIME_CHECKOUT="${RUNTIME_CHECKOUT:-/tmp/sglang-checkout-runtime}" +ROUTER_SRC="$RUNTIME_CHECKOUT/sgl-model-gateway/bindings/python" +WHEEL_DIR="${SGLANG_ROUTER_WHEEL_DIR:-/tmp/sglang-router-wheels}" + +if [[ ! -f "$ROUTER_SRC/pyproject.toml" ]]; then + echo "[checkout-router] ERROR: invalid router checkout: $ROUTER_SRC" >&2 + exit 1 +fi + +echo "[checkout-router] building sglang-router from $ROUTER_SRC" +export CARGO_BUILD_JOBS="${CARGO_BUILD_JOBS:-4}" +python3 -m maturin --version >/dev/null 2>&1 \ + || python3 -m pip install --no-cache-dir "maturin<1.14" + +# Match the ROCm image build recipe when compiling from the checkout copy. +if [[ -f "$RUNTIME_CHECKOUT/sgl-model-gateway/Cargo.toml" ]]; then + sed -i -E 's|^(smg-[a-zA-Z-]+)\s*=\s*"~1\.0\.0"|\1 = "=1.0.0"|' \ + "$RUNTIME_CHECKOUT/sgl-model-gateway/Cargo.toml" +fi + +rm -rf "$WHEEL_DIR" +mkdir -p "$WHEEL_DIR" +( + cd "$ROUTER_SRC" + ulimit -n 65536 || true + python3 -m maturin build --release --features vendored-openssl --out "$WHEEL_DIR" +) + +python3 -m pip uninstall -y sglang-router || true +python3 -m pip install --force-reinstall --no-deps "$WHEEL_DIR"/*.whl + +python3 - <<'PY' +import importlib.metadata +import sglang_router +import sglang_router.sglang_router_rs as rs +from sglang_router.sglang_router_rs import Router + +print(f"[checkout-router] sglang_router_file={sglang_router.__file__}") +print( + "[checkout-router] sglang_router_version=" + f"{importlib.metadata.version('sglang-router')}" +) +print(f"[checkout-router] sglang_router_rs_file={rs.__file__}") +print(f"[checkout-router] Router={Router}") +PY +EOF + cat > "$WORKDIR/prefill_entry.sh" </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 \ From 595125d91e598702892394d355d183b3969ccb4d Mon Sep 17 00:00:00 2001 From: yctseng0211 Date: Wed, 8 Jul 2026 11:02:53 -0500 Subject: [PATCH 5/8] [AMD][DI][CI] decode-metadata probe + recipe for Kimi disagg non-MTP GSM8K drop --- .../debug_utils/disagg_decode_meta_probe.py | 240 ++++++++++++++++++ .../srt/layers/attention/aiter_backend.py | 6 + .../srt/layers/attention/triton_backend.py | 6 + scripts/ci/slurm/nightly-configs.yaml | 19 ++ .../kimik26/1k1k/1p1d-metadump.yaml | 88 +++++++ 5 files changed, 359 insertions(+) create mode 100644 python/sglang/srt/debug_utils/disagg_decode_meta_probe.py create mode 100644 scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-metadump.yaml diff --git a/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py b/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py new file mode 100644 index 000000000000..488c2bdf62cc --- /dev/null +++ b/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py @@ -0,0 +1,240 @@ +"""Env-gated, read-only probe that dumps DECODE-mode attention metadata. + +Purpose +------- +Localize the Kimi-K2.6 2N-1P1D disagg non-MTP GSM8K drop (~0.88 vs 0.944 +single-node). Prior analysis narrowed the cause to the *decode-mode* read of +the MORI-transferred prefix: both the triton and the aiter decode kernels are +fed the SAME scheduler-built metadata (``kv_indptr`` / ``kv_indices`` / +``seq_lens`` / ``num_kv_splits``) derived from ``req_to_token``, and that is +their only shared input. The verify/extend path uses a different metadata set +(``qo_indptr`` / ``custom_mask``), which is why MTP verify over the same KV is +correct. This probe dumps exactly that shared decode metadata plus the KV-pool +row norms at the attended slots, on REAL decode steps (not warmup), so a fixed +prompt can be diffed single-node vs disagg to find the first divergence. + +Enable +------ + SGLANG_DEBUG_DISAGG_DECODE_META=1 + +Optional knobs: + SGLANG_DEBUG_DISAGG_DECODE_META_LAYERS=0 comma list of layer_ids to dump + SGLANG_DEBUG_DISAGG_DECODE_META_STEPS=8 max decode forwards to dump + SGLANG_DEBUG_DISAGG_DECODE_META_MINLEN=16 skip reqs with seq_len below this + (filters out the tiny warmup probe) + SGLANG_DEBUG_DISAGG_DECODE_META_KVNORM=1 also dump KV-pool row norms at the + attended slots (default on) + +How to read the diff +-------------------- +Run the SAME fixed prompt (bs=1) on single-node and on disagg, then compare: + * seq_len at each step -> must match; a mismatch is an off-by-one in the + disagg decode seq_len bookkeeping. + * KV-norm head/tail sequence for the prefix -> same prompt must give the same + norms regardless of physical slots; a mismatch means the transferred KV + values are wrong. + * zero_rows -> must be 0; >0 means the attended prefix slots were never + written (transfer/placement problem, on a real request this time). + * idx==rtt -> sanity; kv_indices is a copy of req_to_token[req, :seq_len]. + +The probe is wrapped in try/except and skips CUDA/HIP graph capture, so it can +never break a run. +""" + +from __future__ import annotations + +import logging +import os +import threading + +import torch + +logger = logging.getLogger(__name__) + + +def _env_flag(name: str, default: str = "0") -> bool: + return os.environ.get(name, default).lower() not in ("0", "", "false", "no") + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, str(default))) + except (TypeError, ValueError): + return default + + +_ENABLED = _env_flag("SGLANG_DEBUG_DISAGG_DECODE_META") +_MAX_STEPS = _env_int("SGLANG_DEBUG_DISAGG_DECODE_META_STEPS", 8) +_MIN_SEQ_LEN = _env_int("SGLANG_DEBUG_DISAGG_DECODE_META_MINLEN", 16) +_KVNORM = _env_flag("SGLANG_DEBUG_DISAGG_DECODE_META_KVNORM", "1") + + +def _parse_layers() -> set: + raw = os.environ.get("SGLANG_DEBUG_DISAGG_DECODE_META_LAYERS", "0") + out = set() + for tok in raw.split(","): + tok = tok.strip() + if tok: + try: + out.add(int(tok)) + except ValueError: + pass + return out or {0} + + +_LAYERS = _parse_layers() +_FIRST_LAYER = min(_LAYERS) + +_lock = threading.Lock() +_step = 0 + + +def _ht(values, k: int = 6) -> str: + """Compact head/tail view of a 1-D sequence.""" + vals = list(values) + if len(vals) <= 2 * k: + return str(vals) + return f"{vals[:k]}...{vals[-k:]}" + + +def maybe_dump_decode_meta(tag: str, backend, layer, forward_batch) -> None: + """Dump decode metadata for the first few real decode forwards. + + ``tag`` distinguishes the backend ("triton" / "aiter"). ``backend`` is the + attention backend instance (must expose ``forward_metadata``, + ``req_to_token_pool`` and ``token_to_kv_pool``). Safe no-op unless the + ``SGLANG_DEBUG_DISAGG_DECODE_META`` env var is set. + """ + if not _ENABLED: + return + try: + _dump_impl(tag, backend, layer, forward_batch) + except Exception as exc: # never break the model forward + logger.warning("[DDM] probe failed (%s): %r", tag, exc) + + +def _dump_impl(tag: str, backend, layer, forward_batch) -> None: + global _step + + # Never run during CUDA/HIP graph capture: the .cpu()/.item() syncs below + # are illegal while tracing a graph. + try: + from sglang.srt.model_executor.runner import get_is_capture_mode + + if get_is_capture_mode(): + return + except Exception: + pass + + layer_id = getattr(layer, "layer_id", 0) + if layer_id not in _LAYERS: + return + + fwd_mode = forward_batch.forward_mode + is_decode = fwd_mode.is_decode() or fwd_mode.is_idle() + is_verify = fwd_mode.is_target_verify() + if not (is_decode or is_verify): + return + + seq_lens = getattr(forward_batch, "seq_lens_cpu", None) + if seq_lens is None: + seq_lens = forward_batch.seq_lens.detach().to("cpu") + seq_lens = seq_lens.tolist() + bs = len(seq_lens) + + # Skip the tiny warmup probe ("The capital of France is", seq_len ~6-9). + if max(seq_lens, default=0) < _MIN_SEQ_LEN: + return + + # One global step per qualifying decode forward; keyed on the first dumped + # layer so all target layers of the same forward share a step id. + with _lock: + if layer_id == _FIRST_LAYER: + _step += 1 + cur_step = _step + if cur_step > _MAX_STEPS: + return + + fmeta = getattr(backend, "forward_metadata", None) + kv_indptr = getattr(fmeta, "kv_indptr", None) + kv_indices = getattr(fmeta, "kv_indices", None) + num_kv_splits = getattr(fmeta, "num_kv_splits", None) + + req_pool_indices = forward_batch.req_pool_indices.detach().to("cpu").tolist() + req_to_token = backend.req_to_token_pool.req_to_token + + kv_indptr_cpu = ( + kv_indptr[: bs + 1].detach().to("cpu").tolist() + if kv_indptr is not None + else None + ) + + mode_name = "VERIFY" if is_verify else "DECODE" + lines = [] + for b in range(bs): + seq_len = int(seq_lens[b]) + if seq_len < _MIN_SEQ_LEN: + continue + rp = int(req_pool_indices[b]) + + # req_to_token slot mapping for this request's tokens. + rtt = req_to_token[rp, :seq_len].detach().to("cpu") + + # kv_indices slice for this request (CSR). + idx_slice = None + n_kv = None + idx_str = "n/a" + match_str = "n/a" + if kv_indptr_cpu is not None and kv_indices is not None: + lo, hi = int(kv_indptr_cpu[b]), int(kv_indptr_cpu[b + 1]) + n_kv = hi - lo + idx_slice = kv_indices[lo:hi].detach().to("cpu") + idx_str = _ht(idx_slice.tolist()) + if idx_slice.numel() == rtt.numel(): + match_str = str(bool(torch.equal(idx_slice.to(rtt.dtype), rtt))) + else: + match_str = f"LEN_MISMATCH({idx_slice.numel()} vs {rtt.numel()})" + + nsplit_str = "n/a" + if num_kv_splits is not None: + try: + nsplit_str = str(int(num_kv_splits[b])) + except Exception: + nsplit_str = "?" + + block = ( + f"[DDM step={cur_step} tag={tag} L{layer_id} mode={mode_name}] " + f"rp={rp} seq_len={seq_len} kv_indptr=[{kv_indptr_cpu[b] if kv_indptr_cpu else '?'}," + f"{kv_indptr_cpu[b + 1] if kv_indptr_cpu else '?'}] n_kv={n_kv} " + f"nsplit={nsplit_str} idx==rtt:{match_str}\n" + f" kv_idx = {idx_str}\n" + f" rtt = {_ht(rtt.tolist())}" + ) + + # KV-pool row norms at the attended slots (real decode step, unlike the + # warmup-only #30433 dump). Same prompt => same norms across runs. + if _KVNORM and layer_id == _FIRST_LAYER: + slots = idx_slice if idx_slice is not None else rtt + block += "\n" + _kv_norm_report(backend, layer_id, slots) + + lines.append(block) + + if lines: + logger.info("\n".join(lines)) + + +def _kv_norm_report(backend, layer_id: int, slots: torch.Tensor) -> str: + try: + pool = backend.token_to_kv_pool + kbuf = pool.get_key_buffer(layer_id) + dev_slots = slots.to(device=kbuf.device, dtype=torch.long) + rows = kbuf[dev_slots].reshape(dev_slots.shape[0], -1).float() + norms = rows.norm(dim=-1) + zero_rows = int((norms < 1e-6).sum().item()) + norms_list = [round(x, 3) for x in norms.detach().to("cpu").tolist()] + return ( + f" KV L{layer_id}: zero_rows={zero_rows}/{norms.numel()} " + f"norm={_ht(norms_list)}" + ) + except Exception as exc: + return f" KV L{layer_id}: norm probe failed: {exc!r}" diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index c3dea8ab355f..0830dafd4606 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -2458,6 +2458,12 @@ def forward_decode( num_kv_splits = self.forward_metadata.num_kv_splits + from sglang.srt.debug_utils.disagg_decode_meta_probe import ( + maybe_dump_decode_meta, + ) + + maybe_dump_decode_meta("aiter", self, layer, forward_batch) + o = self._mla_decode_fwd_with_head_pad( q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), k_buffer.view(-1, 1, 1, layer.qk_head_dim), diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 0be715ac1de9..4a3ad5879121 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -1751,6 +1751,12 @@ def forward_decode( o = cp_lse_ag_out_rs_mha(o_for_decode, local_lse, group) return o.reshape(-1, layer.tp_q_head_num * layer.v_head_dim).to(q.dtype) + from sglang.srt.debug_utils.disagg_decode_meta_probe import ( + maybe_dump_decode_meta, + ) + + maybe_dump_decode_meta("triton", self, layer, forward_batch) + self.decode_attention_fwd( q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), self.token_to_kv_pool.get_key_buffer(layer.layer_id), diff --git a/scripts/ci/slurm/nightly-configs.yaml b/scripts/ci/slurm/nightly-configs.yaml index 80b5f4b6ce30..ac943c3a0fc8 100644 --- a/scripts/ci/slurm/nightly-configs.yaml +++ b/scripts/ci/slurm/nightly-configs.yaml @@ -356,3 +356,22 @@ 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 + +# DIAGNOSTIC (DO NOT MERGE): non-MTP base + decode-metadata probe + eager decode. +# Not part of the default nightly signal in spirit; run explicitly via +# workflow_dispatch `configs=kimik26-fp8-1k1k-1p1d-metadump`. See the recipe. +kimik26-fp8-mi355x-metadump-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] + config_file: scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-metadump.yaml diff --git a/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-metadump.yaml b/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-metadump.yaml new file mode 100644 index 000000000000..8e172d13f956 --- /dev/null +++ b/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-metadump.yaml @@ -0,0 +1,88 @@ +# MI355X Kimi-K2.6 (FP8) 2N 1P1D — DECODE-metadata probe variant (DO NOT MERGE). +# +# Same as 1p1d.yaml (non-MTP base) but wired to dump the decode-mode attention +# metadata for the MORI-transferred prefix, to localize the non-MTP GSM8K drop +# (~0.88 vs single-node 0.944). Isolates ONE thing vs the base: +# * model.env += SGLANG_DEBUG_DISAGG_DECODE_META (+MINLEN/STEPS) -> the probe +# in {triton,aiter} forward_decode logs kv_indptr / kv_indices / seq_lens / +# num_kv_splits and the KV-pool row norms at the attended slots. +# * server_args += --disable-cuda-graph -> decode runs EAGER so forward_decode +# is entered in Python (with HIP graph, decode is replayed and the probe +# never fires). +# MINLEN=128 skips the model warmup AND the short "capital of France" PD probe, +# so the dump budget lands on real GSM8K decode steps (long transferred prefix). +# Perf sweep trimmed to conc=1; only the GSM8K gate matters (it aborts at ~0.88). +# +# Trigger (config name = kimik26-fp8-1k1k-1p1d-metadump): +# gh workflow run "Nightly Test (AMD MI355X 2N 1P1D Disagg)" \ +# --ref amd/kimik26-disagg-decodemeta \ +# -f configs=kimik26-fp8-1k1k-1p1d-metadump + +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: + env: + SGLANG_USE_AITER: 1 + SGLANG_ROCM_FUSED_DECODE_MLA: 0 + # Decode-metadata probe (see debug_utils/disagg_decode_meta_probe.py). + SGLANG_DEBUG_DISAGG_DECODE_META: 1 + SGLANG_DEBUG_DISAGG_DECODE_META_MINLEN: 128 + SGLANG_DEBUG_DISAGG_DECODE_META_STEPS: 16 + SGLANG_DEBUG_DISAGG_DECODE_META_KVNORM: 1 + server_args: + - --model-loader-extra-config + - '{"enable_multithread_load": true}' + - --watchdog-timeout + - 1200 + - --reasoning-parser + - kimi_k2 + - --tool-call-parser + - kimi_k2 + # Force eager decode so forward_decode is entered in Python (probe fires). + - --disable-cuda-graph + +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: + # Diagnostic: only the GSM8K gate matters (the probe dumps during it, and the + # gate aborts at ~0.88 before the sweep). Trim the sweep to conc=1. + concurrencies: [1] + 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 From b329fd5696b3f516b5ab89be40268f0c45a42ef2 Mon Sep 17 00:00:00 2001 From: yctseng0211 Date: Wed, 8 Jul 2026 12:07:11 -0500 Subject: [PATCH 6/8] [AMD][DI][CI] Extend decode-meta probe to prefill (EXTEND) for MORI transfer-fidelity check --- .../debug_utils/disagg_decode_meta_probe.py | 27 ++-- .../srt/layers/attention/aiter_backend.py | 10 ++ test/manual/test_kimi_k26_decodemeta_mi35x.py | 118 ++++++++++++++++++ 3 files changed, 148 insertions(+), 7 deletions(-) create mode 100644 test/manual/test_kimi_k26_decodemeta_mi35x.py diff --git a/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py b/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py index 488c2bdf62cc..0604ed163844 100644 --- a/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py +++ b/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py @@ -133,7 +133,13 @@ def _dump_impl(tag: str, backend, layer, forward_batch) -> None: fwd_mode = forward_batch.forward_mode is_decode = fwd_mode.is_decode() or fwd_mode.is_idle() is_verify = fwd_mode.is_target_verify() - if not (is_decode or is_verify): + is_extend = ( + fwd_mode.is_extend() and not is_verify and not fwd_mode.is_draft_extend_v2() + ) + # DECODE/VERIFY fire on the decode worker; EXTEND fires on the prefill + # worker (its prompt extend) -> one disagg run yields both the + # prefill-stored and decode-received latent norms for the same tokens. + if not (is_decode or is_verify or is_extend): return seq_lens = getattr(forward_batch, "seq_lens_cpu", None) @@ -169,7 +175,7 @@ def _dump_impl(tag: str, backend, layer, forward_batch) -> None: else None ) - mode_name = "VERIFY" if is_verify else "DECODE" + mode_name = "VERIFY" if is_verify else ("EXTEND" if is_extend else "DECODE") lines = [] for b in range(bs): seq_len = int(seq_lens[b]) @@ -190,7 +196,12 @@ def _dump_impl(tag: str, backend, layer, forward_batch) -> None: n_kv = hi - lo idx_slice = kv_indices[lo:hi].detach().to("cpu") idx_str = _ht(idx_slice.tolist()) - if idx_slice.numel() == rtt.numel(): + # kv_indices == req_to_token[req, :seq_len] only for DECODE; EXTEND's + # kv_indices covers the prefix (empty for a pure prefill), so the + # equality check is decode-only to avoid a spurious LEN_MISMATCH. + if not is_decode: + match_str = "n/a(non-decode)" + elif idx_slice.numel() == rtt.numel(): match_str = str(bool(torch.equal(idx_slice.to(rtt.dtype), rtt))) else: match_str = f"LEN_MISMATCH({idx_slice.numel()} vs {rtt.numel()})" @@ -211,11 +222,13 @@ def _dump_impl(tag: str, backend, layer, forward_batch) -> None: f" rtt = {_ht(rtt.tolist())}" ) - # KV-pool row norms at the attended slots (real decode step, unlike the - # warmup-only #30433 dump). Same prompt => same norms across runs. + # KV-pool row norms at req_to_token[req, :seq_len] (real step, unlike the + # warmup-only #30433 dump). This is the SAME quantity on both sides: + # what the prefill worker stored (EXTEND) vs what the decode worker reads + # (DECODE). Same prompt (shared 8-shot prefix) => identical norm sequence + # iff the MORI transfer preserved the latent values. if _KVNORM and layer_id == _FIRST_LAYER: - slots = idx_slice if idx_slice is not None else rtt - block += "\n" + _kv_norm_report(backend, layer_id, slots) + block += "\n" + _kv_norm_report(backend, layer_id, rtt) lines.append(block) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index d8a81212ebb8..054549a80d6d 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -1957,6 +1957,16 @@ def forward_extend( v_descale, ) + from sglang.srt.debug_utils.disagg_decode_meta_probe import ( + maybe_dump_decode_meta, + ) + + # Prefill-side latent dump: KV is now written to the pool above, so this + # captures what the prefill worker stored for the prompt. Compared with + # the decode-worker DECODE dump over the same tokens, it isolates whether + # the MORI transfer preserved the latent values. + maybe_dump_decode_meta("aiter-extend", self, layer, forward_batch) + if self.use_mla: max_q_len = self.forward_metadata.max_q_len max_kv_len = self.forward_metadata.max_kv_len diff --git a/test/manual/test_kimi_k26_decodemeta_mi35x.py b/test/manual/test_kimi_k26_decodemeta_mi35x.py new file mode 100644 index 000000000000..0d1d68a49169 --- /dev/null +++ b/test/manual/test_kimi_k26_decodemeta_mi35x.py @@ -0,0 +1,118 @@ +"""MI35x Kimi-K2.6 single-node DECODE-metadata probe (diagnostic, DO NOT MERGE). + +Single-node counterpart to the disagg `1p1d-metadump` recipe, for the non-MTP +GSM8K drop investigation (disagg ~0.88 vs single-node 0.944). It launches the +SAME server config as test/registered/amd/accuracy/mi35x/test_kimi_k26_eval_mi35x.py +(TP8, aiter prefill / triton decode) but EAGER (--disable-cuda-graph) and with +the decode-metadata probe enabled, then runs a few GSM8K 8-shot questions at +parallel=1 (bs=1) so the `[DDM]` dump has the same shape as the disagg run. + +Purpose: the disagg probe already showed the transferred prefix is present +(zero_rows=0), correctly indexed (idx==rtt), and seq_len is monotonic -- so the +only remaining variable is the transferred KV *values*. The 8-shot GSM8K prefix +is identical here and in the disagg gate, so the per-token latent norm sequence +must match if the values are preserved. Compare this run's step=1 `KV L0 norm` +head against the disagg run's: + + KV L0 norm head = [19.617, 22.029, 17.145, 23.779, 17.698, 19.234] (kv_idx = first 6 prefix tokens) + + * norms MATCH -> transferred KV values are faithful; the bug is in the decode + compute / missing prefill-established state, not the KV read. + * norms DIFFER -> the MORI transfer alters the latent values (layout / stride / + dtype / partial-head), i.e. a transfer-fidelity bug. + +Run (single MI35x node, inside the ROCm container): + python3 test/manual/test_kimi_k26_decodemeta_mi35x.py +then: grep -A6 "\[DDM" on the server stdout. + +Manual-only diagnostic (lives under test/manual/, not registered for CI). +""" + +import os +import unittest +from types import SimpleNamespace + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +KIMI_K26_MODEL_PATH = "moonshotai/Kimi-K2.6" +SERVER_LAUNCH_TIMEOUT = 5400 +TP_SIZE = 8 + + +class TestKimiK26DecodeMetaMI35x(CustomTestCase): + """Single-node decode-metadata probe run (diagnostic).""" + + @classmethod + def setUpClass(cls): + cls.base_url = DEFAULT_URL_FOR_TEST + + def test_kimi_k26_decode_metadata_probe(self): + other_args = [ + "--tp", + str(TP_SIZE), + "--decode-attention-backend", + "triton", + "--prefill-attention-backend", + "aiter", + "--trust-remote-code", + "--model-loader-extra-config", + '{"enable_multithread_load": true}', + "--watchdog-timeout", + "1200", + # Eager decode so forward_decode is entered in Python (the probe + # never fires under HIP graph replay). + "--disable-cuda-graph", + ] + env = os.environ.copy() + env["SGLANG_USE_AITER"] = "1" + env["SGLANG_ROCM_FUSED_DECODE_MLA"] = "0" + # Decode-metadata probe (see debug_utils/disagg_decode_meta_probe.py). + env["SGLANG_DEBUG_DISAGG_DECODE_META"] = "1" + env["SGLANG_DEBUG_DISAGG_DECODE_META_MINLEN"] = "128" + env["SGLANG_DEBUG_DISAGG_DECODE_META_STEPS"] = "16" + env["SGLANG_DEBUG_DISAGG_DECODE_META_KVNORM"] = "1" + + process = popen_launch_server( + KIMI_K26_MODEL_PATH, + self.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=other_args, + env=env, + ) + + try: + requests.get(self.base_url + "/flush_cache") + + # parallel=1 -> bs=1 decode, so the [DDM] dump matches the disagg + # run's shape. A handful of questions is enough; only the first + # request's first STEPS decode steps are dumped. Accuracy is not the + # point here, so no threshold assertion. + args = SimpleNamespace( + num_shots=8, + data_path=None, + num_questions=8, + parallel=1, + max_new_tokens=32, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print( + f"[decodemeta] single-node accuracy over " + f"{args.num_questions} q = {metrics['accuracy']:.3f} " + f"(diagnostic; grep '[DDM' in server log for the metadata dump)" + ) + finally: + kill_process_tree(process.pid) + + +if __name__ == "__main__": + unittest.main() From 7d598f964335e320c4ceb32646c010750a6b907e Mon Sep 17 00:00:00 2001 From: yctseng0211 Date: Wed, 8 Jul 2026 12:47:46 -0500 Subject: [PATCH 7/8] [AMD][DI][CI] Sample fixed absolute KV positions to check whole-prefix transfer fidelity --- .../debug_utils/disagg_decode_meta_probe.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py b/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py index 0604ed163844..1b8d15e680f5 100644 --- a/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py +++ b/python/sglang/srt/debug_utils/disagg_decode_meta_probe.py @@ -244,10 +244,23 @@ def _kv_norm_report(backend, layer_id: int, slots: torch.Tensor) -> str: rows = kbuf[dev_slots].reshape(dev_slots.shape[0], -1).float() norms = rows.norm(dim=-1) zero_rows = int((norms < 1e-6).sum().item()) - norms_list = [round(x, 3) for x in norms.detach().to("cpu").tolist()] + norms_cpu = norms.detach().to("cpu") + norms_list = [round(x, 3) for x in norms_cpu.tolist()] + # Fixed ABSOLUTE-position samples. These token positions fall in the + # shared 8-shot GSM8K prefix, so they are the SAME token on the prefill + # (EXTEND) and decode (DECODE) side and across questions/runs -> the norm + # at each position must match iff the MORI transfer preserved that page. + # Spans multiple pages (page_size=256) to catch per-page transfer bugs + # that head/tail (page 0 + current token) would miss. + n = norms_cpu.numel() + samples = { + p: round(float(norms_cpu[p]), 3) + for p in (0, 128, 256, 512, 768, 1024) + if p < n + } return ( - f" KV L{layer_id}: zero_rows={zero_rows}/{norms.numel()} " - f"norm={_ht(norms_list)}" + f" KV L{layer_id}: zero_rows={zero_rows}/{n} " + f"norm={_ht(norms_list)} samples@abs={samples}" ) except Exception as exc: return f" KV L{layer_id}: norm probe failed: {exc!r}" From 13cf17af25c1136a1b8c0472808ed533c469724b Mon Sep 17 00:00:00 2001 From: yctseng0211 Date: Wed, 8 Jul 2026 18:46:02 -0500 Subject: [PATCH 8/8] [AMD][DI][CI] Force decode nsplit=1 to test split-KV reduction as Kimi disagg drop cause --- .../srt/layers/attention/triton_backend.py | 7 ++ scripts/ci/slurm/nightly-configs.yaml | 19 ++++ .../mi355x-fp8/kimik26/1k1k/1p1d-nsplit1.yaml | 91 +++++++++++++++++++ 3 files changed, 117 insertions(+) create mode 100644 scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-nsplit1.yaml diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 4a3ad5879121..7766c70f6137 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -63,6 +63,13 @@ def _mla_decode_kv_splits_cap( base_max_kv_splits: int, sm_count: int, max_context_len: int ) -> int: + # Diagnostic override: skip the MLA split floor so + # --triton-attention-num-kv-splits can force nsplit as low as 1. Used to test + # whether the decode split-KV online-softmax reduction over a long + # transferred prefix is the source of the non-MTP disagg accuracy drop + # (extend/verify uses no split-KV and is correct at ~0.95). + if get_bool_env_var("SGLANG_DEBUG_MLA_DECODE_NO_SPLIT_CAP", "false"): + return base_max_kv_splits if sm_count <= 0: return base_max_kv_splits sm_cap = next_power_of_2(sm_count) diff --git a/scripts/ci/slurm/nightly-configs.yaml b/scripts/ci/slurm/nightly-configs.yaml index ac943c3a0fc8..47d3c23b63ac 100644 --- a/scripts/ci/slurm/nightly-configs.yaml +++ b/scripts/ci/slurm/nightly-configs.yaml @@ -375,3 +375,22 @@ kimik26-fp8-mi355x-metadump-sglang: search-space: - conc-list: [1] config_file: scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-metadump.yaml + +# DIAGNOSTIC (DO NOT MERGE): non-MTP + force decode nsplit=1 to test whether the +# split-KV reduction is the source of the drop. Run via +# workflow_dispatch `configs=kimik26-fp8-1k1k-1p1d-nsplit1`. +kimik26-fp8-mi355x-nsplit1-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] + config_file: scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-nsplit1.yaml diff --git a/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-nsplit1.yaml b/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-nsplit1.yaml new file mode 100644 index 000000000000..ddb5e033f8f2 --- /dev/null +++ b/scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d-nsplit1.yaml @@ -0,0 +1,91 @@ +# MI355X Kimi-K2.6 (FP8) 2N 1P1D — force decode nsplit=1 (diagnostic, DO NOT MERGE). +# +# Tests whether the decode split-KV online-softmax reduction is the source of +# the non-MTP GSM8K drop. Transfer + metadata + KV values are already CONFIRMED +# correct (probe runs 28957394001/28962093651/28963785464), so the only place +# left is the decode compute: decode uses decode_attention_fwd (split-KV) while +# verify uses extend_attention_fwd (no split-KV) and is correct at ~0.95. At the +# gate's bs=1 the MLA floor gives nsplit~130; forcing nsplit=1 removes the +# cross-split fp reduction. +# +# * accuracy recovers to ~0.94 -> the split-KV reduction is the culprit. +# * still ~0.88 -> rule it out; escalate to decode-via-extend. +# +# Same as 1p1d-metadump.yaml (probe + eager) plus: +# env += SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS (fill num_kv_splits with max) +# += SGLANG_DEBUG_MLA_DECODE_NO_SPLIT_CAP (skip the MLA nsplit floor) +# args += --triton-attention-num-kv-splits 1 +# The probe's `nsplit=` field confirms it actually became 1. +# +# Trigger (config name = kimik26-fp8-1k1k-1p1d-nsplit1): +# gh workflow run "Nightly Test (AMD MI355X 2N 1P1D Disagg)" \ +# --ref amd/kimik26-disagg-decodemeta \ +# -f configs=kimik26-fp8-1k1k-1p1d-nsplit1 + +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: + env: + SGLANG_USE_AITER: 1 + SGLANG_ROCM_FUSED_DECODE_MLA: 0 + # Force decode nsplit=1 (see triton_backend _mla_decode_kv_splits_cap). + SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS: true + SGLANG_DEBUG_MLA_DECODE_NO_SPLIT_CAP: true + # Decode-metadata probe (confirms nsplit=1 in the dump). + SGLANG_DEBUG_DISAGG_DECODE_META: 1 + SGLANG_DEBUG_DISAGG_DECODE_META_MINLEN: 128 + SGLANG_DEBUG_DISAGG_DECODE_META_STEPS: 16 + SGLANG_DEBUG_DISAGG_DECODE_META_KVNORM: 1 + server_args: + - --model-loader-extra-config + - '{"enable_multithread_load": true}' + - --watchdog-timeout + - 1200 + - --reasoning-parser + - kimi_k2 + - --tool-call-parser + - kimi_k2 + # Force a single KV split on the decode kernel. + - --triton-attention-num-kv-splits + - 1 + # Eager decode so forward_decode runs in Python (probe fires). + - --disable-cuda-graph + +runtime: + image: lmsysorg/sglang-rocm:v0.5.13.post1-rocm720-mi35x-20260623 + # Kimi uses split attention backends (aiter prefill / triton decode). + prefill_attention_backend: aiter + decode_attention_backend: triton + 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: + concurrencies: [1] + num_prompts_factor: 4 + random_range_ratio: 1.0 + accuracy: + enabled: true + num_shots: 8 + num_questions: 1319 + threshold: 0.92