qwen3_5 CUDA: a captured bucket's padding rows no longer go through the routed experts (35B, every multi-row forward)

zig/build/qwen_contract.zig

@@ -14,7 +14,7 @@ pub fn host(b: *std.Build, core: *std.Build.Module, all: *std.Build.Step) void {
 pub fn images(b: *std.Build, runner: *std.Build.Step.Compile, controls: ?std.Build.LazyPath, h3: ?std.Build.LazyPath, kernels: anytype, files: []const ?std.Build.LazyPath) void {
     const install = @import("kernel_contract.zig").cuda(b, runner, controls, h3);
     for (kernels, files) |k, file| {
-        if (std.mem.eql(u8, k.name, "moe") or std.mem.eql(u8, k.name, "qwen36_mtp")) {
+        if (std.mem.eql(u8, k.name, "moe") or std.mem.eql(u8, k.name, "experts") or std.mem.eql(u8, k.name, "qwen36_mtp")) {
             if (file) |image| install.dependOn(&b.addInstallFile(image, b.fmt("fatbin/{s}.fatbin", .{k.name})).step);
         }
     }

zig/exl3_prompt_host.zig

@@ -56,7 +56,7 @@ fn run(rows: usize, prompt: bool, shared: bool) !void {
     const s: cuda.Stream = .{ .d = &d, .handle = null };
     var k: ex.Experts = .{ .mod = undefined, .group = function(&d, 1), .rot_in = function(&d, 2), .grouped = @splat(function(&d, 3)), .gateup = function(&d, 4), .down_combine = function(&d, 5) };
     if (comptime @hasField(ex.Experts, "prompt")) k.prompt = .{ .rotate = function(&d, 11), .gateup = function(&d, 12), .down = function(&d, 13) };
-    const ek: glm.experts.Kernels = .{ .nvfp4 = undefined, .router = .{ .plan_small = function(&d, 20), .plan_rank = function(&d, 21), .plan_offsets = function(&d, 22), .plan_scatter = function(&d, 23) } };
+    const ek: glm.experts.Kernels = .{ .nvfp4 = undefined, .router = .{ .plan_small = function(&d, 20), .plan_rank = function(&d, 21), .plan_offsets = function(&d, 22), .plan_scatter = function(&d, 23), .plan_rows = function(&d, 24), .plan_rank_rows = function(&d, 25), .plan_scatter_rows = function(&d, 26) } };
     const w: glm.exl3.Experts = .{ .gt = 100, .ut = 200, .dt = 300, .k2 = 400, .suh_g = 500, .suh_u = 600, .svh_g = 700, .svh_u = 800, .suh_d = 900, .svh_d = 1000, .shx = shared };
     const sc: glm.exl3.Scratch = .{ .picks = 0x100000000, .xg = 0, .xu = 0, .z = 0 };
     const plan: cuda.grouped.Plan = .{ .items = 10000, .members = 20000, .counts = 30000, .rank = 40000, .hist = 50000 };

zig/kernels/cuda/experts.cu

@@ -228,8 +228,36 @@ __global__ void __launch_bounds__(PLAN_THREADS)
   }
 }
 
+// plan_kernel over the first *rows rows' pairs (per a row): a captured bucket's padding rows get no items.
 __global__ void __launch_bounds__(PLAN_THREADS)
-    plan_rank(const int* __restrict__ picks, int P, int E, int* __restrict__ rank, int* __restrict__ hist) {
+    plan_rows(const int* __restrict__ picks, int P, int E, int T, int* __restrict__ members,
+              int* __restrict__ items, int* __restrict__ counts, const int* __restrict__ rows, int per) {
+  P = min(P, *rows * per);
+  __shared__ int pk[PSMALL];
+  __shared__ int cnt[EMAX];
+  __shared__ int off[EMAX];
+  const int tid = threadIdx.x;
+  for (int e = tid; e < EMAX; e += PLAN_THREADS) cnt[e] = 0;
+  __syncthreads();
+  for (int p = tid; p < P; p += PLAN_THREADS) {
+    const int e = picks[p];
+    pk[p] = e;
+    atomicAdd(&cnt[e], 1);
+  }
+  __syncthreads();
+  const int o = place_items(tid < E ? cnt[tid] : 0, E, T, items, counts);
+  if (tid < E) off[tid] = o;
+  __syncthreads();
+  for (int p = tid; p < P; p += PLAN_THREADS) {
+    const int e = pk[p];
+    int rank = 0;
+    for (int q = 0; q < p; ++q) rank += pk[q] == e;
+    members[off[e] + rank] = p;
+  }
+}
+
+__device__ __forceinline__ void rank_pairs(const int* __restrict__ picks, int P, int E, int* __restrict__ rank,
+                                           int* __restrict__ hist) {
   __shared__ int cnt[EMAX];
   const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
   for (int e = tid; e < E; e += PLAN_THREADS) cnt[e] = 0;
@@ -251,6 +279,18 @@ __global__ void __launch_bounds__(PLAN_THREADS)
   for (int x = tid; x < E; x += PLAN_THREADS) hist[(size_t)blockIdx.x * E + x] = cnt[x];
 }
 
+__global__ void __launch_bounds__(PLAN_THREADS)
+    plan_rank(const int* __restrict__ picks, int P, int E, int* __restrict__ rank, int* __restrict__ hist) {
+  rank_pairs(picks, P, E, rank, hist);
+}
+
+// plan_rank over the first *rows rows' pairs (per a row): a captured bucket's padding pairs count nowhere.
+__global__ void __launch_bounds__(PLAN_THREADS)
+    plan_rank_rows(const int* __restrict__ picks, int P, int E, int* __restrict__ rank, int* __restrict__ hist,
+                   const int* __restrict__ rows, int per) {
+  rank_pairs(picks, min(P, *rows * per), E, rank, hist);
+}
+
 __global__ void __launch_bounds__(PLAN_THREADS)
     plan_offsets(int nblk, int E, int T, int* __restrict__ hist, int* __restrict__ items, int* __restrict__ counts) {
   const int tid = threadIdx.x;
@@ -273,6 +313,14 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const
   if (p < P) members[hist[(size_t)(p / PLAN_THREADS) * E + picks[p]] + rank[p]] = p;
 }
 
+// plan_scatter over the first *rows rows' pairs (per a row).
+__global__ void plan_scatter_rows(const int* __restrict__ picks, int P, int E, const int* __restrict__ rank,
+                                  const int* __restrict__ hist, int* __restrict__ members, const int* __restrict__ rows,
+                                  int per) {
+  const int p = blockIdx.x * blockDim.x + threadIdx.x;
+  if (p < min(P, *rows * per)) members[hist[(size_t)(p / PLAN_THREADS) * E + picks[p]] + rank[p]] = p;
+}
+
 } // namespace tf_experts
 
 // Decode form for groups of 64: up with relu^2 (epilogue 1), down to fp32 (epilogue 0).

zig/kernels/cuda/moe.cu

@@ -14,10 +14,12 @@ constexpr int KC = 256;   // inputs a shared-memory chunk holds
 
 }  // namespace
 
-// OUT[r][e] = fp32 x[r] . w[e] over D in index order (one fma chain a pair); grid (rows / 16, E / 16), block 256.
+// OUT[r][e] = fp32 x[r] . w[e] in index order (an fma chain a pair), grid (R / 16, E / 16); skips rows >= P[2].
 extern "C" __global__ void __launch_bounds__(256) tf_moe_router(const __nv_bfloat16* __restrict__ X, int x_stride,
                                                                   const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT,
-                                                                  int R, int D, int E) {
+                                                                  int R, int D, int E, const int* __restrict__ P) {
+    if (P != nullptr) R = min(R, P[2]);
+    if (static_cast<int>(blockIdx.x) * TILE >= R) return;
     __shared__ __nv_bfloat16 xs[TILE][KC + 2];
     __shared__ __nv_bfloat16 ws[TILE][KC + 2];
     const int row0 = blockIdx.x * TILE, e0 = blockIdx.y * TILE;
@@ -47,10 +49,12 @@ __device__ __forceinline__ bool beats(float v, int i, float bv, int bi) { return
 
 }  // namespace
 
-// Row r by a warp: TOPK experts by logit (ties to the lower id), weights bf16(exp(l - top) / sum), the shared slot.
+// Row r by a warp: TOPK experts by logit (ties to lower id), bf16(exp(l - top) / sum), shared slot; r < P[2].
 extern "C" __global__ void __launch_bounds__(32) tf_moe_topk(const float* __restrict__ L, int* __restrict__ PICK,
-                                                              float* __restrict__ WT, int NE, int TOPK) {
+                                                              float* __restrict__ WT, int NE, int TOPK,
+                                                              const int* __restrict__ P) {
     const int64_t r = blockIdx.x;
+    if (P != nullptr && r >= P[2]) return;
     const int lane = threadIdx.x;
     const int NL = NE + 1, SLOTS = TOPK + 1;  // L[r][NE], the shared gate's logit, weighs slot TOPK by bf16(sigmoid(bf16(it)))
     float v[MAXC];
@@ -97,10 +101,12 @@ extern "C" __global__ void __launch_bounds__(32) tf_moe_topk(const float* __rest
     }
 }
 
-// OUT[r][d] = bf16(sum over slots in order of fma(w[r][s], y[r][s][d])); y fp32 (decode) or bf16 (prompt rows).
+// OUT[r][d] = bf16(sum over slots in order of fma(w[r][s], y[r][s][d])), y fp32 or bf16; rows >= P[2] skipped.
 extern "C" __global__ void __launch_bounds__(256) tf_moe_combine(const void* __restrict__ Y, int y_f32, const float* __restrict__ WT,
-                                                                   __nv_bfloat16* __restrict__ OUT, int S, int D) {
+                                                                   __nv_bfloat16* __restrict__ OUT, int S, int D,
+                                                                   const int* __restrict__ P) {
     const int64_t r = blockIdx.x;
+    if (P != nullptr && r >= P[2]) return;
     for (int d = blockIdx.y * 256 + threadIdx.x; d < D; d += gridDim.y * 256) {
         float acc = 0.0f;
         for (int s = 0; s < S; ++s) {

zig/src/core/contract/qwen_served_records.zig

@@ -1,12 +1,15 @@
-//! The Qwen CUDA engine's MoE and MTP launches, read by kernel-contract-limits from every SM's image.
+//! The Qwen CUDA engine's MoE, plan and MTP launches, read by kernel-contract-limits from every SM's image.
 const std = @import("std");
 const registry = @import("registry.zig");
 
 /// nvcc 13.0's usage, the compiler of the release fatbins.
 pub const records = [_]registry.CudaRecord{
-    .{ .launch = .{ .name = "tf_moe_router", .symbol = "tf_moe_router", .image = "moe", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 40, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 } } },
+    .{ .launch = .{ .name = "tf_moe_router", .symbol = "tf_moe_router", .image = "moe", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 } } },
     .{ .launch = .{ .name = "tf_moe_topk", .symbol = "tf_moe_topk", .image = "moe", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 35, .static = 0, .bound = 32 }, .{ .regs = 35, .static = 0, .bound = 32 }, .{ .regs = 35, .static = 0, .bound = 32 }, .{ .regs = 34, .static = 0, .bound = 32 }, .{ .regs = 34, .static = 0, .bound = 32 } } },
     .{ .launch = .{ .name = "tf_moe_combine", .symbol = "tf_moe_combine", .image = "moe", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 } } },
+    .{ .launch = .{ .name = "experts plan_rows", .symbol = "_ZN10tf_experts9plan_rowsEPKiiiiPiS2_S2_S1_i", .image = "experts", .threads = 1024, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 12672, .bound = 1024 }, .{ .regs = 38, .static = 12672, .bound = 1024 }, .{ .regs = 38, .static = 12672, .bound = 1024 }, .{ .regs = 32, .static = 12672, .bound = 1024 }, .{ .regs = 32, .static = 12672, .bound = 1024 } } },
+    .{ .launch = .{ .name = "experts plan_rank_rows", .symbol = "_ZN10tf_experts14plan_rank_rowsEPKiiiPiS2_S1_i", .image = "experts", .threads = 1024, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 4096, .bound = 1024 }, .{ .regs = 40, .static = 4096, .bound = 1024 }, .{ .regs = 40, .static = 4096, .bound = 1024 }, .{ .regs = 42, .static = 4096, .bound = 1024 }, .{ .regs = 42, .static = 4096, .bound = 1024 } } },
+    .{ .launch = .{ .name = "experts plan_scatter_rows", .symbol = "_ZN10tf_experts17plan_scatter_rowsEPKiiiS1_S1_PiS1_i", .image = "experts", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } },
     .{ .launch = .{ .name = "tf_q36m_norm_rope", .symbol = "tf_q36m_norm_rope", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 36, .static = 0, .bound = 32 }, .{ .regs = 37, .static = 0, .bound = 32 }, .{ .regs = 39, .static = 0, .bound = 32 }, .{ .regs = 42, .static = 0, .bound = 32 }, .{ .regs = 42, .static = 0, .bound = 32 } } },
     .{ .launch = .{ .name = "tf_q36m_attn", .symbol = "tf_q36m_attn", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 50, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 51, .static = 0, .bound = 32 }, .{ .regs = 54, .static = 0, .bound = 32 }, .{ .regs = 54, .static = 0, .bound = 32 } } },
     .{ .launch = .{ .name = "tf_q36m_attn_tree", .symbol = "tf_q36m_attn_tree", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 53, .static = 0, .bound = 32 }, .{ .regs = 54, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 57, .static = 0, .bound = 32 }, .{ .regs = 57, .static = 0, .bound = 32 } } },
