Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
376 changes: 376 additions & 0 deletions benchmarks/single_node/agentic/glm5.2_fp8_b300_sglang_mtp.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,376 @@
#!/usr/bin/env bash
set -euo pipefail
set -x

# Agentic trace replay benchmark for GLM-5.2 FP8 on B300 using SGLang with
# EAGLE/MTP speculative decoding. First GLM-5.2 FP8 AgentX recipe on B300; it
# is spec-decode only, per the AgentX policy that agentic recipes are run and
# published with speculative decoding enabled rather than as an STP/MTP A/B
# (MODELS.md: GLM-5.2 agentic non-MTP is deprecated after 2026-08-03).
#
# Port of agentic/glm5.2_fp8_b200_sglang_mtp.sh (the FP8 B200 sibling, itself a
# port of the validated NVFP4 B200 script) with the B300 deltas the NVFP4 B300
# sibling (agentic/glm5.2_fp4_b300_sglang_mtp.sh) carries, marked "B300:"
# below: --mem-fraction-static 0.85 on the 288 GB part and a fixed 270 GB/rank
# HiCache target pool at every concurrency. The FP8 deltas, marked "FP8:", are
# the checkpoint (zai-org/GLM-5.2-FP8, ~756 GB of block-quantized e4m3 weights
# against ~465 GB for GLM-5.2-NVFP4) and --quantization fp8 in place of
# modelopt_fp4. Serve flags are otherwise identical to both siblings so the
# FP8/NVFP4 curves on B300 and the FP8 curves across B200/B300 stay comparable.
#
# Server flags follow the SGLang cookbook GLM-5.x single-node recipes
# (https://docs.sglang.io/cookbook/autoregressive/GLM/GLM-5.2; the published
# GLM-5.1-FP8 cookbook entry uses the same EAGLE shape with quantization: fp8):
# DP_ATTENTION=false -> low-latency arm (TP8, fp8 KV, cutedsl bf16 GEMM)
# DP_ATTENTION=true -> high-throughput DEP arm (TP8 + DP8 attention-DP +
# EP_SIZE expert-parallel MoE via --ep-size)
# Only the low-latency arm is wired into the master config for this MTP recipe
# (see the entry comment on glm5.2-fp8-b300-sglang-agentic-mtp); the DEP branch
# is kept intact so the throughput arm can be added without re-deriving it.
#
# Required env vars:
# MODEL, TP, CONC, KV_OFFLOADING, TOTAL_CPU_DRAM_GB, RESULT_DIR, DURATION,
# EP_SIZE, DP_ATTENTION
#
# KV_OFFLOADING=dram requires KV_OFFLOAD_BACKEND=hicache.

source "$(dirname "$0")/../../benchmark_lib.sh"

check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE DP_ATTENTION

if [[ -n "${SLURM_JOB_ID:-}" ]]; then
echo "JOB $SLURM_JOB_ID running on ${SLURMD_NODENAME:-unknown}"
fi

# B300: runners/launch_b300-dsxe.sh exports MODEL_PATH as
# <root>/GLM-5.2-FP8, under /scratch/models when the checkpoint is in its
# STAGED_MODELS list and under the writable models dir otherwise. GLM-5.2-FP8
# is not staged on DSXE today, so the day-zero run downloads it there; keep the
# HF repo id separate from MODEL so that download works regardless of how the
# launcher rewrites MODEL.
# FP8: the upstream zai-org release; the golden AL below was measured on it.
HF_MODEL_ID="${HF_MODEL_ID:-zai-org/GLM-5.2-FP8}"

