Hugging Face's logo

Hayabala
/
SenseNova-6.8-Code2Video-EXL3

Image-Text-to-Text
Transformers
Safetensors
English
Chinese
qwen3_5
code-video
video-generation
agent
claude-code
qwen3.8
exl3
exllamav3
6-bit
quantization
conversational
Model card Files
xet
 Community
SenseNova-6.8-Code2Video-EXL3 / scripts /quantize.py
Hayabala's picture
Hayabala
EXL3 6.0bpw quant of SenseNova-6.8-Code2Video (no MTP layer)
a6355f3 verified
Raw History Blame Contribute Delete
90.2 kB
import torch
import torch.nn.functional as F
import math
import os
from ....ext import exllamav3_ext as ext
from ....util.progress import ProgressBar
from ....util.memory import free_mem, list_gpu_tensors
from ....util.hadamard import get_hadamard_dt
from .hessian_packed import unpack
from ....util import cuda_sync_active
from ....util.tensor import save_tensor_image
from functools import lru_cache
import threading
# Constant
had_k, had_n = 128, 128
codebook_scale = 1.24371088
codebook_mcg_mult = 0xCBAC1FED
codebook_mul1_mult = 0x83DCD12D
def quant_codebook(quant_args: dict) -> tuple[bool, bool]:
"""(mcg, mul1) codebook selection of quant_args: {"mcg": True} or {"mul1": True} selects that codebook, neither
(or False) the default 3INST. Encoder, packing and the stored marker all read the selection through here, so the
trellis and the marker naming its codebook cannot disagree"""
mcg, mul1 = bool(quant_args.get("mcg")), bool(quant_args.get("mul1"))
assert not (mcg and mul1), "quant_args selects both the mcg and mul1 codebooks"
return mcg, mul1
def codebook_markers(quant_args: dict) -> dict[str, torch.Tensor]:
"""The marker tensor naming the selected codebook, as stored next to the trellis ({} for 3INST). Safetensors
has no uint32, and since the multipliers are locked the stored values are never read: only the presence of
the marker tells inference which codebook a tensor uses"""
mcg, mul1 = quant_codebook(quant_args)
if mcg:
return {"mcg": torch.tensor(codebook_mcg_mult, dtype = torch.uint32).view(torch.int)}
if mul1:
return {"mul1": torch.tensor(codebook_mul1_mult, dtype = torch.uint32).view(torch.int)}
return {}
@lru_cache
def tensor_core_perm(device):
"""
Return the 16x16 tile permutation expected by the EXL3 tensor-core quantization kernels.
The cached index maps row-major tile elements into the lane/interleave order used by the CUDA encoder.
"""
perm_a = [0] * 256
for t in range(32):
r0 = (t % 4) * 2
r1 = r0 + 1
r2 = r0 + 8
r3 = r0 + 9
c0 = t // 4
c1 = c0 + 8
perm_a[t * 8 + 0] = r0 * 16 + c0
perm_a[t * 8 + 1] = r1 * 16 + c0
perm_a[t * 8 + 2] = r2 * 16 + c0
perm_a[t * 8 + 3] = r3 * 16 + c0
perm_a[t * 8 + 4] = r0 * 16 + c1
perm_a[t * 8 + 5] = r1 * 16 + c1
perm_a[t * 8 + 6] = r2 * 16 + c1
perm_a[t * 8 + 7] = r3 * 16 + c1
return torch.tensor(perm_a, dtype = torch.int, device = device)
@lru_cache
def tensor_core_perm_i(device):
return torch.argsort(tensor_core_perm(device))
@lru_cache
def get_temp_buffers(device, K: int, tile_len: int = 256, cb: int = 0):
index = torch.device(device).index if not isinstance(device, int) else device
if index is None:
index = torch.cuda.current_device()
# The extension picks the kernel per architecture / K / codebook and reports the wave it can
# keep resident (blocks per SM x SMs); the dense specializations keep costs in shared memory
# and store one branch bit (K = 1) or byte (K >= 2) per state
optimized, wave = ext.quantize_tiles_scratch(index, K, cb == 1, cb == 2, tile_len)
if optimized:
history_cols = (65536 >> K) // (8 if K == 1 else 1)
temp_costs = torch.empty((1,), dtype = torch.half, device = device)
temp_edges = torch.empty((wave, tile_len, history_cols), dtype = torch.uint8, device = device)
return temp_costs, temp_edges
# The kernel runs one block per tile and caps each wave at min(temp_costs.size(0), 2 * SMs). At K >= 4 the
# temp buffers are cheap enough to size for full occupancy on large GPUs (+17% throughput on GB202 at big
# batches); at lower K, temp_edges is multiple GB already and 256 stays the cap
max_batch_size = 256
if K >= 4:
mp_count = torch.cuda.get_device_properties(device).multi_processor_count
# The kernel may run up to 1024 / threads blocks per SM (see quantize_tiles_kernel.cuh)
blocks_per_sm = {7: 4, 8: 8}.get(K, 2)
max_batch_size = max(256, blocks_per_sm * mp_count)
temp_costs = torch.zeros((max_batch_size, 2, 65536 >> K), dtype = torch.half, device = device)
temp_edges = torch.zeros((max_batch_size, tile_len, 65536 >> K), dtype = torch.short, device = device)
return temp_costs, temp_edges
def frac_k(K):
"""(KA, MASK) for a half-integer bitrate K = KA + 0.5: alternating KA- and (KA + 1)-bit trellis
steps (MASK 0xAAAA, period 16). None for integer K. Half steps keep every tile a whole number of
32-bit words (streamable with int4 loads); the kernels accept any period-16 pattern, but only the
alternating one is instantiated."""
if isinstance(K, int) or float(K).is_integer():
return None
ka = int(K)
assert abs(ka + 0.5 - K) < 1e-9, f"fractional bitrate {K} must be a half-integer"
return ka, 0xAAAA
def trellis_words(K) -> int:
"""uint16 words per 256-weight tile at bitrate K (16 * K, an integer for half-integer K)"""
w = 16 * K
assert float(w).is_integer(), f"bitrate {K}: tile is not a whole number of uint16 words"
return int(w)
@lru_cache
def get_temp_buffers_frac(device, ka: int):
edges = 65536 >> ka
temp_costs = torch.zeros((128, 2, edges), dtype = torch.half, device = device)
temp_edges = torch.zeros((128, 256, edges), dtype = torch.short, device = device)
return temp_costs, temp_edges
def quantize_tiles_frac(tiles, quant_args: dict):
"""Half-integer bitrate (KA + 0.5): alternating pattern of KA / KA+1-bit steps. mul1 codebook only"""
tiles = tiles.contiguous()
assert tiles.shape[1] == 256 and tiles.dtype == torch.float
assert quant_codebook(quant_args)[1], "fractional bitrates require the mul1 codebook"
ka, mask = frac_k(quant_args["K"])
quantized_tiles = torch.zeros_like(tiles)
quantized_idx = torch.zeros_like(tiles, dtype = torch.short)
temp_costs, temp_edges = get_temp_buffers_frac(tiles.device, ka)
ext.quantize_tiles_frac(tiles, quantized_tiles, quantized_idx, temp_costs, temp_edges, ka, mask)
return quantized_tiles, quantized_idx
def quantize_tiles(tiles, quant_args: dict):
"""
Quantize a batch of 16x16 tiles on the current device.
tiles is shaped (num_tiles, 256) in the kernel's expected element order. The CUDA extension returns both the
reconstructed float tile values and the short encoded indices used later for packing. Length-160 rows
(n-gram embedding vectors, mul1 codebook only) are accepted too and quantized as single tail-biting rings.
"""
if frac_k(quant_args["K"]) is not None:
return quantize_tiles_frac(tiles, quant_args)
tiles = tiles.contiguous()
assert tiles.shape[1] in (256, 160)
assert tiles.dtype == torch.float
K = quant_args["K"]
mcg, mul1 = quant_codebook(quant_args)
cb = 2 if mul1 else 1 if mcg else 0
quantized_tiles = torch.zeros_like(tiles)
quantized_idx = torch.zeros_like(tiles, dtype = torch.short)
temp_costs, temp_edges = get_temp_buffers(tiles.device, K, tiles.shape[1], cb)
ext.quantize_tiles(
tiles,
quantized_tiles,
quantized_idx,
temp_costs,
temp_edges,
K,
mcg,
mul1,
)
return quantized_tiles, quantized_idx
@lru_cache
def get_quant_stream(device):
return torch.cuda.Stream(device = device)
@lru_cache
def arch_prior_speed(device: int) -> float:
"""Initial relative quant-throughput guess by architecture, used to seed device splits
before any measurements exist. From measured conversion throughput: Ampere lands at about
half of Ada, which is about 80% of Blackwell (Hopper assumed equal to Blackwell); anything
older is pessimistically assumed half of Ampere."""
if torch.version.hip:
return 1.0 # no measurements (HIP's major is the gfx generation)
major, minor = torch.cuda.get_device_capability(device)
if major >= 9: return 5.0 # Hopper, Blackwell
if major == 8: return 4.0 if minor == 9 else 2.0 # Ada / Ampere
return 1.0
class AutoSplit:
"""Learned per-device speed ratios for heterogeneous multi-GPU conversion splits. Each
workload kind (parallel-quant threads, calibration row workers, tile fan-out slices) is
tracked separately, since their per-unit costs differ; within a kind, only the ratios
between devices matter. Speeds are EMA-damped so one noisy module (JIT warmup, thermal
excursions) doesn't swing the split. Until every active device has a measurement for the
kind, the split falls back to the per-architecture prior."""
def __init__(self):
self.lock = threading.Lock()
self.speeds = {} # kind -> {device: ema work/sec}
def report(self, kind: str, device: int, work: float, elapsed_s: float):
if work <= 0 or elapsed_s < 0.05:
return
s = work / elapsed_s
with self.lock:
d = self.speeds.setdefault(kind, {})
prev = d.get(device)
d[device] = s if prev is None else 0.7 * prev + 0.3 * s
def ratios(self, kind: str, devices: list) -> list:
with self.lock:
d = self.speeds.get(kind)
if d is not None and all(dev in d for dev in devices):
return [d[dev] for dev in devices]
# Incomplete measurements: architecture priors (mixing measured speeds with unit-less
# priors would skew the split, so it is one or the other)
return [arch_prior_speed(dev) for dev in devices]
def has_measured(self, kind: str, devices: list) -> bool:
with self.lock:
d = self.speeds.get(kind)
return d is not None and all(dev in d for dev in devices)
auto_split = AutoSplit()
# Pending timing events from the last measured quantize_tiles_multigpu call, and per-device
# accumulators: single calls are ~ms-scale, so busy time is aggregated over many calls and
# reported to the auto split once enough has been observed
_tiles_calls = 0
_tiles_timing = None
_tiles_acc = {} # device -> [tiles, seconds]
pinned_tiles: torch.Tensor | None = None
pinned_q_tiles: torch.Tensor | None = None
pinned_q_idx: torch.Tensor | None = None
def get_pinned(num_tiles: int):
global pinned_tiles, pinned_q_tiles, pinned_q_idx
if pinned_tiles is None or pinned_tiles.shape[0] < num_tiles:
pinned_tiles = torch.empty((num_tiles, 256), device = "cpu", dtype = torch.float, pin_memory = True)
pinned_q_tiles = torch.empty((num_tiles, 256), device = "cpu", dtype = torch.float, pin_memory = True)
pinned_q_idx = torch.empty((num_tiles, 256), device = "cpu", dtype = torch.int16, pin_memory = True)
return pinned_tiles[:num_tiles, :], pinned_q_tiles[:num_tiles, :], pinned_q_idx[:num_tiles, :]
def quantize_tiles_multigpu(tiles, quant_args: dict):
"""
Quantize tiles across multiple GPUs using pinned host memory for asynchronous fan-out/fan-in.
The input starts on the first device, is copied once to pinned CPU memory, split by device_ratios or evenly, and
each GPU quantizes its slice on a per-device stream. Results are copied back through pinned memory and gathered
on the first device.
"""
if frac_k(quant_args["K"]) is not None:
return quantize_tiles_frac(tiles, quant_args)
devices = quant_args["devices"]
if len(devices) == 1:
return quantize_tiles(tiles, quant_args)
# Get pinned buffers
pin_tiles, pin_q_tiles, pin_q_idx = get_pinned(tiles.shape[0])
# Copy input tiles to pinned memory. Input is always on the first device in the split
copy_input_event = torch.cuda.Event(blocking = False)
main_stream = get_quant_stream(devices[0])
with torch.cuda.stream(main_stream):
tiles = tiles.contiguous()
pin_tiles.copy_(tiles, non_blocking = True)
copy_input_event.record(main_stream)
# Create split slices for input tiles, output tiles and output indices
ratios = quant_args.get("device_ratios")
if ratios:
s = sum(ratios)
split_sizes = [tiles.shape[0] * r / s for r in ratios]
split_sizes = [round(s / 16) * 16 for s in split_sizes]
split_sizes[-1] += tiles.shape[0] - sum(split_sizes)
else:
split_sizes = [tiles.shape[0] // len(devices)] * len(devices)
split_sizes[-1] += tiles.shape[0] - sum(split_sizes)
# Account for negative splits (edge case if too many GPUs and/or tensor too small)
for i in range(len(split_sizes) - 2, -1, -1):
if split_sizes[i + 1] < 0:
split_sizes[i] += split_sizes[i + 1]
split_sizes[i + 1] = 0
# Harvest the previous measured call's per-device busy times (the events have long
# completed by the next call) into the accumulators, and report once every device has
# enough observed time for a meaningful speed. Per-device event pairs, since elapsed_time
# cannot cross devices; a new measured call is fenced periodically
global _tiles_calls, _tiles_timing
if _tiles_timing is not None:
p_devs, p_sizes, p_starts, p_ends = _tiles_timing
if all(e.query() for e in p_ends):
for dv, sz, e0, e1 in zip(p_devs, p_sizes, p_starts, p_ends):
acc = _tiles_acc.setdefault(dv, [0, 0.0])
acc[0] += sz
acc[1] += e0.elapsed_time(e1) / 1000
_tiles_timing = None
if all(_tiles_acc.get(dv, (0, 0.0))[1] > 0.25 for dv in devices):
for dv in devices:
acc = _tiles_acc.pop(dv)
auto_split.report("quant_tiles", dv, acc[0], acc[1])
_tiles_calls += 1
measure = (
_tiles_timing is None and _tiles_calls % 4 == 0 and
all(s > 0 for s in split_sizes) and tiles.shape[0] >= 256
)
m_starts = []
pin_split_tiles = torch.split(pin_tiles, split_sizes)
pin_split_q_tiles = torch.split(pin_q_tiles, split_sizes)
pin_split_q_idx = torch.split(pin_q_idx, split_sizes)
slice_done_events = []
for i, device in enumerate(devices):
stream = get_quant_stream(device)
with torch.cuda.stream(stream):
# Wait for input in host memory
if i > 0:
stream.wait_event(copy_input_event)
if measure:
e0 = torch.cuda.Event(enable_timing = True)
e0.record(stream)
m_starts.append(e0)
if split_sizes[i] > 0:
# Asynchronously copy the slice from the pinned buffer to device memory
dev_tiles = pin_split_tiles[i].to(device, non_blocking = True)
# Preallocate output tensors on the device.
dev_q_tiles = torch.empty_like(dev_tiles, device = device)
dev_q_idx = torch.empty_like(dev_tiles, dtype = torch.short, device = device)
# Work buffers
K = quant_args["K"]
mcg, mul1 = quant_codebook(quant_args)
temp_costs, temp_edges = get_temp_buffers(device, K, 256, 2 if mul1 else 1 if mcg else 0)
ext.quantize_tiles(
dev_tiles,
dev_q_tiles,
dev_q_idx,
temp_costs,
temp_edges,
K,
mcg,
mul1
)
# Async copy back to pinned memory
pin_split_q_tiles[i].copy_(dev_q_tiles, non_blocking = True)
pin_split_q_idx[i].copy_(dev_q_idx, non_blocking = True)
# Finished slice
evt = torch.cuda.Event(blocking = False, enable_timing = measure)
slice_done_events.append(evt)
evt.record(stream)
if measure:
_tiles_timing = (list(devices), list(split_sizes), m_starts, slice_done_events)
# Copy pinned buffers to original device
with torch.cuda.stream(main_stream):
for evt in slice_done_events:
main_stream.wait_event(evt)
q_tiles = torch.empty_like(tiles, device = devices[0])
q_idx = torch.empty_like(tiles, dtype = torch.short, device = devices[0])
q_tiles.copy_(pin_q_tiles, non_blocking = True)
q_idx.copy_(pin_q_idx, non_blocking = True)
return q_tiles, q_idx
def quantize_tiles_multigpu_sync(tiles, quant_args: dict):
"""
Simpler synchronized multi-GPU tile quantization path.
This variant explicitly copies tile chunks to each device, synchronizes around the work, then gathers results
back to the first device. It is easier to reason about than the pinned asynchronous path but offers less overlap.
"""
devices = quant_args["devices"]
if len(devices) == 1:
return quantize_tiles(tiles, quant_args)
tiles = tiles.contiguous()
split_sizes = [tiles.shape[0] // len(devices)] * len(devices)
split_sizes[-1] += tiles.shape[0] - sum(split_sizes)
split_tiles = torch.split(tiles, split_sizes)
tiles_per_device = [chunk.to(device) for chunk, device in zip(split_tiles, devices)]
torch.cuda.synchronize()
q_tiles_per_device = []
q_idx_per_device = []
for dev_tiles, device in zip(tiles_per_device, devices):
with torch.cuda.stream(get_quant_stream(device)):
dev_q_tiles, dev_q_idx = quantize_tiles(dev_tiles, quant_args)
q_tiles_per_device.append(dev_q_tiles)
q_idx_per_device.append(dev_q_idx)
for device in devices:
torch.cuda.synchronize(device)
q_tiles_per_device = [x.to(devices[0]) for x in q_tiles_per_device]
q_idx_per_device = [x.to(devices[0]) for x in q_idx_per_device]
quantized_tiles = torch.cat(q_tiles_per_device, dim = 0)
quantized_idx = torch.cat(q_idx_per_device, dim = 0)
return quantized_tiles, quantized_idx
# hipBLASLt 7.2 on gfx1201 (reached through rocBLAS for torch fp32 matmuls) silently returns wrong rows from
# row 2^19 or 2^20 onward (algorithm-dependent) in fp32 GEMMs; https://github.com/ROCm/rocm-libraries/issues/13243.
# The right-side rotation flattens to k * n / had_dim rows, which only the lm_head of large-vocab models exceeds
# (2048 x 248320 -> ~4M rows), corrupting it. Keep every GEMM at or below 2^19 rows; chunking along k is exact,
# so results are unchanged everywhere else
_HAD_MAX_GEMM_ROWS = 1 << 19
def preapply_had_l(x: torch.Tensor, had_dim):
k, n = x.shape
x_dtype = x.dtype
x = x.to(torch.float)
had = get_hadamard_dt(had_dim, x.device, x.dtype, 1 / math.sqrt(had_dim))
x = (had @ x.view(-1, had_dim, n)).view(k, n)
x = x.to(x_dtype)
return x
def preapply_had_r(x: torch.Tensor, had_dim):
k, n = x.shape
x_dtype = x.dtype
x = x.to(torch.float)
had = get_hadamard_dt(had_dim, x.device, x.dtype, 1 / math.sqrt(had_dim))
rows = max(1, _HAD_MAX_GEMM_ROWS * had_dim // n)
if rows >= k:
x = (x.view(k, -1, had_dim) @ had).view(k, n)
else:
x = torch.cat([
(x[i:i + rows].reshape(-1, had_dim) @ had).view(-1, n)
for i in range(0, k, rows)
], dim = 0)
x = x.to(x_dtype)
return x
def blockwise_preapply_had_l_(x: torch.Tensor, had_dim):
k, n = x.shape
assert k % had_dim == 0
assert x.dtype == torch.float
had = get_hadamard_dt(had_dim, x.device, x.dtype, 1 / math.sqrt(had_dim))
num_blocks = k // had_dim
for i in range(num_blocks):
start = i * had_dim
end = start + had_dim
block = x[start:end, :] # shape (had_dim, n)
block_transformed = had @ block.view(had_dim, n)
x[start:end, :] = block_transformed
def blockwise_preapply_had_r_(x: torch.Tensor, had_dim):
k, n = x.shape
assert n % had_dim == 0
assert x.dtype == torch.float
had = get_hadamard_dt(had_dim, x.device, x.dtype, 1 / math.sqrt(had_dim))
num_blocks = n // had_dim
for i in range(num_blocks):
start = i * had_dim
end = start + had_dim
block = x[:, start:end] # shape (k, had_dim)
block_transformed = block @ had
x[:, start:end] = block_transformed
def save_failed_cholesky(H: torch.Tensor, quant_args: dict, debug_info: dict | None):
"""Dump a Hessian that failed to decompose, before any retry damping is added, so the
failure can be studied offline. H arrives sign-flipped and Hadamard-transformed with the
base sigma_reg damping applied; su (in the dump) and the block Hadamard are orthogonal, so
the captured matrix is exactly recoverable by applying the inverse transforms."""
debug_dir = quant_args.get("debug_dir")
if not debug_dir:
return
try:
os.makedirs(debug_dir, exist_ok = True)
key = (debug_info or {}).get("key") or "unknown"
path = os.path.join(debug_dir, f"cholesky_fail_{key}.pt")
torch.save({
"H": H.detach().cpu().clone(),
"sigma_reg": quant_args.get("sigma_reg", 0.025),
**{k: (v.detach().cpu().clone() if isinstance(v, torch.Tensor) else v)
for k, v in (debug_info or {}).items()},
}, path)
print(f" !! Saved failing Hessian to {path}")
except Exception as e:
print(f" !! Failed to save Hessian debug dump: {e}")
def block_ldl(H: torch.Tensor, b: int, quant_args: dict, verbose: bool, debug_info: dict | None = None):
n, _ = H.shape
assert (n % b == 0)
m = n // b
# Cholesky factorization: H = L @ L.T
# Try on GPU first
num_cholesky_retries = 0
retry_cpu = False
while True:
try:
L = torch.linalg.cholesky(H)
# H is not needed after this, move to CPU. Then overwrite H's GPU storage with L, since we can't otherwise
# free up that VRAM as the tensor is referenced by the parent frame
H_cpu = H.cpu()
H.copy_(L) # VRAM copy, tiny overhead
L = H
H = H_cpu
break
except torch._C._LinAlgError as e:
num_cholesky_retries += 1
if num_cholesky_retries == 1:
# Capture the matrix as it first failed (retry damping has never actually
# recovered a failing decomposition, so the initial state is the interesting one)
save_failed_cholesky(H, quant_args, debug_info)
if num_cholesky_retries > 10:
print(" ## Cholesky decomp. failed, number of retries exceeded")
raise e
print(f" !! Cholesky decomp. failed, increasing diagonal damping, attempt {num_cholesky_retries}/10")
H.diagonal().add_(2.0 * quant_args.get("sigma_reg", 0.025) * H.diagonal().mean())
continue
# Fall back on CPU factorization
except Exception as e:
if e.__class__.__name__ == "OutOfMemoryError" or "CUDA out of memory" in str(e) or "HIP out of memory" in str(e):
retry_cpu = True
break
else:
raise e
if retry_cpu:
print(f" !! Out of memory on {str(H.device)}, trying CPU fallback")
free_mem()
H_cpu = H.cpu()
L_cpu = torch.linalg.cholesky(H_cpu)
# This is ugly, but overwrite H in VRAM to avoid allocating a new tensor, then replace reference with CPU copy
H.copy_(L_cpu)
del L_cpu
L = H
H = H_cpu
# Get blocks along diagonal of L: DL.shape = (m, b, b)
DL = torch.diagonal(L.reshape(m, b, m, b), dim1 = 0, dim2 = 2).permute(2, 0, 1)
# Compute D as D[i] = DL[i] @ DL[i].T for each diagonal block i (don't actually end up needing this)
# D = DL @ DL.transpose(1, 2)
# Invert each diagonal block
DL = torch.linalg.inv(DL)
# Multiply each block's column with its inverse
L = L.view(n, m, b)
for i in range(m):
L[:, i, :] = L[:, i, :] @ DL[i, :, :] # TODO: Could maybe be L[m * b:, i, :]?
L = L.reshape(n, n).contiguous()
# Insert block identity matrices along the diagonal.
# TODO: Figure out if this is necessary. Diagonal blocks should already be identities after previous step
L_block = L.view(m, b, m, b).permute(0, 2, 1,3)
dr = torch.arange(m)
L_block[dr, dr] = torch.stack([torch.eye(b, device = L.device, dtype = H.dtype)] * m)
return L, H # , D.to(DL.device)
def ldlq(
weight: torch.Tensor,
L: torch.Tensor,
quant_args: dict,
pb: ProgressBar | None = None
):
"""
:param weight:
Input weights, shape (k, n). If device is "cpu", result is collected on CPU as well, saving a bunch of
VRAM but adding a little PCIe overhead and many sync points
:param L:
LDL decomposition of regularized H
:param quant_args:
dict:
- K: bitrate
- buf_size_k: buffer size for LDLQ, along k
:param pb:
Optional ProgressPar to update, k // 16 steps
:return:
tuple:
- quantized weight, shape (k, n)
- indices (unpacked), shape (k // 16, n // 16, 256), uint16_t
"""
devices = quant_args["devices"]
for device in devices:
torch.cuda.synchronize(device)
main_stream = get_quant_stream(devices[0])
with torch.cuda.stream(main_stream):
devices = quant_args["devices"]
device = L.device
assert device == torch.device(devices[0])
buffer_device = weight.device
size_k, size_n = weight.shape # Row-major
assert size_k % 16 == 0
assert size_n % 128 == 0
tiles_k = size_k // 16
tiles_n = size_n // 16
buf_size_k = max(quant_args.get("buf_size_k", 128), 16)
assert buf_size_k % 16 == 0
assert size_k % buf_size_k == 0
p_row = 0
# Work buffers
prod_cache = torch.zeros((size_k, size_n), dtype = torch.float, device = device)
weight_q = torch.zeros((size_k, size_n), dtype = torch.float, device = buffer_device)
encoded = torch.zeros((tiles_k, tiles_n, 256), dtype = torch.short, device = buffer_device)
for j in range(size_k, 0, -buf_size_k):
i = j - buf_size_k
# Current span is rows i:j
b_weight = weight[i:j].to(device)
b_weight_q = weight_q[i:j] if device == buffer_device else \
torch.zeros_like(weight_q[i:j], device = device)
b_encoded = encoded[i // 16 : j // 16] if device == buffer_device else \
torch.zeros_like(encoded[i // 16 : j // 16], device = device)
b_prod_cache = prod_cache[i:j]
b_L = L[i:j]
# Iterate over rows of blocks in current span
for bj in range(buf_size_k, 0, -16):
bi = bj - 16
# Error so far for the current span
bb_err = b_weight[bj:] - b_weight_q[bj:]
# Corresponding slice of LDL decomposition of H
bb_L = b_L[bj:, i + bi:i + bj]
# Input tiles for quantization
compensation_term = b_prod_cache[bi:bj]
compensation_term.addmm_(bb_L.T, bb_err, alpha = 1.0, beta = 1.0)
rows = b_weight[bi:bj] + compensation_term
tiles = rows.reshape(16, tiles_n, 16).permute(1, 0, 2).reshape(tiles_n, 256)
# Pre-permute to tensor core layout
tiles = tiles[:, tensor_core_perm(device)]
# Quantize
quant_w, quant_i = quantize_tiles_multigpu(tiles, quant_args)
# Undo permutation on reconstructed tiles, but keep indices in tensor core layout
quant_w = quant_w[:, tensor_core_perm_i(device)]
# Store result
quant_w = quant_w.reshape(tiles_n, 16, 16).permute(1, 0, 2).reshape(16, size_n)
b_weight_q[bi:bj] = quant_w
b_encoded[bi // 16 : bj // 16] = quant_i.unsqueeze(0)
# Update progress
if pb:
p_row += 1
pb.update(p_row)
# Collect output
if device != buffer_device:
weight_q[i:j] = b_weight_q.to(buffer_device)
encoded[i // 16 : j // 16] = b_encoded.to(buffer_device)
# Cache error term for the rest of the matrix
b_err = b_weight - b_weight_q
prod_cache.addmm_(b_L.T, b_err, alpha = 1.0, beta = 1.0)
for device in devices:
torch.cuda.synchronize(device)
return weight_q, encoded
def ldlq_2hess(
weight: torch.Tensor,
L_in: torch.Tensor,
L_out: torch.Tensor,
quant_args: dict,
pb: ProgressBar | None = None
):
"""
LDLQ with error feedback along both axes, for a Kronecker-factored Hessian H ~ H_out (x) H_in (YAQA, Tseng et al.
2025). Objective tr(E^T H_in E H_out) with E = W - Q, shape (k, n); L_in (k, k) and L_out (n, n) are the block-16
LDL factors (zero diagonal, as for ldlq()). The 16x16 tiles are quantized along anti-diagonals from the far corner,
every tile of a diagonal being independent given the earlier ones. With L' = I + L the input of tile (a, c) is
W[a, c] + (L_in'^T E L_out')[a, c], E restricted to the tiles already quantized
kept as a running matrix F that receives one rank-(16 * tiles) update per diagonal:
F += L_in'[rows, :]^T . blockdiag(dE) . L_out'[cols, :]
:return:
tuple, as ldlq(): quantized weight (k, n), indices (k // 16, n // 16, 256)
"""
devices = quant_args["devices"]
for device in devices:
torch.cuda.synchronize(device)
main_stream = get_quant_stream(devices[0])
with torch.cuda.stream(main_stream):
device = L_in.device
weight = weight.to(device)
size_k, size_n = weight.shape
assert size_k % 16 == 0 and size_n % 16 == 0
tiles_k, tiles_n = size_k // 16, size_n // 16
Lk = L_in.clone(); Lk.diagonal().fill_(1.0) # unit lower triangular
Ln = L_out.clone(); Ln.diagonal().fill_(1.0)
F = torch.zeros((size_k, size_n), dtype = torch.float, device = device)
weight_q = torch.zeros_like(weight)
encoded = torch.zeros((tiles_k, tiles_n, 256), dtype = torch.short, device = device)
W4 = weight.view(tiles_k, 16, tiles_n, 16)
Q4 = weight_q.view(tiles_k, 16, tiles_n, 16)
F4 = F.view(tiles_k, 16, tiles_n, 16)
perm, perm_i = tensor_core_perm(device), tensor_core_perm_i(device)
ar16 = torch.arange(16, device = device)
steps = tiles_k + tiles_n - 1
for step, s in enumerate(range(steps - 1, -1, -1)):
a = torch.arange(max(0, s - (tiles_n - 1)), min(tiles_k - 1, s) + 1, device = device)
c = s - a
tiles = (W4[a, :, c, :] + F4[a, :, c, :]).reshape(-1, 256) # (T, 16, 16) row major
quant_w, quant_i = quantize_tiles_multigpu(tiles[:, perm].contiguous(), quant_args)
quant_w = quant_w[:, perm_i].view(-1, 16, 16)
Q4[a, :, c, :] = quant_w
encoded[a, c] = quant_i
if s > 0:
dE = W4[a, :, c, :] - quant_w # (T, 16, 16)
rows = (a.unsqueeze(1) * 16 + ar16).flatten() # (16 T,)
cols = (c.unsqueeze(1) * 16 + ar16) # (T, 16)
# Only tiles above and to the left can still be pending: restrict the update to that corner
k_hi = int(a.max().item()) * 16 + 16
n_hi = int(c.max().item()) * 16 + 16
right = torch.bmm(dE, Ln[cols][:, :, :n_hi]) # (T, 16, n_hi)
F[:k_hi, :n_hi].addmm_(Lk[rows, :k_hi].T, right.reshape(-1, n_hi))
if pb and step % 4 == 0:
pb.update(min(tiles_k, step * tiles_k // steps))
for device in devices:
torch.cuda.synchronize(device)
return weight_q, encoded
def prepare_H_out(H_out: torch.Tensor, sv: torch.Tensor, quant_args: dict, verbose: bool, device):
"""
Output-side Hessian factor in the quantizer's rotated basis, with its block LDL factor. The rotated weight relates
to the original by E_o = S_u P_k E_r P_n S_v, so tr(E_o^T H_in E_o H_out) = tr(E_r^T [..] E_r [P_n S_v H_out S_v P_n]):
the full output scale vector sv (signs and channel scales) goes into the transform.
"""
Ho = H_out.to(device, torch.float).clone()
s = sv.flatten().to(device, torch.float)
Ho *= s.unsqueeze(0)
Ho *= s.unsqueeze(1)
blockwise_preapply_had_r_(Ho, had_n)
blockwise_preapply_had_l_(Ho, had_n)
Ho /= Ho.diagonal().mean().clamp_min(1e-30)
Ho.diagonal().add_(quant_args.get("sigma_reg_out", quant_args.get("sigma_reg", 0.025)))
L_out, Ho = block_ldl(Ho, 16, quant_args, verbose)
L_out.diagonal().fill_(0.0)
return L_out.to(device), Ho.to(device)
def fallback_quant(
weight: torch.Tensor,
q_device: torch.Tensor,
quant_args: dict,
pb: ProgressBar | None = None
):
"""
Perform the same quantization as ldlq() but without an LDL decomposition
:param weight:
Input weights, shape (k, n). If device is "cpu", result is collected on CPU as well, saving a bunch of
VRAM but adding a little PCIe overhead and many sync points
:param q_device:
Target device
:param quant_args:
dict:
- K: bitrate
- buf_size_k: buffer size for faux-LDLQ, along k
:param pb:
Optional ProgressPar to update, k // 16 steps
:return:
tuple:
- quantized weight, shape (k, n)
- indices (unpacked), shape (k // 16, n // 16, 256), uint16_t
"""
devices = quant_args["devices"]
for device in devices:
torch.cuda.synchronize(device)
main_stream = get_quant_stream(devices[0])
with torch.cuda.stream(main_stream):
devices = quant_args["devices"]
device = weight.device
assert device == torch.device(devices[0])
buffer_device = weight.device
size_k, size_n = weight.shape # Row-major
assert size_k % 16 == 0
assert size_n % 128 == 0
tiles_k = size_k // 16
tiles_n = size_n // 16
buf_size_k = max(quant_args.get("buf_size_k", 128), 16)
assert buf_size_k % 16 == 0
assert size_n % buf_size_k == 0
p_row = 0
# Work buffers
weight_q = torch.zeros((size_k, size_n), dtype = torch.float, device = buffer_device)
encoded = torch.zeros((tiles_k, tiles_n, 256), dtype = torch.short, device = buffer_device)
# Accumulate sum of squared error on-device to avoid per-block CPU syncs
mse_sum = torch.zeros((), dtype=torch.float64, device=device)
for j in range(size_k, 0, -buf_size_k):
i = j - buf_size_k
# Current span is rows i:j
b_weight = weight[i:j].to(device)
b_weight_q = weight_q[i:j] if device == buffer_device else \
torch.zeros_like(weight_q[i:j], device = device)
b_encoded = encoded[i // 16 : j // 16] if device == buffer_device else \
torch.zeros_like(encoded[i // 16 : j // 16], device = device)
# Iterate over rows of blocks in current span
for bj in range(buf_size_k, 0, -16):
bi = bj - 16
# Input tiles for quantization
rows = b_weight[bi:bj]
tiles = rows.reshape(16, tiles_n, 16).permute(1, 0, 2).reshape(tiles_n, 256)
# Pre-permute to tensor core layout
tiles = tiles[:, tensor_core_perm(device)]
# Quantize
quant_w, quant_i = quantize_tiles_multigpu(tiles, quant_args)
# Undo permutation on reconstructed tiles, but keep indices in tensor core layout
quant_w = quant_w[:, tensor_core_perm_i(device)]
# Restore row-major block layout
quant_w = quant_w.reshape(tiles_n, 16, 16).permute(1, 0, 2).reshape(16, size_n)
# Accumulate squared error for this block
mse_sum += torch.sum(
(rows.float() - quant_w.float()).pow(2),
dtype=torch.float64,
)
# Store result
b_weight_q[bi:bj] = quant_w
b_encoded[bi // 16 : bj // 16] = quant_i.unsqueeze(0)
# Update progress
if pb:
p_row += 1
pb.update(p_row)
# Collect output
if device != buffer_device:
weight_q[i:j] = b_weight_q.to(buffer_device)
encoded[i // 16 : j // 16] = b_encoded.to(buffer_device)
for device in devices:
torch.cuda.synchronize(device)
mse = (mse_sum / weight.numel()).item()
return weight_q, encoded, mse
def block_trace(A, B, block_size = 1024):
"""
Compute trace(A.T @ B @ A) in column blocks of B to bound the temporary
"""
total = 0.0
for j_start in range(0, B.shape[1], block_size):
j_end = min(j_start + block_size, B.shape[1])
B_block = B[:, j_start:j_end]
A_j_block = A[j_start:j_end, :]
partial = torch.einsum("ik,ij,jk->", A, B_block, A_j_block)
total += partial.item()
return total
def ldlq_batched(
weights: torch.Tensor,
Ls: torch.Tensor,
quant_args: dict,
pb: ProgressBar | None = None
):
"""
LDLQ over a stack of same-shape tensors with per-tensor L, batched so each 16-row step quantizes all
tensors' tiles in one quantize_tiles call. The recursion is identical to ldlq() per batch element; only
the matmuls become bmm. All buffers stay on-device (this path is only used for groups of small tensors).
:param weights:
Input weights, shape (B, k, n), float, on the quant device
:param Ls:
LDL decompositions of the regularized Hessians, shape (B, k, k), same device
:return:
tuple:
- quantized weights, shape (B, k, n)
- indices (unpacked), shape (B, k // 16, n // 16, 256), int16
"""
devices = quant_args["devices"]
for device in devices:
torch.cuda.synchronize(device)
main_stream = get_quant_stream(devices[0])
with torch.cuda.stream(main_stream):
device = weights.device
assert device == torch.device(devices[0]) and Ls.device == device
B, size_k, size_n = weights.shape
assert Ls.shape == (B, size_k, size_k)
assert size_k % 16 == 0
assert size_n % 128 == 0
tiles_k = size_k // 16
tiles_n = size_n // 16
buf_size_k = max(quant_args.get("buf_size_k", 128), 16)
assert buf_size_k % 16 == 0
assert size_k % buf_size_k == 0
p_row = 0
perm = tensor_core_perm(device)
perm_i = tensor_core_perm_i(device)
prod_cache = torch.zeros((B, size_k, size_n), dtype = torch.float, device = device)
weight_q = torch.zeros_like(weights)
encoded = torch.zeros((B, tiles_k, tiles_n, 256), dtype = torch.short, device = device)
for j in range(size_k, 0, -buf_size_k):
i = j - buf_size_k
for bj in range(buf_size_k, 0, -16):
bi = bj - 16
gi, gj = i + bi, i + bj
# Error so far for the remaining rows of the current span
bb_err = weights[:, gj:j] - weight_q[:, gj:j]
# Corresponding slice of the LDL decompositions
bb_L = Ls[:, gj:j, gi:gj]
# Input tiles for quantization; out-of-place equivalent of ldlq()'s in-place accumulation
# (the mutated prod_cache rows are never read again there)
compensation = torch.baddbmm(prod_cache[:, gi:gj], bb_L.transpose(1, 2), bb_err)
rows = weights[:, gi:gj] + compensation
tiles = rows.reshape(B, 16, tiles_n, 16).permute(0, 2, 1, 3).reshape(B * tiles_n, 256)
tiles = tiles[:, perm]
quant_w, quant_i = quantize_tiles_multigpu(tiles, quant_args)
quant_w = quant_w[:, perm_i]
quant_w = quant_w.reshape(B, tiles_n, 16, 16).permute(0, 2, 1, 3).reshape(B, 16, size_n)
weight_q[:, gi:gj] = quant_w
encoded[:, gi // 16] = quant_i.view(B, tiles_n, 256)
if pb:
p_row += 1
pb.update(p_row)
# Cache error term for the rest of the matrices
b_err = weights[:, i:j] - weight_q[:, i:j]
prod_cache.baddbmm_(Ls[:, i:j].transpose(1, 2), b_err)
for device in devices:
torch.cuda.synchronize(device)
return weight_q, encoded
finalize_capture_H_mutex = threading.Lock()
def finalize_capture_H(H_data: dict, quant_args: dict, verbose: bool):
with finalize_capture_H_mutex:
if H_data["H"].is_meta:
H_data["L"] = None
H_data["finalized"] = True
H_data["diag"] = None
H_data["q_fallback"] = True
H = H_data["H"]
k = H.shape[0]
su = (torch.randn(k, device = H_data["device"]).sign() + 1e-5).sign().to(torch.float).unsqueeze(1)
H_data["su"] = su
return True, None, None, su, None
# Unswap H, and reconstruct the square matrix from the capture's block-packed upper triangle
if "H_swap_device" in H_data:
H_data["H"] = H_data["H"].to(H_data["H_swap_device"])
del H_data["H_swap_device"]
if H_data.get("packed"):
H_data["H"] = unpack(H_data["H"], H_data["n"])
H_data["packed"] = False
H = H_data["H"]
if H_data["finalized"]:
return H_data["q_fallback"], H, H_data["L"], H_data["su"], H_data["diag"]
# Mean of samples summed up during forward pass
# Switch to uncalibrated fallback if no input activations or diagonal is too small (few activations)
count = H_data["count"]
if count == 0:
q_fallback = True
diag_mean = 0.0
else:
H /= count
diag_mean = torch.diag(H).mean().item()
# A non-finite Hessian (e.g. fp16 activation overflow reaching the capture) cannot
# be repaired by damping, and NaN would also defeat the < comparison below and turn
# the whole diagonal non-finite through the regularization term
if not math.isfinite(diag_mean):
inf_nan = H_data.get("inf_nan")
counts = f" (captured {inf_nan[0].item():,} inf, {inf_nan[1].item():,} NaN activation values)" \
if inf_nan is not None else ""
print(f" !! Non-finite Hessian for {H_data.get('first_key')}, using uncalibrated fallback{counts}")
q_fallback = True
diag_mean = 0.0
else:
q_fallback = diag_mean < 1e-20
# Regularize diagonal
H.diagonal().add_(quant_args.get("sigma_reg", 0.025) * diag_mean)
# Some tests
diag = H.diagonal().clone()
if verbose:
print(f" - H min/max: {H.min().item():.6f} {H.max().item():.6f}")
print(f" - H mean/std: {H.mean().item():.6f} {H.std().item():.6f}")
print(f" - H diag min/max: {diag.min():.6f} {diag.max():.6f} ")
# Random sign flips for input channel, fixed for the first linear layer to quantize with this H
k = H.shape[0]
su = (torch.randn(k, device = H.device).sign() + 1e-5).sign().to(torch.float).unsqueeze(1)
H_data["su"] = su
# Input had
H *= su.T
blockwise_preapply_had_r_(H, had_k)
H *= su
blockwise_preapply_had_l_(H, had_k)
# Get block LDL decomposition of H, zero diagonal. With scaled Hessians every tensor factors its
# own copy after regularize (see scaled_ldl), so nothing is kept here
if q_fallback or quant_args.get("scaled_hessian"):
L = None
else:
L, H = block_ldl(H, 16, quant_args, verbose, debug_info = {
"key": H_data.get("first_key"),
"count": H_data.get("count"),
"num_total": H_data.get("num_total"),
"inf_nan": H_data.get("inf_nan"),
"su": su,
})
dr = torch.arange(k)
L[dr, dr] = 0
H_data["L"] = L
# H is no longer needed except to compute proxy error so move to CPU
H = H.cpu()
H_data["H"] = H.cpu()
H_data["finalized"] = True
H_data["diag"] = diag
H_data["q_fallback"] = q_fallback
return q_fallback, H, L, su, diag
def pack_trellis(encoded: torch.Tensor, quant_args: dict) -> torch.Tensor:
K = quant_args["K"]
shape = encoded.shape
assert len(shape) == 3 and shape[2] == 256
assert encoded.dtype == torch.int16
packed_shape = (shape[0], shape[1], trellis_words(K))
packed = torch.zeros(packed_shape, dtype = torch.int16, device = encoded.device)
fk = frac_k(K)
if fk is not None:
ext.pack_trellis_frac(packed, encoded.contiguous(), fk[0], fk[1])
else:
ext.pack_trellis(packed, encoded.contiguous(), K)
# unpacked = torch.zeros_like(encoded)
# ext.unpack_trellis(unpacked, packed, K)
# assert torch.equal(unpacked, encoded)
return packed
def pack_signs(signs: torch.Tensor, quant_args: dict) -> torch.Tensor:
signs = signs.half().flatten().contiguous()
assert signs.shape[0] % 16 == 0
packed = torch.zeros(signs.shape[0] // 16, dtype = torch.int16, device = signs.device)
ext.pack_signs(packed, signs)
return packed
# LDLQ error feedback inflates the tiles the quantizer actually sees relative to the regularized
# weight (accumulated compensation grows along the recursion). Mean stream RMS over the recursion
# relative to the first block, measured on Qwen3-8B (q/up/down, layer 10): K=1 1.06-1.10, 1.5
# 1.03-1.05, 2 1.01-1.02, 2.5 ~1.01, 3 ~1.005, >= 4 ~1.0. The global scale search runs on the
# pre-feedback weight, so its sample tiles are inflated by the expected drift and the scale it
# picks lands the recursion's mean input at the codebook's operating point
LDLQ_DRIFT = {1: 1.08, 1.5: 1.035, 2: 1.018, 2.5: 1.009, 3: 1.004}
def ldlq_drift(K) -> float:
return LDLQ_DRIFT.get(float(K), 1.0)
def sample_scale_tiles(weight_r: torch.Tensor, width: int = 3) -> torch.Tensor:
"""
Sample tiles for the global scale search: a wrapped diagonal, guaranteeing every tile row and column is
sampled at least once (outliers are typically whole input or output channels), plus the tiles with the
highest and lowest RMS as explicit outlier insurance for anything a diagonal of the given width misses.
"""
device = weight_r.device
tiles_k = weight_r.shape[0] // 16
tiles_n = weight_r.shape[1] // 16
w4 = weight_r.view(tiles_k, 16, tiles_n, 16)
diag_len = max(tiles_k, tiles_n)
ii = torch.arange(diag_len, device = device).repeat_interleave(width)
ww = torch.arange(width, device = device).repeat(diag_len)
kk = ii % tiles_k
nn = (ii + ww) % tiles_n
# Tiles with extreme RMS, by flat tile index
num_x = max(8, (diag_len * width) // 16)
tile_ms = w4.square().mean(dim = (1, 3)).flatten()
num_x = min(num_x, (tile_ms.shape[0] + 1) // 2)
hi = torch.topk(tile_ms, num_x).indices
lo = torch.topk(tile_ms, num_x, largest = False).indices
xk = torch.cat((hi, lo)) // tiles_n
xn = torch.cat((hi, lo)) % tiles_n
tiles = w4[torch.cat((kk, xk)), :, torch.cat((nn, xn)), :].reshape(-1, 256)
return tiles[:, tensor_core_perm(device)].contiguous()
def g_scale_search_batch(
samples: list[torch.Tensor],
quant_args: dict,
) -> list[tuple[float, torch.Tensor]]:
"""
Global scale search over a group of tensors' tile samples, batching all evaluations into as few
quantize_tiles calls as possible (two host syncs per group, regardless of group size).
The sampled error curve is flat near the optimum but not perfectly unimodal, so a sequential golden-section
search can converge on a local wiggle well away from the true minimum. Instead: a coarse grid over the full
range on a subsample of the tiles, then a fine grid around the coarse minimum on the full sample, refined by
parabolic interpolation; the curve's flatness makes the grid spacing precise enough.
"""
n_t = len(samples)
device = samples[0].device
max_tiles = 65536
def eval_pairs(pairs, tile_sets):
# pairs: list of (tensor_idx, scale); one mse per pair, kernel calls chunked to the batch limit
out = torch.empty(len(pairs), dtype = torch.float, device = device)
i = 0
while i < len(pairs):
j, tot = i, 0
while j < len(pairs) and tot + tile_sets[pairs[j][0]].shape[0] <= max_tiles:
tot += tile_sets[pairs[j][0]].shape[0]
j += 1
j = max(j, i + 1)
batch = torch.cat([tile_sets[t] * s for t, s in pairs[i:j]])
quant_w, _ = quantize_tiles_multigpu(batch, quant_args)
offset = 0
for p, (t, s) in enumerate(pairs[i:j]):
cnt = tile_sets[t].shape[0]
out[i + p] = (quant_w[offset : offset + cnt] / s - tile_sets[t]).square().mean()
offset += cnt
i = j
return out
# Stage 1: coarse grid over the full search range, on a subsample of each tensor's tiles
coarse = [0.1 + 0.2 * i for i in range(10)] # 0.1 .. 1.9, same range as the old golden-section bracket
subs = [s[::3] for s in samples]
pairs1 = [(t, s) for t in range(n_t) for s in coarse]
mse1 = eval_pairs(pairs1, subs).view(n_t, len(coarse))
centers = [coarse[c] for c in mse1.argmin(dim = 1).tolist()]
# Stage 2: fine grid around each tensor's coarse minimum, on its full sample
step = 0.075
fine = [[c + step * (i - 2) for i in range(5)] for c in centers]
pairs2 = [(t, s) for t in range(n_t) for s in fine[t]]
mse2 = eval_pairs(pairs2, samples).view(n_t, 5)
mse2_h = mse2.tolist()
results = []
for t in range(n_t):
best = min(range(5), key = lambda i: mse2_h[t][i])
if 0 < best < 4:
y0, y1, y2 = mse2_h[t][best - 1], mse2_h[t][best], mse2_h[t][best + 1]
denom = y0 - 2.0 * y1 + y2
offset = 0.5 * (y0 - y2) / denom if denom > 0 else 0.0
offset = max(-0.5, min(0.5, offset))
else:
offset = 0.0
best_scale = max(fine[t][best] + offset * step, 0.01)
results.append((best_scale, mse2[t, best]))
return results
def g_scale_gss(
weight_r: torch.Tensor,
verbose: bool,
quant_args: dict,
width: int = 3,
pb: ProgressBar = None
):
devices = quant_args["devices"]
for device in devices:
torch.cuda.synchronize(device)
main_stream = get_quant_stream(devices[0])
# TODO: Figure out why Torch always initializes cuda:0 when exiting this CM, even when it's not used
with torch.cuda.stream(main_stream):
tiles = sample_scale_tiles(weight_r, width)
tiles *= ldlq_drift(quant_args["K"])
if pb:
pb.update(50)
best_scale, best_mse = g_scale_search_batch([tiles], quant_args)[0]
if pb:
pb.update(100)
if verbose:
print(f" - scale search: min = {best_scale:.6f}, mse: {best_mse.item():.6f}")
for device in devices:
torch.cuda.synchronize(device)
return best_scale, best_mse
def block_rms(x: torch.Tensor, dim: int, keepdim: bool = False, blocksize: int = 32):
"""
Compute blockwise x.square().mean(dim, keepdim).sqrt()
"""
n = x.size(dim)
sq = None
for block in torch.split(x, blocksize, dim = dim):
block_sq = block.square().sum(dim = dim, keepdim = keepdim)
if sq is None:
sq = block_sq
else:
sq += block_sq
mean_sq = sq / n
return mean_sq.sqrt()
def block_rms_n(x: torch.Tensor, dim: int = 0, blocksize: int = 32):
"""
Compute blockwise x.square().mean().sqrt()
"""
n = 0
sq = None
for block in torch.split(x, blocksize, dim = dim):
block_sq = block.square().sum()
n += block.numel()
if sq is None:
sq = block_sq
else:
sq += block_sq
mean_sq = sq / n
return mean_sq.sqrt()
def block_nmse(x: torch.Tensor, y: torch.Tensor, dim: int = 0, blocksize: int = 32):
"""
Compute blockwise (x - y).square().mean().item() / y.square().mean().item()
"""
sq = None
diff_sq = None
for block_x, block_y in zip(torch.split(x, blocksize, dim = dim), torch.split(y, blocksize, dim = dim)):
block_sq = block_y.square().sum()
block_diff_sq = (block_x - block_y).square().sum()
if sq is None:
sq = block_sq
diff_sq = block_diff_sq
else:
sq += block_sq
diff_sq += block_diff_sq
return diff_sq.item() / (sq.item() + 1e-20)
def out_scales_ratio(out_channel_scales: torch.Tensor, out_sensitivity: torch.Tensor) -> float | None:
"""
Predicted loss with output channel scales relative to the loss without them.
Quantization error is uniform in the rotated domain, so the scales don't change how much error there is but
where it lands: with scales every output channel gets the same relative error (error energy proportional to
the channel's energy), without them the same absolute error. For a per-channel sensitivity s (loss per unit
of error energy) and channel energy e, the ratio of the two losses is E[s * e] / (E[s] * E[e]): below one
when the sensitive channels are the low-energy ones, above one when they are the high-energy ones.
Returns None when the sensitivity carries no information (wrong size, non-finite or all zero)
"""
e = out_channel_scales.flatten().double().square()
s = out_sensitivity.flatten().to(e.device).double()
if s.shape[0] > e.shape[0]:
s = s[:e.shape[0]]
elif s.shape[0] < e.shape[0]:
s = F.pad(s, (0, e.shape[0] - s.shape[0]))
den = s.mean() * e.mean()
if not torch.isfinite(s).all() or not torch.isfinite(den) or den.item() <= 0.0:
return None
return ((s * e).mean() / den).item()
# In auto mode, output scales are dropped for tensors whose sensitivity predicts a loss above this with them,
# relative to without
out_scales_max_ratio = 1.05
# Added to the mean-normalized sensitivity when scales are derived from it, so channels the estimate saw
# nothing of don't take scales without bound
out_sensitivity_floor = 0.02
def sensitivity_out_scales(out_channel_scales: torch.Tensor, out_sensitivity: torch.Tensor) -> torch.Tensor | None:
"""
Output channel scales that minimize the predicted loss for a per-channel sensitivity.
With scale t on a channel its error energy is proportional to t^2, and the total error in the rotated domain
to sum(sigma^2 / t^2) for channel RMS sigma. For sensitivity s (loss per unit of error energy) the loss
(sum s t^2) (sum sigma^2 / t^2) is minimal at t ~ (sigma^2 / s)^(1/4). Plain output scales, t = sigma, are
that optimum for s ~ 1 / sigma^2.
out_channel_scales: the plain scales (channel RMS), (1, n). Returns the scales in the same shape, mean one over
the nonzero channels, or None when the sensitivity carries no information (wrong size, non-finite or all zero)
"""
sigma = out_channel_scales.flatten().double()
s = out_sensitivity.flatten().to(sigma.device).double()
if s.shape[0] > sigma.shape[0] or not torch.isfinite(s).all():
return None
s = s.clamp(min = 0.0)
mean = s.mean().item()
if mean <= 0.0:
return None
# Channels past the sensitivity's length are padding
s = F.pad(s / mean, (0, sigma.shape[0] - s.shape[0]), value = 1.0) + out_sensitivity_floor
t = (sigma.square() / s).pow(0.25)
live = t > 0
if not live.any():
return None
return (t / t[live].mean()).to(out_channel_scales.dtype).view_as(out_channel_scales)
def regularize(
weight: torch.Tensor,
su: torch.Tensor,
sv: torch.Tensor,
quant_args: dict,
verbose: bool,
H_diag: torch.Tensor | None,
pb: ProgressBar | None,
skip_g_scale: bool = False,
q_fallback: bool = False
):
"""
Transform weights into the distribution expected by EXL3 tile quantization.
The routine chooses whether to apply output-channel scaling, folds output and input sign/scale vectors into the
matrix, applies blockwise Hadamard transforms, and optionally searches for a global scale that minimizes sample
quantization error. It returns the transformed weight plus the scale metadata needed to reconstruct the original
linear layer behavior after quantization.
Output-channel scaling is forced on or off by quant_args["apply_out_scales"] (True/False). With None (auto)
it is on, except for tensors that come with a per-channel output sensitivity in quant_args["out_sensitivity"]
(gate and up projections of gated MLPs, see modules/mlp.py) and whose sensitive channels are the high-energy
ones (see out_scales_ratio).
A tensor that comes with the diagonal of its output-side Hessian in quant_args["out_hessian_diag"] takes its
scales from that instead (see sensitivity_out_scales), and quant_args["out_scales_hessian"] reports it
"""
force_out_scales = quant_args["apply_out_scales"]
quant_args["out_scales_hessian"] = False
# dist_ref = torch.empty((512,), dtype = torch.float, device = weight.device)
# dist_r = torch.empty_like(dist_ref)
def jsd(h1, h2):
m = (h1 + h2) / 2
eps = 1e-12
js = F.kl_div((h1 + eps).log(), m, reduction = "sum") + \
F.kl_div((h2 + eps).log(), m, reduction = "sum")
return js / 2
apply_out_scales = True if force_out_scales is None else force_out_scales
# Apply output scales
out_channel_scales = block_rms(weight, dim = 0, keepdim = True)
mean = out_channel_scales.mean().item()
if mean > 1e-30:
out_channel_scales /= mean
quant_args["zeros"] = False
out_hessian_diag = quant_args.get("out_hessian_diag")
out_sensitivity = quant_args.get("out_sensitivity")
hessian_scales = None
if out_hessian_diag is not None and force_out_scales is None and not q_fallback:
hessian_scales = sensitivity_out_scales(out_channel_scales, out_hessian_diag)
if hessian_scales is not None:
ratio = out_scales_ratio(out_channel_scales, out_hessian_diag)
if verbose and ratio is not None:
print(f" - out scales loss ratio: {ratio:.6f} (output Hessian)")
out_channel_scales = hessian_scales
quant_args["out_scales_hessian"] = True
elif out_sensitivity is not None and not q_fallback:
ratio = out_scales_ratio(out_channel_scales, out_sensitivity)
if verbose and ratio is not None:
print(f" - out scales loss ratio: {ratio:.6f}")
if force_out_scales is None and ratio is not None:
apply_out_scales = ratio <= out_scales_max_ratio
else:
quant_args["zeros"] = True
if force_out_scales is not None:
apply_out_scales = True
zero_out_scales = out_channel_scales.abs() < 1e-30
if apply_out_scales:
out_channel_scales[zero_out_scales] = 0.1
sv = (sv * out_channel_scales + 1e-10).float()
if verbose:
out_channel_std = out_channel_scales.std().item()
out_channel_mean = out_channel_scales.mean().item()
print(f" - out ch scales std/mean: {out_channel_std:.6f} {out_channel_mean:.6f}")
# Output sign flips (and scales)
weight /= sv
# Force zero output channels to zero
sv[zero_out_scales] = 0.0
# Output hadamard transform
blockwise_preapply_had_r_(weight, had_n)
# Input sign flips and scales
in_channel_scales = block_rms(weight, dim = 1, keepdim = True)
in_channel_scales[in_channel_scales.abs() < 1e-30] = 0.1
su = (su * in_channel_scales / (-codebook_scale) + 1e-10).float() # mustn't be inplace
weight /= su
blockwise_preapply_had_l_(weight, had_k)
# Determine best scale for matrix by test quantizing a sample of tiles along a wrapped diagonal
if not skip_g_scale:
g_scale, mse_scale = g_scale_gss(weight, False, quant_args, pb = pb)
else:
g_scale = 1.0
mse_scale = None
weight *= g_scale
su /= g_scale
# ext.test_distribution(weight_os, dist_r, dist_ref, -3.8, 3.8)
# js_os = jsd(dist_r, dist_ref)
if verbose:
print(f" - su/sv std: {su.std().item():.6f} {sv.std().item():.6f}")
print(f" - global scale: {g_scale:.6f}")
if mse_scale is not None: print(f" - sample mse: {mse_scale.item():.6f}")
print(f" - apply_out_scales: {str(apply_out_scales)}")
return apply_out_scales, weight, g_scale, su, sv
def unrotate_H(H_rot: torch.Tensor, su_signs: torch.Tensor) -> torch.Tensor:
"""Undo finalize_capture_H's input transform on the (damped) Hessian: H_rot = P (S H S) P^T with
the blockwise 128-Hadamard P (its own inverse) and the sign flips S"""
H = H_rot.float().clone()
blockwise_preapply_had_r_(H, had_k)
blockwise_preapply_had_l_(H, had_k)
s = su_signs.flatten().sign().float().to(H.device)
H *= s.unsqueeze(0)
H *= s.unsqueeze(1)
return H
# Torch loads its CUDA linalg library on the first linalg call of the process and swaps the dispatch
# stubs while doing so; two threads making that first call together fail with "lazy wrapper should
# be called at most once". finalize_capture_H's factorization runs under its mutex, but scaled_ldl
# runs on the device worker threads, so the library is loaded once under a lock before they start
_linalg_load_lock = threading.Lock()
_linalg_loaded = False
def load_linalg(device):
global _linalg_loaded
if _linalg_loaded:
return
with _linalg_load_lock:
if not _linalg_loaded:
torch.linalg.cholesky(torch.eye(16, device = device))
_linalg_loaded = True
def original_H(H_data: dict, device, cache: bool = False) -> torch.Tensor:
"""The damped Hessian in the original basis, on device (undoing finalize_capture_H's rotation). With
`cache` the copy is kept in the H_data's per-device cache, for a Hessian many tensors share"""
dev_cache = H_data.setdefault("dev_cache", {}) if cache else None
H = dev_cache.get(("H_orig", device.index)) if cache else None
if H is None:
H = unrotate_H(H_data["H"].to(device), H_data["su"])
if cache:
dev_cache[("H_orig", device.index)] = H
return H
def scaled_ldl(H_orig: torch.Tensor, su: torch.Tensor, quant_args: dict) -> torch.Tensor:
"""
Block LDL of the Hessian in the domain LDLQ actually quantizes in. finalize_capture_H factors the
sign-flipped, rotated Hessian once per input (shared by q/k/v, gate/up, experts), but regularize
then divides each tensor's rows by that tensor's own input scales, so an error left in a heavily
scaled row comes back at full size in the output while the shared factorization weights every row
alike. The metric for the regularized weight is P diag(su) H diag(su) P with the full su (signs and
scales, the global scale cancels), which takes one factorization per tensor. H_orig is the damped
Hessian in the original basis (original_H), left untouched. Returns L with a zero diagonal, as
ldlq() expects
"""
load_linalg(H_orig.device)
s = su.view(-1).float().to(H_orig.device)
H = H_orig * s.unsqueeze(0)
H *= s.unsqueeze(1)
blockwise_preapply_had_r_(H, had_k)
blockwise_preapply_had_l_(H, had_k)
L, _ = block_ldl(H, 16, quant_args, False)
dr = torch.arange(L.shape[0], device = L.device)
L[dr, dr] = 0
return L
def damping_share(quant_args: dict) -> float:
"""finalize_capture_H adds sigma_reg x the mean diagonal to H; as a share of the damped mean diagonal"""
r = quant_args.get("sigma_reg", 0.025)
return r / (1.0 + r)
def refit_scales(weight: torch.Tensor, weight_q: torch.Tensor, H: torch.Tensor, su: torch.Tensor, sv: torch.Tensor,
rounds: int = 2, chunk: int = 1024, damping: float = 0.0):
"""
Free post-quantization polish: refit the stored per-channel fp16 scales (suh on the input side, svh on
the output side) to the quantized weight in the Hessian metric, holding the trellis fixed. With
W, Q (k, n) in the original basis: per output column c_n = (q_n^T H w_n) / (q_n^T H q_n) (closed form);
per input row r solves ((Q Q^T) o H) r = rowsum(Q o (H W)); alternated. Returns the rescaled weight_q
and the output error before / after, trace(E H E^T) / trace(W H W^T) = |X E|^2 / |X W|^2 over the
calibration rows. The fits use H as given (damped); `damping` is the share of H's mean diagonal that
finalize_capture_H added, taken back out of the diagonal for the reported error only
"""
dev = weight_q.device
H = H.to(dev)
if damping > 0.0:
d_off = damping * H.diagonal().mean().item()
H_err = H.clone()
H_err.diagonal().sub_(d_off)
else:
H_err = H
W = weight if weight.device == dev else None # large originals may sit on the CPU; stream columns
def cols(x, a, b):
return (x[:, a:b] if x.device == dev else x[:, a:b].to(dev)).float()
k, n = weight_q.shape
Q = weight_q
su = su.float().to(dev).view(k, 1)
sv = sv.float().to(dev).view(1, n)
# H W once (streamed over columns), reused by both steps; den = trace(W H W^T)
den = 0.0
for a in range(0, n, chunk):
b = min(a + chunk, n)
wc = cols(weight, a, b)
HWc = H @ wc
den += (wc * HWc).sum().item()
den_err = den
if H_err is not H:
den_err = 0.0
for a in range(0, n, chunk):
b = min(a + chunk, n)
wc = cols(weight, a, b)
den_err += (wc * (H_err @ wc)).sum().item()
def err(Qx):
e = 0.0
for a in range(0, n, chunk):
b = min(a + chunk, n)
E = cols(weight, a, b) - Qx[:, a:b]
e += (E * (H_err @ E)).sum().item()
return e / max(den_err, 1e-30)
e_before = err(Q)
for _ in range(rounds):
# output-side scales
num = torch.zeros(n, device = dev); dn = torch.zeros(n, device = dev)
for a in range(0, n, chunk):
b = min(a + chunk, n)
wc = cols(weight, a, b)
HWc = H @ wc
HQ = H @ Q[:, a:b]
num[a:b] = (Q[:, a:b] * HWc).sum(0)
dn[a:b] = (Q[:, a:b] * HQ).sum(0)
c = torch.where(dn > 1e-30, num / dn.clamp(min = 1e-30), torch.ones_like(num))
Q.mul_(c.view(1, n))
sv = sv * c.view(1, n)
# input-side scales
A = (Q @ Q.T) * H
b_ = torch.zeros(k, dtype=torch.float, device=dev)
for a in range(0, n, chunk):
b = min(a + chunk, n)
wc = cols(weight, a, b)
HWc = H @ wc
b_.add_((Q[:, a:b] * HWc).sum(1))
A.diagonal().add_(1e-6 * A.diagonal().mean())
r = torch.linalg.solve(A, b_)
r = torch.where(torch.isfinite(r) & (r > 0), r, torch.ones_like(r))
Q.mul_(r.view(k, 1))
su = su * r.view(k, 1)
e_after = err(Q)
return Q, su, sv, e_before, e_after
def quantize_exl3(
weight: torch.Tensor,
H_data: dict,
quant_args: dict,
return_weight_q: bool,
progress_str: str | None = None,
verbose: bool = False,
swap_to_device: torch.device | None = None,
save_reg: str = None
):
"""
:param weight:
Input tensor, row major shape (in_features, out_features)
:param H_data:
Dictionary of hessian tensor and related data, as collected by Linear wrapper class. May be reused between
linear layers with the same input (e.g. Q, K and V projections)
:param quant_args:
dict:
- K: bitrate
- seed: integer seed for random sign flips etc.
- sigma_reg: regularization factor
:param return_weight_q:
Return quantized weight
:param progress_str:
Show progress bar during quantization
:param verbose:
Dump extra stats
:param swap_to_device:
Device a CPU-swapped weight came from. Accepted for API compatibility; the work always
happens on quant_args["devices"][0] and the weight is moved there
:param save_reg:
Save regularized tensor as image to the provided path
:return:
tuple:
- quantized weight
- proxy error: trace(err @ H @ err.T) / (W @ H @ W.T)
- quantized and packed tensors
"""
progress_text = None if not progress_str else progress_str.replace("<step>", "Preparing")
with (ProgressBar(progress_text, 100) as pb):
assert weight.dtype == torch.float
tiles_k = weight.shape[0] // 16
if "seed" in quant_args:
torch.manual_seed(quant_args["seed"])
devices = quant_args["devices"]
device = torch.device(devices[0])
if weight.device != device:
weight = weight.to(device)
k, n = weight.shape
# Get H, LDL decomp. and input/output sign flips
q_fallback, H, L, su, H_diag = finalize_capture_H(H_data, quant_args, verbose)
if H is not None and H.is_cuda:
H = H.to(device)
if L is not None and L.is_cuda:
L = L.to(device)
if su.is_cuda:
su = su.to(device)
if H_diag is not None and H_diag.is_cuda:
H_diag = H_diag.to(device)
sv = (torch.randn(n, device = device).sign() + 1e-5).sign().to(torch.float).unsqueeze(0)
# Move stored L to CPU (if not already), move working L to device
if H_data["L"] is not None:
H_data["L"] = H_data["L"].cpu()
if L is not None:
L = L.to(device)
if verbose:
weight_copy = weight.cpu()
# The post-LDLQ scale refit compares against the original weight (regularize works in place)
weight_orig = weight.clone() if weight.numel() <= 5e8 else weight.cpu()
weight_r = weight
del weight
if verbose:
rms = block_rms_n(weight_r, dim = 0)
print(f" - input tensor rms: {rms:.6f}")
# Regularization
apply_out_scales, weight_r, g_scale, su, sv = regularize(
weight_r,
su,
sv,
quant_args,
verbose,
H_diag,
pb,
q_fallback = q_fallback
)
if save_reg:
save_tensor_image(weight_r, save_reg)
if verbose:
rms = weight_r.square().mean().sqrt()
print(f" - regularized rms: {rms:.6f}")
progress_text = None if not progress_str else progress_str.replace("<step>", "Quantizing")
pb.update(0)
pb.new_task(progress_text, tiles_k)
# Select device for work buffers (CPU is slower for small tensors but saves a lot of VRAM on big ones)
# TODO: Use pynvml or mem_get_info to predict whether CPU buffer is needed
if weight_r.numel() > 5e8:
weight_r = weight_r.cpu()
# The input scales regularize folded into the weight belong in the Hessian too (see scaled_ldl)
if not q_fallback and quant_args.get("scaled_hessian"):
L = scaled_ldl(original_H(H_data, device), su, quant_args)
# Quantize. An external output-side Hessian factor (YAQA-style Kronecker approximation, see ldlq_2hess)
# switches to error feedback along both axes
H_out_r = None
if not q_fallback and quant_args.get("H_out") is not None:
L_out, H_out_r = prepare_H_out(quant_args["H_out"], sv, quant_args, verbose, device)
weight_q, encoded_q = ldlq_2hess(weight_r.to(device), L, L_out, quant_args, pb)
del L, L_out
elif not q_fallback:
weight_q, encoded_q = ldlq(weight_r, L, quant_args, pb) #zxc
del L
else:
weight_q, encoded_q, mse_err = fallback_quant(weight_r, device, quant_args, pb) # zxc
pb.update(tiles_k)
# Metrics
if not q_fallback:
try:
E = weight_r - weight_q # may run on CPU
W = weight_r
Hd = H.to(device)
weight_r = None
E = E.to(device)
if H_out_r is not None:
# Two-sided proxy tr(E^T H_in E H_out) / tr(W^T H_in W H_out)
W = W.to(device)
num = ((Hd @ E) * (E @ H_out_r)).sum().item()
den = ((Hd @ W) * (W @ H_out_r)).sum().item()
if verbose:
print(f" - one-sided proxy: {block_trace(E, Hd) / max(block_trace(W, Hd), 1e-8):.6f}")
else:
num = block_trace(E, Hd)
E = None
W = W.to(device)
den = block_trace(W, Hd)
E = None
W = None
Hd = None
H_out_r = None
proxy_err = num / max(den, 1e-8)
except torch.OutOfMemoryError:
weight_r = None
E = None
W = None
Hd = None
proxy_err = -1.0
else:
proxy_err = mse_err
# free_mem()
weight_q = weight_q.to(device)
weight_q = preapply_had_l(weight_q, had_k)
weight_q *= su
weight_q = preapply_had_r(weight_q, had_n)
weight_q *= sv
# Polish: refit the per-channel fp16 scales to the quantized weight (free at inference)
if not q_fallback and H is not None and not quant_args.get("no_refit"):
H_orig = unrotate_H(H, H_data["su"])
weight_q, su_f, sv_f, e_before, e_after = refit_scales(weight_orig, weight_q, H_orig, su, sv,
damping = damping_share(quant_args))
del H_orig
su = su_f.view(-1, 1).to(su.dtype)
sv = sv_f.view(1, -1).to(sv.dtype)
if verbose:
print(f" - regularized-domain proxy err: {proxy_err:.6f}")
print(f" - scale refit: proxy err {e_before:.6f} -> {e_after:.6f} ({100 * (1 - e_after / max(e_before, 1e-30)):.2f}%)")
# Report the output error of what is stored, in the original basis. The regularized-domain
# figure weights every row and column alike, which is off by the input and output scales
proxy_err = e_after
del weight_orig
if verbose:
weight = weight_copy.to(device)
nmse = block_nmse(weight_q, weight)
print(f" - quant nmse: {nmse:.6f}")
# Compile packed tensor
suh = su.flatten().contiguous().to(dtype = torch.half, copy = True)
svh = sv.flatten().contiguous().to(dtype = torch.half, copy = True)
trellis = pack_trellis(encoded_q.to(device), quant_args) if not quant_args.get("no_pack") else encoded_q
out_tensors = {
# "scale": weight_scale.to(dtype = torch.float, copy = True),
# "su": pack_signs(su, quant_args),
"suh": suh,
# "sv": pack_signs(sv, quant_args),
"svh": svh,
"trellis": trellis,
}
out_tensors.update(codebook_markers(quant_args))
quant_args.update({
"apply_out_scales": apply_out_scales,
"g_scale": g_scale,
"q_fallback": q_fallback,
})
return weight_q, proxy_err, out_tensors
# Pinned staging buffers for _WeightStager, two per target device so an upload can be in
# flight while the next host-side copy fills the other buffer. Grown to the largest tensor
# seen; a worker thread owns its device, so per-device keying is contention-free
_stage_bufs = {}
_stage_bufs_lock = threading.Lock()
def _get_stage_buf(device_index: int, slot: int, nbytes: int):
key = (device_index, slot)
with _stage_bufs_lock:
buf = _stage_bufs.get(key)
if buf is None or buf[0].numel() < nbytes:
buf = (torch.empty(nbytes, dtype = torch.uint8, pin_memory = True),
torch.cuda.Event())
_stage_bufs[key] = buf
return buf
class _WeightStager:
"""
Stages CPU-swapped weights to the device through pinned buffers on a side stream, one
tensor ahead of consumption, casting to fp32 on the device. Replaces the former CPU-side
.float() + pageable synchronous upload: half the PCIe bytes (checkpoints are fp16/bf16),
pinned bandwidth, and the copy overlaps the previous tensor's regularization.
"""
def __init__(self, device: torch.device):
self.device = device
self.copy_stream = torch.cuda.Stream(device = device)
self.slot = 0
self.pending = {}
def prefetch(self, t: int, w: torch.Tensor):
if t in self.pending:
return
if w.is_cuda:
self.pending[t] = (w if w.device == self.device else w.to(self.device), None)
return
nbytes = w.numel() * w.element_size()
buf, reuse_ev = _get_stage_buf(self.device.index, self.slot, nbytes)
self.slot ^= 1
# The buffer's previous async upload must have drained before the host memcpy refills it
reuse_ev.synchronize()
pin = buf[:nbytes].view(w.dtype).view(w.shape)
pin.copy_(w)
with torch.cuda.stream(self.copy_stream):
# dev_w MUST be allocated on the copy stream.
dev_w = torch.empty(w.shape, dtype = w.dtype, device = self.device)
dev_w.copy_(pin, non_blocking = True)
reuse_ev.record(self.copy_stream)
ev = torch.cuda.Event()
ev.record(self.copy_stream)
self.pending[t] = (dev_w, ev)
def get(self, t: int, w: torch.Tensor) -> torch.Tensor:
self.prefetch(t, w)
dev_w, ev = self.pending.pop(t)
if ev is not None:
torch.cuda.current_stream(self.device).wait_event(ev)
# The block belongs to the copy stream's pool; mark its use on the consuming stream
# so its eventual free isn't recycled under a later async upload
dev_w.record_stream(torch.cuda.current_stream(self.device))
return dev_w.float() if dev_w.dtype != torch.float else dev_w
def quantize_exl3_batch(
weights: list[torch.Tensor],
H_datas: list[dict],
quant_args_list: list[dict],
progress_str: str | None = None,
verbose: bool = False,
):
"""
Quantize a group of same-K linears together, batching the scale search and the LDLQ recursion so the
quantize_tiles kernel sees large batches instead of one small tensor's worth of tiles per step.
Two layouts, chosen by whether the H_data dicts are the same object:
- shared Hessian (e.g. all gate/up projections of a block-sparse MLP, or fused q/k/v): the regularized
weights are concatenated along out_features and run through a single ldlq() pass with the shared L.
This is exact: LDLQ treats columns independently given L, and su (drawn per qmap) is identical
- distinct Hessians over same-shape tensors (e.g. per-expert down projections): weights are stacked and
run through ldlq_batched() with per-tensor L
RNG is re-seeded per tensor (finalize's su draw on first call for the qmap, then the per-tensor sv draw),
so per-tensor results should match quantize_exl3() in the wider matmuls. Tensors whose Hessian fails to
finalize (q_fallback) are routed through quantize_exl3() individually.
:return:
list of (proxy_err, out_tensors) per input tensor; quant_args_list entries are updated in place like
quantize_exl3 updates quant_args
"""
n_t = len(weights)
qa0 = quant_args_list[0]
devices = qa0["devices"]
device = torch.device(devices[0])
shared_H = all(hd is H_datas[0] for hd in H_datas)
assert shared_H or all(w.shape == weights[0].shape for w in weights)
assert all(w.shape[0] == weights[0].shape[0] for w in weights)
assert all(qa["K"] == qa0["K"] for qa in quant_args_list)
size_k = weights[0].shape[0]
tiles_k = size_k // 16
results: list = [None] * n_t
progress_text = None if not progress_str else progress_str.replace("<step>", "Preparing")
with ProgressBar(progress_text, 100) as pb:
# Finalize Hessians, replicating the serial path's per-tensor RNG stream. Fallback tensors are
# handled individually by quantize_exl3
finalized = []
batch_idx = []
for t in range(n_t):
qa = quant_args_list[t]
if "seed" in qa:
torch.manual_seed(qa["seed"])
# Unswap and factor on the device doing the quantization rather than on the capture
# device, which would otherwise factor every tensor of the layer for all threads
if "H_swap_device" in H_datas[t]:
H_datas[t]["H_swap_device"] = device
q_fallback, H, L, su, H_diag = finalize_capture_H(H_datas[t], qa, verbose)
if q_fallback:
finalized.append(None)
else:
finalized.append((H, L, su, H_diag))
batch_idx.append(t)
for t in range(n_t):
if finalized[t] is None:
results[t] = quantize_exl3(weights[t].float(), H_datas[t], quant_args_list[t], False, None, verbose)[1:]
if not batch_idx:
return results
# Regularize each tensor with the scale search deferred, then search all scales in one
# batch. Weights arrive in checkpoint precision on the CPU; the stager uploads tensor
# t+1 while tensor t regularizes and casts to fp32 on the device
stager = _WeightStager(device)
stager.prefetch(batch_idx[0], weights[batch_idx[0]])
regs = {}
origs = {}
for bi, t in enumerate(batch_idx):
qa = quant_args_list[t]
if "seed" in qa:
torch.manual_seed(qa["seed"])
H, L, su, H_diag = finalized[t]
weight = stager.get(t, weights[t])
if bi + 1 < len(batch_idx):
stager.prefetch(batch_idx[bi + 1], weights[batch_idx[bi + 1]])
if su.is_cuda:
su = su.to(device)
if H_diag is not None and H_diag.is_cuda:
H_diag = H_diag.to(device)
sv = (torch.randn(weight.shape[1], device = device).sign() + 1e-5).sign().to(torch.float).unsqueeze(0)
# Original kept for the post-LDLQ scale refit (regularize works in place)
origs[t] = weight.clone() if weight.numel() <= 5e7 else weight.cpu()
apply_out_scales, weight_r, _, su, sv = regularize(
weight, su, sv, qa, verbose, H_diag, None, skip_g_scale = True)
regs[t] = [weight_r, su, sv, apply_out_scales]
weights[t] = None
samples = [sample_scale_tiles(regs[t][0]) * ldlq_drift(qa0["K"]) for t in batch_idx]
scales = g_scale_search_batch(samples, qa0)
del samples
g_scales = {}
for t, (g_scale, _) in zip(batch_idx, scales):
regs[t][0] *= g_scale
regs[t][1] /= g_scale
g_scales[t] = g_scale
pb.update(100)
progress_text = None if not progress_str else progress_str.replace("<step>", "Quantizing")
pb.new_task(progress_text, tiles_k)
# Quantize
scaled = qa0.get("scaled_hessian")
# A shared H_data serves many groups spread over several device threads; cache the
# device copies of L (and H, below) in the dict so each device pays the transfer
# once per layer instead of once per group. A benign race can duplicate a copy;
# the loser's tensor is simply collected
dev_cache = H_datas[batch_idx[0]].setdefault("dev_cache", {}) if shared_H else None
if shared_H and not scaled:
L = dev_cache.get(("L", device.index))
if L is None:
L = finalized[batch_idx[0]][1].to(device)
dev_cache[("L", device.index)] = L
widths = [regs[t][0].shape[1] for t in batch_idx]
weight_r_cat = torch.cat([regs[t][0] for t in batch_idx], dim = 1)
for t in batch_idx:
regs[t][0] = None
weight_q_cat, encoded_cat = ldlq(weight_r_cat, L, qa0, pb)
L = None # device copy stays cached in H_data for the layer's remaining groups
weight_rs = list(torch.split(weight_r_cat, widths, dim = 1))
weight_qs = list(torch.split(weight_q_cat, widths, dim = 1))
encodeds = list(torch.split(encoded_cat, [w // 16 for w in widths], dim = 1))
else:
# Per-tensor L: same-shape tensors stack into one batched pass. With scaled Hessians a shared
# H no longer means a shared L, so the concatenated layout above is replaced by stacks per
# shape (gate/up of a group of experts are one stack; q/k/v fall into two). Each stack's L
# is factored here and released with the pass, so the extra memory is one stack's worth
by_shape = {}
for t in batch_idx:
by_shape.setdefault(tuple(regs[t][0].shape), []).append(t)
outs = {}
for ts in by_shape.values():
if scaled:
Ls = torch.stack([
scaled_ldl(original_H(H_datas[t], device, cache = shared_H), regs[t][1], quant_args_list[t])
for t in ts
])
else:
Ls = torch.stack([finalized[t][1].to(device) for t in ts])
weight_r_stack = torch.stack([regs[t][0] for t in ts])
for t in ts:
regs[t][0] = None
weight_q_stack, encoded_stack = ldlq_batched(weight_r_stack, Ls, qa0, pb if len(by_shape) == 1 else None)
del Ls
for t, wr, wq, enc in zip(ts, weight_r_stack.unbind(0), weight_q_stack.unbind(0), encoded_stack.unbind(0)):
outs[t] = (wr, wq, enc)
weight_rs = [outs[t][0] for t in batch_idx]
weight_qs = [outs[t][1] for t in batch_idx]
encodeds = [outs[t][2] for t in batch_idx]
del outs
pb.update(tiles_k)
# Like quantize_exl3, park the stored factor on the CPU once the recursion has run: a layer
# of distinct per-expert factors would otherwise accumulate on the device until the layer
# is done. Tensors sharing the Hessian keep the per-device copy cached above
for t in batch_idx:
hd = H_datas[t]
if hd.get("L") is not None and hd["L"].is_cuda:
hd["L"] = hd["L"].cpu()
# Per-tensor metrics and packing
Hd = None
for bi, t in enumerate(batch_idx):
qa = quant_args_list[t]
_, su, sv, apply_out_scales = regs[t]
if shared_H:
if Hd is None:
Hd = dev_cache.get(("H", device.index))
if Hd is None:
Hd = finalized[t][0].to(device)
dev_cache[("H", device.index)] = Hd
else:
Hd = finalized[t][0].to(device)
try:
E = weight_rs[bi] - weight_qs[bi]
num = block_trace(E, Hd)
E = None
den = block_trace(weight_rs[bi], Hd)
proxy_err = num / max(den, 1e-8)
except torch.OutOfMemoryError:
E = None
proxy_err = -1.0
weight_rs[bi] = None
# Polish: refit the per-channel fp16 scales to the quantized weight (free at inference)
if not qa.get("no_refit"):
wq = weight_qs[bi].to(device)
wq = preapply_had_l(wq, had_k)
wq *= su
wq = preapply_had_r(wq, had_n)
wq *= sv
H_orig = unrotate_H(Hd, H_datas[t]["su"])
_, su_f, sv_f, e_before, e_after = refit_scales(origs[t], wq, H_orig, su, sv,
damping = damping_share(qa))
del wq, H_orig
su = su_f.view(-1, 1).to(su.dtype)
sv = sv_f.view(1, -1).to(sv.dtype)
if verbose:
print(f" - scale refit: proxy err {e_before:.6f} -> {e_after:.6f} ({100 * (1 - e_after / max(e_before, 1e-30)):.2f}%)")
proxy_err = e_after
origs[t] = None
weight_qs[bi] = None
suh = su.flatten().contiguous().to(dtype = torch.half, copy = True)
svh = sv.flatten().contiguous().to(dtype = torch.half, copy = True)
trellis = pack_trellis(encodeds[bi].contiguous(), qa)
encodeds[bi] = None
out_tensors = {
"suh": suh,
"svh": svh,
"trellis": trellis,
}
out_tensors.update(codebook_markers(qa))
qa.update({
"apply_out_scales": apply_out_scales,
"g_scale": g_scales[t],
"q_fallback": False,
})
results[t] = (proxy_err, out_tensors)
return results