diff --git a/zig/kernels/hip/ops/act.hip b/zig/kernels/hip/ops/act.hip index 509978952d..dc088637c8 100644 --- a/zig/kernels/hip/ops/act.hip +++ b/zig/kernels/hip/ops/act.hip @@ -338,6 +338,24 @@ __global__ void moe_combine_kernel(const float* __restrict__ y, const float* __r narrow(acc, out + static_cast(row) * d + c); } +// A shared expert beside the routed stack: its row is not in y, so wts' last slot (tf_moe_select's sigmoid) weights it. +extern "C" __global__ void tf_moe_combine_gate(const float* __restrict__ y, const float* __restrict__ wts, + const float* __restrict__ shared, void* __restrict__ out, int kind, + int slots, int d) { + const int row = blockIdx.y; + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= d) return; + float acc = 0.f; + for (int s = 0; s < slots - 1; ++s) { + acc += y[(static_cast(row) * slots + s) * d + c] * wts[row * slots + s]; + } + if (shared != nullptr) acc += shared[static_cast(row) * d + c] * wts[row * slots + slots - 1]; + const long long at = static_cast(row) * d + c; + if (kind == 0) narrow(acc, static_cast(out) + at); + else if (kind == 1) narrow(acc, static_cast<__half*>(out) + at); + else narrow(acc, reinterpret_cast(out) + at); +} + template __global__ void gdn_gate_kernel(const T* __restrict__ a, const T* __restrict__ b, const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ gate, float* __restrict__ beta, diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig index 6d8ee06f85..51cce7ecc3 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/zig/src/core/contract/hip_records.zig @@ -239,6 +239,20 @@ const own_records = [_]Record{ .{ .vgprs = 5, .sgprs = 16, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, }, }, + .{ + .launch = .{ + .name = "moe shared gate tf_moe_combine_gate", + .symbol = "tf_moe_combine_gate", + .image = "act", + .threads = 256, + }, + .usage = .{ + .{ .vgprs = 9, .sgprs = 23, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 10, .sgprs = 23, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 9, .sgprs = 23, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 9, .sgprs = 20, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + }, + }, .{ .launch = .{ .name = "add rms tf_add_rms_f16", .symbol = "tf_add_rms_f16", .image = "act", .threads = 1024 }, .usage = .{ diff --git a/zig/src/core/quant/experts.zig b/zig/src/core/quant/experts.zig index 4c76cfa10a..12d95715ad 100644 --- a/zig/src/core/quant/experts.zig +++ b/zig/src/core/quant/experts.zig @@ -13,7 +13,10 @@ const Allocator = std.mem.Allocator; const refuse = types.refuse; const join = types.join; -/// A layer's E + 1 experts, shared last: `fused` is gate then up along N per expert, or up alone when not `gated`. +/// The shared expert's own projections, at its own (bits, group): gate and up joined along N, then its down. +pub const Separate = struct { gate_up: mlx.Host, down: mlx.Host }; + +/// A layer's experts, shared last in `fused`/`down`, or beside them in `separate` at another width (then E, not E+1). pub const Experts = struct { fused: quant.Host, gated: bool, @@ -22,8 +25,10 @@ pub const Experts = struct { width: usize, /// The model width (rows of the down). dims: usize, - /// E + 1. + /// E + 1 when the shared expert is in the stack, E when it runs beside it. count: usize, + /// The shared expert at its own width; null when it stacks in `fused` and `down`. + separate: ?Separate = null, }; /// The router rows [E + 1, D] with the shared gate last, in bf16 and widened to `rows32`; `remap` is a rank's own ids. @@ -51,10 +56,10 @@ fn routerRows(a: Allocator, t: anytype, key: []const u8) !struct { data: []f32, return .{ .data = data, .rows = weight.shape[0], .cols = scale.shape[1] * group }; } -/// The routed stack and the shared expert's tensor for one suffix, as `_shared_last` checks them. +/// The routed stack's and the shared expert's tensor for one suffix; `same` says they share (bits, group). const Pair = struct { mine: Tensor, one: Tensor }; -fn pair(t: anytype, stack_key: []const u8, shared: []const u8, suffix: []const u8) !Pair { +fn pair(t: anytype, stack_key: []const u8, shared: []const u8, suffix: []const u8, same: bool) !Pair { var b1: [256]u8 = undefined; var b2: [256]u8 = undefined; var mine = try t.get(join(&b1, &.{ stack_key, suffix })); @@ -65,14 +70,26 @@ fn pair(t: anytype, stack_key: []const u8, shared: []const u8, suffix: []const u one.rank -= 1; one.shape = .{ one.shape[1], one.shape[2], one.shape[3], one.shape[4], 1 }; } - if (mine.rank != one.rank + 1 or !std.mem.eql(usize, mine.shape[1..mine.rank], one.shape[0..one.rank])) + const inner = one.rank -| 1; + if (mine.rank != one.rank + 1 or !std.mem.eql(usize, mine.shape[1 .. 1 + inner], one.shape[0..inner])) + return refuse(stack_key, "does not stack with the shared expert"); + if (same and mine.shape[one.rank] != one.shape[inner]) return refuse(stack_key, "does not stack with the shared expert"); return .{ .mine = mine, .one = one }; } -/// `fused` words, scales and biases of `srcs` stacked per expert, each expert's slices side by side along N. -fn stack(a: Allocator, srcs: []const Pair, dtype: DType) !Tensor { - const count = std.math.add(usize, srcs[0].mine.shape[0], 1) catch return error.UnexpectedTensor; +/// Whether the shared expert of `key` packs at the stack's width `w`, so its tensors stack into it. +fn sameWidth(t: anytype, key: []const u8, w: mlx.Width) !bool { + const one = try t.width(key); + return one.bits == w.bits and one.group == w.group; +} + +/// `srcs` stacked per expert along N; `with_shared` appends the shared expert as expert E, else routed experts only. +fn stack(a: Allocator, srcs: []const Pair, dtype: DType, comptime with_shared: bool) !Tensor { + const count = if (with_shared) + std.math.add(usize, srcs[0].mine.shape[0], 1) catch return error.UnexpectedTensor + else + srcs[0].mine.shape[0]; const rest = srcs[0].mine.shape[2]; var rows: usize = 0; for (srcs) |s| rows = std.math.add(usize, rows, s.mine.shape[1]) catch return error.UnexpectedTensor; @@ -83,8 +100,9 @@ fn stack(a: Allocator, srcs: []const Pair, dtype: DType) !Tensor { var at: usize = 0; for (0..count) |e| for (srcs) |s| { const n = s.mine.shape[1] * rest; - const tensor = if (e + 1 < count) s.mine else s.one; - const src = if (e + 1 < count) + const shared_row = with_shared and e + 1 == count; + const tensor = if (!shared_row) s.mine else s.one; + const src = if (!shared_row) tensor.bytes[e * n * tensor.dtype.size() ..][0 .. n * tensor.dtype.size()] else tensor.bytes; @@ -113,20 +131,88 @@ fn tableDtype(srcs: []const Pair, other: []const Pair) DType { /// One projection's tensors, checked as `_affine_side` does against (bits, group). const Part = struct { words: Pair, scales: Pair, biases: Pair }; -fn part(t: anytype, stack_key: []const u8, shared: []const u8, w: mlx.Width) !Part { +fn part(t: anytype, stack_key: []const u8, shared: []const u8, w: mlx.Width, same: bool) !Part { const p: Part = .{ - .words = try pair(t, stack_key, shared, ".weight"), - .scales = try pair(t, stack_key, shared, ".scales"), - .biases = try pair(t, stack_key, shared, ".biases"), + .words = try pair(t, stack_key, shared, ".weight", same), + .scales = try pair(t, stack_key, shared, ".scales", same), + .biases = try pair(t, stack_key, shared, ".biases", same), }; _ = mlx.geometry.shape(p.words.mine, p.scales.mine, p.biases.mine, w.bits, w.group, true) catch return refuse(stack_key, "packed stack has incompatible or unrepresentable geometry"); - _ = mlx.geometry.shape(p.words.one, p.scales.one, p.biases.one, w.bits, w.group, false) catch - return refuse(shared, "shared expert has incompatible or unrepresentable geometry"); return p; } -fn side(a: Allocator, parts: []const Part, w: mlx.Width) !quant.Host { +/// Input columns of an MLX host: its group tables' columns times the group. +fn cols(h: mlx.Host) usize { + return h.scales.shape[1] * h.group; +} + +/// The routed stack's input columns: its gate/up tables' columns times the stack's group. +fn colsStack(p: Part, w: mlx.Width) usize { + return p.scales.mine.shape[2] * w.group; +} + +/// A projection's (rows, columns) as the fallback's preflight reads them from the group tables. +const Side = struct { rows: usize, columns: usize }; + +/// One affine side as the fallback checks it itself: (rows, columns) of a device-readable packed projection. +fn sideShape(t: anytype, key: []const u8) ?Side { + var b: [256]u8 = undefined; + var o: [256]u8 = undefined; + const w = t.width(key) catch return null; + const words = t.get(join(&b, &.{ key, ".weight" })) catch return null; + const scales = t.get(join(&b, &.{ key, ".scales" })) catch return null; + const biases = t.get(join(&o, &.{ key, ".biases" })) catch return null; + if (mlx.geometry.shape(words, scales, biases, w.bits, w.group, false)) |_| {} else |_| return null; + return .{ .rows = words.shape[0], .columns = scales.shape[1] * w.group }; +} + +/// The shared expert at its own widths; a side no format reads, or widths off the model's, refuses unloaded. +fn separate(a: Allocator, t: anytype, keys: []const []const u8, width: usize, dims: usize) !Separate { + var sides: [3]?Side = undefined; + for (keys, 0..) |key, i| sides[i] = sideShape(t, key); + for (sides) |maybe| if (maybe == null) return error.UnexpectedTensor; + const g = sides[1].?; + const u = sides[0].?; + const d = sides[2].?; + // gate and up one shape and the model's inputs; the down back from its gate's rows to the model width + if (g.rows != u.rows or g.columns != u.columns or g.columns != width) return error.UnexpectedTensor; + if (d.rows != dims or d.columns != g.rows) return error.UnexpectedTensor; + const gate = try mlx.read(a, t, keys[1]); + const up = try mlx.read(a, t, keys[0]); + const down = try mlx.read(a, t, keys[2]); + if (gate.bits != up.bits or gate.group != up.group) return refuse(keys[1], "shared expert's gate and up differ"); + return .{ .gate_up = try joinHost(a, gate, up), .down = down }; +} + +/// Two rank-2 hosts of one width joined along N (gate rows first, then up), as the experts' fused side is. +fn joinHost(a: Allocator, gate: mlx.Host, up: mlx.Host) !mlx.Host { + return .{ + .words = try stackHosts(a, gate.words, up.words), + .scales = try stackHosts(a, gate.scales, up.scales), + .biases = try stackHosts(a, gate.biases, up.biases), + .bits = gate.bits, + .group = gate.group, + }; +} + +/// Two rank-2 tensors of the same row length concatenated along N. +fn stackHosts(a: Allocator, top: Tensor, under: Tensor) !Tensor { + if (top.dtype != under.dtype or top.rank != 2 or under.rank != 2 or top.shape[1] != under.shape[1]) + return error.UnexpectedTensor; + const bytes = try a.alloc(u8, top.bytes.len + under.bytes.len); + @memcpy(bytes[0..top.bytes.len], top.bytes); + @memcpy(bytes[top.bytes.len..], under.bytes); + return .{ + .dtype = top.dtype, + .rank = 2, + .shape = .{ top.shape[0] + under.shape[0], top.shape[1], 1, 1, 1 }, + .bytes = bytes, + }; +} + +/// The stacked host tensors of `parts`: the routed experts, and the shared one last when `with_shared` is set. +fn side(a: Allocator, parts: []const Part, w: mlx.Width, comptime with_shared: bool) !quant.Host { var words: [2]Pair = undefined; var scales: [2]Pair = undefined; var biases: [2]Pair = undefined; @@ -138,9 +224,9 @@ fn side(a: Allocator, parts: []const Part, w: mlx.Width) !quant.Host { const n = parts.len; const dtype = tableDtype(scales[0..n], biases[0..n]); return .{ .mlx = .{ - .words = try stack(a, words[0..n], .i32), - .scales = try stack(a, scales[0..n], dtype), - .biases = try stack(a, biases[0..n], dtype), + .words = try stack(a, words[0..n], .i32, with_shared), + .scales = try stack(a, scales[0..n], dtype, with_shared), + .biases = try stack(a, biases[0..n], dtype, with_shared), .bits = w.bits, .group = w.group, } }; @@ -167,20 +253,31 @@ pub fn experts(a: Allocator, t: anytype, prefix: []const u8) !Experts { const up_w = try t.width(routed_keys[0]); const down_w = try t.width(routed_keys[2]); if (gated and !std.meta.eql(try t.width(routed_keys[1]), up_w)) return refuse(mine, "gate and up differ in width"); - const up = try part(t, routed_keys[0], shared_keys[0], up_w); - const down = try part(t, routed_keys[2], shared_keys[2], down_w); - var fused_parts: [2]Part = undefined; - var n: usize = 0; + // one width throughout stacks the shared expert in as expert E; another leaves it beside the stack + var joined = try sameWidth(t, shared_keys[0], up_w); + if (joined and gated) joined = try sameWidth(t, shared_keys[1], up_w); + if (joined) joined = try sameWidth(t, shared_keys[2], down_w); + const up = try part(t, routed_keys[0], shared_keys[0], up_w, joined); + const down = try part(t, routed_keys[2], shared_keys[2], down_w, joined); + var fused_parts: [2]Part = .{ up, up }; + var n: usize = 1; if (gated) { - fused_parts[0] = try part(t, routed_keys[1], shared_keys[1], up_w); - n = 1; + fused_parts[0] = try part(t, routed_keys[1], shared_keys[1], up_w, joined); + n = 2; if (!std.mem.eql(usize, fused_parts[0].words.mine.shape[0..3], up.words.mine.shape[0..3])) return refuse(mine, "gate and up differ in shape"); } - fused_parts[n] = up; - n += 1; - const fused = try side(a, fused_parts[0..n], up_w); - const d = try side(a, &.{down}, down_w); + if (!joined) return .{ + .fused = try side(a, fused_parts[0..n], up_w, false), + .gated = gated, + .down = try side(a, &.{down}, down_w, false), + .width = up.words.mine.shape[1], + .dims = down.words.mine.shape[1], + .count = up.words.mine.shape[0], + .separate = try separate(a, t, shared_keys[0..3], colsStack(up, up_w), down.words.mine.shape[1]), + }; + const fused = try side(a, fused_parts[0..n], up_w, true); + const d = try side(a, &.{down}, down_w, true); return .{ .fused = fused, .gated = gated, @@ -209,7 +306,9 @@ pub fn routed(a: Allocator, t: anytype, prefix: []const u8, spec: Shape) !Routed at += 1; }; const ex = try experts(a, t, prefix); - if (ex.count != rows) return refuse(prefix, "router and experts disagree on E + 1"); + // a shared expert beside the stack has no router row of its own to count against E + 1: the gate row is its own + if (ex.count != rows and !(ex.count + 1 == rows and ex.separate != null)) + return refuse(prefix, "router and experts disagree on E + 1"); return .{ .router = .{ .dtype = .bf16, .rank = 2, .shape = .{ rows, gate.cols, 1, 1, 1 }, .bytes = router }, .rows32 = .{ .dtype = .f32, .rank = 2, .shape = .{ rows, gate.cols, 1, 1, 1 }, .bytes = rows32 }, @@ -220,15 +319,19 @@ pub fn routed(a: Allocator, t: anytype, prefix: []const u8, spec: Shape) !Routed /// A rank's contiguous routed experts, the shared one on rank 0, and the remap to its own ids (-1: another rank's). pub fn sliceRank(a: Allocator, x: *Routed, r: slice.Rank) slice.Error!void { - const total = x.experts.count - 1; + // a shared expert beside the stack is not part of the stack's count, so the stack holds E routed experts there + const total = x.experts.count - @intFromBool(x.experts.separate == null); const part_n = try slice.even(total, r.world, "experts"); x.experts.fused = try quant.sliceStack(a, x.experts.fused, r, part_n); x.experts.down = try quant.sliceStack(a, x.experts.down, r, part_n); - x.experts.count = part_n + @intFromBool(r.rank == 0); + x.experts.count = part_n + @intFromBool(r.rank == 0 and x.experts.separate == null); + // the shared expert beside the stack belongs to rank 0 alone, where it is run beside that rank's launch + if (x.experts.separate) |s| x.experts.separate = if (r.rank == 0) s else null; const remap = try a.alloc(i32, total + 1); @memset(remap, -1); for (0..part_n) |i| remap[r.rank * part_n + i] = @intCast(i); - if (r.rank == 0) remap[total] = @intCast(part_n); + // the shared expert is this rank's last stacked expert only when it stacks; beside the stack it is -1 (no item) + if (r.rank == 0 and x.experts.separate == null) remap[total] = @intCast(part_n); x.remap = .{ .dtype = .i32, .rank = 1, .shape = .{ total + 1, 1, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(remap) }; } @@ -295,3 +398,154 @@ test "two routed experts and the shared one stack in order, and a rank keeps its @memcpy(std.mem.sliceAsBytes(&remap), r.remap.?.bytes); try std.testing.expectEqualSlices(i32, &.{ -1, 0, -1 }, &remap); } + +/// A table whose projections answer their own (bits, group), so a mixed-width checkpoint can be described. +const Mixed = struct { + map: std.StringHashMapUnmanaged(Tensor) = .empty, + widths: std.StringHashMapUnmanaged(mlx.Width) = .empty, + + pub fn has(m: *const Mixed, key: []const u8) bool { + return m.map.contains(key); + } + + pub fn get(m: *const Mixed, key: []const u8) !Tensor { + return m.map.get(key) orelse error.MissingTensor; + } + + pub fn width(m: *const Mixed, key: []const u8) !mlx.Width { + return m.widths.get(key) orelse .{ .bits = 4, .group = 32 }; + } +}; + +/// One affine projection of `rows` x `columns` at `w` from `seed`; `stacked` gives the routed stack shapes. +fn affineInto( + a: Allocator, + f: *Mixed, + key: []const u8, + rows: usize, + columns: usize, + w: mlx.Width, + seed: usize, + stacked: bool, +) !void { + const words_cols = columns * w.bits / 32; + const words = try a.alloc(u32, rows * words_cols); + for (words, 0..) |*x, i| x.* = @intCast(seed + i); + // a stacked projection is (E, N, K words), a lone one (N, K words); its tables add K / group + const rank: u8 = 2 + @as(u8, @intFromBool(stacked)); + const shape = [5]usize{ if (stacked) 2 else rows, if (stacked) rows / 2 else words_cols, words_cols, 1, 1 }; + const name = try std.fmt.allocPrint(a, "{s}.weight", .{key}); + try f.map.put(a, name, wordTensor(rank, shape, std.mem.sliceAsBytes(words))); + const tables = try a.alloc(u16, rows * (columns / w.group)); + @memset(tables, 0x3f80); + const tshape = [5]usize{ shape[0], if (stacked) rows / 2 else columns / w.group, columns / w.group, 1, 1 }; + for ([_][]const u8{ "scales", "biases" }) |suffix| + try f.map.put(a, try std.fmt.allocPrint(a, "{s}.{s}", .{ key, suffix }), .{ .dtype = .bf16, .rank = rank, .shape = tshape, .bytes = std.mem.sliceAsBytes(tables) }); + if (w.bits != 4 or w.group != 32) try f.widths.put(a, key, w); +} + +/// A packed int32 tensor of `shape`'s first three dimensions over `bytes`. +fn wordTensor(rank: u8, shape: [5]usize, bytes: []const u8) Tensor { + return .{ .dtype = .u32, .rank = rank, .shape = shape, .bytes = bytes }; +} + +/// Two routed experts of one width, stacked (2, 64, ...): gate and up of 64 inputs, down of the hidden 64. +fn routedInto(a: Allocator, f: *Mixed, w: mlx.Width) !void { + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }, 0..) |name, i| { + try affineInto(a, f, try std.fmt.allocPrint(a, "mlp.switch_mlp.{s}", .{name}), 2 * 64, 64, w, 1000 * i, true); + } +} + +/// The router's E rows and the shared gate's row, both over the model's width of 64. +fn routerInto(a: Allocator, f: *Mixed) !void { + const g = try a.alloc(f32, 2 * 64); + for (g, 0..) |*x, i| x.* = @as(f32, @floatFromInt(@as(i32, @intCast(i)) - 64)); + try f.map.put(a, "mlp.gate.weight", .{ .dtype = .f32, .rank = 2, .shape = .{ 2, 64, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(g) }); + const s = try a.alloc(f32, 64); + @memset(s, 0.5); + try f.map.put(a, "mlp.shared_expert_gate.weight", .{ .dtype = .f32, .rank = 2, .shape = .{ 1, 64, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(s) }); +} + +// A layer whose shared expert is 8-bit g64 over 4-bit g32 routed experts: it keeps its own width beside the stack. +test "a shared expert of another width than the routed stack keeps its own and runs beside it" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f: Mixed = .{}; + try routedInto(a, &f, .{ .bits = 4, .group = 32 }); + const wide: mlx.Width = .{ .bits = 8, .group = 64 }; + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }, 0..) |name, i| { + const key = try std.fmt.allocPrint(a, "mlp.shared_expert.{s}", .{name}); + try affineInto(a, &f, key, 64, 64, wide, 500 + 2000 * i, false); + } + try routerInto(a, &f); + const ex = try experts(a, &f, "mlp."); + // the stack holds the two routed experts alone, at their own width; the shared one is separate, at 8-bit g64 + try std.testing.expectEqual(@as(usize, 2), ex.count); + try std.testing.expectEqual(@as(u8, 4), ex.fused.mlx.bits); + try std.testing.expectEqual(@as(u16, 32), ex.down.mlx.group); + const s = ex.separate orelse return error.TestUnexpectedResult; + try std.testing.expectEqual(@as(u8, 8), s.gate_up.bits); + try std.testing.expectEqual(@as(u16, 64), s.gate_up.group); + try std.testing.expectEqual(@as(u8, 8), s.down.bits); + // gate then up along N: 128 rows of 16 words, the gate's last row and the up's first back to back + var join_at: [2]u32 = undefined; + @memcpy(std.mem.sliceAsBytes(&join_at), s.gate_up.words.bytes[4 * (16 * 64 - 1) ..][0..8]); + try std.testing.expectEqualSlices(u32, &.{ 1523, 2500 }, &join_at); + try std.testing.expectEqual(@as(usize, 128 * 16 * 4), s.gate_up.words.bytes.len); + try std.testing.expectEqual(@as(usize, 64 * 16 * 4), s.down.words.bytes.len); + // the routed stack holds gate and up of experts 0 and 1 only: 2 * 128 rows of 8 words at 4-bit, no shared row + try std.testing.expectEqual(@as(usize, 2 * (2 * 64) * 8 * 4), ex.fused.mlx.words.bytes.len); + try std.testing.expectEqual(@as(u32, 1000 + 1023), lastU32(ex.fused.mlx.words)); + try std.testing.expectEqual(@as(u8, 4), ex.down.mlx.bits); + // rank 0 keeps the shared expert beside its own, rank 1 has none; the router's rows still count E + 1 + var r = try routed(a, &f, "mlp.", .{ .experts = 2, .hidden = 64, .top_k = 1 }); + try sliceRank(a, &r, .{ .rank = 1, .world = 2 }); + try std.testing.expectEqual(@as(usize, 1), r.experts.count); + try std.testing.expectEqual(@as(?Separate, null), r.experts.separate); + var remap: [3]i32 = undefined; + @memcpy(std.mem.sliceAsBytes(&remap), r.remap.?.bytes); + try std.testing.expectEqualSlices(i32, &.{ -1, 0, -1 }, &remap); + var first: Routed = try routed(a, &f, "mlp.", .{ .experts = 2, .hidden = 64, .top_k = 1 }); + try sliceRank(a, &first, .{ .rank = 0, .world = 2 }); + try std.testing.expect(first.experts.separate != null); +} + +/// The last u32 of a host's packed words. +fn lastU32(t: Tensor) u32 { + var v: u32 = undefined; + @memcpy(std.mem.asBytes(&v), t.bytes[t.bytes.len - 4 ..]); + return v; +} + +// A shared expert at a width the device formats cannot represent (7-bit codes) is still refused, not run beside. +test "a shared expert at a width no format represents still refuses" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f: Mixed = .{}; + try routedInto(a, &f, .{ .bits = 4, .group = 32 }); + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |name| { + const key = try std.fmt.allocPrint(a, "mlp.shared_expert.{s}", .{name}); + try affineInto(a, &f, key, 64, 64, .{ .bits = 7, .group = 64 }, 500, false); + } + try routerInto(a, &f); + try std.testing.expectError(error.UnexpectedTensor, experts(a, &f, "mlp.")); +} + +// A shared expert whose input width differs from the model's is refused too: its rows cannot multiply `x`. +test "a shared expert of the wrong shape still refuses" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f: Mixed = .{}; + try routedInto(a, &f, .{ .bits = 4, .group = 32 }); + const narrow: mlx.Width = .{ .bits = 8, .group = 32 }; + for ([_][]const u8{ "gate_proj", "up_proj" }) |name| { + const key = try std.fmt.allocPrint(a, "mlp.shared_expert.{s}", .{name}); + try affineInto(a, &f, key, 64, 32, narrow, 500, false); + } + try affineInto(a, &f, "mlp.shared_expert.down_proj", 64, 64, narrow, 900, false); + try routerInto(a, &f); + try std.testing.expectError(error.UnexpectedTensor, experts(a, &f, "mlp.")); +} diff --git a/zig/src/families/flashnext/hip/load.zig b/zig/src/families/flashnext/hip/load.zig index 91923df52e..bbfe72446f 100644 --- a/zig/src/families/flashnext/hip/load.zig +++ b/zig/src/families/flashnext/hip/load.zig @@ -192,6 +192,11 @@ const Loader = struct { .count = r.experts.count, .width = r.experts.width, .dims = r.experts.dims, + // a shared expert of another width than the stack: its own projections, run beside the routed launch + .shared = if (r.experts.separate) |s| hip.moe.Shared{ + .gate_up = try quant.view(try quant.upload(u, .{ .mlx = s.gate_up }), false), + .down = try quant.view(try quant.upload(u, .{ .mlx = s.down }), false), + } else null, }, }; } diff --git a/zig/src/families/qwen3_5/backend/hip/forward/window.zig b/zig/src/families/qwen3_5/backend/hip/forward/window.zig index 6bfb659146..814d7085d5 100644 --- a/zig/src/families/qwen3_5/backend/hip/forward/window.zig +++ b/zig/src/families/qwen3_5/backend/hip/forward/window.zig @@ -54,8 +54,12 @@ pub fn forward(ops: Ops, m: *const view.Model, r: Round, trace: ?fwd.Trace) fwd. .full => |f| try attentionRows(o, m, f, r, index, normed), .linear => |l| try linearRows(o, m, l, r, index, normed), }; - // the launches of a layer's two tails merge when the rows stay on this rank - const merge = o.fused() and y.kind == m.act; + // the launches of a layer's two tails merge when the rows stay on this rank and every slot is in `y` + const stacked = switch (mlp) { + .moe => |rt| rt.experts.shared == null, + .dense => true, + }; + const merge = o.fused() and y.kind == m.act and stacked; const last = index + 1 == m.layers.len; const next_norm = if (last) m.final_norm else inputNorm(m.layers[index + 1]); const dest = if (last) out else normed; diff --git a/zig/src/families/qwen3_5/weights/checkpoint.zig b/zig/src/families/qwen3_5/weights/checkpoint.zig index b6b16ca165..72adf5a2b0 100644 --- a/zig/src/families/qwen3_5/weights/checkpoint.zig +++ b/zig/src/families/qwen3_5/weights/checkpoint.zig @@ -160,9 +160,19 @@ fn subConv(a: std.mem.Allocator, t: *const Table, base: []const u8) !Tensor { return projection.conv(a, t, projection.join(&buf, &.{ base, "conv1d.weight" })); } +/// A layer's routed experts, refusing one whose shared expert runs beside the stack: this bridge stacks it in `fused`. +fn expertStack(a: std.mem.Allocator, t: *const Table, mlp: []const u8, s: config.Spec) !host.Routed { + const r = try experts.routed(a, t, mlp, .{ .experts = s.experts, .hidden = s.hidden, .top_k = s.top_k }); + if (r.experts.separate != null) { + std.log.warn("{s}: a shared expert beside the routed stack is unsupported here", .{mlp}); + return error.UnexpectedTensor; // this bridge stacks the shared expert + } + return r; +} + /// The routed MLP when the config has experts, else gate, up and down. fn readMlp(a: std.mem.Allocator, t: *const Table, mlp: []const u8, s: config.Spec) !host.Mlp { - if (s.experts != 0) return .{ .routed = try experts.routed(a, t, mlp, .{ .experts = s.experts, .hidden = s.hidden, .top_k = s.top_k }) }; + if (s.experts != 0) return .{ .routed = try expertStack(a, t, mlp, s) }; return .{ .dense = .{ .gate = try sub(a, t, mlp, "gate_proj"), .up = try sub(a, t, mlp, "up_proj"), @@ -217,7 +227,7 @@ fn readMtp(a: std.mem.Allocator, t: *const Table, s: config.Spec) !host.Mtp { .head = head, .input_norm = try projection.float(a, t, layer ++ "input_layernorm.weight"), .post_norm = try projection.float(a, t, layer ++ "post_attention_layernorm.weight"), - .mlp = if (routed) .{ .routed = try experts.routed(a, t, mlp, .{ .experts = s.experts, .hidden = s.hidden, .top_k = s.top_k }) } else .{ .dense = .{ + .mlp = if (routed) .{ .routed = try expertStack(a, t, mlp, s) } else .{ .dense = .{ .gate = try projection.projection(a, t, mlp ++ "gate_proj"), .up = try projection.projection(a, t, mlp ++ "up_proj"), .down = try projection.projection(a, t, mlp ++ "down_proj"), diff --git a/zig/src/families/qwen3_5/weights/checkpoint_test.zig b/zig/src/families/qwen3_5/weights/checkpoint_test.zig index 02bc4f4483..47798c35e8 100644 --- a/zig/src/families/qwen3_5/weights/checkpoint_test.zig +++ b/zig/src/families/qwen3_5/weights/checkpoint_test.zig @@ -314,3 +314,61 @@ test "an MTP head shipped zero-centered gets back the 1 the trunk has baked in; try std.testing.expectEqual(@as(f64, 1.25), convert.load(.f32, t.bytes, 3)); } } + +/// The same MoE layer with the shared expert at 8-bit and its own width: its projections cannot stack into the stack. +const mixed_config = + \\{"model_type": "qwen3_5_moe", "tie_word_embeddings": true, + \\ "quantization": {"group_size": 32, "bits": 4, "mode": "affine", + \\ "language_model.model.layers.0.mlp.gate": {"group_size": 32, "bits": 8}, + \\ "language_model.model.layers.0.mlp.shared_expert_gate": {"group_size": 32, "bits": 8}, + \\ "language_model.model.layers.0.mlp.shared_expert.gate_proj": {"group_size": 32, "bits": 8}, + \\ "language_model.model.layers.0.mlp.shared_expert.up_proj": {"group_size": 32, "bits": 8}, + \\ "language_model.model.layers.0.mlp.shared_expert.down_proj": {"group_size": 32, "bits": 8}}, + \\ "text_config": {"hidden_size": 64, "num_hidden_layers": 1, "num_attention_heads": 2, + \\ "num_key_value_heads": 1, "head_dim": 32, "linear_num_key_heads": 1, "linear_num_value_heads": 2, + \\ "linear_key_head_dim": 32, "linear_value_head_dim": 16, "linear_conv_kernel_dim": 4, "vocab_size": 16, + \\ "full_attention_interval": 1, "num_experts": 2, "num_experts_per_tok": 1, "moe_intermediate_size": 32, + \\ "shared_expert_intermediate_size": 32, "eos_token_id": 1}} +; + +test "a mixed-width MoE checkpoint: this bridge stacks the shared expert, so it refuses one beside the stack" { + var img = Image.init(); + defer img.deinit(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const base = L ++ "layers.0."; + const mlp = base ++ "mlp."; + try trunk(&img); + try norms(&img, base); + try attention(&img, base); + try img.affine(mlp ++ "gate", 2, 64, 8, 32); + try img.affine(mlp ++ "shared_expert_gate", 1, 64, 8, 32); + // the routed stack at 4-bit (rank 3, experts first); the shared expert at 8-bit, its own rank-2 projections + const stack = [_]struct { suffix: []const u8, words: [3]usize, table: [3]usize }{ + .{ .suffix = "gate_proj", .words = .{ 2, 32, 8 }, .table = .{ 2, 32, 2 } }, + .{ .suffix = "up_proj", .words = .{ 2, 32, 8 }, .table = .{ 2, 32, 2 } }, + .{ .suffix = "down_proj", .words = .{ 2, 64, 4 }, .table = .{ 2, 64, 1 } }, + }; + for (stack) |p| for ([_][]const u8{ "weight", "scales", "biases" }) |part| { + var buf: [160]u8 = undefined; + const w = std.mem.eql(u8, part, "weight"); + try img.put( + try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.{s}", .{ mlp, p.suffix, part }), + if (w) "U32" else "BF16", + if (w) &p.words else &p.table, + ); + }; + try img.affine(mlp ++ "shared_expert.gate_proj", 32, 64, 8, 32); + try img.affine(mlp ++ "shared_expert.up_proj", 32, 64, 8, 32); + try img.affine(mlp ++ "shared_expert.down_proj", 64, 32, 8, 32); + try img.write(tmp.dir, "model.safetensors"); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = mixed_config }); + const root = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer gpa.free(root); + + var ck = try checkpoint.Checkpoint.open(gpa, io, root); + defer ck.close(); + std.testing.log_level = .err; // the refusal is logged as a warning, as it is meant to be + defer std.testing.log_level = .warn; + try std.testing.expectError(error.UnexpectedTensor, ck.layer(0)); +} diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index 78dc1da06f..2cc4fd2afc 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -31,6 +31,7 @@ test { _ = @import("kernel_tests/graph_split.zig"); _ = @import("kernel_tests/flashnext.zig"); _ = @import("kernel_tests/flashnext_attn.zig"); + _ = @import("kernel_tests/moe_gate.zig"); _ = @import("kernel_tests/nemotron/moe.zig"); _ = @import("kernel_tests/nemotron/pick.zig"); _ = @import("kernel_tests/nemotron/attention.zig"); diff --git a/zig/src/hip/kernel_tests/moe_gate.zig b/zig/src/hip/kernel_tests/moe_gate.zig new file mode 100644 index 0000000000..b2402f1b4d --- /dev/null +++ b/zig/src/hip/kernel_tests/moe_gate.zig @@ -0,0 +1,82 @@ +//! tf_moe_combine_gate against tf_moe_combine: the shared expert as y's last slot, and separate, byte for byte. +const std = @import("std"); +const rig = @import("rig.zig"); +const Rig = rig.Rig; +const gpa = rig.gpa; +const at = rig.at; + +/// A device address as the read-only fp32 pointer the combine launchers take. +fn f(b: anytype) ?[*]const f32 { + return @ptrFromInt(at(b)); +} + +/// A device address as the kernel's output pointer. +fn p(b: anytype) ?*anyopaque { + return @ptrFromInt(at(b)); +} + +/// `kind`'s element size: fp32 at four bytes, fp16 and bf16 at two. +fn size(kind: c_int) usize { + return if (kind == 0) 4 else 2; +} + +/// `rows` rows of `slots` - 1 routed rows plus a shared row, summed by both kernels at every output kind. +fn equal(t: *Rig, rng: *rig.Rng, rows: usize, slots: usize, d: usize) !void { + const routed = slots - 1; + const share = try gpa.alloc(f32, rows * d); // the shared expert's row, `wts`' last slot weighting it + defer gpa.free(share); + for (share) |*v| v.* = @floatCast(rng.unit() * 3 - 1.5); + const hy = try gpa.alloc(f32, rows * slots * d); + defer gpa.free(hy); + for (0..rows) |r| for (0..d) |c| { + for (0..routed) |s| hy[(r * slots + s) * d + c] = @floatCast((rng.unit() + rng.unit()) * 2 - 2); + // the slot the gate must never read from y: poison, so a gate summing every slot could not match + hy[(r * slots + slots - 1) * d + c] = 1e30; + }; + const hw = try gpa.alloc(f32, rows * slots); + defer gpa.free(hw); + for (hw) |*v| v.* = @floatCast(0.5 + rng.unit() / 2); + // the combine's y: y's routed slots, the shared expert in the last slot of every row + const stack = try gpa.alloc(f32, rows * slots * d); + defer gpa.free(stack); + for (0..rows) |r| for (0..d) |c| { + for (0..routed) |s| stack[(r * slots + s) * d + c] = hy[(r * slots + s) * d + c]; + stack[(r * slots + routed) * d + c] = share[r * d + c]; + }; + var y = try t.upload(hy); // the gate's y: the shared row's slot is poison, so the gate must skip it + defer y.free(); + var stacked = try t.upload(stack); + defer stacked.free(); + var shared = try t.upload(share); + defer shared.free(); + var wts = try t.upload(hw); + defer wts.free(); + for ([_]c_int{ 0, 1, 2 }) |kind| { + var a = try t.alloc(rows * d * size(kind)); + defer a.free(); + var b = try t.alloc(rows * d * size(kind)); + defer b.free(); + const r: c_int = @intCast(rows); + const s: c_int = @intCast(slots); + const w: c_int = @intCast(d); + try t.on.tf_moe_combine(f(stacked), f(wts), p(a), kind, r, s, w, t.stream.handle); + try t.on.tf_moe_combine_gate(f(y), f(wts), f(shared), p(b), kind, r, s, w, t.stream.handle); + try t.stream.synchronize(); + const ga = try rig.download(u8, a); + defer gpa.free(ga); + const gb = try rig.download(u8, b); + defer gpa.free(gb); + try std.testing.expectEqualSlices(u8, ga, gb); + } +} + +test "tf_moe_combine_gate equals tf_moe_combine with the shared expert as y's last slot" { + const t = try Rig.open(1 << 20); + defer t.close(); + var rng: rig.Rng = .{ .state = 0x5ECA1E }; + for ([_]usize{ 1, 3, 16 }) |rows| { + try equal(t, &rng, rows, 4, 2048); + try equal(t, &rng, 2, 3, 511); // a width that does not fill a block's columns + try equal(t, &rng, 1, 2, 64); // one routed slot beside the shared expert + } +} diff --git a/zig/src/hip/launches.zig b/zig/src/hip/launches.zig index ace7c08e7c..2913645e23 100644 --- a/zig/src/hip/launches.zig +++ b/zig/src/hip/launches.zig @@ -60,6 +60,8 @@ pub const Launcher = struct { moe_select: Function, moe_act: [2]Function, moe_combine: Triple, + /// the same sum with a shared expert run beside the stack, weighted by wts' last slot + moe_combine_gate: Function, gdn_gate: Triple, gated_delta_tile: [2]Function, // dk 128, 16 gated_delta_wave: [2]Function, @@ -241,6 +243,7 @@ pub const Launcher = struct { try act.function(ns ++ "18moe_combine_kernelI6__halfEEvPKfS4_PT_ii"), try act.function(ns ++ "18moe_combine_kernelI12hip_bfloat16EEvPKfS4_PT_ii"), }; + l.moe_combine_gate = try act.function("tf_moe_combine_gate"); l.gdn_gate = .{ try act.function(ns ++ "15gdn_gate_kernelIfEEvPKT_S4_PKfS6_PfS7_ii"), try act.function(ns ++ "15gdn_gate_kernelI6__halfEEvPKT_S5_PKfS7_PfS8_ii"), @@ -309,6 +312,7 @@ pub const Launcher = struct { pub const tf_moe_select = moe.tf_moe_select; pub const tf_moe_act = moe.tf_moe_act; pub const tf_moe_combine = moe.tf_moe_combine; + pub const tf_moe_combine_gate = moe.tf_moe_combine_gate; pub const tf_moe_route = moe.tf_moe_route; pub const tf_causal = attention.tf_causal; pub const tf_attn_gate = attention.tf_attn_gate; diff --git a/zig/src/hip/launches/moe.zig b/zig/src/hip/launches/moe.zig index 2cccb06e86..366164baab 100644 --- a/zig/src/hip/launches/moe.zig +++ b/zig/src/hip/launches/moe.zig @@ -106,6 +106,19 @@ pub fn tf_moe_combine(l: *const Launcher, y: CF, wts: CF, out: P, kind: c_int, r try l.go(l.moe_combine[tri(kind)], dim(cdiv(d, 256), r, 1), dim(256, 1, 1), 0, s, &a); } +/// The slots' sum plus a shared expert run beside the stack, weighted by wts' last slot; `shared` null adds nothing. +pub fn tf_moe_combine_gate(l: *const Launcher, y: CF, wts: CF, shared: CF, out: P, kind: c_int, r: c_int, slots: c_int, d: c_int, s: S) Error!void { + var a: Args = .{}; + try a.add(ad(y)); + try a.add(ad(wts)); + try a.add(ad(shared)); + try a.add(ad(out)); + try a.add(kind); + try a.add(slots); + try a.add(d); + try l.go(l.moe_combine_gate, dim(cdiv(d, 256), r, 1), dim(256, 1, 1), 0, s, &a); +} + pub fn tf_moe_route(l: *const Launcher, picks: CI, pairs: c_int, experts: c_int, tile: c_int, members: I, items: I, capacity: c_int, s: S) Error!void { var a: Args = .{}; try a.add(ad(picks)); diff --git a/zig/src/hip/moe.zig b/zig/src/hip/moe.zig index 0ddd802a9a..cfe1b43c37 100644 --- a/zig/src/hip/moe.zig +++ b/zig/src/hip/moe.zig @@ -8,14 +8,19 @@ const Tensor = ops.Tensor; const Kind = ops.Kind; const Projection = ops.Projection; -/// A layer's experts, shared expert last: gate and up stacked as one (E + 1, 2 NI, ...) weight, down (E + 1, D, ...). +/// A shared expert of its own width, run beside the routed stack: gate and up as one (2 NI, K) projection, then down. +pub const Shared = struct { gate_up: Projection, down: Projection }; + +/// A layer's experts, shared last in the stack, or beside it in `shared` at its own bits and group (`count` is then E). pub const Experts = struct { fused: Projection, down: Projection, + /// E + 1 when the shared expert is the stack's last expert, E when it is `shared`. count: usize, width: usize, dims: usize, limit: f32 = 0, + shared: ?Shared = null, }; /// The router's rows in fp32 (E + 1, D), the shared expert's gate last, and the experts. @@ -92,21 +97,42 @@ pub fn run(o: Ops, kind: Kind, r: Routed, x: Tensor, rows: usize, prefill: bool) const act_t: Tensor = .{ .ptr = act, .kind = kind }; try o.moeAct(both, act_t, pairs, ex.width, ex.limit); const y = try o.projectRouted(act_t, ex.down, items, count, members, pairs, 1, @min(tile, rows)); + // the shared pair's id is another rank's here: its row reads zero, and rank 0 adds its own gate's branch try o.moeZeroForeign(y, local, ex.count, pairs, ex.dims); const share = try o.arena.of(f32, rows * ex.dims); + if (ex.shared) |sh| { + const s = try sharedRun(o, sh, x, rows, kind, ex.limit); + try o.moeCombineGate(y, wts, s, .{ .ptr = share, .kind = .f32 }, rows, slots, ex.dims); + return .{ .ptr = share, .kind = .f32 }; + } try o.moeCombine(y, wts, .{ .ptr = share, .kind = .f32 }, rows, slots, ex.dims); return .{ .ptr = share, .kind = .f32 }; } const p = try expertParts(o, kind, r, x, rows, logits, pick, wts, items, members, cap); const out = try o.arena.take(rows * ex.dims * kind.size()); + if (ex.shared) |sh| { + const s = try sharedRun(o, sh, x, rows, kind, ex.limit); + try o.moeCombineGate(p.y, wts, s, .{ .ptr = out, .kind = kind }, rows, slots, ex.dims); + return .{ .ptr = out, .kind = kind }; + } try o.moeCombine(p.y, wts, .{ .ptr = out, .kind = kind }, rows, slots, ex.dims); return .{ .ptr = out, .kind = kind }; } -/// The routed slots a row (fp32) and their weights before the sum: a single rank only. +/// down(silu(gate) * up) of `x`'s rows at the shared expert's own widths: (rows, dims) fp32, as the combine reads it. +fn sharedRun(o: Ops, sh: Shared, x: Tensor, rows: usize, kind: Kind, limit: f32) ops.Error!u64 { + const width = sh.gate_up.n / 2; + const both = try o.project(x, sh.gate_up, rows, true); + const act: Tensor = .{ .ptr = try o.arena.take(rows * width * kind.size()), .kind = kind }; + try o.moeAct(both.ptr, act, rows, width, limit); + return (try o.project(act, sh.down, rows, true)).ptr; +} + +/// The routed slots and weights before the sum: one rank, and only where the shared expert stacks in `y`. pub const Parts = struct { y: u64, wts: u64, slots: usize }; pub fn parts(o: Ops, kind: Kind, r: Routed, x: Tensor, rows: usize) ops.Error!Parts { + std.debug.assert(r.experts.shared == null); const slots = r.top_k + 1; const size = if (r.exact) rows else bucket(rows); const cap = capacity(size * slots, r.rows); @@ -136,6 +162,8 @@ fn expertParts(o: Ops, kind: Kind, r: Routed, x: Tensor, rows: usize, logits: u6 tile = tileFor(o, rows); try o.moeRoute(pick, pairs, r.rows, tile, members, items, cap); } + // beside the stack the shared pair's id is the stack's first id past its experts: zero its item, never read its row + if (ex.shared != null) try o.moeForeignItems(items, count, ex.count); // the gate and up's activation is the launch's epilogue when the tile takes the shape const act_t: Tensor = if (try o.projectRoutedAct(x, ex.fused, items, count, members, pairs, slots, @min(tile, rows), ex.limit)) |t| t else blk: { const both = try o.projectRouted(x, ex.fused, items, count, members, pairs, slots, @min(tile, rows)); diff --git a/zig/src/hip/ops/moe.zig b/zig/src/hip/ops/moe.zig index effaebbe81..44eebab57c 100644 --- a/zig/src/hip/ops/moe.zig +++ b/zig/src/hip/ops/moe.zig @@ -40,3 +40,9 @@ pub fn moeAct(o: Ops, both: u64, out: Tensor, pairs: usize, width: usize, limit: pub fn moeCombine(o: Ops, y: u64, wts: u64, out: Tensor, r: usize, slots: usize, d: usize) Error!void { try o.l.tf_moe_combine(f(y), f(wts), p(out.ptr), @backingInt(out.kind), int(r), int(slots), int(d), o.stream); } + +/// The same sum with a shared expert run beside the stack: `shared` (rows, d) fp32, or 0 where it adds nothing. +pub fn moeCombineGate(o: Ops, y: u64, wts: u64, shared: u64, out: Tensor, r: usize, slots: usize, d: usize) Error!void { + if (slots < 2) return error.BadShape; + try o.l.tf_moe_combine_gate(f(y), f(wts), if (shared == 0) null else f(shared), p(out.ptr), @backingInt(out.kind), int(r), int(slots), int(d), o.stream); +} diff --git a/zig/src/hip/ops/ops.zig b/zig/src/hip/ops/ops.zig index 04a17e8410..da652de284 100644 --- a/zig/src/hip/ops/ops.zig +++ b/zig/src/hip/ops/ops.zig @@ -90,6 +90,7 @@ pub const Ops = struct { pub const moeRoute = moe.moeRoute; pub const moeAct = moe.moeAct; pub const moeCombine = moe.moeCombine; + pub const moeCombineGate = moe.moeCombineGate; pub const argmaxRows = draw.argmaxRows; pub const topkRows = draw.topkRows; pub const tokenProb = draw.tokenProb; diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index 5b07bfefb3..642b527fa9 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -16,6 +16,7 @@ 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 gate_k = @import("contract_moe_gate.zig"); const g = core.contract.gates; @@ -136,6 +137,10 @@ 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()); + var gates: [3]gate_k.Gate = undefined; + for (&gates, std.enums.values(gate_k.Which)) |*m, w| m.* = .init(a, &r, s, &l, w); + defer for (&gates) |*m| m.deinit(); + for (&gates) |*m| try kernels.append(a, m.kernel()); var t: g.Tally = .{}; const t0 = ms(io); try g.runAll(a, kernels.items, 0x6b63, &t); diff --git a/zig/tests/hip/contract_moe_gate.zig b/zig/tests/hip/contract_moe_gate.zig new file mode 100644 index 0000000000..6327a35980 --- /dev/null +++ b/zig/tests/hip/contract_moe_gate.zig @@ -0,0 +1,214 @@ +//! tf_moe_combine_gate under the contract on HIP: a shared expert beside the stack, bit-equal to tf_moe_combine. +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 = contract.rows; + +/// The combine's output kind: fp16, bf16 or fp32. +pub const Which = enum { gate_f16, gate_bf16, gate_f32 }; + +pub fn kindOf(w: Which) kit.Kind { + return switch (w) { + .gate_f16 => .f16, + .gate_bf16 => .bf16, + .gate_f32 => .f32, + }; +} + +fn none(_: c.Shape) usize { + return 0; +} + +/// Each routed slot's product and add in order, then the shared expert's product and add. +fn slotChain(s: c.Shape) usize { + return 2 * s.k + 1; +} + +const odd = [_]f64{ std.math.nan(f64), std.math.inf(f64), -std.math.inf(f64) }; +const names = [3][]const u8{ + "hip tf_moe_combine_gate (fp16)", + "hip tf_moe_combine_gate (bf16)", + "hip tf_moe_combine_gate (fp32)", +}; + +/// Flash Next's 2,560-wide hidden at its 10 picks plus a shared expert, a decode row to a prompt chunk. +pub fn spec(which: Which) c.Spec { + const sums = &[_]c.Shape{ + .{ .rows = 1, .n = 2560, .k = 11 }, + .{ .rows = 16, .n = 2560, .k = 11 }, + .{ .rows = 512, .n = 2560, .k = 11 }, + }; + return .{ + .name = names[@backingInt(which)], + .accum = .f32, + .depth = slotChain, + .out = kit.elem(kindOf(which)), + .replaces = "tf_moe_combine", + .tiles = .{ .{}, .{ .size = 256 }, .{} }, + .served = sums, + }; +} + +/// shape.rows rows of shape.k slots (shape.k - 1 routed and the shared expert) over shape.n columns. +pub const Gate = struct { + rig: rig_.Rig, + which: Which, + l: *const hip.Launcher, + k: kit.Kind, + arena: std.heap.ArenaAllocator, + shape: c.Shape = undefined, + x: []f64 = &.{}, // y [rows][slots][n]: the routed slots, the last one never read + w: []f64 = &.{}, // weights [rows][slots], the shared expert's last + shared: []f64 = &.{}, // the shared expert's rows [rows][n] + in: [3]u64 = @splat(0), // y, wts, shared + stack: u64 = 0, // y with the shared expert in its last slot: tf_moe_combine's input + out: u64 = 0, + base: u64 = 0, + bytes: []u8 = &.{}, // y's bytes as the kind's width, so both kernels read the same routed rows + + pub fn init(a: std.mem.Allocator, rt: *const hip.Runtime, s: hip.Stream, l: *const hip.Launcher, w: Which) Gate { + return .{ .rig = .{ .a = a, .r = rt, .s = s }, .which = w, .l = l, .k = kindOf(w), .arena = .init(a) }; + } + + pub fn deinit(self: *Gate) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Gate) g.Kernel { + return .{ .spec = spec(self.which), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Gate); + + pub fn prepare(self: *Gate, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k < 2) 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 s = shape; + self.x = try a.alloc(f64, s.rows * s.k * s.n); + self.w = try a.alloc(f64, s.rows * s.k); + self.shared = try a.alloc(f64, s.rows * s.n); + for (self.x) |*v| v.* = @as(f32, @floatCast(rng.float(f64) * 4 - 2)); + for (self.w) |*v| v.* = @as(f32, @floatCast(rng.float(f64))); + for (self.shared) |*v| v.* = @as(f32, @floatCast(rng.float(f64) * 2 - 1)); + if (fill == .stress) self.stress(rng); + if (fill == .nonfinite) self.plant(); + // y's last slot holds the shared expert for the combine; the gate's y keeps its own there, unread + const stack = try a.dupe(f64, self.x); + for (0..s.rows) |r| for (0..s.n) |col| { + stack[(r * s.k + s.k - 1) * s.n + col] = self.shared[r * s.n + col]; + }; + self.bytes = try kit.f32Bytes(a, self.x); + self.in[0] = try self.rig.input(self.bytes); + self.in[1] = try self.rig.input(try kit.f32Bytes(a, self.w)); + self.in[2] = try self.rig.input(try kit.f32Bytes(a, self.shared)); + self.stack = try self.rig.input(try kit.f32Bytes(a, stack)); + self.out = try self.rig.output(s.rows * s.n * kit.size(self.k)); + self.base = try self.rig.scratch(s.rows * s.n * kit.size(self.k)); + return true; + } + + /// Gate 4: integers whose routed slots undo each other in pairs, so only fp32 sums hold the bound. + fn stress(self: *Gate, rng: std.Random) void { + const s = self.shape; + const half = (s.k - 1) / 2; + for (0..s.rows) |r| { + // weights repeat every `half` slots, so slot j and slot j + half cancel at the same weight + for (0..s.k) |j| self.w[r * s.k + j] = @floatFromInt(200 + j % half); + for (0..s.n) |col| { + self.shared[r * s.n + col] = @floatFromInt(rng.uintLessThan(u32, 8)); + for (0..s.k - 1) |j| { + const y = &self.x[(r * s.k + j) * s.n + col]; + const nudge: f64 = if (j == half) 1 else 0; + const big: f64 = @floatFromInt(500 + rng.uintLessThan(u32, 600)); + y.* = if (j < half) big else nudge - self.x[(r * s.k + j - half) * s.n + col]; + } + } + } + } + + /// NaN in a routed slot of row 0, +Inf and -Inf in two more, and an infinite shared row in the last row. + fn plant(self: *Gate) void { + const s = self.shape; + const at = [_]usize{ 0, s.n + 3, 2 * s.n + 9 }; // (slot * n + column): all of row 0's three slots + for (odd, at) |v, i| self.x[i] = v; + const last = self.shared.len - 1; + self.shared[last] = std.math.inf(f64); + } + + pub fn run(self: *Gate) !void { + const s = self.shape; + try self.l.tf_moe_combine_gate( + kit.ptr(?[*]const f32, self.in[0]), + kit.ptr(?[*]const f32, self.in[1]), + kit.ptr(?[*]const f32, self.in[2]), + kit.ptr(?*anyopaque, self.out), + @backingInt(self.k), + @intCast(s.rows), + @intCast(s.k), + @intCast(s.n), + self.rig.s.handle, + ); + try self.rig.s.synchronize(); + } + + /// tf_moe_combine on the same rows with the shared expert stacked as y's last slot. + pub fn replaced(self: *Gate) ![]const u8 { + const s = self.shape; + try self.l.tf_moe_combine( + kit.ptr(?[*]const f32, self.stack), + kit.ptr(?[*]const f32, self.in[1]), + kit.ptr(?*anyopaque, self.base), + @backingInt(self.k), + @intCast(s.rows), + @intCast(s.k), + @intCast(s.n), + self.rig.s.handle, + ); + try self.rig.s.synchronize(); + return self.rig.bytes(self.base, s.rows * s.n * kit.size(self.k)); + } + + pub fn count(self: *Gate) usize { + return self.shape.rows * self.shape.n; + } + + /// The routed slots' sum plus the shared expert's, in float64, with the terms' magnitude for the bound. + pub fn reference(self: *Gate, want: []f64, mag: []f64) !void { + const s = self.shape; + @memset(mag, 0); + for (0..s.rows) |r| for (0..s.n) |col| { + var acc: f64 = 0; + var abs: f64 = 0; + for (0..s.k - 1) |j| { + const t = self.x[(r * s.k + j) * s.n + col] * self.w[r * s.k + j]; + acc += t; + abs += @abs(t); + } + const sh = self.shared[r * s.n + col] * self.w[r * s.k + s.k - 1]; + want[r * s.n + col] = acc + sh; + mag[r * s.n + col] = abs + @abs(sh); + }; + } + + pub fn simulate(self: *Gate, acc: c.Accum, order: g.Order, out: []f64) !void { + const s = self.shape; + var t: [16]f64 = undefined; + for (0..s.rows) |r| for (0..s.n) |col| { + for (0..s.k - 1) |j| t[j] = self.x[(r * s.k + j) * s.n + col] * self.w[r * s.k + j]; + t[s.k - 1] = self.shared[r * s.n + col] * self.w[r * s.k + s.k - 1]; + out[r * s.n + col] = rows.sumIn(acc, order, t[0..s.k]); + }; + } +};