nvfp4.py

"""NVFP4: E2M1 elements, per-16 E4M3 block scales, two-level weight scaling, and the scale swizzle.

THE FORMAT, and every constant here is derived rather than remembered.

  E2M1 element -- 1 sign, 2 exponent (bias 1), 1 mantissa, with subnormals:
        e=0 subnormal : m * 2^0 * 0.5         -> {0, 0.5}
        e>0 normal    : (1 + 0.5m) * 2^(e-1)  -> {1, 1.5}, {2, 3}, {4, 6}
      magnitudes {0, 0.5, 1, 1.5, 2, 3, 4, 6}, max 6. The 3-bit magnitude field is exactly the
      index into that sorted table, because (e, m) enumerates it in order -- which is what makes
      encoding a bucketize and nothing more.
      NON-UNIFORM, and that is the point: fine near zero, 33% steps at the top, where INT4's
      uniform grid steps 14%. It buys resolution where weights actually live.

  block  16 consecutive elements along the CONTRACTION axis share one E4M3 scale.
         (MXFP4 is the sibling: block 32, E8M0 power-of-two scale.)

  two levels, for weights. E4M3 tops out at 448, so a lone block scale cannot span a whole
  tensor's dynamic range. A per-tensor fp32 factor is applied first:
        global = 448 * 6 / amax_tensor
        block  = amax_block * global / 6   in [0, 448]  -> representable in E4M3
        element= x * global / block        in [-6, 6]   -> representable in E2M1
        dequant: x ~= element * block / global
      Each step is checked by construction, so a tensor cannot silently overflow either container.

THE SCALE SWIZZLE IS MANDATORY, not an optimisation. `torch._scaled_mm_v2` rejects row-major block
scales outright: "scale_a must be swizzled to SWIZZLE_32_4_4 format". The layout tiles the scale
matrix 128 rows x 4 columns and rearranges each tile as (32, 4, 4). `swizzle_scales` implements it
and `verify_nvfp4_gemm` checks it against an fp32 reference -- a wrong permutation still produces a
well-formed GEMM of the wrong numbers, so it is verified, never assumed.
"""
from __future__ import annotations

import torch

FP4_MAX = 6.0
E4M3_MAX = 448.0
BLOCK = 16

# sorted E2M1 magnitudes; index == the 3-bit magnitude field
_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
# midpoints, for round-to-nearest
_BOUND = tuple((_E2M1[i] + _E2M1[i + 1]) / 2 for i in range(7))
_CACHE: dict = {}


def _tab(device):
    k = str(device)
    if k not in _CACHE:
        _CACHE[k] = (torch.tensor(_E2M1, device=device, dtype=torch.float32),
                     torch.tensor(_BOUND, device=device, dtype=torch.float32))
    return _CACHE[k]


