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_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, ImplicitlyCopyable, ImplicitlyDeletable, 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, _safe=False], mbar: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], 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, _safe=False], mbar: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], 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)](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)) -> 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, _safe=False], max_smem: Pointer[Scalar[Self.AccumType], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], li_smem: Pointer[Scalar[Self.AccumType], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], out_smem: Pointer[Scalar[output_dtype], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], 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, _safe=False]] = None, q_scale_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]] = None, attn_sink_log2: Float32 = min_or_neg_inf[DType.float32]())

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, _safe=False], 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])

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact