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 function

matmul_sm100_blockwise_scaled_fp8

def matmul_sm100_blockwise_scaled_fp8[*, transpose_b: Bool, umma_shape: IndexList[Int(3)], block_tile_shape: IndexList[Int(3)], a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c: TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type], a_scales: TileTensor[Storage=a_scales.Storage, address_space=a_scales.address_space, linear_idx_type=a_scales.linear_idx_type], b_scales: TileTensor[Storage=b_scales.Storage, address_space=b_scales.address_space, linear_idx_type=b_scales.linear_idx_type], ctx: DeviceContext)

Enqueues a blockwise scaled FP8 GEMM on SM100 with 1D A-scales and 2D B-scales onto the supplied device context.

Validates the operand dtypes, ranks, and scale granularities, builds the TMA tile descriptors for A, B, and A-scales, and launches the matmul_sm100_blockwise_scaled_fp8_1d2d_wrapper kernel over a two-dimensional grid covering the M and N dimensions.

Parameters:

transpose_b (Bool): Whether B is loaded transposed; must be True (only the transposed-B path is supported).
umma_shape (IndexList[Int(3)]): 3D MMA instruction shape (MMA_M, MMA_N, MMA_K) used by the TCgen05 tensor core operations.
block_tile_shape (IndexList[Int(3)]): 3D block tile shape (BM, BN, BK) giving the per-CTA tile dimensions along the M, N, and K axes. BK must be 64 or 128.
a_swizzle (TensorMapSwizzle): TMA swizzle mode applied to A shared memory loads (defaults to TensorMapSwizzle.SWIZZLE_128B).
b_swizzle (TensorMapSwizzle): TMA swizzle mode applied to B shared memory loads (defaults to TensorMapSwizzle.SWIZZLE_128B).
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional elementwise epilogue lambda applied to each output element of c in place of a direct store (defaults to None).

Args:

c (TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Rank-2 output TileTensor accumulating the scaled GEMM result.
a (TileTensor[Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Rank-2 TileTensor of A operands with float8_e4m3fn elements.
b (TileTensor[Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type]): Rank-2 TileTensor of B operands with float8_e4m3fn elements.
a_scales (TileTensor[Storage=a_scales.Storage, address_space=a_scales.address_space, linear_idx_type=a_scales.linear_idx_type]): Rank-2 TileTensor of A scale factors with float32 elements and 1D scaling granularity along K.
b_scales (TileTensor[Storage=b_scales.Storage, address_space=b_scales.address_space, linear_idx_type=b_scales.linear_idx_type]): Rank-2 TileTensor of B scale factors with float32 elements and 2D scaling granularity along N and K.
ctx (DeviceContext): Device context used to enqueue the kernel launch.