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 abe92589aa705..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 @@ -93,15 +105,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 +145,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/.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/.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 diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 09e307e124316..c1a05811b49fc 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -77,7 +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_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) @@ -107,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/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 3a480033e77af..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 @@ -401,7 +402,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() @@ -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 ef4b3c7f6959a..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() @@ -386,7 +391,6 @@ target_link_libraries(onnxruntime_providers_cuda_plugin PRIVATE CUDA::cublas CUDA::cublasLt CUDA::cufft - CUDA::nvrtc CUDA::cuda_driver cudnn_frontend Boost::mp11 diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 3f6976c3c8955..f50baf6f993a6 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) @@ -325,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/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index a09a80a4732ed..6e382826de37d 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) @@ -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. | @@ -309,7 +311,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..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 @@ -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/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/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..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,7 +16,9 @@ */ #pragma once +#include #include +#include #include #ifdef ENABLE_FP8 #include @@ -60,13 +62,28 @@ inline int getMaxSharedMemoryPerBlockOptin() { 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") { - 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/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/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/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/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/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/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 406eee21c059f..e3a1a44237b6c 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,146 @@ 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) + + +@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/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index caff34b1c79d4..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): @@ -1804,6 +1943,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 +2058,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 +2073,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 +2202,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 +2234,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( 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/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..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": @@ -1822,7 +1874,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) < ( 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..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: '52-real;61-real;75-real;86-real;89-real;90-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/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-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 8995d7134013a..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,10 +61,15 @@ parameters: displayName: 'Linux aarch64 docker base image' default: '' -- name: cmake_cuda_archs +- name: cmake_x64_cuda_archs type: string - displayName: 'CMAKE_CUDA_ARCHITECTURES' - default: '52-real;61-real;75-real;86-real;89-real;90-virtual' + 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 @@ -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-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..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,24 +30,41 @@ 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 - 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 + - 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)" 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..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 @@ -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 @@ -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 2db052a15f22e..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=( @@ -71,7 +70,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 @@ -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_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..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 @@ -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 @@ -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"