LLM Architecture Components: RoPE, GQA, SwiGLU, and More

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:

Casts input to float32 for numerical stability.
Computes RMS: sqrt(mean(x^2) + eps).
Normalizes: x / RMS.
Scales by the learned weight gamma (initialized to ones).
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.

Get Started

Core Concepts

Building Models

Advanced Topics

Reference

LLM Architecture Components: RoPE, GQA, SwiGLU, and More

Token Embeddings

`nn/embedding`

`nn/embedding-from-pretrained`

RMSNorm

Rotary Position Embeddings (RoPE)

`torch/precompute-rope-freqs`

`torch/apply-rope`

Grouped-Query Attention (GQA)

Forward Pass

SwiGLU Feed-Forward Block

Fused Scaled Dot-Product Attention

Causal Attention Masks

Autoregressive Generation

KV Cache

Pattern from `modern_llama.clj`

Pattern from `nanochat.clj`

Complete Llama Block Example

LLM Capability Table

Reference Examples

Modern Llama

NanoChat

Build docs developers (and LLMs) love

Get Started

Core Concepts

Building Models

Advanced Topics

Reference

Documentation Index

​Token Embeddings

​nn/embedding

​nn/embedding-from-pretrained

​RMSNorm

​Rotary Position Embeddings (RoPE)

​torch/precompute-rope-freqs

​torch/apply-rope

​Grouped-Query Attention (GQA)

​Forward Pass

​SwiGLU Feed-Forward Block

​Fused Scaled Dot-Product Attention

​Causal Attention Masks

​Autoregressive Generation

​KV Cache

​Pattern from modern_llama.clj

​Pattern from nanochat.clj

​Complete Llama Block Example

​LLM Capability Table

​Reference Examples

Modern Llama

NanoChat

Build docs developers (and LLMs) love

Token Embeddings

`nn/embedding`

`nn/embedding-from-pretrained`

RMSNorm

Rotary Position Embeddings (RoPE)

`torch/precompute-rope-freqs`

`torch/apply-rope`

Grouped-Query Attention (GQA)

Forward Pass

SwiGLU Feed-Forward Block

Fused Scaled Dot-Product Attention

Causal Attention Masks

Autoregressive Generation

KV Cache

Pattern from `modern_llama.clj`

Pattern from `nanochat.clj`

Complete Llama Block Example

LLM Capability Table

Reference Examples