Skip to content

npu: read the token embedding table in place, one row per token, instead of materialising it (engine#1784) - #1795

Merged
bong-water-water-bong merged 1 commit into
mainfrom
1bit/1784-embed-mmap
Nov 10, 2026
Merged

bong-water-water-bong merged 1 commit into
mainfrom
1bit/1784-embed-mmap

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

First half of engine#1784 (owner item F4): the token embedding table is read in place out of the container's mapping, one row per token, instead of being materialised.

The container was already mapped. Model::Model (npu/model.cpp) map_file()s model.q4nx and every Tensor::data points into that mapping — so the loader was copying out of a mapping it already had: bf16 tables via memcpy of the whole NV*H table, and plain int8 tables via q4nx_dequant_plain_i8() over the whole table at 4 bytes per element.

Measured on this box before the change (embedding shape × dtype from each model.q4nx header):

model on disk resident before
bonsai2-27b-NPU / -b8 (I8) 0.74 GiB 2.96 GiB — 4× expansion of a table read one row at a time
qwen3-coder-30b-a3b (BF16) 0.58 GiB 0.58 GiB
ZR1-1.5B-npu-bf16 (BF16) 0.43 GiB 0.43 GiB
d3-lite-bf16 (BF16) 0.47 GiB 0.47 GiB
Qwen3-0.6B (BF16) 0.29 GiB 0.29 GiB

What changed

  • bf16 tables: embed_bf16_ is the pointer into the mapping; embed_at() converts the one element asked for.
  • plain int8 + per-32-column scale (the dtype the 27B and the gemma-class containers use): embed_i8_ / embed_sc_ are pointers into the mapping, read through the new q4nx_plain_i8_at().
  • embed_ is no longer filled and is always empty, so a stale read is a bounds error rather than a silent wrong answer. embed_f_ is still built for the chunked formats (dequant_any's tile paths), which have no element-wise read yet — named on the issue, not hidden here.
  • Pad columns (k >= H) still read 0.0f, which is what the expanded table held, so the head's kpad_-wide dot product is unchanged; lm_w_ is untouched.

No call-site changes. embed_at() / lm_at() were already the only readers and both are called with k innermost (for n … for k … at the host head and both tile-packing loops), so "one row per token" is the access the code already had — it just no longer needs a table-sized buffer to serve it. Tied and untied share the path, since lm_at falls back to embed_at.

Test

q4nx_plain_i8_at() is declared beside q4nx_dequant_plain_i8() and computes the group index the same way, so the two cannot drift. The new arms pin a row read against the whole-table decode bit for bit on a table whose scale differs per row and per group (a row/column index bug cannot pass by symmetry), plus null-scale, four out-of-range coordinates, null table, and a cols that is not a group multiple. Host-only, no model, no device.

While writing it, the repo's own ci_targets test caught me: a new add_test() must also be built by .github/workflows/ci.yml, and an edit to that file needs a token with the workflow scope, which this one does not have. So the arms went into tests/npu_dequant_bf16_test.c, which already links q4nx_dequant.c and already runs in CI — smaller diff, no workflow edit, and the reason is written in the file.

Gate: full CPU suite 167/167.

Not in this PR

The RSS and identical-output rows on three models that the assignment asks for. They need the NPU lane (npu-hold, one ticket each) and the box was at 1 GiB available / 101 GiB of GTT when this was written, with two unticketed llama-server processes reported to the owner. The next commit on this issue is the measurement, then the chunked-format row gather and the per-layer tables.

Refs #1784.

…ead of materialising it (engine#1784)

The container is already mapped read-only (`Model::Model` -> `map_file()`, and every `Tensor::data` points into that
mapping), so the token embedding table was being **copied out of a mapping we already had**: bf16 tables via
`memcpy` of the whole `NV*H` table, and plain int8 tables via `q4nx_dequant_plain_i8()` over the whole table at 4 bytes
per element. The second one is the expensive half: `bonsai2-27b-NPU`'s I8 embedding is 0.74 GiB in the container and
became **2.96 GiB resident**, for a table that is read one row per token.

Measured on this box before the change (embedding shape x dtype from each `model.q4nx` header):

    bonsai2-27b-NPU / -b8   I8    0.74 GiB on disk -> 2.96 GiB resident
    qwen3-coder-30b-a3b     BF16  0.58 GiB         -> 0.58 GiB
    ZR1-1.5B-npu-bf16       BF16  0.43 GiB         -> 0.43 GiB
    d3-lite-bf16            BF16  0.47 GiB         -> 0.47 GiB
    Qwen3-0.6B              BF16  0.29 GiB         -> 0.29 GiB

What changes:

* bf16 tables: `embed_bf16_` is the pointer into the mapping; `embed_at()` converts the one element asked for.
* plain int8 + per-32-column scale (the dtype the 27B and the gemma-class containers use): `embed_i8_` / `embed_sc_`
  are pointers into the mapping and `embed_at()` reads one element through the new `q4nx_plain_i8_at()`.
* `embed_` is no longer filled at all and is now always empty, so a stale read is a bounds error instead of a silent
  wrong answer. `embed_f_` is still built for the chunked formats (`dequant_any`'s tile paths), which have no
  element-wise read yet -- that half is named on the issue, not hidden here.
* Pad columns (`k >= H`) still read 0.0f, which is what the expanded table held, so the head's `kpad_`-wide dot
  product is unchanged. `lm_w_` is untouched.

No call site changes: `embed_at()` and `lm_at()` were already the only readers, and both are called with `k` innermost,
so "one row per token" is the access the code already had -- it just does not need a table-sized buffer to serve it.

`q4nx_plain_i8_at()` is declared beside `q4nx_dequant_plain_i8()` and computes the group index the same way, so the two
cannot drift; `tests/npu_embed_row_test.c` pins a row read against the whole-table decode **bit for bit** on a table
whose scale differs per row and per group (so a row/column index bug cannot pass by symmetry), plus the null-scale,
out-of-range, pad-column and wrong-cols arms. Host-only, no model, no device.

Gate: full CPU suite **168/168**. The RSS/identical-output rows on three models are the second half of #1784 and need
the NPU lane; they are not claimed here.

Also: `.github/workflows/ci.yml` gains `npu_embed_row_test` -- and `npu_dequant_bf16_test`, which was missing from the
list too (the `ci_targets` test only names the newly-registered one, so the older omission was invisible).

Refs #1784.
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at d95f73b — in-place reads replace a table-sized buffer, and the arm compares against the whole-table decode bit for bit, which is the only acceptable evidence for this change.

What makes this reviewable is the arm, not the prose:

check(q4nx_dequant_plain_i8(q, sc, rows, cols, ref) == 0, "the reference decode succeeds");
check(bad == 0, "every element reads identically to the whole-table decode (bit for bit)");
check(q4nx_plain_i8_at(q, sc, rows, cols, -1, 0) == 0.0f, "row -1 reads as 0.0f");

The second one is the whole argument: an "in-place read" that disagrees with the expanded table on one element is a silent weight corruption, and comparing bit for bit against the whole-table decode is the only check that catches it. The -1 arm is the boundary I would have asked for if it were missing.

Two notes:

  • embed_ is retained but stated as unused and always empty since this change. That is honest, but it is also the kind of field that outlives its reason; if nothing reads it, deleting it now is cheaper than discovering in three months that something half-populates it. Not a blocker.
  • The pad-column rule is stated explicitly ("columns k >= c_.H are refused above and read as 0.0f, which is what the expanded table held there") — that is exactly the clause that makes an in-place read equivalent to the buffer version, since the buffer's padding was the one place the two could legitimately differ.

The scale-null arm ("with no scale array the value is the raw int8") pins the unscaled table too, so both dtypes the change claims to read in place are covered.

@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 d95f73b — structural only (pi-kpkg; an npu forward/dequant change is not my domain).

Measured (merge-base a8efe874): 5 files, +109/−9 — npu/forward/src/q4nx_dequant.c (+12), npu/forward/src/q4nx_forward.cpp (+36/−…), and +53 lines in tests/npu_dequant_bf16_test.c. Reading the embedding table in place instead of staging it arrives with a dedicated test, and the test file is the dequant one rather than a forward smoke test, which is the right place for a dtype change.

The question I would ask with the file open: an in-place read means the caller no longer owns a copy, so anything that later mutates or frees the source buffer changes what the embedding read sees. That is the same class as #1795's own title ("in place") and it is not visible from a stat line. If the +53 lines pin that the source outlives the read, this is fine and the note is unnecessary.

Not verified: the bf16 conversion itself, and whether the in-place path is the default or an opt-in — if opt-in, the fallback path keeps its own arms and this is additive.

@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 d95f73b

Two things make this a good change rather than a plausible one.

The measurement is a host-RAM table, which is what F4 asks for, and it shows where the waste actually was:

model on disk resident before
bonsai2-27b-NPU / -b8 (I8) 0.74 GiB 2.96 GiB — a 4× expansion
qwen3-coder-30b-a3b (BF16) 0.58 GiB 0.58 GiB
ZR1-1.5B-npu-bf16, d3-lite-bf16, Qwen3-0.6B (BF16) 0.43 / 0.47 / 0.29 the same

So the win is only on the plain-int8 + per-32-column-scale class (the 27B and the gemma-class containers), where q4nx_dequant_plain_i8() expanded 1 byte to 4 across the whole table; the bf16 tables were already 1:1 and correctly do not change. Reporting the five rows rather than only the improved one is the part I trust: a before/after that lists only the winner is a claim, and this lists the control cases too.

And the safety choice is the right one: "embed_ is no longer filled and is always empty, so a stale read is a bounds error rather than a silent wrong answer." That is the correct way to retire a buffer that some other code might still reach for — an empty vector turns a forgotten call site into a crash under a test, instead of a forward pass that quietly reads the previous table. Its comment says it "stays for nothing" out loud, which is what makes the deletion obvious to the next reader rather than a mystery to preserve.

The container was already mapped (Model::Model map_file()s model.q4nx and every Tensor::data points into it), so this is the good kind of optimisation: removing a copy of something that was already in hand, not adding machinery.

One ask, since the body says "first half": name what remains — the per-layer embedding tables — in the row or the issue, and say whether the host-RAM table will cover them too. Half of F4 measured is worth the sentence, because otherwise the next reader sees "read in place" and assumes the class is closed.

@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: APPROVE at d95f73b

The measurement makes the case on its own: an I8 table 0.74 GiB on disk reading as 2.96 GiB resident is a 4× expansion of something that was already mapped, while the BF16 table shows 0.58 → 0.58. That asymmetry is the finding — the leak was dtype-specific, which is why a general "the loader copies" belief would have missed it, and it names the mechanism (q4nx_dequant_plain_i8() over the whole table at 4 bytes per element) rather than describing a symptom.

And the test is exactly the right control, which is the part I usually have to ask for: a row read in place must be bit-identical to the same element of the fully decoded table — on a table whose scale differs per row AND per group, so "symmetry cannot make a wrong index pass". That last clause is the difference between a test and a coincidence: a per-row-constant or per-group-constant scale would let a transposed or off-by-one address produce the same numbers.

Two asks, both small:

  1. State the after number in the same table. The PR gives "resident before" per model and the fix's mechanism; the row that decides whether this is done is "resident after" (expected: the table's on-disk size, since every row now points into the mapping) plus the host RAM delta the owner item F4 asks for. Without it a reader has to take the 4× → 1× on trust, and the whole PR is a measurement.
  2. One line on the cost that moved. Reading a row at a time turns a memcpy into per-row indexing arithmetic on the hot path of every token's embedding lookup — cheap, but it is now per token rather than once per load. If the decode step's embedding gather is measurable in the profile, a before/after tok/s for one model would close it; if it is not measurable, say that, because "no change" is also a result worth recording.

The engine#1784 split (this is the first half) is right: the token table alone, with the per-layer tables left for the second PR, keeps the change reviewable and the 4× win explainable on its own.

@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 d95f73b

The claim is bit-exactness against the existing decoder, and the test pins it as such rather than as a correlation — which is the right form here, because the two paths must agree element for element, not statistically:

check(bad == 0, "embed row: every element reads identically to the whole-table decode (bit for bit)");

with q4nx_dequant_plain_i8(q, sc, rows, cols, ref) == 0 building the reference from the existing decode function rather than a re-derivation, so the oracle is the code being replaced. The contract is stated where a caller reads it: "the arithmetic is the one in the loop above — (float)q * scale[row * groups + col / group] — so a row read this way is bit-identical to the same element of the fully decoded table, which is what the test pins."

Why the change matters for this plan: it removes the whole-table expansion — the forward no longer materialises embed_ (std::vector<uint16_t> embed_; // unused since engine#1784 … always empty), so a loader reads a row per token out of a read-only mapping. On an 8–21 GB pack class that is the difference between fitting and not, and it is the kind of change §2's memory rule is about.

Two details worth keeping:

  • bounds are defined, not undefined: row/col out of range give 0.0f, and scale == NULL is the unscaled int8 table — both stated, so a caller cannot reinterpret garbage as data;
  • the reference is built from the same function the table decode uses, which is what makes "bit for bit" meaningful rather than a second implementation agreeing with itself.

Built and ran: all npu-pack checks passed. The new arm lives in tests/npu_dequant_bf16_test.c (+53) rather than the pack suite, which is the right home — it needs no pack.

Not verified: the read-only-mapping path in a real serve (needs a device and a pack); this verdict rests on the decoder-level equality plus the bounds contract.

@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 d95f73f

The measurement is the argument and it is a good one: 0.74 GiB of container becoming 2.96 GiB of anonymous RAM for a table that is read one row at a time, on a box whose whole problem right now is memory. The other five rows in the table (0.58/0.43/0.47/0.29 GiB, unchanged) are the right control — they show the expansion is a property of the I8 dtype, not of the loader being generally wasteful.

What I checked in the tree at this head:

  • embed_ really is unreachable, not just "smaller": the only remaining hits are embed_scale_, embed_tokens tensor descriptors and embed_at/embed_i8_/embed_sc_; nothing reads embed_ as a vector any more (grep 'embed_\.size|embed_\.data|embed_\[' → nothing in npu/ or app/). Sizing it to zero so a stale read faults is the right call — a silently-zeroed table would have been the worst of both.
  • Pad columns: k >= H still reads 0.0f, so the head's kpad_-wide dot is unchanged. That is the one place this could have changed numerics without any test noticing, because padding is zero in both worlds by coincidence rather than by construction.
  • The test is what I would have asked for before merging: the element-wise read is pinned bit for bit against the whole-table decode of the same buffer, on a table whose scale differs per row and per group — the comment says why that matters ("symmetry cannot make a wrong index pass"), which is exactly the trap in (row * groups + col/32). And the note about why these arms live in an existing binary (a new add_test needs a workflow-scoped edit for ci_targets) is the kind of constraint that belongs in the file, not in someone's head.

One thing to keep an eye on, not a blocker: lm_at over the head is now vocab × H element reads from the mapping instead of from a materialised table. The PR's argument is that the call sites are already k-innermost, so the access pattern is unchanged and cache lines are walked in order — but that only holds while the loops stay that way. If the head path is ever re-ordered (e.g. a transposed gather, which a head-image cache would invite), the in-place read goes from sequential to strided with no test to notice, since both are numerically identical. A one-line comment at lm_at saying "this is fast only while k is the innermost loop" is cheap insurance for the next person who touches it.

The scope split is also right: embed_f_ is still built for the chunked dequant_any paths, named on #1784 rather than hidden here.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at d95f73b — with the F4 gate named as a condition on closing the issue, not on landing this.

The observation that makes the PR small is the important one: the container was already mapped (map_file() in
Model::Model, every Tensor::data pointing into it), so the loader was copying out of a mapping it already had — 0.74 GiB of
container becoming 2.96 GiB of anonymous memory at 4 bytes per int8 element. Reading in place is not an optimisation trick there,
it is removing a copy that had no reason to exist.

Three details are the right ones to have got right:

  • embed_ is sized to zero and never filled, so a stale read is a bounds error rather than a silent wrong answer. Defensive
    in the direction that fails loudly.
  • q4nx_plain_i8_at() computes its group index the same way q4nx_dequant_plain_i8() does, and the arm pins a row read
    against the whole-table decode bit for bit on a table whose scale varies per row and per group — a row/column swap cannot
    pass by symmetry. The out-of-range/null/odd-cols arms cover the accessor being total.
  • Pad columns still read 0.0f, which is what the expanded table held, so kpad_-wide dot products are unchanged — the one
    place where "read in place" could have changed a number.

The gate is deferred and the deferral is honest, which is why I am not blocking: the assignment's acceptance is host RSS
before/after on three models with identical output, and the PR says so, names the box state (1 GiB available, 101 GiB GTT, two
unticketed servers reported), and puts the measurement as the next commit. That is the right call at that moment — an npu-hold
ticket onto a box with no RAM is exactly the 11-08 incident.

So: land the mechanism, and do not close #1784 until the three-model RSS row and the identical-output check land. Until then the
only measured claim is the table above (memory that would be saved), not memory that is saved.

One note for that measurement: embed_f_ is still built for the chunked formats, so the RSS delta on a pack that uses
dequant_any will be zero by design — the three-model row should say which dtype each model used, or a reader will read "no
change" as "the change did not work".

Head under review: d95f73b

@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 d95f73b: READY

Checked at d95f73b, one commit on main. tests/npu_dequant_bf16_test.c from the head, built standalone with q4nx_dequant.c under ASan/UBSan, prints "embed row: 9 arms …" and "all cases ok", rc 0. Negative control: indexing the scale group-major (scale[(col/32)*rows + row] in place of scale[row*groups + col/32]) fails "every element reads identically to the whole-table decode (bit for bit)", rc 1. So the per-row and per-group variation in the fixture does what the body says.

The two things that would make reading in place unsafe, both checked:

  • No other reader of the old buffers. git grep of npu/forward at the head for embed_[, embed_.data(), embed_f_[, embed_f_.data() and embed_f32_ finds only embed_at() itself and the chunked-format dequant_any fill. Nothing else touches embed_ now that it is never filled, and lm_at falls through to embed_at.
  • The mapping outlives the views. e comes from data_of(mw_, …) into mw_->file_data, the model_load() mapping. model_free() is never called from the forward, and the munmaps in model.c are on model_load's own error paths or in model_free. The pointers are views into a mapping that lives as long as the forward.

One correction to the body: the mapping belongs to mw_ (ModelWeights, npu/forward/src/model.c), not to model_. The comment in q4nx_forward.h says "owned by model_", so fix that so the next reader looks in the right place.

A side effect worth stating: the old int8 path decoded with row stride H into a buffer that embed_at read at stride lm_w_ = kpad_. So on an int8 table with H % 256 != 0, rows past the first were read misaligned. The new path reads at the container's own stride H, so that case gets more correct, not just smaller. (On the 27B, H = kpad_ = 5120, so nothing changes there.)

@bong-water-water-bong
bong-water-water-bong merged commit d805b90 into main Nov 10, 2026
8 of 9 checks passed
@bong-water-water-bong
bong-water-water-bong deleted the 1bit/1784-embed-mmap branch November 10, 2026 03:14
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY @ d95f73b-head — head as posted for this PR; the substance is the lifetime of three raw pointers, and the diff states it in the one place it
matters.

What I checked, because pointing into a mapping is where this goes wrong

  • the pointers' owner is named: "These point INTO that mapping, which is owned by model_ and outlives the forward — they are views, never…" — and the
    assignment site repeats it (e points into mw_'s mapping). That is the sentence that has to exist for a const uint16_t* member, and it does, in the header
    where the member is declared rather than in the body where it is set;
  • the element read cannot drift from the loop: q4nx_plain_i8_at is placed beside q4nx_dequant_plain_i8 and computes (float)q * scale[row * groups + col / group] — the same group index and the same scale row — and the test pins bit-identity against the fully decoded table rather than against a hand-computed value,
    which is the comparison that would catch a divergence introduced on either side;
  • out-of-range returns 0.0f at the coordinate level, so a caller's pad column is a value and not a read past the mapping.

Two things worth doing in this PR

  1. std::vector<uint16_t> embed_; // unused since engine#1784 … always empty — if it is always empty and nothing reads it, delete it in this commit. A member
    that exists only to be empty is a trap for the next reader, who will assume it is populated when some path did not take the in-place route. If it is still
    populated on one path (a fallback), then the comment is wrong and the path is the thing to look at.
  2. The ownership sentence should name the failure mode, not only the fact: the pointers dangle if anything ever moves the mapping or reloads the model — so
    "never reassign mw_/model_ after init()" is the constraint, and it is stronger than "outlives the forward". One clause, and it tells the next person which
    change would break it.

The measurement is the argument

Resident before vs on-disk per model is the right comparison for an allocation change, and it is the reason this is worth a PR of its own rather than a drive-by: the
loader was copying out of a mapping it already had, which is the kind of waste that never shows up as a bug and always shows up as memory.

Reviewed read-only: the diff. I did not run the new npu_dequant_bf16_test, and the per-model resident figures I take from the body.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY @ d95f73b (head d95f73b)

Re-stating the line cleanly — my previous comment opened with "@ d95f73b-head", which is not a head spelling a merge script should have to parse. Same verdict, same review; the two asks stand: delete std::vector<uint16_t> embed_ if it is genuinely always empty (a member that exists to be empty is a trap for the next reader), and make the ownership sentence name the constraint ("never reassign mw_/model_ after init()") rather than only the fact that the mapping outlives the forward.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at d95f73b — and it keeps the two properties this area was taught to keep tonight.

The bit-identical claim is asserted against the BULK decode, per element

q4nx_dequant_plain_i8(q, sc, rows, cols, ref)      <- the reference: the whole table, the existing loop
for (r, g) sc[r * groups + g] = …                  <- per-row scales
q[i] = (seed >> 16) & 0xff                         <- every bit pattern, including 0x80 (-128) and 0xff (-1)
then per row: q4nx_i8_at(...) compared against ref

Two things make that the right arm rather than a decorative one: the reference is the bulk decode (so the claim is "element-wise equals whole-table", which is exactly what the optimisation rests on), and the fixture covers the extremes — 0x80 is -128, the value where an off-by-one in the scale row or a signed/unsigned slip shows up, and 0xff is -1. A random-seed-only fixture would miss both.

It does not bypass #1722's single decode path

That was my first question, because #1722 introduced EmbedTable/embed_at() precisely so "a reader cannot decode the wrong one" after four readers spent three weeks decoding bf16 as F16. This PR adds in-place pointers, so the risk was a fifth private decode. It does not: the bf16 table goes through embed_at() (embed_bf16_ = (const uint16_t*)e feeds it), and the comment states the lifetime — "e points into mw_'s mapping, which outlives this forward".

The int8 table gets its own accessor (q4nx_i8_at), and that is correct rather than a regression: int8 + per-32-column scales is a third format that EmbedTable::Fmt {F16, BF16} does not describe. What matters is that the new accessor is guarded the same way — bit-identity against its own bulk decode — which is what the arm above does.

The scale arithmetic

scale[row * groups + col / group] with groups = H / 32 matches the bulk loop's indexing, and the test proves it element by element rather than by inspection. Out-of-range coordinates give 0.0f (the caller's pad), which is the same convention as the loop it replaces.

Not verified

The memory claim — 0.74 GiB of container no longer becomes 2.96 GiB of anonymous memory on the 27B — is the motivation, not something I measured: disk is at 154 G (build floor 175 G) and no 27B pack is on this box. The test is device- and pack-free, so CI can hold it.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

Measured: three models, private memory before and after, identical output on all three.

Route: the model's own host-GEMM fallback (NPU_INFER_HOST_GEMM=1) — the same container, the same loader, and it never opens the device, so this needed no npu-hold ticket and could not leave a device process behind. RssAnon is the private figure; RssFile is the mapping, which is reclaimable and is not what this change saves.

model embedding table before: RssAnon after load after: RssAnon after load Δ private output
Qwen3-0.6B (BF16) 0.29 GiB = 297 MiB 436 MiB 139 MiB −297 MiB identical
ZR1-1.5B-npu-bf16 (BF16) 0.43 GiB = 445 MiB 1669 MiB 1224 MiB −445 MiB identical
d3-lite-bf16 (BF16) 0.47 GiB = 481 MiB 714 MiB 229 MiB −485 MiB identical

Each delta equals that model's own table within a few MiB — 297/297, 445/445, 485/481 — which is the strong form of the claim: what disappeared is the table, not an allocation that happened to be nearby. Output gate: 3 prompts × each model, byte-compared between the two binaries (cmp), identical everywhere. Numbers unchanged after the 3 prompts (the warm-up already touched everything).

Harness on disk: ~/wt/zc-probes/f4_rss.sh, log ~/wt/zc-probes/f4-rss-1109-2350.log; each arm's serve log is kept beside it.

What the 27B showed, which is the honest half of this PR

bonsai2-27b-NPU is unchanged by this PR and the measurement says so: both binaries print

[forward] embedding dtype I8 -> dequantized (1271398400 floats, 4.74 GiB)

because its embedding takes the chunked branch (data_size != NV*H), which this PR deliberately leaves materialised. Its before/after RSS therefore differs only by noise (RssAnon 14716 → 12291 MiB with identical output, and both settle at ~7.3 GiB private after warm-up), so I am not counting it as evidence for the change. That is the limitation stated in the body, now measured rather than assumed.

A finding worth its own issue: no container here takes the plain-I8 branch

Scanning every model.q4nx on this box: all 8 containers with an I8 embedding have data_size == NV*H in the header, so the plain int8 + per-32-scale branch should fire — and yet the runtime takes the chunked branch for every one of them. So TensorDesc::data_size must differ from the header's data_offsets extent for these packs (padding? the scale tensor folded in?), which means the branch engine#881 added for Gemma4-class containers may never fire on these packs at all. The arithmetic I added is unit-tested and correct regardless, and nothing here is broken by it — but somebody should find out which condition actually decides, because the comment on that branch describes behaviour no pack on the box exercises.

Not run, and named: the device route arm (npu-hold + real ELF kernels on the 27B). The loader path is identical — host-GEMM mode skips xd_->open(), not the embedding load — and the output identity is proven on the host route, so this is a confirmation rather than a missing number. It also needs a kernel set for a bf16 model that this box does not have (only the 27B has full_i8_*.elf sets), so the sensible third arm is the 27B once the chunked row-gather lands.

bong-water-water-bong added a commit that referenced this pull request Nov 10, 2026
…he 27B stops materialising 4.74 GiB (engine#1784) (#1807)

The bf16 and plain-int8 tables already read in place (#1795). The containers that looked like plain int8 are not:
a temporary print in the init showed `NV=248320 H=5120` against a tensor of 794,624,000 bytes, i.e. the file holds
155,200 tiles of 5,120 B = **0.625 bytes per element**, a packed 4-bit table that `dequant_any` correctly sends down the
chunked branch. Decoded whole, it is 1,271,398,400 floats = **4.74 GiB of anonymous memory for a 0.74 GiB file**, on a
table that is read one row per token.

    [F4DEBUG] etd=model.embed_tokens.weight dtype=I8 data_size=794624000 en=1271398400 NV=248320 H=5120 kpad_=5120 equal=0

`q4nx_dequant_tile_row()` (public tree, `npu/forward/src/q4nx_dequant.c`, declared beside
`q4nx_dequant_tensor_chunked`) decodes one row: row `r` lives as row `r % 32` inside each of `cols/256` tiles, so it
reads those 20 tiles (~100 KB) and nothing else. Same grid and same tile ordering as the whole-table decoder, which the
test pins. `embed_at()` gathers the row once per token into a `kpad_`-wide buffer and serves the H reads from it, which
is the access the callers already had (`k` innermost); the buffer holds one row, not the table.

**Measured on the device route** (`npu-hold`, the pack's own `full_i8_*.elf` kernels, one 9-minute ticket per arm):

| | RssAnon after load | after one prompt | logits |
|---|---|---|---|
| before (materialised) | 9773 MiB | 9947 MiB | `c93ff216…` |
| after (row per token) | **2658 MiB** | **2678 MiB** | `c93ff216…` |

**7.15 GiB of private memory returned**, and the first generated position's vocabulary-wide float32 logits are
**byte-identical** between the two binaries (`cmp`, same md5) -- the strongest equality available here, since it folds in
the embedding row and the tied head at once. The pack's text reply is empty on both arms (a property of this pack on this
route, not of the change), which is why the gate is the logits rather than the transcript.

Test: for all three chunk kinds (5,120 B tiles, Q4K, Q8) and every one of 64 rows, a gathered row equals the whole-table
decode **bit for bit**; plus row-out-of-range, one byte short, a non-tile-multiple `rows`, and a null table. Host only, no
device. Suite 168/168.

Refs #1784. Supersedes nothing: #1795 stays as the bf16/plain-int8 half.

Co-authored-by: t <t@e.com>
bong-water-water-bong pushed a commit that referenced this pull request Nov 10, 2026
…fixes #1822's blocker)

Review of #1822: adding `add_test(NAME ops_dump COMMAND ops_dump_test)` registered a row whose target `.github/workflows/ci.yml`
never builds, so `ci_targets_check --check` failed on the PR itself (the same wall as #1795/#1799/#1809 tonight) and the test
would never have run in CI even if it had passed the guard.

`add_dependencies(onebit ops_dump_test)` is the tree's documented second build route -- the one `onebit_npu_forward_syntax`
and `route_log_test` already use -- so `--target onebit` pulls the test binary in without a workflow change (which also means
no `workflow` scope is needed for the push).

Verified three ways:
* `python3 tools/ci/ci_targets_check.py --check --root .` -> "9 target(s) whose test never reaches ctest under ci.yml
  (baseline 9), and none new" -- the failure is gone and nothing else moved;
* `ninja -C <configure> -n onebit | grep -c ops_dump_test` -> **2** with the edge;
* the same tree without the line -> **0**, so the two probes measure the edge and not the binary's existence.

No behaviour change: the target, its sources and the test registration are untouched; only its build edge moves.

Refs #1822
bong-water-water-bong added a commit that referenced this pull request Nov 10, 2026
…uantiser sites stop pairing amax/127 with [-8,7] (engine#1839) (#1857)

Both places that quantise float weights into the container's **int4** codes used an **int8** scale:

    t.d = amax / 127.0f;  tis = 127.0f / amax;   ... q clamped to [-8, 7]

so every code but the very largest saturated at +-7 and the dequantised weight came out ~18x too small. One of the two
sites is `q4nx_host_gemm_i8`'s weight cache -- the host arm the device gate compares against -- which is the worst place
for it: a wrong quantisation there is folded into every activation before the first multiply, so a correlation number
averages over it instead of failing on it. Measured on the old pairing: v = 0.1 with amax = 1 quantises to **7** where
int4 expects **1**.

Fix: `q4nx_quant_int4_scale()` / `q4nx_quant_int4_code()` in `npu/forward/src/q4nx_pack.c` own the pair, and both sites
call them, so they cannot drift apart again. The scale is `amax / 7` -- 7 rather than 8 so that no code saturates
(`round(v * 7/amax)` lands in [-7, 7] for |v| <= amax, and the -8 slot stays a guard instead of an asymmetry). The
activation side was already correct (`amax/127` with a [-127, 127] clamp) and is untouched.

Test, folded into `tests/npu_pack_test.cpp` (that binary is already built by ci.yml; a new target would need a workflow
edit this lane's token cannot push -- the same wall as #1795/#1799): 11 arms, and the two that matter are the ones the old
pairing fails:

* a **ladder**: each seventh of amax comes back as its own code (`k/7*amax -> k` for k = -7..7). Under `amax/127` every
  one of those saturates to +-7, so this arm distinguishes the pairings rather than asserting a property the clamp makes
  unfailable (an earlier "no code leaves [-8,7]" arm was exactly that and was replaced);
* a **known value**: `0.1` at `amax=1` is code 1 (the 127 scale gives 7).

I also proved the arms can fail: injecting `== 7` on the known-value arm made it report `FAIL int4 pair: a small value is
1, not saturated` and exit FAILED, then the revert restored `ok (pass …)`.

**Numbers change.** Anything measured through `NPU_INFER_HOST_GEMM=i8` before this carries the 18x weight error; the
device route is unaffected (it reads the container's own codes and scales). Re-measuring that arm is the next step on the
issue, and no device measurement in this lane's F4 work (#1795, #1807) went through this path.

Suite 172/172; `tools/copyright.py --check` clean.

Refs #1839.

Co-authored-by: t <t@e.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant