Skip to content

fix(embed): respect physical batch limits for non-causal embeddings - #2387

Open
original4422 wants to merge 2 commits into
abetlen:mainfrom
original4422:fix/noncausal-embedding-microbatches
Open

original4422 wants to merge 2 commits into
abetlen:mainfrom
original4422:fix/noncausal-embedding-microbatches

Conversation

@original4422

Copy link
Copy Markdown

Summary

Fixes #1762. A 513-token BGE-M3 input with n_batch=513 currently reaches a native assertion because the omitted n_ubatch stays at 512.

omitted n_ubatch: embedding -> n_batch; generation -> 512
explicit n_ubatch: retain the requested limit
non-causal/encoder embedding: pack within the native physical batch
single sequence exceeding an explicit physical limit: raise ValueError

Resolve defaults before native context creation, including n_ctx=0 and logical batch clamping. Keep logical truncation unchanged, preserve causal microbatch splitting, and cache the fixed long-context model for the native CI regressions. No vendor changes.

Evidence

  • Before: original 0.3.36 / vendor 0c1e57098, real BGE-M3 F16, actual non-causal context: one 513-token sequence and two sequences totaling 256+257 tokens both terminate isolated processes with SIGABRT (encoder requires n_ubatch >= n_tokens). The initial regression set has 5 failures and 3 passing controls.
  • After: complete python -m pytest --full-trace -v: 41 passed. This includes real long embeddings, aggregate batches, explicit limits/zero, n_ctx=0, clamping, causal generation/embedding controls, and the existing model tests. Long embeddings retain every token and have finite output values.
  • python -m ruff check llama_cpp tests and python -m ruff format --check llama_cpp tests: pass; git diff --check: pass.

Local environment: macOS arm64, Python 3.12.14, original native build with Metal available; the new regression cases explicitly use CPU model layers (n_gpu_layers=0). BGE-M3 F16 revision: 9379ce497e8814b200f2dc0d18eb4045426dcb8c.

Merge Danger

Door: two-way

Blast Radius: embeddings

Omitted embedding microbatches now follow the logical batch size. Explicit physical limits remain supported; a sequence that cannot fit raises a Python error instead of aborting the process. Ordinary generation retains its 512-token default microbatch.

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.

Llama.embed crashes when n_batch > 512

1 participant