[bugfix] cutlass_mla_decode is incorrectly declared when CUDA_VERSION…

[cleanerleon]

… < 12040

Motivation

when CUDA_VERSION < 12040, cutlass_mla_decode is declared with 8 parameters, but in "include/sgl_kernel_ops.h," it is declared with 9 parameters.

Modifications

do TORCH_CHECK inside function cutlass_mla_decode, and do not declare cutlass_mla_decode twice.

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

[gemini-code-assist]

Summary of Changes

Hello @cleanerleon, I'm Gemini Code Assist¹! 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 a bug where the cutlass_mla_decode function was incorrectly declared with a differing number of parameters depending on the CUDA version. The fix refactors the conditional compilation logic to ensure that the TORCH_CHECK for CUDA version compatibility is performed inside the cutlass_mla_decode and cutlass_mla_get_workspace_size functions, eliminating a redundant declaration and resolving the parameter mismatch issue for older CUDA versions.

Highlights

• Refactored CUDA Version Checks: The conditional compilation block that previously declared placeholder functions for CUDA_VERSION < 12040 has been removed, streamlining the codebase.
• Runtime CUDA Version Validation: The TORCH_CHECK statements, which enforce a minimum CUDA version of 12.4, are now placed inside the cutlass_mla_decode and cutlass_mla_get_workspace_size functions, ensuring the check occurs upon function invocation rather than at compile time.
• Eliminated Redundant Declarations: This change resolves a parameter mismatch bug by removing an unnecessary duplicate declaration of cutlass_mla_decode that was active for older CUDA versions, ensuring consistent function signatures.

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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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

1. 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. ↩

[gemini-code-assist]

Code Review

This pull request correctly fixes a compilation issue for CUDA versions older than 12.4 by refactoring how cutlass_mla_decode and cutlass_mla_get_workspace_size are defined. Instead of separate function definitions based on CUDA version, which led to a signature mismatch, the functions now have a single definition with an internal version check. This is a solid improvement. I've added one suggestion to further improve code maintainability by using a preprocessor macro for the version check, which would make the code cleaner and less error-prone.

[gemini-code-assist]

To improve readability and maintainability, consider defining a macro for the CUDA version check at the top of the file. This avoids repeating the complex preprocessor condition throughout the file.

For example:

#if defined(CUDA_VERSION) && CUDA_VERSION >= 12040
#define SGL_HAS_CUTLASS_MLA_SUPPORT
#endif

You can then use this macro to simplify the conditional compilation blocks:

1. This line becomes #ifdef SGL_HAS_CUTLASS_MLA_SUPPORT.
2. The check inside cutlass_mla_decode becomes #ifndef SGL_HAS_CUTLASS_MLA_SUPPORT.
3. The check inside cutlass_mla_get_workspace_size also becomes #ifndef SGL_HAS_CUTLASS_MLA_SUPPORT.