Flash Next CUDA: the hyper-connection normalize over 256-column slices

zig/kernels/cuda/flashnext_hc.cu

@@ -411,6 +411,23 @@ extern "C" __global__ void fn_hc_normalize(const __nv_bfloat16* H, const float*
     }
 }
 
+// fn_hc_normalize a block a (row, stream, 256-column slice): each block sums the stream the same way, so the same bits.
+extern "C" __global__ void __launch_bounds__(256) fn_hc_normalize_cols(const __nv_bfloat16* H, const float* PSS, const float* SCALE,
+                                                                       __nv_bfloat16* OUT, float eps, int D, int SS) {
+    __shared__ float rinv;
+    const int r = blockIdx.x, stream = blockIdx.y, nc = D / 256;
+    if (threadIdx.x == 0) {
+        float total = 0.0f;
+        for (int c = 0; c < nc; ++c) total = __fadd_rn(total, PSS[(static_cast<long long>(r) * nc + c) * SS + stream]);
+        rinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(total, static_cast<float>(D)), eps)));
+    }
+    __syncthreads();
+    const int d = blockIdx.z * 256 + threadIdx.x;
+    const int col = stream * D + d;
+    const long long at = static_cast<long long>(r) * SS * D + col;
+    OUT[at] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(bf(H[at]), rinv), SCALE[col]));
+}
+
 // Dense up-projection values are already BF16: sigmoid, product and stream mean keep their original rounding seams.
 extern "C" __global__ void fn_hc_dense_mix(const __nv_bfloat16* UP, const __nv_bfloat16* NORMED,
                                             __nv_bfloat16* MIXED, float* XSM, int D, int SS) {
@@ -433,6 +450,30 @@ extern "C" __global__ void fn_hc_dense_mix(const __nv_bfloat16* UP, const __nv_b
     }
 }
 
+// fn_hc_dense_mix, a warp a 32-column group and eight a block: the same ops, lane 0 summing its group in order.
+extern "C" __global__ void __launch_bounds__(256) fn_hc_dense_mix8(const __nv_bfloat16* UP, const __nv_bfloat16* NORMED,
+                                                                   __nv_bfloat16* MIXED, float* XSM, int D, int SS) {
+    __shared__ float row[8][32];
+    const int r = blockIdx.x, warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
+    const int group = blockIdx.y * 8 + warp, d = group * 32 + lane;
+    if (group >= D / 32) return;
+    float total = 0.0f;
+    for (int stream = 0; stream < SS; ++stream) {
+        const long long at = (static_cast<long long>(r) * SS + stream) * D + d;
+        const float value = rbf(__fmul_rn(bsig(bf(UP[at])), bf(NORMED[at])));
+        total = __fadd_rn(total, value);
+    }
+    const float mixed = rbf(__fdiv_rn(total, static_cast<float>(SS)));
+    MIXED[static_cast<long long>(r) * D + d] = __float2bfloat16_rn(mixed);
+    row[warp][lane] = mixed;
+    __syncwarp();
+    if (lane == 0) {
+        float sum = 0.0f;
+        for (int i = 0; i < 32; ++i) sum = __fadd_rn(sum, row[warp][i]);
+        XSM[static_cast<long long>(r) * (D / 32) + group] = sum;
+    }
+}
+
 // Shared expert uses its own BF16 projection outputs and rounding boundaries.
 extern "C" __global__ void fn_shared_swiglu(const __nv_bfloat16* GATE, const __nv_bfloat16* UP,
                                              __nv_bfloat16* ACT, int W) {

zig/src/core/contract/registry.zig

@@ -32,6 +32,8 @@ pub const cuda = [_]CudaRecord{
     .{ .launch = .{ .name = "flashnext_hc fn_hc_dense_mix", .symbol = "fn_hc_dense_mix", .image = "flashnext_hc", .threads = 32 }, .usage = .{ .{ .regs = 22, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 } } },
     .{ .launch = .{ .name = "flashnext_hc fn_hc_down", .symbol = "fn_hc_down", .image = "flashnext_hc", .threads = 128, .dynamic = 11136 }, .usage = .{ .{ .regs = 56, .static = 64, .bound = 128 }, .{ .regs = 54, .static = 64, .bound = 128 }, .{ .regs = 54, .static = 64, .bound = 128 }, .{ .regs = 60, .static = 64, .bound = 128 }, .{ .regs = 60, .static = 64, .bound = 128 } } },
     .{ .launch = .{ .name = "flashnext_hc fn_hc_normalize", .symbol = "fn_hc_normalize", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 }, .{ .regs = 38, .static = 16, .bound = 0 } } },
+    .{ .launch = .{ .name = "flashnext_hc fn_hc_dense_mix8", .symbol = "fn_hc_dense_mix8", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 1024, .bound = 256 }, .{ .regs = 34, .static = 1024, .bound = 256 }, .{ .regs = 34, .static = 1024, .bound = 256 }, .{ .regs = 35, .static = 1024, .bound = 256 }, .{ .regs = 34, .static = 1024, .bound = 256 } } },
+    .{ .launch = .{ .name = "flashnext_hc fn_hc_normalize_cols", .symbol = "fn_hc_normalize_cols", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 16, .bound = 256 }, .{ .regs = 40, .static = 16, .bound = 256 }, .{ .regs = 40, .static = 16, .bound = 256 }, .{ .regs = 40, .static = 16, .bound = 256 }, .{ .regs = 38, .static = 16, .bound = 256 } } },
     .{ .launch = .{ .name = "flashnext_hc fn_hc_reduce_act", .symbol = "fn_hc_reduce_act", .image = "flashnext_hc", .threads = 352 }, .usage = .{ .{ .regs = 32, .static = 4096, .bound = 1024 }, .{ .regs = 63, .static = 4096, .bound = 1024 }, .{ .regs = 63, .static = 4096, .bound = 1024 }, .{ .regs = 58, .static = 4096, .bound = 1024 }, .{ .regs = 58, .static = 4096, .bound = 1024 } } },
     .{ .launch = .{ .name = "flashnext_hc fn_hc_upmix", .symbol = "fn_hc_upmix", .image = "flashnext_hc", .threads = 256, .dynamic = 18688 }, .usage = .{ .{ .regs = 69, .static = 2048, .bound = 256 }, .{ .regs = 64, .static = 2048, .bound = 256 }, .{ .regs = 64, .static = 2048, .bound = 256 }, .{ .regs = 78, .static = 2048, .bound = 256 }, .{ .regs = 78, .static = 2048, .bound = 256 } } },
     .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback0", .symbol = "fn_hc_writeback0", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 256, .bound = 256 }, .{ .regs = 30, .static = 256, .bound = 256 }, .{ .regs = 30, .static = 256, .bound = 256 }, .{ .regs = 33, .static = 256, .bound = 256 }, .{ .regs = 37, .static = 256, .bound = 256 } } },

zig/src/families/flashnext/cuda_glue.zig

@@ -28,6 +28,8 @@ const Fns = struct {
     wb_moe: cuda.Function,
     wb_slices: cuda.Function,
     normalize: cuda.Function,
+    normalize_cols: cuda.Function, // fn_hc_normalize's bits over 256-column slices (dims a multiple of 256)
+    dense_mix8: cuda.Function, // fn_hc_dense_mix's bits, a warp a 32-column group, eight a block
     dense_mix: cuda.Function,
     shared_swiglu: cuda.Function,
     shared_scatter: cuda.Function,
@@ -99,6 +101,8 @@ pub const Glue = struct {
             .wb_moe = try g.hc.function("fn_hc_writeback_moe"),
             .wb_slices = try g.hc.function("fn_hc_writeback_slices"),
             .normalize = try g.hc.function("fn_hc_normalize"),
+            .normalize_cols = try g.hc.function("fn_hc_normalize_cols"),
+            .dense_mix8 = try g.hc.function("fn_hc_dense_mix8"),
             .dense_mix = try g.hc.function("fn_hc_dense_mix"),
             .shared_swiglu = try g.hc.function("fn_shared_swiglu"),
             .shared_scatter = try g.hc.function("fn_shared_scatter"),
@@ -228,6 +232,8 @@ pub const Glue = struct {
         args.add(eps);
         args.add(int(dims));
         args.add(int(streams));
+        // a block a 256-column slice: the same stream sum in each, so fn_hc_normalize's bits with more blocks
+        if (dims % 256 == 0) return cuda.launch.launch(g.f.normalize_cols, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(streams), .z = @intCast(dims / 256) }, .block = .{ .x = 256 } }, s, &args);
         try cuda.launch.launch(g.f.normalize, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(streams) }, .block = .{ .x = 256 } }, s, &args);
     }
 
@@ -236,7 +242,8 @@ pub const Glue = struct {
         for ([_]u64{ up, norm, mixed, sums }) |v| args.add(v);
         args.add(int(dims));
         args.add(int(streams));
-        try cuda.launch.launch(g.f.dense_mix, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(dims / 32) }, .block = .{ .x = 32 } }, s, &args);
+        // eight 32-column groups a block, a warp each: fn_hc_dense_mix's bits with an eighth of its blocks
+        try cuda.launch.launch(g.f.dense_mix8, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((dims / 32 + 7) / 8) }, .block = .{ .x = 256 } }, s, &args);
     }
 
     pub fn sharedSwiGLU(g: *const Glue, s: cuda.Stream, gate: u64, up: u64, act: u64, rows: usize, width: usize) !void {

zig/tests/cuda/contract_hc_mix.zig

@@ -0,0 +1,163 @@
+//! fn_hc_dense_mix8 under the contract: fn_hc_dense_mix's bits (rows and group sums); control: its inputs swapped.
+const std = @import("std");
+const cuda = @import("cuda");
+const core = @import("core");
+const rig_ = @import("contract_rig.zig");
+const c = core.contract.spec;
+const g = core.contract.gates;
+const contract = core.contract;
+
+pub const Kind = enum { mix8, control };
+
+const rule = "as fn_hc_dense_mix gives them: NaN and infinite up rows and normed rows";
+
+fn depth(s: c.Shape) usize {
+    return s.k + 8;
+}
+
+pub fn spec(kind: Kind) c.Spec {
+    return .{
+        .name = if (kind == .mix8) "cuda Flash Next fn_hc_dense_mix8 (replaces fn_hc_dense_mix)" else "control: fn_hc_dense_mix8 with its inputs swapped posing as fn_hc_dense_mix",
+        .accum = .f32,
+        .depth = depth,
+        .ops = 4,
+        .mids = 2,
+        .out = .bf16,
+        .replaces = "fn_hc_dense_mix",
+        .nan = .{ .stated = rule },
+        .tiles = .{ .{ .size = 1 }, .{ .size = 256, .step = 32, .min = 32 }, .{} },
+        .served = &.{ .{ .rows = 1, .n = 2560, .k = 4 }, .{ .rows = 17, .n = 2560, .k = 4 }, .{ .rows = 3, .n = 96, .k = 4 } },
+        // each value sums four bf16-rounded terms and is itself rounded to bf16: no accumulator's error shows at gate 4
+        .sums = .{ .f16 = 0, .bf16 = 0, .why = "four bf16-rounded terms a value, rounded to bf16 again" },
+    };
+}
+
+pub const Mix = struct {
+    rig: rig_.Rig,
+    kind: Kind,
+    new_fn: cuda.Function,
+    old_fn: cuda.Function,
+    arena: ?std.heap.ArenaAllocator = null,
+    shape: c.Shape = undefined,
+    up: []u16 = &.{},
+    normed: []u16 = &.{},
+    dup: u64 = 0,
+    dnormed: u64 = 0,
+    out: u64 = 0,
+    base: u64 = 0,
+    msg: [160]u8 = undefined,
+
+    pub fn deinit(self: *Mix) void {
+        if (self.arena) |*arena| arena.deinit();
+        self.rig.deinit();
+    }
+
+    const vt = g.hooks(Mix);
+
+    pub fn kernel(self: *Mix) g.Kernel {
+        return .{ .spec = spec(self.kind), .ctx = self, .vt = &vt, .control = if (self.kind == .control) .shapes else null };
+    }
+
+    fn outBytes(self: *const Mix) usize {
+        return self.shape.rows * self.shape.n * 2 + self.shape.rows * (self.shape.n / 32) * 4;
+    }
+
+    pub fn prepare(self: *Mix, shape: c.Shape, fill: g.Fill, seed: u64) !bool {
+        if (shape.rows == 0 or shape.rows > 64 or shape.n == 0 or shape.n % 32 != 0 or shape.k == 0 or shape.k > 8) return false;
+        self.rig.reset();
+        if (self.arena) |*arena| arena.deinit();
+        self.arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa);
+        const a = self.arena.?.allocator();
+        self.shape = shape;
+        const total = shape.rows * shape.k * shape.n;
+        self.up = try a.alloc(u16, total);
+        self.normed = try a.alloc(u16, total);
+        var prng = std.Random.DefaultPrng.init(seed);
+        const rng = prng.random();
+        for (self.up) |*v| v.* = c.toBf16(rng.float(f32) * 8 - 4);
+        for (self.normed) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2);
+        if (fill == .nonfinite) {
+            self.up[0] = c.toBf16(std.math.nan(f32));
+            self.normed[total - 1] = c.toBf16(std.math.inf(f32));
+            if (total > 40) self.normed[33] = c.toBf16(-std.math.inf(f32));
+        }
+        self.dup = try self.rig.input(std.mem.sliceAsBytes(self.up));
+        self.dnormed = try self.rig.input(std.mem.sliceAsBytes(self.normed));
+        self.out = try self.rig.output(self.outBytes());
+        self.base = try self.rig.scratch(self.outBytes());
+        return true;
+    }
+
+    fn launch(self: *Mix, new: bool, swap: bool, out: u64) !void {
+        var args: cuda.Args = .{};
+        for ([_]u64{ if (swap) self.dnormed else self.dup, if (swap) self.dup else self.dnormed, out, out + self.shape.rows * self.shape.n * 2 }) |v| args.add(v);
+        args.add(@as(c_int, @intCast(self.shape.n)));
+        args.add(@as(c_int, @intCast(self.shape.k)));
+        const groups = self.shape.n / 32;
+        const cfg: cuda.launch.Config = if (new) .{ .grid = .{ .x = @intCast(self.shape.rows), .y = @intCast((groups + 7) / 8) }, .block = .{ .x = 256 } } else .{ .grid = .{ .x = @intCast(self.shape.rows), .y = @intCast(groups) }, .block = .{ .x = 32 } };
+        try cuda.launch.launch(if (new) self.new_fn else self.old_fn, cfg, self.rig.s, &args);
+        try self.rig.s.synchronize();
+    }
+
+    pub fn run(self: *Mix) !void {
+        try self.launch(true, self.kind == .control, self.out);
+    }
+
+    pub fn replaced(self: *Mix) ![]const u8 {
+        try self.launch(false, false, self.base);
+        return self.rig.bytes(self.base, self.outBytes());
+    }
+
+    pub fn check(self: *Mix, _: g.Gate) !?[]const u8 {
+        const a = self.rig.gpu.gpa;
+        const got = try a.dupe(u8, try self.rig.outputBytes());
+        defer a.free(got);
+        const want = try self.replaced();
+        if (std.mem.indexOfDiff(u8, got, want)) |at| return try std.fmt.bufPrint(&self.msg, "byte {d} of {d} differs from fn_hc_dense_mix's", .{ at, got.len });
+        for (0..self.shape.rows * self.shape.n) |i| if (c.classOf(c.read(.bf16, got, i)) != .finite) return null;
+        return try std.fmt.bufPrint(&self.msg, "no output is non-finite: the planted values prove nothing", .{});
+    }
+
+    pub fn count(self: *Mix) usize {
+        return self.shape.rows * self.shape.n;
+    }
+
+    fn values(self: *const Mix, acc: ?c.Accum, order: g.Order, out: []f64, mag: ?[]f64) void {
+        const n = self.shape.n;
+        const ss = self.shape.k;
+        var t: [8]f64 = undefined;
+        for (0..self.shape.rows) |r| for (0..n) |d| {
+            for (0..ss) |st| {
+                const at = (r * ss + st) * n + d;
+                const u: f64 = c.fromBf16(self.up[at]);
+                const sig = 1.0 / (1.0 + @exp(-u));
+                t[st] = sig * @as(f64, c.fromBf16(self.normed[at]));
+            }
+            const total = if (acc) |ac| contract.rows.sumIn(ac, order, t[0..ss]) else blk: {
+                var x: f64 = 0;
+                for (t[0..ss]) |v| x += v;
+                break :blk x;
+            };
+            out[r * n + d] = total / @as(f64, @floatFromInt(ss));
+            if (mag) |m| {
+                var x: f64 = 0;
+                for (t[0..ss]) |v| x += @abs(v);
+                m[r * n + d] = x / @as(f64, @floatFromInt(ss));
+            }
+        };
+    }
+
+    /// The bf16 roundings of each sigmoid and term (half an ulp each, a bit more) and of the mean, on the sums' bound.
+    pub fn bounds(_: *Mix, want: []const f64, mag: []const f64, allow: []f64) !void {
+        const half = std.math.ldexp(@as(f64, 1), -8);
+        for (allow, want, mag) |*l, w, m| l.* += 2.01 * half * m + half * @abs(w) + std.math.floatMin(f32);
+    }
+
+    pub fn reference(self: *Mix, want: []f64, mag: []f64) !void {
+        self.values(null, .sequential, want, mag);
+    }
+
+    pub fn simulate(self: *Mix, acc: c.Accum, order: g.Order, out: []f64) !void {
+        self.values(acc, order, out, null);
+    }
+};