Documentation
¶
Overview ¶
Package llama provides a LLaMA model implementation for the Born ML framework.
It supports LLaMA 2 and LLaMA 3 architectures loaded from GGUF files. The model implements generate.LLMModel and can be used directly with generate.TextGenerator for text generation.
Example — loading and running inference:
backend := cpu.New()
model, err := llama.LoadGGUF("llama-3-8b-Q4_0.gguf", backend)
if err != nil {
log.Fatal(err)
}
tok, _ := tokenizer.LoadHF("tokenizer.json")
gen := generate.NewTextGenerator(model, tok, generate.DefaultSamplingConfig())
text, _ := gen.Generate("Hello, world", generate.DefaultGenerateConfig())
fmt.Println(text)
Index ¶
- type AttentionFunc
- type Config
- type Layer
- func (l *Layer[B]) DebugForward(x *tensor.Tensor[float32, B], cache *nn.KVCache[B], startPos int) (out, attnContrib, ffnContrib *tensor.Tensor[float32, B])
- func (l *Layer[B]) Forward(x *tensor.Tensor[float32, B], cache *nn.KVCache[B], startPos int) *tensor.Tensor[float32, B]
- func (l *Layer[B]) Parameters() []*nn.Parameter[B]
- type Model
- func (m *Model[B]) CPUEmbedData() []float32
- func (m *Model[B]) EnableCPUEmbedding()
- func (m *Model[B]) Forward(input *tensor.RawTensor, cache generate.KVCache, startPos int) *tensor.RawTensor
- func (m *Model[B]) Parameters() []*nn.Parameter[B]
- func (m *Model[B]) Release()
- func (m *Model[B]) SetAttentionFunc(fn AttentionFunc[B])
- func (m *Model[B]) VocabSize() int
- type ModelCache
- type Option
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
This section is empty.
Types ¶
type AttentionFunc ¶
type AttentionFunc[B tensor.Backend] func( q, k, v *tensor.Tensor[float32, B], mask *tensor.Tensor[float32, B], scale float32, ) (*tensor.Tensor[float32, B], *tensor.Tensor[float32, B])
AttentionFunc is a function that computes scaled-dot-product attention.
It replaces the default ScaledDotProductAttention inside each transformer layer. Useful for research experiments (Flash Attention, sliding-window attention, etc.).
Parameters:
- q: Query tensor [batch, n_heads, seq_q, head_dim]
- k: Key tensor [batch, n_heads, seq_k, head_dim]
- v: Value tensor [batch, n_heads, seq_k, head_dim]
- mask: Optional causal mask (nil for no mask)
- scale: Pre-computed 1/sqrt(head_dim)
Returns:
- output: Attended values [batch, n_heads, seq_q, head_dim]
- weights: Attention weights (may be nil if not needed by the implementation)
type Config ¶
type Config struct {
// VocabSize is the vocabulary size (number of token embeddings).
VocabSize int
// HiddenSize is the embedding/hidden dimension (d_model).
// Also called embedding_length in GGUF metadata.
HiddenSize int
// NumLayers is the number of transformer blocks.
// Also called block_count in GGUF metadata.
NumLayers int
// NumHeads is the number of query attention heads.
// Also called attention.head_count in GGUF metadata.
NumHeads int
// NumKVHeads is the number of key/value heads for GQA.
// Also called attention.head_count_kv in GGUF metadata.
// When NumKVHeads < NumHeads, the model uses Grouped Query Attention (GQA).
NumKVHeads int
// HeadDim is the dimension per attention head (HiddenSize / NumHeads).
HeadDim int
// FFNSize is the feed-forward network intermediate dimension.
// Also called feed_forward_length in GGUF metadata.
FFNSize int
// MaxSeqLen is the maximum context length.
// Also called context_length in GGUF metadata.
MaxSeqLen int
// RopeTheta is the base frequency for Rotary Position Embeddings.
// Defaults to 10000.0. LLaMA 3 uses 500000.0.
RopeTheta float64
// NormEpsilon is the epsilon for RMSNorm numerical stability.
// Defaults to 1e-5.
NormEpsilon float32
}
Config holds the hyperparameters for a LLaMA model.
These values are read from GGUF metadata and determine the model architecture. All fields have sensible defaults matching TinyLlama-1.1B.
func ConfigFromGGUF ¶
ConfigFromGGUF extracts model configuration from a parsed GGUF file.
It reads architecture-specific metadata keys following the GGUF specification. All fields absent from the metadata fall back to safe defaults.
func DefaultConfig ¶
func DefaultConfig() Config
DefaultConfig returns a Config matching TinyLlama-1.1B defaults.
Useful as a starting point for custom configurations or unit tests.
type Layer ¶
type Layer[B tensor.Backend] struct { AttnNorm *nn.RMSNorm[B] // Pre-attention layer norm. QProj *nn.Linear[B] // Query projection [hidden, n_q_heads * head_dim]. KProj *nn.Linear[B] // Key projection [hidden, n_kv_heads * head_dim]. VProj *nn.Linear[B] // Value projection [hidden, n_kv_heads * head_dim]. OProj *nn.Linear[B] // Output projection [n_q_heads * head_dim, hidden]. Rope *nn.RotaryEncoding[B] FFNNorm *nn.RMSNorm[B] // Pre-FFN layer norm. FFN *nn.SwiGLUFFN[B] // SwiGLU feed-forward network. // contains filtered or unexported fields }
Layer represents a single LLaMA transformer block.
Each layer applies:
- Pre-attention RMSNorm
- Self-attention (GQA) with RoPE
- Residual connection
- Pre-FFN RMSNorm
- SwiGLU FFN
- Residual connection
func (*Layer[B]) DebugForward ¶
func (l *Layer[B]) DebugForward( x *tensor.Tensor[float32, B], cache *nn.KVCache[B], startPos int, ) (out, attnContrib, ffnContrib *tensor.Tensor[float32, B])
DebugForward is like Forward but returns intermediate attention and FFN contributions for diagnostic inspection.
Input x: [batch, seq_len, hidden_size]. Returns out (final output), attnContrib, and ffnContrib tensors.
func (*Layer[B]) Forward ¶
func (l *Layer[B]) Forward( x *tensor.Tensor[float32, B], cache *nn.KVCache[B], startPos int, ) *tensor.Tensor[float32, B]
Forward computes the output of a single transformer layer.
Input x: [batch, seq_len, hidden_size]. Returns output with the same shape.
cache and startPos are used for incremental KV-cache decoding during inference. Pass cache=nil for training or non-cached inference.
func (*Layer[B]) Parameters ¶
Parameters returns all trainable parameters in this layer.
type Model ¶
type Model[B tensor.Backend] struct { // Embed is the token embedding table [vocab_size, hidden_size]. Embed *nn.Embedding[B] // Layers are the transformer decoder blocks. Layers []*Layer[B] // Norm is the final RMSNorm before the output projection. Norm *nn.RMSNorm[B] // Head is the language-model head that projects to vocabulary logits. // LLaMA may share weights between Embed and Head (tie_word_embeddings). Head *nn.Linear[B] // Config holds the model hyperparameters. Config Config // contains filtered or unexported fields }
Model is a LLaMA causal language model.
It implements generate.LLMModel and can be used with generate.TextGenerator for text generation. Weights can be loaded from GGUF via LoadGGUF.
Type parameter B is the computation backend (e.g., cpu.Backend, webgpu.Backend).
Example:
backend := cpu.New()
model, err := llama.LoadGGUF("model.gguf", backend)
if err != nil {
log.Fatal(err)
}
// model implements generate.LLMModel
cache := llama.NewModelCache(model.Config, model.Config.MaxSeqLen, backend)
gen := generate.NewTextGenerator(model, tok, generate.DefaultSamplingConfig())
gen.SetCache(cache)
func LoadGGUF ¶
LoadGGUF loads a LLaMA model from a GGUF file.
It reads architecture metadata, constructs a Model with the correct shape, and loads all weights — dequantizing quantized tensors (Q4_0, Q4_K, Q8_0, …) transparently.
Supported GGUF architectures: llama, mistral (uses same weight layout).
Example:
backend := cpu.New()
model, err := llama.LoadGGUF("llama-3-8b-Q4_0.gguf", backend)
if err != nil {
log.Fatal(err)
}
defer model.Close() // Release file handle (not needed for CPU backend, but good practice)
func NewModel ¶
NewModel creates an uninitialised LLaMA Model with random weights.
Weights must be loaded via LoadStateDict or by calling LoadGGUF. opts allow injecting a custom attention function (see WithAttentionFunc).
func (*Model[B]) CPUEmbedData ¶ added in v0.9.19
CPUEmbedData returns the CPU-side embedding data, or nil if not enabled.
func (*Model[B]) EnableCPUEmbedding ¶ added in v0.9.19
func (m *Model[B]) EnableCPUEmbedding()
EnableCPUEmbedding moves the embedding and LM head weights to CPU memory.
This is needed when the vocab is too large for a single WebGPU buffer (> 256 MB). After calling this, Forward() will:
- Look up token embeddings on CPU (fast: O(seq × hidden) index copy)
- Transfer the small hidden tensor [batch, seq, hidden] to GPU
- Run transformer layers on GPU
- Transfer final hidden states back to CPU
- Compute LM head matmul on CPU (only last position's logits)
Safe to call once after LoadGGUF. No-op if already enabled. Note: LoadGGUF automatically calls this when the backend is GPU and the embedding exceeds 256 MB, so manual calls are usually unnecessary.
func (*Model[B]) Forward ¶
func (m *Model[B]) Forward( input *tensor.RawTensor, cache generate.KVCache, startPos int, ) *tensor.RawTensor
Forward performs a forward pass and returns logits.
This method satisfies the generate.LLMModel interface.
Parameters:
- input: Token ids [batch, seq_len] as a *tensor.RawTensor (int32).
- cache: Optional *ModelCache[B] for incremental KV-cache decoding. Pass nil for full-sequence forward (training or non-cached inference). The cache must have been created with NewModelCache for this model.
- startPos: Position offset for RoPE (0 for full-sequence; equals number of tokens already in the cache for incremental decoding).
Returns logits [batch, seq_len, vocab_size] as a *tensor.RawTensor. When CPU embedding is enabled, returns [1, 1, vocab_size] (last position only) since the sampler only needs the last token's logits.
func (*Model[B]) Parameters ¶
Parameters returns all trainable parameters of the model.
func (*Model[B]) Release ¶ added in v0.9.2
func (m *Model[B]) Release()
Release frees all GPU buffers held by the model's parameters. Call when the model is no longer needed to reclaim GPU memory immediately instead of waiting for GC. Safe to call multiple times.
func (*Model[B]) SetAttentionFunc ¶
func (m *Model[B]) SetAttentionFunc(fn AttentionFunc[B])
SetAttentionFunc replaces the attention scoring function on all layers. Pass nil to restore default ScaledDotProductAttention.
type ModelCache ¶
ModelCache holds per-layer KV caches for a LLaMA model.
Each transformer layer requires its own independent KV cache; using a single cache for all layers would mix keys and values from different layers.
Pass a ModelCache to Model.Forward to enable incremental decoding. Create one with NewModelCache or let the model create one via NewCache.
func NewModelCache ¶
func NewModelCache[B tensor.Backend](cfg Config, maxSeqLen int, backend B) *ModelCache[B]
NewModelCache creates a per-layer KV cache for the given model configuration.
Parameters:
- cfg: Model configuration (used to determine cache shape per layer)
- maxSeqLen: Maximum sequence length (may differ from cfg.MaxSeqLen for memory control)
- backend: Computation backend
func (*ModelCache[B]) Clear ¶
func (c *ModelCache[B]) Clear()
Clear resets all per-layer caches. Satisfies the generate.KVCache interface so ModelCache can be passed to TextGenerator.
type Option ¶
Option is a functional option for configuring a LLaMA Model.
func WithAttentionFunc ¶
func WithAttentionFunc[B tensor.Backend](fn AttentionFunc[B]) Option[B]
WithAttentionFunc replaces the default scaled-dot-product attention with a custom implementation. Use for research (e.g., Flash Attention, local windows).
Source Files