if (

not _AIMDO_FILE_SLICE_LOAD

or tensor.device.type != "cpu"

or device is None

or device.type != "cuda"

):

return tensor.to(device, non_blocking=True)

size = tensor.numel() * tensor.element_size()

if size == 0:

return tensor.to(device, non_blocking=True)

hostbuf = comfy_aimdo.host_buffer.HostBuffer(size)

host_tensor = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)

host_view = host_tensor[:size].view(dtype=tensor.dtype).view(tensor.shape)

if comfy.memory_management.read_tensor_file_slice_into(tensor, host_view):

out = torch.empty_like(tensor, device=device)

out.copy_(host_view, non_blocking=False)

return out

"""

generator = torch.Generator(device=x.device)

generator.manual_seed(seed)

# Scale to INT8 range — move scale to x's device to handle CPU-stored scales

if isinstance(scale, torch.Tensor):

scale = scale.to(x.device)

x_scaled = x / scale

# Stochastic rounding

x_floor = torch.floor(x_scaled)

fraction = x_scaled - x_floor

del x_scaled # High-precision input no longer needed

# Speed optimization: Create random values directly on the target device

random_vals = torch.rand(x_floor.shape, generator=generator, device=x.device, dtype=x_floor.dtype)

x_rounded = torch.where(random_vals < fraction, x_floor + 1, x_floor)

del random_vals

del fraction

del x_floor

"""Forward with dynamic per-token activation quantization."""

# --- FAST PATH: Triton Fused Kernel ---

if _TRITON_AVAILABLE and _use_triton and x.device.type != "cpu":

return triton_int8_linear(x, weight, weight_scale, bias, compute_dtype)

# --- SLOW PATH: Standard PyTorch ---

# Quantize activations per row (dynamic)

x_8, x_scale = quantize_int8_axiswise(x, dim=-1)

# INT8 Matmul (Outputs Int32)

res = torch._int_mm(x_8, weight.T)

# Dequantize: (res * weight_scale * x_scale)

# Note: Creating intermediate Float tensors here is VRAM heavy

res_scaled = res.float().mul_(weight_scale * x_scale).to(compute_dtype)

if bias is not None:

res_scaled = res_scaled + bias.to(compute_dtype)

"""Forward with dynamic per-token activation quantization and per-row weight quantization.

# --- FAST PATH: Triton Fused Kernel (per-row) ---

if _TRITON_AVAILABLE and _use_triton and x.device.type != "cpu":

return triton_int8_linear_per_row(x, weight, weight_scale, bias, compute_dtype)

# --- SLOW PATH: Standard PyTorch ---

x_8, x_scale = quantize_int8_axiswise(x, dim=-1)

# INT8 Matmul (Outputs Int32)

res = torch._int_mm(x_8, weight.T) # [batch, out_features]

# Dequantize with per-row weight scales

# res[i,j] = sum_k(x_8[i,k] * weight[j,k]) * x_scale[i] * weight_scale[j]

# Broadcasting: res * x_scale * weight_scale.T

res_scaled = res.float().mul_(x_scale).mul_(weight_scale.T).to(compute_dtype)

if bias is not None:

res_scaled = res_scaled + bias.to(compute_dtype)

"""

excluded_names = []

dynamic_quantize = False # Manual toggle for on-the-fly quantization

enable_convrot = False # Toggle for ConvRot Hadamard rotation

use_triton = True # Toggle for Triton fused kernel (mirrors _use_triton)

compute_dtype = None # Optional override for INT8 activation/output compute dtype

_is_prequantized = False # Keep this as a status flag, but don't use for detection

lora_mode = "None" # None/Stochastic bake into INT8 weights; Dynamic applies LoRA at inference

dynamic_lora = False # If True, apply LoRA dynamically at inference; if False, bake into INT8 weights at load time

lora_patches = {} # Map of model_key -> patch list (from load_lora)

lora_strength = 1.0

dynamic_load_device = None # Set by the loader when Aimdo should avoid a full CPU staging copy

skeleton_meta_init = False # Temporary mode for LoRA key-map discovery

_auto_compute_dtype_by_device = {}

@staticmethod

def _default_compute_dtype(x: Tensor) -> torch.dtype:

if x.dtype in (torch.float16, torch.bfloat16):

return x.dtype

return torch.float16

class Linear(manual_cast.Linear):

def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):

if getattr(Int8TensorwiseOps, "skeleton_meta_init", False):

nn.Module.__init__(self)

self.in_features = in_features

self.out_features = out_features

tensor_kwargs = {"device": "meta"}

if dtype is not None:

tensor_kwargs["dtype"] = dtype

self.weight = nn.Parameter(torch.empty((out_features, in_features), **tensor_kwargs), requires_grad=False)

self.bias = nn.Parameter(torch.empty((out_features,), **tensor_kwargs), requires_grad=False) if bias else None

self.weight_comfy_model_dtype = dtype

self.bias_comfy_model_dtype = dtype

# Preserve ComfyUI's Windows/Aimdo lazy-init path. The base

# disable_weight_init.Linear only takes this path for classes

# that do not override _load_from_state_dict; this INT8 class

# does override it, so calling super() would allocate full

# skeleton weights during Pre-LoRA key-map discovery.

elif comfy.model_management.WINDOWS and comfy.memory_management.aimdo_enabled:

nn.Module.__init__(self)

self.in_features = in_features

self.out_features = out_features

self.weight = None

self.bias = None

self.comfy_need_lazy_init_bias = bias

self.weight_comfy_model_dtype = dtype

self.bias_comfy_model_dtype = dtype

else:

super().__init__(in_features, out_features, bias, device, dtype)

self.register_buffer('weight_scale', None)

self._is_quantized = False

self._is_per_row = False # Track quantization granularity

self._use_convrot = False # Track if ConvRot was applied

self._weight_scale_scalar = None # For scalar (non-tensor) scales

self.compute_dtype = None

self.comfy_cast_weights = False

self.lora_patches = [] # List of (down_scaled, up, start, size) set by INT8ModelPatcher

def reset_parameters(self):

return None

@staticmethod

def _normalize_lora_key(key):

if not isinstance(key, str):

return key

for p in ["diffusion_model.", "model.diffusion_model.", "model.", "transformer."]:

if key.startswith(p):

return key[len(p):]

return key

@staticmethod

def _is_bias_key(key):

return isinstance(key, str) and key.endswith(".bias")

@staticmethod

def _format_lora_patches(patches):

formatted = []

for patch in patches or []:

if len(patch) == 4:

v, offset, function, strength = patch

else:

v, offset, function = patch

strength = getattr(Int8TensorwiseOps, "lora_strength", 1.0)

formatted.append((strength, v, 1.0, offset, function))

return formatted

def _apply_int8_lora_patches(self, tensor, key, patches, device):

if not patches or tensor.dtype == torch.int8:

return tensor

temp_dtype = comfy.model_management.lora_compute_dtype(device)

tensor_temp = tensor_to_device_file_slice(tensor, device).to(dtype=temp_dtype)

return comfy.lora.calculate_weight(self._format_lora_patches(patches), tensor_temp, key)

def finalize_pending_int8(self):

pending = getattr(self, "_pending_int8_finalize", None)

if pending is None:

return False

weight_key = pending["weight_key"]

device = pending.get("device")

if device is None:

device = torch.device("cuda") if torch.cuda.is_available() else self.weight.device

weight_tensor = self.weight.detach()

weight_tensor = self._apply_int8_lora_patches(weight_tensor, weight_key, pending.get("lora_patches"), device)

if pending["quantize"]:

if not hasattr(Int8TensorwiseOps, '_logged_otf'):

print(f"INT8 Fast: Quantizing on-the-fly (ConvRot: {pending.get('enable_convrot', False)})")

Int8TensorwiseOps._logged_otf = True

w_gpu = tensor_to_device_file_slice(weight_tensor, device).float()

self._use_convrot = False

if pending.get("enable_convrot", False) and self.in_features % CONVROT_GROUP_SIZE == 0:

try:

from .convrot import build_hadamard, rotate_weight

H = build_hadamard(CONVROT_GROUP_SIZE, device=w_gpu.device, dtype=w_gpu.dtype)

w_gpu = rotate_weight(w_gpu, H, group_size=CONVROT_GROUP_SIZE)

self._use_convrot = True

# Stamp the active groupsize so INT8ModelSave can

# round-trip it deterministically.

self._convrot_groupsize = CONVROT_GROUP_SIZE

except ImportError as e:

logging.warning(f"INT8 Fast: ConvRot enabled but convrot module error: {e}")

q_weight, q_scale = quantize_int8_axiswise(w_gpu, dim=1)

self.weight = nn.Parameter(q_weight.cpu(), requires_grad=False)

self.register_buffer('weight_scale', q_scale.cpu())

self._weight_scale_scalar = None

self._is_quantized = True

self._is_per_row = True

del w_gpu, q_weight, q_scale

else:

self.weight = nn.Parameter(weight_tensor.cpu(), requires_grad=False)

self.weight_comfy_model_dtype = self.weight.dtype

if self.weight_scale is not None:

self.weight_scale_comfy_model_dtype = self.weight_scale.dtype

if self.bias is not None:

self.bias_comfy_model_dtype = self.bias.dtype

delattr(self, "_pending_int8_finalize")

return True

def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):

weight_key = prefix + "weight"

# Utility to normalize keys by stripping common prefixes

def normalize_key(key):

return self._normalize_lora_key(key)

def apply_lora_patches(tensor, key):

if self._is_bias_key(key) or not Int8TensorwiseOps.lora_patches or tensor.dtype == torch.int8:

return tensor

nk = normalize_key(key)

patches = Int8TensorwiseOps.lora_patches.get(nk)

if patches:

# Track applied patches

if not hasattr(Int8TensorwiseOps, 'applied_lora_patches'):

Int8TensorwiseOps.applied_lora_patches = set()

Int8TensorwiseOps.applied_lora_patches.add(nk)

# Print only if multiple sub-patches map to the same layer

# if "weight" in key and len(patches) > 1:

# print(f"INT8 Fast: Baking multiple LoRA parts into {nk} ({len(patches)} sub-patches)")

# ComfyUI dynamically patches during inference using lora_compute_dtype()

# On most modern GPUs, this evaluates to torch.float16.

# We simulate that exact intermediate cast here to achieve a 1:1 binary match.

device = getattr(Int8TensorwiseOps, "dynamic_load_device", None)

if device is None:

device = tensor.device

result_temp = self._apply_int8_lora_patches(tensor, key, patches, device)

return result_temp.to(tensor.dtype)

return tensor

def source_tensor(tensor):

if tensor is not None and getattr(Int8TensorwiseOps, "dynamic_load_device", None) is not None:

return tensor.cpu()

return tensor

scale_key = prefix + "weight_scale"

input_scale_key = prefix + "input_scale"

bias_key = prefix + "bias"

def pop_metadata(sd, p, k):

v = sd.pop(p + k, None)

if v is not None: return v

v = sd.pop("model." + p + k, None)

if v is not None: return v

if p.startswith("model."):

v = sd.pop(p[6:] + k, None)

if v is not None: return v

if p.startswith("diffusion_model."):

v = sd.pop("diffusion_model." + p + k, None)

if v is not None: return v

return None

weight_scale = pop_metadata(state_dict, prefix, "weight_scale")

comfy_quant_tensor = pop_metadata(state_dict, prefix, "comfy_quant")

weight_tensor = state_dict.pop(weight_key, None)

bias_tensor = state_dict.pop(bias_key, None)

# Pop input_scale to clean state_dict, but ignore it

_ = state_dict.pop(input_scale_key, None)

quant_conf_parsed = None

w4a8_handled = False

if comfy_quant_tensor is not None:

try:

import json

quant_conf_parsed = json.loads(bytes(comfy_quant_tensor.tolist()).decode('utf-8'))

if quant_conf_parsed.get("convrot", False):

self._use_convrot = True

Int8TensorwiseOps.enable_convrot = True # Propagate globally for LoRA

if "convrot_groupsize" in quant_conf_parsed:

self._convrot_groupsize = int(quant_conf_parsed["convrot_groupsize"])

Int8TensorwiseOps._global_convrot_groupsize = self._convrot_groupsize

except Exception:

pass

if quant_conf_parsed and quant_conf_parsed.get("format") == "asym_w4a8_int8":

if not _W4A8_AVAILABLE:

raise ImportError("w4a8_int8.py not found next to int8_quant.py")

if weight_tensor is None or weight_tensor.dtype != torch.int8:

logging.warning(f"W4A8: {weight_key} missing packed int8 weight, skipping")

else:

codebook = pop_metadata(state_dict, prefix, "weight_codebook")

s_channel = pop_metadata(state_dict, prefix, "weight_s_channel")

s_rel = pop_metadata(state_dict, prefix, "weight_s_rel")

if codebook is None or s_channel is None or s_rel is None:

logging.warning(f"W4A8: {weight_key} missing codebook/s_channel/s_rel, skipping")

else:

# Lazy load: keep everything packed/compact at load time (cheap,

# matches INT4 ConvRot's pattern). Dequant happens once per device

# on first forward(), on GPU, and is cached from then on.

self._is_w4a8 = True

self._w4a8_group_size = int(quant_conf_parsed.get("group_size", 16))

self._w4a8_codebook = codebook

self._w4a8_s_channel = s_channel

self._w4a8_s_rel = s_rel

self.weight = nn.Parameter(weight_tensor, requires_grad=False)

self._is_quantized = True

self._is_per_row = True

Int8TensorwiseOps._is_prequantized = True

logging.debug(

f"W4A8 (lazy) Loaded: {weight_key} packed {tuple(weight_tensor.shape)} "

f"group_size={self._w4a8_group_size}"

)

Int8TensorwiseOps._w4a8_layer_count = getattr(Int8TensorwiseOps, "_w4a8_layer_count", 0) + 1

w4a8_handled = True

weight_tensor = None # handled -- skip generic weight assignment below

pending_weight_lora_patches = None

if weight_tensor is not None and weight_tensor.dtype != torch.int8:

pending_weight_lora_patches = Int8TensorwiseOps.lora_patches.get(normalize_key(weight_key))

defer_weight_lora = (

getattr(Int8TensorwiseOps, "dynamic_load_device", None) is not None

and pending_weight_lora_patches

)

# Apply LoRA patches to weight and bias once. With Aimdo, large

# weight patches are deferred until KSampler/model load time so

# the loader node stays cheap and VBAR geometry is finalized once.

if weight_tensor is not None and not defer_weight_lora:

weight_tensor = apply_lora_patches(weight_tensor, weight_key)

if bias_tensor is not None:

bias_tensor = apply_lora_patches(bias_tensor, bias_key)

if weight_tensor is not None:

if weight_tensor.dtype == torch.int8 and weight_scale is not None:

# Load Quantized

self._is_quantized = True

self.weight = nn.Parameter(weight_tensor, requires_grad=False)

Int8TensorwiseOps._is_prequantized = True # Found a quantized layer

# Optional explicit hint from saved comfy_quant

per_row_hint = None

if isinstance(quant_conf_parsed, dict) and "per_row" in quant_conf_parsed:

per_row_hint = bool(quant_conf_parsed["per_row"])

if isinstance(weight_scale, torch.Tensor):

if weight_scale.numel() == 1:

# Scalar scale — store as float for speed

self._weight_scale_scalar = weight_scale.float().item()

self.register_buffer('weight_scale', weight_scale.float().reshape(1))

self._weight_scale_scalar = None

elif weight_scale.dim() == 2 and weight_scale.shape[1] == 1:

self.register_buffer('weight_scale', weight_scale.float())

self._weight_scale_scalar = None

self._is_per_row = True if per_row_hint is None else per_row_hint

else:

self.register_buffer('weight_scale', weight_scale.float())

self._weight_scale_scalar = None

self._is_per_row = False if per_row_hint is None else per_row_hint

else:

self.weight_scale = nn.Parameter(

torch.tensor(float(weight_scale), dtype=torch.float32),

requires_grad=False

)

self.weight_scale = None

self._is_per_row = False if per_row_hint is None else per_row_hint

elif weight_tensor.dtype in (torch.float16, torch.bfloat16, torch.float32):

# Load High-Precision

is_excluded = any(ex in prefix for ex in Int8TensorwiseOps.excluded_names)

is_dim1 = self.in_features == 1 or self.out_features == 1 or weight_tensor.ndim == 1

should_quantize = not (is_excluded or is_dim1 or not Int8TensorwiseOps.dynamic_quantize)

defer_finalize = (

getattr(Int8TensorwiseOps, "dynamic_load_device", None) is not None

and (should_quantize or pending_weight_lora_patches)

)

if defer_finalize:

self._is_quantized = False

self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False)

self._pending_int8_finalize = {

"weight_key": weight_key,

"quantize": should_quantize,

"lora_patches": pending_weight_lora_patches,

"device": getattr(Int8TensorwiseOps, "dynamic_load_device", None),

"enable_convrot": getattr(Int8TensorwiseOps, "enable_convrot", False),

}

if pending_weight_lora_patches:

if not hasattr(Int8TensorwiseOps, 'applied_lora_patches'):

Int8TensorwiseOps.applied_lora_patches = set()

Int8TensorwiseOps.applied_lora_patches.add(normalize_key(weight_key))

elif not should_quantize:

self._is_quantized = False

self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False)

else:

# Quantize on the fly

device = getattr(Int8TensorwiseOps, "dynamic_load_device", None)

if device is None:

device = torch.device("cuda") if torch.cuda.is_available() else weight_tensor.device

# Log the first time we quantize in this loader pass

if not hasattr(Int8TensorwiseOps, '_logged_otf'):

print(f"INT8 Fast: Quantizing on-the-fly (ConvRot: {getattr(Int8TensorwiseOps, 'enable_convrot', False)})")

Int8TensorwiseOps._logged_otf = True

# Cast to float32 before rotation and scale computation

w_gpu = weight_tensor.to(device, non_blocking=True).float()

self._use_convrot = False

if getattr(Int8TensorwiseOps, "enable_convrot", False) and self.in_features % CONVROT_GROUP_SIZE == 0:

try:

from .convrot import build_hadamard, rotate_weight

H = build_hadamard(CONVROT_GROUP_SIZE, device=w_gpu.device, dtype=w_gpu.dtype)

w_gpu = rotate_weight(w_gpu, H, group_size=CONVROT_GROUP_SIZE)

self._use_convrot = True

# Stamp the active groupsize on the module so

# INT8ModelSave can serialize it deterministically

# (instead of relying on the loader's default).

self._convrot_groupsize = CONVROT_GROUP_SIZE

except ImportError as e:

logging.warning(f"INT8 Fast: ConvRot enabled but convrot module error: {e}")

q_weight, q_scale = quantize_int8_axiswise(w_gpu, dim=1)

q_weight = q_weight.cpu()

q_scale = q_scale.cpu()

self.weight = nn.Parameter(q_weight, requires_grad=False)

self.register_buffer('weight_scale', q_scale)

self._weight_scale_scalar = None

self._is_quantized = True

self._is_per_row = True

del w_gpu, q_weight, q_scale

else:

self._is_quantized = False

self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False)

elif not w4a8_handled:

missing_keys.append(weight_key)

# Assign bias if it exists (already patched if needed)

if bias_tensor is not None:

self.bias = nn.Parameter(source_tensor(bias_tensor), requires_grad=False)

else:

self.bias = None

# Update archived model dtypes so VBAR geometry uses the correct

# sizes. archive_model_dtypes runs before state_dict loading, so

# weight_comfy_model_dtype is stale (e.g. bfloat16 instead of int8).

# Without this, VBAR allocates 2x the needed memory and the cast

# buffer path misinterprets int8 data as bfloat16.

if self.weight is not None:

self.weight_comfy_model_dtype = self.weight.dtype

if self.weight_scale is not None:

self.weight_scale_comfy_model_dtype = self.weight_scale.dtype

if self.bias is not None:

self.bias_comfy_model_dtype = self.bias.dtype

def _get_weight_scale(self):

return self.weight_scale

def convert_weight(self, _weight, inplace=False):

if not self._is_quantized:

return _weight

return self.weight

def set_weight(self, out_weight, inplace_update=False, seed=0, return_weight=False, **kwargs):

if not self._is_quantized:

new_weight = out_weight.to(self.weight.dtype)

if return_weight:

return new_weight

if inplace_update:

self.weight.data.copy_(new_weight)

else:

self.weight = nn.Parameter(new_weight, requires_grad=False)

return

if out_weight.dtype == torch.int8:

if return_weight:

return out_weight

if inplace_update:

self.weight.data.copy_(out_weight)

else:

self.weight = nn.Parameter(out_weight, requires_grad=False)

return

# Re-quantize if fallback occurred

new_weight = quantize_int8(out_weight, self._get_weight_scale())

if return_weight:

return new_weight

if inplace_update:

self.weight.data.copy_(new_weight)

else:

self.weight = nn.Parameter(new_weight, requires_grad=False)

def set_bias(self, out_bias, inplace_update=False, seed=0, return_weight=False, **kwargs):

if out_bias is None: return None

new_bias = out_bias

if return_weight:

return new_bias

if inplace_update:

if self.bias is not None:

self.bias.data.copy_(new_bias)

else:

self.bias = nn.Parameter(new_bias, requires_grad=False)

def forward(self, x: Tensor) -> Tensor:

"""Fast forward using torch._int_mm for quantized weights."""

# Check if ComfyUI needs to manage weight transfer (VBAR, offloading, LoRA patches, etc.)

# This mirrors the base class check in disable_weight_init.Linear.forward()

need_cast = self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0

if not self._is_quantized:

if need_cast:

weight, bias, offload_stream = cast_bias_weight(self, x, offloadable=True)

out = F.linear(x, weight, bias)

uncast_bias_weight(self, weight, bias, offload_stream)

return out

else:

if x.device != self.weight.device or x.dtype != self.weight.dtype:

weight = self.weight.to(device=x.device, dtype=x.dtype)

bias = self.bias.to(device=x.device, dtype=x.dtype) if self.bias is not None else None

return F.linear(x, weight, bias)

return F.linear(x, self.weight, self.bias)

# INT8 quantized path

if need_cast:

# VBAR / offload / lowvram path

weight, bias, offload_stream = cast_bias_weight(

self, input=None, dtype=torch.int8, device=x.device,

bias_dtype=x.dtype, offloadable=True

)

# Guard: cast_bias_weight may not move to device on ROCm builds

if weight is not None and weight.device != x.device:

weight = weight.to(x.device, non_blocking=True)

if bias is not None and bias.device != x.device:

bias = bias.to(x.device, non_blocking=True)

else:

# Fast path: weights already on GPU, no functions to apply

weight = self.weight

bias = self.bias

offload_stream = None

# Guard: if weight ended up on CPU (low-VRAM offload without need_cast),

# move it to the activation device before hitting Triton

if weight is not None and weight.device != x.device:

weight = weight.to(x.device, non_blocking=True)

if bias is not None and bias.device != x.device:

bias = bias.to(x.device, non_blocking=True)

if getattr(self, "_is_w4a8", False):

codebook = self._w4a8_codebook.to(weight.device, non_blocking=True)

s_channel = self._w4a8_s_channel.to(weight.device, non_blocking=True)

s_rel = self._w4a8_s_rel.to(weight.device, non_blocking=True)

dequant_float = dequant_asym_w4a8_to_float(

weight, codebook, s_channel, s_rel, self._w4a8_group_size

)

weight, w_scale = quantize_int8_axiswise(dequant_float, dim=1)

del dequant_float

else:

w_scale = self._get_weight_scale()

if isinstance(w_scale, torch.Tensor) and w_scale.device != x.device:

w_scale = w_scale.to(x.device, non_blocking=True)

compute_dtype = Int8TensorwiseOps.compute_dtype

if compute_dtype is None:

compute_dtype = Int8TensorwiseOps._default_compute_dtype(x)

x_shape = x.shape

x_2d = x.reshape(-1, x_shape[-1])

if x_2d.dtype != compute_dtype:

x_2d = x_2d.to(compute_dtype)

if getattr(self, "_use_convrot", False):

from .convrot import build_hadamard, rotate_activation

group_size = getattr(self, "_convrot_groupsize", CONVROT_GROUP_SIZE)

H = build_hadamard(group_size, device=x.device, dtype=x_2d.dtype)

x_2d = rotate_activation(x_2d, H, group_size=group_size)

# Sync the loader toggle to the module-level flag read by the forward fns

import sys as _sys

_mod = _sys.modules[__name__]

_mod._use_triton = Int8TensorwiseOps.use_triton

if self._is_per_row:

y = int8_forward_dynamic_per_row(x_2d, weight, w_scale, bias, compute_dtype)

else:

y = int8_forward_dynamic(x_2d, weight, w_scale, bias, compute_dtype)

# Dynamic LoRA Path — handles split QKV via per-patch offsets

for lora_down, lora_up, lora_start, lora_size in self.lora_patches:

lD = lora_down.to(x.device, non_blocking=True)

lU = lora_up.to(x.device, non_blocking=True)

lora_x = F.linear(x_2d.to(lD.dtype), lD)

lora_y = F.linear(lora_x, lU) # [batch, slice_size or full_out]

if lora_start is not None:

y[:, lora_start:lora_start + lora_size] = (

y[:, lora_start:lora_start + lora_size] + lora_y.to(y.dtype)

)

else:

y = y + lora_y.to(y.dtype)

if need_cast:

uncast_bias_weight(self, weight, bias, offload_stream)

return y.reshape(*x_shape[:-1], y.shape[-1])

# Pass-through for other layers

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

@classmethod

def conv_nd(cls, dims, *args, **kwargs):

if dims == 2: return cls.Conv2d(*args, **kwargs)

elif dims == 3: return cls.Conv3d(*args, **kwargs)

is_lowvram_patch = True

def __init__(self, key, patches, module, lora_mode):

self.key = key

self.patches = patches

self.module = module

self.lora_mode = lora_mode

self.prepared_patches = None

def memory_required(self):

if not _use_new_prefetch:

return 0

counter = [0]

for patch in self.patches[self.key]:

comfy.lora.prefetch_prepared_value(patch[1], counter, None, None, False)

return counter[0]

def prepare(self, *args, **kwargs):

if _use_new_prefetch:

# 0.22.0+ signature: prepare(self, destination, stream, copy=True, commit=True)

destination = args[0] if len(args) > 0 else kwargs.get("destination")

stream = args[1] if len(args) > 1 else kwargs.get("stream")

copy = args[2] if len(args) > 2 else kwargs.get("copy", True)

commit = args[3] if len(args) > 3 else kwargs.get("commit", True)

counter = [0]

prepared_patches = [

(patch[0], comfy.lora.prefetch_prepared_value(patch[1], counter, destination, stream, copy), patch[2], patch[3], patch[4])

for patch in self.patches[self.key]

]

if commit:

self.prepared_patches = prepared_patches

return prepared_patches

else:

# 0.21.1- signature: prepare(self, allocate_buffer, stream)

allocate_buffer = args[0] if len(args) > 0 else kwargs.get("allocate_buffer")

stream = args[1] if len(args) > 1 else kwargs.get("stream")

self.prepared_patches = [

(patch[0], comfy.lora.prefetch_prepared_value(patch[1], allocate_buffer, stream), patch[2], patch[3], patch[4])

for patch in self.patches[self.key]

]

return self.prepared_patches

def clear_prepared(self):

self.prepared_patches = None

def __call__(self, weight):

patches = self.prepared_patches if self.prepared_patches is not None else self.patches[self.key]

scale = self.module._get_weight_scale()

if isinstance(scale, torch.Tensor):

scale = scale.to(weight.device)

weight_float = dequantize(weight, scale)

use_convrot = getattr(self.module, "_use_convrot", False)

if use_convrot:

group_size = getattr(self.module, "_convrot_groupsize", CONVROT_GROUP_SIZE)

try:

from .convrot import build_hadamard, rotate_weight

H = build_hadamard(group_size, device=weight.device, dtype=weight_float.dtype)

weight_float = rotate_weight(weight_float, H, group_size=group_size)

except ImportError:

use_convrot = False

patched_weight_float = comfy.lora.calculate_weight(

patches,

weight_float,

self.key,

intermediate_dtype=weight_float.dtype,

)

if use_convrot:

patched_weight_float = rotate_weight(patched_weight_float, H, group_size=group_size)

if self.lora_mode == "Stochastic":

return stochastic_round_int8_delta(

patched_weight_float,

scale,

seed=comfy.utils.string_to_seed(self.key),

)

"""

def finalize_pending_int8(self):

finalized = 0

for module in self.model.modules():

finalize = getattr(module, "finalize_pending_int8", None)

if finalize is not None and finalize():

finalized += 1

if finalized > 0:

self.size = 0

#logging.info(f"INT8 Fast: Finalized {finalized} deferred INT8 layer(s) at model load time.")

if torch.cuda.is_available():

torch.cuda.empty_cache()

def patch_weight_to_device(self, key, device_to=None, inplace_update=False, return_weight=False, force_cast=False):

if key not in self.patches and not force_cast:

return super().patch_weight_to_device(key, device_to, inplace_update, return_weight, force_cast)

# Check if this is one of our INT8 modules

module_path = key.rsplit('.', 1)[0]

try:

module = comfy.utils.get_attr(self.model, module_path)

except AttributeError:

module = None

is_int8_module = hasattr(module, "_is_quantized") and module._is_quantized

patches = self.patches.get(key, [])

if is_int8_module and Int8TensorwiseOps.Linear._is_bias_key(key):

return comfy.utils.get_attr(self.model, key) if return_weight else None

if is_int8_module:

if not Int8TensorwiseOps.dynamic_lora:

# --- BAKE-IN LORA PATH (Dequant → Patch → Quant) ---

# Works with the native ComfyUI LoRA Loader (and also INT8LoraLoader).

