Pantheon/docs/architecture/llama-forward.md

Llama Forward Pass

The Llama forward pass (src/models/llama.c) is shared by Llama, Qwen2, Mistral, and Mixtral architectures. Architecture-specific differences (attention biases, sliding window, MoE routing, metadata keys) are handled by config flags, not separate code paths.

Phi-3 has its own forward pass in src/models/phi3.c due to fused QKV/gate+up weights and NeoX-style RoPE. See the architecture differences table below.

Forward Pass Flow

token (uint32)
  │
  ├── embed: tok_embd[token] → hidden[n_embd]
  │
  ├── for each layer (0..n_layer-1):
  │   │
  │   ├── RMSNorm(hidden, attn_norm)
  │   │
  │   ├── Attention:
  │   │   ├── Q = Wq @ normed  (+bias for Qwen2)
  │   │   ├── K = Wk @ normed  (+bias for Qwen2)
  │   │   ├── V = Wv @ normed  (+bias for Qwen2)
  │   │   ├── RoPE(Q, K, pos, rope_theta)
  │   │   ├── Write K,V to KV cache at pos
  │   │   ├── GQA: for each query head
  │   │   │   ├── scores = Q_head @ K_cache^T / sqrt(head_dim)
  │   │   │   ├── causal mask (future tokens = -inf)
  │   │   │   ├── softmax(scores)
  │   │   │   └── output = scores @ V_cache
  │   │   └── O = Wo @ concat(heads)  (+bias for Qwen2)
  │   │
  │   ├── residual: hidden += O
  │   │
  │   ├── RMSNorm(hidden, ffn_norm)
  │   │
  │   ├── FFN (standard) or MoE FFN (Mixtral):
  │   │   │
  │   │   ├── Standard SwiGLU (n_experts == 0):
  │   │   │   ├── gate = W_gate @ normed
  │   │   │   ├── up   = W_up   @ normed
  │   │   │   ├── act  = silu(gate) * up
  │   │   │   └── down = W_down @ act
  │   │   │
  │   │   └── MoE routing (n_experts > 0, Mixtral):
  │   │       ├── gate_logits = W_router @ normed    → [n_experts]
  │   │       ├── probs = softmax(gate_logits)
  │   │       ├── top_k = select top n_experts_used
  │   │       ├── weights = renormalize(selected probs)
  │   │       ├── for each selected expert e:
  │   │       │   ├── gate_e = expert[e].W_gate @ normed
  │   │       │   ├── up_e   = expert[e].W_up   @ normed
  │   │       │   ├── act_e  = silu(gate_e) * up_e
  │   │       │   └── out_e  = expert[e].W_down @ act_e
  │   │       └── down = Σ (weight_e * out_e)
  │   │
  │   └── residual: hidden += down
  │
  ├── RMSNorm(hidden, output_norm)
  │
  └── logits = W_output @ normed  →  float[n_vocab]

Key Operations

RMSNorm

rms = sqrt(mean(x^2) + eps)
out = (x / rms) * weight

Implemented as a single SIMD kernel (pn_vrms_norm_f32). Default eps = 1e-5, configurable per model via GGUF metadata.

RoPE (Rotary Position Embeddings)

Applies learned rotation to Q and K vectors based on position:

for each pair (i, i+1) in head:
    freq = 1.0 / (theta ^ (2i / head_dim))
    angle = pos * freq
    q[i], q[i+1] = rotate(q[i], q[i+1], angle)

rope_theta defaults to 10000 (Llama). Qwen2.5 uses 1,000,000 for extended context.

Grouped Query Attention (GQA)

GQA reduces KV heads relative to query heads. Each KV head is shared across n_head / n_kv_head query heads:

Model	Q Heads	KV Heads	Ratio
SmolLM2-135M	9	3	3:1
Qwen2.5-0.5B	14	2	7:1
Mistral-7B	32	8	4:1
Phi-3-mini	32	32	1:1 (MHA)

The forward pass maps query head h to KV head h * n_kv_head / n_head.

SwiGLU FFN

gate = W_gate @ x
up   = W_up   @ x
out  = W_down @ (silu(gate) * up)

The silu(gate) * up step is fused into a single SIMD kernel (pn_vsilu_mul_f32).

Sliding Window (Mistral)

When sliding_window > 0, the KV cache operates as a circular buffer:

• Write position: pos % max_seq (wraps around)
• Readable length: min(pos + 1, max_seq) (prevents reading uninitialized slots)
• max_seq = min(sliding_window, n_ctx)

This limits memory to a fixed window size regardless of sequence length. Standard Llama/Qwen2 use sliding_window = 0 (full context).

Mixture of Experts (Mixtral)

When n_experts > 0, the FFN sublayer is replaced with expert routing:

normed = RMSNorm(hidden)
gate_logits = W_router @ normed                   → [n_experts]
probs = softmax(gate_logits)
top_idx, top_wt = top_k(probs, n_experts_used)    → select best experts
top_wt = renormalize(top_wt)                       → weights sum to 1

