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
block_scaled_matmul_small_bn
Implements the small-BN block-scaled matmul kernel for SM100 (Blackwell B200).
Provides the warp-specialized TMA + UMMA kernel and its host-side launch helpers for block-scaled (MXFP8 / NVFP4) matrix multiplication targeting the small-N MMA shapes that require a cp.async SFB load path instead of TMA.
Structsβ
- β
B200BlockScaledMatmulSmem: Owns the shared-memory buffers and mbarrier pools for one kernel block. - β
WarpRole: Encodes the warp role assignments used by the warp-specialized kernel.
Functionsβ
- β
blackwell_block_scaled_matmul_tma_umma_warp_specialized: Launch small-BN block-scaled FP8 matmul kernel on SM100. - β
blackwell_block_scaled_tma_umma_warp_specialized_kernel: Warp-specialized TMA + UMMA kernel for small-BN block-scaled matmul on SM100. - β
consumer_main_loop: TileTensor overload of consumer_main_loop for block-scaled MMA (small BN). - β
load_AB_SFA: Issues multicast TMA loads for one K-group of A, B, and SFA into SMEM.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!