Skip to main content

TOPK_REDUCTION_WGSL

Constant TOPK_REDUCTION_WGSL 

Source
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).