fix: Fix cute dsl moe failure with nvidia-cutlass-dsl >= 4.4.0#2735
fix: Fix cute dsl moe failure with nvidia-cutlass-dsl >= 4.4.0#2735nv-yunzheq merged 6 commits intoflashinfer-ai:mainfrom
Conversation
📝 WalkthroughWalkthroughAdds version-gated monkey-patches for older Cutlass versions and replaces synchronization object creation for TCGen05Mma in newer Cutlass (>=4.4.0) to restore compatibility with both legacy and latest Cutlass releases. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related issues
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
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 Tip CodeRabbit can generate a title for your PR based on the changes with custom instructions.Set the |
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 addresses compatibility issues with 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
|
|
/bot run |
There was a problem hiding this comment.
Code Review
This pull request addresses compatibility issues with nvidia-cutlass-dsl >= 4.4.0. The changes correctly make monkey-patches conditional based on the cutlass version, ensuring they only apply to older versions. However, in flashinfer/fused_moe/cute_dsl/blackwell/custom_pipeline.py, the fix for newer versions is applied unconditionally. This could potentially break compatibility with older versions of cutlass-dsl. I've suggested making these changes conditional as well to maintain broader compatibility, which seems to be the intent of this PR.
| # Directly create MbarrierArray for TCGen05Mma consumer, since | ||
| # PipelineAsync._make_sync_object does not handle TCGen05Mma in cutlass >= 4.4.0. | ||
| sync_object_empty = MbarrierArray( | ||
| barrier_storage=barrier_storage.align(min_align=8) + num_stages, | ||
| num_stages=num_stages, | ||
| agent=consumer, | ||
| ) |
There was a problem hiding this comment.
The change to directly use MbarrierArray fixes the issue for cutlass-dsl >= 4.4.0. However, this is an unconditional change. If MbarrierArray is not available or behaves differently in older versions of cutlass-dsl, this could break backward compatibility. The other files in this PR use a version check (hasattr(cutlass, "__version__")) to apply changes conditionally. To ensure compatibility with both older and newer versions, it would be safer to wrap this change in a version check.
if hasattr(cutlass, "__version__"):
# Directly create MbarrierArray for TCGen05Mma consumer, since
# PipelineAsync._make_sync_object does not handle TCGen05Mma in cutlass >= 4.4.0.
sync_object_empty = MbarrierArray(
barrier_storage=barrier_storage.align(min_align=8) + num_stages,
num_stages=num_stages,
agent=consumer,
)
else:
sync_object_empty = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
)| # Directly create MbarrierArray for TCGen05Mma producer, since | ||
| # PipelineAsync._make_sync_object does not handle TCGen05Mma in cutlass >= 4.4.0. | ||
| sync_object_full = MbarrierArray( | ||
| barrier_storage=barrier_storage.align(min_align=8), | ||
| num_stages=num_stages, | ||
| agent=producer, | ||
| ) |
There was a problem hiding this comment.
Similar to the previous change, this direct use of MbarrierArray is unconditional. To ensure backward compatibility with older cutlass-dsl versions, it's safer to use a version check, as done in other files in this PR.
if hasattr(cutlass, "__version__"):
# Directly create MbarrierArray for TCGen05Mma producer, since
# PipelineAsync._make_sync_object does not handle TCGen05Mma in cutlass >= 4.4.0.
sync_object_full = MbarrierArray(
barrier_storage=barrier_storage.align(min_align=8),
num_stages=num_stages,
agent=producer,
)
else:
sync_object_full = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8), num_stages, producer
)| # Directly create MbarrierArray for TCGen05Mma consumer, since | ||
| # PipelineAsync._make_sync_object does not handle TCGen05Mma in cutlass >= 4.4.0. | ||
| sync_object_empty = MbarrierArray( | ||
| barrier_storage=barrier_storage.align(min_align=8) + num_stages, | ||
| num_stages=num_stages, | ||
| agent=consumer, | ||
| ) |
There was a problem hiding this comment.
Again, this unconditional change to use MbarrierArray might break compatibility with older cutlass-dsl versions. A conditional check based on the cutlass version would be a more robust solution, consistent with the approach in other files.
if hasattr(cutlass, "__version__"):
# Directly create MbarrierArray for TCGen05Mma consumer, since
# PipelineAsync._make_sync_object does not handle TCGen05Mma in cutlass >= 4.4.0.
sync_object_empty = MbarrierArray(
barrier_storage=barrier_storage.align(min_align=8) + num_stages,
num_stages=num_stages,
agent=consumer,
)
else:
sync_object_empty = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
)There was a problem hiding this comment.
🧹 Nitpick comments (1)
flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py (1)
305-312: Feature-detect the scheduler API instead ofcutlass.__version__.CUTLASS 4.2.1 documents
PersistentTileSchedulerParams.__init__withoutswizzle_size/raster_along_m, while 4.3.3 and 4.3.4 already document both parameters. That makeshasattr(cutlass, "__version__")a brittle proxy for the capability you actually need here. Checking the constructor signature or the specific scheduler attributes would make this gate line up with the supported API surface. (docs.nvidia.com)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py` around lines 305 - 312, The current gate uses hasattr(cutlass, "__version__") which is brittle; instead feature-detect the scheduler API by inspecting the PersistentTileSchedulerParams constructor signature for the new parameters (e.g., "swizzle_size" and "raster_along_m") or by checking for the presence of the exact scheduler attributes/methods (e.g., cutlass.utils.PersistentTileSchedulerParams.__init__ signature and cutlass.utils.StaticPersistentTileScheduler._get_cluster_work_idx_with_fastdivmod or a FastDivmod helper). Replace the hasattr(cutlass, "__version__") check with a conditional that uses inspect.signature(cutlass.utils.PersistentTileSchedulerParams.__init__) to test for the parameter names or uses hasattr to detect the new API surface, and only apply the monkey-patches (hooked_PersistentTileSchedulerParams_init and hooked_get_cluster_work_idx_with_fastdivmod) when those features are missing.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In
`@flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py`:
- Around line 305-312: The current gate uses hasattr(cutlass, "__version__")
which is brittle; instead feature-detect the scheduler API by inspecting the
PersistentTileSchedulerParams constructor signature for the new parameters
(e.g., "swizzle_size" and "raster_along_m") or by checking for the presence of
the exact scheduler attributes/methods (e.g.,
cutlass.utils.PersistentTileSchedulerParams.__init__ signature and
cutlass.utils.StaticPersistentTileScheduler._get_cluster_work_idx_with_fastdivmod
or a FastDivmod helper). Replace the hasattr(cutlass, "__version__") check with
a conditional that uses
inspect.signature(cutlass.utils.PersistentTileSchedulerParams.__init__) to test
for the parameter names or uses hasattr to detect the new API surface, and only
apply the monkey-patches (hooked_PersistentTileSchedulerParams_init and
hooked_get_cluster_work_idx_with_fastdivmod) when those features are missing.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7222970a-5a70-4d18-8f7b-7473d5db5aea
📒 Files selected for processing (3)
flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.pyflashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.pyflashinfer/fused_moe/cute_dsl/blackwell/custom_pipeline.py
|
[SUCCESS] Pipeline #45750903: 10/20 passed |
|
I cancelled the pr test because the ci won't pass before #2781 lands, and please re-trigger the test after that pr get merged |
📌 Description
Fix issue #2693
Remove the patch when cute dsl version is order than 4.4.0
🔍 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