| """Bit-exact Q2_K block representation, packing and reconstruction. |
| |
| Layout (84 bytes per 256 weights):: |
| |
| scales[16] uint8 low nibble = Ls[j], high nibble = Lm[j] |
| qs[64] uint8 four 2-bit codes per byte |
| d fp16 super-block scale |
| dmin fp16 super-block minimum scale |
| |
| The decoder reconstructs:: |
| |
| x_hat[16*j + k] = d * Ls[j] * L[16*j + k] - dmin * Lm[j] |
| |
| with ``L in {0,1,2,3}`` and ``Ls, Lm in {0..15}``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import numpy as np |
|
|
| from .fp16 import from_fp16_bits, nearest_int, to_fp16_bits |
|
|
| QK_K = 256 |
| NS = QK_K // 16 |
| NQ = 4 |
|
|
|
|
| @dataclass(slots=True) |
| class Q2KBlock: |
| scales: np.ndarray |
| qs: np.ndarray |
| d: np.uint16 |
| dmin: np.uint16 |
|
|
| @property |
| def d_float(self) -> np.float32: |
| return np.float32(from_fp16_bits(np.uint16(self.d))) |
|
|
| @property |
| def dmin_float(self) -> np.float32: |
| return np.float32(from_fp16_bits(np.uint16(self.dmin))) |
|
|
| def unpack_scales(self) -> tuple[np.ndarray, np.ndarray]: |
| scales = np.asarray(self.scales, dtype=np.uint8) |
| return scales & 0x0F, scales >> 4 |
|
|
| def dequantize(self) -> np.ndarray: |
| ls, lm = self.unpack_scales() |
| return dequantize_params( |
| self.d_float, self.dmin_float, ls, lm, unpack_codes(self.qs) |
| ) |
|
|
| def error(self, x: np.ndarray) -> float: |
| e = np.asarray(x, dtype=np.float64) - self.dequantize().astype(np.float64) |
| return float(np.dot(e, e)) |
|
|
|
|
| def unpack_codes(qs: np.ndarray) -> np.ndarray: |
| qs = np.asarray(qs, dtype=np.uint8) |
| codes = np.empty(QK_K, dtype=np.uint8) |
| for base in (0, 128): |
| q = qs[base // 4 : base // 4 + 32] |
| for shift_level in range(4): |
| lo = base + 32 * shift_level |
| codes[lo : lo + 32] = (q >> np.uint8(2 * shift_level)) & np.uint8(3) |
| return codes |
|
|
|
|
| def pack_codes(codes: np.ndarray) -> np.ndarray: |
| codes = np.asarray(codes, dtype=np.uint8) |
| if codes.shape != (QK_K,): |
| raise ValueError("codes must have shape (256,)") |
| qs = np.empty(QK_K // 4, dtype=np.uint8) |
| for base in (0, 128): |
| q = ( |
| codes[base : base + 32] |
| | (codes[base + 32 : base + 64] << np.uint8(2)) |
| | (codes[base + 64 : base + 96] << np.uint8(4)) |
| | (codes[base + 96 : base + 128] << np.uint8(6)) |
| ) |
| qs[base // 4 : base // 4 + 32] = q |
| return qs |
|
|
|
|
| def pack_scales(ls: np.ndarray, lm: np.ndarray) -> np.ndarray: |
| return (np.asarray(ls, dtype=np.uint8) & 0x0F) | ( |
| (np.asarray(lm, dtype=np.uint8) & 0x0F) << np.uint8(4) |
| ) |
|
|
|
|
| def dequantize_params( |
| d, dmin, ls: np.ndarray, lm: np.ndarray, codes: np.ndarray |
| ) -> np.ndarray: |
| ls = np.asarray(ls, dtype=np.float32) |
| lm = np.asarray(lm, dtype=np.float32) |
| codes = np.asarray(codes, dtype=np.float32) |
| d_sub = np.repeat(np.float32(d) * ls, 16) |
| m_sub = np.repeat(np.float32(dmin) * lm, 16) |
| return d_sub * codes - m_sub |
|
|
|
|
| def codes_from_params( |
| x: np.ndarray, d, dmin, ls: np.ndarray, lm: np.ndarray |
| ) -> np.ndarray: |
| x = np.asarray(x, dtype=np.float32) |
| d_sub = np.repeat(np.float32(d) * np.asarray(ls, dtype=np.float32), 16) |
| m_sub = np.repeat(np.float32(dmin) * np.asarray(lm, dtype=np.float32), 16) |
| safe = d_sub >= np.float32(1.0e-15) |
| codes = np.zeros(QK_K, dtype=np.uint8) |
| if np.any(safe): |
| q = nearest_int((x[safe] + m_sub[safe]) / d_sub[safe]) |
| codes[safe] = np.clip(q, 0, 3).astype(np.uint8) |
| return codes |
|
|
|
|
| def derive_subscales( |
| scales: np.ndarray, mins: np.ndarray, d, dmin |
| ) -> tuple[np.ndarray, np.ndarray]: |
| scales = np.asarray(scales, dtype=np.float32) |
| mins = np.asarray(mins, dtype=np.float32) |
| if float(d) > 1.0e-15: |
| ls = np.clip(nearest_int(scales / np.float32(d)), 0, 15).astype(np.uint8) |
| else: |
| ls = np.zeros(NS, dtype=np.uint8) |
| if float(dmin) > 1.0e-15: |
| lm = np.clip(nearest_int(mins / np.float32(dmin)), 0, 15).astype(np.uint8) |
| else: |
| lm = np.zeros(NS, dtype=np.uint8) |
| return ls, lm |
|
|
|
|
| def weighted_sse( |
| x: np.ndarray, |
| w: np.ndarray | None, |
| d, |
| dmin, |
| ls: np.ndarray, |
| lm: np.ndarray, |
| codes: np.ndarray, |
| ) -> float: |
| xh = dequantize_params(d, dmin, ls, lm, codes) |
| e = np.asarray(x, dtype=np.float64) - xh.astype(np.float64) |
| if w is None: |
| return float(np.dot(e, e)) |
| return float(np.dot(np.asarray(w, dtype=np.float64), e * e)) |
|
|
|
|
| def make_block(d, dmin, ls, lm, codes) -> Q2KBlock: |
| return Q2KBlock( |
| scales=pack_scales(ls, lm), |
| qs=pack_codes(codes), |
| d=np.uint16(to_fp16_bits(np.float32(d))), |
| dmin=np.uint16(to_fp16_bits(np.float32(dmin))), |
| ) |
|
|
|
|
| def blocks_to_arrays(blocks: list[Q2KBlock]) -> tuple[np.ndarray, np.ndarray, np.ndarray]: |
| n = len(blocks) |
| scales = np.empty((n, NS), dtype=np.uint8) |
| qs = np.empty((n, QK_K // 4), dtype=np.uint8) |
| dm = np.empty((n, 2), dtype=np.uint16) |
| for i, b in enumerate(blocks): |
| scales[i] = b.scales |
| qs[i] = b.qs |
| dm[i, 0] = b.d |
| dm[i, 1] = b.dmin |
| return scales, qs, dm |
|
|
|
|
| def blocks_to_bytes(scales: np.ndarray, qs: np.ndarray, dm: np.ndarray) -> bytes: |
| """Serialize Q2_K block arrays to GGUF's 84-byte block layout.""" |
| scales = np.asarray(scales, dtype=np.uint8) |
| qs = np.asarray(qs, dtype=np.uint8) |
| dm = np.asarray(dm, dtype=np.uint16) |
| n = scales.shape[0] |
| out = np.empty((n, 84), dtype=np.uint8) |
| out[:, :16] = scales |
| out[:, 16:80] = qs |
| out[:, 80:82] = np.ascontiguousarray(dm[:, 0]).view(np.uint8).reshape(n, 2) |
| out[:, 82:84] = np.ascontiguousarray(dm[:, 1]).view(np.uint8).reshape(n, 2) |
| return out.tobytes() |
|
|
|
|
| def dequantize_blocks(blocks: list[Q2KBlock]) -> np.ndarray: |
| if not blocks: |
| return np.empty(0, dtype=np.float32) |
| return np.concatenate([b.dequantize() for b in blocks]) |
|
|
|
|
| def codes_from_params_batch( |
| x: np.ndarray, d: np.ndarray, dmin: np.ndarray, ls: np.ndarray, lm: np.ndarray |
| ) -> np.ndarray: |
| x = np.asarray(x, dtype=np.float32) |
| n = x.shape[0] |
| d_sub = np.repeat( |
| np.asarray(d, dtype=np.float32)[:, None] * np.asarray(ls, dtype=np.float32), 16, axis=1 |
| ) |
| m_sub = np.repeat( |
| np.asarray(dmin, dtype=np.float32)[:, None] * np.asarray(lm, dtype=np.float32), 16, axis=1 |
| ) |
| safe = d_sub >= np.float32(1.0e-15) |
| codes = np.zeros((n, QK_K), dtype=np.uint8) |
| if np.any(safe): |
| q = nearest_int((x[safe] + m_sub[safe]) / d_sub[safe]) |
| codes[safe] = np.clip(q, 0, 3).astype(np.uint8) |
| return codes |
|
|
|
|
| def dequantize_params_batch( |
| d: np.ndarray, dmin: np.ndarray, ls: np.ndarray, lm: np.ndarray, codes: np.ndarray |
| ) -> np.ndarray: |
| n = codes.shape[0] |
| d_sub = np.repeat( |
| np.asarray(d, dtype=np.float32)[:, None] * np.asarray(ls, dtype=np.float32), 16, axis=1 |
| ) |
| m_sub = np.repeat( |
| np.asarray(dmin, dtype=np.float32)[:, None] * np.asarray(lm, dtype=np.float32), 16, axis=1 |
| ) |
| return d_sub * np.asarray(codes, dtype=np.float32) - m_sub |
|
|
|
|
| def weighted_sse_batch( |
| x: np.ndarray, |
| w: np.ndarray | None, |
| d: np.ndarray, |
| dmin: np.ndarray, |
| ls: np.ndarray, |
| lm: np.ndarray, |
| codes: np.ndarray, |
| ) -> np.ndarray: |
| xh = dequantize_params_batch(d, dmin, ls, lm, codes) |
| e = np.asarray(x, dtype=np.float64) - xh.astype(np.float64) |
| if w is None: |
| return np.sum(e * e, axis=1) |
| return np.sum(np.asarray(w, dtype=np.float64) * e * e, axis=1) |
|
|
|
|
| def pin_q2k_dc_residual(d, m, ls, lm, codes, X): |
| """Zero-DC pin: re-solve (d, m) on the exact ``sum(e) == 0`` line. |
| |
| With e_i = X_i - d*ls_g*c_i + m*lm_g, requiring Σe = 0 makes m a line |
| m(d) = (d·Σa - ΣX)/Σb (a = ls·c, b = lm); subtracting the residual mean |
| along that line is a translation that can only *lower* SSE, so the |
| closed-form optimum d* = Σ(u·v)/Σ(v²) dominates the input state on both |
| counts and leaves dc ~ 0 (bounded below by the fp16 storage of the |
| ``d``/``m`` fields). Blocks without a usable translation dof (Σb == 0 |
| or Σv² == 0) keep their input parameters. |
| """ |
| ls64 = np.asarray(ls, dtype=np.float64) |
| lm64 = np.asarray(lm, dtype=np.float64) |
| c64 = np.asarray(codes, dtype=np.float64) |
| a_el = np.repeat(ls64, 16, axis=1) * c64 |
| b_el = np.repeat(lm64, 16, axis=1) |
| sx0 = np.sum(np.asarray(X, dtype=np.float64), axis=1) |
| sa = a_el.sum(axis=1) |
| sb = b_el.sum(axis=1) |
| sb_safe = np.where(sb > 0.0, sb, np.nan) |
| u = np.asarray(X, dtype=np.float64) - b_el * (sx0 / sb_safe)[:, None] |
| v = a_el - b_el * (sa / sb_safe)[:, None] |
| vv = np.sum(v * v, axis=1) |
| uv = np.sum(u * v, axis=1) |
| ok = np.isfinite(vv) & (vv > 0.0) & np.isfinite(uv) & (sb > 0.0) |
| d_new = np.where(ok, uv / np.where(ok, vv, 1.0), np.asarray(d, dtype=np.float64)) |
| m_new = (d_new * sa - sx0) / np.where(ok, sb, 1.0) |
| m_new = np.where(ok, m_new, np.asarray(m, dtype=np.float64)) |
| |
| d_new = from_fp16_bits(to_fp16_bits(d_new.astype(np.float32))).astype(np.float64) |
| m_new = from_fp16_bits(to_fp16_bits(m_new.astype(np.float32))).astype(np.float64) |
| return d_new, m_new |
|
|
