Module topk 

STELLAR §A A1a — GPU top-K reduction: CPU oracle, host merge, and the WGSL kernel.

Decode is memory-bandwidth-bound; the sentinel/sampler only needs the high-probability mass, not the 49 k near-zero tail. Instead of reading back the full logit vector (~196 KB/token) and doing a CPU argmax, the GPU reduces each block of the vocabulary to its top-K candidates; the host merges those into the global top-K. The CPU oracle here is the byte-for-byte reference the on-device kernel is verified against (see topk_gpu.rs), exactly as ternary.rs anchors ternary_gpu.rs.

Contract (must match shaders/topk_reduction.wgsl): NaN → −∞ (never selected); ties broken toward the LOWER token id (deterministic); K=1 == argmax.

Structs

TopKItem

One top-K entry: a token id and its raw logit.

Constants

TOPK_BLOCK_SIZE

Block size each workgroup reduces — must equal MAX_BLOCK in the WGSL (var<workgroup> cap).

TOPK_MAX_K

Largest K the host paths support (kept generous; the kernel itself is K-agnostic per round).

TOPK_REDUCTION_WGSL

The block-reduction kernel (auto bind-group layout; entry topk_block).

Functions

merge_block_candidates

Merge per-block GPU candidates (num_blocks × k pairs) into the global top-K. Blocks cover disjoint index ranges, so candidate ids are unique. Drops −∞ entries and any id in masked (the governance/sieve veto — “a masked token never returned”).

topk_cpu

CPU reference top-K over a full logit vector. Drops −∞ (e.g. masked) entries.

topk_params_bytes

16-byte Params uniform: n, k, block_size, cand_base (cand_base=0 for single-chunk).

topk_params_bytes_with_base

Like topk_params_bytes with a non-zero candidate write base for multi-chunk mega-pass.