From 28b05ee4ed44556c9f743cba9a85c3d22c6619cd Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Mon, 13 Jul 2026 07:52:48 -0600 Subject: [PATCH] Added convrotbitnet 1.58 bit quantization --- scripts/test_quantizations.py | 2 +- toolkit/util/convrot_quant.py | 182 ++++++++++++++++++++++++++++++++- ui/src/app/jobs/new/options.ts | 1 + 3 files changed, 183 insertions(+), 2 deletions(-) diff --git a/scripts/test_quantizations.py b/scripts/test_quantizations.py index 034d56a2..f6f85bc2 100644 --- a/scripts/test_quantizations.py +++ b/scripts/test_quantizations.py @@ -48,7 +48,7 @@ VRAM_TOKENS = 4096 QTYPES = [ "bf16", "qfloat8", "float8", "orbit4", "orbitvq4", "convrot8", "convrot4", "convrotint7", "convrotint6", "convrotint5", "convrotint4", "convrotint3", - "convrotint2", + "convrotint2", "convrotbitnet", ] STACK_KEY = f"{VRAM_BLOCKS}-block stack" diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index f9ed9adb..3211dbbf 100644 --- a/toolkit/util/convrot_quant.py +++ b/toolkit/util/convrot_quant.py @@ -51,7 +51,7 @@ import torch.nn.functional as F from toolkit.print import print_acc from toolkit.util.ostris_quant import OstrisQuantizer -CONVROT_QTYPES = ("convrot4", "convrot8") + tuple( +CONVROT_QTYPES = ("convrot4", "convrot8", "convrotbitnet") + tuple( f"convrotint{b}" for b in range(2, 9) ) @@ -61,6 +61,8 @@ def get_convrot_quantizer(qtype: str): return ConvRotQuantizer(rot_size=256) if qtype == "convrot8": return ConvRotInt8Quantizer(rot_size=256) + if qtype == "convrotbitnet": + return ConvRotBitNetQuantizer(rot_size=256) if qtype.startswith("convrotint"): bits = int(qtype[len("convrotint"):]) if 2 <= bits <= 8: @@ -1189,6 +1191,38 @@ def unpack_intn_rows_grouped( return torch.round(codes * ratio).clamp_(-127, 127).to(torch.int8) +# --- convrotbitnet: 1.6-bit base-3 storage of the ternary width -------------- +# +# BitNet-b1.58 style: the codes and scales are EXACTLY convrotint2's (ternary +# {-1,0,1}, MSE-optimal group scales), only the storage differs — 5 ternary +# codes per byte (3^5 = 243 <= 256) = 1.6 bits/weight instead of 2.0. Rows are +# padded to a multiple of 5 codes; the pad is never read back. + + +def pack_ternary_rows(q: torch.Tensor) -> torch.Tensor: + """Pack int8 ternary codes (rows, K) into (rows, ceil(K/5)) uint8 base-3.""" + rows, K = q.shape + bpr = -(-K // 5) + u = q.to(torch.int16) + 1 + if bpr * 5 != K: + u = F.pad(u, (0, bpr * 5 - K)) + w3 = torch.tensor([1, 3, 9, 27, 81], device=q.device, dtype=torch.int16) + return (u.view(rows, bpr, 5) * w3).sum(-1).to(torch.uint8) + + +def unpack_ternary_rows_grouped( + packed: torch.Tensor, gratio: torch.Tensor, rows: int, cols: int +) -> torch.Tensor: + """Torch fallback: base-3 bytes -> ternary codes -> row-grid int8.""" + bpr = packed.shape[1] + p3 = torch.tensor([1, 3, 9, 27, 81], device=packed.device, dtype=torch.int16) + u = (packed.to(torch.int16).unsqueeze(-1) // p3) % 3 + codes = (u - 1).reshape(rows, bpr * 5)[:, :cols].float() + group = cols // gratio.shape[1] + ratio = gratio.repeat_interleave(group, dim=1) + return torch.round(codes * ratio).clamp_(-127, 127).to(torch.int8) + + _intn_kernel = None @@ -1268,6 +1302,46 @@ def _get_intn_grouped_kernel(): return _intn_grouped_kernel +_bitnet_kernel = None + + +def _get_bitnet_kernel(): + global _bitnet_kernel + if _bitnet_kernel is not None: + return _bitnet_kernel + import triton + import triton.language as tl + from triton.language.extra import libdevice + + @triton.jit + def bitnet_unpack_kernel( + p_ptr, r_ptr, o_ptr, n_bytes, bpr, K, group, ngprow, BLOCK: tl.constexpr + ): + # one thread-lane per packed byte -> 5 output codes. a byte's 5 codes can + # straddle a k-group boundary (group % 5 != 0), so the scale ratio is + # gathered per output element rather than per byte. + pid = tl.program_id(0) + i = pid * BLOCK + tl.arange(0, BLOCK) + m = i < n_bytes + b = tl.load(p_ptr + i, mask=m, other=0).to(tl.int32) + row = i // bpr + bcol = i - row * bpr + j = tl.arange(0, 8) + col = bcol[:, None] * 5 + j[None, :] + cm = m[:, None] & (j[None, :] < 5) & (col < K) + p3 = tl.where( + j < 1, 1, tl.where(j < 2, 3, tl.where(j < 3, 9, tl.where(j < 4, 27, 81))) + ) + code = ((b[:, None] // p3[None, :]) % 3 - 1).to(tl.float32) + ratio = tl.load(r_ptr + row[:, None] * ngprow + col // group, mask=cm, other=1.0) + v = libdevice.rint(code * ratio) + v = tl.minimum(tl.maximum(v, -127.0), 127.0) + tl.store(o_ptr + row[:, None] * K + col, v.to(tl.int8), mask=cm) + + _bitnet_kernel = bitnet_unpack_kernel + return _bitnet_kernel + + def _unpack_intn_impl(packed: torch.Tensor, bits: int, rows: int, cols: int) -> torch.Tensor: if _triton_available() and packed.is_cuda: out = torch.empty(rows, cols, device=packed.device, dtype=torch.int8) @@ -1384,6 +1458,85 @@ _intn_linear_ste_op.register_autograd( ) +def _unpack_bitnet_impl( + packed: torch.Tensor, gratio: torch.Tensor, rows: int, cols: int +) -> torch.Tensor: + if _triton_available() and packed.is_cuda: + out = torch.empty(rows, cols, device=packed.device, dtype=torch.int8) + bpr = packed.shape[1] + n_bytes = rows * bpr + ngprow = gratio.shape[1] + BLOCK = 256 + kernel = _get_bitnet_kernel() + kernel[(-(-n_bytes // BLOCK),)]( + packed, gratio, out, n_bytes, bpr, cols, cols // ngprow, ngprow, + BLOCK=BLOCK, num_warps=4, + ) + return out + return unpack_ternary_rows_grouped(packed, gratio, rows, cols) + + +@torch.library.custom_op("ostris::convrot_bitnet_unpack", mutates_args=()) +def _bitnet_unpack_op( + packed: torch.Tensor, gratio: torch.Tensor, rows: int, cols: int +) -> torch.Tensor: + return _unpack_bitnet_impl(packed, gratio, rows, cols) + + +@_bitnet_unpack_op.register_fake +def _bitnet_unpack_fake(packed, gratio, rows, cols): + return torch.empty(rows, cols, device=packed.device, dtype=torch.int8) + + +@torch.library.custom_op("ostris::convrot_bitnet_linear_ste", mutates_args=()) +def _bitnet_linear_ste_op( + x2d: torch.Tensor, + packed: torch.Tensor, + gratio: torch.Tensor, + w_scales_u8: torch.Tensor, + bias: Optional[torch.Tensor], + out_dtype: str, +) -> torch.Tensor: + m, K = x2d.shape + rows = w_scales_u8.view(torch.float32).numel() + qdata = _unpack_bitnet_impl(packed, gratio, rows, K) + aq, a_s = _int8_act_quant_padded(x2d) + i32 = torch._int_mm(aq, qdata.t()) + return _int8_epilogue( + i32[:m], a_s[:m], w_scales_u8.view(torch.float32), bias, + getattr(torch, out_dtype), + ) + + +@_bitnet_linear_ste_op.register_fake +def _bitnet_linear_ste_fake(x2d, packed, gratio, w_scales_u8, bias, out_dtype): + return torch.empty( + x2d.shape[0], w_scales_u8.view(torch.float32).numel(), + device=x2d.device, dtype=getattr(torch, out_dtype), + ) + + +def _bitnet_linear_ste_setup(ctx, inputs, output): + x2d, packed, gratio, w_scales_u8, bias, out_dtype = inputs + ctx.cols = x2d.shape[1] + ctx.save_for_backward(packed, gratio, w_scales_u8) + + +def _bitnet_linear_ste_backward(ctx, grad): + packed, gratio, w_scales_u8 = ctx.saved_tensors + w_scales = w_scales_u8.view(torch.float32) + # opaque op, not the raw impl: compile traces registered backwards with + # fake tensors + qdata = _bitnet_unpack_op(packed, gratio, w_scales.numel(), ctx.cols) + w = qdata.to(grad.dtype) * w_scales.to(grad.dtype).unsqueeze(1) + return grad @ w, None, None, None, None, None + + +_bitnet_linear_ste_op.register_autograd( + _bitnet_linear_ste_backward, setup_context=_bitnet_linear_ste_setup +) + + class ConvRotIntNQuantizer(ConvRotInt8Quantizer): """ConvRot W{n}A8 backend for n in 2..8: convrot8 numerics with an n-bit weight grid and bitpacked storage. One instance per bit-width, shareable @@ -1468,3 +1621,30 @@ class ConvRotIntNQuantizer(ConvRotInt8Quantizer): if gratio is not None: module.crn_gratio = gratio.view(torch.uint8) + +class ConvRotBitNetQuantizer(ConvRotIntNQuantizer): + """BitNet-b1.58 style W1.58A8: convrotint2's exact ternary codes and + MSE-optimal group scales, stored base-3 at 5 codes/byte (1.6 bits/weight + instead of 2.0). Numerically bit-identical to convrotint2.""" + + def __init__(self, rot_size: int = 256): + super().__init__(bits=2, rot_size=rot_size) + + def _quantize_rotated(self, w_rot: torch.Tensor, in_features: int): + group = self._group_for(in_features) + q, gscales = _optimal_group_quantize(w_rot, self.bits, group) + gratio, rscales = _intn_group_ratio_and_rscales(gscales, self.bits) + return pack_ternary_rows(q), rscales, gratio.contiguous() + + def _qdata(self, module) -> torch.Tensor: + return _bitnet_unpack_op( + module.crn_qdata, module.crn_gratio.view(torch.float32), + module.out_features, module.in_features, + ) + + def _linear_ste(self, module, x2d: torch.Tensor, out_dtype: str) -> torch.Tensor: + return _bitnet_linear_ste_op( + x2d, module.crn_qdata, module.crn_gratio.view(torch.float32), + module.crn_scales, module.bias, out_dtype, + ) + diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts index 255f5013..d907964d 100644 --- a/ui/src/app/jobs/new/options.ts +++ b/ui/src/app/jobs/new/options.ts @@ -1246,6 +1246,7 @@ export const quantizationOptions: SelectOption[] = [ { value: 'convrotint4', label: '4bit convrot' }, { value: 'convrotint3', label: '3bit convrot' }, { value: 'convrotint2', label: '2bit convrot' }, + { value: 'convrotbitnet', label: '1.58bit convrot (bitnet)' }, { value: 'uint7', label: '7 bit' }, { value: 'uint6', label: '6 bit' }, { value: 'uint5', label: '5 bit' },