[Common] Experimental CuTeDSL MXFP8 backends in C++ via TVM-FFI

[kainzhong]

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=1 to enable) and transparently falls back to the CUDA kernels for any unsupported config or shape (use NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN=1 to enable warning for unsupported cases).

How it works

• Python CuTeDSL kernels (transformer_engine/common/CuTeDSL/) register a factory as a TVM-FFI global.
• On a matching MXFP8 quantize, the C++ dispatcher builds a config key, asks Python to JIT-compile + register a kernel for that config (cached), then invokes it with DLTensor views of the TE tensors.
• Any miss (backend disabled, unsupported config, non-32-aligned shape) returns false → existing mxfp8::quantize CUDA path runs unchanged.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Added utility classes and functions to retrieve TVM-FFI global functions using strings
• Ported MXFP8 CUDA C++ quantization code to CuTeDSL

Breaking changes:

• Building TE from python now requires apache-tvm-ffi package installed. If building the C++ lib, you can pass NVTE_WITH_CUTEDSL to cmake to opt-out.
• apache-tvm-ffi and nvidia-cutlass-dsl packages are TE's dependency now.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 4/5

[High risk] Adds optional CuTeDSL backend with new build-time and runtime dependencies.

The PR is not ready to merge while the new device-count limit can fail quantization.

Findings

1.  Device count limit breaks quantization ▶
2.  Backend tests always skip ▶

Summary

The PR adds an opt-in CuTeDSL MXFP8 quantization path through TVM-FFI, with build controls and cross-backend tests. Since the previous review, kernel first-launch state was changed to use a fixed-size device array; that introduces a failure on hosts with more than 32 visible devices.

_{Reviews (84) · Last reviewed commit: "nit"}

[ptrendx]

Kind of random, but this function is quite slow. You could do the same much faster with raw char*
manipulation.

[kainzhong]

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?

[ptrendx]

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).

[kainzhong]

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.

[timmoon10]

concat_strings is more ergonomic and it doesn't need the bool ? "1" : "0" pattern. It also gets rid of some, but not all, of the heap allocations.

TransformerEngine/transformer_engine/common/util/string.h

Line 47 in a30aee5

inline std::string concat_strings(const Ts &...args) {

[ptrendx]

One thing to note also is the cache key should include at least the SM arch of the device that was used for the compilation - otherwise there is a chance in the heterogeneous gpu setup to try running the kernel compiled for 1 architecture on another architecture.

[kainzhong]

@fheinecke I modified this so GitHub CI can pass. Is this OK?

[timmoon10]

Nit: This is a quickstart for naive users, who shouldn't need to be aware of advanced configurations like NVTE_WITH_CUTEDSL=0.

Suggested change

	# apache-tvm-ffi is required to build the CuTeDSL bridge when using --no-build-isolation.
	# If that's not wanted, set NVTE_WITH_CUTEDSL=0 to build without CuTeDSL and tvm-ffi bindings.
	pip3 install "apache-tvm-ffi>=0.1.12"
	pip3 install apache-tvm-ffi>=0.1.12 # Install build-time dependencies

[timmoon10]

Nit: This is a minimal build guide for naive users. Documenting advanced options here is distracting.

Suggested change

	# Required by the default CuTeDSL-enabled build when using --no-build-isolation.
	pip install "apache-tvm-ffi>=0.1.12"
	pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable
	Set ``NVTE_WITH_CUTEDSL=0`` on the second command to build without the optional CuTeDSL bridge
	and apache-tvm-ffi as its python build dependency.
	pip install apache-tvm-ffi>=0.1.12 # Build-time dependencies
	pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable

[greptile-apps]

Backend tests always skip

This job sets only the fallback-warning variable, not NVTE_ENABLE_CUTEDSL_BACKEND=1. The new PyTorch test module skips every CuTeDSL comparison when that enable variable is unset, so L0 collects the suite without exercising the new backend. The same issue occurs in qa/L0_jax_unittest/test.sh:31, whose new JAX comparison module uses the same skip gate. This leaves the feature without its intended L0 correctness coverage.

[greptile-apps]

Device count limit breaks quantization

If a host has more than 32 visible CUDA devices, the first supported CuTeDSL quantization throws while creating its kernel state. The call fails instead of launching the kernel or falling back to CUDA. The previous dynamically sized array handled the device count returned by CUDA.

[kainzhong]

Currently a host won't have so many CUDA devices plus this workaround will be removed once I know how to fix CuTeDSL's end

[greptile-apps]

Understood. Given that hosts with more than 32 visible CUDA devices are not an expected deployment scenario for this experimental backend, and this workaround will be removed once the CuTeDSL limitation is resolved, I’ll treat the risk as accepted.