fix(amd): preserve legacy AgentX fabric semantics

Author: cquil11

benchmarks/multi_node/agentic_srt.sh

@@ -42,6 +42,60 @@ done
 resolve_trace_source
 install_agentic_deps
 
+# Preserve the legacy DP-attention replay contract. The SGLang router uses this
+# header to keep every request in one AgentX correlation tree on a stable DP
+# route, which is important for both session continuity and prefix-cache reuse.
+if [[ "${PREFILL_DP_ATTN:-false}" == "true" ]]; then
+    export AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID=true
+fi
+
+# Reset every advertised SGLang worker before each concurrency point, matching
+# the retired amd_utils trace replay. /flush_cache covers GPU radix + host
+# HiCache; the storage-backend endpoint is best-effort because L3 is optional.
+clear_agentic_worker_caches() {
+    local timeout_seconds="${FLUSH_DRAIN_TIMEOUT:-120}"
+    local metrics_csv="${AIPERF_SERVER_METRICS_URLS:-}"
+    if [[ -z "$metrics_csv" ]]; then
+        echo "[clear_caches] WARN: AIPERF_SERVER_METRICS_URLS unset; skipping cache flush" >&2
+        return 0
+    fi
+
+    local -a metrics_urls
+    IFS=',' read -r -a metrics_urls <<< "$metrics_csv"
+    local metrics_url base_url start response code flushed
+    for metrics_url in "${metrics_urls[@]}"; do
+        [[ -n "$metrics_url" ]] || continue
+        base_url="${metrics_url%/metrics}"
+        start=$(date +%s)
+        flushed=0
+        response=""
+        while :; do
+            response=$(curl -sf -m 10 -X POST "${base_url}/flush_cache" 2>/dev/null || true)
+            if grep -qi "Cache flushed" <<< "$response"; then
+                flushed=1
+                break
+            fi
+            if (( $(date +%s) - start >= timeout_seconds )); then
+                break
+            fi
+            sleep 3
+        done
+        if (( flushed )); then
+            echo "[clear_caches] ${base_url}: L1+L2 flushed"
+        else
+            echo "[clear_caches] WARN ${base_url}: L1+L2 flush not confirmed after ${timeout_seconds}s" >&2
+        fi
+
+        code=$(curl -s -m 60 -o /dev/null -w '%{http_code}' \
+            -X POST "${base_url}/hicache/storage-backend/clear" 2>/dev/null || true)
+        if [[ "$code" == "200" ]]; then
+            echo "[clear_caches] ${base_url}: L3 store cleared"
+        else
+            echo "[clear_caches] ${base_url}: L3 clear http=${code:-000} (optional backend unavailable)"
+        fi
+    done
+}
+
 wait_for_agentic_servers_idle() {
     local timeout_seconds="${AIPERF_DRAIN_TIMEOUT_SECONDS:-1800}"
     local poll_seconds="${AIPERF_DRAIN_POLL_SECONDS:-10}"
@@ -125,6 +179,9 @@ for index in "${!CONCURRENCIES[@]}"; do
     mkdir -p "$RESULT_DIR"
 
     echo "Running agentic concurrency $concurrency of: ${CONCURRENCIES[*]}"
+    if [[ "${CLEAR_CACHE_BETWEEN_CONC:-1}" == "1" ]]; then
+        clear_agentic_worker_caches
+    fi
     build_replay_cmd "$RESULT_DIR"
     run_agentic_replay_and_write_outputs "$RESULT_DIR"
 

perf-changelog.yaml

@@ -6670,4 +6670,5 @@
     - "Add small-model aggregate and 1P1D disaggregated correctness lanes for the shared MI355X srt-slurm launcher."
     - "Use the native SGLang router for request routing and AMD MoRI for disaggregated KV transfer, without Dynamo, NIXL, or the deleted amd_utils control plane."
     - "Remove the active and deprecated MI355X matrix entries that depended on amd_utils; configuration-specific child PRs reintroduce the active production lanes as native srt-slurm recipes."
+    - "Preserve the retired AgentX launch contract in the shared srt path: route DP-attention sessions by correlation ID, flush every advertised worker cache before each concurrency point, and run the MI355X RDMA device/QoS/DCQCN plus bare-process GPU drain preflights on every allocated host."
   pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2542

runners/check_mi355x_rdma.sh

@@ -0,0 +1,52 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Fast, per-node fabric preflight for MI355X srt-slurm allocations. This keeps
+# the meaningful QoS/DCQCN gate from the retired amd_utils launcher without its
+# Docker or job-control plumbing.
+
+log() { printf '[%s] %s\n' "$(hostname -s)" "$*"; }
+fail() { log "RDMA preflight failed: $*" >&2; exit 1; }
+
+expected_devices="${IBDEVICES:-rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7}"
+IFS=',' read -r -a devices <<< "$expected_devices"
+for device in "${devices[@]}"; do
+    [[ -d "/sys/class/infiniband/${device}" ]] || fail "missing device ${device}"
+done
+log "found all ${#devices[@]} expected RDMA devices: ${expected_devices}"
+
+if ! command -v nicctl >/dev/null 2>&1; then
+    log "nicctl is unavailable; device presence passed, QoS/DCQCN checks skipped"
+    exit 0
+fi
+
+probe=$(sudo -n nicctl show version firmware 2>&1 || true)
+if grep -qiE 'No AMD NICs|Invalid card handle|Failed to get NIC' <<< "$probe"; then
+    fail "nicctl cannot access the AMD NICs"
+fi
+
+qos=$(sudo -n nicctl show qos 2>/dev/null) || fail "nicctl show qos failed"
+classification=$(awk '/Classification type/ {print $NF; exit}' <<< "$qos")
+[[ "$classification" == "DSCP" ]] || fail "classification is ${classification:-unset}, expected DSCP"
+
+priorities=$(awk '/PFC no-drop priorities/ {print $NF; exit}' <<< "$qos")
+bitmap=$(awk '/PFC priority bitmap/ {print $NF; exit}' <<< "$qos")
+[[ -n "$priorities" ]] || fail "PFC no-drop priorities are missing"
+[[ -n "$bitmap" && "$bitmap" != "0x0" ]] || fail "PFC is disabled"
+IFS=',' read -r -a priority_values <<< "$priorities"
+for priority in "${priority_values[@]}"; do
+    priority="${priority//[^0-9]/}"
+    [[ -n "$priority" ]] || fail "invalid PFC priority list: ${priorities}"
+    (( bitmap & (1 << priority) )) || fail "PFC bitmap ${bitmap} does not cover priority ${priority}"
+done
+
+dcqcn=$(sudo -n nicctl show dcqcn 2>/dev/null) || fail "nicctl show dcqcn failed"
+device_count=$(grep -c 'ROCE device' <<< "$dcqcn" || true)
+(( device_count > 0 )) || fail "no RoCE devices reported by nicctl"
+if grep 'Status' <<< "$dcqcn" | grep -qv 'Enabled'; then
+    fail "DCQCN is disabled on at least one RoCE device"
+fi
+cnp_count=$(awk '/DSCP value used for CNP/ {print $NF}' <<< "$dcqcn" | sort -u | grep -c . || true)
+(( cnp_count == 1 )) || fail "CNP DSCP is inconsistent across NICs"
+
+log "RDMA QoS/DCQCN preflight passed"

runners/launch_mi355x-amds-srt.sh

@@ -249,6 +249,14 @@ if (
     decode_environment.setdefault(
         "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK", str(dispatch_tokens)
     )
+    # The retired launcher also exposed its harness-level dispatch budget to
+    # the server environment after scaling it by the MTP draft width. Keep that
+    # auxiliary value distinct from the model-specific per-rank SGLang pin.
+    if "MORI_MAX_DISPATCH_TOKENS_DECODE" in decode_environment:
+        mtp_size = int(os.environ.get("DECODE_MTP_SIZE", "0"))
+        decode_environment["MORI_MAX_DISPATCH_TOKENS_DECODE"] = str(
+            dispatch_tokens * (mtp_size + 1)
+        )
     decode_environment["SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD"] = str(
         2 * dispatch_tokens
     )

runners/prepare_mi355x_srt_host.sh

@@ -1,6 +1,18 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
+"$(dirname "${BASH_SOURCE[0]}")/check_mi355x_rdma.sh"
+
+# Preserve the legacy bare-process GPU drain gate. Slurm owns the node, but a
+# process left outside the prior job's container can still retain VRAM and make
+# the next model load fail much later with a misleading OOM.
+(
+    unset IS_AGENTIC SCENARIO_TYPE
+    # shellcheck source=benchmarks/benchmark_lib.sh
+    source "$(dirname "${BASH_SOURCE[0]}")/../benchmarks/benchmark_lib.sh"
+    wait_for_amd_gpu_clean
+)
+
 # Some MI355X experiments reserve large 2 MiB HugeTLB pools and leave the
 # reservation behind after their Slurm allocation exits. Those free hugepages
 # remain unavailable to ordinary host allocations, which can make a later

utils/evals/test_batched_eval.py

@@ -272,4 +272,5 @@ def test_amd_srt_launcher_preserves_mori_dispatch_pin() -> None:
     assert 'decode_environment.setdefault(' in launcher
     assert '"SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK", str(dispatch_tokens)' in launcher
     assert '"SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD"' in launcher
-    assert "mtp_size + 1" not in launcher
+    assert 'if "MORI_MAX_DISPATCH_TOKENS_DECODE" in decode_environment:' in launcher
+    assert "dispatch_tokens * (mtp_size + 1)" in launcher