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 function

fa4_softmax

def fa4_softmax[QScaleType: OptionalPointer, KScaleType: OptionalPointer, qkv_dtype: DType, rope_dtype_: Optional[DType], scale_dtype_: Optional[DType], output_type: DType, MaskType: MHAMask, //, KVLUTType: MHAOperand, config: FA4Config[qkv_dtype, rope_dtype_=rope_dtype_, scale_dtype_=scale_dtype_], ValidLengthType: OptionalPointer, SinkType: OptionalPointer, _is_cache_length_accurate: Bool, MaxSeqLenType: OptionallyStaticInt, output_nonempty: Bool = False, single_softmax_wg: Bool = False](smem: SM100AttentionSMem[config], tmem_addr: UInt32, score_row: UInt32, seq_info: SeqInfo, mask: MaskType, num_keys: UInt32, scale: Float32, max_seq_len: UInt32, ragged_tma_store: RaggedTMA3DTile[output_type, BM=(config // config), BN=config.ov_depth, middle_dim=ragged_tma_store.middle_dim, group=config.group if config.fuse_gqa else Int(1), tma_blocks_per_op=ragged_tma_store.tma_blocks_per_op], sink_weights: SinkType, q_scale: QScaleType = NullPointer(), k_scale: KScaleType = NullPointer())