llama

A high-performance LLAMA2 inference implementation in CHICKEN Scheme, based on Andrej Karpathy's llama2.c and its OCaml port llama2.ml.

Description

This egg provides a complete implementation of the LLAMA2 transformer architecture for text generation. It features modular components and uses BLAS integration for high performance.

Requirements

System Dependencies

BLAS library (OpenBLAS, Intel MKL, or system BLAS)

CHICKEN Extensions

srfi-1 - List library
srfi-4 - Numeric vectors
srfi-42 - Comprehensions
srfi-69 - Hash tables
vector-lib - Vector utilities
blas - BLAS bindings
endian-blob - Endian-aware blob operations
endian-port - Endian-aware port operations
random-mtzig - Random number generation
getopt-long - Command-line option parsing
test - Unit testing

Installation

# Install system dependencies (Ubuntu/Debian)
sudo apt-get install chicken-bin libchicken-dev libopenblas-dev

# Install CHICKEN extensions
chicken-install llama

# Download example model (15M parameters, ~60MB) and tokenizer.bin
wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin
wget https://github.com/iraikov/llama-chicken/raw/refs/heads/main/tokenizer.bin

Quick Start

# Basic text generation
llama-cli -c stories15M.bin -p "Once upon a time"

# Creative generation with temperature
llama-cli -c stories15M.bin -t 0.8 -s 100 -p "The meaning of life is"

# Verify model integrity
llama-cli -c stories15M.bin --verify-checkpoint

API

Data Types

config

configrecord: Model configuration parameters.

make-config dim hidden-dim n-layers n-heads n-kv-heads vocab-size seq-len shared-weightsprocedure

Creates a new configuration object.

dim: Model embedding dimension
hidden-dim: FFN hidden layer dimension
n-layers: Number of transformer layers
n-heads: Number of attention heads
n-kv-heads: Number of key-value heads
vocab-size: Vocabulary size
seq-len: Maximum sequence length
shared-weights: Whether to share input/output embeddings

config-dim configprocedure
config-hidden-dim configprocedure
config-n-layers configprocedure
config-n-heads configprocedure
config-n-kv-heads configprocedure
config-vocab-size configprocedure
config-seq-len configprocedure
config-shared-weights configprocedure: Accessors for configuration fields.

transformer-weights

transformer-weightsrecord: Container for all model parameters including embeddings, attention weights, FFN weights, and RoPE frequencies.

make-transformer-weights token-embedding-table rms-att-weight wq wk wv wo rms-ffn-weight w1 w2 w3 rms-final-weight freq-cis-real freq-cis-imag wclsprocedure: Creates a new transformer weights object with all parameter matrices.

run-state

run-staterecord: Runtime state for transformer computation including hidden states, attention caches, and output logits.

make-run-state x xb q k v att key-cache value-cache xb2 hb hb2 logitsprocedure: Creates a new runtime state object.

run-state-x stateprocedure
run-state-logits stateprocedure
run-state-key-cache stateprocedure
run-state-value-cache stateprocedure: Accessors for runtime state fields.

args

argsrecord: Runtime configuration for text generation runs.

make-args checkpoint tokenizer temperature steps prompt seedprocedure

Creates text generation arguments.

checkpoint: Path to model checkpoint file
tokenizer: Path to tokenizer file
temperature: Sampling temperature (0.0 = greedy)
steps: Number of tokens to generate
prompt: Input text prompt
seed: Random seed (optional)

High-Level Functions

run argsprocedure

Main inference function. Takes an args object and performs text generation.

(define args (make-args "model.bin" "tokenizer.bin" 0.8 100 "Hello world" #f))
(run args)

transformer token pos config state weightsprocedure

Run transformer forward pass for a single token.

token: Token ID to process
pos: Position in sequence
config: Model configuration
state: Runtime state (modified in place)
weights: Model parameters

Returns the updated state.

bpe-encode text vocab vocab-scoresprocedure

Tokenize text using Byte-Pair Encoding.

text: Input text string
vocab: List of vocabulary strings
vocab-scores: List of BPE merge scores

Returns list of token IDs

Transformer Components

The modular architecture provides fine-grained control:

token-embedding-lookup state weights tokenprocedure: Load token embedding into state.

get-rope-frequencies weights pos head-sizeprocedure: Extract RoPE frequency rows for given position. Returns two values: real and imaginary frequency vectors.

attention-rmsnorm state weights layer-idx configprocedure: Apply RMS normalization for attention layer.

compute-qkv state weights layer-idx configprocedure: Compute Query, Key, Value matrices for given layer.

apply-rope state config freq-real freq-imagprocedure: Apply Rotary Position Embedding to Q and K vectors.

cache-kv state layer-idx pos configprocedure: Store current key and value vectors in attention cache.

compute-attention state layer-idx pos configprocedure: Compute multi-head attention scores and apply to values.

attention-output state weights layer-idx configprocedure: Apply final attention output projection.

ffn-rmsnorm state weights layer-idx configprocedure: Apply RMS normalization for feed-forward network.

compute-ffn-w1w3 state weights layer-idx configprocedure: Compute first part of FFN: W1(x) and W3(x).

apply-swiglu state configprocedure: Apply SwiGLU activation: SiLU(W1(x)) * W3(x).

ffn-output state weights layer-idx configprocedure: Apply final FFN linear transformation.

process-transformer-layer state weights layer-idx pos config freq-real freq-imagprocedure: Process complete transformer layer (attention + FFN blocks).

final-rmsnorm state weightsprocedure: Apply final RMS normalization before classification.

compute-logits state weights configprocedure: Compute final classification logits.

Utility Functions

rmsnorm output input weightsprocedure: RMS normalization with learnable weights.

matmul output input matrix rows colsprocedure: Matrix-vector multiplication using BLAS.

softmax output input sizeprocedure: Softmax activation with numerical stability.

accum target sourceprocedure: Vector accumulation for residual connections.

argmax vectorprocedure: Find index of maximum element (greedy sampling).

sample probabilities random-stateprocedure: Probabilistic sampling from probability distribution.

verify-checkpoint-data checkpoint-file #!optional detailedprocedure: Load and analyze checkpoint file, printing weight statistics.

Command-Line Interface

The llama-cli command provides easy access to text generation:

llama-cli [options]

Options:
  -h, --help            Show help message
  -c, --checkpoint FILE Model checkpoint file (required)  
  -k, --tokenizer FILE  Tokenizer file (default: tokenizer.bin)
  -t, --temperature NUM Sampling temperature (default: 0.0)
  -s, --steps NUM       Number of tokens to generate (default: 256)
  -p, --prompt TEXT     Input prompt text (default: empty)
  --seed NUM            Random seed for sampling
  --verify-checkpoint   Verify checkpoint integrity

Examples

Basic Usage

(import llama)

;; Simple text generation
(define args (make-args "stories15M.bin" "tokenizer.bin" 0.5 50 "Once upon a time" #f))
(run args)

Interactive REPL Usage

(import llama)

;; Load model components
(define config (make-config 288 768 6 6 6 32000 256 #t))
(define weights (load-checkpoint "stories15M.bin"))
(define state (make-run-state ...))

;; Generate single token
(transformer 1 0 config state weights)
(argmax (run-state-logits state))

;; Custom sampling with temperature
(define logits (run-state-logits state))
(do ((i 0 (+ i 1)))
    ((= i (f32vector-length logits)))
  (f32vector-set! logits i (/ (f32vector-ref logits i) 0.8)))

(define probs (softmax (make-f32vector 32000) logits 32000))
(sample probs random-state)

Batch Processing

;; Process multiple prompts
(define prompts '("Hello world" "The meaning of life" "Once upon a time"))

(for-each (lambda (prompt)
            (printf "Prompt: ~A~%" prompt)
            (let ((args (make-args "stories15M.bin" "tokenizer.bin" 0.5 50 prompt #f)))
              (run args)
              (newline)))
          prompts)

Component-Level Usage

;; Fine-grained control over generation
(define (custom-generation token config state weights)
  ;; Custom attention processing
  (attention-rmsnorm state weights 0 config)
  (compute-qkv state weights 0 config)
  
  ;; Skip some layers for faster inference  
  (let-values (((freq-real freq-imag) (get-rope-frequencies weights 0 2)))
    (process-transformer-layer state weights 0 0 config freq-real freq-imag)
    (process-transformer-layer state weights 2 0 config freq-real freq-imag))
  
  ;; Custom final processing
  (final-rmsnorm state weights)
  (compute-logits state weights config))