For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo module
mha_mask
Mask types and the MHAMask trait for multi-head attention kernels.
Defines the MHAMask trait and concrete implementations including
CausalMask, NullMask, SlidingWindowCausalMask, and MaterializedMask.
Masks encode which query-key pairs are visible and determine per-tile
iteration strategies used by prefill and decode kernels.
comptime values
MASK_VALUE
comptime MASK_VALUE = -10000
Structs
-
AndMask: Mask that's the AND of two masks. If both masks mask off an element, the element is masked off. -
CausalMask: MHA causal mask ensures a token is only affected by previous tokens. -
CausalPaddingMask: Causal mask combined with padding: a position (seq_id, head, q, k) is visible only when q >= k (causal) AND k < valid_lengths[seq_id] (padding). -
ChunkedMask: Mask implementing Chunked attention. -
MaskName: A canonical string name identifying a mask type. -
MaskStrategy: Bit-flag enum that selects the masking strategy for a tile iteration set. -
MaterializedMask: Mask that's backed by a materialized tensor. -
NullMask: Mask that's effectively a noop. -
OrMask: Mask that's the OR of two masks. If either mask masks off an element, the element is masked off. -
SlidingWindowCausalMask: Mask implementing Sliding Window attention. -
SlidingWindowNonCausalMask: Non-causal sliding-window attention mask. -
TileMaskStatus: A tile's masking status.
Traits
-
MHAMask: The MHAMask trait describes masks for MHA kernels, such as the causal mask.
Functions
-
ChunkedCausalMask: Mask implementing Chunked Causal attention for Llama4 models. -
naively_compute_total_iters: Count the non-fully-masked KV tile iterations for a query row by linear scan. -
naively_get_first_nonempty_mask_col: Find the first KV tile column whose mask status is notFULL_MASK.