Skip to main content

topk_gpu

Function topk_gpu 

Source
pub fn topk_gpu(
    device: &Device,
    queue: &Queue,
    logits: &[f32],
    k: usize,
    block_size: usize,
) -> Vec<TopKItem>
Expand description

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