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

rms_norm_block_scaled

Fused RMSNorm + MXFP8 block-scaled ("quantize-on-write") quantization.

SM100 (B200) / SM120. Computes RMSNorm over each row and, in the same kernel, quantizes the normalized activation to MXFP8 (E4M3 values + E8M0 per-32-element block scales) in the exact interleaved rank-5 scale layout the SM100 block-scaled GEMM consumes. Fusing the quant into the norm epilogue removes the HBM round-trip of the bf16 normalized activation that a standalone quantize_dynamic_block_scaled would otherwise re-read.

The result is numerically equivalent to the unfused pipeline rms_norm (bf16 store) -> quantize_dynamic_block_scaled[MXFP8]:

• The normalized value is rounded to bf16 (in_dtype) BEFORE the per-block max and quantization, exactly as the unfused path stores a bf16 intermediate and the standalone quant re-reads it.
• The per-block E8M0 scale (group_max / 448 cast to float8_e8m0fnu), the reciprocal-scale multiply, and the E4M3 rounding reproduce quantize_dynamic_scaled_fp4fp8_kernel bit-for-bit (linalg/fp4_quantization.mojo).
• The rank-5 interleaved scale write and the align_up(rows, 128) padding-row zeroing reuse set_scale_factor / SF_ATOM_M / SF_ATOM_K / SF_MN_GROUP_SIZE from linalg/fp4_utils.mojo, so the downstream GEMM is a drop-in consumer.

Functions

• ​rms_norm_fused_quantize_dynamic_block_scaled: Fuses RMSNorm and MXFP8 block-scaled quantization in a single pass.

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