[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 44 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. The C++ dispatcher (
Confidence Score: 3/5The CuTeDSL path is opt-in and falls back to CUDA on any failure, but the device-capability detection is locked in at module-import time from device 0 rather than the active compute device, making the backend unreliable on mixed-generation multi-GPU systems. The device-capability detection (cuda.core.Device() defaulting to device 0, frozen by @lru_cache) is shared by both utils.py and utils_fp8.py. It locks the Blackwell gate and the e8m0 converter choice at module-import time rather than at JIT-compile time. On multi-GPU systems with mixed GPU generations this can silently disable the CuTeDSL backend or compile Blackwell PTX for a non-Blackwell device. The test file is also fragile without NVTE_ENABLE_CUTEDSL_QUANT_BACKEND set at import, producing misleading assertion failures rather than clean skips. transformer_engine/common/CuTeDSL/utils.py and utils_fp8.py (device-capability detection and e8m0 converter selection), and tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py (missing env-var guard). Important Files Changed
Sequence DiagramsequenceDiagram
participant PY as Python import (common/__init__.py)
participant TVMFFI as TVM-FFI global registry
participant CPP as C++ dispatcher (quantize.cuh)
participant BRIDGE as TVMFFICentral (tvm_ffi_bridge.h)
participant KERNEL as CuTeDSL kernel (quantize_mxfp8.py)
participant CUDA as CUDA C++ kernel
Note over PY: NVTE_ENABLE_CUTEDSL_QUANT_BACKEND=1
PY->>PY: _load_tvm_ffi_library() → ctypes.CDLL(libtvm_ffi.so, RTLD_GLOBAL)
PY->>TVMFFI: register_global_func(get_mxfp8_quantization_function)
Note over CPP: nvte_quantize() called
CPP->>BRIDGE: mxfp8_quantize_cutedsl(config, tensors, stream)
BRIDGE->>BRIDGE: shape alignment check (% 32)
BRIDGE->>BRIDGE: lazyload_function(config) — cache miss
BRIDGE->>TVMFFI: GetGlobal(get_mxfp8_quantization_function)
TVMFFI-->>BRIDGE: entrypoint fn
BRIDGE->>KERNEL: entrypoint(fn_name, dtype, fp8_dtype, flags)
KERNEL->>KERNEL: check Blackwell (device_compute_capability lru_cached)
KERNEL->>KERNEL: cute.compile(kernel_obj, fake_tensors)
KERNEL->>TVMFFI: register_global_func(fn_name, compiled)
KERNEL-->>BRIDGE: True
BRIDGE->>TVMFFI: GetGlobal(fn_name)
BRIDGE->>BRIDGE: "cache[key] = compiled fn"
BRIDGE->>KERNEL: compiled_fn(tensors, stream)
BRIDGE-->>CPP: true
Note over CPP: unsupported config or backend disabled
CPP->>BRIDGE: mxfp8_quantize_cutedsl(config, ...)
BRIDGE-->>CPP: false (nullopt)
CPP->>CUDA: mxfp8::quantize(tensors, stream)
Reviews (17): Last reviewed commit: "fix" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| # apache-tvm-ffi: headers for the C++ API (Module / Function / TensorView) | ||
| # and libtvm_ffi.so for symbol resolution. Used by tvm_ffi_bridge.h / | ||
| # applyTVMFunction. Python registers AOT-compiled CuTeDSL kernels into | ||
| # the global registry; TE C++ looks them up via Function::GetGlobalRequired. |
There was a problem hiding this comment.
Comment is too verbose and talks about things not needed in the context of the build system.
| # rpath pinned to the pip install dir so the loader finds libtvm_ffi.so | ||
| # without LD_LIBRARY_PATH at runtime. | ||
| extra_link_args = [f"-Wl,-rpath,{tvm_ffi_lib_dir}"] |
There was a problem hiding this comment.
This will (maybe) work in the sdist installation, but what if we decide to do the binary wheel?
All of the other such libraries are linked normally and either loaded using LD_LIBRARY_PATH or
being found and loaded at runtime before we load libtransformer_engine.so.
There was a problem hiding this comment.
Also, what about JAX? It will also use those functions - shouldn't those changes land on the common side instead?
There was a problem hiding this comment.
OK I think I need to delete this from the pytorch side and make them dependencies in common instead
| "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?
| Registration is explicit: call :func:`register_cutedsl_backends` to expose the | ||
| CuTeDSL kernel entrypoints (e.g. ``get_mxfp8_quantization_function``) as | ||
| TVM-FFI global functions. The C++ dispatcher (init_cutedsl_extension in the | ||
| PyTorch extension) imports this package and calls it once per process; it then | ||
| probes the names via ``tvm::ffi::Function::GetGlobal`` — finding one means the | ||
| CuTeDSL toolchain is available and the kernel may be compiled on demand, not | ||
| finding it means a plain C++ environment and the dispatcher falls back to the | ||
| CUDA C++ kernel. |
There was a problem hiding this comment.
This comment is too verbose and mostly belongs to the register_cutedsl_backends function instead.
| """Tanh-approximation GELU. Constants and operator grouping match TE's | ||
| `transformer_engine/common/util/math.h::gelu` exactly (factored form | ||
| `x · (0.5 + 0.5·tanh(x·(a + b·x²)))`) so quantized output is bit-exact | ||
| against the C++ fused IS_ACT path. Uses `cute.math.tanh(fastmath=False)` | ||
| rather than the `tanh.approx.f32` PTX intrinsic — TE compiles activation | ||
| kernels without `--use_fast_math` by default, so its `tanhf` is the | ||
| IEEE-precise expansion.""" |
There was a problem hiding this comment.
This raises a question - we do actually have a knob to use fast math for activations - I assume that
there is no equivalent for that here? We should add that.
There was a problem hiding this comment.
Ah yes I just found NVTE_BUILD_ACTIVATION_WITH_FAST_MATH. Fixed my code
| def act_silu(x: Float32) -> Float32: | ||
| """SiLU/Swish: x · σ(x) = x / (1 + e^-x). | ||
| Matches TE's `silu` (`val / (1 + expf(-val))`).""" | ||
| return x / (Float32(1.0) + cute.math.exp(-x, fastmath=True)) |
There was a problem hiding this comment.
I don't think expf (compiled without fast math) is the same as cute.math.exp(fastmath=True) - is it?
| """Quick GELU: x · σ(1.702·x). Matches TE `qgelu_with_alpha(val, 1.702)` = | ||
| `cval · (1 / (1 + expf(-1.702·cval)))` (multiply by sigmoid, not a divide).""" | ||
| z = Float32(1.702) * x | ||
| return x * (Float32(1.0) / (Float32(1.0) + cute.math.exp(-z, fastmath=True))) |
There was a problem hiding this comment.
So why are we not reusing the sigmoid function here?
There was a problem hiding this comment.
Ah Claude Code wrote that. Refactored this
| 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?
| 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.
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>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
948fab5 to
2930b1b
Compare
…it__.py Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
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
2 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: