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
matmul_swiglu_kernels
SM100 fused GEMM+SwiGLU kernel (monolithic warp-specialized).
Follows the same monolithic structure as the BF16/FP8 dense kernel: free
functions for each phase (load, MMA, epilogue) and a single top-level kernel
function that dispatches to them via WarpRole checks.
Two epilogue modes controlled by register_swiglu:
- False (default): TMEM → FP32 SwiGLU → half SMEM → TMA store to GMEM.
- True: TMEM → FP32 SwiGLU → BF16 stores directly to GMEM.
Weight layout: natural row-interleaved (gate, up) pairs in both swap modes.
- non-swap (AB_swapped=False): (gate, up) at adjacent N positions.
- swap (AB_swapped=True): (gate, up) at adjacent M positions; warp.shuffle_xor(_, 4) brings the up value into the gate-owner lane.
Structs
-
SwiGLUKernelConstants: Compile-time constants for TMA descriptor creation and kernel launch.
Functions
-
blackwell_swiglu_warp_specialized_kernel: Fused GEMM+SwiGLU warp-specialized kernel for SM100. -
consumer_main_loop: Consumer main loop for BF16 MMA on SM100. -
load_AB: TMA producer phase that loads A and B tiles from GMEM into SMEM via multicast.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!