"""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)