forked from Karylab-cklius/vllm
Woosuk Kwon
@@ -0,0 +1,210 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Benchmark: Lamport all-gather vs NCCL."""
|
||||
|
||||
import ctypes
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
_cudart = ctypes.CDLL("libcudart.so")
|
||||
IPC = 64
|
||||
|
||||
|
||||
def _cc(r):
|
||||
if r:
|
||||
raise RuntimeError(f"err {r}")
|
||||
|
||||
|
||||
def ipc_buf(sz, rank, ws):
|
||||
p = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
|
||||
_cc(_cudart.cudaMemset(p, 0, sz))
|
||||
_cc(_cudart.cudaDeviceSynchronize())
|
||||
h = (ctypes.c_byte * IPC)()
|
||||
_cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
|
||||
ah = [None] * ws
|
||||
dist.all_gather_object(ah, bytes(h))
|
||||
ptrs = []
|
||||
for i in range(ws):
|
||||
if i == rank:
|
||||
ptrs.append(p.value)
|
||||
else:
|
||||
hh = (ctypes.c_byte * IPC)(*ah[i])
|
||||
pp = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
|
||||
ptrs.append(pp.value)
|
||||
return ptrs
|
||||
|
||||
|
||||
def gpu_timer(fn, warmup=20, repeats=200):
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
s = torch.cuda.Event(enable_timing=True)
|
||||
e = torch.cuda.Event(enable_timing=True)
|
||||
s.record()
|
||||
for _ in range(repeats):
|
||||
fn()
|
||||
e.record()
|
||||
torch.cuda.synchronize()
|
||||
return s.elapsed_time(e) / repeats * 1000
|
||||
|
||||
|
||||
def gpu_timer_graph(fn, warmup=20, repeats=200):
|
||||
"""Time with CUDA graph to exclude CPU overhead."""
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g):
|
||||
fn()
|
||||
for _ in range(5):
|
||||
g.replay()
|
||||
torch.cuda.synchronize()
|
||||
s = torch.cuda.Event(enable_timing=True)
|
||||
e = torch.cuda.Event(enable_timing=True)
|
||||
s.record()
|
||||
for _ in range(repeats):
|
||||
g.replay()
|
||||
e.record()
|
||||
torch.cuda.synchronize()
|
||||
return s.elapsed_time(e) / repeats * 1000
|
||||
|
||||
|
||||
def main():
|
||||
dist.init_process_group("nccl")
|
||||
rank = dist.get_rank()
|
||||
ws = dist.get_world_size()
|
||||
torch.cuda.set_device(rank)
|
||||
dev = f"cuda:{rank}"
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
if rank == 0:
|
||||
from moe_allgather import _load_lib
|
||||
|
||||
lib = _load_lib()
|
||||
dist.barrier()
|
||||
if rank != 0:
|
||||
from moe_allgather import _load_lib
|
||||
|
||||
lib = _load_lib()
|
||||
dist.barrier()
|
||||
from moe_allgather import MoeAllGather
|
||||
|
||||
max_size = 8 * 1024 * 1024
|
||||
bp = ipc_buf(max_size, rank, ws)
|
||||
dist.barrier()
|
||||
|
||||
class FakeCA:
|
||||
pass
|
||||
|
||||
ca = FakeCA()
|
||||
ca.rank = rank
|
||||
ca.world_size = ws
|
||||
ca.device = torch.device(dev)
|
||||
ca.buffer_ptrs = bp
|
||||
ca.max_size = max_size
|
||||
ag = MoeAllGather(ca)
|
||||
dist.barrier()
|
||||
|
||||
configs = [
|
||||
("1tok", 1),
|
||||
("2tok", 2),
|
||||
("4tok", 4),
|
||||
("8tok", 8),
|
||||
("16tok", 16),
|
||||
("32tok", 32),
|
||||
("64tok", 64),
|
||||
("128tok", 128),
|
||||
("256tok", 256),
|
||||
]
|
||||
topk = 8
|
||||
hd = 3584
|
||||
sd = 448
|
||||
|
||||
if rank == 0:
|
||||
print(f"world_size={ws}, max_per_rank={ag.max_per_rank} bytes")
|
||||
print(
|
||||
f"{'config':<12} {'lamport':>10} {'lamp_graph':>10} "
|
||||
f"{'nccl_1ag':>10} {'nccl_graph':>10} {'speedup':>8}"
|
||||
)
|
||||
print("-" * 65)
|
||||
|
||||
for name, N in configs:
|
||||
# Check if data fits in buffer.
|
||||
cursor = 0
|
||||
per_tok = topk * 4 + topk * 4 + hd + sd
|
||||
cursor = N * per_tok
|
||||
cursor = (cursor + 15) & ~15
|
||||
if cursor > ag.max_per_rank:
|
||||
if rank == 0:
|
||||
print(f"{name:<12} {'skip (too large)':>40}")
|
||||
continue
|
||||
|
||||
ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
|
||||
wt = torch.randn(N, topk, dtype=torch.float32, device=dev).abs()
|
||||
hs = torch.randint(0, 255, (N, hd), dtype=torch.uint8, device=dev)
|
||||
sc = torch.randint(0, 255, (N, sd), dtype=torch.uint8, device=dev)
|
||||
inputs = [ids, wt, hs, sc]
|
||||
|
||||
# Custom Lamport kernel.
|
||||
c_outs = [
|
||||
torch.empty(N * ws, *t.shape[1:], dtype=t.dtype, device=dev) for t in inputs
|
||||
]
|
||||
|
||||
def run_lamport():
|
||||
lib.moe_all_gather(
|
||||
ag._buf_ptrs_ptr,
|
||||
ag._counters_ptr,
|
||||
rank,
|
||||
ws,
|
||||
ag.seg_capacity,
|
||||
ag.rank_stride,
|
||||
inputs,
|
||||
c_outs,
|
||||
)
|
||||
|
||||
lam_us = gpu_timer(run_lamport)
|
||||
|
||||
# Lamport with CUDA graph.
|
||||
try:
|
||||
lam_g_us = gpu_timer_graph(run_lamport)
|
||||
except Exception as ex:
|
||||
lam_g_us = float("nan")
|
||||
if rank == 0:
|
||||
print(f" [graph capture failed: {ex}]")
|
||||
|
||||
# NCCL 1×AG (concat into one tensor).
|
||||
cat_inp = torch.cat(
|
||||
[t.reshape(N, -1).contiguous().view(torch.uint8) for t in inputs],
|
||||
dim=1,
|
||||
).contiguous()
|
||||
cat_out = torch.empty(N * ws, cat_inp.shape[1], dtype=torch.uint8, device=dev)
|
||||
|
||||
def run_nccl():
|
||||
dist.all_gather_into_tensor(cat_out, cat_inp)
|
||||
|
||||
nccl_us = gpu_timer(run_nccl)
|
||||
|
||||
# NCCL with CUDA graph.
|
||||
try:
|
||||
nccl_g_us = gpu_timer_graph(run_nccl)
|
||||
except Exception:
|
||||
nccl_g_us = float("nan")
|
||||
|
||||
if rank == 0:
|
||||
speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
|
||||
print(
|
||||
f"{name:<12} {lam_us:>9.1f}µ {lam_g_us:>9.1f}µ "
|
||||
f"{nccl_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x"
|
||||
)
|
||||
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,137 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Test for the Lamport-based MoE all-gather kernel."""
|
||||
|
||||
import ctypes
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
_cudart = ctypes.CDLL("libcudart.so")
|
||||
IPC = 64
|
||||
|
||||
|
||||
def _cc(r):
|
||||
if r:
|
||||
raise RuntimeError(f"CUDA err {r}")
|
||||
|
||||
|
||||
def ipc_buf(sz, rank, ws):
|
||||
p = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
|
||||
_cc(_cudart.cudaMemset(p, 0, sz))
|
||||
_cc(_cudart.cudaDeviceSynchronize())
|
||||
h = (ctypes.c_byte * IPC)()
|
||||
_cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
|
||||
ah = [None] * ws
|
||||
dist.all_gather_object(ah, bytes(h))
|
||||
ptrs = []
|
||||
for i in range(ws):
|
||||
if i == rank:
|
||||
ptrs.append(p.value)
|
||||
else:
|
||||
hh = (ctypes.c_byte * IPC)(*ah[i])
|
||||
pp = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
|
||||
ptrs.append(pp.value)
|
||||
return ptrs
|
||||
|
||||
|
||||
def main():
|
||||
dist.init_process_group("nccl")
|
||||
rank = dist.get_rank()
|
||||
ws = dist.get_world_size()
|
||||
torch.cuda.set_device(rank)
|
||||
dev = f"cuda:{rank}"
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
if rank == 0:
|
||||
from moe_allgather import MoeAllGather, _load_lib
|
||||
|
||||
_load_lib()
|
||||
dist.barrier()
|
||||
from moe_allgather import MoeAllGather, _load_lib
|
||||
|
||||
_load_lib()
|
||||
dist.barrier()
|
||||
|
||||
max_size = 8 * 1024 * 1024
|
||||
bp = ipc_buf(max_size, rank, ws)
|
||||
dist.barrier()
|
||||
|
||||
# Build a fake ca_comm-like object.
|
||||
class FakeCA:
|
||||
pass
|
||||
|
||||
ca = FakeCA()
|
||||
ca.rank = rank
|
||||
ca.world_size = ws
|
||||
ca.device = torch.device(dev)
|
||||
ca.buffer_ptrs = bp
|
||||
ca.max_size = max_size
|
||||
# meta_ptrs not needed for Lamport approach
|
||||
ag = MoeAllGather(ca)
|
||||
dist.barrier()
|
||||
|
||||
errors = 0
|
||||
|
||||
# Test with various token counts.
|
||||
for N in [1, 4, 16, 64]:
|
||||
topk = 8
|
||||
hd = 3584
|
||||
sd = 448
|
||||
ids = (
|
||||
torch.arange(N * topk, dtype=torch.int32, device=dev) + rank * 1000
|
||||
).reshape(N, topk)
|
||||
wt = torch.ones(N, topk, dtype=torch.float32, device=dev) * (rank + 1) * 0.1
|
||||
hs = torch.full((N, hd), rank + 1, dtype=torch.uint8, device=dev)
|
||||
sc = torch.full((N, sd), rank + 1, dtype=torch.uint8, device=dev)
|
||||
|
||||
ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
|
||||
|
||||
for src in range(ws):
|
||||
s, e = src * N, (src + 1) * N
|
||||
exp_ids = (
|
||||
torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
|
||||
).reshape(N, topk)
|
||||
if not torch.equal(ids_g[s:e], exp_ids):
|
||||
print(f"[{rank}] FAIL ids src={src} N={N}")
|
||||
errors += 1
|
||||
exp_wt = torch.full(
|
||||
(N, topk), (src + 1) * 0.1, dtype=torch.float32, device=dev
|
||||
)
|
||||
if not torch.allclose(wt_g[s:e], exp_wt):
|
||||
print(f"[{rank}] FAIL wt src={src} N={N}")
|
||||
errors += 1
|
||||
exp_hs = torch.full((N, hd), src + 1, dtype=torch.uint8, device=dev)
|
||||
if not torch.equal(hs_g[s:e], exp_hs):
|
||||
print(f"[{rank}] FAIL hs src={src} N={N}")
|
||||
errors += 1
|
||||
exp_sc = torch.full((N, sd), src + 1, dtype=torch.uint8, device=dev)
|
||||
if not torch.equal(sc_g[s:e], exp_sc):
|
||||
print(f"[{rank}] FAIL sc src={src} N={N}")
|
||||
errors += 1
|
||||
|
||||
# Without scales.
|
||||
ids_g2, wt_g2, hs_g2, _ = ag.gather(ids, wt, hs)
|
||||
for src in range(ws):
|
||||
s, e = src * N, (src + 1) * N
|
||||
exp_ids = (
|
||||
torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
|
||||
).reshape(N, topk)
|
||||
if not torch.equal(ids_g2[s:e], exp_ids):
|
||||
print(f"[{rank}] FAIL no-sc ids src={src} N={N}")
|
||||
errors += 1
|
||||
|
||||
dist.barrier()
|
||||
print(
|
||||
f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'} (ws={ws})"
|
||||
)
|
||||
dist.destroy_process_group()
|
||||
return errors
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,173 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Stress test: random data, check bitwise correctness against NCCL."""
|
||||
|
||||
import ctypes
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
_cudart = ctypes.CDLL("libcudart.so")
|
||||
IPC = 64
|
||||
|
||||
|
||||
def _cc(r):
|
||||
if r:
|
||||
raise RuntimeError(f"CUDA err {r}")
|
||||
|
||||
|
||||
def ipc_buf(sz, rank, ws):
|
||||
p = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
|
||||
_cc(_cudart.cudaMemset(p, 0, sz))
|
||||
_cc(_cudart.cudaDeviceSynchronize())
|
||||
h = (ctypes.c_byte * IPC)()
|
||||
_cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
|
||||
ah = [None] * ws
|
||||
dist.all_gather_object(ah, bytes(h))
|
||||
ptrs = []
|
||||
for i in range(ws):
|
||||
if i == rank:
|
||||
ptrs.append(p.value)
|
||||
else:
|
||||
hh = (ctypes.c_byte * IPC)(*ah[i])
|
||||
pp = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
|
||||
ptrs.append(pp.value)
|
||||
return ptrs
|
||||
|
||||
|
||||
def main():
|
||||
dist.init_process_group("nccl")
|
||||
rank = dist.get_rank()
|
||||
ws = dist.get_world_size()
|
||||
torch.cuda.set_device(rank)
|
||||
dev = f"cuda:{rank}"
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
if rank == 0:
|
||||
from moe_allgather import _load_lib
|
||||
|
||||
_load_lib()
|
||||
dist.barrier()
|
||||
from moe_allgather import MoeAllGather, _load_lib
|
||||
|
||||
_load_lib()
|
||||
dist.barrier()
|
||||
|
||||
max_size = 8 * 1024 * 1024
|
||||
bp = ipc_buf(max_size, rank, ws)
|
||||
dist.barrier()
|
||||
|
||||
class FakeCA:
|
||||
pass
|
||||
|
||||
ca = FakeCA()
|
||||
ca.rank = rank
|
||||
ca.world_size = ws
|
||||
ca.device = torch.device(dev)
|
||||
ca.buffer_ptrs = bp
|
||||
ca.max_size = max_size
|
||||
ag = MoeAllGather(ca)
|
||||
dist.barrier()
|
||||
|
||||
topk = 8
|
||||
hd = 3584
|
||||
sd = 448
|
||||
errors = 0
|
||||
total_checks = 0
|
||||
sentinel_collisions = 0
|
||||
|
||||
for trial in range(200):
|
||||
# All ranks must use the same N for NCCL reference.
|
||||
N_tensor = torch.randint(1, 65, (1,), device=dev)
|
||||
dist.broadcast(N_tensor, src=0)
|
||||
N = N_tensor.item()
|
||||
# Random data including possible sentinel values
|
||||
ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
|
||||
wt = torch.randn(N, topk, dtype=torch.float32, device=dev)
|
||||
hs = torch.randint(0, 256, (N, hd), dtype=torch.uint8, device=dev)
|
||||
sc = torch.randint(0, 256, (N, sd), dtype=torch.uint8, device=dev)
|
||||
|
||||
# Count sentinel patterns in hidden_states (as uint32 view)
|
||||
hs_u32 = hs.view(torch.int32)
|
||||
sentinel_collisions += (hs_u32 == 0x80000000).sum().item()
|
||||
|
||||
# Custom kernel
|
||||
ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
|
||||
|
||||
# NCCL reference
|
||||
ids_ref = torch.empty(N * ws, topk, dtype=torch.int32, device=dev)
|
||||
wt_ref = torch.empty(N * ws, topk, dtype=torch.float32, device=dev)
|
||||
hs_ref = torch.empty(N * ws, hd, dtype=torch.uint8, device=dev)
|
||||
sc_ref = torch.empty(N * ws, sd, dtype=torch.uint8, device=dev)
|
||||
dist.all_gather_into_tensor(ids_ref, ids)
|
||||
dist.all_gather_into_tensor(wt_ref, wt)
|
||||
dist.all_gather_into_tensor(hs_ref, hs)
|
||||
dist.all_gather_into_tensor(sc_ref, sc)
|
||||
|
||||
# Compare
|
||||
if not torch.equal(ids_g, ids_ref):
|
||||
mismatches = (ids_g != ids_ref).sum().item()
|
||||
if trial < 5 or mismatches > 0:
|
||||
print(f"[{rank}] trial={trial} ids MISMATCH: {mismatches} elements")
|
||||
errors += 1
|
||||
if not torch.equal(wt_g, wt_ref):
|
||||
# Check for -0 vs +0 differences
|
||||
bit_diff = wt_g.view(torch.int32) != wt_ref.view(torch.int32)
|
||||
neg_zero_mask = wt_ref.view(torch.int32) == 0x80000000
|
||||
real_errors = bit_diff & ~neg_zero_mask
|
||||
if real_errors.any():
|
||||
print(
|
||||
f"[{rank}] trial={trial} wt MISMATCH (non-negzero): {real_errors.sum().item()}"
|
||||
)
|
||||
errors += 1
|
||||
if not torch.equal(hs_g, hs_ref):
|
||||
mismatches = (hs_g != hs_ref).sum().item()
|
||||
# Check if mismatches are due to sentinel collision
|
||||
hs_g_u32 = hs_g.view(torch.int32)
|
||||
hs_ref_u32 = hs_ref.view(torch.int32)
|
||||
diff_mask = hs_g_u32 != hs_ref_u32
|
||||
sentinel_mask = (hs_ref_u32 == 0x80000000) & diff_mask
|
||||
non_sentinel = diff_mask & ~sentinel_mask
|
||||
if non_sentinel.any():
|
||||
print(
|
||||
f"[{rank}] trial={trial} hs NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
|
||||
)
|
||||
errors += 1
|
||||
elif sentinel_mask.any():
|
||||
if trial < 3:
|
||||
print(
|
||||
f"[{rank}] trial={trial} hs sentinel collision: "
|
||||
f"{sentinel_mask.sum().item()} words (expected rare)"
|
||||
)
|
||||
if not torch.equal(sc_g, sc_ref):
|
||||
mismatches = (sc_g != sc_ref).sum().item()
|
||||
sc_g_u32 = sc_g.view(torch.int32) if sc_g.numel() % 4 == 0 else None
|
||||
if sc_g_u32 is not None:
|
||||
sc_ref_u32 = sc_ref.view(torch.int32)
|
||||
diff_mask = sc_g_u32 != sc_ref_u32
|
||||
sentinel_mask = (sc_ref_u32 == 0x80000000) & diff_mask
|
||||
non_sentinel = diff_mask & ~sentinel_mask
|
||||
if non_sentinel.any():
|
||||
print(
|
||||
f"[{rank}] trial={trial} sc NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
|
||||
)
|
||||
errors += 1
|
||||
|
||||
total_checks += 1
|
||||
|
||||
dist.barrier()
|
||||
print(
|
||||
f"[rank {rank}] {total_checks} trials, {errors} real errors, "
|
||||
f"{sentinel_collisions} sentinel patterns in hs data. "
|
||||
f"{'PASSED' if errors == 0 else 'FAILED'}"
|
||||
)
|
||||
dist.destroy_process_group()
|
||||
return errors
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,204 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Test + benchmark for Lamport MoE reduce-scatter."""
|
||||
|
||||
import ctypes
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
_cudart = ctypes.CDLL("libcudart.so")
|
||||
IPC = 64
|
||||
|
||||
|
||||
def _cc(r):
|
||||
if r:
|
||||
raise RuntimeError(f"CUDA err {r}")
|
||||
|
||||
|
||||
def ipc_buf(sz, rank, ws):
|
||||
p = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
|
||||
_cc(_cudart.cudaMemset(p, 0, sz))
|
||||
_cc(_cudart.cudaDeviceSynchronize())
|
||||
h = (ctypes.c_byte * IPC)()
|
||||
_cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
|
||||
ah = [None] * ws
|
||||
dist.all_gather_object(ah, bytes(h))
|
||||
ptrs = []
|
||||
for i in range(ws):
|
||||
if i == rank:
|
||||
ptrs.append(p.value)
|
||||
else:
|
||||
hh = (ctypes.c_byte * IPC)(*ah[i])
|
||||
pp = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
|
||||
ptrs.append(pp.value)
|
||||
return ptrs
|
||||
|
||||
|
||||
def gpu_timer_graph(fn, warmup=20, repeats=200):
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g):
|
||||
fn()
|
||||
for _ in range(5):
|
||||
g.replay()
|
||||
torch.cuda.synchronize()
|
||||
s = torch.cuda.Event(enable_timing=True)
|
||||
e = torch.cuda.Event(enable_timing=True)
|
||||
s.record()
|
||||
for _ in range(repeats):
|
||||
g.replay()
|
||||
e.record()
|
||||
torch.cuda.synchronize()
|
||||
return s.elapsed_time(e) / repeats * 1000
|
||||
|
||||
|
||||
def main():
|
||||
dist.init_process_group("nccl")
|
||||
rank = dist.get_rank()
|
||||
ws = dist.get_world_size()
|
||||
torch.cuda.set_device(rank)
|
||||
dev = f"cuda:{rank}"
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
if rank == 0:
|
||||
from moe_reduce_scatter import _load_lib
|
||||
|
||||
_load_lib()
|
||||
dist.barrier()
|
||||
if rank != 0:
|
||||
from moe_reduce_scatter import _load_lib
|
||||
|
||||
_load_lib()
|
||||
dist.barrier()
|
||||
from moe_reduce_scatter import MoeReduceScatter
|
||||
|
||||
max_size = 8 * 1024 * 1024
|
||||
bp = ipc_buf(max_size, rank, ws)
|
||||
dist.barrier()
|
||||
|
||||
class FakeCA:
|
||||
pass
|
||||
|
||||
ca = FakeCA()
|
||||
ca.rank = rank
|
||||
ca.world_size = ws
|
||||
ca.device = torch.device(dev)
|
||||
ca.buffer_ptrs = bp
|
||||
ca.max_size = max_size
|
||||
rs = MoeReduceScatter(ca)
|
||||
dist.barrier()
|
||||
|
||||
D = 7168 # DeepSeek V3 hidden_dim
|
||||
errors = 0
|
||||
|
||||
# ---- Correctness tests ----
|
||||
for N_per_rank in [1, 4, 16]:
|
||||
N_total = N_per_rank * ws
|
||||
|
||||
# Each rank gets a deterministic input.
|
||||
torch.manual_seed(42)
|
||||
# All ranks create the SAME "ground truth" inputs for each rank.
|
||||
all_inputs = [
|
||||
torch.randn(N_total, D, dtype=torch.bfloat16, device=dev) for _ in range(ws)
|
||||
]
|
||||
# This rank's input is all_inputs[rank].
|
||||
my_input = all_inputs[rank]
|
||||
|
||||
# Custom reduce-scatter.
|
||||
custom_out = rs.reduce_scatter(my_input)
|
||||
|
||||
# NCCL reference: reduce_scatter_tensor.
|
||||
nccl_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
|
||||
dist.reduce_scatter_tensor(nccl_out, my_input)
|
||||
|
||||
# bf16 summation order differs between our kernel and NCCL,
|
||||
# giving ~1-2 ULP differences. Use generous tolerance.
|
||||
max_diff = (custom_out.float() - nccl_out.float()).abs().max().item()
|
||||
if not torch.allclose(custom_out, nccl_out, atol=0.125, rtol=0.01):
|
||||
mismatches = (
|
||||
((custom_out.float() - nccl_out.float()).abs() > 0.125).sum().item()
|
||||
)
|
||||
print(
|
||||
f"[{rank}] N_per_rank={N_per_rank} MISMATCH: "
|
||||
f"max_diff={max_diff:.6f}, mismatches={mismatches}"
|
||||
)
|
||||
errors += 1
|
||||
else:
|
||||
if rank == 0:
|
||||
print(f" N_per_rank={N_per_rank}: PASS (max_diff={max_diff:.6f})")
|
||||
|
||||
# ---- Benchmark ----
|
||||
if rank == 0:
|
||||
print(f"\nworld_size={ws}, max_per_rank={rs.max_per_rank} bytes")
|
||||
print(
|
||||
f"{'config':<12} {'lamport':>10} {'lamp_graph':>10} "
|
||||
f"{'nccl':>10} {'nccl_graph':>10} {'speedup':>8}"
|
||||
)
|
||||
print("-" * 65)
|
||||
|
||||
configs = [
|
||||
("1tok", 1),
|
||||
("2tok", 2),
|
||||
("4tok", 4),
|
||||
("8tok", 8),
|
||||
("16tok", 16),
|
||||
("32tok", 32),
|
||||
("64tok", 64),
|
||||
("128tok", 128),
|
||||
("256tok", 256),
|
||||
]
|
||||
|
||||
for name, N_per_rank in configs:
|
||||
N_total = N_per_rank * ws
|
||||
input_bytes = N_total * D * 2 # bf16
|
||||
if input_bytes > rs.max_per_rank:
|
||||
if rank == 0:
|
||||
print(f"{name:<12} {'skip (too large)':>40}")
|
||||
continue
|
||||
|
||||
inp = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
|
||||
c_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
|
||||
n_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
|
||||
|
||||
def run_lamport():
|
||||
rs.reduce_scatter(inp)
|
||||
|
||||
def run_nccl():
|
||||
dist.reduce_scatter_tensor(n_out, inp)
|
||||
|
||||
from bench_moe_allgather import gpu_timer
|
||||
|
||||
lam_us = gpu_timer(run_lamport)
|
||||
try:
|
||||
lam_g_us = gpu_timer_graph(run_lamport)
|
||||
except Exception:
|
||||
lam_g_us = float("nan")
|
||||
|
||||
nccl_us = gpu_timer(run_nccl)
|
||||
try:
|
||||
nccl_g_us = gpu_timer_graph(run_nccl)
|
||||
except Exception:
|
||||
nccl_g_us = float("nan")
|
||||
|
||||
if rank == 0:
|
||||
speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
|
||||
print(
|
||||
f"{name:<12} {lam_us:>9.1f}µ {lam_g_us:>9.1f}µ "
|
||||
f"{nccl_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x"
|
||||
)
|
||||
|
||||
dist.barrier()
|
||||
print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
|
||||
dist.destroy_process_group()
|
||||
return errors
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,270 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Test + benchmark: fused RS + residual + RMSNorm vs separate kernels."""
|
||||
|
||||
import ctypes
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
_cudart = ctypes.CDLL("libcudart.so")
|
||||
IPC = 64
|
||||
|
||||
|
||||
def _cc(r):
|
||||
if r:
|
||||
raise RuntimeError(f"CUDA err {r}")
|
||||
|
||||
|
||||
def ipc_buf(sz, rank, ws):
|
||||
p = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
|
||||
_cc(_cudart.cudaMemset(p, 0, sz))
|
||||
_cc(_cudart.cudaDeviceSynchronize())
|
||||
h = (ctypes.c_byte * IPC)()
|
||||
_cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
|
||||
ah = [None] * ws
|
||||
dist.all_gather_object(ah, bytes(h))
|
||||
ptrs = []
|
||||
for i in range(ws):
|
||||
if i == rank:
|
||||
ptrs.append(p.value)
|
||||
else:
|
||||
hh = (ctypes.c_byte * IPC)(*ah[i])
|
||||
pp = ctypes.c_void_p()
|
||||
_cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
|
||||
ptrs.append(pp.value)
|
||||
return ptrs
|
||||
|
||||
|
||||
def rms_norm_ref(x, gamma, eps):
|
||||
"""Reference RMSNorm in fp32."""
|
||||
xf = x.float()
|
||||
rms = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
|
||||
return (xf * rms * gamma.float()).to(x.dtype)
|
||||
|
||||
|
||||
def gpu_timer(fn, warmup=20, repeats=200):
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
s = torch.cuda.Event(enable_timing=True)
|
||||
e = torch.cuda.Event(enable_timing=True)
|
||||
s.record()
|
||||
for _ in range(repeats):
|
||||
fn()
|
||||
e.record()
|
||||
torch.cuda.synchronize()
|
||||
return s.elapsed_time(e) / repeats * 1000
|
||||
|
||||
|
||||
def gpu_timer_graph(fn, warmup=20, repeats=200):
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g):
|
||||
fn()
|
||||
for _ in range(5):
|
||||
g.replay()
|
||||
torch.cuda.synchronize()
|
||||
s = torch.cuda.Event(enable_timing=True)
|
||||
e = torch.cuda.Event(enable_timing=True)
|
||||
s.record()
|
||||
for _ in range(repeats):
|
||||
g.replay()
|
||||
e.record()
|
||||
torch.cuda.synchronize()
|
||||
return s.elapsed_time(e) / repeats * 1000
|
||||
|
||||
|
||||
def main():
|
||||
dist.init_process_group("nccl")
|
||||
rank = dist.get_rank()
|
||||
ws = dist.get_world_size()
|
||||
torch.cuda.set_device(rank)
|
||||
dev = f"cuda:{rank}"
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
# Compile fused kernel
|
||||
if rank == 0:
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
load(
|
||||
name="moe_rs_fused_kernel",
|
||||
sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
|
||||
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
||||
verbose=False,
|
||||
)
|
||||
dist.barrier()
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
fused_lib = load(
|
||||
name="moe_rs_fused_kernel",
|
||||
sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
|
||||
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
# Also compile separate RS kernel for comparison
|
||||
from moe_reduce_scatter import MoeReduceScatter
|
||||
|
||||
max_size = 8 * 1024 * 1024
|
||||
bp = ipc_buf(max_size, rank, ws)
|
||||
dist.barrier()
|
||||
|
||||
class FakeCA:
|
||||
pass
|
||||
|
||||
ca = FakeCA()
|
||||
ca.rank = rank
|
||||
ca.world_size = ws
|
||||
ca.device = torch.device(dev)
|
||||
ca.buffer_ptrs = bp
|
||||
ca.max_size = max_size
|
||||
|
||||
rs_separate = MoeReduceScatter(ca)
|
||||
|
||||
# Fused kernel setup (uses same buffer layout as MoeReduceScatter)
|
||||
half_size = (max_size // 2) & ~15
|
||||
buf_offset = half_size
|
||||
fused_seg_cap = (half_size // 2) & ~15
|
||||
fused_rank_stride = (fused_seg_cap // ws) & ~15
|
||||
|
||||
fused_buf_ptrs = torch.zeros(8, dtype=torch.int64, device=dev)
|
||||
for i in range(ws):
|
||||
fused_buf_ptrs[i] = bp[i] + buf_offset
|
||||
fused_counters = torch.zeros(3, dtype=torch.int32, device=dev)
|
||||
|
||||
# Init sentinels for fused kernel's buffer region
|
||||
fused_lib.lamport_init(bp[rank] + buf_offset, half_size)
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier()
|
||||
|
||||
D = 7168
|
||||
eps = 1e-6
|
||||
gamma = torch.randn(D, dtype=torch.bfloat16, device=dev).abs() + 0.5
|
||||
errors = 0
|
||||
|
||||
# ---- Correctness ----
|
||||
for N_per_rank in [1, 4]:
|
||||
N_total = N_per_rank * ws
|
||||
torch.manual_seed(42 + rank)
|
||||
moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
|
||||
residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
|
||||
|
||||
# Reference: NCCL RS + add + norm
|
||||
rs_ref = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
|
||||
dist.reduce_scatter_tensor(rs_ref, moe_out)
|
||||
ref_residual = residual + rs_ref
|
||||
ref_normed = rms_norm_ref(ref_residual, gamma, eps)
|
||||
|
||||
# Fused kernel
|
||||
normed_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
|
||||
residual_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
|
||||
fused_lib.moe_rs_fused(
|
||||
fused_buf_ptrs.data_ptr(),
|
||||
fused_counters.data_ptr(),
|
||||
rank,
|
||||
ws,
|
||||
fused_seg_cap,
|
||||
fused_rank_stride,
|
||||
moe_out,
|
||||
residual,
|
||||
gamma,
|
||||
normed_out,
|
||||
residual_out,
|
||||
eps,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Compare
|
||||
max_diff_res = (residual_out.float() - ref_residual.float()).abs().max().item()
|
||||
max_diff_norm = (normed_out.float() - ref_normed.float()).abs().max().item()
|
||||
|
||||
ok = max_diff_res < 0.125 and max_diff_norm < 0.125
|
||||
if rank == 0:
|
||||
print(
|
||||
f" N_per_rank={N_per_rank}: {'PASS' if ok else 'FAIL'} "
|
||||
f"(res_diff={max_diff_res:.4f}, norm_diff={max_diff_norm:.4f})"
|
||||
)
|
||||
if not ok:
|
||||
errors += 1
|
||||
|
||||
# ---- Benchmark ----
|
||||
if rank == 0:
|
||||
print(f"\nBenchmark: D={D}, world_size={ws}")
|
||||
print(
|
||||
f"{'config':<10} {'fused':>10} {'fused_g':>10} "
|
||||
f"{'RS+norm':>10} {'RS+norm_g':>10} {'speedup':>8}"
|
||||
)
|
||||
print("-" * 58)
|
||||
|
||||
for N_per_rank in [1, 2, 4, 8]:
|
||||
N_total = N_per_rank * ws
|
||||
input_bytes = N_total * D * 2
|
||||
if input_bytes > fused_rank_stride:
|
||||
if rank == 0:
|
||||
print(f"{N_per_rank}tok skip (too large)")
|
||||
continue
|
||||
|
||||
moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
|
||||
residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
|
||||
normed_out = torch.empty_like(residual)
|
||||
residual_out = torch.empty_like(residual)
|
||||
rs_out = torch.empty_like(residual)
|
||||
|
||||
# Fused
|
||||
def run_fused():
|
||||
fused_lib.moe_rs_fused(
|
||||
fused_buf_ptrs.data_ptr(),
|
||||
fused_counters.data_ptr(),
|
||||
rank,
|
||||
ws,
|
||||
fused_seg_cap,
|
||||
fused_rank_stride,
|
||||
moe_out,
|
||||
residual,
|
||||
gamma,
|
||||
normed_out,
|
||||
residual_out,
|
||||
eps,
|
||||
)
|
||||
|
||||
# Separate: RS + add + norm (our Lamport RS + triton-like ops)
|
||||
def run_separate():
|
||||
rs_separate.reduce_scatter(moe_out)
|
||||
# Simulate add + RMSNorm (in practice this is a fused triton kernel)
|
||||
tmp = residual + rs_out
|
||||
torch.rsqrt(tmp.float().pow(2).mean(-1, keepdim=True) + eps)
|
||||
|
||||
fused_us = gpu_timer(run_fused)
|
||||
try:
|
||||
fused_g = gpu_timer_graph(run_fused)
|
||||
except Exception:
|
||||
fused_g = float("nan")
|
||||
|
||||
sep_us = gpu_timer(run_separate)
|
||||
try:
|
||||
sep_g = gpu_timer_graph(run_separate)
|
||||
except Exception:
|
||||
sep_g = float("nan")
|
||||
|
||||
if rank == 0:
|
||||
speedup = sep_g / fused_g if fused_g > 0 else float("nan")
|
||||
print(
|
||||
f"{N_per_rank}tok {fused_us:>9.1f}µ {fused_g:>9.1f}µ "
|
||||
f"{sep_us:>9.1f}µ {sep_g:>9.1f}µ {speedup:>7.2f}x"
|
||||
)
|
||||
|
||||
dist.barrier()
|
||||
print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
|
||||
dist.destroy_process_group()
|
||||
return errors
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user