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