quantgap 0.1.0__py3-none-any.whl

quantgap/probes/registry.py

@@ -0,0 +1,34 @@
+"""Probe set versions and their hashes. A probe set is frozen once released."""
+
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+
+from quantgap.models import ProbeSeq
+
+_DATA = Path(__file__).parent / "data"
+AVAILABLE = {"v1": "probes_v1.jsonl"}
+
+
+def probes_path(version: str = "v1") -> Path:
+    if version not in AVAILABLE:
+        raise ValueError(
+            f"Unknown probe set {version!r}. Available: {', '.join(sorted(AVAILABLE))}."
+        )
+    return _DATA / AVAILABLE[version]
+
+
+def probes_sha256(version: str = "v1") -> str:
+    return hashlib.sha256(probes_path(version).read_bytes()).hexdigest()
+
+
+def load_probes(version: str = "v1", path: Path | None = None) -> list[ProbeSeq]:
+    p = path or probes_path(version)
+    rows = []
+    with open(p, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                rows.append(ProbeSeq.model_validate_json(line))
+    return rows

quantgap/probes/render.py

@@ -0,0 +1,158 @@
+"""Turn probe rows into token ids and a score mask, using the reference model's tokenizer.
+
+Position `pos` predicts token `pos + 1`. A position is scored when token `pos + 1` is real
+content: not a special token, and for chat rows, inside the assistant reply.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+
+from quantgap.models import ProbeSeq
+
+MAX_LEN = 512
+GREEDY_PROMPT_TOKENS = 64
+_BOUNDARIES = ("\n", ". ", "。", "।", "؟", "? ", "! ", "! ", "}", ";")
+FALLBACK_TEMPLATE_NOTE = "no chat template on tokenizer: used plain 'User:/Assistant:' text"
+
+
+@dataclass
+class RenderedSeq:
+    seq_id: int
+    group: str
+    token_ids: list[int]
+    score_mask: list[bool]  # per token: is this token a scored target
+    n_scored: int = 0
+    notes: list[str] = field(default_factory=list)
+
+
+@dataclass
+class RenderedGreedy:
+    seq_id: int
+    group: str
+    prompt_ids: list[int]
+    lang: str = "en"
+
+
+def _encode(tok: Any, text: str) -> list[int]:
+    return list(tok(text, add_special_tokens=True)["input_ids"])
+
+
+def _truncate_text(tok: Any, text: str, max_len: int) -> str:
+    """Longest prefix of `text` that fits in max_len tokens, cut at a line or sentence end."""
+    if len(_encode(tok, text)) <= max_len:
+        return text
+    lo, hi = 1, len(text)
+    while lo < hi:  # binary search on characters
+        mid = (lo + hi + 1) // 2
+        if len(_encode(tok, text[:mid])) <= max_len:
+            lo = mid
+        else:
+            hi = mid - 1
+    cut = text[:lo]
+    best = max((cut.rfind(b) + len(b) for b in _BOUNDARIES if cut.rfind(b) > 0), default=0)
+    if best >= int(len(cut) * 0.6):
+        cut = cut[:best]
+    return cut.rstrip()
+
+
+def _special_ids(tok: Any) -> set[int]:
+    return {int(i) for i in getattr(tok, "all_special_ids", [])}
+
+
+def _ids_and_start(tok: Any, full: str, prefix: str, add_special: bool) -> tuple[list[int], int]:
+    """Token ids of `full` and the index of the first token that lies entirely after `prefix`.
+
+    A BPE token can straddle the boundary (for example " The" glues the space of the prompt to
+    the first word of the reply). Such a token counts as context and is not scored.
+    """
+    try:
+        enc = tok(full, add_special_tokens=add_special, return_offsets_mapping=True)
+        ids = list(enc["input_ids"])
+        offs = enc["offset_mapping"]
+        start = next((i for i, (a, _) in enumerate(offs) if a >= len(prefix)), len(ids))
+        return ids, start
+    except (NotImplementedError, TypeError, ValueError):  # slow tokenizer: no offsets
+        ids = list(tok(full, add_special_tokens=add_special)["input_ids"])
+        plen = len(tok(prefix, add_special_tokens=add_special)["input_ids"])
+        return ids, plen
+
+
+def _chat_ids(tok: Any, messages: list[dict[str, str]]) -> tuple[list[int], int, list[str]]:
+    """Return (ids, assistant_start_index, notes)."""
+    notes: list[str] = []
+    if getattr(tok, "chat_template", None):
+        prefix = tok.apply_chat_template(messages[:-1], add_generation_prompt=True, tokenize=False)
+        full = tok.apply_chat_template(messages, tokenize=False)
+        # The template already wrote any BOS token as text, so do not add special tokens again.
+        ids, start = _ids_and_start(tok, full, prefix, add_special=False)
+        if not full.startswith(prefix):
+            notes.append("chat prefix not a prefix of full render: scored all non-special tokens")
+            start = 0
+        return ids, start, notes
+    notes.append(FALLBACK_TEMPLATE_NOTE)
+    prefix = _plain_chat_prefix(messages[:-1]) + "Assistant:"
+    full = prefix + " " + messages[-1]["content"]
+    ids, start = _ids_and_start(tok, full, prefix, add_special=True)
+    return ids, start, notes
+
+
+def _plain_chat_prefix(messages: list[dict[str, str]]) -> str:
+    return "".join(
+        f"{'User' if m['role'] == 'user' else 'Assistant'}: {m['content']}\n" for m in messages
+    )
+
+
+def render_probes(
+    probes: list[ProbeSeq], tok: Any, max_len: int = MAX_LEN
+) -> tuple[list[RenderedSeq], list[RenderedGreedy]]:
+    """Render every probe with `tok`. Deterministic: same inputs give identical ids."""
+    special = _special_ids(tok)
+    seqs: list[RenderedSeq] = []
+    greedy: list[RenderedGreedy] = []
+    for p in probes:
+        notes: list[str] = []
+        if p.kind == "chat":
+            assert p.messages, f"probe {p.seq_id}: chat row without messages"
+            msgs = p.messages
+            ids, start, notes = _chat_ids(tok, msgs)
+            if len(ids) > max_len:
+                # Shorten the assistant reply, never the prompt.
+                reply = msgs[-1]["content"]
+                budget = max_len - start
+                short = _truncate_text(tok, reply, max(8, budget - 8))
+                msgs = msgs[:-1] + [{"role": msgs[-1]["role"], "content": short}]
+                ids, start, notes = _chat_ids(tok, msgs)
+                ids = ids[:max_len]
+        else:
+            assert p.text, f"probe {p.seq_id}: text row without text"
+            ids = _encode(tok, _truncate_text(tok, p.text, max_len))
+            start = 0
+        mask = [False] * len(ids)
+        for t in range(1, len(ids)):
+            mask[t] = t >= start and ids[t] not in special
+        seqs.append(
+            RenderedSeq(p.seq_id, p.group, ids, mask, n_scored=int(np.sum(mask)), notes=notes)
+        )
+        if p.greedy:
+            if p.kind == "chat":
+                gm = [m for m in (p.messages or [])[:-1]]
+                if getattr(tok, "chat_template", None):
+                    s = tok.apply_chat_template(gm, add_generation_prompt=True, tokenize=False)
+                    gids = list(tok(s, add_special_tokens=False)["input_ids"])
+                else:
+                    s = "".join(
+                        f"{'User' if m['role'] == 'user' else 'Assistant'}: {m['content']}\n"
+                        for m in gm
+                    )
+                    gids = _encode(tok, s + "Assistant:")
+            else:
+                gtext = p.greedy_prompt or p.text or ""
+                gids = _encode(tok, gtext)
+                if p.greedy_prompt is None:
+                    gids = gids[:GREEDY_PROMPT_TOKENS]
+            greedy.append(RenderedGreedy(p.seq_id, p.group, gids, p.lang))
+    return seqs, greedy

quantgap/report/chart.py

@@ -0,0 +1,92 @@
+"""Quality-vs-size chart. One series (the quants), one hue, direct labels, no legend box.
+
+Lower and further left is better: a smaller file with less drift. The line joins the quants no
+other quant beats on both size and KLD (the "best for its size" frontier).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+
+from quantgap.models import ScoreResult  # noqa: E402
+
+INK = "#1f2937"
+MUTED = "#6b7280"
+GRID = "#e5e7eb"
+MARK = "#2563eb"
+SURFACE = "#ffffff"
+FLAG = "#b45309"
+
+
+def pareto(results: list[ScoreResult]) -> list[ScoreResult]:
+    """Quants not beaten by a smaller-or-equal file with lower-or-equal KLD."""
+    out: list[ScoreResult] = []
+    best = float("inf")
+    for r in sorted(results, key=lambda r: (r.quant.size_bytes, r.metrics.kld_mean)):
+        if r.metrics.kld_mean < best:
+            out.append(r)
+            best = r.metrics.kld_mean
+    return out
+
+
+def draw_chart(
+    results: list[ScoreResult],
+    out: str | Path,
+    title: str | None = None,
+    budget_bytes: int | None = None,
+) -> Path:
+    out = Path(out)
+    fig, ax = plt.subplots(figsize=(8, 5), dpi=160)
+    fig.patch.set_facecolor(SURFACE)
+    ax.set_facecolor(SURFACE)
+    rs = sorted(results, key=lambda r: r.quant.size_bytes)
+    xs = [r.quant.size_bytes / 1e9 for r in rs]
+    ys = [r.metrics.kld_mean for r in rs]
+    lo = [r.metrics.kld_mean - r.metrics.kld_mean_ci95[0] for r in rs]
+    hi = [r.metrics.kld_mean_ci95[1] - r.metrics.kld_mean for r in rs]
+    ax.errorbar(
+        xs, ys, yerr=[[max(0, v) for v in lo], [max(0, v) for v in hi]],
+        fmt="none", ecolor=MARK, elinewidth=1, alpha=0.45, capsize=0,
+    )  # fmt: skip
+    front = pareto(rs)
+    ax.plot(
+        [r.quant.size_bytes / 1e9 for r in front], [r.metrics.kld_mean for r in front],
+        color=MARK, linewidth=1.5, alpha=0.5, zorder=2,
+    )  # fmt: skip
+    ax.scatter(xs, ys, s=64, color=MARK, edgecolor=SURFACE, linewidth=2, zorder=3)
+    labels = [r.quant.label for r in rs]
+    for r, x, y in zip(rs, xs, ys, strict=True):
+        label = r.quant.label
+        if labels.count(r.quant.label) > 1 and r.quant.publisher:  # same label from two publishers
+            label += f" ({r.quant.publisher})"
+        label += " *" if r.approximate else ""
+        ax.annotate(label, (x, y), textcoords="offset points", xytext=(7, 6), fontsize=8, color=INK)
+    if budget_bytes:
+        ax.axvline(budget_bytes / 1e9, color=MUTED, linewidth=1, linestyle=(0, (4, 3)))
+        ax.annotate(
+            "memory budget", (budget_bytes / 1e9, 1), xycoords=("data", "axes fraction"),
+            xytext=(-4, -12), textcoords="offset points", ha="right", fontsize=8, color=MUTED,
+        )  # fmt: skip
+    ax.set_xlabel("File size (GB)", color=MUTED, fontsize=9)
+    ax.set_ylabel("KLD mean vs original (lower is better)", color=MUTED, fontsize=9)
+    if min(ys) > 0 and max(ys) / max(min(ys), 1e-12) > 30:
+        ax.set_yscale("log")
+    ax.grid(axis="y", color=GRID, linewidth=0.8)
+    ax.set_axisbelow(True)
+    for side in ("top", "right", "left"):
+        ax.spines[side].set_visible(False)
+    ax.spines["bottom"].set_color(GRID)
+    ax.tick_params(colors=MUTED, labelsize=8, length=0)
+    pack = results[0].pack_name or results[0].pack_id
+    ax.set_title(title or f"Quality vs size  |  pack {pack}", loc="left", color=INK, fontsize=11)
+    if any(r.approximate for r in rs):
+        fig.text(0.01, 0.01, "* approximate (server top-N)", fontsize=7, color=FLAG)
+    fig.tight_layout()
+    fig.savefig(out, facecolor=SURFACE)
+    plt.close(fig)
+    return out

quantgap/report/jsonout.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+from quantgap.models import CompareReport, ScoreResult
+
+
+def safe_name(r: ScoreResult) -> str:
+    base = f"{(r.pack_name or r.pack_id).split('@')[0]}-{r.quant.label}"
+    return re.sub(r"[^A-Za-z0-9._-]+", "_", base)
+
+
+def write_result(r: ScoreResult, path: str | Path) -> Path:
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    p.write_text(r.to_json() + "\n", encoding="utf-8")
+    return p
+
+
+def write_report(rep: CompareReport, path: str | Path) -> Path:
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    p.write_text(rep.model_dump_json(indent=2, by_alias=True) + "\n", encoding="utf-8")
+    return p

quantgap/report/table.py

@@ -0,0 +1,82 @@
+"""Rich tables for the terminal."""
+
+from __future__ import annotations
+
+from rich.console import Console
+from rich.table import Table
+
+from quantgap.models import ScoreResult
+from quantgap.util import human_size
+
+
+def print_result(c: Console, r: ScoreResult) -> None:
+    m = r.metrics
+    c.print(f"[bold]quantgap {r.quantgap}[/bold]  pack: {r.pack_name}")
+    flags = " ".join(
+        f"[yellow]{t}[/yellow]"
+        for t, on in (
+            ("quick", r.quick),
+            ("approximate", r.approximate),
+            ("non-deterministic", r.nondeterministic),
+        )
+        if on
+    )
+    c.print(
+        f"quant: {r.quant.source}  [bold]{r.quant.label}[/bold]  {human_size(r.quant.size_bytes)}  "
+        f"{r.backend.name}  {flags}"
+    )
+    t = Table(show_header=True, header_style="bold", box=None, pad_edge=False)
+    t.add_column("metric")
+    t.add_column("value", justify="right")
+    t.add_column("note")
+    t.add_row(
+        "KLD mean",
+        f"{m.kld_mean:.4f}",
+        f"95% CI [{m.kld_mean_ci95[0]:.4f}, {m.kld_mean_ci95[1]:.4f}]  (lower is better)",
+    )
+    t.add_row("KLD worst 1%", f"{m.kld_p99:.3f}", "99th percentile of positions")
+    t.add_row("KLD worst 0.1%", f"{m.kld_p999:.3f}", "99.9th percentile")
+    t.add_row("top-1 agree", f"{m.top1_agree * 100:.1f} %", "same best guess as the original")
+    t.add_row("ppl ratio", f"{m.ppl_ratio:.4f}", "1.0 means no change")
+    d = m.damage
+    if d:
+        t.add_row("loops", f"{d.loops} / {d.n_prompts}", f"original: {d.loops_ref}")
+        t.add_row(
+            "wrong script", f"{d.wrong_script} / {d.n_prompts}", "script not in the original's text"
+        )
+        t.add_row("valid JSON", f"{d.json_valid} / {d.json_n}", f"original: {d.json_valid_ref}")
+        t.add_row(
+            "first divergence",
+            f"{d.first_divergence_median:.0f} tokens",
+            "median, vs the original's greedy text",
+        )
+    c.print(t)
+    c.print("by group:  " + "  ".join(f"{g} {gm.kld_mean:.4f}" for g, gm in m.by_group.items()))
+    for w in r.warnings:
+        c.print(f"[yellow]note:[/yellow] {w}")
+
+
+def print_compare(c: Console, results: list[ScoreResult]) -> None:
+    t = Table(show_header=True, header_style="bold")
+    for col, j in (("quant", "left"), ("size", "right"), ("KLD mean", "right"), ("95% CI", "right"),
+                   ("worst 1%", "right"), ("top-1", "right"), ("ppl", "right"), ("loops", "right"),
+                   ("wrong script", "right"), ("JSON ok", "right")):  # fmt: skip
+        t.add_column(col, justify=j)
+    for r in sorted(results, key=lambda r: r.metrics.kld_mean):
+        m, d = r.metrics, r.metrics.damage
+        t.add_row(
+            r.quant.label + (" *" if r.approximate else "") + (" (quick)" if r.quick else ""),
+            human_size(r.quant.size_bytes),
+            f"{m.kld_mean:.4f}",
+            f"{m.kld_mean_ci95[0]:.4f}-{m.kld_mean_ci95[1]:.4f}",
+            f"{m.kld_p99:.3f}",
+            f"{m.top1_agree * 100:.1f}%",
+            f"{m.ppl_ratio:.3f}",
+            f"{d.loops}/{d.loops_ref}" if d else "-",
+            f"{d.wrong_script}" if d else "-",
+            f"{d.json_valid}/{d.json_n}" if d else "-",
+        )
+    c.print(t)
+    if any(r.approximate for r in results):
+        c.print("* approximate (server top-N)")
+    c.print("loops shows quant/original. Lower KLD is better.")

quantgap/score/__init__.py

@@ -0,0 +1,4 @@
+from quantgap.score.kld import kld_full, kld_topk, ppl_ratio
+from quantgap.score.stats import bootstrap_mean_ci, percentile
+
+__all__ = ["bootstrap_mean_ci", "kld_full", "kld_topk", "percentile", "ppl_ratio"]

quantgap/score/compare.py

@@ -0,0 +1,33 @@
+"""Build a comparison report from several results (one pack only, spec E7)."""
+
+from __future__ import annotations
+
+from quantgap.models import CompareReport, ScoreResult
+from quantgap.score.recommend import ASSUMPTION, recommend
+
+
+class MixedPacksError(ValueError):
+    pass
+
+
+def build_report(results: list[ScoreResult], budget_bytes: int | None = None) -> CompareReport:
+    if not results:
+        raise ValueError("Nothing to compare.")
+    ids = {r.pack_id for r in results}
+    if len(ids) > 1:
+        raise MixedPacksError(
+            "These results come from different packs, so their numbers cannot share one "
+            f"table: {', '.join(sorted(ids))}. Score all quants against the same pack."
+        )
+    rec_text = None
+    if budget_bytes:
+        rec = recommend(results, budget_bytes)
+        rec_text = " ".join([rec.sentence, *rec.warnings])
+    return CompareReport(
+        pack_id=results[0].pack_id,
+        pack_name=results[0].pack_name,
+        results=results,
+        budget_bytes=budget_bytes,
+        recommendation=rec_text,
+        assumptions=ASSUMPTION if budget_bytes else None,
+    )

quantgap/score/damage.py

@@ -0,0 +1,103 @@
+"""Visible-damage checks on greedy generations (section 7.4 of the spec).
+
+Plain words: KLD can look small while a model still does something obviously wrong, like
+repeating itself forever, dropping a Chinese character into English text, or breaking JSON.
+We run short greedy generations and count those failures, always next to the reference's own.
+"""
+
+from __future__ import annotations
+
+import json
+from collections import Counter
+from statistics import median
+
+LOOP_NGRAM = 8
+LOOP_REPEATS = 4
+TAIL_WINDOW = 64
+TAIL_MIN_DISTINCT = 12
+
+# Unicode ranges for the scripts that most often leak into the wrong place.
+_SCRIPTS: dict[str, list[tuple[int, int]]] = {
+    "cjk": [(0x4E00, 0x9FFF), (0x3400, 0x4DBF), (0x3040, 0x30FF), (0xAC00, 0xD7AF)],
+    "cyrillic": [(0x0400, 0x052F)],
+    "arabic": [(0x0600, 0x06FF), (0x0750, 0x077F), (0xFB50, 0xFDFF), (0xFE70, 0xFEFF)],
+    "devanagari": [(0x0900, 0x097F)],
+}
+# Groups where only the reference's own scripts are expected.
+SCRIPT_CHECK_GROUPS = ("general", "code", "math", "json")
+
+
+def is_loop(ids: list[int]) -> bool:
+    """True if any 8-token run repeats 4+ times, or the last 64 tokens use fewer than 12 tokens."""
+    n = len(ids)
+    if n >= LOOP_NGRAM * LOOP_REPEATS:
+        counts = Counter(tuple(ids[i : i + LOOP_NGRAM]) for i in range(n - LOOP_NGRAM + 1))
+        if counts.most_common(1)[0][1] >= LOOP_REPEATS:
+            return True
+    return n >= TAIL_WINDOW and len(set(ids[-TAIL_WINDOW:])) < TAIL_MIN_DISTINCT
+
+
+def scripts_in(text: str) -> set[str]:
+    found: set[str] = set()
+    for ch in text:
+        c = ord(ch)
+        if c < 0x0400:
+            continue
+        for name, ranges in _SCRIPTS.items():
+            if name not in found and any(lo <= c <= hi for lo, hi in ranges):
+                found.add(name)
+    return found
+
+
+def has_wrong_script(quant_text: str, ref_text: str) -> bool:
+    """True if the quant text uses a script the reference text does not."""
+    return bool(scripts_in(quant_text) - scripts_in(ref_text))
+
+
+def first_json_block(text: str) -> str | None:
+    """The first balanced {...} block in the text, or None."""
+    start = text.find("{")
+    if start < 0:
+        return None
+    depth, in_str, esc = 0, False, False
+    for i in range(start, len(text)):
+        ch = text[i]
+        if in_str:
+            if esc:
+                esc = False
+            elif ch == "\\":
+                esc = True
+            elif ch == '"':
+                in_str = False
+            continue
+        if ch == '"':
+            in_str = True
+        elif ch == "{":
+            depth += 1
+        elif ch == "}":
+            depth -= 1
+            if depth == 0:
+                return text[start : i + 1]
+    return None
+
+
+def json_valid(text: str) -> bool:
+    block = first_json_block(text)
+    if block is None:
+        return False
+    try:
+        return isinstance(json.loads(block), dict)
+    except ValueError:
+        return False
+
+
+def first_divergence(quant_ids: list[int], ref_ids: list[int]) -> int:
+    """Index of the first token where the two outputs differ (the length if one is a prefix)."""
+    for i, (a, b) in enumerate(zip(quant_ids, ref_ids, strict=False)):
+        if a != b:
+            return i
+    return min(len(quant_ids), len(ref_ids))
+
+
+def median_or_zero(values: list[int]) -> float:
+    return float(median(values)) if values else 0.0

quantgap/score/kld.py

@@ -0,0 +1,73 @@
+"""Top-k KLD with a tail bucket, top-1 agreement, perplexity ratio.
+
+Plain words: at each position the reference model gives probabilities to its k best next tokens
+and the rest of the probability ("the tail") is lumped into one bucket. We ask the quant for
+its probabilities on the same k tokens, lump its leftover into a tail bucket too, and measure
+how far apart the two sets of numbers are. Lumping can only make two distributions look
+closer, so this number is a lower bound of the true KLD (see docs/method.md).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+EPS = 1e-12  # floor for the quant's probabilities before taking a log
+TAIL_SKIP = 1e-9  # ignore the reference tail term when its mass is this small
+
+
+@dataclass
+class KLResult:
+    kl: np.ndarray  # (n,) per-position KL_k
+    clamps: int  # how many probabilities had to be raised to EPS
+
+
+def kld_topk(p_logp: np.ndarray, q_logp: np.ndarray) -> KLResult:
+    """Per-position top-k + tail KLD.
+
+    p_logp: (n, k) reference log probabilities of its own top-k ids.
+    q_logp: (n, k) quant log probabilities of the *same* ids. NaN means "the backend could not
+        tell" (a server that returns only its own top-N). Such ids are moved into the tail bucket
+        for both distributions, which keeps the result a valid lower bound.
+    """
+    p_logp = np.asarray(p_logp, dtype=np.float64)
+    q_logp = np.asarray(q_logp, dtype=np.float64)
+    valid = ~np.isnan(q_logp)
+    p_lp = np.where(valid, p_logp, 0.0)
+    p = np.where(valid, np.exp(p_logp), 0.0)
+    q = np.where(valid, np.exp(np.where(valid, q_logp, 0.0)), 0.0)
+
+    clamped = valid & (q < EPS)
+    q_safe = np.maximum(q, EPS)
+    head = np.where(valid, p * (p_lp - np.log(q_safe)), 0.0).sum(axis=1)
+
+    p_tail = np.clip(1.0 - p.sum(axis=1), 0.0, 1.0)
+    q_tail_raw = 1.0 - q.sum(axis=1)
+    q_tail = np.maximum(q_tail_raw, EPS)
+    use_tail = p_tail >= TAIL_SKIP
+    tail = np.where(use_tail, p_tail * (np.log(np.maximum(p_tail, EPS)) - np.log(q_tail)), 0.0)
+    # A clearly negative tail means the quant's probabilities add up to more than 1: broken.
+    broken_tail = use_tail & (q_tail_raw < -1e-6)
+    clamps = int(clamped.sum() + broken_tail.sum())
+    return KLResult(head + tail, clamps)
+
+
+def log_softmax(x: np.ndarray) -> np.ndarray:
+    """Stable log-softmax in float64 (subtract the max first)."""
+    x = np.asarray(x, dtype=np.float64)
+    m = x.max(axis=-1, keepdims=True)
+    z = x - m
+    return z - np.log(np.exp(z).sum(axis=-1, keepdims=True))
+
+
+def kld_full(p_logits: np.ndarray, q_logits: np.ndarray) -> np.ndarray:
+    """Exact KLD over the whole vocabulary (used by tests and `doctor --selftest`)."""
+    p_lp = log_softmax(p_logits)
+    q_lp = log_softmax(q_logits)
+    return (np.exp(p_lp) * (p_lp - q_lp)).sum(axis=-1)
+
+
+def ppl_ratio(logp_target_quant: np.ndarray, logp_target_ref: np.ndarray) -> float:
+    """exp(mean(-log q)) / exp(mean(-log p)) over scored positions. 1.0 means no change."""
+    return float(np.exp(np.mean(logp_target_ref - logp_target_quant)))