Skip to content

Commit 77a21c0

Browse files
cquil11cquil11
authored
feat(speedbench): add Muse Spark Open DFlash2 SGLang AL collector (#4128)
1 parent 68145e7 commit 77a21c0

3 files changed

Lines changed: 230 additions & 5 deletions

File tree

‎.github/workflows/speedbench-al.yml‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,15 @@ on: # zizmor: ignore[concurrency-limits]
2525
required: false
2626
type: string
2727
default: 'dsv4'
28+
framework:
29+
description: "Serving framework; selects the <model-prefix>_fp4_b300_<framework>.sh collector"
30+
required: false
31+
type: choice
32+
options: [vllm, sglang]
33+
default: 'vllm'
2834
image:
29-
# Must be vLLM >= 0.30: thinking-on cells rely on --chat-template-kwargs (vllm-project/vllm#44244).
30-
description: "vLLM container image (vLLM >= 0.30)"
35+
# vLLM collectors need vLLM >= 0.30: thinking-on cells rely on --chat-template-kwargs (vllm-project/vllm#44244).
36+
description: "Container image (vLLM >= 0.30 for vLLM collectors)"
3137
required: false
3238
type: string
3339
default: 'vllm/vllm-openai:v0.30.0'
@@ -94,14 +100,14 @@ env:
94100
MODEL: ${{ inputs.model }}
95101
MODEL_PREFIX: ${{ inputs.model-prefix }}
96102
PRECISION: fp4
97-
FRAMEWORK: vllm
103+
FRAMEWORK: ${{ inputs.framework }}
98104
EXP_NAME: ${{ inputs.model-prefix }}_speedbench
99105
IMAGE: ${{ inputs.image }}
100106
TP: '8'
101107
EP_SIZE: '1'
102108
DP_ATTENTION: 'false'
103109
SPEC_DECODING: mtp
104-
BENCH_SCRIPT_OVERRIDE: benchmarks/single_node/speedbench/${{ inputs.model-prefix }}_fp4_b300_vllm.sh
110+
BENCH_SCRIPT_OVERRIDE: benchmarks/single_node/speedbench/${{ inputs.model-prefix }}_fp4_b300_${{ inputs.framework }}.sh
105111
SALLOC_TIME_LIMIT: ${{ inputs.salloc-time }}
106112
# Matrix-collector tunables (propagated into the container via srun --export=ALL).
107113
MTP_LIST: ${{ inputs.mtp-list }}
@@ -143,9 +149,10 @@ jobs:
143149
entry=$(jq -cn \
144150
--arg runner "$RUNNER" \
145151
--arg model_prefix "$MODEL_PREFIX" \
152+
--arg framework "$FRAMEWORK" \
146153
'[{
147154
runner: $runner,
148-
framework: "vllm",
155+
framework: $framework,
149156
"model-prefix": $model_prefix,
150157
precision: "fp4",
151158
"spec-decoding": "mtp"

‎inferencex-e2e/benchmarks/single_node/speedbench/musespark_fp4_b300_sglang.sh‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
#!/usr/bin/env bash
2+
3+
# Muse Spark Open B300 SGLang SPEED-Bench AL matrix collector for DFlash2 speculative
4+
# decoding (RadixArk/Muse-Spark-Open-DFlash2).
5+
#
6+
# For each DFlash2 speculative-token count, measure the REAL acceptance length (AL) on one
7+
# SPEED-Bench category and emit a YAML matrix in the golden_al_distribution shape. Server
8+
# flags follow the verified B300 NVFP4 cell of the SGLang Muse Spark Open cookbook (TP4/EP4),
9+
# minus hierarchical cache, which only matters for prefix reuse.
10+
#
11+
# Muse Spark Open has no MTP head, so the speculative draft is the external DFlash2 model.
12+
# DFlash's --speculative-num-draft-tokens is the verification width (current token + drafts),
13+
# so MTP_LIST level N runs with N + 1 draft tokens and keys stay num_speculative_tokens.
14+
# The checkpoint has no thinking-off mode: reasoning_effort only takes low..xhigh, and the
15+
# SGLang client applies the chat template's default (xhigh), so only thinking_on is collected.
16+
#
17+
# AL is SGLang's avg_spec_accept_length (accepted tokens per verify step, bonus token
18+
# included, cumulative since server start), the same quantity as vLLM's
19+
# 1 + accepted/drafts. Each cell restarts the server, so it covers one cell only.
20+
#
21+
# Dispatch this collector through speedbench-al.yml with framework=sglang.
22+
23+
set -o pipefail
24+
repo_root="$(cd "$(dirname "$0")/../../.." && pwd)"
25+
source "$repo_root/benchmarks/check_env.sh"
26+
check_env_vars \
27+
CATEGORY MODEL MODEL_PATH MTP_LIST OUT_YAML PORT SPEEDBENCH_OUTPUT_LEN THINKING_MODES
28+
29+
if [[ "$THINKING_MODES" != "on" ]]; then
30+
echo "CRITICAL: Muse Spark Open has no thinking-off mode; dispatch with thinking-modes=on (got '$THINKING_MODES')"
31+
exit 1
32+
fi
33+
34+
SERVE_MODEL="${MODEL_PATH}"
35+
DRAFT_MODEL="${DRAFT_MODEL:-RadixArk/Muse-Spark-Open-DFlash2}"
36+
37+
MODEL_KEY="$(basename "$SERVE_MODEL" | tr '[:upper:]' '[:lower:]')"
38+
# AL is a per-draft accept/reject property independent of batch size, so batch the
39+
# SPEED-Bench pass to cut wall-clock, as the DSV4 DSpark collector does.
40+
CONCURRENCY="32"
41+
TEMPERATURE="1.0"
42+
TOP_P="0.95"
43+
44+
SPEEDBENCH_DIR="/workspace/speed_bench_data"
45+
RESULTS_DIR="/workspace/speedbench_results"
46+
47+
mkdir -p "$RESULTS_DIR"
48+
nvidia-smi
49+
GPU_NAME="$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)"
50+
51+
if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then
52+
if [[ ! -w "$(dirname "$MODEL_PATH")" ]]; then
53+
echo "CRITICAL: $MODEL_PATH is empty and $(dirname "$MODEL_PATH") is not writable."
54+
echo "The cluster's models.entries names this checkpoint but the weights were never staged."
55+
exit 1
56+
fi
57+
echo "=== $MODEL_PATH is empty; downloading $MODEL (~484 GB, first run only) ==="
58+
hf download "$MODEL" --local-dir "$MODEL_PATH"
59+
fi
60+
hf download "$DRAFT_MODEL"
61+
62+
echo "=== Downloading SPEED-Bench dataset ==="
63+
pip install -q datasets tiktoken
64+
curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \
65+
| python3 - --config qualitative --output_dir "$SPEEDBENCH_DIR"
66+
67+
if [[ ! -f "$SPEEDBENCH_DIR/qualitative.jsonl" ]]; then
68+
echo "CRITICAL: SPEED-Bench download failed — $SPEEDBENCH_DIR/qualitative.jsonl not found"
69+
exit 1
70+
fi
71+
72+
# SGLang's speed-bench loader only accepts the throughput-split categories on the CLI and
73+
# reads prompts from "turns", while prepare.py writes "messages". Filter the qualitative
74+
# split to $CATEGORY and restore "turns" from the user messages.
75+
CATEGORY_JSONL="$SPEEDBENCH_DIR/qualitative_${CATEGORY}.jsonl"
76+
python3 - "$SPEEDBENCH_DIR/qualitative.jsonl" "$CATEGORY_JSONL" "$CATEGORY" <<'EOF'
77+
import json, sys
78+
src, dst, category = sys.argv[1:]
79+
with open(src) as f, open(dst, "w") as out:
80+
for line in f:
81+
row = json.loads(line)
82+
if row.get("category") == category:
83+
row["turns"] = [m["content"] for m in row["messages"] if m["role"] == "user"]
84+
out.write(json.dumps(row) + "\n")
85+
EOF
86+
NUM_PROMPTS=$(wc -l < "$CATEGORY_JSONL")
87+
if [[ "$NUM_PROMPTS" -eq 0 ]]; then
88+
echo "CRITICAL: no SPEED-Bench qualitative rows for category '$CATEGORY'"
89+
exit 1
90+
fi
91+
92+
SERVER_PID=""
93+
_descendants() {
94+
local pid="$1" child
95+
for child in $(pgrep -P "$pid" 2>/dev/null || true); do
96+
echo "$child"
97+
_descendants "$child"
98+
done
99+
}
100+
cleanup_server() {
101+
if [[ -n "$SERVER_PID" ]]; then
102+
local descendants
103+
descendants=$(_descendants "$SERVER_PID")
104+
kill "$SERVER_PID" 2>/dev/null || true
105+
wait "$SERVER_PID" 2>/dev/null || true
106+
local pid
107+
for pid in $descendants; do
108+
kill -9 "$pid" 2>/dev/null || true
109+
done
110+
local waited=0
111+
while [[ $waited -lt 120 ]]; do
112+
local used
113+
used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | sort -rn | head -1)
114+
if [[ -z "$used" || "$used" -lt 2000 ]]; then break; fi
115+
sleep 3; waited=$((waited + 3))
116+
done
117+
SERVER_PID=""
118+
fi
119+
}
120+
trap 'cleanup_server' EXIT
121+
122+
declare -A AL_RESULT
123+
124+
run_cell() {
125+
local mode="$1" mtp="$2"
126+
127+
echo ""
128+
echo "=========================================="
129+
echo " Cell: thinking=$mode DFLASH=$mtp category=$CATEGORY"
130+
echo "=========================================="
131+
132+
local serve_args=(
133+
--model-path "$SERVE_MODEL"
134+
--host 0.0.0.0 --port "$PORT"
135+
--tp-size 4 --ep-size 4
136+
--quantization modelopt_fp4
137+
--moe-runner-backend flashinfer_trtllm
138+
--kv-cache-dtype fp8_e4m3
139+
--mem-fraction-static 0.8
140+
--trust-remote-code
141+
--enable-layernorm-sp
142+
--enable-megamoe-for-layernorm-sp
143+
--speculative-algorithm DFLASH
144+
--speculative-draft-model-path "$DRAFT_MODEL"
145+
--speculative-draft-attention-backend fa4
146+
--speculative-num-draft-tokens "$((mtp + 1))"
147+
--mm-feature-transport cpu
148+
--reasoning-parser muse
149+
--tool-call-parser muse
150+
)
151+
152+
local server_log="$RESULTS_DIR/server_${mode}_mtp${mtp}.log"
153+
python3 -m sglang.launch_server "${serve_args[@]}" > "$server_log" 2>&1 &
154+
SERVER_PID=$!
155+
156+
if ! PYTHONSAFEPATH=1 PYTHONPATH="$repo_root${PYTHONPATH:+:$PYTHONPATH}" python3 -m infx.bench wait \
157+
--url "http://0.0.0.0:${PORT}/health" --pid "$SERVER_PID" --log "$server_log"; then
158+
echo " -> server failed to start (thinking=$mode dflash=$mtp), recording N/A"
159+
AL_RESULT["${mode}_${mtp}"]="N/A"
160+
cleanup_server
161+
return
162+
fi
163+
164+
local result_file="$RESULTS_DIR/speedbench_${mode}_mtp${mtp}.jsonl"
165+
python3 -m sglang.bench_serving \
166+
--backend sglang \
167+
--port "$PORT" \
168+
--model "$SERVE_MODEL" \
169+
--dataset-name speed-bench \
170+
--dataset-path "$CATEGORY_JSONL" \
171+
--speed-bench-output-len "$SPEEDBENCH_OUTPUT_LEN" \
172+
--num-prompts "$NUM_PROMPTS" \
173+
--max-concurrency "$CONCURRENCY" \
174+
--temperature "$TEMPERATURE" \
175+
--top-p "$TOP_P" \
176+
--disable-ignore-eos \
177+
--output-file "$result_file" \
178+
--output-details
179+
180+
local al
181+
al=$(python3 -c '
182+
import json, sys
183+
with open(sys.argv[1]) as f:
184+
al = json.loads(f.readlines()[-1]).get("accept_length")
185+
print(f"{al:.2f}" if al else "N/A")
186+
' "$result_file" 2>/dev/null || echo "N/A")
187+
echo " -> thinking=$mode DFLASH=$mtp AL=$al"
188+
AL_RESULT["${mode}_${mtp}"]="$al"
189+
190+
cleanup_server
191+
}
192+
193+
for mtp in $MTP_LIST; do
194+
run_cell on "$mtp"
195+
done
196+
197+
{
198+
echo "# Acceptance Length (AL) reference values measured with SPEED-Bench."
199+
echo "# dataset: $CATEGORY | temperature: $TEMPERATURE | top_p: $TOP_P | output_len: $SPEEDBENCH_OUTPUT_LEN"
200+
echo "# thinking_on: chat template default (reasoning_effort xhigh)"
201+
echo "# draft: $DRAFT_MODEL"
202+
echo "# Measured on $MODEL_KEY ($GPU_NAME, SGLang DFlash2), per num_speculative_tokens."
203+
echo "# Auto-generated by benchmarks/single_node/speedbench/musespark_fp4_b300_sglang.sh (speedbench-al.yml)."
204+
echo "#"
205+
echo "# key = num_speculative_tokens (DFlash2 draft tokens); value = golden AL"
206+
echo "${MODEL_KEY}:"
207+
echo " thinking_on:"
208+
for mtp in $MTP_LIST; do
209+
echo " $mtp: ${AL_RESULT[on_${mtp}]:-N/A}"
210+
done
211+
} > "$OUT_YAML"
212+
213+
echo ""
214+
echo "=========================================="
215+
echo " SPEED-Bench AL matrix written to: $OUT_YAML"
216+
echo "=========================================="
217+
cat "$OUT_YAML"

‎inferencex-e2e/configs/runners.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -474,6 +474,7 @@ clusters:
474474
exclusive: true
475475
gres: "gpu:{gpus}"
476476
volumes:
477+
hf-home: {path: ~/.cache/huggingface}
477478
hf-hub-cache: {path: /data/home/sa-shared/gharunners/hf-hub-cache}
478479
aiperf-cache: {path: /data/home/sa-shared/gharunners/aiperf-cache}
479480
scratch: {path: /scratch/models, visibility: node-local}

0 commit comments

Comments
 (0)