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

mha_1q

Implements single-query (decode) multi-head attention for NVIDIA SM100 (Blackwell) GPUs using warp-specialized UMMA and tensor-memory (TMEM) accumulators.

comptime values

logger

comptime logger = Logger(stdout, prefix=String(""), source_location=False)

Structs

• ​MMAOperandOffsetFn: Computes the shared-memory layout and byte offsets for MMA operand tiles, bridging typed tile layouts to legacy MMA descriptors.
• ​RegisterAccumulatorDescription: Holds the number of MMA fragments and the per-thread fragment size for a register-resident accumulator tile.
• ​RegisterAccumulatorLayout: Describes how UMMA accumulator fragments are distributed across the threads of a consumer warp group.
• ​SM100TensorAccumulatorSS: Manages a shared-shared (SS) UMMA accumulator pipeline for SM100, coordinating MMA, TMEM, and barrier synchronization between producer and consumer warps.
• ​SM100TensorAccumulatorTS: Manages a tensor-shared (TS) UMMA accumulator pipeline for SM100, coordinating MMA between a TMEM A-operand and a shared-memory B-operand.
• ​TMemAccumulator: Implements an AccumulatorTile backed by tensor memory (TMEM), storing UMMA results at a given TMEM address.
• ​TMemOperand: Implements a WriteableMMAOperandDescriptor backed by tensor memory (TMEM), used as the A operand of a tensor-shared (TS) UMMA.
• ​UMMADescriptorSS: Holds two shared-memory descriptors for the A and B operands of a shared-shared (SS) UMMA.
• ​UMMADescriptorTS: Pairs a TMEM A-operand descriptor with a shared-memory B-operand descriptor for a tensor-shared (TS) UMMA.

Traits

• ​AccumulatorTile: Describes a UMMA accumulator tile that can be allocated, copied to and from, and viewed as rows of fragments.
• ​DescriptorPair: Provides access to the A and B operand descriptors for a shared-shared (SS) UMMA operation.
• ​DescriptorPairTS: Provides access to the A (tensor-memory) and B (shared-memory) operand descriptors for a tensor-shared (TS) UMMA operation.
• ​WriteableMMAOperandDescriptor: Describes an MMA operand whose source data can be written from a local LayoutTensor into the operand's backing store (e.g. tensor memory).

Functions

• ​local_tensor_type: Returns an uninitialized local-address-space LayoutTensor used for compile-time type inference of register tiles.
• ​mha_sm100_dispatch: Dispatches single-query (decode) multi-head attention to the SM100 kernel, selecting TMA tiles, scheduler, and partition configuration.

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