"""GGUF k-quant packers - torch vectorized, float32 -> packed GGUF bytes.

Reverse of dequant_gguf_slice (gguf_packed.py). Each packer is a direct
translation of llama.cpp ggml-quants.c quantize_row_*_ref:

  Q8_0 (34 B):  d(f16) + qs[32] (int8)              y = d*q8              block 32
  Q2_K (84 B):  scales[16] + qs[64] + d(f16) + dmin(f16)
                y = d*sc*q - dmin*m  (q 0..3)       16 sub-blocks of 16
  Q3_K (110 B): hmask[32] + qs[64] + scales[12] + d(f16)
                y = d*sc*(q - (hm ? 0 : 4))         16 sub-blocks of 16
  Q4_K (144 B): d(f16) + dmin(f16) + scales[12] + qs[128]
                y = d*sc*q - dmin*m  (q 0..15)      8 sub-blocks of 32
  Q5_K (176 B): d(f16) + dmin(f16) + scales[12] + qh[32] + qs[128]
                y = d*sc*(q+16*bit5) - dmin*m       8 sub-blocks of 32
  Q6_K (210 B): ql[128] + qh[64] + scales[16] + d(f16)
                y = d*sc*(q6-32)                    16 sub-blocks of 16

i-quants are NOT packed here - kquant/iq.py quantize_iq*_s() already returns
packed GGUF bytes.
"""

from __future__ import annotations

import torch

QK_K = 256
EPS = 1e-15


def _nearest_int(x: torch.Tensor) -> torch.Tensor:
    """Round half away from zero (C nearest_int / llm_rintf)."""
    return torch.where(x >= 0, torch.floor(x + 0.5), torch.ceil(x - 0.5))


def _f16_bytes(t: torch.Tensor) -> torch.Tensor:
    """fp32 [..., n] -> uint8 [..., n, 2] little-endian f16 bytes."""
    u = (t.half().view(torch.int16).to(torch.int64)) & 0xFFFF
    return torch.stack([u & 0xFF, (u >> 8) & 0xFF], dim=-1).to(torch.uint8)


def _cat_blocks(segments: list[torch.Tensor], out_f: int, nblk: int,
                bpb: int) -> torch.Tensor:
    """Concatenate per-block byte segments into one [out_f, nblk*bpb] row buffer.

    Each segment is [out_f, nblk, seg_len] in the GGUF block layout; the result
    is row-major [out_f, nblk*bpb] (no Python row loop). Every segment is cast
    to uint8 first: a single int64 segment would otherwise promote the whole
    cat to int64 and silently widen the packed bytes 8x.
    """
    segs = [s.to(torch.uint8).reshape(out_f, nblk, -1) for s in segments]
    return torch.cat(segs, dim=-1).reshape(out_f, nblk * bpb)