ffn_out = 0
for each selected expert e:
    gate = expert[e].W_gate @ normed
    up   = expert[e].W_up   @ normed
    hid  = silu(gate) * up
    out  = expert[e].W_down @ hid
    ffn_out += top_wt[e] * out

hidden += ffn_out

Expert weight formats in GGUF:

Merged (most common): 3D tensors sliced into per-expert views (zero-copy mmap)

• blk.N.ffn_gate_exps.weight [n_experts, n_ff, n_embd]
• blk.N.ffn_up_exps.weight [n_experts, n_ff, n_embd]
• blk.N.ffn_down_exps.weight [n_experts, n_embd, n_ff] (transposed)

Split (older): individual tensors per expert

• blk.N.ffn_gate.E.weight [n_ff, n_embd] for E=0..n_experts-1

The loader probes for merged format first, falls back to split.

Architecture Differences

Shared Llama Forward Pass (llama.c)

Feature	Llama	Qwen2	Mistral	Mixtral
Attention biases	No	Yes (Q,K,V,O)	No	No
Sliding window	No	No	Yes (4096)	No
MoE experts	No	No	No	Yes (n_experts, top-k)
GGUF key prefix	llama.*	qwen2.*	llama.*	llama.*
GGUF arch string	llama	qwen2	mistral	llama (detected via expert_count)
RoPE theta	10000	1000000	10000	10000
RoPE style	Interleaved	Interleaved	Interleaved	Interleaved
Q/K/V weights	Separate	Separate	Separate	Separate
FFN weights	Separate	Separate	Separate	Per-expert (merged or split)
Tied embeddings	model-dependent	model-dependent	No	No

Note: Mistral GGUFs use llama.* metadata keys (not mistral.*), but general.architecture = "mistral". Mixtral GGUFs use general.architecture = "llama" and are detected via llama.expert_count > 0.

Phi-3 (phi3.c — separate forward pass)

Feature	Phi-3
GGUF key prefix	phi3.*
Q/K/V weights	Fused attn_qkv.weight [3*n_embd, n_embd]
Gate/Up weights	Fused ffn_up.weight [2*n_ff, n_embd]
RoPE style	NeoX (split-half pairs: head[i], head[i+half])
RoPE dim	Configurable via rope.dimension_count (96 for mini)
Attention	MHA (n_head == n_kv_head = 32)
Tied embeddings	No (separate output.weight)

Phi-3 uses its own forward pass because the fused weight layout requires different matvec + split logic, and NeoX RoPE pairs elements differently from the interleaved style used by Llama/Qwen2/Mistral.

Weight Storage

Two modes, selected automatically based on quant type support:

Quantized (default)

All projection weights (Wq, Wk, Wv, Wo, W_gate, W_up, W_down, W_output) stay in mmap'd quantized form as pn_qweight structs. Matrix-vector multiply happens via pn_qmv() which dequantizes one block at a time.

Embeddings and norms are always F32 (owned memory), since they're accessed by index (embeddings) or element-wise (norms).

F32 legacy

All weights dequantized to tm_tensor on load. Fallback for unsupported quant types. Uses standard tm_matmul for projections.

Per-Token Memory

All intermediate buffers are allocated from a bump arena that resets after each token:

Buffer	Size	Used For
normed	n_embd	RMSNorm output
q, k, v	n_embd	Q/K/V projections
attn_out	n_embd	Attention output
ffn_gate, ffn_up	n_ff	FFN intermediates
ffn_act	n_ff	SiLU activation
ffn_out	n_embd	FFN output
logits_raw	n_vocab	LM head output
scores	max_seq per head	Attention scores

For MoE layers (Mixtral), additional per-layer buffers:

Buffer	Size	Used For
gate_logits	n_experts	Router output
ffn_out	n_embd	Weighted expert accumulator
exp_gate	n_ff	Per-expert gate projection
exp_up	n_ff	Per-expert up projection
exp_hid	n_ff	Per-expert SiLU activation
exp_out	n_embd	Per-expert down projection

Arena is sized at state creation: sum of all buffers * sizeof(float) * 1.25 (25% safety margin). Typical sizes: 1.5 KB (SmolLM2-135M), 11 KB (Mistral-7B), 19 KB (Mixtral_7Bx2).

Batched Multi-Sequence Decode

pn_llama_forward_batch() decodes one token for each of B independent sequences in a single pass, for bulk throughput (e.g. transcribing many pages at once). Each states[b] carries its own KV cache and position, so the sequences can have different lengths and each attends only to its own cache; row b's input is tokens[b] or, for image-token positions, an embd_rows[b] override. It writes [B, n_vocab] row-major logits and advances every state's pos by one.

