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: DType, scale_dtype: 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](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, group=config.group if config.fuse_gqa else Int(1)], sink_weights: SinkType, q_scale: QScaleType = NullPointer(), k_scale: KScaleType = NullPointer())

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit