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

bmm_sm100_blockwise_scaled_fp8

def bmm_sm100_blockwise_scaled_fp8[c_type: DType, a_type: DType, b_type: DType, a_scales_type: DType, b_scales_type: DType, *, 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[c_type: DType, width: SIMDLength, rank: Int, *, alignment: Int = Int(1)](IndexList[rank], SIMD[c_type, width]) capturing thin -> None] = None, b_scaling_block_n: Int = Int(128)](c_: TileTensor[c_type, Storage=c_.Storage, address_space=c_.address_space, linear_idx_type=c_.linear_idx_type], a_: TileTensor[a_type, Storage=a_.Storage, address_space=a_.address_space, linear_idx_type=a_.linear_idx_type], b_: TileTensor[b_type, Storage=b_.Storage, address_space=b_.address_space, linear_idx_type=b_.linear_idx_type], a_scales_: TileTensor[a_scales_type, Storage=a_scales_.Storage, address_space=a_scales_.address_space, linear_idx_type=a_scales_.linear_idx_type], b_scales_: TileTensor[b_scales_type, Storage=b_scales_.Storage, address_space=b_scales_.address_space, linear_idx_type=b_scales_.linear_idx_type], ctx: DeviceContext)

Computes a batched blockwise scaled FP8 matrix multiplication on SM100 (Blackwell) hardware by constructing TMA descriptors for the inputs and scales and enqueuing the blockwise FP8 kernel per batch slice.

Parameters:

c_type (DType): Output tensor element dtype.
a_type (DType): LHS input tensor element dtype.
b_type (DType): RHS input tensor element dtype.
a_scales_type (DType): LHS scales tensor element dtype.
b_scales_type (DType): RHS scales tensor element dtype.
transpose_b (Bool): Whether the RHS input is transposed (must be True).
umma_shape (IndexList[Int(3)]): UMMA instruction shape (m, n, k) used by the kernel.
block_tile_shape (IndexList[Int(3)]): CTA tile shape (BM, BN, BK).
a_swizzle (TensorMapSwizzle): TMA swizzle mode for the LHS input tensor.
b_swizzle (TensorMapSwizzle): TMA swizzle mode for the RHS input tensor.
elementwise_lambda_fn (Optional[def[c_type: DType, width: SIMDLength, rank: Int, *, alignment: Int = Int(1)](IndexList[rank], SIMD[c_type, width]) capturing thin -> None]): Optional epilogue applied to each output element.
b_scaling_block_n (Int): N-direction scale block size for the RHS scales.

Args:

c_ (TileTensor[c_type, Storage=c_.Storage, address_space=c_.address_space, linear_idx_type=c_.linear_idx_type]): Rank-3 output tensor of shape (batch, m, n).
a_ (TileTensor[a_type, Storage=a_.Storage, address_space=a_.address_space, linear_idx_type=a_.linear_idx_type]): Rank-3 LHS input tensor of shape (batch, m, k).
b_ (TileTensor[b_type, Storage=b_.Storage, address_space=b_.address_space, linear_idx_type=b_.linear_idx_type]): Rank-3 RHS input tensor of shape (batch, k, n).
a_scales_ (TileTensor[a_scales_type, Storage=a_scales_.Storage, address_space=a_scales_.address_space, linear_idx_type=a_scales_.linear_idx_type]): Rank-3 LHS scales tensor.
b_scales_ (TileTensor[b_scales_type, Storage=b_scales_.Storage, address_space=b_scales_.address_space, linear_idx_type=b_scales_.linear_idx_type]): Rank-3 RHS scales tensor.
ctx (DeviceContext): Device context used to enqueue the kernel.