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_scale_write_output

def fa4_scale_write_output[output_type: DType, //, config: FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_], output_swizzle_mode: TensorMapSwizzle, tma_bpo: Int, zero_fill: Bool = False](local_row: UInt32, local_warp_idx: UInt32, warp_group_idx: UInt32, inv_row_sum: Float32, o_smem_arg: Pointer[Scalar[output_type], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], o_tmem_arg: TMemTile[DType.float32, (config // config), config.padded_ov_depth], ragged_tma_store: RaggedTMA3DTile[output_type, output_swizzle_mode, 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=tma_bpo], num_output_rows: Int32, out_head_idx: UInt32, out_row_idx: UInt32)