Fix some errors for specific quants. Enable saving quantizations.

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", "convrotbitnet",
+    "convrotint2", "convrotbitnet", "convrotcomfyw4a4",
 ]
 
 STACK_KEY = f"{VRAM_BLOCKS}-block stack"

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", "convrotbitnet") + tuple(
+CONVROT_QTYPES = ("convrot4", "convrot8", "convrotbitnet", "convrotcomfyw4a4") + tuple(
     f"convrotint{b}" for b in range(2, 9)
 )
 
@@ -63,6 +63,8 @@ def get_convrot_quantizer(qtype: str):
         return ConvRotInt8Quantizer(rot_size=256)
     if qtype == "convrotbitnet":
         return ConvRotBitNetQuantizer(rot_size=256)
+    if qtype == "convrotcomfyw4a4":
+        return ConvRotComfyW4A4Quantizer()
     if qtype.startswith("convrotint"):
         bits = int(qtype[len("convrotint"):])
         if 2 <= bits <= 8:
@@ -334,7 +336,8 @@ def _get_kernel():
 def _launch_nvfp4_kernel(x, packed, scales, pts, blocked_scales: bool):
     rows, K = x.shape
     n_col_tiles = -(-(K // BLOCK) // 4)
-    BLOCK_K = 2048 if K >= 2048 else K
+    # triton block shapes must be powers of 2; loads/stores are masked on offs < K
+    BLOCK_K = min(2048, 1 << (K - 1).bit_length())
     grid = (rows, -(-K // BLOCK_K))
     _get_kernel()[grid](
         x, packed, scales, pts,
@@ -617,6 +620,28 @@ class ConvRotQuantizer(OstrisQuantizer):
     def _pts(module) -> torch.Tensor:
         return module.cr_pts.view(torch.float32).reshape(())
 
+    def _rot(self, module) -> int:
+        return module.cr_rot_size
+
+    def fake_quant_rotated_weight(self, module, w_rot: torch.Tensor) -> torch.Tensor:
+        """dequant(quant(w_rot)) on the deployed e2m1 grid with the module's STORED
+        e4m3 block scales (no scale re-optimization) — the value half of the QAT
+        straight-through estimator. Returns float32."""
+        rows, K = w_rot.shape
+        s = module.cr_scales.view(torch.float8_e4m3fn).float() * self._pts(module)
+        denom = s.unsqueeze(-1)
+        safe = torch.where(denom > 0, denom, torch.ones_like(denom))
+        z = (w_rot.float().view(rows, K // BLOCK, BLOCK) / safe).clamp(-F4_MAX, F4_MAX)
+        edges = _cached(
+            _edges_cache, str(w_rot.device),
+            lambda: torch.tensor(_E2M1_EDGES, device=w_rot.device),
+        )
+        vals = torch.tensor(
+            [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=w_rot.device
+        )
+        recon = vals[torch.bucketize(z.abs(), edges)] * torch.sign(z) * safe
+        return recon.view(rows, K)
+
     def _dequantize_rotated(self, module, dtype: torch.dtype) -> torch.Tensor:
         return dequantize_nvfp4(
             module.cr_qdata,
@@ -631,6 +656,27 @@ class ConvRotQuantizer(OstrisQuantizer):
         w = self._dequantize_rotated(module, torch.float32)
         return rotate(w, module.cr_rot_size)  # self-inverse
 
+    @torch.no_grad()
+    def requantize_codes_(self, module, fp_weight: torch.Tensor) -> None:
+        """Re-quantize only the codes on the module's STORED scales — the grid a
+        QAT run trains against. Re-optimizing the scales here would re-grid every
+        weight and destroy the code adjustments training made."""
+        w = fp_weight.to(device=module.cr_qdata.device, dtype=torch.float32)
+        w_rot = rotate(w, module.cr_rot_size)
+        rows, K = w_rot.shape
+        s = module.cr_scales.view(torch.float8_e4m3fn).float() * self._pts(module)
+        denom = s.unsqueeze(-1)
+        safe = torch.where(denom > 0, denom, torch.ones_like(denom))
+        z = (w_rot.view(rows, K // BLOCK, BLOCK) / safe).clamp(-F4_MAX, F4_MAX)
+        edges = _cached(
+            _edges_cache, str(w_rot.device),
+            lambda: torch.tensor(_E2M1_EDGES, device=w_rot.device),
+        )
+        z = z.reshape(rows, K)
+        mag = torch.bucketize(z.abs(), edges).to(torch.uint8)
+        codes = mag | ((z < 0).to(torch.uint8) << 3)
+        module.cr_qdata = ((codes[:, 1::2] << 4) | codes[:, ::2]).contiguous()
+
     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)
@@ -704,12 +750,13 @@ class ConvRotQuantizer(OstrisQuantizer):
 # ---------------- convrot8: W8A8 int8 backend ----------------
 
 
-def quantize_int8_rows(x: torch.Tensor):
+def quantize_int8_rows(x: torch.Tensor, qmax: int = 127):
-    """Symmetric per-row int8 quantization. Returns (int8 (rows, K), fp32 scales (rows,))."""
+    """Symmetric per-row integer quantization to [-qmax, qmax] (int8 storage).
+    Returns (int8 (rows, K), fp32 scales (rows,))."""
     xf = x.float()
-    scales = xf.abs().amax(dim=1) / 127.0
+    scales = xf.abs().amax(dim=1) / qmax
     scales = torch.where(scales > 0, scales, torch.ones_like(scales))
-    q = torch.round(xf / scales.unsqueeze(1)).clamp_(-127, 127).to(torch.int8)
+    q = torch.round(xf / scales.unsqueeze(1)).clamp_(-qmax, qmax).to(torch.int8)
     return q, scales
 
 
@@ -725,7 +772,9 @@ def _get_int8_kernels():
     from triton.language.extra import libdevice
 
     @triton.jit
-    def int8_act_quant_kernel(x_ptr, q_ptr, s_ptr, K, BLOCK_K: tl.constexpr):
+    def int8_act_quant_kernel(
+        x_ptr, q_ptr, s_ptr, K, QMAX: tl.constexpr, BLOCK_K: tl.constexpr
+    ):
         row = tl.program_id(0)
         base = row * K
         acc = tl.zeros((BLOCK_K,), tl.float32)
@@ -734,14 +783,14 @@ def _get_int8_kernels():
             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)
+        scale = tl.where(amax > 0, amax / QMAX, 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)
+            q = tl.minimum(tl.maximum(q, -1.0 * QMAX), 1.0 * QMAX)
             tl.store(q_ptr + base + offs, q.to(tl.int8), mask=mask)
         tl.store(s_ptr + row, scale)
 
@@ -776,7 +825,7 @@ def _get_int8_kernels():
 # nodes with known output shapes (see _nvfp4_act_quant_op). rows are padded to a
 # multiple of 32 inside the op for torch._int_mm; callers slice the mm output.
 @torch.library.custom_op("ostris::convrot_int8_act_quant", mutates_args=())
-def _int8_act_quant_op(x: torch.Tensor) -> list[torch.Tensor]:
+def _int8_act_quant_op(x: torch.Tensor, qmax: int) -> list[torch.Tensor]:
     rows, K = x.shape
     rows_pad = -(-rows // 32) * 32
     x = x.contiguous()
@@ -786,12 +835,16 @@ def _int8_act_quant_op(x: torch.Tensor) -> list[torch.Tensor]:
         q[rows:].zero_()
         scales[rows:].fill_(1.0)
     kernel, _ = _get_int8_kernels()
-    kernel[(rows,)](x, q, scales, K, BLOCK_K=2048 if K >= 2048 else K, num_warps=8)
+    # triton block shapes must be powers of 2; loads/stores are masked on offs < K
+    block_k = min(2048, 1 << (K - 1).bit_length())
+    kernel[(rows,)](
+        x, q, scales, K, QMAX=qmax, BLOCK_K=block_k, num_warps=8
+    )
     return [q, scales]
 
 
 @_int8_act_quant_op.register_fake
-def _int8_act_quant_fake(x):
+def _int8_act_quant_fake(x, qmax):
     rows, K = x.shape
     rows_pad = -(-rows // 32) * 32
     return [
@@ -832,22 +885,22 @@ def _int8_epilogue_fake(i32, a_scales, w_scales, bias, out_dtype):
     return torch.empty(m, n, device=i32.device, dtype=getattr(torch, out_dtype))
 
 
-def quantize_int8_rows_fused(x: torch.Tensor):
+def quantize_int8_rows_fused(x: torch.Tensor, qmax: int = 127):
     """Triton path of quantize_int8_rows: one extra read of x instead of the
     multi-kernel torch chain. Falls back to the torch ops."""
     if not (_triton_available() and x.is_cuda):
-        return quantize_int8_rows(x)
+        return quantize_int8_rows(x, qmax)
-    q, scales = _int8_act_quant_op(x)
+    q, scales = _int8_act_quant_op(x, qmax)
     rows = x.shape[0]
     return q[:rows], scales[:rows]
 
 
-def _int8_act_quant_padded(x: torch.Tensor):
+def _int8_act_quant_padded(x: torch.Tensor, qmax: int = 127):
     """Act quant with rows padded to a multiple of 32 for torch._int_mm."""
     if _triton_available() and x.is_cuda:
-        q, scales = _int8_act_quant_op(x)
+        q, scales = _int8_act_quant_op(x, qmax)
         return q, scales
-    q, scales = quantize_int8_rows(x)
+    q, scales = quantize_int8_rows(x, qmax)
     rows = q.shape[0]
     rows_pad = -(-rows // 32) * 32
     if rows_pad != rows:
@@ -867,10 +920,11 @@ def _int8_linear_ste_op(
     qdata: torch.Tensor,
     w_scales_u8: torch.Tensor,
     bias: Optional[torch.Tensor],
+    act_qmax: int,
     out_dtype: str,
 ) -> torch.Tensor:
     m = x2d.shape[0]
-    aq, a_s = _int8_act_quant_padded(x2d)
+    aq, a_s = _int8_act_quant_padded(x2d, act_qmax)
     i32 = torch._int_mm(aq, qdata.t())
     return _int8_epilogue(
         i32[:m], a_s[:m], w_scales_u8.view(torch.float32), bias,
@@ -879,14 +933,14 @@ def _int8_linear_ste_op(
 
 
 @_int8_linear_ste_op.register_fake
-def _int8_linear_ste_fake(x2d, qdata, w_scales_u8, bias, out_dtype):
+def _int8_linear_ste_fake(x2d, qdata, w_scales_u8, bias, act_qmax, out_dtype):
     return torch.empty(
         x2d.shape[0], qdata.shape[0], device=x2d.device, dtype=getattr(torch, out_dtype)
     )
 
 
 def _int8_linear_ste_setup(ctx, inputs, output):
-    x2d, qdata, w_scales_u8, bias, out_dtype = inputs
+    x2d, qdata, w_scales_u8, bias, act_qmax, out_dtype = inputs
     ctx.save_for_backward(qdata, w_scales_u8)
 
 
@@ -894,7 +948,7 @@ def _int8_linear_ste_backward(ctx, grad):
     qdata, w_scales_u8 = ctx.saved_tensors
     w_scales = w_scales_u8.view(torch.float32).to(grad.dtype)
     w = qdata.to(grad.dtype) * w_scales.unsqueeze(1)
-    return grad @ w, None, None, None, None
+    return grad @ w, None, None, None, None, None
 
 
 _int8_linear_ste_op.register_autograd(
@@ -960,6 +1014,10 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
     per-output-channel symmetric int8 with torch._int_mm. One instance per qtype,
     shareable across modules."""
 
+    # activation quantization range (per-token symmetric [-act_qmax, act_qmax]);
+    # the comfy w4a4 subclass narrows this to 7
+    act_qmax = 127
+
     def __init__(self, rot_size: int = 256):
         self.rot_size = rot_size
 
@@ -1006,9 +1064,17 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
         """Hardware STE linear for the training path. For int8 the saved qdata is
         the resident buffer itself, so autograd holds only a free reference."""
         return _int8_linear_ste_op(
-            x2d, self._qdata(module), self._scales_u8(module), module.bias, out_dtype
+            x2d, self._qdata(module), self._scales_u8(module), module.bias,
+            self.act_qmax, out_dtype,
         )
 
+    def fake_quant_rotated_weight(self, module, w_rot: torch.Tensor) -> torch.Tensor:
+        """dequant(quant(w_rot)) on the deployed int8 grid with the module's STORED
+        per-row scales — the value half of the QAT straight-through estimator.
+        Returns float32."""
+        s = self._scales(module).unsqueeze(1)
+        return torch.round(w_rot.float() / s).clamp_(-127, 127) * s
+
     def _dequantize_rotated(self, module, dtype: torch.dtype) -> torch.Tensor:
         w = self._qdata(module).float() * self._scales(module).unsqueeze(1)
         return w.to(dtype)
@@ -1024,6 +1090,15 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
         module.cr8_qdata = q
         module.cr8_scales = scales.view(torch.uint8)
 
+    @torch.no_grad()
+    def requantize_codes_(self, module, fp_weight: torch.Tensor) -> None:
+        """Re-quantize only the codes on the module's STORED scales — the grid a
+        QAT run trains against (see ConvRotQuantizer.requantize_codes_)."""
+        w = fp_weight.to(device=module.cr8_qdata.device, dtype=torch.float32)
+        w_rot = rotate(w, self._rot(module))
+        s = self._scales(module).unsqueeze(1)
+        module.cr8_qdata = torch.round(w_rot / s).clamp_(-127, 127).to(torch.int8)
+
     def forward(self, module, x: torch.Tensor) -> torch.Tensor:
         rot = self._rot(module)
         in_f, out_f = module.in_features, module.out_features
@@ -1047,7 +1122,7 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
             # 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())
+                aq, a_s = quantize_int8_rows(x2d.detach(), self.act_qmax)
                 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()
@@ -1057,7 +1132,9 @@ 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(rotate(x, rot).reshape(-1, in_f))
+            aq, a_s = _int8_act_quant_padded(
+                rotate(x, rot).reshape(-1, in_f), self.act_qmax
+            )
             i32 = torch._int_mm(aq, self._qdata(module).t())
             out = _int8_epilogue(i32[:m], a_s[:m], self._scales(module), module.bias, x.dtype)
             return out.reshape(*x.shape[:-1], out_f)
@@ -1413,12 +1490,13 @@ def _intn_linear_ste_op(
     w_scales_u8: torch.Tensor,
     bias: Optional[torch.Tensor],
     bits: int,
+    act_qmax: int,
     out_dtype: str,
 ) -> torch.Tensor:
     m, K = x2d.shape
     rows = w_scales_u8.view(torch.float32).numel()
     qdata = _unpack_intn_grouped_impl(packed, gratio, bits, rows, K)
-    aq, a_s = _int8_act_quant_padded(x2d)
+    aq, a_s = _int8_act_quant_padded(x2d, act_qmax)
     i32 = torch._int_mm(aq, qdata.t())
     return _int8_epilogue(
         i32[:m], a_s[:m], w_scales_u8.view(torch.float32), bias,
@@ -1427,7 +1505,7 @@ def _intn_linear_ste_op(
 
 
 @_intn_linear_ste_op.register_fake
-def _intn_linear_ste_fake(x2d, packed, gratio, w_scales_u8, bias, bits, out_dtype):
+def _intn_linear_ste_fake(x2d, packed, gratio, w_scales_u8, bias, bits, act_qmax, out_dtype):
     return torch.empty(
         x2d.shape[0], w_scales_u8.view(torch.float32).numel(),
         device=x2d.device, dtype=getattr(torch, out_dtype),
@@ -1435,7 +1513,7 @@ def _intn_linear_ste_fake(x2d, packed, gratio, w_scales_u8, bias, bits, out_dtyp
 
 
 def _intn_linear_ste_setup(ctx, inputs, output):
-    x2d, packed, gratio, w_scales_u8, bias, bits, out_dtype = inputs
+    x2d, packed, gratio, w_scales_u8, bias, bits, act_qmax, out_dtype = inputs
     ctx.bits = bits
     ctx.cols = x2d.shape[1]
     ctx.save_for_backward(packed, gratio, w_scales_u8)
@@ -1450,7 +1528,7 @@ def _intn_linear_ste_backward(ctx, grad):
         packed, gratio, ctx.bits, 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, None
+    return grad @ w, None, None, None, None, None, None, None
 
 
 _intn_linear_ste_op.register_autograd(
@@ -1495,12 +1573,13 @@ def _bitnet_linear_ste_op(
     gratio: torch.Tensor,
     w_scales_u8: torch.Tensor,
     bias: Optional[torch.Tensor],
+    act_qmax: int,
     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)
+    aq, a_s = _int8_act_quant_padded(x2d, act_qmax)
     i32 = torch._int_mm(aq, qdata.t())
     return _int8_epilogue(
         i32[:m], a_s[:m], w_scales_u8.view(torch.float32), bias,
@@ -1509,7 +1588,7 @@ def _bitnet_linear_ste_op(
 
 
 @_bitnet_linear_ste_op.register_fake
-def _bitnet_linear_ste_fake(x2d, packed, gratio, w_scales_u8, bias, out_dtype):
+def _bitnet_linear_ste_fake(x2d, packed, gratio, w_scales_u8, bias, act_qmax, out_dtype):
     return torch.empty(
         x2d.shape[0], w_scales_u8.view(torch.float32).numel(),
         device=x2d.device, dtype=getattr(torch, out_dtype),
@@ -1517,7 +1596,7 @@ def _bitnet_linear_ste_fake(x2d, packed, gratio, w_scales_u8, bias, out_dtype):
 
 
 def _bitnet_linear_ste_setup(ctx, inputs, output):
-    x2d, packed, gratio, w_scales_u8, bias, out_dtype = inputs
+    x2d, packed, gratio, w_scales_u8, bias, act_qmax, out_dtype = inputs
     ctx.cols = x2d.shape[1]
     ctx.save_for_backward(packed, gratio, w_scales_u8)
 
@@ -1529,7 +1608,7 @@ def _bitnet_linear_ste_backward(ctx, grad):
     # 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
+    return grad @ w, None, None, None, None, None, None
 
 
 _bitnet_linear_ste_op.register_autograd(
@@ -1608,7 +1687,7 @@ class ConvRotIntNQuantizer(ConvRotInt8Quantizer):
         # in the backward
         return _intn_linear_ste_op(
             x2d, module.crn_qdata, gratio.view(torch.float32),
-            module.crn_scales, module.bias, module.crn_bits, out_dtype,
+            module.crn_scales, module.bias, module.crn_bits, self.act_qmax, out_dtype,
         )
 
     def requantize_(self, module, fp_weight: torch.Tensor) -> None:
@@ -1621,6 +1700,46 @@ class ConvRotIntNQuantizer(ConvRotInt8Quantizer):
         if gratio is not None:
             module.crn_gratio = gratio.view(torch.uint8)
 
+    def _codes_on_stored_scales(self, module, w_rot: torch.Tensor) -> torch.Tensor:
+        """n-bit codes of a rotated weight on the module's stored scales (fp32)."""
+        qmax = (1 << (module.crn_bits - 1)) - 1
+        rs = self._scales(module).unsqueeze(1)
+        wf = w_rot.float()
+        gratio = getattr(module, "crn_gratio", None)
+        if gratio is None:
+            return torch.round(wf / rs).clamp_(-qmax, qmax)
+        gratio = gratio.view(torch.float32)
+        group = module.in_features // gratio.shape[1]
+        gs = gratio.repeat_interleave(group, dim=1) * rs
+        safe = torch.where(gs > 0, gs, torch.ones_like(gs))
+        return torch.round(wf / safe).clamp_(-qmax, qmax)
+
+    def fake_quant_rotated_weight(self, module, w_rot: torch.Tensor) -> torch.Tensor:
+        """Like the int8 version but on the n-bit grid: quantize with the stored
+        group scales, then re-express on the row int8 grid exactly like the
+        deployed unpack does. Returns float32."""
+        rs = self._scales(module).unsqueeze(1)
+        codes = self._codes_on_stored_scales(module, w_rot)
+        gratio = getattr(module, "crn_gratio", None)
+        if gratio is None:
+            return codes * rs
+        gratio = gratio.view(torch.float32)
+        group = module.in_features // gratio.shape[1]
+        ratio = gratio.repeat_interleave(group, dim=1)
+        code8 = torch.round(codes * ratio).clamp_(-127, 127)
+        return code8 * rs
+
+    def _pack(self, q: torch.Tensor) -> torch.Tensor:
+        return pack_intn_rows(q, self.bits)
+
+    @torch.no_grad()
+    def requantize_codes_(self, module, fp_weight: torch.Tensor) -> None:
+        """Re-quantize only the codes on the module's STORED scales — the grid a
+        QAT run trains against (see ConvRotQuantizer.requantize_codes_)."""
+        w = fp_weight.to(device=module.crn_qdata.device, dtype=torch.float32)
+        codes = self._codes_on_stored_scales(module, rotate(w, self._rot(module)))
+        module.crn_qdata = self._pack(codes.to(torch.int8))
+
 
 class ConvRotBitNetQuantizer(ConvRotIntNQuantizer):
     """BitNet-b1.58 style W1.58A8: convrotint2's exact ternary codes and
@@ -1636,6 +1755,9 @@ class ConvRotBitNetQuantizer(ConvRotIntNQuantizer):
         gratio, rscales = _intn_group_ratio_and_rscales(gscales, self.bits)
         return pack_ternary_rows(q), rscales, gratio.contiguous()
 
+    def _pack(self, q: torch.Tensor) -> torch.Tensor:
+        return pack_ternary_rows(q)
+
     def _qdata(self, module) -> torch.Tensor:
         return _bitnet_unpack_op(
             module.crn_qdata, module.crn_gratio.view(torch.float32),
@@ -1645,6 +1767,126 @@ class ConvRotBitNetQuantizer(ConvRotIntNQuantizer):
     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,
+            module.crn_scales, module.bias, self.act_qmax, out_dtype,
         )
 
+
+# ---------------- convrotcomfyw4a4: ComfyUI convrot_w4a4 compatible --------------
+#
+# Matches comfy_kitchen's TensorCoreConvRotW4A4Layout numerics exactly so a
+# trained model exports to a checkpoint ComfyUI loads and runs natively:
+#   - regular-Hadamard rotation, group size fixed at 256 (same R4 kronecker
+#     matrix as the other backends; comfy rejects other sizes)
+#   - weights: symmetric per-row int4, scale = row absmax / 7, RTN clamp [-7, 7]
+#   - activations: per-token int4 the same way (act_qmax = 7)
+# Storage here stays our bitpacked layout (bits=4); export_comfy_convrot_w4a4
+# repacks to comfy's nibble pairs and writes weight/weight_scale/comfy_quant
+# keys. Integer GEMM results are identical (int4 values through _int_mm
+# accumulate exactly like their int4 MMA).
+
+
+class ConvRotComfyW4A4Quantizer(ConvRotIntNQuantizer):
+    """ComfyUI-compatible ConvRot W4A4 (see block comment above)."""
+
+    act_qmax = 7
+    COMFY_GROUPSIZE = 256
+
+    def __init__(self, rot_size: int = 256):
+        super().__init__(bits=4, rot_size=self.COMFY_GROUPSIZE)
+
+    def _rot_for(self, d: int) -> int:
+        return self.COMFY_GROUPSIZE
+
+    def can_quantize(self, module: torch.nn.Linear) -> bool:
+        d = module.in_features
+        if d % self.COMFY_GROUPSIZE != 0 or module.out_features % 8 != 0:
+            if d not in _skip_warned:
+                _skip_warned.add(d)
+                print_acc(
+                    f"ConvRot(comfy): skipping linears with in_features={d} "
+                    f"(comfy convrot_w4a4 needs in divisible by {self.COMFY_GROUPSIZE}, "
+                    f"out by 8)"
+                )
+            return False
+        return True
+
+    def _quantize_rotated(self, w_rot: torch.Tensor, in_features: int):
+        # their format only fixes the DECODE (w = code * row_scale), not how the
+        # codes/scales are chosen — so use MSE-optimal per-row scales (sweep +
+        # LS refit with one group spanning the whole row) instead of their
+        # amax/7. group scales can't be expressed: folding 4-bit group grids
+        # onto the 4-bit row grid has no headroom (same inversion as the 8-bit
+        # fold), so per-row optimal is the ceiling of the format
+        q, gscales = _optimal_group_quantize(w_rot, self.bits, w_rot.shape[1])
+        return pack_intn_rows(q, self.bits), gscales.reshape(-1), None
+
+
+def export_comfy_convrot_w4a4(module, prefix: str) -> dict:
+    """State-dict entries for one convrotcomfyw4a4-quantized OstrisLinear in
+    ComfyUI's checkpoint format: packed int4 weight (low nibble = even column),
+    fp32 per-row weight_scale, and the comfy_quant JSON marker."""
+    import json
+
+    q = module.ostris_quantizer
+    if not isinstance(q, ConvRotComfyW4A4Quantizer):
+        raise ValueError(
+            f"export_comfy_convrot_w4a4 needs a convrotcomfyw4a4 module, got "
+            f"qtype '{getattr(q, 'qtype', None)}'"
+        )
+    codes = _unpack_intn_impl(
+        module.crn_qdata, module.crn_bits, module.out_features, module.in_features
+    )
+    lo = codes[:, 0::2].to(torch.int32) & 0x0F
+    hi = codes[:, 1::2].to(torch.int32) & 0x0F
+    packed = (lo | (hi << 4)).to(torch.int8)
+    conf = {
+        "format": "convrot_w4a4",
+        "convrot_groupsize": module.crn_rot_size,
+    }
+    entries = {
+        f"{prefix}weight": packed.contiguous(),
+        f"{prefix}weight_scale": q._scales(module).contiguous(),
+        f"{prefix}comfy_quant": torch.tensor(
+            list(json.dumps(conf).encode("utf-8")), dtype=torch.uint8
+        ),
+    }
+    if module.bias is not None:
+        entries[f"{prefix}bias"] = module.bias.detach()
+    return entries
+
+
+# ---------------- QAT: train the quantized weights themselves ----------------
+
+
+def convrot_qat_forward(module, x: torch.Tensor) -> torch.Tensor:
+    """Quantization-aware training forward for an OstrisLinear with a convrot
+    backend and a trainable full-precision `qat_master` parameter attached.
+
+    The weight enters the matmul through a straight-through fake-quant on the
+    DEPLOYED grid using the module's stored scales (refresh them by calling
+    module.requantize_(master) periodically), so the forward numerics track what
+    the saved quantized model will compute while gradients flow to the master.
+    Activations are fake-quanted like the no-hardware training path, also with a
+    straight-through estimate, so gradients pass through to the input when it
+    carries grad (end-to-end training; layer-local distillation feeds detached
+    inputs and the act STE is then value-identical)."""
+    q = module.ostris_quantizer
+    rot = q._rot(module)
+    in_f, out_f = module.in_features, module.out_features
+    w_rot = rotate(module.qat_master.to(x.dtype), rot)
+    with torch.no_grad():
+        w_fq = q.fake_quant_rotated_weight(module, w_rot.detach()).to(x.dtype)
+    w_ste = w_rot + (w_fq - w_rot).detach()
+    x2d = rotate(x, rot).reshape(-1, in_f)
+    with torch.no_grad():
+        if isinstance(q, ConvRotQuantizer):
+            pts1 = torch.ones((), device=x.device, dtype=torch.float32)
+            aq, a_scales, _ = quantize_nvfp4(x2d.detach(), pts=pts1)
+            x_dq = dequantize_nvfp4(aq, a_scales, pts1, x2d.shape[0], in_f, x.dtype)
+        else:
+            aq, a_scales = quantize_int8_rows(x2d.detach(), q.act_qmax)
+            x_dq = (aq.float() * a_scales.unsqueeze(1)).to(x.dtype)
+    x_ste = x2d + (x_dq - x2d).detach()
+    out = F.linear(x_ste, w_ste, module.bias)
+    return out.reshape(*x.shape[:-1], out_f)
+

toolkit/util/ostris_quant.py

@@ -29,7 +29,7 @@ class OstrisQuantizer:
     """
 
     # the qtype string this instance was resolved from (stamped by
-    # get_ostris_quantizer); pre-quantized saves need it to restore the backend
+    # get_ostris_quantizer); quantized saves need it to restore the backend
     qtype: Optional[str] = None
 
     def can_quantize(self, module: torch.nn.Linear) -> bool:
@@ -60,6 +60,9 @@ class OstrisQuantizer:
         return F.linear(x, w, module.bias)
 
 
+_wrong_device_warned = False
+
+
 class OstrisLinear(torch.nn.Linear):
     """A linear layer whose weight is quantized by an OstrisQuantizer backend.
 
@@ -87,6 +90,24 @@ class OstrisLinear(torch.nn.Linear):
         return w
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.is_cuda and not hasattr(self, "_layer_memory_manager"):
+            # a module left behind on the wrong device (usually cpu after a
+            # low_vram load) would run its dequant/matmul on cpu threads and
+            # silently hammer the cpu. move it to the input's device once.
+            # memory-managed modules are excluded: they keep cpu buffers by
+            # design and stage them to the gpu per forward.
+            buf = next((b for b in self._buffers.values() if b is not None), None)
+            if buf is not None and buf.device != x.device:
+                global _wrong_device_warned
+                if not _wrong_device_warned:
+                    _wrong_device_warned = True
+                    from toolkit.print import print_acc
+                    print_acc(
+                        f"OstrisLinear: quantized weights found on {buf.device} while the "
+                        f"input is on {x.device}; moving them to {x.device}. This usually "
+                        f"means something left the model behind after a low_vram load."
+                    )
+                self.to(x.device)
         return self.ostris_quantizer.forward(self, x)
 
     @torch.no_grad()
@@ -117,10 +138,132 @@ def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
     elif qtype in CONVROT_QTYPES:
         quantizer = get_convrot_quantizer(qtype)
     if quantizer is not None:
+        # quantized saves read this back to restore the backend on load
         quantizer.qtype = qtype
     return quantizer
 
 
+# metadata key shared with toolkit/models/classes/_mixin.py save_quantized
+QUANT_LAYERS_METADATA_KEY = "aitk_quantization"
+
+
+@torch.no_grad()
+def save_quantized_layers(
+    modules: Dict[str, "OstrisLinear"],
+    file_path: str,
+    metadata: Optional[Dict[str, str]] = None,
+    extra_tensors: Optional[Dict[str, torch.Tensor]] = None,
+) -> None:
+    """Save a set of quantized linears (keyed by their submodule path in the
+    target model) as a single safetensors file: backend buffers + bias per
+    module, with the qtype/attrs needed to restore them recorded in the file
+    metadata. Same layout as OstrisModelMixin.save_quantized, but partial —
+    apply with load_quantized_layers."""
+    import json
+    from safetensors.torch import save_file
+
+    quant_map = {}
+    state_dict = {}
+    for name, module in modules.items():
+        if module.ostris_quantizer.qtype is None:
+            raise ValueError(
+                f"Cannot save quantized module '{name}': its quantizer has no "
+                "qtype recorded (was it created through get_ostris_quantizer?)."
+            )
+        entry = {
+            "qtype": module.ostris_quantizer.qtype,
+            "dtype": str(module.ostris_orig_dtype).replace("torch.", ""),
+            "buffers": [],
+            "attrs": {},
+        }
+        if module.bias is not None:
+            state_dict[f"{name}.bias"] = module.bias
+        for buf_name, buf in module._buffers.items():
+            if buf is None:
+                continue
+            state_dict[f"{name}.{buf_name}"] = buf
+            entry["buffers"].append(buf_name)
+        for attr, value in vars(module).items():
+            if attr.startswith("_") or attr in (
+                "training", "in_features", "out_features",
+                "ostris_quantizer", "ostris_orig_dtype",
+            ):
+                continue
+            if isinstance(value, (bool, int, float, str)):
+                entry["attrs"][attr] = value
+        quant_map[name] = entry
+
+    if extra_tensors:
+        # training-state extras (e.g. QAT master weights); load_quantized_layers
+        # ignores them, so deployment loads are unaffected
+        state_dict.update(extra_tensors)
+    state_dict = {
+        k: v.detach().to("cpu", copy=True).contiguous() for k, v in state_dict.items()
+    }
+    meta = dict(metadata or {})
+    meta[QUANT_LAYERS_METADATA_KEY] = json.dumps(
+        {"modules": quant_map, "layers_only": True}
+    )
+    save_file(state_dict, file_path, metadata=meta)
+
+
+@torch.no_grad()
+def load_quantized_layers(root: torch.nn.Module, file_path: str) -> int:
+    """Apply a file written by save_quantized_layers onto a model: restores the
+    backend buffers (and bias) of each recorded module. Target modules may
+    already be OstrisLinear (model quantized on load — buffers are replaced) or
+    still plain nn.Linear (converted in place, no quantization math needed).
+    Returns the number of modules restored."""
+    import json
+    from safetensors import safe_open
+    from safetensors.torch import load_file
+
+    with safe_open(file_path, framework="pt", device="cpu") as f:
+        meta = f.metadata() or {}
+    if QUANT_LAYERS_METADATA_KEY not in meta:
+        raise ValueError(f"{file_path} has no quantized-layer metadata")
+    quant_map = json.loads(meta[QUANT_LAYERS_METADATA_KEY])["modules"]
+    state_dict = load_file(file_path)
+
+    for name, entry in quant_map.items():
+        module = root.get_submodule(name)
+        quantizer = get_ostris_quantizer(entry["qtype"])
+        if quantizer is None:
+            raise ValueError(f"Unknown qtype '{entry['qtype']}' in {file_path}")
+        # figure out the device the module currently lives on
+        ref = next(
+            (t for t in module._parameters.values() if t is not None),
+            next((t for t in module._buffers.values() if t is not None), None),
+        )
+        device = ref.device if ref is not None else torch.device("cpu")
+        if isinstance(module, OstrisLinear):
+            # replacing an existing quantized state (possibly another backend):
+            # its buffers are exclusively backend state, drop them all
+            module._buffers.clear()
+        else:
+            if "weight" in module._parameters:
+                del module._parameters["weight"]
+            module.ostris_orig_dtype = getattr(torch, entry["dtype"])
+            module.__class__ = OstrisLinear
+        for buf_name in entry["buffers"]:
+            module.register_buffer(
+                buf_name,
+                state_dict.pop(f"{name}.{buf_name}").to(device),
+                persistent=False,
+            )
+        for attr, value in entry.get("attrs", {}).items():
+            setattr(module, attr, value)
+        module.ostris_quantizer = quantizer
+        bias_key = f"{name}.bias"
+        if bias_key in state_dict and module.bias is not None:
+            module.bias.data.copy_(
+                state_dict.pop(bias_key).to(device, module.bias.dtype)
+            )
+        if module.bias is not None:
+            module.bias.requires_grad_(False)
+    return len(quant_map)
+
+
 @torch.no_grad()
 def convert_linear_to_ostris(
     module: torch.nn.Linear, quantizer: OstrisQuantizer