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

grouped_matmul

Provides grouped matrix multiplication kernels for CPU, AMD, and NVIDIA GPU targets.

Functions

dispatch_amd_matmul_by_block_shape: Dispatches to the best kernel configuration based on runtime M dimension.
grouped_matmul: TileTensor implementation of grouped_matmul.
grouped_matmul_amd: Launches the AMD grouped matmul kernel, selecting the best block-tile configuration for the runtime M dimension via dispatch_amd_matmul_by_block_shape.
grouped_matmul_amd_kernel_launcher: Computes the AMD GPU grouped matmul by dispatching per-expert tiles through AMDMatmul, with separate zero-fill handling for inactive (expert_id == -1) blocks.
grouped_matmul_kernel_sm100: Computes the SM100 (Blackwell) grouped matmul using TMA async copies into shared memory and tcgen05 tensor-memory MMA accumulation, with an optional elementwise epilogue.
grouped_matmul_rowwise_dynamic_scaled_fp8: Grouped (ragged MoE) FP8 matmul with rowwise weight + per-token act scales.
grouped_matmul_rowwise_scaled_fp8_kernel: Computes the naive grouped FP8 matmul with rowwise weight scales and per-token activation scales, accumulating in fp32 and applying a single post-reduction scale.
grouped_matmul_sm100: Launches the SM100 grouped matmul kernel with TMA descriptors for ragged MoE matrix multiplication on Blackwell GPUs.
grouped_matmul_vendor: TileTensor primary implementation of grouped_matmul_vendor.
naive_epilogue: Launches the naive_epilogue_kernel to apply an elementwise epilogue function across the output tensor c.
naive_epilogue_kernel: Applies an elementwise epilogue function to each SIMD-vectorized element of the output tensor c.
naive_grouped_matmul: TileTensor primary implementation of naive_grouped_matmul.
naive_grouped_matmul_kernel: Computes one element per thread of the grouped matmul product C[a_offsets[z]:a_offsets[z+1], :] = A[...] @ B[expert_ids[z], :, :].T for each active expert z, with an optional elementwise epilogue.

Functions​

Functions