[Core] Use flashinfer sampling kernel when available #7137
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 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:
🚀 |
| @@ -223,8 +224,7 @@ def _prepare_seq_groups( | |||
|
|
|||
| if seq_group_metadata.is_prompt: | |||
| if sampling_params.seed is not None: | |||
| generator = torch.Generator(device=device).manual_seed( | |||
| sampling_params.seed) | |||
| generator = np.random.default_rng(seed=sampling_params.seed) | |||
There was a problem hiding this comment.
Switched to NumPy generator because the per-request seed random number generation is faster at CPU unless writing a dedicated GPU kernel. Also, it can overlap with other GPU work, reducing the overall latency.
|
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 = Noneyapf: $ 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 |
|
@simon-mo what do we do if isort and yapf are conflicting? |
|
Yapf ignore |
|
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: 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:
|
|
@peng1999 thanks so much for the notification. |
|
@comaniac I think it's now ready to be reviewed. |
|
@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. @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:
|
Yeah ideally we should leverage logprobs for all tests with this issue, but I don't expect it to be done in this PR |
Co-authored-by: Michael Goin <michael@neuralmagic.com>
Head branch was pushed to by a user without write access
@peng1999 we improved the implementation of top-k renorm kernels (flashinfer-ai/flashinfer#456) and the issues might be resolved in the next release. |
Add an option to check whether there are nan inputs. This PR also removes all `eps` arguments in renorm kernels: previously we pre-set a eps constant to determine when to stop the binary search, however, this might not be accuracy when vocabulary size grows (e.g. >= 1e6 in llama3 where our eps might be set to 1e-5). In this PR, we implement a loop variant which do not rely on any external eps, and it can help us address some of the issues such as vllm-project/vllm#7137 (comment) .
Co-authored-by: Michael Goin <michael@neuralmagic.com>
Co-authored-by: Michael Goin <michael@neuralmagic.com>
Co-authored-by: Michael Goin <michael@neuralmagic.com>
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 |
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_pfunction 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:
this PR:
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:
format.shto format your code.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-requiredand 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:
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.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!