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

peel_mask

def peel_mask[num_sets: Int, //, mask_strategies: StaticTuple[MaskStrategy, num_sets], load_fn: def[mask_strategy: MaskStrategy](UInt32) capturing thin -> Float32](mut mask_iters: StaticTuple[UInt32, num_sets], kv_row: UInt32) -> Float32

Determine which mask strategy applies to the peeled first iteration.

Walks through mask sets to find the first with remaining iterations, calls load_fn with the corresponding strategy, and decrements the counter. Prevents UInt32 underflow when early sets are empty (e.g. SlidingWindowCausalMask with num_sets=3 and small sequences).

Parameters:

num_sets (Int): Number of mask sets to walk through (inferred).
mask_strategies (StaticTuple[MaskStrategy, num_sets]): MaskStrategy per set, in evaluation order; the first set with remaining iterations supplies the strategy.
load_fn (def[mask_strategy: MaskStrategy](UInt32) capturing thin -> Float32): Callback invoked as load_fn[strategy](kv_row) to load the mask value for the selected strategy.

Args:

mask_iters (StaticTuple[UInt32, num_sets]): Remaining iteration count per set; the selected set's count is decremented in place.
kv_row (UInt32): KV row index forwarded to load_fn.

Returns:

Float32