hip: decode router over the bf16 rows the checkpoint holds

docs/kernel-contract.md

@@ -175,6 +175,9 @@ each row then L2-normed, against float64's ranges at their bf16 points and expf'
 CUDA also checks `dense_rows`' scaled epilogue (Dory's logits, fp32(bf16(x . w)) times a per-column scale), bit-equal to its
 unscaled tile's outputs rounded and scaled on the host.
 HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row (fp16 and bf16 out).
+HIP also checks the router over the rows a checkpoint already holds in bf16 (`tf_router_bf16`, and its fp16-activation
+form), bit-equal to the fp32-rows router it replaces at 1 to 16 rows, with both kinds' non-finite plants through gate 1
+and the exact sums of gate 4.
 
 RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then
 `tf_vit_pointwise`, byte for byte.

zig/kernels/hip/decode/decode.hip

@@ -28,6 +28,58 @@ __device__ __forceinline__ float4 widen4(uint2 v) {
 constexpr int kRouterWaves = 4;
 constexpr int kRouterSteps = 16;  // float4 loads a lane keeps in flight
 constexpr int kRouterRows = 16;
+constexpr int kRouterGroups = 4;  // bf16 weight groups a lane keeps in flight
+
+// The bf16 rows' router: 16 rows a block, a wave an expert, each lane summing (lane * 4 + 128 * g) x, y, z, w.
+template <int KIND>
+__device__ void router_bf16(const void* x, const void* rows, float* logits, int r, int d, int e) {
+    const int lane = threadIdx.x & 31;
+    const int expert = blockIdx.x * kRouterWaves + (threadIdx.x >> 5);
+    if (expert >= e) return;
+    const int first = blockIdx.y * kRouterRows;
+    const uint16_t* xb = static_cast<const uint16_t*>(x) + static_cast<long long>(first) * d;
+    const uint16_t* wb = static_cast<const uint16_t*>(rows) + static_cast<long long>(expert) * d;
+    logits += static_cast<long long>(first) * e;
+    const int rows_left = r - first;
+    const int live = rows_left < kRouterRows ? rows_left : kRouterRows;
+    float acc[kRouterRows];
+#pragma unroll
+    for (int k = 0; k < kRouterRows; ++k) acc[k] = 0.f;
+    // a lane's group: four values a count, strided by the block's width, so its sums meet the fp32 kernel's
+    const int groups = (d + 127) / 128;
+    for (int g = 0; g < groups; g += kRouterGroups) {
+        uint2 wv[kRouterGroups];
+#pragma unroll
+        for (int p = 0; p < kRouterGroups; ++p) {
+            const int at = lane * 4 + 128 * (g + p);
+            wv[p] = (g + p < groups && at < d) ? *reinterpret_cast<const uint2*>(wb + at) : make_uint2(0u, 0u);
+        }
+#pragma unroll
+        for (int p = 0; p < kRouterGroups; ++p) {
+            const int at = lane * 4 + 128 * (g + p);
+            if (g + p >= groups || at >= d) continue;
+            // the checkpoint holds these rows in bf16 whatever the activations are
+            const float4 w = widen4<2>(make_uint2(wv[p].x, wv[p].y));
+#pragma unroll
+            for (int k = 0; k < kRouterRows; ++k) {
+                if (k >= live) continue;
+                const uint16_t* xr = xb + static_cast<long long>(k) * d + at;
+                const float4 xv = widen4<KIND>(*reinterpret_cast<const uint2*>(xr));
+                acc[k] = fmaf(xv.x, w.x, acc[k]);
+                acc[k] = fmaf(xv.y, w.y, acc[k]);
+                acc[k] = fmaf(xv.z, w.z, acc[k]);
+                acc[k] = fmaf(xv.w, w.w, acc[k]);
+            }
+        }
+    }
+#pragma unroll
+    for (int k = 0; k < kRouterRows; ++k) {
+        float v = acc[k];
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1) v += __shfl_xor(v, mask, 32);
+        if (lane == 0 && k < live) logits[static_cast<long long>(k) * e + expert] = v;
+    }
+}
 
 // logits[r, e] = x[r] . rows[e], 16 rows a block, a wave an expert; d is a multiple of 4, weight rows 16-byte aligned.
 template <int KIND>
@@ -398,6 +450,17 @@ extern "C" __global__ void __launch_bounds__(32 * kRouterWaves)
     router_decode<2>(x, rows, logits, r, d, e);
 }
 
+// fp16 activations read bf16 rows: only the activations follow the kind.
+extern "C" __global__ void __launch_bounds__(32 * kRouterWaves)
+    tf_router_bf16_f16(const void* x, const void* rows, float* logits, int r, int d, int e) {
+    router_bf16<1>(x, rows, logits, r, d, e);
+}
+
+extern "C" __global__ void __launch_bounds__(32 * kRouterWaves)
+    tf_router_bf16(const void* x, const void* rows, float* logits, int r, int d, int e) {
+    router_bf16<2>(x, rows, logits, r, d, e);
+}
+
 extern "C" __global__ void __launch_bounds__(kTailThreads)
     tf_tail(void* x, const void* y, const float* wts, const float* weight, void* normed, int kind, int slots, int width,
             float eps) {

zig/src/core/contract/hip_records.zig

@@ -463,7 +463,8 @@ const own_records = [_]Record{
     },
 } ++ @import("hip_records_pair.zig").records;
 
-pub const records = own_records ++ @import("hip_records_stream.zig").records;
+pub const records = own_records ++ @import("hip_records_stream.zig").records ++
+    @import("hip_records_router.zig").records;
 
 /// Gate 5's planted HIP controls: a launch past its __launch_bounds__, and 72 KiB of LDS where a workgroup has 64.
 pub const controls = [_]Record{

zig/src/core/contract/hip_records_router.zig

@@ -0,0 +1,31 @@
+//! Gate 5's records of the router over the stored bf16 rows (decode.hip: tf_router_bf16).
+const std = @import("std");
+const limits = @import("limits.zig");
+const Record = @import("hip_records.zig").Record;
+
+/// The stored-bf16 router, `suffix` naming its activation kind's instantiation; 4 experts a block, 128 threads.
+fn routeTile(comptime name: []const u8, comptime symbol: []const u8) Record {
+    return .{
+        .launch = .{ .name = name, .symbol = symbol, .image = "decode", .threads = 128 },
+        .usage = .{
+            .{ .vgprs = 35, .sgprs = 64, .lds = 0, .private = 0, .bound = 128, .wave = 32 },
+            .{ .vgprs = 35, .sgprs = 64, .lds = 0, .private = 0, .bound = 128, .wave = 32 },
+            .{ .vgprs = 35, .sgprs = 64, .lds = 0, .private = 0, .bound = 128, .wave = 32 },
+            .{ .vgprs = 35, .sgprs = 65, .lds = 0, .private = 0, .bound = 128, .wave = 32 },
+        },
+    };
+}
+
+/// The same launch in bf16, whose weights are the wider activation's rows.
+pub const records = [_]Record{
+    routeTile("router tf_router_bf16_f16", "tf_router_bf16_f16"),
+    .{
+        .launch = .{ .name = "router tf_router_bf16", .symbol = "tf_router_bf16", .image = "decode", .threads = 128 },
+        .usage = .{
+            .{ .vgprs = 36, .sgprs = 64, .lds = 0, .private = 0, .bound = 128, .wave = 32 },
+            .{ .vgprs = 36, .sgprs = 64, .lds = 0, .private = 0, .bound = 128, .wave = 32 },
+            .{ .vgprs = 36, .sgprs = 64, .lds = 0, .private = 0, .bound = 128, .wave = 32 },
+            .{ .vgprs = 36, .sgprs = 65, .lds = 0, .private = 0, .bound = 128, .wave = 32 },
+        },
+    },
+};

zig/src/families/flashnext/hip/load.zig

@@ -182,6 +182,7 @@ const Loader = struct {
         const u = l.up.uploader();
         return .{
             .rows32 = (try u.tensor(r.rows32)).ptr,
+            .rows16 = (try u.tensor(r.router)).ptr,
             .top_k = r.top_k,
             .rows = r.router.shape[0],
             .remap = if (r.remap) |t| (try u.tensor(t)).ptr else 0,

zig/src/hip/kernel_tests/decode.zig

@@ -123,6 +123,83 @@ test "the merged router launch is no further from float64 than twice the launch
     for ([_]usize{ 1, 2, 4, 8 }) |rows| try router(t, &rng, rows, 2048, 257);
 }
 
+/// The router over the stored bf16 rows and over the fp32 rows widened from them: the same logits' bits.
+fn routerRows(t: *Rig, rng: *rig.Rng, kind: c_int, rows: usize, d: usize, experts: usize, planted: bool) !void {
+    const fits = struct {
+        fn bits(k: c_int, v: f32) u16 {
+            return if (k == 1) rig.f16Bits(v) else rig.bf16Bits(v);
+        }
+        fn value(k: c_int, b: u16) f64 {
+            return if (k == 1) rig.f16Value(b) else rig.bf16Value(b);
+        }
+    };
+    const hx = try gpa.alloc(u16, rows * d);
+    defer gpa.free(hx);
+    for (hx) |*v| v.* = fits.bits(kind, rng.unit());
+    const hw16 = try gpa.alloc(u16, experts * d);
+    defer gpa.free(hw16);
+    for (hw16) |*v| v.* = rig.bf16Bits(rng.unit() / 8.0);
+    const hw32 = try gpa.alloc(f32, experts * d);
+    defer gpa.free(hw32);
+    for (hw16, hw32) |b, *v| v.* = @as(f32, @bitCast(@as(u32, b) << 16));
+    // a checkpoint's rows are the widened bf16 ones; a plant the same in both, so the sums must still match bit for bit
+    if (planted) {
+        hx[d / 2] = fits.bits(kind, 1.0);
+        hw16[d] = rig.bf16Bits(-0.0);
+        hw32[d] = -0.0;
+        hw16[2 * d + 1] = 0x7f80;
+        hw32[2 * d + 1] = @bitCast(@as(u32, 0x7f80) << 16);
+    }
+    var x = try t.upload(hx);
+    defer x.free();
+    var w16 = try t.upload(hw16);
+    defer w16.free();
+    var w32 = try t.upload(hw32);
+    defer w32.free();
+    const size = rows * experts * 4;
+    var got: [2]@import("../memory.zig").DeviceBuffer = .{ try t.alloc(size), try t.alloc(size) };
+    defer for (&got) |*b| b.free();
+    // over the fp32 rows widened from the same bf16 values, and over the bf16 rows themselves
+    const s = t.stream.handle;
+    const rows_c, const d_c, const ex_c = .{ @as(c_int, @intCast(rows)), @as(c_int, @intCast(d)), @as(c_int, @intCast(experts)) };
+    try std.testing.expect(try t.on.routerWindow(p(x), kind, @ptrFromInt(at(w32)), @ptrFromInt(at(got[0])), rows_c, d_c, ex_c, s));
+    try std.testing.expect(try t.on.routerBf16(p(x), p(w16), kind, @ptrFromInt(at(got[1])), rows_c, d_c, ex_c, s));
+    try t.stream.synchronize();
+    const a = try rig.download(u8, got[0]);
+    defer gpa.free(a);
+    const b = try rig.download(u8, got[1]);
+    defer gpa.free(b);
+    try std.testing.expectEqualSlices(u8, a, b);
+    if (planted) return;
+    // and both are the float64 sum of the same values: the comparison above is not two equally empty buffers
+    const logits = std.mem.bytesAsSlice(f32, b);
+    for (0..rows) |r| for (0..experts) |e| {
+        var y: f64 = 0;
+        var norm: f64 = 0;
+        for (0..d) |i| {
+            const term = fits.value(kind, hx[r * d + i]) * @as(f64, hw32[e * d + i]);
+            y += term;
+            norm += @abs(term);
+        }
+        try std.testing.expect(@abs(@as(f64, logits[r * experts + e]) - y) / norm <= 1e-5);
+    };
+}
+
+test "the router over the stored bf16 rows writes the bits the router over the fp32 rows writes" {
+    const t = try Rig.open(1 << 20);
+    defer t.close();
+    var rng: rig.Rng = .{ .state = 0x9E3779B97F4A7C15 };
+    // both activation kinds: the weights are bf16 whichever one the chip's activations are
+    for ([_]c_int{ 1, 2 }) |kind| for ([_]usize{ 1, 2, 4, 8, 16 }) |rows| {
+        try routerRows(t, &rng, kind, rows, 2048, 257, false);
+    };
+    for ([_]c_int{ 1, 2 }) |kind| {
+        try routerRows(t, &rng, kind, 4, 2560, 513, false);
+        try routerRows(t, &rng, kind, 3, 2560, 513, true);
+        try routerRows(t, &rng, kind, 16, 2048, 257, true);
+    }
+}
+
 test "the merged residual-and-norm tails are no further from float64 than twice the launches they replace" {
     const t = try Rig.open(1 << 20);
     defer t.close();

zig/src/hip/launches.zig

@@ -107,6 +107,8 @@ pub const Launcher = struct {
 
     const Decode = struct {
         router: [2]Function, // fp16, bf16
+        /// The router over the stored bf16 rows: the same logits' bits from half the bytes.
+        router_bf16: [2]Function, // fp16, bf16
         tail: Function,
         conv_split: Function,
         rms2: Function,
@@ -168,7 +170,7 @@ pub const Launcher = struct {
         const pl = l.mods[@backingInt(kernels.Group.plan)];
         l.fuse = choices.fuse;
         l.wide = choices.wide;
-        l.dec = .{ .router = .{ try dec.function("tf_router_decode_f16"), try dec.function("tf_router_decode_bf16") }, .tail = try dec.function("tf_tail"), .conv_split = try dec.function("tf_conv_split"), .rms2 = try dec.function("tf_rms2"), .gnorm_out = try dec.function("tf_gnorm_out"), .select = try dec.function("tf_select_decode") };
+        l.dec = .{ .router = .{ try dec.function("tf_router_decode_f16"), try dec.function("tf_router_decode_bf16") }, .router_bf16 = .{ try dec.function("tf_router_bf16_f16"), try dec.function("tf_router_bf16") }, .tail = try dec.function("tf_tail"), .conv_split = try dec.function("tf_conv_split"), .rms2 = try dec.function("tf_rms2"), .gnorm_out = try dec.function("tf_gnorm_out"), .select = try dec.function("tf_select_decode") };
         l.plan = .{
             .kv_write = try pl.function("tf_plan_kv_write"),
             .kv_write_act = try pl.function("tf_plan_kv_write_act"),
@@ -303,6 +305,7 @@ pub const Launcher = struct {
     pub const tf_rope_decode = rope.tf_rope_decode;
     pub const tf_rope_prefill = rope.tf_rope_prefill;
     pub const routerTile = moe.routerTile;
+    pub const routerBf16 = moe.routerBf16;
     pub const routerWith = moe.routerWith;
     pub const routerWindow = moe.routerWindow;
     pub const tf_moe_router = moe.tf_moe_router;

zig/src/hip/launches/moe.zig

@@ -53,6 +53,20 @@ pub fn routerWindow(l: *const Launcher, x: C, kind: c_int, rows: CF, logits: F,
     return true;
 }
 
+/// A lane round's router over the stored bf16 rows, which are the fp32 rows' values held narrower.
+pub fn routerBf16(l: *const Launcher, x: C, bf16: C, kind: c_int, logits: F, r: c_int, d: c_int, e: c_int, s: S) Error!bool {
+    if (!l.fuse or r < 1 or r > 16 or @rem(d, 4) != 0 or ad(x) % 8 != 0 or ad(bf16) % 8 != 0) return false;
+    var a: Args = .{};
+    try a.add(ad(x));
+    try a.add(ad(bf16));
+    try a.add(ad(logits));
+    try a.add(r);
+    try a.add(d);
+    try a.add(e);
+    try l.go(l.dec.router_bf16[if (kind == 1) 0 else 1], dim(cdiv(e, 4), cdiv(r, 16), 1), dim(128, 1, 1), 0, s, &a);
+    return true;
+}
+
 pub fn tf_moe_router(l: *const Launcher, x: C, kind: c_int, rows: CF, logits: F, r: c_int, d: c_int, e: c_int, s: S) Error!void {
     // a prompt's rows take the 64 x 64 tiles, a round's few rows a wave an expert; each logit's sum is the same
     var a: Args = .{};

zig/src/hip/moe.zig

@@ -21,6 +21,8 @@ pub const Experts = struct {
 /// The router's rows in fp32 (E + 1, D), the shared expert's gate last, and the experts.
 pub const Routed = struct {
     rows32: u64,
+    /// The same rows as the bf16 they are held in (the fp32 rows' values widened), which a round reads instead.
+    rows16: u64 = 0,
     experts: Experts,
     top_k: usize,
     /// The router's rows, E + 1: the model's expert count, whatever share of the experts this rank holds.
@@ -76,7 +78,7 @@ pub fn run(o: Ops, kind: Kind, r: Routed, x: Tensor, rows: usize, prefill: bool)
     const wts = try o.arena.of(f32, size * slots);
     const items = try o.arena.of(i32, cap * 3);
     const members = try o.arena.of(i32, size * slots);
-    try o.moeRouter(x, r.rows32, logits, rows, ex.dims, r.rows);
+    try o.moeRouterRows(x, r.rows32, r.rows16, logits, rows, ex.dims, r.rows);
     const count: usize = cap;
     var tile: usize = undefined;
     if (r.remap != 0) {
@@ -115,7 +117,7 @@ pub fn parts(o: Ops, kind: Kind, r: Routed, x: Tensor, rows: usize) ops.Error!Pa
     const wts = try o.arena.of(f32, size * slots);
     const items = try o.arena.of(i32, cap * 3);
     const members = try o.arena.of(i32, size * slots);
-    try o.moeRouter(x, r.rows32, logits, rows, r.experts.dims, r.rows);
+    try o.moeRouterRows(x, r.rows32, r.rows16, logits, rows, r.experts.dims, r.rows);
     const p = try expertParts(o, kind, r, x, rows, logits, pick, wts, items, members, cap);
     return .{ .y = p.y, .wts = wts, .slots = slots };
 }

zig/src/hip/ops/moe.zig

@@ -16,7 +16,16 @@ pub fn moeTail(o: Ops, x: Tensor, y: u64, wts: u64, weight: u64, normed: Tensor,
     try z.tf_tail(p(x.ptr), f(y), f(wts), f(weight), p(normed.ptr), @backingInt(x.kind), int(rows), int(slots), int(width), eps, o.stream);
 }
 
+/// The logits over the fp32 rows, or over the bf16 rows a checkpoint holds too (the same values, half the bytes).
 pub fn moeRouter(o: Ops, x: Tensor, rows32: u64, logits: u64, r: usize, d: usize, e: usize) Error!void {
+    return moeRouterRows(o, x, rows32, 0, logits, r, d, e);
+}
+
+/// `rows16`: the router rows as bf16 when the model holds them (0 keeps the fp32 rows).
+pub fn moeRouterRows(o: Ops, x: Tensor, rows32: u64, rows16: u64, logits: u64, r: usize, d: usize, e: usize) Error!void {
+    if (o.window and rows16 != 0) if (o.l.routerBf16(p(x.ptr), p(rows16), @backingInt(x.kind), f(logits), int(r), int(d), int(e), o.stream)) |ran| {
+        if (ran) return;
+    } else |err| return err;
     if (o.window) if (o.l.routerWindow(p(x.ptr), @backingInt(x.kind), f(rows32), f(logits), int(r), int(d), int(e), o.stream)) |ran| {
         if (ran) return;
     } else |err| return err;

zig/src/hip/ops/ops.zig

@@ -86,6 +86,7 @@ pub const Ops = struct {
     pub const copyCols = elementwise.copyCols;
     pub const moeTail = moe.moeTail;
     pub const moeRouter = moe.moeRouter;
+    pub const moeRouterRows = moe.moeRouterRows;
     pub const moeSelect = moe.moeSelect;
     pub const moeRoute = moe.moeRoute;
     pub const moeAct = moe.moeAct;

zig/tests/hip/contract_main.zig

@@ -16,6 +16,8 @@ const relu2_k = @import("contract_relu2.zig");
 const kv_k = @import("contract_kv_write.zig");
 const argmax_k = @import("contract_argmax.zig");
 const pair_k = @import("contract_relu2_pair.zig");
+const router_k = @import("contract_router.zig");
+const kit = @import("contract_kit.zig");
 
 const g = core.contract.gates;
 
@@ -136,6 +138,11 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize {
     var pair = pair_k.Pair.init(a, &r, s, &l);
     defer pair.deinit();
     try kernels.append(a, pair.kernel());
+    // both activation kinds: the rows are bf16 whichever one the chip's activations are
+    var routers: [2]router_k.Router = undefined;
+    for (&routers, [_]kit.Kind{ .f16, .bf16 }) |*x, k| x.* = router_k.Router.init(a, &r, s, &l, k);
+    defer for (&routers) |*x| x.deinit();
+    for (&routers) |*x| try kernels.append(a, x.kernel());
     var t: g.Tally = .{};
     const t0 = ms(io);
     try g.runAll(a, kernels.items, 0x6b63, &t);