def _make_qkx2(nmax: int, x: torch.Tensor, weights: torch.Tensor, rmin: float,
               rdelta: float, nstep: int, use_mad: bool
               ) -> tuple[torch.Tensor, torch.Tensor]:
    """Vectorized make_qkx2_quants (ggml-quants.c, line 799).

    Args:
        nmax: max level (15 Q4_K, 31 Q5_K, 3 Q2_K).
        x: [..., n] values.
        weights: [..., n] per-value weights.
    Returns:
        (scale, the_min) each [..., 1].
    """
    n = x.shape[-1]
    min_v = x.min(dim=-1, keepdim=True).values
    max_v = x.max(dim=-1, keepdim=True).values
    min_v = torch.where(min_v > 0, torch.zeros_like(min_v), min_v)  # if min > 0: min = 0

    sum_w = weights.sum(dim=-1, keepdim=True)
    sum_x = (weights * x).sum(dim=-1, keepdim=True)

    # Initial quantization with iscale = nmax/(max-min).
    iscale = nmax / (max_v - min_v).clamp_min(EPS)
    scale = 1.0 / iscale
    l0 = _nearest_int(iscale * (x - min_v)).clamp(0, nmax)
    diff = scale * l0 + min_v - x
    best_error = (weights * (diff.abs() if use_mad else diff * diff)).sum(dim=-1, keepdim=True)

    # Candidate sweep: is in 0..nstep, iscale = (rmin + rdelta*is + nmax)/(max-min).
    is_grid = rmin + rdelta * torch.arange(nstep + 1, device=x.device, dtype=x.dtype) + nmax
    isc_all = is_grid.view(-1, *([1] * x.dim())) / (max_v - min_v).clamp_min(EPS).unsqueeze(0)
    l_all = _nearest_int(isc_all * (x - min_v).unsqueeze(0)).clamp(0, nmax)  # [S, ..., n]
    wl = weights.unsqueeze(0) * l_all
    sum_l = wl.sum(dim=-1, keepdim=True)
    sum_l2 = (wl * l_all).sum(dim=-1, keepdim=True)
    sum_xl = (weights.unsqueeze(0) * l_all * x.unsqueeze(0)).sum(dim=-1, keepdim=True)

    D = sum_w.unsqueeze(0) * sum_l2 - sum_l * sum_l
    this_scale = (sum_w.unsqueeze(0) * sum_xl - sum_x.unsqueeze(0) * sum_l) / D.clamp_min(EPS)
    this_min = (sum_l2 * sum_x.unsqueeze(0) - sum_l * sum_xl) / D.clamp_min(EPS)
    # if this_min > 0: this_min = 0; this_scale = sum_xl / sum_l2
    pos_min = this_min > 0
    this_scale = torch.where(pos_min, sum_xl / sum_l2.clamp_min(EPS), this_scale)
    this_min = torch.where(pos_min, torch.zeros_like(this_min), this_min)

    diff2 = this_scale * l_all + this_min - x.unsqueeze(0)
    cur_error = (weights.unsqueeze(0) * (diff2.abs() if use_mad else diff2 * diff2)).sum(dim=-1, keepdim=True)

    better = cur_error < best_error                       # [S, ..., 1]
    any_better = better.any(dim=0)
    best_idx = better.int().argmax(dim=0)                 # [..., 1]
    scale = torch.where(any_better, this_scale.gather(0, best_idx.unsqueeze(0)).squeeze(0), scale)
    min_f = torch.where(any_better, this_min.gather(0, best_idx.unsqueeze(0)).squeeze(0), min_v)
    return scale, -min_f


def _pack_q2_q3(l: torch.Tensor) -> torch.Tensor:
    """Pack 256 x 2-bit levels into 64 bytes (llama.cpp q2_K/q3_K qs layout).

    qs[j/4 + l] = L[j+l] | L[j+l+32]<<2 | L[j+l+64]<<4 | L[j+l+96]<<6, j = 0, 128.
    l: [..., 256] uint8 in [0, 3].
    """
    v = l.to(torch.int64)
    lo = v[..., 0:32]; q1 = v[..., 32:64]; q2 = v[..., 64:96]; q3 = v[..., 96:128]
    lo2 = v[..., 128:160]; q12 = v[..., 160:192]; q22 = v[..., 192:224]; q32 = v[..., 224:256]
    first = lo | (q1 << 2) | (q2 << 4) | (q3 << 6)      # [..., 32]
    second = lo2 | (q12 << 2) | (q22 << 4) | (q32 << 6)
    return torch.cat([first, second], dim=-1).to(torch.uint8)  # [..., 64]