def quantize_nvfp4(x: torch.Tensor, block: int = BLOCK, two_level: bool = True):
    """[.., K] -> (packed float4_e2m1fn_x2 [.., K/2], e4m3 scales [.., K/block], global fp32).

    `two_level=False` drops the per-tensor factor, which is what SVDQuant's config does for
    ACTIVATIONS (one level, block 16, E4M3). Weights use both levels.
    """
    assert x.shape[-1] % block == 0, f"last dim {x.shape[-1]} not a multiple of block {block}"
    vals, bounds = _tab(x.device)
    xf = x.detach().float()
    amax = xf.abs().amax().clamp(min=1e-12)
    glob = (E4M3_MAX * FP4_MAX / amax) if two_level else torch.ones((), device=x.device)
    xg = xf * glob

    shp = xg.shape
    g = xg.reshape(*shp[:-1], shp[-1] // block, block)
    bmax = g.abs().amax(dim=-1, keepdim=True)                       # [.., nb, 1]
    bs = (bmax / FP4_MAX).clamp(min=1e-12, max=E4M3_MAX)
    bs_q = bs.to(torch.float8_e4m3fn).float()                       # round the SCALE first ...
    bs_q = torch.where(bs_q > 0, bs_q, torch.full_like(bs_q, 2.0 ** -9))
    e = g / bs_q                                                     # ... then divide by the
    e = e.clamp(-FP4_MAX, FP4_MAX)                                   #     rounded scale, so the
    # encode: nearest E2M1 magnitude, sign in bit 3                  #     decoder can reproduce it
    idx = torch.bucketize(e.abs().contiguous(), bounds)
    code = (idx | (torch.signbit(e).to(torch.uint8) << 3)).to(torch.uint8)
    code = code.reshape(*shp)
    lo, hi = code[..., 0::2], code[..., 1::2]                        # first element in LOW bits
    # VIEWED AS float4_e2m1fn_x2, not left as raw bytes: `_scaled_mm_v2` type-checks its operands
    # and rejects Byte outright ("Expected mat_a to be Float8 or Float4_x2 matrix got Byte"). The
    # storage is identical; the dtype is what tells the kernel how to read it.
    packed = (lo | (hi << 4)).contiguous().view(torch.float4_e2m1fn_x2)
    scales = bs.reshape(*shp[:-1], shp[-1] // block).to(torch.float8_e4m3fn)
    return packed, scales, glob.float()


def dequantize_nvfp4(packed: torch.Tensor, scales: torch.Tensor, glob, block: int = BLOCK):
    """Exact inverse of `quantize_nvfp4`. This is the reference the hardware must agree with."""
    vals, _ = _tab(packed.device)
    if packed.dtype != torch.uint8:
        packed = packed.view(torch.uint8)
    lo = (packed & 0x0F)
    hi = (packed >> 4) & 0x0F
    K = packed.shape[-1] * 2
    code = torch.empty(*packed.shape[:-1], K, dtype=torch.uint8, device=packed.device)
    code[..., 0::2] = lo
    code[..., 1::2] = hi
    mag = vals[(code & 0x07).long()]
    sign = torch.where((code & 0x08) > 0, -1.0, 1.0)
    e = mag * sign
    s = scales.float().unsqueeze(-1)
    x = (e.reshape(*e.shape[:-1], K // block, block) * s).reshape(*e.shape[:-1], K)
    gv = glob if torch.is_tensor(glob) else torch.tensor(glob, device=packed.device)
    return x / gv


def swizzle_scales(sf: torch.Tensor) -> torch.Tensor:
    """[M, K/16] E4M3 block scales -> the SWIZZLE_32_4_4 layout `_scaled_mm_v2` requires.

    Tiles of 128 rows x 4 scale-columns, each stored as (32, 4, 4): row r of a tile goes to
    (r % 32, (r // 32) % 4) and column c to (c % 4). M and K/16 are zero-padded up to the tile
    quantum, which is why the returned buffer can be larger than the input.
    """
    M, S = sf.shape
    Mp = (M + 127) // 128 * 128
    Sp = (S + 3) // 4 * 4
    pad = torch.zeros(Mp, Sp, dtype=sf.dtype, device=sf.device)
    pad[:M, :S] = sf
    t = pad.reshape(Mp // 128, 4, 32, Sp // 4, 4)      # (tile_m, r//32, r%32, tile_k, c%4)
    t = t.permute(0, 3, 2, 1, 4)                        # (tile_m, tile_k, 32, 4, 4)
    return t.reshape(-1).contiguous().view(torch.float8_e4m3fn)


_ONES: dict = {}


def recip_scale(g, device) -> torch.Tensor:
    """The 1-element fp32 reciprocal `_scaled_mm_v2` wants for the TensorWise level.

    Build this ONCE, at load, and hand the same tensor to every call. Computing it per forward
    allocates and launches for a constant, which at this model's shapes costs several times the
    GEMM: 13 us of FP4 GEMM behind 40+ us of scalar bookkeeping.
    """
    if g is None:
        k = str(device)
        if k not in _ONES:
            _ONES[k] = torch.ones(1, device=device, dtype=torch.float32)
        return _ONES[k]
    return (1.0 / (g if torch.is_tensor(g) else torch.tensor(g))).reshape(1).float().to(device)


def nvfp4_mm(a_packed, a_scale_sw, b_packed, b_scale_sw, a_rglob=None, b_rglob=None,
             out_dtype=torch.bfloat16, b_packed_t=None):
    """A[M,K] x B[K,N] on the FP4 tensor cores. Packed operands, pre-swizzled block scales.

    `a_rglob` / `b_rglob` are the RECIPROCALS of the per-tensor global scales, as 1-element fp32
    tensors from `recip_scale` -- precomputed, not derived here. They ride inside the kernel as a
    second TensorWise scale level; dividing the output afterwards instead costs two elementwise
    kernels that at these shapes exceed the GEMM itself. Both operands must carry the level or the
    configuration is rejected, so a one-level activation passes a reciprocal of 1.

    `b_packed_t` lets a caller hand in the transposed view once instead of re-taking it per call.
    """
    ST = torch._C._ScalingType
    SW = torch._C._SwizzleType
    bw, tw = int(ST.BlockWise1x16.value), int(ST.TensorWise.value)
    swz, nos = int(SW.SWIZZLE_32_4_4.value), int(SW.NO_SWIZZLE.value)
    ra = recip_scale(None, a_packed.device) if a_rglob is None else a_rglob
    rb = recip_scale(None, a_packed.device) if b_rglob is None else b_rglob
    bt = b_packed.t() if b_packed_t is None else b_packed_t
    return torch._scaled_mm_v2(a_packed, bt, [a_scale_sw, ra], [bw, tw], [swz, nos],
                               [b_scale_sw, rb], [bw, tw], [swz, nos], None, out_dtype)