Conversation
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughThe test in Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🧰 Additional context used🧬 Code graph analysis (1)tests/attention/test_trtllm_gen_mla.py (1)
🔇 Additional comments (3)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello @nvmbreughe, 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 enhances the testing infrastructure for TRT-LLM batch decode MLA by refactoring an existing test into a more modular and reusable helper function. This change facilitates the addition of a new, dedicated test for DSR1 MLA, enabling comprehensive validation across a wider range of sequence lengths and batch configurations. The primary goal is to ensure the robustness and correctness of the DSR1 MLA implementation within the TRT-LLM framework. Highlights
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the test_trtllm_batch_decode_mla test by splitting it into a helper function and a test wrapper. It also introduces a new test, test_dsr1_trtllm_mla, to validate Multi-head Latent Attention (MLA) with different sequence lengths. The changes are well-structured and improve the test suite's coverage and maintainability. My feedback focuses on adhering to Python's naming conventions for consistency and readability.
| dynamic_scale: bool, | ||
| enable_pdl: bool, | ||
| backend: str, | ||
| MAX_SEQ_LEN: int, |
There was a problem hiding this comment.
Per PEP 8, function parameter names should be in lowercase_with_underscores. Please rename MAX_SEQ_LEN to max_seq_len for consistency. You will also need to update its usage within the function body.
| MAX_SEQ_LEN: int, | |
| max_seq_len: int, |
| @pytest.mark.parametrize("MAX_SEQ_LEN", [1024, 8960]) | ||
| def test_dsr1_trtllm_mla( | ||
| batch_size: int, | ||
| scale: float, | ||
| dtype: torch.dtype, | ||
| page_size: int, | ||
| q_len_per_request: int, | ||
| dynamic_scale: bool, | ||
| enable_pdl: bool, | ||
| backend: str, | ||
| MAX_SEQ_LEN: int, | ||
| ): | ||
| trtllm_batch_decode_mla( | ||
| batch_size, | ||
| scale, | ||
| dtype, | ||
| page_size, | ||
| q_len_per_request, | ||
| dynamic_scale, | ||
| enable_pdl, | ||
| backend, | ||
| MAX_SEQ_LEN, | ||
| ) |
There was a problem hiding this comment.
Following PEP 8 style guidelines, parameter and variable names should be in lowercase_with_underscores. Please rename MAX_SEQ_LEN to max_seq_len in the parametrize decorator, the test function signature, and the call to the helper function.
@pytest.mark.parametrize("max_seq_len", [1024, 8960])
def test_dsr1_trtllm_mla(
batch_size: int,
scale: float,
dtype: torch.dtype,
page_size: int,
q_len_per_request: int,
dynamic_scale: bool,
enable_pdl: bool,
backend: str,
max_seq_len: int,
):
trtllm_batch_decode_mla(
batch_size,
scale,
dtype,
page_size,
q_len_per_request,
dynamic_scale,
enable_pdl,
backend,
max_seq_len,
)
|
/bot run |
|
[FAILED] Pipeline #38686503: 6/18 passed |
|
/bot run |
<!-- .github/pull_request_template.md --> ## 📌 Description Added DSR1 MLA test, and split up the trtllm_batch_decode_mla function. ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 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 - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Improved test suite for batch decoding by making maximum sequence length configurable, adding parameterized runs across short and long lengths, and introducing a compatibility wrapper to preserve legacy behavior. This enhances coverage and validation across varied sequence-length scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Zihao Ye <expye@outlook.com>
4 participants
📌 Description
Added DSR1 MLA test, and split up the trtllm_batch_decode_mla function.
🔍 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
Summary by CodeRabbit