Documentation
¶
Overview ¶
Package tensor is the layer that turns a model into device work.
You describe a computation once as a graph of tensors, compile it into a Plan, and submit that plan with different inputs as often as you like. The plan owns its intermediate memory and its device graph; you own the buffers you named.
Why compile once ¶
A transformer runs the same shapes thousands of times. Deciding which kernel to use, where the intermediates live, and which barriers are needed is work that depends on the shapes and not on the numbers, so it is done once:
rt, _ := tensor.NewRuntime(dev)
b := rt.NewBuilder("mlp")
x := tensor.Input(b, tensor.ValueDesc{Name: "x", DType: accel.F32, Shape: tensor.Shape{4, 8}})
w := tensor.Weight(b, tensor.ValueDesc{Name: "w", DType: accel.F32, Shape: tensor.Shape{4, 8}})
tensor.Output(b, "y", tensor.Mul(b, x, w))
plan, err := b.Compile(rt, tensor.CompileOptions{Label: "mlp"})
Then submit it as often as you like, rebinding what changed.
Errors do not interrupt you ¶
An operator that cannot be built records the problem and returns a poisoned tensor, so model code has no error branch per line. A poisoned tensor produces poisoned results and no further errors: one mistake gives one diagnostic, and Builder.Compile returns them all together, each naming the operator, the operand, and the line of your code that made it.
What it does not do ¶
No quantization, no sampling, no automatic plan cache, and no scheduler across requests. Those are deliberately above this package. See specs/007-tensor-layer.md.
Index ¶
- Constants
- func DeclareSamplingScalars(b *Builder, o SamplingOptions, prefix string)
- func LinearAttention(b *Builder, q, k, v *Tensor, s *State, o LinearOptions) (*Tensor, *State)
- func Output(b *Builder, name string, x *Tensor)
- func Scalar(b *Builder, d ScalarDesc)
- type AttentionOptions
- type Bindings
- type BlockPool
- type Buckets
- type Builder
- type CompileOptions
- type DType
- type Identity
- type Int4
- type KernelSelection
- type LinearOptions
- type Plan
- type PlanCache
- type PortDesc
- type PortKind
- type Quantized
- type Runtime
- type SamplingOptions
- type ScalarDesc
- type ScalarKind
- type ScalarValue
- type Shape
- type SoftmaxOptions
- type State
- type StateDesc
- type Stream
- type Tensor
- func Add(b *Builder, x, y *Tensor) *Tensor
- func Argmax(b *Builder, logits *Tensor) *Tensor
- func Attention(b *Builder, q *Tensor, k, v *State, opts AttentionOptions) *Tensor
- func Broadcast(b *Builder, x *Tensor, shape Shape) *Tensor
- func Cast(b *Builder, x *Tensor, to DType) *Tensor
- func Contiguous(b *Builder, x *Tensor) *Tensor
- func GatherRows(b *Builder, table, ids *Tensor) *Tensor
- func GroupedMatMul(b *Builder, x, w, counts *Tensor) *Tensor
- func GroupedMatVec(b *Builder, x, w, counts *Tensor) *Tensor
- func Input(b *Builder, d ValueDesc) *Tensor
- func Int4MatMul(b *Builder, a *Tensor, w Int4) *Tensor
- func Int4MatVec(b *Builder, a *Tensor, w Int4) *Tensor
- func Linear(b *Builder, x, w, bias *Tensor) *Tensor
- func MatMul(b *Builder, x, w *Tensor) *Tensor
- func Mul(b *Builder, x, y *Tensor) *Tensor
- func Permute(b *Builder, x *Tensor, axes ...int) *Tensor
- func QuantGatherRows(b *Builder, table Quantized, ids *Tensor) *Tensor
- func QuantMatMul(b *Builder, x *Tensor, w Quantized) *Tensor
- func RMSNorm(b *Builder, x, gain *Tensor, eps float32) *Tensor
- func ReadState(b *Builder, s *State) *Tensor
- func Reshape(b *Builder, x *Tensor, shape Shape) *Tensor
- func RoPE(b *Builder, x *Tensor, rotaryDim int, baseName string, positions *Tensor) *Tensor
- func Sample(b *Builder, logits, draws *Tensor, history, counts *State, o SamplingOptions, ...) *Tensor
- func SampleCategorical(b *Builder, weights, draws *Tensor) *Tensor
- func Scale(b *Builder, x *Tensor, scalarName string) *Tensor
- func SiLU(b *Builder, x *Tensor) *Tensor
- func Slice(b *Builder, x *Tensor, axis, start, end int) *Tensor
- func Softmax(b *Builder, x *Tensor, opts SoftmaxOptions) *Tensor
- func SwiGLU(b *Builder, gate, value *Tensor) *Tensor
- func TopKMask(b *Builder, weights *Tensor, k int) *Tensor
- func TopPMask(b *Builder, weights *Tensor, p float32) *Tensor
- func Transpose(b *Builder, x *Tensor, axisA, axisB int) *Tensor
- func Weight(b *Builder, d ValueDesc) *Tensor
- type ValueDesc
Constants ¶
const MinTemperature = 1e-3
MinTemperature is the smallest positive temperature SamplingOptions.Validate accepts.
A guardrail and only a guardrail. It is not the point at which the arithmetic stops working -- that cliff is far below it and is refused separately -- it is the point below which a caller almost certainly meant to ask for greedy decoding and should be told to say so.
const TopMaxRounds = testkernels.TopMaxRounds
TopMaxRounds is how many entries either truncation can keep.
Re-exported from the corpus so a caller can write the check their own configuration has to pass, rather than discovering the bound by watching a top-k keep fewer entries than it was asked for. specs/028-sampling.md states it: both masks walk the distribution one entry per round, so this is a real limit and not a buffer size.
Variables ¶
This section is empty.
Functions ¶
func DeclareSamplingScalars ¶
func DeclareSamplingScalars(b *Builder, o SamplingOptions, prefix string)
DeclareSamplingScalars declares the scalars Sample reads for this policy.
It exists so that the names SamplingOptions.Scalars produces and the names the graph reads come from one place. Declaring them by hand at the call site is how a step ends up binding a value nobody reads, which is silent.
func LinearAttention ¶
LinearAttention steps a gated delta recurrence and returns this step's output.
What it is, and why it is not Attention ¶
A linear-attention layer carries a **matrix per sequence per head** rather than a cache per position:
S_t = S_{t-1}(alpha_t I - beta_t k_t k_t^T) + beta_t v_t k_t^T, o_t = S_t q_t
The state does not grow with context, which is the entire appeal — a 262K-context model has no KV cache for these layers — and it is also why nothing here takes Lengths or Pages: there are no positions to address.
The state is an ordinary State ¶
tensor.State with shape [slots, heads, valueDim, keyDim]: its leading axis is the sequence slot rather than a position, which specs/043-per-row-values.md §9's correction established. So a hybrid model — three of these layers for every softmax one — holds two States of different shapes in one graph and needs nothing else.
The recurrence is sequential, and that is visible in the cost ¶
Token t needs the state token t-1 left behind, so this is a scan rather than a batch of independent rows. The parallelism is over sequences and heads. specs/047-linear-attention.md §4 records the chunked parallel form as deliberately not built: it is what makes the layer fast rather than expressible, and this kernel is the reference it would be checked against. # It returns two things, because it writes two
The step's output *and* the next version of the state. A recurrence reads the state and writes it in one kernel, so unlike a KV cache -- where ScatterRows writes and Attention reads, as two nodes -- there is no way to separate them. Returning the version is what lets a later operator say which contents it meant, and specs/007-tensor-layer.md makes that distinction meaningful rather than decorative: an operator handed the version from *before* this step reads what was there before it.
func Output ¶
Output declares a value the caller wants written into a bound buffer.
Undeclared intermediates are inaccessible after compilation, which is what lets the planner alias them.
func Scalar ¶
func Scalar(b *Builder, d ScalarDesc)
Scalar declares a named per-step value an operator may read.
Declared rather than inferred from use, so a misspelled operator argument is an error instead of a new value nobody binds. specs/007-tensor-layer.md gives that reason and it is worth keeping: what it prevents is a plan that compiles, runs, and reads zero.
The value may change on every submission. What may not change is anything structural: an attribute that alters a shape, a layout, or which kernel is selected needs another plan, because the barriers and the transient layout were computed from it.
Types ¶
type AttentionOptions ¶
type AttentionOptions struct {
// Lengths is a u32 tensor holding, per sequence, how much of that
// sequence's cache holds real tokens.
//
// A tensor rather than a scalar, and specs/043-per-row-values.md draws the
// line: a value every row of a dispatch shares is a uniform, and a value
// that differs per row is device data. Cache lengths are genuinely
// independent across a batch -- that is what continuous batching *is* --
// so no scalar can express them, and no scheduling arrangement makes one
// correct.
//
// One sequence binds a one-element tensor. That is the same path, not a
// special case.
Lengths *Tensor
// Pages is an optional u32 page table: entry [s][i] is the physical block
// holding sequence s's i-th logical block.
//
// Nil means the cache is contiguous, which is the same thing with an
// identity table and a block size of one. It is nil-able rather than
// required because the indirection is a real cost in the innermost loop of
// decode, and a contiguous cache should not pay it -- the operator selects
// the kernel and reports which in [Plan.Selections].
//
// Paging is not a second kind of cache. A State addressed through a page
// table is the same State; what differs is the binding.
Pages *Tensor
// Block is how many positions one physical block holds, and is required
// when Pages is set.
Block int
// ScaleName is a declared f32 scalar, conventionally 1/sqrt(headDim). Named
// rather than computed here so a caller can use a different convention
// without a different plan.
ScaleName string
// BaseName is a declared u32 scalar, required for a prefill: the position
// of its first query token within the cache. It decides what the causal
// mask hides, so a prefill that extends an existing cache masks correctly
// rather than letting its first token see nothing.
//
// Still a scalar because a prefill is one sequence: a ragged step derives
// each token's position from [AttentionOptions.QueryExtents] and Lengths
// instead, so there is no row for this to differ across.
BaseName string
// QueryExtents is a u32 tensor holding, per sequence, how many query tokens
// that sequence contributes to this step. Setting it makes q flat --
// [sum(QueryExtents), qHeads, headDim] -- so sequences may contribute
// different numbers of tokens, which is what lets a prefill chunk share a
// dispatch with decode steps.
//
// specs/046-segmented-extents.md. This is the segmented extent that spec
// states, and attention is its first caller rather than its owner: the same
// primitive is what a grouped GEMM's per-expert token counts are.
//
// # It is not Lengths
//
// Both are u32 and one per sequence, and they are different numbers.
// Lengths is how many cached positions a sequence attends *over*;
// QueryExtents is how many query tokens it contributes *this step*. A step
// mixing a 512-token chunk with three decodes has QueryExtents
// [512, 1, 1, 1] and Lengths whatever those four caches hold.
//
// # A count of zero is legal
//
// A sequence admitted this step with nothing to contribute yet is an
// ordinary member of the batch. It occupies no rows of q and is not an
// error.
//
// # Rows of q past the total are padding
//
// The extents may sum to fewer rows than q holds. The extra rows belong to
// no sequence: they attend nothing and their output is zero, so a bucketed
// batch can pad q to a plan shape rather than inflating a real sequence's
// extent to absorb the difference.
//
// This is not checked here, and cannot be. The sum lives in a tensor, so it
// is device data by specs/043-per-row-values.md §2 and no value at record
// time can be compared against q.shape[0]. The kernel enforces it, which is
// why the behaviour is defined rather than left to a caller's discipline --
// specs/046-segmented-extents.md §1 property 3 and its correction.
QueryExtents *Tensor
}
AttentionOptions carries what attention needs beyond its operands.
type Bindings ¶
type Bindings struct {
Buffers map[string]accel.BufferView
Scalars map[string]ScalarValue
}
Bindings is everything a submission needs from the caller.
type BlockPool ¶
type BlockPool struct {
// contains filtered or unexported fields
}
BlockPool hands out fixed-size pieces of one KV cache so sequences of different lengths can share it.
pool, err := tensor.NewBlockPool(blocks, positionsPerBlock) pages, err := pool.Grow(nil, 40) // enough blocks for 40 positions // bind pages as Attention's AttentionOptions.Pages pool.Free(pages)
Why it is public now and was not ¶
It lived in tensor/internal/pagetable, whose doc said it was internal "until an operator accepts a page table": six exported declarations with no consumer would have let a caller build a pool, allocate blocks, and have nowhere to put the result. Attention takes a page table through AttentionOptions.Pages, so the condition that doc named has been met and this re-exports as it said it would. specs/030-paged-kv.md §4.1.
A sequence's page table is the list of physical blocks it holds, and the blocks need not be adjacent — which is the whole point. A cache sized for the longest sequence a server will ever see holds a fraction of what it reserved when the sequence is short, and that reservation is device memory nothing else can use.
It does not evict ¶
BlockPool.Grow fails when the pool is empty rather than choosing a victim. Choosing one is a policy question about which sequence matters, and a wrong answer silently truncates somebody's context — which is the same reason specs/029-plan-cache.md refuses to truncate a prompt. A caller who wants eviction implements the policy they can defend and frees the pages themselves.
specs/030-paged-kv.md has the addressing.
func NewBlockPool ¶
NewBlockPool divides a cache of blocks*positions into blocks.
func (*BlockPool) Free ¶
Free returns a sequence's blocks to the pool.
Freeing a page twice is refused rather than ignored: it would hand one block to two sequences, and the symptom is one sequence reading another's tokens — which reads as a model producing text from the wrong conversation.
type Buckets ¶
type Buckets struct {
// contains filtered or unexported fields
}
Buckets is a sorted set of sequence lengths a caller compiles plans for.
A prompt runs in the smallest bucket that fits, and the extra positions are padding. Padding needs no mask: a causal prefill's query position attends to at most its own position, and padding sits after the real tokens, so a real position's window never reaches it. What the padded rows compute is discarded.
The trade is arithmetic for compilation. Fewer buckets waste more work per request; more buckets hold more plans and more memory. That is a policy, which is why it is a caller's list rather than a rule here. The sizes are unexported so a literal cannot bypass NewBuckets. A `Buckets{512, 128, 256}` written by hand is unsorted, and Buckets.For searches — so a size-100 request returned 512, the first element, rather than 128. Silently, with a plausible answer.
func NewBuckets ¶
NewBuckets sorts a bucket set and refuses a duplicate. It is the only way to build one; see Buckets.
Refuses rather than collapses, and the distinction is worth the sentence: a repeated size is a caller's list saying something they did not mean, and silently collapsing it would compile one plan where they asked for two and report nothing. Every other malformed set here is refused for the same reason, so de-duplicating would be the one place a mistake was tidied away.
func (Buckets) For ¶
For returns the smallest bucket that holds n tokens.
A prompt longer than the largest bucket is an error rather than a truncation. Truncating changes what the model was asked and produces a plausible answer to a different question, which is the worst failure a serving path can have.
type Builder ¶
type Builder struct {
// contains filtered or unexported fields
}
Builder records one tensor graph.
It collects errors rather than returning them, so model code has no error branch per line. See the package doc.
func (*Builder) Compile ¶
func (b *Builder) Compile(rt *Runtime, opts CompileOptions) (*Plan, error)
Compile turns the recorded graph into a plan.
Every collected error is returned together rather than the first one, because a model with three mistakes should take one compile to find all three. Each names the operator, the operand, and the line that recorded it.
func (*Builder) Err ¶
Err reports every diagnostic collected so far, joined.
Compile returns the same thing; this exists so a caller building a model in pieces can check without compiling.
func (*Builder) Identity ¶
Identity is a digest over everything about a recorded graph that changes what compiling it produces.
It exists so a PlanCache can tell one model from another, and specs/007-tensor-layer.md is explicit that a shape is not enough: two different models over the same shapes are different plans, and returning one for the other produces a confident wrong answer.
What it covers: every operator in record order, the kernel each selected and that kernel's own digest, every operand's shape, dtype and layout, and every port and scalar's name and kind. What it deliberately does not cover: the *values* a caller binds, which are what a plan is reused across.
type CompileOptions ¶
type CompileOptions struct {
Label string
}
CompileOptions carries what a plan needs beyond the graph.
type Int4 ¶
type Int4 struct {
// Codes is u32, eight weights per element, low nibble first. See
// [quant.Int4Quantize] — a caller does not pack this by hand.
Codes *Tensor
// Scales and Zeros are f16, one each per [quant.Int4Group] weights.
Scales *Tensor
Zeros *Tensor
// Weights is how many weights Codes holds, which the element count cannot
// give: eight per word means a matrix of 8n and one of 8n-3 pack into the
// same number of words. specs/048-int4.md §4.
Weights int
}
Int4 is a weight matrix stored as packed 4-bit codes with a scale and a zero point per group.
Three planes rather than two, and the third is what makes it four bits. Quantized's int8 is symmetric, so a scale is enough; at four bits the codes have to be spent where the weights actually are, which needs an offset as well. specs/048-int4.md §1.
Bundled for Quantized's reason: binding one matrix's codes against another matrix's scales compiles, runs, and produces a matrix of noise.
type KernelSelection ¶
KernelSelection reports which kernel an operator became, and why.
Reported rather than merely decided, because specs/007-tensor-layer.md makes fused attention a *selection* rather than a capability: a caller who cannot see which they got cannot explain a performance cliff or a numeric difference, and would have to guess from timings.
type LinearOptions ¶
type LinearOptions struct {
// Alpha and Beta are f32 tensors with one entry per token, in the flat
// order q is in. Alpha decays the state and Beta writes into it: the
// recurrence is S <- alpha*S + beta*k*(v - S k)^T.
Alpha *Tensor
Beta *Tensor
// QueryExtents is a u32 tensor holding, per sequence, how many tokens that
// sequence contributes to this step. The same segmented extent
// [AttentionOptions.QueryExtents] takes, and for the same reason: a decode
// step is one token per sequence, a prefill is many, and a mixed step is
// both at once.
//
// A count of zero is legal. specs/046-segmented-extents.md §1.
QueryExtents *Tensor
}
LinearOptions configures a gated delta layer.
specs/047-linear-attention.md. Every field is per token or per sequence and none is a scalar, which is specs/043-per-row-values.md §2 applied without argument: a gate differs per token and an extent differs per sequence.
type Plan ¶
type Plan struct {
// contains filtered or unexported fields
}
It owns one device graph and the transient memory the planner chose; the caller owns every buffer they named. Immutable, and with the graph's one-submission-in-flight restriction.
func (*Plan) Memory ¶
func (p *Plan) Memory() accel.GraphMemory
Memory reports what the plan's intermediates cost after the graph planner aliased them.
Straight from the device graph rather than counted here, which is the point: the intermediates went through the recorder, so they got specs/017's aliasing and this number is evidence of it rather than a second accounting.
func (*Plan) Ports ¶
Ports reports every external buffer this plan expects, in declaration order.
A caller binds by name, so this is how they discover what to bind without having kept the builder. Sorted by declaration rather than by name, because declaration order is the order a reader of the model saw them.
func (*Plan) Scalars ¶
func (p *Plan) Scalars() []ScalarDesc
Scalars reports every named per-step value this plan reads.
func (*Plan) Selections ¶
func (p *Plan) Selections() []KernelSelection
Selections reports which kernel each operator became, and why.
func (*Plan) Submit ¶
Submit binds and runs the plan, returning without waiting.
The whole binding set is validated before anything is bound, so a submission with one bad name leaves the plan exactly as it was. A failure comes back through an already-signalled fence rather than as a second return value, which matches specs/003-command-graph.md and means a caller checks one thing.
One submission at a time ¶
specs/007-tensor-layer.md gives a plan "the Graph's one-submission-in-flight restriction", and this is where it has to be enforced rather than inherited. Binding happens here and synchronously; the submission is handed to the queue's serial stream and runs later. So a second Submit before the first has run would rebind the slots underneath it -- and the graph's own in-flight check does not catch that, because the graph is not marked in flight until its worker reaches it.
That was not theoretical. Two submissions with different inputs and different outputs, back to back, produced *one* result: the first submission wrote into the second's output buffer, and both fences reported success. A silently lost result is the worst failure available here, so this refuses instead.
type PlanCache ¶
type PlanCache struct {
// contains filtered or unexported fields
}
PlanCache keeps compiled plans so a caller can vary shape without paying for compilation on every request.
cache := tensor.NewPlanCache(rt)
defer cache.Close()
plan, err := cache.Compile(func(b *tensor.Builder) { buildModel(b, n) },
tensor.CompileOptions{Label: "prefill"})
The second call with the same model and the same shapes returns the first call's plan.
What it does not do ¶
It never evicts. A plan owns transient device memory, and a cache that freed one on its own would free memory a caller might be about to submit against — so it grows until PlanCache.Close. A caller who needs a bound gives it a bounded set of shapes, which is what Buckets is for.
It is also not a substitute for holding a plan in a variable. A decode loop runs one shape and should keep its plan; this is for the case where the shape varies, and a lookup still costs a digest and a map probe.
specs/029-plan-cache.md has the reasoning, and specs/007-tensor-layer.md has the requirement its key satisfies.
func NewPlanCache ¶
NewPlanCache returns a cache over one runtime.
func (*PlanCache) Compile ¶
func (c *PlanCache) Compile(record func(*Builder), opts CompileOptions) (*Plan, error)
Compile returns the plan for a recorded graph, compiling it once.
The graph is recorded by the callback rather than passed in, because the point of a hit is not to record it: a caller who built the graph to ask for its key would have paid most of what the cache saves.
type PortKind ¶
type PortKind uint8
PortKind says how a caller supplies an external value.
const ( // PortInput varies on every submission. PortInput PortKind = iota // PortWeight may be rebound between submissions and is read-only. PortWeight // PortState is caller-owned read-write storage, never aliased by the // planner. PortState // PortOutput is a value the caller wants written somewhere they can read. PortOutput )
type Quantized ¶
type Quantized struct {
// Quants is i8, one per weight, shaped like the matrix.
Quants *Tensor
// Scales is f16, one per [quant.Int8Block] weights of the flattened
// matrix.
Scales *Tensor
}
Quantized is a weight stored as quants and scales.
Two tensors rather than one, because that is what the device holds: specs/001-device-resources.md types a buffer by dtype and an interleaved block struct has no dtype, so a quantized matrix is two planes. Bundling them in one value keeps a caller from binding a matrix's quants against another matrix's scales, which would compile, run, and produce a matrix of noise.
type Runtime ¶
type Runtime struct {
// contains filtered or unexported fields
}
Runtime owns one device and the pipelines compiled for it.
Pipelines are cached here rather than on a plan because two plans over the same model share nearly all of them, and compiling MSL is a call into the device compiler that takes milliseconds. There is deliberately no plan cache: specs/007-tensor-layer.md puts that above this package, because a key that looked only at shapes would be wrong in ways nobody could see.
func NewRuntime ¶
NewRuntime prepares a device for tensor work.
func (*Runtime) Close ¶
Close releases the runtime's pipelines.
Every plan built from it must be closed first, which is checked rather than assumed: a plan outliving its runtime would hold a pipeline nobody owns.
func (*Runtime) NewBuilder ¶
NewBuilder starts recording one tensor graph.
A builder belongs to one goroutine. It records rather than executes, so nothing here touches the device until Compile.
type SamplingOptions ¶
type SamplingOptions struct {
// Temperature of 0 is greedy, and is a different graph rather than a small
// number: see [SamplingOptions.Validate].
Temperature float32
// TopK and TopP are off at 0. "Off" removes the node from the graph, which
// is the only way to turn truncation off -- a TopK equal to the vocabulary
// is not off, because the mask's round count is bounded and it would
// silently keep [TopMaxRounds].
TopK int
TopP float32
// Repetition divides a positive logit and multiplies a non-positive one.
// Both 0 and 1 are off.
Repetition float32
// Presence is subtracted once from any token that occurred, Frequency once
// per occurrence. 0 is off for both.
Presence float32
Frequency float32
}
SamplingOptions is one sequence's policy.
The zero value is greedy: argmax, no penalties, no truncation. That is the policy a caller gets by passing nothing, and it is the one with no way to surprise them.
There is no generator in here ¶
specs/039-sampling-policy.md section 2. Copying this struct copies numbers. The draw comes from a Stream the caller holds per sequence and reaches the graph as a tensor, because it is a per-row value and specs/043-per-row-values.md makes those tensors rather than scalars.
func (SamplingOptions) Greedy ¶
func (o SamplingOptions) Greedy() bool
Greedy reports whether this policy takes the argmax rather than a draw.
func (SamplingOptions) Penalised ¶
func (o SamplingOptions) Penalised() bool
Penalised reports whether any penalty is configured.
Exported because it is the question a caller has to answer before deciding whether to declare the history and counts states at all, and re-deriving it from three fields at the call site is how the two answers drift apart.
func (SamplingOptions) Scalars ¶
func (o SamplingOptions) Scalars(prefix string, n, historyCap uint32) (map[string]ScalarValue, error)
Scalars is what one decode step binds.
Every value here is a number a step changes without changing the plan: specs/039-sampling-policy.md section 7 prices them at nothing, because Submit rewrites every uniform on every submission anyway. What is *structural* -- which nodes exist, and the k and p the masks were recorded with -- is not here, and changing one of those is a different plan.
n is how much of the history ring is filled: min(tokens so far, capacity). It is refused above the capacity rather than clamped, because a clamp turns a caller's off-by-one into a penalty over a window they did not ask for.
func (SamplingOptions) Validate ¶
func (o SamplingOptions) Validate() error
Validate refuses a policy rather than clamping it.
This is the first layer that can report an error at all: a kernel cannot, which is why TopMaxRounds truncates down there and the draw clamps. So the refusals belong here, and they refuse rather than repair, because a policy quietly changed into a different one produces plausible tokens and no evidence.
type ScalarDesc ¶
type ScalarDesc struct {
Name string
Kind ScalarKind
}
ScalarDesc declares a named runtime value.
type ScalarKind ¶
type ScalarKind uint8
ScalarKind is the type of a named per-step value.
const ( // ScalarU32 is an unsigned integer, such as a sequence length. ScalarU32 ScalarKind = iota // ScalarF32 is a float, such as a softmax scale or a RoPE base. ScalarF32 )
func (ScalarKind) String ¶
func (k ScalarKind) String() string
type ScalarValue ¶
type ScalarValue struct {
Kind ScalarKind
U32 uint32
F32 float32
}
ScalarValue is one named value, supplied at submission.
func (ScalarValue) String ¶
func (v ScalarValue) String() string
type Shape ¶
type Shape []int
Shape is a tensor's extent, outermost dimension first.
The last dimension varies fastest, which is row-major and is what every kernel in the corpus indexes. All dimensions are positive concrete integers: there is no symbolic shape in v0, so inference is arithmetic rather than a solver.
type SoftmaxOptions ¶
type SoftmaxOptions struct {
// Axis is normalized; it must be the last, which is the only axis the
// registered kernel reduces over.
Axis int
}
SoftmaxOptions carries the attributes softmax takes.
type State ¶
type State struct {
// contains filtered or unexported fields
}
State is a version of caller-owned mutable storage.
A value rather than a handle: ScatterRows returns the next version, and an operator reading an *earlier* version is reading what was there before the write. Holding on to an old version is therefore meaningful rather than a mistake, and the DAG records which one each reader meant.
func LayerState ¶
LayerState is a compile-time view of one layer's slice of a state.
The version chain and the binding identity are the parent's, which is what makes a per-layer cache one buffer rather than N: a model with thirty-two layers binds one KV tensor and each layer addresses its own slice.
func NewState ¶
Persistent declares caller-owned mutable storage.
Never transient and never aliased by the planner: the caller owns the buffer and its contents outlive the submission, which is the whole point of a cache.
func ScatterRows ¶
ScatterRows writes rows into a state at runtime indices and returns the next version.
The next version rather than nothing, so a reader downstream names *which* state it meant. An operator holding the version before this write reads what was there, and the graph orders the two because their byte ranges overlap.
type StateDesc ¶
type StateDesc struct {
Name string
DType DType
// Shape is the whole extent, including the sequence capacity. A KV cache is
// [capacity, heads, headDim].
Shape Shape
}
StateDesc declares caller-owned mutable storage.
type Stream ¶
type Stream struct{ Seed uint64 }
Stream is one sequence's source of draws.
It is a value, and that is the design rather than an accident. The shape everyone writes first is a policy struct holding a *rand.Rand, and Go copies a struct silently -- on assignment, by value, into a map -- so two sequences end up sharing one generator, interleaving their draws, and neither reproduces. A *rand.Rand is also not safe for concurrent use, so two goroutines holding copies of that struct are a data race the detector finds only if a test happens to run them together.
Copying a Stream copies a number. There is nothing to share.
func Derive ¶
Derive gives sequence seq of a batch its own stream from one root seed.
seq+1 rather than seq so that sequence 0 is not the root seed itself, which would make a one-sequence batch and an unbatched run of the same seed draw the same numbers -- true today and silently false the moment the derivation changes.
func (Stream) Draw ¶
Draw returns the uniform in [0,1) for token step of this sequence.
The step is the token index, not a draw counter ¶
The caller already holds this number: it is the position the KV cache writes at. So the sampler stores nothing per sequence, resuming a sequence at token N costs nothing, and exactly one draw is defined per token whether or not that step used it. Turning temperature off for one step does not shift every later token, which it would if the stream had a position to advance.
Why the result can never be 1.0 ¶
Twenty-four bits of the finalized state, divided by 2^24. Both the numerator and the divisor are exact in f32, the largest numerator is 2^24-1, and so the largest result is exactly 0.99999994.
The obvious spelling, float32(rng.Float64()), rounds up to exactly 1.0 for about one input in 2^24. specs/028-sampling.md's walk clamps a draw of 1.0 down, so nothing crashes: the last token in the vocabulary silently receives the extra mass, on one step in sixteen million, and every differential still passes because both backends clamp identically. This division makes the generator and that backstop agree by construction rather than by luck.
type Tensor ¶
type Tensor struct {
// contains filtered or unexported fields
}
Tensor is an immutable logical value in one builder's graph.
Immutable because the graph is a DAG of values rather than a sequence of mutations: an operator returns a new tensor and never writes into an operand, which is what lets the planner decide where intermediates live and which of them can share memory.
func Argmax ¶
Argmax writes the index of the largest logit in each row.
Greedy decoding, and the operator specs/039-sampling-policy.md section 3 puts behind a temperature of zero: T = 0 selects this rather than a softmax at a tiny T, because with a logit gap smaller than T the softmax is not one-hot at all and the walk samples the runner-up.
Ties go to the lowest index ¶
Equal logits are ordinary -- an untrained model produces them everywhere and a trained one produces them at saturation. specs/028-sampling.md section 3 states the rule because the alternative is not "some index" but a *different* index on each backend: a tree reduction's answer depends on which lane compared which pair, and two backends reducing at different widths would disagree about a token.
The result is u32, which needs nothing new: GatherRows and ScatterRows already require u32 ids, so a sampled token feeds the next step's embedding lookup directly.
func Attention ¶
func Attention(b *Builder, q *Tensor, k, v *State, opts AttentionOptions) *Tensor
Attention scores a query against a cached key/value pair.
Why the cache is State and the query is a Tensor ¶
The query is this step's; the cache is every step's. That asymmetry is the whole shape of decoding, and expressing it in the types means a caller cannot accidentally write the query or read a stale cache: a State carries a version and a Tensor does not.
Fusion is not a selection, because there is nothing to select between ¶
specs/007-tensor-layer.md said fused attention is "runtime kernel selection, not a device capability", with the composed definition -- score MatMul, Softmax, value MatMul -- as both the correctness reference and the fallback. The reference half holds and the fallback half does not: several query heads share one key/value head, so the composed form needs a matrix multiply per head, and specs/025-tensor-operators.md multiplies two matrices with no leading axes broadcast. The composition exists only at kvHeads == 1, which no model this serves uses. 007 is corrected; the fused kernels are the only path, and specs/044-unbounded-context.md is why they can be.
The composed reference still runs, in the corpus tests, over the shapes it can express. That is what keeps it a reference.
func Broadcast ¶
Broadcast expands size-one axes to a larger shape.
The expansion is a **zero stride**, which is the whole trick: every index along that axis reads the same element, so nothing is materialized and nothing is copied. A kernel that indexes contiguously cannot read it, which is why lowering refuses a broadcast operand it cannot express as a repeated contiguous run.
func Cast ¶
Cast converts between storage formats.
An operator rather than an implicit rule at every boundary, and specs/007-tensor-layer.md's reason is the one that matters: a conversion costs a pass over the data and changes the numbers, so it is something a caller writes rather than something that happens to them. `Add` refusing two dtypes and `Cast` existing are the same decision seen from two sides.
f16 to f32 is exact; f32 to f16 rounds to nearest-even and a value outside f16's range becomes an infinity rather than a saturated maximum, because a silently clamped weight is a plausible weight.
bf16 widens and does not narrow ¶
A checkpoint ships bf16 -- Qwen3 does -- and bf16 to f32 is exact: bf16 is f32's top half, the same eight-bit exponent with sixteen zero bits below, so the conversion is a shift. Going to f16 instead is the one lossy step in this pipeline, because f16 carries a five-bit exponent where bf16 carries f32's and a bf16 value can be outside f16's range entirely. Only the widening is registered, so a caller who wants f16 writes both casts and sees where the error enters -- which is the reason this operator exists at all rather than happening implicitly at a boundary.
func Contiguous ¶
Contiguous packs a strided view into fresh contiguous storage.
Why a caller reaches for this ¶
A view is free. Permute, Transpose, Slice and Broadcast change strides and copy nothing, which is what makes a head split cost nothing. But a kernel that indexes contiguously cannot read one, so a transposed operand reaches a matmul as a refusal — and until this existed there was no way to convert it. specs/025-tensor-operators.md named this operator in four error messages, one of which told a caller to insert it.
Why it is explicit and not automatic ¶
Inserting the copy silently is the choice specs/007-tensor-layer.md declines: *"a copy nobody asked for is a cost nobody can see."* A transformer's reshaping is free precisely because nothing materializes behind the caller's back, and an operator that quietly packed would make the free case and the expensive case look identical in the source.
So a strided operand is still refused, and the refusal names this. Calling it is how a caller says the copy is worth it.
What it costs ¶
One element read and one written, per element, plus the index arithmetic: each destination element decomposes its linear index into coordinates and dots them with the source strides. An operand that is already contiguous is returned unchanged rather than copied, so calling this defensively is free.
func GatherRows ¶
Rows gathers whole rows of a table by index.
This is an embedding lookup, and the out-of-range rule is the interesting part: specs/007-tensor-layer.md makes an out-of-range id a caller error, and the kernel writes zeros rather than reading outside the table. Zeros are a plausible embedding, so this is the one place where a diagnostic would be worth more than a safe answer -- and where the corpus kernel gives the safe answer because a GPU has no other option.
func GroupedMatMul ¶
GroupedMatMul multiplies each expert's tokens by that expert's matrix.
out[t][n] = Σₖ x[t][k] · w[e(t)][k][n]
Why a separate operator from [GroupedMatVec] ¶
The shapes a decode step and a prefill have. A decode routes one token to a few experts and reads those matrices once, so a workgroup per (token, column) is right. A prefill has many tokens per expert, and reading an expert's matrix once per token wastes the bandwidth a mixture-of-experts layer is built to save — so this puts a workgroup on an (expert, column tile) and walks that expert's whole segment through shared tiles.
Each weight is read once per block of testkernels.TileM tokens rather than once per token. specs/049-grouped-gemm.md §5.
It takes the same inputs ¶
x is [Σ counts, K] ordered by expert, w is [E, K, N], counts is one per expert — GroupedMatVec's arguments exactly, so switching between the two is a one-word edit and not a re-plumbing.
func GroupedMatVec ¶
GroupedMatVec multiplies each token by the weight matrix its segment names.
out[t][n] = Σₖ x[t][k] · w[e(t)][k][n]
What it is for ¶
A mixture-of-experts layer: E weight matrices and a router that sends each token to a few of them. The naive form — run every expert and mask — is expressible with MatMul today and does E/k times the work, which inverts the reason such a layer exists.
specs/049-grouped-gemm.md. The shape is [046](046-segmented-extents.md)'s segmented extent with the row being an expert rather than a sequence, so this operator adds no concept: counts is one count per expert, and the offsets are derived here as they are there.
Tokens arrive ordered by expert ¶
x is [Σ counts, K], the tokens of expert 0 then expert 1 and so on. Producing that order from a routing table is a sort of a few thousand small integers, which is cheaper on the host than the dispatch it precedes — the same argument specs/046-segmented-extents.md §1.1 makes for deriving the offsets.
An expert with no tokens is ordinary ¶
Its count is zero and it contributes nothing. That is not an edge case here: with top-2-of-8 routing, six experts get nothing on every single token.
Rows of x past the total are padding ¶
The counts may sum to fewer tokens than x holds. Those rows routed nowhere: they read no weights and their output is zero. The sum is device data, so this is enforced in the kernel rather than refused here -- specs/046-segmented-extents.md §1 property 3.
func Int4MatMul ¶
Int4MatMul multiplies a batch of f32 activations by a packed 4-bit matrix.
out[m][n] = Σₖ a[m][k] · ((code(k·N+n) − z[g]) · s[g])
Why a separate operator from [Int4MatVec] ¶
The shapes a decode step and a prefill have, and they want different kernels for the reason specs/048-int4.md §5 gives. A decode reads the whole model to produce one token, so its matvec is bound by how fast the weights arrive. A prefill has many tokens against the same matrix, so the matrix is read once per tile of tokens and the unpacking is amortised over the tile rather than repeated per token.
Taking a vector here would work and would be slower than Int4MatVec, and taking a matrix there would not compile. A caller should see which one they wrote, which is QuantMatMul's argument one width down.
The accuracy is Int4MatVec's: the same representation, the same reconstruction, and specs/048-int4.md §3's bound covers both.
func Int4MatVec ¶
Int4MatVec multiplies f32 activations by a packed 4-bit weight matrix.
out[n] = Σₖ a[k] · ((code(k·N+n) − z[g]) · s[g])
Why a separate operator ¶
QuantMatMul's reason, one width down: the cost is different and a caller should see which one they wrote. What differs here is also the *accuracy*, and not in one direction — specs/048-int4.md §3 states the bound as a group's range over 30 where int8's is a peak over 254, so a matrix whose weights cluster away from zero is represented better by these four bits than by eight symmetric ones, and one centred on zero is about seventeen times worse.
A vector, not a matrix ¶
This is the shape a decode step has, which is where the memory pressure four-bit weights exist to relieve is felt: a decode reads the whole model to produce one token. A prefill wants a tiled form over the same representation and specs/048-int4.md §5 records it as not built.
func Linear ¶
Linear multiplies and adds a bias in one kernel.
An authored epilogue rather than MatMul followed by Add, which is a selection specs/010-kernel-corpus.md registers: the composed form is correct and writes the whole result twice. The composed form remains the reference.
func MatMul ¶
MatMul multiplies two matrices, accumulating in f32.
The result is f32 whatever the operands are, because that is what accumulating in f32 means: rounding the sum back to f16 at the end would throw away the precision the accumulation was for, and a caller who wants it narrow can convert.
func Permute ¶
Permute reorders axes.
The strides move with them, so this is bookkeeping: a permuted tensor reads the same bytes in a different order, and only becomes a copy if something downstream needs it contiguous.
func QuantGatherRows ¶
QuantRows gathers rows of a quantized table.
An embedding table is the largest single tensor in a small model and the one quantization helps most: every token reads one row of it and nothing else, so the whole table sits in memory to serve one row at a time.
func QuantMatMul ¶
QuantMatMul multiplies f16 or f32 activations by a quantized weight matrix.
A separate operator rather than MatMul taking a different argument type, because the cost is different and a caller should see which one they wrote. specs/027-quantization.md states the error budget, and it is not the unquantized one: the result is within a derived distance of what the unquantized product would have given, and that distance is proportional to the largest weight in each block.
Why the activation may be either width ¶
Because the weight's width is not the activation's business. int8 is what a model reaches for when it is too large to hold otherwise, and every other operator in this package produces f32 -- so requiring the activation to be f16 put a Cast in front of every projection of the configuration least able to afford one (accel issue 14). The quants and scales keep their widths in both cases: a weight is loaded from a file and an activation is produced by the graph, and only the second is free to be wide.
func RMSNorm ¶
RMSNorm normalizes each row by its root mean square and scales by a gain.
The mean square and the reciprocal square root are computed in f32 whatever the storage, which specs/007-tensor-layer.md requires: a normalization accumulated in f16 loses the thing it is measuring.
func ReadState ¶
ReadState returns the tensor a state version holds.
Why an older version is refused ¶
specs/007-tensor-layer.md makes a write return the *next version*, so holding an older one is meaningful. Expressing that on the device would mean copying the previous contents aside before the write, because both versions live in one caller-owned buffer -- and v0 does not do that.
The alternative was to let an old version read the new contents, which is what the first implementation did: the version chain compiled, ordered nothing, and a test that deliberately read the stale version still passed. A distinction that cannot be violated is not a distinction, so this refuses.
func Reshape ¶
Reshape reinterprets a tensor's extent without moving anything.
Legal only on a contiguous operand, which is not a limitation of this implementation but of what reshaping means: a strided view's elements are not adjacent, so a different extent over them describes different elements. The error says so rather than saying "unsupported", because the fix is reshaping before the transpose or slice that made it strided, and a reader should not have to guess that.
func RoPE ¶
RoPE applies rotary position embedding in place on a copy.
In place on a *copy*, because the registered kernel rotates its buffer and a tensor is an immutable value. So this lowers to a copy into a transient followed by a dispatch over it, and the copy is reported: it is real work and hiding it would make a rotation look free.
Why positions are a tensor and the base is a scalar ¶
specs/043-per-row-values.md draws the line: a value every row of a dispatch shares is a uniform, and a value that differs per row is device data. The frequency base is a property of the model; the position is a property of the sequence.
This took a scalar offset and computed `row + offset`, which is exactly right for one sequence and exactly wrong for two. In a batched decode the row index is the *slot*, so slot 0 rotates at the offset and slot 1 at offset+1 — meaning one member of the batch is ever rotated at its own cache length. The output stays finite and fluent; long-range coherence degrades in a way that reads as "the model is a bit weak" rather than as a bug.
A single sequence binds a one-row positions tensor. That is the same path, not a special case.
func Sample ¶
func Sample(b *Builder, logits, draws *Tensor, history, counts *State, o SamplingOptions, prefix string) *Tensor
Sample records one sequence's sampling policy and returns the chosen token.
The order, and why not another ¶
logits ─▶ penalties ─▶ ×1/T ─▶ softmax ─▶ top-k ─▶ top-p ─▶ draw ─▶ token
(optional) (optional) (optional)
└─▶ argmax ─▶ token when Temperature is 0
Penalties act before temperature because subtracting a after scaling by 1/T is subtracting a·T before it: two knobs a caller turns independently must not multiply. Truncation acts after the softmax because top-p is a mass threshold, and top-k joins it there so both masks act on the values the sampler will actually walk -- f32 rounding can make two distinct logits equal probabilities, and a top-k over logits would then keep a different boundary entry than the walk sees. Top-k precedes top-p because top-p is relative to its own input's total, so each bound is one the other cannot violate.
Nothing renormalizes and there is never a second softmax. A mask leaves the weights summing below one, which invites a fix; specs/028-sampling.md deleted that fix and made the walk compare against draw × total instead. A softmax over a mask's output is near-uniform over the whole vocabulary, because exp(0) = 1 for every dropped entry.
What the caller owns ¶
draws is one uniform per row, from a Stream. history and counts are caller-owned state, required exactly when a penalty is configured and refused when one is not, so a caller cannot bind storage nothing reads.
counts is [vocab]u32 and is rebuilt from the history every step rather than carried: this returns the version after that rebuild, so a caller holding the old one is reading what was there before, which specs/007-tensor-layer.md makes meaningful rather than a mistake.
func SampleCategorical ¶
SampleCategorical draws one index per row from the weights that row holds.
out[r] = min{ i : Σⱼ≤ᵢ weights[r][j] > draws[r] × Σⱼ weights[r][j] }
The draws are a tensor ¶
One per row, and specs/043-per-row-values.md section 2 is the reason: a value that differs per row is device data. A batch sharing one draw stays perfectly reproducible and stops being independent, which is a failure no reproducibility test can see.
Each draw is a uniform in [0, 1). A draw outside it is clamped rather than refused, because a kernel cannot report an error and an unclamped draw reads past the end of the row; specs/028-sampling.md section 4 states that and the clamp is unconditional, so correctness does not depend on a build mode.
The weights need not be normalized ¶
The walk compares against the draw times the row's own total, so this is correct for any vector of non-negative weights and not only for a distribution. That is what lets TopKMask and TopPMask compose in front of it without a renormalizing pass -- and specs/039-sampling-policy.md section 5 forbids putting a second Softmax there, since exp(0) = 1 for every dropped entry would make the mask do nothing.
A zero-weight entry can never be drawn, because its partial sum does not increase and the comparison is strict. That is what makes a mask a mask.
The walk is in index order ¶
A parallel prefix sum would be faster and would put the boundary somewhere else when two weights are equal. specs/028-sampling.md section 3 takes the reproducible answer over the fast one, which is the trade specs/008-numerics.md makes when it forbids a tolerance.
func Scale ¶
Scale multiplies x by a named runtime f32 scalar.
A named scalar rather than a Go float32, because the value changes every step and a compiled plan must not have to be rebuilt for it. An attribute that changed a shape, a layout, or which kernel is selected would be different: that needs another plan, and specs/007-tensor-layer.md draws the line there.
func SiLU ¶
SiLU returns x * sigmoid(x), elementwise.
Evaluated in f32 whatever the storage dtype, which specs/007-tensor-layer.md requires and specs/008-numerics.md explains: an activation evaluated in f16 loses accuracy where it matters most, near zero.
func Slice ¶
Slice narrows one axis to a half-open range.
Unit step only, which keeps a slice a view: a strided step would still be a view, and nothing in v0 needs one, so it is absent rather than untested.
func Softmax ¶
func Softmax(b *Builder, x *Tensor, opts SoftmaxOptions) *Tensor
Softmax normalizes each row into a distribution.
The maximum subtraction, the exponentiation, the sum and the division all happen in f32. Subtracting the maximum first is not an optimization: without it, exp overflows for any row whose largest value exceeds about 88.
func SwiGLU ¶
SwiGLU returns SiLU(gate) * value, elementwise.
One authored kernel rather than two operators composed, which is a selection and not a fusion pass: specs/010-kernel-corpus.md registers it, and the composed definition remains the correctness reference. The shapes must be equal rather than merely broadcastable, because the fused kernel indexes all three operands together.
func TopKMask ¶
TopKMask keeps the k largest entries of each row and zeroes the rest.
A mask over weights rather than a step fused into the draw: kept entries carry their input value and dropped ones carry zero, so the result feeds SampleCategorical directly. specs/039-sampling-policy.md section 5 places it after Softmax and before TopPMask, and that order is not re-derived here: top-p is a mass threshold so it must see probabilities, and top-p is relative to its input's own total so it composes after a k that has already cut.
Exactly k, including at a tie ¶
The comparison is lexicographic on (value, index) descending, so an entry ties only with itself and "the k largest" is exactly k entries whatever the data. A threshold search would be fewer rounds and would keep however many happened to sit above wherever the bisection stopped -- which is k only when nothing ties, and ties near the tail of a distribution are the normal case.
k is refused rather than clamped ¶
The kernel clamps its round count to TopMaxRounds because a kernel cannot report an error. This is the first layer that can, so a k above the bound is refused here: silently running top-128 when a caller asked for top-vocab changes what a model samples without changing what it reports, which specs/039-sampling-policy.md section 5 calls out by name.
func TopPMask ¶
TopPMask keeps the smallest set of largest entries whose mass reaches p.
The nucleus. The same walk TopKMask performs with a different stopping rule: it accumulates weight rather than counting entries, and the entry that crosses the threshold is kept -- which is what makes the set the *smallest* one reaching p rather than the largest one below it.
The mass is a fraction of the row's own total rather than of one, so this composes after a TopKMask and on unnormalized weights, for the same reason SampleCategorical scales its draw by the total.
p is refused rather than clamped, and zero is not "off" ¶
A p of zero makes the threshold zero, so the walk never advances its frontier, the mask keeps nothing, and SampleCategorical over an all-zero row finds no partial sum above its target and returns the last index. That is a plausible token id from a row that was entirely erased. The way to say "no truncation" is to leave the operator out of the graph, which specs/039-sampling-policy.md section 5 states as "off means the node is absent".
It costs the same whatever p is ¶
All TopMaxRounds rounds run. Stopping the loop early would put a barrier in non-uniform control flow, which specs/002-compute-model.md section 3.1 and specs/018-cooperative-lowering.md forbid, so only the frontier's advance is conditional. p = 0.5 costs what p = 0.99 costs, and a top-k followed by a top-p costs k + 128 workgroup reductions. Enabling this by default on a false cost model is a design mistake rather than only a slow step.
Source Files