"""Full HPC V2 Q3_K quantizer (CUDA GPU + Vectorized CPU).

Implements the complete HPC-Quantize V2 decision stack for llama.cpp/GGML
``Q3_K`` (256 weights per 110-byte block = 3.4375 bpw).  The Q3_K hierarchy is
the Q4_0 pipeline applied twice:

* **Level A — 16-weight sub-blocks** (the per-sub-block scale ``s_j`` search):
  runs the exact Q4_0 decision stack on every 16-weight sub-block with the
  signed 3-bit grid ``[-4, +3]``:
  1. Geometry-Preserving Analog + WLS seeds (``+d`` and asymmetric ``-d``).
  2. 24-candidate scale lattice & extended spectral objective.
  3. 6-state sequential Sieve & 6-state Viterbi DP across the sub-block chain.
  4. Asymmetric sign-flip fp16-ULP polish.
  5. Closed-form 1-D quadratic Vesica + DC code projection (``sse_budget``).

* **Level B — 256-weight super-blocks** (the GGML ``(d, sc)`` fit, no
  shortcuts relative to Q4_0):
  1. **Geometry-Preserving Analog + WLS Seeds** for the super-scale ``d``:
     four seeds (``+max``/``-max`` grid midpoint ``/31.5`` and the iso-moment
     companded analog ``+/-``), each refined by 4 rounds of weighted least
     squares over the 6-bit sub-block scales with per-sub-block importance,
     fp16-rounded and scored by the true 3-bit encode SSE (``+d`` and ``-d``).
  2. **24-Candidate Super-Scale Lattice & Extended Spectral Objective**:
     ``cand_d = fp16(wls_d * MULTS_D)``, each candidate jointlyLS derives the
     signed 6-bit scales ``sc ∈ [-32, 31]`` and 3-bit codes ``[-4, +3]`` and is
     scored on importance-weighted SSE + D6/Vesica fold pyramid + DC penalty
     over the full 256-weight block.
  3. **6-State Sequential Sieve & 6-State Viterbi DP** over the 24-candidate
     lattice with free-energy bin reduction and greedy override.
  4. **Asymmetric Sign-Flip fp16-ULP Polish**: ``±max(ulp_radius, 4)`` fp16
     ULPs around both ``+d_sel`` and ``-d_sel``, scored on exact encode SSE.
  5. **Closed-Form Vesica/DC Code Projection within ``sse_budget``**:
     a bounded gradient project-and-accept step on the 3-bit codes using the
     effective per-weight scale ``d*sc_j`` and the strict monotone
     ``(SSE + spectral)`` decrease test.

The final ``block_q3_K`` payload is emitted in the exact GGML layout
(``hmask[32] + qs[64] + scales[12] + fp16 d``, 110 bytes).
"""

from __future__ import annotations

import logging
from math import ceil
import numpy as np

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

from .anneal import thermal_refine_q3_k, thermal_refine_q3_k_cuda
from .config import HPCConfig
from .cuda_backend import cuda_is_available
from .factored import MULTS_D
from .fp16 import from_fp16_bits, nearest_int, to_fp16_bits
from .sieve import sieve_measure_chain, subblock_chain_edges
from .spectral import pin_dc_codes_uniform, pin_dc_codes_uniform_cuda
from .viterbi import robust_temperature

QK3_K = 256
Q3_K_BLOCK_BYTES = 110
Q3_K_SUBBLOCK = 16
Q3_K_N_SUBBLOCKS = QK3_K // Q3_K_SUBBLOCK
Q3_SC_MIN = -32
Q3_SC_MAX = 31
Q3_CODE_MIN = -4
Q3_CODE_MAX = 3
Q3_N_CAND = 24
Q3_N_STATES = 6

_log = logging.getLogger(__name__)

# Chunk sizes (in 256-weight super-blocks).  The CUDA default matches the
# Q4_0 reference chunk (131072 blocks * 32 weights = 4,194,304 weights per
# chunk): a Q3_K block is 8x wider, so 131072 / 8 = 16384 blocks keeps the same
# per-chunk GPU footprint while amortising the per-chunk CPU work (two
# Sieve/Viterbi passes, ~512 Python-loop steps each) and kernel-launch overhead.
# The CPU default stays small (numpy temporaries stay cache-resident).
Q3_K_CHUNK_BLOCKS_CUDA = 16384
Q3_K_CHUNK_BLOCKS_CPU = 1024
# Empirical peak VRAM per 256-weight super-block on the CUDA pipeline (same
# 4 GiB per reference chunk as Q4_0, over an 8x-wider block — identical
# 32768 bytes/block); used to auto-scale the GPU chunk to the free memory
# actually available on the device.
Q3_K_CUDA_BYTES_PER_BLOCK = 32768
Q3_K_CUDA_VRAM_FRACTION = 0.45


def _to_dev(x, dev: "torch.device") -> "torch.Tensor":
    """float32 tensor on ``dev``; no copy if ``x`` is already a device tensor."""
    if isinstance(x, torch.Tensor):
        return x.to(dev, dtype=torch.float32)
    return torch.from_numpy(np.ascontiguousarray(x)).to(dev, dtype=torch.float32)


def _nearest_int_t(value):
    """GGML round-to-nearest-even for torch tensors."""
    shifted = value + 12582912.0
    return (shifted.view(torch.int32) & 0x007FFFFF) - 0x00400000


def _unpack_q3_k_scales(scales: np.ndarray) -> np.ndarray:
    """Decode the 12-byte GGML Q3_K scale field into signed values [-32, 31]."""
    scales = np.asarray(scales, dtype=np.uint8)
    B = scales.shape[0]
    vals = np.empty((B, Q3_K_N_SUBBLOCKS), dtype=np.uint8)
    for j in range(8):
        vals[:, j] = scales[:, j] & np.uint8(0x0F)
        vals[:, j + 8] = scales[:, j] >> np.uint8(4)
    for j in range(Q3_K_N_SUBBLOCKS):
        hi = (scales[:, 8 + (j % 4)] >> np.uint8(2 * (j // 4))) & np.uint8(0x03)
        vals[:, j] |= hi << np.uint8(4)
    return vals.astype(np.int16) - 32


def _pack_q3_k_scales(scales: np.ndarray) -> np.ndarray:
    """Encode signed 6-bit scales [-32, 31] into the 12-byte GGML field."""
    sc = np.asarray(scales, dtype=np.int16).reshape(-1, Q3_K_N_SUBBLOCKS)
    B = sc.shape[0]
    vals = np.clip(sc + 32, 0, 63).astype(np.uint8)
    out = np.zeros((B, 12), dtype=np.uint8)
    for j in range(8):
        out[:, j] = vals[:, j] & np.uint8(0x0F)
        out[:, j] |= (vals[:, j + 8] & np.uint8(0x0F)) << np.uint8(4)
    for j in range(Q3_K_N_SUBBLOCKS):
        hi = (vals[:, j] >> np.uint8(4)) & np.uint8(0x03)
        out[:, 8 + (j % 4)] |= hi << np.uint8(2 * (j // 4))
    return out


def _pack_q3_k_codes(codes: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Pack signed 3-bit codes [-4, 3] into GGML hmask[32] + qs[64]."""
    c = np.asarray(codes, dtype=np.int16).reshape(-1, QK3_K)
    B = c.shape[0]
    low = (c & 3).astype(np.uint8)
    hmask = np.zeros((B, QK3_K // 8), dtype=np.uint8)
    high = c >= 0
    for bit in range(8):
        group = high[:, bit * 32 : (bit + 1) * 32]
        hmask |= group.astype(np.uint8) << np.uint8(bit)
    qs = np.empty((B, QK3_K // 4), dtype=np.uint8)
    qs[:, :32] = (
        low[:, 0:32]
        | (low[:, 32:64] << np.uint8(2))
        | (low[:, 64:96] << np.uint8(4))
        | (low[:, 96:128] << np.uint8(6))
    )
    qs[:, 32:] = (
        low[:, 128:160]
        | (low[:, 160:192] << np.uint8(2))
        | (low[:, 192:224] << np.uint8(4))
        | (low[:, 224:256] << np.uint8(6))
    )
    return hmask, qs


def dequantize_q3_k_bytes(raw: bytes | np.ndarray, n_blocks: int) -> np.ndarray:
    """Decode GGML Q3_K blocks into float32 weights."""
    blocks = np.frombuffer(raw, dtype=np.uint8).reshape(n_blocks, Q3_K_BLOCK_BYTES)
    hmask = blocks[:, : QK3_K // 8]
    qs = blocks[:, QK3_K // 8 : QK3_K // 8 + QK3_K // 4]
    sc = _unpack_q3_k_scales(blocks[:, QK3_K // 8 + QK3_K // 4 : Q3_K_BLOCK_BYTES - 2])
    d = blocks[:, Q3_K_BLOCK_BYTES - 2 :].copy().view("<f2").reshape(n_blocks, 1)

    out = np.empty((n_blocks, QK3_K), dtype=np.float32)
    for sb in range(Q3_K_N_SUBBLOCKS):
        group = sb // 2
        l0 = (sb % 2) * 16
        l = np.arange(l0, l0 + 16)
        shift = np.uint8(2 * (group % 4))
        qoff = 0 if group < 4 else 32
        low = (qs[:, qoff + l] >> shift) & np.uint8(0x03)
        high = (hmask[:, l] >> np.uint8(group)) & np.uint8(0x01)
        code = low.astype(np.int16) - np.where(high != 0, 0, 4).astype(np.int16)
        scale = d * sc[:, sb : sb + 1].astype(np.float32)
        out[:, sb * 16 : (sb + 1) * 16] = scale * code.astype(np.float32)
    return out.reshape(-1)


def _q3_sieve_and_viterbi(cand_errors: np.ndarray, cfg: HPCConfig) -> np.ndarray:
    """6-state sequential Sieve + 6-state Viterbi DP over a Q3 candidate lattice."""
    B = cand_errors.shape[0]
    if cfg.optimizer != "hpc" or B < 2:
        return np.argmin(cand_errors, axis=1).astype(np.int64)

    max_gb = min(int(cfg.graph_max_blocks), 512)
    stride = max(1, ceil(B / max_gb))
    gb = ceil(B / stride)
    temp = max(float(robust_temperature(cand_errors, cfg)), 1.0e-12)

    pad = gb * stride - B
    ce = cand_errors.astype(np.float64)
    if pad > 0:
        ce = np.concatenate([ce, np.full((pad, Q3_N_CAND), np.inf, dtype=np.float64)], axis=0)
    ce_grouped = ce.reshape(gb, stride, Q3_N_STATES, 4)

    group_min = np.min(ce_grouped, axis=(1, 3))
    finite = np.isfinite(ce_grouped)
    counts = np.sum(finite, axis=(1, 3))
    valid = (counts > 0) & np.isfinite(group_min)
    gm_b = np.where(np.isfinite(group_min), group_min, 0.0)[:, None, :, None]
    shifted = np.where(finite, (ce_grouped - gm_b) / (2.0 * temp), np.inf)
    sums = np.sum(np.exp(-shifted), axis=(1, 3))
    mean_exp = np.maximum(sums / np.maximum(counts, 1), 1.0e-300)
    bin_err = np.where(valid, group_min - 2.0 * temp * np.log(mean_exp), np.inf)

    agg = np.mean(np.where(finite, ce_grouped, 0.0), axis=1).reshape(gb, Q3_N_CAND)
    mn = np.min(agg, axis=1, keepdims=True)
    coeff = np.exp(-np.clip((agg - mn) / (2.0 * temp), 0.0, 80.0)).reshape(gb, Q3_N_STATES, 4)
    amps = np.sum(coeff, axis=2)
    norms = np.linalg.norm(amps, axis=1, keepdims=True)
    amps = np.where(norms > 1.0e-30, amps / np.maximum(norms, 1.0e-30), 1.0 / np.sqrt(6.0))

    marg, _ = sieve_measure_chain(
        amps,
        subblock_chain_edges(gb),
        same_penalty=cfg.same_penalty,
        log_slack=cfg.log_slack,
        deterministic=True,
    )
    m_tot = np.sum(marg, axis=1, keepdims=True)
    pm = np.where(m_tot > 1.0e-30, marg / np.maximum(m_tot, 1.0e-30), 1.0 / 6.0)
    log_prior = np.log(pm + 1.0e-30)

    fin_err = bin_err[np.isfinite(bin_err)]
    scale_err = max(float(fin_err.mean()) if fin_err.size else 1.0e-10, 1.0e-20)
    local = np.where(np.isfinite(bin_err), bin_err - cfg.viterbi_beta * scale_err * log_prior, np.inf)

    states = np.arange(Q3_N_STATES)
    trans = cfg.viterbi_alpha * scale_err * np.abs(states[None, :] - states[:, None])
    dp = np.full((gb, Q3_N_STATES), np.inf, dtype=np.float64)
    back = np.zeros((gb, Q3_N_STATES), dtype=np.int64)
    dp[0] = local[0]
    for gi in range(1, gb):
        totals = dp[gi - 1][:, None] + trans
        best_sp = np.argmin(totals, axis=0)
        best_val = totals[best_sp, states]
        dp[gi] = np.where(np.isfinite(best_val), best_val + local[gi], np.inf)
        back[gi] = best_sp

    path = np.zeros(gb, dtype=np.int64)
    if np.any(np.isfinite(dp[-1])):
        path[-1] = int(np.argmin(dp[-1]))
        for gi in range(gb - 2, -1, -1):
            path[gi] = back[gi + 1, path[gi + 1]]

    ce_bins = cand_errors.reshape(B, Q3_N_STATES, 4)
    bin_local_arg = np.argmin(ce_bins, axis=2)
    blk_gi = np.arange(B) // stride
    chosen_bin = path[blk_gi]
    chosen_cand = chosen_bin * 4 + bin_local_arg[np.arange(B), chosen_bin]

    global_best = np.argmin(cand_errors, axis=1).astype(np.int64)
    cur_e = cand_errors[np.arange(B), chosen_cand]
    g_e = cand_errors[np.arange(B), global_best]
    return np.where(g_e < cur_e * cfg.greedy_override_ratio, global_best, chosen_cand)


# ---------------------------------------------------------------------------
# Level A: HPC Q4_0 stack on 16-weight sub-blocks -> signed fp16 scales s_j
# ---------------------------------------------------------------------------


def _subblock_scales_hpc(
    X: np.ndarray, W: np.ndarray, cfg: HPCConfig
) -> np.ndarray:
    """Run the HPC 3-bit pipeline on (M, 16) sub-blocks and return signed scales."""
    M = X.shape[0]
    X_f32 = X.astype(np.float32)
    W_f32 = W.astype(np.float32)
    lam_dc = (float(cfg.dc_lambda) / Q3_K_SUBBLOCK) if cfg.use_spectral else 0.0
    lam_vw = (float(cfg.vw_lambda) / Q3_K_SUBBLOCK) if cfg.use_spectral else 0.0
    half = Q3_K_SUBBLOCK // 2

    amax = np.max(np.abs(X_f32), axis=1)
    d_range = amax / 3.5
    wsum = np.maximum(np.sum(W_f32, axis=1, keepdims=True), 1.0e-30)
    mu = np.sum(W_f32 * X_f32, axis=1, keepdims=True) / wsum
    var = np.sum(W_f32 * (X_f32 - mu) ** 2, axis=1, keepdims=True) / wsum
    sig = np.sqrt(np.maximum(var, 1.0e-30))
    x_comp = mu + sig * 1.85 * np.tanh((X_f32 - mu) / (1.85 * sig))
    d_analog = np.max(np.abs(x_comp), axis=1).astype(np.float32) / 3.5

    wls_d = d_range.copy()
    best_seed_sse = np.full(M, np.inf, dtype=np.float64)
    for d_start in (d_range, -d_range, d_analog, -d_analog):
        d = d_start.astype(np.float64)
        for _ in range(4):
            act = np.abs(d) > 1.0e-15
            sd = np.where(act, d, 1.0)[:, None]
            q = np.where(act[:, None], np.clip(nearest_int((X_f32 / sd).astype(np.float32)) + 4, 0, 7), 4)
            c = (q - 4).astype(np.float64)
            num = np.sum(W_f32 * X_f32 * c, axis=1, dtype=np.float64)
            den = np.sum(W_f32 * c * c, axis=1, dtype=np.float64)
            valid = den > 1.0e-30
            d = np.where(valid, num / np.where(valid, den, 1.0), 0.0)
        df = from_fp16_bits(to_fp16_bits(d.astype(np.float32)))
        act = np.abs(df) > 1.0e-15
        sdf = np.where(act, df, 1.0)[:, None]
        q = np.where(act[:, None], np.clip(nearest_int(X_f32 / sdf) + 4, 0, 7), 4)
        e = X_f32 - df[:, None] * (q.astype(np.float32) - 4.0)
        sse = np.sum(W_f32 * e * e, axis=1, dtype=np.float64)
        take = sse < best_seed_sse
        best_seed_sse = np.where(take, sse, best_seed_sse)
        wls_d = np.where(take, df, wls_d)

    cand_d16 = to_fp16_bits(wls_d[:, None] * MULTS_D[None, :])
    cand_d = from_fp16_bits(cand_d16)
    act_c = np.abs(cand_d) > 1.0e-15
    sd_c = np.where(act_c, cand_d, 1.0)[:, :, None]
    q_c = np.where(
        act_c[:, :, None],
        np.clip(nearest_int(X_f32[:, None, :] / sd_c) + 4, 0, 7),
        4,
    ).astype(np.float32)
    e_c = X_f32[:, None, :] - cand_d[:, :, None] * (q_c - 4.0)
    sse_c = np.sum(W_f32[:, None, :] * e_c * e_c, axis=2)
    if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0):
        v_c = e_c[:, :, :half] + e_c[:, :, half:]
        dc_c = np.sum(e_c, axis=2)
        cand_errors = (sse_c + lam_vw * np.sum(v_c * v_c, axis=2) + lam_dc * dc_c * dc_c).astype(np.float32)
    else:
        cand_errors = sse_c.astype(np.float32)

    best_cand = _q3_sieve_and_viterbi(cand_errors, cfg)
    rows = np.arange(M)
    d_sel = cand_d[rows, best_cand]

    rad = max(int(cfg.ulp_radius), 4)
    deltas = np.arange(-rad, rad + 1, dtype=np.int32)
    pos_b = np.clip(to_fp16_bits(np.abs(d_sel)).astype(np.int32)[:, None] + deltas[None, :], 0, 0x7BFF).astype(np.uint16)
    neg_b = (pos_b | np.uint16(0x8000)).astype(np.uint16)
    all_b = np.concatenate([pos_b, neg_b], axis=1)
    all_d = from_fp16_bits(all_b)
    act_u = np.abs(all_d) > 1.0e-15
    sd_u = np.where(act_u, all_d, 1.0)[:, :, None]
    q_u = np.where(
        act_u[:, :, None],
        np.clip(nearest_int(X_f32[:, None, :] / sd_u) + 4, 0, 7),
        4,
    ).astype(np.uint8)
    e_u = X_f32[:, None, :] - all_d[:, :, None] * (q_u.astype(np.float32) - 4.0)
    sse_u = np.sum(W_f32[:, None, :] * e_u * e_u, axis=2, dtype=np.float64)
    pick_u = np.argmin(sse_u, axis=1)
    best_d = all_d[rows, pick_u]
    best_q = q_u[rows, pick_u].astype(np.int64)

    if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0):
        c_f = (best_q - 4).astype(np.float64)
        e = X_f32.astype(np.float64) - best_d[:, None].astype(np.float64) * c_f
        sse_live = np.sum(W_f32 * e * e, axis=1, dtype=np.float64)
        sse_cap = sse_live * (1.0 + cfg.sse_budget)
        v = e[:, :half] + e[:, half:]
        dc = np.sum(e, axis=1)
        d_f64 = best_d.astype(np.float64)[:, None]
        denom = (W_f32 + lam_vw + lam_dc) * (d_f64 * d_f64)
        act_b = np.abs(d_f64) > 1.0e-15
        s_den = np.where(act_b & (denom > 1.0e-30), denom, 1.0)
        v_full = np.concatenate([v, v], axis=1)
        grad = d_f64 * (W_f32 * e + lam_vw * v_full + lam_dc * dc[:, None])
        dq = np.where(act_b, nearest_int((grad / s_den).astype(np.float32)), 0)
        q_new = np.clip(best_q + dq, 0, 7)
        e_new = X_f32.astype(np.float64) - d_f64 * (q_new - 4).astype(np.float64)
        sse_new = np.sum(W_f32 * e_new * e_new, axis=1, dtype=np.float64)
        v_new = e_new[:, :half] + e_new[:, half:]
        dc_new = np.sum(e_new, axis=1)
        spec_old = lam_vw * np.sum(v * v, axis=1) + lam_dc * dc * dc
        spec_new = lam_vw * np.sum(v_new * v_new, axis=1) + lam_dc * dc_new * dc_new
        accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-15)
        best_q = np.where(accept[:, None], q_new, best_q)

    return best_d.astype(np.float32)


def _subblock_scales_hpc_cuda(
    X_src, W_src, dev: "torch.device", cfg: HPCConfig, *, as_tensor: bool = False
):
    """Run the HPC 3-bit sub-block pipeline on CUDA and return signed scales.

    ``X_src`` / ``W_src`` may be numpy arrays or device tensors.  With
    ``as_tensor=True`` the scales stay on ``dev`` (no D2H copy).
    """
    M = X_src.shape[0]
    lam_dc = (float(cfg.dc_lambda) / Q3_K_SUBBLOCK) if cfg.use_spectral else 0.0
    lam_vw = (float(cfg.vw_lambda) / Q3_K_SUBBLOCK) if cfg.use_spectral else 0.0
    half = Q3_K_SUBBLOCK // 2

    with torch.no_grad():
        X = _to_dev(X_src, dev)
        W = _to_dev(W_src, dev)

        amax = torch.max(torch.abs(X), dim=1).values
        d_range = amax / 3.5
        wsum = torch.clamp(torch.sum(W, dim=1, keepdim=True), min=1.0e-30)
        mu = torch.sum(W * X, dim=1, keepdim=True) / wsum
        var = torch.sum(W * (X - mu) ** 2, dim=1, keepdim=True) / wsum
        sig = torch.sqrt(torch.clamp(var, min=1.0e-30))
        x_comp = mu + sig * 1.85 * torch.tanh((X - mu) / (1.85 * sig))
        d_analog = torch.max(torch.abs(x_comp), dim=1).values / 3.5

        wls_d = d_range.clone()
        best_seed_sse = torch.full((M,), float("inf"), dtype=torch.float32, device=dev)
        for d_start in (d_range, -d_range, d_analog, -d_analog):
            d = d_start.clone()
            for _ in range(4):
                act = torch.abs(d) > 1.0e-15
                sd = torch.where(act, d, torch.ones_like(d))[:, None]
                sh = (X / sd) + 12582912.0
                qi = ((sh.view(torch.int32) & 0x007FFFFF) - 0x00400000 + 4).clamp(0, 7)
                q = torch.where(act[:, None], qi, torch.full_like(qi, 4))
                c = (q - 4).to(torch.float32)
                num = torch.sum(W * X * c, dim=1)
                den = torch.sum(W * c * c, dim=1)
                valid = den > 1.0e-30
                d = torch.where(valid, num / torch.where(valid, den, torch.ones_like(den)), torch.zeros_like(den))
            df = d.to(torch.float16).to(torch.float32)
            act = torch.abs(df) > 1.0e-15
            sdf = torch.where(act, df, torch.ones_like(df))[:, None]
            sh = (X / sdf) + 12582912.0
            qi = ((sh.view(torch.int32) & 0x007FFFFF) - 0x00400000 + 4).clamp(0, 7)
            q = torch.where(act[:, None], qi, torch.full_like(qi, 4))
            e = X - df[:, None] * (q.to(torch.float32) - 4.0)
            sse = torch.sum(W * e * e, dim=1)
            take = sse < best_seed_sse
            best_seed_sse = torch.where(take, sse, best_seed_sse)
            wls_d = torch.where(take, df, wls_d)

        mults = torch.tensor(MULTS_D, dtype=torch.float32, device=dev)
        cand_d = (wls_d[:, None] * mults[None, :]).to(torch.float16).to(torch.float32)
        act_c = torch.abs(cand_d) > 1.0e-15
        sd_c = torch.where(act_c, cand_d, torch.ones_like(cand_d))[:, :, None]
        sh_c = (X[:, None, :] / sd_c) + 12582912.0
        qi_c = ((sh_c.view(torch.int32) & 0x007FFFFF) - 0x00400000 + 4).clamp(0, 7)
        q_c = torch.where(act_c[:, :, None], qi_c, torch.full_like(qi_c, 4)).to(torch.float32)
        e_c = X[:, None, :] - cand_d[:, :, None] * (q_c - 4.0)
        sse_c = torch.sum(W[:, None, :] * e_c * e_c, dim=2)
        if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0):
            v_c = e_c[:, :, :half] + e_c[:, :, half:]
            dc_c = torch.sum(e_c, dim=2)
            cand_err_t = sse_c + lam_vw * torch.sum(v_c * v_c, dim=2) + lam_dc * dc_c * dc_c
        else:
            cand_err_t = sse_c

        best_cand_np = _q3_sieve_and_viterbi(cand_err_t.cpu().numpy(), cfg)
        best_cand_t = torch.from_numpy(best_cand_np).to(dev)
        rows = torch.arange(M, device=dev)
        d_sel = cand_d[rows, best_cand_t]

        rad = max(int(cfg.ulp_radius), 4)
        deltas = torch.arange(-rad, rad + 1, dtype=torch.int32, device=dev)
        abs_bits = torch.abs(d_sel).to(torch.float16).view(torch.int16).to(torch.int32) & 0x7FFF
        pos_b = torch.clamp(abs_bits[:, None] + deltas[None, :], 0, 0x7BFF)
        neg_b = pos_b | 0x8000
        all_b = torch.cat([pos_b, neg_b], dim=1)
        all_d = all_b.to(torch.int16).view(torch.float16).to(torch.float32)
        act_u = torch.abs(all_d) > 1.0e-15
        sd_u = torch.where(act_u, all_d, torch.ones_like(all_d))[:, :, None]
        sh_u = (X[:, None, :] / sd_u) + 12582912.0
        qi_u = ((sh_u.view(torch.int32) & 0x007FFFFF) - 0x00400000 + 4).clamp(0, 7)
        q_u = torch.where(act_u[:, :, None], qi_u, torch.full_like(qi_u, 4)).to(torch.uint8)
        e_u = X[:, None, :] - all_d[:, :, None] * (q_u.to(torch.float32) - 4.0)
        sse_u = torch.sum(W[:, None, :] * e_u * e_u, dim=2)
        pick_u = torch.argmin(sse_u, dim=1)
        best_d = all_d[rows, pick_u]
        best_q = q_u[rows, pick_u].to(torch.int64)

        if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0):
            d_col = best_d[:, None]
            e = X - d_col * (best_q.to(torch.float32) - 4.0)
            sse_live = torch.sum(W * e * e, dim=1)
            sse_cap = sse_live * (1.0 + float(cfg.sse_budget))
            v = e[:, :half] + e[:, half:]
            dc = torch.sum(e, dim=1)
            denom = (W + lam_vw + lam_dc) * (d_col * d_col)
            act_b = torch.abs(d_col) > 1.0e-15
            s_den = torch.where(act_b & (denom > 1.0e-30), denom, torch.ones_like(denom))
            v_full = torch.cat([v, v], dim=1)
            grad = d_col * (W * e + lam_vw * v_full + lam_dc * dc[:, None])
            sh_dq = (grad / s_den) + 12582912.0
            dq = torch.where(
                act_b,
                ((sh_dq.view(torch.int32) & 0x007FFFFF) - 0x00400000).to(torch.int64),
                torch.zeros_like(best_q),
            )
            q_new = torch.clamp(best_q + dq, 0, 7)
            e_new = X - d_col * (q_new.to(torch.float32) - 4.0)
            sse_new = torch.sum(W * e_new * e_new, dim=1)
            v_new = e_new[:, :half] + e_new[:, half:]
            dc_new = torch.sum(e_new, dim=1)
            spec_old = lam_vw * torch.sum(v * v, dim=1) + lam_dc * dc * dc
            spec_new = lam_vw * torch.sum(v_new * v_new, dim=1) + lam_dc * dc_new * dc_new
            accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-12)
            best_q = torch.where(accept[:, None], q_new, best_q)

        if as_tensor:
            return best_d
        return best_d.cpu().numpy().astype(np.float32)


# ---------------------------------------------------------------------------
# Level B: GGML (d, sc, code) super-block fit — full Q4_0 decision stack
# ---------------------------------------------------------------------------


def _fit_q3_k_blocks(
    X: np.ndarray, W: np.ndarray, sub_scales: np.ndarray, cfg: HPCConfig
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Full HPC V2 super-block fit for (B, 256) blocks on CPU."""
    B = X.shape[0]
    X3 = X.astype(np.float32).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
    W3 = W.astype(np.float32).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
    Xf = X3.reshape(B, QK3_K)
    Wf = W3.reshape(B, QK3_K)
    s = sub_scales.reshape(B, Q3_K_N_SUBBLOCKS).astype(np.float32)
    rows = np.arange(B)
    lam_dc = (float(cfg.dc_lambda) / QK3_K) if cfg.use_spectral else 0.0
    lam_vw = (float(cfg.vw_lambda) / QK3_K) if cfg.use_spectral else 0.0
    Wj = W3.sum(axis=2)  # (B, 16) per-sub-block importance

    def _encode(d: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """JointLS quantize s with scale d: sc in [-32,31], codes in [-4,3]."""
        act = np.abs(d) > 1.0e-15
        safe_d = np.where(act, d, 1.0)
        sc = np.where(
            act[:, None],
            np.clip(nearest_int((s / safe_d[:, None]).astype(np.float32)), Q3_SC_MIN, Q3_SC_MAX),
            0,
        )
        scale = d[:, None] * sc.astype(np.float32)
        act_s = np.abs(scale) > 1.0e-15
        safe_s = np.where(act_s, scale, 1.0)
        code = np.where(
            act_s[:, :, None],
            np.clip(nearest_int((X3 / safe_s[:, :, None]).astype(np.float32)), Q3_CODE_MIN, Q3_CODE_MAX),
            0,
        )
        e = X3 - scale[:, :, None] * code.astype(np.float32)
        sse = np.sum(W3 * e * e, axis=(1, 2), dtype=np.float64)
        return sc, code, sse

    # 1. Geometry-Preserving Analog + WLS Seeds for the super-scale d:
    amax = np.max(np.abs(s), axis=1)
    d_range = amax / 31.5
    wsum = np.maximum(np.sum(Wj, axis=1, keepdims=True), 1.0e-30)
    mu = np.sum(Wj * s, axis=1, keepdims=True) / wsum
    var = np.sum(Wj * (s - mu) ** 2, axis=1, keepdims=True) / wsum
    sig = np.sqrt(np.maximum(var, 1.0e-30))
    s_comp = mu + sig * 1.85 * np.tanh((s - mu) / (1.85 * sig))
    d_analog = (np.max(np.abs(s_comp), axis=1) / 31.5).astype(np.float32)

    wls_d = d_range.copy()
    best_seed_sse = np.full(B, np.inf, dtype=np.float64)
    for d_start in (d_range, -d_range, d_analog, -d_analog):
        d = d_start.astype(np.float64)
        for _ in range(4):
            act = np.abs(d) > 1.0e-15
            sd = np.where(act, d, 1.0)[:, None]
            sc = np.where(
                act[:, None],
                np.clip(nearest_int((s / sd).astype(np.float32)), Q3_SC_MIN, Q3_SC_MAX),
                0,
            )
            c = sc.astype(np.float64)
            num = np.sum(Wj * s * c, axis=1, dtype=np.float64)
            den = np.sum(Wj * c * c, axis=1, dtype=np.float64)
            valid = den > 1.0e-30
            d = np.where(valid, num / np.where(valid, den, 1.0), 0.0)
        df = from_fp16_bits(to_fp16_bits(d.astype(np.float32)))
        _, _, sse = _encode(df)
        take = sse < best_seed_sse
        best_seed_sse = np.where(take, sse, best_seed_sse)
        wls_d = np.where(take, df, wls_d)

    # 2. 24-Candidate Super-Scale Lattice with Extended Spectral Objective:
    cand_d16 = to_fp16_bits(wls_d[:, None] * MULTS_D[None, :])  # (B, 24)
    cand_d = from_fp16_bits(cand_d16)
    act_c = np.abs(cand_d) > 1.0e-15
    safe_c = np.where(act_c, cand_d, 1.0)
    sc_c = np.where(
        act_c[:, :, None],
        np.clip(nearest_int((s[:, None, :] / safe_c[:, :, None]).astype(np.float32)), Q3_SC_MIN, Q3_SC_MAX),
        0,
    )
    scale_c = cand_d[:, :, None] * sc_c.astype(np.float32)  # (B, 24, 16)
    act_s = np.abs(scale_c) > 1.0e-15
    safe_s = np.where(act_s, scale_c, 1.0)
    code_c = np.where(
        act_s[:, :, :, None],
        np.clip(nearest_int((X3[:, None] / safe_s[:, :, :, None]).astype(np.float32)), Q3_CODE_MIN, Q3_CODE_MAX),
        0,
    ).astype(np.int16)
    e_f = (X3[:, None] - scale_c[:, :, :, None] * code_c.astype(np.float32)).reshape(B, Q3_N_CAND, QK3_K)
    sse_c = np.sum(Wf[:, None, :] * e_f * e_f, axis=2)
    if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0):
        half = QK3_K // 2
        v_c = e_f[:, :, :half] + e_f[:, :, half:]
        dc_c = np.sum(e_f, axis=2)
        cand_errors = (sse_c + lam_vw * np.sum(v_c * v_c, axis=2) + lam_dc * dc_c * dc_c).astype(np.float32)
    else:
        cand_errors = sse_c.astype(np.float32)

    # 3. 6-State Sieve + 6-State Viterbi DP over the super-scale lattice:
    best_cand = _q3_sieve_and_viterbi(cand_errors, cfg)
    d_sel = cand_d[rows, best_cand]

    # 4. Asymmetric Sign-Flip fp16-ULP Polish (+/-8 ULPs around +d_sel and -d_sel):
    rad = max(int(cfg.ulp_radius), 4)
    deltas = np.arange(-rad, rad + 1, dtype=np.int32)
    pos_b = np.clip(
        to_fp16_bits(np.abs(d_sel)).astype(np.int32)[:, None] + deltas[None, :],
        0,
        0x7BFF,
    ).astype(np.uint16)
    neg_b = (pos_b | np.uint16(0x8000)).astype(np.uint16)
    all_b = np.concatenate([pos_b, neg_b], axis=1)
    all_d = from_fp16_bits(all_b)
    act_u = np.abs(all_d) > 1.0e-15
    safe_u = np.where(act_u, all_d, 1.0)
    sc_u = np.where(
        act_u[:, :, None],
        np.clip(nearest_int((s[:, None, :] / safe_u[:, :, None]).astype(np.float32)), Q3_SC_MIN, Q3_SC_MAX),
        0,
    )
    scale_u = all_d[:, :, None] * sc_u.astype(np.float32)
    act_su = np.abs(scale_u) > 1.0e-15
    safe_su = np.where(act_su, scale_u, 1.0)
    code_u = np.where(
        act_su[:, :, :, None],
        np.clip(nearest_int((X3[:, None] / safe_su[:, :, :, None]).astype(np.float32)), Q3_CODE_MIN, Q3_CODE_MAX),
        0,
    )
    e_u = X3[:, None] - scale_u[:, :, :, None] * code_u.astype(np.float32)
    sse_u = np.sum(W3[:, None] * e_u * e_u, axis=(2, 3), dtype=np.float64)
    pick_u = np.argmin(sse_u, axis=1)
    best_d16 = all_b[rows, pick_u]
    best_d = all_d[rows, pick_u]

    # Final (sc, code) derived from the polished super-scale d:
    act_b = np.abs(best_d) > 1.0e-15
    safe_b = np.where(act_b, best_d, 1.0)
    sc = np.where(
        act_b[:, None],
        np.clip(nearest_int((s / safe_b[:, None]).astype(np.float32)), Q3_SC_MIN, Q3_SC_MAX),
        0,
    )
    scale = best_d[:, None] * sc.astype(np.float32)
    act_s = np.abs(scale) > 1.0e-15
    safe_s = np.where(act_s, scale, 1.0)
    code = np.where(
        act_s[:, :, None],
        np.clip(nearest_int((X3 / safe_s[:, :, None]).astype(np.float32)), Q3_CODE_MIN, Q3_CODE_MAX),
        0,
    )

    # 5. Closed-Form 1-D Quadratic Vesica + DC Code Projection within sse_budget:
    if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0):
        d_w = np.repeat(best_d[:, None].astype(np.float64) * sc.astype(np.float64), Q3_K_SUBBLOCK, axis=1)
        e = Xf.astype(np.float64) - d_w * code.reshape(B, QK3_K)
        sse_live = np.sum(Wf * e * e, axis=1, dtype=np.float64)
        sse_cap = sse_live * (1.0 + cfg.sse_budget)
        half = QK3_K // 2
        v = e[:, :half] + e[:, half:]
        dc = np.sum(e, axis=1)
        denom = (Wf + lam_vw + lam_dc) * (d_w * d_w)
        act_bw = np.abs(d_w) > 1.0e-15
        s_den = np.where(act_bw & (denom > 1.0e-30), denom, 1.0)
        v_full = np.concatenate([v, v], axis=1)
        grad = d_w * (Wf * e + lam_vw * v_full + lam_dc * dc[:, None])
        dq = np.where(act_bw, nearest_int((grad / s_den).astype(np.float32)), 0)
        code_f = code.reshape(B, QK3_K)
        code_new = np.clip(code_f + dq, Q3_CODE_MIN, Q3_CODE_MAX)
        e_new = Xf.astype(np.float64) - d_w * code_new
        sse_new = np.sum(Wf * e_new * e_new, axis=1, dtype=np.float64)
        v_new = e_new[:, :half] + e_new[:, half:]
        dc_new = np.sum(e_new, axis=1)
        spec_old = lam_vw * np.sum(v * v, axis=1) + lam_dc * dc * dc
        spec_new = lam_vw * np.sum(v_new * v_new, axis=1) + lam_dc * dc_new * dc_new
        accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-15)
        code = np.where(accept[:, None], code_new, code_f).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)

    # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing):
    if cfg.therm_passes > 0:
        scale_rp = best_d[:, None] * sc.astype(np.float32)
        thermal_c, _ = thermal_refine_q3_k(Xf, Wf, scale_rp, code.reshape(B, QK3_K), cfg)
        code = thermal_c.astype(np.int16).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)

    return best_d16.astype(np.uint16), sc.astype(np.int16), code.astype(np.int16)


def _fit_q3_k_blocks_cuda(
    X_src,
    W_src,
    sub_scales,
    dev: "torch.device",
    cfg: HPCConfig,
    *,
    as_tensor: bool = False,
    therm_seed: int | None = None,
):
    """Full HPC V2 super-block fit for (B, 256) blocks on CUDA.

    Inputs may be numpy arrays or device tensors.  Default return is the numpy
    triple ``(d16 uint16, sc int16, code int16)``.  With ``as_tensor=True`` the
    results stay on ``dev`` as ``(d16, d, sc, code)`` where ``d`` is the fp16-
    decoded float32 super-scale, ``sc`` is (B, 16) and ``code`` is (B, 16, 16).
    """
    B = X_src.shape[0]
    lam_dc = (float(cfg.dc_lambda) / QK3_K) if cfg.use_spectral else 0.0
    lam_vw = (float(cfg.vw_lambda) / QK3_K) if cfg.use_spectral else 0.0

    with torch.no_grad():
        X3 = _to_dev(X_src, dev).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
        W3 = _to_dev(W_src, dev).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
        Xf = X3.reshape(B, QK3_K)
        Wf = W3.reshape(B, QK3_K)
        s = _to_dev(sub_scales, dev).reshape(B, Q3_K_N_SUBBLOCKS)
        rows = torch.arange(B, device=dev)
        Wj = W3.sum(dim=2)  # (B, 16) per-sub-block importance

        def _encode(d: "torch.Tensor"):
            """JointLS quantize s with scale d: sc in [-32,31], codes in [-4,3]."""
            act = torch.abs(d) > 1.0e-15
            safe_d = torch.where(act, d, torch.ones_like(d))
            q_sc = _nearest_int_t(s / safe_d[:, None]).clamp(Q3_SC_MIN, Q3_SC_MAX)
            sc = torch.where(act[:, None], q_sc, torch.zeros_like(q_sc))
            scale = d[:, None] * sc.to(torch.float32)
            act_s = torch.abs(scale) > 1.0e-15
            safe_s = torch.where(act_s, scale, torch.ones_like(scale))
            q_cd = _nearest_int_t(X3 / safe_s[:, :, None]).clamp(Q3_CODE_MIN, Q3_CODE_MAX)
            code = torch.where(act_s[:, :, None], q_cd, torch.zeros_like(q_cd))
            e = X3 - scale[:, :, None] * code.to(torch.float32)
            sse = torch.sum(W3 * e * e, dim=(1, 2))
            return sc, code, sse

        # 1. Geometry-Preserving Analog + WLS Seeds for the super-scale d:
        amax = torch.max(torch.abs(s), dim=1).values
        d_range = amax / 31.5
        wsum = torch.clamp(torch.sum(Wj, dim=1, keepdim=True), min=1.0e-30)
        mu = torch.sum(Wj * s, dim=1, keepdim=True) / wsum
        var = torch.sum(Wj * (s - mu) ** 2, dim=1, keepdim=True) / wsum
        sig = torch.sqrt(torch.clamp(var, min=1.0e-30))
        s_comp = mu + sig * 1.85 * torch.tanh((s - mu) / (1.85 * sig))
        d_analog = torch.max(torch.abs(s_comp), dim=1).values / 31.5

        wls_d = d_range.clone()
        best_seed_sse = torch.full((B,), float("inf"), dtype=torch.float32, device=dev)
        for d_start in (d_range, -d_range, d_analog, -d_analog):
            d = d_start.clone()
            for _ in range(4):
                act = torch.abs(d) > 1.0e-15
                sd = torch.where(act, d, torch.ones_like(d))[:, None]
                qi = _nearest_int_t(s / sd).clamp(Q3_SC_MIN, Q3_SC_MAX)
                sc = torch.where(act[:, None], qi, torch.zeros_like(qi))
                c = sc.to(torch.float32)
                num = torch.sum(Wj * s * c, dim=1)
                den = torch.sum(Wj * c * c, dim=1)
                valid = den > 1.0e-30
                d = torch.where(valid, num / torch.where(valid, den, torch.ones_like(den)), torch.zeros_like(den))
            df = d.to(torch.float16).to(torch.float32)
            _, _, sse = _encode(df)
            take = sse < best_seed_sse
            best_seed_sse = torch.where(take, sse, best_seed_sse)
            wls_d = torch.where(take, df, wls_d)

        # 2. 24-Candidate Super-Scale Lattice with Extended Spectral Objective:
        mults = torch.tensor(MULTS_D, dtype=torch.float32, device=dev)
        cand_d = (wls_d[:, None] * mults[None, :]).to(torch.float16).to(torch.float32)
        act_c = torch.abs(cand_d) > 1.0e-15
        safe_c = torch.where(act_c, cand_d, torch.ones_like(cand_d))
        q_sc_c = _nearest_int_t(s[:, None, :] / safe_c[:, :, None]).clamp(Q3_SC_MIN, Q3_SC_MAX)
        sc_c = torch.where(act_c[:, :, None], q_sc_c, torch.zeros_like(q_sc_c))
        scale_c = cand_d[:, :, None] * sc_c.to(torch.float32)
        act_s = torch.abs(scale_c) > 1.0e-15
        safe_s = torch.where(act_s, scale_c, torch.ones_like(scale_c))
        q_cd_c = _nearest_int_t(X3[:, None] / safe_s[:, :, :, None]).clamp(Q3_CODE_MIN, Q3_CODE_MAX)
        code_c = torch.where(act_s[:, :, :, None], q_cd_c, torch.zeros_like(q_cd_c))
        e_f = (X3[:, None] - scale_c[:, :, :, None] * code_c.to(torch.float32)).reshape(B, Q3_N_CAND, QK3_K)
        sse_c = torch.sum(Wf[:, None, :] * e_f * e_f, dim=2)
        if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0):
            half = QK3_K // 2
            v_c = e_f[:, :, :half] + e_f[:, :, half:]
            dc_c = torch.sum(e_f, dim=2)
            cand_err_t = sse_c + lam_vw * torch.sum(v_c * v_c, dim=2) + lam_dc * dc_c * dc_c
        else:
            cand_err_t = sse_c

        # 3. 6-State Sieve + 6-State Viterbi DP over the super-scale lattice:
        best_cand_np = _q3_sieve_and_viterbi(cand_err_t.cpu().numpy(), cfg)
        best_cand_t = torch.from_numpy(best_cand_np).to(dev)
        d_sel = cand_d[rows, best_cand_t]

        # 4. Asymmetric Sign-Flip fp16-ULP Polish (+/-8 ULPs around +d_sel and -d_sel):
        rad = max(int(cfg.ulp_radius), 4)
        deltas = torch.arange(-rad, rad + 1, dtype=torch.int32, device=dev)
        abs_bits = torch.abs(d_sel).to(torch.float16).view(torch.int16).to(torch.int32) & 0x7FFF
        pos_b = torch.clamp(abs_bits[:, None] + deltas[None, :], 0, 0x7BFF)
        neg_b = pos_b | 0x8000
        all_b = torch.cat([pos_b, neg_b], dim=1)
        all_d = all_b.to(torch.int16).view(torch.float16).to(torch.float32)
        act_u = torch.abs(all_d) > 1.0e-15
        safe_u = torch.where(act_u, all_d, torch.ones_like(all_d))
        q_sc_u = _nearest_int_t(s[:, None, :] / safe_u[:, :, None]).clamp(Q3_SC_MIN, Q3_SC_MAX)
        sc_u = torch.where(act_u[:, :, None], q_sc_u, torch.zeros_like(q_sc_u))
        scale_u = all_d[:, :, None] * sc_u.to(torch.float32)
        act_su = torch.abs(scale_u) > 1.0e-15
        safe_su = torch.where(act_su, scale_u, torch.ones_like(scale_u))
        q_cd_u = _nearest_int_t(X3[:, None] / safe_su[:, :, :, None]).clamp(Q3_CODE_MIN, Q3_CODE_MAX)
        code_u = torch.where(act_su[:, :, :, None], q_cd_u, torch.zeros_like(q_cd_u))
        e_u = X3[:, None] - scale_u[:, :, :, None] * code_u.to(torch.float32)
        sse_u = torch.sum(W3[:, None] * e_u * e_u, dim=(2, 3))
        pick_u = torch.argmin(sse_u, dim=1)
        best_d16 = all_b[rows, pick_u]
        best_d = all_d[rows, pick_u]

        # Final (sc, code) derived from the polished super-scale d:
        act_b = torch.abs(best_d) > 1.0e-15
        safe_b = torch.where(act_b, best_d, torch.ones_like(best_d))
        q_bs = _nearest_int_t(s / safe_b[:, None]).clamp(Q3_SC_MIN, Q3_SC_MAX)
        sc = torch.where(act_b[:, None], q_bs, torch.zeros_like(q_bs))
        scale = best_d[:, None] * sc.to(torch.float32)
        act_s = torch.abs(scale) > 1.0e-15
        safe_s = torch.where(act_s, scale, torch.ones_like(scale))
        q_bd = _nearest_int_t(X3 / safe_s[:, :, None]).clamp(Q3_CODE_MIN, Q3_CODE_MAX)
        code = torch.where(act_s[:, :, None], q_bd, torch.zeros_like(q_bd))
        code_f = code.reshape(B, QK3_K)

        # 5. Closed-Form Vesica/DC Code Projection within sse_budget:
        if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0):
            d_w = (best_d[:, None] * sc.to(torch.float32)).repeat_interleave(Q3_K_SUBBLOCK, dim=1)
            e_f64 = Xf - d_w * code_f.to(torch.float32)
            sse_live = torch.sum(Wf * e_f64 * e_f64, dim=1)
            sse_cap = sse_live * (1.0 + float(cfg.sse_budget))
            half = QK3_K // 2
            v = e_f64[:, :half] + e_f64[:, half:]
            dc = torch.sum(e_f64, dim=1)
            denom = (Wf + lam_vw + lam_dc) * (d_w * d_w)
            act_bw = torch.abs(d_w) > 1.0e-15
            s_den = torch.where(act_bw & (denom > 1.0e-30), denom, torch.ones_like(denom))
            v_full = torch.cat([v, v], dim=1)
            grad = d_w * (Wf * e_f64 + lam_vw * v_full + lam_dc * dc[:, None])
            dq = torch.where(
                act_bw,
                _nearest_int_t(grad / s_den),
                torch.zeros_like(code_f),
            )
            code_new = torch.clamp(code_f + dq, Q3_CODE_MIN, Q3_CODE_MAX)
            e_new = Xf - d_w * code_new.to(torch.float32)
            sse_new = torch.sum(Wf * e_new * e_new, dim=1)
            v_new = e_new[:, :half] + e_new[:, half:]
            dc_new = torch.sum(e_new, dim=1)
            spec_old = lam_vw * torch.sum(v * v, dim=1) + lam_dc * dc * dc
            spec_new = lam_vw * torch.sum(v_new * v_new, dim=1) + lam_dc * dc_new * dc_new
            accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-12)
            code = torch.where(accept[:, None], code_new, code_f).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)

        # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing):
        #    On CUDA the anneal keeps every tensor on-device — the CPU/numpy
        #    chain previously forced a full D2H→CPU→D2H round-trip per chunk
        #    and made the GPU idle for the whole refinery.
        if cfg.therm_passes > 0:
            scale_rp = best_d[:, None] * sc.to(torch.float32)   # (B, 16) per-superblock
            thermal_c, _ = thermal_refine_q3_k_cuda(
                Xf, None if Wf is None else Wf, scale_rp,
                code.reshape(B, QK3_K), cfg, seed=therm_seed,
            )
            # thermal_refine_q3_k_cuda already returns source-signed codes [-4, 3]
            code = thermal_c.reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK).to(torch.int16)

        if as_tensor:
            return best_d16, best_d, sc, code.reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
        return (
            best_d16.cpu().numpy().astype(np.uint16),
            sc.cpu().numpy().astype(np.int16),
            code.cpu().numpy().astype(np.int16),
        )


