[TRTLLM-13445][test] Add SageAttention backend-routing and SAGE recip…#15746
Conversation
📝 WalkthroughWalkthroughAdds a ChangesTRTLLM-SAGE Attention Test Coverage
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/unittest/_torch/visual_gen/test_visual_gen_args.py (1)
102-102: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new parameterized test signature.
Line 102 adds a new test function without an explicit return type. Please add
-> Nonehere as well so the new test keeps the same typing bar as the rest of the repo.As per coding guidelines, "Always annotate functions with return types (use
Noneif the function does not return anything)."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_visual_gen_args.py` at line 102, The new parameterized test function test_supported_quant_config_sage is missing an explicit return type annotation; update its signature to include -> None so it matches the repo’s typing conventions. Apply this change directly on the test method definition and keep the existing parameters unchanged.Source: Coding guidelines
tests/unittest/_torch/visual_gen/test_attention_integration.py (1)
296-330: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations to the new helper and test methods.
Lines 296-330 introduce three new Python functions without explicit return types. Please annotate
_build_sage_routed_attentionand both test methods so the new test code stays aligned with the repo's typing rules.Suggested update
-def _build_sage_routed_attention(qkv_mode: QKVMode): +def _build_sage_routed_attention( + qkv_mode: QKVMode, +) -> tuple[Attention, QuantAttentionConfig]: """Build a TRTLLM-SAGE-configured Attention for backend-routing checks.""" @@ -class TestSageAttentionBackendRouting: - def test_self_attention_uses_trtllm_sage_backend(self): +class TestSageAttentionBackendRouting: + def test_self_attention_uses_trtllm_sage_backend(self) -> None: @@ - def test_cross_attention_with_sage_config_falls_back_to_vanilla(self): + def test_cross_attention_with_sage_config_falls_back_to_vanilla(self) -> None:As per coding guidelines, "Always annotate functions with return types (use
Noneif the function does not return anything)."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_attention_integration.py` around lines 296 - 330, The new helper and test methods in TestSageAttentionBackendRouting are missing explicit return annotations, which violates the repo’s typing rule. Update _build_sage_routed_attention to declare its returned tuple type, and mark both test_self_attention_uses_trtllm_sage_backend and test_cross_attention_with_sage_config_falls_back_to_vanilla as returning None to keep the test module consistent with the rest of the codebase.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_attention_integration.py`:
- Around line 296-330: The new helper and test methods in
TestSageAttentionBackendRouting are missing explicit return annotations, which
violates the repo’s typing rule. Update _build_sage_routed_attention to declare
its returned tuple type, and mark both
test_self_attention_uses_trtllm_sage_backend and
test_cross_attention_with_sage_config_falls_back_to_vanilla as returning None to
keep the test module consistent with the rest of the codebase.
In `@tests/unittest/_torch/visual_gen/test_visual_gen_args.py`:
- Line 102: The new parameterized test function test_supported_quant_config_sage
is missing an explicit return type annotation; update its signature to include
-> None so it matches the repo’s typing conventions. Apply this change directly
on the test method definition and keep the existing parameters unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9abf4225-205a-477e-beb5-903639125c54
📒 Files selected for processing (2)
tests/unittest/_torch/visual_gen/test_attention_integration.pytests/unittest/_torch/visual_gen/test_visual_gen_args.py
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #56589 [ run ] triggered by Bot. Commit: |
|
PR_Github #56589 [ run ] completed with state |
…e coverage Add unit coverage for the SageAttention integration (dev epic TRTLLM-12126): - TestSageAttentionBackendRouting: verify TRTLLM+quant_attention_config routes self-attention (FUSE_QKV) to the SAGE path and that cross-attention (SEPARATE_QKV) falls back to VANILLA. - Parametrize the quant-config validator test over all five supported SAGE recipes. Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com>
yingguo-trt
force-pushed
the
yiguo/TRTLLM-13445-sageattention-qa
branch
from
61400f9 to
0888939
Compare
|
/bot run |
1 similar comment
|
/bot run |
|
/bot run |
|
PR_Github #57296 [ run ] triggered by Bot. Commit: |
|
PR_Github #57296 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57413 [ run ] triggered by Bot. Commit: |
|
PR_Github #57413 [ run ] completed with state |
NVIDIA#15746) Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com>
4 participants
…e coverage
Add unit coverage for the SageAttention integration (dev epic TRTLLM-12126):
Summary by CodeRabbit
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.