feat: add BF16 grouped GEMM MoE kernels#415
Open
saltyminty wants to merge 4 commits into
Open
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdds experimental SM100 BF16 grouped GEMM APIs for unfused, GLU, dGLU, and Wgrad operations. Each operation supports dense (contiguous weight tensors) and discrete (per-expert pointer arrays) MoE weight layouts, unified backend dispatch between BF16 and legacy block-scaled variants, persistent tile scheduling, CUTLASS/CuTe kernels, tensor validation, and compilation caching. Includes CUDA kernel implementations, descriptor-first lifecycle facades, extensive documentation, and BF16 numerical test fixtures. ChangesSM100 grouped GEMM APIs
Estimated code review effort: 5 (Critical) | ~120 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
Author
|
@cudnn-ci-bot run |
|
Backend pipeline not launched Reason: PR base branch 'main' does not match 'develop' |
saltyminty
force-pushed
the
fix/mingyangw/integrate-bf16-moe-grouped-gemm-kernels-into-cudnn-frontend
branch
from
July 22, 2026 22:05
54de6f4 to
a06ad15
Compare
Collaborator
Author
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-415-a06ad15 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
python/cudnn/grouped_gemm/moe_kernel_helpers.py (1)
809-809: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer a
defover an assigned lambda (Ruff E731).A local
def is_power_of_2(x): return x > 0 and (x & (x - 1)) == 0reads better and matches the existingis_power_of_2helpers already defined elsewhere in the package.🤖 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/grouped_gemm/moe_kernel_helpers.py` at line 809, Replace the assigned lambda with a local def named is_power_of_2, preserving its positive power-of-two check exactly and matching the existing helper style elsewhere in the package.Source: Linters/SAST tools
python/cudnn/grouped_gemm/grouped_gemm_unfused/api.py (1)
21-21: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant
__all__shadowed later in the module.This
__all__ = ["GroupedGemmBf16API"]is overwritten by the__all__at Line 381, so it has no effect and misleadingly impliesGroupedGemmBf16APIis exported. Remove this early assignment.♻️ Proposed cleanup
-__all__ = ["GroupedGemmBf16API"] - - from ._bf16_api import GroupedGemmBf16API🤖 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/grouped_gemm/grouped_gemm_unfused/api.py` at line 21, Remove the early __all__ assignment near GroupedGemmBf16API; retain the later module-level __all__ definition as the sole export declaration.
🤖 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.
Inline comments:
In `@python/cudnn/__init__.py`:
- Around line 277-278: The module’s __all__ currently exports every
non-underscore global, including internal imports and helper names. Update
_EAGER_PUBLIC_NAMES in python/cudnn/__init__.py to use an explicit
intended-public symbol list or exclude the identified helpers/imports, then
append Graph and wrapper as before.
In `@python/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.py`:
- Around line 1254-1274: In the non-BF16 launch path of execute, record both
b_ptrs and sfb_ptrs on current_stream via the existing _record_pointer_stream
mechanism before invoking self._compiled_kernel. Match the BF16 path’s
pointer-stream recording behavior while leaving kernel arguments unchanged.
In `@python/cudnn/grouped_gemm/grouped_gemm_wgrad/_bf16_api.py`:
- Line 49: Update the input_order annotations in
python/cudnn/grouped_gemm/grouped_gemm_wgrad/_bf16_api.py:49 and
python/cudnn/grouped_gemm/grouped_gemm_wgrad/_blockscaled_api.py:48 to use
Python 3.9-compatible Union[WGradInputOrder, str] syntax, or add the future
annotations import in both modules. Preserve the existing defaults and
annotation meaning.
In `@test/python/fe_api/test_grouped_gemm_dglu_bf16_utils.py`:
- Line 4: Remove the internal “!35” token from the module docstring and the
corresponding docstrings near the affected test utilities, replacing it with
descriptive wording such as “the reference oracle” while preserving their
meaning.
---
Nitpick comments:
In `@python/cudnn/grouped_gemm/grouped_gemm_unfused/api.py`:
- Line 21: Remove the early __all__ assignment near GroupedGemmBf16API; retain
the later module-level __all__ definition as the sole export declaration.
In `@python/cudnn/grouped_gemm/moe_kernel_helpers.py`:
- Line 809: Replace the assigned lambda with a local def named is_power_of_2,
preserving its positive power-of-two check exactly and matching the existing
helper style elsewhere in the package.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ff467305-d18f-4386-b3f1-7358bbedd70c
📒 Files selected for processing (36)
README.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.mddocs/fe-oss-apis/overview.mdpython/cudnn/__init__.pypython/cudnn/grouped_gemm/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/_bf16_api.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/api.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/moe_grouped_gemm_dglu_dbias.pypython/cudnn/grouped_gemm/grouped_gemm_glu/_bf16_api.pypython/cudnn/grouped_gemm/grouped_gemm_glu/_blockscaled_api.pypython/cudnn/grouped_gemm/grouped_gemm_glu/api.pypython/cudnn/grouped_gemm/grouped_gemm_glu/moe_grouped_gemm_glu_bias.pypython/cudnn/grouped_gemm/grouped_gemm_unfused/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_unfused/_bf16_api.pypython/cudnn/grouped_gemm/grouped_gemm_unfused/api.pypython/cudnn/grouped_gemm/grouped_gemm_unfused/moe_grouped_gemm.pypython/cudnn/grouped_gemm/grouped_gemm_utils.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/_bf16_api.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/_blockscaled_api.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/api.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_grouped_gemm_wgrad.pypython/cudnn/grouped_gemm/moe_kernel_helpers.pypython/cudnn/grouped_gemm/moe_sched_extension.pypython/cudnn/grouped_gemm/moe_utils.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_glu.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.pytest/python/fe_api/test_grouped_gemm_bf16.pytest/python/fe_api/test_grouped_gemm_bf16_utils.pytest/python/fe_api/test_grouped_gemm_dglu_bf16_utils.pytest/python/fe_api/test_grouped_gemm_glu_bf16_utils.pytest/python/fe_api/test_grouped_gemm_wgrad_bf16_utils.py
Anerudhan
added
cat-feature
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
.pre-commit-config.yaml (1)
18-18: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winRetain Black’s safety check unless the tradeoff is intentional.
--fastskips Black’s AST equivalence check, reducing protection against formatter-induced semantic changes in pre-commit and CI. Remove it unless the performance benefit is required and the pinned Black revision is verified to support this tradeoff.Proposed change
- args: ["--line-length", "160", "--fast"] + args: ["--line-length", "160"]🤖 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 @.pre-commit-config.yaml at line 18, Update the Black hook configuration’s args to remove the “--fast” option, preserving Black’s default AST equivalence safety check while retaining the existing line-length setting.
🤖 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 @.pre-commit-config.yaml:
- Line 18: Update the Black hook configuration’s args to remove the “--fast”
option, preserving Black’s default AST equivalence safety check while retaining
the existing line-length setting.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 59399c70-de0b-4ab7-95e4-c2158b710c0c
📒 Files selected for processing (4)
.pre-commit-config.yamltest/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_glu.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py
🚧 Files skipped from review as they are similar to previous changes (3)
- test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py
- test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py
- test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py
Collaborator
Author
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-415-dba72a4 |
Collaborator
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-415-dba72a4 |
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
3 participants
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.
Summary
Kernel provenance
Validation
PYTHONPATH=python:test/python python -m pytest test/python/fe_api/test_grouped_gemm_glu.py::test_grouped_gemm_glu_wrapper_bf16 test/python/fe_api/test_grouped_gemm_dglu.py::test_grouped_gemm_dglu_wrapper_bf16 -q(6 passed on SM100)pre-commit(Black and Black-Jupyter passed; Clang format had no applicable files)Notes
Summary by CodeRabbit
black --fast.