Mojo function
depth512_softmax
depth512_softmax[MaskType: MHAMask, qkv_dtype: DType, output_type: DType, config: Depth512SM100Config[qkv_dtype], page_size: Int](smem: Depth512AttentionSMem[config], score_row: UInt32, num_keys: UInt32, mask: MaskType, scale: Float32, ragged_tma_store: RaggedTMA3DTile[output_type, config.swizzle_mode, Depth512SM100Config[qkv_dtype].BM, config.ov_depth], num_output_rows: Int32, out_head_idx: UInt32, out_row_idx: UInt32)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!