[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 48 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
Confidence Score: 4/5The new CuTeDSL dispatch path falls back to CUDA for any unsupported config, so existing quantize calls are not broken. Several open issues from prior rounds remain unresolved including reduce_dbias running on unwritten workspace when the kernel is nooped, device-0 arch query in multi-GPU environments, and a test-collection crash when tvm_ffi is absent. The fallback-to-CUDA logic is sound and the env-var gate is now correctly applied at registration time. The main risks are the reduce_dbias noop divergence from the CUDA path, the device-0 arch selection which could compile Blackwell PTX for the wrong device, and the hard apache-tvm-ffi build dependency. Files Needing Attention: transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh (reduce_dbias noop handling), transformer_engine/common/CuTeDSL/utils_fp8.py (device-0 arch selection at import time), tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py (unconditional tvm_ffi import) Important Files Changed
Sequence DiagramsequenceDiagram
participant User as User code
participant Cpp as quantize.cuh dispatcher
participant Bridge as TVMFFICentral
participant Python as Python factory
participant Cache as kernel cache
participant CUDA as CUDA C++ kernel
User->>Cpp: nvte_quantize
Cpp->>Bridge: mxfp8_quantize_cutedsl config
Bridge->>Bridge: check env var and libtvm_ffi.so
Bridge->>Cache: lookup config key
alt cache miss first call
Bridge->>Python: get_mxfp8_quantization_function fn_name dtype
Python->>Python: validate config
Python->>Python: device_compute_capability device 0
Python->>Python: cute.compile JIT kernel
Python->>Python: register_global_func fn_name compiled
Python-->>Bridge: True or False
Bridge->>Cache: emplace key fn or nullopt
else cache hit
Cache-->>Bridge: fn or nullopt
end
alt CuTeDSL supported
Bridge-->>Cpp: tvm ffi Function
Cpp->>Cpp: zero_scales_kernel if swizzled
Cpp->>Python: invoke compiled kernel tensors stream
Cpp->>Cpp: reduce_dbias if WITH_DBIAS
Cpp-->>User: true done
else unsupported or backend disabled
Bridge-->>Cpp: nullopt or false
Cpp->>CUDA: mxfp8 quantize
CUDA-->>User: result
end
Reviews (25): 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.
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>
…it__.py Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci 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 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
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: