llm

package module
v0.1.0 Latest Latest
Warning

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

 Go to latest Published: Nov 3, 2026 License: BSD-3-Clause Imports: 14 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.
  • One submit per token. Every dispatch of the forward pass is recorded into one compute pass; sampling (argmax) runs on the GPU and only the 4-byte token is read back. Prompt tokens go in submits of 8.
  • 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 tolerances (vendor-sensitive tests are marked)
  • Tokenizer: ids identical to llama.cpp's /tokenize on the fixed prompts.
  • CPU path against llama.cpp (TestOracleLlamaCpp, one thread, greedy): a token differing from llama.cpp's is accepted only at a near-tie of llama.cpp's own distribution (gap at most 0.25 nats) after at least 6 equal tokens; step-1 log-probability differences against its top-10 stay within 0.5 nats. Measured: the ggml-exact mode (Q8_0 activations, f16 KV cache) does not reproduce llama.cpp to 1e-3, so no tighter bound is claimed.
  • GPU against the CPU f32 path (vendor-sensitive): logits within 1e-3 relative to the largest logit, and equal greedy tokens. Fused multiply-adds and a different exp on another vendor move the last bits.
  • GPU against llama.cpp (vendor-sensitive): as the CPU path above.

Results (2026-11-03, Apple M4 Max, Chromium, Metal, f32 kernels)

SmolLM2-135M-Instruct Q4_0, one container, greedy, 128-token prompt and 128 generated tokens, five interleaved rounds against llama.cpp 0.8.0 on one CPU core (llama-bench -t 1 -ngl 0, the go-pkgx bottle, openssl 3.5.8), 1-minute load 6.4 to 7.7:

prefill tok/s (p50, min-max) decode tok/s (p50, min-max)
container, WebGPU, f32 311 (302-318) 171 (160-176)
llama.cpp CPU, one core 1286 (1278-1289) 529 (521-531)
WebLLM q4f32_1 (baseline, go-wasmcontainers/gpu docs/llm-baseline) 2656 224

Slower than the one-core CPU reference: an open defect, not yet optimised (faithful port first). Where the decode time goes (per token, 405 dispatches): 0.13-0.2 ms recording in the guest, 3.8 ms (short context) to 5.5-6 ms (context of 128 to 192) in submit, of which about 0.4 ms is the crossing (measured separately, docs/llm-design.md). So the GPU work itself, about 9 us per dispatch of tiny kernels (a 576-wide row is 72 words for 64 threads), and an attention kernel with one workgroup per head that walks the context serially, are the cause. Prefill is a decode pass per prompt token (8 per submit): no matrix-matrix kernel yet.

Correctness: on four fixed prompts the container's 32 greedy tokens equal llama.cpp's on two; the other two diverge at the 10th and 22nd token at a near-tie of llama.cpp's own distribution (0.16 and 0.12 nats), exactly as the Go CPU path (f32) does. GPU against CPU path: logits within 1.6e-6 relative, 16 of 16 tokens equal.

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 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

This section is empty.

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

	// 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
	// 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
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.