refactor: refactor xqa interface #1701
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @yzh119, 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 focuses on enhancing the usability and consistency of the XQA attention mechanism. It streamlines the xqa function by intelligently inferring numerous operational parameters from the input tensor shapes, thereby reducing boilerplate and potential for error. The changes also broaden the supported data types to include bfloat16 and ensure the new interface is properly documented.
Highlights
- API Refactoring: The XQA (eXtended Query Attention) API has been refactored to align its terminology and parameter names with other attention backends, improving consistency across the codebase.
- Simplified Interface: The public
xqafunction's interface has been significantly simplified. Many parameters, such as batch size, number of query heads, head dimension, and maximum sequence length, are now automatically inferred from the shapes of the input tensors, reducing the number of explicit arguments required. - BFloat16 Support: The XQA kernels now support
bfloat16data type in addition tofloat16, providing more flexibility for mixed-precision training and inference. - Documentation Update: The XQA API has been added to the project's documentation, making it easier for developers to understand and utilize the functionality.
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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.
| 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 pull request 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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request significantly improves the xqa interface by refactoring the API to be more consistent with other backends and simplifying it by inferring parameters from tensor shapes. The addition of documentation is also a welcome improvement.
My review identified a critical issue in the sliding window logic which would cause incorrect behavior when the sliding window is disabled. I've also pointed out a minor type hint inconsistency for an optional parameter that should be corrected for better code clarity and correctness.
Overall, these are great changes that enhance the usability of the library. Once the identified issues are addressed, this PR will be in excellent shape.
| max_seq_len = num_pages_per_seq * page_size | ||
|
|
||
| # Determine if sliding window is used | ||
| use_sliding_window = sliding_win_size >= 0 |
There was a problem hiding this comment.
There's a logic error in determining use_sliding_window. According to the docstring, a sliding_win_size of 0 should disable the sliding window. However, with the current logic sliding_win_size >= 0, it will be enabled. This causes the kernel to be compiled with SLIDING_WINDOW=1 and when called with slidingWinSize=0, it incorrectly skips all tokens in the sequence. The condition should be sliding_win_size > 0 to correctly disable sliding window when the size is 0.
| use_sliding_window = sliding_win_size >= 0 | |
| use_sliding_window = sliding_win_size > 0 |
| seqLen: torch.Tensor, | ||
| batchSize: int, | ||
| kvCacheScale: torch.Tensor, | ||
| sinks: torch.Tensor, |
There was a problem hiding this comment.
The type hint for the sinks parameter should be Optional[torch.Tensor] to match its usage. The public xqa function allows sinks to be None, and the C++ backend is designed to handle this. This type hint should be updated to reflect that it's an optional parameter to avoid confusion and potential issues with static analysis tools.
| sinks: torch.Tensor, | |
| sinks: Optional[torch.Tensor], |
| seqLen: torch.Tensor, | ||
| batchSize: int, | ||
| kvCacheScale: torch.Tensor, | ||
| sinks: torch.Tensor, |
There was a problem hiding this comment.
Similar to the xqa stub, the type hint for the sinks parameter in _fake_xqa should be Optional[torch.Tensor] to accurately reflect that it can be None.
| sinks: torch.Tensor, | |
| sinks: Optional[torch.Tensor], |
Signed-off-by: Qidi Sang <[email protected]>
|
Moved to #1769 |
1 participant
📌 Description
This PR
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes