Skip to main content

Module gpu_sample

Module gpu_sample 

Source
Expand description

ADR-040 §26 iter-M — GPU-side first-max argmax + threshold candidate collect.

Replaces the host full-vocab argmax + candidate-threshold scans (~0.92ms/step on the autoregressive critical path) with one GPU dispatch that reads back only the per-slot top1 + the few threshold candidates. The host keeps the cheap F64 rerank over those candidates (Metal has no f64). Byte-matches the host argmax_f32_first_max (first-max, lower-index tie-break) + the finalize threshold scan (logits >= top1_val - 0.5f).

Statics§

GPU_SAMPLE_SHADER_SOURCE

Functions§

dispatch_gpu_sample_argmax_candidates
Dispatch GPU argmax+candidate-collect over [n_slots, vocab] logits.
register