Blame for cuda_backend.py

"""CUDA GPU-accelerated backend for HPC-Quantize V2 (PyTorch CUDA).

Executes the entire HPC-Quantize V2 pipeline in batched FP32 GPU tensor kernels:

* Batched GGML seed solvers (``make_qkx2``, ``make_qp``) and 5-step WLS;
* Batched Geometry-Preserving Weight Analogs (Sub-block Iso-Moment,
  Vesica Gram-Preserving, and Cross-Sub-block Scale-Lattice);
* Batched Coupled Orthogonal-Basis Bilinear Q2_K Factorization with rank-1
  1-D integer refinement and 36-seed candidate lattice evaluation;
* Sequential 6-state Sieve & 36-state Viterbi DP + vectorized row DC carry;
* Batched Algebraic Phase-4 Coupled Assembly, Vesica+DC augmented 2x2 Cramer
  (d, dmin) solve, O(1) quadratic fp16-ULP polish, closed-form 1-D quadratic
  Vesica/DC code projection, and row-parallel GPU DC residual nudge.
"""

from __future__ import annotations

import time
import numpy as np

try:
    import torch
except ImportError:  # pragma: no cover
    torch = None  # type: ignore[assignment]

from .config import HPCConfig
from .factored import MULTS_D, MULTS_M, N_CAND_D, N_CAND_M, TOTAL_CANDIDATES
from .pipeline import QuantizedTensor, TensorStats, _phase3_select
from .q2k import NS, QK_K, Q2KBlock, pin_q2k_dc_residual
from .seeds import seed_clamp_value
from .spectral import fold_weights


def cuda_is_available() -> bool:
    return bool(torch is not None and torch.cuda.is_available())


def cuda_device_name(device: str = "cuda") -> str:
    if not cuda_is_available():
        return "CPU"
    idx = torch.device(device).index or 0
    return str(torch.cuda.get_device_name(idx))


def _nearest_int_t(val: "torch.Tensor") -> "torch.Tensor":
    """Bit-exact GGML round-to-nearest-even on GPU."""
    shifted = val.to(torch.float32) + 12582912.0
    bits = shifted.view(torch.int32)
    return ((bits & 0x007FFFFF) - 0x00400000).to(torch.int64)


def _round_fp16_t(val: "torch.Tensor") -> "torch.Tensor":
    return val.to(torch.float16).to(torch.float32)


def _to_fp16_bits_t(val: "torch.Tensor") -> "torch.Tensor":
    return val.to(torch.float16).view(torch.int16).to(torch.int32) & 0xFFFF


def _from_fp16_bits_t(bits: "torch.Tensor") -> "torch.Tensor":
    return bits.to(torch.int16).view(torch.float16).to(torch.float32)


def _codes_from_params_t(
    X: "torch.Tensor",
    d: "torch.Tensor",
    m: "torch.Tensor",
    ls: "torch.Tensor",
    lm: "torch.Tensor",
) -> "torch.Tensor":
    d_sub = (d[:, None].to(torch.float32) * ls.to(torch.float32)).repeat_interleave(16, dim=1)
    m_sub = (m[:, None].to(torch.float32) * lm.to(torch.float32)).repeat_interleave(16, dim=1)
    safe = d_sub >= 1.0e-15
    safe_d = torch.where(safe, d_sub, torch.ones_like(d_sub))
    q = _nearest_int_t((X.to(torch.float32) + m_sub) / safe_d).clamp(0, 3)
    return torch.where(safe, q, torch.zeros_like(q)).to(torch.uint8)


def _dequantize_t(
    d: "torch.Tensor",
    m: "torch.Tensor",
    ls: "torch.Tensor",
    lm: "torch.Tensor",
    codes: "torch.Tensor",
) -> "torch.Tensor":
    d_sub = (d[:, None].to(torch.float32) * ls.to(torch.float32)).repeat_interleave(16, dim=1)
    m_sub = (m[:, None].to(torch.float32) * lm.to(torch.float32)).repeat_interleave(16, dim=1)
    return d_sub * codes.to(torch.float32) - m_sub


def _objective_t(
    X: "torch.Tensor",
    W: "torch.Tensor",
    d: "torch.Tensor",
    m: "torch.Tensor",
    ls: "torch.Tensor",
    lm: "torch.Tensor",
    codes: "torch.Tensor",
    carry: "torch.Tensor | None",
    cfg: HPCConfig,
) -> tuple["torch.Tensor", "torch.Tensor"]:
    xh = _dequantize_t(d, m, ls, lm, codes)
    e = X - xh
    sse = torch.sum(W * e * e, dim=1, dtype=torch.float32)
    if not cfg.use_spectral or (cfg.dc_lambda == 0.0 and cfg.vw_lambda == 0.0):
        return sse, sse
    fw = fold_weights(QK_K, cfg)
    pen = torch.zeros_like(sse)
    src = e
    m_len = QK_K // 2
    lvl = 0
    while m_len >= 1:
        dst = src[:, :m_len] + src[:, m_len : 2 * m_len]
        if m_len == 1 and carry is not None:
            dst = dst + carry[:, None].to(torch.float32)
        if fw[lvl] != 0.0:
            pen = pen + float(fw[lvl]) * torch.sum(dst * dst, dim=1)
        src = dst
        m_len //= 2
        lvl += 1
    return sse + pen, sse


def _make_qkx2_cuda(x: "torch.Tensor", w: "torch.Tensor", nmax: int = 3):
    x = x.to(torch.float32)
    w = w.to(torch.float32)
    xmin = torch.clamp(torch.min(x, dim=1).values, max=0.0)
    xmax = torch.max(x, dim=1).values
    span = xmax - xmin
    degenerate = (xmax == xmin) | (span <= 1.0e-30)
    sum_w = torch.sum(w, dim=1)
    sum_x = torch.sum(w * x, dim=1)
    safe_span = torch.where(span > 1.0e-30, span, torch.ones_like(span))
    iscale = nmax / safe_span
    scale = 1.0 / iscale
    codes = _nearest_int_t(iscale[:, None] * (x - xmin[:, None])).clamp(0, nmax)
    best_mad = torch.sum(w * torch.abs(scale[:, None] * codes.to(torch.float32) + xmin[:, None] - x), dim=1)
    best_scale = scale.clone()
    best_min = xmin.clone()
    best_codes = codes.clone()

    steps = torch.arange(16, dtype=torch.float32, device=x.device)[None, :]  # (1, 16)
    try_iscale = (-0.5 + 0.1 * steps + nmax) / safe_span[:, None]            # (N, 16)
    aux = _nearest_int_t(try_iscale[:, :, None] * (x - xmin[:, None])[:, None, :]).clamp(0, nmax)
    aux_f = aux.to(torch.float32)
    w_e = w[:, None, :]
    x_e = x[:, None, :]
    sl = torch.sum(w_e * aux_f, dim=2)
    sl2 = torch.sum(w_e * aux_f * aux_f, dim=2)
    sxl = torch.sum(w_e * aux_f * x_e, dim=2)
    det = sum_w[:, None] * sl2 - sl * sl
    valid = det > 1.0e-30
    s_det = torch.where(valid, det, torch.ones_like(det))
    this_scale = (sum_w[:, None] * sxl - sum_x[:, None] * sl) / s_det
    this_min = (sl2 * sum_x[:, None] - sl * sxl) / s_det
    pos_min = this_min > 0.0
    this_min = torch.where(pos_min, torch.zeros_like(this_min), this_min)
    s_sl2 = torch.where(sl2 > 1.0e-30, sl2, torch.ones_like(sl2))
    this_scale = torch.where(pos_min & (sl2 > 1.0e-30), sxl / s_sl2, this_scale)
    mad = torch.sum(w_e * torch.abs(this_scale[:, :, None] * aux_f + this_min[:, :, None] - x_e), dim=2)
    mad = torch.where(valid, mad, torch.full_like(mad, float("inf")))
    min_mad, best_k = torch.min(mad, dim=1)
    take = min_mad < best_mad
    rows = torch.arange(x.shape[0], device=x.device)
    best_scale = torch.where(take, this_scale[rows, best_k], best_scale)
    best_min = torch.where(take, this_min[rows, best_k], best_min)
    best_codes = torch.where(take[:, None], aux[rows, best_k], best_codes)

    best_scale = torch.where(degenerate, torch.zeros_like(best_scale), best_scale)
    best_codes = torch.where(degenerate[:, None], torch.zeros_like(best_codes), best_codes)
    best_min = torch.where(degenerate, xmin, best_min)
    return best_scale, -best_min, best_codes.to(torch.uint8)


