1882 lines
86 KiB
Python
1882 lines
86 KiB
Python
"""INT-Crush ops factory — builds custom manual_cast ops for quantized inference.
|
|
|
|
Returns QuantizedTensor weights for VBAR/memory management, but the
|
|
forward() bypasses mixed_precision_ops overhead — going directly from
|
|
packed weights to Triton kernels (or PyTorch fallback) without the
|
|
per-input QuantizedTensor.from_float / cast_bias_weight round-trip.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
from typing import Any, Callable
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from ._quant_utils import rotate_activations
|
|
from .quant_layout import IntCrushInt4Layout, IntCrushInt8Layout
|
|
from ._orbitquant import (
|
|
rpbh_rotate,
|
|
rpbh_rotate_fast,
|
|
codebook_dequant,
|
|
codebook_dequant_fast,
|
|
codebook_dequant_int4,
|
|
codebook_dequant_int8,
|
|
adaln_rtn_dequant,
|
|
parse_orbitquant_metadata,
|
|
load_orbitquant_globals,
|
|
get_quant_scheme_from_metadata,
|
|
infer_quant_scheme_from_scale,
|
|
infer_orbitquant_metadata_from_state_dict,
|
|
# RPBH uniform-grid (non-OrbitQuant) support
|
|
parse_rpbh_metadata,
|
|
load_rpbh_globals,
|
|
infer_rpbh_metadata_from_state_dict,
|
|
ORBITQUANT_NORM_EPS,
|
|
)
|
|
from . import _triton_runtime as _tr
|
|
from ._cuda_graph import CudaGraphRunner
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
# ── Module-level constants ───────────────────────────────────────────────────
|
|
|
|
# Layer-name substrings that disable Hadamard rotation (embedding, norm, head,
|
|
# projection layers are not rotated — they have isotropic weight distributions).
|
|
_NO_ROTATION_NAME_SUBSTRINGS: tuple[str, ...] = (
|
|
"embed", "norm", "modulation", "output", "lm_head", "proj_out",
|
|
)
|
|
|
|
# Safetensors state-dict key suffixes used by INT-Crush quantized layers.
|
|
# All keys are prefixed by the module path (e.g. "model.layer.weight_scale").
|
|
_SD_KEY_WEIGHT = "weight"
|
|
_SD_KEY_SCALE = "weight_scale"
|
|
_SD_KEY_ZP = "weight_zp"
|
|
_SD_KEY_PERM = "weight_perm"
|
|
_SD_KEY_SMOOTH = "weight_smooth"
|
|
_SD_KEY_SMOOTHROT = "weight_smoothrot_factors"
|
|
_SD_KEY_L1 = "weight_L1"
|
|
_SD_KEY_L2 = "weight_L2"
|
|
_SD_KEY_QUANT = "comfy_quant"
|
|
_SD_KEY_BIAS = "bias"
|
|
|
|
# OrbitQuant: M threshold for fused codebook GEMM vs dequant-then-cuBLAS.
|
|
# Set via env: INTCRUSH_OQ_FUSED_THRESHOLD=N (0 disables fused path).
|
|
import os as _os_dispatch
|
|
_OQ_FUSED_GEMM_M_THRESHOLD: int = int(
|
|
_os_dispatch.environ.get("INTCRUSH_OQ_FUSED_THRESHOLD", "512")
|
|
)
|
|
|
|
# Valid float dtypes for weight scales.
|
|
_VALID_SCALE_DTYPES = (torch.float16, torch.float32, torch.bfloat16)
|
|
|
|
# ── OrbitQuant metadata cache ────────────────────────────────────────────────
|
|
# Bridges safetensors metadata to _load_from_state_dict, which can't see
|
|
# global tensors (RPBH params, codebooks) during per-module loading.
|
|
_orbitquant_metadata_cache: dict | None = None
|
|
|
|
|
|
def _set_orbitquant_metadata_cache(metadata: dict | None) -> None:
|
|
"""Store safetensors metadata for OrbitQuant detection during loading."""
|
|
global _orbitquant_metadata_cache
|
|
_orbitquant_metadata_cache = metadata
|
|
|
|
|
|
def _get_orbitquant_metadata_cache() -> dict | None:
|
|
"""Retrieve cached OrbitQuant metadata (read-only, cleared by loader)."""
|
|
return _orbitquant_metadata_cache
|
|
|
|
|
|
# ── RPBH metadata cache (mirrors OrbitQuant cache) ─────────────────────────
|
|
# Bridges safetensors metadata to _load_from_state_dict for RPBH uniform-grid
|
|
# models (format version 1/3 + int_crush.rpbh=true).
|
|
_rpbh_metadata_cache: dict | None = None
|
|
|
|
|
|
def _set_rpbh_metadata_cache(metadata: dict | None) -> None:
|
|
"""Store safetensors metadata for RPBH detection during loading."""
|
|
global _rpbh_metadata_cache
|
|
_rpbh_metadata_cache = metadata
|
|
|
|
|
|
def _get_rpbh_metadata_cache() -> dict | None:
|
|
"""Retrieve cached RPBH metadata (read-only, cleared by loader)."""
|
|
return _rpbh_metadata_cache
|
|
|
|
|
|
# ── Lazy ComfyUI imports (cached on first call) ─────────────────────────────
|
|
|
|
_QuantizedTensor = None
|
|
_cast_bias_weight = None
|
|
_uncast_bias_weight = None
|
|
|
|
|
|
def _get_qt():
|
|
global _QuantizedTensor
|
|
if _QuantizedTensor is None:
|
|
from comfy.quant_ops import QuantizedTensor
|
|
_QuantizedTensor = QuantizedTensor
|
|
return _QuantizedTensor
|
|
|
|
|
|
def _get_cast():
|
|
global _cast_bias_weight, _uncast_bias_weight
|
|
if _cast_bias_weight is None:
|
|
from comfy.ops import cast_bias_weight, uncast_bias_weight
|
|
_cast_bias_weight = cast_bias_weight
|
|
_uncast_bias_weight = uncast_bias_weight
|
|
return _cast_bias_weight, _uncast_bias_weight
|
|
|
|
|
|
def _move_adapter_weights(adapter: Any, device: torch.device, dtype: torch.dtype) -> None:
|
|
"""Move all tensor entries in ``adapter.weights`` to *device*/*dtype*.
|
|
|
|
Non-tensor entries (floats, None, lists) are left untouched.
|
|
"""
|
|
moved: list[Any] = []
|
|
for w in adapter.weights:
|
|
if w is not None and isinstance(w, torch.Tensor):
|
|
moved.append(w.to(device=device, dtype=dtype))
|
|
else:
|
|
moved.append(w)
|
|
adapter.weights = type(adapter.weights)(moved)
|
|
|
|
|
|
def _adapter_weights_fit(adapter: Any, device: torch.device) -> bool:
|
|
"""Check adapter residency against total device VRAM.
|
|
|
|
High-rank LoRAs on large DiTs can exceed 1 GiB of matrices. On
|
|
low-VRAM cards with VBAR-resident weights, keeping them GPU-resident
|
|
starves activation allocations.
|
|
|
|
Gate residency on total VRAM, not instantaneous free memory. VBAR
|
|
makes free memory unreliable at first forward because model pages
|
|
keep faulting in. When the adapter does not fit, the standard-LoRA
|
|
path casts up and down transiently per forward.
|
|
"""
|
|
total = sum(
|
|
w.numel() * w.element_size()
|
|
for w in adapter.weights
|
|
if isinstance(w, torch.Tensor)
|
|
)
|
|
try:
|
|
dev = torch.device(device)
|
|
total_vram = torch.cuda.get_device_properties(dev).total_memory
|
|
except Exception:
|
|
return True
|
|
return total * 5 <= total_vram
|
|
|
|
|
|
def _apply_lokr(
|
|
cache: dict[str, Any],
|
|
x_raw: torch.Tensor,
|
|
out: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""In-place LoKr Kronecker accumulation (combined or decomposed factors)."""
|
|
scale = cache["scale"]
|
|
device = x_raw.device
|
|
dtype = x_raw.dtype
|
|
N = x_raw.shape[0]
|
|
|
|
w1 = cache.get("w1")
|
|
w2_t = cache.get("w2_t")
|
|
w1_a = cache.get("w1_a")
|
|
w1_b = cache.get("w1_b")
|
|
w2_a_t = cache.get("w2_a_t")
|
|
w2_b_t = cache.get("w2_b_t")
|
|
|
|
# Determine dims from whichever form is active.
|
|
if w1 is not None:
|
|
out_l, in_m = w1.shape
|
|
elif w1_a is not None:
|
|
out_l = w1_a.shape[0]
|
|
in_m = w1_b.shape[1]
|
|
else:
|
|
raise ValueError("LoKr cache has no w1 factors")
|
|
|
|
if w2_t is not None:
|
|
in_n, out_k = w2_t.shape
|
|
elif w2_b_t is not None:
|
|
in_n = w2_b_t.shape[0]
|
|
out_k = w2_a_t.shape[1]
|
|
else:
|
|
raise ValueError("LoKr cache has no w2 factors")
|
|
|
|
# Dynamic chunk sizing (cap 256 MB).
|
|
el_sz = 2 if dtype == torch.float16 else 4
|
|
bytes_per_token = out_l * (in_n + out_k) * el_sz
|
|
try:
|
|
free, _ = torch.cuda.mem_get_info(device)
|
|
budget = min(free // 4, 256 * 1024 * 1024)
|
|
except Exception:
|
|
budget = 32 * 1024 * 1024
|
|
chunk = max(1, min(N, budget // max(bytes_per_token, 1)))
|
|
|
|
# Transient device copies; released when this call returns.
|
|
if w1 is not None:
|
|
w1 = w1.to(device)
|
|
else:
|
|
w1_a = w1_a.to(device)
|
|
w1_b = w1_b.to(device)
|
|
if w2_t is not None:
|
|
w2_t = w2_t.to(device)
|
|
else:
|
|
w2_a_t = w2_a_t.to(device)
|
|
w2_b_t = w2_b_t.to(device)
|
|
|
|
if not out.is_contiguous():
|
|
out = out.contiguous()
|
|
out_3d = out.view(N, out_l, out_k)
|
|
|
|
for s in range(0, N, chunk):
|
|
e = min(s + chunk, N)
|
|
X = x_raw[s:e].reshape(e - s, in_m, in_n)
|
|
|
|
# w1 contraction → [C, out_l, in_n]
|
|
if w1 is not None:
|
|
temp = w1 @ X
|
|
else:
|
|
temp = w1_a @ (w1_b @ X)
|
|
|
|
# w2 contraction: out += scale * (temp @ w2^T)
|
|
out_slice = out_3d[s:e].reshape(-1, out_k)
|
|
if w2_t is not None:
|
|
out_slice.addmm_(temp.view(-1, in_n), w2_t,
|
|
beta=1.0, alpha=scale)
|
|
else:
|
|
mid = temp.view(-1, in_n).matmul(w2_b_t)
|
|
out_slice.addmm_(mid, w2_a_t, beta=1.0, alpha=scale)
|
|
del mid
|
|
del temp, X, out_slice
|
|
|
|
del out_3d
|
|
return out
|
|
|
|
|
|
# ── Ops factory ──────────────────────────────────────────────────────────────
|
|
|
|
_ops_cache: dict[tuple, type] = {}
|
|
|
|
|
|
def make_intcrush_ops(
|
|
quant_format: str,
|
|
rot_size: int,
|
|
use_pytorch: bool = False,
|
|
use_w4a16: bool = False,
|
|
metadata: dict | None = None,
|
|
kernel_backend: str = "auto",
|
|
mp_plan: dict | None = None,
|
|
) -> type:
|
|
"""Build a manual_cast ops class with lean INT-Crush forward.
|
|
|
|
Args:
|
|
quant_format: "int4_crush", "int8_crush", or "orbitquant"
|
|
rot_size: Hadamard rotation group size
|
|
use_pytorch: force PyTorch fallback (no Triton kernels)
|
|
use_w4a16: disable W4A8 path, use W4A16 unpack+GEMM instead
|
|
metadata: safetensors file metadata (used for OrbitQuant detection)
|
|
kernel_backend: "auto" (CUDA→Triton→PyTorch), "triton", "cuda", or "pytorch"
|
|
mp_plan: per-layer mixed precision plan {name: {"w": 4|8, "a": 4|8}} or None
|
|
"""
|
|
# Cache key must distinguish models whose class-level metadata differs.
|
|
# _orbitquant_metadata / _rpbh_metadata / _mp_plan are captured at class
|
|
# creation and read back by _load_from_state_dict; sharing a class across
|
|
# models with different quantization config silently reuses stale metadata.
|
|
# Hash the full int_crush.* metadata so distinct configs get distinct classes.
|
|
if metadata:
|
|
_meta_items = sorted(
|
|
(k, str(v)) for k, v in metadata.items()
|
|
if isinstance(k, str) and k.startswith("int_crush.")
|
|
)
|
|
_meta_key = hash(tuple(_meta_items))
|
|
else:
|
|
_meta_key = ""
|
|
_mp_key = json.dumps(mp_plan, sort_keys=True) if mp_plan else ""
|
|
cache_key = (quant_format, rot_size, use_pytorch, use_w4a16,
|
|
_meta_key, kernel_backend, _mp_key)
|
|
if cache_key in _ops_cache:
|
|
return _ops_cache[cache_key]
|
|
from comfy.ops import manual_cast
|
|
|
|
class IntCrushOps(manual_cast):
|
|
_orbitquant_metadata: dict | None = metadata # class-level attribute
|
|
_orbitquant_model_path: str = "" # set during first _load_from_state_dict
|
|
_rpbh_metadata: dict | None = None # class-level RPBH metadata cache
|
|
_mp_plan: dict | None = mp_plan # mixed precision plan (per-layer W/A bit-widths)
|
|
|
|
class Linear(manual_cast.Linear):
|
|
"""INT-Crush Linear: QuantizedTensor storage with lean forward()."""
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self._intcrush_is_quantized: bool = False
|
|
self._intcrush_rot_need: bool = True
|
|
self._intcrush_rot_size: int = rot_size
|
|
self._intcrush_perm: torch.Tensor | None = None
|
|
self._intcrush_smooth: torch.Tensor | None = None
|
|
self._intcrush_smoothrot_factors: torch.Tensor | None = None
|
|
self._intcrush_L1: torch.Tensor | None = None
|
|
self._intcrush_L2: torch.Tensor | None = None
|
|
self._intcrush_use_pytorch: bool = use_pytorch
|
|
self._intcrush_use_w4a16: bool = use_w4a16
|
|
self._intcrush_kernel_backend: str = kernel_backend
|
|
self._intcrush_w_in: int | None = None
|
|
self.quant_format: str | None = None
|
|
self.layout_type: str | None = None
|
|
# ── INT-Crush adapters (LoRA / LoKr / LoHa / OFT / BOFT) ──
|
|
# Ordered stack entries; chained LoRAs each contribute their
|
|
# own scaled residual.
|
|
self._intcrush_adapters: list[dict[str, Any]] = []
|
|
# ── OrbitQuant state ──
|
|
self._intcrush_is_orbitquant: bool = False
|
|
self._intcrush_oq_perm: torch.Tensor | None = None # RPBH permutation
|
|
self._intcrush_oq_signs: torch.Tensor | None = None # RPBH signs
|
|
self._intcrush_oq_h: int | None = None # Hadamard block size
|
|
self._intcrush_oq_codebook: torch.Tensor | None = None # Lloyd-Max centroids [2^b]
|
|
self._intcrush_oq_bits: int = 4 # Weight bit-width
|
|
self._intcrush_oq_d: int | None = None # Original input dim
|
|
self._intcrush_oq_d_padded: int | None = None # Padded input dim
|
|
self._intcrush_oq_quant_scheme: str | None = None # "orbitquant_codebook" or "adaln_rtn"
|
|
self._intcrush_oq_codebook_dev: torch.Tensor | None = None # codebook on GPU (lazy, one-time)
|
|
# ── RPBH uniform-grid state (non-OrbitQuant RPBH) ──
|
|
self._intcrush_is_rpbh: bool = False
|
|
self._intcrush_rpbh_perm: torch.Tensor | None = None # RPBH permutation
|
|
self._intcrush_rpbh_signs: torch.Tensor | None = None # RPBH signs
|
|
self._intcrush_rpbh_h: int | None = None # Hadamard block size
|
|
self._intcrush_rpbh_d: int | None = None # Original input dim
|
|
self._intcrush_rpbh_d_padded: int | None = None # Padded input dim
|
|
# ── Mixed precision state ──
|
|
self._intcrush_mp_weight_bits: int | None = None
|
|
self._intcrush_mp_act_bits: int | None = None
|
|
# CUDA graph runner for the INT8 two-kernel path; None until first use.
|
|
self._graph_runner = None
|
|
|
|
|
|
def _apply_adapter(self, x_raw: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
|
|
"""Apply adapter residuals in stack order.
|
|
|
|
Supports all ComfyUI adapter types (LoRA, LoKr, LoHa,
|
|
OFT, BOFT, GLoRA). Chained LoRAs each contribute their
|
|
own scaled residual; entries apply in attachment order.
|
|
|
|
Standard LoRA uses the memory-efficient ``addmm_`` path
|
|
(only a ``[N, rank]`` intermediate). LoKr uses a custom
|
|
in-place path that accumulates directly into *out*
|
|
without materialising the full delta, keeping the
|
|
decomposed factors resident. All other adapter types go
|
|
through the adapter's own ``h()``/``g()`` methods.
|
|
|
|
Args:
|
|
x_raw: 2-D input ``[N, in_features]`` (unrotated).
|
|
out: 2-D output ``[N, out_features]`` from quantised
|
|
GEMM (+ optional SVD residual).
|
|
"""
|
|
if not self._intcrush_adapters:
|
|
return out
|
|
|
|
# Don't drop _graph_runner — the graph only captures the GEMM.
|
|
|
|
device = x_raw.device
|
|
dtype = x_raw.dtype
|
|
|
|
for entry in self._intcrush_adapters:
|
|
adapter = entry["adapter"]
|
|
strength = entry["strength"]
|
|
|
|
# Move adapter weights to device/dtype on first use or
|
|
# relocation. Under VRAM pressure leave masters on the
|
|
# host and cast transiently per forward. LoKr is exempt:
|
|
# its cache owns the GPU copies, so moving the adapter
|
|
# too would double residency.
|
|
if (not entry["ready"]
|
|
or entry["device"] != device
|
|
or entry["dtype"] != dtype):
|
|
if adapter.name != "lokr":
|
|
entry["transient"] = not _adapter_weights_fit(adapter, device)
|
|
if not entry["transient"]:
|
|
_move_adapter_weights(adapter, device, dtype)
|
|
entry["ready"] = True
|
|
entry["device"] = device
|
|
entry["dtype"] = dtype
|
|
# Drop stale LoKr cache on device/dtype change.
|
|
entry["lokr_cache"] = None
|
|
|
|
# ── Standard LoRA: fast in-place addmm_ path ──────────
|
|
# Only allocates [N, rank] — never materialises the full
|
|
# [N, out_features] delta. This is critical for VRAM.
|
|
if adapter.name == "lora":
|
|
v = adapter.weights
|
|
up, down, alpha = v[0], v[1], v[2]
|
|
if entry["transient"]:
|
|
up = up.to(device=device, dtype=dtype)
|
|
down = down.to(device=device, dtype=dtype)
|
|
rank = down.shape[0]
|
|
scale = ((float(alpha) / rank) if alpha else 1.0)
|
|
scale *= strength
|
|
mid = F.linear(x_raw, down) # [N, rank]
|
|
out.addmm_(mid, up.t(), beta=1, alpha=scale)
|
|
continue
|
|
|
|
# ── LoKr: in-place Kronecker accumulation ─────────────
|
|
# Per-entry CPU cache, rebuilt on device/dtype change.
|
|
if adapter.name == "lokr":
|
|
cache = entry["lokr_cache"]
|
|
if cache is None:
|
|
cache = self._setup_lokr_cache(
|
|
adapter, dtype, strength)
|
|
entry["lokr_cache"] = cache
|
|
out = _apply_lokr(cache, x_raw, out)
|
|
continue
|
|
|
|
# ── All other adapters: h() + g() (in-place add) ──────
|
|
adapter.multiplier = strength
|
|
|
|
# h() returns the additive residual (zeros for OFT/BOFT).
|
|
delta = adapter.h(x_raw, out)
|
|
out.add_(delta.to(out.dtype))
|
|
|
|
# g() may return a new tensor (OFT/BOFT rotation).
|
|
out = adapter.g(out)
|
|
|
|
return out
|
|
|
|
def _setup_lokr_cache(self, adapter: Any, dtype: torch.dtype,
|
|
strength: float) -> dict:
|
|
"""Build cached LoKr factors on CPU.
|
|
|
|
Factors stay CPU-resident; _apply_lokr casts them to the
|
|
compute device transiently per forward, so nothing LoKr
|
|
holds remains on GPU between forwards.
|
|
"""
|
|
v = adapter.weights
|
|
w1, w2 = v[0], v[1]
|
|
alpha = v[2]
|
|
w1_a, w1_b = v[3], v[4]
|
|
w2_a, w2_b = v[5], v[6]
|
|
t2 = v[7] if len(v) > 7 else None
|
|
|
|
cache: dict[str, Any] = {}
|
|
|
|
# ── w1: combined [out_l, in_m] or decomposed ───────────
|
|
if w1 is not None:
|
|
cache["w1"] = w1.to(dtype=dtype)
|
|
cache["w1_a"] = None
|
|
cache["w1_b"] = None
|
|
elif w1_a is not None and w1_b is not None:
|
|
cache["w1"] = None
|
|
cache["w1_a"] = w1_a.to(dtype=dtype) # [out_l, rank]
|
|
cache["w1_b"] = w1_b.to(dtype=dtype) # [rank, in_m]
|
|
else:
|
|
raise ValueError("LoKr adapter has neither w1 nor (w1_a, w1_b)")
|
|
|
|
# ── w2: combined [out_k, in_n], tensor- or matrix-decomp ─
|
|
if w2 is not None:
|
|
cache["w2_t"] = w2.to(dtype=dtype).t().contiguous() # [in_n, out_k]
|
|
cache["w2_a_t"] = None
|
|
cache["w2_b_t"] = None
|
|
elif t2 is not None:
|
|
# 4-D tensor decomposition — flatten to 2-D before .t().
|
|
t2_d = t2.to(dtype=dtype)
|
|
w2a_d = w2_a.to(dtype=dtype)
|
|
w2b_d = w2_b.to(dtype=dtype)
|
|
w2_full = torch.einsum("i j k l, j r, i p -> p r k l",
|
|
t2_d, w2b_d, w2a_d)
|
|
w2_full = w2_full.reshape(
|
|
w2_full.shape[0] * w2_full.shape[1],
|
|
w2_full.shape[2] * w2_full.shape[3]).contiguous() # [out_k, in_n]
|
|
cache["w2_t"] = w2_full.t().contiguous() # [in_n, out_k]
|
|
cache["w2_a_t"] = None
|
|
cache["w2_b_t"] = None
|
|
del t2_d, w2a_d, w2b_d, w2_full
|
|
elif w2_a is not None and w2_b is not None:
|
|
cache["w2_t"] = None
|
|
cache["w2_b_t"] = w2_b.to(dtype=dtype).t().contiguous()
|
|
cache["w2_a_t"] = w2_a.to(dtype=dtype).t().contiguous()
|
|
else:
|
|
raise ValueError("LoKr adapter has neither w2 nor (w2_a, w2_b)")
|
|
|
|
# ── scale ──────────────────────────────────────────────
|
|
if w1_a is not None:
|
|
rank = w1_b.shape[0]
|
|
elif w2_a is not None:
|
|
rank = w2_b.shape[0]
|
|
else:
|
|
rank = None
|
|
|
|
if rank is not None and alpha is not None:
|
|
cache["scale"] = (float(alpha) / rank) * strength
|
|
else:
|
|
cache["scale"] = strength
|
|
|
|
return cache
|
|
|
|
def _prep_activation(self, x: torch.Tensor, pad_to_w_in: bool = True, apply_perm: bool = True) -> torch.Tensor:
|
|
"""Common activation prep: reshape → SmoothRot → (RPBH|Hadamard) → pad → permute.
|
|
|
|
Args:
|
|
x: raw input tensor [..., in_features]
|
|
pad_to_w_in: pad feature dim to match packed weight width
|
|
apply_perm: apply PermuQuant channel permutation
|
|
"""
|
|
x_2d = x.reshape(-1, x.shape[-1])
|
|
|
|
# SmoothRot: 1/s BEFORE rotation (both RPBH and Hadamard).
|
|
smoothrot_factors = self._intcrush_smoothrot_factors
|
|
if smoothrot_factors is not None:
|
|
if smoothrot_factors.device != x_2d.device:
|
|
smoothrot_factors = smoothrot_factors.to(
|
|
device=x_2d.device, dtype=x_2d.dtype)
|
|
self._intcrush_smoothrot_factors = smoothrot_factors
|
|
x_2d = x_2d / smoothrot_factors
|
|
|
|
# ── RPBH rotation (uniform-grid path) ──
|
|
# Takes precedence over group Hadamard. RPBH pads internally
|
|
# to d_padded via rpbh_rotate.
|
|
if self._intcrush_is_rpbh:
|
|
perm = self._intcrush_rpbh_perm
|
|
signs = self._intcrush_rpbh_signs
|
|
if perm is not None and perm.device != x_2d.device:
|
|
perm = perm.to(device=x_2d.device)
|
|
self._intcrush_rpbh_perm = perm
|
|
if signs is not None and signs.device != x_2d.device:
|
|
signs = signs.to(device=x_2d.device)
|
|
self._intcrush_rpbh_signs = signs
|
|
if self._intcrush_use_pytorch:
|
|
x_2d = rpbh_rotate(x_2d, perm, signs, self._intcrush_rpbh_h)
|
|
else:
|
|
x_2d = rpbh_rotate_fast(x_2d, perm, signs, self._intcrush_rpbh_h)
|
|
elif self._intcrush_rot_need:
|
|
x_2d = rotate_activations(x_2d, self._intcrush_rot_size)
|
|
|
|
if pad_to_w_in:
|
|
w_in = self._intcrush_w_in
|
|
if x_2d.shape[-1] < w_in:
|
|
x_2d = F.pad(x_2d, (0, w_in - x_2d.shape[-1]))
|
|
|
|
if apply_perm:
|
|
perm = self._intcrush_perm
|
|
if perm is not None:
|
|
x_2d = x_2d[..., perm]
|
|
|
|
return x_2d
|
|
|
|
def _apply_residuals(self, x: torch.Tensor, out: torch.Tensor,
|
|
x_prep: torch.Tensor | None = None) -> torch.Tensor:
|
|
"""Apply SVD low-rank residual and standard adapter residual."""
|
|
svd_L1 = self._intcrush_L1
|
|
svd_L2 = self._intcrush_L2
|
|
if svd_L1 is not None and svd_L2 is not None:
|
|
# Rebind to self after device/dtype move so subsequent
|
|
# forwards reuse the cached tensor (avoids re-allocation
|
|
# when the model has been relocated to a different device).
|
|
if svd_L1.device != x.device or svd_L1.dtype != x.dtype:
|
|
svd_L1 = svd_L1.to(device=x.device, dtype=x.dtype)
|
|
self._intcrush_L1 = svd_L1
|
|
if svd_L2.device != x.device or svd_L2.dtype != x.dtype:
|
|
svd_L2 = svd_L2.to(device=x.device, dtype=x.dtype)
|
|
self._intcrush_L2 = svd_L2
|
|
svd_x = x_prep if x_prep is not None else x.reshape(-1, x.shape[-1])
|
|
mid = svd_x.to(svd_L1.dtype) @ svd_L2.T
|
|
out.addmm_(mid, svd_L1.T.to(out.dtype), beta=1, alpha=1)
|
|
|
|
# _apply_adapter may return a new tensor — use the return value.
|
|
out = self._apply_adapter(x.reshape(-1, x.shape[-1]), out)
|
|
return out
|
|
|
|
def _finish_forward(self, out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
|
|
"""Convert dtype and reshape to original input shape.
|
|
|
|
Note: uncast (VBAR unpin) is handled by forward()'s
|
|
try/finally — not here — so that the pin is always
|
|
released even when _apply_residuals raises.
|
|
"""
|
|
out = out.to(x.dtype)
|
|
return out.reshape(*x.shape[:-1], -1)
|
|
|
|
# ── forward-path helpers ──────────────────────────────────────────
|
|
|
|
def _setup_casting(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None, Callable[..., None] | None, type]:
|
|
"""Resolve weight/bias/uncast for the forward pass.
|
|
|
|
Handles VBAR casting, low-VRAM offload, and the fast path
|
|
where no casting is needed. Returns (weight, bias, uncast, QT).
|
|
"""
|
|
QT = _get_qt()
|
|
wfn = self.weight_function
|
|
bfn = self.bias_function
|
|
need_cast = (
|
|
self.comfy_cast_weights
|
|
or wfn or bfn
|
|
or getattr(self, "weight_lowvram_function", None) is not None
|
|
or getattr(self, "bias_lowvram_function", None) is not None
|
|
)
|
|
|
|
if need_cast:
|
|
weight_dtype = (
|
|
self.weight._params.orig_dtype
|
|
if isinstance(self.weight, QT) else x.dtype
|
|
)
|
|
_cbw, _ucbw = _get_cast()
|
|
weight, bias, offload_stream = _cbw(
|
|
self, input=None, dtype=weight_dtype,
|
|
device=x.device, bias_dtype=x.dtype, offloadable=True,
|
|
)
|
|
uncast = lambda: _ucbw(self, weight, bias, offload_stream)
|
|
else:
|
|
weight = self.weight
|
|
bias = self.bias
|
|
# One-time move: if weight is still on CPU (deferred
|
|
# from state-dict load), bring it to device now.
|
|
# QuantizedTensor.to() moves _qdata, scale, perm, zp
|
|
# atomically via __torch_dispatch__ → _handle_to.
|
|
if weight.device != x.device:
|
|
weight = weight.to(device=x.device)
|
|
self.weight = nn.Parameter(weight, requires_grad=False)
|
|
uncast = None
|
|
|
|
return weight, bias, uncast, QT
|
|
|
|
def _forward_predequant(self, weight: torch.Tensor, bias: torch.Tensor | None,
|
|
x: torch.Tensor) -> torch.Tensor:
|
|
"""Forward when weight is already dequantized (e.g. wrapper patch ran)."""
|
|
x_2d = self._prep_activation(x, pad_to_w_in=False, apply_perm=False)
|
|
# Dequantized weights are float32 (e.g. LowVramPatch dequantizes
|
|
# and applies LoRA in cast_bias_weight) while compute runs in the
|
|
# model dtype.
|
|
if weight.dtype != x_2d.dtype:
|
|
weight = weight.to(x_2d.dtype)
|
|
if bias is not None and bias.dtype != x_2d.dtype:
|
|
bias = bias.to(x_2d.dtype)
|
|
out = F.linear(x_2d, weight, bias)
|
|
out = self._apply_residuals(x, out, x_prep=x_2d)
|
|
return self._finish_forward(out, x)
|
|
|
|
# ── OrbitQuant forward paths ──
|
|
|
|
def _rpbh_rotate(self, x_2d: torch.Tensor) -> torch.Tensor:
|
|
"""RPBH-rotate an activation, honoring use_pytorch + device laziness."""
|
|
device = x_2d.device
|
|
perm = self._intcrush_oq_perm
|
|
signs = self._intcrush_oq_signs
|
|
if perm is not None and perm.device != device:
|
|
perm = perm.to(device=device)
|
|
self._intcrush_oq_perm = perm
|
|
if signs is not None and signs.device != device:
|
|
signs = signs.to(device=device)
|
|
self._intcrush_oq_signs = signs
|
|
|
|
if self._intcrush_use_pytorch:
|
|
return rpbh_rotate(x_2d, perm, signs, self._intcrush_oq_h)
|
|
return rpbh_rotate_fast(x_2d, perm, signs, self._intcrush_oq_h)
|
|
|
|
def _dequant_codebook(self, qdata: torch.Tensor, weight: Any) -> torch.Tensor:
|
|
"""Dequantize a codebook weight to float."""
|
|
device = qdata.device
|
|
cb = self._intcrush_oq_codebook
|
|
if cb is not None and cb.device != device:
|
|
cb = cb.to(device=device)
|
|
self._intcrush_oq_codebook = cb
|
|
|
|
d_padded = self._intcrush_oq_d_padded
|
|
if self._intcrush_use_pytorch:
|
|
return codebook_dequant(
|
|
qdata, weight._params.scale, cb,
|
|
self._intcrush_oq_d, d_padded, self._intcrush_oq_bits,
|
|
)
|
|
return codebook_dequant_fast(
|
|
qdata, weight._params.scale, cb,
|
|
self._intcrush_oq_d, d_padded, self._intcrush_oq_bits,
|
|
)
|
|
|
|
def _oom_safe_matmul(
|
|
self,
|
|
x: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
bias: torch.Tensor | None,
|
|
) -> torch.Tensor:
|
|
# Memory-aware matmul: out = x @ weight.T (+ bias).
|
|
M, K = x.shape
|
|
N = weight.shape[0]
|
|
device = x.device
|
|
dtype = x.dtype
|
|
|
|
# Defensive dtype match (callers already pre-cast).
|
|
if weight.dtype != dtype:
|
|
weight = weight.to(dtype)
|
|
|
|
out_bytes = M * N * x.element_size()
|
|
try:
|
|
free, _ = torch.cuda.mem_get_info(device)
|
|
except Exception:
|
|
free = 0
|
|
|
|
# Only pay empty_cache() when we genuinely lack room for the
|
|
# output (plus a 2x slack for the matmul workspace).
|
|
if free < out_bytes * 2:
|
|
try:
|
|
torch.cuda.empty_cache()
|
|
free, _ = torch.cuda.mem_get_info(device)
|
|
except Exception:
|
|
free = 0
|
|
|
|
# Pre-allocate the output once (the dominant allocation).
|
|
out = torch.empty((M, N), dtype=dtype, device=device)
|
|
|
|
# Chunk over M to bound the cuBLAS workspace.
|
|
chunk = M
|
|
if free > 0:
|
|
bytes_per_row = (K + N) * x.element_size()
|
|
denom = max(bytes_per_row * 4, 1)
|
|
chunk = max(1, min(M, free // denom))
|
|
|
|
w_t = weight.t() # view [K, N]
|
|
for s in range(0, M, chunk):
|
|
e = min(s + chunk, M)
|
|
torch.mm(x[s:e], w_t, out=out[s:e])
|
|
|
|
if bias is not None:
|
|
out.add_(bias.to(dtype))
|
|
return out
|
|
|
|
def _forward_orbitquant(
|
|
self,
|
|
qdata: torch.Tensor,
|
|
weight: Any,
|
|
bias: torch.Tensor | None,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""OrbitQuant forward: RPBH rotate → GEMM.
|
|
|
|
Codebook layers: fused codebook GEMM or dequant-then-cuBLAS.
|
|
AdaLN layers: NO rotation — weights are stored UNROTATED.
|
|
"""
|
|
x_2d = x.reshape(-1, x.shape[-1])
|
|
scheme = self._intcrush_oq_quant_scheme
|
|
|
|
if scheme == "adaln_rtn":
|
|
# AdaLN weights are UNROTATED — do NOT rotate x.
|
|
w_float = adaln_rtn_dequant(
|
|
qdata,
|
|
weight._params.scale,
|
|
self._intcrush_oq_d,
|
|
self._intcrush_oq_bits,
|
|
)
|
|
# Cast to the compute dtype BEFORE the matmul so the fp32
|
|
# dequant tensor is freed (refcount drops to 0) rather than
|
|
# coexisting with an fp16 cast copy during the GEMM.
|
|
if w_float.dtype != x_2d.dtype:
|
|
w_float = w_float.to(x_2d.dtype)
|
|
out = self._oom_safe_matmul(x_2d, w_float, bias)
|
|
del w_float
|
|
else:
|
|
x_rot = self._rpbh_rotate(x_2d)
|
|
M = x_rot.shape[0]
|
|
|
|
# Fused codebook GEMM for small M, dequant-then-cuBLAS for large M.
|
|
_use_fused = (
|
|
_tr.TRITON_CODEBOOK_GEMM
|
|
and not self._intcrush_use_pytorch
|
|
and M <= _OQ_FUSED_GEMM_M_THRESHOLD
|
|
and self._intcrush_oq_quant_scheme == "orbitquant_codebook"
|
|
)
|
|
|
|
if _use_fused:
|
|
# Lazy one-time device move for codebook.
|
|
dev = x_rot.device
|
|
oq_cb = self._intcrush_oq_codebook_dev
|
|
if oq_cb is None:
|
|
oq_cb = self._intcrush_oq_codebook.to(device=dev)
|
|
self._intcrush_oq_codebook_dev = oq_cb
|
|
# Norms come from Params which may be on CPU.
|
|
oq_norms = weight._params.scale.to(device=dev)
|
|
out = _tr.fused_codebook_gemm_dequant(
|
|
x_rot, qdata, oq_norms, oq_cb,
|
|
bias=bias, bits=self._intcrush_oq_bits,
|
|
)
|
|
else:
|
|
w_float = self._dequant_codebook(qdata, weight)
|
|
if w_float.dtype != x_rot.dtype:
|
|
w_float = w_float.to(x_rot.dtype)
|
|
out = self._oom_safe_matmul(x_rot, w_float, bias)
|
|
|
|
# Apply adapter residuals (LoRA, etc.) on raw activation x.
|
|
out = self._apply_residuals(x, out, x_prep=None)
|
|
|
|
return self._finish_forward(out, x)
|
|
|
|
def _forward_predequant_oq(
|
|
self,
|
|
weight: torch.Tensor,
|
|
bias: torch.Tensor | None,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""OrbitQuant forward when weight is already dequantized to float.
|
|
|
|
Codebook layers: rotate activation for cancellation.
|
|
AdaLN layers: weight is UNROTATED — do NOT rotate.
|
|
"""
|
|
x_2d = x.reshape(-1, x.shape[-1])
|
|
scheme = self._intcrush_oq_quant_scheme
|
|
|
|
if scheme == "adaln_rtn":
|
|
out = self._oom_safe_matmul(x_2d, weight, bias)
|
|
else:
|
|
x_rot = self._rpbh_rotate(x_2d)
|
|
out = self._oom_safe_matmul(x_rot, weight, bias)
|
|
|
|
out = self._apply_residuals(x, out, x_prep=None)
|
|
return self._finish_forward(out, x)
|
|
|
|
def _forward_int4_w4a8(self, qdata: torch.Tensor, weight: Any,
|
|
bias: torch.Tensor | None, x: torch.Tensor) -> torch.Tensor:
|
|
"""INT4 Path 1: W4A8 — unpack INT4→INT8, dynamic-quantize activations,
|
|
fused INT8 GEMM + dequant (fastest when all Triton kernels available)."""
|
|
x_2d = self._prep_activation(x)
|
|
x_int8, s_a = _tr.dynamic_quantize_activation(x_2d)
|
|
needs_zp = weight._params.zp is not None
|
|
x_2d_sum = x_2d.sum(dim=-1, keepdim=True) if needs_zp else None
|
|
w_int8 = _tr.unpack_int4_to_int8(qdata, self._intcrush_w_in)
|
|
scale_flat = weight._params.scale.reshape(-1).float().contiguous()
|
|
_gemm_fn = _tr.fused_w4a8_gemm_dequant if _tr.TRITON_W4A8_GEMM else _tr.fused_int8_gemm_dequant
|
|
out = _gemm_fn(
|
|
x_int8, w_int8, scale_flat, s_a,
|
|
bias=bias, out_dtype=x.dtype,
|
|
)
|
|
del x_int8, w_int8
|
|
if needs_zp:
|
|
zp_cor = (scale_flat * weight._params.zp.reshape(-1).float()).to(out.dtype)
|
|
out.addmm_(x_2d_sum.to(out.dtype), zp_cor.unsqueeze(0), beta=1, alpha=-1)
|
|
|
|
out = self._apply_residuals(x, out, x_prep=x_2d)
|
|
return self._finish_forward(out, x)
|
|
|
|
def _forward_int4_w4a16(self, qdata: torch.Tensor, weight: Any,
|
|
bias: torch.Tensor | None, x: torch.Tensor) -> torch.Tensor:
|
|
"""INT4 Path 2: W4A16 — Triton unpack to float16, then cuBLAS GEMM."""
|
|
x_2d = self._prep_activation(x)
|
|
scale_flat = weight._params.scale.reshape(-1).float().contiguous()
|
|
needs_zp = weight._params.zp is not None
|
|
x_2d_sum = x_2d.sum(dim=-1, keepdim=True) if needs_zp else None
|
|
weight_f16 = _tr.unpack_int4_to_float16(qdata, scale_flat, self._intcrush_w_in)
|
|
# Cast activation to match weight dtype (handles bfloat16 input with float16 weights)
|
|
out = F.linear(x_2d.to(weight_f16.dtype), weight_f16,
|
|
bias.to(weight_f16.dtype) if bias is not None else None)
|
|
del weight_f16
|
|
if needs_zp:
|
|
zp_cor = (scale_flat * weight._params.zp.reshape(-1).float()).to(out.dtype)
|
|
out.addmm_(x_2d_sum.to(out.dtype), zp_cor.unsqueeze(0), beta=1, alpha=-1)
|
|
|
|
out = self._apply_residuals(x, out, x_prep=x_2d)
|
|
return self._finish_forward(out, x)
|
|
|
|
def _forward_int4_pytorch(self, qdata: torch.Tensor, weight: Any,
|
|
bias: torch.Tensor | None,
|
|
x: torch.Tensor) -> torch.Tensor:
|
|
"""INT4 Path 3: PyTorch fallback — full dequant to float, then F.linear."""
|
|
w_float = IntCrushInt4Layout.dequantize(qdata, weight._params).to(x.device, x.dtype)
|
|
x_2d = self._prep_activation(x, pad_to_w_in=True)
|
|
out = F.linear(x_2d, w_float, bias.to(w_float.dtype) if bias is not None else None)
|
|
del w_float
|
|
out = self._apply_residuals(x, out, x_prep=x_2d)
|
|
return self._finish_forward(out, x)
|
|
|
|
def _forward_int8(self, qdata: torch.Tensor, weight: Any,
|
|
bias: torch.Tensor | None, x: torch.Tensor) -> torch.Tensor:
|
|
"""INT8 forward: small-batch PyTorch fallback or Triton two-kernel path."""
|
|
w_scale = weight._params.scale
|
|
zp = weight._params.zp
|
|
|
|
# Determine M and compute_dtype from raw x (before prep).
|
|
# prep_activation is a view-based reshape + CUDA ops, so M is
|
|
# the same before and after prep.
|
|
M = x.numel() // x.shape[-1]
|
|
compute_dtype = (x.dtype if x.dtype in (torch.float16, torch.bfloat16)
|
|
else torch.bfloat16)
|
|
|
|
# PyTorch fallback: honour explicit kernel_backend="pytorch"
|
|
# even when Triton is available. _forward_int8 otherwise ignores
|
|
# _intcrush_use_pytorch and falls into the Triton path below.
|
|
if self._intcrush_use_pytorch:
|
|
x_2d = self._prep_activation(x)
|
|
if self._intcrush_smoothrot_factors is None:
|
|
smooth = self._intcrush_smooth
|
|
if smooth is not None:
|
|
if smooth.device != x_2d.device:
|
|
smooth = smooth.to(device=x_2d.device, dtype=x_2d.dtype)
|
|
self._intcrush_smooth = smooth
|
|
x_2d = x_2d / smooth
|
|
w_scale_2d = w_scale.reshape(-1, 1) if w_scale.ndim == 1 else w_scale
|
|
w_float = qdata.to(compute_dtype) * w_scale_2d.to(compute_dtype)
|
|
if zp is not None:
|
|
w_float = w_float - (zp.reshape(-1, 1).to(compute_dtype)
|
|
* w_scale_2d.to(compute_dtype))
|
|
out = F.linear(x_2d, w_float, bias)
|
|
del x_2d, w_float
|
|
out = self._apply_residuals(x, out)
|
|
return self._finish_forward(out, x)
|
|
|
|
# ── CUDA Graph path ──
|
|
# Requirements: Triton available, no zero-point,
|
|
# no VBAR/casting (bias must be static), graphs enabled.
|
|
_can_graph = (
|
|
_tr.TRITON_AVAILABLE
|
|
and not self._intcrush_use_pytorch
|
|
and zp is None
|
|
and not self._intcrush_is_rpbh
|
|
and not (self.comfy_cast_weights or self.weight_function
|
|
or self.bias_function or self.weight_lowvram_function
|
|
or self.bias_lowvram_function)
|
|
and _tr.USE_CUDA_GRAPHS
|
|
)
|
|
if _can_graph:
|
|
if self._graph_runner is None:
|
|
self._graph_runner = CudaGraphRunner(self, safe_no_clone=False)
|
|
out = self._graph_runner.run(x, bias, compute_dtype)
|
|
# _apply_residuals uses the raw x (not the prepped x_2d),
|
|
# which is what the current code already does.
|
|
out = self._apply_residuals(x, out)
|
|
return self._finish_forward(out, x)
|
|
|
|
# ── Original paths (prep happens here) ──
|
|
x_2d = self._prep_activation(x)
|
|
|
|
# Old SmoothQuant: 1/s AFTER Hadamard (only for non-SmoothRot layers).
|
|
if self._intcrush_smoothrot_factors is None:
|
|
smooth = self._intcrush_smooth
|
|
if smooth is not None:
|
|
if smooth.device != x_2d.device:
|
|
smooth = smooth.to(device=x_2d.device, dtype=x_2d.dtype)
|
|
self._intcrush_smooth = smooth
|
|
x_2d = x_2d / smooth
|
|
|
|
# No Triton: dequant weights to float and use F.linear.
|
|
if not _tr.TRITON_AVAILABLE:
|
|
w_scale_2d = w_scale.reshape(-1, 1) if w_scale.ndim == 1 else w_scale
|
|
w_float = qdata.to(compute_dtype) * w_scale_2d.to(compute_dtype)
|
|
if zp is not None:
|
|
w_float = w_float - (zp.reshape(-1, 1).to(compute_dtype)
|
|
* w_scale_2d.to(compute_dtype))
|
|
out = F.linear(x_2d, w_float, bias)
|
|
del x_2d, w_float
|
|
elif _tr.USE_FUSED_W8A8 and _tr.TRITON_FUSED_QGD:
|
|
# Fused single-kernel: quantize fp16→int8 + GEMM + dequant.
|
|
needs_zp = zp is not None
|
|
x_2d_sum = x_2d.sum(dim=-1, keepdim=True) if needs_zp else None
|
|
out = _tr.fused_quant_int8_gemm_dequant(
|
|
x_2d, qdata, w_scale,
|
|
bias=bias, out_dtype=compute_dtype,
|
|
)
|
|
del x_2d
|
|
if needs_zp:
|
|
scale_flat = w_scale.reshape(-1).float().contiguous()
|
|
zp_cor = (scale_flat * zp.reshape(-1).float()).to(out.dtype)
|
|
out.addmm_(x_2d_sum.to(out.dtype), zp_cor.unsqueeze(0), beta=1, alpha=-1)
|
|
else:
|
|
# Two-kernel path (quantize → GEMM) is always faster than
|
|
# the fused quant+GEMM kernel for batch > 16: the fused kernel
|
|
# reads fp16 twice (abs-max pass + quantize pass), while this
|
|
# path reads fp16 once, then int8 on the GEMM (half bandwidth).
|
|
needs_zp = zp is not None
|
|
x_2d_sum = x_2d.sum(dim=-1, keepdim=True) if needs_zp else None
|
|
x_int8, s_a = _tr.dynamic_quantize_activation(x_2d)
|
|
del x_2d
|
|
out = _tr.fused_int8_gemm_dequant(
|
|
x_int8, qdata, w_scale, s_a,
|
|
bias=bias, out_dtype=compute_dtype,
|
|
)
|
|
del x_int8
|
|
if needs_zp:
|
|
scale_flat = w_scale.reshape(-1).float().contiguous()
|
|
zp_cor = (scale_flat * zp.reshape(-1).float()).to(out.dtype)
|
|
out.addmm_(x_2d_sum.to(out.dtype), zp_cor.unsqueeze(0), beta=1, alpha=-1)
|
|
|
|
out = self._apply_residuals(x, out)
|
|
return self._finish_forward(out, x)
|
|
|
|
# ── comfy-kitchen CUDA forward paths ────────────────────────────
|
|
|
|
def _resolve_ck_backend(self) -> str:
|
|
"""Resolve 'auto' kernel_backend to a concrete backend name.
|
|
|
|
Returns 'cuda', 'triton', or 'pytorch'.
|
|
"""
|
|
backend = self._intcrush_kernel_backend
|
|
if backend == "pytorch":
|
|
return "pytorch"
|
|
if backend == "cuda":
|
|
return "cuda" if _tr.can_use_cuda_int8() else "pytorch"
|
|
if backend == "triton":
|
|
return "triton" if _tr.TRITON_AVAILABLE else "pytorch"
|
|
# auto: prefer CUDA → Triton → PyTorch
|
|
if _tr.can_use_cuda_int8():
|
|
return "cuda"
|
|
if _tr.TRITON_AVAILABLE:
|
|
return "triton"
|
|
return "pytorch"
|
|
|
|
def _resolve_ck_backend_int4(self) -> str:
|
|
"""Resolve 'auto' kernel_backend for INT4 (needs both unpack + GEMM).
|
|
|
|
Returns 'cuda', 'triton', or 'pytorch'.
|
|
"""
|
|
backend = self._intcrush_kernel_backend
|
|
if backend == "pytorch":
|
|
return "pytorch"
|
|
if backend == "cuda":
|
|
return "cuda" if _tr.can_use_cuda_int4() else "pytorch"
|
|
if backend == "triton":
|
|
return "triton" if _tr.TRITON_AVAILABLE else "pytorch"
|
|
# auto
|
|
if _tr.can_use_cuda_int4():
|
|
return "cuda"
|
|
if _tr.TRITON_AVAILABLE:
|
|
return "triton"
|
|
return "pytorch"
|
|
|
|
def _ck_prep_smoothrot(self, x_2d: torch.Tensor) -> torch.Tensor:
|
|
"""Apply SmoothRot factors (1/s BEFORE rotation) in-place if present.
|
|
|
|
Returns x_2d (possibly modified). Safe to call on any device.
|
|
"""
|
|
smoothrot = self._intcrush_smoothrot_factors
|
|
if smoothrot is not None:
|
|
if smoothrot.device != x_2d.device:
|
|
smoothrot = smoothrot.to(device=x_2d.device, dtype=x_2d.dtype)
|
|
self._intcrush_smoothrot_factors = smoothrot
|
|
x_2d = x_2d / smoothrot
|
|
return x_2d
|
|
|
|
def _ck_prep_pad(self, x_2d: torch.Tensor) -> torch.Tensor:
|
|
"""Pad feature dim to match packed weight width."""
|
|
w_in = self._intcrush_w_in
|
|
if w_in is not None and x_2d.shape[-1] < w_in:
|
|
x_2d = F.pad(x_2d, (0, w_in - x_2d.shape[-1]))
|
|
return x_2d
|
|
|
|
def _forward_int4_ck(
|
|
self,
|
|
qdata: torch.Tensor,
|
|
weight: Any,
|
|
bias: torch.Tensor | None,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""INT4 forward via comfy-kitchen CUDA kernels.
|
|
|
|
Two sub-paths selected at runtime:
|
|
|
|
* **Fused ConvRot** (rot_size=256, no perm/RPBH/smooth):
|
|
SmoothRot → pad → ``int8_linear(convrot=True)`` — the Hadamard
|
|
rotation, per-row activation quantization, and INT8 GEMM all
|
|
happen inside a single CUDA kernel.
|
|
|
|
* **Separate rotation** (all other cases): ``_prep_activation``
|
|
handles SmoothRot + rotation + pad + permute, then unpack
|
|
INT4→INT8, then ``int8_linear(convrot=False)``.
|
|
"""
|
|
from comfy_kitchen.backends.cuda import (
|
|
int8_linear as ck_int8_linear,
|
|
prepare_int4_weight_for_int8_linear as ck_unpack_int4,
|
|
)
|
|
|
|
scale_flat = weight._params.scale.reshape(-1).float().contiguous()
|
|
zp = weight._params.zp
|
|
|
|
has_perm = self._intcrush_perm is not None
|
|
has_rpbh = self._intcrush_is_rpbh
|
|
has_smooth = self._intcrush_smooth is not None
|
|
|
|
use_fused = _tr.can_use_fused_convrot(
|
|
self._intcrush_rot_size, has_perm, has_rpbh, has_smooth,
|
|
)
|
|
|
|
if use_fused:
|
|
# ── Fused ConvRot: rotation + quantize + GEMM in one call ──
|
|
x_2d = x.reshape(-1, x.shape[-1])
|
|
x_2d = self._ck_prep_smoothrot(x_2d)
|
|
x_2d = self._ck_prep_pad(x_2d)
|
|
|
|
# Unpack INT4 → INT8 (CUDA kernel). Cast to int8 for
|
|
# dtype safety — comfy-kitchen stores packed int4 as int8.
|
|
w_int8 = ck_unpack_int4(qdata.to(torch.int8).contiguous())
|
|
|
|
out = ck_int8_linear(
|
|
x_2d.contiguous(),
|
|
w_int8,
|
|
scale_flat,
|
|
bias=bias,
|
|
out_dtype=x.dtype,
|
|
convrot=True,
|
|
convrot_groupsize=256,
|
|
)
|
|
else:
|
|
# ── Separate rotation path ──
|
|
x_2d = self._prep_activation(x)
|
|
|
|
w_int8 = ck_unpack_int4(qdata.to(torch.int8).contiguous())
|
|
|
|
out = ck_int8_linear(
|
|
x_2d.contiguous(),
|
|
w_int8,
|
|
scale_flat,
|
|
bias=bias,
|
|
out_dtype=x.dtype,
|
|
convrot=False,
|
|
)
|
|
|
|
# Zero-point correction (asymmetric quantization)
|
|
if zp is not None:
|
|
x_2d_sum = x_2d.sum(dim=-1, keepdim=True)
|
|
zp_cor = (scale_flat * zp.reshape(-1).float()).to(out.dtype)
|
|
out.addmm_(x_2d_sum.to(out.dtype), zp_cor.unsqueeze(0),
|
|
beta=1, alpha=-1)
|
|
|
|
out = self._apply_residuals(x, out, x_prep=x_2d)
|
|
return self._finish_forward(out, x)
|
|
|
|
def _forward_int8_ck(
|
|
self,
|
|
qdata: torch.Tensor,
|
|
weight: Any,
|
|
bias: torch.Tensor | None,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""INT8 forward via comfy-kitchen CUDA kernels.
|
|
|
|
Same two sub-paths as the INT4 CUDA forward (fused ConvRot vs
|
|
separate rotation), plus SmoothQuant handling for non-SmoothRot
|
|
layers in the separate-rotation path.
|
|
"""
|
|
from comfy_kitchen.backends.cuda import int8_linear as ck_int8_linear
|
|
|
|
w_scale = weight._params.scale.float().contiguous()
|
|
zp = weight._params.zp
|
|
|
|
has_perm = self._intcrush_perm is not None
|
|
has_rpbh = self._intcrush_is_rpbh
|
|
has_smooth = self._intcrush_smooth is not None
|
|
|
|
use_fused = _tr.can_use_fused_convrot(
|
|
self._intcrush_rot_size, has_perm, has_rpbh, has_smooth,
|
|
)
|
|
|
|
if use_fused:
|
|
# ── Fused ConvRot ──
|
|
x_2d = x.reshape(-1, x.shape[-1])
|
|
x_2d = self._ck_prep_smoothrot(x_2d)
|
|
x_2d = self._ck_prep_pad(x_2d)
|
|
|
|
out = ck_int8_linear(
|
|
x_2d.contiguous(),
|
|
qdata.contiguous(),
|
|
w_scale,
|
|
bias=bias,
|
|
out_dtype=x.dtype,
|
|
convrot=True,
|
|
convrot_groupsize=256,
|
|
)
|
|
else:
|
|
# ── Separate rotation ──
|
|
x_2d = self._prep_activation(x)
|
|
|
|
# SmoothQuant: 1/s AFTER Hadamard (non-SmoothRot layers only)
|
|
if self._intcrush_smoothrot_factors is None:
|
|
smooth = self._intcrush_smooth
|
|
if smooth is not None:
|
|
if smooth.device != x_2d.device:
|
|
smooth = smooth.to(device=x_2d.device,
|
|
dtype=x_2d.dtype)
|
|
self._intcrush_smooth = smooth
|
|
x_2d = x_2d / smooth
|
|
|
|
out = ck_int8_linear(
|
|
x_2d.contiguous(),
|
|
qdata.contiguous(),
|
|
w_scale,
|
|
bias=bias,
|
|
out_dtype=x.dtype,
|
|
convrot=False,
|
|
)
|
|
|
|
# Zero-point correction
|
|
if zp is not None:
|
|
x_2d_sum = x_2d.sum(dim=-1, keepdim=True)
|
|
scale_flat = w_scale.reshape(-1).float().contiguous()
|
|
zp_cor = (scale_flat * zp.reshape(-1).float()).to(out.dtype)
|
|
out.addmm_(x_2d_sum.to(out.dtype), zp_cor.unsqueeze(0),
|
|
beta=1, alpha=-1)
|
|
|
|
out = self._apply_residuals(x, out)
|
|
return self._finish_forward(out, x)
|
|
|
|
# ── state-dict loading helpers ────────────────────────────────────
|
|
|
|
@staticmethod
|
|
def _detect_weight_format(weight: torch.Tensor | None, scale: torch.Tensor | None) -> str | None:
|
|
"""Detect INT-Crush format by weight/scale dtype signature.
|
|
|
|
Returns "int4_crush" / "int8_crush" / None.
|
|
"""
|
|
if weight is None or scale is None:
|
|
return None
|
|
if scale.dtype not in _VALID_SCALE_DTYPES:
|
|
return None
|
|
if weight.dtype == torch.uint8:
|
|
return "int4_crush"
|
|
if weight.dtype == torch.int8:
|
|
return "int8_crush"
|
|
return None
|
|
|
|
def _pop_intcrush_tensors(self, state_dict: dict, prefix: str) -> dict[str, Any]:
|
|
"""Pop all INT-Crush tensors from *state_dict* and return them as a dict.
|
|
|
|
The returned dict has keys: weight, scale, zp, perm, smooth,
|
|
smoothrot_factors, L1, L2, bias. All popped keys are
|
|
returned so they can be removed from missing_keys.
|
|
"""
|
|
result: dict[str, Any] = {}
|
|
result["weight"] = state_dict.pop(prefix + _SD_KEY_WEIGHT)
|
|
result["scale"] = state_dict.pop(prefix + _SD_KEY_SCALE)
|
|
result["zp"] = state_dict.pop(prefix + _SD_KEY_ZP, None)
|
|
result["perm"] = state_dict.pop(prefix + _SD_KEY_PERM, None)
|
|
result["smooth"] = state_dict.pop(prefix + _SD_KEY_SMOOTH, None)
|
|
result["smoothrot_factors"] = state_dict.pop(prefix + _SD_KEY_SMOOTHROT, None)
|
|
result["L1"] = state_dict.pop(prefix + _SD_KEY_L1, None)
|
|
result["L2"] = state_dict.pop(prefix + _SD_KEY_L2, None)
|
|
result["bias"] = state_dict.pop(prefix + _SD_KEY_BIAS, None)
|
|
state_dict.pop(prefix + _SD_KEY_QUANT, None)
|
|
return result
|
|
|
|
def _store_intcrush_attrs(
|
|
self,
|
|
tensors: dict[str, Any],
|
|
params: Any,
|
|
device: torch.device,
|
|
) -> None:
|
|
"""Store quantization metadata and auxiliary tensors as instance attributes.
|
|
|
|
Tensors are kept on CPU (dtype-cast only) so that ComfyUI's
|
|
ModelPatcher/aimdo can manage VRAM placement. They are moved
|
|
to the active device lazily on first forward via _prep_activation
|
|
and _apply_residuals.
|
|
"""
|
|
self._intcrush_perm = params.perm
|
|
self._intcrush_w_in = self.in_features
|
|
self._intcrush_smooth = (
|
|
tensors["smooth"].float()
|
|
if tensors["smooth"] is not None else None
|
|
)
|
|
# SmoothRot factors: applied BEFORE Hadamard (1/s → R).
|
|
# Distinct from _intcrush_smooth which is applied AFTER (R → /s).
|
|
self._intcrush_smoothrot_factors = (
|
|
tensors["smoothrot_factors"].float()
|
|
if tensors["smoothrot_factors"] is not None else None
|
|
)
|
|
# SVD low-rank factors: FP16 branch absorbed before quantization.
|
|
self._intcrush_L1 = (
|
|
tensors["L1"].to(dtype=torch.float16)
|
|
if tensors["L1"] is not None else None
|
|
)
|
|
self._intcrush_L2 = (
|
|
tensors["L2"].to(dtype=torch.float16)
|
|
if tensors["L2"] is not None else None
|
|
)
|
|
|
|
@staticmethod
|
|
def _clear_known_missing_keys(
|
|
missing_keys: list[str],
|
|
prefix: str,
|
|
) -> None:
|
|
"""Remove INT-Crush state-dict keys from missing_keys so PyTorch
|
|
doesn't complain about legitimately absent standard keys."""
|
|
known_suffixes = (
|
|
_SD_KEY_WEIGHT, _SD_KEY_SCALE, _SD_KEY_ZP, _SD_KEY_PERM,
|
|
_SD_KEY_SMOOTH, _SD_KEY_SMOOTHROT, _SD_KEY_L1, _SD_KEY_L2,
|
|
_SD_KEY_BIAS, _SD_KEY_QUANT,
|
|
)
|
|
for suffix in known_suffixes:
|
|
k = prefix + suffix
|
|
if k in missing_keys:
|
|
missing_keys.remove(k)
|
|
|
|
# ── PyTorch Module overrides ──────────────────────────────────────
|
|
|
|
def _load_from_state_dict(
|
|
self,
|
|
state_dict: dict,
|
|
prefix: str,
|
|
local_metadata: Any,
|
|
strict: bool,
|
|
missing_keys: list[str],
|
|
unexpected_keys: list[str],
|
|
error_msgs: list[str],
|
|
) -> None:
|
|
from comfy.quant_ops import QuantizedTensor, get_layout_class, QUANT_ALGOS
|
|
|
|
weight = state_dict.get(prefix + _SD_KEY_WEIGHT)
|
|
scale = state_dict.get(prefix + _SD_KEY_SCALE)
|
|
fmt = self._detect_weight_format(weight, scale)
|
|
|
|
if fmt is None:
|
|
# Not an INT-Crush layer — delegate to default loader.
|
|
super()._load_from_state_dict(
|
|
state_dict, prefix, local_metadata,
|
|
strict, missing_keys, unexpected_keys, error_msgs)
|
|
return
|
|
|
|
# ── OrbitQuant detection ──
|
|
oq_metadata = getattr(type(self), '_orbitquant_metadata', None)
|
|
if oq_metadata is None:
|
|
oq_metadata = _get_orbitquant_metadata_cache()
|
|
if oq_metadata is not None:
|
|
type(self)._orbitquant_metadata = oq_metadata
|
|
quant_scheme = get_quant_scheme_from_metadata(oq_metadata, prefix)
|
|
|
|
if quant_scheme is None and scale is not None and oq_metadata is not None:
|
|
# No per-layer metadata key — infer from scale shape.
|
|
# Requires oq_metadata to avoid false-positives on non-OrbitQuant files.
|
|
quant_scheme = infer_quant_scheme_from_scale(scale)
|
|
|
|
if quant_scheme in ("orbitquant_codebook", "adaln_rtn"):
|
|
self._load_orbitquant_layer(
|
|
state_dict, prefix, oq_metadata,
|
|
weight, scale, quant_scheme,
|
|
missing_keys, unexpected_keys,
|
|
)
|
|
return
|
|
|
|
# ── RPBH uniform-grid detection (non-OrbitQuant RPBH) ──
|
|
rpbh_metadata = getattr(type(self), '_rpbh_metadata', None)
|
|
if rpbh_metadata is None:
|
|
rpbh_metadata = _get_rpbh_metadata_cache()
|
|
if rpbh_metadata is not None:
|
|
type(self)._rpbh_metadata = rpbh_metadata
|
|
|
|
is_rpbh_layer = (
|
|
rpbh_metadata is not None
|
|
and parse_rpbh_metadata(rpbh_metadata) is not None
|
|
)
|
|
|
|
if is_rpbh_layer:
|
|
self._load_rpbh_layer(
|
|
state_dict, prefix, rpbh_metadata,
|
|
weight, scale, fmt,
|
|
missing_keys, unexpected_keys,
|
|
)
|
|
return
|
|
|
|
# ── INT-Crush layer: wrap as QuantizedTensor ──
|
|
qconfig = QUANT_ALGOS[fmt]
|
|
layout_type = qconfig["comfy_tensor_layout"]
|
|
layout_cls = get_layout_class(layout_type)
|
|
tensors = self._pop_intcrush_tensors(state_dict, prefix)
|
|
|
|
device = getattr(self, 'weight', None)
|
|
device = device.device if device is not None else torch.device("cpu")
|
|
|
|
# Determine whether this layer needs Hadamard rotation based on name.
|
|
name_lower = prefix.lower()
|
|
rot_need = (
|
|
rot_size > 0 and not any(
|
|
p in name_lower for p in _NO_ROTATION_NAME_SUBSTRINGS
|
|
)
|
|
)
|
|
self._intcrush_rot_need = rot_need
|
|
self._intcrush_rot_size = rot_size
|
|
self._intcrush_is_quantized = True
|
|
|
|
# ── Mixed precision: read per-layer bit-widths ──
|
|
mp_plan = getattr(type(self), '_mp_plan', None)
|
|
if mp_plan:
|
|
bare_name = prefix.rstrip(".")
|
|
mp_entry = mp_plan.get(bare_name)
|
|
if mp_entry is None:
|
|
# Fallback: strip common prefixes
|
|
for strip_prefix in ("model.", "model.diffusion_model."):
|
|
if bare_name.startswith(strip_prefix):
|
|
mp_entry = mp_plan.get(
|
|
bare_name[len(strip_prefix):])
|
|
if mp_entry is not None:
|
|
break
|
|
if mp_entry is not None:
|
|
self._intcrush_mp_weight_bits = mp_entry["w"]
|
|
self._intcrush_mp_act_bits = mp_entry["a"]
|
|
_log.debug("mp_plan: %s -> W%dA%d",
|
|
bare_name, mp_entry["w"], mp_entry["a"])
|
|
else:
|
|
_log.debug("mp_plan: no entry for '%s'", bare_name)
|
|
|
|
perm = tensors["perm"]
|
|
zp = tensors["zp"]
|
|
|
|
if fmt == "int4_crush":
|
|
params = layout_cls.Params(
|
|
scale=tensors["scale"].float(),
|
|
orig_dtype=torch.float16,
|
|
orig_shape=(self.out_features, self.in_features),
|
|
rot_need=rot_need,
|
|
rot_size=rot_size,
|
|
perm=perm if perm is not None else None,
|
|
zp=zp if zp is not None else None,
|
|
)
|
|
else:
|
|
params = layout_cls.Params(
|
|
scale=tensors["scale"].float().reshape(-1),
|
|
orig_dtype=torch.float16,
|
|
orig_shape=(self.out_features, self.in_features),
|
|
rot_need=rot_need,
|
|
rot_size=rot_size,
|
|
perm=perm if perm is not None else None,
|
|
zp=zp if zp is not None else None,
|
|
)
|
|
|
|
self._store_intcrush_attrs(tensors, params, device)
|
|
|
|
# Load to CPU. ComfyUI's model loader calls .to(device) on
|
|
# the model for full-VRAM mode, and cast_bias_weight handles
|
|
# per-layer VBAR demand-paging for low-VRAM mode. This avoids
|
|
# forcing all weights resident before the ModelPatcher can
|
|
# decide. QuantizedTensor.to(device) moves _qdata, scale,
|
|
# perm, and zp atomically via __torch_dispatch__ → _handle_to.
|
|
self.weight = nn.Parameter(
|
|
QuantizedTensor(
|
|
tensors["weight"].to(dtype=qconfig["storage_t"]),
|
|
layout_type, params,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
self.quant_format = fmt
|
|
self.layout_type = layout_type
|
|
|
|
bias_tensor = tensors["bias"]
|
|
if bias_tensor is not None:
|
|
self.bias = nn.Parameter(
|
|
bias_tensor.to(device=device), requires_grad=False,
|
|
)
|
|
else:
|
|
self.bias = None
|
|
|
|
if perm is not None:
|
|
self.register_parameter("weight_perm",
|
|
nn.Parameter(perm, requires_grad=False))
|
|
if zp is not None:
|
|
self.register_parameter("weight_zp",
|
|
nn.Parameter(zp, requires_grad=False))
|
|
|
|
self._clear_known_missing_keys(missing_keys, prefix)
|
|
|
|
def _load_orbitquant_layer(
|
|
self,
|
|
state_dict: dict,
|
|
prefix: str,
|
|
metadata: dict,
|
|
weight: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
quant_scheme: str,
|
|
missing_keys: list[str],
|
|
unexpected_keys: list[str],
|
|
) -> None:
|
|
"""Load an OrbitQuant-quantized layer."""
|
|
from .quant_layout import OrbitQuantLayout
|
|
from ._orbitquant import infer_orbitquant_metadata_from_state_dict
|
|
|
|
# If metadata is None, try to infer from state dict globals
|
|
if metadata is None:
|
|
metadata = infer_orbitquant_metadata_from_state_dict(state_dict)
|
|
if metadata is None:
|
|
raise ValueError(
|
|
"OrbitQuant layer detected (scale shape indicates codebook/AdaLN) "
|
|
"but no metadata found. Ensure the safetensors file has "
|
|
"int_crush.format_version=4 metadata."
|
|
)
|
|
# Store inferred metadata on the ops class for subsequent layers
|
|
type(self)._orbitquant_metadata = metadata
|
|
|
|
oq_meta = parse_orbitquant_metadata(metadata)
|
|
if oq_meta is None:
|
|
raise ValueError("OrbitQuant metadata not found in file")
|
|
|
|
model_path = metadata.get("int_crush.orbitquant_dims", "default")
|
|
|
|
# Load globals (idempotent — cached after first call)
|
|
globals_cache = load_orbitquant_globals(state_dict, metadata, model_path)
|
|
|
|
type(self)._orbitquant_model_path = model_path
|
|
|
|
# Determine d from the module's declared in_features
|
|
d = self.in_features
|
|
d_padded = oq_meta["d_padded_map"].get(d, d)
|
|
bits = oq_meta["codebook_bits"]
|
|
|
|
# Get RPBH params for this dimension.
|
|
# AdaLN-RTN layers are UNROTATED — no RPBH needed.
|
|
if quant_scheme == "adaln_rtn":
|
|
rpbh = globals_cache["rpbh"].get(d)
|
|
if rpbh is None:
|
|
# AdaLN may share d with a codebook layer (common for
|
|
# hidden dim); if not, use a sensible fallback.
|
|
rpbh = {
|
|
"perm": None, "signs": None, "h": 1, "d_padded": d,
|
|
}
|
|
codebook = None # AdaLN-RTN uses group scales, not a codebook
|
|
else:
|
|
rpbh = globals_cache["rpbh"].get(d)
|
|
if rpbh is None:
|
|
raise ValueError(
|
|
f"OrbitQuant: no RPBH params for d={d}. "
|
|
f"Available dims: {list(globals_cache['rpbh'].keys())}"
|
|
)
|
|
codebook = globals_cache["codebooks"].get((d, bits))
|
|
if codebook is None:
|
|
raise ValueError(
|
|
f"OrbitQuant: no codebook for (d={d}, bits={bits}). "
|
|
f"Available: {list(globals_cache['codebooks'].keys())}"
|
|
)
|
|
|
|
# Pop tensors
|
|
tensors = self._pop_intcrush_tensors(state_dict, prefix)
|
|
|
|
# Store OrbitQuant state
|
|
self._intcrush_is_quantized = True
|
|
self._intcrush_is_orbitquant = True
|
|
self._intcrush_oq_perm = rpbh["perm"]
|
|
self._intcrush_oq_signs = rpbh["signs"]
|
|
self._intcrush_oq_h = rpbh["h"]
|
|
self._intcrush_oq_codebook = codebook
|
|
self._intcrush_oq_bits = bits
|
|
self._intcrush_oq_d = d
|
|
self._intcrush_oq_d_padded = d_padded
|
|
self._intcrush_oq_quant_scheme = quant_scheme
|
|
self._intcrush_rot_need = False # OrbitQuant uses RPBH, not group Hadamard
|
|
|
|
# Build Params — include codebook for dequantize() support
|
|
layout_cls = OrbitQuantLayout
|
|
params = layout_cls.Params(
|
|
scale=tensors["scale"].float(),
|
|
orig_dtype=torch.float16,
|
|
orig_shape=(self.out_features, d),
|
|
quant_scheme=quant_scheme,
|
|
d=d,
|
|
d_padded=d_padded,
|
|
bits=bits,
|
|
codebook=codebook, # enables dequantize() for LoRA/VBAR/fallback
|
|
)
|
|
|
|
# Register as QuantizedTensor
|
|
QT = _get_qt()
|
|
self.weight = nn.Parameter(
|
|
QT(tensors["weight"], "OrbitQuantLayout", params),
|
|
requires_grad=False,
|
|
)
|
|
self.quant_format = "orbitquant"
|
|
self.layout_type = "OrbitQuantLayout"
|
|
|
|
# Bias
|
|
if tensors["bias"] is not None:
|
|
self.bias = nn.Parameter(tensors["bias"], requires_grad=False)
|
|
else:
|
|
self.bias = None
|
|
|
|
self._clear_known_missing_keys(missing_keys, prefix)
|
|
|
|
# Remove OrbitQuant global keys from unexpected_keys (first layer only)
|
|
for key in list(unexpected_keys):
|
|
if key.startswith("_rpbh_") or key.startswith("_orbitquant_"):
|
|
unexpected_keys.remove(key)
|
|
|
|
def _load_rpbh_layer(
|
|
self,
|
|
state_dict: dict,
|
|
prefix: str,
|
|
metadata: dict,
|
|
weight: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
fmt: str,
|
|
missing_keys: list[str],
|
|
unexpected_keys: list[str],
|
|
) -> None:
|
|
"""Load an RPBH uniform-grid layer (INT4/INT8 packed + RPBH rotation)."""
|
|
from comfy.quant_ops import QuantizedTensor, get_layout_class, QUANT_ALGOS
|
|
|
|
# If metadata is None, try to infer from state dict globals
|
|
if metadata is None:
|
|
metadata = infer_rpbh_metadata_from_state_dict(state_dict)
|
|
if metadata is None:
|
|
raise ValueError(
|
|
"RPBH layer detected but no metadata found. "
|
|
"Ensure the safetensors file has int_crush.rpbh=true metadata."
|
|
)
|
|
type(self)._rpbh_metadata = metadata
|
|
|
|
rpbh_meta = parse_rpbh_metadata(metadata)
|
|
if rpbh_meta is None:
|
|
raise ValueError("RPBH metadata not found in file")
|
|
|
|
model_path = metadata.get("int_crush.rpbh_dims", "default")
|
|
|
|
# Load globals (idempotent — cached after first call)
|
|
globals_cache = load_rpbh_globals(state_dict, metadata, model_path)
|
|
|
|
# Determine d from the module's declared in_features
|
|
d = self.in_features
|
|
d_padded = rpbh_meta["d_padded_map"].get(d, d)
|
|
|
|
# Get RPBH params for this dimension
|
|
rpbh = globals_cache["rpbh"].get(d)
|
|
if rpbh is None:
|
|
raise ValueError(
|
|
f"RPBH: no RPBH params for d={d}. "
|
|
f"Available dims: {list(globals_cache['rpbh'].keys())}"
|
|
)
|
|
|
|
# Pop standard INT-Crush tensors (weight, scale, zp, perm, smooth, etc.)
|
|
tensors = self._pop_intcrush_tensors(state_dict, prefix)
|
|
|
|
device = getattr(self, 'weight', None)
|
|
device = device.device if device is not None else torch.device("cpu")
|
|
|
|
# Store RPBH state — replaces group Hadamard rotation.
|
|
self._intcrush_is_quantized = True
|
|
self._intcrush_is_rpbh = True
|
|
self._intcrush_rpbh_perm = rpbh["perm"]
|
|
self._intcrush_rpbh_signs = rpbh["signs"]
|
|
self._intcrush_rpbh_h = rpbh["h"]
|
|
self._intcrush_rpbh_d = d
|
|
self._intcrush_rpbh_d_padded = d_padded
|
|
|
|
# RPBH replaces group Hadamard rotation — disable it.
|
|
self._intcrush_rot_need = False
|
|
self._intcrush_rot_size = rpbh_meta["block_size"]
|
|
|
|
# w_in must equal d_padded (the padded weight width).
|
|
self._intcrush_w_in = d_padded
|
|
|
|
# Build Params using the standard layout class (IntCrushInt4Layout or IntCrushInt8Layout).
|
|
qconfig = QUANT_ALGOS[fmt]
|
|
layout_type = qconfig["comfy_tensor_layout"]
|
|
layout_cls = get_layout_class(layout_type)
|
|
|
|
perm = tensors["perm"]
|
|
zp = tensors["zp"]
|
|
|
|
params = layout_cls.Params(
|
|
scale=tensors["scale"].float() if fmt == "int4_crush" else tensors["scale"].float().reshape(-1),
|
|
orig_dtype=torch.float16,
|
|
orig_shape=(self.out_features, self.in_features),
|
|
rot_need=False, # RPBH handles rotation, not group Hadamard
|
|
rot_size=0, # no group Hadamard
|
|
perm=perm if perm is not None else None,
|
|
zp=zp if zp is not None else None,
|
|
)
|
|
|
|
self._store_intcrush_attrs(tensors, params, device)
|
|
|
|
# Load to CPU — same as standard INT-Crush path.
|
|
self.weight = nn.Parameter(
|
|
QuantizedTensor(
|
|
tensors["weight"].to(dtype=qconfig["storage_t"]),
|
|
layout_type, params,
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
self.quant_format = fmt
|
|
self.layout_type = layout_type
|
|
|
|
bias_tensor = tensors["bias"]
|
|
if bias_tensor is not None:
|
|
self.bias = nn.Parameter(
|
|
bias_tensor.to(device=device), requires_grad=False,
|
|
)
|
|
else:
|
|
self.bias = None
|
|
|
|
if perm is not None:
|
|
self.register_parameter("weight_perm",
|
|
nn.Parameter(perm, requires_grad=False))
|
|
if zp is not None:
|
|
self.register_parameter("weight_zp",
|
|
nn.Parameter(zp, requires_grad=False))
|
|
|
|
self._clear_known_missing_keys(missing_keys, prefix)
|
|
|
|
# Remove RPBH global keys from unexpected_keys (first layer only)
|
|
for key in list(unexpected_keys):
|
|
if key.startswith("_rpbh_perm_") or key.startswith("_rpbh_signs_"):
|
|
unexpected_keys.remove(key)
|
|
|
|
def convert_weight(self, weight: torch.Tensor, inplace: bool = False) -> torch.Tensor:
|
|
"""Keep QuantizedTensor as-is for ModelPatcher compatibility."""
|
|
from comfy.quant_ops import QuantizedTensor
|
|
if self._intcrush_is_quantized and isinstance(weight, QuantizedTensor):
|
|
return weight
|
|
# OrbitQuant may have a plain float weight after set_weight
|
|
if getattr(self, '_intcrush_is_orbitquant', False):
|
|
return weight
|
|
if isinstance(weight, QuantizedTensor):
|
|
return weight.dequantize()
|
|
return weight
|
|
|
|
def set_weight(self, out_weight: torch.Tensor, inplace_update: bool = False, seed: int = 0,
|
|
return_weight: bool = False, **kwargs: Any) -> torch.Tensor | None:
|
|
"""Requantize float weights, preserving rotation and permutation params.
|
|
|
|
``inplace_update`` is accepted for ModelPatcher API compat and
|
|
ignored (in-place mutation of a QuantizedTensor would corrupt
|
|
the packed buffer). ``seed`` is forwarded to stochastic rounding;
|
|
``return_weight`` returns the weight instead of assigning it.
|
|
"""
|
|
if self._graph_runner is not None:
|
|
self._graph_runner = None
|
|
|
|
if not self._intcrush_is_quantized:
|
|
new_weight = out_weight.to(self.weight.dtype)
|
|
if return_weight:
|
|
return new_weight
|
|
self.weight = nn.Parameter(new_weight, requires_grad=False)
|
|
return
|
|
|
|
from comfy.quant_ops import QuantizedTensor, get_layout_class
|
|
import dataclasses
|
|
|
|
if isinstance(out_weight, QuantizedTensor):
|
|
if return_weight:
|
|
return out_weight
|
|
self.weight = nn.Parameter(out_weight, requires_grad=False)
|
|
return
|
|
|
|
# OrbitQuant: can't requantize without RPBH + codebook.
|
|
# Store as a plain float weight — forward will use _forward_predequant_oq.
|
|
if getattr(self, '_intcrush_is_orbitquant', False):
|
|
new_weight = out_weight.to(dtype=torch.float16)
|
|
if return_weight:
|
|
return new_weight
|
|
self.weight = nn.Parameter(new_weight, requires_grad=False)
|
|
return
|
|
|
|
layout_cls = get_layout_class(self.layout_type)
|
|
packed, params = layout_cls.quantize(
|
|
out_weight.float(), scale="recalculate",
|
|
stochastic_rounding=seed,
|
|
)
|
|
# quantize() stamps orig_dtype from the fp32 patch tensor.
|
|
# Restore the pre-patch compute dtype, as ComfyUI's ops.set_weight
|
|
# does with .to(self.weight.dtype).
|
|
params = dataclasses.replace(
|
|
params,
|
|
orig_dtype=self.weight.dtype,
|
|
rot_need=self._intcrush_rot_need,
|
|
rot_size=self._intcrush_rot_size,
|
|
perm=self._intcrush_perm,
|
|
)
|
|
qt = QuantizedTensor(packed, self.layout_type, params)
|
|
if return_weight:
|
|
return qt
|
|
self.weight = nn.Parameter(qt, requires_grad=False)
|
|
|
|
def set_bias(self, out_bias: torch.Tensor | None, inplace_update: bool = False, seed: int = 0,
|
|
return_weight: bool = False, **kwargs: Any) -> torch.Tensor | None:
|
|
"""Replace the bias. ``inplace_update`` and ``seed`` are accepted
|
|
for ModelPatcher API compat and have no effect."""
|
|
if out_bias is None:
|
|
return None
|
|
if return_weight:
|
|
return out_bias
|
|
self.bias = nn.Parameter(out_bias, requires_grad=False)
|
|
|
|
@torch.no_grad()
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
|
# ── Non-quantized fallback ──
|
|
if not self._intcrush_is_quantized:
|
|
return super().forward(x)
|
|
|
|
# ── ComfyUI weight casting (VBAR + low-VRAM) ──
|
|
weight, bias, uncast, QT = self._setup_casting(x)
|
|
|
|
# try/finally guarantees VBAR pages are always unpinned,
|
|
# even when _apply_residuals (LoRA) raises.
|
|
try:
|
|
# ── Weight already dequantized (e.g. wrapper patch ran) ──
|
|
if not isinstance(weight, QT):
|
|
if self._intcrush_is_orbitquant:
|
|
return self._forward_predequant_oq(weight, bias, x)
|
|
return self._forward_predequant(weight, bias, x)
|
|
|
|
qdata = weight._qdata
|
|
|
|
# ── OrbitQuant forward ──
|
|
if self._intcrush_is_orbitquant:
|
|
return self._forward_orbitquant(qdata, weight, bias, x)
|
|
|
|
# ── INT4 forward paths ──
|
|
if isinstance(weight._params, IntCrushInt4Layout.Params):
|
|
backend = self._resolve_ck_backend_int4()
|
|
|
|
if backend == "cuda":
|
|
return self._forward_int4_ck(qdata, weight, bias, x)
|
|
|
|
if backend == "triton":
|
|
_have_w4a8_kernel = (
|
|
(_tr.TRITON_W4A8_GEMM or _tr.TRITON_INT8_GEMM)
|
|
and _tr.TRITON_DYNQUANT
|
|
and _tr.TRITON_INT4_INT8_UNPACK
|
|
)
|
|
if (not self._intcrush_use_pytorch and not self._intcrush_use_w4a16
|
|
and (self._intcrush_rot_need or self._intcrush_is_rpbh)
|
|
and _have_w4a8_kernel):
|
|
return self._forward_int4_w4a8(qdata, weight, bias, x)
|
|
|
|
if not self._intcrush_use_pytorch and _tr.TRITON_INT4_UNPACK:
|
|
return self._forward_int4_w4a16(qdata, weight, bias, x)
|
|
|
|
# pytorch fallback (also catches failed backend attempts)
|
|
return self._forward_int4_pytorch(qdata, weight, bias, x)
|
|
|
|
# ── INT8 forward paths ──
|
|
if isinstance(weight._params, IntCrushInt8Layout.Params):
|
|
backend = self._resolve_ck_backend()
|
|
|
|
if backend == "cuda":
|
|
return self._forward_int8_ck(qdata, weight, bias, x)
|
|
|
|
if backend == "triton":
|
|
return self._forward_int8(qdata, weight, bias, x)
|
|
|
|
# pytorch fallback
|
|
self._intcrush_use_pytorch = True
|
|
return self._forward_int8(qdata, weight, bias, x)
|
|
|
|
# Unknown layout — delegate to standard manual_cast forward.
|
|
return super().forward(x)
|
|
|
|
finally:
|
|
if uncast is not None:
|
|
uncast()
|
|
|
|
class GroupNorm(manual_cast.GroupNorm):
|
|
pass
|
|
|
|
class LayerNorm(manual_cast.LayerNorm):
|
|
pass
|
|
|
|
class Conv2d(manual_cast.Conv2d):
|
|
pass
|
|
|
|
class Conv3d(manual_cast.Conv3d):
|
|
pass
|
|
|
|
class ConvTranspose2d(manual_cast.ConvTranspose2d):
|
|
pass
|
|
|
|
class Embedding(manual_cast.Embedding):
|
|
pass
|
|
|
|
_ops_cache[cache_key] = IntCrushOps
|
|
return IntCrushOps
|
|
|
|
|
|
__all__ = ["make_intcrush_ops"]
|