glm cuda: a decode step's select gives each row a 1,024-thread block

docs/kernel-contract.md

@@ -182,7 +182,9 @@ CUDA also checks GLM-5.3-Flash's KV page kernels on a page table out of order: t
 (`tf_glm_lat_write`, `tf_glm_pool_keys`), the index scores and the selected tokens' rows (`tf_glm_scores`,
 `tf_glm_dsa_tokens`), and the MLA chunks through their merge (`tf_glm_chunks_16` and `tf_glm_chunks_32`, dense and
 sparse). Its prompt indexer (`tf_glm_index_scores`, `tf_glm_index_pools`) is bit-equal to `tf_glm_scores` and
-`tf_glm_select_rows`; the pools' stated rule ranks any NaN above +inf, as the GPU's `s + 0` makes every NaN one positive NaN.
+`tf_glm_select_rows`, and its decode select at a 1,024-thread block a row (`tf_glm_select_rows_wide`) to
+`tf_glm_select_rows`; both stated rules rank any NaN above +inf, as the GPU's `s + 0` makes every NaN one
+positive NaN.
 
 GLM's Metal window projections register narrow and wide rows at 2, 3, 4, 5, 6 and 8 bits, plus both gathered paths.
 Independent FP64 linear cases accompany byte replacements against frozen pre-window-optimization kernels.
@@ -203,7 +205,7 @@ and FP32-output paths, FP8G's 32-row cluster tile with FP32 output, the common E
 decode-sized latent top-k and NVFP4 up experts, the Mamba scan and Super's NVFP4 prompt down kernels, three of
 GLM-5.3-Flash's packed matmuls and three of its DFlash2 launches, and the training launches it compiles differently,
 and nvcc 13.3 keeps 13.2's and changes `tf_nemo_topk` on sm_120 and sm_121, four of GLM-5.3-Flash's KV page
