hip: load a shared expert wider than the routed experts beside the stack

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<long long>(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<long long>(row) * slots + s) * d + c] * wts[row * slots + s];
+    }
+    if (shared != nullptr) acc += shared[static_cast<long long>(row) * d + c] * wts[row * slots + slots - 1];
+    const long long at = static_cast<long long>(row) * d + c;
+    if (kind == 0) narrow(acc, static_cast<float*>(out) + at);
+    else if (kind == 1) narrow(acc, static_cast<__half*>(out) + at);
+    else narrow(acc, reinterpret_cast<hip_bfloat16*>(out) + at);
+}
+
 template <typename T>
 __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,

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 = .{

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,
             },
         };
     }

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;

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"),

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));
+}

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");

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
+    }
+}

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;

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));

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));