[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 47 commits into
Conversation
daee750 to
218cd24
Compare
2f8c8da to
8448930
Compare
|
Benchmark on ptyche (B200 GPU, ARM CPU): |
Greptile SummaryThis PR introduces an experimental, opt-in CuTeDSL backend for MXFP8 quantization in Transformer Engine, bridging JIT-compiled CUTLASS Python DSL kernels into the existing C++ MXFP8
Confidence Score: 3/5Not yet safe to merge: several correctness and deployment issues remain unresolved. The dispatch shim calls lazyload_function (which may trigger tens-of-seconds of NVCC JIT) before the cheap workspace-size-query early-return, stalling latency-sensitive callers on first use. The reduce_dbias step after the TVM-FFI kernel call runs unconditionally — if the noop flag caused the GPU kernel to skip writing the workspace, reduce_dbias sums uninitialized memory into dbias_tensor. Both apache-tvm-ffi and nvidia-cutlass-dsl are now hard install_requires for every TE user and a CMake FATAL_ERROR with no opt-out flag, despite the feature being opt-in. Files Needing Attention: quantize_mxfp8_cutedsl.cuh (workspace-query JIT stall, reduce_dbias noop ordering), setup.py + pyproject.toml + CMakeLists.txt (unconditional hard dependencies), quantize_mxfp8.py (device-0 hardcode in capability check) Important Files Changed
Sequence DiagramsequenceDiagram
participant App as Application
participant QCUH as quantize.cuh (dispatch)
participant CDSL as quantize_mxfp8_cutedsl.cuh
participant TVMC as TVMFFICentral (singleton)
participant PY as Python: get_mxfp8_quantization_function
participant CUDA as CUDA C++ mxfp8::quantize
App->>QCUH: nvte_quantize(input, output, config)
QCUH->>CDSL: "mxfp8_quantize_cutedsl<IS_DBIAS,...>(...)"
CDSL->>CDSL: shape alignment check (% 32)
alt shape not 32-aligned
CDSL-->>QCUH: false (fall back)
QCUH->>CUDA: mxfp8::quantize(...)
else shape OK
CDSL->>TVMC: lazyload_function(config)
alt cache hit
TVMC-->>CDSL: "optional<Function>"
else cache miss (first call)
TVMC->>PY: retrieve_func_from_python(key)
PY->>PY: cute.compile(kernel_obj, ...) [JIT via NVCC]
PY->>PY: tvm_ffi.register_global_func(key, compiled)
PY-->>TVMC: true/false
TVMC->>TVMC: cache_.emplace(key, fn_or_nullopt)
TVMC-->>CDSL: "optional<Function>"
end
alt function found
CDSL->>CDSL: zero_scales_kernel (if swizzled+padding)
CDSL->>PY: "(*fn)(mX, mO_row, mS_row, ..., stream)"
PY-->>CDSL: kernel executes on GPU
opt WITH_DBIAS
CDSL->>CDSL: reduce_dbias(workspace to dbias_tensor)
end
CDSL-->>QCUH: true
else no function (unsupported)
CDSL-->>QCUH: false
QCUH->>CUDA: mxfp8::quantize(...)
end
end
Reviews (23): Last reviewed commit: "warn only once if dlopen fails" | 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
| std::optional<tvm::ffi::Function> mxfp8_quant_func_opt = | ||
| tvm_ffi_bridge::TVMFFICentral::getInstance().lazyload_function(config); | ||
| if (!mxfp8_quant_func_opt.has_value()) { | ||
| return false; | ||
| } | ||
|
|
||
| // When only WITH_DBIAS is true, we use a larger tile size (align with CUDA C++ implementation) | ||
| const bool cast_dbias_only = config.with_dbias && !config.with_dact && !config.with_act; | ||
| const size_t chunk_rows = cast_dbias_only ? 128 : 64; // input rows reduced per CTA | ||
| // Each CTA writes one partial-dbias row, so the workspace (and the cross-CTA | ||
| // reduction below) has ceil(M / chunk_rows) rows. | ||
| const size_t workspace_rows = (flat_m + chunk_rows - 1) / chunk_rows; | ||
|
|
||
| // dbias workspace-size query, mirroring mxfp8::quantize: the framework first | ||
| // calls with an unallocated workspace to learn its shape, allocates a buffer of | ||
| // that shape, then calls again to run. The kernel writes per-row-block partial | ||
| // dbias into this workspace; reducing it to the final dbias is a separate step. | ||
| if (config.with_dbias && workspace_tensor != nullptr && workspace_tensor->data.dptr == nullptr) { | ||
| workspace_tensor->data.shape = {workspace_rows, flat_n}; | ||
| workspace_tensor->data.dtype = DType::kFloat32; | ||
| return true; | ||
| } |
There was a problem hiding this comment.
lazyload_function(config) — which may compile the CUDA kernel (tens of seconds of NVCC JIT) — runs before the workspace-size-query early-return on line 175. The very first call with workspace_tensor->data.dptr == nullptr for a with_dbias config blocks the calling thread for the full JIT duration before returning only the workspace dimensions. The CUDA C++ path computes this immediately from DIVUP(rows, CHUNK_DIM_Y) without any compilation.
In frameworks that call the workspace-size query in a latency-sensitive context (e.g., during graph capture or model init), this causes an unexpected multi-second stall on first use. Moving the workspace-size query before the lazyload_function call (or checking the cache hit before compiling) would decouple the size calculation from JIT and keep the workspace query cheap for cached configs.
| # 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>
for more information, see https://pre-commit.ci 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: