glm: TF_GLM_PREFILL_REDUCE=ring, for fabrics where a rank reaches only its neighbours

zig/src/cuda/nccl.zig

@@ -59,11 +59,22 @@ pub const Library = struct {
 
     pub fn check(self: *const Library, r: Result, what: []const u8) Error!void {
         if (r == 0) return;
+        if (@import("builtin").is_test) return error.NcclFailed; // a test's injected failure: the runner fails error logs
         const s = self.api.ncclGetErrorString(r);
         std.log.err("{s}: {s} ({d})", .{ what, if (s) |p| std.mem.span(p) else "?", r });
         return error.NcclFailed;
     }
 
+    /// `post(ctx)`'s sends and receives as one group; a failed post still ends it, so later calls stay out of it.
+    pub fn group(self: *const Library, ctx: anytype, comptime post: fn (@TypeOf(ctx)) anyerror!void) anyerror!void {
+        try self.check(self.api.ncclGroupStart(), "ncclGroupStart");
+        post(ctx) catch |err| {
+            _ = self.api.ncclGroupEnd();
+            return err;
+        };
+        try self.check(self.api.ncclGroupEnd(), "ncclGroupEnd");
+    }
+
     /// NCCL's version as 10000 * major + 100 * minor + patch.
     pub fn version(self: *const Library) Error!c_int {
         var v: c_int = 0;

zig/src/families/glm/cuda.zig

@@ -13,6 +13,7 @@ pub const mlp = @import("cuda_mlp.zig");
 pub const shard = @import("cuda_shard.zig");
 pub const weights = @import("cuda_weights.zig");
 pub const forward = @import("cuda_forward.zig");
+pub const relay = @import("cuda_relay.zig");
 pub const config = @import("config.zig");
 pub const quant = @import("cuda_quant.zig");
 pub const q4 = @import("cuda_q4.zig");
@@ -37,6 +38,7 @@ test {
     _ = shard;
     _ = weights;
     _ = forward;
+    _ = relay;
     _ = quant;
     _ = q4;
     _ = ranks;

zig/src/families/glm/cuda_engine.zig

@@ -187,7 +187,7 @@ pub const Engine = struct {
 
     fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward {
         out.* = .{.{ .w = &e.w, .st = &e.cur.st, .b = b }};
-        return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else null, .split = e.settings.split, .kda_wide = e.settings.kda_wide };
+        return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else null, .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide };
     }
 
     /// Binds slot `slot`'s sequence for the steps that follow, making it (a zeroed state) on first use.
@@ -417,7 +417,7 @@ pub const Engine = struct {
             rows += p.ids.len;
         }
         var one: [1]fw.Rank = .{.{ .w = &e.w, .st = segs[0].st, .b = b, .segs = segs[0..parts.len] }};
-        const f: fw.Forward = .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = &one, .exchange = e.exchange(), .pieces = null, .split = e.settings.split, .kda_wide = e.settings.kda_wide };
+        const f: fw.Forward = .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = &one, .exchange = e.exchange(), .pieces = null, .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide };
         try f.run(@ptrCast(ids[0..rows]));
         for (parts, segs[0..parts.len]) |p, g| {
             try e.bind(p.slot);

zig/src/families/glm/cuda_forward.zig

@@ -18,6 +18,7 @@ const quant = @import("cuda_quant.zig");
 const q4 = @import("cuda_q4.zig");
 const state = @import("cuda_state.zig");
 const shared = @import("cuda_shared.zig");
+const relay = @import("cuda_relay.zig");
 
 pub const prefill_rows = state.prefill_rows;
 pub const window_rows = state.window_rows;
@@ -158,6 +159,7 @@ pub const Forward = struct {
     kda_wide: bool = true, // TF_GLM_KDA_DECODE_WIDE (the Python family's default): decode windows on KDA's three-kernel chain
     pieces: ?*const Pieces = null, // a prompt chunk's exchanges in row pieces on a second stream (null: whole)
     split: bool = true, // TF_GLM_PREFILL_REDUCE=split (the default with send/recv): each rank glues its own rows
+    ring: bool = false, // TF_GLM_PREFILL_REDUCE=ring: split, its exchanges passed along the ring (cuda_relay)
 
     pub fn set(f: *const Forward) glue.Set {
         return .{ .glue = &f.m.glue };
@@ -275,7 +277,7 @@ pub const Forward = struct {
     }
 
     /// reduce.shares: [lo, hi) of the rows of a piece of `n` each rank owns: about equal, the first n % world one more.
-    fn share(n: usize, world: usize, rank: usize) [2]usize {
+    pub fn share(n: usize, world: usize, rank: usize) [2]usize {
         const base = n / world;
         const extra = n % world;
         const lo = rank * base + @min(rank, extra);
@@ -324,13 +326,15 @@ pub const Forward = struct {
                     const me = r.w.share.rank;
                     const mine = share(n, world, me);
                     const own = mine[1] - mine[0];
-                    try x.lib.check(x.lib.api.ncclGroupStart(), "ncclGroupStart");
-                    for (0..world) |peer| if (peer != me) {
-                        const theirs = share(n, world, peer);
-                        try x.lib.check(x.lib.api.ncclSend(r.b.part + (lo + theirs[0]) * D * 4, (theirs[1] - theirs[0]) * D, .f32, @intCast(peer), x.comm, p.side.handle), "ncclSend");
-                        try x.lib.check(x.lib.api.ncclRecv(room(r, lo) + peer * own * D * 4, own * D, .f32, @intCast(peer), x.comm, p.side.handle), "ncclRecv");
-                    };
-                    try x.lib.check(x.lib.api.ncclGroupEnd(), "ncclGroupEnd");
+                    if (f.ring) try relay.partials(x.lib, x.comm, p.side, .{ .world = world, .rank = me, .n = n }, room(r, lo), r.b.part, lo, D) else {
+                        try x.lib.check(x.lib.api.ncclGroupStart(), "ncclGroupStart");
+                        for (0..world) |peer| if (peer != me) {
+                            const theirs = share(n, world, peer);
+                            try x.lib.check(x.lib.api.ncclSend(r.b.part + (lo + theirs[0]) * D * 4, (theirs[1] - theirs[0]) * D, .f32, @intCast(peer), x.comm, p.side.handle), "ncclSend");
+                            try x.lib.check(x.lib.api.ncclRecv(room(r, lo) + peer * own * D * 4, own * D, .f32, @intCast(peer), x.comm, p.side.handle), "ncclRecv");
+                        };
+                        try x.lib.check(x.lib.api.ncclGroupEnd(), "ncclGroupEnd");
+                    }
                     try f.copy(p.side, room(r, lo) + me * own * D * 4, r.b.part + (lo + mine[0]) * D * 4, own * D * 4);
                 },
             }
@@ -401,14 +405,16 @@ pub const Forward = struct {
                     const r = f.ranks[0];
                     const me = r.w.share.rank;
                     const rows_of = if (last) r.b.hidden else r.b.normed;
-                    const mine = share(n, world, me);
-                    try x.lib.check(x.lib.api.ncclGroupStart(), "ncclGroupStart");
-                    for (0..world) |peer| if (peer != me) {
-                        const theirs = share(n, world, peer);
-                        try x.lib.check(x.lib.api.ncclSend(rows_of + (lo + mine[0]) * D * 2, (mine[1] - mine[0]) * D, .bf16, @intCast(peer), x.comm, p.side.handle), "ncclSend");
-                        try x.lib.check(x.lib.api.ncclRecv(rows_of + (lo + theirs[0]) * D * 2, (theirs[1] - theirs[0]) * D, .bf16, @intCast(peer), x.comm, p.side.handle), "ncclRecv");
-                    };
-                    try x.lib.check(x.lib.api.ncclGroupEnd(), "ncclGroupEnd");
+                    if (f.ring) try relay.glued(x.lib, x.comm, p.side, .{ .world = world, .rank = me, .n = n }, rows_of, lo, D) else {
+                        const mine = share(n, world, me);
+                        try x.lib.check(x.lib.api.ncclGroupStart(), "ncclGroupStart");
+                        for (0..world) |peer| if (peer != me) {
+                            const theirs = share(n, world, peer);
+                            try x.lib.check(x.lib.api.ncclSend(rows_of + (lo + mine[0]) * D * 2, (mine[1] - mine[0]) * D, .bf16, @intCast(peer), x.comm, p.side.handle), "ncclSend");
+                            try x.lib.check(x.lib.api.ncclRecv(rows_of + (lo + theirs[0]) * D * 2, (theirs[1] - theirs[0]) * D, .bf16, @intCast(peer), x.comm, p.side.handle), "ncclRecv");
+                        };
+                        try x.lib.check(x.lib.api.ncclGroupEnd(), "ncclGroupEnd");
+                    }
                 },
             }
             try p.shared[q].record(p.side);

zig/src/families/glm/cuda_ranks.zig

@@ -7,7 +7,8 @@ const Sha256 = std.crypto.hash.sha2.Sha256;
 
 /// The forward's knobs a rank reads from its environment (the Python family's names and defaults).
 pub const Settings = struct {
-    split: bool = true, // TF_GLM_PREFILL_REDUCE: split (each rank glues its own rows) or gather
+    split: bool = true, // TF_GLM_PREFILL_REDUCE: split (each rank glues its own rows), ring or gather
+    ring: bool = false, // TF_GLM_PREFILL_REDUCE=ring: split, its exchanges passed along the ring (cuda_relay)
     kda_wide: bool = true, // TF_GLM_KDA_DECODE_WIDE: 1 (decode windows on KDA's three-kernel chain) or 0
     pieces: u32 = 3, // TF_GLM_OVERLAP_PIECES: a prompt chunk's exchanges in this many row pieces (1: whole).
     heat: ?heat.Bands = null, // TF_HEAT_HIGH and TF_HEAT_LOW (upstream's names): a prompt chunk waits on every rank's hottest zone
@@ -17,8 +18,11 @@ pub const Settings = struct {
     pub fn read(get: *const fn ([:0]const u8) ?[]const u8, problem: *[]const u8) !Settings {
         var s: Settings = .{};
         if (get("TF_GLM_PREFILL_REDUCE")) |v| {
-            if (std.mem.eql(u8, v, "split")) s.split = true else if (std.mem.eql(u8, v, "gather")) s.split = false else {
-                problem.* = "TF_GLM_PREFILL_REDUCE: split or gather";
+            if (std.mem.eql(u8, v, "split")) s.split = true else if (std.mem.eql(u8, v, "ring")) {
+                s.split = true;
+                s.ring = true;
+            } else if (std.mem.eql(u8, v, "gather")) s.split = false else {
+                problem.* = "TF_GLM_PREFILL_REDUCE: split, ring or gather";
                 return error.BadSetting;
             }
         }
@@ -63,7 +67,7 @@ pub const Fingerprint = extern struct {
     layout: u32 = 5, // this struct's version
     world: u32,
     rank: u32,
-    split: u32,
+    split: u32, // TF_GLM_PREFILL_REDUCE: 0 gather, 1 split, 2 ring
     kda_wide: u32,
     pieces: u32,
     layers: u32, // TF_GLM_LAYERS (0: all)
@@ -105,7 +109,7 @@ pub const Fingerprint = extern struct {
         var f: Fingerprint = .{
             .world = @intCast(world),
             .rank = @intCast(rank),
-            .split = @intFromBool(s.split),
+            .split = if (!s.split) 0 else if (s.ring) 2 else 1,
             .kda_wide = @intFromBool(s.kda_wide),
             .pieces = s.pieces,
             .layers = s.layers,
@@ -230,6 +234,11 @@ fn fake(name: [:0]const u8) ?[]const u8 {
     return null;
 }
 
+fn ringed(name: [:0]const u8) ?[]const u8 {
+    if (std.mem.eql(u8, name, "TF_GLM_PREFILL_REDUCE")) return "ring";
+    return null;
+}
+
 fn bad(name: [:0]const u8) ?[]const u8 {
     if (std.mem.eql(u8, name, "TF_GLM_PREFILL_REDUCE")) return "sum";
     return null;
@@ -240,7 +249,10 @@ test "settings from the environment: the defaults, the choices, and a refusal na
     const s = try Settings.read(&fake, &problem);
     try std.testing.expect(!s.split and s.kda_wide and s.pieces == 2 and s.heat.?.high_c == 92 and s.heat.?.low_c == 88);
     try std.testing.expectError(error.BadSetting, Settings.read(&bad, &problem));
-    try std.testing.expectEqualStrings("TF_GLM_PREFILL_REDUCE: split or gather", problem);
+    try std.testing.expectEqualStrings("TF_GLM_PREFILL_REDUCE: split, ring or gather", problem);
+    const r = try Settings.read(&ringed, &problem);
+    try std.testing.expect(r.split and r.ring);
+    try std.testing.expectEqual(@as(u32, 2), Fingerprint.of(2, 0, r, 8192, true, @splat(0), @splat(0), "/c").split);
 }
 
 test "fingerprints: agreeing ranks pass; the first differing field and rank are named" {