File size: 4,702 Bytes
e9c8366 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """fp4 — FP4 (E2M1) 4-bit floating-point format primitives.
FP4 E2M1: 1 sign + 2 exponent + 1 mantissa. Bias=1. 16 levels (including
signed zero). Per-group FP8 (E4M3) scale (group_size=16) + F32 global scale.
NVFP4 (NVIDIA): weights AND activations in FP4, per-16 FP8 scales, F32 global
weight_scale_2 + input_scale. Storage: 0.5 byte/weight + 1 byte FP8 scale per
16 weights + 4 bytes global.
MXFP4 (OCP MX): FP4 E2M1 + E8M0 (8-bit exponent only) shared block scale
(block_size=32). Single scale per block.
"""
from __future__ import annotations
import torch
# FP4 E2M1: sign(1) + exp(2) + mant(1). bias=1.
# 16 levels: code 0..15, bit 3=sign, bits 2-1=exp, bit 0=mant.
# Special: code 0 = +0, code 8 = -0.
# exp=0 (codes 0,1,8,9): subnormal: sign * (mant/2) * 2^(1-bias) = sign * mant/2 * 2^0 = sign*mant/2
# code 0: +0, code 1: +0.5, code 8: -0, code 9: -0.5
# exp=1 (codes 2,3,10,11): sign * (1 + mant/2) * 2^(1-1) = sign * (1+mant/2)
# code 2: +1.0, code 3: +1.5, code 10: -1.0, code 11: -1.5
# exp=2 (codes 4,5,12,13): sign * (1+mant/2) * 2^(2-1) = sign * (1+mant/2)*2
# code 4: +2.0, code 5: +3.0, code 12: -2.0, code 13: -3.0
# exp=3 (codes 6,7,14,15): sign * (1+mant/2) * 2^(3-1) = sign * (1+mant/2)*4
# code 6: +4.0, code 7: +6.0, code 14: -4.0, code 15: -6.0
FP4_E2M1_LUT: torch.Tensor = torch.tensor([
0.0, # 0: +0
0.5, # 1: +0.5 (subnormal)
1.0, # 2: +1.0
1.5, # 3: +1.5
2.0, # 4: +2.0
3.0, # 5: +3.0
4.0, # 6: +4.0
6.0, # 7: +6.0
-0.0, # 8: -0
-0.5, # 9: -0.5
-1.0, # 10: -1.0
-1.5, # 11: -1.5
-2.0, # 12: -2.0
-3.0, # 13: -3.0
-4.0, # 14: -4.0
-6.0, # 15: -6.0
], dtype=torch.float32)
def quantize_fp4(w: torch.Tensor, group_size: int = 16) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize weights to FP4 E2M1 indices + per-group absmax scale.
Args:
w: float [out, in] (or flat). in is padded to be divisible by group_size.
group_size: elements per scale group (NVFP4=16, MXFP4=32).
Returns:
(idx, scale): idx int64 [out, in_padded] values [0,15], scale fp32 [out, num_groups].
"""
w = w.detach().float()
out_features = w.shape[0]
in_features = w.shape[1] if w.dim() > 1 else w.numel()
if w.dim() > 1:
flat = w
else:
flat = w.reshape(1, -1)
out_features = 1
pad = (group_size - (flat.shape[1] % group_size)) % group_size
if pad > 0:
flat = torch.nn.functional.pad(flat, (0, pad))
in_padded = flat.shape[1]
num_groups = in_padded // group_size
grouped = flat.reshape(out_features, num_groups, group_size)
# Per-group absmax scale.
scale = grouped.abs().amax(dim=2).clamp(min=1e-8)
scale_exp = scale.unsqueeze(2).expand_as(grouped)
w_norm = grouped / scale_exp
# Nearest FP4 level.
lut = FP4_E2M1_LUT.to(w_norm.device)
diff = w_norm.unsqueeze(-1) - lut
idx = diff.abs().argmin(dim=-1).to(torch.int64)
idx = idx.reshape(out_features, in_padded)
return idx, scale
def dequantize_fp4(
idx: torch.Tensor, scale: torch.Tensor, group_size: int = 16, in_features: int | None = None
) -> torch.Tensor:
"""Reconstruct float weight from FP4 indices + per-group scale."""
out_features = idx.shape[0]
in_padded = idx.shape[1]
num_groups = in_padded // group_size
lut = FP4_E2M1_LUT.to(idx.device)
w_norm = lut[idx] # [out, in_padded]
w_grouped = w_norm.reshape(out_features, num_groups, group_size)
scale_exp = scale.to(torch.float32).unsqueeze(2).expand_as(w_grouped)
w_deq = (w_grouped * scale_exp).reshape(out_features, in_padded)
if in_features is not None and in_features < in_padded:
w_deq = w_deq[:, :in_features]
return w_deq
def pack_fp4(idx: torch.Tensor) -> torch.Tensor:
"""Pack two 4-bit FP4 indices into one int8 byte."""
assert idx.shape[-1] % 2 == 0, "in_features must be even for packing"
i = idx.to(torch.int16)
packed = (i[..., 0::2] << 4) | (i[..., 1::2] & 0x0F)
return packed.to(torch.int8)
def unpack_fp4(packed: torch.Tensor) -> torch.Tensor:
"""Unpack int8 bytes into two 4-bit FP4 indices."""
p = packed.to(torch.int16)
high = (p >> 4) & 0x0F
low = p & 0x0F
out = torch.stack([high, low], dim=-1).reshape(*p.shape[:-1], -1)
return out.to(torch.int64)
# E8M0: 8-bit exponent only (no sign, no mantissa). Used as MXFP block scale.
# Value = 2^(code - 127). code 0..255.
def _make_e8m0_lut() -> torch.Tensor:
lut = torch.zeros(256, dtype=torch.float64)
for code in range(256):
lut[code] = 2.0 ** (code - 127)
return lut.float()
E8M0_LUT: torch.Tensor = _make_e8m0_lut() |