scaled_dot_product_attention_op

→TensorImplPtr

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

source edit

Python wrappers

pylucid.nn.functional.scaled_dot_product_attentionfunction

Run scaled dot-product attention and return only the output tensor.

Public entry point that wires Lucid's autograd system to the SDPA kernel. Internally calls ScaledDotProductAttentionBackward::forward and, when GradMode::is_enabled() and any of q, k, v require a gradient, attaches a fresh ScaledDotProductAttentionBackward node so the output participates in BPTT.

Math

$\text{Attn}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V$ .

References

Vaswani et al., "Attention Is All You Need" (NeurIPS 2017).

See Also

scaled_dot_product_attention_with_weights_op : Variant that also returns the post-softmax attention weight matrix W.

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); nullptr skips the additive term entirely.

scaledouble

Multiplicative factor applied to

QK^\top

before the softmax.

is_causalbool

If true, an engine-side upper-triangular -inf mask is added so each query position attends only to earlier keys.

Returns

TensorImplPtr

Output tensor of shape (..., L_q, d_v).