From 0a36982c921a70acfbc0c6a0c08c909a362cbc7e Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 11 Oct 2026 08:18:09 +0900 Subject: [PATCH] cuda: a reduce-scatter in rank order over NCCL send and receive (rank_scatter.zig) nccl.zig binds ncclSend and ncclRecv beside the collectives. rank_scatter.zig cuts a [rows][width] fp32 partial by whole rows (the first rows % world ranks one more, as canon.partition cuts equal ranks), sends each peer its rows and receives its own from each peer in one group, then rank_sum.cu adds the terms rank 0 first, rounding each add. Every rank so holds the bits an all-gather followed by a rank-order sum gives, with each rank receiving only its own rows. NCCL's ncclReduceScatter is not used: it adds in its ring's order, not rank order. tf-cuda-test rank-scatter [id] checks every rank's rows against an all-gather plus a host rank-order sum (rows 1, 4, 64, 2047 and 2048 by 4096, three seeds), the ranks joined by NCCL_COMM_ID or by rank 0's unique id in a file, and times both exchanges. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/rank_sum.cu | 11 +++ zig/src/cuda/kernels.zig | 1 + zig/src/cuda/nccl.zig | 2 + zig/src/cuda/rank_scatter.zig | 80 +++++++++++++++++ zig/src/cuda/root.zig | 2 + zig/tests/cuda/main.zig | 3 + zig/tests/cuda/rank_scatter_tests.zig | 125 ++++++++++++++++++++++++++ 8 files changed, 225 insertions(+) create mode 100644 zig/kernels/cuda/rank_sum.cu create mode 100644 zig/src/cuda/rank_scatter.zig create mode 100644 zig/tests/cuda/rank_scatter_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index d4075393bc..a56c8e4992 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -28,6 +28,7 @@ const kernels = [_]Kernel{ .{ .name = "exlr8", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 lane projections, the Metal kernels' order .{ .name = "exlr8_dense", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 dense matrices by row, the Metal order .{ .name = "nvfp4_experts", .flags = &.{"-O3"} }, // nvfp4/experts.cu's device code, the NVFP4 expert kernels + .{ .name = "rank_sum", .flags = &.{"-O3"} }, // ours: the rank-order sum of a reduce-scatter (rank_scatter.zig) .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws .{ .name = "train", .flags = &.{"-O3"} }, // ours: Sliding Weights' learning step (train.metal's kernels) diff --git a/zig/kernels/cuda/rank_sum.cu b/zig/kernels/cuda/rank_sum.cu new file mode 100644 index 0000000000..c8ed24cd6a --- /dev/null +++ b/zig/kernels/cuda/rank_sum.cu @@ -0,0 +1,11 @@ +// The rank-order sum of a reduce-scatter (rank_scatter.zig): rank 0's fp32 term first, then 1, 2, ... rounded each add. + +extern "C" __global__ void __launch_bounds__(256) + tf_rank_sum(const float* __restrict__ peers, const float* __restrict__ own, int world, int me, long long n, + float* __restrict__ out) { + for (long long i = blockIdx.x * 256LL + threadIdx.x; i < n; i += (long long)gridDim.x * 256) { + float acc = me == 0 ? own[i] : peers[i]; + for (int r = 1; r < world; ++r) acc = __fadd_rn(acc, r == me ? own[i] : peers[(long long)r * n + i]); + out[i] = acc; + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 63201f48d3..436643ff5e 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -28,6 +28,7 @@ pub const fp8_experts: []const u8 = if (available) &Blob("fatbin_fp8_experts").b pub const exlr8: []const u8 = if (available) &Blob("fatbin_exlr8").bytes else &.{}; pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").bytes else &.{}; pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts").bytes else &.{}; +pub const rank_sum: []const u8 = if (available) &Blob("fatbin_rank_sum").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &.{}; pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..b0d527ce75 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -25,6 +25,8 @@ pub const Api = struct { ncclAllReduce: *const fn (D, D, usize, DataType, RedOp, Comm, abi.Stream) callconv(.c) R, ncclAllGather: *const fn (D, D, usize, DataType, Comm, abi.Stream) callconv(.c) R, ncclBroadcast: *const fn (D, D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclSend: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclRecv: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, ncclGroupStart: *const fn () callconv(.c) R, ncclGroupEnd: *const fn () callconv(.c) R, }; diff --git a/zig/src/cuda/rank_scatter.zig b/zig/src/cuda/rank_scatter.zig new file mode 100644 index 0000000000..b5a74cf329 --- /dev/null +++ b/zig/src/cuda/rank_scatter.zig @@ -0,0 +1,80 @@ +//! A reduce-scatter by rows in rank order: an all-gather plus a rank-order sum's bits (not ncclReduceScatter's order). +const std = @import("std"); +const nccl = @import("nccl.zig"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); + +const Stream = stream_.Stream; + +/// Rank `rank`'s rows [lo, hi): whole rows, the first `rows % world` ranks one more (canon.partition's equal cut). +pub fn share(rows: usize, world: usize, rank: usize) [2]usize { + const base = rows / world; + const extra = rows % world; + const lo = rank * base + @min(rank, extra); + return .{ lo, lo + base + @intFromBool(rank < extra) }; +} + +/// Bytes of the receive buffer `run` needs: a slot of the largest share for every rank. +pub fn scratchBytes(rows: usize, world: usize, width: usize) usize { + return world * ((rows + world - 1) / world) * width * 4; +} + +pub const RankScatter = struct { + mod: module.Module, + sum: module.Function, + + pub fn load(d: *const driver.Driver) !RankScatter { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try module.Module.load(d, kernels.rank_sum); + errdefer mod.unload(); + return .{ .mod = mod, .sum = try mod.function("tf_rank_sum") }; + } + + pub fn unload(k: *RankScatter) void { + k.mod.unload(); + } + + /// `partial` [rows][width] fp32 summed over the ranks into `out` [own rows][width]: a send/recv group, then a sum. + pub fn run(k: *const RankScatter, lib: *const nccl.Library, comm: nccl.Comm, world: usize, rank: usize, s: Stream, partial: u64, rows: usize, width: usize, scratch: u64, out: u64) !void { + const mine = share(rows, world, rank); + const n = (mine[1] - mine[0]) * width; + try lib.check(lib.api.ncclGroupStart(), "ncclGroupStart"); + for (0..world) |peer| { + if (peer == rank) continue; + const theirs = share(rows, world, peer); + const count = (theirs[1] - theirs[0]) * width; + if (count > 0) try lib.check(lib.api.ncclSend(partial + theirs[0] * width * 4, count, .f32, @intCast(peer), comm, s.handle), "ncclSend"); + if (n > 0) try lib.check(lib.api.ncclRecv(scratch + peer * n * 4, n, .f32, @intCast(peer), comm, s.handle), "ncclRecv"); + } + try lib.check(lib.api.ncclGroupEnd(), "ncclGroupEnd"); + if (n == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ scratch, partial + mine[0] * width * 4 }) |v| a.add(v); + a.add(@as(c_int, @intCast(world))); + a.add(@as(c_int, @intCast(rank))); + a.add(@as(i64, @intCast(n))); + a.add(out); + const blocks: u32 = @intCast(@min((n + 255) / 256, 4096)); + try launch_.launch(k.sum, .{ .grid = .{ .x = blocks, .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &a); + } +}; + +test "shares cover the rows in rank order, whole rows, the first rows % world ranks one more" { + for ([_]usize{ 0, 1, 4, 64, 2047, 2048 }) |rows| for (1..6) |world| { + var at: usize = 0; + for (0..world) |r| { + const s = share(rows, world, r); + try std.testing.expectEqual(at, s[0]); + try std.testing.expectEqual(rows / world + @intFromBool(r < rows % world), s[1] - s[0]); + at = s[1]; + } + try std.testing.expectEqual(rows, at); + }; + try std.testing.expectEqual([2]usize{ 0, 683 }, share(2048, 3, 0)); + try std.testing.expectEqual([2]usize{ 1366, 2048 }, share(2048, 3, 2)); + try std.testing.expectEqual([2]usize{ 683, 1365 }, share(2047, 3, 1)); + try std.testing.expectEqual(@as(usize, 3 * 683 * 4096 * 4), scratchBytes(2048, 3, 4096)); +} diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 6ad88717e7..2c2362db9f 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -33,6 +33,7 @@ pub const grouped = @import("grouped.zig"); pub const qlinear = @import("qlinear.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); +pub const rank_scatter = @import("rank_scatter.zig"); pub const modelopt = @import("modelopt.zig"); test { @@ -41,6 +42,7 @@ test { _ = @import("qlinear.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); + _ = @import("rank_scatter.zig"); _ = launch; _ = abi; _ = aot; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 7ab32f226c..ec5c2bd91c 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -7,6 +7,7 @@ const runtime_tests = @import("runtime_tests.zig"); const bench = @import("bench.zig"); const oracle_tests = @import("oracle_tests.zig"); const libs_tests = @import("libs_tests.zig"); +const rank_scatter_tests = @import("rank_scatter_tests.zig"); const sample_tests = @import("sample_tests.zig"); const glue_tests = @import("glue_tests.zig"); const window_profile = @import("window_profile.zig"); @@ -37,6 +38,7 @@ const usage = \\ symbols every gdn fatbin instantiation resolves by its listed symbol \\ cublaslt bf16 GEMM through cuBLASLt against an fp64 reference \\ nccl one-rank NCCL all-reduce and all-gather + \\ rank-scatter [id] rank-order reduce-scatter (NCCL_COMM_ID or an id file) against an all-gather \\ gdn-replay replay_kernel bits against the Python oracle's fixture \\ gdn-tree tree_kernel bits against the Python oracle's fixture \\ triton a Triton cubin's bits against the Python oracle's fixture @@ -105,6 +107,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "symbols")) return runtime_tests.symbols(gpu); if (std.mem.eql(u8, cmd, "cublaslt")) return libs_tests.cublaslt(gpu); if (std.mem.eql(u8, cmd, "nccl")) return libs_tests.nccl(gpu); + if (std.mem.eql(u8, cmd, "rank-scatter")) return rank_scatter_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10), try std.fmt.parseInt(usize, try arg(rest, 1), 10), if (rest.len > 2) rest[2] else null); if (std.mem.eql(u8, cmd, "gdn-replay")) return oracle_tests.gdnReplay(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "gdn-tree")) return oracle_tests.gdnTree(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "triton")) return oracle_tests.tritonKernel(gpu, try arg(rest, 0)); diff --git a/zig/tests/cuda/rank_scatter_tests.zig b/zig/tests/cuda/rank_scatter_tests.zig new file mode 100644 index 0000000000..9c432318d8 --- /dev/null +++ b/zig/tests/cuda/rank_scatter_tests.zig @@ -0,0 +1,125 @@ +//! The rank-order reduce-scatter on real ranks: each rank's rows against an all-gather plus a host rank-order sum. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +const rs = cuda.rank_scatter; +const width = 4096; +const rows_cases = [_]usize{ 1, 4, 64, 2047, 2048 }; +const reps = 20; + +/// Rank `rank`'s partial rows for a case: normal values scaled over nine binades, so the adds round. +fn partial(gpa: std.mem.Allocator, seed: u64, rank: usize, rows: usize) ![]f32 { + var rng = std.Random.DefaultPrng.init(seed *% 1000003 +% rank *% 7919 +% rows); + const out = try gpa.alloc(f32, rows * width); + for (out, 0..) |*v, i| v.* = rng.random().floatNorm(f32) * std.math.ldexp(@as(f32, 1.0), @as(i32, @intCast(i % 9)) - 4); + return out; +} + +/// A communicator: one rank alone, or `world` ranks on NCCL_COMM_ID or on rank 0's unique id in `id_path` (as hex). +fn comm(lib: *const cuda.nccl.Library, gpu: Gpu, rank: usize, world: usize, id_path: ?[]const u8) !cuda.nccl.Comm { + var c: cuda.nccl.Comm = null; + if (world == 1) { + const devices = [1]c_int{gpu.ctx.device}; + try lib.check(lib.api.ncclCommInitAll(@ptrCast(&c), 1, &devices), "ncclCommInitAll"); + return c; + } + var id: cuda.nccl.UniqueId = undefined; + if (id_path) |path| { + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + const hex = std.fmt.bytesToHex(id.internal, .lower); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = &hex }); + } else { + var hex: [256]u8 = undefined; + const got = try std.Io.Dir.cwd().readFile(gpu.io, path, &hex); + try check.expect(got.len == 256, "unique id file {s}: {d} hex digits", .{ path, got.len }); + _ = try std.fmt.hexToBytes(&id.internal, &hex); + } + } else { + try check.expect(std.c.getenv("NCCL_COMM_ID") != null, "rank-scatter: NCCL_COMM_ID or an id file", .{}); + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + } + try lib.check(lib.api.ncclCommInitRank(&c, @intCast(world), id, @intCast(rank)), "ncclCommInitRank"); + return c; +} + +pub fn run(gpu: Gpu, rank: usize, world: usize, id_path: ?[]const u8) !void { + const a = gpu.gpa; + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + const c = try comm(&lib, gpu, rank, world, id_path); + defer _ = lib.api.ncclCommDestroy(c); + var k = try rs.RankScatter.load(gpu.d); + defer k.unload(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + for (rows_cases) |rows| for ([_]u64{ 1, 2, 3 }) |seed| { + const mine = rs.share(rows, world, rank); + const own = mine[1] - mine[0]; + const host = try partial(a, seed, rank, rows); + defer a.free(host); + var dp = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer dp.free(); + var scratch = try cuda.DeviceBuffer.alloc(gpu.d, @max(rs.scratchBytes(rows, world, width), 4)); + defer scratch.free(); + var out = try cuda.DeviceBuffer.alloc(gpu.d, @max(own * width * 4, 4)); + defer out.free(); + var gathered = try cuda.DeviceBuffer.alloc(gpu.d, world * rows * width * 4); + defer gathered.free(); + try k.run(&lib, c, world, rank, s, dp.ptr, rows, width, scratch.ptr, out.ptr); + try lib.check(lib.api.ncclAllGather(dp.ptr, gathered.ptr, rows * width, .f32, c, s.handle), "ncclAllGather"); + try s.synchronize(); + const g = try check.download(gpu, gathered); + defer a.free(g); + const gf = std.mem.bytesAsSlice(f32, @as([]align(4) u8, @alignCast(g))); + const want = try a.alloc(f32, own * width); + defer a.free(want); + for (want, 0..) |*w, i| { + const at = mine[0] * width + i; + var acc = gf[at]; + for (1..world) |r| acc += gf[r * rows * width + at]; + w.* = acc; + } + if (own > 0) { + const got = try check.download(gpu, out); + defer a.free(got); + try check.sameBytes("rank-order reduce-scatter", got[0 .. own * width * 4], std.mem.sliceAsBytes(want)); + } + }; + check.pass("rank-scatter rank {d} of {d}: rows {any} x {d} (seeds 1-3), every slice byte-equal to an all-gather plus a rank-order sum", .{ rank, world, rows_cases, width }); + try timing(gpu, &lib, c, &k, s, rank, world); +} + +/// Median ms of the reduce-scatter (sends, receives and the sum) and of an all-gather of the same partials, 2,048 rows. +fn timing(gpu: Gpu, lib: *const cuda.nccl.Library, c: cuda.nccl.Comm, k: *const rs.RankScatter, s: cuda.Stream, rank: usize, world: usize) !void { + const rows = 2048; + const host = try partial(gpu.gpa, 9, rank, rows); + defer gpu.gpa.free(host); + var dp = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer dp.free(); + var scratch = try cuda.DeviceBuffer.alloc(gpu.d, rs.scratchBytes(rows, world, width)); + defer scratch.free(); + var out = try cuda.DeviceBuffer.alloc(gpu.d, rows * width * 4); + defer out.free(); + var gathered = try cuda.DeviceBuffer.alloc(gpu.d, world * rows * width * 4); + defer gathered.free(); + var t0 = try cuda.Event.init(gpu.d, true); + defer t0.deinit(); + var t1 = try cuda.Event.init(gpu.d, true); + defer t1.deinit(); + var ms: [2][reps]f32 = undefined; + for (0..reps + 3) |i| for (0..2) |which| { + try t0.record(s); + if (which == 0) + try k.run(lib, c, world, rank, s, dp.ptr, rows, width, scratch.ptr, out.ptr) + else + try lib.check(lib.api.ncclAllGather(dp.ptr, gathered.ptr, rows * width, .f32, c, s.handle), "ncclAllGather"); + try t1.record(s); + try t1.synchronize(); + if (i >= 3) ms[which][i - 3] = try cuda.Event.elapsedMs(t0, t1); + }; + for (&ms) |*m| std.mem.sort(f32, m, {}, std.sort.asc(f32)); + check.pass("rank-scatter rank {d} of {d}, {d} x {d} fp32: reduce-scatter {d:.3} ms, all-gather alone {d:.3} ms (medians of {d})", .{ rank, world, rows, width, ms[0][reps / 2], ms[1][reps / 2], reps }); +}