Mojo function
apply_mask
apply_mask[dtype: DType, BN: Int, MaskType: MHAMask, ScoreModType: ScoreModTrait, //, *, use_score_mod: Bool, decoding: Bool = False](srow: LayoutTensor[dtype, Layout.row_major(BN), MutableAnyOrigin, address_space=AddressSpace(5)], mask: MaskType, score_mod: ScoreModType, scale_log2e: Scalar[dtype], mask_status: TileMaskStatus, *, prompt_idx: UInt32, q_head_idx: UInt32, kv_tile_start_row: UInt32, seq_len: UInt32, max_seq_len: UInt32, num_keys: UInt32, row: UInt32, start_pos: UInt32)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!