From df85c7598f66041432079ca2da2f3780995bcdb3 Mon Sep 17 00:00:00 2001 From: gilby Date: Thu, 29 Oct 2026 19:10:11 +0100 Subject: [PATCH] hip: decode router over the bf16 rows the checkpoint holds The decode-window router read the router weights as fp32 rows widened at load. tf_router_bf16 (and its fp16-activation form) reads the bf16 rows the checkpoint already stores, prefetching four groups, with each lane's multiply-add order unchanged, so its outputs equal the fp32-rows router's bytes. Registered with the HIP kernel contract as that router's replacement. Refs #1603 --- docs/kernel-contract.md | 3 + zig/kernels/hip/decode/decode.hip | 63 ++++++ zig/src/core/contract/hip_records.zig | 3 +- zig/src/core/contract/hip_records_router.zig | 31 +++ zig/src/families/flashnext/hip/load.zig | 1 + zig/src/hip/kernel_tests/decode.zig | 77 +++++++ zig/src/hip/launches.zig | 5 +- zig/src/hip/launches/moe.zig | 14 ++ zig/src/hip/moe.zig | 6 +- zig/src/hip/ops/moe.zig | 9 + zig/src/hip/ops/ops.zig | 1 + zig/tests/hip/contract_main.zig | 7 + zig/tests/hip/contract_router.zig | 214 +++++++++++++++++++ 13 files changed, 430 insertions(+), 4 deletions(-) create mode 100644 zig/src/core/contract/hip_records_router.zig create mode 100644 zig/tests/hip/contract_router.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1ed43ac146..504f01551a 100644 --- a/docs/kernel-contract.md +++ b/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. diff --git a/zig/kernels/hip/decode/decode.hip b/zig/kernels/hip/decode/decode.hip index 42ba9fd3da..19e2eff375 100644 --- a/zig/kernels/hip/decode/decode.hip +++ b/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 +__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(x) + static_cast(first) * d; + const uint16_t* wb = static_cast(rows) + static_cast(expert) * d; + logits += static_cast(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(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(k) * d + at; + const float4 xv = widen4(*reinterpret_cast(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(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 @@ -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) { diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig index 6d8ee06f85..199d4d239e 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/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{ diff --git a/zig/src/core/contract/hip_records_router.zig b/zig/src/core/contract/hip_records_router.zig new file mode 100644 index 0000000000..f16010c4d4 --- /dev/null +++ b/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 }, + }, + }, +}; diff --git a/zig/src/families/flashnext/hip/load.zig b/zig/src/families/flashnext/hip/load.zig index 91923df52e..2b92642de5 100644 --- a/zig/src/families/flashnext/hip/load.zig +++ b/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, diff --git a/zig/src/hip/kernel_tests/decode.zig b/zig/src/hip/kernel_tests/decode.zig index 5c85b799ca..f9c6e7c743 100644 --- a/zig/src/hip/kernel_tests/decode.zig +++ b/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(); diff --git a/zig/src/hip/launches.zig b/zig/src/hip/launches.zig index ace7c08e7c..0a7801cd87 100644 --- a/zig/src/hip/launches.zig +++ b/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; diff --git a/zig/src/hip/launches/moe.zig b/zig/src/hip/launches/moe.zig index 2cccb06e86..17bfc68edd 100644 --- a/zig/src/hip/launches/moe.zig +++ b/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 = .{}; diff --git a/zig/src/hip/moe.zig b/zig/src/hip/moe.zig index 0ddd802a9a..eec8d726d2 100644 --- a/zig/src/hip/moe.zig +++ b/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 }; } diff --git a/zig/src/hip/ops/moe.zig b/zig/src/hip/ops/moe.zig index effaebbe81..6afbf7fa54 100644 --- a/zig/src/hip/ops/moe.zig +++ b/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; diff --git a/zig/src/hip/ops/ops.zig b/zig/src/hip/ops/ops.zig index 04a17e8410..3960bc2568 100644 --- a/zig/src/hip/ops/ops.zig +++ b/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; diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index 5b07bfefb3..1577a2e69b 100644 --- a/zig/tests/hip/contract_main.zig +++ b/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); diff --git a/zig/tests/hip/contract_router.zig b/zig/tests/hip/contract_router.zig new file mode 100644 index 0000000000..4c8caa2154 --- /dev/null +++ b/zig/tests/hip/contract_router.zig @@ -0,0 +1,214 @@ +//! The router over the stored bf16 rows on HIP: the fp32-data router's logits, bit for bit. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const kit = @import("contract_kit.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows_k = contract.rows; + +/// A lane's groups of four, 128 values apart, then the wave's five-step tree. +fn depth(s: c.Shape) usize { + return (s.k + 127) / 128 + 5; +} + +const names = [2][]const u8{ + "hip tf_router_bf16_f16 (the stored bf16 rows)", + "hip tf_router_bf16 (the stored bf16 rows)", +}; + +pub fn spec(k: kit.Kind) c.Spec { + return .{ + .name = names[@intFromBool(k == .bf16)], + .accum = .f32, + .depth = depth, + .out = .f32, + .replaces = if (k == .bf16) "tf_router_decode_bf16" else "tf_router_decode_f16", + // a round's rows: 16 a block of row groups, and a lane's group of four over the hidden width + .tiles = .{ .{ .size = 16 }, .{}, .{ .size = 128, .step = 4, .min = 4 } }, + .served = &.{ + .{ .rows = 16, .n = 513, .k = 2560 }, + .{ .rows = 4, .n = 513, .k = 2560 }, + .{ .rows = 1, .n = 513, .k = 2560 }, + }, + }; +} + +/// shape.rows rows of shape.k, and shape.n router rows as a checkpoint holds them (bf16) and widened to fp32. +pub const Router = struct { + rig: rig_.Rig, + l: *const hip.Launcher, + k: kit.Kind, + arena: std.heap.ArenaAllocator, + shape: c.Shape = undefined, + x: []f64 = &.{}, + w: []f64 = &.{}, + in: u64 = 0, + rows16: u64 = 0, + rows32: u64 = 0, + out: u64 = 0, + base: u64 = 0, + terms: []f64 = &.{}, + + pub fn init(a: std.mem.Allocator, rt: *const hip.Runtime, s: hip.Stream, l: *const hip.Launcher, k: kit.Kind) Router { + return .{ .rig = .{ .a = a, .r = rt, .s = s }, .l = l, .k = k, .arena = .init(a) }; + } + + pub fn deinit(self: *Router) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Router) g.Kernel { + return .{ .spec = spec(self.k), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Router); + + /// The rows of the router the engine serves: the experts and the shared expert's gate. + fn experts(self: *const Router) usize { + return self.shape.n; + } + + pub fn prepare(self: *Router, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 16 or shape.n == 0 or shape.k < 4 or shape.k % 4 != 0) return false; + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + self.shape = shape; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const n = shape.rows * shape.k; + const e = shape.n; + self.x = try a.alloc(f64, n); + self.w = try a.alloc(f64, e * shape.k); + self.terms = try a.alloc(f64, @max(shape.k, 1)); + // random weights under 1/8 keep the logits small; the stress rows are integers gate 4 can sum exactly + for (self.w) |*v| v.* = switch (fill) { + .stress => 0, + else => kit.round(.bf16, (rng.float(f64) * 2 - 1) / 8), + }; + for (self.x) |*v| v.* = switch (fill) { + .stress => 0, + else => kit.round(self.k, rng.float(f64) * 2 - 1), + }; + if (fill == .stress) self.stress(rng) else if (fill == .nonfinite) self.plant(); + const xb = try a.alloc(u8, n * kit.size(self.k)); + for (self.x, 0..) |v, i| kit.put(self.k, xb, i, v); + const w16 = try a.alloc(u8, 2 * e * shape.k); + const w32 = try a.alloc(u8, 4 * e * shape.k); + for (self.w, 0..) |v, i| { + const h = c.toBf16(@floatCast(v)); + std.mem.writeInt(u16, w16[2 * i ..][0..2], h, .little); + std.mem.writeInt(u32, w32[4 * i ..][0..4], @as(u32, h) << 16, .little); + } + self.in = try self.rig.input(xb); + self.rows16 = try self.rig.input(w16); + self.rows32 = try self.rig.input(w32); + self.out = try self.rig.output(shape.rows * e * 4); + self.base = try self.rig.scratch(shape.rows * e * 4); + return true; + } + + /// Gate 4's rows: integers a lane sums exactly in fp32, mirrored so the row's own result stays small. + fn stress(self: *Router, rng: std.Random) void { + const k = self.shape.k; + const h = k / 2; + for (0..self.shape.rows) |r| for (0..h) |i| { + const v: f64 = @floatFromInt(1 + rng.uintLessThan(u32, 128)); + self.x[r * k + i] = v; + self.x[r * k + h + i] = v + @as(f64, @floatFromInt(@intFromBool(i < 8 and k > 16))); + }; + if (k % 2 == 1) for (0..self.shape.rows) |r| { + self.x[r * k + k - 1] = 0; + }; + for (0..self.shape.n) |e| for (0..h) |i| { + const v: f64 = @floatFromInt(1 + rng.uintLessThan(u32, 7)); + self.w[e * k + i] = v; + self.w[e * k + h + i] = -v; + }; + if (k % 2 == 1) for (0..self.shape.n) |e| { + self.w[e * k + k - 1] = 0; + }; + } + + /// Gate 1: a NaN expert, an infinite expert each way, and a NaN input row. + fn plant(self: *Router) void { + const k = self.shape.k; + const e = self.shape.n; + self.w[k / 3] = std.math.nan(f64); + self.w[1 * k + k / 2] = std.math.inf(f64); + if (e > 2 and k / 2 + 1 < k) self.w[2 * k + k / 2 + 1] = -std.math.inf(f64); + if (self.shape.rows > 1) self.x[k] = std.math.nan(f64); + } + + /// One row's terms against one expert's row, into `t`. + fn termsOf(self: *Router, r: usize, e: usize, t: []f64) void { + const k = self.shape.k; + for (t, 0..) |*o, i| o.* = self.x[r * k + i] * self.w[e * k + i]; + } + + pub fn run(self: *Router) !void { + const s = self.shape; + const ran = try self.l.routerBf16( + @ptrFromInt(self.in), + @ptrFromInt(self.rows16), + @backingInt(self.k), + @ptrFromInt(self.out), + @intCast(s.rows), + @intCast(s.k), + @intCast(s.n), + self.rig.s.handle, + ); + if (!ran) return error.BadShape; + try self.rig.s.synchronize(); + } + + /// The router over the fp32 rows widened from the same bf16 values. + pub fn replaced(self: *Router) ![]const u8 { + const s = self.shape; + const ran = try self.l.routerWindow( + @ptrFromInt(self.in), + @backingInt(self.k), + @ptrFromInt(self.rows32), + @ptrFromInt(self.base), + @intCast(s.rows), + @intCast(s.k), + @intCast(s.n), + self.rig.s.handle, + ); + if (!ran) return error.BadShape; + return self.rig.bytes(self.base, s.rows * s.n * 4); + } + + pub fn count(self: *Router) usize { + return self.shape.rows * self.shape.n; + } + + pub fn reference(self: *Router, want: []f64, mag: []f64) !void { + const t = self.terms; + for (0..self.shape.rows) |r| for (0..self.shape.n) |e| { + self.termsOf(r, e, t); + var sum: f64 = 0; + var abs: f64 = 0; + for (t) |v| { + sum += v; + abs += @abs(v); + } + want[r * self.shape.n + e] = sum; + mag[r * self.shape.n + e] = abs; + }; + } + + /// The logits if each lane's sums ran in `acc` in `order`, the wave's tree exact. + pub fn simulate(self: *Router, acc: c.Accum, order: rows_k.Order, out: []f64) !void { + const t = self.terms; + for (0..self.shape.rows) |r| for (0..self.shape.n) |e| { + self.termsOf(r, e, t); + out[r * self.shape.n + e] = rows_k.sumIn(acc, order, t); + }; + } +};