Segmented 8-Bit Floating Point Tensor Compression

In modern large-scale neural network training, memory efficiency is paramount. One powerful technique for reducing memory consumption is low-precision quantization, where high-precision floating-point values are converted to more compact representations while preserving model accuracy.

FP8 (8-bit Floating Point) is a cutting-edge numerical format gaining traction in production systems like those powering massive language models. The E4M3 variant of FP8 uses 1 sign bit, 4 exponent bits, and 3 mantissa bits, providing a dynamic range suitable for neural network activations. The maximum representable value in E4M3 format is 448.

Block-wise Quantization improves upon naïve tensor-level quantization by partitioning the input tensor into smaller segments (blocks) and computing an independent scale factor for each block. This approach offers several advantages:

1. Enhanced Precision: Each block captures local value distributions, preventing outliers in one region from degrading precision across the entire tensor
2. Adaptive Dynamic Range: Blocks with smaller values get smaller scale factors, allowing full utilization of the FP8 representable range
3. Production Readiness: This technique is actively used in systems like Kimi K2, DeepSeek, and other frontier AI models

Quantization Process: For each block of values:

1. Compute the maximum absolute value within the block
2. Calculate the scale factor as: scale = max_abs_value / 448
3. Quantize each value by dividing by the scale: quantized_value = round(original_value / scale)

Your Task: Implement a function that performs block-wise FP8 quantization on an input tensor. The function should divide the tensor into blocks, compute scale factors, and return the quantized representation for each block.