HOTFIX: Skip mamba Stochastic Rounding tests on sm_120#2699
HOTFIX: Skip mamba Stochastic Rounding tests on sm_120#2699bkryu merged 2 commits intoflashinfer-ai:mainfrom
Conversation
ishovkun
requested review from
bkryu,
jimmyzho,
kahyunnam,
nv-yunzheq and
yzh119
as code owners
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdded a new public utility function Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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, 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 a hotfix to prevent test failures related to stochastic rounding on specific GPU architectures, particularly SM120 (Blackwell lite). It achieves this by adding a new utility function to precisely determine hardware support for a critical PTX instruction ( Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
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 introduces a helper function is_cvt_rs_supported to correctly detect hardware support for the cvt.rs.f16x2.f32 PTX instruction, which is not available on all SM >= 100 architectures (e.g., SM120). This function is then used in Mamba tests to conditionally skip or adjust tests that rely on Triton's stochastic rounding, which lacks a software fallback unlike FlashInfer's implementation. The changes correctly address the test failures on unsupported hardware. I have one suggestion to make the new helper function more robust against future architecture changes, given the feature is non-forward-compatible.
Do not assume the forward compatibility of cv.rs Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
flashinfer/utils.py (1)
589-599: Cache this public utility to follow module-level API caching convention.Please add
@functools.cacheto this API helper.Proposed fix
+@functools.cache -def is_cvt_rs_supported(device: torch.device = None) -> bool: +def is_cvt_rs_supported(device: Optional[torch.device] = None) -> bool:As per coding guidelines: "
flashinfer/*.py: Use@functools.cachedecorator on Python API functions to implement two-level module caching (Python in-memory and file-level on disk)".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/utils.py` around lines 589 - 599, The public helper is_cvt_rs_supported lacks module-level in-memory caching; add the `@functools.cache` decorator to the is_cvt_rs_supported function to follow the module caching convention and ensure you import functools at top if not already present; keep the function semantics unchanged, only prepend `@functools.cache` to is_cvt_rs_supported so repeated calls reuse the cached result.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@flashinfer/utils.py`:
- Around line 589-599: The current is_cvt_rs_supported function returns True for
any major >= 10 except 12, which may enable unsupported future SMs; restrict the
support check to known supported families only by changing the predicate in
is_cvt_rs_supported (which calls get_compute_capability) to explicitly return
True only for the supported majors (e.g., 10 and 11) and False otherwise,
ensuring the function does not assume forward compatibility for unknown SM
majors.
---
Nitpick comments:
In `@flashinfer/utils.py`:
- Around line 589-599: The public helper is_cvt_rs_supported lacks module-level
in-memory caching; add the `@functools.cache` decorator to the is_cvt_rs_supported
function to follow the module caching convention and ensure you import functools
at top if not already present; keep the function semantics unchanged, only
prepend `@functools.cache` to is_cvt_rs_supported so repeated calls reuse the
cached result.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4d521cca-be48-43eb-a4c4-78661caa1a91
📒 Files selected for processing (4)
flashinfer/utils.pytests/mamba/test_philox_rounding.pytests/mamba/test_selective_state_update_mtp.pytests/mamba/test_selective_state_update_stp.py
|
/bot run |
…#2699) <!-- .github/pull_request_template.md --> ## 📌 Description I added a checker `is_cvt_rs_supported` that is used to skip mamba tests that use Triton is a reference. FlashInfer implementation already has fallback software-emulated SR emulation, whereas the Triton reference does not, which leads to test failures. Previously, we only checked that the main SM version is larger or equal than 100. ## 🔍 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 - [x] 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`. - [x] 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 * **New Features** * Added a utility to detect CUDA compute capability and whether a device supports a specific runtime rounding path. * **Tests** * Updated tests to use the new runtime capability check, improving when GPU-dependent stochastic rounding and related references are exercised or skipped. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
3 participants
📌 Description
I added a checker
is_cvt_rs_supportedthat is used to skip mamba tests that use Triton is a reference.FlashInfer implementation already has fallback software-emulated SR emulation, whereas the Triton reference does not, which leads to test failures. Previously, we only checked that the main SM version is larger or equal than 100.
🔍 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
New Features
Tests