layer.MultiHeadAttention — Causal Self-Attention Layer

MultiHeadAttention is the core attention layer implementing scaled dot-product multi-head causal self-attention with optional KV caching. It contains four linear projection layers (Wq, Wk, Wv, Wo) and applies RoPE (Rotary Position Embedding) to both Q and K before computing attention scores.

MultiHeadAttention()

func MultiHeadAttention(embedDim, numOfHeads, headDim int, rope function.RoPEFunc, useCache ...bool) *MultiHeadAttentionT

Constructs and returns a new *MultiHeadAttentionT with four initialized Linear weight matrices.

embedDim

int

required

Input and output embedding dimension E. The Wq/Wk/Wv projections read from this dimension, and the Wo output projection writes back to it.

numOfHeads

int

required

Number of attention heads H. The full hidden dimension H*D is split evenly across heads after projection.

headDim

int

required

Per-head key/query/value dimension D. Typically set to embedDim / numOfHeads. The scaling factor 1/sqrt(D) is applied to the raw attention scores.

rope

function.RoPEFunc

required

The rotary position embedding function. Construct it with function.RoPE(theta, embedDim, contextLen) and share the same instance across all blocks in a model.

useCache

...bool

Optional variadic flag. Pass true to enable KV caching for autoregressive (token-by-token) generation. Defaults to false (no caching; used during training).

The constructor initializes the following internal weight matrices:

Layers: L.Layers{
    "Wq": Linear(E, H*D),  // query projection
    "Wk": Linear(E, H*D),  // key projection
    "Wv": Linear(E, H*D),  // value projection
    "Wo": Linear(H*D, E),  // output projection
}

Example — standard training-mode construction (no cache):

embedDim   := 512
numOfHeads := 8
headDim    := 64   // embedDim / numOfHeads
theta      := 1000.0
contextLen := 10

rope := function.RoPE(theta, embedDim, contextLen)
mha  := layer.MultiHeadAttention(embedDim, numOfHeads, headDim, rope)

Example — generation-mode construction (with KV cache):

mha := layer.MultiHeadAttention(embedDim, numOfHeads, headDim, rope, true)

Forward()

func (l *MultiHeadAttentionT) Forward(x ...*variable.Variable) []*variable.Variable

Runs the full causal self-attention computation. First() is a convenience wrapper that calls Forward and returns only the first (and only) output variable.

Tensor shape flow

Step	Operation	Shape
Input	`x`	`(B, C, E)`
After Wq / Wk / Wv	linear projection	`(B, C, H*D)`
After reshape	split into heads	`(B, C, H, D)`
After transpose	head-first layout	`(B, H, C, D)`
After RoPE	rotary embeddings on Q, K	`(B, H, C, D)`
`QKᵀ / sqrt(D)`	raw attention scores	`(B, H, C, C)`
With KV cache	scores against full cache	`(B, H, C, C+offset)`
After causal mask + softmax	attention weights	`(B, H, C, C)` or `(B, H, C, C+offset)` with cache
After `weights @ V`	attended values	`(B, H, C, D)`
After transpose + reshape	concatenate heads	`(B, C, H*D)`
After Wo	output projection	`(B, C, E)`

Where B = batch size, C = sequence (context) length, E = embedDim, H = numOfHeads, D = headDim.

When caching is active and this is not the first call, the attention scores have shape (B, H, C, C+offset) because K and V are concatenated with previously cached values. The causal mask is only applied on the first call.

Causal masking

The lower-triangular causal mask is applied when useCache is false, or when caching is enabled but this is the very first call (i.e., kCache == nil). The mask fills upper-triangle positions with -inf before softmax, preventing each position from attending to future tokens.

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)

KV Cache behavior

When useCache = true, MultiHeadAttentionT maintains an internal KV cache across successive Forward calls. This is the standard approach for efficient autoregressive text generation.

Call	`kCache` before	Action	`offset` after
1st (`kCache == nil`)	`nil`	initialize cache with K, V; apply causal mask	`C`
2nd+	`(B, H, prev, D)`	concat new K/V onto cache; skip causal mask	`offset + C`

On subsequent calls the full cached K and V are used in the attention computation, so the model attends over the complete context seen so far without reprocessing prior tokens.

Call ClearCache() between independent generation requests. Leaving stale cache entries will corrupt attention for the next sequence.

ClearCache()

func (l *MultiHeadAttentionT) ClearCache()

Resets the KV cache and position offset to their zero values:

kCache → nil
vCache → nil
offset → 0

Must be called before generating a new, independent sequence when useCache = true.

mha.ClearCache()

Model

Tokenizer

Layers

Data & Utilities

layer.MultiHeadAttention — Causal Self-Attention Layer

MultiHeadAttention()

Forward()

Tensor shape flow

Causal masking

KV Cache behavior

ClearCache()

Build docs developers (and LLMs) love

Model

Tokenizer

Layers

Data & Utilities

Documentation Index

​MultiHeadAttention()

​Forward()

​Tensor shape flow

​Causal masking

​KV Cache behavior

​ClearCache()

Build docs developers (and LLMs) love

MultiHeadAttention()

Forward()

Tensor shape flow

Causal masking

KV Cache behavior

ClearCache()