[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI - #3137
Open
kainzhong wants to merge 96 commits into
Open
[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI#3137kainzhong wants to merge 96 commits into
kainzhong wants to merge 96 commits into
Conversation
kainzhong
force-pushed
the
cutedsl_mxfp8_common
branch
from
June 27, 2026 08:27
daee750 to
218cd24
Compare
kainzhong
force-pushed
the
cutedsl_mxfp8_common
branch
4 times, most recently
from
July 8, 2026 21:44
2f8c8da to
8448930
Compare
Collaborator
Author
|
Benchmark on ptyche (B200 GPU, ARM CPU): |
kainzhong
marked this pull request as ready for review
July 8, 2026 23:25
Contributor
Greptile SummaryThis PR introduces an experimental, opt-in CuTeDSL backend for MXFP8 quantization that JIT-compiles Python kernels (via
Confidence Score: 4/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant PY as Python __init__.py
participant TVM as tvm_ffi registry
participant CPP as C++ nvte_quantize
participant CACHE as TVMFFIConfigCache
participant JIT as cute.compile()
participant CUDA as CUDA C++ fallback
Note over PY: NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=1
PY->>TVM: register_global_func("get_mxfp8_quantization_function", ...)
PY->>TVM: load_tvm_ffi_library() via dlopen
CPP->>CPP: "mxfp8_quantize_cutedsl<IS_DBIAS,...>()"
CPP->>CACHE: config.get_kernel() → cache.get_or_load(config)
CACHE->>CACHE: shared_lock: check map_[id]
alt cache miss
CACHE->>TVM: Function::GetGlobal("get_mxfp8_quantization_function")
CACHE->>TVM: entrypoint(fn_name, dtype, fp8_dtype, rowwise, ...)
TVM->>JIT: compile_cutedsl_function_from_cfg(cfg)
JIT-->>TVM: cute.compile() → tvm_ffi registered fn
TVM-->>CACHE: returns True / False
CACHE->>CACHE: unique_lock: map_.emplace(id, fn_opt)
end
CACHE-->>CPP: "optional<Function>"
alt kernel available
CPP->>TVM: "(*fn)(mX, mO_row, mS_row, mO_col, mS_col, mAmax, noop, mActInput, mWorkspace, stream)"
TVM-->>CPP: kernel runs on GPU
opt WITH_DBIAS
CPP->>CUDA: reduce_dbias(workspace → dbias)
end
else no kernel (unsupported shape/dtype/device)
CPP->>CUDA: return false → CUDA C++ kernel path
end
Reviews (48): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
ptrendx
reviewed
Jul 9, 2026
| "importlib-metadata>=1.0", | ||
| "packaging", | ||
| "apache-tvm-ffi>=0.1.12", | ||
| "nvidia-cutlass-dsl>=4.2.0", |
Member
There was a problem hiding this comment.
Due to other things (like cudnn frontend CuTeDSL kernels), I'm pretty sure we need a later version
of that package (4.4.2 I think?). Adding @ksivaman to comment.
Collaborator
Author
There was a problem hiding this comment.
I'll change this to 4.4.2
Comment on lines
+21
to
+22
| GTEST_SKIPs the mismatched half), non-32-divisible shapes are omitted (the | ||
| dispatcher can never route them to CuTeDSL), and a missing kernel registration |
Member
There was a problem hiding this comment.
Why are the non-32-divisible shapes omitted? Is this a limitation of the cutedsl implementation?
Collaborator
Author
There was a problem hiding this comment.
Because my CuTeDSL kernels are compiled with
sym_M = cute.sym_int32(divisibility=MXFP8_BLOCK_SCALING_SIZE)
sym_N = cute.sym_int32(divisibility=MXFP8_BLOCK_SCALING_SIZE)
So it assumes 32-divisible shape. Maybe non-32-divisible can be supported as well. I'll run some benchmarks and see if it hurts performance but I think normally people wouldn't use these weird shapes?
| # | ||
| # See LICENSE for license information. | ||
|
|
||
| """Cross-backend bit-exactness tests for the CuTeDSL MXFP8 quantize kernels. |
Member
There was a problem hiding this comment.
This makes sense in this initial stage, but I would explicitly mark this file as temporary, since
ultimately we will want to standardize on this backend.
Collaborator
Author
There was a problem hiding this comment.
I could port the CUDA C++ tests to python and make this a standalone test instead of comparing with CUDA kernel's output, but then I thought since we already validated CUDA implementation it would be easier to just make that the reference and compare the result instead.
If we want to standardize on this then should this be python MXFP8 reference implementation on its own?
Comment on lines
+44
to
+70
| std::string to_key() const { | ||
| std::string key; | ||
| key.reserve(56); | ||
| key.append("cutedsl_mxfp8_") | ||
| .append(te_dtype_to_str(dtype)) | ||
| .append("_") | ||
| .append(te_dtype_to_str(fp8_dtype)) | ||
| .append("_") | ||
| .append(rowwise ? "1" : "0") | ||
| .append("_") | ||
| .append(colwise ? "1" : "0") | ||
| .append("_") | ||
| .append(swizzled ? "1" : "0") | ||
| .append("_") | ||
| .append(with_amax ? "1" : "0") | ||
| .append("_") | ||
| .append(with_dbias ? "1" : "0") | ||
| .append("_") | ||
| .append(with_dact ? "1" : "0") | ||
| .append("_") | ||
| .append(with_act ? "1" : "0") | ||
| .append("_") | ||
| .append(with_noop ? "1" : "0") | ||
| .append("_") | ||
| .append(activation_to_str(activation)); | ||
| return key; | ||
| } |
Member
There was a problem hiding this comment.
Kind of random, but this function is quite slow. You could do the same much faster with raw char*
manipulation.
Collaborator
Author
There was a problem hiding this comment.
Emmmm but I reserved 56 chars before I do append. I don't know if char* will be faster than this since they both don't require resizing the string?
Member
There was a problem hiding this comment.
You still need to create those additional 1-letter strings in this version. At the very least you could make a "1_" and "0_" strings upfront and use those instead (also, you don't even need the underscore there between those 1s and 0s).
Collaborator
Author
There was a problem hiding this comment.
Ah OK I just made some changes. Now every quantization config owns their cache and the cache key is uint32 now. This to_key is now only used to build the function name used when registering the function to TVM-FFI registry and it happens only once when you request a not yet ready kernel. Later we will fetch it from C++ cache with uint32 cache key which is more efficient.
kainzhong
force-pushed
the
cutedsl_mxfp8_common
branch
from
July 10, 2026 00:34
948fab5 to
2930b1b
Compare
13 tasks
kainzhong
force-pushed
the
cutedsl_mxfp8_common
branch
from
July 24, 2026 18:33
af55445 to
87adfe7
Compare
kainzhong
force-pushed
the
cutedsl_mxfp8_common
branch
from
July 27, 2026 20:09
6e55eef to
c47fc5c
Compare
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>
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>
for more information, see https://pre-commit.ci
Collaborator
Author
noop flag won't fire with dbias. |
Contributor
|
You're right. With The question is what the intended noop semantics are when
The CUDA C++ reference path checks the noop tensor before dbias-related work (as the earlier resolved Tip: You can customize Greptile's behavior for this repo with |
Collaborator
Author
|
@greptileai in the latest main branch noop is ignored with dbias already: |
Contributor
|
You're right. Looking at the code:
The behaviors are aligned. My concern was unfounded — no issue here. |
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
kainzhong
commented
Aug 21, 2026
| # or if the shape is already divisible by the tile size (so no masking is needed) | ||
| if cutlass.const_expr(self.ACT_NEED_MASKING): | ||
| skip_masking = ( | ||
| mX.shape[0] % self.general_divisible_kernel._TILE_ROWS == 0 |
Collaborator
Author
There was a problem hiding this comment.
Don't need to check both directions if we are not quantizing both directions. Only check the direction you quantize
ptrendx
reviewed
Aug 21, 2026
Comment on lines
+105
to
+121
| option(NVTE_WITH_CUTEDSL | ||
| "Build the CuTeDSL quantize backend bridge (requires the apache-tvm-ffi Python package)" ON) | ||
| if(NVTE_WITH_CUTEDSL) | ||
| execute_process( | ||
| COMMAND ${Python_EXECUTABLE} -c "import tvm_ffi.libinfo as li; print(li.find_include_path())" | ||
| OUTPUT_VARIABLE TVM_FFI_INCLUDE_DIR | ||
| OUTPUT_STRIP_TRAILING_WHITESPACE | ||
| RESULT_VARIABLE TVM_FFI_INCLUDE_QUERY) | ||
| if(NOT TVM_FFI_INCLUDE_QUERY EQUAL 0) | ||
| message(FATAL_ERROR | ||
| "Could not import the tvm_ffi Python package (with '${Python_EXECUTABLE}'), " | ||
| "whose headers Transformer Engine needs to compile the CuTeDSL quantize " | ||
| "backend bridge (common/tvm_ffi_bridge.h). Install it into this Python " | ||
| "environment: `pip install apache-tvm-ffi`, or configure with " | ||
| "-DNVTE_WITH_CUTEDSL=OFF to build the C++ library without the backend.") | ||
| endif() | ||
| endif() |
Member
There was a problem hiding this comment.
@fheinecke Could you review this part (+pyproject/setup.py changes) - will this work under build isolation?
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci
Comment on lines
+105
to
+120
| option(NVTE_WITH_CUTEDSL | ||
| "Build the CuTeDSL quantize backend bridge (requires the apache-tvm-ffi Python package)" ON) | ||
| if(NVTE_WITH_CUTEDSL) | ||
| execute_process( | ||
| COMMAND ${Python_EXECUTABLE} -c "import tvm_ffi.libinfo as li; print(li.find_include_path())" | ||
| OUTPUT_VARIABLE TVM_FFI_INCLUDE_DIR | ||
| OUTPUT_STRIP_TRAILING_WHITESPACE | ||
| RESULT_VARIABLE TVM_FFI_INCLUDE_QUERY) | ||
| if(NOT TVM_FFI_INCLUDE_QUERY EQUAL 0) | ||
| message(FATAL_ERROR | ||
| "Could not import the tvm_ffi Python package (with '${Python_EXECUTABLE}'), " | ||
| "whose headers Transformer Engine needs to compile the CuTeDSL quantize " | ||
| "backend bridge (common/tvm_ffi_bridge.h). Install it into this Python " | ||
| "environment: `pip install apache-tvm-ffi`, or configure with " | ||
| "-DNVTE_WITH_CUTEDSL=OFF to build the C++ library without the backend.") | ||
| endif() |
Contributor
There was a problem hiding this comment.
NVTE_WITH_CUTEDSL=ON default breaks builds without apache-tvm-ffi
The option advertises itself as opt-in but defaults to ON, and the failure path is a hard FATAL_ERROR rather than a graceful downgrade. Any existing CMake-based CI pipeline or development environment that does not have apache-tvm-ffi installed will now fail at configure time with no way to proceed without either installing the new package or explicitly passing -DNVTE_WITH_CUTEDSL=OFF. Defaulting the option to OFF — or downgrading the missing-package diagnostic to a WARNING that automatically sets NVTE_WITH_CUTEDSL OFF — would preserve backward compatibility for builds that don't need the CuTeDSL backend.
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.
Description
Adds an experimental, opt-in CuTeDSL backend for MXFP8 quantization. MXFP8 nvte_quantize calls can be routed to JIT-compiled CuTeDSL (CUTLASS Python DSL) kernels instead of the existing CUDA C++ kernels, bridged into the C++ dispatcher via apache-tvm-ffi (https://github.com/apache/tvm-ffi).
It's off by default (use
NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=1to enable) and transparently falls back to the CUDA kernels for any unsupported config or shape (useNVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN=1to enable warning for unsupported cases).How it works
TODO:
Type of change
Changes
Breaking changes:
NVTE_WITH_CUTEDSLto cmake to opt-out.Checklist: