Documentation

LeanMlir.E4M3Quant

Pure-Lean E4M3 (fp8) fake-quant over raw-f32 ByteArrays #

The quantizer the Lean fp8 trainer (MainMnistLinearE4M3Verified) wraps around the verified linear train-step kernel. It implements the same E4M3 (1-4-3, bias 7) round-to-nearest grid as the numpy oracle (scripts/mnist_e4m3_demo.py:to_e4m3): subnormals on the e = −6 grid (step 2⁻⁹), saturating at ±448.

This is the host-side "operand byte preparation" half of the §3b render-tie (LeanMlir/Proofs/Float/E4M3Fold.lean): fp8 = fp32 arithmetic on operands projected onto the E4M3 grid, with fp32 accumulate inside the kernel. No fp8 hardware or fp8 StableHLO type is needed — q runs here, in Lean, before the verified matmul sees the bytes.

Values move through F32.read (extern f32→Float) and are re-encoded f32→Float32 → 4 little-endian bytes; the master weights stay fp32 across the gradient update.

source
@[inline]
def F32E4M3.e4m3Max :
Float

E4M3 largest finite magnitude: S.1111.110 = 2⁸·1.75.

Equations
Instances For
    source
    def F32E4M3.roundE4M3 (x : Float) :
    Float

    Round one value to the E4M3 grid (round-to-nearest, subnormals on the e = −6 grid, saturating). Mirrors to_e4m3 in the numpy oracle.

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      source
      @[inline]
      def F32E4M3.pushF32LE (acc : ByteArray) (x : Float) :
      ByteArray

      Append x as 4 little-endian f32 bytes (narrowing f64 → Float32).

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        source
        def F32E4M3.quantPerTensor (ba : ByteArray) :
        ByteArray

        Per-tensor E4M3 quant (one scale s = max|·|/448). Round-trips through the grid: returns s · q(vᵢ/s) as f32 bytes (the dequantized operand).

        Equations
        • One or more equations did not get rendered due to their size.
        Instances For
          source
          def F32E4M3.quantPerColumn (ba : ByteArray) (d0 d1 : Nat) :
          ByteArray

          Per-output-column E4M3 quant for a row-major [d0 × d1] matrix (the "block scale" sWⱼ of §3b: each output column scaled independently).

          Equations
          • One or more equations did not get rendered due to their size.
          Instances For
            source
            def F32E4M3.addDelta (master wOut wq : ByteArray) :
            ByteArray

            fp32-master gradient-delta update. The fused kernel returns wOut = wq − lr·∇ (it updated the quantized operand). Applying the same gradient to the fp32 master is master + (wOut − wq) = master − lr·∇.

            Equations
            • One or more equations did not get rendered due to their size.
            Instances For
              source
              def F32E4M3.quantPerLeadingBlock (ba : ByteArray) (nBlocks blockSize : Nat) :
              ByteArray

              Per-leading-block E4M3 quant: nBlocks contiguous blocks of blockSize elements, each with its own scale. A conv kernel [oc, ic, k, k] (row-major) is oc blocks of ic·k·k, so this is per-output-channel quant.

              Equations
              • One or more equations did not get rendered due to their size.
              Instances For
                source
                def F32E4M3.quantPackedParams (params : ByteArray) (specs : Array (Array Nat × Nat)) :
                ByteArray

                Quantize a packed param buffer slot-by-slot, per the (dims, initKind) layout (VerifiedNetSpec.toSpecs, func-arg order). initKind == 0 weights: dense [ic, oc] → quantPerColumn (per output column), conv [oc, ic, k, k] → quantPerLeadingBlock (per output channel). Biases / γ / β (initKind 1 or 2) are copied (kept fp32). This is the packed analogue of the per-tensor/per-column quant the linear trainer does — the weight operand prep for the verified MLP/CNN train step.

                Equations
                • One or more equations did not get rendered due to their size.
                Instances For