> ## Documentation Index > Fetch the complete documentation index at: https://antlobach-clorch-182a83cb.mintlify.site/llms.txt > Use this file to discover all available pages before exploring further. # LLM Architecture Components: RoPE, GQA, SwiGLU, and More > RMSNorm, RoPE, grouped-query attention, SwiGLU, fused SDPA, KV cache, and autoregressive generation: the full set of LLM components available in Clorch. Clorch ships a full set of LLM-relevant architecture components that map directly onto modern transformer designs like Llama-3. All components are implemented as standard `IModule` values: they compose with `nn/forward`, participate in `nn/parameters`, transfer between devices with `nn/to`, and work within `nn/defmodel` blocks. This page describes each component and shows how they fit together in a complete model. ```clojure theme={null} (require '[clorch.nn :as nn] '[clorch.nn.functional :as F] '[clorch.torch :as t]) ``` *** ## Token Embeddings ### `nn/embedding` Creates a learnable embedding table with `num-embeddings` rows and `embedding-dim` columns. ```clojure theme={null} (def tok-emb (nn/embedding 50257 128 :_freeze false)) ;; Forward: input is an integer tensor of token IDs (nn/forward tok-emb idx) ;; idx shape [B, T] → output shape [B, T, 128] ``` **Options:** | Key | Default | Description | | - | - | - | | `:padding-idx` | `nil` | Index whose embedding is zeroed and not updated | | `:max-norm` | `nil` | Renormalize embeddings exceeding this norm | | `:scale-grad-by-freq` | `false` | Scale gradients by inverse frequency | | `:sparse` | `false` | Use sparse gradient updates | | `:_freeze` | `true` | Freeze weights (set `false` to train) | ### `nn/embedding-from-pretrained` Initialize an embedding table from an existing tensor (e.g., GloVe or word2vec weights). ```clojure theme={null} (def emb (nn/embedding-from-pretrained pretrained-weights :freeze false)) ``` *** ## RMSNorm `nn/rmsnorm` implements Root Mean Square Layer Normalization, which is preferred over LayerNorm in most modern LLMs (including Llama) because it omits the mean-centering step. ```clojure theme={null} (def norm (nn/rmsnorm 128)) (def norm (nn/rmsnorm 128 :eps 1e-6)) (nn/forward norm x) ;; x shape [B, T, D] → normalized [B, T, D] ``` Internally, the forward pass: 1. Casts input to float32 for numerical stability. 2. Computes RMS: `sqrt(mean(x^2) + eps)`. 3. Normalizes: `x / RMS`. 4. Scales by the learned weight `gamma` (initialized to ones). 5. Casts result back to the input dtype. *** ## Rotary Position Embeddings (RoPE) RoPE encodes position information directly into query and key tensors by rotating them with sinusoidal frequencies, enabling the model to generalize to longer sequences than it was trained on. ### `torch/precompute-rope-freqs` Precomputes the cosine and sine embedding matrices for a given head dimension and sequence length. Returns a two-element vector `[cos-emb sin-emb]`. ```clojure theme={null} (def head-dim (quot emb-dim n-heads)) ;; Returns [cos-emb, sin-emb], each shape [seq-len, head-dim] (def freqs (t/precompute-rope-freqs head-dim 64)) ;; Custom theta (default 10000.0) (def freqs (t/precompute-rope-freqs head-dim 64 :theta 500000.0)) ``` ### `torch/apply-rope` Applies the precomputed rotation to query or key tensors. The input must have shape `[B, T, n-heads, head-dim]`. ```clojure theme={null} (let [[cos-emb sin-emb] freqs] (def q-rotated (t/apply-rope q cos-emb sin-emb)) (def k-rotated (t/apply-rope k cos-emb sin-emb))) ``` The rotation applies the standard formula: `[x1, x2] → [x1*cos − x2*sin, x1*sin + x2*cos]`. *** ## Grouped-Query Attention (GQA) `nn/GroupedQueryAttention` implements multi-head attention where the number of key-value heads (`n-kv-heads`) can be less than the number of query heads (`n-heads`). When they are equal it degrades to standard multi-head attention; with `n-kv-heads = 1` it becomes multi-query attention. ```clojure theme={null} (def gqa (nn/GroupedQueryAttention dim ;; model dimension n-heads ;; number of query heads n-kv-heads ;; number of K/V heads (≤ n-heads) context-len ;; maximum sequence length drop-rate)) ;; attention dropout ``` ### Forward Pass The forward pass accepts either a tensor or a map with optional RoPE frequencies, causal mask, and KV cache: ```clojure theme={null} ;; Plain forward (nn/forward gqa x) ;; With RoPE and causal mask (nn/forward gqa {:x x :mask mask :freqs freqs}) ;; With KV cache for autoregressive generation (nn/forward gqa {:x x :freqs freqs :kv-cache my-cache-atom}) ``` The module handles K/V head repetition internally: if `n-kv-heads < n-heads`, it repeats the K and V tensors with `torch/repeat-interleave` to match the query head count. *** ## SwiGLU Feed-Forward Block `nn/SwiGLU` implements the gated linear unit feed-forward block used in Llama-style models. It applies a SiLU-gated projection that selectively amplifies features: ```text theme={null} output = W3(SiLU(W1(x)) ⊙ W2(x)) ``` ```clojure theme={null} (def ffwd (nn/SwiGLU dim (* 4 dim))) (nn/forward ffwd x) ;; x shape [B, T, dim] → [B, T, dim] ``` SwiGLU has three linear projections (`w1`, `w2`, `w3`) versus the two in a standard MLP. The hidden dimension is typically `4 × dim`. *** ## Fused Scaled Dot-Product Attention `F/scaled-dot-product-attention` wraps LibTorch's fused SDPA dispatcher. On CUDA, LibTorch automatically selects Flash Attention, memory-efficient attention, or the math kernel based on dtype, shape, mask presence, and hardware capability. ```clojure theme={null} ;; Basic call (F/scaled-dot-product-attention query key value) ;; With attention mask (for causal or custom patterns) (F/scaled-dot-product-attention query key value :attention-mask causal-mask :dropout-p 0.0 :causal? false :scale nil ;; defaults to 1/sqrt(head-dim) :enable-gqa? false) ``` Query, key, and value tensors should have shape `[B, n-heads, T, head-dim]`. The release suite verifies fused SDPA as part of the CUDA path. Using this function on a CUDA device will engage Flash Attention automatically when the hardware and dtype support it. *** ## Causal Attention Masks Most autoregressive models require a causal mask that prevents each position from attending to future positions. ```clojure theme={null} ;; Build a boolean mask: True where attention should be blocked (def seq-len 64) (def mask (t/eq (t/tril (t/ones [seq-len seq-len])) 0)) ;; mask[i,j] = true when j > i (future positions) ;; In the attention module, apply with masked-fill (def masked-scores (t/masked-fill scores mask t/-inf)) ``` For KV-cache generation where only one new token is processed at a time, no mask is needed (the K/V cache tracks history implicitly). *** ## Autoregressive Generation `nn/generate` runs a greedy autoregressive generation loop without gradient tracking. ```clojure theme={null} (def output-ids (nn/generate model start-ids ;; initial token IDs, shape [1, T] 100 ;; max-new-tokens context-size)) ;; context window length ``` It samples the next token from the softmax of the last position's logits using `torch/multinomial`. When the sequence exceeds `context-size`, it crops to the most recent `context-size` tokens before each forward pass. *** ## KV Cache The KV cache stores past key and value tensors to avoid recomputing them on every generation step. In Clorch's Llama-style examples, caches are plain Clojure atoms holding maps. ### Pattern from `modern_llama.clj` ```clojure theme={null} ;; One atom per transformer block (def cache (atom {})) ;; Prefill: process a prompt of 10 tokens (def freqs1 (t/precompute-rope-freqs head-dim 10)) (def out1 (nn/forward block {:x x1 :freqs freqs1 :kv-cache cache})) ;; cache now holds {:k-prev [1,10,n-kv-heads,head-dim] ;; :v-prev [1,10,n-kv-heads,head-dim]} ;; Decode: process one new token, provide only its RoPE slice (def all-freqs (t/precompute-rope-freqs head-dim 11)) (def freqs2 [(t/ix (first all-freqs) [10 11]) (t/ix (second all-freqs) [10 11])]) (def out2 (nn/forward block {:x x2 :freqs freqs2 :kv-cache cache})) ;; cache now holds concatenated K/V for 11 positions ``` Inside `GroupedQueryAttention`, the cache update looks like: ```clojure theme={null} (let [{:keys [k-prev v-prev]} @kv-cache k-curr (if k-prev (t/cat [k-prev k] 1) k) v-curr (if v-prev (t/cat [v-prev v] 1) v)] (reset! kv-cache {:k-prev k-curr :v-prev v-curr}) [k-curr v-curr]) ``` ### Pattern from `nanochat.clj` The NanoChat example allocates one cache atom per transformer block and manages their lifetimes explicitly with `torch/retain!` and `torch/release!` to prevent native memory leaks during long generation loops: ```clojure theme={null} (def caches (vec (repeatedly n-layers #(atom {})))) ;; After each step, retain new cache tensors and release old ones (when caches (run! #(t/retain! @%) caches)) ;; ... (when old-caches (run! t/release! old-caches)) ``` *** ## Complete Llama Block Example The following block definition from `examples/modern_llama.clj` assembles all LLM components into a standard Llama-style transformer block: ```clojure theme={null} (nn/defmodel ModernLlamaBlock [dim n-heads n-kv-heads context-len drop-rate] [sa (nn/GroupedQueryAttention dim n-heads n-kv-heads context-len drop-rate) ffwd (nn/SwiGLU dim (* 4 dim)) ln1 (nn/rmsnorm dim) ln2 (nn/rmsnorm dim)] (forward [input] (let [{:keys [x mask freqs kv-cache]} (if (map? input) input {:x input}) x (t/add x (nn/forward sa {:x (nn/forward ln1 x) :mask mask :freqs freqs :kv-cache kv-cache})) x (t/add x (nn/forward ffwd (nn/forward ln2 x)))] x))) ``` *** ## LLM Capability Table | Capability | Status | Clorch surface | | - | - | - | | Token embeddings | ✅ Implemented | `nn/embedding`, `nn/embedding-from-pretrained` | | RMSNorm | ✅ Implemented | `nn/rmsnorm` | | Rotary position embeddings | ✅ Implemented | `t/precompute-rope-freqs`, `t/apply-rope` | | Grouped-query attention | ✅ Implemented | `nn/GroupedQueryAttention` | | SwiGLU feed-forward | ✅ Implemented | `nn/SwiGLU` | | Flash/fused SDPA | ✅ Implemented | `F/scaled-dot-product-attention` | | Causal attention masks | ✅ Implemented | Tensor masking + Llama/GPT examples | | KV cache | ✅ Implemented | `examples/modern_llama.clj`, `examples/nanochat.clj` | | Autoregressive generation | ✅ Implemented | `nn/generate`, NanoChat generation loop | | Llama-style blocks | ✅ Implemented | `examples/modern_llama.clj`, `examples/nanochat.clj` | | Mixed-precision training | ✅ Implemented | `clorch.amp` | | Quantized LLM inference | ❌ Not complete | End-to-end quantization pending | | FSDP / tensor parallelism | ❌ Not complete | NCCL DDP only today | *** ## Reference Examples `examples/modern_llama.clj` — a standalone demonstration of a single Llama block with RoPE, GQA, SwiGLU, and KV caching. Includes prefill and single-token decode to verify the cache shape transitions. `examples/nanochat.clj` — a compact Llama-3-style model trained on a small text corpus. Demonstrates full training, checkpointing, streaming generation with KV caching, and an interactive chat loop. This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.