Skip to content

Batch embedding for 10-20x speedup on short documents - #13

Closed
eiderman wants to merge 1 commit into
masterfrom
eider/batch-embedding
Closed

eiderman wants to merge 1 commit into
masterfrom
eider/batch-embedding

Conversation

@eiderman

Copy link
Copy Markdown
Collaborator

Embedding short documents one at a time wastes GPU/CPU cycles on per-call overhead. The T5 encoder can process multiple sequences in a single forward pass if they're padded to equal length.

Changes

  • Add embed_batch(): pads short texts to equal length and runs them through a single forward pass. Long texts (>2048 tokens) take the standard windowed stride path. embed() is now a thin wrapper.
  • Extract filter_and_normalize() to eliminate duplicated MIN_NORM filtering + L2 normalization.
  • Replace the single-doc Gatherer iterator in embed_chunks with a loop that collects rows into batches of 32 before embedding.

embed_batch() pads short texts to equal length and runs them through a
single forward pass. Long texts (>2048 tokens) take a windowed stride
path. embed() is now a thin wrapper. Shared filter_and_normalize()
eliminates the duplicated MIN_NORM + L2 normalization logic.

In embed_chunks, replace the single-doc Gatherer iterator with a loop
that collects rows into batches of 32 before embedding.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Comment thread src/quantized_t5.rs
let (scores, position_bias) = match position_bias {
Some(position_bias) => {
let scores = crate::fast_ops::fast_add(&scores, position_bias)?;
let scores = scores.broadcast_add(position_bias)?;

@jacobgorm jacobgorm Apr 22, 2026 •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am fairly sure our fast_add was indeed a lot faster than broadcast_add in candle

@jacobgorm

Copy link
Copy Markdown
Collaborator

I tried testing this change with:

make nfcorpus

But it hangs for me

@jacobgorm

Copy link
Copy Markdown
Collaborator

I've pushed fully self-contained nfcorpus scoring support to master

@eiderman eiderman closed this Apr 23, 2026
@eiderman
eiderman deleted the eider/batch-embedding branch April 23, 2026 16:33
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.

2 participants