feat(agentic): add GLM-5.2 MI325X full-context sweeps / 新增 GLM-5.2 MI325X 全上下文 AgentX 扫描

benchmarks/multi_node/amd_utils/env.sh

@@ -236,12 +236,18 @@ else
     export SGLANG_HEALTH_CHECK_TIMEOUT=600
 
     # GLM-5: uses NSA (not MLA), needs fused-decode-MLA disabled + fast loading
-    if [[ "$MODEL_NAME" == "GLM-5-FP8" ]]; then
+    MODEL_CONFIG_KEY="${MODEL_YAML_KEY:-$MODEL_NAME}"
+    if [[ "$MODEL_CONFIG_KEY" == "GLM-5-FP8" || "$MODEL_CONFIG_KEY" == "GLM-5.2-FP8" ]]; then
         export SGLANG_ROCM_FUSED_DECODE_MLA=0
         export ROCM_QUICK_REDUCE_QUANTIZATION=INT4
         export SAFETENSORS_FAST_GPU=1
     fi
 
+    if [[ "$MODEL_CONFIG_KEY" == "GLM-5.2-FP8" ]]; then
+        export SGLANG_DSA_FUSE_TOPK=false
+        export SGLANG_OPT_USE_TOPK_V2=false
+    fi
+
     # Disable allocating memory in one pass
     export MORI_SHMEM_MODE=ISOLATION
 

benchmarks/multi_node/amd_utils/job.slurm

@@ -42,14 +42,14 @@ if [[ -z "${DOCKER_IMAGE_NAME:-}" ]]; then
 fi
 
 MODEL_NAME="${MODEL_NAME:-None}"
-# Resolve the models.yaml entry the same way server_sglang.sh does: agentic runs
-# (IS_AGENTIC) use the '<model>-AgentX' recipe, non-agentic disaggregated runs use
-# '<model>-DI'. Fall back to the bare model name if the variant key is absent.
-# MODEL_NAME itself is left unchanged so env.sh/server_sglang.sh still see the base name.
-if [[ "${IS_AGENTIC:-0}" == "1" || "${IS_AGENTIC:-}" == "true" ]]; then
-    MODEL_YAML_KEY="${MODEL_NAME}-AgentX"
-else
-    MODEL_YAML_KEY="${MODEL_NAME}-DI"
+# Respect an explicit recipe key (needed when MODEL_NAME is a cache path), then
+# follow the standard AgentX/DI naming convention used by current AMD recipes.
+if [[ -z "${MODEL_YAML_KEY:-}" ]]; then
+    if [[ "${IS_AGENTIC:-0}" == "1" || "${IS_AGENTIC:-}" == "true" ]]; then
+        MODEL_YAML_KEY="${MODEL_NAME}-AgentX"
+    else
+        MODEL_YAML_KEY="${MODEL_NAME}-DI"
+    fi
 fi
 if ! grep -q "^${MODEL_YAML_KEY}:" "$MODELS_YAML"; then
     if grep -q "^${MODEL_NAME}:" "$MODELS_YAML"; then
@@ -117,7 +117,7 @@ export MODEL_DIR
 
 if [[ "$ENGINE" == "vllm-disagg" ]]; then
     # vLLM: Extract hf_dir from models.yaml, search multiple paths, resolve HF cache snapshots
