bug: Embeddings from non-causal models silently diverge when the input is longer than the context's batch size #674

Issue description

With a non-causal embedding model (EmbeddingGemma), LlamaEmbeddingContext.getEmbeddingFor() returns a vector that does not match llama.cpp's own embedding of the same input once the input is longer than the context's batch size. No error or warning is raised.

The default embedding context is affected. createEmbeddingContext({contextSize: 2048}) gets batchSize 512, because the default batch size is Math.min(contextSize * sequences, 512). Inputs longer than the batch size but within the context size are accepted, and their embeddings diverge (see the table below).

This is unchanged on master as of 0ad4a86: the default is still in src/evaluator/LlamaContext/LlamaContext.ts (getDefaultContextBatchSize), and LlamaEmbeddingContext.getEmbeddingFor() only checks the input length against the context size.

Expected Behavior

For an input that fits in the context, getEmbeddingFor() should return the same embedding as llama.cpp's llama-embedding for that input. If the input cannot be evaluated correctly with the context's batch size, it should throw.

Actual Behavior

Below the batch size (210 and 430 tokens), the result matches llama-embedding (cosine 1.0000). Above it (650, 870 and 1090 tokens), the result diverges. Setting batchSize: 2048 on the same context makes all five inputs match.

Model ggml-org/embeddinggemma-300M-GGUF (embeddinggemma-300M-Q8_0.gguf), contextSize: 2048. The reference is llama-embedding -c 2048 -b 2048 -ub 2048 --pooling mean --embd-normalize 2 on the same GGUF and text.

My guess at the cause (I have not traced it through the code): the sequence is evaluated in several llama_decode calls of at most batchSize tokens. With non-causal attention, tokens in one call cannot attend to tokens in a later call, so the pooled vector is not the embedding of the whole input. llama.cpp's check that n_ubatch >= n_tokens for non-causal attention never fires, because each call on its own is within the batch size.

Causal models do not seem to be affected: running the same script with Qwen3-Embedding-0.6B-Q8_0.gguf (causal, pooling_type 3 / last) gives cosine 1.00000 between the default and batchSize: 2048 at every length in the table.

Steps to reproduce

// node repro.mjs <path to embeddinggemma-300M-Q8_0.gguf> [batchSize]
// Minimal repro: embeddings of inputs longer than the context's batch size, for a non-causal model.
import {getLlama} from "node-llama-cpp";
import {writeFileSync} from "node:fs";

const [modelPath, batchSizeArg] = process.argv.slice(2);
const llama = await getLlama();
const model = await llama.loadModel({modelPath});
const context = await model.createEmbeddingContext({
    contextSize: 2048,
    ...(batchSizeArg ? {batchSize: Number(batchSizeArg)} : {})
});
console.log("batchSize:", context._llamaContext.batchSize);

const words = ["river", "mountain", "library", "engine", "garden", "theory", "signal", "harvest", "lantern", "orbit"];
const makeText = (n) => Array.from({length: n}, (_, i) =>
    `Sentence ${i} mentions the ${words[i % 10]} and the ${words[(i * 7 + 3) % 10]}.`).join(" ");

const out = {};
for (const n of [20, 40, 60, 80, 100]) {
    const text = makeText(n);
    const tokens = model.tokenize(text).length;
    const {vector} = await context.getEmbeddingFor(text);
    out[n] = {tokens, text, vector: Array.from(vector)};
    console.log(`sentences=${n} tokens=${tokens}`);
}
writeFileSync(`out-${batchSizeArg ?? "default"}.json`, JSON.stringify(out));
process.exit(0);

1. node repro.mjs embeddinggemma-300M-Q8_0.gguf and node repro.mjs embeddinggemma-300M-Q8_0.gguf 2048
2. Write each text to a file and run llama-embedding -m embeddinggemma-300M-Q8_0.gguf -f <file> -c 2048 -b 2048 -ub 2048 --pooling mean --embd-normalize 2 --embd-output-format json
3. Compare the vectors with cosine similarity

My Environment

npx --yes node-llama-cpp inspect gpu output:

OS: macOS 25.6.0 (arm64)
Node: 24.13.1 (arm64)

node-llama-cpp: 3.22.1
Prebuilt binaries: v0.5.0

Metal: available

Metal device: Apple M1 Pro
Metal used VRAM: 0% (384KB/24.96GB)
Metal free VRAM: 99.99% (24.96GB/24.96GB)
Metal unified memory: 24.96GB (100%)

Reference: llama-embedding version 0.5.0 (build 11146, commit 7fe450e19), Homebrew.

Additional Context

• I also see this with EmbeddingGemma 2 (ggml-org/embeddinggemma-2-GGUF, Q8_0). That setup is different from the one above: node-llama-cpp 3.20.0 built from source against llama.cpp b11501, with two local build fixes. The first replaces common_batch_add, which current llama.cpp no longer provides; the open PR fix: preserve batch insertion compatibility #670 in this repository fixes the same build failure. The second uses llama_model_n_embd_out instead of llama_model_n_embd for the size of the returned vector, as 3.22.1 already does. Neither touches batching. For 40 inputs of more than 600 tokens, compared with the sentence-transformers fp32 reference (google/embeddinggemma-2), the median cosine was 0.843 with the default batch size and 0.9999 with batchSize: 2048. Ten inputs of 213 to 377 tokens matched the reference with the default batch size (cosine 0.9998 or higher).
• Downstream, this affects qmd, which creates its embedding context without batchSize and embeds chunks of up to ~900 tokens. qmd users reported the wrong vectors in issue Embedding inputs at/over the context's physical batch boundary (512 tokens) silently store deterministic garbage vectors tobi/qmd#897, and PR fix(embed): keep the whole input in one embedding batch (#897) tobi/qmd#993 fixes it on the qmd side by passing batchSize.
• Possible fixes: default batchSize to contextSize for embedding contexts of non-causal models, or throw in getEmbeddingFor() when the input is longer than batchSize for such models.

Relevant Features Used

• Metal support

Are you willing to resolve this issue by submitting a Pull Request?

No, I don't have the time and I'm okay to wait for the community / maintainers to resolve this issue.