Skip to content

fix: preserve sequence offsets in unpooled batched embeddings - #2395

Open
linjiaye929 wants to merge 4 commits into
abetlen:mainfrom
linjiaye929:codex/unpooled-embedding-offsets
Open

linjiaye929 wants to merge 4 commits into
abetlen:mainfrom
linjiaye929:codex/unpooled-embedding-offsets

Conversation

@linjiaye929

Copy link
Copy Markdown

Problem

With LLAMA_POOLING_TYPE_NONE, Llama.embed() advances the flat output pointer by a token count after each sequence, although the pointer indexes floats. When multiple sequences share a decode batch and the embedding dimension is greater than one, later sequences read the wrong rows.

The pinned native API contract specifies contiguous output in batch order with shape n_outputs * n_embd.
The Python slicing code therefore needs to advance by size * n_embd floats, rather than size.

Change

  • Correct that one offset increment.
  • Add deterministic regression coverage for multiple sequences, batch rollover, single-sequence batches, and normalized output.
  • Add an Unreleased changelog entry.

Pooled embeddings and the public API are unchanged.

Validation

Validated against Python commit 1652066e0af45f2313b339670ef9555e8a54e545 and its unmodified native dependency 0c1e57098bba43ac29e6e3b677cdceebdd22334f on Linux x86-64, Python 3.12.14, GCC 14, CPU-only, with two parallel build jobs.
The native library was built from source and imported normally.

Commands were run in an isolated environment with HF_HUB_OFFLINE=1 and HF_HUB_DISABLE_TELEMETRY=1:

  • python -m pytest tests/test_embedding_offsets.py -q: baseline 4 failed, 2 passed; corrected code 6 passed.
  • python -m pytest tests/test_embedding_offsets.py tests/test_llama_grammar.py tests/test_llama_chat_format.py tests/test_llama_speculative.py tests/test_llama.py::test_llama_cpp_version -q: 13 passed.
  • make lint: passed, including formatting checks for all 30 Python/test files.

The regression uses the real Python method with a deterministic ctypes float buffer and mocked decode/output-getter boundaries. It does not exercise model inference.
Full make test was not run because its model-download fixtures and vocabulary GGUF asset were unavailable within this validation scope. No model assets were downloaded, and GPU/backend-specific behavior was not tested.

AI assistance

AI assistance was used for investigation, the patch, regression tests, and this description. The results above distinguish the automated checks actually run from model-dependent coverage that remains unverified.

@linjiaye929
linjiaye929 marked this pull request as ready for review October 31, 2026 03:37

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant