From 47ef124133841184afb680c775f31e8db65e6619 Mon Sep 17 00:00:00 2001 From: plotarmordev <299844489+plotarmordev@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:13:09 +0800 Subject: [PATCH] radio CUDA: fuse the tower's bias, activation and residual launches The RADIO tower ran a bias kernel and then a pointwise kernel at each of its bias sites. One kernel now does both, sharing the pointwise arithmetic so every output byte is the one the two-kernel path wrote; the kernel contract registers it as replacing tf_vit_bias then tf_vit_pointwise and compares bytes on every variant. --- docs/kernel-contract.md | 5 +- zig/kernels/cuda/vit.cu | 23 +++- zig/src/core/contract.zig | 1 + zig/src/core/contract/registry.zig | 1 + zig/src/core/contract/vit_records.zig | 18 +++ zig/src/cuda/vit.zig | 22 ++++ zig/src/families/radio/tower.zig | 9 +- zig/tests/cuda/contract_tests.zig | 13 +++ zig/tests/cuda/contract_vit_glue.zig | 152 ++++++++++++++++++++++++++ 9 files changed, 231 insertions(+), 13 deletions(-) create mode 100644 zig/src/core/contract/vit_records.zig create mode 100644 zig/tests/cuda/contract_vit_glue.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index bc1e3f6768..cca00df498 100644 --- a/docs/kernel-contract.md +++ b/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`), 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 diff --git a/zig/kernels/cuda/vit.cu b/zig/kernels/cuda/vit.cu index fafb08e9e0..054b3dde82 100644 --- a/zig/kernels/cuda/vit.cu +++ b/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(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(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. diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index 7aa6f3d1bd..27b912a3cf 100644 --- a/zig/src/core/contract.zig +++ b/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; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 2f307c6753..87332ab6b4 100644 --- a/zig/src/core/contract/registry.zig +++ b/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 } } }, diff --git a/zig/src/core/contract/vit_records.zig b/zig/src/core/contract/vit_records.zig new file mode 100644 index 0000000000..80ed9d03e3 --- /dev/null +++ b/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)); +} diff --git a/zig/src/cuda/vit.zig b/zig/src/cuda/vit.zig index a62067e0bd..5584ff35a9 100644 --- a/zig/src/cuda/vit.zig +++ b/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, or vision_kernel without slots: 16 rows a warp. const Full = struct { dim: usize, diff --git a/zig/src/families/radio/tower.zig b/zig/src/families/radio/tower.zig index cabd18cc47..34a61fcd70 100644 --- a/zig/src/families/radio/tower.zig +++ b/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); diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 86aa6eb966..c52788fba4 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/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(); diff --git a/zig/tests/cuda/contract_vit_glue.zig b/zig/tests/cuda/contract_vit_glue.zig new file mode 100644 index 0000000000..9f0a024203 --- /dev/null +++ b/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"; + } +};