From d0181a8ec2cd6814866c73dac0b4af3cac3eb552 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 26 Oct 2026 18:53:20 +0000 Subject: [PATCH 1/6] make cuda-bindings dependency explicit Signed-off-by: Kaining Zhong --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 2dabf37b322..7e04fcdab29 100644 --- a/setup.py +++ b/setup.py @@ -135,6 +135,7 @@ def setup_requirements() -> Tuple[List[str], List[str]]: "pydantic", "importlib-metadata>=1.0", "packaging", + "cuda-bindings", ] test_reqs: List[str] = ["pytest>=8.2.1"] From a91e1d942ac40b438a43645e1c026588abad0b77 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 26 Oct 2026 18:53:41 +0000 Subject: [PATCH 2/6] make CuTeDSL fallback warn or fail Signed-off-by: Kaining Zhong --- docs/envvars.rst | 10 ++-- .../cast/mxfp8/quantize_mxfp8_cutedsl.cuh | 6 +-- transformer_engine/common/tvm_ffi_bridge.h | 54 +++++++++++++------ 3 files changed, 47 insertions(+), 23 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 09e2b462780..c1e1b0b62ec 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -153,11 +153,13 @@ General It also requires ``apache-tvm-ffi``, ``nvidia-cutlass-dsl`` to be installed and you are using CPython 3.10 or later. If these requirements are not met, Transformer Engine will fall back to using CUDA C++ kernels if possible. -.. envvar:: NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN +.. envvar:: NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN - :Type: ``int`` (0 or 1) - :Default: ``0`` - :Description: Warn when TE falls back to the CUDA C++ kernels instead of dispatching to available CuTeDSL kernels. Useful to verify if the CuTeDSL path is taken. + :Type: ``str`` (``WARN`` or ``ERROR``) + :Default: None + :Description: Control what happens when TE fails dispatch to a CuTeDSL kernel and would fall back to the CUDA C++ kernel. + ``WARN`` emits a warning explaining why CuTeDSL kernels are unavailable and falls back to CUDA kernels. ``ERROR`` raises an error and fails instead. + This includes the case where the CuTeDSL backend is disabled (see :envvar:`NVTE_ENABLE_CUTEDSL_BACKEND`). .. envvar:: NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh index 431812db143..e8897dc92ba 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh @@ -121,7 +121,7 @@ inline bool mxfp8_quantize_cutedsl(const MXFP8QuantConfig &config, const Tensor const size_t flat_n = input_tensor->flat_last_dim(); if (flat_n % 16 != 0) { // Shape not supported because the last dimension of output is not 16 bytes aligned which is required by TMA. - maybe_warn_cutedsl_not_chosen("the last dimension of input is not 16 bytes aligned."); + report_cutedsl_not_chosen("the last dimension of input is not 16 bytes aligned."); return false; } @@ -229,12 +229,12 @@ bool mxfp8_quantize_cutedsl(const Tensor *input_tensor, const Tensor *act_input_ Tensor *workspace_tensor, const bool use_2d_quantization, cudaStream_t stream) { if (!tvm_ffi_bridge::TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { - maybe_warn_cutedsl_not_chosen("the CuTeDSL backend is disabled."); + report_cutedsl_not_chosen("the CuTeDSL backend is disabled."); return false; } // TODO(kainingz): port 2D quantization to CuTeDSL if (use_2d_quantization) { - maybe_warn_cutedsl_not_chosen("2D quantization is not supported."); + report_cutedsl_not_chosen("2D quantization is not supported."); return false; } checkCuDriverContext(stream); diff --git a/transformer_engine/common/tvm_ffi_bridge.h b/transformer_engine/common/tvm_ffi_bridge.h index d57fd8f03da..f36e0591e51 100644 --- a/transformer_engine/common/tvm_ffi_bridge.h +++ b/transformer_engine/common/tvm_ffi_bridge.h @@ -240,7 +240,7 @@ class TVMFFICentral { "`bool retrieve_func_from_python(const std::string&) const`, " "and `std::optional get_kernel() const`."); if (!cutedsl_backend_enabled_.load(std::memory_order_relaxed)) { - maybe_warn_not_chosen( + report_not_chosen( "the CuTeDSL backend is disabled, so no kernel is available for " "config `", cfg.to_key(), "`. Set NVTE_ENABLE_CUTEDSL_BACKEND=1 to enable it."); @@ -252,7 +252,7 @@ class TVMFFICentral { // Warn once rather than spamming every quantize call while TVM-FFI is unavailable. static std::once_flag warned; std::call_once(warned, [this] { - maybe_warn_not_chosen( + report_not_chosen( "Failed to load tvm-ffi function because libtvm_ffi.so is not successfully loaded."); }); return std::nullopt; @@ -261,7 +261,7 @@ class TVMFFICentral { std::optional fn = cfg.retrieve_func_from_python(key) ? tvm::ffi::Function::GetGlobal(key) : std::nullopt; if (!fn) { - maybe_warn_not_chosen("no TVM-FFI kernel is registered for config `", key, "`."); + report_not_chosen("no TVM-FFI kernel is registered for config `", key, "`."); } return fn; } @@ -292,7 +292,11 @@ class TVMFFICentral { return cutedsl_backend_enabled_.load(std::memory_order_relaxed); } - bool get_warn_cutedsl_backend_not_chosen() const { return warn_cutedsl_backend_not_chosen_; } + // Behavior when a CuTeDSL kernel is not chosen and TE falls back to the CUDA C++ kernel, + // controlled by NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN. + enum class ReportNotChosenLevel { kSilent, kWarn, kError }; + + ReportNotChosenLevel get_report_not_chosen_level() const { return report_not_chosen_level_; } // These are C ABI calls: TVM reports failures through its return code and thread-local error, // not C++ exceptions. noexcept prevents an unexpected exception from crossing that ABI boundary. @@ -302,11 +306,17 @@ class TVMFFICentral { void move_error(TVMFFIObjectHandle *result) const noexcept { move_error_(result); } - // Optionally emit a warning explaining why the CuTeDSL backend was not chosen for this config. + // Optionally warn or fail when CuTeDSL kernels are not chosen based on NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN. template - void maybe_warn_not_chosen(const Args &...reason) const { - if (warn_cutedsl_backend_not_chosen_) { - NVTE_WARN("The CuTeDSL kernel is not chosen because ", reason...); + void report_not_chosen(const Args &...reason) const { + switch (report_not_chosen_level_) { + case ReportNotChosenLevel::kSilent: + break; + case ReportNotChosenLevel::kWarn: + NVTE_WARN("The CuTeDSL kernel is not chosen because ", reason...); + break; + case ReportNotChosenLevel::kError: + NVTE_ERROR("The CuTeDSL kernel is not chosen because ", reason...); } } @@ -315,7 +325,7 @@ class TVMFFICentral { TVMFFICentral() : tvm_ffi_available_(false), cutedsl_backend_enabled_(is_cutedsl_backend_enabled()), - warn_cutedsl_backend_not_chosen_(warn_if_cutedsl_backend_not_chosen()) { + report_not_chosen_level_(read_report_not_chosen_level()) { set_tvm_ffi_available(prepare_tvm_ffi()); } @@ -356,14 +366,26 @@ class TVMFFICentral { return transformer_engine::getenv("NVTE_ENABLE_CUTEDSL_BACKEND"); } - static bool warn_if_cutedsl_backend_not_chosen() { - return transformer_engine::getenv("NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN"); + static ReportNotChosenLevel read_report_not_chosen_level() { + const std::string value = + transformer_engine::getenv("NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN"); + if (value.empty()) { + return ReportNotChosenLevel::kSilent; + } + if (value == "WARN") { + return ReportNotChosenLevel::kWarn; + } + if (value == "ERROR") { + return ReportNotChosenLevel::kError; + } + NVTE_ERROR("Invalid value `", value, + "` for NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN. Expected WARN or ERROR."); } std::atomic tvm_ffi_available_; // libtvm_ffi.so loaded; false disables the backend std::mutex tvm_ffi_init_mutex_; std::atomic cutedsl_backend_enabled_; - const bool warn_cutedsl_backend_not_chosen_; + const ReportNotChosenLevel report_not_chosen_level_; void *tvm_ffi_handle_ = nullptr; // These are symbols dynamically loaded from TVM-FFI @@ -390,7 +412,7 @@ class TVMFFIConfigCache { std::optional get_or_load(const Config &cfg) { TVMFFICentral ¢ral = TVMFFICentral::getInstance(); // Checked ahead of the cache so that toggling the backend off still disables - // already-resolved configs; load_tvm_ffi_function emits the "disabled" warning. + // already-resolved configs; load_tvm_ffi_function reports the "disabled" reason. if (!central.get_cutedsl_backend_enabled()) { return central.load_tvm_ffi_function(cfg); } @@ -426,10 +448,10 @@ class TVMFFIConfigCache { std::unordered_map> map_; }; -// Optionally emit a warning explaining why the CuTeDSL backend was not chosen for this config. +// Optionally warn or fail when CuTeDSL kernels are not chosen based on NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN. template -inline void maybe_warn_cutedsl_not_chosen(const Args &...reason) { - TVMFFICentral::getInstance().maybe_warn_not_chosen(reason...); +inline void report_cutedsl_not_chosen(const Args &...reason) { + TVMFFICentral::getInstance().report_not_chosen(reason...); } } // namespace tvm_ffi_bridge From c18318882f80d315010a63792ce6810c7f23a6c8 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 26 Oct 2026 19:03:45 +0000 Subject: [PATCH 3/6] don't warn or fail for the unsupported 2d quant case Signed-off-by: Kaining Zhong --- transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh index e8897dc92ba..ed591980b6a 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh @@ -234,7 +234,7 @@ bool mxfp8_quantize_cutedsl(const Tensor *input_tensor, const Tensor *act_input_ } // TODO(kainingz): port 2D quantization to CuTeDSL if (use_2d_quantization) { - report_cutedsl_not_chosen("2D quantization is not supported."); + // report_cutedsl_not_chosen("2D quantization is not supported."); return false; } checkCuDriverContext(stream); From 1b1a9b4a50fb54205284bdc8a9e5d2e4c95795c0 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 26 Oct 2026 20:14:57 +0000 Subject: [PATCH 4/6] fail if CuTeDSL kerenls are unavailable Signed-off-by: Kaining Zhong --- qa/L0_cppunittest/test.sh | 2 ++ qa/L0_jax_unittest/test.sh | 2 ++ qa/L0_pytorch_unittest/test.sh | 2 ++ 3 files changed, 6 insertions(+) diff --git a/qa/L0_cppunittest/test.sh b/qa/L0_cppunittest/test.sh index c7499282f44..27224a7579a 100755 --- a/qa/L0_cppunittest/test.sh +++ b/qa/L0_cppunittest/test.sh @@ -12,6 +12,8 @@ mkdir -p "$XML_LOG_DIR" TE_LIB_PATH=$(pip3 show transformer-engine | grep -E "Location:|Editable project location:" | tail -n 1 | awk '{print $NF}') export LD_LIBRARY_PATH=$TE_LIB_PATH:$LD_LIBRARY_PATH +export NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN=ERROR + # Set parallelization parameters NUM_PHYSICAL_CORES=$(nproc) NUM_PARALLEL_JOBS=4 diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index 339887719b1..8276b379b56 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -21,6 +21,8 @@ FAILED_CASES="" export NVTE_JAX_TEST_TIMING=1 +export NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN=ERROR + pip3 install "nltk>=3.8.2,!=3.10.1" || error_exit "Failed to install nltk" pip3 install pytest==8.2.1 pytest-timeout==2.4.0 || error_exit "Failed to install pytest dependencies" diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f45e9369d10..cfcf2351823 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -27,6 +27,8 @@ export NVTE_FLASH_ATTN_V2=1 export NVTE_FLASH_ATTN_V3=0 export NVTE_FLASH_ATTN_V4=0 +export NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN=ERROR + pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_optional_flash_attn_import.xml $TE_PATH/tests/pytorch/test_optional_flash_attn_import.py || test_fail "test_optional_flash_attn_import.py" From 30e3237ee8891ddeeb568c190009ad902b97cdbf Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Mon, 26 Oct 2026 20:36:03 +0000 Subject: [PATCH 5/6] only warn or fail when CuTeDSL is enabled Signed-off-by: Kaining Zhong --- docs/envvars.rst | 2 +- qa/L0_cppunittest/test.sh | 1 + qa/L0_jax_unittest/test.sh | 1 + qa/L0_pytorch_unittest/test.sh | 1 + .../common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh | 1 - transformer_engine/common/tvm_ffi_bridge.h | 12 ++++++------ 6 files changed, 10 insertions(+), 8 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index c1e1b0b62ec..c8dffeb14e2 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -159,7 +159,7 @@ General :Default: None :Description: Control what happens when TE fails dispatch to a CuTeDSL kernel and would fall back to the CUDA C++ kernel. ``WARN`` emits a warning explaining why CuTeDSL kernels are unavailable and falls back to CUDA kernels. ``ERROR`` raises an error and fails instead. - This includes the case where the CuTeDSL backend is disabled (see :envvar:`NVTE_ENABLE_CUTEDSL_BACKEND`). + Only effective when the CuTeDSL backend is enabled (see :envvar:`NVTE_ENABLE_CUTEDSL_BACKEND`). .. envvar:: NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB diff --git a/qa/L0_cppunittest/test.sh b/qa/L0_cppunittest/test.sh index 27224a7579a..8011c314bd9 100755 --- a/qa/L0_cppunittest/test.sh +++ b/qa/L0_cppunittest/test.sh @@ -12,6 +12,7 @@ mkdir -p "$XML_LOG_DIR" TE_LIB_PATH=$(pip3 show transformer-engine | grep -E "Location:|Editable project location:" | tail -n 1 | awk '{print $NF}') export LD_LIBRARY_PATH=$TE_LIB_PATH:$LD_LIBRARY_PATH +# If CuTeDSL backend is not enabled, this will be ignored export NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN=ERROR # Set parallelization parameters diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index 8276b379b56..101be267418 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -21,6 +21,7 @@ FAILED_CASES="" export NVTE_JAX_TEST_TIMING=1 +# If CuTeDSL backend is not enabled, this will be ignored export NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN=ERROR pip3 install "nltk>=3.8.2,!=3.10.1" || error_exit "Failed to install nltk" diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index cfcf2351823..7393ebef276 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -27,6 +27,7 @@ export NVTE_FLASH_ATTN_V2=1 export NVTE_FLASH_ATTN_V3=0 export NVTE_FLASH_ATTN_V4=0 +# If CuTeDSL backend is not enabled, this will be ignored export NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN=ERROR pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh index ed591980b6a..6380dda21da 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8_cutedsl.cuh @@ -229,7 +229,6 @@ bool mxfp8_quantize_cutedsl(const Tensor *input_tensor, const Tensor *act_input_ Tensor *workspace_tensor, const bool use_2d_quantization, cudaStream_t stream) { if (!tvm_ffi_bridge::TVMFFICentral::getInstance().get_cutedsl_backend_enabled()) { - report_cutedsl_not_chosen("the CuTeDSL backend is disabled."); return false; } // TODO(kainingz): port 2D quantization to CuTeDSL diff --git a/transformer_engine/common/tvm_ffi_bridge.h b/transformer_engine/common/tvm_ffi_bridge.h index f36e0591e51..1bb3074c9dc 100644 --- a/transformer_engine/common/tvm_ffi_bridge.h +++ b/transformer_engine/common/tvm_ffi_bridge.h @@ -240,10 +240,6 @@ class TVMFFICentral { "`bool retrieve_func_from_python(const std::string&) const`, " "and `std::optional get_kernel() const`."); if (!cutedsl_backend_enabled_.load(std::memory_order_relaxed)) { - report_not_chosen( - "the CuTeDSL backend is disabled, so no kernel is available for " - "config `", - cfg.to_key(), "`. Set NVTE_ENABLE_CUTEDSL_BACKEND=1 to enable it."); return std::nullopt; } // Only check if libtvm_ffi.so is loaded if user enables the CuTeDSL backend. @@ -308,7 +304,11 @@ class TVMFFICentral { // Optionally warn or fail when CuTeDSL kernels are not chosen based on NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN. template + // No-op while the CuTeDSL backend is disabled, since falling back is then expected. void report_not_chosen(const Args &...reason) const { + if (!get_cutedsl_backend_enabled()) { + return; + } switch (report_not_chosen_level_) { case ReportNotChosenLevel::kSilent: break; @@ -412,9 +412,9 @@ class TVMFFIConfigCache { std::optional get_or_load(const Config &cfg) { TVMFFICentral ¢ral = TVMFFICentral::getInstance(); // Checked ahead of the cache so that toggling the backend off still disables - // already-resolved configs; load_tvm_ffi_function reports the "disabled" reason. + // already-resolved configs. if (!central.get_cutedsl_backend_enabled()) { - return central.load_tvm_ffi_function(cfg); + return std::nullopt; } // Otherwise try the cache first, and ask Python to compile/register the kernel if not found. const uint32_t id = cfg.to_id(); From 62dd9ec72f05d63c0f049fbc23c186fcd6c60cea Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 28 Oct 2026 18:44:52 +0000 Subject: [PATCH 6/6] remove cuda-bindings from setup.py Signed-off-by: Kaining Zhong --- setup.py | 1 - 1 file changed, 1 deletion(-) diff --git a/setup.py b/setup.py index 7e04fcdab29..2dabf37b322 100644 --- a/setup.py +++ b/setup.py @@ -135,7 +135,6 @@ def setup_requirements() -> Tuple[List[str], List[str]]: "pydantic", "importlib-metadata>=1.0", "packaging", - "cuda-bindings", ] test_reqs: List[str] = ["pytest>=8.2.1"]