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

naively_get_first_nonempty_mask_col

def naively_get_first_nonempty_mask_col[MaskType: MHAMask, //, BM: Int, BN: Int](mask: MaskType, seq_id: UInt32, q_row: UInt32) -> UInt32

Find the first KV tile column whose mask status is not FULL_MASK.

Scans KV column tiles starting at 0 with stride BN, calling mask.status() until a tile that is not fully masked is found. Used as a fallback start_column() implementation for mask types that do not have a closed-form expression.

Parameters:

MaskType (MHAMask): Concrete MHAMask implementation to query.
BM (Int): Query tile height (number of query rows per block).
BN (Int): Key tile width (number of key columns per block).

Args:

mask (MaskType): The mask instance.
seq_id (UInt32): Sequence/batch index.
q_row (UInt32): Starting query-row index of the current query tile.

Returns:

UInt32: The column index of the first non-fully-masked KV tile, aligned to a multiple of BN.