scaled_dot_product_attention

→Tensor

scaled_dot_product_attention(query: Tensor, key: Tensor, value: Tensor, attn_mask: Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False)

source edit

Implementing kernel

Scaled dot-product attention — the core of every Transformer block.

Computes the attention-weighted aggregation of value vectors using query-key dot products as similarity scores:

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}} + M\right) V

where $M$ is an optional additive mask (-inf to disallow attention at a position, 0 to allow). The $1/\sqrt{d_k}$ factor keeps the softmax in a usable temperature regime as the head dimension $d_k$ grows — without it, large dot products saturate the softmax into a near one-hot distribution and gradients vanish.

Parameters

queryTensor

Shape (B, H, T, E). B batch, H heads, T query positions, E head dimension

d_k

keyTensor

Shape (B, H, S, E).

valueTensor

Shape (B, H, S, V). V may differ from E.

attn_maskTensor= None

Mask broadcast-compatible with (B, H, T, S), in either of two forms. Additive (any float dtype): the values are added to the scores, so use large negatives (or -inf) at positions to mask out. Boolean: a keep-mask — True attends, False is masked out — converted internally to the additive form it stands for, so the two spellings are exactly equivalent on every path, fused or not. Combined with is_causal, both apply. A query row whose every key is masked out (False or -inf throughout) attends to nothing: its output row is 0 and no gradient flows through it, on every device and dtype. A large finite negative is not a mask-out — such a row still normalises, to the mean of the values.

dropout_pfloat= 0.0

Dropout probability applied to attention weights during training. Default 0.0.

is_causalbool= False

If True, apply an upper-triangular causal mask so each query position only attends to keys at the same or earlier positions (autoregressive decoder self-attention). For a non-square score matrix the triangle is aligned top-left — query i sees keys 0..i — so decoding against a KV cache wants an explicit mask.

scalefloat= None

Override the default

1/\sqrt{d_k}

scale factor.

enable_gqabool= False

Grouped-query attention: key and value may have fewer heads than query, as long as theirs divides its; each key / value head then serves that many consecutive query heads. Default False.

Returns

Tensor

Attention output of shape (B, H, T, V). The leading axes of query, key and value broadcast against each other, so a key of batch 1 serves every query in the batch.

Notes

Introduced in Attention Is All You Need (Vaswani et al., 2017). The implementation uses the log-sum-exp form of softmax for numerical stability under aggressive masking, and fuses the scale into the score matrix prior to softmax. Causal masking enables efficient autoregressive decoding when combined with a key/value cache.

Examples

>>> import lucid
>>> from lucid.nn.functional import scaled_dot_product_attention
>>> q = lucid.randn(2, 8, 16, 64)          # (B, H, T, E)
>>> k = lucid.randn(2, 8, 16, 64)
>>> v = lucid.randn(2, 8, 16, 64)
>>> out = scaled_dot_product_attention(q, k, v, is_causal=True)
>>> out.shape
(2, 8, 16, 64)

Used by 2

scaled_dot_product_attention(query: Tensor, key: Tensor, value: Tensor, attn_mask: Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False)

>>> import lucid >>> from lucid.nn.functional import scaled_dot_product_attention >>> q = lucid.randn(2, 8, 16, 64) # (B, H, T, E) >>> k = lucid.randn(2, 8, 16, 64) >>> v = lucid.randn(2, 8, 16, 64) >>> out = scaled_dot_product_attention(q, k, v, is_causal=True) >>> out.shape (2, 8, 16, 64)