radio CUDA: fuse the tower's bias, activation and residual launches

docs/kernel-contract.md

@@ -111,8 +111,9 @@ Mamba scan of a verify window (its host twin's bytes), qmm_group's cluster tile,
 bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out, qmmf's FP4 tiles with fp32 out (K slices met in a
 cluster, fused, or by `reduce_kernel<true>`), the checkpoint math's FP8 128-row prompt tiles (both slice meetings,
 bit-equal to the 64-row tiles) and NVFP4 experts: the decode-sized forms and the prompt tiles (up to down's NVFP4
-rows, fp32 down and the down summing two width shares, each bit-equal to the kernels it replaces; sm_120 up), and the
-reduce-scatter's `tf_rank_sum`.
+rows, fp32 down and the down summing two width shares, each bit-equal to the kernels it replaces; sm_120 up), the
+reduce-scatter's `tf_rank_sum`, and RADIO's fused bias, activation and residual (`tf_vit_bias_act`, bit-equal to
+`tf_vit_bias` then `tf_vit_pointwise`).
 
 ## CUDA records per nvcc release
 

zig/kernels/cuda/vit.cu

@@ -38,6 +38,14 @@ __device__ __forceinline__ float gelu_erf(float x) {
     return __fmul_rn(__fmul_rn(x, 0.5f), __fadd_rn(1.0f, erff(__fmul_rn(x, alpha))));
 }
 
+// bf16(act(v)), then bf16(R[at] + v) when R; v is already rounded to bf16.
+__device__ __forceinline__ __nv_bfloat16 pointwise_at(float v, const __nv_bfloat16* R, int64_t at, int act) {
+    if (act == 1) v = rbf(gelu_tanh(v));
+    if (act == 2) v = rbf(gelu_erf(v));
+    if (act == 3) v = rbf(v > 0.0f ? __fmul_rn(v, v) : 0.0f);
+    return __float2bfloat16_rn(R ? __fadd_rn(bf(R[at]), v) : v);
+}
+
 }  // namespace
 
 // Block a row: y = bf16((x - mean) * rstd * w + b), mean and variance in two fixed-order passes.
@@ -86,11 +94,16 @@ extern "C" __global__ void __launch_bounds__(NT) tf_vit_pointwise(const __nv_bfl
                                                                       int64_t count, int act) {
     const int64_t at = static_cast<int64_t>(blockIdx.x) * NT + threadIdx.x;
     if (at >= count) return;
-    float v = bf(X[at]);
-    if (act == 1) v = rbf(gelu_tanh(v));
-    if (act == 2) v = rbf(gelu_erf(v));
-    if (act == 3) v = rbf(v > 0.0f ? __fmul_rn(v, v) : 0.0f);
-    O[at] = __float2bfloat16_rn(R ? __fadd_rn(bf(R[at]), v) : v);
+    O[at] = pointwise_at(bf(X[at]), R, at, act);
+}
+
+// Block (row, column block): bias then pointwise, keeping the intermediate bf16 rounding in a register.
+extern "C" __global__ void __launch_bounds__(NT) tf_vit_bias_act(const float* __restrict__ ACC, const __nv_bfloat16* __restrict__ BIAS,
+                                                                 const __nv_bfloat16* R, __nv_bfloat16* O, int N, int act) {
+    const int c = blockIdx.y * NT + threadIdx.x;
+    if (c >= N) return;
+    const int64_t at = static_cast<int64_t>(blockIdx.x) * N + c;
+    O[at] = pointwise_at(rbf(__fadd_rn(ACC[at], bf(BIAS[c]))), R, at, act);
 }
 
 // Block an output row (i, j) of an h x w grid: rows (2i, 2j), (2i, 2j+1), (2i+1, 2j), (2i+1, 2j+1) side by side.

zig/src/core/contract.zig

@@ -13,6 +13,7 @@ pub const cubin = @import("contract/cubin.zig");
 test {
     _ = @import("contract/qwen_served_records.zig");
     _ = @import("contract/delta_slots_records.zig");
+    _ = @import("contract/vit_records.zig");
     _ = spec;
     _ = gates;
     _ = rows;

zig/src/core/contract/registry.zig

@@ -118,6 +118,7 @@ pub const cuda = [_]CudaRecord{
     .{ .launch = .{ .name = "vit tf_vit_layer_norm", .symbol = "tf_vit_layer_norm", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 38, .static = 32, .bound = 256 }, .{ .regs = 38, .static = 32, .bound = 256 }, .{ .regs = 38, .static = 32, .bound = 256 }, .{ .regs = 40, .static = 32, .bound = 256 }, .{ .regs = 40, .static = 32, .bound = 256 } } },
     .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } },
     .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } },
+    @import("vit_records.zig").record,
     .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } },
     .{ .launch = .{ .name = "qwen_vision tf_qv_embed", .symbol = "tf_qv_embed", .image = "qwen_vision", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 33, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } },
     .{ .launch = .{ .name = "qwen_vision tf_qv_layernorm", .symbol = "tf_qv_layernorm", .image = "qwen_vision", .threads = 256 }, .usage = .{ .{ .regs = 23, .static = 132, .bound = 256 }, .{ .regs = 26, .static = 132, .bound = 256 }, .{ .regs = 26, .static = 132, .bound = 256 }, .{ .regs = 23, .static = 132, .bound = 256 }, .{ .regs = 23, .static = 132, .bound = 256 } } },

zig/src/core/contract/vit_records.zig

@@ -0,0 +1,18 @@
+//! RADIO's fused epilogue launch, measured with the release fatbins' nvcc.
+const std = @import("std");
+const registry = @import("registry.zig");
+
+pub const record: registry.CudaRecord = .{
+    .launch = .{ .name = "vit tf_vit_bias_act", .symbol = "tf_vit_bias_act", .image = "vit", .threads = 256 },
+    .usage = .{ .{ .regs = 15, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } },
+};
+
+test "RADIO fused epilogue is in the live registry and fits every SM" {
+    var found: usize = 0;
+    for (registry.all_cuda) |active| if (std.mem.eql(u8, active.launch.symbol, record.launch.symbol)) {
+        try std.testing.expectEqualDeep(record, active);
+        found += 1;
+    };
+    try std.testing.expectEqual(@as(usize, 1), found);
+    try std.testing.expectEqual(@as(usize, 0), try registry.checkCuda(std.testing.allocator, &.{record}, true));
+}

zig/src/cuda/vit.zig

@@ -26,6 +26,7 @@ pub const Kernels = struct {
     layer_norm_fn: module.Function,
     pointwise_fn: module.Function,
     bias_fn: module.Function,
+    bias_act_fn: module.Function,
     unshuffle_fn: module.Function,
     cast_fn: module.Function,
     position_fn: module.Function,
@@ -39,6 +40,7 @@ pub const Kernels = struct {
             .layer_norm_fn = try mod.function("tf_vit_layer_norm"),
             .pointwise_fn = try mod.function("tf_vit_pointwise"),
             .bias_fn = try mod.function("tf_vit_bias"),
+            .bias_act_fn = try mod.function("tf_vit_bias_act"),
             .unshuffle_fn = try mod.function("tf_vit_unshuffle"),
             .cast_fn = try mod.function("tf_vit_cast"),
             .position_fn = try mod.function("tf_vit_position_add"),
@@ -76,6 +78,15 @@ pub const Kernels = struct {
         try go(k.bias_fn, s, rows, blocks(n, threads), &a);
     }
 
+    /// o [rows, n] = bf16(act(bf16(acc + b)) + r), acc f32; r may be o.
+    pub fn biasAct(k: *const Kernels, s: Stream, acc: u64, b: u64, r: ?u64, o: u64, rows: usize, n: usize, act: Act) !void {
+        var a: Args = .{};
+        for ([_]u64{ acc, b, r orelse 0, o }) |p| a.add(p);
+        a.add(int(n));
+        a.add(@backingInt(act));
+        try go(k.bias_act_fn, s, rows, blocks(n, threads), &a);
+    }
+
     /// y [h/2 * w/2, 4c] = the projector's pixel unshuffle of x [h * w, c], an even grid.
     pub fn unshuffle(k: *const Kernels, s: Stream, x: u64, y: u64, h: usize, w: usize, c: usize) !void {
         if (h % 2 != 0 or w % 2 != 0) return error.RowShape;
@@ -168,6 +179,17 @@ pub fn linear(g: *Gemms, k: *const Kernels, s: Stream, acc: u64, x: u64, w: u64,
     try k.bias(s, acc, b, y, rows, n);
 }
 
+/// out = bf16(act(bf16(x . w^T + b)) + r); y holds the intermediate only when the GEMM fuses bias.
+pub fn linearAct(g: *Gemms, k: *const Kernels, s: Stream, acc: u64, x: u64, w: u64, b: u64, y: u64, rows: usize, n: usize, kk: usize, r: ?u64, out: u64, act: Act) !void {
+    const p = try g.plan(rows, n, kk, .bf16, .bias);
+    if (p.epilogue == .bias) {
+        try p.runBias(x, w, b, y, g.workspace.ptr, workspace_bytes, s.handle);
+        return k.pointwise(s, y, r, out, rows * n, act);
+    }
+    try g.run(s, x, w, acc, rows, n, kk, .f32);
+    try k.biasAct(s, acc, b, r, out, rows, n, act);
+}
+
 /// prefill_attention.cu's full_kernel<D, WARPS, slots>, or vision_kernel<D, WARPS> without slots: 16 rows a warp.
 const Full = struct {
     dim: usize,

zig/src/families/radio/tower.zig

@@ -125,13 +125,10 @@ pub const Tower = struct {
             try k.layerNorm(s, b.h, l.norm1, l.norm1_bias, b.x, rows, hd, c.eps);
             try vit.linear(g, k, s, b.acc, b.x, l.qkv, l.qkv_bias, b.wide, rows, 3 * hd, hd);
             try t.at.full(s, b.wide, b.a, rows, c.heads, dim, scale);
-            try vit.linear(g, k, s, b.acc, b.a, l.proj, l.proj_bias, b.x, rows, hd, hd);
-            try k.pointwise(s, b.x, b.h, b.h, rows * hd, .none);
+            try vit.linearAct(g, k, s, b.acc, b.a, l.proj, l.proj_bias, b.x, rows, hd, hd, b.h, b.h, .none);
             try k.layerNorm(s, b.h, l.norm2, l.norm2_bias, b.x, rows, hd, c.eps);
-            try vit.linear(g, k, s, b.acc, b.x, l.fc1, l.fc1_bias, b.wide, rows, c.ffn, hd);
-            try k.pointwise(s, b.wide, null, b.wide, rows * c.ffn, .gelu);
-            try vit.linear(g, k, s, b.acc, b.wide, l.fc2, l.fc2_bias, b.x, rows, hd, c.ffn);
-            try k.pointwise(s, b.x, b.h, b.h, rows * hd, .none);
+            try vit.linearAct(g, k, s, b.acc, b.x, l.fc1, l.fc1_bias, b.wide, rows, c.ffn, hd, null, b.wide, .gelu);
+            try vit.linearAct(g, k, s, b.acc, b.wide, l.fc2, l.fc2_bias, b.x, rows, hd, c.ffn, b.h, b.h, .none);
         }
         if (!c.final_norm) return feats; // the patch rows: cls and registers dropped
         try k.layerNorm(s, feats, w.final, w.final_bias, b.x, n, hd, c.eps);

zig/tests/cuda/contract_tests.zig

@@ -23,6 +23,7 @@ const ck_k = @import("contract_ck.zig");
 const scan_k = @import("contract_scan.zig");
 const lane8_k = @import("contract_ck_lane.zig");
 const latent_k = @import("contract_latent.zig");
+const vit_k = @import("contract_vit_glue.zig");
 
 const round_k = @import("contract_tree_round.zig");
 
@@ -89,6 +90,18 @@ pub fn run(gpu: check.Gpu) !void {
     defer rank_sum.unload();
     var kernels: std.ArrayList(g.Kernel) = .empty;
     defer kernels.deinit(a);
+    var vit = try cuda.vit.Kernels.load(gpu.d);
+    defer vit.unload();
+    var glue: [12]vit_k.Glue = undefined;
+    for (&glue, 0..) |*item, i| item.* = .{
+        .rig = .{ .gpu = gpu, .s = s },
+        .k = &vit,
+        .act = @fromBackingInt(@intCast(i % 4)),
+        .residual = i >= 4,
+        .in_place = i >= 8,
+    };
+    defer for (&glue) |*item| item.deinit();
+    for (&glue) |*item| try kernels.append(a, item.kernel());
     var row_list: [7]rows_k.Row = undefined;
     var nrows: usize = 0;
     defer for (row_list[0..nrows]) |*r| r.deinit();

zig/tests/cuda/contract_vit_glue.zig

@@ -0,0 +1,152 @@
+//! Compare the fused epilogue against the original launches on the same guarded inputs.
+const std = @import("std");
+const cuda = @import("cuda");
+const core = @import("core");
+const c = core.contract.spec;
+const g = core.contract.gates;
+const Rig = @import("contract_rig.zig").Rig;
+
+fn depth(_: c.Shape) usize {
+    return 2;
+}
+
+fn spec(act: cuda.vit.Act, residual: bool, in_place: bool) c.Spec {
+    const names = comptime blk: {
+        var result: [3][4][]const u8 = undefined;
+        for (.{ "no residual", "residual", "in-place residual" }, 0..) |mode, i| {
+            for (.{ "none", "gelu_tanh", "gelu", "relu2" }, 0..) |activation, j| {
+                result[i][j] = "cuda vit tf_vit_bias_act " ++ activation ++ " " ++ mode;
+            }
+        }
+        break :blk result;
+    };
+    return .{
+        .name = names[if (in_place) @as(usize, 2) else if (residual) 1 else 0][@intCast(@backingInt(act))],
+        .accum = .f32,
+        .depth = depth,
+        .out = .bf16,
+        .in_place = in_place,
+        .replaces = "tf_vit_bias then tf_vit_pointwise",
+        .nan = .{ .stated = "byte-identical to bias then pointwise, including non-finite inputs" },
+        .tiles = .{ .{}, .{ .size = 256 }, .{} },
+        .served = &.{ .{ .rows = 13322, .n = 1280, .k = 1 }, .{ .rows = 13322, .n = 5120, .k = 1 } },
+        .sums = .{ .f16 = 0, .bf16 = 0, .why = "pointwise rounding boundaries are checked against the replaced kernels" },
+    };
+}
+
+pub const Glue = struct {
+    rig: Rig,
+    k: *const cuda.vit.Kernels,
+    act: cuda.vit.Act,
+    residual: bool,
+    in_place: bool,
+    shape: c.Shape = undefined,
+    acc: u64 = 0,
+    bias: u64 = 0,
+    r: u64 = 0,
+    out: u64 = 0,
+    mid: u64 = 0,
+    base: u64 = 0,
+    msg: [200]u8 = undefined,
+
+    pub fn deinit(self: *Glue) void {
+        self.rig.deinit();
+    }
+
+    const vt = g.hooks(Glue);
+
+    pub fn kernel(self: *Glue) g.Kernel {
+        return .{ .spec = spec(self.act, self.residual, self.in_place), .ctx = self, .vt = &vt };
+    }
+
+    pub fn prepare(self: *Glue, sh: c.Shape, fill: g.Fill, seed: u64) !bool {
+        if (sh.k != 1) return false;
+        self.rig.reset();
+        self.shape = sh;
+        const a = self.rig.gpu.gpa;
+        const acc = try a.alloc(f32, self.count());
+        defer a.free(acc);
+        const bias = try a.alloc(u16, sh.n);
+        defer a.free(bias);
+        const residual = try a.alloc(u16, self.count());
+        defer a.free(residual);
+        var prng = std.Random.DefaultPrng.init(seed);
+        const rng = prng.random();
+        for (acc, residual, 0..) |*v, *r, i| {
+            v.* = rng.float(f32) * 16 - 8;
+            r.* = c.toBf16(rng.float(f32) * 8 - 4);
+            if (fill == .nonfinite) v.* = switch (i % 7) {
+                0 => std.math.nan(f32),
+                1 => std.math.inf(f32),
+                2 => -std.math.inf(f32),
+                3 => -0.0,
+                else => v.*,
+            };
+        }
+        for (bias) |*b| b.* = c.toBf16(rng.float(f32) * 2 - 1);
+        if (fill == .nonfinite) {
+            bias[0] = c.toBf16(std.math.nan(f32));
+            if (sh.n > 1) bias[1] = c.toBf16(std.math.inf(f32));
+            for (residual, 0..) |*r, i| {
+                if (i % 11 == 4) r.* = c.toBf16(std.math.inf(f32));
+                if (i % 11 == 5) r.* = c.toBf16(-std.math.inf(f32));
+            }
+        }
+        if (fill == .stress) {
+            @memset(acc, 8);
+            @memset(bias, c.toBf16(0.25));
+            @memset(residual, c.toBf16(-8));
+        }
+        self.acc = try self.rig.input(std.mem.sliceAsBytes(acc));
+        self.bias = try self.rig.input(std.mem.sliceAsBytes(bias));
+        self.r = try self.rig.input(std.mem.sliceAsBytes(residual));
+        self.out = try self.rig.output(self.count() * 2);
+        self.mid = try self.rig.scratch(self.count() * 2);
+        self.base = try self.rig.scratch(self.count() * 2);
+        return true;
+    }
+
+    pub fn count(self: *Glue) usize {
+        return self.shape.rows * self.shape.n;
+    }
+
+    pub fn run(self: *Glue) !void {
+        if (self.in_place) try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = self.out, .len = self.count() * 2 }).copyFrom(0, self.r, self.count() * 2, self.rig.s.handle);
+        const residual: ?u64 = if (self.in_place) self.out else if (self.residual) self.r else null;
+        try self.k.biasAct(self.rig.s, self.acc, self.bias, residual, self.out, self.shape.rows, self.shape.n, self.act);
+        try self.rig.s.synchronize();
+    }
+
+    pub fn replaced(self: *Glue) ![]const u8 {
+        try self.k.bias(self.rig.s, self.acc, self.bias, self.mid, self.shape.rows, self.shape.n);
+        try self.k.pointwise(self.rig.s, self.mid, if (self.residual) self.r else null, self.base, self.count(), self.act);
+        return self.rig.bytes(self.base, self.count() * 2);
+    }
+
+    pub fn reference(self: *Glue, want: []f64, mag: []f64) !void {
+        const bytes = try self.replaced();
+        for (want, mag, 0..) |*w, *m, i| {
+            w.* = c.read(.bf16, bytes, i);
+            m.* = @abs(w.*);
+        }
+    }
+
+    pub fn simulate(self: *Glue, acc: c.Accum, order: g.Order, out: []f64) !void {
+        const biased = core.contract.rows.sumIn(acc, order, &.{ 8, 0.25 });
+        var v = c.fromBf16(c.toBf16(@floatCast(biased)));
+        // Both GELU forms round to the identity on these positive stress inputs.
+        if (self.act == .relu2) v = c.fromBf16(c.toBf16(v * v));
+        const result = if (self.residual) core.contract.rows.sumIn(acc, order, &.{ -8, v }) else v;
+        @memset(out, c.fromBf16(c.toBf16(@floatCast(result))));
+    }
+
+    pub fn check(self: *Glue, _: 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 tf_vit_bias then tf_vit_pointwise", .{ at, got.len });
+        for (0..self.count()) |i| if (c.classOf(c.read(.bf16, got, i)) != .finite) return null;
+        return "no output is non-finite: the planted values prove nothing";
+    }
+};