ag_flash_attention: Multi-head attention in a single fused operation

In ggmlR: 'GGML' Tensor Operations for Machine Learning

Multi-head attention in a single fused operation

Description

Computes scaled dot-product attention over all heads at once with ggml_flash_attn_ext(), instead of the per-head loop ag_multihead_attention uses. Both the forward pass and its gradient are one operation each, so the tape holds a single node rather than the dozens a head loop records.

Usage

ag_flash_attention(q, k, v, n_heads, scale = NULL, mask = NULL, causal = FALSE)

Arguments

q, k, v	ag_tensors of shape [d_model, seq]. k and v may have a different sequence length from q (cross-attention), but must match each other.
n_heads	Number of attention heads. Must divide d_model.
scale	Softmax scale. Defaults to 1/sqrt(d_model / n_heads).
mask	Optional attention mask, [seq_kv, seq_q] – keys index the ROWS and queries the columns, which is the transpose of the usual "row = query" convention and is not caught by shape checking when the two lengths match. Either logical (TRUE where attention is allowed) or numeric (0 to attend, -Inf to block). The same mask is used for the gradient.
causal	TRUE builds the causal mask for you: query j attends to keys 1..j. Cannot be combined with mask.

Details

q, k and v are [d_model, seq] matrices whose rows are split into n_heads contiguous blocks – the same head layout ag_multihead_attention() uses. Projections are not included: apply W_q, W_k, W_v before the call and W_o after it.

Value

An ag_tensor of shape [d_model, seq_q].

Examples

d_model <- 16L; seq_len <- 8L
q <- ag_param(matrix(runif(d_model * seq_len, -1, 1), d_model, seq_len))
k <- ag_param(matrix(runif(d_model * seq_len, -1, 1), d_model, seq_len))
v <- ag_param(matrix(runif(d_model * seq_len, -1, 1), d_model, seq_len))
with_grad_tape({
  out  <- ag_flash_attention(q, k, v, n_heads = 4L)
  loss <- ag_mse_loss(out, ag_tensor(matrix(0, d_model, seq_len)))
})
backward(loss)

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