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.

Multi-head attention in itsubaki/gpt projects the input hidden state into query, key, and value tensors through independent linear maps, reshapes them to expose the head dimension, applies Rotary Position Embeddings (RoPE) to queries and keys, computes scaled dot-product attention scores with a causal lower-triangular mask, and projects the concatenated head outputs back to the embedding dimension. An optional KV cache enables efficient single-token generation by reusing previously computed keys and values.

​
Constructor

func MultiHeadAttention(
    embedDim, numOfHeads, headDim int,
    rope function.RoPEFunc,
    useCache ...bool,
) *MultiHeadAttentionT {
    E, H, D := embedDim, numOfHeads, headDim
    return &MultiHeadAttentionT{
        numOfHeads: numOfHeads,
        headDim:    headDim,
        rope:       rope,
        useCache:   len(useCache) > 0 && useCache[0],
        Layers: L.Layers{
            "Wq": Linear(E, H*D),
            "Wk": Linear(E, H*D),
            "Wv": Linear(E, H*D),
            "Wo": Linear(H*D, E),
        },
    }
}
Each weight matrix is initialized from Normal(0, 0.02):
MatrixShapeRole
Wq[E, H×D]Projects input to queries
Wk[E, H×D]Projects input to keys
Wv[E, H×D]Projects input to values
Wo[H×D, E]Projects concatenated head outputs back to embedding dim

​
Full forward pass

func (l *MultiHeadAttentionT) Forward(x ...*variable.Variable) []*variable.Variable {
    v, shape := x[0], x[0].Shape()
    B, C, H, D := shape[0], shape[1], l.numOfHeads, l.headDim

    // 1. Project Q, K, V — shape (B, C, H*D)
    Q := l.Layers["Wq"].First(v)
    K := l.Layers["Wk"].First(v)
    V := l.Layers["Wv"].First(v)

    // 2. Reshape and transpose to (B, H, C, D)
    Q = F.Transpose(0, 2, 1, 3)(F.Reshape(B, C, H, D)(Q))
    K = F.Transpose(0, 2, 1, 3)(F.Reshape(B, C, H, D)(K))
    V = F.Transpose(0, 2, 1, 3)(F.Reshape(B, C, H, D)(V))

    // 3. Apply RoPE to Q and K
    Q = l.rope(l.offset)(Q)
    K = l.rope(l.offset)(K)

    // 4. KV cache logic
    isFirstCall := l.kCache == nil
    if l.useCache {
        if isFirstCall {
            l.kCache = K
            l.vCache = V
        } else {
            l.kCache = F.Concat(2)(l.kCache, K) // (B, H, C+cache, D)
            l.vCache = F.Concat(2)(l.vCache, V) // (B, H, C+cache, D)
        }
        K = l.kCache
        V = l.vCache
        l.offset += C
    }

    // 5. Scaled dot-product: QK^T / sqrt(D)
    Kt     := F.Transpose(0, 1, 3, 2)(K)                    // (B, H, D, C)
    scores := F.MatMul(Q, Kt)                               // (B, H, C, D) @ (B, H, D, C) -> (B, H, C, C)
    scores  = F.MulC(1.0/math.Sqrt(float64(D)), scores)    // scale

    // 6. Causal mask — applied only on the first cache call or when caching is off
    if !l.useCache || isFirstCall {
        mask := tensor.Tril(tensor.Ones[float64](C, C))
        cond := func(m float64) bool { return m == 0 }
        scores = F.MaskFill(mask, cond, math.Inf(-1))(scores)
    }

    // 7. Softmax → weights → V → reshape → Wo
    weights := F.Softmax(-1)(scores)          // (B, H, C, C)
    hidden  := F.MatMul(weights, V)           // (B, H, C, C) @ (B, H, C, D) -> (B, H, C, D)
    hidden   = F.Transpose(0, 2, 1, 3)(hidden) // (B, C, H, D)
    hidden   = F.Reshape(B, C, H*D)(hidden)   // (B, C, H*D)
    output  := l.Layers["Wo"].First(hidden)   // (B, C, E)

    return []*variable.Variable{output}
}

​
Step-by-step breakdown

Step 1 — Project Q, K, V. Each of the three independent linear layers maps the (B, C, E) input to (B, C, H×D). Step 2 — Reshape and transpose. The H×D dimension is split into (H, D) then transposed to (B, H, C, D) so that each head’s sequence is contiguous in memory for the subsequent matmul. Step 3 — Apply RoPE. The same RoPEFunc closure rotates pairs of head-dimension coordinates in Q and K using position-dependent angles. Keys and queries receive identical rotation at the same positions, so relative position information is encoded in their dot product. Step 4 — KV cache. See the KV caching section below. Step 5 — Scaled dot-product. K is transposed from (B, H, C, D) to (B, H, D, C) and multiplied with Q of shape (B, H, C, D) to produce attention scores (B, H, C, C). Dividing by sqrt(D) prevents the dot products from growing too large in high-dimensional spaces, which would push the softmax into regions with vanishing gradients. Step 6 — Causal mask. A lower-triangular mask of ones is built with tensor.Tril. Positions where the mask value is 0 (the upper triangle) are filled with math.Inf(-1) before the softmax, ensuring each token can only attend to itself and earlier tokens. Step 7 — Output projection. After softmax, the weight matrix is multiplied with V, the heads are concatenated by transposing back and reshaping to (B, C, H×D), and finally Wo maps back to (B, C, E).

​
Causal mask

mask := tensor.Tril(tensor.Ones[float64](C, C))
cond := func(m float64) bool { return m == 0 }
scores = F.MaskFill(mask, cond, math.Inf(-1))(scores)
tensor.Tril returns a C×C lower-triangular matrix with ones on and below the main diagonal. F.MaskFill replaces every score at a position where the mask is zero with negative infinity. After softmax those positions become exactly zero, which prevents any information flow from future tokens to the current token — enforcing the autoregressive constraint during both training and generation.

​
KV caching

The offset counter and the kCache / vCache fields enable incremental generation: 
type MultiHeadAttentionT struct {
    numOfHeads int
    headDim    int
    rope       function.RoPEFunc
    offset     int
    useCache   bool
    kCache     *variable.Variable
    vCache     *variable.Variable
    L.Layers
}
StateBehaviour
kCache == nil (first call)Full prompt is processed; causal mask is applied; cache is initialized to K and V
kCache != nil (subsequent calls)Single new token; new K and V are appended along the sequence axis; causal mask is skipped because the new token can attend to all cached positions
offsetTracks the total number of tokens already processed; passed to RoPEFunc so new tokens receive the correct rotational angle
The concatenation is along axis 2 (the context/sequence dimension): 
l.kCache = F.Concat(2)(l.kCache, K) // (B, H, C+cache, D)
l.vCache = F.Concat(2)(l.vCache, V) // (B, H, C+cache, D)

​
Clearing the cache

func (l *MultiHeadAttentionT) ClearCache() {
    l.kCache = nil
    l.vCache = nil
    l.offset = 0
}
Setting both cache pointers to nil and resetting the offset counter returns the layer to its initial state, making isFirstCall true on the next forward pass.
The KV cache accumulates computation graph nodes when autograd tracking is active. Use the cache only during inference / generation — never during training — or you risk memory growth from unexpectedly retained backward graphs. Always call ClearCache() before starting a new generation sequence.

Build docs developers (and LLMs) love

Get started for freeTalk to us