Blame for seeds.py

"""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),
    }