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 module
bmm
Provides batched matrix multiplication (BMM) for CPU and GPU targets.
comptime values
elementwise_epilogue_type
comptime elementwise_epilogue_type = def[c_type: DType, width: SIMDLength, rank: Int, *, alignment: Int = Int(1)](IndexList[rank], SIMD[c_type, width]) capturing thin -> None
logger
comptime logger = Logger(stdout, prefix=String(""), source_location=False)
Functions
-
batched_matmul: TileTensor primary implementation ofbatched_matmul. -
batched_matmul_dynamic_scaled_fp8: Dispatches a batched blockwise scaled FP8 matrix multiplication to the SM100 blockwise kernel on Blackwell hardware or falls back to the naive per-batch implementation on H100. -
batched_matmul_dynamic_scaled_fp8_naive: Computes a batched blockwise scaled FP8 matrix multiplication using a naive per-batch loop that calls the 2D blockwise scaled FP8 kernel for each batch slice. -
batched_matmul_kernel_gpu: Computes a single batch slice of a batched matrix multiplication on the GPU by dispatching to the multistage GEMM kernel on NVIDIA or theAMDMatmulkernel on AMD hardware. -
batched_matmul_shape: Compute the output shape of abatch_matmuloperation, and assert the inputs are compatible. -
bmm_sm100_blockwise_scaled_fp8: 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. -
naive_batched_matmul_kernel: Computes one element per thread of a batched matrix multiplication using a naive scalar accumulation loop over the contraction dimension.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!