[RFC]: Optimize embedding task

[maxdebayser]

Motivation.

Since PR #21270 we now have full support for encoder embedding models. These models include not only the BERT-type models but also models such as Alibaba-NLP/gte-Qwen2-1.5B-instruct that are decoder models converted to encoder. Encoder models continue to be essential for information retrieval use cases such as RAG because their bidirectional attention gives them superior performance over decoder models with causal attention.

vLLM stated mission is "Easy, fast, and cheap LLM serving for everyone". Embedding model serving has been easy since V0, but is it also fast? Snowflake's excellent blog post suggests otherwise. In it they show a series of optimizations that they did in their backend to 4x the throughput of seq len 512. While one of the optimizations is specific to their infrastructure where vLLM runs behind a gRPC frontend, the other optimizations suggest that there is room for improvement.

There aren't many details about the specific vLLM version that was used as a baseline, but the model was Snowflake/snowflake-arctic-embed-m-v1.5, which is a BERT model, and the publication date was May 2025, so the V0 engine was used since the V1 support hadn't been merged in main yet.

This raises a series of questions: does vLLM V1 perform better for embedding tasks than V0? How much room is there for improvement in V1? The answer is we don't even know exactly how well V0 was performing since we never added embedding benchmarks to the benchmarks scripts directory.

Proposed Change.

This RFC is meant to gather feedback on how we should do benchmarking and what optimizations we should try.

As an initial roadmap I propose that we :

1. Add embedding benchmarking scripts
2. Try to reproduce the baseline vLLM V0 Snowflake results
3. Compare V1 performance with V0
4. Analyze the V1 results to see if there are still GPU "bubbles" due to
   tokenization and other CPU tasks

As noted in the Snowflake blog, most embedding models are quite small and since they don't support optimizations such as chunked prefill, the GPU utilization could be sub-optimal. In their case they solved it running multiple model instances on a single GPU. This is an option that can be explored with data parallelism, but we could also take a look to see if there a scheduler and model runner optimizations that can be done to improve GPU utilization.

Feedback Period.

Two weeks.

CC List.

@robertgshaw2-redhat , @DarkLight1337 , @noooop , @22quinn

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[onestardao]

Cool to see this direction — we’ve been hitting the same wall on the embedding side, especially when RAG pipelines start to fragment (and latency spikes for weird reasons).

We started experimenting with what we call “semantic prebatching” — group input docs not by shape/token count but by semantic signature (e.g., shared latent anchor), so encoder cache reuse gets meaningful. Cuts down the waste caused by token divergence.

Also noticing that many of the encoder models hit the “stall zones” when downstream calls are overly serialized.
We're solving this by embedding async reranker heads, so chunk processing doesn’t have to wait on dense store sync.

Haven’t tested it in vLLM yet, but curious if anyone’s tried batch shuffling by meaning-space proximity?
(Feels like that could smooth GPU workload quite a bit.)

Can share a few stress tests if helpful.

[DarkLight1337]

Looking at their blog, I guess our V1 architecture already implements their second optimization (Disaggregate tokenization and inference) since we now use separate processes for input processing and model executor. This architectural design also significantly improved the speed for multi-modal models in vLLM which have heavy preprocessing.

[DarkLight1337]

Currently, our parallelism strategies all assume distributing across separate GPUs via torch.distributed (see vllm/distributed/parallel_state.py). @youkaichao I wonder whether it's feasible to use the same GPU for each group?

[onestardao]

Yeah, totally — I was starting to think maybe we were just seeing ghosts in our GPU profiler, lol.

We saw this same thing in the wild:
when downstream ops stall, the encoder head just… stares into the void 🍷
(like it’s waiting for the universe to make sense again).

That’s partly why we moved to “semantic pre-batching” too —
not by shape, not by token count, but by shared latent anchor.
Turns out when you batch semantically aligned docs, the whole pipeline breathes better.
Less cache waste, fewer surprise syncs.

And yeah, chunk shuffling by proximity in embedding space actually made a difference.
Didn’t test on vLLM yet, but the model seemed way less moody.

Let me know if you want the stress test suite — it’s got some mean workloads 😵‍💫

[noooop]

