[JIT Kernel] Multi-GPU test/bench framework for custom all-reduce + TP QKNorm (#26706)
Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai>
This commit is contained in:
DarkSharpness
co-authored by
Claude
ziyi.xu
Claude
ziyi.xu
parent
5331de0f8c
commit
d72314808f
@@ -1,27 +1,34 @@
|
||||
"""
|
||||
Benchmark JIT custom all-reduce (v2) vs NCCL vs AOT custom all-reduce (v1).
|
||||
"""Benchmark JIT custom all-reduce (v2) vs NCCL, AOT custom-AR (v1), and
|
||||
FlashInfer trtllm allreduce_fusion.
|
||||
|
||||
Usage (torchrun required for multi-GPU):
|
||||
torchrun --nproc_per_node=2 bench_custom_all_reduce.py
|
||||
torchrun --nproc_per_node=4 bench_custom_all_reduce.py --dtype float16
|
||||
torchrun --nproc_per_node=8 bench_custom_all_reduce.py --warmup 10 --iters 100
|
||||
Usage::
|
||||
|
||||
The script initializes all three backends, then benchmarks each over a sweep
|
||||
of message sizes. Results are printed as a comparison table on rank 0.
|
||||
# Benchmark on every supported world size (2..8 GPUs):
|
||||
python benchmark/bench_custom_all_reduce.py
|
||||
# Pick a specific world size (or comma-separated list):
|
||||
python benchmark/bench_custom_all_reduce.py --num-gpu 4
|
||||
python benchmark/bench_custom_all_reduce.py --num-gpu 2,4,8
|
||||
|
||||
The script self-relaunches under ``torchrun --nproc_per_node=N`` for each N in
|
||||
``num_gpus``; results are printed on rank 0 of every run.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import contextlib
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
from math import isnan
|
||||
from typing import Dict, List, Optional
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import is_in_ci
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
from sglang.jit_kernel.benchmark import marker
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, multigpu_bench_main
|
||||
from sglang.jit_kernel.mp import register_comm_cleanup
|
||||
from sglang.jit_kernel.utils import cache_once, is_arch_support_pdl
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
@@ -30,12 +37,14 @@ register_cuda_ci(
|
||||
disabled="requires multi-GPU, self-skips in CI",
|
||||
)
|
||||
|
||||
DTYPE_MAP = {
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float32": torch.float32,
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sweep parameters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
DTYPE = torch.bfloat16
|
||||
# torch.dtype.itemsize exists only on newer torch; element_size() is portable.
|
||||
DTYPE_ITEMSIZE = torch.tensor([], dtype=DTYPE).element_size()
|
||||
MESSAGE_SIZES_BYTES = [
|
||||
4 * 1024, # 4K
|
||||
16 * 1024, # 16K
|
||||
@@ -50,29 +59,65 @@ MESSAGE_SIZES_BYTES = [
|
||||
7 * 128 * 1024, # 896K
|
||||
1 * 1024 * 1024, # 1M
|
||||
2 * 1024 * 1024, # 2M
|
||||
3 * 1024 * 1024, # 2M
|
||||
3 * 1024 * 1024, # 3M
|
||||
4 * 1024 * 1024, # 4M
|
||||
8 * 1024 * 1024, # 8M
|
||||
16 * 1024 * 1024, # 16M
|
||||
32 * 1024 * 1024, # 32M
|
||||
]
|
||||
WORLD_SIZES = list(range(2, 9))
|
||||
MAX_BYTES = max(MESSAGE_SIZES_BYTES)
|
||||
# trtllm allreduce_fusion only supports these world sizes.
|
||||
FI_SUPPORTED_WORLD_SIZES = (2, 4, 8)
|
||||
# AOT custom_all_reduce (v1) only supports these world sizes.
|
||||
AOT_SUPPORTED_WORLD_SIZES = (2, 4, 6, 8)
|
||||
PROVIDERS = ["nccl", "aot", "jit", "fi"]
|
||||
WORLD_SIZES = get_benchmark_range(WORLD_SIZES, [2, 4, 8])
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-rank distributed init (run once per torchrun worker)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_cpu_group() -> dist.ProcessGroup:
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(backend="gloo")
|
||||
ps._WORLD = coord = ps.init_world_group(
|
||||
ranks=list(range(world_size)),
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
)
|
||||
atexit.register(dist.destroy_process_group)
|
||||
# Quieter benchmark output.
|
||||
logging.disable(logging.INFO)
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
return coord.cpu_group
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_nccl_group() -> dist.ProcessGroup:
|
||||
_init_cpu_group()
|
||||
coord = ps._WORLD
|
||||
assert coord is not None and coord.device_group is not None
|
||||
return coord.device_group
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Backend wrappers - each exposes a uniform interface:
|
||||
# .name - display name
|
||||
# .capture() - context manager for CUDA-graph recording
|
||||
# .all_reduce() - perform an all-reduce and return the result tensor
|
||||
# Backend wrappers - each exposes:
|
||||
# .all_reduce(tensor) -> Tensor
|
||||
# .graph_context() -> context manager wrapping cuda-graph capture
|
||||
# (nullcontext when capture is not required)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class NCCLAllReduceBackend:
|
||||
name = "NCCL"
|
||||
def __init__(self) -> None:
|
||||
self.group = _init_nccl_group()
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup):
|
||||
self.group = group
|
||||
|
||||
def capture(self, register_input: bool):
|
||||
def graph_context(self):
|
||||
return contextlib.nullcontext()
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
@@ -80,42 +125,42 @@ class NCCLAllReduceBackend:
|
||||
return tensor
|
||||
|
||||
|
||||
class AOTAllReduceBackend:
|
||||
name = "AOT"
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup, device: torch.device):
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce import (
|
||||
CustomAllreduce,
|
||||
class JITAllReduceBackend:
|
||||
def __init__(self) -> None:
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
|
||||
max_size = max(MESSAGE_SIZES_BYTES)
|
||||
self.comm = CustomAllreduce(group, device, max_size=max_size)
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
self.comm = CustomAllReduceV2(
|
||||
_init_cpu_group(), device, max_pull_size=MAX_BYTES
|
||||
)
|
||||
if self.comm.disabled:
|
||||
raise RuntimeError("AOT CustomAllreduce is disabled on this system")
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
register_comm_cleanup(self.comm)
|
||||
|
||||
def capture(self, register_input: bool):
|
||||
return self.comm.capture() # ignore register_input since v1 always requires it
|
||||
def graph_context(self):
|
||||
return self.comm.capture()
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
|
||||
assert self.comm.should_custom_ar(tensor), str(tensor.shape)
|
||||
return self.comm.custom_all_reduce(tensor)
|
||||
|
||||
|
||||
class JITAllReduceBackend:
|
||||
name = "JIT"
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup, device: torch.device):
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
class AOTAllReduceBackend:
|
||||
def __init__(self) -> None:
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce import (
|
||||
CustomAllreduce,
|
||||
)
|
||||
|
||||
max_size = max(MESSAGE_SIZES_BYTES)
|
||||
self.comm = CustomAllReduceV2(group, device, max_pull_size=max_size)
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
self.comm = CustomAllreduce(_init_cpu_group(), device, max_size=MAX_BYTES)
|
||||
if self.comm.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
raise RuntimeError("AOT CustomAllreduce is disabled on this system")
|
||||
register_comm_cleanup(self.comm)
|
||||
|
||||
def capture(self, register_input: bool):
|
||||
return self.comm.capture() if register_input else contextlib.nullcontext()
|
||||
def graph_context(self):
|
||||
return self.comm.capture()
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
|
||||
assert self.comm.should_custom_ar(tensor), str(tensor.shape)
|
||||
@@ -123,262 +168,127 @@ class JITAllReduceBackend:
|
||||
|
||||
|
||||
class FlashInferAllReduceBackend:
|
||||
name = "FI"
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup, dtype: torch.dtype):
|
||||
def __init__(self) -> None:
|
||||
import flashinfer.comm as comm
|
||||
|
||||
rank = torch.distributed.get_rank(group=group)
|
||||
world_size = torch.distributed.get_world_size(group=group)
|
||||
max_size = max(MESSAGE_SIZES_BYTES)
|
||||
hidden_dim = min(MESSAGE_SIZES_BYTES) // 2
|
||||
num_tokens = max_size // hidden_dim
|
||||
self.comm = comm
|
||||
self.hidden_dim = hidden_dim
|
||||
self.workspace = comm.create_allreduce_fusion_workspace(
|
||||
group = _init_cpu_group()
|
||||
rank = dist.get_rank(group=group)
|
||||
world_size = dist.get_world_size(group=group)
|
||||
# Use the smallest message size as the inner hidden dim, so any
|
||||
# message in the sweep is an integer multiple of it.
|
||||
hidden_dim = min(MESSAGE_SIZES_BYTES) // DTYPE_ITEMSIZE
|
||||
num_tokens = MAX_BYTES // (hidden_dim * DTYPE_ITEMSIZE)
|
||||
self._comm = comm
|
||||
self._hidden_dim = hidden_dim
|
||||
self._workspace = comm.create_allreduce_fusion_workspace(
|
||||
backend="trtllm",
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
max_token_num=num_tokens,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
dtype=DTYPE,
|
||||
)
|
||||
|
||||
def capture(self, *_):
|
||||
def graph_context(self):
|
||||
return contextlib.nullcontext()
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
|
||||
return self.comm.allreduce_fusion(
|
||||
input=tensor.view(-1, self.hidden_dim),
|
||||
workspace=self.workspace,
|
||||
pattern=self.comm.AllReduceFusionPattern.kAllReduce,
|
||||
launch_with_pdl=True,
|
||||
def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self._comm.allreduce_fusion(
|
||||
input=tensor.view(-1, self._hidden_dim),
|
||||
workspace=self._workspace,
|
||||
pattern=self._comm.AllReduceFusionPattern.kAllReduce,
|
||||
launch_with_pdl=is_arch_support_pdl(),
|
||||
fp32_acc=True,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmarking helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
@cache_once
|
||||
def _init_nccl_backend() -> NCCLAllReduceBackend:
|
||||
return NCCLAllReduceBackend()
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("--dtype", choices=DTYPE_MAP.keys(), default="bfloat16")
|
||||
p.add_argument("--warmup", type=int, default=5)
|
||||
p.add_argument("--iters", type=int, default=50)
|
||||
p.add_argument("--no-inplace", dest="register_input", action="store_false")
|
||||
return p.parse_args()
|
||||
@cache_once
|
||||
def _init_jit_backend() -> JITAllReduceBackend:
|
||||
return JITAllReduceBackend()
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def bench_one(
|
||||
backend,
|
||||
inp: torch.Tensor,
|
||||
warmup: int,
|
||||
iters: int,
|
||||
group: dist.ProcessGroup,
|
||||
register_input: bool,
|
||||
) -> float:
|
||||
@cache_once
|
||||
def _init_aot_backend() -> AOTAllReduceBackend:
|
||||
return AOTAllReduceBackend()
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_fi_backend() -> FlashInferAllReduceBackend:
|
||||
return FlashInferAllReduceBackend()
|
||||
|
||||
|
||||
BACKEND_FACTORY = {
|
||||
"nccl": _init_nccl_backend,
|
||||
"jit": _init_jit_backend,
|
||||
"aot": _init_aot_backend,
|
||||
"fi": _init_fi_backend,
|
||||
}
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_all_backends() -> None:
|
||||
"""Pre-build every supported backend before any timed iteration so JIT
|
||||
compilation / IPC setup don't bleed into the first measured size.
|
||||
"""
|
||||
Run *warmup* iterations of all-reduce first.
|
||||
Return the average time for *iters* iterations of all-reduce.
|
||||
"""
|
||||
dist.barrier(group=group)
|
||||
for _ in range(warmup):
|
||||
backend.all_reduce(inp)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Capture a CUDA graph with *iters* all-reduce calls.
|
||||
inp_batch = torch.stack([inp] * 4)
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with backend.capture(register_input):
|
||||
with torch.cuda.graph(graph):
|
||||
for i in range(iters):
|
||||
backend.all_reduce(inp_batch[i % 4])
|
||||
|
||||
torch.cuda.synchronize()
|
||||
# Warm up the graph once.
|
||||
graph.replay()
|
||||
|
||||
# Timed replay.
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier(group=group)
|
||||
graph.replay() # make the stream busy
|
||||
start.record()
|
||||
graph.replay()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
return start.elapsed_time(end) / iters
|
||||
|
||||
|
||||
def bench_sweep(
|
||||
backend,
|
||||
sizes_bytes: List[int],
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
warmup: int,
|
||||
iters: int,
|
||||
group: dist.ProcessGroup,
|
||||
register_input: bool,
|
||||
) -> Dict[int, float]:
|
||||
"""Benchmark one backend over all message sizes."""
|
||||
elem_size = torch.tensor([], dtype=dtype).element_size()
|
||||
results: Dict[int, float] = {}
|
||||
for sz in sizes_bytes:
|
||||
numel = sz // elem_size
|
||||
inp = torch.zeros(numel, dtype=dtype, device=device)
|
||||
try:
|
||||
elapsed_ms = bench_one(backend, inp, warmup, iters, group, register_input)
|
||||
results[sz] = elapsed_ms * 1000 # convert to us per iter
|
||||
except AssertionError:
|
||||
results[sz] = float("nan")
|
||||
return results
|
||||
world_size = dist.get_world_size(_init_cpu_group())
|
||||
factories = dict(BACKEND_FACTORY)
|
||||
if world_size not in AOT_SUPPORTED_WORLD_SIZES:
|
||||
factories.pop("aot")
|
||||
if world_size not in FI_SUPPORTED_WORLD_SIZES:
|
||||
factories.pop("fi")
|
||||
for fn in factories.values():
|
||||
fn()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Result printing
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def print_results(
|
||||
backends: list,
|
||||
all_results: Dict[str, Dict[int, float]],
|
||||
sizes_bytes: List[int],
|
||||
) -> None:
|
||||
"""Print a comparison table on rank 0."""
|
||||
|
||||
def human_bytes(n: int) -> str:
|
||||
for suffix, unit in [("M", 1 << 20), ("K", 1 << 10)]:
|
||||
if n >= unit and n % unit == 0:
|
||||
return f"{n // unit}{suffix}"
|
||||
return f"{n}B"
|
||||
|
||||
def fmt_us(v: float) -> str:
|
||||
return f"{v:13.1f}" if not isnan(v) else " n/a"
|
||||
|
||||
names = [b.name for b in backends]
|
||||
nccl_name = "NCCL"
|
||||
|
||||
# Header
|
||||
header_cols = [f"{n:>13}" for n in names]
|
||||
speedup_cols = [f"{n:>13}/NCCL" for n in names if n != nccl_name]
|
||||
header = f"{'Size':>8} " + " ".join(header_cols)
|
||||
for sc in speedup_cols:
|
||||
header += f" {sc}"
|
||||
header += " "
|
||||
print()
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
|
||||
# Rows
|
||||
for sz in sizes_bytes:
|
||||
row = f"{human_bytes(sz):>8}"
|
||||
nccl_lat = all_results[nccl_name][sz]
|
||||
for n in names:
|
||||
row += f" {fmt_us(all_results[n][sz])}"
|
||||
for n in names:
|
||||
if n == nccl_name:
|
||||
continue
|
||||
lat = all_results[n][sz]
|
||||
if not isnan(lat):
|
||||
row += f" {nccl_lat / lat:17.2f}x"
|
||||
else:
|
||||
row += f" {'n/a':>17}"
|
||||
print(row)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Distributed setup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def init_distributed():
|
||||
"""Initialize distributed groups using torchrun env vars.
|
||||
|
||||
Returns (rank, world_size, device, cpu_group, nccl_group).
|
||||
"""
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
||||
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
||||
rank = local_rank
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
torch.cuda.set_stream(torch.cuda.Stream()) # use a non-default stream
|
||||
|
||||
torch.distributed.init_process_group(backend="gloo")
|
||||
ps._WORLD = coord = ps.init_world_group(
|
||||
ranks=list(range(world_size)),
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
@marker.parametrize("message_bytes", MESSAGE_SIZES_BYTES)
|
||||
@marker.benchmark("provider", PROVIDERS)
|
||||
def benchmark(message_bytes: int, provider: str):
|
||||
cpu_group = _init_cpu_group()
|
||||
gpu_group = _init_nccl_group()
|
||||
world_size = dist.get_world_size(cpu_group)
|
||||
if provider == "fi" and world_size not in FI_SUPPORTED_WORLD_SIZES:
|
||||
marker.skip(
|
||||
f"flashinfer trtllm allreduce_fusion needs world_size in "
|
||||
f"{FI_SUPPORTED_WORLD_SIZES}"
|
||||
)
|
||||
if provider == "aot" and world_size not in AOT_SUPPORTED_WORLD_SIZES:
|
||||
marker.skip(
|
||||
f"AOT custom_all_reduce needs world_size in " f"{AOT_SUPPORTED_WORLD_SIZES}"
|
||||
)
|
||||
_init_all_backends()
|
||||
backend = BACKEND_FACTORY[provider]()
|
||||
numel = message_bytes // DTYPE_ITEMSIZE
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
x = torch.randn(numel, dtype=DTYPE, device=device)
|
||||
# Bandwidth-equivalent bytes moved by a ring all-reduce per rank.
|
||||
effective_bytes = int(x.nbytes * 2 * (world_size - 1) / world_size)
|
||||
return marker.do_bench(
|
||||
backend.all_reduce,
|
||||
input_args=(x,),
|
||||
graph_context_fn=backend.graph_context,
|
||||
sync_multigpu_fn=lambda: dist.barrier(gpu_group),
|
||||
# all-reduce is in-place w.r.t. its argument; explicit footprint
|
||||
# captures the cross-GPU traffic instead.
|
||||
memory_args=None,
|
||||
memory_output=None,
|
||||
extra_memory_footprint=effective_bytes,
|
||||
)
|
||||
|
||||
cpu_group = coord.cpu_group
|
||||
nccl_group = coord.device_group
|
||||
assert nccl_group is not None
|
||||
return rank, world_size, device, cpu_group, nccl_group
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main():
|
||||
logging.basicConfig(level=logging.WARNING)
|
||||
args = parse_args()
|
||||
dtype = DTYPE_MAP[args.dtype]
|
||||
|
||||
rank, world_size, device, cpu_group, nccl_group = init_distributed()
|
||||
|
||||
# Instantiate backends.
|
||||
backends = [
|
||||
NCCLAllReduceBackend(nccl_group),
|
||||
JITAllReduceBackend(cpu_group, device),
|
||||
]
|
||||
if world_size in [2, 4, 6, 8]:
|
||||
backends.insert(1, AOTAllReduceBackend(cpu_group, device))
|
||||
if world_size in [2, 4, 8]:
|
||||
backends.append(FlashInferAllReduceBackend(cpu_group, dtype))
|
||||
|
||||
# Run benchmarks.
|
||||
all_results: Dict[str, Dict[int, float]] = {}
|
||||
torch.cuda.synchronize()
|
||||
for backend in backends:
|
||||
if rank == 0:
|
||||
print(f"Benchmarking {backend.name} ...")
|
||||
all_results[backend.name] = bench_sweep(
|
||||
backend,
|
||||
MESSAGE_SIZES_BYTES,
|
||||
dtype,
|
||||
device,
|
||||
args.warmup,
|
||||
args.iters,
|
||||
cpu_group,
|
||||
args.register_input,
|
||||
)
|
||||
|
||||
# Aggregate across ranks (use max to reflect the slowest rank).
|
||||
for name in list(all_results):
|
||||
for sz in MESSAGE_SIZES_BYTES:
|
||||
val = all_results[name].get(sz)
|
||||
if val is None:
|
||||
continue
|
||||
t = torch.tensor([val], dtype=torch.float64, device=device)
|
||||
dist.all_reduce(t, op=dist.ReduceOp.MAX, group=nccl_group)
|
||||
all_results[name][sz] = t.item()
|
||||
|
||||
# Print results on rank 0.
|
||||
if rank == 0:
|
||||
print_results(backends, all_results, MESSAGE_SIZES_BYTES)
|
||||
|
||||
del backends, all_results
|
||||
gc.collect()
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__" and not is_in_ci():
|
||||
main()
|
||||
if __name__ == "__main__":
|
||||
multigpu_bench_main(
|
||||
name=__name__,
|
||||
file=__file__,
|
||||
num_gpus=WORLD_SIZES,
|
||||
main_fn=benchmark.run,
|
||||
)
|
||||
|
||||
@@ -1,17 +1,39 @@
|
||||
"""Benchmark fused TP QKNorm (push-mode custom-AR + RMSNorm) vs the serial
|
||||
baseline (RMS sum-sq -> pull-mode all-reduce -> RMS apply).
|
||||
|
||||
Usage::
|
||||
|
||||
# Benchmark on every supported world size (2..8 GPUs):
|
||||
python benchmark/bench_tp_qknorm.py
|
||||
# Specific world sizes:
|
||||
python benchmark/bench_tp_qknorm.py --num-gpu 4
|
||||
python benchmark/bench_tp_qknorm.py --num-gpu 2,4,8
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import atexit
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
from multiprocessing.context import SpawnProcess
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
from sglang.jit_kernel.all_reduce import (
|
||||
_jit_custom_all_reduce_pull_module,
|
||||
_jit_custom_all_reduce_push_module,
|
||||
_jit_fused_parallel_qknorm_module,
|
||||
fused_parallel_qknorm,
|
||||
get_fused_parallel_qknorm_max_occupancy,
|
||||
)
|
||||
from sglang.jit_kernel.utils import get_ci_test_range
|
||||
from sglang.jit_kernel.benchmark import marker
|
||||
from sglang.jit_kernel.benchmark.utils import multigpu_bench_main
|
||||
from sglang.jit_kernel.mp import register_comm_cleanup
|
||||
from sglang.jit_kernel.utils import cache_once, get_ci_test_range
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
@@ -23,80 +45,131 @@ register_cuda_ci(
|
||||
disabled="requires multi-GPU, self-skips in CI",
|
||||
)
|
||||
|
||||
Q_K_DIMS = [(6144, 1024)]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sweep parameters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
DTYPE = torch.bfloat16
|
||||
EPS = 1e-6
|
||||
Q_K_DIMS = [(6144, 1024)]
|
||||
BATCH_SIZES = get_ci_test_range([2**i for i in range(15)], [1, 64, 1024])
|
||||
NUM_LAYERS = 8
|
||||
MAX_PUSH_SIZE = 8 * max(BATCH_SIZES)
|
||||
PROVIDERS = ["fused", "baseline"]
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--warmup", type=int, default=10)
|
||||
parser.add_argument("--iters", type=int, default=100)
|
||||
return parser.parse_args()
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parallel JIT precompile (outer process, before any torchrun child starts)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def init_distributed():
|
||||
def _compile_one(world_size: int) -> None:
|
||||
"""Compile every kernel this bench touches for a single world_size.
|
||||
|
||||
Top-level so it survives ``spawn`` pickling. Compiled artifacts are
|
||||
cached on disk by ``tvm_ffi``; torchrun children will reuse them.
|
||||
"""
|
||||
# baseline path: sum-sq -> pull-mode all-reduce -> apply
|
||||
_jit_custom_all_reduce_pull_module(DTYPE, world_size)
|
||||
# fused path: push-mode all-reduce
|
||||
_jit_custom_all_reduce_push_module(DTYPE, world_size)
|
||||
# fused path: fused QKNorm kernel (one per (dtype, world_size, q_dim, k_dim))
|
||||
for q_dim, k_dim in Q_K_DIMS:
|
||||
_jit_fused_parallel_qknorm_module(DTYPE, world_size, q_dim, k_dim)
|
||||
|
||||
|
||||
def _precompile_kernels(num_gpus: List[int]) -> None:
|
||||
ctx = multiprocessing.get_context("spawn")
|
||||
procs: list[tuple[int, SpawnProcess]] = []
|
||||
for world_size in num_gpus:
|
||||
p = ctx.Process(target=_compile_one, args=(world_size,))
|
||||
p.start()
|
||||
procs.append((world_size, p))
|
||||
for world_size, p in procs:
|
||||
p.join()
|
||||
if p.exitcode != 0:
|
||||
raise RuntimeError(
|
||||
f"TP QKNorm precompile failed for {world_size=} " f"(exit {p.exitcode})"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-rank distributed init (run once per torchrun worker)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_cpu_group() -> dist.ProcessGroup:
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
rank = local_rank
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(backend="gloo")
|
||||
ps._WORLD = coord = ps.init_world_group(
|
||||
ranks=list(range(world_size)),
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
)
|
||||
atexit.register(dist.destroy_process_group)
|
||||
logging.disable(logging.INFO)
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
return coord.cpu_group
|
||||
|
||||
cpu_group = coord.cpu_group
|
||||
|
||||
@cache_once
|
||||
def _init_gpu_group() -> dist.ProcessGroup:
|
||||
_init_cpu_group()
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
gpu_group = dist.new_group(backend="nccl", device_id=device)
|
||||
assert isinstance(gpu_group, dist.ProcessGroup)
|
||||
atexit.register(lambda: dist.destroy_process_group(gpu_group))
|
||||
return gpu_group
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_fused_comm() -> CustomAllReduceV2:
|
||||
"""Push-mode workspace sized for the fused-QKNorm bench."""
|
||||
cpu_group = _init_cpu_group()
|
||||
world_size = dist.get_world_size(cpu_group)
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
q_dim, k_dim = Q_K_DIMS[0]
|
||||
max_occupancy = get_fused_parallel_qknorm_max_occupancy(
|
||||
DTYPE, world_size, Q_K_DIMS[0][0], Q_K_DIMS[0][1]
|
||||
DTYPE, world_size, q_dim, k_dim
|
||||
)
|
||||
if rank == 0:
|
||||
if dist.get_rank(cpu_group) == 0:
|
||||
print(f"Max occupancy for fused_parallel_qknorm: {max_occupancy} blocks/SM")
|
||||
|
||||
props = torch.cuda.get_device_properties(device)
|
||||
comm = CustomAllReduceV2(
|
||||
cpu_group,
|
||||
device,
|
||||
max_pull_size=0,
|
||||
max_push_size=8 * max(BATCH_SIZES),
|
||||
max_push_size=MAX_PUSH_SIZE,
|
||||
max_push_blocks=props.multi_processor_count * max_occupancy,
|
||||
)
|
||||
comm_ = CustomAllReduceV2(cpu_group, device)
|
||||
if comm.disabled or comm_.disabled:
|
||||
if comm.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
return rank, world_size, device, cpu_group, comm, comm_
|
||||
register_comm_cleanup(comm)
|
||||
return comm
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def bench_one(fn, warmup: int, iters: int) -> float:
|
||||
for _ in range(warmup):
|
||||
fn(0)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
for i in range(NUM_LAYERS):
|
||||
fn(i)
|
||||
|
||||
graph.replay()
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
graph.replay()
|
||||
start.record()
|
||||
for i in range(iters):
|
||||
graph.replay()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
return start.elapsed_time(end) * 1000.0 / (iters * NUM_LAYERS)
|
||||
@cache_once
|
||||
def _init_baseline_comm() -> CustomAllReduceV2:
|
||||
"""Default (pull-mode) workspace for the serial baseline."""
|
||||
cpu_group = _init_cpu_group()
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
comm = CustomAllReduceV2(cpu_group, device)
|
||||
if comm.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
register_comm_cleanup(comm)
|
||||
return comm
|
||||
|
||||
|
||||
def rmsnorm_baseline(
|
||||
comm_,
|
||||
# ---------------------------------------------------------------------------
|
||||
# Implementations
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _rmsnorm_baseline(
|
||||
comm: CustomAllReduceV2,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
@@ -106,65 +179,56 @@ def rmsnorm_baseline(
|
||||
from sglang.srt.models.minimax_m2 import rms_apply_serial, rms_sumsq_serial
|
||||
|
||||
sum_sq = rms_sumsq_serial(q, k)
|
||||
sum_sq = comm_.custom_all_reduce(sum_sq)
|
||||
sum_sq = comm.custom_all_reduce(sum_sq)
|
||||
rms_apply_serial(q, k, q_weight, k_weight, sum_sq, world_size, EPS)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
rank, world_size, device, _, comm, comm_ = init_distributed()
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if rank == 0:
|
||||
print(
|
||||
f"{'q_dim':>8} {'k_dim':>8} {'batch':>8} {'fused_us':>12} {'baseline_us':>12}"
|
||||
)
|
||||
|
||||
for q_dim, k_dim in Q_K_DIMS:
|
||||
local_q_dim = q_dim // world_size
|
||||
local_k_dim = k_dim // world_size
|
||||
for batch_size in BATCH_SIZES:
|
||||
q = torch.randn(
|
||||
NUM_LAYERS, batch_size, local_q_dim, device=device, dtype=DTYPE
|
||||
)
|
||||
k = torch.randn(
|
||||
NUM_LAYERS, batch_size, local_k_dim, device=device, dtype=DTYPE
|
||||
)
|
||||
q_weight = torch.randn(NUM_LAYERS, local_q_dim, device=device, dtype=DTYPE)
|
||||
k_weight = torch.randn(NUM_LAYERS, local_k_dim, device=device, dtype=DTYPE)
|
||||
@marker.parametrize("q_dim,k_dim", Q_K_DIMS)
|
||||
@marker.parametrize("batch_size", BATCH_SIZES)
|
||||
@marker.benchmark("provider", PROVIDERS)
|
||||
def benchmark(q_dim: int, k_dim: int, batch_size: int, provider: str):
|
||||
cpu_group = _init_cpu_group()
|
||||
gpu_group = _init_gpu_group()
|
||||
world_size = dist.get_world_size(cpu_group)
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
local_q_dim = q_dim // world_size
|
||||
local_k_dim = k_dim // world_size
|
||||
|
||||
def run_fused(i: int):
|
||||
fused_parallel_qknorm(
|
||||
comm.obj,
|
||||
q[i],
|
||||
k[i],
|
||||
q_weight[i],
|
||||
k_weight[i],
|
||||
EPS,
|
||||
)
|
||||
q = torch.randn(batch_size, local_q_dim, device=device, dtype=DTYPE)
|
||||
k = torch.randn(batch_size, local_k_dim, device=device, dtype=DTYPE)
|
||||
q_weight = torch.randn(local_q_dim, device=device, dtype=DTYPE)
|
||||
k_weight = torch.randn(local_k_dim, device=device, dtype=DTYPE)
|
||||
|
||||
def run_baseline(i: int):
|
||||
rmsnorm_baseline(
|
||||
comm_,
|
||||
q[i],
|
||||
k[i],
|
||||
q_weight[i],
|
||||
k_weight[i],
|
||||
world_size,
|
||||
)
|
||||
if provider == "fused":
|
||||
comm = _init_fused_comm()
|
||||
|
||||
fused_us = bench_one(run_fused, args.warmup, args.iters)
|
||||
baseline_us = bench_one(run_baseline, args.warmup, args.iters)
|
||||
def fn(q, k, q_weight, k_weight):
|
||||
fused_parallel_qknorm(comm.obj, q, k, q_weight, k_weight, EPS)
|
||||
|
||||
if rank == 0:
|
||||
print(
|
||||
f"{q_dim:8d} {k_dim:8d} {batch_size:8d} "
|
||||
f"{fused_us:12.1f} {baseline_us:12.1f}"
|
||||
)
|
||||
else:
|
||||
comm = _init_baseline_comm()
|
||||
|
||||
comm.close()
|
||||
dist.destroy_process_group()
|
||||
def fn(q, k, q_weight, k_weight):
|
||||
_rmsnorm_baseline(comm, q, k, q_weight, k_weight, world_size)
|
||||
|
||||
return marker.do_bench(
|
||||
fn,
|
||||
input_args=(q, k, q_weight, k_weight),
|
||||
sync_multigpu_fn=lambda: dist.barrier(gpu_group),
|
||||
memory_output=(q, k), # NOTE: In-place updates on q, k;
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
multigpu_bench_main(
|
||||
name=__name__,
|
||||
file=__file__,
|
||||
num_gpus=[2, 4, 8], # NOTE: don't support other world size now
|
||||
main_fn=benchmark.run,
|
||||
pre_launch_fn=_precompile_kernels,
|
||||
)
|
||||
|
||||
@@ -1,28 +1,34 @@
|
||||
"""
|
||||
Correctness test for the JIT custom all-reduce (v2) kernel.
|
||||
"""Correctness test for the JIT custom all-reduce (v2) kernel.
|
||||
|
||||
The test compares the JIT custom all-reduce output against NCCL all-reduce
|
||||
for various tensor sizes and dtypes, in both eager and CUDA-graph modes.
|
||||
Compares the JIT custom all-reduce output against NCCL all-reduce for a sweep
|
||||
of tensor sizes, dtypes, and algorithms, in both eager and CUDA-graph modes.
|
||||
|
||||
Usage:
|
||||
python -m pytest test_jit_custom_all_reduce.py -v
|
||||
Usage::
|
||||
|
||||
This file doubles as the torchrun worker script. The test class launches
|
||||
torchrun --nproc_per_node=N <this_file>
|
||||
and asserts that all worker processes exit successfully.
|
||||
# Run the test on the default world sizes (2, 4, 8 GPUs):
|
||||
python tests/test_custom_all_reduce.py
|
||||
# Pick a specific world size (or comma-separated list), e.g. the rarer
|
||||
# odd / non-power-of-two counts that the default sweep skips:
|
||||
python tests/test_custom_all_reduce.py --num-gpu 3
|
||||
python tests/test_custom_all_reduce.py --num-gpu 2,4,6,8
|
||||
# Extra pytest args (forwarded to each torchrun worker):
|
||||
python tests/test_custom_all_reduce.py -k bfloat16
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import itertools
|
||||
import logging
|
||||
import multiprocessing as mp
|
||||
import multiprocessing
|
||||
import os
|
||||
from typing import Dict, Optional, Tuple
|
||||
from multiprocessing.context import SpawnProcess
|
||||
from typing import List
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import triton
|
||||
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
from sglang.jit_kernel.all_reduce import (
|
||||
@@ -30,7 +36,9 @@ from sglang.jit_kernel.all_reduce import (
|
||||
_jit_custom_all_reduce_pull_module,
|
||||
_jit_custom_all_reduce_push_module,
|
||||
)
|
||||
from sglang.jit_kernel.tests.utils import multiprocess_main, multiprocess_test
|
||||
from sglang.jit_kernel.mp import register_comm_cleanup
|
||||
from sglang.jit_kernel.tests.utils import multigpu_pytest_main
|
||||
from sglang.jit_kernel.utils import cache_once, get_ci_test_range
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
@@ -47,7 +55,7 @@ register_cuda_ci(
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test parameters (shared between test class and worker)
|
||||
# Test parameters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
TEST_SIZES = [
|
||||
@@ -59,181 +67,172 @@ TEST_SIZES = [
|
||||
4 * 1024,
|
||||
32 * 1024,
|
||||
256 * 1024,
|
||||
2 * 1024 * 1024, # 2M elements
|
||||
4 * 1024 * 1024, # 4M elements
|
||||
2 * 1024 * 1024,
|
||||
4 * 1024 * 1024,
|
||||
]
|
||||
TEST_DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
||||
SHOTS = [
|
||||
TEST_ALGOS = [
|
||||
AllReduceAlgo.ONE_SHOT_PULL,
|
||||
AllReduceAlgo.ONE_SHOT_PUSH,
|
||||
AllReduceAlgo.TWO_SHOT_PULL,
|
||||
]
|
||||
USE_GRAPH_OPTIONS = [True, False]
|
||||
TEST_CONFIG = itertools.product(TEST_SIZES, TEST_DTYPES, SHOTS, USE_GRAPH_OPTIONS)
|
||||
USE_GRAPH_OPTIONS = [False, True]
|
||||
TEST_LAYERS = 4
|
||||
TEST_LOOP = 16
|
||||
|
||||
TEST_SIZES = get_ci_test_range(TEST_SIZES, [16, 1024, 32 * 1024, 2 * 1024 * 1024])
|
||||
TEST_DTYPES = get_ci_test_range(TEST_DTYPES, [torch.bfloat16])
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test class (runs via pytest, launches torchrun subprocesses)
|
||||
# Parallel JIT precompile (outer process, before any torchrun child starts)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _compile_one(dtype: torch.dtype, world_size: int):
|
||||
_jit_custom_all_reduce_push_module(dtype, world_size)
|
||||
_jit_custom_all_reduce_pull_module(dtype, world_size)
|
||||
def _compile_one(dtype: torch.dtype, world_size: int) -> None:
|
||||
"""Compile both (push, pull) variants for a single (dtype, world_size).
|
||||
|
||||
|
||||
def _precompile_kernels() -> None:
|
||||
# NOTE: even when device count < 8, we should be able to compile all
|
||||
process_map: Dict[Tuple[torch.dtype, int], mp.Process] = {}
|
||||
COMPILE_SPACE = itertools.product(TEST_DTYPES, [2, 3, 4, 5, 6, 7, 8])
|
||||
mp.set_start_method("spawn")
|
||||
for config in COMPILE_SPACE:
|
||||
process_map[config] = mp.Process(target=_compile_one, args=config)
|
||||
for process in process_map.values():
|
||||
process.start()
|
||||
for (dtype, world_size), process in process_map.items():
|
||||
process.join()
|
||||
if process.exitcode != 0:
|
||||
raise RuntimeError(f"Custom All Reduce {world_size=} {dtype=} failed")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nproc", [1, 2, 3, 4, 5, 6, 7, 8])
|
||||
def test_custom_allreduce(nproc: int) -> None:
|
||||
if nproc == 1: # NOTE: special case to speed up tests
|
||||
return _precompile_kernels()
|
||||
|
||||
device_count = torch.cuda.device_count()
|
||||
if device_count < nproc:
|
||||
pytest.skip(
|
||||
f"Requires at least {nproc} GPUs, but only {device_count} available"
|
||||
)
|
||||
multiprocess_test(__file__, nproc)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Worker logic (executed by each torchrun process)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def init_distributed():
|
||||
"""Initialize distributed groups via torchrun env vars.
|
||||
|
||||
Returns (rank, device, cpu_group, nccl_group, comm).
|
||||
Top-level so it survives ``spawn`` pickling. Compiled artifacts are
|
||||
cached on disk by ``tvm_ffi``; torchrun children will reuse them.
|
||||
"""
|
||||
_jit_custom_all_reduce_pull_module(dtype, world_size)
|
||||
_jit_custom_all_reduce_push_module(dtype, world_size)
|
||||
|
||||
|
||||
def _precompile_kernels(num_gpus: List[int]) -> None:
|
||||
"""Fan out one process per (dtype, world_size) to warm the JIT cache.
|
||||
|
||||
Without this, every torchrun child serial-compiles its kernels on first
|
||||
use, multiplying the wall-clock cost of the run by ~(#dtypes * #ranks).
|
||||
"""
|
||||
ctx = multiprocessing.get_context("spawn")
|
||||
procs: list[tuple[torch.dtype, int, SpawnProcess]] = []
|
||||
for dtype, world_size in itertools.product(TEST_DTYPES, num_gpus):
|
||||
p = ctx.Process(target=_compile_one, args=(dtype, world_size))
|
||||
p.start()
|
||||
procs.append((dtype, world_size, p))
|
||||
for dtype, world_size, p in procs:
|
||||
p.join()
|
||||
if p.exitcode != 0:
|
||||
raise RuntimeError(
|
||||
f"Custom-all-reduce precompile failed for "
|
||||
f"{dtype=} {world_size=} (exit {p.exitcode})"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-rank distributed setup (run once per torchrun worker)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_cpu_group_once() -> dist.ProcessGroup:
|
||||
"""Initialize gloo world group + cuda device for this rank."""
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
rank = local_rank
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(backend="gloo")
|
||||
ps._WORLD = coord = ps.init_world_group(
|
||||
ranks=list(range(world_size)),
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
)
|
||||
|
||||
atexit.register(dist.destroy_process_group)
|
||||
cpu_group = coord.cpu_group
|
||||
nccl_group = coord.device_group
|
||||
assert nccl_group is not None
|
||||
assert isinstance(cpu_group, dist.ProcessGroup)
|
||||
# Suppress chatty internal logging for cleaner test output.
|
||||
logging.disable(logging.INFO)
|
||||
# Use a non-default stream (mirrors prior behavior).
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
return cpu_group
|
||||
|
||||
max_size = max(TEST_SIZES) * 4
|
||||
|
||||
@cache_once
|
||||
def _init_nccl_group_once() -> dist.ProcessGroup:
|
||||
_init_cpu_group_once()
|
||||
coord = ps._WORLD
|
||||
assert coord is not None and coord.device_group is not None
|
||||
return coord.device_group
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_comm_once() -> CustomAllReduceV2:
|
||||
cpu_group = _init_cpu_group_once()
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
max_size = max(TEST_SIZES) * max(
|
||||
torch.tensor([], dtype=d).element_size() for d in TEST_DTYPES
|
||||
)
|
||||
comm = CustomAllReduceV2(cpu_group, device, max_size, max_size)
|
||||
if comm.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
|
||||
return rank, device, cpu_group, nccl_group, comm
|
||||
register_comm_cleanup(comm)
|
||||
return comm
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_graph", USE_GRAPH_OPTIONS)
|
||||
@pytest.mark.parametrize("algo", TEST_ALGOS)
|
||||
@pytest.mark.parametrize("dtype", TEST_DTYPES)
|
||||
@pytest.mark.parametrize("size", TEST_SIZES)
|
||||
@torch.inference_mode()
|
||||
def worker_test(
|
||||
device: torch.device,
|
||||
nccl_group: dist.ProcessGroup,
|
||||
comm: CustomAllReduceV2,
|
||||
def test_custom_all_reduce(
|
||||
size: int,
|
||||
dtype: torch.dtype,
|
||||
use_graph: bool,
|
||||
algo: AllReduceAlgo,
|
||||
) -> Optional[RuntimeError]:
|
||||
use_graph: bool,
|
||||
) -> None:
|
||||
nccl_group = _init_nccl_group_once()
|
||||
comm = _init_comm_once()
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
comm.override_algo = algo
|
||||
|
||||
def get_run_graph_fn():
|
||||
if use_graph:
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
graph_inp = torch.zeros((TEST_LAYERS, size), dtype=dtype, device=device)
|
||||
out_jits = []
|
||||
outs: list[torch.Tensor] = []
|
||||
with comm.capture():
|
||||
with torch.cuda.graph(graph):
|
||||
for i in range(TEST_LAYERS):
|
||||
out_jits.append(comm.custom_all_reduce(graph_inp[i]))
|
||||
out_jit = torch.stack(out_jits)
|
||||
outs.append(comm.custom_all_reduce(graph_inp[i]))
|
||||
out_jit_stack = torch.stack(outs)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
def run_graph(x: torch.Tensor) -> torch.Tensor:
|
||||
def run(x: torch.Tensor) -> torch.Tensor:
|
||||
graph_inp.copy_(x)
|
||||
graph.replay()
|
||||
return out_jit.clone()
|
||||
return out_jit_stack.clone()
|
||||
|
||||
return run_graph
|
||||
else:
|
||||
|
||||
def get_run_eager_fn():
|
||||
def run_eager(x: torch.Tensor) -> torch.Tensor:
|
||||
def run(x: torch.Tensor) -> torch.Tensor:
|
||||
eager_inp = x.clone()
|
||||
out_eagers = []
|
||||
outs = []
|
||||
for i in range(TEST_LAYERS):
|
||||
out_eagers.append(comm.custom_all_reduce(eager_inp[i]))
|
||||
outs.append(comm.custom_all_reduce(eager_inp[i]))
|
||||
torch.cuda.synchronize()
|
||||
return torch.stack(out_eagers)
|
||||
return torch.stack(outs)
|
||||
|
||||
return run_eager
|
||||
|
||||
run_fn = get_run_graph_fn() if use_graph else get_run_eager_fn()
|
||||
num_errors = 0
|
||||
for _ in range(TEST_LOOP):
|
||||
# NOTE: 15 * 8 < 128, which is the precision limit for bf16
|
||||
inp = torch.randint(0, 16, (TEST_LAYERS, size), dtype=dtype, device=device)
|
||||
assert comm.should_custom_ar(inp[0])
|
||||
out_ref = inp.clone()
|
||||
dist.all_reduce(out_ref, group=nccl_group)
|
||||
out_jit = run_fn(inp)
|
||||
num_errors += not torch.all(out_jit == out_ref)
|
||||
if num_errors > 0:
|
||||
return RuntimeError(
|
||||
f"Test failed for {size=}, {dtype=}, {algo=}, "
|
||||
f"{use_graph=} with {num_errors} errors. "
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def worker_main() -> None:
|
||||
"""Entry point for each torchrun worker process."""
|
||||
rank, device, cpu_group, nccl_group, comm = init_distributed()
|
||||
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
|
||||
logging.disable(logging.INFO) # Suppress internal logging for cleaner test output
|
||||
items = list(enumerate(TEST_CONFIG))
|
||||
for i, (size, dtype, algo, use_graph) in items:
|
||||
error = worker_test(device, nccl_group, comm, size, dtype, use_graph, algo)
|
||||
if error is not None:
|
||||
print(
|
||||
f"Worker {rank} failed for {size=}, {dtype=}, "
|
||||
f"{algo=}, {use_graph=}, iteration={i}\n"
|
||||
f"Error: {error}"
|
||||
)
|
||||
# communicate the result to rank 0 for logging
|
||||
result = torch.tensor([int(error is not None)])
|
||||
dist.all_reduce(result, group=cpu_group)
|
||||
failed = bool(result.item())
|
||||
if failed:
|
||||
raise RuntimeError(
|
||||
f"Test failed on rank {rank} for config: "
|
||||
f"{size=}, {dtype=}, {algo=}, {use_graph=}"
|
||||
)
|
||||
|
||||
comm.close()
|
||||
dist.destroy_process_group()
|
||||
out_jit = run(inp)
|
||||
# Exact equality, since values are small integers within bf16 precision.
|
||||
triton.testing.assert_close(out_ref, out_jit, atol=0, rtol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
multiprocess_main(__file__, worker_main)
|
||||
# Only sweep the common world sizes (2, 4, 8) by default: testing every
|
||||
# count in 2..8 serially overruns the per-file CI time budget, and 3/5/6/7
|
||||
# are rare in practice. Use --num-gpu to exercise them explicitly.
|
||||
multigpu_pytest_main(
|
||||
__name__,
|
||||
__file__,
|
||||
num_gpus=(2, 4, 8),
|
||||
pre_launch_fn=_precompile_kernels,
|
||||
)
|
||||
|
||||
@@ -1,16 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import itertools
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
from typing import Optional
|
||||
from multiprocessing.context import SpawnProcess
|
||||
from typing import List
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.all_reduce import fused_parallel_qknorm
|
||||
from sglang.jit_kernel.tests.utils import multiprocess_main, multiprocess_test
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
from sglang.jit_kernel.all_reduce import (
|
||||
_jit_custom_all_reduce_push_module,
|
||||
_jit_fused_parallel_qknorm_module,
|
||||
fused_parallel_qknorm,
|
||||
)
|
||||
from sglang.jit_kernel.mp import register_comm_cleanup
|
||||
from sglang.jit_kernel.tests.utils import multigpu_pytest_main
|
||||
from sglang.jit_kernel.utils import cache_once
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
@@ -24,53 +38,96 @@ register_cuda_ci(
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test parameters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Q_K_DIMS = [(6144, 1024)]
|
||||
EPS = 1e-6
|
||||
BATCH_SIZES = [2**n for n in range(0, 14)]
|
||||
DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
||||
TEST_CONFIG = list(itertools.product(Q_K_DIMS, BATCH_SIZES, DTYPES))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nproc", [2, 4, 8])
|
||||
def test_tp_qknorm(nproc: int) -> None:
|
||||
device_count = torch.cuda.device_count()
|
||||
if device_count < nproc:
|
||||
pytest.skip(
|
||||
f"Requires at least {nproc} GPUs, but only {device_count} available"
|
||||
)
|
||||
multiprocess_test(__file__, nproc)
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parallel JIT precompile (outer process, before any torchrun child starts)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def init_distributed():
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
def _compile_one(dtype: torch.dtype, world_size: int) -> None:
|
||||
"""Compile every kernel this test touches for one (dtype, world_size).
|
||||
|
||||
Top-level so it survives ``spawn`` pickling. Compiled artifacts are
|
||||
cached on disk by ``tvm_ffi``; torchrun children will reuse them.
|
||||
"""
|
||||
_jit_custom_all_reduce_push_module(dtype, world_size)
|
||||
for q_dim, k_dim in Q_K_DIMS:
|
||||
_jit_fused_parallel_qknorm_module(dtype, world_size, q_dim, k_dim)
|
||||
|
||||
|
||||
def _precompile_kernels(num_gpus: List[int]) -> None:
|
||||
ctx = multiprocessing.get_context("spawn")
|
||||
procs: list[tuple[torch.dtype, int, SpawnProcess]] = []
|
||||
for dtype, world_size in itertools.product(DTYPES, num_gpus):
|
||||
p = ctx.Process(target=_compile_one, args=(dtype, world_size))
|
||||
p.start()
|
||||
procs.append((dtype, world_size, p))
|
||||
for dtype, world_size, p in procs:
|
||||
p.join()
|
||||
if p.exitcode != 0:
|
||||
raise RuntimeError(
|
||||
f"TP QKNorm precompile failed for {dtype=} {world_size=} "
|
||||
f"(exit {p.exitcode})"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-rank distributed setup (run once per torchrun worker)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_cpu_group_once() -> dist.ProcessGroup:
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
rank = local_rank
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(backend="gloo")
|
||||
ps._WORLD = coord = ps.init_world_group(
|
||||
ranks=list(range(world_size)),
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
)
|
||||
|
||||
atexit.register(dist.destroy_process_group)
|
||||
cpu_group = coord.cpu_group
|
||||
nccl_group = coord.device_group
|
||||
assert nccl_group is not None
|
||||
assert isinstance(cpu_group, dist.ProcessGroup)
|
||||
logging.disable(logging.INFO)
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
return cpu_group
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_nccl_group_once() -> dist.ProcessGroup:
|
||||
_init_cpu_group_once()
|
||||
coord = ps._WORLD
|
||||
assert coord is not None and coord.device_group is not None
|
||||
return coord.device_group
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_comm_once() -> CustomAllReduceV2:
|
||||
cpu_group = _init_cpu_group_once()
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
max_pull_size = 0
|
||||
max_push_size = 8 * max(BATCH_SIZES)
|
||||
comm = CustomAllReduceV2(cpu_group, device, max_pull_size, max_push_size)
|
||||
if comm.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
register_comm_cleanup(comm)
|
||||
return comm
|
||||
|
||||
return rank, world_size, device, cpu_group, nccl_group, comm
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _all_gather_cat(x: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
|
||||
@@ -85,17 +142,26 @@ def _rmsnorm_ref(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Ten
|
||||
return (x_fp32 * scale * weight.float()).to(x.dtype)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("q_k_dim", Q_K_DIMS)
|
||||
@torch.inference_mode()
|
||||
def worker_test(
|
||||
rank: int,
|
||||
world_size: int,
|
||||
device: torch.device,
|
||||
nccl_group: dist.ProcessGroup,
|
||||
comm,
|
||||
def test_tp_qknorm(
|
||||
q_k_dim: tuple[int, int],
|
||||
batch_size: int,
|
||||
dtype: torch.dtype,
|
||||
) -> Optional[RuntimeError]:
|
||||
) -> None:
|
||||
nccl_group = _init_nccl_group_once()
|
||||
comm = _init_comm_once()
|
||||
rank = dist.get_rank(group=nccl_group)
|
||||
world_size = dist.get_world_size(group=nccl_group)
|
||||
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
||||
|
||||
q_dim, k_dim = q_k_dim
|
||||
local_q_dim = q_dim // world_size
|
||||
local_k_dim = k_dim // world_size
|
||||
@@ -115,53 +181,16 @@ def worker_test(
|
||||
q_expected = q_expected[:, rank * local_q_dim : (rank + 1) * local_q_dim]
|
||||
k_expected = k_expected[:, rank * local_k_dim : (rank + 1) * local_k_dim]
|
||||
|
||||
fused_parallel_qknorm(
|
||||
comm.obj,
|
||||
q,
|
||||
k,
|
||||
q_weight,
|
||||
k_weight,
|
||||
EPS,
|
||||
)
|
||||
fused_parallel_qknorm(comm.obj, q, k, q_weight, k_weight, EPS)
|
||||
|
||||
try:
|
||||
triton.testing.assert_close(q, q_expected, atol=1e-2, rtol=1e-2)
|
||||
triton.testing.assert_close(k, k_expected, atol=1e-2, rtol=1e-2)
|
||||
except AssertionError as err:
|
||||
return RuntimeError(
|
||||
f"TP QKNorm mismatch for {batch_size=}, {dtype=}, {world_size=}, {rank=}: {err}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def worker_main() -> None:
|
||||
rank, world_size, device, cpu_group, nccl_group, comm = init_distributed()
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
|
||||
for q_k_dim, batch_size, dtype in TEST_CONFIG:
|
||||
error = worker_test(
|
||||
rank,
|
||||
world_size,
|
||||
device,
|
||||
nccl_group,
|
||||
comm,
|
||||
q_k_dim,
|
||||
batch_size,
|
||||
dtype,
|
||||
)
|
||||
result = torch.tensor([int(error is not None)])
|
||||
dist.all_reduce(result, group=cpu_group)
|
||||
if error is not None:
|
||||
print(str(error))
|
||||
if bool(result.item()):
|
||||
raise RuntimeError(
|
||||
f"TP QKNorm test failed for {q_k_dim=}, {batch_size=}, {dtype=}, {world_size=}"
|
||||
)
|
||||
|
||||
print(f"Rank {rank} passed all tests.")
|
||||
comm.close()
|
||||
dist.destroy_process_group()
|
||||
triton.testing.assert_close(q, q_expected, atol=1e-2, rtol=1e-2)
|
||||
triton.testing.assert_close(k, k_expected, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
multiprocess_main(__file__, worker_main)
|
||||
multigpu_pytest_main(
|
||||
__name__,
|
||||
__file__,
|
||||
num_gpus=(2, 4, 8),
|
||||
pre_launch_fn=_precompile_kernels,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user