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.
Implemented traits
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, element_size=p_reg_tile.element_size], not_last_iter: Bool, cache_start_pos: UInt32 = UInt32(0))
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!