From dd7332f0fb0f9e4f164000acb6a64093c2c4d4a9 Mon Sep 17 00:00:00 2001
From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com>
Date: Sun, 25 Oct 2026 23:08:19 +0900
Subject: [PATCH 1/3] glm cuda: a decode step's select gives each row a
 1,024-thread block

tf_glm_select_rows runs a 128-thread block a row. A decode step has a few rows, so over a long context's
bucket four warps a row read every pool while most SMs idle. tf_glm_select_rows_wide runs the same radix
select with 1,024 threads: T only sets how many pools a pass reads (8 a thread), the digits are counts and
the picks leave through an index-order scan, so every T gives the same pools in the same order. cuda_sparse
launches it for at most 48 rows (one wave on GB10's 48 SMs) over buckets of 4,096 pools or more.

glm-ref checks both kernels' bytes and the host's rule from the 1,024- to the 131,072-pool bucket at 1 to
128 rows, and prints both launches' times.
---
 zig/kernels/cuda/glm_glue_index.cuh  | 42 ++++++++++++------
 zig/src/families/glm/cuda_glue.zig   |  1 +
 zig/src/families/glm/cuda_sparse.zig | 11 ++++-
 zig/tests/cuda/glm_ref_index.zig     | 65 ++++++++++++++++++++++++++++
 zig/tests/cuda/glm_ref_tests.zig     |  1 +
 5 files changed, 106 insertions(+), 14 deletions(-)

diff --git a/zig/kernels/cuda/glm_glue_index.cuh b/zig/kernels/cuda/glm_glue_index.cuh
index f0cd03c1f9..00473d4db5 100644
--- a/zig/kernels/cuda/glm_glue_index.cuh
+++ b/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,7 +173,7 @@ 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) {
@@ -179,10 +181,24 @@ extern "C" __global__ void __launch_bounds__(128)
       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);
+}
diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig
index 5f044f0b20..520fb1525f 100644
--- a/zig/src/families/glm/cuda_glue.zig
+++ b/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,
diff --git a/zig/src/families/glm/cuda_sparse.zig b/zig/src/families/glm/cuda_sparse.zig
index cccd010e59..fe7c391382 100644
--- a/zig/src/families/glm/cuda_sparse.zig
+++ b/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);
diff --git a/zig/tests/cuda/glm_ref_index.zig b/zig/tests/cuda/glm_ref_index.zig
index f9d027cf7e..9514cf62cd 100644
--- a/zig/tests/cuda/glm_ref_index.zig
+++ b/zig/tests/cuda/glm_ref_index.zig
@@ -96,3 +96,68 @@ pub fn selection(r: *Rig, set: glm.glue.Set, ops: cuda.Module) !void {
     }
     check.pass("DSA indexer {d} rows at {d}: scores within float64 tolerance, top-{d} pools equal the host's rule exactly", .{ rows, pos, sp.topk_pools });
 }
+
+/// tf_glm_select_rows_wide against tf_glm_select_rows (bytes) and the host's rule, and both launches' times.
+pub fn selectionWide(r: *Rig, set: glm.glue.Set) !void {
+    var prng = std.Random.DefaultPrng.init(0x5e1ec7); // its own draws: r.rng's order stays the later checks'
+    const rng = prng.random();
+    var t0 = try cuda.Event.init(r.gpu.d, true);
+    defer t0.deinit();
+    var t1 = try cuda.Event.init(r.gpu.d, true);
+    defer t1.deinit();
+    const reps = 20;
+    const cases = [_][2]usize{ .{ 1500, 1 }, .{ 6000, 4 }, .{ 12000, 1 }, .{ 24000, 1 }, .{ 32768, 1 }, .{ 136000, 1 }, .{ 136000, 4 }, .{ 136000, 32 }, .{ 136000, 48 }, .{ 136000, 64 }, .{ 136000, 128 }, .{ 500000, 1 } };
+    for (cases) |c| {
+        const pos = c[0];
+        const rows = c[1];
+        const np = sp.poolBucket(pos, rows, 1 << 20);
+        const host = try r.a.alloc(f32, rows * np);
+        for (0..rows) |row| {
+            const vis = (pos + row + 1) / 4;
+            for (host[row * np ..][0..np], 0..) |*v, p| {
+                const level: f32 = @floatFromInt(@as(i32, @intCast(rng.uintLessThan(u32, 4096))) - 2048);
+                v.* = if (p >= vis) -std.math.inf(f32) else if (level == 0 and rng.boolean()) -0.0 else level * 0.25;
+            }
+        }
+        const scores = try r.dev(std.mem.sliceAsBytes(host));
+        const p0: i32 = @intCast(pos);
+        const at = try r.dev(std.mem.asBytes(&p0));
+        const outs = [2]u64{ try r.zeros(rows * sp.topk_pools * 8), try r.zeros(rows * sp.topk_pools * 8) };
+        var ms: [2]f64 = undefined;
+        for ([2]cuda.Function{ set.glue.select_rows, set.glue.select_rows_wide }, [2]u32{ 128, 1024 }, outs, 0..) |f, threads, out, i| {
+            var a: cuda.Args = .{};
+            inline for (.{ scores, out, @as(i64, @intCast(np)), at, @as(c_int, sp.topk_pools) }) |v| a.add(v);
+            const cfg: cuda.launch.Config = .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = threads } };
+            try cuda.launch.launch(f, cfg, r.s, &a);
+            try t0.record(r.s);
+            for (0..reps) |_| try cuda.launch.launch(f, cfg, r.s, &a);
+            try t1.record(r.s);
+            try t1.synchronize();
+            ms[i] = @as(f64, try cuda.Event.elapsedMs(t0, t1)) / reps;
+        }
+        const narrow = try r.back(outs[0], rows * sp.topk_pools * 8);
+        const wide = try r.back(outs[1], rows * sp.topk_pools * 8);
+        if (std.mem.indexOfDiff(u8, narrow, wide)) |b|
+            try check.expect(false, "select {d} rows at {d}: byte {d} of the wide kernel's differs from tf_glm_select_rows'", .{ rows, pos, b });
+        // the host's rule on the first and the last row
+        const picks = std.mem.bytesAsSlice(i64, wide);
+        for ([2]usize{ 0, rows - 1 }) |row| {
+            const vis = (pos + row + 1) / 4;
+            const limit = @min(np, @max(vis, sp.topk_pools));
+            const keys = try r.a.alloc(u32, limit);
+            for (keys, 0..) |*k, p| k.* = orderKey(host[row * np + p]);
+            const idx = try r.a.alloc(u32, limit);
+            for (idx, 0..) |*v, p| v.* = @intCast(p);
+            std.mem.sort(u32, idx, keys, struct {
+                fn lt(k: []const u32, x: u32, y: u32) bool {
+                    return k[x] > k[y] or (k[x] == k[y] and x < y);
+                }
+            }.lt);
+            const n = @min(limit, sp.topk_pools);
+            std.mem.sort(u32, idx[0..n], {}, std.sort.asc(u32));
+            for (idx[0..n], picks[row * sp.topk_pools ..][0..n], 0..) |want, have, i|
+                try check.expect(have == want, "select {d} rows at {d}, row {d} pick {d}: pool {d}, the host's rule gives {d}", .{ rows, pos, row, i, have, want });
+        }
+        check.pass("DSA select {d} rows at {d} ({d}-pool bucket): wide equals tf_glm_select_rows byte for byte and the host's rule; {d:.4} ms vs {d:.4} ms a launch", .{ rows, pos, np, ms[1], ms[0] });
+    }
+}
diff --git a/zig/tests/cuda/glm_ref_tests.zig b/zig/tests/cuda/glm_ref_tests.zig
index 8cf5b5fa7d..14e163d5ef 100644
--- a/zig/tests/cuda/glm_ref_tests.zig
+++ b/zig/tests/cuda/glm_ref_tests.zig
@@ -204,6 +204,7 @@ pub fn run(gpu: Gpu) !void {
     var ops = try cuda.Module.load(gpu.d, cuda.kernels.glm_ops);
     defer ops.unload();
     try @import("glm_ref_index.zig").selection(&r, set, ops);
+    try @import("glm_ref_index.zig").selectionWide(&r, set);
     try @import("glm_ref_hc.zig").gammas(&r, set); // last: its draws from r.rng follow every check above
     check.pass("glm-ref: every projection within {e:.3} of its float64 tolerance", .{r.worst});
 }

