Fix for https://github.com/sgl-project/sglang/issues/22072 / 修复 SGLang MoRI 分离式推理在 MI325X 上的挂起问题

benchmarks/multi_node/amd_utils/env.sh

@@ -62,6 +62,19 @@ set +x
 
 export NCCL_IB_HCA=${NCCL_IB_HCA:-$IBDEVICES}
 
+# =============================================================================
+# Shared: runner hardware family
+# =============================================================================
+# Reduce the runner label to its GPU family, mirroring
+# infx/matrix/generate.py::_hardware_family. RUNNER_TYPE reaches the container
+# either abstract ("mi325x") or fleet-scoped ("cluster:mi325x-amds",
+# "cluster:mi325x-tw"); configs/runners.yaml maps both onto the same
+# mi325x-amds_NN machines, so the family is what SKU-specific settings gate on.
+# Exported for server_sglang.sh, which sources this file.
+RUNNER_FAMILY="${RUNNER_TYPE:-}"
+RUNNER_FAMILY="${RUNNER_FAMILY#cluster:}"
+export RUNNER_FAMILY="${RUNNER_FAMILY%%-*}"
+
 # =============================================================================
 # MoRI-specific environment
 # =============================================================================
@@ -128,6 +141,28 @@ else
     fi
 fi
 
+# Fan the detected class out to every transport that touches this fabric. MoRI's KV
+# transfer and NCCL's own collectives have to request the same PFC-protected class,
+# or the lossless guarantee only covers half the inter-node traffic.
+#
+# SL is the 802.1p priority: TC is the IP ToS byte (DSCP<<2) and priority is DSCP>>3,
+# so SL = TC>>5 (e.g. TC=104 -> DSCP 26/AF31 -> SL 3). bnxt_re REJECTS inconsistent
+# DSCP/SL pairs ("Given DSCP N and/or SL M not mapping to lossless queue") and
+# SILENTLY downgrades to the best-effort queue, which surfaces under load as
+# RETRY_EXC_ERR / stalled KV transfers, so every path needs the SL filled in - not
+# just the runner-set one.
+#
+# When detection found nothing these stay unset: requesting a class the fabric does
+# not map to its lossless queue is worse than leaving the transports on the default.
+if [[ -n "${MORI_RDMA_TC:-}" ]]; then
+    [[ -z "${MORI_RDMA_SL:-}" ]] && export MORI_RDMA_SL=$(( MORI_RDMA_TC >> 5 ))
+    export MORI_IO_TC="${MORI_IO_TC:-$MORI_RDMA_TC}"
+    export MORI_IO_SL="${MORI_IO_SL:-$MORI_RDMA_SL}"
+    export NCCL_IB_TC="${NCCL_IB_TC:-$MORI_RDMA_TC}"
+    export NCCL_IB_SL="${NCCL_IB_SL:-$MORI_RDMA_SL}"
+    echo "[INFO] RDMA QoS: MORI_RDMA_TC=$MORI_RDMA_TC MORI_RDMA_SL=$MORI_RDMA_SL MORI_IO_TC=$MORI_IO_TC MORI_IO_SL=$MORI_IO_SL NCCL_IB_TC=$NCCL_IB_TC NCCL_IB_SL=$NCCL_IB_SL"
+fi
+
 # =============================================================================
 # Engine-specific environment
 # =============================================================================
@@ -255,6 +290,24 @@ else
     export MORI_MAX_DISPATCH_TOKENS_PREFILL=8192
     export MORI_MAX_DISPATCH_TOKENS_DECODE=512
 
+    # DeepSeek-R1-0528 on MI325X only: 512 undersizes the decode MoRI MoE dispatch
+    # buffer for the conc-32 DP+EP cross-node all-to-all, which stalls under load.
+    # Scoped to this model+SKU because the value also scales the decode
+    # --chunked-prefill-size (via the models.yaml formula) and the inter-kernel
+    # switch threshold below, so it must not move for other SKUs or models.
+    #
+    # RUNNER_FAMILY is the GPU family derived from RUNNER_TYPE at the top of this
+    # file. configs/*-master.yaml declares the label per entry ("runner: mi325x" or
+    # "runner: cluster:mi325x-amds"), job.slurm forwards RUNNER_TYPE into the
+    # container, and benchmark_lib.sh records it as the result "hw" field.
+    # Deliberately not RUNNER_NAME: that is the GitHub runner instance
+    # (mi325x-amds_00) used host-side to pick the launcher, and it is never
+    # forwarded into the container.
+    if [[ "$MODEL_NAME" == "DeepSeek-R1-0528" ]] && [[ "$RUNNER_FAMILY" == "mi325x" ]]; then
+        export MORI_MAX_DISPATCH_TOKENS_DECODE=4096
+        echo "[INFO] $RUNNER_TYPE + $MODEL_NAME: MORI_MAX_DISPATCH_TOKENS_DECODE=$MORI_MAX_DISPATCH_TOKENS_DECODE (conc-32 DP+EP cross-node dispatch buffer)"
+    fi
+
     export MORI_MOE_MAX_INPUT_TOKENS_PREFILL=32768
     export MORI_MOE_MAX_INPUT_TOKENS_DECODE=2703
 

