Skip to content

forward-serve: --dump-logits writes the raw distribution, so a KLD comparison is possible at all - #940

Merged
bong-water-water-bong merged 1 commit into
mainfrom
1bit/671-dump-logits
Nov 4, 2026
Merged

bong-water-water-bong merged 1 commit into
mainfrom
1bit/671-dump-logits

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

What

forward-serve gains an off-by-default --dump-logits <dir> flag that writes each generated position's raw
float32 distribution over the whole vocabulary as step_<n>.bin.

Root cause / why it is needed

Measuring a KL divergence against another stack (engine#671 gate (b)) needs the distribution, and neither existing route
can produce it:

  • the OpenAI API caps top_logprobs — {"error":{"message":"top_logprobs: must be 0 to 100"}};
  • 1bit npu-run --dump-logits does dump logits, but it drives the old lane family (layer_ctx1.elf, layer_pdi),
    not the full_i8_* per-shape family that every current container uses, so on an imported OGA pack it fails at load
    with cannot read …/layer_ctx1.elf.

The dump is taken immediately before sample_token(), which takes logits by const reference, so what lands on disk is
the model's own distribution — before temperature, top-k/top-p and the GBNF mask narrow it.

Proof (measured, not asserted)

app/forward_serve.cpp, one PR, no behaviour change when the flag is absent (nothing is written unless asked; a bad
directory is refused at startup rather than silently ignored).

Against AMD's recorded distribution for the same pack and the same prompts (~/ogaload/window/oracle_logits.npz,
captured through OGA on AMD's own NPU stack), our imported container in host-GEMM mode:

prompt KL(AMD‖ours) argmax our rank for AMD's token AMD p / our p
"What is the capital of France?" 0.0003 nats 6722 / 6722 1 0.999 / 0.999
"Name the first three prime numbers." 0.0006 nats 1156 / 1156 1 0.994 / 0.991
"Explain in one sentence why the sky is blue." 0.0025 nats 12884 / 12884 1 0.951 / 0.963

Alignment was proven rather than assumed: AMD's capture is the distribution after its first generated token, so the
comparison uses our step_1 and requires our step_0 argmax to equal AMD's first output token (785) — it does, on all
three prompts. (My own first run compared step_0 and reported 11-21 nats, which is what a one-position offset looks
like, and is the reason the alignment check is now printed on every row.)

Checks: builds with ONEBIT_NPU=ON, ninja onebit clean, existing ctest NPU set unaffected; tools/copyright.py --check on the touched file. Closes nothing on its own — it is the vehicle for #671 gate (b).

Notes

  • The flag is a measurement tool, so it is deliberately not wired into any release row.
  • It writes one file per generated position, so it is suitable for short reference runs, not for serving traffic;
    the docstring in the struct says so.

…mparison is possible at all

The OpenAI API caps top_logprobs at 100, and `1bit npu-run --dump-logits` drives the old lane
family (layer_ctx1.elf / layer_pdi) rather than the per-shape full_i8_* family every current
container uses, so neither route can produce a vocabulary-wide distribution for an imported pack.

The dump is taken immediately before sample_token(), which takes logits by const reference, so
it is the model's own distribution -- before temperature, top-k/top-p and the GBNF mask narrow
it.  Off by default; a bad directory is refused at startup rather than ignored.

Measured against AMD's recorded distribution for the same pack and prompts (engine#671 gate b):
KL(AMD||ours) = 0.0003 / 0.0006 / 0.0025 nats, argmax agreeing at every recorded position.
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

Addendum: the same measurement re-run from this branch's own build (~/wt/i671kld/build-npu/1bit, i.e. origin/main + this patch — not the tree the flag was first prototyped in) reproduces the table exactly: KL = 0.0003 / 0.0006 / 0.0025 nats, argmax 6722/1156/12884 agreeing, our rank for AMD's token = 1 on all three, alignment check (step_0 argmax == AMD's first output token 785) passing on all three. Copyright check rc 0.

@github-actions

github-actions Bot commented Nov 4, 2026 •

Copy link
Copy Markdown

PR Reviewer Guide 🔍

(Review updated until commit 41eb612)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶

671 - Partially compliant

Compliant requirements:

  • Support for loading Ryzen AI OGA NPU model packs on Linux through the engine
  • Implementation of parsing quantized weights into Q4NX container
  • Self-check dequantization capability
  • Per-shape full ELF generation capability
  • Gate on first pack with greedy output matching vendor stack
  • KL divergence comparison capability via --dump-logits flag

Non-compliant requirements:

  • Second pack of different family from public collection
  • ORT execution provider integration (out of scope)

Requires further human verification:

  • Verification of KL divergence measurements against vendor stack
  • Testing with second pack from public collection
  • Performance benchmarks against vendor stack
⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪
🧪 No relevant tests
🔒 No security concerns identified
⚡ Recommended focus areas for review

Potential File System Access Issue

The code creates directories for dumping logits without proper error handling for potential race conditions or permission issues. If two processes try to create the same directory simultaneously, or if the directory creation fails due to permissions, the application might not behave as expected. The error code check is present but may not cover all failure scenarios.

if (!dump_logits.empty()) {
    std::error_code ec;
    std::filesystem::create_directories(dump_logits, ec);
    if (ec || !std::filesystem::is_directory(dump_logits))
        throw std::runtime_error("--dump-logits: cannot use " + dump_logits + " as a directory");
    e.dump_logits = dump_logits;
    std::fprintf(stderr, "1bit forward-serve: dumping raw logits to %s\n", dump_logits.c_str());
}
No Validation of Directory Write Permissions

The code checks if the directory exists and is a directory, but does not verify that the application has write permissions to that directory. This could lead to runtime errors when trying to write files if the directory is read-only or if the application lacks necessary permissions.

if (!dump_logits.empty()) {
    std::error_code ec;
    std::filesystem::create_directories(dump_logits, ec);
    if (ec || !std::filesystem::is_directory(dump_logits))
        throw std::runtime_error("--dump-logits: cannot use " + dump_logits + " as a directory");
    e.dump_logits = dump_logits;
    std::fprintf(stderr, "1bit forward-serve: dumping raw logits to %s\n", dump_logits.c_str());
}

@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 41eb61291. The dump is taken from the right place, it compiles with the real NPU flags, and the
oracle half of the proof table checks out from the artefact. Two notes, neither a gate.

It compiles where CI does not. app/forward_serve.cpp at this head, compiled with the project's own
ONEBIT_NPU flags (taken from an existing -DONEBIT_NPU=ON compile_commands.json, source tree pointed at this
worktree): rc 0, including the new --dump-logits block. tools/copyright.py --check rc 0.

The distribution is the model's, not the sampler's. sample_token(const std::vector<float> & logits, …) takes
it by const reference (app/npu_sampling.h:88) and the dump copy runs immediately before the call, with the grammar
mask passed separately (gram ? &allowed : nullptr), so the file is pre-temperature, pre-top-k/p and pre-GBNF. The
top_logprobs claim is the code's own: npu_sampling.h:295 is literally
"top_logprobs: must be 0 to " + kMaxTopLogprobs, and the cap is why the API route cannot serve gate (b).

The AMD side of the table verifies from ~/ogaload/window/oracle_logits.npz (which is on the box):

prompt 0  argmax 6722   softmax p(argmax) = 0.9987   top3 (6722 .9987, 3146 .0003, 4226 .0001)
prompt 1  argmax 1156   softmax p          = 0.9936   top3 (1156 .9936, 2326 .0052, 10250 .0007)
prompt 2  argmax 12884  softmax p          = 0.9512   top3 (12884 .9512, 6303 .0224, 9237 .0093)

6722 / 1156 / 12884 and 0.999 / 0.994 / 0.951 are the table's AMD column exactly, and out_i == logits_i
byte for byte in the npz, so the two names in that file are one distribution. What I could not reproduce is the
KL column (0.0003 / 0.0006 / 0.0025) and our probabilities: see the first note.

Note 1 (the one that costs this tool a device token): forward-serve takes the NPU device lock even in
host-GEMM mode.
take_npu_device_lock(streams) is at app/forward_serve.cpp:752, and the NPU_INFER_HOST_GEMM
check that decides whether a device is opened at all is at :758. So the very run this flag exists for — the
host-GEMM reference dump — cannot start while another process holds the lock. It could not here: the owner's Bonsai
serve holds /tmp/1bit-npu-device.lock, so I did not run it (lane rule), which is why the KL column above is yours
and not mine. Moving the host_gemm decision above the lock, or giving it an explicit opt-out, would make the
reference dump runnable beside a serve — and it is exactly the kind of run that should not need the device.

Note 2 (small): the std::fopen in the loop is silent on failure — a directory that goes away mid-run produces
a clean answer and no logits. The startup check covers the ordinary case, so a one-line warning there would only
matter for a long run.

Everything else is right: off by default (dump_logits empty ⇒ nothing written), a bad directory refused at
startup rather than ignored, one step_<n>.bin per generated position, and the file is vocabulary × float32.
mergeStateStatus BLOCKED is CI only (build and pr_agent were still running); the code review above is at this
head.

@bong-water-water-bong bong-water-water-bong added the tprs-ready TPRS: review verdict READY posted by the standing reviewer; merge gate satisfied label Nov 4, 2026
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

CI is green at 41eb61291 (build + Analyze ×5 + CodeQL SUCCESS, hw_gate SKIPPED; the UNSTABLE state is pr_agent cancelled by its own superseding run), so the READY above stands at this head and tprs-ready is on. Base is main and #938 is already in (54df1d42c), so no stacking remains.

@github-actions

github-actions Bot commented Nov 4, 2026

Copy link
Copy Markdown

Persistent review updated to latest commit 41eb612

@bong-water-water-bong
bong-water-water-bong merged commit b6dbd71 into main Nov 4, 2026
11 of 12 checks passed
@bong-water-water-bong
bong-water-water-bong deleted the 1bit/671-dump-logits branch November 4, 2026 01:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Review effort 3/5 tprs-ready TPRS: review verdict READY posted by the standing reviewer; merge gate satisfied

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant