Transformer & Modern Large Language Models
Transformer & Modern Large Language Models
Located at:
netcl.nn.modern,netcl.nn.transformer,netcl.ops.attention,netcl.ops.norm_rope,netcl.text
netcl contains an end-to-end, high-performance pipeline for modern decoder-only language models (architectures modeled after LLaMA, Mistral, and GPT).
Instead of legacy Post-LayerNorm designs, netcl adheres to modern architectural best practices:
- Pre-LN Residuals: Guarantees gradient stability through dozens of stacked transformer layers.
- RMSNorm: Reduces memory bandwidth requirements by 50% compared to classical LayerNorm.
- Rotary Position Embeddings (RoPE): Relative position encoding without learnable absolute position matrices.
- SwiGLU: Gated feed-forward networks delivering higher model capacity per parameter.
- flash_attention: Streaming attention with online Softmax tiling without materializing quadratic $(S \times S)$ attention maps in VRAM.
- KVCache: $O(1)$ key/value state updates for low-latency autoregressive token generation.
[!TIP] Theory & Curriculum Link: For the first-principles mathematical derivation of attention, Q/K/V routing, and FlashAttention tiling, see Curriculum: Tokens & Attention and Curriculum: Hardware & FlashAttention.
1. Quick Example: Building a Modern Decoder Block
A complete implementation of a LLaMA-style decoder block with pre-normalization, Rotary Attention, and SwiGLU:
import numpy as np
import netcl.autograd as ag
from netcl.core.device import manager
from netcl.core.tensor import Tensor
from netcl.nn.modern import DecoderBlock, KVCache
q = manager.default("auto").queue
# 1. Hyperparameters
d_model = 256 # Model hidden dimensionality
n_head = 4 # Number of attention heads (head_dim = 256 // 4 = 64)
batch_size = 2
seq_len = 16
# 2. Instantiate decoder block (bundles RMSNorm, RotaryAttention, and SwiGLU)
block = DecoderBlock(
dim=d_model,
n_head=n_head,
rope_base=10000.0, # Base frequency for rotary embeddings
qk_norm=True, # Prevents attention logit growth at high learning rates
queue=q,
)
# 3. Input batch: (Batch, Sequence_Length, Hidden_Dimension)
x_host = np.random.randn(batch_size, seq_len, d_model).astype(np.float32)
x = Tensor.from_host(q, x_host)
# 4. Differentiable forward pass under autograd
with ag.Tape() as tape:
out = block(ag.tensor(x))
print(f"Decoder output shape: {out.value.shape}") # (2, 16, 256)
The block evaluates:
$$\text{SwiGLU}(x) = (\mathbf{x} \mathbf{W}_{\text{gate}} \odot \text{SiLU}(\mathbf{x} \mathbf{W}_{\text{up}})) \mathbf{W}_{\text{down}}$$
Where RoPE encodes token position $m$ using 2D rotation matrices with base frequencies $\theta_i = 10000^{-2(i-1)/d}$.
2. Fast Autoregressive Generation with KVCache
When generating text, the initial prompt is processed in parallel (Prefill). Subsequent tokens are generated one by one (Decode).
Without caching, generating each new token requires recomputing attention across all previous tokens ($O(S^2)$ cost). With KVCache, the compute cost per generated token drops to $O(1)$:
# 1. Allocate KV-Cache for 1 layer and maximum capacity of 2048 tokens
kv_cache = KVCache(
n_layer=1,
batch=1,
n_head=n_head,
head_dim=64,
capacity=2048,
queue=q,
)
kv_cache.reset()
# 2. Prompt Prefill: process all prompt tokens concurrently
prompt_x = Tensor.from_host(q, np.random.randn(1, 10, d_model).astype(np.float32))
hidden = block(ag.tensor(prompt_x), cache=kv_cache, layer=0)
kv_cache.advance(n=10)
# 3. Autoregressive Generation: feed exactly one token at a time
for step in range(5):
single_token = Tensor.from_host(q, np.random.randn(1, 1, d_model).astype(np.float32))
hidden = block(ag.tensor(single_token), cache=kv_cache, layer=0)
kv_cache.advance(n=1)
print(f"Token {step + 1} generated. Active cache length: {kv_cache.length}")
3. Memory-Efficient Execution with flash_attention
Standard attention materializes an attention map of shape $(B, H, S, S)$:
$$\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q} \mathbf{K}^T}{\sqrt{d_k}} + \mathbf{M}\right) \mathbf{V}$$
For long sequences ($S = 4,096$), this intermediate matrix consumes multiple gigabytes of VRAM per layer. netcl.ops.attention.flash_attention streams across keys and values using GPU local memory without materializing the $(S \times S)$ matrix:
from netcl.ops.attention import flash_attention
# Inputs: (Batch, Heads, Sequence_Length, Head_Dimension)
B, H, S, D = 2, 4, 1024, 64
q_t = Tensor.from_host(q, np.random.randn(B, H, S, D).astype(np.float32))
k_t = Tensor.from_host(q, np.random.randn(B, H, S, D).astype(np.float32))
v_t = Tensor.from_host(q, np.random.randn(B, H, S, D).astype(np.float32))
# Execute streaming causal attention with online Softmax:
out, lse = flash_attention(q_t, k_t, v_t, causal=True)
print(f"FlashAttention output shape: {out.shape}") # (2, 4, 1024, 64)
4. Text Tokenization: BPETokenizer
Language models process integer token IDs rather than raw ASCII strings. netcl.text provides a lossless Byte-Level BPE tokenizer:
from netcl.text import BPETokenizer
# 1. Initialize tokenizer and train vocabulary on a text corpus
tok = BPETokenizer(special_tokens=["<|endoftext|>", "<|im_start|>", "<|im_end|>"])
corpus = "Here is sample training text for the language model running on OpenCL with netcl."
tok.train(corpus, vocab_size=300)
# 2. Encode text into token IDs
tokens = tok.encode("Language model on OpenCL!")
print("Token IDs:", tokens)
# 3. Losslessly decode back into text
reconstructed = tok.decode(tokens)
print("Decoded text:", reconstructed)
# 4. Save and load tokenizer state
tok.save("my_tokenizer.json")
loaded_tok = BPETokenizer.load("my_tokenizer.json")
Practical Engineering Rules
- Even Head Dimensions: Because
RoPErotates coordinate pairs in 2D planes,d_model // num_headsmust always be an even integer (e.g. 32, 64, 128). - Enable QK-Norm: Always set
RotaryAttention(..., qk_norm=True). This normalizes query and key vectors before computing the dot product, preventing attention logits from exploding at high learning rates. - Weight Tying: When vocabulary size is large ($V = 32,000$), use
netcl.nn.modern.TiedLinearto share weights between the token embedding table and the final unembedding projection layer. This eliminates up to 30% of total model parameters.
Related Documentation
- Text API Reference: Tokenizer training, regex splitting, and chat formatters.
- Curriculum: Tokens & Attention: Detailed mathematical derivation of attention.
- Curriculum: Hardware & FlashAttention: Online Softmax and memory tiling mechanics.