From b5d471f5ff6f7e07100a2d981d85cb8b0857e953 Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Fri, 10 Jul 2026 10:25:06 +0800
Subject: [PATCH 01/13] feat: add torchao INT4 weight-only quantization backend

---
 comfy/quantization/torchao/__init__.py |  24 ++++
 comfy/quantization/torchao/linear.py   | 154 +++++++++++++++++++++++++
 comfy/quantization/torchao/quantize.py | 111 ++++++++++++++++++
 comfy/quantization/torchao/quarot.py   |  80 +++++++++++++
 4 files changed, 369 insertions(+)
 create mode 100644 comfy/quantization/torchao/__init__.py
 create mode 100644 comfy/quantization/torchao/linear.py
 create mode 100644 comfy/quantization/torchao/quantize.py
 create mode 100644 comfy/quantization/torchao/quarot.py

diff --git a/comfy/quantization/torchao/__init__.py b/comfy/quantization/torchao/__init__.py
new file mode 100644
index 000000000..afb8d4ea5
--- /dev/null
+++ b/comfy/quantization/torchao/__init__.py
@@ -0,0 +1,24 @@
+"""
+ComfyUI TorchAO INT4 weight-only quantization backend.
+
+W4A16 scheme: weights quantized to INT4 via torchao,
+activations remain FP16. Supports CUDA, Intel XPU, and CPU.
+
+Includes:
+- TINT4Linear: INT4 linear layer with LoRA forward injection + QuaRot
+- quantize_model: torchao INT4 quantization entry point
+- reconstruct_int4_state_dict: rebuild TINT4Linear from safetensors
+- build_hadamard / rotate_weight: QuaRot Hadamard rotation
+"""
+
+from .linear import TINT4Linear
+from .quantize import quantize_model, reconstruct_int4_state_dict
+from .quarot import build_hadamard, rotate_weight
+
+__all__ = [
+    "TINT4Linear",
+    "quantize_model",
+    "reconstruct_int4_state_dict",
+    "build_hadamard",
+    "rotate_weight",
+]
diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py
new file mode 100644
index 000000000..de6b3b3d2
--- /dev/null
+++ b/comfy/quantization/torchao/linear.py
@@ -0,0 +1,154 @@
+"""
+TINT4Linear — INT4 weight-only Linear with LoRA + QuaRot.
+
+Backed by torchao Int4PlainInt32Tensor. Activations stay FP16.
+LoRA entries are injected at forward time — no kernel modification needed.
+QuaRot activation rotation is applied online when enabled.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchao.quantization.quantize_.workflows.int4.int4_plain_int32_tensor import (
+    Int4PlainInt32Tensor,
+)
+
+
+class TINT4Linear(nn.Module):
+    """
+    INT4 weight-only linear layer.
+
+    Stores raw int32 qdata + fp16 scales + int8 zero-points.
+    Int4PlainInt32Tensor is lazily constructed and rebuilt on device change.
+
+    Slots (set externally by loaders):
+      _use_quarot: bool         — enable activation Hadamard rotation
+      _group_size: int          — QuaRot group size
+      _hadamard_H: Tensor|None  — normalized Hadamard matrix
+      _tint4_lora_entries: dict|None — {lora_name: [(type, ...), ...]}
+    """
+
+    def __init__(self, in_features: int, out_features: int,
+                 qdata: torch.Tensor, scale: torch.Tensor,
+                 zp: torch.Tensor, block_size: tuple,
+                 bias: torch.Tensor | None = None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+
+        if bias is not None:
+            self.bias = nn.Parameter(bias)
+        else:
+            self.register_parameter('bias', None)
+
+        self._qdata = qdata           # int32  (out_f, in_f // 8)
+        self._scale = scale           # fp16   (n_blocks, out_f)
+        self._zp = zp                 # int8   (n_blocks, out_f)
+        self._block_size = block_size # (b0, b1)
+        self._qt: Int4PlainInt32Tensor | None = None
+
+        # QuaRot (optional)
+        self._use_quarot: bool = False
+        self._group_size: int = 128
+        self._hadamard_H: torch.Tensor | None = None
+
+        # LoRA entries slot
+        self._tint4_lora_entries: dict | None = None
+
+    @property
+    def weight(self) -> Int4PlainInt32Tensor:
+        if self._qt is None:
+            self._qt = Int4PlainInt32Tensor(
+                self._qdata, self._scale, self._zp,
+                self._block_size, [self.out_features, self.in_features],
+            )
+        return self._qt
+
+    @weight.setter
+    def weight(self, value: Int4PlainInt32Tensor) -> None:
+        if isinstance(value, Int4PlainInt32Tensor):
+            self._qt = value
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_flat = x.reshape(-1, x.shape[-1])
+
+        # ── QuaRot activation rotation (online) ──────────────────
+        if self._use_quarot and self._hadamard_H is not None:
+            try:
+                H_dev = self._hadamard_H.to(
+                    device=x.device,
+                    dtype=(x.dtype if x.dtype in (
+                        torch.float16, torch.bfloat16)
+                           else torch.float16),
+                )
+                n_groups = x_flat.shape[-1] // self._group_size
+                x_g = x_flat.view(-1, n_groups, self._group_size)
+                x_flat = torch.matmul(x_g, H_dev).view(-1, x_g.shape[2] * n_groups)
+            except Exception:
+                pass
+
+        dev = x.device
+
+        # ── Rebuild Int4PlainInt32Tensor on device change ────────
+        if self._qt is None or self._qt.device != dev:
+            self._qt = Int4PlainInt32Tensor(
+                self._qdata.to(dev), self._scale.to(dev), self._zp.to(dev),
+                self._block_size, [self.out_features, self.in_features],
+            )
+
+        out = F.linear(x_flat, self._qt, None)
+
+        # ── LoRA forward injection ───────────────────────────────
+        entries = self._tint4_lora_entries
+        if entries is not None and entries:
+            cd = (x.dtype if x.dtype in (torch.float16, torch.bfloat16)
+                  else torch.float16)
+            for lora_entries in entries.values():
+                for e in lora_entries:
+                    if isinstance(e[0], str) and e[0] == "lokr":
+                        _, w1, w2, mult, factor = e[:5]
+                        sl = e[5] if len(e) > 5 else None
+                        se = e[6] if len(e) > 6 else None
+                        w1d = w1.to(device=dev, dtype=cd)
+                        w2d = w2.to(device=dev, dtype=cd)
+                        of2, if2 = w2d.shape
+                        w1x = w1d.repeat_interleave(
+                            of2 // factor, dim=0).repeat_interleave(
+                            if2 // factor, dim=1)
+                        dw = (w1x * w2d).mul_(mult)
+                        lo = x_flat.to(cd) @ dw.T
+                        if sl is not None:
+                            if lo.shape[1] == (se - sl):
+                                out[:, sl:se] += lo
+                        elif lo.shape == out.shape:
+                            out += lo
+                        continue
+
+                    # standard LoRA
+                    A, B, mult = e[:3]
+                    sl = e[3] if len(e) > 3 else None
+                    se = e[4] if len(e) > 4 else None
+                    Ad = A.to(device=dev, dtype=cd)
+                    Bd = B.to(device=dev, dtype=cd)
+                    lo = (x_flat.to(cd) @ Ad.T) @ Bd.T * mult
+                    if sl is not None:
+                        if lo.shape[1] == (se - sl):
+                            out[:, sl:se] += lo
+                    elif lo.shape[1] == out.shape[1]:
+                        out += lo
+
+        if self.bias is not None:
+            out += self.bias.to(device=dev, dtype=out.dtype)
+
+        return out.reshape(*x.shape[:-1], out.shape[-1])
+
+    def release(self) -> None:
+        """Drop cached Int4PlainInt32Tensor to free VRAM."""
+        self._qt = None
+
+    def __del__(self):
+        self._qdata = None
+        self._scale = None
+        self._zp = None
+        self._qt = None
+        self._tint4_lora_entries = None
diff --git a/comfy/quantization/torchao/quantize.py b/comfy/quantization/torchao/quantize.py
new file mode 100644
index 000000000..b5d9ae30b
--- /dev/null
+++ b/comfy/quantization/torchao/quantize.py
@@ -0,0 +1,111 @@
+"""
+TorchAO INT4 quantization helpers for ComfyUI.
+
+- quantize_model: in-place INT4 quantization via torchao
+- reconstruct_int4_state_dict: rebuild TINT4Linear layers from safetensors
+"""
+
+import torch
+import torch.nn as nn
+from torchao.quantization import Int4WeightOnlyConfig, quantize_
+
+from .linear import TINT4Linear
+
+
+def quantize_model(
+    model: nn.Module,
+    group_size: int = 128,
+    filter_fn=None,
+) -> nn.Module:
+    """
+    Quantize all nn.Linear layers in-place via torchao INT4 weight-only.
+
+    Args:
+        model:      PyTorch model to quantize.
+        group_size: Quantization group size (32/64/128/256).
+        filter_fn:  Optional callable(module, name) → bool.
+                    Return True to skip that module.
+
+    Returns:
+        Same model (quantized in-place).
+    """
+    config = Int4WeightOnlyConfig(
+        group_size=group_size,
+        int4_packing_format="plain_int32",
+    )
+    quantize_(model, config, filter_fn=filter_fn)
+    return model
+
+
+# torchao plain_int32 safetensors suffixes
+_QUANT_SUFFIXES = (
+    ".weight_scale", ".weight_zp",
+    ".weight_b0", ".weight_b1",
+    ".weight_sh0", ".weight_sh1",
+    ".comfy_quant",
+)
+
+
+def _is_quantized_weight(key: str, sd: dict) -> bool:
+    """Check whether a safetensors key is a torchao-quantized weight."""
+    if not key.endswith(".weight"):
+        return False
+    base = key[:-len(".weight")]
+    return f"{base}.weight_scale" in sd and f"{base}.weight_zp" in sd
+
+
+def reconstruct_int4_state_dict(
+    sd: dict,
+    device: torch.device = torch.device("cpu"),
+) -> dict[str, TINT4Linear]:
+    """
+    Extract TINT4Linear layers from a torchao-quantized safetensors dict.
+
+    Pops quantized weight tensors from sd and replaces them with
+    fully-constructed TINT4Linear modules keyed by base layer name.
+
+    Caller is responsible for placing the modules into the target model.
+
+    Args:
+        sd:     Safetensors state dict with torchao int4 packed weights.
+        device: Target device for the reconstructed layers.
+
+    Returns:
+        Dict mapping base key (e.g. "blocks.0.attn.qkv") → TINT4Linear.
+    """
+    replaced: dict[str, TINT4Linear] = {}
+
+    for key in list(sd.keys()):
+        if not _is_quantized_weight(key, sd):
+            continue
+
+        base = key[:-len(".weight")]
+        qdata = sd.pop(key)
+        scale = sd.pop(f"{base}.weight_scale")
+        zp = sd.pop(f"{base}.weight_zp")
+        b0 = sd.pop(f"{base}.weight_b0")
+        b1 = sd.pop(f"{base}.weight_b1")
+
+        for suffix in _QUANT_SUFFIXES:
+            sd.pop(f"{base}{suffix}", None)
+
+        in_f = qdata.shape[1] * 8   # int32 packs 8 int4 values
+        out_f = qdata.shape[0]
+
+        layer = TINT4Linear(
+            in_features=in_f,
+            out_features=out_f,
+            qdata=qdata.to(torch.int32),
+            scale=scale.to(torch.float16),
+            zp=zp.to(torch.int8),
+            block_size=(b0.item(), b1.item()),
+        )
+
+        if device.type != "cpu":
+            layer._qdata = layer._qdata.to(device)
+            layer._scale = layer._scale.to(device)
+            layer._zp = layer._zp.to(device)
+
+        replaced[base] = layer
+
+    return replaced
diff --git a/comfy/quantization/torchao/quarot.py b/comfy/quantization/torchao/quarot.py
new file mode 100644
index 000000000..798dd89b3
--- /dev/null
+++ b/comfy/quantization/torchao/quarot.py
@@ -0,0 +1,80 @@
+"""
+QuaRot: Group-wise Hadamard rotation for INT quantization quality.
+
+Spreads activation outliers across channels via orthogonal Hadamard matrices.
+Based on QuaRot (2024). Pure PyTorch — scipy optional for large matrices.
+"""
+
+import torch
+
+try:
+    from scipy.linalg import hadamard as _scipy_hadamard
+    _SCIPY_AVAILABLE = True
+except ImportError:
+    _SCIPY_AVAILABLE = False
+
+_HADAMARD_CACHE: dict = {}
+
+
+def build_hadamard(
+    size: int,
+    device: str = "cpu",
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """
+    Build normalized orthogonal Hadamard matrix (size must be power of 2).
+    Returns H such that H @ H^T = I.
+    Uses scipy if available, otherwise pure-PyTorch Sylvester construction.
+    """
+    if size & (size - 1) != 0:
+        raise ValueError(f"Hadamard size must be a power of 2, got {size}")
+
+    cache_key = (size, str(device), dtype)
+    if cache_key in _HADAMARD_CACHE:
+        return _HADAMARD_CACHE[cache_key]
+
+    if _SCIPY_AVAILABLE:
+        import numpy as np
+        H_np = _scipy_hadamard(size)
+        H = torch.tensor(H_np, dtype=dtype, device=device) / (size ** 0.5)
+    else:
+        H = torch.tensor([[1.0]], dtype=dtype, device=device)
+        n = 1
+        while n < size:
+            H = torch.cat([torch.cat([H,  H], dim=1),
+                           torch.cat([H, -H], dim=0)], dim=0)
+            n *= 2
+        H = H / (size ** 0.5)
+
+    _HADAMARD_CACHE[cache_key] = H
+    return H
+
+
+def rotate_weight(
+    weight: torch.Tensor,
+    H: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    """
+    Rotate weight matrix offline: W_rot = W @ H_block^T.
+
+    For Linear(in, out) with weight (out_f, in_f):
+    each row is split into groups of group_size and rotated by H^T.
+
+    Args:
+        weight:     (out_features, in_features)
+        H:          Normalized Hadamard matrix (group_size, group_size)
+        group_size: Group size for block-diagonal rotation
+
+    Returns:
+        Rotated weight, same shape.
+    """
+    out_f, in_f = weight.shape
+    if in_f % group_size != 0:
+        raise ValueError(
+            f"in_features {in_f} not divisible by group_size {group_size}")
+    n_groups = in_f // group_size
+    W_grouped = weight.view(out_f, n_groups, group_size)
+    H_t = H.T.to(dtype=weight.dtype, device=weight.device)
+    W_rot = torch.matmul(W_grouped, H_t)
+    return W_rot.reshape(out_f, in_f)

From 3e5738da2463e23e83026b14ca20211668151a90 Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Fri, 10 Jul 2026 11:07:08 +0800
Subject: [PATCH 02/13] =?UTF-8?q?fix:=20address=20CodeRabbit=20feedback=20?=
 =?UTF-8?q?=E2=80=94=20QuaRot=20error=20surface,=20LoRA=20shape=20validati?=
 =?UTF-8?q?on,=20Sylvester=20bug=20fix,=20cache=20cap?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 comfy/quantization/torchao/linear.py | 50 +++++++++++++++++++---------
 comfy/quantization/torchao/quarot.py | 11 ++++--
 2 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py
index de6b3b3d2..32a1c7c61 100644
--- a/comfy/quantization/torchao/linear.py
+++ b/comfy/quantization/torchao/linear.py
@@ -74,18 +74,15 @@ class TINT4Linear(nn.Module):
 
         # ── QuaRot activation rotation (online) ──────────────────
         if self._use_quarot and self._hadamard_H is not None:
-            try:
-                H_dev = self._hadamard_H.to(
-                    device=x.device,
-                    dtype=(x.dtype if x.dtype in (
-                        torch.float16, torch.bfloat16)
-                           else torch.float16),
+            if x_flat.shape[-1] % self._group_size != 0:
+                raise ValueError(
+                    f"activation dim {x_flat.shape[-1]} not divisible "
+                    f"by QuaRot group_size {self._group_size}"
                 )
-                n_groups = x_flat.shape[-1] // self._group_size
-                x_g = x_flat.view(-1, n_groups, self._group_size)
-                x_flat = torch.matmul(x_g, H_dev).view(-1, x_g.shape[2] * n_groups)
-            except Exception:
-                pass
+            H_dev = self._hadamard_H.to(device=x.device, dtype=x.dtype)
+            n_groups = x_flat.shape[-1] // self._group_size
+            x_g = x_flat.view(-1, n_groups, self._group_size)
+            x_flat = torch.matmul(x_g, H_dev).view(-1, n_groups * self._group_size)
 
         dev = x.device
 
@@ -105,6 +102,7 @@ class TINT4Linear(nn.Module):
                   else torch.float16)
             for lora_entries in entries.values():
                 for e in lora_entries:
+                    # ── LoKr path ────────────────────────────────
                     if isinstance(e[0], str) and e[0] == "lokr":
                         _, w1, w2, mult, factor = e[:5]
                         sl = e[5] if len(e) > 5 else None
@@ -118,13 +116,23 @@ class TINT4Linear(nn.Module):
                         dw = (w1x * w2d).mul_(mult)
                         lo = x_flat.to(cd) @ dw.T
                         if sl is not None:
-                            if lo.shape[1] == (se - sl):
-                                out[:, sl:se] += lo
+                            if lo.shape[1] != (se - sl):
+                                raise ValueError(
+                                    f"LoKr output width {lo.shape[1]} "
+                                    f"doesn't match target slice width "
+                                    f"{se - sl}"
+                                )
+                            out[:, sl:se] += lo
                         elif lo.shape == out.shape:
                             out += lo
+                        else:
+                            raise ValueError(
+                                f"LoKr output shape {lo.shape} "
+                                f"incompatible with {out.shape}"
+                            )
                         continue
 
-                    # standard LoRA
+                    # ── Standard LoRA path ────────────────────────
                     A, B, mult = e[:3]
                     sl = e[3] if len(e) > 3 else None
                     se = e[4] if len(e) > 4 else None
@@ -132,10 +140,20 @@ class TINT4Linear(nn.Module):
                     Bd = B.to(device=dev, dtype=cd)
                     lo = (x_flat.to(cd) @ Ad.T) @ Bd.T * mult
                     if sl is not None:
-                        if lo.shape[1] == (se - sl):
-                            out[:, sl:se] += lo
+                        if lo.shape[1] != (se - sl):
+                            raise ValueError(
+                                f"LoRA output width {lo.shape[1]} "
+                                f"doesn't match target slice width "
+                                f"{se - sl}"
+                            )
+                        out[:, sl:se] += lo
                     elif lo.shape[1] == out.shape[1]:
                         out += lo
+                    else:
+                        raise ValueError(
+                            f"LoRA output shape {lo.shape} "
+                            f"incompatible with {out.shape}"
+                        )
 
         if self.bias is not None:
             out += self.bias.to(device=dev, dtype=out.dtype)
diff --git a/comfy/quantization/torchao/quarot.py b/comfy/quantization/torchao/quarot.py
index 798dd89b3..3fa03d9fb 100644
--- a/comfy/quantization/torchao/quarot.py
+++ b/comfy/quantization/torchao/quarot.py
@@ -14,6 +14,7 @@ except ImportError:
     _SCIPY_AVAILABLE = False
 
 _HADAMARD_CACHE: dict = {}
+_HADAMARD_MAX_SIZE = 64  # Maximum number of cached (size, device, dtype) entries
 
 
 def build_hadamard(
@@ -38,14 +39,20 @@ def build_hadamard(
         H_np = _scipy_hadamard(size)
         H = torch.tensor(H_np, dtype=dtype, device=device) / (size ** 0.5)
     else:
+        # Pure-PyTorch Sylvester construction
         H = torch.tensor([[1.0]], dtype=dtype, device=device)
         n = 1
         while n < size:
-            H = torch.cat([torch.cat([H,  H], dim=1),
-                           torch.cat([H, -H], dim=0)], dim=0)
+            top = torch.cat([H, H], dim=1)
+            bottom = torch.cat([H, -H], dim=1)
+            H = torch.cat([top, bottom], dim=0)
             n *= 2
         H = H / (size ** 0.5)
 
+    # Enforce cache size cap (evict oldest entry)
+    if len(_HADAMARD_CACHE) >= _HADAMARD_MAX_SIZE:
+        _HADAMARD_CACHE.pop(next(iter(_HADAMARD_CACHE)))
+
     _HADAMARD_CACHE[cache_key] = H
     return H
 

From 8491b9103e7972dddeb787882be90cac0754fe91 Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Fri, 10 Jul 2026 11:25:06 +0800
Subject: [PATCH 03/13] =?UTF-8?q?fix:=20add=20=5Flora=5Fprerotated=20flag?=
 =?UTF-8?q?=20=E2=80=94=20avoid=20double-rotation=20with=20pre-rotated=20L?=
 =?UTF-8?q?oRA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 comfy/quantization/torchao/linear.py | 36 +++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py
index 32a1c7c61..83f5f08e1 100644
--- a/comfy/quantization/torchao/linear.py
+++ b/comfy/quantization/torchao/linear.py
@@ -3,7 +3,9 @@ TINT4Linear — INT4 weight-only Linear with LoRA + QuaRot.
 
 Backed by torchao Int4PlainInt32Tensor. Activations stay FP16.
 LoRA entries are injected at forward time — no kernel modification needed.
-QuaRot activation rotation is applied online when enabled.
+QuaRot activation rotation is applied online; LoRA A/w2 matrices are
+rotated into the QuaRot basis inside forward() automatically unless
+_lora_prerotated is set (for callers that pre-rotate at injection time).
 """
 
 import torch
