Primitives

attention(Q, K, V; causal = false, kwargs...) -> O

Allocating forward; see attention!. Autodiff entry point (ChainRulesCore/Mooncake rules differentiate attention).

attention!(O, Q, K, V; causal = false, checkpoints = nothing, kwargs...)

Fused multi-head attention with GQA support. For additive bias or other arbitrary variants use flex_attention! (e.g. BiasScore).

• Q: (Dk, SeqLen_Q, Heads, Batch)
• K: (Dk, SeqLen_K, Heads_KV, Batch)
• V: (Dv, SeqLen_K, Heads_KV, Batch)
• O: (Dv, SeqLen_Q, Heads, Batch)

O/Q/K/V may be bare arrays or Shadowed (only the primals are read). Pass checkpoints materialized from checkpoints(space, attention!, Q, K, V) to save the softmax statistics M/L for ∇attention!. Compute precision is set by tensorcore (default from eltype(Q)); accumulation is Float32. Other keywords: causal; input_pos, absolute position of the first query; k_lengths/q_lengths, optional per-batch valid lengths; TILE_M/TILE_N. Returns nothing.

∇attention(Ō, Q, K, V, O; checkpoints, causal, kwargs...) -> (Q̄, K̄, V̄)

Allocating backward of attention!: from the output gradient Ō, the inputs, the output O, and the forward's checkpoints (M/L), return the input gradients. Thin wrapper over ∇attention!.

∇attention!(O::Shadowed, Q::Shadowed, K::Shadowed, V::Shadowed;
            checkpoints, scratch = …, causal, kwargs...)

In-place backward of attention!: reads the output gradient O.shadow and the primals, and writes input gradients into Q.shadow/K.shadow/V.shadow (overwrite). checkpoints must carry M/L from the forward; scratch is the transient backward workspace (scratch(∇attention!, …) descriptions, materialized against a Space/Frame). causal/input_pos must match the forward. Returns nothing.

decode_attention!(O, Q, K, V; lengths, n_splits = 8, TILE_N = 64)

Split-KV (Flash-Decoding) attention for single-token batched decode: one query vector per head and sequence, with the KV cache processed in n_splits parallel slices.

• Q: (Dk, Heads, Batch)
• K: (Dk, SeqLen_K, Heads_KV, Batch)
• V: (Dv, SeqLen_K, Heads_KV, Batch)
• O: (Dv, Heads, Batch)
• lengths: (Batch,), valid KV length per sequence

softmax(X) -> Y

Allocating softmax; see softmax!. Autodiff entry point (ChainRulesCore/Mooncake rules differentiate softmax).

softmax!(Y, X)

Numerically stable softmax over each column of X, in place. Y/X may be bare arrays or Shadowed (only the primals are read). Returns nothing.

∇softmax(Ȳ, Y) -> X̄

Allocating backward of softmax!, from the output gradient Ȳ and the forward output Y. Thin wrapper over ∇softmax!.

∇softmax!(Y::Shadowed, X::Shadowed)

In-place backward of softmax!, computed from the forward output. Reads the output Y.primal and its gradient Y.shadow, and writes the input gradient into X.shadow (overwrite). softmax has no checkpoints or scratch — the output is all the backward needs. Returns nothing.

rms_norm(X, W; eps, offset = 0f0, kwargs...) -> Y

Allocating forward; see rms_norm!. This is the autodiff entry point (ChainRulesCore/Mooncake rules differentiate rms_norm).

rms_norm!(Y, X, W; eps, offset = 0f0, TILE_M = 256, checkpoints = nothing)

RMS-normalize each column of X in place: y = x * rstd * (w + offset) with rstd = 1/√(mean(x²) + eps).

• X, Y: (M, N)
• W: (M,)

Y/X/W may be bare arrays or Shadowed (only the primals are read). Pass checkpoints materialized from checkpoints(space, rms_norm!, X, W) to save rstd into checkpoints.Rstd for ∇rms_norm!; the default nothing saves nothing (inference). Returns nothing.

∇rms_norm(Ȳ, X, W; checkpoints, offset = 0f0, kwargs...) -> (X̄, W̄)

Allocating backward of rms_norm!. checkpoints must carry Rstd from the forward; offset must match. Thin wrapper over ∇rms_norm!.

∇rms_norm!(Y::Shadowed, X::Shadowed, W::Shadowed; checkpoints, scratch = …, offset = 0f0, …)

In-place backward of rms_norm!. Reads the output gradient from Y.shadow and the inputs from X.primal/W.primal, and writes the input gradients into X.shadow/W.shadow (overwrite). checkpoints must carry checkpoints.Rstd from the forward; scratch is the transient backward workspace (scratch(∇rms_norm!, …) descriptions, materialized against a Space/Frame), allocated on the fly if not supplied. Returns nothing.

layer_norm(X, W, B; eps, kwargs...) -> Y

Allocating forward; see layer_norm!. Autodiff entry point (ChainRulesCore/Mooncake rules differentiate layer_norm).

layer_norm!(Y, X, W, B; eps, TILE_M = 256, checkpoints = nothing)

Layer-normalize each column of X in place: y = (x - mean) * rstd * w + b with rstd = 1/√(var + eps).

• X, Y: (M, N)
• W, B: (M,)

Y/X/W/B may be bare arrays or Shadowed (only the primals are read). Pass checkpoints materialized from checkpoints(space, layer_norm!, X, W, B) to save Mean/Rstd for ∇layer_norm!; the default nothing saves nothing (inference). Returns nothing.

∇layer_norm(Ȳ, X, W, B; checkpoints, kwargs...) -> (X̄, W̄, B̄)

Allocating backward of layer_norm!. checkpoints must carry Mean/Rstd from the forward. Thin wrapper over ∇layer_norm!.

∇layer_norm!(Y::Shadowed, X::Shadowed, W::Shadowed, B::Shadowed; checkpoints, scratch = …, …)

In-place backward of layer_norm!. Reads the output gradient from Y.shadow and the inputs from X.primal/W.primal, and writes the input gradients into X.shadow/W.shadow/B.shadow (overwrite). checkpoints must carry Mean/Rstd from the forward; scratch is the transient backward workspace (scratch(∇layer_norm!, …) descriptions, materialized against a Space/Frame), allocated on the fly if not supplied. Returns nothing.

flex_attention(Q, K, V; score_mod = NoOpScore(), mask_mod = FullMask(), kwargs...) -> O

Allocating forward; see flex_attention!. Autodiff through flex_attention yields Q/K/V and score_mod-parameter gradients (mask_mod fixed config). For score-mod param grads without the AD engine, call ∇flex_attention! with a grad_shadow ∂score_mod directly.

flex_attention!(O, Q, K, V; score_mod = NoOpScore(), mask_mod = FullMask(), checkpoints = nothing, kwargs...)

FlexAttention forward: fused multi-head attention whose variant is given by two mods — score_mod rewrites the attention scores and mask_mod decides which query-key pairs attend.

• Q: (Dk, SeqLen_Q, Heads, Batch)
• K: (Dk, SeqLen_K, Heads_KV, Batch)
• V: (Dv, SeqLen_K, Heads_KV, Batch)
• O: (Dv, SeqLen_Q, Heads, Batch)

Heads must be a multiple of Heads_KV (GQA). O/Q/K/V may be bare arrays or Shadowed (only the primals are read). Pass checkpoints materialized from checkpoints(space, flex_attention!, Q, K, V) to save the softmax statistics M/L for ∇flex_attention!.

Keywords: block_mask, a precomputed BlockMask; block_sparse = true, in-kernel block skipping for analytic masks; input_pos = 0; qk_scale = 1/√Dk; tensorcore compute precision (Float32 accumulation); TILE_M/TILE_N. Returns nothing.

∇flex_attention(Ō, Q, K, V, O; checkpoints, score_mod, mask_mod, ∂score_mod = nothing, kwargs...) -> (Q̄, K̄, V̄)

Allocating backward of flex_attention!. Thin wrapper over ∇flex_attention!; score-mod parameter grads go through ∂score_mod.

∇flex_attention!(O::Shadowed, Q::Shadowed, K::Shadowed, V::Shadowed;
                 checkpoints, scratch = …, score_mod, mask_mod, ∂score_mod = nothing, kwargs...)

In-place backward of flex_attention!: reads O.shadow and the primals, writes input gradients into Q.shadow/K.shadow/V.shadow (overwrite). checkpoints must carry M/L; scratch is the transient workspace (scratch(∇flex_attention!, …) descriptions, materialized against a Space/Frame). score_mod/mask_mod must match the forward. Score-mod parameter gradients accumulate into ∂score_mod, a grad_shadow of score_mod (nothing to skip). The precomputed-BlockMask path has no backward. Returns nothing.

FullMask()

No masking — every query attends to every key. The default mask_mod.

CausalMask()

Each query attends to keys at or before its own position (q ≥ kv). Positions are absolute, so input_pos shifts the queries.

SlidingWindowMask(w)

Causal attention within a window of w: q - w ≤ kv ≤ q.

PrefixMask(len)

Keys at positions kv < len are visible to every query — a building block; see prefix_lm.

DocumentMask(doc)
DocumentMask(doc_q, doc_kv)

Attention only within the same document: 0-based position p belongs to document doc[p + 1], and a pair attends iff its ids match. Pass separate vectors when queries and keys come from different sequences.

AndMask(a, b)

Conjunction of two masks — attend where both allow. Construct with a & b.

OrMask(a, b)

Disjunction of two masks — attend where either allows. Construct with a | b.

prefix_lm(len)

Prefix-LM mask: bidirectional attention over the first len positions, causal after. Equals CausalMask() | PrefixMask(len).

NoOpScore()

Identity score mod — leaves scores untouched. The default score_mod.

SoftCapScore(cap)

Logit soft-capping: s -> cap * tanh(s / cap). cap receives no gradient.

AliBiScore(slopes)

ALiBi positional bias: adds slopes[h] * (q - kv) to the scores, one slope per query head.

BiasScore(bias)

Adds bias to the scores, broadcast over heads/batch when those dims are smaller. Indexed by local query position (input_pos does not shift it).

• bias: (SeqLen_K, SeqLen_Q, BiasHeads, BiasBatch)

ComposeScore(a, b)

Apply score mod a, then b. Construct with b ∘ a.

PairFeatureScore(op, q_features, k_features, pair_proj)

Adds a per-head projection of pair_feature outputs to the scores. Features are indexed by local query position (input_pos does not shift them).

• q_features: (F, SeqLen_Q, Batch)
• k_features: (F, SeqLen_K, Batch)
• pair_proj: (Heads, PD)

pair_feature(op, qvals::NTuple{F}, kvals::NTuple{F}) -> NTuple{PD}

Compute PD pair features from the per-position feature values of one (query, key) pair. Implement this for your op type using broadcast dots (.+, .*, exp., …) over the entries — the entries are scalars on the host and tiles on the device, and broadcasting lifts the same definition to both. PD must equal size(pair_proj, 2) of the PairFeatureScore.

NO generic fallback method on purpose: an error("…$(typeof(op))…") fallback drags vararg string MethodInstances into device-code inference, which cuTile's compiler cache cannot handle (lattice error on Vararg argtypes). A missing implementation surfaces as a plain MethodError instead.

∇pair_feature(op, qvals, kvals, dphi::NTuple{PD}) -> (dq, dk)

VJP of pair_feature: given cotangents dphi for the PD pair features, return cotangent tuples for qvals and kvals (length F each). Write it in the same broadcast style as pair_feature; each entry must be a broadcast result involving the inputs (0f0 .* qvals[f] for a zero gradient, not a bare scalar).

Only required with grad_shadow(m; feature_grads = true); projection gradients never need it.

grad_shadow(x)

Build a zeroed gradient accumulator for x — the accumulate counterpart to a Shadowed shadow (which is overwrite, hence uninitialized). The single factory for gradient containers, dispatched on what x is:

• an array → a zeroed copy (fill!(similar(x), 0));
• a score mod → the same struct with every field recursively grad_shadowed (array fields zeroed, composed sub-mods recursed, non-differentiable fields such as BiasScore.nheads left as-is);
• anything else → returned unchanged (a non-differentiable leaf).

Pass a score-mod shadow as ∂score_mod to ∇flex_attention! and read the accumulated gradients from its fields afterwards. Gradients accumulate in place, so a shadow is valid for ONE backward call — rebuild (or re-zero) to reuse.

BlockMask

Precomputed coarse block sparsity: per query block, the KV blocks holding any unmasked element, and which of those are fully unmasked. Built with build_block_mask; pass as block_mask to flex_attention!. For masks with no analytic geometry (e.g. DocumentMask). No backward support.

build_block_mask(keep, n_q, n_kv; TILE_M = 64, TILE_N = 64, input_pos = 0) -> BlockMask

Build a BlockMask on the host from the predicate keep(q, kv) -> Bool (0-based absolute positions). keep must agree with the device mask_mod — e.g. (q, kv) -> hmask(adapt(Array, mod), q, kv) — and TILE_M/TILE_N must match the kernel launch.

hmask(mask_mod, q, kv; b = 1, h = 1) -> Bool

Evaluate a mask mod on the host at 0-based scalar positions — for CPU references and tests. For mods carrying device arrays, evaluate a host copy: hmask(adapt(Array, mod), q, kv).

hscore(score_mod, s, q, kv; b = 1, h = 1)

Evaluate a score mod on the host: the modified score for scalar score s at 0-based positions (q, kv). For mods carrying device arrays, evaluate a host copy: hscore(adapt(Array, mod), s, q, kv).