[Performance]: logit bias implementation uses a slow for loop

[iacolippo]

Proposal to improve performance

We were running logit bias with a big dictionary and noticed a significant slowdown in generation. We looked into the implementation and saw that it uses a for loop

vllm/vllm/entrypoints/openai/logits_processors.py

Line 48 in a79b122

for token_id, bias in logit_bias.items():

Any specific reason this is done this way? From quick tests it seems that a scatter_add would be significantly faster. If this has not been considered before, I'll spend some time to make a proper benchmark and a PR.

On my Mac, with a -100 bias on 40k tokens out of 150k

In [48]: %timeit f_for(x, logit_bias)
106 ms ± 2.92 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [49]: %timeit f_scatter(x, logit_bias)
3.74 ms ± 13.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

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.

[DarkLight1337]

cc @njhill

[NickLucche]

Just speculating, but maybe the tensor conversion was deemed costly for small dictionaries?
Anyways it might be worth to give it a try with the torch-friendly logit_bias_logits_processor implementation and see how it does with small dictionaries (I have no doubts at the scale you reported the for loop one pales in comparison).

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[njhill]

@iacolippo FWIW we'll be making sure to have efficient implementations for all of the sampling ops in V1. Preparing a more efficient logit bias impl in this PR.

[xu-song]

Here is an efficient implementation in V0 #13334

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!