experts3 CUDA: a split decode pair form, a column block a CTA with its K groups shared by the warps

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 <int M>
+struct Part {
+  float p[M][NTW][4][32];
+  uint32_t sb[M][NTW][2][32];
+  float xs[2][32];
+};
+
+template <int M, int BITS>
+__device__ __forceinline__ void part_stage(Part<M>& out, const Stage<M, BITS>& st, int lane) {
+  out.xs[0][lane] = group_sum<XV>(st.xa);
+  out.xs[1][lane] = group_sum<XV>(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<BITS>(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 <int M>
+__device__ __forceinline__ void chain(float (&acc)[M][4], const Part<M>& 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 <int M, int D, int W, int BITS>
+__device__ __forceinline__ void split_loop(float (&acc)[M][4], Part<M> (&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<M, BITS> st[D];
+#pragma unroll
+  for (int d = 0; d < D; ++d)
+    if (d * W + w < KG) load_stage<M, BITS>(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<M, BITS>(parts[w], st[d], lane);
+          if (g + D * W < KG) load_stage<M, BITS>(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<M>(acc, parts[s], w, lane, v1);
+        __syncthreads();
+      }
+    }
+  }
+}
+
+// epilogue's SwiGLU stores for tile j alone.
+template <int M>
+__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 <int M, int EPI, int D, int WARPS>
+__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<M> 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<M, D, WARPS, 4>(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<M, D, WARPS, 3>(acc, parts, blk, KG, w, lane, t, x0, x1, v0, v1);
+    }
+    store_tile<M>(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);

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

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.