537 lines
17 KiB
Python
537 lines
17 KiB
Python
"""
|
|
Benchmark fused allreduce+rmsnorm on AMD with correctness checks.
|
|
|
|
This script targets the same fused op used by SGLang:
|
|
`tensor_model_parallel_fused_allreduce_rmsnorm`.
|
|
|
|
It reports:
|
|
- eager mode latency (prefill-like)
|
|
- graph mode latency (decode-like)
|
|
- fused availability (whether fused path returns non-None)
|
|
- correctness (fused output matches split allreduce + rmsnorm reference)
|
|
|
|
Usage example:
|
|
torchrun --nproc_per_node=8 \
|
|
benchmark/kernels/all_reduce/benchmark_fused_ar_rms_amd.py \
|
|
--dtype bfloat16 \
|
|
--prefill-shapes 2048x8192,8192x8192 \
|
|
--decode-shapes 1x8192,4x8192,16x8192 \
|
|
--warmup 10 --iters 30 --repeats 5
|
|
"""
|
|
|
|
import argparse
|
|
import csv
|
|
import os
|
|
import statistics
|
|
from typing import Dict, List, Optional, Sequence, Tuple
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn.functional as F
|
|
|
|
from sglang.srt.distributed.communication_op import (
|
|
tensor_model_parallel_all_reduce,
|
|
tensor_model_parallel_fused_allreduce_rmsnorm,
|
|
)
|
|
from sglang.srt.distributed.parallel_state import (
|
|
destroy_distributed_environment,
|
|
destroy_model_parallel,
|
|
graph_capture,
|
|
init_distributed_environment,
|
|
initialize_model_parallel,
|
|
set_custom_all_reduce,
|
|
)
|
|
|
|
Shape = Tuple[int, int]
|
|
|
|
|
|
def parse_shapes(raw: str) -> List[Shape]:
|
|
shapes: List[Shape] = []
|
|
for item in [x.strip() for x in raw.split(",") if x.strip()]:
|
|
if "x" not in item:
|
|
raise ValueError(f"Invalid shape '{item}', expected MxN format.")
|
|
m_str, n_str = item.split("x", 1)
|
|
m = int(m_str)
|
|
n = int(n_str)
|
|
if m <= 0 or n <= 0:
|
|
raise ValueError(f"Invalid shape '{item}', both dims must be positive.")
|
|
shapes.append((m, n))
|
|
if not shapes:
|
|
raise ValueError("Empty shape list is not allowed.")
|
|
return shapes
|
|
|
|
|
|
def dtype_from_name(name: str) -> torch.dtype:
|
|
mapping = {
|
|
"float16": torch.float16,
|
|
"fp16": torch.float16,
|
|
"bfloat16": torch.bfloat16,
|
|
"bf16": torch.bfloat16,
|
|
}
|
|
if name not in mapping:
|
|
raise ValueError(f"Unsupported dtype: {name}")
|
|
return mapping[name]
|
|
|
|
|
|
def check_close(
|
|
a: torch.Tensor, b: torch.Tensor, dtype: torch.dtype
|
|
) -> Tuple[bool, str]:
|
|
if dtype == torch.bfloat16:
|
|
rtol, atol = 2e-2, 1.25e-1
|
|
else:
|
|
rtol, atol = 1e-2, 2e-2
|
|
try:
|
|
torch.testing.assert_close(a, b, rtol=rtol, atol=atol)
|
|
return True, "PASS"
|
|
except AssertionError:
|
|
max_diff = torch.max(torch.abs(a - b)).item()
|
|
mean_diff = torch.mean(torch.abs(a - b)).item()
|
|
return False, f"FAIL(max={max_diff:.6f},mean={mean_diff:.6f})"
|
|
|
|
|
|
def _measure_us(
|
|
fn,
|
|
warmup: int,
|
|
iters: int,
|
|
repeats: int,
|
|
device: torch.device,
|
|
) -> Tuple[float, Dict[str, float]]:
|
|
for _ in range(warmup):
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
|
|
start_event = torch.cuda.Event(enable_timing=True)
|
|
end_event = torch.cuda.Event(enable_timing=True)
|
|
samples_us: List[float] = []
|
|
|
|
for _ in range(max(1, repeats)):
|
|
_barrier(device)
|
|
torch.cuda.synchronize()
|
|
start_event.record()
|
|
for _ in range(iters):
|
|
fn()
|
|
end_event.record()
|
|
end_event.synchronize()
|
|
samples_us.append(start_event.elapsed_time(end_event) * 1000.0 / iters)
|
|
|
|
sorted_samples = sorted(samples_us)
|
|
p50 = float(statistics.median(sorted_samples))
|
|
p95 = float(sorted_samples[int((len(sorted_samples) - 1) * 0.95)])
|
|
return p50, {
|
|
"p50_us": p50,
|
|
"p95_us": p95,
|
|
"min_us": float(sorted_samples[0]),
|
|
"max_us": float(sorted_samples[-1]),
|
|
}
|
|
|
|
|
|
def _barrier(device: torch.device):
|
|
try:
|
|
dist.barrier(device_ids=[device.index])
|
|
except TypeError:
|
|
dist.barrier()
|
|
|
|
|
|
def _mean_across_ranks(value: float, device: torch.device) -> float:
|
|
t = torch.tensor([value], dtype=torch.float64, device=device)
|
|
dist.all_reduce(t, op=dist.ReduceOp.SUM)
|
|
t /= dist.get_world_size()
|
|
return float(t.item())
|
|
|
|
|
|
def _all_true_across_ranks(value: bool, device: torch.device) -> bool:
|
|
t = torch.tensor([1 if value else 0], dtype=torch.int32, device=device)
|
|
dist.all_reduce(t, op=dist.ReduceOp.MIN)
|
|
return bool(int(t.item()))
|
|
|
|
|
|
def _make_inputs(
|
|
shape: Shape,
|
|
dtype: torch.dtype,
|
|
seed: int,
|
|
residual_mode: str,
|
|
rank: int,
|
|
device: torch.device,
|
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
m, n = shape
|
|
torch.manual_seed(seed + rank * 17)
|
|
x = torch.randn((m, n), dtype=torch.float32, device=device).to(dtype)
|
|
if residual_mode == "self":
|
|
residual = x.clone()
|
|
elif residual_mode == "random":
|
|
residual = torch.randn((m, n), dtype=torch.float32, device=device).to(dtype)
|
|
elif residual_mode == "zero":
|
|
residual = torch.zeros((m, n), dtype=dtype, device=device)
|
|
else:
|
|
raise ValueError(f"Unknown residual_mode: {residual_mode}")
|
|
weight = torch.randn((n,), dtype=torch.float32, device=device).to(dtype)
|
|
return x, residual, weight
|
|
|
|
|
|
def _split_reference(
|
|
x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
ar_out = tensor_model_parallel_all_reduce(x.clone())
|
|
residual_out = ar_out + residual
|
|
out = F.rms_norm(
|
|
input=residual_out,
|
|
normalized_shape=(residual_out.shape[-1],),
|
|
weight=weight,
|
|
eps=eps,
|
|
)
|
|
return out, residual_out
|
|
|
|
|
|
def bench_eager(
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
eps: float,
|
|
warmup: int,
|
|
iters: int,
|
|
repeats: int,
|
|
) -> Dict[str, object]:
|
|
split_fn = lambda: _split_reference(x, residual, weight, eps)
|
|
split_us, split_stats = _measure_us(split_fn, warmup, iters, repeats, x.device)
|
|
|
|
fused_probe = tensor_model_parallel_fused_allreduce_rmsnorm(
|
|
x.clone(), residual.clone(), weight, eps
|
|
)
|
|
fused_available = fused_probe is not None
|
|
|
|
fused_us: Optional[float] = None
|
|
fused_stats: Optional[Dict[str, float]] = None
|
|
if fused_available:
|
|
fused_fn = lambda: tensor_model_parallel_fused_allreduce_rmsnorm(
|
|
x, residual, weight, eps
|
|
)
|
|
fused_us, fused_stats = _measure_us(fused_fn, warmup, iters, repeats, x.device)
|
|
|
|
ref_out, ref_residual = _split_reference(x, residual, weight, eps)
|
|
if fused_available:
|
|
fused_out, fused_residual = tensor_model_parallel_fused_allreduce_rmsnorm(
|
|
x.clone(), residual.clone(), weight, eps
|
|
)
|
|
out_ok, out_detail = check_close(fused_out, ref_out, x.dtype)
|
|
res_ok, res_detail = check_close(fused_residual, ref_residual, x.dtype)
|
|
correctness_ok = out_ok and res_ok
|
|
correctness_detail = f"out={out_detail}, residual={res_detail}"
|
|
else:
|
|
correctness_ok = True
|
|
correctness_detail = "SKIP(fused_unavailable)"
|
|
|
|
return {
|
|
"split_us": split_us,
|
|
"split_stats": split_stats,
|
|
"fused_available": fused_available,
|
|
"fused_us": fused_us,
|
|
"fused_stats": fused_stats,
|
|
"correctness_ok": correctness_ok,
|
|
"correctness_detail": correctness_detail,
|
|
}
|
|
|
|
|
|
def bench_graph(
|
|
x: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
eps: float,
|
|
warmup: int,
|
|
iters: int,
|
|
repeats: int,
|
|
) -> Dict[str, object]:
|
|
split_x = x.clone()
|
|
split_res = residual.clone()
|
|
split_graph_out: Optional[torch.Tensor] = None
|
|
|
|
with graph_capture() as gc:
|
|
split_graph = torch.cuda.CUDAGraph()
|
|
with torch.cuda.graph(split_graph, stream=gc.stream):
|
|
split_graph_out, _ = _split_reference(split_x, split_res, weight, eps)
|
|
|
|
def split_replay():
|
|
split_graph.replay()
|
|
|
|
split_us, split_stats = _measure_us(split_replay, warmup, iters, repeats, x.device)
|
|
|
|
fused_probe = tensor_model_parallel_fused_allreduce_rmsnorm(
|
|
x.clone(), residual.clone(), weight, eps
|
|
)
|
|
fused_available = fused_probe is not None
|
|
|
|
fused_us: Optional[float] = None
|
|
fused_stats: Optional[Dict[str, float]] = None
|
|
fused_graph_out: Optional[torch.Tensor] = None
|
|
fused_graph_residual: Optional[torch.Tensor] = None
|
|
|
|
if fused_available:
|
|
fused_x = x.clone()
|
|
fused_res = residual.clone()
|
|
with graph_capture() as gc:
|
|
fused_graph = torch.cuda.CUDAGraph()
|
|
with torch.cuda.graph(fused_graph, stream=gc.stream):
|
|
fused_graph_out, fused_graph_residual = (
|
|
tensor_model_parallel_fused_allreduce_rmsnorm(
|
|
fused_x, fused_res, weight, eps
|
|
)
|
|
)
|
|
|
|
def fused_replay():
|
|
fused_graph.replay()
|
|
|
|
fused_us, fused_stats = _measure_us(
|
|
fused_replay, warmup, iters, repeats, x.device
|
|
)
|
|
|
|
ref_out, ref_residual = _split_reference(x, residual, weight, eps)
|
|
if (
|
|
fused_available
|
|
and fused_graph_out is not None
|
|
and fused_graph_residual is not None
|
|
):
|
|
fused_graph.replay()
|
|
torch.cuda.synchronize()
|
|
out_ok, out_detail = check_close(fused_graph_out, ref_out, x.dtype)
|
|
res_ok, res_detail = check_close(fused_graph_residual, ref_residual, x.dtype)
|
|
correctness_ok = out_ok and res_ok
|
|
correctness_detail = f"out={out_detail}, residual={res_detail}"
|
|
else:
|
|
correctness_ok = True
|
|
correctness_detail = "SKIP(fused_unavailable)"
|
|
|
|
return {
|
|
"split_us": split_us,
|
|
"split_stats": split_stats,
|
|
"fused_available": fused_available,
|
|
"fused_us": fused_us,
|
|
"fused_stats": fused_stats,
|
|
"correctness_ok": correctness_ok,
|
|
"correctness_detail": correctness_detail,
|
|
}
|
|
|
|
|
|
def _shape_bytes(shape: Shape, dtype: torch.dtype) -> int:
|
|
m, n = shape
|
|
return m * n * torch.tensor([], dtype=dtype).element_size()
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description="Benchmark fused allreduce+rmsnorm (prefill eager + decode graph)."
|
|
)
|
|
parser.add_argument(
|
|
"--dtype",
|
|
type=str,
|
|
default="bf16",
|
|
choices=["fp16", "bf16", "float16", "bfloat16"],
|
|
)
|
|
parser.add_argument("--eps", type=float, default=1e-6)
|
|
parser.add_argument("--seed", type=int, default=1234)
|
|
parser.add_argument(
|
|
"--residual-mode",
|
|
type=str,
|
|
default="self",
|
|
choices=["self", "random", "zero"],
|
|
help="Use residual=x (self) to match aiter test behavior by default.",
|
|
)
|
|
parser.add_argument(
|
|
"--prefill-shapes",
|
|
type=str,
|
|
default="2048x8192,8192x8192,16384x8192",
|
|
help="Comma-separated MxN shapes for eager mode.",
|
|
)
|
|
parser.add_argument(
|
|
"--decode-shapes",
|
|
type=str,
|
|
default="1x8192,2x8192,4x8192,8x8192,16x8192",
|
|
help="Comma-separated MxN shapes for graph mode.",
|
|
)
|
|
parser.add_argument("--warmup", type=int, default=10)
|
|
parser.add_argument("--iters", type=int, default=30)
|
|
parser.add_argument("--repeats", type=int, default=5)
|
|
parser.add_argument(
|
|
"--mode",
|
|
type=str,
|
|
default="both",
|
|
choices=["eager", "graph", "both"],
|
|
)
|
|
parser.add_argument(
|
|
"--csv-out",
|
|
type=str,
|
|
default=None,
|
|
help="Optional output CSV path (written on rank 0 only).",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
dtype = dtype_from_name(args.dtype)
|
|
rank = int(os.environ.get("RANK", "0"))
|
|
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
|
local_rank = int(os.environ.get("LOCAL_RANK", str(rank)))
|
|
torch.cuda.set_device(local_rank % torch.cuda.device_count())
|
|
device = torch.device(f"cuda:{local_rank % torch.cuda.device_count()}")
|
|
|
|
set_custom_all_reduce(True)
|
|
init_distributed_environment(
|
|
world_size=world_size,
|
|
rank=rank,
|
|
local_rank=local_rank,
|
|
distributed_init_method="env://",
|
|
backend="nccl",
|
|
)
|
|
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
|
|
|
prefill_shapes = parse_shapes(args.prefill_shapes)
|
|
decode_shapes = parse_shapes(args.decode_shapes)
|
|
|
|
if rank == 0:
|
|
print(
|
|
"Config: "
|
|
f"world_size={world_size}, dtype={dtype}, residual_mode={args.residual_mode}, "
|
|
f"warmup={args.warmup}, iters={args.iters}, repeats={args.repeats}"
|
|
)
|
|
|
|
run_modes: Sequence[str]
|
|
if args.mode == "both":
|
|
run_modes = ("eager", "graph")
|
|
else:
|
|
run_modes = (args.mode,)
|
|
csv_rows: List[Dict[str, object]] = []
|
|
|
|
for mode in run_modes:
|
|
shapes = prefill_shapes if mode == "eager" else decode_shapes
|
|
if rank == 0:
|
|
phase_name = "prefill(eager)" if mode == "eager" else "decode(graph)"
|
|
print("\n" + "=" * 120)
|
|
print(f"Mode: {phase_name}")
|
|
print(
|
|
"| Shape | Input bytes/rank | Split p50 (us) | Fused p50 (us) | Speedup | Fused available | Correctness |"
|
|
)
|
|
print(
|
|
"|:------|-----------------:|---------------:|---------------:|--------:|:----------------|:------------|"
|
|
)
|
|
|
|
for shape in shapes:
|
|
x, residual, weight = _make_inputs(
|
|
shape=shape,
|
|
dtype=dtype,
|
|
seed=args.seed,
|
|
residual_mode=args.residual_mode,
|
|
rank=rank,
|
|
device=device,
|
|
)
|
|
|
|
if mode == "eager":
|
|
metrics = bench_eager(
|
|
x=x,
|
|
residual=residual,
|
|
weight=weight,
|
|
eps=args.eps,
|
|
warmup=args.warmup,
|
|
iters=args.iters,
|
|
repeats=args.repeats,
|
|
)
|
|
else:
|
|
metrics = bench_graph(
|
|
x=x,
|
|
residual=residual,
|
|
weight=weight,
|
|
eps=args.eps,
|
|
warmup=args.warmup,
|
|
iters=args.iters,
|
|
repeats=args.repeats,
|
|
)
|
|
|
|
split_us = _mean_across_ranks(float(metrics["split_us"]), device)
|
|
fused_available = _all_true_across_ranks(
|
|
bool(metrics["fused_available"]), device
|
|
)
|
|
correctness_ok = _all_true_across_ranks(
|
|
bool(metrics["correctness_ok"]), device
|
|
)
|
|
|
|
fused_us: Optional[float] = None
|
|
if fused_available and metrics["fused_us"] is not None:
|
|
fused_us = _mean_across_ranks(float(metrics["fused_us"]), device)
|
|
|
|
if rank == 0:
|
|
m, n = shape
|
|
shape_str = f"{m}x{n}"
|
|
bytes_per_rank = _shape_bytes(shape, dtype)
|
|
if fused_us is not None and fused_us > 0:
|
|
speedup = split_us / fused_us
|
|
speedup_str = f"{speedup:.3f}x"
|
|
fused_str = f"{fused_us:.1f}"
|
|
else:
|
|
speedup_str = "N/A"
|
|
fused_str = "N/A"
|
|
correctness_text = (
|
|
"PASS" if correctness_ok else str(metrics["correctness_detail"])
|
|
)
|
|
print(
|
|
f"| {shape_str} | {bytes_per_rank} | {split_us:.1f} | {fused_str} | "
|
|
f"{speedup_str} | {str(fused_available)} | {correctness_text} |"
|
|
)
|
|
csv_rows.append(
|
|
{
|
|
"mode": mode,
|
|
"shape": shape_str,
|
|
"m": m,
|
|
"n": n,
|
|
"bytes_per_rank": bytes_per_rank,
|
|
"split_p50_us": split_us,
|
|
"fused_p50_us": fused_us if fused_us is not None else "",
|
|
"speedup_split_over_fused": (
|
|
split_us / fused_us
|
|
if fused_us is not None and fused_us > 0
|
|
else ""
|
|
),
|
|
"fused_available": fused_available,
|
|
"correctness_ok": correctness_ok,
|
|
"correctness_detail": correctness_text,
|
|
"dtype": str(dtype),
|
|
"world_size": world_size,
|
|
"residual_mode": args.residual_mode,
|
|
"warmup": args.warmup,
|
|
"iters": args.iters,
|
|
"repeats": args.repeats,
|
|
}
|
|
)
|
|
|
|
if rank == 0 and args.csv_out:
|
|
os.makedirs(os.path.dirname(args.csv_out) or ".", exist_ok=True)
|
|
fieldnames = [
|
|
"mode",
|
|
"shape",
|
|
"m",
|
|
"n",
|
|
"bytes_per_rank",
|
|
"split_p50_us",
|
|
"fused_p50_us",
|
|
"speedup_split_over_fused",
|
|
"fused_available",
|
|
"correctness_ok",
|
|
"correctness_detail",
|
|
"dtype",
|
|
"world_size",
|
|
"residual_mode",
|
|
"warmup",
|
|
"iters",
|
|
"repeats",
|
|
]
|
|
with open(args.csv_out, "w", newline="", encoding="utf-8") as f:
|
|
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
writer.writerows(csv_rows)
|
|
print(f"\nSaved CSV to: {args.csv_out}")
|
|
|
|
_barrier(device)
|
|
destroy_model_parallel()
|
|
destroy_distributed_environment()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|