Updated decorator to support unspecified default #2026
Conversation
nvmbreughe
requested review from
bkryu,
cyx-6,
wenscarl and
yzh119
as code owners
|
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. WalkthroughAllowed SM110 for the FP4 Cutlass GEMM requirement; refactored Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Caller
participant Decorator as backend_requirement(wrapper)
participant Binder as signature.bind + apply_defaults
participant Support as _is_problem_size_supported
participant Target as original function
Caller->>Decorator: call decorated_func(...args, **kwargs)
Decorator->>Binder: bind args/kwargs to signature (apply defaults)
Binder-->>Decorator: bound_kwargs (positional args collapsed)
Decorator->>Support: _is_problem_size_supported(**bound_kwargs)
alt supported
Decorator->>Target: call original func with original *args, **kwargs
Target-->>Caller: return result
else not supported
Decorator-->>Caller: raise error (unsupported config)
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Poem
Pre-merge checks and finishing touches✅ 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🪛 Ruff (0.14.3)flashinfer/utils.py1033-1035: Avoid specifying long messages outside the exception class (TRY003) ⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
🔇 Additional comments (4)
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 refines the 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 updates the backend_requirement decorator to correctly handle default values for the backend parameter and adds support for SM 110 compute capability to the cutlass backend for mm_fp4. The change to the decorator, however, does not correctly handle positional arguments for the backend parameter. I've suggested a more robust implementation using inspect.signature().bind() to correctly resolve the backend argument regardless of whether it's passed as a positional or keyword argument.
flashinfer/utils.py
Outdated
| def get_backend(args, kwargs): | ||
| # backend may not be specified, but could have a default value | ||
| sig = inspect.signature(func) | ||
| backend_parameter = sig.parameters.get("backend") | ||
| if ( | ||
| backend_parameter | ||
| and backend_parameter.default != inspect.Parameter.empty | ||
| ): | ||
| backend = kwargs.get("backend", backend_parameter.default) | ||
| else: | ||
| backend = kwargs.get("backend") | ||
| return backend |
There was a problem hiding this comment.
The new get_backend function is a good improvement to handle default values for the backend parameter. However, it only considers keyword arguments (kwargs) and ignores positional arguments (args), which are passed to it but not used. This can lead to incorrect behavior when backend is provided as a positional argument.
A more robust approach is to use inspect.signature().bind() to correctly resolve the backend argument from both positional and keyword arguments, while also handling default values gracefully. This ensures the correct backend is always identified.
| def get_backend(args, kwargs): | |
| # backend may not be specified, but could have a default value | |
| sig = inspect.signature(func) | |
| backend_parameter = sig.parameters.get("backend") | |
| if ( | |
| backend_parameter | |
| and backend_parameter.default != inspect.Parameter.empty | |
| ): | |
| backend = kwargs.get("backend", backend_parameter.default) | |
| else: | |
| backend = kwargs.get("backend") | |
| return backend | |
| def get_backend(args, kwargs): | |
| # backend may not be specified, but could have a default value | |
| sig = inspect.signature(func) | |
| try: | |
| bound_args = sig.bind(*args, **kwargs) | |
| bound_args.apply_defaults() | |
| return bound_args.arguments.get("backend") | |
| except TypeError: | |
| # Fallback for safety, though it's unlikely to be needed if the call is valid. | |
| return kwargs.get("backend") |
There was a problem hiding this comment.
If the function signature does not indicate a default backend, get_backend will return None. We should handle this edge case in the rest of the decorators, right?
There was a problem hiding this comment.
Should be fixed with apply_defaults
There was a problem hiding this comment.
@nvmbreughe, I think the changes are straightforward, but we may want to add unit tests to check the unspecified backend case. Can we add these?
|
Verified by not providing "backend" to mm_fp4. LGTM. Thanks for the quick fix! |
bkryu
left a comment
bkryu
left a comment
There was a problem hiding this comment.
Thanks for adding the unit tests. LTGM!
4 participants
📌 Description
Updated decorator to support unspecified default. This was causing issues when calling mm_fp4 without backend specified.
Also added SM 110 as a supported backend on the cutlass backend (mm_fp4)
🔍 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
Bug Fixes
Tests