cuda: the staged NVFP4 prompt entries leave an expert's items alone (skip), as the decode form does

zig/kernels/cuda/nvfp4_experts.cu

@@ -291,11 +291,12 @@ __device__ __forceinline__ void prompt_body(const __nv_bfloat16* __restrict__ X,
                                             const uint4* __restrict__ W, const float* __restrict__ scale, int KG,
                                             int NB, const int* __restrict__ items, const int* __restrict__ counts,
                                             const int* __restrict__ members, void* __restrict__ out, int N,
-                                            float limit) {
+                                            float limit, int skip) {
   static_assert(RT == 4, "the dispatch below spells out 0 .. 4 row tiles");
   const int nbt = (NB + WN - 1) / WN, it = blockIdx.x / nbt, cbt = blockIdx.x - it * nbt;
   if (it >= __ldg(counts)) return;
   const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2);
+  if (e == skip) return;  // an expert the family computes itself, left alone as the decode form does
   const int nt = max(0, min(RT, (cnt - 16 * RT * ((threadIdx.x >> 5) / WN) + 15) >> 4));
 #define TF_TILES(NT) \
   prompt_tiles<M, EPI, RT, WM, WN, NT>(X, x_stride, slots, W, scale, KG, NB, members, out, N, limit, cbt, e, first, cnt)
@@ -321,9 +322,9 @@ template __global__ void nvfp4_expert_kernel<1, 1, 4>(const __nv_bfloat16* __res
   extern "C" __global__ void __launch_bounds__(WN * 32) NAME(                                                        \
       const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W,                     \
       const float* __restrict__ scale, int KG, int NB, const int* __restrict__ items, const int* __restrict__ counts, \
-      const int* __restrict__ members, void* __restrict__ out, int N, float limit) {                                 \
+      const int* __restrict__ members, void* __restrict__ out, int N, float limit, int skip) {                       \
     tf_nvfp4_experts::prompt_body<M, EPI, 4, 1, WN>(X, x_stride, slots, W, scale, KG, NB, items, counts, members,    \
-                                                    out, N, limit);                                                  \
+                                                    out, N, limit, skip);                                            \
   }
 TF_NVFP4_PROMPT(tf_nvfp4_prompt_relu2, 1, 1, 8)
 TF_NVFP4_PROMPT(tf_nvfp4_prompt_swiglu, 2, 2, 4)

zig/src/cuda/experts.zig