-    DISK_DIR_NAME=$(awk '/^'"$MODEL_NAME"':/{found=1; next}
+    DISK_DIR_NAME=$(awk -v key="$MODEL_YAML_KEY" '$0 == key ":" {found=1; next}
         found && /^[^ ]/{exit}
         found && /hf_dir:/{gsub(/[" ]/, "", $2); print $2; exit}' "$MODELS_YAML")
     DISK_DIR_NAME="${DISK_DIR_NAME:-$MODEL_NAME}"
@@ -193,7 +193,7 @@ else
     }
 
     # Extract hf_dir from models.yaml (same as vllm-disagg path above)
-    SGL_DISK_DIR_NAME=$(awk '/^'"$MODEL_NAME"':/{found=1; next}
+    SGL_DISK_DIR_NAME=$(awk -v key="$MODEL_YAML_KEY" '$0 == key ":" {found=1; next}
         found && /^[^ ]/{exit}
         found && /hf_dir:/{gsub(/[" ]/, "", $2); print $2; exit}' "$MODELS_YAML")
     SGL_DISK_DIR_NAME="${SGL_DISK_DIR_NAME:-$MODEL_NAME}"
@@ -243,6 +243,11 @@ echo "NUM_NODES: $NUM_NODES (xP=$xP + yD=$yD)"
 FULL_NODELIST=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
 SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n $NUM_NODES)
 SELECTED_NODELIST_STR=$(echo "$SELECTED_NODES" | tr '\n' ',' | sed 's/,$//')
+SELECTED_NODE_COUNT=$(echo "$SELECTED_NODES" | sed '/^$/d' | wc -l)
+if [[ "$SELECTED_NODE_COUNT" -ne "$NUM_NODES" ]]; then
+    echo "FATAL ERROR: allocation exposes ${SELECTED_NODE_COUNT} nodes, expected ${NUM_NODES}" >&2
+    exit 1
+fi
 
 # Docker privilege detection — evaluated per-node since group membership varies.
 # Exported as a snippet so every srun participant resolves it locally.
@@ -267,16 +272,27 @@ echo "Selected nodes: $SELECTED_NODELIST_STR"
 
 USER_NAME=$(whoami)
 MASTER_NODE=$(echo "$SELECTED_NODES" | head -n 1)
-NODE0_ADDR=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$MASTER_NODE" bash -c 'ip route get 1.1.1.1')
-NODE0_ADDR=$(echo "$NODE0_ADDR" | awk '/src/ {print $7}')
-
 IPS=()
 for NODE in $SELECTED_NODES; do
-    IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'ip route get 1.1.1.1')
-    IP=$(echo "$IP" | awk '/src/ {print $7}')
+    if ! IP_ROUTE=$(srun --nodes=1 --ntasks=1 --time=00:05:00 --nodelist="$NODE" \
+        bash -c 'ip route get 1.1.1.1'); then
+        echo "FATAL ERROR: unable to launch IP-resolution task on $NODE" >&2
+        exit 1
+    fi
+    IP=$(echo "$IP_ROUTE" | awk '/src/ {print $7; exit}')
+    if [[ -z "$IP" ]]; then
+        echo "FATAL ERROR: no source IP resolved for $NODE" >&2
+        exit 1
+    fi
     IPS+=("$IP")
 done
 
+if [[ "${#IPS[@]}" -ne "$NUM_NODES" ]]; then
+    echo "FATAL ERROR: resolved ${#IPS[@]} node IPs, expected ${NUM_NODES}" >&2
+    exit 1
+fi
+NODE0_ADDR="${IPS[0]}"
+
 echo "Node IPs: ${IPS[*]}"
 
 DOCKER_MOUNT_PATH="/workspace"
@@ -332,7 +348,8 @@ export SPEC_DECODING="${SPEC_DECODING:-}"
 export IS_MULTINODE="${IS_MULTINODE:-false}"
 
 SANITIZED_USER=$(echo "$USER_NAME" | tr -c 'a-zA-Z0-9_.-' '_')
-export DOCKER_CONT_NAME="container_${ENGINE}_${SANITIZED_USER}_${MODEL_NAME}_${SLURM_JOB_ID}"
+SANITIZED_MODEL=$(printf '%s' "$MODEL_NAME" | tr -c 'a-zA-Z0-9_.-' '_')
+export DOCKER_CONT_NAME="container_${ENGINE}_${SANITIZED_USER}_${SANITIZED_MODEL}_${SLURM_JOB_ID}"
 
 # vLLM external router container.
 # NOTE: vllm/vllm-router only retains ~16 recent nightlies on Docker Hub; older
@@ -384,6 +401,7 @@ DOCKER_ENV_COMMON=(
     -e NODE0_ADDR=\$NODE0_ADDR
     -e MODEL_DIR=/models
     -e MODEL_NAME=\$MODEL_NAME
+    -e MODEL_YAML_KEY=\$MODEL_YAML_KEY
     -e GPUS_PER_NODE=\$GPUS_PER_NODE
     -e xP=\$xP
     -e yD=\$yD
@@ -445,6 +463,17 @@ DOCKER_ENV_COMMON=(
     -e DECODE_ENABLE_DP=\$DECODE_ENABLE_DP
     -e DECODE_MTP_SIZE=\$DECODE_MTP_SIZE
     -e IS_MULTINODE=\$IS_MULTINODE
+    -e SCENARIO_TYPE=\$SCENARIO_TYPE
+    -e CONC=\$CONC
+    -e CONC_LIST
+    -e ROUTER_METADATA
+    -e KV_P2P_TRANSFER=\$KV_P2P_TRANSFER
+    -e MODEL=\$MODEL
+    -e WEKA_LOADER_OVERRIDE=\$WEKA_LOADER_OVERRIDE
+    -e SYNC_BARRIER_TIMEOUT
+    -e HF_HUB_CACHE=/models
+    -e IBDEVICES
+    -e MORI_RDMA_TC
     -e DRY_RUN=\${DRY_RUN:-0}
 )
 
@@ -539,12 +568,27 @@ if [[ -n "${CLIENT_IMAGE:-}" ]]; then
     srun --nodelist="$SELECTED_NODELIST_SRUN" bash -c 'eval "$DOCKER_CMD_DETECT"; $DOCKER_CMD pull '"$CLIENT_IMAGE"' >/dev/null 2>&1 || true' 2>/dev/null || true
 fi
 
+echo "Ensuring Docker image is present on every selected node..."
+srun \
+  --nodelist="$SELECTED_NODELIST_SRUN" \
+  --kill-on-bad-exit=1 \
+  --unbuffered \
+  bash -c '
+set -euo pipefail
+eval "$DOCKER_CMD_DETECT"
+if ! $DOCKER_CMD image inspect "$DOCKER_IMAGE_NAME" >/dev/null 2>&1; then
+    echo "[docker-pull] $(hostname): $DOCKER_IMAGE_NAME"
+    $DOCKER_CMD pull "$DOCKER_IMAGE_NAME"
+fi
+'
+
+SERVER_RUN_RC=0
 srun \
   --nodelist="$SELECTED_NODELIST_SRUN" \
   --kill-on-bad-exit=1 \
   --signal=TERM@30 \
   --unbuffered \
-  bash -lc "
+  bash -c "
 set -euo pipefail
 
 echo \"Rank \$SLURM_PROCID on \$(hostname)\"
@@ -635,6 +679,8 @@ fi  # end: if ENGINE == atom-disagg
 \$DOCKER_CMD ps -aq --filter \"$CONT_FILTER\" | xargs -r \$DOCKER_CMD rm -f || true
 \$DOCKER_CMD ps -aq | xargs -r \$DOCKER_CMD stop || true
 
+echo \"[docker-run] image=$DOCKER_IMAGE_NAME container=$DOCKER_CONT_NAME\"
+
 # Start vLLM external router container on node 0
 if [[ \"$ENGINE\" == \"vllm-disagg\" && \"$ROUTER_TYPE\" == \"vllm-router\" && \"\$SLURM_PROCID\" == \"0\" ]]; then
     \$DOCKER_CMD rm -f \"$ROUTER_CONT_NAME\" 2>/dev/null || true
@@ -694,7 +740,7 @@ fi
     -v ${MODEL_DIR}:/models \
     -v \$HOME/.ssh:/root/.ssh \
     --shm-size 128G \
-    -v /tmp:/run_logs \
+    -v ${BENCHMARK_LOGS_DIR}:/run_logs \
     -v ${BENCHMARK_LOGS_DIR}:/benchmark_logs \
     -v ${DI_REPO_DIR}:${DOCKER_MOUNT_PATH} \
     -v ${HICACHE_MC_CONFIG}:/config/hicache_mc.env:ro \
@@ -717,7 +763,7 @@ DOCKER_EXIT_CODE=\$?
 echo \"[rank 0] Main container exited (rc=\$DOCKER_EXIT_CODE). Stopping vllm-router...\"
 \$DOCKER_CMD rm -f \"$ROUTER_CONT_NAME\" 2>/dev/null || true
 exit \$DOCKER_EXIT_CODE
-"
+" || SERVER_RUN_RC=$?
 
 if [[ "${KEEP_CONTAINERS}" != "1" ]]; then
     srun --nodelist="$SELECTED_NODELIST_SRUN" bash -c 'eval "$DOCKER_CMD_DETECT"; $DOCKER_CMD rm -f '"$DOCKER_CONT_NAME"' '"$CLIENT_CONT_NAME"' 2>/dev/null || true'
@@ -729,3 +775,8 @@ if [[ "${KEEP_CONTAINERS}" != "1" ]]; then
         '
     fi
 fi
+
+if [[ "$SERVER_RUN_RC" -ne 0 ]]; then
+    echo "ERROR: distributed server step failed with exit code $SERVER_RUN_RC" >&2
+    exit "$SERVER_RUN_RC"
+fi

benchmarks/multi_node/amd_utils/models.yaml

@@ -264,6 +264,38 @@ GLM-5-FP8:
       chunked_prefill_size: 262144
       cuda_graph_bs_range: "1-128"
 
+GLM-5.2-FP8:
+  base_flags: "--decode-log-interval 1000 --log-level info --watchdog-timeout 3600 --load-balance-method round_robin --served-model-name zai-org/GLM-5.2-FP8 --enable-metrics --disaggregation-transfer-backend mori --tool-call-parser glm47 --reasoning-parser glm45 --dsa-prefill-backend tilelang --dsa-decode-backend tilelang --dsa-topk-backend torch --kv-cache-dtype bfloat16 --context-length 1048576 --max-total-tokens 1048576 --model-loader-extra-config '{\\\"enable_multithread_load\\\": true, \\\"num_threads\\\": 8}'"
+  mtp_flags: ""
+  dp_flags: "--enable-dp-attention --moe-dense-tp-size 1 --enable-dp-lm-head"
+  ep_flags: "--moe-a2a-backend mori"
+  prefill:
+    mem_fraction_static: 0.85
+    disable_radix_cache: false
+    dp:
+      max_running_requests: 16
+      chunked_prefill_size: 131072
+      cuda_graph_bs: "1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16"
+    no_dp:
+      max_running_requests: 16
+      chunked_prefill_size: 131072
+      cuda_graph_bs_range: "1-16"
+  decode:
+    mem_fraction_static: 0.85
+    prefill_round_robin_balance: true
+    dp:
+      max_running_requests: 16
+      chunked_prefill_size: 131072
+      cuda_graph_bs_range: "1-16"
+    ep_only:
+      max_running_requests: 16
+      chunked_prefill_size: 131072
+      cuda_graph_bs_range: "1-16"
+    no_dp:
+      max_running_requests: 16
+      chunked_prefill_size: 131072
+      cuda_graph_bs_range: "1-16"
+
 DeepSeek-R1-0528-MXFP4-Preview:
   base_flags: "--decode-log-interval 1000 --log-level warning --watchdog-timeout 3600 --load-balance-method round_robin --kv-cache-dtype fp8_e4m3 --attention-backend aiter --disaggregation-transfer-backend mori"
   mtp_flags: "--speculative-algorithm NEXTN --speculative-eagle-topk 1"

benchmarks/multi_node/amd_utils/server_sglang.sh

@@ -81,12 +81,13 @@ import yaml, sys, os
 
 config_path = '${MODELS_YAML}'
 model_name = '${MODEL_NAME}'
+explicit_model_key = '${MODEL_YAML_KEY:-}'.strip()
 
 # Select the models.yaml recipe variant by run type: agentic runs (IS_AGENTIC)
 # use the '<model>-AgentX' entry, non-agentic disaggregated runs use '<model>-DI'.
 # Fall back to the bare model name if the variant-specific key is absent.
 is_agentic = '${IS_AGENTIC:-0}'.strip().lower() in ('1', 'true')
-model_key = f'{model_name}-AgentX' if is_agentic else f'{model_name}-DI'
+model_key = explicit_model_key or (f'{model_name}-AgentX' if is_agentic else f'{model_name}-DI')
 
 with open(config_path) as f:
     models = yaml.safe_load(f)
@@ -604,7 +605,18 @@ wait_or_die() {            # $1 = server pid to watch; rest = blocking command
     "$@" & local cmd=$!
     while kill -0 "$cmd" 2>/dev/null; do
         kill -0 "$watch" 2>/dev/null || {
-            echo "FATAL: $(hostname) local sglang server (pid $watch) died; tearing down job" >&2
+            local server_rc=0
+            wait "$watch" || server_rc=$?
+            sleep 1
+            echo "FATAL: $(hostname) local sglang server (pid $watch) died with rc=$server_rc; tearing down job" >&2
+            local log_dir="/run_logs/slurm_job-${SLURM_JOB_ID}"
+            local log_file
+            for log_file in "$log_dir"/*_"$(hostname)".log; do
+                [[ -f "$log_file" ]] || continue
+                echo "===== Local SGLang log: $log_file =====" >&2
+                tail -n 200 "$log_file" >&2 || true
+                echo "===== End local SGLang log =====" >&2
+            done
             kill "$cmd" 2>/dev/null || true
             return 1
         }

benchmarks/multi_node/amd_utils/setup_deps.sh

@@ -78,6 +78,72 @@ install_amd_quark() {
     _SETUP_INSTALLED+=("amd-quark")
 }
 
+# ---------------------------------------------------------------------------
+# SGLang: Patch aiter gluon pa_mqa_logits — fix 2D → 3D instr_shape for
+# Triton >= 3.5. The base non-preshuffle variant omitted the conditional used
+# by the other variants; GLM-5.2 exercises this path on gfx942.
+# ---------------------------------------------------------------------------
+patch_gluon_pa_mqa_logits_instr_shape() {
+    python3 -c '
+import os, sys
+
+target = "/sgl-workspace/aiter/aiter/ops/triton/gluon/pa_mqa_logits.py"
+if not os.path.isfile(target):
+    print("[SETUP] gluon pa_mqa_logits.py not found, skipping")
+    sys.exit(0)
+
+src = open(target).read()
+if "[PATCHED] 3D instr_shape for base gluon variant" in src:
+    print("[SETUP] gluon pa_mqa_logits 3D instr_shape patch already applied")
+    sys.exit(0)
+
+old = """\
+    mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout(
+        version=CDNA_VERSION,
+        instr_shape=[16, 16],
+        transposed=False,
+        warps_per_cta=[1, NumWarps],
+    )
+    mfma_layout_a: gl.constexpr = gl.DotOperandLayout(
+        operand_index=0, parent=mfma_layout, k_width=16
+    )
+    mfma_layout_b: gl.constexpr = gl.DotOperandLayout(
+        operand_index=1, parent=mfma_layout, k_width=16
+    )"""
+
+new = """\
+    # [PATCHED] 3D instr_shape for base gluon variant
+    if _Use_2d_instr_shape_mfma_layout:
+        mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout(
+            version=CDNA_VERSION,
+            instr_shape=[16, 16],
+            transposed=False,
+            warps_per_cta=[1, NumWarps],
+        )
+    else:
+        mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout(
+            version=CDNA_VERSION,
+            instr_shape=[16, 16, 32],
+            transposed=False,
+            warps_per_cta=[1, NumWarps],
+        )
+    mfma_layout_a: gl.constexpr = gl.DotOperandLayout(
+        operand_index=0, parent=mfma_layout, k_width=16
+    )
+    mfma_layout_b: gl.constexpr = gl.DotOperandLayout(
+        operand_index=1, parent=mfma_layout, k_width=16
+    )"""
+
+if old not in src:
+    print("[SETUP] WARN: gluon pa_mqa_logits pattern not found — aiter version may have changed")
+    sys.exit(0)
+
+open(target, "w").write(src.replace(old, new, 1))
+print("[SETUP] Patched: gluon pa_mqa_logits 3D instr_shape for base variant")
+'
+    _SETUP_INSTALLED+=("gluon-instr-shape-fix")
+}
+
 # ---------------------------------------------------------------------------
 # SGLang: Install latest transformers for GLM model type support.
 #
@@ -119,6 +185,7 @@ if [[ "$ENGINE" == "vllm-disagg" ]]; then
     export PATH="${UCX_HOME}/bin:/usr/local/bin/etcd:/root/.cargo/bin:${PATH}"
     export LD_LIBRARY_PATH="${UCX_HOME}/lib:${RIXL_HOME}/lib:${RIXL_HOME}/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}"
 else
+    patch_gluon_pa_mqa_logits_instr_shape
     install_transformers_glm5
 fi