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

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

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 = Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]

ProducerPipeline

comptime ProducerPipeline = RolePipeline[_, producer_sub_stages=_, consumer_sub_stages=_, cta_group=_]

SharedMemPointer

comptime SharedMemPointer[type: AnyType] = Pointer[type, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]

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: Manages a paired set of producer/consumer mbarriers for pipeline synchronization.
• ​RolePipeline: Unified producer/consumer pipeline for barrier synchronization.
• ​SM100TensorAccumulator: Performs the C = A @ B tensor contraction on SM100 using tcgen05.mma instructions.
• ​StagedPipeline: Unified pipeline for K, V, and KV tile barrier management.
• ​STMatrixLayout: Layout for using st_matrix for writing the final accumulator to smem.
• ​STMatrixOffsets: Precomputed TMEM and local-fragment offsets for one st_matrix repeat column.
• ​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: Represents a tile in SM100 tensor memory (TMEM) and provides async load/store helpers.

Functions

• ​add_ftz: Returns the flush-to-zero sum of two float32 values.
• ​add_ftz_rm: Returns the round-to-nearest-even flush-to-zero sum of two f32x2 vectors.
• ​apply_mask: Applies bitmask, computed, and out-of-bounds masking strategies to a row of BN attention scores.
• ​apply_oob_mask: Applies the out-of-bounds key mask to a pair of attention scores.
• ​break_into_powers_of_two: Calls func for each power-of-two-sized chunk of N, plus a final pow_two=0 call for pipeline cleanup.
• ​bulk_mma: Issues a full-tile SS (both operands in SMEM) non-warp-specialized tcgen05.mma contraction.
• ​bulk_mma_partial: Issues a partial-K TS contraction for a partially-loaded last KV tile, non-warp-specialized.
• ​bulk_mma_ss_partial: Issues a partial-K SS contraction for a partially-loaded last KV tile, non-warp-specialized.
• ​bulk_mma_ws: Issues a full-tile SS (both operands in SMEM) warp-specialized tcgen05.mma.ws contraction.
• ​bulk_mma_ws_partial: Issues a partial-K SS warp-specialized contraction for a partially-loaded last KV tile.
• ​bulk_mma_ws_ts: Issues a full-tile TS (A in TMEM, B in SMEM) warp-specialized tcgen05.mma.ws contraction.
• ​bulk_mma_ws_ts_partial: Issues a partial-K TS warp-specialized contraction for a partially-loaded last KV tile.
• ​cluster_remote_smem_addr: Map a local .shared byte address to peer peer_rank's window in the cluster.
• ​clusters_per_wave: Number of size-cluster_size thread-block clusters that fit on the target Blackwell datacenter GPU in ONE wave, honoring GPC co-residency.
• ​combine_pack_o_row: LSE-combine own * scale_own + peer * scale_peer over n f32 O lanes, cast to the 2-byte output_type, and pack into n // 2 u32 lanes.
• ​cumulative_power_of_two: Returns the cumulative sum of the first i power-of-two components of N.
• ​elect_mma_arrive: Arrive at the mbar pointer for the MMA instruction.
• ​exp2_emulation: Computes 2^x for an f32x2 vector via a degree-3 polynomial approximation.
• ​expect_bytes_pred: Issue mbarrier.arrive.expect_tx.shared::cta.b64 predicated on pred != 0.
• ​extract_power_of_two: Returns the i-th power-of-two component when decomposing N into decreasing powers of two.
• ​fma_ftz:
• ​intrin: Wraps a ternary float32 PTX intrinsic (e.g. max.f32).
• ​intrin_ftz: Wraps a flush-to-zero (FTZ) binary float32 PTX intrinsic.
• ​intrin_ftz_x2: Wraps a flush-to-zero (FTZ) binary f32x2 PTX intrinsic.
• ​llvm_opaque_tid: Returns the opaque thread ID via the llvm.nvvm.read.ptx.sreg.tid.x intrinsic.
• ​load_cluster_smem: Load width elements from peer peer_rank's shared memory at local_ptr.
• ​mask_select8: Masks 8 contiguous scores against one byte of a 32-column bitmask.
• ​max_ftz: Returns the flush-to-zero maximum of two float32 values.
• ​maximum: Reduces BN float32 scores into width lane-maxima using FTZ max.
• ​mul_ftz: Returns the flush-to-zero product of two float32 values.
• ​o_store_tma_blocks_per_op: Box size (swizzle-granularity blocks per batched O-store TMA).
• ​pack_row: Cast the w f32 O lanes o_vals[start : start + w] to output_type and pack them into one 16 B SWIZZLE_NONE store register (exactly four u32).
• ​peel_mask: Determine which mask strategy applies to the peeled first iteration.
• ​scale_pack_o_row: Scale the w f32 O lanes o_vals[start : start + w] by inv_row_sum, cast to the 2-byte output_type, and pack into w // 2 u32 lanes (the row-major 16 B SWIZZLE_NONE store register).
• ​splitk_partition_idx: This CTA's split-K partition index [0, splitk_partitions).
• ​splitk_window: Front-loaded balanced split of the combined K-tile range [0, T).
• ​st_shared_v4_b32: Explicit 16 B st.shared.v4.b32 (one STS.128) to dst[elem_off].
• ​store_cluster_smem: Store val into peer peer_rank's shared memory at local_ptr.
• ​sub_ftz: Returns the flush-to-zero difference of two float32 values.
• ​sum: Reduces a BN-element local tensor into a width-2 SIMD vector via vectorized accumulation.

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