File size: 4,768 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 | """fp8 — FP8 (E4M3 / E5M2) floating-point 8-bit format primitives.
FP8 E4M3: 1 sign + 4 exponent + 3 mantissa. Bias=7. Wider mantissa → more
precision, narrower range. Used for weights/activations forward (NVIDIA H100+).
FP8 E5M2: 1 sign + 5 exponent + 2 mantissa. Bias=15. Wider range, less precision.
Used for gradients / heavy-tailed distributions.
Both encodable/decodable in pure PyTorch (no hardware FP8 tensor cores needed).
Dequant → fp32/fp16 matmul (like int8/NF4 strategies).
"""
from __future__ import annotations
import math
import torch
def _make_fp8_lut(exp_bits: int, mant_bits: int, bias: int) -> torch.Tensor:
"""Build the 256-entry lookup table for an FP8 format.
Returns fp32 tensor [256] of dequantized values for each 8-bit code.
Code layout: bit 7 = sign, bits 6..(7-exp) = exponent, low bits = mantissa.
Special: all-zero exponent with zero mantissa = +0; all-zero exp with nonzero
mantissa = subnormal; all-ones exponent = NaN/Inf (we map to ±448 for E4M3,
±57344 for E5M2 — max finite, to avoid NaN in dequant).
"""
lut = torch.zeros(256, dtype=torch.float32)
exp_mask = (1 << exp_bits) - 1
mant_mask = (1 << mant_bits) - 1
max_exp = exp_mask # all-ones exponent
max_finite = 0.0
# Compute max finite value (exp = max_exp - 1, mant = all ones).
emax = (max_exp - 1) - bias
max_finite = (2.0 ** emax) * (1.0 + mant_mask / (2.0 ** mant_bits))
for code in range(256):
sign = -1.0 if (code & 0x80) else 1.0
exp = (code >> mant_bits) & exp_mask
mant = code & mant_mask
if exp == 0:
if mant == 0:
val = 0.0 # signed zero
else:
# Subnormal: val = sign * mant / 2^mant_bits * 2^(1-bias)
val = sign * (mant / (2.0 ** mant_bits)) * (2.0 ** (1 - bias))
elif exp == max_exp:
# Inf/NaN — map to max finite (avoid NaN propagation).
val = sign * max_finite
else:
# Normalized: val = sign * (1 + mant/2^mb) * 2^(exp - bias)
val = sign * (1.0 + mant / (2.0 ** mant_bits)) * (2.0 ** (exp - bias))
lut[code] = val
return lut
# Precomputed LUTs (256 entries each, fp32).
FP8_E4M3_LUT: torch.Tensor = _make_fp8_lut(exp_bits=4, mant_bits=3, bias=7)
FP8_E5M2_LUT: torch.Tensor = _make_fp8_lut(exp_bits=5, mant_bits=2, bias=15)
def quantize_fp8(w: torch.Tensor, lut: torch.Tensor, scale: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize a float tensor to FP8 codes + per-group/per-channel scale.
Args:
w: float tensor [out, in] (or any shape).
lut: 256-entry FP8 LUT (FP8_E4M3_LUT or FP8_E5M2_LUT).
scale: precomputed scale [out] (per-channel) or scalar. If None, computed
as absmax(w) / max(abs(lut)).
Returns:
(codes, scale) where codes: uint8 [out, in], scale: fp32 [out] or scalar.
"""
w = w.detach().float()
if scale is None:
max_lut = lut.abs().amax().clamp(min=1e-12)
if w.dim() > 1:
max_abs = w.abs().amax(dim=tuple(range(1, w.dim())))
scale = (max_abs / max_lut).clamp(min=1e-12)
else:
scale = (w.abs().amax() / max_lut).clamp(min=1e-12).reshape(1)
# Normalize by scale (broadcast depends on scale shape).
if scale.numel() == 1:
w_norm = w / scale
else:
reshape = [1] * w.dim()
reshape[0] = w.shape[0]
w_norm = w / scale.reshape(reshape)
# Nearest LUT entry.
# Clip to LUT range to avoid out-of-range.
lut_max = lut.abs().amax().item()
w_norm = w_norm.clamp(-lut_max, lut_max)
diff = w_norm.unsqueeze(-1) - lut.to(w_norm.device)
codes = diff.abs().argmin(dim=-1).to(torch.uint8)
return codes, scale.to(torch.float32)
def dequantize_fp8(codes: torch.Tensor, scale: torch.Tensor, lut: torch.Tensor) -> torch.Tensor:
"""Reconstruct float tensor from FP8 codes + scale.
Args:
codes: uint8 [out, in] (or any shape).
scale: fp32 [out] (per-channel) or scalar.
lut: 256-entry FP8 LUT.
Returns:
fp32 tensor same shape as codes.
"""
lut_dev = lut.to(codes.device)
w = lut_dev[codes.long()] # [shape], fp32
if scale.numel() == 1:
return w * scale.to(torch.float32)
reshape = [1] * w.dim()
reshape[0] = w.shape[0]
return w * scale.to(torch.float32).reshape(reshape)
def pack_fp8(codes: torch.Tensor) -> torch.Tensor:
"""FP8 codes are already 1 byte each — no packing needed. Returns int8 view."""
return codes.view(torch.int8)
def unpack_fp8(packed: torch.Tensor) -> torch.Tensor:
"""Unpack int8 view back to uint8 codes."""
return packed.view(torch.uint8) |