chore: upgrade sgl-kernel v0.1.8.post2 #7186
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @Alcanderian, 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 implements a fix for a scaling problem within the CUTLASS MLA attention backend. Specifically, it ensures that the correct scaling factor is applied during the decoding process by passing the necessary parameter to the kernel function.
Highlights
- Fix CUTLASS MLA Scaling: Corrected the
forward_decodefunction in the CUTLASS MLA attention backend to pass the appropriatesm_scaleparameter to the underlying kernel call. This addresses a potential scaling issue during the decoding phase.
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 pull request aims to fix the scaling factor used in the CUTLASS MLA backend by passing layer.scaling as sm_scale to the cutlass_mla_decode kernel. This is a positive change for correctness, especially if YARN scaling or other dynamic scaling factors are involved.
The main point of attention for this PR is its dependency on an update to the sgl-kernel library. The cutlass_mla_decode function in sgl-kernel needs to be updated to accept and use this new sm_scale parameter. The PR description correctly notes that a new version of sgl-kernel needs to be bumped, but its reference to PR #7184 (this PR itself) as a prerequisite is circular and should be clarified, ideally by linking to the specific sgl-kernel PR or commit that introduces the necessary changes. Merging this PR without the corresponding kernel update will lead to runtime errors.
| @@ -280,6 +280,7 @@ def forward_decode( | |||
| seq_lens=forward_batch.seq_lens.to(torch.int32), | |||
| page_table=self.forward_metadata.block_kv_indices, | |||
| workspace=self.forward_metadata.workspace, | |||
| sm_scale=layer.scaling, | |||
There was a problem hiding this comment.
This change introduces sm_scale=layer.scaling to the cutlass_mla_decode call. This is a crucial fix for ensuring the correct scaling factor, potentially including YARN adjustments, is used in the CUTLASS MLA kernel.
However, this change is critically dependent on an upcoming update to the sgl-kernel library. Specifically, the sgl_kernel.attention.cutlass_mla_decode Python wrapper and the underlying C++ kernel must be modified to:
- Accept a new
sm_scaleparameter. - Utilize this
sm_scaleparameter internally, instead of the currently hardcoded scale calculation in the C++ kernel.
The current sgl-kernel code (as per the provided context) does not support this sm_scale parameter. If this PR is merged without the corresponding sgl-kernel update and version bump, it will lead to a TypeError at runtime due to an unexpected keyword argument.
Please ensure the following:
- The
sgl-kernelchanges are made and tested. - The new Python signature for
sgl_kernel.attention.cutlass_mla_decodeis compatible with this call (e.g.,sm_scaleis added as a keyword argument, andnum_kv_splitsis handled correctly). - The
sgl-kernelversion is bumped and integrated intosglangconcurrently with or before this PR. - The PR description's dependency note (currently referencing PR [fix] fix cutlass_mla_backend with cuda_graph and add sm_scale for sgl-kernel cutlass_mla #7184, which is this PR itself) should be clarified to point to the relevant
sgl-kernelpull request or commit to avoid confusion.
Co-authored-by: zhyncs <me@zhyncs.com>
Motivation
should be merged after #7184 and new version of sgl-kernel bumped
Modifications
Checklist