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
- 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]) Forward ¶
func (m *Model[B]) Forward( input *tensor.RawTensor, cache interface{ Clear() }, 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.
func (*Model[B]) Parameters ¶
Parameters returns all trainable parameters of the model.
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