Mojo module

mha_fa3_utils

comptime values

KVTMATile

comptime KVTMATile[dtype: DType, swizzle_mode: TensorMapSwizzle, *, BN: Int, depth: Int, BK: Int = depth] = TMATensorTile[dtype, _split_last_layout[dtype](IndexList[3, DType.int64](BN, 1, BK, Tuple[]()), swizzle_mode, True), _ragged_desc_layout[dtype](IndexList[3, DType.int64](BN, 1, BK, Tuple[]()), swizzle_mode)]

Parameters

• ​dtype (DType):
• ​swizzle_mode (TensorMapSwizzle):
• ​BN (Int):
• ​depth (Int):
• ​BK (Int):

QTMATile

comptime QTMATile[dtype: DType, swizzle_mode: TensorMapSwizzle, *, BM: Int, depth: Int, group: Int, decoding: Bool] = TMATensorTile[dtype, _split_last_layout[dtype](q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding](), swizzle_mode, True), _ragged_desc_layout[dtype](q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding](), swizzle_mode)]

Parameters

• ​dtype (DType):
• ​swizzle_mode (TensorMapSwizzle):
• ​BM (Int):
• ​depth (Int):
• ​group (Int):
• ​decoding (Bool):

Structs

• ​MHAPosition: Position of the MHA-kernel. When decoding=False, q_head_stride == q_num_heads. When decoding=True, q_head_stride == 1.
• ​NonNullPointer:
• ​NullPointer:
• ​Pack:
• ​PositionSummary:

Traits

• ​OptionalPointer:

Functions

• ​get_q_head_idx:
• ​get_seq_info:
• ​kv_coord:
• ​output_reg_to_smem:
• ​output_reg_to_smem_st_matrix:
• ​produce:
• ​q_coord: Returns the coordinates for a tma load on the Q matrix. This load can be 3D, 4D, or 5D.
• ​q_gmem_shape:
• ​q_smem_shape:
• ​q_tma:

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