Repository navigation
npu: read the token embedding table in place, one row per token, instead of materialising it (engine#1784) - #1795
Conversation
…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.
|
VERDICT: READY at 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 Two notes:
The |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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:
- 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.
- One line on the cost that moved. Reading a row at a time turns a
memcpyinto 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
left a comment
There was a problem hiding this comment.
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/colout of range give0.0f, andscale == NULLis 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
left a comment
There was a problem hiding this comment.
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 areembed_scale_,embed_tokenstensor descriptors andembed_at/embed_i8_/embed_sc_; nothing readsembed_as a vector any more (grep 'embed_\.size|embed_\.data|embed_\['→ nothing innpu/orapp/). 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 >= Hstill reads0.0f, so the head'skpad_-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 newadd_testneeds aworkflow-scoped edit forci_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.
|
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 ( Three details are the right ones to have got right:
The gate is deferred and the deferral is honest, which is why I am not blocking: the assignment's acceptance is host RSS So: land the mechanism, and do not close #1784 until the three-model RSS row and the identical-output check land. Until then the One note for that measurement: Head under review: d95f73b |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
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 grepofnpu/forwardat the head forembed_[,embed_.data(),embed_f_[,embed_f_.data()andembed_f32_finds onlyembed_at()itself and the chunked-formatdequant_anyfill. Nothing else touchesembed_now that it is never filled, andlm_atfalls through toembed_at. - The mapping outlives the views.
ecomes fromdata_of(mw_, …)intomw_->file_data, themodel_load()mapping.model_free()is never called from the forward, and themunmaps inmodel.care onmodel_load's own error paths or inmodel_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.)
|
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 What I checked, because pointing into a mapping is where this goes wrong
Two things worth doing in this PR
The measurement is the argumentResident 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 Reviewed read-only: the diff. I did not run the new |
|
VERDICT: READY @ d95f73b (head 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 |
|
VERDICT: READY at The bit-identical claim is asserted against the BULK decode, per elementTwo 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 — It does not bypass #1722's single decode pathThat was my first question, because #1722 introduced The int8 table gets its own accessor ( The scale arithmetic
Not verifiedThe 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. |
|
Measured: three models, private memory before and after, identical output on all three. Route: the model's own host-GEMM fallback (
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 ( Harness on disk: What the 27B showed, which is the honest half of this PR
because its embedding takes the chunked branch ( A finding worth its own issue: no container here takes the plain-I8 branchScanning every Not run, and named: the device route arm ( |
…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>
…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
…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>
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()smodel.q4nxand everyTensor::datapoints into that mapping — so the loader was copying out of a mapping it already had: bf16 tables viamemcpyof the wholeNV*Htable, and plain int8 tables viaq4nx_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.q4nxheader):bonsai2-27b-NPU/-b8(I8)qwen3-coder-30b-a3b(BF16)ZR1-1.5B-npu-bf16(BF16)d3-lite-bf16(BF16)Qwen3-0.6B(BF16)What changed
embed_bf16_is the pointer into the mapping;embed_at()converts the one element asked for.embed_i8_/embed_sc_are pointers into the mapping, read through the newq4nx_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.k >= H) still read0.0f, which is what the expanded table held, so the head'skpad_-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 withkinnermost (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, sincelm_atfalls back toembed_at.Test
q4nx_plain_i8_at()is declared besideq4nx_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 acolsthat is not a group multiple. Host-only, no model, no device.While writing it, the repo's own
ci_targetstest caught me: a newadd_test()must also be built by.github/workflows/ci.yml, and an edit to that file needs a token with theworkflowscope, which this one does not have. So the arms went intotests/npu_dequant_bf16_test.c, which already linksq4nx_dequant.cand 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 unticketedllama-serverprocesses 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.