Repository navigation
npu: Embedding-Gemma-300M-NPU2 serves /v1/embeddings on the per-op forward (refs #1786) - #2314
Conversation
…rward (refs #1786) FastFlowLM's Embedding-Gemma pack is a Gemma3TextModel with use_bidirectional_attention and the sentence-transformers head. It now serves through `1bit serve --embedding -m <dir> --device npu|cpu`: - loader: the unprefixed names (layers.N.*, embed_tokens, norm) get `model.`; - dequant: an unquantised BF16 [out, in] projection is read into the padded grid; - semantics: use_bidirectional_attention makes every layer, sliding ones included, read the whole input in one pass (an input past half the window is refused by name); a gemma3 pack with BF16 projections is an HF copy whose norms are raw w, so it runs as 1 + w; - Q4nxNpuForward::embed_pooled(): one batched pass at position 0, the final normed rows mean-pooled, dense_2 and dense_3, L2; - embed_serve: a q4nx encoder kind beside the ternary container, ids framed by tokenizer_config's add_bos_token / add_eos_token. Against transformers' Gemma3TextModel on unsloth/embeddinggemma-300m: cos 0.99999 on the host, 0.9991 on the NPU (5 inputs in 0.88 s). tests/embed_forward_ci.py pins it on a fabricated container (cos 1.0000000); the sliding-causal and no-offset mutants fail it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
ready for review at 3985fd9 |
PR Reviewer Guide 🔍(Review updated until commit 6982044)Here are some key observations to aid the review process:
✅ Resolved findingsnpu/forward/src/q4nx_forward.cpp:4417-4435Incorrect Handling of Sentence Transformers Head The app/embed_serve.cpp:244-252Potential Race Condition The npu/forward/src/q4nx_forward.cpp:4443-4447Incorrect Handling of Norms The npu/forward/src/q4nx_forward.cpp:200-213Incorrect Handling of BF16 Projections The npu/forward/src/q4nx_forward.cpp:4408-4466Potential Memory Issue The |
|
VERDICT: READY at The three gaps, in the order that matters
Gap 2 is the one I would have expected to be missed and silently shipped: the endpoint returns vectors either way, and the only way to notice is that retrieval quality is poor — which is exactly the kind of defect that survives a smoke test. Why the endpoint claim matters" Two asks
And the load refusal is worth keeping visible"stopped at load ( |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 3985fd9
The test is the reason to land this, and it is the standard I would ask for: a numpy Gemma3 encoder over the same weights, cos 0.99999 on the host and 0.999093 on the NPU, with each rule's own discriminating number
next to it -- "the norm rule alone: 0.09 on the real pack, the sliding layers left causal: 0.70 to 0.90". Arms that name what a wrong implementation would score are the ones that catch a regression; a single cosine would not
tell the norm rule from the pool rule.
One thing to answer before a second gemma3 pack arrives, and it is not blocking because this PR only ADDS a case:
if (sem_.gemma3_bf16_norm_probe && !sem_.rmsnorm_offset) {
const TensorDesc* q0 = model_tensor_by_name(mw_, "model.layers.0.self_attn.q_proj.weight");
if (q0 && !strcmp(q0->dtype, "BF16")) { sem_.rmsnorm_offset = 1; … }
}
if (sem_.rmsnorm_offset_probe && !in_norm_.empty() && c_.H > 0) {
sem_.rmsnorm_offset = q4nx_rmsnorm_offset_from_norm(w0.data(), c_.H); // the DIRECT instrument, 15 lines down
}q_proj's dtype is a proxy for "the norms are HF's raw w", and the direct instrument already exists in the same file: q4nx_rmsnorm_offset_from_norm() decides from the norm weight's own mean and minimum, with the
reason written on it. For this pack the two agree (raw w -> offset 1), so nothing is wrong today. But a gemma3 pack that is quantised and still raw w reads not BF16 and keeps offset 0 -- which your own measurement
prices at cos 0.09. Two selectors for one question is the shape that drifts, and the direct one is the answer. If the ordering forces the proxy (in_norm_ not populated for gemma3's unprefixed layout at that point), say
so in the comment -- that is the sentence a later lane needs.
And the negative branch has no arm: the fixture ships BF16 projections, so only the proxy's true branch is exercised. The false branch is unchanged behaviour (gemma3 got offset 0 before this PR too), which is why I would
not hold it -- but a second fixture with quantised projections and (1 + w) norms would pin the boundary, and the harness already generates the container in Python, so a dtype knob is the whole change.
Checked and it holds:
bidirectionalextends the existing non-causal mode rather than adding a second one (s->bidirectional = 0; if (get_num(json, "use_bidirectional_attention", &d)) s->bidirectional = d != 0;), so an absent key keeps a causal
gemma3 -- absent stays distinguishable fromfalse.- The fixture is not a stub: unprefixed names,
embed_tokenspre-scaled bysqrt(H), norms as HF's raww, a sliding then a full layer, an intermediate size off the 256 grid (384 -> 512), and a second request past half the
window must be refused -- that last arm is the one that keeps the packed-window bound honest rather than implied. usage.prompt_tokensis 83 and equals the reference ids' total, so the tokeniser path is pinned to the same numbers as the embeddings.ctest -j 6without device tests 278/278, every forward arm by hand, and the NPU row run inside annpu-holdwith amdxdna clean before and after -- the check this box requires.- The kernel set is named: "locally built and not shipped (4 designs,
QKV K768 N1536,O K768 N1024,GU K768 N2…)" -- so the device row's provenance is stated instead of implied, which is what I want on a device number.
Not run: the device row and the real pack (yours, in an npu-hold ticket).
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT at 3985fd9: CHANGES
The build is red in Test, and the gate names both the site and the fix — this is one call swap, not a redesign.
139/283 Test #139: file_create_check ... ***Failed
tests/embed_pooled_test.cpp:52: fopen for writing creates 0666 and follows a symlink:
FILE* f = std::fopen(out …
file_create_check: 1 site(s). Use onebit_fopen_private_wb / onebit_fopen_shared_wb
(npu/forward/include/file_create.h)
std::fopen(path, "wb") creates at 0666 before the umask and follows a symlink at the final component. On a box whose /tmp is a tmpfs shared by every lane, a dump written that way is world-writable and can be redirected. file_create.h already has the reviewed helpers and its own words are the argument: "a file that holds activations, weights or a request is a dump and takes the 0600 one."
So: onebit_fopen_private_wb(out) at tests/embed_pooled_test.cpp:52. That is where the other test that writes a dump ended up (#2254 earlier today), and it is the same one-line shape.
Worth saying plainly, because it is the good news in a red build: this is the third instance of the class tonight (#2254, and two runtime sites before it), and this one was caught mechanically, by a gate, with the remedy in the message — rather than by a reviewer noticing. file_create_check_self passes in the same run, so the gate is testing itself too. Earlier today I asked for a test_file_create helper so the class would end rather than the instance; someone built the checker instead, and this run is it working.
Everything else in the PR compiled: 282 of 283 tests pass, so the embeddings path itself is not implicated by this failure — fix the one site and the pipeline should reach the end for the first time.
…te_wb (refs #1786) CI's file_create_check (and CodeQL) flagged the plain fopen: 0666 and a followed symlink. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
CI at ready for review at b76acb6 |
|
Persistent review updated to latest commit b76acb6 |
|
VERDICT: READY at Checked in order of how easy each would have been to get subtly wrong:
One note, not a blocker: the |
|
VERDICT: READY at Where the arms actually liveThe test surface here is
Still open from my last review (unchanged)
Why the PR is sound as it standsThe three gaps are correctly diagnosed and each is fixed at its own layer: the loader learns an embedding family (rather than the check being deleted), the forward learns bidirectional attention (a different path, not a flag), and the endpoint learns a second head family ( |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at b76acb6
Carried from 3985fd9, where the substance was reviewed; the only delta is one line in the test, and it is the house convention applied.
- FILE* f = std::fopen(out_path, "wb");
+ FILE* f = onebit_fopen_private_wb(out_path); // 0600, no symlink follow (file_create.h)npu/forward/include/file_create.h already documents this as the rule for dumps and probes ("0600: dumps and probes (an activation or weight that holds…"), and batch_check_state_test.cpp and
npu_pack_import_test.cpp use it -- so the pooled-embedding fixture was the one writing a 0644 file through a symlink-following fopen while its neighbours did not.
Everything from my verdict at 3985fd9 stands unchanged: the numpy Gemma3 encoder reference at cos 0.99999 host / 0.999093 NPU, the per-rule discriminating numbers ("the norm rule alone: 0.09 on the real pack, the
sliding layers left causal: 0.70 to 0.90"), the fixture's off-grid intermediate size (384 -> 512), the past-half-window request that must be refused, usage.prompt_tokens 83 matching the reference ids, ctest -j 6 278/278,
and the device row inside an npu-hold with amdxdna clean before and after.
And the two things I raised stay raised, neither blocking: the gemma3 norm convention is still decided by the q_proj-dtype proxy while q4nx_rmsnorm_offset_from_norm() sits 15 lines below as the direct instrument
(they agree for this pack; a quantised pack with raw w would read not BF16 and keep offset 0, which your own measurement prices at cos 0.09), and the proxy's false branch still has no arm.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT at b76acb6: READY
I tested at 3985fd9. The one commit since, b76acb69, changes only how embed_pooled_test writes its output (onebit_fopen_private_wb, 3 lines), and none of what I ran depends on that. The merged tree with main changes the PR's 13 files (+576 −5). Each of the 5 removed lines is replaced by its edited form: max_tokens, the embed_serve::container branch, the two l_full sites and the dump guard. I used an ONEBIT_NPU host build (NPU_INFER_HOST_GEMM=1), with embed_pooled_test compiled against that tree's libonebit_npu_forward.a. Under 175 G free I could not configure a separate host-only tree.
Arms and mutants (the numbers in the body reproduce):
| run | result |
|---|---|
embed_forward at head |
PASS, cos 1.0000000, and the 9-token input is refused by name |
mutant: (l_full || bidir_sliding_) → (l_full || false), 2 sites |
FAIL, cos 0.9905108 |
| mutant: BF16 norm probe off | FAIL, cos 0.2364964 |
| mutant: over-long refusal off | FAIL, "embedded 9 ids -> 256 floats" |
batch_check_state arms: default, qwen2, qwen2_untied, lfm2, lfm3, qwen2vl, qwen3_5, hy_v3 |
0 FAIL each |
The BF16 norm rule does not catch the existing gemma3 packs.
- I scanned every gemma3/gemma family container on the box with a config.
gemma3-270m-npu2staysrmsnorm_offset=0and passes.gemma4-e2b-npu2staysrmsnorm_offset=0.- Only
Embedding-Gemma-300M-NPU2switches tormsnorm_offset=1 <- gemma3 with BF16 projections.
gemma4-e2bfails 3 "holds exactly N positions" checks, and it fails them identically on main at9a589d7, which I rebuilt to check. Logits still match (rel-L1 0, argmax 107). So that predates this PR. It is gemma4's KV layout, not this change, and is worth a note on #1786's gemma4 row.
Real pack, host route, through embed_pooled_test with <bos> text <eos> ids from the pack's own tokenizer.json. These are five sentences of mine, not the PR's:
- The vector has 768 dims and unit norm.
- cos 0.838 for "The cat sat on the mat." vs "A cat was sitting on the rug.".
- cos 0.706 for the two stock-market paraphrases.
- No cross-topic pair is above 0.266.
That is a sanity check, not a reference. I have no transformers or torch here, so the 0.999990 / 0.999093 rows against unsloth/embeddinggemma-300m are the author's, and the NPU row is too.
Code read against sentence-transformers' EmbeddingGemma:
- Mean pooling runs over every row, special tokens included, after the final norm.
dense_2thendense_3, with no activation and no bias, then L2. A pack without the dense head pools to H.- The refusal at
M > window/2is the conservative side of the bidirectional band. HF's band is about half the window either way, and 256 tokens sit inside it on either convention.
One thing to keep in mind, not blocking: "gemma3 + BF16 q_proj ⇒ raw w" is a rule about provenance. A future BF16 gemma3 pack written by a converter that does add the 1 would be double-offset, with no error. The load line names the rule, which is what someone would need to find it, and the mutant shows the encoder arm depends on it.
CI: the ci build job at this head passed, 100% tests passed, 0 tests failed out of 283, and embed_forward (test 265) ran and passed in CI.
|
Persistent review updated to latest commit 6982044 |
|
VERDICT: READY at What I got wrongI asked for an asymmetry arm: encode
So What the driver already pins, and it is more than I assumed
Still open, and it is the one that matters across two familiesThe The three original gaps remain correctly fixedload refusal taught to know the embedding family (rather than the check deleted), bidirectional attention as its own mode, and the sentence-transformers head ( |
|
The shape of this is right — refuse, don't approximate — and two things I would want answered before it lands. What is right, and it is the harder half of an embeddings route:
1. Nothing checks that an embedding's length matches the warm-up's. if (v.empty()) return fail(500, …);
data.push_back({{"object","embedding"},{"index",i},{"embedding",v}});An embedding whose length drifted — a forward bug, a mixed container, an id past the vocabulary — would be returned with a different dimension and stored by the client as if it belonged to the same space. One line makes it a named 500 instead: if (static_cast<int>(v.size()) != m.dim_)
return fail(500, "input[" + std::to_string(i) + "] produced " + std::to_string(v.size()) +
" dimensions, the model's are " + std::to_string(m.dim_));2. The tokenizer runs outside the lock that guards the forward. Could not run it: the route needs an |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 6982044 — this is the pattern the rest of us should be copying, and I have one observation about the CI number that is not a blocker.
Why it is strong, in order of how much it would have caught:
- The real-pack reference is
transformers' ownGemma3TextModel, not a re-derivation. 0.999990 host / 0.999093 on the NPU inside annpu-hold, with amdxdna clean before and after — quoting the ticket discipline in the measurement is the detail that makes the NPU number usable by someone else. - Mutants on the real pack, with the numbers: norm rule off → cos 0.09, sliding layers left causal → 0.70–0.90. That is the difference between "correctness was tested" and "correctness was shown to be able to fail" — most PRs I have read tonight stop at the first.
- The CI arm is hermetic and still independent: a two-layer container in the pack's own layout (BF16, unprefixed names, pre-scaled embedding, raw norms, one sliding and one full layer, IM off the 256 grid, plus the dense head) against a numpy Gemma3 encoder. It covers the four things that break this model class — the
model.prefix aliasing, thesqrt(H)pre-scale, the norm convention, and bidirectional attention — rather than the one that happened to break. - The refusal paths are arms too: a 9-token input refused by name, a 403-token input getting a 400 that names the cap.
The one observation. The CI arm reports cos 1.0000000 against the numpy reference, which is exact. Exactness is possible here (same weights, same float32 operations), but it also means the numpy encoder and embed_pooled() agree to the last bit — which is consistent with the reference reproducing the implementation's operation order (accumulation order, GELU form, pooling divisor) rather than only its mathematics. That is not a defect and it does not weaken the real-pack evidence, which is the number that counts. It does mean the hermetic arm would not catch a mathematically equivalent but differently ordered implementation, so I would describe it in the docstring as "holds the implementation, on the same weights, to a numpy encoder written to the same order" rather than as an independent check — the real-pack row is the independent one.
Nothing else holds. The model_alias_unprefixed_names guard ("only where the prefixed name is not already taken") is the right shape for a layout repair, and returning the count so the log can say how many were renamed is what makes it debuggable when a new pack variant shows up.
…d medium, and it bit #2183 tonight (#2331) `registration_loss_check` covers what a branch removes from the BUILD. The same lost update happens to documentation and nothing looked at it. Measured on engine#2183 at `407c0c5`: `docs/npu.md` is 3,914 lines on main and 3,358 on the branch, and the merge result is missing a section that had merged THAT MORNING in #2314 -- the rebase had resolved `docs/npu.md` by taking the branch's side. Three decisions keep this from crying wolf, each from a false positive seen today: * compare against the MERGE BASE, not against `--base` itself. A branch that forked before other lanes landed their sections shows them as missing when measured against today's main -- my first version reported 3 lost sections on #2183 where 1 had gone, the same artefact that flagged 38 innocent PRs earlier today. * the set is REPO-WIDE, not per file: a section moved between documents is not a loss. * a heading is identified by its TEXT, not its level: promoting `### Foo` to `## Foo` is an edit, not a loss. Measured on #2183's head, with the merge base fixed: rc 1 "A pack with both families: a 17-row chunk no longer pays for 128 rows (#1959)" <- in docs/npu.md, gone from the branch "Sentence embeddings from a q4nx encoder (Embedding-Gemma-300M-NPU2, engine#1786)" <- the one I found by hand 2 of 642 section(s) at 3e3e264 are missing **The second one is the tool earning its keep on its first run**: I had reported only the Embedding-Gemma section, because my hand check compared `###` headings and the other is a `####`. A reviewer who fixes one and leaves the other has done half the work. Four self-test arms: a removed section is a finding; a section MOVED to another document is not; a section whose LEVEL changed is not; and an unreadable base is rc 2, never a pass. `ctest -R doc_section_loss` -> Passed 0.06 s; `copyright --check` rc 0; `ci_targets_check --check` rc 0. Co-authored-by: lane <lane@local>
Refs #1786 (the
Embedding-Gemma-300M-NPU2row).What was missing. FastFlowLM's
Embedding-Gemma-300M-NPU2is aGemma3TextModelwithuse_bidirectional_attention: true, plus the sentence-transformers head (dense_2768→3072,dense_33072→768). The per-op forward stopped at load (no embed_tokens/lm_head.weight). Even past that, it encodes causally, one token at a time, and/v1/embeddingsserved only the ternary BERT container.Change.
npu/forward/src/model.c:model_alias_unprefixed_names()prefixes the pack's unprefixed names (layers.N.*,embed_tokens,norm) withmodel., only where that name is free.dequant_any(): an unquantised BF16[out, in]projection (this pack ships every weight that way) is copied into the forward's padded grid (IM 1152 → 1280), with zeros in the padding.q4nx_semantics:bidirectional, fromuse_bidirectional_attention, puts the forward in its existing non-causal mode and extends it to sliding layers. That equals the band a sliding layer was trained with only while the input fits half the window, so a longer input is refused by name.w: every norm here equalsunsloth/embeddinggemma-300m's to bf16 rounding. It runs as1 + w. FastFlowLM's quantised gemma3 packs, which store1 + w, keep the current rule.Q4nxNpuForward::embed_pooled(): one batched pass at position 0, with the LM head skipped; the final normed rows are mean-pooled, thendense_2,dense_3and L2.app/embed_serve: a q4nx encoder kind beside the ternary container, on the same flags and surface (--embedding -m <dir> --device npu|cpu). Ids are framed fromtokenizer_config.json'sadd_bos_token/add_eos_token(<bos> text <eos>), which gives the reference tokenizer's ids exactly. The input cap (256 tokens) is answered as a 400.docs/npu.md: a section beside "Sentence embeddings".Measured. The reference is transformers' own
Gemma3TextModelonunsloth/embeddinggemma-300m, with mean pooling and both dense modules as sentence-transformers applies them. Five sentences, the same ids, through1bit serve --embedding:--device cpu(host GEMM)--device npu(npu-hold, amdxdna clean before and after)The host similarity matrix is the reference's to three decimals.
usage.prompt_tokensis 83, the reference ids' total. A 403-token input gets a 400 naming the cap. Mutants on the real pack: with the norm rule off, cos 0.09; with the sliding layers left causal, 0.70–0.90.Test.
tests/embed_forward_ci.py(embed_forward, registered in the host-only CI configuration beside thebatch_check_statearms) writes a two-layer container in this pack's layout: BF16, unprefixed names, pre-scaled embedding, raw norms, one sliding and one full layer, IM off the 256 grid, plus the dense head. It holdsembed_pooled()to a numpy Gemma3 encoder on the same weights: cos 1.0000000. A 9-token input (window 16) must be refused by name. Mutants: sliding layers causal → cos 0.9905 FAIL; norm rule off → 0.2365 FAIL.On this box:
ctest -j 6without the device tests, 278/278; every forward arm by hand ("", lfm2, lfm3, qwen2, qwen2vl, qwen2_untied, qwen3_5, hy_v3, gptoss) rc 0.The kernel set for the device row is locally built and not shipped (4 designs,
QKV K768 N1536,O K768 N1024,GU K768 N2560,D K1280 N1024).🤖 Generated with Claude Code