llm

package module
v0.3.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Nov 4, 2026 License: BSD-3-Clause Imports: 17 Imported by: 0

README

llm

A small-LLM inference library in Go for WASI containers: a wasip1 guest reads a GGUF model from a read-only mount (/models/NAME.gguf) and generates text on the GPU through the batched command stream of go-wasmcontainers/gpu (wasi_webgpu_batch), in the browser first. Design and plan: docs/llm-design.md in go-wasmcontainers/gpu.

Package What
gguf bounded GGUF v3 header parser over an io.ReaderAt; fuzzed, corpus committed
llm (root) byte-level BPE tokenizer (gpt2 model, smollm pre-tokenizer), llama config, CPU forward pass (the oracle and native fallback), GPU forward pass and its WGSL f32 kernels
examples/generate the container: run, check (GPU against the CPU path), bench

First model: SmolLM2-135M-Instruct Q4_0 (Q4_0, Q4_1 and Q8_0 tensors).

Design points

  • Streaming load. Weights go from the file to the GPU in 16 MiB pieces, repacked per block (18-byte Q4_0 blocks become aligned u32 quants and f32 scales); the model never sits in guest memory.
  • Eight tokens per submit. Every dispatch of the forward pass is recorded into compute passes; sampling (argmax) runs on the GPU and only the sampled tokens are read back, once per 8. Prompt tokens go through the layers in chunks of up to 64.
  • Vendor-neutral f32. No subgroups, no shader-f16. The workgroup size (64, lower if the adapter's invocation or workgroup-memory limit is lower), the 2-D grids and the buffer sizes come from the adapter limits the host reports. Reductions are tree reductions in workgroup memory: the order of summation depends on the workgroup size, not on the hardware. RoPE angles are computed on the CPU in double precision, so no adapter's sin/cos accuracy enters. Only Apple Metal has been run; no NVIDIA or AMD result is claimed.

Tests

go test ./...                                  # parser, refusals, fuzz corpus, tokenizer-free units
LLM_MODELS=DIR go test -run Oracle -v .        # needs llama-server (the go-pkgx llama.cpp bottle) and
                                               # DIR/SmolLM2-135M-Instruct-Q4_0.gguf (pinned in go-wasmcontainers-llm-corpus)

Model files are never committed; tests that need them skip without LLM_MODELS. The GPU oracle runs in the browser from go-wasmcontainers/gpu (LLM_REPO=... go test -tags e2e ./e2e -run TestLLMOracle).

Oracle (vendor-sensitive tests are marked)
  • Tokenizer: ids identical to llama.cpp's /tokenize on the fixed prompts.
  • Strict oracle: llama.cpp on the same weights dequantised to F32. llama-quantize --allow-requantize SmolLM2-135M-Instruct-Q4_0.gguf OUT.gguf F32 (llama.cpp's own tool, so the weights are identical) run with --cache-type-k f32 --cache-type-v f32 -fa off: f32 activations, f32 cache. TestOracleDequantizedF32: the Go CPU f32 path gives 32 of 32 greedy tokens on all four prompts and step-1 log-probability differences of at most 1.4e-5 nats (tolerance 1e-4). The GPU container is held to the same tokens (vendor-sensitive: it differs from the CPU path by 1.6e-6 relative in logits on Apple Metal; another vendor's fused multiply-adds and exp move the last bits, and a token flips only at a gap that small).
  • Why llama.cpp's Q4_0 CPU run differs (about 0.1 to 0.4 nats at step 1). On CPU ggml quantises the activations to Q8_0 (Q8_1 for Q4_1 weights) before every quantised matrix product (vec_dot_q4_0_q8_0), our f32 path does not. The Go CPU path can emulate it (ExactGGML: Q8_0/Q8_1 activations, f16 KV cache) and that is not enough to match to 1e-3: such a path amplifies rounding by about 10^6 (TestQuantisedActivationsAmplifyRounding: summing the RMSNorm squares in float64 instead of float32, a 1e-7 change, moves the emulated logits by 0.12 to 0.24 nats and the f32 path by 1e-5), so only a bit-exact port of ggml's CPU kernels could match. The default f16 KV cache adds about 1e-3 to 5e-3. So the Q4_0 run is compared only loosely (TestOracleLlamaCpp): a token differing is accepted at a near-tie of llama.cpp's own distribution (at most 0.25 nats, after 6 equal tokens) and step-1 differences within 0.5 nats; on two of four prompts the f32 path diverges from it at the 10th and 22nd token.

Results (2026-11-04, v0.3.0, Apple M4 Max, Chromium, Metal, f32 kernels)

SmolLM2-135M-Instruct Q4_0, one container, greedy, 128-token prompt and 128 generated tokens, interleaved rounds, 1-minute load 3.5 to 5.5. Profile, every change with its tok/s, loop budgets and the structural limit: docs/perf/2026-11-04.md.

prefill tok/s (p50, min-max) decode tok/s (p50, min-max)
container v0.3.0, WebGPU, f32 5359 (5152-5472) 580 (572-601)
container v0.2.0 (same rounds) 312 (306-321) 154 (149-161)
llama.cpp 0.8.0, one CPU core (llama-bench -t 1 -ngl 0, go-pkgx bottle) 1268 (1040-1281) 514 (499-525)
WebLLM q4f32_1 (go-wasmcontainers/gpu docs/llm-baseline, not re-run) 2656 224

v0.2.0 was slower than the one-core CPU reference (171 against 529 decode): 405 tiny dependent dispatches and a serial attention kernel. v0.3.0 fuses a layer into five dispatches (RMSNorm, QKV, RoPE and KV append; attention in chunks merged in the output projection; RMSNorm, gate, up and SiLU; down), computes 8 tokens per decode submit and runs prompt chunks of up to 64 tokens through the layers together. Models whose layers are not Q4_0 (down may be Q4_1) keep the v0.2.0 kernels. Firefox and non-Apple adapters are not measured.

Correctness: the strict oracle (llama.cpp on the same weights dequantised to F32) gives 32 of 32 greedy tokens on all four prompts. GPU against CPU path: logits within 1.5e-6 relative, 16 of 16 tokens equal, on a 4-token and a 187-token prompt.

Documentation

Index

Constants

View Source 
const ChunkSize = 16 << 20

ChunkSize is the size of one weight upload: tensors are read from the model file and sent to the GPU in pieces of at most this many bytes, so the whole model never sits in guest memory.

View Source 
const MaxBatch = 16

MaxBatch bounds Batch.

View Source 
const PieceSize = 4 << 20

PieceSize is the size of one read from the model file while loading.

View Source 
const PrefillChunk = 8

PrefillChunk is the number of prompt tokens recorded per submit: a submit must stay far below the host's 2 s limit.

Variables

View Source 
var DecodeBatch = 8

DecodeBatch is the number of tokens a decode submit computes (Generate): the crossing and the read-back are paid once for that many tokens. After the end-of-sequence token or a stop from emit, at most DecodeBatch-1 tokens have been computed for nothing.

View Source 
var NoFused bool

NoFused turns the fused layer kernels off (the v0.2.0 path, for A/B runs).

View Source 
var PrefillBatch = 64

PrefillBatch is the number of prompt tokens the fused layers process together in one pass: every weight matrix is read once per chunk and the dispatches are shared. A chunk of 64 tokens takes about 12 ms of GPU time on an Apple M4 Max, far below the 500 ms budget of a submit.

View Source 
var TPR = 8

TPR is the number of invocations that share one output row of the fused matrix-vector kernels (a power of two dividing the workgroup size; the workgroup then computes WG/TPR rows). It trades the reduction depth against the rows per workgroup.

Functions

func Argmax 

func Argmax(l []float32) int32

Argmax returns the index of the largest logit (the first on ties, as llama.cpp's greedy sampler does).

func F16ToF32 

func F16ToF32(h uint16) float32

F16ToF32 converts an IEEE binary16 value, exactly.

func F32ToF16 

func F32ToF16(f float32) uint16

F32ToF16 converts to binary16, round to nearest even.

func RoundF16 

func RoundF16(f float32) float32

RoundF16 rounds f through binary16 (what an f16 KV cache stores).

func Shaders 

func Shaders(wg int) map[string]string

Shaders returns every kernel's WGSL for workgroup size wg, by name: for tools that validate them (cmd/wgsldump).

Types

type CPU 

type CPU struct {
	Cfg       Config
	ExactGGML bool
	// F16KV alone stores the KV cache in binary16 (what llama.cpp's default
	// cache type does) without quantising activations.
	F16KV bool

	// MaxLayers, when non-zero, runs only that many layers and returns nil
	// (debugging); LastX is the residual stream after the layers run.
	MaxLayers int
	LastX     []float32
	// contains filtered or unexported fields
}

CPU is the forward pass in Go: the equality oracle for the WGSL kernels and the native fallback. It holds every weight in memory.

ExactGGML reproduces llama.cpp's CPU backend: activations are quantised to Q8_0 (Q8_1 for Q4_1 weights) before every quantised matrix product and the KV cache is stored in binary16. Without it the products use f32 activations and an f32 cache, which is what the GPU kernels compute.

func LoadCPU 

func LoadCPU(f *gguf.File, r io.ReaderAt, maxCtx int) (*CPU, error)

LoadCPU reads every tensor of f from r.

func (*CPU) Forward 

func (m *CPU) Forward(tok int32) []float32

Forward runs one token at the next position and returns the logits.

func (*CPU) Pos 

func (m *CPU) Pos() int

Pos is the number of tokens in the cache.

func (*CPU) Reset 

func (m *CPU) Reset()

Reset clears the KV cache.

type Config 

type Config struct {
	Layers, Embd, FF, Heads, KVHeads, HeadDim, Vocab, Ctx int
	RopeBase, Eps                                         float32
}

Config is the llama-architecture hyperparameters read from GGUF metadata (llama.* keys; gguf.md "LLM" section).

func ConfigFrom 

func ConfigFrom(f *gguf.File) (Config, error)

ConfigFrom reads and checks the metadata. Only architecture "llama" with rope dimension equal to the head size and no rope scaling is supported.

type GPU 

type GPU struct {
	Cfg Config

	Stats Stats

	// Layers, when non-zero, runs only that many layers (debugging; no
	// logits and no sampling are then produced).
	Layers int
	Skip   map[string]bool
	// contains filtered or unexported fields
}

GPU is the forward pass on the GPU for the llama architecture, in f32.

func NewGPU 

func NewGPU(dev *webgpu.Device, lim Limits, f *gguf.File, r io.ReaderAt, ctx int) (*GPU, error)

NewGPU creates the kernels, allocates every buffer and streams the weights of the model to the GPU. ctx is the KV cache length in tokens.

func (*GPU) Generate 

func (g *GPU) Generate(prompt []int32, n int, eos int32, emit func(int32) bool) (Timing, error)

Generate decodes greedily: the prompt in submits of PrefillChunk tokens, then one submit per token whose only read-back is the 4-byte sampled token (argmax on the GPU). emit is called with each generated token; Generate stops after n tokens, at the end-of-sequence token (which is not emitted) or when emit returns false.

func (*GPU) MemoryUsed 

func (g *GPU) MemoryUsed() uint64

MemoryUsed is the GPU memory the model, its cache and its scratch hold.

func (*GPU) Peek 

func (g *GPU) Peek(name string, n int) ([]float32, error)

Peek reads the first n floats of an activation buffer ("x", "h", "q", "k", "v", "att", "gate", "up", "logits"): for tests, never in the decode loop.

func (*GPU) PrefillLogits 

func (g *GPU) PrefillLogits(prompt []int32) ([]float32, error)

PrefillLogits runs the prompt and returns the logits of its last token (for tests: reading 4*vocab bytes per token is what Generate avoids).

type Limits 

type Limits struct {
	MaxInvocationsPerWorkgroup uint32
	MaxWorkgroupStorage        uint32 // bytes
	MaxWorkgroupsPerDimension  uint32
	MaxStorageBuffers          uint32 // per shader stage
	MaxBindingSize             uint64 // maxStorageBufferBindingSize
	MaxBufferSize              uint64
	MemoryBudget               uint64 // GPU memory the container may hold (capset max_memory)
}

Limits are the adapter limits the kernels adapt to (the capset's DeviceLimits, go-virtio/gpu wgpubatch). Nothing in the kernels assumes a vendor's values: the workgroup size and the grids are derived from these.

func DefaultLimits 

func DefaultLimits() Limits

DefaultLimits are WebGPU's defaults (the guaranteed minimum).

func (Limits) WorkgroupSize 

func (l Limits) WorkgroupSize() int

WorkgroupSize is the workgroup size of the reductions: 64 (a size every WebGPU adapter supports, a multiple of the 32-wide and 64-wide waves of current hardware but not relying on either), reduced when the adapter's invocation or workgroup-memory limit is lower. A power of two.

type Stats 

type Stats struct {
	Submits    int
	Dispatches int
	Tokens     int
}

Stats counts what a run did.

type Timing 

type Timing struct {
	Prefill, Decode time.Duration
	PromptTokens    int
	Generated       int
	Submits         int
	Dispatches      int
	// Decode time split: the guest recording the pass (Go, in wasm), and
	// Submit (the crossing, the replay, the GPU, the read-back).
	Record, Submit time.Duration
}

Timing of a generation.

type Tokenizer 

type Tokenizer struct {
	EOS    int32
	BOS    int32
	AddBOS bool
	// contains filtered or unexported fields
}

Tokenizer is the byte-level BPE ("gpt2") tokenizer stored in GGUF metadata (tokenizer.ggml.tokens, .merges, .token_type), with the "smollm" pre-tokenizer of llama.cpp's llama-vocab.cpp: digits split one by one, then the GPT-2 pattern 's|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S).

func NewTokenizer 

func NewTokenizer(f *gguf.File) (*Tokenizer, error)

NewTokenizer builds the tokenizer from a parsed header.

func (*Tokenizer) Decode 

func (t *Tokenizer) Decode(ids []int32) string

Decode returns the text of tokens. Special tokens are written as they are.

func (*Tokenizer) Encode 

func (t *Tokenizer) Encode(text string) []int32

Encode tokenizes text. Special tokens written in the text are matched as such (llama.cpp's parse_special).

func (*Tokenizer) Piece 

func (t *Tokenizer) Piece(id int32) []byte

Piece is the bytes of one token.

func (*Tokenizer) VocabSize 

func (t *Tokenizer) VocabSize() int

VocabSize is the number of tokens.

Source Files

View all Source files

Directories

Path Synopsis
cmd
ggufdump command
Command ggufdump prints the tensors of a GGUF file: name, ggml type, dims.
 Command ggufdump prints the tensors of a GGUF file: name, ggml type, dims.
wgsldump command
Command wgsldump prints the WGSL kernels as JSON ({"name": "source"}), for validation in a browser (go-wasmcontainers/gpu e2e TestLLMShaders).
 Command wgsldump prints the WGSL kernels as JSON ({"name": "source"}), for validation in a browser (go-wasmcontainers/gpu e2e TestLLMShaders).
examples
generate command
Command generate is a WASI container that generates text with a GGUF llama model on the GPU, through the batched command stream (wasi_webgpu_batch).
 Command generate is a WASI container that generates text with a GGUF llama model on the GPU, through the batched command stream (wasi_webgpu_batch).
Package gguf parses the header of a GGUF file (ggml docs/gguf.md, version 3): metadata key/values and tensor infos.
 Package gguf parses the header of a GGUF file (ggml docs/gguf.md, version 3): metadata key/values and tensor infos.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.