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

fp4_utils

Provides low-level FP4 encode/decode utilities and scale-factor helpers.

comptime values

E2M1_TO_FLOAT32

comptime E2M1_TO_FLOAT32 = SIMD(SIMD(0), SIMD(0.5), SIMD(1), SIMD(1.5), SIMD(2), SIMD(3), SIMD(4), SIMD(6), SIMD(-0.0), SIMD(-0.5), SIMD(-1), SIMD(-1.5), SIMD(-2), SIMD(-3), SIMD(-4), SIMD(-6), __list_literal__=NoneType(None))

E4M3_MAXABS_RECIP

comptime E4M3_MAXABS_RECIP = 0.002232142857142857

FP4_E2M1_MANTISSA_WIDTH

comptime FP4_E2M1_MANTISSA_WIDTH = 1

FP4_E2M1_MAX_EXPONENT

comptime FP4_E2M1_MAX_EXPONENT = 2

MXFP4_SF_DTYPE

comptime MXFP4_SF_DTYPE = DType.float8_e8m0fnu

MXFP4_SF_VECTOR_SIZE

comptime MXFP4_SF_VECTOR_SIZE = 32

MXFP8_SF_DTYPE

comptime MXFP8_SF_DTYPE = DType.float8_e8m0fnu

MXFP8_SF_VECTOR_SIZE

comptime MXFP8_SF_VECTOR_SIZE = 32

NVFP4_SF_DTYPE

comptime NVFP4_SF_DTYPE = DType.float8_e4m3fn

NVFP4_SF_VECTOR_SIZE

comptime NVFP4_SF_VECTOR_SIZE = 16

SF_ATOM_K

comptime SF_ATOM_K = 4

SF_ATOM_M

comptime SF_ATOM_M = Tuple(Int(32), Int(4))

SF_K_GROUP_SIZE

comptime SF_K_GROUP_SIZE[SF_VECTOR_SIZE: Int] = (Int(4) * SF_VECTOR_SIZE)

Parameters

• ​SF_VECTOR_SIZE (Int):

SF_MN_GROUP_SIZE

comptime SF_MN_GROUP_SIZE = ((load_from_mem SF_ATOM_M.__getitem_param__[Int(0)]()) * (load_from_mem SF_ATOM_M.__getitem_param__[Int(1)]()))

W4A8_A_DTYPE

comptime W4A8_A_DTYPE = DType.float8_e4m3fn

W4A8_B_DTYPE

comptime W4A8_B_DTYPE = DType.uint8

Functions

• ​block_scaled_operands_compatible: Reports whether A and B may be fed to one block-scaled UMMA.
• ​block_scaled_umma_kind: Selects the SM100 UMMA kind matching the operand and scale-factor types.
• ​cast_f4e2m1x2_to_fp16x2: Converts two FP4 E2M1 nibbles packed in one byte to two float16 values using SM100 PTX.
• ​cast_float_to_fp4e2m1_amd: Converts up to eight floating-point values to packed FP4 E2M1 using AMD CDNA4 intrinsics.
• ​cast_fp32_to_fp4e2m1: Converts eight float32 values to a packed FP4 E2M1 word using SM100 PTX instructions.
• ​cast_fp_to_fp4e2m1: Rounds each floating-point element to the nearest FP4 E2M1 representable value.
• ​cast_uint_to_fp4e2m1: Unpacks FP4 E2M1 nibbles packed inside unsigned integers to floating-point values.
• ​compute_mxfp4_even_scale: Computes the OCP MXFP4 E8M0 scale using even-mode rounding.
• ​convert_ref_scales_to_mxfp8_format: Converts reference float32 block scales into the 5D MXFP8 E8M0 scale-factor layout.
• ​decode_e2m1_to_bf16: Decodes E2M1 nibbles to bfloat16 with branch-free bit arithmetic.
• ​decode_e2m1_to_f16: Decodes E2M1 nibbles to float16 by exponent injection (Preston's trick).
• ​decode_e2m1_to_f32: Decodes E2M1 nibbles to float32 with branch-free bit arithmetic.
• ​decode_e2m1_to_f32_inject: Decodes E2M1 nibbles to float32 by exponent injection (Preston's trick).
• ​get_batched_scale_factor: Loads a scale factor from a 6D batched LayoutTensor at the given batch, row, and column.
• ​get_scale_factor: Loads a scale factor from a 5D non-batched LayoutTensor at the given row and column.
• ​get_scaling_kind: Selects the UMMA kind for a same-format block-scaled matmul.
• ​is_w4a8_operand_pair: Reports whether (activations, weights) is the mixed W4A8 pair.
• ​set_batched_scale_factor: Stores a scale factor into a 6D batched LayoutTensor at the given batch, row, and column.
• ​set_scale_factor: Stores a scale factor into a 5D non-batched LayoutTensor at the given row and column.