ag_multihead_attention: Create a Multi-Head Attention layer

In ggmlR: 'GGML' Tensor Operations for Machine Learning

Create a Multi-Head Attention layer

Description

Implements scaled dot-product multi-head attention as in "Attention Is All You Need" (Vaswani et al., 2017).

Usage

ag_multihead_attention(d_model, n_heads, dropout = 0, bias = TRUE)

Arguments

d_model	Model (embedding) dimension
n_heads	Number of attention heads. d_model must be divisible by n_heads.
dropout	Attention dropout probability (default 0, applied in training mode only)
bias	Logical: add bias to output projection (default TRUE)

Details

Calling convention (mirrors PyTorch nn.MultiheadAttention):

• layer$forward(q) — self-attention (k = v = q)

• layer$forward(q, k, v) — cross-attention

Tensor layout: [d_model, seq_len] — columns are tokens, consistent with the rest of the ag_* API.

Forward pass:

  Q = W_q %*% q            [d_k * n_heads, seq_len]
  K = W_k %*% k            [d_k * n_heads, seq_len]
  V = W_v %*% v            [d_v * n_heads, seq_len]

  for each head h:
    q_h = Q[h*d_k+1 : (h+1)*d_k, ]    [d_k, seq_len]
    k_h = K[h*d_k+1 : (h+1)*d_k, ]    [d_k, seq_len]
    v_h = V[h*d_v+1 : (h+1)*d_v, ]    [d_v, seq_len]
    A_h = softmax(t(q_h) %*% k_h / sqrt(d_k))   [seq_len, seq_len]
    if causal_mask: A_h[i,j] = 0 for j > i
    head_h = v_h %*% A_h             [d_v, seq_len]

  concat = rbind(head_1, ..., head_H)  [d_v*n_heads, seq_len]
  out    = W_o %*% concat + b_o      [d_model, seq_len]

Value

An ag_multihead_attention environment with $forward(q, k, v, causal_mask) and $parameters()

Choosing between this and ag_flash_attention()

This layer loops over heads in R. Since ag_* has no slice operation, extracting one head is a matmul against a selector matrix, so a 4-head block records around 52 tape nodes for the attention core – against one for ag_flash_attention, which fuses all heads into a single ggml_flash_attn_ext() call and its gradient into one ggml_flash_attn_back(). Measured on the attention core alone, the fused op is 15-37x faster depending on shape (inst/scripts/measure_ag_flash_attn.R).

Prefer ag_flash_attention() for plain and causal multi-head attention. This layer remains the fallback for what the fused op cannot do:

• attention dropout (dropout > 0);

• packaged projections – it owns W_q, W_k, W_v, W_o and b_o, whereas ag_flash_attention() takes already-projected Q/K/V and leaves the output projection to the caller;

• shapes or options ggml_flash_attn_back() rejects – ALiBi, logit softcap and attention sinks all assert.

Which API this belongs to

This is an ag_* layer: it runs on the dynamic autograd path, where a tape records each operation as it executes and backward differentiates through R closures. It is not a layer for ggml_model_sequential or ggml_model and cannot be passed to them; the two APIs do not share layer objects. For attention inside a compiled model, build the block from nn_* layers instead (see the transformer encoder section of the README).

The practical consequences of being on this path: gradients and optimizer state are R doubles regardless of ag_dtype, so budget 8 bytes per scalar (ag_estimate_training_memory); the tape holds every activation until zero_grad() (ag_tape_memory); and ag_device("gpu") is usually slower than the CPU for individual ops, because each one pays a separate upload and download.

See Also

ag_flash_attention for the fused implementation.

Examples

# Self-attention
mha <- ag_multihead_attention(64L, 8L)
x   <- ag_tensor(matrix(rnorm(64 * 10), 64, 10))  # [d_model=64, seq_len=10]
out <- mha$forward(x)                              # [64, 10]

# Cross-attention
q   <- ag_tensor(matrix(rnorm(64 * 10), 64, 10))
kv  <- ag_tensor(matrix(rnorm(64 * 15), 64, 15))
out <- mha$forward(q, kv, kv)

# Causal (GPT-style)
out <- mha$forward(x, causal_mask = TRUE)

ggmlR documentation built on Oct. 23, 2026, 1:07 a.m.