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_sm100

Provides warp-specialized persistent grouped GEMM kernels for SM100 (B200) GPUs.

Structs

• ​WarpRole: Enumerates the role each warp plays in the warp-specialized kernel.

Functions

• ​blackwell_tma_umma_warp_specialized_kernel: Implements the warp-specialized persistent grouped GEMM kernel for SM100.
• ​consumer_main_loop: Performs the MMA consumer step for one pipeline stage.
• ​grouped_matmul_sm100_persistent: Launches the persistent grouped GEMM kernel for SM100 from host tensors.
• ​load_AB: Loads A and B tiles from global memory into shared memory via TMA multicast.
• ​load_AB_cuda_core: CUDA core fallback for load_AB when K*sizeof < 16 bytes.
• ​multi_stage_store_C: Drains accumulated results from tensor memory and stores them to global memory.
• ​stsm_helper: Stores a register fragment to shared memory using the stmatrix instruction.
• ​zero_output: Zero-fills an output tile for skipped or invalid expert assignments.
• ​zero_output_epilogue: Zero an inactive group's output THROUGH the elementwise epilogue.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit