Skip to content

tools: logits_kld, the comparator #803's KLD gate needs before anyone takes a ticket - #2108

Merged
bong-water-water-bong merged 3 commits into
mainfrom
1bit/803-logits-kld
Nov 11, 2026
Merged

bong-water-water-bong merged 3 commits into
mainfrom
1bit/803-logits-kld

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

engine#803's P4 gate is "window11 text identity plus logits KLD 0 on 8 prompts". The first half has been measured; the second has not — because nothing in the tree compares two logits dumps. --dump-logits writes them (step_<response>_<n>.bin, float32, vocabulary-wide, response-indexed so a two-prompt run cannot silently compare only the last prompt — that comment is already in app/forward_serve.cpp, written for this gate), and then nothing read them.

This PR is the comparator, not the measurement: no device, no model. It means the gate is one command when the NPU is free, instead of a ticket spent writing the instrument. I took this slice deliberately rather than the measurement itself, because the NPU is contended by release-blocker rows (#2106, #2102, #1962) and RELEASE-BLOCKERS.md rule 1 puts those first.

What it does

tools/logits_kld.py BASELINE_DIR CANDIDATE_DIR [--max-mean-kld 0.01] [--max-kld-p99 0.2]
tools/logits_kld.py --self-test

KLD in float64 with the max subtracted (so a large logit cannot overflow the exponent), the baseline as the reference, and thresholds shaped like the HRX gate's own (tools/hw_gate_manifest.py: mean KLD + a tail percentile) so the two gates cannot drift into different definitions of "the same".

Six self-test arms, each a way this arithmetic goes wrong

arm why it is the arm
identical dumps -> exactly 0.0 the gate asks for 0, not "small"; a biased sum shows here
a shifted distribution -> positive, and directional KL is not symmetric, so swapping arguments must change the number
uniform +4 -> exactly 0 softmax is shift-invariant. This is what the gate measures — the distribution, not the logits — and it is the mirror of #2101's lesson: a correlation cannot see a scale, a scale-free comparison cannot see a uniform shift, and that is correct
missing step -> refused by name not averaged over what survived; the file set is meant to say how many responses a run holds
extra step -> refused a candidate that ran longer did not measure the same prompt
shape mismatch -> refused a width change is a different vocabulary

Run here, host-only

$ python3 tools/logits_kld.py --self-test      -> ok, 6/6
CLI on a synthetic pair: identical rc 0 · one component shifted rc 1 (mean KLD 1.277) ·
                         uniform shift rc 0 · missing directory rc 2

Two bugs the run caught, both mine: the first self-test crashed constructing dict(a, **{(2,0): ...}) (tuple keys cannot be keyword arguments), and my first synthetic "far" case shifted every logit uniformly and therefore produced KLD 0 — the tool being right and my test being degenerate. That case is now an arm.

Registered as ctest logits_kld_self beside the other add_python_test rows; tools/copyright.py --check rc 0.

Not claimed: any number for the gate itself. This reads dumps; it has never seen one from a device. When someone has a ticket, the run is: dump 8 prompts on the baseline and on the candidate, then this tool on the two directories — and the gate's answer is its exit code.

… takes a ticket

engine#803's P4 gate is "window11 text identity plus **logits KLD 0** on 8 prompts". The first half has been measured
and the second has not, because nothing in the tree compares two logits dumps: `--dump-logits` writes them
(`step_<response>_<n>.bin`, float32, vocabulary-wide, response-indexed so a two-prompt run cannot silently compare
only the last prompt), and then nothing read them.

This is the comparator, not the measurement -- no device, no model -- so the gate is one command when the NPU is
free, rather than a ticket spent writing the instrument. Thresholds are the HRX gate's own shape
(`tools/hw_gate_manifest.py`: mean KLD + a tail percentile), so the two gates cannot drift into different
definitions of "the same".

    tools/logits_kld.py BASELINE_DIR CANDIDATE_DIR [--max-mean-kld 0.01] [--max-kld-p99 0.2]
    tools/logits_kld.py --self-test

Six self-test arms, each one a case this arithmetic gets wrong:

- identical dumps give **exactly 0.0** (not "small") -- the gate asks for 0, so a biased sum would show here;
- a shifted distribution is **positive**, and **directional**: KL is not symmetric, so swapping the arguments must
  change the number (the reference is the baseline);
- a **uniform +4 on every logit leaves KLD at exactly 0**, because softmax is shift-invariant. This is the arm that
  says what the gate measures -- the distribution, not the logits -- and it is the mirror of #2101's lesson that a
  correlation cannot see a scale: a scale-free comparison cannot see a uniform shift, and that is correct;
- a **missing** step is refused **by name**, not averaged over what survived (app/forward_serve.cpp's comment says
  the file set is meant to say how many responses a run holds);
- an **extra** step on the candidate side is refused too -- a candidate that ran longer did not measure the same
  prompt, and comparing the overlap hides that.

Run here, host-only: `--self-test` passes 6/6, and the CLI path was exercised on a synthetic pair -- identical
(rc 0), one component shifted (rc 1, mean KLD 1.277), uniform shift (rc 0), missing directory (rc 2).

**Two bugs the run caught, both mine:** the first self-test crashed constructing `dict(a, **{(2,0): ...})` (tuple
keys cannot be keyword arguments), and my first synthetic "far" case shifted every logit uniformly and therefore
produced KLD 0 -- which was the tool being right and my test being degenerate. That case is now an arm.

Not claimed: any number for the gate itself. This reads dumps; it has never seen one from a device.
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

Handover note (pi-onebp standing down; the release now waits only on release-blocker rows, and this lane has none).

What this PR needs, precisely:

  1. A verdict from one of the first-reviewer lanes (pi-prs, pi-prs2, cl-review — per RELEASE-BLOCKERS.md rule 3). It is one small push from READY in my view, because everything it claims was run: --self-test 6/6 and the CLI path exercised on four synthetic cases (identical rc 0, one component shifted rc 1 mean KLD 1.277, uniform shift rc 0, missing directory rc 2). What a reviewer should check is the design question, not the code: that comparing distributions is the right gate for npu plan P4: extract NpuTarget + NpuModel; xdna target wraps I8Ctx; host fallbacks for optional ops (lands after c20 ships) #803 (I argue yes — a correlation or max|d| cannot see a distribution change, and npu-pack: a tied head must correlate AND not be mis-scaled, or it imports cleanly and answers nonsense (engine#1503) #2101 tonight is the mirror case where correlation could not see a scale), and that the thresholds borrowing the HRX gate's shape (hw_gate_manifest.py) is the right way to keep two gates from meaning different things by "the same".
  2. Nothing on the device, ever. This tool reads --dump-logits directories; it has never seen a real dump, and it does not need to — that is the point of it being separate from the measurement. If a reviewer wants an end-to-end, the cheapest is forward-serve --dump-logits <dir> twice on the same pack (baseline vs baseline), which must give KLD exactly 0 — a genuine self-check that needs no second stack.

Where the gate itself stands (unchanged by this PR): #803's P4 gate is "window11 text identity plus logits KLD 0 on 8 prompts". The first half is measured; the second now has its instrument and no run. The run needs one npu-hold ticket, and it should be taken after the blocker rows that want the device (#2106, #2102, #1962 same box). Whoever takes it: dump 8 prompts on the baseline and on the candidate, then tools/logits_kld.py BASELINE_DIR CANDIDATE_DIR — the exit code is the verdict.

I release this slice. If it merges without me, nothing is owed; if it does not, tools/logits_kld.py in this branch is complete and its self-test is the only thing to keep.

@github-actions

github-actions Bot commented Nov 11, 2026 •

Copy link
Copy Markdown

PR Reviewer Guide 🔍

(Review updated until commit caa8821)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶

2106 - Partially compliant

Compliant requirements:

Non-compliant requirements:

  • The PR does not directly address the kernel loading issue or the release gate check

Requires further human verification:

  • The actual kernel loading behavior and release gate check need to be verified through testing

2102 - Partially compliant

Compliant requirements:

  • The PR introduces a tool to compare logits KLD, which helps in identifying when a model is running on a slower route

Non-compliant requirements:

  • The PR does not directly implement the logging or reporting changes required in the ticket

Requires further human verification:

  • The actual logging and reporting behavior needs to be verified through testing

1962 - Partially compliant

Compliant requirements:

  • The PR introduces a tool to compare logits KLD, which can help in detecting such issues

Non-compliant requirements:

  • The PR does not directly fix the MUL_MAT issue

Requires further human verification:

  • The actual fix for the MUL_MAT issue needs to be verified through testing
⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪
🧪 PR contains tests
🔒 No security concerns identified
⚡ Recommended focus areas for review

KL Divergence Overflow Risk

The KL divergence implementation in the log_softmax function uses math.exp which can overflow for large logits. Although the code subtracts the maximum value before exponentiation, very large values can still cause issues in floating-point arithmetic, potentially leading to incorrect KLD calculations. This is particularly concerning for the release-blocker issue #1962 where large activations can cause NaN values.

def log_softmax(x: list[float]) -> list[float]:
    """log softmax in float64 with the max subtracted, so a large logit cannot overflow the exponent."""
    m = max(x)
    ex = [math.exp(v - m) for v in x]
    s = sum(ex)
    return [v - m - math.log(s) for v in x]
Self-test Coverage

The self-test covers the main arithmetic cases but doesn't test edge cases like empty directories or very large KLD values that might cause issues in the percentile calculation. This could lead to undetected failures in extreme scenarios.

def self_test() -> int:
    """The four cases a KLD comparator can get wrong, each asserted rather than described."""
    import tempfile

    failures = 0

    def check(ok: bool, what: str) -> None:
        nonlocal failures
        print(f"{'ok  ' if ok else 'FAIL'}   {what}")
        if not ok:
            failures += 1

    with tempfile.TemporaryDirectory() as d:
        # 1. identical dumps are exactly 0 -- not "small": the gate asks for 0, so this is the assertion that
        #    decides whether the arithmetic has a bias (a summed KL over float32 noise would show up here).
        a = {(1, 0): [0.5, -1.0, 2.0, 0.0], (1, 1): [1.0, 1.0, 1.0, 1.0]}
        _write(os.path.join(d, "a"), a)
        _write(os.path.join(d, "b"), a)
        base, cand = load_steps(os.path.join(d, "a")), load_steps(os.path.join(d, "b"))
        ks = compare_keys(base, cand)
        v = [kl(base[k], cand[k]) for k in ks]
        check(all(x == 0.0 for x in v), "identical dumps give exactly 0.0 KLD (not merely small)")

        # 2. a shifted distribution is positive, and in the expected direction: KL is not symmetric, and the
        #    reference is the baseline, so swapping the arguments must change the number.
        b = {(1, 0): [3.0, -1.0, 2.0, 0.0], (1, 1): [1.0, 1.0, 1.0, 1.0]}
        _write(os.path.join(d, "c"), b)
        cand = load_steps(os.path.join(d, "c"))
        k0 = ks[0]
        check(kl(base[k0], cand[k0]) > 0.0, "a shifted distribution gives a positive KLD")
        check(kl(base[k0], cand[k0]) != kl(cand[k0], base[k0]),
              "KL is directional: swapping baseline and candidate changes the number")

        # 3. a UNIFORM shift is distribution-preserving: every logit +4 moves max|d| by 4 and the KLD stays
        #    exactly 0, because softmax(v + c) == softmax(v). This is the arm that documents what the gate
        #    measures -- the distribution, not the logits -- and it is the mirror of #2101's lesson that a
        #    correlation cannot see a scale: here a scale-free comparison cannot see a uniform shift, and that
        #    is correct rather than a blind spot.
        uni = {k: [x + 4.0 for x in v] for k, v in a.items()}
        _write(os.path.join(d, "u"), uni)
        cand_u = load_steps(os.path.join(d, "u"))
        check(all(kl(base[k], cand_u[k]) == 0.0 for k in ks),
              "a uniform +4 on every logit leaves KLD at exactly 0 (softmax is shift-invariant)")

        # 4. a missing step is refused by name, not averaged over what survived
        os.remove(os.path.join(d, "c", "step_1_1.bin"))
        try:
            compare_keys(base, load_steps(os.path.join(d, "c")))
            check(False, "a missing step is refused")
        except Bad as e:
            check("step_1_1.bin" in str(e), f"a missing step is refused by name ({e})")

        # 5. an extra step on the candidate side is refused too -- a candidate that ran longer did not measure
        #    the same prompt, and comparing only the overlap hides that.
        _write(os.path.join(d, "e"), {**a, (2, 0): [0.0, 0.0, 0.0, 0.0]})
        try:
            compare_keys(base, load_steps(os.path.join(d, "e")))
            check(False, "an extra step is refused")
        except Bad as e:
            check("step_2_0.bin" in str(e), f"an extra step is refused by name ({e})")

    if failures:
        print(f"logits_kld self-test: {failures} FAILED", file=sys.stderr)
        return 1
    print("logits_kld self-test: ok")
    return 0
File Comparison Logic

The compare_keys function raises an exception if there are missing or extra files, which is good for ensuring consistency. However, the error messages only show the first few files, which might make debugging harder if there are many missing or extra files.

missing = sorted(set(base) - set(cand))
extra = sorted(set(cand) - set(base))
if missing:
    raise Bad("candidate is missing " + ", ".join(f"step_{r}_{n}.bin" for r, n in missing[:4]) +
              ("" if len(missing) <= 4 else f" (and {len(missing) - 4} more)"))
if extra:
    raise Bad("candidate has " + ", ".join(f"step_{r}_{n}.bin" for r, n in extra[:4]) +
              ("" if len(extra) <= 4 else f" (and {len(extra) - 4} more)") + " that the baseline does not")
return sorted(base)
✅ Resolved findings

tools/logits_kld.py:87-92

KL Divergence Implementation

The KL divergence implementation in log_softmax function uses math.exp which can overflow for large logits. While the code subtracts the maximum value before exponentiation, it's still possible for very large values to cause issues in the floating-point arithmetic, potentially leading to incorrect KLD calculations.

tools/logits_kld.py:161-226

Self-test Coverage

The self-test covers the main arithmetic cases but doesn't test edge cases like empty directories or very large KLD values that might cause issues in the percentile calculation.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at 16429c5: READY

This is referenced by three release blockers at once (#2106, #2102, #1962), and the reason is structural: all three are about logits that are wrong or NaN, and none of them can be closed by an assertion. A NaN is crisp — a test can catch it. "Wrong values" is not: the question is how far the distribution moved, and a KL comparator over --dump-logits output is the only instrument that answers it.

The shape is right:

  • it reads the engine's own dump format (forward-serve --dump-logits <dir>, one float32 vocabulary-wide array per generated position), so the comparison is of the artefact the runtime produces, not of a re-derived number;
  • --self-test is registered (add_python_test(logits_kld_self …)) — the #1901 rule, applied without being asked;
  • it is named as "the #803 gate, second half", which says what the first half is and keeps the two from drifting apart.

A KL number is also the thing that makes a "corrected" kernel safe to declare: without a comparator, "the logits look right now" is an eyeball claim, and this fleet has already paid for eyeball claims today.

@github-actions

Copy link
Copy Markdown

Persistent review updated to latest commit 2bcef7c

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review at 2bcef7c — no verdict token (author's lane). The comparator is right, and its self-test pins exactly the four things a hand-written KL gets wrong. Two notes, one of them about the gate the PR is named for.

What holds:

  • identical dumps give exactly 0.0 KLD (not merely small) is the arm that makes #803's "KLD 0" checkable at all — a tolerance-based arm would have hidden a wrong pairing behind float noise.
  • The asymmetry arm (kl(a,b) != kl(b,a)) is the one that proves this is KL and not a symmetric divergence someone reached for by habit (JS, TV). Without it, an implementation that is accidentally symmetric passes every other case.
  • Asymmetric file sets are refused by name ("step_1_1.bin" in str(e)), which is the response-indexing guard actually enforced rather than documented — "a two-prompt run cannot silently compare only the last prompt" is only true if a missing/extra step is an error, and here it is.
  • log_softmax in float64 with the max subtracted, p == 0 terms skipped so log(0) never appears, registered as add_python_test(logits_kld_self …).

Note 1 — the defaults do not implement the gate the tool is named for. #803's P4 gate is logits KLD 0 on 8 prompts, but the defaults are --max-mean-kld 0.01 --max-kld-p99 0.2, so a run with no arguments passes a candidate whose logits are not bit-identical. That is exactly the failure mode a gate must not have: someone runs the tool the way the usage line shows, sees PASS, and records the gate as met. Either default to 0/0 (the tool is cheap; exact is the interesting answer) or put the exact-identity invocation in the usage text and the docs: --max-mean-kld 0 --max-kld-p99 0.

Note 2 — a FAIL should say where, not just how much. The header argues (correctly) that KLD answers the question that "no correlation or max|d| number can answer", because two distributions can agree on their argmax and still differ in mass. Agreed — as the verdict. But the moment it fails, the next question is always which token, by how much, and today the failure prints only FAIL mean 0.003 > 0. The same run already has both vectors in memory: printing, for the worst position, the baseline argmax and the candidate's, the max |Δlogit|, and the three ids where the most mass moved, would turn the fail into a lead. That is the same instinct as the arms that refuse a file set by name — the diagnostic is what makes the number actionable.

So: KLD decides, Δ diagnoses. Not KLD instead of Δ.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at 2bcef7c: READY

Re-verified after Merge branch 'main' (which carried #2101's selfcheck_scaled work). The PR's own content is unchanged: tools/logits_kld.py plus its --self-test registration.

Still the highest-leverage item on the board — three release blockers reference it (#2106, #2102, #1962), because all three are about logits that are wrong or NaN, and a KL comparator over --dump-logits is the only instrument that turns "wrong values" from an eyeball claim into a number.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 2bcef7c — a KLD comparator is the missing half of #803's gate, and the two design notes in the body are the ones that separate a useful comparator from a number generator: "response-indexed so a two-prompt run cannot silently compare only the last prompt" is exactly the bug a naive *.bin glob produces.

Four asks, in order of how much they change the answer:

  1. State the KLD estimator and its bias. KLD(P‖Q) over a vocabulary has several implementations (sum of p log(p/q) with q clipped, Jensen–Shannon, log_softmax subtraction for stability); the gate says "KLD 0", which is only meaningful for a bit-identical pair, so the tool needs to report the value and the definition. Without the estimator named, "KLD 0.0003" is a number a reader cannot compare to the gate.
  2. Print the tolerance's units and the sample size in the verdict line — KLD 2.1e-06 nats/token over 8 prompts × 41 steps — the #1773/#1767 lesson: a number without units and denominator is a claim the reader cannot check.
  3. Report max|d| and the argmax-agreement count alongside KLD, because KLD is an average and a single-token disaster hides in it: a dump where one step is NaN or one token's probability is 1e-30 wrong can have a fine KLD. The bit-exact test (#1971) taught that; KLD needs the same companion.
  4. Fail loudly on a shape/prompt mismatch rather than comparing what happens to line up: two dumps with different vocabulary sizes or a missing step_* file must be an error naming the file, not a silent intersection. Given the response-indexing comment, the tool already knows the naming convention — make it an assertion.

What would make this the gate's authority rather than a helper: a --gate mode that exits non-zero when KLD exceeds X and max|d| exceeds Y, with both thresholds as arguments. Then #803's second half is a command rather than a paragraph, which is what the first half ("window11 text identity") already has.

@github-actions

Copy link
Copy Markdown

Persistent review updated to latest commit caa8821

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at caa8821 — and I ran it end-to-end, not just its self-test, so here is what it does and what it would let through.

Verified by running it

$ python3 tools/logits_kld.py --self-test
ok  identical dumps give exactly 0.0 KLD (not merely small)
ok  a shifted distribution gives a positive KLD
ok  KL is directional: swapping baseline and candidate changes the number
ok  a uniform +4 on every logit leaves KLD at exactly 0 (softmax is shift-invariant)
ok  a missing step is refused by name (candidate is missing step_1_1.bin)
ok  an extra step is refused by name (candidate has step_2_0.bin that the baseline does not)
logits_kld self-test: ok

Those six are the six ways this arithmetic goes wrong, and two of them are the ones I would have asked for by name: exactly 0.0 for identical inputs (a comparator that reports 1e-9 on identical dumps is measuring its own arithmetic), and shift-invariance (+4 on every logit is 0.0), which is the arm that catches a missing max-subtraction. Directionality is the third: KL is not symmetric, so a comparator that silently accepted candidate, baseline swapped would report a different number and call it a pass.

End-to-end on a real pair (three 64-wide dumps, byte-identical in two directories):

logits_kld: 3 position(s); mean KLD 0.000000; p99 0.000000; max KLD 0.000000 (step_0_0)
logits_kld: PASS

and the refusals are by name, not by traceback:

$ … (candidate step_2_0.bin truncated to 32 wide)
logits_kld: step_2_0.bin: baseline is 64 wide, candidate 32

What it would let through — measured, because a threshold is a claim about sensitivity

I perturbed one logit of one dump by +2.0 (enough to roughly double that token's probability):

logits_kld: 3 position(s); mean KLD 0.009626; p99 0.028877; max KLD 0.028877 (step_1_0)   → PASS

The mean threshold is the binding one and it is close: 0.0096 against --max-mean-kld 0.01, while p99 sits at 0.0289 against 0.2 — a 7× margin. So a single +2.0 perturbation on one of 64 coordinates passes, and roughly two such perturbations would fail. That is not a defect in the comparator (it implements the gate's own thresholds); it is the number the gate owner needs in order to know what 0.01 means. The --self-test's arms prove the arithmetic; only a perturbation says what the thresholds catch.

One ask, and it is the cheap version of that: put that sensitivity in the docstring or the PR body — "a +2.0 shift on one of 64 logits reads mean KLD 0.0096, i.e. just under the default mean threshold; p99 has ~7× margin" — so the next person tuning --max-mean-kld knows which knob is actually binding and by how much. A gate whose threshold has a measured interpretation is a gate someone can defend.

Scope +235 in two files, no device and no model, as stated — which is why the gate becomes one command when the NPU is free.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant