From 17aa12477ab2c510df3b676520cb93071430b515 Mon Sep 17 00:00:00 2001 From: "Gareth.Y" <91742523+linjiaye929@users.noreply.github.com> Date: Sat, 31 Oct 2026 11:52:37 +0900 Subject: [PATCH 1/4] fix: preserve sequence offsets in unpooled batched embeddings --- llama_cpp/llama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 14e2f8500..af5e3b4d3 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -1131,7 +1131,7 @@ def decode_batch(seq_sizes: List[int]): internals.normalize_embedding(e) for e in embedding ] data.append(embedding) - pos += size + pos += size * n_embd else: for i in range(len(seq_sizes)): ptr = llama_cpp.llama_get_embeddings_seq(self._ctx.ctx, i) From 092cd89672b010b538d77b3ca5d45af16053c1af Mon Sep 17 00:00:00 2001 From: "Gareth.Y" <91742523+linjiaye929@users.noreply.github.com> Date: Sat, 31 Oct 2026 11:53:45 +0900 Subject: [PATCH 2/4] test: cover unpooled embedding offsets across batch boundaries --- tests/test_embedding_offsets.py | 59 +++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 tests/test_embedding_offsets.py diff --git a/tests/test_embedding_offsets.py b/tests/test_embedding_offsets.py new file mode 100644 index 000000000..c47e5431c --- /dev/null +++ b/tests/test_embedding_offsets.py @@ -0,0 +1,59 @@ +import ctypes +from types import SimpleNamespace + +import pytest + +from llama_cpp import llama as llama_module + + +@pytest.mark.parametrize("normalize", [False, True]) +@pytest.mark.parametrize("n_batch,n_seq_max", [(32, 4), (4, 4), (32, 1)]) +def test_unpooled_embedding_offsets(monkeypatch, normalize, n_batch, n_seq_max): + """Each sequence must read its own token rows from the flat native buffer.""" + sequences = [] + buffers = [] + tokens_by_text = {b"first": [1, 2], b"second": [3, 4], b"third": [5]} + n_embd = 3 + + def add_sequence(tokens, seq_id, logits_all): + assert seq_id == len(sequences) + assert logits_all + sequences.append(tokens) + + def decode(batch): + values = [ + float(token * n_embd + dim) + for tokens in sequences + for token in tokens + for dim in range(n_embd) + ] + buffers[:] = [(ctypes.c_float * len(values))(*values)] + + model = SimpleNamespace( + n_embd=lambda: n_embd, + n_batch=n_batch, + context_params=SimpleNamespace(n_seq_max=n_seq_max, embeddings=True), + pooling_type=lambda: llama_module.llama_cpp.LLAMA_POOLING_TYPE_NONE, + verbose=False, + tokenize=lambda text: tokens_by_text[text], + _batch=SimpleNamespace(reset=sequences.clear, add_sequence=add_sequence), + _ctx=SimpleNamespace(ctx=None, kv_cache_clear=lambda: None, decode=decode), + reset=lambda: None, + ) + monkeypatch.setattr( + llama_module.llama_cpp, + "llama_get_embeddings", + lambda ctx: ctypes.cast(buffers[0], ctypes.POINTER(ctypes.c_float)), + ) + output, count = llama_module.Llama.embed( + model, ["first", "second", "third"], normalize=normalize, return_count=True + ) + assert count == 5 + assert [len(sequence) for sequence in output] == [2, 2, 1] + for sequence, tokens in zip(output, tokens_by_text.values()): + for row, token in zip(sequence, tokens): + expected = [float(token * n_embd + dim) for dim in range(n_embd)] + if normalize: + norm = sum(value * value for value in expected) ** 0.5 + expected = [value / norm for value in expected] + assert row == pytest.approx(expected) From 5ffc528fe095ed6054e5a2b450b9e3fe82a16bfa Mon Sep 17 00:00:00 2001 From: "Gareth.Y" <91742523+linjiaye929@users.noreply.github.com> Date: Sat, 31 Oct 2026 11:54:55 +0900 Subject: [PATCH 3/4] docs: record unpooled batched embedding fix --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..c59cd130f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +- fix: preserve sequence offsets in unpooled batched embeddings by @linjiaye929 + ## [0.3.36] - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098 From 08fefca91d55d7065cbd158db046467d9630d1e4 Mon Sep 17 00:00:00 2001 From: "Gareth.Y" <91742523+linjiaye929@users.noreply.github.com> Date: Sat, 31 Oct 2026 12:32:50 +0900 Subject: [PATCH 4/4] docs: link embedding offset fix to upstream PR #2395 --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c59cd130f..7ad10819a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -- fix: preserve sequence offsets in unpooled batched embeddings by @linjiaye929 +- fix: preserve sequence offsets in unpooled batched embeddings by @linjiaye929 in #2395 ## [0.3.36]