Documentation
¶
Overview ¶
Package argmax generates the WGSL of a greedy sampler in two stages, so that a large vocabulary is not walked by one workgroup: Groups workgroups each reduce a slice of the logits to its (largest, lowest index among equals), then one workgroup reduces the partials and writes the token. Ties go to the lowest index, like llama.cpp's greedy sampler.
Stage 1 bindings: 0 logits (read f32), 1 partial values (read_write f32, Groups), 2 partial indices (read_write u32, Groups); dispatch Groups x 1. Stage 2 bindings: 0 values, 1 indices (read), 2 tokens (read_write u32), 3 out (read_write u32), 4 state (read u32); dispatch 1 x 1. It writes the token to tokens[state[1]+1] and to out[0].
Index ¶
Constants ¶
const ( Access1 = "rww" Access2 = "rrwwr" )
Access strings of the two stages, one letter per binding.
const Groups = 64
Groups is the number of workgroups of stage 1.
Variables ¶
This section is empty.
Functions ¶
Types ¶
This section is empty.
Source Files