github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add convrot4 and convrot8 quantizations to the ui. A lot of improvements to the convrot quantization methods.

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-11 13:49:41 -06:00
parent 933ca1c517
commit 095d6e7418
6 changed files with 557 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

226
toolkit/util/convrot_quant.py
View File

@@ -66,11 +66,9 @@ 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()
@@ -162,15 +160,22 @@ def dequantize_nvfp4(
K: int,
dtype: torch.dtype,
) -> torch.Tensor:
# single-pass triton path when available: the torch chain below is ~7 full-size
# elementwise passes with fp32 intermediates, which made every convrot4 training
# backward pay a dequant cost comparable to the gradient matmul itself
if _triton_available() and packed.is_cuda and dtype in (torch.bfloat16, torch.float16, torch.float32):
return _fp4_dequant_op(
packed, scales.view(torch.uint8), pts.reshape(1).view(torch.uint8),
str(dtype).split(".")[-1],
)
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
# the lookup table is built inline (NOT module-cached): this function runs inside
# custom-op backwards, which torch.compile traces with fake tensors where a
# pre-existing real tensor is illegal; an in-trace constructed constant is fine
vals = torch.tensor(
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], 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)
@@ -229,7 +234,9 @@ def _get_kernel():
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)
# clamp to the e4m3 max so an oversized block can't overflow the fp8
# conversion (matters now that activations use a fixed pts=1)
scale8 = tl.minimum(amax / (6.0 * pts), 448.0).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
@@ -298,8 +305,11 @@ def _nvfp4_act_quant_op(x: torch.Tensor) -> list[torch.Tensor]:
rows, K = x.shape
rows_pad = -(-rows // 16) * 16
x = x.contiguous()
pts = x.float().abs().amax() / (F4_MAX * F8_E4M3_MAX)
pts = torch.where(pts > 0, pts, torch.ones_like(pts))
# activations use a FIXED per-tensor scale of 1: their block scales fit the
# e4m3 range natively (unlike tiny weight magnitudes, which keep dynamic pts),
# measured quality-neutral on realistic distributions — and it removes a full
# activation read (global amax) plus a device sync from every forward
pts = torch.ones((), device=x.device)
packed = torch.empty(rows_pad, K // 2, device=x.device, dtype=torch.uint8)
if rows_pad != rows:
packed[rows:].zero_()
@@ -333,11 +343,12 @@ def quantize_nvfp4_fused(x: torch.Tensor, blocked_scales: bool = False):
rows, K = x.shape
if not (_triton_available() and x.is_cuda and K % 16 == 0):
if blocked_scales:
# match the custom op: rows padded to a multiple of 16 for _scaled_mm
# match the custom op: rows padded to a multiple of 16 for _scaled_mm,
# fixed pts=1 for activations (see _nvfp4_act_quant_op)
rows_pad = -(-rows // 16) * 16
if rows_pad != rows:
x = F.pad(x, (0, 0, 0, rows_pad - rows))
packed, scales, pts = quantize_nvfp4(x)
packed, scales, pts = quantize_nvfp4(x, pts=torch.ones((), device=x.device))
return packed, to_blocked(scales), pts
return quantize_nvfp4(x)
if blocked_scales:
@@ -353,6 +364,77 @@ def quantize_nvfp4_fused(x: torch.Tensor, blocked_scales: bool = False):
return packed, scales, pts
# ---------------- fp4 dequant kernel (backward hot path) ----------------
_dequant_kernel = None
def _get_dequant_kernel():
global _dequant_kernel
if _dequant_kernel is not None:
return _dequant_kernel
import triton
import triton.language as tl
@triton.jit
def nvfp4_dequant_kernel(
q_ptr, s_ptr, pts_ptr, out_ptr, K,
BLOCK_B: tl.constexpr,
):
row = tl.program_id(0)
pid_k = tl.program_id(1)
pts = tl.load(pts_ptr)
offs_b = pid_k * BLOCK_B + tl.arange(0, BLOCK_B)
bmask = offs_b < K // 2
byte = tl.load(q_ptr + row * (K // 2) + offs_b, mask=bmask, other=0)
codes = tl.interleave(byte & 15, byte >> 4) # (2*BLOCK_B,), column order
m = (codes & 7).to(tl.float32)
# arithmetic e2m1 decode ([0, .5, 1, 1.5, 2, 3, 4, 6]), exact
mag = tl.where(m < 2, m * 0.5, tl.exp2(tl.floor(m / 2) - 1) * (1 + (m % 2) * 0.5))
v = tl.where((codes & 8) > 0, -mag, mag)
n_s: tl.constexpr = (2 * BLOCK_B) // 16
offs_s = pid_k * n_s + tl.arange(0, n_s)
s = tl.load(s_ptr + row * (K // 16) + offs_s, mask=offs_s < K // 16, other=0.0)
vb = tl.reshape(v, (n_s, 16)) * (s.to(tl.float32) * pts)[:, None]
out = tl.reshape(vb, (2 * BLOCK_B,))
offs_v = pid_k * (2 * BLOCK_B) + tl.arange(0, 2 * BLOCK_B)
tl.store(out_ptr + row * K + offs_v, out.to(out_ptr.dtype.element_ty), mask=offs_v < K)
_dequant_kernel = nvfp4_dequant_kernel
return _dequant_kernel
# custom op so the kernel stays opaque where it matters most: inside the fp4
# training op's registered backward, which torch.compile traces with fake tensors
@torch.library.custom_op("ostris::convrot_fp4_dequant", mutates_args=())
def _fp4_dequant_op(
packed: torch.Tensor,
scales_u8: torch.Tensor,
pts_u8: torch.Tensor,
out_dtype: str,
) -> torch.Tensor:
rows, half = packed.shape
out = torch.empty(rows, half * 2, device=packed.device, dtype=getattr(torch, out_dtype))
kernel = _get_dequant_kernel()
block_b = 1024
grid = (rows, -(-half // block_b))
kernel[grid](
packed.contiguous(),
scales_u8.view(torch.float8_e4m3fn),
pts_u8.view(torch.float32),
out, half * 2,
BLOCK_B=block_b, num_warps=4,
)
return out
@_fp4_dequant_op.register_fake
def _fp4_dequant_fake(packed, scales_u8, pts_u8, out_dtype):
rows, half = packed.shape
return torch.empty(rows, half * 2, device=packed.device, dtype=getattr(torch, out_dtype))
# ---------------- backend ----------------
@@ -379,6 +461,67 @@ def _fp4_gemm_supported(device) -> bool:
return supported
# the fp4 training-path linear: forward VALUE is the real fp4 tensor-core gemm
# (bit-identical to the inference path), gradient is the straight-through estimate
# d y / d x_rot ~= dequant(W'). the backward re-dequantizes the weight from the fp4
# codes instead of saving a bf16 copy per layer (F.linear in the old fake-quant path
# retained a dequantized weight for every layer of the graph — a full bf16 model of
# extra train-step vram) and x is not saved at all.
@torch.library.custom_op("ostris::convrot_fp4_linear_ste", mutates_args=())
def _fp4_linear_ste_op(
x2d: torch.Tensor,
qdata: torch.Tensor,
scales_u8: torch.Tensor,
scales_blocked_u8: torch.Tensor,
pts_u8: torch.Tensor,
bias: Optional[torch.Tensor],
out_dtype: str,
) -> torch.Tensor:
m = x2d.shape[0]
aq, a_scales_blocked, a_pts = quantize_nvfp4_fused(x2d, blocked_scales=True)
out = torch._scaled_mm(
aq.view(torch.float4_e2m1fn_x2),
qdata.view(torch.float4_e2m1fn_x2).t(),
a_scales_blocked.view(torch.float8_e4m3fn),
scales_blocked_u8.view(torch.float8_e4m3fn),
out_dtype=getattr(torch, out_dtype),
)
if out.shape[0] != m:
out = out[:m]
s = (a_pts * pts_u8.view(torch.float32).reshape(())).to(out.dtype)
if bias is not None:
return torch.addcmul(bias, out, s)
return out * s
@_fp4_linear_ste_op.register_fake
def _fp4_linear_ste_fake(x2d, qdata, scales_u8, scales_blocked_u8, pts_u8, bias, out_dtype):
return torch.empty(
x2d.shape[0], qdata.shape[0], device=x2d.device, dtype=getattr(torch, out_dtype)
)
def _fp4_linear_ste_setup(ctx, inputs, output):
x2d, qdata, scales_u8, scales_blocked_u8, pts_u8, bias, out_dtype = inputs
ctx.save_for_backward(qdata, scales_u8, pts_u8)
def _fp4_linear_ste_backward(ctx, grad):
qdata, scales_u8, pts_u8 = ctx.saved_tensors
out_f, in_half = qdata.shape
w = dequantize_nvfp4(
qdata, scales_u8.view(torch.float8_e4m3fn),
pts_u8.view(torch.float32).reshape(()),
out_f, in_half * 2, grad.dtype,
)
return grad @ w, None, None, None, None, None, None
_fp4_linear_ste_op.register_autograd(
_fp4_linear_ste_backward, setup_context=_fp4_linear_ste_setup
)
class ConvRotQuantizer(OstrisQuantizer):
"""ConvRot W4A4 backend. One instance per qtype, shareable across modules."""
@@ -451,16 +594,24 @@ class ConvRotQuantizer(OstrisQuantizer):
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]
m = x.numel() // in_f
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
# training path, gated on requires_grad alone (not is_grad_enabled) so
# both passes of gradient checkpointing take the same branch
if _fp4_gemm_supported(x.device):
# fp4 tensor-core forward (bit-identical to the inference path)
# with a straight-through analytic backward
x2d = rotate(x, rot).reshape(-1, in_f)
out = _fp4_linear_ste_op(
x2d, module.cr_qdata, module.cr_scales,
module.cr_scales_blocked, module.cr_pts, module.bias,
str(x.dtype).split(".")[-1],
)
return out.reshape(*x.shape[:-1], out_f)
# no fp4 hardware: straight-through fake-quant of the activation and a
# differentiable bf16 matmul against the dequantized rotated weight
x2d = rotate(x, rot).reshape(-1, in_f)
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)
@@ -471,8 +622,14 @@ class ConvRotQuantizer(OstrisQuantizer):
if _fp4_gemm_supported(x.device):
# row padding for _scaled_mm happens inside the act-quant op (compile
# safety); slice the mm output back to m rows (a contiguous prefix)
aq, a_scales_blocked, a_pts = quantize_nvfp4_fused(x2d, blocked_scales=True)
# safety); slice the mm output back to m rows (a contiguous prefix).
# NOTE: no fused rotate+quant here (unlike convrot8): the e2m1 packing
# needs tl.reshape/tl.split of dot-derived tensors, which this triton
# backend miscompiles when the dot sits in/after a loop (element order
# scrambles). the int8 kernel avoids those primitives and fuses safely.
aq, a_scales_blocked, a_pts = quantize_nvfp4_fused(
rotate(x, rot).reshape(-1, in_f), blocked_scales=True
)
out = torch._scaled_mm(
aq.view(torch.float4_e2m1fn_x2),
module.cr_qdata.view(torch.float4_e2m1fn_x2).t(),
@@ -491,7 +648,7 @@ class ConvRotQuantizer(OstrisQuantizer):
# no fp4 hardware: dequantized matmul (correct, no speedup)
w = self._dequantize_rotated(module, x.dtype)
out = F.linear(x2d, w, module.bias)
out = F.linear(rotate(x, rot).reshape(-1, in_f), w, module.bias)
return out.reshape(*x.shape[:-1], out_f)
@@ -803,22 +960,26 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
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]
m = x.numel() // in_f
if x.requires_grad:
# training: gated on requires_grad alone so both gradient-checkpoint
# passes take the same branch
if _int8_gemm_supported(x.device):
# int8 tensor-core forward (bit-identical to the inference path)
# with a straight-through analytic backward
# with a straight-through analytic backward. the rotation stays a
# cublas matmul: an in-kernel tl.dot rotation was tried and measured
# SLOWER on every shape (small-tile dots at low tensor-core
# utilization, computed twice for the amax and quant passes, cost
# more than the activation round-trips they saved)
out = _int8_linear_ste_op(
x2d, module.cr8_qdata, module.cr8_scales, module.bias,
rotate(x, rot).reshape(-1, in_f),
module.cr8_qdata, module.cr8_scales, module.bias,
str(x.dtype).split(".")[-1],
)
return out.reshape(*x.shape[:-1], out_f)
# no int8 hardware: straight-through fake-quant + bf16 matmul
x2d = rotate(x, rot).reshape(-1, in_f)
with torch.no_grad():
aq, a_s = quantize_int8_rows(x2d.detach())
x_dq = (aq.float() * a_s.unsqueeze(1)).to(x.dtype)
@@ -830,11 +991,12 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
if _int8_gemm_supported(x.device):
# row padding for _int_mm happens inside the act-quant op (compile
# safety); slice the mm output back to m rows (a contiguous prefix)
aq, a_s = _int8_act_quant_padded(x2d)
aq, a_s = _int8_act_quant_padded(rotate(x, rot).reshape(-1, in_f))
i32 = torch._int_mm(aq, module.cr8_qdata.t())
out = _int8_epilogue(i32[:m], a_s[:m], self._scales(module), module.bias, x.dtype)
return out.reshape(*x.shape[:-1], out_f)
w = self._dequantize_rotated(module, x.dtype)
out = F.linear(x2d, w, module.bias)
out = F.linear(rotate(x, rot).reshape(-1, in_f), w, module.bias)
return out.reshape(*x.shape[:-1], out_f)

19
toolkit/util/ostris_quant.py
View File

@@ -28,6 +28,10 @@ class OstrisQuantizer:
other methods. One backend instance may be shared by many modules.
"""
# the qtype string this instance was resolved from (stamped by
# get_ostris_quantizer); pre-quantized saves need it to restore the backend
qtype: Optional[str] = None
def can_quantize(self, module: torch.nn.Linear) -> bool:
"""Whether this backend can quantize the given linear (e.g. shape constraints)."""
return True
@@ -105,13 +109,16 @@ def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
from toolkit.util.orbit_vq_quant import ORBIT_VQ_QTYPES, OrbitVQQuantizer
from toolkit.util.convrot_quant import CONVROT_QTYPES, get_convrot_quantizer
quantizer = None
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
quantizer = OrbitQuantizer(ORBIT_QTYPES[qtype])
elif qtype in ORBIT_VQ_QTYPES:
quantizer = OrbitVQQuantizer(**ORBIT_VQ_QTYPES[qtype])
elif qtype in CONVROT_QTYPES:
quantizer = get_convrot_quantizer(qtype)
if quantizer is not None:
quantizer.qtype = qtype
return quantizer
@torch.no_grad()