From 3ca1e303c75590ba7ec0df767410d10f930e48e5 Mon Sep 17 00:00:00 2001 From: "Riel St. Amand" <54451475+CanadaApollo6@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:53:19 -0400 Subject: [PATCH] experts3 CUDA: a split decode pair form, a column block a CTA with its K groups shared by the warps For calls with few column blocks (one to three token rows) the pair kernel leaves most warps idle. The split form gives one CTA a column block: its warps compute the per-group mma terms side by side, then each warp runs its tile's fma chain over them in group order, so the sums are expert_kernel's bits. run() takes it below a per-GPU bound set from the pair kernel's resident warps; down and the prompt kernels are unchanged. --- zig/kernels/cuda/experts3.cu | 140 ++++++++++++++++++++++++++++- zig/src/core/contract/registry.zig | 1 + zig/src/cuda/experts3.zig | 13 ++- zig/tests/cuda/experts3_tests.zig | 110 ++++++++++++++++++++--- zig/tests/cuda/main.zig | 4 +- 5 files changed, 253 insertions(+), 15 deletions(-) diff --git a/zig/kernels/cuda/experts3.cu b/zig/kernels/cuda/experts3.cu index b9ec5b243f..8ff2f1a87c 100644 --- a/zig/kernels/cuda/experts3.cu +++ b/zig/kernels/cuda/experts3.cu @@ -161,6 +161,141 @@ __global__ void __launch_bounds__(WARPS * 32) } } +// One group's terms for the chain: a lane's 4 mma steps a tile, its scale and bias words, its rows' input sums. +template +struct Part { + float p[M][NTW][4][32]; + uint32_t sb[M][NTW][2][32]; + float xs[2][32]; +}; + +template +__device__ __forceinline__ void part_stage(Part& out, const Stage& st, int lane) { + out.xs[0][lane] = group_sum(st.xa); + out.xs[1][lane] = group_sum(st.xb); +#pragma unroll + for (int m = 0; m < M; ++m) { + float p[NTW][4]; +#pragma unroll + for (int j = 0; j < NTW; ++j) p[j][0] = p[j][1] = p[j][2] = p[j][3] = 0.f; +#pragma unroll + for (int ks = 0; ks < Geo<64>::KS; ++ks) { + const uint32_t a0 = comp(st.xa[ks >> 1], 2 * (ks & 1)); + const uint32_t a2 = comp(st.xa[ks >> 1], 2 * (ks & 1) + 1); + const uint32_t a1 = comp(st.xb[ks >> 1], 2 * (ks & 1)); + const uint32_t a3 = comp(st.xb[ks >> 1], 2 * (ks & 1) + 1); + const int sh = (ks & 1) ? 8 : 0; +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const uint32_t word = nibword(st.w[m], j * 2 + (ks >> 1)); + mma(p[j], a0, a1, a2, a3, nib2(word, sh), nib2(word, sh + 4)); + } + } +#pragma unroll + for (int j = 0; j < NTW; ++j) { +#pragma unroll + for (int r = 0; r < 4; ++r) out.p[m][j][r][lane] = p[j][r]; + out.sb[m][j][0][lane] = comp(st.sb[m][j >> 2], j & 3); + out.sb[m][j][1][lane] = comp(st.sb[m][(NTW + j) >> 2], (NTW + j) & 3); + } + } +} + +// Tile j's chain over one group's terms, as compute_stage runs it. +template +__device__ __forceinline__ void chain(float (&acc)[M][4], const Part& in, int j, int lane, bool hi) { + const float xsa = in.xs[0][lane], xsb = in.xs[1][lane]; +#pragma unroll + for (int m = 0; m < M; ++m) { + const uint32_t sp = in.sb[m][j][0][lane], bp = in.sb[m][j][1][lane]; + const float s0 = lo_f(sp), s1 = hi_f(sp), b0 = lo_f(bp), b1 = hi_f(bp); + float(&a)[4] = acc[m]; + a[0] = fmaf(xsa, b0, fmaf(in.p[m][j][0][lane], s0, a[0])); + a[1] = fmaf(xsa, b1, fmaf(in.p[m][j][1][lane], s1, a[1])); + if (hi) { + a[2] = fmaf(xsb, b0, fmaf(in.p[m][j][2][lane], s0, a[2])); + a[3] = fmaf(xsb, b1, fmaf(in.p[m][j][3][lane], s1, a[3])); + } + } +} + +// Warp w takes groups w, w + W, ...; after each round of W groups, warp j chains tile j over them in group order. +template +__device__ __forceinline__ void split_loop(float (&acc)[M][4], Part (&parts)[W], const uint4* blk, int KG, int w, + int lane, int t, const __nv_bfloat16* x0, const __nv_bfloat16* x1, bool v0, + bool v1) { + Stage st[D]; +#pragma unroll + for (int d = 0; d < D; ++d) + if (d * W + w < KG) load_stage(st[d], blk, d * W + w, lane, t, x0, x1, v0, v1); + const int rounds = (KG + W - 1) / W; + for (int i0 = 0; i0 < rounds; i0 += D) { +#pragma unroll + for (int d = 0; d < D; ++d) { + const int i = i0 + d; + if (i < rounds) { + const int g = i * W + w; + if (g < KG) { + part_stage(parts[w], st[d], lane); + if (g + D * W < KG) load_stage(st[d], blk, g + D * W, lane, t, x0, x1, v0, v1); + } + __syncthreads(); + const int n = min(W, KG - i * W); + for (int s = 0; s < n; ++s) chain(acc, parts[s], w, lane, v1); + __syncthreads(); + } + } + } +} + +// epilogue's SwiGLU stores for tile j alone. +template +__device__ __forceinline__ void store_tile(const float (&acc)[M][4], __nv_bfloat16* out, int N, int col, int pr0, + int pr1, bool v0, bool v1, float limit) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + if (!(h ? v1 : v0)) continue; + const size_t row = (size_t)(h ? pr1 : pr0) * N; + const float o0 = swiglu(acc[0][2 * h], acc[M - 1][2 * h], limit); + const float o1 = swiglu(acc[0][2 * h + 1], acc[M - 1][2 * h + 1], limit); + *reinterpret_cast<__nv_bfloat162*>(out + row + col) = __floats2bfloat162_rn(o0, o1); + } +} + +// expert_kernel's sums with a column block a CTA, its groups shared by the warps: for few blocks, as at one row. +template +__global__ void __launch_bounds__(WARPS * 32) + split_kernel(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, int KG, + int NB, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, void* __restrict__ out, int N, float limit, + const uint4* __restrict__ W4, int e4) { + static_assert(WARPS == NTW && M == 2 && EPI == 2, "the SwiGLU pair, a warp a tile in the chain"); + __shared__ Part parts[WARPS]; + const int lane = threadIdx.x & 31, gq = lane >> 2, t = lane & 3, w = threadIdx.x >> 5; + const int units = __ldg(counts) * NB; + for (int unit = blockIdx.x; unit < units; unit += gridDim.x) { + const int it = unit / NB, cb = unit - it * NB; + const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + const bool v0 = gq < cnt, v1 = gq + 8 < cnt; + const int pr0 = v0 ? __ldg(members + first + gq) : 0; + const int pr1 = v1 ? __ldg(members + first + gq + 8) : 0; + const int r0 = slots ? pr0 / slots : pr0, r1 = slots ? pr1 / slots : pr1; + const __nv_bfloat16* x0 = X + (size_t)r0 * x_stride + t * 16; + const __nv_bfloat16* x1 = X + (size_t)r1 * x_stride + t * 16; + float acc[M][4]; +#pragma unroll + for (int m = 0; m < M; ++m) acc[m][0] = acc[m][1] = acc[m][2] = acc[m][3] = 0.f; + if (e == e4) { // block-uniform: a unit is one expert's column block + const uint4* blk = W4 + (size_t)cb * KG * (M * Lane<4>::BLOCK); + split_loop(acc, parts, blk, KG, w, lane, t, x0, x1, v0, v1); + } else { + const uint4* blk = W + ((size_t)e * NB + cb) * (size_t)KG * (M * Lane<3>::BLOCK); + split_loop(acc, parts, blk, KG, w, lane, t, x0, x1, v0, v1); + } + store_tile(acc, reinterpret_cast<__nv_bfloat16*>(out), N, cb * COLS + 2 * t + 8 * w, pr0, pr1, v0, v1, limit); + } +} + } // namespace tf_experts3 // MLX 3-bit words (an LSB-first bitstream, 3 bytes a run of 8 codes) to MLX 4-bit words: a thread an output word. @@ -200,10 +335,13 @@ extern "C" __global__ void tf_experts3_planes(const uint32_t* __restrict__ in, u } } -// Decode form for groups of 64: the SwiGLU pair form (M = 2, epilogue 2) and down to fp32 (epilogue 0). +// Decode form for groups of 64: the SwiGLU pair form (M = 2, epilogue 2) and its split form, down to fp32 (epilogue 0). template __global__ void tf_experts3::expert_kernel<2, 2, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, const uint4* __restrict__, int); template __global__ void tf_experts3::expert_kernel<1, 0, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, const uint4* __restrict__, int); +template __global__ void tf_experts3::split_kernel<2, 2, 2, 4>(const __nv_bfloat16* __restrict__, int, int, + const uint4* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, + void* __restrict__, int, float, const uint4* __restrict__, int); diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 2f307c6753..6a1bb15dfd 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -217,6 +217,7 @@ pub const cuda = [_]CudaRecord{ // Affine-3 experts: the decode pair and down, the prompt pair and down at their shared bytes, the repacks .{ .launch = .{ .name = "experts3 decode pair", .symbol = "_ZN11tf_experts313expert_kernelILi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3", .threads = 128 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "experts3 decode down", .symbol = "_ZN11tf_experts313expert_kernelILi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3", .threads = 128 }, .usage = .{ .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 197, .static = 0, .bound = 128 }, .{ .regs = 197, .static = 0, .bound = 128 }, .{ .regs = 190, .static = 0, .bound = 128 }, .{ .regs = 190, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts3 decode split", .symbol = "_ZN11tf_experts312split_kernelILi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3", .threads = 128 }, .usage = .{ .{ .regs = 254, .static = 25600, .bound = 128 }, .{ .regs = 255, .static = 25600, .bound = 128 }, .{ .regs = 255, .static = 25600, .bound = 128 }, .{ .regs = 255, .static = 25600, .bound = 128 }, .{ .regs = 255, .static = 25600, .bound = 128 } } }, .{ .launch = .{ .name = "experts3_prefill pair", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi2ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 46080, .opt_in = true }, .usage = .{ .{ .regs = 179, .static = 0, .bound = 256 }, .{ .regs = 179, .static = 0, .bound = 256 }, .{ .regs = 179, .static = 0, .bound = 256 }, .{ .regs = 146, .static = 0, .bound = 256 }, .{ .regs = 146, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "experts3_prefill down", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 47104, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } }, diff --git a/zig/src/cuda/experts3.zig b/zig/src/cuda/experts3.zig index 7ae0e3bb72..25fc233a1f 100644 --- a/zig/src/cuda/experts3.zig +++ b/zig/src/cuda/experts3.zig @@ -18,12 +18,15 @@ fn int(x: usize) c_int { pub const Affine3Experts = struct { pair_fn: Function, + split_fn: Function, // the pair form with a column block a CTA and its groups shared by the CTA's warps down_fn: Function, pre_pair_fn: Function, pre_down_fn: Function, widen_fn: Function, planes_fn: Function, resident: [2]usize, // blocks the decode pair and down kernels fill: per SM times SMs + split_resident: usize, // blocks split_fn fills + split_below: usize, // pair column blocks below which split_fn runs: 7/8 of the warps pair_fn keeps resident pub const words = 224; // uint32 a 3-bit block: a 2-bit plane (128), a 1-bit plane (64), then scales and biases (32) /// The prompt kernels' dynamic shared bytes: the larger of the 3-bit and 4-bit (shared expert) stage sets. @@ -32,6 +35,7 @@ pub const Affine3Experts = struct { pub const symbols = struct { pub const pair = "_ZN11tf_experts313expert_kernelILi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i"; + pub const split = "_ZN11tf_experts312split_kernelILi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i"; pub const down = "_ZN11tf_experts313expert_kernelILi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i"; pub const pre_pair = "_ZN19tf_experts3_prefill14prefill_kernelILi2ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i"; pub const pre_down = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i"; @@ -43,6 +47,7 @@ pub const Affine3Experts = struct { pub fn resolve(decode_mod: Module, prompt_mod: Module, sms: usize) !Affine3Experts { var a: Affine3Experts = undefined; a.pair_fn = try decode_mod.function(symbols.pair); + a.split_fn = try decode_mod.function(symbols.split); a.down_fn = try decode_mod.function(symbols.down); a.widen_fn = try decode_mod.function(symbols.widen); a.planes_fn = try decode_mod.function(symbols.planes); @@ -51,6 +56,8 @@ pub const Affine3Experts = struct { try a.pre_pair_fn.allowDynamicShared(pair_prompt_smem); try a.pre_down_fn.allowDynamicShared(prompt_smem); a.resident = .{ @max(1, try a.pair_fn.occupancy(128, 0)) * sms, @max(1, try a.down_fn.occupancy(128, 0)) * sms }; + a.split_resident = @max(1, try a.split_fn.occupancy(128, 0)) * sms; + a.split_below = 4 * a.resident[0] * 7 / 8; return a; } @@ -73,11 +80,13 @@ pub const Affine3Experts = struct { /// Decode rows: `pair` takes token rows through SwiGLU over the pair tables, else pair rows down to fp32 sums. pub fn run(e: *const Affine3Experts, s: Stream, pair: bool, in: experts.Rows, l: experts.Layer, p: Plan, out: u64, items: usize) !void { const v = view(l, pair); - const grid = @min((items * v.nb + 3) / 4, e.resident[@intFromBool(!pair)]); + const units = items * v.nb; + const split = pair and units < e.split_below; + const grid = if (split) @min(units, e.split_resident) else @min((units + 3) / 4, e.resident[@intFromBool(!pair)]); if (grid < 1) return; var a = args(in.x, in.stride, if (pair) in.slots else 0, v.w, v.kg, v.nb, p, out, v.n, v.w4, l.shared); const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128 }, .shared = 0 }; - try launch_.launch(if (pair) e.pair_fn else e.down_fn, cfg, s, &a); + try launch_.launch(if (split) e.split_fn else if (pair) e.pair_fn else e.down_fn, cfg, s, &a); } /// Prompt rows, 64 pairs x 128 columns a CTA: `pair` SwiGLU over the pair tables, else down to bf16 sums. diff --git a/zig/tests/cuda/experts3_tests.zig b/zig/tests/cuda/experts3_tests.zig index d845fa354b..5089b8991c 100644 --- a/zig/tests/cuda/experts3_tests.zig +++ b/zig/tests/cuda/experts3_tests.zig @@ -124,6 +124,10 @@ fn tables(gpu: Gpu, k: Kernels, s: cuda.Stream, parts: []const [2]Codes) !Tables } const Shape = struct { experts: usize = 6, ni: usize = 128, dims: usize = 256, rows: usize = 40, slots: usize = 3 }; +/// Qwen3.5-35B-A3B's expert widths at a verify window of 4 rows: 8 routed picks and the shared expert a row. +const wide: Shape = .{ .experts = 16, .ni = 512, .dims = 2048, .rows = 4, .slots = 9 }; +/// 5 K groups: the split form's last round has fewer groups than warps. +const short: Shape = .{ .dims = 320 }; /// A plan for `picks` at `tile`, on buffers that live as long as the returned Plan's owner. const PlanBufs = struct { @@ -153,9 +157,8 @@ const PlanBufs = struct { }; /// The SwiGLU pair form and down, decode and prompt rows: the 3-bit tables give the 4-bit tables' bytes. -fn sameBits(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random) !void { +fn sameBits(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random, sh: Shape) !void { const a = gpu.gpa; - const sh: Shape = .{}; const e = sh.experts; var cs: [3][2]Codes = undefined; for (&cs, [_][2]usize{ .{ sh.ni, sh.dims }, .{ sh.ni, sh.dims }, .{ sh.dims, sh.ni } }) |*c, nk| c.* = .{ try Codes.init(a, r, e, nk[0], nk[1], 3), try Codes.init(a, r, 1, nk[0], nk[1], 4) }; @@ -167,11 +170,11 @@ fn sameBits(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random) !void { const pairs = sh.rows * sh.slots; const picks = try a.alloc(i32, pairs); defer a.free(picks); - for (0..sh.rows) |row| { - picks[row * sh.slots] = @intCast(r.uintLessThan(usize, e)); - picks[row * sh.slots + 1] = @intCast(if (row % 4 == 0) 0 else r.uintLessThan(usize, e)); - picks[row * sh.slots + 2] = @intCast(e); // the shared expert, as the engine's last slot - } + for (0..sh.rows) |row| for (0..sh.slots) |i| { + // the shared expert in the last slot, as the engine places it; expert 0 in slot 1 of every fourth row + const pick = if (i == sh.slots - 1) e else if (i == 1 and row % 4 == 0) 0 else r.uintLessThan(usize, e); + picks[row * sh.slots + i] = @intCast(pick); + }; const x = try a.alloc(u16, sh.rows * sh.dims); defer a.free(x); for (x) |*v| v.* = randBf16(r, 1.0); @@ -206,7 +209,8 @@ fn sameBits(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random) !void { try check.sameBytes(if (prompt) "prompt SwiGLU rows" else "decode SwiGLU rows", outs[1][0], outs[0][0]); const used = pairs * sh.dims * @as(usize, if (prompt) 2 else 4); try check.sameBytes(if (prompt) "prompt down sums" else "decode down sums", outs[1][1][0..used], outs[0][1][0..used]); - check.pass("experts3 {s}: {d} pairs over {d} 3-bit experts and a 4-bit shared one, SwiGLU rows and down sums byte-equal to affine-4 on the same codes", .{ form, pairs, e }); + const pair_form = if (prompt) "" else if (k.a3.split_below == 0) ", the pair a column block a warp" else ", the pair a column block a CTA"; + check.pass("experts3 {s} ({d} by {d}{s}): {d} pairs over {d} 3-bit experts and a 4-bit shared one, SwiGLU rows and down sums byte-equal to affine-4 on the same codes", .{ form, sh.ni, sh.dims, pair_form, pairs, e }); } } @@ -272,7 +276,85 @@ fn projections(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random) !void { } } -pub fn run(gpu: Gpu) !void { +/// Decode up (both pair forms) and down at the wide shape over 64 experts, fresh picks a call: us a call at 1-4 rows. +fn timeDecode(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random, reps: usize) !void { + const a = gpu.gpa; + const sh: Shape = .{ .experts = 64, .ni = wide.ni, .dims = wide.dims, .slots = wide.slots }; + const e = sh.experts; + var cs: [3][2]Codes = undefined; + for (&cs, [_][2]usize{ .{ sh.ni, sh.dims }, .{ sh.ni, sh.dims }, .{ sh.dims, sh.ni } }) |*c, nk| c.* = .{ try Codes.init(a, r, e, nk[0], nk[1], 3), try Codes.init(a, r, 1, nk[0], nk[1], 4) }; + defer for (cs) |c| for (c) |x| x.deinit(a); + var pair = try tables(gpu, k, s, &.{ cs[0], cs[1] }); + defer pair.deinit(); + var down = try tables(gpu, k, s, &.{cs[2]}); + defer down.deinit(); + const l3: cuda.experts.Layer = .{ .format = .affine3, .up = pair.t3.ptr, .down = down.t3.ptr, .width = sh.ni, .dims = sh.dims, .experts = e + 1, .swiglu = true, .shared = @intCast(e), .shared_up = pair.shared4.ptr, .shared_down = down.shared4.ptr }; + var a3 = k.a3; + const g: cuda.experts.Grouped = .{ .affine3 = &a3 }; + const x = try a.alloc(u16, 4 * sh.dims); + defer a.free(x); + for (x) |*v| v.* = randBf16(r, 1.0); + var dx = try upload(gpu, std.mem.sliceAsBytes(x)); + defer dx.free(); + var act = try cuda.DeviceBuffer.alloc(gpu.d, 4 * sh.slots * sh.ni * 2); + defer act.free(); + var y = try cuda.DeviceBuffer.alloc(gpu.d, 4 * sh.slots * sh.dims * 4); + defer y.free(); + var t0 = try cuda.Event.init(gpu.d, true); + defer t0.deinit(); + var t1 = try cuda.Event.init(gpu.d, true); + defer t1.deinit(); + for ([_]usize{ 1, 2, 3, 4 }) |rows| { + var plans: [8]PlanBufs = undefined; + var made: usize = 0; + defer for (plans[0..made]) |*p| p.deinit(); + for (&plans) |*p| { + var picks: [4 * 9]i32 = undefined; + for (0..rows) |row| { + var taken: [64]bool = @splat(false); + for (0..sh.slots - 1) |i| { + var pick = r.uintLessThan(usize, e); + while (taken[pick]) pick = r.uintLessThan(usize, e); + taken[pick] = true; + picks[row * sh.slots + i] = @intCast(pick); + } + picks[row * sh.slots + sh.slots - 1] = @intCast(e); + } + p.* = try PlanBufs.init(gpu, k, s, picks[0 .. rows * sh.slots], e + 1, 16); + made += 1; + } + try s.synchronize(); + const rows_in: cuda.experts.Rows = .{ .x = dx.ptr, .stride = sh.dims, .slots = sh.slots }; + const act_in: cuda.experts.Rows = .{ .x = act.ptr, .stride = sh.ni }; + var us: [3]f64 = undefined; + for (0..3) |form| { + a3.split_below = if (form == 1) std.math.maxInt(usize) else 0; + // the calls replayed as one captured graph, as the engine's forwards run them + try cuda.graph.beginCapture(s, .thread_local); + for (0..reps) |i| { + const p = plans[i % plans.len]; + if (form < 2) { + try g.decode(s, .up, rows_in, l3, p.plan(), act.ptr, p.items_n, -1); + } else try g.decode(s, .down_f32, act_in, l3, p.plan(), y.ptr, p.items_n, -1); + } + var graph = try cuda.graph.endCapture(s); + defer graph.deinit(); + var exec = try graph.instantiate(); + defer exec.deinit(); + for (0..2) |round| { + try t0.record(s); + try exec.launchOn(s); + try t1.record(s); + try t1.synchronize(); + if (round == 1) us[form] = @as(f64, try cuda.Event.elapsedMs(t0, t1)) * 1000 / @as(f64, @floatFromInt(reps)); + } + } + const auto = if (plans[0].items_n * (sh.ni / 32) < k.a3.split_below) "a block a CTA" else "a block a warp"; + std.debug.print("TIME experts3 decode ({d} by {d}, {d} picks a row) at {d} rows: up a block a warp {d:.1} us, a block a CTA {d:.1} us (taken: {s}); down {d:.1} us\n", .{ sh.ni, sh.dims, sh.slots, rows, us[0], us[1], auto, us[2] }); + } +} + +pub fn run(gpu: Gpu, reps: usize) !void { var prng = std.Random.DefaultPrng.init(0x3b17); const r = prng.random(); var mods: [5]cuda.Module = undefined; @@ -286,6 +368,14 @@ pub fn run(gpu: Gpu) !void { const k: Kernels = .{ .a3 = try A3.resolve(mods[3], mods[4], sms), .a4 = try A4.resolve(mods[0], mods[1], mods[2], sms), .router = try cuda.grouped.Router.resolve(mods[0]) }; var s = try cuda.Stream.init(gpu.d, true); defer s.deinit(); - try sameBits(gpu, k, s, r); + // both decode pair forms: a column block a warp, and a column block a CTA + for ([_]usize{ 0, std.math.maxInt(usize) }) |below| { + var kf = k; + kf.a3.split_below = below; + try sameBits(gpu, kf, s, r, .{}); + try sameBits(gpu, kf, s, r, wide); + try sameBits(gpu, kf, s, r, short); + } try projections(gpu, k, s, r); + if (reps > 0) try timeDecode(gpu, k, s, r, reps); } diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 837f1bb070..849be6bc57 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -170,7 +170,7 @@ const usage = \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) \\ nvfp4 NVFP4 projections against recorded lane-matmul and prompt-GEMM fixture bytes \\ moe softmax routing and SwiGLU experts against host references - \\ experts3 affine-3 packing and outputs against affine-4 and fp64 + \\ experts3 [reps] affine-3 packing and outputs against affine-4 and fp64; reps > 0 times decode \\ ggml Bolt's IQ2_XXS experts and Q8_0 / IQ2_XXS dense rows on real tensors against fp64 (tools/zig/bolt_ggml_fixture.py) \\ ggml-bench IQ2_XXS expert and Q8_0 lm_head decode GB/s (synthetic blocks) \\ modelopt ModelOpt repack, FP8 and BF16 lane rows and relu^2 NVFP4 experts against fp64 @@ -292,7 +292,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-vision")) return glm_vision_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "moe")) return moe_tests.run(gpu); - if (std.mem.eql(u8, cmd, "experts3")) return experts3_tests.run(gpu); + if (std.mem.eql(u8, cmd, "experts3")) return experts3_tests.run(gpu, if (rest.len > 0) try std.fmt.parseInt(usize, rest[0], 10) else 0); if (std.mem.eql(u8, cmd, "ggml")) return ggml_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ggml-bench")) return ggml_tests.bench(gpu); if (std.mem.eql(u8, cmd, "checkpoint")) return ck_tests.run(gpu);