Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/itsubaki/gpt/llms.txt

Use this file to discover all available pages before exploring further.

Block assembles a complete transformer block with pre-normalization (RMSNorm), multi-head causal self-attention, a residual connection, a second RMSNorm, a SwiGLU feed-forward network, and a final residual connection. It is the core repeating unit stacked to form a full GPT model.

Block()

func Block(embedDim, numOfHeads int, rope function.RoPEFunc, useCache ...bool) *BlockT
Constructs and returns a new *BlockT. The per-head dimension is derived automatically as headDim = embedDim / numOfHeads.
embedDim
int
required
Hidden dimension used by all sub-layers: the RMSNorm scale vectors, all attention projection matrices, and the SwiGLU FFN weight matrices.
numOfHeads
int
required
Number of attention heads passed through to MultiHeadAttention. embedDim must be divisible by numOfHeads.
rope
function.RoPEFunc
required
Shared RoPE function constructed with function.RoPE(theta, embedDim, contextLen). The same instance should be passed to every block in the model.
useCache
...bool
Optional variadic flag forwarded to MultiHeadAttention. Pass true to enable KV caching for autoregressive generation.
The constructor creates the following sub-layers:
Layers: L.Layers{
    "norm1": RMSNorm(embedDim),
    "norm2": RMSNorm(embedDim),
    "attn":  MultiHeadAttention(embedDim, numOfHeads, headDim, rope, useCache...),
    "ffn":   SwiGLU(embedDim),
}
Example:
embedDim   := 64
numOfHeads := 8
theta      := 1000.0
contextLen := 30

rope  := function.RoPE(theta, embedDim, contextLen)
block := layer.Block(embedDim, numOfHeads, rope)

x      := variable.Randn([]int{2, 30, 64}) // (B=2, C=30, E=64)
output := block.First(x)
// output.Shape() → [2 30 64]

Forward() — residual stream

func (l *BlockT) Forward(x ...*variable.Variable) []*variable.Variable
Applies the standard pre-norm transformer block computation with two residual connections. First() is a convenience wrapper that returns only the first output variable.
x0 := l.Layers["norm1"].First(x...)  // pre-norm (RMSNorm)
x1 := l.Layers["attn"].First(x0)     // causal multi-head attention
x2 := F.Add(x[0], x1)                // 1st residual: skip over attention
x3 := l.Layers["norm2"].First(x2)    // pre-norm (RMSNorm)
x4 := l.Layers["ffn"].First(x3)      // SwiGLU feed-forward network
x5 := F.Add(x2, x4)                  // 2nd residual: skip over FFN
The output x5 has the same shape as the input x[0]: (B, C, E).
The residual connections bypass the normalization layers, not just the sub-layer computations. The pre-norm position (applying norm before attention/FFN rather than after) is the convention used in modern LLMs and helps stabilize training at depth.

SwiGLU feed-forward network

func SwiGLU(xDim int) *SwiGLUT
The default FFN used inside each Block. Constructs a gated feed-forward network with three Linear layers and a SiLU (sigmoid-weighted linear unit) gate.
xDim
int
required
Input and output dimension of the FFN, equal to embedDim of the enclosing block.
The hidden dimension is set to xDim * 8 / 3 (following the LLaMA convention for SwiGLU sizing):
hiddenDim := int(xDim * 8 / 3)

Layers: L.Layers{
    "W": Linear(xDim, hiddenDim),  // gate branch
    "V": Linear(xDim, hiddenDim),  // value branch
    "O": Linear(hiddenDim, xDim),  // output projection
}
The forward pass computes a gated activation where the gate is applied element-wise:
a     := l.Layers["W"].First(x...)          // gate: (B, C, hiddenDim)
b     := l.Layers["V"].First(x...)          // value: (B, C, hiddenDim)
gated := F.Mul(F.Mul(a, F.Sigmoid(a)), b)   // SiLU(a) ⊙ b
o     := l.Layers["O"].First(gated)         // output: (B, C, xDim)
This implements the SwiGLU activation: output = (a ⊙ sigmoid(a)) ⊙ b, projected back to xDim by O. The SiLU gate a ⊙ sigmoid(a) replaces the simpler GELU gate used in the original GPT-2 FFN.

FFN (alternative)

func FFN(xDim, hiddenDim int) *FFNT
A simpler two-layer feed-forward network with GELU activation. Not used by Block by default, but available as an alternative FFN.
xDim
int
required
Input and output dimension.
hiddenDim
int
required
Intermediate hidden dimension between the two linear layers. Typically set to 4 * xDim following the original transformer convention.
Layers: L.Layers{
    "l1": Linear(xDim, hiddenDim),
    "l2": Linear(hiddenDim, xDim),
}
Forward pass:
x0 := l.Layers["l1"].First(x...)  // linear up-projection
x1 := F.GELU(x0)                  // GELU non-linearity
x2 := l.Layers["l2"].First(x1)    // linear down-projection
Use FFN when you need compatibility with GPT-2-style architectures. For LLaMA-style models use the default SwiGLU, which is what Block constructs automatically.

ClearCache()

func (l *BlockT) ClearCache()
Delegates to the internal attention layer’s ClearCache(), resetting its KV cache and position offset:
l.Layers["attn"].(*MultiHeadAttentionT).ClearCache()
Call this between independent generation sequences when the block was constructed with useCache = true.

Params()

func (l *BlockT) Params() L.Parameters
Returns all learnable parameters from every sub-layer, keyed by dot-separated names that combine the sub-layer name with the parameter name. For example:
KeyDescription
"norm1.gamma"RMSNorm scale vector before attention
"norm2.gamma"RMSNorm scale vector before FFN
"attn.Wq.w"Query projection weight matrix
"attn.Wk.w"Key projection weight matrix
"attn.Wv.w"Value projection weight matrix
"attn.Wo.w"Output projection weight matrix
"ffn.W.w"SwiGLU gate branch weight
"ffn.V.w"SwiGLU value branch weight
"ffn.O.w"SwiGLU output projection weight
These keys are used when saving and loading model checkpoints.

Build docs developers (and LLMs) love