Block-Wise FP8 (E4M3-style) Quantization

The problem, from first principles

[03-groupwise-int4-quantization]'s INT4 uses a fixed absolute step size within each group, so small values within a group still lose relative precision compared to the group's largest value. FP8 (E4M3: 1 sign bit, 4 exponent bits, 3 mantissa bits) keeps a floating-point-style exponent instead, so its relative precision stays roughly constant across magnitudes. Implement a simplified block-wise FP8 quantization scheme: split a weight matrix into fixed-size 2D blocks, compute one scale per block, and quantize each block's values into a simulated 8-bit floating-point format with 4 exponent bits and 3 mantissa bits.

From theory to code

scale[block] = max(|W[block]|) / 448          -- FP8 E4M3 max representable magnitude
x_scaled = W[block] / scale[block]
x_fp8 = round_to_3_mantissa_bits(x_scaled)     -- simulated via frexp/ldexp
x_hat = x_fp8 * scale[block]

Constraints

• Rows and columns must each be evenly divisible by block_rows/block_cols respectively.
• One independent scale per 2D block, from that block's max absolute value divided by 448 (E4M3's max magnitude).
• Simulate 3-mantissa-bit rounding on the scaled values using frexp/ldexp (round mantissa to nearest 1/8 step).
• Return the simulated-FP8 quantized values (still floats, since FP8 isn't a native NumPy dtype), the per-block scale matrix, and the dequantized reconstruction.

Hints