For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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.