def _make_qp_cuda(values: "torch.Tensor", weights: "torch.Tensor", nmax: int = 15):
    x = values.to(torch.float32)
    w = weights.to(torch.float32)
    xmax = torch.max(x, dim=1).values
    dead = xmax < 1.0e-15
    safe_max = torch.where(dead, torch.ones_like(xmax), xmax)

    steps = torch.arange(-4, 5, dtype=torch.float32, device=x.device)[None, :]  # (1, 9)
    t_is = (0.1 * steps + nmax) / safe_max[:, None]                             # (N, 9)
    t_sc = 1.0 / t_is
    t_cd = _nearest_int_t(t_is[:, :, None] * x[:, None, :]).clamp(0, nmax).to(torch.float32)
    mse = torch.sum(w[:, None, :] * (x[:, None, :] - t_sc[:, :, None] * t_cd) ** 2, dim=2)
    best_k = torch.argmin(mse, dim=1)
    rows = torch.arange(x.shape[0], device=x.device)
    best_iscale = t_is[rows, best_k]

    codes = _nearest_int_t(best_iscale[:, None] * x).clamp(0, nmax)
    cf = codes.to(torch.float32)
    sumlx = torch.sum(w * x * cf, dim=1)
    suml2 = torch.sum(w * cf * cf, dim=1)

    for _ in range(3):
        changed = False
        for i in range(x.shape[1]):
            wi = w[:, i]
            xi = x[:, i]
            li = codes[:, i].to(torch.float32)
            slx = sumlx - wi * xi * li
            sl2 = suml2 - wi * li * li
            valid = (slx > 0.0) & (sl2 > 0.0)
            s_slx = torch.where(valid, slx, torch.ones_like(slx))
            new_l = _nearest_int_t(xi * sl2 / s_slx).clamp(0, nmax)
            nl_f = new_l.to(torch.float32)
            trial = valid & (new_l != codes[:, i])
            slx2 = slx + wi * xi * nl_f
            sl2_2 = sl2 + wi * nl_f * nl_f
            better = trial & (slx2 * slx2 * suml2 > sumlx * sumlx * sl2_2)
            if torch.any(better):
                codes[:, i] = torch.where(better, new_l, codes[:, i])
                sumlx = torch.where(better, slx2, sumlx)
                suml2 = torch.where(better, sl2_2, suml2)
                changed = True
        if not changed:
            break

    s_l2 = torch.where(suml2 > 1.0e-30, suml2, torch.ones_like(suml2))
    d = torch.where((suml2 > 1.0e-30) & (~dead), sumlx / s_l2, torch.zeros_like(sumlx))
    codes = torch.where(dead[:, None], torch.zeros_like(codes), codes)
    return d, codes.to(torch.uint8)


def _block_clamp_cuda(X: "torch.Tensor", cfg: HPCConfig) -> "torch.Tensor":
    """Per-block kurtosis-gated outlier clamp thresholds on GPU (shape (B, 1, 1))."""
    b2 = X * X
    mean_sq = torch.mean(b2, dim=1)
    mean_4 = torch.mean(b2 * b2, dim=1)
    kurt = torch.where(mean_sq > 1.0e-30, mean_4 / torch.clamp(mean_sq * mean_sq, min=1.0e-30), torch.full_like(mean_sq, 3.0))
    sigma = torch.sqrt(torch.clamp(mean_sq, min=0.0))
    mult = torch.where(
        kurt <= cfg.clamp_kurtosis_off,
        torch.full_like(kurt, 1.0e9),
        torch.where(
            kurt <= cfg.clamp_kurtosis_soft,
            torch.full_like(kurt, float(cfg.clamp_sigma_soft)),
            torch.full_like(kurt, float(cfg.clamp_sigma_hard)),
        ),
    )
    clamp = torch.where(mean_sq > 1.0e-30, sigma * mult, torch.full_like(sigma, float("inf")))
    return clamp[:, None, None]


def _build_seeds_cuda(X: "torch.Tensor", W: "torch.Tensor | None", cfg: HPCConfig):
    B = X.shape[0]
    blocks = X.reshape(B, 16, 16)
    per_blk_clamp = _block_clamp_cuda(X, cfg)
    clipped = torch.clamp(blocks, min=-per_blk_clamp, max=per_blk_clamp)
    if W is None:
        sigma_sq = torch.mean(blocks * blocks, dim=(1, 2), keepdim=True)
        weights = torch.sqrt(sigma_sq + clipped * clipped)
    else:
        weights = W.reshape(B, 16, 16)
    scales, mins, codes = _make_qkx2_cuda(clipped.reshape(-1, 16), weights.reshape(-1, 16), 3)
    scales = scales.reshape(B, 16)
    mins = mins.reshape(B, 16)
    codes = codes.reshape(B, 256)
    sw = torch.sum(weights, dim=2)
    d, ls = _make_qp_cuda(scales, sw, 15)
    dmin, lm = _make_qp_cuda(mins, sw, 15)
    return {
        "scales": scales,
        "mins": mins,
        "codes": codes,
        "sw": sw,
        "d": d,
        "dmin": dmin,
        "ls": ls,
        "lm": lm,
        "clamp": per_blk_clamp.reshape(B, 1),
    }


def _phase2_wls_cuda(X: "torch.Tensor", W: "torch.Tensor", seeds: dict):
    per_blk_clamp = seeds["clamp"]
    clipped = torch.clamp(X, min=-per_blk_clamp, max=per_blk_clamp)
    base_d = seeds["d"].clone()
    base_m = seeds["dmin"].clone()
    ls = seeds["ls"].clone()
    lm = seeds["lm"].clone()
    scales = seeds["scales"]
    mins = seeds["mins"]

    for _ in range(5):
        codes = _codes_from_params_t(clipped, base_d, base_m, ls, lm).to(torch.float32)
        A = ls.to(torch.float32).repeat_interleave(16, dim=1) * codes
        Bm = lm.to(torch.float32).repeat_interleave(16, dim=1)
        Saa = torch.sum(W * A * A, dim=1)
        Sab = torch.sum(W * A * Bm, dim=1)
        Sbb = torch.sum(W * Bm * Bm, dim=1)
        Sxa = torch.sum(W * clipped * A, dim=1)
        Sxb = torch.sum(W * clipped * Bm, dim=1)
        det = Saa * Sbb - Sab * Sab
        valid = det > 1.0e-30
        s_det = torch.where(valid, det, torch.ones_like(det))
        d_new = (Sbb * Sxa - Sab * Sxb) / s_det
        m_new = (Sab * Sxa - Saa * Sxb) / s_det
        upd_d = valid & (d_new > 0.0) & torch.isfinite(d_new)
        upd_m = valid & (m_new > 0.0) & torch.isfinite(m_new)
        base_d = torch.where(upd_d, _round_fp16_t(d_new), base_d)
        base_m = torch.where(upd_m, _round_fp16_t(m_new), base_m)
        sd = torch.where(base_d > 1.0e-15, base_d, torch.ones_like(base_d))
        sm = torch.where(base_m > 1.0e-15, base_m, torch.ones_like(base_m))
        ls = torch.where(
            base_d[:, None] > 1.0e-15,
            _nearest_int_t(scales / sd[:, None]).clamp(0, 15),
            torch.zeros_like(ls, dtype=torch.int64),
        ).to(torch.uint8)
        lm = torch.where(
            base_m[:, None] > 1.0e-15,
            _nearest_int_t(mins / sm[:, None]).clamp(0, 15),
            torch.zeros_like(lm, dtype=torch.int64),
        ).to(torch.uint8)

    return base_d, base_m


