[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI - #3137
[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI#3137kainzhong wants to merge 55 commits into
Conversation
daee750 to
218cd24
Compare
2f8c8da to
8448930
Compare
|
Benchmark on ptyche (B200 GPU, ARM CPU): |
Greptile SummaryThis PR adds an experimental, opt-in CuTeDSL (CUTLASS Python DSL) backend for MXFP8 quantization, bridged into the C++ dispatcher via Apache TVM-FFI. When
Confidence Score: 3/5The core fallback to existing CUDA kernels is safe, but the new CuTeDSL dispatch path has multiple behavioral issues — including unconditional reduce_dbias after a nooped kernel and swizzled-scale zeroing logic — that have been flagged across several review rounds without confirmed fixes in the current diff. The mandatory promotion of apache-tvm-ffi and nvidia-cutlass-dsl as hard install requirements is a significant breaking change for all TE users. Several issues flagged in prior rounds appear unaddressed in the current diff: the noop+dbias divergence from the CUDA path, the double-compilation race under concurrent first calls, and the no-opt-out mandatory dependency on tvm-ffi at both build and install time. The workspace-size query committing to a layout before confirming compilability is a new fragility. The CUDA fallback path itself is unchanged and sound, so production workloads not opting into the env var are unaffected. Files Needing Attention: quantize_mxfp8_cutedsl.cuh (noop+dbias semantics, workspace-size query ordering), tvm_ffi_bridge.h (concurrent compilation, permanent failure caching), setup.py and pyproject.toml (mandatory dependency promotion), CMakeLists.txt (no opt-out build flag) Important Files Changed
Sequence DiagramsequenceDiagram
participant Py as Python (import TE)
participant Init as common/__init__.py
participant Cpp as C++ quantize()
participant Central as TVMFFICentral
participant Cache as TVMFFIConfigCache
participant PyKernel as Python get_mxfp8_quantization_function
participant CUDA as mxfp8::quantize (CUDA)
Py->>Init: import transformer_engine
Init->>Init: _load_tvm_ffi_library() [if env var set]
Init->>Init: _register_cutedsl_backends()
Py->>Cpp: nvte_quantize(input, output)
Cpp->>Central: getInstance() [lazy init: dlopen libtvm_ffi.so]
Cpp->>Cache: get_or_load(MXFP8QuantConfig)
Cache->>Cache: check map_ (cache miss)
Cache->>Central: load_tvm_ffi_function(cfg)
Central->>PyKernel: entrypoint(fn_name, dtype, ...) via TVM-FFI
PyKernel->>PyKernel: compile_cutedsl_function_from_cfg(cfg)
PyKernel->>PyKernel: tvm_ffi.register_global_func(fn_name, compiled)
PyKernel-->>Central: True
Central->>Central: GetGlobal(key)
Central-->>Cache: optional Function
Cache->>Cache: map_.emplace(id, fn)
Cache-->>Cpp: optional Function
alt CuTeDSL kernel found
Cpp->>Cpp: invoke mxfp8_quant_func
Cpp->>Cpp: reduce_dbias (if with_dbias)
else CuTeDSL not available
Cpp->>CUDA: mxfp8::quantize(...)
end
Reviews (29): Last reviewed commit: "fix" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| "importlib-metadata>=1.0", | ||
| "packaging", | ||
| "apache-tvm-ffi>=0.1.12", | ||
| "nvidia-cutlass-dsl>=4.2.0", |
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.
There was a problem hiding this comment.
I'll change this to 4.4.2
| GTEST_SKIPs the mismatched half), non-32-divisible shapes are omitted (the | ||
| dispatcher can never route them to CuTeDSL), and a missing kernel registration |
There was a problem hiding this comment.
Why are the non-32-divisible shapes omitted? Is this a limitation of the cutedsl implementation?
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. |
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.
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?
| 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; | ||
| } |
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.
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?
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).
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.
| if(NOT TVM_FFI_CMAKE_QUERY EQUAL 0) | ||
| message(FATAL_ERROR | ||
| "Could not import the tvm_ffi Python package (with '${Python_EXECUTABLE}'), " | ||
| "which Transformer Engine requires to build the CuTeDSL quantize backend " | ||
| "bridge (common/tvm_ffi_bridge.h). Install it into this Python environment: " | ||
| "`pip install apache-tvm-ffi`.") | ||
| endif() | ||
| find_package(tvm_ffi CONFIG REQUIRED PATHS "${TVM_FFI_CMAKE_DIR}") | ||
|
|
There was a problem hiding this comment.
The tvm_ffi detection is unconditional: if the package is absent the build hard-fails with FATAL_ERROR, with no NVTE_ENABLE_CUTEDSL CMake option to disable it. Every other optional feature in this file that has a build-time cost (NVTE_ENABLE_NVSHMEM, NVTE_WITH_CUBLASMP, etc.) is guarded by an option() flag. Without an analogous guard here, any environment where apache-tvm-ffi is not installed — CI images, embedded build systems, custom wheels — cannot build Transformer Engine at all, even if the CuTeDSL backend is never used at runtime.
948fab5 to
2930b1b
Compare
af55445 to
87adfe7
Compare
| # 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 | ||
|
|
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.
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>
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>
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>
| except Exception as e: # pylint: disable=broad-exception-caught | ||
| logger.error( | ||
| "CuTeDSL MXFP8 kernel compilation failed, falling back to the CUDA C++ kernel: %s", | ||
| e, | ||
| ) | ||
| return False | ||
| tvm_ffi.register_global_func(fn_name, compiled, override=True) | ||
|
|
||
| return True |
There was a problem hiding this comment.
register_global_func uncaught exception crashes quantize instead of falling back
compile_cutedsl_function_from_cfg is correctly wrapped in a try/except Exception, but tvm_ffi.register_global_func(fn_name, compiled, override=True) is called outside that block. If it raises (e.g., a TVM-FFI internal error, an invalid function handle, or a name-space conflict), the Python exception propagates through TVM-FFI's (*entrypoint)(...) invocation in C++ (retrieve_func_from_python), which is not guarded there either. The exception travels up through load_tvm_ffi_function → get_or_load → get_kernel → mxfp8_quantize_cutedsl and crashes the quantize operation instead of gracefully returning false to let the CUDA kernel handle the call. Moving the register_global_func call inside the existing try block (or wrapping it in its own) preserves the established fallback contract.
3 participants
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
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: