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
MmaOpSM100_BlockScaled_SS
struct MmaOpSM100_BlockScaled_SS[c_type: DType, a_type: DType, b_type: DType, sfa_dtype: DType, sfb_dtype: DType, scaling_kind: UMMAKind, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], /, *, accum_type: DType = DType.float32, cta_group: Int = Int(1), cluster_shape: IndexList[Int(3)] = Index[Int, Int, Int](Int(1), Int(1), Int(1)), a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, transpose_b: Bool = False, enable_small_sfb: Bool = False]
SM100 (Blackwell) warp-specialized MMA operation for block-scaled GEMM (MXFP4, MXFP8, NVFP4).
Extends MmaOpSM100_SS with per-block scale factor handling for quantized
formats. Copies scale factors from shared memory to TMEM via tcgen05_cp
and issues UMMA instructions with the appropriate scale index per K slice.
Parameters
- c_type (
DType): Output accumulator element type. - a_type (
DType): A operand element type (float8_e4m3fn or uint8 for FP4). - b_type (
DType): B operand element type; must matcha_type. - sfa_dtype (
DType): Scale factor dtype for A; must matchsfb_dtype. - sfb_dtype (
DType): Scale factor dtype for B. - scaling_kind (
UMMAKind): UMMA scaling kind (MXF8F6F4, MXF4, or MXF4NVF4). - block_tile_shape (
IndexList[Int(3)]): (BM, BN, BK) shape of the shared-memory tile. - mma_shape (
IndexList[Int(3)]): (MMA_M, MMA_N, MMA_K) shape of a single UMMA instruction. - accum_type (
DType): Accumulator precision; defaults to float32. - cta_group (
Int): Number of CTAs collaborating on one tile (1 or 2). - cluster_shape (
IndexList[Int(3)]): CTA cluster shape for multicast; defaults to (1, 1, 1). - a_swizzle (
TensorMapSwizzle): Swizzle mode for the A operand. - b_swizzle (
TensorMapSwizzle): Swizzle mode for the B operand. - transpose_b (
Bool): Must be True; SM100 UMMA always uses transposed B. - enable_small_sfb (
Bool): Allow SFB loading for MMA_N < 64 via cooperative tcgen05_st.
Fields
- idesc (
UMMAInsDescriptor[scaling_kind]): - mask (
UInt16):
Implemented traits
AnyType,
Copyable,
Defaultable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
Methods
__init__
def __init__() -> Self
mma
def mma(self, a: TileTensor[Storage=a.Storage, address_space=AddressSpace.SHARED, linear_idx_type=a.linear_idx_type], b: TileTensor[Storage=b.Storage, address_space=AddressSpace.SHARED, linear_idx_type=b.linear_idx_type], sfa_smem: TileTensor[Storage=sfa_smem.Storage, address_space=AddressSpace.SHARED, linear_idx_type=sfa_smem.linear_idx_type], sfb_smem: TileTensor[Storage=sfb_smem.Storage, address_space=AddressSpace.SHARED, linear_idx_type=sfb_smem.linear_idx_type], c_tmem: UInt32, sfa_tmem: UInt32, sfb_tmem: UInt32, init_c: Bool, sfb_tmem_adj: UInt32 = UInt32(0))
TileTensor overload for block-scaled MMA input tiles.
Creates MMA descriptors directly from swizzle parameters.
commit
def commit(self, ptr_mbar: Pointer[address_space=AddressSpace.SHARED, _safe=False])
wait
def wait(self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!