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 struct

MLA_SM100_Decode_Sparse

struct MLA_SM100_Decode_Sparse[q_type: DType, KVLUTType: MHAOperand, output_type: DType, SplitAccumType: OptionalPointer, MaskType: MHAMask, config: MLA_SM100_Decode_Config, ValidLengthType: OptionalPointer, _is_cache_length_accurate: Bool = False, ragged: Bool = False, AttnSinkPtrType: OptionalPointer = NullPointer[.float32], has_extra_kv: Bool = False, TopkLengthsPtrType: OptionalPointer = NullPointer[.int32], Engine: TensorEngine = DefaultEngine]

Sparse MLA decoding kernel for SM100 with FP8 KV cache and gather4 TMA.

Extends the SM100 MLA decode kernel to sparse attention by loading KV tiles through gather4 TMA using per-tile sparse indices. Uses four warpgroups plus four individual warps: a convert warpgroup performs FP8-to-BF16 conversion in shared memory before the QK and PV MMA warps consume the data. Supports blockwise FP8 scaling, attention sinks, an extra KV cache for always-attend tokens, and variable top-k lengths.

Parameters

• ​q_type (DType): DType of the query tensor and softmax accumulator.
• ​KVLUTType (MHAOperand): MHAOperand describing the paged KV cache layout and element type; its dtype is the FP8 KV element type.
• ​output_type (DType): DType of the output tensor written via TMA store.
• ​SplitAccumType (OptionalPointer): OptionalPointer selecting the split-K accumulation buffer dtype used by the combine kernel.
• ​MaskType (MHAMask): MHAMask selecting the attention mask; only NullMask and CausalMask are supported.
• ​config (MLA_SM100_Decode_Config): MLA_SM100_Decode_Config with tile sizes, stage counts, and swizzle parameters for the decode kernel.
• ​ValidLengthType (OptionalPointer): OptionalPointer selecting the dtype of the per-batch valid-length buffer.
• ​_is_cache_length_accurate (Bool): Whether the cached sequence length is exact, enabling tighter bounds (defaults to False).
• ​ragged (Bool): Whether ragged batching is enabled, so each batch may have a different query length (defaults to False).
• ​AttnSinkPtrType (OptionalPointer): OptionalPointer type of the per-head attention sink; non-null applies the sink (defaults to NullPointer).
• ​has_extra_kv (Bool): Whether a separate extra KV cache holds always-attend tokens (defaults to False).
• ​TopkLengthsPtrType (OptionalPointer): OptionalPointer type of topk_lengths and extra_topk_lengths; non-null reads the sparse top-k length per batch instead of using the stride (defaults to NullPointer).
• ​Engine (TensorEngine): Engine policy of the scalar_args tile operand (defaults to DefaultEngine[element_width=1]).

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

AccumType

comptime AccumType = get_accum_type[q_type]()

bf16_bytes_per_element

comptime bf16_bytes_per_element = size_of[DType.bfloat16]()

bf16_type

comptime bf16_type = DType.bfloat16

BlockElems

comptime BlockElems = (config * config)

bytes_per_element

comptime bytes_per_element = size_of[q_type]()

Common_MLA_Op

comptime Common_MLA_Op = MLA_SM100_Decode_Common[q_type, KVLUTType, output_type, SplitAccumType, MaskType, config, ValidLengthType, _is_cache_length_accurate, ragged]

fp8_bytes_per_element

comptime fp8_bytes_per_element = size_of[DType.float8_e4m3fn]()

fp8_type

comptime fp8_type = DType.float8_e4m3fn

gather4_indices_bytes

comptime gather4_indices_bytes = (config * size_of[Int32]())

gather4_num_4row_chunks

comptime gather4_num_4row_chunks = (config // Int(4))

has_attn_sink

comptime has_attn_sink = not AttnSinkPtrType.is_null.__bool__()

has_variable_topk

comptime has_variable_topk = not TopkLengthsPtrType.is_null.__bool__()

kv_nope_type

comptime kv_nope_type = KVLUTType.dtype

kv_rope_type

comptime kv_rope_type = DType.bfloat16

kv_type

comptime kv_type = KVLUTType.dtype

KVStageElems

comptime KVStageElems = ((config // config) * Int((mul config.BM, config.BN_QK)))

KVStageTotalBytes

comptime KVStageTotalBytes = (Int((mul (config // config), size_of[DType.float8_e4m3fn](), config.BM, config.BN_QK)) + Int((mul (config // config), size_of[DType.bfloat16](), config.BM, config.BN_QK)))

load2cvt_num_producer

comptime load2cvt_num_producer = (Int(1) + Int(32) if (xor identical((xor (lt config.scale_block_size, 1), True), False), True) else Int(0))

nope_gather4_box_w

comptime nope_gather4_box_w = _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle.SWIZZLE_NONE]()

nope_gather4_num_col_groups

comptime nope_gather4_num_col_groups = ceildiv((config // Int(8)), _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))]())

nope_gather4_tile_width

comptime nope_gather4_tile_width = (config // Int(8))

NopeStageBytes

comptime NopeStageBytes = (Int((mul (config // config), config.BM, config.BN_QK)) * size_of[DType.float8_e4m3fn]())

NopeStageElems

comptime NopeStageElems = ((config // config) * Int((mul config.BM, config.BN_QK)))

NumNopeBlocks

comptime NumNopeBlocks = (config // config)

NumQKBlocks

comptime NumQKBlocks = (config // config)

NumRopeBlocks

comptime NumRopeBlocks = (config // config)

NumVOBlocks

comptime NumVOBlocks = (config // config)

output_tile_width

comptime output_tile_width = ((config // Int(2)) * (Int(4) // size_of[output_type]()))

rope_gather4_box_w

comptime rope_gather4_box_w = _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle.SWIZZLE_128B]()

rope_gather4_num_col_groups

comptime rope_gather4_num_col_groups = ceildiv(config.rope_depth, _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle(Int32(3))]())

rope_gather4_tile_width

comptime rope_gather4_tile_width = (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2))

RopeStageBytes

comptime RopeStageBytes = (Int((mul (config // config), config.BM, config.BN_QK)) * size_of[DType.bfloat16]())

RopeStageElems

comptime RopeStageElems = ((config // config) * Int((mul config.BM, config.BN_QK)))

UMMAPVSS

comptime UMMAPVSS = DecodeSM100PVSS[q_type, MLA_SM100_Decode_Sparse[q_type, KVLUTType, output_type, SplitAccumType, MaskType, config, ValidLengthType, _is_cache_length_accurate, ragged, AttnSinkPtrType, has_extra_kv, TopkLengthsPtrType, Engine].AccumType, config=config]

UMMAQKTSS

comptime UMMAQKTSS = DecodeSM100QKTSS[q_type, MLA_SM100_Decode_Sparse[q_type, KVLUTType, output_type, SplitAccumType, MaskType, config, ValidLengthType, _is_cache_length_accurate, ragged, AttnSinkPtrType, has_extra_kv, TopkLengthsPtrType, Engine].AccumType, config=config]

Methods

kernel

static def kernel(q_tma: TMATensorTile[q_type, Int(2), IndexList(config, config, __list_literal__=NoneType(None)), _default_desc_shape[Int(2), q_type, IndexList(config, config, __list_literal__=NoneType(None)), config.swizzle_mode]()], k_nope_tma: TMATensorTile[.int64, Int(2), IndexList(config, _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))](), __list_literal__=NoneType(None))], k_rope_tma: TMATensorTile[.bfloat16, Int(2), IndexList(config, _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle(Int32(3))](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle(Int32(3))](), __list_literal__=NoneType(None))], o_tma: TMATensorTile[output_type, Int(3), IndexList(Int(1), config, (config // Int(4)), __list_literal__=NoneType(None)), _default_desc_shape[Int(3), output_type, IndexList(Int(1), config, (config // Int(4)), __list_literal__=NoneType(None)), config.swizzle_mode]()], kv_lut: KVLUTType, scale: Float32, mla_decode_pack: MLA_Decode_Pack[ValidLengthType, MaskType, SplitAccumType], d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin]], indices_stride_dev: Int32, topk_lengths: TopkLengthsPtrType, scales_ptr: Pointer[Float32, MutAnyOrigin], attn_sink_ptr: AttnSinkPtrType, extra_k_nope_tma: TMATensorTile[.int64, Int(2), IndexList(config, _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))](), __list_literal__=NoneType(None))], extra_k_rope_tma: TMATensorTile[.bfloat16, Int(2), IndexList(config, _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle(Int32(3))](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle(Int32(3))](), __list_literal__=NoneType(None))], extra_kv_lut: KVLUTType, extra_d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin]], extra_topk_lengths: TopkLengthsPtrType, extra_indices_stride_dev: Int32, extra_scales_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin]], scalar_args: TileTensor[.int64, Layout[TypeList[ComptimeInt[Int(3)]](), TypeList[ComptimeInt[Int(1)]]()], MutAnyOrigin, Engine=Engine])

idx_producer

static def idx_producer(idx_bars: DecodeSM100MiscMBars[Int(2), Int(32), Int(32)], idx_smem_base: Pointer[Int32, MutAnyOrigin, address_space=AddressSpace.SHARED], kv_lut: KVLUTType, d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin]], topk: Int, scales_ptr: Pointer[Float32, MutAnyOrigin], scale_smem_base: Pointer[UInt8, MutAnyOrigin, address_space=AddressSpace.SHARED], offset_position: OffsetPosition[config, KVLUTType, ragged, _is_cache_length_accurate, ValidLengthType, config.decoding_warp_split_k, True, has_extra_kv, TopkLengthsPtrType], num_orig_blocks: Int, extra_kv_lut: KVLUTType, extra_d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin]], extra_topk: Int, extra_scales_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin]])

Index transform producer running on warp 11 (32 threads).

For each KV tile, transforms d_indices to TMA rows and (when blockwise) loads FP32 scales to scale SMEM. Signals idx_bars when each tile's data is ready. Runs 1 tile ahead of warp 8.

Args:

• ​idx_bars (DecodeSM100MiscMBars[Int(2), Int(32), Int(32)]): Double-buffered barrier pipeline between this producer (warp 11) and the index consumer (warp 8).
• ​idx_smem_base (Pointer[Int32, MutAnyOrigin, address_space=AddressSpace.SHARED]): Base SMEM pointer for the double-buffered transformed gather4 row indices, BN_QK Int32 per stage.
• ​kv_lut (KVLUTType): MHAOperand for the original paged KV cache, used to convert encoded indices to TMA rows.
• ​d_indices (OptionalReg[Pointer[Int32, MutAnyOrigin]]): Per-tile sparse indices into the original KV cache, encoded as physical_block * page_size + offset; -1 marks invalid entries.
• ​topk (Int): Number of selected KV tokens in the original cache for this batch.
• ​scales_ptr (Pointer[Float32, MutAnyOrigin]): Pointer to FP32 blockwise scales for the original KV cache.
• ​scale_smem_base (Pointer[UInt8, MutAnyOrigin, address_space=AddressSpace.SHARED]): Base SMEM pointer for the double-buffered e8m0 scales, used when config.scale_block_size > 0.
• ​offset_position (OffsetPosition[config, KVLUTType, ragged, _is_cache_length_accurate, ValidLengthType, config.decoding_warp_split_k, True, has_extra_kv, TopkLengthsPtrType]): Precomputed per-CTA position and offset state for the current batch and split.
• ​num_orig_blocks (Int): Number of BN_QK-sized KV tiles belonging to the original cache, equals ceildiv(topk, BN_QK).
• ​extra_kv_lut (KVLUTType): MHAOperand for the extra paged KV cache holding always-attend tokens.
• ​extra_d_indices (OptionalReg[Pointer[Int32, MutAnyOrigin]]): Per-tile sparse indices into the extra KV cache, encoded as physical_block * page_size + offset; -1 marks invalid entries.
• ​extra_topk (Int): Number of selected KV tokens in the extra cache for this batch.
• ​extra_scales_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin]]): Pointer to FP32 blockwise scales for the extra KV cache, or None when blockwise scaling is off.

load

static def load(q_tma: TMATensorTile[q_type, Int(2), IndexList(config, config, __list_literal__=NoneType(None)), _default_desc_shape[Int(2), q_type, IndexList(config, config, __list_literal__=NoneType(None)), config.swizzle_mode]()], k_nope_tma: TMATensorTile[.int64, Int(2), IndexList(config, _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))](), __list_literal__=NoneType(None))], k_rope_tma: TMATensorTile[.bfloat16, Int(2), IndexList(config, _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle(Int32(3))](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle(Int32(3))](), __list_literal__=NoneType(None))], q_smem: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED], kv_nope_smem_fp8: Pointer[Float8_e4m3fn, MutAnyOrigin, address_space=AddressSpace.SHARED], kv_rope_smem_bf16_base: Pointer[BFloat16, MutAnyOrigin, address_space=AddressSpace.SHARED], mbar_q: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED], kv_load2cvt_pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), Int((add Int(32) if (xor identical((xor (lt config.scale_block_size, 1), True), False), True) else Int(0), 1)), (_resolve_warpgroup_size() + Int(2))], offset_position: OffsetPosition[config, KVLUTType, ragged, _is_cache_length_accurate, ValidLengthType, config.decoding_warp_split_k, True, has_extra_kv, TopkLengthsPtrType], idx_bars: DecodeSM100MiscMBars[Int(2), Int(32), Int(32)], idx_smem_base: Pointer[Int32, MutAnyOrigin, address_space=AddressSpace.SHARED], num_orig_blocks: Int, topk: Int, extra_k_nope_tma: TMATensorTile[.int64, Int(2), IndexList(config, _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))](), __list_literal__=NoneType(None))], extra_k_rope_tma: TMATensorTile[.bfloat16, Int(2), IndexList(config, _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle(Int32(3))](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.bfloat16, (Int((add (mul config.rope_depth, 2), config.padded_depth)) // Int(2)), TensorMapSwizzle(Int32(3))](), __list_literal__=NoneType(None))], extra_topk: Int)

convertFP8ToBF16

static def convertFP8ToBF16(kv_nope_smem_fp8: Pointer[Float8_e4m3fn, MutAnyOrigin, address_space=AddressSpace.SHARED], kv_smem_bf16: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED], kv_load2cvt_pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), Int((add Int(32) if (xor identical((xor (lt config.scale_block_size, 1), True), False), True) else Int(0), 1)), (_resolve_warpgroup_size() + Int(2))], kv_cvt2mma_pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), _resolve_warpgroup_size(), Int(2)], num_k_tiles: Int, scale_smem_base: Pointer[UInt8, MutAnyOrigin, address_space=AddressSpace.SHARED])

mmaQK

static def mmaQK(tmem_addr: UInt32, q_smem: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED], kv_smem: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED], mbar_q: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED], s_bars: DecodeSM100MiscMBars[Int(2), Int(1), _resolve_warpgroup_size()], kv_cvt2mma_pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), _resolve_warpgroup_size(), Int(2)], kv_load2cvt_pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), Int((add Int(32) if (xor identical((xor (lt config.scale_block_size, 1), True), False), True) else Int(0), 1)), (_resolve_warpgroup_size() + Int(2))], offset_position: OffsetPosition[config, KVLUTType, ragged, _is_cache_length_accurate, ValidLengthType, config.decoding_warp_split_k, True, has_extra_kv, TopkLengthsPtrType])

mmaPV

static def mmaPV(tmem_addr: UInt32, kv_smem: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED], p_bars: DecodeSM100MiscMBars[Int(2), _resolve_warpgroup_size(), Int(1)], o_bars: DecodeSM100MiscMBars[Int(2), Int(1), _resolve_warpgroup_size()], kv_cvt2mma_pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), _resolve_warpgroup_size(), Int(2)], kv_load2cvt_pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), Int((add Int(32) if (xor identical((xor (lt config.scale_block_size, 1), True), False), True) else Int(0), 1)), (_resolve_warpgroup_size() + Int(2))], offset_position: OffsetPosition[config, KVLUTType, ragged, _is_cache_length_accurate, ValidLengthType, config.decoding_warp_split_k, True, has_extra_kv, TopkLengthsPtrType])