# A non-empty directory is NOT a staged checkpoint. The NVFP4 B200 sibling found
# /lustre/fsw/gharunners/models/GLM-5.2-NVFP4 holding config.json,
# generation_config.json, hf_quant_config.json, chat_template.jinja, README.md
# and .quant_summary.txt and NOTHING else -- an aborted or metadata-only pull.
# An `ls -A` emptiness guard accepts that, so all five cells of run 30729467646
# skipped the download and went straight to serve; SGLang read config.json
# fine, then died in AutoTokenizer.from_pretrained with "Couldn't instantiate
# the backend tokenizer" because neither the tokenizer files nor a single
# weight shard were on disk. Check for a COMPLETE checkpoint instead: the
# tokenizer, the shard index, and every shard the index names.
checkpoint_is_complete() {
local dir="$1"
[[ -d "$dir" ]] || return 1
[[ -f "$dir/tokenizer_config.json" ]] || return 1
[[ -f "$dir/tokenizer.json" || -f "$dir/tokenizer.model" ]] || return 1
[[ -f "$dir/model.safetensors.index.json" ]] || return 1
CKPT_DIR="$dir" python3 - <<'PYEOF'
import json, os, sys
d = os.environ["CKPT_DIR"]
with open(os.path.join(d, "model.safetensors.index.json")) as fh:
shards = sorted(set(json.load(fh)["weight_map"].values()))
missing = [s for s in shards if not os.path.isfile(os.path.join(d, s))]
if missing:
print(f"{len(missing)}/{len(shards)} shards missing, e.g. {missing[:3]}", file=sys.stderr)
sys.exit(1)
PYEOF
}

if [[ -n "${MODEL_PATH:-}" ]]; then
if ! checkpoint_is_complete "$MODEL_PATH"; then
# Every concurrency of this sweep runs as its own allocation against
# the same shared path, so serialize: one cell pulls the ~756 GB
# checkpoint (141 shards) and the rest wait on it rather than five
# racing writers. `hf download` resumes into a partially-populated
# --local-dir, so a metadata-only stub is fine to download on top of.
mkdir -p "$MODEL_PATH"
MODEL_DOWNLOAD_LOCK="${MODEL_PATH%/}.download.lock"
echo "Checkpoint at $MODEL_PATH is incomplete; acquiring $MODEL_DOWNLOAD_LOCK"
exec 9>"$MODEL_DOWNLOAD_LOCK"
flock -w "${MODEL_DOWNLOAD_LOCK_TIMEOUT:-21600}" 9 || {
echo "Error: timed out waiting for another cell to stage $MODEL_PATH" >&2
exit 1
}
if checkpoint_is_complete "$MODEL_PATH"; then
echo "Another cell staged $MODEL_PATH while we waited"
else
hf download "$HF_MODEL_ID" --local-dir "$MODEL_PATH"
fi
flock -u 9
exec 9>&-
checkpoint_is_complete "$MODEL_PATH" || {
echo "Error: $MODEL_PATH is still incomplete after hf download $HF_MODEL_ID." >&2
exit 1
}
fi
else
hf download "$HF_MODEL_ID"
export MODEL_PATH="$HF_MODEL_ID"
fi
nvidia-smi

resolve_trace_source
install_agentic_deps

SERVER_LOG="$RESULT_DIR/server.log"
mkdir -p "$RESULT_DIR"

CACHE_ARGS=()
if require_agentic_kv_offload_backend hicache; then
# HiCache extends RadixAttention: prefixes evicted from the HBM KV pool
# spill to a pinned host pool instead of being recomputed. On the
# 1M-context agentic corpus the live working set outgrows HBM past
# conc 8 (TP8) and the radix hit rate collapses to <0.1 against a ~0.97
# theoretical ceiling, so every turn re-prefills its whole history; the
# host tier restores those hits at C2C bandwidth.
# GLM-5.2 is DSA/MLA-family (attention_backend=dsa): every TP rank holds
# complete per-token KV.
#
# B300: use the NVFP4 B300 sibling's fixed 270 GB/rank target pool at every
# concurrency (#2651) rather than the B200 recipe's ratio mode at c1-c8.
# From SGLang v0.5.16 --hicache-size sizes only the target KV host pool;
# the coupled DSA indexer host pool (~38.73 GB/rank) is allocated on top.
# cluster:b300-dsxe advertises 3,977,095 MiB of host DRAM and this config
# exposes 80% (~3.18 TB); 270 GB/rank target + indexer across TP8 uses
# about 2.47 TB, leaving startup headroom. Overridable via HICACHE_SIZE.
DEFAULT_HICACHE_SIZE=270
MAX_HICACHE_SIZE=270
HICACHE_SIZE="${HICACHE_SIZE:-$DEFAULT_HICACHE_SIZE}"
if ! [[ "$HICACHE_SIZE" =~ ^[0-9]+$ ]]; then
echo "Error: HICACHE_SIZE must be a positive integer, got $HICACHE_SIZE" >&2
exit 1
fi
if awk -v s="$HICACHE_SIZE" -v cap="$MAX_HICACHE_SIZE" 'BEGIN { exit !(s <= 0 || s > cap) }'; then
echo "Error: HICACHE_SIZE=$HICACHE_SIZE must be in (0, $MAX_HICACHE_SIZE]" >&2
exit 1
fi
HICACHE_WRITE_POLICY="${HICACHE_WRITE_POLICY:-write_back}"
HICACHE_IO_BACKEND="${HICACHE_IO_BACKEND:-direct}"
HICACHE_MEM_LAYOUT="${HICACHE_MEM_LAYOUT:-page_first_direct}"
echo "HiCache CPU tier: conc=$CONC, target_size=$HICACHE_SIZE GB, total_capacity=${TOTAL_CPU_DRAM_GB} GB, write_policy=$HICACHE_WRITE_POLICY, io_backend=$HICACHE_IO_BACKEND, mem_layout=$HICACHE_MEM_LAYOUT"
CACHE_ARGS=(
--enable-hierarchical-cache
--hicache-size "$HICACHE_SIZE"
--hicache-write-policy "$HICACHE_WRITE_POLICY"
--hicache-io-backend "$HICACHE_IO_BACKEND"
--hicache-mem-layout "$HICACHE_MEM_LAYOUT"
)
fi

# With attention-DP, front the DP ranks with sglang-router using consistent
# hashing on the AIPerf correlation id so multi-turn sessions stay on the DP
# rank that holds their radix-cache prefix.
USE_SGLANG_ROUTER=false
SGLANG_BACKEND_PORT="$PORT"
ROUTER_LOG="$RESULT_DIR/router.log"
if [ "$DP_ATTENTION" = "true" ]; then
USE_SGLANG_ROUTER=true
export AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID=true
SGLANG_BACKEND_PORT=$((PORT + 1))
SGLANG_ROUTER_METRICS_PORT=$((PORT + 10000))
fi

# MTP: GLM-5.2 ships its own nextn head (num_nextn_predict_layers=1), so EAGLE
# runs off the checkpoint with no external draft model. num-steps 3 /
# eagle-topk 1 / num-draft-tokens 4 is 3 speculative tokens per verification
# step -- the same shape the NVFP4 B200/B300 siblings and the GLM-5.1-FP8
# cookbook speculative-mtp entry use, and the draft length whose golden AL is
# pinned below.
SPEC_ARGS=(
--speculative-algorithm EAGLE
--speculative-num-steps 3
--speculative-eagle-topk 1
--speculative-num-draft-tokens 4
)

PARALLEL_ARGS=(--tp "$TP" --ep-size "$EP_SIZE")
CHUNKED_PREFILL_SIZE=8192
if [ "$DP_ATTENTION" = "true" ]; then
# chunked-prefill-size is a whole-engine budget split across DP ranks:
# the cookbook HT cell's 8192 becomes 1,024 tokens/rank/step under dp8,
# which starves prefill on the 1M-context agentic corpus (observed: a
# conc-256 warmup could not drain within AIPerf's 1800s grace period
# while KV usage sat at ~0.01). Use the cookbook's own dp8 lever from
# the B200 cells (32768 = ~4096/rank).
CHUNKED_PREFILL_SIZE=32768
PARALLEL_ARGS+=(
--dp "$TP"
--enable-dp-attention
--tokenizer-worker-num "$TP"
--dist-init-addr "127.0.0.1:$((PORT + 2000))"
)
# Carried over from the NVFP4 sibling, where the draft MoE is bf16
# (hf_quant_config excludes model.layers.78*) and inheriting the target
# model's FlashInfer all-to-all dies at init with "Pre-permute function
# for flashinfer to triton is not registered". FP8: GLM-5.2-FP8 quantizes
# the nextn experts like every other layer (modules_to_not_convert lists
# only layer-78 norms and biases), so the draft MoE takes the FP8 runner
# and this pin is likely unnecessary here; it is kept so the DEP arm, if
# wired, starts from the configuration that is known to boot. Only
# relevant once expert parallelism puts an a2a in the MoE path -- the
# plain-TP arm below has none.
SPEC_ARGS+=(
--speculative-moe-a2a-backend none
--speculative-moe-runner-backend triton
)
else
# Cookbook low-latency levers; the DP-attention cell omits them.
PARALLEL_ARGS+=(
--kv-cache-dtype fp8_e4m3
--bf16-gemm-backend cutedsl
--max-prefill-tokens 8192
)
fi

# AgentX concurrency counts live session trees, not individual requests.
# Allow subagent fan-out to exceed CONC without clipping request bursts.
MAX_RUNNING_REQUESTS=$((2 * CONC))
GRAPH_ARGS=()
if [ "$DP_ATTENTION" != "true" ]; then
# Cookbook low-latency captures graphs up to its request cap; the
# DP-attention cell leaves the CUDA-graph batch list at SGLang defaults.
# --cuda-graph-max-bs counts requests, not verification tokens: SGLang's
# spec-decode graph runner scales each captured batch by
# --speculative-num-draft-tokens itself.
CUDA_GRAPH_MAX_BS=$MAX_RUNNING_REQUESTS
[ "$CUDA_GRAPH_MAX_BS" -gt 64 ] && CUDA_GRAPH_MAX_BS=64
GRAPH_ARGS=(--cuda-graph-max-bs "$CUDA_GRAPH_MAX_BS")
fi

# B300: 288 GB HBM3e per GPU. The NVFP4 B300 sibling runs 0.85, leaving ~43 GB
# of non-static headroom for the EAGLE draft head's verification-batch
# activations, the extra CUDA-graph capture at 4 draft tokens, and GLM-5.2's
# DSA indexer temporaries; that headroom is a fraction of the card and does
# not depend on the checkpoint.
# FP8: what does change is the KV pool inside the static share. The ~756 GB
# checkpoint is ~94.5 GB/GPU across TP8 (NVFP4: ~58 GB/GPU), so the fp8 KV
# pool is roughly 150 GB/GPU here against ~187 GB/GPU on NVFP4 B300 (and
# ~55 GB/GPU on FP8 B200). HiCache absorbs the difference as host spill.
MEM_FRACTION_STATIC="${MEM_FRACTION_STATIC:-0.85}"

export PYTHONNOUSERSITE=1
export TORCH_CUDA_ARCH_LIST=10.0
# Each concurrency in a full sweep is a separate Slurm allocation, while the
# DSXE runner home directory is shared. Keep SGLang's FlashInfer autotune, Triton,
# Inductor, and CUDA JIT caches allocation-local so concurrent cells cannot
# overwrite the same per-rank runtime-cache files. Non-Slurm launchers can
# provide an explicit SGLANG_CACHE_DIR override.
if [[ -n "${SLURM_JOB_ID:-}" ]]; then
export SGLANG_CACHE_DIR="${SGLANG_CACHE_DIR:-/tmp/sglang-cache-${SLURM_JOB_ID}}"
fi
# Agentic warmup dispatches hundreds of large prompts at once; allow up to
# 15 minutes of TCP progress before AIPerf declares a connection dead.
export AIPERF_HTTP_TCP_USER_TIMEOUT=900000
# AIPerf pins one pooled keep-alive connection per session (client-side
# keep-alive 300s) while uvicorn's default SGLANG_TIMEOUT_KEEP_ALIVE is 5s;
# inter-turn idle gaps (capped at 10s) can reuse a socket exactly as the
# server closes it -> ECONNRESET -> terminal warmup failure. Outlast the
# client pool so the race cannot occur.
export SGLANG_TIMEOUT_KEEP_ALIVE=900

# AgentX pins acceptance to the committed golden AL so submissions are compared
# on system performance at a fixed acceptance target rather than on draft-head
# quality (golden_al_distribution/README.md). 2.99 is the GLM-5.2 curve at
# num_speculative_tokens=3, thinking_on
# (golden_al_distribution/glm5.2_mtp.yaml, SPEED-Bench coding, run 28058352479).
# FP8: that curve was measured on this very checkpoint (glm-5.2-fp8), so no
# cross-precision assumption is involved here.
#
# SGLANG_SIMULATE_ACC_TOKEN_MODE only exists from SGLang v0.5.16. An older
# image would silently honor ACC_LEN/ACC_METHOD and ignore the token-mode half
# of the contract.
#
# EVAL_ONLY leaves simulated acceptance off: it commits drafted tokens
# regardless of the target logits, so generated text is wrong and the eval
# would score ~0.
if [ "${EVAL_ONLY:-false}" != "true" ]; then
export SGLANG_SIMULATE_ACC_LEN=2.99
export SGLANG_SIMULATE_ACC_METHOD=match-expected
export SGLANG_SIMULATE_ACC_TOKEN_MODE=real-draft-token
fi