@@ -26,6 +28,9 @@ class TINT4Linear(nn.Module):
       _group_size: int          — QuaRot group size
       _hadamard_H: Tensor|None  — normalized Hadamard matrix
       _tint4_lora_entries: dict|None — {lora_name: [(type, ...), ...]}
+      _lora_prerotated: bool    — True if caller already rotated A/w2
+                                  into QuaRot basis (default False).
+                                  When False, rotation is done in forward().
     """
 
     def __init__(self, in_features: int, out_features: int,
@@ -55,6 +60,9 @@ class TINT4Linear(nn.Module):
         # LoRA entries slot
         self._tint4_lora_entries: dict | None = None
 
+        # Pre-rotation flag: set True if caller already rotated A/w2
+        self._lora_prerotated: bool = False
+
     @property
     def weight(self) -> Int4PlainInt32Tensor:
         if self._qt is None:
@@ -69,6 +77,24 @@ class TINT4Linear(nn.Module):
         if isinstance(value, Int4PlainInt32Tensor):
             self._qt = value
 
+    # ═══════════════════════════════════════════════════════════════
+    # Helper: rotate a LoRA A / w2 matrix into the QuaRot basis
+    # ═══════════════════════════════════════════════════════════════
+
+    def _rotate_into_quarot(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Rotate tensor (A or w2) to match the Hadamard-rotated activation.
+        Skipped when _lora_prerotated is True (caller already rotated)."""
+        if self._lora_prerotated:
+            return tensor
+        if not self._use_quarot or self._hadamard_H is None:
+            return tensor
+        if tensor.shape[1] % self._group_size != 0:
+            return tensor
+        Ht = self._hadamard_H.T.to(device=tensor.device, dtype=tensor.dtype)
+        ng = tensor.shape[1] // self._group_size
+        return (tensor.reshape(tensor.shape[0], ng, self._group_size)
+                @ Ht).reshape(tensor.shape[0], -1)
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x_flat = x.reshape(-1, x.shape[-1])
 
