Mojo module

mha_sm100_2q

Aliases

LocalTensor

alias LocalTensor[dtype: DType, layout: Layout, element_layout: Layout = Layout(IntTuple(1), IntTuple(1))] = LayoutTensor[dtype, layout, MutableAnyOrigin, address_space=AddressSpace(5), element_layout=element_layout]

Parameters

• ​dtype (DType):
• ​layout (Layout):
• ​element_layout (Layout):

MBarType

alias MBarType = UnsafePointer[SharedMemBarrier, address_space=AddressSpace(3)]

SharedMemPointer

alias SharedMemPointer[type: AnyType] = UnsafePointer[type, address_space=AddressSpace(3)]

Parameters

• ​type (AnyType):

SharedMemTensor

alias SharedMemTensor[dtype: DType, layout: Layout] = LayoutTensor[dtype, layout, MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=DType.int32, linear_idx_type=DType.int32, alignment=128]

Parameters

• ​dtype (DType):
• ​layout (Layout):

Structs

• ​ConsumerPipeline:
• ​FA4Config:
• ​FA4MiscMBars:
• ​KVConsumerPipeline: Pipeline for managing the consumption of K and V. This follows the order of Tri Dao and Cutlass implementations (modulo any rotation of the ops through the iterations).
• ​KVPipeline: KVPipeline has num_kv_stages * num_mma_stages stages. num_kv_stages refers to how many K and V tiles we pipeline for performing the S = Q@K' and O += P@V MMAs. Each of these MMAs is broken up into num_mma_stages pipelined MMAs. We set step=False for all but the last MMA that completes the operation. An alternative implementation would separate the two, and potentially allow for more overall stages at the cost of slightly more bookkeeping.
• ​KVProducerPipeline:
• ​MBarPipeline:
• ​ProducerPipeline:
• ​SM100MHA2Q:
• ​SM100TensorAccumulatorSS:
• ​SM100TensorAccumulatorTS:
• ​STMatrixLayout: Layout for using st_matrix for writing the final accumulator to smem.
• ​STMatrixOffsets:
• ​TMADestination:
• ​TMemTile:

Functions

• ​apply_mask:
• ​break_into_powers_of_two:
• ​cumulative_power_of_two:
• ​extract_power_of_two:
• ​maximum:
• ​mha_sm100_dispatch:
• ​scale_write_output:
• ​sum:

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