From 15f1df36e227d719086b108bbbcd2ab19be700d6 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 15 Apr 2026 06:51:25 +0000 Subject: [PATCH] Add tests Signed-off-by: Woosuk Kwon --- .../custom_ar/bench_moe_allgather.py | 210 ++++++++++++++ .../custom_ar/test_moe_allgather.py | 137 +++++++++ .../custom_ar/test_moe_allgather_stress.py | 173 +++++++++++ .../custom_ar/test_moe_reduce_scatter.py | 204 +++++++++++++ csrc/standalone/custom_ar/test_rs_fused.py | 270 ++++++++++++++++++ 5 files changed, 994 insertions(+) create mode 100644 csrc/standalone/custom_ar/bench_moe_allgather.py create mode 100644 csrc/standalone/custom_ar/test_moe_allgather.py create mode 100644 csrc/standalone/custom_ar/test_moe_allgather_stress.py create mode 100644 csrc/standalone/custom_ar/test_moe_reduce_scatter.py create mode 100644 csrc/standalone/custom_ar/test_rs_fused.py diff --git a/csrc/standalone/custom_ar/bench_moe_allgather.py b/csrc/standalone/custom_ar/bench_moe_allgather.py new file mode 100644 index 00000000000..7c9916d6ca4 --- /dev/null +++ b/csrc/standalone/custom_ar/bench_moe_allgather.py @@ -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() diff --git a/csrc/standalone/custom_ar/test_moe_allgather.py b/csrc/standalone/custom_ar/test_moe_allgather.py new file mode 100644 index 00000000000..0d2285752b5 --- /dev/null +++ b/csrc/standalone/custom_ar/test_moe_allgather.py @@ -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()) diff --git a/csrc/standalone/custom_ar/test_moe_allgather_stress.py b/csrc/standalone/custom_ar/test_moe_allgather_stress.py new file mode 100644 index 00000000000..ce6d4b9b76f --- /dev/null +++ b/csrc/standalone/custom_ar/test_moe_allgather_stress.py @@ -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()) diff --git a/csrc/standalone/custom_ar/test_moe_reduce_scatter.py b/csrc/standalone/custom_ar/test_moe_reduce_scatter.py new file mode 100644 index 00000000000..db3ae7672d5 --- /dev/null +++ b/csrc/standalone/custom_ar/test_moe_reduce_scatter.py @@ -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()) diff --git a/csrc/standalone/custom_ar/test_rs_fused.py b/csrc/standalone/custom_ar/test_rs_fused.py new file mode 100644 index 00000000000..8b01fcc4474 --- /dev/null +++ b/csrc/standalone/custom_ar/test_rs_fused.py @@ -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())