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_oob_mask

def apply_oob_mask[*, mask_strategy: MaskStrategy, apply_log2e_after_mask: Bool](s_arg: SIMD[DType.float32, SIMDLength(2)], *, prompt_idx: UInt32, q_head_idx: UInt32, kv_tile_start_row: Int32, max_seq_len: UInt32, num_keys: Int32, score_row: Int32, score_col: Int32) -> SIMD[DType.float32, SIMDLength(2)]

Applies the out-of-bounds key mask to a pair of attention scores.

Scores for columns at or beyond num_keys are replaced with MASK_VALUE; optionally scales by log2e before masking.

Parameters:

mask_strategy (MaskStrategy): MaskStrategy bitset selecting which masking strategies to apply; the out-of-bounds clip runs only when OUT_OF_BOUNDS is set.
apply_log2e_after_mask (Bool): Whether to multiply the scores by log2e before masking.

Args:

s_arg (SIMD[DType.float32, SIMDLength(2)]): Pair of attention scores to mask.
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.
max_seq_len (UInt32): Maximum sequence length.
num_keys (Int32): Number of valid keys; columns at or beyond this index are masked with MASK_VALUE.
score_row (Int32): Row index of the score in the query dimension.
score_col (Int32): Starting column index of the score pair; columns from this index onward are compared to num_keys.

Returns:

SIMD[DType.float32, SIMDLength(2)]