Add version-adaptive nvvm.atomicrmw wrapper for cutlass-dsl 4.5.x compat#411
Open
Anerudhan wants to merge 1 commit into
Open
Add version-adaptive nvvm.atomicrmw wrapper for cutlass-dsl 4.5.x compat#411Anerudhan wants to merge 1 commit into
Anerudhan wants to merge 1 commit into
Conversation
nvidia-cutlass-dsl 4.5.x generates atomicrmw(res, op, ptr, a, ...) with a required explicit result type; 4.6.0+ infers the result and removed the parameter, so the same call site cannot work across the >=4.5.0 range the pyproject now allows. Introduce cudnn/nvvm_compat.py with an atomicrmw wrapper that detects the installed builder signature once at import and forwards accordingly, and route all 13 raw nvvm.atomicrmw call sites (grouped_gemm, discrete grouped_gemm, gemm_amax, gemm_dsrelu, gemm_swiglu, sdpa) through it. This is the cudnn-frontend equivalent of the atomicrmw fix on rachitg/cute_dsl_kernel_library@fix/support_dsl_47. Validated on sm_100: - 4.5.0 (res branch): dglu fp8 + glu pass (dglu fp8 also proves the issue NVIDIA#397 NVVM failure is absent on 4.5's libNVVM) - 4.6.1 (no-res branch): glu / dsrelu fp8 / gemm_amax fp8 pass; dglu fp8 still blocked by the 4.6.x libNVVM regression (NVIDIA#397) - 4.7.0a0 internal nightly: full dglu suite incl. all fp8 variants passes - 4.7's libNVVM resolves NVIDIA#397 with no kernel changes needed Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
📝 WalkthroughWalkthroughThe change adds a version-adaptive NVVM ChangesNVVM atomic compatibility
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
python/cudnn/nvvm_compat.py (2)
13-13: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winHandle non-introspectable builders explicitly.
inspect.signature()can raiseTypeErrororValueError; because this runs during module import, an unsupported binding currently produces an opaque import failure. Catch these exceptions and raise a targeted compatibility error, or verify that every supported Cutlass DSL build guarantees an inspectable signature. (docs.python.org)🤖 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 `@python/cudnn/nvvm_compat.py` at line 13, Update the _ATOMICRMW_TAKES_RES initialization to handle TypeError and ValueError from inspect.signature(nvvm.atomicrmw) during module import, and raise a targeted compatibility error that clearly identifies the unsupported or non-introspectable NVVM binding.
13-20: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd regression tests for both
atomicrmwsignatures.Mock builders with and without the
resparameter, then verify argument forwarding, especiallyloc/ip, and confirm that the wrapper returns the builder result unchanged.🤖 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 `@python/cudnn/nvvm_compat.py` around lines 13 - 20, Add regression tests for the atomicrmw wrapper covering builders whose signatures include and omit res. Mock each builder, invoke atomicrmw with representative res, op, ptr, a, loc, and ip values, and assert the correct forwarded arguments and that the wrapper returns each builder’s result unchanged.
🤖 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 `@python/cudnn/nvvm_compat.py`:
- Line 13: Update the _ATOMICRMW_TAKES_RES initialization to handle TypeError
and ValueError from inspect.signature(nvvm.atomicrmw) during module import, and
raise a targeted compatibility error that clearly identifies the unsupported or
non-introspectable NVVM binding.
- Around line 13-20: Add regression tests for the atomicrmw wrapper covering
builders whose signatures include and omit res. Mock each builder, invoke
atomicrmw with representative res, op, ptr, a, loc, and ip values, and assert
the correct forwarded arguments and that the wrapper returns each builder’s
result unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 47286d40-3d08-4e1b-90a2-f1a073f67cdf
📒 Files selected for processing (10)
python/cudnn/discrete_grouped_gemm/discrete_kernel_utils.pypython/cudnn/discrete_grouped_gemm/moe_persistent_scheduler.pypython/cudnn/gemm_amax/dense_blockscaled_gemm_persistent_amax.pypython/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.pypython/cudnn/gemm_swiglu/dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.pypython/cudnn/grouped_gemm/moe_kernel_helpers.pypython/cudnn/grouped_gemm/moe_persistent_scheduler.pypython/cudnn/grouped_gemm/utils.pypython/cudnn/nvvm_compat.pypython/cudnn/sdpa/utils.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
nvidia-cutlass-dsl 4.5.x generates atomicrmw(res, op, ptr, a, ...) with a required explicit result type; 4.6.0+ infers the result and removed the parameter, so the same call site cannot work across the >=4.5.0 range the pyproject now allows.
Introduce cudnn/nvvm_compat.py with an atomicrmw wrapper that detects the installed builder signature once at import and forwards accordingly, and route all 13 raw nvvm.atomicrmw call sites (grouped_gemm, discrete grouped_gemm, gemm_amax, gemm_dsrelu, gemm_swiglu, sdpa) through it.
This is the cudnn-frontend equivalent of the atomicrmw fix on rachitg/cute_dsl_kernel_library@fix/support_dsl_47.
Validated on sm_100:
Before submitting
pre-commit runand committed any formatting changes.Affected area
Summary
Why
Related issues
API and compatibility impact
Testing
Summary by CodeRabbit