cuda: a reduce-scatter in rank order over NCCL send and receive

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)

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;
+  }
+}

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 &.{};

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,
 };

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));
+}

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;

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 <rank> <n> [id]  rank-order reduce-scatter (NCCL_COMM_ID or an id file) against an all-gather
     \\  gdn-replay <dir>          replay_kernel bits against the Python oracle's fixture
     \\  gdn-tree <dir>            tree_kernel bits against the Python oracle's fixture
     \\  triton <dir>              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));

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 });
+}