123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
import csv
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from sglang.test.ci.ci_register import register_amd_ci
|
|
|
|
register_amd_ci(est_time=240, suite="stage-c-test-large-8-gpu-amd")
|
|
|
|
|
|
class TestAiterAllreduceFusionAmd(unittest.TestCase):
|
|
def test_fused_ar_rms_benchmark(self):
|
|
if not torch.cuda.is_available():
|
|
self.skipTest("CUDA/ROCm device is not available.")
|
|
if torch.cuda.device_count() < 8:
|
|
self.skipTest("This test requires at least 8 GPUs.")
|
|
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
benchmark_script = (
|
|
repo_root
|
|
/ "benchmark"
|
|
/ "kernels"
|
|
/ "all_reduce"
|
|
/ "benchmark_fused_ar_rms_amd.py"
|
|
)
|
|
self.assertTrue(
|
|
benchmark_script.exists(),
|
|
f"Missing benchmark script: {benchmark_script}",
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory(prefix="aiter_fused_ar_rms_") as tmpdir:
|
|
csv_path = Path(tmpdir) / "fused_ar_rms_check.csv"
|
|
cmd = [
|
|
sys.executable,
|
|
"-m",
|
|
"torch.distributed.run",
|
|
"--standalone",
|
|
"--nproc_per_node=8",
|
|
str(benchmark_script),
|
|
"--dtype",
|
|
"bf16",
|
|
"--prefill-shapes",
|
|
# Include both <=64MiB and >64MiB shapes to verify default gate behavior.
|
|
"128x7168,512x7168,2048x7168,4096x7168,5120x7168",
|
|
"--decode-shapes",
|
|
"1x7168,8x7168,64x7168,512x7168",
|
|
"--warmup",
|
|
"3",
|
|
"--iters",
|
|
"15",
|
|
"--repeats",
|
|
"2",
|
|
"--csv-out",
|
|
str(csv_path),
|
|
]
|
|
|
|
env = os.environ.copy()
|
|
result = subprocess.run(
|
|
cmd,
|
|
cwd=str(repo_root),
|
|
env=env,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
timeout=1200,
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
self.fail(
|
|
"Benchmark command failed.\n"
|
|
f"Return code: {result.returncode}\n"
|
|
f"Command: {' '.join(cmd)}\n"
|
|
f"Output:\n{result.stdout}"
|
|
)
|
|
|
|
self.assertTrue(csv_path.exists(), f"CSV output not found: {csv_path}")
|
|
|
|
with open(csv_path, "r", encoding="utf-8") as f:
|
|
rows = list(csv.DictReader(f))
|
|
|
|
self.assertGreater(len(rows), 0, "CSV contains no rows.")
|
|
|
|
eager_rows = [r for r in rows if r["mode"] == "eager"]
|
|
graph_rows = [r for r in rows if r["mode"] == "graph"]
|
|
self.assertGreater(len(eager_rows), 0, "Missing eager rows in CSV.")
|
|
self.assertGreater(len(graph_rows), 0, "Missing graph rows in CSV.")
|
|
|
|
# Correctness should always pass regardless of fused availability.
|
|
bad_rows = [r for r in rows if r["correctness_ok"] != "True"]
|
|
self.assertEqual(
|
|
[],
|
|
bad_rows,
|
|
f"Found correctness failures: {bad_rows}",
|
|
)
|
|
|
|
# We should see fused path active for small shapes in both modes.
|
|
self.assertTrue(
|
|
any(r["fused_available"] == "True" for r in eager_rows),
|
|
"Expected at least one eager row with fused_available=True.",
|
|
)
|
|
self.assertTrue(
|
|
any(r["fused_available"] == "True" for r in graph_rows),
|
|
"Expected at least one graph row with fused_available=True.",
|
|
)
|
|
|
|
# Default gate should reject at least one oversized eager shape.
|
|
large_eager_rows = [
|
|
r for r in eager_rows if int(r["bytes_per_rank"]) > 64 * 1024 * 1024
|
|
]
|
|
self.assertTrue(
|
|
any(r["fused_available"] == "False" for r in large_eager_rows),
|
|
"Expected fused fallback for oversized eager shape(s) under default gate.",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|