Mojo function

fa4_softmax

fa4_softmax[KVLUTType: MHAOperand, output_type: DType, MaskType: MHAMask, config: FA4Config, ValidLengthType: OptionalPointer, SinkType: OptionalPointer, _is_cache_length_accurate: Bool, MaxSeqLenType: OptionallyStaticInt](mbars: FA4MiscMBars[num_qk_stages=config.num_qk_stages, num_pv_stages=config.num_pv_stages, num_kv_stages=config.num_kv_stages, use_order_barriers=EnableForcedOrdering], score_row: UInt32, seq_info: SeqInfo, mask: MaskType, num_keys: UInt32, scale: Float32, max_seq_len: UInt32, ragged_tma_store: RaggedTMA3DTile[output_type, config.swizzle_mode, (config // 2), config.depth], sink_weights: SinkType)

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