Tiny add sanity checks for DeepGEMM inputs #7157
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @fzyzcjy, 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 introduces small, optional sanity checks within the DeepGEMM W8A8 FP8 linear layer implementation. The checks ensure that the quantization scales used are valid according to the UE8M0 format, helping to catch potential issues during development or debugging when enabled via an environment variable.
Highlights
- Sanity Checks: Added optional sanity checks for the input scales (
x_scaleandweight_scale) in thedeepgemm_w8a8_block_fp8_linear_with_fallbackfunction. - UE8M0 Format Verification: The new checks verify that the scales conform to the expected UE8M0 format (powers of 2) by comparing them to their ceiling values in that format.
- Conditional Execution: These sanity checks are guarded by the
SGLANG_W8A8_DEEPGEMM_SANITY_CHECK_UE8M0environment variable, allowing them to be enabled or disabled. - Helper Functions: Introduced
_check_ue8m0to perform the assertion andceil_to_ue8m0(copied from DeepGEMM) to calculate the target UE8M0 value.
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 introduces sanity checks for DeepGEMM inputs by adding a check to ensure that x_scale and weight_scale conform to UE8M0 characteristics when the SGLANG_W8A8_DEEPGEMM_SANITY_CHECK_UE8M0 environment variable is set. The changes include the functions _check_ue8m0 and ceil_to_ue8m0.
| def _check_ue8m0(name, x): | ||
| x_ceil = ceil_to_ue8m0(x) | ||
| assert torch.all(x == x_ceil), f"{name=} {x=} {x_ceil=}" |
There was a problem hiding this comment.
Adding type hints to the _check_ue8m0 function improves code readability.
| def _check_ue8m0(name, x): | |
| x_ceil = ceil_to_ue8m0(x) | |
| assert torch.all(x == x_ceil), f"{name=} {x=} {x_ceil=}" | |
| def _check_ue8m0(name: str, x: torch.Tensor) -> None: | |
| x_ceil = ceil_to_ue8m0(x) | |
| assert torch.all(x == x_ceil), f"{name=} {x=} {x_ceil=}" |
Motivation
useful for debugging
Modifications
Checklist