Fix AWQ Dequant and Weight Loading of deepseek v2

[AniZpZ]

Motivation

related issue: #6312
There is an incompatibility between AWQ and the new features, which makes AWQ unusable in the latest SGLang. This PR aims to fix the issue.

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Hello @AniZpZ, 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!

Summary of Changes

Hello! Gemini or gemini-code-assist here, providing a summary of this pull request. This PR aims to address an incompatibility issue between AWQ (Activation-aware Weight Quantization) and recent features in SGLang, specifically impacting the usability of AWQ in the latest version. The changes involve modifications to the DeepSeek v2 model's weight loading logic to correctly handle weight concatenation for AWQ and moe_wna16 quantization. Additionally, the CUDA kernel for AWQ dequantization is updated to correctly calculate grid dimensions using ceiling division and includes boundary checks to prevent out-of-bounds memory access, ensuring robustness for various weight matrix dimensions. The test suite for AWQ dequantization is also updated with new parameters to cover relevant edge cases.

Highlights

• AWQ Incompatibility Fix: Addresses a known issue (related to #6312) where AWQ was unusable due to incompatibilities with new SGLang features.
• DeepSeek v2 Weight Loading: Modifies the load_weights function for the DeepSeek v2 model to adjust the concatenation dimension for weights when using AWQ or moe_wna16 quantization.
• AWQ Dequantization Kernel Fixes: Updates the CUDA kernel for AWQ dequantization (dequantize_weights) to include a boundary check based on the actual weight dimensions and corrects the kernel launch grid calculation to use ceiling division, ensuring all elements are processed.
• AWQ Dequantization Tests: Updates the test parameters for the AWQ dequantization comparison test to include dimensions that might expose issues with the previous grid calculation logic.

Changelog

• python/sglang/srt/models/deepseek_v2.py
  • Modified load_weights to conditionally set the concatenation dimension (cat_dim) to 1 for AWQ and moe_wna16 quantization when fusing q_a_proj_weight and kv_a_proj_weight.
• sgl-kernel/csrc/gemm/awq_kernel.cu
  • Added qweight_rows parameter to the dequantize_weights CUDA kernel signature.
  • Added a boundary check if (col >= qweight_cols || row >= qweight_rows) return; inside the dequantize_weights kernel.
  • Changed the calculation of x_blocks and y_blocks in awq_dequantize to use ceiling division (size + threads - 1) / threads.
  • Passed the new qweight_rows parameter to the dequantize_weights kernel launch calls.
• sgl-kernel/tests/test_awq_dequant.py
  • Updated the lists of qweight_row and qweight_col parameters used in the test_awq_dequant_compare_implementations test.

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.

---

Weights are quantized,
Kernel grid now covers all,
AWQ works again.

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

The pull request addresses an incompatibility issue between AWQ and new features in SGLang, specifically related to dequantization and weight loading in deepseek v2 models. The changes include modifications to both Python and CUDA code to ensure correct AWQ operation. Overall, the changes seem well-structured and targeted at resolving the identified issue. However, there are a few areas that could benefit from further scrutiny and refinement.

Summary of Findings

• Readability and Maintainability: The repeated condition in python/sglang/srt/models/deepseek_v2.py should be extracted to a variable for better readability.
• Kernel Argument Usage: Verify that the qweight_rows argument added to the dequantize_weights kernel in sgl-kernel/csrc/gemm/awq_kernel.cu is actually used within the kernel.
• Test Coverage: Consider re-enabling the larger test values in sgl-kernel/tests/test_awq_dequant.py to ensure robustness across a wider range of input sizes.

Merge Readiness

The pull request appears to address the stated issue and includes necessary code modifications. However, the high severity comments regarding the usage of qweight_rows in the CUDA kernel should be addressed before merging. Additionally, consider the medium severity comments to improve the code's readability and test coverage. I am unable to approve this pull request, and recommend that others review and approve this code before merging.

[gemini-code-assist]

It's good that qweight_rows is being passed to the dequantize_weights kernel. Double check that the kernel uses this value to ensure that the dequantization is performed correctly for the new deepseek v2 models.

[gemini-code-assist]

It's good that qweight_rows is being passed to the dequantize_weights kernel. Double check that the kernel uses this value to ensure that the dequantization is performed correctly for the new deepseek v2 models.

[gemini-code-assist]

The condition self.quant_config.get_name() == "awq" or self.quant_config.get_name() == "moe_wna16" is repeated. Consider extracting this to a variable for better readability and maintainability.

is_awq_or_moe_wna16 = self.quant_config.get_name() == "awq" or self.quant_config.get_name() == "moe_wna16"
                            cat_dim = 0
                            if is_awq_or_moe_wna16:
                                cat_dim = 1

[gemini-code-assist]

Adding qweight_rows as an argument to the dequantize_weights kernel is a good step to ensure correct dequantization. However, it's crucial to verify that qweight_rows is actually being used within the kernel to influence the dequantization process. If it's not used, it's an unnecessary argument.

[gemini-code-assist]

This check if (col >= qweight_cols || row >= qweight_rows) return; is important to prevent out-of-bounds access. However, consider adding an assertion or a log message (if the condition is met unexpectedly) to help debug potential issues with the kernel launch configuration.

[gemini-code-assist]

The calculation of x_blocks and y_blocks uses ceiling division to ensure all elements are processed. This is good for correctness. However, it's important to ensure that the kernel handles the case where qweight_cols or qweight_rows is not a multiple of x_num_threads or y_num_threads efficiently. Consider adding a comment explaining this.

[gemini-code-assist]

The commented-out test parameters suggest that larger values were previously used. It would be beneficial to re-enable these larger values in the test suite to ensure that the fix is robust across a wider range of input sizes. If there's a reason these were commented out (e.g., resource constraints), it should be documented.