Mojo struct

SM100MLA

struct SM100MLA[KVLUTType: MHAOperand, KRopeType: MHAOperand, output_type: DType, MaskType: MHAMask, SchedulerType: MHATileScheduler, config: MLAConfig, ValidLengthType: OptionalPointer, SinkType: OptionalPointer, KVRowOffsetsType: OptionalPointer, MaxSeqLenType: OptionallyStaticInt, PartitionType: MHAPartitionScheme, _ndbuffer_mha_operand: Bool]

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

accum_type

comptime accum_type = DType.float32

BK0

comptime BK0 = SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].depth if (SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].qkv_type == DType.bfloat16) if (KRopeType.dtype == DType.float8_e4m3fn) else (KRopeType.dtype == DType.float8_e4m3fn) or (KRopeType.dtype == SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].qkv_type) else SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].nope_depth

BM

comptime BM = config.BM

BN

comptime BN = config.BN

cache_depth

comptime cache_depth = config.cache_depth

cta_group

comptime cta_group = 1

depth

comptime depth = config.depth

group

comptime group = config.group

KVPipelineType

comptime KVPipelineType = StagedPipeline[config.num_kv_stages, config.num_qk_stages]

MiscMBarsType

comptime MiscMBarsType = FA4MiscMBars[num_qk_stages=config.num_qk_stages, num_pv_stages=config.num_pv_stages, num_kv_stages=config.num_kv_stages, separate_kv=False]

MMA_M

comptime MMA_M = (config // 2)

nope_depth

comptime nope_depth = config.nope_depth

nope_mma_kind

comptime nope_mma_kind = UMMAKind.KIND_F16 if SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].qkv_type.is_half_float() else UMMAKind.KIND_F8F6F4

num_m_mmas

comptime num_m_mmas = 2

num_q_heads

comptime num_q_heads = config.num_q_heads

num_qk_stages

comptime num_qk_stages = config.num_qk_stages

padded_depth

comptime padded_depth = config.padded_depth

page_size

comptime page_size = KVLUTType.page_size

PositionType

comptime PositionType = MHAPosition[config.BM, config.BN, config.depth, config.padded_depth, config.num_q_heads, config.group, False]

qkv_dt_size

comptime qkv_dt_size = size_of[SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].qkv_type]()

qkv_type

comptime qkv_type = KVLUTType.dtype

rope_depth

comptime rope_depth = config.rope_depth

rope_mma_kind

comptime rope_mma_kind = UMMAKind.KIND_F16 if KRopeType.dtype.is_half_float() else UMMAKind.KIND_F8F6F4

simd_size

comptime simd_size = simd_width_of[SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].qkv_type]()

UMMA0RopeType

comptime UMMA0RopeType = SM100TensorAccumulatorSS[KRopeType.dtype, DType.float32, SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].MMA_M, SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].BN, SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].rope_depth, mma_kind=SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].rope_mma_kind, swizzle_a=config.rope_swizzle_mode, swizzle_b=config.rope_swizzle_mode, num_stages=SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].num_qk_stages]

UMMA0Type

comptime UMMA0Type = SM100TensorAccumulatorSS[SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].qkv_type, DType.float32, SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].MMA_M, SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].BN, SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].BK0, mma_kind=SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].nope_mma_kind, swizzle_a=config.qkv_swizzle_mode, swizzle_b=config.qkv_swizzle_mode, num_stages=SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].num_qk_stages]

UMMA1Type

comptime UMMA1Type = SM100TensorAccumulatorTS[SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].qkv_type, DType.float32, SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].MMA_M, SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].nope_depth, SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].BN, config.qkv_swizzle_mode, mma_kind=SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].nope_mma_kind, transpose_b=False, num_stages=SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].num_qk_stages]

Methods

softmax

static softmax[apply_scale: Bool = False](tmem_addr: UInt32, warp_idx: UInt32, mbars: FA4MiscMBars[num_qk_stages=config.num_qk_stages, num_pv_stages=config.num_pv_stages, num_kv_stages=config.num_kv_stages, separate_kv=False], score_row: UInt32, seq_info: SeqInfo, mask: MaskType, num_keys: UInt32, scale: Float32, max_seq_len: UInt32, ragged_tma_store: RaggedTMA3DTile[output_type, config.output_swizzle_mode, (config // 2), SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].nope_depth], o_smem: UnsafePointer[Scalar[output_type], MutAnyOrigin, address_space=AddressSpace.SHARED], correction_smem_arg: UnsafePointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED], q_scale_smem_arg: OptionalReg[UnsafePointer[Scalar[KVLUTType.scale_dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]] = None, k_scale_smem_arg: OptionalReg[UnsafePointer[Scalar[KVLUTType.scale_dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]] = None)

correction

static correction(tmem_addr: UInt32, mbars: FA4MiscMBars[num_qk_stages=config.num_qk_stages, num_pv_stages=config.num_pv_stages, num_kv_stages=config.num_kv_stages, separate_kv=False], score_row: UInt32, num_keys: UInt32, mask: MaskType, correction_smem_arg: UnsafePointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED])

mask_status

static mask_status(mask: MaskType, score_row: UInt32, kv_row: UInt32) -> TileMaskStatus

Returns:

TileMaskStatus

descriptor_q

static descriptor_q(q_smem: UnsafePointer[Scalar[SM100MLA[KVLUTType, KRopeType, output_type, MaskType, SchedulerType, config, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType, _ndbuffer_mha_operand].qkv_type], MutAnyOrigin, address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair

Returns:

MMASmemDescriptorPair

descriptor_q_rope

static descriptor_q_rope(q_smem: UnsafePointer[Scalar[KRopeType.dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair

Returns:

MMASmemDescriptorPair

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