|
| 1 | +#!/usr/bin/env bash |
| 2 | +set -euo pipefail |
| 3 | +set -x |
| 4 | + |
| 5 | +# Agentic trace replay benchmark for GLM-5.2 NVFP4 on B300 using SGLang. |
| 6 | +# |
| 7 | +# Server flags follow the SGLang cookbook B300 NVFP4 single-node recipes |
| 8 | +# (https://docs.sglang.io/cookbook/autoregressive/GLM/GLM-5.2), STP only: |
| 9 | +# the cookbook's EAGLE MTP variants are intentionally not wired up yet. |
| 10 | +# DP_ATTENTION=false -> low-latency arm (TP8, fp8 KV, cutedsl bf16 GEMM) |
| 11 | +# DP_ATTENTION=true -> high-throughput arm (TP8 + DP8 attention-DP) |
| 12 | +# |
| 13 | +# Required env vars: |
| 14 | +# MODEL, TP, CONC, KV_OFFLOADING, TOTAL_CPU_DRAM_GB, RESULT_DIR, DURATION, |
| 15 | +# EP_SIZE, DP_ATTENTION |
| 16 | + |
| 17 | +source "$(dirname "$0")/../../benchmark_lib.sh" |
| 18 | + |
| 19 | +check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE DP_ATTENTION |
| 20 | + |
| 21 | +if [[ "$KV_OFFLOADING" != "none" ]]; then |
| 22 | + echo "Error: KV_OFFLOADING=$KV_OFFLOADING is not supported by this recipe" >&2 |
| 23 | + exit 1 |
| 24 | +fi |
| 25 | + |
| 26 | +if [[ -n "${SLURM_JOB_ID:-}" ]]; then |
| 27 | + echo "JOB $SLURM_JOB_ID running on ${SLURMD_NODENAME:-unknown}" |
| 28 | +fi |
| 29 | + |
| 30 | +# `hf download` creates the target dir if missing and is itself idempotent. |
| 31 | +# When MODEL_PATH is unset (stand-alone runs), fall back to the HF_HUB_CACHE. |
| 32 | +# Either way, MODEL_PATH is what the server is launched with. |
| 33 | +if [[ -n "${MODEL_PATH:-}" ]]; then |
| 34 | + if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then |
| 35 | + hf download "$MODEL" --local-dir "$MODEL_PATH" |
| 36 | + fi |
| 37 | +else |
| 38 | + hf download "$MODEL" |
| 39 | + export MODEL_PATH="$MODEL" |
| 40 | +fi |
| 41 | +nvidia-smi |
| 42 | + |
| 43 | +resolve_trace_source |
| 44 | +install_agentic_deps |
| 45 | + |
| 46 | +SERVER_LOG="$RESULT_DIR/server.log" |
| 47 | +mkdir -p "$RESULT_DIR" |
| 48 | + |
| 49 | +# With attention-DP, front the DP ranks with sglang-router using consistent |
| 50 | +# hashing on the AIPerf correlation id so multi-turn sessions stay on the DP |
| 51 | +# rank that holds their radix-cache prefix. |
| 52 | +USE_SGLANG_ROUTER=false |
| 53 | +SGLANG_BACKEND_PORT="$PORT" |
| 54 | +ROUTER_LOG="$RESULT_DIR/router.log" |
| 55 | +if [ "$DP_ATTENTION" = "true" ]; then |
| 56 | + USE_SGLANG_ROUTER=true |
| 57 | + export AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID=true |
| 58 | + SGLANG_BACKEND_PORT=$((PORT + 1)) |
| 59 | + SGLANG_ROUTER_METRICS_PORT=$((PORT + 10000)) |
| 60 | +fi |
| 61 | + |
| 62 | +PARALLEL_ARGS=(--tp "$TP" --ep-size "$EP_SIZE") |
| 63 | +CHUNKED_PREFILL_SIZE=8192 |
| 64 | +if [ "$DP_ATTENTION" = "true" ]; then |
| 65 | + # chunked-prefill-size is a whole-engine budget split across DP ranks: |
| 66 | + # the cookbook HT cell's 8192 becomes 1,024 tokens/rank/step under dp8, |
| 67 | + # which starves prefill on the 1M-context agentic corpus (observed: a |
| 68 | + # conc-256 warmup could not drain within AIPerf's 1800s grace period |
| 69 | + # while KV usage sat at ~0.01). Use the cookbook's own dp8 lever from |
| 70 | + # the B200 cells (32768 = ~4096/rank). |
| 71 | + CHUNKED_PREFILL_SIZE=32768 |
| 72 | + # At conc 512 the saturation working set outlives the default 1800s |
| 73 | + # warmup drain grace: the drain converges healthily (~0.45 req/s, zero |
| 74 | + # errors) but needs ~2500s end to end. 3600 is a maximum wait, not a |
| 75 | + # fixed sleep — lower-conc DPA points still finish as fast as they drain. |
| 76 | + export AGENTIC_WARMUP_GRACE_PERIOD=3600 |
| 77 | + PARALLEL_ARGS+=( |
| 78 | + --dp "$TP" |
| 79 | + --enable-dp-attention |
| 80 | + --tokenizer-worker-num "$TP" |
| 81 | + --dist-init-addr "127.0.0.1:$((PORT + 2000))" |
| 82 | + ) |
| 83 | +else |
| 84 | + # Cookbook low-latency levers; the DP-attention cell omits them. |
| 85 | + PARALLEL_ARGS+=( |
| 86 | + --kv-cache-dtype fp8_e4m3 |
| 87 | + --bf16-gemm-backend cutedsl |
| 88 | + --max-prefill-tokens 8192 |
| 89 | + ) |
| 90 | +fi |
| 91 | + |
| 92 | +# AgentX concurrency counts live session trees, not individual requests. |
| 93 | +# Allow subagent fan-out to exceed CONC without clipping request bursts. |
| 94 | +MAX_RUNNING_REQUESTS=$((2 * CONC)) |
| 95 | +GRAPH_ARGS=() |
| 96 | +if [ "$DP_ATTENTION" != "true" ]; then |
| 97 | + # Cookbook low-latency captures graphs up to its request cap; the |
| 98 | + # DP-attention cell leaves the CUDA-graph batch list at SGLang defaults. |
| 99 | + CUDA_GRAPH_MAX_BS=$MAX_RUNNING_REQUESTS |
| 100 | + [ "$CUDA_GRAPH_MAX_BS" -gt 64 ] && CUDA_GRAPH_MAX_BS=64 |
| 101 | + GRAPH_ARGS=(--cuda-graph-max-bs "$CUDA_GRAPH_MAX_BS") |
| 102 | +fi |
| 103 | + |
| 104 | +export PYTHONNOUSERSITE=1 |
| 105 | +export TORCH_CUDA_ARCH_LIST=10.0 |
| 106 | +# Agentic warmup dispatches hundreds of large prompts at once; allow up to |
| 107 | +# 15 minutes of TCP progress before AIPerf declares a connection dead. |
| 108 | +export AIPERF_HTTP_TCP_USER_TIMEOUT=900000 |
| 109 | +# AIPerf pins one pooled keep-alive connection per session (client-side |
| 110 | +# keep-alive 300s) while uvicorn's default SGLANG_TIMEOUT_KEEP_ALIVE is 5s; |
| 111 | +# inter-turn idle gaps (capped at 10s) can reuse a socket exactly as the |
| 112 | +# server closes it -> ECONNRESET -> terminal warmup failure. Outlast the |
| 113 | +# client pool so the race cannot occur. |
| 114 | +export SGLANG_TIMEOUT_KEEP_ALIVE=900 |
| 115 | + |
| 116 | +SGLANG_CMD=( |
| 117 | + python3 -m sglang.launch_server |
| 118 | + --model-path "$MODEL_PATH" |
| 119 | + --served-model-name "$MODEL" |
| 120 | + --host 0.0.0.0 |
| 121 | + --port "$SGLANG_BACKEND_PORT" |
| 122 | + --trust-remote-code |
| 123 | + "${PARALLEL_ARGS[@]}" |
| 124 | + --quantization modelopt_fp4 |
| 125 | + # GLM-5.2 emits the GLM-4.7-style <tool_call>/<arg_key>/<arg_value> format; |
| 126 | + # the glm47 parser is required for structured message.tool_calls (glm45 |
| 127 | + # leaves calls as raw text). Without it the SWE-bench mini-swe-agent eval |
| 128 | + # dies with RepeatedFormatError ("No tool calls found in the response") on |
| 129 | + # every instance and scores 0. Reasoning parser keeps hybrid-thinking |
| 130 | + # output in reasoning_content instead of polluting content. Neither flag |
| 131 | + # affects trace-replay throughput (pre-canned replay discards live |
| 132 | + # responses). |
| 133 | + --tool-call-parser glm47 |
| 134 | + --reasoning-parser glm45 |
| 135 | + --chunked-prefill-size "$CHUNKED_PREFILL_SIZE" |
| 136 | + --mem-fraction-static 0.85 |
| 137 | + --max-running-requests "$MAX_RUNNING_REQUESTS" |
| 138 | + "${GRAPH_ARGS[@]}" |
| 139 | + --watchdog-timeout 1800 |
| 140 | + --enable-metrics |
| 141 | +) |
| 142 | + |
| 143 | +printf '%q ' "${SGLANG_CMD[@]}" | tee "$RESULT_DIR/sglang_command.txt" |
| 144 | +printf '\n' | tee -a "$RESULT_DIR/sglang_command.txt" |
| 145 | + |
| 146 | +echo "Starting SGLang server for B300..." |
| 147 | +"${SGLANG_CMD[@]}" > "$SERVER_LOG" 2>&1 & |
| 148 | +SERVER_PID=$! |
| 149 | +echo "Server PID: $SERVER_PID" |
| 150 | + |
| 151 | +wait_for_server_ready --port "$SGLANG_BACKEND_PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID" |
| 152 | + |
| 153 | +if [ "$USE_SGLANG_ROUTER" = "true" ]; then |
| 154 | + echo "Starting SGLang router on port $PORT for $TP DP ranks..." |
| 155 | + python3 -m sglang_router.launch_router \ |
| 156 | + --worker-urls "http://localhost:$SGLANG_BACKEND_PORT" \ |
| 157 | + --policy consistent_hashing \ |
| 158 | + --request-id-headers x-correlation-id \ |
| 159 | + --dp-aware \ |
| 160 | + --host 0.0.0.0 \ |
| 161 | + --port "$PORT" \ |
| 162 | + --prometheus-host 127.0.0.1 \ |
| 163 | + --prometheus-port "$SGLANG_ROUTER_METRICS_PORT" \ |
| 164 | + --connect-timeout-secs 900 \ |
| 165 | + --request-timeout-secs 14400 \ |
| 166 | + --disable-health-check \ |
| 167 | + --disable-retries > "$ROUTER_LOG" 2>&1 & |
| 168 | + ROUTER_PID=$! |
| 169 | + echo "Router PID: $ROUTER_PID" |
| 170 | + wait_for_server_ready --port "$PORT" --server-log "$ROUTER_LOG" --server-pid "$ROUTER_PID" |
| 171 | +fi |
| 172 | + |
| 173 | +if [ "${EVAL_ONLY}" = "true" ]; then |
| 174 | + # GLM-5.2's chat template defaults to reasoning_effort=Max when the |
| 175 | + # client passes no chat_template_kwargs (mini-swe-agent doesn't), and the |
| 176 | + # heavy thinking burns the default 75-step budget: on the 23-instance |
| 177 | + # slice, 12/23 trajectories exited LimitsExceeded unsubmitted while 10 of |
| 178 | + # the 11 that submitted resolved. Double the step budget for this recipe; |
| 179 | + # other recipes keep the shared 75 default. |
| 180 | + export SWEBENCH_AGENT_STEP_LIMIT=150 |
| 181 | + run_eval --port "$PORT" |
| 182 | +else |
| 183 | + build_replay_cmd "$RESULT_DIR" |
| 184 | + REPLAY_CMD+=" --server-metrics http://localhost:$SGLANG_BACKEND_PORT/metrics" |
| 185 | + run_agentic_replay_and_write_outputs "$RESULT_DIR" |
| 186 | +fi |
0 commit comments