Skip to content

Keep the tail when truncating long states for the encoder - #6

Open
ssrajadh wants to merge 1 commit into
Contrastive-LM:mainfrom
ssrajadh:fix-truncation-side
Open

ssrajadh wants to merge 1 commit into
Contrastive-LM:mainfrom
ssrajadh:fix-truncation-side

Conversation

@ssrajadh

Copy link
Copy Markdown

Embedder sends truncate_prompt_tokens but not which side to cut, so vLLM falls back to the tokenizer default. For the pooling runner that default became right in vLLM 0.13 (vllm/tokenizers/registry.py). Through 0.12, every runner truncated from the left. On a current vLLM (#3 is on 0.30.0), a state over 2,048 tokens keeps its first 2,048 tokens. state_text puts the question last, so the question is what gets dropped. Training cuts the other way: Recipe.state_ids in train/embed_utils.py keeps the tail.

This PR sends truncation_side: "left" whenever truncate_prompt_tokens is set. The side names what gets cut, so left drops tokens from the start and keeps the last 2,048, question included. The change is two lines in Embedder._fetch.

vLLM accepts truncation_side on embedding requests from 0.18. Older versions allow unknown request fields, so they ignore it instead of rejecting the request, which means 0.13 to 0.17 still cut from the right. If you want that gap closed too, I can add a vllm>=0.18 pin to the vllm extra in this PR.

Testing

I don't have an NVIDIA GPU, so I tested against a Qwen3-8B encoder on MLX that follows vLLM's truncate_prompt_tokens / truncation_side semantics and matches the output reported in #3. The table shows cosine similarity between Embedder.embed() and the same state cut to its last 2,048 tokens:

State Before After
2,762 tokens -0.0006 1.0000
2,757 tokens -0.0167 1.0000
17 tokens 1.0000 1.0000
12 tokens 1.0000 1.0000

The two long states come from DecisionBench medium, where 139 of the 293 question states are over 2,048 tokens. Shorter states don't change. Neither do the DeepSWE and Terminal-Bench results, because bon_eval.py scores precomputed embeddings. The MLX encoder in #5 always keeps the tail, so with this change the two backends agree on long states.

Embedder sends truncate_prompt_tokens without a side, so vLLM falls back to
the tokenizer default. For the pooling runner that is 'right' since vLLM 0.13,
which keeps the first N tokens and drops the question that state_text appends
last. The training recipe (Recipe.state_ids) keeps the tail, so send
truncation_side='left' explicitly. vLLM accepts the field from 0.18 and
ignores unknown fields before that.

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