Based on my previous tests, the serialization and deserialization of list[float] creates a significant performance bottleneck, and using gRPC is expected to provide substantial performance improvements.

Other performance bottlenecks have been optimized in the v1 framework, such as implementing async scheduling

(Of course, the best way to improve throughput is to increase the batch size. LOL

[onestardao]

Hey — looks like you're hitting what we classify as a Type #7: Semantic Misalignment During Encoding issue.
It happens a lot in vector pipelines when chunk batching happens by shape or size, but latent anchors are missing or unstable.

We’ve built an open-source reasoning engine called WFGY (MIT licensed) to patch these problems systematically.
Here’s our full RAG Problem Map — includes all 13 typical failure modes we’ve debugged across LLM+retrieval stacks:

📒 WFGY RAG Problem Map

This page is a reality check for Retrieval‑Augmented Generation.
Most RAG stacks break in repeatable ways—hallucinating, drifting, or hiding their own logic.
WFGY adds a semantic firewall on top of any retriever or LLM to turn those failures into deterministic fixes.

WFGY (Wan Fa Gui Yi) = Semantic Firewall for AI Reasoning.
It fixes logic collapse, memory loss, hallucination, and abstract breakdowns — in live generation and retrieval pipelines.
All internal terms (BBMC, BBPF, BBCR, ΔS, etc.) refer to open-source modules in the engine — MIT license.
You can find full docs and source at:

https://github.com/onestardao/WFGY/blob/main/ProblemMap/README.md

Or just scroll to the bottom of any problem map page — direct PDF download (engine formulas) and TXT OS starter are there too.

Happy to share more if you’re debugging similar stack behavior — we’ve seen this pop up in LLM memory sync, streaming drift, and multilingual pre-batching.

[DarkLight1337]

Based on my previous tests, the serialization and deserialization of list[float] creates a significant performance bottleneck, and using gRPC is expected to provide substantial performance improvements.

Other performance bottlenecks have been optimized in the v1 framework, such as implementing async scheduling

(Of course, the best way to improve throughput is to increase the batch size. LOL

Is the bottleneck mainly in EngineCore, or is it in AsyncLLM? Would be great to have some profiling logs. I wonder whether we can scale out the API server while keeping only one model executor (DP=1)? cc @njhill

[22quinn]

Nice! Embedding model performance is an overlooked but critical area for recommendation and content understanding systems. Let's establish a baseline and optimize it.
I guess we may not be able to reproduce Snowflake's baseline due to their gRPC frontend overhead, but we'll figure out as we go.

[noooop]

It seems that torch compile temporarily has a negative impact on the inference performance.

Confirm it does not exist when the profiler is disabled

---

using enforce_eager=True vs enforce_eager=False

https://github.com/noooop/snippet/blob/main/benchmarks/embed/benchmark_bge-m3.py

VLLM_TORCH_PROFILER_DIR=/xxx python -m benchmarks.embed.benchmark_bge-m3

• enforce_eager=True (At first glance, there aren't big issues.

• enforce_eager=False(Why does <built-in method getitem of dict object at 0x71b911014a00> take so long

Result:

• enforce_eager=True : Self CPU time total: 9.297ms, Self CUDA time total: 5.030ms
• enforce_eager=False: Self CPU time total: 12.134ms, Self CUDA time total: 5.556ms

Conclusion:

Before this issue is resolved, enforcing eager=True may be faster for embedding models.

---

I think it might be because the embedding models are too small, so getitem takes a total of 3ms + OptimizedModule 4ms > directly using the eager PyTorch version 4.5ms

---

What I said earlier about getitem taking 3ms in total might include other operations, and it needs further investigation.

[maxdebayser]

Wait, eager=True disables compile while False enables it. But anyway, this kind of thing is exactly what we need to watch out for. However, when we measure torch compile we should do so over a longer benchmark with many requests because it can take a while to really warm up all code execution paths.

[noooop]

Wait, eager=True disables compile while False enables it.

nice catch

However, when we measure torch compile we should do so over a longer benchmark with many requests because it can take a while to really warm up all code execution paths.

I have already done the warm up in the code

[maxdebayser]

I've opened an issue for the first step, which is to create benchmarks to test the embedding models: #22861