# All patches are applied in float space via ComfyUI's standard mechanism,

# then the result is re-quantized back to INT8.

# Identify current weight in the model

current_weight = comfy.utils.get_attr(self.model, key)

scale = module._get_weight_scale()

if device_to is None:

device_to = current_weight.device

# ALWAYS use the weight from backup as the source if it exists to prevent additive stacking.

# If it doesn't exist, this is the first patch, so create it from the current model weight.

if key not in self.backup:

import collections

BackupEntry = collections.namedtuple('Dimension', ['weight', 'inplace_update'])

self.backup[key] = BackupEntry(

weight=current_weight.to(device=self.offload_device, copy=inplace_update),

inplace_update=inplace_update,

)

source_weight = current_weight

else:

# Use existing backup as source

source_weight = self.backup[key].weight

# 1. Dequantize to float (move scale to device_to since it lives on CPU)

if isinstance(scale, torch.Tensor):

scale = scale.to(device_to)

weight_float = dequantize(source_weight.to(device_to), scale)

# 2. Handle ConvRot: de-rotate into weight space before patching

use_convrot = getattr(module, "_use_convrot", False)

if use_convrot:

group_size = getattr(module, "_convrot_groupsize", CONVROT_GROUP_SIZE)

try:

from .convrot import build_hadamard, rotate_weight

H = build_hadamard(group_size, device=device_to, dtype=weight_float.dtype)

weight_float = rotate_weight(weight_float, H, group_size=group_size)

except ImportError:

pass

# 3. Patch in float space using ComfyUI's standard mechanism.

# calculate_weight handles LoRA, LoHA, LoKR, DoRA, etc.

patches_list = self.patches.get(key, [])

patched_weight_float = comfy.lora.calculate_weight(patches_list, weight_float, key)

# 4. Handle ConvRot: re-rotate

if use_convrot:

patched_weight_float = rotate_weight(patched_weight_float, H, group_size=group_size)

# 5. Re-quantize back to INT8 using the original scale

if getattr(Int8TensorwiseOps, "lora_mode", "None") == "Stochastic":

patched_weight_int8 = stochastic_round_int8_delta(patched_weight_float, scale)

else:

patched_weight_int8 = quantize_int8(patched_weight_float, scale)

# 6. Move back to original device and store

patched_weight_int8 = patched_weight_int8.to(current_weight.device)

if return_weight:

return patched_weight_int8

if inplace_update:

current_weight.data.copy_(patched_weight_int8)