SGLANG_CMD=(
python3 -m sglang.launch_server
--model-path "$MODEL_PATH"
--served-model-name "$MODEL"
--host 0.0.0.0
--port "$SGLANG_BACKEND_PORT"
--trust-remote-code
"${PARALLEL_ARGS[@]}"
# FP8: zai-org/GLM-5.2-FP8 ships quantization_config quant_method=fp8 with
# 128x128 weight blocks and dynamic e4m3 activations; pass the method
# explicitly as the Qwen3.5 FP8 B200/B300 SGLang siblings do.
--quantization fp8
# GLM-5.2 emits the GLM-4.7-style <tool_call>/<arg_key>/<arg_value> format;
# the glm47 parser is required for structured message.tool_calls (glm45
# leaves calls as raw text). Without it the SWE-bench mini-swe-agent eval
# dies with RepeatedFormatError ("No tool calls found in the response") on
# every instance and scores 0. Reasoning parser keeps hybrid-thinking
# output in reasoning_content instead of polluting content. Neither flag
# affects trace-replay throughput (pre-canned replay discards live
# responses).
--tool-call-parser glm47
--reasoning-parser glm45
--chunked-prefill-size "$CHUNKED_PREFILL_SIZE"
--mem-fraction-static "$MEM_FRACTION_STATIC"
--max-running-requests "$MAX_RUNNING_REQUESTS"
"${SPEC_ARGS[@]}"
"${GRAPH_ARGS[@]}"
"${CACHE_ARGS[@]}"
--watchdog-timeout 1800
--enable-metrics
)

printf '%q ' "${SGLANG_CMD[@]}" | tee "$RESULT_DIR/sglang_command.txt"
printf '\n' | tee -a "$RESULT_DIR/sglang_command.txt"

{
echo "=== SGLANG_SIMULATE_ACC_* env vars at launch (empty => real verification) ==="
env | grep -E '^SGLANG_SIMULATE_ACC_' | sort || true
echo "============================================================================"
} | tee "$SERVER_LOG"

echo "Starting SGLang server for B300..."
"${SGLANG_CMD[@]}" >> "$SERVER_LOG" 2>&1 &
SERVER_PID=$!
echo "Server PID: $SERVER_PID"

wait_for_server_ready --port "$SGLANG_BACKEND_PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID"

if [ "$USE_SGLANG_ROUTER" = "true" ]; then
echo "Starting SGLang router on port $PORT for $TP DP ranks..."
python3 -m sglang_router.launch_router \
--worker-urls "http://localhost:$SGLANG_BACKEND_PORT" \
--policy consistent_hashing \
--request-id-headers x-correlation-id \
--dp-aware \
--host 0.0.0.0 \
--port "$PORT" \
--prometheus-host 127.0.0.1 \
--prometheus-port "$SGLANG_ROUTER_METRICS_PORT" \
--connect-timeout-secs 900 \
--request-timeout-secs 14400 \
--disable-health-check \
--disable-retries > "$ROUTER_LOG" 2>&1 &
ROUTER_PID=$!
echo "Router PID: $ROUTER_PID"
wait_for_server_ready --port "$PORT" --server-log "$ROUTER_LOG" --server-pid "$ROUTER_PID"
fi

if [ "${EVAL_ONLY}" = "true" ]; then
# GLM-5.2's chat template defaults to reasoning_effort=Max when the
# client passes no chat_template_kwargs (mini-swe-agent doesn't), and the
# heavy thinking burns the default 75-step budget: on the 23-instance
# slice, 12/23 trajectories exited LimitsExceeded unsubmitted while 10 of
# the 11 that submitted resolved. Double the step budget for this recipe;
# other recipes keep the shared 75 default.
export SWEBENCH_AGENT_STEP_LIMIT=150
run_eval --port "$PORT"
else
build_replay_cmd "$RESULT_DIR"
REPLAY_CMD+=" --server-metrics http://localhost:$SGLANG_BACKEND_PORT/metrics"
run_agentic_replay_and_write_outputs "$RESULT_DIR"
fi
Loading
Loading