For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo module

block_scaled_quantization

Provides block-scaled quantization kernels for NVFP4, MXFP4, and MXFP8.

comptime values

logger

comptime logger = Logger(stdout, prefix=String(""), source_location=False)

Functions

• ​block_scaled_matmul: Dispatches the SM100 block-scaled matmul for NVFP4 or MXFP8, selecting between the Mojo heuristic kernel and the vendor fallback.
• ​block_scaled_matmul_with_epilogue: Our sm100 block scaled matmul kernel still does not support fusion of elementwise operations. This is a temporary implementation that uses our sm100 block scaled matmul kernel and dispatch a separate epilogue kernel to apply the elementwise operations. Callers must allocate c; when an elementwise_lambda_fn is supplied the matmul result is written into c and then read back by the lambda.
• ​block_scales_interleave: Reinterleaves rank-2 scale factors into the 5D TCGEN layout on SM100 hardware.
• ​block_scales_interleave_fp4: Launches the SM100 kernel that reinterleaves rank-2 scale factors into the 5D TCGEN layout.
• ​block_scales_interleave_fp4_kernel: GPU kernel that reinterleaves rank-2 scale factors into the 5D TCGEN layout.
• ​grouped_matmul_block_scaled_amd: Launches the grouped per-expert MXFP4 block-scaled matmul kernel on AMD CDNA4.
• ​grouped_matmul_block_scaled_amd_kernel:
• ​grouped_quantize_dynamic_scaled_fp4_async: Launches the grouped per-expert quantization kernel for NVFP4/MXFP4/MXFP8 on SM100 hardware.
• ​grouped_quantize_dynamic_scaled_fp4_async_kernel: GPU kernel that quantizes per-expert BF16 activation tiles to NVFP4/MXFP4/MXFP8 with TMA-based scale-factor stores.
• ​matmul_dynamic_block_scaled_amd: Launches the AMD CDNA4 MXFP4 block-scaled matmul kernel.
• ​matmul_dynamic_block_scaled_amd_kernel:
• ​naive_block_scaled_matmul: Reference block-scaled matmul that emulates TCGEN scale-factor accumulation on SM100 hardware.
• ​naive_block_scaled_matmul_kernel: Naive GPU kernel that emulates a block-scaled matmul using TCGEN scale factors.
• ​quantize_dynamic_block_scaled: Dispatches dynamic block-scaled quantization to the appropriate hardware-specific kernel.
• ​quantize_dynamic_block_scaled_mxfp4: Launches the AMD CDNA4 MXFP4 quantization kernel over a flat input buffer.
• ​quantize_dynamic_block_scaled_mxfp4_kernel:
• ​quantize_dynamic_scaled_async_fp4_kernel: GPU kernel that quantizes BF16 tiles to NVFP4 using TMA async copies and warp-specialized PDL.
• ​quantize_dynamic_scaled_fp4_async: Launches the TMA-based NVFP4 quantization kernel on SM100 hardware.
• ​quantize_dynamic_scaled_fp4fp8: Launches the SM100 kernel that quantizes BF16 input to NVFP4, MXFP4, or MXFP8 with per-block scale factors.
• ​quantize_dynamic_scaled_fp4fp8_kernel:
• ​quantize_mx_amd: Quantize BF16 activations to MXFP4 or MXFP8 on AMD CDNA4 (MI355X).
• ​quantize_mxfp8_lane_group: Quantizes one thread's slice of an MX block to MXFP8, cooperatively.