Mojo function

fa4_softmax

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](smem: SM100AttentionSMem[config], score_row: UInt32, seq_info: SeqInfo, mask: MaskType, num_keys: UInt32, scale: Float32, max_seq_len: UInt32, ragged_tma_store: RaggedTMA3DTile[output_type, ragged_tma_store.swizzle_mode, (config // 2), config.ov_depth], 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