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

mla_decode_utils

Provides shared utilities for SM100 MLA decode attention kernels.

Defines TMA tile helpers, pipeline producer/consumer structs, MMA tensor accumulator descriptors, and the common softmax/correction/store logic reused across the BF16, native FP8, and per-token-scale decode backends.

comptime values

QOTMATile

comptime QOTMATile[dtype: DType, BM: Int, BK: Int, swizzle_mode: TensorMapSwizzle] = TMATensorTile[dtype, Int(2), IndexList(BM, BK, __list_literal__=NoneType(None)), _default_desc_shape[Int(2), dtype, IndexList(BM, BK, __list_literal__=NoneType(None)), swizzle_mode]()]

Parameters

• ​dtype (DType):
• ​BM (Int):
• ​BK (Int):
• ​swizzle_mode (TensorMapSwizzle):

ScalesTMATile

comptime ScalesTMATile[BN_QK: Int] = TMATensorTile[DType.float32, Int(2), IndexList(Int(1), BN_QK, __list_literal__=NoneType(None))]

Parameters

• ​BN_QK (Int):

Structs

• ​DecodeCConsumer: Consumer side of the single-stage C pipeline between softmax and correction.
• ​DecodeCProducer: Producer side of the single-stage C pipeline between softmax and correction.
• ​DecodeKVConsumer: Consumer side of the decode KV pipeline that waits for and releases KV stages.
• ​DecodeKVProducer: Producer side of the decode KV pipeline that loads KV tiles via TMA.
• ​DecodeOConsumer: Consumer side of the two-stage O pipeline between MMA and correction.
• ​DecodeOProducer: Producer side of the two-stage O pipeline between MMA and correction.
• ​DecodeOutConsumer: Consumer side of the output writeback pipeline that waits for and releases output stages.
• ​DecodeOutProducer: Producer side of the output writeback pipeline that stages output tiles in SMEM for TMA store.
• ​DecodePConsumer: Consumer side of the two-stage P pipeline between softmax and MMA.
• ​DecodePConsumerN: N-stage parameterized consumer side of the P pipeline between softmax and MMA.
• ​DecodePProducer: Producer side of the two-stage P pipeline between softmax and MMA.
• ​DecodePProducerN: N-stage parameterized producer side of the P pipeline between softmax and MMA.
• ​DecodeSConsumer: Consumer side of the two-stage S pipeline between MMA and softmax.
• ​DecodeSConsumerN: N-stage parameterized consumer side of the S pipeline between MMA and softmax.
• ​DecodeSM100MiscMBars: Manages a generic producer/consumer mbarrier pair for the S, P, C, and O pipelines.
• ​DecodeSM100PVSS: Tensor accumulator for the PV MMA with both P and V operands in SMEM.
• ​DecodeSM100PVSS_FP8: Tensor accumulator for the native FP8 PV MMA with both P and V in FP8 SMEM.
• ​DecodeSM100QKTSS: Tensor accumulator for the QK^T MMA with both Q and K operands in SMEM.
• ​DecodeSM100QKTSS_Content_FP8: Tensor accumulator for the content-only FP8 QK^T MMA used by the per-token-scale rope-aware kernel.
• ​DecodeSM100QKTSS_FP8: Tensor accumulator for the native FP8 QK^T MMA with both Q and K in FP8 SMEM.
• ​DecodeSM100QKTSS_Rope_BF16: Tensor accumulator for the rope-only BF16 QK^T MMA used by the per-token-scale rope-aware kernel.
• ​DecodeSM100QKTTS: Tensor accumulator for the QK^T MMA with Q in TMEM and K in SMEM.
• ​DecodeSProducer: Producer side of the two-stage S pipeline between MMA and softmax.
• ​DecodeSProducerN: N-stage parameterized producer side of the S pipeline between MMA and softmax.
• ​KVCvt2MmaConsumer: Consumes BF16 KV tiles from the convert producer for the MMA pipeline.
• ​KVCvt2MmaProducer: Produces converted BF16 KV tiles for the MMA consumer pipeline.
• ​KVLoad2CvtConsumer: Consumer side of the FP8-to-BF16 load-and-convert KV pipeline.
• ​KVLoad2CvtProducer: Producer side of the FP8-to-BF16 load-and-convert KV pipeline.
• ​KVPipelineGeneric: KVPipeline has num_kv_stages * num_qk_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_qk_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.
• ​MLA_Decode_Pack: Bundles the mask, valid-length, and split-K accumulator pointers passed to decode kernels.
• ​MLA_SM100_Decode_Common: Provides the shared softmax, correction, and store logic for SM100 MLA decode kernels.
• ​MLA_SM100_Decode_Config: Holds the tile sizes, swizzle modes, and SMEM/TMEM layout for an SM100 MLA decode kernel.
• ​OffsetPosition: Computes and stores per-CTA row offsets and KV key ranges for the decode kernel.
• ​OutPipeline: OutPipeline has num_out_stages stages. num_out_stages refers to how many output stages we pipeline for performing the output store.

Functions

• ​clamped_index_coordinate: Builds a four-component index coordinate with the key index clamped to the last valid key.
• ​cvt_fp8x16_from_u32x4_to_bf16x16_packed_2xu32x4: Converts 16 FP8 bytes (one v4.b32 load) to 16 packed BF16 values.
• ​cvt_fp8x8_from_2xu32_to_bf16x8_packed_u32x4: Converts eight FP8 values packed in two uint32 registers to eight BF16 values packed in four uint32 registers.
• ​e8m0_to_bf16_broadcast: Convert an e8m0 scale byte to a bf16 value broadcast into both halves of a uint32.
• ​hmul2_bf16x8_by_scalar: Multiply 8 packed bf16 values (in 4 uint32 registers) by a bf16x2 scalar broadcast.
• ​ld_shared_v4_u32: Loads four contiguous uint32 values from shared memory at a byte offset.
• ​num_matrix_view_rows_decode: TileTensor overload of num_matrix_view_rows_decode.
• ​st_shared_v4_b32_at_bf16_elem_off: Stores four uint32 values to shared memory at a BF16 element offset.
• ​st_shared_v4_b32_at_fp8_elem_off: Stores four uint32 values to shared memory at an FP8 element offset.
• ​tma_tile_qo: Creates a TMA descriptor for the Q or output tensor used in MLA decode.
• ​tma_tile_scales: Create a TMA descriptor for per-token float32 scales.
• ​write_bf16x2_row_to_smem_chunked: Chunked write with optional scaling. Reduces register pressure.
• ​write_fp8_row_to_smem_chunked: Writes float32 data to SMEM as FP8 with FP8-byte swizzle.

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