Skip to content

[Common] Experimental CuTeDSL MXFP4 backend#3223

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

[Common] Experimental CuTeDSL MXFP4 backend#3223
janekb04 wants to merge 26 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, one commit at a time:

  • Add the CuTe DSL NVFP4 backend dispatch path to quantize.cuh (done)
    • Route the optimized 1D NVFP4 quantize-transpose path through the new CuTe DSL backend, falling back to the existing CUDA kernel when the backend declines to handle the case.
    • Add the FP4 dtype mapping to the TVM FFI bridge (tvm_ffi_bridge.h) and an empty stub header for the backend.
  • Add the TVM FFI bridge scaffolding (done)
    • Establish the general structure and file layout of the backend, copied from the MXFP8 backend.
    • The C++ side is scaffolded in quantize_transpose_nvfp4_cutedsl.cuh, mirroring quantize_mxfp8_cutedsl.cuh.
    • The Python side is scaffolded in quantize_transpose_nvfp4.py, mirroring quantize_mxfp8.py.
    • Kernel instantiation parameters and validation are left as todos.
  • Adapt the CuTe DSL scaffolding to the NVFP4 case (done)
    • Populate NVFP4QuantizeConfig with the kernel instantiation parameters (stochastic rounding, fast math, row-scaled NVFP4 and transpose flags) on both the C++ and Python sides, mirroring quantize_transpose_nvfp4.cuh:1441. The config values are derived from the quantization config and output tensor and forwarded across the TVM FFI boundary.
    • Add the input/output validation checks mirroring the CUDA implementation, guarding the backend behind FP4_TYPE_SUPPORTED.
  • 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 22 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
janekb04 added 2 commits July 22, 2026 20:39
Route the optimized 1D NVFP4 quantize-transpose path through the new
CuTe DSL backend, falling back to the existing CUDA kernel when the
CuTe DSL backend declines to handle the case. Adds the FP4 dtype
mapping to the TVM FFI bridge and an empty stub header for the backend.

Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Add the C++ and Python scaffolding for dispatching NVFP4
quantize-transpose to the CuTe DSL backend via TVM FFI, mirroring the
general structure and file layout of the MXFP8 backend
(quantize_mxfp8_cutedsl.cuh and quantize_mxfp8.py). Kernel
instantiation parameters and validation are left as todos.

Signed-off-by: Jan Bielak <jbielak@nvidia.com>
Populate NVFP4QuantizeConfig with the kernel instantiation parameters
(stochastic rounding, fast math, row-scaled NVFP4 and transpose flags)
on both the C++ and Python sides, and add the input/output validation
mirroring the CUDA implementation in quantize_transpose_nvfp4.cuh. The
config values are derived from the quantization config and output
tensor and forwarded across the TVM FFI boundary, and the backend is
guarded behind FP4_TYPE_SUPPORTED.

Signed-off-by: Jan Bielak <jbielak@nvidia.com>
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