Skip to main content

Module plan

Module plan 

Source
Expand description

A step graph lowered to launches of the kernel library. The lane decides no fusion of its own: the step, lifted into the one IR, is optimized by the shared table under the lane’s shim (LaneShim), which answers whether a fused class has a launch; each op the optimizer leaves takes the launches its link names (link), a class the lane runs whole one launch, a class it runs as its parts each part’s. Then the lane’s schedule sets how the launches run — four rows a warp for a wide product (rows_a_warp), the shared-row attention, the split counts, a batch’s sequences as one launch (sequences) — and every tensor is given its buffers.

A row read at a literal index of a computed tensor (row 4 of bxo) is no launch at all: a view into the tensor — unless the tensor is a table written in place before the row’s last reader, where the row is copied when taken. A view a projection reads quantized is a view into its tensor’s quantized form too (when it starts on a whole block), so the tensor is quantized once, in its producer, for every row a projection reads of it.

Re-exports§

pub use crate::tensor_graph::graph::SeqIo;
pub use crate::tensor_graph::graph::StepIo;

Structs§

Advance
What the arg-max that closes a step writes besides its result: the token the next decode step reads (every form’s arg-max writes the decode step’s; a batch’s sequence item writes item item of its tokens) and the position the sequence’s next step reads; and the run generated so far, which only the host reads.
Buffer
One tensor’s buffers.
ConvPart
One causal convolution of a merged launch (Launch::ConvParts): its input, weights, ring and output, its channels, and whether each head of its output is normed to unit length.
Dst
Where a kernel writes its result: a buffer, or the rows of a table of rows pitch wide from a starting row — the whole of it, or (unit) a part of every unit of it.
Fusion
A fusion rule that fired: its name and the graph ops it joined.
GemmPart
A part of a GEMM in parts: its matrix, bias, destination and outputs a row.
Part
A part of every unit of a wider tensor: unit u’s values start u·pitch + offset in.
PreNorm
The norm a norm launch takes of its added operand in the same block before adding it (sandwich_norm): the weight and epsilon of the norm that operand’s own launch took.
RopeSide
One rotation of a merged launch (Launch::RopeParts): its input (read in place as a part of every unit of a wider tensor where it names one), its norm, its destination, and its heads, per_row of them a row.
Schedule
A step lowered: its launches in order (with the graph ops each computes — none for a quantize), its buffers, and the fusions that made the launches.
SharedExperts
The shared experts a mixture takes into its own launches (shared_expert): one SwiGLU FFN of dff, its gate, up and down projections read in place in their formats.
Split
An attention’s queries and keys in two parts, split at own: a query’s first own values its unit’s in q (q_pitch apart a (row, head)), the rest its unit’s in q_tail; a key’s first own its group’s in the keys, the rest its row’s in k_tail, every group’s alike — an expanded latent attention’s queries, each head’s unrotated part read in place in the query projection and its rotated part, and its keys, each head’s expansion and the rotated key the heads share, read in place in the cache.

Enums§

DType
The element type of a buffer.
Fill
What a buffer holds and how it is filled before the first step.
Launch
One launch of a library kernel. A producer names the tensor it computes (out); f32 when it writes the tensor’s f32 buffer, quant when it writes its quantized form. Row counts are the block’s (1 for a decode step). A part is where a launch reads its input x unit by unit: a part of every unit of a wider tensor (none, the whole of x).
Router
How an expert mixture’s router weights are held: quantized checkpoint rows of a format, or f32 values (an expert’s row of the input’s width at a time) — the router reads the input f32 then, under every numerics.
RowSel
Which row of a table a write starts at: the row a parameter names (1-based), or a fixed row.
Touch
How a launch touches a tensor: through its f32 (or i32) buffer, or its quantized activation.

Constants§

GEMV_ROWS
The most rows the GEMV kernel multiplies in one launch (its warp per output row holds an accumulator per row); a projection of more rows runs on the tensor cores, or (under the precise numerics) as slabs of this many rows.
MAX_CHUNK
The largest prefill chunk a lane compiles.
QBLOCK
The values of an activation’s quantized form that share one scale.
RULES
The steps of a lane a plan can hold back (plan’s held), in the order the lane takes them: the shared rules its kernels link, by their names in the shared table (crate::tensor_graph::fuse); rows_a_warp, the schedule’s four rows a warp for a wide product; the decode step’s products of one activation as one GEMV in parts (gemv_parts, the lane’s own rule in the table); and sequences, a batch’s launches over its sequences as one. The arg-max’s advance closes the step and is always taken; the activation quantizer’s fusion is the fuse_quant knob.

Functions§

attended_span
The most positions one query of an attention over a cache of len rows attends: the cache, or its window.
gates_batch_rows
Whether op is a batch’s gated attentions (attn_output_gate over its sequences’ rows), or their fold with the quantize of the product: the op the lane runs as each sequence’s attention and its row’s set.
kind
A launch’s kind, as a census counts it.
plan
The step graph g lowered under numerics and knobs: constants holds the placement of every constant the graph declares, io the step’s carried names, held the steps of RULES held back — a shared rule, or a step of the lane’s schedule.
quant_in_place
A view offset elements into a tensor reads the tensor’s quantized form in place: it starts on a whole block.