def _quantize_q3_k_hpc_cuda(
    X: np.ndarray, W: np.ndarray, dev: "torch.device", cfg: HPCConfig,
    block_start: int = 0,
) -> tuple[bytes, float]:
    """Full HPC V2 Q3_K pipeline on CUDA for a chunk of (B, 256) blocks.

    X/W are uploaded once.  Level A scales, Level B fit, the zero-DC pin and the
    raw-SSE reduction all stay on the device; the only D2H copies are the
    Sieve/Viterbi candidate-error matrices (sequential DP, runs in numpy) and
    the final (d16, sc, code) needed for byte packing.
    """
    B = X.shape[0]

    with torch.no_grad():
        X_t = torch.from_numpy(np.ascontiguousarray(X, dtype=np.float32)).to(dev)
        W_t = torch.from_numpy(np.ascontiguousarray(W, dtype=np.float32)).to(dev)

        sub_scales_t = _subblock_scales_hpc_cuda(
            X_t.reshape(B * Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK),
            W_t.reshape(B * Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK),
            dev,
            cfg,
            as_tensor=True,
        )
        d16_t, d_t, sc_t, code_t = _fit_q3_k_blocks_cuda(
            X_t, W_t, sub_scales_t, dev, cfg, as_tensor=True,
            therm_seed=(
                cfg.therm_seed
                ^ ((0x9E3779B97F4A7C15 * block_start) & 0x3FFFFFFFFFFFFFFF)
            ),
        )
        del sub_scales_t, W_t

        # Zero-DC Pin (deterministic code-flip exchange within pin_budget), on
        # the device.  ``d_t`` is already the fp16-decoded super-scale, so no
        # uint16 numpy -> torch round-trip is needed (torch's uint16 support is
        # version-dependent and fails on older releases).
        if cfg.pin_dc:
            codes_pin, _ = pin_dc_codes_uniform_cuda(
                X_t,
                d_t,
                code_t.reshape(B, QK3_K).to(torch.int64),
                sc_t.to(torch.float32).repeat_interleave(Q3_K_SUBBLOCK, dim=1),
                Q3_CODE_MIN,
                Q3_CODE_MAX,
                budget_frac=float(cfg.pin_budget),
            )
            code_t = codes_pin.reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)

        deq = d_t[:, None, None] * sc_t[:, :, None].to(torch.float32) * code_t.to(torch.float32)
        raw_sse = float(torch.sum((X_t - deq.reshape(B, QK3_K)) ** 2, dtype=torch.float64).item())

        d16 = d16_t.cpu().numpy().astype(np.uint16)
        sc = sc_t.cpu().numpy().astype(np.int16)
        code = code_t.cpu().numpy().astype(np.int16)

    hmask, qs = _pack_q3_k_codes(code)
    scales = _pack_q3_k_scales(sc)
    out = np.empty((B, Q3_K_BLOCK_BYTES), dtype=np.uint8)
    out[:, : QK3_K // 8] = hmask
    out[:, QK3_K // 8 : QK3_K // 8 + QK3_K // 4] = qs
    out[:, QK3_K // 8 + QK3_K // 4 : Q3_K_BLOCK_BYTES - 2] = scales
    out[:, Q3_K_BLOCK_BYTES - 2 :] = d16.astype("<u2").view(np.uint8).reshape(B, 2)
    return out.tobytes(), raw_sse


def _quantize_q3_k_hpc_cpu(
    X: np.ndarray, W: np.ndarray, cfg: HPCConfig
) -> tuple[bytes, float]:
    """Full HPC V2 Q3_K pipeline on CPU for a chunk of (B, 256) blocks."""
    B = X.shape[0]
    X_f32 = X.astype(np.float32)
    W_f32 = W.astype(np.float32)

    sub_scales = _subblock_scales_hpc(
        X_f32.reshape(B * Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK),
        W_f32.reshape(B * Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK),
        cfg,
    )
    d16, sc, code = _fit_q3_k_blocks(X_f32, W_f32, sub_scales, cfg)

    # Zero-DC Pin (deterministic code-flip exchange within pin_budget):
    if cfg.pin_dc:
        code_pin, _ = pin_dc_codes_uniform(
            X_f32,
            from_fp16_bits(d16),
            code.reshape(B, QK3_K).astype(np.int64),
            np.repeat(sc, Q3_K_SUBBLOCK, axis=1),
            Q3_CODE_MIN,
            Q3_CODE_MAX,
            budget_frac=float(cfg.pin_budget),
        )
        code = code_pin.reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK).astype(np.int16)

    hmask, qs = _pack_q3_k_codes(code)
    scales = _pack_q3_k_scales(sc)
    out = np.empty((B, Q3_K_BLOCK_BYTES), dtype=np.uint8)
    out[:, : QK3_K // 8] = hmask
    out[:, QK3_K // 8 : QK3_K // 8 + QK3_K // 4] = qs
    out[:, QK3_K // 8 + QK3_K // 4 : Q3_K_BLOCK_BYTES - 2] = scales
    out[:, Q3_K_BLOCK_BYTES - 2 :] = d16.astype("<u2").view(np.uint8).reshape(B, 2)

    d = from_fp16_bits(d16).astype(np.float32)
    deq = (d[:, None, None] * sc[:, :, None].astype(np.float32) * code.astype(np.float32)).reshape(B, QK3_K)
    raw_sse = float(np.sum((X_f32 - deq) ** 2, dtype=np.float64))
    return out.tobytes(), raw_sse


def quantize_tensor_q3_k(
    weights: np.ndarray,
    importance: np.ndarray | None = None,
    cfg: HPCConfig | None = None,
    chunk_blocks: int | None = None,
) -> tuple[bytes, int, float]:
    """Quantize a tensor to GGML Q3_K using the full HPC V2 pipeline.

    ``chunk_blocks`` (256-weight super-blocks per chunk) defaults to
    ``Q3_K_CHUNK_BLOCKS_CUDA`` on the GPU path (same weights-per-chunk as the
    Q4_0 reference) and ``Q3_K_CHUNK_BLOCKS_CPU`` on the CPU path.
    """
    cfg = cfg or HPCConfig()
    x = np.asarray(weights, dtype=np.float32).ravel()
    if x.size == 0 or x.size % QK3_K != 0:
        raise ValueError("Q3_K requires a non-empty weight count divisible by 256")
    n_blocks = x.size // QK3_K
    X = x.reshape(n_blocks, QK3_K)

    if importance is not None:
        imp = np.asarray(importance, dtype=np.float32).ravel()
        if imp.size != x.size:
            raise ValueError("importance length must match weight count")
        W_raw = imp.reshape(n_blocks, QK3_K)
    else:
        W_raw = None

    cuda_ok = cfg.device != "cpu" and cuda_is_available()
    use_cuda = cuda_ok and (
        cfg.device.startswith("cuda") or (cfg.device == "auto" and n_blocks >= 64)
    )
    dev = torch.device("cuda" if cfg.device == "auto" else cfg.device) if use_cuda else None
    if chunk_blocks is None:
        if use_cuda and dev is not None:
            # VRAM-aware default: size the chunk from the free memory at
            # start-up, never below the reference chunk (which always fits
            # even on small cards).
            try:
                free_bytes, _ = torch.cuda.mem_get_info(dev)
                chunk_blocks = max(
                    Q3_K_CHUNK_BLOCKS_CUDA,
                    int(Q3_K_CUDA_VRAM_FRACTION * free_bytes) // Q3_K_CUDA_BYTES_PER_BLOCK,
                )
            except Exception:
                chunk_blocks = Q3_K_CHUNK_BLOCKS_CUDA
        else:
            chunk_blocks = Q3_K_CHUNK_BLOCKS_CPU
    chunk_blocks = max(1, int(chunk_blocks))
    _log.info(
        "Q3_K routing: cfg.device=%r cuda_is_available=%s n_blocks=%d -> %s (chunk_blocks=%d)",
        cfg.device, cuda_ok if cfg.device != "cpu" else "n/a", n_blocks,
        f"CUDA[{dev}]" if use_cuda else "CPU", chunk_blocks,
    )

    chunks: list[bytes] = []
    total_sse = 0.0
    for start in range(0, n_blocks, chunk_blocks):
        stop = min(start + chunk_blocks, n_blocks)
        Xb = X[start:stop]
        sig2 = np.mean(Xb * Xb, axis=1, keepdims=True)
        if W_raw is None:
            Wb = np.sqrt(sig2 + Xb * Xb)
        else:
            wr = W_raw[start:stop]
            Wb = np.sqrt(np.maximum(wr, 1.0e-6)) * np.sqrt(sig2 + Xb * Xb) if cfg.temper_imatrix else wr
        wm = np.mean(Wb, axis=1, keepdims=True)
        Wb = np.where(wm > 1.0e-30, Wb / np.maximum(wm, 1.0e-30), 1.0).astype(np.float32)

        if use_cuda and dev is not None:
            raw_b, sse_b = _quantize_q3_k_hpc_cuda(Xb, Wb, dev, cfg, block_start=start)
        else:
            raw_b, sse_b = _quantize_q3_k_hpc_cpu(Xb, Wb, cfg)
        chunks.append(raw_b)
        total_sse += sse_b

    rmse = float(np.sqrt(total_sse / x.size))
    return b"".join(chunks), n_blocks, rmse