Mojo struct
SM100MHA2Q
@register_passable(trivial)
struct SM100MHA2Q[KVLUTType: MHAOperand, output_type: DType, MaskType: MHAMask, ScoreModType: ScoreModTrait, SchedulerType: MHATileScheduler, config: FA4Config, use_score_mod: Bool, ValidLengthType: OptionalPointer, SinkType: OptionalPointer, KVRowOffsetsType: OptionalPointer, _is_cache_length_accurate: Bool, MaxSeqLenType: OptionallyStaticInt, PartitionType: MHAPartitionScheme, descriptor_shape: IndexList[3], remaining_global_dim_rank: Int]
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
accum_type
comptime accum_type = get_accum_type[SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_type]()
BM
comptime BM = config.BM
BN
comptime BN = config.BN
cta_group
comptime cta_group = 1
depth
comptime depth = config.depth
group
comptime group = config.group
k_bytes
comptime k_bytes = SIMD[DType.uint32, 1]((SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].swizzle_granularity * config)).__rmul__[DType.uint32, 1](SIMD[DType.uint32, 1](SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_dt_size))
k_elements
comptime k_elements = SIMD[DType.uint32, 1]((SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].swizzle_granularity * config))
KVPipelineType
comptime KVPipelineType = KVPipeline[config.num_kv_stages, config.num_mma_stages]
MMA_K
comptime MMA_K = 16
MMA_M
comptime MMA_M = (config // 2)
num_m_mmas
comptime num_m_mmas = 2
num_mma_stages
comptime num_mma_stages = config.num_mma_stages
num_q_heads
comptime num_q_heads = config.num_q_heads
OPipelineType
comptime OPipelineType = MBarPipeline[2]
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, _is_decoding[MaxSeqLenType]()]
qkv_dt_size
comptime qkv_dt_size = size_of[SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_type]()
qkv_type
comptime qkv_type = KVLUTType.dtype
qo_bytes
comptime qo_bytes = SIMD[DType.uint32, 1]((SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_dt_size * SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qo_elements))
qo_elements
comptime qo_elements = (SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].padded_depth * SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].MMA_M)
ragged
comptime ragged = ValidLengthType.is_null.__bool__().__invert__()
simd_size
comptime simd_size = simd_width_of[SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_type]()
swizzle_granularity
comptime swizzle_granularity = (config.swizzle_mode.bytes() // SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_dt_size)
UMMA0Type
comptime UMMA0Type = SM100TensorAccumulatorSS[SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_type, SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].accum_type, SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].MMA_M, SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].BN, SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].depth, swizzle_a=config.swizzle_mode, swizzle_b=config.swizzle_mode, num_stages=SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].num_mma_stages]
UMMA1Type
comptime UMMA1Type = SM100TensorAccumulatorTS[SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_type, SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].accum_type, SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].MMA_M, config.padded_depth, SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].BN, config.swizzle_mode, transpose_b=False, num_stages=SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].num_mma_stages]
v_bytes_per_mma
comptime v_bytes_per_mma = SIMD[DType.uint32, 1](((SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_dt_size * 16) * config))
Methods
kernel
static kernel(q_tma_op: TMATensorTile[KVLUTType.dtype, tile_layout_k_major[KVLUTType.dtype, (config // 2), config.BK0, config.swizzle_mode](), _tma_desc_tile_layout[KVLUTType.dtype, 2, IndexList[2, DType.int64]((config // 2), config.BK0, Tuple[]()), swizzle_mode=config.swizzle_mode]()], k_tma_op: TMATensorTile[KVLUTType.dtype, tile_layout_k_major[KVLUTType.dtype, config.BN, config.padded_depth, config.swizzle_mode](), _tma_desc_tile_layout[KVLUTType.dtype, 2, IndexList[2, DType.int64](config.BN, config.padded_depth, Tuple[]()), swizzle_mode=config.swizzle_mode]()], v_tma_op: TMATensorTile[KVLUTType.dtype, tile_layout_mn_major[KVLUTType.dtype, config.padded_depth, config.BN, config.swizzle_mode](), _tma_desc_tile_layout[KVLUTType.dtype, 2, IndexList[2, DType.int64](config.BN, config.padded_depth, Tuple[]()), False, config.swizzle_mode](), False], o_ptr_arg: LegacyUnsafePointer[Scalar[output_type]], ragged_tma_store: RaggedTensorMap[output_type, descriptor_shape, remaining_global_dim_rank, config.swizzle_mode], kv_lut: KVLUTType, scale: Float32, batch_size: UInt32, num_keys_arg: UInt32, pack: Pack[MaskType, ScoreModType, SchedulerType, ValidLengthType, SinkType, KVRowOffsetsType, MaxSeqLenType, PartitionType])
mask_status
static mask_status(mask: MaskType, score_row: UInt32, kv_row: UInt32) -> TileMaskStatus
Returns:
TileMaskStatus
scale_write_output
static scale_write_output(local_row: UInt32, inv_row_sum: Scalar[SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].accum_type], o_smem: LegacyUnsafePointer[Scalar[output_type], address_space=AddressSpace.SHARED], o_tmem: TMemTile[SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].accum_type, (SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].BM // 2), SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].padded_depth], o_ptr: LegacyUnsafePointer[Scalar[output_type]], ragged_tma_store: RaggedTensorMap[output_type, descriptor_shape, remaining_global_dim_rank, config.swizzle_mode], warp_group_idx: UInt32, consumer_mbar: LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED], current_seq: Int, num_output_rows: Int32)
softmax
static softmax(tmem_addr: UInt32, warp_idx: UInt32, mbars: FA4MiscMBars, o_mbar: LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED], score_row: UInt32, seq_info: SeqInfo, mask: MaskType, num_keys: UInt32, scale: Scalar[SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].accum_type], score_mod: ScoreModType, max_seq_len: UInt32, o_ptr_arg: LegacyUnsafePointer[Scalar[output_type]], ragged_tma_store: RaggedTensorMap[output_type, descriptor_shape, remaining_global_dim_rank, config.swizzle_mode], o_smem: LegacyUnsafePointer[Scalar[output_type], address_space=AddressSpace.SHARED], sink_weights: SinkType)
correction
static correction(tmem_addr: UInt32, mbars: FA4MiscMBars, o_mbar: LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED], score_row: UInt32, num_keys: UInt32, mask: MaskType)
load
static load(mbars: FA4MiscMBars, kv_pipeline_arg: KVPipeline[config.num_kv_stages, config.num_mma_stages], score_row: UInt32, num_keys: UInt32, seq_info: SeqInfo, max_seq_len: MaxSeqLenType, mask: MaskType, q_tma_op: TMATensorTile[KVLUTType.dtype, tile_layout_k_major[KVLUTType.dtype, (config // 2), config.BK0, config.swizzle_mode](), _tma_desc_tile_layout[KVLUTType.dtype, 2, IndexList[2, DType.int64]((config // 2), config.BK0, Tuple[]()), swizzle_mode=config.swizzle_mode]()], k_tma_op: TMATensorTile[KVLUTType.dtype, tile_layout_k_major[KVLUTType.dtype, config.BN, config.padded_depth, config.swizzle_mode](), _tma_desc_tile_layout[KVLUTType.dtype, 2, IndexList[2, DType.int64](config.BN, config.padded_depth, Tuple[]()), swizzle_mode=config.swizzle_mode]()], v_tma_op: TMATensorTile[KVLUTType.dtype, tile_layout_mn_major[KVLUTType.dtype, config.padded_depth, config.BN, config.swizzle_mode](), _tma_desc_tile_layout[KVLUTType.dtype, 2, IndexList[2, DType.int64](config.BN, config.padded_depth, Tuple[]()), False, config.swizzle_mode](), False], kv_lut: KVLUTType, q_smem: LegacyUnsafePointer[Scalar[KVLUTType.dtype], address_space=AddressSpace.SHARED])
descriptor_q
static descriptor_q(q_smem: LegacyUnsafePointer[Scalar[SM100MHA2Q[KVLUTType, output_type, MaskType, ScoreModType, SchedulerType, config, use_score_mod, ValidLengthType, SinkType, KVRowOffsetsType, _is_cache_length_accurate, MaxSeqLenType, PartitionType, descriptor_shape, remaining_global_dim_rank].qkv_type], address_space=AddressSpace.SHARED]) -> MMASmemDescriptorPair
Returns:
mma
static mma(tmem_addr: UInt32, mbars: FA4MiscMBars, kv_pipeline_arg: KVPipeline[config.num_kv_stages, config.num_mma_stages], o_mbar: LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED], score_row: UInt32, num_keys: UInt32, mask: MaskType, q_smem: LegacyUnsafePointer[Scalar[KVLUTType.dtype], address_space=AddressSpace.SHARED])
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!