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 package
gpu
Provides the GPU matmul dispatch entry point and hardware-specific kernel selection.
comptime values
logger
comptime logger = Logger(stdout, prefix=String(""), source_location=False)
Packages
-
amd: Provides the AMD GPU backend implementations for matmuls. -
amd_rdna: Provides the AMD RDNA GPU backend implementations for matmuls. -
apple: Provides the Apple silicon GPU backend implementations for matmuls. -
sm100: Provides the Nvidia Blackwell backend implementations for matmuls. -
sm100_structured: SM100 Structured Kernels - Blackwell matmul implementation. -
sm80: Provides the CPU Hopper backend implementations for matmuls. -
sm90: Provides the Nvidia Hopper backend implementations for matmuls.
Modules
-
profiler: Provides warp-level profiling infrastructure for Blackwell warp-specialized matmul kernels. -
tile_scheduler: Provides the persistent tile scheduler for GPU matmul kernels. -
tile_scheduler_splitk: Provides the split-K tile scheduler for GPU matmul kernels on SM90.
Functions
-
matmul_kernel: Matrix Multiplication using shared memory. This version loads blocks of size tile_size x tile_size from A and B and updates a tile_size x tile_size in C. The thread block should have shape (tile_size, tile_size, 1). Each thread is mapped one element in C. The grid should have shape (N/tile_size, M/tile_size, 1). N is the first dimension for coalesced access. -
matmul_kernel_naive: Computes one output element of a matrix multiply per thread by iterating over the K dimension. -
multistage_gemm: TileTensor overload ofmultistage_gemm. Converts to LayoutTensor and dispatches to a GEMM kernel. -
split_k_reduce: Reduces a split-K workspace into the output tensor by summing across K partitions.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!