Python wrappers
Public Python APIs implemented by this engine symbol.Autograd node for scaled dot-product attention with optional additive mask and post-softmax causal constraint.
Implements the canonical SDPA operation from Vaswani et al. (2017),
,
where are query / key / value projections shaped
(..., L, D) and is an additive mask (-inf at disallowed
positions, 0 elsewhere). The node saves three edges
{Q, K, V} and a fourth Storage slot saved_weights_ holding the
pre-multiplication attention weight matrix
, where s = scale_.
On the MLX (GPU) path the fused mlx::core::fast::scaled_dot_product_attention
kernel does not return — instead it returns a {1}-shape
placeholder. apply() detects saved_weights_.shape().size() == 1
and re-computes from the saved
before back-propagating. On the Accelerate (CPU) path the
kernel returns the genuine weight matrix and the recompute branch is
skipped.
Math
Gradients flow back through the softmax Jacobian (per row), giving
where .
References
Vaswani et al., "Attention Is All You Need" (NeurIPS 2017).
Attributes
schema_v1OpSchema"scaled_dot_product_attention", version 1, AmpPolicy::ForceFP32, autograd-aware). AMP forces the kernel to FP32 because the softmax-exp accumulation is numerically unstable at half precision.scale_doubles (typically ) multiplied into before the softmax. Default: 1.0.orig_q_shape_ShapeQ recorded at forward; used by apply() to restore leading batch dims on the gradient tensor.orig_k_shape_ShapeK.orig_v_shape_ShapeV.saved_weights_StorageW. Genuine on the CPU path; on the GPU path this is a {1}-shape placeholder and the backward recomputes W from the saved Q, K and scale_.Notes
The mask is additive: callers must already have converted boolean
"where to mask" patterns into floats containing -inf where the
position should be excluded and 0 elsewhere (see
lucid.nn.modules.attention._to_additive_mask). Causal masking
is implemented engine-side when is_causal=true by adding an
upper-triangular -inf matrix before the softmax, so the caller
does not need to materialise the mask in Python.
Static methods
1forward
→TensorImplPtrint forward(const int & q, const int & k, const int & v, const int & attn_mask_or_null, double scale, bool is_causal)Run the SDPA forward and return the output tensor.
Dispatches to IBackend::sdpa_forward after validating that the
leading batch dims of q, k and v agree, that
q.shape[-1] == k.shape[-1] (the dot-product key dim d_k),
and that k.shape[-2] == v.shape[-2] (the source-sequence
length L_k). On the autograd-recording path this also
attaches a new ScaledDotProductAttentionBackward node with
saved_weights_ populated.
Parameters
qTensorImplPtr(..., L_q, d_k) with at least 2 dimensions. All leading "batch" dims must broadcast with the corresponding dims of k and v.kTensorImplPtr(..., L_k, d_k); the last dim must equal q.shape[-1].vTensorImplPtr(..., L_k, d_v); the penultimate dim must equal k.shape[-2].attn_mask_or_nullTensorImplPtr(..., L_q, L_k). Pass nullptr to omit (no additive term is materialised).scaledoubleQ @ K^T before the softmax. Typically .is_causalbooltrue, an upper-triangular -inf mask is added engine-side, restricting position i to attend only to positions j <= i.Returns
TensorImplPtrAttention output of shape (..., L_q, d_v).
Raises
ErrorBuilderMethods
1Compute gradients with respect to Q, K and V from the upstream gradient of the output.
Reconstructs the saved weight matrix on the GPU path (where
saved_weights_ is a {1}-shape placeholder) by recomputing
, then dispatches to
IBackend::sdpa_backward for the three input gradients.
Parameters
grad_outStorage(..., L_q, d_v).Returns
std::vector<Storage>Three Storage objects in edge order {dQ, dK, dV}, each reshaped back to orig_q_shape_ / orig_k_shape_ / orig_v_shape_.