pub const TOPK_REDUCTION_WGSL: &str = "// topk_reduction.wgsl \u{2014} STELLAR \u{a7}A A1a: GPU block top-K reduction (decision D4/D5).\n//\n// Each workgroup reduces a contiguous BLOCK of `logits` to its local top-K\n// (value, global-index) candidates via K iterative parallel-argmax passes\n// (NOT a full-vocab bitonic sort \u{2014} we need top-K with K << N). The host merges\n// every block\'s K candidates into the global top-K. This replaces the per-token\n// 196 KB full-logit readback + CPU argmax with a ~(num_blocks \u{d7} K)-pair readback.\n//\n// Semantics: NaN \u{2192} -inf (never selected); ties broken toward the LOWER index\n// (deterministic). K=1 degenerates to a parallel argmax.\n//\n// Bindings (auto layout, native): 0=logits(read) 1=params(uniform)\n// 2=cand_val(read_write) 3=cand_idx(read_write).\n\nconst WG: u32 = 64u;\nconst MAX_BLOCK: u32 = 1024u;\nconst NEG_INF: f32 = -3.4028235e38;\n\nstruct Params {\n n: u32, // total number of logits\n k: u32, // top-K per block (caller guarantees k >= 1)\n block_size: u32, // elements per workgroup (<= MAX_BLOCK)\n /// Base index into cand_val/cand_idx for this chunk (multi-chunk mega-pass).\n cand_base: u32,\n};\n\n@group(0) @binding(0) var<storage, read> logits: array<f32>;\n@group(0) @binding(1) var<uniform> params: Params;\n@group(0) @binding(2) var<storage, read_write> cand_val: array<f32>;\n@group(0) @binding(3) var<storage, read_write> cand_idx: array<u32>;\n\n// Block logits, mutated (masked) across the K rounds.\nvar<workgroup> s_val: array<f32, MAX_BLOCK>;\n// Per-thread reduction scratch.\nvar<workgroup> r_val: array<f32, WG>;\nvar<workgroup> r_idx: array<u32, WG>;\n\n@compute @workgroup_size(WG)\nfn topk_block(\n @builtin(workgroup_id) wg: vec3<u32>,\n @builtin(local_invocation_id) lid: vec3<u32>,\n) {\n let tid = lid.x;\n let blk = wg.x;\n let base = blk * params.block_size;\n let bsize = min(params.block_size, MAX_BLOCK);\n\n // Load block into shared memory; NaN and out-of-range \u{2192} -inf.\n var i = tid;\n loop {\n if (i >= bsize) { break; }\n let g = base + i;\n var v = NEG_INF;\n if (g < params.n) {\n let raw = logits[g];\n if (raw == raw) { v = raw; } // NaN != NaN \u{21d2} leaves -inf\n }\n s_val[i] = v;\n i = i + WG;\n }\n workgroupBarrier();\n\n // K iterative parallel-argmax passes.\n var round = 0u;\n loop {\n if (round >= params.k) { break; }\n\n // Each thread scans its strided slice for a local best (lowest idx on ties).\n var best_v = NEG_INF;\n var best_i = 0u;\n var j = tid;\n loop {\n if (j >= bsize) { break; }\n let v = s_val[j];\n if (v > best_v) { best_v = v; best_i = j; }\n j = j + WG;\n }\n r_val[tid] = best_v;\n r_idx[tid] = best_i;\n workgroupBarrier();\n\n // Tree reduction; on equal values keep the lower index.\n var stride = WG / 2u;\n loop {\n if (stride == 0u) { break; }\n if (tid < stride) {\n let ov = r_val[tid + stride];\n let oi = r_idx[tid + stride];\n let cv = r_val[tid];\n let ci = r_idx[tid];\n if (ov > cv || (ov == cv && oi < ci)) {\n r_val[tid] = ov;\n r_idx[tid] = oi;\n }\n }\n workgroupBarrier();\n stride = stride / 2u;\n }\n\n // Thread 0 emits the winner (global index) and masks it for the next round.\n if (tid == 0u) {\n let widx = r_idx[0];\n let out = params.cand_base + blk * params.k + round;\n cand_val[out] = r_val[0];\n cand_idx[out] = base + widx;\n s_val[widx] = NEG_INF;\n }\n workgroupBarrier();\n round = round + 1u;\n }\n}\n";Expand description
The block-reduction kernel (auto bind-group layout; entry topk_block).