Skip to content

GH-1392: Add padding strategies for SentenceVectorsDL.embedAll and pad on a GPU by default - #1474

Draft
krickert wants to merge 1 commit into
feat/GH-1391-embed-batch-capfrom
feat/GH-1392-embed-padding
Draft

krickert wants to merge 1 commit into
feat/GH-1391-embed-batch-capfrom
feat/GH-1392-embed-padding

Conversation

@krickert

@krickert krickert commented Oct 11, 2026 •

Copy link
Copy Markdown
Contributor

Draft until #1468 merges; this branch is stacked on it.

I've tried all 3 of these strategies before, and a embeddings fine-tuner this is one aspect to get a max vertical scaling. As you can see from below, you can save a lot of compute time and energy consumption when digging this deep.

SentenceVectorsDL.embedAll groups texts by exact tokenized length. On natural text almost every length differs, so most inferences hold one or two texts, and on a GPU a call of mixed lengths is no faster than calling embed in a loop.

PaddingStrategy sets how embedAll batches texts of different lengths:

  • EXACT_LENGTH does not pad and runs one inference per tokenized length. It is the default on the CPU.
  • LONGEST orders texts by length and pads each batch to its longest text. It is the default on a GPU when the vocabulary has a padding token.
  • MAX_LENGTH pads every text to the maximum length, for models or accelerators that need a fixed shape.

A new constructor takes the strategy, and the onnx provider takes a padding option. The padding id is read from the vocabulary ([PAD], or for RoBERTa). Padded positions are 0 in the attention mask and mean pooling skips them, so the vectors match up to floating-point rounding when the model honors the mask. Every strategy stays within 16384 token positions per inference. The machine learning chapter documents the option and the GPU default.

SentenceVectorsDLBenchmark on all-MiniLM-L6-v2, RTX 4080 SUPER, texts of 4 to 40 words, embedAll time per call with EXACT_LENGTH and with the GPU default:

Texts EXACT_LENGTH LONGEST
8 3603 us 862 us
32 11431 us 2152 us
128 23815 us 7443 us

Tests: batching under each strategy, the token position cap with padding, the padding id from the vocabulary, the device default, and SentenceVectorsDLPaddingEdgesTest for the cap boundaries under every strategy, input order with interleaved lengths, a RoBERTa vocabulary, padding id 0, CLS pooling and unit length scaling, the pooled-output model, single embed under MAX_LENGTH, and CUDA against the CPU (skipped without a CUDA provider or device). tiny-masked.onnx, generated by gen_tiny_masked_model.py next to it, sums the ids times the attention mask, so a test can tell masked padding from unmasked.

Fixes #1392

@krickert
krickert force-pushed the feat/GH-1391-embed-batch-cap branch from dc68fd7 to f2a2206 Compare October 11, 2026 22:39
@krickert
krickert force-pushed the feat/GH-1392-embed-padding branch from cc4df6c to 3d639ac Compare October 11, 2026 22:40
@krickert
krickert force-pushed the feat/GH-1391-embed-batch-cap branch from f2a2206 to e540ab2 Compare October 12, 2026 01:20
…d on a GPU by default

embedAll grouped texts by exact tokenized length, so on natural text
most inferences took one or two texts. PaddingStrategy.LONGEST orders
texts by length and pads each batch to its longest member, MAX_LENGTH
pads every input to the maximum length, and EXACT_LENGTH behaves as
before. Under every strategy no inference takes more than 16384 token
positions, except a single row wider than that, which runs alone. The
onnx provider gets a padding option.

A session built without a strategy uses LONGEST on a GPU when the
vocabulary has a padding token ([PAD], or <pad> for RoBERTa), and
EXACT_LENGTH otherwise; the CPU default is unchanged. Padded positions
are 0 in the attention mask and mean pooling skips them, so vectors
match up to floating-point rounding when the model honors the mask.

tiny-masked.onnx, generated by gen_tiny_masked_model.py next to it, sums
the ids times the attention mask so tests can tell masked padding from
unmasked.

SentenceVectorsDLBenchmark on all-MiniLM-L6-v2, RTX 4080 SUPER, texts of
4 to 40 words, embedAll time per call before and after:

  8 texts:     3603 us ->  862 us
  32 texts:   11431 us -> 2152 us
  128 texts:  23815 us -> 7443 us

Fixes #1392

Signed-off-by: Kristian Rickert <kristian@apache.org>
Co-authored-by: Claude Code 2.1.296 (Opus 5.5)
@krickert
krickert force-pushed the feat/GH-1392-embed-padding branch from 3d639ac to bc0c2b1 Compare October 12, 2026 01:21
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