[JIT Kernel][Feature] Support JIT custom all reduce (rewrite as v2) (#19880)
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
DarkSharpness
@@ -17,7 +17,7 @@ PenaltyReturnTypeOnItsOwnLine: 100 # Keeps return type with function name
|
||||
IncludeCategories:
|
||||
- Regex: '^<sgl_kernel/.*\.h>$'
|
||||
Priority: 0
|
||||
- Regex: '^<sgl_kernel/impl/.*>$'
|
||||
- Regex: '^<sgl_kernel/.*/.*>$'
|
||||
Priority: 2
|
||||
- Regex: '^<sgl_kernel/.*\.cuh>$'
|
||||
Priority: 1
|
||||
|
||||
192
python/sglang/jit_kernel/all_reduce.py
Normal file
192
python/sglang/jit_kernel/all_reduce.py
Normal file
@@ -0,0 +1,192 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import enum
|
||||
from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, cast
|
||||
|
||||
import torch
|
||||
import tvm_ffi
|
||||
|
||||
from sglang.jit_kernel.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
|
||||
|
||||
class ConfigResult(NamedTuple):
|
||||
num_blocks: int
|
||||
num_threads: int
|
||||
|
||||
|
||||
class AllReduceAlgo(enum.Enum):
|
||||
ONE_SHOT_PUSH = enum.auto()
|
||||
ONE_SHOT_PULL = enum.auto()
|
||||
TWO_SHOT_PULL = enum.auto()
|
||||
|
||||
def is_push(self) -> bool:
|
||||
return self == AllReduceAlgo.ONE_SHOT_PUSH
|
||||
|
||||
@property
|
||||
def shot(self) -> int:
|
||||
return 2 if self == AllReduceAlgo.TWO_SHOT_PULL else 1
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
CUSTOM_AR_HANDLE = List[int]
|
||||
CUSTOM_AR_PAIR = Tuple[int, CUSTOM_AR_HANDLE]
|
||||
|
||||
class CustomAllReduceObj:
|
||||
def __init__(
|
||||
self,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
pull_buffer_bytes: int,
|
||||
push_buffer_bytes: int,
|
||||
graph_input_count: int,
|
||||
*,
|
||||
max_pull_blocks: Optional[int] = None,
|
||||
max_push_blocks: Optional[int] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Create a CustomAllReduceObj instance.
|
||||
|
||||
:param rank: The rank of the current process.
|
||||
:param world_size: The total number of processes in the group.
|
||||
:param pull_buffer_bytes: The size of the buffer (in bytes) used for pull-based all-reduce.
|
||||
:param push_buffer_bytes: The size of the buffer (in bytes) used for push-based all-reduce.
|
||||
:param graph_input_count: The maximum number of inputs in all CUDA graphs.
|
||||
:param max_pull_blocks: The maximum number of thread blocks to launch for pull-based all-reduce.
|
||||
If None, it will be determined by the implementation.
|
||||
:param max_push_blocks: The maximum number of thread blocks to launch for push-based all-reduce.
|
||||
If None, it will be determined by the implementation.
|
||||
"""
|
||||
|
||||
@property
|
||||
def world_size(self) -> int: ...
|
||||
def share_storage(self) -> CUSTOM_AR_HANDLE: ...
|
||||
def share_graph_inputs(self) -> List[CUSTOM_AR_PAIR]: ...
|
||||
def post_init(self, handles: List[CUSTOM_AR_HANDLE]) -> None: ...
|
||||
def register_inputs(self, handles: List[List[CUSTOM_AR_PAIR]]) -> None: ...
|
||||
def set_cuda_graph_capture(self, is_capturing: bool) -> None: ...
|
||||
def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None: ...
|
||||
def all_reduce(
|
||||
self, input: torch.Tensor, algo: AllReduceAlgo
|
||||
) -> tvm_ffi.Tensor: ...
|
||||
def config_pull(
|
||||
self, num_blocks: int = -1, num_threads: int = -1
|
||||
) -> ConfigResult:
|
||||
"""
|
||||
Configure the CUDA kernel's grid and block dimensions.
|
||||
This provides only the upper bound of the configuration,
|
||||
and the actual launch configuration may be determined by implementation.
|
||||
Note that push-based all-reduce can not be configured currently.
|
||||
|
||||
:param num_blocks: The maximum number of thread blocks to launch. -1 means no limit.
|
||||
:param num_threads: The maximum number of threads per block. -1 means no limit.
|
||||
|
||||
:return: The previous configuration as a ConfigResult named tuple.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int):
|
||||
args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
"custom_all_reduce",
|
||||
*args,
|
||||
extra_ldflags=["-lcuda"],
|
||||
cuda_files=["distributed/custom_all_reduce_pull.cuh"],
|
||||
cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int):
|
||||
args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
"custom_all_reduce",
|
||||
*args,
|
||||
extra_ldflags=["-lcuda"],
|
||||
cuda_files=["distributed/custom_all_reduce_push.cuh"],
|
||||
cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def get_custom_all_reduce_cls() -> type[CustomAllReduceObj]:
|
||||
module = load_jit(
|
||||
"custom_all_reduce_base",
|
||||
extra_ldflags=["-lcuda"],
|
||||
cuda_files=["distributed/custom_all_reduce_base.cuh"],
|
||||
cuda_wrappers=[("register_once", "register_custom_all_reduce")],
|
||||
)
|
||||
module.register_once()
|
||||
device = torch.cuda.current_device()
|
||||
props = torch.cuda.get_device_properties(device)
|
||||
NUM_CTA = props.multi_processor_count
|
||||
MAX_THREADS = 512
|
||||
|
||||
@tvm_ffi.register_object("sgl.CustomAllReduce")
|
||||
class CustomAllReduceObjReal(tvm_ffi.Object):
|
||||
def __init__(
|
||||
self,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
pull_buffer_bytes: int,
|
||||
push_buffer_bytes: int,
|
||||
graph_input_count: int,
|
||||
*,
|
||||
max_pull_blocks: Optional[int] = None,
|
||||
max_push_blocks: Optional[int] = None,
|
||||
) -> None:
|
||||
self.__ffi_init__(
|
||||
rank,
|
||||
world_size,
|
||||
NUM_CTA if max_pull_blocks is None else max_pull_blocks,
|
||||
NUM_CTA if max_push_blocks is None else max_push_blocks,
|
||||
pull_buffer_bytes,
|
||||
push_buffer_bytes,
|
||||
graph_input_count,
|
||||
)
|
||||
self._world_size = world_size
|
||||
self._pull_config = ConfigResult(NUM_CTA, MAX_THREADS)
|
||||
self.configure_pull(*self._pull_config) # type: ignore
|
||||
|
||||
@property
|
||||
def world_size(self) -> int:
|
||||
return self._world_size
|
||||
|
||||
def all_reduce(
|
||||
self,
|
||||
input: torch.Tensor,
|
||||
algo: AllReduceAlgo,
|
||||
) -> tvm_ffi.Tensor:
|
||||
compile_fn = (
|
||||
_jit_custom_all_reduce_push_module
|
||||
if algo.is_push()
|
||||
else _jit_custom_all_reduce_pull_module
|
||||
)
|
||||
module = compile_fn(input.dtype, self._world_size)
|
||||
return module.all_reduce(self, input, algo.shot)
|
||||
|
||||
def config_pull(
|
||||
self, num_blocks: int = -1, num_threads: int = -1
|
||||
) -> ConfigResult:
|
||||
old_config = self._pull_config
|
||||
num_blocks = num_blocks if num_blocks != -1 else old_config.num_blocks
|
||||
num_threads = num_threads if num_threads != -1 else old_config.num_threads
|
||||
new_config = ConfigResult(num_blocks, num_threads)
|
||||
if new_config != old_config:
|
||||
result = ConfigResult(*self.configure_pull(*new_config)) # type: ignore
|
||||
assert result == self._pull_config
|
||||
self._pull_config = new_config
|
||||
return old_config
|
||||
|
||||
def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None:
|
||||
self.free_ipc_handles() # type: ignore
|
||||
torch.distributed.barrier(group=tp_cpu_group)
|
||||
self.free_storage() # type: ignore
|
||||
|
||||
return cast(type["CustomAllReduceObj"], CustomAllReduceObjReal)
|
||||
377
python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py
Normal file
377
python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py
Normal file
@@ -0,0 +1,377 @@
|
||||
"""
|
||||
Benchmark JIT custom all-reduce (v2) vs NCCL vs AOT custom all-reduce (v1).
|
||||
|
||||
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
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
from math import isnan
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import is_in_ci
|
||||
|
||||
DTYPE_MAP = {
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float32": torch.float32,
|
||||
}
|
||||
|
||||
MESSAGE_SIZES_BYTES = [
|
||||
4 * 1024, # 4K
|
||||
16 * 1024, # 16K
|
||||
64 * 1024, # 64K
|
||||
128 * 1024, # 128K
|
||||
3 * 64 * 1024, # 192K
|
||||
4 * 64 * 1024, # 256K
|
||||
3 * 128 * 1024, # 384K
|
||||
4 * 128 * 1024, # 512K
|
||||
5 * 128 * 1024, # 640K
|
||||
6 * 128 * 1024, # 768K
|
||||
7 * 128 * 1024, # 896K
|
||||
1 * 1024 * 1024, # 1M
|
||||
2 * 1024 * 1024, # 2M
|
||||
3 * 1024 * 1024, # 2M
|
||||
4 * 1024 * 1024, # 4M
|
||||
8 * 1024 * 1024, # 8M
|
||||
16 * 1024 * 1024, # 16M
|
||||
32 * 1024 * 1024, # 32M
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class NCCLAllReduceBackend:
|
||||
name = "NCCL"
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup):
|
||||
self.group = group
|
||||
|
||||
def capture(self, register_input: bool):
|
||||
return contextlib.nullcontext()
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
dist.all_reduce(tensor, group=self.group)
|
||||
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,
|
||||
)
|
||||
|
||||
max_size = max(MESSAGE_SIZES_BYTES)
|
||||
self.comm = CustomAllreduce(group, device, max_size=max_size)
|
||||
if self.comm.disabled:
|
||||
raise RuntimeError("AOT CustomAllreduce is disabled on this system")
|
||||
|
||||
def capture(self, register_input: bool):
|
||||
return self.comm.capture() # ignore register_input since v1 always requires it
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
max_size = max(MESSAGE_SIZES_BYTES)
|
||||
self.comm = CustomAllReduceV2(group, device, max_pull_size=max_size)
|
||||
if self.comm.disabled:
|
||||
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
||||
|
||||
def capture(self, register_input: bool):
|
||||
return self.comm.capture() if register_input else contextlib.nullcontext()
|
||||
|
||||
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:
|
||||
name = "FI"
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup, dtype: torch.dtype):
|
||||
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(
|
||||
backend="trtllm",
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
max_token_num=num_tokens,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def capture(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,
|
||||
fp32_acc=True,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmarking helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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()
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def bench_one(
|
||||
backend,
|
||||
inp: torch.Tensor,
|
||||
warmup: int,
|
||||
iters: int,
|
||||
group: dist.ProcessGroup,
|
||||
register_input: bool,
|
||||
) -> float:
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Result printing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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",
|
||||
)
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,27 @@
|
||||
#include <sgl_kernel/ffi.h>
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
#include <sgl_kernel/distributed/custom_all_reduce.cuh>
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
inline void register_custom_all_reduce() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
using Class = host::distributed::CustomAllReduceBase;
|
||||
refl::ObjectDef<Class>()
|
||||
.def(refl::init<uint32_t, uint32_t, uint32_t, uint32_t, int64_t, int64_t, int64_t>(), "__init__")
|
||||
.def("share_storage", &Class::share_storage)
|
||||
.def("share_graph_inputs", &Class::share_graph_inputs)
|
||||
.def("post_init", &Class::post_init)
|
||||
.def("register_inputs", &Class::register_inputs)
|
||||
.def("set_cuda_graph_capture", &Class::set_cuda_graph_capture)
|
||||
.def("free_ipc_handles", &Class::free_ipc_handles)
|
||||
.def("free_storage", &Class::free_storage)
|
||||
.def("configure_pull", &Class::configure_pull);
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
// Partially migrated from AOT kernel:
|
||||
// https://github.com/sgl-project/sglang/blob/v0.5.9/sgl-kernel/csrc/allreduce/custom_all_reduce.cu
|
||||
// Which was originally adapted from:
|
||||
// https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/custom_all_reduce.cu
|
||||
// We redesign the controller interface to minimize control plane traffic,
|
||||
// and fuse the reduce-scatter and broadcast in the 2-shot all reduce
|
||||
#include <sgl_kernel/ffi.h>
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
#include <sgl_kernel/distributed/common.cuh>
|
||||
#include <sgl_kernel/distributed/custom_all_reduce.cuh>
|
||||
|
||||
#include <bit>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
namespace {
|
||||
|
||||
using device::distributed::PullController;
|
||||
using host::distributed::AllReduceData;
|
||||
using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
|
||||
|
||||
struct AllReduceParams {
|
||||
void* __restrict__ output;
|
||||
uint32_t rank;
|
||||
uint32_t num_items; // NOTE: support at most 4G, but that's too much
|
||||
};
|
||||
|
||||
[[maybe_unused]]
|
||||
SGL_DEVICE void prefetch_uniform_ptr(const void* ptr) {
|
||||
asm volatile("prefetchu.L1 [%0];" ::"l"(ptr) : "memory");
|
||||
}
|
||||
|
||||
#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1)
|
||||
|
||||
template <bool kBroadcast, typename DType, uint32_t kNumGPU>
|
||||
SGL_DEVICE void all_reduce_impl(const AllReduceParams& params, DType* (&input)[kNumGPU]) {
|
||||
using namespace device;
|
||||
|
||||
constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
|
||||
using DType2 = packed_t<DType>;
|
||||
using Storage = AlignedVector<DType2, kVecSize>;
|
||||
const auto& [output, rank, num_items] = params;
|
||||
|
||||
for (auto i = blockIdx.x;; i += gridDim.x) {
|
||||
const auto offset = i * blockDim.x + threadIdx.x;
|
||||
if (offset * kVecSize * 2 >= num_items) break;
|
||||
Storage storage[kNumGPU];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i) {
|
||||
storage[i].load(input[i], offset);
|
||||
}
|
||||
const Storage result = distributed::reduce_impl(storage);
|
||||
if constexpr (kBroadcast) {
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i) {
|
||||
result.store(input[i], offset);
|
||||
}
|
||||
} else {
|
||||
result.store(output, offset);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
|
||||
CUSTOM_AR_KERNEL void all_reduce_one_shot_kernel(
|
||||
const AllReduceData* __restrict__ data,
|
||||
const AllReduceParams __grid_constant__ params,
|
||||
const PullController __grid_constant__ ctrl) {
|
||||
/// NOTE: we assume the data array is ready before the previous kernel
|
||||
DType* input[kNumGPU];
|
||||
prefetch_uniform_ptr(data);
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i)
|
||||
input[i] = static_cast<DType*>(data->input[i]);
|
||||
device::PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
|
||||
all_reduce_impl</*kBroadcast=*/false>(params, input);
|
||||
|
||||
device::PDLTriggerSecondary<kUsePDL>();
|
||||
ctrl.sync</*kFence=*/0, /*kStart=*/0>(params.rank, kNumGPU);
|
||||
}
|
||||
|
||||
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
|
||||
CUSTOM_AR_KERNEL void all_reduce_two_shot_kernel(
|
||||
const AllReduceData* __restrict__ data,
|
||||
const AllReduceParams __grid_constant__ params,
|
||||
const PullController __grid_constant__ ctrl) {
|
||||
// get the range of this rank
|
||||
using device::kWarpThreads, device::div_ceil;
|
||||
|
||||
prefetch_uniform_ptr(data);
|
||||
DType* input[kNumGPU];
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i)
|
||||
input[i] = static_cast<DType*>(data->input[i]);
|
||||
|
||||
constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
|
||||
const uint32_t num_items = params.num_items;
|
||||
const uint32_t total_vec = num_items / (kVecSize * 2); // must be divisible here
|
||||
const uint32_t vec_per_rank = div_ceil(div_ceil(total_vec, kNumGPU), kWarpThreads) * kWarpThreads;
|
||||
const uint32_t local_vec_start = min(params.rank * vec_per_rank, total_vec);
|
||||
const uint32_t local_vec_finish = min(local_vec_start + vec_per_rank, total_vec);
|
||||
const uint32_t local_start = local_vec_start * kVecSize * 2;
|
||||
const uint32_t local_length = (local_vec_finish - local_vec_start) * kVecSize * 2;
|
||||
const auto local_params = AllReduceParams{
|
||||
.output = nullptr, // this is not used for 2-shot all reduce
|
||||
.rank = params.rank,
|
||||
.num_items = local_length,
|
||||
};
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i)
|
||||
input[i] += local_start;
|
||||
|
||||
device::PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
|
||||
all_reduce_impl</*kBroadcast=*/true>(local_params, input);
|
||||
|
||||
device::PDLTriggerSecondary<kUsePDL>();
|
||||
ctrl.sync</*kFence=*/1, /*kStart=*/0>(params.rank, kNumGPU);
|
||||
}
|
||||
|
||||
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
|
||||
struct CustomAllReducePull : public CustomAllReduceBase {
|
||||
static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
|
||||
static constexpr auto one_shot_kernel = all_reduce_one_shot_kernel<DType, kNumGPU, kUsePDL>;
|
||||
static constexpr auto two_shot_kernel = all_reduce_two_shot_kernel<DType, kNumGPU, kUsePDL>;
|
||||
static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
|
||||
|
||||
tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) {
|
||||
using namespace host;
|
||||
const bool use_2shot = (shot == 2);
|
||||
const auto device = input.device();
|
||||
const auto input_ptr = input.data_ptr();
|
||||
const auto buffer_ptr = get_pull_buffer(m_storage);
|
||||
const auto num_items_int64 = input.numel();
|
||||
const auto num_items = static_cast<uint32_t>(num_items_int64);
|
||||
const auto items_per_block = m_cta_size * kVecSize * 2;
|
||||
const auto needed_blocks = div_ceil(num_items, items_per_block);
|
||||
const auto num_blocks = std::min(needed_blocks, m_num_cta);
|
||||
const auto kernel = use_2shot ? two_shot_kernel : one_shot_kernel;
|
||||
// only 1-shot + graph capture need extra output buffer
|
||||
const auto output = (m_is_graph_capturing && !use_2shot) ? ffi::empty_like(input) : input;
|
||||
const auto params = AllReduceParams{
|
||||
.output = use_2shot ? nullptr : output.data_ptr(),
|
||||
.rank = m_rank,
|
||||
.num_items = num_items,
|
||||
};
|
||||
|
||||
RuntimeCheck(input.IsContiguous(), "Input tensor must be contiguous");
|
||||
RuntimeCheck(m_num_gpu == kNumGPU, "Mismatch GPU count");
|
||||
RuntimeCheck(shot == 1 || shot == 2, "Invalid shot count: ", shot);
|
||||
RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported");
|
||||
RuntimeCheck(is_type<DType>(input.dtype()), "Input dtype mismatch");
|
||||
RuntimeCheck(std::bit_cast<intptr_t>(input_ptr) % 16 == 0, "Input pointer is not properly aligned");
|
||||
RuntimeCheck(m_pull_ctrl.has_value(), "Controller is not initialized");
|
||||
RuntimeCheck(static_cast<int64_t>(num_items) == num_items_int64, "Number of items exceeds 4G limit");
|
||||
|
||||
const auto& ctrl = *m_pull_ctrl;
|
||||
const auto stream = LaunchKernel::resolve_device(device);
|
||||
auto launch = LaunchKernel{num_blocks, m_cta_size, stream};
|
||||
launch.enable_pdl(kUsePDL);
|
||||
const auto check_capturing = [&] {
|
||||
if (!m_is_graph_capturing) return false; // override to avoid cudaRT call overhead
|
||||
cudaStreamCaptureStatus status;
|
||||
RuntimeDeviceCheck(cudaStreamIsCapturing(stream, &status));
|
||||
return status == cudaStreamCaptureStatusActive;
|
||||
};
|
||||
if (check_capturing()) {
|
||||
// no-op if not really capturing, we're in a dummy run
|
||||
const auto data_ptr = allocate_graph_capture_input(input_ptr);
|
||||
/// NOTE: we assume when the graph is replayed, the data_ptr should be ready
|
||||
launch(kernel, data_ptr, params, ctrl);
|
||||
} else {
|
||||
// 1.copy the input to the buffer
|
||||
const auto input_bytes = static_cast<int64_t>(sizeof(DType) * num_items);
|
||||
RuntimeCheck(input_bytes <= m_pull_buffer_bytes, "Input is too large, num items: ", num_items);
|
||||
RuntimeDeviceCheck(cudaMemcpyAsync(buffer_ptr, input_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
// 2. launch the all reduce kernel
|
||||
const auto data_ptr = get_data_ptr(); // use default buffer
|
||||
launch(kernel, data_ptr, params, ctrl);
|
||||
if (use_2shot) { // 3. copy the reduced result back to the output, because 2-shot doesn't write to output
|
||||
RuntimeDeviceCheck(cudaMemcpyAsync(input_ptr, buffer_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
}
|
||||
}
|
||||
return output;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
|
||||
tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
|
||||
using Impl = CustomAllReducePull<DType, kNumGPU, kUsePDL>;
|
||||
return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
@@ -0,0 +1,253 @@
|
||||
// Partially adapted from:
|
||||
// https://github.com/flashinfer-ai/flashinfer/blob/v0.6.4/include/flashinfer/comm/trtllm_allreduce_fusion.cuh
|
||||
// We simplify the lamport design and minimize the ring buffer count (from 3 -> 2)
|
||||
#include <sgl_kernel/ffi.h>
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
#include <sgl_kernel/distributed/common.cuh>
|
||||
#include <sgl_kernel/distributed/custom_all_reduce.cuh>
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
namespace {
|
||||
|
||||
using device::distributed::PushController;
|
||||
using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
|
||||
|
||||
struct AllReducePushData {
|
||||
void* __restrict__ buffer[device::distributed::kMaxNumGPU];
|
||||
const void* input;
|
||||
void* output;
|
||||
uint32_t rank;
|
||||
uint32_t num_items;
|
||||
uint32_t buffer_bytes;
|
||||
uint32_t epoch_bytes;
|
||||
};
|
||||
|
||||
#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1)
|
||||
|
||||
template <typename T>
|
||||
struct fp_trait {};
|
||||
|
||||
// TODO: support more dtypes
|
||||
template <>
|
||||
struct fp_trait<bf16_t> {
|
||||
using type = uint16_t;
|
||||
[[maybe_unused]]
|
||||
static constexpr uint16_t pos_zero = 0x0000u;
|
||||
[[maybe_unused]]
|
||||
static constexpr uint16_t neg_zero = 0x8000u;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct fp_trait<fp16_t> {
|
||||
using type = uint16_t;
|
||||
[[maybe_unused]]
|
||||
static constexpr uint16_t pos_zero = 0x0000u;
|
||||
[[maybe_unused]]
|
||||
static constexpr uint16_t neg_zero = 0x8000u;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct fp_trait<float> {
|
||||
using type = uint32_t;
|
||||
[[maybe_unused]]
|
||||
static constexpr uint32_t pos_zero = 0x00000000u;
|
||||
[[maybe_unused]]
|
||||
static constexpr uint32_t neg_zero = 0x80000000u;
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
SGL_DEVICE void clear_pos_zero(DType& val) {
|
||||
using Trait = fp_trait<DType>;
|
||||
const auto ptr = reinterpret_cast<typename Trait::type*>(&val);
|
||||
if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero;
|
||||
}
|
||||
|
||||
template <typename DType>
|
||||
SGL_DEVICE bool is_pos_zero(const DType& val) {
|
||||
using Trait = fp_trait<DType>;
|
||||
const auto ptr = reinterpret_cast<const typename Trait::type*>(&val);
|
||||
return *ptr == Trait::pos_zero;
|
||||
}
|
||||
|
||||
template <typename DType>
|
||||
SGL_DEVICE DType get_pos_zero() {
|
||||
using Trait = fp_trait<DType>;
|
||||
const auto value = Trait::pos_zero;
|
||||
return *reinterpret_cast<const DType*>(&value);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SGL_DEVICE void ld_global_volatile_16B(T& x, const void* addr, int64_t offset) {
|
||||
static_assert(alignof(T) == 16 && sizeof(T) == 16);
|
||||
addr = device::pointer::offset<T>(addr, offset);
|
||||
uint4 val;
|
||||
asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];"
|
||||
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
|
||||
: "l"(addr));
|
||||
x = *reinterpret_cast<const T*>(&val);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SGL_DEVICE void st_global_volatile_16B(const T& x, void* addr, int64_t offset) {
|
||||
static_assert(alignof(T) == 16 && sizeof(T) == 16);
|
||||
const uint4 val = *reinterpret_cast<const uint4*>(&x);
|
||||
addr = device::pointer::offset<T>(addr, offset);
|
||||
asm volatile(
|
||||
"st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr));
|
||||
}
|
||||
|
||||
template <typename DType, uint32_t kNumGPU>
|
||||
SGL_DEVICE void push_impl(DType* (&push_buf)[kNumGPU], const void* data, uint32_t num_items) {
|
||||
using namespace device;
|
||||
constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
|
||||
using Storage = AlignedVector<packed_t<DType>, kVecSize>;
|
||||
|
||||
for (auto i = blockIdx.x;; i += gridDim.x) {
|
||||
const auto offset = i * blockDim.x + threadIdx.x;
|
||||
if (offset * kVecSize * 2 >= num_items) break;
|
||||
Storage vec;
|
||||
vec.load(data, offset);
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kVecSize; ++j) {
|
||||
clear_pos_zero(vec[j].x);
|
||||
clear_pos_zero(vec[j].y);
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i) {
|
||||
st_global_volatile_16B(vec, push_buf[i], offset);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, uint32_t kNumGPU>
|
||||
SGL_DEVICE void poll_impl(DType* (&poll_buf)[kNumGPU], void* data, uint32_t num_items) {
|
||||
using namespace device;
|
||||
constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
|
||||
using Storage = AlignedVector<packed_t<DType>, kVecSize>;
|
||||
|
||||
for (auto i = blockIdx.x;; i += gridDim.x) {
|
||||
const auto offset = i * blockDim.x + threadIdx.x;
|
||||
if (offset * kVecSize * 2 >= num_items) break;
|
||||
Storage storage[kNumGPU];
|
||||
|
||||
while (true) {
|
||||
bool has_pos_zero = false;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i) {
|
||||
ld_global_volatile_16B(storage[i], poll_buf[i], offset);
|
||||
#pragma unroll
|
||||
for (auto j = 0; j < kVecSize; ++j) {
|
||||
has_pos_zero |= is_pos_zero(storage[i][j].x);
|
||||
has_pos_zero |= is_pos_zero(storage[i][j].y);
|
||||
}
|
||||
}
|
||||
if (!has_pos_zero) break;
|
||||
}
|
||||
|
||||
const Storage result = distributed::reduce_impl(storage);
|
||||
result.store(data, offset);
|
||||
|
||||
Storage pos_zeros;
|
||||
pos_zeros.fill({get_pos_zero<DType>(), get_pos_zero<DType>()});
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i) {
|
||||
pos_zeros.store(poll_buf[i], offset);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
|
||||
CUSTOM_AR_KERNEL void all_reduce_one_shot_push_kernel(
|
||||
const AllReducePushData __grid_constant__ params, //
|
||||
const PushController __grid_constant__ ctrl) {
|
||||
using namespace device;
|
||||
|
||||
const auto [buffer, input, output, rank, num_items, buffer_bytes, epoch_bytes] = params;
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
// Phase 1: Push data from input to all ranks' buffers
|
||||
const auto epoch_offset = ctrl.epoch() * epoch_bytes;
|
||||
DType* push_buf[kNumGPU];
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i) {
|
||||
push_buf[i] = static_cast<DType*>(pointer::offset(buffer[i], rank * buffer_bytes, epoch_offset));
|
||||
}
|
||||
push_impl(push_buf, input, num_items);
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
// Phase 2: Poll local data
|
||||
DType* poll_buf[kNumGPU];
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i) {
|
||||
poll_buf[i] = static_cast<DType*>(pointer::offset(buffer[rank], i * buffer_bytes, epoch_offset));
|
||||
}
|
||||
poll_impl(poll_buf, output, num_items);
|
||||
ctrl.exit();
|
||||
}
|
||||
|
||||
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
|
||||
struct CustomAllReducePush : public CustomAllReduceBase {
|
||||
static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
|
||||
static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
|
||||
|
||||
tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) {
|
||||
using namespace host;
|
||||
const auto device = input.device();
|
||||
const auto input_ptr = input.data_ptr();
|
||||
const auto num_items_int64 = input.numel();
|
||||
const auto num_items = static_cast<uint32_t>(num_items_int64);
|
||||
const auto num_blocks = m_max_num_cta_push; // must be constant to ensure correctness
|
||||
const auto num_threads = [&] {
|
||||
for (const auto t : {128u, 256u, 512u}) {
|
||||
if (t * num_blocks * 2 * kVecSize >= num_items) return t;
|
||||
}
|
||||
return 1024u;
|
||||
}();
|
||||
const auto output = input;
|
||||
AllReducePushData params;
|
||||
for (uint32_t i = 0; i < kNumGPU; ++i) {
|
||||
params.buffer[i] = get_push_buffer(m_peer_storage[i]);
|
||||
}
|
||||
params.input = input_ptr;
|
||||
params.output = input_ptr;
|
||||
params.rank = m_rank;
|
||||
params.num_items = num_items;
|
||||
params.buffer_bytes = m_push_buffer_bytes;
|
||||
params.epoch_bytes = kNumGPU * params.buffer_bytes;
|
||||
|
||||
RuntimeCheck(input.IsContiguous(), "Input must be contiguous");
|
||||
RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch");
|
||||
RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported");
|
||||
RuntimeCheck(is_type<DType>(input.dtype()), "Input dtype mismatch");
|
||||
RuntimeCheck(std::bit_cast<intptr_t>(input_ptr) % 16 == 0, "Input pointer is not properly aligned");
|
||||
RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized");
|
||||
RuntimeCheck(shot == 1, "Push all-reduce only supports 1-shot, got: ", shot);
|
||||
RuntimeCheck(static_cast<int64_t>(num_items) == num_items_int64, "Number of items exceeds 4G limit");
|
||||
|
||||
const auto input_bytes = static_cast<int64_t>(sizeof(DType) * num_items_int64);
|
||||
RuntimeCheck(input_bytes <= m_push_buffer_bytes, "Input is too large, num items: ", num_items);
|
||||
|
||||
const auto kernel = all_reduce_one_shot_push_kernel<DType, kNumGPU, kUsePDL>;
|
||||
LaunchKernel(num_blocks, num_threads, device) //
|
||||
.enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl);
|
||||
return output;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType, uint32_t kNumGPU, bool kUsePDL>
|
||||
tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
|
||||
using Impl = CustomAllReducePush<DType, kNumGPU, kUsePDL>;
|
||||
return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
@@ -40,8 +40,6 @@ struct Vec {
|
||||
|
||||
using I4 = Vec<int, 4>;
|
||||
|
||||
using host::div_ceil;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
|
||||
// No support for async
|
||||
#else
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
#pragma once
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
namespace device::distributed {
|
||||
|
||||
inline constexpr uint32_t kMaxNumGPU = 8;
|
||||
|
||||
struct alignas(128) Semaphore {
|
||||
public:
|
||||
constexpr Semaphore() : m_flag(0), m_counter(0) {}
|
||||
|
||||
template <bool kFence>
|
||||
SGL_DEVICE uint32_t get() const {
|
||||
uint32_t val;
|
||||
if constexpr (kFence) {
|
||||
asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag));
|
||||
} else {
|
||||
asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag));
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
template <bool kFence>
|
||||
SGL_DEVICE uint32_t add(uint32_t val) {
|
||||
uint32_t old_val;
|
||||
if constexpr (kFence) {
|
||||
asm volatile("atom.release.sys.global.add.u32 %0, [%1], %2;" : "=r"(old_val) : "l"(&m_flag), "r"(val));
|
||||
} else {
|
||||
asm volatile("atom.global.add.u32 %0, [%1], %2;" : "=r"(old_val) : "l"(&m_flag), "r"(val));
|
||||
}
|
||||
return old_val;
|
||||
}
|
||||
|
||||
// Only called by the owning GPU - plain load is sufficient
|
||||
SGL_DEVICE uint32_t get_counter() const {
|
||||
return m_counter;
|
||||
}
|
||||
|
||||
// Only called by the owning GPU - plain store is sufficient
|
||||
SGL_DEVICE void set_counter(uint32_t val) {
|
||||
m_counter = val;
|
||||
}
|
||||
|
||||
private:
|
||||
uint32_t m_flag;
|
||||
uint32_t m_counter;
|
||||
};
|
||||
|
||||
struct PullController {
|
||||
public:
|
||||
PullController(void** signals, uint32_t num_gpu) {
|
||||
for (uint32_t i = 0; i < num_gpu; ++i) {
|
||||
m_signals[i] = static_cast<Semaphore*>(signals[i]);
|
||||
}
|
||||
}
|
||||
|
||||
/// Synchronize all GPUs.
|
||||
/// When kFence is true, establishes happens-before across GPUs using
|
||||
/// release/acquire semantics, ensuring prior writes are visible system-wide.
|
||||
template <bool kFence, bool kStart>
|
||||
SGL_DEVICE void sync(uint32_t rank, uint32_t num_gpu) const {
|
||||
// For fenced sync: ensure all threads in this block have completed their writes,
|
||||
// so the signaling thread's release carries them transitively.
|
||||
static_assert(!(kFence && kStart), "Start stage does not need to wait fence");
|
||||
if constexpr (kFence || !kStart) __syncthreads();
|
||||
constexpr auto kStage = kStart ? 1 : 2;
|
||||
const auto warp_id = threadIdx.x / kWarpThreads;
|
||||
const auto lane_id = threadIdx.x % kWarpThreads;
|
||||
if (lane_id == 0 && warp_id < num_gpu) {
|
||||
auto& signal = m_signals[warp_id][blockIdx.x];
|
||||
signal.add<kFence>(1);
|
||||
if (warp_id == rank) {
|
||||
const auto target = num_gpu * kStage;
|
||||
/// NOTE: correctness here:
|
||||
/// - base is only read/updated locally by the owning GPU
|
||||
const auto base = signal.get_counter();
|
||||
while (signal.get<kFence>() - base < target)
|
||||
;
|
||||
if constexpr (!kStart) {
|
||||
signal.set_counter(base + target);
|
||||
}
|
||||
}
|
||||
}
|
||||
if constexpr (kStart) __syncthreads();
|
||||
}
|
||||
|
||||
private:
|
||||
Semaphore* __restrict__ m_signals[kMaxNumGPU];
|
||||
};
|
||||
|
||||
struct PushController {
|
||||
public:
|
||||
static constexpr int64_t kNumStages = 2;
|
||||
|
||||
PushController(void* ptr) : m_local_signal(static_cast<Semaphore*>(ptr)) {}
|
||||
|
||||
SGL_DEVICE uint32_t epoch() const {
|
||||
return m_local_signal[blockIdx.x].get_counter();
|
||||
}
|
||||
|
||||
SGL_DEVICE void exit() const {
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0) {
|
||||
auto& signal = m_local_signal[blockIdx.x];
|
||||
const auto epoch = signal.get_counter();
|
||||
signal.set_counter((epoch + 1) % kNumStages);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
Semaphore* m_local_signal;
|
||||
};
|
||||
|
||||
} // namespace device::distributed
|
||||
@@ -0,0 +1,349 @@
|
||||
#pragma once
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
#include <sgl_kernel/distributed/common.cuh>
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/container/tuple.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
|
||||
#include <array>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <functional>
|
||||
#include <numeric>
|
||||
#include <optional>
|
||||
#include <span>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
namespace host::distributed {
|
||||
|
||||
using device::distributed::PullController, device::distributed::PushController;
|
||||
|
||||
struct AllReduceData {
|
||||
constexpr AllReduceData() {}
|
||||
void* __restrict__ input[device::distributed::kMaxNumGPU];
|
||||
};
|
||||
|
||||
using ExternHandle = tvm::ffi::Array<char>;
|
||||
|
||||
inline ExternHandle to_extern_handle(void* ptr) {
|
||||
ExternHandle array;
|
||||
cudaIpcMemHandle_t handle;
|
||||
RuntimeDeviceCheck(cudaIpcGetMemHandle(&handle, ptr));
|
||||
for (size_t i = 0; i < sizeof(handle); ++i) {
|
||||
array.push_back(handle.reserved[i]);
|
||||
}
|
||||
return array;
|
||||
}
|
||||
|
||||
inline void* from_extern_handle(const ExternHandle& array) {
|
||||
cudaIpcMemHandle_t handle;
|
||||
RuntimeCheck(array.size() == sizeof(handle), "Invalid IPC handle size: ", array.size());
|
||||
for (size_t i = 0; i < sizeof(handle); ++i) {
|
||||
handle.reserved[i] = array[i];
|
||||
}
|
||||
void* ptr;
|
||||
RuntimeDeviceCheck(cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess));
|
||||
return ptr;
|
||||
}
|
||||
|
||||
struct HandleHash {
|
||||
std::size_t operator()(const cudaIpcMemHandle_t& handle) const {
|
||||
return std::hash<std::string_view>{}({handle.reserved, sizeof(handle.reserved)});
|
||||
}
|
||||
};
|
||||
|
||||
struct HandleEqual {
|
||||
bool operator()(const cudaIpcMemHandle_t& a, const cudaIpcMemHandle_t& b) const {
|
||||
return std::memcmp(a.reserved, b.reserved, sizeof(a.reserved)) == 0;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* \brief The control plane of the custom all-reduce implementation.
|
||||
* It manages the internal state and synchronization of the participating GPUs.
|
||||
*/
|
||||
struct CustomAllReduceBase : public tvm::ffi::Object {
|
||||
public:
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.CustomAllReduce", CustomAllReduceBase, tvm::ffi::Object);
|
||||
|
||||
static constexpr bool _type_mutable = true;
|
||||
using InputPair = tvm::ffi::Tuple<int64_t, ExternHandle>; // (offset, ipc handle)
|
||||
|
||||
CustomAllReduceBase(
|
||||
uint32_t rank,
|
||||
uint32_t num_gpu,
|
||||
uint32_t max_num_cta_pull,
|
||||
uint32_t max_num_cta_push,
|
||||
int64_t pull_buffer_size,
|
||||
int64_t push_buffer_size,
|
||||
int64_t graph_buffer_count)
|
||||
: m_pull_buffer_bytes(pull_buffer_size),
|
||||
m_push_buffer_bytes(push_buffer_size),
|
||||
m_graph_buffer_count(graph_buffer_count),
|
||||
m_rank(rank),
|
||||
m_num_gpu(num_gpu),
|
||||
m_max_num_cta_pull(max_num_cta_pull),
|
||||
m_max_num_cta_push(max_num_cta_push),
|
||||
// default config for pull kernel, can be updated by `configure()`
|
||||
m_num_cta(max_num_cta_pull),
|
||||
m_cta_size(256) {
|
||||
RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes()));
|
||||
RuntimeCheck(rank < num_gpu, "Invalid rank: ", rank);
|
||||
const int64_t kU32Max = static_cast<int64_t>(std::numeric_limits<uint32_t>::max());
|
||||
const int64_t push_buffer_size_all = push_all_ranks_bytes();
|
||||
RuntimeCheck(pull_buffer_size <= kU32Max, "Buffer size is too large: ", pull_buffer_size);
|
||||
RuntimeCheck(push_buffer_size_all <= kU32Max, "Push buffer size is too large: ", push_buffer_size_all);
|
||||
}
|
||||
|
||||
ExternHandle share_storage() {
|
||||
return to_extern_handle(m_storage);
|
||||
}
|
||||
|
||||
tvm::ffi::Array<InputPair> share_graph_inputs() {
|
||||
tvm::ffi::Array<InputPair> result;
|
||||
const auto new_inputs_count = registered_count() - m_cum_registered_count;
|
||||
RuntimeCheck(new_inputs_count >= 0, "Invalid new count: ", new_inputs_count);
|
||||
result.reserve(new_inputs_count);
|
||||
std::unordered_map<void*, ExternHandle> ipc_cache;
|
||||
const auto get_handle = [&](void* ptr) -> ExternHandle {
|
||||
const auto it = ipc_cache.find(ptr);
|
||||
if (it != ipc_cache.end()) return it->second;
|
||||
const auto handle = to_extern_handle(ptr);
|
||||
ipc_cache.try_emplace(ptr, handle);
|
||||
return handle;
|
||||
};
|
||||
for (const auto ptr : std::span(m_graph_capture_inputs).subspan(m_cum_registered_count)) {
|
||||
// note: must share the base address of each allocation, or we get wrong address
|
||||
void* base_ptr;
|
||||
const auto cu_result = cuPointerGetAttribute(&base_ptr, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, (CUdeviceptr)ptr);
|
||||
RuntimeCheck(cu_result == CUDA_SUCCESS, "failed to get pointer attr");
|
||||
const auto offset = reinterpret_cast<char*>(ptr) - reinterpret_cast<char*>(base_ptr);
|
||||
result.push_back(InputPair{offset, get_handle(base_ptr)});
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
void post_init(tvm::ffi::Array<ExternHandle> ipc_storages) {
|
||||
RuntimeCheck(ipc_storages.size() == m_num_gpu, "Invalid array size: ", ipc_storages.size());
|
||||
m_peer_storage.resize(m_num_gpu);
|
||||
for (const auto i : irange(m_num_gpu)) {
|
||||
if (i == m_rank) {
|
||||
m_peer_storage[i] = m_storage;
|
||||
} else {
|
||||
m_peer_storage[i] = from_extern_handle(ipc_storages[i]);
|
||||
}
|
||||
}
|
||||
|
||||
// set signal buffer to zero
|
||||
const auto pull_signal = get_pull_signal(m_storage);
|
||||
RuntimeDeviceCheck(cudaMemset(pull_signal, 0, pull_signal_bytes()));
|
||||
|
||||
// update the pull controller and data pointer
|
||||
RuntimeCheck(!m_pull_ctrl.has_value(), "Controller is already initialized");
|
||||
m_pull_ctrl.emplace(m_peer_storage.data(), m_num_gpu);
|
||||
AllReduceData data;
|
||||
for (const auto i : irange(m_num_gpu)) {
|
||||
data.input[i] = get_pull_buffer(m_peer_storage[i]);
|
||||
}
|
||||
const auto default_data_ptr = get_data_ptr();
|
||||
RuntimeDeviceCheck(cudaMemcpy(default_data_ptr, &data, sizeof(AllReduceData), cudaMemcpyHostToDevice));
|
||||
|
||||
// update the push controller and data pointer
|
||||
RuntimeCheck(!m_push_ctrl.has_value(), "Controller is already initialized");
|
||||
const auto push_signal = get_push_signal(m_storage);
|
||||
RuntimeDeviceCheck(cudaMemset(push_signal, 0, push_signal_bytes()));
|
||||
m_push_ctrl.emplace(push_signal);
|
||||
const auto push_buffer = get_push_buffer(m_storage);
|
||||
RuntimeDeviceCheck(cudaMemset(push_buffer, 0, push_all_ranks_bytes()));
|
||||
}
|
||||
|
||||
void register_inputs(tvm::ffi::Array<tvm::ffi::Array<InputPair>> ipc_graph_inputs) {
|
||||
RuntimeCheck(ipc_graph_inputs.size() == m_num_gpu);
|
||||
const auto new_registered_count = registered_count() - m_cum_registered_count;
|
||||
RuntimeCheck(new_registered_count >= 0, "Invalid registered count: ", new_registered_count);
|
||||
if (new_registered_count == 0) return; // avoid `m_get_data_ptr()` out-of-bounds
|
||||
std::vector<AllReduceData> data;
|
||||
data.resize(new_registered_count);
|
||||
const auto open_cached = [&](const ExternHandle& h) -> void* {
|
||||
RuntimeCheck(h.size() == sizeof(cudaIpcMemHandle_t), "Invalid IPC handle size: ", h.size());
|
||||
cudaIpcMemHandle_t handle;
|
||||
for (size_t i = 0; i < sizeof(handle); ++i)
|
||||
handle.reserved[i] = h[i];
|
||||
const auto [it, success] = m_ipc_cache.try_emplace(handle, nullptr);
|
||||
if (success) {
|
||||
void* ptr;
|
||||
RuntimeDeviceCheck(cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess));
|
||||
it->second = ptr;
|
||||
}
|
||||
return it->second;
|
||||
};
|
||||
for (const auto i : irange(ipc_graph_inputs.size())) {
|
||||
const auto& array = ipc_graph_inputs[i];
|
||||
RuntimeCheck(int64_t(array.size()) == new_registered_count);
|
||||
if (i == m_rank) {
|
||||
for (const auto j : irange(new_registered_count)) {
|
||||
data[j].input[i] = m_graph_capture_inputs[m_cum_registered_count + j];
|
||||
}
|
||||
} else {
|
||||
for (const auto j : irange(new_registered_count)) {
|
||||
/// NOTE: structural binding will cause intern compiler error...
|
||||
const auto elem = array[j];
|
||||
const auto offset = get<0>(elem);
|
||||
const auto ipc_handle = get<1>(elem);
|
||||
data[j].input[i] = pointer::offset(open_cached(ipc_handle), offset);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const auto new_registered_bytes = sizeof(AllReduceData) * new_registered_count;
|
||||
const auto dst_ptr = get_data_ptr(m_cum_registered_count);
|
||||
m_cum_registered_count += new_registered_count;
|
||||
RuntimeDeviceCheck(cudaMemcpy(dst_ptr, data.data(), new_registered_bytes, cudaMemcpyHostToDevice));
|
||||
}
|
||||
|
||||
void set_cuda_graph_capture(bool enabled) {
|
||||
m_is_graph_capturing = enabled;
|
||||
}
|
||||
|
||||
void free_ipc_handles() {
|
||||
for (const auto& pair : m_ipc_cache) {
|
||||
host::RuntimeDeviceCheck(cudaIpcCloseMemHandle(pair.second));
|
||||
}
|
||||
m_ipc_cache.clear();
|
||||
}
|
||||
|
||||
void free_storage() {
|
||||
host::RuntimeDeviceCheck(cudaFree(m_storage));
|
||||
m_storage = nullptr;
|
||||
}
|
||||
|
||||
tvm::ffi::Tuple<uint32_t, uint32_t> configure_pull(uint32_t num_cta, uint32_t cta_size) {
|
||||
using host::RuntimeCheck;
|
||||
const auto min_cta_size = m_num_gpu * device::kWarpThreads;
|
||||
RuntimeCheck(num_cta > 0 && num_cta <= m_max_num_cta_pull, "Invalid number of CTAs: ", num_cta);
|
||||
RuntimeCheck(cta_size >= min_cta_size, "Block size must be at least ", min_cta_size);
|
||||
const auto old_num_cta = m_num_cta;
|
||||
const auto old_block_size = m_cta_size;
|
||||
m_num_cta = num_cta;
|
||||
m_cta_size = cta_size;
|
||||
return tvm::ffi::Tuple<uint32_t, uint32_t>{old_num_cta, old_block_size};
|
||||
}
|
||||
|
||||
protected:
|
||||
AllReduceData* allocate_graph_capture_input(void* data_ptr) {
|
||||
const auto count = registered_count();
|
||||
RuntimeCheck(count < m_graph_buffer_count, "Graph buffer overflow, increase `graph_buffer_count`!");
|
||||
m_graph_capture_inputs.push_back(data_ptr);
|
||||
return get_data_ptr(count);
|
||||
}
|
||||
AllReduceData* get_data_ptr(int64_t which = -1) {
|
||||
const auto count = registered_count();
|
||||
RuntimeCheck(which >= -1 && which < count, "Invalid graph buffer index: ", which, ", count: ", count);
|
||||
const auto start = get_pull_params(m_storage);
|
||||
return static_cast<AllReduceData*>(start) + (1 + which);
|
||||
}
|
||||
int64_t registered_count() const {
|
||||
return static_cast<int64_t>(m_graph_capture_inputs.size());
|
||||
}
|
||||
int64_t pull_signal_bytes() const {
|
||||
return sizeof(device::distributed::Semaphore) * m_max_num_cta_pull;
|
||||
}
|
||||
int64_t push_signal_bytes() const {
|
||||
return sizeof(device::distributed::Semaphore) * m_max_num_cta_push;
|
||||
}
|
||||
int64_t params_bytes() const {
|
||||
return sizeof(AllReduceData) * (1 + m_graph_buffer_count); // 1 for default
|
||||
}
|
||||
int64_t push_all_ranks_bytes() const {
|
||||
return PushController::kNumStages * m_num_gpu * m_push_buffer_bytes;
|
||||
}
|
||||
int64_t storage_bytes() const {
|
||||
// | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) |
|
||||
return _get_offset_impl(5);
|
||||
}
|
||||
void* get_pull_signal(void* ptr) const {
|
||||
return pointer::offset(ptr, _get_offset_impl(0));
|
||||
}
|
||||
void* get_push_signal(void* ptr) const {
|
||||
return pointer::offset(ptr, _get_offset_impl(1));
|
||||
}
|
||||
void* get_pull_params(void* ptr) const {
|
||||
return pointer::offset(ptr, _get_offset_impl(2));
|
||||
}
|
||||
void* get_pull_buffer(void* ptr) const {
|
||||
return pointer::offset(ptr, _get_offset_impl(3));
|
||||
}
|
||||
void* get_push_buffer(void* ptr) const {
|
||||
return pointer::offset(ptr, _get_offset_impl(4));
|
||||
}
|
||||
int64_t _get_offset_impl(int64_t which) const {
|
||||
const int64_t offset_map[5] = {
|
||||
/*[0]=*/pull_signal_bytes(),
|
||||
/*[1]=*/push_signal_bytes(),
|
||||
/*[2]=*/params_bytes(),
|
||||
/*[3]=*/m_pull_buffer_bytes,
|
||||
/*[4]=*/push_all_ranks_bytes(),
|
||||
};
|
||||
RuntimeCheck(which >= 0 && which <= 5, "Invalid offset index: ", which);
|
||||
return std::accumulate(offset_map, offset_map + which, int64_t(0));
|
||||
}
|
||||
|
||||
const int64_t m_pull_buffer_bytes;
|
||||
const int64_t m_push_buffer_bytes;
|
||||
const int64_t m_graph_buffer_count;
|
||||
const uint32_t m_rank;
|
||||
const uint32_t m_num_gpu;
|
||||
const uint32_t m_max_num_cta_pull;
|
||||
const uint32_t m_max_num_cta_push;
|
||||
// these 2 config should only affect pull kernel
|
||||
uint32_t m_num_cta;
|
||||
uint32_t m_cta_size;
|
||||
// other states
|
||||
bool m_is_graph_capturing = false;
|
||||
int64_t m_cum_registered_count = 0;
|
||||
std::optional<PullController> m_pull_ctrl;
|
||||
std::optional<PushController> m_push_ctrl;
|
||||
void* m_storage = nullptr;
|
||||
std::vector<void*> m_graph_capture_inputs;
|
||||
std::vector<void*> m_peer_storage;
|
||||
std::unordered_map<cudaIpcMemHandle_t, void*, HandleHash, HandleEqual> m_ipc_cache;
|
||||
};
|
||||
|
||||
struct CustomAllReduceRef : public tvm::ffi::ObjectRef {
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CustomAllReduceRef, tvm::ffi::ObjectRef, CustomAllReduceBase);
|
||||
};
|
||||
|
||||
} // namespace host::distributed
|
||||
|
||||
namespace device::distributed {
|
||||
|
||||
template <typename DType2, size_t N, uint32_t M>
|
||||
SGL_DEVICE auto reduce_impl(AlignedVector<DType2, N> (&storage)[M]) -> AlignedVector<DType2, N> {
|
||||
fp32x2_t acc[N] = {};
|
||||
#pragma unroll // unroll num gpu
|
||||
for (uint32_t i = 0; i < M; ++i) {
|
||||
#pragma unroll // unroll vec
|
||||
for (uint32_t j = 0; j < N; ++j) {
|
||||
const auto [x, y] = cast<fp32x2_t>(storage[i][j]);
|
||||
auto& [x_acc, y_acc] = acc[j];
|
||||
x_acc += x;
|
||||
y_acc += y;
|
||||
}
|
||||
}
|
||||
|
||||
AlignedVector<DType2, N> result;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < N; ++j) {
|
||||
result[j] = cast<DType2>(acc[j]);
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace device::distributed
|
||||
104
python/sglang/jit_kernel/include/sgl_kernel/ffi.h
Normal file
104
python/sglang/jit_kernel/include/sgl_kernel/ffi.h
Normal file
@@ -0,0 +1,104 @@
|
||||
#pragma once
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
#include <tvm/ffi/extra/c_env_api.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
|
||||
namespace host::ffi {
|
||||
|
||||
using tvm::ffi::Tensor, tvm::ffi::TensorView, tvm::ffi::ShapeView;
|
||||
|
||||
inline Tensor empty(ShapeView shape, DLDataType dtype, DLDevice device) {
|
||||
return Tensor::FromEnvAlloc(::TVMFFIEnvTensorAlloc, shape, dtype, device);
|
||||
}
|
||||
|
||||
inline Tensor empty_like(TensorView tensor) {
|
||||
return empty(tensor.shape(), tensor.dtype(), tensor.device());
|
||||
}
|
||||
|
||||
struct _dummy_deleter {
|
||||
void operator()(void*) const {}
|
||||
};
|
||||
|
||||
// template <typename Fn = _dummy_deleter>
|
||||
|
||||
template <typename Fn>
|
||||
struct FromBlobContext {
|
||||
[[no_unique_address]] Fn deleter;
|
||||
int64_t dimension;
|
||||
int64_t* get_shape() {
|
||||
return reinterpret_cast<int64_t*>(this + 1);
|
||||
}
|
||||
int64_t* get_stride() {
|
||||
return this->get_shape() + dimension;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Fn = _dummy_deleter>
|
||||
inline Tensor from_blob(
|
||||
void* data,
|
||||
ShapeView shape,
|
||||
DLDataType dtype,
|
||||
DLDevice device,
|
||||
Fn&& deleter = {},
|
||||
std::optional<ShapeView> stride = {},
|
||||
uint64_t byte_offset = 0) {
|
||||
using Context = FromBlobContext<std::decay_t<Fn>>;
|
||||
const auto ndim = shape.size();
|
||||
const auto ctx = [&] {
|
||||
auto ptr = std::malloc(sizeof(Context) + sizeof(int64_t) * ndim * 2);
|
||||
auto ctx = static_cast<Context*>(ptr);
|
||||
std::construct_at(ctx, std::forward<Fn>(deleter), static_cast<int64_t>(ndim));
|
||||
stdr::copy_n(shape.data(), ndim, ctx->get_shape());
|
||||
if (stride.has_value()) {
|
||||
RuntimeCheck(stride->size() == ndim, "Stride ndim mismatch!");
|
||||
stdr::copy_n(stride->data(), ndim, ctx->get_stride());
|
||||
} else {
|
||||
int64_t stride_val = 1;
|
||||
for (const auto i : irange(ndim)) {
|
||||
const auto j = ndim - 1 - i;
|
||||
ctx->get_stride()[j] = stride_val;
|
||||
stride_val *= shape[j];
|
||||
}
|
||||
}
|
||||
return ctx;
|
||||
}();
|
||||
const auto tensor = DLTensor{
|
||||
.data = data,
|
||||
.device = device,
|
||||
.ndim = static_cast<int32_t>(ndim),
|
||||
.dtype = dtype,
|
||||
.shape = ctx->get_shape(),
|
||||
.strides = ctx->get_stride(),
|
||||
.byte_offset = byte_offset,
|
||||
};
|
||||
const auto blob_deleter = [](DLManagedTensor* self) {
|
||||
auto ctx = static_cast<Context*>(self->manager_ctx);
|
||||
ctx->deleter(self->dl_tensor.data);
|
||||
std::destroy_at(ctx);
|
||||
std::free(ctx);
|
||||
};
|
||||
auto managed_tensor = DLManagedTensor{tensor, ctx, blob_deleter};
|
||||
return Tensor::FromDLPack(&managed_tensor);
|
||||
}
|
||||
|
||||
template <typename Fn = _dummy_deleter>
|
||||
inline Tensor from_blob_like(
|
||||
void* data,
|
||||
TensorView t,
|
||||
Fn&& deleter = {},
|
||||
bool is_contiguous = false, // if override to true, the stride will be ignored
|
||||
uint64_t byte_offset = 0) {
|
||||
const auto stride = is_contiguous ? std::nullopt : std::optional{t.strides()};
|
||||
return from_blob(data, t.shape(), t.dtype(), t.device(), std::forward<Fn>(deleter), stride, byte_offset);
|
||||
}
|
||||
|
||||
} // namespace host::ffi
|
||||
@@ -142,6 +142,11 @@ SGL_DEVICE void PDLTriggerSecondary() {
|
||||
#endif
|
||||
}
|
||||
|
||||
template <std::integral T, std::integral U>
|
||||
SGL_DEVICE constexpr auto div_ceil(T a, U b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Load data with the specified type and offset from a void pointer.
|
||||
* \tparam T The type to load.
|
||||
|
||||
@@ -80,15 +80,11 @@ struct AlignedVector {
|
||||
|
||||
public:
|
||||
/// \brief Vectorized load from `ptr` at the given element `offset`.
|
||||
template <typename U>
|
||||
SGL_DEVICE void load(const U* ptr, std::size_t offset = 0) {
|
||||
static_assert(std::is_same_v<U, T> || std::is_same_v<U, void>);
|
||||
SGL_DEVICE void load(const void* ptr, int64_t offset = 0) {
|
||||
m_storage = reinterpret_cast<const storage_t*>(ptr)[offset];
|
||||
}
|
||||
/// \brief Vectorized store to `ptr` at the given element `offset`.
|
||||
template <typename U>
|
||||
SGL_DEVICE void store(U* ptr, std::size_t offset = 0) const {
|
||||
static_assert(std::is_same_v<U, T> || std::is_same_v<U, void>);
|
||||
SGL_DEVICE void store(void* ptr, int64_t offset = 0) const {
|
||||
reinterpret_cast<storage_t*>(ptr)[offset] = m_storage;
|
||||
}
|
||||
/// \brief Fill all N elements with the same `value`.
|
||||
|
||||
227
python/sglang/jit_kernel/tests/test_custom_all_reduce.py
Normal file
227
python/sglang/jit_kernel/tests/test_custom_all_reduce.py
Normal file
@@ -0,0 +1,227 @@
|
||||
"""
|
||||
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.
|
||||
|
||||
Usage:
|
||||
python -m pytest test_jit_custom_all_reduce.py -v
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from tqdm import tqdm
|
||||
|
||||
import sglang.srt.distributed.parallel_state as ps
|
||||
from sglang.jit_kernel.all_reduce import AllReduceAlgo
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
||||
CustomAllReduceV2,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test parameters (shared between test class and worker)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
TEST_SIZES = [
|
||||
16,
|
||||
32,
|
||||
512,
|
||||
1024,
|
||||
1024 + 16, # weird case
|
||||
4 * 1024,
|
||||
32 * 1024,
|
||||
256 * 1024,
|
||||
2 * 1024 * 1024, # 2M elements
|
||||
4 * 1024 * 1024, # 4M elements
|
||||
]
|
||||
TEST_DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
||||
SHOTS = [
|
||||
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)
|
||||
TEST_LAYERS = 2
|
||||
TEST_LOOP = 16
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test class (runs via pytest, launches torchrun subprocesses)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_torchrun(nproc: int, timeout: int = 300) -> None:
|
||||
"""Launch this script as a torchrun worker and assert success."""
|
||||
cmd = [
|
||||
"torchrun",
|
||||
f"--nproc_per_node={nproc}",
|
||||
__file__,
|
||||
]
|
||||
os.environ["DISABLE_PBAR"] = "1"
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
assert result.returncode == 0, (
|
||||
f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n"
|
||||
f"{result.stdout}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nproc", [2, 3, 4, 5, 6, 7, 8])
|
||||
def test_custom_allreduce(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"
|
||||
)
|
||||
run_torchrun(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).
|
||||
"""
|
||||
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)
|
||||
|
||||
dist.init_process_group(backend="gloo")
|
||||
ps._WORLD = coord = ps.init_world_group(
|
||||
ranks=list(range(world_size)),
|
||||
local_rank=local_rank,
|
||||
backend="nccl",
|
||||
)
|
||||
|
||||
cpu_group = coord.cpu_group
|
||||
nccl_group = coord.device_group
|
||||
assert nccl_group is not None
|
||||
|
||||
max_size = max(TEST_SIZES) * 4
|
||||
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
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def worker_test(
|
||||
device: torch.device,
|
||||
nccl_group: dist.ProcessGroup,
|
||||
comm: CustomAllReduceV2,
|
||||
size: int,
|
||||
dtype: torch.dtype,
|
||||
use_graph: bool,
|
||||
algo: AllReduceAlgo,
|
||||
) -> Optional[RuntimeError]:
|
||||
comm.override_algo = algo
|
||||
|
||||
def get_run_graph_fn():
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
graph_inp = torch.zeros((TEST_LAYERS, size), dtype=dtype, device=device)
|
||||
out_jits = []
|
||||
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)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
def run_graph(x: torch.Tensor) -> torch.Tensor:
|
||||
graph_inp.copy_(x)
|
||||
graph.replay()
|
||||
return out_jit.clone()
|
||||
|
||||
return run_graph
|
||||
|
||||
def get_run_eager_fn():
|
||||
def run_eager(x: torch.Tensor) -> torch.Tensor:
|
||||
eager_inp = x.clone()
|
||||
out_eagers = []
|
||||
for i in range(TEST_LAYERS):
|
||||
out_eagers.append(comm.custom_all_reduce(eager_inp[i]))
|
||||
torch.cuda.synchronize()
|
||||
return torch.stack(out_eagers)
|
||||
|
||||
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)
|
||||
torch.cuda.synchronize()
|
||||
nccl_group.barrier().wait()
|
||||
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()
|
||||
world_size = dist.get_world_size()
|
||||
|
||||
torch.cuda.set_stream(torch.cuda.Stream())
|
||||
|
||||
logging.disable(logging.INFO) # Suppress internal logging for cleaner test output
|
||||
items = list(enumerate(TEST_CONFIG))
|
||||
disable_pbar = os.environ.get("DISABLE_PBAR", "0") == "1" or rank != 0
|
||||
pbar = tqdm(items, desc=f"Testing {world_size} GPUs", disable=disable_pbar)
|
||||
for i, (size, dtype, algo, use_graph) in pbar:
|
||||
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)], device=device)
|
||||
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()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
worker_main()
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
import ctypes
|
||||
import logging
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from functools import partial
|
||||
from typing import Any, List, Optional, Union
|
||||
@@ -15,11 +14,9 @@ import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
|
||||
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
|
||||
from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
|
||||
gpu_p2p_access_check,
|
||||
is_full_nvlink,
|
||||
can_use_custom_all_reduce_with_nvlink,
|
||||
is_weak_contiguous,
|
||||
)
|
||||
from sglang.srt.distributed.parallel_state import in_the_same_node_as
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import (
|
||||
get_bool_env_var,
|
||||
@@ -36,20 +33,6 @@ _is_musa = is_musa()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _can_p2p(rank: int, world_size: int) -> bool:
|
||||
# SGLANG_SKIP_P2P_CHECK can be set to False in sglang
|
||||
SGLANG_SKIP_P2P_CHECK = os.getenv("SGLANG_SKIP_P2P_CHECK", "0") == "1"
|
||||
for i in range(world_size):
|
||||
if i == rank:
|
||||
continue
|
||||
if SGLANG_SKIP_P2P_CHECK:
|
||||
logger.info("Skipping P2P check and trusting the driver's P2P report.")
|
||||
return torch.cuda.can_device_access_peer(rank, i)
|
||||
if not gpu_p2p_access_check(rank, i):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class CustomAllreduce:
|
||||
_SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
|
||||
_MAX_CAR_SIZE = 8192 * 1024
|
||||
@@ -86,35 +69,8 @@ class CustomAllreduce:
|
||||
# e.g. in a non-cuda environment
|
||||
return
|
||||
|
||||
self.group = group
|
||||
|
||||
assert (
|
||||
dist.get_backend(group) != dist.Backend.NCCL
|
||||
), "CustomAllreduce should be attached to a non-NCCL group."
|
||||
|
||||
if not all(in_the_same_node_as(group, source_rank=0)):
|
||||
# No need to initialize custom allreduce for multi-node case.
|
||||
logger.warning(
|
||||
"Custom allreduce is disabled because this process group"
|
||||
" spans across nodes."
|
||||
)
|
||||
return
|
||||
|
||||
rank = dist.get_rank(group=self.group)
|
||||
world_size = dist.get_world_size(group=self.group)
|
||||
if world_size == 1:
|
||||
# No need to initialize custom allreduce for single GPU case.
|
||||
return
|
||||
|
||||
if world_size not in CustomAllreduce._SUPPORTED_WORLD_SIZES:
|
||||
logger.warning(
|
||||
"Custom allreduce is disabled due to an unsupported world"
|
||||
" size: %d. Supported world sizes: %s. To silence this "
|
||||
"warning, specify disable_custom_all_reduce=True explicitly.",
|
||||
world_size,
|
||||
str(CustomAllreduce._SUPPORTED_WORLD_SIZES),
|
||||
)
|
||||
return
|
||||
rank = dist.get_rank(group=group)
|
||||
world_size = dist.get_world_size(group=group)
|
||||
|
||||
if isinstance(device, int):
|
||||
device = torch.device(f"cuda:{device}")
|
||||
@@ -123,46 +79,16 @@ class CustomAllreduce:
|
||||
# now `device` is a `torch.device` object
|
||||
assert isinstance(device, torch.device)
|
||||
self.device = device
|
||||
full_nvlink = can_use_custom_all_reduce_with_nvlink(
|
||||
group=group,
|
||||
device=device,
|
||||
supported_world_size=self._SUPPORTED_WORLD_SIZES,
|
||||
cls_name="CustomAllreduce",
|
||||
)
|
||||
if full_nvlink is None:
|
||||
return # fail to get nvlink status
|
||||
|
||||
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
|
||||
if cuda_visible_devices:
|
||||
device_ids = list(map(int, cuda_visible_devices.split(",")))
|
||||
else:
|
||||
device_ids = list(range(torch.cuda.device_count()))
|
||||
|
||||
physical_device_id = device_ids[device.index]
|
||||
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
|
||||
gather_list = [
|
||||
torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
|
||||
]
|
||||
dist.all_gather(gather_list, tensor, group=self.group)
|
||||
physical_device_ids = [t.item() for t in gather_list]
|
||||
|
||||
# test nvlink first, this will filter out most of the cases
|
||||
# where custom allreduce is not supported
|
||||
# this checks hardware and driver support for NVLink
|
||||
if _is_cuda or _is_hip or _is_musa:
|
||||
full_nvlink = is_full_nvlink(physical_device_ids, world_size)
|
||||
|
||||
if world_size > 2 and not full_nvlink:
|
||||
logger.warning(
|
||||
"Custom allreduce is disabled because it's not supported on"
|
||||
" more than two PCIe-only GPUs. To silence this warning, "
|
||||
"specify disable_custom_all_reduce=True explicitly."
|
||||
)
|
||||
return
|
||||
# test P2P capability, this checks software/cudaruntime support
|
||||
# this is expensive to compute at the first time
|
||||
# then we cache the result
|
||||
# On AMD GPU, p2p is always enabled between XGMI connected GPUs
|
||||
if not _is_hip and not _can_p2p(rank, world_size):
|
||||
logger.warning(
|
||||
"Custom allreduce is disabled because your platform lacks "
|
||||
"GPU P2P capability or P2P test failed. To silence this "
|
||||
"warning, specify disable_custom_all_reduce=True explicitly."
|
||||
)
|
||||
return
|
||||
|
||||
self.group = group
|
||||
self.max_size = max_size
|
||||
self.rank = rank
|
||||
self.world_size = world_size
|
||||
@@ -458,7 +384,17 @@ def dispatch_custom_allreduce():
|
||||
|
||||
On AMD with 1-stage AR enabled, use sglang's CustomAllreduce (has deterministic_all_reduce method).
|
||||
Otherwise use AiterCustomAllreduce if available.
|
||||
|
||||
Set SGLANG_USE_JIT_ALL_REDUCE=1 to use the JIT-compiled v2 implementation.
|
||||
"""
|
||||
# HARDCODED: opt-in flag for v2 JIT all-reduce.
|
||||
# Set SGLANG_USE_JIT_ALL_REDUCE=1 to enable.
|
||||
if _is_cuda and get_bool_env_var("SGLANG_USE_JIT_ALL_REDUCE", default="false"):
|
||||
from .custom_all_reduce_v2 import CustomAllReduceV2
|
||||
|
||||
logger.debug("[AR] Using CustomAllReduceV2 (JIT-compiled)")
|
||||
return CustomAllReduceV2
|
||||
|
||||
if _is_cuda or _is_musa:
|
||||
return CustomAllreduce
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@ import torch.multiprocessing as mp
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
|
||||
from sglang.srt.distributed.parallel_state import in_the_same_node_as
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_musa
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -384,7 +385,92 @@ def is_weak_contiguous(inp: torch.Tensor):
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["gpu_p2p_access_check"]
|
||||
def can_p2p(rank: int, world_size: int) -> bool:
|
||||
# SGLANG_SKIP_P2P_CHECK can be set to False in sglang
|
||||
SGLANG_SKIP_P2P_CHECK = os.getenv("SGLANG_SKIP_P2P_CHECK", "0") == "1"
|
||||
for i in range(world_size):
|
||||
if i == rank:
|
||||
continue
|
||||
if SGLANG_SKIP_P2P_CHECK:
|
||||
logger.info("Skipping P2P check and trusting the driver's P2P report.")
|
||||
return torch.cuda.can_device_access_peer(rank, i)
|
||||
if not gpu_p2p_access_check(rank, i):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def can_use_custom_all_reduce_with_nvlink(
|
||||
group: torch.distributed.ProcessGroup,
|
||||
device: torch.device,
|
||||
supported_world_size: List[int],
|
||||
cls_name: str,
|
||||
) -> Optional[bool]: # None if fail; otherwise return whether NVLink is available
|
||||
assert (
|
||||
dist.get_backend(group) != dist.Backend.NCCL
|
||||
), f"{cls_name} should be attached to a non-NCCL group."
|
||||
|
||||
rank = dist.get_rank(group=group)
|
||||
world_size = dist.get_world_size(group=group)
|
||||
|
||||
# No need to initialize custom allreduce for single GPU case.
|
||||
if world_size == 1:
|
||||
return
|
||||
|
||||
# No need to initialize custom allreduce for multi-node case.
|
||||
if not all(in_the_same_node_as(group, source_rank=0)):
|
||||
logger.warning(
|
||||
f"{cls_name} is disabled because this process group" " spans across nodes."
|
||||
)
|
||||
return
|
||||
|
||||
# For not supported world size, we disable custom allreduce.
|
||||
if world_size not in supported_world_size:
|
||||
logger.warning(
|
||||
f"{cls_name} is disabled due to an unsupported world"
|
||||
f" size: {world_size}. Supported world sizes: {supported_world_size}. "
|
||||
"To silence this warning, specify disable_custom_all_reduce=True explicitly.",
|
||||
)
|
||||
return
|
||||
|
||||
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
|
||||
if cuda_visible_devices:
|
||||
device_ids = list(map(int, cuda_visible_devices.split(",")))
|
||||
else:
|
||||
device_ids = list(range(torch.cuda.device_count()))
|
||||
physical_device_id = device_ids[device.index]
|
||||
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
|
||||
gather_list = [
|
||||
torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
|
||||
]
|
||||
dist.all_gather(gather_list, tensor, group=group)
|
||||
physical_device_ids = [int(t) for t in gather_list]
|
||||
full_nvlink = is_full_nvlink(physical_device_ids, world_size)
|
||||
|
||||
# test nvlink first, this will filter out most of the cases
|
||||
# where custom allreduce is not supported
|
||||
# this checks hardware and driver support for NVLink
|
||||
if world_size > 2 and not full_nvlink:
|
||||
logger.warning(
|
||||
f"{cls_name} is disabled because it's not supported on"
|
||||
" more than two PCIe-only GPUs. To silence this warning, "
|
||||
"specify disable_custom_all_reduce=True explicitly."
|
||||
)
|
||||
return
|
||||
|
||||
# test P2P capability, this checks software/cudaruntime support
|
||||
# this is expensive to compute at the first time
|
||||
# then we cache the result
|
||||
# On AMD GPU, p2p is always enabled between XGMI connected GPUs
|
||||
if not _is_hip and not can_p2p(rank, world_size):
|
||||
logger.warning(
|
||||
f"{cls_name} is disabled because your platform lacks "
|
||||
"GPU P2P capability or P2P test failed. To silence this "
|
||||
"warning, specify disable_custom_all_reduce=True explicitly."
|
||||
)
|
||||
return
|
||||
|
||||
return full_nvlink
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
batch_src, batch_tgt, output_file = pickle.loads(sys.stdin.buffer.read())
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Dict, List, Optional, TypeVar
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.distributed import ProcessGroup
|
||||
|
||||
from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_custom_all_reduce_cls
|
||||
from sglang.srt.distributed import is_in_piecewise_cuda_graph
|
||||
from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
|
||||
can_use_custom_all_reduce_with_nvlink,
|
||||
is_weak_contiguous,
|
||||
)
|
||||
from sglang.srt.utils import is_sm100_supported, log_info_on_rank0
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
INF = 1 << 60
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModeConfig:
|
||||
one_shot_push_threshold: int # below this, use one-shot push
|
||||
one_shot_pull_threshold: int # below this, use one-shot pull
|
||||
|
||||
|
||||
class CustomAllReduceV2:
|
||||
def __init__(
|
||||
self,
|
||||
group: ProcessGroup,
|
||||
device: torch.device,
|
||||
max_pull_size: Optional[int] = None,
|
||||
max_push_size: Optional[int] = None,
|
||||
) -> None:
|
||||
_init_config()
|
||||
self.disabled = True
|
||||
full_nvlink = can_use_custom_all_reduce_with_nvlink(
|
||||
group=group,
|
||||
device=device,
|
||||
supported_world_size=list(THRESHOLD_2_SHOT_MAP.keys()),
|
||||
cls_name="CustomAllReduceV2",
|
||||
)
|
||||
if full_nvlink != True:
|
||||
return
|
||||
|
||||
self.group = group
|
||||
self.rank = dist.get_rank(group=self.group)
|
||||
self.world_size = dist.get_world_size(group=self.group)
|
||||
self.override_shot(None)
|
||||
if max_pull_size is None:
|
||||
max_pull_size = 16 * 1024 * 1024 # default to 16MB
|
||||
if max_push_size is None:
|
||||
max_push_size = self.config.one_shot_push_threshold
|
||||
max_push_size = min(max_push_size, max_pull_size)
|
||||
self.max_pull_size = max_pull_size
|
||||
self.max_push_size = max_push_size
|
||||
self.override_algo: Optional[AllReduceAlgo] = None
|
||||
self.obj = get_custom_all_reduce_cls()(
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
pull_buffer_bytes=self.max_pull_size,
|
||||
push_buffer_bytes=self.max_push_size,
|
||||
graph_input_count=131072,
|
||||
)
|
||||
self._post_init_obj()
|
||||
self.disabled = False
|
||||
log_info_on_rank0(logger, "Custom allreduce v2 initialized successfully")
|
||||
|
||||
def override_shot(self, shot: int | None):
|
||||
if shot is None:
|
||||
self.config = THRESHOLD_2_SHOT_MAP[self.world_size]
|
||||
else:
|
||||
assert shot in (1, 2)
|
||||
threshold = INF if shot == 1 else 0
|
||||
self.config = replace(self.config, one_shot_pull_threshold=threshold)
|
||||
|
||||
@contextmanager
|
||||
def capture(self):
|
||||
try:
|
||||
self.obj.set_cuda_graph_capture(True)
|
||||
yield
|
||||
finally:
|
||||
self.obj.set_cuda_graph_capture(False)
|
||||
if not self.disabled:
|
||||
# cannot call when graph is capturing
|
||||
assert (
|
||||
torch.cuda.is_current_stream_capturing() == False
|
||||
), "Cannot register graph inputs while capturing CUDA graph"
|
||||
pairs = self.obj.share_graph_inputs()
|
||||
handles = [handle for _, handle in pairs]
|
||||
offsets = [offset for offset, _ in pairs]
|
||||
handles_all = self._share_list(handles)
|
||||
offsets_all = self._share_list(offsets)
|
||||
result = [list(zip(o, h)) for o, h in zip(offsets_all, handles_all)]
|
||||
self.obj.register_inputs(result)
|
||||
log_info_on_rank0(logger, f"Registering {len(pairs)} cuda graph addresses")
|
||||
|
||||
def should_custom_ar(self, inp: torch.Tensor) -> bool:
|
||||
"""Check if the input tensor is suitable for custom all-reduce."""
|
||||
if self.disabled:
|
||||
return False
|
||||
inp_size = inp.numel() * inp.element_size()
|
||||
# custom allreduce requires input byte size to be multiples of 16
|
||||
if inp_size % 16 != 0:
|
||||
return False
|
||||
if not is_weak_contiguous(inp):
|
||||
return False
|
||||
return inp_size <= self.max_pull_size
|
||||
|
||||
def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor:
|
||||
if is_in_piecewise_cuda_graph(): # disable inplace optimization
|
||||
try:
|
||||
self.obj.set_cuda_graph_capture(False)
|
||||
return self._all_reduce(input)
|
||||
finally:
|
||||
self.obj.set_cuda_graph_capture(True)
|
||||
return self._all_reduce(input)
|
||||
|
||||
def close(self):
|
||||
if not self.disabled and hasattr(self, "obj"):
|
||||
self.obj.free(self.group)
|
||||
|
||||
def _all_reduce(self, input: torch.Tensor) -> torch.Tensor:
|
||||
"""Perform the actual all-reduce via JIT kernel."""
|
||||
algo = self._determine_algo(input)
|
||||
return torch.from_dlpack(self.obj.all_reduce(input, algo))
|
||||
|
||||
def _determine_algo(self, input: torch.Tensor) -> AllReduceAlgo:
|
||||
if self.override_algo is not None:
|
||||
return self.override_algo
|
||||
input_bytes = input.numel() * input.element_size()
|
||||
if input_bytes <= self.config.one_shot_push_threshold:
|
||||
return AllReduceAlgo.ONE_SHOT_PUSH
|
||||
if input_bytes <= self.config.one_shot_pull_threshold:
|
||||
return AllReduceAlgo.ONE_SHOT_PULL
|
||||
else:
|
||||
return AllReduceAlgo.TWO_SHOT_PULL
|
||||
|
||||
def _post_init_obj(self):
|
||||
handles = [self.obj.share_storage()]
|
||||
result = self._share_list(handles)
|
||||
assert all(len(r) == 1 for r in result)
|
||||
result = [h[0] for h in result]
|
||||
self.obj.post_init(result)
|
||||
|
||||
def _share_list(self, input: List[T]) -> List[List[T]]:
|
||||
input_tensor = torch.tensor(input, dtype=torch.int64, device="cpu")
|
||||
gather_list = [torch.empty_like(input_tensor) for _ in range(self.world_size)]
|
||||
dist.all_gather(gather_list, input_tensor, group=self.group)
|
||||
return [g.tolist() for g in gather_list]
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
|
||||
def _init_config():
|
||||
global THRESHOLD_2_SHOT_MAP
|
||||
KB, MB = 1024, 1024 * 1024
|
||||
|
||||
if is_sm100_supported():
|
||||
# NOTE: This result is based on benchmarks on B200 GPUs
|
||||
THRESHOLD_2_SHOT_MAP = {
|
||||
2: ModeConfig(4 * MB, INF),
|
||||
3: ModeConfig(4 * MB, 4 * MB),
|
||||
4: ModeConfig(2 * MB, 2 * MB),
|
||||
5: ModeConfig(2 * MB, 2 * MB),
|
||||
6: ModeConfig(1 * MB, 1 * MB),
|
||||
7: ModeConfig(896 * KB, 896 * KB),
|
||||
8: ModeConfig(720 * KB, 720 * KB),
|
||||
}
|
||||
else:
|
||||
# NOTE: This result is based on benchmarks on H200 GPUs
|
||||
THRESHOLD_2_SHOT_MAP = {
|
||||
2: ModeConfig(2 * MB, INF),
|
||||
3: ModeConfig(512 * KB, 512 * KB),
|
||||
4: ModeConfig(384 * KB, 256 * KB),
|
||||
5: ModeConfig(256 * KB, 256 * KB),
|
||||
6: ModeConfig(192 * KB, 192 * KB),
|
||||
7: ModeConfig(192 * KB, 192 * KB),
|
||||
8: ModeConfig(160 * KB, 160 * KB),
|
||||
}
|
||||
# TODO: tune on more GPUs, e.g A100
|
||||
|
||||
|
||||
THRESHOLD_2_SHOT_MAP: Dict[int, ModeConfig] = {}
|
||||
Reference in New Issue
Block a user