Neural Network primitives from NNlib.jl

Flux re-exports all of the functions exported by the NNlib package. This includes activation functions, described on their own page. Many of the functions on this page exist primarily as the internal implementation of Flux layer, but can also be used independently.

Attention

Primitives for the Attention layer.

NNlib.scaled_dot_product_attention — Function

scaled_dot_product_attention(query, key, value, [bias]; [fdrop, mask, scale, is_causal, q_offset, q_lengths, kv_lengths, q_segments, kv_segments, maxlen])
sdpa(query, key, value, [bias]; kws...)

Multihead dot product attention used in transformer architectures, with an explicit head axis.

The input arrays must have the shape (head_dim, nheads, seq_len, batch_size...), that is a feature (head) dimension, a heads dimension, a sequence-length dimension, then an arbitrary number of batch dimensions or none. The number of heads is taken from the second dimension of each input, so no nheads keyword is needed.

This is the reverse of the (batch_size, seq_len, nheads, head_dim) layout of JAX and Flax. It is not the reverse of PyTorch's (batch_size, nheads, seq_len, head_dim): an array ported from PyTorch by reversing its dimensions has size (head_dim, seq_len, nheads, batch_size) and still needs permutedims(x, (1, 3, 2, 4)). A missing permutation cannot be detected, since the array is a valid input whose heads are taken as the sequence.

Grouped-query attention (GQA) is supported: if key and value have fewer heads than query, they are shared across query heads. The number of query heads must then be divisible by the number of key/value heads.

Returns the attention output array of size (v_head_dim, nheads, q_len, batch_size...).

See also scaled_dot_product_attention_with_weights if you need the attention weights too.

Arguments

query: Query array of size (qk_head_dim, nheads, q_len, batch_size...).
key: Key array of size (qk_head_dim, nkvheads, kv_len, batch_size...).
value: Value array of size (v_head_dim, nkvheads, kv_len, batch_size...).
bias: Either nothing or an array broadcastable to size (kv_len, q_len, nheads, batch_size...). It will be added to the attention logits before applying the softmax. Default nothing.
fdrop: A dropout function or layer to be applied on the attention weights right after the softmax. Default identity (no dropout).
mask: Either nothing or a boolean array broadcastable to size (kv_len, q_len, nheads, batch_size...). The mask is applied to the attention logits just before the softmax. See make_causal_mask for creating causal masks. Default nothing.
scale: The multiplier used to scale the queryᵀ * key dot products before the softmax. Default nothing, meaning 1 / √(qk_head_dim).
is_causal: If true, a causal mask is applied so that each query position can only attend to key positions up to and including its own. Equivalent to mask = make_causal_mask(query; dims=3) for equal query and key lengths, but a fused implementation does not have to materialize the mask. If :bottom_right, the mask is aligned with the last valid key instead of the first: the queries are taken as the last positions of the keys, as when decoding with a KV cache, and q_offset is kv_lengths .- q_lengths, these being the sizes of the arrays where not given. This is the causal mask of FlashAttention and PyTorch's causal_lower_right. Default false.
q_offset: With is_causal=true, the number of keys that precede the first query: the query i attends to the keys j ≤ q_offset + i. An integer, or an integer array with one entry for each batch element. It counts the keys given, which with a sliding-window cache is not the absolute position. It has no effect without is_causal, and is set by is_causal=:bottom_right. Default 0.
q_lengths, kv_lengths: Either nothing or the number of valid positions of each query and of each key/value sequence in the batch, the following ones being padding: an integer array with one entry for each batch element, or an integer for all of them. The keys past kv_lengths are not attended to, and the output is zero at the queries past q_lengths. Each can be given without the other. Default nothing.
q_segments, kv_segments: Either nothing or the lengths of the sequences packed one after the other along the sequence dimension of query and of key and value, as an integer vector. See below. kv_segments defaults to q_segments.
maxlen: With q_segments, either nothing or an integer at least as large as every segment. It saves reading the lengths of the segments on the host where they are on a GPU, which waits for the device, and stands in for them where they cannot be read, as the traced ones of Reactant. Default nothing.

mask, is_causal and kv_lengths can be combined: a key is then attended to only where all of them allow it.

Packed sequences

With q_segments, the inputs hold sequences of different lengths one after the other, without padding: they have size (head_dim, nheads, total_len), without batch dimensions, and the output has the size of query. The s-th segment of query attends to the s-th segment of key and value only, as a batch element does in a padded batch: is_causal applies within each segment, and q_offset, q_lengths and kv_lengths take one entry for each segment. The lengths of the segments have to add up to the sequence lengths of the arrays. A bias or a mask is not supported.

This is how documents packed in a sequence are kept from attending to each other, and how a batch of sequences of different lengths avoids padding: layers other than the attention act on the (features, total_len) array of all the positions.

Fused implementation

Without an fdrop, the output is computed tile by tile, with the softmax taken online over the tiles of keys: the attention weights are never materialized, so that memory grows linearly rather than quadratically with the sequence length. This is what Arrays and GPU arrays of Float16, Float32 and Float64 take, and on the GPU those of BFloat16, with any of the arguments above. The bias and the mask, if any, have to be on the device of the inputs.

Any other call takes the materialized implementation, which forms the weights. NNlib.is_fused_attention tells which of the two a call takes.

The fused implementation has a gradient with respect to query, key, value and scale, which recomputes the weights tile by tile and does not materialize them either. It is what automatic differentiation based on ChainRules (as that of Zygote) and the reverse mode of Enzyme use. Under them a call with a bias takes the materialized implementation, which gives the gradient with respect to the bias too, and so does every call under the other automatic differentiation packages. A query that attends to no key, whose output is NaN, takes no part in the fused gradient.

With CUDA.jl and cuDNN.jl loaded, the Float16 and BFloat16 CuArrays that cuDNN's own fused kernel has an engine for are handed to it: it runs on tensor cores, and has a gradient. It supports scale, is_causal, q_lengths, kv_lengths and grouped-query attention, and requires:

no bias, fdrop or mask (that is nothing, identity and nothing);
with is_causal, a causal mask aligned with the first key: the default q_offset of zero, and for :bottom_right as many queries as keys, without q_lengths or kv_lengths;
a head dimension that is a multiple of 8 and at most 256;
a value with the same head dimension as query.

Examples

# (head_dim, nheads, seq_len, batch)
q = rand(Float32, 16, 8, 20, 4)
k = v = rand(Float32, 16, 8, 30, 4)
y = scaled_dot_product_attention(q, k, v)
# size(y) == (16, 8, 20, 4)

# grouped-query attention: 8 query heads sharing 2 key/value heads
q = rand(Float32, 16, 8, 20, 4)
k = v = rand(Float32, 16, 2, 30, 4)
y = scaled_dot_product_attention(q, k, v)

# a padded batch, whose first key/value sequence has 12 valid positions out of 30
y = scaled_dot_product_attention(q, k, v; kv_lengths=[12, 30, 30, 30])

# decoding with a KV cache of 64 positions: 3 new queries for each batch element,
# appended to the 10, 20, 30 and 40 keys already cached
q = rand(Float32, 16, 8, 3, 4)
k = v = rand(Float32, 16, 8, 64, 4)
y = scaled_dot_product_attention(q, k, v; is_causal=:bottom_right, kv_lengths=[13, 23, 33, 43])
# the same, with the offsets of the queries
y = scaled_dot_product_attention(q, k, v; is_causal=true, q_offset=[10, 20, 30, 40],
                                 kv_lengths=[13, 23, 33, 43])

# three documents of 5, 12 and 3 positions packed in one sequence, each attending
# causally to itself only
x = rand(Float32, 16, 8, 20)
y = scaled_dot_product_attention(x, x, x; is_causal=true, q_segments=[5, 12, 3])
# size(y) == (16, 8, 20)