Added initial experimental support for convrot4 and convrot8 quantization.
This commit is contained in:
Jaret Burkett
700
toolkit/util/convrot_quant.py
Normal file
700
toolkit/util/convrot_quant.py
Normal file
@@ -0,0 +1,700 @@
|
||||
"""
|
||||
ConvRot quantization backends (convrot4 / convrot8 qtypes).
|
||||
|
||||
convrot4 is the paper's W4A4 NVFP4 method described below. convrot8 pairs the same
|
||||
rotation with per-token / per-output-channel symmetric int8 (W8A8) and
|
||||
torch._int_mm: near-lossless (~1% weight error), and the fast path runs on any int8
|
||||
tensor-core gpu (Ampere+), not just Blackwell. The rotation is what makes the coarse
|
||||
per-row scales safe — it spreads outliers so a whole row shares one scale without
|
||||
clipping damage (the classic SmoothQuant failure mode).
|
||||
|
||||
Implements "ConvRot: Rotation-Based Plug-and-Play 4-bit Quantization for Diffusion
|
||||
Transformers" (arXiv:2512.03673) as an OstrisQuantizer backend, self-contained on
|
||||
top of torch (no torchao version requirements).
|
||||
|
||||
Method: weights and activations are rotated with a block *regular* Hadamard
|
||||
transform (R4 = [[1,1,1,-1],[1,1,-1,1],[1,-1,1,1],[-1,1,1,1]]/2 Kronecker-powered
|
||||
to rot_size, a power of 4, default 256). Unlike the standard Hadamard whose all-ones
|
||||
row concentrates the block mean into one coordinate, the regular Hadamard has
|
||||
constant row sums, smoothing row-wise and column-wise outliers symmetrically. The
|
||||
rotation is folded into the weight offline and applied to the activation at runtime,
|
||||
so it cancels in the matmul. Both sides are then quantized to NVFP4 (fp4 e2m1 values,
|
||||
fp8 e4m3 scale per 16 elements, plus one fp32 per-tensor scale) and multiplied with
|
||||
the Blackwell fp4 tensor cores via torch._scaled_mm — a real ~5-6x gemm speedup, ~2x
|
||||
at the layer level after rotation + activation-quant overhead.
|
||||
|
||||
Paths:
|
||||
- inference (no grad): rotate -> fused triton nvfp4 activation quant ->
|
||||
hardware fp4 gemm. Requires sm_100+ (Blackwell); otherwise falls back to the
|
||||
dequantized matmul below.
|
||||
- training (grad enabled): rotate -> straight-through fake-quant of the
|
||||
activation (so adapters train against the same W4A4 numerics that deployment
|
||||
uses) -> bf16 matmul with the dequantized rotated weight. Fully differentiable
|
||||
w.r.t. the input.
|
||||
|
||||
Everything is deterministic: the rotation is a fixed matrix (no randomness at all)
|
||||
and quantization is pure rounding.
|
||||
|
||||
Quantized state attached to each module:
|
||||
cr_qdata packed e2m1 codes (uint8, out x in/2; low nibble = even element)
|
||||
cr_scales e4m3 block scales (out x in/16)
|
||||
cr_scales_blocked the same scales pre-swizzled for torch._scaled_mm
|
||||
cr_pts fp32 per-tensor scale (scalar)
|
||||
cr_rot / module.cr_rot_size rotation block size
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from toolkit.print import print_acc
|
||||
from toolkit.util.ostris_quant import OstrisQuantizer
|
||||
|
||||
CONVROT_QTYPES = ("convrot4", "convrot8")
|
||||
|
||||
|
||||
def get_convrot_quantizer(qtype: str):
|
||||
if qtype == "convrot4":
|
||||
return ConvRotQuantizer(rot_size=256)
|
||||
if qtype == "convrot8":
|
||||
return ConvRotInt8Quantizer(rot_size=256)
|
||||
return None
|
||||
|
||||
|
||||
F4_MAX = 6.0
|
||||
F8_E4M3_MAX = 448.0
|
||||
BLOCK = 16 # nvfp4 scale block
|
||||
|
||||
_E2M1_VALS = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
|
||||
_E2M1_EDGES = [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0]
|
||||
|
||||
_hadamard_cache = {}
|
||||
_vals_cache = {}
|
||||
_edges_cache = {}
|
||||
_skip_warned = set()
|
||||
|
||||
|
||||
def _cached(cache, key, build):
|
||||
if key not in cache:
|
||||
cache[key] = build()
|
||||
return cache[key]
|
||||
|
||||
|
||||
def regular_hadamard(rot_size: int, device, dtype=torch.bfloat16) -> torch.Tensor:
|
||||
"""The ConvRot rotation: Kronecker powers of the 4x4 regular Hadamard matrix,
|
||||
orthonormal. Symmetric and orthogonal, so it is its own inverse."""
|
||||
key = (rot_size, str(device), dtype)
|
||||
|
||||
def build():
|
||||
r4 = torch.tensor(
|
||||
[[1.0, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]],
|
||||
dtype=torch.float64,
|
||||
)
|
||||
h = r4.clone()
|
||||
while h.shape[0] < rot_size:
|
||||
h = torch.kron(h, r4)
|
||||
if h.shape[0] != rot_size:
|
||||
raise ValueError(f"rot_size {rot_size} is not a power of 4")
|
||||
return (h / rot_size**0.5).to(device=device, dtype=dtype)
|
||||
|
||||
return _cached(_hadamard_cache, key, build)
|
||||
|
||||
|
||||
def largest_pow4_divisor(d: int) -> int:
|
||||
h = 1
|
||||
while d % (h * 4) == 0:
|
||||
h *= 4
|
||||
return h
|
||||
|
||||
|
||||
def rotate(x: torch.Tensor, rot_size: int) -> torch.Tensor:
|
||||
"""Apply the block regular-Hadamard rotation along the last dim (self-inverse)."""
|
||||
if rot_size == 1:
|
||||
return x
|
||||
h = regular_hadamard(rot_size, x.device, x.dtype)
|
||||
shape = x.shape
|
||||
xb = x.reshape(-1, shape[-1] // rot_size, rot_size)
|
||||
return torch.matmul(xb, h).reshape(shape)
|
||||
|
||||
|
||||
def to_blocked(m: torch.Tensor) -> torch.Tensor:
|
||||
"""Rearrange an (R, C) scale matrix into the swizzled layout torch._scaled_mm
|
||||
expects for block-scaled fp4 (cublas 128x4-tile layout)."""
|
||||
rows, cols = m.shape
|
||||
rb, cb = -(-rows // 128), -(-cols // 4)
|
||||
if (rows, cols) != (rb * 128, cb * 4):
|
||||
padded = torch.zeros(rb * 128, cb * 4, device=m.device, dtype=m.dtype)
|
||||
padded[:rows, :cols] = m
|
||||
m = padded
|
||||
blocks = m.view(rb, 128, cb, 4).permute(0, 2, 1, 3)
|
||||
return blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16).flatten()
|
||||
|
||||
|
||||
def quantize_nvfp4(x: torch.Tensor, pts: Optional[torch.Tensor] = None):
|
||||
"""Quantize (rows, K) to nvfp4. Returns (packed uint8 (rows, K/2),
|
||||
e4m3 scales (rows, K/16), fp32 per-tensor scale)."""
|
||||
rows, K = x.shape
|
||||
xf = x.float()
|
||||
if pts is None:
|
||||
pts = xf.abs().amax() / (F4_MAX * F8_E4M3_MAX)
|
||||
pts = torch.where(pts > 0, pts, torch.ones_like(pts))
|
||||
xb = xf.view(rows, K // BLOCK, BLOCK)
|
||||
scales = (xb.abs().amax(dim=-1) / (F4_MAX * pts)).to(torch.float8_e4m3fn)
|
||||
denom = (scales.float() * pts).unsqueeze(-1)
|
||||
z = (xb / torch.where(denom > 0, denom, torch.ones_like(denom))).clamp(
|
||||
-F4_MAX, F4_MAX
|
||||
)
|
||||
edges = _cached(
|
||||
_edges_cache, str(x.device), lambda: torch.tensor(_E2M1_EDGES, device=x.device)
|
||||
)
|
||||
mag = torch.bucketize(z.abs(), edges).to(torch.uint8)
|
||||
codes = (mag | ((z < 0).to(torch.uint8) << 3)).view(rows, K)
|
||||
packed = ((codes[:, 1::2] << 4) | codes[:, ::2]).contiguous()
|
||||
return packed, scales, pts
|
||||
|
||||
|
||||
def dequantize_nvfp4(
|
||||
packed: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
pts: torch.Tensor,
|
||||
rows: int,
|
||||
K: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
codes = torch.stack([packed & 15, packed >> 4], dim=-1).view(rows, K)
|
||||
vals = _cached(
|
||||
_vals_cache,
|
||||
str(packed.device),
|
||||
lambda: torch.tensor(_E2M1_VALS, device=packed.device),
|
||||
)
|
||||
mag = torch.index_select(vals, 0, (codes & 7).flatten().to(torch.int32)).view(
|
||||
rows, K
|
||||
)
|
||||
v = mag * torch.where((codes & 8) > 0, -1.0, 1.0)
|
||||
v = v.view(rows, K // BLOCK, BLOCK) * (scales.float() * pts).unsqueeze(-1)
|
||||
return v.view(rows, K).to(dtype)
|
||||
|
||||
|
||||
# ---------------- fused triton activation quant ----------------
|
||||
|
||||
_triton_ok = None
|
||||
|
||||
|
||||
def _triton_available() -> bool:
|
||||
global _triton_ok
|
||||
if _triton_ok is None:
|
||||
try:
|
||||
import triton # noqa: F401
|
||||
import triton.language as tl # noqa: F401
|
||||
|
||||
_triton_ok = True
|
||||
except Exception:
|
||||
_triton_ok = False
|
||||
print_acc(
|
||||
"ConvRot: triton is not available. The fused activation-quant kernel is "
|
||||
"disabled and activations will be quantized with plain torch ops instead "
|
||||
"— inference gets slower (most of the fp4 speedup is lost), but quality "
|
||||
"and training are unaffected."
|
||||
)
|
||||
return _triton_ok
|
||||
|
||||
|
||||
_kernel = None
|
||||
|
||||
|
||||
def _get_kernel():
|
||||
global _kernel
|
||||
if _kernel is not None:
|
||||
return _kernel
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
@triton.jit
|
||||
def nvfp4_act_quant_kernel(
|
||||
x_ptr,
|
||||
out_ptr,
|
||||
scale_ptr,
|
||||
pts_ptr,
|
||||
K,
|
||||
n_col_tiles,
|
||||
BLOCK_K: tl.constexpr,
|
||||
BLOCKED_SCALES: tl.constexpr,
|
||||
):
|
||||
pid_m = tl.program_id(0)
|
||||
pid_k = tl.program_id(1)
|
||||
pts = tl.load(pts_ptr)
|
||||
offs = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
|
||||
mask = offs < K
|
||||
x = tl.load(x_ptr + pid_m * K + offs, mask=mask, other=0.0).to(tl.float32)
|
||||
xb = tl.reshape(x, (BLOCK_K // 16, 16))
|
||||
amax = tl.max(tl.abs(xb), axis=1)
|
||||
scale8 = (amax / (6.0 * pts)).to(tl.float8e4nv)
|
||||
denom = scale8.to(tl.float32) * pts
|
||||
denom = tl.where(denom > 0, denom, 1.0)
|
||||
# note: triton fp32 division on this backend is ~1ulp off ieee (even with
|
||||
# tl.fdiv ieee_rounding=True), so values landing exactly on a code boundary
|
||||
# can round to the adjacent code vs the torch path. ties are equidistant, so
|
||||
# this changes nothing quantitatively; activation codes are transient (never
|
||||
# stored), and the kernel itself is deterministic.
|
||||
z = xb / denom[:, None]
|
||||
z = tl.minimum(tl.maximum(z, -6.0), 6.0)
|
||||
az = tl.abs(z)
|
||||
# strict > so exact midpoints go to the lower code, matching torch.bucketize
|
||||
code = (
|
||||
(az > 0.25).to(tl.uint8)
|
||||
+ (az > 0.75).to(tl.uint8)
|
||||
+ (az > 1.25).to(tl.uint8)
|
||||
+ (az > 1.75).to(tl.uint8)
|
||||
+ (az > 2.5).to(tl.uint8)
|
||||
+ (az > 3.5).to(tl.uint8)
|
||||
+ (az > 5.0).to(tl.uint8)
|
||||
)
|
||||
code = code | ((z < 0).to(tl.uint8) << 3)
|
||||
lo, hi = tl.split(tl.reshape(code, (BLOCK_K // 2, 2)))
|
||||
byte = lo | (hi << 4)
|
||||
offs_b = pid_k * (BLOCK_K // 2) + tl.arange(0, BLOCK_K // 2)
|
||||
tl.store(out_ptr + pid_m * (K // 2) + offs_b, byte, mask=offs_b < K // 2)
|
||||
s_idx = pid_k * (BLOCK_K // 16) + tl.arange(0, BLOCK_K // 16)
|
||||
if BLOCKED_SCALES:
|
||||
# store straight into the cublas 128x4-tile swizzle (see to_blocked)
|
||||
r_t = pid_m // 128
|
||||
r_in = pid_m % 128
|
||||
c_t = s_idx // 4
|
||||
c = s_idx % 4
|
||||
offs_s = (
|
||||
((r_t * n_col_tiles + c_t) * 32 + (r_in % 32)) * 16
|
||||
+ (r_in // 32) * 4
|
||||
+ c
|
||||
)
|
||||
else:
|
||||
offs_s = pid_m * (K // 16) + s_idx
|
||||
tl.store(scale_ptr + offs_s, scale8, mask=s_idx < K // 16)
|
||||
|
||||
_kernel = nvfp4_act_quant_kernel
|
||||
return _kernel
|
||||
|
||||
|
||||
def quantize_nvfp4_fused(x: torch.Tensor, blocked_scales: bool = False):
|
||||
"""Triton path of quantize_nvfp4 for the inference hot loop: one read of x,
|
||||
writes packed codes + e4m3 scales (row-major, or directly in the swizzled
|
||||
layout torch._scaled_mm wants when blocked_scales=True). Falls back to the
|
||||
torch ops (row-major only)."""
|
||||
rows, K = x.shape
|
||||
if not (_triton_available() and x.is_cuda and K % 16 == 0):
|
||||
packed, scales, pts = quantize_nvfp4(x)
|
||||
return (
|
||||
(packed, to_blocked(scales), pts)
|
||||
if blocked_scales
|
||||
else (packed, scales, pts)
|
||||
)
|
||||
pts = x.float().abs().amax() / (F4_MAX * F8_E4M3_MAX)
|
||||
pts = torch.where(pts > 0, pts, torch.ones_like(pts))
|
||||
x = x.contiguous()
|
||||
packed = torch.empty(rows, K // 2, device=x.device, dtype=torch.uint8)
|
||||
n_col_tiles = -(-(K // BLOCK) // 4)
|
||||
if blocked_scales:
|
||||
# zero-init: rows are padded to 128-tiles and the pad region must be zero
|
||||
scales = torch.zeros(
|
||||
(-(-rows // 128)) * 128 * n_col_tiles * 4,
|
||||
device=x.device,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
)
|
||||
else:
|
||||
scales = torch.empty(
|
||||
rows, K // BLOCK, device=x.device, dtype=torch.float8_e4m3fn
|
||||
)
|
||||
BLOCK_K = 2048 if K >= 2048 else K
|
||||
grid = (rows, -(-K // BLOCK_K))
|
||||
_get_kernel()[grid](
|
||||
x,
|
||||
packed,
|
||||
scales,
|
||||
pts,
|
||||
K,
|
||||
n_col_tiles,
|
||||
BLOCK_K=BLOCK_K,
|
||||
BLOCKED_SCALES=blocked_scales,
|
||||
num_warps=4,
|
||||
)
|
||||
return packed, scales, pts
|
||||
|
||||
|
||||
# ---------------- backend ----------------
|
||||
|
||||
|
||||
_warned_no_fp4 = False
|
||||
|
||||
|
||||
def _fp4_gemm_supported(device) -> bool:
|
||||
global _warned_no_fp4
|
||||
device = torch.device(device)
|
||||
supported = (
|
||||
device.type == "cuda"
|
||||
and torch.cuda.is_available()
|
||||
and torch.cuda.get_device_capability(device)[0] >= 10 # Blackwell
|
||||
)
|
||||
if not supported and not _warned_no_fp4:
|
||||
_warned_no_fp4 = True
|
||||
print_acc(
|
||||
f"ConvRot: no fp4 tensor-core support on this device ({device}; needs an "
|
||||
"NVIDIA Blackwell GPU, sm_100+). Inference falls back to dequantized bf16 "
|
||||
"matmuls: correct output but NO speedup, and inference activations stay "
|
||||
"unquantized (W4A16 numerics instead of W4A4). The training path is "
|
||||
"unaffected (it always simulates W4A4 via fake-quant)."
|
||||
)
|
||||
return supported
|
||||
|
||||
|
||||
class ConvRotQuantizer(OstrisQuantizer):
|
||||
"""ConvRot W4A4 backend. One instance per qtype, shareable across modules."""
|
||||
|
||||
def __init__(self, rot_size: int = 256):
|
||||
self.rot_size = rot_size
|
||||
|
||||
def _rot_for(self, d: int) -> int:
|
||||
return min(self.rot_size, largest_pow4_divisor(d))
|
||||
|
||||
def can_quantize(self, module: torch.nn.Linear) -> bool:
|
||||
d = module.in_features
|
||||
rot = self._rot_for(d)
|
||||
if d % BLOCK != 0 or module.out_features % BLOCK != 0 or rot < 16:
|
||||
if d not in _skip_warned:
|
||||
_skip_warned.add(d)
|
||||
print_acc(
|
||||
f"ConvRot: skipping linears with in_features={d} "
|
||||
f"(needs in/out divisible by 16 and a power-of-4 block >= 16)"
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None:
|
||||
rot = self._rot_for(module.in_features)
|
||||
w_rot = rotate(weight_fp32, rot)
|
||||
packed, scales, pts = quantize_nvfp4(w_rot)
|
||||
# scales/pts are stored as uint8 byte views: nn.Module._apply dtype-casts
|
||||
# every floating buffer (module.to(dtype=...) would silently convert the
|
||||
# e4m3 scales to bf16 and fp32 pts to bf16, corrupting them). integer
|
||||
# buffers are only ever moved, never cast.
|
||||
module.register_buffer("cr_qdata", packed, persistent=False)
|
||||
module.register_buffer("cr_scales", scales.view(torch.uint8), persistent=False)
|
||||
module.register_buffer(
|
||||
"cr_scales_blocked", to_blocked(scales).view(torch.uint8), persistent=False
|
||||
)
|
||||
module.register_buffer(
|
||||
"cr_pts",
|
||||
pts.detach().clone().reshape(1).view(torch.uint8),
|
||||
persistent=False,
|
||||
)
|
||||
module.cr_rot_size = rot
|
||||
|
||||
@staticmethod
|
||||
def _pts(module) -> torch.Tensor:
|
||||
return module.cr_pts.view(torch.float32).reshape(())
|
||||
|
||||
def _dequantize_rotated(self, module, dtype: torch.dtype) -> torch.Tensor:
|
||||
return dequantize_nvfp4(
|
||||
module.cr_qdata,
|
||||
module.cr_scales.view(torch.float8_e4m3fn),
|
||||
self._pts(module),
|
||||
module.out_features,
|
||||
module.in_features,
|
||||
dtype,
|
||||
)
|
||||
|
||||
def dequantize(self, module) -> torch.Tensor:
|
||||
w = self._dequantize_rotated(module, torch.float32)
|
||||
return rotate(w, module.cr_rot_size) # self-inverse
|
||||
|
||||
def requantize_(self, module, fp_weight: torch.Tensor) -> None:
|
||||
w = fp_weight.to(device=module.cr_qdata.device, dtype=torch.float32)
|
||||
w_rot = rotate(w, module.cr_rot_size)
|
||||
packed, scales, pts = quantize_nvfp4(w_rot)
|
||||
module.cr_qdata = packed
|
||||
module.cr_scales = scales.view(torch.uint8)
|
||||
module.cr_scales_blocked = to_blocked(scales).view(torch.uint8)
|
||||
module.cr_pts = pts.detach().clone().reshape(1).view(torch.uint8)
|
||||
|
||||
def forward(self, module, x: torch.Tensor) -> torch.Tensor:
|
||||
rot = module.cr_rot_size
|
||||
in_f, out_f = module.in_features, module.out_features
|
||||
x_rot = rotate(x, rot)
|
||||
x2d = x_rot.reshape(-1, in_f)
|
||||
m = x2d.shape[0]
|
||||
|
||||
if x.requires_grad:
|
||||
# training path: straight-through fake-quant of the activation so
|
||||
# adapters see the same W4A4 numerics as deployment, then a
|
||||
# differentiable bf16 matmul against the dequantized rotated weight.
|
||||
# gated on requires_grad alone (not is_grad_enabled) so both passes of
|
||||
# gradient checkpointing take the same branch and recompute identically
|
||||
with torch.no_grad():
|
||||
aq, a_scales, a_pts = quantize_nvfp4(x2d.detach())
|
||||
x_dq = dequantize_nvfp4(aq, a_scales, a_pts, m, in_f, x.dtype)
|
||||
w = self._dequantize_rotated(module, x.dtype)
|
||||
x_ste = x2d + (x_dq - x2d).detach()
|
||||
out = F.linear(x_ste, w, module.bias)
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
|
||||
if _fp4_gemm_supported(x.device):
|
||||
pad = (-m) % BLOCK
|
||||
if pad:
|
||||
x2d = F.pad(x2d, (0, 0, 0, pad))
|
||||
aq, a_scales_blocked, a_pts = quantize_nvfp4_fused(x2d, blocked_scales=True)
|
||||
out = torch._scaled_mm(
|
||||
aq.view(torch.float4_e2m1fn_x2),
|
||||
module.cr_qdata.view(torch.float4_e2m1fn_x2).t(),
|
||||
a_scales_blocked.view(torch.float8_e4m3fn),
|
||||
module.cr_scales_blocked.view(torch.float8_e4m3fn),
|
||||
out_dtype=x.dtype,
|
||||
)
|
||||
if pad:
|
||||
out = out[:m]
|
||||
s = (a_pts * self._pts(module)).to(x.dtype)
|
||||
if module.bias is not None:
|
||||
out = torch.addcmul(module.bias, out, s)
|
||||
else:
|
||||
out = out * s
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
|
||||
# no fp4 hardware: dequantized matmul (correct, no speedup)
|
||||
w = self._dequantize_rotated(module, x.dtype)
|
||||
out = F.linear(x2d, w, module.bias)
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
|
||||
|
||||
# ---------------- convrot8: W8A8 int8 backend ----------------
|
||||
|
||||
|
||||
def quantize_int8_rows(x: torch.Tensor):
|
||||
"""Symmetric per-row int8 quantization. Returns (int8 (rows, K), fp32 scales (rows,))."""
|
||||
xf = x.float()
|
||||
scales = xf.abs().amax(dim=1) / 127.0
|
||||
scales = torch.where(scales > 0, scales, torch.ones_like(scales))
|
||||
q = torch.round(xf / scales.unsqueeze(1)).clamp_(-127, 127).to(torch.int8)
|
||||
return q, scales
|
||||
|
||||
|
||||
_int8_kernels = None
|
||||
|
||||
|
||||
def _get_int8_kernels():
|
||||
global _int8_kernels
|
||||
if _int8_kernels is not None:
|
||||
return _int8_kernels
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from triton.language.extra import libdevice
|
||||
|
||||
@triton.jit
|
||||
def int8_act_quant_kernel(x_ptr, q_ptr, s_ptr, K, BLOCK_K: tl.constexpr):
|
||||
row = tl.program_id(0)
|
||||
base = row * K
|
||||
acc = tl.zeros((BLOCK_K,), tl.float32)
|
||||
for k0 in range(0, K, BLOCK_K):
|
||||
offs = k0 + tl.arange(0, BLOCK_K)
|
||||
v = tl.load(x_ptr + base + offs, mask=offs < K, other=0.0).to(tl.float32)
|
||||
acc = tl.maximum(acc, tl.abs(v))
|
||||
amax = tl.max(acc, axis=0)
|
||||
scale = tl.where(amax > 0, amax / 127.0, 1.0)
|
||||
for k0 in range(0, K, BLOCK_K):
|
||||
offs = k0 + tl.arange(0, BLOCK_K)
|
||||
mask = offs < K
|
||||
v = tl.load(x_ptr + base + offs, mask=mask, other=0.0).to(tl.float32)
|
||||
# rint = round-half-to-even, matching torch.round in the reference path
|
||||
q = libdevice.rint(v / scale)
|
||||
q = tl.minimum(tl.maximum(q, -127.0), 127.0)
|
||||
tl.store(q_ptr + base + offs, q.to(tl.int8), mask=mask)
|
||||
tl.store(s_ptr + row, scale)
|
||||
|
||||
@triton.jit
|
||||
def int8_epilogue_kernel(
|
||||
i_ptr,
|
||||
as_ptr,
|
||||
ws_ptr,
|
||||
b_ptr,
|
||||
o_ptr,
|
||||
N,
|
||||
HAS_BIAS: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
cb = tl.program_id(1)
|
||||
offs = cb * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
mask = offs < N
|
||||
acc = tl.load(i_ptr + row * N + offs, mask=mask, other=0).to(tl.float32)
|
||||
a_s = tl.load(as_ptr + row)
|
||||
w_s = tl.load(ws_ptr + offs, mask=mask, other=0.0)
|
||||
out = acc * (a_s * w_s)
|
||||
if HAS_BIAS:
|
||||
out += tl.load(b_ptr + offs, mask=mask, other=0.0).to(tl.float32)
|
||||
tl.store(o_ptr + row * N + offs, out.to(o_ptr.dtype.element_ty), mask=mask)
|
||||
|
||||
_int8_kernels = (int8_act_quant_kernel, int8_epilogue_kernel)
|
||||
return _int8_kernels
|
||||
|
||||
|
||||
def quantize_int8_rows_fused(x: torch.Tensor):
|
||||
"""Triton path of quantize_int8_rows: one extra read of x instead of the
|
||||
multi-kernel torch chain. Falls back to the torch ops."""
|
||||
rows, K = x.shape
|
||||
if not (_triton_available() and x.is_cuda):
|
||||
return quantize_int8_rows(x)
|
||||
x = x.contiguous()
|
||||
q = torch.empty(rows, K, device=x.device, dtype=torch.int8)
|
||||
scales = torch.empty(rows, device=x.device, dtype=torch.float32)
|
||||
kernel, _ = _get_int8_kernels()
|
||||
kernel[(rows,)](x, q, scales, K, BLOCK_K=2048 if K >= 2048 else K, num_warps=8)
|
||||
return q, scales
|
||||
|
||||
|
||||
def _int8_epilogue(
|
||||
i32: torch.Tensor,
|
||||
a_scales: torch.Tensor,
|
||||
w_scales: torch.Tensor,
|
||||
bias,
|
||||
out_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""out = i32 * a_scales[:, None] * w_scales[None, :] (+ bias), in out_dtype."""
|
||||
m, n = i32.shape
|
||||
if _triton_available() and i32.is_cuda:
|
||||
out = torch.empty(m, n, device=i32.device, dtype=out_dtype)
|
||||
_, kernel = _get_int8_kernels()
|
||||
grid = (m, -(-n // 1024))
|
||||
kernel[grid](
|
||||
i32,
|
||||
a_scales,
|
||||
w_scales,
|
||||
bias if bias is not None else i32,
|
||||
out,
|
||||
n,
|
||||
HAS_BIAS=bias is not None,
|
||||
BLOCK_N=1024,
|
||||
num_warps=4,
|
||||
)
|
||||
return out
|
||||
out = i32.float() * w_scales
|
||||
out = out * a_scales.unsqueeze(1)
|
||||
if bias is not None:
|
||||
out = out + bias.float()
|
||||
return out.to(out_dtype)
|
||||
|
||||
|
||||
_int8_mm_ok = None
|
||||
|
||||
|
||||
def _int8_gemm_supported(device) -> bool:
|
||||
global _int8_mm_ok
|
||||
device = torch.device(device)
|
||||
if device.type != "cuda" or not torch.cuda.is_available():
|
||||
supported = False
|
||||
else:
|
||||
if _int8_mm_ok is None:
|
||||
try:
|
||||
a = torch.zeros(32, 64, dtype=torch.int8, device=device)
|
||||
b = torch.zeros(64, 32, dtype=torch.int8, device=device)
|
||||
torch._int_mm(a, b)
|
||||
_int8_mm_ok = True
|
||||
except Exception:
|
||||
_int8_mm_ok = False
|
||||
supported = _int8_mm_ok
|
||||
global _warned_no_int8
|
||||
if not supported and not _warned_no_int8:
|
||||
_warned_no_int8 = True
|
||||
print_acc(
|
||||
f"ConvRot: int8 matmul (torch._int_mm) is not usable on this device "
|
||||
f"({device}). Inference falls back to dequantized bf16 matmuls: correct "
|
||||
"output but NO speedup, and inference activations stay unquantized "
|
||||
"(W8A16 numerics instead of W8A8). The training path is unaffected "
|
||||
"(it always simulates W8A8 via fake-quant)."
|
||||
)
|
||||
return supported
|
||||
|
||||
|
||||
_warned_no_int8 = False
|
||||
|
||||
|
||||
class ConvRotInt8Quantizer(OstrisQuantizer):
|
||||
"""ConvRot W8A8 backend: shared regular-Hadamard rotation + per-token /
|
||||
per-output-channel symmetric int8 with torch._int_mm. One instance per qtype,
|
||||
shareable across modules."""
|
||||
|
||||
def __init__(self, rot_size: int = 256):
|
||||
self.rot_size = rot_size
|
||||
|
||||
def _rot_for(self, d: int) -> int:
|
||||
return min(self.rot_size, largest_pow4_divisor(d))
|
||||
|
||||
def can_quantize(self, module: torch.nn.Linear) -> bool:
|
||||
d = module.in_features
|
||||
if d % BLOCK != 0 or module.out_features % 8 != 0 or self._rot_for(d) < 16:
|
||||
if d not in _skip_warned:
|
||||
_skip_warned.add(d)
|
||||
print_acc(
|
||||
f"ConvRot: skipping linears with in_features={d} "
|
||||
f"(needs in divisible by 16, out by 8, and a power-of-4 block >= 16)"
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None:
|
||||
rot = self._rot_for(module.in_features)
|
||||
q, scales = quantize_int8_rows(rotate(weight_fp32, rot))
|
||||
module.register_buffer("cr8_qdata", q, persistent=False)
|
||||
# fp32 scales stored as a uint8 byte view (see convrot4: module.to(dtype=...)
|
||||
# would otherwise cast them)
|
||||
module.register_buffer("cr8_scales", scales.view(torch.uint8), persistent=False)
|
||||
module.cr8_rot_size = rot
|
||||
|
||||
@staticmethod
|
||||
def _scales(module) -> torch.Tensor:
|
||||
return module.cr8_scales.view(torch.float32)
|
||||
|
||||
def _dequantize_rotated(self, module, dtype: torch.dtype) -> torch.Tensor:
|
||||
w = module.cr8_qdata.float() * self._scales(module).unsqueeze(1)
|
||||
return w.to(dtype)
|
||||
|
||||
def dequantize(self, module) -> torch.Tensor:
|
||||
return rotate(
|
||||
self._dequantize_rotated(module, torch.float32), module.cr8_rot_size
|
||||
)
|
||||
|
||||
def requantize_(self, module, fp_weight: torch.Tensor) -> None:
|
||||
w = fp_weight.to(device=module.cr8_qdata.device, dtype=torch.float32)
|
||||
q, scales = quantize_int8_rows(rotate(w, module.cr8_rot_size))
|
||||
module.cr8_qdata = q
|
||||
module.cr8_scales = scales.view(torch.uint8)
|
||||
|
||||
def forward(self, module, x: torch.Tensor) -> torch.Tensor:
|
||||
rot = module.cr8_rot_size
|
||||
in_f, out_f = module.in_features, module.out_features
|
||||
x_rot = rotate(x, rot)
|
||||
x2d = x_rot.reshape(-1, in_f)
|
||||
m = x2d.shape[0]
|
||||
|
||||
if x.requires_grad:
|
||||
# training: straight-through fake-quant so adapters see deployment
|
||||
# W8A8 numerics; differentiable bf16 matmul. gated on requires_grad
|
||||
# alone so both gradient-checkpoint passes take the same branch
|
||||
with torch.no_grad():
|
||||
aq, a_s = quantize_int8_rows(x2d.detach())
|
||||
x_dq = (aq.float() * a_s.unsqueeze(1)).to(x.dtype)
|
||||
w = self._dequantize_rotated(module, x.dtype)
|
||||
x_ste = x2d + (x_dq - x2d).detach()
|
||||
out = F.linear(x_ste, w, module.bias)
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
|
||||
if _int8_gemm_supported(x.device):
|
||||
pad = (-m) % 32
|
||||
if pad:
|
||||
x2d = F.pad(x2d, (0, 0, 0, pad))
|
||||
aq, a_s = quantize_int8_rows_fused(x2d)
|
||||
i32 = torch._int_mm(aq, module.cr8_qdata.t())
|
||||
out = _int8_epilogue(i32, a_s, self._scales(module), module.bias, x.dtype)
|
||||
if pad:
|
||||
out = out[:m]
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
|
||||
w = self._dequantize_rotated(module, x.dtype)
|
||||
out = F.linear(x2d, w, module.bias)
|
||||
return out.reshape(*x.shape[:-1], out_f)
|
||||
@@ -93,7 +93,9 @@ class OstrisLinear(torch.nn.Linear):
|
||||
# emit a plain full precision weight so full-model saves need no special casing
|
||||
destination[prefix + "weight"] = self.dequantize_weight()
|
||||
if self.bias is not None:
|
||||
destination[prefix + "bias"] = self.bias if keep_vars else self.bias.detach()
|
||||
destination[prefix + "bias"] = (
|
||||
self.bias if keep_vars else self.bias.detach()
|
||||
)
|
||||
|
||||
|
||||
def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
|
||||
@@ -101,16 +103,21 @@ def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
|
||||
does not belong to a custom backend. Add new backends here."""
|
||||
from toolkit.util.orbit_quant import ORBIT_QTYPES, OrbitQuantizer
|
||||
from toolkit.util.orbit_vq_quant import ORBIT_VQ_QTYPES, OrbitVQQuantizer
|
||||
from toolkit.util.convrot_quant import CONVROT_QTYPES, get_convrot_quantizer
|
||||
|
||||
if qtype in ORBIT_QTYPES:
|
||||
return OrbitQuantizer(ORBIT_QTYPES[qtype])
|
||||
if qtype in ORBIT_VQ_QTYPES:
|
||||
return OrbitVQQuantizer(**ORBIT_VQ_QTYPES[qtype])
|
||||
if qtype in CONVROT_QTYPES:
|
||||
return get_convrot_quantizer(qtype)
|
||||
return None
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def convert_linear_to_ostris(module: torch.nn.Linear, quantizer: OstrisQuantizer) -> bool:
|
||||
def convert_linear_to_ostris(
|
||||
module: torch.nn.Linear, quantizer: OstrisQuantizer
|
||||
) -> bool:
|
||||
"""Quantize an nn.Linear in place (class swap). Returns True if the module was
|
||||
converted (or already was), False if it is not a candidate."""
|
||||
if isinstance(module, OstrisLinear):
|
||||
|
||||
Reference in New Issue
Block a user