diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 66f4015c760..40696d8344c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -23,6 +23,7 @@ Changelog *Speculative Decoding* - Add the **external draft** speculative-decoding mode: train a pretrained standalone draft model against a cached base model, instead of grafting a draft head onto the target. Select it with ``mtsp.convert(draft, [("external", cfg)])`` or the ``speculative_external`` recipe; the draft keeps its own embeddings and lm_head and exports as an ordinary HF checkpoint. +- Add a matched A|B speculative-decoding benchmark: ``examples/specdec_bench/compare_ab.py`` summarizes target-alone vs target + drafter runs (speedup, acceptance length, greedy match), and the launcher's ``specdec_ab_dflash2.yaml`` runs both arms in one job. *Megatron Framework (M-LM / M-Bridge)* diff --git a/examples/specdec_bench/compare_ab.py b/examples/specdec_bench/compare_ab.py new file mode 100644 index 00000000000..7cac353f93b --- /dev/null +++ b/examples/specdec_bench/compare_ab.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compare matched specdec_bench runs: target alone (A) vs target + draft (B). + +Reads run.py --save_dir outputs laid out as /r/{A,B}_c/, or +/{A,B}_c/ for a single repeat: + timing.json "Output TPS", "TTFT Time", "Request Generation Step Time" + acceptance_rate.json "Average_AL", "Category_AL" + *responses.jsonl generated text per question and turn (mtbench_, specbench_ or + plain responses.jsonl, depending on the dataset) + configuration.json "temperature", else runtime_params.sampling_kwargs.temperature + ("engine default" when sampling_kwargs set none); CONFIG_KEYS + +Writes one summary with a row per (repeat, concurrency): speedup = Output TPS B / Output +TPS A, AL of both arms, step time, TTFT, result-entry and turn counts, and each arm's +CONFIG_KEYS so a gate can confirm what actually ran. + +At concurrency 1 and temperature 0 the row also gets an exact-match comparison of A and B +outputs. Above concurrency 1, batching alone breaks bitwise equality; with sampling, a +mismatch is expected. Neither reports a match. The same row's "a_self_match" compares A with +A at the next concurrency of the repeat, also at temperature 0: how far the target alone +drifts with batch shape, the noise floor for the A-vs-B rate (some hybrid targets agree +with themselves on few turns). + +A row lists "problems" when an arm is missing, has no configuration.json, a result file has +other than exactly one entry (run.py appends to an existing save_dir), an arm produced no +responses, or the arms produced different numbers of turns. Any problem makes the CLI exit 1. + +Usage: + python3 compare_ab.py [--json ab_summary.json] +""" + +import argparse +import json +import os +import re +import sys +from glob import glob + +ARM_DIR = re.compile(r"([AB])_c(\d+)") +REPEAT_DIR = re.compile(r"r(\d+)") +# Copied per arm from configuration.json: which algorithm, drafter and speculative-token +# count ran (--block_size for DFLASH/DSPARK, --draft_length otherwise), and the settings +# both arms must share for a matched comparison (target, engine, GPU, parallelism, data). +CONFIG_KEYS = ( + "speculative_algorithm", + "draft_model_dir", + "block_size", + "draft_length", + "num_requests", + "model_dir", + "engine", + "engine_version", + "gpu", + "tp_size", + "ep_size", + "dataset", + "dataset_path", + "output_length", + "concurrency", +) + + +def _load_entries(path): + """All entries of a run.py result file (Metric.write appends one per run).""" + if not os.path.exists(path): + return [] + with open(path) as f: + data = json.load(f) + return data if isinstance(data, list) else [data] + + +def _mean(stat): + """specdec_bench stores statistics as {"mean": "0.1234", ...} with string values.""" + if isinstance(stat, dict): + return float(stat["mean"]) + return float(stat) if stat is not None else None + + +def _responses(arm_dir): + out = {} + for path in glob(os.path.join(arm_dir, "*responses.jsonl")): + with open(path) as f: + for line in f: + row = json.loads(line) + out[row["question_id"]] = row["choices"][0]["turns"] + return out + + +def _configuration(arm_dir): + path = os.path.join(arm_dir, "configuration.json") + if not os.path.exists(path): + return None + with open(path) as f: + return json.load(f) + + +def _temperature(config): + """Sampling temperature with run.py's precedence: --temperature, then runtime_params.""" + if config is None: + return None + if config.get("temperature") is not None: + return config["temperature"] + runtime_params = config.get("runtime_params") or {} + if "sampling_kwargs" not in runtime_params: + return 0 + # sampling_kwargs reach the engine as-is; without a temperature the engine default + # applies (1.0 for vLLM, TRT-LLM and AutoDeploy), which this file cannot see. + return (runtime_params["sampling_kwargs"] or {}).get("temperature", "engine default") + + +def load_arm(arm_dir): + """Summarize one run.py save_dir; metrics come from the last entry of each file.""" + timing = _load_entries(os.path.join(arm_dir, "timing.json")) + acceptance = _load_entries(os.path.join(arm_dir, "acceptance_rate.json")) + responses = _responses(arm_dir) + config = _configuration(arm_dir) + arm = { + "timing_entries": len(timing), + "acceptance_entries": len(acceptance), + "responses": len(responses), + "turns": sum(len(turns) for turns in responses.values()), + "temperature": _temperature(config), + "config": {key: config.get(key) for key in CONFIG_KEYS} if config else None, + "_responses": responses, + } + if timing: + arm["tps"] = float(timing[-1]["Output TPS"]) + arm["ttft_s"] = _mean(timing[-1].get("TTFT Time")) + arm["step_s"] = _mean(timing[-1].get("Request Generation Step Time")) + if acceptance: + arm["al"] = acceptance[-1].get("Average_AL") + arm["category_al"] = acceptance[-1].get("Category_AL") + return arm + + +def compare_outputs(a, b): + """Exact-match rate over (question, turn) pairs present in both arms.""" + pairs = exact = first_pairs = first_exact = 0 + for qid in sorted(set(a) & set(b)): + for turn, (ta, tb) in enumerate(zip(a[qid], b[qid])): + pairs += 1 + exact += ta == tb + if turn == 0: + first_pairs += 1 + first_exact += ta == tb + total_turns = sum(map(len, a.values())) + sum(map(len, b.values())) + return { + "pairs": pairs, + "unpaired_turns": total_turns - 2 * pairs, + "exact_match_rate": exact / pairs if pairs else None, + "first_turn_pairs": first_pairs, + "first_turn_exact_match_rate": first_exact / first_pairs if first_pairs else None, + } + + +def _arm_dirs(root): + """Yield (repeat, arm, concurrency, path) for every {A,B}_c dir under root.""" + for name in sorted(os.listdir(root)): + path = os.path.join(root, name) + if not os.path.isdir(path): + continue + if m := ARM_DIR.fullmatch(name): + yield 0, m.group(1), int(m.group(2)), path + elif r := REPEAT_DIR.fullmatch(name): + for arm_name in sorted(os.listdir(path)): + arm_path = os.path.join(path, arm_name) + if (m := ARM_DIR.fullmatch(arm_name)) and os.path.isdir(arm_path): + yield int(r.group(1)), m.group(1), int(m.group(2)), arm_path + + +def _problems(a, b): + problems = [] + for name, arm in (("A", a), ("B", b)): + if arm is None: + problems.append(f"arm {name} missing") + continue + if arm["config"] is None: + problems.append(f"arm {name} has no configuration.json") + if arm["timing_entries"] != 1 or arm["acceptance_entries"] != 1: + problems.append( + f"arm {name} has {arm['timing_entries']} timing / " + f"{arm['acceptance_entries']} acceptance entries, expected 1" + ) + if not arm["responses"]: + problems.append(f"arm {name} produced no responses") + if a and b and a["turns"] != b["turns"]: + problems.append(f"turn count differs: A {a['turns']}, B {b['turns']}") + return problems + + +def summarize(root): + """Build the A|B summary for every (repeat, concurrency) found under root.""" + arms = {} + for repeat, arm, conc, path in _arm_dirs(root): + arms[(repeat, conc, arm)] = load_arm(path) + + rows = [] + for repeat, conc in sorted({(r, c) for r, c, _ in arms}): + a, b = arms.get((repeat, conc, "A")), arms.get((repeat, conc, "B")) + row = {"repeat": repeat, "concurrency": conc, "problems": _problems(a, b)} + for name, arm in (("a", a), ("b", b)): + if arm is None: + continue + for key in ("tps", "ttft_s", "step_s", "al", "timing_entries", "responses", "turns"): + row[f"{key}_{name}"] = arm.get(key) + row[f"temperature_{name}"] = arm["temperature"] + row[f"config_{name}"] = arm["config"] + if b and b.get("category_al"): + row["category_al_b"] = b["category_al"] + if row.get("tps_a") and row.get("tps_b"): + row["speedup"] = row["tps_b"] / row["tps_a"] + # Exact match is only meaningful for greedy decoding without batching effects, so + # both arms must record temperature 0; an unknown temperature gets no match. + if ( + conc == 1 + and a is not None + and b is not None + and a["temperature"] == b["temperature"] == 0 + ): + row["match"] = compare_outputs(a["_responses"], b["_responses"]) + higher = sorted(c for r, c, n in arms if r == repeat and n == "A" and c > conc) + control = arms[(repeat, higher[0], "A")] if higher else None + if control is not None and control["temperature"] == 0: + row["a_self_match"] = { + "against_concurrency": higher[0], + **compare_outputs(a["_responses"], control["_responses"]), + } + rows.append(row) + return { + "root": os.path.abspath(root), + "ok": bool(rows) and not any(r["problems"] for r in rows), + "rows": rows, + } + + +def _fmt(value, fmt): + return format(value, fmt) if isinstance(value, (int, float)) else "-" + + +def main(argv=None): + """CLI entry point: print a table, optionally write JSON; exit 1 on any problem.""" + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("root", help="Directory holding the A/B save_dirs") + parser.add_argument("--json", dest="json_path", help="Write the summary to this file") + args = parser.parse_args(argv) + + summary = summarize(args.root) + if args.json_path: + with open(args.json_path, "w") as f: + json.dump(summary, f, indent=2) + if not summary["rows"]: + print(f"no {{A,B}}_c directories under {args.root}", file=sys.stderr) + return 2 + + print("| Repeat | C | AL (B) | TPS A | TPS B | Speedup | Turns A/B | Exact match (C=1, T=0) |") + print("| --- | --- | --- | --- | --- | --- | --- | --- |") + for r in summary["rows"]: + match = r.get("match") or {} + print( + f"| {r['repeat']} | {r['concurrency']} | {_fmt(r.get('al_b'), '.3f')} | " + f"{_fmt(r.get('tps_a'), '.1f')} | {_fmt(r.get('tps_b'), '.1f')} | " + f"{_fmt(r.get('speedup'), '.2f')} | {r.get('turns_a', '-')}/{r.get('turns_b', '-')} | " + f"{_fmt(match.get('exact_match_rate'), '.1%')} |" + ) + for r in summary["rows"]: + for problem in r["problems"]: + print(f"PROBLEM repeat {r['repeat']} concurrency {r['concurrency']}: {problem}") + return 0 if summary["ok"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/examples/specdec_bench/test_compare_ab.py b/tests/examples/specdec_bench/test_compare_ab.py new file mode 100644 index 00000000000..2efeb763633 --- /dev/null +++ b/tests/examples/specdec_bench/test_compare_ab.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""compare_ab.py: speedup, AL and exact match from matched A|B run.py outputs.""" + +import json + +import pytest +from compare_ab import main, summarize + + +def _write_arm( + arm_dir, + tps_entries=(100.0,), + al=1.0, + turns_by_question=None, + responses_file="mtbench_responses.jsonl", + config=None, +): + turns_by_question = turns_by_question or {1: ["x", "y"]} + arm_dir.mkdir(parents=True) + timing = [ + { + "Output TPS": tps, + "TTFT Time": {"mean": "0.0300"}, + "Request Generation Step Time": {"mean": "0.0040"}, + } + for tps in tps_entries + ] + (arm_dir / "timing.json").write_text(json.dumps(timing)) + (arm_dir / "acceptance_rate.json").write_text( + json.dumps([{"Average_AL": al, "Category_AL": {"math": al}}]) + ) + if config is not False: + (arm_dir / "configuration.json").write_text( + json.dumps({"temperature": 0.0, **(config or {})}) + ) + lines = [ + json.dumps({"question_id": q, "choices": [{"turns": turns}]}) + for q, turns in turns_by_question.items() + ] + (arm_dir / responses_file).write_text("\n".join(lines) + "\n") + + +@pytest.mark.parametrize( + ("layout", "responses_file"), + [("repeats", "mtbench_responses.jsonl"), ("flat", "responses.jsonl")], +) +def test_summary_reports_speedup_and_matches_only_at_concurrency_1( + tmp_path, layout, responses_file +): + """Both directory layouts; plain responses.jsonl is what non-MT-Bench datasets write.""" + base = tmp_path / "r1" if layout == "repeats" else tmp_path + a_turns = {1: ["x", "y"], 2: ["z", "w"]} + b_turns = {1: ["x", "y"], 2: ["z", "DIFFERENT"]} + config_a = {"speculative_algorithm": "NONE"} + config_b = { + "speculative_algorithm": "DFLASH", + "block_size": 7, + "gpu": "NVIDIA GB300", + "serving_config": {}, + } + a_turns_c8 = {1: ["x", "y"], 2: ["DRIFT", "w"]} # the target alone drifts with batch shape + for conc in (1, 8): + turns = a_turns if conc == 1 else a_turns_c8 + _write_arm(base / f"A_c{conc}", [100.0 * conc], 1.0, turns, responses_file, config_a) + _write_arm(base / f"B_c{conc}", [140.0 * conc], 1.77, b_turns, responses_file, config_b) + + summary = summarize(str(tmp_path)) + rows = {r["concurrency"]: r for r in summary["rows"]} + assert summary["ok"] + assert rows[1]["repeat"] == (1 if layout == "repeats" else 0) + for conc in (1, 8): + assert rows[conc]["speedup"] == pytest.approx(1.4) + assert rows[conc]["al_b"] == pytest.approx(1.77) + assert rows[conc]["turns_a"] == rows[conc]["turns_b"] == 4 + # What each arm ran, for the gate; only CONFIG_KEYS are copied. + assert rows[conc]["config_a"]["speculative_algorithm"] == "NONE" + assert rows[conc]["config_b"]["speculative_algorithm"] == "DFLASH" + assert rows[conc]["config_b"]["block_size"] == 7 + assert rows[conc]["config_b"]["gpu"] == "NVIDIA GB300" + assert "serving_config" not in rows[conc]["config_b"] + assert rows[1]["match"] == { + "pairs": 4, + "unpaired_turns": 0, + "exact_match_rate": 0.75, + "first_turn_pairs": 2, + "first_turn_exact_match_rate": 1.0, + } + assert rows[1]["a_self_match"]["against_concurrency"] == 8 + assert rows[1]["a_self_match"]["exact_match_rate"] == 0.75 + assert rows[1]["a_self_match"]["first_turn_exact_match_rate"] == 0.5 + # Batching alone breaks bitwise equality above concurrency 1, so no match there. + assert "match" not in rows[8] + assert "a_self_match" not in rows[8] + + +@pytest.mark.parametrize( + ("config", "temperature"), + [ + ({"temperature": 0.6}, 0.6), + ({"temperature": None, "runtime_params": {"sampling_kwargs": {"temperature": 0.6}}}, 0.6), + # The engine then samples at its own default (1.0 for vLLM), not greedily. + ({"temperature": None, "runtime_params": {"sampling_kwargs": {"top_p": 0.9}}}, None), + # No configuration.json: the temperature is unknown, so greedy is not assumed. + (False, None), + ], +) +def test_no_match_reported_when_sampling(tmp_path, config, temperature): + """With temperature > 0 a mismatch is expected and the rate would mean nothing.""" + for name in ("A", "B"): + _write_arm(tmp_path / f"{name}_c1", config=config if config is False else None) + if config is not False: + (tmp_path / f"{name}_c1" / "configuration.json").write_text(json.dumps(config)) + (row,) = summarize(str(tmp_path))["rows"] + assert "match" not in row + if temperature is not None: + assert row["temperature_b"] == pytest.approx(temperature) + + +@pytest.mark.parametrize( + ("arms", "problem"), + [ + ({"A": {}, "B": {}}, None), + ({"A": {}}, "arm B missing"), + ({"A": {"tps_entries": (50.0, 100.0)}, "B": {}}, "arm A has 2 timing"), + ({"A": {}, "B": {"turns_by_question": {1: ["x"]}}}, "turn count differs"), + ({"A": {"config": False}, "B": {}}, "arm A has no configuration.json"), + ], +) +def test_cli_exit_code_reflects_problems(tmp_path, capsys, arms, problem): + """A missing arm, appended results or a turn mismatch fails the CLI, not just the JSON.""" + for name, kwargs in arms.items(): + _write_arm(tmp_path / f"{name}_c1", **kwargs) + out = tmp_path / "ab_summary.json" + assert main([str(tmp_path), "--json", str(out)]) == (0 if problem is None else 1) + (row,) = json.loads(out.read_text())["rows"] + if problem is None: + assert row["problems"] == [] + else: + assert any(problem in p for p in row["problems"]) + assert problem in capsys.readouterr().out diff --git a/tools/launcher/common/specdec_bench/ab_compare.sh b/tools/launcher/common/specdec_bench/ab_compare.sh new file mode 100755 index 00000000000..6e08ef5eeaa --- /dev/null +++ b/tools/launcher/common/specdec_bench/ab_compare.sh @@ -0,0 +1,128 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source ${SCRIPT_DIR}/../service_utils.sh + +trap 'error_handler $0 $LINENO' ERR +trap 'exit_handler' EXIT + +################################################################################################### +# Matched A|B speculative-decoding benchmark in one job, so both arms share the node, image, +# prompts and sampling. For each repeat and concurrency it runs examples/specdec_bench/run.py +# twice: A = target alone (--speculative_algorithm NONE) and B = target + draft, A first on +# odd repeats and B first on even ones, so run order does not favor one arm. It then +# writes ${SAVE_ROOT}/ab_summary.json via examples/specdec_bench/compare_ab.py (speedup, AL, +# and A-vs-B exact match at concurrency 1 and temperature 0). Fails if an arm fails or the +# summary lists problems (missing arm, appended results, turn count mismatch). +# +# Required env: +# HF_MODEL_CKPT target model +# DRAFT_MODEL exported draft checkpoint +# SAVE_ROOT output root; results land in ${SAVE_ROOT}/r/{A,B}_c/ +# Optional env: +# ENGINE run.py --engine (default: VLLM) +# SPEC_ALGO run.py --speculative_algorithm for arm B (default: DFLASH) +# NUM_SPEC_TOKENS speculative tokens; --block_size for DFLASH/DSPARK, else --draft_length +# (default: 7) +# CONCURRENCIES comma-separated list (default: "1,8,32"); spaces also work +# REPEATS how many times to run the whole sweep (default: 1) +# TRTLLM_LAUNCH_SCRIPT command prefix for run.py, e.g. trtllm-llmapi-launch +# Arguments are passed to every run.py call (dataset, output length, TP size, ...). + +: "${HF_MODEL_CKPT:?HF_MODEL_CKPT is required}" +: "${DRAFT_MODEL:?DRAFT_MODEL is required}" +: "${SAVE_ROOT:?SAVE_ROOT is required}" +# Checked up front: otherwise arm A runs for minutes before arm B fails to load the draft. +# Only local paths: a Hugging Face repo ID (org/name) is resolved by the engine. +case "$DRAFT_MODEL" in + /*|./*|../*) + if [ ! -d "$DRAFT_MODEL" ]; then + report_result "FAIL: ab_compare: DRAFT_MODEL=${DRAFT_MODEL} is not a directory" + exit 1 + fi + ;; +esac +ENGINE=${ENGINE:-VLLM} +SPEC_ALGO=${SPEC_ALGO:-DFLASH} +NUM_SPEC_TOKENS=${NUM_SPEC_TOKENS:-7} +CONCURRENCIES=${CONCURRENCIES:-1,8,32} +REPEATS=${REPEATS:-1} +BENCH=modules/Model-Optimizer/examples/specdec_bench + +case "$SPEC_ALGO" in + DFLASH|DSPARK) SPEC_ARGS=(--block_size "$NUM_SPEC_TOKENS") ;; + *) SPEC_ARGS=(--draft_length "$NUM_SPEC_TOKENS") ;; +esac + +if ! pip show boto3 >/dev/null 2>&1 || ! pip show datasets >/dev/null 2>&1 || \ + ! pip show seaborn >/dev/null 2>&1; then + if ! pip install -r ${BENCH}/requirements.txt; then + report_result "FAIL: ab_compare: pip install requirements.txt failed" + exit 1 + fi +fi + +# run.py appends to existing result files, so a reused directory would mix runs. +if [ -n "$(ls -A "$SAVE_ROOT" 2>/dev/null)" ]; then + report_result "FAIL: ab_compare: SAVE_ROOT=${SAVE_ROOT} is not empty" + exit 1 +fi +mkdir -p "$SAVE_ROOT" +python3 -c 'import json, socket, subprocess, sys +gpus = subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True).stdout.splitlines() +json.dump({"host": socket.gethostname(), "gpus": gpus}, open(sys.argv[1], "w"), indent=2)' \ + "${SAVE_ROOT}/env.json" + +# Stops the job at the first failed arm: a sweep with a hole cannot be compared. +run_arm() { # arm (A or B), repeat, concurrency, then the run.py args shared by every arm + local arm=$1 repeat=$2 conc=$3 + shift 3 + local save_dir="${SAVE_ROOT}/r${repeat}/${arm}_c${conc}" + local arm_args=(--speculative_algorithm NONE) + if [ "$arm" = B ]; then + arm_args=(--speculative_algorithm "$SPEC_ALGO" --draft_model_dir "$DRAFT_MODEL" "${SPEC_ARGS[@]}") + fi + echo "=== ${save_dir} ===" + if ! ${TRTLLM_LAUNCH_SCRIPT:-} python3 ${BENCH}/run.py \ + --model_dir "${HF_MODEL_CKPT}" \ + --tokenizer "${HF_MODEL_CKPT}" \ + --engine "${ENGINE}" \ + --save_dir "${save_dir}" \ + "${arm_args[@]}" --concurrency "$conc" "$@"; then + report_result "FAIL: ab_compare: arm ${arm} failed (repeat ${repeat}, concurrency ${conc})" + exit 1 + fi +} + +for repeat in $(seq 1 "$REPEATS"); do + order="A B" + [ $((repeat % 2)) -eq 0 ] && order="B A" + for conc in ${CONCURRENCIES//,/ }; do + for arm in $order; do + run_arm "$arm" "$repeat" "$conc" "$@" + done + done +done + +# compare_ab.py exits 1 when a row lists problems (missing arm, appended results, turn +# count mismatch); the summary is written either way. +if ! python3 ${BENCH}/compare_ab.py "$SAVE_ROOT" --json "${SAVE_ROOT}/ab_summary.json"; then + report_result "FAIL: ab_compare: problems in ${SAVE_ROOT}/ab_summary.json" + exit 1 +fi +report_result "PASS: ab_compare: ${SAVE_ROOT}/ab_summary.json" diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/specdec_ab_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/specdec_ab_dflash2.yaml new file mode 100644 index 00000000000..174e6d796fb --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/specdec_ab_dflash2.yaml @@ -0,0 +1,58 @@ +# Matched A|B benchmark of a Qwen3-8B DFlash2 drafter on vLLM: Qwen3-8B alone (A) vs +# Qwen3-8B + drafter (B), in one job on one node, so the two arms share GPU, image, +# prompts and sampling. See common/specdec_bench/ab_compare.sh. +# +# Output in <>: +# r/{A,B}_c/ run.py results per arm +# ab_summary.json per (repeat, concurrency): Output TPS of both arms, +# speedup B/A, AL, step time, TTFT, response counts, and +# A-vs-B exact match at concurrency 1 and temperature 0; +# the job fails if any row lists problems +# env.json host and GPUs the job ran on +# +# draft_model must point at an exported drafter, typically the output of +# hf_online_dflash2.yaml in an earlier experiment: +# uv run launch.py --yaml examples/Qwen/Qwen3-8B/specdec_ab_dflash2.yaml --yes \ +# pipeline.global_vars.draft_model=/cicd//dflash2_bs16/exported-checkpoint-final +# +# Expect speedup > 1 at every concurrency, shrinking as concurrency grows, and exact match +# close to, but not always at, 100% at concurrency 1: the drafter's verification pass batches +# tokens differently, so rare numeric ties can flip. Plan roughly 45 min per repeat on one +# GPU; throughput at high concurrency is noisy, so use REPEATS >= 2 when comparing drafters. + +job_name: Qwen3-8B_DFlash2_ab_vllm +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3-8B + hf_data: /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl + draft_model: /cicd//dflash2_bs16/exported-checkpoint-final + output_dir: /scratchspace/ab + + task_0: + script: common/specdec_bench/ab_compare.sh + args: + - --mtbench <> + - --temperature 0 + - --output_length 4096 + - --tp_size 1 + - --ep_size 1 + environment: + - HF_MODEL_CKPT: <> + - DRAFT_MODEL: <> + - SAVE_ROOT: <> + - SPEC_ALGO: DFLASH + # Positions past 7 accept ~0 for a 2000-step drafter (7 and 15 give the same AL). + - NUM_SPEC_TOKENS: "7" + # Commas, not spaces: nemo_run emits unquoted `export FOO=value`. + - CONCURRENCIES: "1,8,32" + - REPEATS: "1" + # Node-local: a shared Triton cache breaks when jobs on different nodes write it. + - TRITON_CACHE_DIR: /tmp/triton-cache + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: "04:00:00" + # A release with DFlash2 serving (vllm#52816). + container: vllm/vllm-openai:v0.31.0