@@ -341,7 +341,6 @@ pub const Nvfp4Experts = struct {
     /// Prompt rows on a plan of at most `items` items of 64 pairs: up (relu^2 or SwiGLU by the layer), bf16 sums down.
     pub fn runPrompt(e: *const Nvfp4Experts, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void {
         if (epi == .down_f32) return error.EpilogueNotBuilt;
-        if (skip >= 0) return error.OptionNotBuilt;
         const gate = epi == .up;
         const parts: usize = if (gate and !l.relu2) 2 else 1;
         const nb = (if (gate) l.width else l.dims) / cols;
@@ -357,6 +356,7 @@ pub const Nvfp4Experts = struct {
         for ([_]u64{ p.items, p.counts, p.members, out }) |v| a.add(v);
         a.add(@as(i32, @intCast(if (gate) l.width else l.dims)));
         a.add(if (gate) l.limit else @as(f32, 0.0));
+        a.add(skip);
         const fi: usize = if (!gate) 2 else if (l.relu2) 0 else 1;
         const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(items * ((nb + prompt_wn[fi] - 1) / prompt_wn[fi])), .y = 1, .z = 1 }, .block = .{ .x = @intCast(prompt_wn[fi] * 32), .y = 1, .z = 1 }, .shared = promptSmem(parts, prompt_wn[fi]) };
         try launch_.launch(e.prompt_fns[fi], cfg, s, &a);

zig/tests/cuda/quant_ref_experts.zig

@@ -120,6 +120,41 @@ fn batch(r: *Rig, comptime label: []const u8, g: experts.Grouped, router: groupe
     return .{ .act = act, .y = y };
 }
 
+/// Prompt items of expert `skip` (the shared one, as a family that computes it itself passes) are left alone.
+fn skipped(r: *Rig, comptime label: []const u8, g: experts.Grouped, router: grouped.Router, l: experts.Layer, x: []const u16, rows: usize) !void {
+    const pairs = rows * slots;
+    const pk = try picks(r, rows);
+    const pl = try plan(r, router, try r.dev(std.mem.sliceAsBytes(pk)), pairs, experts.promptTile(l.format));
+    const fill = try r.a.alloc(u8, pairs * @max(NI, D) * 2);
+    @memset(fill, 0xA5);
+    const skip: i32 = E - 1;
+    const up_in: experts.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(x[0 .. rows * D])), .stride = D, .slots = slots };
+    const a0 = try r.zeros(pairs * NI * 2);
+    const a1 = try r.dev(fill[0 .. pairs * NI * 2]);
+    try g.prompt(r.s, .up, up_in, l, pl.p, a0, pl.items, -1);
+    try g.prompt(r.s, .up, up_in, l, pl.p, a1, pl.items, skip);
+    try leftAlone(label ++ " prompt gate-up", pk, skip, try r.back(a0, pairs * NI * 2), try r.back(a1, pairs * NI * 2), NI * 2, fill);
+    const down_in: experts.Rows = .{ .x = a0, .stride = NI };
+    const y0 = try r.zeros(pairs * D * 2);
+    const y1 = try r.dev(fill[0 .. pairs * D * 2]);
+    try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y0, pl.items, -1);
+    try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y1, pl.items, skip);
+    try leftAlone(label ++ " prompt down", pk, skip, try r.back(y0, pairs * D * 2), try r.back(y1, pairs * D * 2), D * 2, fill);
+}
+
+/// Each pair's row: the skipped expert's still the fill, every other one equal to the run without skip.
+fn leftAlone(name: []const u8, pk: []const i32, skip: i32, plain: []const u8, got: []const u8, row: usize, fill: []const u8) !void {
+    var left: usize = 0;
+    var differ: usize = 0;
+    for (pk, 0..) |e, p| {
+        const b = got[p * row ..][0..row];
+        if (e == skip) {
+            if (std.mem.eql(u8, b, fill[0..row])) left += 1 else differ += 1;
+        } else if (!std.mem.eql(u8, plain[p * row ..][0..row], b)) differ += 1;
+    }
+    try check.expect(differ == 0 and left > 0, "{s}: expert {d}'s {d} pairs left alone, every other pair as without skip ({d} differ)", .{ name, skip, left, differ });
+}
+
 const Fp8M = struct { codes: []const u8, inv: []const f32, n: usize, k: usize };
 fn fp8Val(c: Fp8M, e: usize, row: usize, col: usize) f64 {
     const nb = c.n / 128;
@@ -181,6 +216,7 @@ pub fn run(r: *Rig) !void {
         const M = Matrix(Fp8M);
         for ([_]usize{ 1, 3 }) |rows| _ = try batch(r, "fp8", g, router, l, x, rows, false, M{ .ctx = gm, .value = fp8Val }, M{ .ctx = um, .value = fp8Val }, M{ .ctx = dm, .value = fp8Val });
         _ = try batch(r, "fp8", g, router, l, x, 300, true, M{ .ctx = gm, .value = fp8Val }, M{ .ctx = um, .value = fp8Val }, M{ .ctx = dm, .value = fp8Val });
+        try skipped(r, "fp8", g, router, l, x, 300);
     }
     // NVFP4
     {
@@ -207,6 +243,7 @@ pub fn run(r: *Rig) !void {
         const M = Matrix(Fp4M);
         for ([_]usize{ 1, 3, 40 }) |rows| _ = try batch(r, "nvfp4", g, router, l, x, rows, false, M{ .ctx = gm, .value = fp4Val }, M{ .ctx = um, .value = fp4Val }, M{ .ctx = dm, .value = fp4Val });
         _ = try batch(r, "nvfp4", g, router, l, x, 300, true, M{ .ctx = gm, .value = fp4Val }, M{ .ctx = um, .value = fp4Val }, M{ .ctx = dm, .value = fp4Val });
+        try skipped(r, "nvfp4", g, router, l, x, 300);
     }
-    check.pass("experts fp8g and nvfp4 (decode, prompt): gate-up and down within float64 tolerance", .{});
+    check.pass("experts fp8g and nvfp4 (decode, prompt): gate-up and down within float64 tolerance; a skipped expert's prompt items left alone", .{});
 }