benchmarks/multi_node/amd_utils/server_sglang.sh

@@ -189,6 +189,16 @@ print(f'DECODE_CUDA_GRAPH_BS_NO_DP_END=\"{e}\"')
 
 echo "Loaded model configuration for: $MODEL_NAME"
 
+# DeepSeek-R1-0528 on MI325X only: raised for the MI325X PD disagg tests (#22072).
+# Kept out of models.yaml because that entry is shared by every AMD SKU, and a higher
+# static fraction is the riskier setting on the smaller-HBM parts (MI300X 192GB vs
+# MI325X 256GB). Must be applied here rather than in env.sh: the models.yaml loader
+# above hard-assigns PREFILL_MEM_FRACTION_STATIC, clobbering anything env.sh sets.
+if [[ "$MODEL_NAME" == "DeepSeek-R1-0528" ]] && [[ "$RUNNER_FAMILY" == "mi325x" ]]; then
+    PREFILL_MEM_FRACTION_STATIC=0.9
+    echo "[INFO] $RUNNER_TYPE + $MODEL_NAME: PREFILL_MEM_FRACTION_STATIC=$PREFILL_MEM_FRACTION_STATIC (overrides the models.yaml 0.8)"
+fi
+
 # Compute DP-dependent prefill parameters
 if [[ "$PREFILL_ENABLE_DP" == "true" ]]; then
     prefill_cuda_graph_bs=($PREFILL_CUDA_GRAPH_BS_DP)
@@ -240,14 +250,24 @@ fi
 
 # When both DP and EP are enabled, override max-running-requests and dispatch tokens
 if [[ "$DECODE_ENABLE_DP" == "true" ]] && [[ "$DECODE_ENABLE_EP" == "true" ]]; then
-    decode_max_running_requests=$BENCH_MAX_CONC_VALUE
     decode_dp_ranks=$DECODE_TP_SIZE
-    MORI_MAX_DISPATCH_TOKENS_DECODE=$((BENCH_MAX_CONC_VALUE / decode_dp_ranks))
-    # MORI_MOE_MAX_INPUT_TOKENS_DECODE=$((MORI_MAX_DISPATCH_TOKENS_DECODE * decode_dp_ranks * 7 / 10))
-    # Update derived variable
-    SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD=$((MORI_MAX_DISPATCH_TOKENS_DECODE * 2))
-    export SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD
-    echo "[DP+EP override] Decode: max-running-requests=$decode_max_running_requests, DISPATCH_TOKENS=$MORI_MAX_DISPATCH_TOKENS_DECODE, MOE_MAX_INPUT=$MORI_MOE_MAX_INPUT_TOKENS_DECODE, INTER_KERNEL_SWITCH=$SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD"
+    # max-running-requests is split across DP ranks (sglang's req_to_token_pool is
+    # per-rank = max_running_requests // dp_ranks). It must be >= dp_ranks, else the
+    # per-rank pool floors to 0 and get_batch_sizes_to_capture collapses capture_bs
+    # to [0] (AssertionError). This happens when bench concurrency < dp_ranks.
+    # Old heuristic set max_running=BENCH_MAX_CONC (=conc) and then shrank the MoRI
+    # dispatch tokens to max_running/dp_ranks (=2 at conc32/DP16). That gives each DP
+    # rank only ~2 request slots and a 2-token MoE all-to-all dispatch buffer, which
+    # starves the cross-node EP path under load (TTFT ~11s, mass timeouts). The patch
+    # reference instead keeps a generous per-rank pool (conc*TP) and the env.sh MoRI
+    # dispatch / MOE token sizes (4096 / 2703). Floor at dp_ranks to keep capture_bs>0.
+    decode_max_running_requests=$((BENCH_MAX_CONC_VALUE * decode_dp_ranks))
+    if (( decode_max_running_requests < decode_dp_ranks )); then
+        decode_max_running_requests=$decode_dp_ranks
+    fi
+    # Keep MORI_MAX_DISPATCH_TOKENS_DECODE / MOE_MAX_INPUT / INTER_KERNEL_SWITCH from
+    # env.sh (do NOT shrink them to max_running/dp_ranks).
+    echo "[DP+EP override] Decode: max-running-requests=$decode_max_running_requests, DISPATCH_TOKENS=$MORI_MAX_DISPATCH_TOKENS_DECODE (env), MOE_MAX_INPUT=$MORI_MOE_MAX_INPUT_TOKENS_DECODE (env)"
 fi
 
 # Build the composed config strings (equivalent to the old MODEL_PREFILL_CONFIGS / MODEL_DECODE_CONFIGS)