[Performance] V1 Pooling Models E2E Performance Optimization

[noooop]

Purpose

V1 Pooling Models E2E Performance Optimization

#  https://github.com/noooop/snippet/blob/main/benchmarks/embed/profile.py
VLLM_TORCH_PROFILER_DIR=/xxx python profile.py

main CPU: 67.134ms CUDA: 46.417ms -> this pr CPU: 39.740ms CUDA: 39.890ms

benchmarks:

• sentence_transformers: https://github.com/noooop/snippet/blob/main/benchmarks/embed/st.py
• v0: https://github.com/noooop/snippet/blob/main/benchmarks/embed/v0.py
• v1: https://github.com/noooop/snippet/blob/main/benchmarks/embed/v1.py

Result:

• X-axis: Throughput (token/s)
• Y-axis: Latency, Time needed for one step (ms)
• The curve lower right is better ↘

Frankly speaking,
this optimizations only show significant improvements when handling many small requests—they have almost no impact on large requests (since the backbone network latency dominates in such cases).

Of course, a faster implementation is always cooler.

Please click here for long details.
↓↓↓↓↓↓↓↓↓

1. replace_roberta_positions takes too long

6d8f55e: CPU: 67.134ms CUDA: 46.417ms -> CPU: 53.261ms CUDA: 42.820ms

Correct Testing

• pytest -s -vvv tests/models/language/pooling/test_baai.py
  • BAAI/bge-m3 for XLMRobertaModel
  • BAAI/bge-reranker-base for XLMRobertaForSequenceClassification

2. reduce cuda sync

8ff2418: CPU: 53.261ms CUDA: 42.820ms -> CPU: 43.016ms CUDA: 42.516ms

3. Use as many batch operations as possible

13e44df: CPU: 43.016ms CUDA: 42.516ms -> CPU: 42.463ms CUDA: 40.629ms

4. non_blocking seq_lens (Use seq_lens_cpu directly later

876cb9a: CPU: 42.463ms CUDA: 40.629ms -> CPU: 39.934ms CUDA: 40.128ms

5. remove prompt_len == hidden_states.shape[0] cuda sync by using prompt_len (cpu tenser) == num_scheduled_tokens (cpu tenser)

93d0a95: CPU: 39.934ms CUDA: 40.128ms -> CPU: 39.509ms CUDA: 40.094ms

6. _decode_token_type_ids uses ones_like to avoid getting shape which causes cuda sync

f649899: CPU: 39.509ms CUDA: 40.094ms -> CPU: 39.317ms CUDA: 40.109ms

7. non_blocking torch.zeros in _build_encoder_only_attn_metadata

d7fa9a8: CPU: 39.317ms CUDA: 40.109ms -> CPU: 39.182ms CUDA: 39.981ms

8. using pooling_cursor Rather than torch.split(hidden_states, num_scheduled_tokens_list)

2f10806: CPU: 39.182ms CUDA: 39.981ms -> CPU: 39.740ms CUDA: 39.890ms

Test Plan

keep CI green

Test Result

(Optional) Documentation Update

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[dosubot]

Related Documentation

No published documentation to review for changes on this repository.
Write your first living document

^{How did I do? Any feedback?}  

[maxdebayser]

@noooop , nice work! There are two optimizations in this PR: the roberta position ids and removing the split of the hidden states. Do you have benchmark results on how much each individually changes the performance compared to the main branch?

[noooop]

@noooop , nice work! There are two optimizations in this PR: the roberta position ids and removing the split of the hidden states. Do you have benchmark results on how much each individually changes the performance compared to the main branch?

Please click detail for long details.

[noooop]

these optimizations only show significant improvements when handling many small requests—they have almost no impact on large requests (since the backbone network latency dominates in such cases).

[noooop]

@DarkLight1337

By the way, since #22878 Pooling models mteb test uses enforce_eager, mteb_test_embed_models has become less (never) flaky.

I'll reset the threshold back to MTEB_EMBED_TOL = 1e-4 next month.

I'm not 100% sure that flaky is caused by torch.compile. It is very likely caused by it.

cc @robertgshaw2-redhat

[DarkLight1337]

Tests pass so this should be good to go, thanks!