The batched path is restricted to the quantized, non-MoE Llama/Qwen weights. For MoE, F32 legacy, or ternary weights it returns without mutating any state so the caller can fall back to per-sequence pn_llama_forward. Results are bit-identical to running B separate sequential decodes (verified by test_batch_decode_correct: a B=2 batch reproduces each sequence's standalone greedy stream with no cross-contamination).

Fused AMX-panel M=B matmul (pn_qmm)

The projection and FFN matmuls in the batched path go through pn_qmm (src/core/qweight.c), a fused quantized M=B kernel: out[b, r] = sum_c dequant(W[r, c]) * acts[b, c]. Each weight row is dequantized once into a small L2-resident panel and multiplied against all B activation vectors, so the quantized weight is read and dequantized once per step while the matmul runs at AMX throughput (the panel GEMM lowers to AMX under BLAS). This avoids both the NEON-dot ceiling of the single-vector pn_qmv and the full-weight F32 copy to DRAM that a plain BLAS GEMM would need.

Because the per-token matmul FLOPs are irreducible, throughput is compute-bound and plateaus rather than scaling linearly with B. Measured on LightOnOCR 1B Q4_K, M4, against 38.8 tok/s single-stream pn_qmv:

path	B=16	B=32
full-BLAS F32 copy	49.7	46.5 (DRAM-bound, regresses)
NEON fused dot	58.4	-
AMX panel (pn_qmm)	72.5	84.3

The AMX panel gives roughly 2.2x aggregate CPU throughput and settles near ~85 tok/s. B=1 is intentionally suboptimal (the panel adds a scratch pass over the fused pn_qmv), so single sequences should use pn_llama_forward. pn_qmm falls back to a per-b pn_qmv loop above PN_QMM_MAX_B (64) or for ternary or GPU-resident weights.

Flash (Blocked) Prefill Attention

Prefill over a long image-token sequence was the bulk-OCR bottleneck. The original per-query attention recomputed over the whole K/V cache once per query position, giving O(P^2) DRAM traffic across the ~1369 image-token queries of a page.

prefill_flash_range processes a block of queries (PF_BQ=32) together and streams K and V in blocks (PF_BK=256) that stay cache-resident and are reused across the whole query block. This cuts K/V reads from about P times to about P/BQ. The math is unchanged from the per-query path (full-precision softmax, with the scores tile kept L2-resident), and a per-query fallback is retained for the sliding-window wraparound case. It is parallelized over query heads.

Measured: prefill ~4.3 to ~3.2 s/page (~1.33x), bulk OCR ~6.3 to ~5.2 s/page steady state. It is pure CPU, so it helps regardless of the vision backend and without contending for the GPU (unlike GPU prefill below). Output stays correct and batched decode stays bit-identical.

Optional GPU Paths (off by default)

Two Metal offload paths exist behind build flags and environment switches. Both are off by default because, on the current hardware, they lose to the CPU AMX path or contend with the vision GPU. The shipped app builds Pantheon without USE_METAL, so both are inert there.

• GPU batched decode (PN_BATCH_GPU, gated at runtime by pn_qmm_set_gpu_decode): uploads the Q4_K decoder weights once (pn_qmv_gpu_upload) and runs the M=B GEMM on the GPU (fm_sgemm_q4_k). Bit-identical to the CPU path, but not a win: pn_qmm stages, dispatches, syncs, and copies back per weight, so a step does ~168 CPU/GPU round-trips and the per-call sync latency dominates (~68 tok/s at B=16 vs ~86 tok/s for the CPU AMX panel). Auto-batching cannot hide it because each GEMM's result is immediately consumed by a CPU op (bias, RoPE, attention). Beating AMX would need a GPU-resident decode (the whole layer chain on GPU, one sync per step); the flag is kept as the seam for that.

• GPU F16 prefill (PN_GPU_PREFILL): uploads the Q4_K decoder weights once and runs each prefill projection as an F16 MPS GEMM (fm_gemm_q4_k_f16) instead of the CPU dequant+AMX path. Also a net loss for bulk OCR, for two reasons: prefill is attention-bound, not GEMM-bound (the O(P^2) decoder attention over the image tokens dominates, so GPU-accelerating only the GEMMs is ~1.3x, not the ~5x a GEMM microbench suggests), and it contends with Core ML vision for the single GPU, breaking the CPU-compute / GPU-vision overlap that makes bulk fast (~6.3 s/page went to ~8-10 s/page). It is kept as an opt-in for the CPU-vision case, where the GPU is otherwise idle. Setting it also forces decode back onto the CPU AMX path (via the pn_qmm_set_gpu_decode gate) so the prefill upload does not drag decode onto the slower GPU route.

ABI note: pn_qweight.gpu is an unconditional struct field, present in every build (just NULL and unused without USE_METAL). It must not be wrapped in #ifdef USE_METAL: a conditional field makes the struct layout depend on the build flag, so a translation unit compiled the other way reads later fields at the wrong offset.