[fix] fix cutlass_mla_backend with cuda_graph and add sm_scale for sgl-kernel cutlass_mla #7184
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 the cutlass_mla attention mechanism to ensure it works correctly when CUDA graphs are enabled. The core change involves standardizing the handling and default value of the num_kv_splits parameter, which was previously causing issues in the CUDA graph context.
Highlights
- Fix CUDA Graph Compatibility: I've addressed an issue where the
cutlass_mlaattention backend was not correctly handling thenum_kv_splitsparameter when used in conjunction with CUDA graphs, leading to potential errors or incorrect behavior. - Explicit Parameter Passing: The Python code in
cutlass_mla_backend.pynow explicitly passesnum_kv_splits=1when calculating workspace size and calling the underlying decode kernel, ensuring the correct value is used in CUDA graph capture and execution contexts. - Default Value Alignment: The default value for the
num_kv_splitsparameter has been updated from-1to1in both the C++ kernel declarations (sgl_kernel_ops.h) and the Python wrapper functions (attention.py), aligning the default behavior across the kernel interface.
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 addresses an issue with cutlass_mla when using CUDA graphs by explicitly setting the num_kv_splits parameter to 1 in various parts of the codebase. This change is consistently applied across the Python backend, Python wrapper, and C++ header/fallback. While this fixes the CUDA graph compatibility, the C++ kernel code suggests that num_kv_splits=1 might have performance implications for certain workloads compared to the default auto-tuning behavior. It would be beneficial to add unit tests specifically for this fix, particularly in scenarios where the performance impact might be noticeable (large context, small batch size), and to enhance the PR description with more context about the root cause of the issue and the rationale behind forcing num_kv_splits=1.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces an sm_scale parameter to the CUTLASS MLA kernel for configurable attention scaling and changes the default for num_kv_splits to 1. The review focuses on clarifying a specific benchmark value and ensuring the implications of the num_kv_splits default change are well-understood and documented.
…l-kernel cutlass_mla (sgl-project#7184)
Motivation
set default num_kv_splits to 1, avoiding cuda graph issue
Modifications
Checklist