From da2373648a4bad2ef524e1e50c83a2361f145a69 Mon Sep 17 00:00:00 2001 From: gilby Date: Thu, 29 Oct 2026 17:37:06 +0100 Subject: [PATCH 1/5] hip: load a shared expert wider than the routed experts beside the stack The fused routed+shared expert stack needs one bit width and group size, so a checkpoint whose shared expert is wider (oQ: 8-bit g64 over 4-bit g32 routed experts) was refused. When the widths differ, the shared expert now loads as its own projections at its own width and runs beside the routed launch; tf_moe_combine_gate adds it inside the same fp32 slot accumulation, weighted by the router's shared gate. Checkpoints whose widths match keep the fused stack, unchanged. Fixes #1599 --- zig/kernels/hip/ops/act.hip | 18 + zig/src/core/quant/experts.zig | 326 ++++++++++++++++-- zig/src/families/flashnext/hip/load.zig | 5 + .../qwen3_5/backend/hip/forward/window.zig | 8 +- zig/src/hip/launches.zig | 4 + zig/src/hip/launches/moe.zig | 13 + zig/src/hip/moe.zig | 32 +- zig/src/hip/ops/moe.zig | 6 + zig/src/hip/ops/ops.zig | 1 + 9 files changed, 375 insertions(+), 38 deletions(-) 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/quant/experts.zig b/zig/src/core/quant/experts.zig index 4c76cfa10a..e1cce3380d 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,89 @@ 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 +225,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 +254,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 +307,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 +320,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 +399,157 @@ 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/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; From c05f48f706357cc4664483f2c14046a5d6bb0652 Mon Sep 17 00:00:00 2001 From: gilby Date: Fri, 30 Oct 2026 15:02:28 +0100 Subject: [PATCH 2/5] zig fmt --- zig/src/core/quant/experts.zig | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/zig/src/core/quant/experts.zig b/zig/src/core/quant/experts.zig index e1cce3380d..12d95715ad 100644 --- a/zig/src/core/quant/experts.zig +++ b/zig/src/core/quant/experts.zig @@ -142,7 +142,6 @@ fn part(t: anytype, stack_key: []const u8, shared: []const u8, w: mlx.Width, sam return p; } - /// 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; @@ -441,8 +440,7 @@ fn affineInto( @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) }); + 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); } @@ -462,12 +460,10 @@ fn routedInto(a: Allocator, f: *Mixed, w: mlx.Width) !void { 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) }); + 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) }); + 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. From 974e039c033e9da3303e2339e28ed543ec795775 Mon Sep 17 00:00:00 2001 From: gilby Date: Fri, 30 Oct 2026 16:53:09 +0100 Subject: [PATCH 3/5] hip: register tf_moe_combine_gate with the kernel contract Contract test at all three output kinds, bit-equal to tf_moe_combine with the shared expert as y's last slot, and usage records for every served chip. --- zig/src/core/contract/hip_records.zig | 14 ++ zig/tests/hip/contract_main.zig | 5 + zig/tests/hip/contract_moe_gate.zig | 214 ++++++++++++++++++++++++++ 3 files changed, 233 insertions(+) create mode 100644 zig/tests/hip/contract_moe_gate.zig 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/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]); + }; + } +}; From 51e844738d5d6a16d89e9216417824ccf2d4e604 Mon Sep 17 00:00:00 2001 From: gilby Date: Fri, 30 Oct 2026 17:38:36 +0100 Subject: [PATCH 4/5] qwen3_5: refuse a shared expert the HIP index leaves beside the routed stack core.quant.experts now stacks a same-width shared expert into the routed stack and leaves a wider one beside it as separate. This bridge always stacks, so a separate shared expert would be dropped and the router's last pick would index past the stack. Refuse it at the HIP host index, where both Bridge.init and fromCheckpointRank consume the stack, and test that a mixed-width qwen3_5-shaped expert set is refused. --- .../families/qwen3_5/weights/checkpoint.zig | 14 ++++- .../qwen3_5/weights/checkpoint_test.zig | 58 +++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) 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)); +} From 8f17d1f649f5cbacc3b5b1da610c2a990b6382e0 Mon Sep 17 00:00:00 2001 From: gilby Date: Fri, 30 Oct 2026 18:28:32 +0100 Subject: [PATCH 5/5] hip: test tf_moe_combine_gate against tf_moe_combine byte for byte Place the shared expert's row as the combine kernel's last slot and compare the two kernels' outputs at fp32, fp16 and bf16. The gate's own y keeps a poison value in that slot, so a gate that summed every slot instead of reading the shared operand could not match. Decode rows, a width that does not fill a block's columns, and a single routed slot beside the shared one. --- zig/src/hip/kernel_tests.zig | 1 + zig/src/hip/kernel_tests/moe_gate.zig | 82 +++++++++++++++++++++++++++ 2 files changed, 83 insertions(+) create mode 100644 zig/src/hip/kernel_tests/moe_gate.zig 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 + } +}