Add an mxfp8-input version of the gemm_proj_rope_mxfp8 kernel. - #438
Conversation
This also modifies the existing kernel (renaming it to _bf16in) and its API so that the correct version of the kernel is selected.
|
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 (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThe fused GEMM+RoPE+MXFP8 API now supports separate BF16-input and MXFP8-input SM100 paths. It adds dtype-based wrapper dispatch, a dedicated MXFP8 CUTLASS kernel, updated exports and validation, expanded tests, and documentation for both input formats. ChangesGEMM RoPE MXFP8 input variants
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant Wrapper
participant API
participant Kernel
Caller->>Wrapper: provide x, w, and optional scales
Wrapper->>API: dispatch by input dtype
API->>Kernel: compile and execute selected path
Kernel-->>Wrapper: produce row/column MXFP8 outputs and scales
Wrapper-->>Caller: return output tensors
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (7)
python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py (3)
642-644: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnused host-side computation.
acc_shape/tCtAcc_fakeare computed here but never used — the kernel recomputes them fromtiled_mma(Lines 280-281). Safe to delete.🤖 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_mxfp8in.py` around lines 642 - 644, Remove the unused host-side acc_shape and tCtAcc_fake computations near NUM_TMEM_COLS; the kernel already recomputes the required fragments from tiled_mma, so leave NUM_TMEM_COLS and the surrounding setup unchanged.
392-398: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead assignment.
tCtSFB_mma = tCtSFBon Line 392 is unconditionally overwritten on Line 398; drop it.🧹 Proposed cleanup
tCtAcc = tCtAcc_base[(None, None, None, acc_empty.index)] - tCtSFB_mma = tCtSFB offset = cutlass.Int32((n_idx % 2) * 2)🤖 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_mxfp8in.py` around lines 392 - 398, Remove the redundant initial assignment to tCtSFB_mma in the surrounding tensor setup; retain the later cute.make_tensor assignment using shifted_sfb_ptr as the sole initialization.
109-145: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicated inline-asm reciprocal snippet.
Lines 122-130 in
_e8m0are a byte-for-byte copy of_e8m0_inv; have_e8m0call it so the PTX exists in one place.♻️ Proposed dedup
sbyte32 = cutlass.Int32(packed_i16) & cutlass.Int32(0xFF) - inv_bits = _llvm_dialect.inline_asm( - _mlir_ir.F32Type.get(), - [_Int32Sym(sbyte32).ir_value()], - "{ .reg .s32 t; sub.s32 t, 254, $1; shl.b32 t, t, 23; mov.b32 $0, t; }", - "=f,r", - has_side_effects=False, - is_align_stack=False, - asm_dialect=_llvm_dialect.AsmDialect.AD_ATT, - ) - return cutlass.Float32(inv_bits), sbyte32 + return _e8m0_inv(sbyte32), sbyte32Note this requires moving
_e8m0_invabove_e8m0.🤖 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_mxfp8in.py` around lines 109 - 145, Move _e8m0_inv above _e8m0, then replace the duplicated inv_bits inline-assembly block in _e8m0 with a call to _e8m0_inv(sbyte32). Preserve _e8m0’s returned Float32 reciprocal value and sbyte32 unchanged, leaving the reciprocal PTX implementation solely in _e8m0_inv.test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8.py (2)
29-40: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider demoting the 4096-token MXFP8 sweep out of L0.
Both MXFP8-input tests compile a fresh SM100 kernel and run the full PyTorch reference oracle for
tokens=2048and4096. Keeping2048at L0 and moving4096to a higher level would keep the fast tier fast.As per coding guidelines: "Mark every new Python test with a level from
L0throughL4; keepL0tests fast and place large parameter sweeps at higher levels."🤖 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/gemm/test_gemm_proj_rope_mxfp8.py` around lines 29 - 40, The parameterized MXFP8-input tests test_gemm_proj_rope_mxfp8_mxfp8in_compile_execute and test_gemm_proj_rope_mxfp8_mxfp8in_wrapper should keep the 2048-token case at L0 while assigning the 4096-token case to an appropriate higher test level, preserving both existing test variants.Source: Coding guidelines
149-178: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winWrapper contract test: consider asserting on the message, not just
AssertionError.All three cases only assert the exception type. Since
AssertionErroris also raised by unrelated internal asserts on the same call path,pytest.raises(AssertionError, match=...)would pin each case to its intended contract (dtype agreement, missing scales, scales-with-bf16).🤖 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/gemm/test_gemm_proj_rope_mxfp8.py` around lines 149 - 178, The test test_wrapper_rejects_dtype_mismatch_and_missing_scales should validate each wrapper contract via the AssertionError message, not only its type. Add distinct pytest.raises match patterns for dtype disagreement, missing FP8 scales, and scales supplied with BF16 inputs, while preserving the existing inputs and calls.python/cudnn/__init__.py (1)
319-320: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemoved
GemmProjRopeMxfp8Sm100is a breaking export change.The lazy table no longer resolves the old name, so
cudnn.GemmProjRopeMxfp8Sm100now raisesAttributeError. Intentional per the PR objectives and the API is experimental; consider a temporary alias entry pointing at the bf16-input class if any downstream pins the old name.🤖 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/__init__.py` around lines 319 - 320, Add a temporary lazy-export alias for GemmProjRopeMxfp8Sm100 in the export table, resolving it to the existing GemmProjRopeMxfp8Bf16InSm100 class so cudnn.GemmProjRopeMxfp8Sm100 remains available while preserving the current exports.python/cudnn/gemm_proj_rope_mxfp8/api.py (1)
449-450: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueModule-level caches grow without bound.
_bf16in_obj_cache/_mxfp8in_obj_cacheretain compiled kernels for every distinct (shape, device) key for the process lifetime. Fine for fixed training shapes, but worth a size bound or eviction if callers vary token counts. Consistent with sibling packages, so purely optional.🤖 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 449 - 450, Bound the module-level _bf16in_obj_cache and _mxfp8in_obj_cache so compiled kernels do not accumulate indefinitely across distinct shape/device keys. Add a bounded-cache or eviction policy consistent with sibling packages, while preserving cache reuse for recent entries.
🤖 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 `@docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md`:
- Around line 23-27: Document the w_scale orientation consistently at both
affected sites in docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md: lines
23-27 and 118-131. Update the w_scale shape to state (Q_LORA // BLOCK,
NUM_HEADS·HEAD_DIM) when w_out_in=False and (NUM_HEADS·HEAD_DIM, Q_LORA //
BLOCK) when w_out_in=True, matching the wrapper’s transpose behavior and the
existing w documentation.
In `@python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py`:
- Around line 482-484: Update the address calculations in the mxfp8 input/output
path around cs_row_bytes, row_bytes, cos_base, sin_base, and the corresponding
mQrow/mQcol pointers so they use actual tensor strides rather than assuming
contiguous layouts. Alternatively, extend check_support() in api.py to reject
non-contiguous cos, sin, and output tensors before mark_layout_dynamic; apply
the same correction to the later duplicated calculations.
- Around line 65-67: Update the MXFP8 path to prevent a hardcoded K=1536 scale
stride: either derive K_SCALE_WORDS and the _relay_scale_words strides from the
actual K dimension/scale tensor shape, or add an explicit k_dim == 1536
requirement in the MXFP8 check_support validation. Ensure unsupported K
dimensions are rejected before compilation and scale indexing remains correct.
- Around line 204-217: Update the MXFP8 support validation in check_support to
require an even num_heads value, preventing the SFB relay path from processing
an out-of-range final head; preserve the existing sfb_row_base and relay-copy
logic in the kernel.
---
Nitpick comments:
In `@python/cudnn/__init__.py`:
- Around line 319-320: Add a temporary lazy-export alias for
GemmProjRopeMxfp8Sm100 in the export table, resolving it to the existing
GemmProjRopeMxfp8Bf16InSm100 class so cudnn.GemmProjRopeMxfp8Sm100 remains
available while preserving the current exports.
In `@python/cudnn/gemm_proj_rope_mxfp8/api.py`:
- Around line 449-450: Bound the module-level _bf16in_obj_cache and
_mxfp8in_obj_cache so compiled kernels do not accumulate indefinitely across
distinct shape/device keys. Add a bounded-cache or eviction policy consistent
with sibling packages, while preserving cache reuse for recent entries.
In `@python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py`:
- Around line 642-644: Remove the unused host-side acc_shape and tCtAcc_fake
computations near NUM_TMEM_COLS; the kernel already recomputes the required
fragments from tiled_mma, so leave NUM_TMEM_COLS and the surrounding setup
unchanged.
- Around line 392-398: Remove the redundant initial assignment to tCtSFB_mma in
the surrounding tensor setup; retain the later cute.make_tensor assignment using
shifted_sfb_ptr as the sole initialization.
- Around line 109-145: Move _e8m0_inv above _e8m0, then replace the duplicated
inv_bits inline-assembly block in _e8m0 with a call to _e8m0_inv(sbyte32).
Preserve _e8m0’s returned Float32 reciprocal value and sbyte32 unchanged,
leaving the reciprocal PTX implementation solely in _e8m0_inv.
In `@test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8.py`:
- Around line 29-40: The parameterized MXFP8-input tests
test_gemm_proj_rope_mxfp8_mxfp8in_compile_execute and
test_gemm_proj_rope_mxfp8_mxfp8in_wrapper should keep the 2048-token case at L0
while assigning the 4096-token case to an appropriate higher test level,
preserving both existing test variants.
- Around line 149-178: The test
test_wrapper_rejects_dtype_mismatch_and_missing_scales should validate each
wrapper contract via the AssertionError message, not only its type. Add distinct
pytest.raises match patterns for dtype disagreement, missing FP8 scales, and
scales supplied with BF16 inputs, while preserving the existing inputs and
calls.
🪄 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: 48d1b477-252a-4344-a1d3-941447b55f04
📒 Files selected for processing (8)
docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.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_bf16in.pypython/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.pytest/python/fe_api/gemm/test_gemm_proj_rope_mxfp8.pytest/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_utils.py
💤 Files with no reviewable changes (1)
- python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8_bf16in.py
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md (2)
87-114: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winImport the classes used by the examples.
Both class API snippets reference
GemmProjRopeMxfp8Bf16InSm100andGemmProjRopeMxfp8Mxfp8InSm100without importing them, so copying either example raisesNameError. Add imports from the documented public package exports.🤖 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 `@docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md` around lines 87 - 114, Add imports for GemmProjRopeMxfp8Bf16InSm100 and GemmProjRopeMxfp8Mxfp8InSm100 from the documented public package exports before their respective class API examples, so both snippets can run without NameError.
102-113: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winMake the MXFP8 class layout contract explicit.
This direct class example omits
w_out_in, while its comment claims TE-native[out, in]weights. The documented default at Line 130 applies to the wrapper, not clearly to the class API. Passw_out_in=Trueexplicitly or document the class default and expectedw_code/w_scalelayout to prevent callers from supplying a shape the class interprets incorrectly.🤖 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 `@docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md` around lines 102 - 113, The MXFP8 class example must make the weight-layout contract explicit. Update the GemmProjRopeMxfp8Mxfp8InSm100 constructor example to pass w_out_in=True when using TE-native [out, in] weights, or clearly document the class default and the expected w_code/w_scale layout; ensure the surrounding comment matches that behavior.
🤖 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.
Outside diff comments:
In `@docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md`:
- Around line 87-114: Add imports for GemmProjRopeMxfp8Bf16InSm100 and
GemmProjRopeMxfp8Mxfp8InSm100 from the documented public package exports before
their respective class API examples, so both snippets can run without NameError.
- Around line 102-113: The MXFP8 class example must make the weight-layout
contract explicit. Update the GemmProjRopeMxfp8Mxfp8InSm100 constructor example
to pass w_out_in=True when using TE-native [out, in] weights, or clearly
document the class default and the expected w_code/w_scale layout; ensure the
surrounding comment matches that behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7279ce38-08a2-4ed8-a7eb-2e17d1d56a47
📒 Files selected for processing (4)
docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.mdpython/cudnn/gemm_proj_rope_mxfp8/api.pypython/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.pytest/python/fe_api/gemm/test_gemm_proj_rope_mxfp8.py
🚧 Files skipped from review as they are similar to previous changes (3)
- python/cudnn/gemm_proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py
- test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8.py
- python/cudnn/gemm_proj_rope_mxfp8/api.py
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-438-d8ff4d0 |
Anerudhan
added
orig-nv-eng
3 participants
This also modifies the existing kernel (renaming it to _bf16in) and its API so that the correct version of the kernel is selected.
Before submitting
pre-commit runand committed any formatting changes.Affected area
Summary
Updating the gemm_proj_rope_mxfp8 with a second kernel specialized for mxfp8 inputs (instead of bf16).
Why
Some training recipes perform the projections in bf16, and some in mxfp8. We want to support fusion for both.
API and compatibility impact
This does change the name of the existing kernel (for better disambiguation), but the wrapper name remains the same.
Summary by CodeRabbit
x/wdtype, with MXFP8 requiringx_scale/w_scale, and producing dual-direction MXFP8 outputs.