285 lines
9.7 KiB
Python
285 lines
9.7 KiB
Python
from __future__ import annotations
|
|
|
|
import atexit
|
|
import contextlib
|
|
import logging
|
|
import os
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
import sglang.srt.distributed.parallel_state as ps
|
|
from sglang.kernels.jit.benchmark import marker
|
|
from sglang.kernels.jit.benchmark.utils import get_benchmark_range, multigpu_bench_main
|
|
from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl
|
|
from sglang.kernels.ops.communication.mp import register_comm_cleanup
|
|
from sglang.srt.runtime_context import get_parallel
|
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
|
|
|
register_cuda_ci(
|
|
est_time=120,
|
|
stage="base-b-kernel-benchmark",
|
|
runner_config="1-gpu-large",
|
|
disabled="requires multi-GPU, self-skips in CI",
|
|
)
|
|
register_amd_ci(est_time=120, stage="jit-kernel-benchmark", runner_config="amd")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Sweep parameters
|
|
# ---------------------------------------------------------------------------
|
|
|
|
DTYPE = torch.bfloat16
|
|
DTYPE_ITEMSIZE = DTYPE.itemsize
|
|
MESSAGE_SIZES_KB = [2**x for x in range(2, 17)]
|
|
MESSAGE_SIZES_KB += [192, 384, 640, 768, 896, 1536, 3072]
|
|
MESSAGE_SIZES_KB.sort()
|
|
WORLD_SIZES = list(range(2, 9)) + [16]
|
|
MAX_BYTES = max(MESSAGE_SIZES_KB) * 1024
|
|
# 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)
|
|
# jit-eager times the naive-loop dispatch (eager heuristics); jit-graph
|
|
# captures the calls in a CUDA graph (graph heuristics + pointer table).
|
|
PROVIDERS = ["nccl", "aot", "jit-eager", "jit-graph", "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",
|
|
)
|
|
# The context answers a handle from what was stated on it, not from
|
|
# this module global, so a rank stood up by hand says so itself.
|
|
get_parallel().override_permanently(world_group=coord)
|
|
atexit.register(dist.destroy_process_group)
|
|
torch.cuda.set_stream(torch.cuda.Stream())
|
|
return coord.cpu_group
|
|
|
|
|
|
@cache_once
|
|
def _init_nccl_group() -> dist.ProcessGroup:
|
|
_init_cpu_group()
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
device_group = torch.distributed.new_group(
|
|
backend="nccl",
|
|
device_id=torch.device(f"cuda:{local_rank}"),
|
|
)
|
|
assert isinstance(device_group, dist.ProcessGroup)
|
|
return device_group
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Backend wrappers - each exposes:
|
|
# .all_reduce(tensor) -> Tensor
|
|
# .graph_context() -> context manager wrapping cuda-graph capture
|
|
# (nullcontext when capture is not required)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class NCCLAllReduceBackend:
|
|
def __init__(self) -> None:
|
|
self.group = _init_nccl_group()
|
|
|
|
def graph_context(self):
|
|
return contextlib.nullcontext()
|
|
|
|
def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
|
|
dist.all_reduce(tensor, group=self.group)
|
|
return tensor
|
|
|
|
|
|
class JITAllReduceBackend:
|
|
def __init__(self) -> None:
|
|
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
|
CustomAllReduceV2,
|
|
)
|
|
|
|
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
|
|
# tuned workspace sizes, capped at the sweep maximum
|
|
self.comm = CustomAllReduceV2(_init_cpu_group(), device, max_size=MAX_BYTES)
|
|
if self.comm.disabled:
|
|
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
|
# keep the whole sweep on the custom-AR path: the tuned config would
|
|
# otherwise send the largest sizes back to NCCL
|
|
self.comm.uncap_pull_thresholds()
|
|
register_comm_cleanup(self.comm)
|
|
|
|
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 AOTAllReduceBackend:
|
|
def __init__(self) -> None:
|
|
from sglang.srt.distributed.device_communicators.custom_all_reduce import (
|
|
CustomAllreduce,
|
|
)
|
|
|
|
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("AOT CustomAllreduce is disabled on this system")
|
|
register_comm_cleanup(self.comm)
|
|
|
|
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 FlashInferAllReduceBackend:
|
|
def __init__(self) -> None:
|
|
import flashinfer.comm as comm
|
|
|
|
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 = 1024 * min(MESSAGE_SIZES_KB) // 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,
|
|
)
|
|
|
|
def graph_context(self):
|
|
return contextlib.nullcontext()
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
@cache_once
|
|
def _init_nccl_backend() -> NCCLAllReduceBackend:
|
|
return NCCLAllReduceBackend()
|
|
|
|
|
|
@cache_once
|
|
def _init_jit_backend() -> JITAllReduceBackend:
|
|
return JITAllReduceBackend()
|
|
|
|
|
|
@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-eager": _init_jit_backend,
|
|
"jit-graph": _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.
|
|
"""
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
if local_rank == 0: # NOTE: log some verbose info on initialization
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
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()
|
|
|
|
# reset level to warning
|
|
logging.getLogger().setLevel(logging.WARNING)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Benchmark
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@marker.parametrize("message_KB", MESSAGE_SIZES_KB)
|
|
@marker.benchmark("provider", PROVIDERS)
|
|
def benchmark(message_KB: 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 {AOT_SUPPORTED_WORLD_SIZES}"
|
|
)
|
|
_init_all_backends()
|
|
backend = BACKEND_FACTORY[provider]()
|
|
message_bytes = message_KB * 1024
|
|
numel = message_bytes // DTYPE_ITEMSIZE
|
|
device_id = int(os.environ["LOCAL_RANK"])
|
|
device = torch.device(f"cuda:{device_id}")
|
|
x = torch.randn(numel, dtype=DTYPE, device=device)
|
|
ctx_fn = backend.graph_context if not provider.endswith("eager") else None
|
|
# 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=ctx_fn,
|
|
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,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
multigpu_bench_main(
|
|
name=__name__,
|
|
file=__file__,
|
|
num_gpus=WORLD_SIZES,
|
|
main_fn=benchmark.run,
|
|
)
|