For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

apply_mask

def apply_mask[BN: Int, MaskType: MHAMask, //, *, mask_strategy: MaskStrategy, skip_scale: Bool = False](mut srow: Array[Float32, BN], mask: MaskType, scale_log2e: Float32, *, prompt_idx: UInt32, q_head_idx: UInt32, kv_tile_start_row: Int32, max_seq_len: UInt32, num_keys: Int32, score_row: Int32)

Applies bitmask, computed, and out-of-bounds masking strategies to a row of BN attention scores.

Scales by scale_log2e (unless skip_scale), then applies the mask strategy: the bitmask path uses mask_select8 per 32-column batch, the computed path calls mask.mask, and both paths apply the out-of-bounds clip via apply_oob_mask.

Parameters:

BN (Int): Number of scores in the row; must be a multiple of 32 for the bitmask path.
MaskType (MHAMask): MHAMask type providing the apply_log2e_after_mask flag and the mask_bits or mask masking primitives.
mask_strategy (MaskStrategy): MaskStrategy bitset selecting which masking strategies to apply (BITMASK, COMPUTED, OUT_OF_BOUNDS).
skip_scale (Bool): Whether to skip the scale_log2e pre-scaling (defaults to False).

Args:

srow (Array[Float32, BN]): Row of BN attention scores to mask in place.
mask (MaskType): Mask object providing bitmask or computed mask values.
scale_log2e (Float32): Softmax scale factor in log2 base, applied to scores before masking unless skip_scale is set.
prompt_idx (UInt32): Index of the prompt in the batch.
q_head_idx (UInt32): Index of the query head.
kv_tile_start_row (Int32): Starting row of the KV tile in the key dimension, the absolute column offset of the first key in this tile.
max_seq_len (UInt32): Maximum sequence length.
num_keys (Int32): Number of valid keys; columns at or beyond this index are out of bounds.
score_row (Int32): Row index of the score in the query dimension.