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

block_scaled_matmul

def block_scaled_matmul[c_type: DType, a_type: DType, b_type: DType, scales_dtype: DType, //, *, SF_VECTOR_SIZE: Int, transpose_b: Bool = True, transpose_a: Bool = False, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, elementwise_compute_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]] = None, pdl_level: PDLLevel = PDLLevel.OFF, _trace_description: StringSlice[ImmStaticOrigin] = StringSlice(""), target: StringSlice[ImmStaticOrigin] = StringSlice("cpu")](c_device: TileTensor[c_type, Storage=c_device.Storage, linear_idx_type=c_device.linear_idx_type], a_device: TileTensor[a_type, Storage=a_device.Storage, linear_idx_type=a_device.linear_idx_type], b_device: TileTensor[b_type, Storage=b_device.Storage, linear_idx_type=b_device.linear_idx_type], a_scales_device: TileTensor[scales_dtype, Storage=a_scales_device.Storage, linear_idx_type=a_scales_device.linear_idx_type], b_scales_device: TileTensor[scales_dtype, Storage=b_scales_device.Storage, linear_idx_type=b_scales_device.linear_idx_type], tensor_sf: Float32, ctx: DeviceContext)

Dispatches the SM100 block-scaled matmul for NVFP4 or MXFP8, selecting between the Mojo heuristic kernel and the vendor fallback.

For small M or MXFP8 shapes the Mojo heuristic_and_outliers_dispatch kernel is tried first; on a dispatch miss the vendor (cuBLASLt) block-scaled matmul is used as a fallback.

Parameters:

c_type (DType): Element type of the output matrix (inferred).
a_type (DType): Element type of the LHS input matrix (inferred).
b_type (DType): Element type of the RHS input matrix (inferred).
scales_dtype (DType): Element type of the block scale-factor tensors (inferred).
SF_VECTOR_SIZE (Int): Number of elements covered by each block scale factor: 16 for NVFP4 or 32 for MXFP8.
transpose_b (Bool): Whether b_device is stored transposed (defaults to True).
transpose_a (Bool): Whether a_device is stored transposed (defaults to False).
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional epilogue lambda applied to the matmul result after it is written to c_device (defaults to None).
elementwise_compute_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]]): Optional compute lambda that transforms the matmul result in-flight before storage (defaults to None).
pdl_level (PDLLevel): Programmatic dependency launch level for the heuristic dispatch kernel (defaults to PDLLevel.OFF).
_trace_description (StringSlice[ImmStaticOrigin]): Extra label appended to the trace event for profiling (defaults to "").
target (StringSlice[ImmStaticOrigin]): Trace target string identifying the device in profiling output (defaults to "cpu").