pub enum MaskKind {
None,
Causal,
SlidingWindow(usize),
Custom,
Bias,
}Expand description
What kind of attention mask the kernel should apply.
Borrowed from MAX’s nn/attention/mha_mask.mojo pattern (#20 in
PLAN.md): one attention kernel handles all variants by branching on
the mask kind, instead of forcing every caller to materialize a mask
tensor. The win is two-fold:
None— single unpadded sequence: no mask load, no per-key compare in the inner loop.Causal— autoregressive decode: kernel generates the upper- triangular fill from(qi, ki)directly; noseq²mask tensor ever exists.
Custom is the existing path — read mask values from the 4th input.
Variants§
None
No masking — every position attends to every position.
Causal
Causal (autoregressive) — position qi attends only to ki <= qi.
SlidingWindow(usize)
Sliding window — position qi attends to ki ∈ [qi - w, qi].
Custom
Read mask values from the input tensor (default; matches BERT
padding-mask behavior). Tensor shape [batch, key_len] with
1.0 = valid, <0.5 = ignored.
Bias
Additive per-head, per-query bias tensor
[batch, num_heads, query_len, key_len] added to the
QK^T · scale scores before softmax. Lets DETR-style boxRPB
and other learned position biases reuse the fast Op::Attention
path instead of decomposing into matmul + add + softmax + matmul.