Skip to content

Custom floating point, in PyTorch

Quantize tensors to arbitrary FP layouts — including OFP8 / MX FP4 and FP8 — with CUDA kernels and autograd.

  • Any layout

    Sign, exponent, mantissa, bias, and reserved-NaN behavior are yours to set.

  • Trainable

    Round uses a straight-through estimator so gradients flow through quantization.

  • Block-scaled

    Any element codebook with a shared per-block scale.

Install

pip install torch-floating-point --no-build-isolation

git clone https://github.com/SamirMoustafa/torch-floating-point.git
cd torch-floating-point
pip install --no-build-isolation -e .

Requires Python 3.10+ and PyTorch 2.4+. Install that torch first, then this package with --no-build-isolation. A C++ compiler is required. CUDA kernels compile when that torch is a CUDA build and nvcc is available (pip torch does not ship nvcc); otherwise a CPU extension is built. FORCE_CPU=1 / FORCE_CUDA=1 override the default.

Quick start

import torch
from floating_point import FloatingPoint, Round

fp8 = FloatingPoint(sign_bits=1, exponent_bits=4, mantissa_bits=3, bias=7, bits=8)
x = torch.randn(8, requires_grad=True)

y = Round(fp8)(x)
y.sum().backward()

Train with quantized weights using the same rounder:

import torch.nn as nn
from floating_point import FloatingPoint, Round


class FloatPointLinear(nn.Module):
    def __init__(self, inn, out, fp):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out, inn))
        self.bias = nn.Parameter(torch.randn(out))
        self.rounder = Round(fp)

    def forward(self, x):
        return nn.functional.linear(x, self.rounder(self.weight), self.bias)

Next: Formats · Explorer · block scaling · Autograd · Example · API

Cite

@software{moustafa2025torchfloatingpoint,
  title={Torch Floating Point: A PyTorch library for custom floating point quantization},
  author={Samir Moustafa},
  year={2025},
  url={https://github.com/SamirMoustafa/torch-floating-point}
}