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 struct

MaskTileOp

struct MaskTileOp[accum_type: DType, token_gen: Bool, mma_shape: IndexList[Int(3)], num_m_mmas: Int, num_n_mmas: Int, mask_t: MHAMask, group: Int, mma_m: Int, use_exp2: Bool = False, num_heads_per_token: Int = group, valid_rows: Int = group, mla_mode: Bool = False, num_warps_m: Int = Int(1)]

Apply an MHAMask to the per-lane QK score registers in place.

Parameters

accum_type (DType): The element type of the QK score registers and scale.
token_gen (Bool): Whether the kernel runs in token-generation (decode) mode versus prefill.
mma_shape (IndexList[Int(3)]): The MMA tile shape (M, N, K) used to compute per-fragment row and column offsets.
num_m_mmas (Int): Number of MMAs along the M (query-row) dimension of the score tile.
num_n_mmas (Int): Number of MMAs along the N (key-column) dimension of the score tile.
mask_t (MHAMask): The MHAMask type to apply to the score registers.
group (Int): GQA group size; the number of query heads sharing one KV head.
mma_m (Int): MMA M extent (32 or 16 on gfx950) that selects the lane and fragment geometry.
use_exp2 (Bool): Whether the consumer softmax uses base-2 exponentiation (defaults to False).
num_heads_per_token (Int): Token-fold geometry; H heads per token mapping a folded row to its token and head, where token = row // H and head = row % H (defaults to group).
valid_rows (Int): Number of live query rows M = H * S across all heads and sequence positions (defaults to group).
mla_mode (Bool): Whether the token fold applies, for AMD MLA decode only; when False the fold arithmetic is comptime-dead (defaults to False).
num_warps_m (Int): Number of warps splitting the MMA M dimension (= BM // WM), which controls the dead-row guard strategy (defaults to 1).

Implemented traits

AnyType, ImplicitlyDeletable

`comptime` members

`FragmentLayoutT`

comptime FragmentLayoutT = Layout[*?, *?]

`output_frag_size`

comptime output_frag_size = MaskTileOp[accum_type, token_gen, mma_shape, num_m_mmas, num_n_mmas, mask_t, group, mma_m, use_exp2, num_heads_per_token, valid_rows, mla_mode, num_warps_m].FragmentLayoutT.static_product

`WarpLayoutT`

comptime WarpLayoutT = Layout[*?, *?]

Methods

`apply`

static def apply(masked: Bool, kv_tile_start_row: UInt32, kv_tile_num_rows: UInt32, start_pos: UInt32, seq_len: UInt32, num_keys: UInt32, mask_block_row: UInt32, mask_warp_row: UInt32, mask_warp_col: UInt32, scale: Scalar[accum_type], mask: mask_t, p_reg_tile: TileTensor[accum_type, Storage=p_reg_tile.Storage, address_space=p_reg_tile.address_space, linear_idx_type=p_reg_tile.linear_idx_type], not_last_iter: Bool, cache_start_pos: UInt32 = UInt32(0))

Parameters​

Implemented traits​

comptime members​

FragmentLayoutT​

output_frag_size​

WarpLayoutT​

Methods​

apply​