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_Common

struct MLA_SM100_Decode_Common[q_type: DType, KVLUTType: MHAOperand, output_dtype: DType, SplitAccumType: OptionalPointer, MaskType: MHAMask, config: MLA_SM100_Decode_Config, ValidLengthType: OptionalPointer, _is_cache_length_accurate: Bool = False, ragged: Bool = False]

Provides the shared softmax, correction, and store logic for SM100 MLA decode kernels.

Parameters

• ​q_type (DType): Element type of the Q and P operands stored in SMEM.
• ​KVLUTType (MHAOperand): MHAOperand providing the KV cache tensor and its element type.
• ​output_dtype (DType): Element type of the output tile written back via TMA (must be bfloat16).
• ​SplitAccumType (OptionalPointer): OptionalPointer type wrapping the split-K LSE accumulator buffer.
• ​MaskType (MHAMask): MHAMask type applied to the attention scores.
• ​config (MLA_SM100_Decode_Config): Decode config supplying tile sizes, swizzle modes, and TMEM/SMEM layout.
• ​ValidLengthType (OptionalPointer): OptionalPointer type wrapping the per-batch valid-sequence-length tensor.
• ​_is_cache_length_accurate (Bool): When False, the kernel adds the local sequence length to the cache length to compute the total key count (defaults to False).
• ​ragged (Bool): When True, the valid-lengths tensor is interpreted as input row offsets enabling ragged batching (defaults to False).

Implemented traits

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

comptime members

AccumType

comptime AccumType = get_accum_type[q_type]()

BlockElems

comptime BlockElems = (config * config)

bytes_per_element

comptime bytes_per_element = size_of[q_type]()

kv_type

comptime kv_type = KVLUTType.dtype

KVStageElems

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

NumQKBlocks

comptime NumQKBlocks = (config // config)

NumVOBlocks

comptime NumVOBlocks = (config // config)

O_M

comptime O_M = (config * Int(2))

O_N

comptime O_N = (config // Int(2))

output_tile_width

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

S_M

comptime S_M = (config * Int(2))

S_N

comptime S_N = (config // Int(2))

UMMAPVSS

comptime UMMAPVSS = DecodeSM100PVSS[q_type, MLA_SM100_Decode_Common[q_type, KVLUTType, output_dtype, SplitAccumType, MaskType, config, ValidLengthType, _is_cache_length_accurate, ragged].AccumType, config=config]

UMMAQKTSS

comptime UMMAQKTSS = DecodeSM100QKTSS[q_type, MLA_SM100_Decode_Common[q_type, KVLUTType, output_dtype, SplitAccumType, MaskType, config, ValidLengthType, _is_cache_length_accurate, ragged].AccumType, config=config]

Methods

pdl_early_exit

static def pdl_early_exit[fold_q: Bool = False](split_idx: Int, batch_idx: Int, max_seq_len: Int, out_row_offset: Int, batch_size: Int, lse_accum_split_ptr: SplitAccumType, o_tma: TMATensorTile[output_dtype, Int(2), IndexList(config, (config // Int(4)), __list_literal__=NoneType(None)), _default_desc_shape[Int(2), output_dtype, IndexList(config, (config // Int(4)), __list_literal__=NoneType(None)), config.swizzle_mode]()], seq_idx_fold: UInt32 = UInt32(0))

load_kv

static def load_kv(tma: TMATensorTile[Self.kv_type, Int(3), _padded_shape[Int(3), Self.kv_type, IndexList(config, Int(1), config, __list_literal__=NoneType(None)), config.kv_tma_swizzle_mode](), _ragged_shape[Int(3), Self.kv_type, IndexList(config, Int(1), config, __list_literal__=NoneType(None)), config.kv_tma_swizzle_mode]()], smem: Pointer[Scalar[Self.kv_type], MutAnyOrigin, address_space=AddressSpace.SHARED], mbar: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED], col_start: Int, row_start: Int)

load_q

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]()], smem: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED], mbar: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED], col_start: Int, row_start: Int)

apply_mask

static def apply_mask[half_load: Int, NonCausalMask: Bool, CausalMask: Bool, fold_q_num_heads: Int = Int(0), SlidingWindowSize: Int = Int(0), SparseCausalLogical: Bool = False](tiles_done: Int, col0: Int, num_keys: Int, s_row: TileTensor[Self.AccumType, Layout[*(), *()], MutUntrackedOrigin, address_space=AddressSpace.LOCAL], mask: MaskType, prompt_idx: UInt32, q_head_idx: UInt32, score_row: UInt32, cache_len: Int, start_pos: UInt32, cache_start_pos: UInt32, kv_start_row: Int = Int(0), logical_indices: OptionalReg[Pointer[Int32, MutAnyOrigin]] = None, logical_indices_stride: Int = Int(0), q_token_idx: Int = Int(0), logical_indices_len: Int = Int(0)) -> Scalar[Self.AccumType]

Returns:

Scalar[Self.AccumType]

Softmax

static def Softmax[native_fp8: Bool = False, num_sp_stages: Int = Int(2), fp8_p_stage_stride: Int = Int(0), has_per_token_scales: Bool = False, has_attn_sink: Bool = False, _op_sparse: Bool = False, _op_has_extra_kv: Bool = False, _op_has_variable_topk: Bool = False, fold_q: Bool = False, q_len_fold: Int = Int(1)](tmem_addr: UInt32, s_bars: DecodeSM100MiscMBars[num_sp_stages, Int(1), _resolve_warpgroup_size()], p_bars: DecodeSM100MiscMBars[num_sp_stages, _resolve_warpgroup_size(), Int(1)], p_smem_ptr: Pointer[Scalar[q_type], MutAnyOrigin, address_space=AddressSpace.SHARED], max_smem: Pointer[Scalar[Self.AccumType], MutAnyOrigin, address_space=AddressSpace.SHARED], li_smem: Pointer[Scalar[Self.AccumType], MutAnyOrigin, address_space=AddressSpace.SHARED], out_smem: Pointer[Scalar[output_dtype], MutAnyOrigin, address_space=AddressSpace.SHARED], c_bars: DecodeSM100MiscMBars[Int(1), _resolve_warpgroup_size(), _resolve_warpgroup_size()], corr_done_bars: DecodeSM100MiscMBars[Int(2), _resolve_warpgroup_size(), _resolve_warpgroup_size()], out_pipeline: OutPipeline[((config // (config // Int(4))) // Int(2) if (xor (eq ((config // Int(2)) // (config // Int(4))), 0), True) else Int(1)), _resolve_warpgroup_size(), Int(1)], offset_position: OffsetPosition[config, KVLUTType, ragged, _is_cache_length_accurate, ValidLengthType, config.decoding_warp_split_k, _op_sparse, _op_has_extra_kv, _op_has_variable_topk], scale: Float32, mask: MaskType, prompt_idx: UInt32, lse_accum_split_ptr: SplitAccumType, batch_size: Int, scale_k_smem: OptionalReg[Pointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED]] = None, q_scale_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin]] = None, attn_sink_log2: Float32 = min_or_neg_inf[DType.float32](), logical_indices: OptionalReg[Pointer[Int32, MutAnyOrigin]] = None, logical_indices_stride: Int = Int(0), logical_indices_len: Int = Int(0))

Correction

static def Correction[_op_sparse: Bool = False, _op_has_extra_kv: Bool = False, _op_has_variable_topk: Bool = False](tmem_addr: UInt32, o_bars: DecodeSM100MiscMBars[Int(2), Int(1), _resolve_warpgroup_size()], c_bars: DecodeSM100MiscMBars[Int(1), _resolve_warpgroup_size(), _resolve_warpgroup_size()], corr_done_bars: DecodeSM100MiscMBars[Int(2), _resolve_warpgroup_size(), _resolve_warpgroup_size()], offset_position: OffsetPosition[config, KVLUTType, ragged, _is_cache_length_accurate, ValidLengthType, config.decoding_warp_split_k, _op_sparse, _op_has_extra_kv, _op_has_variable_topk])

store

static def store[_op_sparse: Bool = False, _op_has_extra_kv: Bool = False, _op_has_variable_topk: Bool = False, fold_q: Bool = False, q_len_fold: Int = Int(1)](out_pipeline: OutPipeline[((config // (config // Int(4))) // Int(2) if (xor (eq ((config // Int(2)) // (config // Int(4))), 0), True) else Int(1)), _resolve_warpgroup_size(), Int(1)], out_smem: Pointer[Scalar[output_dtype], MutAnyOrigin, address_space=AddressSpace.SHARED], o_tma: TMATensorTile[output_dtype, Int(2), IndexList(config, (config // Int(4)), __list_literal__=NoneType(None)), _default_desc_shape[Int(2), output_dtype, IndexList(config, (config // Int(4)), __list_literal__=NoneType(None)), config.swizzle_mode]()], offset_position: OffsetPosition[config, KVLUTType, ragged, _is_cache_length_accurate, ValidLengthType, config.decoding_warp_split_k, _op_sparse, _op_has_extra_kv, _op_has_variable_topk])