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.

The itsubaki/gpt model is a decoder-only transformer that performs a single forward pass: integer token IDs are looked up in an embedding table, passed through a configurable stack of transformer blocks, normalized by a final RMSNorm, and projected by a bias-free linear layer to produce a logit vector over the vocabulary. The highest-probability token at the last position is sampled to generate the next token autoregressively.

Architecture diagram

Token IDs

Embedding

┌─────────────────────────────┐
│ Transformer Block × N       │
│                             │
│ RMSNorm                     │
│   ↓                         │
│ Multi-Head Attention + RoPE │
│   ↓                         │
│ Residual                    │
│   ↓                         │
│ RMSNorm                     │
│   ↓                         │
│ SwiGLU                      │
│   ↓                         │
│ Residual                    │
└─────────────────────────────┘

RMSNorm

Linear

Logits

Major components

ComponentRole
EmbeddingsMaps each token ID to a dense vector of size EmbedDim
Transformer blocksNumOfBlocks stacked pre-norm blocks; each contains causal attention and a SwiGLU FFN
RMSNorm (final)Normalizes the residual stream before projection
Linear (unembed)Projects from EmbedDim back to VocabSize; no bias

Embeddings layer

Embeddings(vocabSize, embedDim) holds a single weight matrix of shape [vocabSize, embedDim] initialized from a normal distribution with σ = 0.02. During the forward pass, token IDs are used as row indices to select embedding vectors, which are then reshaped to (B, C, EmbedDim).

Transformer blocks

NumOfBlocks identical BlockT instances are registered under the keys block[0]block[N-1]. A single RoPEFunc closure — precomputed once in NewGPT — is shared across every block’s attention layer, so the cos/sin tables are allocated only once regardless of depth.

Final RMSNorm

After the last transformer block, the residual stream is normalized by RMSNorm(embedDim). This layer holds a single learnable gamma vector of length EmbedDim (initialized to ones) and applies y = gamma * x / rms with ε = 1e-5.

Linear unembedding

Linear(embedDim, vocabSize) projects the normalized hidden state to logits of shape (B, C, VocabSize). Unlike many implementations this layer carries no bias; the weight matrix is tied to the vocabulary dimension.

The GPT struct

type GPT struct {
    VocabSize     int
    MaxContextLen int
    EmbedDim      int
    NumOfHeads    int
    NumOfBlocks   int
    Theta         float64
    M.Model
}
FieldDescription
VocabSizeNumber of tokens in the vocabulary (e.g. 1000)
MaxContextLenMaximum sequence length; also bounds the RoPE table (e.g. 256)
EmbedDimHidden dimension for embeddings and residual stream (e.g. 192)
NumOfHeadsNumber of attention heads; headDim = EmbedDim / NumOfHeads
NumOfBlocksDepth of the transformer stack (e.g. 6)
ThetaRoPE base frequency, typically 10000

Constructor

func NewGPT(
    vocabSize int,
    maxContextLen int,
    embedDim int,
    numOfHeads int,
    numOfBlocks int,
    theta float64,
    useCache ...bool,
) *GPT
useCache is an optional variadic flag. When true, every attention layer inside each block allocates a KV cache that is populated incrementally during autoregressive generation. Pass false (or omit it) for training.

Forward pass

func (m *GPT) Forward(ids *variable.Variable) *variable.Variable {
    x := m.L["embed"].First(ids)
    for i := range m.NumOfBlocks {
        x = m.L[fmt.Sprintf("block[%d]", i)].First(x)
    }

    x = m.L["norm"].First(x)
    logits := m.L["unembed"].First(x) // (B, C, VocabSize)
    return logits
}
The loop iterates over NumOfBlocks transformer blocks using sequential string keys. The embedding lookup, all block computations, the final norm, and the linear projection are all differentiable operations in the autograd engine, so a single call to logits.Backward() propagates gradients through the entire model.

KV caching

When useCache: true is passed to NewGPT, each MultiHeadAttentionT layer maintains a running cache of key and value tensors. During the first generation step the full prompt is processed normally with a causal mask; subsequent single-token steps append to the cache and skip the mask. Use ClearCache() to reset between independent generation calls:
func (m *GPT) ClearCache() {
    for i := range m.NumOfBlocks {
        m.L[fmt.Sprintf("block[%d]", i)].(*L.BlockT).ClearCache()
    }
}
Call ClearCache() before starting each new generation sequence. Forgetting to reset leaves stale keys and values from a previous context in the cache.

Transformer Block

Pre-norm, causal attention, residual connections, and SwiGLU FFN inside each block.

Attention

Multi-head causal self-attention with RoPE and optional KV cache.

Positional Encoding

Rotary Position Embeddings (RoPE): precomputed cos/sin tables and backward pass.

Quickstart

Train and run text generation end-to-end.

Build docs developers (and LLMs) love