def _pack_scales_min_k4(sc: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
    """Pack 8 (scale, min) 6-bit pairs into 12 bytes (get_scale_min_k4 layout).

    j<4: q[j]=sc[j], q[j+4]=m[j]; j>=4: q[j+4]=(sc&0xF)|((m&0xF)<<4),
         q[j-4] |= (sc>>4)<<6, q[j] |= (m>>4)<<6.
    """
    sc = sc.clamp(0, 63).to(torch.int64)
    m = m.clamp(0, 63).to(torch.int64)
    q = torch.zeros(*sc.shape[:-1], 12, dtype=torch.int64, device=sc.device)
    q[..., 0:4] = sc[..., 0:4]
    q[..., 4:8] = m[..., 0:4]
    q[..., 8:12] = (sc[..., 4:8] & 0xF) | ((m[..., 4:8] & 0xF) << 4)
    q[..., 0:4] |= (sc[..., 4:8] >> 4) << 6
    q[..., 4:8] |= (m[..., 4:8] >> 4) << 6
    return q.to(torch.uint8)


def _unpack_scales_min_k4(q: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Mirror of _gguf_bits.unpack_scales_k4 for packed [..., 12] int64 bytes."""
    q = q.to(torch.int64)
    sc = torch.zeros(*q.shape[:-1], 8, dtype=torch.int64, device=q.device)
    m = torch.zeros_like(sc)
    sc[..., 0:4] = q[..., 0:4] & 0x3F
    m[..., 0:4] = q[..., 4:8] & 0x3F
    sc[..., 4:8] = (q[..., 8:12] & 0x0F) | ((q[..., 0:4] >> 6) << 4)
    m[..., 4:8] = (q[..., 8:12] >> 4) | ((q[..., 4:8] >> 6) << 4)
    return sc, m


def _pad(w: torch.Tensor, block: int) -> tuple[torch.Tensor, int]:
    out_f, in_f = w.shape
    nblk = (in_f + block - 1) // block
    wp = torch.zeros(out_f, nblk * block, dtype=torch.float32, device=w.device)
    wp[:, :in_f] = w
    return wp, nblk


# ---------------------------------------------------------------------------
# Q8_0 (llama.cpp:276) - block 32.
# ---------------------------------------------------------------------------

def pack_q8_0(w: torch.Tensor) -> torch.Tensor:
    out_f, in_f = w.shape
    wp, nblk = _pad(w, 32)
    x = wp.view(out_f, nblk, 32)
    amax = x.abs().amax(dim=-1).clamp_min(EPS)
    d = amax / 127.0
    qs = _nearest_int(x / d.unsqueeze(-1)).clamp(-127, 127).to(torch.int8).view(torch.uint8)
    return _cat_blocks([_f16_bytes(d), qs], out_f, nblk, 34)


# ---------------------------------------------------------------------------
# Q2_K (llama.cpp:891) - 16 sub-blocks of 16.
# ---------------------------------------------------------------------------

def pack_q2_k(w: torch.Tensor) -> torch.Tensor:
    out_f, in_f = w.shape
    wp, nblk = _pad(w, QK_K)
    x = wp.view(out_f, nblk, 16, 16)
    scale, m = _make_qkx2(3, x, x.abs(), -0.5, 0.1, 15, True)   # [.., 16, 1]
    scale = scale.squeeze(-1)
    m = m.squeeze(-1)
    max_scale = scale.amax(dim=-1).clamp_min(EPS)
    max_min = m.amax(dim=-1).clamp_min(EPS)

    ls = _nearest_int((15.0 / max_scale).unsqueeze(-1) * scale).clamp(0, 15).to(torch.int64)
    lm = _nearest_int((15.0 / max_min).unsqueeze(-1) * m).clamp(0, 15).to(torch.int64)
    scales = (ls & 0xF) | ((lm & 0xF) << 4)                    # [out, nb, 16]
    d = max_scale / 15.0
    dmin = max_min / 15.0

    # Re-quantize with stored d/dmin/scales.
    d_f = d.unsqueeze(-1).unsqueeze(-1)
    dm_f = dmin.unsqueeze(-1).unsqueeze(-1)
    sc16 = (scales & 0xF).to(torch.float32).unsqueeze(-1)
    m16 = (scales >> 4).to(torch.float32).unsqueeze(-1)
    l = _nearest_int((x + dm_f * m16) / (d_f * sc16).clamp_min(EPS)).clamp(0, 3).to(torch.int64)
    qs = _pack_q2_q3(l.view(out_f, nblk, QK_K))

    return _cat_blocks([scales.to(torch.uint8), qs,
                        _f16_bytes(d), _f16_bytes(dmin)],
                       out_f, nblk, 84)


# ---------------------------------------------------------------------------
# Q3_K (llama.cpp:1229) - 16 sub-blocks of 16.
# ---------------------------------------------------------------------------

def _make_q3_quants(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Vectorized make_q3_quants(16, 4, x, L, do_rmse=true).

    x: [..., 16]. Returns (scale, L) with L in [-4, 3] (int), scale [..., 1].
    """
    amax = x.abs().amax(dim=-1)
    max_v = x.gather(-1, x.abs().argmax(dim=-1, keepdim=True)).squeeze(-1)  # signed max|x| element
    max_s = torch.where(amax < EPS, torch.zeros_like(max_v), max_v)
    iscale = -4.0 / max_s.clamp_min(EPS)
    l = _nearest_int(iscale.unsqueeze(-1) * x).clamp(-4, 3).to(torch.float32)
    w_ = x * x
    sumlx = (w_ * x * l).sum(dim=-1, keepdim=True)
    suml2 = (w_ * l * l).sum(dim=-1, keepdim=True)
    for _ in range(5):
        slx = sumlx - w_ * x * l                     # [..., 16] per-element
        sl2 = suml2 - w_ * l * l                     # [..., 16] per-element
        new_l = _nearest_int(x * sl2 / torch.where(slx > 0, slx, torch.ones_like(slx))).clamp(-4, 3)
        slx2 = slx + w_ * x * new_l                  # [..., 16]
        sl2_new = sl2 + w_ * new_l * new_l           # [..., 16]
        better = (slx > 0) & (sl2_new > 0) & (slx2 * slx2 * suml2 > sumlx * sumlx * sl2_new) & (new_l != l)
        l = torch.where(better, new_l, l)
        # Global sums update incrementally over accepted elements.
        sumlx = sumlx + (better * (slx2 - slx)).sum(dim=-1, keepdim=True)
        suml2 = suml2 + (better * (sl2_new - sl2)).sum(dim=-1, keepdim=True)
    scale = torch.where(suml2 > 0, sumlx / suml2.clamp_min(EPS), torch.zeros_like(sumlx))
    return scale, l


def pack_q3_k(w: torch.Tensor) -> torch.Tensor:
    out_f, in_f = w.shape
    wp, nblk = _pad(w, QK_K)
    x = wp.view(out_f, nblk, 16, 16)
    scale, _l = _make_q3_quants(x)                              # [.., 16, 1] per sub-block

    sc4 = scale.squeeze(-1)                                     # [out, nb, 16] signed
    max_sc = sc4.gather(-1, sc4.abs().argmax(dim=-1, keepdim=True)).squeeze(-1)  # [out, nb]
    max_sc = torch.where(sc4.abs().amax(dim=-1) < EPS, torch.zeros_like(max_sc), max_sc)

    isc = -32.0 / torch.where(max_sc.abs() < EPS, torch.ones_like(max_sc), max_sc)
    l8 = (_nearest_int(isc.unsqueeze(-1) * sc4).clamp(-32, 31) + 32).to(torch.int64)  # [.., 16] 0..63
    l4 = l8 & 0xF
    hi = (l8 >> 4) & 0x3
    scales = torch.zeros(*l8.shape[:-1], 12, dtype=torch.int64, device=w.device)
    scales[..., 0:8] = l4[..., 0:8]
    scales[..., 0:8] |= l4[..., 8:16] << 4
    hib = hi.view(*hi.shape[:-1], 4, 4).transpose(-1, -2)  # [.., j%4, j//4]
    hi_sh = (hib << (2 * torch.arange(4, device=w.device)).view(1, 1, 1, 4)).sum(-1)
    scales[..., 8:12] = hi_sh
    d = -max_sc / 32.0

    # Re-quantize levels with the actual stored scales (aux unpacking, same as
    # dequantize_row_q3_K / dequant_q3_k_packed_rows).
    s0 = (scales[..., 0] | (scales[..., 1] << 8) | (scales[..., 2] << 16) | (scales[..., 3] << 24)).to(torch.int64)
    s1 = (scales[..., 4] | (scales[..., 5] << 8) | (scales[..., 6] << 16) | (scales[..., 7] << 24)).to(torch.int64)
    s2 = (scales[..., 8] | (scales[..., 9] << 8) | (scales[..., 10] << 16) | (scales[..., 11] << 24)).to(torch.int64)
    kmask1 = 0x03030303
    kmask2 = 0x0F0F0F0F
    tmp = s2
    aux0 = (s0 & kmask2) | (((tmp >> 0) & kmask1) << 4)
    aux1 = (s1 & kmask2) | (((tmp >> 2) & kmask1) << 4)
    aux2 = ((s0 >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4)
    aux3 = ((s1 >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4)
    aux = torch.stack([aux0, aux1, aux2, aux3], dim=-1)        # [..., 4]
    byte = (aux.unsqueeze(-1) >> (8 * torch.arange(4, device=w.device))).reshape(
        *scales.shape[:-1], 16) & 0xFF
    sc8 = torch.where(byte >= 128, byte - 256, byte).to(torch.float32)
    sc = (sc8 - 32)
    dl = d.unsqueeze(-1).unsqueeze(-1) * sc.unsqueeze(-1)
    dl_safe = torch.where(dl == 0, torch.ones_like(dl), dl)   # C: if (!d) continue
    Lq = _nearest_int(x / dl_safe).clamp(-4, 3).to(torch.int64)  # [-4, 3]

    # hmask: bit (j%32) set when level in 4..7.
    lq_shift = Lq + 4
    is_hi = (lq_shift > 3).to(torch.int64)
    lq_final = (lq_shift - 4 * is_hi).reshape(out_f, nblk, QK_K)
    is_hi_flat = is_hi.reshape(out_f, nblk, QK_K)
    idx = torch.arange(QK_K, device=w.device)
    hmask = torch.zeros(out_f, nblk, 32, dtype=torch.int64, device=w.device)
    hmask.scatter_add_(-1, (idx % 32).view(1, 1, -1).expand(out_f, nblk, QK_K),
                       is_hi_flat << (idx // 32).view(1, 1, -1))
    qs = _pack_q2_q3(lq_final)

    return _cat_blocks([hmask.to(torch.uint8), qs, scales.to(torch.uint8),
                        _f16_bytes(d)], out_f, nblk, 110)


# ---------------------------------------------------------------------------
# Q4_K (llama.cpp:1457) - 8 sub-blocks of 32.
# ---------------------------------------------------------------------------

def pack_q4_k(w: torch.Tensor) -> torch.Tensor:
    out_f, in_f = w.shape
    wp, nblk = _pad(w, QK_K)
    x = wp.view(out_f, nblk, 8, 32)
    av_x = (x * x).mean(dim=-1, keepdim=True).sqrt()            # [.., 8, 1]
    weights = av_x + x.abs()
    scale, m = _make_qkx2(15, x, weights, -1.0, 0.1, 20, False)
    scale = scale.squeeze(-1)
    m = m.squeeze(-1)
    max_scale = scale.amax(dim=-1).clamp_min(EPS)
    max_min = m.amax(dim=-1).clamp_min(EPS)
    ls = _nearest_int((63.0 / max_scale).unsqueeze(-1) * scale).clamp(0, 63)
    lm = _nearest_int((63.0 / max_min).unsqueeze(-1) * m).clamp(0, 63)
    q = _pack_scales_min_k4(ls, lm)                             # [out, nb, 12]
    d = max_scale / 63.0
    dmin = max_min / 63.0

    d_f = d.unsqueeze(-1).unsqueeze(-1)
    dm_f = dmin.unsqueeze(-1).unsqueeze(-1)
    sc, mm = _unpack_scales_min_k4(q)
    dl = d_f * sc.to(torch.float32).unsqueeze(-1)
    ml = dm_f * mm.to(torch.float32).unsqueeze(-1)
    l = _nearest_int((x + ml) / dl.clamp_min(EPS)).clamp(0, 15).to(torch.int64)
    l2 = l.reshape(out_f, nblk, 4, 2, 32)
    qs = (l2[..., 0, :] | (l2[..., 1, :] << 4)).reshape(out_f, nblk, 128)

    return _cat_blocks([_f16_bytes(d), _f16_bytes(dmin), q, qs],
                       out_f, nblk, 144)


# ---------------------------------------------------------------------------
# Q5_K (llama.cpp:1644) - 8 sub-blocks of 32.
# ---------------------------------------------------------------------------

def pack_q5_k(w: torch.Tensor) -> torch.Tensor:
    out_f, in_f = w.shape
    wp, nblk = _pad(w, QK_K)
    x = wp.view(out_f, nblk, 8, 32)
    av_x = (x * x).mean(dim=-1, keepdim=True).sqrt()
    weights = av_x + x.abs()
    scale, m = _make_qkx2(31, x, weights, -0.5, 0.1, 15, False)
    scale = scale.squeeze(-1)
    m = m.squeeze(-1)
    max_scale = scale.amax(dim=-1).clamp_min(EPS)
    max_min = m.amax(dim=-1).clamp_min(EPS)
    ls = _nearest_int((63.0 / max_scale).unsqueeze(-1) * scale).clamp(0, 63)
    lm = _nearest_int((63.0 / max_min).unsqueeze(-1) * m).clamp(0, 63)
    q = _pack_scales_min_k4(ls, lm)
    d = max_scale / 63.0
    dmin = max_min / 63.0

    d_f = d.unsqueeze(-1).unsqueeze(-1)
    dm_f = dmin.unsqueeze(-1).unsqueeze(-1)
    sc, mm = _unpack_scales_min_k4(q)
    dl = d_f * sc.to(torch.float32).unsqueeze(-1)
    ml = dm_f * mm.to(torch.float32).unsqueeze(-1)
    l = _nearest_int((x + ml) / dl.clamp_min(EPS)).clamp(0, 31).to(torch.int64)

    # ql + qh: for g in 0..3 (groups of 64): pair of 32-elem sub-blocks.
    # Each sub-block s (0..7) contributes bit s when its level > 15.
    hi_s = (l > 15).to(torch.int64)                     # [out, nblk, 8, 32]
    qh = (hi_s * (1 << torch.arange(8, device=w.device)).view(1, 1, 8, 1)).sum(-2)  # [out, nblk, 32]
    lo1 = torch.where(hi_s.bool(), l - 16, l)           # [out, nblk, 8, 32]
    ql = (lo1[..., 0::2] | (lo1[..., 1::2] << 4)).reshape(out_f, nblk, 128)

    return _cat_blocks([_f16_bytes(d), _f16_bytes(dmin), q,
                        qh.to(torch.uint8), ql], out_f, nblk, 176)


# ---------------------------------------------------------------------------
# Q6_K (llama.cpp:1869) - 16 sub-blocks of 16.
# ---------------------------------------------------------------------------

def _make_qx_quants_rmse(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Vectorized make_qx_quants(16, 32, x, L, rmse_type=1, qw=NULL) (ggml-quants.c:628).

    RMSE-weighted scale: scale = sumlx/suml2 with w = x^2.
    x: [..., 16]. Returns (scale [..., 1], L in [-32, 31] as float).
    """
    amax = x.abs().amax(dim=-1)
    max_v = x.gather(-1, x.abs().argmax(dim=-1, keepdim=True)).squeeze(-1)
    max_s = torch.where(amax < EPS, torch.zeros_like(max_v), max_v)
    iscale = -32.0 / max_s.clamp_min(EPS)
    l = _nearest_int(iscale.unsqueeze(-1) * x).clamp(-32, 31).to(torch.float32)
    w_ = x * x
    sumlx = (w_ * x * l).sum(dim=-1, keepdim=True)
    suml2 = (w_ * l * l).sum(dim=-1, keepdim=True)
    scale = torch.where(suml2 > 0, sumlx / suml2.clamp_min(EPS), torch.zeros_like(sumlx))
    return scale, l


def pack_q6_k(w: torch.Tensor) -> torch.Tensor:
    out_f, in_f = w.shape
    wp, nblk = _pad(w, QK_K)
    x = wp.view(out_f, nblk, 16, 16)
    scale, _l = _make_qx_quants_rmse(x)                         # [.., 16, 1]

    sc4 = scale.squeeze(-1)                                     # [out, nb, 16] signed
    max_sc = sc4.gather(-1, sc4.abs().argmax(dim=-1, keepdim=True)).squeeze(-1)
    max_sc = torch.where(sc4.abs().amax(dim=-1) < EPS, torch.zeros_like(max_sc), max_sc)

    isc = -128.0 / torch.where(max_sc.abs() < EPS, torch.ones_like(max_sc), max_sc)
    sc16 = _nearest_int(isc.unsqueeze(-1) * sc4).clamp(-127, 127).to(torch.int64)
    d = -max_sc / 128.0

    dl = d.unsqueeze(-1).unsqueeze(-1) * sc16.to(torch.float32).unsqueeze(-1)
    dl_safe = torch.where(dl == 0, torch.ones_like(dl), dl)   # C: if (!d) continue
    l = (_nearest_int(x / dl_safe).clamp(-32, 31).to(torch.int64) + 32).to(torch.int64)  # 0..63

    lo = l & 0xF
    hi = (l >> 4) & 0x3
    seg = l.view(out_f, nblk, 2, 8, 16)                 # [.., h, 8 subs, 16]
    seg_lo = seg & 0xF
    seg_hi = (seg >> 4) & 0x3
    ql_h = torch.stack([
        seg_lo[..., 0, :] | (seg_lo[..., 4, :] << 4),
        seg_lo[..., 1, :] | (seg_lo[..., 5, :] << 4),
        seg_lo[..., 2, :] | (seg_lo[..., 6, :] << 4),
        seg_lo[..., 3, :] | (seg_lo[..., 7, :] << 4),
    ], dim=-2)                                   # [out, nblk, 2, 4, 16]
    ql = ql_h.reshape(out_f, nblk, 128)
    qh_h = torch.stack([
        seg_hi[..., 0, :] | (seg_hi[..., 2, :] << 2)
        | (seg_hi[..., 4, :] << 4) | (seg_hi[..., 6, :] << 6),
        seg_hi[..., 1, :] | (seg_hi[..., 3, :] << 2)
        | (seg_hi[..., 5, :] << 4) | (seg_hi[..., 7, :] << 6),
    ], dim=-2)                                   # [out, nblk, 2, 2, 16]
    qh = qh_h.reshape(out_f, nblk, 64)

    return _cat_blocks([ql, qh, sc16.to(torch.uint8), _f16_bytes(d)],
                       out_f, nblk, 210)


_GGUF_PACK = {
    8: pack_q8_0,
    10: pack_q2_k,
    11: pack_q3_k,
    12: pack_q4_k,
    13: pack_q5_k,
    14: pack_q6_k,
}


def pack_gguf_bytes(w: torch.Tensor, gguf_dtype: int) -> torch.Tensor:
    """Pack fp32 [out, in] weight into GGUF bytes for a Q*_K / Q8_0 dtype."""
    fn = _GGUF_PACK.get(gguf_dtype)
    if fn is None:
        raise ValueError(f"pack_gguf_bytes: no packer for GGML dtype {gguf_dtype}")
    return fn(w)


def packer_for(gguf_dtype: int):
    """Return the packer callable for a GGML dtype, or None for i-quants."""
    return _GGUF_PACK.get(gguf_dtype)


def supported_pack_dtypes() -> list[int]:
    return sorted(_GGUF_PACK.keys())