@@ -17,7 +20,7 @@ pub const cuda_132 = [_]registry.Override{
     .{ .name = "tf_q36m_attn", .usage = .{ .{ .regs = 50, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 51, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 } } },
 };
 
-test "the six served MoE and MTP launches are in the live registry and fit every SM by their records" {
+test "the nine served MoE, plan and MTP launches are in the live registry and fit every SM by their records" {
     for (records) |record| {
         var found: usize = 0;
         for (registry.all_cuda) |active| if (std.mem.eql(u8, active.launch.symbol, record.launch.symbol)) {

zig/src/cuda/grouped.zig

@@ -22,42 +22,68 @@ pub const Plan = struct { members: u64, items: u64, counts: u64, rank: u64, hist
 /// experts.cu's plan instances (cuobjdump -symbols of the experts fatbin).
 pub const symbols = struct {
     pub const plan = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_";
+    pub const plan_rows = "_ZN10tf_experts9plan_rowsEPKiiiiPiS2_S2_S1_i";
     pub const rank = "_ZN10tf_experts9plan_rankEPKiiiPiS2_";
+    pub const rank_rows = "_ZN10tf_experts14plan_rank_rowsEPKiiiPiS2_S1_i";
     pub const offsets = "_ZN10tf_experts12plan_offsetsEiiiPiS0_S0_";
     pub const scatter = "_ZN10tf_experts12plan_scatterEPKiiiS1_S1_Pi";
+    pub const scatter_rows = "_ZN10tf_experts17plan_scatter_rowsEPKiiiS1_S1_PiS1_i";
 };
 
 pub const Router = struct {
     plan_small: Function,
+    plan_rows: Function,
     plan_rank: Function,
     plan_offsets: Function,
     plan_scatter: Function,
+    plan_rank_rows: Function,
+    plan_scatter_rows: Function,
 
     /// The plan kernels from a loaded experts module (kernels.experts).
     pub fn resolve(mod: Module) !Router {
-        return .{ .plan_small = try mod.function(symbols.plan), .plan_rank = try mod.function(symbols.rank), .plan_offsets = try mod.function(symbols.offsets), .plan_scatter = try mod.function(symbols.scatter) };
+        return .{ .plan_small = try mod.function(symbols.plan), .plan_rows = try mod.function(symbols.plan_rows), .plan_rank = try mod.function(symbols.rank), .plan_offsets = try mod.function(symbols.offsets), .plan_scatter = try mod.function(symbols.scatter), .plan_rank_rows = try mod.function(symbols.rank_rows), .plan_scatter_rows = try mod.function(symbols.scatter_rows) };
     }
 
     fn go(f: Function, s: Stream, grid: usize, block: u32, a: *launch_.Args) !void {
         try launch_.launch(f, .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = block } }, s, a);
     }
 
+    /// route over the first `rows.*` rows' pairs (`per` a row; `rows` a device int, 0: all), at any width.
+    pub fn routeRows(r: Router, s: Stream, picks: u64, pairs: usize, count: usize, tile: usize, p: Plan, rows: u64, per: usize) !void {
+        if (rows == 0) return r.route(s, picks, pairs, count, tile, p);
+        if (pairs > small) return r.wide(s, picks, pairs, count, tile, p, .{ .rows = rows, .per = per });
+        var a: launch_.Args = .{};
+        a.add(picks);
+        for ([_]usize{ pairs, count, tile }) |v| a.add(@as(c_int, @intCast(v)));
+        for ([_]u64{ p.members, p.items, p.counts, rows }) |v| a.add(v);
+        a.add(@as(c_int, @intCast(per)));
+        return go(r.plan_rows, s, 1, 1024, &a);
+    }
+
     /// experts.route: `pairs` int32 expert ids at `picks` grouped into items of at most `tile` pairs (16 decode).
     pub fn route(r: Router, s: Stream, picks: u64, pairs: usize, count: usize, tile: usize, p: Plan) !void {
+        if (pairs > small) return r.wide(s, picks, pairs, count, tile, p, null);
         var a: launch_.Args = .{};
-        if (pairs <= small) {
-            a.add(picks);
-            for ([_]usize{ pairs, count, tile }) |v| a.add(@as(c_int, @intCast(v)));
-            for ([_]u64{ p.members, p.items, p.counts }) |v| a.add(v);
-            return go(r.plan_small, s, 1, 1024, &a);
-        }
+        a.add(picks);
+        for ([_]usize{ pairs, count, tile }) |v| a.add(@as(c_int, @intCast(v)));
+        for ([_]u64{ p.members, p.items, p.counts }) |v| a.add(v);
+        return go(r.plan_small, s, 1, 1024, &a);
+    }
+
+    /// The plan ranked in blocks of 1024 pairs; `live` (a device row count and pairs a row) caps the pairs planned.
+    fn wide(r: Router, s: Stream, picks: u64, pairs: usize, count: usize, tile: usize, p: Plan, live: ?struct { rows: u64, per: usize }) !void {
         const nblk = (pairs + 1023) / 1024;
+        var a: launch_.Args = .{};
         a.add(picks);
         a.add(@as(c_int, @intCast(pairs)));
         a.add(@as(c_int, @intCast(count)));
         a.add(p.rank);
         a.add(p.hist);
-        try go(r.plan_rank, s, nblk, 1024, &a);
+        if (live) |l| {
+            a.add(l.rows);
+            a.add(@as(c_int, @intCast(l.per)));
+        }
+        try go(if (live != null) r.plan_rank_rows else r.plan_rank, s, nblk, 1024, &a);
         var b: launch_.Args = .{};
         for ([_]usize{ nblk, count, tile }) |v| b.add(@as(c_int, @intCast(v)));
         for ([_]u64{ p.hist, p.items, p.counts }) |v| b.add(v);
@@ -67,7 +93,11 @@ pub const Router = struct {
         c.add(@as(c_int, @intCast(pairs)));
         c.add(@as(c_int, @intCast(count)));
         for ([_]u64{ p.rank, p.hist, p.members }) |v| c.add(v);
-        try go(r.plan_scatter, s, (pairs + 255) / 256, 256, &c);
+        if (live) |l| {
+            c.add(l.rows);
+            c.add(@as(c_int, @intCast(l.per)));
+        }
+        try go(if (live != null) r.plan_scatter_rows else r.plan_scatter, s, (pairs + 255) / 256, 256, &c);
     }
 };