[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI - #3137
Open
kainzhong wants to merge 64 commits into
Open
[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI#3137kainzhong wants to merge 64 commits into
kainzhong wants to merge 64 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 adds an experimental, opt-in CuTeDSL backend for MXFP8 quantization, bridged into the C++ dispatcher via apache-tvm-ffi. Python CuTeDSL kernels register as TVM-FFI globals, are JIT-compiled on first use (and cached), and are invoked via
Confidence Score: 4/5The experimental CuTeDSL path is cleanly additive and falls back to established CUDA kernels on any miss, so non-CuTeDSL code paths are unaffected. However, several correctness interactions (noop + dbias reduction, multi-GPU device targeting) remain under active discussion across review rounds and have not all been resolved. The dispatch logic, JIT cache, and Python-side fallbacks are carefully structured and correctly ordered (workspace query now precedes kernel loading). Several substantive concerns from prior review rounds still appear unresolved in the changed files, particularly around noop semantics with dbias workspace and hard-wiring to device 0 in multi-GPU environments. Files Needing Attention: transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh (noop + reduce_dbias interaction, workspace-query ownership), transformer_engine/common/CuTeDSL/cast/mxfp8/quantize_mxfp8.py (device_compute_capability hardcoded to device 0), setup.py and CMakeLists.txt (unconditional hard dependencies) Important Files Changed
Sequence DiagramsequenceDiagram
participant User as User
participant Init as common/__init__.py
participant Dispatcher as quantize.cuh
participant CuTeDSL_C as mxfp8_quantize_cutedsl
participant Cache as TVMFFIConfigCache
participant Python as get_mxfp8_quantization_function
participant CUDA as mxfp8::quantize
User->>Init: import transformer_engine
Init->>Init: "_load_tvm_ffi_library() [if NVTE_ENABLE_CUTEDSL=1]"
Init->>Init: _register_cutedsl_backends()
User->>Dispatcher: nvte_quantize()
Dispatcher->>CuTeDSL_C: mxfp8_quantize_cutedsl(...)
alt Backend disabled or shape not 32-aligned
CuTeDSL_C-->>Dispatcher: false
Dispatcher->>CUDA: mxfp8::quantize(...)
else Backend enabled, shape OK
CuTeDSL_C->>Cache: get_or_load(config)
alt Cache hit
Cache-->>CuTeDSL_C: tvm::ffi::Function
else Cache miss
Cache->>Python: retrieve_func_from_python(key)
Python->>Python: compile_cutedsl_function_from_cfg [NVCC JIT]
Python->>Python: tvm_ffi.register_global_func(key, compiled)
Python-->>Cache: true
Cache-->>CuTeDSL_C: tvm::ffi::Function
end
CuTeDSL_C->>CuTeDSL_C: "invoke (*fn)(&mX, &mO_row, ...)"
CuTeDSL_C-->>Dispatcher: true
end
Reviews (32): 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
Comment on lines
+1
to
+10
| # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # | ||
| # See LICENSE for license information. | ||
|
|
||
| """Cross-backend bit-exactness tests for the CuTeDSL MXFP8 quantize kernels.""" | ||
|
|
||
| import ctypes | ||
| import os | ||
| from typing import Callable, NamedTuple, Optional | ||
|
|
Contributor
There was a problem hiding this comment.
import tvm_ffi raises ImportError for users without the package
The test file imports tvm_ffi at module level before the pytestmark skip guard is evaluated. Pytest collects all test modules regardless of environment; users without apache-tvm-ffi installed will see a collection error instead of a clean skip. The import should be moved inside the test body or placed under a try/except ImportError guard that sets tvm_ffi_available = False, similar to how the test already conditionally sets cutedsl_enabled.
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>
Perform noop tensor check on device 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
force-pushed
the
cutedsl_mxfp8_common
branch
from
July 27, 2026 20:10
c47fc5c to
4815598
Compare
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
…ot dispatched 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>
kainzhong
force-pushed
the
cutedsl_mxfp8_common
branch
from
August 6, 2026 00:59
6090a71 to
ce9ef37
Compare
for more information, see https://pre-commit.ci
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
3 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: see if we can enable CuTeDSL kernel in the C++ MXFP8 tests and also test this on JAX
Type of change
Changes
Breaking changes:
NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=1. Otherwise they should be fine (pythonwouldn't load tvm-ffi library to make it available in C++ and register CuTeDSL kernels, and C++ wouldn't be able to usedlopento loadlibtvm_ffi.soloaded from python so it will fall back to CUDA kernels)Checklist: