Mojo struct
MaskTileOp
struct MaskTileOp[accum_type: DType, token_gen: Bool, mma_shape: IndexList[3], num_m_mmas: Int, num_n_mmas: Int, mask_t: MHAMask, group: Int, mma_m: Int, use_exp2: Bool = False]
Apply an MHAMask to the per-lane QK score registers in place.
Implemented traits
AnyType,
ImplicitlyDestructible
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].FragmentLayoutT.static_product
WarpLayoutT
comptime WarpLayoutT = Layout[*?, *?]
Methods
apply
static 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, p_reg_tile.LayoutType, p_reg_tile.origin, 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!