[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI - #3137
Open
kainzhong wants to merge 65 commits into
Open
[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI#3137kainzhong wants to merge 65 commits into
kainzhong wants to merge 65 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 (CUTLASS Python DSL) backend for MXFP8 quantization, bridged into the existing C++ dispatcher via
Confidence Score: 3/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant User as Python/C++ caller
participant Disp as quantize_fwd/bwd_helper
participant CuTeC as mxfp8_quantize_cutedsl (C++)
participant Central as TVMFFICentral
participant Cache as TVMFFIConfigCache
participant PyFFI as TVM-FFI (Python bridge)
participant JIT as CuTeDSL JIT (get_mxfp8_quantization_function)
participant CUDA as mxfp8::quantize (CUDA C++)
User->>Disp: nvte_quantize(input, output, ...)
Disp->>CuTeC: "mxfp8_quantize_cutedsl<...>(tensors, stream)"
CuTeC->>Central: get_cutedsl_backend_enabled()?
alt backend disabled
Central-->>CuTeC: false
CuTeC-->>Disp: false
Disp->>CUDA: mxfp8::quantize(...)
else backend enabled
CuTeC->>Cache: get_or_load(MXFP8QuantConfig)
Cache->>Central: load_tvm_ffi_function(cfg)
Central->>PyFFI: call "get_mxfp8_quantization_function"(key, params)
PyFFI->>JIT: compile_cutedsl_function_from_cfg(cfg)
alt compilation succeeds
JIT-->>PyFFI: compiled kernel
PyFFI->>PyFFI: register_global_func(key, compiled)
PyFFI-->>Central: True
Central->>PyFFI: GetGlobal(key) → Function
Central-->>Cache: "std::optional<Function>"
Cache-->>CuTeC: mxfp8_quant_func
CuTeC->>PyFFI: "(*func)(DLTensors, stream)"
PyFFI-->>CuTeC: kernel done
CuTeC-->>Disp: true (success)
else compilation fails / unsupported config
JIT-->>PyFFI: False (caught by try/except)
PyFFI-->>Central: False
Central-->>Cache: std::nullopt
Cache-->>CuTeC: nullopt
CuTeC-->>Disp: false
Disp->>CUDA: mxfp8::quantize(...)
end
end
Reviews (33): Last reviewed commit: "Merge branch 'main' into cutedsl_mxfp8_c..." | 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
kainzhong
commented
Aug 14, 2026
| kernel_obj = kernel_class(cfg) | ||
| # M, N must be divisible by the MXFP8 scale-block size (MXFP8_BLOCK_SCALING_SIZE = 32) — the | ||
| # same alignment the CUDA C++ kernel requires. | ||
| sym_M = cute.sym_int32(divisibility=MXFP8_BLOCK_SCALING_SIZE) |
Collaborator
Author
There was a problem hiding this comment.
Try remove this limitation. We pad SF so 128 & 4 divisibility is fine but we shouldn't require it for input.
See if TMA's padding can solve this issue for you (but beware of when there's activation you'll nned to do boundary check
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:
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: