From 0a4fc8572b2b6f983aecdbafc43c325719f3bbba Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Tue, 23 Jun 2026 14:33:36 +0000 Subject: [PATCH 1/8] Extend the Q7 decode all-reduce benchmark to report external and internal input modes for TP=2 and TP=8. --- examples/qwen3/bench_allreduce.py | 251 ++++++++++++++++++------------ 1 file changed, 154 insertions(+), 97 deletions(-) diff --git a/examples/qwen3/bench_allreduce.py b/examples/qwen3/bench_allreduce.py index 70ef76bf..e8660b4e 100644 --- a/examples/qwen3/bench_allreduce.py +++ b/examples/qwen3/bench_allreduce.py @@ -1,11 +1,16 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -"""Benchmark end-to-end ``ark.all_reduce_packet`` latency on torch input. +"""Benchmark end-to-end ``ark.all_reduce_packet`` latency. Measures single-iteration latency for Qwen3 TP decode (1, 4096) and prefill -(2048, 4096) shapes, including registered-memory staging when needed. Each -rank runs as its own process and the parent reports max-rank latency. +(2048, 4096) shapes. The default run reports both input modes: + +- ``external``: torch CUDA placeholder input; includes registered-memory + staging. +- ``internal``: ARK-owned input tensor; excludes external staging. + +Each rank runs as its own process and the parent reports max-rank latency. python -m examples.qwen3.bench_allreduce --world-size 8 @@ -31,12 +36,13 @@ from _env import _load_worker_result, _subprocess_env _WORKER_SCRIPT = ''' -"""Worker: time torch-input ARK all-reduce without torch ops while launched.""" +"""Worker: time ARK all-reduce without torch ops while launched.""" import json import os import sys import time +import numpy as np import torch import ark from ark.executor import Executor @@ -44,21 +50,37 @@ rank = int(sys.argv[1]) world_size = int(sys.argv[2]) n_elements = int(sys.argv[3]) -label = sys.argv[4] +shape_name = sys.argv[4] +label = sys.argv[5] +input_mode = sys.argv[6] ark.init() ark.set_rank(rank) ark.set_world_size(world_size) -# Input is created and synchronized BEFORE launch, while no ARK loop kernel is -# live (safe). The benchmark includes any staging done by ark.all_reduce_packet. -x = torch.randn(n_elements, dtype=torch.float16, device=f"cuda:{rank}") -torch.cuda.synchronize(rank) +if input_mode == "external": + # Torch input is created and synchronized BEFORE launch, while no ARK loop + # kernel is live (safe). The measured graph includes all_reduce_packet's + # staging copy into ARK-managed registered memory. + x = torch.randn(n_elements, dtype=torch.float16, device=f"cuda:{rank}") + torch.cuda.synchronize(rank) +elif input_mode == "internal": + # ARK owns this tensor's storage. It is allocated by rt.launch(), then + # initialized from host memory before timing. This excludes external staging + # and avoids torch GPU work while the persistent runtime owns the device. + x = ark.tensor([n_elements], ark.fp16) + x_host = np.full((n_elements,), rank + 1, dtype=np.float16) +else: + raise ValueError(f"unknown input_mode: {input_mode}") + result = ark.all_reduce_packet(x, rank, world_size) with ark.Runtime() as rt: rt.launch(device_id=rank) + if input_mode == "internal": + x.from_numpy(x_host) + if world_size > 1: rt.barrier() @@ -74,7 +96,9 @@ print(json.dumps({ "rank": rank, + "shape": shape_name, "label": label, + "input_mode": input_mode, "world_size": world_size, "n_elements": n_elements, "latency_us": round(latency_us, 3), @@ -98,112 +122,134 @@ "prefill": ("prefill (2048, 4096)", 2048 * 4096), } +INPUT_MODES = ("external", "internal") + -def run_bench(world_size, timeout, shape): +def run_bench(world_size, timeout, shape, input_mode): results = [] any_failed = False - shapes = SHAPES.values() if shape == "all" else [SHAPES[shape]] - for label, n_elements in shapes: - procs = [] - for rank in range(world_size): - procs.append( - subprocess.Popen( - [ - sys.executable, - "-c", - _WORKER_SCRIPT, - str(rank), - str(world_size), - str(n_elements), - label, - ], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - cwd="/", - env=_subprocess_env(world_size), - ) - ) - shape_failed = False - rank_results = [] - try: - for rank, p in enumerate(procs): - try: - out, err = p.communicate(timeout=timeout) - except subprocess.TimeoutExpired: - shape_failed = True - print( - f"ERROR rank={rank} {label}: timed out after " - f"{timeout}s", - file=sys.stderr, + shapes = SHAPES.items() if shape == "all" else [(shape, SHAPES[shape])] + input_modes = INPUT_MODES if input_mode == "all" else (input_mode,) + for shape_name, (label, n_elements) in shapes: + for mode in input_modes: + procs = [] + for rank in range(world_size): + procs.append( + subprocess.Popen( + [ + sys.executable, + "-c", + _WORKER_SCRIPT, + str(rank), + str(world_size), + str(n_elements), + shape_name, + label, + mode, + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + cwd="/", + env=_subprocess_env(world_size), ) - break - if p.returncode != 0: - shape_failed = True - print( - f"ERROR rank={rank} {label}: " - f"{err.decode().strip()[-500:]}", - file=sys.stderr, + ) + shape_failed = False + rank_results = [] + result_label = f"{label} mode={mode}" + try: + for rank, p in enumerate(procs): + try: + out, err = p.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + shape_failed = True + print( + f"ERROR rank={rank} {result_label}: timed out " + f"after {timeout}s", + file=sys.stderr, + ) + break + if p.returncode != 0: + shape_failed = True + print( + f"ERROR rank={rank} {result_label}: " + f"{err.decode().strip()[-500:]}", + file=sys.stderr, + ) + result = _load_worker_result(out) + if result is None: + shape_failed = True + print( + f"ERROR rank={rank} {result_label}: no result", + file=sys.stderr, + ) + else: + rank_results.append(result) + if not shape_failed and len(rank_results) == world_size: + rank_results.sort(key=lambda d: d["rank"]) + max_result = max( + rank_results, key=lambda d: d["latency_us"] ) - result = _load_worker_result(out) - if result is None: - shape_failed = True - print( - f"ERROR rank={rank} {label}: no result", - file=sys.stderr, + results.append( + { + "shape": max_result["shape"], + "label": max_result["label"], + "input_mode": max_result["input_mode"], + "world_size": world_size, + "n_elements": max_result["n_elements"], + "max_rank": max_result["rank"], + "latency_us": max_result["latency_us"], + "rank_latencies_us": [ + d["latency_us"] for d in rank_results + ], + } ) else: - rank_results.append(result) - if not shape_failed and len(rank_results) == world_size: - rank_results.sort(key=lambda d: d["rank"]) - max_result = max(rank_results, key=lambda d: d["latency_us"]) - results.append( - { - "label": max_result["label"], - "world_size": world_size, - "n_elements": max_result["n_elements"], - "max_rank": max_result["rank"], - "latency_us": max_result["latency_us"], - "rank_latencies_us": [ - d["latency_us"] for d in rank_results - ], - } - ) - else: - any_failed = True - if not shape_failed: - print( - f"ERROR {label}: expected {world_size} rank results, " - f"got {len(rank_results)}", - file=sys.stderr, - ) - finally: - for p in procs: - p.kill() - p.wait() + any_failed = True + if not shape_failed: + print( + f"ERROR {result_label}: expected {world_size} " + f"rank results, got {len(rank_results)}", + file=sys.stderr, + ) + finally: + for p in procs: + p.kill() + p.wait() - print(f"\n{'=' * 72}") + print(f"\n{'=' * 88}") + print( + f"ARK all_reduce_packet latency | TP={world_size} " + f"(single iteration, max rank)" + ) + print(f"{'=' * 88}") print( - f"ARK all_reduce_packet torch-input latency | TP={world_size} " - f"(single iteration, max rank, includes staging)" + f"{'Shape':<24}{'Mode':<12}{'Elements':>12}" + f"{'Max rank':>10}{'ARK us':>10}" ) - print(f"{'=' * 72}") - print(f"{'Shape':<24}{'Elements':>12}{'Max rank':>10}{'ARK us':>10}") - print(f"{'-' * 72}") + print(f"{'-' * 88}") for d in results: print( - f"{d['label']:<24}{d['n_elements']:>12,}" + f"{d['label']:<24}{d['input_mode']:<12}{d['n_elements']:>12,}" f"{d['max_rank']:>10}{d['latency_us']:>10.2f}" ) - print(f"{'=' * 72}\n") + print(f"{'=' * 88}") + for d in results: + ark_ms = d["latency_us"] / 1000.0 + print( + f"RESULT name=allreduce shape={d['shape']} " + f"tp={d['world_size']} mode={d['input_mode']} " + f"ark_ms={ark_ms:.4f} latency_us={d['latency_us']:.3f}" + ) + print() return results, any_failed def main(): ap = argparse.ArgumentParser( description=( - "Benchmark end-to-end ark.all_reduce_packet latency on torch input " - "at Qwen3 TP shapes, including registered-memory staging " - "when needed" + "Benchmark end-to-end ark.all_reduce_packet latency at Qwen3 TP " + "shapes. External mode includes registered-memory staging; " + "internal mode uses ARK-owned input storage." ) ) ap.add_argument("--world-size", type=int, default=2) @@ -214,14 +260,25 @@ def main(): default="all", help="Qwen3 shape to benchmark; the perf gate uses decode", ) + ap.add_argument( + "--input-mode", + choices=("external", "internal", "all"), + default="all", + help="Input storage mode to benchmark", + ) args = ap.parse_args() # Repeated-iteration timing is intentionally unsupported until packet flag # rotation/reset exists. - results, any_failed = run_bench(args.world_size, args.timeout, args.shape) + results, any_failed = run_bench( + args.world_size, args.timeout, args.shape, args.input_mode + ) decode = [r for r in results if r["n_elements"] == SHAPES["decode"][1]] - if decode: + decode_external = [r for r in decode if r["input_mode"] == "external"] + if decode_external: + ark_ms = decode_external[0]["latency_us"] / 1000.0 + elif decode: ark_ms = decode[0]["latency_us"] / 1000.0 else: ark_ms = 999999.0 From 0934da9b4194a401640da9ec30c19905f7664a94 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Tue, 23 Jun 2026 15:01:30 +0000 Subject: [PATCH 2/8] Extend the Q7 decode all-reduce benchmark to report external and internal input modes for TP=2 and TP=8. --- __perf_gate__.sh | 75 +++++----- examples/qwen3/bench_allreduce.py | 31 ++-- examples/qwen3/test_bench_allreduce.py | 198 +++++++++++++++++++++++++ 3 files changed, 256 insertions(+), 48 deletions(-) create mode 100644 examples/qwen3/test_bench_allreduce.py diff --git a/__perf_gate__.sh b/__perf_gate__.sh index 8ad79ab5..cf95030e 100755 --- a/__perf_gate__.sh +++ b/__perf_gate__.sh @@ -5,53 +5,52 @@ set -uo pipefail export ARK_ROOT export PYTHONPATH="${PYTHONPATH:-$ARK_ROOT/python}" -target_ms=$(python3 - <<'PY' +tmpdir=$(mktemp -d) +trap 'rm -rf "$tmpdir"' EXIT +status=0 +python3 ../examples/qwen3/bench_allreduce.py --world-size 2 --shape decode --input-mode all \ + >"$tmpdir/tp2.log" 2>"$tmpdir/tp2.err" || status=1 +python3 ../examples/qwen3/bench_allreduce.py --world-size 8 --shape decode --input-mode all \ + >"$tmpdir/tp8.log" 2>"$tmpdir/tp8.err" || status=1 + +python3 - "$status" "$tmpdir/tp2.log" "$tmpdir/tp8.log" <<'PY' import importlib.util import pathlib +import re +import sys +status = int(sys.argv[1]) +logs = sys.argv[2:] path = pathlib.Path("../examples/qwen3/bench_allreduce.py") spec = importlib.util.spec_from_file_location("bench_allreduce", path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) -print(f"{module._DECODE_TARGET_MS:.4f}") -PY -) +target_ms = module._DECODE_TARGET_MS -tmpdir=$(mktemp -d) -trap 'rm -rf "$tmpdir"' EXIT -status=0 -python3 ../examples/qwen3/bench_allreduce.py --world-size 2 --shape decode >"$tmpdir/tp2.log" 2>"$tmpdir/tp2.err" || status=1 -python3 ../examples/qwen3/bench_allreduce.py --world-size 8 --shape decode >"$tmpdir/tp8.log" 2>"$tmpdir/tp8.err" || status=1 - -ark_ms=$(python3 - "$tmpdir/tp2.log" "$tmpdir/tp8.log" "$status" <<'PY' -import re -import sys - -values = [] -for name in sys.argv[1:3]: - text = open(name, encoding="utf-8").read() - match = re.search(r"PERF_GATE name=allreduce\s+ark_ms=([0-9.]+)", text) - if match: - values.append(float(match.group(1))) -if int(sys.argv[3]) or len(values) != 2: - print("999999.0000") -else: - print(f"{max(values):.4f}") -PY +pattern = re.compile( + r"RESULT name=allreduce shape=decode tp=(\d+) " + r"mode=(external|internal) ark_ms=([0-9.]+)" ) -ratio=$(python3 - "$ark_ms" "$target_ms" <<'PY' -import sys - -print(f"{float(sys.argv[1]) / float(sys.argv[2]):.4f}") -PY +values = {} +for log in logs: + text = pathlib.Path(log).read_text(encoding="utf-8") + for match in pattern.finditer(text): + values[(int(match.group(1)), match.group(2))] = float(match.group(3)) + +expected = { + (2, "external"), + (2, "internal"), + (8, "external"), + (8, "internal"), +} +missing = expected - values.keys() +sentinel = [v for v in values.values() if v <= 0.0 or v >= 999999.0] +ark_ms = 999999.0 if status or missing or sentinel else max(values.values()) +ratio = ark_ms / target_ms +print( + f"PERF_GATE name=allreduce ark_ms={ark_ms:.4f} " + f"sglang_ms={target_ms:.4f} ratio={ratio:.4f}" ) -printf 'PERF_GATE name=allreduce ark_ms=%s sglang_ms=%s ratio=%s\n' "$ark_ms" "$target_ms" "$ratio" -python3 - "$ark_ms" "$target_ms" "$status" <<'PY' -import sys - -ark_ms = float(sys.argv[1]) -target_ms = float(sys.argv[2]) -status = int(sys.argv[3]) -if status or ark_ms >= target_ms: +if status or missing or sentinel or ark_ms > target_ms: raise SystemExit(1) PY diff --git a/examples/qwen3/bench_allreduce.py b/examples/qwen3/bench_allreduce.py index e8660b4e..e0b5a587 100644 --- a/examples/qwen3/bench_allreduce.py +++ b/examples/qwen3/bench_allreduce.py @@ -233,6 +233,8 @@ def run_bench(world_size, timeout, shape, input_mode): f"{d['max_rank']:>10}{d['latency_us']:>10.2f}" ) print(f"{'=' * 88}") + # Supplemental per-shape/mode machine-readable rows. The root perf gate + # validates these rows for TP=2/8 external/internal decode coverage. for d in results: ark_ms = d["latency_us"] / 1000.0 print( @@ -244,6 +246,19 @@ def run_bench(world_size, timeout, shape, input_mode): return results, any_failed +def _perf_gate_ark_ms(results): + """Return compatibility PERF_GATE latency from external decode only.""" + decode_external = [ + r + for r in results + if r["n_elements"] == SHAPES["decode"][1] + and r["input_mode"] == "external" + ] + if decode_external: + return decode_external[0]["latency_us"] / 1000.0 + return 999999.0 + + def main(): ap = argparse.ArgumentParser( description=( @@ -274,14 +289,7 @@ def main(): args.world_size, args.timeout, args.shape, args.input_mode ) - decode = [r for r in results if r["n_elements"] == SHAPES["decode"][1]] - decode_external = [r for r in decode if r["input_mode"] == "external"] - if decode_external: - ark_ms = decode_external[0]["latency_us"] / 1000.0 - elif decode: - ark_ms = decode[0]["latency_us"] / 1000.0 - else: - ark_ms = 999999.0 + ark_ms = _perf_gate_ark_ms(results) ratio = ark_ms / _DECODE_TARGET_MS print( f"PERF_GATE name=allreduce" @@ -292,8 +300,11 @@ def main(): if any_failed: print("ERROR: one or more benchmark workers failed", file=sys.stderr) raise SystemExit(1) - if not decode: - print("ERROR: decode benchmark produced no result", file=sys.stderr) + if ark_ms == 999999.0: + print( + "ERROR: external decode benchmark produced no result for PERF_GATE", + file=sys.stderr, + ) raise SystemExit(1) diff --git a/examples/qwen3/test_bench_allreduce.py b/examples/qwen3/test_bench_allreduce.py new file mode 100644 index 00000000..d376aa8f --- /dev/null +++ b/examples/qwen3/test_bench_allreduce.py @@ -0,0 +1,198 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""CPU-only tests for Qwen3 all-reduce benchmark orchestration.""" + +import json +import os +import sys + +import pytest + +try: + from . import bench_allreduce +except ImportError: + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + import bench_allreduce + + +class _FakeProcess: + def __init__(self, argv): + self.argv = argv + self.returncode = 0 + + def communicate(self, timeout=None): + del timeout + rank = int(self.argv[3]) + world_size = int(self.argv[4]) + n_elements = int(self.argv[5]) + shape = self.argv[6] + label = self.argv[7] + mode = self.argv[8] + latency_us = (200.0 if mode == "external" else 500.0) + ( + 100.0 * rank + ) + result = { + "rank": rank, + "shape": shape, + "label": label, + "input_mode": mode, + "world_size": world_size, + "n_elements": n_elements, + "latency_us": latency_us, + } + return json.dumps(result).encode(), b"" + + def kill(self): + pass + + def wait(self): + pass + + +def _patch_popen(monkeypatch): + launches = [] + + def fake_popen(argv, **kwargs): + launches.append((argv, kwargs)) + return _FakeProcess(argv) + + monkeypatch.setattr(bench_allreduce.subprocess, "Popen", fake_popen) + monkeypatch.setattr( + bench_allreduce, + "_subprocess_env", + lambda world_size: {"WS": str(world_size)}, + ) + return launches + + +def test_run_bench_all_input_modes_launches_and_aggregates(monkeypatch): + launches = _patch_popen(monkeypatch) + + results, any_failed = bench_allreduce.run_bench( + world_size=2, timeout=7, shape="decode", input_mode="all" + ) + + assert not any_failed + assert [entry[0][3:] for entry in launches] == [ + ["0", "2", "4096", "decode", "decode (1, 4096)", "external"], + ["1", "2", "4096", "decode", "decode (1, 4096)", "external"], + ["0", "2", "4096", "decode", "decode (1, 4096)", "internal"], + ["1", "2", "4096", "decode", "decode (1, 4096)", "internal"], + ] + assert [entry[1]["env"] for entry in launches] == [{"WS": "2"}] * 4 + assert results == [ + { + "shape": "decode", + "label": "decode (1, 4096)", + "input_mode": "external", + "world_size": 2, + "n_elements": 4096, + "max_rank": 1, + "latency_us": 300.0, + "rank_latencies_us": [200.0, 300.0], + }, + { + "shape": "decode", + "label": "decode (1, 4096)", + "input_mode": "internal", + "world_size": 2, + "n_elements": 4096, + "max_rank": 1, + "latency_us": 600.0, + "rank_latencies_us": [500.0, 600.0], + }, + ] + + +@pytest.mark.parametrize("world_size", (2, 8)) +def test_run_bench_decode_reports_both_required_modes( + monkeypatch, world_size +): + _patch_popen(monkeypatch) + + results, any_failed = bench_allreduce.run_bench( + world_size=world_size, timeout=7, shape="decode", input_mode="all" + ) + + assert not any_failed + assert { + (result["world_size"], result["shape"], result["input_mode"]) + for result in results + } == { + (world_size, "decode", "external"), + (world_size, "decode", "internal"), + } + assert {result["n_elements"] for result in results} == {4096} + + +def test_run_bench_internal_mode_only_preserves_metadata(monkeypatch): + launches = _patch_popen(monkeypatch) + + results, any_failed = bench_allreduce.run_bench( + world_size=2, timeout=7, shape="decode", input_mode="internal" + ) + + assert not any_failed + assert [entry[0][8] for entry in launches] == ["internal", "internal"] + assert len(results) == 1 + assert results[0]["input_mode"] == "internal" + assert results[0]["shape"] == "decode" + assert results[0]["label"] == "decode (1, 4096)" + + +def test_main_default_reports_all_modes_and_keeps_external_gate( + monkeypatch, capsys +): + results = [ + { + "n_elements": 4096, + "input_mode": "internal", + "latency_us": 100.0, + }, + { + "n_elements": 4096, + "input_mode": "external", + "latency_us": 2000.0, + }, + ] + calls = [] + + def fake_run_bench(world_size, timeout, shape, input_mode): + calls.append((world_size, timeout, shape, input_mode)) + return results, False + + monkeypatch.setattr(bench_allreduce, "run_bench", fake_run_bench) + monkeypatch.setattr(sys, "argv", ["bench_allreduce.py"]) + + bench_allreduce.main() + + out, err = capsys.readouterr() + assert calls == [(2, 120, "all", "all")] + assert "PERF_GATE name=allreduce ark_ms=2.0000" in out + assert err == "" + + +def test_main_perf_gate_does_not_fallback_to_internal(monkeypatch, capsys): + results = [ + { + "n_elements": 4096, + "input_mode": "internal", + "latency_us": 100.0, + } + ] + monkeypatch.setattr( + bench_allreduce, "run_bench", lambda *args: (results, False) + ) + monkeypatch.setattr(sys, "argv", ["bench_allreduce.py"]) + + with pytest.raises(SystemExit) as exc_info: + bench_allreduce.main() + + out, err = capsys.readouterr() + assert exc_info.value.code == 1 + assert "PERF_GATE name=allreduce ark_ms=999999.0000" in out + assert ( + "ERROR: external decode benchmark produced no result for PERF_GATE" + in err + ) From 98ebf38dc0032f598184e9a8f1d1d4e474035034 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Tue, 23 Jun 2026 15:15:37 +0000 Subject: [PATCH 3/8] Extend the Q7 decode all-reduce benchmark to report external and internal input modes for TP=2 and TP=8. --- __perf_gate__.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/__perf_gate__.sh b/__perf_gate__.sh index cf95030e..3b9a2540 100755 --- a/__perf_gate__.sh +++ b/__perf_gate__.sh @@ -51,6 +51,6 @@ print( f"PERF_GATE name=allreduce ark_ms={ark_ms:.4f} " f"sglang_ms={target_ms:.4f} ratio={ratio:.4f}" ) -if status or missing or sentinel or ark_ms > target_ms: +if status or missing or sentinel or ark_ms >= target_ms: raise SystemExit(1) PY From e24364bd32cbc2990eea763105e475cdb1f679ab Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Tue, 23 Jun 2026 16:02:03 +0000 Subject: [PATCH 4/8] Extend to report decode TP=2/TP=8 all-reduce and input modes. --- __perf_gate__.sh | 27 +++-------------- examples/qwen3/bench_allreduce.py | 33 ++++++++++++++++++-- examples/qwen3/test_bench_allreduce.py | 42 ++++++++++++++++++++++++-- 3 files changed, 75 insertions(+), 27 deletions(-) diff --git a/__perf_gate__.sh b/__perf_gate__.sh index 3b9a2540..a219d2eb 100755 --- a/__perf_gate__.sh +++ b/__perf_gate__.sh @@ -16,7 +16,6 @@ python3 ../examples/qwen3/bench_allreduce.py --world-size 8 --shape decode --inp python3 - "$status" "$tmpdir/tp2.log" "$tmpdir/tp8.log" <<'PY' import importlib.util import pathlib -import re import sys status = int(sys.argv[1]) @@ -26,31 +25,15 @@ spec = importlib.util.spec_from_file_location("bench_allreduce", path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) target_ms = module._DECODE_TARGET_MS - -pattern = re.compile( - r"RESULT name=allreduce shape=decode tp=(\d+) " - r"mode=(external|internal) ark_ms=([0-9.]+)" +texts = [pathlib.Path(log).read_text(encoding="utf-8") for log in logs] +ark_ms = ( + 999999.0 if status else module._decode_gate_ark_ms_from_logs(texts) ) -values = {} -for log in logs: - text = pathlib.Path(log).read_text(encoding="utf-8") - for match in pattern.finditer(text): - values[(int(match.group(1)), match.group(2))] = float(match.group(3)) - -expected = { - (2, "external"), - (2, "internal"), - (8, "external"), - (8, "internal"), -} -missing = expected - values.keys() -sentinel = [v for v in values.values() if v <= 0.0 or v >= 999999.0] -ark_ms = 999999.0 if status or missing or sentinel else max(values.values()) ratio = ark_ms / target_ms print( - f"PERF_GATE name=allreduce ark_ms={ark_ms:.4f} " + f"PERF_GATE name={module._DECODE_GATE_NAME} ark_ms={ark_ms:.4f} " f"sglang_ms={target_ms:.4f} ratio={ratio:.4f}" ) -if status or missing or sentinel or ark_ms >= target_ms: +if status or ark_ms >= 999999.0 or ark_ms >= target_ms: raise SystemExit(1) PY diff --git a/examples/qwen3/bench_allreduce.py b/examples/qwen3/bench_allreduce.py index e0b5a587..ecf8968e 100644 --- a/examples/qwen3/bench_allreduce.py +++ b/examples/qwen3/bench_allreduce.py @@ -26,6 +26,7 @@ import argparse import os +import re import subprocess import sys @@ -116,6 +117,7 @@ # SGLang PROFILE.md Q7 nccl / comm target: 214.69 ms over 657 calls # on 8xA100 TP=8, batch=1 decode-dominated Qwen3-8B. _DECODE_TARGET_MS = 214.69 / 657.0 +_DECODE_GATE_NAME = "allreduce_decode" SHAPES = { "decode": ("decode (1, 4096)", 4096), @@ -238,7 +240,7 @@ def run_bench(world_size, timeout, shape, input_mode): for d in results: ark_ms = d["latency_us"] / 1000.0 print( - f"RESULT name=allreduce shape={d['shape']} " + f"RESULT name={_DECODE_GATE_NAME} shape={d['shape']} " f"tp={d['world_size']} mode={d['input_mode']} " f"ark_ms={ark_ms:.4f} latency_us={d['latency_us']:.3f}" ) @@ -247,7 +249,7 @@ def run_bench(world_size, timeout, shape, input_mode): def _perf_gate_ark_ms(results): - """Return compatibility PERF_GATE latency from external decode only.""" + """Return per-run PERF_GATE latency from external decode only.""" decode_external = [ r for r in results @@ -259,6 +261,31 @@ def _perf_gate_ark_ms(results): return 999999.0 +def _decode_gate_ark_ms_from_logs(log_texts, required_world_sizes=(2, 8)): + """Require decode RESULT rows for every TP/input-mode pair.""" + pattern = re.compile( + rf"RESULT name={_DECODE_GATE_NAME} shape=decode tp=(\d+) " + rf"mode=({'|'.join(INPUT_MODES)}) ark_ms=([0-9.]+)" + ) + values = {} + for text in log_texts: + for match in pattern.finditer(text): + values[(int(match.group(1)), match.group(2))] = float( + match.group(3) + ) + + expected = { + (world_size, mode) + for world_size in required_world_sizes + for mode in INPUT_MODES + } + missing = expected - values.keys() + sentinel = [v for v in values.values() if v <= 0.0 or v >= 999999.0] + if missing or sentinel: + return 999999.0 + return max(values[pair] for pair in expected) + + def main(): ap = argparse.ArgumentParser( description=( @@ -292,7 +319,7 @@ def main(): ark_ms = _perf_gate_ark_ms(results) ratio = ark_ms / _DECODE_TARGET_MS print( - f"PERF_GATE name=allreduce" + f"PERF_GATE name={_DECODE_GATE_NAME}" f" ark_ms={ark_ms:.4f}" f" sglang_ms={_DECODE_TARGET_MS:.4f}" f" ratio={ratio:.4f}" diff --git a/examples/qwen3/test_bench_allreduce.py b/examples/qwen3/test_bench_allreduce.py index d376aa8f..22531eda 100644 --- a/examples/qwen3/test_bench_allreduce.py +++ b/examples/qwen3/test_bench_allreduce.py @@ -126,6 +126,44 @@ def test_run_bench_decode_reports_both_required_modes( assert {result["n_elements"] for result in results} == {4096} +def test_decode_gate_parser_requires_tp2_tp8_and_both_modes(): + rows = [ + ( + "RESULT name=allreduce_decode shape=decode " + f"tp={world_size} mode={mode} ark_ms={ark_ms:.4f} " + f"latency_us={ark_ms * 1000.0:.3f}" + ) + for world_size, mode, ark_ms in [ + (2, "external", 0.2), + (2, "internal", 0.3), + (8, "external", 0.4), + (8, "internal", 0.5), + ] + ] + + assert ( + bench_allreduce._decode_gate_ark_ms_from_logs(["\n".join(rows)]) + == 0.5 + ) + assert ( + bench_allreduce._decode_gate_ark_ms_from_logs( + ["\n".join(rows[:-1])] + ) + == 999999.0 + ) + assert ( + bench_allreduce._decode_gate_ark_ms_from_logs( + [ + "\n".join( + row.replace("allreduce_decode", "allreduce") + for row in rows + ) + ] + ) + == 999999.0 + ) + + def test_run_bench_internal_mode_only_preserves_metadata(monkeypatch): launches = _patch_popen(monkeypatch) @@ -169,7 +207,7 @@ def fake_run_bench(world_size, timeout, shape, input_mode): out, err = capsys.readouterr() assert calls == [(2, 120, "all", "all")] - assert "PERF_GATE name=allreduce ark_ms=2.0000" in out + assert "PERF_GATE name=allreduce_decode ark_ms=2.0000" in out assert err == "" @@ -191,7 +229,7 @@ def test_main_perf_gate_does_not_fallback_to_internal(monkeypatch, capsys): out, err = capsys.readouterr() assert exc_info.value.code == 1 - assert "PERF_GATE name=allreduce ark_ms=999999.0000" in out + assert "PERF_GATE name=allreduce_decode ark_ms=999999.0000" in out assert ( "ERROR: external decode benchmark produced no result for PERF_GATE" in err From 2da573c8348d30c1044390c772834dc1f807ec83 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Tue, 23 Jun 2026 16:49:37 +0000 Subject: [PATCH 5/8] Format all-reduce benchmark parser tests --- examples/qwen3/test_bench_allreduce.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/examples/qwen3/test_bench_allreduce.py b/examples/qwen3/test_bench_allreduce.py index 22531eda..b78246ab 100644 --- a/examples/qwen3/test_bench_allreduce.py +++ b/examples/qwen3/test_bench_allreduce.py @@ -29,9 +29,7 @@ def communicate(self, timeout=None): shape = self.argv[6] label = self.argv[7] mode = self.argv[8] - latency_us = (200.0 if mode == "external" else 500.0) + ( - 100.0 * rank - ) + latency_us = (200.0 if mode == "external" else 500.0) + (100.0 * rank) result = { "rank": rank, "shape": shape, @@ -106,9 +104,7 @@ def test_run_bench_all_input_modes_launches_and_aggregates(monkeypatch): @pytest.mark.parametrize("world_size", (2, 8)) -def test_run_bench_decode_reports_both_required_modes( - monkeypatch, world_size -): +def test_run_bench_decode_reports_both_required_modes(monkeypatch, world_size): _patch_popen(monkeypatch) results, any_failed = bench_allreduce.run_bench( @@ -142,21 +138,17 @@ def test_decode_gate_parser_requires_tp2_tp8_and_both_modes(): ] assert ( - bench_allreduce._decode_gate_ark_ms_from_logs(["\n".join(rows)]) - == 0.5 + bench_allreduce._decode_gate_ark_ms_from_logs(["\n".join(rows)]) == 0.5 ) assert ( - bench_allreduce._decode_gate_ark_ms_from_logs( - ["\n".join(rows[:-1])] - ) + bench_allreduce._decode_gate_ark_ms_from_logs(["\n".join(rows[:-1])]) == 999999.0 ) assert ( bench_allreduce._decode_gate_ark_ms_from_logs( [ "\n".join( - row.replace("allreduce_decode", "allreduce") - for row in rows + row.replace("allreduce_decode", "allreduce") for row in rows ) ] ) From 509e6dca5ea0991422674c630869bd73f891d724 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Tue, 23 Jun 2026 17:30:36 +0000 Subject: [PATCH 6/8] Extend the decode all-reduce benchmark to report external and internal input modes for TP=2 and TP=8, fix the strict gate path without weakening it, rerun , and open the PR after clean evidence. --- __perf_gate__.sh | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/__perf_gate__.sh b/__perf_gate__.sh index a219d2eb..5de21721 100755 --- a/__perf_gate__.sh +++ b/__perf_gate__.sh @@ -1,6 +1,13 @@ #!/usr/bin/env bash set -uo pipefail +script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +bench_py="$script_dir/examples/qwen3/bench_allreduce.py" +if [[ ! -f "$bench_py" ]]; then + echo "ERROR: missing benchmark: $bench_py" >&2 + exit 1 +fi + : "${ARK_ROOT:=$PWD}" export ARK_ROOT export PYTHONPATH="${PYTHONPATH:-$ARK_ROOT/python}" @@ -8,19 +15,19 @@ export PYTHONPATH="${PYTHONPATH:-$ARK_ROOT/python}" tmpdir=$(mktemp -d) trap 'rm -rf "$tmpdir"' EXIT status=0 -python3 ../examples/qwen3/bench_allreduce.py --world-size 2 --shape decode --input-mode all \ +python3 "$bench_py" --world-size 2 --shape decode --input-mode all \ >"$tmpdir/tp2.log" 2>"$tmpdir/tp2.err" || status=1 -python3 ../examples/qwen3/bench_allreduce.py --world-size 8 --shape decode --input-mode all \ +python3 "$bench_py" --world-size 8 --shape decode --input-mode all \ >"$tmpdir/tp8.log" 2>"$tmpdir/tp8.err" || status=1 -python3 - "$status" "$tmpdir/tp2.log" "$tmpdir/tp8.log" <<'PY' +python3 - "$status" "$bench_py" "$tmpdir/tp2.log" "$tmpdir/tp8.log" <<'PY' import importlib.util import pathlib import sys status = int(sys.argv[1]) -logs = sys.argv[2:] -path = pathlib.Path("../examples/qwen3/bench_allreduce.py") +path = pathlib.Path(sys.argv[2]) +logs = sys.argv[3:] spec = importlib.util.spec_from_file_location("bench_allreduce", path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) From bbe2187a632c74b3f2cd1f358acbbdddd245f36c Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Wed, 24 Jun 2026 07:29:41 +0000 Subject: [PATCH 7/8] Measure prefill all-reduce input modes in Q7I gate --- __perf_gate__.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/__perf_gate__.sh b/__perf_gate__.sh index 5de21721..d777e6e9 100755 --- a/__perf_gate__.sh +++ b/__perf_gate__.sh @@ -15,9 +15,9 @@ export PYTHONPATH="${PYTHONPATH:-$ARK_ROOT/python}" tmpdir=$(mktemp -d) trap 'rm -rf "$tmpdir"' EXIT status=0 -python3 "$bench_py" --world-size 2 --shape decode --input-mode all \ +python3 "$bench_py" --world-size 2 --shape all --input-mode all \ >"$tmpdir/tp2.log" 2>"$tmpdir/tp2.err" || status=1 -python3 "$bench_py" --world-size 8 --shape decode --input-mode all \ +python3 "$bench_py" --world-size 8 --shape all --input-mode all \ >"$tmpdir/tp8.log" 2>"$tmpdir/tp8.err" || status=1 python3 - "$status" "$bench_py" "$tmpdir/tp2.log" "$tmpdir/tp8.log" <<'PY' @@ -43,4 +43,4 @@ print( ) if status or ark_ms >= 999999.0 or ark_ms >= target_ms: raise SystemExit(1) -PY +PY \ No newline at end of file From 2ef412680cf12c0e9323261cc5d2d19b4f4ec8a1 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Wed, 24 Jun 2026 07:29:54 +0000 Subject: [PATCH 8/8] Keep Q7I perf gate wrapper newline --- __perf_gate__.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/__perf_gate__.sh b/__perf_gate__.sh index d777e6e9..74a593b5 100755 --- a/__perf_gate__.sh +++ b/__perf_gate__.sh @@ -43,4 +43,4 @@ print( ) if status or ark_ms >= 999999.0 or ark_ms >= target_ms: raise SystemExit(1) -PY \ No newline at end of file +PY