From 638de4f6a2f64dfdd47ede75d05b58485969698b Mon Sep 17 00:00:00 2001
From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com>
Date: Sun, 25 Oct 2026 23:08:20 +0900
Subject: [PATCH 2/3] contract: tf_glm_select_rows_wide replaces
 tf_glm_select_rows byte for byte

The wide select registers as a replacement of the 128-thread kernel (its bytes on the same scores), with a
stated NaN rule checked against the host's picks at gate 1: order keys descending, -0 as +0 and any NaN
above +inf, as the GPU's s + 0 gives one positive NaN. The index scores and the select register through one
owned group, so the runner stays within the source-size limit, as the Gemma tower's do. Gate 5's records
come from kernel-contract-limits over sm_80 to sm_121 fatbins of nvcc 13.0 and 13.3; 13.3 differs on
sm_120 and sm_121.
---
 docs/kernel-contract.md                    |   6 +-
 zig/src/core/contract/glm_cuda_records.zig |   2 +
 zig/tests/cuda/contract_glm_index.zig      |  23 +++
 zig/tests/cuda/contract_glm_select.zig     | 158 +++++++++++++++++++++
 zig/tests/cuda/contract_tests.zig          |   6 +-
 5 files changed, 190 insertions(+), 5 deletions(-)
 create mode 100644 zig/tests/cuda/contract_glm_select.zig

diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md
index af412e6887..1ed43ac146 100644
--- a/docs/kernel-contract.md
+++ b/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
diff --git a/zig/src/core/contract/glm_cuda_records.zig b/zig/src/core/contract/glm_cuda_records.zig
index 540ca63d80..c1a948bd8d 100644
--- a/zig/src/core/contract/glm_cuda_records.zig
+++ b/zig/src/core/contract/glm_cuda_records.zig
@@ -43,6 +43,7 @@ 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_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 +65,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 } } },
 };
