Documentation
¶
Overview ¶
Package q4 generates the WGSL of a matrix-vector product for ggml Q4_0 and Q4_1 weights, f32 activations, with no subgroups and no f16.
The weights are repacked by the caller: the 16 quant bytes of a block as one aligned vec4<u32> (array<vec4<u32>>, one element per block, row-major), the block scale (and for Q4_1 the minimum) as f32 arrays. A workgroup of WG invocations computes RPW = WG/TPR output rows; TPR invocations share a row and each takes the blocks b = lane, lane+TPR, ... of it; the input vector is staged once per workgroup in workgroup memory (in tiles when it does not fit MaxShared) and the TPR partial sums of a row are added by a tree reduction in workgroup memory. Sizes come from the adapter's limits, never from a vendor's values; the summation order depends on WG and TPR only.
The z axis of the dispatch is a batch of independent input vectors (tokens): vector tk is read at tk*N floats and row r of vector tk written at tk*Rows+r, so a chunk of vectors reads every weight once.
Callers add fused prologues and epilogues through Hooks (an RMSNorm in front, a SiLU, RoPE or attention merge behind); Store and Acc are the plain ones. The shaders are generated per model: dimensions are WGSL constants.
Index ¶
Constants ¶
const ( EpiStore = ` if (ln == 0u && valid) { y[tk * ROWS + row] = red[li]; }` EpiAcc = ` if (ln == 0u && valid) { y[tk * ROWS + row] = y[tk * ROWS + row] + red[li]; }` )
EpiAcc and EpiStore are the plain epilogues (the y binding is last).
const Prelude = `` /* 615-byte string literal not displayed */
Prelude is the workgroup reduction helpers {{WG}} wide: rsum and rmax add or take the maximum of one value per invocation by a tree reduction and return it to all of them. Replace {{WG}} before use.
const RowPlain = `let row = wg * RPW + lr; let valid = row < ROWS;`
RowPlain maps workgroup, row in the workgroup to row = wg*RPW + lr.
Variables ¶
This section is empty.
Functions ¶
Types ¶
type Binding ¶
type Binding struct {
Name, Type string // WGSL name and type, e.g. "pv", "array<vec4<f32>>"
RW bool
}
Binding is an extra storage binding, numbered after the standard ones.
type Hooks ¶
type Hooks struct {
Bindings []Binding // after qs, scale, [mm], [x4]; y is added by Store and Acc
Decl string // module-scope declarations (constants)
Func string // module-scope functions
NoInput bool // no x4 binding: Load computes the input
Load string // expression of type vec4<f32>: input element gi (vec4 index) of vector tk
RowMap string // statements defining `row` (u32) and `valid` (bool)
Epi string // statements after the reduction: red[lr*TPR] is row lr's sum
}
Hooks are the fused parts of a kernel, as WGSL text. In scope: wg (workgroup index in the 2-D grid), tk (z, the vector), lr (row in the workgroup), ln (lane in the row), li, the constants WG TPR RPW N NB ROWS NX EPS, red.
type Spec ¶
type Spec struct {
WG, TPR int // workgroup size and invocations per row (powers of two, TPR | WG)
MaxGroups int // the adapter's workgroups per dimension
N, Rows int // input length (a multiple of 32) and output rows
Eps float32
Q41 bool // Q4_1 (scale d and minimum m) instead of Q4_0
// Norm puts an RMSNorm with weights in front: the scales binding is "aux"
// holding the block scales followed, at float NW, by the N weights.
Norm bool
NW int
Mode Mode
Hooks Hooks
}
Spec describes one kernel.
func (Spec) Tile ¶
Tile returns the staged tile of the input: in vec4s, the number of tiles and the blocks per tile. The tile and the reduction array fit MaxShared.
func (Spec) Workgroups ¶
Workgroups is the number of workgroups for a vector (before the z axis) in the 1-D count; a Custom mode whose workgroups do not map RPW rows each passes its own count to Source.
Source Files