-launches and its two prompt-indexer launches. `kernel-contract` and
+launches, its two prompt-indexer launches and its wide decode select. `kernel-contract` and
 `kernel-contract-limits` read the release from each image's `.note.nv.tkinfo`, check the image against the newest set
 at or below that release (13.0's when an image names none), and print which set they used. A release with no set of
 its own is checked against the one below it, so its differences show as stale records: re-derive them with

zig/kernels/cuda/glm_glue_index.cuh

@@ -70,8 +70,9 @@ __device__ __forceinline__ uint32_t order_key(float s) {
 
 namespace glm_index {
 
-// 4 warps' exclusive prefix of v in thread order and its total, via slot of `wsum`; its barrier orders slot reuse.
-__device__ __forceinline__ int block_scan(int v, int (*wsum)[4], int slot, int* total) {
+// W warps' exclusive prefix of v in thread order and its total, via slot of `wsum`; its barrier orders slot reuse.
+template <int W>
+__device__ __forceinline__ int block_scan(int v, int (*wsum)[W], int slot, int* total) {
   const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
   int inc = v;
   for (int o = 1; o < 32; o <<= 1) {
@@ -82,18 +83,18 @@ __device__ __forceinline__ int block_scan(int v, int (*wsum)[4], int slot, int*
   __syncthreads();
   int before = inc - v;
   for (int w = 0; w < warp; ++w) before += wsum[slot][w];
-  *total = wsum[slot][0] + wsum[slot][1] + wsum[slot][2] + wsum[slot][3];
+  int all = 0;
+  for (int w = 0; w < W; ++w) all += wsum[slot][w];
+  *total = all;
   return before;
 }
 
-}  // namespace glm_index
-
-// sparse._select_rows (VISIBLE), a row per 4-warp block: K best visible pools by 8-bit radix select, ties to lower.
-extern "C" __global__ void __launch_bounds__(128)
-    tf_glm_select_rows(const float* scores, int64_t* out, int64_t np, const int* pos, int k) {
+// sparse._select_rows (VISIBLE), a row per T-thread block: K best visible pools by 8-bit radix select, ties to lower.
+template <int T>
+__device__ __forceinline__ void select_row(const float* scores, int64_t* out, int64_t np, const int* pos, int k) {
   __shared__ int hist[256];
   __shared__ uint32_t sh_prefix;
-  __shared__ int sh_need, wsum[2][4];
+  __shared__ int sh_need, wsum[2][T / 32];
   const int64_t r = blockIdx.x;
   const float* row = scores + r * np;
   const int t = threadIdx.x, lane = t & 31, warp = t >> 5;
@@ -112,12 +113,13 @@ extern "C" __global__ void __launch_bounds__(128)
       for (int j = 0; j < 8; ++j) u[j] = i0 + j < limit ? glm_index::order_key(row[i0 + j]) : 0u;
     }
   };
+  // T only sets the pools a pass reads (8 a thread): digit counts and index-order scans give every T the same picks
   uint32_t prefix = 0, fixed = 0;
   int need = k;
   for (int shift = 24; shift >= 0; shift -= 8) {
     for (int b = t; b < 256; b += blockDim.x) hist[b] = 0;
     __syncthreads();
-    for (int64_t c = 0; c < limit; c += 1024) {
+    for (int64_t c = 0; c < limit; c += 8 * T) {
       uint32_t u[8];
       keys(c, u);
       for (int j = 0; j < 8; ++j) {
@@ -159,7 +161,7 @@ extern "C" __global__ void __launch_bounds__(128)
     __syncthreads();
   }
   int run_take = 0, run_eq = 0;
-  for (int64_t c = 0; c < limit; c += 1024) {
+  for (int64_t c = 0; c < limit; c += 8 * T) {
     uint32_t u[8];
     keys(c, u);
     bool eq[8], gt[8];
@@ -171,18 +173,32 @@ extern "C" __global__ void __launch_bounds__(128)
       eqs += eq[j];
     }
     int all_eq, all_take;
-    int seen = run_eq + glm_index::block_scan(eqs, wsum, 0, &all_eq);
+    int seen = run_eq + block_scan<T / 32>(eqs, wsum, 0, &all_eq);
     bool take[8];
     int takes = 0;
     for (int j = 0; j < 8; ++j) {
       take[j] = gt[j] || (eq[j] && seen < need);
       seen += eq[j];
       takes += take[j];
     }
-    int at = run_take + glm_index::block_scan(takes, wsum, 1, &all_take);
+    int at = run_take + block_scan<T / 32>(takes, wsum, 1, &all_take);
     for (int j = 0; j < 8; ++j)
       if (take[j]) out[r * k + at++] = c + 8 * t + j;
     run_eq += all_eq;
     run_take += all_take;
   }
 }
+
+}  // namespace glm_index
+
+// sparse._select_rows (VISIBLE), a row per 4-warp block: K best visible pools by 8-bit radix select, ties to lower.
+extern "C" __global__ void __launch_bounds__(128)
+    tf_glm_select_rows(const float* scores, int64_t* out, int64_t np, const int* pos, int k) {
+  glm_index::select_row<128>(scores, out, np, pos, k);
+}
+
+// tf_glm_select_rows a row per 1,024-thread block, for a few rows over many pools (a decode step's): its bytes.
+extern "C" __global__ void __launch_bounds__(1024)
+    tf_glm_select_rows_wide(const float* scores, int64_t* out, int64_t np, const int* pos, int k) {
+  glm_index::select_row<1024>(scores, out, np, pos, k);
+}

zig/src/core/contract/glm_cuda_records.zig

@@ -43,6 +43,8 @@ pub const cuda = [_]registry.CudaRecord{
     .{ .launch = .{ .name = "glm_glue tf_glm_scores", .symbol = "tf_glm_scores", .image = "glm_glue", .threads = 128, .min_sm = 80 }, .usage = .{ .{ .regs = 126, .static = 32768, .bound = 128 }, .{ .regs = 131, .static = 32768, .bound = 128 }, .{ .regs = 131, .static = 32768, .bound = 128 }, .{ .regs = 155, .static = 32768, .bound = 128 }, .{ .regs = 155, .static = 32768, .bound = 128 } } },
     .{ .launch = .{ .name = "glm_index tf_glm_index_scores", .symbol = "tf_glm_index_scores", .image = "glm_index", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 128, .static = 32768, .bound = 256 }, .{ .regs = 128, .static = 32768, .bound = 256 }, .{ .regs = 128, .static = 32768, .bound = 256 }, .{ .regs = 124, .static = 32768, .bound = 256 }, .{ .regs = 124, .static = 32768, .bound = 256 } } },
     .{ .launch = .{ .name = "glm_index tf_glm_index_pools", .symbol = "tf_glm_index_pools", .image = "glm_index", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 62, .static = 33864, .bound = 256 }, .{ .regs = 64, .static = 33864, .bound = 256 }, .{ .regs = 64, .static = 33864, .bound = 256 }, .{ .regs = 59, .static = 33864, .bound = 256 }, .{ .regs = 59, .static = 33864, .bound = 256 } } },
+    .{ .launch = .{ .name = "glm_glue tf_glm_select_rows", .symbol = "tf_glm_select_rows", .image = "glm_glue", .threads = 128, .min_sm = 80 }, .usage = .{ .{ .regs = 45, .static = 1072, .bound = 128 }, .{ .regs = 54, .static = 1072, .bound = 128 }, .{ .regs = 54, .static = 1072, .bound = 128 }, .{ .regs = 48, .static = 1072, .bound = 128 }, .{ .regs = 48, .static = 1072, .bound = 128 } } },
+    .{ .launch = .{ .name = "glm_glue tf_glm_select_rows_wide", .symbol = "tf_glm_select_rows_wide", .image = "glm_glue", .threads = 1024, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 1296, .bound = 1024 }, .{ .regs = 61, .static = 1296, .bound = 1024 }, .{ .regs = 61, .static = 1296, .bound = 1024 }, .{ .regs = 54, .static = 1296, .bound = 1024 }, .{ .regs = 61, .static = 1296, .bound = 1024 } } },
     .{ .launch = .{ .name = "glm_glue tf_glm_chunks_16", .symbol = "tf_glm_chunks_16", .image = "glm_glue", .threads = 256, .dynamic = 83328, .opt_in = true, .min_sm = 80 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 256 }, .{ .regs = 64, .static = 0, .bound = 256 }, .{ .regs = 64, .static = 0, .bound = 256 }, .{ .regs = 123, .static = 0, .bound = 256 }, .{ .regs = 123, .static = 0, .bound = 256 } } },
     .{ .launch = .{ .name = "glm_glue tf_glm_chunks_32", .symbol = "tf_glm_chunks_32", .image = "glm_glue", .threads = 256, .dynamic = 101120, .opt_in = true, .min_sm = 80 }, .usage = .{ .{ .regs = 150, .static = 0, .bound = 256 }, .{ .regs = 150, .static = 0, .bound = 256 }, .{ .regs = 150, .static = 0, .bound = 256 }, .{ .regs = 130, .static = 0, .bound = 256 }, .{ .regs = 130, .static = 0, .bound = 256 } } },
     .{ .launch = .{ .name = "glm_ops tf_glm_dsa_tokens", .symbol = "tf_glm_dsa_tokens", .image = "glm_ops", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 } } },
@@ -64,6 +66,7 @@ pub const cuda_133 = [_]registry.Override{
     .{ .name = "glm_glue tf_glm_scores", .usage = .{ .{ .regs = 122, .static = 32768, .bound = 128 }, .{ .regs = 129, .static = 32768, .bound = 128 }, .{ .regs = 129, .static = 32768, .bound = 128 }, .{ .regs = 157, .static = 32768, .bound = 128 }, .{ .regs = 157, .static = 32768, .bound = 128 } } },
     .{ .name = "glm_index tf_glm_index_scores", .usage = .{ .{ .regs = 128, .static = 32768, .bound = 256 }, .{ .regs = 128, .static = 32768, .bound = 256 }, .{ .regs = 128, .static = 32768, .bound = 256 }, .{ .regs = 126, .static = 32768, .bound = 256 }, .{ .regs = 126, .static = 32768, .bound = 256 } } },
     .{ .name = "glm_index tf_glm_index_pools", .usage = .{ .{ .regs = 62, .static = 33864, .bound = 256 }, .{ .regs = 64, .static = 33864, .bound = 256 }, .{ .regs = 64, .static = 33864, .bound = 256 }, .{ .regs = 54, .static = 33864, .bound = 256 }, .{ .regs = 54, .static = 33864, .bound = 256 } } },
+    .{ .name = "glm_glue tf_glm_select_rows_wide", .usage = .{ .{ .regs = 32, .static = 1296, .bound = 1024 }, .{ .regs = 61, .static = 1296, .bound = 1024 }, .{ .regs = 61, .static = 1296, .bound = 1024 }, .{ .regs = 62, .static = 1296, .bound = 1024 }, .{ .regs = 62, .static = 1296, .bound = 1024 } } },
     .{ .name = "glm_glue tf_glm_chunks_32", .usage = .{ .{ .regs = 150, .static = 0, .bound = 256 }, .{ .regs = 150, .static = 0, .bound = 256 }, .{ .regs = 150, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 } } },
     .{ .name = "glm_ops tf_glm_dsa_tokens", .usage = .{ .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 } } },
 };

zig/src/families/glm/cuda_glue.zig

@@ -44,6 +44,7 @@ pub const Kernels = struct {
     index_write: cuda.Function,
     pool_keys: cuda.Function,
     select_rows: cuda.Function,
+    select_rows_wide: cuda.Function,
     bmm_decode: cuda.Function,
     bmm_rows: cuda.Function,
     bmm_tile: cuda.Function,

zig/src/families/glm/cuda_sparse.zig

@@ -19,6 +19,14 @@ fn cdiv(a: usize, b: usize) usize {
     return (a + b - 1) / b;
 }
 
+/// At most wide_rows rows (one wave on 48 SMs) over wide_pools pools or more: 1,024 threads a row, the same bytes.
+const wide_rows = 48;
+const wide_pools = 4096;
+
+fn wideSelect(rows: usize, np: usize) bool {
+    return rows <= wide_rows and np >= wide_pools;
+}
+
 /// sparse.index_update: the window's LayerNormed index keys and gates into ik and ig at pos, completed pools into pk.
 pub fn indexUpdate(set: glue.Set, stream: cuda.Stream, k_raw: u64, k_stride: usize, gate: u64, ln_w: u64, ln_b: u64, ape: u64, ik: u64, ig: u64, pk: u64, t: latent.Table, pos: u64, rows: usize) !void {
     var a: cuda.Args = .{};
@@ -89,7 +97,8 @@ pub fn selectTokensProfile(set: glue.Set, ops: cuda.Module, stream: cuda.Stream,
         };
         var sel: cuda.Args = .{};
         inline for (.{ s.scores, s.pools + a * topk_pools * 8, @as(i64, @intCast(np_max)), at, @as(c_int, topk_pools) }) |v| sel.add(v);
-        try cuda.launch.launch(set.glue.select_rows, .{ .grid = .{ .x = @intCast(n) }, .block = .{ .x = 128 } }, stream, &sel);
+        const wide = wideSelect(n, np_max);
+        try cuda.launch.launch(if (wide) set.glue.select_rows_wide else set.glue.select_rows, .{ .grid = .{ .x = @intCast(n) }, .block = .{ .x = if (wide) 1024 else 128 } }, stream, &sel);
         stamp.end();
     }
     const stamp = prof.span(profile, stream, .index_select);

zig/tests/cuda/contract_glm_index.zig

@@ -3,8 +3,10 @@ const std = @import("std");
 const cuda = @import("cuda");
 const core = @import("core");
 const glm = @import("glm_cuda");
+const check = @import("check.zig");
 const rig_ = @import("contract_rig.zig");
 const pages = @import("contract_glm_pages.zig");
+const select = @import("contract_glm_select.zig");
 
 const contract = core.contract;
 const c = contract.spec;
@@ -326,3 +328,24 @@ pub const Tokens = struct {
         return null;
     }
 };
+
+/// The index scores and the decode select, registered together so the runner keeps one entry for GLM's index.
+pub const Index = struct {
+    scores: Scores,
+    select: select.Select,
+
+    pub fn init(gpu: check.Gpu, s: cuda.Stream, fns: glm.glue.Kernels) Index {
+        return .{ .scores = .{ .rig = .{ .gpu = gpu, .s = s }, .f = fns.scores }, .select = .{ .rig = .{ .gpu = gpu, .s = s }, .wide = fns.select_rows_wide, .narrow = fns.select_rows } };
+    }
+
+    /// After init's value sits at its final address: each kernel's context points into it.
+    pub fn append(self: *Index, a: std.mem.Allocator, kernels: *std.ArrayList(g.Kernel)) !void {
+        try kernels.append(a, self.scores.kernel());
+        try kernels.append(a, self.select.kernel());
+    }
+
+    pub fn deinit(self: *Index) void {
+        self.select.deinit();
+        self.scores.deinit();
+    }
+};