diff --git a/zig/tests/cuda/contract_glm_index.zig b/zig/tests/cuda/contract_glm_index.zig
index 3c5736ac1e..4ee444259f 100644
--- a/zig/tests/cuda/contract_glm_index.zig
+++ b/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();
+    }
+};
diff --git a/zig/tests/cuda/contract_glm_select.zig b/zig/tests/cuda/contract_glm_select.zig
new file mode 100644
index 0000000000..e28f3ce49c
--- /dev/null
+++ b/zig/tests/cuda/contract_glm_select.zig
@@ -0,0 +1,158 @@
+//! GLM's decode select under the contract on CUDA: tf_glm_select_rows_wide, bit-equal to tf_glm_select_rows.
+const std = @import("std");
+const cuda = @import("cuda");
+const core = @import("core");
+const glm = @import("glm_cuda");
+const rig_ = @import("contract_rig.zig");
+
+const contract = core.contract;
+const c = contract.spec;
+const g = contract.gates;
+const sparse = glm.sparse;
+
+fn noSums(_: c.Shape) usize {
+    return 0;
+}
+
+const rule = "tf_glm_select_rows' bytes, and the host's picks: order keys descending, -0 as +0 and any NaN above +inf " ++
+    "(the GPU's s + 0 gives its one positive NaN), ties to the lower pool, the K picks in pool order";
+
+/// A decode step's few rows over long contexts' buckets, where cuda_sparse launches the wide form; edges at one pass.
+pub fn spec() c.Spec {
+    const k = sparse.topk_pools;
+    const served = &[_]c.Shape{ .{ .rows = 1, .n = 4096, .k = k }, .{ .rows = 4, .n = 65536, .k = k }, .{ .rows = 48, .n = 131072, .k = k } };
+    return .{ .name = "cuda glm select rows tf_glm_select_rows_wide", .accum = .none, .depth = noSums, .out = .raw, .replaces = "tf_glm_select_rows", .nan = .{ .stated = rule }, .tiles = .{ .{}, .{ .size = 8192, .min = k }, .{} }, .served = served };
+}
+
+/// sparse._order_key as the GPU computes it: s + 0 turns -0 into +0 and every NaN into one positive NaN.
+fn orderKey(s: f32) u32 {
+    const b: i32 = if (std.math.isNan(s)) std.math.maxInt(i32) else @bitCast(s + 0.0);
+    return @bitCast(b ^ ((b >> 31) | std.math.minInt(i32)));
+}
+
+fn before(keys: []const u32, x: u32, y: u32) bool {
+    return keys[x] > keys[y] or (keys[x] == keys[y] and x < y);
+}
+
+/// Each row's K pools of an n-pool bucket (shape: rows x n x K); the narrow kernel's bytes are its baseline.
+pub const Select = struct {
+    rig: rig_.Rig,
+    wide: cuda.Function,
+    narrow: cuda.Function,
+    shape: c.Shape = undefined,
+    host: []f32 = &.{},
+    at: usize = 0,
+    in: u64 = 0,
+    pos: u64 = 0,
+    out: u64 = 0,
+    base: u64 = 0,
+    msg: [160]u8 = undefined,
+
+    pub fn deinit(self: *Select) void {
+        self.rig.gpu.gpa.free(self.host);
+        self.rig.deinit();
+    }
+
+    pub fn kernel(self: *Select) g.Kernel {
+        return .{ .spec = spec(), .ctx = self, .vt = &vt };
+    }
+
+    const vt = g.hooks(Select);
+
+    fn visible(self: *const Select, r: usize) usize {
+        return (self.at + r + 1) / sparse.pool;
+    }
+
+    fn outBytes(self: *const Select) usize {
+        return self.shape.rows * self.shape.k * 8;
+    }
+
+    /// Quarter steps in [-512, 512) with -0 among the zeros, -inf past the visible pools; non-finite: rows 0-3 plants.
+    pub fn prepare(self: *Select, shape: c.Shape, fill: g.Fill, seed: u64) !bool {
+        if (shape.k != sparse.topk_pools or shape.n < shape.k or shape.rows == 0) return false;
+        self.rig.reset();
+        const a = self.rig.gpu.gpa;
+        a.free(self.host);
+        self.host = &.{};
+        self.shape = shape;
+        // a bucket's three quarters visible; past an edge's n every pool is (a limit of n), at 512 fewer than K are
+        self.at = if (std.math.isPowerOfTwo(shape.n)) 3 * shape.n else 4 * shape.n;
+        self.host = try a.alloc(f32, shape.rows * shape.n);
+        var prng = std.Random.DefaultPrng.init(seed);
+        const rng = prng.random();
+        for (0..shape.rows) |r| {
+            const vis = self.visible(r);
+            for (self.host[r * shape.n ..][0..shape.n], 0..) |*v, p| {
+                const level: f32 = @floatFromInt(@as(i32, @intCast(rng.uintLessThan(u32, 4096))) - 2048);
+                v.* = if (p >= vis) -std.math.inf(f32) else if (level == 0 and rng.boolean()) -0.0 else level * 0.25;
+            }
+        }
+        if (fill == .nonfinite) {
+            const neg_nan: f32 = @bitCast(@as(u32, 0xffc00000));
+            const put = [_]f32{ std.math.nan(f32), neg_nan, std.math.inf(f32), -std.math.inf(f32) };
+            for (put, 0..) |v, r| if (r < shape.rows) {
+                const step = @min(self.visible(r), shape.n) / 8;
+                for (0..8) |e| self.host[r * shape.n + e * step + r] = v;
+            };
+        }
+        const p: i32 = @intCast(self.at);
+        self.in = try self.rig.input(std.mem.sliceAsBytes(self.host));
+        self.pos = try self.rig.input(std.mem.asBytes(&p));
+        self.out = try self.rig.output(self.outBytes());
+        self.base = try self.rig.scratch(self.outBytes());
+        return true;
+    }
+
+    pub fn run(self: *Select) !void {
+        try self.launch(self.wide, 1024, self.out);
+        try self.rig.s.synchronize();
+    }
+
+    fn launch(self: *Select, f: cuda.Function, threads: u32, out: u64) !void {
+        const x = self.shape;
+        var a: cuda.Args = .{};
+        inline for (.{ self.in, out, @as(i64, @intCast(x.n)), self.pos, @as(c_int, @intCast(x.k)) }) |v| a.add(v);
+        try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(x.rows) }, .block = .{ .x = threads } }, self.rig.s, &a);
+    }
+
+    /// tf_glm_select_rows on the same scores, a 128-thread block a row as cuda_sparse launches it.
+    pub fn replaced(self: *Select) ![]const u8 {
+        try self.launch(self.narrow, 128, self.base);
+        return self.rig.bytes(self.base, self.outBytes());
+    }
+
+    pub fn count(self: *Select) usize {
+        return self.outBytes();
+    }
+
+    pub fn reference(_: *Select, want: []f64, mag: []f64) !void {
+        @memset(want, 0);
+        @memset(mag, 0);
+    }
+
+    /// tf_glm_select_rows' bytes, then every row's picks against the host's rule over its first max(visible, K) pools.
+    pub fn check(self: *Select, _: g.Gate) !?[]const u8 {
+        const a = self.rig.gpu.gpa;
+        const x = self.shape;
+        const got = try a.dupe(u8, try self.rig.outputBytes());
+        defer a.free(got);
+        if (std.mem.indexOfDiff(u8, got, try self.replaced())) |at| return try std.fmt.bufPrint(&self.msg, "byte {d} of {d} differs from tf_glm_select_rows'", .{ at, got.len });
+        const keys = try a.alloc(u32, x.n);
+        defer a.free(keys);
+        const idx = try a.alloc(u32, x.n);
+        defer a.free(idx);
+        for (0..x.rows) |r| {
+            const limit = @min(x.n, @max(self.visible(r), x.k));
+            for (keys[0..limit], self.host[r * x.n ..][0..limit]) |*k, s| k.* = orderKey(s);
+            for (idx[0..limit], 0..) |*v, p| v.* = @intCast(p);
+            std.mem.sort(u32, idx[0..limit], keys[0..limit], before);
+            const picks = idx[0..x.k];
+            std.mem.sort(u32, picks, {}, std.sort.asc(u32));
+            for (picks, 0..) |want, i| {
+                const have = std.mem.readInt(i64, got[(r * x.k + i) * 8 ..][0..8], .little);
+                if (have != want) return try std.fmt.bufPrint(&self.msg, "row {d} pick {d}: pool {d}, the host's rule gives {d}", .{ r, i, have, want });
+            }
+        }
+        return null;
+    }
+};
diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig
index 65af43f23c..fd53c75789 100644
--- a/zig/tests/cuda/contract_tests.zig
+++ b/zig/tests/cuda/contract_tests.zig
@@ -535,9 +535,9 @@ pub fn run(gpu: check.Gpu) !void {
     var glm_pool = glm_pages_k.PoolKeys{ .rig = .{ .gpu = gpu, .s = s }, .f = glm_fns.pool_keys };
     defer glm_pool.deinit();
     try kernels.append(a, glm_pool.kernel());
-    var glm_scores = glm_index_k.Scores{ .rig = .{ .gpu = gpu, .s = s }, .f = glm_fns.scores };
-    defer glm_scores.deinit();
-    try kernels.append(a, glm_scores.kernel());
+    var glm_index = glm_index_k.Index.init(gpu, s, glm_fns);
+    defer glm_index.deinit();
+    try glm_index.append(a, &kernels);
     const glm_fast = try @import("contract_glm_pools.zig").Fast.create(a, gpu, s, glm_fns, &kernels);
     defer glm_fast.destroy(a);
     var glm_ops = try cuda.Module.load(gpu.d, kk.glm_ops);

From b8f755d96f57b9dcb72786e3a5c36f6f3f6b9125 Mon Sep 17 00:00:00 2001
From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com>
Date: Tue, 27 Oct 2026 23:19:48 +0900
Subject: [PATCH 3/3] contract: tf_glm_select_rows (128 threads) records its
 launch, a changed kernel

It is now select_row<128>, and it still serves every prompt chunk and every decode step outside the wide
rule, so gate 5 checks it as it checks the wide form. The numbers come from kernel-contract-limits over
sm_80 to sm_121 fatbins of this tree: nvcc 13.0 and 13.3 give the same registers and static shared bytes on
all five SMs, so 13.3's set needs no entry for it.
---
 zig/src/core/contract/glm_cuda_records.zig | 1 +
 1 file changed, 1 insertion(+)

diff --git a/zig/src/core/contract/glm_cuda_records.zig b/zig/src/core/contract/glm_cuda_records.zig
index c1a948bd8d..6d23ed6650 100644
--- a/zig/src/core/contract/glm_cuda_records.zig
+++ b/zig/src/core/contract/glm_cuda_records.zig
@@ -43,6 +43,7 @@ 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 } } },