373 lines
15 KiB
Python
373 lines
15 KiB
Python
"""QuantizedTensor layout classes for INT-Crush INT4/INT8 formats.
|
|
|
|
These bridge INT-Crush's packed weight format with ComfyUI's QuantizedTensor
|
|
dispatch system, enabling VBAR demand-paging, async offload, and LoRA
|
|
compatibility through ComfyUI's standard cast_bias_weight pipeline.
|
|
|
|
Each layout class defines:
|
|
- Params dataclass: metadata stored alongside packed weights
|
|
- state_dict_tensors(): serialization layout for safetensors
|
|
- dequantize(): full dequantization (used for LoRA patching and fallback)
|
|
- quantize(): float → packed (used by QuantizedTensor.from_float() for LoRA requantization)
|
|
- supports_fast_matmul(): check for fused Triton kernel availability
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from dataclasses import dataclass
|
|
|
|
from ._quant_utils import unpack_int4, INT4_SCALE_DIVISOR, INT4_MIN, INT4_MAX, INT4_TWOS_COMPLEMENT_BIAS
|
|
from ._quant_utils import INT8_MIN, INT8_MAX, INT8_SCALE_DIVISOR
|
|
|
|
__all__ = ["IntCrushInt4Layout", "IntCrushInt8Layout", "OrbitQuantLayout", "register_intcrush_layouts"]
|
|
|
|
# BaseLayoutParams is required by QuantizedTensor.__tensor_flatten__ which
|
|
# calls dataclasses.fields(self._params) to discover tensor sub-fields.
|
|
try:
|
|
from comfy_kitchen.tensor.base import BaseLayoutParams
|
|
except ImportError:
|
|
# Minimal fallback when comfy_kitchen is unavailable (e.g. unit tests).
|
|
# The real BaseLayoutParams is a @dataclass(frozen=True) with three fields:
|
|
# scale: torch.Tensor
|
|
# orig_dtype: torch.dtype
|
|
# orig_shape: tuple[int, ...]
|
|
@dataclass(frozen=True)
|
|
class BaseLayoutParams:
|
|
scale: torch.Tensor
|
|
orig_dtype: torch.dtype
|
|
orig_shape: tuple
|
|
|
|
def __post_init__(self):
|
|
pass
|
|
|
|
|
|
class IntCrushInt4Layout:
|
|
"""Layout class for INT-Crush packed INT4 (uint8, 2 values per byte).
|
|
|
|
Stores weights as packed uint8 with per-row FP16 scales. Optionally
|
|
supports zero-point (asymmetric quantization) and PermuQuant channel
|
|
permutation indices.
|
|
"""
|
|
|
|
@dataclass(frozen=True)
|
|
class Params(BaseLayoutParams):
|
|
rot_need: bool = True
|
|
rot_size: int = 256
|
|
perm: torch.Tensor | None = None
|
|
zp: torch.Tensor | None = None
|
|
|
|
def _tensor_fields(self) -> list[str]:
|
|
fields = ["scale"]
|
|
if self.perm is not None:
|
|
fields.append("perm")
|
|
if self.zp is not None:
|
|
fields.append("zp")
|
|
return fields
|
|
|
|
def _validate_tensor_fields(self):
|
|
if isinstance(self.scale, torch.Tensor):
|
|
object.__setattr__(self, "scale", self.scale.to(dtype=torch.float32, non_blocking=True))
|
|
|
|
@classmethod
|
|
def state_dict_tensors(cls, qdata, params):
|
|
tensors = {"": qdata, "weight_scale": params.scale}
|
|
if params.perm is not None:
|
|
tensors["weight_perm"] = params.perm
|
|
if params.zp is not None:
|
|
tensors["weight_zp"] = params.zp
|
|
return tensors
|
|
|
|
@classmethod
|
|
def dequantize(cls, qdata, params):
|
|
"""Full dequantization to float — used for LoRA patching and PyTorch fallback."""
|
|
in_features = params.orig_shape[1]
|
|
unpacked = unpack_int4(qdata, in_features)
|
|
if params.zp is not None:
|
|
return (unpacked.float() - params.zp.reshape(-1, 1).float()) * params.scale.reshape(-1, 1).float()
|
|
return unpacked.float() * params.scale.reshape(-1, 1).float()
|
|
|
|
@classmethod
|
|
def supports_fast_matmul(cls):
|
|
# comfy-kitchen CUDA (cuBLASLt / CUTLASS / GEMV — sm75+)
|
|
try:
|
|
from comfy_kitchen.backends.cuda import _EXT_AVAILABLE
|
|
if _EXT_AVAILABLE:
|
|
return True
|
|
except ImportError:
|
|
pass
|
|
# Triton fallback
|
|
try:
|
|
from .kernels.triton_int8_gemm import fused_int8_gemm_dequant # noqa: F401 -- probe-by-import
|
|
from .kernels.triton_quantize import dynamic_quantize_activation # noqa: F401
|
|
from .kernels.triton_int4_to_int8_unpack import unpack_int4_to_int8 # noqa: F401
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
@classmethod
|
|
def quantize(cls, tensor, scale=None, stochastic_rounding=0, inplace_ops=False):
|
|
"""Quantize a float tensor to packed INT4 with per-row scales.
|
|
|
|
Used by QuantizedTensor.from_float() during LoRA requantization.
|
|
The input tensor must already be in rotated+permuted space.
|
|
"""
|
|
max_vals = tensor.abs().amax(dim=1, keepdim=True)
|
|
if scale is None or (isinstance(scale, str) and scale == "recalculate"):
|
|
scales = (max_vals / INT4_SCALE_DIVISOR).clamp(min=1e-8).to(torch.float16)
|
|
elif isinstance(scale, torch.Tensor):
|
|
scales = scale.to(torch.float16)
|
|
else:
|
|
scales = torch.tensor([scale], dtype=torch.float16, device=tensor.device)
|
|
|
|
w_scaled = tensor / scales.to(tensor.dtype)
|
|
int_rounded = w_scaled.round().clamp(INT4_MIN, INT4_MAX).to(torch.int8)
|
|
# Pack two INT4 values per byte (low nibble = even index).
|
|
K = int_rounded.shape[-1]
|
|
if K % 2 != 0:
|
|
pad = torch.zeros(*int_rounded.shape[:-1], 1, dtype=torch.int8, device=int_rounded.device)
|
|
int_rounded = torch.cat([int_rounded, pad], dim=-1)
|
|
q_i8 = torch.where(int_rounded < 0, INT4_TWOS_COMPLEMENT_BIAS + int_rounded, int_rounded).to(torch.uint8)
|
|
packed = q_i8[..., 0::2] | (q_i8[..., 1::2] << 4)
|
|
|
|
orig_dtype = tensor.dtype
|
|
orig_shape = tuple(tensor.shape)
|
|
params = cls.Params(
|
|
scale=scales.reshape(-1), orig_dtype=orig_dtype, orig_shape=orig_shape,
|
|
)
|
|
return packed, params
|
|
|
|
@classmethod
|
|
def requantize_kwargs(cls, qtensor):
|
|
# No options are preserved. VBAR's requantize_from_float uses
|
|
# scale="recalculate", and rot/perm live on the module, not the params.
|
|
return {}
|
|
|
|
|
|
class IntCrushInt8Layout:
|
|
"""Layout class for INT-Crush INT8 (plain int8) with per-row FP16 scales."""
|
|
|
|
@dataclass(frozen=True)
|
|
class Params(BaseLayoutParams):
|
|
rot_need: bool = True
|
|
rot_size: int = 256
|
|
perm: torch.Tensor | None = None
|
|
zp: torch.Tensor | None = None
|
|
|
|
def _tensor_fields(self) -> list[str]:
|
|
fields = ["scale"]
|
|
if self.perm is not None:
|
|
fields.append("perm")
|
|
if self.zp is not None:
|
|
fields.append("zp")
|
|
return fields
|
|
|
|
def _validate_tensor_fields(self):
|
|
if isinstance(self.scale, torch.Tensor):
|
|
object.__setattr__(self, "scale", self.scale.to(dtype=torch.float32, non_blocking=True))
|
|
|
|
@classmethod
|
|
def state_dict_tensors(cls, qdata, params):
|
|
tensors = {"": qdata, "weight_scale": params.scale}
|
|
if params.perm is not None:
|
|
tensors["weight_perm"] = params.perm
|
|
if params.zp is not None:
|
|
tensors["weight_zp"] = params.zp
|
|
return tensors
|
|
|
|
@classmethod
|
|
def dequantize(cls, qdata, params):
|
|
"""Full dequantization to float — used for LoRA patching and PyTorch fallback."""
|
|
if params.zp is not None:
|
|
return (qdata.float() - params.zp.reshape(-1, 1).float()) * params.scale.reshape(-1, 1).float()
|
|
return qdata.float() * params.scale.reshape(-1, 1).float()
|
|
|
|
@classmethod
|
|
def supports_fast_matmul(cls):
|
|
# comfy-kitchen CUDA (cuBLASLt / CUTLASS / GEMV — sm75+)
|
|
try:
|
|
from comfy_kitchen.backends.cuda import _EXT_AVAILABLE
|
|
if _EXT_AVAILABLE:
|
|
return True
|
|
except ImportError:
|
|
pass
|
|
# Triton fallback
|
|
try:
|
|
from .kernels.triton_int8_gemm import fused_int8_gemm_dequant # noqa: F401 -- probe-by-import
|
|
from .kernels.triton_quantize import dynamic_quantize_activation # noqa: F401
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
@classmethod
|
|
def quantize(cls, tensor, scale=None, stochastic_rounding=0, inplace_ops=False):
|
|
"""Quantize a float tensor to INT8 with per-row scales.
|
|
|
|
Used by QuantizedTensor.from_float() during LoRA requantization.
|
|
The input tensor must already be in rotated space.
|
|
"""
|
|
max_vals = tensor.abs().amax(dim=1, keepdim=True)
|
|
if scale is None or (isinstance(scale, str) and scale == "recalculate"):
|
|
scales = (max_vals / INT8_SCALE_DIVISOR).clamp(min=1e-8).to(torch.float32)
|
|
elif isinstance(scale, torch.Tensor):
|
|
scales = scale.to(torch.float32)
|
|
else:
|
|
scales = torch.tensor([scale], dtype=torch.float32, device=tensor.device)
|
|
|
|
w_scaled = tensor / scales.to(tensor.dtype)
|
|
int_rounded = w_scaled.round().clamp(INT8_MIN, INT8_MAX).to(torch.int8)
|
|
|
|
orig_dtype = tensor.dtype
|
|
orig_shape = tuple(tensor.shape)
|
|
params = cls.Params(
|
|
scale=scales.reshape(-1), orig_dtype=orig_dtype, orig_shape=orig_shape,
|
|
)
|
|
return int_rounded, params
|
|
|
|
@classmethod
|
|
def requantize_kwargs(cls, qtensor):
|
|
# No options are preserved. VBAR's requantize_from_float uses
|
|
# scale="recalculate", and rot/perm live on the module, not the params.
|
|
return {}
|
|
|
|
|
|
class OrbitQuantLayout:
|
|
"""Layout class for OrbitQuant codebook-based quantization.
|
|
|
|
Stores packed codebook indices (uint8 for INT4, int8 for INT8) with
|
|
per-row L2 norms (bfloat16). AdaLN layers use per-group RTN scales
|
|
(float16) instead.
|
|
|
|
The codebook is stored in Params.codebook so that dequantize() can
|
|
access it for LoRA patching, VBAR demand-paging, and PyTorch fallback.
|
|
"""
|
|
|
|
@dataclass(frozen=True)
|
|
class Params(BaseLayoutParams):
|
|
quant_scheme: str = "orbitquant_codebook" # or "adaln_rtn"
|
|
d: int = 0 # Original input dimension
|
|
d_padded: int = 0 # Padded input dimension
|
|
bits: int = 4 # Weight bit-width
|
|
codebook: torch.Tensor | None = None # Lloyd-Max centroids [2^b]
|
|
|
|
def _tensor_fields(self) -> list[str]:
|
|
fields = ["scale"]
|
|
if self.codebook is not None:
|
|
fields.append("codebook")
|
|
return fields
|
|
|
|
def _validate_tensor_fields(self):
|
|
if isinstance(self.scale, torch.Tensor):
|
|
object.__setattr__(self, "scale",
|
|
self.scale.to(dtype=torch.float32, non_blocking=True))
|
|
# codebook stays float32 (already correct dtype from converter)
|
|
|
|
@classmethod
|
|
def state_dict_tensors(cls, qdata, params):
|
|
tensors = {"": qdata, "weight_scale": params.scale}
|
|
# codebook is stored globally, not per-layer — not in state_dict_tensors
|
|
return tensors
|
|
|
|
@classmethod
|
|
def dequantize(cls, qdata, params):
|
|
"""Full dequantization to float — used for LoRA patching, VBAR,
|
|
and PyTorch fallback.
|
|
|
|
Produces weights in the ROTATED basis (W @ Π^T). The forward path
|
|
must also rotate activations for cancellation. For LoRA patching,
|
|
this is fine — LoRA is applied on raw activations, and the output
|
|
is in the original (unrotated) space because the rotation cancels.
|
|
|
|
Uses the Triton codebook dequant (fp32 output) when available for
|
|
speed; falls back to PyTorch otherwise.
|
|
"""
|
|
from ._orbitquant import codebook_dequant, adaln_rtn_dequant
|
|
|
|
if params.quant_scheme == "adaln_rtn":
|
|
return adaln_rtn_dequant(qdata, params.scale, params.d, params.bits)
|
|
|
|
if params.codebook is None:
|
|
raise RuntimeError(
|
|
"OrbitQuant codebook not available in Params. "
|
|
"This should not happen — the loader stores it."
|
|
)
|
|
# Try Triton fp32 output for exact LoRA parity.
|
|
try:
|
|
from . import _triton_runtime as _tr
|
|
if _tr.TRITON_CODEBOOK_DEQUANT and qdata.is_cuda:
|
|
out = torch.empty(
|
|
qdata.shape[0], params.d_padded,
|
|
dtype=torch.float32, device=qdata.device,
|
|
)
|
|
return _tr.codebook_dequant_triton(
|
|
qdata, params.scale, params.codebook,
|
|
params.d_padded, params.bits, out=out,
|
|
)
|
|
except Exception:
|
|
pass
|
|
return codebook_dequant(
|
|
qdata, params.scale, params.codebook,
|
|
params.d, params.d_padded, params.bits,
|
|
)
|
|
|
|
@classmethod
|
|
def supports_fast_matmul(cls):
|
|
# RPBH rotation is the recurring per-step cost; fast matmul is
|
|
# meaningful once the RPBH kernel is loaded.
|
|
try:
|
|
from . import _triton_runtime as _tr
|
|
return bool(_tr.TRITON_RPBH_ROTATE)
|
|
except Exception:
|
|
return False
|
|
|
|
@classmethod
|
|
def quantize(cls, tensor, scale=None, stochastic_rounding=0, inplace_ops=False):
|
|
"""Quantize a float tensor to OrbitQuant format.
|
|
|
|
Used by QuantizedTensor.from_float() during LoRA requantization.
|
|
NOT YET IMPLEMENTED — requires RPBH rotation + codebook quantization.
|
|
For now, the caller (set_weight) handles this by storing as plain float.
|
|
"""
|
|
raise NotImplementedError(
|
|
"OrbitQuant requantization not yet implemented. "
|
|
"The weight will be stored as plain float and the forward path "
|
|
"will use _forward_predequant_oq."
|
|
)
|
|
|
|
|
|
def register_intcrush_layouts():
|
|
"""Register INT-Crush INT4/INT8 formats with ComfyUI's quantization system.
|
|
|
|
Idempotent and safe to call even when comfy_kitchen is unavailable (no-op).
|
|
"""
|
|
import logging
|
|
_log = logging.getLogger(__name__)
|
|
|
|
try:
|
|
from comfy.quant_ops import register_layout_class, QUANT_ALGOS
|
|
except ImportError:
|
|
_log.warning("INT-Crush: comfy.quant_ops not available; layout registration skipped")
|
|
return
|
|
|
|
try:
|
|
register_layout_class("IntCrushInt4Layout", IntCrushInt4Layout)
|
|
register_layout_class("IntCrushInt8Layout", IntCrushInt8Layout)
|
|
register_layout_class("OrbitQuantLayout", OrbitQuantLayout)
|
|
except Exception:
|
|
_log.warning("INT-Crush: layout class registration failed", exc_info=True)
|
|
|
|
QUANT_ALGOS.setdefault("int4_crush", {
|
|
"storage_t": torch.uint8,
|
|
"parameters": {"weight_scale", "weight_zp", "weight_perm"},
|
|
"comfy_tensor_layout": "IntCrushInt4Layout",
|
|
})
|
|
QUANT_ALGOS.setdefault("int8_crush", {
|
|
"storage_t": torch.int8,
|
|
"parameters": {"weight_scale", "weight_perm", "weight_zp"},
|
|
"comfy_tensor_layout": "IntCrushInt8Layout",
|
|
})
|
|
QUANT_ALGOS.setdefault("orbitquant", {
|
|
"storage_t": torch.uint8, # INT4 default; overridden for INT8
|
|
"parameters": {"weight_scale"},
|
|
"comfy_tensor_layout": "OrbitQuantLayout",
|
|
})
|