Skip to content

[Common] Add gated MXFP8 CuTeDSL kernel - #3840

Open
kainzhong wants to merge 50 commits into
NVIDIA:mainfrom
kainzhong:gated_mxfp8_cutedsl
Open

kainzhong wants to merge 50 commits into
NVIDIA:mainfrom
kainzhong:gated_mxfp8_cutedsl

Conversation

@kainzhong

Copy link
Copy Markdown
Collaborator

Description

Add CuTeDSL port of CUDA's MXFP8 gated quantization kernel

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

  • Add CuTeDSL MXFP8 gated kernel

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

kainzhong and others added 30 commits October 7, 2026 07:22
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
kainzhong and others added 14 commits October 19, 2026 21:23
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Conflicts arose because the base branch was squash-merged as NVIDIA#3642.
Resolved by 3-way merging against the pre-gated tip (bfdb240) instead
of the stale fork point. CuTeDSL/__init__.py: keep the entrypoint loop
and register get_mxfp8_requantization_function from NVIDIA#3560 in it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
@greptile-apps

greptile-apps Bot commented Nov 9, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerView in GreptileConfidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable new issues were established.

Summary

Adds a CuTeDSL backend for forward and backward gated MXFP8 quantization.

  • MXFP8 gated activations can run through CuTeDSL.

No actionable new issues were established. GPU tests were not run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Gated MXFP8 call] --> B{CuTeDSL enabled and supported?}
    B -->|No| C[CUDA kernel]
    B -->|Yes| D[Load or compile CuTeDSL kernel]
    D --> E{Kernel available?}
    E -->|No| C
    E -->|Yes| F[Apply gated activation]
    F --> G[Quantize values and write scales]
Loading

Reviews (6) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." · Reviewed by Greptile

maybe_warn_cutedsl_not_chosen was renamed to report_cutedsl_not_chosen in NVIDIA#3746, so
this file no longer compiled with NVTE_WITH_CUTEDSL.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
ptrendx
ptrendx previously approved these changes Nov 11, 2026
@ptrendx

ptrendx commented Nov 11, 2026

Copy link
Copy Markdown
Member

/te-ci

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
kainzhong and others added 2 commits November 13, 2026 22:38
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

# Conflicts:
#	tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py
#	transformer_engine/common/CuTeDSL/__init__.py

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants