|
| 1 | +#!/usr/bin/env bash |
| 2 | +set -eo pipefail |
| 3 | +set -x |
| 4 | + |
| 5 | +# Agentic trace replay benchmark for DeepSeek-V4-Pro FP4 on MI355X using SGLang. |
| 6 | +# |
| 7 | +# KV_OFFLOADING=dram requires KV_OFFLOAD_BACKEND=hicache. |
| 8 | +# |
| 9 | +# Required env vars: |
| 10 | +# MODEL, TP, CONC, KV_OFFLOADING, TOTAL_CPU_DRAM_GB, RESULT_DIR |
| 11 | +# |
| 12 | +# KV_OFFLOADING=dram requires one of these. |
| 13 | +# KV_OFFLOAD_BACKEND=hicache. |
| 14 | + |
| 15 | +source "$(dirname "$0")/../../benchmark_lib.sh" |
| 16 | + |
| 17 | +check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE DP_ATTENTION |
| 18 | + |
| 19 | +if [[ -n "$SLURM_JOB_ID" ]]; then |
| 20 | + echo "JOB $SLURM_JOB_ID running on $SLURMD_NODENAME" |
| 21 | +fi |
| 22 | + |
| 23 | +# ROCR/HIP visibility under slurm cgroups. |
| 24 | +if [ -n "$ROCR_VISIBLE_DEVICES" ]; then |
| 25 | + export HIP_VISIBLE_DEVICES="$ROCR_VISIBLE_DEVICES" |
| 26 | +fi |
| 27 | + |
| 28 | +if [[ -n "$MODEL_PATH" ]]; then |
| 29 | + if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then |
| 30 | + hf download "$MODEL" --local-dir "$MODEL_PATH" |
| 31 | + fi |
| 32 | +else |
| 33 | + hf download "$MODEL" |
| 34 | + export MODEL_PATH="$MODEL" |
| 35 | +fi |
| 36 | +rocm-smi || true |
| 37 | +amd-smi || true |
| 38 | + |
| 39 | +# ---- Resolve traces and install deps ---------------------------------------- |
| 40 | +resolve_trace_source |
| 41 | +install_agentic_deps |
| 42 | + |
| 43 | +# ---- Server config ---------------------------------------------------------- |
| 44 | +SERVER_LOG="$RESULT_DIR/server.log" |
| 45 | +mkdir -p "$RESULT_DIR" |
| 46 | + |
| 47 | +CACHE_ARGS=() |
| 48 | +if agentic_kv_offload_enabled; then |
| 49 | + # HiCache config — https://lmsysorg.mintlify.app/cookbook/autoregressive/DeepSeek/DeepSeek-V4 |
| 50 | + case "$KV_OFFLOAD_BACKEND" in |
| 51 | + hicache) |
| 52 | + HICACHE_RATIO=4 |
| 53 | + HICACHE_WRITE_POLICY="write_through" |
| 54 | + HICACHE_IO_BACKEND="direct" |
| 55 | + HICACHE_MEM_LAYOUT="page_first_direct" |
| 56 | + CACHE_ARGS=( |
| 57 | + --enable-hierarchical-cache |
| 58 | + --hicache-ratio "$HICACHE_RATIO" |
| 59 | + --hicache-write-policy "$HICACHE_WRITE_POLICY" |
| 60 | + --hicache-io-backend "$HICACHE_IO_BACKEND" |
| 61 | + --hicache-mem-layout "$HICACHE_MEM_LAYOUT" |
| 62 | + ) |
| 63 | + echo "HiCache DSv4 CPU tier: ratio=$HICACHE_RATIO, write_policy=$HICACHE_WRITE_POLICY, io_backend=$HICACHE_IO_BACKEND, mem_layout=$HICACHE_MEM_LAYOUT" |
| 64 | + ;; |
| 65 | + *) |
| 66 | + echo "Error: unsupported KV_OFFLOAD_BACKEND '$KV_OFFLOAD_BACKEND' (expected: hicache)" >&2 |
| 67 | + exit 1 |
| 68 | + ;; |
| 69 | + esac |
| 70 | +fi |
| 71 | +# ---- Client config ---------------------------------------------------------- |
| 72 | +export AIPERF_HTTP_TCP_USER_TIMEOUT=1000000 |
| 73 | + |
| 74 | +# ---- LLM server config ---------------------------------------------------------- |
| 75 | +USE_SGLANG_ROUTER=false |
| 76 | +SGLANG_BACKEND_PORT="$PORT" |
| 77 | +ROUTER_LOG="$RESULT_DIR/router.log" |
| 78 | +MEM_FRACTION_STATIC=0.90 |
| 79 | +CHUNKED_PREFILL_SIZE=8192 |
| 80 | +PARALLEL_ARGS=(--tensor-parallel-size "$TP") |
| 81 | +if [ "$DP_ATTENTION" = "true" ]; then |
| 82 | + USE_SGLANG_ROUTER=true |
| 83 | + export AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID=true |
| 84 | + SGLANG_BACKEND_PORT=$((PORT + 1)) |
| 85 | + SGLANG_ROUTER_METRICS_PORT=$((PORT + 10000)) |
| 86 | + SGLANG_ROUTER_CMD=(python3 -m sglang_router.launch_router) |
| 87 | + |
| 88 | + export SGLANG_SHARED_EXPERT_TP1=1 |
| 89 | + export SGLANG_DP_SHARED_EXPERT_LOCAL=1 |
| 90 | + export SGLANG_DP_USE_GATHERV=1 |
| 91 | + export SGLANG_DP_USE_REDUCE_SCATTER=1 |
| 92 | + export GPU_MAX_HW_QUEUES=5 |
| 93 | + |
| 94 | + CHUNKED_PREFILL_SIZE=$((8192 * TP)) |
| 95 | + PARALLEL_ARGS+=( |
| 96 | + --dp "$TP" |
| 97 | + --enable-dp-attention |
| 98 | + --enable-prefill-delayer |
| 99 | + ) |
| 100 | +fi |
| 101 | + |
| 102 | +if [ "$EP_SIZE" -gt 1 ]; then |
| 103 | + PARALLEL_ARGS+=(--ep-size "$EP_SIZE") |
| 104 | +fi |
| 105 | + |
| 106 | +# SGLang treats max-running-requests as a global DPA limit and partitions it |
| 107 | +# internally. CUDA graph capture is per scheduler, so only its batch size is |
| 108 | +# divided across DP ranks. |
| 109 | +MAX_RUNNING_REQUESTS=$((2 * CONC)) |
| 110 | +CUDA_GRAPH_MAX_BS=$CONC |
| 111 | +[ "$CUDA_GRAPH_MAX_BS" -gt 128 ] && CUDA_GRAPH_MAX_BS=128 |
| 112 | + |
| 113 | +# Simulated acceptance-length (AL) settings. |
| 114 | +export SGLANG_DEFAULT_THINKING=1 |
| 115 | +export SGLANG_DSV4_REASONING_EFFORT=high |
| 116 | +export SGLANG_SIMULATE_ACC_LEN=2.49 |
| 117 | +export SGLANG_SIMULATE_ACC_METHOD=match-expected |
| 118 | +export SGLANG_SIMULATE_ACC_TOKEN_MODE=real-draft-token |
| 119 | + |
| 120 | +export SGLANG_USE_ROCM700A=0 |
| 121 | +export SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton |
| 122 | +export AITER_BF16_FP8_MOE_BOUND=0 |
| 123 | + |
| 124 | +export SGLANG_ENABLE_UNIFIED_RADIX_TREE=1 |
| 125 | +export SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS=1 |
| 126 | + |
| 127 | +METRICS_ARGS=(--enable-metrics) |
| 128 | +SPEC_ARGS=( |
| 129 | + --speculative-algorithm EAGLE |
| 130 | + --speculative-num-steps 3 |
| 131 | + --speculative-eagle-topk 1 |
| 132 | + --speculative-num-draft-tokens 4 |
| 133 | +) |
| 134 | + |
| 135 | +if [ ${#SPEC_ARGS[@]} -gt 0 ]; then |
| 136 | + MEM_FRACTION_STATIC=$(awk "BEGIN {printf \"%.2f\", $MEM_FRACTION_STATIC - 0.10}") |
| 137 | +fi |
| 138 | + |
| 139 | +SGLANG_CMD=( |
| 140 | + python3 -m sglang.launch_server |
| 141 | + --model-path "$MODEL_PATH" |
| 142 | + --served-model-name "$MODEL" |
| 143 | + --host 0.0.0.0 |
| 144 | + --port "$SGLANG_BACKEND_PORT" |
| 145 | + --trust-remote-code |
| 146 | + "${PARALLEL_ARGS[@]}" |
| 147 | + --attention-backend compressed |
| 148 | + --cuda-graph-max-bs-decode "$CUDA_GRAPH_MAX_BS" |
| 149 | + --max-running-requests "$MAX_RUNNING_REQUESTS" |
| 150 | + --mem-fraction-static "$MEM_FRACTION_STATIC" |
| 151 | + --swa-full-tokens-ratio 0.10 |
| 152 | + --page-size 256 |
| 153 | + --kv-cache-dtype fp8_e4m3 |
| 154 | + --chunked-prefill-size "$CHUNKED_PREFILL_SIZE" |
| 155 | + --disable-shared-experts-fusion |
| 156 | + --tool-call-parser deepseekv4 |
| 157 | + --reasoning-parser deepseek-v4 |
| 158 | + --chat-template "$(dirname "$0")/../chat_templates/deepseek_v4_thinking.jinja" |
| 159 | + --watchdog-timeout 1800 |
| 160 | + "${METRICS_ARGS[@]}" |
| 161 | + "${SPEC_ARGS[@]}" |
| 162 | + "${CACHE_ARGS[@]}" |
| 163 | +) |
| 164 | + |
| 165 | +printf '%q ' "${SGLANG_CMD[@]}" | tee "$RESULT_DIR/sglang_command.txt" |
| 166 | +printf '\n' | tee -a "$RESULT_DIR/sglang_command.txt" |
| 167 | + |
| 168 | +{ |
| 169 | + echo "=== SGLANG_* env vars at launch ===" |
| 170 | + env | grep -E '^SGLANG_' | sort |
| 171 | + echo "===================================" |
| 172 | +} | tee "$SERVER_LOG" |
| 173 | + |
| 174 | +echo "Starting SGLang server for MI355X..." |
| 175 | +"${SGLANG_CMD[@]}" >> "$SERVER_LOG" 2>&1 & |
| 176 | +SERVER_PID=$! |
| 177 | +echo "Server PID: $SERVER_PID" |
| 178 | + |
| 179 | +wait_for_server_ready --port "$SGLANG_BACKEND_PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID" |
| 180 | + |
| 181 | +if [ "$USE_SGLANG_ROUTER" = "true" ]; then |
| 182 | + echo "Starting SGLang router on port $PORT for $TP DP ranks..." |
| 183 | + "${SGLANG_ROUTER_CMD[@]}" \ |
| 184 | + --worker-urls "http://localhost:$SGLANG_BACKEND_PORT" \ |
| 185 | + --policy consistent_hashing \ |
| 186 | + --request-id-headers x-correlation-id \ |
| 187 | + --dp-aware \ |
| 188 | + --host 0.0.0.0 \ |
| 189 | + --port "$PORT" \ |
| 190 | + --prometheus-host 127.0.0.1 \ |
| 191 | + --prometheus-port "$SGLANG_ROUTER_METRICS_PORT" \ |
| 192 | + --connect-timeout-secs 900 \ |
| 193 | + --request-timeout-secs 14400 \ |
| 194 | + --disable-health-check \ |
| 195 | + --disable-retries > "$ROUTER_LOG" 2>&1 & |
| 196 | + ROUTER_PID=$! |
| 197 | + echo "Router PID: $ROUTER_PID" |
| 198 | + wait_for_server_ready --port "$PORT" --server-log "$ROUTER_LOG" --server-pid "$ROUTER_PID" |
| 199 | +fi |
| 200 | + |
| 201 | +# ---- Run benchmark ---------------------------------------------------------- |
| 202 | +build_replay_cmd "$RESULT_DIR" |
| 203 | +REPLAY_CMD+=" --server-metrics http://localhost:$SGLANG_BACKEND_PORT/metrics" |
| 204 | + |
| 205 | +run_agentic_replay_and_write_outputs "$RESULT_DIR" |
0 commit comments