GEMM+RoPE+MXFP8 fusion#367
Merged
Anerudhan merged 6 commits intoJul 10, 2026
Merged
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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 (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a new SM100 fused projection path combining GEMM, per-head YARN RoPE, and rowwise/columnwise MXFP8 quantization, with kernel, APIs, exports, documentation, and tests. ChangesFused GEMM+RoPE+MXFP8 kernel and API
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Anerudhan
added
cat-feature
Collaborator
![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: 5
🧹 Nitpick comments (1)
python/cudnn/gemm_proj_rope_mxfp8/api.py (1)
187-195: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the full public wrapper contract.
The docstring should state the SM100-only requirement, required dtypes for
w/cos/sin, contiguity/alignment expectations, and exact output scale shapes. As per path instructions,python/cudnn/**: “Focus on documentation”.🤖 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/gemm_proj_rope_mxfp8/api.py` around lines 187 - 195, The public wrapper docstring in the allocate/compile/run entrypoint should document the full contract, not just the basic tensor shapes. Update the docstring for the wrapper that takes x, w, cos, sin, w_out_in, and stream to explicitly mention the SM100-only requirement, the required dtypes for w/cos/sin, any contiguity and alignment constraints, and the exact output scale tensor shapes. Keep the wording aligned with the existing API symbols so callers can understand the supported inputs and outputs without reading the implementation.Source: Path instructions
🤖 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/gemm_proj_rope_mxfp8/__init__.py`:
- Around line 10-15: The __all__ export list in the module initializer is not
sorted, which triggers Ruff RUF022. Reorder the entries in __all__ in
alphabetical order while keeping the same exported symbols, and verify the
module-level export list in __init__.py remains the only place needing this
adjustment.
In `@python/cudnn/gemm_proj_rope_mxfp8/api.py`:
- Around line 66-92: The check_support method currently validates only dtypes,
token tiling, and CUDA capability, but it should also enforce the full kernel
contract before compile. Add shape, layout, contiguity, and device checks in
check_support for x_desc, w_desc, cos_desc, sin_desc, q_fp8_row_desc,
q_fp8_col_desc, q_scales_row_desc, and q_scales_col_desc so they match the
hardcoded NUM_HEADS * HEAD_DIM and QK_ROPE assumptions, are contiguous where
required, and all tensors are on the current CUDA device before launch.
- Around line 209-216: The cache key in the GEMM rotary projection path is
missing several specialization inputs, so a cached kernel can be reused for
incompatible tensors and bypass support checks. Update the cache_key
construction in the relevant api function to include all
tensor/layout/device-sensitive fields used by the specialization, especially the
sin dtype, rotary-table shapes, and any other fields that affect kernel
selection alongside x, w, cos, and w_out_in.
In `@python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py`:
- Around line 592-629: The FE API test suite does not cover
gemm_proj_rope_mxfp8_reference, so add a test under test/python/fe_api that
exercises this oracle through the wrapper in both w_out_in modes and with token
counts that are valid multiples of BLOCK. Use the gemm_proj_rope_mxfp8_reference
function and its corresponding FE API entrypoint to verify the contract for the
row/column outputs and scale shapes, and include cases that match the supported
token-block alignment.
- Around line 553-586: The exported run() entrypoint currently bypasses the
wrapper’s input validation, so direct callers can hit the (_cache key based on
tokens and w_out_in) with tensors that do not match the expected shape or dtype
specialization. Add the same shape/dtype/layout checks from the wrapper into
run() itself, or make those checks internal and always enforced before
from_dlpack() and cutlass.compile() are reached, using the run() and _cache
symbols to keep the specialization safe.
---
Nitpick comments:
In `@python/cudnn/gemm_proj_rope_mxfp8/api.py`:
- Around line 187-195: The public wrapper docstring in the allocate/compile/run
entrypoint should document the full contract, not just the basic tensor shapes.
Update the docstring for the wrapper that takes x, w, cos, sin, w_out_in, and
stream to explicitly mention the SM100-only requirement, the required dtypes for
w/cos/sin, any contiguity and alignment constraints, and the exact output scale
tensor shapes. Keep the wording aligned with the existing API symbols so callers
can understand the supported inputs and outputs without reading the
implementation.
🪄 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: f235a440-f409-48b6-911c-07f3bc9270e3
📒 Files selected for processing (3)
python/cudnn/gemm_proj_rope_mxfp8/__init__.pypython/cudnn/gemm_proj_rope_mxfp8/api.pypython/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
test/python/fe_api/test_gemm_proj_rope_mxfp8.py (1)
56-59: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueStray string literals used as section headers.
These bare strings aren't docstrings (not the first statement of a def/class/module) and have no runtime effect beyond acting as no-op expression statements. Regular
#comments would be more idiomatic for section separators.Also applies to: 102-104
🤖 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 `@test/python/fe_api/test_gemm_proj_rope_mxfp8.py` around lines 56 - 59, Replace the stray bare string literals in the test module with regular section comments; they are not valid docstrings in this location and only create no-op expressions. Update the affected separator text in the file around the GemmProjRopeMxfp8 test section, and also the other noted block, so the intent remains clear without using standalone string literals.
🤖 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/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py`:
- Around line 545-557: The current validation in run() only checks tokens,
projected output dim, and rope tensor shapes, but not the GEMM K dimension from
x.shape[1], so a bad input can still reach the fused kernel and fail opaquely.
Add an explicit K-dimension check in run() using the same expected K derived in
check_support()/kernel setup, and mirror that validation in check_support() so
both entry paths reject mismatches early with a clear ValueError. Use the
existing run and check_support methods in gemm_proj_rope_mxfp8.py to keep the
validation logic consistent.
- Around line 574-586: The run() entrypoint currently uses the active CUDA
device for current_stream() and HardwareInfo() while only caching by
x.device.index, so it can silently mix devices. Update run() to bind execution
with torch.cuda.device(x.device) before any stream or hardware queries, and add
a fast validation that all input tensors/devices match x.device. Keep the cache
key logic in run() and ensure the compiled function lookup and launch are
performed under the correct device context.
In `@test/python/fe_api/test_gemm_proj_rope_mxfp8_utils.py`:
- Around line 37-38: The skip message in the compute capability gate is
reporting the wrong value, which can mislead debugging. Update the pytest.skip
call in the gating logic to reference the actual compute capability check used
by this test, not just the major version, so the message matches the
`compute_capability < 100` condition in `test_gemm_proj_rope_mxfp8_utils`.
Ensure the displayed unsupported capability value reflects the real
`major*10+minor` result.
---
Nitpick comments:
In `@test/python/fe_api/test_gemm_proj_rope_mxfp8.py`:
- Around line 56-59: Replace the stray bare string literals in the test module
with regular section comments; they are not valid docstrings in this location
and only create no-op expressions. Update the affected separator text in the
file around the GemmProjRopeMxfp8 test section, and also the other noted block,
so the intent remains clear without using standalone string literals.
🪄 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: 1cca45f1-1c77-4d4a-ae23-5e93978fc329
📒 Files selected for processing (8)
docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.mddocs/fe-oss-apis/overview.mdpython/cudnn/__init__.pypython/cudnn/gemm_proj_rope_mxfp8/__init__.pypython/cudnn/gemm_proj_rope_mxfp8/api.pypython/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.pytest/python/fe_api/test_gemm_proj_rope_mxfp8.pytest/python/fe_api/test_gemm_proj_rope_mxfp8_utils.py
✅ Files skipped from review due to trivial changes (3)
- docs/fe-oss-apis/overview.md
- python/cudnn/gemm_proj_rope_mxfp8/init.py
- docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cudnn/gemm_proj_rope_mxfp8/api.py
Comment on lines
+545
to
+557
| tokens = x.shape[0] | ||
| proj_dim = w.shape[0] if w_out_in else w.shape[1] | ||
| # Lightweight structural validation (run() is a public entry that skips check_support); | ||
| # catch the common misuses with a clear error instead of an opaque CUTLASS/CUDA crash. | ||
| if tokens % TILE_M != 0: | ||
| raise ValueError(f"tokens ({tokens}) must be a multiple of TILE_M ({TILE_M})") | ||
| if proj_dim % HEAD_DIM != 0: | ||
| raise ValueError(f"weight projected dim ({proj_dim}) must be a multiple of HEAD_DIM ({HEAD_DIM})") | ||
| if tuple(cos.shape) != (tokens, QK_ROPE) or tuple(sin.shape) != (tokens, QK_ROPE): | ||
| raise ValueError(f"cos/sin must both be ({tokens}, {QK_ROPE}); got {tuple(cos.shape)}, {tuple(sin.shape)}") | ||
| # Number of heads is derived from the weight's projected dimension (compile-time Constexpr); | ||
| # HEAD_DIM is the fixed per-head width the epilogue is specialized for. | ||
| num_heads = proj_dim // HEAD_DIM |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the target file and inspect the relevant section with line numbers.
git ls-files python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py
wc -l python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py
sed -n '500,620p' python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py | cat -n
# Check for any validation helpers or shape checks elsewhere in the file/module.
rg -n "check_support|ValueError|shape\\[1\\]|ndim|tokens|proj_dim|k_dim|w_out_in" python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py
# Look for tests in the requested area.
git ls-files 'test/python/fe_api/**' | rg 'gemm_proj_rope_mxfp8|rope|mxfp8|cudnn'
rg -n "gemm_proj_rope_mxfp8|rope|mxfp8" test/python/fe_apiRepository: NVIDIA/cudnn-frontend
Length of output: 13855
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the wrapper and support path around run().
sed -n '1,260p' python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py | cat -n | sed -n '1,260p'
# Inspect the test helper that constructs shapes/layouts.
sed -n '1,220p' test/python/fe_api/test_gemm_proj_rope_mxfp8_utils.py | cat -n
# Inspect the actual tests for any negative cases around mismatched dimensions.
sed -n '1,220p' test/python/fe_api/test_gemm_proj_rope_mxfp8.py | cat -nRepository: NVIDIA/cudnn-frontend
Length of output: 23497
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search for any shared input-shape validation around this API.
rg -n "shape\\[1\\].*w|w.shape\\[[01]\\].*x.shape|must match weight input dim|x and w must both be 2D|check_support\\(|ValueError\\(" python/cudnn test/python/fe_api | sed -n '1,220p'
# Look for the public wrapper/object class and its support checks.
rg -n "class GemmProjRopeMxfp8Sm100|def check_support|def execute|def compile|def wrapper_sm100" python/cudnn -n
# If the class is generated or spread across files, inspect the surrounding symbols.
fd -a "gemm_proj_rope_mxfp8" python/cudnnRepository: NVIDIA/cudnn-frontend
Length of output: 41603
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the API class that calls into run().
sed -n '1,280p' python/cudnn/gemm_proj_rope_mxfp8/api.py | cat -n
sed -n '280,420p' python/cudnn/gemm_proj_rope_mxfp8/api.py | cat -n
# Check whether tests exercise mismatched K or rely on check_support for shape validation.
rg -n "mismatch|shape|K dim|input dim|check_support\\(|execute\\(" test/python/fe_api/test_gemm_proj_rope_mxfp8.py test/python/fe_api/test_gemm_proj_rope_mxfp8_utils.pyRepository: NVIDIA/cudnn-frontend
Length of output: 17277
Validate the GEMM K dimension in run() and check_support(). run() bypasses support checks, and the API path only validates the projected output dim; a mismatched x.shape[1] can still reach the fused kernel and fail opaquely.
🤖 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/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8.py` around lines 545 -
557, The current validation in run() only checks tokens, projected output dim,
and rope tensor shapes, but not the GEMM K dimension from x.shape[1], so a bad
input can still reach the fused kernel and fail opaquely. Add an explicit
K-dimension check in run() using the same expected K derived in
check_support()/kernel setup, and mirror that validation in check_support() so
both entry paths reject mismatches early with a clear ValueError. Use the
existing run and check_support methods in gemm_proj_rope_mxfp8.py to keep the
validation logic consistent.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@test/python/fe_api/test_gemm_proj_rope_mxfp8.py`:
- Around line 166-167: The current finiteness check in the GEMM/ROPE FP8 test
only verifies the raw FP8 payloads, so it can miss bad E8M0 scale handling;
update the assertions in test_gemm_proj_rope_mxfp8 to use the shared
dequantization helper on both qr/qc and their scale tensors, then assert
finiteness on the dequantized outputs rather than on qr.float() and qc.float()
alone.
🪄 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: 08a1452d-9cc8-4eb3-b756-fdab7bcd62e2
📒 Files selected for processing (1)
test/python/fe_api/test_gemm_proj_rope_mxfp8.py
saltyminty
reviewed
Jul 10, 2026
| self.tokens = int(sample_x.shape[0]) | ||
| # Heads derived from the weight's projected dim (compile-time Constexpr for the kernel); | ||
| # check_support() validates that this divides evenly. | ||
| proj_dim = int(sample_w.shape[0] if self.w_out_in else sample_w.shape[1]) |
Collaborator
There was a problem hiding this comment.
I don't believe we ever validate sample_w.shape[1] to be equal to be NUM_HEADS (or sample_w or other input tensors in general), is that something we need?
Contributor
Author
There was a problem hiding this comment.
That NUM_HEADS constant was actually not used. I've removed it now: it's inferred from the weight dimensions.
Collaborator
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-367-2c89b1e |
Anerudhan
approved these changes
Jul 10, 2026
5 tasks
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
4 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.
This PR introduces a new kernel: a fused GEMM+RoPE+MXFP8 operation intended to be used with in forward direction of DeepseekV3 training.
Summary by CodeRabbit