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 module
attention_utils
Shared SM100 attention primitives used by both MHA and MLA kernels.
This module contains generic SM100 (Blackwell) GPU primitives including:
- TMEM access helpers (TMemTile, STMatrixLayout)
- Pipeline synchronization (StagedPipeline, RolePipeline, etc.)
- FTZ arithmetic (add_ftz, sub_ftz, mul_ftz, etc.)
- Barrier helpers (FA4MiscMBars)
- MMA building blocks (bulk_mma, SM100TensorAccumulator)
- Masking utilities (apply_mask, apply_oob_mask)
comptime values
ConsumerPipeline
comptime ConsumerPipeline = RolePipeline[_, False, _, _, _]
FP32_EXP_BIAS
comptime FP32_EXP_BIAS = 127
KConsumerPipeline
comptime KConsumerPipeline = TMAConsumerPipeline[_, _]
KPipeline
comptime KPipeline = StagedPipeline[_, _]
KProducerPipeline
comptime KProducerPipeline = TMAProducerPipeline[_, _]
KVPipeline
comptime KVPipeline = StagedPipeline[_, _]
LocalLT
comptime LocalLT[dtype: DType, layout: Layout, element_layout: Layout = Layout(IntTuple(Int(1)), IntTuple(Int(1)))] = LayoutTensor[dtype, layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=element_layout]
Parameters
LocalTensor
comptime LocalTensor[dtype: DType, layout: Layout[shape_types, stride_types]] = TileTensor[dtype, Layout[shape_types, stride_types], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]
Parameters
- dtype (
DType): - layout (
Layout[shape_types, stride_types]):
MBarType
comptime MBarType = UnsafePointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED]
ProducerPipeline
comptime ProducerPipeline = RolePipeline[_, producer_sub_stages=_, consumer_sub_stages=_, cta_group=_]
SharedMemPointer
comptime SharedMemPointer[type: AnyType] = UnsafePointer[type, MutAnyOrigin, address_space=AddressSpace.SHARED]
Parameters
- type (
AnyType):
SharedMemTensor
comptime SharedMemTensor[dtype: DType, layout: Layout[shape_types, stride_types]] = TileTensor[dtype, Layout[shape_types, stride_types], MutUntrackedOrigin, address_space=AddressSpace.SHARED]
Parameters
- dtype (
DType): - layout (
Layout[shape_types, stride_types]):
VConsumerPipeline
comptime VConsumerPipeline = TMAConsumerPipeline[_, _, False]
VPipeline
comptime VPipeline = StagedPipeline[_]
VProducerPipeline
comptime VProducerPipeline = TMAProducerPipeline[_, _, False]
Structs
-
FA4MiscMBars: Manages all mbarrier resources for FA4. -
MBarPipeline: -
RolePipeline: Unified producer/consumer pipeline for barrier synchronization. -
SM100TensorAccumulator: -
StagedPipeline: Unified pipeline for K, V, and KV tile barrier management. -
STMatrixLayout: Layout for usingst_matrixfor writing the final accumulator to smem. -
STMatrixOffsets: -
TMAConsumerPipeline: Unified consumer pipeline for K and V TMA consumption. -
TMADestination: Pairs a shared memory TileTensor with a barrier for TMA operations. -
TMAProducerPipeline: Unified producer pipeline for K and V TMA loading. -
TMemTile:
Functions
-
add_ftz: -
add_ftz_rm: -
apply_mask: -
apply_oob_mask: -
break_into_powers_of_two: -
bulk_mma: -
bulk_mma_partial: -
bulk_mma_ss_partial: -
bulk_mma_ws: -
bulk_mma_ws_partial: -
bulk_mma_ws_ts: -
bulk_mma_ws_ts_partial: -
cumulative_power_of_two: -
elect_mma_arrive: Arrive at the mbar pointer for the MMA instruction. -
exp2_emulation: -
expect_bytes_pred: Issuembarrier.arrive.expect_tx.shared::cta.b64predicated onpred != 0. -
extract_power_of_two: -
fma_ftz: -
intrin: -
intrin_ftz: -
intrin_ftz_x2: -
llvm_opaque_tid: -
max_ftz: -
maximum: -
mul_ftz: -
peel_mask: Determine which mask strategy applies to the peeled first iteration. -
sub_ftz: -
sum:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!