Skip to content

npu: Embedding-Gemma-300M-NPU2 serves /v1/embeddings on the per-op forward (refs #1786) - #2314

Merged
bong-water-water-bong merged 3 commits into
mainfrom
cl-flm/1786-embeddinggemma
Nov 13, 2026
Merged

bong-water-water-bong merged 3 commits into
mainfrom
cl-flm/1786-embeddinggemma

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

Refs #1786 (the Embedding-Gemma-300M-NPU2 row).

What was missing. FastFlowLM's Embedding-Gemma-300M-NPU2 is a Gemma3TextModel with use_bidirectional_attention: true, plus the sentence-transformers head (dense_2 768→3072, dense_3 3072→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/embeddings served 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) with model., 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, from use_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.
    • A gemma3 pack whose projections are BF16 is an HF copy whose norms are raw w: every norm here equals unsloth/embeddinggemma-300m's to bf16 rounding. It runs as 1 + w. FastFlowLM's quantised gemma3 packs, which store 1 + 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, then dense_2, dense_3 and 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 from tokenizer_config.json's add_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 Gemma3TextModel on unsloth/embeddinggemma-300m, with mean pooling and both dense modules as sentence-transformers applies them. Five sentences, the same ids, through 1bit serve --embedding:

route min cos to the reference 5 inputs
--device cpu (host GEMM) 0.999990
--device npu (npu-hold, amdxdna clean before and after) 0.999093 0.88 s

The host similarity matrix is the reference's to three decimals. usage.prompt_tokens is 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 the batch_check_state arms) 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 holds embed_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 6 without 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

…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>
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

ready for review at 3985fd9

@github-actions

github-actions Bot commented Nov 12, 2026 •

Copy link
Copy Markdown

PR Reviewer Guide 🔍

(Review updated until commit 6982044)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶

1786 - Partially compliant

Compliant requirements:

  • The Embedding-Gemma-300M-NPU2 model source imports with zero gaps and its self-check passes against a named reference
  • The Embedding-Gemma-300M-NPU2 model source serves on --device npu
  • The Embedding-Gemma-300M-NPU2 model source answers a known-answer prompt with the model's own chat template, 3 identical requests returning identical text
  • The Embedding-Gemma-300M-NPU2 model source decode tok/s and time-to-first-token measured under npu-hold

Non-compliant requirements:

  • The row names the kernel set and the engine commit that produced it

Requires further human verification:

  • The row names the kernel set and the engine commit that produced it
⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪
🧪 PR contains tests
🔒 No security concerns identified
⚡ Recommended focus areas for review

Incorrect Handling of BF16 Projections

The dequant_any function in Q4nxNpuForward assumes that BF16 projections are always properly aligned and sized. If the data is not aligned or the size is incorrect, it could lead to memory corruption or incorrect results, especially when dealing with padded grids.

static int dequant_any(const void* data, size_t nbytes, int rows, int cols, float* out, const TensorDesc* desc = nullptr) {
    if (!data || !out || rows <= 0 || cols <= 0) return -1;
    // engine#1786: an unquantised projection (FastFlowLM's Embedding-Gemma-300M-NPU2 ships every weight as BF16
    // [out, in]). Copied into the forward's grid, which may be padded past the declared shape (IM 1152 -> 1280): the
    // pad rows and columns are zero, the layout every chunked reader below writes.
    if (desc && desc->ndim == 2 && !strcmp(desc->dtype, "BF16")) {
        const int64_t r0 = desc->shape[0], c0 = desc->shape[1];
        if (r0 <= 0 || c0 <= 0 || r0 > rows || c0 > cols || nbytes < (size_t)r0 * (size_t)c0 * 2) return -1;
        const uint16_t* w = (const uint16_t*)data;
        for (int r = 0; r < rows; r++)
            for (int c = 0; c < cols; c++)
                out[(size_t)r * cols + c] = (r < r0 && c < c0) ? bf16_to_f32(w[(size_t)r * (size_t)c0 + (size_t)c]) : 0.0f;
        return 0;
    }
Potential Memory Issue in embed_pooled

The embed_pooled function in Q4nxNpuForward allocates memory for last_hidden_rows_ and fnorm vectors without checking for potential overflow or excessive memory usage, especially when handling large inputs. This could lead to memory exhaustion or performance degradation.

bool Q4nxNpuForward::embed_pooled(const std::vector<int>& ids, std::vector<float>& out) {
    out.clear();
    if (!noncausal_attn_) {
        err_ = "embed_pooled(): this model is causal (no use_bidirectional_attention, no noncausal plan); a sentence "
               "embedding pools an encoder's rows";
        return false;
    }
    if (ids.empty()) { err_ = "embed_pooled(): empty input"; return false; }
    const int H = c_.H;
    // The sentence-transformers head, as stored: dense_2 [D2, H] then dense_3 [D3, D2], both linear (identity
    // activation). Read once, f32; a pack without them pools to H.
    auto load_dense = [&](const char* name, int in_w, std::vector<float>& w, int& rows) -> bool {
        if (!w.empty()) return true;
        TensorDesc* d = model_tensor_by_name(mw_, name);
        if (!d) { rows = 0; return true; }
        if (d->ndim != 2 || d->shape[1] != in_w || strcmp(d->dtype, "BF16") != 0) {
            err_ = std::string("embed_pooled(): ") + name + " is not a BF16 [out, " + std::to_string(in_w) + "] matrix";
            return false;
        }
        rows = (int)d->shape[0];
        const uint16_t* raw = (const uint16_t*)model_tensor_data(mw_, d);
        if (!raw || (size_t)d->data_size < (size_t)rows * (size_t)in_w * 2) { err_ = std::string("embed_pooled(): ") + name + " is short"; return false; }
        w.resize((size_t)rows * (size_t)in_w);
        for (size_t i = 0; i < w.size(); i++) w[i] = bf16_to_f32(raw[i]);
        return true;
    };
    if (!load_dense("dense_2.weight", H, dense2_, dense2_out_)) return false;
    if (dense2_out_ > 0 && !load_dense("dense_3.weight", dense2_out_, dense3_, dense3_out_)) return false;
    std::vector<float> unused;
    hidden_only_ = true;
    const bool ok = step_batch(ids, 0, unused);
    hidden_only_ = false;
    if (!ok) return false;
    const int M = (int)ids.size();
    if (last_hidden_rows_.size() != (size_t)M * (size_t)H) { err_ = "embed_pooled(): row count"; return false; }
    std::vector<double> mean((size_t)H, 0.0);
    for (int m = 0; m < M; m++)
        for (int k = 0; k < H; k++) mean[(size_t)k] += last_hidden_rows_[(size_t)m * H + k];
    std::vector<float> v((size_t)H);
    for (int k = 0; k < H; k++) v[(size_t)k] = (float)(mean[(size_t)k] / M);
    auto linear = [](const std::vector<float>& w, int rows, int in_w, const std::vector<float>& x) {
        std::vector<float> y((size_t)rows);
        for (int r = 0; r < rows; r++) {
            double s = 0.0;
            const float* wr = &w[(size_t)r * (size_t)in_w];
            for (int k = 0; k < in_w; k++) s += (double)wr[k] * x[(size_t)k];
            y[(size_t)r] = (float)s;
        }
        return y;
    };
    if (dense2_out_ > 0) v = linear(dense2_, dense2_out_, H, v);
    if (dense3_out_ > 0) v = linear(dense3_, dense3_out_, dense2_out_, v);
    double n2 = 0.0;
    for (float x : v) n2 += (double)x * x;
    const double inv = n2 > 0 ? 1.0 / std::sqrt(n2) : 0.0;
    for (float& x : v) x = (float)(x * inv);
    out = std::move(v);
    return true;
}
Incorrect Bidirectional Attention Handling

The handling of bidirectional attention in step_batch_impl may not correctly account for the sliding window constraint when bidir_sliding_ is true. The check for input length against sem_.sliding_window / 2 might not be sufficient to prevent incorrect behavior for inputs that are just under the limit but still cause issues.

if (bidir_sliding_ && sem_.sliding_window > 0 && M > sem_.sliding_window / 2) {
    err_ = "use_bidirectional_attention: a " + std::to_string(M) + "-token input is longer than half the " +
           std::to_string(sem_.sliding_window) + "-token sliding window, where a sliding layer's band is no " +
           "longer the whole prompt; send at most " + std::to_string(sem_.sliding_window / 2) + " tokens";
    return false;
}
✅ Resolved findings

npu/forward/src/q4nx_forward.cpp:4417-4435

Incorrect Handling of Sentence Transformers Head

The embed_pooled function in Q4nxNpuForward assumes that the sentence-transformers head (dense_2 and dense_3) is always present and correctly formatted. If these tensors are missing or incorrectly formatted, the function may fail or produce incorrect results.

app/embed_serve.cpp:244-252

Potential Race Condition

The embeddings function in Backend uses a mutex (m.mu) to protect access to m.fwd->embed_pooled, but it does not ensure that the embed_pooled function itself is thread-safe or that all operations within it are properly synchronized, which could lead to race conditions.

npu/forward/src/q4nx_forward.cpp:4443-4447

Incorrect Handling of Norms

The embed_pooled function in Q4nxNpuForward assumes that the norms are correctly formatted and that the rmsnorm_offset is correctly set. If these assumptions are incorrect, the function may produce incorrect results.

npu/forward/src/q4nx_forward.cpp:200-213

Incorrect Handling of BF16 Projections

The dequant_any function in Q4nxNpuForward handles BF16 projections, but it assumes that the data is always aligned and properly sized. If the data is not aligned or the size is incorrect, it could lead to memory corruption or incorrect results.

npu/forward/src/q4nx_forward.cpp:4408-4466

Potential Memory Issue

The embed_pooled function in Q4nxNpuForward allocates memory for last_hidden_rows_ and fnorm vectors without checking for potential overflow or excessive memory usage, especially when handling large inputs. This could lead to memory exhaustion or performance degradation.

Comment thread tests/embed_pooled_test.cpp Fixed
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 3985fd9 — an embedding model needs three things a causal LM does not, and the PR names all three before any of the code.

The three gaps, in the order that matters

"a Gemma3TextModel with use_bidirectional_attention: true, plus the sentence-transformers head (dense_2 768→3072, dense_3 3072→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/embeddings served only the ternary BERT container."

  1. a load refusal — an embedding model has no lm_head.weight, so a forward that requires one refuses a valid pack. That is the right refusal (by name) but the wrong requirement, and the distinction matters: the fix is not to relax the check but to know that the family has no head;
  2. bidirectional attention — an encoder without a causal mask, which is a different attention path, not a flag on the same one. A causal encoder on an embedding model produces plausible vectors that are quietly wrong for the task (each token sees only its prefix);
  3. the pooling/head — dense_2/dense_3 with a 768→3072→768 sandwich, which is sentence-transformers' convention and nothing a causal LM has.

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

"/v1/embeddings served only the ternary BERT container" — so this is the second embedding family, and the first with a transformer head. The risk in a second implementation is that the two disagree about a detail the client depends on: normalisation (L2 vs not), the pooling rule (CLS vs mean), and the vector's declared dimension. Two endpoints serving different conventions under one route is the kind of thing a downstream index discovers after ingesting a million vectors.

Two asks

  1. state the pooling and normalisation explicitly — which token(s), mean or CLS, L2-normalised or not, and whether that matches what the BERT container does. One table, family | pooling | normalisation | dimension, is the only way a client can trust the route across two models — and it is the same discipline as npu: the per-op forward reads gemma3 text decoders; Medgemma-4B answers on the host (Refs #1786) #2242's three config facts and npu: the per-op forward reads LFM3, so LFM3-1.7B-NPU2 loads and answers (Refs #1786) #2254's four rules;
  2. give the asymmetry test — a bidirectional encoder must have emb("the cat sat") != emb("sat cat the") and a token's vector must change when the sentence continues past it. The second is the one that catches a causal mask left on: with a causal mask, the first token's vector is identical for "a" and "a b c", so it is a one-line assertion that fails exactly when the mask is wrong.

And the load refusal is worth keeping visible

"stopped at load (no embed_tokens/lm_head.weight)" — a refusal that names the missing tensor is what made this a day's work instead of a week's; the fix should add an embedding family to the loader's known shapes, not delete the check, so the next headless model still gets a named refusal rather than passing into a forward that assumes a head.

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

  • bidirectional extends 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 from false.
  • The fixture is not a stub: unprefixed names, embed_tokens pre-scaled by sqrt(H), norms as HF's raw w, 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_tokens is 83 and equals the reference ids' total, so the tokeniser path is pinned to the same numbers as the embeddings.
  • ctest -j 6 without device tests 278/278, every forward arm by hand, and the NPU row run inside an npu-hold with 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 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 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>
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

CI at 3985fd9 failed one test, file_create_check: tests/embed_pooled_test.cpp:52 used a plain fopen (0666, follows a symlink), which CodeQL flagged too. It now writes through onebit_fopen_private_wb. The same CI run passed embed_forward (265/283). This head passes file_create_check and ci_targets locally, and embed_forward is still cos 1.0000000.

ready for review at b76acb6

@github-actions

Copy link
Copy Markdown

Persistent review updated to latest commit b76acb6

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at e538cd681 — the pooling arithmetic and the refusal paths are both right, and the docs describe what the code does.

Checked in order of how easy each would have been to get subtly wrong:

  1. Mean pooling is unmasked, and that is correct here — mean[k] += last_hidden_rows_[m*H+k] over M rows then /= M, with ids
    being exactly the request's tokens. There is no padded batch to mask. Worth one comment at the accumulation, because the moment a
    batched path reuses embed_pooled the missing mask becomes a silent wrong vector rather than a compile error.
  2. The causal refusal is by name and by plan: embed_pooled() errors with "this model is causal (no use_bidirectional_attention,
    no noncausal plan)" instead of returning a causal-averaged vector, and the caller separately refuses when config.json says
    bidirectional but the forward is causal. Two independent checks for the same property is the right amount for a path whose whole
    premise is bidirectionality.
  3. The head is applied in the order sentence-transformers applies it: mean after the final norm, then dense_2 and dense_3
    (linear, no bias), then L2 normalisation (n2 sum, 1/sqrt, scale). A consumer doing cosine similarity downstream gets a unit
    vector, which is what /v1/embeddings clients assume.
  4. The cap is a stated 400, not a silent truncation: input over 256 tokens is refused naming the cap, well under the 512 window.
  5. The verification is against numpy, not against itself: the CI case runs the bidirectional gemma3 layout (BF16 projections,
    unprefixed names, pre-scaled embeddings, raw norms, one sliding and one full layer) through embed_pooled() and compares with a
    reference. The body also records what the rules buy: cosine 0.09 to 0.28 with them, 0.09 with the norm rule off.

One note, not a blocker: the 1 + w norm offset is chosen per pack ("FastFlowLM's quantised gemma3 packs store 1 + w; this one is a
copy of the HF checkpoint"), and a non-sentence-transformers gemma3 BF16 pack would get the wrong offset silently. The load line that
prints rmsnorm_offset=1 <- gemma3 is the mitigation, and it is the right one — I would only add that phrase to docs/npu.md where the
pack rules are listed, so it is discoverable without reading a log.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at b76acb6 — carried from 3985fd9; and having read the arms, the asymmetry check I asked for is one assertion in a file that already exists.

Where the arms actually live

The test surface here is tests/embed_forward_ci.py — a Python driver that builds a synthetic container from real pack names ("names with no model. prefix (layers.N.*, embed_tokens, norm, dense_2, dense_3)"), not a set of C++ REQUIREs. Two consequences worth stating:

  • the asymmetry / prefix assertion I asked for belongs in that script and is a few lines: encode "a" and "a b c", compare the first token's vector. With bidirectional attention it must differ; with a causal mask left on it is identical, because the first position sees only itself in both cases. That single comparison is the whole test for the mask, and a smoke test of "the endpoint returns 768 floats" cannot substitute;
  • "(1 + w) RMSNorm everywhere, q/k norms, rotate-half RoPE on the whole head (base 10000 …)" — the driver already pins the numeric recipe, so the mask is the remaining unpinned rule, and it sits naturally beside them.

Still open from my last review (unchanged)

  • a family | pooling | normalisation | dimension table — pooling appears once in the body, normalis/asymmetr not at all; with two embedding families on one route (this and the ternary BERT container), a client that ingests vectors from both needs the conventions stated, not inferred;
  • the exit-contract detail the driver already has (77 when the container is not set, stdout for a requested refusal) is exactly right and worth keeping in the body — it is the SKIP-vs-FAIL discipline applied to a harness.

Why the PR is sound as it stands

The 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 (dense_2/dense_3). Nothing here is a workaround, and the refusal "no embed_tokens/lm_head.weight" survives as a named diagnostic for the next headless model.

@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 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 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 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-npu2 stays rmsnorm_offset=0 and passes.
    • gemma4-e2b-npu2 stays rmsnorm_offset=0.
    • Only Embedding-Gemma-300M-NPU2 switches to rmsnorm_offset=1 <- gemma3 with BF16 projections.
  • gemma4-e2b fails 3 "holds exactly N positions" checks, and it fails them identically on main at 9a589d7, 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_2 then dense_3, with no activation and no bias, then L2. A pack without the dense head pools to H.
  • The refusal at M > window/2 is 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.

@github-actions

Copy link
Copy Markdown

Persistent review updated to latest commit 6982044

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 6982044 — and reading the driver corrected my own ask, which I would rather say than repeat.

What I got wrong

I asked for an asymmetry arm: encode "a" and "a b c", the first token's vector must differ. The driver says:

"the sliding layers left causal: 0.70 to 0.90" — and "a second request one token past half the window must be refused"

So use_bidirectional_attention applies to the full layers; the sliding layers stay causal by design, which is Gemma3's architecture, not an oversight. My test would therefore fail correctly on a sliding layer — and a lane "fixing" the failure would break the architecture. The arm has to be scoped: the first token's vector must change across prefixes on a full (bidirectional) layer, and must not on a sliding layer. That is a better test than the one I proposed, and I would not have known it without reading the driver.

What the driver already pins, and it is more than I assumed

Still open, and it is the one that matters across two families

The family | pooling | normalisation | dimension table. pooling appears once in the body; normalis and asymmetr not at all. With /v1/embeddings now serving both this model and the ternary BERT container, a client that ingests from both needs the conventions stated. That is not a stylistic ask: L2-normalised vs not changes cosine and dot-product rankings, and mean-pooled vs CLS changes what "similar" means — two families disagreeing under one route is discovered after a million vectors are indexed.

The three original gaps remain correctly fixed

load refusal taught to know the embedding family (rather than the check deleted), bidirectional attention as its own mode, and the sentence-transformers head (dense_2 768→3072, dense_3 3072→768). Nothing in the correction above changes any of that.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

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:

  • The cap is checked after tokenisation and before the forward, and the 400 names both numbers: input[i] is N tokens; this encoder takes at most M. A silently truncated embedding would be a plausible vector for the wrong text — the worst kind of wrong for anything downstream that stores it.
  • if (!m.fwd->noncausal()) refuses when config.json declares use_bidirectional_attention but the forward is causal: an encoder whose attention is the wrong direction still returns unit-length vectors, so nothing downstream would ever notice. That check is the difference between serving an encoder and serving noise-with-the-right-shape.
  • The warm-up runs before the route answers, and its failure is a startup refusal (encoder warm-up failed: …), not a first-request 500.
  • The forward call is behind lock_guard, and a failed embed_pooled reports the forward's own last_error() rather than "internal error".

1. Nothing checks that an embedding's length matches the warm-up's. dim_ is set once from the probe vector, and per-input vectors are pushed into the response as they come:

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. m.tok->encode(...) is called before lock_guard g(m.mu), so two concurrent requests encode at the same time while sharing one tokenizer. If that tokenizer is thread-safe, fine — but then the reason the lock exists is worth one word (// the forward is stateful; the tokenizer is not), and if it is not thread-safe this is a data race that will show up as a mangled input under load rather than as an error.

Could not run it: the route needs an ONEBIT_NPU build and the Embedding-Gemma-300M-NPU2 pack, neither of which is available in this lane's tree (and disk here is under the new-configure threshold), so this is a reading of the request path and its refusals rather than of a served response. The numpy-reference arm (embed_pooled_test against a Gemma3 encoder, with the off-grid intermediate 384→512) is the right arm to have; it is the one I would run first when the build is affordable.

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

  1. The real-pack reference is transformers' own Gemma3TextModel, not a re-derivation. 0.999990 host / 0.999093 on the NPU inside an npu-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.
  2. 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.
  3. 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, the sqrt(H) pre-scale, the norm convention, and bidirectional attention — rather than the one that happened to break.
  4. 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.

@bong-water-water-bong
bong-water-water-bong merged commit 3230a4b into main Nov 13, 2026
10 checks passed
@bong-water-water-bong
bong-water-water-bong deleted the cl-flm/1786-embeddinggemma branch November 13, 2026 00:32
bong-water-water-bong added a commit that referenced this pull request Nov 13, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants