Skip to content

[Common] Experimental CuTeDSL MXFP4 backend#3223

Draft
janekb04 wants to merge 25 commits into
NVIDIA:mainfrom
janekb04:cutedsl_nvfp4
Draft

[Common] Experimental CuTeDSL MXFP4 backend#3223
janekb04 wants to merge 25 commits into
NVIDIA:mainfrom
janekb04:cutedsl_nvfp4

Conversation

@janekb04

@janekb04 janekb04 commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Description

This PR adds a CuTe DSL NVFP4 quantization backend. It is based on the work in #3137.

In this PR only a CuTe DSL version of quantize_transpose_tuned_1D is implemented. It handles the case of 1D bfloat16 to NVFP4 quantization, mirroring the capabilities of the original CUDA implementation in quantize_transpose_nvfp4_tuned_1D.cuh.

To be rebased once #3137 merges.

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

Support for the CuTe DSL backend is implemented in a top-down manner.

  • Add the path to call the CuTE DSL backend to quantize.cuh (first commit, done)
  • Add the TVM FFI bridge scaffolding (second commit, done)
    • The C++ side of the TVM FFI bridge is implemented in quantize_transpose_nvfp4_cutedsl.cuh, analogically to quantize_mxfp8_cutedsl.cuh.
    • The Python side of the TVM FFI bridge is implemented in quantize_transpose_nvfp4.py, analogically to quantize_mxfp8.py.
  • Add the implementation of the CuTe DSL version of the kernel and fill in the TVM FFI scaffolding (todo)

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 24 commits July 10, 2026 00:31
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci

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>
for more information, see https://pre-commit.ci

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci

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>
for more information, see https://pre-commit.ci

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
…it__.py

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
@github-actions github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Jul 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants