ggml_flash_attn_back: Flash Attention Backward (Graph)

In ggmlR: 'GGML' Tensor Operations for Machine Learning

Flash Attention Backward (Graph)

Description

Backward pass for ggml_flash_attn_ext: given the gradient of the attention output, produces the gradients of Q, K and V in one pass.

Usage

ggml_flash_attn_back(ctx, q, k, v, mask, d, scale)

Arguments

ctx	GGML context
q	Query tensor [DK, N, H, B], as passed to the forward pass
k	Key tensor [DK, M, H_kv, B], as passed to the forward pass
v	Value tensor [DV, M, H_kv, B], as passed to the forward pass
mask	Attention mask [M, N, H_m, B_m] of type F16, or NULL
d	Gradient of the forward output, shape [DV, H, N, B]
scale	Softmax scale used in the forward pass (e.g. 1/sqrt(DK))

Details

Upstream ggml ships this operation as a stub that aborts, left behind when flash_attn was replaced by flash_attn_ext; ggmlR implements it so attention can be trained through the fused op. It is also wired into the graph autodiff, so ggml_build_backward_expand reaches it on its own — calling this directly is only needed when building a backward graph by hand.

The three gradients come back packed head to tail in one contiguous 1D tensor: grad_q, then grad_k, then grad_v, each padded to GGML_MEM_ALIGN. Take a view of the slice you need.

Not supported: max_bias (ALiBi), logit_softcap and attention sinks. Q, K, V and d must be F32.

Value

Tensor holding grad_q, grad_k and grad_v packed contiguously

See Also

ggml_flash_attn_ext for the forward pass.

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