[BugFix] fix 3 issues: (1) using metadata for causal-conv1d, (2) indexing overflow in v1 vLLM, and (3) init_states in v0 #20838
Conversation
vLLM Signed-off-by: Tuan M. Hoang-Trong <tmhoangt@us.ibm.com>
|
👋 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 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 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @thoangtrvn, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request addresses two critical bug fixes related to the Mamba model implementation within vLLM. It ensures that the causal convolution operation correctly leverages metadata for performance and resolves an issue with the initialization of states for v0 vLLM models, preventing runtime errors during model evaluation with certain architectures like Bamba.
Highlights
- Causal Convolution Metadata Usage: The
causal_conv1dkernel inmamba_mixer2.pyhas been updated to correctly utilizemamba2_metadatafor accelerated computation. Previously, this metadata was not being passed or used effectively, impacting performance. - V0 vLLM Initial States Correction: Fixed a bug in the construction of
initial_statesfor v0 vLLM models withinmamba_mixer2.py. Thehas_initial_states_ptensor was incorrectly assumed to be prefill-only, leading to aRuntimeErrordue to tensor size mismatch. The logic now correctly sliceshas_initial_states_pto[:num_prefills]whenVLLM_USE_V1is false, ensuring proper state initialization for mixed prefill-decode batches.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This PR fixes the causal-conv1d kernel by utilizing metadata for accelerated computation and resolves a tensor shape mismatch in constructing initial_states for v0 vLLM, which was causing a runtime error. The changes effectively address the described issues.
|
|
Signed-off-by: Tuan M. Hoang-Trong <tmhoangt@us.ibm.com>
|
@cyang49 pointed out that the zeros_like leads to 84us overhead in E2E running, so we can replace with a cheaper one. |
Signed-off-by: Tuan M. Hoang-Trong <tmhoangt@us.ibm.com>
|
Added indexing fix to the second kernel, similar to what proposed in the recent PR #20938, when the kernel is used in v1 VLLM which can have a much larger cache size and indexing can be overflow in 32-bit. |
thoangtrvn
changed the title
…xing overflow in v1 vLLM, and (3) init_states in v0 (vllm-project#20838) Signed-off-by: Tuan M. Hoang-Trong <tmhoangt@us.ibm.com> Co-authored-by: Tuan M. Hoang-Trong <tmhoangt@us.ibm.com> Signed-off-by: Himanshu Jaju <hj@mistral.ai>
Purpose
This PR fixes
initial_statesfor v0 vLLM (detected when running lm_eval with Bamba model) from previous PR [v1] Support mamba2 #19327Test Plan
In the main, running this is ok
In the main, running this would trigger below error
The reason is that in the codepath for v0 vLLM of mamba-mixer2,
has_initial_states_pis not actually the prefill-only data; but it is the data of mixed prefill-decode batch. Therefore, we need to applyhas_initial_states_p[:num_prefills,...]With the PR, running this is ok now:
Test Result
(Optional) Documentation Update