Fix tensor shapes for DeepEP and DeepGEMM assertions #19546
Conversation
Signed-off-by: Piotr Tarasiewicz <ptarasiewicz@nvidia.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 @ptarasiewiczNV, 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 assertion errors encountered when using VLLM with specific DeepEP low-latency and DeepGEMM configurations. The fix involves ensuring that the input tensors provided to the DeepGEMM and DeepEP kernels have the exact shapes expected by those kernels, which was causing failures previously.
Highlights
- DeepGEMM Tensor Shape Fix: Adjusted the shape of the
outputtensor before passing it to thedg.m_grouped_gemm_fp8_fp8_bf16_nt_maskedkernel inbatched_deep_gemm_moe.pyto match the expected input shape, resolving an assertion error. - DeepEP Tensor Shape Fix: Reshaped the
fused_expert_outputtensor based on expert grouping and token dispatch size before callingself.buffer.low_latency_combineindeepep_ll_prepare_finalize.py, fixing a related assertion error.
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 configureGemini 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 addresses assertion errors in DeepEP and DeepGEMM by reshaping tensors to meet kernel expectations. The changes appear correct and targeted.
To further improve the PR, please consider:
- Adding Code Comments: For both
batched_deep_gemm_moe.pyanddeepep_ll_prepare_finalize.py, adding comments explaining the necessity of the.view()operations would greatly enhance code clarity and maintainability. These comments should ideally state what shape the downstream kernel expects. - PR Description Checklist: The PR description is missing details for the 'Test Plan' and 'Test Result' sections. Providing this information would help reviewers and future maintainers understand how the fix was validated and the impact of the changes.
Overall, the core logic changes seem sound.
| @@ -120,6 +120,8 @@ def apply( | |||
| a2q = a2q.view(E, max_num_tokens, -1) | |||
| a2q_scale = a2q_scale.view(E, max_num_tokens, -1) | |||
|
|
|||
| output = output.view(E, max_num_tokens, -1) | |||
There was a problem hiding this comment.
Consider adding a comment here to explain why this reshape is necessary. For example, clarifying that this view aligns the output tensor with the shape expected by the dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked kernel would improve code clarity for future readers.
| _, _, num_max_dispatch_tokens_per_rank, _, num_experts = self.handle | ||
| fused_expert_output = fused_expert_output.view( | ||
| num_experts // self.buffer.group_size, | ||
| self.buffer.group_size * num_max_dispatch_tokens_per_rank, -1) |
There was a problem hiding this comment.
It would be helpful to add a comment explaining this reshape. For instance, mentioning that this view operation transforms fused_expert_output to match the input shape requirements of the self.buffer.low_latency_combine kernel, possibly referencing the expected dimensions (e.g., (num_experts / group_size, group_size * num_max_dispatch_tokens_per_rank, hidden_size)), would enhance readability and maintainability.
| @@ -120,6 +120,8 @@ def apply( | |||
| a2q = a2q.view(E, max_num_tokens, -1) | |||
| a2q_scale = a2q_scale.view(E, max_num_tokens, -1) | |||
|
|
|||
| output = output.view(E, max_num_tokens, -1) | |||
There was a problem hiding this comment.
If the output shape is incorrect it should be fixed in workspace_shapes. I think @varun-sundar-rabindranath already has some fixes for this in another PR.
There was a problem hiding this comment.
Thanks @ptarasiewiczNV - This is the fix I have #19515 - In addition to resolving this it fixes a few other issues as well. Can you please take a look and see if it works for you ? much appreciated 🙌 Thanks.
There was a problem hiding this comment.
Thanks, this looks great. I will be able to test it tomorrow, but definitely this is the proper direction so I will close this PR
|
Closing in favor of #19515 |
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.When running DEP16 DSR1 on 2x8xH100 with
VLLM_ALL2ALL_BACKEND="deepep_low_latency"andVLLM_USE_DEEP_GEMM=1there were 2 separate assertion errors.In DeepGEMM m_grouped_gemm_fp8_fp8_bf16_nt_masked and in DeepEP low_latency_combine.
Both of those can be solved by simply changing the view of the tensor to match the expected shapes by the kernels.
Purpose
Fix assertion errors.
Test Plan
Node 0:
Node 1:
Test Result
No errors during model initialization, correct completion outputs.
(Optional) Documentation Update