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.

(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.

(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).

(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.

(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].

(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].

(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.

(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:

;; 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:

output = W3(SiLU(W1(x)) ⊙ W2(x))

(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.

;; 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.

;; 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.

(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

;; 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:

(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:

(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:

(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

Modern Llama

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.

NanoChat

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.