IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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 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, has_attn_sink: Bool = False, has_extra_kv: Bool = False, has_variable_topk: Bool = False]

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).
  • ​has_attn_sink (Bool): Whether an attention sink is applied, keeping leading tokens always attended (defaults to False).
  • ​has_extra_kv (Bool): Whether a separate extra KV cache holds always-attend tokens (defaults to False).
  • ​has_variable_topk (Bool): Whether the sparse top-k length varies per batch, read from topk_lengths (defaults to False).

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, 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))

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 (eq (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, has_attn_sink, has_extra_kv, has_variable_topk].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, has_attn_sink, has_extra_kv, has_variable_topk].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[DType.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[DType.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(2), IndexList(config, (config // Int(4)), __list_literal__=NoneType(None)), _default_desc_shape[Int(2), output_type, IndexList(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, _safe=False]], indices_stride: Int, topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]], scales_ptr: Pointer[Float32, MutAnyOrigin, _safe=False], attn_sink_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]], extra_k_nope_tma: TMATensorTile[DType.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[DType.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, _safe=False]], extra_topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]], extra_indices_stride: Int, extra_scales_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]], scalar_args: TileTensor[DType.int64, Layout[*?, *?], MutAnyOrigin])

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, _safe=False], kv_lut: KVLUTType, d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]], topk: Int, scales_ptr: Pointer[Float32, MutAnyOrigin, _safe=False], scale_smem_base: Pointer[UInt8, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], offset_position: OffsetPosition[config, KVLUTType, ragged, _is_cache_length_accurate, ValidLengthType, config.decoding_warp_split_k, True, has_extra_kv, has_variable_topk], num_orig_blocks: Int, extra_kv_lut: KVLUTType, extra_d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]], extra_topk: Int, extra_scales_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]])

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:

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[DType.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[DType.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, _safe=False], kv_nope_smem_fp8: Pointer[Float8_e4m3fn, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], kv_rope_smem_bf16_base: Pointer[BFloat16, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], mbar_q: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], kv_load2cvt_pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), Int((add Int(32) if (xor (eq (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, has_variable_topk], idx_bars: DecodeSM100MiscMBars[Int(2), Int(32), Int(32)], idx_smem_base: Pointer[Int32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], num_orig_blocks: Int, topk: Int, extra_k_nope_tma: TMATensorTile[DType.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[DType.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, _safe=False], kv_smem_bf16: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], kv_load2cvt_pipe: KVPipelineGeneric[config.num_kv_stages, Int(1), Int((add Int(32) if (xor (eq (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, _safe=False])

mmaQK​

static def mmaQK(tmem_addr: UInt32, q_smem: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], kv_smem: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], mbar_q: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], 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 (eq (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, has_variable_topk])

mmaPV​

static def mmaPV(tmem_addr: UInt32, kv_smem: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], 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 (eq (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, has_variable_topk])

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact