| """Batched GGML seed quantizers and activation-aware seed weighting. |
| |
| The C engine performs ``make_qkx2_quants`` once per 16-weight sub-block and |
| ``make_qp_quants`` once per 16-sub-block superblock. V2 batches both across |
| every sub-block / superblock in the tensor with NumPy, which removes the |
| per-block Python overhead while preserving the exact search rules. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from .config import HPCConfig |
| from .fp16 import nearest_int |
|
|
|
|
| def make_qkx2_batch( |
| x: np.ndarray, w: np.ndarray, nmax: int = 3 |
| ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: |
| """Batched ``make_qkx2_quants``. Returns (scale, min_abs, codes).""" |
| x = np.asarray(x, dtype=np.float64) |
| w = np.asarray(w, dtype=np.float64) |
| rows, n = x.shape |
| xmin = x.min(axis=1) |
| xmax = x.max(axis=1) |
| xmin = np.minimum(xmin, 0.0) |
| span = xmax - xmin |
| degenerate = (xmax == xmin) | (span <= 1.0e-30) |
|
|
| sum_w = w.sum(axis=1) |
| sum_x = (w * x).sum(axis=1) |
|
|
| safe_span = np.where(span > 1.0e-30, span, 1.0) |
| iscale = nmax / safe_span |
| scale = 1.0 / iscale |
| codes = np.clip(nearest_int(iscale[:, None] * (x - xmin[:, None])), 0, nmax) |
| best_mad = np.sum(w * np.abs(scale[:, None] * codes + xmin[:, None] - x), axis=1) |
| best_scale = scale.copy() |
| best_min = xmin.copy() |
| best_codes = codes.copy() |
|
|
| for step in range(16): |
| try_iscale = (-0.5 + 0.1 * step + nmax) / safe_span |
| aux = np.clip(nearest_int(try_iscale[:, None] * (x - xmin[:, None])), 0, nmax) |
| sl = (w * aux).sum(axis=1) |
| sl2 = (w * aux * aux).sum(axis=1) |
| sxl = (w * aux * x).sum(axis=1) |
| det = sum_w * sl2 - sl * sl |
| valid = det > 1.0e-30 |
| safe_det = np.where(valid, det, 1.0) |
| this_scale = (sum_w * sxl - sum_x * sl) / safe_det |
| this_min = (sl2 * sum_x - sl * sxl) / safe_det |
| positive_min = this_min > 0.0 |
| this_min = np.where(positive_min, 0.0, this_min) |
| this_scale = np.where( |
| positive_min & (sl2 > 1.0e-30), sxl / np.where(sl2 > 1.0e-30, sl2, 1.0), this_scale |
| ) |
| mad = np.sum(w * np.abs(this_scale[:, None] * aux + this_min[:, None] - x), axis=1) |
| take = valid & (mad < best_mad) |
| if np.any(take): |
| best_mad = np.where(take, mad, best_mad) |
| best_scale = np.where(take, this_scale, best_scale) |
| best_min = np.where(take, this_min, best_min) |
| best_codes[take] = aux[take] |
|
|
| best_scale = np.where(degenerate, 0.0, best_scale) |
| best_codes[degenerate] = 0 |
| best_min = np.where(degenerate, xmin, best_min) |
| return best_scale, -best_min, best_codes.astype(np.uint8) |
|
|
|
|
| def make_qp_batch( |
| values: np.ndarray, weights: np.ndarray, nmax: int = 15 |
| ) -> tuple[np.ndarray, np.ndarray]: |
| """Batched ``make_qp_quants``. Factors positive values into d * code.""" |
| x = np.asarray(values, dtype=np.float64) |
| w = np.asarray(weights, dtype=np.float64) |
| rows, n = x.shape |
| xmax = x.max(axis=1) |
| dead = xmax < 1.0e-15 |
| safe_max = np.where(dead, 1.0, xmax) |
|
|
| iscale = nmax / safe_max |
| codes = np.clip(nearest_int(iscale[:, None] * x), 0, nmax) |
| scale = 1.0 / iscale |
| best_mse = np.sum(w * (x - scale[:, None] * codes) ** 2, axis=1) |
| best_iscale = iscale.copy() |
|
|
| for step in range(-4, 5): |
| if step == 0: |
| continue |
| trial_iscale = (0.1 * step + nmax) / safe_max |
| trial_scale = 1.0 / trial_iscale |
| trial_codes = np.clip(nearest_int(trial_iscale[:, None] * x), 0, nmax) |
| mse = np.sum(w * (x - trial_scale[:, None] * trial_codes) ** 2, axis=1) |
| take = mse < best_mse |
| if np.any(take): |
| best_mse = np.where(take, mse, best_mse) |
| best_iscale = np.where(take, trial_iscale, best_iscale) |
|
|
| codes = np.clip(nearest_int(best_iscale[:, None] * x), 0, nmax) |
| sumlx = (w * x * codes).sum(axis=1) |
| suml2 = (w * codes * codes).sum(axis=1) |
|
|
| for _ in range(5): |
| changed = False |
| for i in range(n): |
| wi = w[:, i] |
| xi = x[:, i] |
| li = codes[:, i].astype(np.float64) |
| slx = sumlx - wi * xi * li |
| sl2 = suml2 - wi * li * li |
| valid = (slx > 0.0) & (sl2 > 0.0) |
| safe_slx = np.where(valid, slx, 1.0) |
| new_l = np.clip(nearest_int(xi * sl2 / safe_slx), 0, nmax) |
| trial = valid & (new_l != codes[:, i]) |
| if not np.any(trial): |
| continue |
| slx2 = slx + wi * xi * new_l |
| sl2_2 = sl2 + wi * new_l * new_l |
| better = trial & (slx2 * slx2 * suml2 > sumlx * sumlx * sl2_2) |
| if np.any(better): |
| codes[better, i] = new_l[better].astype(np.uint8) |
| sumlx = np.where(better, slx2, sumlx) |
| suml2 = np.where(better, sl2_2, suml2) |
| changed = True |
| if not changed: |
| break |
|
|
| d = np.where(suml2 > 1.0e-30, sumlx / np.where(suml2 > 1.0e-30, suml2, 1.0), 0.0) |
| d = np.where(dead, 0.0, d) |
| codes[dead] = 0 |
| return d, codes.astype(np.uint8) |
|
|
|
|
| def seed_clamp_value(weights: np.ndarray, cfg: HPCConfig) -> float: |
| x = np.asarray(weights, dtype=np.float64).ravel() |
| mean_sq = float(np.mean(x * x)) |
| if mean_sq <= 1.0e-30: |
| return float("inf") |
| mean_fourth = float(np.mean(x * x * x * x)) |
| kurtosis = mean_fourth / (mean_sq * mean_sq) |
| sigma = float(np.sqrt(mean_sq)) |
| if kurtosis <= cfg.clamp_kurtosis_off: |
| return float("inf") |
| if kurtosis <= cfg.clamp_kurtosis_soft: |
| return cfg.clamp_sigma_soft * sigma |
| return cfg.clamp_sigma_hard * sigma |
|
|
|
|
| def seed_weights( |
| block: np.ndarray, importance: np.ndarray | None, sigma_sq: float, clamp_val: float |
| ) -> np.ndarray: |
| if importance is not None: |
| return np.asarray(importance, dtype=np.float64) |
| return np.sqrt(sigma_sq + np.minimum(block, clamp_val) ** 2) |
|
|
|
|
| def block_clamp_values(blocks: np.ndarray, cfg: HPCConfig) -> np.ndarray: |
| """Per-block kurtosis-gated outlier clamp thresholds (shape (n_blocks, 1, 1)).""" |
| b = np.asarray(blocks, dtype=np.float64).reshape(-1, QK_K_SEED := 256) |
| b2 = b * b |
| mean_sq = np.mean(b2, axis=1) |
| mean_fourth = np.mean(b2 * b2, axis=1) |
| kurt = np.where(mean_sq > 1.0e-30, mean_fourth / np.maximum(mean_sq * mean_sq, 1.0e-30), 3.0) |
| sigma = np.sqrt(np.maximum(mean_sq, 0.0)) |
| mult = np.where( |
| kurt <= cfg.clamp_kurtosis_off, |
| 1.0e9, |
| np.where(kurt <= cfg.clamp_kurtosis_soft, cfg.clamp_sigma_soft, cfg.clamp_sigma_hard), |
| ) |
| clamp = np.where(mean_sq > 1.0e-30, sigma * mult, np.inf) |
| return clamp[:, None, None] |
|
|
|
|
| def build_seeds( |
| x: np.ndarray, |
| importance: np.ndarray | None, |
| cfg: HPCConfig, |
| clamp_val: float | np.ndarray = float("inf"), |
| ) -> dict[str, np.ndarray]: |
| """Phase-1 seeds for all blocks: per-sub (scale,min) and shared (d,dmin).""" |
| n_blocks = x.shape[0] |
| blocks = np.asarray(x, dtype=np.float64).reshape(n_blocks, 16, 16) |
| per_blk_clamp = block_clamp_values(blocks, cfg) |
| clipped = np.clip(blocks, -per_blk_clamp, per_blk_clamp) |
| if importance is None: |
| sigma_sq = np.mean(blocks * blocks, axis=(1, 2), keepdims=True) |
| weights = np.sqrt(sigma_sq + clipped * clipped) |
| else: |
| weights = np.asarray(importance, dtype=np.float64).reshape(n_blocks, 16, 16) |
|
|
| flat_x = clipped.reshape(-1, 16) |
| flat_w = weights.reshape(-1, 16) |
| scales, mins, codes = make_qkx2_batch(flat_x, flat_w, nmax=3) |
| scales = scales.reshape(n_blocks, 16) |
| mins = mins.reshape(n_blocks, 16) |
| codes = codes.reshape(n_blocks, 256) |
| sw = weights.sum(axis=2) |
| d, ls = make_qp_batch(scales, sw, nmax=15) |
| dmin, lm = make_qp_batch(mins, sw, nmax=15) |
| return { |
| "scales": scales.astype(np.float32), |
| "mins": mins.astype(np.float32), |
| "codes": codes.astype(np.uint8), |
| "sw": sw.astype(np.float32), |
| "d": d.astype(np.float32), |
| "dmin": dmin.astype(np.float32), |
| "ls": ls.astype(np.uint8), |
| "lm": lm.astype(np.uint8), |
| "weights": weights.astype(np.float32), |
| } |
|
|
