From dbcd1649f17f9a8a1b3c15481c6ac9fa9824f5a8 Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:15:28 -0700 Subject: [PATCH 01/11] Fix libcudart.so.13 hard dependency in pybind module breaking import on CPU-only Linux (#29590) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description `onnxruntime-gpu` 1.27 introduced a hard `NEEDED libcudart.so.13` entry in `onnxruntime_pybind11_state.so`, causing `ImportError` at `import onnxruntime` on CPU-only Linux machines — before any provider is selected. **Root cause:** `cmake/onnxruntime_python.cmake` was changed to compile `fpA_intB_gemm_adaptor.cu` and `fpA_intB_gemm_preprocessors_impl.cu` directly into `onnxruntime_pybind11_state.so` and link `CUDA::cudart` (dynamic). This embeds a load-time CUDA dependency in the Python module itself. **Fix:** Move the CUDA weight-preprocessing entry point (`pack_weights_for_cuda_mixed_gemm`) out of the main pybind module and into a **standalone extension module**, `onnxruntime_cuda_quant_preprocess`, that links `CUDA::cudart` on its own. The main `onnxruntime_pybind11_state.so` no longer compiles or links any CUDA code, so `import onnxruntime` has no `libcudart` dependency. The new module is imported **lazily** by `onnxruntime/python/tools/quantization/cuda_quantizer.py` only when weight prepacking is actually requested — never at `import onnxruntime` time. These preprocessing APIs are **offline-only** helpers: they are used by quantization tooling and model builders to produce prepacked weight initializers ahead of time, and are not part of the inference runtime hot path. Because nothing in the runtime imports them, isolating them into a separate, on-demand DLL has no runtime cost and cleanly keeps CUDA out of the base `import onnxruntime` path. **Why not the provider bridge:** An earlier iteration routed the call through the `ProviderInfo_CUDA` virtual interface (`TryGetProviderInfo_CUDA()`). That does not work for the CUDA-EP-as-plugin build (`onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON`): `cuda_provider_factory.cc` is excluded from the plugin sources and there is no provider bridge, so `TryGetProviderInfo_CUDA()` returns `nullptr` and the call throws. The standalone module has no such dependency and works for **both** the legacy in-tree CUDA EP build and the plugin build. ### Key Changes | File | Change | |---|---| | `onnxruntime/python/onnxruntime_pybind_cuda_quant.cc` | **New.** Self-contained `pack_weights_for_cuda_mixed_gemm` (device malloc + transpose/convert + arch permutation) and a `PYBIND11_MODULE(onnxruntime_cuda_quant_preprocess, …)` entry point. | | `cmake/onnxruntime_python.cmake` | Add the `onnxruntime_cuda_quant_preprocess` module target (built when `onnxruntime_USE_CUDA AND NOT WIN32`, compiling the two `fpA_intB` `.cu` files + `CUDA::cudart` + cutlass, hidden visibility) and copy it into `onnxruntime/capi/`. Main pybind module keeps no CUDA sources/links. | | `onnxruntime/python/onnxruntime_pybind_quant.cc` | Remove the `USE_CUDA` `PackWeightsForMixedGemm` and its registration. The CPU-only `pack_fp4_weights_for_cuda_moe_gemm` stays in the main module. | | `onnxruntime/core/providers/cuda/cuda_provider_factory.{h,cc}` | Revert the `PackWeightsForMixedGemm` `ProviderInfo_CUDA` addition (no longer needed; absent in plugin builds). | | `onnxruntime/python/tools/quantization/cuda_quantizer.py` | `_get_pack_weights_for_cuda_mixed_gemm()` now imports `onnxruntime.capi.onnxruntime_cuda_quant_preprocess` lazily; add `has_cuda_weight_prepacking()` capability helper. | | `setup.py` | Package `onnxruntime_cuda_quant_preprocess.so` in the Linux/macOS wheels. | | `onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py` | Point the prepacked-weight parity test and its skip guard at the new module. | | `docs/contrib_ops/cuda/matmul_nbits.md` | Update the offline-packer code snippets to import the new module. | ### Motivation and Context `import onnxruntime` must succeed on CPU-only machines even when the GPU wheel is installed. CUDA dependency errors should surface only when a CUDA provider is explicitly loaded/selected, or when offline CUDA weight prepacking is explicitly requested. This restores the 1.26 behavior where `onnxruntime_pybind11_state.so` had no `NEEDED libcudart.so.*` entry, and — unlike the provider-bridge approach — it also works in the CUDA-EP-as-plugin build. ### Testing Notes - Built both modules in the CUDA build; `readelf -d onnxruntime_pybind11_state.so` shows **no** `libcudart` `NEEDED` entry, while `onnxruntime_cuda_quant_preprocess.so` has `NEEDED libcudart.so.13`. - `import onnxruntime` and lazy loading of `onnxruntime.capi.onnxruntime_cuda_quant_preprocess` both succeed; `has_cuda_weight_prepacking()` returns `True` on a CUDA machine. - `test_op_matmulnbits_prepacked_cuda.py` passes (INT4/INT8 prepacked-vs-runtime parity), confirming the relocated packer produces byte-identical prepacked weights. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Tianlei Wu --- cmake/CMakeLists.txt | 1 + cmake/onnxruntime_python.cmake | 93 +++++-- docs/contrib_ops/cuda/matmul_nbits.md | 6 +- docs/cuda_plugin_ep/QUICK_START.md | 2 +- docs/cuda_plugin_ep/cuda_plugin_ep_design.md | 8 +- .../python/onnxruntime_pybind_cuda_quant.cc | 166 +++++++++++++ .../python/onnxruntime_pybind_mlvalue.cc | 29 --- .../python/onnxruntime_pybind_ortvalue.cc | 9 - .../python/onnxruntime_pybind_quant.cc | 143 +---------- .../tools/quantization/cuda_quantizer.py | 226 ++++++++++++++++-- .../test_op_matmulnbits_prepacked_cuda.py | 44 +++- setup.py | 9 + tools/ci_build/build.py | 4 +- 13 files changed, 519 insertions(+), 221 deletions(-) create mode 100644 onnxruntime/python/onnxruntime_pybind_cuda_quant.cc diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 09e307e124316..ad446214cfc8f 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -77,6 +77,7 @@ cmake_dependent_option(onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS "Build with CUD cmake_dependent_option(onnxruntime_USE_CUDA_NHWC_OPS "Build CUDA with NHWC op support" ON "onnxruntime_USE_CUDA" OFF) cmake_dependent_option(onnxruntime_BUILD_CUDA_EP_AS_PLUGIN "Build CUDA EP as a separate plugin shared library instead of the legacy in-tree provider" OFF "onnxruntime_USE_CUDA" OFF) +option(onnxruntime_BUILD_CUDA_QUANT_PREPROCESS "Build CUDA weight-packing module onnxruntime_cuda_quant_preprocess.so" ON) option(onnxruntime_CUDA_MINIMAL "Build CUDA without any operations apart from memcpy ops. Usefuel for a very minial TRT build" OFF) option(onnxruntime_ENABLE_CUDA_LINE_NUMBER_INFO "When building with CUDA support, generate device code line number information." OFF) option(onnxruntime_USE_OPENVINO "Build with OpenVINO support" OFF) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 3f6976c3c8955..14b01f1f25473 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -20,6 +20,13 @@ file(GLOB onnxruntime_pybind_srcs CONFIGURE_DEPENDS ${onnxruntime_pybind_srcs_pattern} ) +# onnxruntime_pybind_cuda_quant.cc is compiled into the standalone +# onnxruntime_cuda_quant_preprocess extension module (see below), not into +# onnxruntime_pybind11_state. It includes and links CUDA::cudart, +# so compiling it into the main pybind module would break CPU-only builds and +# re-introduce the hard libcudart dependency this design avoids. +list(REMOVE_ITEM onnxruntime_pybind_srcs ${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_cuda_quant.cc) + if(onnxruntime_ENABLE_TRAINING) list(REMOVE_ITEM onnxruntime_pybind_srcs ${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_module.cc) endif() @@ -231,22 +238,11 @@ target_link_libraries(onnxruntime_pybind11_state PRIVATE Python::NumPy ) -# Starting with Python 3.8 on Windows, PATH environment variable are no longer used to resolve DLL dependencies -# for extension modules or libraries loaded via ctypes. -# To avoid package import issues, we do not link pybind module against the CUDA runtime on Windows, instead of -# os.add_dll_directory() to deal with CUDA paths. -if (onnxruntime_USE_CUDA AND NOT WIN32) - target_sources(onnxruntime_pybind11_state PRIVATE - "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" - "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu" - ) - include(cutlass) - target_include_directories(onnxruntime_pybind11_state PRIVATE ${cutlass_SOURCE_DIR}/include) - target_link_libraries(onnxruntime_pybind11_state PRIVATE CUDA::cudart) -endif() -if (onnxruntime_USE_CUDA AND WIN32) - target_compile_definitions(onnxruntime_pybind11_state PRIVATE ORT_NO_CUDA_IN_PYBIND) -endif() +# The CUDA quantization helpers (pack_weights_for_cuda_mixed_gemm) are built into a +# separate extension module (onnxruntime_cuda_quant_preprocess) that is imported on +# demand. Do NOT compile CUDA source files directly into onnxruntime_pybind11_state or +# link CUDA::cudart from it: that would create a hard libcudart.so dependency that +# prevents importing the Python module on CPU-only machines. set(onnxruntime_pybind11_state_dependencies ${onnxruntime_EXTERNAL_DEPENDENCIES} @@ -315,6 +311,71 @@ else() set_target_properties(onnxruntime_pybind11_state PROPERTIES SUFFIX ".so") endif() +# --------------------------------------------------------------------------- +# Standalone CUDA weight-preprocessing extension module. +# +# The CUDA weight-packing kernels (pack_weights_for_cuda_mixed_gemm) are compiled +# into their OWN Python extension module instead of onnxruntime_pybind11_state. +# This keeps the hard libcudart dependency out of the main pybind module so that +# `import onnxruntime` still works on CPU-only machines. +# +# Production weight packing is done in PyTorch (cuda_quantizer.py); this module is +# retained only as a byte-parity oracle for that PyTorch packer. It is gated by +# onnxruntime_BUILD_CUDA_QUANT_PREPROCESS. +# +# It does NOT go through the provider bridge / ProviderInfo_CUDA, so it works for +# both the legacy in-tree CUDA EP build and the CUDA-EP-as-plugin build. +# +# Not built on Windows: matching the previous behavior where CUDA runtime was not +# linked into Python extension modules (DLL search path constraints since +# Python 3.8), so pack_weights_for_cuda_mixed_gemm was unavailable there. +if (onnxruntime_USE_CUDA AND NOT WIN32 AND onnxruntime_BUILD_CUDA_QUANT_PREPROCESS) + onnxruntime_add_shared_library_module(onnxruntime_cuda_quant_preprocess + "${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_cuda_quant.cc" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu" + ) + include(cutlass) + onnxruntime_add_include_to_target(onnxruntime_cuda_quant_preprocess Python::Module onnxruntime_common) + target_include_directories(onnxruntime_cuda_quant_preprocess PRIVATE + ${ONNXRUNTIME_ROOT} + ${pybind11_INCLUDE_DIRS} + ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} + ${cutlass_SOURCE_DIR}/include + ${cutlass_SOURCE_DIR}/tools/util/include + ) + target_compile_definitions(onnxruntime_cuda_quant_preprocess PRIVATE USE_CUDA) + target_link_libraries(onnxruntime_cuda_quant_preprocess PRIVATE + onnxruntime_common + Boost::mp11 + safeint_interface + ${ABSEIL_LIBS} + CUDA::cudart + ${pybind11_lib} + Python::NumPy + ) + if (NOT MSVC) + target_compile_options(onnxruntime_cuda_quant_preprocess PRIVATE "-fvisibility=hidden") + endif() + set_target_properties(onnxruntime_cuda_quant_preprocess PROPERTIES PREFIX "" SUFFIX ".so" FOLDER "ONNXRuntime") + if (APPLE) + set_target_properties(onnxruntime_cuda_quant_preprocess PROPERTIES + INSTALL_RPATH "@loader_path" + BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH_USE_LINK_PATH FALSE) + elseif (NOT CMAKE_SYSTEM_NAME MATCHES "AIX") + target_link_options(onnxruntime_cuda_quant_preprocess PRIVATE "LINKER:-rpath=\$ORIGIN") + endif() + # Place the module next to the main pybind module inside onnxruntime/capi. + add_custom_command( + TARGET onnxruntime_cuda_quant_preprocess POST_BUILD + COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/capi + COMMAND ${CMAKE_COMMAND} -E copy + $ + $/onnxruntime/capi/ + ) +endif() + # Generate version_info.py in Windows build. # Has to be done before onnxruntime_python_srcs is set. if (WIN32) diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index a09a80a4732ed..3f8e1f1d8f616 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -87,9 +87,9 @@ step is **not** performed. The offline CUDA packer exposed through Python produces this layout: ```python -from onnxruntime.capi import _pybind_state as _pybind +from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant -prepacked_flat = _pybind.pack_weights_for_cuda_mixed_gemm( +prepacked_flat = _cuda_quant.pack_weights_for_cuda_mixed_gemm( q_weight.reshape(N, -1), N, K, bits, 80 ) prepacked_b = np.asarray(prepacked_flat, dtype=np.int8).view(np.uint8).reshape(q_weight.shape) @@ -309,7 +309,7 @@ present. `ComputeInternal` then: GEMV profiling helpers, e.g. `profile_qmoe_gemv.sh`). - CUDA prepacked-weight parity tests: [onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py](../../../onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py). - These use `_pybind_state.pack_weights_for_cuda_mixed_gemm(..., 80)` to produce + These use `onnxruntime_cuda_quant_preprocess.pack_weights_for_cuda_mixed_gemm(..., 80)` to produce `weight_prepacked=1` initializers and compare their outputs against runtime fpA_intB prepacking for int4/int8 and GEMV/GEMM-shaped `M` values. - Constructor failure tests for unsupported prepacked configurations live in diff --git a/docs/cuda_plugin_ep/QUICK_START.md b/docs/cuda_plugin_ep/QUICK_START.md index ab7b4308f13e5..7b971d621083f 100644 --- a/docs/cuda_plugin_ep/QUICK_START.md +++ b/docs/cuda_plugin_ep/QUICK_START.md @@ -146,7 +146,7 @@ sess = ort.InferenceSession( **Python `OrtValue` host/device copies:** -`OrtValue.update_inplace()` and `OrtValue.numpy()` work with CUDA plugin tensors after the plugin has been registered. On Linux, the ONNX Runtime Python binding links the CUDA runtime and can fall back to direct `cudaMemcpy` if the legacy CUDA provider bridge is unavailable. On Windows, the Python binding is built with `ORT_NO_CUDA_IN_PYBIND`, so it cannot call CUDA runtime APIs directly; host/device copies must use the data-transfer implementation registered by the CUDA plugin library. If `OrtValue.update_inplace()` fails with a message about the CUDA provider interface or an unsupported GPU device, verify that the plugin library is registered before creating or updating CUDA `OrtValue` objects. +`OrtValue.update_inplace()` and `OrtValue.numpy()` work with CUDA plugin tensors after the plugin has been registered. The Python binding cannot call CUDA runtime APIs directly; host/device copies must use the data-transfer implementation registered by the CUDA plugin library. If `OrtValue.update_inplace()` fails with a message about the CUDA provider interface or an unsupported GPU device, verify that the plugin library is registered before creating or updating CUDA `OrtValue` objects. ### External GPU Allocator Options diff --git a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md index a06f1de95b99e..7dfa2bd2d8667 100644 --- a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md +++ b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md @@ -344,10 +344,12 @@ This is intentionally conservative and correct for the plugin EP's first sync in The Python `OrtValue` helpers (`update_inplace()` for host-to-device and `numpy()` for device-to-host) historically reached CUDA copies through the legacy provider bridge (`GetProviderInfo_CUDA()`). That bridge requires the provider shared library to export `GetProvider()`, which the CUDA plugin intentionally does not export. -The fallback path is platform-specific: +To keep working when the bridge is absent (as with the plugin EP), the pybind can reach CUDA copies two ways: the legacy provider bridge (`TryGetProviderInfo_CUDA()`) and a plugin-registered `OrtDataTransfer` copy function (`CreateDataTransferMemCpy()`, backed by the plugin EP's `IDataTransfer`). It tries whichever is available and throws if neither is, in which case a CUDA `OrtValue` copy cannot be performed. -- On non-Windows CUDA builds, `onnxruntime_pybind11_state` links `CUDA::cudart`. If `TryGetProviderInfo_CUDA()` fails, pybind can copy directly with `cudaMemcpy`; host-to-device copies synchronize the default stream, matching `ProviderInfo_CUDA::cudaMemcpy_HostToDevice()`. -- On Windows CUDA builds, pybind is compiled with `ORT_NO_CUDA_IN_PYBIND` and does not link CUDA runtime APIs. If `TryGetProviderInfo_CUDA()` fails, pybind must obtain an `OrtDataTransfer` copy function from the registered plugin EP. Without a registered plugin data-transfer implementation, CUDA `OrtValue.update_inplace()` cannot copy host data into the plugin-owned device tensor. +The two code paths differ only in which mechanism they try first, and this does not change the outcome (exactly one applies in a given build): + +- `OrtValue.update_inplace(numpy_array)` / `OrtValue.numpy()` (in `onnxruntime_pybind_ortvalue.cc`) try the provider bridge first, then fall back to the plugin `OrtDataTransfer`. +- `OrtValue.update_inplace(OrtValue)` (`UpdateOrtValueInplace` in `onnxruntime_pybind_mlvalue.cc`) tries the plugin `OrtDataTransfer` first, then falls back to the built-in CUDA provider copy functions. ### 5.2 Handle Access Path diff --git a/onnxruntime/python/onnxruntime_pybind_cuda_quant.cc b/onnxruntime/python/onnxruntime_pybind_cuda_quant.cc new file mode 100644 index 0000000000000..3edadd786a8c9 --- /dev/null +++ b/onnxruntime/python/onnxruntime_pybind_cuda_quant.cc @@ -0,0 +1,166 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Standalone CUDA weight-preprocessing extension module. +// +// This module is intentionally kept SEPARATE from onnxruntime_pybind11_state so +// that `import onnxruntime` never triggers a load-time dependency on the CUDA +// runtime (libcudart). The CUDA weight packing kernels below link CUDA::cudart, +// so this module has a hard libcudart dependency -- but it is imported lazily by +// onnxruntime.python.tools.quantization.cuda_quantizer only when CUDA weight +// prepacking is actually requested. +// +// This approach works for both the legacy in-tree CUDA EP build and the +// CUDA-EP-as-plugin build (onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON), because it +// does not rely on the provider bridge / ProviderInfo_CUDA interface (which is +// not available in plugin builds). + +#include +#include + +#include + +#include +#include +#include + +#include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" +#include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" + +namespace py = pybind11; + +namespace { + +void ThrowIfCudaError(cudaError_t status, const char* expression) { + if (status != cudaSuccess) { + std::ostringstream oss; + oss << expression << " failed: " << cudaGetErrorString(status); + throw std::runtime_error(oss.str()); + } +} + +struct CudaDeleter { + void operator()(void* p) const { + if (p) cudaFree(p); + } +}; + +using CudaPtr = std::unique_ptr; + +// Preprocess quantized weights for CUDA mixed-precision GEMM kernels (FpA_IntB format). +// +// MatMulNBits/QMoE stores quantized weights in (N, K) layout: +// - N = number of output channels (columns in weight matrix W) +// - K = number of input features (rows in weight matrix W) +// - For 4-bit: shape is (N, K/2) bytes where each byte packs 2 elements +// - For 8-bit: shape is (N, K) bytes +// +// FpA_IntB GEMM kernels expect weights in (K, N) layout (transposed) for efficient +// memory access during matrix multiplication. This function: +// 1. Transposes from (N, K) to (K, N) layout +// 2. Converts unsigned quantized values to signed int8 with zero-point adjustment +// - 4-bit: uint4 -> int8 with zero_point=8 (range [0,15] -> [-8,7]) +// - 8-bit: uint8 -> int8 with zero_point=128 (range [0,255] -> [-128,127]) +// 3. Applies architecture-specific row permutation for optimized tensor core access +// +// Input: q_weights - Quantized weights from MatMulNBits in (N, K) layout +// Output: Preprocessed weights in (K, N) layout ready for fpA_intB GEMM kernels +py::array_t PackWeightsForMixedGemm( + py::array_t q_weights, + int32_t N, + int32_t K, + int32_t bits, + int32_t force_arch = -1) { + py::buffer_info q_weights_buf = q_weights.request(); + + if (bits != 4 && bits != 8) { + throw std::invalid_argument("bits must be 4 or 8"); + } + if (N <= 0 || K <= 0) { + throw std::invalid_argument("N and K must be positive"); + } + if (bits == 4 && K % 2 != 0) { + throw std::invalid_argument("K must be even for 4-bit packed weights"); + } + if (q_weights_buf.ndim != 2 || q_weights_buf.shape[0] != N || q_weights_buf.shape[1] != K / (8 / bits)) { + throw std::invalid_argument("q_weights must have shape (N, K / (8 / bits))"); + } + + int n = static_cast(N); + int k = static_cast(K); + + size_t packed_weight_bytes = static_cast(n) * static_cast(k) / (8 / bits); + py::array_t processed_weights({static_cast(packed_weight_bytes)}); + py::buffer_info processed_weights_buf = processed_weights.request(); + + auto make_cuda_ptr = [](size_t bytes) -> CudaPtr { + void* p = nullptr; + ThrowIfCudaError(cudaMalloc(&p, bytes), "cudaMalloc"); + return CudaPtr(p); + }; + + auto packed_transposed_weight_space = make_cuda_ptr(packed_weight_bytes); + int8_t* packed_transposed_weight = reinterpret_cast(packed_transposed_weight_space.get()); + + auto fpA_intB_weight_buffer_ = make_cuda_ptr(packed_weight_bytes); + int8_t* preprocessed_weight = reinterpret_cast(fpA_intB_weight_buffer_.get()); + + const uint8_t* blob_data_cpu = static_cast(q_weights_buf.ptr); + + auto blob_data_gpu_buf = make_cuda_ptr(packed_weight_bytes); + uint8_t* blob_data_gpu = reinterpret_cast(blob_data_gpu_buf.get()); + + cudaStream_t stream = cudaStreamLegacy; + ThrowIfCudaError(cudaMemcpyAsync(blob_data_gpu, blob_data_cpu, packed_weight_bytes, cudaMemcpyHostToDevice, stream), + "cudaMemcpyAsync host-to-device"); + + if (bits == 4) { + ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( + stream, packed_transposed_weight, blob_data_gpu, n, k); + } else { + // 8 bits + ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( + stream, packed_transposed_weight, blob_data_gpu, n, k); + } + + using ::onnxruntime::llm::kernels::weight_only::QuantType; + QuantType quant_type = bits == 4 ? QuantType::W4_A16 : QuantType::W8_A16; + + int sm = force_arch; + if (sm < 0) { + int device_id = 0; + ThrowIfCudaError(cudaGetDevice(&device_id), "cudaGetDevice"); + cudaDeviceProp device_prop; + ThrowIfCudaError(cudaGetDeviceProperties(&device_prop, device_id), "cudaGetDeviceProperties"); + sm = device_prop.major * 10 + device_prop.minor; + } + sm = ::onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(sm); + + auto permutation_map_buffer = make_cuda_ptr(32 * sizeof(int32_t)); + + ::onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( + stream, + sm, + preprocessed_weight, + packed_transposed_weight, + reinterpret_cast(permutation_map_buffer.get()), + {static_cast(k), static_cast(n)}, + quant_type); + + ThrowIfCudaError(cudaGetLastError(), "preprocess CUDA kernel launch"); + ThrowIfCudaError(cudaMemcpyAsync(processed_weights_buf.ptr, preprocessed_weight, packed_weight_bytes, + cudaMemcpyDeviceToHost, stream), + "cudaMemcpyAsync device-to-host"); + ThrowIfCudaError(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); + + return processed_weights; +} + +} // namespace + +PYBIND11_MODULE(onnxruntime_cuda_quant_preprocess, m) { + m.doc() = "CUDA weight-only quantization preprocessing helpers (loaded on demand)."; + m.def("pack_weights_for_cuda_mixed_gemm", &PackWeightsForMixedGemm, + "Pack quantized weights for CUDA mixed-precision GEMM (FpA_IntB format)", + py::arg("q_weights"), py::arg("N"), py::arg("K"), py::arg("bits"), py::arg("force_arch") = -1); +} diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc index 10e55259f834e..5ea5f42958926 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc @@ -23,10 +23,6 @@ #include "core/framework/kernel_registry.h" #include "core/framework/provider_options_utils.h" -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) -#include -#endif - #ifdef USE_DML using Microsoft::WRL::ComPtr; @@ -184,23 +180,6 @@ int32_t GetTensorProtoType(const OrtValue& ort_value) { } #ifdef USE_CUDA -namespace { - -#if !defined(ORT_NO_CUDA_IN_PYBIND) -void CudaRuntimeMemCpy(void* dst, const void* src, size_t num_bytes, cudaMemcpyKind kind) { - const auto copy_result = cudaMemcpy(dst, src, num_bytes, kind); - ORT_ENFORCE(copy_result == cudaSuccess, "cudaMemcpy failed: ", cudaGetErrorString(copy_result)); - - if (kind == cudaMemcpyHostToDevice) { - // Match ProviderInfo_CUDA::cudaMemcpy_HostToDevice: cudaMemcpy() uses the default - // stream, and pageable host-to-device copies can return before DMA to device is done. - const auto sync_result = cudaStreamSynchronize(0); - ORT_ENFORCE(sync_result == cudaSuccess, "cudaStreamSynchronize failed: ", cudaGetErrorString(sync_result)); - } -} -#endif - -} // namespace void CpuToCudaMemCpy(void* dst, const void* src, size_t num_bytes) { if (TryGetProviderInfo_CUDA() != nullptr) { @@ -208,11 +187,7 @@ void CpuToCudaMemCpy(void* dst, const void* src, size_t num_bytes) { return; } -#if !defined(ORT_NO_CUDA_IN_PYBIND) - CudaRuntimeMemCpy(dst, src, num_bytes, cudaMemcpyHostToDevice); -#else ORT_THROW("CUDA provider interface is not available for host-to-device copy."); -#endif } void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { @@ -221,11 +196,7 @@ void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { return; } -#if !defined(ORT_NO_CUDA_IN_PYBIND) - CudaRuntimeMemCpy(dst, src, num_bytes, cudaMemcpyDeviceToHost); -#else ORT_THROW("CUDA provider interface is not available for device-to-host copy."); -#endif } const std::unordered_map* GetCudaToHostMemCpyFunction(const OrtDevice& device) { diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index cf7f86a0b9e41..7bf9325cf2208 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -205,7 +205,6 @@ void addOrtValueMethods(pybind11::module& m) { #ifdef USE_CUDA if (device.Vendor() == OrtDevice::VendorIds::NVIDIA) { MemCpyFunc cpu_to_device_copy_fn = CpuToCudaMemCpy; -#if defined(ORT_NO_CUDA_IN_PYBIND) if (TryGetProviderInfo_CUDA() != nullptr) { if (!IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); @@ -217,12 +216,6 @@ void addOrtValueMethods(pybind11::module& m) { "Unsupported GPU device: Cannot find the supported GPU device."); } } -#else - if (TryGetProviderInfo_CUDA() != nullptr && - !IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { - throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); - } -#endif onnxruntime::python::CopyDataToTensor( py_values, @@ -467,12 +460,10 @@ void addOrtValueMethods(pybind11::module& m) { switch (device.Vendor()) { #ifdef USE_CUDA case OrtDevice::VendorIds::NVIDIA: -#if defined(ORT_NO_CUDA_IN_PYBIND) if (TryGetProviderInfo_CUDA() == nullptr) { return GetPyObjFromTensor(*ml_value, nullptr, nullptr, /*zero_copy_non_owning=*/true); } -#endif return GetPyObjFromTensor(*ml_value, nullptr, GetCudaToHostMemCpyFunction(device), /*zero_copy_non_owning=*/true); #endif diff --git a/onnxruntime/python/onnxruntime_pybind_quant.cc b/onnxruntime/python/onnxruntime_pybind_quant.cc index 7220153b4fa17..b2b6ed77c6296 100644 --- a/onnxruntime/python/onnxruntime_pybind_quant.cc +++ b/onnxruntime/python/onnxruntime_pybind_quant.cc @@ -9,11 +9,6 @@ #include "contrib_ops/cpu/quantization/dequantize_blockwise_bnb4.h" #include "core/util/thread_utils.h" -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) -#include -#include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" -#include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" -#endif #include #include #include @@ -147,132 +142,6 @@ void QuantizeMatMulBnb4Blockwise( tp.get()); } -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) -namespace cuda { -void ThrowIfCudaError(cudaError_t status, const char* expression) { - if (status != cudaSuccess) { - std::ostringstream oss; - oss << expression << " failed: " << cudaGetErrorString(status); - throw std::runtime_error(oss.str()); - } -} - -struct CudaDeleter { - void operator()(void* p) const { - if (p) cudaFree(p); - } -}; - -using CudaPtr = std::unique_ptr; - -// Preprocess quantized weights for CUDA mixed-precision GEMM kernels (FpA_IntB format). -// -// MatMulNBits/QMoE stores quantized weights in (N, K) layout: -// - N = number of output channels (columns in weight matrix W) -// - K = number of input features (rows in weight matrix W) -// - For 4-bit: shape is (N, K/2) bytes where each byte packs 2 elements -// - For 8-bit: shape is (N, K) bytes -// -// FpA_IntB GEMM kernels expect weights in (K, N) layout (transposed) for efficient -// memory access during matrix multiplication. This function: -// 1. Transposes from (N, K) to (K, N) layout -// 2. Converts unsigned quantized values to signed int8 with zero-point adjustment -// - 4-bit: uint4 -> int8 with zero_point=8 (range [0,15] -> [-8,7]) -// - 8-bit: uint8 -> int8 with zero_point=128 (range [0,255] -> [-128,127]) -// 3. Applies architecture-specific row permutation for optimized tensor core access -// -// Input: q_weights - Quantized weights from MatMulNBits in (N, K) layout -// Output: Preprocessed weights in (K, N) layout ready for fpA_intB GEMM kernels -py::array_t PackWeightsForMixedGemm( - py::array_t q_weights, - int32_t N, - int32_t K, - int32_t bits, - int32_t force_arch = -1) { - py::buffer_info q_weights_buf = q_weights.request(); - - if (bits != 4 && bits != 8) { - throw std::invalid_argument("bits must be 4 or 8"); - } - if (N <= 0 || K <= 0) { - throw std::invalid_argument("N and K must be positive"); - } - if (bits == 4 && K % 2 != 0) { - throw std::invalid_argument("K must be even for 4-bit packed weights"); - } - if (q_weights_buf.ndim != 2 || q_weights_buf.shape[0] != N || q_weights_buf.shape[1] != K / (8 / bits)) { - throw std::invalid_argument("q_weights must have shape (N, K / (8 / bits))"); - } - - int n = static_cast(N); - int k = static_cast(K); - - size_t packed_weight_bytes = static_cast(n) * static_cast(k) / (8 / bits); - py::array_t processed_weights({static_cast(packed_weight_bytes)}); - py::buffer_info processed_weights_buf = processed_weights.request(); - - auto make_cuda_ptr = [](size_t bytes) -> CudaPtr { - void* p = nullptr; - ThrowIfCudaError(cudaMalloc(&p, bytes), "cudaMalloc"); - return CudaPtr(p); - }; - - auto packed_transposed_weight_space = make_cuda_ptr(packed_weight_bytes); - int8_t* packed_transposed_weight = reinterpret_cast(packed_transposed_weight_space.get()); - - auto fpA_intB_weight_buffer_ = make_cuda_ptr(packed_weight_bytes); - int8_t* preprocessed_weight = reinterpret_cast(fpA_intB_weight_buffer_.get()); - - const uint8_t* blob_data_cpu = static_cast(q_weights_buf.ptr); - - auto blob_data_gpu_buf = make_cuda_ptr(packed_weight_bytes); - uint8_t* blob_data_gpu = reinterpret_cast(blob_data_gpu_buf.get()); - - cudaStream_t stream = cudaStreamLegacy; - ThrowIfCudaError(cudaMemcpyAsync(blob_data_gpu, blob_data_cpu, packed_weight_bytes, cudaMemcpyHostToDevice, stream), - "cudaMemcpyAsync host-to-device"); - - if (bits == 4) { - ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( - stream, packed_transposed_weight, blob_data_gpu, n, k); - } else { - // 8 bits - ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( - stream, packed_transposed_weight, blob_data_gpu, n, k); - } - - using ::onnxruntime::llm::kernels::weight_only::QuantType; - QuantType quant_type = bits == 4 ? QuantType::W4_A16 : QuantType::W8_A16; - - int sm = force_arch; - if (sm < 0) { - int device_id = 0; - ThrowIfCudaError(cudaGetDevice(&device_id), "cudaGetDevice"); - cudaDeviceProp device_prop; - ThrowIfCudaError(cudaGetDeviceProperties(&device_prop, device_id), "cudaGetDeviceProperties"); - sm = device_prop.major * 10 + device_prop.minor; - } - sm = ::onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(sm); - - auto permutation_map_buffer = make_cuda_ptr(32 * sizeof(int32_t)); - - ::onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( - stream, - sm, - preprocessed_weight, - packed_transposed_weight, - reinterpret_cast(permutation_map_buffer.get()), - {static_cast(k), static_cast(n)}, - quant_type); - - ThrowIfCudaError(cudaGetLastError(), "preprocess CUDA kernel launch"); - ThrowIfCudaError(cudaMemcpyAsync(processed_weights_buf.ptr, preprocessed_weight, packed_weight_bytes, cudaMemcpyDeviceToHost, stream), - "cudaMemcpyAsync device-to-host"); - ThrowIfCudaError(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); - - return processed_weights; -} - // Pack FP4 (MXFP4) weights for MoE GEMM kernels. // // Input: q_weights in [N, K/2] layout (FP4 packed 2 per byte along K dimension, row-major) @@ -280,6 +149,7 @@ py::array_t PackWeightsForMixedGemm( // // Unlike INT4 which requires architecture-specific row permutation and interleaving, // FP4 (SM90+ TMA path) only needs a simple transpose at the nibble level. +// This function is CPU-only and does not require CUDA to be present. py::array_t PackFP4WeightsForMoE( py::array_t q_weights, int32_t N, @@ -300,7 +170,7 @@ py::array_t PackFP4WeightsForMoE( int K_half = K / 2; int N_half = N / 2; size_t out_size = static_cast(K) * static_cast(N_half); - py::array_t output({static_cast(out_size)}); + py::array_t output(static_cast(out_size)); py::buffer_info out_buf = output.request(); uint8_t* dst = static_cast(out_buf.ptr); std::memset(dst, 0, out_size); @@ -327,8 +197,6 @@ py::array_t PackFP4WeightsForMoE( return output; } -} // namespace cuda -#endif void CreateQuantPybindModule(py::module& m) { m.def("quantize_matmul_2bits", &QuantizeMatMulNBitsBlockwise); @@ -343,14 +211,9 @@ void CreateQuantPybindModule(py::module& m) { m.def("quantize_qdq_matmul_2bits", &QuantizeQDQMatMulNBitsBlockwise); m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) - m.def("pack_weights_for_cuda_mixed_gemm", &cuda::PackWeightsForMixedGemm, - "Pack quantized weights for CUDA mixed-precision GEMM (FpA_IntB format)", - py::arg("q_weights"), py::arg("N"), py::arg("K"), py::arg("bits"), py::arg("force_arch") = -1); - m.def("pack_fp4_weights_for_cuda_moe_gemm", &cuda::PackFP4WeightsForMoE, + m.def("pack_fp4_weights_for_cuda_moe_gemm", &PackFP4WeightsForMoE, "Pack FP4 (MXFP4) weights for CUDA MoE GEMM: transpose [N,K/2] to column-major [K,N/2]", py::arg("q_weights"), py::arg("N"), py::arg("K")); -#endif } } // namespace python diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py index 6b03485280d62..2ab532b3d0e58 100644 --- a/onnxruntime/python/tools/quantization/cuda_quantizer.py +++ b/onnxruntime/python/tools/quantization/cuda_quantizer.py @@ -7,9 +7,13 @@ """CUDA weight-only quantization helpers. This module contains small Python utilities for producing the weight layouts -consumed by CUDA weight-only kernels. The helpers deliberately wrap the same C++ -pybind entry points used by runtime prepacking so tests and model builders can -generate byte-identical quantized weights. +consumed by CUDA weight-only kernels. The blockwise quantizers wrap the same C++ +pybind entry points used by runtime prepacking, and the mixed-GEMM weight packer +is a PyTorch reimplementation of the runtime CUDA packing, so tests and model +builders can generate byte-identical quantized weights. The PyTorch packer runs +on CUDA when a device is available and falls back to CPU otherwise, which is the +only option on platforms where the standalone CUDA packer is not built (Windows). +A GPU-gated parity test validates it against that standalone CUDA packer. Two storage families are exposed: @@ -24,10 +28,14 @@ from __future__ import annotations +import functools +import logging from typing import TYPE_CHECKING import numpy as np +_logger = logging.getLogger(__name__) + if TYPE_CHECKING: import torch @@ -43,20 +51,207 @@ def _get_torch(): def _get_pack_weights_for_cuda_mixed_gemm(): - """Return the CUDA mixed-GEMM weight prepacker from the ORT pybind module.""" + """Return the standalone CUDA mixed-GEMM weight packer (parity oracle). + + Production packing uses the PyTorch implementation (``_pack_weights_for_cuda_mixed_gemm``). + This standalone packer lives in ``onnxruntime.capi.onnxruntime_cuda_quant_preprocess``, a + separate extension module that links the CUDA runtime (built only on non-Windows CUDA + builds). It is imported lazily here (never at ``import onnxruntime`` time) and is used by + the parity test to validate the PyTorch packer byte-for-byte. + """ try: - from onnxruntime.capi import _pybind_state as _pybind # noqa: PLC0415 + from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant # noqa: PLC0415 except ImportError as e: raise ImportError( - "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." + "The standalone CUDA weight packer (onnxruntime_cuda_quant_preprocess) is unavailable; " + "it is built only on non-Windows onnxruntime-gpu CUDA builds." ) from e try: - return _pybind.pack_weights_for_cuda_mixed_gemm + return _cuda_quant.pack_weights_for_cuda_mixed_gemm except AttributeError as e: - raise ImportError( - "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." - ) from e + raise ImportError("onnxruntime_cuda_quant_preprocess is missing pack_weights_for_cuda_mixed_gemm.") from e + + +def has_cuda_weight_prepacking() -> bool: + """Return True if mixed-GEMM weight prepacking is available. + + Prepacking is implemented with PyTorch (CUDA when available, CPU otherwise), so it is + available whenever torch is importable. Callers use this to skip prepack code paths + (and tests) when torch is unavailable. + """ + try: + _get_torch() + except ImportError: + return False + return True + + +@functools.lru_cache(maxsize=1) +def _warn_cpu_prepack_once() -> None: + _logger.warning( + "CUDA device is not available; packing mixed-GEMM weights on CPU with PyTorch. " + "This is correct but significantly slower for large Mixture-of-Experts models. " + "Pack on a CUDA-enabled machine for best performance." + ) + + +def _prepack_device(): + """Pick the torch device for mixed-GEMM weight packing (CUDA if available, else CPU).""" + torch = _get_torch() + if torch.cuda.is_available(): + return torch.device("cuda") + _warn_cpu_prepack_once() + return torch.device("cpu") + + +def _preprocess_weights_for_mixed_gemm_torch(tensor, bits: int, sm: int): + """PyTorch port of the runtime CUDA ``preprocess_weights_for_mixed_gemm``. + + ``tensor`` is a signed int8 weight in ``(K, N/pack)`` packed row-major layout on any + device. Returns the CUTLASS mixed-GEMM layout with the same shape/dtype/device. This + mirrors ``preprocess_weights_for_mixed_gemm_cuda`` (permute_B_rows -> subbyte_transpose + -> interleave_column_major -> add_bias_and_interleave) so its output is byte-identical + to the standalone CUDA packer, for both the SM80 (Ampere) and SM90 (Hopper) layouts. + """ + torch = _get_torch() + bits_a = 16 # fp16/bf16 activations + bits_b = 4 if bits == 4 else 8 + + if tensor.dim() == 2: + tensor = tensor.unsqueeze(0) + + permutation_map = { + "16_8": [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15], + "16_4": [ + 0, + 1, + 8, + 9, + 16, + 17, + 24, + 25, + 2, + 3, + 10, + 11, + 18, + 19, + 26, + 27, + 4, + 5, + 12, + 13, + 20, + 21, + 28, + 29, + 6, + 7, + 14, + 15, + 22, + 23, + 30, + 31, + ], + } + mma_shape_n = 8 + b_rows_per_mma = 8 * 16 // bits_b + + num_experts, num_rows, num_cols = tensor.shape[0], tensor.shape[1], tensor.shape[2] + if num_rows % b_rows_per_mma != 0 or num_cols % mma_shape_n != 0: + raise ValueError( + f"weight shape (rows={num_rows}, packed_cols={num_cols}) is incompatible with mixed-GEMM " + f"packing (rows must be a multiple of {b_rows_per_mma}, packed cols a multiple of {mma_shape_n})." + ) + + # permute_B_rows_for_mixed_gemm + if sm < 100: + pmap = permutation_map[f"{bits_a}_{bits_b}"] + row_idx = [(r // b_rows_per_mma) * b_rows_per_mma + pmap[r % b_rows_per_mma] for r in range(num_rows)] + tensor = tensor[:, row_idx, :] + + # subbyte_transpose + original_shape = tensor.shape + if bits_b == 4: + u = tensor.view(torch.uint8) + high = (u >> 4).permute(0, 2, 1).unsqueeze(2) + low = ((u << 4) >> 4).permute(0, 2, 1).unsqueeze(2) + merged = torch.cat([low, high], dim=2).reshape(u.shape[0], -1, u.shape[1]) + merged = merged[:, :, 0::2] + merged[:, :, 1::2] * 16 + tensor = merged.view(torch.int8).reshape(original_shape) + else: + tensor = tensor.permute(0, 2, 1).reshape(original_shape) + + # interleave_column_major_tensor + interleave = bits_a // bits_b + if interleave > 1 and sm < 90: + rows_per_tile = 128 * 8 // bits_a + elts_in_int32 = 32 // bits_b + if num_rows % elts_in_int32 != 0 or num_rows % rows_per_tile != 0: + raise ValueError(f"num_rows ({num_rows}) is incompatible with column-interleave tiling.") + tensor = tensor.reshape( + num_experts, -1, interleave, num_rows // rows_per_tile, rows_per_tile * 4 // elts_in_int32 + ) + tensor = tensor.permute(0, 1, 3, 2, 4).reshape(original_shape) + + # add_bias_and_interleave_quantized_tensor_inplace + if bits_b == 8: + t = tensor.to(torch.int64) # widen so the +128 rebias cannot overflow int8 + t += -256 * (t > 127).to(torch.int64) + 128 + t = t.reshape(-1, 4)[:, [0, 2, 1, 3]].reshape(original_shape) + tensor = t.to(torch.uint8).view(torch.int8) + else: + u = tensor.view(torch.uint8) + high = (u >> 4).unsqueeze(-1) + low = ((u << 4) >> 4).unsqueeze(-1) + merged = torch.cat([low, high], dim=-1).reshape(u.shape[0], u.shape[1], -1) + merged = merged.reshape(-1, 8)[:, [0, 2, 4, 6, 1, 3, 5, 7]].reshape(merged.shape) + merged = merged.to(torch.int16) + merged += -16 * (merged > 7).to(torch.int16) + 8 + merged = merged[:, :, 0::2] + merged[:, :, 1::2] * 16 + tensor = merged.to(torch.uint8).view(torch.int8) + + return tensor.squeeze(0).contiguous() + + +def _pack_weights_for_cuda_mixed_gemm(q_weights, n: int, k: int, bits: int, force_arch: int = 80) -> np.ndarray: + """PyTorch implementation of the CUDA ``pack_weights_for_cuda_mixed_gemm``. + + ``q_weights`` is ORT's unsigned MatMulNBits/QMoE storage ``(N, K/pack)`` (uint8). Returns + a flat ``int8`` numpy array with the CUTLASS mixed-GEMM layout, byte-identical to the + standalone CUDA packer. Runs on CUDA when available, otherwise on CPU. + """ + torch = _get_torch() + bits = int(bits) + force_arch = int(force_arch) + if bits not in (4, 8): + raise ValueError(f"bits must be 4 or 8, got {bits}.") + if force_arch not in (80, 90): + raise ValueError(f"force_arch must be 80 (SM80) or 90 (SM90), got {force_arch}.") + pack = 8 // bits + device = _prepack_device() + + q = torch.as_tensor(np.ascontiguousarray(q_weights)).view(torch.uint8).reshape(n, k // pack).to(device) + + # Front-end adaptor: transpose ORT (N, K) -> (K, N) and convert unsigned -> signed int8. + if bits == 4: + low = (q & 0x0F).to(torch.int16) + high = (q >> 4).to(torch.int16) + unpacked = torch.empty((n, k), dtype=torch.int16, device=device) + unpacked[:, 0::2] = low + unpacked[:, 1::2] = high + signed_t = (unpacked - 8).transpose(0, 1).contiguous() # (K, N), zero point 8 + packed_t = ((signed_t[:, 0::2] & 0x0F) | ((signed_t[:, 1::2] & 0x0F) << 4)).to(torch.uint8).view(torch.int8) + else: + signed_t = (q.to(torch.int16) - 128).transpose(0, 1).contiguous() # (K, N), zero point 128 + packed_t = signed_t.to(torch.uint8).view(torch.int8) + + out = _preprocess_weights_for_mixed_gemm_torch(packed_t.contiguous(), bits, force_arch) + return out.reshape(-1).cpu().numpy() def _get_quantize_matmul_nbits(): @@ -146,8 +341,7 @@ def qmoe_per_channel_quantize( When ``prepack`` is true, returned weights have shape ``[K, N/pack]``. Otherwise, returned weights keep raw per-channel storage ``[N, K/pack]``. - CUDA prepacking requires ``pack_weights_for_cuda_mixed_gemm`` from an - onnxruntime-gpu CUDA build. + Prepacking uses PyTorch (CUDA when available, CPU otherwise). """ torch = _get_torch() @@ -159,14 +353,12 @@ def qmoe_per_channel_quantize( if not prepack: return qweight, scales - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() - n, k = weights.shape pack = 8 // int(bits) if n % pack != 0: raise ValueError(f"N ({n}) must be divisible by {pack} for CUDA QMoE prepacked weights.") - packed = pack_weights_for_cuda_mixed_gemm(qweight.numpy(), n, k, int(bits), force_arch) + packed = _pack_weights_for_cuda_mixed_gemm(qweight.numpy(), n, k, int(bits), force_arch) packed = np.asarray(packed).view(np.uint8).reshape(k, n // pack) return torch.from_numpy(np.ascontiguousarray(packed)), scales @@ -330,7 +522,7 @@ def matmulnbits_prepacked_blockwise_quantize( unsigned_full_range=unsigned_full_range, ) - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() + pack_weights_for_cuda_mixed_gemm = _pack_weights_for_cuda_mixed_gemm packed = pack_weights_for_cuda_mixed_gemm(qweight.reshape(n, -1).numpy(), n, k, bits, force_arch) packed = np.asarray(packed).view(np.uint8).reshape(qweight.shape) packed = torch.from_numpy(np.ascontiguousarray(packed)) @@ -375,7 +567,7 @@ def qmoe_prepacked_blockwise_quantize( unsigned_full_range=unsigned_full_range, ) - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() + pack_weights_for_cuda_mixed_gemm = _pack_weights_for_cuda_mixed_gemm packed = pack_weights_for_cuda_mixed_gemm(qweight.reshape(n, -1).numpy(), n, k, bits, force_arch) packed = np.asarray(packed).view(np.uint8).reshape(k, n // pack) torch = _get_torch() diff --git a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 406eee21c059f..9d3a378f6db19 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -16,6 +16,12 @@ import onnxruntime as ort from onnxruntime.capi import _pybind_state as _pybind +from onnxruntime.quantization.cuda_quantizer import _pack_weights_for_cuda_mixed_gemm + +try: + from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant +except ImportError: + _cuda_quant = None @contextmanager @@ -32,7 +38,7 @@ def set_env(name: str, value: str): @unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") -@unittest.skipUnless(hasattr(_pybind, "pack_weights_for_cuda_mixed_gemm"), "fpA_intB weight packer is unavailable") +@unittest.skipUnless(_cuda_quant is not None, "fpA_intB weight packer is unavailable") class TestMatMulNBitsPrepackedCuda(unittest.TestCase): def _quantize_weight(self, weight: np.ndarray, bits: int, block_size: int): k, n = weight.shape @@ -118,7 +124,7 @@ def _check_prepacked_parity( bias = rng.normal(0.0, 1.0, size=(n,)).astype(np.float16) if has_bias else None q_weight, scales = self._quantize_weight(weight, bits, block_size) - prepacked_flat = _pybind.pack_weights_for_cuda_mixed_gemm(q_weight.reshape(n, -1), n, k, bits, force_arch) + prepacked_flat = _cuda_quant.pack_weights_for_cuda_mixed_gemm(q_weight.reshape(n, -1), n, k, bits, force_arch) prepacked_weight = np.asarray(prepacked_flat, dtype=np.int8).view(np.uint8).reshape(q_weight.shape) raw_model = self._make_model((m, k), q_weight, scales, bits, block_size, weight_prepacked=0, bias=bias) @@ -171,5 +177,39 @@ def test_int8_sm90_prepacked_weight_matches_runtime_prepack(self): self._check_sm90_parity(bits=8, block_size=128, m=32) +@unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") +@unittest.skipUnless(_cuda_quant is not None, "standalone CUDA weight packer (parity oracle) is unavailable") +class TestCudaQuantizerTorchPackerParity(unittest.TestCase): + """Validate the PyTorch mixed-GEMM packer in cuda_quantizer.py against the CUDA oracle. + + ``cuda_quantizer._pack_weights_for_cuda_mixed_gemm`` (PyTorch, used in production, and the + only option on Windows where the standalone module is not built) must be byte-identical to + the standalone ``onnxruntime_cuda_quant_preprocess.pack_weights_for_cuda_mixed_gemm`` (the + CUDA code the runtime prepack uses). This test is the guard against silent drift; it only + runs where the oracle is built (non-Windows CUDA). + """ + + def _check(self, bits: int, force_arch: int, n: int, k: int): + pack = 8 // bits + rng = np.random.default_rng(20260708 + bits * 100 + force_arch + n + k) + q = rng.integers(0, 256, size=(n, k // pack), dtype=np.uint8) + oracle = np.asarray(_cuda_quant.pack_weights_for_cuda_mixed_gemm(q, n, k, bits, force_arch), dtype=np.int8) + torch_out = _pack_weights_for_cuda_mixed_gemm(q, n, k, bits, force_arch).astype(np.int8) + self.assertEqual(oracle.shape, torch_out.shape, f"shape mismatch bits={bits} arch={force_arch} N={n} K={k}") + np.testing.assert_array_equal( + torch_out, oracle, err_msg=f"byte mismatch bits={bits} arch={force_arch} N={n} K={k}" + ) + + def test_torch_packer_matches_cuda_oracle(self): + # Cover both weight bit-widths, both mixed-GEMM layouts (SM80/SM90), and a GPT-OSS-20B + # MoE shape (fused gate+up FC1 [5760, 2880] and down FC2 [2880, 2880]). + shapes = [(256, 256), (512, 256), (256, 512), (5760, 2880), (2880, 2880), (128, 128)] + for bits in (4, 8): + for force_arch in (80, 90): + for n, k in shapes: + with self.subTest(bits=bits, force_arch=force_arch, n=n, k=k): + self._check(bits, force_arch, n, k) + + if __name__ == "__main__": unittest.main() diff --git a/setup.py b/setup.py index 62ced38819f2c..58aadc75b5010 100644 --- a/setup.py +++ b/setup.py @@ -375,6 +375,7 @@ def finalize_options(self): if platform.system() == "Linux" or platform.system() == "AIX": libs = [ "onnxruntime_pybind11_state.so", + "onnxruntime_cuda_quant_preprocess.so", "libdnnl.so.2", "libmklml_intel.so", "libmklml_gnu.so", @@ -388,6 +389,13 @@ def finalize_options(self): dl_libs.append(providers_cann) dl_libs.append(providers_qnn) dl_libs.append("libonnxruntime.so*") + # onnxruntime_cuda_quant_preprocess.so is a standalone CUDA extension module used only as a + # byte-parity oracle for the PyTorch weight packer. It is built (and thus present here) only + # when the CMake option onnxruntime_BUILD_CUDA_QUANT_PREPROCESS is ON. The glob-based filters below + # drop missing files, so listing it here is a no-op when it was not built. It must be listed in + # dl_libs (not just libs) so that manylinux test wheels include it: the manylinux packaging path + # builds "data" from dl_libs only (see the is_manylinux block below). + dl_libs.append("onnxruntime_cuda_quant_preprocess.so") # DNNL, TensorRT, OpenVINO, and QNN EPs are built as shared libs libs.extend(["libonnxruntime_providers_shared.so"]) libs.extend(["libonnxruntime_providers_dnnl.so"]) @@ -422,6 +430,7 @@ def finalize_options(self): elif platform.system() == "Darwin": libs = [ "onnxruntime_pybind11_state.so", + "onnxruntime_cuda_quant_preprocess.so", "libdnnl.2.dylib", "mimalloc.so", "libonnxruntime*.dylib", diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 8317018d33c64..4061e4fafbe7d 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -1822,7 +1822,9 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs): if not args.disable_contrib_ops: run_subprocess( - [sys.executable, "-m", "unittest", "discover", "-s", "quantization"], cwd=cwd, dll_path=dll_path + [sys.executable, "-m", "unittest", "discover", "-s", "quantization", "-v"], + cwd=cwd, + dll_path=dll_path, ) if args.enable_transformers_tool_test and (sys.version_info.major, sys.version_info.minor) < ( From cccd3ef1296dbd025f3f7713ae990c15c69a1de6 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Jul 2026 11:07:13 -0700 Subject: [PATCH 02/11] [CUDA] Fix XQA GroupQueryAttention cudaErrorInvalidValue on Blackwell (sm_120) (#29706) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary GroupQueryAttention's XQA decode kernel failed on consumer Blackwell GPUs (RTX 50-series, sm_120) with `cudaErrorInvalidValue`, while working fine on A100 (sm_80) and H200 (sm_90). This adds a runtime shared-memory capability check so XQA is only selected when the device can actually satisfy the kernel's dynamic shared-memory request, and otherwise falls back to cuDNN SDPA / Flash. A100/H200 behavior is unchanged. ### Root cause XQA bakes its shared-memory layout at compile time from `__CUDA_ARCH__`: - sm_80 / sm_87 / sm_90 use the large K/V-tile layout (`preferedKHeadPartBytes=128`, `cacheVTileSeqLen=64`) → up to ~140 KB of dynamic shared memory for head_size 128/256. - sm_86 / sm_89 / sm_120 use the small layout (64 / 32) → ~78–96 KB. Release/packaging binaries are built with a maximum arch of `90-virtual` (compute_90 PTX only, no native sm_120 SASS). On sm_120 the driver JIT-compiles that sm_90 PTX, so the kernel's `smemSize` carries the Hopper value (~140 KB). `launchMHA` then calls `cudaFuncSetAttribute(..., cudaFuncAttributeMaxDynamicSharedMemorySize, size)`, which exceeds sm_120's ~99 KB per-block opt-in limit (`sharedMemPerBlockOptin`) and returns `cudaErrorInvalidValue`. A100 (163 KB) and H200 (227 KB) have enough room, so they were unaffected. ### Key changes | File | Change | |---|---| | `xqa/xqa_impl_gen.cuh` | Add `GetSmemSize()` host helper that reads the per-kernel `smemSize` device symbol (accurate even for a PTX kernel JIT-compiled for the running SM). | | `xqa/xqa_loader_fp16_impl.cuh` | Add `GetXQAKernelSmemBytes(group_size)` head-dim dispatcher. The non-quantized fp16 footprint is an upper bound for the int8/fp8/bf16 variants (smaller cache element), so one query covers all XQA paths. | | `xqa/xqa_loader_fp16.cu`, `xqa/xqa_loader.h` | Expose `GetXQARequiredSharedMemoryBytes(device_prop, head_size, num_heads, kv_num_heads)`; a single non-templated entry point used by both the fp16 and bf16 GQA kernels. | | `group_query_attention.cc`, `group_query_attention.h` | Gate XQA selection on `required_smem <= device_prop.sharedMemPerBlockOptin`; fall back to cuDNN SDPA / Flash when it does not fit. Result is cached per node (`xqa_shared_memory_ok_`) since head_size/group are constant. | | `xqa/mha_impl.cuh` | Defensive backstop in `launchMHA`: if the requested shared memory still exceeds the device limit, throw an actionable message (which SM to build for / how to disable XQA) instead of the opaque `cudaErrorInvalidValue`. | ### CUDA graph safety `GetXQARequiredSharedMemoryBytes` uses `cudaMemcpyFromSymbol`, which synchronizes and is illegal during CUDA graph capture. The query is: - **cached** per node, so it runs at most once; - **guarded** with `onnxruntime::llm::common::isCapturing(Stream(context))` so the synchronizing call is only issued when the compute stream is not capturing; - resolved during ORT's non-captured warm-up run(s) before capture begins. If the value is somehow still unresolved while capturing, XQA is conservatively skipped for that run (safe fallback) without caching, so a later non-capturing run can resolve it. Warm-up and capture therefore make the same XQA/fallback decision, keeping the captured graph consistent with replay. ### Testing notes - Built the affected TUs (GQA dispatcher + fp16/bf16/int8/fp8 XQA loaders) with `CMAKE_CUDA_ARCHITECTURES="80;90"` (the configuration that reproduces the failure); all compile cleanly. - To validate the fix end-to-end, run a fp16/bf16 GQA decode workload on an sm_120 GPU (e.g. RTX 5090): it should now run (via fallback) instead of returning `cudaErrorInvalidValue`. Set `ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1` to confirm the selected backend. - To actually run XQA (the fast path) on Blackwell, build with native arch `120` in `CMAKE_CUDA_ARCHITECTURES` (and `100` for datacenter Blackwell). With a native sm_120 cubin the layout is ~80 KB and fits, so XQA is selected. --- .../cuda/bert/group_query_attention.cc | 32 +++++++++++++++++++ .../cuda/bert/group_query_attention.h | 5 +++ .../contrib_ops/cuda/bert/xqa/mha_impl.cuh | 15 +++++++++ .../cuda/bert/xqa/xqa_impl_gen.cuh | 25 +++++++++++++++ .../contrib_ops/cuda/bert/xqa/xqa_loader.h | 12 +++++++ .../cuda/bert/xqa/xqa_loader_fp16.cu | 28 ++++++++++++++++ .../cuda/bert/xqa/xqa_loader_fp16_impl.cuh | 31 ++++++++++++++++++ .../cuda/llm/common/cuda_runtime_utils.h | 4 +-- 8 files changed, 150 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 7cfb7c73fd016..287eaa0e97da2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -20,6 +20,7 @@ #include "contrib_ops/cuda/bert/unfused_attention.h" #include "contrib_ops/cuda/utils/dump_cuda_tensor.h" #include "contrib_ops/cpu/utils/debug_macros.h" +#include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; @@ -469,6 +470,37 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons data.use_xqa = (is_non_quantized_supported || is_int8_quantized_supported || is_fp8_quantized_supported); + if (data.use_xqa) { + // Consumer Blackwell (sm_120) and some other devices expose a smaller per-block opt-in + // shared-memory limit than the SM the XQA kernel was compiled / JIT-compiled for (sm_80 and + // sm_90 use a larger K/V-tile layout, ~140 KB for head_size 128/256). Launching XQA there + // fails with cudaErrorInvalidValue because cudaFuncSetAttribute requests more shared memory + // than the device allows. Query the actual kernel requirement (read from the device symbol, + // so it is accurate even for a PTX kernel JIT-compiled for this SM) and fall back to cuDNN + // SDPA / Flash when it does not fit. Cached per node because head_size and group size are + // constant for a given GQA node (avoids a device->host copy on every decode step). + int xqa_smem_ok = xqa_shared_memory_ok_.load(std::memory_order_relaxed); + if (xqa_smem_ok < 0) { + // GetXQARequiredSharedMemoryBytes issues a cudaMemcpyFromSymbol, which synchronizes and is + // therefore illegal during CUDA graph capture (it would invalidate the capture). The result + // is invariant for a node (fixed head_size / group size / device), and ORT performs at least + // one non-captured warm-up run before capturing, so this normally resolves and caches during + // warm-up. Guard defensively: only issue the synchronizing query when the compute stream is + // not capturing. If we somehow reach here for the first time while capturing, conservatively + // skip XQA for this run instead of breaking the capture, and leave the cache unresolved so a + // later non-capturing run can determine it. + if (!onnxruntime::llm::common::isCapturing(Stream(context))) { + const size_t required_smem = onnxruntime::contrib::cuda::GetXQARequiredSharedMemoryBytes( + device_prop, parameters.head_size, parameters.num_heads, parameters.kv_num_heads); + xqa_smem_ok = (required_smem == 0 || required_smem <= device_prop.sharedMemPerBlockOptin) ? 1 : 0; + xqa_shared_memory_ok_.store(xqa_smem_ok, std::memory_order_relaxed); + } else { + xqa_smem_ok = 0; // capturing (or status query failed) and not yet resolved -> fall back + } + } + data.use_xqa = (xqa_smem_ok != 0); + } + if (data.use_xqa) { size_t xqa_internal_bytes = onnxruntime::contrib::cuda::GetXQAScratchSize( GetDeviceProp(), diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h index d219bd83474d2..dc0d5977c0129 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include "core/providers/cuda/cuda_kernel.h" #include "contrib_ops/cuda/bert/group_query_attention_impl.h" @@ -52,6 +53,10 @@ class GroupQueryAttention final : public CudaKernel { // FP32 head_sink cached in PrePack for the XQA path (empty when head_sink is not a constant initializer). IAllocatorUniquePtr xqa_head_sink_; int xqa_head_sink_count_ = 0; // Number of elements in xqa_head_sink_ (0 when not prepacked). + // Cached result of the XQA shared-memory fit check for this node (-1 unknown, 0 does not fit, + // 1 fits). head_size and group size are constant for a given GQA node, so the (device-symbol) + // query is done once and reused to avoid a per-decode-step device->host copy. + mutable std::atomic xqa_shared_memory_ok_{-1}; const AttentionKernelOptions* kernel_options_; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh index 810c33adb5364..90f7aea9aedf8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh @@ -30,6 +30,7 @@ #ifndef GENERATE_CUBIN #include "hostUtils.h" #include +#include #ifndef NDEBUG #include #endif @@ -2491,6 +2492,20 @@ void launchMHA(cudaDeviceProp const& prop, uint32_t nbKHeads, static uint32_t const hostSmemSize = [&]() { uint32_t size; checkCuda(cudaMemcpyFromSymbol(&size, smemSize, sizeof(smemSize))); + // Defensive backstop: the kernel's shared-memory footprint is fixed at compile time for its + // target SM (sm_80/sm_90 use a larger K/V-tile layout than sm_86/sm_89/sm_120). When such a + // kernel is JIT-compiled from PTX onto a device with a smaller per-block opt-in limit (e.g. + // consumer Blackwell sm_120, ~99 KB), cudaFuncSetAttribute below returns cudaErrorInvalidValue. + // The GQA dispatcher already skips XQA in that case (see GetXQARequiredSharedMemoryBytes); this + // guard turns any remaining mismatch into an actionable message instead of an opaque CUDA error. + if (size > prop.sharedMemPerBlockOptin) { + throw std::runtime_error( + "XQA kernel requires " + std::to_string(size) + + " bytes of shared memory per block but this GPU allows only " + + std::to_string(prop.sharedMemPerBlockOptin) + + " bytes. Build ONNX Runtime with the device's native CUDA architecture (e.g. add 120 to " + "CMAKE_CUDA_ARCHITECTURES for sm_120 / RTX 50-series) or disable XQA (ORT_ENABLE_XQA=0)."); + } checkCuda(cudaFuncSetAttribute(kernel_mha, cudaFuncAttributeMaxDynamicSharedMemorySize, size)); return size; }(); diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh index a759c3d7f9760..4db0ed8a4c242 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh @@ -134,5 +134,30 @@ inline Status Launch( return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA is only supported on Ampere (SM80) or newer GPUs."); #endif } + +#ifndef GENERATE_CUBIN +// Host helper: dynamic shared-memory size (bytes) this kernel instantiation requests at launch, +// read from the device `smemSize` symbol. Because it is read from the loaded module, the value +// reflects the ACTUAL kernel that will run -- including a kernel JIT-compiled from PTX for a +// different SM, whose shared-memory layout (and therefore smemSize) was fixed at PTX-generation +// time. The GQA dispatcher uses this to avoid selecting XQA on devices whose per-block opt-in +// shared memory is smaller than the kernel needs, which would otherwise fail at launch with +// cudaErrorInvalidValue (e.g. consumer Blackwell sm_120 running a kernel JIT'd from sm_90 PTX, +// where the Hopper layout needs ~140 KB but sm_120 allows only ~99 KB). Returns 0 if the value +// cannot be queried. +inline size_t GetSmemSize() { +#ifdef XQA_HAS_SM80_TARGET + uint32_t size = 0; + if (cudaMemcpyFromSymbol(&size, smemSize, sizeof(smemSize)) != cudaSuccess) { + (void)cudaGetLastError(); // clear any sticky error so it does not leak to the next CUDA call + return 0; + } + return static_cast(size); +#else + return 0; +#endif +} +#endif // GENERATE_CUBIN + #undef XQA_HAS_SM80_TARGET } // namespace NAMESPACE_NAME diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h index 562bbab9a0693..a66abc958fe47 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h @@ -54,6 +54,18 @@ size_t GetXQAScratchSize( XqaQuantType kv_quant_type, bool is_bf16 = false); +// Dynamic shared memory (bytes) the XQA decode kernel requires for the given head size and +// query/KV group size on this device. The value is read from the loaded kernel module, so it is +// accurate even when the kernel is a PTX image JIT-compiled for the running SM. Returns 0 when it +// cannot be determined (e.g. unsupported head size / group size). Callers should skip XQA when the +// returned value exceeds device_prop.sharedMemPerBlockOptin (otherwise the launch fails with +// cudaErrorInvalidValue, e.g. consumer Blackwell sm_120 running a kernel built only for sm_90). +size_t GetXQARequiredSharedMemoryBytes( + const cudaDeviceProp& device_prop, + int head_size, + int num_heads, + int kv_num_heads); + } // namespace cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu index b744c23a906a8..61d059b854736 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu @@ -35,6 +35,8 @@ Status LaunchXQAKernelImpl( void* workspace, size_t workspace_size); +size_t GetXQAKernelSmemBytes(int group_size); + } // namespace H128 namespace H64 { @@ -61,6 +63,8 @@ Status LaunchXQAKernelImpl( void* workspace, size_t workspace_size); +size_t GetXQAKernelSmemBytes(int group_size); + } // namespace H64 namespace H256 { @@ -87,6 +91,8 @@ Status LaunchXQAKernelImpl( void* workspace, size_t workspace_size); +size_t GetXQAKernelSmemBytes(int group_size); + } // namespace H256 // Dispatcher Implementation @@ -178,6 +184,28 @@ size_t GetXQAScratchSize( return roundUp(semaphore_size, 128) + scratch_size; } +size_t GetXQARequiredSharedMemoryBytes( + const cudaDeviceProp& device_prop, + int head_size, + int num_heads, + int kv_num_heads) { + if (device_prop.major < 8 || kv_num_heads <= 0) { + return 0; + } + const int group_size = num_heads / kv_num_heads; + // The dtype (fp16 vs bf16) does not change the shared-memory footprint (both are 2-byte + // elements), so the fp16 kernels are queried for both. The non-quantized kernel is an upper + // bound for the int8/fp8 variants, so a single query covers all XQA paths. + if (head_size == 256) { + return H256::GetXQAKernelSmemBytes(group_size); + } else if (head_size == 128) { + return H128::GetXQAKernelSmemBytes(group_size); + } else if (head_size == 64) { + return H64::GetXQAKernelSmemBytes(group_size); + } + return 0; +} + // Instantiate template for half template Status LaunchXQAKernel( const cudaDeviceProp& device_prop, diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh index e861fb83ef098..09431f50f63cf 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh @@ -226,6 +226,37 @@ Status LaunchXQAKernelImpl( } } +#ifndef GENERATE_CUBIN +// Returns the dynamic shared-memory (bytes) the non-quantized XQA kernel for this head dim and +// group size requests at launch (read from the device symbol, so it is accurate even for a PTX +// kernel JIT-compiled for the running SM). Used by the GQA dispatcher to skip XQA when the +// device's per-block opt-in shared memory is too small. The non-quantized kernel is an upper +// bound for the int8/fp8 variants (smaller cache element), so this single query also guards the +// quantized paths. Not marked inline: it is defined in exactly one TU per head-dim namespace +// (xqa_loader_fp16_{64,128,256}.cu) and called from the head-size dispatcher TU, so it needs an +// externally linkable definition. +size_t GetXQAKernelSmemBytes(int group_size) { + switch (group_size) { + case 1: + return grp1_fp16::GetSmemSize(); + case 2: + return grp2_fp16::GetSmemSize(); + case 4: + return grp4_fp16::GetSmemSize(); + case 5: + return grp5_fp16::GetSmemSize(); + case 8: + return grp8_fp16::GetSmemSize(); + case 16: + return grp16_fp16::GetSmemSize(); + case 32: + return grp32_fp16::GetSmemSize(); + default: + return 0; + } +} +#endif // GENERATE_CUBIN + // Instantiate template for half } // namespace HEAD_DIM_NAMESPACE diff --git a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h index 8b9e35adc106e..a3419b88b8d01 100644 --- a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h +++ b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h @@ -22,6 +22,7 @@ #include #endif #include "core/providers/cuda/shared_inc/cuda_call.h" +#include "core/platform/env_var_utils.h" namespace onnxruntime::llm::common { inline int getDevice() { @@ -61,8 +62,7 @@ inline std::optional isCudaLaunchBlocking() { thread_local bool firstCall = true; thread_local std::optional result = std::nullopt; if (!firstCall) { - char const* env = std::getenv("CUDA_LAUNCH_BLOCKING"); - if (env != nullptr && std::string(env) == "1") { + if (ParseEnvironmentVariableWithDefault("CUDA_LAUNCH_BLOCKING", 0) == 1) { result = true; } else { result = false; From 866cc784c51f8260628803282c931c4e33ad8ee7 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Jul 2026 11:20:27 -0700 Subject: [PATCH 03/11] [CUDA] Do not link nvrtc (#29705) To avoid hard dependency on nvrtc dll even when it is not used for some models. --- cmake/onnxruntime_providers_cuda.cmake | 2 +- cmake/onnxruntime_providers_cuda_plugin.cmake | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake index 3a480033e77af..f486b89ec0206 100644 --- a/cmake/onnxruntime_providers_cuda.cmake +++ b/cmake/onnxruntime_providers_cuda.cmake @@ -401,7 +401,7 @@ endif() target_compile_definitions(${target} PRIVATE NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING) target_include_directories(${target} PRIVATE ${CUDNN_INCLUDE_DIR}) - target_link_libraries(${target} PRIVATE CUDA::cublasLt CUDA::cublas cudnn_frontend CUDA::curand CUDA::cufft CUDA::cudart CUDA::nvrtc CUDA::cuda_driver + target_link_libraries(${target} PRIVATE CUDA::cublasLt CUDA::cublas cudnn_frontend CUDA::curand CUDA::cufft CUDA::cudart CUDA::cuda_driver ${ABSEIL_LIBS} ${ONNXRUNTIME_PROVIDERS_SHARED} Boost::mp11 safeint_interface) endif() diff --git a/cmake/onnxruntime_providers_cuda_plugin.cmake b/cmake/onnxruntime_providers_cuda_plugin.cmake index ef4b3c7f6959a..0f3b42a3f549f 100644 --- a/cmake/onnxruntime_providers_cuda_plugin.cmake +++ b/cmake/onnxruntime_providers_cuda_plugin.cmake @@ -386,7 +386,6 @@ target_link_libraries(onnxruntime_providers_cuda_plugin PRIVATE CUDA::cublas CUDA::cublasLt CUDA::cufft - CUDA::nvrtc CUDA::cuda_driver cudnn_frontend Boost::mp11 From facf1fb1a836d38d8c0519bbfe9d3d76b118531b Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Jul 2026 11:34:47 -0700 Subject: [PATCH 04/11] [CUDA] Update cuda arch list for packages of cuda 12.8 (#29711) Drop 52-real; 90-virtual Add 120-real; 120-virtual Ensure 86-real is included Q: Why not add 100-real to cuda 12.8 build? A: We assume that those machines will have cuda 13.x for best performance. Q: Why drops 52-real A: Many applications require float16 support, while 52-real cannot support it. --- .../azure-pipelines/c-api-noopenmp-packaging-pipelines.yml | 2 +- .../github/azure-pipelines/custom-nuget-packaging-pipeline.yml | 2 +- tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml | 2 +- .../github/azure-pipelines/py-cuda-packaging-pipeline.yml | 2 +- .../azure-pipelines/stages/plugin-cuda-packaging-stage.yml | 2 +- .../github/azure-pipelines/stages/plugin-linux-cuda-stage.yml | 2 +- .../github/azure-pipelines/stages/plugin-win-cuda-stage.yml | 2 +- tools/ci_build/github/linux/build_cuda_c_api_package.sh | 2 +- tools/ci_build/github/linux/build_linux_python_package.sh | 2 +- tools/ci_build/github/linux/build_nodejs_package.sh | 2 +- tools/ci_build/github/linux/build_tensorrt_c_api_package.sh | 2 +- 11 files changed, 11 insertions(+), 11 deletions(-) diff --git a/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml b/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml index 8ed58864abe8d..43ebfa94e26bc 100644 --- a/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml +++ b/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml @@ -95,7 +95,7 @@ variables: value: $(Agent.TempDirectory)\9.14.0.64_cuda13 - name: CudaArchs ${{ if eq(parameters.CudaVersion, '12.8') }}: - value: '75-real;86-real;89-real;90-virtual' + value: '61-real;75-real;86-real;89-real;120-real;120-virtual' ${{ if eq(parameters.CudaVersion, '13.0') }}: value: '75-real;80-real;86-real;89-real;90-real;100-real;120-real;120-virtual' diff --git a/tools/ci_build/github/azure-pipelines/custom-nuget-packaging-pipeline.yml b/tools/ci_build/github/azure-pipelines/custom-nuget-packaging-pipeline.yml index aa29e6c9279e7..24070a8f1da06 100644 --- a/tools/ci_build/github/azure-pipelines/custom-nuget-packaging-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/custom-nuget-packaging-pipeline.yml @@ -34,7 +34,7 @@ variables: value: $(Agent.TempDirectory)\v13.0 - name: CmakeCudaArchitectures ${{ if eq(parameters.CudaVersion, '12.8') }}: - value: 52-real;61-real;75-real;86-real;89-real;90-virtual + value: 61-real;75-real;86-real;89-real;120-real;120-virtual ${{ if eq(parameters.CudaVersion, '13.0') }}: value: 75-real;80-real;86-real;89-real;90-real;100-real;120-real;120-virtual diff --git a/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml b/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml index 5608b59f9f2c1..a6c840736acb8 100644 --- a/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml @@ -133,7 +133,7 @@ extends: ${{ if eq(parameters.cuda_version, '12.8') }}: python_package_name: 'onnxruntime-ep-cuda12' docker_base_image: 'onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda12_x64_almalinux8_gcc14:20251017.1' - cmake_cuda_archs: '52-real;61-real;75-real;86-real;89-real;90-virtual' + cmake_cuda_archs: '61-real;75-real;86-real;89-real;120-real;120-virtual' ${{ if eq(parameters.cuda_version, '13.0') }}: python_package_name: 'onnxruntime-ep-cuda13' docker_base_image: 'onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda13_x64_almalinux8_gcc14:20251107.1' diff --git a/tools/ci_build/github/azure-pipelines/py-cuda-packaging-pipeline.yml b/tools/ci_build/github/azure-pipelines/py-cuda-packaging-pipeline.yml index 51f644cb5656a..1ce32d276bf90 100644 --- a/tools/ci_build/github/azure-pipelines/py-cuda-packaging-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/py-cuda-packaging-pipeline.yml @@ -64,5 +64,5 @@ extends: cmake_build_type: ${{ parameters.cmake_build_type }} cuda_version: '12.8' cudnn_folder: '' - cmake_cuda_archs: '52-real;61-real;75-real;86-real;89-real;90-virtual' + cmake_cuda_archs: '61-real;75-real;86-real;89-real;120-real;120-virtual' docker_base_image: 'onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda12_x64_almalinux8_gcc14:20251017.1' diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml index 8995d7134013a..84260c2ee7c43 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml @@ -54,7 +54,7 @@ parameters: - name: cmake_cuda_archs type: string displayName: 'CMAKE_CUDA_ARCHITECTURES' - default: '52-real;61-real;75-real;86-real;89-real;90-virtual' + default: '61-real;75-real;86-real;89-real;120-real;120-virtual' - name: python_version type: string diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml index db8591bf40a9e..e665ac35dd8ed 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml @@ -43,7 +43,7 @@ parameters: - name: cmake_cuda_archs type: string - default: '52-real;61-real;75-real;86-real;89-real;90-virtual' + default: '61-real;75-real;86-real;89-real;120-real;120-virtual' - name: python_version type: string diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml index ce611dad92f92..ef292c2ae3c72 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml @@ -28,7 +28,7 @@ parameters: - name: cmake_cuda_archs type: string - default: '52-real;61-real;75-real;86-real;89-real;90-virtual' + default: '61-real;75-real;86-real;89-real;120-real;120-virtual' stages: - stage: Win_plugin_cuda_x64_Build diff --git a/tools/ci_build/github/linux/build_cuda_c_api_package.sh b/tools/ci_build/github/linux/build_cuda_c_api_package.sh index b06fbc70c8ca9..45f07b754bf29 100755 --- a/tools/ci_build/github/linux/build_cuda_c_api_package.sh +++ b/tools/ci_build/github/linux/build_cuda_c_api_package.sh @@ -2,7 +2,7 @@ set -e -x if [ "$CUDA_VERSION" == "12.8" ]; then - CUDA_ARCHS="60-real;70-real;75-real;80-real;90a-real;90-virtual" + CUDA_ARCHS="60-real;70-real;75-real;80-real;86-real;90-real;120-real;120-virtual" elif [ "$CUDA_VERSION" == "13.0" ]; then CUDA_ARCHS="75-real;80-real;86-real;89-real;90-real;100-real;120-real;120-virtual" else diff --git a/tools/ci_build/github/linux/build_linux_python_package.sh b/tools/ci_build/github/linux/build_linux_python_package.sh index 2db052a15f22e..803597bd68cb6 100755 --- a/tools/ci_build/github/linux/build_linux_python_package.sh +++ b/tools/ci_build/github/linux/build_linux_python_package.sh @@ -71,7 +71,7 @@ fi if [ "$BUILD_DEVICE" == "GPU" ]; then if [ "$CUDA_VERSION" == "12.8" ]; then - CUDA_ARCHS="60-real;70-real;75-real;80-real;86-real;90a-real;90-virtual" + CUDA_ARCHS="60-real;70-real;75-real;80-real;86-real;90-real;120-real;120-virtual" elif [ "$CUDA_VERSION" == "13.0" ]; then CUDA_ARCHS="75-real;80-real;86-real;89-real;90-real;100-real;120-real;120-virtual" else diff --git a/tools/ci_build/github/linux/build_nodejs_package.sh b/tools/ci_build/github/linux/build_nodejs_package.sh index eb9b8cda9f482..d085fc61808f0 100755 --- a/tools/ci_build/github/linux/build_nodejs_package.sh +++ b/tools/ci_build/github/linux/build_nodejs_package.sh @@ -2,7 +2,7 @@ set -e -x if [ "$CUDA_VERSION" == "12.8" ]; then - CUDA_ARCHS="60-real;70-real;75-real;80-real;90a-real;90-virtual" + CUDA_ARCHS="60-real;70-real;75-real;80-real;86-real;90-real;120-real;120-virtual" elif [ "$CUDA_VERSION" == "13.0" ]; then CUDA_ARCHS="75-real;80-real;86-real;89-real;90-real;100-real;120-real;120-virtual" else diff --git a/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh b/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh index 8f03577bb546e..da02024e9eb1a 100755 --- a/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh +++ b/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh @@ -2,7 +2,7 @@ set -e -x if [ "$CUDA_VERSION" == "12.8" ]; then - CUDA_ARCHS="60-real;70-real;75-real;80-real;90a-real;90-virtual" + CUDA_ARCHS="60-real;70-real;75-real;80-real;86-real;90-real;120-real;120-virtual" elif [ "$CUDA_VERSION" == "13.0" ]; then CUDA_ARCHS="75-real;80-real;86-real;89-real;90-real;100-real;120-real;120-virtual" else From 083799e2a6a375447abff46177975431ffb7a25f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Jul 2026 12:07:17 -0700 Subject: [PATCH 05/11] [CUDA] Add cuDNN-free ArgMax/ArgMin/ReduceSum and fix LogSoftmax on plugin EP (#29620) ### Summary Phase 2 of the CUDA plugin execution provider "no-cuDNN" work. It lets single last-axis `ArgMax`/`ArgMin` run through a small custom CUDA kernel instead of cuDNN, fixes `LogSoftmax` classification in the plugin adapter, and adds a non-throwing cuDNN handle accessor so reduction kernels can fall back gracefully when cuDNN is disabled. ### Key Changes | Area | Change | |---|---| | `reduction_functions.cu` / `.h` | New `arg_min_max_last_axis` kernel (instantiated for `half`, `float`, `double`) that computes ArgMax/ArgMin indices over the last dimension of a row-major matrix without cuDNN. | | `reduction_ops.cc` | In `ReduceComputeCore`, route a single last-axis ArgMax/ArgMin (`CUDNN_REDUCE_TENSOR_FLATTENED_INDICES`) to the custom kernel when shapes fit `int`; otherwise fall through to the existing cuDNN path. `ReduceKernel::ComputeImpl` now uses `TryGetCudnnHandle`. | | `cuda_kernel.h` (native) / `cuda_kernel_adapter.h` (plugin) | Add `TryGetCudnnHandle`, which returns the cuDNN handle when available and `nullptr` otherwise (instead of throwing at handle acquisition). | | `softmax.h` | Detect `LogSoftmax` from `node.OpType()` instead of `info.GetKernelDef().OpName()`, so the plugin EP adapter classifies it correctly. | | `test_cuda_plugin_ep.py` | Add `LogSoftmax` and `ArgMin` tests; drop the `@requires_cudnn` gate from `ArgMax`, `ReduceMean`, `ReduceSum`; reduce over the last axis to exercise the cuDNN-free paths. | | `docs/cuda_plugin_ep/QUICK_START.md` | Drop `ArgMax` and reductions from the list of ops that still require cuDNN. | ### Correctness Notes - `select_last_index == 1` is already rejected on the CUDA EP, so the kernel keeping the first matching index (strict `>` / `<`) is spec-correct for the supported case. - The custom path guards `n > 0`, returns early for `m == 0`, computes the row offset in `int64_t`, and only engages when `m` and `n` fit in `int` (`gsl::narrow_cast`); larger tensors fall back to cuDNN. ### Testing - `python -m pytest onnxruntime/test/python/transformers/test_cuda_plugin_ep.py -k "log_softmax or argmax or argmin or reduce_mean or reduce_sum"` - Plugin no-cuDNN validation: `bash .env/cuda_130_plugin_no_cudnn.sh --build --test_plugin` - `onnxruntime_provider_test --gtest_filter='*Reduce*:*ArgM*'` --- .github/workflows/linux_cuda_no_cudnn.yml | 145 +++++++++-- docs/cuda_plugin_ep/QUICK_START.md | 2 +- onnxruntime/core/providers/cuda/cuda_kernel.h | 4 + .../core/providers/cuda/math/softmax.h | 2 +- .../cuda/plugin/cuda_kernel_adapter.h | 26 +- .../cuda/reduction/reduction_functions.cu | 232 ++++++++++++++++++ .../cuda/reduction/reduction_functions.h | 13 + .../providers/cuda/reduction/reduction_ops.cc | 54 +++- .../cpu/reduction/reduction_ops_test.cc | 25 ++ .../test_cases/reduction_functions_test.cc | 93 +++++++ .../transformers/test_cuda_plugin_ep.py | 69 +++++- 11 files changed, 632 insertions(+), 33 deletions(-) diff --git a/.github/workflows/linux_cuda_no_cudnn.yml b/.github/workflows/linux_cuda_no_cudnn.yml index abe92589aa705..47f8b33f15bbb 100644 --- a/.github/workflows/linux_cuda_no_cudnn.yml +++ b/.github/workflows/linux_cuda_no_cudnn.yml @@ -93,15 +93,37 @@ jobs: ldd /build/Release/Release/libonnxruntime_providers_cuda.so | tee /tmp/ldd.txt ! grep -i cudnn /tmp/ldd.txt' - - name: Run no-cuDNN CUDA EP smoke test + - name: Run no-cuDNN CUDA EP op smoke test run: | - docker run --rm --gpus all \ + # The build image runs as a non-root user by default, but this step + # deletes the system cuDNN libraries and rebuilds the dynamic linker + # cache (both require root). Run this disposable test container as root. + docker run --rm --gpus all --user root \ -v "${{ runner.temp }}/Release:/build/Release" \ "${{ steps.build_docker_image_step.outputs.full-image-name }}" \ bash -lc 'set -e PATH=/opt/python/cp312-cp312/bin:$PATH LD_LIBRARY_PATH=/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-} export PATH LD_LIBRARY_PATH + + # The build image contains cuDNN. Remove its runtime libraries from this disposable + # test container so the smoke test proves that the provider works when cuDNN is + # physically unavailable, rather than merely relying on enable_cudnn=0. + for root in /usr /opt /lib /lib64; do + if [ -e "$root" ]; then + find "$root" -name "libcudnn*.so*" -print -delete 2>/dev/null || true + fi + done + ldconfig + if ldconfig -p | grep -qi cudnn; then + echo "cuDNN remains available in the dynamic linker cache" >&2 + exit 1 + fi + if find /usr /opt /lib /lib64 -name "libcudnn*.so*" -print -quit 2>/dev/null | grep -q .; then + echo "cuDNN runtime libraries remain in the no-cuDNN test container" >&2 + exit 1 + fi + WHEEL_PATH=$(find /build/Release/Release/dist -type f -name "onnxruntime_gpu-*.whl" | head -n 1) if [ -z "$WHEEL_PATH" ]; then echo "No built onnxruntime GPU wheel found under /build/Release/Release/dist" >&2 @@ -111,21 +133,114 @@ jobs: python -m pip install --no-cache-dir --force-reinstall --no-deps numpy onnx "$WHEEL_PATH" python - <<"PY" import numpy as np - import onnx import onnxruntime as ort from onnx import TensorProto, helper - x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]) - y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3]) - node = helper.make_node("Add", ["x", "x"], ["y"]) - graph = helper.make_graph([node], "cuda_no_cudnn_smoke", [x], [y]) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) - model.ir_version = 10 - providers = [("CUDAExecutionProvider", {"enable_cudnn": "0"})] - sess = ort.InferenceSession(model.SerializeToString(), providers=providers) - data = np.arange(6, dtype=np.float32).reshape(2, 3) - result = sess.run(None, {"x": data})[0] - np.testing.assert_allclose(result, data + data) - print("CUDA no-cuDNN smoke test passed") + + + def run_op(name, model, feeds, expected, rtol=1e-4, atol=1e-4): + model.ir_version = 10 + options = ort.SessionOptions() + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + sess = ort.InferenceSession(model.SerializeToString(), sess_options=options, providers=providers) + got = sess.run(None, feeds)[0] + np.testing.assert_allclose(got, expected, rtol=rtol, atol=atol) + print("[no-cuDNN] " + name + " passed") + + + def make_model(nodes, inputs, outputs, opset, initializers=None): + graph = helper.make_graph(nodes, "no_cudnn_smoke", inputs, outputs, initializer=initializers or []) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + + + def f32_in(shape): + return helper.make_tensor_value_info("x", TensorProto.FLOAT, shape) + + + data = np.random.rand(2, 3).astype(np.float32) + + # Elementwise baseline (does not depend on cuDNN). + run_op( + "Add", + make_model( + [helper.make_node("Add", ["x", "x"], ["y"])], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + 21, + ), + {"x": data}, + data + data, + ) + + # LogSoftmax over the last axis (softmax kernel, no cuDNN). + shifted = data - data.max(axis=1, keepdims=True) + log_softmax = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True)) + run_op( + "LogSoftmax", + make_model( + [helper.make_node("LogSoftmax", ["x"], ["y"], axis=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + 13, + ), + {"x": data}, + log_softmax, + ) + + # ReduceMean over the last axis (matrix-reduction path, no cuDNN). + run_op( + "ReduceMean", + make_model( + [helper.make_node("ReduceMean", ["x"], ["y"], axes=[1], keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 1])], + 13, + ), + {"x": data}, + data.mean(axis=1, keepdims=True), + ) + + # ReduceSum over a middle axis, which requires the general no-cuDNN path. + reduce_sum_data = np.random.rand(2, 3, 4).astype(np.float32) + run_op( + "ReduceSum", + make_model( + [helper.make_node("ReduceSum", ["x", "axes"], ["y"], keepdims=1)], + [f32_in([2, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 1, 4])], + 13, + initializers=[helper.make_tensor("axes", TensorProto.INT64, [1], [1])], + ), + {"x": reduce_sum_data}, + reduce_sum_data.sum(axis=1, keepdims=True), + ) + + # ArgMax over the last axis (custom cuDNN-free kernel). + run_op( + "ArgMax", + make_model( + [helper.make_node("ArgMax", ["x"], ["y"], axis=1, keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2, 1])], + 13, + ), + {"x": data}, + np.argmax(data, axis=1).reshape(2, 1).astype(np.int64), + ) + + # ArgMin over the last axis (custom cuDNN-free kernel). + run_op( + "ArgMin", + make_model( + [helper.make_node("ArgMin", ["x"], ["y"], axis=1, keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2, 1])], + 13, + ), + {"x": data}, + np.argmin(data, axis=1).reshape(2, 1).astype(np.int64), + ) + + print("CUDA no-cuDNN op smoke tests passed") PY' diff --git a/docs/cuda_plugin_ep/QUICK_START.md b/docs/cuda_plugin_ep/QUICK_START.md index 7b971d621083f..793e25685af8b 100644 --- a/docs/cuda_plugin_ep/QUICK_START.md +++ b/docs/cuda_plugin_ep/QUICK_START.md @@ -27,7 +27,7 @@ For local Linux CUDA 13 validation, use the no-cuDNN helper script. It keeps `CU bash .env/cuda_130_plugin_no_cudnn.sh --build --test_plugin ``` -The test mode sets `ORT_TEST_CUDA_PLUGIN_EP=1` and `ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1`, which passes `enable_cudnn=0` to plugin sessions and skips plugin tests for operators that still require cuDNN, such as Conv, ConvTranspose, BatchNormalization, InstanceNormalization, LRN, ArgMax, reductions, Einsum, and cuDNN-backed pooling paths. +The test mode sets `ORT_TEST_CUDA_PLUGIN_EP=1` and `ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1`, which passes `enable_cudnn=0` to plugin sessions and skips plugin tests for operators that still require cuDNN, such as Conv, ConvTranspose, BatchNormalization, InstanceNormalization, LRN, Einsum, and cuDNN-backed pooling paths. ## Minimum ONNX Runtime Version diff --git a/onnxruntime/core/providers/cuda/cuda_kernel.h b/onnxruntime/core/providers/cuda/cuda_kernel.h index 85d629764da48..6152b5f888eb2 100644 --- a/onnxruntime/core/providers/cuda/cuda_kernel.h +++ b/onnxruntime/core/providers/cuda/cuda_kernel.h @@ -133,6 +133,10 @@ class CudaKernel : public OpKernel { return RequireCudnnHandle(GetCudnnHandle(static_cast(ctx->GetComputeStream()))); } + inline cudnnHandle_t TryGetCudnnHandle(OpKernelContext* ctx) const { + return GetCudnnHandle(static_cast(ctx->GetComputeStream())); + } + static inline cudnnHandle_t GetCudnnHandle(onnxruntime::CudaStream* stream) { return stream ? stream->cudnn_handle_ : nullptr; } diff --git a/onnxruntime/core/providers/cuda/math/softmax.h b/onnxruntime/core/providers/cuda/math/softmax.h index c0c0818042c15..09b6eeec834b5 100644 --- a/onnxruntime/core/providers/cuda/math/softmax.h +++ b/onnxruntime/core/providers/cuda/math/softmax.h @@ -47,7 +47,7 @@ class Softmax final : public CudaKernel { } } - log_softmax_ = info.GetKernelDef().OpName() == "LogSoftmax"; + log_softmax_ = node.OpType() == "LogSoftmax"; } Status ComputeInternal(OpKernelContext* context) const override; diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index b1cc748be754a..7c6cdd0e4dc80 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -1064,8 +1064,30 @@ class CudaKernel : public OpKernel { std::string("cuDNN is unavailable or disabled for CUDA Plugin Execution Provider: ") + onnxruntime::cuda::CudnnLibrary::Get().Error())); } - if (handle != nullptr && stream != nullptr) { - CUDNN_CALL_THROW(cudnnSetStream(handle, stream)); + // Bind the shared handle to the current compute stream. cudaStream_t 0/nullptr is the default + // stream, which is still a valid stream to bind, so do this unconditionally to avoid leaving + // the handle bound to a stale stream from a previous call. + CUDNN_CALL_THROW(cudnnSetStream(handle, stream)); + return handle; + } + + cudnnHandle_t TryGetCudnnHandle(OpKernelContext* ctx) const { + auto stream = Stream(ctx); + auto handle = GetCudnnHandle(stream); + if (handle != nullptr) { + return handle; + } + + handle = DefaultCudnnHandle(); + if (handle != nullptr) { + // Bind the shared handle to the current compute stream. cudaStream_t 0/nullptr is the default + // stream, which is still a valid stream to bind, so do this unconditionally to avoid leaving + // the handle bound to a stale stream from a previous call. + // Keep this accessor non-throwing: if the stream cannot be bound, treat it as "no cuDNN handle" + // so callers can fall back to a cuDNN-free path instead of failing. + if (!CUDNN_CALL(cudnnSetStream(handle, stream)).IsOK()) { + return nullptr; + } } return handle; } diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index ed97507f87641..afc2d9b061a55 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -4,6 +4,8 @@ #include "core/providers/cuda/reduction/reduction_functions.h" #include +#include +#include #include #include @@ -513,6 +515,236 @@ INSTANTIATE_REDUCE_MATRIX_COLUMNS(double); INSTANTIATE_REDUCE_MATRIX_COLUMNS(BFloat16); #undef INSTANTIATE_REDUCE_MATRIX_COLUMNS +namespace detail { +constexpr int kMaxReduceRank = 16; + +struct ReduceSumNdMetadata { + int output_segment_count{}; + int reduction_segment_count{}; + int64_t output_segment_sizes[kMaxReduceRank]{}; + int64_t output_segment_strides[kMaxReduceRank]{}; + int64_t reduction_segment_sizes[kMaxReduceRank]{}; + int64_t reduction_segment_strides[kMaxReduceRank]{}; + int64_t output_count{}; + int64_t reduction_count{}; +}; + +template +struct SumState { + T sum{}; + + __device__ __forceinline__ void Add(T value) { sum += value; } + __device__ __forceinline__ T Result() const { return sum; } +}; + +template <> +struct SumState { + double sum{}; + double correction{}; + + // Neumaier compensation is used instead of Kahan so that independently accumulated + // thread partials can be merged without losing a small value between large values. + __device__ __forceinline__ void Add(double value) { + const double next = __dadd_rn(sum, value); + const double error = fabs(sum) >= fabs(value) + ? __dadd_rn(__dsub_rn(sum, next), value) + : __dadd_rn(__dsub_rn(value, next), sum); + correction = __dadd_rn(correction, error); + sum = next; + } + + __device__ __forceinline__ double Result() const { return __dadd_rn(sum, correction); } +}; + +template +struct MergeSumState { + __device__ __forceinline__ SumState operator()(SumState lhs, const SumState& rhs) const { + lhs.Add(rhs.sum); + if constexpr (std::is_same_v) { + lhs.Add(rhs.correction); + } + return lhs; + } +}; + +template +__device__ __forceinline__ T CastReduceSumResult(TAccum value) { + if constexpr (std::is_integral_v) { + const double value_as_double = static_cast(value); + const double max_value = static_cast(std::numeric_limits::max()); + const double min_value = static_cast(std::numeric_limits::min()); + if (value_as_double >= max_value) return std::numeric_limits::max(); + if (value_as_double <= min_value) return std::numeric_limits::min(); + } + return static_cast(value); +} + +template +__global__ void reduce_sum_nd_kernel(const T* input, T* output, ReduceSumNdMetadata metadata) { + using TAccum = std::conditional_t, double, AccumulationType_t>; + using BlockReduce = cub::BlockReduce, BlockSize>; + __shared__ typename BlockReduce::TempStorage reduce_storage; + __shared__ int64_t input_base; + + // One cooperative block reduces each output. Grid-striding keeps the launch bounded for + // large outputs while retaining enough blocks to saturate the device. + for (int64_t output_index = blockIdx.x; + output_index < metadata.output_count; + output_index += gridDim.x) { + if (threadIdx.x == 0) { + int64_t remaining = output_index; + int64_t base = 0; + for (int segment = metadata.output_segment_count - 1; segment >= 0; --segment) { + const int64_t coordinate = segment == 0 ? remaining : remaining % metadata.output_segment_sizes[segment]; + if (segment != 0) remaining /= metadata.output_segment_sizes[segment]; + base += coordinate * metadata.output_segment_strides[segment]; + } + input_base = base; + } + __syncthreads(); + + SumState thread_sum{}; + for (int64_t reduction_index = threadIdx.x; reduction_index < metadata.reduction_count; + reduction_index += BlockSize) { + int64_t remaining = reduction_index; + int64_t input_index = input_base; + // Adjacent reduced dimensions are collapsed on the host, so this loop performs + // divisions only at reduced/non-reduced boundaries rather than once per rank. + for (int segment = metadata.reduction_segment_count - 1; segment >= 0; --segment) { + const int64_t coordinate = segment == 0 ? remaining : remaining % metadata.reduction_segment_sizes[segment]; + if (segment != 0) remaining /= metadata.reduction_segment_sizes[segment]; + input_index += coordinate * metadata.reduction_segment_strides[segment]; + } + thread_sum.Add(static_cast(input[input_index])); + } + + const SumState block_sum = BlockReduce(reduce_storage).Reduce(thread_sum, MergeSumState{}); + if (threadIdx.x == 0) { + output[output_index] = CastReduceSumResult(block_sum.Result()); + } + __syncthreads(); // reduce_storage is reused by the next grid-stride iteration. + } +} +} // namespace detail + +template +Status reduce_sum_nd(cudaStream_t stream, const T* input, T* output, + gsl::span dims, gsl::span axes) { + ORT_RETURN_IF_NOT(dims.size() <= detail::kMaxReduceRank, + "The general CUDA ReduceSum kernel supports ranks up to ", detail::kMaxReduceRank, "."); + + detail::ReduceSumNdMetadata metadata; + const int rank = gsl::narrow_cast(dims.size()); + std::array strides{}; + SafeInt stride = 1; + for (int axis = rank - 1; axis >= 0; --axis) { + ORT_RETURN_IF_NOT(dims[axis] > 0, "ReduceSum dimensions must be positive."); + strides[axis] = static_cast(stride); + stride *= dims[axis]; + } + + std::array reduced{}; + if (axes.empty()) { + for (int axis = 0; axis < rank; ++axis) reduced[axis] = true; + } else { + for (int64_t axis : axes) { + if (axis < 0) axis += rank; + ORT_RETURN_IF_NOT(axis >= 0 && axis < rank, "ReduceSum axis is out of range."); + ORT_RETURN_IF_NOT(!reduced[axis], "ReduceSum axes must not contain duplicates."); + reduced[axis] = true; + } + } + + SafeInt output_count = 1; + SafeInt reduction_count = 1; + for (int axis = 0; axis < rank;) { + const bool is_reduced = reduced[axis]; + SafeInt segment_size = 1; + int last_axis = axis; + do { + segment_size *= dims[axis]; + last_axis = axis++; + } while (axis < rank && reduced[axis] == is_reduced); + + if (is_reduced) { + const int segment = metadata.reduction_segment_count++; + metadata.reduction_segment_sizes[segment] = static_cast(segment_size); + metadata.reduction_segment_strides[segment] = strides[last_axis]; + reduction_count *= static_cast(segment_size); + } else { + const int segment = metadata.output_segment_count++; + metadata.output_segment_sizes[segment] = static_cast(segment_size); + metadata.output_segment_strides[segment] = strides[last_axis]; + output_count *= static_cast(segment_size); + } + } + metadata.output_count = static_cast(output_count); + metadata.reduction_count = static_cast(reduction_count); + + constexpr int block_size = 256; + constexpr int max_blocks = 65535; + const int grid_size = static_cast(std::min(max_blocks, metadata.output_count)); + detail::reduce_sum_nd_kernel<<>>(input, output, metadata); + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE_REDUCE_SUM_ND(T) \ + template Status reduce_sum_nd(cudaStream_t stream, const T* input, T* output, \ + gsl::span dims, gsl::span axes) +INSTANTIATE_REDUCE_SUM_ND(half); +INSTANTIATE_REDUCE_SUM_ND(float); +INSTANTIATE_REDUCE_SUM_ND(double); +INSTANTIATE_REDUCE_SUM_ND(BFloat16); +INSTANTIATE_REDUCE_SUM_ND(int32_t); +INSTANTIATE_REDUCE_SUM_ND(int64_t); +#undef INSTANTIATE_REDUCE_SUM_ND + +namespace detail { +template +__global__ void arg_min_max_last_axis_kernel(const TIn* input, int64_t* output, int m, int n) { + const int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= m) return; + + const int64_t row_offset = static_cast(row) * n; + TIn best_value = input[row_offset]; + int64_t best_index = 0; + for (int i = 1; i < n; ++i) { + const TIn value = input[row_offset + i]; + if constexpr (IsArgMax) { + if (value > best_value) { + best_value = value; + best_index = i; + } + } else { + if (value < best_value) { + best_value = value; + best_index = i; + } + } + } + + output[row] = best_index; +} +} // namespace detail + +template +Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* output, int m, int n) { + // The kernel reads input[row_offset] unconditionally, so a non-empty reduction axis is required. + if (m == 0 || n <= 0) return Status::OK(); + constexpr int block_size = 256; + const int grid_size = (m + block_size - 1) / block_size; + detail::arg_min_max_last_axis_kernel<<>>(input, output, m, n); + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(T) \ + template Status arg_min_max_last_axis(cudaStream_t stream, const T* input, int64_t* output, int m, int n); \ + template Status arg_min_max_last_axis(cudaStream_t stream, const T* input, int64_t* output, int m, int n) +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(half); +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(float); +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(double); +#undef INSTANTIATE_ARG_MIN_MAX_LAST_AXIS + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index 30b0afe4fe799..2e73a73a353f7 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -103,6 +103,19 @@ Status reduce_matrix_rows(cudaStream_t stream, const TIn* input, TOut* output, i template Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output, int m, int n, void* buffer, size_t buffer_size); +/** + * Computes ReduceSum for an arbitrary set of axes. This is the general fallback + * for axis layouts that cannot be flattened into one of the optimized matrix + * reductions above. + */ +template +Status reduce_sum_nd(cudaStream_t stream, const T* input, T* output, + gsl::span dims, gsl::span axes); + +/** Computes ArgMax/ArgMin indices over the last dimension in a row-major matrix. */ +template +Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* output, int m, int n); + /** Apply unary elementwise division. */ template void UnaryDiv(cudaStream_t stream, const T* input, T* output, T denominator, size_t count); diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index a8019cda5c411..ed6195c0bbb8b 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -299,10 +299,6 @@ Status PrepareForReduce(const Tensor* X, const int64_t rank = gsl::narrow(input_shape.NumDimensions()); prepare_reduce_metadata.input_count = input_shape.Size(); - if (rank > 8) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "cuDNN only supports up to 8-D tensors in reduction"); - } - const auto input_dims = input_shape.GetDims(); std::vector reduced(rank, false); if (axes.size() > 0) { @@ -481,6 +477,41 @@ Status ReduceComputeCore(const AllocatorPtr& gpu_allocator, const CudaKernel* ke } } + if constexpr (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES) { + if (axes.size() == 1) { + const int64_t rank = input_shape.NumDimensions(); + const int64_t axis = HandleNegativeAxis(axes[0], rank); + if (axis == rank - 1) { + const int64_t m = input_shape.SizeToDimension(axis); + const int64_t n = input_shape[axis]; + if (n > 0 && m <= std::numeric_limits::max() && n <= std::numeric_limits::max()) { + if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MAX) { + return arg_min_max_last_axis(stream, reinterpret_cast(input.Data()), + output.MutableData(), gsl::narrow_cast(m), + gsl::narrow_cast(n)); + } + if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MIN) { + return arg_min_max_last_axis(stream, reinterpret_cast(input.Data()), + output.MutableData(), gsl::narrow_cast(m), + gsl::narrow_cast(n)); + } + } + } + } + } + + // Preserve the optimized matrix reductions above and the existing cuDNN path when available. + // Without cuDNN, use a general CUDA kernel for plain ReduceSum axis layouts that cannot be + // represented as a contiguous matrix reduction. + if constexpr (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_NO_INDICES) { + if ((cudnn_handle == nullptr || input_shape.NumDimensions() > 8) && + cudnn_reduce_op == CUDNN_REDUCE_TENSOR_ADD && + !calculate_log && !calculate_sqt && !log_sum_exp) { + return reduce_sum_nd(stream, reinterpret_cast(input.Data()), + reinterpret_cast(output.MutableData()), input_shape.GetDims(), axes); + } + } + // This reduction keep adding values to this buffer. If a non-zero value, say 1000, is here, the sum will start with 1000. // Therefore zeroing out the memory is required CUDA_RETURN_IF_ERROR(cudaMemsetAsync(output.MutableDataRaw(), 0, output.SizeInBytes(), stream)); @@ -785,7 +816,7 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe const bool fast_reduction = fast_reduction_ && !ctx->GetUseDeterministicCompute(); return ReduceComputeCore(AllocatorPtr{}, this, *X, prepare_reduce_metadata, *Y, cudnn_reduce_op, axes, calculate_log_, calculate_sqt_, log_sum_exp_, fast_reduction, - Stream(ctx), GetComputeStream(ctx), GetCudnnHandle(ctx)); + Stream(ctx), GetComputeStream(ctx), TryGetCudnnHandle(ctx)); } #define SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(T) \ @@ -797,9 +828,8 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe const Tensor* X = ctx->Input(0); \ TensorShapeVector axes; \ size_t num_inputs = ctx->InputCount(); \ - if (num_inputs == 2) { \ - const Tensor* axes_tensor = ctx->Input(1); \ - ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); \ + const Tensor* axes_tensor = num_inputs == 2 ? ctx->Input(1) : nullptr; \ + if (axes_tensor != nullptr) { \ ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor."); \ auto nDims = static_cast(axes_tensor->Shape()[0]); \ const auto* data = axes_tensor->Data(); \ @@ -858,6 +888,14 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe return Status::OK(); \ } \ \ + if constexpr (std::is_same_v || std::is_same_v) { \ + if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_ADD && \ + !calculate_log_ && !calculate_sqt_ && !log_sum_exp_) { \ + return reduce_sum_nd(Stream(ctx), reinterpret_cast(X->Data()), \ + reinterpret_cast(Y->MutableData()), X->Shape().GetDims(), axes); \ + } \ + } \ + \ CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes(), Stream(ctx))); \ \ size_t indices_bytes = 0; \ diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index 95f5ae30016e2..13d6016d09de9 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc @@ -2863,6 +2863,31 @@ TEST(ReductionOpTest, ReduceSum_int64) { test.Run(); } +#if defined(USE_CUDA) +TEST(ReductionOpTest, ReduceSum_int64_omitted_optional_axes) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3}, {1, 2, 3}); + test.AddOptionalInputEdge(); + test.AddOutput("reduced", {}, {6}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(ReductionOpTest, ReduceSum_int64_cancellation) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + const int64_t large = int64_t{1} << 53; + test.AddInput("data", {3}, {large, 1, -large}); + test.AddInput("axes", {1}, {0}); + test.AddOutput("reduced", {}, {1}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + TEST(ReductionOpTest, ReduceSum_default_axes_keepdims) { OpTester test("ReduceSum"); test.AddAttribute("keepdims", (int64_t)1); diff --git a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc index 593255b9e9c23..410527f0a6ea6 100644 --- a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc @@ -3,7 +3,9 @@ #include "gtest/gtest.h" +#include #include +#include #include "core/providers/cuda/shared_inc/cuda_utils.h" #include "core/common/optional.h" @@ -238,6 +240,97 @@ TEST(ReductionFunctionsTest, ReduceColumnsToColumnRepeated) { TestReduceColumnsToColumnRepeated(17, 8192, 100, 2e-4f); } +TEST(ReductionFunctionsTest, ReduceSumNdMiddleAndMultipleAxes) { + const std::vector dims{2, 3, 4, 2}; + const std::vector axes{1, 3}; + std::vector input(48); + std::iota(input.begin(), input.end(), 1.0f); + std::vector expected(8, 0.0f); + for (int64_t d0 = 0; d0 < dims[0]; ++d0) { + for (int64_t d1 = 0; d1 < dims[1]; ++d1) { + for (int64_t d2 = 0; d2 < dims[2]; ++d2) { + for (int64_t d3 = 0; d3 < dims[3]; ++d3) { + expected[d0 * dims[2] + d2] += input[((d0 * dims[1] + d1) * dims[2] + d2) * dims[3] + d3]; + } + } + } + } + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + CheckDeviceValues(expected.size(), d_output.get(), expected.data(), 1e-6f); +} + +TEST(ReductionFunctionsTest, ReduceSumNdIntegerSaturation) { + const std::vector dims{2, 3, 2}; + const std::vector axes{1}; + const int32_t big = 1'100'000'000; + const std::vector input(12, big); + const std::vector expected(4, std::numeric_limits::max()); + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(int32_t), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + std::vector actual(expected.size()); + cudaMemcpy(actual.data(), d_output.get(), actual.size() * sizeof(int32_t), cudaMemcpyDeviceToHost); + EXPECT_EQ(actual, expected); +} + +TEST(ReductionFunctionsTest, ReduceSumNdLargeReductionSmallOutput) { + const std::vector dims{2, 131072, 3}; + const std::vector axes{1}; + std::vector input(TensorShape(dims).Size(), 1.0f); + const std::vector expected(6, 131072.0f); + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + CheckDeviceValues(expected.size(), d_output.get(), expected.data(), 0.0f); +} + +TEST(ReductionFunctionsTest, ReduceSumNdInt64Cancellation) { + const std::vector dims{1, 3, 1}; + const std::vector axes{1}; + const int64_t large = int64_t{1} << 53; + const std::vector input{large, 1, -large}; + const std::vector expected{1}; + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(int64_t), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + std::vector actual(expected.size()); + cudaMemcpy(actual.data(), d_output.get(), actual.size() * sizeof(int64_t), cudaMemcpyDeviceToHost); + EXPECT_EQ(actual, expected); +} + +TEST(ReductionFunctionsTest, ReduceSumNdRank9) { + const std::vector dims{2, 2, 2, 2, 2, 2, 2, 2, 2}; + const std::vector axes{1, 3, 5, 7}; + std::vector input(TensorShape(dims).Size(), 1.0f); + const std::vector expected(32, 16.0f); + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + CheckDeviceValues(expected.size(), d_output.get(), expected.data(), 0.0f); +} + TEST(ReductionFunctionsTest, BufferOffsets) { const int m = 2048; const int n = 1024; diff --git a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index caff34b1c79d4..c95a03dc50d8e 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -1804,6 +1804,25 @@ def expected(f): result = _run_model_test(target_device, "Softmax", model, feed, expected) self.assertEqual(result, TEST_PASS, "Softmax test failed") + def test_op_log_softmax(self): + target_device = get_cuda_plugin_device() + model = _make_simple_model( + "LogSoftmax", + [("X", TensorProto.FLOAT, [2, 5])], + [("Y", TensorProto.FLOAT, [2, 5])], + attrs={"axis": 1}, + opset=13, + ) + feed = {"X": np.random.rand(2, 5).astype(np.float32)} + + def expected(f): + x = f["X"] + shifted = x - np.max(x, axis=1, keepdims=True) + return shifted - np.log(np.sum(np.exp(shifted), axis=1, keepdims=True)) + + result = _run_model_test(target_device, "LogSoftmax", model, feed, expected) + self.assertEqual(result, TEST_PASS, "LogSoftmax test failed") + def test_op_relu(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -1900,7 +1919,6 @@ def test_op_flatten(self): result = _run_model_test(target_device, "Flatten", model, feed, lambda f: f["X"].reshape(2, 12)) self.assertEqual(result, TEST_PASS, "Flatten test failed") - @requires_cudnn def test_op_argmax(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -1916,6 +1934,21 @@ def test_op_argmax(self): ) self.assertEqual(result, TEST_PASS, "ArgMax test failed") + def test_op_argmin(self): + target_device = get_cuda_plugin_device() + model = _make_simple_model( + "ArgMin", + [("X", TensorProto.FLOAT, [3, 5])], + [("Y", TensorProto.INT64, [3, 1])], + attrs={"axis": 1, "keepdims": 1}, + opset=13, + ) + feed = {"X": np.random.rand(3, 5).astype(np.float32)} + result = _run_model_test( + target_device, "ArgMin", model, feed, lambda f: np.argmin(f["X"], axis=1).reshape(3, 1) + ) + self.assertEqual(result, TEST_PASS, "ArgMin test failed") + def test_op_topk(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -2030,23 +2063,21 @@ def expected(f): result = _run_model_test(target_device, "ConvTranspose", model, feed, expected) self.assertEqual(result, TEST_PASS, "ConvTranspose test failed") - @requires_cudnn def test_op_reduce_mean(self): target_device = get_cuda_plugin_device() model = _make_simple_model( "ReduceMean", [("X", TensorProto.FLOAT, [3, 4, 5])], - [("Y", TensorProto.FLOAT, [3, 1, 5])], - attrs={"axes": [1], "keepdims": 1}, + [("Y", TensorProto.FLOAT, [3, 4, 1])], + attrs={"axes": [2], "keepdims": 1}, opset=13, ) feed = {"X": np.random.rand(3, 4, 5).astype(np.float32)} result = _run_model_test( - target_device, "ReduceMean", model, feed, lambda f: np.mean(f["X"], axis=1, keepdims=True) + target_device, "ReduceMean", model, feed, lambda f: np.mean(f["X"], axis=2, keepdims=True) ) self.assertEqual(result, TEST_PASS, "ReduceMean test failed") - @requires_cudnn def test_op_reduce_sum(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -2064,6 +2095,32 @@ def test_op_reduce_sum(self): ) self.assertEqual(result, TEST_PASS, "ReduceSum test failed") + def _run_reduce_sum_integer_last_axis(self, onnx_dtype, np_dtype): + # Mirrors the qwen attention_mask usage: a rank-2 integer ReduceSum over the last axis. + # Integer ReduceSum takes a specialized path that does not use the float matrix fast path, + # so without cuDNN it must fall back to the general native kernel. Covering it here keeps the + # no-cuDNN CI honest for the exact layout that broke real models. + target_device = get_cuda_plugin_device() + model = _make_simple_model( + "ReduceSum", + [("X", onnx_dtype, [2, 8]), ("axes", TensorProto.INT64, [1])], + [("Y", onnx_dtype, [2, 1])], + attrs={"keepdims": 1}, + opset=13, + ) + axes_init = helper.make_tensor("axes", TensorProto.INT64, [1], [1]) + model.graph.initializer.append(axes_init) + feed = {"X": np.arange(16, dtype=np_dtype).reshape(2, 8)} + return _run_model_test(target_device, "ReduceSum", model, feed, lambda f: np.sum(f["X"], axis=1, keepdims=True)) + + def test_op_reduce_sum_int64_last_axis(self): + result = self._run_reduce_sum_integer_last_axis(TensorProto.INT64, np.int64) + self.assertEqual(result, TEST_PASS, "ReduceSum int64 test failed") + + def test_op_reduce_sum_int32_last_axis(self): + result = self._run_reduce_sum_integer_last_axis(TensorProto.INT32, np.int32) + self.assertEqual(result, TEST_PASS, "ReduceSum int32 test failed") + def test_op_gather_nd(self): target_device = get_cuda_plugin_device() model = _make_simple_model( From 97816ad6cd372373a0b73b11c89b2e7a157b6974 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Jul 2026 14:47:59 -0700 Subject: [PATCH 06/11] Enable Spectre-mitigated MSVC libs for BinSkim builds (#29624) ### Description This updates the Windows BinSkim-compliant build flags so `/Qspectre` builds also link against the MSVC Spectre-mitigated CRT/STL static libraries. `/Qspectre` only affects ONNX Runtime's own object files; BinSkim BA2024 can still report violations when the default non-Spectre `libcmt.lib`, `libcpmt.lib`, or `libvcruntime.lib` are linked into the final binary. ### Motivation and Context Release validation reported BinSkim BA2024 (`EnableSpectreMitigations`) warnings for `onnxruntime.dll` even when ORT was built with `--use_binskim_compliant_compile_flags`. The warning identified MSVC runtime and STL static libraries as the non-mitigated modules. This change makes the build option select the Spectre-mitigated MSVC library directory when it is available from the Visual Studio toolset. ### Key Changes - Adds `get_msvc_spectre_lib_dir()` to locate `%VCToolsInstallDir%\lib\spectre\` for the target Windows architecture. - Appends a quoted `/LIBPATH:` linker flag whenever Windows BinSkim flags enable `/Qspectre` and AddressSanitizer is not enabled. - Emits a warning when the Spectre-mitigated MSVC libraries cannot be found, with guidance to install the Visual Studio "C++ Spectre-mitigated libs" component. - Preserves the existing ASAN behavior because ASAN libraries do not have Spectre-mitigated variants. ### Testing - `python3 -m ruff check tools/ci_build/build.py` - `python3 -m ruff format --check tools/ci_build/build.py` `lintrunner -a tools/ci_build/build.py` was also attempted. It found the repository config and applied no file changes, but the local environment could not execute the Ruff lintrunner adapters because `python` is not available on PATH; the direct `python3 -m ruff` checks above passed. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- tools/ci_build/build.py | 52 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 4061e4fafbe7d..3d501c16079a8 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -320,6 +320,42 @@ def generate_vcpkg_install_options(build_dir, args): return vcpkg_install_options +def get_msvc_spectre_lib_dir(args): + """Return the directory that holds the MSVC Spectre-mitigated CRT/STL static libraries for the + target architecture, or None if it cannot be located. + + The /Qspectre compile flag only mitigates ONNX Runtime's own object files. The prebuilt MSVC + CRT/STL static libraries (libcmt.lib, libcpmt.lib, libvcruntime.lib) that get linked into the + binaries also need to be the Spectre-mitigated variants, otherwise BinSkim BA2024 + (EnableSpectreMitigations) still fails. Those variants ship in the "C++ Spectre-mitigated libs" + Visual Studio component under %VCToolsInstallDir%\\lib\\spectre\\. + """ + vctools_dir = os.environ.get("VCToolsInstallDir") # noqa: SIM112 + if not vctools_dir: + return None + if args.arm: + arch = "arm" + elif args.arm64: + arch = "arm64" + elif args.arm64ec: + arch = "arm64ec" + elif args.x86: + arch = "x86" + else: + # Default to the target architecture selected by vcvarsall.bat (x86, x64, arm, arm64), + # falling back to x64 which is what the official Windows release packages use. + arch = os.environ.get("VSCMD_ARG_TGT_ARCH", "x64") + spectre_dir = Path(vctools_dir) / "lib" / "spectre" / arch + if spectre_dir.is_dir(): + return str(spectre_dir) + # Some toolsets do not ship a dedicated arm64ec folder; those reuse the arm64 Spectre libraries. + if args.arm64ec: + fallback = Path(vctools_dir) / "lib" / "spectre" / "arm64" + if fallback.is_dir(): + return str(fallback) + return None + + def generate_build_tree( cmake_path, source_dir, @@ -1137,6 +1173,22 @@ def generate_build_tree( # Address Sanitizer libs do not have a Qspectre version. So they two cannot be both enabled. if not args.enable_address_sanitizer: cflags += ["/Qspectre"] + # /Qspectre only mitigates ONNX Runtime's own object files. The prebuilt MSVC + # CRT/STL static libraries (libcmt.lib, libcpmt.lib, libvcruntime.lib) that are + # linked into the binaries also have to be the Spectre-mitigated variants, + # otherwise BinSkim BA2024 (EnableSpectreMitigations) still fails. Prepend the + # Spectre lib directory to the linker search path so those libraries are + # resolved ahead of the default (non-mitigated) CRT libraries. + spectre_lib_dir = get_msvc_spectre_lib_dir(args) + if spectre_lib_dir is not None: + ldflags = [f'/LIBPATH:"{spectre_lib_dir}"', *ldflags] + else: + log.warning( + "Could not locate the MSVC Spectre-mitigated CRT/STL libraries. The " + "resulting binaries may fail BinSkim BA2024 (EnableSpectreMitigations). " + "Install the 'C++ Spectre-mitigated libs' component from the Visual " + "Studio installer and build from a Developer Command Prompt." + ) if config == "Release": cflags += ["/O2", "/Ob2", "/DNDEBUG"] elif config == "RelWithDebInfo": From 3ca3a33df843418def66be175515d79175ad3769 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Jul 2026 17:14:58 -0700 Subject: [PATCH 07/11] [BUILD] CUDA_QUANT_PREPROCESS off by default and Adjust CI (#29687) This pull request updates the build and CI configuration for CUDA-related workflows and the main CMake options. The main changes are the addition of new CMake build flags to enable CUDA quantization preprocessing, improved formatting for build flags, and a change in the default for the CUDA quant preprocess build option. These updates improve clarity, make it easier to customize builds, and ensure that the CUDA quant preprocess module is only built when explicitly requested. **Build configuration changes:** * Added the `--cmake_extra_defines onnxruntime_BUILD_CUDA_QUANT_PREPROCESS=ON` flag to the CUDA build jobs in `.github/workflows/linux_cuda_ci.yml` and `.github/workflows/linux_cuda_plugin_ci.yml`, enabling the CUDA quantization preprocessing module during CI builds. [[1]](diffhunk://#diff-04806013a5e7991ba5606145885d9b0fcd99a7df1f3bb96a2d38fc724ccd9b2aL32-R46) [[2]](diffhunk://#diff-04806013a5e7991ba5606145885d9b0fcd99a7df1f3bb96a2d38fc724ccd9b2aL114-R138) [[3]](diffhunk://#diff-64cd92765a9461e73a80c6b0401fe1960170834725a7e8a2ea153aff6d8f8388R45) * Added the `--cmake_extra_defines onnxruntime_QUICK_BUILD=ON` and `--cmake_extra_defines onnxruntime_USE_FPA_INTB_GEMM=OFF` flags to the CUDA no-cudnn build job for faster builds and to disable a specific GEMM implementation. [[1]](diffhunk://#diff-4e310144ab53bd9b6e48c7ceba29ad2c310724645278a28b85f7ba3a453c4980L35-R47) [[2]](diffhunk://#diff-04806013a5e7991ba5606145885d9b0fcd99a7df1f3bb96a2d38fc724ccd9b2aL114-R138) **Formatting and maintainability:** * Reformatted long `extra_build_flags` strings in workflow YAML files to use multi-line lists for improved readability and easier maintenance. [[1]](diffhunk://#diff-04806013a5e7991ba5606145885d9b0fcd99a7df1f3bb96a2d38fc724ccd9b2aL32-R46) [[2]](diffhunk://#diff-04806013a5e7991ba5606145885d9b0fcd99a7df1f3bb96a2d38fc724ccd9b2aL114-R138) [[3]](diffhunk://#diff-4e310144ab53bd9b6e48c7ceba29ad2c310724645278a28b85f7ba3a453c4980L35-R47) **CMake option default change:** * Changed the default value of the `onnxruntime_BUILD_CUDA_QUANT_PREPROCESS` CMake option from `ON` to `OFF` in `cmake/CMakeLists.txt`, so the CUDA quantization preprocessing module is only built when explicitly enabled. --- .github/workflows/linux_cuda_ci.yml | 28 ++++++++++++++++++++-- .github/workflows/linux_cuda_no_cudnn.yml | 14 ++++++++++- .github/workflows/linux_cuda_plugin_ci.yml | 2 +- cmake/CMakeLists.txt | 4 ++-- 4 files changed, 42 insertions(+), 6 deletions(-) diff --git a/.github/workflows/linux_cuda_ci.yml b/.github/workflows/linux_cuda_ci.yml index 89dcf718e8bf5..a4c463a21dfb8 100644 --- a/.github/workflows/linux_cuda_ci.yml +++ b/.github/workflows/linux_cuda_ci.yml @@ -29,7 +29,21 @@ jobs: dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cuda docker_build_args: '--build-arg BASEIMAGE=onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda13_x64_almalinux8_gcc14:20251107.1' docker_image_repo: onnxruntimecuda13manylinuxbuild - extra_build_flags: '--use_binskim_compliant_compile_flags --build_wheel --parallel --nvcc_threads 4 --flash_nvcc_threads 4 --cuda_version=13.0 --cuda_home=/usr/local/cuda-13.0 --cudnn_home=/usr/local/cuda-13.0 --enable_cuda_profiling --build_java --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 onnxruntime_BUILD_UNIT_TESTS=ON onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON' + extra_build_flags: >- + --use_binskim_compliant_compile_flags + --build_wheel + --parallel + --build_java + --nvcc_threads 4 + --flash_nvcc_threads 4 + --cuda_version=13.0 + --cuda_home=/usr/local/cuda-13.0 + --cudnn_home=/usr/local/cuda-13.0 + --enable_cuda_profiling + --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 + --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=ON + --cmake_extra_defines onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON + --cmake_extra_defines onnxruntime_BUILD_CUDA_QUANT_PREPROCESS=ON python_path_prefix: 'PATH=/opt/python/cp310-cp310/bin:$PATH' run_tests: false # <<< Do not run tests in this job upload_build_output: true # <<< Upload the build/Release directory @@ -111,5 +125,15 @@ jobs: build_config: Release mode: 'test' # Set mode to test execution_providers: 'cuda' - extra_build_flags: '--use_binskim_compliant_compile_flags --cuda_version=13.0 --cuda_home=/usr/local/cuda-13.0 --cudnn_home=/usr/local/cuda-13.0 --enable_cuda_profiling --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 onnxruntime_BUILD_UNIT_TESTS=ON onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON' + extra_build_flags: >- + --use_binskim_compliant_compile_flags + --cuda_version=13.0 + --cuda_home=/usr/local/cuda-13.0 + --cudnn_home=/usr/local/cuda-13.0 + --enable_cuda_profiling + --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 + --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=ON + --cmake_extra_defines onnxruntime_QUICK_BUILD=ON + --cmake_extra_defines onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON + --cmake_extra_defines onnxruntime_BUILD_CUDA_QUANT_PREPROCESS=ON python_path_prefix: 'PATH=/opt/python/cp310-cp310/bin:$PATH' diff --git a/.github/workflows/linux_cuda_no_cudnn.yml b/.github/workflows/linux_cuda_no_cudnn.yml index 47f8b33f15bbb..3194c3d2581f9 100644 --- a/.github/workflows/linux_cuda_no_cudnn.yml +++ b/.github/workflows/linux_cuda_no_cudnn.yml @@ -32,7 +32,19 @@ jobs: dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cuda docker_build_args: '--build-arg BASEIMAGE=onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda13_x64_almalinux8_gcc14:20251107.1' docker_image_repo: onnxruntimecuda13manylinuxbuild - extra_build_flags: '--use_binskim_compliant_compile_flags --build_wheel --parallel --nvcc_threads 4 --flash_nvcc_threads 4 --cuda_version=13.0 --cuda_home=/usr/local/cuda-13.0 --cudnn_home=/usr/local/cuda-13.0 --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 onnxruntime_BUILD_UNIT_TESTS=ON' + extra_build_flags: >- + --use_binskim_compliant_compile_flags + --build_wheel + --parallel + --nvcc_threads 4 + --flash_nvcc_threads 4 + --cuda_version=13.0 + --cuda_home=/usr/local/cuda-13.0 + --cudnn_home=/usr/local/cuda-13.0 + --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 + --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=ON + --cmake_extra_defines onnxruntime_QUICK_BUILD=ON + --cmake_extra_defines onnxruntime_USE_FPA_INTB_GEMM=OFF python_path_prefix: 'PATH=/opt/python/cp312-cp312/bin:$PATH' run_tests: false upload_build_output: true diff --git a/.github/workflows/linux_cuda_plugin_ci.yml b/.github/workflows/linux_cuda_plugin_ci.yml index 4fafcd619c5f9..221dca79ac147 100644 --- a/.github/workflows/linux_cuda_plugin_ci.yml +++ b/.github/workflows/linux_cuda_plugin_ci.yml @@ -42,6 +42,7 @@ jobs: --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=ON --cmake_extra_defines onnxruntime_QUICK_BUILD=ON --cmake_extra_defines onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON + --cmake_extra_defines onnxruntime_BUILD_CUDA_QUANT_PREPROCESS=ON python_path_prefix: 'PATH=/opt/python/cp312-cp312/bin:$PATH' run_tests: false upload_build_output: true @@ -145,4 +146,3 @@ jobs: cd /onnxruntime_src/onnxruntime/test/python/transformers python test_cuda_plugin_ep.py " - diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index ad446214cfc8f..c8bf4a7bb87ff 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -77,8 +77,8 @@ cmake_dependent_option(onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS "Build with CUD cmake_dependent_option(onnxruntime_USE_CUDA_NHWC_OPS "Build CUDA with NHWC op support" ON "onnxruntime_USE_CUDA" OFF) cmake_dependent_option(onnxruntime_BUILD_CUDA_EP_AS_PLUGIN "Build CUDA EP as a separate plugin shared library instead of the legacy in-tree provider" OFF "onnxruntime_USE_CUDA" OFF) -option(onnxruntime_BUILD_CUDA_QUANT_PREPROCESS "Build CUDA weight-packing module onnxruntime_cuda_quant_preprocess.so" ON) -option(onnxruntime_CUDA_MINIMAL "Build CUDA without any operations apart from memcpy ops. Usefuel for a very minial TRT build" OFF) +option(onnxruntime_BUILD_CUDA_QUANT_PREPROCESS "Build CUDA weight-packing module onnxruntime_cuda_quant_preprocess.so" OFF) +option(onnxruntime_CUDA_MINIMAL "Build CUDA without any operations apart from memcpy ops. Useful for a very minimal TRT build" OFF) option(onnxruntime_ENABLE_CUDA_LINE_NUMBER_INFO "When building with CUDA support, generate device code line number information." OFF) option(onnxruntime_USE_OPENVINO "Build with OpenVINO support" OFF) option(onnxruntime_USE_COREML "Build with CoreML support" OFF) From 2d2aa85f01f8b27b780a0d6bae5ff3a21e016472 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Jul 2026 17:59:09 -0700 Subject: [PATCH 08/11] [CUDA] Fix null allocator passed to plugin EP kernel PrePack (#29658) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description When ONNX Runtime is built with the CUDA execution provider as a plugin (`onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON`), the EP-API op-kernel adapter (`ep::adapter::KernelImpl::PrePackWeightImpl`) received a valid `OrtAllocator*` from the framework but discarded it and forwarded a **null** `AllocatorPtr{}` into the wrapped kernel's `PrePack()`. Any CUDA kernel that pre-packs a constant weight (e.g. `MatMulNBits`, `Conv`, `GroupQueryAttention`, quantized MoE) then allocated scratch through that null allocator and crashed during session initialization: ``` IAllocator::ValidateAllocator(const T&) [with T = std::shared_ptr] allocator != nullptr was false ``` This surfaced end to end as an ONNX Runtime GenAI `og.Model(...)` failure on a gpt-oss-20b (`MatMulNBits`) model when the CUDA EP was loaded as a plugin: the trivial init session succeeded, but the first real model session crashed while pre-packing quantized weights. ### Key Changes | File | Change | |---|---| | `include/onnxruntime/ep/adapter/op_kernel.h` | `PrePackWeightImpl` now wraps the incoming `OrtAllocator*` and forwards a valid `AllocatorPtr` to `OpKernel::PrePack` instead of a null `AllocatorPtr{}`. | | `include/onnxruntime/ep/adapter/allocator.h` | Add a **non-owning** `IAllocatorWrappingOrtAllocator(OrtAllocator*)` constructor. The framework owns the pre-pack allocator, so the wrapper must not take ownership (an owning `Ort::Allocator` would release it on destruction). Calls now go through the raw `OrtAllocator*` function pointers directly, preserving the `Reserve`→`Alloc` (version ≥ 18) and `GetStats`/`AllocOnStream` (version ≥ 23) fallbacks. | | `onnxruntime/test/python/transformers/test_cuda_plugin_ep.py` | Add `test_registration_matmul_nbits_prepack`: builds a fp16 `MatMulNBits` model with a runtime-prepacked (`weight_prepacked=0`) quantized weight, so weight pre-packing (`MatMulNBits::PrePack_B` → `IAllocator::MakeUniquePtr(alloc, ...)`) runs during session creation. This crashed before the fix and now passes. | ### Motivation and Context The pre-pack allocator is provided and owned by the framework for the duration of the `PrePack` call. The legacy (in-tree) CUDA EP received it correctly; only the plugin op-kernel adapter dropped it. The non-owning wrapper matches the lifetime contract used elsewhere in the adapter (e.g. `KernelInfoGetAllocator`) and keeps the CUDA-EP-as-plugin build behaviorally identical to the in-tree EP. ### Testing - New `test_registration_matmul_nbits_prepack` in `test_cuda_plugin_ep.py` passes on a CUDA-EP-as-plugin build (`ORT_TEST_CUDA_PLUGIN_EP=1`) and skips gracefully when the device lacks fpA_intB GEMM support. It re-raises (fails) if the `allocator != nullptr` assertion recurs. - Verified the model that originally reproduced the crash now creates and runs end to end through the CUDA plugin EP (gpt-oss-20b, `MatMulNBits`), including a 100-sample MMLU sanity run. - Existing `test_cuda_plugin_ep.py` registration tests continue to pass. --- include/onnxruntime/ep/adapter/allocator.h | 47 ++++-- include/onnxruntime/ep/adapter/op_kernel.h | 6 +- .../transformers/test_cuda_plugin_ep.py | 139 ++++++++++++++++++ 3 files changed, 176 insertions(+), 16 deletions(-) diff --git a/include/onnxruntime/ep/adapter/allocator.h b/include/onnxruntime/ep/adapter/allocator.h index 1fb78f81fce19..01a36f2026963 100644 --- a/include/onnxruntime/ep/adapter/allocator.h +++ b/include/onnxruntime/ep/adapter/allocator.h @@ -25,35 +25,47 @@ class IAllocatorWrappingOrtAllocator final : public IAllocator { public: explicit IAllocatorWrappingOrtAllocator(Ort::Allocator ort_allocator) : IAllocator(*(EnsureOrtAllocatorHasValue(ort_allocator).GetInfo())), - ort_allocator_(std::move(ort_allocator)) { + owned_ort_allocator_(std::move(ort_allocator)), + ort_allocator_(owned_ort_allocator_) { + } + + // Wraps an OrtAllocator without taking ownership. The caller must keep it alive for the + // lifetime of this IAllocator. This is used for allocators passed to PrePackWeight, which are + // provided and owned by the ORT framework/caller for the duration of the PrePack call (this + // wrapper must not release them). + explicit IAllocatorWrappingOrtAllocator(OrtAllocator* ort_allocator) + : IAllocator(*EnsureOrtAllocatorHasValue(ort_allocator)->Info(ort_allocator)), + owned_ort_allocator_(nullptr), + ort_allocator_(ort_allocator) { } void* Alloc(size_t size) override { - return ort_allocator_.Alloc(size); + return ort_allocator_->Alloc(ort_allocator_, size); } void Free(void* p) override { - ort_allocator_.Free(p); + ort_allocator_->Free(ort_allocator_, p); } void* Reserve(size_t size) override { - return ort_allocator_.Reserve(size); + if (ort_allocator_->version >= 18 && ort_allocator_->Reserve != nullptr) { + return ort_allocator_->Reserve(ort_allocator_, size); + } + return ort_allocator_->Alloc(ort_allocator_, size); } bool IsStreamAware() const override { static constexpr uint32_t kOrtAllocatorAllocOnStreamMinVersion = 23; - const OrtAllocator* raw = ort_allocator_; - return raw->version >= kOrtAllocatorAllocOnStreamMinVersion && raw->AllocOnStream != nullptr; + return ort_allocator_->version >= kOrtAllocatorAllocOnStreamMinVersion && ort_allocator_->AllocOnStream != nullptr; } void* AllocOnStream(size_t size, Stream* stream) override { static constexpr uint32_t kOrtAllocatorAllocOnStreamMinVersion = 23; - OrtAllocator* raw = ort_allocator_; - if (raw->version >= kOrtAllocatorAllocOnStreamMinVersion && raw->AllocOnStream != nullptr) { - return raw->AllocOnStream(raw, size, reinterpret_cast(stream)); + if (ort_allocator_->version >= kOrtAllocatorAllocOnStreamMinVersion && ort_allocator_->AllocOnStream != nullptr) { + return ort_allocator_->AllocOnStream(ort_allocator_, size, reinterpret_cast(stream)); } - return raw->Alloc(raw, size); + return ort_allocator_->Alloc(ort_allocator_, size); } void GetStats(AllocatorStats* stats) override { @@ -62,10 +74,11 @@ class IAllocatorWrappingOrtAllocator final : public IAllocator { // GetStats was added in OrtAllocator version 23. For older allocators the function pointer // may be uninitialized, so we must not call through it. - const OrtAllocator* raw = ort_allocator_; - if (raw->version < 23 || !raw->GetStats) return; + if (ort_allocator_->version < 23 || !ort_allocator_->GetStats) return; - Ort::KeyValuePairs kvps = ort_allocator_.GetStats(); + OrtKeyValuePairs* stats_kvps = nullptr; + Ort::ThrowOnError(ort_allocator_->GetStats(ort_allocator_, &stats_kvps)); + Ort::KeyValuePairs kvps{stats_kvps}; std::vector keys, values; kvps.GetKeyValuePairs(keys, values); const size_t n = keys.size() < values.size() ? keys.size() : values.size(); @@ -82,7 +95,13 @@ class IAllocatorWrappingOrtAllocator final : public IAllocator { return ort_allocator; } - Ort::Allocator ort_allocator_; + static OrtAllocator* EnsureOrtAllocatorHasValue(OrtAllocator* ort_allocator) { + ORT_ENFORCE(ort_allocator != nullptr, "OrtAllocator must be non-null."); + return ort_allocator; + } + + Ort::Allocator owned_ort_allocator_; + OrtAllocator* ort_allocator_{}; ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(IAllocatorWrappingOrtAllocator); }; diff --git a/include/onnxruntime/ep/adapter/op_kernel.h b/include/onnxruntime/ep/adapter/op_kernel.h index 1f103b64a443e..baddee86ab3a2 100644 --- a/include/onnxruntime/ep/adapter/op_kernel.h +++ b/include/onnxruntime/ep/adapter/op_kernel.h @@ -9,6 +9,7 @@ #include #include +#include #include #include "core/framework/allocator.h" @@ -222,14 +223,15 @@ struct KernelImpl : OrtKernelImpl { static OrtStatus* ORT_API_CALL PrePackWeightImpl(_In_ OrtKernelImpl* this_ptr, _In_ const OrtValue* weight, int input_index, - _In_ OrtAllocator* /* allocator */, + _In_ OrtAllocator* allocator, _In_opt_ OrtSharedPrePackedWeightCache* /* prepacked_weight_cache */, _Out_ bool* is_packed) noexcept { Status status; ORT_TRY { auto* kernel_impl = static_cast(this_ptr)->impl_.get(); const auto tensor = CreateTensorFromApiValue(const_cast(weight)); - status = kernel_impl->PrePack(tensor, input_index, AllocatorPtr{}, *is_packed, nullptr); + auto allocator_wrapper = std::make_shared(allocator); + status = kernel_impl->PrePack(tensor, input_index, std::move(allocator_wrapper), *is_packed, nullptr); } ORT_CATCH(const std::exception& ex) { ORT_HANDLE_EXCEPTION([&]() { diff --git a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index c95a03dc50d8e..137aaaee2f2da 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -5,6 +5,7 @@ import os import tempfile import unittest +from contextlib import contextmanager import numpy as np import torch @@ -105,6 +106,42 @@ def is_cuda_mempool_unsupported_error(exc: Exception) -> bool: ) +# The fpA_intB MatMulNBits pre-pack path (which triggers the EP-provided allocator use this test +# guards) requires a compute capability >= 7.5 (Turing) device; see MatMulNBits::MatMulNBits. +FPA_INTB_MIN_COMPUTE_CAPABILITY = (7, 5) + + +def get_cuda_compute_capability(device): + """Return the device (major, minor) CUDA compute capability from plugin EP metadata, or None.""" + value = device.ep_metadata.get("cuda_compute_capability") + if not value: + return None + try: + major_str, _, minor_str = value.partition(".") + return (int(major_str), int(minor_str)) + except (TypeError, ValueError): + return None + + +def is_fpa_intb_unsupported_error(exc: Exception) -> bool: + """True if the error indicates the fpA_intB MatMulNBits path is unsupported for this device/config.""" + message = str(exc) + return "fpA_intB" in message + + +@contextmanager +def scoped_env(name: str, value: str): + old_value = os.environ.get(name) + os.environ[name] = value + try: + yield + finally: + if old_value is None: + os.environ.pop(name, None) + else: + os.environ[name] = old_value + + def _create_session_options(session_config=None): sess_options = onnxrt.SessionOptions() if session_config: @@ -220,6 +257,52 @@ def create_gemm_model(model_path, alpha=1.0, beta=1.0, transA=0, transB=0): save(model_def, model_path) +def create_matmul_nbits_model(model_path, k=64, n=64, bits=4, block_size=32): + # Create a MatMulNBits (com.microsoft) model with a runtime-prepacked (weight_prepacked=0) + # fp16 quantized weight. During session initialization the CUDA kernel pre-packs the constant + # weight via MatMulNBits::PrePack_B, which allocates scratch through the EP-provided allocator + # (IAllocator::MakeUniquePtr(alloc, ...)). This is a regression guard: the CUDA-EP-as-plugin + # adapter previously forwarded a null allocator into PrePack, so this model crashed at session + # creation with "IAllocator::ValidateAllocator ... allocator != nullptr was false". + k_blocks = (k + block_size - 1) // block_size + blob_size = block_size * bits // 8 + + # Deterministic quantized weight and positive scales. Exact numerics are not asserted here; + # the test only requires that pre-packing runs (and no longer crashes) end to end. + b = np.full((n, k_blocks, blob_size), 0x88, dtype=np.uint8) + scales = np.full((n, k_blocks), 0.02, dtype=np.float16) + + node_def = helper.make_node( + "MatMulNBits", + ["A", "B", "scales"], + ["Y"], + domain="com.microsoft", + K=k, + N=n, + bits=bits, + block_size=block_size, + weight_prepacked=0, + ) + graph_def = helper.make_graph( + [node_def], + "test-model-matmulnbits", + [helper.make_tensor_value_info("A", TensorProto.FLOAT16, [2, k])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT16, [2, n])], + initializer=[ + helper.make_tensor("B", TensorProto.UINT8, b.shape, b.tobytes(), raw=True), + helper.make_tensor("scales", TensorProto.FLOAT16, scales.shape, scales.tobytes(), raw=True), + ], + ) + opset_onnx = OperatorSetIdProto() + opset_onnx.version = 21 + opset_ms = OperatorSetIdProto() + opset_ms.domain = "com.microsoft" + opset_ms.version = 1 + model_def = helper.make_model(graph_def, producer_name="onnx-example", opset_imports=[opset_onnx, opset_ms]) + model_def.ir_version = 10 + save(model_def, model_path) + + def create_conv_model(model_path): # Create a simple Conv model: Y = Conv(X, W) node_def = helper.make_node("Conv", ["X", "W"], ["Y"], pads=[1, 1, 1, 1], strides=[1, 1], dilations=[1, 1], group=1) @@ -719,6 +802,62 @@ def test_registration_conv(self): result = run_operator_test(target_device, create_conv_model, inputs, _expected_conv) self.assertTrue(result, "Conv plugin registration test failed") + def test_registration_matmul_nbits_prepack(self): + # Regression guard for the CUDA-EP-as-plugin pre-pack allocator bug: a fp16 MatMulNBits + # weight is pre-packed during session initialization, which allocates scratch through the + # EP-provided allocator. The plugin op-kernel adapter previously forwarded a null allocator + # into PrePack, crashing session creation with an IAllocator::ValidateAllocator failure. + target_device = get_cuda_plugin_device() + + # The fpA_intB pre-pack path only runs on SM >= 7.5. Skip deterministically on older GPUs + # so the test does not silently pass without exercising the allocator-forwarding path. + compute_capability = get_cuda_compute_capability(target_device) + if compute_capability is None: + self.skipTest("CUDA plugin EP device metadata did not include cuda_compute_capability") + if compute_capability < FPA_INTB_MIN_COMPUTE_CAPABILITY: + self.skipTest( + "MatMulNBits fpA_intB pre-pack requires compute capability >= 7.5, but device reports " + f"{compute_capability[0]}.{compute_capability[1]}" + ) + + with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as tmp: + model_path = tmp.name + try: + create_matmul_nbits_model(model_path) + sess_options = _create_session_options() + sess_options.add_provider_for_devices([target_device], _plugin_provider_options()) + + # Force the fpA_intB path so weight pre-packing (MatMulNBits::PrePack_B) runs during + # session creation, which is where the null-allocator crash used to happen. Given the + # deterministic capability check above, session creation is expected to succeed, so any + # exception here (including the "allocator != nullptr" assertion) is a genuine regression + # and is deliberately propagated instead of skipped. + with scoped_env("ORT_FPA_INTB_GEMM", "1"): + sess = onnxrt.InferenceSession(model_path, sess_options=sess_options) + + assigned_nodes, assignment_info = _get_assigned_nodes(sess, CUDA_PLUGIN_EP_NAME) + self.assertTrue( + assigned_nodes, + f"{CUDA_PLUGIN_EP_NAME} was assigned no nodes. " + f"Assignments: {_format_assignment_summary(assignment_info)}", + ) + + a = np.random.rand(2, 64).astype(np.float16) + # Some devices/configs may still lack a valid fpA_intB tactic at runtime; skip only + # for that known-unsupported case and re-raise anything else so real regressions + # (e.g. incorrect kernels or the allocator assertion) are not hidden. + try: + res = sess.run(None, {"A": a}) + except Exception as exc: + if is_fpa_intb_unsupported_error(exc): + self.skipTest(f"fpA_intB MatMulNBits not supported on this device: {exc}") + raise + self.assertEqual(res[0].shape, (2, 64)) + self.assertTrue(np.isfinite(res[0].astype(np.float32)).all(), "MatMulNBits produced non-finite output") + finally: + if os.path.exists(model_path): + os.remove(model_path) + # ---- Provider options tests ---- def test_provider_options_valid(self): From 7271228be030bdec00b957867120d3dd7c986750 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Jul 2026 18:11:12 -0700 Subject: [PATCH 09/11] Add Windows ARM64 CUDA plugin package and align CUDA metadata/artifact naming (#28896) ### Description This PR adds Windows ARM64 support to CUDA plugin packaging and updates related build/packaging logic for correctness and consistency across architectures. In response to review feedback, it also: - Corrects CMake comments to match actual Windows ARM64 CUDA toolkit search behavior. - Prioritizes architecture-specific cuDNN DLL search paths before generic fallback paths. - Aligns Windows ARM64 Python artifact naming with the ARM64 CUDA toolkit version. - Extends packaging metadata with per-platform CUDA version fields (including win-arm64) and updates metadata-reading template validation/exports accordingly. ### Motivation and Context Windows ARM64 CUDA plugin packaging requires architecture-aware handling for toolkit/cuDNN discovery and artifact metadata. These updates prevent cross-architecture ambiguity (especially for CUDA 13.x win-arm64 vs x64), improve downstream artifact selection reliability, and keep metadata semantics consistent with produced artifacts. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- cmake/external/cuDNN.cmake | 14 +- .../external/onnxruntime_external_deps.cmake | 21 ++ cmake/onnxruntime_common.cmake | 9 +- cmake/onnxruntime_providers_cuda.cmake | 7 +- cmake/onnxruntime_providers_cuda_plugin.cmake | 9 +- cmake/onnxruntime_python.cmake | 12 +- .../core/providers/cuda/cu_inc/cub.cuh | 16 +- .../Microsoft.ML.OnnxRuntime.EP.Cuda.csproj | 8 + plugin-ep-cuda/csharp/README.md | 2 + plugin-ep-cuda/csharp/pack_nuget.py | 1 + .../azure-pipelines/plugin-cuda-pipeline.yml | 36 ++-- .../plugin-cuda-nuget-packaging-stage.yml | 11 ++ .../stages/plugin-cuda-packaging-stage.yml | 92 ++++++--- .../stages/plugin-win-cuda-stage.yml | 184 +++++++++++++----- .../read-cuda-plugin-metadata-steps.yml | 14 +- 15 files changed, 333 insertions(+), 103 deletions(-) diff --git a/cmake/external/cuDNN.cmake b/cmake/external/cuDNN.cmake index 2c755649d1360..c15e82ba2e5d6 100644 --- a/cmake/external/cuDNN.cmake +++ b/cmake/external/cuDNN.cmake @@ -12,10 +12,22 @@ string(REGEX MATCH "#define CUDNN_MAJOR [1-9]+" macrodef "${cudnn_version_header string(REGEX MATCH "[1-9]+" CUDNN_MAJOR_VERSION "${macrodef}") function(find_cudnn_library NAME) + if(WIN32) + # Since CUDA 13.x, Windows cuDNN import libraries live under an architecture-specific + # subdirectory (lib/x64 for win-x64, lib/arm64 for win-arm64). A single cuDNN package + # may ship both arches, so only search the subdirectory matching the current target + # platform — otherwise find_library could pick the wrong arch (it returns the first + # match in PATH_SUFFIXES order). + if(onnxruntime_target_platform STREQUAL "ARM64" OR onnxruntime_target_platform STREQUAL "ARM64EC") + set(_cudnn_arch_suffixes lib/arm64 lib/${onnxruntime_CUDA_VERSION}/arm64) + else() + set(_cudnn_arch_suffixes lib/x64 lib/${onnxruntime_CUDA_VERSION}/x64) + endif() + endif() find_library( ${NAME}_LIBRARY ${NAME} "lib${NAME}.so.${CUDNN_MAJOR_VERSION}" HINTS $ENV{CUDNN_PATH} ${CUDNN_PATH} ${Python_SITEARCH}/nvidia/cudnn ${CUDAToolkit_LIBRARY_DIR} - PATH_SUFFIXES lib64 lib/x64 lib lib/${onnxruntime_CUDA_VERSION}/x64 + PATH_SUFFIXES lib64 ${_cudnn_arch_suffixes} lib REQUIRED ) diff --git a/cmake/external/onnxruntime_external_deps.cmake b/cmake/external/onnxruntime_external_deps.cmake index ed3b0aa8192a7..9edea3042dbfb 100644 --- a/cmake/external/onnxruntime_external_deps.cmake +++ b/cmake/external/onnxruntime_external_deps.cmake @@ -880,6 +880,27 @@ endif() set(onnxruntime_LINK_DIRS) if (onnxruntime_USE_CUDA) + # Work around a CMake limitation (present through at least CMake 3.31 and current + # upstream master) when building natively on a Windows-on-ARM64 host. FindCUDAToolkit + # only sets the Windows import-library search suffix when the host is x64: + # + # if(CMAKE_HOST_SYSTEM_NAME STREQUAL "Windows") + # if(CMAKE_HOST_SYSTEM_PROCESSOR STREQUAL "AMD64") + # set(_CUDAToolkit_win_search_dirs lib/x64) + # set(_CUDAToolkit_win_stub_search_dirs lib/x64/stubs) + # + # On an ARM64 host the suffix is left empty, so find_library() for cudart only looks in + # "lib64" and never finds /lib/.../cudart.lib. find_package(CUDAToolkit) then + # fails with: Could NOT find CUDAToolkit (missing: CUDA_CUDART). Pre-seed the (internal) + # search-suffix variables with win-arm64 import-library locations (lib/arm64 and + # lib/arm64/stubs) so the toolkit's cudart.lib can be found. FindCUDAToolkit unsets + # these at the end, so this only affects the search below and is a no-op once CMake + # gains native WoA support. + if(CMAKE_HOST_SYSTEM_NAME STREQUAL "Windows" AND CMAKE_HOST_SYSTEM_PROCESSOR STREQUAL "ARM64") + set(_CUDAToolkit_win_search_dirs lib/arm64) + set(_CUDAToolkit_win_stub_search_dirs lib/arm64/stubs) + endif() + find_package(CUDAToolkit REQUIRED) # cuDNN is not needed for minimal CUDA builds (e.g., TensorRT-only builds) diff --git a/cmake/onnxruntime_common.cmake b/cmake/onnxruntime_common.cmake index b081e22e8b3f4..678f567922ba7 100644 --- a/cmake/onnxruntime_common.cmake +++ b/cmake/onnxruntime_common.cmake @@ -94,18 +94,23 @@ else() "${ONNXRUNTIME_ROOT}/core/platform/device_discovery_default.cc") endif() +# Raw /bigobj is a cl.exe option. Do not apply it to CUDA sources; nvcc treats a +# standalone /bigobj as an input file on Windows ARM64 CUDA 13.1. +set(onnxruntime_msvc_bigobj_compile_option + "$<$>,$>>:/bigobj>") + if(onnxruntime_target_platform STREQUAL "ARM64EC") if (MSVC) link_directories("$ENV{VCINSTALLDIR}/Tools/MSVC/$ENV{VCToolsVersion}/lib/ARM64EC") link_directories("$ENV{VCINSTALLDIR}/Tools/MSVC/$ENV{VCToolsVersion}/ATLMFC/lib/ARM64EC") link_libraries(softintrin.lib) - add_compile_options("$<$>:/bigobj>") + add_compile_options("${onnxruntime_msvc_bigobj_compile_option}") endif() endif() if(onnxruntime_target_platform STREQUAL "ARM64") if (MSVC) - add_compile_options("$<$>:/bigobj>") + add_compile_options("${onnxruntime_msvc_bigobj_compile_option}") endif() endif() diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake index f486b89ec0206..f9487fb2c8b5e 100644 --- a/cmake/onnxruntime_providers_cuda.cmake +++ b/cmake/onnxruntime_providers_cuda.cmake @@ -364,8 +364,9 @@ target_compile_options(${target} PRIVATE "$<$:/Zc:preprocessor>") target_compile_options(${target} PRIVATE "$<$:SHELL:-Xcompiler /Zc:__cplusplus>") target_compile_options(${target} PRIVATE "$<$:SHELL:-Xcompiler /Zc:preprocessor>") - target_compile_options(${target} PRIVATE "$<$:SHELL:-Xcompiler /bigobj>") - target_compile_options(${target} PRIVATE "$<$:/bigobj>") + # Pass /bigobj to the CUDA host compiler using dash spelling. Raw /bigobj is excluded + # from global ARM64 CUDA options in onnxruntime_common.cmake because nvcc parses it as input. + target_compile_options(${target} PRIVATE "$<$:-Xcompiler=-bigobj>") # /permissive is required for CUTLASS cute headers and to work around MSVC template resolution # issues with abseil headers when compiled through nvcc. # See https://github.com/NVIDIA/cutlass/issues/3065 @@ -449,7 +450,7 @@ target_compile_definitions(${target} PRIVATE COMPILE_HOPPER_TMA_GROUPED_GEMMS) endif() if (MSVC) - target_compile_options(${target} PRIVATE "$<$:SHELL:-Xcompiler /bigobj>") + # Do NOT add another /bigobj here: the MSVC block above already forwards it to cl. target_compile_options(${target} PRIVATE "$<$:SHELL:-Xcompiler /wd4172>") endif() endif() diff --git a/cmake/onnxruntime_providers_cuda_plugin.cmake b/cmake/onnxruntime_providers_cuda_plugin.cmake index 0f3b42a3f549f..a2755f80c6027 100644 --- a/cmake/onnxruntime_providers_cuda_plugin.cmake +++ b/cmake/onnxruntime_providers_cuda_plugin.cmake @@ -174,7 +174,9 @@ if (MSVC) "$<$:SHELL:-Xcompiler /wd4211>" "$<$:SHELL:-Xcompiler /Zc:__cplusplus>" "$<$:SHELL:-Xcompiler /Zc:preprocessor>" - "$<$:SHELL:-Xcompiler /bigobj>" + # Pass /bigobj to the CUDA host compiler using dash spelling. Raw /bigobj is excluded + # from global ARM64 CUDA options in onnxruntime_common.cmake because nvcc parses it as input. + "$<$:-Xcompiler=-bigobj>" ) target_compile_options(onnxruntime_providers_cuda_plugin PRIVATE @@ -232,7 +234,10 @@ if (MSVC) "$<$:SHELL:-Xcompiler /wd4211>" "$<$:SHELL:-Xcompiler /Zc:__cplusplus>" "$<$:SHELL:-Xcompiler /Zc:preprocessor>" - "$<$:SHELL:-Xcompiler /bigobj>" + # Unlike the options explicitly paired with -Xcompiler above, the raw /bigobj inherited + # from global compile options is parsed by nvcc as an input file on ARM64. Exclude that raw + # option in onnxruntime_common.cmake and forward its dash-spelled equivalent explicitly. + "$<$:-Xcompiler=-bigobj>" ) endif() diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 14b01f1f25473..f50baf6f993a6 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -386,11 +386,19 @@ if (WIN32) if(onnxruntime_CUDNN_HOME) # may have x64 in the path # may have a path with CUDA toolkit version if multiple installed on the machine + # Since CUDA 13.x, Windows cuDNN DLLs live under an architecture-specific subdirectory + # (bin/x64 for win-x64, bin/arm64 for win-arm64). A single cuDNN package may ship both + # arches, so only search the subdirectory matching the current target platform. + if(onnxruntime_target_platform STREQUAL "ARM64" OR onnxruntime_target_platform STREQUAL "ARM64EC") + set(_cudnn_arch "arm64") + else() + set(_cudnn_arch "x64") + endif() set(CUDNN_SEARCH_PATHS + "${onnxruntime_CUDNN_HOME}/bin/${_cudnn_arch}/cudnn64_*.dll" + "${onnxruntime_CUDNN_HOME}/bin/${onnxruntime_CUDA_VERSION}/${_cudnn_arch}/cudnn64_*.dll" "${onnxruntime_CUDNN_HOME}/bin/cudnn64_*.dll" - "${onnxruntime_CUDNN_HOME}/bin/x64/cudnn64_*.dll" "${onnxruntime_CUDNN_HOME}/bin/${onnxruntime_CUDA_VERSION}/cudnn64_*.dll" - "${onnxruntime_CUDNN_HOME}/bin/${onnxruntime_CUDA_VERSION}/x64/cudnn64_*.dll" ) set(CUDNN_DLL_PATH "") foreach(search_path ${CUDNN_SEARCH_PATHS}) diff --git a/onnxruntime/core/providers/cuda/cu_inc/cub.cuh b/onnxruntime/core/providers/cuda/cu_inc/cub.cuh index e88f780fd30b9..4eedf303c5ed8 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/cub.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/cub.cuh @@ -6,14 +6,26 @@ // Wrapper include for . // Macro definition of `__out` (SAL annotation) conflicts with parameter name `__out` in -// /include/cccl/cuda/__ptx/instructions/generated/tcgen05_ld.h. -// As a workaround, undefine `__out` before including cub/cub.cuh. +// /include/cccl/cuda/__ptx/instructions/generated/tcgen05_ld.h (and the +// other tcgen05_*.h PTX instruction headers). As a workaround, undefine `__out` before +// including the CCCL headers. #if defined(_MSC_VER) #pragma push_macro("__out") #undef __out // CCCL cub/config.cuh has a #pragma warning(pop) without matching push in CUDA v13.3. #pragma warning(push) #pragma warning(disable : 4193) + +// Depending on the CUDA toolkit version, the CCCL PTX instruction headers are not always +// reached through (e.g. they may be pulled in later by another CUDA header, +// after the macro has been restored below). Parse the PTX umbrella here, while `__out` is +// undefined, so those headers are processed safely regardless of include order. Guarded by +// __has_include so toolkits without are unaffected. +#if defined(__has_include) +#if __has_include() +#include +#endif +#endif #endif #include diff --git a/plugin-ep-cuda/csharp/Microsoft.ML.OnnxRuntime.EP.Cuda/Microsoft.ML.OnnxRuntime.EP.Cuda.csproj b/plugin-ep-cuda/csharp/Microsoft.ML.OnnxRuntime.EP.Cuda/Microsoft.ML.OnnxRuntime.EP.Cuda.csproj index bed4cfa8e2130..b34cb8ee82460 100644 --- a/plugin-ep-cuda/csharp/Microsoft.ML.OnnxRuntime.EP.Cuda/Microsoft.ML.OnnxRuntime.EP.Cuda.csproj +++ b/plugin-ep-cuda/csharp/Microsoft.ML.OnnxRuntime.EP.Cuda/Microsoft.ML.OnnxRuntime.EP.Cuda.csproj @@ -42,6 +42,14 @@ Condition="Exists('runtimes\win-x64\native')" /> + + + + (RID, list of native binary filenames expected in the source dir). PLATFORMS: dict[str, tuple[str, tuple[str, ...]]] = { "win_x64": ("win-x64", ("onnxruntime_providers_cuda.dll",)), + "win_arm64": ("win-arm64", ("onnxruntime_providers_cuda.dll",)), "linux_x64": ("linux-x64", ("libonnxruntime_providers_cuda.so",)), "linux_aarch64": ("linux-arm64", ("libonnxruntime_providers_cuda.so",)), } diff --git a/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml b/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml index a6c840736acb8..4b713fbc1a3d7 100644 --- a/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml @@ -21,23 +21,31 @@ parameters: type: boolean default: true +- name: build_windows_arm64 + displayName: 'Build Windows ARM64 (CUDA 13.x only)' + type: boolean + default: true + - name: build_linux_x64 displayName: 'Build Linux x64' type: boolean default: true - name: build_linux_aarch64 - displayName: 'Build Linux aarch64 (CUDA 13.0 only)' + displayName: 'Build Linux aarch64 (CUDA 13.x only)' type: boolean - default: false + default: true +# CUDA major.minor series for the build. The exact toolkit minor can differ per platform, +# so the CUDA 13 option uses a major-only label. It is translated to concrete toolkit +# versions when invoking the packaging stage below. - name: cuda_version displayName: 'CUDA Version' type: string - default: '12.8' + default: '13.x' values: - '12.8' - - '13.0' + - '13.x' - name: package_type displayName: 'Package Type' @@ -64,9 +72,9 @@ variables: # Non-dev package versions (release, RC) must use Release build type - name: invalidBuildTypeConfig value: ${{ and(ne(parameters.package_type, 'dev'), ne(parameters.cmake_build_type, 'Release')) }} - # aarch64 is only available for CUDA 13.0 + # aarch64 is only available for CUDA 13.x - name: invalidAArch64Config - value: ${{ and(eq(parameters.build_linux_aarch64, true), ne(parameters.cuda_version, '13.0')) }} + value: ${{ and(eq(parameters.build_linux_aarch64, true), ne(parameters.cuda_version, '13.x')) }} extends: template: v1/1ES.Official.PipelineTemplate.yml@1esPipelines @@ -117,25 +125,29 @@ extends: displayName: 'ERROR: Non-dev package version requires Release build type' - ${{ if eq(variables['invalidAArch64Config'], 'True') }}: - script: | - echo "##vso[task.logissue type=error]Linux aarch64 build is only available for CUDA 13.0." + echo "##vso[task.logissue type=error]Linux aarch64 build is only available for CUDA 13.x." exit 1 - displayName: 'ERROR: aarch64 requires CUDA 13.0' + displayName: 'ERROR: aarch64 requires CUDA 13.x' - ${{ else }}: - template: stages/plugin-cuda-packaging-stage.yml parameters: build_windows_x64: ${{ parameters.build_windows_x64 }} + build_windows_arm64: ${{ parameters.build_windows_arm64 }} build_linux_x64: ${{ parameters.build_linux_x64 }} build_linux_aarch64: ${{ parameters.build_linux_aarch64 }} - cuda_version: ${{ parameters.cuda_version }} + cuda_version: ${{ replace(parameters.cuda_version, '13.x', '13.0') }} + arm64_cuda_version: '13.1' package_type: ${{ parameters.package_type }} version_file: ${{ variables.epVersionFile }} cmake_build_type: ${{ parameters.cmake_build_type }} ${{ if eq(parameters.cuda_version, '12.8') }}: python_package_name: 'onnxruntime-ep-cuda12' docker_base_image: 'onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda12_x64_almalinux8_gcc14:20251017.1' - cmake_cuda_archs: '61-real;75-real;86-real;89-real;120-real;120-virtual' - ${{ if eq(parameters.cuda_version, '13.0') }}: + cmake_x64_cuda_archs: '61-real;75-real;86-real;89-real;120-real;120-virtual' + cmake_arm64_cuda_archs: '61-real;75-real;86-real;89-real;120-real;120-virtual' + ${{ if eq(parameters.cuda_version, '13.x') }}: python_package_name: 'onnxruntime-ep-cuda13' docker_base_image: 'onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda13_x64_almalinux8_gcc14:20251107.1' docker_base_image_aarch64: 'onnxruntimebuildcache.azurecr.io/public/azureml/onnxruntime_build_cuda13_aarch64_almalinux9_gcc14:20260323.1' - cmake_cuda_archs: '75-real;80-real;86-real;89-real;90-real;100-real;120-real;120-virtual' + cmake_x64_cuda_archs: '75-real;80-real;86-real;89-real;90-real;120-real;120-virtual' + cmake_arm64_cuda_archs: '110-real;120-real;121-real;120-virtual' diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-nuget-packaging-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-nuget-packaging-stage.yml index 02d5d08495c7c..ae7ec2fcd5823 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-nuget-packaging-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-nuget-packaging-stage.yml @@ -17,6 +17,7 @@ parameters: type: object default: win_x64: false + win_arm64: false linux_x64: false linux_aarch64: false @@ -26,6 +27,8 @@ stages: dependsOn: - ${{ if eq(parameters.platforms.win_x64, true) }}: - Win_plugin_cuda_x64_Build + - ${{ if eq(parameters.platforms.win_arm64, true) }}: + - Win_plugin_cuda_arm64_Build - ${{ if eq(parameters.platforms.linux_x64, true) }}: - Linux_plugin_cuda_x64 - ${{ if eq(parameters.platforms.linux_aarch64, true) }}: @@ -81,6 +84,13 @@ stages: artifactName: cuda_plugin_win_x64 targetPath: '$(Build.BinariesDirectory)\artifacts\win_x64' + - ${{ if eq(parameters.platforms.win_arm64, true) }}: + - task: DownloadPipelineArtifact@2 + displayName: 'Download win-arm64 artifacts' + inputs: + artifactName: cuda_plugin_win_arm64 + targetPath: '$(Build.BinariesDirectory)\artifacts\win_arm64' + - ${{ if eq(parameters.platforms.linux_x64, true) }}: - task: DownloadPipelineArtifact@2 displayName: 'Download linux-x64 artifacts' @@ -111,6 +121,7 @@ stages: # time and resolve to a boolean value 'True' or 'False'. Compare case-insensitively. platforms_enabled = { "win_x64": "${{ parameters.platforms.win_x64 }}".lower() == "true", + "win_arm64": "${{ parameters.platforms.win_arm64 }}".lower() == "true", "linux_x64": "${{ parameters.platforms.linux_x64 }}".lower() == "true", "linux_aarch64": "${{ parameters.platforms.linux_aarch64 }}".lower() == "true", } diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml index 84260c2ee7c43..5f9837fddb69e 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml @@ -4,6 +4,11 @@ parameters: type: boolean default: true +- name: build_windows_arm64 + displayName: 'Build Windows ARM64' + type: boolean + default: false + - name: build_linux_x64 displayName: 'Build Linux x64' type: boolean @@ -19,6 +24,11 @@ parameters: displayName: 'CUDA Version' default: '12.8' +- name: arm64_cuda_version + type: string + displayName: 'Windows ARM64 CUDA Version' + default: '13.1' + - name: package_type displayName: 'Package Type' type: string @@ -51,11 +61,16 @@ parameters: displayName: 'Linux aarch64 docker base image' default: '' -- name: cmake_cuda_archs +- name: cmake_x64_cuda_archs type: string - displayName: 'CMAKE_CUDA_ARCHITECTURES' + displayName: 'CMAKE_CUDA_ARCHITECTURES for x64' default: '61-real;75-real;86-real;89-real;120-real;120-virtual' +- name: cmake_arm64_cuda_archs + type: string + displayName: 'CMAKE_CUDA_ARCHITECTURES for ARM64' + default: '110-real;120-real;121-real;120-virtual' + - name: python_version type: string displayName: 'Python version' @@ -110,9 +125,14 @@ stages: $metadata = [ordered]@{ cuda_version = '${{ parameters.cuda_version }}' cuda_version_no_dot = '${{ replace(parameters.cuda_version, '.', '') }}' + cuda_version_win_x64 = '${{ parameters.cuda_version }}' + cuda_version_win_x64_no_dot = '${{ replace(parameters.cuda_version, '.', '') }}' + cuda_version_win_arm64 = '${{ parameters.arm64_cuda_version }}' + cuda_version_win_arm64_no_dot = '${{ replace(parameters.arm64_cuda_version, '.', '') }}' docker_base_image_linux_x64 = '${{ parameters.docker_base_image }}' python_artifact_linux_x64 = 'cuda_plugin_python_linux_x64_cuda${{ replace(parameters.cuda_version, '.', '') }}' python_artifact_win_x64 = 'cuda_plugin_python_win_x64_cuda${{ replace(parameters.cuda_version, '.', '') }}' + python_artifact_win_arm64 = 'cuda_plugin_python_win_arm64_cuda${{ replace(parameters.arm64_cuda_version, '.', '') }}' nuget_artifact = 'cuda_plugin_nuget' } @@ -125,8 +145,22 @@ stages: - ${{ if eq(parameters.build_windows_x64, true) }}: - template: plugin-win-cuda-stage.yml parameters: + arch: 'x64' cuda_version: ${{ parameters.cuda_version }} - cmake_cuda_archs: ${{ parameters.cmake_cuda_archs }} + cmake_cuda_archs: ${{ parameters.cmake_x64_cuda_archs }} + package_version: ${{ parameters.package_type }} + version_file: ${{ parameters.version_file }} + python_package_name: ${{ parameters.python_package_name }} + cmake_build_type: ${{ parameters.cmake_build_type }} + + # Windows ARM64 + - ${{ if eq(parameters.build_windows_arm64, true) }}: + - template: plugin-win-cuda-stage.yml + parameters: + arch: 'arm64' + cuda_version: ${{ parameters.cuda_version }} + arm64_cuda_version: ${{ parameters.arm64_cuda_version }} + cmake_cuda_archs: ${{ parameters.cmake_arm64_cuda_archs }} package_version: ${{ parameters.package_type }} version_file: ${{ parameters.version_file }} python_package_name: ${{ parameters.python_package_name }} @@ -140,7 +174,7 @@ stages: arch: 'x64' machine_pool: 'onnxruntime-Ubuntu2404-AMD-CPU' cuda_version: ${{ parameters.cuda_version }} - cmake_cuda_archs: ${{ parameters.cmake_cuda_archs }} + cmake_cuda_archs: ${{ parameters.cmake_x64_cuda_archs }} package_version: ${{ parameters.package_type }} version_file: ${{ parameters.version_file }} python_package_name: ${{ parameters.python_package_name }} @@ -150,7 +184,7 @@ stages: docker_python_exe_path: ${{ parameters.docker_python_exe_path }} artifact_name: cuda_plugin_linux_x64 - # Linux aarch64 (CUDA 13.0 only) + # Linux aarch64 (CUDA 13.x only) - ${{ if eq(parameters.build_linux_aarch64, true) }}: - template: plugin-linux-cuda-stage.yml parameters: @@ -158,7 +192,7 @@ stages: arch: 'aarch64' machine_pool: 'onnxruntime-linux-ARM64-CPU-2019' cuda_version: ${{ parameters.cuda_version }} - cmake_cuda_archs: ${{ parameters.cmake_cuda_archs }} + cmake_cuda_archs: ${{ parameters.cmake_arm64_cuda_archs }} package_version: ${{ parameters.package_type }} version_file: ${{ parameters.version_file }} python_package_name: ${{ parameters.python_package_name }} @@ -176,23 +210,23 @@ stages: DoEsrp: true platforms: win_x64: ${{ parameters.build_windows_x64 }} + win_arm64: ${{ parameters.build_windows_arm64 }} linux_x64: ${{ parameters.build_linux_x64 }} linux_aarch64: ${{ parameters.build_linux_aarch64 }} # Create zip packages for Foundry Local consumption. - - ${{ if or(eq(parameters.build_windows_x64, true), eq(parameters.build_linux_x64, true)) }}: + - ${{ if or(eq(parameters.build_windows_x64, true), eq(parameters.build_windows_arm64, true), eq(parameters.build_linux_x64, true), eq(parameters.build_linux_aarch64, true)) }}: - stage: Package_Foundry_Local_CUDA_Zips displayName: 'Package Foundry Local CUDA Plugin-EP Zips' dependsOn: - ${{ if eq(parameters.build_windows_x64, true) }}: - Win_plugin_cuda_x64_Build + - ${{ if eq(parameters.build_windows_arm64, true) }}: + - Win_plugin_cuda_arm64_Build - ${{ if eq(parameters.build_linux_x64, true) }}: - Linux_plugin_cuda_x64 - # TODO: win-arm64 and linux-aarch64 are still being developed. - # - ${{ if eq(parameters.build_windows_arm64, true) }}: - # - Win_plugin_cuda_arm64_Build - # - ${{ if eq(parameters.build_linux_aarch64, true) }}: - # - Linux_plugin_cuda_aarch64 + - ${{ if eq(parameters.build_linux_aarch64, true) }}: + - Linux_plugin_cuda_aarch64 # NOTE: macOS arm64 (Apple Silicon) does not have CUDA support since # Apple devices do not use Nvidia GPUs. jobs: @@ -222,6 +256,13 @@ stages: artifactName: cuda_plugin_win_x64 targetPath: $(Build.SourcesDirectory)/cuda-plugin-win-x64 + - ${{ if eq(parameters.build_windows_arm64, true) }}: + - task: DownloadPipelineArtifact@2 + displayName: 'Download cuda_plugin_win_arm64' + inputs: + artifactName: cuda_plugin_win_arm64 + targetPath: $(Build.SourcesDirectory)/cuda-plugin-win-arm64 + - ${{ if eq(parameters.build_linux_x64, true) }}: - task: DownloadPipelineArtifact@2 displayName: 'Download cuda_plugin_linux_x64' @@ -229,20 +270,12 @@ stages: artifactName: cuda_plugin_linux_x64 targetPath: $(Build.SourcesDirectory)/cuda-plugin-linux-x64 - # TODO: win-arm64 and linux-aarch64 are still being developed. - # - ${{ if eq(parameters.build_windows_arm64, true) }}: - # - task: DownloadPipelineArtifact@2 - # displayName: 'Download cuda_plugin_win_arm64' - # inputs: - # artifactName: cuda_plugin_win_arm64 - # targetPath: $(Build.SourcesDirectory)/cuda-plugin-win-arm64 - - # - ${{ if eq(parameters.build_linux_aarch64, true) }}: - # - task: DownloadPipelineArtifact@2 - # displayName: 'Download cuda_plugin_linux_aarch64' - # inputs: - # artifactName: cuda_plugin_linux_aarch64 - # targetPath: $(Build.SourcesDirectory)/cuda-plugin-linux-aarch64 + - ${{ if eq(parameters.build_linux_aarch64, true) }}: + - task: DownloadPipelineArtifact@2 + displayName: 'Download cuda_plugin_linux_aarch64' + inputs: + artifactName: cuda_plugin_linux_aarch64 + targetPath: $(Build.SourcesDirectory)/cuda-plugin-linux-aarch64 - task: PowerShell@2 displayName: 'Create version.json and zip packages for each platform' @@ -261,10 +294,9 @@ stages: # NOTE: macOS arm64 (Apple Silicon) is excluded — no CUDA support since Apple devices do not use Nvidia GPUs. $platforms = @( @{ name = 'win-x64'; dir = '$(Build.SourcesDirectory)/cuda-plugin-win-x64' }, - @{ name = 'linux-x64'; dir = '$(Build.SourcesDirectory)/cuda-plugin-linux-x64' } - # TODO: win-arm64 and linux-aarch64 are still being developed. - # @{ name = 'win-arm64'; dir = '$(Build.SourcesDirectory)/cuda-plugin-win-arm64' }, - # @{ name = 'linux-aarch64'; dir = '$(Build.SourcesDirectory)/cuda-plugin-linux-aarch64' } + @{ name = 'win-arm64'; dir = '$(Build.SourcesDirectory)/cuda-plugin-win-arm64' }, + @{ name = 'linux-x64'; dir = '$(Build.SourcesDirectory)/cuda-plugin-linux-x64' }, + @{ name = 'linux-aarch64'; dir = '$(Build.SourcesDirectory)/cuda-plugin-linux-aarch64' } ) $resolvedVersion = $null diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml index ef292c2ae3c72..e81b7510c7183 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml @@ -1,4 +1,11 @@ parameters: +- name: arch + type: string + default: x64 + values: + - x64 + - arm64 + - name: package_version type: string default: dev @@ -23,6 +30,16 @@ parameters: type: string default: '12.8' +# win-arm64 uses a separate CUDA toolkit and cuDNN build than win-x64. NVIDIA only +# ships Windows-on-ARM CUDA for 13.x, hosted under a 'win-arm64/' blob path prefix. +- name: arm64_cuda_version + type: string + default: '13.1' + +- name: arm64_cudnn_folder + type: string + default: '9.24.0.17_cuda13' + - name: python_package_name type: string @@ -31,16 +48,23 @@ parameters: default: '61-real;75-real;86-real;89-real;120-real;120-virtual' stages: - - stage: Win_plugin_cuda_x64_Build + - stage: Win_plugin_cuda_${{ parameters.arch }}_Build dependsOn: [] jobs: - - job: Win_plugin_cuda_x64_Build + - job: Win_plugin_cuda_${{ parameters.arch }}_Build timeoutInMinutes: 360 workspace: clean: all pool: - name: onnxruntime-Win-CPU-VS2022-Latest - os: windows + ${{ if eq(parameters.arch, 'arm64') }}: + name: onnxruntime-qnn-windows-vs-2022-arm64 + os: windows + hostArchitecture: Arm64 + demands: + - Agent.Version -equals 4.264.2 + ${{ else }}: + name: onnxruntime-Win-CPU-VS2022-Latest + os: windows templateContext: sdl: codeSignValidation: @@ -54,7 +78,7 @@ stages: outputs: - output: pipelineArtifact targetPath: $(Build.ArtifactStagingDirectory) - artifactName: cuda_plugin_win_x64 + artifactName: cuda_plugin_win_${{ parameters.arch }} variables: - template: ../templates/common-variables.yml - name: GRADLE_OPTS @@ -68,7 +92,7 @@ stages: - template: ../templates/setup-build-tools.yml parameters: - host_cpu_arch: 'x64' + host_cpu_arch: ${{ parameters.arch }} python_version: ${{ parameters.python_version }} - template: ../templates/set-nightly-build-option-variable-step.yml @@ -85,26 +109,46 @@ stages: env: TMPDIR: "$(Agent.TempDirectory)" - - task: AzureCLI@2 - displayName: 'Download CUDA SDK v${{ parameters.cuda_version }}' - inputs: - azureSubscription: AIInfraBuildOnnxRuntimeOSS - scriptType: 'batch' - scriptLocation: 'inlineScript' - inlineScript: | - set AZCOPY_AUTO_LOGIN_TYPE=AZCLI - azcopy.exe cp --recursive https://lotusscus.blob.core.windows.net/models/cuda_sdk/v${{ parameters.cuda_version }} "$(Agent.TempDirectory)" - - # Since CUDA 13.0, CUDA DLLs are in bin\x64 folder instead of bin folder for Windows. - - powershell: | - Write-Host "Adding CUDA to PATH" - Write-Host "##vso[task.prependpath]$env:AGENT_TEMPDIRECTORY\v${{ parameters.cuda_version }}\bin" - Write-Host "##vso[task.prependpath]$env:AGENT_TEMPDIRECTORY\v${{ parameters.cuda_version }}\bin\x64" - Write-Host "##vso[task.prependpath]$env:AGENT_TEMPDIRECTORY\v${{ parameters.cuda_version }}\extras\CUPTI\lib64" - displayName: 'Add CUDA to PATH' - - # Download cuDNN separately for CUDA 13.0 - - ${{ if eq(parameters.cuda_version, '13.0') }}: + # CUDA SDK download. win-x64 toolkits live at cuda_sdk/v; win-arm64 toolkits + # live at cuda_sdk/win-arm64/v (NVIDIA Windows-on-ARM CUDA is 13.x only). + - ${{ if ne(parameters.arch, 'arm64') }}: + - task: AzureCLI@2 + displayName: 'Download CUDA SDK v${{ parameters.cuda_version }}' + inputs: + azureSubscription: AIInfraBuildOnnxRuntimeOSS + scriptType: 'batch' + scriptLocation: 'inlineScript' + inlineScript: | + set AZCOPY_AUTO_LOGIN_TYPE=AZCLI + azcopy.exe cp --recursive https://lotusscus.blob.core.windows.net/models/cuda_sdk/v${{ parameters.cuda_version }} "$(Agent.TempDirectory)" + + # Since CUDA 13.0, CUDA DLLs are in bin\x64 folder instead of bin folder for Windows. + - powershell: | + Write-Host "Adding CUDA to PATH" + Write-Host "##vso[task.prependpath]$env:AGENT_TEMPDIRECTORY\v${{ parameters.cuda_version }}\bin" + Write-Host "##vso[task.prependpath]$env:AGENT_TEMPDIRECTORY\v${{ parameters.cuda_version }}\bin\x64" + Write-Host "##vso[task.prependpath]$env:AGENT_TEMPDIRECTORY\v${{ parameters.cuda_version }}\extras\CUPTI\lib64" + displayName: 'Add CUDA to PATH' + + - ${{ else }}: + # The win-arm64 agents (onnxruntime-qnn-windows-vs-2022-arm64) do not have the Azure + # CLI installed, so the AzureCLI@2 task cannot be used here (it fails with "Azure CLI + # 2.x is not installed on this machine"). azcopy is on PATH and can read this storage + # account directly, matching the QNN SDK download in + # templates/jobs/download_win_qnn_sdk.yml. + - powershell: | + azcopy.exe cp --recursive https://lotusscus.blob.core.windows.net/models/cuda_sdk/win-arm64/v${{ parameters.arm64_cuda_version }} "$(Agent.TempDirectory)" + displayName: 'Download CUDA SDK win-arm64 v${{ parameters.arm64_cuda_version }}' + + - powershell: | + Write-Host "Adding CUDA to PATH" + Write-Host "##vso[task.prependpath]$env:AGENT_TEMPDIRECTORY\v${{ parameters.arm64_cuda_version }}\bin" + Write-Host "##vso[task.prependpath]$env:AGENT_TEMPDIRECTORY\v${{ parameters.arm64_cuda_version }}\bin\arm64" + Write-Host "##vso[task.prependpath]$env:AGENT_TEMPDIRECTORY\v${{ parameters.arm64_cuda_version }}\extras\CUPTI\lib64" + displayName: 'Add CUDA to PATH' + + # Download cuDNN separately for CUDA 13.0 (win-x64) + - ${{ if and(ne(parameters.arch, 'arm64'), eq(parameters.cuda_version, '13.0')) }}: - task: AzureCLI@2 displayName: 'Download cuDNN for CUDA 13.0' inputs: @@ -115,8 +159,16 @@ stages: set AZCOPY_AUTO_LOGIN_TYPE=AZCLI azcopy.exe cp --recursive https://lotusscus.blob.core.windows.net/models/cudnn_sdk/$(cuda13_cudnn_folder) "$(Agent.TempDirectory)" - # CUDA 12.x build (no separate cuDNN) - - ${{ if ne(parameters.cuda_version, '13.0') }}: + # Download cuDNN for win-arm64 (hosted under cudnn_sdk/win-arm64/). + # Uses azcopy directly because the win-arm64 agents do not have the Azure CLI installed + # (see the CUDA SDK download step above). + - ${{ if eq(parameters.arch, 'arm64') }}: + - powershell: | + azcopy.exe cp --recursive https://lotusscus.blob.core.windows.net/models/cudnn_sdk/win-arm64/${{ parameters.arm64_cudnn_folder }} "$(Agent.TempDirectory)" + displayName: 'Download cuDNN for win-arm64 (${{ parameters.arm64_cudnn_folder }})' + + # CUDA 12.x build (no separate cuDNN). win-x64 only — win-arm64 always uses CUDA 13.x. + - ${{ if and(ne(parameters.arch, 'arm64'), ne(parameters.cuda_version, '13.0')) }}: - task: PythonScript@0 displayName: 'Build CUDA Plugin (CUDA ${{ parameters.cuda_version }})' inputs: @@ -143,8 +195,8 @@ stages: $(TelemetryOption) workingDirectory: '$(Build.BinariesDirectory)' - # CUDA 13.0 build (separate cuDNN folder) - - ${{ if eq(parameters.cuda_version, '13.0') }}: + # CUDA 13.0 build (separate cuDNN folder). win-x64 only — win-arm64 handled below. + - ${{ if and(ne(parameters.arch, 'arm64'), eq(parameters.cuda_version, '13.0')) }}: - task: PythonScript@0 displayName: 'Build CUDA Plugin (CUDA ${{ parameters.cuda_version }})' inputs: @@ -172,6 +224,36 @@ stages: $(TelemetryOption) workingDirectory: '$(Build.BinariesDirectory)' + # win-arm64 build (CUDA 13.x with a separate cuDNN folder). Runs natively on the + # ARM64 agent, so build.py auto-selects the ARM64 toolset (-A ARM64). + - ${{ if eq(parameters.arch, 'arm64') }}: + - task: PythonScript@0 + displayName: 'Build CUDA Plugin (win-arm64 CUDA ${{ parameters.arm64_cuda_version }})' + inputs: + scriptPath: '$(Build.SourcesDirectory)\tools\ci_build\build.py' + arguments: >- + --config ${{ parameters.cmake_build_type }} + --build_dir $(Build.BinariesDirectory) + --skip_submodule_sync + --cmake_generator "$(VSGenerator)" + --parallel + --nvcc_threads 1 + --flash_nvcc_threads 1 + --use_vcpkg + --use_vcpkg_ms_internal_asset_cache + --use_binskim_compliant_compile_flags + --update + --build + --use_cuda + --cuda_home="$(Agent.TempDirectory)\v${{ parameters.arm64_cuda_version }}" + --cudnn_home="$(Agent.TempDirectory)\${{ parameters.arm64_cudnn_folder }}" + --skip_tests + --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES="${{ parameters.cmake_cuda_archs }}" + --cmake_extra_defines onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON + --cmake_extra_defines $(PluginEpVersionDefine) + $(TelemetryOption) + workingDirectory: '$(Build.BinariesDirectory)' + # Esrp signing - template: ../templates/win-esrp-dll.yml parameters: @@ -203,19 +285,29 @@ stages: type nul > "$(Build.ArtifactStagingDirectory)\version\$(PluginPackageVersion)" displayName: 'Create version marker file' - - job: Win_plugin_cuda_x64_Python_Package - dependsOn: Win_plugin_cuda_x64_Build + - job: Win_plugin_cuda_${{ parameters.arch }}_Python_Package + dependsOn: Win_plugin_cuda_${{ parameters.arch }}_Build timeoutInMinutes: 30 workspace: clean: all pool: - name: onnxruntime-Win-CPU-VS2022-Latest - os: windows + ${{ if eq(parameters.arch, 'arm64') }}: + name: onnxruntime-qnn-windows-vs-2022-arm64 + os: windows + hostArchitecture: Arm64 + demands: + - Agent.Version -equals 4.264.2 + ${{ else }}: + name: onnxruntime-Win-CPU-VS2022-Latest + os: windows templateContext: outputs: - output: pipelineArtifact targetPath: '$(Build.ArtifactStagingDirectory)\python' - artifactName: cuda_plugin_python_win_x64_cuda${{ replace(parameters.cuda_version, '.', '') }} + ${{ if eq(parameters.arch, 'arm64') }}: + artifactName: cuda_plugin_python_win_arm64_cuda${{ replace(parameters.arm64_cuda_version, '.', '') }} + ${{ else }}: + artifactName: cuda_plugin_python_win_x64_cuda${{ replace(parameters.cuda_version, '.', '') }} variables: - template: ../templates/common-variables.yml steps: @@ -225,7 +317,7 @@ stages: - template: ../templates/setup-build-tools.yml parameters: - host_cpu_arch: 'x64' + host_cpu_arch: ${{ parameters.arch }} python_version: ${{ parameters.python_version }} - template: ../templates/set-nightly-build-option-variable-step.yml @@ -239,18 +331,14 @@ stages: - task: DownloadPipelineArtifact@2 displayName: 'Download plugin build artifacts' inputs: - artifactName: cuda_plugin_win_x64 + artifactName: cuda_plugin_win_${{ parameters.arch }} targetPath: '$(Build.BinariesDirectory)\plugin_artifacts' - - task: PowerShell@2 + - powershell: | + python -m pip install -r "$(Build.SourcesDirectory)\plugin-ep-cuda\python\requirements-build-wheel.txt" + python "$(Build.SourcesDirectory)\plugin-ep-cuda\python\build_wheel.py" ` + --binary_dir "$(Build.BinariesDirectory)\plugin_artifacts\bin" ` + --version "$(PluginPythonPackageVersion)" ` + --package_name "${{ parameters.python_package_name }}" ` + --output_dir "$(Build.ArtifactStagingDirectory)\python" displayName: 'Build Python wheel' - inputs: - targetType: inline - pwsh: true - script: | - python -m pip install -r "$(Build.SourcesDirectory)\plugin-ep-cuda\python\requirements-build-wheel.txt" - python "$(Build.SourcesDirectory)\plugin-ep-cuda\python\build_wheel.py" ` - --binary_dir "$(Build.BinariesDirectory)\plugin_artifacts\bin" ` - --version "$(PluginPythonPackageVersion)" ` - --package_name "${{ parameters.python_package_name }}" ` - --output_dir "$(Build.ArtifactStagingDirectory)\python" diff --git a/tools/ci_build/github/azure-pipelines/templates/read-cuda-plugin-metadata-steps.yml b/tools/ci_build/github/azure-pipelines/templates/read-cuda-plugin-metadata-steps.yml index 3b367abbcdc09..433e5d4aabffb 100644 --- a/tools/ci_build/github/azure-pipelines/templates/read-cuda-plugin-metadata-steps.yml +++ b/tools/ci_build/github/azure-pipelines/templates/read-cuda-plugin-metadata-steps.yml @@ -21,9 +21,14 @@ steps: $requiredProperties = @( 'cuda_version', 'cuda_version_no_dot', + 'cuda_version_win_x64', + 'cuda_version_win_x64_no_dot', + 'cuda_version_win_arm64', + 'cuda_version_win_arm64_no_dot', 'docker_base_image_linux_x64', 'python_artifact_linux_x64', 'python_artifact_win_x64', + 'python_artifact_win_arm64', 'nuget_artifact' ) @@ -35,11 +40,18 @@ steps: Write-Host "Detected CUDA version: $($metadata.cuda_version)" Write-Host "Linux Python artifact: $($metadata.python_artifact_linux_x64)" - Write-Host "Windows Python artifact: $($metadata.python_artifact_win_x64)" + Write-Host "Windows x64 Python artifact: $($metadata.python_artifact_win_x64)" + Write-Host "Windows ARM64 Python artifact: $($metadata.python_artifact_win_arm64)" Write-Host "##vso[task.setvariable variable=CudaVersion]$($metadata.cuda_version)" Write-Host "##vso[task.setvariable variable=CudaVersionNoDot]$($metadata.cuda_version_no_dot)" + Write-Host "##vso[task.setvariable variable=CudaVersionWinX64]$($metadata.cuda_version_win_x64)" + Write-Host "##vso[task.setvariable variable=CudaVersionWinX64NoDot]$($metadata.cuda_version_win_x64_no_dot)" + Write-Host "##vso[task.setvariable variable=CudaVersionWinArm64]$($metadata.cuda_version_win_arm64)" + Write-Host "##vso[task.setvariable variable=CudaVersionWinArm64NoDot]$($metadata.cuda_version_win_arm64_no_dot)" Write-Host "##vso[task.setvariable variable=CudaDockerBaseImage]$($metadata.docker_base_image_linux_x64)" Write-Host "##vso[task.setvariable variable=CudaPluginPythonLinuxArtifactName]$($metadata.python_artifact_linux_x64)" Write-Host "##vso[task.setvariable variable=CudaPluginPythonWinArtifactName]$($metadata.python_artifact_win_x64)" + Write-Host "##vso[task.setvariable variable=CudaPluginPythonWinX64ArtifactName]$($metadata.python_artifact_win_x64)" + Write-Host "##vso[task.setvariable variable=CudaPluginPythonWinArm64ArtifactName]$($metadata.python_artifact_win_arm64)" Write-Host "##vso[task.setvariable variable=CudaPluginNuGetArtifactName]$($metadata.nuget_artifact)" From e5feaccba9635495663ddbbe2a9873bd1a561ca4 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 17 Jul 2026 03:58:13 -0700 Subject: [PATCH 10/11] Enable fpA_intB GEMM in CUDA builds and add configurable options (#29622) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Enable fpA_intB GEMM in CUDA builds and add configurable options ### Summary This PR turns the CUDA fpA_intB (weight-only, FP activation × int weight) MatMulNBits path on by default in CUDA builds, replaces the ambiguous `ORT_FPA_INTB_GEMM` bitmask with a simple on/off flag, and adds session-config keys so the path and its autotuning buckets can be controlled per session. It also makes the CUTLASS tactic profiler CUDA-graph safe and configurable. ### Motivation The fpA_intB kernels were previously gated behind `onnxruntime_USE_FPA_INTB_GEMM=OFF` and an `ORT_FPA_INTB_GEMM` integer bitmask (`0x01=all`, `0x02=GEMV`, `0x04=int4`, `0x08=int8`). The bitmask was ambiguous (e.g. `6` could read as "int4 + GEMV" or "GEMV for int4 and int8") and the GEMM/GEMV kernels actually share one weight layout, so splitting them was never valid. Shipping the kernels by default and exposing a plain enable flag plus per-session config makes the feature usable and tunable without rebuilding. ### Key Changes #### Build enablement | File | Change | |------|--------| | [cmake/CMakeLists.txt](cmake/CMakeLists.txt) | `onnxruntime_USE_FPA_INTB_GEMM` becomes a `cmake_dependent_option`, defaulting **ON** when `onnxruntime_USE_CUDA` is enabled (OFF otherwise). | | [tools/ci_build/github/linux/build_cuda_c_api_package.sh](tools/ci_build/github/linux/build_cuda_c_api_package.sh), [build_linux_python_package.sh](tools/ci_build/github/linux/build_linux_python_package.sh), [build_tensorrt_c_api_package.sh](tools/ci_build/github/linux/build_tensorrt_c_api_package.sh) | Flip packaging builds from `onnxruntime_USE_FPA_INTB_GEMM=OFF` to `ON`. | #### Option simplification (bitmask → boolean) - Removed `kFpAIntBGemmOption_All/Gemv/Int4/Int8` and the old bitmask parsing. - Added `ParseFpAIntBEnabled`: `""` / `"0"` / `"off"` → disabled; any other value → enabled (numeric non-zero still works for back-compat). - The enable flag now only governs nodes **without** prepacked weights. A prepacked weight is already stored in the fpA_intB layout, so the choice was fixed at export time; the constructor forces the path on for prepacked nodes and only `ORT_ENFORCE`s that the shape/hardware actually support it. - GEMV is no longer independently toggleable: it is enabled whenever supported, since GEMM and GEMV share the same weight layout. #### New session-config keys (EP-agnostic, config wins over env) | Config key | Env fallback | Meaning | |------------|--------------|---------| | `ep.cuda.fpa_intb_gemm` | `ORT_FPA_INTB_GEMM` | Enable/disable the fpA_intB path (`0`/`off` vs `1`/`on`). | | `ep.cuda.fpa_intb_profile_m` | `ORT_FPA_INTB_PROFILE_M` | Comma-separated initial profile-M buckets (e.g. `"1,8,64,512"`); empty uses the default bucket set. | These are read by both the built-in CUDA EP and the CUDA plugin EP via `OpKernelInfo::GetConfigOptions()`. #### Profiler: CUDA-graph-safe, in-memory autotuning - Added `getBestConfigOrProfile()` for lazy single-bucket profiling outside CUDA-graph capture; during capture the kernel falls back to a pure lookup (`getBestConfig`) because profiling launches kernels, records/synchronizes events, and allocates scratch — all illegal during capture. - Added configurable profile-M buckets: `ParseProfileMList`, `setProfileMOverride`, `getProfileMBuckets`, plus `kEnvProfileM` and `kDefaultProfileMaxM` (default max M lowered to `2048`). - Clearer error when an M bucket was not profiled before capture ("run a warmup inference outside capture first"). #### Docs - [docs/contrib_ops/cuda/matmul_nbits.md](docs/contrib_ops/cuda/matmul_nbits.md): `ORT_FPA_INTB_GEMM` documented as int/string on/off; clarified prepacked-weight strictness. ### Testing Notes - New: [onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py](onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py) — exercises the prepacked fpA_intB path and the boolean/numeric back-compat values of the enable flag. - To verify locally (CUDA, SM ≥ 75): - Build with CUDA (fpA_intB now defaults ON): `./build.sh --use_cuda ...` - Run: `python -m pytest onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py` - Sanity-check config override: set `ep.cuda.fpa_intb_gemm=0` on a non-prepacked node and confirm the path is skipped; confirm a prepacked node still forces the path on. --- cmake/CMakeLists.txt | 2 +- docs/contrib_ops/cuda/matmul_nbits.md | 8 +- .../cuda/llm/common/cuda_runtime_utils.h | 25 ++- .../cuda/llm/fpA_intB_gemm_profiler.cc | 76 ++++++-- .../cuda/llm/fpA_intB_gemm_profiler.h | 31 +++- .../contrib_ops/cuda/llm/gemm_profiler.h | 170 +++++++++++++----- .../cuda/quantization/matmul_nbits.cc | 53 +++--- .../cuda/quantization/matmul_nbits.h | 93 +++++++--- .../test/contrib_ops/matmul_4bits_test.cc | 11 +- .../test_op_matmulnbits_prepacked_cuda.py | 107 +++++++++++ .../github/linux/build_cuda_c_api_package.sh | 2 +- .../linux/build_linux_python_package.sh | 3 +- .../linux/build_tensorrt_c_api_package.sh | 2 +- 13 files changed, 467 insertions(+), 116 deletions(-) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index c8bf4a7bb87ff..c1a05811b49fc 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -108,7 +108,7 @@ option(onnxruntime_USE_LEAN_ATTENTION "Build lean attention kernel for scaled do cmake_dependent_option(onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION "Build memory efficient attention kernel for scaled dot product attention" ON "onnxruntime_USE_CUDA" OFF) option(onnxruntime_USE_FP4_QMOE "Build CUDA QMoE FP4 kernels" OFF) option(onnxruntime_USE_FP8_QMOE "Build CUDA QMoE FP8 kernels" OFF) -option(onnxruntime_USE_FPA_INTB_GEMM "Build FpA IntB gemm cuda kernels" OFF) +cmake_dependent_option(onnxruntime_USE_FPA_INTB_GEMM "Build FpA IntB gemm cuda kernels" ON "onnxruntime_USE_CUDA" OFF) option(onnxruntime_USE_INT4_KV_CACHE "Build cuda kernels for int4 kv cache" OFF) option(onnxruntime_USE_FP8_KV_CACHE "Build cuda kernels for fp8 kv cache" ON) option(onnxruntime_QUICK_BUILD "Speed up build by skipping some kernels for faster development" OFF) diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index 3f8e1f1d8f616..6e382826de37d 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -263,8 +263,10 @@ Prepacked weights are intentionally strict: - If ORT was built without `onnxruntime_USE_FPA_INTB_GEMM=ON`, any nonzero `weight_prepacked` value throws during kernel construction. -- If `ORT_FPA_INTB_GEMM` is unset or `0`, any nonzero `weight_prepacked` value - throws instead of silently falling back to a raw-layout path. +- Any nonzero `weight_prepacked` value forces the fpA_intB path on, so the enable + flag (`ep.cuda.fpa_intb_gemm` session config, or the `ORT_FPA_INTB_GEMM` env + var) is ignored for prepacked weights — the layout choice was fixed at export + time and cannot be turned off at run time. - Nonzero `weight_prepacked` requires FP16 or BF16 input `A`, because only the CUDA fpA_intB path consumes this layout. - `weight_prepacked` must match the layout the selected kernel expects: `1` is @@ -293,7 +295,7 @@ present. `ComputeInternal` then: | Variable | Type / default | Effect | |----------|----------------|--------| | `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION` | bool, `0` | Disable the router GEMV specialization (§4.2); shapes fall back to the generic GEMV / dequant path. Useful for A/B benchmarking. | -| `ORT_FPA_INTB_GEMM` | int bitmask, `0` | Enable the CUTLASS weight-only path (§6). `0x01` = all, `0x02` = CUDA GEMV, `0x04` = int4, `0x08` = int8. `0` disables it. | +| `ORT_FPA_INTB_GEMM` | int/string, `0` | Enable the CUTLASS weight-only path (§6). `0` or `off` disables it, otherwise enables it. | | `ORT_MATMULNBITS_FORCE_CHUNKED` | int, `0` | Force the chunked dequant+GEMM fallback (§5) regardless of the size heuristic. | | `ORT_MATMULNBITS_CHUNK_SIZE` | int64, `32768` | Target rows per chunk in the chunked fallback. Values `< 1` reset to the default. | diff --git a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h index a3419b88b8d01..94cb5909a523c 100644 --- a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h +++ b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h @@ -16,13 +16,14 @@ */ #pragma once +#include #include +#include #include #ifdef ENABLE_FP8 #include #endif #include "core/providers/cuda/shared_inc/cuda_call.h" -#include "core/platform/env_var_utils.h" namespace onnxruntime::llm::common { inline int getDevice() { @@ -61,12 +62,28 @@ inline int getMaxSharedMemoryPerBlockOptin() { inline std::optional isCudaLaunchBlocking() { thread_local bool firstCall = true; thread_local std::optional result = std::nullopt; - if (!firstCall) { - if (ParseEnvironmentVariableWithDefault("CUDA_LAUNCH_BLOCKING", 0) == 1) { - result = true; + if (firstCall) { + // Read the env var directly here instead of via core/platform/env_var_utils.h. + // This is a leaf CUDA header pulled in by kernel headers (e.g. compute_occupancy.h) + // BEFORE the SHARED_PROVIDER bridge (provider_api.h) is established. env_var_utils.h + // unconditionally includes core/common/logging/logging.h while SHARED_PROVIDER is + // undefined, which then clashes with the logging stubs provider_api.h defines later in + // the same translation unit. Read the variable directly to keep this header self-contained. + // std::getenv is avoided on MSVC because it raises C4996 (treated as an error); _dupenv_s + // is the supported Windows-safe replacement. +#if defined(_WIN32) + char* env = nullptr; + size_t env_len = 0; + if (_dupenv_s(&env, &env_len, "CUDA_LAUNCH_BLOCKING") == 0 && env != nullptr) { + result = std::string(env) == "1"; } else { result = false; } + std::free(env); +#else + char const* env = std::getenv("CUDA_LAUNCH_BLOCKING"); + result = (env != nullptr && std::string(env) == "1"); +#endif firstCall = false; } return result; diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc index e5b15856a6c05..1b291e92bfafe 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc @@ -16,7 +16,15 @@ */ #if USE_FPA_INTB_GEMM #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" + +#include +#include +#include +#include + #include "contrib_ops/cuda/llm/common/workspace.h" +#include "core/common/parse_string.h" +#include "core/common/string_utils.h" using namespace onnxruntime::llm::common; using namespace onnxruntime::llm::kernels::cutlass_kernels; @@ -58,7 +66,7 @@ void WeightOnlyGroupwiseQuantGemmPluginProfiler::runTactic( onnxruntime::llm::kernels::fpA_intB_gemv::kernel_launcher(mArch, params, stream); } else { // run CUTLASS kernel - int const wsSize = mRunner->getWorkspaceSize(m, originalN, k); + int const wsSize = static_cast(mRunner->getWorkspaceSize(m, originalN, k)); if (mQuantBits == 8) { mRunner->gemm(actPtr, reinterpret_cast(weightPtr), inputScalesPtr, zerosPtr, biasesPtr, outputPtr, m, originalN, k, mGroupSize, tactic, workspacePtr, wsSize, stream); @@ -71,17 +79,17 @@ void WeightOnlyGroupwiseQuantGemmPluginProfiler::runTactic( size_t WeightOnlyGroupwiseQuantGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) { // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int const originalN = mQuantBits == 8 ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO; + int const originalN = static_cast(mQuantBits == 8 ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO); std::vector workspaces = { - maxM * k * sizeof(half), // A - k * n * sizeof(half), // B - k * originalN * sizeof(half) / mGroupSize, // scales - k * originalN * sizeof(half) / mGroupSize, // zeros - originalN * sizeof(half), // biases - maxM * originalN * sizeof(half), // C - mRunner->getWorkspaceSize(maxM, originalN, k) // workspace + /* A */ maxM * k * sizeof(half), + /* B */ k * n * sizeof(half), + /* scales */ k * originalN * sizeof(half) / mGroupSize, + /* zeros */ k * originalN * sizeof(half) / mGroupSize, + /* biases */ originalN * sizeof(half), + /* C */ maxM * originalN * sizeof(half), + mRunner->getWorkspaceSize(static_cast(maxM), originalN, static_cast(k)) // workspace }; - return calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); + return calculateTotalWorkspaceSize(workspaces.data(), static_cast(workspaces.size())); } std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::getTactics( @@ -97,5 +105,53 @@ bool WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic(int m, int /*n*/, i return true; } +std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMList(const std::string& value) { + std::vector result; + if (value.empty()) { + return result; + } + std::set unique; + for (const auto token : onnxruntime::utils::SplitString(value, ",", true)) { + const std::string trimmed_token = onnxruntime::utils::TrimString(token); + if (trimmed_token.empty()) { + continue; + } + int m = 0; + if (TryParseStringWithClassicLocale(trimmed_token, m) && m > 0) { + unique.insert(m); + } + } + result.assign(unique.begin(), unique.end()); + return result; +} + +std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::getProfileMBuckets( + int minM, int maxM, bool /*hasWeightOnlyCudaKernel*/) const { + int const lo = std::max(1, minM); + int const hi = std::max(lo, maxM); + + std::set buckets; + + if (!mProfileMOverride.empty()) { + for (int m : mProfileMOverride) { + buckets.insert(std::min(std::max(lo, m), hi)); + } + } else { + // Small default bucket set clamped to [lo, hi]. + static const int kDefault[] = {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048}; + for (int m : kDefault) { + if (m >= lo && m <= hi) { + buckets.insert(m); + } + } + } + + // Always include the decode bucket (M=1) and the top bucket so both extremes are tuned. + buckets.insert(lo); + buckets.insert(hi); + + return std::vector(buckets.begin(), buckets.end()); +} + } // namespace onnxruntime::llm::kernels::weight_only #endif diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h index b1336f45cab27..eabd6b0e78a75 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h @@ -16,17 +16,18 @@ */ #pragma once -#include "contrib_ops/cuda/llm/gemm_profiler.h" -#include "contrib_ops/cuda/llm/fpA_intB_gemm/fpA_intB_gemm.h" -#include "contrib_ops/cuda/llm/fpA_intB_gemv/fpA_intB_gemv.h" - #include #include #include #include #include +#include #include +#include "contrib_ops/cuda/llm/gemm_profiler.h" +#include "contrib_ops/cuda/llm/fpA_intB_gemm/fpA_intB_gemm.h" +#include "contrib_ops/cuda/llm/fpA_intB_gemv/fpA_intB_gemv.h" + using WeightOnlyGemmRunner = onnxruntime::llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunnerInterface; using WeightOnlyGemmRunnerPtr = std::shared_ptr; using KernelType = onnxruntime::llm::kernels::fpA_intB_gemv::KernelType; @@ -43,12 +44,31 @@ constexpr int32_t INT4_BITS = 4; constexpr int32_t FP16_INT4_RATIO = FP16_BITS / INT4_BITS; constexpr int32_t FP16_INT8_RATIO = FP16_BITS / INT8_BITS; +// Comma-separated list of M buckets to profile for MatMulNBits/fpA_intB. Overrides the default +// reduced bucket set. Example: ORT_FPA_INTB_PROFILE_M="1,8,64,512". +constexpr const char* kEnvProfileM = "ORT_FPA_INTB_PROFILE_M"; + +// Default top M that bounds the initial profile sweep when no override is given. Larger runtime +// M values are handled by lazy single-bucket profiling. +constexpr int kDefaultProfileMaxM = 2048; + class WeightOnlyGroupwiseQuantGemmPluginProfiler : public GemmPluginProfiler { public: using Config = onnxruntime::llm::cutlass_extensions::CutlassGemmConfig; + // Parses a comma-separated list of M buckets (e.g. "1,8,64,512") into a sorted, de-duplicated, + // positive list (empty when the string is empty/blank). Used for the ep.cuda.fpa_intb_profile_m + // session-config key and the ORT_FPA_INTB_PROFILE_M env var, both resolved by the kernel. + static std::vector ParseProfileMList(const std::string& value); + + // Overrides the initial profile M-bucket set for this profiler instance (per session). An empty + // list keeps the built-in default bucket set. Resolved by the kernel from session config / env. + void setProfileMOverride(std::vector ms) { + mProfileMOverride = std::move(ms); + } + void setQuant(int bits, bool has_bias, bool has_zeros) { mQuantBits = bits; mHasBiases = has_bias; @@ -74,6 +94,8 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler bool checkTactic(int m, int n, int k, Config const& tactic) const override; + std::vector getProfileMBuckets(int minM, int maxM, bool hasWeightOnlyCudaKernel) const override; + private: bool mHasBiases; bool mHasZeros; @@ -81,6 +103,7 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler int mGroupSize; KernelType mCudaKernelType; int mArch; + std::vector mProfileMOverride; }; } // namespace onnxruntime::llm::kernels::weight_only diff --git a/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h index d9c19cb734e91..28d570c51f910 100644 --- a/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h @@ -18,6 +18,7 @@ #include +#include #include #include #include @@ -110,10 +111,10 @@ class GemmPluginProfiler { using MProfileMap = std::unordered_map>; using MProfileMapPtr = std::shared_ptr; + // requires shared ownership to read from *this + using reader_lock = std::shared_lock; // requires exclusive ownership to write to *this - using reader_lock = std::unique_lock; - // requires shared ownership to read from other - using writer_lock = std::shared_lock; + using writer_lock = std::unique_lock; // Struct of continuing map if GEMMs to the best profiles for different Ms struct MNKProfileMap { @@ -168,6 +169,13 @@ class GemmPluginProfiler { std::optional getBestConfig(int m, GemmIdType const& gemmId) const; + // Like getBestConfig, but if the requested M bucket has not been profiled yet, profiles it + // lazily (single bucket) and inserts it into the in-process map. This briefly blocks the caller + // but guarantees a tuned tactic for any runtime M, which is what makes the reduced first-time M + // sweep safe. Must not be called while the compute stream is being captured into a CUDA graph + // (the caller is responsible for using getBestConfig instead during capture). + std::optional getBestConfigOrProfile(int m, GemmIdType const& gemmId); + virtual int getMaxProfileM() const; protected: @@ -183,10 +191,16 @@ class GemmPluginProfiler { virtual void initTmpData(int m, int n, int k, char* workspace, size_t size, cudaStream_t stream); + // Returns the ordered set of M buckets to profile during the initial sweep, given the + // (rounded) profile range [minM, maxM]. The default reproduces the historical dense sweep. + // Subclasses may override to profile a smaller, configurable bucket set. + virtual std::vector getProfileMBuckets(int minM, int maxM, bool hasWeightOnlyCudaKernel) const; + private: - std::optional profileTacticsForProblem(int m, int n, int k, std::vector const& tactics); + std::optional profileTacticsForProblem(int m, int n, int k, std::vector const& tactics, + char* workspace, cudaStream_t stream); - float profileTacticForProblem(int m, int n, int k, Config const& tactic); + float profileTacticForProblem(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t stream); int nextPowerOfTwo(int v) const { --v; @@ -206,14 +220,14 @@ class GemmPluginProfiler { private: MNKProfileMapPtr mMNKProfileMap{}; - onnxruntime::IAllocatorUniquePtr mWorkspaceTmp{nullptr}; - - cudaStream_t mStream; - GemmDims mDims{}; bool mSkip{false}; + // Remembered from the initial profileTactics call so lazy single-bucket profiling can + // reproduce the same tactic candidate set. + bool mHasWeightOnlyCudaKernel{false}; + onnxruntime::AllocatorPtr mAllocator; }; @@ -277,8 +291,10 @@ void GemmPluginProfiler::profileT mRunner = runner; mType = type; + mDims = dims; + mHasWeightOnlyCudaKernel = hasWeightOnlyCudaKernel; - int const maxM = std::min(nextPowerOfTwo(dims.maxM), getMaxProfileM()); + int const maxM = std::min(nextPowerOfTwo(static_cast(dims.maxM)), getMaxProfileM()); size_t workspace_bytes = computeTmpSize(maxM, dims.n, dims.k); @@ -293,11 +309,13 @@ void GemmPluginProfiler::profileT auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); bool isAllocated{false}; + onnxruntime::IAllocatorUniquePtr workspace_tmp{nullptr}; + cudaStream_t stream; auto profileTactics = [&](int m, int n, int k) { if (mProfileMap->count(m) == 0) { if (!isAllocated) { - this->mWorkspaceTmp = onnxruntime::IAllocator::MakeUniquePtr(mAllocator, workspace_bytes, true); + workspace_tmp = onnxruntime::IAllocator::MakeUniquePtr(mAllocator, workspace_bytes, true); #if ORT_LLM_VERBOSE AllocatorStats stats; this->mAllocator->GetStats(&stats); @@ -307,43 +325,27 @@ void GemmPluginProfiler::profileT isAllocated = true; } - initTmpData(m, n, k, this->mWorkspaceTmp.get(), workspace_bytes, this->mStream); + initTmpData(m, n, k, workspace_tmp.get(), workspace_bytes, stream); auto tactics = this->getTactics(m, n, k); // Profile different tactics for particular m and insert best config to the map - mProfileMap->insert({m, this->profileTacticsForProblem(m, n, k, tactics)}); + mProfileMap->insert({m, this->profileTacticsForProblem(m, n, k, tactics, workspace_tmp.get(), stream)}); } }; - CUDA_CALL_THROW(cudaStreamCreate(&mStream)); - - int const startMinMRounded = nextPowerOfTwo(dims.minM); + CUDA_CALL_THROW(cudaStreamCreate(&stream)); - if (hasWeightOnlyCudaKernel) { - // Profile tactics for finer granularity of M, - // if CUDA kernel is enabled for weight-only plugins - int minM = dims.minM; - for (int m = std::max(1, minM); m < std::min(16, maxM); m += 1) { - profileTactics(m, dims.n, dims.k); - } - - for (int m = 16; m < maxM; m *= 2) { - profileTactics(m, dims.n, dims.k); - } - } else { - // Profile tactics for CUTLASS kernel only - for (int m = std::max(1, startMinMRounded); m < maxM; m *= 2) { - profileTactics(m, dims.n, dims.k); - } + // Profile the (possibly reduced) set of M buckets. Any unprofiled runtime M is handled + // later by lazy single-bucket profiling in getBestConfigOrProfile. + for (int m : getProfileMBuckets(static_cast(dims.minM), maxM, hasWeightOnlyCudaKernel)) { + profileTactics(m, static_cast(dims.n), static_cast(dims.k)); } - profileTactics(maxM, dims.n, dims.k); - if (isAllocated) { // Free tmp data - mWorkspaceTmp.reset(); + workspace_tmp.reset(); } - CUDA_CALL_THROW(cudaStreamDestroy(mStream)); + CUDA_CALL_THROW(cudaStreamDestroy(stream)); } template @@ -372,9 +374,93 @@ std::optional GemmPluginProfiler +std::vector GemmPluginProfiler::getProfileMBuckets( + int minM, int maxM, bool hasWeightOnlyCudaKernel) const { + // Default: reproduce the historical dense sweep so any other users of this template keep + // their behavior. Subclasses may override this to profile a smaller bucket set. + std::vector buckets; + if (hasWeightOnlyCudaKernel) { + for (int m = std::max(1, minM); m < std::min(16, maxM); m += 1) { + buckets.push_back(m); + } + for (int m = 16; m < maxM; m *= 2) { + buckets.push_back(m); + } + } else { + for (int m = std::max(1, nextPowerOfTwo(minM)); m < maxM; m *= 2) { + buckets.push_back(m); + } + } + buckets.push_back(maxM); + return buckets; +} + +template +std::optional GemmPluginProfiler::getBestConfigOrProfile( + int m, GemmIdType const& gemmId) { + if (mSkip) { + return std::nullopt; + } + + int const target = std::min(std::max(1, nextPowerOfTwo(m)), getMaxProfileM()); + + // Fast path: an already-profiled (exact or rounded) bucket under a shared read lock. + { + reader_lock lock(mMNKProfileMap->mutex); + if (mMNKProfileMap->existsMProfileMap(gemmId)) { + auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); + if (mProfileMap->count(m) > 0) { + return mProfileMap->at(m); + } + if (mProfileMap->count(target) > 0) { + return mProfileMap->at(target); + } + } + } + + // We can only profile lazily if the profiling context from construction is available. + if (mRunner == nullptr || mAllocator == nullptr || !mDims.isInitialized()) { + ORT_LLM_LOG_WARNING("Cannot lazily profile an unprofiled M bucket: profiler context is unavailable."); + return std::nullopt; + } + + int const n = static_cast(mDims.n); + int const k = static_cast(mDims.k); + size_t const workspace_bytes = computeTmpSize(target, n, k); + + cudaStream_t stream; + CUDA_CALL_THROW(cudaStreamCreate(&stream)); + auto workspace_tmp = onnxruntime::IAllocator::MakeUniquePtr(mAllocator, workspace_bytes, true); + initTmpData(target, n, k, workspace_tmp.get(), workspace_bytes, stream); + + auto tactics = this->getTactics(target, n, k); + auto best = this->profileTacticsForProblem(target, n, k, tactics, workspace_tmp.get(), stream); + + workspace_tmp.reset(); + CUDA_CALL_THROW(cudaStreamDestroy(stream)); + + writer_lock lock(mMNKProfileMap->mutex); + if (!mMNKProfileMap->existsMProfileMap(gemmId)) { + mMNKProfileMap->createMProfileMap(gemmId); + } + auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); + + if (mProfileMap->count(m) > 0) { + return mProfileMap->at(m); + } + if (mProfileMap->count(target) > 0) { + return mProfileMap->at(target); + } + + mProfileMap->insert({target, best}); + + return best; +} + template std::optional GemmPluginProfiler::profileTacticsForProblem( - int m, int n, int k, std::vector const& tactics) { + int m, int n, int k, std::vector const& tactics, char* workspace, cudaStream_t stream) { ORT_LLM_LOG_ENTRY(); float bestTime = std::numeric_limits::max(); @@ -394,7 +480,7 @@ std::optional GemmPluginProfiler 1 if constexpr (std::is_same_v) { @@ -441,15 +527,13 @@ std::optional GemmPluginProfiler float GemmPluginProfiler::profileTacticForProblem( - int m, int n, int k, Config const& tactic) { + int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t stream) { constexpr int warmup = 5; constexpr int runs = 10; - cudaStream_t stream = mStream; - // Warmup the execution for (int i = 0; i < warmup; ++i) { - runTactic(m, n, k, tactic, mWorkspaceTmp.get(), stream); + runTactic(m, n, k, tactic, workspace, stream); } cudaEvent_t start; @@ -461,7 +545,7 @@ float GemmPluginProfiler::profile // Profile GEMM for (int i = 0; i < runs; ++i) { - runTactic(m, n, k, tactic, mWorkspaceTmp.get(), stream); + runTactic(m, n, k, tactic, workspace, stream); } CUDA_CALL_THROW(cudaEventRecord(stop, stream)); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index abd43dc6f1f6c..2630d76b6e4dc 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -17,6 +17,7 @@ #include "contrib_ops/cuda/llm/fpA_intB_gemm/fpA_intB_gemm.h" #include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" #include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" +#include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" #endif #include "contrib_ops/cuda/llm/common/logger.h" #include "contrib_ops/cpu/quantization/matmul_nbits_helper.h" @@ -109,8 +110,8 @@ void MatMulNBits::InitGemmProfiler(int sm) { } gemmProfiler_->setCudaKernelType(cuda_kernel_type, sm); - gemmProfiler_->setQuant(nbits_, has_bias_, has_zero_points_); - gemmProfiler_->setGroupSize(block_size_); + gemmProfiler_->setQuant(static_cast(nbits_), has_bias_, has_zero_points_); + gemmProfiler_->setGroupSize(static_cast(block_size_)); auto allocator = this->Info().GetAllocator(OrtMemType::OrtMemTypeDefault); gemmProfiler_->setAllocator(allocator); @@ -119,15 +120,15 @@ void MatMulNBits::InitGemmProfiler(int sm) { template void MatMulNBits::RunGemmProfile(bool hasWeightOnlyCudaKernel, int min_m, int max_m) { // Number of 16-bit elements after casting int8/int4 to fp16. - int n_16b = N_ / (nbits_ == 8 ? 2 : 4); + int n_16b = static_cast(N_ / (nbits_ == 8 ? 2 : 4)); // Include the packing/kernel SM in the GEMM id so the SM80-compatibility and native SM90 kernels // (which need different tactics) do not share profiled configs for the same (N, K, dtype). const int kernel_sm = FpAIntBPackingSmForKernel(); if constexpr (std::is_same_v) { - gemmId_ = GemmIdCore(n_16b, K_, onnxruntime::llm::nvinfer::DataType::kHALF, kernel_sm); + gemmId_ = GemmIdCore(n_16b, static_cast(K_), onnxruntime::llm::nvinfer::DataType::kHALF, kernel_sm); } else if constexpr (std::is_same_v) { - gemmId_ = GemmIdCore(n_16b, K_, onnxruntime::llm::nvinfer::DataType::kBF16, kernel_sm); + gemmId_ = GemmIdCore(n_16b, static_cast(K_), onnxruntime::llm::nvinfer::DataType::kBF16, kernel_sm); } GemmDims dims = {min_m, max_m, n_16b, K_}; @@ -205,10 +206,10 @@ Status MatMulNBits::PrePack_B([[maybe_unused]] const Tensor& tensor, if (nbits_ == 4) { // Transpose the weight and add default zero point. onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( - stream, packed_transposed_weight, blob_data, n, k); + stream, packed_transposed_weight, blob_data, static_cast(n), static_cast(k)); } else { onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( - stream, packed_transposed_weight, blob_data, n, k); + stream, packed_transposed_weight, blob_data, static_cast(n), static_cast(k)); } using onnxruntime::llm::kernels::weight_only::QuantType; @@ -251,7 +252,7 @@ Status MatMulNBits::PrePack_Scale([[maybe_unused]] const Tensor& tensor, CudaT* transposed_scales = reinterpret_cast(fpA_intB_scale_buffer_.get()); onnxruntime::llm::kernels::fpA_intB_gemv::launch_transpose_scale_kernel( - stream, reinterpret_cast(tensor.Data()), transposed_scales, n, k_blocks); + stream, reinterpret_cast(tensor.Data()), transposed_scales, static_cast(n), static_cast(k_blocks)); CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); DUMP_TENSOR_INIT(); @@ -287,16 +288,16 @@ Status MatMulNBits::PrePack_ZeroPoint([[maybe_unused]] const Tensor& tensor, if (nbits_ == 4) { onnxruntime::llm::kernels::fpA_intB_gemv::launch_scaled_zero_point_kernel( stream, reinterpret_cast(zero_points_data), - transposed_scales, scaled_zero_points, n, k_blocks, default_zero_point); + transposed_scales, scaled_zero_points, static_cast(n), static_cast(k_blocks), default_zero_point); } else { onnxruntime::llm::kernels::fpA_intB_gemv::launch_scaled_zero_point_kernel( stream, reinterpret_cast(zero_points_data), - transposed_scales, scaled_zero_points, n, k_blocks, default_zero_point); + transposed_scales, scaled_zero_points, static_cast(n), static_cast(k_blocks), default_zero_point); } } else { // zero point is not uint8_t type onnxruntime::llm::kernels::fpA_intB_gemv::launch_scaled_zero_point_kernel( stream, reinterpret_cast(zero_points_data), - transposed_scales, scaled_zero_points, n, k_blocks, default_zero_point); + transposed_scales, scaled_zero_points, static_cast(n), static_cast(k_blocks), default_zero_point); } CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); @@ -371,18 +372,30 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { if constexpr (std::is_same::value || std::is_same::value) { if (has_fpA_intB_gemm_) { // We expect weight/scale/zero_point(optional) inputs are initializers and have been prepacked. - // User could disable it by setting ORT_FPA_INTB_GEMM=0 if those tensors cannot be prepacked (It is rare). + // A non-prepacked node can opt out by setting session config ep.cuda.fpa_intb_gemm=0 (or env + // ORT_FPA_INTB_GEMM=0) if those tensors cannot be prepacked (it is rare). const bool has_fpA_intB_weight = is_prepacked_weight_ || weight_prepacked_ != kMatMulNBitsWeightNotPrepacked; ORT_ENFORCE(has_fpA_intB_weight && is_prepacked_scale_ && (is_prepacked_zero_point_ || !has_zero_points_), "To use fpA_intB_gemm, prepacking must be done on weight, scale and zero point."); const void* fpA_intB_weight = is_prepacked_weight_ ? fpA_intB_weight_buffer_.get() : static_cast(blob_data); - auto const bestTactic = gemmProfiler_->getBestConfig(m, gemmId_); + // During CUDA graph capture we must not lazily profile: profiling launches kernels, records + // and synchronizes events, and allocates/frees scratch, all of which are illegal while the + // compute stream is being captured. Fall back to a lookup of an already-profiled bucket + // (warmup runs before capture populate these); only outside capture do we allow lazy + // single-bucket profiling. Note: a null cudaStream_t is the default stream (a valid capture + // target under per-thread default streams), so query the capture status unconditionally. + const bool stream_is_capturing = onnxruntime::llm::common::isCapturing(stream); + auto const bestTactic = stream_is_capturing ? gemmProfiler_->getBestConfig(m, gemmId_) + : gemmProfiler_->getBestConfigOrProfile(m, gemmId_); if (!bestTactic.has_value()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, - "No valid fpA_intB MatMulNBits tactic for M=", m, - ", N=", n, ", K=", k); + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "No valid fpA_intB MatMulNBits tactic for M=", m, ", N=", n, ", K=", k, + stream_is_capturing + ? ". The M bucket was not profiled before CUDA graph capture; run a warmup inference outside capture first." + : ""); } // Env-gated diagnostics (ORT_FPA_INTB_DEBUG=1): dump the selected tactic, the kernel path @@ -428,7 +441,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { a_data, pre_quant_scale_ptr, fpA_intB_weight, fpA_intB_scale_buffer_.get(), has_zero_points_ ? fpA_intB_zero_buffer_.get() : nullptr, bias_data, out_data, - alpha, m, n, k, block_size_, cuda_kernel_type, apply_alpha_in_advance); + alpha, m, n, k, static_cast(block_size_), cuda_kernel_type, apply_alpha_in_advance); // Launch the GEMV with the arch the weights were PACKED for (FpAIntBPackingSmForKernel), // not the raw device SM. The GEMV interleave layout is arch-dependent: arch in [90,100) @@ -450,7 +463,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { 1.f, out_data, m, n, k, - block_size_, + static_cast(block_size_), *bestTactic, reinterpret_cast(workspace_buffer.get()), workspace_size, @@ -627,7 +640,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { // Chunked dequant+GEMM: scratch buffer is already sized for one chunk. const int64_t chunk_n = chunk_target_rows; - auto* out_data = reinterpret_cast(Y->MutableData()); + auto* chunk_out_data = reinterpret_cast(Y->MutableData()); for (int64_t n_start = 0; n_start < N_; n_start += chunk_n) { const int64_t n_end = std::min(n_start + chunk_n, N_); @@ -673,7 +686,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { reinterpret_cast(a_data), // A [M, K] helper.Lda(transa), &zero, - out_data + n_start, // C[:, n_start] — strided output + chunk_out_data + n_start, // C[:, n_start] — strided output SafeInt(helper.Ldc()), // ldc = N (full output stride) GetDeviceProp(), UseTF32())); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index 9a7bfe2895582..9a4edb988d2fd 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -7,7 +7,10 @@ // pre-packed and block-compacted into int4 // #pragma once +#include +#include #include "core/common/safeint.h" +#include "core/common/string_utils.h" #include "core/providers/cuda/cuda_kernel.h" #include "core/providers/cuda/shared_inc/fpgeneric.h" #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" @@ -34,15 +37,45 @@ using onnxruntime::llm::kernels::weight_only::WeightOnlyGroupwiseQuantGemmPlugin using GemmProfilerPtr = std::shared_ptr; using WeightOnlyGemmRunnerPtr = std::shared_ptr; -// Environment variable to configure fpA_intB_gemm for experiments. Set it to 0 to disable, 1 to eanble all. +// Environment variable to enable/disable the fpA_intB path: unset/0/off to disable, other value to enable. +// This only affects nodes whose weights are NOT prepacked (see the constructor). constexpr const char* kFpAIntBGemmOption = "ORT_FPA_INTB_GEMM"; -constexpr int kFpAIntBGemmOption_All = 0x01; -constexpr int kFpAIntBGemmOption_Gemv = 0x02; -constexpr int kFpAIntBGemmOption_Int4 = 0x04; -constexpr int kFpAIntBGemmOption_Int8 = 0x08; + constexpr int64_t kMatMulNBitsWeightNotPrepacked = 0; constexpr int64_t kMatMulNBitsWeightPrepackedSm80 = 1; constexpr int64_t kMatMulNBitsWeightPrepackedSm90 = 2; + +// Session-option config keys. These are readable by BOTH the built-in CUDA EP and the CUDA plugin +// EP: every kernel is created via KernelRegistryManager::CreateKernel, which injects the +// session-level ConfigOptions, and the plugin CUDA EP wraps a CUDAExecutionProvider that reuses this +// same kernel. Each key overrides its ORT_* environment-variable equivalent (config wins). +// ep.cuda.fpa_intb_gemm <-> ORT_FPA_INTB_GEMM (0/off, 1/on) +// ep.cuda.fpa_intb_profile_m <-> ORT_FPA_INTB_PROFILE_M (initial profile M buckets) +constexpr const char* kConfigFpAIntBGemm = "ep.cuda.fpa_intb_gemm"; +constexpr const char* kConfigFpAIntBProfileM = "ep.cuda.fpa_intb_profile_m"; + +// Resolves a setting from the session config first (per-session, EP-agnostic), then the environment +// variable, else empty. Session config wins so a model/session can override a process-wide env var. +inline std::string ResolveFpAIntBConfigOrEnv(const OpKernelInfo& info, const char* config_key, + const char* env_key) { + const auto from_config = info.GetConfigOptions().GetConfigEntry(config_key); + if (from_config.has_value()) { + return *from_config; + } + return ParseEnvironmentVariableWithDefault(env_key, ""); +} + +// Parses the fpA_intB enable flag. "on" enables the full fpA_intB path (the CUTLASS GEMM and, where +// supported, the GEMV decode kernel; they share one weight layout and cannot be split). Accepts +// (case-insensitive): "" / "0" / "off" -> disabled; otherwise, enabled (a non-zero numeric value +// still enables, for backward compatibility). +inline bool ParseFpAIntBEnabled(const std::string& value) { + const std::string lowered = onnxruntime::utils::GetLowercaseString(onnxruntime::utils::TrimString(value)); + if (lowered.empty() || lowered == "0" || lowered == "off") { + return false; + } + return true; +} #endif template @@ -106,43 +139,55 @@ class MatMulNBits final : public CudaKernel { } if constexpr (std::is_same::value || std::is_same::value) { - int option = ParseEnvironmentVariableWithDefault(kFpAIntBGemmOption, 0); - ORT_ENFORCE(!(weight_prepacked_ != kMatMulNBitsWeightNotPrepacked && option == 0), - "weight_prepacked requires the fpA_intB path, but ORT_FPA_INTB_GEMM is off for this node"); + const bool prepacked = weight_prepacked_ != kMatMulNBitsWeightNotPrepacked; + // The enable flag (session config ep.cuda.fpa_intb_gemm, else ORT_FPA_INTB_GEMM env) only + // chooses the path for weights that are NOT prepacked. A prepacked weight is already stored in + // the fpA_intB layout, so the choice was made at export time and cannot be turned off here. + const bool enable_fpa_intb = + prepacked || + ParseFpAIntBEnabled(ResolveFpAIntBConfigOrEnv(info, kConfigFpAIntBGemm, kFpAIntBGemmOption)); // Note: a fused bias (input[5]) is fully supported by the fpA_intB GEMV, CUTLASS SM80/SM90 // GEMM (EpilogueOpBias), and the tactic profiler, so bias-bearing nodes (e.g. gpt-oss // qkv_proj/o_proj) are eligible. Only g_idx/reorder remains unsupported by this path. - if ((option & (static_cast(nbits_) | kFpAIntBGemmOption_All)) != 0 && + if (enable_fpa_intb && (block_size_ == 32 || block_size_ == 64 || block_size_ == 128) && (nbits_ == 4 || nbits_ == 8) && !has_g_idx_ && N_ % (nbits_ == 8 ? 32 : 64) == 0 && K_ % block_size_ == 0 && sm_ >= 75) { - if ((option & (kFpAIntBGemmOption_Gemv | kFpAIntBGemmOption_All)) != 0) { - using onnxruntime::llm::kernels::fpA_intB_gemv::KernelType; - KernelType cuda_kernel_type; - if constexpr (std::is_same::value) { - cuda_kernel_type = (nbits_ == 8) ? KernelType::FP16Int8Groupwise : KernelType::FP16Int4Groupwise; - } else if constexpr (std::is_same::value) { - cuda_kernel_type = (nbits_ == 8) ? KernelType::BF16Int8Groupwise : KernelType::BF16Int4Groupwise; - } - - if (onnxruntime::llm::kernels::fpA_intB_gemv::is_supported(sm_, cuda_kernel_type)) { - has_fpA_intB_gemv_ = true; - } + // The CUTLASS GEMM and the GEMV decode kernel consume the same fpA_intB weight layout, so + // enable GEMV whenever it is supported; a node cannot mix fpA_intB and legacy layouts. + using onnxruntime::llm::kernels::fpA_intB_gemv::KernelType; + KernelType cuda_kernel_type; + if constexpr (std::is_same::value) { + cuda_kernel_type = (nbits_ == 8) ? KernelType::FP16Int8Groupwise : KernelType::FP16Int4Groupwise; + } else if constexpr (std::is_same::value) { + cuda_kernel_type = (nbits_ == 8) ? KernelType::BF16Int8Groupwise : KernelType::BF16Int4Groupwise; + } + if (onnxruntime::llm::kernels::fpA_intB_gemv::is_supported(sm_, cuda_kernel_type)) { + has_fpA_intB_gemv_ = true; } InitGemmProfiler(FpAIntBPackingSmForKernel()); - constexpr int max_m = 8291; + // Initial profile M buckets from session config (ep.cuda.fpa_intb_profile_m) with + // ORT_FPA_INTB_PROFILE_M env fallback; empty -> profiler uses its default bucket set. + std::vector profile_m = WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMList( + ResolveFpAIntBConfigOrEnv(info, kConfigFpAIntBProfileM, + onnxruntime::llm::kernels::weight_only::kEnvProfileM)); + gemmProfiler_->setProfileMOverride(profile_m); + + int max_m = profile_m.empty() ? onnxruntime::llm::kernels::weight_only::kDefaultProfileMaxM + : profile_m.back(); RunGemmProfile(has_fpA_intB_gemv_, 1, max_m); has_fpA_intB_gemm_ = true; } - if (weight_prepacked_ != kMatMulNBitsWeightNotPrepacked) { + if (prepacked) { ORT_ENFORCE(has_fpA_intB_gemm_, - "weight_prepacked requires the fpA_intB path, but it is disabled or unsupported for this node"); + "weight_prepacked requires the fpA_intB path, but it is unsupported for this node " + "(check bits, block_size, N/K alignment, g_idx, and compute capability >= 7.5)"); ORT_ENFORCE(weight_prepacked_ == RequiredWeightPrepackedFormat(), "weight_prepacked=", weight_prepacked_, " does not match the format required by the selected fpA_intB kernel: ", RequiredWeightPrepackedFormat()); diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 51895938ce3aa..2927477a66244 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -984,8 +984,13 @@ TEST(MatMulNBits, Fp16_Int4_NoZeroPoint_Bias_Prepacked) { RunTest(opts, std::move(eps)); } -TEST(MatMulNBits, Fp16_Int4_PrepackedWeightRequiresFpAIntBGemm) { - ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_FPA_INTB_GEMM", "0"}}}; +// A prepacked weight (weight_prepacked!=0) forces the fpA_intB path on regardless of the enable +// flag, and the constructor ORT_ENFORCEs that the path is actually supported for the node. Here the +// block_size (256) is outside the fpA_intB-supported set {32, 64, 128}, so kernel construction is +// rejected up front with "weight_prepacked requires the fpA_intB path, but it is unsupported ...", +// even though ORT_FPA_INTB_GEMM is enabled. +TEST(MatMulNBits, Fp16_Int4_PrepackedWeightRejectedWhenFpAIntBUnsupported) { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_FPA_INTB_GEMM", "1"}}}; auto cuda_ep = DefaultCudaExecutionProvider(); if (!cuda_ep) { @@ -994,7 +999,7 @@ TEST(MatMulNBits, Fp16_Int4_PrepackedWeightRequiresFpAIntBGemm) { TestOptions opts{}; opts.M = 1, opts.N = 256, opts.K = 1024; - opts.block_size = 64; + opts.block_size = 256; opts.disable_cpu_ep_fallback = true; opts.weight_prepacked = 1; opts.expected_failure = "weight_prepacked requires"; diff --git a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 9d3a378f6db19..e3a1a44237b6c 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -211,5 +211,112 @@ def test_torch_packer_matches_cuda_oracle(self): self._check(bits, force_arch, n, k) +@unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") +@unittest.skipUnless(hasattr(_pybind, "quantize_matmul_4bits"), "MatMulNBits 4-bit quantizer is unavailable") +class TestFpAIntBConfigKeys(unittest.TestCase): + """Session-config keys ep.cuda.fpa_intb_gemm / ep.cuda.fpa_intb_profile_m. + + These do not need the offline weight packer (pack_weights_for_cuda_mixed_gemm), so they run in + more build configurations than TestMatMulNBitsPrepackedCuda. They cover: the config key enabling + the fpA_intB path (on/off only), session config overriding the ORT_FPA_INTB_GEMM env var, the + profile-M key being accepted, and env-var backward compatibility. + """ + + def setUp(self): + # Make sure no env override leaks in from the process / other tests. + for name in ("ORT_FPA_INTB_GEMM", "ORT_FPA_INTB_PROFILE_M"): + os.environ.pop(name, None) + + def _quantize_weight(self, weight: np.ndarray, bits: int, block_size: int): + k, n = weight.shape + k_blocks = (k + block_size - 1) // block_size + blob_size = block_size * bits // 8 + q_weight = np.zeros((n, k_blocks, blob_size), dtype=np.uint8) + scales = np.zeros((n, k_blocks), dtype=np.float16) + zero_points = np.zeros((n, (k_blocks + 1) // 2), dtype=np.uint8) + _pybind.quantize_matmul_4bits(q_weight, weight, scales, zero_points, block_size, n, k, True) + return q_weight, np.abs(scales) + + def _make_model(self, m, k, n, q_weight, scales, bits, block_size, weight_prepacked=0) -> ModelProto: + node = helper.make_node( + "MatMulNBits", + ["A", "B", "scales"], + ["Y"], + domain="com.microsoft", + K=k, + N=n, + bits=bits, + block_size=block_size, + weight_prepacked=weight_prepacked, + ) + graph = helper.make_graph( + [node], + "fpa_intb_config_keys_test", + [helper.make_tensor_value_info("A", TensorProto.FLOAT16, [m, k])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT16, [m, n])], + initializer=[ + numpy_helper.from_array(q_weight, name="B"), + numpy_helper.from_array(scales, name="scales"), + ], + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ) + model.ir_version = 10 + return model + + def _run(self, model: ModelProto, a: np.ndarray, config: dict[str, str] | None = None) -> np.ndarray: + so = ort.SessionOptions() + for key, value in (config or {}).items(): + so.add_session_config_entry(key, value) + sess = ort.InferenceSession(model.SerializeToString(), so, providers=["CUDAExecutionProvider"]) + return sess.run(None, {"A": a})[0] + + def _make_int4_case(self, m=32, k=256, n=512, block_size=64): + rng = np.random.default_rng(2024) + a = rng.normal(0.0, 0.25, size=(m, k)).astype(np.float16) + weight = rng.normal(0.0, 0.25, size=(k, n)).astype(np.float16) + q_weight, scales = self._quantize_weight(weight, 4, block_size) + model = self._make_model(m, k, n, q_weight, scales, 4, block_size) + return model, a, q_weight, scales + + def test_config_key_enables_fpa_intb(self): + # On fpA_intB-capable hardware (compute capability >= 7.5) the baseline (no config) runs the + # standard dequant path -- for a non-prepacked node the enable flag defaults to disabled -- + # while the config key selects the fpA_intB path; the two paths must stay numerically + # equivalent. On sm < 75 both fall back to the dequant path, so this asserts equivalence + # rather than the switch itself (the prepacked tests force and exercise the fpA_intB kernel). + # Only on/off is accepted. + model, a, _, _ = self._make_int4_case() + ref = self._run(model, a) + for value in ("1", "on", "all", "true"): + out = self._run(model, a, {"ep.cuda.fpa_intb_gemm": value}) + np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2, err_msg=f"value={value}") + + def test_profile_m_config_key_accepted(self): + model, a, _, _ = self._make_int4_case() + ref = self._run(model, a) + out = self._run(model, a, {"ep.cuda.fpa_intb_gemm": "1", "ep.cuda.fpa_intb_profile_m": "1,8,32"}) + np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2) + + def test_session_config_overrides_env(self): + # env var says off, session config says on -> the session config must win. + model, a, _, _ = self._make_int4_case() + ref = self._run(model, a) + with set_env("ORT_FPA_INTB_GEMM", "0"): + out = self._run(model, a, {"ep.cuda.fpa_intb_gemm": "1"}) + np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2) + + def test_env_var_backward_compatible(self): + model, a, _, _ = self._make_int4_case() + ref = self._run(model, a) + # "1" plus a legacy non-zero numeric value (previously a bitmask) both mean "enabled" now. + for value in ("1", "4"): + with set_env("ORT_FPA_INTB_GEMM", value): + out = self._run(model, a) + np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2, err_msg=f"env={value}") + + if __name__ == "__main__": unittest.main() diff --git a/tools/ci_build/github/linux/build_cuda_c_api_package.sh b/tools/ci_build/github/linux/build_cuda_c_api_package.sh index 45f07b754bf29..2b02d03bd24c6 100755 --- a/tools/ci_build/github/linux/build_cuda_c_api_package.sh +++ b/tools/ci_build/github/linux/build_cuda_c_api_package.sh @@ -23,5 +23,5 @@ docker run -e SYSTEM_COLLECTIONURI --rm \ --parallel --nvcc_threads 1 --flash_nvcc_threads 1 \ --use_cuda --cuda_version=$CUDA_VERSION --cuda_home=/usr/local/cuda-$CUDA_VERSION --cudnn_home=/usr/local/cuda-$CUDA_VERSION \ --skip_tests --use_vcpkg --use_vcpkg_ms_internal_asset_cache \ ---cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}' 'onnxruntime_USE_FPA_INTB_GEMM=OFF' \ +--cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}' 'onnxruntime_USE_FPA_INTB_GEMM=ON' \ && cd /build/Release && make install DESTDIR=/build/installed" diff --git a/tools/ci_build/github/linux/build_linux_python_package.sh b/tools/ci_build/github/linux/build_linux_python_package.sh index 803597bd68cb6..c8f2802342827 100755 --- a/tools/ci_build/github/linux/build_linux_python_package.sh +++ b/tools/ci_build/github/linux/build_linux_python_package.sh @@ -6,7 +6,6 @@ set -e -x mkdir -p /build/dist EXTRA_ARG="" -ENABLE_CACHE=false # Put 3.12 at the last because Ubuntu 24.04 use python 3.12 and we will upload the intermediate build files of this # config to Azure DevOps Artifacts and download them to a Ubuntu 24.04 machine to run the tests. PYTHON_EXES=( @@ -88,7 +87,7 @@ if [ "$BUILD_DEVICE" == "GPU" ]; then #Enable CUDA EP. BUILD_ARGS+=("--use_cuda" "--cuda_version=$SHORT_CUDA_VERSION" "--cuda_home=$CUDA_HOME" "--cudnn_home=$CUDA_HOME") BUILD_ARGS+=("--nvcc_threads=1" "--flash_nvcc_threads=1") - BUILD_ARGS+=("--cmake_extra_defines" "CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}" "onnxruntime_USE_FPA_INTB_GEMM=OFF") + BUILD_ARGS+=("--cmake_extra_defines" "CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}" "onnxruntime_USE_FPA_INTB_GEMM=ON") # Enable TRT EP only if TensorRT is installed. if [ -f /usr/include/NvInfer.h ]; then BUILD_ARGS+=("--use_tensorrt" "--tensorrt_home=/usr") diff --git a/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh b/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh index da02024e9eb1a..bd78da6555fa9 100755 --- a/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh +++ b/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh @@ -23,5 +23,5 @@ docker run -e SYSTEM_COLLECTIONURI --rm --volume /data/onnx:/data/onnx:ro --volu --skip_submodule_sync --use_binskim_compliant_compile_flags --build_shared_lib --build_java --build_nodejs \ --parallel --nvcc_threads 1 --flash_nvcc_threads 1 \ --use_tensorrt --cuda_version=$CUDA_VERSION --cuda_home=/usr/local/cuda-$CUDA_VERSION --cudnn_home=/usr --tensorrt_home=/usr \ - --cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}' 'onnxruntime_USE_FPA_INTB_GEMM=OFF' \ + --cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}' 'onnxruntime_USE_FPA_INTB_GEMM=ON' \ --use_vcpkg --use_vcpkg_ms_internal_asset_cache && cd /build/Release && make install DESTDIR=/build/installed" From e79aa8faf7c7d6a05a9478c5f0f783ec88612d05 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Jul 2026 05:50:30 -0700 Subject: [PATCH 11/11] [Build] Use GPU pool to unblock CI temporarily (#29731) Builds are blocked by onnxruntime-github-vs2022-latest. Try unblock our limited PRs for release. Co-authored-by: GitHub Copilot --- .github/workflows/publish-csharp-apidocs.yml | 2 +- .github/workflows/windows_cuda.yml | 2 +- .github/workflows/windows_cuda_no_cudnn.yml | 2 +- .github/workflows/windows_cuda_plugin.yml | 2 +- .github/workflows/windows_openvino.yml | 2 +- .github/workflows/windows_qnn_x64.yml | 2 +- .github/workflows/windows_tensorrt.yml | 2 +- .github/workflows/windows_x64_debug_build_x64_debug.yml | 2 +- .github/workflows/windows_x64_release_build_x64_release.yml | 2 +- ...generic_interface_build_x64_release_ep_generic_interface.yml | 2 +- .../workflows/windows_x64_release_vitisai_build_x64_release.yml | 2 +- .github/workflows/windows_x64_release_xnnpack.yml | 2 +- .github/workflows/windows_x86.yml | 2 +- 13 files changed, 13 insertions(+), 13 deletions(-) diff --git a/.github/workflows/publish-csharp-apidocs.yml b/.github/workflows/publish-csharp-apidocs.yml index ac9d9a88a2144..4bcfec8d65719 100644 --- a/.github/workflows/publish-csharp-apidocs.yml +++ b/.github/workflows/publish-csharp-apidocs.yml @@ -27,7 +27,7 @@ jobs: build: runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=publish-csharp-apidocs-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] env: diff --git a/.github/workflows/windows_cuda.yml b/.github/workflows/windows_cuda.yml index 9776c22fedabb..4edaffedd2727 100644 --- a/.github/workflows/windows_cuda.yml +++ b/.github/workflows/windows_cuda.yml @@ -21,7 +21,7 @@ jobs: name: Windows GPU CUDA CI Pipeline runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=windows-cuda-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] steps: diff --git a/.github/workflows/windows_cuda_no_cudnn.yml b/.github/workflows/windows_cuda_no_cudnn.yml index 0f2103cc80dd5..b5d81b3c412ee 100644 --- a/.github/workflows/windows_cuda_no_cudnn.yml +++ b/.github/workflows/windows_cuda_no_cudnn.yml @@ -24,7 +24,7 @@ jobs: name: Windows CUDA Plugin EP Build without cuDNN runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=windows-cuda-plugin-no-cudnn-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] steps: diff --git a/.github/workflows/windows_cuda_plugin.yml b/.github/workflows/windows_cuda_plugin.yml index 1e700cc04dab4..1343325dbeb75 100644 --- a/.github/workflows/windows_cuda_plugin.yml +++ b/.github/workflows/windows_cuda_plugin.yml @@ -20,7 +20,7 @@ jobs: name: Windows CUDA Plugin EP Build runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=windows-cuda-plugin-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] steps: diff --git a/.github/workflows/windows_openvino.yml b/.github/workflows/windows_openvino.yml index 36cc1aebad8af..da495814ed8b8 100644 --- a/.github/workflows/windows_openvino.yml +++ b/.github/workflows/windows_openvino.yml @@ -20,7 +20,7 @@ jobs: name: Windows OpenVINO CI Pipeline runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=BUILD_OPENVINO_EP-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] timeout-minutes: 240 diff --git a/.github/workflows/windows_qnn_x64.yml b/.github/workflows/windows_qnn_x64.yml index 5560de89040de..7d4ca871d8203 100644 --- a/.github/workflows/windows_qnn_x64.yml +++ b/.github/workflows/windows_qnn_x64.yml @@ -20,7 +20,7 @@ jobs: name: Windows x64 QNN CI Pipeline (${{ matrix.QnnLibKind }}) runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=build_test_qnn_ep-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}-${{matrix.QnnLibKind}}" ] timeout-minutes: 120 diff --git a/.github/workflows/windows_tensorrt.yml b/.github/workflows/windows_tensorrt.yml index 3ad0076de6d52..f557b978b5e75 100644 --- a/.github/workflows/windows_tensorrt.yml +++ b/.github/workflows/windows_tensorrt.yml @@ -21,7 +21,7 @@ jobs: name: Windows GPU TensorRT CI Pipeline runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=windows-tensorrt-build-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] steps: diff --git a/.github/workflows/windows_x64_debug_build_x64_debug.yml b/.github/workflows/windows_x64_debug_build_x64_debug.yml index 134326b8f57de..9da5e19f77174 100644 --- a/.github/workflows/windows_x64_debug_build_x64_debug.yml +++ b/.github/workflows/windows_x64_debug_build_x64_debug.yml @@ -15,7 +15,7 @@ jobs: build_x64_debug: runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=windows-x64-debug-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] timeout-minutes: 300 diff --git a/.github/workflows/windows_x64_release_build_x64_release.yml b/.github/workflows/windows_x64_release_build_x64_release.yml index 7acbdcd404ad3..3185b6b821eed 100644 --- a/.github/workflows/windows_x64_release_build_x64_release.yml +++ b/.github/workflows/windows_x64_release_build_x64_release.yml @@ -15,7 +15,7 @@ jobs: build_x64_release: runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=windows-x64-release-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] timeout-minutes: 300 diff --git a/.github/workflows/windows_x64_release_ep_generic_interface_build_x64_release_ep_generic_interface.yml b/.github/workflows/windows_x64_release_ep_generic_interface_build_x64_release_ep_generic_interface.yml index 8cb5bad664921..d1c000ffab52c 100644 --- a/.github/workflows/windows_x64_release_ep_generic_interface_build_x64_release_ep_generic_interface.yml +++ b/.github/workflows/windows_x64_release_ep_generic_interface_build_x64_release_ep_generic_interface.yml @@ -15,7 +15,7 @@ jobs: build_x64_release_ep_generic_interface: runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=build_x64_release_ep_generic_interface-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] timeout-minutes: 300 diff --git a/.github/workflows/windows_x64_release_vitisai_build_x64_release.yml b/.github/workflows/windows_x64_release_vitisai_build_x64_release.yml index a3200b21cf658..789250059c0a0 100644 --- a/.github/workflows/windows_x64_release_vitisai_build_x64_release.yml +++ b/.github/workflows/windows_x64_release_vitisai_build_x64_release.yml @@ -15,7 +15,7 @@ jobs: build_x64_release_vitisai: runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=build_x64_release_vitisai-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] timeout-minutes: 300 diff --git a/.github/workflows/windows_x64_release_xnnpack.yml b/.github/workflows/windows_x64_release_xnnpack.yml index 1e55bd376a7e5..122b1f75c8864 100644 --- a/.github/workflows/windows_x64_release_xnnpack.yml +++ b/.github/workflows/windows_x64_release_xnnpack.yml @@ -15,7 +15,7 @@ jobs: build_x64_release_xnnpack: runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=build_x64_release_xnnpack-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] timeout-minutes: 300 diff --git a/.github/workflows/windows_x86.yml b/.github/workflows/windows_x86.yml index 897b4f712ffdd..c4e30211a23e0 100644 --- a/.github/workflows/windows_x86.yml +++ b/.github/workflows/windows_x86.yml @@ -15,7 +15,7 @@ jobs: build_x86_release: runs-on: [ "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", + "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", "JobId=windows-x86-release-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] timeout-minutes: 300