vllm-project
[Core] Use flashinfer sampling kernel when available
Flashinfer contains a combined kernel top_k_top_p_sampling_from_probs, and it is way faster than the sorting kernels used currently. This will eliminate the timely _apply_top_k_top_p function and reduce the GPU time of sampler.
main:
this PR:
End-to-end test (input length=1024, batch size=2048, model=qwen2-1.5b):
main:
Processed prompts: 100%|█████████████| 2048/2048 [01:08<00:00, 29.86it/s, est. speed input: 30574.62 toks/s, output: 3821.83 toks/s]
this PR:
Processed prompts: 100%|█████████████| 2048/2048 [01:02<00:00, 32.61it/s, est. speed input: 33396.84 toks/s, output: 4174.60 toks/s]
BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE
PR Checklist (Click to Expand)
Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.
PR Title and Classification
Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:
[Bugfix]for bug fixes.[CI/Build]for build or continuous integration improvements.[Doc]for documentation fixes and improvements.[Model]for adding a new model or improving an existing model. Model name should appear in the title.[Frontend]For changes on the vLLM frontend (e.g., OpenAI API server,LLMclass, etc.)[Kernel]for changes affecting CUDA kernels or other compute kernels.[Core]for changes in the core vLLM logic (e.g.,LLMEngine,AsyncLLMEngine,Scheduler, etc.)[Hardware][Vendor]for hardware-specific changes. Vendor name should appear in the prefix (e.g.,[Hardware][AMD]).[Misc]for PRs that do not fit the above categories. Please use this sparingly.
Note: If the PR spans more than one category, please include all relevant prefixes.
Code Quality
The PR need to meet the following code quality standards:
- We adhere to Google Python style guide and Google C++ style guide.
- Pass all linter checks. Please use
format.shto format your code. - The code need to be well-documented to ensure future contributors can easily understand the code.
- Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
- Please add documentation to
docs/source/if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.
Notes for Large Changes
Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.
What to Expect for the Reviews
The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:
- After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
- After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
- After the review, the reviewer will put an
action-requiredlabel on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR. - Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.
Thank You
Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!
👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.
Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).
To run full CI, you can do one of these:
- Comment
/readyon the PR - Add
readylabel to the PR - Enable auto-merge.
🚀
![github-actions[bot] avatar](https://avatars.githubusercontent.com/in/15368?v=4 "Published by a pub.dev verified publisher"){kind=small}
The isort and yapf do not agree on the formatting of the import 🤯
isort:
$ isort --check --diff vllm/model_executor/layers/sampler.py
ERROR: /mnt/workspace/pgw/vllm/vllm/model_executor/layers/sampler.py Imports are incorrectly sorted and/or formatted.
--- /mnt/workspace/pgw/vllm/vllm/model_executor/layers/sampler.py:before 2024-08-05 17:48:41.674248
+++ /mnt/workspace/pgw/vllm/vllm/model_executor/layers/sampler.py:after 2024-08-05 17:53:58.759998
@@ -24,8 +24,8 @@
from vllm.utils import async_numpy_to_tensor
try:
- from flashinfer.sampling import (top_k_top_p_sampling_from_probs as
- flashinfer_top_k_top_p_sampling)
+ from flashinfer.sampling import (
+ top_k_top_p_sampling_from_probs as flashinfer_top_k_top_p_sampling)
except ImportError:
flashinfer_top_k_top_p_sampling = None
yapf:
$ yapf --diff vllm/model_executor/layers/sampler.py
--- vllm/model_executor/layers/sampler.py (original)
+++ vllm/model_executor/layers/sampler.py (reformatted)
@@ -24,8 +24,8 @@
from vllm.utils import async_numpy_to_tensor
try:
- from flashinfer.sampling import (
- top_k_top_p_sampling_from_probs as flashinfer_top_k_top_p_sampling)
+ from flashinfer.sampling import (top_k_top_p_sampling_from_probs as
+ flashinfer_top_k_top_p_sampling)
except ImportError:
flashinfer_top_k_top_p_sampling = None
Switched the numpy generator back to torch generator, because the numpy generator is too slow for the original sampler. The flashinfer sampler will become slower but still has improvment in end-to-end test:
Processed prompts: 100%|█████████████| 2048/2048 [01:03<00:00, 32.36it/s, est. speed input: 33138.13 toks/s, output: 4142.27 toks/s]
In the future, we can write a dedicated kernel to improve the random seed generation.
I noticed that the top-k-top-p sampler in flashinfer has incompatible semantics compared to vLLM's. vLLM applys top-k first and then top-p (See #1885), while flashinfer applys top-k and top-p filter simultaneously.
Example:
probs = torch.arange(0.2, 0.6, step=0.05, dtype=torch.float32, device="cuda").softmax(dim=-1)[None, :]
# probs: [[0.1042, 0.1096, 0.1152, 0.1211, 0.1273, 0.1339, 0.1407, 0.1479]]
uniform_samples = torch.tensor([[0.51, 0.6, 0.8, 0.9]], device="cuda")
top_k = torch.tensor([4] * 4, device="cuda")
top_p = torch.tensor([0.5] * 4, device="cuda")
top_k_top_p_sampling_from_probs(probs.expand(4, -1).contiguous(), uniform_samples, top_k, top_p)
# output:
# [tensor([4, 5, 6, 7], device='cuda:0', dtype=torch.int32),
# tensor([True, True, True, True], device='cuda:0')]
while in vllm only 6 and 7 can be selected:
from vllm.model_executor.layers.sampler import _apply_top_k_top_p
_apply_top_k_top_p(probs, p=top_p[:1], k=top_k[:1])
# output: [[ -inf, -inf, -inf, -inf, -inf, -inf, 0.1407, 0.1479]]
Given this, the fallback strategy becomes meaningless. I propose to do the following:
- We keep the original sampler behaviour (the huggingface behaviour) as default
- Add a
VLLM_FLASHINFER_SAMPLERflag for users who accept the changed behaviour to enable faster sampling - Clearly document the behaviour difference.
@peng1999 thanks so much for the notification. I can add the implementations for first top-k then top-p to flashinfer as well.
@peng1999 I add an option filter_apply_order to flashinfer's top_k_top_p_sampling_from_prob to align with huggingface vllm's behavior, another API top_k_top_p_sampling_from_logits have similar functionality but applies to pre-softmax logits: https://github.com/flashinfer-ai/flashinfer/pull/431
Update: https://github.com/flashinfer-ai/flashinfer/pull/431 was merged and top_k_first will be the default option.
v0.1.4 was released: https://github.com/flashinfer-ai/flashinfer/releases/tag/v0.1.4
Flashinfer 0.1.4 now enables more controls to do sampling. Now the fallback mechanism in _top_k_top_p_multinomial_with_flashinfer is totally implemented by kernels in flashinfer library. The new kernels also ensures the top-k-top-p semantics is the same as vLLM's.
@peng1999 please let me know when this is available for the final review and I'll try to get this in asap. Thanks
@peng1999 can you look into the CI failure?
It seems that when GPTQ is enabled, the logits between inference has some small difference even if the same seed is given. And flashinfer's algorithm is sensitive to small differences, hence the output mismatch.
The following is top-9 probs value and its index using batch=10, seed=523, input="Hello, my name is" on Qwen2-1.5B-GPTQ-Int4 model.
tensor([[0.0141, 0.0109, 0.0092, 0.0078, 0.0077, 0.0072, 0.0068, 0.0065, 0.0063],
[0.0142, 0.0110, 0.0093, 0.0078, 0.0078, 0.0071, 0.0068, 0.0065, 0.0063],
[0.0141, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0068, 0.0065, 0.0063],
[0.0142, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0067, 0.0065, 0.0063],
[0.0141, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0068, 0.0066, 0.0063],
[0.0142, 0.0110, 0.0093, 0.0078, 0.0078, 0.0071, 0.0068, 0.0065, 0.0063],
[0.0141, 0.0110, 0.0092, 0.0078, 0.0077, 0.0071, 0.0067, 0.0066, 0.0063],
[0.0141, 0.0110, 0.0093, 0.0078, 0.0078, 0.0071, 0.0068, 0.0066, 0.0063],
[0.0142, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0068, 0.0066, 0.0063],
[0.0141, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0068, 0.0066, 0.0063]],
device='cuda:0')
tensor([[20445, 3757, 6798, 730, 7937, 619, 386, 23016, 350],
[20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266],
[20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266],
[20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266],
[20445, 3757, 6798, 730, 7937, 619, 386, 23016, 350],
[20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266],
[20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266],
[20445, 3757, 6798, 7937, 730, 619, 386, 23016, 350],
[20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266],
[20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266]],
device='cuda:0')
@comaniac Is this behaviour of gptq expected?
That's understandable. We could disable FlashInfer sampling in this case. Meanwhile we may want to note somewhere to encourage users to disable it when they found discrepancy (and unacceptable) outputs.
Also cc @robertgshaw2-neuralmagic @mgoin for the GPTQ accuracy issue visibility.
The PP test and 2-Node test seems to have the same problem as GPTQ. But I have problem running the test locally to verify my guess.
Hmm I suppose this would be the case everywhere then...I'll then suggest the following:
- We disable FlashInfer sampling by default and use the env variable to enable it in this PR.
- We could have follow-up PRs to improve GPTQ tests. Specifically they should check logprobs instead of tokens.
- Then we re-enable FlashInfer sampling by default.
- To avoid changing the API (env variable name, we could use "USE_FLASHINFER_SAMPLER=0/1".
@comaniac it looked like he shared top-logprobs already from the gptq test? If it isn't using logprobs, I agree we should change that
Yeah ideally we should leverage logprobs for all tests with this issue, but I don't expect it to be done in this PR
And flashinfer's algorithm is sensitive to small differences, hence the output mismatch.
@peng1999 we improved the implementation of top-k renorm kernels (https://github.com/flashinfer-ai/flashinfer/pull/456) and the issues might be resolved in the next release.
@peng1999 can you look into the CI failure?
It seems that when GPTQ is enabled, the logits between inference has some small difference even if the same seed is given. And flashinfer's algorithm is sensitive to small differences, hence the output mismatch.
The following is top-9 probs value and its index using batch=10, seed=523, input="Hello, my name is" on Qwen2-1.5B-GPTQ-Int4 model.
tensor([[0.0141, 0.0109, 0.0092, 0.0078, 0.0077, 0.0072, 0.0068, 0.0065, 0.0063], [0.0142, 0.0110, 0.0093, 0.0078, 0.0078, 0.0071, 0.0068, 0.0065, 0.0063], [0.0141, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0068, 0.0065, 0.0063], [0.0142, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0067, 0.0065, 0.0063], [0.0141, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0068, 0.0066, 0.0063], [0.0142, 0.0110, 0.0093, 0.0078, 0.0078, 0.0071, 0.0068, 0.0065, 0.0063], [0.0141, 0.0110, 0.0092, 0.0078, 0.0077, 0.0071, 0.0067, 0.0066, 0.0063], [0.0141, 0.0110, 0.0093, 0.0078, 0.0078, 0.0071, 0.0068, 0.0066, 0.0063], [0.0142, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0068, 0.0066, 0.0063], [0.0141, 0.0110, 0.0092, 0.0078, 0.0078, 0.0071, 0.0068, 0.0066, 0.0063]], device='cuda:0') tensor([[20445, 3757, 6798, 730, 7937, 619, 386, 23016, 350], [20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266], [20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266], [20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266], [20445, 3757, 6798, 730, 7937, 619, 386, 23016, 350], [20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266], [20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266], [20445, 3757, 6798, 7937, 730, 619, 386, 23016, 350], [20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266], [20445, 3757, 6798, 7937, 730, 619, 386, 23016, 11266]], device='cuda:0')@comaniac Is this behaviour of gptq expected?
Hi @peng1999 would you mind trying flashinfer v0.1.6 and see if the problem still exists?
Hi @peng1999 would you mind trying flashinfer v0.1.6 and see if the problem still exists?
Unfortunately the test still not passed using flashinfer v0.1.6.
Tested locally using VLLM_USE_FLASHINFER_SAMPLER=1 pytest -x tests/quantization/test_cpu_offload.py -k test_cpu_offload_gptq:
E AssertionError: Results for model='Qwen/Qwen2-1.5B-Instruct-GPTQ-Int4' are not the same with arg1=['--quantization', 'gptq'] and arg2=['--quantization', 'gptq', '--cpu-offload-gb', '1']. arg1_result={'test': 'seeded_sampling', 'text': ' Amanda Bunnell.', 'finish_reason': 'length', 'usage': CompletionUsage(completion_tokens=5, prompt_tokens=5, total_tokens=10)} != arg2_result={'test': 'seeded_sampling', 'text': ' Ruth Yaireen.', 'finish_reason': 'length', 'usage': CompletionUsage(completion_tokens=5, prompt_tokens=5, total_tokens=10)}
tests/utils.py:275: AssertionError