def _factor_1d_cuda(
    V: "torch.Tensor",
    W: "torch.Tensor",
    nmax: int = 15,
    iters: int = 5,
    refine: bool = True,
) -> tuple["torch.Tensor", "torch.Tensor"]:
    """Batched FP32 1-D factorization across B blocks on GPU with rank-1 sweep."""
    V = torch.clamp(V.to(torch.float32), min=0.0)
    W = torch.clamp(W.to(torch.float32), min=0.0)
    B, n = V.shape
    vmax = torch.max(V, dim=1).values
    active = vmax >= 1.0e-30
    d = torch.where(active, vmax / nmax, torch.zeros_like(vmax))
    sd = torch.where(d > 1.0e-30, d, torch.ones_like(d))
    codes = torch.where(
        active[:, None] & (d[:, None] > 1.0e-30),
        _nearest_int_t(V / sd[:, None]).clamp(0, nmax),
        torch.zeros((B, n), dtype=torch.int64, device=V.device),
    )
    for _ in range(iters):
        cf = codes.to(torch.float32)
        den = torch.sum(W * cf * cf, dim=1)
        valid = active & (den > 1.0e-30)
        s_den = torch.where(valid, den, torch.ones_like(den))
        dn = torch.sum(W * V * cf, dim=1) / s_den
        d = torch.where(valid & (dn > 0.0), dn, d)
        sd = torch.where(d > 1.0e-30, d, torch.ones_like(d))
        codes = torch.where(
            active[:, None] & (d[:, None] > 1.0e-30),
            _nearest_int_t(V / sd[:, None]).clamp(0, nmax),
            torch.zeros_like(codes),
        )

    if refine:
        opts = torch.arange(nmax + 1, dtype=torch.float32, device=V.device)[None, :]
        opts_sq = opts * opts
        wv = W * V
        cf = codes.to(torch.float32)
        sum_lv = torch.sum(wv * cf, dim=1)
        sum_l2 = torch.sum(W * cf * cf, dim=1)
        s_l2 = torch.where(sum_l2 > 1.0e-30, sum_l2, torch.ones_like(sum_l2))
        best_gain = torch.where(sum_l2 > 1.0e-30, (sum_lv * sum_lv) / s_l2, torch.zeros_like(sum_lv))
        rows = torch.arange(B, device=V.device)

        for i in range(n):
            wi = W[:, i : i + 1]
            wvi = wv[:, i : i + 1]
            li = codes[:, i : i + 1].to(torch.float32)
            slv = (sum_lv[:, None] - wvi * li) + wvi * opts
            sl2 = (sum_l2[:, None] - wi * (li * li)) + wi * opts_sq
            valid = (wi > 0.0) & (sl2 > 1.0e-30) & (slv > 0.0)
            gains = torch.where(valid, (slv * slv) / torch.where(valid, sl2, torch.ones_like(sl2)), -torch.ones_like(slv))
            max_g, k_idx = torch.max(gains, dim=1)
            improve = active & (k_idx != codes[:, i]) & (max_g > best_gain + 1.0e-12)
            if torch.any(improve):
                best_gain = torch.where(improve, max_g, best_gain)
                sum_lv = torch.where(improve, slv[rows, k_idx], sum_lv)
                sum_l2 = torch.where(improve, sl2[rows, k_idx], sum_l2)
                codes[:, i] = torch.where(improve, k_idx, codes[:, i])

        s_l2 = torch.where(sum_l2 > 1.0e-30, sum_l2, torch.ones_like(sum_l2))
        d = torch.where(active & (sum_l2 > 1.0e-30), sum_lv / s_l2, d)

    return d, codes.to(torch.uint8)


def _analogs_cuda(X: "torch.Tensor", W: "torch.Tensor") -> list["torch.Tensor"]:
    B = X.shape[0]
    xb = X.reshape(B, NS, 16)
    wb = W.reshape(B, NS, 16)
    wsum = torch.clamp(torch.sum(wb, dim=2, keepdim=True), min=1.0e-30)
    mu = torch.sum(wb * xb, dim=2, keepdim=True) / wsum
    centered = xb - mu
    var = torch.sum(wb * centered * centered, dim=2, keepdim=True) / wsum
    sigma = torch.sqrt(torch.clamp(var, min=1.0e-30))
    active = sigma >= 1.0e-15
    z = centered / sigma
    q_ideal = torch.linspace(-1.35, 1.35, 16, dtype=torch.float32, device=X.device)
    order = torch.argsort(z, dim=2, stable=True)
    z_quant = torch.empty_like(z)
    z_quant.scatter_(2, order, q_ideal[None, None, :].expand_as(z))

    analogs = []
    for tau, beta, q_blend in ((1.60, 0.25, 0.30), (2.15, 0.15, 0.15)):
        w_rel = torch.clamp(wb / (wsum / 16.0), min=1.0e-12) ** beta
        z_comp = tau * torch.tanh((z * w_rel) / tau)
        z_a = (1.0 - q_blend) * z_comp + q_blend * z_quant
        mu_a = torch.sum(wb * z_a, dim=2, keepdim=True) / wsum
        z_a = z_a - mu_a
        var_a = torch.sum(wb * z_a * z_a, dim=2, keepdim=True) / wsum
        scale = torch.where(var_a > 1.0e-30, torch.sqrt(var / torch.clamp(var_a, min=1.0e-30)), torch.zeros_like(var_a))
        sub_a = torch.where(active, mu + z_a * scale, xb)
        orig_min = torch.min(xb, dim=2, keepdim=True).values
        new_min = torch.min(sub_a, dim=2, keepdim=True).values
        shift = torch.where((orig_min <= 0.0) & (new_min > 0.0), new_min, torch.zeros_like(new_min))
        analogs.append((sub_a - shift).reshape(B, QK_K))

    return analogs


def _sub_ls_cuda(X: "torch.Tensor", W: "torch.Tensor", codes: "torch.Tensor"):
    B = X.shape[0]
    q = codes.to(torch.float32).reshape(B, NS, 16)
    xb = X.reshape(B, NS, 16)
    wb = W.reshape(B, NS, 16)
    Sqq = torch.sum(wb * q * q, dim=2)
    Swq = torch.sum(wb * q, dim=2)
    Sw = torch.sum(wb, dim=2)
    Sqx = torch.sum(wb * q * xb, dim=2)
    Sx = torch.sum(wb * xb, dim=2)
    det = Sqq * Sw - Swq * Swq
    valid = det > 1.0e-30
    s_det = torch.where(valid, det, torch.ones_like(det))
    s_sw = torch.where(Sw > 1.0e-30, Sw, torch.ones_like(Sw))
    sj = torch.clamp(torch.where(valid, (Sqx * Sw - Swq * Sx) / s_det, torch.zeros_like(det)), min=0.0)
    tj = torch.where(valid, (sj * Swq - Sx) / s_sw, torch.zeros_like(det))
    bad_tj = tj < 0.0
    tj = torch.where(bad_tj, torch.zeros_like(tj), tj)
    s_sqq = torch.where(Sqq > 1.0e-30, Sqq, torch.ones_like(Sqq))
    fallback = torch.clamp(torch.where(Sqq > 1.0e-30, Sqx / s_sqq, torch.zeros_like(Sqq)), min=0.0)
    sj = torch.where(bad_tj, fallback, sj)
    return sj, tj, Sqq, Swq, Sw, Sx


def _factor_candidates_cuda(
    X: "torch.Tensor",
    W: "torch.Tensor",
    scales: "torch.Tensor",
    mins: "torch.Tensor",
    base_d: "torch.Tensor",
    base_m: "torch.Tensor",
    cfg: HPCConfig,
) -> np.ndarray:
    """Vectorized FP32 36-seed candidate lattice across B blocks on GPU."""
    B = X.shape[0]
    device = X.device
    nd, nm = int(cfg.n_cand_d), int(cfg.n_cand_m)
    di_idx = [a if nd == N_CAND_D else (a * N_CAND_D) // nd for a in range(nd)]
    mi_idx = [b if nm == N_CAND_M else (b * N_CAND_M) // nm for b in range(nm)]
    md = torch.tensor(MULTS_D[di_idx], dtype=torch.float32, device=device)
    mm = torch.tensor(MULTS_M[mi_idx], dtype=torch.float32, device=device)

    ds_mat = _round_fp16_t(base_d[:, None] * md[None, :])
    ms_mat = _round_fp16_t(base_m[:, None] * mm[None, :])
    ds = ds_mat.repeat_interleave(nm, dim=1)
    ms = ms_mat.repeat(1, nd)
    K = ds.shape[1]

    sc = scales[:, None, :]
    mn = mins[:, None, :]
    s_ds = torch.where(ds[:, :, None] > 1.0e-30, ds[:, :, None], torch.ones_like(ds[:, :, None]))
    s_ms = torch.where(ms[:, :, None] > 1.0e-30, ms[:, :, None], torch.ones_like(ms[:, :, None]))
    ls0 = torch.where(
        ds[:, :, None] > 1.0e-30,
        _nearest_int_t(sc / s_ds).clamp(0, 15),
        torch.zeros((B, K, NS), dtype=torch.int64, device=device),
    ).to(torch.float32)
    lm0 = torch.where(
        ms[:, :, None] > 1.0e-30,
        _nearest_int_t(mn / s_ms).clamp(0, 15),
        torch.zeros((B, K, NS), dtype=torch.int64, device=device),
    ).to(torch.float32)

    X_e = X[:, None, :]
    W_e = W[:, None, :]
    fw = fold_weights(QK_K, cfg)

    def _eval_k(ls_k: "torch.Tensor", lm_k: "torch.Tensor"):
        d_sub = (ds[:, :, None] * ls_k).repeat_interleave(16, dim=2)
        m_sub = (ms[:, :, None] * lm_k).repeat_interleave(16, dim=2)
        safe = d_sub >= 1.0e-15
        sd = torch.where(safe, d_sub, torch.ones_like(d_sub))
        q = torch.where(
            safe,
            _nearest_int_t((X_e + m_sub) / sd).clamp(0, 3),
            torch.zeros((B, K, QK_K), dtype=torch.int64, device=device),
        ).to(torch.float32)
        e = X_e - (d_sub * q - m_sub)
        sse = torch.sum(W_e * e * e, dim=2)
        if cfg.use_spectral and (cfg.dc_lambda != 0.0 or cfg.vw_lambda != 0.0):
            pen = torch.zeros_like(sse)
            src = e
            m_len = QK_K // 2
            lvl = 0
            while m_len >= 1:
                dst = src[:, :, :m_len] + src[:, :, m_len : 2 * m_len]
                if fw[lvl] != 0.0:
                    pen = pen + float(fw[lvl]) * torch.sum(dst * dst, dim=2)
                src = dst
                m_len //= 2
                lvl += 1
            return sse + pen, q
        return sse, q

    obj0, q0 = _eval_k(ls0, lm0)
    q_sub = q0.reshape(B, K, NS, 16)
    xb = X.reshape(B, 1, NS, 16)
    wb = W.reshape(B, 1, NS, 16)
    Sqq = torch.sum(wb * q_sub * q_sub, dim=3)
    Swq = torch.sum(wb * q_sub, dim=3)
    Sw = torch.sum(wb, dim=3)
    Sqx = torch.sum(wb * q_sub * xb, dim=3)
    Sx = torch.sum(wb * xb, dim=3)
    det = Sqq * Sw - Swq * Swq
    valid = det > 1.0e-30
    s_det = torch.where(valid, det, torch.ones_like(det))
    s_sw = torch.where(Sw > 1.0e-30, Sw, torch.ones_like(Sw))
    sj = torch.clamp(torch.where(valid, (Sqx * Sw - Swq * Sx) / s_det, torch.zeros_like(det)), min=0.0)
    tj = torch.where(valid, (sj * Swq - Sx) / s_sw, torch.zeros_like(det))
    bad_tj = tj < 0.0
    tj = torch.where(bad_tj, torch.zeros_like(tj), tj)
    s_sqq = torch.where(Sqq > 1.0e-30, Sqq, torch.ones_like(Sqq))
    sj = torch.where(bad_tj, torch.clamp(torch.where(Sqq > 1.0e-30, Sqx / s_sqq, torch.zeros_like(Sqq)), min=0.0), sj)

    ls1 = torch.where(
        ds[:, :, None] > 1.0e-30,
        _nearest_int_t(sj / s_ds).clamp(0, 15),
        torch.zeros((B, K, NS), dtype=torch.int64, device=device),
    ).to(torch.float32)
    t_tgt = torch.clamp((ds[:, :, None] * ls1 * Swq - Sx) / s_sw, min=0.0) if cfg.factor_coupled_q2k else tj
    lm1 = torch.where(
        ms[:, :, None] > 1.0e-30,
        _nearest_int_t(t_tgt / s_ms).clamp(0, 15),
        torch.zeros((B, K, NS), dtype=torch.int64, device=device),
    ).to(torch.float32)
    obj1, _ = _eval_k(ls1, lm1)
    best_k = torch.minimum(obj0, obj1).cpu().numpy()

    out = np.full((B, TOTAL_CANDIDATES), np.inf, dtype=np.float32)
    if nd == N_CAND_D and nm == N_CAND_M:
        out[:, :] = best_k
    else:
        for a, di in enumerate(di_idx):
            cb = int(di) >> 2
            for b, mi in enumerate(mi_idx):
                fb = int(mi) >> 2
                col = best_k[:, a * nm + b]
                for aa in range(4):
                    for bb in range(4):
                        idx = (cb * 4 + aa) * N_CAND_M + fb * 4 + bb
                        out[:, idx] = np.minimum(out[:, idx], col)
    return out


def _standalone_coupled_solve_cuda(
    X: "torch.Tensor", W: "torch.Tensor", cfg: HPCConfig
) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor"]:
    B = X.shape[0]
    Sw = torch.sum(W.reshape(B, NS, 16), dim=2)
    s_Sw = torch.clamp(Sw, min=1.0e-30)

    sources = [X]
    if cfg.use_weight_analogs:
        sources.extend(_analogs_cuda(X, W))

    best_obj = torch.full((B,), float("inf"), dtype=torch.float32, device=X.device)
    best_d = torch.zeros((B,), dtype=torch.float32, device=X.device)
    best_m = torch.zeros((B,), dtype=torch.float32, device=X.device)
    best_ls = torch.zeros((B, NS), dtype=torch.uint8, device=X.device)
    best_lm = torch.zeros((B, NS), dtype=torch.uint8, device=X.device)
    best_codes = torch.zeros((B, QK_K), dtype=torch.uint8, device=X.device)

    for src in sources:
        sb = src.reshape(B, NS, 16)
        xmin = torch.clamp(torch.min(sb, dim=2).values, max=0.0)
        xmax = torch.max(sb, dim=2).values
        span = xmax - xmin
        s0 = torch.where(span > 1.0e-30, span / 3.0, torch.zeros_like(span))
        t0 = -xmin
        s_s0 = torch.where(s0 > 1.0e-30, s0, torch.ones_like(s0))
        q0 = torch.where(
            s0[:, :, None] > 1.0e-30,
            _nearest_int_t((sb + t0[:, :, None]) / s_s0[:, :, None]).clamp(0, 3),
            torch.zeros((B, NS, 16), dtype=torch.int64, device=X.device),
        ).reshape(B, QK_K).to(torch.uint8)

        codes = q0
        for _ in range(3):
            s, t, Sqq, Swq, _, Sx = _sub_ls_cuda(X, W, codes)
            q_bar = Swq / s_Sw
            Vqq = torch.clamp(Sqq - Swq * q_bar, min=1.0e-12 * Sw)
            d, ls = _factor_1d_cuda(s, Vqq, 15, 4, cfg.factor_refine_1d)
            t_c = torch.clamp((d[:, None] * ls.to(torch.float32) * Swq - Sx) / s_Sw, min=0.0)
            m, lm = _factor_1d_cuda(t_c, Sw, 15, 4, cfg.factor_refine_1d)
            d = _round_fp16_t(d)
            m = _round_fp16_t(m)
            codes = _codes_from_params_t(X, d, m, ls, lm)
            _, sse_cur = _objective_t(X, W, d, m, ls, lm, codes, None, cfg)
            take = sse_cur < best_obj - 1.0e-12
            if torch.any(take):
                best_obj = torch.where(take, sse_cur, best_obj)
                best_d = torch.where(take, d, best_d)
                best_m = torch.where(take, m, best_m)
                best_ls = torch.where(take[:, None], ls, best_ls)
                best_lm = torch.where(take[:, None], lm, best_lm)
                best_codes = torch.where(take[:, None], codes, best_codes)

    return best_d, best_m, best_ls, best_lm, best_codes, best_obj


def _extended_dm_solve_and_ulp_cuda(
    X: "torch.Tensor",
    W: "torch.Tensor",
    ls: "torch.Tensor",
    lm: "torch.Tensor",
    codes: "torch.Tensor",
    d_cur: "torch.Tensor",
    m_cur: "torch.Tensor",
    carry: "torch.Tensor",
    cfg: HPCConfig,
) -> tuple["torch.Tensor", "torch.Tensor"]:
    B = X.shape[0]
    A = ls.to(torch.float32).repeat_interleave(16, dim=1) * codes.to(torch.float32)
    Bm = lm.to(torch.float32).repeat_interleave(16, dim=1)
    Saa = torch.sum(W * A * A, dim=1)
    Sab = torch.sum(W * A * Bm, dim=1)
    Sbb = torch.sum(W * Bm * Bm, dim=1)
    Sxa = torch.sum(W * X * A, dim=1)
    Sxb = torch.sum(W * X * Bm, dim=1)

    det = Saa * Sbb - Sab * Sab
    valid = det > 1.0e-30
    s_det = torch.where(valid, det, torch.ones_like(det))
    d_new = (Sbb * Sxa - Sab * Sxb) / s_det
    m_new = (Sab * Sxa - Saa * Sxb) / s_det
    d_opt = torch.where(valid & (d_new > 0.0) & torch.isfinite(d_new), d_new, d_cur)
    m_opt = torch.where(valid & (m_new >= 0.0) & torch.isfinite(m_new), m_new, m_cur)

    rad = max(int(cfg.ulp_radius), 1)
    deltas = torch.arange(-rad, rad + 1, dtype=torch.int32, device=X.device)
    d16_c = _to_fp16_bits_t(d_opt)
    m16_c = _to_fp16_bits_t(m_opt)
    d_bits = torch.clamp(d16_c[:, None] + deltas[None, :], 0, 0x7BFF)
    m_bits = torch.clamp(m16_c[:, None] + deltas[None, :], 0, 0x7BFF)
    dd = _from_fp16_bits_t(d_bits)[:, :, None]
    mm = _from_fp16_bits_t(m_bits)[:, None, :]

    q_form = (
        Saa[:, None, None] * (dd * dd)
        - 2.0 * Sab[:, None, None] * (dd * mm)
        + Sbb[:, None, None] * (mm * mm)
        - 2.0 * Sxa[:, None, None] * dd
        + 2.0 * Sxb[:, None, None] * mm
    ).reshape(B, -1)
    best_idx = torch.argmin(q_form, dim=1)
    K = deltas.numel()
    rows = torch.arange(B, device=X.device)
    return dd[rows, best_idx // K, 0], mm[rows, 0, best_idx % K]


def _refine_trajectory_cuda(
    X: "torch.Tensor",
    W: "torch.Tensor",
    Xb: "torch.Tensor",
    Wb: "torch.Tensor",
    Sw: "torch.Tensor",
    Sx: "torch.Tensor",
    safe_Sw: "torch.Tensor",
    d0: "torch.Tensor",
    m0: "torch.Tensor",
    ls0: "torch.Tensor",
    lm0: "torch.Tensor",
    codes0: "torch.Tensor",
    carry: "torch.Tensor",
    cfg: HPCConfig,
):
    B = X.shape[0]
    device = X.device
    best_d, best_m = d0.clone(), m0.clone()
    best_ls, best_lm = ls0.clone(), lm0.clone()
    best_codes = codes0.clone()
    _, best_sse = _objective_t(X, W, best_d, best_m, best_ls, best_lm, best_codes, carry, cfg)
    d, m, ls, lm, codes = best_d.clone(), best_m.clone(), best_ls.clone(), best_lm.clone(), best_codes.clone()

    ls_deltas = torch.tensor([-3, -2, -1, 0, 1, 2, 3], dtype=torch.int64, device=device)
    lm_deltas = torch.tensor([-1, 0, 1], dtype=torch.int64, device=device).repeat(7)
    b_idx = torch.arange(B, device=device)[:, None]
    s_idx = torch.arange(NS, device=device)[None, :]

    for _ in range(cfg.refine_passes):
        d_new, m_new = _extended_dm_solve_and_ulp_cuda(X, W, ls, lm, codes, d, m, carry, cfg)
        _, sse_dm = _objective_t(X, W, d_new, m_new, ls, lm, codes, carry, cfg)
        imp_dm = sse_dm < best_sse - 1.0e-12
        if torch.any(imp_dm):
            best_sse = torch.where(imp_dm, sse_dm, best_sse)
            best_d = torch.where(imp_dm, d_new, best_d)
            best_m = torch.where(imp_dm, m_new, best_m)
            best_ls = torch.where(imp_dm[:, None], ls, best_ls)
            best_lm = torch.where(imp_dm[:, None], lm, best_lm)
            best_codes = torch.where(imp_dm[:, None], codes, best_codes)
        d, m = best_d.clone(), best_m.clone()

        codes_re = _codes_from_params_t(X, d, m, ls, lm)
        _, sse_re = _objective_t(X, W, d, m, ls, lm, codes_re, carry, cfg)
        imp_re = sse_re < best_sse - 1.0e-12
        if torch.any(imp_re):
            best_sse = torch.where(imp_re, sse_re, best_sse)
            best_codes = torch.where(imp_re[:, None], codes_re, best_codes)
        codes = best_codes.clone()

        q = codes.to(torch.float32).reshape(B, NS, 16)
        Sqq = torch.sum(Wb * q * q, dim=2)
        Swq = torch.sum(Wb * q, dim=2)
        Sqx = torch.sum(Wb * q * Xb, dim=2)
        q_bar = Swq / safe_Sw
        Vqq = torch.clamp(Sqq - Swq * q_bar, min=0.0)
        Cqx = Sqx - Sx * q_bar
        s_star = torch.where(
            Vqq > 1.0e-30,
            torch.clamp(Cqx / torch.clamp(Vqq, min=1.0e-30), min=0.0),
            d[:, None] * ls.to(torch.float32),
        )
        sd = torch.where(d[:, None] > 1.0e-15, d[:, None], torch.ones((B, 1), dtype=torch.float32, device=device))
        sm = torch.where(m[:, None] > 1.0e-15, m[:, None], torch.ones((B, 1), dtype=torch.float32, device=device))
        ls_center = torch.where(
            d[:, None] > 1.0e-15,
            _nearest_int_t(s_star / sd).clamp(0, 15),
            torch.zeros((B, NS), dtype=torch.int64, device=device),
        )
        c_ls_base = torch.clamp(ls_center[:, :, None] + ls_deltas[None, None, :], 0, 15).to(torch.float32)
        d_sub_base = d[:, None, None] * c_ls_base
        t_opt_base = torch.clamp((d_sub_base * Swq[:, :, None] - Sx[:, :, None]) / safe_Sw[:, :, None], min=0.0)
        lm_center = torch.where(
            m[:, None, None] > 1.0e-15,
            _nearest_int_t(t_opt_base / sm[:, :, None]).clamp(0, 15),
            torch.zeros_like(c_ls_base, dtype=torch.int64),
        )
        cand_ls = c_ls_base.repeat_interleave(3, dim=2)
        cand_lm = torch.clamp(lm_center.repeat_interleave(3, dim=2) + lm_deltas[None, None, :], 0, 15).to(torch.float32)

        d_sub_c = d[:, None, None] * cand_ls
        m_sub_c = m[:, None, None] * cand_lm
        safe_dc = d_sub_c >= 1.0e-15
        sd_c = torch.where(safe_dc, d_sub_c, torch.ones_like(d_sub_c))
        q_c = torch.where(
            safe_dc[:, :, :, None],
            _nearest_int_t((Xb[:, :, None, :] + m_sub_c[:, :, :, None]) / sd_c[:, :, :, None]).clamp(0, 3),
            torch.zeros((B, NS, 21, 16), dtype=torch.int64, device=device),
        ).to(torch.float32)
        e_c = Xb[:, :, None, :] - (d_sub_c[:, :, :, None] * q_c - m_sub_c[:, :, :, None])
        err_c = torch.sum(Wb[:, :, None, :] * e_c * e_c, dim=3)
        pick = torch.argmin(err_c, dim=2)
        ls_try = cand_ls[b_idx, s_idx, pick].to(torch.uint8)
        lm_try = cand_lm[b_idx, s_idx, pick].to(torch.uint8)
        codes_try = _codes_from_params_t(X, d, m, ls_try, lm_try)
        d_try, m_try = _extended_dm_solve_and_ulp_cuda(X, W, ls_try, lm_try, codes_try, d, m, carry, cfg)
        _, sse_try = _objective_t(X, W, d_try, m_try, ls_try, lm_try, codes_try, carry, cfg)
        imp_sub = sse_try < best_sse - 1.0e-12
        if torch.any(imp_sub):
            best_sse = torch.where(imp_sub, sse_try, best_sse)
            best_d = torch.where(imp_sub, d_try, best_d)
            best_m = torch.where(imp_sub, m_try, best_m)
            best_ls = torch.where(imp_sub[:, None], ls_try, best_ls)
            best_lm = torch.where(imp_sub[:, None], lm_try, best_lm)
            best_codes = torch.where(imp_sub[:, None], codes_try, best_codes)
        d, m, ls, lm, codes = best_d.clone(), best_m.clone(), best_ls.clone(), best_lm.clone(), best_codes.clone()

    return best_d, best_m, best_ls, best_lm, best_codes, best_sse


def _compute_dc_carry_fast(
    X_np: np.ndarray,
    scales_seed: np.ndarray,
    mins_seed: np.ndarray,
    best: np.ndarray,
    base_d: np.ndarray,
    base_m: np.ndarray,
    cfg: HPCConfig,
    row_width: int | None,
) -> np.ndarray:
    """Vectorized residual DC computation + fast row recurrence."""
    n = X_np.shape[0]
    di = best // N_CAND_M
    mi = best % N_CAND_M
    d = (base_d * MULTS_D[di]).astype(np.float16).astype(np.float32)
    m = (base_m * MULTS_M[mi]).astype(np.float16).astype(np.float32)
    sd = np.where(d[:, None] > 1.0e-15, d[:, None], 1.0)
    sm = np.where(m[:, None] > 1.0e-15, m[:, None], 1.0)
    ls = np.where(d[:, None] > 1.0e-15, np.clip(np.rint(scales_seed / sd), 0, 15), 0).astype(np.float32)
    lm = np.where(m[:, None] > 1.0e-15, np.clip(np.rint(mins_seed / sm), 0, 15), 0).astype(np.float32)
    d_sub = np.repeat(d[:, None] * ls, 16, axis=1)
    m_sub = np.repeat(m[:, None] * lm, 16, axis=1)
    safe = d_sub >= 1.0e-15
    q = np.where(safe, np.clip(np.rint((X_np + m_sub) / np.where(safe, d_sub, 1.0)), 0, 3), 0.0)
    dc_res = np.sum(X_np - (d_sub * q - m_sub), axis=1, dtype=np.float64)

    bpr = row_width // QK_K if row_width and row_width % QK_K == 0 else 0
    carry = np.zeros(n, dtype=np.float32)
    decay = float(cfg.dc_decay)
    cum = bool(cfg.carry_cumulative)
    if bpr > 0 and n % bpr == 0:
        n_rows = n // bpr
        res_2d = dc_res.reshape(n_rows, bpr)
        car_2d = np.zeros((n_rows, bpr), dtype=np.float64)
        rolling = np.zeros(n_rows, dtype=np.float64)
        for col in range(bpr):
            car_2d[:, col] = decay * rolling
            rolling = (car_2d[:, col] + res_2d[:, col]) if cum else res_2d[:, col]
        carry = car_2d.reshape(n).astype(np.float32)
    else:
        rolling_s = 0.0
        for i in range(n):
            if bpr > 0 and i % bpr == 0:
                rolling_s = 0.0
            carry[i] = np.float32(decay * rolling_s)
            rolling_s = (carry[i] + dc_res[i]) if cum else dc_res[i]
    return carry


def _apply_dc_nudge_cuda(
    X_np: np.ndarray,
    W_np: np.ndarray | None,
    final_d: np.ndarray,
    final_m: np.ndarray,
    final_ls: np.ndarray,
    final_lm: np.ndarray,
    final_codes: np.ndarray,
    cfg: HPCConfig,
    row_width: int | None,
    dev: "torch.device",
) -> np.ndarray:
    """Row-parallel GPU DC residual nudge across all rows simultaneously."""
    n = X_np.shape[0]
    bpr = row_width // QK_K if row_width and row_width % QK_K == 0 else 0
    if bpr <= 0 or n % bpr != 0:
        return final_codes

    n_rows = n // bpr
    codes_out = final_codes.copy().reshape(n_rows, bpr, QK_K)
    X_r = X_np.reshape(n_rows, bpr, QK_K)
    W_r = None if W_np is None else W_np.reshape(n_rows, bpr, QK_K)
    d_r = final_d.reshape(n_rows, bpr)
    m_r = final_m.reshape(n_rows, bpr)
    ls_r = final_ls.reshape(n_rows, bpr, NS)
    lm_r = final_lm.reshape(n_rows, bpr, NS)

    decay = float(cfg.dc_decay)
    cum = bool(cfg.carry_cumulative)
    row_batch = 1024
    q_opts = torch.arange(4, dtype=torch.float32, device=dev)[None, :, None]  # (1, 4, 1)

    for r0 in range(0, n_rows, row_batch):
        r1 = min(r0 + row_batch, n_rows)
        R = r1 - r0
        rolling = torch.zeros(R, dtype=torch.float32, device=dev)
        X_batch = torch.from_numpy(X_r[r0:r1]).to(dev)
        W_batch = (
            torch.ones_like(X_batch)
            if W_r is None
            else torch.from_numpy(W_r[r0:r1]).to(dev)
        )
        d_batch = torch.from_numpy(d_r[r0:r1]).to(dev)
        m_batch = torch.from_numpy(m_r[r0:r1]).to(dev)
        ls_batch = torch.from_numpy(ls_r[r0:r1]).to(dev, dtype=torch.float32)
        lm_batch = torch.from_numpy(lm_r[r0:r1]).to(dev, dtype=torch.float32)
        cd_batch = torch.from_numpy(codes_out[r0:r1]).to(dev, dtype=torch.int64)
        rows = torch.arange(R, device=dev)

        for col in range(bpr):
            carry = decay * rolling
            x_c = X_batch[:, col]
            w_c = W_batch[:, col]
            d_c = d_batch[:, col]
            m_c = m_batch[:, col]
            ls_c = ls_batch[:, col]
            lm_c = lm_batch[:, col]
            cd_c = cd_batch[:, col]

            d_sub = (d_c[:, None] * ls_c).repeat_interleave(16, dim=1)
            m_sub = (m_c[:, None] * lm_c).repeat_interleave(16, dim=1)
            active = d_sub >= 1.0e-15
            e = x_c - (d_sub * cd_c.to(torch.float32) - m_sub)
            sse = torch.sum(w_c * e * e, dim=1)
            cap = sse * (1.0 + float(cfg.sse_budget))
            res = torch.sum(e, dim=1) + carry

            for _ in range(2):
                e_new = x_c[:, None, :] - (d_sub[:, None, :] * q_opts - m_sub[:, None, :])
                dsse = w_c[:, None, :] * (e_new * e_new - e[:, None, :] * e[:, None, :])
                mask = (
                    active[:, None, :]
                    & (q_opts.to(torch.int64) != cd_c[:, None, :])
                    & ((sse[:, None, None] + dsse) <= cap[:, None, None])
                )
                new_res = res[:, None, None] + (e_new - e[:, None, :])
                impr = torch.where(
                    mask,
                    torch.abs(res)[:, None, None] - torch.abs(new_res),
                    torch.full_like(new_res, -float("inf")),
                ).reshape(R, -1)
                best_imp, flat_idx = torch.max(impr, dim=1)
                do_upd = best_imp > 1.0e-12
                if not torch.any(do_upd):
                    break
                q_pick = flat_idx // QK_K
                k_pick = flat_idx % QK_K
                old_e = e[rows, k_pick]
                chosen_e = e_new[rows, q_pick, k_pick]
                de = chosen_e - old_e
                res = torch.where(do_upd, res + de, res)
                sse = torch.where(
                    do_upd,
                    sse + w_c[rows, k_pick] * (chosen_e * chosen_e - old_e * old_e),
                    sse,
                )
                e[rows, k_pick] = torch.where(do_upd, chosen_e, old_e)
                cd_c[rows, k_pick] = torch.where(do_upd, q_pick, cd_c[rows, k_pick])

            cd_batch[:, col] = cd_c
            dc_res = torch.sum(e, dim=1)
            rolling = (carry + dc_res) if cum else dc_res

        codes_out[r0:r1] = cd_batch.cpu().numpy().astype(np.uint8)

    return codes_out.reshape(n, QK_K)


# VRAM-aware CUDA batching for the Q2_K factored pipeline.  The candidate
# kernels materialize several live (B, K, 256) fp32 tensors (K = n_cand_d *
# n_cand_m), so the per-block footprint scales with K (~6 fp32 temporaries of
# 256 weights each).  The floor is the historical 4096-block batch (known to
# fit even on small cards); on large GPUs the batch grows with the free VRAM.
Q2_K_CUDA_FLOOR_BLOCKS = 4096
Q2_K_CUDA_VRAM_FRACTION = 0.45
Q2_K_CUDA_FP32_TEMPS_PER_CANDIDATE = 6


def quantize_cuda(
    weights: np.ndarray,
    importance: np.ndarray | None,
    row_width: int | None,
    cfg: HPCConfig,
    status_callback=None,
    stop_event=None,
    device: str = "cuda",
) -> QuantizedTensor:
    """Run the full HPC-Quantize V2 pipeline on a CUDA GPU."""
    if not cuda_is_available():
        raise RuntimeError("PyTorch CUDA is not available on this system")

    def _emit(phase: str, done: int, total: int, detail: str, **extra) -> None:
        if status_callback is not None:
            status_callback({"phase": phase, "done": done, "total": total, "detail": detail, **extra})

    def _check_stop() -> None:
        if stop_event is not None and stop_event.is_set():
            raise InterruptedError("Quantization cancelled by user")

    x = np.asarray(weights, dtype=np.float32)
    if x.size == 0:
        raise ValueError("weight tensor is empty")
    if x.size % QK_K != 0:
        raise ValueError("weight count must be a multiple of 256")
    if x.ndim == 2 and row_width is None and x.shape[1] % QK_K == 0:
        row_width = int(x.shape[1])
    x = x.ravel()
    n_blocks = x.size // QK_K
    X_np = x.reshape(n_blocks, QK_K)

    if importance is not None:
        imp = np.asarray(importance, dtype=np.float32).ravel()
        if imp.size != x.size:
            raise ValueError("importance must match the weight count")
        if np.any(imp < 0.0) or not np.all(np.isfinite(imp)):
            raise ValueError("importance must be finite and non-negative")
        W_np = imp.reshape(n_blocks, QK_K)
    else:
        W_np = None

    sigma = float(np.sqrt(np.mean(x.astype(np.float64) ** 2)))
    clamp_val = seed_clamp_value(x, cfg)
    dev = torch.device(device)
    gpu_name = cuda_device_name(device)
    batch_blocks = Q2_K_CUDA_FLOOR_BLOCKS
    try:
        # VRAM-aware default: scale the batch to free memory at start-up,
        # never below the floor.  Footprint per block: K live (B, K, 256)
        # fp32 temporaries drive the peak (see _factor_candidates_cuda).
        k_eff = int(cfg.n_cand_d) * int(cfg.n_cand_m)
        bytes_per_block = (
            Q2_K_CUDA_FP32_TEMPS_PER_CANDIDATE * k_eff * QK_K * 4
        )
        free_bytes, _ = torch.cuda.mem_get_info(dev)
        batch_blocks = max(
            Q2_K_CUDA_FLOOR_BLOCKS,
            int(Q2_K_CUDA_VRAM_FRACTION * free_bytes) // max(bytes_per_block, 1),
        )
    except Exception:
        batch_blocks = Q2_K_CUDA_FLOOR_BLOCKS

    candidate_errors = np.empty((n_blocks, TOTAL_CANDIDATES), dtype=np.float32)
    scales_seed = np.empty((n_blocks, NS), dtype=np.float32)
    mins_seed = np.empty((n_blocks, NS), dtype=np.float32)
    base_d = np.empty((n_blocks,), dtype=np.float32)
    base_m = np.empty((n_blocks,), dtype=np.float32)
    solo_d = np.empty((n_blocks,), dtype=np.float32)
    solo_m = np.empty((n_blocks,), dtype=np.float32)
    solo_ls = np.empty((n_blocks, NS), dtype=np.uint8)
    solo_lm = np.empty((n_blocks, NS), dtype=np.uint8)
    solo_codes = np.empty((n_blocks, QK_K), dtype=np.uint8)
    solo_obj = np.empty((n_blocks,), dtype=np.float64)

    t0 = time.time()
    with torch.no_grad():
        for start in range(0, n_blocks, batch_blocks):
            _check_stop()
            stop = min(start + batch_blocks, n_blocks)
            X_t = torch.from_numpy(X_np[start:stop]).to(dev)
            if W_np is None:
                W_t = torch.ones_like(X_t)
            else:
                W_raw_t = torch.from_numpy(W_np[start:stop]).to(dev)
                if cfg.temper_imatrix:
                    sig2_t = torch.mean(X_t * X_t, dim=1, keepdim=True)
                    W_t = torch.sqrt(torch.clamp(W_raw_t, min=1.0e-6)) * torch.sqrt(sig2_t + X_t * X_t)
                    wm_t = torch.mean(W_t, dim=1, keepdim=True)
                    W_t = torch.where(wm_t > 1.0e-30, W_t / torch.clamp(wm_t, min=1.0e-30), torch.ones_like(W_t))
                else:
                    W_t = W_raw_t
            seeds_t = _build_seeds_cuda(X_t, None if W_np is None else W_t, cfg)
            bd_t, bm_t = _phase2_wls_cuda(X_t, W_t, seeds_t)

            scales_seed[start:stop] = seeds_t["scales"].cpu().numpy()
            mins_seed[start:stop] = seeds_t["mins"].cpu().numpy()
            base_d[start:stop] = bd_t.cpu().numpy()
            base_m[start:stop] = bm_t.cpu().numpy()

            candidate_errors[start:stop] = _factor_candidates_cuda(
                X_t, W_t, seeds_t["scales"], seeds_t["mins"], bd_t, bm_t, cfg
            )
            sd_t, sm_t, sls_t, slm_t, sc_t, sobj_t = _standalone_coupled_solve_cuda(X_t, W_t, cfg)
            _, sse_Solo = _objective_t(X_t, W_t, sd_t, sm_t, sls_t, slm_t, sc_t, None, cfg)
            solo_d[start:stop] = sd_t.cpu().numpy()
            solo_m[start:stop] = sm_t.cpu().numpy()
            solo_ls[start:stop] = sls_t.cpu().numpy()
            solo_lm[start:stop] = slm_t.cpu().numpy()
            solo_codes[start:stop] = sc_t.cpu().numpy()
            solo_obj[start:stop] = sse_Solo.cpu().numpy()

            elapsed = max(time.time() - t0, 1.0e-6)
            rate = stop / elapsed
            eta = (n_blocks - stop) / max(rate, 1.0e-9)
            _emit(
                f"CUDA Phase 1–2b ({gpu_name}): Seeds + Analogs + Factored Field",
                stop,
                n_blocks,
                f"blk {stop}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
                f"mean_obj={float(np.mean(solo_obj[:stop])):.5e}",
                rate=rate,
                eta=eta,
                sigma=sigma,
            )

        _check_stop()
        _emit(
            "Phase 3: 6-State Sieve & 36-State Viterbi DP",
            0,
            n_blocks,
            f"Running Sieve + 36-state Viterbi sequence optimization ({cfg.optimizer.upper()})",
            sigma=sigma,
        )
        best, temperature, _ = _phase3_select(candidate_errors, cfg)
        carry = _compute_dc_carry_fast(
            X_np, scales_seed, mins_seed, best, base_d, base_m, cfg, row_width
        )

        final_d = np.empty((n_blocks,), dtype=np.float32)
        final_m = np.empty((n_blocks,), dtype=np.float32)
        final_ls = np.empty((n_blocks, NS), dtype=np.uint8)
        final_lm = np.empty((n_blocks, NS), dtype=np.uint8)
        final_codes = np.empty((n_blocks, QK_K), dtype=np.uint8)
        factor_floor_used = 0
        t_p4 = time.time()

        for start in range(0, n_blocks, batch_blocks):
            _check_stop()
            stop = min(start + batch_blocks, n_blocks)
            B_c = stop - start
            X_t = torch.from_numpy(X_np[start:stop]).to(dev)
            if W_np is None:
                W_t = torch.ones_like(X_t)
            else:
                W_raw_t = torch.from_numpy(W_np[start:stop]).to(dev)
                if cfg.temper_imatrix:
                    sig2_t = torch.mean(X_t * X_t, dim=1, keepdim=True)
                    W_t = torch.sqrt(torch.clamp(W_raw_t, min=1.0e-6)) * torch.sqrt(sig2_t + X_t * X_t)
                    wm_t = torch.mean(W_t, dim=1, keepdim=True)
                    W_t = torch.where(wm_t > 1.0e-30, W_t / torch.clamp(wm_t, min=1.0e-30), torch.ones_like(W_t))
                else:
                    W_t = W_raw_t
            Xb = X_t.reshape(B_c, NS, 16)
            Wb = W_t.reshape(B_c, NS, 16)
            Sw = torch.sum(Wb, dim=2)
            Sx = torch.sum(Wb * Xb, dim=2)
            safe_Sw = torch.clamp(Sw, min=1.0e-30)
            carry_t = torch.from_numpy(carry[start:stop]).to(dev)

            b_slice = best[start:stop]
            di = b_slice // N_CAND_M
            mi = b_slice % N_CAND_M
            d_v = torch.from_numpy(base_d[start:stop] * MULTS_D[di]).to(dev)
            m_v = torch.from_numpy(base_m[start:stop] * MULTS_M[mi]).to(dev)
            d_v = _round_fp16_t(d_v)
            m_v = _round_fp16_t(m_v)
            sc_s = torch.from_numpy(scales_seed[start:stop]).to(dev)
            mn_s = torch.from_numpy(mins_seed[start:stop]).to(dev)
            sd_v = torch.where(d_v[:, None] > 1.0e-15, d_v[:, None], torch.ones_like(d_v[:, None]))
            sm_v = torch.where(m_v[:, None] > 1.0e-15, m_v[:, None], torch.ones_like(m_v[:, None]))
            ls_v = torch.where(
                d_v[:, None] > 1.0e-15,
                _nearest_int_t(sc_s / sd_v).clamp(0, 15),
                torch.zeros((B_c, NS), dtype=torch.int64, device=dev),
            ).to(torch.uint8)
            lm_v = torch.where(
                m_v[:, None] > 1.0e-15,
                _nearest_int_t(mn_s / sm_v).clamp(0, 15),
                torch.zeros((B_c, NS), dtype=torch.int64, device=dev),
            ).to(torch.uint8)
            codes_v = _codes_from_params_t(X_t, d_v, m_v, ls_v, lm_v)

            d1, m1, ls1, lm1, c1, obj1 = _refine_trajectory_cuda(
                X_t, W_t, Xb, Wb, Sw, Sx, safe_Sw, d_v, m_v, ls_v, lm_v, codes_v, carry_t, cfg
            )

            if cfg.use_factor_floor:
                d_f = torch.from_numpy(solo_d[start:stop]).to(dev)
                m_f = torch.from_numpy(solo_m[start:stop]).to(dev)
                ls_f = torch.from_numpy(solo_ls[start:stop]).to(dev)
                lm_f = torch.from_numpy(solo_lm[start:stop]).to(dev)
                c_f = torch.from_numpy(solo_codes[start:stop]).to(dev)
                d2, m2, ls2, lm2, c2, obj2 = _refine_trajectory_cuda(
                    X_t, W_t, Xb, Wb, Sw, Sx, safe_Sw, d_f, m_f, ls_f, lm_f, c_f, carry_t, cfg
                )
                take2 = obj2 < obj1 - 1.0e-12
                factor_floor_used += int(torch.sum(take2).item())
                d1 = torch.where(take2, d2, d1)
                m1 = torch.where(take2, m2, m1)
                ls1 = torch.where(take2[:, None], ls2, ls1)
                lm1 = torch.where(take2[:, None], lm2, lm1)
                c1 = torch.where(take2[:, None], c2, c1)

            final_d[start:stop] = d1.cpu().numpy()
            final_m[start:stop] = m1.cpu().numpy()
            final_ls[start:stop] = ls1.cpu().numpy()
            final_lm[start:stop] = lm1.cpu().numpy()
            final_codes[start:stop] = c1.cpu().numpy()

            elapsed = max(time.time() - t_p4, 1.0e-6)
            rate = stop / elapsed
            eta = (n_blocks - stop) / max(rate, 1.0e-9)
            _emit(
                f"CUDA Phase 4 ({gpu_name}): Algebraic Coupled Bilinear Assembly",
                stop,
                n_blocks,
                f"blk {stop}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) floors(fac={factor_floor_used})",
                rate=rate,
                eta=eta,
                sigma=sigma,
                temperature=temperature,
                cand_floors=0,
                fac_floors=factor_floor_used,
            )

        if cfg.apply_dc_carry and (cfg.dc_lambda > 0.0 or cfg.dc_decay > 0.0):
            _check_stop()
            final_codes = _apply_dc_nudge_cuda(
                X_np, W_np, final_d, final_m, final_ls, final_lm, final_codes, cfg, row_width, dev
            )

        if cfg.pin_dc:
            _check_stop()
            t_pin = time.time()
            _emit(
                f"Phase 4.7: Zero-DC Pin ({gpu_name})",
                0,
                n_blocks,
                "re-solving (d, m) on the exact dc=0 constraint line "
                "(exact-JS zero-coherent-residual pin, SSE-nondecreasing)",
                sigma=sigma,
                temperature=temperature,
            )
            final_d, final_m = pin_q2k_dc_residual(
                final_d.copy(), final_m.copy(), final_ls, final_lm, final_codes, X_np.astype(np.float64)
            )
            _emit(
                f"Phase 4.7: Zero-DC Pin ({gpu_name})",
                n_blocks,
                n_blocks,
                f"done in {time.time() - t_pin:.1f}s",
                sigma=sigma,
                temperature=temperature,
            )

    scales = (final_ls & 0x0F) | ((final_lm & 0x0F) << np.uint8(4))
    qs = np.empty((n_blocks, QK_K // 4), dtype=np.uint8)
    for base in (0, 128):
        qs[:, base // 4 : base // 4 + 32] = (
            final_codes[:, base : base + 32]
            | (final_codes[:, base + 32 : base + 64] << np.uint8(2))
            | (final_codes[:, base + 64 : base + 96] << np.uint8(4))
            | (final_codes[:, base + 96 : base + 128] << np.uint8(6))
        )
    dm = np.empty((n_blocks, 2), dtype=np.uint16)
    dm[:, 0] = final_d.astype(np.float16).view(np.uint16)
    dm[:, 1] = final_m.astype(np.float16).view(np.uint16)

    d_f32 = dm[:, 0].view(np.float16).astype(np.float32)
    m_f32 = dm[:, 1].view(np.float16).astype(np.float32)
    d_sub = np.repeat(d_f32[:, None] * final_ls.astype(np.float32), 16, axis=1)
    m_sub = np.repeat(m_f32[:, None] * final_lm.astype(np.float32), 16, axis=1)
    xh_all = d_sub * final_codes.astype(np.float32) - m_sub
    err_mat = X_np.astype(np.float64) - xh_all.astype(np.float64)
    total_sse = float(np.sum(err_mat * err_mat))
    weighted = total_sse if W_np is None else float(np.sum(W_np.astype(np.float64) * err_mat * err_mat))
    rmse = float(np.sqrt(total_sse / x.size))

    blocks = [
        Q2KBlock(scales=scales[i], qs=qs[i], d=np.uint16(dm[i, 0]), dmin=np.uint16(dm[i, 1]))
        for i in range(n_blocks)
    ]

    stats = TensorStats(
        n_blocks=n_blocks,
        total_sse=total_sse,
        weighted_sse=weighted,
        rmse=rmse,
        sigma=sigma,
        rmse_over_sigma=rmse / sigma if sigma > 1.0e-15 else 0.0,
        temperature=temperature,
        factor_floor_used=factor_floor_used,
        candidate_floor_used=0,
    )
    _emit(
        f"Phase 5: Tensor Complete ({gpu_name})",
        n_blocks,
        n_blocks,
        f"Done: blocks={n_blocks} rmse={rmse:.6e} rmse/sigma={stats.rmse_over_sigma:.5f} "
        f"wSSE={weighted:.5e} temp={temperature:.3e} floors(fac={factor_floor_used})",
        rmse=rmse,
        sigma=sigma,
        temperature=temperature,
        cand_floors=0,
        fac_floors=factor_floor_used,
    )
    return QuantizedTensor(
        blocks=blocks,
        scales=scales,
        qs=qs,
        dm=dm,
        best_candidate=best,
        factor_objectives=solo_obj,
        config=cfg,
        stats=stats,
    )