Skip to main content

Module topk_gpu

Module topk_gpu 

Source
Expand description

STELLAR §A A1a — native GPU dispatch of the top-K reduction + on-device parity.

Runs shaders/topk_reduction.wgsl on a real wgpu device: each workgroup reduces a block of the logits to its top-K candidates, the host merges to the global top-K. The #[test] below verifies it on silicon against the CPU oracle topk::topk_cpu (skips cleanly with no adapter). This is the reusable core that the decode-loop splice (behind QUALIA_LLM_GPU_TOPK) will call instead of the full-logit-readback dispatch_output_argmax_chunked.

Native only — mirrors ternary_gpu.rs.

Functions§

topk_gpu
Reduce logits to its global top-K on the GPU. block_size is elements per workgroup (clamped to TOPK_BLOCK_SIZE, the WGSL var<workgroup> cap). Blocking (native readback).