From 3a576e32520b81d6ba708e2c3b5951ec3257ef24 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 11 Oct 2026 09:36:20 +0900 Subject: [PATCH] cuda: the staged NVFP4 prompt entries leave an expert's items alone (skip), as the decode form does A family that computes one expert itself (GLM-5.3-Flash stacks its BF16 shared expert as the last expert id) plans every pair together and has the grouped kernels skip that id. The decode form already takes skip; Grouped.prompt refused it with OptionNotBuilt for NVFP4. The prompt entries now take skip and return before staging any rows of that expert's items. Nemotron passes -1, so its bits do not change. quant-ref now checks it: with the shared expert skipped, its pairs' prompt rows keep what was there and every other pair's row equals the run without skip, for fp8 and nvfp4. --- zig/kernels/cuda/nvfp4_experts.cu | 7 ++--- zig/src/cuda/experts.zig | 2 +- zig/tests/cuda/quant_ref_experts.zig | 39 +++++++++++++++++++++++++++- 3 files changed, 43 insertions(+), 5 deletions(-) diff --git a/zig/kernels/cuda/nvfp4_experts.cu b/zig/kernels/cuda/nvfp4_experts.cu index 0a92b042f4..be724dead8 100644 --- a/zig/kernels/cuda/nvfp4_experts.cu +++ b/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(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(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) diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index c940202bea..0f8c06aa6f 100644 --- a/zig/src/cuda/experts.zig +++ b/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); diff --git a/zig/tests/cuda/quant_ref_experts.zig b/zig/tests/cuda/quant_ref_experts.zig index 4b9dad472d..db4beeceeb 100644 --- a/zig/tests/cuda/quant_ref_experts.zig +++ b/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", .{}); }