scaled_dot_product_attention_with_weights_op

→std::vector<TensorImplPtr>

int scaled_dot_product_attention_with_weights_op(const int & q, const int & k, const int & v, const int & attn_mask_or_null, double scale, bool is_causal)

source edit

Run scaled dot-product attention and return {output, weights}.

Identical numerics to scaled_dot_product_attention_op but exposes the post-softmax attention weight matrix $W = \text{softmax}(QK^\top \cdot s + M)$ shaped (..., L_q, L_k) as a second return value. Useful for visualisation, head analysis, or any caller that needs to inspect or further process attention scores.

When autograd is required, two ScaledDotProductAttentionBackward nodes are wired — one for the output tensor (the normal SDPA backward) and a second, independent node attached to the weights tensor so callers may use the weights in downstream loss terms without conflating gradients between the two outputs.

See Also

scaled_dot_product_attention_op : Lighter variant returning only the contracted output, when attention weights are not needed.

Parameters

qTensorImplPtr

Query tensor of shape (..., L_q, d_k).

kTensorImplPtr

Key tensor of shape (..., L_k, d_k).

vTensorImplPtr

Value tensor of shape (..., L_k, d_v).

attn_mask_or_nullTensorImplPtr

Optional additive mask broadcastable to (..., L_q, L_k).

scaledouble

Multiplicative factor applied to

QK^\top

before the softmax.

is_causalbool

If true, an upper-triangular -inf mask is applied engine-side.

Returns

std::vector<TensorImplPtr>

A two-element vector {output, weights} where output has shape (..., L_q, d_v) and weights has shape (..., L_q, L_k).

Notes

On the MLX (GPU) path the fused kernel does not produce a dense weight matrix; this entry point therefore recomputes $\text{softmax}(QK^\top \cdot s + M)$ to materialise the weights for the caller. On the Accelerate (CPU) path the weight matrix is a natural byproduct of the manual softmax and is returned directly.