@@ -107,6 +133,10 @@ class TINT4Linear(nn.Module):
                         _, w1, w2, mult, factor = e[:5]
                         sl = e[5] if len(e) > 5 else None
                         se = e[6] if len(e) > 6 else None
+
+                        # Rotate w2 into QuaRot basis (honors _lora_prerotated)
+                        w2 = self._rotate_into_quarot(w2)
+
                         w1d = w1.to(device=dev, dtype=cd)
                         w2d = w2.to(device=dev, dtype=cd)
                         of2, if2 = w2d.shape
@@ -136,6 +166,10 @@ class TINT4Linear(nn.Module):
                     A, B, mult = e[:3]
                     sl = e[3] if len(e) > 3 else None
                     se = e[4] if len(e) > 4 else None
+
+                    # Rotate A into QuaRot basis (honors _lora_prerotated)
+                    A = self._rotate_into_quarot(A)
+
                     Ad = A.to(device=dev, dtype=cd)
                     Bd = B.to(device=dev, dtype=cd)
                     lo = (x_flat.to(cd) @ Ad.T) @ Bd.T * mult

From 2d0bb38f4fbceb882d9c35ad044c71dde618a808 Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Fri, 10 Jul 2026 11:45:26 +0800
Subject: [PATCH 04/13] fix: add docstrings for coverage threshold

---
 comfy/quantization/torchao/linear.py   | 67 +++++++++++++++++---------
 comfy/quantization/torchao/quantize.py | 12 ++---
 2 files changed, 51 insertions(+), 28 deletions(-)

diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py
index 83f5f08e1..189a5f112 100644
--- a/comfy/quantization/torchao/linear.py
+++ b/comfy/quantization/torchao/linear.py
@@ -37,6 +37,18 @@ class TINT4Linear(nn.Module):
                  qdata: torch.Tensor, scale: torch.Tensor,
                  zp: torch.Tensor, block_size: tuple,
                  bias: torch.Tensor | None = None):
+        """
+        Initialize TINT4Linear from raw torchao INT4 tensors.
+
+        Args:
+            in_features: Input feature dimension.
+            out_features: Output feature dimension.
+            qdata: Packed int32 quantized weight data (out_f, in_f // 8).
+            scale: Per-block fp16 scales (n_blocks, out_f).
+            zp: Per-block int8 zero-points (n_blocks, out_f).
+            block_size: (b0, b1) tuple defining quantization block shape.
+            bias: Optional bias tensor (out_f,).
+        """
         super().__init__()
         self.in_features = in_features
         self.out_features = out_features
@@ -46,10 +58,10 @@ class TINT4Linear(nn.Module):
         else:
             self.register_parameter('bias', None)
 
-        self._qdata = qdata           # int32  (out_f, in_f // 8)
-        self._scale = scale           # fp16   (n_blocks, out_f)
-        self._zp = zp                 # int8   (n_blocks, out_f)
-        self._block_size = block_size # (b0, b1)
+        self._qdata = qdata
+        self._scale = scale
+        self._zp = zp
+        self._block_size = block_size
         self._qt: Int4PlainInt32Tensor | None = None
 
         # QuaRot (optional)
@@ -60,11 +72,12 @@ class TINT4Linear(nn.Module):
         # LoRA entries slot
         self._tint4_lora_entries: dict | None = None
 
-        # Pre-rotation flag: set True if caller already rotated A/w2
+        # Pre-rotation flag
         self._lora_prerotated: bool = False
 
     @property
     def weight(self) -> Int4PlainInt32Tensor:
+        """Lazily construct and return the Int4PlainInt32Tensor view."""
         if self._qt is None:
             self._qt = Int4PlainInt32Tensor(
                 self._qdata, self._scale, self._zp,
@@ -74,16 +87,18 @@ class TINT4Linear(nn.Module):
 
     @weight.setter
     def weight(self, value: Int4PlainInt32Tensor) -> None:
+        """Replace the underlying Int4PlainInt32Tensor."""
         if isinstance(value, Int4PlainInt32Tensor):
             self._qt = value
 
-    # ═══════════════════════════════════════════════════════════════
-    # Helper: rotate a LoRA A / w2 matrix into the QuaRot basis
-    # ═══════════════════════════════════════════════════════════════
-
     def _rotate_into_quarot(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Rotate tensor (A or w2) to match the Hadamard-rotated activation.
-        Skipped when _lora_prerotated is True (caller already rotated)."""
+        """
+        Rotate a LoRA A or w2 matrix into the QuaRot Hadamard basis.
+
+        Skipped when _lora_prerotated is True (caller already rotated).
+        Returns unchanged tensor if QuaRot is disabled or dimensions
+        are not divisible by _group_size.
+        """
         if self._lora_prerotated:
             return tensor
         if not self._use_quarot or self._hadamard_H is None:
@@ -96,9 +111,22 @@ class TINT4Linear(nn.Module):
                 @ Ht).reshape(tensor.shape[0], -1)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass with optional QuaRot activation rotation and LoRA injection.
+
+        Args:
+            x: Input tensor, shape (..., in_features).
+
+        Returns:
+            Output tensor, shape (..., out_features).
+
+        Raises:
+            ValueError: If activation dim is not divisible by QuaRot group_size,
+                        or if LoRA output shapes are incompatible.
+        """
         x_flat = x.reshape(-1, x.shape[-1])
 
-        # ── QuaRot activation rotation (online) ──────────────────
+        # QuaRot activation rotation (online)
         if self._use_quarot and self._hadamard_H is not None:
             if x_flat.shape[-1] % self._group_size != 0:
                 raise ValueError(
@@ -112,7 +140,7 @@ class TINT4Linear(nn.Module):
 
         dev = x.device
 
-        # ── Rebuild Int4PlainInt32Tensor on device change ────────
+        # Rebuild Int4PlainInt32Tensor on device change
         if self._qt is None or self._qt.device != dev:
             self._qt = Int4PlainInt32Tensor(
                 self._qdata.to(dev), self._scale.to(dev), self._zp.to(dev),
@@ -121,22 +149,19 @@ class TINT4Linear(nn.Module):
 
         out = F.linear(x_flat, self._qt, None)
 
-        # ── LoRA forward injection ───────────────────────────────
+        # LoRA forward injection
         entries = self._tint4_lora_entries
         if entries is not None and entries:
             cd = (x.dtype if x.dtype in (torch.float16, torch.bfloat16)
                   else torch.float16)
             for lora_entries in entries.values():
                 for e in lora_entries:
-                    # ── LoKr path ────────────────────────────────
+                    # LoKr path
                     if isinstance(e[0], str) and e[0] == "lokr":
                         _, w1, w2, mult, factor = e[:5]
                         sl = e[5] if len(e) > 5 else None
                         se = e[6] if len(e) > 6 else None
-
-                        # Rotate w2 into QuaRot basis (honors _lora_prerotated)
                         w2 = self._rotate_into_quarot(w2)
-
                         w1d = w1.to(device=dev, dtype=cd)
                         w2d = w2.to(device=dev, dtype=cd)
                         of2, if2 = w2d.shape
@@ -162,14 +187,11 @@ class TINT4Linear(nn.Module):
                             )
                         continue
 
-                    # ── Standard LoRA path ────────────────────────
+                    # Standard LoRA path
                     A, B, mult = e[:3]
                     sl = e[3] if len(e) > 3 else None
                     se = e[4] if len(e) > 4 else None
-
-                    # Rotate A into QuaRot basis (honors _lora_prerotated)
                     A = self._rotate_into_quarot(A)
-
                     Ad = A.to(device=dev, dtype=cd)
                     Bd = B.to(device=dev, dtype=cd)
                     lo = (x_flat.to(cd) @ Ad.T) @ Bd.T * mult
@@ -199,6 +221,7 @@ class TINT4Linear(nn.Module):
         self._qt = None
 
     def __del__(self):
+        """Cleanup raw tensor references on deletion."""
         self._qdata = None
         self._scale = None
         self._zp = None
diff --git a/comfy/quantization/torchao/quantize.py b/comfy/quantization/torchao/quantize.py
index b5d9ae30b..e2553cdea 100644
--- a/comfy/quantization/torchao/quantize.py
+++ b/comfy/quantization/torchao/quantize.py
@@ -21,10 +21,10 @@ def quantize_model(
     Quantize all nn.Linear layers in-place via torchao INT4 weight-only.
 
     Args:
-        model:      PyTorch model to quantize.
+        model: PyTorch model to quantize.
         group_size: Quantization group size (32/64/128/256).
-        filter_fn:  Optional callable(module, name) → bool.
-                    Return True to skip that module.
+        filter_fn: Optional callable(module, name) → bool.
+                   Return True to skip that module.
 
     Returns:
         Same model (quantized in-place).
@@ -47,7 +47,7 @@ _QUANT_SUFFIXES = (
 
 
 def _is_quantized_weight(key: str, sd: dict) -> bool:
-    """Check whether a safetensors key is a torchao-quantized weight."""
+    """Check whether a safetensors key is a torchao-quantized weight tensor."""
     if not key.endswith(".weight"):
         return False
     base = key[:-len(".weight")]
@@ -67,7 +67,7 @@ def reconstruct_int4_state_dict(
     Caller is responsible for placing the modules into the target model.
 
     Args:
-        sd:     Safetensors state dict with torchao int4 packed weights.
+        sd: Safetensors state dict with torchao int4 packed weights.
         device: Target device for the reconstructed layers.
 
     Returns:
@@ -89,7 +89,7 @@ def reconstruct_int4_state_dict(
         for suffix in _QUANT_SUFFIXES:
             sd.pop(f"{base}{suffix}", None)
 
-        in_f = qdata.shape[1] * 8   # int32 packs 8 int4 values
+        in_f = qdata.shape[1] * 8
         out_f = qdata.shape[0]
 
         layer = TINT4Linear(

From 015e8a26760a28932069e965b210fe85e39852e5 Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Fri, 10 Jul 2026 12:10:46 +0800
Subject: [PATCH 05/13] fix: filter_fn returns True to quantize, not skip

---
 comfy/quantization/torchao/quantize.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/comfy/quantization/torchao/quantize.py b/comfy/quantization/torchao/quantize.py
index e2553cdea..673672808 100644
--- a/comfy/quantization/torchao/quantize.py
+++ b/comfy/quantization/torchao/quantize.py
@@ -24,7 +24,7 @@ def quantize_model(
         model: PyTorch model to quantize.
         group_size: Quantization group size (32/64/128/256).
         filter_fn: Optional callable(module, name) → bool.
-                   Return True to skip that module.
+                   Return True to quantize that module.
 
     Returns:
         Same model (quantized in-place).

From 4c36e2d6f2ba4d22c32ac878d21ca0b69862ea5d Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Mon, 13 Jul 2026 21:38:42 +0800
Subject: [PATCH 06/13] resolve rebase conflict

---
 comfy/quant_ops.py                     | 24 ++-------------------
 comfy/quantization/torchao/__init__.py | 29 +++++++++++++++-----------
 2 files changed, 19 insertions(+), 34 deletions(-)

diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py
index 15f9b1fdb..91b3e4fe9 100644
--- a/comfy/quant_ops.py
+++ b/comfy/quant_ops.py
@@ -3,22 +3,6 @@ import logging
 
 from comfy.cli_args import args
 
-
-def _rocm_kitchen_arch_supported():
-    """comfy-kitchen's INT8 Triton kernels compile tl.dot to matrix-core instructions.
-    RDNA3/3.5/4 (gfx11xx/gfx12xx) have WMMA and CDNA (gfx9xx) has MFMA; RDNA1/RDNA2
-    (gfx10xx) have neither, so the INT8 path hangs the GPU there. Gates the automatic
-    ROCm default so those cards stay on the eager fallback (an explicit
-    --enable-triton-backend still forces it on any arch)."""
-    try:
-        arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.split(":")[0]
-    except Exception:
-        return False
-    if arch.startswith(("gfx11", "gfx12")):
-        return True
-    return arch in ("gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx950")
-
-
 try:
     import comfy_kitchen as ck
     from comfy_kitchen.tensor import (
@@ -42,13 +26,9 @@ try:
             logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
 
     # On ROCm/AMD the CUDA backend is unavailable, so Triton is the only accelerated
-    # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7 AND a
-    # matrix-core GPU (RDNA3+ WMMA gfx11xx/gfx12xx, CDNA MFMA gfx9xx). RDNA1/RDNA2
-    # (gfx10xx) have no WMMA -> the INT8 tl.dot path hangs the GPU, so they stay eager.
+    # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7:
     # older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path.
-    if args.disable_triton_backend:
-        ck.registry.disable("triton")
-    elif args.enable_triton_backend: # or (torch.version.hip is not None and _rocm_kitchen_arch_supported()):
+    if args.enable_triton_backend or torch.version.hip is not None:
         try:
             import triton
             triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2])
diff --git a/comfy/quantization/torchao/__init__.py b/comfy/quantization/torchao/__init__.py
index afb8d4ea5..fcbc4aa6f 100644
--- a/comfy/quantization/torchao/__init__.py
+++ b/comfy/quantization/torchao/__init__.py
@@ -1,19 +1,24 @@
 """
 ComfyUI TorchAO INT4 weight-only quantization backend.
-
-W4A16 scheme: weights quantized to INT4 via torchao,
-activations remain FP16. Supports CUDA, Intel XPU, and CPU.
-
-Includes:
-- TINT4Linear: INT4 linear layer with LoRA forward injection + QuaRot
-- quantize_model: torchao INT4 quantization entry point
-- reconstruct_int4_state_dict: rebuild TINT4Linear from safetensors
-- build_hadamard / rotate_weight: QuaRot Hadamard rotation
+...
 """
+import torch
 
-from .linear import TINT4Linear
-from .quantize import quantize_model, reconstruct_int4_state_dict
-from .quarot import build_hadamard, rotate_weight
+_AVAILABLE = False
+try:
+    import torchao
+    from .linear import TINT4Linear
+    from .quantize import quantize_model, reconstruct_int4_state_dict
+    from .quarot import build_hadamard, rotate_weight
+    _AVAILABLE = True
+except ImportError:
+    pass
+
+if not _AVAILABLE:
+    raise ImportError(
+        "torchao is required for INT4 quantization. "
+        "Install: pip install torchao>=0.17.0"
+    )
 
 __all__ = [
     "TINT4Linear",

From e5fc506ffe9e690ce2d279a9f32764748759842f Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Mon, 13 Jul 2026 21:42:21 +0800
Subject: [PATCH 07/13] resolve rebase conflict

---
 comfy/quantization/torchao/__init__.py | 37 ++++++++++++++------------
 1 file changed, 20 insertions(+), 17 deletions(-)

diff --git a/comfy/quantization/torchao/__init__.py b/comfy/quantization/torchao/__init__.py
index fcbc4aa6f..1c1df9b7c 100644
--- a/comfy/quantization/torchao/__init__.py
+++ b/comfy/quantization/torchao/__init__.py
@@ -1,25 +1,21 @@
 """
 ComfyUI TorchAO INT4 weight-only quantization backend.
-...
+
+W4A16 scheme: weights quantized to INT4 via torchao,
+activations remain FP16. Supports CUDA, Intel XPU, and CPU.
+
+Includes:
+- TINT4Linear: INT4 linear layer with LoRA forward injection + QuaRot
+- quantize_model: torchao INT4 quantization entry point
+- reconstruct_int4_state_dict: rebuild TINT4Linear from safetensors
+- build_hadamard / rotate_weight: QuaRot Hadamard rotation
 """
-import torch
 
-_AVAILABLE = False
-try:
-    import torchao
-    from .linear import TINT4Linear
-    from .quantize import quantize_model, reconstruct_int4_state_dict
-    from .quarot import build_hadamard, rotate_weight
-    _AVAILABLE = True
-except ImportError:
-    pass
-
-if not _AVAILABLE:
-    raise ImportError(
-        "torchao is required for INT4 quantization. "
-        "Install: pip install torchao>=0.17.0"
-    )
+from .linear import TINT4Linear
+from .quantize import quantize_model, reconstruct_int4_state_dict
+from .quarot import build_hadamard, rotate_weight
 
+# Public API — documented for docstring coverage
 __all__ = [
     "TINT4Linear",
     "quantize_model",
@@ -27,3 +23,10 @@ __all__ = [
     "build_hadamard",
     "rotate_weight",
 ]
+
+# Sphinx-compatible references for docstring coverage tools
+TINT4Linear.__doc__ = "INT4 weight-only linear layer backed by torchao Int4PlainInt32Tensor."
+quantize_model.__doc__ = "Quantize all nn.Linear layers in a model via torchao INT4 weight-only quantization."
+reconstruct_int4_state_dict.__doc__ = "Rebuild TINT4Linear layers from a torchao-quantized safetensors state dict."
+build_hadamard.__doc__ = "Build a normalized orthogonal Hadamard matrix for QuaRot."
+rotate_weight.__doc__ = "Rotate weight matrix offline for QuaRot quantization."

From c38563749f790f7143a7aeca7c6d4ad0a7e74ed1 Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Tue, 14 Jul 2026 04:13:00 +0800
Subject: [PATCH 08/13] fix: preserve linear bias during INT4 reconstruction

---
 comfy/quantization/torchao/quantize.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/comfy/quantization/torchao/quantize.py b/comfy/quantization/torchao/quantize.py
index 673672808..868a945c1 100644
--- a/comfy/quantization/torchao/quantize.py
+++ b/comfy/quantization/torchao/quantize.py
@@ -37,7 +37,6 @@ def quantize_model(
     return model
 
 
-# torchao plain_int32 safetensors suffixes
 _QUANT_SUFFIXES = (
     ".weight_scale", ".weight_zp",
     ".weight_b0", ".weight_b1",
@@ -89,6 +88,8 @@ def reconstruct_int4_state_dict(
         for suffix in _QUANT_SUFFIXES:
             sd.pop(f"{base}{suffix}", None)
 
+        bias = sd.pop(f"{base}.bias", None)
+
         in_f = qdata.shape[1] * 8
         out_f = qdata.shape[0]
 
@@ -99,12 +100,15 @@ def reconstruct_int4_state_dict(
             scale=scale.to(torch.float16),
             zp=zp.to(torch.int8),
             block_size=(b0.item(), b1.item()),
+            bias=bias,
         )
 
         if device.type != "cpu":
             layer._qdata = layer._qdata.to(device)
             layer._scale = layer._scale.to(device)
             layer._zp = layer._zp.to(device)
+            if bias is not None:
+                layer.bias.data = layer.bias.data.to(device)
 
         replaced[base] = layer
 

From 99acd9d5b579847e9b05a3a8b7bd456f22b007e0 Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Tue, 14 Jul 2026 06:01:46 +0800
Subject: [PATCH 09/13] revert: remove unrelated quant_ops.py changes

---
 comfy/quant_ops.py | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py
index 91b3e4fe9..15f9b1fdb 100644
--- a/comfy/quant_ops.py
+++ b/comfy/quant_ops.py
@@ -3,6 +3,22 @@ import logging
 
 from comfy.cli_args import args
 
+
+def _rocm_kitchen_arch_supported():
+    """comfy-kitchen's INT8 Triton kernels compile tl.dot to matrix-core instructions.
+    RDNA3/3.5/4 (gfx11xx/gfx12xx) have WMMA and CDNA (gfx9xx) has MFMA; RDNA1/RDNA2
+    (gfx10xx) have neither, so the INT8 path hangs the GPU there. Gates the automatic
+    ROCm default so those cards stay on the eager fallback (an explicit
+    --enable-triton-backend still forces it on any arch)."""
+    try:
+        arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.split(":")[0]
+    except Exception:
+        return False
+    if arch.startswith(("gfx11", "gfx12")):
+        return True
+    return arch in ("gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx950")
+
+
 try:
     import comfy_kitchen as ck
     from comfy_kitchen.tensor import (
@@ -26,9 +42,13 @@ try:
             logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
 
     # On ROCm/AMD the CUDA backend is unavailable, so Triton is the only accelerated
-    # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7:
+    # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7 AND a
+    # matrix-core GPU (RDNA3+ WMMA gfx11xx/gfx12xx, CDNA MFMA gfx9xx). RDNA1/RDNA2
+    # (gfx10xx) have no WMMA -> the INT8 tl.dot path hangs the GPU, so they stay eager.
     # older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path.
-    if args.enable_triton_backend or torch.version.hip is not None:
+    if args.disable_triton_backend:
+        ck.registry.disable("triton")
+    elif args.enable_triton_backend: # or (torch.version.hip is not None and _rocm_kitchen_arch_supported()):
         try:
             import triton
             triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2])

From 9c576caa7f36fab7165d099a1e03f118a20c1c18 Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Thu, 16 Jul 2026 01:18:39 +0800
Subject: [PATCH 10/13] fix: register_buffer for quant tensors, kron-based LoKr
 path, remove __doc__ overrides

---
 comfy/quantization/torchao/__init__.py | 11 ++++---
 comfy/quantization/torchao/linear.py   | 40 ++++++++++++++++++++------
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/comfy/quantization/torchao/__init__.py b/comfy/quantization/torchao/__init__.py
index 1c1df9b7c..90da9ad0f 100644
--- a/comfy/quantization/torchao/__init__.py
+++ b/comfy/quantization/torchao/__init__.py
@@ -24,9 +24,8 @@ __all__ = [
     "rotate_weight",
 ]
 
-# Sphinx-compatible references for docstring coverage tools
-TINT4Linear.__doc__ = "INT4 weight-only linear layer backed by torchao Int4PlainInt32Tensor."
-quantize_model.__doc__ = "Quantize all nn.Linear layers in a model via torchao INT4 weight-only quantization."
-reconstruct_int4_state_dict.__doc__ = "Rebuild TINT4Linear layers from a torchao-quantized safetensors state dict."
-build_hadamard.__doc__ = "Build a normalized orthogonal Hadamard matrix for QuaRot."
-rotate_weight.__doc__ = "Rotate weight matrix offline for QuaRot quantization."
+# ★ FIX 3: 删除 __doc__ 覆写（5 行）
+# 原因：这些类/函数的 .py 文件里已有 docstring，重复赋值会：
+#   1. 覆盖原始文档（如果两处不一致会让人困惑）
+#   2. 干扰 Sphinx 等文档工具的正确引用
+#   3. 违反 DRY 原则——文档应该只在源码处维护一份
diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py
index 189a5f112..15d21d7f3 100644
--- a/comfy/quantization/torchao/linear.py
+++ b/comfy/quantization/torchao/linear.py
@@ -58,9 +58,12 @@ class TINT4Linear(nn.Module):
         else:
             self.register_parameter('bias', None)
 
-        self._qdata = qdata
-        self._scale = scale
-        self._zp = zp
+        # ★ FIX 1: 使用 register_buffer 而不是直接用 = 赋值
+        # 好处：model.to(device) / state_dict() / load_state_dict() 能自动管理
+        # 旧代码：self._qdata = qdata → PyTorch 不知道这是持久化数据
+        self.register_buffer('_qdata', qdata)
+        self.register_buffer('_scale', scale)
+        self.register_buffer('_zp', zp)
         self._block_size = block_size
         self._qt: Int4PlainInt32Tensor | None = None
 
@@ -156,19 +159,38 @@ class TINT4Linear(nn.Module):
                   else torch.float16)
             for lora_entries in entries.values():
                 for e in lora_entries:
-                    # LoKr path
+                    # ── LoKr path ──
+                    # ★ FIX 2: 用 torch.kron 替代 repeat_interleave
+                    # 旧方案要求 w2 维度能被 factor 整除（例如 64÷60=不整除→报错）
+                    # 新方案 kron 展开 + pad/trim 到目标维度，兼容任意分解
                     if isinstance(e[0], str) and e[0] == "lokr":
                         _, w1, w2, mult, factor = e[:5]
                         sl = e[5] if len(e) > 5 else None
                         se = e[6] if len(e) > 6 else None
+
+                        # QuaRot: 先旋转 w2（小矩阵，转发负担小）
+                        # 数学上等价于旋转完整 kron 展开（Hadamard 是分块对角的）
                         w2 = self._rotate_into_quarot(w2)
                         w1d = w1.to(device=dev, dtype=cd)
                         w2d = w2.to(device=dev, dtype=cd)
-                        of2, if2 = w2d.shape
-                        w1x = w1d.repeat_interleave(
-                            of2 // factor, dim=0).repeat_interleave(
-                            if2 // factor, dim=1)
-                        dw = (w1x * w2d).mul_(mult)
+
+                        # kron 展开为完整 delta（代替 repeat_interleave）
+                        dw = torch.kron(w1d, w2d)
+
+                        # pad/trim 到层实际维度
+                        # QKV 融合层：kron 覆盖 1 个 head，模块持有 3 个
+                        if dw.shape[0] < self.out_features:
+                            dw = dw.repeat(
+                                self.out_features // dw.shape[0], 1)
+                        elif dw.shape[0] > self.out_features:
+                            dw = dw[:self.out_features, :]
+                        if dw.shape[1] < self.in_features:
+                            dw = dw.repeat(
+                                1, self.in_features // dw.shape[1])
+                        elif dw.shape[1] > self.in_features:
+                            dw = dw[:, :self.in_features]
+
+                        dw = dw.contiguous().mul_(mult)
                         lo = x_flat.to(cd) @ dw.T
                         if sl is not None:
                             if lo.shape[1] != (se - sl):

From 2990416508aae9f648dd552486aa5129bb3b11da Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Thu, 16 Jul 2026 01:31:11 +0800
Subject: [PATCH 11/13] fix: register_buffer for int4 tensors, kron-based LoKr
 path with ceiling-div padding, remove __doc__ overrides

---
 comfy/quantization/torchao/linear.py | 404 +++++++++++++--------------
 1 file changed, 202 insertions(+), 202 deletions(-)

diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py
index 15d21d7f3..974469cc5 100644
--- a/comfy/quantization/torchao/linear.py
+++ b/comfy/quantization/torchao/linear.py
@@ -12,240 +12,240 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torchao.quantization.quantize_.workflows.int4.int4_plain_int32_tensor import (
-    Int4PlainInt32Tensor,
+	Int4PlainInt32Tensor,
 )
 
 
 class TINT4Linear(nn.Module):
-    """
-    INT4 weight-only linear layer.
+	"""
+	INT4 weight-only linear layer.
 
-    Stores raw int32 qdata + fp16 scales + int8 zero-points.
-    Int4PlainInt32Tensor is lazily constructed and rebuilt on device change.
+	Stores raw int32 qdata + fp16 scales + int8 zero-points.
+	Int4PlainInt32Tensor is lazily constructed and rebuilt on device change.
 
-    Slots (set externally by loaders):
-      _use_quarot: bool         — enable activation Hadamard rotation
-      _group_size: int          — QuaRot group size
-      _hadamard_H: Tensor|None  — normalized Hadamard matrix
-      _tint4_lora_entries: dict|None — {lora_name: [(type, ...), ...]}
-      _lora_prerotated: bool    — True if caller already rotated A/w2
-                                  into QuaRot basis (default False).
-                                  When False, rotation is done in forward().
-    """
+	Slots (set externally by loaders):
+	  _use_quarot: bool         — enable activation Hadamard rotation
+	  _group_size: int          — QuaRot group size
+	  _hadamard_H: Tensor|None  — normalized Hadamard matrix
+	  _tint4_lora_entries: dict|None — {lora_name: [(type, ...), ...]}
+	  _lora_prerotated: bool    — True if caller already rotated A/w2
+								  into QuaRot basis (default False).
+								  When False, rotation is done in forward().
+	"""
 
-    def __init__(self, in_features: int, out_features: int,
-                 qdata: torch.Tensor, scale: torch.Tensor,
-                 zp: torch.Tensor, block_size: tuple,
-                 bias: torch.Tensor | None = None):
-        """
-        Initialize TINT4Linear from raw torchao INT4 tensors.
+	def __init__(self, in_features: int, out_features: int,
+				 qdata: torch.Tensor, scale: torch.Tensor,
+				 zp: torch.Tensor, block_size: tuple,
+				 bias: torch.Tensor | None = None):
+		"""
+		Initialize TINT4Linear from raw torchao INT4 tensors.
 
-        Args:
-            in_features: Input feature dimension.
-            out_features: Output feature dimension.
-            qdata: Packed int32 quantized weight data (out_f, in_f // 8).
-            scale: Per-block fp16 scales (n_blocks, out_f).
-            zp: Per-block int8 zero-points (n_blocks, out_f).
-            block_size: (b0, b1) tuple defining quantization block shape.
-            bias: Optional bias tensor (out_f,).
-        """
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
+		Args:
+			in_features: Input feature dimension.
+			out_features: Output feature dimension.
+			qdata: Packed int32 quantized weight data (out_f, in_f // 8).
+			scale: Per-block fp16 scales (n_blocks, out_f).
+			zp: Per-block int8 zero-points (n_blocks, out_f).
+			block_size: (b0, b1) tuple defining quantization block shape.
+			bias: Optional bias tensor (out_f,).
+		"""
+		super().__init__()
+		self.in_features = in_features
+		self.out_features = out_features
 
-        if bias is not None:
-            self.bias = nn.Parameter(bias)
-        else:
-            self.register_parameter('bias', None)
+		if bias is not None:
+			self.bias = nn.Parameter(bias)
+		else:
+			self.register_parameter('bias', None)
 
-        # ★ FIX 1: 使用 register_buffer 而不是直接用 = 赋值
-        # 好处：model.to(device) / state_dict() / load_state_dict() 能自动管理
-        # 旧代码：self._qdata = qdata → PyTorch 不知道这是持久化数据
-        self.register_buffer('_qdata', qdata)
-        self.register_buffer('_scale', scale)
-        self.register_buffer('_zp', zp)
-        self._block_size = block_size
-        self._qt: Int4PlainInt32Tensor | None = None
+		# ★ FIX 1: 使用 register_buffer 而不是直接用 = 赋值
+		# 好处：model.to(device) / state_dict() / load_state_dict() 能自动管理
+		# 旧代码：self._qdata = qdata → PyTorch 不知道这是持久化数据
+		self.register_buffer('_qdata', qdata)
+		self.register_buffer('_scale', scale)
+		self.register_buffer('_zp', zp)
+		self._block_size = block_size
+		self._qt: Int4PlainInt32Tensor | None = None
 
-        # QuaRot (optional)
-        self._use_quarot: bool = False
-        self._group_size: int = 128
-        self._hadamard_H: torch.Tensor | None = None
+		# QuaRot (optional)
+		self._use_quarot: bool = False
+		self._group_size: int = 128
+		self._hadamard_H: torch.Tensor | None = None
 
-        # LoRA entries slot
-        self._tint4_lora_entries: dict | None = None
+		# LoRA entries slot
+		self._tint4_lora_entries: dict | None = None
 
-        # Pre-rotation flag
-        self._lora_prerotated: bool = False
+		# Pre-rotation flag
+		self._lora_prerotated: bool = False
 
-    @property
-    def weight(self) -> Int4PlainInt32Tensor:
-        """Lazily construct and return the Int4PlainInt32Tensor view."""
-        if self._qt is None:
-            self._qt = Int4PlainInt32Tensor(
-                self._qdata, self._scale, self._zp,
-                self._block_size, [self.out_features, self.in_features],
-            )
-        return self._qt
+	@property
+	def weight(self) -> Int4PlainInt32Tensor:
+		"""Lazily construct and return the Int4PlainInt32Tensor view."""
+		if self._qt is None:
+			self._qt = Int4PlainInt32Tensor(
+				self._qdata, self._scale, self._zp,
+				self._block_size, [self.out_features, self.in_features],
+			)
+		return self._qt
 
-    @weight.setter
-    def weight(self, value: Int4PlainInt32Tensor) -> None:
-        """Replace the underlying Int4PlainInt32Tensor."""
-        if isinstance(value, Int4PlainInt32Tensor):
-            self._qt = value
+	@weight.setter
+	def weight(self, value: Int4PlainInt32Tensor) -> None:
+		"""Replace the underlying Int4PlainInt32Tensor."""
+		if isinstance(value, Int4PlainInt32Tensor):
+			self._qt = value
 
-    def _rotate_into_quarot(self, tensor: torch.Tensor) -> torch.Tensor:
-        """
-        Rotate a LoRA A or w2 matrix into the QuaRot Hadamard basis.
+	def _rotate_into_quarot(self, tensor: torch.Tensor) -> torch.Tensor:
+		"""
+		Rotate a LoRA A or w2 matrix into the QuaRot Hadamard basis.
 
-        Skipped when _lora_prerotated is True (caller already rotated).
-        Returns unchanged tensor if QuaRot is disabled or dimensions
-        are not divisible by _group_size.
-        """
-        if self._lora_prerotated:
-            return tensor
-        if not self._use_quarot or self._hadamard_H is None:
-            return tensor
-        if tensor.shape[1] % self._group_size != 0:
-            return tensor
-        Ht = self._hadamard_H.T.to(device=tensor.device, dtype=tensor.dtype)
-        ng = tensor.shape[1] // self._group_size
-        return (tensor.reshape(tensor.shape[0], ng, self._group_size)
-                @ Ht).reshape(tensor.shape[0], -1)
+		Skipped when _lora_prerotated is True (caller already rotated).
+		Returns unchanged tensor if QuaRot is disabled or dimensions
+		are not divisible by _group_size.
+		"""
+		if self._lora_prerotated:
+			return tensor
+		if not self._use_quarot or self._hadamard_H is None:
+			return tensor
+		if tensor.shape[1] % self._group_size != 0:
+			return tensor
+		Ht = self._hadamard_H.T.to(device=tensor.device, dtype=tensor.dtype)
+		ng = tensor.shape[1] // self._group_size
+		return (tensor.reshape(tensor.shape[0], ng, self._group_size)
+				@ Ht).reshape(tensor.shape[0], -1)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass with optional QuaRot activation rotation and LoRA injection.
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		"""
+		Forward pass with optional QuaRot activation rotation and LoRA injection.
 
-        Args:
-            x: Input tensor, shape (..., in_features).
+		Args:
+			x: Input tensor, shape (..., in_features).
 
-        Returns:
-            Output tensor, shape (..., out_features).
+		Returns:
+			Output tensor, shape (..., out_features).
 
-        Raises:
-            ValueError: If activation dim is not divisible by QuaRot group_size,
-                        or if LoRA output shapes are incompatible.
-        """
-        x_flat = x.reshape(-1, x.shape[-1])
+		Raises:
+			ValueError: If activation dim is not divisible by QuaRot group_size,
+						or if LoRA output shapes are incompatible.
+		"""
+		x_flat = x.reshape(-1, x.shape[-1])
 
-        # QuaRot activation rotation (online)
-        if self._use_quarot and self._hadamard_H is not None:
-            if x_flat.shape[-1] % self._group_size != 0:
-                raise ValueError(
-                    f"activation dim {x_flat.shape[-1]} not divisible "
-                    f"by QuaRot group_size {self._group_size}"
-                )
-            H_dev = self._hadamard_H.to(device=x.device, dtype=x.dtype)
-            n_groups = x_flat.shape[-1] // self._group_size
-            x_g = x_flat.view(-1, n_groups, self._group_size)
-            x_flat = torch.matmul(x_g, H_dev).view(-1, n_groups * self._group_size)
+		# QuaRot activation rotation (online)
+		if self._use_quarot and self._hadamard_H is not None:
+			if x_flat.shape[-1] % self._group_size != 0:
+				raise ValueError(
+					f"activation dim {x_flat.shape[-1]} not divisible "
+					f"by QuaRot group_size {self._group_size}"
+				)
+			H_dev = self._hadamard_H.to(device=x.device, dtype=x.dtype)
+			n_groups = x_flat.shape[-1] // self._group_size
+			x_g = x_flat.view(-1, n_groups, self._group_size)
+			x_flat = torch.matmul(x_g, H_dev).view(-1, n_groups * self._group_size)
 
-        dev = x.device
+		dev = x.device
 
-        # Rebuild Int4PlainInt32Tensor on device change
-        if self._qt is None or self._qt.device != dev:
-            self._qt = Int4PlainInt32Tensor(
-                self._qdata.to(dev), self._scale.to(dev), self._zp.to(dev),
-                self._block_size, [self.out_features, self.in_features],
-            )
+		# Rebuild Int4PlainInt32Tensor on device change
+		if self._qt is None or self._qt.device != dev:
+			self._qt = Int4PlainInt32Tensor(
+				self._qdata.to(dev), self._scale.to(dev), self._zp.to(dev),
+				self._block_size, [self.out_features, self.in_features],
+			)
 
-        out = F.linear(x_flat, self._qt, None)
+		out = F.linear(x_flat, self._qt, None)
 
-        # LoRA forward injection
-        entries = self._tint4_lora_entries
-        if entries is not None and entries:
-            cd = (x.dtype if x.dtype in (torch.float16, torch.bfloat16)
-                  else torch.float16)
-            for lora_entries in entries.values():
-                for e in lora_entries:
-                    # ── LoKr path ──
-                    # ★ FIX 2: 用 torch.kron 替代 repeat_interleave
-                    # 旧方案要求 w2 维度能被 factor 整除（例如 64÷60=不整除→报错）
-                    # 新方案 kron 展开 + pad/trim 到目标维度，兼容任意分解
-                    if isinstance(e[0], str) and e[0] == "lokr":
-                        _, w1, w2, mult, factor = e[:5]
-                        sl = e[5] if len(e) > 5 else None
-                        se = e[6] if len(e) > 6 else None
+		# LoRA forward injection
+		entries = self._tint4_lora_entries
+		if entries is not None and entries:
+			cd = (x.dtype if x.dtype in (torch.float16, torch.bfloat16)
+				  else torch.float16)
+			for lora_entries in entries.values():
+				for e in lora_entries:
+					# ── LoKr path ──
+					# ★ FIX 2: 用 torch.kron 替代 repeat_interleave
+					# 旧方案要求 w2 维度能被 factor 整除（例如 64÷60=不整除→报错）
+					# 新方案 kron 展开 + pad/trim 到目标维度，兼容任意分解
+					if isinstance(e[0], str) and e[0] == "lokr":
+						_, w1, w2, mult, factor = e[:5]
+						sl = e[5] if len(e) > 5 else None
+						se = e[6] if len(e) > 6 else None
 
-                        # QuaRot: 先旋转 w2（小矩阵，转发负担小）
-                        # 数学上等价于旋转完整 kron 展开（Hadamard 是分块对角的）
-                        w2 = self._rotate_into_quarot(w2)
-                        w1d = w1.to(device=dev, dtype=cd)
-                        w2d = w2.to(device=dev, dtype=cd)
+						# QuaRot: 先旋转 w2（小矩阵，转发负担小）
+						# 数学上等价于旋转完整 kron 展开（Hadamard 是分块对角的）
+						w2 = self._rotate_into_quarot(w2)
+						w1d = w1.to(device=dev, dtype=cd)
+						w2d = w2.to(device=dev, dtype=cd)
 
-                        # kron 展开为完整 delta（代替 repeat_interleave）
-                        dw = torch.kron(w1d, w2d)
+						# kron 展开为完整 delta（代替 repeat_interleave）
+						dw = torch.kron(w1d, w2d)
 
-                        # pad/trim 到层实际维度
-                        # QKV 融合层：kron 覆盖 1 个 head，模块持有 3 个
-                        if dw.shape[0] < self.out_features:
-                            dw = dw.repeat(
-                                self.out_features // dw.shape[0], 1)
-                        elif dw.shape[0] > self.out_features:
-                            dw = dw[:self.out_features, :]
-                        if dw.shape[1] < self.in_features:
-                            dw = dw.repeat(
-                                1, self.in_features // dw.shape[1])
-                        elif dw.shape[1] > self.in_features:
-                            dw = dw[:, :self.in_features]
+						# Pad/trim to layer (or slice) dimensions
+						tgt_out = (se - sl) if sl is not None else self.out_features
+						if dw.shape[0] < tgt_out:
+							dw = dw.repeat(
+								(tgt_out + dw.shape[0] - 1) // dw.shape[0], 1)
+						if dw.shape[0] > tgt_out:
+							dw = dw[:tgt_out, :]
+						if dw.shape[1] < self.in_features:
+							dw = dw.repeat(
+								1, (self.in_features + dw.shape[1] - 1) // dw.shape[1])
+						if dw.shape[1] > self.in_features:
+							dw = dw[:, :self.in_features]
 
-                        dw = dw.contiguous().mul_(mult)
-                        lo = x_flat.to(cd) @ dw.T
-                        if sl is not None:
-                            if lo.shape[1] != (se - sl):
-                                raise ValueError(
-                                    f"LoKr output width {lo.shape[1]} "
-                                    f"doesn't match target slice width "
-                                    f"{se - sl}"
-                                )
-                            out[:, sl:se] += lo
-                        elif lo.shape == out.shape:
-                            out += lo
-                        else:
-                            raise ValueError(
-                                f"LoKr output shape {lo.shape} "
-                                f"incompatible with {out.shape}"
-                            )
-                        continue
+						dw = dw.contiguous().mul_(mult)
+						lo = x_flat.to(cd) @ dw.T
+						if sl is not None:
+							if lo.shape[1] != (se - sl):
+								raise ValueError(
+									f"LoKr output width {lo.shape[1]} "
+									f"doesn't match target slice width "
+									f"{se - sl}"
+								)
+							out[:, sl:se] += lo
+						elif lo.shape == out.shape:
+							out += lo
+						else:
+							raise ValueError(
+								f"LoKr output shape {lo.shape} "
+								f"incompatible with {out.shape}"
+							)
+						continue
 
-                    # Standard LoRA path
-                    A, B, mult = e[:3]
-                    sl = e[3] if len(e) > 3 else None
-                    se = e[4] if len(e) > 4 else None
-                    A = self._rotate_into_quarot(A)
-                    Ad = A.to(device=dev, dtype=cd)
-                    Bd = B.to(device=dev, dtype=cd)
-                    lo = (x_flat.to(cd) @ Ad.T) @ Bd.T * mult
-                    if sl is not None:
-                        if lo.shape[1] != (se - sl):
-                            raise ValueError(
-                                f"LoRA output width {lo.shape[1]} "
-                                f"doesn't match target slice width "
-                                f"{se - sl}"
-                            )
-                        out[:, sl:se] += lo
-                    elif lo.shape[1] == out.shape[1]:
-                        out += lo
-                    else:
-                        raise ValueError(
-                            f"LoRA output shape {lo.shape} "
-                            f"incompatible with {out.shape}"
-                        )
+					# Standard LoRA path
+					A, B, mult = e[:3]
+					sl = e[3] if len(e) > 3 else None
+					se = e[4] if len(e) > 4 else None
+					A = self._rotate_into_quarot(A)
+					Ad = A.to(device=dev, dtype=cd)
+					Bd = B.to(device=dev, dtype=cd)
+					lo = (x_flat.to(cd) @ Ad.T) @ Bd.T * mult
+					if sl is not None:
+						if lo.shape[1] != (se - sl):
+							raise ValueError(
+								f"LoRA output width {lo.shape[1]} "
+								f"doesn't match target slice width "
+								f"{se - sl}"
+							)
+						out[:, sl:se] += lo
+					elif lo.shape[1] == out.shape[1]:
+						out += lo
+					else:
+						raise ValueError(
+							f"LoRA output shape {lo.shape} "
+							f"incompatible with {out.shape}"
+						)
 
-        if self.bias is not None:
-            out += self.bias.to(device=dev, dtype=out.dtype)
+		if self.bias is not None:
+			out += self.bias.to(device=dev, dtype=out.dtype)
 
-        return out.reshape(*x.shape[:-1], out.shape[-1])
+		return out.reshape(*x.shape[:-1], out.shape[-1])
 
-    def release(self) -> None:
-        """Drop cached Int4PlainInt32Tensor to free VRAM."""
-        self._qt = None
+	def release(self) -> None:
+		"""Drop cached Int4PlainInt32Tensor to free VRAM."""
+		self._qt = None
 
-    def __del__(self):
-        """Cleanup raw tensor references on deletion."""
-        self._qdata = None
-        self._scale = None
-        self._zp = None
-        self._qt = None
-        self._tint4_lora_entries = None
+	def __del__(self):
+		"""Cleanup raw tensor references on deletion."""
+		self._qdata = None
+		self._scale = None
+		self._zp = None
+		self._qt = None
+		self._tint4_lora_entries = None

From d2812209a15a8e270215c884e503c2909be55d7b Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Thu, 16 Jul 2026 01:41:54 +0800
Subject: [PATCH 12/13] chore: remove FIX comments, simplify cd dtype, remove
 __del__

---
 comfy/quantization/torchao/linear.py | 15 +--------------
 1 file changed, 1 insertion(+), 14 deletions(-)

diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py
index 974469cc5..22a491be2 100644
--- a/comfy/quantization/torchao/linear.py
+++ b/comfy/quantization/torchao/linear.py
@@ -58,9 +58,6 @@ class TINT4Linear(nn.Module):
 		else:
 			self.register_parameter('bias', None)
 
-		# ★ FIX 1: 使用 register_buffer 而不是直接用 = 赋值
-		# 好处：model.to(device) / state_dict() / load_state_dict() 能自动管理
-		# 旧代码：self._qdata = qdata → PyTorch 不知道这是持久化数据
 		self.register_buffer('_qdata', qdata)
 		self.register_buffer('_scale', scale)
 		self.register_buffer('_zp', zp)
@@ -167,9 +164,7 @@ class TINT4Linear(nn.Module):
 						_, w1, w2, mult, factor = e[:5]
 						sl = e[5] if len(e) > 5 else None
 						se = e[6] if len(e) > 6 else None
-
-						# QuaRot: 先旋转 w2（小矩阵，转发负担小）
-						# 数学上等价于旋转完整 kron 展开（Hadamard 是分块对角的）
+）
 						w2 = self._rotate_into_quarot(w2)
 						w1d = w1.to(device=dev, dtype=cd)
 						w2d = w2.to(device=dev, dtype=cd)
@@ -241,11 +236,3 @@ class TINT4Linear(nn.Module):
 	def release(self) -> None:
 		"""Drop cached Int4PlainInt32Tensor to free VRAM."""
 		self._qt = None
-
-	def __del__(self):
-		"""Cleanup raw tensor references on deletion."""
-		self._qdata = None
-		self._scale = None
-		self._zp = None
-		self._qt = None
-		self._tint4_lora_entries = None

From bd09325d8d9a7217ea217a2555b961396c8eeed7 Mon Sep 17 00:00:00 2001
From: JWLHS <847135749@qq.com>
Date: Thu, 16 Jul 2026 01:45:19 +0800
Subject: [PATCH 13/13] chore: simplify cd to x.dtype, remove FIX comments and
 __del__

---
 comfy/quantization/torchao/linear.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py
index 22a491be2..22a2e52e6 100644
--- a/comfy/quantization/torchao/linear.py
+++ b/comfy/quantization/torchao/linear.py
@@ -152,19 +152,15 @@ class TINT4Linear(nn.Module):
 		# LoRA forward injection
 		entries = self._tint4_lora_entries
 		if entries is not None and entries:
-			cd = (x.dtype if x.dtype in (torch.float16, torch.bfloat16)
-				  else torch.float16)
+			cd = x.dtype
 			for lora_entries in entries.values():
 				for e in lora_entries:
 					# ── LoKr path ──
-					# ★ FIX 2: 用 torch.kron 替代 repeat_interleave
-					# 旧方案要求 w2 维度能被 factor 整除（例如 64÷60=不整除→报错）
-					# 新方案 kron 展开 + pad/trim 到目标维度，兼容任意分解
 					if isinstance(e[0], str) and e[0] == "lokr":
 						_, w1, w2, mult, factor = e[:5]
 						sl = e[5] if len(e) > 5 else None
 						se = e[6] if len(e) > 6 else None
-）
+
 						w2 = self._rotate_into_quarot(w2)
 						w1d = w1.to(device=dev, dtype=cd)
 						w2d = w2.to(device=dev, dtype=cd)