diff --git a/.github/workflows/labeler.yml b/.github/workflows/labeler.yml index f05ba4a6517c0..08a7573a06552 100644 --- a/.github/workflows/labeler.yml +++ b/.github/workflows/labeler.yml @@ -15,7 +15,7 @@ jobs: ] timeout-minutes: 120 steps: - - uses: github/issue-labeler@c1b0f9f52a63158c4adc09425e858e87b32e9685 # v3.4 + - uses: github/issue-labeler@ccca1e16fcb66762200e616fc29b57b940c1cb67 # v3.5 with: repo-token: "${{ secrets.GITHUB_TOKEN }}" configuration-path: .github/labeler.yml diff --git a/.github/workflows/linux_ci.yml b/.github/workflows/linux_ci.yml index ad2b35e5e96a6..e94c71f146ed7 100644 --- a/.github/workflows/linux_ci.yml +++ b/.github/workflows/linux_ci.yml @@ -2,10 +2,11 @@ # It leverages a reusable workflow (`reusable_linux_build.yml`) to handle the core build and test logic # within Docker containers, ensuring a consistent environment. # -# The workflow consists of five parallel jobs targeting different combinations: +# The workflow consists of parallel jobs targeting different combinations: # - build-linux-x64-debug: Builds/tests Debug config on Linux x64 (AMD CPU pool), enables AddressSanitizer. # - build-linux-x64-release: Builds/tests Release config on Linux x64 (AMD CPU pool), includes wheel/nuget/benchmark flags. # - orttraining-linux-ci-pipeline: Builds/tests Release config with Training enabled (--enable_training) on Linux x64 (AMD CPU pool). +# - build-linux-x64-no-exceptions: Builds MinSizeRel config with exceptions disabled (--disable_exceptions), without unit tests, on Linux x64 (AMD CPU pool). # - build-linux-arm64-debug: Builds/tests Debug config on Linux arm64 (ARM CPU pool); ASan is disabled due to excessive runtime. Includes wheel build. # - build-linux-arm64-release: Builds/tests Release config on Linux arm64 (ARM CPU pool), includes wheel/benchmark flags. # @@ -79,7 +80,7 @@ jobs: dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cpu docker_image_repo: onnxruntimecpubuildpythonx64 # MATERIALIZE=ON needs no explicit pre-build pip install on this leg: the manylinux cp314 - # image already ships onnx==1.22.0 + numpy==2.4.2 (the exact gate pins) via its + # image already ships the pinned onnx + numpy versions via its # requirements.txt (install_deps.sh installs it into cp311-cp314). Intentional, not an # omission -- unlike the QNN/android legs which run on bare agents and pip-install the pins. extra_build_flags: '--use_binskim_compliant_compile_flags --build_wheel --build_nuget --enable_transformers_tool_test --cmake_extra_defines onnxruntime_BUILD_BENCHMARKS=ON --cmake_extra_defines onnxruntime_MATERIALIZE_ONNX_NODE_TESTS=ON' @@ -103,6 +104,22 @@ jobs: secrets: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + build-linux-x64-no-exceptions: + name: Build Linux x64 MinSizeRel (Exceptions Disabled) + uses: ./.github/workflows/reusable_linux_build.yml + with: + pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + build_config: MinSizeRel + architecture: x64 + dockerfile_path: tools/ci_build/github/linux/docker/inference/x86_64/default/cpu/Dockerfile + docker_image_repo: onnxruntimecpubuildcix64 + # TODO: fix the unit test code (unguarded try/catch and gsl::narrow) to build with exceptions disabled, then build and run the tests here. + extra_build_flags: '--use_binskim_compliant_compile_flags --disable_exceptions --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=OFF' + run_tests: false + job_identifier: build-linux-x64-no-exceptions + secrets: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # --- arm64 Builds --- build-linux-arm64-debug: name: Build Linux arm64 Debug diff --git a/.github/workflows/linux_cuda_ci.yml b/.github/workflows/linux_cuda_ci.yml index d85f8c4f3b340..af9257757a0d1 100644 --- a/.github/workflows/linux_cuda_ci.yml +++ b/.github/workflows/linux_cuda_ci.yml @@ -127,6 +127,7 @@ jobs: mode: 'test' # Set mode to test execution_providers: 'cuda' extra_build_flags: >- + --enable_pybind --use_binskim_compliant_compile_flags --cuda_version=13.0 --cuda_home=/usr/local/cuda-13.0 @@ -137,4 +138,4 @@ jobs: --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' + python_path_prefix: 'export ORT_RUN_CUDA_TACTIC_CACHE_TESTS=1 PATH=/opt/python/cp310-cp310/bin:$PATH' diff --git a/.github/workflows/linux_cuda_plugin_ci.yml b/.github/workflows/linux_cuda_plugin_ci.yml index f2eac7508c6cc..3e83e8aa3ee0b 100644 --- a/.github/workflows/linux_cuda_plugin_ci.yml +++ b/.github/workflows/linux_cuda_plugin_ci.yml @@ -136,6 +136,8 @@ jobs: -v ${{ runner.temp }}/Release:/build/Release \ -e NVIDIA_VISIBLE_DEVICES=all \ -e ORT_TEST_CUDA_PLUGIN_EP=1 \ + -e ORT_RUN_CUDA_TACTIC_CACHE_TESTS=1 \ + -e ORT_CUDA_PLUGIN_PATH=/build/Release/Release/libonnxruntime_providers_cuda.so \ ${{ steps.build_docker_image_step.outputs.full-image-name }} \ bash -c " set -ex @@ -144,6 +146,8 @@ jobs: python -m pip install /build/Release/Release/dist/onnxruntime*.whl python -m pip install numpy onnx python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + cd /build/Release/Release + python /onnxruntime_src/onnxruntime/test/python/onnxruntime_test_python_fpa_intb_tune.py cd /onnxruntime_src/onnxruntime/test/python/transformers python test_cuda_plugin_ep.py " diff --git a/.github/workflows/title-only-labeler.yml b/.github/workflows/title-only-labeler.yml index 7a79d468cf7b6..d3a3d91f6bd40 100644 --- a/.github/workflows/title-only-labeler.yml +++ b/.github/workflows/title-only-labeler.yml @@ -15,7 +15,7 @@ jobs: ] timeout-minutes: 120 steps: - - uses: github/issue-labeler@c1b0f9f52a63158c4adc09425e858e87b32e9685 # v3.4 + - uses: github/issue-labeler@ccca1e16fcb66762200e616fc29b57b940c1cb67 # v3.5 with: repo-token: "${{ secrets.GITHUB_TOKEN }}" configuration-path: .github/title-only-labeler.yml diff --git a/ThirdPartyNotices.txt b/ThirdPartyNotices.txt index 75af6ad7e3db5..409183748a915 100644 --- a/ThirdPartyNotices.txt +++ b/ThirdPartyNotices.txt @@ -19,76 +19,6 @@ General Public License. _____ -Intel Math Kernel Library (Intel MKL) - -Intel Simplified Software License (Version April 2018) - -Copyright (c) 2018 Intel Corporation. - -Use and Redistribution. You may use and redistribute the software (the “Software”), without modification, -provided the following conditions are met: - -* Redistributions must reproduce the above copyright notice and the following terms of use in the Software -and in the documentation and/or other materials provided with the distribution. - -* Neither the name of Intel nor the names of its suppliers may be used to endorse or promote products -derived from this Software without specific prior written permission. - -* No reverse engineering, decompilation, or disassembly of this Software is permitted. - -Limited patent license. Intel grants you a world-wide, royalty-free, non-exclusive license under patents it now -or hereafter owns or controls to make, have made, use, import, offer to sell and sell (“Utilize”) this Software, -but solely to the extent that any such patent is necessary to Utilize the Software alone. The patent license -shall not apply to any combinations which include this software. No hardware per se is licensed hereunder. - -Third party and other Intel programs. “Third Party Programs” are the files listed in the “third-party-programs.txt” -text file that is included with the Software and may include Intel programs under separate license terms. -Third Party Programs, even if included with the distribution of the Materials, are governed by -separate license terms and those license terms solely govern your use of those programs. - -DISCLAIMER. THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, -INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS -FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT ARE DISCLAIMED. THIS SOFTWARE IS -NOT INTENDED FOR USE IN SYSTEMS OR APPLICATIONS WHERE FAILURE OF THE SOFTWARE -MAY CAUSE PERSONAL INJURY OR DEATH AND YOU AGREE THAT YOU ARE FULLY RESPONSIBLE FOR ANY -CLAIMS, COSTS, DAMAGES, EXPENSES, AND ATTORNEYS’ FEES ARISING OUT OF ANY SUCH USE, -EVEN IF ANY CLAIM ALLEGES THAT INTEL WAS NEGLIGENT REGARDING THE DESIGN OR MANUFACTURE OF -THE MATERIALS. - -LIMITATION OF LIABILITY. IN NO EVENT WILL INTEL BE LIABLE FOR ANY DIRECT, INDIRECT, -INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT -LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR -PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, -WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY -OF SUCH DAMAGE. YOU AGREE TO INDEMNIFY AND HOLD INTEL HARMLESS AGAINST ANY CLAIMS -AND EXPENSES RESULTING FROM YOUR USE OR UNAUTHORIZED USE OF THE SOFTWARE. - -No support. Intel may make changes to the Software, at any time without notice, and is not obligated to -support, update or provide training for the Software. - -Termination. Intel may terminate your right to use the Software in the event of your breach of this Agreement -and you fail to cure the breach within a reasonable period of time. - -Feedback. Should you provide Intel with comments, modifications, corrections, enhancements or other input -(“Feedback”) related to the Software Intel will be free to use, disclose, reproduce, license or otherwise -distribute or exploit the Feedback in its sole discretion without any obligations or restrictions of any kind, -including without limitation, intellectual property rights or licensing obligations. - -Compliance with laws. You agree to comply with all relevant laws and regulations governing your use, -transfer, import or export (or prohibition thereof) of the Software. - -Governing law. All disputes will be governed by the laws of the United States of America and the State of -Delaware without reference to conflict of law principles and subject to the exclusive jurisdiction of the state or -federal courts sitting in the State of Delaware, and each party agrees that it submits to the personal -jurisdiction and venue of those courts and waives any objections. The United Nations Convention on -Contracts for the International Sale of Goods (1980) is specifically excluded and will not apply to the -Software. - -*Other names and brands may be claimed as the property of others. - -_____ - protocolbuffers/protobuf Copyright 2008 Google Inc. All rights reserved. diff --git a/cgmanifests/webgpu/cgmanifest.webgpu.json b/cgmanifests/webgpu/cgmanifest.webgpu.json index 1f45f3b6a72e6..81044f6eeeb83 100644 --- a/cgmanifests/webgpu/cgmanifest.webgpu.json +++ b/cgmanifests/webgpu/cgmanifest.webgpu.json @@ -6,9 +6,9 @@ "component": { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git", - "tag": "v20260916.214021" + "tag": "v20261013.193955" } }, "comments": "runtime; WebGPU EP root dependency pinned in cmake/deps.txt and patched by cmake/external/onnxruntime_external_deps.cmake.", @@ -38,7 +38,7 @@ "component": { "type": "git", "git": { - "commitHash": "2d78d7ce09549caca58fc28b442bbad3a19da432", + "commitHash": "af5aede14a0b4013be6890148d34e0fd7b1fac69", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/third_party/abseil-cpp" } }, @@ -47,7 +47,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -60,7 +60,7 @@ "component": { "type": "git", "git": { - "commitHash": "398c57e93a383373546bc6ae5ac2052b6155f2ba", + "commitHash": "18a42733eaf8bd7b03567f7b876ca2f2e1e7bc0e", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/third_party/protobuf" } }, @@ -69,7 +69,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -82,7 +82,7 @@ "component": { "type": "git", "git": { - "commitHash": "04fd3caa1e8267e4d95c806cad901181728e1006", + "commitHash": "86f980c731e62ae4eaf383d320449d71687936bf", "repositoryUrl": "https://github.com/KhronosGroup/SPIRV-Headers.git" } }, @@ -91,7 +91,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -101,7 +101,7 @@ "component": { "type": "git", "git": { - "commitHash": "25318da789b94854908a15b4c718145e4e63fc87", + "commitHash": "911d1eaf87fd7763b2365d7e339816fa31ec3f4b", "repositoryUrl": "https://github.com/KhronosGroup/SPIRV-Tools.git" } }, @@ -110,7 +110,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -120,7 +120,7 @@ "component": { "type": "git", "git": { - "commitHash": "ee2ec5fd83dafce291024683b50dc89219333076", + "commitHash": "c46850864f4661461b0f6cb9922c058ffea4915e", "repositoryUrl": "https://github.com/KhronosGroup/Vulkan-Headers.git" } }, @@ -129,7 +129,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -139,7 +139,7 @@ "component": { "type": "git", "git": { - "commitHash": "930a38bce146cf85c5bd7cb00fa33a66c640c0c6", + "commitHash": "90738f6edfa19561a978fc5f97543ba37075ba9d", "repositoryUrl": "https://github.com/KhronosGroup/Vulkan-Utility-Libraries.git" } }, @@ -148,7 +148,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -167,7 +167,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -177,7 +177,7 @@ "component": { "type": "git", "git": { - "commitHash": "13395eebe853e811d274f7bd9b71fb7d9b5a5851", + "commitHash": "456ae730017b9b4bd3371927abc82627fd51feef", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/third_party/zlib" } }, @@ -186,7 +186,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -196,7 +196,7 @@ "component": { "type": "git", "git": { - "commitHash": "0fb0e3e45aeab84781a0c515298b40363f2bf312", + "commitHash": "4dd6a9f723c6ffdb44147b073ea23b6f9a0d5751", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/base/allocator/partition_allocator.git" } }, @@ -205,7 +205,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -215,7 +215,7 @@ "component": { "type": "git", "git": { - "commitHash": "397ffb2264b4650c97b2418ec9a0f07228719317", + "commitHash": "9757d44faa47de7e8d1bb731c30551ec608624f9", "repositoryUrl": "https://github.com/microsoft/DirectXShaderCompiler.git" } }, @@ -224,7 +224,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -243,7 +243,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -253,7 +253,7 @@ "component": { "type": "git", "git": { - "commitHash": "c202b4a9dac30e789ed6e3b2354efa94357a56f3", + "commitHash": "64e87f9fbdaf6f53f6cff8f15e4b06f17cead514", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/buildtools" } }, @@ -263,7 +263,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -273,7 +273,7 @@ "component": { "type": "git", "git": { - "commitHash": "70510081984cfcdb14a15b3e08dfe9776dc7ed37", + "commitHash": "9f796802e5d633c96dbe97e0b6683d5e62594f9f", "repositoryUrl": "https://chromium.googlesource.com/external/github.com/llvm/llvm-project/clang/tools/clang-format.git" } }, @@ -283,7 +283,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -293,7 +293,7 @@ "component": { "type": "git", "git": { - "commitHash": "e5285b2d6e5a51de33b19d41cf0b76cef329cdee", + "commitHash": "e1705face979e29370ff393177db19397782f568", "repositoryUrl": "https://chromium.googlesource.com/chromium/tools/depot_tools.git" } }, @@ -303,7 +303,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -323,7 +323,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -333,7 +333,7 @@ "component": { "type": "git", "git": { - "commitHash": "14024f8f89421ba73a3d28c9b8f9fbd745d92dff", + "commitHash": "09351f6ec00c2b65f2d03585ae6da07098b024dc", "repositoryUrl": "https://chromium.googlesource.com/external/github.com/llvm/llvm-project/libcxxabi.git" } }, @@ -343,7 +343,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -353,7 +353,7 @@ "component": { "type": "git", "git": { - "commitHash": "91b3328166da1db7db15cb49c714d59198e56cab", + "commitHash": "5dbe1703ba9e6669cfcbcfc6fa6085b5115dd7cc", "repositoryUrl": "https://chromium.googlesource.com/external/github.com/llvm/llvm-project/libc.git" } }, @@ -363,7 +363,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -383,7 +383,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -393,7 +393,7 @@ "component": { "type": "git", "git": { - "commitHash": "9eae7c82c57b9e64a85bcf3a246821dffd9d4d5e", + "commitHash": "cd7fbbb1c416a259adde421e3105046c43482a2c", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/build" } }, @@ -403,7 +403,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -413,7 +413,7 @@ "component": { "type": "git", "git": { - "commitHash": "7fb0953a310546e26de3fc0bc1e87b7618998f57", + "commitHash": "ac2cea8c4e7f8dea698c56dc3d9c261d16860896", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/tools/clang" } }, @@ -423,7 +423,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -433,7 +433,7 @@ "component": { "type": "git", "git": { - "commitHash": "0732f7097d066a81e673f90862994f76886fc638", + "commitHash": "1fd0e28956b24e0261f4272a0e4709da3af6d789", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/tools/memory" } }, @@ -443,7 +443,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -463,7 +463,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -483,7 +483,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -493,7 +493,7 @@ "component": { "type": "git", "git": { - "commitHash": "661c2e320f2e8d2bf13eb8acb035156da4c0b187", + "commitHash": "46d00d031ad908563088220923f4d963aba2eb7d", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/tools/mb" } }, @@ -503,7 +503,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -513,7 +513,7 @@ "component": { "type": "git", "git": { - "commitHash": "c643e243d3c79b5c411a4232f264b718cb740c83", + "commitHash": "45289eb985f1c323b05d5f42d1fd0b477e33acb6", "repositoryUrl": "https://chromium.googlesource.com/chromium/src/testing" } }, @@ -523,7 +523,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -533,7 +533,7 @@ "component": { "type": "git", "git": { - "commitHash": "adc3407201e641a73a551d2f1181c684c33eea56", + "commitHash": "9951014982324338ea932dfdba259aeb1cca70f7", "repositoryUrl": "https://chromium.googlesource.com/external/github.com/llvm/llvm-project/compiler-rt/lib/fuzzer.git" } }, @@ -543,7 +543,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -563,7 +563,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -573,7 +573,7 @@ "component": { "type": "git", "git": { - "commitHash": "471aa8aec31a4de79a96b72815180bc9dc71c420", + "commitHash": "9f2dadde7e63acbe93c38de5d2b50c08b1d33ea0", "repositoryUrl": "https://chromium.googlesource.com/catapult.git" } }, @@ -583,7 +583,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -603,7 +603,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -623,7 +623,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -643,7 +643,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -663,7 +663,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -683,7 +683,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -703,7 +703,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -716,7 +716,7 @@ "component": { "type": "git", "git": { - "commitHash": "9fda9a70a68cc860213fd0181f1a1020dee10896", + "commitHash": "9649fea17b77e672ae19687bf0222109b6270e0d", "repositoryUrl": "https://chromium.googlesource.com/angle/angle" } }, @@ -726,7 +726,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -739,7 +739,7 @@ "component": { "type": "git", "git": { - "commitHash": "694585a05946e1ed49b6bd577ca6537cbb57f025", + "commitHash": "0f87bb2d3742c2992fd1a7dcb2f051eebf79299c", "repositoryUrl": "https://swiftshader.googlesource.com/SwiftShader" } }, @@ -749,7 +749,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -759,7 +759,7 @@ "component": { "type": "git", "git": { - "commitHash": "0bd3b32840486c11b8b35ad64028302ec3a4582b", + "commitHash": "f21a172bc64be6f27c5c4c2a68c0b2bf2a07d2af", "repositoryUrl": "https://chromium.googlesource.com/vulkan-deps" } }, @@ -769,7 +769,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -779,7 +779,7 @@ "component": { "type": "git", "git": { - "commitHash": "81f3df51ab1cf14ea35ff35f921241625cfbd7ca", + "commitHash": "8c90c518c2c18af98be121527309baea5b3282f0", "repositoryUrl": "https://github.com/KhronosGroup/glslang.git" } }, @@ -789,7 +789,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -802,7 +802,7 @@ "component": { "type": "git", "git": { - "commitHash": "e146980907071e728176acfb8612641d25aacf09", + "commitHash": "6b45d76280a9525740c38c1796060eadb9267488", "repositoryUrl": "https://github.com/KhronosGroup/Vulkan-Loader.git" } }, @@ -812,7 +812,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -822,7 +822,7 @@ "component": { "type": "git", "git": { - "commitHash": "462d9819e5953e064e1dcdc04d3edc5fc6bc9431", + "commitHash": "a967170a2b8a1ded4dbdf82d533caaec89184efe", "repositoryUrl": "https://github.com/KhronosGroup/Vulkan-Tools.git" } }, @@ -832,7 +832,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -842,7 +842,7 @@ "component": { "type": "git", "git": { - "commitHash": "80c4ae25ee567668585e53a9c98e9fdc41b03456", + "commitHash": "b38b8fd1535d9e48fe8e2759001cc339d4fe4234", "repositoryUrl": "https://github.com/KhronosGroup/Vulkan-ValidationLayers.git" } }, @@ -852,7 +852,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -865,7 +865,7 @@ "component": { "type": "git", "git": { - "commitHash": "1cdd228e34966dd6b95bd203e9f84faba0f371a1", + "commitHash": "6af574a14089ccfee87efe230ebcdd8742859813", "repositoryUrl": "https://github.com/KhronosGroup/OpenGL-Registry.git" } }, @@ -875,7 +875,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -888,7 +888,7 @@ "component": { "type": "git", "git": { - "commitHash": "5961a7fe64cf8a126890ced6f13d69e0a1e1b83e", + "commitHash": "db3425b8246136faccb5e2782b5694960bd6edf1", "repositoryUrl": "https://github.com/KhronosGroup/EGL-Registry.git" } }, @@ -898,7 +898,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -911,7 +911,7 @@ "component": { "type": "git", "git": { - "commitHash": "32cc9f03ad14135b5f6f833dba83e7e80edc256f", + "commitHash": "cc11994e92512d0a3715e9d28876ec48accb971a", "repositoryUrl": "https://github.com/gpuweb/cts.git" } }, @@ -921,7 +921,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -941,7 +941,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -961,7 +961,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -981,7 +981,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -991,7 +991,7 @@ "component": { "type": "git", "git": { - "commitHash": "e95743d3940e0ff3c267ab55ced9ae6120c7d416", + "commitHash": "22507c95ccb9e705eddc7cab8f28aa46b134f5fe", "repositoryUrl": "https://github.com/gpuweb/gpuweb.git" } }, @@ -1001,7 +1001,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -1011,7 +1011,7 @@ "component": { "type": "git", "git": { - "commitHash": "b5ff182caa90e53293f47939716281342c0812ba", + "commitHash": "a1dc77409bb2b8bd80540cbf4b8db219e123e909", "repositoryUrl": "https://github.com/webgpu-native/webgpu-headers.git" } }, @@ -1021,7 +1021,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -1041,7 +1041,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } @@ -1061,7 +1061,7 @@ { "type": "git", "git": { - "commitHash": "202b044ac73a82216bf1e8e706138bb7aedd3a15", + "commitHash": "f63db59722e8c1c0da9a3bcea0e40fb5c817a059", "repositoryUrl": "https://github.com/google/dawn.git" } } diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index d7fa631df9750..95b547c89dff9 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -254,8 +254,7 @@ option(onnxruntime_MINIMAL_BUILD "Exclude as much as possible from the build. Su option(onnxruntime_CLIENT_PACKAGE_BUILD "Enables default settings that are more appropriate for client/on-device workloads." OFF) option(onnxruntime_ENABLE_SESSION_THREADPOOL_CALLBACKS "Enable per-session thread pool work callbacks." OFF) cmake_dependent_option(onnxruntime_DISABLE_RTTI "Disable RTTI" ON "NOT onnxruntime_ENABLE_PYTHON;NOT onnxruntime_USE_CUDA" OFF) -# For now onnxruntime_DISABLE_EXCEPTIONS will only work with onnxruntime_MINIMAL_BUILD, more changes (ONNX, non-CPU EP, ...) are required to run this standalone -cmake_dependent_option(onnxruntime_DISABLE_EXCEPTIONS "Disable exception handling. Requires onnxruntime_MINIMAL_BUILD currently." ON "onnxruntime_MINIMAL_BUILD;NOT onnxruntime_ENABLE_PYTHON" OFF) +cmake_dependent_option(onnxruntime_DISABLE_EXCEPTIONS "Disable exception handling." ${onnxruntime_MINIMAL_BUILD} "NOT onnxruntime_ENABLE_PYTHON" OFF) # Even when onnxruntime_DISABLE_ABSEIL is ON, ONNX Runtime still needs to link to abseil. option(onnxruntime_DISABLE_ABSEIL "Do not use Abseil data structures in ONNX Runtime source code. Redefine Inlined containers to STD containers." OFF) diff --git a/cmake/adjust_global_compile_flags.cmake b/cmake/adjust_global_compile_flags.cmake index 6e548acbce3a0..c9bf3462fc72e 100644 --- a/cmake/adjust_global_compile_flags.cmake +++ b/cmake/adjust_global_compile_flags.cmake @@ -153,13 +153,7 @@ else() endif() endif() -# If this is only enabled in an onnxruntime_ORT_MODEL_FORMAT_ONLY build we don't need ONNX changes -# as we (currently) only pull in data_type_utils.cc/h which doesn't throw if (onnxruntime_DISABLE_EXCEPTIONS) - if (NOT onnxruntime_MINIMAL_BUILD) - message(FATAL_ERROR "onnxruntime_MINIMAL_BUILD required for onnxruntime_DISABLE_EXCEPTIONS") - endif() - if (onnxruntime_ENABLE_PYTHON) # pybind11 highly depends on C++ exceptions. message(FATAL_ERROR "onnxruntime_ENABLE_PYTHON must be disabled for onnxruntime_DISABLE_EXCEPTIONS") @@ -178,7 +172,10 @@ if (onnxruntime_DISABLE_EXCEPTIONS) string(APPEND CMAKE_CXX_FLAGS " /wd4834 /wd4702") add_compile_definitions("_HAS_EXCEPTIONS=0") else() - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-exceptions -fno-unwind-tables -fno-asynchronous-unwind-tables") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-exceptions") + if (onnxruntime_MINIMAL_BUILD) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-unwind-tables -fno-asynchronous-unwind-tables") + endif() endif() endif() diff --git a/cmake/deps.txt b/cmake/deps.txt index 467f2c2d5d1a3..b31831f7eef7e 100644 --- a/cmake/deps.txt +++ b/cmake/deps.txt @@ -34,10 +34,9 @@ microsoft_gsl;https://github.com/microsoft/GSL/archive/refs/tags/v4.2.1.zip;1094 microsoft_wil;https://github.com/microsoft/wil/archive/refs/tags/v1.0.250325.1.zip;826c8bd47c2258ec61b8b218e031e5b33d27f761 mimalloc;https://github.com/microsoft/mimalloc/archive/refs/tags/v2.1.1.zip;d5ee7d34223d0567892db5179849939c8769dc41 mp11;https://github.com/boostorg/mp11/archive/refs/tags/boost-1.82.0.zip;9bc9e01dffb64d9e0773b2e44d2f22c51aace063 -# ONNX pinned to the 1.22.0 final release (tag v1.22.0). rc2->final changed only CI -# workflow yml + VERSION_NUMBER (no operator-schema, opset, or testdata change), so it is -# pure version plumbing. The 7 CI requirements.txt files are pinned to onnx==1.22.0 (now on PyPI). -onnx;https://github.com/onnx/onnx/archive/refs/tags/v1.22.0.zip;2b2cd58ac7a26df5371266149e0c76776330cdf1 +# ONNX 1.23.2 final release (opset 28). +# The 7 CI requirements.txt files are pinned to onnx==1.23.2. +onnx;https://github.com/onnx/onnx/archive/refs/tags/v1.23.2.zip;8846de32ab45e653c2cd651d8d6c4c299fb1583d # Use the latest commit of 10.9-GA onnx_tensorrt;https://github.com/onnx/onnx-tensorrt/archive/d5dce67db7c2e64b07e055571f5ec06f7f254de2.zip;01114d3b67650857281fa50faa2e412130a63b69 protobuf;https://github.com/protocolbuffers/protobuf/releases/download/v33.6/protobuf-33.6.tar.gz;051fee744fcca04dab37eb009d49aa3e9546acda @@ -59,7 +58,7 @@ deep_gemm;https://github.com/deepseek-ai/DeepGEMM/archive/559d79fb6994a58b8a15b4 extensions;https://github.com/microsoft/onnxruntime-extensions/archive/c24b7bab0c12f53da76d0c31b03b9f0f8ec8f3b4.zip;239063aee4946a9af147b473a4c3da78ba7413b4 directx_headers;https://github.com/microsoft/DirectX-Headers/archive/refs/tags/v1.613.1.zip;47653509a3371eabb156360f42faf582f314bf2e cudnn_frontend;https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.27.0.zip;1e4c9a464d3437e388ab0163f3be068dba783c08 -dawn;https://github.com/google/dawn/archive/refs/tags/v20260916.214021.zip;514e914f23a0787213c89eed879f74b5e275690d +dawn;https://github.com/google/dawn/archive/refs/tags/v20261013.193955.zip;f62c128edcc33e0ba6b090d6550db10c4f237a42 dawn_agility_sdk;https://www.nuget.org/api/v2/package/Microsoft.Direct3D.D3D12/1.721.3-preview;fa5f5fc8d0c8c209cfbb530be634960d20595841 kleidiai;https://github.com/ARM-software/kleidiai/archive/refs/tags/v1.30.0.tar.gz;36aebdd5bd28c40b1943eb93498ae42e703510e5 # kleidiai-qmx is pinned to a specific commit as there are no tagged releases. When an appropriate tagged release becomes available, diff --git a/cmake/external/libprotobuf-mutator b/cmake/external/libprotobuf-mutator index 7a2ed51a6b682..4373c81cd7d50 160000 --- a/cmake/external/libprotobuf-mutator +++ b/cmake/external/libprotobuf-mutator @@ -1 +1 @@ -Subproject commit 7a2ed51a6b682a83e345ff49fc4cfd7ca47550db +Subproject commit 4373c81cd7d50a65f7fbd8dc3779c0c3b1dc9d72 diff --git a/cmake/external/onnx b/cmake/external/onnx index 2bb50465112fe..e1ac8e74b620b 160000 --- a/cmake/external/onnx +++ b/cmake/external/onnx @@ -1 +1 @@ -Subproject commit 2bb50465112feca9003e1ed654d77f01ff1415ca +Subproject commit e1ac8e74b620bbdc7ba07f0440806bd93f651a28 diff --git a/cmake/onnxruntime_providers_cuda_plugin.cmake b/cmake/onnxruntime_providers_cuda_plugin.cmake index 465799f0bcc90..87d3a3af13975 100644 --- a/cmake/onnxruntime_providers_cuda_plugin.cmake +++ b/cmake/onnxruntime_providers_cuda_plugin.cmake @@ -427,11 +427,18 @@ if(NOT onnxruntime_DISABLE_CONTRIB_OPS) endif() # LLM OBJECT library: SM75+ (backward compatible with fpA_intB_gemv/gemm which support SM75). - # FP4 QMoE sources are compiled separately at native SM120 below. Keep the broad LLM target - # on virtual SM120 PTX under MSVC to avoid CCCL tcgen05 host-compile failures. + # FP4 QMoE sources are compiled separately at native SM120 below. Under MSVC, keep virtual + # SM120 PTX and add non-"a" native SASS to avoid CCCL tcgen05 host-compile failures. if(_cuda_plugin_llm_srcs) + set(_plugin_llm_extra_compile_options) if(MSVC) onnxruntime_filter_cuda_archs(_plugin_llm_cuda_architectures MIN_SM 75 REPLACE_SM120_REAL_WITH_VIRTUAL) + if("120" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG OR "120a" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG) + # Virtual compute_120 PTX alone needs a driver as new as the toolkit to JIT; add native non-"a" SASS + # (no tcgen05 path) so the ordinary LLM kernels (fpA_intB GEMV/GEMM) run without JIT. + list(APPEND _plugin_llm_extra_compile_options + "$<$:SHELL:-gencode=arch=compute_120,code=sm_120>") + endif() else() onnxruntime_filter_cuda_archs(_plugin_llm_cuda_architectures MIN_SM 75) endif() @@ -441,7 +448,7 @@ if(NOT onnxruntime_DISABLE_CONTRIB_OPS) PARENT onnxruntime_providers_cuda_plugin CUDA_ARCHITECTURES "${_plugin_llm_cuda_architectures}" NVCC_THREADS "${onnxruntime_plugin_nvcc_threads}" - COMPILE_OPTIONS ${_cuda_plugin_shared_compile_options} + COMPILE_OPTIONS ${_cuda_plugin_shared_compile_options} ${_plugin_llm_extra_compile_options} SOURCES ${_cuda_plugin_llm_srcs}) endif() endif() diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 2c8ff7567702c..fb976e6ffbc92 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -348,6 +348,9 @@ if (onnxruntime_USE_CUDA AND NOT WIN32 AND onnxruntime_BUILD_CUDA_QUANT_PREPROCE ${cutlass_SOURCE_DIR}/tools/util/include ) target_compile_definitions(onnxruntime_cuda_quant_preprocess PRIVATE USE_CUDA) + # Offline packing must honor force_arch=90 even when the build excludes SM90 GEMM kernels. + # Otherwise get_arch_for_mixed_gemm_weight_preprocess silently substitutes the SM80 layout. + target_compile_options(onnxruntime_cuda_quant_preprocess PRIVATE "-UEXCLUDE_SM_90") target_link_libraries(onnxruntime_cuda_quant_preprocess PRIVATE onnxruntime_common Boost::mp11 diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index a50658db46486..4e630b5bb45d2 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1120,6 +1120,7 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P if (NOT onnxruntime_DISABLE_CONTRIB_OPS) list(APPEND onnxruntime_test_providers_cuda_plugin_internal_test_src + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/moe_gemm_int2_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp4_to_fp8_kernel_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc" @@ -1153,11 +1154,16 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P if (NOT onnxruntime_DISABLE_CONTRIB_OPS) list(APPEND onnxruntime_providers_cuda_plugin_ut_impl_src + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/gemm_tactic_cache.cc" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels_fp16_uint2.cu" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/moe/qmoe_kernels.cu" ) + if (onnxruntime_USE_FP4_QMOE) + list(APPEND onnxruntime_providers_cuda_plugin_ut_impl_src + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu") + endif() endif() onnxruntime_add_object_library(onnxruntime_providers_cuda_plugin_ut_impl ${onnxruntime_providers_cuda_plugin_ut_impl_src}) @@ -1176,6 +1182,7 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P ORT_USE_EP_API_ADAPTERS=1) target_compile_options(onnxruntime_providers_cuda_plugin_ut_impl PRIVATE ${_cuda_plugin_shared_compile_options} + "$<$:--expt-relaxed-constexpr>" "$<$:-Wno-unused-parameter>" "$<$:SHELL:--threads \"${onnxruntime_plugin_nvcc_threads}\">") if (CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 13.0) @@ -1677,6 +1684,44 @@ block() onnxruntime_apply_test_target_workarounds(${onnxruntime_provider_test_target}) onnxruntime_set_plugin_ep_test_environment(onnxruntime_provider_test) + if (onnxruntime_USE_CUDA AND onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND + NOT onnxruntime_DISABLE_CONTRIB_OPS) + set(fp8_decode_filter "CUDA_EP_Unittest.QMoEFp8Gemv*") + if (onnxruntime_BUILD_CUDA_EP_AS_PLUGIN) + set(fp8_decode_outer_filter "${fp8_decode_filter}") + else() + set(fp8_decode_outer_filter "CUDA_EP_Unittest.All") + endif() + foreach(native_decode 0 1) + add_test(NAME onnxruntime_fp8_native_decode_${native_decode} + COMMAND $ + "--gtest_filter=${fp8_decode_outer_filter}" + WORKING_DIRECTORY $) + set_tests_properties(onnxruntime_fp8_native_decode_${native_decode} PROPERTIES + ENVIRONMENT "GTEST_FILTER=${fp8_decode_filter};ORT_QMOE_FP8_GEMV_NATIVE_DECODE=${native_decode}" + RUN_SERIAL TRUE) + endforeach() + endif() + + if (onnxruntime_USE_CUDA AND onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND + onnxruntime_USE_FP4_QMOE AND NOT onnxruntime_DISABLE_CONTRIB_OPS) + set(nvfp4_scale_filter "CUDA_EP_Unittest.QMoENvfp4Gemv*") + if (onnxruntime_BUILD_CUDA_EP_AS_PLUGIN) + set(nvfp4_scale_outer_filter "${nvfp4_scale_filter}") + else() + set(nvfp4_scale_outer_filter "CUDA_EP_Unittest.All") + endif() + foreach(native_scale 0 1) + add_test(NAME onnxruntime_nvfp4_native_scale_${native_scale} + COMMAND $ + "--gtest_filter=${nvfp4_scale_outer_filter}" + WORKING_DIRECTORY $) + set_tests_properties(onnxruntime_nvfp4_native_scale_${native_scale} PROPERTIES + ENVIRONMENT "GTEST_FILTER=${nvfp4_scale_filter};ORT_NVFP4_GEMV_NATIVE_SCALE=${native_scale}" + RUN_SERIAL TRUE) + endforeach() + endif() + if (onnxruntime_USE_WEBGPU AND onnxruntime_USE_EP_API_ADAPTERS AND NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_REDUCED_OPS_BUILD AND NOT IOS AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten") @@ -1989,13 +2034,13 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) # coverage gap this feature exists to close. # --------------------------------------------------------------------------- if (onnxruntime_MATERIALIZE_ONNX_NODE_TESTS AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten") - # onnx version pin: derive from the archive URL in cmake/deps.txt (single source of truth). - string(REGEX MATCH "v([0-9]+\\.[0-9]+\\.[0-9]+)" _onnx_url_ver "${DEP_URL_onnx}") + # Derive the expected wheel version from the archive URL so this also works with installed ONNX packages. + string(REGEX MATCH "v([0-9]+\\.[0-9]+\\.[0-9]+)" _onnx_url_version "${DEP_URL_onnx}") if(CMAKE_MATCH_1) set(_onnx_pinned_version ${CMAKE_MATCH_1}) else() - message(FATAL_ERROR "Could not parse the pinned onnx version from DEP_URL_onnx='${DEP_URL_onnx}' " - "(expected a vX.Y.Z tag). Fix cmake/deps.txt or this regex.") + message(FATAL_ERROR "Could not parse the pinned ONNX version from DEP_URL_onnx='${DEP_URL_onnx}' " + "(expected a vX.Y.Z tag). Fix cmake/deps.txt or this parser.") endif() # Python interpreter is not guaranteed for static test-only builds (the top-level @@ -2023,27 +2068,10 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) " OR reconfigure with -Donnxruntime_MATERIALIZE_ONNX_NODE_TESTS=OFF (node-test coverage will be dropped).\n" " Details: ${_onnx_err}") endif() - # onnx version gate: HARD FAIL on a genuine mismatch, but RC / pre-release AWARE. - # ONNX's opset-bump workflow ships wheels like 1.23.0rc1 or 1.23.0.dev20240101 whose - # COMPILED opset registry already matches the formal 1.23.0 tag, so we compare on the - # RELEASE BASE (major.minor.micro) rather than the raw string. This mirrors - # materialize_onnx_node_tests.py::_release_base EXACTLY (regex ^(\d+)\.(\d+)\.(\d+), with a - # raw-string fallback when there is no leading X.Y.Z) so the cmake and Python layers agree: - # an rcN/.devN wheel of the pinned tag passes, while a real major/minor/micro mismatch - # (e.g. 1.21.x, or 1.23.0 when pinned at 1.22.0) still FATALs. Both sides are normalized; - # _onnx_pinned_version is already a clean X.Y.Z (parsed from the deps.txt vX.Y.Z tag), so - # normalizing it is a no-op kept only for symmetry with the Python two-sided compare. - string(REGEX MATCH "^[0-9]+\\.[0-9]+\\.[0-9]+" _onnx_ver_base "${_onnx_ver}") - if(_onnx_ver_base STREQUAL "") - set(_onnx_ver_base "${_onnx_ver}") - endif() - string(REGEX MATCH "^[0-9]+\\.[0-9]+\\.[0-9]+" _onnx_pin_base "${_onnx_pinned_version}") - if(_onnx_pin_base STREQUAL "") - set(_onnx_pin_base "${_onnx_pinned_version}") - endif() - if(NOT _onnx_ver_base STREQUAL _onnx_pin_base) + # The source and wheel are both final releases, so require exact version parity. + if(NOT _onnx_ver STREQUAL _onnx_pinned_version) message(FATAL_ERROR - "onnx ${_onnx_ver} (release base ${_onnx_ver_base}) != pinned ${_onnx_pinned_version} " + "onnx ${_onnx_ver} != pinned ${_onnx_pinned_version} " "(cmake/deps.txt). A mismatched wheel bakes the wrong opset/IR into the materialized " "corpus (silent drift).\n" " Fix: pip install onnx==${_onnx_pinned_version}") @@ -2121,6 +2149,11 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) VERBATIM) add_custom_target(onnx_node_tests_materialized ALL DEPENDS ${_materialized_node_root}/.stamp) + # Keep the corpus available for targeted onnx_test_runner builds as well as + # normal ALL builds. The runner is the direct C++ consumer of this artifact. + if(TARGET onnx_test_runner) + add_dependencies(onnx_test_runner onnx_node_tests_materialized) + endif() if (NOT onnxruntime_REDUCED_OPS_BUILD) # First-class ctest over the materialized node corpus (the durable replacement for the @@ -3144,4 +3177,26 @@ if (onnxruntime_BUILD_SHARED_LIB AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten" ) endif() +if(WIN32 AND onnxruntime_USE_1DS_TELEMETRY) + if(NOT GDK_PLATFORM AND NOT CMAKE_SYSTEM_NAME STREQUAL "WindowsStore") + set(telemetry_lifecycle_test_dir "${TEST_SRC_DIR}/platform/windows_telemetry/lifecycle") + onnxruntime_add_shared_library(onnxruntime_telemetry_shutdown_test_library + "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test_library.cc" + "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test_library.def") + target_link_libraries(onnxruntime_telemetry_shutdown_test_library PRIVATE onnxruntime_common mat) + onnxruntime_add_executable(onnxruntime_telemetry_shutdown_test + "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test.cc") + add_dependencies(onnxruntime_telemetry_shutdown_test onnxruntime_telemetry_shutdown_test_library) + foreach(mode process_exit normal_destruction explicit_unload) + add_test(NAME onnxruntime_telemetry_${mode}_test + COMMAND onnxruntime_telemetry_shutdown_test + "$" ${mode} + WORKING_DIRECTORY "$") + set_tests_properties(onnxruntime_telemetry_${mode}_test PROPERTIES TIMEOUT 15) + endforeach() + set_target_properties(onnxruntime_telemetry_shutdown_test_library onnxruntime_telemetry_shutdown_test + PROPERTIES FOLDER "ONNXRuntimeTest") + endif() +endif() + include(onnxruntime_fuzz_test.cmake) diff --git a/cmake/patches/onnx/onnx.patch b/cmake/patches/onnx/onnx.patch index 3f3f0aad774e1..cc0005bdf1ba0 100644 --- a/cmake/patches/onnx/onnx.patch +++ b/cmake/patches/onnx/onnx.patch @@ -1,8 +1,8 @@ diff --git a/CMakeLists.txt b/CMakeLists.txt -index 791e24b..5f861a1 100644 +index 6f1d56949..8a2b78c97 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt -@@ -46,6 +46,7 @@ option(ONNX_DISABLE_EXCEPTIONS "Disable exception handling." OFF) +@@ -53,6 +53,7 @@ option(ONNX_DISABLE_EXCEPTIONS "Disable exception handling." OFF) option(ONNX_DISABLE_STATIC_REGISTRATION "Disable static registration for ONNX operator schemas." OFF) option(ONNX_USE_UNITY_BUILD "Enable Unity (Jumbo) build for" OFF) option(ONNX_PYODIDE_BUILD "Build ONNX for Pyodide/Emscripten (wasm32)." OFF) @@ -10,7 +10,19 @@ index 791e24b..5f861a1 100644 option(ONNX_INSTALL "Install ONNX targets, headers, and CMake config files" ON) if(WIN32) option(ONNX_USE_MSVC_STATIC_RUNTIME "Build with MSVC static runtime. Ignored if CMAKE_MSVC_RUNTIME_LIBRARY is defined." OFF) -@@ -515,7 +516,22 @@ add_dependencies(onnx_proto onnx_proto_gen) +@@ -135,7 +136,10 @@ if(ONNX_DISABLE_EXCEPTIONS) + string(REGEX REPLACE "/EHsc" "/EHs-c-" CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}") + add_definitions(-D_HAS_EXCEPTIONS=0) + else() +- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-exceptions -fno-unwind-tables -fno-asynchronous-unwind-tables") ++ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-exceptions") ++ if(ONNX_MINIMAL_BUILD) ++ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-unwind-tables -fno-asynchronous-unwind-tables") ++ endif() + endif() + endif() + +@@ -562,7 +566,23 @@ add_dependencies(onnx_proto onnx_proto_headers) # Core library. Hand-written sources come from target_sources(onnx ...) in the # subdirectories; the protobuf objects are pulled in from onnx_proto. add_library(onnx $) @@ -22,6 +34,7 @@ index 791e24b..5f861a1 100644 + # the full operator/schema sources via add_subdirectory(onnx). + target_sources(onnx PRIVATE + "${ONNX_ROOT}/onnx/defs/data_type_utils.cc" ++ "${ONNX_ROOT}/onnx/defs/doc_strings.cc" + ) + # Ensure ONNX_ML is treated as ON for minimal build consistency with ORT's file + set(ONNX_ML ON CACHE BOOL "Enable traditional ML API." FORCE) @@ -34,11 +47,37 @@ index 791e24b..5f861a1 100644 add_dependencies(onnx onnx_proto) set_target_properties(onnx PROPERTIES CXX_VISIBILITY_PRESET hidden) add_onnx_global_defines(onnx) +diff --git a/onnx/defs/doc_strings.cc b/onnx/defs/doc_strings.cc +index 976280392..8c22e55ff 100644 +--- a/onnx/defs/doc_strings.cc ++++ b/onnx/defs/doc_strings.cc +@@ -1705,6 +1705,8 @@ const char kDoc_DequantizeLinear_ver24[] = ""; + const char kDoc_RandomNormal_ver1[] = ""; + const char kDoc_Round_ver11[] = ""; + const char kDoc_SpaceToDepth_ver1[] = ""; ++const char kDoc_SpaceToDepth_ver28[] = ""; ++const char kDoc_DepthToSpace_ver28[] = ""; + const char kDoc_InstanceNormalization_ver6[] = ""; + const char kDoc_ThresholdedRelu_ver10[] = ""; + const char kDoc_Acosh_ver9[] = ""; +diff --git a/onnx/defs/nn/defs.cc b/onnx/defs/nn/defs.cc +index 4b58256e6..1a8103e7c 100644 +--- a/onnx/defs/nn/defs.cc ++++ b/onnx/defs/nn/defs.cc +@@ -2692,7 +2692,7 @@ ONNX_OPERATOR_SET_SCHEMA( + // positive value. + axis += input_ndim; + } +- if (axis < 0) { ++ if (axis < 0 || axis >= input_ndim) { + fail_shape_inference( + "Unexpected axis value (", + axis, diff --git a/onnx/defs/nn/old.cc b/onnx/defs/nn/old.cc -index 3887169..8432613 100644 +index cbb1532c0..6597a2ef9 100644 --- a/onnx/defs/nn/old.cc +++ b/onnx/defs/nn/old.cc -@@ -4078,7 +4078,6 @@ ONNX_OPERATOR_SET_SCHEMA( +@@ -4118,7 +4118,6 @@ ONNX_OPERATOR_SET_SCHEMA( GroupNormalization, 18, OpSchema() @@ -46,146 +85,40 @@ index 3887169..8432613 100644 .SetDoc(GroupNormalization_ver18_doc) .Attr("epsilon", "The epsilon value to use to avoid division by zero.", AttributeProto::FLOAT, 1e-5f) .Attr( -diff --git a/onnx/defs/nn/defs.cc b/onnx/defs/nn/defs.cc -index 0c271d5..6d64da7 100644 ---- a/onnx/defs/nn/defs.cc -+++ b/onnx/defs/nn/defs.cc -@@ -2670,7 +2670,7 @@ ONNX_OPERATOR_SET_SCHEMA( - // positive value. - axis += input_ndim; - } -- if (axis < 0) { -+ if (axis < 0 || axis >= input_ndim) { - fail_shape_inference( - "Unexpected axis value (", - axis, -diff --git a/onnx/common/assertions.h b/onnx/common/assertions.h ---- a/onnx/common/assertions.h -+++ b/onnx/common/assertions.h -@@ -27,9 +27,9 @@ - #endif - std::string barf(const char* fmt, ...); - --[[noreturn]] void throw_assert_error(std::string& /*msg*/); -+[[noreturn]] void throw_assert_error(const std::string& /*msg*/); - --[[noreturn]] void throw_tensor_error(std::string& /*msg*/); -+[[noreturn]] void throw_tensor_error(const std::string& /*msg*/); - - } // namespace ONNX_NAMESPACE - -diff --git a/onnx/common/assertions.cc b/onnx/common/assertions.cc ---- a/onnx/common/assertions.cc -+++ b/onnx/common/assertions.cc -@@ -42,11 +42,11 @@ - } - --void throw_assert_error(std::string& msg) { -+void throw_assert_error(const std::string& msg) { - ONNX_THROW_EX(assert_error(msg)); - } - --void throw_tensor_error(std::string& msg) { -+void throw_tensor_error(const std::string& msg) { - ONNX_THROW_EX(tensor_error(msg)); +diff --git a/onnx/defs/printer.cc b/onnx/defs/printer.cc +index 1d43bccb0..b6beed6f3 100644 +--- a/onnx/defs/printer.cc ++++ b/onnx/defs/printer.cc +@@ -9,6 +9,9 @@ + #include + #include + #include ++#include ++#include ++#include + #include + #include + #include +@@ -31,6 +34,14 @@ std::vector ParseRawIntData(const TensorProto& tensor) { + // point, the shortest form that round-trips through the parser. + template + void PrintNumber(std::ostream& os, T value) { ++#if defined(__APPLE__) && defined(__ENVIRONMENT_IPHONE_OS_VERSION_MIN_REQUIRED__) && \ ++ __ENVIRONMENT_IPHONE_OS_VERSION_MIN_REQUIRED__ < 160300 ++ // std::to_chars is unavailable when targeting iOS releases before 16.3. ++ std::ostringstream value_stream; ++ value_stream.imbue(std::locale::classic()); ++ value_stream << std::setprecision(std::numeric_limits::max_digits10) << value; ++ os << value_stream.str(); ++#else + if constexpr (std::is_floating_point_v) { + // Normalize NaN: MSVC's to_chars emits payload forms like "nan(ind)" + // that the parser does not accept. +@@ -42,6 +53,7 @@ void PrintNumber(std::ostream& os, T value) { + std::array buf{}; + const auto res = std::to_chars(buf.data(), buf.data() + buf.size(), value); + os.write(buf.data(), res.ptr - buf.data()); ++#endif } - - } // namespace ONNX_NAMESPACE -diff --git a/onnx/common/tensor.h b/onnx/common/tensor.h ---- a/onnx/common/tensor.h -+++ b/onnx/common/tensor.h -@@ -51,13 +51,13 @@ - /// size_from_dim() cannot handle this case, while elem_num() handles it correctly - int64_t elem_num() const { -- return safe_dim_product(sizes_, [](const char* msg) { throw tensor_error(msg); }); -+ return safe_dim_product(sizes_, [](const char* msg) { throw_tensor_error(msg); }); - } - int64_t size_from_dim(int dim) const { - if (dim < 0) { - dim += static_cast(sizes_.size()); - } - ONNX_ASSERT(dim >= 0 && (size_t)dim < sizes_.size()) -- return safe_dim_product(sizes_.begin() + dim, sizes_.end(), [](const char* msg) { throw tensor_error(msg); }); -+ return safe_dim_product(sizes_.begin() + dim, sizes_.end(), [](const char* msg) { throw_tensor_error(msg); }); - } - - int32_t elem_type() const { -diff --git a/onnx/common/scoped_resource.h b/onnx/common/scoped_resource.h ---- a/onnx/common/scoped_resource.h -+++ b/onnx/common/scoped_resource.h -@@ -25,18 +25,18 @@ - // Usage: --// using ScopedFd = ScopedResource<-1, close_fd>; --// ScopedFd guard(fd); // destructor calls close_fd(fd) -+// using ScopedFd = ScopedResource; -+// ScopedFd guard(fd); // destructor calls FdTraits::close(fd) - // int raw = guard.release(); // relinquish ownership --template -+template - class ScopedResource { -- using T = decltype(Invalid); -+ using T = typename Traits::type; - T val_; - - public: - explicit ScopedResource(T v) : val_(v) {} - ~ScopedResource() { -- if (val_ != Invalid) { -- Close(val_); -+ if (val_ != Traits::invalid()) { -+ Traits::close(val_); - } - } - T get() const { - return val_; -@@ -44,7 +44,7 @@ - } - T release() { - T tmp = val_; -- val_ = Invalid; -+ val_ = Traits::invalid(); - return tmp; - } - ScopedResource(const ScopedResource&) = delete; -@@ -55,20 +55,32 @@ - #ifdef _WIN32 --inline void close_handle(HANDLE h) { -- CloseHandle(h); --} --using ScopedHandle = ScopedResource; -+struct HandleTraits { -+ using type = HANDLE; -+ static HANDLE invalid() { -+ return INVALID_HANDLE_VALUE; -+ } -+ static void close(HANDLE h) { -+ CloseHandle(h); -+ } -+}; -+using ScopedHandle = ScopedResource; - #endif - --inline void close_fd(int fd) { -+struct FdTraits { -+ using type = int; -+ static constexpr int invalid() { -+ return -1; -+ } -+ static void close(int fd) { - #ifdef _WIN32 -- _close(fd); -+ _close(fd); - #else -- close(fd); -+ ::close(fd); - #endif --} --using ScopedFd = ScopedResource<-1, close_fd>; -+ } -+}; -+using ScopedFd = ScopedResource; - - // RAII guard that invokes a callable on destruction (scope exit). - template - class ScopeExit { - F fn_; + + bool IsValidIdentifier(const std::string& str) { diff --git a/cmake/vcpkg-ports/onnx/binskim.patch b/cmake/vcpkg-ports/onnx/binskim.patch index 3f3f0aad774e1..cc0005bdf1ba0 100644 --- a/cmake/vcpkg-ports/onnx/binskim.patch +++ b/cmake/vcpkg-ports/onnx/binskim.patch @@ -1,8 +1,8 @@ diff --git a/CMakeLists.txt b/CMakeLists.txt -index 791e24b..5f861a1 100644 +index 6f1d56949..8a2b78c97 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt -@@ -46,6 +46,7 @@ option(ONNX_DISABLE_EXCEPTIONS "Disable exception handling." OFF) +@@ -53,6 +53,7 @@ option(ONNX_DISABLE_EXCEPTIONS "Disable exception handling." OFF) option(ONNX_DISABLE_STATIC_REGISTRATION "Disable static registration for ONNX operator schemas." OFF) option(ONNX_USE_UNITY_BUILD "Enable Unity (Jumbo) build for" OFF) option(ONNX_PYODIDE_BUILD "Build ONNX for Pyodide/Emscripten (wasm32)." OFF) @@ -10,7 +10,19 @@ index 791e24b..5f861a1 100644 option(ONNX_INSTALL "Install ONNX targets, headers, and CMake config files" ON) if(WIN32) option(ONNX_USE_MSVC_STATIC_RUNTIME "Build with MSVC static runtime. Ignored if CMAKE_MSVC_RUNTIME_LIBRARY is defined." OFF) -@@ -515,7 +516,22 @@ add_dependencies(onnx_proto onnx_proto_gen) +@@ -135,7 +136,10 @@ if(ONNX_DISABLE_EXCEPTIONS) + string(REGEX REPLACE "/EHsc" "/EHs-c-" CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}") + add_definitions(-D_HAS_EXCEPTIONS=0) + else() +- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-exceptions -fno-unwind-tables -fno-asynchronous-unwind-tables") ++ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-exceptions") ++ if(ONNX_MINIMAL_BUILD) ++ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-unwind-tables -fno-asynchronous-unwind-tables") ++ endif() + endif() + endif() + +@@ -562,7 +566,23 @@ add_dependencies(onnx_proto onnx_proto_headers) # Core library. Hand-written sources come from target_sources(onnx ...) in the # subdirectories; the protobuf objects are pulled in from onnx_proto. add_library(onnx $) @@ -22,6 +34,7 @@ index 791e24b..5f861a1 100644 + # the full operator/schema sources via add_subdirectory(onnx). + target_sources(onnx PRIVATE + "${ONNX_ROOT}/onnx/defs/data_type_utils.cc" ++ "${ONNX_ROOT}/onnx/defs/doc_strings.cc" + ) + # Ensure ONNX_ML is treated as ON for minimal build consistency with ORT's file + set(ONNX_ML ON CACHE BOOL "Enable traditional ML API." FORCE) @@ -34,11 +47,37 @@ index 791e24b..5f861a1 100644 add_dependencies(onnx onnx_proto) set_target_properties(onnx PROPERTIES CXX_VISIBILITY_PRESET hidden) add_onnx_global_defines(onnx) +diff --git a/onnx/defs/doc_strings.cc b/onnx/defs/doc_strings.cc +index 976280392..8c22e55ff 100644 +--- a/onnx/defs/doc_strings.cc ++++ b/onnx/defs/doc_strings.cc +@@ -1705,6 +1705,8 @@ const char kDoc_DequantizeLinear_ver24[] = ""; + const char kDoc_RandomNormal_ver1[] = ""; + const char kDoc_Round_ver11[] = ""; + const char kDoc_SpaceToDepth_ver1[] = ""; ++const char kDoc_SpaceToDepth_ver28[] = ""; ++const char kDoc_DepthToSpace_ver28[] = ""; + const char kDoc_InstanceNormalization_ver6[] = ""; + const char kDoc_ThresholdedRelu_ver10[] = ""; + const char kDoc_Acosh_ver9[] = ""; +diff --git a/onnx/defs/nn/defs.cc b/onnx/defs/nn/defs.cc +index 4b58256e6..1a8103e7c 100644 +--- a/onnx/defs/nn/defs.cc ++++ b/onnx/defs/nn/defs.cc +@@ -2692,7 +2692,7 @@ ONNX_OPERATOR_SET_SCHEMA( + // positive value. + axis += input_ndim; + } +- if (axis < 0) { ++ if (axis < 0 || axis >= input_ndim) { + fail_shape_inference( + "Unexpected axis value (", + axis, diff --git a/onnx/defs/nn/old.cc b/onnx/defs/nn/old.cc -index 3887169..8432613 100644 +index cbb1532c0..6597a2ef9 100644 --- a/onnx/defs/nn/old.cc +++ b/onnx/defs/nn/old.cc -@@ -4078,7 +4078,6 @@ ONNX_OPERATOR_SET_SCHEMA( +@@ -4118,7 +4118,6 @@ ONNX_OPERATOR_SET_SCHEMA( GroupNormalization, 18, OpSchema() @@ -46,146 +85,40 @@ index 3887169..8432613 100644 .SetDoc(GroupNormalization_ver18_doc) .Attr("epsilon", "The epsilon value to use to avoid division by zero.", AttributeProto::FLOAT, 1e-5f) .Attr( -diff --git a/onnx/defs/nn/defs.cc b/onnx/defs/nn/defs.cc -index 0c271d5..6d64da7 100644 ---- a/onnx/defs/nn/defs.cc -+++ b/onnx/defs/nn/defs.cc -@@ -2670,7 +2670,7 @@ ONNX_OPERATOR_SET_SCHEMA( - // positive value. - axis += input_ndim; - } -- if (axis < 0) { -+ if (axis < 0 || axis >= input_ndim) { - fail_shape_inference( - "Unexpected axis value (", - axis, -diff --git a/onnx/common/assertions.h b/onnx/common/assertions.h ---- a/onnx/common/assertions.h -+++ b/onnx/common/assertions.h -@@ -27,9 +27,9 @@ - #endif - std::string barf(const char* fmt, ...); - --[[noreturn]] void throw_assert_error(std::string& /*msg*/); -+[[noreturn]] void throw_assert_error(const std::string& /*msg*/); - --[[noreturn]] void throw_tensor_error(std::string& /*msg*/); -+[[noreturn]] void throw_tensor_error(const std::string& /*msg*/); - - } // namespace ONNX_NAMESPACE - -diff --git a/onnx/common/assertions.cc b/onnx/common/assertions.cc ---- a/onnx/common/assertions.cc -+++ b/onnx/common/assertions.cc -@@ -42,11 +42,11 @@ - } - --void throw_assert_error(std::string& msg) { -+void throw_assert_error(const std::string& msg) { - ONNX_THROW_EX(assert_error(msg)); - } - --void throw_tensor_error(std::string& msg) { -+void throw_tensor_error(const std::string& msg) { - ONNX_THROW_EX(tensor_error(msg)); +diff --git a/onnx/defs/printer.cc b/onnx/defs/printer.cc +index 1d43bccb0..b6beed6f3 100644 +--- a/onnx/defs/printer.cc ++++ b/onnx/defs/printer.cc +@@ -9,6 +9,9 @@ + #include + #include + #include ++#include ++#include ++#include + #include + #include + #include +@@ -31,6 +34,14 @@ std::vector ParseRawIntData(const TensorProto& tensor) { + // point, the shortest form that round-trips through the parser. + template + void PrintNumber(std::ostream& os, T value) { ++#if defined(__APPLE__) && defined(__ENVIRONMENT_IPHONE_OS_VERSION_MIN_REQUIRED__) && \ ++ __ENVIRONMENT_IPHONE_OS_VERSION_MIN_REQUIRED__ < 160300 ++ // std::to_chars is unavailable when targeting iOS releases before 16.3. ++ std::ostringstream value_stream; ++ value_stream.imbue(std::locale::classic()); ++ value_stream << std::setprecision(std::numeric_limits::max_digits10) << value; ++ os << value_stream.str(); ++#else + if constexpr (std::is_floating_point_v) { + // Normalize NaN: MSVC's to_chars emits payload forms like "nan(ind)" + // that the parser does not accept. +@@ -42,6 +53,7 @@ void PrintNumber(std::ostream& os, T value) { + std::array buf{}; + const auto res = std::to_chars(buf.data(), buf.data() + buf.size(), value); + os.write(buf.data(), res.ptr - buf.data()); ++#endif } - - } // namespace ONNX_NAMESPACE -diff --git a/onnx/common/tensor.h b/onnx/common/tensor.h ---- a/onnx/common/tensor.h -+++ b/onnx/common/tensor.h -@@ -51,13 +51,13 @@ - /// size_from_dim() cannot handle this case, while elem_num() handles it correctly - int64_t elem_num() const { -- return safe_dim_product(sizes_, [](const char* msg) { throw tensor_error(msg); }); -+ return safe_dim_product(sizes_, [](const char* msg) { throw_tensor_error(msg); }); - } - int64_t size_from_dim(int dim) const { - if (dim < 0) { - dim += static_cast(sizes_.size()); - } - ONNX_ASSERT(dim >= 0 && (size_t)dim < sizes_.size()) -- return safe_dim_product(sizes_.begin() + dim, sizes_.end(), [](const char* msg) { throw tensor_error(msg); }); -+ return safe_dim_product(sizes_.begin() + dim, sizes_.end(), [](const char* msg) { throw_tensor_error(msg); }); - } - - int32_t elem_type() const { -diff --git a/onnx/common/scoped_resource.h b/onnx/common/scoped_resource.h ---- a/onnx/common/scoped_resource.h -+++ b/onnx/common/scoped_resource.h -@@ -25,18 +25,18 @@ - // Usage: --// using ScopedFd = ScopedResource<-1, close_fd>; --// ScopedFd guard(fd); // destructor calls close_fd(fd) -+// using ScopedFd = ScopedResource; -+// ScopedFd guard(fd); // destructor calls FdTraits::close(fd) - // int raw = guard.release(); // relinquish ownership --template -+template - class ScopedResource { -- using T = decltype(Invalid); -+ using T = typename Traits::type; - T val_; - - public: - explicit ScopedResource(T v) : val_(v) {} - ~ScopedResource() { -- if (val_ != Invalid) { -- Close(val_); -+ if (val_ != Traits::invalid()) { -+ Traits::close(val_); - } - } - T get() const { - return val_; -@@ -44,7 +44,7 @@ - } - T release() { - T tmp = val_; -- val_ = Invalid; -+ val_ = Traits::invalid(); - return tmp; - } - ScopedResource(const ScopedResource&) = delete; -@@ -55,20 +55,32 @@ - #ifdef _WIN32 --inline void close_handle(HANDLE h) { -- CloseHandle(h); --} --using ScopedHandle = ScopedResource; -+struct HandleTraits { -+ using type = HANDLE; -+ static HANDLE invalid() { -+ return INVALID_HANDLE_VALUE; -+ } -+ static void close(HANDLE h) { -+ CloseHandle(h); -+ } -+}; -+using ScopedHandle = ScopedResource; - #endif - --inline void close_fd(int fd) { -+struct FdTraits { -+ using type = int; -+ static constexpr int invalid() { -+ return -1; -+ } -+ static void close(int fd) { - #ifdef _WIN32 -- _close(fd); -+ _close(fd); - #else -- close(fd); -+ ::close(fd); - #endif --} --using ScopedFd = ScopedResource<-1, close_fd>; -+ } -+}; -+using ScopedFd = ScopedResource; - - // RAII guard that invokes a callable on destruction (scope exit). - template - class ScopeExit { - F fn_; + + bool IsValidIdentifier(const std::string& str) { diff --git a/cmake/vcpkg-ports/onnx/fix-dependency-protobuf.patch b/cmake/vcpkg-ports/onnx/fix-dependency-protobuf.patch index 4a438def6cf1b..12e1a7a662a25 100644 --- a/cmake/vcpkg-ports/onnx/fix-dependency-protobuf.patch +++ b/cmake/vcpkg-ports/onnx/fix-dependency-protobuf.patch @@ -2,14 +2,14 @@ diff --git a/CMakeLists.txt b/CMakeLists.txt index 4a7df7d71..257d77e61 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt -@@ -159,6 +159,7 @@ if(NOT DEFINED CMAKE_POSITION_INDEPENDENT_CODE) +@@ -198,6 +198,7 @@ if(NOT DEFINED CMAKE_POSITION_INDEPENDENT_CODE) set(CMAKE_POSITION_INDEPENDENT_CODE ON) endif() +find_package(protobuf CONFIG REQUIRED) list(APPEND CMAKE_MODULE_PATH ${ONNX_ROOT}/cmake/external) - if(NOT ONNX_BUILD_CUSTOM_PROTOBUF) - if((ONNX_USE_LITE_PROTO AND TARGET protobuf::libprotobuf-lite) OR ((NOT ONNX_USE_LITE_PROTO) AND TARGET protobuf::libprotobuf)) + if(ONNX_USE_ASAN OR ONNX_USE_TSAN OR ONNX_USE_TYPESAN) + find_package(Sanitizer REQUIRED) diff --git a/cmake/ONNXConfig.cmake.in b/cmake/ONNXConfig.cmake.in index aa1ef47d9..d450b51e5 100644 --- a/cmake/ONNXConfig.cmake.in diff --git a/cmake/vcpkg-ports/onnx/portfile.cmake b/cmake/vcpkg-ports/onnx/portfile.cmake index 7a266e1a07663..c1391ed9a5aed 100644 --- a/cmake/vcpkg-ports/onnx/portfile.cmake +++ b/cmake/vcpkg-ports/onnx/portfile.cmake @@ -3,8 +3,8 @@ vcpkg_check_linkage(ONLY_STATIC_LIBRARY) vcpkg_from_github( OUT_SOURCE_PATH SOURCE_PATH REPO onnx/onnx - REF "v1.22.0" - SHA512 13fafff073a8e0bcf67fd06195da979c3c6273b3f5fdd25f5e5c39ad6af20eefe107f13dc73d52b6021e986d87401e46121e00bba8fbd09f098334e34f78462f + REF "v1.23.2" + SHA512 fc93f19acc87fa0f4b75897c4f7c7e3e14d29f078f6e2b1a2806c5e2454789abdb44067437bc073d612fb48a564d3e5616fbe5e254d021ccec3e301170b1781d PATCHES fix-cmakelists.patch fix-dependency-protobuf.patch diff --git a/cmake/vcpkg-ports/onnx/vcpkg.json b/cmake/vcpkg-ports/onnx/vcpkg.json index c399d4c1ab087..888a94bfc7f40 100644 --- a/cmake/vcpkg-ports/onnx/vcpkg.json +++ b/cmake/vcpkg-ports/onnx/vcpkg.json @@ -1,6 +1,6 @@ { "name": "onnx", - "version-semver": "1.22.0", + "version-semver": "1.23.2", "port-version": 0, "description": "Open standard for machine learning interoperability", "homepage": "https://onnx.ai", diff --git a/csharp/test/Microsoft.ML.OnnxRuntime.Tests.NetCoreApp/InferenceTest.netcore.cs b/csharp/test/Microsoft.ML.OnnxRuntime.Tests.NetCoreApp/InferenceTest.netcore.cs index c0475bb6102c1..14d4fffd55f56 100644 --- a/csharp/test/Microsoft.ML.OnnxRuntime.Tests.NetCoreApp/InferenceTest.netcore.cs +++ b/csharp/test/Microsoft.ML.OnnxRuntime.Tests.NetCoreApp/InferenceTest.netcore.cs @@ -491,9 +491,6 @@ private static Dictionary GetSkippedModels(DirectoryInfo modelsD { "test_mul_uint8", "Could not find an implementation for Mul(14) node with name" }, - { "test_bitshift_right_uint16", "Could not find an implementation for BitShift(11) nodeMeta with name ''"}, - { "test_bitshift_left_uint16", "Could not find an implementation for BitShift(11)"}, - { "test_pow_types_float32_uint64", "Could not find an implementation for Pow(15) node with name ''"}, { "test_pow_types_float32_uint32", "Could not find an implementation for Pow(15) node with name ''"}, diff --git a/docs/C_API_Guidelines.md b/docs/C_API_Guidelines.md index 3416d0a50455d..9b2fac933933d 100644 --- a/docs/C_API_Guidelines.md +++ b/docs/C_API_Guidelines.md @@ -81,3 +81,5 @@ Use types that fall into established patterns. For example, we use int64_t for d ### 9. Adding a new API Follow these guidelines and instructions in the source code. "Rules on how to add a new Ort API version" in [onnxruntime_c_api.cc](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/session/onnxruntime_c_api.cc). + +If a new API needs to be proved out before it's ready to add as a stable API, consider adding it as an experimental C API first. See the [experimental C API guidelines](./Experimental_C_API_Guidelines.md) for details. diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 567b727336196..2703de7cc1fae 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -42,6 +42,7 @@ Do not modify directly.* * com.microsoft.FusedMatMulActivation * com.microsoft.GatedAdd * com.microsoft.GatedDeltaNet + * com.microsoft.GatedDeltaNetStateReplay * com.microsoft.GatedRMSNorm * com.microsoft.GatedRelativePositionBias * com.microsoft.GatherBlockQuantized @@ -125,6 +126,7 @@ Do not modify directly.* * com.microsoft.SkipSimplifiedLayerNormalization * com.microsoft.Snpe * com.microsoft.SparseAttention + * com.microsoft.SparseAttentionIndexer * com.microsoft.SparsePagedAttention * com.microsoft.SparseToDenseMatMul * com.microsoft.Tokenizer @@ -900,9 +902,9 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
epsilon : float (default is 1e-05)
+
epsilon : float
Epsilon added before reciprocal square root.
-
num_branches : int (default is 0)
+
num_branches : int
Number of branches for flattened input. Omit or set to zero for grouped input.
@@ -925,10 +927,10 @@ This version of the operator has been available since version 1 of the 'com.micr #### Type Constraints
-
M : tensor(float), tensor(float16), tensor(bfloat16)
-
Constrain scale to floating-point tensors.
T : tensor(float), tensor(float16), tensor(bfloat16)
Constrain input and output to floating-point tensors.
+
M : tensor(float), tensor(float16), tensor(bfloat16)
+
Constrain scale to floating-point tensors.
@@ -1701,21 +1703,21 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.DynamicSparseAttention** Model-neutral sparse grouped-query attention with a contiguous main KV cache. - + `selected_indices` and `selected_counts` are produced by an external selector. The operator never computes selection scores or TopK indices. The first `selected_counts[q]` entries in each selected-index row are valid; remaining entries must be -1. Valid entries must be unique, non-negative request-local positions in the selected source. - + `attention_mode="selected_only"` attends only to selected entries. `attention_mode="local_plus_selected"` jointly normalizes causally valid main-cache entries in `local_window_size` and selected auxiliary entries. An optional per-query-head sink can participate in either mode; it contributes to the shared softmax denominator but has no value vector. A row with no entries and no sink produces zero output. - + Supported mode/source combinations: - + - `selected_only` + `main`: Qwen4-Exp QSA-compatible execution. - `local_plus_selected` + `auxiliary`: DeepSeek V4 CSA-compatible execution and requires `local_window_size > 0`. - + The main cache uses BNSH layout and ordinary contiguous append semantics. Selection changes reads, not cache writes. Auxiliary K/V use BNSH layout and are read-only. When `auxiliary_kv_shared=1`, auxiliary_value may be omitted and auxiliary_key is used as both K and V. DeepSeek-style post-attention inverse/conjugate RoPE and output projection remain @@ -2372,7 +2374,7 @@ This version of the operator has been available since version 1 of the 'com.micr Alternatively, key and value may be omitted and query contains packed QKV with shape `[total_tokens, 2 * num_heads_q * head_size_qk + num_heads_v * head_size_v]`. `initial_state` is required for packed QKV to determine the head counts and head sizes. - + The leading token axis may instead be spelled as an explicit `[batch_size, sequence_length]` pair, making query/key/value (and the output) rank 4 and decay/beta rank 3. The memory layout is identical; the rank-4 spelling exists so an exporter can round-trip a `[B, S, H*D]` @@ -2506,6 +2508,78 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.GatedDeltaNetStateReplay** + + Replay one request's captured GatedDeltaNet transitions into an inactive recurrent-state row. + The WebGPU implementation uses FP32 backing tensors and CPU metadata; it does not recompute + normalization, gates, beta, or retrieval. It waits for checked GPU completion before returning. + + `metadata` is int64 [11], ordered as: + source_offset, destination_offset, decay_offset, key_offset, delta_offset, + Hv, Dv, Dk, Hk, capacity, kept_count. + Offsets are FP32-element offsets into the full backing tensors, not storage-binding byte offsets. + All offsets are nonnegative; dimensions and capacity are positive; 1 <= kept_count <= capacity. + The caller must ensure that all kept transitions were actually captured. + + The source_state selected by source_offset must be the state snapshot from immediately before + the captured GatedDeltaNet transitions began. Using the producer's final state or another + same-shaped snapshot is invalid because captured delta is state-dependent. + The capsule geometry (Hv, Dv, Dk, Hk) and capacity must match the producer that generated it. + capacity is the original capsule capture/layout capacity. kept_count only selects the leading + transitions to replay; kept_count must NOT be substituted for capacity when deriving the + decay/key/delta section offsets. + + State rows have layout [Hv, Dv, Dk]. The capsule contains decay [capacity, Hv], + key [capacity, Hk, Dk], and delta [capacity, Hv, Dv] starting at their respective offsets. + For every state element (h,v,i): + hk = floor(h * Hk / Hv) + s = source[h,v,i] + for t = 0 .. kept_count-1: + s *= decay[t,h] + s += key[t,hk,i] * delta[t,h,v] + destination[h,v,i] = s + Unused capsule tail entries are not read. All other destination elements are preserved. + + The output must be explicitly prebound to the same full allocation as destination_backing. + Destination backing must be distinct from source and capsule backing. This first implementation + supports one descriptor, uint32 shader indexing and no graph capture. + Source and capsule may share read-only backing; storage bindings are deduplicated only when + their BufferViews have the same Program binding owner, not merely equal raw buffer handles. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Inputs + +
+
source_backing : T
+
Full rank-1 FP32 backing tensor containing the state snapshot from immediately before capture.
+
capsule_backing : T
+
Full rank-1 FP32 compact-update backing tensor.
+
destination_backing : T
+
Full rank-1 FP32 inactive destination-state backing tensor.
+
metadata : TI
+
CPU int64 [11] containing offsets, geometry, capacity and kept_count.
+
+ +#### Outputs + +
+
destination_out : T
+
Same backing allocation and shape as destination_backing.
+
+ +#### Type Constraints + +
+
T : tensor(float)
+
State and capsule tensors are FP32.
+
TI : tensor(int64)
+
CPU replay metadata is int64.
+
+ + ### **com.microsoft.GatedRMSNorm** Gated RMS normalization as used by Mamba2 / gated DeltaNet attention outputs: @@ -3120,29 +3194,29 @@ This version of the operator has been available since version 1 of the 'com.micr This operator implements grouped-query attention with past state (KV cache) support. It also supports optional float8, int8 or int4 quantization for the KV cache to reduce memory footprint. - + **Cache Format:** The past and present KV cache tensors are expected in a BNSH format: `(batch_size, num_heads, cache_sequence_length, head_size)`, where `cache_sequence_length` is the length of the cached key/value sequences, or the maximum sequence length when past and present buffer sharing is used. - + **Windowed KV Cache (`sliding_window_cache` attribute):** When `sliding_window_cache` is 1, the past/present buffers are window-sized instead of full-length and the operator evicts internally. Let `C` be the cache capacity (dimension 2 of `past_key`, which is also the sequence dimension of `present_key`), `W` be `local_window_size`, and `T` be the absolute number of tokens processed so far by this batch entry, i.e. `seqlens_k[b] + 1`. The scalar `total_sequence_length` input is only the batch maximum of `T`; the layout below is per batch entry, so a ragged batch gets a different resident range per entry. `C` must be at least `W`. - + After a step, rows `[0, L)` of `present_key` and `present_value` hold the `L` most recent positions in increasing position order, so row `i` holds absolute position `T - L + i`. The retained positions are always physically contiguous and start at row 0; the layout never wraps around, so a ring-buffer layout cannot be exposed through these outputs. Rows `[L, C)` are unspecified. The resident count `L` is a function of `T` alone: - + ``` G = C - W + 1 L(T) = T if T <= C L(T) = T - G * ceil((T - C) / G) otherwise ``` - + Hence `min(T, W) <= L(T) <= min(T, C)`: the whole window stays resident, and eviction reclaims `G` positions at once rather than one position per step, so consumers must not assume that the cache is kept full at `min(T, C)`. - + Because `L` depends only on `T`, the resulting layout is independent of how the tokens were split into steps: a multi-token step of `S` tokens (speculative decoding, chunked prefill) leaves exactly the layout that the same tokens would produce one at a time. Any `S >= 1` is accepted, including `S > C`; a step that would evict positions it still has to read is staged internally, so the capacity does not have to cover the step. When past context is present, the existing operator restriction still applies: `sequence_length > 1` requires `batch_size == 1`. - + An execution provider may accept only part of the `C >= W` range. A configuration with `C < W` (equivalently, `W > C`) is invalid and is rejected with `INVALID_ARGUMENT`. The CUDA implementation requires `C == W`, so there `G` is 1 and `L(T)` is `min(T, C)`; a larger capacity is rejected. The CPU implementation accepts any `C >= W`, and slack above the window amortizes compaction over `G` steps. - + To drop the last `k` tokens, for example after rejecting speculative draft tokens, re-run with the smaller `total_sequence_length` and `seqlens_k` and leave the buffer untouched. That is exact when `L(T - k) == L(T) - k`, which callers can evaluate with the formula above. Otherwise the shorter layout needs positions that have already been evicted, and the window has to be re-materialized. - + **Quantization:** When quantization is enabled, `past_key` and `past_value` inputs can be of type `float8e4m3fn`, `uint8` or `int8`. The corresponding `k_scale` and `v_scale` tensors must be provided. The operator will output `present_key` and `present_value` in same format as the `past_key` and `past_value`. @@ -3274,7 +3348,7 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
num_branches : int (default is 0)
+
num_branches : int
Number of branches for flattened streams. Omit or set to zero for grouped streams.
@@ -3301,10 +3375,10 @@ This version of the operator has been available since version 1 of the 'com.micr #### Type Constraints
-
M : tensor(float), tensor(float16), tensor(bfloat16)
-
Constrain mixing weights to floating-point tensors.
T : tensor(float), tensor(float16), tensor(bfloat16)
Constrain streams, block output, and output to floating-point tensors.
+
M : tensor(float), tensor(float16), tensor(bfloat16)
+
Constrain mixing weights to floating-point tensors.
@@ -3323,9 +3397,9 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
num_branches : int (default is 0)
+
num_branches : int
Number of branches for flattened streams. Omit or set to zero for grouped streams.
-
reduction_scale : float (default is 1)
+
reduction_scale : float
Multiplier applied to the branch reduction.
@@ -3348,10 +3422,10 @@ This version of the operator has been available since version 1 of the 'com.micr #### Type Constraints
-
M : tensor(float), tensor(float16), tensor(bfloat16)
-
Constrain mixing weights to floating-point tensors.
T : tensor(float), tensor(float16), tensor(bfloat16)
Constrain streams and output to floating-point tensors.
+
M : tensor(float), tensor(float16), tensor(bfloat16)
+
Constrain mixing weights to floating-point tensors.
@@ -3756,7 +3830,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.MatMulBlockQuantizedFp8Weight** Block-scaled FP8 (E4M3) matrix multiplication with optional FP8 activation quantization. - + The weight tensor B has shape [N, K] with one FP32 scale per `block_size` consecutive K values (`b_scale` of shape [N, ceil(K / block_size)]). The scaled weight value is `B_scaled[n, k] = fp8_e4m3(B[n, k]) * b_scale[n, k / block_size]`. @@ -4377,7 +4451,7 @@ This version of the operator has been available since version 1 of the 'com.micr
activation_beta : float
Beta parameter used in activation function.
activation_type : string
-
Activation function to use. Choose from relu, gelu, silu, swiglu and identity. Default is relu
+
Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu
k : int
Number of top experts to select from expert pool
normalize_routing_weights : int
@@ -4623,17 +4697,6 @@ This version of the operator has been available since version 1 of the 'com.micr - head_offsets, when provided, adds a fixed per-output-head offset after the modulo by the head's vocabulary size, letting all heads across all n-gram orders share one flat embedding table. - Optional inputs add packed-sequence and Qwen4-Exp-style n-gram embedding support: - - - eos_token_id, when provided together with reset_on_eos != 0, causes causal history to reset at EOS - boundaries: any shifted position at or before the most recent EOS strictly before the current - position is replaced with eos_token_id instead of the real token. - - segment_ids, when provided, additionally resets causal history at any position whose segment id - differs from the immediately preceding position's segment id within input_ids. Segment boundaries - are not checked against past_ids history. - - head_offsets, when provided, adds a fixed per-output-head offset after the modulo by the head's - vocabulary size, letting all heads across all n-gram orders share one flat embedding table. - #### Version This version of the operator has been available since version 1 of the 'com.microsoft' operator set. @@ -5033,7 +5096,7 @@ This version of the operator has been available since version 1 of the 'com.micr [total_tokens, ...] axis instead of a dense [batch_size, sequence_length, ...] axis. It selects, for every packed query token, the sparse-attention candidates that a following SparsePagedAttention (or similar) operator is allowed to read. - + Unlike SparseAttentionIndexer, this operator: * takes packed query/key tensors plus cumulative_sequence_lengths (request boundaries) and past_sequence_lengths (per-request past length) instead of a dense batch and a dense mask; @@ -5044,10 +5107,10 @@ This version of the operator has been available since version 1 of the 'com.micr rejected as a deterministic no-op on state rather than truncated or allowed to corrupt memory; * additionally emits selected_counts, the exact number of active (non -1) entries per query, so that no downstream consumer needs to scan selected_indices for its query's true count. - + Both policy_mode values keep the semantics of SparseAttentionIndexer, applied independently to each request's own packed token range and fixed-capacity state slice: - + policy_mode = "qsa" ("query sparse attention" token indexer) Processes each request's new tokens sequentially: appends raw indexer keys to the generic pending buffer, and whenever it reaches compress_ratio tokens, mean-pools it, applies RMSNorm @@ -5057,8 +5120,9 @@ This version of the operator has been available since version 1 of the 'com.micr sum_h ReLU(q_h . k), the token_budget / compress_ratio highest scoring blocks are kept, and their token indices are emitted (request-local logical positions, i.e. the same numbering as past_sequence_lengths + local offset) followed by the causally visible tokens of the trailing - incomplete block. - + incomplete block. QSA positions are the request-local logical cache positions derived from + past_sequence_lengths and cumulative_sequence_lengths; position_ids must be omitted. + policy_mode = "csa" ("compressed sparse attention" block indexer) Applies the same window-plan arithmetic as SparseAttentionIndexer (overlap/leftover/new window count) independently per request, using that request's own buffer_length and new token count; @@ -5066,7 +5130,7 @@ This version of the operator has been available since version 1 of the 'com.micr rotated and appended to key_state. Queries are scored against every causally visible compressed entry with sum_h w_h * ReLU(q_h . k) and the index_topk highest scoring entry indices are emitted. - + Common contract: * selected_indices is int32 with a fixed capacity that only depends on attributes: token_budget + compress_ratio - 1 for "qsa" (values are request-local token positions into the @@ -5076,7 +5140,8 @@ This version of the operator has been available since version 1 of the 'com.micr with attention_mode="local_plus_selected", selected_kv_source="auxiliary"; key_state is layout-compatible with a [batch_size, capacity, 1, head_size] auxiliary cache when K = V). Unused entries are -1 and selected_counts holds the exact number of used entries. - * key_norm_weight is the effective RMSNorm multiplier, exactly as in SparseAttentionIndexer. + * query_norm_weight and key_norm_weight are the effective RMSNorm multipliers, exactly as in + SparseAttentionIndexer. * Accumulation, pooling, softmax, normalization and scoring are performed in float32 and the result is rounded once to the tensor element type. * Ties in the top-k selection are broken by the smaller entry index, and the emitted entries are @@ -5086,7 +5151,7 @@ This version of the operator has been available since version 1 of the 'com.micr * cumulative_sequence_lengths, past_sequence_lengths and past_state_lengths are read directly by the device kernel; a zero-token request row (a repeated cumulative offset) is valid and simply contributes no query rows for that request. - + OgaEngine integration note: this operator only defines the ORT contrib op; wiring past_key_state / past_kv_buffer / past_gate_buffer / past_state_lengths as Engine-managed, per-request fixed-size state (analogous to a paged auxiliary cache) is expected to happen in the @@ -5141,7 +5206,7 @@ This version of the operator has been available since version 1 of the 'com.micr
cumulative_sequence_lengths : M
Device-resident packed request boundaries with shape (batch_size + 1); cumulative_sequence_lengths[0] must be 0 and cumulative_sequence_lengths[batch_size] must equal total_tokens. Request b owns rows [cumulative_sequence_lengths[b], cumulative_sequence_lengths[b + 1]) of query/key (a repeated offset is a valid zero-token row).
past_sequence_lengths : M
-
Device-resident number of tokens already processed for each request before this call, with shape (batch_size). Used as the default absolute query position when position_ids is omitted (policy_mode 'qsa'), and to validate state consistency.
+
Device-resident number of tokens already processed for each request before this call, with shape (batch_size). Used to derive request-local logical query positions for policy_mode 'qsa' and to validate state consistency.
gate (optional) : T
Only for policy_mode 'csa': gate projection of the new tokens with shape (total_tokens, 2 * head_size).
position_bias (optional) : T
@@ -5149,7 +5214,7 @@ This version of the operator has been available since version 1 of the 'com.micr
head_weights (optional) : T
Only for policy_mode 'csa': per-head score weights with shape (total_tokens, num_heads).
position_ids (optional) : I
-
Optional for policy_mode 'qsa', required for policy_mode 'csa': absolute position of every packed query, with shape (total_tokens).
+
Only for policy_mode 'csa': absolute position of every packed query, with shape (total_tokens).
past_key_state : T
Generic fixed-capacity state: policy_mode 'qsa' stores prepared complete-block keys; policy_mode 'csa' stores compressed keys. Shape is (batch_size, state_capacity, head_size) and never changes across calls.
past_kv_buffer : T
@@ -5207,7 +5272,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.PackedSparseAttentionIndexerMerge** - Stable union of a mapped base selection and [range_start, range_end). Valid prefixes contain nonnegative indices. Padding is ignored. Outputs preserve first-occurrence order and are padded with -1. Status is 0 for success, 1 for invalid device metadata, or 2 for capacity overflow; failures return count 0 and all -1. The caller supplies a common index namespace and guarantees causal visibility. + Stable union of a mapped base selection and either [range_start, range_end) or an additional selection prefix. Valid prefixes contain nonnegative indices. Padding is ignored. Outputs preserve first-occurrence order and are padded with -1. Status is 0 for success, 1 for invalid device metadata, or 2 for capacity overflow; failures return count 0 and all -1. The caller supplies a common index namespace and guarantees causal visibility. #### Version @@ -5236,9 +5301,9 @@ This version of the operator has been available since version 1 of the 'com.micr
range_ends (optional) : T
Exclusive ends not less than starts, shape (N).
additional_indices (optional) : T
-
Reserved; must be absent for append_range.
+
Required for append_indices: selections of shape (N, A); absent for append_range.
additional_counts (optional) : T
-
Reserved; must be absent for append_range.
+
Required for append_indices: valid prefix lengths of shape (N); absent for append_range.
#### Outputs @@ -6196,13 +6261,15 @@ This version of the operator has been available since version 1 of the 'com.micr A positive block_size selects square block scaling with float32, float16, or bfloat16 fc*_scales tensors shaped [num_experts, ceil(N / block_size), ceil(K / block_size)]: dequantized_weight[e, n, k] = float(weight[e, n, k]) * scale[e, n / block_size, k / block_size] - Partial edge blocks are allowed. No zero points or activation scales are used. + Partial edge blocks are allowed. No zero points or activation scales are accepted. Without a positive block_size, FP8 uses the legacy per-expert fc*_global_scale inputs instead. Block-scaled FP8 does not use global scales. Activations retain the input type (weight-only quantization). On CUDA SM80-or-newer GPUs, SiLU/SwiGLU block-scaled FP8 uses fused weight-only GEMV/GEMM by default. These kernels decode and scale FP8 tiles on chip without materializing full FP16/BF16 weight buffers. Set ORT_ENABLE_FP8_FUSED=0 before session creation to select the compact dense-dequant fallback. Other activation types and unsupported routing sizes retain the fallback. + The WebGPU block-FP8 kernel supports block_size=128 with float32 scales and rejects projections + that exceed 32-bit shader addressing or the device's per-dimension dispatch limit. Packed byte dimensions are computed as logical_element_count * effective_expert_weight_bits / 8. Weight rows must be byte-aligned. Zero-point rows are padded to a whole byte when necessary. @@ -6238,9 +6305,9 @@ This version of the operator has been available since version 1 of the 'com.micr
activation_beta : float
Beta parameter used in activation function.
activation_type : string
-
Activation function to use. Choose from relu, gelu, silu, swiglu and identity. Default is relu
+
Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu
block_size : int
-
For integer quantization, size of each quantization block along the K (input feature) dimension. Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). Both hidden_size and inter_size must be divisible by the block size. The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. For FP8 ('fp8'), a positive value instead specifies square blocks along both N and K, with floating-point scales shaped [E, ceil(N/block_size), ceil(K/block_size)]; partial blocks are allowed. Without a positive value, FP8 uses per-expert global scales. For integer quantization ('int'), omitting block_size means there is no blocking and a whole column shares one scaling factor.
+
For integer quantization, size of each quantization block along the K (input feature) dimension. Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). For integer and FP4 quantization, both hidden_size and inter_size must be divisible by the block size. FP8 with block_size=128 supports partial 128x128 tiles. The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. For FP8 ('fp8'), a positive value instead specifies square blocks along both N and K, with floating-point scales shaped [E, ceil(N/block_size), ceil(K/block_size)]; partial blocks are allowed. Without a positive value, FP8 uses per-expert global scales. For integer quantization ('int'), omitting block_size means there is no blocking and a whole column shares one scaling factor.
expert_weight_bits : int
Number of bits used in quantized weights. Supported values are 2, 4, and 8. Default is 4 bits
fc1_expert_weight_bits : int
@@ -7285,7 +7352,7 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
alpha : float (default is 1)
+
alpha : float
Scale used when the optional scale input is absent.
@@ -7308,10 +7375,10 @@ This version of the operator has been available since version 1 of the 'com.micr #### Type Constraints
-
M : tensor(float), tensor(float16), tensor(bfloat16)
-
Constrain the optional scale to floating-point tensors.
T : tensor(float), tensor(float16), tensor(bfloat16)
Constrain input and output to floating-point tensors.
+
M : tensor(float), tensor(float16), tensor(bfloat16)
+
Constrain the optional scale to floating-point tensors.
@@ -8352,6 +8419,8 @@ This version of the operator has been available since version 1 of the 'com.micr subsequent calls to preserve boundaries across chunked prefill and decode calls. - head_offsets, when provided, adds a fixed per-output-head offset after the modulo. Addition wraps in the input id type on overflow. + - capture_count enables compact prefix state capture when state_update_capacity is positive. For + each request, state_update contains the trailing state after each captured prefix token. #### Version @@ -8368,9 +8437,11 @@ This version of the operator has been available since version 1 of the 'com.micr
Compressed tokenizer id used to pad causal shifts before the beginning of a request's sequence.
reset_on_eos : int
When non-zero and eos_token_id is provided, reset causal n-gram history at EOS boundaries. Default is 0.
+
state_update_capacity : int
+
Static number of compact contiguous-prefix states to expose per request. Valid range is [0, 8]. capture_count is required exactly when this is positive.
-#### Inputs (4 - 9) +#### Inputs (4 - 10)
input_ids : M
@@ -8391,9 +8462,11 @@ This version of the operator has been available since version 1 of the 'com.micr
Optional token-major segment ids with shape (total_tokens), used to reset causal history at segment boundaries within each packed request.
past_segment_ids (optional) : S
Optional segment ids corresponding to past_ids, with shape (batch_size, max_ngram_size - 1). Thread present_segment_ids from the previous call into this input to preserve segment boundaries across calls.
+
capture_count (optional) : S
+
Optional device int32 tensor with shape (batch_size). For each request, captures that many local prefix states, clamped to the request length and state_update_capacity. Required exactly when state_update_capacity is positive.
-#### Outputs (1 - 3) +#### Outputs (1 - 4)
hash_ids : M
@@ -8402,6 +8475,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Trailing max_ngram_size - 1 ids of past_ids followed by each request's own tokens, with shape (batch_size, max_ngram_size - 1). Feed this back as past_ids on the next call.
present_segment_ids (optional) : S
Trailing max_ngram_size - 1 segment ids corresponding to present_ids, with shape (batch_size, max_ngram_size - 1). Feed this back as past_segment_ids on the next call.
+
state_update (optional) : M
+
Optional trailing id states after each captured prefix token, with shape (batch_size, state_update_capacity, max_ngram_size - 1). Unused slots are pad_id.
#### Type Constraints @@ -8696,3 +8771,4 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
+ diff --git a/docs/Experimental_C_API_Guidelines.md b/docs/Experimental_C_API_Guidelines.md new file mode 100644 index 0000000000000..90e4201e12a58 --- /dev/null +++ b/docs/Experimental_C_API_Guidelines.md @@ -0,0 +1,92 @@ +# Experimental C API Guidelines + +## Overview + +ONNX Runtime (ORT) provides a mechanism for exposing experimental C API functions. Experimental APIs are a proving ground for new functionality. They may eventually be promoted to the stable ORT API or abandoned, so their availability is not guaranteed across releases. While an experimental function is available under a given name, however, its signature does not change, and its ABI-relevant types and behavioral contract remain backward compatible. + +Experimental functions are not part of the stable `OrtApi` struct. Instead, they are resolved by name at runtime through the stable C API [`OrtApi::GetExperimentalFunction`](https://onnxruntime.ai/docs/api/c/struct_ort_api.html#ace33338eb9175cdc92b52f61c9aff9a6), which was introduced in ORT 1.28. Given a name, it returns a generic function pointer, or `NULL` if the named function is not available in the running ORT. A non-null result means the function is present and can be used; a null result means it is absent and the caller must fall back or fail gracefully. This keeps experimental functions off the stable ABI while still making them reachable from any ORT build that supports them. + +Every experimental function has a name that ends in a `_SinceV` suffix, where `N` is the ORT API version in which that function was first introduced (for example, `OrtApi_ExperimentalApiTest_SinceV28`). The suffix is part of the identity of the function: it makes each name unique across API versions (which correspond to ORT minor releases). Backward-compatible fixes and behavioral extensions may retain the existing name; a signature change or an incompatible ABI or behavioral-contract change requires a new name. + +The lookup names, the function-signature typedefs, and any auxiliary types the experimental functions require are provided by a companion C header, [onnxruntime_experimental_c_api.h](../include/onnxruntime/core/session/onnxruntime_experimental_c_api.h). C++ consumers should instead use the companion C++ header, [onnxruntime_experimental_cxx_api.h](../include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h), which builds on the C header and adds typed accessors in the `Ort::Experimental` namespace. + +## User Guidelines + +These guidelines are for consumers who want to *call* experimental functions. + +### Expect availability to vary + +There is no promise that an experimental function will exist in any future (or past) version of ORT. A newer experimental header may also remove its typedef and name constant. By choosing to use an experimental API, you are committing to adapting your code when updating the header — either following the function to its stable replacement after promotion, or accounting for its removal. At runtime, do not assume that a function declared by the header used to build your application is present in the loaded ORT library. + +### Check availability at runtime + +Because an experimental function may be absent, always confirm that the lookup returned a usable pointer before calling it. In C, this means checking the pointer returned by `GetExperimentalFunction` against `NULL`. The C++ header offers two accessor flavors per function so you can choose how absence is handled: + +- A nullable accessor (`Get__Fn`) returns the typed function pointer, or `nullptr` if the function is not available in this build. Use it when the function is optional and you want to branch on availability. +- A throwing accessor (`Get__FnOrThrow`) returns a guaranteed-non-null typed function pointer, or throws `Ort::Exception` with `ORT_NOT_IMPLEMENTED` if the function is not available. Use it when the function is required and its absence is a hard error. + +### Use the experimental headers for names and signature typedefs + +Do not hard-code experimental function names or re-declare their signatures by hand. The companion headers are the source of truth: + +- The C header supplies, for each experimental function, a function-pointer typedef (`OrtExperimental__Fn`) and a name constant (`kOrtExperimental__FnName`). Look the function up with the name constant, then cast the returned generic pointer to the matching typedef before calling it. Calling through any other type is undefined behavior. +- The C++ header wraps this pattern in the typed `Ort::Experimental` accessors described above, so you get the correct type without a manual cast, plus any C++ wrapper types associated with the experimental functions. + +The headers provide the known function names and their signatures; the runtime you are calling into will either have a given function or not, which is exactly what the runtime lookup tells you. For C and C++ usage examples, see [onnxruntime_experimental_c_api.h](../include/onnxruntime/core/session/onnxruntime_experimental_c_api.h). + +## Implementer Guidelines + +These guidelines are for ORT contributors who want to *add* an experimental function. + +### Declare the function in one place + +Experimental functions are declared in a single X-macro include file, [`onnxruntime_experimental_c_api.inc`](../include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc), that serves as the source of truth for the consumer headers and the runtime registration table. Add one `ORT_EXPERIMENTAL_API(SinceVersion, ReturnType, Name, Params...)` entry, where: + +- `SinceVersion` is the current ORT API version, written as a numeric literal (for example, `28`) — not a macro such as `ORT_API_VERSION`. +- `ReturnType` is the function's return type, typically `OrtStatusPtr`. +- `Name` is the base name (`` below) without the version suffix, by convention prefixed with the stable struct the function is destined for. +- `Params` is the parameter list, using the same SAL annotations as other ORT C API declarations. + +The lookup name (`_SinceV`), the typedefs, the name constants, the C++ accessors, and the registration-table entry are all generated mechanically from this one entry. Implement the function body under the `OrtExperimentalApis` namespace in whatever source location is appropriate. + +### Follow the naming convention + +Experimental function names follow the pattern `_SinceV`. `` consists of `_`. The expanded pattern is `__SinceV`. + +The `` prefix (for example, `OrtApi`, `OrtEpApi`, or `OrtCompileApi`) indicates the stable struct the function is intended to be promoted into. + +`` is the intended name of the C API function. + +The `_SinceV` suffix records the ORT API version in which the experimental function first appeared and guarantees the name is unique by construction. This suffix is automatically appended via the `ORT_EXPERIMENTAL_API` X-macro. + +At promotion, using the placeholders above, the `` struct gets a new function, ``, appended to it. + +### Place auxiliary types in the appropriate header + +If an experimental function needs a new auxiliary type, declare it alongside the function's other artifacts: opaque handle types belong in the C header, and any C++ wrapper types belong in the C++ header, in the `Ort::Experimental` namespace. + +### Document the expected promotion timeline + +Document the expected promotion or abandonment timeline in the new API's documentation. E.g., "Promotion to stable API expected by version X." This is an expectation, not a contract, and it can be updated as plans change. + +### Keep a given experimental API name backward compatible + +Because an experimental API name may be available across more than one ORT release while consumers fetch the function by name, its signature must not change, and its ABI-relevant auxiliary types and behavioral contract must remain backward compatible for the lifetime of that name. Compatible bug fixes and behavioral extensions are allowed. For a signature change or an incompatible ABI or behavioral-contract change, introduce a *new* name rather than altering the existing one. This should be straightforward to do in a later API version, since the introducing API version is baked into the name. + +### Add tests + +Add test coverage for new experimental functions, exercising both the lookup mechanism and the function's behavior. Test coverage is expected just as for any other API function. + +## Experimental API Lifecycle + +1. **Add.** Add one experimental API entry with the current ORT API version and implement the function. +2. **Update.** A backward-compatible fix or behavioral extension may retain the existing name. For a signature change or an incompatible ABI or behavioral-contract change, introduce a new experimental API entry at the current ORT API version — which produces a new lookup name — rather than modifying the existing one. The old and new entries may coexist. Never reuse a name. +3. **Promote or remove.** + - **Promote:** Add the function to the stable API struct (append-only, dropping the `_` prefix and the `_SinceV` suffix from the name) and remove the experimental entry. + - **Remove:** Delete the experimental entry. Because the versioned name is unique and never reused, no separate retirement tracking is needed. + +## References + +- [onnxruntime_experimental_c_api.h](../include/onnxruntime/core/session/onnxruntime_experimental_c_api.h) — C names, typedefs, and auxiliary type declarations. +- [onnxruntime_experimental_cxx_api.h](../include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h) — C++ `Ort::Experimental` accessors and wrapper types. +- [test_experimental_api.cc](../onnxruntime/test/shared_lib/test_experimental_api.cc) — test coverage and example usage for experimental API function lookup and the C++ accessors. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 59ca1b3f29411..9786f3169bac0 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -63,13 +63,15 @@ The **OpSet Version** column uses the following notation: |||[9, 13]|**T** = tensor(double), tensor(float)| |||[7, 8]|**T** = tensor(double), tensor(float)| |BitCast|*in* input:**T1**
*out* output:**T2**|26+|**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| -|BitShift|*in* X:**T**
*in* Y:**T**
*out* Z:**T**|11+|**T** = tensor(uint32), tensor(uint64), tensor(uint8)| +|BitShift|*in* X:**T**
*in* Y:**T**
*out* Z:**T**|28+|**T** = tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|||[11, 27]|**T** = tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |BitwiseAnd|*in* A:**T**
*in* B:**T**
*out* C:**T**|18+|**T** = tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |BitwiseNot|*in* X:**T**
*out* Y:**T**|18+|**T** = tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |BitwiseOr|*in* A:**T**
*in* B:**T**
*out* C:**T**|18+|**T** = tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |BitwiseXor|*in* A:**T**
*in* B:**T**
*out* C:**T**|18+|**T** = tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |BlackmanWindow|*in* size:**T1**
*out* output:**T2**|17+|**T1** = tensor(int32), tensor(int64)
**T2** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| -|Cast|*in* input:**T1**
*out* output:**T2**|25+|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(float8e8m0), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(float8e8m0), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)| +|Cast|*in* input:**T1**
*out* output:**T2**|28+|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(float8e8m0), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(float8e8m0), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)| +|||[25, 27]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(float8e8m0), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(float8e8m0), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)| |||24|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(float8e8m0), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(float8e8m0), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)| |||23|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)| |||[21, 22]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint2), tensor(uint32), tensor(uint4), tensor(uint64), tensor(uint8)| @@ -287,7 +289,8 @@ The **OpSet Version** column uses the following notation: |||12|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint32), tensor(uint64), tensor(uint8)| |||[8, 11]|**T** = tensor(double), tensor(float)| |||[6, 7]|**T** = tensor(float)| -|Mod|*in* A:**T**
*in* B:**T**
*out* C:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|Mod|*in* A:**T**
*in* B:**T**
*out* C:**T**|28+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|||[13, 27]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[10, 12]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |Mul|*in* A:**T**
*in* B:**T**
*out* C:**T**|14+|**T** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||13|**T** = tensor(double), tensor(float), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| @@ -477,7 +480,7 @@ The **OpSet Version** column uses the following notation: |||[1, 21]|**T** = tensor(float)| |Softsign|*in* input:**T**
*out* output:**T**|22+|**T** = tensor(float)| |||[1, 21]|**T** = tensor(float)| -|SpaceToDepth|*in* input:**T**
*out* output:**T**|13+|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| +|SpaceToDepth|*in* input:**T**
*out* output:**T**|[13, 27]|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| |||[1, 12]|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| |Split|*in* input:**T**
*in* split:**T**
*out* outputs...:**T**

or

*in* input:**T**
*in* split:**tensor(int64)**
*out* outputs:**T**

or

*in* input:**T**
*out* outputs:**T**|18+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[13, 17]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| @@ -693,7 +696,8 @@ The **OpSet Version** column uses the following notation: |||14|**T** = tensor(double), tensor(float), tensor(float16)
**U** = tensor(double), tensor(float), tensor(float16)| |||[9, 13]|**T** = tensor(double), tensor(float), tensor(float16)| |||[7, 8]|**T** = tensor(double), tensor(float), tensor(float16)| -|Cast|*in* input:**T1**
*out* output:**T2**|25+|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|Cast|*in* input:**T1**
*out* output:**T2**|28+|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|||[25, 27]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[23, 24]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[21, 22]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[19, 20]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| @@ -866,7 +870,8 @@ The **OpSet Version** column uses the following notation: |Min|*in* data_0:**T**
*out* min:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| |||12|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| |||[6, 11]|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| -|Mod|*in* A:**T**
*in* B:**T**
*out* C:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| +|Mod|*in* A:**T**
*in* B:**T**
*out* C:**T**|28+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| +|||[13, 27]|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| |||[10, 12]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| |Mul|*in* A:**T**
*in* B:**T**
*out* C:**T**|14+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||13|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| @@ -1017,7 +1022,7 @@ The **OpSet Version** column uses the following notation: |||[1, 21]|**T** = tensor(double), tensor(float), tensor(float16)| |Softsign|*in* input:**T**
*out* output:**T**|22+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |||[1, 21]|**T** = tensor(double), tensor(float), tensor(float16)| -|SpaceToDepth|*in* input:**T**
*out* output:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16)| +|SpaceToDepth|*in* input:**T**
*out* output:**T**|[13, 27]|**T** = tensor(double), tensor(float), tensor(float16)| |||[1, 12]|**T** = tensor(double), tensor(float), tensor(float16)| |Split|*in* input:**T**
*in* split:**T**
*out* outputs...:**T**

or

*in* input:**T**
*in* split:**tensor(int64)**
*out* outputs:**T**

or

*in* input:**T**
*out* outputs:**T**|18+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[13, 17]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| @@ -1200,7 +1205,7 @@ The **OpSet Version** column uses the following notation: |||10|**I** = tensor(int64)
**T** = tensor(float), tensor(float16)| |||[8, 9]|**I** = tensor(int64)
**T** = tensor(float), tensor(float16)| |||[1, 7]|**T** = tensor(float), tensor(float16)| -|SpaceToDepth|*in* input:**T**
*out* output:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16)| +|SpaceToDepth|*in* input:**T**
*out* output:**T**|[13, 27]|**T** = tensor(double), tensor(float), tensor(float16)| |||[1, 12]|**T** = tensor(double), tensor(float), tensor(float16)| | | | | @@ -1242,12 +1247,13 @@ The **OpSet Version** column uses the following notation: |||14+|**T** = tensor(float), tensor(float16)| |||9+|**T** = tensor(float), tensor(float16)| |||7+|**T** = tensor(float), tensor(float16)| -|BitShift|*in* X:**T**
*in* Y:**T**
*out* Z:**T**|11+|**T** = tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|BitShift|*in* X:**T**
*in* Y:**T**
*out* Z:**T**|[11, 27]|**T** = tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |BitwiseAnd|*in* A:**T**
*in* B:**T**
*out* C:**T**|18+|**T** = tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |BitwiseNot|*in* X:**T**
*out* Y:**T**|18+|**T** = tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |BitwiseOr|*in* A:**T**
*in* B:**T**
*out* C:**T**|18+|**T** = tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |BitwiseXor|*in* A:**T**
*in* B:**T**
*out* C:**T**|18+|**T** = tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| -|Cast|*in* input:**T1**
*out* output:**T2**|21+|**T1** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|Cast|*in* input:**T1**
*out* output:**T2**|28+|**T1** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|||21+|**T1** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||19+|**T1** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||13+|**T1** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||9+|**T1** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| @@ -1409,7 +1415,8 @@ The **OpSet Version** column uses the following notation: |||12+|**T** = tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||8+|**T** = tensor(float), tensor(float16)| |||6+|**T** = tensor(float), tensor(float16)| -|Mod|*in* A:**T**
*in* B:**T**
*out* C:**T**|13+|**T** = tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint8)| +|Mod|*in* A:**T**
*in* B:**T**
*out* C:**T**|28+|**T** = tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint8)| +|||13+|**T** = tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint8)| |||10+|**T** = tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint8)| |Mul|*in* A:**T**
*in* B:**T**
*out* C:**T**|14+|**T** = tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||13+|**T** = tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| diff --git a/docs/Validate_CUDA_Package.md b/docs/Validate_CUDA_Package.md new file mode 100644 index 0000000000000..0b438fd5ce18c --- /dev/null +++ b/docs/Validate_CUDA_Package.md @@ -0,0 +1,101 @@ +# Validate CUDA kernel execution in a Python package + +A Python wheel can import successfully and list `CUDAExecutionProvider` while +lacking executable CUDA kernels for the installed GPU. CPU architecture and GPU +compute capability are separate requirements: an aarch64 wheel also needs CUDA +code compatible with its target GPU. + +`onnxruntime/test/python/onnxruntime_test_python_cuda_kernel.py` is a standalone +package smoke test. It generates an FP32 `Mul` model and synthetic inputs in +memory, so no downloaded models, datasets, credentials, or application code are +needed. It checks two input sizes, exact results, and profiling events proving +that `Mul` ran on the requested execution provider. Graph optimizations and CPU +fallback are disabled for the CUDA test. Missing CUDA support is a failure, +not a skipped test. + +## Run against an installed wheel + +Use a machine with a supported NVIDIA GPU, driver, and the CUDA/cuDNN libraries +required by the wheel. See the [CUDA execution provider requirements](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements). +Create an isolated environment and install the wheel under test: + +```bash +python3 -m venv /tmp/ort-cuda-check +. /tmp/ort-cuda-check/bin/activate +ORT_WHEEL=/absolute/path/to/the/candidate.whl +python -m pip install numpy onnx "$ORT_WHEEL" +``` + +Set `ORT_WHEEL` to the actual wheel filename. Do not install the CPU and GPU +distributions together: they both provide the `onnxruntime` Python module. +Copy the single test file into a directory outside the source checkout and run it +there to ensure Python imports the installed package: + +```bash +cp onnxruntime/test/python/onnxruntime_test_python_cuda_kernel.py /tmp/ort-cuda-check/ +cd /tmp/ort-cuda-check +python onnxruntime_test_python_cuda_kernel.py -v +``` + +Both `test_mul_cpu` and `test_mul_cuda` must pass. A CPU-only sanity check is +available, but does not establish GPU support: + +```bash +python onnxruntime_test_python_cuda_kernel.py TestCudaKernel.test_mul_cpu -v +``` + +To select a GPU on a multi-GPU machine, set `CUDA_VISIBLE_DEVICES` before starting +Python. The test uses the default device visible to ONNX Runtime. + +## Reproduce an architecture coverage failure + +For example, run the CUDA test on an SM100 GPU with a Linux aarch64 CUDA 13 wheel +that omits SM100 kernels. The expected failure is a CUDA kernel launch error such +as `cudaErrorNoKernelImageForDevice` during `Mul`. Repeat the same test on the same +GPU with a wheel built with SM100 coverage. The corrected package must pass both +the result and provider checks. + +Use a separate virtual environment for each wheel. Record the wheel SHA-256, +ONNX Runtime version/build information, Python version, CPU architecture, GPU +model/compute capability, driver, CUDA/cuDNN versions, and test output. Useful +commands include: + +```bash +sha256sum "$ORT_WHEEL" +uname -m +getconf PAGESIZE +nvidia-smi --query-gpu=name,compute_cap,driver_version --format=csv,noheader +python -c 'import platform, onnxruntime as ort; print(platform.python_version()); print(ort.__version__); print(ort.get_build_info()); print(ort.__file__); print(ort.get_available_providers())' +``` + +An unavailable execution provider or missing CUDA library is an environment or +package failure, not evidence of the particular missing-kernel regression. +Likewise, CPU success alone is not a substitute for the CUDA test. + +## Source builds and coverage limits + +For a focused native aarch64 SM100 build, first follow the +[CUDA build prerequisites](https://onnxruntime.ai/docs/build/eps.html#cuda) and +set `CUDA_HOME` and `CUDNN_HOME` to compatible CUDA 13 and cuDNN installations. +Then run from the source checkout: + +```bash +python tools/ci_build/build.py --config Release --build_dir build/sm100 \ + --update --build --parallel 8 --build_wheel --use_cuda \ + --cuda_home "$CUDA_HOME" --cudnn_home "$CUDNN_HOME" \ + --cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=100-real' +``` + +Install the resulting wheel from `build/sm100/Release/dist/` into the isolated +environment above. This focused build requests SM100; validating the full +packaging architecture list still requires the corresponding packaging build. + +The existing Python test runner invokes this module when `--use_cuda` and Python +tests are enabled. CMake copies it with the other Python test files. A packaging +job that only builds wheels does not execute it; validate the produced wheel on +the target GPU separately using the commands above. + +This test checks basic compiled CUDA kernel execution. It does not qualify every +operator, every supported architecture, plugin execution providers, or +application-level numerical accuracy. Run it on each GPU architecture relevant +to a package change, alongside the existing provider and application tests. diff --git a/docs/annotated_partitioning/attention_workspace_estimation.md b/docs/annotated_partitioning/attention_workspace_estimation.md index 04899befead52..5d61455f59227 100644 --- a/docs/annotated_partitioning/attention_workspace_estimation.md +++ b/docs/annotated_partitioning/attention_workspace_estimation.md @@ -240,24 +240,34 @@ Neither adapter changes runtime `GetScratchBuffer()` calls or allocation topology. The CPU `total_sequence_length` scalar and past/present aliasing are not -available through `WorkspaceInputShape`. For non-windowed GQA, the scalar can -exceed the past-cache sequence dimension and directly scale Flash, MEA, and -unfused workspace. Non-windowed execution can also preserve one full past -tensor when exactly one past/present pair aliases. The current adapter therefore -reports non-windowed GQA as unavailable rather than treating the past shape as -a total-KV bound or omitting alias-preservation scratch. - -Successful estimates are currently limited to sliding-window cache nodes. -Their final cache capacity remains `C`, while a multi-token step uses a -transient staged/effective attention extent of `C + S`; a single-token step -uses `C`. This shape-only bound remains sound even when the scalar -`total_sequence_length` is much larger. Runtime requires both past/present -pairs to alias before staging or compaction, excluding the partial-alias -preservation path. Sliding-window attention bias remains -unsupported by runtime and is also unavailable to the estimator. If cuDNN can -be reached anywhere in a graph-free bounded domain, aggregation remains -unavailable because cuDNN's allocator-based workspace has no sound graph-free -oracle. Level 1 then uses the generic fallback and Level 2 emits no requirement. +available through `WorkspaceInputShape`. For non-windowed GQA, the positive +caller-declared `ep.cuda.gqa_workspace_max_total_sequence_length` value supplies +the missing total-KV bound. The estimator requires that envelope to cover the +query-length bound and sizes cache-dependent work at the greater of the +allocated past-cache capacity and the envelope. Because graph analysis cannot +prove aliasing, the aggregate also includes the worst-case full-past +preservation scratch. The declaration is an estimation contract, not a +runtime-enforced input limit or a no-OOM guarantee. + +Successful estimates can cover both sliding-window nodes and positive-envelope +non-windowed nodes. Sliding-window cache capacity remains `C`, while a +multi-token step uses a transient staged/effective attention extent of `C + S`; +a single-token step uses `C`. Runtime requires both past/present pairs to alias +before sliding-window staging or compaction, excluding the partial-alias +preservation path. Estimation remains unavailable for missing or invalid shape +facts, an unset or undersized non-windowed envelope, cache capacity outside the +runtime int32 ABI, attention bias, or any reachable cuDNN route because cuDNN's +allocator-based workspace has no sound graph-free oracle. Under capacity-aware +CUDA placement, unavailable GQA proof declines CUDA assignment; generic +fallback remains available to unrelated operators, and profile data cannot +substitute for the missing GQA proof. Level 2 emits no requirement when its +estimate is unavailable. + +The CUDA plugin's host-side accounting path does not yet receive analytical +estimates, so capacity-aware placement declines new GQA assignments there even +when an envelope is declared. A fused group containing such a node is declined +as a whole. Placement without a resource accountant, previously accounted +assignments, and unrelated operators or plugin EPs retain their existing behavior. Unlike the log-only PA/PMHA Level-1 probe, a successful GQA Level-1 estimate is passed to the #31962 resource accountant as `runtime_workspace_bytes` and can diff --git a/docs/contrib_ops/PackedSparseAttentionIndexerMerge.md b/docs/contrib_ops/PackedSparseAttentionIndexerMerge.md index 8dbe76d542ad1..dd4cf123b6b32 100644 --- a/docs/contrib_ops/PackedSparseAttentionIndexerMerge.md +++ b/docs/contrib_ops/PackedSparseAttentionIndexerMerge.md @@ -32,9 +32,13 @@ every status before publishing dependent results. A status output does not prevent downstream nodes from executing in the same graph. CUDA uses one cooperative block per query with invocation-local scratch for -stable duplicate removal. WebGPU uses one invocation per query, with a -workgroup hash table for capacities up to 4096 and an output-prefix scan for -larger capacities. Neither kernel modifies the base tensors or retains a +stable duplicate removal. Hash scratch is bounded by the candidate width and +the output capacity plus one overflow-detection entry, rather than the output +capacity alone. Successful rows use at most half of the hash slots. +WebGPU uses one cooperative workgroup per query. Parallel first-occurrence +checks and ordered chunk compaction preserve stable output order without a +capacity-dependent workgroup hash table or a serial large-capacity fallback. +Neither kernel modifies the base tensors or retains a selection between runs. Regression tests: `PackedSparseAttentionIndexerMerge.*` in diff --git a/docs/contrib_ops/cuda/gemm_profiler_cache.md b/docs/contrib_ops/cuda/gemm_profiler_cache.md new file mode 100644 index 0000000000000..db72ee973c5d0 --- /dev/null +++ b/docs/contrib_ops/cuda/gemm_profiler_cache.md @@ -0,0 +1,343 @@ +# fpA_intB / QMoE GEMM Tactic Autotune Cache — Design & Plan + +## 1. Scope and motivation + +This PR implements persistent tactic caching for **MatMulNBits fpA_intB only**. +QMoE caching and automatic model-sidecar selection are future work. The reduced +initial M sweep and lazy profiling already exist in the base branch; persistence +avoids repeating that work in a new process. + +## 2. Current implementation + +| Concern | Behavior | +|---|---| +| In-process identity | Shape, activation dtype, packing SM, quantization bits/group size, zero points, bias, GEMV support, and GPU name | +| Initial sweep | Small configurable M set, default `{1,2,4,...,2048}`; missing runtime buckets are profiled lazily | +| Persistence | Opt-in MatMulNBits TSV file; no file I/O when no location is configured | +| Reuse guard | GPU name, SM, CUDA runtime, ORT version, and tactic-selection version | +| Loaded tactics | Candidate membership plus an untimed execution on profiler scratch checks problem-specific constraints and compiled kernel support before reuse; buckets above the session's launch limit are skipped | +| Failed profiling | No negative results are saved; old negative rows are ignored and retried | +| Flush boundary | Construction and lazy results are staged in memory, then batched at CUDA EP teardown for both built-in and plugin EPs | +| Shared maps | Newly added buckets observed on exact or rounded hits, including CUDA graph capture, are staged once per attached profiler, including different cache prefixes | +| Concurrency | Only locally changed rows are merged into a complete disk snapshot; read errors preserve the old file and dirty rows; all failed writes clean up their temporary file | + +### Startup validation measurement + +On an H200 with CUDA 13.0 and the compact fpA_intB build, a 750 MB synthetic model +with 32 distinct MatMulNBits shapes (`K=8192`, `N=4096,4160,...,6080`) and five M +buckets (`1,8,64,512,2048`) took 3.01 s to construct cold. Three fresh-process +cached runs took 1.79–1.96 s (median 1.82 s); forcing 64-row chunks reduced cached +startup to 1.74 s by skipping the larger validation buckets. These are total +session-construction times, including model loading and weight packing, not +isolated validation timings. All problem-specific execution checks remain enabled; +stream/scratch reuse is deferred pending evidence that validation dominates startup. + +## 3. Goals / non-goals + +**Goals (including future work)** +- Reduce first-time tuning by profiling a small, configurable set of M buckets. +- Persist tuned tactics to disk and reuse across sessions. +- Deduplicate identical shapes across nodes (single cache entry per unique problem). +- Reuse only on the *same* GPU model + SM + toolkit/runtime version. +- Optionally attach the cache to the model. +- Provide an offline tuning tool. +- Keep persistent filesystem writes opt-in unless an explicit cache directory/prefix is configured. + +**Non-goals** +- Changing kernel implementations or the tactic search space. +- Sharing tactics across different GPUs or CUDA/ORT versions. +- Persisting prepacked weights (a separate concern). + +## 3.1 Priority + +Implement the cache in this order: + +1. **MatMulNBits fpA_intB first.** This path pays the broad first-time profile sweep at kernel + construction and is the clearest startup-time win. +2. **QMoE grouped GEMM second.** Production QMoE currently calls `MoeGemmProfiler` for grouped-GEMM + tactic selection, including integer W4A16/W8A16 grouped GEMM and FP4/FP8 modes when they route + through the grouped-GEMM runner. The MoE GEMV fast path itself does not need profiling, but the + current QMoE compute path selects grouped-GEMM tactics before the runner can fall through to GEMV. + Therefore QMoE caching is useful, but lower priority than MatMulNBits unless QMoE startup tuning + shows up as the dominant user-visible cost. + +## 4. Cache split and serialization format + +MatMulNBits and QMoE should use **separate cache tables**. They can be implemented by the same +`gemm_tactic_cache` utility, but the persisted rows should not be forced into one sparse schema: + +- MatMulNBits is a single GEMM/GEMV problem keyed by fpA_intB packed shape and quantization mode. +- QMoE is a grouped-GEMM problem keyed by expert topology, activation, bias, parallelism, and GEMM + role (`Gemm1` or `Gemm2`). + +Physical files (only the MatMulNBits table is implemented): + +``` +.matmulnbits_fpa_intb.tsv +.qmoe_gemm.tsv # planned +``` + +For a user cache directory, `` is derived from the hardware/build signature. For a +model sidecar, `` is derived from the model path when the model path is known, or from +an explicit `ORT_CUDA_GEMM_TACTIC_CACHE_PREFIX` / session config value. + +The cache is small and **naturally tabular**: one row per `(problem_key, M-bucket) → tactic`. The +chosen format is **TSV** (tab-separated) because ORT core has no JSON library, the Python offline +tool must read/write the same file, and `cudaDeviceProp.name` contains spaces (and occasionally +punctuation) that make CSV quoting fragile. + +| Format | C++ parse | Python | Debuggable | Evolution | Decision | +|---|---|---|---|---|---| +| **TSV** | trivial split | `csv` stdlib | excellent | map by header name | **Chosen** | +| CSV | trivial | `csv` stdlib | good | good | Rejected (name spaces/commas → quoting) | +| JSON (hand-rolled) | needs a mini-parser (no dep) | native | good | excellent (nested/optional) | Fallback | +| Binary (flatbuffers/proto) | build complexity | needs bindings | none | careful | Rejected (overkill for KB data) | +| INI/key=value | easy | manual | ok | good | Rejected (awkward for repeated rows) | + +**MatMulNBits file layout** (`.matmulnbits_fpa_intb.tsv`): + +``` +# ort_cuda_gemm_tactic_cache v1 +# table matmulnbits_fpa_intb +# tactic_selection_version 4 +# device_name NVIDIA A100-SXM4-80GB +# sm 80 +# multiprocessor_count 108 +# cuda_runtime 13000 +# cuda_driver 13000 +# ort_version 1.23.0 +# ort_git_commit +# ort_build_config Release +n_16b k activation_dtype weight_type bits block_size has_zero_points zero_point_dtype gemv_enabled has_bias packing_sm paired_gemv_mode wave_aware_gemv m_bucket valid_config sm_version tile80 tile90 tile100 tile120 split_k_style split_k stages cluster mainloop epilogue tma enable_cuda_kernel cuda_kernel_variant +12288 4096 half uint4b_t 4 64 1 uint4b_t 1 0 80 0 0 1 1 80 23 0 0 0 1 1 3 0 0 0 0 1 0 +12288 4096 half uint4b_t 4 64 1 uint4b_t 1 0 80 0 0 64 1 80 17 0 0 0 0 -1 4 0 0 0 0 0 0 +... +``` + +**QMoE file layout** (`.qmoe_gemm.tsv # planned`): + +``` +# ort_cuda_gemm_tactic_cache v1 +# table qmoe_gemm +# tactic_selection_version 2 +# device_name NVIDIA A100-SXM4-80GB +# sm 80 +# multiprocessor_count 108 +# cuda_runtime 13000 +# cuda_driver 13000 +# ort_version 1.23.0 +# ort_git_commit +# ort_build_config Release +schema_version n k activation_dtype weight_dtype quant_format gemm_type num_experts top_k activation_type bias group_size parallelism_mode need_weights m_bucket valid_config sm_version tile80 tile90 tile100 tile120 split_k_style split_k stages cluster mainloop epilogue tma enable_cuda_kernel +1 2880 2880 half uint4b_t mxfp4 Gemm1 32 4 Swiglu 0 32 default 0 512 1 80 23 0 0 0 0 1 1 3 0 0 0 0 0 +... +``` + +- `#`-prefixed lines carry a cache format version, table name, and hardware/build signature. +- One header row names columns; readers map by **column name** so appended columns are + backward/forward compatible (unknown columns ignored; missing columns defaulted). +- Enum/config fields are stored as their integer values (mirror of `CutlassGemmConfig`). +- Free-text fields must not contain raw tabs or newlines. Use percent-encoding for `%`, `\t`, and + `\n` when writing and decode on read. +- `valid_config=0` is understood for format compatibility but ignored on load. Only successful + tactics are persisted, because a profiling failure can be temporary. +- Writes are atomic: acquire a lock, reload the current file, merge the new rows, write to a temp + file, then `rename` and release the lock. The lock prevents lost updates from two sessions tuning + different shapes concurrently. + +## 5. Hardware signature + +Purpose: reject reuse across different GPU models / SM / toolkit versions. + +``` +signature = { + device_name: cudaDeviceProp.name // e.g. "NVIDIA GeForce RTX 4090" + sm: major*10 + minor // e.g. 80, 89, 90 + multiprocessor_count: cudaDeviceProp.multiProcessorCount + cuda_runtime: CUDART_VERSION + cuda_driver: cudaDriverGetVersion() + ort_version: ORT_VERSION + ort_git_commit: ORT_GIT_COMMIT or "unknown" + ort_build_config: ORT_BUILD_CONFIG or build config string +} +``` + +- Used both to name the cache file and as a stored guard. On load, a mismatch on + `device_name` / `sm` / `cuda_runtime` / `ort_version` rejects the file (forces re-profiling). + `multiprocessor_count`, `cuda_driver`, `ort_git_commit`, and `ort_build_config` are recorded for + diagnostics only; every loaded CUTLASS tactic is also re-validated against the current runner. +- `tactic_selection_version` (a code constant, not hardware) is also strict. Bump it whenever the + profiler's selection logic changes (for example, the L2-resident GEMV bias), because dev builds + keep the same `ort_version` across commits. +- Rationale: RTX 4090 and RTX 4060 are both `sm_89` but perform differently, so the **device name**, + not just SM, is the primary discriminator — matching the requirement. + +## 6. Cache keys (problem keys) + +Use op-specific keys rather than a single sparse superset. + +``` +matmulnbits_key = { + n_16b, k, activation_dtype, weight_type, bits, block_size, + has_zero_points, zero_point_dtype, gemv_enabled, has_bias, packing_sm, + paired_gemv_mode, wave_aware_gemv +} + +qmoe_key = { + n, k, activation_dtype, weight_dtype, quant_format, gemm_type, + num_experts, top_k, activation_type, bias, group_size, + parallelism_mode, need_weights +} +``` + +Each key maps to a set of `(m_bucket → optional)`. Many nodes with the same shape +collapse to one entry. For `MatMulNBits`, `n_16b` is the value used by the existing `GemmIdCore`. +For QMoE, the key includes the profiler parameters that affect `GemmProfilerBackend::init` so that +cache reuse remains conservative. + +MatMulNBits keys separate paired-K modes (disabled, autotuned, forced) and wave-aware GEMV +launches. `cuda_kernel_variant` preserves the selected GEMV implementation (0 = default, +1 = paired-K, 2–7 = the six symmetric INT4 M=1 decode geometries). Selection version 4 +rejects older selections made before the streaming-L2 decode profiler and its new candidates. +Loaded decode variants must still pass the current layout, dtype, quantization, M-bucket, and +grid-limit checks before they are validated on the device. + +## 7. Implemented architecture + +The shared module `contrib_ops/cuda/llm/gemm_tactic_cache.{h,cc}` provides hardware +signatures, TSV encoding, config parsing, and the `MatMulNBitsTacticCache` wrapper. +There is no QMoE wrapper yet. + +Location precedence is session prefix, session directory, environment prefix, +then environment directory. This selects one file, not a sequence of fallback +files. No location means in-process caching only. + +```mermaid +flowchart TD + A[Kernel construction] --> B{Persistent cache configured?} + B -- yes --> C[Load matching signature and validate tactics] + B -- no --> D[Profile missing initial buckets] + C --> D + D --> E[Stage successful tactics in memory] + F[Lazy profiling or new shared-map hit] --> E + E --> G[EP teardown: snapshot registry, release registry lock] + G --> H[Lock file, read complete snapshot, merge dirty rows, atomic replace] +``` + +## 8. Original roadmap and future extensions + +This roadmap includes unimplemented QMoE and model-sidecar work. Section 2 is the +authoritative description of shipped behavior; phases below are not an options list. + +### Phase 1 — Foundation (`gemm_tactic_cache.{h,cc}`) +Hardware signature, TSV (de)serialization, `CutlassGemmConfig` serialize/parse, problem key, +file lock, atomic file write, and two typed cache wrappers: MatMulNBits and QMoE. Self-contained and +unit-testable. + +### Phase 2 — Wire disk cache into both profilers *(depends on 1)* +- *Strategy*: Wire the `gemm_tactic_cache` utility independently into `GemmPluginProfiler` and `MoeGemmProfiler` rather than attempting a large refactor. The shared cache utility provides a simple `GetConfig(...)` / `PutConfig(...)` API. +- `gemm_profiler.h`: before the M sweep in `profileTactics`, load matching entries and skip if hit; + after profiling, merge into the disk cache. Re-enable the serialize stubs. +- `matmul_nbits.{h,cc}`: pass `bits / block_size / zero-point mode / gemv_enabled / has_bias / packing_sm` + into the MatMulNBits key. +- `moe_quantization.cc` + `moe_gemm_profiler.{h,cc}`: use the same disk cache + key; QMoE gains + cross-node dedup via the shared cache. Add a process-global MoE profiler manager + mirroring `s_profilerManager` for in-session dedup before the first write to prevent redundant file I/O for models with many MoE nodes. +- Cache directory/prefix is opt-in via env var or session config. If unset, do not write persistent + files. + +### Phase 3 — Reduce first-time M sweep *(parallel with 2)* +- Replace the fixed `1…8291` sweep with a small default bucket set + `{1,2,4,8,16,32,64,128,256,512,1024,2048}` clamped to `[minM, maxM]`, unioned with exact decode + `M=1` and `bucketM(maxM)`. +- If runtime asks for an unprofiled bucket, **lazily profile that single bucket** and merge it into the cache. This briefly blocks inference but guarantees optimal performance and handles unprofiled edge cases safely. +- `ORT_FPA_INTB_PROFILE_M` (comma list) overrides the set. Derive `MatMulNBits` `max_m` from + dims/env instead of the `8291` literal. + +### Phase 4 — Optional store-with-model *(depends on 2)* +- Sidecars use the same cache prefix scheme: + - `".matmulnbits_fpa_intb.tsv"` + - `".qmoe_gemm.tsv"` +- Writing is opt-in via session config or env var (`ep.cuda.gemm_tactic_cache_to_model` / + `ORT_CUDA_GEMM_TACTIC_CACHE_TO_MODEL`). The CUDA EP records the model path it sees in + `GetCapability` (`GraphViewer::ModelPath()` for the built-in EP, `Graph_GetModelPath` for the plugin + EP), and kernels use it as the cache prefix. Models loaded from bytes have no path, so the explicit + prefix/dir settings still apply to them. + +### Phase 4b — Load on create, save on close *(depends on 2)* +- Kernel construction loads matching rows before profiling and re-reads the file first if it changed + since this process last read or wrote it (last-write time + size). +- Every tuned tactic (construction sweep and lazy buckets) is only staged in memory. The single disk + write happens when the session closes (CUDA EP teardown), with a last-chance write at process exit + for sessions that are never released. +- `ep.cuda.gemm_tactic_cache_mode` / `ORT_CUDA_GEMM_TACTIC_CACHE_MODE` selects `load_save` (default), + `load` (read-only), or `save` (re-tune and overwrite). +- Session close writes every dirty cache in the process, not only the closing session's caches. The + write is dirty-guarded and merges with the file, so writing another live session's rows early is + harmless. + +### Phase 5 — Offline tuning tool *(depends on 4)* +- `onnxruntime/python/tools/fpa_intb_tune.py` (+ CLI): `--model`, `--output-prefix`, `--m-values`, + `--no-inference`. Sets session options, builds a CUDA EP session, runs dummy inferences over + the requested M values to force profiling, releases the session to flush the cache, prints the + path and tuned shapes. + +### Phase 6 — Docs + tests *(depends on 2–5)* +- Document env vars, sidecar format, hardware guard, and the offline tool in + `docs/contrib_ops/cuda/matmul_nbits.md` and the QMoE docs. +- Unit tests: `CutlassGemmConfig` serialize/parse round-trip; signature-mismatch rejection; + TSV column-name mapping with an appended column. +- A100 validation: 2nd session reuses cache (profiling skipped, tune time ≈ 0); outputs unchanged; + many same-shape MoE nodes yield one entry per unique shape. +- The fresh-process Python integration test also runs with the CUDA plugin wheel. It checks + session-specific selected-tactic reports so import-time plugin log filtering cannot hide cache + reuse or bucket filtering. Built-in CUDA additionally checks the INFO diagnostics. +- CUDA Python CI installs CPU-only PyTorch for weight-packer parity tests when the transformers + suite is disabled; compact builds skip only prepacked cases outside their supported contract. +- The standalone CUDA weight-packer oracle supports both SM80 and SM90 offline layouts even when + the build excludes SM90 compute kernels. Explicit `force_arch=90` must not substitute SM80 bytes. + +## 9. Supported environment variables & options + +| Name | Effect | +|---|---| +| `ORT_CUDA_GEMM_TACTIC_CACHE_DIR` | Directory for persistent cache files. Unset means persistent cache disabled. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_PREFIX` | Explicit file prefix. Writes `.matmulnbits_fpa_intb.tsv`. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_TO_MODEL` | `1` → store the cache next to the model as `.matmulnbits_fpa_intb.tsv` when the model path is known. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_MODE` | `load_save` (default), `load` (read-only), or `save` (re-tune and overwrite). | +| `ORT_FPA_INTB_PROFILE_M` | Comma-separated M buckets to profile (overrides the default set for MatMulNBits/fpA_intB). | +| `ep.cuda.gemm_tactic_cache_dir` | Session-option equivalent of `ORT_CUDA_GEMM_TACTIC_CACHE_DIR`. | +| `ep.cuda.gemm_tactic_cache_prefix` | Session-option equivalent of `ORT_CUDA_GEMM_TACTIC_CACHE_PREFIX`. | +| `ep.cuda.gemm_tactic_cache_to_model` | Session-option equivalent of `ORT_CUDA_GEMM_TACTIC_CACHE_TO_MODEL`. | +| `ep.cuda.gemm_tactic_cache_mode` | Session-option equivalent of `ORT_CUDA_GEMM_TACTIC_CACHE_MODE`. | +| `ep.cuda.gemm_tactic_cache_tuning_results_prefix` | Optional separate TSV output of this session's selected tactics (including valid cache hits), flushed at EP teardown. The tuning tool uses a fresh temporary prefix to verify that the current model selected successful tactics and that their full settings reached the main cache. No environment-variable fallback. | + +## 10. Original roadmap file inventory (includes future work) + +**Add** +- `contrib_ops/cuda/llm/gemm_tactic_cache.h` / `.cc` +- `onnxruntime/python/tools/fpa_intb_tune.py` + +**Modify** +- `contrib_ops/cuda/llm/gemm_profiler.h` — sweep reduction, cache hooks, re-enable serialize +- `contrib_ops/cuda/llm/cutlass_extensions/gemm_configs.h` — config (de)serialization support +- `contrib_ops/cuda/quantization/matmul_nbits.{h,cc}` — key metadata, `max_m`, sidecar wiring +- `contrib_ops/cuda/moe/moe_quantization.cc`, `llm/moe_gemm/moe_gemm_profiler.{h,cc}` — cache + bucket cut +- `docs/contrib_ops/cuda/matmul_nbits.md`, QMoE docs + +Note: normal and plugin CUDA provider builds glob `contrib_ops/cuda/*.cc`, so no explicit CMake edit +is expected for `gemm_tactic_cache.cc`. The new `.cc` still needs a `#if defined(USE_CUDA) && defined(ORT_USE_CUDA_GEMM)` (or analogous) guard +and must avoid dependencies on non-plugin CUDA EP internals. + +## 11. Roadmap decisions (not all implemented) + +1. **QMoE in-process dedup**: Added to Phase 2. A global MoE profiler manager is necessary to prevent duplicate profiling in the very first session before the cache is cleanly merged. +2. **Cache miss strategy**: Added to Phase 3. Lazy profile of the exact missing bucket is the chosen strategy over nearest-neighbor fallback. +3. **Profiler integration**: Handled via simple utility APIs in both `GemmPluginProfiler` and `MoeGemmProfiler` rather than a unified base class to minimize refactoring risk. + +## 12. Open considerations + +1. **Strictness of driver/MP checks** — the guard is strict on device name, SM, runtime, and ORT + version; driver, MP count, git commit, and build config stay diagnostic unless validation shows + they must be strict. +2. **JSON fallback** — switch to hand-rolled JSON only if nested/optional structure is later required. diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index 46775da75200d..c8e0c1d023da3 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -237,6 +237,80 @@ devices decline the fused path and use generic dequantization plus cuBLAS. --- +### 4.5 Opt-in FP16 raw INT4 tail outputs + +`ORT_MATMULNBITS_UNALIGNED4_FAST_PATH=1` optionally extends the existing raw +INT4 fused kernels to output widths `N % 8 == 4`, including `N=324`. +It defaults to **off**. Eligibility is restricted to FP16, `bits=4`, +`block_size=32`, column-wise quantization, `N >= 4`, `K > 0` divisible by 32, +and `1 <= M <= 4`, with no zero points, bias or `g_idx`. The fpA_intB path +and forced chunking are excluded. The kernels themselves support `M <= 8`, but +for `M = 5..8` the batched tail kernel was slower than the fallback in +measurements (see below), so larger `M` keeps the existing dequantize + GEMM path. +Existing aligned matrices and default-off dispatch remain unchanged; this +selector does not turn on fpA_intB. + +For M=1, the existing shared-scale kernel is instantiated with tail masking. +The final CTA loads only valid scale rows and fills invalid shared slots with +zero. All eight warps participate in the scale-loading barrier before invalid +output warps exit; those warps never read weights or store output. For M=2..4, +the existing register-batched kernel uses one output per warp, with whole +invalid warps exiting before tensor reads. That kernel has no CTA barrier. +Valid outputs retain the original kernel's unpacking, FP16 accumulation order +and warp reduction. No packing, persistent allocation, expert routing, or +model precision change is introduced. + +Measured on one H200 (CUDA 13.0, FP16, `block_size=32`, device-resident I/O, +median of 7x500 `Run`s, which include per-run host launch overhead). Off is the +existing dequantize + GEMM fallback; on is this selector: + +| N x K | M=1 off / on (us) | M=2 | M=3 | M=4 | M=5 | M=8 | +| --- | --- | --- | --- | --- | --- | --- | +| 324 x 10240 | 24.7 / 19.3 | 28.9 / 17.7 | 28.4 / 20.4 | 26.3 / 20.7 | 26.6 / 38.8 | 26.3 / 42.6 | +| 324 x 4096 | 21.7 / 16.4 | 25.6 / 16.0 | 25.6 / 17.1 | 23.6 / 17.3 | 23.6 / 24.8 | 23.0 / 25.7 | +| 1028 x 4096 | 23.0 / 16.2 | 28.3 / 15.8 | 28.4 / 16.9 | 26.3 / 17.1 | 25.8 / 24.4 | 25.9 / 26.0 | +| 4100 x 2048 | 26.9 / 16.1 | 30.8 / 17.1 | 31.0 / 18.4 | 31.0 / 19.0 | 30.7 / 25.5 | 29.6 / 26.8 | + +The selector is therefore limited to `M <= 4`. These are operator-level +timings, not end-to-end model throughput. + +When explicitly enabled, each kernel prints at most one successful dispatch +receipt per M1/batched branch: + +```text +[MatMulNBitsUnaligned4] node="projection" M=1 K=10240 N=324 type=FP16 bits=4 branch=M1_SHARED_TAIL +[MatMulNBitsUnaligned4] node="projection" M=4 K=10240 N=324 type=FP16 bits=4 branch=BATCHED_REGISTER_TAIL +``` + +Atomic flags prevent duplicate lines from concurrent Runs. A receipt follows +the launch and a successful non-synchronizing CUDA error peek; it is not proof +of completed asynchronous execution. Capture can record this launch without +new cache/event dependencies; graph replay does not re-enter host dispatch. + +The numerical reference is the **original aligned raw INT4 kernel**, padding +N324 to N328 and comparing the first 324 FP16 output bytes (similarly N4 to 8 +and N12 to 16). This is **not** a promise of byte parity with the previous +dequantization-plus-cuBLAS fallback, whose arithmetic differs. Full-model +token/quality parity and an independently audited performance/memory pilot +are required before accepting the selector for a run. + +GPU validation (CUDA EP or the CUDA plugin EP; skipped when CUDA is unavailable): + +```bash +python -m unittest onnxruntime/test/python/transformers/test_matmul_nbits_unaligned4.py -v +``` + +The paired test requires 144 byte comparisons over M1,2,3,4; N4,12,324; +K32,320,10240; and constant, signed-zero, adversarial and realish tensors. +Independent poison rows in the padded reference must not affect valid outputs. +Because `InferenceSession` allocates tensors from the ORT CUDA arena rather +than exact-sized standalone buffers, this session-level test does not provide +exact-boundary OOB detection under `compute-sanitizer`; verifying strict +allocation boundaries requires a direct kernel launch harness with exact-sized +input and output buffers and prefix/suffix poison guards. + +--- + ## 5. Fallback Path — Dequantize + GEMM When neither the fpA_intB path nor the fused GEMV applies (e.g. `M > 1`, @@ -531,6 +605,132 @@ Prepacked weights are intentionally strict: --- +### 6.3 fpA_intB Tactic Autotune Cache + +The fpA_intB profiler (§6) tunes CUTLASS/CUDA tactics per problem shape. By +default this is in-process only and re-run every session. When a cache location +is configured, tuned tactics are **persisted to disk** and reused across +sessions, and the first-time tuning cost is reduced. Design details: +[gemm_profiler_cache.md](gemm_profiler_cache.md). + +**Reduced first-time sweep + lazy fallback.** Instead of profiling a dense +`M = 1 … 8291` sweep at construction, the profiler now tunes a small bucket set +(`{1,2,4,…,2048}` by default, clamped to the problem range and always including +`M=1` and the top bucket). If a runtime `M` maps to an unprofiled bucket, that +single bucket is profiled lazily on first use and kept in the in-process map for +the rest of the session. Lazy profiling is skipped +while a CUDA graph is being captured (profiling kernels/events/allocations are +illegal during capture), so run a +warmup inference **before** capture to tune any `M` buckets your captured graph +needs. Set `ORT_FPA_INTB_PROFILE_M` (comma-separated) to override the bucket set; +its maximum also sets the initial +profile range. + +**Load on create, save on close.** When a session is created, each +`MatMulNBits` kernel loads matching tactics from the cache before profiling, so +only missing buckets are tuned. If another session or process changed the cache +file since this process last read or wrote it, the file is read again first. +Newly tuned tactics (from the construction-time sweep and from lazy profiling) +are only staged in memory while the session runs. They are written to disk once, +when the session closes (`~CUDAExecutionProvider` for the built-in EP, `~CudaEp` +for the CUDA plugin EP), so neither session creation nor inference does any +cache file writes. Tactics of sessions still alive at process exit are written +when the provider library is unloaded. + +**Enabling persistence.** Persistence is opt-in. Provide a directory, an +explicit file prefix, or enable the model sidecar, via environment variable or +session-config option: + +| Env var | Session-config key | Effect | +|---------|--------------------|--------| +| `ORT_CUDA_GEMM_TACTIC_CACHE_DIR` | `ep.cuda.gemm_tactic_cache_dir` | Directory for the cache file; the file name is derived from the hardware signature. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_PREFIX` | `ep.cuda.gemm_tactic_cache_prefix` | Explicit file prefix; writes `.matmulnbits_fpa_intb.tsv`. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_TO_MODEL` | `ep.cuda.gemm_tactic_cache_to_model` | `1` stores the cache next to the model as `.matmulnbits_fpa_intb.tsv`. Ignored for models loaded from bytes. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_MODE` | `ep.cuda.gemm_tactic_cache_mode` | `load_save` (default): load on create, save on close. `load`: read-only, never writes (e.g. a pre-tuned cache shipped with an application). `save`: ignore existing rows, re-tune, and save on close; the session profiles every shape again even if another session in the same process already tuned it (nodes with the same shape inside the session still share one tuning). Other values fail session creation. | + +Resolution order: session-config prefix → session-config dir → model sidecar → +env prefix → env dir. A non-empty prefix wins over a directory, and an explicit +session-config location overrides the sidecar and both env locations. Each +setting is read from session config first and falls back to its env var. When +no location is set, nothing is +written and behavior is unchanged. Cache instances are process-global and keyed +by their resolved location, so all `MatMulNBits` nodes/sessions that resolve to +the same location share one cache, while sessions configured with different +locations each get their own. Missing cache directories are created on the first +write. + +**Stale-tactic guard.** On load, each CUTLASS tactic read from disk is checked +against the tactics the current runner can actually dispatch (`getConfigs()`); +any row that no longer matches is dropped and that bucket is re-profiled, so a +parseable-but-incompatible cache row is never handed to the kernel. +Validation skips buckets above the session's fpA_intB launch limit when M chunking is enabled. + +**File format.** A tab-separated file (`.matmulnbits_fpa_intb.tsv`) with +`#`-prefixed header lines carrying a format version and a hardware/build +signature (GPU name, SM, CUDA runtime, ORT version/commit/build config), a +column-name row, and one row per `(problem shape, M bucket) → tactic`. Readers +map by column name, so appended columns are backward/forward compatible. + +**Hardware guard.** On load, the file's signature must strictly match the current +`device_name`, `sm`, `cuda_runtime`, and `ort_version`; otherwise the file is +ignored and shapes are re-profiled. This prevents reusing tactics across different +GPUs or toolkit/ORT versions (e.g. RTX 4090 vs RTX 4060, which share `sm_89`). +`ort_git_commit`, `ort_build_config`, `multiprocessor_count`, and `cuda_driver` +are recorded in the header for diagnostics but are **not** part of the guard: a +tactic is only a config selected among `getConfigs()`, so a cross-commit or +cross-build reuse can at worst pick a slightly suboptimal (never incorrect) +tactic, and the stale-tactic guard above re-validates every CUTLASS tactic +against the current runner anyway. The header also carries a +`tactic_selection_version` that must match exactly; it is bumped whenever the +profiler's selection logic changes, so caches tuned under an older policy are +re-profiled even when the ORT version is unchanged. + +**Offline tuning tool.** To pre-populate a cache without running your own +workload, use +[onnxruntime/python/tools/fpa_intb_tune.py](../../../onnxruntime/python/tools/fpa_intb_tune.py): + +```bash +python -m onnxruntime.tools.fpa_intb_tune \ + --model model.onnx \ + --output-prefix /path/to/cache/mymodel \ + --m-values 1,8,16,32,64,128,256,512,1024,2048 +``` + +It enables the fpA_intB path, the cache prefix, and the profile bucket set through +session config entries, creates a CUDA-EP session (which profiles the bucket set during +kernel construction), best-effort runs dummy inferences at each `M`, releases the session +so all staged buckets are flushed, validates the cache signature against the active CUDA +device and installed wheel, then prints the cache path and a summary of tuned shapes. +Buckets already in the cache are reused; pass `--retune` to profile every bucket +again (`ep.cuda.gemm_tactic_cache_mode=save`). +CPU fallback, missing files, incompatible signatures, and caches without successful +tactics are errors. BF16 dummy inputs use explicitly typed OrtValues. +The tool also requires a fresh session-specific record of selected tactics and verifies that +each full tactic configuration was persisted, so an unrelated cache or an unsaved correction +cannot report success. Valid cache hits succeed without rewriting the original +file. The tool currently requires a built-in CUDA Python package; CPU packages with a registered +CUDA plugin are rejected with an explicit build-metadata diagnostic. The cache itself supports +both built-in and plugin providers. + +**onnxruntime-genai integration.** Two options: + +1. *Environment variable* (no config change): export + `ORT_CUDA_GEMM_TACTIC_CACHE_DIR=/path` before launching genai. +2. *genai_config.json* (recommended): add the session-config key under the + decoder's `session_options`. genai forwards any unrecognized `session_options` + entry to `AddConfigEntry`, so it reaches the CUDA EP: + + ```json + "decoder": { + "session_options": { + "ep.cuda.gemm_tactic_cache_dir": "/path/to/cache", + "provider_options": [ { "cuda": { "enable_cuda_graph": "1" } } ] + } + } + ``` + +--- + ## 7. Bias Handling Only the router specialization (§4.3) fuses bias inside the GEMV. For every other @@ -551,6 +751,12 @@ present. `ComputeInternal` then: |----------|----------------|--------| | `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION` | bool, `0` | Disable the router GEMV specialization (§4.3); shapes fall back to the generic GEMV / dequant path. Useful for A/B benchmarking. | | `ORT_FPA_INTB_GEMM` | int/string, `0` | Enable the CUTLASS weight-only path (§6). `0` or `off` disables it, otherwise enables it. | +| `ORT_MATMULNBITS_UNALIGNED4_FAST_PATH` | `0` or `1`, `0` | Opt in to FP16 scale-only raw INT4 tail-output kernels for N%8=4, block32 and M1..4 (§4.5). | +| `ORT_FPA_INTB_PROFILE_M` | comma list, unset | Override the M buckets profiled for the fpA_intB tactic cache (§6.1). The maximum value also bounds the initial profile range. Session-config equivalent: `ep.cuda.fpa_intb_profile_m`. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_DIR` | path, unset | Directory for the persistent fpA_intB tactic cache (§6.1). Unset means the cache is in-process only. Session-config equivalent: `ep.cuda.gemm_tactic_cache_dir`. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_PREFIX` | path prefix, unset | Explicit cache file prefix (§6.1); writes `.matmulnbits_fpa_intb.tsv`. Session-config equivalent: `ep.cuda.gemm_tactic_cache_prefix`. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_TO_MODEL` | bool, `0` | Store the fpA_intB tactic cache next to the model (§6.1). Session-config equivalent: `ep.cuda.gemm_tactic_cache_to_model`. | +| `ORT_CUDA_GEMM_TACTIC_CACHE_MODE` | string, `load_save` | `load_save`, `load` (read-only), or `save` (re-tune) for the fpA_intB tactic cache (§6.1). Session-config equivalent: `ep.cuda.gemm_tactic_cache_mode`. | | `ORT_MATMULNBITS_FORCE_CHUNKED` | int, `0` | Force the chunked dequant+GEMM fallback (§5) regardless of the size heuristic, and bypass the fpA_intB M-chunking size condition (§6.2). | | `ORT_MATMULNBITS_CHUNK_SIZE` | int64, `32768` | Target rows per chunk in the chunked fallback. Values `< 1` reset to the default. | | `ORT_MATMULNBITS_M_CHUNK_SIZE` | int, `0` | Max rows of `A` per fpA_intB launch (§6.2). `0` disables M chunking. Overridden by the `ep.cuda.matmul_nbits_m_chunk_size` session config entry. Also applies to the CUDA plugin EP. | @@ -620,6 +826,11 @@ present. `ComputeInternal` then: [onnxruntime/test/contrib_ops/matmul_2bits_test.cc](../../../onnxruntime/test/contrib_ops/matmul_2bits_test.cc). Run it from `onnxruntime_provider_test` with `--gtest_filter=MatMul2BitsCuda.*`. +- fpA_intB tactic cache (§6.3) unit tests: + [onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc](../../../onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc) + (config serialize/parse round-trip, signature-mismatch rejection, appended-column + tolerance, store/load and merge). Build with `onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON` + and run `./onnxruntime_provider_test --gtest_filter=CUDA_EP_Unittest.All`. - GEMV profiling baselines and methodology are recorded in [qmoe_gemv_experiments.md](qmoe_gemv_experiments.md). - To compare the router specialization against the generic path, run the same diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md index cc18492d78228..0dbcb7b0dacb0 100644 --- a/docs/contrib_ops/cuda/moe_qmoe.md +++ b/docs/contrib_ops/cuda/moe_qmoe.md @@ -82,7 +82,7 @@ input tokens → router (top-k softmax) → permute by expert | `expert_weight_bits` (QMoE only) | int | 4 | 4 (INT4/MXFP4) or 8 (INT8/FP8). | | `block_size` (QMoE only) | int | -1 | Group size for INT4/INT8 group-wise quantization. -1 = per-output-channel. FP4/WFP4AFP8 normalize an omitted value to 32 and require 32 when present; NVFP4 normalizes an omitted value to 16 and requires 16 when present. FP8 uses square blocks for positive values, per-expert globals otherwise. | | `quant_type` (QMoE only) | string | `"int"` | `"int"`, `"fp4"`, `"nvfp4"`, `"fp8"`, `"wfp4afp8"`. See [§3](#3-quantization-modes). | -| `weights_prepacked` (QMoE only) | int | -1 | Tri-state, only meaningful when `quant_type="int"`. The prepacked layouts selected by `-1` and `1` are **EP-determined**. `-1` (default): the INT4/INT8 `fc1`/`fc2` initializers are already prepacked in the EP's default layout (e.g. from `pack_weights_for_cuda_mixed_gemm` for the CUDA EP). `1`: already prepacked in an alternate EP-selected layout. `0`: the initializers are raw `[E, N, K/pack]` tensors (as produced by `quantize_matmul_{4,8}bits`) and the kernel runs the CUTLASS layout transform in `PrePack()`. **Note:** the CUDA EP INT4/INT8 MoE GEMM always runs the Ampere (SM80) kernel — even on SM90 — so it consumes the SM80 `fpA_intB` layout on all architectures; `-1` and `1` are therefore equivalent for the CUDA EP today, and `1` is reserved for a possible future Hopper-specific layout. See [§5.1](#51-weights-input-2--5--8). | +| `weights_prepacked` (QMoE only) | int | -1 | Layout selector for integer quantization and CUDA NVFP4. Integer semantics are unchanged: `-1` and `1` select EP-determined prepacked layouts (both SM80 `fpA_intB` on CUDA today); `0` selects raw `[E, N, K/pack]` weights transformed in `PrePack()`. For CUDA NVFP4, omitted/`-1` and `0` select legacy N-packed bytes; `1` selects offline K-packed row-major bytes. Mode `2` is unsupported for NVFP4. See [§5.1](#51-weights-input-2--5--8). | ### 2.2 Type Constraints @@ -333,6 +333,27 @@ pointers are read at compute time instead. > layout transform for INT weights. FP4 / FP8 / WFP4AFP8 weight handling is > unaffected. +**CUDA NVFP4** uses a separate `weights_prepacked` contract: + +| `weights_prepacked` | Physical packed-byte storage | +|---------------------|------------------------------| +| Omitted / `-1` | Legacy N-packed `[E, K, N/2]` | +| `0` | Explicit legacy N-packed `[E, K, N/2]` | +| `1` | Offline K-packed row-major `[E, N, K/2]` | + +Logical ONNX dimensions remain `[E, K, N/2]` in every mode. Mode `1` changes +only the physical weight-byte layout: E2M1 codes, raw E4M3 block scales +`[E, N, K/16]`, and FP32 global scales `[E]` retain their values. Mode `2` is +unsupported. Integer-quantization semantics are unchanged. + +Mode `1` shares the offline payload across eligible native GEMV decode, +Blackwell native GEMM prefill, and dense dequantization fallback without a +decode-only weight repack or expanded GEMV scale bank. Blackwell native GEMM +still derives its TMA-swizzled block scales; the dense fallback reads the +retained raw scale inputs, not that swizzled buffer. Disabling prepacking with +`session.disable_prepacking` preserves correct offline-layout dequantization +and eligible GEMV execution; native prefill requires prepacking. + MXFP4 weights must be packed by `pack_fp4_weights_for_cuda_moe_gemm`. FP8 weights are stored as raw e4m3 bytes (no packing). @@ -809,6 +830,7 @@ debug switches. | `ORT_ENABLE_FP4_GEMV` | on | Fused MXFP4 GEMV decode kernel. Set to `0` to force the dequant-to-dense fallback (debugging/bisecting). Active in the SM80+ fallback regime, and as the decode arm when native CUTLASS prefill is enabled. | | `ORT_FP4_GEMV_AUTOTUNE` | `0` | Opt-in per-shape autotune of the GEMV CtaN/Threads tiling. Enabling it synchronizes the first uncached inference for each shape. | | `ORT_FP4_GEMV_AUTOTUNE_LOG` | `0` | Set to `1` to log the chosen GEMV configs per shape. | +| `ORT_FP4_GEMV_DECODE_KERNEL` | `1` | NVFP4 offline K-packed (`weights_prepacked=1`) decode GEMV mode: `1` (default) enables `MoeGemvFp4DecodeKernel` in the measured SM90 decode envelope (`expanded_num_rows <= 20`, `N, K <= 2560`, and per-projection `|global_scale| * 448 < 3.98`); `0` forces `MoeGemvFp4RawKPackedKernel`; `2` opts into the full kernel envelope (`SM90+`, `expanded_num_rows <= 80`, `K <= 8192`, including the generic scale decoder). Outputs are bit-identical across all modes. Read once per process. See [§9b.2](#9b2-fused-gemv-decode-group-16). | | `ORT_FP4_GEMV_INTERLEAVED` | `0` | **Experimental, opt-in.** Routes the MXFP4 decode GEMV through the `ColumnMajorInterleaved` weight layout (`kInterleave=4`, `kStepK=32`) with dtype-conditional accumulation. fp16 gets faster decode; bf16 stays accuracy-safe. Default off keeps the shipping `ColumnMajor` path byte-for-byte unchanged. See [§9.10](#910-interleaved-gemv-layout--dtype-conditional-accumulation). | | `ORT_FP4_GEMV_INTERLEAVED_HALFACC` | `0` | **Override.** When `ORT_FP4_GEMV_INTERLEAVED=1`, forces 16-bit accumulation for *both* fp16 and bf16, overriding the dtype-conditional policy; regresses bf16 accuracy, so it is off by default. | | `ORT_FP4_SM80_GEMM` | `1` | Routes SM80+ (including SM120/SM121) FP4 prefill through the fused-dequant grouped GEMM. Set to `0` to force dense fallback for debugging or comparison. Decode routes through the fused MXFP4 GEMV, reading the *same* pre-packed buffer as prefill. In this regime the raw e2m1 initializers are released after `PrePack`. See [§9.11](#911-single-copy-sm80-weights-prefill--decode-share-one-buffer). | @@ -1124,18 +1146,75 @@ compile-time template (`static_assert((CtaK/kInterleave) % GroupSize == 0)`). NV - `is_moe_gemv_fp4_supported` accepts `group_size ∈ {16, 32}`; the dispatch instantiates the `GroupSize=16` cases in `dispatch_moe_gemv_group_size` / `dispatch_moe_gemv_interleaved_swiglu_group_size` ([moe_gemv_device.cuh](onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh)). -- Prepacked NVFP4 uses the non-interleaved `ColumnMajor` layout; the opt-in interleaved path +- Legacy prepacked NVFP4 uses the non-interleaved `ColumnMajor` layout; the opt-in interleaved path ([§9.10](#910-interleaved-gemv-layout--dtype-conditional-accumulation)) is MXFP4-only because its `kStepK=32` tile is tied to the block-32 scale layout. -- `QMoECombineNvfp4ScalesForGemv` ([qmoe_kernels.cu](onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu)) +- For legacy prepacked NVFP4, `QMoECombineNvfp4ScalesForGemv` ([qmoe_kernels.cu](onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu)) decodes the `float8e4m3fn` block scales, folds in the per-expert FP32 global scale, and rewrites `[E, n, k/16] → [E, k/16, n]` in the activation dtype (`TypeA`) that the GEMV expects. - The decode gate ([moe_quantization.cc](onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc)) fires - when `expanded = num_tokens·top_k ∈ (0, 64]`, SwiGLU is fused, and both FC1 + when `expanded = num_tokens·top_k ∈ (0, 64]` for legacy N-packed storage or + `(0, 80]` for offline K-packed storage, SwiGLU is fused, and both FC1 (`n=2·inter`, `k=hidden`) and FC2 (`n=hidden`, `k=inter`) satisfy `n,k ≥ 512` and group-16 block alignment. The fused expert-map prologue supports `top_k ∈ {1, 2, 4, 6, 8, 10}` and up to 1022 experts. - Qwen Flash (`hidden=2560`, `inter=640`, `top_k=10`) qualifies for one through six tokens. + Qwen Flash (`hidden=2560`, `inter=640`, `top_k=10`) qualifies for one through six tokens + with legacy storage, or one through eight tokens with offline K-packed storage. + +#### Decode kernel for offline K-packed storage + +For `weights_prepacked=1` (row-major `[E, N, K/2]`) the GEMV launcher first tries +`MoeGemvFp4DecodeKernel` ([moe_gemv_fp4.cu](onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu)), +a decode-specialised replacement for `MoeGemvFp4RawKPackedKernel`. Each thread owns two adjacent +output columns and the K-lane pair `(2j, 2j + 1)`, so one 128-bit load covers two 16-element blocks; +several K tiles are kept in flight, FP4 is decoded in registers (table-free for scales below 4), and +the activation row is staged once per CTA in shared memory as FP32. + +The kernel is **bit-identical** to `MoeGemvFp4RawKPackedKernel`: it keeps the same K-lane count +(`KLanes = 4` for `K <= 1024`, else `8`), the same lane assignment (`lane = block % KLanes`), the same +separate even/odd FP32 accumulators per lane, and the same `acc_even + acc_odd` followed by an xor tree +over `KLanes / 2 .. 1`, so outputs do not depend on which of the two kernels runs. + +By default (`ORT_FP4_GEMV_DECODE_KERNEL=1`), dispatch is restricted to the H200-measured decode envelope: +- `SM90` (`sm == 90`) with a selected CUDA image compiled for `SM90` or newer (`attributes.ptxVersion >= 90`; + forward-compatible builds using pre-SM90 PTX such as `CMAKE_CUDA_ARCHITECTURES=80` fall back to + `MoeGemvFp4RawKPackedKernel`), +- FP16 or BF16, at most 20 expanded rows, `K % 32 == 0`, `K <= 2560`, even `N <= 2560`, and 16/2/16-byte aligned + weight/scale/activation pointers, +- Per-projection table-free scale eligibility (`|global_scale| * 448 < 3.98` for every expert in the projection), + recorded independently for FC1 and FC2 during `PrePack`. Projections whose scales exceed this bound stay on + `MoeGemvFp4RawKPackedKernel` by default rather than taking the slower generic `Fp4I2FConverter` branch. +- When `ORT_NVFP4_GEMV_NATIVE_SCALE=1` is also enabled on `SM90` (see [§9b.3](#9b3-fast-e2m1--halfbf16-decode)), + FP16 projections with `K <= 1024` (such as Qwen3.8 FC2, `K=640`) and any projection that stays on + `MoeGemvFp4RawKPackedKernel` use the 4/8-lane native-scale `MoeGemvFp4RawKPackedKernel<..., NativeScale=true>` + (`8.03 us` vs. `8.41 us` for FP16 FC2 at 10 expanded rows, and `17.50 / 9.90 us` vs. `18.54 / 10.30 us` on + large-scale fallback), while remaining bit-identical. +- Steps aside when `ORT_NVFP4_FC1_EXPERT_REUSE=1` selects the two-row reuse kernel. + +Set `ORT_FP4_GEMV_DECODE_KERNEL=0` before the first GEMV invocation (latched once per process) to force +`MoeGemvFp4RawKPackedKernel`, or `ORT_FP4_GEMV_DECODE_KERNEL=2` to opt into the full kernel envelope +(`SM90+`, up to 80 expanded rows, `K <= 8192`, any even `N`, and the generic `Fp4I2FConverter` fallback for +`|global_scale| * 448 >= 3.98`). + +On H200 (`SM90`, CUDA 13.0) for the Qwen3.8 shape (`hidden=2560`, `inter=640`, 512 experts, `top_k=10`, +`weights_prepacked=1`), the FC1 (`N=1280`, `K=2560`, fused SwiGLU) and FC2 (`N=2560`, `K=640`) launches +average (over 300 launches per CUDA graph, 10 graph replays; repeated-replay variation within ±0.15 us): +- BF16, 10 expanded rows: `12.13` and `7.49 us` vs. `18.51` and `10.27 us` (`1.47x` faster FC1+FC2). +- BF16, 20 expanded rows: `20.79` and `12.20 us` vs. `36.05` and `17.23 us` (`1.62x` faster FC1+FC2). +- FP16, 10 expanded rows: `13.77` and `8.41 us` vs. `17.55` and `9.44 us` (`1.22x` faster FC1+FC2). +- FP16, 20 expanded rows: `24.71` and `14.28 us` vs. `34.01` and `16.80 us` (`1.30x` faster FC1+FC2). + +To keep the 512-expert weights (`800 MiB` FC1 + `100 MiB` scales; `400 MiB` FC2 + `50 MiB` scales) cold in +H200's `50 MiB` L2 cache during CUDA-graph timing, consecutive launches in the captured graph advance +`permuted_row_to_expert` (or `router_probs` at the ORT node level) through 64 disjoint routing permutations +covering all 512 experts (`ceil(512 / 10) = 52` launches stream the full `900 MiB` FC1 / `450 MiB` FC2 bank +before any expert repeats). To reproduce the end-to-end ORT QMoE node comparison on H200: + +```bash +python onnxruntime/test/python/transformers/benchmark_qwen38_mtp_kernels.py --baseline-stdin --kernel nvfp4_decode < onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu +ORT_FP4_GEMV_DECODE_KERNEL=0 pytest onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py -k test_nvfp4_offline_row_major_decode_gemv_bit_exact +ORT_FP4_GEMV_DECODE_KERNEL=1 pytest onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py -k test_nvfp4_offline_row_major_decode_gemv_bit_exact +``` #### Raw-layout memory option @@ -1175,6 +1254,44 @@ the raw kernel, so profiling them only adds synchronization and repeated work. W ### 9b.3 Fast E2M1 → half/bf16 decode +For offline K-packed NVFP4 GEMV, `ORT_NVFP4_GEMV_NATIVE_SCALE=1` optionally +uses native E4M3-to-half conversion on compiled SM90+ code for block scales. +It defaults to false. Set it before initializing the runtime; the launch +specializations cache its value. Older compiled architectures retain software +conversion. Existing K-lane selection and optional expert-row reuse are retained. +Global scaling, activation-dtype rounding, weight decoding, accumulation, +routing, and prefill are unchanged. This scalar conversion does not introduce +FP8 activations or native FP8 matrix multiplication. + +CUDA internal-test builds with FP4 QMoE register +`onnxruntime_nvfp4_native_scale_0` and `onnxruntime_nvfp4_native_scale_1` in CTest. +These launch fresh processes with the flag explicitly disabled/enabled, because +the launcher caches it. The native-specific test skips unless FP8 helpers and the +loaded kernel's binary and PTX compilation targets support SM90+; a newer device +alone is insufficient. Software/fallback numerical coverage remains separate. +The all-scale-code regression exercises FP16/BF16 with K=16, 80, and 1040. +Nonzero contributions in every K block, varying scales, and an independent +full-reduction reference cover all four/eight K lanes and later iterations. + +CUDA plugin internal-test builds compile the FP4 launcher implementation into +their test implementation target, with relaxed constexpr enabled for its CUTLASS +headers. The standalone kernel benchmark preserves device-function templates and +extracts only the global kernel body, excluding host-only inspection helpers. +Its source-generation unit tests run without CuPy or GPU execution. + +For native K-packed NVFP4 verification, `ORT_NVFP4_FC1_EXPERT_REUSE=1` enables +experimental two-row FC1 weight and scale reuse. Set it before the first GEMV +invocation; it is read once per launcher specialization. The dispatch is limited +to fused SwiGLU with 1280 output columns, reduction size 2560, at least 40 expanded +rows, and available sorted expert routing maps. Within that path, only expert +groups with at least four rows reuse a decoded tile; other groups compute rows +individually. Stored FP4 codes/scales and activation-dtype rounding are unchanged. + +The default is `0`: local A100 tests improved repeated-expert groups but regressed +sparse expert distributions. Enable it only after measuring the target model's +expert overlap and end-to-end latency on the target GPU. It does not enable native +FP4 tensor-core computation or dequantize the checkpoint into a dense weight bank. + Profiling the NVFP4 decode GEMV on H200 (SM90) at the Qwen shapes showed the FC1/FC2 kernels are **ALU-pipeline bound** (ncu: ALU ≈ 79%, DRAM ≈ 7%), i.e. the per-element FP4 dequantization — not memory or tiling — dominates. The original `Fp4I2FConverter::decode` used a `float` lookup table, @@ -1423,7 +1540,7 @@ GroupedGemm with EpilogueOpDefault: QMoE op only needs to construct `QuantParams::FP8(dequant_fc1, nullptr, dequant_fc2)` from the per-expert global scales (inputs 15/16). -### 10.3 Dequant fallback (SM<90 or square-block scales) +### 10.3 Block-scaled FP8 dispatch and dequant fallback `LaunchQMoEDequantizeFp8Weights` decodes weights into BF16/FP16 and the dense A16 runner is used. @@ -1440,9 +1557,16 @@ Positive `block_size=B` selects square-block FP8, including the official CUDA converts the small scale grids to float32 on the execution stream. - Inputs 15/16 and all zero points must be omitted. Nonpositive `block_size` retains the legacy per-expert global-scale contract and native dispatch. -- Block FP8 always uses dequantization plus the dense FP16/BF16 runner, even on - H200. The native FP8 epilogue supports only one scale per expert and cannot - correctly apply these two-dimensional blocks. +- On SM80-or-newer GPUs, SiLU/SwiGLU block FP8 uses fused weight-only GEMV/GEMM + by default for supported routing sizes. These kernels decode FP8 tiles and + apply square-block scales on chip, without full FP16/BF16 weight buffers. + SiLU without FC3 retains FC1 accumulators in FP32 through bias and activation, + then narrows to the activation dtype, matching the dense epilogue. Gated + projections retain their existing FP16/BF16 intermediate rounding. + Set `ORT_ENABLE_FP8_FUSED=0` before session creation to select the compact + dense-dequant fallback. Other activation types and unsupported routing sizes + also retain this fallback. The legacy native FP8 epilogue supports only one + scale per expert and cannot apply these two-dimensional blocks. - Separate FC1 gate / FC3 up projections are supported with `activation_type="silu"` and `swiglu_fusion=0`. FC1/FC3 biases are currently rejected in this mode. Fused SwiGLU accepts interleaved (`1`) or contiguous gate/up halves (`2`); @@ -1451,7 +1575,7 @@ Positive `block_size=B` selects square-block FP8, including the official dense CUDA runner. They need not be divisible by `B`; partial scale blocks work. The official shape is BF16, hidden size 2560, intermediate size 640, 512 experts, -top-10 routing, and SiLU with separate gate/up projections. Block FP8 routes +top-10 routing, and SiLU with separate gate/up projections. The dense fallback routes before dequantizing, compacts the selected expert IDs on the GPU, and decodes each selected expert once per row tile. Repeated selections share one dense slot. Biases follow the same mapping; routing records retain original expert IDs. @@ -1468,7 +1592,7 @@ The existing `ep.cuda.qmoe_row_tile_size` session option also bounds FP8 weight scratch during prefill. Buffers are reused on the same stream between tiles; experts selected in multiple tiles are decoded again. Without tiling, a large prefill may still reserve all `E` slots. Legacy global-scale FP8 dispatch is -unchanged. This remains a dense-GEMM fallback, not native block-FP8 execution. +unchanged. This fallback remains dense-GEMM execution, not the fused block-FP8 path. Router logits use the existing softmax/top-k semantics and `normalize_routing_weights` controls renormalization. @@ -1515,6 +1639,36 @@ Un-permute → weighted sum → output --- +### 10.6 Optional native E4M3 GEMV decoding + +`ORT_QMOE_FP8_GEMV_NATIVE_DECODE=1` enables native E4M3-to-half conversion for +block-scaled FP8 weights in the fused FP16/BF16 GEMV path on compiled SM90+ +code. The flag defaults to false. Set it before initializing the runtime; +launch specializations cache its value. Older compiled architectures retain +software conversion. GEMM/prefill continues to use the existing decoder. + +CUDA internal-test builds register `onnxruntime_fp8_native_decode_0` and +`onnxruntime_fp8_native_decode_1` in CTest. Each flag value runs in a fresh process. +The graph cases inspect the captured kernel function to verify the requested +template specialization, and native-specific coverage additionally checks the +FP8 build gate and the loaded kernel's binary/PTX targets before running. It skips +explicitly when native instructions are unavailable instead of claiming fallback +execution as native coverage. +All K elements contribute to the numerical reference; the three K blocks have +different scales and the partial tail has a significant nonzero activation. +This detects dropped blocks, missing tails, and reuse of the first block's scale. + +This changes only scalar byte decoding. Block scaling, activation-dtype +rounding, accumulation order, row mappings, and gate/up layouts are retained. +Finite E4M3FN values are exactly representable in FP16, so this conversion does +not add quantization. NaN propagation is preserved; NaN-payload equivalence +and full-model output identity are not promised. + +CUDA kernel regressions cover all 256 bytes through GEMV with mapped rows, +separate and fused gate/up layouts, 16- and 128-element scale blocks, partial +K blocks, FP32/FP16/BF16 scale tensors, and CUDA graph capture/replay. +Run them in separate processes with the native flag disabled and enabled. + ## 11. WFP4AFP8 Details `quant_type="wfp4afp8"` pairs MXFP4 weights with FP8 e4m3 activations. Unlike diff --git a/docs/contrib_ops/webgpu/matmul_block_quantized_fp8_weight.md b/docs/contrib_ops/webgpu/matmul_block_quantized_fp8_weight.md new file mode 100644 index 0000000000000..9a728f17edfed --- /dev/null +++ b/docs/contrib_ops/webgpu/matmul_block_quantized_fp8_weight.md @@ -0,0 +1,27 @@ +# WebGPU block-quantized FP8 dense MatMul + +`com.microsoft::MatMulBlockQuantizedFp8Weight` supports FP16 activations and +row-major E4M3FN weights with FP32 block scales shaped +`[N, ceil(K / block_size)]`. It supports optional activation quantize/dequantize, +bias, rank-one and batched inputs, partial blocks, and non-word-aligned weight rows. +BF16 activation storage is not supported by this WebGPU kernel. + +When activation scaling is supplied, quantization uses saturating E4M3FN +round-to-nearest, ties-to-even semantics. NaN activations remain NaN rather than +being replaced by the maximum finite value. Weight decoding preserves finite +E4M3FN subnormals before multiplication. + +Prefill can use 16x16x16 FP16 subgroup matrices with FP32 accumulation when the +adapter exposes the required configuration, subgroup-size requirements, and the +instance exposes the WGSL `subgroup_id` language feature. Other configurations +use the scalar FP32-reduction path. This is FP8 storage with FP16 widening, not +native FP8 matrix multiplication. + +Shared shader generation includes `@builtin(subgroup_id)` only for programs +that explicitly request it. Both dense FP8 and routed block-FP8 QMoE matrix +programs check the language feature before selecting that path. Ordinary +operators and scalar fallbacks do not acquire this language-feature requirement. + +The focused `MatMulBlockQuantizedFp8WeightTest.*` tests disable CPU EP fallback. +Matrix-specific cases check capabilities and skip explicitly on unsupported +adapters; scalar cases remain separate. diff --git a/docs/design/webgpu_matmul_execution_planner.md b/docs/design/webgpu_matmul_execution_planner.md new file mode 100644 index 0000000000000..23935af9f5c32 --- /dev/null +++ b/docs/design/webgpu_matmul_execution_planner.md @@ -0,0 +1,94 @@ +# WebGPU MatMul Execution Planner Design + +## Goal + +Make every `MatMulComputeDispatcher` implementation path explicit and independently testable. Selection policy must be separated from execution, retain the existing common rules as a fallback, and allow vendor-specific policy to override any performance heuristic. + +## Algorithms + +Introduce `MatMulAlgorithm` with six concrete values: + +- `SubgroupMatrix`: the common subgroup-matrix implementation in `subgroup_matrix_matmul.cc`. +- `Gemv`: the vendor-neutral FP16 specialization for packed M dimensions 1–8 in `MatMulGemvProgram`. +- `Naive`: `MatMulNaiveProgram`. +- `Subgroup`: the common subgroup implementation currently located in `vendor/intel/math/matmul.cc`. +- `Packed`: the generic packed `MatMulProgram` without Split-K. +- `PackedSplitK`: the same packed program with Split-K initialization and atomic accumulation. + +There is no `Auto` algorithm value. Automatic versus forced selection is represented by an optional forced algorithm, which prevents an unset policy state from being confused with an executable implementation. + +## Selection Architecture + +Add a `MatMulExecutionPlanner` base class. Automatic selection uses this order: + +1. Select `Naive` when `K == 0`. This is a correctness rule and cannot be overridden by vendor policy. +2. Ask the virtual `SelectVendorAlgorithm` policy hook for any algorithm. A vendor may use its own thresholds for every algorithm, not only vendor-specific implementations. +3. If the vendor returns no selection, call the private, non-virtual `SelectCommonAlgorithm` fallback. It selects `SubgroupMatrix` when supported, then `Gemv` for its compatible FP16 shape, `Naive` when `N < 8 && K < 8`, `PackedSplitK` when the existing `SplitKConfig::UseSplitK` rule succeeds, and otherwise `Packed`. + +The `Gemv` implementation requires packed batch size 1, packed `M` in the range 1–8, nonzero `K`, FP16 inputs and output, `N` divisible by four, no bias, and no fused activation. For batched `A` with unbatched `B`, `A`'s batch dimensions are folded into packed `M`. Automatic policy additionally limits GEMV to `2048 <= K <= 8192` and `16 <= N <= 64`; those ranges are measured performance thresholds rather than implementation prerequisites. Forced selection bypasses the performance thresholds while retaining the implementation constraints. The planner keeps `SubgroupMatrix` ahead of `Gemv`, preserving the precedence that existed when GEMV was introduced. + +An Intel-derived planner implements the vendor hook. It preserves the original policy by selecting `SubgroupMatrix` first when applicable, then selecting the common `Subgroup` implementation under Intel-specific shape thresholds. Other vendors use the base planner unchanged. Future vendor policies can derive from the base planner, override as many performance ranges as needed, and return no selection to delegate the remaining ranges to the common fallback without adding vendor conditionals to `MatMulComputeDispatcher`. + +The subgroup implementation is not inherently Intel-specific: forcing it requires the WebGPU `Subgroups` feature and a supported subgroup size. Its current source location under `vendor/intel` is historical and should move to the common math directory in a follow-up, while Intel-specific automatic-selection thresholds and tuning remain under `vendor/intel`. + +The subgroup shader implements subgroup-size branches for 8, 16, and 32 lanes. A fixed-size adapter may use any of those sizes directly. An adapter reporting a size range must expose subgroup-size control so the execution plan can select a supported size explicitly; otherwise the subgroup algorithm is unavailable. + +The planner accepts immutable problem facts rather than mutable policy decisions: logical and packed dimensions, batch sizes, adapter architecture, input types, packing and layout facts, deterministic-compute state, activation and bias facts, and device capabilities. It also owns a copy of the immutable `SplitKConfig` selected for the adapter. The private common fallback evaluates `SplitKConfig::UseSplitK` directly from the problem facts instead of receiving a precomputed decision from `MatMulComputeDispatcher`. A vendor may ignore that common recommendation and apply independent thresholds from the raw facts. This keeps the planner deterministic and unit-testable without a WebGPU device. + +`SplitKConfig` contains only generic Split-K eligibility evaluation and data. Adapter routing is handled by a small generic factory, while Intel architecture profiles and their measured threshold tables live under `vendor/intel`. `WebGpuContext` owns the selected configuration so GEMM and MatMul use the same profile; the MatMul planner receives that configuration when it is created. A future vendor can add its own profile builder and factory route without adding conditions to `MatMulComputeDispatcher` or changing the generic evaluator. + +## Execution Configuration + +Algorithm selection and execution tuning are separate decisions. `SelectAlgorithm` returns only the selected enum and is useful when policy is the caller's concern. `CreateExecutionPlan` selects the enum and its typed configuration together, first asking the protected `SelectVendorConfiguration` hook for tuning and then using private common defaults when the vendor declines. It validates configuration compatibility, host-known algorithm prerequisites, and dispatch arithmetic before publishing the output plan. A failed request leaves the caller's plan unchanged. The tuning hook runs for both automatic and forced algorithms, so the test-only factory options control the implementation path without disabling real device tuning. + +The packed configuration initially contains workgroup size, elements per thread, inner tile size, and Split-K size. `ApplyMatMulPacked` consumes those values directly and includes shader-affecting values in its cache key. The common configuration preserves the existing `8x8x1` workgroup, `4x1x1` or `4x4x1` elements-per-thread rule, inner tile size 32, and adapter Split-K size. A vendor may replace any of these values without changing `MatMulComputeDispatcher` or the packed implementation. + +Packed tuning is validated while the execution plan is created. The current shader requires both Z workgroup dimensions to remain one because Z identifies a batch or Split-K slice rather than a tiled output axis. Split-K sizes must be greater than one and aligned to the inner tile so adjacent workgroups cannot overlap their K ranges. Dispatch counts use widened, overflow-safe arithmetic and are range-checked before conversion to WebGPU's 32-bit dimensions. + +Configuration is represented by an algorithm-specific variant rather than a bag of unrelated optional fields. Plan creation rejects an algorithm/configuration type mismatch before returning success. Empty configuration types reserve the same typed boundary for algorithms whose tuning remains inside their existing implementation. In particular, subgroup-matrix MatMul already receives a vendor-specific `SubgroupMatrixTilingSelector`; that existing selector continues to choose tile M, tile N, and split K without coupling the planner to device or shader classes. + +## Forced Test Selection + +`WebGpuExecutionProviderTestOptions` carries an optional `MatMulAlgorithm` through the internal `CreateForTesting` factory path. Production provider-option parsing cannot populate it. The selected value is stored in the execution provider and exposed read-only through `ComputeContextBase`. + +When set, the planner chooses the requested enum before applying heuristic rules and validates its hard prerequisites before returning the plan. Unsupported device features, data types, layouts, deterministic-compute settings, or other correctness constraints produce a descriptive failure naming the forced algorithm; forced mode never silently falls back. + +Heuristic thresholds are not hard prerequisites. For example, forcing subgroup bypasses Intel's current `M/N/K` performance thresholds while still requiring subgroup support. Forcing GEMV bypasses its automatic `N/K` tuning window while retaining its packed batch/M, FP16, vector-width, bias, and activation constraints. Forcing Split-K bypasses performance thresholds while still requiring a usable Split-K configuration, non-deterministic compute, compatible packing/activation, and supported bias layout. + +## Dispatch and Implementation Boundaries + +Introduce `MatMulComputeDispatcher` as the single compute entry point below the MatMul, pointwise Conv, and contrib Attention kernels. Each kernel owns one dispatcher for its lifetime. The three concepts have distinct responsibilities: + +- `MatMulExecutionPlanner` is a pure policy object. It consumes immutable invocation facts and creates a host-validated plan without running device work. +- `MatMulExecutionPlan` is an immutable per-invocation value containing the selected algorithm and its typed tuning configuration. +- `MatMulComputeDispatcher` gathers runtime facts, asks the planner for a plan, validates live device resources, and invokes the selected implementation. + +The dispatcher lazily creates adapter-dependent state from the first compute context and then owns: + +- one `MatMulExecutionPlanner`, which contains selection policy and immutable adapter tuning data; and +- an optional `SubgroupMatrixMatMulImpl`, which contains only the subgroup-matrix implementation's persistent device state, including cached padded constant weights. + +The dispatcher does not store a current or previously selected algorithm. Selection occurs for every invocation because shapes, input properties, activation, and bias can differ between calls. The session's forced-test configuration, when present, participates in each selection without becoming mutable dispatcher state. The resulting `MatMulExecutionPlan` is local to that invocation. + +For each call, `MatMulComputeDispatcher::Compute` computes shared shape and capability facts once, asks the planner for one validated execution plan, and switches directly to the matching implementation. MatMul, pointwise Conv, and contrib Attention delegate through this same entry point rather than coordinating the planner and implementations themselves. + +The planner remains pure policy: it creates plans but does not create device programs, cache tensor data, or execute kernels. `SubgroupMatrixMatMulImpl` is named for the implementation it owns and is reached only when the plan selects `SubgroupMatrix`; its applicability query is non-mutating, and its execution method does not communicate selection through a `handled` output. This removes trial execution as a dispatch mechanism. + +The other implementations remain focused stateless functions or program builders unless they acquire persistent state in the future. The naive, subgroup, packed, and packed Split-K paths therefore do not receive empty polymorphic wrapper classes. The generic packed helper takes an explicit Split-K mode and packed configuration; it does not re-run selection or tuning heuristics. WebGPU's existing program cache continues to own reusable compiled programs. + +This replaces `MatMulOptImplCache` and the generic `MatMulOptImpl` interface. If another algorithm later needs persistent implementation state, the dispatcher can own a separately named backend for that algorithm without changing planner policy or pretending that one object represents whichever algorithm happened to be selected most recently. + +## Compatibility + +With no forcing option or vendor override, the common planner remains equivalent to the previous selection order. The MatMul, pointwise Conv, and contrib Attention operator APIs and model semantics do not change; only their internal MatMul compute ownership moves behind the dispatcher. + +Forced selection is available only through the internal test factory and cannot be set through session or provider options. + +## Testing + +- Add device-independent planner unit tests covering every common branch, forced-over-vendor precedence, the zero-K correctness guard, vendor-over-common precedence, Intel override, default fallback, GEMV precedence, independent vendor thresholds, common packed defaults, vendor tuning of a forced algorithm, Split-K profile routing, and current Intel architecture boundaries. +- Add typed configuration tests for every algorithm. +- Add WebGPU MatMul tests that choose shapes which normally select a different path, force a compatible algorithm, and verify numerical output. Hardware-specific forced algorithms are tested only when their hard capabilities are present; strict-failure tests cover unsupported forced choices. +- Run hardware-backed tests on the build's default WebGPU backend. Hardware-specific tests inspect the selected adapter's capabilities and skip unsupported algorithms. +- Verify the selected adapter exposes subgroup size control, f16, and the cooperative/subgroup-matrix configuration required by the 8x16x16 kernel before claiming subgroup-matrix execution coverage. +- Run planner tests and hardware-backed MatMul tests on the local Intel Arc Vulkan adapter. macOS-arm64 Metal CI remains additional cross-backend coverage; lavapipe is not used because it cannot execute MatMul reliably. diff --git a/docs/design/webgpu_paged_attention.md b/docs/design/webgpu_paged_attention.md index 423cdab58d4d6..8a3405f2f5e77 100644 --- a/docs/design/webgpu_paged_attention.md +++ b/docs/design/webgpu_paged_attention.md @@ -190,7 +190,7 @@ split-reduce threshold): `ApplyFlashAttention` consult one shared predicate: ```cpp -bool ShouldRunFusedPagedPrefill(context, is_fp16, max_seqlen_q, +bool CanApplyFlashAttentionPagedPrefill(context, is_fp16, max_seqlen_q, head_size, block_size); ``` @@ -230,7 +230,7 @@ a padded BSNH scratch: `[token_count, 1, N, H]` view plus `cumulative_seqlens_q`; the fused paged-prefill shader reads Q rows through `q_row = cumulative_seqlens_q[b] + q_idx` (`q_varlen` template variant). Only safe when - `ShouldRunFusedPagedPrefill` returns true. + `CanApplyFlashAttentionPagedPrefill` returns true. Skipping unpack + repack removes two dispatches (~300–500 µs of CPU dispatch cost per Run on D3D12) plus the padded scratch allocation @@ -242,7 +242,7 @@ Two full paged K/V reads plus two contiguous K/V writes per layer per step (gather), Q/output shuffles proportional to `token_count`, and a scratch allocation of `2 * B * kv_num_heads * max_kv_len * head_size + B * max_seqlen_q * hidden_size` bytes. Fallback is exercised on configurations -that `ShouldRunFusedPagedPrefill` rejects (fp32, `head_size > 256`, or +that `CanApplyFlashAttentionPagedPrefill` rejects (fp32, `head_size > 256`, or `block_size < max_k_step`); direct paged paths cover the common case on every WebGPU adapter. @@ -355,7 +355,7 @@ preserving the single-kernel prefill algorithm and workgroup Q tiling of `FlashAttentionProgram`. It supports fp16, both BSNH Q and packed varlen Q (`q_varlen` template variant), variable-Q-length causal masking via `seqlen_k` + `seqlens_q`, and no attention bias / head sink / TurboQuant. -`ShouldRunFusedPagedPrefill` (see §4.3) is the shared selection gate. +`CanApplyFlashAttentionPagedPrefill` (see §4.3) is the shared selection gate. Measured wins on the microbench matrix vs. the #31611 gather-then-flash cascade: **1.01×–1.25×** on uniform prefill (geomean ~1.13×) and **1.14×–1.73×** on varlen prefill (geomean ~1.29×). @@ -453,7 +453,7 @@ onnxruntime/contrib_ops/webgpu/bert/ The direct paged decode / fused paged prefill programs are declared in `flash_attention.h` alongside the dense-FA programs so they can share the same `ApplyFlashAttention` dispatch surface, adapter cache hints, and -`ShouldRunFusedPagedPrefill` selection helper. +`CanApplyFlashAttentionPagedPrefill` selection helper. Shared helper (already exists on CUDA, refactor location in Phase 1): @@ -503,7 +503,7 @@ ComputeInternal: # Route selection: use_direct_paged_decode = (max_seqlen_q < 32) - use_direct_paged_prefill = ShouldRunFusedPagedPrefill(context, is_fp16, + use_direct_paged_prefill = CanApplyFlashAttentionPagedPrefill(context, is_fp16, max_seqlen_q, head_size, block_size) use_direct_paged_attention = use_direct_paged_decode or use_direct_paged_prefill @@ -529,7 +529,7 @@ ComputeInternal: ApplyFlashAttention() on paged cache directly (no gather) RunRepackOutput() - # (b) Fallback (rejected by ShouldRunFusedPagedPrefill, e.g. fp32, + # (b) Fallback (rejected by CanApplyFlashAttentionPagedPrefill, e.g. fp32, # head_size > 256, or block_size < max_k_step): else: RunGatherKV() # -> [B, kv_num_heads, max_kv_len, head_size] @@ -571,7 +571,7 @@ Feature guards (v1 rejects with `NOT_IMPLEMENTED` and a specific message): | WebGPU graph capture forbids host-visible reads mid-graph. | Consume `attention_metadata` (Phase 2) instead of D→H syncing `cumulative_seqlens_q`. | | Subgroup width varies (16 Intel, 32 NV/AMD, 64 Qualcomm/Apple). | For the dense `FlashAttentionProgram` the `is_qualcomm`/`is_nvidia`/`is_apple`/`has_subgroups` cache-hint knobs select between shm and subgroup-shuffle variants. The paged prefill / paged decode shaders use no subgroup intrinsics (they run on every adapter that satisfies the dtype/shm-budget/alignment gates), so they do not need those knobs. | | WGSL forbids mixed `i32 + u32` arithmetic. `cumulative_seqlens_q` is stored as `array` but the row index is `u32`; tint reports the resolution failure as an opaque `absl::…raw_hash_map<>::at` at runtime. | Cast the storage-buffer element to `u32` before the addition (`u32(cumulative_seqlens_q[b]) + q_idx`). Applies in both the fused prefill shader's `loadq` and `writeo`. | -| Fused paged prefill assumes one K/V tile lives entirely inside a single paged block, so both `loadk` and `loadv` do one `block_table` lookup per tile. `paged_attention_helper` only enforces `block_size >= 16` (power-of-two), so `block_size = 16` combined with fp16 `head_size <= 128` (`max_k_step = 32`) would splice into a physically-adjacent block that need not be the next entry in `block_table`. | `ShouldRunFusedPagedPrefill` rejects `block_size < max_k_step`; callers fall back to gather-then-flash. Both values are powers of two ≥ 16, so `>=` implies alignment. | +| Fused paged prefill assumes one K/V tile lives entirely inside a single paged block, so both `loadk` and `loadv` do one `block_table` lookup per tile. `paged_attention_helper` only enforces `block_size >= 16` (power-of-two), so `block_size = 16` combined with fp16 `head_size <= 128` (`max_k_step = 32`) would splice into a physically-adjacent block that need not be the next entry in `block_table`. | `CanApplyFlashAttentionPagedPrefill` rejects `block_size < max_k_step`; callers fall back to gather-then-flash. Both values are powers of two ≥ 16, so `>=` implies alignment. | --- diff --git a/docs/python/examples/plot_pipeline.py b/docs/python/examples/plot_pipeline.py index e0bae57b7e2c8..c73d3ba30c211 100644 --- a/docs/python/examples/plot_pipeline.py +++ b/docs/python/examples/plot_pipeline.py @@ -29,12 +29,10 @@ ################################# -# Draw a model with ONNX +# Draw a model with pydot # ++++++++++++++++++++++ -# We use `net_drawer.py `_ -# included in *onnx* package. -# We use *onnx* to load the model -# in a different way than before. +# We use *onnx* to load the model in a different way than before, then +# construct a pydot graph from its nodes and values. from onnx import ModelProto # noqa: E402 @@ -46,18 +44,29 @@ ################################### # We convert it into a graph. -from onnx.tools.net_drawer import GetOpNodeProducer, GetPydotGraph # noqa: E402 +import pydot # noqa: E402 + +pydot_graph = pydot.Dot(model.graph.name, rankdir="LR") +value_nodes = {} +for index, node in enumerate(model.graph.node): + node_name = node.name or f"{node.op_type}_{index}" + op_node = pydot.Node(node_name, label=node.op_type, shape="box") + pydot_graph.add_node(op_node) + for value_name in node.input: + value_node = value_nodes.setdefault(value_name, pydot.Node(value_name, shape="ellipse")) + pydot_graph.add_node(value_node) + pydot_graph.add_edge(pydot.Edge(value_node, op_node)) + for value_name in node.output: + value_node = pydot.Node(value_name, shape="ellipse") + value_nodes[value_name] = value_node + pydot_graph.add_node(value_node) + pydot_graph.add_edge(pydot.Edge(op_node, value_node)) -pydot_graph = GetPydotGraph( - model.graph, name=model.graph.name, rankdir="LR", node_producer=GetOpNodeProducer("docstring") -) pydot_graph.write_dot("graph.dot") ####################################### # Then into an image -import os # noqa: E402 - -os.system("dot -O -Tpng graph.dot") +pydot_graph.write_png("graph.dot.png") ################################ # Which we display... diff --git a/include/onnxruntime/core/framework/resource_accountant.h b/include/onnxruntime/core/framework/resource_accountant.h index 99bea37b982b9..c7a80839b7221 100644 --- a/include/onnxruntime/core/framework/resource_accountant.h +++ b/include/onnxruntime/core/framework/resource_accountant.h @@ -67,11 +67,15 @@ struct WorkspaceEstimateComparisonSummary { struct WorkspaceEstimatorConfig { std::optional cuda_fpa_intb_gemm; std::optional cuda_fpa_intb_profile_m; - // Validated nonnegative int64 envelope; zero means unspecified. - // Reader-only until the GQA workspace estimator consumes the envelope. + // Validated nonnegative, int32-compatible envelope; zero means unspecified. int64_t cuda_gqa_workspace_max_total_sequence_length = 0; }; +// Reads the caller-declared CUDA GQA envelope and validates its runtime ABI range. +Status ParseCudaGqaWorkspaceMaxTotalSequenceLength( + const ConfigOptions& config_options, + int64_t& max_total_sequence_length); + using NodeWorkspaceReservationMap = InlinedHashMap; using WorkspaceReservationMap = InlinedHashMap; diff --git a/include/onnxruntime/core/graph/schema_registry.h b/include/onnxruntime/core/graph/schema_registry.h index ca51e3621b2c6..f60589a6d82e0 100644 --- a/include/onnxruntime/core/graph/schema_registry.h +++ b/include/onnxruntime/core/graph/schema_registry.h @@ -37,7 +37,7 @@ class IOnnxRuntimeOpSchemaCollection : public ONNX_NAMESPACE::ISchemaRegistry { using ISchemaRegistry::GetSchema; - const ONNX_NAMESPACE::OpSchema* GetSchema(const std::string& key, const int maxInclusiveVersion, + const ONNX_NAMESPACE::OpSchema* GetSchema(const std::string& key, int maxInclusiveVersion, const std::string& domain) const final { const ONNX_NAMESPACE::OpSchema* latest_schema = nullptr; int earliest_opset_where_unchanged = std::numeric_limits::max(); diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 4bec5ed2669c7..c38ecd3a6794f 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -224,6 +224,11 @@ typedef enum ONNXTensorElementDataType { ONNX_TENSOR_ELEMENT_DATA_TYPE_INT2, // maps to 4 packed int2 values (size == 1 byte) // Float8E8M0 type introduced in ONNX 1.21. 8-bit float with 8 exponent bits, 0 mantissa bits, no sign bit. ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E8M0, // Non-IEEE floating-point format, all values are powers of two + // Float6 types introduced in ONNX 1.23. + // ORT tensor storage uses one byte per element. TensorProto typed serialization uses int32_data; + // TensorProto raw_data uses bit-packed 6-bit values. + ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E2M3, + ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E3M2, } ONNXTensorElementDataType; // Synced with onnx TypeProto oneof @@ -7580,14 +7585,16 @@ struct OrtApi { /** \brief Retrieve an experimental function pointer by name. * * Experimental functions are not part of the stable ABI and may be added or removed between releases without notice. - * Use the companion header onnxruntime_experimental_c_api.h for typedefs, name constants, and (for C++) typed - * accessors. + * + * Use the companion headers: + * - onnxruntime_experimental_c_api.h for function pointer typedefs and name constants. + * - onnxruntime_experimental_cxx_api.h for C++ typed accessors and wrapper types. * * \param[in] name The null-terminated name of the experimental function to look up. * Names follow the pattern "__SinceV". * Name constants are defined in onnxruntime_experimental_c_api.h. - * \return The function pointer cast to ::OrtExperimentalFnPtr, or nullptr if the function is not available in this - * build. The caller must cast the returned pointer to the correct function pointer type before calling. + * \return The function pointer cast to ::OrtExperimentalFnPtr, or NULL if the function is not available. + * The caller must cast the returned pointer to the correct function pointer type before calling. * Function pointer typedefs are defined in onnxruntime_experimental_c_api.h. * * \since Version 1.28. diff --git a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h index 77a1b7c0f1683..663f8e4d3a0df 100644 --- a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h @@ -473,6 +473,8 @@ typedef struct OrtEpProfilerImpl OrtEpProfilerImpl; * - Concurrency notes: each OrtEpProfilerImpl only receives calls for its specific run. * OrtEpProfilerImpl::EndProfiling must only return EP events for its specific run. * + * Some callbacks are optional. Pointers to unimplemented callbacks must be set to NULL. + * * \since Version 1.25. */ struct OrtEpProfilerImpl { @@ -725,6 +727,9 @@ typedef struct OrtKernelImpl OrtKernelImpl; /** * \brief Contains functions that an OrtEp implements to specify the computation for an operator kernel. + * + * Some callbacks are optional. Pointers to unimplemented callbacks must be set to NULL. + * * \since Version 1.24. */ struct OrtKernelImpl { @@ -2215,6 +2220,9 @@ typedef enum OrtWeightlessSupport { /** * \brief The OrtEp struct provides functions to implement for an execution provider. + * + * Some callbacks are optional. Pointers to unimplemented callbacks must be set to NULL. + * * \since Version 1.22. */ struct OrtEp { @@ -2837,6 +2845,9 @@ typedef OrtStatus* (*ReleaseEpApiFactoryFn)(_In_ OrtEpFactory* factory); /** * \brief The OrtEpFactory provides functions to create and manage execution providers. + * + * Some callbacks are optional. Pointers to unimplemented callbacks must be set to NULL. + * * \since Version 1.22. */ struct OrtEpFactory { diff --git a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.h b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.h index 10dbec14e533f..f4b55d6d652fa 100644 --- a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.h @@ -8,13 +8,15 @@ // // The function pointer typedefs and name constants should be used together with the experimental API lookup function // `OrtApi::GetExperimentalFunction()`. A function's availability should always be checked at runtime (the lookup -// returns nullptr if the function is not present). +// returns NULL if the function is not present). // // C++ API consumers should use the companion header, onnxruntime_experimental_cxx_api.h, which provides typed // accessors for the experimental API functions. // -// IMPORTANT: Experimental functions are NOT part of the stable ABI. They may be added, changed, or removed between -// releases without notice. Anything in this file should be treated as experimental and unstable. +// IMPORTANT: Experimental functions are NOT part of the stable ABI. They may be added or removed between releases +// without notice. If an experimental function with a given name is available across releases, its signature does not +// change, and its ABI-relevant types and behavioral contract remain backward compatible. A signature change or an +// incompatible ABI or behavioral-contract change requires a new name. // // C usage: // OrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_Fn fn = diff --git a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc index 0123b02818584..8ac61f7334f50 100644 --- a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc +++ b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc @@ -21,8 +21,9 @@ // Experimental function lifecycle rules: // - Adding: Add one ORT_EXPERIMENTAL_API line with the current API version and implement it. // - Removing: Delete the line. The versioned name is inherently unique and cannot be reused. -// - Changing a signature: Add a new entry at the current API version (producing a new name) and optionally delete -// the old entry. Both can coexist during a transition period. +// - Updating: Backward-compatible fixes and behavioral extensions may retain the existing name. For a signature +// change or an incompatible ABI or behavioral-contract change, add a new entry at the current API version +// (producing a new name) and optionally delete the old entry. Both can coexist during a transition period. // - Promoting to stable: Move the function to the stable API struct (append-only, drop the _SinceV // suffix). Delete the experimental entry. diff --git a/include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h index 2919b4ddeba09..b2d9e36e9d514 100644 --- a/include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h @@ -8,8 +8,10 @@ // // It is the C++ companion to onnxruntime_experimental_c_api.h. // -// IMPORTANT: Experimental functions are NOT part of the stable ABI. They may be added, changed, or removed between -// releases without notice. Anything in this file should be treated as experimental and unstable. +// IMPORTANT: Experimental functions are NOT part of the stable ABI. They may be added or removed between releases +// without notice. If an experimental function with a given name is available across releases, its signature does not +// change, and its ABI-relevant types and behavioral contract remain backward compatible. A signature change or an +// incompatible ABI or behavioral-contract change requires a new name. // // Two accessor flavors are generated for each experimental function: // diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 02537bb643fe7..1251fae691b7c 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -490,13 +490,16 @@ static const char* const kOrtSessionOptionsCudaFpAIntBGemvPairedK = "ep.cuda.fpa /// Opt in with "1" to wave-aware fp16/int4 M=8 GEMV tiles on sm_12x; "0" or unset keeps the default dispatch. static const char* const kOrtSessionOptionsCudaFpAIntBGemvWaveAware = "ep.cuda.fpa_intb_gemv_wave_aware"; -/// Reserved total KV-length envelope (accumulated past + current tokens) for CUDA -/// GroupQueryAttention workspace estimation. Currently reader-only: the value is validated and -/// forwarded to the Level-1 estimator but is not consumed, so it does not change workspace -/// estimates or partitioning. Non-windowed estimation remains unavailable. -/// A nonnegative decimal int64 is required; "0" or unset (default) means unspecified. -/// Negative, malformed, or overflowing explicit values cause INVALID_ARGUMENT when creating -/// resource accountants. This is not a runtime-enforced input limit or a no-OOM guarantee. +/// Declared total KV-length envelope (accumulated past + current tokens) for CUDA +/// GroupQueryAttention workspace estimation. A positive value enables non-windowed Level-1 +/// estimation using this bound; "0" or unset (default) leaves non-windowed estimation unavailable. +/// The estimate also covers the allocated past cache capacity if it exceeds this envelope. +/// Windowed estimation continues to use the cache capacity instead. +/// A decimal integer in [0, INT32_MAX] is required. Invalid or out-of-ABI values cause +/// INVALID_ARGUMENT during session initialization. Under capacity-aware partitioning, CUDA +/// placement is declined when no proven GQA estimate is available; the generic workspace fallback +/// is not used as a proof bound. This caller-declared envelope is not a runtime-enforced input limit +/// or a no-OOM guarantee. static const char* const kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength = "ep.cuda.gqa_workspace_max_total_sequence_length"; diff --git a/include/onnxruntime/core/session/onnxruntime_type_conversion.h b/include/onnxruntime/core/session/onnxruntime_type_conversion.h index 0941b83b3edc9..d5c818a916a54 100644 --- a/include/onnxruntime/core/session/onnxruntime_type_conversion.h +++ b/include/onnxruntime/core/session/onnxruntime_type_conversion.h @@ -33,7 +33,7 @@ constexpr ONNXTensorElementDataType ToOrtTensorElementDataType(int tensor_proto_ return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT2; default: if (tensor_proto_element_type < 0 || - tensor_proto_element_type > static_cast(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E8M0)) { + tensor_proto_element_type > static_cast(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E3M2)) { return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; } return static_cast(tensor_proto_element_type); diff --git a/js/web/docs/webgl-operators.md b/js/web/docs/webgl-operators.md index 1a8a99b88b2e0..2fbd03963809b 100644 --- a/js/web/docs/webgl-operators.md +++ b/js/web/docs/webgl-operators.md @@ -31,7 +31,7 @@ See [Compatibility](../README.md#Compatibility) for a list of the supported plat | [BitwiseOr](https://github.com/onnx/onnx/blob/main/docs/Operators.md#BitwiseOr) | | | [BitwiseXor](https://github.com/onnx/onnx/blob/main/docs/Operators.md#BitwiseXor) | | | [BlackmanWindow](https://github.com/onnx/onnx/blob/main/docs/Operators.md#BlackmanWindow) | | -| [Cast](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Cast) | [6-8](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-6), [9-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-9), [13-18](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-13), [19-20](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-19), [21-22](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-21), [23](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-23), [24](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-24), [25+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-25) | +| [Cast](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Cast) | [6-8](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-6), [9-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-9), [13-18](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-13), [19-20](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-19), [21-22](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-21), [23](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-23), [24](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-24), [25-27](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-25), [28+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Cast-28) | | [CastLike](https://github.com/onnx/onnx/blob/main/docs/Operators.md#CastLike) | | | [CausalConvWithState](https://github.com/onnx/onnx/blob/main/docs/Operators.md#CausalConvWithState) | | | [Ceil](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Ceil) | [6-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Ceil-6), [13+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Ceil-13) | @@ -53,7 +53,7 @@ See [Compatibility](../README.md#Compatibility) for a list of the supported plat | [CumSum](https://github.com/onnx/onnx/blob/main/docs/Operators.md#CumSum) | | | [DFT](https://github.com/onnx/onnx/blob/main/docs/Operators.md#DFT) | | | [DeformConv](https://github.com/onnx/onnx/blob/main/docs/Operators.md#DeformConv) | | -| [DepthToSpace](https://github.com/onnx/onnx/blob/main/docs/Operators.md#DepthToSpace) | [1-10](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#DepthToSpace-1), [11-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#DepthToSpace-11), [13+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#DepthToSpace-13) | +| [DepthToSpace](https://github.com/onnx/onnx/blob/main/docs/Operators.md#DepthToSpace) | [1-10](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#DepthToSpace-1), [11-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#DepthToSpace-11), [13-27](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#DepthToSpace-13), [28+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#DepthToSpace-28) | | [DequantizeLinear](https://github.com/onnx/onnx/blob/main/docs/Operators.md#DequantizeLinear) | | | [Det](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Det) | | | [Div](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Div) | [7-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Div-7), [13](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Div-13), [14+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Div-14) | @@ -144,7 +144,7 @@ See [Compatibility](../README.md#Compatibility) for a list of the supported plat | [Reciprocal](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Reciprocal) | | | [ReduceL1](https://github.com/onnx/onnx/blob/main/docs/Operators.md#ReduceL1) | | | [ReduceL2](https://github.com/onnx/onnx/blob/main/docs/Operators.md#ReduceL2) | | -| [ReduceLogSum](https://github.com/onnx/onnx/blob/main/docs/Operators.md#ReduceLogSum) | [1-10](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceLogSum-1), [11-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceLogSum-11), [13-17](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceLogSum-13), [18+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceLogSum-18) | +| [ReduceLogSum](https://github.com/onnx/onnx/blob/main/docs/Operators.md#ReduceLogSum) | [1-10](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceLogSum-1), [11-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceLogSum-11), [13-17](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceLogSum-13), [18-27](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceLogSum-18), [28+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceLogSum-28) | | [ReduceLogSumExp](https://github.com/onnx/onnx/blob/main/docs/Operators.md#ReduceLogSumExp) | | | [ReduceMax](https://github.com/onnx/onnx/blob/main/docs/Operators.md#ReduceMax) | [1-10](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMax-1), [11](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMax-11), [12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMax-12), [13-17](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMax-13), [18-19](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMax-18), [20+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMax-20) | | [ReduceMean](https://github.com/onnx/onnx/blob/main/docs/Operators.md#ReduceMean) | [1-10](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMean-1), [11-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMean-11), [13-17](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMean-13), [18+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#ReduceMean-18) | @@ -195,6 +195,7 @@ See [Compatibility](../README.md#Compatibility) for a list of the supported plat | [StringSplit](https://github.com/onnx/onnx/blob/main/docs/Operators.md#StringSplit) | | | [Sub](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Sub) | [7-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Sub-7), [13](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Sub-13), [14+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Sub-14) | | [Sum](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Sum) | [6-7](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Sum-6), [8-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Sum-8), [13+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Sum-13) | +| [SwiGLU](https://github.com/onnx/onnx/blob/main/docs/Operators.md#SwiGLU) | | | [Swish](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Swish) | | | [Tan](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Tan) | [7-21](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Tan-7), [22+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Tan-22) | | [Tanh](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Tanh) | [6-12](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Tanh-6), [13+](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Tanh-13) | diff --git a/js/web/docs/webgpu-operators.md b/js/web/docs/webgpu-operators.md index ee471fc657ce5..6307b1fa24a03 100644 --- a/js/web/docs/webgpu-operators.md +++ b/js/web/docs/webgpu-operators.md @@ -25,7 +25,7 @@ Do not modify directly.* | BatchNormalization | ai.onnx(7-8,9-13,14,15+); com.ms.internal.nhwc(7-8,9-13,14,15+) | | | BiasAdd | com.microsoft(1+) | | | BiasSplitGelu | com.microsoft(1+) | | -| Cast | ai.onnx(6-8,9-12,13-18,19-20,21+) | | +| Cast | ai.onnx(6-8,9-12,13-18,19-20,21-27,28+) | | | Ceil | ai.onnx(6-12,13+) | | | Clip | ai.onnx(6-10,11,12,13+) | | | Concat | ai.onnx(1-3,4-10,11-12,13+) | | diff --git a/js/web/lib/wasm/jsep/backend-webgpu.ts b/js/web/lib/wasm/jsep/backend-webgpu.ts index 753a11a88791b..b7ca6b1c95034 100644 --- a/js/web/lib/wasm/jsep/backend-webgpu.ts +++ b/js/web/lib/wasm/jsep/backend-webgpu.ts @@ -694,6 +694,9 @@ export class WebGpuBackend { } upload(gpuDataId: number, data: Uint8Array): void { + // The allocation planner can reuse a buffer still read by pending commands. Submit those commands before + // queue.writeBuffer overwrites it, including copies that do not increment pendingDispatchNumber. + this.flush(); this.gpuDataManager.upload(gpuDataId, data); } diff --git a/js/web/test/unittests/backends/webgpu/test-upload.ts b/js/web/test/unittests/backends/webgpu/test-upload.ts new file mode 100644 index 0000000000000..a7604cc9136a8 --- /dev/null +++ b/js/web/test/unittests/backends/webgpu/test-upload.ts @@ -0,0 +1,110 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { expect } from 'chai'; +import { env } from 'onnxruntime-common'; + +import { WebGpuBackend } from '../../../../lib/wasm/jsep/backend-webgpu'; + +describe('WebGPU upload ordering', function () { + let backend: WebGpuBackend; + + beforeEach(async function (this: Mocha.Context) { + backend = new WebGpuBackend(); + const adapter = await navigator.gpu?.requestAdapter(); + if (!adapter) { + this.skip(); + } + await backend.initialize( + { ...env, trace: false, wasm: {}, webgpu: { ...env.webgpu, profiling: { mode: 'off' } } }, + adapter, + ); + backend.device.pushErrorScope('validation'); + }); + + afterEach(async function () { + if (!backend.device) { + return; + } + try { + backend.flush(); + expect(await backend.device.popErrorScope()).to.equal(null); + } finally { + backend.dispose(); + backend.device.destroy(); + } + }); + + for (const operation of ['compute', 'memcpy'] as const) { + for (const { name, size, offset } of [ + { name: 'aligned data', size: 16, offset: 0 }, + { name: 'unaligned data offset', size: 16, offset: 1 }, + { name: 'padded data', size: 3, offset: 0 }, + ]) { + it(`preserves pending ${operation} inputs when uploading ${name} to the same buffer`, async function () { + const input = backend.alloc(size); + const firstOutput = backend.alloc(size); + const secondOutput = backend.alloc(size); + const first = new Uint8Array(new ArrayBuffer(size + offset), offset, size).fill(1); + const second = new Uint8Array(new ArrayBuffer(size + offset), offset, size).fill(2); + + const pipeline = + operation === 'compute' + ? backend.device.createComputePipeline({ + layout: 'auto', + compute: { + module: backend.device.createShaderModule({ + code: ` + @group(0) @binding(0) var input: array; + @group(0) @binding(1) var output: array; + @compute @workgroup_size(1) + fn main(@builtin(global_invocation_id) id: vec3) { + output[id.x] = input[id.x]; + } + `, + }), + entryPoint: 'main', + }, + }) + : undefined; + + const recordRead = (output: number) => { + if (!pipeline) { + backend.memcpy(input, output); + // Copies leave the dispatch counter at zero but still need submission before the upload. + expect(backend.pendingDispatchNumber).to.equal(0); + return; + } + const pass = backend.getComputePassEncoder(); + pass.setPipeline(pipeline); + pass.setBindGroup( + 0, + backend.device.createBindGroup({ + layout: pipeline.getBindGroupLayout(0), + entries: [ + { binding: 0, resource: { buffer: backend.gpuDataManager.get(input)!.buffer } }, + { binding: 1, resource: { buffer: backend.gpuDataManager.get(output)!.buffer } }, + ], + }), + ); + pass.dispatchWorkgroups(4); + backend.pendingDispatchNumber++; + }; + + // Model the allocation planner reusing one input buffer for successive CPU EP outputs. + // Neither read is explicitly flushed before the next upload. + backend.upload(input, first); + recordRead(firstOutput); + backend.upload(input, second); + recordRead(secondOutput); + + const firstResult = new Uint8Array(size); + const secondResult = new Uint8Array(size); + await backend.download(firstOutput, () => firstResult); + await backend.download(secondOutput, () => secondResult); + expect(Array.from(firstResult)).to.deep.equal(Array.from(first)); + expect(Array.from(secondResult)).to.deep.equal(Array.from(second)); + }); + } + } +}); diff --git a/js/web/test/unittests/index.ts b/js/web/test/unittests/index.ts index c27d8efb0cc78..79652ba8ac30e 100644 --- a/js/web/test/unittests/index.ts +++ b/js/web/test/unittests/index.ts @@ -4,6 +4,7 @@ 'use strict'; if (typeof window !== 'undefined') { + require('./backends/webgpu/test-upload'); require('./backends/webgl/test-glsl-function-inliner'); require('./backends/webgl/test-conv-new'); require('./backends/webgl/test-pack-unpack'); diff --git a/model_package/src/asset_hasher.cc b/model_package/src/asset_hasher.cc index f30b520e539f0..e835c288873e1 100644 --- a/model_package/src/asset_hasher.cc +++ b/model_package/src/asset_hasher.cc @@ -50,11 +50,8 @@ ModelPackageStatus* ComputeDirectoryAssetUri(const fs::path& source_dir, "ComputeDirectoryAssetUri: cannot iterate '" + source_dir.u8string() + "': " + ec.message()); } - for (; walker != fs::recursive_directory_iterator(); walker.increment(ec)) { - if (ec) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, - "ComputeDirectoryAssetUri: iteration error: " + ec.message()); - } + std::error_code iter_ec; + for (; !iter_ec && walker != fs::recursive_directory_iterator(); walker.increment(iter_ec)) { const fs::directory_entry& de = *walker; if (de.is_symlink(ec)) { return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, @@ -73,6 +70,10 @@ ModelPackageStatus* ComputeDirectoryAssetUri(const fs::path& source_dir, de.path().u8string() + "' (only regular files and directories allowed)."); } } + if (iter_ec) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, + "ComputeDirectoryAssetUri: iteration error: " + iter_ec.message()); + } std::sort(entries.begin(), entries.end(), [](const auto& a, const auto& b) { return a.first < b.first; }); diff --git a/model_package/src/authoring.cc b/model_package/src/authoring.cc index 1d5faa1aaea3e..ea4f8f9548160 100644 --- a/model_package/src/authoring.cc +++ b/model_package/src/authoring.cc @@ -78,16 +78,15 @@ ModelPackageStatus* RefreshSharedAssetsHelper(ModelPackage* pkg) { return mp::RefreshSharedAssets(pkg, mp::PathOptionsFor(pkg)); } -ModelPackageStatus* PostMutate(ModelPackage* pkg, bool refresh_assets = true) try { +ModelPackageStatus* PostMutate(ModelPackage* pkg, bool refresh_assets = true) MODEL_PACKAGE_TRY { mp::DropViewCache(pkg); if (refresh_assets) { if (auto* s = RefreshSharedAssetsHelper(pkg)) return s; } if (auto* s = mp::RefreshPackageMetadata(pkg)) return s; return mp::RefreshExecutorInfoCache(pkg, /*strict_missing_external=*/false); -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH ordered_json& EnsureManifestComponentsObject(ModelPackage* pkg) { if (!pkg->manifest.contains("components")) { @@ -128,7 +127,7 @@ ModelPackageStatus* ModelPackage_New(ModelPackage** out) { ModelPackageStatus* ModelPackage_SetComponentInline(ModelPackage* pkg, const char* name, - const char* component_json) try { + const char* component_json) MODEL_PACKAGE_TRY { if (!pkg) return NullArg("pkg"); if (!name) return NullArg("name"); if (!component_json) return NullArg("component_json"); @@ -156,13 +155,12 @@ ModelPackageStatus* ModelPackage_SetComponentInline(ModelPackage* pkg, } RebuildComponentIndex(pkg); return PostMutate(pkg); -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH ModelPackageStatus* ModelPackage_SetComponentExternal(ModelPackage* pkg, const char* name, - const char* path) try { + const char* path) MODEL_PACKAGE_TRY { if (!pkg) return NullArg("pkg"); if (!name) return NullArg("name"); if (!path) return NullArg("path"); @@ -218,9 +216,8 @@ ModelPackageStatus* ModelPackage_SetComponentExternal(ModelPackage* pkg, } RebuildComponentIndex(pkg); return PostMutate(pkg); -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH ModelPackageStatus* ModelPackage_RemoveComponent(ModelPackage* pkg, const char* name) { if (!pkg) return NullArg("pkg"); @@ -245,7 +242,7 @@ ModelPackageStatus* ModelPackage_RemoveComponent(ModelPackage* pkg, const char* ModelPackageStatus* ModelPackage_SetVariant(ModelPackage* pkg, const char* component_name, const char* variant_name, - const char* variant_json) try { + const char* variant_json) MODEL_PACKAGE_TRY { if (!pkg) return NullArg("pkg"); if (!component_name) return NullArg("component_name"); if (!variant_name) return NullArg("variant_name"); @@ -289,9 +286,8 @@ ModelPackageStatus* ModelPackage_SetVariant(ModelPackage* pkg, // Invalidate cached component JSON. comp->component_json_cache.reset(); return PostMutate(pkg); -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH ModelPackageStatus* ModelPackage_RemoveVariant(ModelPackage* pkg, const char* component_name, @@ -342,7 +338,7 @@ ModelPackageStatus* MutateExecutorInfo(ModelPackage* pkg, const char* component, const char* variant, const char* namespace_, - const ordered_json* new_value /* null = remove */) try { + const ordered_json* new_value /* null = remove */) MODEL_PACKAGE_TRY { if (!pkg) return NullArg("pkg"); if (!component) return NullArg("component"); if (!variant) return NullArg("variant"); @@ -377,9 +373,8 @@ ModelPackageStatus* MutateExecutorInfo(ModelPackage* pkg, if (auto* s = ReparseVariantInPlace(pkg, comp, var)) return s; comp->component_json_cache.reset(); return PostMutate(pkg, /*refresh_assets=*/false); -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH } // namespace @@ -420,7 +415,7 @@ ModelPackageStatus* ModelPackage_AddSharedAsset(ModelPackage* pkg, const char* source_dir, const char* expected_uri_or_null, bool copy_in, - const char** out_uri) try { + const char** out_uri) MODEL_PACKAGE_TRY { if (!pkg) return NullArg("pkg"); if (!source_dir) return NullArg("source_dir"); if (!out_uri) return NullArg("out_uri"); @@ -465,9 +460,8 @@ ModelPackageStatus* ModelPackage_AddSharedAsset(ModelPackage* pkg, } *out_uri = pkg->shared_assets[sit->second]->uri_cache.c_str(); return nullptr; -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH ModelPackageStatus* ModelPackage_RemoveSharedAsset(ModelPackage* pkg, const char* uri) { if (!pkg) return NullArg("pkg"); diff --git a/model_package/src/commit_prune.cc b/model_package/src/commit_prune.cc index 120422b4f9109..809e5473d769b 100644 --- a/model_package/src/commit_prune.cc +++ b/model_package/src/commit_prune.cc @@ -116,22 +116,24 @@ ModelPackageStatus* CopyTreeNoFollow(const fs::path& src, const fs::path& dst) { fs::create_directories(dst, ec); if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, "mkdir '" + dst.u8string() + "': " + ec.message()); - for (fs::recursive_directory_iterator it(src, fs::directory_options::none, ec), end; - it != end; it.increment(ec)) { - if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, - "iterate '" + src.u8string() + "': " + ec.message()); + std::error_code iter_ec; + for (fs::recursive_directory_iterator it(src, fs::directory_options::none, iter_ec), end; + !iter_ec && it != end; it.increment(iter_ec)) { const auto& entry = *it; fs::path rel = fs::relative(entry.path(), src, ec); fs::path target = dst / rel; - if (entry.is_symlink()) { + fs::file_status st = entry.symlink_status(ec); + if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, + "stat '" + entry.path().u8string() + "': " + ec.message()); + if (fs::is_symlink(st)) { return MakeStatus(MODEL_PACKAGE_ERR_SCHEMA, "shared asset source contains a symlink: '" + entry.path().u8string() + "'."); } - if (entry.is_directory()) { + if (fs::is_directory(st)) { fs::create_directories(target, ec); if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, "mkdir '" + target.u8string() + "': " + ec.message()); - } else if (entry.is_regular_file()) { + } else if (fs::is_regular_file(st)) { fs::create_directories(target.parent_path(), ec); fs::copy_file(entry.path(), target, fs::copy_options::overwrite_existing, ec); if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, @@ -144,6 +146,8 @@ ModelPackageStatus* CopyTreeNoFollow(const fs::path& src, const fs::path& dst) { entry.path().u8string() + "'."); } } + if (iter_ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, + "iterate '" + src.u8string() + "': " + iter_ec.message()); if (auto* s = FsyncPath(dst, /*is_dir=*/true)) return s; return nullptr; } @@ -609,17 +613,16 @@ extern "C" { ModelPackageStatus* ModelPackage_Commit(ModelPackage* pkg, const char* dest_root_or_null, - ModelPackageWriteMode mode) try { + ModelPackageWriteMode mode) MODEL_PACKAGE_TRY { if (!pkg) return NullArg("pkg"); if (dest_root_or_null) { return CommitToDestRoot(pkg, fs::u8path(dest_root_or_null), mode); } return CommitInPlace(pkg, mode); -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH -ModelPackageStatus* ModelPackage_Prune(ModelPackage* pkg) try { +ModelPackageStatus* ModelPackage_Prune(ModelPackage* pkg) MODEL_PACKAGE_TRY { if (!pkg) return NullArg("pkg"); if (pkg->package_root.empty()) return nullptr; @@ -636,13 +639,19 @@ ModelPackageStatus* ModelPackage_Prune(ModelPackage* pkg) try { std::error_code ec; fs::path assets_root = pkg->package_root / "shared_assets"; if (fs::is_directory(assets_root, ec)) { - for (const auto& entry : fs::directory_iterator(assets_root, ec)) { - if (ec) break; - if (!entry.is_directory()) continue; + std::error_code open_ec, iter_ec; + for (fs::directory_iterator it(assets_root, open_ec), end; !iter_ec && it != end; it.increment(iter_ec)) { + const auto& entry = *it; + fs::file_status st = entry.status(ec); + if (st.type() == fs::file_type::none) { + return MakeStatus(MODEL_PACKAGE_ERR_IO, "stat '" + entry.path().u8string() + "': " + ec.message()); + } + if (!fs::is_directory(st)) continue; if (!IsTmpName(entry.path())) continue; if (!IsOldEnough(entry.path())) continue; fs::remove_all(entry.path(), ec); } + if (iter_ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, "iterate '" + assets_root.u8string() + "': " + iter_ec.message()); } // Tracked-orphan sweep: components before variants so a component_dir @@ -652,8 +661,7 @@ ModelPackageStatus* ModelPackage_Prune(ModelPackage* pkg) try { SweepOrphanDirs(pkg, &pkg->pending_orphan_variant_dirs, live_dirs); return nullptr; -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH } // extern "C" diff --git a/model_package/src/manifest_parser.cc b/model_package/src/manifest_parser.cc index df6f3289111f3..50ebf373e7798 100644 --- a/model_package/src/manifest_parser.cc +++ b/model_package/src/manifest_parser.cc @@ -315,7 +315,9 @@ ModelPackageStatus* LoadComponentForEntry(const fs::path& manifest_dir, std::error_code ec; if (fs::is_directory(resolved, ec)) { resolved /= kComponentFileName; - if (!fs::exists(resolved)) { + bool has_component = fs::exists(resolved, ec); + if (ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, "component '" + name + "': " + ec.message()); + if (!has_component) { return MakeStatus(MODEL_PACKAGE_ERR_NOT_FOUND, "component '" + name + "': directory has no '" + kComponentFileName + "'."); @@ -385,8 +387,9 @@ ModelPackageStatus* LoadSharedAssets(ModelPackage* pkg, const PathResolverOption std::error_code ec; if (!pkg->package_root.empty() && fs::is_directory(assets_root, ec)) { std::vector discovered; - for (const auto& entry : fs::directory_iterator(assets_root, ec)) { - if (ec) break; + std::error_code iter_ec; + for (fs::directory_iterator it(assets_root, iter_ec), end; !iter_ec && it != end; it.increment(iter_ec)) { + const auto& entry = *it; if (!entry.is_directory(ec)) continue; std::string name = entry.path().filename().u8string(); std::string uri = SharedAssetUriFromDirName(name); @@ -394,6 +397,7 @@ ModelPackageStatus* LoadSharedAssets(ModelPackage* pkg, const PathResolverOption if (!seen.insert(uri).second) continue; discovered.push_back(std::move(uri)); } + if (iter_ec) return MakeStatus(MODEL_PACKAGE_ERR_IO, "iterate '" + assets_root.u8string() + "': " + iter_ec.message()); std::sort(discovered.begin(), discovered.end()); for (auto& uri : discovered) ordered_uris.push_back(std::move(uri)); } diff --git a/model_package/src/model_package_impl.cc b/model_package/src/model_package_impl.cc index 8d66cc7accd3b..b4f4902611607 100644 --- a/model_package/src/model_package_impl.cc +++ b/model_package/src/model_package_impl.cc @@ -189,7 +189,7 @@ void ModelPackageStatus_Release(ModelPackageStatus* s) { ModelPackageStatus* ModelPackage_Open(const char* package_root, const ModelPackageOpenOptions* opts, - ModelPackage** out) try { + ModelPackage** out) MODEL_PACKAGE_TRY { if (!package_root) return NullArg("package_root"); if (!out) return NullArg("out"); @@ -209,9 +209,8 @@ ModelPackageStatus* ModelPackage_Open(const char* package_root, } *out = pkg.release(); return nullptr; -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH void ModelPackage_Close(ModelPackage* pkg) { if (!pkg) return; @@ -285,7 +284,7 @@ ModelPackageStatus* ModelPackage_ResolveStringRef(const ModelPackage* pkg, const char* base_dir, const char* input, bool must_exist, - const char** out_path) try { + const char** out_path) MODEL_PACKAGE_TRY { if (!pkg) return NullArg("pkg"); if (!input) return NullArg("input"); if (!out_path) return NullArg("out_path"); @@ -323,9 +322,8 @@ ModelPackageStatus* ModelPackage_ResolveStringRef(const ModelPackage* pkg, slot = resolved.u8string(); *out_path = slot.c_str(); return nullptr; -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH // ───────────────────────────────────────────────────────────────────────────── // Round-trip JSON getters @@ -383,7 +381,7 @@ ModelPackageStatus* ModelPackage_GetVariantJson(const ModelPackage* pkg, // ───────────────────────────────────────────────────────────────────────────── ModelPackageStatus* ModelPackage_ComputeDirectoryHash(const char* source_dir, - const char** out_uri) try { + const char** out_uri) MODEL_PACKAGE_TRY { if (!source_dir) return NullArg("source_dir"); if (!out_uri) return NullArg("out_uri"); static thread_local std::string slot; @@ -392,8 +390,7 @@ ModelPackageStatus* ModelPackage_ComputeDirectoryHash(const char* source_dir, } *out_uri = slot.c_str(); return nullptr; -} catch (const std::system_error& error) { - return MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); } +MODEL_PACKAGE_CATCH } // extern "C" diff --git a/model_package/src/status_impl.h b/model_package/src/status_impl.h index 6cc1c94238f98..c6e7abaea5e7e 100644 --- a/model_package/src/status_impl.h +++ b/model_package/src/status_impl.h @@ -9,10 +9,25 @@ #include #include +#include #include #include "model_package_api.h" +/// Function-try-block for functions returning ModelPackageStatus*: +/// MODEL_PACKAGE_CATCH turns a std::system_error into MODEL_PACKAGE_ERR_IO. +/// Both expand to nothing when exceptions are disabled. +#if defined(__cpp_exceptions) || defined(_CPPUNWIND) +#define MODEL_PACKAGE_TRY try +#define MODEL_PACKAGE_CATCH \ + catch (const std::system_error& error) { \ + return ::model_package::MakeStatus(MODEL_PACKAGE_ERR_IO, error.what()); \ + } +#else +#define MODEL_PACKAGE_TRY +#define MODEL_PACKAGE_CATCH +#endif + struct ModelPackageStatus { ModelPackageErrorCode code{MODEL_PACKAGE_ERR_INVALID_ARG}; std::string message; diff --git a/model_package/tests/test_asset_hashing.cc b/model_package/tests/test_asset_hashing.cc index 9c983cf096952..842a7c59488e3 100644 --- a/model_package/tests/test_asset_hashing.cc +++ b/model_package/tests/test_asset_hashing.cc @@ -213,6 +213,28 @@ bool test_directory_hash_rejects_symlink() { return true; } +bool test_directory_hash_reports_unreadable_subdirectory() { + Sandbox s; + s.Write("a.txt", "alpha"); + s.Write("locked/b.txt", "beta"); + const fs::path locked = s.root() / "locked"; + std::error_code ec; + fs::permissions(locked, fs::perms::none, ec); + fs::directory_iterator probe(locked, ec); + if (!ec) { + fs::permissions(locked, fs::perms::owner_all, ec); + std::printf("[SKIP] %s (directory permissions not enforced)\n", g_current); + return true; + } + const char* uri = nullptr; + ModelPackageStatus* st = ModelPackage_ComputeDirectoryHash(s.root().u8string().c_str(), &uri); + fs::permissions(locked, fs::perms::owner_all, ec); + CHECK(st != nullptr); + CHECK(ModelPackageStatus_Code(st) == MODEL_PACKAGE_ERR_IO); + ModelPackageStatus_Release(st); + return true; +} + bool test_directory_hash_known_value_single_file() { // Known-answer check: the directory URI hashes a manifest of " \n" // lines, so compute the expected value the same way and compare. @@ -340,6 +362,7 @@ const Test kTests[] = { {"directory_hash_content_change_differs", test_directory_hash_content_change_differs}, {"directory_hash_empty_dirs_ignored", test_directory_hash_empty_dirs_ignored}, {"directory_hash_rejects_symlink", test_directory_hash_rejects_symlink}, + {"directory_hash_reports_unreadable_subdirectory", test_directory_hash_reports_unreadable_subdirectory}, {"directory_hash_known_value_single_file", test_directory_hash_known_value_single_file}, {"directory_hash_sorted_order_independent_of_walk", test_directory_hash_sorted_order_independent_of_walk}, {"directory_hash_uses_forward_slash", test_directory_hash_uses_forward_slash}, diff --git a/model_package/tests/test_commit.cc b/model_package/tests/test_commit.cc index ba6ba3d5923f6..44e796d28260b 100644 --- a/model_package/tests/test_commit.cc +++ b/model_package/tests/test_commit.cc @@ -192,6 +192,48 @@ bool test_commit_pending_shared_asset_copy_in() { return true; } +bool test_commit_shared_asset_copy_in_nested_directory() { + Sandbox s; + s.Write("src_asset/m.onnx", "hello world"); + s.Write("src_asset/weights/w.bin", "weights"); + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), s.path("src_asset").u8string().c_str(), + nullptr, /*copy_in=*/true, &uri)); + std::string uri_copy(uri); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").u8string().c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + fs::path landed = s.path("pkg") / "shared_assets" / ("sha256-" + uri_copy.substr(7)); + CHECK(fs::is_regular_file(landed / "weights" / "w.bin")); + return true; +} + +bool test_commit_reports_unreadable_copy_in_subdirectory() { + Sandbox s; + s.Write("src_asset/m.onnx", "hello world"); + s.Write("src_asset/weights/w.bin", "weights"); + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + + const char* uri = nullptr; + CHECK_OK(ModelPackage_AddSharedAsset(p.get(), s.path("src_asset").u8string().c_str(), + nullptr, /*copy_in=*/true, &uri)); + const fs::path locked = s.path("src_asset") / "weights"; + std::error_code ec; + fs::permissions(locked, fs::perms::none, ec); + fs::directory_iterator probe(locked, ec); + if (!ec) { + fs::permissions(locked, fs::perms::owner_all, ec); + std::printf("[SKIP] %s (directory permissions not enforced)\n", g_current); + return true; + } + ModelPackageStatus* st = ModelPackage_Commit(p.get(), s.path("pkg").u8string().c_str(), + MODEL_PACKAGE_WRITE_PRESERVE); + fs::permissions(locked, fs::perms::owner_all, ec); + CHECK_ERR(st, MODEL_PACKAGE_ERR_IO); + return true; +} + bool test_commit_dense_inlines_external_component() { Sandbox s; PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); @@ -355,6 +397,24 @@ bool test_prune_removes_stale_staging_dirs() { return true; } +bool test_prune_reports_unstatable_shared_assets_entry() { + Sandbox s; + PkgHandle p = MakeAuthoredPkgAt(s.path("pkg")); + CHECK_OK(ModelPackage_Commit(p.get(), s.path("pkg").u8string().c_str(), + MODEL_PACKAGE_WRITE_PRESERVE)); + + fs::path loop = s.path("pkg") / "shared_assets" / "loop"; + std::error_code ec; + fs::create_directories(loop.parent_path(), ec); + fs::create_directory_symlink("loop", loop, ec); + if (ec) { + std::printf("[SKIP] %s (symlink unsupported)\n", g_current); + return true; + } + CHECK_ERR(ModelPackage_Prune(p.get()), MODEL_PACKAGE_ERR_IO); + return true; +} + bool test_commit_accepts_unreferenced_shared_asset() { // Shared assets no longer require an in-manifest reference: AddSharedAsset // signals the user's intent to ship the asset, period. Commit materializes @@ -442,6 +502,8 @@ const Test kTests[] = { {"commit_requires_package_root", test_commit_requires_package_root}, {"commit_external_component_writes_file", test_commit_external_component_writes_file}, {"commit_pending_shared_asset_copy_in", test_commit_pending_shared_asset_copy_in}, + {"commit_shared_asset_copy_in_nested_directory", test_commit_shared_asset_copy_in_nested_directory}, + {"commit_reports_unreadable_copy_in_subdirectory", test_commit_reports_unreadable_copy_in_subdirectory}, {"commit_dense_inlines_external_component", test_commit_dense_inlines_external_component}, {"commit_dense_rejects_external_executor_info", test_commit_dense_rejects_external_executor_info}, {"commit_dest_root_self_contained", test_commit_dest_root_self_contained}, @@ -450,6 +512,7 @@ const Test kTests[] = { {"prune_never_touches_shared_assets", test_prune_never_touches_shared_assets}, {"prune_reclaims_tracked_orphan_variant_dirs", test_prune_reclaims_tracked_orphan_variant_dirs}, {"prune_removes_stale_staging_dirs", test_prune_removes_stale_staging_dirs}, + {"prune_reports_unstatable_shared_assets_entry", test_prune_reports_unstatable_shared_assets_entry}, {"commit_accepts_unreferenced_shared_asset", test_commit_accepts_unreferenced_shared_asset}, {"commit_leaves_no_temp_files", test_commit_leaves_no_temp_files}, {"commit_unicode_paths_and_assets", test_commit_unicode_paths_and_assets}, diff --git a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc index 4b50b240bdea0..23e8b6b02b354 100644 --- a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc @@ -15,6 +15,7 @@ namespace onnxruntime::contrib { BranchwiseRMSNorm, kMSDomain, 1, T, kCpuExecutionProvider, \ KernelDefBuilder() \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("F", DataTypeImpl::GetTensorType()) \ .TypeConstraint("M", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}), \ @@ -47,7 +48,11 @@ template BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) : LayerNormImpl(info, -1, info.GetAttrOrDefault("epsilon", 1e-5f), true), epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - num_branches_(info.GetAttrOrDefault("num_branches", 0)) {} + num_branches_(info.GetAttrOrDefault("num_branches", 0)), + use_onnx_arithmetic_(info.GetAttrOrDefault("use_onnx_arithmetic", 0) != 0) { + const auto mode = info.GetAttrOrDefault("use_onnx_arithmetic", 0); + ORT_ENFORCE(mode == 0 || mode == 1, "use_onnx_arithmetic must be 0 or 1"); +} template Status BranchwiseRMSNorm::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, @@ -71,7 +76,10 @@ Status BranchwiseRMSNorm::Compute(OpKernelContext* context) const { } auto* y = context->Output(0, x->Shape()); - if (scale != nullptr && + auto* widened = context->OutputCount() > 1 ? context->Output(1, x->Shape()) : nullptr; + ORT_RETURN_IF_NOT(widened == nullptr || use_onnx_arithmetic_, + "Y_fp32 requires use_onnx_arithmetic=1"); + if (!use_onnx_arithmetic_ && scale != nullptr && scale->Shape().Size() == params.hidden && scale->GetElementType() == x->GetElementType()) { TensorShapeVector grouped_dims(x->Shape().GetDims().begin(), @@ -98,15 +106,20 @@ Status BranchwiseRMSNorm::Compute(OpKernelContext* context) const { const float value = static_cast(x_data[offset + h]); sum_sq += value * value; } - const float inv_rms = 1.0f / std::sqrt(sum_sq / static_cast(params.hidden) + epsilon_); + const float rms = std::sqrt(sum_sq / static_cast(params.hidden) + epsilon_); + const float inv_rms = 1.0f / rms; for (int64_t h = 0; h < params.hidden; ++h) { const int64_t scale_index = scale != nullptr && scale->Shape().Size() == params.hidden ? h : (group % params.branches) * params.hidden + h; const float weight = scale == nullptr ? 1.0f : LoadFloat(*scale, scale_index); - y_data[offset + h] = - static_cast(static_cast(x_data[offset + h]) * inv_rms * weight); + const float value = static_cast(x_data[offset + h]); + y_data[offset + h] = static_cast( + (use_onnx_arithmetic_ ? value / rms : value * inv_rms) * weight); + if (widened != nullptr) { + widened->MutableData()[offset + h] = static_cast(y_data[offset + h]); + } } }, 0); diff --git a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h index 63284794f3f87..bd58f86db058d 100644 --- a/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h +++ b/onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h @@ -18,6 +18,7 @@ class BranchwiseRMSNorm final : public LayerNormImpl { private: float epsilon_; int64_t num_branches_; + bool use_onnx_arithmetic_; }; } // namespace onnxruntime::contrib diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc index 110acc10a8dee..d42dbae635047 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc @@ -52,6 +52,75 @@ constexpr bool kQNBitForceFp32Supported = true; constexpr bool kQNBitForceFp32Supported = false; #endif +Status ValidatePrePackInitializerShape(const OpKernelInfo& info, + const Tensor& tensor, + int input_idx, + size_t N, + size_t K, + size_t block_size, + size_t nbits, + bool has_zp_arg, + bool has_zp_input) { + const int64_t n = static_cast(N); + const int64_t k = static_cast(K); + const int64_t bs = static_cast(block_size); + const int64_t bits = static_cast(nbits); + const int64_t k_blocks = matmul_nbits_helper::GetKBlocks(k, bs); + const int64_t blob_size = matmul_nbits_helper::GetBlobSize(bs, bits); + const int64_t zp_blob_size_uint8 = matmul_nbits_helper::GetUint8ZeroPointBlobSize(k_blocks, bits); + const int64_t scales_size = SafeInt(n) * k_blocks; + const int64_t packed_zp_size = SafeInt(n) * zp_blob_size_uint8; + + auto validate_scales_shape = [&](const TensorShape& shape) -> Status { + ORT_RETURN_IF_NOT(shape == TensorShape({scales_size}) || shape == TensorShape({n, k_blocks}), + "MatMulNBits PrePack: scales initializer shape ", shape, + " does not match attribute-derived shape [", scales_size, "] or [", + n, ",", k_blocks, "]"); + return Status::OK(); + }; + + auto validate_zero_points_shape = [&](const TensorShape& shape, int32_t element_type) -> Status { + if (element_type == ONNX_NAMESPACE::TensorProto_DataType_UINT8) { + ORT_RETURN_IF_NOT(shape == TensorShape({packed_zp_size}) || + shape == TensorShape({n, zp_blob_size_uint8}), + "MatMulNBits PrePack: zero_points initializer shape ", shape, + " does not match attribute-derived shape [", packed_zp_size, "] or [", + n, ",", zp_blob_size_uint8, "]"); + } else { + ORT_RETURN_IF_NOT(shape == TensorShape({scales_size}) || shape == TensorShape({n, k_blocks}), + "MatMulNBits PrePack: zero_points initializer shape ", shape, + " does not match attribute-derived shape [", scales_size, "] or [", + n, ",", k_blocks, "]"); + } + return Status::OK(); + }; + + const TensorShape& shape = tensor.Shape(); + if (input_idx == InputIndex::B) { + ORT_RETURN_IF_NOT(shape == TensorShape({n, k_blocks, blob_size}), + "MatMulNBits PrePack: B initializer shape ", shape, + " does not match attribute-derived shape [", n, ",", k_blocks, ",", blob_size, + "] (N=", N, ", K=", K, ", bits=", nbits, ", block_size=", block_size, ")"); + + const Tensor* scales_tensor = nullptr; + if (info.TryGetConstantInput(InputIndex::scales, &scales_tensor) && scales_tensor != nullptr) { + ORT_RETURN_IF_ERROR(validate_scales_shape(scales_tensor->Shape())); + } + + const Tensor* zp_tensor = nullptr; + if (has_zp_arg && has_zp_input && + info.TryGetConstantInput(InputIndex::zero_points, &zp_tensor) && zp_tensor != nullptr) { + ORT_RETURN_IF_ERROR(validate_zero_points_shape(zp_tensor->Shape(), zp_tensor->GetElementType())); + } + } else if (input_idx == InputIndex::scales) { + ORT_RETURN_IF_ERROR(validate_scales_shape(shape)); + } else if (input_idx == InputIndex::zero_points) { + ORT_RETURN_IF_ERROR(validate_zero_points_shape(shape, tensor.GetElementType())); + } + + return Status::OK(); +} + // T: A data type. template MLAS_QNBIT_GEMM_COMPUTE_TYPE @@ -275,91 +344,9 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ All /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; - // Validate the incoming initializer's shape against the attribute-derived shape before any of - // the pack routines below dereference tensor.DataRaw(). The MLAS pack routines size their reads - // from the (N, K, bits, block_size) attributes; without this check a crafted model whose - // attributes overstate the real tensor extents would trigger a heap-buffer-overflow READ at - // session initialization. The matching guard in matmul_nbits_helper::CheckInputs is invoked - // from Compute() -- too late, because PrePack has already done the OOB read, and by then the - // original B tensor is passed as nullptr so the Compute-time check never sees it. - // - // When input_idx == B, this guard also validates the constant scales and zero_points - // initializers (looked up via TryGetConstantInput). SessionState::PrepackConstantInitializedTensors - // iterates inputs in index order, so the B PrePack call runs before scales/zero_points are - // prepacked on their own. The B prepack path reads those constant tensors and passes their - // raw data to the MLAS pack routines (MlasLutGemmPack, MlasQNBitGemmPackQuantBData), which size - // their reads from the same (N, K, bits, block_size) attributes. Without validating scales / - // zero_points here, a crafted model with an undersized scales or zero_points buffer would still - // trigger an OOB read inside the B packing pass before each tensor's own PrePack call could - // catch the mismatch. - // - // This validation runs before the early-return guards below (has_g_idx_, unquantized ZP, - // dynamic-ZP-with-LUT, !MlasIsQNBitGemmAvailable). On builds where MLAS QNBit GEMM is not - // available (e.g. Windows x86 32-bit) PrePack would otherwise short-circuit before reaching - // these checks, and the original B tensor is dropped after PrePack so Compute()'s helper-time - // check never sees it. Running the validation first makes session init reject bad-shape models - // consistently across all build configurations. The checks are cheap (a few TensorShape - // equality comparisons) and independent of any MLAS kernel availability. - { - const int64_t n = static_cast(N_); - const int64_t k = static_cast(K_); - const int64_t bs = static_cast(block_size_); - const int64_t bits = static_cast(nbits_); - // Layout derivations live in matmul_nbits_helper.h so this guard and the Compute-time - // CheckInputs stay in lockstep if the canonical packing layout ever changes. - const int64_t k_blocks = matmul_nbits_helper::GetKBlocks(k, bs); - const int64_t blob_size = matmul_nbits_helper::GetBlobSize(bs, bits); - const int64_t zp_blob_size_uint8 = matmul_nbits_helper::GetUint8ZeroPointBlobSize(k_blocks, bits); - - auto validate_scales_shape = [&](const TensorShape& s) -> Status { - // scales may be 1D [n * k_blocks] or 2D [n, k_blocks] for backward compatibility. - ORT_RETURN_IF_NOT(s == TensorShape({n * k_blocks}) || s == TensorShape({n, k_blocks}), - "MatMulNBits PrePack: scales initializer shape ", s, - " does not match attribute-derived shape [", n * k_blocks, "] or [", - n, ",", k_blocks, "]"); - return Status::OK(); - }; - - auto validate_zero_points_shape = [&](const TensorShape& s, int32_t element_type) -> Status { - if (element_type == ONNX_NAMESPACE::TensorProto_DataType_UINT8) { - ORT_RETURN_IF_NOT(s == TensorShape({n * zp_blob_size_uint8}) || s == TensorShape({n, zp_blob_size_uint8}), - "MatMulNBits PrePack: zero_points initializer shape ", s, - " does not match attribute-derived shape [", n * zp_blob_size_uint8, "] or [", - n, ",", zp_blob_size_uint8, "]"); - } else { - ORT_RETURN_IF_NOT(s == TensorShape({n * k_blocks}) || s == TensorShape({n, k_blocks}), - "MatMulNBits PrePack: zero_points initializer shape ", s, - " does not match attribute-derived shape [", n * k_blocks, "] or [", - n, ",", k_blocks, "]"); - } - return Status::OK(); - }; - - const TensorShape& shape = tensor.Shape(); - - if (input_idx == InputIndex::B) { - ORT_RETURN_IF_NOT(shape == TensorShape({n, k_blocks, blob_size}), - "MatMulNBits PrePack: B initializer shape ", shape, - " does not match attribute-derived shape [", n, ",", k_blocks, ",", blob_size, - "] (N=", N_, ", K=", K_, ", bits=", nbits_, ", block_size=", block_size_, ")"); - - // Also validate constant scales / zero_points, which the B prepack path below dereferences - // (via TryGetConstantInput) and hands to MLAS, before their own PrePack calls run. - const Tensor* scales_tensor = nullptr; - if (OpKernel::Info().TryGetConstantInput(InputIndex::scales, &scales_tensor) && scales_tensor != nullptr) { - ORT_RETURN_IF_ERROR(validate_scales_shape(scales_tensor->Shape())); - } - const Tensor* zp_tensor = nullptr; - if (has_zp_arg_ && has_zp_input_ && - OpKernel::Info().TryGetConstantInput(InputIndex::zero_points, &zp_tensor) && zp_tensor != nullptr) { - ORT_RETURN_IF_ERROR(validate_zero_points_shape(zp_tensor->Shape(), zp_tensor->GetElementType())); - } - } else if (input_idx == InputIndex::scales) { - ORT_RETURN_IF_ERROR(validate_scales_shape(shape)); - } else if (input_idx == InputIndex::zero_points) { - ORT_RETURN_IF_ERROR(validate_zero_points_shape(shape, tensor.GetElementType())); - } - } + ORT_RETURN_IF_ERROR(ValidatePrePackInitializerShape(OpKernel::Info(), tensor, input_idx, + N_, K_, block_size_, nbits_, + has_zp_arg_, has_zp_input_)); if (has_g_idx_) { return Status::OK(); @@ -741,6 +728,10 @@ template <> Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ AllocatorPtr alloc, /*out*/ bool& is_packed, /*out*/ PrePackedWeights* prepacked_weights) { + ORT_RETURN_IF_ERROR(ValidatePrePackInitializerShape(OpKernel::Info(), tensor, input_idx, + N_, K_, block_size_, nbits_, + has_zp_arg_, has_zp_input_)); + if (input_idx == InputIndex::scales || input_idx == InputIndex::bias) { auto sptr = tensor.Data(); auto tensor_size = static_cast(tensor.Shape().Size()); diff --git a/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc b/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc index ff52d9d0444a2..11ee80db5c744 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc @@ -104,6 +104,8 @@ QLinearSoftmax::QLinearSoftmax(const OpKernelInfo& info) Status QLinearSoftmax::Compute(OpKernelContext* ctx) const { const auto* X = ctx->Input(0); const auto& X_shape = X->Shape(); + auto* Y = ctx->Output(0, X_shape); + // edge case. one or more dims with value of 0. nothing to do if (X_shape.Size() == 0) { return Status::OK(); @@ -111,8 +113,6 @@ Status QLinearSoftmax::Compute(OpKernelContext* ctx) const { auto axis = static_cast(HandleNegativeAxis(axis_, int64_t(X_shape.NumDimensions()))); - auto* Y = ctx->Output(0, X_shape); - concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool(); const size_t D = onnxruntime::narrow(opset_ < OPSET13 ? X_shape.SizeFromDimension(onnxruntime::narrow(axis)) : X_shape[onnxruntime::narrow(axis)]); EXP_OUT_DTYPE tmp_lookup_table[256]; @@ -131,7 +131,7 @@ common::Status QlinearSoftmaxCPU(size_t N, const T* x_data, T* y_data, const QLinearSoftmax::EXP_OUT_DTYPE* lookup_table, - QLinearSoftmax::EXP_OUT_DTYPE y_scale, + double y_scale, T yzp, onnxruntime::concurrency::ThreadPool* thread_pool); @@ -141,7 +141,7 @@ common::Status QlinearSoftmaxCPU(size_t N, const uint8_t* x_data, uint8_t* y_data, const QLinearSoftmax::EXP_OUT_DTYPE* lookup_table, - QLinearSoftmax::EXP_OUT_DTYPE y_scale, + double y_scale, uint8_t yzp, onnxruntime::concurrency::ThreadPool* thread_pool) { using onnxruntime::TensorOpCost; @@ -187,8 +187,9 @@ common::Status QlinearSoftmaxCPU(size_t N, const size_t vx = *x_t_cur++; const QLinearSoftmax::EXP_OUT_DTYPE vt = shifted_lookuptable[vx]; // simulate round function, and re-quant to uint8 - const uint32_t vq = static_cast(std::nearbyintf(((vt * c_y_scale)) / vsum)) + c_y_zp; - const uint8_t vy = vq > 255 ? static_cast(255) : static_cast(vq); + // Use a wide intermediate and saturate before converting to an integer. + const double vq = std::nearbyint((static_cast(vt) * c_y_scale) / static_cast(vsum)) + c_y_zp; + const uint8_t vy = static_cast(std::clamp(vq, 0.0, 255.0)); *y_t++ = vy; } while (--elements_n != 0); x_t = x_t_cur; @@ -204,7 +205,7 @@ common::Status QlinearSoftmaxCPU(size_t N, const int8_t* x_data, int8_t* y_data, const QLinearSoftmax::EXP_OUT_DTYPE* lookup_table, - QLinearSoftmax::EXP_OUT_DTYPE y_scale, + double y_scale, int8_t yzp, onnxruntime::concurrency::ThreadPool* thread_pool) { using onnxruntime::TensorOpCost; @@ -218,8 +219,8 @@ common::Status QlinearSoftmaxCPU(size_t N, [x_data, y_data, D, y_scale, yzp, &lookup_table](std::ptrdiff_t first, std::ptrdiff_t last) { const auto c_y_scale = y_scale; const auto c_y_zp = yzp; - constexpr int32_t low = std::numeric_limits::lowest(); - constexpr int32_t high = std::numeric_limits::max(); + constexpr double low = std::numeric_limits::lowest(); + constexpr double high = std::numeric_limits::max(); const int8_t* x_t = x_data + first * D; int8_t* y_t = y_data + first * D; @@ -246,7 +247,7 @@ common::Status QlinearSoftmaxCPU(size_t N, const uint8_t vx = uint8_t(adjustment + (*x_t_cur++)); const QLinearSoftmax::EXP_OUT_DTYPE vt = shifted_lookuptable[vx]; // simulate round function, and re-quant to Int8 - const int32_t vq = static_cast(std::nearbyintf(((vt * c_y_scale)) / vsum)) + c_y_zp; + const double vq = std::nearbyint((static_cast(vt) * c_y_scale) / static_cast(vsum)) + c_y_zp; const int8_t vy = static_cast(std::clamp(vq, low, high)); *y_t++ = vy; } while (--elements_n != 0); @@ -276,7 +277,12 @@ Status QLinearSoftmax::ComputeInternal(OpKernelContext* context, const Tensor& i concurrency::ThreadPool* thread_pool) const { const auto* Y_scale_tensor = context->Input(3); const auto* Y_zp_tensor = context->Input(4); - const QLinearSoftmax::EXP_OUT_DTYPE Y_scale = std::floor(1.0F / (*(Y_scale_tensor->Data()))); + const float output_scale = *(Y_scale_tensor->Data()); + double Y_scale = std::floor(1.0F / output_scale); + if (!std::isfinite(Y_scale)) { + // Keep the usual reciprocal rounding, but allow every finite positive float scale. + Y_scale = std::floor(1.0 / static_cast(output_scale)); + } const auto& X_shape = input.Shape(); const size_t N = onnxruntime::narrow(X_shape.SizeToDimension(onnxruntime::narrow(axis))); const size_t D = onnxruntime::narrow(X_shape.SizeFromDimension(onnxruntime::narrow(axis))); diff --git a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_merge_common.h b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_merge_common.h index 30bb3d96659b7..f30c837563882 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_merge_common.h +++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_merge_common.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #ifndef SHARED_PROVIDER @@ -24,9 +25,16 @@ struct Dimensions { int32_t rows; int32_t capacity; int32_t queries; - int32_t hash_capacity; + int64_t hash_capacity; }; +inline int64_t HashCapacity(int64_t output_capacity, int64_t candidate_capacity) { + const int64_t insertion_bound = std::min(output_capacity + 1, candidate_capacity); + int64_t capacity = 1; + while (capacity < 2 * insertion_bound) capacity *= 2; + return capacity; +} + inline Status Validate(const Tensor* base, const Tensor* counts, const Tensor* rows, const Tensor* starts, const Tensor* ends, int64_t output_capacity, Dimensions& dimensions) { @@ -41,10 +49,8 @@ inline Status Validate(const Tensor* base, const Tensor* counts, const Tensor* r ORT_RETURN_IF_NOT(counts->Shape() == TensorShape({cached_rows}) && starts->Shape() == rows->Shape() && ends->Shape() == rows->Shape(), "Counts and range tensors must match their row dimensions"); - int64_t hash_capacity = 1; - while (hash_capacity < output_capacity) hash_capacity *= 2; dimensions = {static_cast(cached_rows), static_cast(base_capacity), - static_cast(queries), static_cast(hash_capacity)}; + static_cast(queries), HashCapacity(output_capacity, base_capacity)}; return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/activation/activations_impl.cu b/onnxruntime/contrib_ops/cuda/activation/activations_impl.cu index e99f380e6fe02..5b76b5d346a65 100644 --- a/onnxruntime/contrib_ops/cuda/activation/activations_impl.cu +++ b/onnxruntime/contrib_ops/cuda/activation/activations_impl.cu @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include +#include #include "activations_impl.h" #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cu_inc/unary_elementwise_impl.cuh" @@ -39,11 +40,22 @@ struct OP_ScaledTanh : public CtxScaledTanh { template struct OP_QuickGelu : public CtxQuickGelu { __device__ __inline__ T operator()(const T& a) const { - T v = a * static_cast(alpha); - T one = static_cast(1.f); - T zero = static_cast(0.f); - T sigmoid = v >= zero ? one / (one + _Exp(-v)) : one - one / (one + _Exp(v)); - return a * sigmoid; + using ComputeT = std::conditional_t, double, float>; + const ComputeT value = static_cast(a); + const ComputeT scaled = value * static_cast(alpha); + if constexpr (std::is_same_v) { + // Float exp enters its subnormal range near -87; keep the final product wide. + if (scaled < -87.0f) { + const double wide_value = static_cast(value); + const double exponent = _Exp(wide_value * static_cast(alpha)); + return static_cast(static_cast(wide_value * exponent / (1.0 + exponent))); + } + } + const ComputeT exponent = _Exp(scaled >= ComputeT(0) ? -scaled : scaled); + // Preserve negative tails and round only after the fused sigmoid/product. + const ComputeT sigmoid = scaled >= ComputeT(0) ? ComputeT(1) / (ComputeT(1) + exponent) + : exponent / (ComputeT(1) + exponent); + return static_cast(value * sigmoid); } }; diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc index ee21a71c2b4aa..1c91e9619364e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc @@ -9,24 +9,29 @@ #include "contrib_ops/cpu/hyper_connection_helper.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" +#include "core/providers/cuda/reduction/reduction_functions.h" namespace onnxruntime::contrib::cuda { using namespace onnxruntime::cuda; -#define REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(T) \ +#define REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(V, T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ - BranchwiseRMSNorm, kMSDomain, 1, T, kCudaExecutionProvider, \ + BranchwiseRMSNorm, kMSDomain, V, T, kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("F", DataTypeImpl::GetTensorType()) \ .TypeConstraint("M", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}), \ BranchwiseRMSNorm); -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(float) -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(MLFloat16) -REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(BFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(1, float) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(1, MLFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(1, BFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(2, float) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(2, MLFloat16) +REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL(2, BFloat16) #undef REGISTER_BRANCHWISE_RMS_NORM_CUDA_KERNEL @@ -34,30 +39,70 @@ template BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) : CudaKernel(info), epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - num_branches_(info.GetAttrOrDefault("num_branches", 0)) {} + num_branches_(info.GetAttrOrDefault("num_branches", 0)), + use_onnx_arithmetic_(info.GetAttrOrDefault("use_onnx_arithmetic", 0) != 0), + input_start_(info.GetAttrOrDefault("input_start", 0)), + input_width_(info.GetAttrOrDefault("input_width", 0)) { + const auto mode = info.GetAttrOrDefault("use_onnx_arithmetic", 0); + ORT_ENFORCE(mode == 0 || mode == 1, "use_onnx_arithmetic must be 0 or 1"); + ORT_ENFORCE(input_start_ >= 0 && input_width_ >= 0, "input_start and input_width must be nonnegative"); + ORT_ENFORCE(input_width_ != 0 || input_start_ == 0, "input_start requires input_width"); + ORT_ENFORCE(input_width_ == 0 || (use_onnx_arithmetic_ && num_branches_ == 0), + "input_width requires use_onnx_arithmetic=1 and grouped input"); +} template Status BranchwiseRMSNorm::ComputeInternal(OpKernelContext* context) const { using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; const auto* x = context->Input(0); const auto* scale = context->Input(1); + auto output_dims = x->Shape().AsShapeVector(); + ORT_RETURN_IF_NOT(!output_dims.empty(), "BranchwiseRMSNorm requires a non-scalar input"); + const int64_t input_stride = output_dims.back(); + if (input_width_ != 0) { + ORT_RETURN_IF_NOT(input_start_ <= input_stride && input_width_ <= input_stride - input_start_, + "input slice exceeds the last dimension"); + output_dims.back() = input_width_; + } + const TensorShape output_shape(output_dims); hyper_connection::StreamShape params; - ORT_RETURN_IF_ERROR(hyper_connection::ResolveStreamShape(x->Shape(), num_branches_, params)); + ORT_RETURN_IF_ERROR(hyper_connection::ResolveStreamShape(output_shape, num_branches_, params)); if (scale != nullptr) { ORT_RETURN_IF_ERROR(hyper_connection::ValidateScale(scale->Shape(), params)); } ORT_RETURN_IF_NOT(params.branches <= std::numeric_limits::max() && - params.hidden <= std::numeric_limits::max(), + params.hidden <= std::numeric_limits::max() && + (input_width_ == 0 || input_stride <= std::numeric_limits::max()), "branch or hidden dimension is too large for CUDA"); - const int64_t groups = x->Shape().Size() / params.hidden; - auto* y = context->Output(0, x->Shape()); + const int64_t groups = output_shape.Size() / params.hidden; + auto* y = context->Output(0, output_shape); + auto* widened = context->OutputCount() > 1 ? context->Output(1, output_shape) : nullptr; + ORT_RETURN_IF_NOT(widened == nullptr || use_onnx_arithmetic_, + "Y_fp32 requires use_onnx_arithmetic=1"); if (groups == 0) { return Status::OK(); } const auto* x_data = reinterpret_cast(x->Data()); auto* y_data = reinterpret_cast(y->MutableData()); + if (use_onnx_arithmetic_) { + ORT_RETURN_IF_NOT(groups <= std::numeric_limits::max(), + "ONNX arithmetic requires the row count to fit in int32"); + const auto buffer_size = compute_reduce_matrix_columns_buffer_size( + static_cast(groups), static_cast(params.hidden)); + auto squares = GetScratchBuffer(groups == 1 ? output_shape.Size() : 0, GetComputeStream(context)); + auto sums = GetScratchBuffer(groups, GetComputeStream(context)); + auto buffer = GetScratchBuffer(buffer_size, GetComputeStream(context)); + return LaunchStandardBranchwiseRMSNorm( + Stream(context), x_data + input_start_, scale == nullptr ? nullptr : scale->DataRaw(), + scale == nullptr ? x->GetElementType() : scale->GetElementType(), y_data, + static_cast(groups), static_cast(params.branches), static_cast(params.hidden), + scale == nullptr || scale->Shape().Size() == params.hidden, epsilon_, + widened == nullptr ? nullptr : widened->MutableData(), + squares.get(), sums.get(), buffer.get(), buffer_size, + static_cast(input_width_ == 0 ? params.hidden : input_stride)); + } return LaunchMixedScaleBranchwiseRMSNorm( Stream(context), x_data, scale == nullptr ? nullptr : scale->DataRaw(), scale == nullptr ? x->GetElementType() : scale->GetElementType(), y_data, diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h index 959fa12b4a46a..c26a8d47b41ff 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h @@ -16,6 +16,9 @@ class BranchwiseRMSNorm final : public onnxruntime::cuda::CudaKernel { private: float epsilon_; int64_t num_branches_; + bool use_onnx_arithmetic_; + int64_t input_start_; + int64_t input_width_; }; } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu index a56e55ce21331..240b057b827e0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu @@ -7,10 +7,12 @@ #include #include #include +#include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" #include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/reduction/reduction_functions.h" namespace onnxruntime::contrib::cuda { namespace { @@ -29,6 +31,34 @@ __device__ __forceinline__ float LoadScale(const void* data, int64_t index, int return __bfloat162float(static_cast(data)[index]); } +template +__global__ void BranchwiseSquareKernel(const T* x, float* squares, int64_t count) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) { + const float value = to_float(x[index]); + squares[index] = __fmul_rn(value, value); + } +} + +template +__global__ void BranchwiseStandardEpilogueKernel( + const T* x, const void* scale, int scale_type, const float* sums, T* y, float* widened, + int64_t count, int branches, int hidden, bool shared_scale, float epsilon, int input_row_stride) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) { + const int64_t group = index / hidden; + const int64_t scale_index = shared_scale ? index % hidden : index % (static_cast(branches) * hidden); + const float weight = scale == nullptr ? 1.0f : LoadScale(scale, scale_index, scale_type); + const float mean = sums[group] / static_cast(hidden); + const float rms = sqrtf(__fadd_rn(mean, epsilon)); + const int64_t source_index = Packed ? group * input_row_stride + index % hidden : index; + const float normalized = to_float(x[source_index]) / rms; + const T rounded = from_float(__fmul_rn(normalized, weight)); + y[index] = rounded; + if (widened != nullptr) widened[index] = to_float(rounded); + } +} + template __global__ void MixedScaleBranchwiseRMSNormKernel(const T* x, const void* scale, int scale_type, T* y, int64_t groups, int branches, int hidden, @@ -77,6 +107,37 @@ __global__ void MixedScaleBranchwiseRMSNormKernel(const T* x, const void* scale, } // namespace +template +Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const void* scale, + int scale_type, T* y, int groups, int branches, + int hidden, bool shared_scale, float epsilon, + float* widened, float* squares, float* sums, void* reduction_buffer, + size_t reduction_buffer_size, int input_row_stride) { + const int64_t count = static_cast(groups) * hidden; + const int64_t blocks = (count + kThreads - 1) / kThreads; + ORT_RETURN_IF_NOT(blocks <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); + // Standard ReduceMean removes singleton dimensions and prefers row reduction for a scalar result. + if (groups == 1) { + BranchwiseSquareKernel<<(blocks), kThreads, 0, stream>>>(x, squares, count); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + ORT_RETURN_IF_ERROR((onnxruntime::cuda::reduce_matrix_rows( + stream, squares, sums, hidden, 1))); + } else { + using ReduceT = std::conditional_t, BFloat16, T>; + ORT_RETURN_IF_ERROR((onnxruntime::cuda::reduce_square_matrix_columns( + stream, reinterpret_cast(x), sums, groups, hidden, + reduction_buffer, reduction_buffer_size, input_row_stride))); + } + if (input_row_stride == hidden) { + BranchwiseStandardEpilogueKernel<<(blocks), kThreads, 0, stream>>>( + x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride); + } else { + BranchwiseStandardEpilogueKernel<<(blocks), kThreads, 0, stream>>>( + x, scale, scale_type, sums, y, widened, count, branches, hidden, shared_scale, epsilon, input_row_stride); + } + return CUDA_CALL(cudaGetLastError()); +} + template Status LaunchMixedScaleBranchwiseRMSNorm(cudaStream_t stream, const T* x, const void* scale, int scale_type, T* y, int64_t groups, int branches, @@ -107,4 +168,12 @@ INSTANTIATE(__nv_bfloat16) #undef INSTANTIATE +#define INSTANTIATE_STANDARD(T) \ + template Status LaunchStandardBranchwiseRMSNorm(cudaStream_t, const T*, const void*, int, T*, \ + int, int, int, bool, float, float*, float*, float*, void*, size_t, int); +INSTANTIATE_STANDARD(float) +INSTANTIATE_STANDARD(half) +INSTANTIATE_STANDARD(__nv_bfloat16) +#undef INSTANTIATE_STANDARD + } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h index 549de28e62bff..a037130d825a2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h @@ -14,4 +14,11 @@ Status LaunchMixedScaleBranchwiseRMSNorm(cudaStream_t stream, const T* x, const int scale_type, T* y, int64_t groups, int branches, int hidden, bool shared_scale, float epsilon); +template +Status LaunchStandardBranchwiseRMSNorm(cudaStream_t stream, const T* x, const void* scale, + int scale_type, T* y, int groups, int branches, + int hidden, bool shared_scale, float epsilon, + float* widened, float* squares, float* sums, void* reduction_buffer, + size_t reduction_buffer_size, int input_row_stride); + } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc index 0fbb58e0f46ef..ccfb284012189 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc @@ -391,6 +391,27 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( aggregate.sized_backends == GQAReachableBackend::None) { aggregate.status = Unavailable("No GQA backend route can be sized."); } + + // A non-windowed step that finds exactly one past/present K/V pair aliasing + // copies the full past cache into preservation scratch before the backend + // runs (group_query_attention.cc separate_past_buffer). That copy coexists + // with the selected route's workspace, so charge its worst case on top of the + // per-route maximum. present_kv_cache_capacity_bound covers the full allocated + // past capacity, which can exceed the active total KV length for a static cache. + // The stored per-head extent (halved for 4-bit) never exceeds head_size_bound. + if (aggregate.status.IsOK() && bounds.account_partial_alias_preservation) { + size_t preservation = 0; + auto status = Mul4(static_cast(bounds.batch_size_bound), + static_cast(bounds.kv_num_heads), + static_cast(bounds.present_kv_cache_capacity_bound), + static_cast(bounds.head_size_bound), preservation); + if (status.IsOK()) status = Mul(preservation, bounds.cache_element_size, preservation); + if (status.IsOK()) { + status = Add(aggregate.total_workspace_bytes, preservation, + aggregate.total_workspace_bytes); + } + if (!status.IsOK()) aggregate.status = status; + } return aggregate; } diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h index 0783f90172a0c..bcb2f92711595 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h @@ -35,10 +35,11 @@ constexpr bool HasGQAReachableBackend(GQAReachableBackend mask, // Graph-free upper bounds and immutable feature facts. Positive dimensions are // componentwise bounds, not a claim that their combination is one runtime shape. -// The aggregate assumes both past/present K/V pairs alias. A caller that permits -// partial aliasing must separately account for preservation scratch. For a -// non-windowed cache, present_kv_cache_capacity_bound must bound the runtime -// present KV length, not only the past input tensor's sequence dimension. +// The aggregate assumes both past/present K/V pairs alias unless +// account_partial_alias_preservation is enabled, in which case it adds the +// worst-case preservation scratch for partial aliasing. For a non-windowed +// cache, present_kv_cache_capacity_bound must bound both the allocated past +// capacity and the active total KV length. struct GQAWorkspaceBounds { size_t qkv_element_size = 0; size_t cache_element_size = 0; @@ -52,6 +53,13 @@ struct GQAWorkspaceBounds { GQAKvQuantizationType k_quantization = GQAKvQuantizationType::None; GQAKvQuantizationType v_quantization = GQAKvQuantizationType::None; bool is_windowed_kv_cache = false; + // Non-windowed GQA permits exactly one past/present K/V pair to alias, which + // makes the runtime copy the full past cache into preservation scratch that + // coexists with the selected backend workspace. Windowed GQA rejects + // non-shared buffers, so both pairs always alias and no copy occurs. When + // true, the aggregate adds the worst-case preservation copy on top of the + // per-route maximum. See group_query_attention.cc (separate_past_buffer). + bool account_partial_alias_preservation = false; bool do_rotary = false; bool is_packed_qkv = false; bool use_qk_norm = false; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc index 19d6eb404bee1..68ad24bbf7f71 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc @@ -199,14 +199,17 @@ std::optional BuildBounds( // WorkspaceInputShape exposes the total_sequence_length scalar's shape, not // its value. For a non-windowed cache that value can exceed past dim 2 and - // directly scales backend workspace. Non-windowed execution can also copy a - // full past tensor when only one past/present pair aliases, but alias state is - // unavailable here. Windowed execution is bounded by the cache shape and the - // runtime requires both past/present pairs to alias before allocating scratch. - // config.max_total_sequence_length carries the externally declared KV-length envelope - // (session option ep.cuda.gqa_workspace_max_total_sequence_length) intended to bound this - // non-windowed path; it is plumbed but not yet consumed here, so behavior is unchanged. - if (!config.sliding_window_cache) return std::nullopt; + // directly scales backend workspace, so a sound bound needs the KV-length + // envelope config.max_total_sequence_length (session option + // ep.cuda.gqa_workspace_max_total_sequence_length). Without it there is no + // sound non-windowed bound, so decline -- this preserves the prior behavior + // whenever the envelope is unset. Windowed execution is bounded by the cache + // shape. The other non-windowed concern -- a full past-tensor copy when only + // one past/present pair aliases -- is charged separately as preservation + // scratch via bounds.account_partial_alias_preservation below; windowed + // execution requires both pairs to alias, so it never incurs that copy. + const bool windowed = config.sliding_window_cache; + if (!windowed && config.max_total_sequence_length <= 0) return std::nullopt; const bool packed = !Present(shapes, kKey); const TensorShape* query = Shape(shapes, kQuery); @@ -235,9 +238,8 @@ std::optional BuildBounds( } if (head_bound < 8 || batch_bound <= 0 || sequence_bound <= 0) return std::nullopt; - // Cache capacity is the only sound bound for total KV length available at - // this API. Without past tensors, the CPU scalar value is unavailable and a - // single-token invocation can name an arbitrarily larger total length. + // Past tensors are required for both cache kinds: they carry the kv-head and + // per-head cache geometry and allocated capacity validated below. if (!Present(shapes, kPastKey)) return std::nullopt; const auto* past_key = Shape(shapes, kPastKey); const auto* past_value = Shape(shapes, kPastValue); @@ -251,7 +253,11 @@ std::optional BuildBounds( return std::nullopt; } if ((*past_key)[2] != (*past_value)[2]) return std::nullopt; - const int64_t capacity_bound = (*past_key)[2]; + // A non-windowed static cache can exceed the active total KV length. MEA + // expands the present capacity, and partial aliasing preserves the full past + // allocation, so bound both the allocated cache and the declared KV length. + const int64_t capacity_bound = + windowed ? (*past_key)[2] : std::max((*past_key)[2], config.max_total_sequence_length); int64_t cache_head_bound = std::min((*past_key)[3], (*past_value)[3]); if (config.kv_cache_bit_width == 4) { if (cache_head_bound > std::numeric_limits::max() / 2) return std::nullopt; @@ -259,13 +265,17 @@ std::optional BuildBounds( } head_bound = std::min(head_bound, cache_head_bound); if (batch_bound <= 0 || capacity_bound <= 0 || head_bound < 8 || - capacity_bound != config.local_window_size || + (windowed && capacity_bound != config.local_window_size) || + (!windowed && (config.max_total_sequence_length < sequence_bound || + capacity_bound > std::numeric_limits::max())) || !ValidateAuxiliaryShapes(config, shapes, batch_bound, sequence_bound, head_bound, capacity_bound)) { return std::nullopt; } - // Runtime rejects attention bias for sliding-window GQA. + // Attention bias is declined for both cache kinds here: sliding-window GQA + // rejects it at runtime, and the non-windowed estimate does not yet model the + // bias-only unfused route. Declining keeps the estimate a sound upper bound. if (Present(shapes, kAttentionBias)) return std::nullopt; const bool k_quantized = config.k_quantization != GQAKvQuantizationType::None; @@ -298,6 +308,7 @@ std::optional BuildBounds( bounds.k_quantization = config.k_quantization; bounds.v_quantization = config.v_quantization; bounds.is_windowed_kv_cache = config.sliding_window_cache; + bounds.account_partial_alias_preservation = !windowed; bounds.do_rotary = config.do_rotary; bounds.is_packed_qkv = packed; bounds.use_qk_norm = Present(shapes, kQNorm); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h index 6d331273e4e4d..7af799f27694e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h @@ -47,9 +47,9 @@ struct GQAWorkspaceEstimateConfig { bool head_sink_is_prepacked = false; // Level-1 possibility; true charges persistent and initialization lifetimes. bool head_sink_may_be_prepacked = false; - // Reserved, currently unconsumed total KV-length envelope from session option - // ep.cuda.gqa_workspace_max_total_sequence_length. Zero means unspecified; - // this field does not enforce a runtime input limit. + // Caller-declared total KV-length envelope from session option + // ep.cuda.gqa_workspace_max_total_sequence_length. It bounds non-windowed + // estimation; zero means unspecified. This is not a runtime input limit. int64_t max_total_sequence_length = 0; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc index 0b38c6f0e5484..722e99b096418 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.cc @@ -15,11 +15,11 @@ namespace onnxruntime { namespace contrib { namespace cuda { -#define REGISTER_KERNEL_TYPED(T) \ +#define REGISTER_KERNEL_TYPED(T, Version) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ RotaryEmbedding, \ kMSDomain, \ - 1, \ + Version, \ T, \ kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ @@ -27,9 +27,12 @@ namespace cuda { .TypeConstraint("M", DataTypeImpl::GetTensorType()), \ RotaryEmbedding); -REGISTER_KERNEL_TYPED(float) -REGISTER_KERNEL_TYPED(MLFloat16) -REGISTER_KERNEL_TYPED(BFloat16) +REGISTER_KERNEL_TYPED(float, 1) +REGISTER_KERNEL_TYPED(MLFloat16, 1) +REGISTER_KERNEL_TYPED(BFloat16, 1) +REGISTER_KERNEL_TYPED(float, 2) +REGISTER_KERNEL_TYPED(MLFloat16, 2) +REGISTER_KERNEL_TYPED(BFloat16, 2) template RotaryEmbedding::RotaryEmbedding(const OpKernelInfo& info) : CudaKernel(info) { @@ -38,6 +41,9 @@ RotaryEmbedding::RotaryEmbedding(const OpKernelInfo& info) : CudaKernel(info) num_heads = static_cast(info.GetAttrOrDefault("num_heads", 0)); interleaved = (info.GetAttrOrDefault("interleaved", 0) == 1); is_packed_batching = (info.GetAttrOrDefault("is_packed_batching", 0) == 1); + const auto mode = info.GetAttrOrDefault("use_fp32_arithmetic", 0); + ORT_ENFORCE(mode == 0 || mode == 1, "use_fp32_arithmetic must be 0 or 1"); + use_fp32_arithmetic = mode == 1; } template @@ -83,7 +89,8 @@ Status RotaryEmbedding::ComputeInternal(OpKernelContext* context) const { parameters.position_ids_format, interleaved, device_prop.maxThreadsPerBlock, - parameters.transposed); + parameters.transposed, + use_fp32_arithmetic); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h index 55f654a2cea8f..4487746579a9a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding.h @@ -23,6 +23,7 @@ class RotaryEmbedding final : public CudaKernel { int rotary_embedding_dim; bool interleaved; bool is_packed_batching; + bool use_fp32_arithmetic; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu index 47c46b73f7b83..55648cc944249 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu @@ -10,6 +10,7 @@ Kernel implementation for rotary embeddings. #include "contrib_ops/cuda/bert/rotary_embedding_impl.h" #include +#include #include "core/common/safeint.h" #include "core/providers/cuda/cu_inc/common.cuh" @@ -20,7 +21,7 @@ namespace onnxruntime { namespace contrib { namespace cuda { -template +template __global__ void RotaryEmbeddingBSNH(T* output, // BxSxNxH const T* input, // BxSxNxH const T* cos_cache, // Mx(H/2) @@ -128,7 +129,14 @@ __global__ void RotaryEmbeddingBSNH(T* output, // BxSxNx } // Use values from shared memory - if constexpr (use_smem) { + if constexpr (use_fp32_arithmetic) { + const float a = static_cast(use_smem ? smem[i] : input_data[i]); + const float b = static_cast(use_smem ? smem[j] : input_data[j]); + const float cosine = static_cast(cos_data[cache_idx]); + const float sine = static_cast(sin_data[cache_idx]); + const float signed_b = sign < T(0) ? -b : b; + output_data[i] = T(__fadd_rn(__fmul_rn(a, cosine), __fmul_rn(signed_b, sine))); + } else if constexpr (use_smem) { output_data[i] = smem[i] * cos_data[cache_idx] + sign * smem[j] * sin_data[cache_idx]; } else { output_data[i] = input_data[i] * cos_data[cache_idx] + sign * input_data[j] * sin_data[cache_idx]; @@ -142,7 +150,8 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, - const int max_threads_per_block, const bool is_input_bnsh_format) { + const int max_threads_per_block, const bool is_input_bnsh_format, + const bool use_fp32_arithmetic) { int sequence_stride; int batch_stride; int heads_stride; @@ -168,7 +177,7 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu rotary_embedding_dim, max_sequence_length, position_ids_format, interleaved, max_threads_per_block, - in_strides, out_strides); + in_strides, out_strides, use_fp32_arithmetic); } template @@ -179,8 +188,8 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, const int max_threads_per_block, - int4 in_strides, int4 out_strides // strides in bnsh coord -) { + int4 in_strides, int4 out_strides, // strides in bnsh coord + const bool use_fp32_arithmetic) { // Note: Current implementation assumes head_size <= max_threads_per_block // because head_size is currently large for LLaMA-2. For smaller head_size // and num_heads values, we can create a block as `block(num_heads, head_size, 1)` @@ -196,20 +205,25 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu assert(head_size <= max_threads_per_block); - if (output == input) { - // In-place operation: use shared memory to avoid read-after-write hazards - size_t smem_size = head_size * sizeof(T); - RotaryEmbeddingBSNH<<>>( - output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, - num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, - interleaved, in_strides, out_strides); - } else { - // Separate buffers: no shared memory needed - RotaryEmbeddingBSNH<<>>( - output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, - num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, - interleaved, in_strides, out_strides); - } + auto launch = [&](auto mode) { + constexpr bool fp32 = decltype(mode)::value; + if (output == input) { + size_t smem_size = head_size * sizeof(T); + RotaryEmbeddingBSNH<<>>( + output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, + num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, + interleaved, in_strides, out_strides); + } else { + RotaryEmbeddingBSNH<<>>( + output, input, cos_cache, sin_cache, position_ids, past_sequence_lengths, sequence_length, + num_heads, head_size, rotary_embedding_dim, max_sequence_length, position_ids_format, + interleaved, in_strides, out_strides); + } + }; + if (use_fp32_arithmetic) + launch(std::true_type{}); + else + launch(std::false_type{}); return CUDA_CALL(cudaGetLastError()); } @@ -220,7 +234,8 @@ template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, float* o const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, - const int max_threads_per_block, const bool is_input_bnsh_format); + const int max_threads_per_block, const bool is_input_bnsh_format, + const bool use_fp32_arithmetic); template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, half* output, const half* input, const int64_t* position_ids, const int* past_sequence_lengths, const half* cos_cache, @@ -228,14 +243,15 @@ template Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, half* out const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, - const int max_threads_per_block, const bool is_input_bnsh_format); + const int max_threads_per_block, const bool is_input_bnsh_format, + const bool use_fp32_arithmetic); template Status LaunchRotaryEmbeddingKernel( cudaStream_t stream, BFloat16* output, const BFloat16* input, const int64_t* position_ids, const int* past_sequence_lengths, const BFloat16* cos_cache, const BFloat16* sin_cache, const int batch_size, const int sequence_length, const int num_heads, const int head_size, const int rotary_embedding_dim, const int max_sequence_length, const int position_ids_format, const bool interleaved, const int max_threads_per_block, - const bool is_input_bnsh_format); + const bool is_input_bnsh_format, const bool use_fp32_arithmetic); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h index 2d81e6b4067af..4422572fc271d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.h @@ -27,7 +27,8 @@ Status LaunchRotaryEmbeddingKernel( const int position_ids_format, const bool interleaved, const int max_threads_per_block, - const bool is_input_bnsh_format); + const bool is_input_bnsh_format, + const bool use_fp32_arithmetic = false); template Status LaunchRotaryEmbeddingKernel( @@ -48,7 +49,8 @@ Status LaunchRotaryEmbeddingKernel( const bool interleaved, const int max_threads_per_block, int4 in_strides, - int4 out_strides); + int4 out_strides, + const bool use_fp32_arithmetic = false); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 546b5b12e178f..fe140646872de 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -160,6 +160,9 @@ class CUDA_ONNX_OP_TYPED_CLASS_NAME(1, MLFloat16, ParametricSoftplus); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, RotaryEmbedding); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, float, RotaryEmbedding); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, MLFloat16, RotaryEmbedding); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, BFloat16, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MRotaryEmbedding); @@ -179,6 +182,9 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, GatedRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, BranchwiseRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, BranchwiseRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, BranchwiseRMSNorm); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, float, BranchwiseRMSNorm); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, MLFloat16, BranchwiseRMSNorm); +class CUDA_MS_OP_TYPED_CLASS_NAME(2, BFloat16, BranchwiseRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, ScaledSiLU); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, ScaledSiLU); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, ScaledSiLU); @@ -524,6 +530,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -543,6 +552,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, 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 69471e6f954f2..925f9cb71347c 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc @@ -30,6 +30,18 @@ using namespace onnxruntime::llm::common; using namespace onnxruntime::llm::kernels::cutlass_kernels; +namespace { +// Compares tactics by their persisted representation, which is exactly what a cache hit stores. +bool SameConfigColumns(const std::optional& a, + const std::optional& b) { + std::vector columns_a; + std::vector columns_b; + onnxruntime::llm::gemm_cache::AppendConfigColumns(columns_a, a); + onnxruntime::llm::gemm_cache::AppendConfigColumns(columns_b, b); + return columns_a == columns_b; +} +} // namespace + namespace onnxruntime::llm::kernels::weight_only { std::optional> ComputeWeightOnlyGemmProfilerBufferSizes( @@ -214,8 +226,14 @@ bool WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic(int m, int n, int k fpA_intB_gemv::IsInt4DecodeGeometryLegal(tactic.cudaKernelVariant, SafeInt(n) * (FP16_BITS / mQuantBits), k, mDecodeInterleave); } - return m < 16 && (tactic.cudaKernelVariant == 0 || - (tactic.cudaKernelVariant == 1 && m >= 5 && m <= 8)); + const bool paired_bucket = m >= 5 && m <= 8; + if (tactic.cudaKernelVariant == 1) { + return mPairedGemvMode != 0 && paired_bucket; + } + return tactic.cudaKernelVariant == 0 && m < 16 && !(mPairedGemvMode == 2 && paired_bucket); + } + if (mPairedGemvMode == 2 && m >= 5 && m <= 8) { + return false; } return true; } @@ -314,5 +332,128 @@ std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::GetInitialProfileMB return std::vector(buckets.begin(), buckets.end()); } +onnxruntime::llm::gemm_cache::MatMulNBitsKey WeightOnlyGroupwiseQuantGemmPluginProfiler::makeCacheKey( + GemmIdCore const& gemmId, bool hasWeightOnlyCudaKernel) const { + onnxruntime::llm::gemm_cache::MatMulNBitsKey key; + key.n_16b = gemmId.n; + key.k = gemmId.k; + key.activation_dtype = (gemmId.dtype == onnxruntime::llm::nvinfer::DataType::kBF16) ? "bfloat16" : "half"; + key.weight_type = (mQuantBits == INT8_BITS) ? "uint8_t" + : (mQuantBits == INT2_BITS) ? "uint2b_t" + : "uint4b_t"; + key.bits = mQuantBits; + key.block_size = mGroupSize; + key.has_zero_points = mHasZeros; + key.zero_point_dtype = mHasZeros ? key.weight_type : "none"; + key.gemv_enabled = hasWeightOnlyCudaKernel; + key.has_bias = mHasBiases; + key.packing_sm = mKernelArch; + key.paired_gemv_mode = gemmId.tag; + key.wave_aware_gemv = gemmId.wave_aware; + return key; +} + +void WeightOnlyGroupwiseQuantGemmPluginProfiler::loadPersistentCache( + GemmIdCore const& gemmId, MProfileMap& map, bool hasWeightOnlyCudaKernel) { + if (mCache == nullptr || !mCacheAccess.load) { + return; + } + auto key = makeCacheKey(gemmId, hasWeightOnlyCudaKernel); + auto buckets = mCache->GetAll(key); + if (buckets.empty()) { + return; + } + + // Validate tactics loaded from disk against the tactics this runner can actually dispatch. A + // parseable-but-incompatible cache row (e.g. hand-edited, or written by a build whose signature + // happens to match but whose tactic set differs) would otherwise be handed straight to the kernel. + // Rejected rows are dropped so the bucket is re-profiled (and the stale row overwritten on staging). + // The CUDA-GEMV tactic is only dispatchable when this runner has a GEMV kernel, so it is rejected + // otherwise, regardless of what the row's key claims. + auto const valid_configs = getTactics(0, gemmId.n, gemmId.k); + auto is_valid_cutlass = [&valid_configs](Config const& c) { + for (auto const& v : valid_configs) { + if (v.sm_version == c.sm_version && v.is_tma_warp_specialized == c.is_tma_warp_specialized && + v.tile_config_sm80 == c.tile_config_sm80 && v.tile_config_sm90 == c.tile_config_sm90 && + v.tile_config_sm100 == c.tile_config_sm100 && v.tile_config_sm120 == c.tile_config_sm120 && + v.split_k_style == c.split_k_style && v.split_k_factor == c.split_k_factor && + v.stages == c.stages && v.cluster_shape == c.cluster_shape && + v.mainloop_schedule == c.mainloop_schedule && v.epilogue_schedule == c.epilogue_schedule) { + return true; + } + } + return false; + }; + + size_t accepted = 0; + for (auto const& [m, config] : buckets) { + if (!config.has_value() || m <= 0 || m > getMaxProfileM() || map.count(m) != 0) { + continue; + } + if (!checkTactic(m, gemmId.n, gemmId.k, *config)) { + ORT_LLM_LOG_WARNING("Dropping unsupported cached fpA_intB tactic from the tactic cache; re-profiling."); + continue; + } + if (config->enableCudaKernel && !hasWeightOnlyCudaKernel) { + ORT_LLM_LOG_WARNING("Dropping cached fpA_intB CUDA-GEMV tactic: no GEMV kernel is available; re-profiling."); + continue; + } + if (!config->enableCudaKernel && !is_valid_cutlass(*config)) { + ORT_LLM_LOG_WARNING("Dropping incompatible cached fpA_intB tactic from the tactic cache; re-profiling."); + continue; + } + if (!validatePersistentTactic(m, gemmId.n, gemmId.k, *config)) { + continue; + } + // Do not clobber tactics already selected in-process this session. + if (map.emplace(m, config).second) { + markCacheBucketStaged(mCache.get(), gemmId, m); + } + ++accepted; + } + if (accepted != 0) { + ORT_LLM_LOG_INFO("Loaded " + std::to_string(accepted) + " validated fpA_intB tactics from " + mCache->FilePath()); + } +} + +bool WeightOnlyGroupwiseQuantGemmPluginProfiler::stageProfiledTactics( + GemmIdCore const& gemmId, MProfileMap const& map, bool hasWeightOnlyCudaKernel) { + if (mCache == nullptr && mTuningResultsCache == nullptr) { + return false; + } + auto key = makeCacheKey(gemmId, hasWeightOnlyCudaKernel); + bool added = false; + for (auto const& [m, config] : map) { + if (!config.has_value()) { + continue; + } + if (mTuningResultsCache != nullptr && + markCacheBucketStaged(mTuningResultsCache.get(), gemmId, m)) { + mTuningResultsCache->Put(key, m, config); + } + if (mCache == nullptr || !mCacheAccess.save) { + continue; + } + if (!markCacheBucketStaged(mCache.get(), gemmId, m)) { + continue; + } + // Skip buckets already recorded with the same tactic (cache hits). A row rejected on load is still + // in the cache, so it must be overwritten by the freshly profiled tactic. + auto const cached = mCache->Get(key, m); + if (cached.has_value() && SameConfigColumns(*cached, config)) { + continue; + } + mCache->Put(key, m, config); + added = true; + } + return added; +} + +void WeightOnlyGroupwiseQuantGemmPluginProfiler::stagePersistentCache( + GemmIdCore const& gemmId, MProfileMap const& map, bool hasWeightOnlyCudaKernel) { + // Construction and lazy/shared-hit paths stage only. The staged tactics are written to disk at + // CUDA EP teardown (FlushMatMulNBitsTacticCaches in matmul_nbits.cc). + stageProfiledTactics(gemmId, map, hasWeightOnlyCudaKernel); +} } // 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 e8b9246ff2f78..ce5eb68365753 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h @@ -16,10 +16,12 @@ */ #pragma once +#include + +#include #include #include #include -#include #include #include #include @@ -30,6 +32,7 @@ #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 "contrib_ops/cuda/llm/gemm_tactic_cache.h" using WeightOnlyGemmRunner = onnxruntime::llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunnerInterface; using WeightOnlyGemmRunnerPtr = std::shared_ptr; @@ -54,16 +57,15 @@ constexpr const char* kEnvProfileM = "ORT_FPA_INTB_PROFILE_M"; // M values are handled by lazy single-bucket profiling. constexpr int kDefaultProfileMaxM = 2048; -// Computes the single temporary CUDA allocation used while profiling tactics. -// `packed_n` is the number of 16-bit elements that hold one packed weight row. -// This pure-math helper is shared by the runtime profiler and partition-time -// memory estimation so their allocation formulas cannot drift. // Maps a measured tactic time to the value compared during tactic selection. For small M, CUTLASS // must beat the CUDA GEMV by 10% when the weight fits in L2, because the profiler then times it // L2-resident while decode streams it from DRAM. float GetWeightOnlyGemmSelectionTime(int m, size_t weight_bytes, size_t l2_cache_bytes, bool is_cuda_kernel, float time); +// Computes the single temporary CUDA allocation used while profiling tactics. +// `packed_n` is the number of 16-bit elements that hold one packed weight row. +// Shared with partition-time memory estimation so the allocation formulas cannot drift. std::optional> ComputeWeightOnlyGemmProfilerBufferSizes( size_t max_m, size_t packed_n, size_t k, int quant_bits, size_t group_size, size_t runner_workspace_bytes, size_t streaming_l2_bytes = 0); @@ -119,6 +121,24 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler mL2CacheBytes = l2CacheBytes; } + // Attaches the process-global persistent tactic cache. A nullptr keeps the + // in-process-only behavior (no disk reads/writes). A non-null `session_profile_map` replaces the + // process-global profile map so this profiler deduplicates shapes within its own session while + // using `cache` for cross-session persistence (and re-tuning without mutating live sessions). + void setPersistentCache(std::shared_ptr cache, + onnxruntime::llm::gemm_cache::CacheAccess access = {}, + MNKProfileMapPtr session_profile_map = nullptr) { + mCache = std::move(cache); + mCacheAccess = access; + if (session_profile_map != nullptr) { + setSelectionTactics(session_profile_map); + } + } + + void setTuningResultsCache(std::shared_ptr cache) { + mTuningResultsCache = std::move(cache); + } + // Paired-K fp16 int4 GEMV tactic for the M = 8 bucket (M = 5..8 at run time). Mode 1 adds it as an extra // candidate, so it is kept only where it beats the default GEMV and the CUTLASS kernels; mode 2 offers // only that tactic (testing and benchmarking); mode 0 disables it. @@ -153,7 +173,22 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler std::vector getProfileMBuckets(int minM, int maxM, bool hasWeightOnlyCudaKernel) const override; + void loadPersistentCache(GemmIdCore const& gemmId, MProfileMap& map, + bool hasWeightOnlyCudaKernel) override; + + // Records buckets in `map` that are new or differ from the cache in memory (no disk write). + void stagePersistentCache(GemmIdCore const& gemmId, MProfileMap const& map, + bool hasWeightOnlyCudaKernel) override; + private: + onnxruntime::llm::gemm_cache::MatMulNBitsKey makeCacheKey(GemmIdCore const& gemmId, + bool hasWeightOnlyCudaKernel) const; + + // Populates the in-memory cache with any buckets in `map` not already recorded (no disk write). + // Returns true if at least one new bucket was staged. + bool stageProfiledTactics(GemmIdCore const& gemmId, MProfileMap const& map, + bool hasWeightOnlyCudaKernel); + bool mHasBiases = false; bool mHasZeros = false; int mQuantBits = 0; @@ -165,6 +200,9 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler size_t mL2CacheBytes = 0; std::atomic mProfileWeightIndex{0}; int mPairedGemvMode = 0; + std::shared_ptr mCache; + onnxruntime::llm::gemm_cache::CacheAccess mCacheAccess; + std::shared_ptr mTuningResultsCache; bool mWaveAwareGemv = false; std::vector mProfileMOverride; }; diff --git a/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h index 99ba6fb238865..4b552a489c691 100644 --- a/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h @@ -19,13 +19,21 @@ #include #include +#include +#include +#include +#include #include +#include #include #include #include #include #include +#include #include +#include +#include #include #include "contrib_ops/cuda/llm/nv_infer_datatype.h" @@ -90,39 +98,55 @@ struct GemmDims { // Unique ID of GEMM // In our case GEMM is uniquely identified by N and K, plus the target SM architecture (so the // SM80-compatibility and native SM90 kernels for the same shape do not share profiled configs). +// The quantization fields keep the process-global profile map aligned with the persistent cache key: +// kernels with different bits/group size/zero points/bias/GEMV support must not share tactics. +// device_name separates GPUs that share a packing SM (e.g. RTX 4090 and RTX 4060) so a tactic profiled +// on one device is never reused, or written to a disk cache, for another. +// device_id also isolates profiling state and allocator context on identical GPUs. class GemmIdCore { public: int n; int k; nvinfer::DataType dtype; - int packing_sm; - int device_sm; - // Distinguishes tactic candidate sets for the same shape. - int tag; - bool wave_aware = false; - int device_id; + int sm; int quant_bits; + int device_id; int group_size; - bool has_bias; bool has_zeros; - - GemmIdCore(int n_, int k_, nvinfer::DataType const& dtype_, int packing_sm_ = 0, int device_sm_ = 0, - bool wave_aware_ = false, int tag_ = 0, - int device_id_ = 0, int quant_bits_ = 0, int group_size_ = 0, - bool has_bias_ = false, bool has_zeros_ = false) - : n(n_), k(k_), dtype(dtype_), packing_sm(packing_sm_), device_sm(device_sm_), tag(tag_), wave_aware(wave_aware_), device_id(device_id_), quant_bits(quant_bits_), group_size(group_size_), has_bias(has_bias_), has_zeros(has_zeros_) { + bool gemv_enabled; + bool has_bias; + std::string device_name; + // Distinguishes optional GEMV tactic sets and launch heuristics for the same shape. + bool wave_aware = false; + int tag = 0; + + GemmIdCore(int n_, int k_, nvinfer::DataType const& dtype_, int sm_ = 0, int bits_ = 0, int group_size_ = 0, + bool has_zeros_ = false, bool gemv_enabled_ = false, bool has_bias_ = false, + std::string device_name_ = std::string(), bool wave_aware_ = false, int tag_ = 0, int device_id_ = 0) + : n(n_), + k(k_), + dtype(dtype_), + sm(sm_), + quant_bits(bits_), + device_id(device_id_), + group_size(group_size_), + has_zeros(has_zeros_), + gemv_enabled(gemv_enabled_), + has_bias(has_bias_), + device_name(std::move(device_name_)), + wave_aware(wave_aware_), + tag(tag_) { } GemmIdCore() : n(-1), k(-1), dtype(nvinfer::DataType::kFLOAT), // dtype does not matter here - packing_sm(0), - device_sm(0), - tag(0), - device_id(0), + sm(0), quant_bits(0), + device_id(0), group_size(0), - has_bias(false), - has_zeros(false) { + has_zeros(false), + gemv_enabled(false), + has_bias(false) { } bool operator==(GemmIdCore const& id) const { @@ -132,8 +156,10 @@ class GemmIdCore { friend std::ostream& operator<<(std::ostream& out, GemmIdCore const& id) { out << "(N;K)=(" << id.n << ";" << id.k << "),"; out << " type=" << static_cast(id.dtype); - out << " packing_sm=" << id.packing_sm; - out << " device_sm=" << id.device_sm; + out << " sm=" << id.sm; + out << " bits=" << id.quant_bits << " group_size=" << id.group_size; + out << " has_zeros=" << id.has_zeros << " gemv=" << id.gemv_enabled << " has_bias=" << id.has_bias; + out << " device=" << id.device_name << " device_id=" << id.device_id; out << " tag=" << id.tag; out << " wave_aware=" << id.wave_aware; return out; @@ -141,10 +167,10 @@ class GemmIdCore { protected: bool isEqual(GemmIdCore const& id) const { - return n == id.n && k == id.k && dtype == id.dtype && - packing_sm == id.packing_sm && device_sm == id.device_sm && device_id == id.device_id && - quant_bits == id.quant_bits && group_size == id.group_size && has_bias == id.has_bias && - has_zeros == id.has_zeros && tag == id.tag && wave_aware == id.wave_aware; + return n == id.n && k == id.k && dtype == id.dtype && sm == id.sm && quant_bits == id.quant_bits && + group_size == id.group_size && has_zeros == id.has_zeros && gemv_enabled == id.gemv_enabled && + has_bias == id.has_bias && device_name == id.device_name && device_id == id.device_id && tag == id.tag && + wave_aware == id.wave_aware; } }; @@ -154,16 +180,14 @@ struct GemmIdCoreHash { auto h1 = std::hash{}(id.n); auto h2 = std::hash{}(id.k); auto h3 = std::hash{}(static_cast(id.dtype)); - auto h4 = std::hash{}(id.packing_sm); - auto h5 = std::hash{}(id.device_sm); - auto h6 = std::hash{}(id.device_id); - auto h7 = std::hash{}(id.quant_bits); - auto h8 = std::hash{}(id.group_size); - auto h9 = std::hash{}(id.has_bias); - auto h10 = std::hash{}(id.has_zeros); - auto h11 = std::hash{}(id.tag); - auto h12 = std::hash{}(id.wave_aware); - return h1 ^ h2 ^ h3 ^ h4 ^ h5 ^ h6 ^ h7 ^ h8 ^ h9 ^ h10 ^ h11 ^ h12; + auto h4 = std::hash{}(id.sm); + auto h5 = std::hash{}((id.quant_bits << 16) ^ (id.group_size << 3) ^ (id.has_bias ? 4 : 0) ^ + (id.has_zeros ? 2 : 0) ^ (id.gemv_enabled ? 1 : 0)); + auto h6 = std::hash{}(id.device_name); + auto h7 = std::hash{}(id.tag); + auto h8 = std::hash{}(id.wave_aware); + auto h9 = std::hash{}(id.device_id); + return h1 ^ h2 ^ h3 ^ h4 ^ h5 ^ h6 ^ h7 ^ h8 ^ h9; } }; @@ -202,6 +226,18 @@ class GemmPluginProfiler { } return iter->second; } + + std::mutex stagedMutex; + std::unordered_map, GemmIdHashType>> + stagedCacheBuckets; + + bool markCacheBucketStaged(const void* cache_id, GemmIdType const& id, int m) { + if (cache_id == nullptr) { + return true; + } + std::lock_guard lock(stagedMutex); + return stagedCacheBuckets[cache_id][id].insert(m).second; + } }; using MNKProfileMapPtr = std::shared_ptr; @@ -220,6 +256,9 @@ class GemmPluginProfiler { void setSelectionTactics(MNKProfileMapPtr const& map) { mMNKProfileMap = map; + std::lock_guard staged_lock(mStagedBucketsMutex); + mStagedBuckets.clear(); + mStagedBucketCount.store(0, std::memory_order_release); } void setSkip(bool skip) { @@ -230,22 +269,54 @@ class GemmPluginProfiler { mAllocator = std::move(allocator); } - void setRunner(RunnerPtr const& runner) { - mRunner = runner; - } - - std::optional getBestConfig(int m, GemmIdType const& gemmId) const; + std::optional getBestConfig(int m, GemmIdType const& gemmId); // 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). + // sweep safe. A newly profiled bucket is staged into the persistent cache in memory only; it + // reaches disk at CUDA EP teardown, never on the inference path. 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; + void setMaxProfileM(int max_m) { + ORT_ENFORCE(max_m > 0, "The GEMM launch limit must be positive."); + mMaxProfileM = std::min(max_m, kMaxProfileM); + } + protected: + // An untimed launch checks the runner's problem-specific constraints and compiled kernel support. + // This is used only during construction, before a disk tactic can enter the shared profile map. + virtual bool validatePersistentTactic(int m, int n, int k, Config const& tactic) { + onnxruntime::IAllocatorUniquePtr workspace{nullptr}; + struct Stream { + cudaStream_t value = nullptr; + ~Stream() { + if (value != nullptr) { + cudaStreamSynchronize(value); + cudaStreamDestroy(value); + } + } + } stream; + try { + const size_t bytes = computeTmpSize(m, n, k); + workspace = onnxruntime::IAllocator::MakeUniquePtr(mAllocator, bytes, true); + CUDA_CALL_THROW(cudaStreamCreate(&stream.value)); + CUDA_CALL_THROW(cudaMemsetAsync(workspace.get(), 0, bytes, stream.value)); + initTmpData(m, n, k, workspace.get(), bytes, stream.value); + runTactic(m, n, k, tactic, workspace.get(), stream.value); + CUDA_CALL_THROW(cudaStreamSynchronize(stream.value)); + return true; + } catch (const std::exception& e) { + cudaGetLastError(); + ORT_LLM_LOG_WARNING(std::string("Rejecting cached GEMM tactic: ") + e.what()); + return false; + } + } + virtual void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) = 0; virtual size_t computeTmpSize(size_t maxM, size_t n, size_t k) = 0; @@ -269,7 +340,55 @@ class GemmPluginProfiler { // Subclasses may override to profile a smaller, configurable bucket set. virtual std::vector getProfileMBuckets(int minM, int maxM, bool hasWeightOnlyCudaKernel) const; + // Optional persistent (disk) tactic cache hooks. Default implementations are no-ops, which + // preserves the in-process-only behavior. Subclasses may override them to load matching + // tactics before the M sweep (populating `map` so profiling is skipped) and to stage profiled + // tactics in memory. Both are invoked while holding the profile-map writer lock. Staged tactics + // reach disk only when the session closes (CUDA EP teardown), never during session creation or + // inference. + virtual void loadPersistentCache(GemmIdType const& /*gemmId*/, MProfileMap& /*map*/, + bool /*hasWeightOnlyCudaKernel*/) {} + + virtual void stagePersistentCache(GemmIdType const& /*gemmId*/, MProfileMap const& /*map*/, + bool /*hasWeightOnlyCudaKernel*/) {} + + virtual bool hasRunnerContext() const { + return mRunner != nullptr; + } + + bool markCacheBucketStaged(const void* cache_id, GemmIdType const& gemmId, int m) { + return mMNKProfileMap == nullptr || mMNKProfileMap->markCacheBucketStaged(cache_id, gemmId, m); + } + private: + void stageNewPersistentBuckets(GemmIdType const& gemmId, MProfileMap const& map) { + // Profile maps only grow. A shared-map addition is staged once per attached profiler, + // including exact and rounded hits discovered by a different session. Only newly + // selected buckets are published so previously staged buckets cannot overwrite later + // retuned choices in a shared persistent cache. + if (mStagedBucketCount.load(std::memory_order_acquire) == map.size()) { + return; + } + std::lock_guard staged_lock(mStagedBucketsMutex); + auto& staged = mStagedBuckets[gemmId]; + MProfileMap new_buckets; + for (auto const& [m, config] : map) { + if (staged.insert(m).second) { + new_buckets.emplace(m, config); + } + } + if (!new_buckets.empty()) { + stagePersistentCache(gemmId, new_buckets, mHasWeightOnlyCudaKernel); + } + if (mStagedBuckets.size() == 1) { + mStagedBucketCount.store(staged.size(), std::memory_order_release); + } + } + + std::atomic mStagedBucketCount{0}; + std::mutex mStagedBucketsMutex; + std::unordered_map, GemmIdHashType> mStagedBuckets; + std::optional profileTacticsForProblem(int m, int n, int k, std::vector const& tactics, char* workspace, cudaStream_t stream); @@ -288,8 +407,10 @@ class GemmPluginProfiler { bool mSkip{false}; + int mMaxProfileM{kMaxProfileM}; + // Remembered from the initial profileTactics call so lazy single-bucket profiling can - // reproduce the same tactic candidate set. + // reproduce the same tactic candidate set and cache key. bool mHasWeightOnlyCudaKernel{false}; onnxruntime::AllocatorPtr mAllocator; @@ -318,8 +439,28 @@ class GemmPluginProfilerManager { return profiler; } + // Returns a profile map private to `session_key` (any address that is stable while the session lives). + // Profilers of the same session share one map, so identical shapes are profiled once per session, + // but the map is never shared with the process-global one or with other sessions. The map is dropped + // when the last profiler holding it goes away, so a later session reusing the key starts empty. + MNKProfileMapPtr getSessionProfileMap(const void* session_key) { + std::lock_guard lock(mSessionMapsMutex); + for (auto it = mSessionMaps.begin(); it != mSessionMaps.end();) { + it = it->second.expired() ? mSessionMaps.erase(it) : std::next(it); + } + auto& weak_map = mSessionMaps[session_key]; + auto map = weak_map.lock(); + if (map == nullptr) { + map = std::make_shared(); + weak_map = map; + } + return map; + } + private: MNKProfileMapPtr mMNKProfileMap{}; + std::mutex mSessionMapsMutex; + std::unordered_map> mSessionMaps; }; } // namespace onnxruntime::llm::kernels::weight_only @@ -333,7 +474,7 @@ GemmPluginProfiler::GemmPluginPro template int GemmPluginProfiler::getMaxProfileM() const { - return kMaxProfileM; + return mMaxProfileM; } template @@ -384,6 +525,9 @@ void GemmPluginProfiler::profileT onnxruntime::IAllocatorUniquePtr workspace_tmp{nullptr}; cudaStream_t stream; + // Populate from the persistent cache (if any) so already-tuned M buckets are skipped below. + loadPersistentCache(gemmId, *mProfileMap, hasWeightOnlyCudaKernel); + auto profileTactics = [&](int m, int n, int k) { if (mProfileMap->count(m) == 0) { if (!isAllocated) { @@ -418,11 +562,14 @@ void GemmPluginProfiler::profileT workspace_tmp.reset(); } CUDA_CALL_THROW(cudaStreamDestroy(stream)); + + // Stage only. EP teardown batches disk writes without holding the profile-map lock. + stageNewPersistentBuckets(gemmId, *mProfileMap); } template std::optional GemmPluginProfiler::getBestConfig( - int m, GemmIdType const& gemmId) const { + int m, GemmIdType const& gemmId) { ORT_LLM_LOG_ENTRY(); reader_lock lock(mMNKProfileMap->mutex); @@ -433,6 +580,10 @@ std::optional GemmPluginProfilergetMProfileMap(gemmId)); + if (mMNKProfileMap->getMProfileMap(gemmId)->count(m) > 0) { return mMNKProfileMap->getMProfileMap(gemmId)->at(m); } else if (mMNKProfileMap->getMProfileMap(gemmId)->count(mRounded) > 0) { @@ -481,6 +632,7 @@ std::optional GemmPluginProfilermutex); if (mMNKProfileMap->existsMProfileMap(gemmId)) { auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); + stageNewPersistentBuckets(gemmId, *mProfileMap); if (mProfileMap->count(m) > 0) { return mProfileMap->at(m); } @@ -491,7 +643,7 @@ std::optional GemmPluginProfiler GemmPluginProfilergetMProfileMap(gemmId); + stageNewPersistentBuckets(gemmId, *mProfileMap); if (mProfileMap->count(m) > 0) { return mProfileMap->at(m); } @@ -526,6 +679,8 @@ std::optional GemmPluginProfilerinsert({target, best}); + // In-memory staging only; the disk write happens at CUDA EP teardown, off the inference path. + stageNewPersistentBuckets(gemmId, *mProfileMap); return best; } diff --git a/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.cc b/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.cc new file mode 100644 index 0000000000000..34bea9ac0c8d5 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.cc @@ -0,0 +1,921 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if USE_FPA_INTB_GEMM + +#include "core/providers/shared_library/provider_api.h" +#ifdef BUILD_CUDA_EP_AS_PLUGIN +#include "core/common/common.h" +#include "core/providers/cuda/plugin/provider_api_shims.h" +#endif + +#include "contrib_ops/cuda/llm/gemm_tactic_cache.h" + +#include + +#include +#include +#include +#include +#include +#include +#include // NOLINT(build/c++17) +#include +#include +#include +#include +#include +#include +#include +#include + +#include "onnxruntime_config.h" + +#if defined(_WIN32) +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#include +#include +#endif + +namespace onnxruntime::llm::gemm_cache { + +namespace { + +using onnxruntime::llm::cutlass_extensions::ClusterShape; +using onnxruntime::llm::cutlass_extensions::CutlassTileConfig; +using onnxruntime::llm::cutlass_extensions::CutlassTileConfigSM100; +using onnxruntime::llm::cutlass_extensions::CutlassTileConfigSM120; +using onnxruntime::llm::cutlass_extensions::CutlassTileConfigSM90; +using onnxruntime::llm::cutlass_extensions::EpilogueScheduleType; +using onnxruntime::llm::cutlass_extensions::MainloopScheduleType; +using onnxruntime::llm::cutlass_extensions::SplitKStyle; + +std::vector SplitTabs(const std::string& line) { + std::vector out; + std::string field; + std::istringstream ss(line); + while (std::getline(ss, field, '\t')) { + out.push_back(field); + } + // std::getline drops a trailing empty field; recover it so column counts stay stable. + if (!line.empty() && line.back() == '\t') { + out.emplace_back(); + } + return out; +} + +std::string JoinTabs(const std::vector& fields) { + std::string out; + for (size_t i = 0; i < fields.size(); ++i) { + if (i != 0) { + out.push_back('\t'); + } + out += fields[i]; + } + return out; +} + +bool ParseInt(const std::string& s, int& value) { + if (s.empty()) { + return false; + } + errno = 0; + char* end = nullptr; + const long long parsed = std::strtoll(s.c_str(), &end, 10); + if (end == s.c_str() || *end != '\0' || errno != 0 || + parsed < std::numeric_limits::min() || parsed > std::numeric_limits::max()) { + return false; + } + value = static_cast(parsed); + return true; +} + +// Extracts the "git-commit-id=..." token from ORT_BUILD_INFO, if present. +std::string ExtractGitCommit(const std::string& build_info) { + const std::string key = "git-commit-id="; + auto pos = build_info.find(key); + if (pos == std::string::npos) { + return "unknown"; + } + pos += key.size(); + auto end = build_info.find_first_of(", ", pos); + if (end == std::string::npos) { + end = build_info.size(); + } + std::string commit = build_info.substr(pos, end - pos); + return commit.empty() ? "unknown" : commit; +} + +std::string CurrentOrtVersion() { +#ifdef ORT_VERSION + return std::string(ORT_VERSION); +#else + return "unknown"; +#endif +} + +std::string CurrentOrtBuildInfo() { +#ifdef ORT_BUILD_INFO + return std::string(ORT_BUILD_INFO); +#else + return std::string(); +#endif +} + +std::string CurrentBuildConfig() { +#ifdef NDEBUG + return "Release"; +#else + return "Debug"; +#endif +} + +// Cache paths are UTF-8; convert explicitly so non-ASCII paths work on Windows. +std::filesystem::path Utf8Path(const std::string& s) { + return std::filesystem::path(std::u8string(reinterpret_cast(s.data()), s.size())); +} + +int CurrentProcessId() { +#if defined(_WIN32) + return static_cast(::GetCurrentProcessId()); +#else + return static_cast(::getpid()); +#endif +} + +// RAII cross-process advisory lock on ".lock". +class ScopedFileLock { + public: + explicit ScopedFileLock(const std::string& path, bool try_only = false) : lock_path_(Utf8Path(path + ".lock")) { +#if defined(_WIN32) + handle_ = ::CreateFileW(lock_path_.c_str(), GENERIC_READ | GENERIC_WRITE, + FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr, + OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr); + if (handle_ != INVALID_HANDLE_VALUE) { + OVERLAPPED ov{}; + const DWORD flags = LOCKFILE_EXCLUSIVE_LOCK | (try_only ? LOCKFILE_FAIL_IMMEDIATELY : 0); + locked_ = ::LockFileEx(handle_, flags, 0, MAXDWORD, MAXDWORD, &ov) != 0; + } +#else + fd_ = ::open(lock_path_.c_str(), O_CREAT | O_RDWR, 0644); + if (fd_ != -1) { + const int operation = LOCK_EX | (try_only ? LOCK_NB : 0); + locked_ = ::flock(fd_, operation) == 0; + } +#endif + } + + ~ScopedFileLock() { +#if defined(_WIN32) + if (handle_ != INVALID_HANDLE_VALUE) { + if (locked_) { + OVERLAPPED ov{}; + ::UnlockFileEx(handle_, 0, MAXDWORD, MAXDWORD, &ov); + } + ::CloseHandle(handle_); + } +#else + if (fd_ != -1) { + if (locked_) { + ::flock(fd_, LOCK_UN); + } + ::close(fd_); + } +#endif + } + + bool locked() const { return locked_; } + + ScopedFileLock(const ScopedFileLock&) = delete; + ScopedFileLock& operator=(const ScopedFileLock&) = delete; + + private: + std::filesystem::path lock_path_; + bool locked_ = false; +#if defined(_WIN32) + HANDLE handle_ = INVALID_HANDLE_VALUE; +#else + int fd_ = -1; +#endif +}; + +} // namespace + +std::string TsvEncode(const std::string& s) { + std::string out; + out.reserve(s.size()); + for (char c : s) { + switch (c) { + case '%': + out += "%25"; + break; + case '\t': + out += "%09"; + break; + case '\n': + out += "%0A"; + break; + case '\r': + out += "%0D"; + break; + default: + out.push_back(c); + break; + } + } + return out; +} + +std::string TsvDecode(const std::string& s) { + std::string out; + out.reserve(s.size()); + for (size_t i = 0; i < s.size(); ++i) { + if (s[i] == '%' && i + 2 < s.size() && + std::isxdigit(static_cast(s[i + 1])) && + std::isxdigit(static_cast(s[i + 2]))) { + out.push_back(static_cast(std::stoi(s.substr(i + 1, 2), nullptr, 16))); + i += 2; + continue; + } + out.push_back(s[i]); + } + return out; +} + +HardwareSignature HardwareSignature::Compute() { + int device = 0; + if (cudaGetDevice(&device) == cudaSuccess) { + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, device) == cudaSuccess) { + return FromDevice(prop.name, prop.major * 10 + prop.minor, prop.multiProcessorCount); + } + } + return FromDevice(std::string(), 0, 0); +} + +HardwareSignature HardwareSignature::FromDevice(std::string device_name, int sm, int multiprocessor_count) { + HardwareSignature sig; + sig.device_name = std::move(device_name); + sig.sm = sm; + sig.multiprocessor_count = multiprocessor_count; + sig.cuda_runtime = CUDART_VERSION; + int driver = 0; + if (cudaDriverGetVersion(&driver) == cudaSuccess) { + sig.cuda_driver = driver; + } + + sig.ort_version = CurrentOrtVersion(); + sig.ort_git_commit = ExtractGitCommit(CurrentOrtBuildInfo()); + sig.ort_build_config = CurrentBuildConfig(); + return sig; +} + +bool HardwareSignature::StrictMatches(const HardwareSignature& other) const { + // ort_git_commit and ort_build_config are intentionally excluded (diagnostic only): reusing a + // tactic across commits or build configs can at worst be slightly suboptimal, never incorrect, + // and loadPersistentCache re-validates each CUTLASS tactic against the current runner anyway. + return device_name == other.device_name && + sm == other.sm && + cuda_runtime == other.cuda_runtime && + ort_version == other.ort_version; +} + +std::string HardwareSignature::FilePrefixToken() const { + std::string token = device_name.empty() ? "unknown_gpu" : device_name; + for (char& c : token) { + if (!std::isalnum(static_cast(c))) { + c = '_'; + } + } + token += "_sm" + std::to_string(sm); + return token; +} + +void AppendConfigColumns(std::vector& row, const std::optional& config) { + if (!config.has_value()) { + // valid_config=0; remaining columns are placeholders (ignored on parse). + // Order: valid_config sm_version tile80 tile90 tile100 tile120 + // split_k_style split_k stages cluster mainloop epilogue tma enable_cuda_kernel cuda_kernel_variant + const char* placeholders[kNumConfigColumns] = { + "0", "0", "0", "0", "0", "0", "0", "-1", "-1", "0", "0", "0", "0", "0", "0"}; + for (const char* p : placeholders) { + row.emplace_back(p); + } + return; + } + + const CutlassGemmConfig& c = *config; + row.emplace_back("1"); // valid_config + row.emplace_back(std::to_string(c.sm_version)); // sm_version + row.emplace_back(std::to_string(static_cast(c.tile_config_sm80))); // tile80 + row.emplace_back(std::to_string(static_cast(c.tile_config_sm90))); // tile90 + row.emplace_back(std::to_string(static_cast(c.tile_config_sm100))); // tile100 + row.emplace_back(std::to_string(static_cast(c.tile_config_sm120))); // tile120 + row.emplace_back(std::to_string(static_cast(c.split_k_style))); // split_k_style + row.emplace_back(std::to_string(c.split_k_factor)); // split_k + row.emplace_back(std::to_string(c.stages)); // stages + row.emplace_back(std::to_string(static_cast(c.cluster_shape))); // cluster + row.emplace_back(std::to_string(static_cast(c.mainloop_schedule))); // mainloop + row.emplace_back(std::to_string(static_cast(c.epilogue_schedule))); // epilogue + row.emplace_back(c.is_tma_warp_specialized ? "1" : "0"); // tma + row.emplace_back(c.enableCudaKernel ? "1" : "0"); // enable_cuda_kernel + row.emplace_back(std::to_string(c.cudaKernelVariant)); // cuda_kernel_variant +} + +std::optional> ParseConfigColumns( + const std::vector& columns, size_t begin) { + if (begin + kNumConfigColumns > columns.size()) { + return std::nullopt; + } + + int vals[kNumConfigColumns]; + for (int i = 0; i < kNumConfigColumns; ++i) { + if (!ParseInt(columns[begin + i], vals[i])) { + return std::nullopt; + } + } + + const int valid_config = vals[0]; + if (valid_config == 0) { + // A profiled bucket with no valid tactic. + return std::optional{std::nullopt}; + } + + CutlassGemmConfig c; + c.sm_version = vals[1]; + c.tile_config_sm80 = static_cast(vals[2]); + c.tile_config_sm90 = static_cast(vals[3]); + c.tile_config_sm100 = static_cast(vals[4]); + c.tile_config_sm120 = static_cast(vals[5]); + c.split_k_style = static_cast(vals[6]); + c.split_k_factor = vals[7]; + c.stages = vals[8]; + c.cluster_shape = static_cast(vals[9]); + c.mainloop_schedule = static_cast(vals[10]); + c.epilogue_schedule = static_cast(vals[11]); + c.is_tma_warp_specialized = vals[12] != 0; + c.enableCudaKernel = vals[13] != 0; + c.cudaKernelVariant = vals[14]; + if (c.cudaKernelVariant < 0 || c.cudaKernelVariant > 7 || (!c.enableCudaKernel && c.cudaKernelVariant != 0)) { + return std::nullopt; + } + return std::optional{c}; +} + +namespace { +std::string ToLower(std::string s) { + std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return static_cast(std::tolower(c)); }); + return s; +} +} // namespace + +std::optional ParseCacheAccess(const std::string& mode) { + const std::string m = ToLower(mode); + if (m.empty() || m == "load_save") { + return CacheAccess{true, true}; + } + if (m == "load") { + return CacheAccess{true, false}; + } + if (m == "save") { + return CacheAccess{false, true}; + } + return std::nullopt; +} + +std::optional ParseCacheFlag(const std::string& value) { + const std::string v = ToLower(value); + if (v.empty() || v == "0" || v == "false" || v == "off") { + return false; + } + if (v == "1" || v == "true" || v == "on") { + return true; + } + return std::nullopt; +} + +bool MatMulNBitsKey::operator==(const MatMulNBitsKey& o) const { + return n_16b == o.n_16b && k == o.k && + activation_dtype == o.activation_dtype && weight_type == o.weight_type && + bits == o.bits && block_size == o.block_size && + has_zero_points == o.has_zero_points && zero_point_dtype == o.zero_point_dtype && + gemv_enabled == o.gemv_enabled && has_bias == o.has_bias && packing_sm == o.packing_sm && + paired_gemv_mode == o.paired_gemv_mode && wave_aware_gemv == o.wave_aware_gemv; +} + +std::size_t MatMulNBitsKeyHash::operator()(const MatMulNBitsKey& k) const { + std::size_t h = std::hash{}(k.n_16b); + auto mix = [&h](std::size_t v) { + h ^= v + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2); + }; + mix(std::hash{}(k.k)); + mix(std::hash{}(k.activation_dtype)); + mix(std::hash{}(k.weight_type)); + mix(std::hash{}(k.bits)); + mix(std::hash{}(k.block_size)); + mix(std::hash{}(k.has_zero_points)); + mix(std::hash{}(k.zero_point_dtype)); + mix(std::hash{}(k.gemv_enabled)); + mix(std::hash{}(k.has_bias)); + mix(std::hash{}(k.packing_sm)); + mix(std::hash{}(k.paired_gemv_mode)); + mix(std::hash{}(k.wave_aware_gemv)); + return h; +} + +namespace { + +// Column names for the MatMulNBits data rows, in write order. +const std::vector& MatMulNBitsColumnNames() { + static const std::vector names = { + "n_16b", "k", "activation_dtype", "weight_type", "bits", + "block_size", "has_zero_points", "zero_point_dtype", "gemv_enabled", "has_bias", "packing_sm", + "paired_gemv_mode", "wave_aware_gemv", + "m_bucket", "valid_config", "sm_version", "tile80", "tile90", "tile100", "tile120", + "split_k_style", "split_k", "stages", "cluster", "mainloop", "epilogue", "tma", + "enable_cuda_kernel", "cuda_kernel_variant"}; + return names; +} + +} // namespace + +MatMulNBitsTacticCache::MatMulNBitsTacticCache(std::string file_path, HardwareSignature signature) + : file_path_(std::move(file_path)), signature_(std::move(signature)) {} + +std::string MatMulNBitsTacticCache::ResolveFilePath(const std::string& config_dir, const std::string& config_prefix, + const HardwareSignature& signature, + const std::string& sidecar_prefix) { + std::string prefix = config_prefix; + std::string dir = config_dir; + // Any session-config location overrides the env vars, so an env prefix cannot shadow a session dir. + if (prefix.empty() && dir.empty() && !sidecar_prefix.empty()) { + std::error_code ec; + const std::filesystem::path sidecar_path = Utf8Path(sidecar_prefix); + const std::filesystem::path abs_sidecar = + sidecar_path.is_relative() ? std::filesystem::absolute(sidecar_path, ec) : sidecar_path; + if (!ec && !abs_sidecar.empty()) { + const auto abs_u8 = abs_sidecar.u8string(); + prefix.assign(abs_u8.begin(), abs_u8.end()); + } else { + prefix = sidecar_prefix; + } + } + if (prefix.empty() && dir.empty()) { + prefix = GetEnvironmentVar(kEnvCachePrefix); + dir = GetEnvironmentVar(kEnvCacheDir); + } + + const std::string suffix = std::string(".") + kTableMatMulNBits + ".tsv"; + if (!prefix.empty()) { + return prefix + suffix; + } + if (!dir.empty()) { + return dir + "/" + signature.FilePrefixToken() + suffix; + } + return std::string(); +} + +std::shared_ptr MatMulNBitsTacticCache::MaybeCreate( + const std::string& config_dir, const std::string& config_prefix) { + HardwareSignature signature = HardwareSignature::Compute(); + std::string file_path = ResolveFilePath(config_dir, config_prefix, signature); + if (file_path.empty()) { + return nullptr; + } + + auto cache = std::make_shared(std::move(file_path), std::move(signature)); + auto status = cache->Load(); + ORT_UNUSED_PARAMETER(status); // A missing/mismatched file simply yields an empty cache. + return cache; +} + +std::optional> MatMulNBitsTacticCache::Get( + const MatMulNBitsKey& key, int m_bucket) const { + std::lock_guard guard(mutex_); + auto it = table_.find(key); + if (it == table_.end()) { + return std::nullopt; + } + auto bucket_it = it->second.find(m_bucket); + if (bucket_it == it->second.end()) { + return std::nullopt; + } + return bucket_it->second; +} + +MatMulNBitsTacticCache::BucketMap MatMulNBitsTacticCache::GetAll(const MatMulNBitsKey& key) const { + std::lock_guard guard(mutex_); + auto it = table_.find(key); + if (it == table_.end()) { + return {}; + } + return it->second; +} + +bool MatMulNBitsTacticCache::IsDirtyLocked(const MatMulNBitsKey& key, int m_bucket) const { + auto key_it = dirty_rows_.find(key); + return key_it != dirty_rows_.end() && key_it->second.find(m_bucket) != key_it->second.end(); +} + +void MatMulNBitsTacticCache::Put(const MatMulNBitsKey& key, int m_bucket, + const std::optional& config) { + if (!config.has_value()) { + return; // A failed profiling attempt can be transient and must not survive a process restart. + } + std::lock_guard guard(mutex_); + table_[key][m_bucket] = config; + dirty_rows_[key][m_bucket] = DirtyRow{config, ++next_dirty_generation_}; +} + +onnxruntime::common::Status MatMulNBitsTacticCache::Load() { + // The stamp is taken before reading so a concurrent writer can only cause a harmless extra reload. + const auto stamp = CurrentFileStamp(); + std::ifstream in(Utf8Path(file_path_)); + if (!in.is_open()) { + std::error_code ec; + const bool exists = std::filesystem::exists(Utf8Path(file_path_), ec); + if (!exists && !ec) { + std::lock_guard guard(mutex_); + known_stamp_ = stamp; + return onnxruntime::common::Status::OK(); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to read gemm tactic cache: ", file_path_); + } + auto status = Load(in); + if (status.IsOK()) { + std::lock_guard guard(mutex_); + known_stamp_ = stamp; + } + return status; +} + +onnxruntime::common::Status MatMulNBitsTacticCache::Load(std::istream& in) { + HardwareSignature file_sig; + file_sig.ort_build_config.clear(); + std::string magic_ok_version; + std::string selection_version; + std::string table_name; + std::vector header_columns; + std::unordered_map loaded; + + std::string line; + while (std::getline(in, line)) { + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + if (line.empty()) { + continue; + } + + if (line[0] == '#') { + auto fields = SplitTabs(line.substr(1)); + // Strip a leading space after '#'. + if (!fields.empty()) { + auto& first = fields[0]; + size_t start = first.find_first_not_of(' '); + first = (start == std::string::npos) ? "" : first.substr(start); + } + if (fields.size() >= 2) { + const std::string& k = fields[0]; + const std::string& v = fields[1]; + if (k == kCacheMagic) { + magic_ok_version = v; + } else if (k == "table") { + table_name = v; + } else if (k == "tactic_selection_version") { + selection_version = v; + } else if (k == "device_name") { + file_sig.device_name = TsvDecode(v); + } else if (k == "sm") { + ParseInt(v, file_sig.sm); + } else if (k == "multiprocessor_count") { + ParseInt(v, file_sig.multiprocessor_count); + } else if (k == "cuda_runtime") { + ParseInt(v, file_sig.cuda_runtime); + } else if (k == "cuda_driver") { + ParseInt(v, file_sig.cuda_driver); + } else if (k == "ort_version") { + file_sig.ort_version = v; + } else if (k == "ort_git_commit") { + file_sig.ort_git_commit = v; + } else if (k == "ort_build_config") { + file_sig.ort_build_config = v; + } + } + continue; + } + + if (header_columns.empty()) { + header_columns = SplitTabs(line); + continue; + } + + // Data row: map by column name so appended columns are tolerated. + auto values = SplitTabs(line); + auto field = [&](const char* name) -> std::string { + for (size_t i = 0; i < header_columns.size() && i < values.size(); ++i) { + if (header_columns[i] == name) { + return values[i]; + } + } + return std::string(); + }; + + MatMulNBitsKey key; + if (!ParseInt(field("n_16b"), key.n_16b) || !ParseInt(field("k"), key.k) || + !ParseInt(field("bits"), key.bits) || !ParseInt(field("block_size"), key.block_size) || + !ParseInt(field("packing_sm"), key.packing_sm)) { + continue; + } + key.activation_dtype = TsvDecode(field("activation_dtype")); + key.weight_type = TsvDecode(field("weight_type")); + int has_zp = 0; + ParseInt(field("has_zero_points"), has_zp); + key.has_zero_points = has_zp != 0; + key.zero_point_dtype = TsvDecode(field("zero_point_dtype")); + int gemv = 0; + ParseInt(field("gemv_enabled"), gemv); + key.gemv_enabled = gemv != 0; + int has_bias = 0; + ParseInt(field("has_bias"), has_bias); + key.has_bias = has_bias != 0; + int wave_aware = 0; + if (!ParseInt(field("paired_gemv_mode"), key.paired_gemv_mode) || + key.paired_gemv_mode < 0 || key.paired_gemv_mode > 2 || + !ParseInt(field("wave_aware_gemv"), wave_aware) || (wave_aware != 0 && wave_aware != 1)) { + continue; + } + key.wave_aware_gemv = wave_aware != 0; + + int m_bucket = 0; + if (!ParseInt(field("m_bucket"), m_bucket)) { + continue; + } + + // Locate the start of the config columns by name so appended key columns do not shift them. + size_t valid_config_index = header_columns.size(); + for (size_t i = 0; i < header_columns.size(); ++i) { + if (header_columns[i] == "valid_config") { + valid_config_index = i; + break; + } + } + auto parsed = ParseConfigColumns(values, valid_config_index); + if (!parsed.has_value() || !parsed->has_value() || m_bucket <= 0) { + continue; + } + loaded[key][m_bucket] = *parsed; + } + + if (in.bad() || !in.eof()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Incomplete read of gemm tactic cache: ", file_path_); + } + + // Reject the file if the format, selection policy, table, or hardware/build signature does not match. + if (magic_ok_version != kCacheFormatVersion || selection_version != kTacticSelectionVersion || + table_name != kTableMatMulNBits || !signature_.StrictMatches(file_sig)) { + return onnxruntime::common::Status::OK(); + } + + std::lock_guard guard(mutex_); + for (auto& [key, buckets] : loaded) { + auto& dest = table_[key]; + for (auto& [m, cfg] : buckets) { + if (!IsDirtyLocked(key, m)) { + dest[m] = cfg; // Refresh clean entries from disk while preserving locally pending updates. + } + } + } + return onnxruntime::common::Status::OK(); +} + +onnxruntime::common::Status MatMulNBitsTacticCache::ReloadIfChanged() { + const auto stamp = CurrentFileStamp(); + { + std::lock_guard guard(mutex_); + if (stamp == known_stamp_) { + return onnxruntime::common::Status::OK(); + } + } + return Load(); +} + +std::optional MatMulNBitsTacticCache::CurrentFileStamp() const { + const std::filesystem::path path = Utf8Path(file_path_); + std::error_code ec; + const auto time = std::filesystem::last_write_time(path, ec); + if (ec) { + return std::nullopt; + } + const auto size = std::filesystem::file_size(path, ec); + if (ec) { + return std::nullopt; + } + return FileStamp{time, size}; +} + +onnxruntime::common::Status MatMulNBitsTacticCache::WriteAllLocked( + const std::unordered_map& table) const { + // Unique per writer so concurrent writers never share a temp file. + static std::atomic tmp_counter{0}; + const std::string tmp_path = file_path_ + ".tmp." + std::to_string(CurrentProcessId()) + "." + + std::to_string(tmp_counter.fetch_add(1)); + const std::filesystem::path tmp_fs_path = Utf8Path(tmp_path); + const std::filesystem::path fs_path = Utf8Path(file_path_); + struct TempFileCleanup { + const std::filesystem::path& path; + ~TempFileCleanup() { + std::error_code ec; + std::filesystem::remove(path, ec); + } + } cleanup{tmp_fs_path}; + { + std::ofstream out(tmp_fs_path, std::ios::trunc); + if (!out.is_open()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Failed to open gemm tactic cache temp file for writing: ", tmp_path); + } + + // Signature / format header lines. + out << "# " << kCacheMagic << '\t' << kCacheFormatVersion << '\n'; + out << "# table\t" << kTableMatMulNBits << '\n'; + out << "# tactic_selection_version\t" << kTacticSelectionVersion << '\n'; + out << "# device_name\t" << TsvEncode(signature_.device_name) << '\n'; + out << "# sm\t" << signature_.sm << '\n'; + out << "# multiprocessor_count\t" << signature_.multiprocessor_count << '\n'; + out << "# cuda_runtime\t" << signature_.cuda_runtime << '\n'; + out << "# cuda_driver\t" << signature_.cuda_driver << '\n'; + out << "# ort_version\t" << signature_.ort_version << '\n'; + out << "# ort_git_commit\t" << signature_.ort_git_commit << '\n'; + out << "# ort_build_config\t" << signature_.ort_build_config << '\n'; + + // Column header. + out << JoinTabs(MatMulNBitsColumnNames()) << '\n'; + + // Data rows (sorted by bucket within each key for stable, diffable output). + for (const auto& [key, buckets] : table) { + for (const auto& [m_bucket, cfg] : buckets) { + std::vector row; + row.reserve(MatMulNBitsColumnNames().size()); + row.emplace_back(std::to_string(key.n_16b)); + row.emplace_back(std::to_string(key.k)); + row.emplace_back(TsvEncode(key.activation_dtype)); + row.emplace_back(TsvEncode(key.weight_type)); + row.emplace_back(std::to_string(key.bits)); + row.emplace_back(std::to_string(key.block_size)); + row.emplace_back(key.has_zero_points ? "1" : "0"); + row.emplace_back(TsvEncode(key.zero_point_dtype)); + row.emplace_back(key.gemv_enabled ? "1" : "0"); + row.emplace_back(key.has_bias ? "1" : "0"); + row.emplace_back(std::to_string(key.packing_sm)); + row.emplace_back(std::to_string(key.paired_gemv_mode)); + row.emplace_back(key.wave_aware_gemv ? "1" : "0"); + row.emplace_back(std::to_string(m_bucket)); + AppendConfigColumns(row, cfg); + out << JoinTabs(row) << '\n'; + } + } + + out.close(); + if (!out.good()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Failed while writing gemm tactic cache temp file: ", tmp_path); + } + } + +#if defined(_WIN32) + if (::MoveFileExW(tmp_fs_path.c_str(), fs_path.c_str(), MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH) == 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Failed to atomically replace gemm tactic cache file: ", file_path_, + ", Windows error: ", ::GetLastError()); + } +#else + if (std::rename(tmp_fs_path.c_str(), fs_path.c_str()) != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Failed to atomically replace gemm tactic cache file: ", file_path_); + } +#endif + return onnxruntime::common::Status::OK(); +} + +onnxruntime::common::Status MatMulNBitsTacticCache::Flush() { + return FlushImpl(/*allow_blocking*/ true); +} + +onnxruntime::common::Status MatMulNBitsTacticCache::TryFlush() { + return FlushImpl(/*allow_blocking*/ false); +} + +onnxruntime::common::Status MatMulNBitsTacticCache::FlushImpl(bool allow_blocking) { + { + std::unique_lock guard(mutex_, std::defer_lock); + if (allow_blocking) { + guard.lock(); + } else if (!guard.try_lock()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Gemm tactic cache mutex is busy: ", file_path_); + } + if (dirty_rows_.empty()) { + return onnxruntime::common::Status::OK(); + } + } + + const std::filesystem::path parent = Utf8Path(file_path_).parent_path(); + if (!parent.empty()) { + std::error_code ec; + std::filesystem::create_directories(parent, ec); + if (ec) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to create gemm tactic cache directory for ", file_path_, + ": ", ec.message()); + } + } + + // Without the lock two writers could load the same snapshot and the last rename would drop the + // other's rows, so fail the flush (the rows stay dirty and are retried at the next flush point). + ScopedFileLock lock(file_path_, /*try_only*/ !allow_blocking); + if (!lock.locked()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to acquire the gemm tactic cache file lock: ", file_path_, ".lock"); + } + + // Reload the current on-disk file so concurrently-written rows are not dropped, + // then overlay only locally changed rows. + MatMulNBitsTacticCache disk(file_path_, signature_); + ORT_RETURN_IF_ERROR(disk.Load()); + + std::unordered_map merged; + { + std::lock_guard disk_guard(disk.mutex_); + merged = disk.table_; + } + decltype(dirty_rows_) pending; + { + std::unique_lock guard(mutex_, std::defer_lock); + if (allow_blocking) { + guard.lock(); + } else if (!guard.try_lock()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Gemm tactic cache mutex is busy: ", file_path_); + } + if (dirty_rows_.empty()) { + return onnxruntime::common::Status::OK(); + } + // Keep pending rows in dirty_rows_ while WriteAllLocked runs so a concurrent Load() or + // ReloadIfChanged() treats the in-flight flush as locally pending and does not overwrite it + // with the pre-rename file contents. + pending = dirty_rows_; + for (const auto& [key, buckets] : pending) { + auto& dest = merged[key]; + for (const auto& [m, dirty] : buckets) { + dest[m] = dirty.config; + } + } + } + + onnxruntime::common::Status status = WriteAllLocked(merged); + if (!status.IsOK()) { + return status; + } + + const auto stamp = CurrentFileStamp(); + std::unique_lock guard(mutex_, std::defer_lock); + if (allow_blocking) { + guard.lock(); + } else if (!guard.try_lock()) { + return status; + } + for (const auto& [key, buckets] : pending) { + auto dirty_key_it = dirty_rows_.find(key); + if (dirty_key_it == dirty_rows_.end()) { + continue; + } + for (const auto& [m, flushed] : buckets) { + auto dirty_it = dirty_key_it->second.find(m); + if (dirty_it != dirty_key_it->second.end() && dirty_it->second.generation == flushed.generation) { + dirty_key_it->second.erase(dirty_it); + } + } + if (dirty_key_it->second.empty()) { + dirty_rows_.erase(dirty_key_it); + } + } + known_stamp_ = stamp; + // Adopt rows from the merged snapshot (including updates to existing clean buckets from other + // writers), since known_stamp_ now hides them from ReloadIfChanged(). Preserve any row dirtied + // while the write was in flight. + for (const auto& [key, buckets] : merged) { + auto& dest = table_[key]; + for (const auto& [m, cfg] : buckets) { + if (!IsDirtyLocked(key, m)) { + dest[m] = cfg; + } + } + } + return status; +} + +} // namespace onnxruntime::llm::gemm_cache + +#endif // USE_FPA_INTB_GEMM diff --git a/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.h b/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.h new file mode 100644 index 0000000000000..82a79d19e11b5 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.h @@ -0,0 +1,253 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// +// Persistent, hardware-keyed tactic cache for weight-only GEMM autotuning. +// +// This is the shared foundation described in +// docs/contrib_ops/cuda/gemm_profiler_cache.md (Phase 1). It provides: +// * HardwareSignature: GPU model / SM / toolkit / build guard used to name and +// validate cache files so tactics are never reused across incompatible setups. +// * (De)serialization of CutlassGemmConfig to/from a fixed set of integer columns. +// * TSV read/write helpers with percent-encoding for free-text fields. +// * MatMulNBitsTacticCache: a typed, disk-backed cache keyed by the fpA_intB +// packed problem shape and quantization mode, with file locking, atomic write, +// and in-memory merge. +// +// Persistence is opt-in. When no cache directory or explicit prefix is configured, +// the cache stays in memory only and never touches disk. +// +// The declarations below are always available in CUDA builds; the implementation in +// gemm_tactic_cache.cc is compiled only when USE_FPA_INTB_GEMM is defined and is only +// referenced from the fpA_intB profiling paths that are themselves so guarded. +// +#pragma once + +#include +#include // NOLINT(build/c++17) +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/common/status.h" +#include "contrib_ops/cuda/llm/cutlass_extensions/gemm_configs.h" + +namespace onnxruntime::llm::gemm_cache { + +using onnxruntime::llm::cutlass_extensions::CutlassGemmConfig; + +// Cache file format version. Bump when the on-disk schema changes incompatibly. +constexpr const char* kCacheMagic = "ort_cuda_gemm_tactic_cache"; +constexpr const char* kCacheFormatVersion = "v1"; + +// Bump whenever profiling or tactic selection changes which tactic wins, so stale caches are rejected. +constexpr const char* kTacticSelectionVersion = "4"; + +constexpr const char* kTableMatMulNBits = "matmulnbits_fpa_intb"; + +// Environment variables (see docs section 9). +constexpr const char* kEnvCacheDir = "ORT_CUDA_GEMM_TACTIC_CACHE_DIR"; +constexpr const char* kEnvCachePrefix = "ORT_CUDA_GEMM_TACTIC_CACHE_PREFIX"; +constexpr const char* kEnvCacheMode = "ORT_CUDA_GEMM_TACTIC_CACHE_MODE"; +constexpr const char* kEnvCacheToModel = "ORT_CUDA_GEMM_TACTIC_CACHE_TO_MODEL"; + +// Session-option config keys (equivalents of the env vars). These can be set via +// SessionOptions::AddConfigEntry, or through onnxruntime-genai's session_options in +// genai_config.json (any unrecognized session_options key is forwarded to AddConfigEntry). +constexpr const char* kSessionConfigCacheDir = "ep.cuda.gemm_tactic_cache_dir"; +constexpr const char* kSessionConfigCachePrefix = "ep.cuda.gemm_tactic_cache_prefix"; +constexpr const char* kSessionConfigCacheMode = "ep.cuda.gemm_tactic_cache_mode"; +constexpr const char* kSessionConfigCacheToModel = "ep.cuda.gemm_tactic_cache_to_model"; + +// Persistence directions used by a session. Parsed from the cache mode: +// "load_save" (default): load at session creation, save newly tuned tactics at session close. +// "load": read-only; never writes (e.g. a pre-tuned cache shipped with an application). +// "save": ignore existing rows, re-tune, and write the results at session close. +struct CacheAccess { + bool load = true; + bool save = true; +}; + +// Returns std::nullopt for an unrecognized mode. An empty mode means "load_save". +std::optional ParseCacheAccess(const std::string& mode); + +// Returns true for "1"/"true"/"on" (case-insensitive), false for empty/"0"/"false"/"off", +// and std::nullopt otherwise. +std::optional ParseCacheFlag(const std::string& value); + +// Optional separate output containing only tactics selected by the tuning session. +constexpr const char* kSessionConfigTuningResultsPrefix = "ep.cuda.gemm_tactic_cache_tuning_results_prefix"; + +// Hardware / build signature used both to name cache files and as a stored guard. +// Reuse is rejected on mismatch of any strict field (see StrictMatches). +struct HardwareSignature { + std::string device_name; // cudaDeviceProp.name, e.g. "NVIDIA A100-SXM4-80GB" + int sm = 0; // major*10 + minor, e.g. 80, 89, 90 + int multiprocessor_count = 0; // diagnostic only + int cuda_runtime = 0; // CUDART_VERSION + int cuda_driver = 0; // cudaDriverGetVersion(), diagnostic only + std::string ort_version; // ORT_VERSION (strict-match field) + std::string ort_git_commit; // parsed from ORT_BUILD_INFO or "unknown" (diagnostic only) + std::string ort_build_config; // "Release" / "Debug" (diagnostic only) + + // Computes the signature for the current CUDA device and build. + static HardwareSignature Compute(); + + // Computes the signature for an already-queried device (avoids a cudaGetDeviceProperties call). + static HardwareSignature FromDevice(std::string device_name, int sm, int multiprocessor_count); + + // Strict reuse guard: device_name, sm, cuda_runtime, and ort_version must all match. + // multiprocessor_count, cuda_driver, ort_git_commit, and ort_build_config are recorded for + // diagnostics only and are NOT part of the guard: a tactic is just a config selected among + // getConfigs(), so a cross-commit/cross-build reuse can at worst pick a slightly suboptimal + // (never incorrect) tactic, and loadPersistentCache re-validates every CUTLASS tactic against + // the current runner and drops any that no longer applies. + bool StrictMatches(const HardwareSignature& other) const; + + // A filesystem-safe token derived from device_name + sm used as a default cache prefix. + std::string FilePrefixToken() const; +}; + +// TSV free-text encoding: percent-encodes '%', '\t', '\n', '\r' so a field never +// breaks the tab/newline structure. Decode reverses it. +std::string TsvEncode(const std::string& s); +std::string TsvDecode(const std::string& s); + +// Fixed integer column layout for a (possibly-absent) tactic, mirroring the doc: +// valid_config sm_version tile80 tile90 tile100 tile120 +// split_k_style split_k stages cluster mainloop epilogue tma enable_cuda_kernel cuda_kernel_variant +constexpr int kNumConfigColumns = 15; + +// Appends the config columns (as decimal strings) for `config` to `row`. +// A missing tactic (std::nullopt) is written with valid_config=0 and defaults elsewhere. +void AppendConfigColumns(std::vector& row, const std::optional& config); + +// Parses the config columns starting at `columns[begin]`. Returns std::nullopt +// (outer) if the columns are malformed; the inner optional is std::nullopt when +// valid_config==0 (a profiled bucket with no valid tactic). +std::optional> ParseConfigColumns( + const std::vector& columns, size_t begin); + +// Problem key for MatMulNBits fpA_intB. Many nodes with the same shape collapse to +// one entry. n_16b is the value used by GemmIdCore (N after casting int weight to fp16). +struct MatMulNBitsKey { + int n_16b = 0; + int k = 0; + std::string activation_dtype; // "half" / "bfloat16" + std::string weight_type; // "uint2b_t" / "uint4b_t" / "uint8_t" + int bits = 0; + int block_size = 0; + bool has_zero_points = false; + std::string zero_point_dtype; // e.g. "uint4b_t" / activation dtype / "none" + bool gemv_enabled = false; + bool has_bias = false; + int packing_sm = 0; + int paired_gemv_mode = 0; + bool wave_aware_gemv = false; + + bool operator==(const MatMulNBitsKey& o) const; +}; + +struct MatMulNBitsKeyHash { + std::size_t operator()(const MatMulNBitsKey& k) const; +}; + +// Disk-backed tactic cache for MatMulNBits fpA_intB. Thread-safe. +// +// Lookup/merge semantics: +// * Get(key, m) -> outer nullopt means "not cached, must profile". +// Only successful tactics are returned; legacy negative rows are ignored. +// * Put(key, m, config) records a successful tactic in memory; nullopt is ignored. +// * Flush() atomically merges locally changed rows into the on-disk file under a +// file lock so concurrent sessions tuning different shapes do not lose updates. +class MatMulNBitsTacticCache { + public: + // Resolves the cache file path, or returns an empty string when persistence is disabled. + // Resolution order: session-config prefix, session-config dir, `sidecar_prefix` (the model path + // when the model sidecar is enabled), env prefix, env dir. + static std::string ResolveFilePath(const std::string& config_dir, const std::string& config_prefix, + const HardwareSignature& signature, const std::string& sidecar_prefix = ""); + + // Returns a configured cache if persistence is enabled, otherwise nullptr (callers keep the + // in-process-only behavior). The file location follows ResolveFilePath. The returned cache has + // already attempted to load matching rows from disk. + static std::shared_ptr MaybeCreate(const std::string& config_dir = "", + const std::string& config_prefix = ""); + + // Test/tooling entry point: build a cache bound to an explicit file path. + MatMulNBitsTacticCache(std::string file_path, HardwareSignature signature); + + // Ordered map from M-bucket to its (possibly-absent) tactic for one problem key. + using BucketMap = std::map>; + + const std::string& FilePath() const { return file_path_; } + const HardwareSignature& Signature() const { return signature_; } + + std::optional> Get(const MatMulNBitsKey& key, int m_bucket) const; + + // Returns all cached M-buckets for `key` (empty if the key is not cached). + BucketMap GetAll(const MatMulNBitsKey& key) const; + + void Put(const MatMulNBitsKey& key, int m_bucket, const std::optional& config); + + // Loads rows from file_path_ whose header signature strictly matches signature_. + // A missing file or mismatched signature leaves the in-memory table untouched and + // returns OK (treated as an empty/rejected cache). + onnxruntime::common::Status Load(); + + // Parses a complete snapshot. No rows are accepted if the stream reports a read error. + onnxruntime::common::Status Load(std::istream& in); + + // Loads again if the file changed on disk since the last Load() or Flush(), refreshing clean + // rows from disk while preserving locally pending updates. + onnxruntime::common::Status ReloadIfChanged(); + + // Atomically merges locally changed rows with the current on-disk file and writes + // it back. No-op (OK) if there is nothing to persist. + onnxruntime::common::Status Flush(); + + // Non-blocking variant of Flush() for shutdown/unload paths where waiting on a mutex or file + // lock held by a stopped worker thread could hang the process. + onnxruntime::common::Status TryFlush(); + + // Test hook used to simulate a worker thread stopped while holding the cache mutex. + std::unique_lock LockMutexForTesting() const { + return std::unique_lock(mutex_); + } + + private: + struct DirtyRow { + std::optional config; + uint64_t generation = 0; + }; + using DirtyBucketMap = std::map; + + bool IsDirtyLocked(const MatMulNBitsKey& key, int m_bucket) const; + + onnxruntime::common::Status FlushImpl(bool allow_blocking); + + // Last write time and size of the file, or std::nullopt if it does not exist. + using FileStamp = std::pair; + std::optional CurrentFileStamp() const; + + onnxruntime::common::Status WriteAllLocked( + const std::unordered_map& table) const; + + std::string file_path_; + HardwareSignature signature_; + + mutable std::mutex mutex_; + std::unordered_map table_; + // Only locally changed rows participate in a flush; loaded rows must not undo another writer's update. + uint64_t next_dirty_generation_ = 0; + std::unordered_map dirty_rows_; + std::optional known_stamp_; +}; + +} // namespace onnxruntime::llm::gemm_cache diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh index dcd4d8a5ce133..db01d095c3e64 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh @@ -33,6 +33,14 @@ constexpr static int ACTIVATION_THREADS_PER_BLOCK = 256; struct QuantParams; +template +struct ReluSquared { + CUTLASS_HOST_DEVICE T operator()(const T& value) const { + const T positive = cutlass::epilogue::thread::ReLu{}(value); + return positive * positive; + } +}; + template class ActFn> __global__ void doGatedActivationKernel(ActivationOutputType* output, const GemmOutputType* gemm_result, const int64_t* num_valid_tokens_ptr, int64_t inter_size, @@ -129,8 +137,9 @@ void doGatedActivation(ActivationOutputType* output, const GemmOutputType* gemm_ case ActivationType::Gelu: return &doGatedActivationKernel; case ActivationType::Relu: - case ActivationType::Relu2: return &doGatedActivationKernel; + case ActivationType::Relu2: + return &doGatedActivationKernel; case ActivationType::Identity: default: return &doGatedActivationKernel; @@ -384,8 +393,9 @@ void doActivation(T* output, const GemmOutputType* gemm_result, const float* fp8 case ActivationType::Geglu: return &doActivationKernel; case ActivationType::Relu: - case ActivationType::Relu2: return &doActivationKernel; + case ActivationType::Relu2: + return &doActivationKernel; case ActivationType::Silu: case ActivationType::Swiglu: case ActivationType::SwigluBias: diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h index b1c242f89497e..864422e91fdf0 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h @@ -1064,8 +1064,7 @@ void MoeGemmRunner::moeGemmBiasAct( runGemm(inputs, hopper_inputs); break; case ActivationType::Relu2: - runGemm(inputs, hopper_inputs); - break; + ORT_THROW("Relu2 requires a separate squared-ReLU activation kernel."); case ActivationType::InvalidType: ORT_THROW("Activation type for fpA_intB must be valid."); break; diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu index 5d4a72f8d7e62..f436c59c98d85 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu @@ -11,6 +11,7 @@ // Shared device-side kernels + launch/dispatch helpers (fpA_intB_gemv namespace). #include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh" #include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" +#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" #include "core/platform/env_var_utils.h" namespace onnxruntime::llm { @@ -161,7 +162,13 @@ struct TypeTag { template using Fp4GemvAccT = std::conditional_t::value, half, float>; +template __device__ __forceinline__ float DecodeE4M3Fn(uint8_t code) { +#if defined(ORT_CUDA_TYPE_HELPER_HAS_FP8) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (NativeDecode) { + return onnxruntime::contrib::cuda::e4m3_to_float(code); + } +#endif const int sign = code & 0x80; const int exponent = (code >> 3) & 0x0f; const int mantissa = code & 0x07; @@ -353,7 +360,7 @@ __global__ void MoeGemvFp4RawNPackedKernel( #endif } -template +template __global__ void MoeGemvFp4RawKPackedKernel( const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales, const T* bias, T* out, @@ -362,9 +369,10 @@ __global__ void MoeGemvFp4RawKPackedKernel( cutlass_kernels::ActivationParams activation_params, const int* permuted_row_to_source_row, int num_rows) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - constexpr int kKLanes = 8; + constexpr int kKLanes = KLanes; constexpr int kBlockSize = 16; - constexpr int kColsPerBlock = 16; + constexpr int kColsPerBlock = 128 / kKLanes; + static_assert(kKLanes == 4 || kKLanes == 8); constexpr unsigned kFullMask = 0xffffffffu; using PackedT = std::conditional_t, half2, __nv_bfloat162>; @@ -383,21 +391,36 @@ __global__ void MoeGemvFp4RawKPackedKernel( } const int column = static_cast(blockIdx.y) * kColsPerBlock + threadIdx.x / kKLanes; + bool second_row = false; + if constexpr (ReuseExpertRows) { + const int first = static_cast(expert_first_token_offset[expert]); + const int end = static_cast(expert_first_token_offset[expert + 1]); + if (end - first >= 4) { + if ((row - first) % 2 != 0) { + return; + } + second_row = row + 1 < end; + } + } const int k_lane = threadIdx.x % kKLanes; const int source_row = permuted_row_to_source_row ? permuted_row_to_source_row[row] % num_rows : row; const T* row_act = act + static_cast(source_row) * reduction_size; + const int second_source_row = second_row + ? (permuted_row_to_source_row ? permuted_row_to_source_row[row + 1] % num_rows : row + 1) + : source_row; + const T* second_act = act + static_cast(second_source_row) * reduction_size; const uint8_t* expert_weight = weight + static_cast(expert) * weight_expert_stride; const uint8_t* expert_scales = block_scales + static_cast(expert) * scale_expert_stride; const float global_scale = global_scales[expert]; const T* expert_bias = EnableBias ? bias + static_cast(expert) * output_columns : nullptr; - float accumulators[2] = {}; + float accumulators[ReuseExpertRows ? 2 : 1][2] = {}; for (int reduction_column = k_lane * kBlockSize; column < output_columns && reduction_column < reduction_size; reduction_column += kKLanes * kBlockSize) { uint2 packed = *reinterpret_cast( expert_weight + static_cast(column) * (reduction_size / 2) + reduction_column / 2); - const T scale = static_cast(DecodeE4M3Fn( + const T scale = static_cast(DecodeE4M3Fn( expert_scales[static_cast(column) * (reduction_size / kBlockSize) + reduction_column / kBlockSize]) * global_scale); @@ -405,18 +428,31 @@ __global__ void MoeGemvFp4RawKPackedKernel( alignas(16) T activations[kBlockSize]; reinterpret_cast(activations)[0] = reinterpret_cast(row_act + reduction_column)[0]; reinterpret_cast(activations)[1] = reinterpret_cast(row_act + reduction_column)[1]; + alignas(16) T second_activations[kBlockSize]; + if constexpr (ReuseExpertRows) { + if (second_row) { + reinterpret_cast(second_activations)[0] = reinterpret_cast(second_act + reduction_column)[0]; + reinterpret_cast(second_activations)[1] = reinterpret_cast(second_act + reduction_column)[1]; + } + } alignas(4) T decoded[kBlockSize]; fiv::Fp4I2FConverter::template convert(&packed, decoded); #pragma unroll for (int pair = 0; pair < kBlockSize / 2; ++pair) { reinterpret_cast(decoded)[pair] = __hmul2(reinterpret_cast(decoded)[pair], *reinterpret_cast(scale_pair)); - accumulators[0] += static_cast(decoded[pair * 2]) * static_cast(activations[pair * 2]); - accumulators[1] += static_cast(decoded[pair * 2 + 1]) * static_cast(activations[pair * 2 + 1]); + accumulators[0][0] += static_cast(decoded[pair * 2]) * static_cast(activations[pair * 2]); + accumulators[0][1] += static_cast(decoded[pair * 2 + 1]) * static_cast(activations[pair * 2 + 1]); + if constexpr (ReuseExpertRows) { + if (second_row) { + accumulators[1][0] += static_cast(decoded[pair * 2]) * static_cast(second_activations[pair * 2]); + accumulators[1][1] += static_cast(decoded[pair * 2 + 1]) * static_cast(second_activations[pair * 2 + 1]); + } + } } } - float accumulator = accumulators[0] + accumulators[1]; + float accumulator = accumulators[0][0] + accumulators[0][1]; #pragma unroll for (int offset = kKLanes / 2; offset > 0; offset /= 2) { accumulator += __shfl_xor_sync(kFullMask, accumulator, offset, kKLanes); @@ -426,6 +462,408 @@ __global__ void MoeGemvFp4RawKPackedKernel( StoreMoeGemvFp4Pair( accumulator, paired_accumulator, expert_bias, out, row, column, output_columns, expert, activation_params); } + if constexpr (ReuseExpertRows) { + if (second_row) { + float second_accumulator = accumulators[1][0] + accumulators[1][1]; +#pragma unroll + for (int offset = kKLanes / 2; offset > 0; offset /= 2) { + second_accumulator += __shfl_xor_sync(kFullMask, second_accumulator, offset, kKLanes); + } + const float second_paired = __shfl_xor_sync(kFullMask, second_accumulator, kKLanes); + if (k_lane == 0 && column < output_columns && column % 2 == 0) { + StoreMoeGemvFp4Pair( + second_accumulator, second_paired, expert_bias, out, row + 1, column, output_columns, expert, activation_params); + } + } + } +#endif +} + +bool Fp4MoeGemvSupportsNativeScale() { +#if defined(ORT_CUDA_TYPE_HELPER_HAS_FP8) + cudaFuncAttributes attributes{}; + CUDA_CALL_THROW(cudaFuncGetAttributes( + &attributes, MoeGemvFp4RawKPackedKernel)); + return attributes.binaryVersion >= 90 && attributes.ptxVersion >= 90; +#else + return false; +#endif +} + +bool Fp4MoeGemvKernelUsesNativeScale(const void* kernel) { + return kernel == reinterpret_cast(MoeGemvFp4RawKPackedKernel) || + kernel == reinterpret_cast(MoeGemvFp4RawKPackedKernel) || + kernel == reinterpret_cast(MoeGemvFp4RawKPackedKernel<__nv_bfloat16, false, false, 4, false, true>) || + kernel == reinterpret_cast(MoeGemvFp4RawKPackedKernel<__nv_bfloat16, false, false, 8, false, true>); +} + +// --------------------------------------------------------------------------------------------- +// Decode-specialised NVFP4 GEMV for the row-major [E, N, K/2] layout (small expanded row counts). +// +// MoeGemvFp4RawKPackedKernel above issues 8-byte weight loads and re-reads/decodes the activation row +// for every 16-element block, which leaves the decode shapes issue/latency bound. This kernel streams the +// weights straight from global memory into registers with 128-bit loads, keeps several K tiles in flight, +// decodes FP4 in registers and stages the activation row once per CTA in shared memory as FP32. +// +// It reproduces MoeGemvFp4RawKPackedKernel's floating-point summation order exactly, so the results are +// bit-identical: K lane L = block (16 elements) % KLanes keeps separate even/odd-element accumulators over +// its blocks, a lane total is acc_even + acc_odd, and lane totals are combined by an xor tree over +// KLanes / 2, ..., 1. Here one thread owns the lane pair (2j, 2j + 1) of two adjacent columns: one 16-byte +// load covers the two blocks, the xor-tree steps over lane distances KLanes / 2 .. 2 become shuffles +// between the KLanes / 2 threads of a column, and the last step (lane distance 1) is an in-thread add. +template +__device__ __forceinline__ float PackedLowToFloat(uint32_t word) { + if constexpr (std::is_same_v) { + return __half2float(__ushort_as_half(static_cast(word & 0xffffu))); + } else { + return __uint_as_float(word << 16); + } +} + +template +__device__ __forceinline__ float PackedHighToFloat(uint32_t word) { + if constexpr (std::is_same_v) { + return __half2float(__ushort_as_half(static_cast(word >> 16))); + } else { + return __uint_as_float(word & 0xffff0000u); + } +} + +__device__ __forceinline__ uint4 LoadWeightStreaming(const void* ptr) { + uint4 value; + asm("ld.global.nc.L1::no_allocate.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(value.x), "=r"(value.y), "=r"(value.z), "=r"(value.w) + : "l"(ptr)); + return value; +} + +// The activation row is staged as FP32 in shared memory and must fit the default 48 KiB limit. +static constexpr int kDecodeMaxK = 8192; +// Default enablement (ORT_FP4_GEMV_DECODE_KERNEL=1) is restricted to the H200-measured decode envelope: +// SM90, up to 20 expanded rows, N and K up to 2560, and verified table-free scale eligibility +// (|global_scale| * 448 < 3.98). Set ORT_FP4_GEMV_DECODE_KERNEL=2 to opt into the full kernel envelope +// (SM90+, up to 80 expanded rows, K <= 8192, any even N, and the generic Fp4I2FConverter scale fallback). +static constexpr int kDecodeDefaultMaxN = 2560; +static constexpr int kDecodeDefaultMaxK = 2560; +static constexpr int64_t kDecodeDefaultMaxExpandedRows = 20; +// Tiling of the two KLanes variants (KLanes = 8 for K > 1024, 4 otherwise): warps per CTA, K tiles per +// prefetch batch and the min-blocks-per-SM hint that caps registers. Tuned on H200 for the Qwen3.8 FC1 +// (N=1280, K=2560, fused SwiGLU) and FC2 (N=2560, K=640) shapes at 10 expanded rows. +static constexpr int kDecodeK8Warps = 2; +static constexpr int kDecodeK8Batch = 2; +static constexpr int kDecodeK8MinBlocks = 2; +static constexpr int kDecodeK4Warps = 2; +static constexpr int kDecodeK4Batch = 2; +static constexpr int kDecodeK4MinBlocks = 4; +// Rows beyond this use the KPacked kernel, which amortizes less per-row work as the grid grows. +static constexpr int64_t kDecodeMaxExpandedRows = 80; + +// Table-free decode: place the three e2m1 magnitude bits (e1 e0 m) at the lowest exponent bits and the +// top mantissa bit of T and the sign at bit 15. The result is fp4_value * 2^-126 for bf16 and +// fp4_value * 2^-14 for fp16 (denormal for e == 0), so multiplying by block_scale * 2^126 (resp. 2^14) +// gives exactly T(fp4_value * block_scale) with a single rounding. Returns the pairs {nibble j, nibble +// j + 4} for j = 0..3. +template +__device__ __forceinline__ void Fp4WordToRaw(uint32_t word, uint32_t (&raw)[4]) { + if constexpr (std::is_same_v) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint32_t nibbles = word >> (4 * j); + raw[j] = ((nibbles & 0x00070007u) << 9) | ((nibbles & 0x00080008u) << 12); + } + } else { + const uint32_t upper = word >> 8; + raw[0] = ((word & 0x000F000Fu) * 0x1040u) & 0x81C081C0u; + raw[1] = ((word & 0x00F000F0u) * 0x0104u) & 0x81C081C0u; + raw[2] = ((upper & 0x000F000Fu) * 0x1040u) & 0x81C081C0u; + raw[3] = ((upper & 0x00F000F0u) * 0x0104u) & 0x81C081C0u; + } +} + +// {2^126, 2^126} in bf16 and {2^14, 2^14} in fp16. +template +static __device__ constexpr uint32_t kFp4RawRescale = std::is_same_v ? 0x74007400u : 0x7E807E80u; + +// Activation unit (32 elements = one lane pair's K slice of a tile) stride in shared memory: 36 floats = +// 9 float4 chunks, so the threads of a warp that read different units hit distinct banks. +static constexpr int kDecodeActUnitStride = 36; +// Every thread owns two adjacent columns, so the SwiGLU pair (n, n + 1) stays inside a thread. +static constexpr int kDecodeColsPerThread = 2; + +template +using PackedPairT = std::conditional_t, half2, __nv_bfloat162>; + +// Accumulates the 32 products of one weight vector per column. chain[c][lane * 2 + parity] is the +// accumulator of the lane (0: elements 0..15, 1: elements 16..31) and element parity. `decode(c, w, pairs)` +// yields the four packed pairs {element j, element j + 4} of 32-bit weight word w of column c; each +// activation chunk is read from shared memory once for all columns. +template +__device__ __forceinline__ void FmaWeightVector(const float* act, float (&chain)[kDecodeColsPerThread][4], + DecodeFn decode) { +#pragma unroll + for (int w = 0; w < 4; ++w) { + uint32_t decoded_pairs[kDecodeColsPerThread][4]; +#pragma unroll + for (int c = 0; c < kDecodeColsPerThread; ++c) { + decode(c, w, decoded_pairs[c]); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + // Elements 8w + 4h .. 8w + 4h + 3 are float4 chunk 2w + h of the unit. + const float4 a = *reinterpret_cast(act + (2 * w + h) * 4); + const float activations[4] = {a.x, a.y, a.z, a.w}; +#pragma unroll + for (int c = 0; c < kDecodeColsPerThread; ++c) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float decoded = h == 0 ? PackedLowToFloat(decoded_pairs[c][i]) + : PackedHighToFloat(decoded_pairs[c][i]); + chain[c][(w / 2) * 2 + (i & 1)] += decoded * activations[i]; + } + } + } + } +} + +// Block scales of the two blocks of one vector: scale_rounded = T(e4m3 * global_scale), packed +// {block A, block B} in T. +template +__device__ __forceinline__ uint32_t RoundedScalePair(uint32_t scale_code, float global_scale) { + uint32_t half_pair; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(half_pair) : "h"(static_cast(scale_code))); + const float scale_a = __half2float(__ushort_as_half(static_cast(half_pair & 0xffffu))) * global_scale; + const float scale_b = __half2float(__ushort_as_half(static_cast(half_pair >> 16))) * global_scale; + if constexpr (std::is_same_v) { + const half2 rounded = __floats2half2_rn(scale_a, scale_b); + return *reinterpret_cast(&rounded); + } else { + const __nv_bfloat162 rounded = __floats2bfloat162_rn(scale_a, scale_b); + return *reinterpret_cast(&rounded); + } +} + +// Operands of one weight vector per column, passed by value to the out-of-line generic decode so the +// caller's register-resident arrays are never address-taken. +struct DecodeVectorState { + uint4 packed[kDecodeColsPerThread]; + uint32_t scale_codes[kDecodeColsPerThread]; + float chain[kDecodeColsPerThread][4]; +}; + +// Generic decode (any T, any block scale): the Fp4I2FConverter path of the KPacked kernel. +template +__device__ __noinline__ DecodeVectorState AccumulateWeightVectorGeneric(DecodeVectorState state, float global_scale, + const float* act) { + uint32_t scale_rounded[kDecodeColsPerThread]; +#pragma unroll + for (int c = 0; c < kDecodeColsPerThread; ++c) { + scale_rounded[c] = RoundedScalePair(state.scale_codes[c], global_scale); + } + FmaWeightVector(act, state.chain, [&](int c, int w, uint32_t (&decoded_pairs)[4]) { + uint32_t word_copy = (&state.packed[c].x)[w]; + alignas(4) T decoded[8]; + fiv::Fp4I2FConverter::template convert<8>(&word_copy, decoded); + const uint32_t scale_splat = __byte_perm(scale_rounded[c], scale_rounded[c], (w / 2) ? 0x3232 : 0x1010); +#pragma unroll + for (int p = 0; p < 4; ++p) { + reinterpret_cast*>(decoded)[p] = __hmul2(reinterpret_cast*>(decoded)[p], + *reinterpret_cast*>(&scale_splat)); + } + const uint16_t* bits = reinterpret_cast(decoded); +#pragma unroll + for (int j = 0; j < 4; ++j) { + decoded_pairs[j] = static_cast(bits[j]) | (static_cast(bits[j + 4]) << 16); + } + }); + return state; +} + +// Fast path: valid when every block scale is below 4 (|global_scale| * 448 < 3.98), so that the rescaled +// block scale (block_scale * 2^126 in bf16, * 2^14 in fp16) stays finite. +template +__device__ __forceinline__ void AccumulateWeightVector( + const uint4 (&packed)[kDecodeColsPerThread], const uint32_t (&scale_codes)[kDecodeColsPerThread], + float global_scale, bool fast_scales, const float* act, float (&chain)[kDecodeColsPerThread][4]) { + if (fast_scales) { + uint32_t scale_raw[kDecodeColsPerThread][2]; +#pragma unroll + for (int c = 0; c < kDecodeColsPerThread; ++c) { + const uint32_t rounded_bits = RoundedScalePair(scale_codes[c], global_scale); + const PackedPairT shifted = __hmul2(*reinterpret_cast*>(&rounded_bits), + *reinterpret_cast*>(&kFp4RawRescale)); + const uint32_t shifted_bits = *reinterpret_cast(&shifted); + scale_raw[c][0] = __byte_perm(shifted_bits, shifted_bits, 0x1010); + scale_raw[c][1] = __byte_perm(shifted_bits, shifted_bits, 0x3232); + } + FmaWeightVector(act, chain, [&](int c, int w, uint32_t (&decoded_pairs)[4]) { + uint32_t raw[4]; + Fp4WordToRaw((&packed[c].x)[w], raw); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const PackedPairT product = __hmul2(*reinterpret_cast*>(&raw[j]), + *reinterpret_cast*>(&scale_raw[c][w / 2])); + decoded_pairs[j] = *reinterpret_cast(&product); + } + }); + return; + } + DecodeVectorState state; +#pragma unroll + for (int c = 0; c < kDecodeColsPerThread; ++c) { + state.packed[c] = packed[c]; + state.scale_codes[c] = scale_codes[c]; +#pragma unroll + for (int q = 0; q < 4; ++q) { + state.chain[c][q] = chain[c][q]; + } + } + state = AccumulateWeightVectorGeneric(state, global_scale, act); +#pragma unroll + for (int c = 0; c < kDecodeColsPerThread; ++c) { +#pragma unroll + for (int q = 0; q < 4; ++q) { + chain[c][q] = state.chain[c][q]; + } + } +} + +// A CTA covers Warps * 32 / (KLanes / 2) * 2 adjacent columns of one expanded row. Each thread loads +// `Batch` K tiles (KLanes blocks of 16 elements per column) at a time and prefetches the next batch while +// it accumulates the current one. +template +__global__ void __launch_bounds__(Warps * 32, MinBlocks) MoeGemvFp4DecodeKernel( + const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales, + const T* bias, T* out, + const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, + int64_t weight_expert_stride, int64_t scale_expert_stride, int n, int k, + cutlass_kernels::ActivationParams activation_params, + const int* permuted_row_to_source_row, int num_rows) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr int kThreadsPerColumn = KLanes / 2; + constexpr int kThreads = Warps * 32; + constexpr int kColsPerCta = kThreads / kThreadsPerColumn * kDecodeColsPerThread; + constexpr int kTileBytes = KLanes * 8; // weight bytes of one column and tile + constexpr unsigned kFullMask = 0xffffffffu; + static_assert(KLanes == 4 || KLanes == 8); + static_assert(Batch >= 1); + + const int row = static_cast(blockIdx.x); + int expert = permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : 0; +#pragma unroll 1 + for (int expert_index = 0; expert_index < num_experts && permuted_row_to_expert == nullptr; ++expert_index) { + if (row >= static_cast(expert_first_token_offset[expert_index + 1])) { + expert = expert_index + 1; + continue; + } + break; + } + if (expert < 0 || expert >= num_experts) { + return; + } + + const int j = threadIdx.x % kThreadsPerColumn; + const int first_column = static_cast(blockIdx.y) * kColsPerCta + + (static_cast(threadIdx.x) / kThreadsPerColumn) * kDecodeColsPerThread; + const bool column_valid = first_column < n; // n is even, so both columns of a thread are valid or neither + const int source_row = permuted_row_to_source_row ? permuted_row_to_source_row[row] % num_rows : row; + const int blocks = k / 16; + const int tiles = (blocks + KLanes - 1) / KLanes; + + const uint8_t* weight_row = weight + static_cast(expert) * weight_expert_stride + + static_cast(first_column) * (k / 2) + j * 16; + const uint8_t* scale_row = block_scales + static_cast(expert) * scale_expert_stride + + static_cast(first_column) * blocks + j * 2; + + // Threads whose lane pair lies past the end of K in a tile load nothing and skip the tile. + const auto tile_valid = [&](int tile) { return column_valid && tile * KLanes + j * 2 < blocks; }; + const auto load_batch = [&](int first_tile, uint4(&packed)[Batch][kDecodeColsPerThread], + uint32_t (&scale_codes)[Batch][kDecodeColsPerThread]) { +#pragma unroll + for (int b = 0; b < Batch; ++b) { + const int tile = first_tile + b; + const bool valid = tile < tiles && tile_valid(tile); +#pragma unroll + for (int c = 0; c < kDecodeColsPerThread; ++c) { + packed[b][c] = valid ? LoadWeightStreaming(weight_row + static_cast(c) * (k / 2) + tile * kTileBytes) + : make_uint4(0, 0, 0, 0); + scale_codes[b][c] = + valid ? *reinterpret_cast(scale_row + static_cast(c) * blocks + tile * KLanes) + : 0u; + } + } + }; + + uint4 packed[Batch][kDecodeColsPerThread]; + uint32_t scale_codes[Batch][kDecodeColsPerThread]; + load_batch(0, packed, scale_codes); + + // Stage the activation row as FP32, one padded 32-element unit per lane pair. + extern __shared__ __align__(16) float act_smem[]; + { + const uint4* source = reinterpret_cast(act + static_cast(source_row) * k); + for (int index = threadIdx.x; index < k / 8; index += kThreads) { + const uint4 raw = source[index]; + const uint32_t words[4] = {raw.x, raw.y, raw.z, raw.w}; + float values[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + values[2 * i] = PackedLowToFloat(words[i]); + values[2 * i + 1] = PackedHighToFloat(words[i]); + } + float* destination = act_smem + (index / 4) * kDecodeActUnitStride + (index % 4) * 8; + *reinterpret_cast(destination) = make_float4(values[0], values[1], values[2], values[3]); + *reinterpret_cast(destination + 4) = make_float4(values[4], values[5], values[6], values[7]); + } + } + const float global_scale = global_scales[expert]; + const bool fast_scales = fabsf(global_scale) * 448.0f < 3.98f; + __syncthreads(); + + float chain[kDecodeColsPerThread][4] = {}; +#pragma unroll 1 + for (int first_tile = 0; first_tile < tiles; first_tile += Batch) { + uint4 current_packed[Batch][kDecodeColsPerThread]; + uint32_t current_codes[Batch][kDecodeColsPerThread]; +#pragma unroll + for (int b = 0; b < Batch; ++b) { +#pragma unroll + for (int c = 0; c < kDecodeColsPerThread; ++c) { + current_packed[b][c] = packed[b][c]; + current_codes[b][c] = scale_codes[b][c]; + } + } + if (first_tile + Batch < tiles) { + load_batch(first_tile + Batch, packed, scale_codes); + } +#pragma unroll + for (int b = 0; b < Batch; ++b) { + const int tile = first_tile + b; + if (tile < tiles && tile_valid(tile)) { + const float* unit_act = act_smem + (tile * (KLanes / 2) + j) * kDecodeActUnitStride; + AccumulateWeightVector(current_packed[b], current_codes[b], global_scale, fast_scales, unit_act, chain); + } + } + } + + // Lane totals (acc_even + acc_odd) of the two lanes of the pair, then the xor tree over lane distances + // KLanes / 2 .. 2 as shuffles between the threads of a column. + float lane_total[kDecodeColsPerThread][2]; +#pragma unroll + for (int c = 0; c < kDecodeColsPerThread; ++c) { + lane_total[c][0] = chain[c][0] + chain[c][1]; + lane_total[c][1] = chain[c][2] + chain[c][3]; +#pragma unroll + for (int offset = kThreadsPerColumn / 2; offset > 0; offset /= 2) { + lane_total[c][0] += __shfl_xor_sync(kFullMask, lane_total[c][0], offset); + lane_total[c][1] += __shfl_xor_sync(kFullMask, lane_total[c][1], offset); + } + } + if (j == 0 && column_valid) { + const T* expert_bias = EnableBias ? bias + static_cast(expert) * n : nullptr; + StoreMoeGemvFp4Pair( + lane_total[0][0] + lane_total[0][1], lane_total[1][0] + lane_total[1][1], expert_bias, out, row, + first_column, n, expert, activation_params); + } #endif } @@ -435,25 +873,137 @@ void LaunchMoeGemvFp4RawNPacked( const T* bias, T* out, const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, int64_t n, int64_t k, cutlass_kernels::ActivationParams activation_params, - const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, bool weights_row_major) { + const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, bool weights_row_major, int sm, + bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr) { + if (decode_kernel_launched != nullptr) { + *decode_kernel_launched = false; + } constexpr int kThreads = 128; constexpr int kColsPerBlock = 16; const int64_t weight_expert_stride = n * k / 2; const int64_t scale_expert_stride = n * (k / 16); const dim3 grid(static_cast(expanded_num_rows), static_cast((n + kColsPerBlock - 1) / kColsPerBlock)); if (weights_row_major) { - if (bias != nullptr) { - MoeGemvFp4RawKPackedKernel<<>>( - act, weight, block_scales, global_scales, bias, out, - expert_first_token_offset, permuted_row_to_expert, num_experts, - weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k), activation_params, - permuted_row_to_source_row, static_cast(num_rows)); + static const bool enable_expert_reuse = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_NVFP4_FC1_EXPERT_REUSE", 0) == 1; + // Decode kernel mode (env ORT_FP4_GEMV_DECODE_KERNEL): + // 0: disabled (always use MoeGemvFp4RawKPackedKernel) + // 1: default (measured SM90 envelope with verified table-free scale eligibility) + // 2: opt-in full kernel envelope (SM90+, up to 80 expanded rows, K <= 8192, including generic scale fallback) + static const int decode_mode = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_DECODE_KERNEL", 1); + const bool reuse_experts = enable_expert_reuse && FusedSwiGlu && n == 1280 && k == 2560 && + expanded_num_rows >= 40 && expert_first_token_offset != nullptr && + permuted_row_to_expert != nullptr; + static const bool native_scale = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_NVFP4_GEMV_NATIVE_SCALE", false); + static const bool native_scale_supported = native_scale && Fp4MoeGemvSupportsNativeScale(); + const bool allow_extended_decode_envelope = decode_mode >= 2; + // When ORT_NVFP4_GEMV_NATIVE_SCALE=1 is enabled on SM90, the 4-lane native-scale KPacked kernel + // (MoeGemvFp4RawKPackedKernel) is slightly faster than MoeGemvFp4DecodeKernel + // for FP16 at K <= 1024 (e.g. Qwen3.8 FC2 K=640: 8.03 vs 8.41 us at top-10, 14.17 vs 14.23 us at top-20) + // while producing bit-identical results. + const bool prefer_native_kpacked_fp16 = + std::is_same_v && native_scale_supported && k <= 1024; + const bool decode_envelope_ok = + allow_extended_decode_envelope + ? (sm >= 90) + : (sm == 90 && decode_scale_eligible && !prefer_native_kpacked_fp16 && + expanded_num_rows <= kDecodeDefaultMaxExpandedRows && + n <= kDecodeDefaultMaxN && k <= kDecodeDefaultMaxK); + // Decode kernel: same summation order and KLanes dispatch as the KPacked kernel below (bit-identical). + // It steps aside for the expert-reuse variant and for shapes/alignments/scales outside the active envelope. + if (!reuse_experts && decode_mode != 0 && decode_envelope_ok && expanded_num_rows <= kDecodeMaxExpandedRows && + k % 32 == 0 && k <= kDecodeMaxK && n % 2 == 0 && reinterpret_cast(weight) % 16 == 0 && + reinterpret_cast(block_scales) % 2 == 0 && reinterpret_cast(act) % 16 == 0) { + const size_t decode_smem = static_cast(k / 32) * kDecodeActUnitStride * sizeof(float); + const auto launch_decode = [&](auto k_lanes, auto warps_tag, auto batch_tag, auto min_blocks_tag) { + constexpr int kKLanes = decltype(k_lanes)::value; + constexpr int kWarps = decltype(warps_tag)::value; + constexpr int kBatch = decltype(batch_tag)::value; + constexpr int kMinBlocks = decltype(min_blocks_tag)::value; + constexpr int kBlockCols = kWarps * 32 / (kKLanes / 2) * kDecodeColsPerThread; + const dim3 decode_grid(static_cast(expanded_num_rows), + static_cast((n + kBlockCols - 1) / kBlockCols)); + const auto launch_bias = [&](auto bias_tag) { + constexpr bool kEnableBias = decltype(bias_tag)::value; + // PTX compiled below SM90 has an empty decode body, even when JIT-loaded on SM90+. + // Image selection depends on the device SM; refresh when a thread switches architectures. + static thread_local int cached_sm = -1; + static thread_local bool available = false; + if (cached_sm != sm) { + cudaFuncAttributes attributes{}; + CUDA_CALL_THROW(cudaFuncGetAttributes( + &attributes, MoeGemvFp4DecodeKernel)); + available = attributes.ptxVersion >= 90; + cached_sm = sm; + } + if (!available) { + return false; + } + MoeGemvFp4DecodeKernel + <<>>( + act, weight, block_scales, global_scales, bias, out, + expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k), + activation_params, permuted_row_to_source_row, static_cast(num_rows)); + return true; + }; + if (bias != nullptr) { + return launch_bias(std::true_type{}); + } else { + return launch_bias(std::false_type{}); + } + }; + bool launched; + if (k <= 1024) { + launched = launch_decode(std::integral_constant{}, std::integral_constant{}, + std::integral_constant{}, std::integral_constant{}); + } else { + launched = launch_decode(std::integral_constant{}, std::integral_constant{}, + std::integral_constant{}, std::integral_constant{}); + } + if (launched) { + if (decode_kernel_launched != nullptr) { + *decode_kernel_launched = true; + } + return; + } + } + const auto launch_k_packed = [&](auto k_lanes, auto reuse_rows, auto native_tag) { + constexpr int kKLanes = decltype(k_lanes)::value; + constexpr bool kReuseRows = decltype(reuse_rows)::value; + constexpr bool kNativeScale = decltype(native_tag)::value; + constexpr int kNativeColsPerBlock = kThreads / kKLanes; + const dim3 native_grid(static_cast(expanded_num_rows), + static_cast((n + kNativeColsPerBlock - 1) / kNativeColsPerBlock)); + if (bias != nullptr) { + MoeGemvFp4RawKPackedKernel<<>>( + act, weight, block_scales, global_scales, bias, out, + expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k), activation_params, + permuted_row_to_source_row, static_cast(num_rows)); + } else { + MoeGemvFp4RawKPackedKernel<<>>( + act, weight, block_scales, global_scales, bias, out, + expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k), activation_params, + permuted_row_to_source_row, static_cast(num_rows)); + } + }; + const auto launch = [&](auto native_tag) { + if (k <= 1024) { + launch_k_packed(std::integral_constant{}, std::false_type{}, native_tag); + } else if (reuse_experts) { + launch_k_packed(std::integral_constant{}, std::true_type{}, native_tag); + } else { + launch_k_packed(std::integral_constant{}, std::false_type{}, native_tag); + } + }; + if (native_scale) { + launch(std::true_type{}); } else { - MoeGemvFp4RawKPackedKernel<<>>( - act, weight, block_scales, global_scales, bias, out, - expert_first_token_offset, permuted_row_to_expert, num_experts, - weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k), activation_params, - permuted_row_to_source_row, static_cast(num_rows)); + launch(std::false_type{}); } return; } @@ -552,6 +1102,27 @@ bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int return is_moe_gemv_fp4_supported(sm, expanded_num_rows, n, k, group_size, MoeGemvConfig::kDefault); } +bool is_moe_gemv_fp4_decode_image_available(int sm, bool is_fp16) { + if (sm < 90) { + return false; + } + const auto check_type = [&](auto type_tag) { + using T = typename decltype(type_tag)::type; + cudaFuncAttributes attributes{}; + CUDA_CALL_THROW(cudaFuncGetAttributes( + &attributes, MoeGemvFp4DecodeKernel)); + return attributes.ptxVersion >= 90; + }; + if (is_fp16) { + return check_type(TypeTag{}); + } +#ifdef ENABLE_BF16 + return check_type(TypeTag<__nv_bfloat16>{}); +#else + return false; +#endif +} + template void launch_moe_gemv_fp4_symmetric(const T* act, const uint8_t* weight, const T* scales, const uint8_t* raw_block_scales, const float* raw_global_scales, @@ -559,8 +1130,11 @@ void launch_moe_gemv_fp4_symmetric(const T* act, const uint8_t* weight, const T* const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, int sm, MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed, - cudaStream_t stream, bool weights_row_major) { - ORT_UNUSED_PARAMETER(sm); + cudaStream_t stream, bool weights_row_major, bool decode_scale_eligible, + bool* decode_kernel_launched) { + if (decode_kernel_launched != nullptr) { + *decode_kernel_launched = false; + } if (raw_n_packed) { ORT_ENFORCE(group_size == 16, "Raw N-packed FP4 GEMV is NVFP4-only."); ORT_ENFORCE(raw_block_scales != nullptr && raw_global_scales != nullptr, @@ -569,7 +1143,7 @@ void launch_moe_gemv_fp4_symmetric(const T* act, const uint8_t* weight, const T* act, weight, raw_block_scales, raw_global_scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, n, k, cutlass_kernels::ActivationParams{}, nullptr, expanded_num_rows, stream, - weights_row_major); + weights_row_major, sm, decode_scale_eligible, decode_kernel_launched); return; } // Interleaved path: ColumnMajorInterleaved layout + dtype-conditional accumulation + smaller @@ -628,8 +1202,11 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( int64_t expanded_num_rows, int64_t inter_size, int64_t k, int group_size, int sm, cutlass_kernels::ActivationParams activation_params, MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed, - const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, bool weights_row_major) { - ORT_UNUSED_PARAMETER(sm); + const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, bool weights_row_major, + bool decode_scale_eligible, bool* decode_kernel_launched) { + if (decode_kernel_launched != nullptr) { + *decode_kernel_launched = false; + } if (raw_n_packed) { ORT_ENFORCE(group_size == 16, "Raw N-packed FP4 GEMV is NVFP4-only."); ORT_ENFORCE(raw_block_scales != nullptr && raw_global_scales != nullptr, @@ -638,7 +1215,7 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( act, weight, raw_block_scales, raw_global_scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, inter_size * 2, k, activation_params, permuted_row_to_source_row, num_rows, stream, - weights_row_major); + weights_row_major, sm, decode_scale_eligible, decode_kernel_launched); return; } // Interleaved path: ColumnMajorInterleaved layout + dtype-conditional accumulation + smaller @@ -694,22 +1271,23 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( template void launch_moe_gemv_fp4_symmetric( const half*, const uint8_t*, const half*, const uint8_t*, const float*, const half*, half*, const int64_t*, const int*, int, - int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, bool, cudaStream_t, bool); + int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, bool, cudaStream_t, bool, bool, bool*); template void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( const half*, const uint8_t*, const half*, const uint8_t*, const float*, const half*, half*, const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, MoeGemvConfig, bool, bool, - const int*, int64_t, cudaStream_t, bool); + const int*, int64_t, cudaStream_t, bool, bool, bool*); #ifdef ENABLE_BF16 template void launch_moe_gemv_fp4_symmetric<__nv_bfloat16>( const __nv_bfloat16*, const uint8_t*, const __nv_bfloat16*, const uint8_t*, const float*, const __nv_bfloat16*, __nv_bfloat16*, - const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, bool, cudaStream_t, bool); + const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, bool, cudaStream_t, + bool, bool, bool*); template void launch_moe_gemv_fp4_symmetric_interleaved_swiglu<__nv_bfloat16>( const __nv_bfloat16*, const uint8_t*, const __nv_bfloat16*, const uint8_t*, const float*, const __nv_bfloat16*, __nv_bfloat16*, const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, - MoeGemvConfig, bool, bool, const int*, int64_t, cudaStream_t, bool); + MoeGemvConfig, bool, bool, const int*, int64_t, cudaStream_t, bool, bool, bool*); #endif } // namespace moe_gemv diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h index ed6d6d72f3a24..4081168fd07cc 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h @@ -7,6 +7,7 @@ #pragma once +#include #include #include @@ -41,6 +42,11 @@ inline constexpr int64_t kMaxProfiledExpandedRowsNvfp4KPacked = 80; // compute dispatch query this so the prepacked weights and the kernel always agree. bool Fp4MoeGemvUseInterleaved(); +// Checks the loaded kernel's compilation target, not just the current device's SM. +bool Fp4MoeGemvSupportsNativeScale(); +// Identifies native-scale specializations for unfused, no-bias K-packed launches. +bool Fp4MoeGemvKernelUsesNativeScale(const void* kernel); + // Shape-derived default tiling for the non-interleaved ColumnMajor FP4 GEMV. Used whenever the // runtime does not have a profiled result for the shape, which is the shipping default because // ORT_FP4_GEMV_AUTOTUNE is off (it synchronizes the inference stream) and is skipped entirely @@ -72,6 +78,27 @@ bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int // no separate ColToRow decode copy of the expert weights has to be allocated. bool is_moe_gemv_fp4_sm80_layout_supported(int64_t n, int64_t k, int group_size); +// Table-free E2M1 decode in MoeGemvFp4DecodeKernel requires every block scale (e4m3 * global_scale, +// with E4M3 max 448) to stay below 4 so the rescaled block scale (block_scale * 2^126 in bf16, +// * 2^14 in fp16) stays finite. Projections whose per-expert global scales exceed this bound fall +// back to MoeGemvFp4RawKPackedKernel by default instead of taking the slower generic decode branch. +inline constexpr float kNvfp4E4M3Max = 448.0f; +inline constexpr float kNvfp4DecodeFastScaleMaxProduct = 3.98f; + +__host__ __device__ inline bool IsNvfp4DecodeFastScale(float global_scale) { +#if defined(__CUDA_ARCH__) + return isfinite(global_scale) && fabsf(global_scale) * kNvfp4E4M3Max < kNvfp4DecodeFastScaleMaxProduct; +#else + return std::isfinite(global_scale) && std::fabs(global_scale) * kNvfp4E4M3Max < kNvfp4DecodeFastScaleMaxProduct; +#endif +} + +// Returns true when the CUDA image selected on the current device (of architecture `sm`) for +// MoeGemvFp4DecodeKernel was compiled with virtual architecture >= SM90. Forward-compatible PTX +// builds below SM90 (e.g. CMAKE_CUDA_ARCHITECTURES=80) compile an empty decode kernel body and +// return false here so callers fall back to MoeGemvFp4RawKPackedKernel. +bool is_moe_gemv_fp4_decode_image_available(int sm, bool is_fp16); + // Launches the MXFP4 (e2m1) MoE GEMV in the non-interleaved ColumnMajor layout. // act: [expanded_num_rows, k] permuted activations (row-major), T = half/bf16 // weight: [num_experts, n, k/2] e2m1 codes packed two per byte (even-K low nibble) @@ -94,7 +121,8 @@ void launch_moe_gemv_fp4_symmetric( const float* raw_global_scales, const T* bias, T* out, const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, int sm, MoeGemvConfig config, bool sm80_pair_interleaved, - bool raw_n_packed, cudaStream_t stream, bool weights_row_major = false); + bool raw_n_packed, cudaStream_t stream, bool weights_row_major = false, + bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr); // Launches the MXFP4 MoE GEMV and fuses interleaved SwiGLU activation. // weight/scales/bias use raw FC1 output width n = 2 * inter_size @@ -107,7 +135,7 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( int64_t inter_size, int64_t k, int group_size, int sm, cutlass_kernels::ActivationParams activation_params, MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed, const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, - bool weights_row_major = false); + bool weights_row_major = false, bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr); } // namespace moe_gemv } // namespace kernels diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu index 98b0f6b56cdcc..43bd9170cc772 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu @@ -2500,11 +2500,23 @@ void CutlassMoeFCRunner(quant_params.groupwise.fc1.weight_zeros) : nullptr, fc1_expert_biases, reinterpret_cast(output), alpha_scale_ptr_array, /*occupancy*/ nullptr, - fc1_activation_type, expanded_num_rows, /*N*/ int64_t(fc1_out_size), + fc1_activation_type == ActivationType::Relu2 + ? ActivationType::Identity + : fc1_activation_type, + expanded_num_rows, /*N*/ int64_t(fc1_out_size), /*K*/ hidden_size, num_experts_per_node, quant_params.groupwise.group_size, bias_is_broadcast, /*use_fused_moe*/ false, stream, activation_params, config}; gemm_runner.moeGemmBiasAct(universal_input, TmaWarpSpecializedGroupedGemmInput{}); + if (fc1_activation_type == ActivationType::Relu2) { + // Preserve the FC1 output-type rounding before squaring; the GEMM already added its bias. + doActivation( + output, reinterpret_cast(output), nullptr, + static_cast(nullptr), bias_is_broadcast, + expert_first_token_offset, num_experts_per_node, inter_size, expanded_num_rows, + fc1_activation_type, quant_params, false, nullptr, stream, activation_params); + } + sync_check_cuda_error(stream); } else { ORT_ENFORCE(!config.is_tma_warp_specialized); diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_base.h b/onnxruntime/contrib_ops/cuda/moe/moe_base.h index 95b54de0f2cbb..3fca3b6153903 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_base.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_base.h @@ -49,6 +49,8 @@ class MoEBase { ORT_ENFORCE(op_kernel_info.GetAttr("activation_type", &activation_type_str).IsOK()); if (activation_type_str == "relu") { activation_type_ = ActivationType::Relu; + } else if (activation_type_str == "relu2") { + activation_type_ = ActivationType::Relu2; } else if (activation_type_str == "gelu") { activation_type_ = ActivationType::Gelu; } else if (activation_type_str == "silu") { diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index a6b9f0cd36fa4..7ab59ba44abd2 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -1453,6 +1453,10 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } } const bool use_packed_int = use_packed_int_gemv || use_packed_int_prefill; + const bool use_fp8_fused = is_block_fp8 && enable_fp8_fused_ && + (kernel_activation_type == ActivationType::Silu || + kernel_activation_type == ActivationType::Swiglu) && + moe_params.num_rows <= 65535 / k_; if (use_int_dequant_fallback && !use_packed_int) { const size_t total_dequant_bytes = SafeInt(int_dequant_fc1_bytes) + int_dequant_fc2_bytes; ORT_RETURN_IF_NOT(total_dequant_bytes <= static_cast(int_dequant_max_scratch_bytes_), @@ -1462,11 +1466,6 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { kQMoEIntDequantMaxScratchBytesConfig, " only for bounded correctness workloads."); } - const bool use_fp8_fused = is_block_fp8 && enable_fp8_fused_ && - (kernel_activation_type == ActivationType::Silu || - kernel_activation_type == ActivationType::Swiglu) && - moe_params.num_rows <= 65535 / k_; - const qmoe::RowTilePlan row_tile_plan = qmoe::MakeRowTilePlan( moe_params.num_rows, effective_row_tile_size, @@ -1497,8 +1496,11 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { std::array runner_tile_configs{}; size_t runner_tile_config_count = 0; size_t workspace_size = 0; - // Standalone FP4 GEMV returns before runMoe and does not consume grouped-GEMM tactics or workspace. - if (!use_fp8_fused && !use_packed_int && !use_nvfp4_packed_prefill && !run_fp4_gemv) { + // The packed INT GEMV path allocates its own scratch below and returns before reaching the + // dense grouped-GEMM tile loop, so profiling/sizing the dense runner here would be pure + // overhead (mutex, two dense-tactic profiling launches, and an unused large workspace + // allocation). Skip it entirely for that path; workspace_size stays 0. + if (!use_packed_int && !use_fp8_fused) { std::lock_guard profiler_lock(mGemmProfilerMutex); // Profiling launches grouped-GEMM kernels, records/synchronizes CUDA events, and @@ -1704,14 +1706,6 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { return fused_routing; }; -#if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) - if (routing_snapshot_) { - auto* pilot = context->GetKernelPilot(); - ORT_RETURN_IF_NOT(pilot, "MoE expert tracking is enabled but its collector is unavailable."); - ORT_RETURN_IF_ERROR(routing_snapshot_->BeginInvocation(pilot->Moe(), static_cast(moe_params.num_experts))); - } -#endif - if (use_fp8_fused) { Tensor* output = context->Output(0, input->Shape()); namespace ck = onnxruntime::llm::kernels::cutlass_kernels; @@ -1722,11 +1716,13 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { const int fc1_n = inter * (fusion ? 2 : 1); const int max_routes = narrow(row_tile_plan.rows_per_tile * k_); const size_t element_size = input->DataType()->Size(); + const size_t projection_element_size = fusion ? element_size : sizeof(float); auto row_map = GetScratchBuffer(max_routes, GetComputeStream(context)); auto sorted_experts = GetScratchBuffer(max_routes, GetComputeStream(context)); auto expert_offsets = GetScratchBuffer(num_experts + 1, GetComputeStream(context)); auto tile_offsets = GetScratchBuffer(num_experts + 1, GetComputeStream(context)); - auto projection = GetScratchBuffer(SafeInt(max_routes) * fc1_n * element_size, GetComputeStream(context)); + const size_t projection_bytes = SafeInt(max_routes) * fc1_n * projection_element_size; + auto projection = GetScratchBuffer(projection_bytes, GetComputeStream(context)); auto activation = GetScratchBuffer(SafeInt(max_routes) * inter * element_size, GetComputeStream(context)); auto result = GetScratchBuffer(SafeInt(max_routes) * hidden * element_size, GetComputeStream(context)); auto scale_type = [](const Tensor* tensor) { @@ -1742,7 +1738,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { route_tile(row_offset, rows); #if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) if (routing_snapshot_) { - ORT_RETURN_IF_ERROR(routing_snapshot_->Capture(expert_indices, SafeInt(routes), stream)); + ORT_RETURN_IF_ERROR(routing_snapshot_->Capture(expert_indices, routes, stream)); } #endif IAllocatorUniquePtr counts; @@ -1785,13 +1781,21 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { params.block_size = narrow(block_size_); params.fusion = fusion; const auto* tile_input = static_cast(input->DataRaw()) + row_offset * hidden; - LaunchQMoEFp8Projection(params, tile_input, static_cast(projection.get()), stream); const auto* fc1_bias = fc1_experts_bias_optional ? static_cast(fc1_experts_bias_optional->DataRaw()) : nullptr; - LaunchQMoEFp8Activation(static_cast(projection.get()), fc1_bias, static_cast(activation.get()), - sorted_experts.get(), routes, inter, fusion, - split_fp8_fc1 || !fusion ? 1.0f : activation_alpha_, - split_fp8_fc1 ? 0.0f : activation_beta_, - split_fp8_fc1 ? std::numeric_limits::infinity() : swiglu_limit_, stream); + auto project_and_activate = [&](auto* projected) { + LaunchQMoEFp8Projection(params, tile_input, projected, stream); + LaunchQMoEFp8Activation(projected, fc1_bias, static_cast(activation.get()), + sorted_experts.get(), routes, inter, fusion, + split_fp8_fc1 || !fusion ? 1.0f : activation_alpha_, + split_fp8_fc1 ? 0.0f : activation_beta_, + split_fp8_fc1 ? std::numeric_limits::infinity() : swiglu_limit_, stream); + }; + if (fusion) { + project_and_activate(static_cast(projection.get())); + } else { + // Match the dense SiLU epilogue: narrow only after bias and activation. + project_and_activate(static_cast(projection.get())); + } params.weights = static_cast(fc2_experts_weights->DataRaw()); params.scales = fc2_scale_tensor->DataRaw(); params.scale_type = scale_type(fc2_scale_tensor); @@ -1820,7 +1824,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { #endif } if (enable_kernel_debug_info_) { - const size_t activation_bytes = SafeInt(max_routes) * (fc1_n + inter + hidden) * element_size; + const size_t activation_bytes = SafeInt(projection_bytes) + + SafeInt(max_routes) * (inter + hidden) * element_size; PrintQMoEKernelDebugInfo(row_tile_plan.rows_per_tile > 8 ? "fp8_fused_gemm" : "fp8_fused_gemv", moe_params.num_rows, row_tile_plan.rows_per_tile, final_tile_rows, max_routes, final_tile_rows * k_, 0, total_scratch_bytes + activation_bytes); @@ -2514,6 +2519,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { const bool is_capturing = capture_status != cudaStreamCaptureStatusNone; const bool do_tune = enable_fp4_gemv_autotune_ && !have_tune && !is_capturing; + bool fc1_used_decode = false; + bool fc2_used_decode = false; auto run_fused = [&](auto* t_ptr) { using T = std::remove_pointer_t; const bool skip_expand = fp4_gemv_skip_expand_; @@ -2533,7 +2540,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { static_cast(fc1_bias), static_cast(p_fc1_buf.get()), p_efto, p_exp, num_experts, expanded, inter, hidden, gemv_group_size, sm_, act_params, cfg, fc1_gemv_sm80_layout, use_raw_nvfp4_gemv, skip_expand ? p_r2u : nullptr, num_rows, stream, - nvfp4_weights_row_major_); + nvfp4_weights_row_major_, nvfp4_fc1_decode_scale_eligible_, &fc1_used_decode); }; auto launch_fc2 = [&](MoeGemvConfig cfg) { gemv::launch_moe_gemv_fp4_symmetric( @@ -2543,7 +2550,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { gemv_fc2_raw_global_scales, static_cast(fc2_bias), static_cast(p_fc2_buf.get()), p_efto, p_exp, num_experts, expanded, hidden, inter, gemv_group_size, sm_, cfg, - fc2_gemv_sm80_layout, use_raw_nvfp4_gemv, stream, nvfp4_weights_row_major_); + fc2_gemv_sm80_layout, use_raw_nvfp4_gemv, stream, nvfp4_weights_row_major_, + nvfp4_fc2_decode_scale_eligible_, &fc2_used_decode); }; if (do_tune) { @@ -2647,6 +2655,15 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { moe_params.num_rows, moe_params.num_rows, expanded, expanded, workspace_size, total_scratch_bytes); + if (nvfp4_weights_row_major_) { + std::cout << "QMoE NVFP4 DecodeImageAvailable=" + << (gemv::is_moe_gemv_fp4_decode_image_available(sm_, is_fp16_) ? 1 : 0) + << " FC1ScaleEligible=" << (nvfp4_fc1_decode_scale_eligible_ ? 1 : 0) + << " FC2ScaleEligible=" << (nvfp4_fc2_decode_scale_eligible_ ? 1 : 0) + << " FC1Kernel=" << (fc1_used_decode ? "decode" : "kpacked") + << " FC2Kernel=" << (fc2_used_decode ? "decode" : "kpacked") + << std::endl; + } } #if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) if (routing_snapshot_) { @@ -2752,8 +2769,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { fc2_weight_data = dequant_fc2_weights.get(); } // FP4 (W4A16) and WFP4AFP8 (W4A8) share the MXFP4 weight format (Float8E8M0 block scales, block 32). - // NVFP4 (W4A16) uses Float8E4M3FN block scales with block 16 and always runs the dequant fallback. - // When the native CUTLASS path is unavailable on the current SM (always for NVFP4), or when native + // NVFP4 (W4A16) uses Float8E4M3FN block scales with block 16. + // When the native CUTLASS path is unavailable on the current SM, or when native // FP4 routes this call to the dense fallback for the large per-expert-M regime, dequantize the E2M1 // weights to FP16/BF16 and run the dense A16 runner. if (!use_fp4_deep_gemm && @@ -2761,16 +2778,18 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { !fp4_sm80_prefill) { // The dequant kernel expects raw [E, n, k_blocks] block scales. When native FP4 is enabled // (this is the large per-expert-M fallback), packed_fp4_*_block_scales_ holds the TMA-swizzled - // layout, so use the raw copy kept in gemv_fp4_*_block_raw_ instead. NVFP4 and the SM<90 - // dequant-only build have no raw copy, so packed_fp4_*_block_scales_ already holds the raw scales. - const void* p_fc1_block_scales = gemv_fp4_fc1_block_raw_ ? gemv_fp4_fc1_block_raw_.get() - : (packed_fp4_fc1_block_scales_ ? packed_fp4_fc1_block_scales_.get() - : (fp4_fc1_block_scales ? fp4_fc1_block_scales->DataRaw() : nullptr)); + // layout, so use the raw copy kept in gemv_fp4_*_block_raw_ instead. NVFP4 retains its raw + // scale inputs; prefer those even when native prepacking also created swizzled scales. + const void* p_fc1_block_scales = is_nvfp4 && fp4_fc1_block_scales ? fp4_fc1_block_scales->DataRaw() + : gemv_fp4_fc1_block_raw_ ? gemv_fp4_fc1_block_raw_.get() + : (packed_fp4_fc1_block_scales_ ? packed_fp4_fc1_block_scales_.get() + : (fp4_fc1_block_scales ? fp4_fc1_block_scales->DataRaw() : nullptr)); const void* p_fc1_global_scale = packed_fc1_global_scale_ ? packed_fc1_global_scale_.get() : (fc1_global_scale ? fc1_global_scale->DataRaw() : nullptr); - const void* p_fc2_block_scales = gemv_fp4_fc2_block_raw_ ? gemv_fp4_fc2_block_raw_.get() - : (packed_fp4_fc2_block_scales_ ? packed_fp4_fc2_block_scales_.get() - : (fp4_fc2_block_scales ? fp4_fc2_block_scales->DataRaw() : nullptr)); + const void* p_fc2_block_scales = is_nvfp4 && fp4_fc2_block_scales ? fp4_fc2_block_scales->DataRaw() + : gemv_fp4_fc2_block_raw_ ? gemv_fp4_fc2_block_raw_.get() + : (packed_fp4_fc2_block_scales_ ? packed_fp4_fc2_block_scales_.get() + : (fp4_fc2_block_scales ? fp4_fc2_block_scales->DataRaw() : nullptr)); const void* p_fc2_global_scale = packed_fc2_global_scale_ ? packed_fc2_global_scale_.get() : (fc2_global_scale ? fc2_global_scale->DataRaw() : nullptr); ORT_RETURN_IF_NOT(p_fc1_block_scales && p_fc1_global_scale && p_fc2_block_scales && p_fc2_global_scale, @@ -3387,11 +3406,37 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, (quant_type_ == "fp4" || quant_type_ == "nvfp4" || quant_type_ == "fp8" || quant_type_ == "wfp4afp8")) { ORT_ENFORCE(tensor.IsDataType() && tensor.Shape().NumDimensions() == 1, "QMoE ", quant_type_, " global scales must be 1-D float tensors."); + const auto is_nvfp4_decode_scale_eligible = [&](const Tensor& scale_tensor) -> bool { + if (quant_type_ != "nvfp4") { + return false; + } + const int64_t count = scale_tensor.Shape().Size(); + if (count <= 0) { + return false; + } + const float* scales = scale_tensor.Data(); + std::unique_ptr host_scales; + if (scale_tensor.Location().device.Type() != OrtDevice::CPU) { + host_scales = std::make_unique(static_cast(count)); + CUDA_CALL_THROW(cudaMemcpyAsync(host_scales.get(), scales, static_cast(count) * sizeof(float), + cudaMemcpyDeviceToHost, stream)); + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + scales = host_scales.get(); + } + for (int64_t i = 0; i < count; ++i) { + if (!onnxruntime::llm::kernels::moe_gemv::IsNvfp4DecodeFastScale(scales[i])) { + return false; + } + } + return true; + }; if (input_idx == 15) { + nvfp4_fc1_decode_scale_eligible_ = is_nvfp4_decode_scale_eligible(tensor); PrePackCopyToGpu(tensor, stream, alloc, packed_fc1_global_scale_, is_packed); fp4_fc1_global_scale_e_ = tensor.Shape()[0]; TryBuildGemvFp4Scales(1, stream, alloc); } else { + nvfp4_fc2_decode_scale_eligible_ = is_nvfp4_decode_scale_eligible(tensor); PrePackCopyToGpu(tensor, stream, alloc, packed_fc2_global_scale_, is_packed); fp4_fc2_global_scale_e_ = tensor.Shape()[0]; TryBuildGemvFp4Scales(2, stream, alloc); diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h index 77525b7433d18..a6453001e43a1 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h @@ -93,6 +93,8 @@ class QMoE final : public CudaKernel, public MoEBase { // 1 is reserved for a possible future Hopper-specific layout (e.g. W4A8). bool weights_prepacked_ = true; bool nvfp4_weights_row_major_ = false; + bool nvfp4_fc1_decode_scale_eligible_ = false; + bool nvfp4_fc2_decode_scale_eligible_ = false; // Cached source weight shapes captured at PrePack time. When the // PrePack hook consumed and released the original int4/int8 or MXFP4 // weight initializers (``is_packed = true``), ``context->Input(2)`` diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu index 5d4c2dd3a4a19..3c91fe4b7a967 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu @@ -10,10 +10,13 @@ #include #include +#include #include "contrib_ops/cuda/llm/moe_gemm/moe_kernels.h" #include "core/common/narrow.h" +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/cu_inc/cub.cuh" +#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" #include "core/providers/cuda/cu_inc/topk_warp_sort.cuh" #include "core/providers/cuda/cuda_common.h" @@ -1026,7 +1029,15 @@ __device__ __forceinline__ float DecodeUE8M0(uint8_t code) { return code == 0 ? 0.0f : exp2f(static_cast(code) - 127); } +template __device__ __forceinline__ float DecodeFloat8E4M3FN(uint8_t code) { +#if defined(ORT_CUDA_TYPE_HELPER_HAS_FP8) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (NativeDecode) { + return e4m3_to_float(code); + } +#endif + // ONNX float8e4m3fn has no infinities. The only NaN payloads are 0x7F/0xFF; + // finite values, including the max finite code 0x7E, use the normal E4M3 formula. const uint32_t sign = static_cast(code & 0x80) << 24; const uint32_t magnitude = code & 0x7F; if (magnitude == 0x7F) { @@ -1536,15 +1547,8 @@ void LaunchQMoECompactExperts( CUDA_CALL_THROW(cudaGetLastError()); } -struct QMoEFp8WeightRow { - const uint8_t* weights; - const void* scales; - int scale_type; - int64_t scale_offset; -}; - -__device__ __forceinline__ QMoEFp8WeightRow ResolveQMoEFp8WeightRow( - const QMoEFp8ProjectionParams& params, int expert, int row) { +template +__device__ __forceinline__ T QMoEFp8Weight(const QMoEFp8ProjectionParams& params, int expert, int row, int col) { const uint8_t* weights = params.weights; const void* scales = params.scales; int scale_type = params.scale_type; @@ -1562,27 +1566,17 @@ __device__ __forceinline__ QMoEFp8WeightRow ResolveQMoEFp8WeightRow( } const int scale_n = (source_n + params.block_size - 1) / params.block_size; const int scale_k = (params.k + params.block_size - 1) / params.block_size; - const int64_t scale_offset = (static_cast(expert) * scale_n + row / params.block_size) * scale_k; - const int64_t weight_offset = (static_cast(expert) * source_n + row) * params.k; - return {weights + weight_offset, scales, scale_type, scale_offset}; -} - -__device__ __forceinline__ float QMoEFp8RowScale(const QMoEFp8WeightRow& row, int block) { - const int64_t index = row.scale_offset + block; - return row.scale_type == 0 ? static_cast(row.scales)[index] - : row.scale_type == 1 ? __half2float(static_cast(row.scales)[index]) - : __bfloat162float(static_cast(row.scales)[index]); -} - -template -__device__ __forceinline__ T QMoEFp8Weight(const QMoEFp8ProjectionParams& params, int expert, int row, int col) { - const auto weight_row = ResolveQMoEFp8WeightRow(params, expert, row); - const float scale = QMoEFp8RowScale(weight_row, col / params.block_size); - return static_cast(DecodeFloat8E4M3FN(weight_row.weights[col]) * scale); + const int64_t scale_index = (static_cast(expert) * scale_n + row / params.block_size) * scale_k + + col / params.block_size; + const int64_t weight_index = (static_cast(expert) * source_n + row) * params.k + col; + const float scale = scale_type == 0 ? static_cast(scales)[scale_index] + : scale_type == 1 ? __half2float(static_cast(scales)[scale_index]) + : __bfloat162float(static_cast(scales)[scale_index]); + return static_cast(DecodeFloat8E4M3FN(weights[weight_index]) * scale); } -template -__global__ void QMoEFp8GemvKernel(QMoEFp8ProjectionParams params, const T* input, T* output) { +template +__global__ void QMoEFp8GemvKernel(QMoEFp8ProjectionParams params, const T* input, OutputT* output) { const int row = static_cast(blockIdx.y); const int expert = params.experts[row]; const int output_col = static_cast(blockIdx.x) * 4 + threadIdx.x / 32; @@ -1590,34 +1584,34 @@ __global__ void QMoEFp8GemvKernel(QMoEFp8ProjectionParams params, const T* input return; } const int input_row = params.row_to_unpermuted ? params.row_to_unpermuted[row] % params.num_rows : row; - const auto weight_row = ResolveQMoEFp8WeightRow(params, expert, output_col); - const T* activation = input + static_cast(input_row) * params.k; float sum = 0.0f; - if (params.block_size == 128) { - for (int block_begin = 0; block_begin < params.k; block_begin += 128) { - const float scale = QMoEFp8RowScale(weight_row, block_begin / 128); -#pragma unroll - for (int offset = 0; offset < 128; offset += 32) { - const int col = block_begin + threadIdx.x % 32 + offset; - if (col < params.k) { - const T weight = static_cast(DecodeFloat8E4M3FN(weight_row.weights[col]) * scale); - sum = fmaf(static_cast(activation[col]), static_cast(weight), sum); - } - } - } - } else { - for (int col = threadIdx.x % 32; col < params.k; col += 32) { - const T weight = static_cast(DecodeFloat8E4M3FN(weight_row.weights[col]) * - QMoEFp8RowScale(weight_row, col / params.block_size)); - sum = fmaf(static_cast(activation[col]), static_cast(weight), sum); - } + for (int col = threadIdx.x % 32; col < params.k; col += 32) { + sum = fmaf(static_cast(input[static_cast(input_row) * params.k + col]), + static_cast(QMoEFp8Weight(params, expert, output_col, col)), sum); } sum = WarpReduceSum(sum); if (threadIdx.x % 32 == 0) { - output[static_cast(row) * params.n + output_col] = static_cast(sum); + output[static_cast(row) * params.n + output_col] = static_cast(sum); } } +bool QMoEFp8GemvSupportsNativeDecode() { +#if defined(ORT_CUDA_TYPE_HELPER_HAS_FP8) + cudaFuncAttributes attributes{}; + CUDA_CALL_THROW(cudaFuncGetAttributes(&attributes, QMoEFp8GemvKernel)); + return attributes.binaryVersion >= 90 && attributes.ptxVersion >= 90; +#else + return false; +#endif +} + +bool QMoEFp8GemvKernelUsesNativeDecode(const void* kernel) { + return kernel == reinterpret_cast(QMoEFp8GemvKernel) || + kernel == reinterpret_cast(QMoEFp8GemvKernel<__nv_bfloat16, true>) || + kernel == reinterpret_cast(QMoEFp8GemvKernel) || + kernel == reinterpret_cast(QMoEFp8GemvKernel<__nv_bfloat16, true, float>); +} + __global__ void QMoEFp8ExpertTilesKernel(const int64_t* expert_offsets, int* tile_offsets, int num_experts) { int tiles = 0; tile_offsets[0] = 0; @@ -1633,8 +1627,8 @@ void LaunchQMoEFp8ExpertTiles(const int64_t* expert_offsets, int* tile_offsets, CUDA_CALL_THROW(cudaGetLastError()); } -template -__global__ void QMoEFp8GemmKernel(QMoEFp8ProjectionParams params, const T* input, T* output) { +template +__global__ void QMoEFp8GemmKernel(QMoEFp8ProjectionParams params, const T* input, OutputT* output) { #if __CUDA_ARCH__ >= 800 namespace wmma = nvcuda::wmma; const int tile = static_cast(blockIdx.y); @@ -1696,27 +1690,33 @@ __global__ void QMoEFp8GemmKernel(QMoEFp8ProjectionParams params, const T* input const int64_t row = row_start + index / 64; const int col = column_start + index % 64; if (row < row_end && col < params.n) { - output[row * params.n + col] = static_cast(output_tile[index]); + output[row * params.n + col] = static_cast(output_tile[index]); } } #endif } -template +template void LaunchQMoEFp8Projection(const QMoEFp8ProjectionParams& params, - const T* input, T* output, cudaStream_t stream) { + const T* input, OutputT* output, cudaStream_t stream) { if (params.tile_offsets) { const int tiles = std::min(params.expanded_rows, std::min(params.num_experts, params.expanded_rows) + params.expanded_rows / 16); - QMoEFp8GemmKernel<<>>(params, input, output); + QMoEFp8GemmKernel<<>>(params, input, output); } else { - QMoEFp8GemvKernel<<>>(params, input, output); + static const bool native_decode = + ParseEnvironmentVariableWithDefault("ORT_QMOE_FP8_GEMV_NATIVE_DECODE", false); + if (native_decode) { + QMoEFp8GemvKernel<<>>(params, input, output); + } else { + QMoEFp8GemvKernel<<>>(params, input, output); + } } CUDA_CALL_THROW(cudaGetLastError()); } -template -__global__ void QMoEFp8ActivationKernel(const T* input, const T* bias, T* output, +template +__global__ void QMoEFp8ActivationKernel(const InputT* input, const T* bias, T* output, const int* experts, int rows, int inter_size, int fusion, float alpha, float beta, float limit) { const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; @@ -1744,8 +1744,8 @@ __global__ void QMoEFp8ActivationKernel(const T* input, const T* bias, T* output output[index] = static_cast(gate / (1.0f + expf(-alpha * gate)) * up); } -template -void LaunchQMoEFp8Activation(const T* input, const T* bias, T* output, +template +void LaunchQMoEFp8Activation(const InputT* input, const T* bias, T* output, const int* experts, int rows, int inter_size, int fusion, float alpha, float beta, float limit, cudaStream_t stream) { const int64_t count = static_cast(rows) * inter_size; @@ -1756,8 +1756,12 @@ void LaunchQMoEFp8Activation(const T* input, const T* bias, T* output, template void LaunchQMoEFp8Projection(const QMoEFp8ProjectionParams&, const half*, half*, cudaStream_t); template void LaunchQMoEFp8Projection<__nv_bfloat16>(const QMoEFp8ProjectionParams&, const __nv_bfloat16*, __nv_bfloat16*, cudaStream_t); +template void LaunchQMoEFp8Projection(const QMoEFp8ProjectionParams&, const half*, float*, cudaStream_t); +template void LaunchQMoEFp8Projection<__nv_bfloat16, float>(const QMoEFp8ProjectionParams&, const __nv_bfloat16*, float*, cudaStream_t); template void LaunchQMoEFp8Activation(const half*, const half*, half*, const int*, int, int, int, float, float, float, cudaStream_t); template void LaunchQMoEFp8Activation<__nv_bfloat16>(const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, const int*, int, int, int, float, float, float, cudaStream_t); +template void LaunchQMoEFp8Activation(const float*, const half*, half*, const int*, int, int, int, float, float, float, cudaStream_t); +template void LaunchQMoEFp8Activation<__nv_bfloat16, float>(const float*, const __nv_bfloat16*, __nv_bfloat16*, const int*, int, int, int, float, float, float, cudaStream_t); template __global__ void QMoEDequantizeFp8WeightsKernel( @@ -1978,16 +1982,6 @@ void LaunchQMoEDequantizeNvfp4WeightsImpl( bool weights_row_major) { ORT_ENFORCE(bias == nullptr || output_bias != nullptr, "QMoE NVFP4 bias gathering requires an output buffer."); constexpr int block = 256; - const int64_t row_major_grid_y = - (static_cast(k) + block * kQMoEDequantizeFp4VecK - 1) / - (block * kQMoEDequantizeFp4VecK); - if (weights_row_major && k % 16 == 0 && row_major_grid_y <= 65535 && num_experts <= 65535) { - const dim3 grid(n, onnxruntime::narrow(row_major_grid_y), num_experts); - QMoEDequantizeNvfp4RowMajorWeightsVecKernel<<>>( - packed_weights, block_scales, global_scales, output, n, k, compact_to_expert, bias, output_bias); - CUDA_CALL_THROW(cudaGetLastError()); - return; - } if (!weights_row_major && !compact_to_expert && !bias && QMoEDequantizeFp4VecApplies<16>(num_experts, n, k)) { const dim3 tile_block(kQMoEDequantizeFp4TileK / kQMoEDequantizeFp4VecK, kQMoEDequantizeFp4TileN); const dim3 tile_grid((n + kQMoEDequantizeFp4TileN - 1) / kQMoEDequantizeFp4TileN, diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h index 7bb9b771a92d2..00a49bddbe72d 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h @@ -293,12 +293,15 @@ struct QMoEFp8ProjectionParams { void LaunchQMoEFp8ExpertTiles(const int64_t* expert_offsets, int* tile_offsets, int num_experts, cudaStream_t stream); -template +template void LaunchQMoEFp8Projection(const QMoEFp8ProjectionParams& params, - const T* input, T* output, cudaStream_t stream); + const T* input, OutputT* output, cudaStream_t stream); -template -void LaunchQMoEFp8Activation(const T* input, const T* bias, T* output, +bool QMoEFp8GemvSupportsNativeDecode(); +bool QMoEFp8GemvKernelUsesNativeDecode(const void* kernel); + +template +void LaunchQMoEFp8Activation(const InputT* input, const T* bias, T* output, const int* experts, int rows, int inter_size, int fusion, float alpha, float beta, float limit, cudaStream_t stream); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched.cuh index 91a686a8e2ce6..25b38c091feff 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched.cuh @@ -12,6 +12,11 @@ namespace onnxruntime { namespace contrib { namespace cuda { +template +bool TryMatMulUnalignedBatched4Bits( + T* output, const T* a_data, const uint8_t* b_data_quant, const T* scales_data, + int m, int n, int k, cudaStream_t stream); + template bool TryMatMulBatched4Bits( T* output, diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched_half.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched_half.cu index 892d43b0520da..75ef6af0f1123 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched_half.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched_half.cu @@ -1,12 +1,42 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "contrib_ops/cuda/quantization/matmul_4bits_batched_impl.cuh" namespace onnxruntime { namespace contrib { namespace cuda { +template +bool TryMatMulUnalignedBatched4Bits( + T* output, const T* a_data, const uint8_t* b_data_quant, const T* scales_data, + int m, int n, int k, cudaStream_t stream) { + static_assert(std::is_same_v); + if (m < 2 || m > 8 || n < 4 || n % 8 != 4 || k <= 0 || k % 32 != 0) { + return false; + } + const int cta_m = m <= 2 ? 2 : m <= 4 ? 4 + : 8; + dim3 threads(onnxruntime::cuda::GPU_WARP_SIZE_HOST, kColsPerThreadBlock); + dim3 blocks(n / kColsPerThreadBlock + 1, (m + cta_m - 1) / cta_m); + if (cta_m == 2) { + MatMulFloat4BatchedKernel<<>>( + output, a_data, b_data_quant, scales_data, nullptr, m, n, k, k / 32); + } else if (cta_m == 4) { + MatMulFloat4BatchedKernel<<>>( + output, a_data, b_data_quant, scales_data, nullptr, m, n, k, k / 32); + } else { + MatMulFloat4BatchedKernel<<>>( + output, a_data, b_data_quant, scales_data, nullptr, m, n, k, k / 32); + } + return true; +} + +template bool TryMatMulUnalignedBatched4Bits( + half*, const half*, const uint8_t*, const half*, int, int, int, cudaStream_t); + template bool TryMatMulBatched4Bits( half*, const half*, const uint8_t*, const half*, const uint8_t*, int, int, int, int, size_t, cudaStream_t); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched_impl.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched_impl.cuh index e5852b4f80352..d6f5abfe134c2 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched_impl.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_batched_impl.cuh @@ -174,7 +174,7 @@ __device__ __forceinline__ float HorizontalAdd(__nv_bfloat162 acc) { #endif } -template +template __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock, 3) MatMulFloat4BatchedKernel( T* output, const T* a_data, @@ -189,6 +189,13 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock, 3) MatMulFloat const int lane_id = threadIdx.x; const int warp_id = WarpUniform(threadIdx.y); const int col_base = (blockIdx.x * kColsPerThreadBlock + warp_id) * CtaN; + if constexpr (TailN) { + static_assert(CtaN == 1 && !has_zero_point, "TailN needs one scale-only output per warp"); + // This register-only kernel has no CTA barriers; the whole invalid warp exits before pointers. + if (col_base >= n) { + return; + } + } const int m_base = blockIdx.y * CtaM; const int valid = m - m_base; constexpr int k_per_iter = kWarpSize * kElementsPerThreadPerIteration; diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_half.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_half.cu index a45a889993613..f2748252e5fad 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_half.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_half.cu @@ -1,12 +1,41 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + +#include "contrib_ops/cuda/quantization/matmul_4bits_batched.cuh" #include "contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh" +#include "contrib_ops/cuda/quantization/matmul_nbits.cuh" namespace onnxruntime { namespace contrib { namespace cuda { +template +bool TryMatMulUnaligned4Bits( + T* output, const T* a_data, const uint8_t* b_data_quant, const T* scales_data, + int m, int n, int k, size_t shared_mem_per_block, cudaStream_t stream) { + static_assert(std::is_same_v); + if (m < 1 || m > 8 || n < 4 || n % 8 != 4 || k <= 0 || k % 32 != 0) { + return false; + } + if (m > 1) { + return TryMatMulUnalignedBatched4Bits(output, a_data, b_data_quant, scales_data, m, n, k, stream); + } + const size_t shared_mem_size = sizeof(T) * (k / 32) * kColsPerThreadBlock; + if (shared_mem_size > shared_mem_per_block) { + return false; + } + dim3 blocks(n / kColsPerThreadBlock + 1, 1); + dim3 threads(onnxruntime::cuda::GPU_WARP_SIZE_HOST, kColsPerThreadBlock); + MatMulFloat4BitsKernelM1<<>>( + output, a_data, b_data_quant, scales_data, nullptr, 1, n, k, k / 32); + return true; +} + +template bool TryMatMulUnaligned4Bits( + half*, const half*, const uint8_t*, const half*, int, int, int, size_t, cudaStream_t); + template bool TryMatMul4BitsM1( half*, const half*, const uint8_t*, const half*, const uint8_t*, int, int, int, size_t, cudaStream_t); } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh index 63c69731196b7..98c8a3435bdf1 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh @@ -14,7 +14,7 @@ namespace onnxruntime { namespace contrib { namespace cuda { -template +template __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloat4BitsKernelM1( T* output, const T* a_data, @@ -37,8 +37,13 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloat4Bi int offset = n_block_id * kColsPerThreadBlock * blocks_per_K; for (int i = warp_id * kWarpSize + lane_id; i < kColsPerThreadBlock * blocks_per_K; i += kColsPerThreadBlock * kWarpSize) { - const int column = n_block_id * kColsPerThreadBlock + i / blocks_per_K; - b_scale_vec[i] = column < n ? scales_data[offset + i] : T{}; + if constexpr (TailN) { + const int remaining_cols = n - n_block_id * kColsPerThreadBlock; + const int valid_cols = remaining_cols < kColsPerThreadBlock ? remaining_cols : kColsPerThreadBlock; + b_scale_vec[i] = i < valid_cols * blocks_per_K ? scales_data[offset + i] : T(0.f); + } else { + b_scale_vec[i] = scales_data[offset + i]; + } } uint8_t* b_zp_vec; @@ -57,6 +62,13 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloat4Bi b_zp_vec += warp_id * b_zp_k * 2; } __syncthreads(); + if constexpr (TailN) { + static_assert(!has_zero_point, "TailN is scale-only"); + // Every thread participates in the scale-loading barrier before an invalid output warp exits. + if (n_id >= n) { + return; + } + } if (n_id >= n) { return; diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index ed45ba2ea928d..5daba70705798 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -5,9 +5,23 @@ #include #include +#include // NOLINT(build/c++17) +#include +#include #include +#include +#include +#include #include +#if defined(_WIN32) +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#include +#endif + #include "core/common/status.h" #include "core/common/float16.h" #include "core/framework/tuning_context.h" @@ -23,6 +37,7 @@ #include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" #include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" #include "contrib_ops/cuda/quantization/matmul_nbits_sm90_validation.h" +#include "contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h" #endif #include "contrib_ops/cuda/llm/common/logger.h" #include "contrib_ops/cpu/quantization/matmul_nbits_helper.h" @@ -117,6 +132,157 @@ using onnxruntime::llm::kernels::weight_only::WeightOnlyGroupwiseQuantGemmPlugin using onnxruntime::llm::kernels::weight_only::WeightTypeId; static GemmPluginProfilerManager s_profilerManager; +namespace { +// Process-global registry of persistent tactic caches, keyed by their resolved file location so that +// identical shapes are tuned once and reused across sessions and nodes that share a location. +struct GlobalTacticCacheRegistry { + std::mutex mutex; + std::unordered_map> caches; +}; + +GlobalTacticCacheRegistry& GetGlobalTacticCacheRegistry() { + // Leaked on purpose: a CUDA EP owned by another static may be destroyed after this library's statics. + static auto* registry = new GlobalTacticCacheRegistry(); + return *registry; +} + +bool IsWindowsDllShutdownInProgress() { +#if defined(_WIN32) + using RtlDllShutdownInProgressFn = BOOLEAN(NTAPI*)(); + static const auto rtl_dll_shutdown_in_progress = reinterpret_cast( + ::GetProcAddress(::GetModuleHandleW(L"ntdll.dll"), "RtlDllShutdownInProgress")); + return rtl_dll_shutdown_in_progress != nullptr && rtl_dll_shutdown_in_progress() != 0; +#else + return false; +#endif +} + +void FlushAllTacticCaches(bool log_failures, bool allow_blocking) { + if (IsWindowsDllShutdownInProgress()) { + return; + } + try { + auto& registry = GetGlobalTacticCacheRegistry(); + std::vector> caches; + { + std::unique_lock lock(registry.mutex, std::defer_lock); + if (allow_blocking) { + lock.lock(); + } else if (!lock.try_lock()) { + return; + } + caches.reserve(registry.caches.size()); + for (const auto& entry : registry.caches) { + caches.push_back(entry.second); + } + } + // Flush outside the registry lock so file I/O never blocks session creation on another thread. + for (const auto& cache : caches) { + if (cache == nullptr) { + continue; + } + if (IsWindowsDllShutdownInProgress()) { + return; + } + try { + auto status = allow_blocking ? cache->Flush() : cache->TryFlush(); + if (!status.IsOK() && log_failures) { + ORT_LLM_LOG_WARNING("Failed to flush MatMulNBits gemm tactic cache: " + status.ErrorMessage()); + } + } catch (...) { + // Swallow: cache persistence is best-effort and must not escape teardown. + } + } + } catch (...) { + // Registry snapshot allocation must not escape teardown either. + } +} + +void OnProcessExitFlushTacticCaches() { + FlushAllTacticCaches(/*log_failures*/ false, /*allow_blocking*/ false); +} + +void RegisterProcessExitTacticCacheFlusher() { +#if defined(_WIN32) + // Static destructors in a Windows DLL run during DLL_PROCESS_DETACH, after ExitProcess has + // already stopped worker threads (which may still hold a registry or cache mutex). Register a + // process-wide CRT exit callback instead so live sessions save during exit() before worker + // threads are stopped, and pin the DLL so the callback address stays valid if the provider is + // unloaded earlier. + static std::once_flag once; + std::call_once(once, []() { + HMODULE module = nullptr; + if (::GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_PIN, + reinterpret_cast(&OnProcessExitFlushTacticCaches), &module)) { + _crt_atexit(&OnProcessExitFlushTacticCaches); + } + }); +#endif +} + +// Saves tactics of sessions still alive at process exit / module unload without blocking during +// Windows DLL shutdown. Does not log: the default logger may be gone. +struct ExitTacticCacheFlusher { + ~ExitTacticCacheFlusher() { FlushAllTacticCaches(/*log_failures*/ false, /*allow_blocking*/ false); } +}; +ExitTacticCacheFlusher s_exit_tactic_cache_flusher; +} // namespace + +// Returns the process-global cache for `file_path` (creating, loading, and registering it on first +// use). Returns nullptr when the location is already bound to a different GPU's signature (e.g. one +// explicit prefix shared by heterogeneous devices). When `load` is set, an existing cache is re-read +// if the file changed since it was last loaded or written, so each new session starts from the +// latest file. Tactics are only staged in memory while sessions run; the disk write happens in +// FlushMatMulNBitsTacticCaches() when a session closes. +static std::shared_ptr GetGlobalMatMulNBitsTacticCache( + std::string file_path, onnxruntime::llm::gemm_cache::HardwareSignature signature, bool load) { + using onnxruntime::llm::gemm_cache::MatMulNBitsTacticCache; + auto& registry = GetGlobalTacticCacheRegistry(); + std::lock_guard lock(registry.mutex); + auto it = registry.caches.find(file_path); + if (it != registry.caches.end()) { + if (!it->second->Signature().StrictMatches(signature)) { + ORT_LLM_LOG_WARNING("fpA_intB tactic cache " + file_path + + " is already used by a different GPU in this process; persistence is disabled for " + + signature.device_name); + return nullptr; + } + if (load) { + static_cast(it->second->ReloadIfChanged()); + } + return it->second; + } + + RegisterProcessExitTacticCacheFlusher(); + auto cache = std::make_shared(file_path, std::move(signature)); + auto status = cache->Load(); + static_cast(status); // A missing/mismatched file simply yields an empty cache. + registry.caches.emplace(std::move(file_path), cache); + return cache; +} + +// Flushes every registered tactic cache to disk when a session closes. This is the only place tuned +// tactics reach disk, which keeps file I/O off session creation and inference. Best-effort and +// dirty-guarded (Flush() is a no-op when nothing new was staged), so calling it once per CUDA EP +// teardown is cheap even when several sessions share the process. Safe to call from a destructor: +// never throws. +void FlushMatMulNBitsTacticCaches() { + FlushAllTacticCaches(/*log_failures*/ true, /*allow_blocking*/ true); +} + +void FlushMatMulNBitsTacticCachesForShutdown() { + FlushAllTacticCaches(/*log_failures*/ false, /*allow_blocking*/ false); +} + +std::shared_ptr GetGlobalMatMulNBitsTacticCacheForTesting( + std::string file_path, onnxruntime::llm::gemm_cache::HardwareSignature signature, bool load) { + return GetGlobalMatMulNBitsTacticCache(std::move(file_path), std::move(signature), load); +} + +std::unique_lock LockMatMulNBitsTacticCacheRegistryForTesting() { + return std::unique_lock(GetGlobalTacticCacheRegistry().mutex); +} + constexpr auto kScaleAndZeros = cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS; constexpr auto kScaleOnly = cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY; @@ -523,6 +689,41 @@ void ValidateSm90PrepackedWeightSupport(int sm, int64_t block_size) { template void MatMulNBits::InitGemmProfiler(int sm) { + // Resolve the persistent tactic cache from session config (falls back to env vars). + namespace gc = onnxruntime::llm::gemm_cache; + const auto& info = this->Info(); + const std::string mode = ResolveFpAIntBConfigOrEnv(info, gc::kSessionConfigCacheMode, gc::kEnvCacheMode); + const std::optional parsed_access = gc::ParseCacheAccess(mode); + ORT_ENFORCE(parsed_access.has_value(), "Invalid ", gc::kSessionConfigCacheMode, " value '", mode, + "'. Expected load_save, load, or save."); + const gc::CacheAccess cache_access = *parsed_access; + const std::string to_model = ResolveFpAIntBConfigOrEnv(info, gc::kSessionConfigCacheToModel, gc::kEnvCacheToModel); + const std::optional use_sidecar = gc::ParseCacheFlag(to_model); + ORT_ENFORCE(use_sidecar.has_value(), "Invalid ", gc::kSessionConfigCacheToModel, " value '", to_model, + "'. Expected 0 or 1."); + std::string sidecar_prefix; + if (*use_sidecar) { + // Empty for models loaded from bytes, in which case the other locations still apply. + std::filesystem::path model_path = this->GetSessionModelPath(); + if (!model_path.empty() && model_path.is_relative()) { + std::error_code ec; + std::filesystem::path abs_path = std::filesystem::absolute(model_path, ec); + if (!ec && !abs_path.empty()) { + model_path = std::move(abs_path); + } + } + const auto model_path_u8 = model_path.u8string(); + sidecar_prefix.assign(model_path_u8.begin(), model_path_u8.end()); + } + + const auto& device_prop = this->GetDeviceProp(); + const gc::HardwareSignature cache_signature = gc::HardwareSignature::FromDevice( + device_prop.name, device_prop.major * 10 + device_prop.minor, device_prop.multiProcessorCount); + std::string cache_file_path = gc::MatMulNBitsTacticCache::ResolveFilePath( + info.GetConfigOptions().GetConfigOrDefault(gc::kSessionConfigCacheDir, ""), + info.GetConfigOptions().GetConfigOrDefault(gc::kSessionConfigCachePrefix, ""), cache_signature, + sidecar_prefix); + gemmProfiler_ = s_profilerManager.createGemmPluginProfiler(/*inference*/ false); using onnxruntime::llm::kernels::fpA_intB_gemv::KernelType; @@ -633,6 +834,23 @@ void MatMulNBits::InitGemmProfiler(int sm) { gemmProfiler_->setGroupSize(static_cast(block_size_)); gemmProfiler_->setDecodeInterleave(weight_prepacked_ == kMatMulNBitsWeightPrepackedSm90 ? 1 : 4); + if (!cache_file_path.empty()) { + // Persisted sessions use a session-local profile map so they deduplicate identical shapes + // within the session while loading new or retuned tactics from the persistent cache rather + // than reusing stale entries from the process-global map. + gemmProfiler_->setPersistentCache( + GetGlobalMatMulNBitsTacticCache(std::move(cache_file_path), cache_signature, cache_access.load), + cache_access, + s_profilerManager.getSessionProfileMap(this->GetSessionKey())); + } + const std::string tuning_results_prefix = + info.GetConfigOptions().GetConfigOrDefault(gc::kSessionConfigTuningResultsPrefix, ""); + if (!tuning_results_prefix.empty()) { + gemmProfiler_->setTuningResultsCache(GetGlobalMatMulNBitsTacticCache( + gc::MatMulNBitsTacticCache::ResolveFilePath("", tuning_results_prefix, cache_signature), cache_signature, + /*load*/ false)); + } + auto allocator = this->Info().GetAllocator(OrtMemType::OrtMemTypeDefault); gemmProfiler_->setAllocator(allocator); } @@ -645,19 +863,11 @@ void MatMulNBits::RunGemmProfile(bool hasWeightOnlyCudaKernel, int min_m, int // Include both the packing/kernel SM and physical device SM so compatible weight layouts can // still select different tactics on different GPU generations. const int kernel_sm = FpAIntBPackingSmForKernel(); - // The tag keeps profiled results with and without the optional paired-K tactic apart. - const int tactic_set_tag = paired_gemv_mode_; - if constexpr (std::is_same_v) { - gemmId_ = GemmIdCore(n_16b, static_cast(K_), onnxruntime::llm::nvinfer::DataType::kHALF, kernel_sm, - sm_, wave_aware_gemv_, tactic_set_tag, - this->GetDeviceId(), static_cast(nbits_), static_cast(block_size_), - has_bias_, has_zero_points_); - } else if constexpr (std::is_same_v) { - gemmId_ = GemmIdCore(n_16b, static_cast(K_), onnxruntime::llm::nvinfer::DataType::kBF16, kernel_sm, - sm_, wave_aware_gemv_, tactic_set_tag, - this->GetDeviceId(), static_cast(nbits_), static_cast(block_size_), - has_bias_, has_zero_points_); - } + const auto dtype = std::is_same_v ? onnxruntime::llm::nvinfer::DataType::kBF16 + : onnxruntime::llm::nvinfer::DataType::kHALF; + gemmId_ = GemmIdCore(n_16b, static_cast(K_), dtype, kernel_sm, static_cast(nbits_), + static_cast(block_size_), has_zero_points_, hasWeightOnlyCudaKernel, has_bias_, + this->GetDeviceProp().name, wave_aware_gemv_, paired_gemv_mode_, this->GetDeviceId()); GemmDims dims = {min_m, max_m, n_16b, K_}; gemmProfiler_->profileTactics(weightOnlyGemmRunner_, gemmId_.dtype, dims, gemmId_, hasWeightOnlyCudaKernel); @@ -946,7 +1156,6 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { int m = SafeInt(helper.M()); int n = SafeInt(helper.N()); int k = SafeInt(helper.K()); - DUMP_TENSOR_INIT(); #if USE_FPA_INTB_GEMM @@ -1104,6 +1313,28 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { } #endif +#if USE_FPA_INTB_GEMM + const bool unaligned4_has_fpA_intB = has_fpA_intB_gemm_; +#else + constexpr bool unaligned4_has_fpA_intB = false; +#endif + if constexpr (std::is_same_v) { + if (IsUnaligned4FastPathEligible( + unaligned4_fast_path_, true, m, n, k, static_cast(nbits_), static_cast(block_size_), + column_wise_quant_blk_, zero_points != nullptr, bias != nullptr, reorder_idx != nullptr, + unaligned4_has_fpA_intB, force_chunked_) && + TryMatMulUnaligned4Bits( + reinterpret_cast(Y->MutableData()), reinterpret_cast(a_data), + blob_data, reinterpret_cast(scales_data), m, n, k, + GetDeviceProp().sharedMemPerBlock, stream)) { + if (unaligned4_log_.ShouldLog(m == 1)) { + CUDA_RETURN_IF_ERROR(cudaPeekAtLastError()); + unaligned4_log_.Log(m == 1, this->Node().Name(), m, k, n); + } + return Status::OK(); + } + } + if ((reorder_idx_data == nullptr) && (!zero_points || !zero_points->IsDataType())) { // First, try the fused fast path. It handles bias only for the GPT-OSS router GEMV // specialization; for other shapes it fails when bias is present. diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh index 00eaef114ce9d..eab444df8a2d2 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh @@ -31,6 +31,11 @@ constexpr bool IsMatMul8BitsSmallM( block_size == 64; } +template +bool TryMatMulUnaligned4Bits( + T* output, const T* a_data, const uint8_t* b_data_quant, const T* scales_data, + int m, int n, int k, size_t shared_mem_per_block, cudaStream_t stream); + template bool TryMatMul2Bits( T* output, diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index d1018d637cfab..f04fae522ac37 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -25,6 +25,7 @@ #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" #include "contrib_ops/cuda/quantization/matmul_nbits_workspace_estimate.h" #include "contrib_ops/cuda/quantization/matmul_nbits_sm90_validation.h" +#include "contrib_ops/cuda/quantization/matmul_nbits_unaligned4.h" #include "core/platform/env_var_utils.h" namespace onnxruntime { @@ -303,6 +304,8 @@ class MatMulNBits final : public CudaKernel { if (m_chunk_size_ > 0) { const int64_t max_unchunked_m = force_chunked_ ? m_chunk_size_ : std::max(m_chunk_size_, m_chunk_min_rows_); + gemmProfiler_->setMaxProfileM(static_cast( + std::min(max_unchunked_m, onnxruntime::llm::kernels::weight_only::kMaxProfileM))); max_m = static_cast(std::min(max_m, max_unchunked_m)); } RunGemmProfile(has_fpA_intB_gemv_, 1, max_m); @@ -409,6 +412,9 @@ class MatMulNBits final : public CudaKernel { bool is_zero_points_scale_same_type_{false}; bool force_chunked_{false}; int64_t chunk_target_rows_{kDefaultChunkTargetRows}; + bool unaligned4_fast_path_{ + ParseUnaligned4FastPathEnabled(ParseEnvironmentVariableWithDefault(kUnaligned4FastPathEnvVar, ""))}; + Unaligned4DispatchLog unaligned4_log_; #if USE_FPA_INTB_GEMM bool has_fpA_intB_gemv_{false}; diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h new file mode 100644 index 0000000000000..e59deb39341e3 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#if USE_FPA_INTB_GEMM && !defined(DISABLE_CONTRIB_OPS) && !defined(USE_CUDA_MINIMAL) +#include "contrib_ops/cuda/llm/gemm_tactic_cache.h" +#endif + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// Flushes all process-global MatMulNBits fpA_intB tactic caches to disk (best-effort, dirty-guarded). +// +// During a session, lazily profiled tactics are only staged into the in-memory caches (no disk I/O); +// this function performs the disk write. It is invoked at CUDA EP teardown: ~CUDAExecutionProvider for +// the built-in EP and ~CudaEp for the CUDA plugin EP, as well as provider shutdown before worker +// threads are stopped. In builds without onnxruntime_USE_FPA_INTB_GEMM there is no tactic cache and +// this is a no-op. +#if USE_FPA_INTB_GEMM && !defined(DISABLE_CONTRIB_OPS) && !defined(USE_CUDA_MINIMAL) +void FlushMatMulNBitsTacticCaches(); +void FlushMatMulNBitsTacticCachesForShutdown(); +std::shared_ptr GetGlobalMatMulNBitsTacticCacheForTesting( + std::string file_path, onnxruntime::llm::gemm_cache::HardwareSignature signature, bool load); +std::unique_lock LockMatMulNBitsTacticCacheRegistryForTesting(); +#else +inline void FlushMatMulNBitsTacticCaches() {} +inline void FlushMatMulNBitsTacticCachesForShutdown() {} +#endif + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_unaligned4.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_unaligned4.h new file mode 100644 index 0000000000000..9ac043fa5549a --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_unaligned4.h @@ -0,0 +1,59 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include + +#include "core/common/parse_string.h" +#include "core/common/string_utils.h" + +namespace onnxruntime::contrib::cuda { + +constexpr const char* kUnaligned4FastPathEnvVar = "ORT_MATMULNBITS_UNALIGNED4_FAST_PATH"; + +// The kernels support M <= 8, but for M in [5, 8] the register-batched tail kernel measured slower than the +// dequantize + GEMM fallback on H200 (N=324, K=10240: ~39-43 us vs ~26 us), so dispatch stops at M = 4. +constexpr int kUnaligned4MaxM = 4; + +inline bool ParseUnaligned4FastPathEnabled(const std::string& value) { + const auto trimmed = onnxruntime::utils::TrimString(value); + int enabled = 0; + ORT_ENFORCE(trimmed.empty() || TryParseStringWithClassicLocale(trimmed, enabled), + "Invalid ", kUnaligned4FastPathEnvVar, ": expected 0 or 1."); + ORT_ENFORCE(enabled == 0 || enabled == 1, "Invalid ", kUnaligned4FastPathEnvVar, ": expected 0 or 1."); + return enabled == 1; +} + +inline bool IsUnaligned4FastPathEligible( + bool enabled, bool fp16, int m, int n, int k, int bits, int block_size, + bool column_wise, bool zero_points, bool bias, bool g_idx, bool optimized, + bool force_chunked) { + return enabled && fp16 && m >= 1 && m <= kUnaligned4MaxM && n >= 4 && n % 8 == 4 && + k > 0 && k % 32 == 0 && bits == 4 && block_size == 32 && column_wise && + !zero_points && !bias && !g_idx && !optimized && !force_chunked; +} + +// The explicitly enabled selector emits one actual successful dispatch receipt per kernel branch. +class Unaligned4DispatchLog { + public: + bool ShouldLog(bool m1) const { + return !(m1 ? m1_logged_ : batched_logged_).load(std::memory_order_relaxed); + } + void Log(bool m1, const std::string& node, int m, int k, int n) const { + auto& logged = m1 ? m1_logged_ : batched_logged_; + bool expected = false; + if (logged.compare_exchange_strong(expected, true, std::memory_order_relaxed)) { + std::fprintf(stderr, + "[MatMulNBitsUnaligned4] node=\"%s\" M=%d K=%d N=%d type=FP16 bits=4 branch=%s\n", + node.c_str(), m, k, n, m1 ? "M1_SHARED_TAIL" : "BATCHED_REGISTER_TAIL"); + } + } + + private: + mutable std::atomic m1_logged_{false}; + mutable std::atomic batched_logged_{false}; +}; + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc index daf1305969249..8e9c6fde2ad20 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -163,6 +163,7 @@ PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo template Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context) const { + bool reuse = false; const Tensor* mode = context->InputCount() > psai::kIndexShareMode ? context->Input(psai::kIndexShareMode) : nullptr; @@ -181,14 +182,16 @@ Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context const int64_t merged_rows = merged_indices->Shape()[0]; ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kMergedCounts), "merged_counts", {merged_rows})); ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kMergedStatus), "merged_status", {merged_rows})); - if (value == 1) return ComputeReuse(context); - ORT_RETURN_IF(merged_rows != 0, "Refresh requires zero merge rows"); - const Tensor* query = context->Input(psai::kQuery); - ORT_RETURN_IF(query == nullptr || query->Shape().NumDimensions() != 2, "Refresh requires rank-2 queries"); - Tensor* status = context->Output(psai::kIndexShareStatus, TensorShape({query->Shape()[0]})); - ORT_RETURN_IF(status == nullptr, "IndexShare status output is required"); - if (status->SizeInBytes() != 0) { - CUDA_RETURN_IF_ERROR(cudaMemsetAsync(status->MutableDataRaw(), 0, status->SizeInBytes(), Stream(context))); + reuse = value == 1; + if (!reuse) { + ORT_RETURN_IF(merged_rows != 0, "Refresh requires zero merge rows"); + const Tensor* query = context->Input(psai::kQuery); + ORT_RETURN_IF(query == nullptr || query->Shape().NumDimensions() != 2, "Refresh requires rank-2 queries"); + Tensor* status = context->Output(psai::kIndexShareStatus, TensorShape({query->Shape()[0]})); + ORT_RETURN_IF(status == nullptr, "IndexShare status output is required"); + if (status->SizeInBytes() != 0) { + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(status->MutableDataRaw(), 0, status->SizeInBytes(), Stream(context))); + } } } else { ORT_RETURN_IF(output_capacity_ != 0, "max_output_entries requires the IndexShare interface"); @@ -224,6 +227,7 @@ Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context "PackedSparseAttentionIndexer: state update capture is only valid for policy_mode 'qsa'"); } + if (reuse) return ComputeReuse(context); return is_qsa ? ComputeQsa(context) : ComputeCsa(context); } @@ -374,7 +378,6 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con params.state_capacity = static_cast(state_capacity); params.buffer_capacity = static_cast(buffer_capacity); params.capacity = static_cast(capacity); - params.has_position_ids = position_ids != nullptr; params.output_capacity = static_cast(output_capacity); params.epsilon = epsilon_; params.scale = has_scale_ ? scale_ : 1.0f / std::sqrt(static_cast(head_size)); @@ -539,7 +542,6 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, con params.state_capacity = static_cast(state_capacity); params.buffer_capacity = static_cast(buffer_capacity); params.capacity = static_cast(capacity); - params.has_position_ids = true; params.output_capacity = static_cast(output_capacity); params.epsilon = epsilon_; params.scale = has_scale_ ? scale_ : 1.0f / std::sqrt(static_cast(head_size)); @@ -552,6 +554,29 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, con auto overflow_flags = GetScratchBuffer(static_cast(std::max(batch_size, 1)), GetComputeStream(context)); + Tensor* reuse_status = nullptr; + if (state_only) { + reuse_status = context->Output(psai::kIndexShareStatus, TensorShape({batch_size})); + ORT_RETURN_IF(reuse_status == nullptr, "IndexShare status output is required"); + for (const auto& pair : {std::pair{psai::kMergedIndices, selected_indices}, + std::pair{psai::kMergedCounts, selected_counts}, + std::pair{psai::kMergedStatus, reuse_status}}) { + const Tensor* input = context->Input(pair.first); + if (input->SizeInBytes() != 0 && input->DataRaw() != pair.second->MutableDataRaw()) { + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(pair.second->MutableDataRaw(), input->DataRaw(), input->SizeInBytes(), + cudaMemcpyDeviceToDevice, Stream(context))); + } + } + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(overflow_flags.get(), 0, std::max(batch_size, 1) * sizeof(int32_t), + Stream(context))); + ORT_RETURN_IF_ERROR(LaunchValidatePackedIndexerReuse( + Stream(context), cumulative_sequence_lengths->Data(), past_sequence_lengths->Data(), + past_state_lengths->Data(), selected_indices->MutableData(), + selected_counts->MutableData(), reuse_status->MutableData(), + static_cast(batch_size), static_cast(output_capacity), static_cast(state_capacity), + static_cast(compress_ratio_), false, overflow_flags.get())); + } + ORT_RETURN_IF_ERROR(LaunchCsaPackedSparseAttentionIndexer( Stream(context), params, reinterpret_cast(query->Data()), @@ -579,12 +604,12 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, con float_workspace.get(), overflow_flags.get())); if (!state_only) return Status::OK(); - const Tensor* merged_status = context->Input(psai::kMergedStatus); - Tensor* reuse_status = context->Output(psai::kIndexShareStatus, merged_status->Shape()); return LaunchValidatePackedIndexerReuse(Stream(context), cumulative_sequence_lengths->Data(), + past_sequence_lengths->Data(), past_state_lengths->Data(), selected_indices->MutableData(), selected_counts->MutableData(), reuse_status->MutableData(), static_cast(batch_size), - static_cast(output_capacity), overflow_flags.get()); + static_cast(output_capacity), static_cast(state_capacity), + static_cast(compress_ratio_), false, overflow_flags.get()); } template @@ -615,6 +640,8 @@ Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) c const Tensor* key = context->Input(psai::kKey); const int64_t query_width = query->Shape()[1] - (is_qsa && key == nullptr ? head_size : 0); ORT_RETURN_IF(query_width <= 0 || query_width % head_size != 0, "Reuse query width must contain complete heads"); + const int64_t num_heads = query_width / head_size; + ORT_RETURN_IF_ERROR(CheckIntDimension("num_heads", num_heads, false)); ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kKeyNormWeight), "key_norm_weight", {head_size})); const int64_t width = is_qsa ? head_size : 2 * head_size; ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kPastKeyState), "past_key_state", {batch, state_capacity_, head_size})); @@ -624,20 +651,34 @@ Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) c if (!is_qsa) { ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kPastGateBuffer), "past_gate_buffer", {batch, psai::GenericBufferCapacity(compress_ratio_), width})); + ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kGate), "gate", {advance ? batch : 0, width})); + ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kPositionBias), "position_bias", {compress_ratio_, width})); + ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kHeadWeights), "head_weights", {0, num_heads})); + ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kPositionIds), "position_ids", {advance ? batch : 0})); } else { - for (int index : {psai::kGate, psai::kPositionBias, psai::kHeadWeights, psai::kPastGateBuffer, psai::kPositionIds}) { - ORT_RETURN_IF(context->Input(index) != nullptr, "QSA reuse does not accept CSA-only inputs"); + Tensor* present_gate_buffer = context->Output(psai::kPresentGateBuffer, TensorShape({0})); + ORT_RETURN_IF(present_gate_buffer != nullptr, "QSA reuse must omit present_gate_buffer"); + const Tensor* capture_count = context->Input(psai::kStateUpdateCaptureCount); + ORT_RETURN_IF(state_update_capacity_ > 0 && capture_count == nullptr, "Reuse requires state_update_capture_count"); + if (capture_count != nullptr) { + ORT_RETURN_IF_ERROR(CheckShape(capture_count, "state_update_capture_count", {batch})); } - if (state_update_capacity_ > 0) { - ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kStateUpdateCaptureCount), "state_update_capture_count", {batch})); + const Tensor* active = context->Input(psai::kStateUpdateActive); + if (active != nullptr) { + ORT_RETURN_IF_ERROR(CheckShape(active, "state_update_active", {1})); } } - for (int index : {psai::kKey, psai::kGate, psai::kHeadWeights}) { - const Tensor* input = context->Input(index); - const int64_t expected_rows = advance && index != psai::kHeadWeights ? batch : 0; - ORT_RETURN_IF(input != nullptr && (input->Shape().NumDimensions() != 2 || input->Shape()[0] != expected_rows), - "Reuse projection rows do not match the state policy"); + if (key != nullptr || !is_qsa) { + ORT_RETURN_IF_ERROR(CheckShape(key, "key", {advance ? batch : 0, width})); } + RotaryCacheShape rotary; + ORT_RETURN_IF_ERROR(CheckRotaryCache(context->Input(psai::kCosCache), + context->Input(psai::kSinCache), batch, rotary)); + ORT_RETURN_IF(rotary.rotary_width <= 0 || + (is_qsa ? rotary.rotary_width % 2 != 0 || rotary.rotary_width > head_size + : rotary.rotary_width > head_size / 2), + "Reuse rotary cache width is incompatible with head_size"); + if (advance) return ComputeCsa(context, true); for (const auto& pair : {std::pair{psai::kMergedIndices, psai::kSelectedIndices}, std::pair{psai::kMergedCounts, psai::kSelectedCounts}, std::pair{psai::kMergedStatus, psai::kIndexShareStatus}, @@ -646,7 +687,6 @@ Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) c std::pair{psai::kPastStateLengths, psai::kPresentStateLengths}, std::pair{psai::kPastGateBuffer, psai::kPresentGateBuffer}}) { if (is_qsa && pair.first == psai::kPastGateBuffer) continue; - if (advance && pair.first >= psai::kPastKeyState && pair.first <= psai::kPastStateLengths) continue; const Tensor* input = context->Input(pair.first); ORT_RETURN_IF(input == nullptr, "Reuse state and selection inputs are required"); Tensor* output = context->Output(pair.second, input->Shape()); @@ -656,7 +696,6 @@ Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) c cudaMemcpyDeviceToDevice, Stream(context))); } } - if (advance) return ComputeCsa(context, true); if (is_qsa) { const Tensor* qsa_norm = context->Input(psai::kQueryNormWeight); ORT_RETURN_IF(qsa_norm == nullptr || qsa_norm->Shape().NumDimensions() != 1, "Reuse requires query norm weights"); @@ -667,10 +706,13 @@ Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) c } return LaunchValidatePackedIndexerReuse(Stream(context), context->Input(psai::kCumulativeSequenceLengths)->Data(), + lengths->Data(), + context->Input(psai::kPastStateLengths)->Data(), context->Output(psai::kSelectedIndices, indices->Shape())->MutableData(), context->Output(psai::kSelectedCounts, counts->Shape())->MutableData(), context->Output(psai::kIndexShareStatus, status->Shape())->MutableData(), - static_cast(batch), static_cast(output_capacity_)); + static_cast(batch), static_cast(output_capacity_), + static_cast(state_capacity_), static_cast(compress_ratio_), is_qsa); } template class PackedSparseAttentionIndexer; diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index 2c57dd3256b7e..6efc9be09c9c5 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -65,6 +65,16 @@ constexpr int kQwenHeadSize = 128; constexpr int kQwenCompressRatio = 4; constexpr int kQwenNumHeads = 4; +template +Status ConfigureDynamicSharedMemory(Kernel kernel, size_t bytes, const char* kernel_name) { + int device = 0; + int max_bytes = 0; + CUDA_RETURN_IF_ERROR(cudaGetDevice(&device)); + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&max_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + ORT_RETURN_IF_ERROR(psai::ValidateDynamicSharedMemory(bytes, static_cast(max_bytes), kernel_name)); + return CUDA_CALL(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_bytes)); +} + bool UseHierarchicalQsaTopK(const PackedSparseAttentionIndexerParams& params) { return params.total_tokens > 0 && params.total_tokens <= kDistributedTopKMaxRows && params.state_capacity >= kDistributedTopKMinBlocks && @@ -573,6 +583,108 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* } } +template +__global__ void QsaPairedScoreTileTopKKernel(const T* present_key_state, const float* query_rotated, + const int32_t* cumulative_sequence_lengths, + const int32_t* past_sequence_lengths, const int64_t* position_ids, + const int32_t* present_state_lengths, const int32_t* overflow_flags, + uint64_t* tile_keys, int tile_count, + PackedSparseAttentionIndexerParams params) { + __shared__ uint64_t keys[2][kHierarchicalTileBlocks]; + __shared__ int block_counts[2]; + const int first_token = static_cast(blockIdx.y) * 2; + const int warp = static_cast(threadIdx.x) / kWarpSize; + const int lane = static_cast(threadIdx.x) % kWarpSize; + if (threadIdx.x < 2) { + const int metadata_token = first_token + static_cast(threadIdx.x); + int count = 0; + if (metadata_token < params.total_tokens && overflow_flags[0] == 0) { + const int64_t position = params.has_position_ids ? position_ids[metadata_token] + : static_cast(past_sequence_lengths[0]) + metadata_token; + const int64_t causal = SaiCausalThreshold(position, kQwenCompressRatio); + const int64_t visible = min(causal, static_cast(present_state_lengths[psai::kKeyStateLength])); + count = static_cast(max(static_cast(0), visible)); + } + block_counts[threadIdx.x] = count; + } + __syncthreads(); + float4 query_values[2][kQwenNumHeads]; +#pragma unroll + for (int query_index = 0; query_index < 2; ++query_index) { +#pragma unroll + for (int head = 0; head < kQwenNumHeads; ++head) { + const int token = first_token + query_index; + query_values[query_index][head] = token < params.total_tokens + ? LoadQsaVector4(query_rotated + (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize + lane * 4) + : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + } + } + const int key_stride = tile_count * kHierarchicalTileBlocks; + for (int query_index = 0; query_index < 2; ++query_index) { + if (first_token + query_index < params.total_tokens) { + const int active = (block_counts[query_index] + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks; + const int64_t row_base = static_cast(first_token + query_index) * key_stride; + for (int index = active * kHierarchicalTileBlocks + blockIdx.x * blockDim.x + threadIdx.x; + index < key_stride; index += gridDim.x * blockDim.x) { + tile_keys[row_base + index] = topk::kPaddingSortKey; + } + } + } + const int visible = max(block_counts[0], block_counts[1]); + const int active_tiles = (visible + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks; + for (int tile = static_cast(blockIdx.x); tile < active_tiles; tile += static_cast(gridDim.x)) { + const int block = tile * kHierarchicalTileBlocks + warp; + const float4 key_value = block < visible + ? LoadQsaVector4(present_key_state + static_cast(block) * kQwenHeadSize + lane * 4) + : make_float4(0.0f, 0.0f, 0.0f, 0.0f); +#pragma unroll + for (int query_index = 0; query_index < 2; ++query_index) { + float total = 0.0f; +#pragma unroll + for (int head = 0; head < kQwenNumHeads; ++head) { + const float4 query_value = query_values[query_index][head]; + float score = query_value.x * key_value.x + query_value.y * key_value.y + + query_value.z * key_value.z + query_value.w * key_value.w; + score = topk::WarpReduceSum(score); + total += fmaxf(score, 0.0f); + } + uint64_t key = topk::kPaddingSortKey; + if (block < block_counts[query_index]) { + const float scaled = total * params.scale; + key = topk::PackStableSortKey(scaled == 0.0f ? 0.0f : scaled, block); + } + if (lane == 0) { + keys[query_index][warp] = key; + } + } + __syncthreads(); + if (threadIdx.x < 2) { + const int query_index = static_cast(threadIdx.x); + for (int index = 1; index < kHierarchicalTileBlocks; ++index) { + const uint64_t key = keys[query_index][index]; + int insertion = index; + while (insertion > 0 && key > keys[query_index][insertion - 1]) { + keys[query_index][insertion] = keys[query_index][insertion - 1]; + --insertion; + } + keys[query_index][insertion] = key; + } + } + __syncthreads(); + if (threadIdx.x < 2 * kHierarchicalTileBlocks) { + const int query_index = static_cast(threadIdx.x) / kHierarchicalTileBlocks; + const int slot = static_cast(threadIdx.x) % kHierarchicalTileBlocks; + if (first_token + query_index < params.total_tokens && + tile * kHierarchicalTileBlocks < block_counts[query_index]) { + const int64_t output_base = static_cast(first_token + query_index) * key_stride + + tile * kHierarchicalTileBlocks; + tile_keys[output_base + slot] = keys[query_index][slot]; + } + } + __syncthreads(); + } +} + __device__ __forceinline__ uint64_t QsaMergeRank(const uint64_t* left, int left_count, const uint64_t* right, int right_count, int rank) { int low = max(0, rank - right_count); @@ -593,6 +705,47 @@ __device__ __forceinline__ uint64_t QsaMergeRank(const uint64_t* left, int left_ return topk::kPaddingSortKey; } +__global__ void QsaCompactTileTopKKernel(const uint64_t* input_keys, uint64_t* output_keys, + int input_stride, int output_stride) { + constexpr int kChunkSize = 2048; + static_assert(kChunkSize % kHierarchicalTileBlocks == 0); + __shared__ uint64_t shared_keys[2][kChunkSize]; + const int chunk_start = static_cast(blockIdx.x) * kChunkSize; + const int64_t input_base = static_cast(blockIdx.y) * input_stride + chunk_start; + const int64_t output_base = static_cast(blockIdx.y) * output_stride + + static_cast(blockIdx.x) * kBoundedTopKMax; + if (input_keys[input_base] == topk::kPaddingSortKey) { + for (int rank = threadIdx.x; rank < kBoundedTopKMax; rank += blockDim.x) { + output_keys[output_base + rank] = topk::kPaddingSortKey; + } + return; + } + for (int index = threadIdx.x; index < kChunkSize; index += blockDim.x) { + shared_keys[0][index] = chunk_start + index < input_stride + ? input_keys[input_base + index] + : topk::kPaddingSortKey; + } + __syncthreads(); + int input_bank = 0; + for (int width = kHierarchicalTileBlocks; width < kChunkSize; width *= 2) { + const int count = min(width, kBoundedTopKMax); + const int output_count = min(count * 2, kBoundedTopKMax); + const int pairs = kChunkSize / (width * 2); + for (int index = threadIdx.x; index < pairs * output_count; index += blockDim.x) { + const int pair_start = (index / output_count) * width * 2; + const int rank = index % output_count; + shared_keys[1 - input_bank][pair_start + rank] = + QsaMergeRank(shared_keys[input_bank] + pair_start, count, + shared_keys[input_bank] + pair_start + width, count, rank); + } + __syncthreads(); + input_bank = 1 - input_bank; + } + for (int rank = threadIdx.x; rank < kBoundedTopKMax; rank += blockDim.x) { + output_keys[output_base + rank] = shared_keys[input_bank][rank]; + } +} + __global__ void QsaMergeTileTopKKernel(const uint64_t* input_keys, uint64_t* output_keys, int key_stride, int list_width, int rows) { const int pairs_per_row = (key_stride + 2 * list_width - 1) / (2 * list_width); @@ -982,6 +1135,7 @@ __global__ void CsaUpdateStateKernel(const T* key, const T* gate, const T* key_n cumulative_sequence_lengths[params.batch_size] != params.total_tokens || req_start < 0 || req_end < req_start || req_end > params.total_tokens || (params.state_only && (req_start != b || req_end != b + 1)) || + (params.state_only && overflow_flags[b] != 0) || past_sequence_lengths[b] < 0 || raw_key_len < 0 || raw_key_len > params.state_capacity || raw_buf_len < 0 || raw_buf_len > params.buffer_capacity; const int req_len = invalid_metadata ? 0 : req_end - req_start; @@ -1282,25 +1436,46 @@ __global__ void CsaSelectKernel(const float* scores, const int32_t* cumulative_s namespace { __global__ void ValidatePackedReuseKernel(const int32_t* cumulative_lengths, int32_t* indices, - int32_t* counts, int32_t* status, int capacity, const int32_t* overflow_flags) { + int32_t* counts, int32_t* status, int capacity, int32_t* overflow_flags, + const int32_t* past_sequence_lengths, const int32_t* past_state_lengths, + int batch_size, int state_capacity, int compress_ratio, bool is_qsa) { const int row = blockIdx.x; - if (cumulative_lengths[row] == row && cumulative_lengths[row + 1] == row + 1 && - (overflow_flags == nullptr || overflow_flags[row] == 0)) return; - for (int column = threadIdx.x; column < capacity; column += blockDim.x) { + __shared__ bool valid; + if (threadIdx.x == 0) { + const int past = past_sequence_lengths[row]; + const int64_t state_offset = static_cast(row) * 2; + const int key_length = past_state_lengths[state_offset]; + const int buffer_length = past_state_lengths[state_offset + 1]; + valid = cumulative_lengths[0] == 0 && cumulative_lengths[batch_size] == batch_size && + cumulative_lengths[row] == row && cumulative_lengths[row + 1] == row + 1 && + status[row] == 0 && counts[row] >= 0 && counts[row] <= capacity && + past >= 0 && key_length >= 0 && key_length <= state_capacity && + buffer_length >= 0 && buffer_length < (is_qsa ? compress_ratio : int64_t{2} * compress_ratio) && + (!is_qsa || (key_length == past / compress_ratio && buffer_length == past % compress_ratio)) && + (overflow_flags == nullptr || overflow_flags[row] == 0); + } + __syncthreads(); + if (valid) return; + for (int64_t column = threadIdx.x; column < capacity; column += blockDim.x) { indices[static_cast(row) * capacity + column] = -1; } if (threadIdx.x == 0) { counts[row] = 0; if (status[row] == 0) status[row] = 1; + if (overflow_flags != nullptr) overflow_flags[row] = 1; } } } // namespace Status LaunchValidatePackedIndexerReuse(cudaStream_t stream, const int32_t* cumulative_lengths, + const int32_t* past_sequence_lengths, const int32_t* past_state_lengths, int32_t* indices, int32_t* counts, int32_t* status, - int batch_size, int capacity, const int32_t* overflow_flags) { + int batch_size, int capacity, int state_capacity, int compress_ratio, + bool is_qsa, int32_t* overflow_flags) { if (batch_size == 0) return Status::OK(); - ValidatePackedReuseKernel<<>>(cumulative_lengths, indices, counts, status, capacity, overflow_flags); + ValidatePackedReuseKernel<<>>( + cumulative_lengths, indices, counts, status, capacity, overflow_flags, past_sequence_lengths, + past_state_lengths, batch_size, state_capacity, compress_ratio, is_qsa); return CUDA_CALL(cudaGetLastError()); } @@ -1354,6 +1529,8 @@ Status LaunchQsaPackedSparseAttentionIndexer( } const size_t value_bytes = static_cast(params.head_size) * sizeof(float); + ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( + QsaUpdateStateKernel, 2 * value_bytes + kThreads * sizeof(float), "QsaUpdateStateKernel")); const int state_blocks = static_cast(std::min(params.batch_size, kSaiMaxGridDimX)); QsaUpdateStateKernel<<>>( key, key_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, past_kv_buffer, @@ -1373,6 +1550,8 @@ Status LaunchQsaPackedSparseAttentionIndexer( const int64_t rotate_rows = static_cast(params.total_tokens) * params.num_heads; const int rotate_blocks = static_cast(std::min(rotate_rows, kSaiMaxGridDimX)); + ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( + PackedRotateQueryKernel, value_bytes + kThreads * sizeof(float), "PackedRotateQueryKernel")); PackedRotateQueryKernel<<>>( query, query_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, position_ids, query_rotated, params); @@ -1491,6 +1670,8 @@ Status LaunchCsaPackedSparseAttentionIndexer( const size_t value_bytes = static_cast(params.head_size) * sizeof(float); const int state_blocks = static_cast(std::min(params.batch_size, kSaiMaxGridDimX)); + ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( + CsaUpdateStateKernel, value_bytes + kThreads * sizeof(float), "CsaUpdateStateKernel")); CsaUpdateStateKernel<<>>( key, gate, key_norm_weight, cos_cache, sin_cache, position_bias, cumulative_sequence_lengths, past_sequence_lengths, past_kv_buffer, past_gate_buffer, past_state_lengths, present_key_state, @@ -1505,6 +1686,8 @@ Status LaunchCsaPackedSparseAttentionIndexer( const int64_t rotate_rows = static_cast(params.total_tokens) * params.num_heads; const int rotate_blocks = static_cast(std::min(rotate_rows, kSaiMaxGridDimX)); + ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( + PackedRotateQueryKernel, value_bytes + kThreads * sizeof(float), "PackedRotateQueryKernel")); PackedRotateQueryKernel<<>>( query, query_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, position_ids, query_rotated, params); diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h index 9e2f812eb450a..ab4dcd4849116 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h @@ -48,8 +48,10 @@ struct PackedSparseAttentionIndexerParams { size_t GetQsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& params); size_t GetCsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& params); Status LaunchValidatePackedIndexerReuse(cudaStream_t stream, const int32_t* cumulative_lengths, + const int32_t* past_sequence_lengths, const int32_t* past_state_lengths, int32_t* indices, int32_t* counts, int32_t* status, - int batch_size, int capacity, const int32_t* overflow_flags = nullptr); + int batch_size, int capacity, int state_capacity, int compress_ratio, + bool is_qsa, int32_t* overflow_flags = nullptr); // `overflow_flags` is a caller-allocated int32 scratch buffer with at least `batch_size` elements // (unused when batch_size == 0). The update kernel writes, per request, whether this call's new // blocks/windows would exceed state_capacity; when it does, the whole step is rejected for that diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge.cc index fd80e2de965e8..00dc75288b395 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge.cc @@ -1,4 +1,5 @@ #include "core/providers/cuda/cuda_kernel.h" +#include "core/common/safeint.h" #include "contrib_ops/cpu/sparse/packed_sparse_attention_indexer_merge_common.h" #include "contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.h" @@ -31,6 +32,8 @@ class PackedSparseAttentionIndexerMerge final : public onnxruntime::cuda::CudaKe additional_counts->Shape() != rows->Shape() || additional->Shape()[1] > INT32_MAX - dimensions.capacity, "Additional selections must have shape [N,A] and counts [N], with bounded capacity"); + dimensions.hash_capacity = indexer_merge::HashCapacity( + capacity_, static_cast(dimensions.capacity) + additional->Shape()[1]); } else { ORT_RETURN_IF(additional || additional_counts, "append_range does not accept additional index inputs"); ORT_RETURN_IF_ERROR(indexer_merge::Validate(base, counts, rows, starts, ends, capacity_, dimensions)); @@ -40,7 +43,7 @@ class PackedSparseAttentionIndexerMerge final : public onnxruntime::cuda::CudaKe auto* status = context->Output(2, {dimensions.queries}); if (dimensions.queries == 0) return Status::OK(); auto workspace = GetScratchBuffer( - static_cast(dimensions.queries) * dimensions.hash_capacity * 2, GetComputeStream(context)); + SafeInt(dimensions.queries) * dimensions.hash_capacity * 2, GetComputeStream(context)); if (append_indices_) { CUDA_RETURN_IF_ERROR(LaunchPackedSparseAttentionIndexerMergeIndices( Stream(context), base->Data(), counts->Data(), rows->Data(), diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.cu index d897f5548d220..5d5fa04c9517c 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.cu @@ -7,13 +7,13 @@ namespace { constexpr int kThreads = 256; -__device__ int HashSlot(int value, int capacity) { - return static_cast((static_cast(value) * 2654435761u) & (capacity - 1)); +__device__ int64_t HashSlot(int value, int64_t capacity) { + return (static_cast(value) * 2654435761u) & (capacity - 1); } -__device__ int FindSlot(const int* keys, int value, int capacity) { - int slot = HashSlot(value, capacity); - for (int probe = 0; probe < capacity; ++probe) { +__device__ int64_t FindSlot(const int* keys, int value, int64_t capacity) { + int64_t slot = HashSlot(value, capacity); + for (int64_t probe = 0; probe < capacity; ++probe) { if (keys[slot] == value || keys[slot] == -1) return slot; slot = (slot + 1) & (capacity - 1); } @@ -36,7 +36,7 @@ __device__ int InclusiveScan(int* flags, int value) { __global__ void MergeRangeKernel( const int* base_indices, const int* base_counts, const int* base_rows, const int* starts, const int* ends, int* indices, int* counts, int* status, - int* workspace, int cached_rows, int base_capacity, int output_capacity, int hash_capacity) { + int* workspace, int cached_rows, int base_capacity, int output_capacity, int64_t hash_capacity) { const int query = blockIdx.x; const int lane = threadIdx.x; int* output = indices + static_cast(query) * output_capacity; @@ -46,10 +46,12 @@ __global__ void MergeRangeKernel( __shared__ int failed; __shared__ int emitted; __shared__ int overlap; + __shared__ int inserted_count; if (lane == 0) { failed = 0; emitted = 0; overlap = 0; + inserted_count = 0; counts[query] = 0; status[query] = 0; } @@ -83,31 +85,35 @@ __global__ void MergeRangeKernel( first[slot] = INT_MAX; } __syncthreads(); - for (int64_t column = lane; column < count; column += kThreads) { - const int value = base[column]; - int slot = HashSlot(value, hash_capacity); - bool inserted = false; - for (int probe = 0; probe < hash_capacity; ++probe) { - const int previous = atomicCAS(keys + slot, -1, value); - if (previous == -1 || previous == value) { - atomicMin(first + slot, static_cast(column)); - inserted = true; - break; + for (int64_t chunk = 0; chunk < count; chunk += kThreads) { + const int64_t column = chunk + lane; + if (column < count) { + const int value = base[column]; + int64_t slot = HashSlot(value, hash_capacity); + bool inserted = false; + for (int64_t probe = 0; probe < hash_capacity; ++probe) { + const int previous = atomicCAS(keys + slot, -1, value); + if (previous == -1 || previous == value) { + atomicMin(first + slot, static_cast(column)); + if (previous == -1 && atomicAdd(&inserted_count, 1) >= output_capacity) atomicExch(&failed, 2); + inserted = true; + break; + } + slot = (slot + 1) & (hash_capacity - 1); } - slot = (slot + 1) & (hash_capacity - 1); + if (!inserted) atomicExch(&failed, 2); + } + __syncthreads(); + if (failed) { + if (lane == 0) status[query] = 2; + return; } - if (!inserted) atomicExch(&failed, 2); - } - __syncthreads(); - if (failed) { - if (lane == 0) status[query] = 2; - return; } for (int64_t chunk = 0; chunk < count; chunk += kThreads) { const int64_t column = chunk + lane; const int value = column < count ? base[column] : -1; - const int slot = column < count ? FindSlot(keys, value, hash_capacity) : -1; + const int64_t slot = column < count ? FindSlot(keys, value, hash_capacity) : -1; const int unique = slot >= 0 && first[slot] == column; const int rank = InclusiveScan(flags, unique); if (unique) { @@ -127,7 +133,7 @@ __global__ void MergeRangeKernel( } for (int64_t chunk = start; chunk < end; chunk += kThreads) { const int64_t value = chunk + lane; - const int slot = value < end ? FindSlot(keys, static_cast(value), hash_capacity) : -1; + const int64_t slot = value < end ? FindSlot(keys, static_cast(value), hash_capacity) : -1; const int unique = value < end && (slot < 0 || keys[slot] != value); const int rank = InclusiveScan(flags, unique); if (unique) output[emitted + rank - 1] = static_cast(value); @@ -142,7 +148,7 @@ __global__ void MergeIndicesKernel( const int* base_indices, const int* base_counts, const int* base_rows, const int* additional_indices, const int* additional_counts, int* indices, int* counts, int* status, int* workspace, - int cached_rows, int base_capacity, int additional_capacity, int output_capacity, int hash_capacity) { + int cached_rows, int base_capacity, int additional_capacity, int output_capacity, int64_t hash_capacity) { const int query = blockIdx.x; const int lane = threadIdx.x; int* output = indices + static_cast(query) * output_capacity; @@ -151,13 +157,15 @@ __global__ void MergeIndicesKernel( __shared__ int flags[kThreads]; __shared__ int failed; __shared__ int emitted; + __shared__ int inserted_count; if (lane == 0) { failed = 0; emitted = 0; + inserted_count = 0; counts[query] = 0; status[query] = 0; } - for (int column = lane; column < output_capacity; column += kThreads) output[column] = -1; + for (int64_t column = lane; column < output_capacity; column += kThreads) output[column] = -1; __syncthreads(); const int row = base_rows[query]; const int added_count = additional_counts[query]; @@ -173,7 +181,7 @@ __global__ void MergeIndicesKernel( const int* base = base_indices + static_cast(row) * base_capacity; const int* added = additional_indices + static_cast(query) * additional_capacity; const int total = base_count + added_count; - for (int column = lane; column < total; column += kThreads) { + for (int64_t column = lane; column < total; column += kThreads) { if ((column < base_count ? base[column] : added[column - base_count]) < 0) atomicExch(&failed, 1); } __syncthreads(); @@ -181,35 +189,39 @@ __global__ void MergeIndicesKernel( if (lane == 0) status[query] = 1; return; } - for (int slot = lane; slot < hash_capacity; slot += kThreads) { + for (int64_t slot = lane; slot < hash_capacity; slot += kThreads) { keys[slot] = -1; first[slot] = INT_MAX; } __syncthreads(); - for (int column = lane; column < total; column += kThreads) { - const int value = column < base_count ? base[column] : added[column - base_count]; - int slot = HashSlot(value, hash_capacity); - bool inserted = false; - for (int probe = 0; probe < hash_capacity; ++probe) { - const int previous = atomicCAS(keys + slot, -1, value); - if (previous == -1 || previous == value) { - atomicMin(first + slot, column); - inserted = true; - break; + for (int64_t chunk = 0; chunk < total; chunk += kThreads) { + const int64_t column = chunk + lane; + if (column < total) { + const int value = column < base_count ? base[column] : added[column - base_count]; + int64_t slot = HashSlot(value, hash_capacity); + bool inserted = false; + for (int64_t probe = 0; probe < hash_capacity; ++probe) { + const int previous = atomicCAS(keys + slot, -1, value); + if (previous == -1 || previous == value) { + atomicMin(first + slot, static_cast(column)); + if (previous == -1 && atomicAdd(&inserted_count, 1) >= output_capacity) atomicExch(&failed, 2); + inserted = true; + break; + } + slot = (slot + 1) & (hash_capacity - 1); } - slot = (slot + 1) & (hash_capacity - 1); + if (!inserted) atomicExch(&failed, 2); + } + __syncthreads(); + if (failed) { + if (lane == 0) status[query] = 2; + return; } - if (!inserted) atomicExch(&failed, 2); - } - __syncthreads(); - if (failed) { - if (lane == 0) status[query] = 2; - return; } - for (int chunk = 0; chunk < total; chunk += kThreads) { - const int column = chunk + lane; + for (int64_t chunk = 0; chunk < total; chunk += kThreads) { + const int64_t column = chunk + lane; const int value = column < total ? (column < base_count ? base[column] : added[column - base_count]) : -1; - const int slot = column < total ? FindSlot(keys, value, hash_capacity) : -1; + const int64_t slot = column < total ? FindSlot(keys, value, hash_capacity) : -1; const int unique = slot >= 0 && first[slot] == column; const int rank = InclusiveScan(flags, unique); if (unique && emitted + rank <= output_capacity) output[emitted + rank - 1] = value; @@ -218,7 +230,7 @@ __global__ void MergeIndicesKernel( __syncthreads(); } if (emitted > output_capacity) { - for (int column = lane; column < output_capacity; column += kThreads) output[column] = -1; + for (int64_t column = lane; column < output_capacity; column += kThreads) output[column] = -1; if (lane == 0) status[query] = 2; } else if (lane == 0) { counts[query] = emitted; @@ -232,7 +244,7 @@ cudaError_t LaunchPackedSparseAttentionIndexerMergeIndices( const int32_t* base_rows, const int32_t* additional_indices, const int32_t* additional_counts, int32_t* indices, int32_t* counts, int32_t* status, int32_t* workspace, int32_t cached_rows, int32_t base_capacity, int32_t additional_capacity, - int32_t query_count, int32_t output_capacity, int32_t hash_capacity) { + int32_t query_count, int32_t output_capacity, int64_t hash_capacity) { if (query_count == 0) return cudaSuccess; MergeIndicesKernel<<>>( base_indices, base_counts, base_rows, additional_indices, additional_counts, @@ -246,7 +258,7 @@ cudaError_t LaunchPackedSparseAttentionIndexerMerge( const int32_t* base_rows, const int32_t* starts, const int32_t* ends, int32_t* indices, int32_t* counts, int32_t* status, int32_t* workspace, int32_t cached_rows, int32_t base_capacity, int32_t query_count, - int32_t output_capacity, int32_t hash_capacity) { + int32_t output_capacity, int64_t hash_capacity) { if (query_count == 0) return cudaSuccess; MergeRangeKernel<<>>( base_indices, base_counts, base_rows, starts, ends, indices, counts, status, diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.h index 293879a7aa0f9..7a80b33916678 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.h @@ -10,13 +10,13 @@ cudaError_t LaunchPackedSparseAttentionIndexerMergeIndices( const int32_t* base_rows, const int32_t* additional_indices, const int32_t* additional_counts, int32_t* indices, int32_t* counts, int32_t* status, int32_t* workspace, int32_t cached_rows, int32_t base_capacity, int32_t additional_capacity, - int32_t query_count, int32_t output_capacity, int32_t hash_capacity); + int32_t query_count, int32_t output_capacity, int64_t hash_capacity); cudaError_t LaunchPackedSparseAttentionIndexerMerge( cudaStream_t stream, const int32_t* base_indices, const int32_t* base_counts, const int32_t* base_rows, const int32_t* starts, const int32_t* ends, int32_t* indices, int32_t* counts, int32_t* status, int32_t* workspace, int32_t cached_rows, int32_t base_capacity, int32_t query_count, - int32_t output_capacity, int32_t hash_capacity); + int32_t output_capacity, int64_t hash_capacity); } // namespace onnxruntime::contrib::cuda \ No newline at end of file diff --git a/onnxruntime/contrib_ops/webgpu/bert/attention.cc b/onnxruntime/contrib_ops/webgpu/bert/attention.cc index 9a6f6062abaa5..7ba3075c6c9fb 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/attention.cc @@ -9,7 +9,6 @@ #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" #include "core/providers/webgpu/webgpu_supported_types.h" #include "core/providers/webgpu/webgpu_utils.h" -#include "core/providers/webgpu/math/matmul.h" using namespace onnxruntime::webgpu; using namespace ::onnxruntime::common; using namespace ONNX_NAMESPACE; @@ -659,7 +658,7 @@ Attention::Attention(const OpKernelInfo& info) Status PrepareQKV(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& parameters, const Tensor* input, const Tensor* weights, const Tensor* bias, Tensor* q, Tensor* k, Tensor* v, - MatMulOptImplCache& matmul_compute_cache, bool weights_are_constant) { + MatMulComputeDispatcher& matmul_compute_dispatcher, bool weights_are_constant) { // Use MatMul to compute packed QKV output: input * weights + bias // Then use SplitPackedQKV to split into Q, K, V in BSD format // Returns Q, K, V in BSD format @@ -673,9 +672,9 @@ Status PrepareQKV(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtte std::vector matmul_inputs = {input, weights, bias}; // Call MatMul: packed_qkv = input * weights + bias - ORT_RETURN_IF_ERROR(onnxruntime::webgpu::ComputeMatMul( - &context, Activation(), matmul_inputs, &packed_qkv, /*is_channels_last=*/true, - matmul_compute_cache, weights_are_constant)); + ORT_RETURN_IF_ERROR(matmul_compute_dispatcher.Compute( + context, Activation(), matmul_inputs, &packed_qkv, /*is_channels_last=*/true, + weights_are_constant)); // Output Q, K, V in BSD format return SplitPackedQKV(context, parameters, &packed_qkv, q, k, v, parameters.hidden_size_); @@ -743,7 +742,7 @@ Status Attention::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) // Compute Q, K, V from input, weights, and bias (returns BSD format) ORT_RETURN_IF_ERROR(PrepareQKV(context, parameters, input, weights, bias, &Q_bsd, &K_bsd, &V_bsd, - matmul_compute_cache_, weights_are_constant_)); + matmul_compute_dispatcher_, weights_are_constant_)); parameters.qkv_format_ = Q_K_V_BSNH; // Check if we can use flash attention diff --git a/onnxruntime/contrib_ops/webgpu/bert/attention.h b/onnxruntime/contrib_ops/webgpu/bert/attention.h index 30864c17e74ef..eb08876a96fdd 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/attention.h @@ -4,7 +4,7 @@ #pragma once #include "core/providers/webgpu/compute_context.h" -#include "core/providers/webgpu/math/matmul.h" +#include "core/providers/webgpu/math/matmul_compute_dispatcher.h" #include "core/providers/webgpu/program.h" #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_kernel.h" @@ -150,7 +150,7 @@ class Attention final : public WebGpuKernel, public onnxruntime::contrib::Attent Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override; private: - mutable MatMulOptImplCache matmul_compute_cache_; + mutable MatMulComputeDispatcher matmul_compute_dispatcher_; bool weights_are_constant_ = false; }; diff --git a/onnxruntime/contrib_ops/webgpu/bert/bias_add.cc b/onnxruntime/contrib_ops/webgpu/bert/bias_add.cc index e53192bbed618..b8693824bce1f 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/bias_add.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/bias_add.cc @@ -61,7 +61,7 @@ Status BiasAdd::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) co BiasAddProgram program{}; program - .AddInputs({{input, ProgramTensorMetadataDependency::None, components}, + .AddInputs({{input, ProgramTensorMetadataDependency::Type, components}, {bias, ProgramTensorMetadataDependency::None, components}, {residual, ProgramTensorMetadataDependency::None, components}}) .AddOutput({output, ProgramTensorMetadataDependency::None, components}) diff --git a/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc b/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc index 4656fc486fccf..788eb01d0337d 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc @@ -23,7 +23,7 @@ ONNX_OPERATOR_KERNEL_EX( Status BiasSplitGeluProgram::GenerateShaderCode(ShaderHelper& shader) const { const ShaderVariableHelper& x = - shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); const ShaderVariableHelper& bias = shader.AddInput("bias", ShaderUsage::UseUniform); const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform); @@ -37,7 +37,7 @@ Status BiasSplitGeluProgram::GenerateShaderCode(ShaderHelper& shader) const { << "let inputOffset = biasIdx + batchIndex * halfChannels * 2;\n" << "let valueLeft = " << x.GetByOffset("inputOffset") << " + " << bias.GetByOffset("biasIdx") << ";\n" << "let valueRight = " << x.GetByOffset("inputOffset + halfChannels") << " + " << bias.GetByOffset("biasIdx + halfChannels") << ";\n" - << "let geluRight = valueRight * 0.5 * (erf_v(valueRight / M_SQRT2) + 1);\n" + << "let geluRight = valueRight * 0.5 * (erf_v(valueRight / x_element_t(M_SQRT2)) + 1);\n" << output.SetByOffset("global_idx", "valueLeft * geluRight"); return Status::OK(); @@ -70,7 +70,7 @@ Status BiasSplitGelu::ComputeInternal(onnxruntime::webgpu::ComputeContext& conte BiasSplitGeluProgram program{}; program - .AddInputs({{input, ProgramTensorMetadataDependency::None, components}, + .AddInputs({{input, ProgramTensorMetadataDependency::Type, components}, {bias, ProgramTensorMetadataDependency::None, components}}) .AddOutput({output, ProgramTensorMetadataDependency::None, components}) .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) @@ -80,4 +80,4 @@ Status BiasSplitGelu::ComputeInternal(onnxruntime::webgpu::ComputeContext& conte } // namespace webgpu } // namespace contrib -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc b/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc index 44747120cc2bc..92af2ee42a95b 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc @@ -15,6 +15,7 @@ ONNX_OPERATOR_KERNEL_EX( BranchwiseRMSNorm, kMSDomain, 1, kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) + .TypeConstraint("F", DataTypeImpl::GetTensorType()) .TypeConstraint("M", WebGpuSupportedFloatTypes()), BranchwiseRMSNorm); @@ -61,7 +62,10 @@ Status BranchwiseRMSNormProgram::GenerateShaderCode(ShaderHelper& shader) const BranchwiseRMSNorm::BranchwiseRMSNorm(const OpKernelInfo& info) : WebGpuKernel(info), epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), - num_branches_(info.GetAttrOrDefault("num_branches", 0)) {} + num_branches_(info.GetAttrOrDefault("num_branches", 0)) { + ORT_ENFORCE(info.GetAttrOrDefault("use_onnx_arithmetic", 0) == 0, + "WebGPU BranchwiseRMSNorm does not support use_onnx_arithmetic=1"); +} Status BranchwiseRMSNorm::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { const auto* x = context.Input(0); diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index addf1280ed177..717abe8274b2b 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -470,18 +470,18 @@ Status FlashAttentionPagedPrefillProgram::GenerateShaderCode(ShaderHelper& shade WGSL_TEMPLATE_VARIABLE(value_cache, value_cache)); } -Status ComputeFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& context, - const Tensor* q, - const Tensor* key_cache, - const Tensor* value_cache, - const Tensor* block_table, - Tensor* output, - const Tensor* seqlen_k, - const Tensor* seqlens_q, - const WebgpuAttentionParameters& parameters, - uint32_t block_size, - uint32_t max_num_blocks_per_seq, - const Tensor* cumulative_seqlens_q) { +Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& context, + const Tensor* q, + const Tensor* key_cache, + const Tensor* value_cache, + const Tensor* block_table, + Tensor* output, + const Tensor* seqlen_k, + const Tensor* seqlens_q, + const WebgpuAttentionParameters& parameters, + uint32_t block_size, + uint32_t max_num_blocks_per_seq, + const Tensor* cumulative_seqlens_q) { ORT_RETURN_IF_NOT(q != nullptr && key_cache != nullptr && value_cache != nullptr && block_table != nullptr, "Paged prefill requires Q, K/V cache, and block_table."); ORT_RETURN_IF_NOT(seqlen_k != nullptr && seqlens_q != nullptr, @@ -931,6 +931,217 @@ Status ComputeFlashAttentionPagedDecodeVxReduce(onnxruntime::webgpu::ComputeCont return context.RunProgram(program); } +namespace { + +bool CanApplyDensePrefill(const Tensor* Q, const WebgpuAttentionParameters& parameters, + onnxruntime::webgpu::ComputeContext& context, + const Tensor* block_table, uint32_t kv_cache_quantization_bits) { + if (block_table != nullptr || parameters.sequence_length_ < 32) { + return false; + } + const auto vendor = context.AdapterInfo().vendor; + const bool is_apple = vendor == std::string_view{"apple"}; + return DensePrefillFitsWorkgroupStorage( + is_apple || vendor == std::string_view{"nvidia"} || !context.HasFeature(wgpu::FeatureName::Subgroups), + Q->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, + parameters.head_size_, kv_cache_quantization_bits, vendor == std::string_view{"qualcomm"}, + is_apple ? 128u : 64u, context.DeviceLimits().maxComputeWorkgroupStorageSize); +} + +bool CanApplyFusedPagedPrefill(const Tensor* Q, const WebgpuAttentionParameters& parameters, + onnxruntime::webgpu::ComputeContext& context, + const Tensor* block_table, uint32_t block_size, + uint32_t kv_cache_quantization_bits, const Tensor* attention_bias, + const Tensor* head_sink, const Tensor* seqlen_k, const Tensor* seqlens_q) { + return block_table != nullptr && kv_cache_quantization_bits == 0 && + attention_bias == nullptr && head_sink == nullptr && parameters.qkv_format_ == Q_K_V_BSNH && + seqlen_k != nullptr && seqlens_q != nullptr && + CanApplyFlashAttentionPagedPrefill(context, + Q->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, + parameters.sequence_length_, parameters.head_size_, static_cast(block_size)); +} + +bool CanApplySplitReduce(const Tensor* Q, const WebgpuAttentionParameters& parameters, + onnxruntime::webgpu::ComputeContext& context, + const Tensor* block_table, uint32_t kv_cache_quantization_bits) { + // Short queries favor split-reduce; dense prefill also falls back when storage does not fit. + return parameters.sequence_length_ < 32 || + (block_table == nullptr && + !CanApplyDensePrefill(Q, parameters, context, block_table, kv_cache_quantization_bits)); +} + +bool CanApplyPagedDecode(const Tensor* block_table, uint32_t kv_cache_quantization_bits) { + return block_table != nullptr && kv_cache_quantization_bits == 0; +} + +Status ApplyDensePrefill(const Tensor* Q, Tensor* present_key, Tensor* present_value, + const Tensor* attention_bias, Tensor* attn_output, + const WebgpuAttentionParameters& parameters, + onnxruntime::webgpu::ComputeContext& context, + const Tensor* seqlen_k, const Tensor* seqlens_q, const Tensor* head_sink, + uint32_t present_sequence_length, uint32_t kv_cache_quantization_bits, + int compressed_head_size_u32, int local_window_size) { + constexpr uint32_t tile_size = 64; + const bool use_seqlen_k = seqlen_k != nullptr; + const bool use_seqlens_q = seqlens_q != nullptr; + const bool kv_cache_quantization_enabled = kv_cache_quantization_bits != 0; + const bool has_local_window = local_window_size > 0; + const bool is_nvidia = context.AdapterInfo().vendor == std::string_view{"nvidia"}; + const bool is_apple = context.AdapterInfo().vendor == std::string_view{"apple"}; + const bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; + const bool has_subgroups = context.HasFeature(wgpu::FeatureName::Subgroups); + // Prefill path: FlashAttentionProgram (single kernel with subgroup shuffles) + bool has_attention_bias = attention_bias != nullptr; + bool is_fp16 = Q->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; + bool q_BNSH = parameters.qkv_format_ == Q_K_V_BNSH; + bool has_head_sink = head_sink != nullptr; + FlashAttentionProgram program{"FlashAttention", + has_attention_bias, + is_qualcomm, + is_fp16, + parameters.head_size_, + parameters.num_heads_, + parameters.is_unidirectional_, + is_nvidia, + is_apple, + has_subgroups, + q_BNSH, + use_seqlen_k, + has_head_sink, + has_local_window, + kv_cache_quantization_bits, + compressed_head_size_u32, + use_seqlens_q}; + program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, 4}, + {present_key, ProgramTensorMetadataDependency::TypeAndRank, + kv_cache_quantization_enabled ? 1 : 4}, + {present_value, ProgramTensorMetadataDependency::TypeAndRank, + kv_cache_quantization_enabled ? 1 : 4}}); + if (has_attention_bias) { + program.AddInputs({{attention_bias, ProgramTensorMetadataDependency::TypeAndRank}}); + } + if (use_seqlen_k) { + program.AddInputs({{seqlen_k, ProgramTensorMetadataDependency::None}}); + } + if (use_seqlens_q) { + program.AddInputs({{seqlens_q, ProgramTensorMetadataDependency::None}}); + } + if (has_head_sink) { + program.AddInputs({{head_sink, ProgramTensorMetadataDependency::Type}}); + } + program.AddOutputs({{attn_output, ProgramTensorMetadataDependency::TypeAndRank, 4}}); + const float alpha = parameters.scale_ == 0.0f ? 1.f / sqrt(static_cast(parameters.head_size_)) + : parameters.scale_; + + // On Apple GPUs, use a larger workgroup size to reduce barrier overhead. + const uint32_t prefill_tile_size = is_apple ? 128 : tile_size; + const uint32_t num_seq_tile = (parameters.sequence_length_ + prefill_tile_size - 1) / prefill_tile_size; + + uint32_t attn_bias_dim0 = 1; + uint32_t attn_bias_dim1 = 1; + uint32_t attn_bias_dim3 = 0; + if (has_attention_bias) { + const auto& bias_shape = attention_bias->Shape(); + attn_bias_dim0 = static_cast(bias_shape[0]); + attn_bias_dim1 = static_cast(bias_shape[1]); + attn_bias_dim3 = static_cast(bias_shape[3]); + } + + program.SetDispatchGroupSize(parameters.batch_size_ * parameters.num_heads_ * num_seq_tile) + .SetWorkgroupSize(prefill_tile_size) + .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, + parameters.is_unidirectional_, is_qualcomm, is_nvidia, is_apple, + has_subgroups, q_BNSH, use_seqlen_k, has_head_sink, has_local_window, + kv_cache_quantization_bits, + compressed_head_size_u32, program.max_k_step(), use_seqlens_q) + .AddUniformVariables({{static_cast(parameters.sequence_length_)}, + {static_cast(parameters.total_sequence_length_)}, + {static_cast(present_sequence_length)}, + {static_cast(parameters.batch_size_)}, + {static_cast(parameters.n_reps)}, + {alpha}, + {num_seq_tile}, + {attn_bias_dim0}, + {attn_bias_dim1}, + {attn_bias_dim3}, + {static_cast(has_local_window ? local_window_size : 0)}}); + + return context.RunProgram(program); +} + +Status ApplySplitReduce(const Tensor* Q, Tensor* present_key, Tensor* present_value, + const Tensor* attention_bias, Tensor* attn_output, + const WebgpuAttentionParameters& parameters, + onnxruntime::webgpu::ComputeContext& context, + const Tensor* seqlen_k, const Tensor* seqlens_q, const Tensor* head_sink, + uint32_t present_sequence_length, uint32_t kv_cache_quantization_bits, + int compressed_head_size_u32, int local_window_size, const Tensor* block_table, + uint32_t block_size, uint32_t max_num_blocks_per_seq, + Tensor* indirect_buffer_ptr, uint32_t m_tile, + const Tensor* total_seqlen) { + constexpr uint32_t tile_size = 64; + const bool use_seqlen_k = seqlen_k != nullptr; + const bool use_seqlens_q = seqlens_q != nullptr; + const bool use_paged_kv_cache = block_table != nullptr; + const bool use_indirect_dispatch = indirect_buffer_ptr != nullptr; + // Split-reduce path (fused QKV + VxReduce). Handles quantized and unquantized caches. + const uint32_t num_total_seq_length_tile = (parameters.total_sequence_length_ + tile_size - 1) / tile_size; + const uint32_t num_present_sequence_length_tile = (present_sequence_length + tile_size - 1) / tile_size; + + const TensorShapeVector metadata_dims({parameters.batch_size_, parameters.num_heads_, + parameters.sequence_length_, num_present_sequence_length_tile, 2}); + const TensorShape metadata_shape(metadata_dims); + Tensor metadata = context.CreateGPUTensor(DataTypeImpl::GetType(), metadata_shape); + + const TensorShapeVector out_split_vx_dims({parameters.batch_size_, parameters.num_heads_, + parameters.sequence_length_, num_present_sequence_length_tile, parameters.head_size_}); + const TensorShape out_split_vx_shape(out_split_vx_dims); + Tensor out_split_vx = context.CreateGPUTensor(Q->DataType(), out_split_vx_shape); + + if (CanApplyPagedDecode(block_table, kv_cache_quantization_bits)) { + ORT_RETURN_IF_ERROR(ComputeFlashAttentionPagedDecodeQKV(context, Q, attention_bias, &out_split_vx, present_key, present_value, + &metadata, seqlen_k, block_table, + parameters, indirect_buffer_ptr, num_total_seq_length_tile, + num_present_sequence_length_tile, tile_size, use_indirect_dispatch, + present_sequence_length, m_tile, use_seqlen_k, total_seqlen, + kv_cache_quantization_bits, + compressed_head_size_u32, + use_seqlens_q, seqlens_q, + block_size, max_num_blocks_per_seq)); + + ORT_RETURN_IF_ERROR(ComputeFlashAttentionPagedDecodeVxReduce(context, &out_split_vx, &metadata, attn_output, seqlen_k, parameters, + num_total_seq_length_tile, + num_present_sequence_length_tile, tile_size, + head_sink, m_tile, use_seqlen_k)); + } else { + // Quantized paged caches cannot be interpreted by the dense decode shader. + if (use_paged_kv_cache) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "FlashAttention (WebGPU): paged KV cache present but the paged " + "decode path was rejected (KV cache quantization enabled). Extend " + "FlashAttentionPagedDecodeQKV to support the requested feature, " + "or gate the feature off at the PagedAttention layer before " + "dispatching FA."); + } + ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeQKV(context, Q, attention_bias, &out_split_vx, present_key, present_value, + &metadata, seqlen_k, + parameters, indirect_buffer_ptr, num_total_seq_length_tile, + num_present_sequence_length_tile, tile_size, use_indirect_dispatch, + present_sequence_length, m_tile, use_seqlen_k, total_seqlen, + kv_cache_quantization_bits, + compressed_head_size_u32, + use_seqlens_q, seqlens_q, local_window_size)); + + ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeVxReduce(context, &out_split_vx, &metadata, attn_output, seqlen_k, parameters, + num_total_seq_length_tile, + num_present_sequence_length_tile, present_sequence_length, tile_size, + head_sink, m_tile, use_seqlen_k)); + } + return Status::OK(); +} + +} // namespace + Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const Tensor* attention_bias, Tensor* output, const Tensor* past_key, Tensor* present_key, const Tensor* past_value, Tensor* present_value, const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context, const Tensor* seqlen_k, @@ -939,9 +1150,7 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co const Tensor* block_table, uint32_t block_size, uint32_t max_num_blocks_per_seq, const Tensor* cumulative_seqlens_q, int local_window_size) { constexpr uint32_t tile_size = 64; - const bool use_seqlens_q = seqlens_q != nullptr; const bool use_paged_kv_cache = block_table != nullptr; - const bool has_local_window = local_window_size > 0; const uint32_t kv_cache_quantization_bits = context.KvCacheQuantizationBits(); const bool kv_cache_quantization_enabled = kv_cache_quantization_bits != 0; @@ -1219,237 +1428,33 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co attn_output = &attn_output_temp; } - // Route between prefill path (FlashAttentionProgram, single kernel) - // and split-reduce decode path (QKV + VxReduce, 2 kernels). - // Split-reduce wins for short Q (sequence_length < 32) across all KV - // cache lengths measured: 1.13x-2.07x faster at total_sequence_length - // 128 / 500 / 2000 on a representative LLM (32 heads, head_size 96). - const bool is_fp16_q = - Q->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; - const bool is_nvidia = context.AdapterInfo().vendor == std::string_view{"nvidia"}; - const bool is_apple = context.AdapterInfo().vendor == std::string_view{"apple"}; - const bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; - const bool has_subgroups = context.HasFeature(wgpu::FeatureName::Subgroups); - const uint32_t dense_prefill_workgroup_size = is_apple ? 128 : tile_size; - const bool dense_prefill_fits_workgroup_storage = - DensePrefillFitsWorkgroupStorage( - is_apple || is_nvidia || !has_subgroups, is_fp16_q, parameters.head_size_, - kv_cache_quantization_bits, is_qualcomm, dense_prefill_workgroup_size, - context.DeviceLimits().maxComputeWorkgroupStorageSize); - const bool use_split_reduce = - parameters.sequence_length_ < 32 || - (!use_paged_kv_cache && !dense_prefill_fits_workgroup_storage); - - if (!use_split_reduce) { - // Ask the shared helper whether the fused paged-prefill shader can run on - // this (adapter, config, shape) triple, then AND in the additional - // "features not yet supported by the paged shader" bits that only the FA - // caller can see (attention_bias, head_sink, KV cache quantization, QKV format, - // varlen-metadata inputs). Keeping the adapter/dtype/shape gate in the - // helper is the anti-drift invariant: PagedAttention uses the same - // predicate to decide whether it can hand FA a packed-varlen Q view. - const bool use_paged_prefill = - use_paged_kv_cache && !kv_cache_quantization_enabled && - attention_bias == nullptr && head_sink == nullptr && - parameters.qkv_format_ == Q_K_V_BSNH && - seqlen_k != nullptr && seqlens_q != nullptr && - ShouldRunFusedPagedPrefill(context, is_fp16_q, parameters.sequence_length_, - parameters.head_size_, static_cast(block_size)); - if (use_paged_prefill) { - ORT_RETURN_IF_ERROR(ComputeFlashAttentionPagedPrefill(context, - Q, - present_key, - present_value, - block_table, - attn_output, - seqlen_k, - seqlens_q, - parameters, - block_size, - max_num_blocks_per_seq, - cumulative_seqlens_q)); - } else { - // Fall-through defensive guard. When use_paged_kv_cache is true, - // past_key / past_value point at the paged KV cache (shape - // (num_blocks, block_size, kv_num_heads, head_size)), not a dense - // BNSH tensor. The dense FlashAttentionProgram below would read them - // as if they were BNSH and silently corrupt the output — bit-for-bit - // wrong, no crash. - // - // Today the AND-chain above holds by construction: PagedAttention v1 - // rejects head_sink / softcap / quantized KV caches / non-SEPARATE-layout at - // input validation and force-sets qkv_format = BSNH, so - // use_paged_prefill collapses to ShouldRunFusedPagedPrefill(). When - // that helper rejects (fp32, block_size < max_k_step, head_size > 256), - // PagedAttention takes the gather-then-flash cascade upstream and - // calls FA with block_table = nullptr, so this branch is never reached - // with use_paged_kv_cache == true. - // - // If a future Phase-2 item (softcap, head_sink, q_norm_weight, ...) - // is added to PagedAttention without matching support in - // FlashAttentionPagedPrefillProgram, this guard makes the mismatch a - // loud CI failure instead of a silent output-corruption bug. - if (use_paged_kv_cache) { - return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "FlashAttention (WebGPU): paged KV cache present but the fused " - "paged-prefill path was rejected by the extra prefill AND-chain " - "(attention_bias / head_sink / KV cache quantization / non-BSNH qkv_format / " - "missing seqlen). Extend FlashAttentionPagedPrefillProgram to " - "support the requested feature, or gate the feature off at the " - "PagedAttention layer before dispatching FA."); - } - // Prefill path: FlashAttentionProgram (single kernel with subgroup shuffles) - bool has_attention_bias = attention_bias != nullptr; - bool is_fp16 = is_fp16_q; - bool q_BNSH = parameters.qkv_format_ == Q_K_V_BNSH; - bool has_head_sink = head_sink != nullptr; - FlashAttentionProgram program{"FlashAttention", - has_attention_bias, - is_qualcomm, - is_fp16, - parameters.head_size_, - parameters.num_heads_, - parameters.is_unidirectional_, - is_nvidia, - is_apple, - has_subgroups, - q_BNSH, - use_seqlen_k, - has_head_sink, - has_local_window, - kv_cache_quantization_bits, - compressed_head_size_u32, - use_seqlens_q}; - // When TQ is active, KV cache is u32-packed — use u32 tensor views for present_key/present_value. - const Tensor* fa_present_key = - kv_cache_quantization_enabled ? quantized_present_key : present_key; - const Tensor* fa_present_value = - kv_cache_quantization_enabled ? quantized_present_value : present_value; - program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, 4}, - {fa_present_key, ProgramTensorMetadataDependency::TypeAndRank, - kv_cache_quantization_enabled ? 1 : 4}, - {fa_present_value, ProgramTensorMetadataDependency::TypeAndRank, - kv_cache_quantization_enabled ? 1 : 4}}); - if (has_attention_bias) { - program.AddInputs({{attention_bias, ProgramTensorMetadataDependency::TypeAndRank}}); - } - if (use_seqlen_k) { - program.AddInputs({{seqlen_k, ProgramTensorMetadataDependency::None}}); - } - if (use_seqlens_q) { - program.AddInputs({{seqlens_q, ProgramTensorMetadataDependency::None}}); - } - if (has_head_sink) { - program.AddInputs({{head_sink, ProgramTensorMetadataDependency::Type}}); - } - program.AddOutputs({{attn_output, ProgramTensorMetadataDependency::TypeAndRank, 4}}); - const float alpha = parameters.scale_ == 0.0f ? 1.f / sqrt(static_cast(parameters.head_size_)) - : parameters.scale_; - - // On Apple GPUs, use a larger workgroup size to reduce barrier overhead. - const uint32_t prefill_tile_size = is_apple ? 128 : tile_size; - const uint32_t num_seq_tile = (parameters.sequence_length_ + prefill_tile_size - 1) / prefill_tile_size; - - uint32_t attn_bias_dim0 = 1; - uint32_t attn_bias_dim1 = 1; - uint32_t attn_bias_dim3 = 0; - if (has_attention_bias) { - const auto& bias_shape = attention_bias->Shape(); - attn_bias_dim0 = static_cast(bias_shape[0]); - attn_bias_dim1 = static_cast(bias_shape[1]); - attn_bias_dim3 = static_cast(bias_shape[3]); - } - - program.SetDispatchGroupSize(parameters.batch_size_ * parameters.num_heads_ * num_seq_tile) - .SetWorkgroupSize(prefill_tile_size) - .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, - parameters.is_unidirectional_, is_qualcomm, is_nvidia, is_apple, - has_subgroups, q_BNSH, use_seqlen_k, has_head_sink, has_local_window, - kv_cache_quantization_bits, - compressed_head_size_u32, program.max_k_step(), use_seqlens_q) - .AddUniformVariables({{static_cast(parameters.sequence_length_)}, - {static_cast(parameters.total_sequence_length_)}, - {static_cast(present_sequence_length)}, - {static_cast(parameters.batch_size_)}, - {static_cast(parameters.n_reps)}, - {alpha}, - {num_seq_tile}, - {attn_bias_dim0}, - {attn_bias_dim1}, - {attn_bias_dim3}, - {static_cast(has_local_window ? local_window_size : 0)}}); - - ORT_RETURN_IF_ERROR(context.RunProgram(program)); - } + Tensor* attention_key = kv_cache_quantization_enabled ? quantized_present_key : present_key; + Tensor* attention_value = kv_cache_quantization_enabled ? quantized_present_value : present_value; + + if (CanApplyFusedPagedPrefill(Q, parameters, context, block_table, block_size, + kv_cache_quantization_bits, attention_bias, head_sink, seqlen_k, seqlens_q)) { + ORT_RETURN_IF_ERROR(ApplyFlashAttentionPagedPrefill(context, Q, attention_key, attention_value, + block_table, attn_output, seqlen_k, seqlens_q, + parameters, block_size, max_num_blocks_per_seq, + cumulative_seqlens_q)); + } else if (CanApplyDensePrefill(Q, parameters, context, block_table, kv_cache_quantization_bits)) { + ORT_RETURN_IF_ERROR(ApplyDensePrefill(Q, attention_key, attention_value, attention_bias, + attn_output, parameters, context, seqlen_k, seqlens_q, + head_sink, present_sequence_length, kv_cache_quantization_bits, + compressed_head_size_u32, local_window_size)); + } else if (CanApplySplitReduce(Q, parameters, context, block_table, kv_cache_quantization_bits)) { + ORT_RETURN_IF_ERROR(ApplySplitReduce(Q, attention_key, attention_value, attention_bias, + attn_output, parameters, context, seqlen_k, seqlens_q, + head_sink, present_sequence_length, kv_cache_quantization_bits, + compressed_head_size_u32, local_window_size, block_table, + block_size, max_num_blocks_per_seq, indirect_buffer_ptr, + m_tile, total_seqlen)); } else { - // Split-reduce path (fused QKV + VxReduce). Handles quantized and unquantized caches. - const uint32_t num_total_seq_length_tile = (parameters.total_sequence_length_ + tile_size - 1) / tile_size; - const uint32_t num_present_sequence_length_tile = (present_sequence_length + tile_size - 1) / tile_size; - - const TensorShapeVector metadata_dims({parameters.batch_size_, parameters.num_heads_, - parameters.sequence_length_, num_present_sequence_length_tile, 2}); - const TensorShape metadata_shape(metadata_dims); - Tensor metadata = context.CreateGPUTensor(DataTypeImpl::GetType(), metadata_shape); - - const TensorShapeVector out_split_vx_dims({parameters.batch_size_, parameters.num_heads_, - parameters.sequence_length_, num_present_sequence_length_tile, parameters.head_size_}); - const TensorShape out_split_vx_shape(out_split_vx_dims); - Tensor out_split_vx = context.CreateGPUTensor(Q->DataType(), out_split_vx_shape); - - Tensor* qkv_present_key = - kv_cache_quantization_enabled ? quantized_present_key : present_key; - Tensor* qkv_present_value = - kv_cache_quantization_enabled ? quantized_present_value : present_value; - - // Phase 2 scaffold: when per-batch Q lengths are provided (PagedAttention path), - // route through duplicated decode programs so KV-page-aware changes stay isolated - // from baseline FlashAttention decode kernels. - const bool use_paged_decode_programs = - use_paged_kv_cache && !kv_cache_quantization_enabled; - if (use_paged_decode_programs) { - ORT_RETURN_IF_ERROR(ComputeFlashAttentionPagedDecodeQKV(context, Q, attention_bias, &out_split_vx, qkv_present_key, qkv_present_value, - &metadata, seqlen_k, block_table, - parameters, indirect_buffer_ptr, num_total_seq_length_tile, - num_present_sequence_length_tile, tile_size, use_indirect_dispatch, - present_sequence_length, m_tile, use_seqlen_k, total_seqlen, - kv_cache_quantization_bits, - compressed_head_size_u32, - use_seqlens_q, seqlens_q, - block_size, max_num_blocks_per_seq)); - - ORT_RETURN_IF_ERROR(ComputeFlashAttentionPagedDecodeVxReduce(context, &out_split_vx, &metadata, attn_output, seqlen_k, parameters, - num_total_seq_length_tile, - num_present_sequence_length_tile, tile_size, - head_sink, m_tile, use_seqlen_k)); - } else { - // Fall-through defensive guard (symmetric to the prefill branch above). - // When use_paged_kv_cache is true, dropping into the dense - // FlashAttentionDecodeQKV shader would misinterpret the paged cache as - // dense BNSH and silently corrupt output. Today - // Paged quantized KV cache is currently unsupported. Keep this guard so - // a future paged-cache integration fails loudly instead of corrupting output. - if (use_paged_kv_cache) { - return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "FlashAttention (WebGPU): paged KV cache present but the paged " - "decode path was rejected (KV cache quantization enabled). Extend " - "FlashAttentionPagedDecodeQKV to support the requested feature, " - "or gate the feature off at the PagedAttention layer before " - "dispatching FA."); - } - ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeQKV(context, Q, attention_bias, &out_split_vx, qkv_present_key, qkv_present_value, - &metadata, seqlen_k, - parameters, indirect_buffer_ptr, num_total_seq_length_tile, - num_present_sequence_length_tile, tile_size, use_indirect_dispatch, - present_sequence_length, m_tile, use_seqlen_k, total_seqlen, - kv_cache_quantization_bits, - compressed_head_size_u32, - use_seqlens_q, seqlens_q, local_window_size)); - - ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeVxReduce(context, &out_split_vx, &metadata, attn_output, seqlen_k, parameters, - num_total_seq_length_tile, - num_present_sequence_length_tile, present_sequence_length, tile_size, - head_sink, m_tile, use_seqlen_k)); - } + // A paged cache must never fall through to a shader that assumes dense BNSH storage. + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "FlashAttention (WebGPU): paged KV cache present but the fused " + "paged-prefill path was rejected by the dtype / shape / attention_bias / " + "head_sink / KV cache quantization / qkv_format / sequence-length gates."); } // Apply the Q4 inverse Hadamard transform: attn_output_temp -> output. @@ -1466,11 +1471,11 @@ bool CanApplyFlashAttention(const WebgpuAttentionParameters& parameters, onnxrun ((context.AdapterInfo().vendor == std::string_view{"qualcomm"} && parameters.head_size_ % 8 == 0) || parameters.head_size_ % 4 == 0); } -bool ShouldRunFusedPagedPrefill(onnxruntime::webgpu::ComputeContext& context, - bool is_fp16, - int max_seqlen_q, - int head_size, - int block_size) { +bool CanApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& context, + bool is_fp16, + int max_seqlen_q, + int head_size, + int block_size) { // The paged prefill shader uses only shared-memory algorithms (no subgroup // intrinsics), so it runs correctly on every WebGPU adapter. No adapter // gate here — the only gates are dtype, shape, shm-budget, and diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index 14e3dea471a09..7c0eaf624718e 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -192,18 +192,18 @@ class FlashAttentionPagedPrefillProgram final : public Program { public: @@ -380,11 +380,11 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co // // The shader uses only shared-memory algorithms (no subgroup intrinsics), so // it runs correctly on every WebGPU adapter — no adapter-class gate. -bool ShouldRunFusedPagedPrefill(onnxruntime::webgpu::ComputeContext& context, - bool is_fp16, - int max_seqlen_q, - int head_size, - int block_size); +bool CanApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& context, + bool is_fp16, + int max_seqlen_q, + int head_size, + int block_size); bool CanApplyFlashAttention(const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context); diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_paged_prefill.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_paged_prefill.wgsl.template index de2b4aefcd0ee..ec123f4ecc397 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_paged_prefill.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_paged_prefill.wgsl.template @@ -24,7 +24,7 @@ // it is preferable. // // Alignment invariant: paged_attention_helper enforces block_size to be a -// power of two >= 16, and ShouldRunFusedPagedPrefill additionally rejects +// power of two >= 16, and CanApplyFlashAttentionPagedPrefill additionally rejects // block_size < max_k_step. max_k_step is 32 (fp16, head <= 128) or 16 // otherwise, and k_start advances by max_k_step starting at 0. Because both // block_size and max_k_step are powers of two and block_size >= max_k_step, @@ -57,7 +57,7 @@ const head_size_vec : u32 = head_size / vec_factor; // K and V tiles in shared memory. Budget = 2 * head_size * max_k_step * // element_size B (two tiles, k and v). Sized for the 16 KiB shm budget -// enforced by ShouldRunFusedPagedPrefill: +// enforced by CanApplyFlashAttentionPagedPrefill: // * fp16 head_size 128, max_k_step 32 -> 2 * 128 * 32 * 2 = 16 KiB // * fp16 head_size 256, max_k_step 16 -> 2 * 256 * 16 * 2 = 16 KiB // (matches the dense FA sizing at shm-path adapters.) diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net_state_replay.cc b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net_state_replay.cc new file mode 100644 index 0000000000000..620e320d517e2 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net_state_replay.cc @@ -0,0 +1,179 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/webgpu/bert/gated_delta_net_state_replay.h" + +#include +#include +#include + +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "core/providers/webgpu/shader_helper.h" + +namespace onnxruntime::contrib::webgpu { + +using namespace onnxruntime::webgpu; + +ONNX_OPERATOR_KERNEL_EX( + GatedDeltaNetStateReplay, + kMSDomain, + 1, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) + .InputMemoryType(OrtMemTypeCPUInput, 3) + .MayInplace(2, 0), + GatedDeltaNetStateReplay); + +namespace { + +constexpr uint32_t kWorkgroupSize = 128; +constexpr uint64_t kMaxIndex = std::numeric_limits::max(); + +class GatedDeltaNetStateReplayProgram final : public Program { + public: + GatedDeltaNetStateReplayProgram() : Program{"GatedDeltaNetStateReplay"} {} + + Status GenerateShaderCode(ShaderHelper& shader) const override { + const auto& source = shader.AddInput("source", ShaderUsage::UseUniform); + const auto& decay = shader.AddInput("decay", ShaderUsage::UseUniform); + const auto& key = shader.AddInput("key", ShaderUsage::UseUniform); + const auto& delta = shader.AddInput("delta", ShaderUsage::UseUniform); + const auto& destination = shader.AddOutput("destination", ShaderUsage::UseUniform); + return WGSL_TEMPLATE_APPLY(shader, "bert/gated_delta_net_state_replay.wgsl.template", + WGSL_TEMPLATE_VARIABLE(decay, decay), + WGSL_TEMPLATE_VARIABLE(delta, delta), + WGSL_TEMPLATE_VARIABLE(destination, destination), + WGSL_TEMPLATE_VARIABLE(key, key), + WGSL_TEMPLATE_VARIABLE(source, source)); + } + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"state_elements", ProgramUniformVariableDataType::Uint32}, + {"num_heads_v", ProgramUniformVariableDataType::Uint32}, + {"head_size_v", ProgramUniformVariableDataType::Uint32}, + {"head_size_k", ProgramUniformVariableDataType::Uint32}, + {"num_heads_k", ProgramUniformVariableDataType::Uint32}, + {"kept_count", ProgramUniformVariableDataType::Uint32}); +}; + +// Enqueue only. The operator's completion boundary is deliberately separate from the dispatch. +Status EnqueueReplay(onnxruntime::webgpu::ComputeContext& context, const Tensor& source, const Tensor& capsule, + Tensor& destination, const std::array& desc, + uint32_t state_elements, uint32_t decay_elements, + uint32_t key_elements, uint32_t delta_elements, + uint32_t dispatch_x, uint32_t dispatch_y) { + GatedDeltaNetStateReplayProgram program; + program + .AddInputs({ProgramInput::BufferView(&source, ProgramTensorMetadataDependency::Type, + TensorShape{state_elements}, desc[0]), + ProgramInput::BufferView(&capsule, ProgramTensorMetadataDependency::Type, + TensorShape{decay_elements}, desc[2]), + ProgramInput::BufferView(&capsule, ProgramTensorMetadataDependency::Type, + TensorShape{key_elements}, desc[3]), + ProgramInput::BufferView(&capsule, ProgramTensorMetadataDependency::Type, + TensorShape{delta_elements}, desc[4])}) + .AddOutput(ProgramOutput::BufferView(&destination, ProgramTensorMetadataDependency::Type, + TensorShape{state_elements}, desc[1])) + .SetDispatchGroupSize(dispatch_x, dispatch_y) + .SetWorkgroupSize(kWorkgroupSize) + .AddUniformVariables({{state_elements}, {desc[5]}, {desc[6]}, {desc[7]}, {desc[8]}, {desc[10]}}); + const auto& limits = context.DeviceLimits(); + const auto segment_count = [&limits](const Tensor& backing) { + return (static_cast(backing.SizeInBytes()) - 1) / limits.maxStorageBufferBindingSize + 1; + }; + uint64_t binding_count = 0; + // Use the program's Tensor-owner identity, not raw buffer equality (wrappers may differ). + for (size_t index = 0; index < program.Inputs().size(); ++index) { + if (program.InputBufferOwner(index) == index) { + binding_count += segment_count(*program.Inputs()[index].tensor); + } + } + for (size_t index = 0; index < program.Outputs().size(); ++index) { + if (program.OutputBufferOwner(index) == index) { + binding_count += segment_count(*program.Outputs()[index].tensor); + } + } + ORT_RETURN_IF_NOT(binding_count <= limits.maxStorageBuffersPerShaderStage && + binding_count + 1 <= limits.maxBindingsPerBindGroup, + "GatedDeltaNetStateReplay backing tensors exceed WebGPU binding limits."); + return context.RunProgram(program); +} + +} // namespace + +Status GatedDeltaNetStateReplay::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + ORT_RETURN_IF(context.IsGraphCaptureEnabled(), "GatedDeltaNetStateReplay does not support graph capture."); + const auto* source = context.Input(0); + const auto* capsule = context.Input(1); + const auto* destination = context.Input(2); + const auto* metadata = context.Input(3); + ORT_RETURN_IF_NOT(source && capsule && destination && metadata, + "GatedDeltaNetStateReplay requires all four inputs."); + for (const auto* backing : {source, capsule, destination}) { + ORT_RETURN_IF_NOT(backing->IsDataType() && backing->Shape().NumDimensions() == 1, + "GatedDeltaNetStateReplay backing tensors must be rank-1 FP32."); + ORT_RETURN_IF_NOT(backing->Shape().Size() > 0 && + static_cast(backing->Shape().Size()) <= kMaxIndex, + "GatedDeltaNetStateReplay backing extent exceeds the supported indexing range."); + } + ORT_RETURN_IF_NOT(metadata->IsDataType() && metadata->Shape() == TensorShape{11}, + "GatedDeltaNetStateReplay metadata must be CPU int64 with shape [11]."); + + std::array desc{}; + const auto* values = metadata->Data(); + for (size_t index = 0; index < desc.size(); ++index) { + ORT_RETURN_IF_NOT(values[index] >= 0 && static_cast(values[index]) <= kMaxIndex, + "GatedDeltaNetStateReplay metadata field ", index, " is outside the uint32 indexing range."); + desc[index] = static_cast(values[index]); + } + const uint64_t hv = desc[5], dv = desc[6], dk = desc[7], hk = desc[8], capacity = desc[9]; + ORT_RETURN_IF_NOT(hv > 0 && dv > 0 && dk > 0 && hk > 0 && capacity > 0, + "GatedDeltaNetStateReplay dimensions and capacity must be positive."); + ORT_RETURN_IF_NOT(desc[10] > 0 && desc[10] <= capacity, + "GatedDeltaNetStateReplay kept_count must be in [1, capacity]."); + ORT_RETURN_IF_NOT(hv * dv <= kMaxIndex && hk * dk <= kMaxIndex && (hv - 1) * hk <= kMaxIndex, + "GatedDeltaNetStateReplay head geometry exceeds the uint32 indexing range."); + const uint64_t state_elements = hv * dv * dk; + const uint64_t decay_elements = capacity * hv; + const uint64_t key_elements = capacity * (hk * dk); + const uint64_t delta_elements = capacity * (hv * dv); + ORT_RETURN_IF_NOT(state_elements <= kMaxIndex && + decay_elements <= kMaxIndex && key_elements <= kMaxIndex && delta_elements <= kMaxIndex, + "GatedDeltaNetStateReplay tensor geometry exceeds the uint32 indexing range."); + const auto check_view = [](const Tensor& backing, uint32_t offset, uint64_t length) -> Status { + ORT_RETURN_IF_NOT(static_cast(offset) + length <= static_cast(backing.Shape().Size()), + "GatedDeltaNetStateReplay view exceeds its backing tensor."); + return Status::OK(); + }; + ORT_RETURN_IF_ERROR(check_view(*source, desc[0], state_elements)); + ORT_RETURN_IF_ERROR(check_view(*destination, desc[1], state_elements)); + ORT_RETURN_IF_ERROR(check_view(*capsule, desc[2], decay_elements)); + ORT_RETURN_IF_ERROR(check_view(*capsule, desc[3], key_elements)); + ORT_RETURN_IF_ERROR(check_view(*capsule, desc[4], delta_elements)); + ORT_RETURN_IF(destination->DataRaw() == source->DataRaw() || destination->DataRaw() == capsule->DataRaw(), + "GatedDeltaNetStateReplay destination backing must be distinct from source and capsule."); + + const auto& limits = context.DeviceLimits(); + const uint64_t groups = (state_elements - 1) / kWorkgroupSize + 1; + const uint64_t dispatch_x = std::min(groups, static_cast(limits.maxComputeWorkgroupsPerDimension)); + const uint64_t dispatch_y = (groups - 1) / dispatch_x + 1; + // Include 2D dispatch padding so the framework's flattened global index cannot wrap. + ORT_RETURN_IF_NOT(dispatch_y <= limits.maxComputeWorkgroupsPerDimension && + dispatch_x * dispatch_y * kWorkgroupSize <= kMaxIndex + 1, + "GatedDeltaNetStateReplay dispatch exceeds the supported indexing range."); + + auto* output = context.Output(0, destination->Shape()); + ORT_RETURN_IF_NOT(output && output->MutableDataRaw() == destination->DataRaw(), + "GatedDeltaNetStateReplay output must be prebound to destination_backing."); + // Do not bind destination_backing as a read-only shader input alongside its writable output. + return context.RunAndWaitChecked([&]() { + return EnqueueReplay(context, *source, *capsule, *output, desc, + static_cast(state_elements), static_cast(decay_elements), + static_cast(key_elements), static_cast(delta_elements), + static_cast(dispatch_x), static_cast(dispatch_y)); + }); +} + +} // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net_state_replay.h b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net_state_replay.h new file mode 100644 index 0000000000000..d0122a1cb9e90 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net_state_replay.h @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/webgpu_kernel.h" + +namespace onnxruntime::contrib::webgpu { + +class GatedDeltaNetStateReplay final : public onnxruntime::webgpu::WebGpuKernel { + public: + explicit GatedDeltaNetStateReplay(const OpKernelInfo& info) : WebGpuKernel(info) {} + + Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override; +}; + +} // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net_state_replay.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net_state_replay.wgsl.template new file mode 100644 index 0000000000000..7ec61fe2cb04f --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net_state_replay.wgsl.template @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#use .getByOffset .setByOffset + +$MAIN { + if (global_idx >= uniforms.state_elements) { + return; + } + let i = global_idx % uniforms.head_size_k; + let v = (global_idx / uniforms.head_size_k) % uniforms.head_size_v; + let h = global_idx / (uniforms.head_size_v * uniforms.head_size_k); + let hk = (h * uniforms.num_heads_k) / uniforms.num_heads_v; + var s = source.getByOffset(global_idx); + for (var t = 0u; t < uniforms.kept_count; t++) { + s *= decay.getByOffset(t * uniforms.num_heads_v + h); + s += key.getByOffset((t * uniforms.num_heads_k + hk) * uniforms.head_size_k + i) + * delta.getByOffset((t * uniforms.num_heads_v + h) * uniforms.head_size_v + v); + } + destination.setByOffset(global_idx, s); +} diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc index 93850dffc6d74..e25c593a149df 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc @@ -1158,9 +1158,9 @@ Status PackedSparseAttentionIndexer::ComputeQsa(onnxruntime::webgpu::ComputeCont ORT_RETURN_IF_ERROR(CheckStorageBindingLimit( context, "QSA selection", {&query_metadata, present_key_state, &rotary_cache, &sequence_metadata, present_state_lengths, &overflow_flags, - selected_indices, selected_counts, position_ids})); - PackedSparseAttentionIndexerQsaSelectProgram select{rotary.batched, position_ids != nullptr}; - select.CacheHint(rotary.batched, position_ids != nullptr) + selected_indices, selected_counts, position_ids})); + PackedSparseAttentionIndexerQsaSelectProgram select{rotary.batched, position_ids != nullptr}; + select.CacheHint(rotary.batched, position_ids != nullptr) .SetWorkgroupSize(kWorkgroupSize) .AddInputs({{&query_metadata, ProgramTensorMetadataDependency::Type}, {present_key_state, ProgramTensorMetadataDependency::Type}, diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.cc b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.cc index 548c24645d54f..f1fc78f849699 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.cc @@ -19,69 +19,100 @@ Status PackedSparseAttentionIndexerMergeProgram::GenerateShaderCode(ShaderHelper const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias); const auto& output_counts = shader.AddOutput("output_counts", ShaderUsage::UseValueTypeAlias); const auto& status = shader.AddOutput("status", ShaderUsage::UseValueTypeAlias); - const bool use_hash = hash_capacity_ <= 4096; - if (use_hash) { - shader.AdditionalImplementation() << "var keys: array;\n"; - } + shader.AdditionalImplementation() + << "var positions: array;\n" + << "var base_offset_shared: u32;\n" + << "var count_shared: u32;\n" + << "var start_shared: i32;\n" + << "var end_shared: i32;\n" + << "var result_shared: u32;\n" + << "var emitted_shared: u32;\n" + << "var overlap_shared: atomic;\n" + << "var invalid_shared: atomic;\n"; auto& body = shader.MainFunctionBody(); - body << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.queries") - << "let offset = global_idx * uniforms.capacity;\n" - << output_counts.SetByOffset("global_idx", "0") << status.SetByOffset("global_idx", "0") - << "for (var column = 0u; column < uniforms.capacity; column++) {\n" + body << "if (workgroup_idx >= uniforms.queries) { return; }\n" + << "let offset = workgroup_idx * uniforms.capacity;\n" + << "for (var column = local_idx; column < uniforms.capacity; column += workgroup_size_x) {\n" << output.SetByOffset("offset + column", "-1") << "}\n" - << "let row = " << rows.GetByOffset("global_idx") << ";\n" - << "let start = " << starts.GetByOffset("global_idx") << ";\n" - << "let end = " << ends.GetByOffset("global_idx") << ";\n" + << "storageBarrier();\n" + << "if (local_idx == 0u) {\n" + << output_counts.SetByOffset("workgroup_idx", "0") << status.SetByOffset("workgroup_idx", "0") + << "let row = " << rows.GetByOffset("workgroup_idx") << ";\n" + << "let start = " << starts.GetByOffset("workgroup_idx") << ";\n" + << "let end = " << ends.GetByOffset("workgroup_idx") << ";\n" << "if (row < 0 || u32(row) >= uniforms.rows || start < 0 || end < start) {\n" - << status.SetByOffset("global_idx", "1") << "return; }\n" + << "result_shared = 1u;\n" + << "} else {\n" << "let count = " << counts.GetByOffset("u32(row)") << ";\n" << "if (count < 0 || u32(count) > uniforms.base_capacity) {\n" - << status.SetByOffset("global_idx", "1") << "return; }\n" - << "let base_offset = u32(row) * uniforms.base_capacity;\n" - << "for (var column = 0u; column < u32(count); column++) { if (" - << base.GetByOffset("base_offset + column") << " < 0) {\n" - << status.SetByOffset("global_idx", "1") << "return; }}\n"; - if (use_hash) { - body << "for (var slot = 0u; slot < " << hash_capacity_ << "u; slot++) { keys[slot] = -1; }\n"; - } - body << "var emitted = 0u; var overlap = 0u; var overflow = false;\n" - << "for (var column = 0u; column < u32(count); column++) {\n" - << "let value = " << base.GetByOffset("base_offset + column") << ";\n" - << "var duplicate = false;\n"; - if (use_hash) { - body << "var slot = (u32(value) * 2654435761u) & " << hash_capacity_ - 1 << "u;\n" - << "for (var probe = 0u; probe < " << hash_capacity_ << "u; probe++) {\n" - << "if (keys[slot] == value) { duplicate = true; break; }\n" - << "if (keys[slot] == -1) { keys[slot] = value; break; }\n" - << "slot = (slot + 1u) & " << hash_capacity_ - 1 << "u; }\n"; - } else { - body << "for (var previous = 0u; previous < emitted; previous++) { if (" - << output.GetByOffset("offset + previous") << " == value) { duplicate = true; break; }}\n"; - } - body << "if (!duplicate) { if (emitted == uniforms.capacity) { overflow = true; break; }\n" - << output.SetByOffset("offset + emitted", "value") - << "emitted++; if (value >= start && value < end) { overlap++; }} }\n" - << "if (overflow || u32(end - start) - overlap > uniforms.capacity - emitted) {\n" - << "for (var column = 0u; column < uniforms.capacity; column++) {\n" + << "result_shared = 1u;\n" + << "} else {\n" + << "base_offset_shared = u32(row) * uniforms.base_capacity;\n" + << "count_shared = u32(count); start_shared = start; end_shared = end;\n" + << "}}}\n" + << "if (workgroupUniformLoad(&result_shared) != 0u) {\n" + << "if (local_idx == 0u) { " << status.SetByOffset("workgroup_idx", "1") << "} return; }\n" + << "let base_offset = workgroupUniformLoad(&base_offset_shared);\n" + << "let count = workgroupUniformLoad(&count_shared);\n" + << "let start = workgroupUniformLoad(&start_shared);\n" + << "let end = workgroupUniformLoad(&end_shared);\n" + << "for (var column = local_idx; column < count; column += workgroup_size_x) { if (" + << base.GetByOffset("base_offset + column") << " < 0) { atomicOr(&invalid_shared, 1u); }}\n" + << "workgroupBarrier();\n" + << "if (local_idx == 0u) { result_shared = atomicLoad(&invalid_shared); }\n" + << "if (workgroupUniformLoad(&result_shared) != 0u) {\n" + << "if (local_idx == 0u) { " << status.SetByOffset("workgroup_idx", "1") << "} return; }\n"; + + // Each chunk is compacted in input order; storage barriers publish its output to range membership scans. + auto emit_compaction = [&](bool base_phase) { + body << "positions[local_idx] = select(0u, 1u, keep);\n" + << "workgroupBarrier();\n" + << "if (local_idx == 0u) {\n" + << "var emitted = emitted_shared;\n" + << "for (var lane = 0u; lane < workgroup_size_x; lane++) {\n" + << "let flag = positions[lane]; positions[lane] = emitted; emitted += flag; }\n" + << "emitted_shared = emitted; }\n" + << "workgroupBarrier();\n" + << "if (keep && positions[local_idx] < uniforms.capacity) {\n" + << output.SetByOffset("offset + positions[local_idx]", "value"); + if (base_phase) { + body << "if (value >= start && value < end) { atomicAdd(&overlap_shared, 1u); }\n"; + } + body << "}\n" + << "storageBarrier();\n" + << "workgroupBarrier();\n"; + }; + + body << "for (var chunk = 0u; chunk < count; chunk += workgroup_size_x) {\n" + << "let column = chunk + local_idx;\n" + << "var value = 0i; var keep = column < count;\n" + << "if (keep) {\n" + << "value = " << base.GetByOffset("base_offset + column") << ";\n" + << "for (var previous = 0u; previous < column; previous++) { if (" + << base.GetByOffset("base_offset + previous") << " == value) { keep = false; break; }}}\n"; + emit_compaction(true); + body << "if (workgroupUniformLoad(&emitted_shared) > uniforms.capacity) { break; }}\n" + << "let base_emitted = workgroupUniformLoad(&emitted_shared);\n" + << "if (local_idx == 0u) {\n" + << "if (base_emitted > uniforms.capacity || " + "u32(end - start) - atomicLoad(&overlap_shared) > uniforms.capacity - base_emitted) {\n" + << "result_shared = 2u; }}\n" + << "if (workgroupUniformLoad(&result_shared) != 0u) {\n" + << "for (var column = local_idx; column < uniforms.capacity; column += workgroup_size_x) {\n" << output.SetByOffset("offset + column", "-1") << "}\n" - << status.SetByOffset("global_idx", "2") << "return; }\n" - << "let base_emitted = emitted;\n" - << "for (var value = start; value < end; value++) {\n" - << "var duplicate = false;\n"; - if (use_hash) { - body << "var slot = (u32(value) * 2654435761u) & " << hash_capacity_ - 1 << "u;\n" - << "for (var probe = 0u; probe < " << hash_capacity_ << "u; probe++) {\n" - << "if (keys[slot] == value) { duplicate = true; break; }\n" - << "if (keys[slot] == -1) { break; }\n" - << "slot = (slot + 1u) & " << hash_capacity_ - 1 << "u; }\n"; - } else { - body << "for (var previous = 0u; previous < base_emitted; previous++) { if (" - << output.GetByOffset("offset + previous") << " == value) { duplicate = true; break; }}\n"; - } - body << "if (!duplicate) {\n" - << output.SetByOffset("offset + emitted", "value") - << "emitted++; }}\n" - << output_counts.SetByOffset("global_idx", "i32(emitted)"); + << "if (local_idx == 0u) { " << status.SetByOffset("workgroup_idx", "2") << "} return; }\n" + << "let range_length = u32(end - start);\n" + << "for (var chunk = 0u; chunk < range_length; chunk += workgroup_size_x) {\n" + << "let column = chunk + local_idx;\n" + << "var value = 0i; var keep = column < range_length;\n" + << "if (keep) {\n" + << "value = start + i32(column);\n" + << "for (var previous = 0u; previous < base_emitted; previous++) { if (" + << output.GetByOffset("offset + previous") << " == value) { keep = false; break; }}}\n"; + emit_compaction(false); + body << "}\n" + << "if (local_idx == 0u) {\n" + << output_counts.SetByOffset("workgroup_idx", "i32(emitted_shared)") << "}\n"; return Status::OK(); } @@ -106,10 +137,9 @@ Status PackedSparseAttentionIndexerMerge::ComputeInternal(ComputeContext& contex } const auto* base_binding = base->Shape().Size() == 0 ? &empty_binding.value() : base; const auto* counts_binding = counts->Shape().Size() == 0 ? &empty_binding.value() : counts; - PackedSparseAttentionIndexerMergeProgram program(dimensions.hash_capacity); - program.SetWorkgroupSize(1) + PackedSparseAttentionIndexerMergeProgram program; + program.SetWorkgroupSize(256) .SetDispatchGroupSize(dimensions.queries) - .CacheHint(std::to_string(dimensions.hash_capacity)) .AddInputs({{base_binding, ProgramTensorMetadataDependency::Type}, {counts_binding, ProgramTensorMetadataDependency::Type}, {rows, ProgramTensorMetadataDependency::Type}, diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.h b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.h index 1ce92ba4b7d56..404ddef014e2b 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.h +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.h @@ -11,16 +11,12 @@ using onnxruntime::webgpu::ComputeContext; class PackedSparseAttentionIndexerMergeProgram final : public Program { public: - explicit PackedSparseAttentionIndexerMergeProgram(int32_t hash_capacity) - : Program{"PackedSparseAttentionIndexerMerge"}, hash_capacity_(hash_capacity) {} + PackedSparseAttentionIndexerMergeProgram() : Program{"PackedSparseAttentionIndexerMerge"} {} Status GenerateShaderCode(ShaderHelper& shader) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"rows", ProgramUniformVariableDataType::Uint32}, {"base_capacity", ProgramUniformVariableDataType::Uint32}, {"queries", ProgramUniformVariableDataType::Uint32}, {"capacity", ProgramUniformVariableDataType::Uint32}); - - private: - int32_t hash_capacity_; }; class PackedSparseAttentionIndexerMerge final : public WebGpuKernel { diff --git a/onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc index 6fa19445edc87..dc12f90a557c5 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc @@ -833,10 +833,10 @@ Status PagedAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& cont query->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; const bool use_direct_paged_prefill = !has_local_window && head_sink == nullptr && - ShouldRunFusedPagedPrefill(context, is_fp16_q, - static_cast(max_seqlen_q), - parameters.head_size, - parameters.block_size); + CanApplyFlashAttentionPagedPrefill(context, is_fp16_q, + static_cast(max_seqlen_q), + parameters.head_size, + parameters.block_size); const bool use_direct_paged_attention = use_direct_paged_decode || use_direct_paged_prefill; // Unpack/Repack fast paths (see full explanation below at the Q/output view // construction). skip_unpack_repack must be computed before the @@ -945,7 +945,7 @@ Status PagedAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& cont // [token_count, 1, N, H] view and cumulative_seqlens_q. Only the // fused paged-prefill shader knows how to consume this // (q_varlen #param), so mode (b) is only safe when - // ShouldRunFusedPagedPrefill would return true — which is exactly + // CanApplyFlashAttentionPagedPrefill would return true — which is exactly // use_direct_paged_prefill (see storage-binding block above). // // If neither mode applies, we fall back to Unpack + BSNH-padded scratch. diff --git a/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h b/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h index f34579543a4cb..0fbfd764eb276 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h +++ b/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h @@ -32,9 +32,9 @@ class SparseAttentionIndexerQsaConcatProgram final {"sequence_length", ProgramUniformVariableDataType::Uint32}, {"past_sequence_length", ProgramUniformVariableDataType::Uint32}, {"total_sequence_length", ProgramUniformVariableDataType::Uint32}, - {"head_size", ProgramUniformVariableDataType::Uint32}, - {"key_row_stride", ProgramUniformVariableDataType::Uint32}, - {"key_offset", ProgramUniformVariableDataType::Uint32}); + {"head_size", ProgramUniformVariableDataType::Uint32}, + {"key_row_stride", ProgramUniformVariableDataType::Uint32}, + {"key_offset", ProgramUniformVariableDataType::Uint32}); private: bool has_past_; @@ -52,7 +52,7 @@ class SparseAttentionIndexerQsaSelectProgram final {"sequence_length", ProgramUniformVariableDataType::Uint32}, {"num_heads", ProgramUniformVariableDataType::Uint32}, {"head_size", ProgramUniformVariableDataType::Uint32}, - {"query_row_stride", ProgramUniformVariableDataType::Uint32}, + {"query_row_stride", ProgramUniformVariableDataType::Uint32}, {"rotary_width", ProgramUniformVariableDataType::Uint32}, {"max_rotary_length", ProgramUniformVariableDataType::Uint32}, {"rotary_cache_batch_stride", ProgramUniformVariableDataType::Uint32}, diff --git a/onnxruntime/contrib_ops/webgpu/fused_conv.cc b/onnxruntime/contrib_ops/webgpu/fused_conv.cc index e6b7ac3ec24d4..e4b59bde631bd 100644 --- a/onnxruntime/contrib_ops/webgpu/fused_conv.cc +++ b/onnxruntime/contrib_ops/webgpu/fused_conv.cc @@ -14,9 +14,8 @@ using onnxruntime::webgpu::Conv; template class FusedConv final : public Conv { public: - FusedConv(const OpKernelInfo& info) : Conv(info) { - ORT_ENFORCE(GetFusedActivationAttr(info, Conv::activation_).IsOK()); - } + // The Conv<_, is_fused = true> base already parses the activation attribute. + FusedConv(const OpKernelInfo& info) : Conv(info) {} }; ONNX_OPERATOR_KERNEL_EX( diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index f6e691cb1d7f4..93182e9e0fdbf 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -136,6 +136,7 @@ Status ApplyBlockFp8ExpertMatMul(ComputeContext& context, const uint32_t scale_k_blocks = (inner - 1) / 128 + 1; uint32_t byte_offset = 0; const bool use_matrix = !indirect_experts && rows >= 128 && cols >= 256 && inner >= 256 && + context.HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) && BlockFp8MatrixDispatchFits(rows, cols, context.DeviceLimits().maxComputeWorkgroupsPerDimension) && input->DataType() == DataTypeImpl::GetType() && @@ -144,6 +145,7 @@ Status ApplyBlockFp8ExpertMatMul(ComputeContext& context, 16, 16, 16, 32, false}}) .has_value(); BlockFp8ExpertMatMulProgram program{bias != nullptr, indirect_experts != nullptr, broadcast_input, use_matrix}; + program.SetUseSubgroupId(use_matrix); program.AddInputs({{input, ProgramTensorMetadataDependency::Type}}); if (indirect_experts) { program.AddInputs({{&raw_weights, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}}) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template new file mode 100644 index 0000000000000..64387a727c0e4 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_activation.wgsl.template @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#use .getByOffset .setByOffset + +const E4M3_MAX: f32 = 448.0; + +$MAIN { + if (global_idx >= uniforms.size) { + return; + } + let scale = input_scale.getByOffset(0u); + let inverse = select(0.0, 1.0 / scale, scale != 0.0); + let input_value = f32(input_a.getByOffset(global_idx)); + let v = input_value * inverse; + // Classify bits so finite-math optimizations cannot fold a NaN self-comparison away. + if ((bitcast(v) & 0x7fffffffu) > 0x7f800000u) { + output.setByOffset(global_idx, f16(v)); + return; + } + let magnitude = min(abs(v), E4M3_MAX); + let magnitude_exponent = (bitcast(magnitude) >> 23u) & 255u; + let step = bitcast((max(magnitude_exponent, 121u) - 3u) << 23u); + let units = magnitude / step; + let lower = floor(units); + let halfway = (lower + 0.5) * step; + let round_up = magnitude > halfway || + (magnitude == halfway && (u32(lower) & 1u) != 0u); + let rounded = min((lower + select(0.0, 1.0, round_up)) * step, E4M3_MAX); + let q = select(rounded, -rounded, v < 0.0); + output.setByOffset(global_idx, f16(q * scale)); +} diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc new file mode 100644 index 0000000000000..75c5c8cde31d2 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.cc @@ -0,0 +1,218 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.h" + +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "core/common/narrow.h" +#include "core/providers/cpu/math/matmul_helper.h" +#include "core/providers/webgpu/math/subgroup_matrix_config.h" +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" +#include "core/providers/webgpu/webgpu_utils.h" +#if !defined(DISABLE_FLOAT8_TYPES) +#include "core/common/float8.h" +#endif + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +using namespace onnxruntime::webgpu; + +namespace { + +class Fp8ActivationProgram final : public Program { + public: + Fp8ActivationProgram() : Program{"Fp8MatMulActivationQdq"} {} + + Status GenerateShaderCode(ShaderHelper& shader) const override { + const auto& input_a = shader.AddInput("input_a", ShaderUsage::UseValueTypeAlias); + const auto& input_scale = shader.AddInput("input_scale", ShaderUsage::UseValueTypeAlias); + const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias); + return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_block_quantized_fp8_activation.wgsl.template", + WGSL_TEMPLATE_VARIABLE(input_a, input_a), + WGSL_TEMPLATE_VARIABLE(input_scale, input_scale), + WGSL_TEMPLATE_VARIABLE(output, output)); + } + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32}); +}; + +class Fp8EmptyReductionProgram final : public Program { + public: + explicit Fp8EmptyReductionProgram(bool has_bias) + : Program{"Fp8MatMulEmptyReduction"}, has_bias_(has_bias) {} + + Status GenerateShaderCode(ShaderHelper& shader) const override { + if (has_bias_) { + const auto& bias = shader.AddInput("bias"); + const auto& output = shader.AddOutput("output"); + shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.size") + << output.SetByOffset("global_idx", bias.GetByOffset("global_idx % uniforms.N")); + } else { + const auto& output = shader.AddOutput("output"); + shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.size") + << output.SetByOffset("global_idx", "f16(0)"); + } + return Status::OK(); + } + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}); + + private: + bool has_bias_; +}; + +class Fp8MatMulProgram final : public Program { + public: + Fp8MatMulProgram(bool has_bias, bool use_matrix) + : Program{"MatMulBlockQuantizedFp8Weight"}, has_bias_(has_bias), use_matrix_(use_matrix) {} + + Status GenerateShaderCode(ShaderHelper& shader) const override { + const auto& input_a = shader.AddInput("input_a", ShaderUsage::UseValueTypeAlias); + const auto& weights = shader.AddInput("weights", ShaderUsage::UseValueTypeAlias); + const auto& scales = shader.AddInput("scales", ShaderUsage::UseValueTypeAlias); + const ShaderVariableHelper* bias = nullptr; + if (has_bias_) { + bias = &shader.AddInput("bias", ShaderUsage::UseValueTypeAlias); + } + const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias); + return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_block_quantized_fp8_weight.wgsl.template", + WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(use_matrix, use_matrix_), + WGSL_TEMPLATE_OPTIONAL_VARIABLE(bias, bias), + WGSL_TEMPLATE_VARIABLE(input_a, input_a), + WGSL_TEMPLATE_VARIABLE(output, output), + WGSL_TEMPLATE_VARIABLE(scales, scales), + WGSL_TEMPLATE_VARIABLE(weights, weights)); + } + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"K", ProgramUniformVariableDataType::Uint32}, + {"K_blocks", ProgramUniformVariableDataType::Uint32}, + {"block_size", ProgramUniformVariableDataType::Uint32}); + + private: + bool has_bias_; + bool use_matrix_; +}; + +} // namespace + +MatMulBlockQuantizedFp8Weight::MatMulBlockQuantizedFp8Weight(const OpKernelInfo& info) + : WebGpuKernel(info), + block_size_(narrow(info.GetAttrOrDefault("block_size", 128))) { + ORT_ENFORCE(block_size_ > 0, "block_size must be positive."); +} + +Status MatMulBlockQuantizedFp8Weight::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + const Tensor* a = context.Input(0); + const Tensor* b = context.Input(1); + const Tensor* scales = context.Input(2); + const Tensor* a_scale = context.Input(3); + const Tensor* bias = context.Input(4); + + const auto& a_shape = a->Shape(); + const auto& b_shape = b->Shape(); + ORT_RETURN_IF_NOT(a_shape.NumDimensions() >= 1 && b_shape.NumDimensions() == 2, + "A must have rank >= 1 and B must have shape [N, K]."); + const int64_t n_dim = b_shape[0]; + const int64_t k_dim = b_shape[1]; + ORT_RETURN_IF_NOT(a_shape[a_shape.NumDimensions() - 1] == k_dim, + "A and B contraction dimensions must match."); + const int64_t k_blocks = (k_dim + block_size_ - 1) / block_size_; + ORT_RETURN_IF_NOT(scales->Shape() == TensorShape({n_dim, k_blocks}), + "b_scale must have shape [N, ceil(K/block_size)]."); + ORT_RETURN_IF_NOT(a_scale == nullptr || a_scale->Shape().Size() == 1, + "a_scale must contain one value."); + ORT_RETURN_IF_NOT(bias == nullptr || bias->Shape() == TensorShape({n_dim}), + "bias must have shape [N]."); + + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(a_shape, b_shape, false, true)); + Tensor* y = context.Output(0, helper.OutputShape()); + if (y->Shape().Size() == 0) { + return Status::OK(); + } + + const uint32_t K = narrow(k_dim); + const uint32_t N = narrow(n_dim); + const uint32_t M = narrow(y->Shape().Size() / n_dim); + const uint32_t K_blocks = narrow(k_blocks); + if (K == 0) { + Fp8EmptyReductionProgram empty_program{bias != nullptr}; + empty_program.SetWorkgroupSize(64); + empty_program.SetDispatchGroupSize((narrow(y->Shape().Size()) + 63u) / 64u); + if (bias) { + empty_program.AddInput({bias, ProgramTensorMetadataDependency::Type}); + } + empty_program.AddOutput({y, ProgramTensorMetadataDependency::Type}) + .AddUniformVariables({{narrow(y->Shape().Size())}, {N}}); + return context.RunProgram(empty_program); + } + const Tensor* activation = a; + Tensor qdq; + if (a_scale) { + qdq = context.CreateGPUTensor(a->DataType(), a_shape); + Fp8ActivationProgram qdq_program; + qdq_program.SetWorkgroupSize(64); + qdq_program.SetDispatchGroupSize((narrow(a_shape.Size()) + 63u) / 64u); + qdq_program.AddInputs({{a, ProgramTensorMetadataDependency::Type}, + {a_scale, ProgramTensorMetadataDependency::Type}}) + .AddOutput({&qdq, ProgramTensorMetadataDependency::Type}) + .AddUniformVariables({{narrow(a_shape.Size())}}); + ORT_RETURN_IF_ERROR(context.RunProgram(qdq_program)); + activation = &qdq; + } + + const bool use_matrix = M >= 8 && + context.HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) && + SelectSubgroupMatrixConfig(context, {{wgpu::SubgroupMatrixComponentType::F16, + wgpu::SubgroupMatrixComponentType::F32, + 16, 16, 16, 32, false}}) + .has_value(); + + auto memory_info = OrtMemoryInfo{ + WEBGPU_BUFFER, + OrtDeviceAllocator, + OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NONE, 0}}; + // Tensor views require void*, but the alias only binds B as read-only storage; no copy or mutation occurs. + Tensor weight_bytes(DataTypeImpl::GetType(), b_shape, const_cast(b->DataRaw()), memory_info); + Fp8MatMulProgram program{bias != nullptr, use_matrix}; + program.SetUseSubgroupId(use_matrix); + program.SetWorkgroupSize(use_matrix ? 128 : 32); + if (use_matrix && context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { + program.SetSubgroupSize(32); + } + program.SetDispatchGroupSize(use_matrix ? (N + 15u) / 16u : N, + use_matrix ? (M + 63u) / 64u : M); + // Flatten rounds the byte count up to u32 words; WebGPU buffers are 16-byte padded by BufferManager. + // The shader extracts only bytes whose offset is within N*K, even in the final partial word. + program.AddInputs({{activation, ProgramTensorMetadataDependency::Type}, + {&weight_bytes, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}, + {scales, ProgramTensorMetadataDependency::Type}}) + .AddOutput({y, ProgramTensorMetadataDependency::Type}) + .AddUniformVariables({{M}, {N}, {K}, {K_blocks}, {block_size_}}); + if (bias) { + program.AddInput({bias, ProgramTensorMetadataDependency::Type}); + } + return context.RunProgram(program); +} + +#if !defined(DISABLE_FLOAT8_TYPES) +ONNX_OPERATOR_KERNEL_EX( + MatMulBlockQuantizedFp8Weight, + kMSDomain, + 1, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()), + MatMulBlockQuantizedFp8Weight); +#endif + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.h new file mode 100644 index 0000000000000..db8954e5d5670 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.h @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/program.h" +#include "core/providers/webgpu/webgpu_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +class MatMulBlockQuantizedFp8Weight final : public onnxruntime::webgpu::WebGpuKernel { + public: + explicit MatMulBlockQuantizedFp8Weight(const OpKernelInfo& info); + Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override; + + private: + uint32_t block_size_; +}; + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template new file mode 100644 index 0000000000000..ba0321cf88832 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_block_quantized_fp8_weight.wgsl.template @@ -0,0 +1,107 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#param has_bias +#param use_matrix +#use .getByOffset .setByOffset + +fn decode_fp8(code: u32) -> f32 { + let exponent = (code >> 3u) & 15u; + let mantissa = code & 7u; + let magnitude = select(f32(mantissa) * 0.001953125, + bitcast(((exponent + 120u) << 23u) | (mantissa << 20u)), + exponent != 0u); + let value = select(magnitude, -magnitude, (code & 128u) != 0u); + return select(value, bitcast(0x7fc00000u | code), exponent == 15u && mantissa == 7u); +} + +fn weight(n: u32, k: u32) -> f16 { + let offset = n * uniforms.K + k; + let code = (weights.getByOffset(offset / 4u) >> ((offset % 4u) * 8u)) & 255u; + return f16(decode_fp8(code) * scales.getByOffset(n * uniforms.K_blocks + k / uniforms.block_size)); +} + +#if use_matrix +// The padded workgroup tiles make every subgroup-matrix load and store in bounds. +var tile_a: array; +var tile_b: array; +var tile_y: array; +#else +var partial: array; +#endif +$MAIN { +#if use_matrix + let lane = local_idx; + let n_tiles = (uniforms.N + 15u) / 16u; + if (workgroup_idx >= n_tiles * ((uniforms.M + 63u) / 64u)) { + return; + } + let m_base = (workgroup_idx / n_tiles) * 64u; + let n_base = (workgroup_idx % n_tiles) * 16u; + var acc: subgroup_matrix_result; + for (var tile = 0u; tile < (uniforms.K + 15u) / 16u; tile++) { + let k_base = tile * 16u; + for (var i = lane; i < 1024u; i += 128u) { + let m = m_base + i / 16u; + let k = k_base + i % 16u; + tile_a[i] = f16(0); + if (m < uniforms.M && k < uniforms.K) { + tile_a[i] = input_a.getByOffset(m * uniforms.K + k); + } + } + for (var i = lane; i < 256u; i += 128u) { + let n = n_base + i / 16u; + let k = k_base + i % 16u; + tile_b[i] = f16(0); + if (n < uniforms.N && k < uniforms.K) { + tile_b[i] = weight(n, k); + } + } + workgroupBarrier(); + let left = subgroupMatrixLoad, row_major>( + &tile_a, subgroup_index * 256u, 16u); + let right = subgroupMatrixLoad, col_major>(&tile_b, 0u, 16u); + acc = subgroupMatrixMultiplyAccumulate(left, right, acc); + workgroupBarrier(); + } + subgroupMatrixStore(&tile_y, subgroup_index * 256u, acc, 16u); + workgroupBarrier(); + for (var i = lane; i < 1024u; i += 128u) { + let m = m_base + i / 16u; + let n = n_base + i % 16u; + if (m < uniforms.M && n < uniforms.N) { + var value = tile_y[i]; +#if has_bias + value += f32(bias.getByOffset(n)); +#endif + output.setByOffset(m * uniforms.N + n, f16(value)); + } + } +#else + let lane = local_idx; + if (workgroup_idx >= uniforms.M * uniforms.N) { + return; + } + let m = workgroup_idx / uniforms.N; + let n = workgroup_idx % uniforms.N; + var sum = 0.0; + for (var k = lane; k < uniforms.K; k += 32u) { + sum += f32(input_a.getByOffset(m * uniforms.K + k)) * f32(weight(n, k)); + } + partial[lane] = sum; + workgroupBarrier(); + for (var stride = 16u; stride > 0u; stride /= 2u) { + if (lane < stride) { + partial[lane] += partial[lane + stride]; + } + workgroupBarrier(); + } + if (lane == 0u) { + var value = partial[0]; +#if has_bias + value += f32(bias.getByOffset(n)); +#endif + output.setByOffset(m * uniforms.N + n, f16(value)); + } +#endif +} diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc index a51cbec200cf8..8f6b9f6763108 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc @@ -4,7 +4,6 @@ #include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "core/providers/webgpu/math/subgroup_matrix_config.h" -#include "core/providers/webgpu/vendor/intel/intel_device_info.h" #include "core/providers/webgpu/webgpu_utils.h" namespace onnxruntime { @@ -62,8 +61,10 @@ Status GenerateShaderCode16x16x16(ShaderHelper& shader, const ShaderVariableHelper& b, const ShaderVariableHelper& scales_b, const ShaderVariableHelper& output, - uint32_t nbits, const SubgroupMatrixConfig& config, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect) { - // Use the 128x128 tile shader for the 16x16x16 config. + uint32_t nbits, const SubgroupMatrixConfig& config, uint32_t tile_size_m, + uint32_t tile_size_n, uint32_t tile_size_k, bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect) { + // Use the configurable tiled shader for the 16x16x16 configs. return WGSL_TEMPLATE_APPLY(shader, "quantization/subgroup_matrix_matmul_nbits_16x16x16_128.wgsl.template", WGSL_TEMPLATE_PARAMETER(has_bias, has_bias), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx), @@ -74,6 +75,9 @@ Status GenerateShaderCode16x16x16(ShaderHelper& shader, WGSL_TEMPLATE_PARAMETER(sg_mat_k, config.K), WGSL_TEMPLATE_PARAMETER(sg_mat_m, config.M), WGSL_TEMPLATE_PARAMETER(sg_mat_n, config.N), + WGSL_TEMPLATE_PARAMETER(tile_size_k, tile_size_k), + WGSL_TEMPLATE_PARAMETER(tile_size_m, tile_size_m), + WGSL_TEMPLATE_PARAMETER(tile_size_n, tile_size_n), WGSL_TEMPLATE_VARIABLE(input_b, b), WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); @@ -84,7 +88,7 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelper& scales_b, const ShaderVariableHelper& output, uint32_t nbits, const SubgroupMatrixConfig& config, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, - bool has_tail_buffer) { + uint32_t tile_m, uint32_t tile_n, bool has_tail_buffer) { return WGSL_TEMPLATE_APPLY(shader, "quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template", WGSL_TEMPLATE_PARAMETER(has_bias, has_bias), WGSL_TEMPLATE_PARAMETER(has_tail_buffer, has_tail_buffer), @@ -96,6 +100,8 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, WGSL_TEMPLATE_PARAMETER(sg_mat_k, config.K), WGSL_TEMPLATE_PARAMETER(sg_mat_m, config.M), WGSL_TEMPLATE_PARAMETER(sg_mat_n, config.N), + WGSL_TEMPLATE_PARAMETER(tile_m, tile_m), + WGSL_TEMPLATE_PARAMETER(tile_n, tile_n), WGSL_TEMPLATE_VARIABLE(input_b, b), WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); @@ -138,9 +144,12 @@ Status SubgroupMatrixMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader if (config_.Is(8, 8, 8)) { return GenerateShaderCode8x8x8(shader, a, b, scales_b, output, nbits_, has_zero_points_, has_bias_, has_weight_idx_, has_weight_idx_indirect_); } else if (config_.Is(8, 16, 16)) { - return GenerateShaderCode8x16x16(shader, b, scales_b, output, nbits_, config_, has_zero_points_, has_bias_, has_weight_idx_, has_weight_idx_indirect_, has_tail_buffer_); + return GenerateShaderCode8x16x16(shader, b, scales_b, output, nbits_, config_, has_zero_points_, has_bias_, has_weight_idx_, has_weight_idx_indirect_, tile_m_, tile_n_, has_tail_buffer_); } else if (config_.Is(16, 16, 16)) { - return GenerateShaderCode16x16x16(shader, b, scales_b, output, nbits_, config_, has_zero_points_, has_bias_, has_weight_idx_, has_weight_idx_indirect_); + return GenerateShaderCode16x16x16(shader, b, scales_b, output, nbits_, config_, tile_m_, + tile_n_, tile_k_, + has_zero_points_, has_bias_, has_weight_idx_, + has_weight_idx_indirect_); } else { return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::NOT_IMPLEMENTED, "Unsupported subgroup matrix config dimensions."); @@ -178,20 +187,72 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te Tensor* y, const uint32_t weight_index, const Tensor* weight_index_indirect) { + const bool has_zero_points = zero_points != nullptr; + const bool has_bias = bias != nullptr; + const bool has_weight_idx_indirect = weight_index_indirect != nullptr; + const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; + // Determine tile sizes first (needed for prepack padding). - uint32_t tile_size_a = 32; - uint32_t tile_size_b = 64; - uint32_t work_group_size = 128; + uint32_t tile_m = 32; + uint32_t tile_n = 64; + uint32_t tile_size_k = 32; + uint32_t workgroup_size = 128; if (config.Is(8, 16, 16)) { - // 8x16x16 config: 8 subgroups, 256 threads, 64x64 tiles - tile_size_a = 64; - work_group_size = 256; + // Default: 64x64 tile size, 256 thread. + tile_m = 64; + tile_n = 64; + workgroup_size = 256; + // Bias output uses a scratch tile, so keep it at 64x64 for workgroup memory limits. + if (!has_bias) { + if (M >= 2048 && N % 256 == 0) { + // Large M (M>=2048): 128x256 tile size, 512 threads. + tile_m = 128; + tile_n = 256; + workgroup_size = 512; + } else if (M >= 256 && N % 128 == 0) { + // Mid-sized M (256<=M<2048): 128x128 tile size, 512 threads. + tile_m = 128; + tile_n = 128; + workgroup_size = 512; + } + } } else if (config.Is(16, 16, 16)) { - // 16x16x16 config: 4 subgroups, 128 threads, 128x128 tiles - tile_size_a = 128; - tile_size_b = 128; - work_group_size = 128; + if (config.subgroupSize == 64) { + // A wave64 subgroup computes a 64x64 micro-tile. Small M distributes waves along N; + // tall tiles share each dequantized B tile across row waves. + if (M <= 64) { + tile_m = 64; + // Avoid a fully padded second column wave when N is exactly 64. + tile_n = N == 64 ? 64 : 128; + tile_size_k = 64; + } else { + // Use four waves when that does not add M padding over a two-wave tile and M is large + // enough to amortize the larger workgroup. + const uint32_t row_tiles_128 = M / 128 + static_cast(M % 128 != 0); + const bool use_four_waves = M > 384 && row_tiles_128 % 2 == 0; + tile_m = use_four_waves ? 256 : 128; + tile_n = 64; + + // Larger K blocks reduce barriers. Four-wave tiles retain enough occupancy for BK=128; + // larger two-wave grids use BK=64, while small two-wave grids favor BK=32. + tile_size_k = M < 256 ? 32 : (use_four_waves || M <= 384 ? 128 : 64); + } + // The K loop and cooperative-matrix loads do not mask a partial BK tile. K is guaranteed + // to be divisible by 32, so step down until BK divides it exactly. + if (K % tile_size_k != 0) { + tile_size_k = tile_size_k == 128 && K % 64 == 0 ? 64 : 32; + } + workgroup_size = (tile_m / 64) * (tile_n / 64) * config.subgroupSize; + } else { + // Wave32: four subgroups in a 2x2 grid. + tile_m = 128; + tile_n = 128; + workgroup_size = 4 * config.subgroupSize; + } } + ORT_ENFORCE(workgroup_size <= context.DeviceLimits().maxComputeWorkgroupSizeX && + workgroup_size <= context.DeviceLimits().maxComputeInvocationsPerWorkgroup, + "Subgroup matrix MatMulNBits workgroup size exceeds device limits: ", workgroup_size); // If applicable, layout optimization of input matrix A(MxK) can be used for SubgroupMatrixLoad. Tensor a_prepack; @@ -207,7 +268,7 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te } // Pad M to workgroup tile size so all subgroups read valid prepacked data. - const uint32_t padded_M = ((M + tile_size_a - 1) / tile_size_a) * tile_size_a; + const uint32_t padded_M = CeilDiv(M, tile_m) * tile_m; const auto dispatch_group_size_x = padded_M / m; ORT_ENFORCE(K % k == 0, "K must be a multiple of ", k); const auto dispatch_group_size_y = K / k; @@ -226,66 +287,54 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te constexpr uint32_t kU32Components = 4; TensorShape y_shape{1, M, N}; - const bool has_zero_points = zero_points != nullptr; - const bool has_bias = bias != nullptr; - const bool has_weight_idx_indirect = weight_index_indirect != nullptr; - const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; // For the no-bias 8x16x16 path, a trailing partial M-tile is written via the same - // branch-free fast path as an in-bounds tile, but into a small [tile_size_a, N] + // branch-free fast path as an in-bounds tile, but into a small [tile_m, N] // padded buffer instead of `y` (which is too short by this point); the crop-copy // dispatch below moves the valid rows into `y`. This keeps the no-bias write-out // free of any bounds-checked workgroup-scratch store. - const bool has_tail_buffer = !has_bias && config.Is(8, 16, 16) && (M % tile_size_a != 0); - SubgroupMatrixMatMulNBitsProgram mul_program{nbits, config, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, has_tail_buffer}; - mul_program.SetWorkgroupSize(work_group_size); + const bool has_tail_buffer = !has_bias && config.Is(8, 16, 16) && (M % tile_m != 0); + SubgroupMatrixMatMulNBitsProgram matmul_program{nbits, config, tile_m, tile_n, tile_size_k, + has_zero_points, has_bias, has_weight_idx, + has_weight_idx_indirect, has_tail_buffer}; + matmul_program.SetWorkgroupSize(workgroup_size); // Pin kernels running on variable-size adapters to the subgroup size they were written for. if (context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { - mul_program.SetSubgroupSize(config.subgroupSize); + matmul_program.SetSubgroupSize(config.subgroupSize); } - uint32_t dispatch_x = (N + tile_size_b - 1) / tile_size_b; - uint32_t num_m_tiles = (M + tile_size_a - 1) / tile_size_a; - uint32_t dispatch_y = num_m_tiles; - // For large M on Intel Xe, cap dispatch_y so each workgroup processes multiple - // M-tiles sequentially, reducing scheduling overhead. - if (M > 2048 && context.AdapterInfo().vendor == std::string_view{"intel"}) { - const uint32_t hw_subgroups = - ::onnxruntime::webgpu::intel::HwSubgroups(std::string_view{context.AdapterInfo().architecture}); - if (hw_subgroups > 0) { - constexpr uint32_t kOccupancyFactor = 16; // empirically tuned on Xe2/Xe3 devices - uint32_t target_wgs = hw_subgroups * kOccupancyFactor / (work_group_size / 32); - dispatch_y = std::min(dispatch_y, (target_wgs + dispatch_x - 1) / dispatch_x); - } - } - uint32_t m_tiles_per_wg = (num_m_tiles + dispatch_y - 1) / dispatch_y; - mul_program.SetDispatchGroupSize(dispatch_x, dispatch_y, 1); - mul_program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(nbits == 4 ? kU32Components : 2 * kU32Components)}, - {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) - .AddUniformVariables({{M}, {N}, {K}, {zero_blocks_per_col}, {weight_index}, {m_tiles_per_wg}}) + uint32_t num_N_tile = CeilDiv(N, tile_n); + uint32_t num_M_tile = CeilDiv(M, tile_m); + matmul_program.SetDispatchGroupSize(num_N_tile, num_M_tile, 1); + + const int input_b_components = static_cast(nbits == 4 ? kU32Components : 2 * kU32Components); + matmul_program.AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}, + {b, ProgramTensorMetadataDependency::TypeAndRank, input_b_components}, + {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + .AddUniformVariables({{M}, {N}, {K}, {zero_blocks_per_col}, {num_N_tile}, {num_M_tile}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, y_shape, 1}) .CacheHint(nbits, static_cast(config.componentType), static_cast(config.resultComponentType), config.M, config.N, config.K, config.subgroupSize, - has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, has_tail_buffer); + tile_m, tile_n, tile_size_k, has_zero_points, has_bias, has_weight_idx, + has_weight_idx_indirect, has_tail_buffer); if (has_zero_points) { - mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); + matmul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } if (bias) { - mul_program.AddInput({bias, ProgramTensorMetadataDependency::None}); + matmul_program.AddInput({bias, ProgramTensorMetadataDependency::None}); } if (has_weight_idx_indirect) { - mul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); + matmul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } Tensor tail_buffer; if (has_tail_buffer) { - tail_buffer = context.CreateGPUTensor(y->DataType(), TensorShape{tile_size_a, N}); - mul_program.AddOutput({&tail_buffer, ProgramTensorMetadataDependency::None}); + tail_buffer = context.CreateGPUTensor(y->DataType(), TensorShape{tile_m, N}); + matmul_program.AddOutput({&tail_buffer, ProgramTensorMetadataDependency::None}); } - ORT_RETURN_IF_ERROR(context.RunProgram(mul_program)); + ORT_RETURN_IF_ERROR(context.RunProgram(matmul_program)); if (has_tail_buffer) { // Only the rows below `tail_rows` in `tail_buffer` were ever written by the @@ -293,8 +342,8 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te // fixed); crop-copy just those into the real output. constexpr uint32_t kTailCopyComponents = 4; ORT_ENFORCE(N % kTailCopyComponents == 0, "N must be a multiple of ", kTailCopyComponents); - const uint32_t tail_rows = M - (num_m_tiles - 1) * tile_size_a; - const uint32_t row_offset = (num_m_tiles - 1) * tile_size_a; + const uint32_t tail_rows = M - (num_M_tile - 1) * tile_m; + const uint32_t row_offset = (num_M_tile - 1) * tile_m; const uint32_t n_vec4 = N / kTailCopyComponents; const uint32_t output_size = tail_rows * n_vec4; const uint32_t output_offset = row_offset * n_vec4; @@ -325,14 +374,26 @@ bool CanApplySubgroupMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& cont bool has_weight_idx_indirect) { config.reset(); - // Subgroup matrix kernels only support 4-bit/8-bit quantization. - if (nbits != 4 && nbits != 8) { + // Subgroup matrix kernels only support 4-bit/8-bit quantization with block_size 32. + if (!((nbits == 4 || nbits == 8) && block_size == 32)) { return false; } - // Dispatch precondition: the subgroup-matrix kernel is reserved for the - // tile-optimized M range without indirect weight indexing. - if (!(M >= kMinMForTileOptimization && !has_weight_idx_indirect)) { + if (M < kMinMForTileOptimization) { + return false; + } + + if (K % 32 != 0 || N % 64 != 0) { + return false; + } + + // TODO: Support batch. + if (batch_count != 1) { + return false; + } + + // TODO: Support weight_idx_indirect. + if (has_weight_idx_indirect) { return false; } @@ -352,15 +413,21 @@ bool CanApplySubgroupMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& cont return false; } - if (block_size != 32 || batch_count != 1 || K % 32 != 0 || N % 64 != 0) { - return false; - } - constexpr auto kF16 = wgpu::SubgroupMatrixComponentType::F16; constexpr auto kF32 = wgpu::SubgroupMatrixComponentType::F32; if (!is_fp16) { config = SelectSubgroupMatrixConfig(context, {{kF32, kF32, 8, 8, 8, 32, false}}); + } else if (context.AdapterInfo().vendor == std::string_view{"amd"} && + context.AdapterInfo().backendType == wgpu::BackendType::D3D12) { + // The AMD D3D12 wave64 kernel is faster for most measured prefill shapes. Tile + // selection inside the kernel remains workload-dependent on M and K. + config = SelectSubgroupMatrixConfig( + context, {{kF16, kF16, 16, 16, 16, 64, true}, + {kF16, kF16, 16, 16, 16, 32, true}, + {kF16, kF16, 8, 16, 16, 32, true}, + {kF16, kF16, 8, 8, 8, 32, false}}); } else { + // Preserve the established preference for other adapters. config = SelectSubgroupMatrixConfig( context, {{kF16, kF16, 16, 16, 16, 32, true}, {kF16, kF16, 8, 16, 16, 32, true}, diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h index 0bbdc1f5257ea..240c476675e72 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h @@ -20,10 +20,16 @@ using namespace onnxruntime::webgpu; class SubgroupMatrixMatMulNBitsProgram final : public Program { public: - SubgroupMatrixMatMulNBitsProgram(uint32_t nbits, SubgroupMatrixConfig config, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool has_tail_buffer) - : Program{"SubgroupMatrixMatMulNBits"}, + SubgroupMatrixMatMulNBitsProgram(uint32_t nbits, SubgroupMatrixConfig config, uint32_t tile_m, + uint32_t tile_n, uint32_t tile_k, + bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect, bool has_tail_buffer) + : Program{config.subgroupSize == 64 ? "SubgroupMatrixMatMulNBitsWave64" : "SubgroupMatrixMatMulNBits"}, nbits_(nbits), config_(config), + tile_m_(tile_m), + tile_n_(tile_n), + tile_k_(tile_k), has_zero_points_(has_zero_points), has_bias_(has_bias), has_weight_idx_{has_weight_idx}, @@ -35,12 +41,16 @@ class SubgroupMatrixMatMulNBitsProgram final : public Program buf_b: array; // 128 * 40 = 5120 elements = 10240 bytes +var buf_b: array; // Small staging buffer for edge-tile output (one 16x16 tile per subgroup) -var coopmat_stage: array; // 16*16*4 = 1024 elements = 2048 bytes -// Total shared memory: 10240 + 2048 = 12288 bytes +var coopmat_stage: array; #if n_bits == 4 -// Dequantize 4-bit quantized B weights into buf_b with padded stride. -// First half: rows 0..63, Second half: rows 64..127 -// 128 threads, 2 threads per N row, each handles 16 K elements. -fn dequant_b_to_buf(n_base: u32, k_idx: u32, n_idx_0: u32, k_chunk: u32) { +// Dequantize one 16-element K chunk from one B row into buf_b. +fn dequant_b_to_buf(n_base: u32, k_idx: u32, n_idx: u32, k_chunk: u32) { #if has_weight_idx #if has_weight_idx_indirect let actual_weight_idx = weight_index_indirect[uniforms.weight_idx]; @@ -75,57 +78,27 @@ fn dequant_b_to_buf(n_base: u32, k_idx: u32, n_idx_0: u32, k_chunk: u32) { #endif let k_offset = k_chunk * 16u; - // First half: rows 0..63 - { - let global_n = n_base + n_idx_0; - let buf_base = n_idx_0 * SHMEM_STRIDE + k_offset; - if (global_n < uniforms.N) { - let packed_idx = u32((global_n * uniforms.K + k_idx + k_offset) / 8) + weight_offset / 8; - let scale_idx = (global_n * uniforms.K + k_idx + k_offset) / kQuantizationBlockSize + scale_offset; - let scale = f16(scales_b.getByOffset(scale_idx)); - let zero = mm_read_zero( - actual_weight_idx, global_n, (k_idx + k_offset) / kQuantizationBlockSize, uniforms.N, uniforms.zero_blocks_per_col); - for (var step: u32 = 0; step < 2u; step++) { - let packed_weights = input_b.getByOffset(packed_idx + step); - let w_lo = (vec4(unpack4xU8(packed_weights & 0x0F0F0F0Fu)) - vec4(zero)) * scale; - let w_hi = (vec4(unpack4xU8((packed_weights >> 4u) & 0x0F0F0F0Fu)) - vec4(zero)) * scale; - let b = buf_base + step * 8u; - buf_b[b] = w_lo[0]; buf_b[b + 1] = w_hi[0]; - buf_b[b + 2] = w_lo[1]; buf_b[b + 3] = w_hi[1]; - buf_b[b + 4] = w_lo[2]; buf_b[b + 5] = w_hi[2]; - buf_b[b + 6] = w_lo[3]; buf_b[b + 7] = w_hi[3]; - } - } else { - for (var i: u32 = 0; i < 16u; i++) { - buf_b[buf_base + i] = f16(0.0); - } + let global_n = n_base + n_idx; + let buf_base = n_idx * SHMEM_STRIDE + k_offset; + if (global_n < uniforms.N) { + let packed_idx = u32((global_n * uniforms.K + k_idx + k_offset) / 8) + weight_offset / 8; + let scale_idx = (global_n * uniforms.K + k_idx + k_offset) / kQuantizationBlockSize + scale_offset; + let scale = f16(scales_b.getByOffset(scale_idx)); + let zero = mm_read_zero( + actual_weight_idx, global_n, (k_idx + k_offset) / kQuantizationBlockSize, uniforms.N, uniforms.zero_blocks_per_col); + for (var step: u32 = 0; step < 2u; step++) { + let packed_weights = input_b.getByOffset(packed_idx + step); + let w_lo = (vec4(unpack4xU8(packed_weights & 0x0F0F0F0Fu)) - vec4(zero)) * scale; + let w_hi = (vec4(unpack4xU8((packed_weights >> 4u) & 0x0F0F0F0Fu)) - vec4(zero)) * scale; + let b = buf_base + step * 8u; + buf_b[b] = w_lo[0]; buf_b[b + 1] = w_hi[0]; + buf_b[b + 2] = w_lo[1]; buf_b[b + 3] = w_hi[1]; + buf_b[b + 4] = w_lo[2]; buf_b[b + 5] = w_hi[2]; + buf_b[b + 6] = w_lo[3]; buf_b[b + 7] = w_hi[3]; } - } - // Second half: rows 64..127 - { - let n_idx_1 = n_idx_0 + 64u; - let global_n = n_base + n_idx_1; - let buf_base = n_idx_1 * SHMEM_STRIDE + k_offset; - if (global_n < uniforms.N) { - let packed_idx = u32((global_n * uniforms.K + k_idx + k_offset) / 8) + weight_offset / 8; - let scale_idx = (global_n * uniforms.K + k_idx + k_offset) / kQuantizationBlockSize + scale_offset; - let scale = f16(scales_b.getByOffset(scale_idx)); - let zero = mm_read_zero( - actual_weight_idx, global_n, (k_idx + k_offset) / kQuantizationBlockSize, uniforms.N, uniforms.zero_blocks_per_col); - for (var step: u32 = 0; step < 2u; step++) { - let packed_weights = input_b.getByOffset(packed_idx + step); - let w_lo = (vec4(unpack4xU8(packed_weights & 0x0F0F0F0Fu)) - vec4(zero)) * scale; - let w_hi = (vec4(unpack4xU8((packed_weights >> 4u) & 0x0F0F0F0Fu)) - vec4(zero)) * scale; - let b = buf_base + step * 8u; - buf_b[b] = w_lo[0]; buf_b[b + 1] = w_hi[0]; - buf_b[b + 2] = w_lo[1]; buf_b[b + 3] = w_hi[1]; - buf_b[b + 4] = w_lo[2]; buf_b[b + 5] = w_hi[2]; - buf_b[b + 6] = w_lo[3]; buf_b[b + 7] = w_hi[3]; - } - } else { - for (var i: u32 = 0; i < 16u; i++) { - buf_b[buf_base + i] = f16(0.0); - } + } else { + for (var i: u32 = 0; i < 16u; i++) { + buf_b[buf_base + i] = f16(0.0); } } } @@ -133,7 +106,7 @@ fn dequant_b_to_buf(n_base: u32, k_idx: u32, n_idx_0: u32, k_chunk: u32) { #if n_bits == 8 // Dequantize 8-bit quantized B weights into buf_b with padded stride. -fn dequant_b_to_buf(n_base: u32, k_idx: u32, n_idx_0: u32, k_chunk: u32) { +fn dequant_b_to_buf(n_base: u32, k_idx: u32, n_idx: u32, k_chunk: u32) { #if has_weight_idx #if has_weight_idx_indirect let actual_weight_idx = weight_index_indirect[uniforms.weight_idx]; @@ -149,57 +122,27 @@ fn dequant_b_to_buf(n_base: u32, k_idx: u32, n_idx_0: u32, k_chunk: u32) { #endif let k_offset = k_chunk * 16u; - // First half: rows 0..63 - { - let global_n = n_base + n_idx_0; - let buf_base = n_idx_0 * SHMEM_STRIDE + k_offset; - if (global_n < uniforms.N) { - let packed_idx = u32((global_n * uniforms.K + k_idx + k_offset) / 8) + weight_offset / 8; - let scale_idx = (global_n * uniforms.K + k_idx + k_offset) / kQuantizationBlockSize + scale_offset; - let scale = f16(scales_b.getByOffset(scale_idx)); - let zero = mm_read_zero( - actual_weight_idx, global_n, (k_idx + k_offset) / kQuantizationBlockSize, uniforms.N, uniforms.zero_blocks_per_col); - for (var step: u32 = 0; step < 2u; step++) { - let packed_weights = input_b.getByOffset(packed_idx + step); - let w_lo = (vec4(unpack4xU8(packed_weights[0])) - vec4(zero)) * scale; - let w_hi = (vec4(unpack4xU8(packed_weights[1])) - vec4(zero)) * scale; - let b = buf_base + step * 8u; - buf_b[b] = w_lo[0]; buf_b[b + 1] = w_lo[1]; - buf_b[b + 2] = w_lo[2]; buf_b[b + 3] = w_lo[3]; - buf_b[b + 4] = w_hi[0]; buf_b[b + 5] = w_hi[1]; - buf_b[b + 6] = w_hi[2]; buf_b[b + 7] = w_hi[3]; - } - } else { - for (var i: u32 = 0; i < 16u; i++) { - buf_b[buf_base + i] = f16(0.0); - } + let global_n = n_base + n_idx; + let buf_base = n_idx * SHMEM_STRIDE + k_offset; + if (global_n < uniforms.N) { + let packed_idx = u32((global_n * uniforms.K + k_idx + k_offset) / 8) + weight_offset / 8; + let scale_idx = (global_n * uniforms.K + k_idx + k_offset) / kQuantizationBlockSize + scale_offset; + let scale = f16(scales_b.getByOffset(scale_idx)); + let zero = mm_read_zero( + actual_weight_idx, global_n, (k_idx + k_offset) / kQuantizationBlockSize, uniforms.N, uniforms.zero_blocks_per_col); + for (var step: u32 = 0; step < 2u; step++) { + let packed_weights = input_b.getByOffset(packed_idx + step); + let w_lo = (vec4(unpack4xU8(packed_weights[0])) - vec4(zero)) * scale; + let w_hi = (vec4(unpack4xU8(packed_weights[1])) - vec4(zero)) * scale; + let b = buf_base + step * 8u; + buf_b[b] = w_lo[0]; buf_b[b + 1] = w_lo[1]; + buf_b[b + 2] = w_lo[2]; buf_b[b + 3] = w_lo[3]; + buf_b[b + 4] = w_hi[0]; buf_b[b + 5] = w_hi[1]; + buf_b[b + 6] = w_hi[2]; buf_b[b + 7] = w_hi[3]; } - } - // Second half: rows 64..127 - { - let n_idx_1 = n_idx_0 + 64u; - let global_n = n_base + n_idx_1; - let buf_base = n_idx_1 * SHMEM_STRIDE + k_offset; - if (global_n < uniforms.N) { - let packed_idx = u32((global_n * uniforms.K + k_idx + k_offset) / 8) + weight_offset / 8; - let scale_idx = (global_n * uniforms.K + k_idx + k_offset) / kQuantizationBlockSize + scale_offset; - let scale = f16(scales_b.getByOffset(scale_idx)); - let zero = mm_read_zero( - actual_weight_idx, global_n, (k_idx + k_offset) / kQuantizationBlockSize, uniforms.N, uniforms.zero_blocks_per_col); - for (var step: u32 = 0; step < 2u; step++) { - let packed_weights = input_b.getByOffset(packed_idx + step); - let w_lo = (vec4(unpack4xU8(packed_weights[0])) - vec4(zero)) * scale; - let w_hi = (vec4(unpack4xU8(packed_weights[1])) - vec4(zero)) * scale; - let b = buf_base + step * 8u; - buf_b[b] = w_lo[0]; buf_b[b + 1] = w_lo[1]; - buf_b[b + 2] = w_lo[2]; buf_b[b + 3] = w_lo[3]; - buf_b[b + 4] = w_hi[0]; buf_b[b + 5] = w_hi[1]; - buf_b[b + 6] = w_hi[2]; buf_b[b + 7] = w_hi[3]; - } - } else { - for (var i: u32 = 0; i < 16u; i++) { - buf_b[buf_base + i] = f16(0.0); - } + } else { + for (var i: u32 = 0; i < 16u; i++) { + buf_b[buf_base + i] = f16(0.0); } } } @@ -211,14 +154,11 @@ $MAIN { let sg_idx = u32(local_idx / sg_size); let sg_tid = local_idx % sg_size; - let warp_r = sg_idx % 2u; // 0 or 1 (row in 2x2 grid) - let warp_c = sg_idx / 2u; // 0 or 1 (col in 2x2 grid) + let warp_r = sg_idx % kSubgroupsM; + let warp_c = sg_idx / kSubgroupsM; - // Each subgroup's A starts at a different M offset in the prepacked layout - // Subgroup sg_idx handles rows: warp_r * WM .. warp_r * WM + WM (within the BM tile) - // But in the prepacked layout, the subgroups are numbered 0..7 (kSgMatCountM = BM/kSgMatM = 8) - // Subgroups 0,1 handle rows 0..63 (warp_r=0), subgroups 2,3 handle rows 64..127 (warp_r=1) - // Within each 64-row band, we have 4 coopmat rows of 16 + // Prepacked A is arranged as 16-row subgroup matrices along M. warp_r selects one + // 64-row band (four matrices); column subgroups with the same warp_r share those A rows. // For prepacked A: sg_mat index = (workgroup_y * kSgMatCountM + first_coopmat_row_in_warp) * sg_mat_count_k let sg_mat_count_k = uniforms.K / kSgMatK; @@ -252,7 +192,10 @@ $MAIN { // Main K-loop for (var k_block: u32 = 0; k_block < uniforms.K; k_block += BK) { // Load B (Q4_K/Q8 weights) to buf_b with padded stride - dequant_b_to_buf(global_base_n, k_block, local_idx / 2u, local_idx % 2u); + const kChunksPerRow: u32 = BK / 16u; + for (var task = local_idx; task < BN * kChunksPerRow; task += workgroup_size_x) { + dequant_b_to_buf(global_base_n, k_block, task / kChunksPerRow, task % kChunksPerRow); + } workgroupBarrier(); // Compute: load B once per k-step, iterate 4 A rows @@ -504,7 +447,7 @@ $MAIN { #else // Non-bias path if (tile_in_bounds) { - // Fast path: entire 128x128 tile is in bounds - direct store to global + // Fast path: entire workgroup tile is in bounds - direct store to global let base = dr * stride + dc; subgroupMatrixStore(&output, base, c00, stride); subgroupMatrixStore(&output, base + 16u, c01, stride); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template index 5163f98bcc6f3..312196af6986b 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template @@ -5,12 +5,12 @@ // Uses subgroupMatrixLoad/Store/MultiplyAccumulate for hardware DPAS acceleration // // A: prepacked layout (from PrepackProgram) - elements arranged for subgroupMatrixLoad -// B: 4-bit or 8-bit quantized, dequantized into SLM (tile_b) +// B: 4-bit or 8-bit quantized, dequantized into SLM (tile_b). // -// Workgroup: 256 threads = 8 subgroups x 32 lanes -// Tile: kTileM(64) x kTileN(64) -// Each subgroup handles kSgMatM(8) M rows x 64 N columns (4 x kSgMatN N blocks) -// K dimension iterated in kTileK(32) blocks +// Workgroup size and tile_m/tile_n (kTileM x kTileN) are caller-configurable; the number of +// subgroups per workgroup must equal kTileM / sg_mat_m (i.e. workgroup size = kTileM / sg_mat_m * 32). +// Each workgroup handles one (M-tile, N-tile) pair, decoded from a flattened workgroup_idx +// K dimension iterated in kChunkSizeK(32) blocks #param n_bits #param has_zero_points @@ -21,254 +21,237 @@ #param sg_mat_k #param sg_mat_m #param sg_mat_n +#param tile_m +#param tile_n #use .getByOffset .setByOffset #include "quantization/matmul_nbits_zero_pt.wgsl.template" -const kQuantizationBlockSize: u32 = 32; +// A pack is the load/dequantization unit for eight consecutive K-elements of one B row. +// It is stored as a u32 for 4-bit weights or a vec2 for 8-bit weights. +// Four packs form a 32-element quantization block and share its scale and zero point. +#if n_bits == 4 +alias PackT = u32; // one pack = one nibble-packed u32 word (8 values) +#elif n_bits == 8 +alias PackT = vec2; // one pack = one byte-packed vec2 (8 values) +#endif + +const kQuantizationBlockSize : u32 = 32; +const kPacksPerBlock : u32 = 4; // Packs per quantization block, both n_bits +const kElementsPerPack : u32 = 8; // K-elements dequantized per pack, both n_bits + +const kTileM : u32 = u32(tile_m); // Tile size along M (output rows per workgroup) +const kTileN : u32 = u32(tile_n); // Tile size along N (output columns per workgroup) +const kChunkSizeK : u32 = 32; // Tile size along K (reduction dimension per iteration) +const kSgMatDimM : u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) +const kSgMatDimN : u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) +const kSgMatDimK : u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) +const kSgMatLeftSize : u32 = kSgMatDimM * kSgMatDimK; // Elements per left (A) subgroup matrix +const kSgMatStrideN : u32 = kSgMatDimN * kChunkSizeK; // Stride between N blocks in tile_b +const kSgMatCountM : u32 = kTileM / kSgMatDimM; // Number of subgroup matrices per tile along M +const kSgMatCountN : u32 = kTileN / kSgMatDimN; // Number of subgroup matrices per tile along N +const kSgMatCountK : u32 = kChunkSizeK / kSgMatDimK; // Number of subgroup_matrix K-substeps per kChunkSizeK block -const kTileM: u32 = 64; // Tile size along M (output rows per workgroup) -const kTileN: u32 = 64; // Tile size along N (output columns per workgroup) -const kTileK: u32 = 32; // Tile size along K (reduction dimension per iteration) -const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) -const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) -const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) -const kSgMatSizeLeft: u32 = kSgMatM * kSgMatK; // Elements per left (A) subgroup matrix -const kSgMatStrideN: u32 = kSgMatN * kTileK; // Stride between N blocks in tile_b -const kSgMatCountM: u32 = kTileM / kSgMatM; // Number of subgroup matrices per tile along M +alias SgMatLeft = subgroup_matrix_left; +alias SgMatRight = subgroup_matrix_right; +alias SgMatResult = subgroup_matrix_result; -// Shared local memory for dequantized B weights [kTileN, kTileK], column-major -var tile_b: array; +// Shared local memory for dequantized B weights [kTileN, kChunkSizeK], column-major +var tile_b : array; #if has_bias // Scratch space for bounds-checked write-out with bias addition, [kTileM, kTileN], // row-major. The no-bias path never needs this - see `has_tail_buffer` below. -var scratch: array; +var scratch : array; #endif -#if n_bits == 4 -// Dequantize 4-bit quantized B weights and store into tile_b [kTileN, kTileK] workgroup memory. -// 256 threads load kTileN(64) x kTileK(32). 4 threads per N, 8 K elements per thread. -// k_chunk_idx: which 8-element chunk of kTileK this thread handles (0..3). -fn dequant_b_to_tile(n_base: u32, k_idx: u32, n_idx: u32, k_chunk_idx: u32) { - let global_n = n_base + n_idx; - if (global_n >= uniforms.N) { - return; - } - let k_offset = k_chunk_idx * 8; -#if has_weight_idx -#if has_weight_idx_indirect - let actual_weight_idx = weight_index_indirect[uniforms.weight_idx]; -#else - let actual_weight_idx = uniforms.weight_idx; -#endif - let weight_offset = actual_weight_idx * uniforms.K * uniforms.N; - let scale_offset = actual_weight_idx * uniforms.N * (uniforms.K / kQuantizationBlockSize); -#else - const actual_weight_idx : u32 = 0; - const weight_offset : u32 = 0; - const scale_offset : u32 = 0; -#endif - let packed_idx = u32((global_n * uniforms.K + k_idx + k_offset) / 8) + weight_offset / 8; - let scale_idx = (global_n * uniforms.K + k_idx + k_offset) / kQuantizationBlockSize + scale_offset; - let scale = f16(scales_b.getByOffset(scale_idx)); - let zero = mm_read_zero( - actual_weight_idx, global_n, (k_idx + k_offset) / kQuantizationBlockSize, uniforms.N, uniforms.zero_blocks_per_col); - let packed_weights = input_b.getByOffset(packed_idx); - let weights_lo = - (vec4(unpack4xU8(packed_weights & 0x0F0F0F0Fu)) - vec4(zero)) * scale; - let weights_hi = - (vec4(unpack4xU8((packed_weights >> 4) & 0x0F0F0F0Fu)) - vec4(zero)) * scale; - let tile_base = n_idx * kTileK + k_offset; - tile_b[tile_base] = weights_lo[0]; - tile_b[tile_base + 1] = weights_hi[0]; - tile_b[tile_base + 2] = weights_lo[1]; - tile_b[tile_base + 3] = weights_hi[1]; - tile_b[tile_base + 4] = weights_lo[2]; - tile_b[tile_base + 5] = weights_hi[2]; - tile_b[tile_base + 6] = weights_lo[3]; - tile_b[tile_base + 7] = weights_hi[3]; +// --- Data Loading Functions --- + +// scale_offset is resolved before the K loop, not recomputed for each load. +// Note: the scales buffer is strided by blocks-per-row (K / kQuantizationBlockSize), which is +// NOT the same as uniforms.zero_blocks_per_col (that's padded to a zero-point byte boundary). +fn load_scale(row : u32, block_idx : u32, scale_offset : u32) -> output_element_t { + let blocks_per_col = uniforms.K / kQuantizationBlockSize; + let offset = row * blocks_per_col + block_idx + scale_offset; + return scales_b.getByOffset(offset); } -#endif -#if n_bits == 8 -// Dequantize 8-bit quantized B weights and store into tile_b [kTileN, kTileK] workgroup memory. -// 256 threads load kTileN(64) x kTileK(32). 4 threads per N, 8 K elements per thread. -// k_chunk_idx: which 8-element chunk of kTileK this thread handles (0..3). -fn dequant_b_to_tile(n_base: u32, k_idx: u32, n_idx: u32, k_chunk_idx: u32) { - let global_n = n_base + n_idx; - if (global_n >= uniforms.N) { - return; - } - let k_offset = k_chunk_idx * 8; -#if has_weight_idx +// Fetches a single pack of raw B data (one u32 word for n_bits=4, one +// vec2 for n_bits=8), addressed by its flat pack index within the row +// (row * packs_per_row + pack_idx), offset by weight_offset (in block units, +// scaled up to pack units here). +fn load_b(row : u32, pack_idx : u32, weight_offset : u32) -> PackT { + let k_of_b = uniforms.K / kQuantizationBlockSize; + let packs_per_row = k_of_b * kPacksPerBlock; + let offset = row * packs_per_row + pack_idx + weight_offset * kPacksPerBlock; + return input_b.getByOffset(offset); +} + +// --- Main Compute --- + +$MAIN { + // Resolve the (possibly indirect) weight index once per workgroup. #if has_weight_idx_indirect - let actual_weight_idx = weight_index_indirect[uniforms.weight_idx]; -#else - let actual_weight_idx = uniforms.weight_idx; -#endif - let weight_offset = actual_weight_idx * uniforms.K * uniforms.N; - let scale_offset = actual_weight_idx * uniforms.N * (uniforms.K / kQuantizationBlockSize); + let actual_weight_idx = weight_index_indirect[uniforms.weight_idx]; +#elif has_weight_idx + let actual_weight_idx = uniforms.weight_idx; #else - const actual_weight_idx : u32 = 0; - const weight_offset : u32 = 0; - const scale_offset : u32 = 0; -#endif - let packed_idx = u32((global_n * uniforms.K + k_idx + k_offset) / 8) + weight_offset / 8; - let scale_idx = (global_n * uniforms.K + k_idx + k_offset) / kQuantizationBlockSize + scale_offset; - let scale = f16(scales_b.getByOffset(scale_idx)); - let zero = mm_read_zero( - actual_weight_idx, global_n, (k_idx + k_offset) / kQuantizationBlockSize, uniforms.N, uniforms.zero_blocks_per_col); - let packed_weights = input_b.getByOffset(packed_idx); - let weights_lo = (vec4(unpack4xU8(packed_weights[0])) - vec4(zero)) * scale; - let weights_hi = (vec4(unpack4xU8(packed_weights[1])) - vec4(zero)) * scale; - let tile_base = n_idx * kTileK + k_offset; - tile_b[tile_base] = weights_lo[0]; - tile_b[tile_base + 1] = weights_lo[1]; - tile_b[tile_base + 2] = weights_lo[2]; - tile_b[tile_base + 3] = weights_lo[3]; - tile_b[tile_base + 4] = weights_hi[0]; - tile_b[tile_base + 5] = weights_hi[1]; - tile_b[tile_base + 6] = weights_hi[2]; - tile_b[tile_base + 7] = weights_hi[3]; -} + let actual_weight_idx = 0u; #endif + let scale_offset = actual_weight_idx * uniforms.N * (uniforms.K / kQuantizationBlockSize); + // One scale per block; load_b converts the shared block offset to pack units. + let weight_offset = scale_offset; -$MAIN { - let global_base_b = workgroup_id.x * kTileN; - let sg_idx = u32(local_idx / sg_size); - let sg_mat_count_k = uniforms.K / kSgMatK; - let num_tiles_m = (uniforms.M + kTileM - 1) / kTileM; + let m_tile = (workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile; - // Zero-initialized accumulator template (used to reset per M-tile iteration). - var sg_mat_zero: subgroup_matrix_result; + let global_base_m = m_tile * kTileM; + let global_base_n = (workgroup_idx % uniforms.num_N_tile) * kTileN; - // Sequential M-loop: each workgroup processes a contiguous block of M-tiles. - let m_start = workgroup_id.y * uniforms.m_tiles_per_wg; - let m_end = min(m_start + uniforms.m_tiles_per_wg, num_tiles_m); - for (var m_tile: u32 = m_start; m_tile < m_end; m_tile++) { - let global_base_a = m_tile * kTileM; - let sg_mat_idx = (m_tile * kSgMatCountM + sg_idx) * sg_mat_count_k; + let sg_idx = u32(local_idx / sg_size); + var sg_mat_offset_a = (m_tile * kSgMatCountM + sg_idx) * (uniforms.K / kSgMatDimK) * kSgMatLeftSize; - var sg_mat_offset_a = sg_mat_idx * kSgMatSizeLeft; + var sg_mat_acc : array; + for (var k_idx : u32 = 0; k_idx < uniforms.K; k_idx += kChunkSizeK) { + let k_block_idx = k_idx / kQuantizationBlockSize; + for (var load_unit = local_idx; load_unit < kTileN * kPacksPerBlock; load_unit += workgroup_size_x) { + let n_idx = load_unit / kPacksPerBlock; + let pack_idx = load_unit % kPacksPerBlock; + let global_n = global_base_n + n_idx; + let tile_b_offset = n_idx * kChunkSizeK + pack_idx * kElementsPerPack; - var sg_mat_c0 = sg_mat_zero; - var sg_mat_c1 = sg_mat_zero; - var sg_mat_c2 = sg_mat_zero; - var sg_mat_c3 = sg_mat_zero; - for (var k_idx: u32 = 0; k_idx < uniforms.K; k_idx += kTileK) { - // Load Phase - dequant_b_to_tile(global_base_b, k_idx, local_idx / 4, local_idx % 4); - workgroupBarrier(); + if (global_n < uniforms.N) { + let packed = load_b(global_n, k_block_idx * kPacksPerBlock + pack_idx, weight_offset); + let scale = load_scale(global_n, k_block_idx, scale_offset); + let zero_point = mm_read_zero(actual_weight_idx, global_n, k_block_idx, + uniforms.N, uniforms.zero_blocks_per_col); +#if n_bits == 4 + let lower = + (vec4(unpack4xU8(packed & 0x0F0F0F0Fu)) - vec4(zero_point)) * scale; + let upper = + (vec4(unpack4xU8((packed >> 4u) & 0x0F0F0F0Fu)) - vec4(zero_point)) * scale; + tile_b[tile_b_offset] = lower[0]; + tile_b[tile_b_offset + 1u] = upper[0]; + tile_b[tile_b_offset + 2u] = lower[1]; + tile_b[tile_b_offset + 3u] = upper[1]; + tile_b[tile_b_offset + 4u] = lower[2]; + tile_b[tile_b_offset + 5u] = upper[2]; + tile_b[tile_b_offset + 6u] = lower[3]; + tile_b[tile_b_offset + 7u] = upper[3]; +#elif n_bits == 8 + let lower = + (vec4(unpack4xU8(packed.x)) - vec4(zero_point)) * scale; + let upper = + (vec4(unpack4xU8(packed.y)) - vec4(zero_point)) * scale; + tile_b[tile_b_offset] = lower[0]; + tile_b[tile_b_offset + 1u] = lower[1]; + tile_b[tile_b_offset + 2u] = lower[2]; + tile_b[tile_b_offset + 3u] = lower[3]; + tile_b[tile_b_offset + 4u] = upper[0]; + tile_b[tile_b_offset + 5u] = upper[1]; + tile_b[tile_b_offset + 6u] = upper[2]; + tile_b[tile_b_offset + 7u] = upper[3]; +#endif + } + } + workgroupBarrier(); - for (var sg_mat_k_idx: u32 = 0; sg_mat_k_idx < kTileK; sg_mat_k_idx += kSgMatK) - { - // Load A from global memory (prepacked layout). - // Syntax: subgroupMatrixLoad(src_ptr, src_offset, src_stride) - var sg_mat_a0: subgroup_matrix_left = - subgroupMatrixLoad, row_major>( - &input_a, sg_mat_offset_a, kSgMatK); - sg_mat_offset_a += kSgMatSizeLeft; + for (var sg_mat_k_idx = 0u; sg_mat_k_idx < kSgMatCountK; sg_mat_k_idx++) { + var sg_mat_a : SgMatLeft = subgroupMatrixLoad(&input_a, + sg_mat_offset_a, + kSgMatDimK); + sg_mat_offset_a += kSgMatLeftSize; - // Load B from shared local memory. - // tile_b [kTileN, kTileK] is stored as column major. - var sg_mat_b0: subgroup_matrix_right = - subgroupMatrixLoad, col_major>( - &tile_b, sg_mat_k_idx, kTileK); - var sg_mat_b1: subgroup_matrix_right = - subgroupMatrixLoad, col_major>( - &tile_b, sg_mat_k_idx + kSgMatStrideN, kTileK); - var sg_mat_b2: subgroup_matrix_right = - subgroupMatrixLoad, col_major>( - &tile_b, sg_mat_k_idx + 2 * kSgMatStrideN, kTileK); - var sg_mat_b3: subgroup_matrix_right = - subgroupMatrixLoad, col_major>( - &tile_b, sg_mat_k_idx + 3 * kSgMatStrideN, kTileK); + var sg_mat_offset_b = sg_mat_k_idx * kSgMatDimK; + for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { + var sg_mat_b : SgMatRight = subgroupMatrixLoad(&tile_b, + sg_mat_offset_b, + kChunkSizeK); + sg_mat_offset_b += kSgMatStrideN; - // Compute Phase - // Syntax: subgroupMatrixMultiplyAccumulate left, right, accumulate -> accumulate - sg_mat_c0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0); - sg_mat_c1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c1); - sg_mat_c2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c2); - sg_mat_c3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c3); - } - workgroupBarrier(); - } + sg_mat_acc[sg_mat_n_idx] = subgroupMatrixMultiplyAccumulate(sg_mat_a, sg_mat_b, sg_mat_acc[sg_mat_n_idx]); + } + } + workgroupBarrier(); + } - // Write out + // Write out #if has_bias - // Store results to scratch workgroup memory, then add bias and write to output. - // scratch layout: [kTileM, kTileN] row-major - let scratch_m_base = sg_idx * kSgMatM; - subgroupMatrixStore(&scratch, scratch_m_base * kTileN, sg_mat_c0, kTileN); - subgroupMatrixStore(&scratch, scratch_m_base * kTileN + kSgMatN, sg_mat_c1, kTileN); - subgroupMatrixStore(&scratch, scratch_m_base * kTileN + 2 * kSgMatN, sg_mat_c2, kTileN); - subgroupMatrixStore(&scratch, scratch_m_base * kTileN + 3 * kSgMatN, sg_mat_c3, kTileN); - workgroupBarrier(); + let scratch_store_base = sg_idx * kSgMatDimM * kTileN; + for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { + subgroupMatrixStore(&scratch, + scratch_store_base + sg_mat_n_idx * kSgMatDimN, + sg_mat_acc[sg_mat_n_idx], + kTileN); + } + workgroupBarrier(); - // 256 threads write 64x64 = 4096 elements. Each thread handles 16 elements. - // Thread mapping: m = local_idx / 4, n_base = (local_idx % 4) * 16 - let out_m = local_idx / 4; - let out_n_base = (local_idx % 4) * 16; - let global_m = global_base_a + out_m; - if (global_m < uniforms.M) { - let global_n_base = global_base_b + out_n_base; - let scratch_base = out_m * kTileN + out_n_base; - let out_base = global_m * uniforms.N + global_n_base; + let ThreadsPerMRow = workgroup_size_x / kTileM; // Threads cooperatively writing each M row + let ElemsPerThreadRow = kTileN / ThreadsPerMRow; // N-elements each thread writes per row + let out_m = local_idx / ThreadsPerMRow; + let out_n_base = (local_idx % ThreadsPerMRow) * ElemsPerThreadRow; + let out_global_m = global_base_m + out_m; + if (out_global_m < uniforms.M) { + let out_global_n_base = global_base_n + out_n_base; + let scratch_base = out_m * kTileN + out_n_base; + let out_base = out_global_m * uniforms.N + out_global_n_base; #if has_weight_idx_indirect - let bias_offset = weight_index_indirect[uniforms.weight_idx] * uniforms.N; + let bias_offset = weight_index_indirect[uniforms.weight_idx] * uniforms.N; #elif has_weight_idx - let bias_offset = uniforms.weight_idx * uniforms.N; + let bias_offset = uniforms.weight_idx * uniforms.N; #else - const bias_offset: u32 = 0; + const bias_offset : u32 = 0; #endif - for (var i: u32 = 0; i < 16; i++) { - if (global_n_base + i < uniforms.N) { - let val = output_element_t(scratch[scratch_base + i]) - + bias[bias_offset + global_n_base + i]; - output.setByOffset(out_base + i, val); - } - } - } + for (var i : u32 = 0; i < ElemsPerThreadRow; i++) { + if (out_global_n_base + i < uniforms.N) { + let val = f16(scratch[scratch_base + i]) + bias[bias_offset + out_global_n_base + i]; + output.setByOffset(out_base + i, val); + } + } + } #else - // A trailing M-tile may be partial (M is not required to be a multiple of - // kTileM). Writing it with subgroupMatrixStore straight to output would run - // past M. That's not just an out-of-bounds risk: a subgroupMatrixStore whose - // footprint leaves the buffer does not degrade to a per-element clamp or zero - // fill - Tint's robustness transform resets the whole store to offset 0 with - // the minimum stride, so it would silently overwrite a different, valid region - // of the buffer (e.g. row 0) instead. + // A trailing M-tile may be partial (M is not required to be a multiple of + // kTileM). Writing it with subgroupMatrixStore straight to output would run + // past M. That's not just an out-of-bounds risk: a subgroupMatrixStore whose + // footprint leaves the buffer does not degrade to a per-element clamp or zero + // fill - Tint's robustness transform resets the whole store to offset 0 with + // the minimum stride, so it would silently overwrite a different, valid region + // of the buffer (e.g. row 0) instead. #if has_tail_buffer - // Every tile - including the tail one - still writes via the same - // branch-free fast path (no per-element bounds check, no workgroup - // scratch); the tail tile targets a [kTileM, N] padded `tail_output` - // buffer instead of `output` (which is too short by this point). - // ApplySubgroupMatrixMatMulNBits crops the valid rows back into the real - // output with a follow-up copy dispatch. - if (global_base_a + kTileM <= uniforms.M) { - let sg_mat_offset_c = global_base_a * uniforms.N + global_base_b + sg_idx * kSgMatM * uniforms.N; - subgroupMatrixStore(&output, sg_mat_offset_c, sg_mat_c0, uniforms.N); - subgroupMatrixStore(&output, sg_mat_offset_c + kSgMatN, sg_mat_c1, uniforms.N); - subgroupMatrixStore(&output, sg_mat_offset_c + 2 * kSgMatN, sg_mat_c2, uniforms.N); - subgroupMatrixStore(&output, sg_mat_offset_c + 3 * kSgMatN, sg_mat_c3, uniforms.N); - } else { - let sg_mat_offset_c = global_base_b + sg_idx * kSgMatM * uniforms.N; - subgroupMatrixStore(&tail_output, sg_mat_offset_c, sg_mat_c0, uniforms.N); - subgroupMatrixStore(&tail_output, sg_mat_offset_c + kSgMatN, sg_mat_c1, uniforms.N); - subgroupMatrixStore(&tail_output, sg_mat_offset_c + 2 * kSgMatN, sg_mat_c2, uniforms.N); - subgroupMatrixStore(&tail_output, sg_mat_offset_c + 3 * kSgMatN, sg_mat_c3, uniforms.N); - } + // Every tile - including the tail one - still writes via the same + // branch-free fast path (no per-element bounds check, no workgroup + // scratch); the tail tile targets a [kTileM, N] padded `tail_output` + // buffer instead of `output` (which is too short by this point). + // ApplySubgroupMatrixMatMulNBits crops the valid rows back into the real + // output with a follow-up copy dispatch. + if (global_base_m + kTileM <= uniforms.M) { + let out_base = (global_base_m + sg_idx * kSgMatDimM) * uniforms.N + global_base_n; + for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { + subgroupMatrixStore(&output, + out_base + sg_mat_n_idx * kSgMatDimN, + sg_mat_acc[sg_mat_n_idx], + uniforms.N); + } + } else { + let out_base = sg_idx * kSgMatDimM * uniforms.N + global_base_n; + for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { + subgroupMatrixStore(&tail_output, + out_base + sg_mat_n_idx * kSgMatDimN, + sg_mat_acc[sg_mat_n_idx], + uniforms.N); + } + } #else - // ApplySubgroupMatrixMatMulNBits only sets has_tail_buffer when M isn't a - // multiple of kTileM, so here M is guaranteed aligned: every tile - - // including the last - is always fully in bounds. - let sg_mat_offset_c = global_base_a * uniforms.N + global_base_b + sg_idx * kSgMatM * uniforms.N; - subgroupMatrixStore(&output, sg_mat_offset_c, sg_mat_c0, uniforms.N); - subgroupMatrixStore(&output, sg_mat_offset_c + kSgMatN, sg_mat_c1, uniforms.N); - subgroupMatrixStore(&output, sg_mat_offset_c + 2 * kSgMatN, sg_mat_c2, uniforms.N); - subgroupMatrixStore(&output, sg_mat_offset_c + 3 * kSgMatN, sg_mat_c3, uniforms.N); + // ApplySubgroupMatrixMatMulNBits only sets has_tail_buffer when M isn't a + // multiple of kTileM, so here M is guaranteed aligned: every tile - + // including the last - is always fully in bounds. + let out_base = (global_base_m + sg_idx * kSgMatDimM) * uniforms.N + global_base_n; + for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { + subgroupMatrixStore(&output, + out_base + sg_mat_n_idx * kSgMatDimN, + sg_mat_acc[sg_mat_n_idx], + uniforms.N); + } #endif #endif - } // end M-tile loop } // MAIN diff --git a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index cc927b08a7195..8277b69261a98 100644 --- a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc @@ -8,6 +8,7 @@ #include "contrib_ops/webgpu/bert/engram_gate.h" #include "contrib_ops/webgpu/bert/gated_add.h" #include "contrib_ops/webgpu/bert/gated_delta_net.h" +#include "contrib_ops/webgpu/bert/gated_delta_net_state_replay.h" #include "contrib_ops/webgpu/bert/group_query_attention.h" #include "contrib_ops/webgpu/bert/branchwise_rms_norm.h" #include "contrib_ops/webgpu/bert/hyper_connection_post_mix.h" @@ -48,6 +49,7 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -57,6 +59,9 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, +#if !defined(DISABLE_FLOAT8_TYPES) + BuildKernelCreateInfo, +#endif BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/framework/debug_node_inputs_outputs_utils.cc b/onnxruntime/core/framework/debug_node_inputs_outputs_utils.cc index 5137c22d6cf61..aaf0f6f73dead 100644 --- a/onnxruntime/core/framework/debug_node_inputs_outputs_utils.cc +++ b/onnxruntime/core/framework/debug_node_inputs_outputs_utils.cc @@ -172,13 +172,14 @@ void DumpTensorToFile(const Tensor& tensor, const std::string& tensor_name, cons const PathString file_path_str = file_path.native(); int output_fd; ORT_THROW_IF_ERROR(Env::Default().FileOpenWr(file_path_str, output_fd)); - try { + ORT_TRY { ORT_ENFORCE( tensor_proto.SerializeToFileDescriptor(output_fd), "Failed to write tensor to file - tensor: ", tensor_name, ", file: ", ToUTF8String(file_path_str)); - } catch (...) { + } + ORT_CATCH(...) { ORT_IGNORE_RETURN_VALUE(Env::Default().FileClose(output_fd)); - throw; + ORT_RETHROW; } ORT_THROW_IF_ERROR(Env::Default().FileClose(output_fd)); } diff --git a/onnxruntime/core/framework/onnxruntime_map_type_info.h b/onnxruntime/core/framework/onnxruntime_map_type_info.h index 09674df403fe3..21e8f3e7a2503 100644 --- a/onnxruntime/core/framework/onnxruntime_map_type_info.h +++ b/onnxruntime/core/framework/onnxruntime_map_type_info.h @@ -32,7 +32,7 @@ constexpr ONNXTensorElementDataType ToONNXTensorElementDataType( consteval bool IsTensorProtoToOrtElementTypeMapBijective() { constexpr size_t ort_element_type_count = - static_cast(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E8M0) + 1; + static_cast(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E3M2) + 1; constexpr size_t onnx_element_type_count = ONNX_NAMESPACE::TensorProto_DataType_DataType_ARRAYSIZE; if constexpr (onnx_element_type_count != ort_element_type_count) { return false; diff --git a/onnxruntime/core/framework/resource_accountant.cc b/onnxruntime/core/framework/resource_accountant.cc index f455b1b2969cd..39510f1202895 100644 --- a/onnxruntime/core/framework/resource_accountant.cc +++ b/onnxruntime/core/framework/resource_accountant.cc @@ -19,6 +19,7 @@ #include #include +#include #include namespace onnxruntime { @@ -121,8 +122,9 @@ class SizeBasedResourceAccountant : public IResourceAccountant { // If profiling statistics are available, uses their non-workspace cost and // profiled temporary allocations. If a Level-1 runtime workspace estimate is // also provided, uses the maximum of the profiled and estimated workspace. - // Without profiling, computes known initializer/output bytes and uses the - // Level-1 runtime workspace or, when unavailable, fallback workspace. + // Without profiling, or when the profile omits a node with a Level-1 estimate, + // computes known initializer/output bytes and uses the Level-1 runtime + // workspace or, when unavailable, fallback workspace. // Initializer bytes are charged separately, so persistent prepack estimates // must include only additional allocations, not storage reused directly from // an initializer (for example, an offline-prepacked weight). Persistent @@ -181,11 +183,12 @@ class SizeBasedResourceAccountant : public IResourceAccountant { return static_cast(resource_count); } - // Preserve the established partial-profile behavior: a node absent from - // the stats file has zero cost. Falling through to ad-hoc accounting - // would change existing partition decisions and mix profile-based costs - // with initializer bookkeeping that the profile path does not use. - return static_cast(0); + // Preserve the established partial-profile behavior when no estimator is + // available. An explicit Level-1 estimate must still affect placement, so + // use the static initializer/output accounting below for that node. + if (!level1_memory_estimate.has_value()) { + return static_cast(0); + } } const auto* graph = node.GetContainingGraph(); @@ -538,6 +541,29 @@ static Status LoadNodeAllocationStats( return Status::OK(); } +Status ParseCudaGqaWorkspaceMaxTotalSequenceLength( + const ConfigOptions& config_options, + int64_t& max_total_sequence_length) { + const auto knob = + config_options.GetConfigEntry(kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength); + if (!knob.has_value()) { + max_total_sequence_length = 0; + return Status::OK(); + } + + int64_t parsed = 0; + if (!TryParseStringWithClassicLocale(*knob, parsed) || parsed < 0 || + parsed > std::numeric_limits::max()) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength, + " must be a decimal integer in [0, INT32_MAX], but got: '", *knob, "'"); + } + + max_total_sequence_length = parsed; + return Status::OK(); +} + Status CreateAccountants( const ConfigOptions& config_options, const std::filesystem::path& model_path, @@ -545,17 +571,8 @@ Status CreateAccountants( WorkspaceEstimatorConfig estimator_config{ config_options.GetConfigEntry(kOrtSessionOptionsCudaFpAIntBGemm), config_options.GetConfigEntry(kOrtSessionOptionsCudaFpAIntBProfileM)}; - if (const auto knob = - config_options.GetConfigEntry(kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength); - knob.has_value()) { - auto& length = estimator_config.cuda_gqa_workspace_max_total_sequence_length; - if (!TryParseStringWithClassicLocale(*knob, length) || length < 0) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength, - " must be a nonnegative decimal int64 value, but got: '", *knob, "'"); - } - } + ORT_RETURN_IF_ERROR(ParseCudaGqaWorkspaceMaxTotalSequenceLength( + config_options, estimator_config.cuda_gqa_workspace_max_total_sequence_length)); std::optional result; // Check if CUDA partitioning settings are provided diff --git a/onnxruntime/core/framework/tensor_type_and_shape.cc b/onnxruntime/core/framework/tensor_type_and_shape.cc index e91d0396968e3..8f9ec22e49a3e 100644 --- a/onnxruntime/core/framework/tensor_type_and_shape.cc +++ b/onnxruntime/core/framework/tensor_type_and_shape.cc @@ -187,6 +187,12 @@ constexpr ONNXTensorElementDataType TensorDataTypeToOnnxRuntimeTensorElementData case o::TensorProto_DataType_FLOAT8E8M0: type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E8M0; break; + case o::TensorProto_DataType_FLOAT6E2M3: + type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E2M3; + break; + case o::TensorProto_DataType_FLOAT6E3M2: + type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E3M2; + break; case o::TensorProto_DataType_FLOAT16: type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; break; diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index cd00b02bfac2f..3916a320980ba 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -2237,6 +2237,8 @@ ONNXTensorElementDataType CApiElementTypeFromProtoType(int type) { #if !defined(DISABLE_FLOAT4_TYPES) CASE_TYPE(FLOAT4E2M1) #endif + CASE_TYPE(FLOAT6E2M3) + CASE_TYPE(FLOAT6E3M2) default: return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; @@ -2712,6 +2714,25 @@ common::Status SparseTensorProtoToDenseTensorProto(const ONNX_NAMESPACE::SparseT ONNX_NAMESPACE::TensorProto_DataType_STRING); } + // CopySparseData also validates the type, but it runs after allocating the dense buffer + // and is skipped entirely for an empty dense tensor. Validate here so unsupported types + // are rejected before either path. + if (indices.data_type() != ONNX_NAMESPACE::TensorProto_DataType_INT8 && + indices.data_type() != ONNX_NAMESPACE::TensorProto_DataType_INT16 && + indices.data_type() != ONNX_NAMESPACE::TensorProto_DataType_INT32 && + indices.data_type() != ONNX_NAMESPACE::TensorProto_DataType_INT64) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_GRAPH, + "Sparse tensor: ", name, + " indices. Should be one of the following types: int8, int16, int32 or int64"); + } + + // Validate embedded indices before allocating the dense zero-filled buffer so malformed + // inline data is rejected before a large representable dense allocation is attempted. + if (!HasExternalData(indices)) { + ORT_RETURN_IF_ERROR(ValidateEmbeddedTensorProtoDataSizeAndShape(indices)); + } + if (dense_elements == 0) { // if there are no elements in the dense tensor, we can return early with an empty tensor proto return status; diff --git a/onnxruntime/core/framework/tensorprotoutils.h b/onnxruntime/core/framework/tensorprotoutils.h index 59bd955c6519e..d1ff89c12120c 100644 --- a/onnxruntime/core/framework/tensorprotoutils.h +++ b/onnxruntime/core/framework/tensorprotoutils.h @@ -58,6 +58,8 @@ inline constexpr std::arrayset_dim_value(present_sequence_length); } else { - // Cannot compute exact present_sequence_length. - if (ctx.getNumInputs() > 6 && past_dims[2].has_dim_value() && past_dims[2].dim_value() == 0) { - // If total_sequence_length is provided and past_key has 0 length, present_key will grow. - // Leave the dimension as dynamic to avoid "Error merging shape info" warning. - present_shape.add_dim(); - } else { - *present_shape.add_dim() = past_dims[2]; - } + // total_sequence_length is a runtime input here, so neither a symbolic nor a fixed past + // length bounds present. Copying the past dimension would let the allocation planner give + // present a past-sized buffer ("Shape mismatch attempting to re-use buffer"). + present_shape.add_dim(); } *present_shape.add_dim() = past_dims[3]; // head_size @@ -488,6 +484,16 @@ void GroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); + const int64_t num_heads = getAttribute(ctx, "num_heads", 0); + const int64_t kv_num_heads = getAttribute(ctx, "kv_num_heads", 0); + if (num_heads <= 0 || num_heads > std::numeric_limits::max()) { + fail_shape_inference("num_heads must be a positive int"); + } + if (kv_num_heads <= 0 || kv_num_heads > std::numeric_limits::max() || + num_heads % kv_num_heads != 0) { + fail_shape_inference("kv_num_heads must be positive and divide num_heads"); + } + if (hasInputShape(ctx, 0)) { const auto& query_shape = getInputShape(ctx, 0); const auto& query_dims = query_shape.dim(); @@ -496,10 +502,9 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex } auto output_shape = query_shape; - const int64_t num_heads = getAttribute(ctx, "num_heads", 0); - const int64_t kv_num_heads = getAttribute(ctx, "kv_num_heads", 0); if (ctx.getInputType(2) == nullptr) { - if (num_heads > 0 && kv_num_heads > 0 && query_dims[2].has_dim_value()) { + output_shape.mutable_dim(2)->Clear(); + if (query_dims[2].has_dim_value()) { const int64_t packed_heads = num_heads + 2 * kv_num_heads; const int64_t packed_hidden_size = query_dims[2].dim_value(); if (packed_hidden_size % packed_heads != 0) { @@ -507,8 +512,7 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex } output_shape.mutable_dim(2)->set_dim_value(num_heads * (packed_hidden_size / packed_heads)); } - } else if (num_heads > 0 && query_dims[2].has_dim_value() && - query_dims[2].dim_value() % num_heads != 0) { + } else if (query_dims[2].has_dim_value() && query_dims[2].dim_value() % num_heads != 0) { fail_shape_inference("Query hidden size must be divisible by the number of query heads"); } updateOutputShape(ctx, 0, output_shape); @@ -528,10 +532,7 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex if (hasInputShape(ctx, 0)) { const auto& query_dims = getInputShape(ctx, 0).dim(); - const int64_t num_heads = getAttribute(ctx, "num_heads", 0); - const int64_t kv_num_heads = getAttribute(ctx, "kv_num_heads", 0); - if (query_dims.size() == 3 && num_heads > 0 && kv_num_heads > 0 && - query_dims[2].has_dim_value()) { + if (query_dims.size() == 3 && query_dims[2].has_dim_value()) { const bool is_packed = ctx.getInputType(2) == nullptr; const int64_t divisor = is_packed ? num_heads + 2 * kv_num_heads : num_heads; if (query_dims[2].dim_value() % divisor != 0) { @@ -541,8 +542,9 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex ONNX_NAMESPACE::TensorShapeProto present_shape; *present_shape.add_dim() = query_dims[0]; present_shape.add_dim()->set_dim_value(kv_num_heads); + const bool past_inputs_omitted = ctx.getInputType(3) == nullptr && ctx.getInputType(4) == nullptr; const auto* total_length_data = ctx.getInputData(10); - if (total_length_data != nullptr) { + if (past_inputs_omitted && total_length_data != nullptr) { const auto total_lengths = ParseData(total_length_data); if (total_lengths.size() != 1) { fail_shape_inference("total_sequence_length must contain exactly one element"); @@ -609,7 +611,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( AttributeProto::INT, static_cast(0)) .Attr("rotary_interleaved", "Whether rotary pairs use interleaved layout.", AttributeProto::INT, static_cast(0)) - .Attr("rotary_offset", "First head channel covered by rotary embedding.", + .Attr("rotary_offset", + "First head channel covered by rotary embedding. Must be nonnegative, a multiple of 8, and leave room " + "for the rotary dimensions within head_size.", AttributeProto::INT, static_cast(0)) .Attr("qk_norm_epsilon", "Epsilon for optional per-head Q/K RMS normalization.", AttributeProto::FLOAT, 1e-6f) @@ -618,7 +622,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("auxiliary_kv_shared", "Use auxiliary_key as both key and value when auxiliary_value is omitted.", AttributeProto::INT, static_cast(0)) .Input(0, "query", - "Query [batch, sequence, num_heads * head_size], or packed QKV.", "T") + "Query [batch, sequence, num_heads * head_size]. When key and value are omitted, each row contains " + "Q, K, then V concatenated with total width (num_heads + 2 * kv_num_heads) * head_size.", + "T") .Input(1, "key", "Current main key [batch, sequence, kv_num_heads * head_size].", "T", OpSchema::Optional) .Input(2, "value", "Current main value [batch, sequence, kv_num_heads * head_size].", "T", @@ -2004,7 +2010,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( static OpSchema PackedSparseAttentionIndexerMergeSchema() { return OpSchema() .SetDoc( - "Stable union of a mapped base selection and [range_start, range_end). " + "Stable union of a mapped base selection and either [range_start, range_end) " + "or an additional selection prefix. " "Valid prefixes contain nonnegative indices. Padding is ignored. Outputs preserve " "first-occurrence order and are padded with -1. Status is 0 for success, 1 for invalid " "device metadata, or 2 for capacity overflow; failures return count 0 and all -1. " @@ -2016,8 +2023,10 @@ static OpSchema PackedSparseAttentionIndexerMergeSchema() { .Input(2, "base_row_indices", "Output-query to cached-row mapping, shape (N).", "T") .Input(3, "range_starts", "Inclusive nonnegative starts, shape (N).", "T", OpSchema::Optional) .Input(4, "range_ends", "Exclusive ends not less than starts, shape (N).", "T", OpSchema::Optional) - .Input(5, "additional_indices", "Reserved; must be absent for append_range.", "T", OpSchema::Optional) - .Input(6, "additional_counts", "Reserved; must be absent for append_range.", "T", OpSchema::Optional) + .Input(5, "additional_indices", "Required for append_indices: selections of shape (N, A); absent for append_range.", + "T", OpSchema::Optional) + .Input(6, "additional_counts", "Required for append_indices: valid prefix lengths of shape (N); absent for append_range.", + "T", OpSchema::Optional) .Output(0, "selected_indices", "Merged indices, shape (N, max_output_entries).", "T") .Output(1, "selected_counts", "Merged prefix lengths, shape (N).", "T") .Output(2, "status", "Per-query device validation status, shape (N).", "T") @@ -2033,6 +2042,34 @@ static OpSchema PackedSparseAttentionIndexerMergeSchema() { if (hasInputShape(ctx, 0) && getInputShape(ctx, 0).dim_size() != 2) { fail_shape_inference("base_indices must have rank 2"); } + const bool append_indices = policy->s() == "append_indices"; + for (size_t index = 3; index <= 6; ++index) { + const bool provided = index < ctx.getNumInputs() && ctx.getInputType(index) != nullptr; + if (provided != (append_indices ? index >= 5 : index <= 4)) { + fail_shape_inference("Merge inputs must match policy_mode"); + } + } + const auto check_rows = [&](size_t input, int rank, size_t reference) { + if (!hasInputShape(ctx, input)) return; + const auto& shape = getInputShape(ctx, input); + if (shape.dim_size() != rank) fail_shape_inference("Merge input ", input, " must have rank ", rank); + if (hasInputShape(ctx, reference)) { + const auto& expected = getInputShape(ctx, reference); + if (expected.dim_size() > 0 && shape.dim(0).has_dim_value() && expected.dim(0).has_dim_value() && + shape.dim(0).dim_value() != expected.dim(0).dim_value()) { + fail_shape_inference("Merge input ", input, " has a mismatched row dimension"); + } + } + }; + check_rows(1, 1, 0); + check_rows(2, 1, 2); + if (append_indices) { + check_rows(5, 2, 2); + check_rows(6, 1, 2); + } else { + check_rows(3, 1, 2); + check_rows(4, 1, 2); + } if (hasInputShape(ctx, 2)) { const auto& rows = getInputShape(ctx, 2); if (rows.dim_size() != 1) fail_shape_inference("base_row_indices must have rank 1"); @@ -2358,7 +2395,6 @@ void SparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex fail_shape_inference("SparseAttentionIndexer: policy_mode '", policy_mode, "' requires exactly ", expected_outputs, " declared outputs, got ", ctx.getNumOutputs()); } - updateOutputElemType(ctx, sai::kSelectedIndices, ONNX_NAMESPACE::TensorProto_DataType_INT32); const bool has_key = SparseAttentionIndexerHasInput(ctx, sai::kKey); @@ -2378,6 +2414,9 @@ void SparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex if (query_width_dim.has_dim_value() && query_width_dim.dim_value() <= 0) { fail_shape_inference("SparseAttentionIndexer: query width must be > 0, got ", query_width_dim.dim_value()); } + if (query_width_dim.has_dim_value() && query_width_dim.dim_value() > std::numeric_limits::max()) { + fail_shape_inference("SparseAttentionIndexer: query width must be no greater than INT_MAX"); + } if (query_norm_shape == nullptr) { return; } @@ -2420,7 +2459,12 @@ void SparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex } else { auto* total_dim = present_shape.add_dim(); if (past_key_shape->dim(1).has_dim_value() && sequence_dim.has_dim_value()) { - total_dim->set_dim_value(past_key_shape->dim(1).dim_value() + sequence_dim.dim_value()); + const int64_t past_length = past_key_shape->dim(1).dim_value(); + const int64_t sequence_length = sequence_dim.dim_value(); + if (past_length > std::numeric_limits::max() - sequence_length) { + fail_shape_inference("SparseAttentionIndexer: present key sequence length exceeds INT64_MAX"); + } + total_dim->set_dim_value(past_length + sequence_length); } } SparseAttentionIndexerAppendDim(present_shape, head_size_dim); @@ -2475,7 +2519,11 @@ void SparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex if (SparseAttentionIndexerHasInput(ctx, sai::kPastSequenceLength)) { *entry_dim = past_compressed_shape->dim(1); } else if (plan_known && past_compressed_shape->dim(1).has_dim_value()) { - entry_dim->set_dim_value(past_compressed_shape->dim(1).dim_value() + plan.new_window_count); + const int64_t past_length = past_compressed_shape->dim(1).dim_value(); + if (past_length > std::numeric_limits::max() - plan.new_window_count) { + fail_shape_inference("SparseAttentionIndexer: present compressed-key length exceeds INT64_MAX"); + } + entry_dim->set_dim_value(past_length + plan.new_window_count); } SparseAttentionIndexerAppendDim(present_shape, head_size_dim); updateOutputShape(ctx, sai::kPresentKey, present_shape); @@ -2512,7 +2560,8 @@ decoders, chosen with the policy_mode attribute: the current window its "Cb" half), normalizes and rotates the entry, appends it to the compressed-key state, scores the queries against every compressed entry with sum_h w_h * ReLU(q_h . k), masks the entries a query may not attend to and emits the index_topk - highest scoring entry indices. + highest scoring entry indices. If a channel's softmax denominator is non-positive or non-finite + (for example, all logits are -inf or any logit is +inf), its pooled value is zero. Common contract: * selected_indices is int32 with a fixed capacity that only depends on attributes: @@ -2718,6 +2767,11 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference "'"); } const bool is_qsa = policy == psai::Policy::kQsa; + const std::string state_policy = getAttribute(ctx, "indexshare_state_policy", std::string("default")); + if ((state_policy != "default" && state_policy != "hold" && state_policy != "advance") || + (is_qsa && state_policy == "advance")) { + fail_shape_inference("PackedSparseAttentionIndexer: invalid indexshare_state_policy for policy_mode"); + } const bool has_indexshare = PackedSparseAttentionIndexerHasInput(ctx, psai::kIndexShareMode); const int64_t output_capacity = getAttribute(ctx, "max_output_entries", static_cast(0)); @@ -2726,7 +2780,8 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference fail_shape_inference("PackedSparseAttentionIndexer: all IndexShare inputs must be provided together"); } } - if (has_indexshare != (output_capacity > 0) || output_capacity > std::numeric_limits::max()) { + if (has_indexshare != (output_capacity > 0) || output_capacity < 0 || + output_capacity > std::numeric_limits::max()) { fail_shape_inference("PackedSparseAttentionIndexer: IndexShare requires a positive max_output_entries"); } if (has_indexshare != ctx.hasOutput(psai::kIndexShareStatus)) { @@ -2779,8 +2834,7 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference } // Strict policy input validation: every fixed slot required by every policy must be provided; - // csa-only slots must be provided iff policy_mode is 'csa'; position_ids is optional for 'qsa' - // and required for 'csa'. + // csa-only slots must be provided iff policy_mode is 'csa'. for (int index : {psai::kQuery, psai::kQueryNormWeight, psai::kKeyNormWeight, psai::kCosCache, psai::kSinCache, psai::kCumulativeSequenceLengths, psai::kPastSequenceLengths, psai::kPastKeyState, psai::kPastKvBuffer, psai::kPastStateLengths}) { @@ -2799,9 +2853,10 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference : " is required when policy_mode is 'csa'"); } } - if (!is_qsa && !PackedSparseAttentionIndexerHasInput(ctx, psai::kPositionIds)) { + if (PackedSparseAttentionIndexerHasInput(ctx, psai::kPositionIds) == is_qsa) { fail_shape_inference("PackedSparseAttentionIndexer: input ", psai::kPositionIds, - " (position_ids) is required when policy_mode is 'csa'"); + is_qsa ? " (position_ids) must be omitted when policy_mode is 'qsa'" + : " (position_ids) is required when policy_mode is 'csa'"); } const bool has_capture_count = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateCaptureCount); const bool has_state_update_active = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateActive); @@ -2884,6 +2939,10 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference PackedSparseAttentionIndexerShape(ctx, psai::kStateUpdateCaptureCount, 1); const auto* state_update_active_shape = PackedSparseAttentionIndexerShape(ctx, psai::kStateUpdateActive, 1); + const auto* mode_shape = PackedSparseAttentionIndexerShape(ctx, psai::kIndexShareMode, 1); + const auto* merged_indices_shape = PackedSparseAttentionIndexerShape(ctx, psai::kMergedIndices, 2); + const auto* merged_counts_shape = PackedSparseAttentionIndexerShape(ctx, psai::kMergedCounts, 1); + const auto* merged_status_shape = PackedSparseAttentionIndexerShape(ctx, psai::kMergedStatus, 1); auto require_equal_dims = [](const ONNX_NAMESPACE::TensorShapeProto* lhs, int lhs_index, const ONNX_NAMESPACE::TensorShapeProto* rhs, int rhs_index, @@ -2903,6 +2962,15 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference }; const int64_t buffer_capacity = psai::GenericBufferCapacity(compress_ratio); + require_dim_value(mode_shape, 0, 1, "indexshare_mode must have shape {1}"); + require_dim_value(merged_indices_shape, 1, output_capacity, "merged_indices width must equal max_output_entries"); + require_equal_dims(merged_indices_shape, 0, merged_counts_shape, 0, "merged_counts must match merged_indices rows"); + require_equal_dims(merged_indices_shape, 0, merged_status_shape, 0, "merged_status must match merged_indices rows"); + if (merged_indices_shape != nullptr && merged_indices_shape->dim(0).has_dim_value() && + merged_indices_shape->dim(0).dim_value() != 0) { + require_equal_dims(merged_indices_shape, 0, past_sequence_shape, 0, + "Nonempty merged_indices must have one row per request"); + } if (!has_indexshare) require_equal_dims(key_shape, 0, query_shape, 0, "key dimension 0 must equal query dimension 0"); require_equal_dims(query_norm_shape, 0, key_norm_shape, 0, "query_norm_weight and key_norm_weight dimensions must match"); @@ -2948,6 +3016,8 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference } } if (!is_qsa) { + require_equal_dims(key_shape, 0, gate_shape, 0, "gate dimension 0 must equal key dimension 0"); + require_equal_dims(key_shape, 0, position_ids_shape, 0, "position_ids dimension 0 must equal key dimension 0"); if (!has_indexshare) require_equal_dims(gate_shape, 0, query_shape, 0, "gate dimension 0 must equal query dimension 0"); require_equal_dims(head_weights_shape, 0, query_shape, 0, "head_weights dimension 0 must equal query dimension 0"); @@ -3097,7 +3167,8 @@ request's own packed token range and fixed-capacity state slice: sum_h ReLU(q_h . k), the token_budget / compress_ratio highest scoring blocks are kept, and their token indices are emitted (request-local logical positions, i.e. the same numbering as past_sequence_lengths + local offset) followed by the causally visible tokens of the trailing - incomplete block. + incomplete block. QSA positions are the request-local logical cache positions derived from + past_sequence_lengths and cumulative_sequence_lengths; position_ids must be omitted. policy_mode = "csa" ("compressed sparse attention" block indexer) Applies the same window-plan arithmetic as SparseAttentionIndexer (overlap/leftover/new window @@ -3219,8 +3290,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Input(7, "past_sequence_lengths", "Device-resident number of tokens already processed for each request before this call, with " - "shape (batch_size). Used as the default absolute query position when position_ids is omitted " - "(policy_mode 'qsa'), and to validate state consistency.", + "shape (batch_size). Used to derive request-local logical query positions for policy_mode 'qsa' " + "and to validate state consistency.", "M") .Input(8, "gate", @@ -3240,8 +3311,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Input(11, "position_ids", - "Optional for policy_mode 'qsa', required for policy_mode 'csa': absolute position of every " - "packed query, with shape (total_tokens).", + "Only for policy_mode 'csa': absolute position of every packed query, with shape (total_tokens).", "I", OpSchema::Optional) .Input(12, @@ -3406,58 +3476,68 @@ constexpr const char* RotaryEmbedding_ver1_doc = R"DOC( RotaryEmbedding is the implementation of rotary positional embeddings (RoPE). The positions are represented as rotation matrices that are multiplied to query and key before the inner product of query and key is taken. )DOC"; +static OpSchema RotaryEmbeddingSchema() { + return OpSchema() + .SetDoc(RotaryEmbedding_ver1_doc) + .Attr("scale", + "Custom scale will be used if specified. Default value is 1.0", + AttributeProto::FLOAT, + OPTIONAL_VALUE) + .Attr("interleaved", + "Indicates whether the input has real and imaginary parts interleaved. " + "Default value is 0 (False), meaning the first half of the input consists of real values " + "and the second half consists of imaginary values.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("rotary_embedding_dim", + "Rotary embedding dimension. Default value is 0.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("num_heads", + "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("is_packed_batching", + "ragged batch inputs or not. Default value is 0", + AttributeProto::INT, + OPTIONAL_VALUE) + .Input(0, + "input", + "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", + "T") + .Input(1, + "position_ids", + "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", + "M") + .Input(2, + "cos_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Input(3, + "sin_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Output(0, + "output", + "tensor with same shape as input.", + "T") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") + .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateShapeFromInputToOutput(ctx, 0, 0); + }); +} + +ONNX_MS_OPERATOR_SET_SCHEMA(RotaryEmbedding, 1, RotaryEmbeddingSchema()); ONNX_MS_OPERATOR_SET_SCHEMA( - RotaryEmbedding, 1, - OpSchema() - .SetDoc(RotaryEmbedding_ver1_doc) - .Attr("scale", - "Custom scale will be used if specified. Default value is 1.0", - AttributeProto::FLOAT, - OPTIONAL_VALUE) - .Attr("interleaved", - "Indicates whether the input has real and imaginary parts interleaved. " - "Default value is 0 (False), meaning the first half of the input consists of real values " - "and the second half consists of imaginary values.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("rotary_embedding_dim", - "Rotary embedding dimension. Default value is 0.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("num_heads", - "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("is_packed_batching", - "ragged batch inputs or not. Default value is 0", - AttributeProto::INT, - OPTIONAL_VALUE) - .Input(0, - "input", - "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", - "T") - .Input(1, - "position_ids", - "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", - "M") - .Input(2, - "cos_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Input(3, - "sin_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Output(0, - "output", - "tensor with same shape as input.", - "T") - .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") - .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - propagateShapeFromInputToOutput(ctx, 0, 0); - })); + RotaryEmbedding, 2, + RotaryEmbeddingSchema() + .Attr("use_fp32_arithmetic", + "Widen input and cache values to float32, round each multiplication and addition in float32, " + "then convert the result to the input type. Values are 0 (legacy arithmetic) or 1. " + "Currently supported by CUDA only.", + AttributeProto::INT, static_cast(0))); constexpr const char* MRotaryEmbedding_ver1_doc = R"DOC( MRotaryEmbedding is the fused implementation of Multimodal Rotary Positional Embeddings (M-RoPE) used by the @@ -4115,10 +4195,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( if (ctx.getNumOutputs() > 1) { propagateElemTypeFromInputToOutput(ctx, 0, 1); } - if (ctx.getNumOutputs() > 2) { - propagateElemTypeFromInputToOutput(ctx, 3, 2); - } - const int64_t max_ngram_size = getAttribute(ctx, "max_ngram_size", int64_t{-1}); const int64_t n_head_per_ngram = getAttribute(ctx, "n_head_per_ngram", int64_t{-1}); if (max_ngram_size < 2) { @@ -4484,13 +4560,12 @@ channels-first layout; only the memory layout differs. constexpr const char* EngramGate_ver1_doc = R"DOC( Fuses the Engram gate. -The op consumes already projected keys and hidden-state queries in either dense -(batch_size, sequence_length, hc_mult, hidden_size) or packed -(total_tokens, hc_mult, hidden_size) layout. The projected value has the corresponding -(batch_size, sequence_length, hidden_size) or (total_tokens, hidden_size) layout and is shared by -every hyper-connection. The key and value projections stay outside the op so they can run on the -execution provider's tuned MatMul (weight prepacking, tensor cores, quantized weights) and so the -value projection is computed once per token instead of once per hyper-connection. +The op consumes already projected keys in (batch_size, sequence_length, hc_mult, hidden_size) layout, +the hidden-state queries in the same layout, an already projected value in +(batch_size, sequence_length, hidden_size) layout that is shared by every hyper-connection, and the two +RMSNorm scales. The key and value projections stay outside the op so they can run on the execution +provider's tuned MatMul (weight prepacking, tensor cores, quantized weights) and so the value +projection is computed once per token instead of once per hyper-connection. It computes the Engram gate: @@ -4513,17 +4588,16 @@ ONNX_MS_OPERATOR_SET_SCHEMA( 1.0e-5f) .Input(0, "key", - "Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size) " - "or (total_tokens, hc_mult, hidden_size).", + "Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size).", "T") .Input(1, "query", - "Hidden-state queries with the same shape as key.", + "Hidden-state queries with shape (batch_size, sequence_length, hc_mult, hidden_size).", "T") .Input(2, "value", "Projected Engram value shared by every hyper-connection, with shape " - "(batch_size, sequence_length, hidden_size) or (total_tokens, hidden_size).", + "(batch_size, sequence_length, hidden_size).", "T") .Input(3, "key_norm_scale", @@ -4541,11 +4615,12 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Output(0, "output", - "Gated value tensor with the same shape as key.", + "Gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size).", "T") .Output(1, "gated_value_normed", - "Optional RMS-normalized gated value tensor with the same shape as key.", + "Optional RMS-normalized gated value tensor with shape " + "(batch_size, sequence_length, hc_mult, hidden_size).", "T", OpSchema::Optional) .TypeConstraint("T", @@ -4559,8 +4634,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( if (hasInputShape(ctx, 0)) { const auto& key_shape = getInputShape(ctx, 0); - if (key_shape.dim_size() != 3 && key_shape.dim_size() != 4) { - fail_shape_inference("EngramGate: key must have rank 3 or 4"); + if (key_shape.dim_size() != 4) { + fail_shape_inference("EngramGate: key must have rank 4"); } propagateShapeFromInputToOutput(ctx, 0, 0); if (ctx.getNumOutputs() > 1) { @@ -4569,39 +4644,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } if (hasInputShape(ctx, 1)) { const auto& query_shape = getInputShape(ctx, 1); - if (query_shape.dim_size() != 3 && query_shape.dim_size() != 4) { - fail_shape_inference("EngramGate: query must have rank 3 or 4"); - } - if (hasInputShape(ctx, 0)) { - const auto& key_shape = getInputShape(ctx, 0); - if (query_shape.dim_size() != key_shape.dim_size()) { - fail_shape_inference("EngramGate: query must have the same rank as key"); - } - for (int i = 0; i < key_shape.dim_size(); ++i) { - if (query_shape.dim(i).has_dim_value() && key_shape.dim(i).has_dim_value() && - query_shape.dim(i).dim_value() != key_shape.dim(i).dim_value()) { - fail_shape_inference("EngramGate: query must have the same shape as key"); - } - } + if (query_shape.dim_size() != 4) { + fail_shape_inference("EngramGate: query must have rank 4"); } } if (hasInputShape(ctx, 2)) { const auto& value_shape = getInputShape(ctx, 2); - if (value_shape.dim_size() != 2 && value_shape.dim_size() != 3) { - fail_shape_inference("EngramGate: value must have rank 2 or 3"); - } - if (hasInputShape(ctx, 0) && value_shape.dim_size() + 1 != getInputShape(ctx, 0).dim_size()) { - fail_shape_inference("EngramGate: value rank must be one less than key rank"); - } - if (hasInputShape(ctx, 0) && value_shape.dim_size() + 1 == getInputShape(ctx, 0).dim_size()) { - const auto& key_shape = getInputShape(ctx, 0); - for (int i = 0; i < value_shape.dim_size(); ++i) { - const int key_dim = i == value_shape.dim_size() - 1 ? i + 1 : i; - if (value_shape.dim(i).has_dim_value() && key_shape.dim(key_dim).has_dim_value() && - value_shape.dim(i).dim_value() != key_shape.dim(key_dim).dim_value()) { - fail_shape_inference("EngramGate: value must match key's token dimensions and hidden size"); - } - } + if (value_shape.dim_size() != 3) { + fail_shape_inference("EngramGate: value must have rank 3"); } } })); @@ -4808,9 +4858,7 @@ transaction-safe only when the whole operator call is unconditionally committed; may select a prefix or roll back must preserve initial_state and replay the compact state update. When state_update_capacity is positive, capture_count is required with shape (batch_size), and -state_update has shape (batch_size, state_update_capacity, channels). The optional CPU int32 -state_update_active input has shape (1). When present and zero, state_update remains all zero; -when omitted or nonzero, capture proceeds normally. +state_update has shape (batch_size, state_update_capacity, channels). For request b, slots [0, clamp(capture_count[b], 0, min(state_update_capacity, sequence_length[b]))) contain the original local input token values. These values represent the append component of each shift-left-and-append state transition. @@ -4886,12 +4934,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "and state_update_capacity. Required exactly when state_update_capacity is positive.", "M", OpSchema::Optional) - .Input(6, - "state_update_active", - "Optional CPU int32 tensor with shape (1). Zero disables compact state update " - "capture for this call; omitted or nonzero enables capture.", - "M", - OpSchema::Optional) .Output(0, "output", "Token-major convolution output with the same shape as input.", @@ -4912,8 +4954,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Constrain input and output types to float tensors.") .TypeConstraint("M", {"tensor(int32)"}, - "Constrain cumulative_sequence_length and capture_count to device int32 tensors, " - "and state_update_active to a CPU int32 tensor.") + "Constrain cumulative_sequence_length and capture_count to device int32 tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); propagateElemTypeFromInputToOutput(ctx, 0, 1); @@ -4974,14 +5015,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( fail_shape_inference("VarlenCausalConvWithState: capture_count must have shape (batch_size)"); } } - if (hasInputShape(ctx, 6)) { - auto& state_update_active_shape = getInputShape(ctx, 6); - if (state_update_active_shape.dim_size() != 1 || - (state_update_active_shape.dim(0).has_dim_value() && - state_update_active_shape.dim(0).dim_value() != 1)) { - fail_shape_inference("VarlenCausalConvWithState: state_update_active must have shape (1)"); - } - } TensorShapeProto state_shape; // batch_size = cumulative_sequence_length.dim(0) - 1 @@ -5195,6 +5228,73 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } })); +ONNX_MS_OPERATOR_SET_SCHEMA( + GatedDeltaNetStateReplay, 1, + OpSchema() + .SetDoc(R"DOC( +Replay one request's captured GatedDeltaNet transitions into an inactive recurrent-state row. +The WebGPU implementation uses FP32 backing tensors and CPU metadata; it does not recompute +normalization, gates, beta, or retrieval. It waits for checked GPU completion before returning. + +`metadata` is int64 [11], ordered as: + source_offset, destination_offset, decay_offset, key_offset, delta_offset, + Hv, Dv, Dk, Hk, capacity, kept_count. +Offsets are FP32-element offsets into the full backing tensors, not storage-binding byte offsets. +All offsets are nonnegative; dimensions and capacity are positive; 1 <= kept_count <= capacity. +The caller must ensure that all kept transitions were actually captured. + +The source_state selected by source_offset must be the state snapshot from immediately before +the captured GatedDeltaNet transitions began. Using the producer's final state or another +same-shaped snapshot is invalid because captured delta is state-dependent. +The capsule geometry (Hv, Dv, Dk, Hk) and capacity must match the producer that generated it. +capacity is the original capsule capture/layout capacity. kept_count only selects the leading +transitions to replay; kept_count must NOT be substituted for capacity when deriving the +decay/key/delta section offsets. + +State rows have layout [Hv, Dv, Dk]. The capsule contains decay [capacity, Hv], +key [capacity, Hk, Dk], and delta [capacity, Hv, Dv] starting at their respective offsets. +For every state element (h,v,i): + hk = floor(h * Hk / Hv) + s = source[h,v,i] + for t = 0 .. kept_count-1: + s *= decay[t,h] + s += key[t,hk,i] * delta[t,h,v] + destination[h,v,i] = s +Unused capsule tail entries are not read. All other destination elements are preserved. + +The output must be explicitly prebound to the same full allocation as destination_backing. +Destination backing must be distinct from source and capsule backing. This first implementation +supports one descriptor, uint32 shader indexing and no graph capture. +Source and capsule may share read-only backing; storage bindings are deduplicated only when +their BufferViews have the same Program binding owner, not merely equal raw buffer handles. +)DOC") + .Input(0, "source_backing", + "Full rank-1 FP32 backing tensor containing the state snapshot from immediately before capture.", "T") + .Input(1, "capsule_backing", "Full rank-1 FP32 compact-update backing tensor.", "T") + .Input(2, "destination_backing", "Full rank-1 FP32 inactive destination-state backing tensor.", "T") + .Input(3, "metadata", "CPU int64 [11] containing offsets, geometry, capacity and kept_count.", "TI") + .Output(0, "destination_out", "Same backing allocation and shape as destination_backing.", "T") + .TypeConstraint("T", {"tensor(float)"}, "State and capsule tensors are FP32.") + .TypeConstraint("TI", {"tensor(int64)"}, "CPU replay metadata is int64.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 2, 0); + for (size_t index = 0; index < 4; ++index) { + if (!hasInputShape(ctx, index)) { + continue; + } + const auto& shape = getInputShape(ctx, index); + if (shape.dim_size() != 1) { + fail_shape_inference("GatedDeltaNetStateReplay inputs must have rank 1."); + } + if (index == 3 && shape.dim(0).has_dim_value() && shape.dim(0).dim_value() != 11) { + fail_shape_inference("GatedDeltaNetStateReplay metadata must have shape [11]."); + } + } + if (hasInputShape(ctx, 2)) { + propagateShapeFromInputToOutput(ctx, 2, 0); + } + })); + constexpr const char* GatedDeltaNet_ver1_doc = R"DOC( Packed (token-major) gated delta network / linear attention with an explicit recurrent state. Implemented by CUDA and native WebGPU execution providers. WebGPU supports float and float16 @@ -5209,6 +5309,7 @@ rather than from attributes: Alternatively, key and value may be omitted and query contains packed QKV with shape `[total_tokens, 2 * num_heads_q * head_size_qk + num_heads_v * head_size_v]`. +`initial_state` is required for packed QKV to determine the head counts and head sizes. The leading token axis may instead be spelled as an explicit `[batch_size, sequence_length]` pair, making query/key/value (and the output) rank 4 and decay/beta rank 3. The memory layout @@ -5363,16 +5464,20 @@ ONNX_MS_OPERATOR_SET_SCHEMA( updateOutputElemType(ctx, 2, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); } - if (!hasInputShape(ctx, 0)) { - return; - } - const auto& query_shape = getInputShape(ctx, 0); - const bool has_key = ctx.getInputType(1) != nullptr; - const bool has_value = ctx.getInputType(2) != nullptr; + const bool has_key = ctx.getNumInputs() > 1 && ctx.getInputType(1) != nullptr; + const bool has_value = ctx.getNumInputs() > 2 && ctx.getInputType(2) != nullptr; if (has_key != has_value) { fail_shape_inference("GatedDeltaNet: key and value must both be present or both be omitted"); } const bool packed_qkv = !has_key; + const bool has_state = ctx.getNumInputs() > 6 && ctx.getInputType(6) != nullptr; + if (packed_qkv && !has_state) { + fail_shape_inference("GatedDeltaNet: initial_state is required for packed QKV"); + } + if (!hasInputShape(ctx, 0)) { + return; + } + const auto& query_shape = getInputShape(ctx, 0); const int rank = query_shape.dim_size(); const int token_dims = rank - (packed_qkv ? 1 : 2); if ((packed_qkv && rank != 2 && rank != 3) || @@ -5394,25 +5499,62 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } const ONNX_NAMESPACE::TensorShapeProto* state_shape = nullptr; - if (hasInputShape(ctx, 6)) { + if (ctx.getNumInputs() > 6 && hasInputShape(ctx, 6)) { state_shape = &getInputShape(ctx, 6); if (state_shape->dim_size() != 4) { fail_shape_inference("GatedDeltaNet: initial_state must have rank 4"); } } + if (packed_qkv) { + if (state_shape == nullptr) { + return; + } + if (query_shape.dim(token_dims).has_dim_value() && + state_shape->dim(1).has_dim_value() && state_shape->dim(2).has_dim_value() && + state_shape->dim(3).has_dim_value()) { + const int64_t packed_size = query_shape.dim(token_dims).dim_value(); + const int64_t num_heads_v = state_shape->dim(1).dim_value(); + const int64_t head_size_v = state_shape->dim(2).dim_value(); + const int64_t head_size_qk = state_shape->dim(3).dim_value(); + constexpr int64_t max_dimension = std::numeric_limits::max(); + if (num_heads_v <= 0 || head_size_v <= 0 || head_size_qk <= 0) { + fail_shape_inference("GatedDeltaNet: packed QKV head counts and head sizes must be positive"); + } + if (num_heads_v > max_dimension / head_size_v || head_size_qk > max_dimension / 2) { + fail_shape_inference("GatedDeltaNet: packed QKV dimensions overflow int64"); + } + const int64_t value_size = num_heads_v * head_size_v; + const int64_t query_key_size = 2 * head_size_qk; + if (packed_size <= value_size || (packed_size - value_size) % query_key_size != 0) { + fail_shape_inference( + "GatedDeltaNet: packed QKV last dimension must be 2 * num_heads_q * head_size_qk + " + "num_heads_v * head_size_v"); + } + const int64_t num_heads_q = (packed_size - value_size) / query_key_size; + if (num_heads_v % num_heads_q != 0) { + fail_shape_inference("GatedDeltaNet: num_heads_v must be a positive multiple of num_heads_q"); + } + } + const bool has_cu_seqlens = ctx.getNumInputs() > 3 && ctx.getInputType(3) != nullptr; + if (!has_cu_seqlens && token_dims == 1 && query_shape.dim(0).has_dim_value() && + state_shape->dim(0).has_dim_value()) { + const int64_t total_tokens = query_shape.dim(0).dim_value(); + const int64_t batch = state_shape->dim(0).dim_value(); + if (batch <= 0 || total_tokens <= 0 || total_tokens % batch != 0) { + fail_shape_inference( + "GatedDeltaNet: total_tokens must be divisible by batch, with positive token count and batch size"); + } + } + } + ONNX_NAMESPACE::TensorShapeProto out_shape; for (int i = 0; i < token_dims; ++i) { *out_shape.add_dim() = query_shape.dim(i); } if (packed_qkv) { - if (state_shape != nullptr) { - *out_shape.add_dim() = state_shape->dim(1); - *out_shape.add_dim() = state_shape->dim(2); - } else { - out_shape.add_dim(); - out_shape.add_dim(); - } + *out_shape.add_dim() = state_shape->dim(1); + *out_shape.add_dim() = state_shape->dim(2); } else { if (query_shape.dim(token_dims).has_dim_value() && value_shape->dim(token_dims).has_dim_value()) { @@ -5622,31 +5764,90 @@ Applies RMS normalization independently to each branch. X may use grouped shape (..., C, H), or flattened shape (..., C * H) when num_branches is specified. The optional scale may have shape (C * H), (C, H), or (H). Arithmetic is performed in float32 and the result is converted to T. +With use_onnx_arithmetic=1, compute FP32 square, mean, epsilon addition, square +root, division, and scale multiplication with separate rounding boundaries. +CUDA reuses the standard fast last-axis ReduceMean summation order. This mode +does not substitute reciprocal square root or change the reduction tree. +Optional Y_fp32 losslessly widens the already T-rounded Y to float32; it does +not bypass the output rounding boundary and requires use_onnx_arithmetic=1. +WebGPU currently supports only the default mode without Y_fp32. )DOC"; +static OpSchema BranchwiseRMSNormSchema() { + return OpSchema() + .SetDoc(BranchwiseRMSNorm_ver1_doc) + .Attr("epsilon", "Epsilon added before reciprocal square root.", + AttributeProto::FLOAT, 1e-5f) + .Attr("num_branches", + "Number of branches for flattened input. Omit or set to zero for grouped input.", + AttributeProto::INT, static_cast(0)) + .Attr("use_onnx_arithmetic", + "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. " + "Supported values are 0 (default) and 1.", + AttributeProto::INT, static_cast(0)) + .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") + .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", + "M", OpSchema::Optional) + .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") + .Output(1, "Y_fp32", + "Optional float32 widening of the already T-rounded Y. " + "Requires use_onnx_arithmetic=1.", + "F", OpSchema::Optional) + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain input and output to floating-point tensors.") + .TypeConstraint("M", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain scale to floating-point tensors.") + .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + const auto start = getAttribute(ctx, "input_start", 0); + const auto width = getAttribute(ctx, "input_width", 0); + if (start < 0 || width < 0 || (width == 0 && start != 0)) { + fail_shape_inference("input_start and input_width must describe a nonnegative slice"); + } + if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 || + getAttribute(ctx, "num_branches", 0) != 0)) { + fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input"); + } + if (hasInputShape(ctx, 0)) { + auto shape = getInputShape(ctx, 0); + if (width != 0) { + if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two"); + auto* last = shape.mutable_dim(shape.dim_size() - 1); + if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) { + fail_shape_inference("input slice exceeds the last dimension"); + } + last->set_dim_value(width); + } + updateOutputShape(ctx, 0, shape); + } + if (ctx.getNumOutputs() > 1) { + if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) { + fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1"); + } + updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + if (hasInputShape(ctx, 0)) { + auto shape = getInputShape(ctx, 0); + if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width); + updateOutputShape(ctx, 1, shape); + } + } + }); +} + +ONNX_MS_OPERATOR_SET_SCHEMA(BranchwiseRMSNorm, 1, BranchwiseRMSNormSchema()); ONNX_MS_OPERATOR_SET_SCHEMA( - BranchwiseRMSNorm, 1, - OpSchema() - .SetDoc(BranchwiseRMSNorm_ver1_doc) - .Attr("epsilon", "Epsilon added before reciprocal square root.", - AttributeProto::FLOAT, 1e-5f) - .Attr("num_branches", - "Number of branches for flattened input. Omit or set to zero for grouped input.", + BranchwiseRMSNorm, 2, + BranchwiseRMSNormSchema() + .Attr("input_start", "CUDA-only starting offset in the last input dimension.", AttributeProto::INT, static_cast(0)) - .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") - .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", - "M", OpSchema::Optional) - .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") - .TypeConstraint("T", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain input and output to floating-point tensors.") - .TypeConstraint("M", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain scale to floating-point tensors.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - propagateShapeFromInputToOutput(ctx, 0, 0); - })); + .Attr("input_width", + "CUDA-only width selected from the last input dimension before normalization. " + "Zero retains the full input. Requires grouped input and use_onnx_arithmetic=1. " + "Outputs retain the leading dimensions and use this last dimension.", + AttributeProto::INT, static_cast(0))); constexpr const char* ScaledSiLU_ver1_doc = R"DOC( Computes SiLU after scaling, with explicit T rounding: @@ -5667,7 +5868,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output to floating-point tensors.") - .TypeConstraint("M", + .TypeConstraint("M", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain the optional scale to floating-point tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { @@ -5698,7 +5899,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain streams and output to floating-point tensors.") - .TypeConstraint("M", + .TypeConstraint("M", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain mixing weights to floating-point tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { @@ -5739,7 +5940,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( constexpr const char* HyperConnectionPostMix_ver1_doc = R"DOC( Mixes existing streams and injects one branch output: Y[..., k, h] = sum_c(stream_mix[..., c, k] * streams[..., c, h]) - + post_mix[..., k, h] * branch_output[..., h]. + + post_mix[..., k, h] * block_output[..., h]. stream_mix is optional and defaults to the identity. post_mix has shape (..., C), (..., C, 1), or (..., C, H). No activation is applied. )DOC"; @@ -5760,7 +5961,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain streams, block output, and output to floating-point tensors.") - .TypeConstraint("M", + .TypeConstraint("M", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain mixing weights to floating-point tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index abc7e84188962..c38c4d1e1936e 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1435,7 +1435,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( MoE, 1, OpSchema() .SetDoc(MoE_ver1_doc) - .Attr("activation_type", "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. Default is relu", AttributeProto::STRING, std::string("relu")) + .Attr("activation_type", "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu", AttributeProto::STRING, std::string("relu")) .Attr("swiglu_fusion", "0: not fused, 1: fused and interleaved. 2: fused and not interleaved.", AttributeProto::INT, static_cast(0)) .Attr("swiglu_limit", "The limit used to clamp in SwiGLU. No clamp when limit is not provided.", AttributeProto::FLOAT, OPTIONAL_VALUE) .Attr("activation_alpha", "Alpha parameter used in activation function.", AttributeProto::FLOAT, 1.0f) @@ -1488,13 +1488,15 @@ constexpr const char* qMoE_ver1_doc = R"DOC( A positive block_size selects square block scaling with float32, float16, or bfloat16 fc*_scales tensors shaped [num_experts, ceil(N / block_size), ceil(K / block_size)]: dequantized_weight[e, n, k] = float(weight[e, n, k]) * scale[e, n / block_size, k / block_size] - Partial edge blocks are allowed. No zero points or activation scales are used. + Partial edge blocks are allowed. No zero points or activation scales are accepted. Without a positive block_size, FP8 uses the legacy per-expert fc*_global_scale inputs instead. Block-scaled FP8 does not use global scales. Activations retain the input type (weight-only quantization). On CUDA SM80-or-newer GPUs, SiLU/SwiGLU block-scaled FP8 uses fused weight-only GEMV/GEMM by default. These kernels decode and scale FP8 tiles on chip without materializing full FP16/BF16 weight buffers. Set ORT_ENABLE_FP8_FUSED=0 before session creation to select the compact dense-dequant fallback. Other activation types and unsupported routing sizes retain the fallback. + The WebGPU block-FP8 kernel supports block_size=128 with float32 scales and rejects projections + that exceed 32-bit shader addressing or the device's per-dimension dispatch limit. Packed byte dimensions are computed as logical_element_count * effective_expert_weight_bits / 8. Weight rows must be byte-aligned. Zero-point rows are padded to a whole byte when necessary. @@ -1521,7 +1523,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema() .SetDoc(qMoE_ver1_doc) .Attr("activation_type", - "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. Default is relu", + "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu", AttributeProto::STRING, std::string("relu")) .Attr("k", @@ -2623,7 +2625,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA(GatherND, 1, .Output(0, "output", "Tensor of rank q-1+r-indices[-1].", "T") .TypeConstraint( "T", - OpSchema::all_tensor_types_ir4(), + OpSchema::all_tensor_types(), "Constrain input and output types to any tensor type.") .TypeConstraint( "Tind", diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 3cc51426adc4d..93b771a425e07 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -91,10 +91,12 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, PagedAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, LinearAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, LinearAttentionGate); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, BranchwiseRMSNorm); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 2, BranchwiseRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ScaledSiLU); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPreMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPostMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNet); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNetStateReplay); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedAdd); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, NGramHashMapping); @@ -113,6 +115,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RemovePadding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RestorePadding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Rfft); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RotaryEmbedding); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 2, RotaryEmbedding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MRotaryEmbedding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GemmaRotaryEmbedding); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SampleOp); @@ -121,10 +124,10 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipGroupNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipLayerNormalization); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipSimplifiedLayerNormalization); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SparseAttention); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SparsePagedAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SparseAttentionIndexer); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, PackedSparseAttentionIndexer); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, PackedSparseAttentionIndexerMerge); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SparsePagedAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SparseToDenseMatMul); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Tokenizer); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, TorchEmbedding); @@ -224,10 +227,12 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); @@ -248,6 +253,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); @@ -257,10 +263,10 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); - fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index 436c68685430c..66d02d0d19c8c 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -2028,6 +2028,12 @@ struct MLAS_HALF_GEMM_DATA_PARAMS { * Bias must be nullptr, and OutputProcessor must be nullptr. */ bool BIsBackendNativePacked = false; + /** + * Unpacked fp16 B is stored as N x K instead of K x N. + * This requires a supporting backend. MlasHalfGemmBatch throws if the + * selected backend does not consume this layout. + */ + bool BIsTransposed = false; }; /** @@ -2765,7 +2771,7 @@ MlasComputeFP16Erf( * @param algo GELU algorithm variant (exact erf or tanh approximation). */ void -MLASCALL +MLASCALL MlasComputeFP16Gelu( const MLAS_FP16* input, MLAS_FP16* output, diff --git a/onnxruntime/core/mlas/lib/halfgemm.cpp b/onnxruntime/core/mlas/lib/halfgemm.cpp index cad0b087aa408..8c3ad5e144259 100644 --- a/onnxruntime/core/mlas/lib/halfgemm.cpp +++ b/onnxruntime/core/mlas/lib/halfgemm.cpp @@ -156,6 +156,11 @@ MlasHalfGemmBatch( std::runtime_error, "backend-native halfgemm packed B is not supported by generic MLAS halfgemm"); } + if (DataParams[gemm_i].BIsTransposed) { + MLAS_THROW_EX( + std::runtime_error, + "transposed unpacked B is not supported by generic MLAS halfgemm"); + } } const MLAS_HALFGEMM_DISPATCH* dispatch = MlasHalfGemmGetDispatch(); diff --git a/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp index 14afd31b27c6f..e63d36e823866 100644 --- a/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp +++ b/onnxruntime/core/mlas/lib/kleidiai/halfgemm_kleidiai.cpp @@ -20,6 +20,7 @@ #include "kai/ukernels/matmul/kai_matmul.h" #include "kai/ukernels/matmul/kai_matmul_pack_rhs.h" #include "kai/ukernels/matmul/pack/kai_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme.h" +#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_x16p2vlx2b_x16_x16_sme.h" namespace { constexpr const char* Sve2p1HalfGemmKernelName = @@ -124,12 +125,15 @@ bool ReadSve2p1PackedRhsMetadata( return true; } -size_t GetPackedRhsSize(KaiHalfGemmBackend backend, size_t N, size_t K) { +size_t GetPackedRhsSize(KaiHalfGemmBackend backend, size_t N, size_t K, bool BIsTransposed) { switch (backend) { case KaiHalfGemmBackend::Sme: + if (BIsTransposed) { + return kai_get_rhs_packed_size_rhs_pack_nxk_x16p2vlx2b_x16_x16_sme(N, K); + } return kai_get_rhs_packed_size_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(N, K); case KaiHalfGemmBackend::Sve2p1: - return GetSve2p1PackedRhsSize(N, K); + return BIsTransposed ? 0 : GetSve2p1PackedRhsSize(N, K); case KaiHalfGemmBackend::None: return 0; } @@ -143,9 +147,14 @@ size_t PackRhs( size_t ldb_bytes, const MLAS_FP16* rhs, const MLAS_FP16* bias, - void* rhs_packed + void* rhs_packed, + bool BIsTransposed ) { if (backend == KaiHalfGemmBackend::Sve2p1) { + if (BIsTransposed) { + return 0; + } + const auto& packer = kai_rhs_pack_kxn_x16p16vsx2bx16_x16_x16_sve(); const kai_matmul_pack_rhs_uker_config config{}; const kai_matmul_pack_rhs_uker_rhs_packed_dim_args packed_shape{N, K}; @@ -164,9 +173,16 @@ size_t PackRhs( } const auto& hgemm = GetKleidiAIHgemmUKernel(); - kai_run_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme( - 1, N, K, hgemm.ukernel.get_nr(), hgemm.ukernel.get_kr(), hgemm.ukernel.get_sr(), ldb_bytes, - rhs, bias, nullptr, rhs_packed, 0, nullptr); + if (BIsTransposed) { + kai_run_rhs_pack_nxk_x16p2vlx2b_x16_x16_sme( + 1, N, K, hgemm.ukernel.get_nr(), hgemm.ukernel.get_kr(), hgemm.ukernel.get_sr(), ldb_bytes, + rhs, bias, nullptr, rhs_packed, 0, nullptr); + } else { + kai_run_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme( + 1, N, K, hgemm.ukernel.get_nr(), hgemm.ukernel.get_kr(), hgemm.ukernel.get_sr(), ldb_bytes, + rhs, bias, nullptr, rhs_packed, 0, nullptr); + } + return 0; } @@ -312,14 +328,17 @@ ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize( size_t N, size_t K ) { - if (TransA != CblasNoTrans || TransB != CblasNoTrans || N == 0 || K == 0) { + const auto backend = SelectKaiHalfGemmBackend(); + + if (TransA != CblasNoTrans || + (TransB != CblasNoTrans && (TransB != CblasTrans || backend != KaiHalfGemmBackend::Sme)) || + N == 0 || K == 0) { return 0; } - const auto backend = SelectKaiHalfGemmBackend(); return backend == KaiHalfGemmBackend::Sve2p1 ? GetSve2p1NativePackedRhsSize(N, K) - : GetPackedRhsSize(backend, N, K); + : GetPackedRhsSize(backend, N, K, TransB == CblasTrans); } bool @@ -333,18 +352,21 @@ ArmKleidiAI::MlasHalfGemmKleidiAIPackB( size_t ldb, void* PackedB ) { - if (TransA != CblasNoTrans || TransB != CblasNoTrans) { + const auto backend = SelectKaiHalfGemmBackend(); + + if (TransA != CblasNoTrans || + (TransB != CblasNoTrans && (TransB != CblasTrans || backend != KaiHalfGemmBackend::Sme))) { return false; } - if (PackedB == nullptr || B == nullptr || N == 0 || K == 0 || ldb < N) { + const size_t minimum_ldb = TransB == CblasTrans ? K : N; + if (PackedB == nullptr || B == nullptr || N == 0 || K == 0 || ldb < minimum_ldb) { return false; } - const auto backend = SelectKaiHalfGemmBackend(); const size_t packed_rhs_size = backend == KaiHalfGemmBackend::Sve2p1 ? GetSve2p1NativePackedRhsSize(N, K) - : GetPackedRhsSize(backend, N, K); + : GetPackedRhsSize(backend, N, K, TransB == CblasTrans); if (packed_rhs_size == 0) { return false; } @@ -362,7 +384,7 @@ ArmKleidiAI::MlasHalfGemmKleidiAIPackB( packed_rhs_data = static_cast(PackedB) + Sve2p1PackedRhsMetadataSize; } - const size_t vector_length = PackRhs(backend, N, K, ldb_bytes, B, zero_bias.data(), packed_rhs_data); + const size_t vector_length = PackRhs(backend, N, K, ldb_bytes, B, zero_bias.data(), packed_rhs_data, TransB == CblasTrans); if (backend == KaiHalfGemmBackend::Sve2p1) { std::memcpy(PackedB, &vector_length, sizeof(vector_length)); } @@ -408,11 +430,17 @@ ArmKleidiAI::MlasHalfGemmBatch( // this override accepts arbitrary MLAS batches. Validate every native-packed // RHS before any entry writes output so a later mismatch cannot partially modify C. bool needs_rhs_packing = false; + size_t packed_rhs_size = 0; for (size_t b = 0; b < BatchN; ++b) { const auto& data = DataParams[b]; if (data.OutputProcessor != nullptr) { return false; } + if (data.BIsTransposed && + (backend != KaiHalfGemmBackend::Sme || data.BIsBackendNativePacked || data.BIsPacked || + data.BIsfp32 || data.ldb < K)) { + return false; + } if (data.BIsBackendNativePacked && (data.ldb != 0 || data.Bias != nullptr)) { return false; } @@ -424,6 +452,11 @@ ArmKleidiAI::MlasHalfGemmBatch( return false; } } + const size_t data_packed_rhs_size = GetPackedRhsSize(backend, N, K, data.BIsTransposed); + if (data_packed_rhs_size == 0) { + return false; + } + packed_rhs_size = std::max(packed_rhs_size, data_packed_rhs_size); // Native-packed RHS is consumed directly below. Only allocate the // runtime RHS packing scratch when at least one batch entry needs it. needs_rhs_packing = needs_rhs_packing || !data.BIsBackendNativePacked; @@ -437,11 +470,6 @@ ArmKleidiAI::MlasHalfGemmBatch( KLEIDIAI_KERNEL_LOG(Sve2p1HalfGemmKernelName); } - const size_t packed_rhs_size = GetPackedRhsSize(backend, N, K); - if (packed_rhs_size == 0) { - return false; - } - // TODO: Plumb MLAS_ACTIVATION through this call site if MLAS_HALF_GEMM_DATA_PARAMS // grows fused activation support. const float clamp_min = -std::numeric_limits::infinity(); @@ -520,7 +548,8 @@ ArmKleidiAI::MlasHalfGemmBatch( rhs_vector_length = PackRhs( backend, N, K, ldb_bytes, rhs_base, - data.Bias != nullptr ? data.Bias : g_kai_half_tls.bias_zero.data(), rhs_packed_buffer); + data.Bias != nullptr ? data.Bias : g_kai_half_tls.bias_zero.data(), rhs_packed_buffer, + data.BIsTransposed); rhs_packed = rhs_packed_buffer; } diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 89ecf9295cf35..33a5bb69c0805 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -2152,6 +2152,54 @@ MlasNchwcCostToWorkIndex( return (BatchGroup * FilterSetCount + Set) * OutputHeight + Row; } +// +// Picks the filter set size for the NCHWc grouped convolution algorithms. +// Normally MaximumFilterSetSize, but when that would not produce enough work +// units (BatchCount * GroupCount * FilterSetCount * OutputHeight) to give +// every thread at least Target units, the size is halved -- repeatedly, down +// to 1 -- so the operation splits more finely. +// +// Pulled out of MLAS_NCHWC_GROUPED_CONV_ALGORITHM::ChooseFilterSetSize as a +// pure function of its inputs (no environment variable read, no object +// state) so the halving dispatch can be exercised directly by unit tests +// across a range of shapes, including ones with ragged block counts. +// + +inline +size_t +MlasNchwcChooseFilterSetSize( + size_t MaximumFilterSetSize, + size_t Target, + ptrdiff_t tids, + size_t BlockSize, + size_t OutputChannels, + size_t BatchCount, + size_t GroupCount, + size_t OutputHeight + ) +{ + if (Target == 0 || tids <= 1) { + return MaximumFilterSetSize; + } + + const size_t Wanted = static_cast(tids) * Target; + const size_t Blocks = OutputChannels / BlockSize; + + size_t Size = MaximumFilterSetSize; + + while (Size > 1) { + const size_t Sets = (Blocks + Size - 1) / Size; + + if (BatchCount * GroupCount * Sets * OutputHeight >= Wanted) { + break; + } + + Size /= 2; + } + + return Size; +} + // // Define the minimum floating point value (and its bit value equivalent) that // has no fractional bits. This number can be used for fast rounding of floating diff --git a/onnxruntime/core/mlas/lib/snchwc.cpp b/onnxruntime/core/mlas/lib/snchwc.cpp index bd4a2c68be63e..0922605a3b310 100644 --- a/onnxruntime/core/mlas/lib/snchwc.cpp +++ b/onnxruntime/core/mlas/lib/snchwc.cpp @@ -18,6 +18,93 @@ Module Name: #include "mlasi.h" #include "snchwc.h" +#include +#include +#include + +namespace { + +// +// Reads an environment variable into a caller supplied buffer. MSVC rejects +// getenv outright, so the secure variant is used there. +// + +bool +MlasNchwcGetEnvironmentVariable(const char* Name, char* Buffer, size_t BufferSize) +{ + assert(BufferSize > 0); +#if defined(_MSC_VER) + size_t Length = 0; + + if (getenv_s(&Length, Buffer, BufferSize, Name) != 0) { + return false; + } + + return (Length != 0); +#else + const char* Value = std::getenv(Name); + + if (Value == nullptr) { + return false; + } + + std::strncpy(Buffer, Value, BufferSize - 1); + Buffer[BufferSize - 1] = '\0'; + + return true; +#endif +} + +// +// Controls whether the output channel blocks are spread evenly over the filter +// sets. Enabled by default; set MLAS_NCHWC_FILTERSET_BALANCE=0 to restore the +// previous ragged-last-set behavior. The toggle exists so that the two +// partitionings can be compared from a single binary, which removes the build +// as a variable when measuring. +// + +bool +MlasNchwcUseBalancedFilterSets(void) +{ + static const bool Enabled = []() { + char Buffer[16]; + + if (!MlasNchwcGetEnvironmentVariable("MLAS_NCHWC_FILTERSET_BALANCE", Buffer, sizeof(Buffer))) { + return true; + } + + return (Buffer[0] != '0'); + }(); + + return Enabled; +} + +// +// Controls whether the filter set size may be reduced below FilterSetSize when +// a convolution would otherwise not produce enough work units to keep every +// thread busy. Work units are (filter set, output line) pairs, so a shallow +// output with few filter sets can leave threads idle. Disabled by default. +// + +size_t +MlasNchwcFilterSetTarget(void) +{ + static const size_t Target = []() { + char Buffer[16]; + + if (!MlasNchwcGetEnvironmentVariable("MLAS_NCHWC_FSS_TARGET", Buffer, sizeof(Buffer))) { + return size_t{0}; + } + + const int Parsed = std::atoi(Buffer); + return (Parsed > 0) ? static_cast(Parsed) : size_t{0}; + }(); + + return Target; +} + +} // namespace + // // Define the base thread context for NCWHc convolution or pooling operations. // @@ -539,10 +626,39 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM // reused for a given set of input inside the kernel. // - static constexpr size_t FilterSetSize = 4; + static constexpr size_t MaximumFilterSetSize = 4; + + // + // Number of output channel blocks handled by one kernel invocation. Larger + // is cheaper per multiply-accumulate because a broadcast input is reused + // across every filter block, but it also produces fewer, coarser work units. + // + const size_t FilterSetSize; const size_t FilterSetCount; + // + // Distribute the output channel blocks evenly over the filter sets. + // + // Assigning FilterSetSize blocks to every set leaves the remainder in the + // final set, so that set can be up to FilterSetSize times cheaper than its + // peers. Work units are partitioned by count rather than by cost, and the + // output line index varies fastest, so a whole ragged filter set is handed + // to a different thread than the full ones and that thread finishes early. + // + // Spreading the blocks evenly costs nothing: the number of filter sets is + // unchanged, so the kernel still broadcasts each input the same number of + // times and issues the same number of multiply-accumulates. Only the + // per-thread balance changes. + // + // Example, 96 channels with a block size of 16: 6 blocks previously split + // as 4 + 2 and now split as 3 + 3. + // + + const bool BalancedFilterSets; + const size_t FilterSetBase; + const size_t FilterSetRemainder; + // // Stores the current output line, filter cluster, and group that this thread // is operating on. @@ -554,15 +670,63 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM size_t WorkRemaining; size_t FilterCount; + // + // Picks the filter set size. Normally the maximum, but when the resulting + // work unit count would not give every thread a reasonable number of units + // the size is halved so the operation splits more finely. + // + + size_t ChooseFilterSetSize(void) const + { + // + // MlasNchwcGetBlockSize() returns one rather than zero on platforms + // without NCHWc support, precisely so this division is always safe. + // + + assert(BlockSize > 0); + + return MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, MlasNchwcFilterSetTarget(), + WorkBlock->tids, BlockSize, OutputChannels, + BatchCount, GroupCount, OutputHeight); + } + MLAS_NCHWC_GROUPED_CONV_ALGORITHM(const MLAS_NCHWC_CONV_WORK_BLOCK* WorkBlock) : MLAS_NCHWC_CONV_ALGORITHM(WorkBlock), - FilterSetCount((OutputChannels + (BlockSize * FilterSetSize) - 1) / (BlockSize * FilterSetSize)) + FilterSetSize(ChooseFilterSetSize()), + FilterSetCount((OutputChannels + (BlockSize * FilterSetSize) - 1) / (BlockSize * FilterSetSize)), + BalancedFilterSets(MlasNchwcUseBalancedFilterSets()), + FilterSetBase((OutputChannels / BlockSize) / FilterSetCount), + FilterSetRemainder((OutputChannels / BlockSize) % FilterSetCount) + { + } + + // + // Returns the index of the first output channel block owned by a filter set. + // + + size_t FilterSetBlockOffset(size_t Set) const { + if (!BalancedFilterSets) { + return Set * FilterSetSize; + } + + return Set * FilterSetBase + std::min(Set, FilterSetRemainder); } void ComputeFilterCount(void) { - FilterCount = std::min(FilterSetSize, (OutputChannels / BlockSize) - FilterSet * FilterSetSize); + if (!BalancedFilterSets) { + FilterCount = std::min(FilterSetSize, (OutputChannels / BlockSize) - FilterSet * FilterSetSize); + return; + } + + // + // FilterSetCount is the ceiling of the block count over FilterSetSize, + // so FilterSetBase never exceeds FilterSetSize and the kernel dispatch + // table is always in range. + // + + FilterCount = FilterSetBase + ((FilterSet < FilterSetRemainder) ? 1 : 0); } void SeekToWork(size_t WorkIndex) @@ -588,14 +752,16 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM Input += BatchGroup * InputChannels * InputSize; Output += BatchGroup * OutputChannels * OutputSize; - Output += BlockSize * FilterSet * FilterSetSize * OutputSize; + const size_t FilterBlockOffset = FilterSetBlockOffset(FilterSet); + + Output += BlockSize * FilterBlockOffset * OutputSize; Filter += Group * OutputChannels * InputChannels * KernelSize; - Filter += BlockSize * FilterSet * FilterSetSize * InputChannels * KernelSize; + Filter += BlockSize * FilterBlockOffset * InputChannels * KernelSize; if (Bias != nullptr) { Bias += Group * OutputChannels; - Bias += BlockSize * FilterSet * FilterSetSize; + Bias += BlockSize * FilterBlockOffset; } // @@ -688,6 +854,12 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM void PrepareWorkWeighted(ptrdiff_t Index) { + // Balanced sets are all equal size, so the uniform partitioner is exact. + if (BalancedFilterSets) { + PrepareWork(Index); + return; + } + const size_t TotalBlockedFilters = OutputChannels / BlockSize; const size_t LastSetFilterCount = TotalBlockedFilters - (FilterSetCount - 1) * FilterSetSize; @@ -719,7 +891,7 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM } }; -constexpr size_t MLAS_NCHWC_GROUPED_CONV_ALGORITHM::FilterSetSize; +constexpr size_t MLAS_NCHWC_GROUPED_CONV_ALGORITHM::MaximumFilterSetSize; // // Implementation of the direct convolution algorithm where the input buffer is @@ -2055,3 +2227,4 @@ MlasPoolAverageIncludePadFloatKernel( } #endif + diff --git a/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx_common.h b/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx_common.h index e2120a7b5a603..6c4fc9c801d64 100644 --- a/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx_common.h +++ b/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx_common.h @@ -1,4 +1,5 @@ #pragma once +#include "core/common/safeint.h" #include "qnbitgemm.h" #include "sqnbitgemm_q8_block.h" @@ -72,21 +73,22 @@ QNBitGemmPackQuantBDataSize( ) { const size_t BlockCountK = MlasDivRoundup(K, BlkLen); + const size_t BlkDataSize = MlasQNBitBlkDataSizeInBytes(BlkBitWidth, BlkLen); + const SafeInt BlockCount = SafeInt(N) * BlockCountK; if (ComputeType == SQNBIT_CompInt8) { - size_t PackedQuantBDataSize = N * BlockCountK * MlasQNBitBlkDataSizeInBytes(BlkBitWidth, BlkLen); - const size_t ScaleSize = N * BlockCountK * sizeof(float); - size_t BlkSumSize = MlasDivRoundup(N, 16) * BlockCountK * 16 * sizeof(float); + SafeInt PackedQuantBDataSize = BlockCount * BlkDataSize; + const SafeInt ScaleSize = BlockCount * sizeof(float); + SafeInt BlkSumSize = SafeInt(MlasDivRoundup(N, 16)) * BlockCountK * 16 * sizeof(float); // avx512 requires alignment on a 64-byte boundary constexpr size_t PackedQuantBDataAlignment = 64; - PackedQuantBDataSize += PackedQuantBDataAlignment - 1; + PackedQuantBDataSize += SafeInt(PackedQuantBDataAlignment) - 1; constexpr size_t BlkSumAlignment = MlasQNBitQuantBBlkSumAlignment(); - BlkSumSize += BlkSumAlignment - 1; + BlkSumSize += SafeInt(BlkSumAlignment) - 1; - return PackedQuantBDataSize + ScaleSize + BlkSumSize; + return static_cast(PackedQuantBDataSize + ScaleSize + BlkSumSize); } else { - const size_t PackedQuantBDataSize = N * BlockCountK * MlasQNBitBlkDataSizeInBytes(BlkBitWidth, BlkLen); - return PackedQuantBDataSize; + return static_cast(BlockCount * BlkDataSize); } } diff --git a/onnxruntime/core/optimizer/attention_fusion_helper.h b/onnxruntime/core/optimizer/attention_fusion_helper.h index 814c2c2ea3259..4e1850bf4a267 100644 --- a/onnxruntime/core/optimizer/attention_fusion_helper.h +++ b/onnxruntime/core/optimizer/attention_fusion_helper.h @@ -400,7 +400,7 @@ bool MatchUnidirMaskSubgraph(const Graph& graph, const Node& add_node, MatchUnid } std::vector path1{ - {0, 0, "Cast", {9, 13, 19, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}, kOnnxDomain}, {0, 0, "Slice", {10, 11, 13}, kOnnxDomain}, // Last Slice {0, 0, "Slice", {10, 11, 13}, kOnnxDomain}, // Mask Slice {0, 1, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index 0dcae0be3e9ac..c315e121bb9d6 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include @@ -675,13 +676,11 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, // Build the TensorProto that corresponds to the computed OrtValue and add it as initializer to the graph. auto* constant_arg_out = node->MutableOutputDefs()[output_idx]; const Tensor& out_tensor = ort_value.Get(); - // Fallback to deep copy for aliased tensors (e.g., Reshape, Unsqueeze) - // to prevent dangling pointers when parent initializers are freed. - const bool use_tensor_buffer = out_tensor.OwnsBuffer(); + assert(out_tensor.OwnsBuffer()); ONNX_NAMESPACE::TensorProto out_tensorproto = utils::TensorToTensorProto( out_tensor, constant_arg_out->Name(), - use_tensor_buffer); + /*use_tensor_buffer=*/true); ONNX_NAMESPACE::TensorShapeProto result_shape; for (auto& dim : out_tensor.Shape().GetDims()) { @@ -689,7 +688,7 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, } constant_arg_out->SetShape(result_shape); - // Inlined data (small or deep-copied tensors) does not need to hold an OrtValue reference. + // Inline data does not need an OrtValue to keep a backing buffer alive. if (!utils::HasExternalData(out_tensorproto)) { ORT_THROW_IF_ERROR(graph.AddInitializedOrtValue(out_tensorproto, OrtValue())); } else { diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 0101a4f8b9da0..045ac89970b31 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -248,8 +248,8 @@ static bool MatchPositionEmbeddingSubgraphsFromGather( std::vector parent_path_1{ {0, 1, "Expand", {8, 13}, kOnnxDomain}, {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, - {0, 0, "Cast", {9, 13, 19, 21, 23, 24, 25}, kOnnxDomain}, - {0, 0, "Squeeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}, kOnnxDomain}, + {0, 0, "Squeeze", {1, 11, 13, 21, 23, 24, 25, 28}, kOnnxDomain}, {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "NonZero", {9, 13}, kOnnxDomain}, {0, 0, "ConstantOfShape", {9, 20, 21, 23, 24, 25}, kOnnxDomain}, @@ -260,7 +260,7 @@ static bool MatchPositionEmbeddingSubgraphsFromGather( std::vector parent_path_2{ {0, 1, "Expand", {8, 13}, kOnnxDomain}, {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, - {0, 0, "Squeeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Squeeze", {1, 11, 13, 21, 23, 24, 25, 28}, kOnnxDomain}, {0, 0, "Transpose", {1, 13, 21, 23, 24, 25}, kOnnxDomain}, {0, 0, "NonZero", {9, 13}, kOnnxDomain}, {0, 0, "ConstantOfShape", {9, 20, 21, 23, 24, 25}, kOnnxDomain}, @@ -272,15 +272,15 @@ static bool MatchPositionEmbeddingSubgraphsFromGather( std::vector parent_path_3{ {0, 1, "Expand", {8, 13}, kOnnxDomain}, {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, - {0, 0, "Range", {11, 27}, kOnnxDomain}, - {0, 1, "Cast", {9, 13, 19, 21, 23, 24, 25}, kOnnxDomain}, + {0, 0, "Range", {11, 27, 28}, kOnnxDomain}, + {0, 1, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}, kOnnxDomain}, {0, 0, "Gather", {1, 11, 13}, kOnnxDomain}, {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; // Path 4 pattern (Path 3 with no "Cast"): std::vector parent_path_4{ {0, 1, "Expand", {8, 13}, kOnnxDomain}, {0, 0, "Unsqueeze", {1, 11, 13, 21, 23, 24, 25}, kOnnxDomain}, - {0, 0, "Range", {11, 27}, kOnnxDomain}, + {0, 0, "Range", {11, 27, 28}, kOnnxDomain}, {0, 1, "Gather", {1, 11, 13}, kOnnxDomain}, {0, 0, "Shape", {1, 13, 15, 19, 21, 23, 24, 25}, kOnnxDomain}}; // Match one of the three path patterns. diff --git a/onnxruntime/core/optimizer/fast_gelu_fusion.cc b/onnxruntime/core/optimizer/fast_gelu_fusion.cc index cda93c9e09e31..8b911a6e63914 100644 --- a/onnxruntime/core/optimizer/fast_gelu_fusion.cc +++ b/onnxruntime/core/optimizer/fast_gelu_fusion.cc @@ -177,7 +177,7 @@ MatchResult FastGeluFusion::CheckSecondFormula(Graph& graph, Node& pow1_node, if (p_cast1_node != nullptr) { Node& cast1_node = *graph.GetNode(p_cast1_node->Index()); // this is fused Cast node, so expect 2 output edges - if (!(graph_utils::IsSupportedOptypeVersionAndDomain(cast1_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) && + if (!(graph_utils::IsSupportedOptypeVersionAndDomain(cast1_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) && cast1_node.InputDefs().size() == 1 && CheckNode(graph, cast1_node, pow1_node.GetExecutionProviderType(), false)) || cast1_node.GetOutputEdgesCount() != 2) { @@ -293,7 +293,7 @@ Status FastGeluFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, if (p_cast3_node == nullptr) continue; Node& cast3_node = *graph.GetNode(p_cast3_node->Index()); - if (!(graph_utils::IsSupportedOptypeVersionAndDomain(cast3_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) && + if (!(graph_utils::IsSupportedOptypeVersionAndDomain(cast3_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) && cast3_node.InputDefs().size() == 1 && CheckNode(graph, cast3_node, node.GetExecutionProviderType(), true))) { continue; diff --git a/onnxruntime/core/optimizer/insert_cast_transformer.cc b/onnxruntime/core/optimizer/insert_cast_transformer.cc index 05f53273447b9..f2983b18d9def 100644 --- a/onnxruntime/core/optimizer/insert_cast_transformer.cc +++ b/onnxruntime/core/optimizer/insert_cast_transformer.cc @@ -25,9 +25,31 @@ bool InsertCastTransformer::NeedInsertCast(const onnxruntime::Node* node, const // We don't cast a node with a subgraph as we'd need to do a lot more checking of the subgraph inputs // (both explicit and implicit) and contents to determine if it was safe to do so. // TODO: a better check is to check does the CPU kernel with float exist or not. - return node->GetExecutionProviderType().empty() && - !node->ContainsSubgraph() && - IsMLFloat16Tensor(*input); + if (!node->GetExecutionProviderType().empty() || node->ContainsSubgraph() || !IsMLFloat16Tensor(*input)) { + return false; + } + + if (const auto* schema = node->Op()) { + const auto& formal_inputs = schema->inputs(); + const auto& input_counts = node->InputArgCount(); + const auto& inputs = node->InputDefs(); + const auto& constraints = schema->typeConstraintMap(); + const auto float_type = ONNX_NAMESPACE::Utils::DataTypeUtils::ToType("tensor(float)"); + size_t actual_index = 0; + for (size_t formal_index = 0; formal_index < formal_inputs.size() && formal_index < input_counts.size(); + ++formal_index) { + const auto constraint = constraints.find(formal_inputs[formal_index].GetTypeStr()); + if (constraint != constraints.end() && constraint->second.first.count(float_type) == 0) { + for (int count = 0; count < input_counts[formal_index]; ++count) { + if (inputs[actual_index + count] && IsMLFloat16Tensor(*inputs[actual_index + count])) { + return false; + } + } + } + actual_index += input_counts[formal_index]; + } + } + return true; } onnxruntime::NodeArg* AddCastNode(onnxruntime::Graph& graph, diff --git a/onnxruntime/core/optimizer/isinf_reducesum_fusion.cc b/onnxruntime/core/optimizer/isinf_reducesum_fusion.cc index 3f963e86acb0d..2185261629e24 100644 --- a/onnxruntime/core/optimizer/isinf_reducesum_fusion.cc +++ b/onnxruntime/core/optimizer/isinf_reducesum_fusion.cc @@ -45,7 +45,7 @@ Status IsInfReduceSumFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l // This Cast can be skipped as we are replacing the subgraph with IsAllFinite, which supports FP16 auto cast1_node_iter = isinf_node.InputNodesBegin(); if (cast1_node_iter != isinf_node.InputNodesEnd() && - graph_utils::IsSupportedOptypeVersionAndDomain(*cast1_node_iter, "Cast", {9, 13, 19, 21, 23, 24, 25}) && + graph_utils::IsSupportedOptypeVersionAndDomain(*cast1_node_iter, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) && cast1_node_iter->GetOutputEdgesCount() == 1) { // check input type of cast node Node& cast1_node = *graph.GetNode(cast1_node_iter->Index()); @@ -65,7 +65,7 @@ Status IsInfReduceSumFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l } Node& cast2_node = *graph.GetNode(cast2_node_itr->Index()); - if (!graph_utils::IsSupportedOptypeVersionAndDomain(cast2_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) || + if (!graph_utils::IsSupportedOptypeVersionAndDomain(cast2_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) || cast2_node.GetOutputEdgesCount() != 1 || graph.NodeProducesGraphOutput(cast2_node)) { continue; diff --git a/onnxruntime/core/optimizer/layer_norm_fusion.cc b/onnxruntime/core/optimizer/layer_norm_fusion.cc index 753b23724d7eb..13b0fcb99e004 100644 --- a/onnxruntime/core/optimizer/layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/layer_norm_fusion.cc @@ -15,8 +15,6 @@ namespace onnxruntime { // LayerNorm supports limited data types. static constexpr std::array supported_data_types{"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}; -// Default epsilon -static constexpr float DEFAULT_LAYERNORM_EPSILON = 1e-5f; static bool IsSupportedDataType(const Node& node, int first_n_inputs = -1) { int input_index = 0; @@ -79,7 +77,8 @@ static std::vector GetAxesFromReduceMeanNode(Node& reduce_mean_node, co return axes_values; }; -static bool TryGetScalarInitializerAsDouble(const Graph& graph, const NodeArg& node_arg, double& value) { +template +static bool TryGetScalarInitializer(const Graph& graph, const NodeArg& node_arg, T& value) { const auto* tensor_proto = graph_utils::GetConstantInitializer(graph, node_arg.Name()); if (tensor_proto == nullptr) { return false; @@ -92,46 +91,70 @@ static bool TryGetScalarInitializerAsDouble(const Graph& graph, const NodeArg& n switch (tensor_proto->data_type()) { case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE: - value = initializer.data()[0]; + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_INT8: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_INT16: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_INT32: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_INT64: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_UINT8: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_UINT16: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_UINT32: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; case ONNX_NAMESPACE::TensorProto_DataType_UINT64: - value = static_cast(initializer.data()[0]); + value = static_cast(initializer.data()[0]); return true; default: return false; } } +static bool TryGetScalarInitializerOrFloatCastAsDouble(const Graph& graph, const NodeArg& node_arg, double& value) { + if (TryGetScalarInitializer(graph, node_arg, value)) { + return true; + } + + const Node* cast_node = graph.GetProducerNode(node_arg.Name()); + if (cast_node == nullptr || + !graph_utils::IsSupportedOptypeVersionAndDomain(*cast_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) || + cast_node->InputDefs().empty() || cast_node->InputDefs()[0] == nullptr) { + return false; + } + + const auto to_it = cast_node->GetAttributes().find("to"); + float cast_value = 0.0f; + if (to_it == cast_node->GetAttributes().end() || + to_it->second.i() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT || + !TryGetScalarInitializer(graph, *cast_node->InputDefs()[0], cast_value)) { + return false; + } + + value = static_cast(cast_value); + return true; +} + static bool IsPowExponentTwo(const Graph& graph, const Node& pow_node) { const auto& pow_inputs = pow_node.InputDefs(); if (pow_inputs.size() < 2 || pow_inputs[1] == nullptr) { @@ -139,18 +162,18 @@ static bool IsPowExponentTwo(const Graph& graph, const Node& pow_node) { } double exponent_value = 0.0; - if (TryGetScalarInitializerAsDouble(graph, *pow_inputs[1], exponent_value)) { + if (TryGetScalarInitializer(graph, *pow_inputs[1], exponent_value)) { return exponent_value == 2.0; } const Node* exponent_input_node = graph_utils::GetInputNode(pow_node, 1); if (exponent_input_node == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*exponent_input_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) || + !graph_utils::IsSupportedOptypeVersionAndDomain(*exponent_input_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) || exponent_input_node->InputDefs().empty() || exponent_input_node->InputDefs()[0] == nullptr) { return false; } - return TryGetScalarInitializerAsDouble(graph, *exponent_input_node->InputDefs()[0], exponent_value) && + return TryGetScalarInitializer(graph, *exponent_input_node->InputDefs()[0], exponent_value) && exponent_value == 2.0; } @@ -334,7 +357,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, if (p_reduce_mean_input_node) { Node& reduce_mean_input_node = *graph.GetNode(p_reduce_mean_input_node->Index()); // If input to the 1st ReduceMean is a Cast, and the Cast has same consumer count as subCnt + 1 - if (graph_utils::IsSupportedOptypeVersionAndDomain(reduce_mean_input_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) && + if (graph_utils::IsSupportedOptypeVersionAndDomain(reduce_mean_input_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) && reduce_mean_input_node.GetExecutionProviderType() == reduce_mean_node.GetExecutionProviderType() && optimizer_utils::CheckOutputEdges(graph, reduce_mean_input_node, static_cast(subCnt) + 1)) { nodes_to_remove.insert(nodes_to_remove.begin(), reduce_mean_input_node); @@ -347,7 +370,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const Node* p_cast1 = nullptr; if (!p_sub_node_dup && sub_node.GetOutputEdgesCount() == 1) { Node& cast_node = *graph.GetNode(sub_node.OutputNodesBegin()->Index()); - if (graph_utils::IsSupportedOptypeVersionAndDomain(cast_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) && + if (graph_utils::IsSupportedOptypeVersionAndDomain(cast_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) && cast_node.GetExecutionProviderType() == reduce_mean_node.GetExecutionProviderType() && optimizer_utils::CheckOutputEdges(graph, cast_node, 2u) && IsSupportedDataType(cast_node)) { p_cast1 = &cast_node; @@ -446,7 +469,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const Node* p_cast2 = graph_utils::FirstParentByType(pow_node, "Cast"); if (p_cast2 != nullptr && p_cast2 != p_cast1) { Node& cast_node = *graph.GetNode(p_cast2->Index()); - if (!graph_utils::IsSupportedOptypeVersionAndDomain(cast_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) || + if (!graph_utils::IsSupportedOptypeVersionAndDomain(cast_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) || cast_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || !optimizer_utils::CheckOutputEdges(graph, cast_node, 1)) { continue; @@ -464,7 +487,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, // can be removed. This is one possible place a Cast Op can exist, that is between Div and Mul nodes. // div --> mul or div --> cast --> mul Node* next_node = graph.GetNode(div_node.OutputNodesBegin()->Index()); - if (graph_utils::IsSupportedOptypeVersionAndDomain(*next_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) && + if (graph_utils::IsSupportedOptypeVersionAndDomain(*next_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) && optimizer_utils::CheckOutputEdges(graph, *next_node, 1)) { nodes_to_remove.push_back(*next_node); next_node = graph.GetNode(next_node->OutputNodesBegin()->Index()); @@ -572,20 +595,21 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, continue; } + const NodeArg* epsilon_input = GetOtherAddInput(add2_node, *reduce_mean2_node.MutableOutputDefs()[0]); + double epsilon = 0.0; + if (epsilon_input == nullptr || + !TryGetScalarInitializerOrFloatCastAsDouble(graph, *epsilon_input, epsilon) || + static_cast(static_cast(epsilon)) != epsilon) { + continue; + } + InlinedVector layer_norm_input_defs{x_input, scale, bias}; Node& layer_norm_node = graph.AddNode(graph.GenerateNodeName(mul_node.Name() + "/LayerNormFusion/"), "LayerNormalization", "fused LayerNorm subgraphs ", layer_norm_input_defs, {}, mul_node, nullptr, kOnnxDomain); - - // Get constant "epsilon" from "Add2" node if available. Else, default value will be used. - double epsilon = 0.0; - if (TryGetScalarInitializerAsDouble(graph, *add2_node.MutableInputDefs()[1], epsilon)) { - layer_norm_node.AddAttribute("epsilon", static_cast(epsilon)); - } else { - layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON); - } + layer_norm_node.AddAttribute("epsilon", static_cast(epsilon)); // The axis definition of layer_norm is ranging from axis to the last dim layer_norm_node.AddAttribute("axis", static_cast(axes_values[0])); @@ -745,7 +769,7 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr if (is_gpu_ep && p_pow_input_node) { Node& pow_input_node = *graph.GetNode(p_pow_input_node->Index()); // If input to Pow is a Cast, and the Cast has 2 consumers only (Pow, Div) - if (graph_utils::IsSupportedOptypeVersionAndDomain(pow_input_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) && + if (graph_utils::IsSupportedOptypeVersionAndDomain(pow_input_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) && pow_input_node.GetExecutionProviderType() == pow_node.GetExecutionProviderType() && optimizer_utils::CheckOutputEdges(graph, pow_input_node, 2)) { nodes_to_remove.insert(nodes_to_remove.begin(), pow_input_node); @@ -755,7 +779,7 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr // div --> mul or div --> cast --> mul Node* next_node = graph.GetNode(div_node.OutputNodesBegin()->Index()); - if (graph_utils::IsSupportedOptypeVersionAndDomain(*next_node, "Cast", {9, 13, 19, 21, 23, 24, 25}) && + if (graph_utils::IsSupportedOptypeVersionAndDomain(*next_node, "Cast", {9, 13, 19, 21, 23, 24, 25, 28}) && optimizer_utils::CheckOutputEdges(graph, *next_node, 1)) { if (!is_gpu_ep) continue; nodes_to_remove.push_back(*next_node); @@ -824,25 +848,17 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr continue; } + double epsilon = 0.0; + if (!TryGetScalarInitializerOrFloatCastAsDouble(graph, *epsilon_input, epsilon) || + static_cast(static_cast(epsilon)) != epsilon) { + continue; + } + InlinedVector layer_norm_input_defs{x_input, scale}; Node& layer_norm_node = graph.AddNode(graph.GenerateNodeName(mul_node.Name() + "/SimplifiedLayerNormFusion/"), "SimplifiedLayerNormalization", "fused LayerNorm subgraphs ", layer_norm_input_defs, {}, mul_node, nullptr, kOnnxDomain); - - // Get constant "epsilon" from "Add" node if available. Else, default value will be used. - const ONNX_NAMESPACE::TensorProto* tensor_proto = - graph_utils::GetConstantInitializer(graph, epsilon_input->Name()); - if (tensor_proto != nullptr && tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { - Initializer initializer{graph, *tensor_proto, graph.ModelPath()}; - // epsilon must be a scalar/1-element tensor; fall back to default otherwise. - if (initializer.size() == 1) { - layer_norm_node.AddAttribute("epsilon", initializer.data()[0]); - } else { - layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON); - } - } else { - layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON); - } + layer_norm_node.AddAttribute("epsilon", static_cast(epsilon)); // Set stash_type to double if any input is double, default value if float. if (x_input->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE || diff --git a/onnxruntime/core/optimizer/matmul_bn_fusion.cc b/onnxruntime/core/optimizer/matmul_bn_fusion.cc index be52e26a2901f..d3a2d35033de7 100644 --- a/onnxruntime/core/optimizer/matmul_bn_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_bn_fusion.cc @@ -2,10 +2,13 @@ // Licensed under the MIT License. #include "core/optimizer/matmul_bn_fusion.h" +#include "core/framework/tensorprotoutils.h" #include "core/graph/graph_utils.h" #include "core/optimizer/initializer.h" #include "core/optimizer/utils.h" +#include + namespace onnxruntime { namespace { @@ -13,6 +16,70 @@ const std::vector> dest = {"BatchNormalization", {1, 6, 7, 9, 14, 15}}; + +bool DimsAfterAxisAreOne(const ONNX_NAMESPACE::TensorShapeProto& shape, int axis) { + for (int i = axis + 1; i < shape.dim_size(); ++i) { + if (!utils::HasDimValue(shape.dim(i)) || shape.dim(i).dim_value() != 1) { + return false; + } + } + return true; +} + +// Checks that column j of the MatMul output ends up in channel j (axis 1) of the BatchNormalization +// input. Otherwise the in-between nodes move the columns to another axis or mix them with other +// dimensions, e.g. MatMul [2, 4] -> Reshape [1, 4, 2, 1], and the fusion would scale the wrong values. +bool MatMulColumnsAreBNChannels(const Graph& graph, const Node& matmul_node, const Node& batch_norm_node) { + const ONNX_NAMESPACE::TensorShapeProto* shape = matmul_node.OutputDefs()[0]->Shape(); + if (shape == nullptr || shape->dim_size() != 2 || !utils::HasDimValue(shape->dim(1))) { + return false; + } + const int64_t num_columns = shape->dim(1).dim_value(); + int column_axis = 1; + + const Node* node = graph.GetNode(matmul_node.OutputNodesBegin()->Index()); + while (node->Index() != batch_norm_node.Index()) { + const ONNX_NAMESPACE::TensorShapeProto* output_shape = node->OutputDefs()[0]->Shape(); + if (output_shape == nullptr) { + return false; + } + + if (node->OpType() == "Transpose") { + const auto* perm_attr = graph_utils::GetNodeAttribute(*node, "perm"); + if (perm_attr == nullptr) { + // The default perm reverses the dimensions. + column_axis = shape->dim_size() - 1 - column_axis; + } else { + const auto& perm = perm_attr->ints(); + const auto it = std::find(perm.begin(), perm.end(), column_axis); + if (it == perm.end()) { + return false; + } + column_axis = static_cast(it - perm.begin()); + } + } else { + // Reshape keeps the row-major order, so the columns stay on one axis only if every axis + // after it has size 1, both before and after the Reshape. + if (!DimsAfterAxisAreOne(*shape, column_axis)) { + return false; + } + column_axis = output_shape->dim_size() - 1; + while (num_columns != 1 && column_axis > 0 && utils::HasDimValue(output_shape->dim(column_axis)) && + output_shape->dim(column_axis).dim_value() == 1) { + --column_axis; + } + if (column_axis < 0 || !utils::HasDimValue(output_shape->dim(column_axis)) || + output_shape->dim(column_axis).dim_value() != num_columns) { + return false; + } + } + + shape = output_shape; + node = graph.GetNode(node->OutputNodesBegin()->Index()); + } + + return column_axis == 1; +} } // namespace bool NodeIsIgnorable(const Graph& graph, const Node& root_node, NodeIndex curr_node_index) { @@ -78,6 +145,7 @@ std::optional MatchPath(const Graph& graph, const Node& root_node, No * - B tensor of MatMul should be constant. * - scale, B, mean, var tensors of BatchNormalization should be constant. * - Every node in the path, except the BatchNormalization, should have only 1 output edge. + * - The in-between nodes should keep the MatMul columns as the BatchNormalization channels (axis 1). */ bool MatmulBNFusion::SatisfyCondition(const Graph& graph, const Node& node, const logging::Logger&) const { if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "MatMul", {1, 9, 13}) || @@ -123,6 +191,10 @@ bool MatmulBNFusion::SatisfyCondition(const Graph& graph, const Node& node, cons return false; } + if (!MatMulColumnsAreBNChannels(graph, node, *batch_norm_node)) { + return false; + } + // First output from BN is required. Others are optional. If any optional outputs exist we can't fuse. const auto& output_defs = batch_norm_node->OutputDefs(); if (output_defs.size() > 1) { diff --git a/onnxruntime/core/optimizer/matmul_scale_fusion.cc b/onnxruntime/core/optimizer/matmul_scale_fusion.cc index 151cd19df63c8..71ddb2dd28230 100644 --- a/onnxruntime/core/optimizer/matmul_scale_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_scale_fusion.cc @@ -3,6 +3,9 @@ #include "core/optimizer/matmul_scale_fusion.h" +#include +#include + #include "onnx/defs/attr_proto_util.h" #include "core/common/optional.h" @@ -16,12 +19,54 @@ namespace onnxruntime { namespace { template -struct ExtractScalarAsFloatDispatchTarget { +bool IsIntegralExactlyRepresentableAsFloat(T value) { + static_assert(std::is_integral_v); + using U = std::make_unsigned_t; + const U unsigned_value = static_cast(value); + const U magnitude = [&]() { + if constexpr (std::is_signed_v) { + return value < 0 ? U{0} - unsigned_value : unsigned_value; + } else { + return unsigned_value; + } + }(); + + size_t significant_bits = 0; + for (U remaining = magnitude; remaining != 0; remaining >>= 1) { + ++significant_bits; + } + + constexpr size_t float_digits = std::numeric_limits::digits; + if (significant_bits <= float_digits) { + return true; + } + + const size_t discarded_bits = significant_bits - float_digits; + return (magnitude & ((U{1} << discarded_bits) - 1)) == 0; +} + +struct ScalarScale { + double value; + bool is_double; +}; + +template +struct ExtractScalarDispatchTarget { Status operator()(const ONNX_NAMESPACE::TensorProto& tensor_proto, const std::filesystem::path& model_path, - float& scalar_float) { + ScalarScale& scale, bool& lossless) { T scalar; ORT_RETURN_IF_ERROR(utils::UnpackTensor(tensor_proto, model_path, &scalar, 1)); - scalar_float = static_cast(scalar); + const float scalar_float = static_cast(scalar); + if constexpr (std::is_same_v) { + lossless = static_cast(scalar_float) == scalar; + } else if constexpr (std::is_integral_v) { + lossless = IsIntegralExactlyRepresentableAsFloat(scalar); + } else { + lossless = true; + } + scale = {std::is_same_v ? static_cast(scalar) + : static_cast(scalar_float), + std::is_same_v}; return Status::OK(); } }; @@ -41,7 +86,7 @@ std::optional GetTensorShape(const NodeArg& node_arg) { // if it does not have any broadcasting effect on the Mul or Div output shape. // Because the dimension lengths can only be 1, we only need to consider additional leading dimensions being prepended. // `max_rank` should be set to the rank of the other Mul or Div input to avoid that. -std::optional GetScalarConstantInitializer(const Graph& graph, const NodeArg& node_arg, size_t max_rank) { +std::optional GetScalarConstantInitializer(const Graph& graph, const NodeArg& node_arg, size_t max_rank) { const auto* initializer = graph_utils::GetConstantInitializer(graph, node_arg.Name()); if (!initializer) { @@ -59,18 +104,20 @@ std::optional GetScalarConstantInitializer(const Graph& graph, const Node return {}; } - float scalar{}; + ScalarScale scale{}; + bool lossless = false; utils::MLTypeCallDispatcher< uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double, BFloat16> dispatcher{initializer->data_type()}; ORT_THROW_IF_ERROR( - (dispatcher.InvokeRet(*initializer, graph.ModelPath(), scalar))); + (dispatcher.InvokeRet( + *initializer, graph.ModelPath(), scale, lossless))); - return {scalar}; + return lossless ? std::optional{scale} : std::nullopt; } // gets the scale value and its input index if node is a fusable scale (Mul or Div by scalar constant) -std::optional> GetScaleFromNode( +std::optional> GetScaleFromNode( const Graph& graph, const Node& scale_node, const InlinedHashSet& excluded_initializer_names) { const auto is_excluded_initializer = @@ -95,7 +142,11 @@ std::optional> GetScaleFromNode( if (!divisor.has_value()) return std::nullopt; - return {std::make_pair(1.0f / *divisor, scale_reciprocal_arg_index)}; + ScalarScale scale = *divisor; + scale.value = scale.is_double + ? 1.0 / scale.value + : static_cast(1.0f / static_cast(scale.value)); + return {std::make_pair(scale, scale_reciprocal_arg_index)}; } if (graph_utils::IsSupportedOptypeVersionAndDomain(scale_node, "Mul", {7, 13, 14})) { @@ -128,7 +179,7 @@ struct ScaleMergeInfo { // the edge from the base node to the original node Node::EdgeConstIterator node_to_merge_edge; // the scale of the original node - float scale; + ScalarScale scale; // the index of the input or output def on the original node // this def is moved to the fused node // for a leading scale (scale -> MatMul), it will be the unscaled input @@ -236,15 +287,34 @@ Status ProcessNode( { ONNX_NAMESPACE::AttributeProto& alpha_attr = fused_node_attrs["alpha"]; - float total_scale = utils::HasFloat(alpha_attr) ? alpha_attr.f() : 1.0f; - auto accumulate_scale = [&total_scale](const ScaleMergeInfo& fusion) { - total_scale *= fusion.scale; - }; - - std::for_each(input_node_merges.begin(), input_node_merges.end(), accumulate_scale); - std::for_each(output_node_merges.begin(), output_node_merges.end(), accumulate_scale); - - alpha_attr = ONNX_NAMESPACE::MakeAttribute("alpha", total_scale); + const bool uses_double_scale = + std::any_of(input_node_merges.begin(), input_node_merges.end(), + [](const ScaleMergeInfo& fusion) { return fusion.scale.is_double; }) || + std::any_of(output_node_merges.begin(), output_node_merges.end(), + [](const ScaleMergeInfo& fusion) { return fusion.scale.is_double; }); + + if (uses_double_scale) { + double total_scale = utils::HasFloat(alpha_attr) ? static_cast(alpha_attr.f()) : 1.0; + auto accumulate_scale = [&total_scale](const ScaleMergeInfo& fusion) { + total_scale *= fusion.scale.value; + }; + std::for_each(input_node_merges.begin(), input_node_merges.end(), accumulate_scale); + std::for_each(output_node_merges.begin(), output_node_merges.end(), accumulate_scale); + + const float fused_scale = static_cast(total_scale); + if (static_cast(fused_scale) != total_scale) { + return Status::OK(); + } + alpha_attr = ONNX_NAMESPACE::MakeAttribute("alpha", fused_scale); + } else { + float total_scale = utils::HasFloat(alpha_attr) ? alpha_attr.f() : 1.0f; + auto accumulate_scale = [&total_scale](const ScaleMergeInfo& fusion) { + total_scale *= static_cast(fusion.scale.value); + }; + std::for_each(input_node_merges.begin(), input_node_merges.end(), accumulate_scale); + std::for_each(output_node_merges.begin(), output_node_merges.end(), accumulate_scale); + alpha_attr = ONNX_NAMESPACE::MakeAttribute("alpha", total_scale); + } } auto get_mutable_node_to_merge = [&graph](const ScaleMergeInfo& merge) -> Node& { diff --git a/onnxruntime/core/optimizer/nchwc_transformer.cc b/onnxruntime/core/optimizer/nchwc_transformer.cc index a00106a973231..99bfe801b35c1 100644 --- a/onnxruntime/core/optimizer/nchwc_transformer.cc +++ b/onnxruntime/core/optimizer/nchwc_transformer.cc @@ -359,6 +359,11 @@ void NchwcTransformerImpl::TransformConv(Node& node) { const int64_t output_channels = conv_W_tensor_proto->dims(0); const int64_t input_channels = conv_W_tensor_proto->dims(1); + // NCHWc Conv requires positive input and output channel counts. + if (output_channels <= 0 || input_channels <= 0) { + return; + } + int64_t group_count; const auto* group_attr = graph_utils::GetNodeAttribute(node, "group"); if (group_attr != nullptr && utils::HasInt(*group_attr)) { diff --git a/onnxruntime/core/optimizer/pad_fusion.cc b/onnxruntime/core/optimizer/pad_fusion.cc index 05310af8cdcad..6310904b8464a 100644 --- a/onnxruntime/core/optimizer/pad_fusion.cc +++ b/onnxruntime/core/optimizer/pad_fusion.cc @@ -195,7 +195,7 @@ bool PadFusion::SatisfyCondition(const Graph& graph, const Node& node, const log } const Node& child_node = *node.OutputNodesBegin(); - if (graph_utils::IsSupportedOptypeVersionAndDomain(child_node, "Cast", {1, 6, 9, 13, 19, 21, 23, 24, 25})) { + if (graph_utils::IsSupportedOptypeVersionAndDomain(child_node, "Cast", {1, 6, 9, 13, 19, 21, 23, 24, 25, 28})) { if (child_node.GetOutputEdgesCount() != 1) { return false; } diff --git a/onnxruntime/core/optimizer/slice_concat_to_space_to_depth_fusion.cc b/onnxruntime/core/optimizer/slice_concat_to_space_to_depth_fusion.cc index 8caea2c150990..424245806f61e 100644 --- a/onnxruntime/core/optimizer/slice_concat_to_space_to_depth_fusion.cc +++ b/onnxruntime/core/optimizer/slice_concat_to_space_to_depth_fusion.cc @@ -3,6 +3,7 @@ #include "core/optimizer/slice_concat_to_space_to_depth_fusion.h" +#include #include #include #include @@ -283,24 +284,11 @@ TypeProto MakeSpaceToDepthOutputTypeProto(const NodeArg& input) { return output_type; } -bool TryMatchSlicePhase(const Graph& graph, - const Node& slice, - const NodeArg& common_input, - const logging::Logger& logger, - NormalizedSliceParams& params, - SlicePhase& phase) { - if (slice.InputDefs().empty() || slice.InputDefs()[0] != &common_input) { - return false; - } - - IntValues starts; - IntValues ends; - IntValues axes; - IntValues steps; - if (!GetSliceInfo(graph, slice, logger, starts, ends, axes, steps)) { - return false; - } - +bool NormalizeSliceParams(const IntValues& starts, + const IntValues& ends, + const IntValues& axes, + const IntValues& steps, + NormalizedSliceParams& params) { params.starts = {0, 0, 0, 0}; params.ends = { std::numeric_limits::max(), @@ -316,15 +304,16 @@ bool TryMatchSlicePhase(const Graph& graph, return false; } - if (axis_seen[onnxruntime::narrow(axis)]) { + const size_t axis_index = onnxruntime::narrow(axis); + if (axis_seen[axis_index]) { return false; } - axis_seen[onnxruntime::narrow(axis)] = true; + axis_seen[axis_index] = true; - params.starts[onnxruntime::narrow(axis)] = starts[i]; - params.ends[onnxruntime::narrow(axis)] = ends[i]; - params.steps[onnxruntime::narrow(axis)] = steps[i]; + params.starts[axis_index] = starts[i]; + params.ends[axis_index] = ends[i]; + params.steps[axis_index] = steps[i]; } for (size_t axis = 0; axis < params.starts.size(); ++axis) { @@ -333,31 +322,132 @@ bool TryMatchSlicePhase(const Graph& graph, } } - if (params.starts[0] != 0 || params.starts[1] != 0 || - params.steps[0] != 1 || params.steps[1] != 1 || - params.steps[kHeightAxis] != kBlockSize || params.steps[kWidthAxis] != kBlockSize) { - return false; - } + return true; +} - if (!IsFullExtentEnd(common_input, 0, params.ends[0]) || - !IsFullExtentEnd(common_input, 1, params.ends[1]) || - !IsFullExtentEnd(common_input, kHeightAxis, params.ends[kHeightAxis]) || - !IsFullExtentEnd(common_input, kWidthAxis, params.ends[kWidthAxis])) { - return false; - } +// A focus branch is not always a single Slice that strides both spatial axes. +// Exporters commonly emit one Slice per axis, so a branch can be a chain of two +// Slice nodes. Two chained nodes is enough to cover height and width. +constexpr size_t kMaxSliceChainLength = 2; - const int64_t h_offset = params.starts[kHeightAxis]; - const int64_t w_offset = params.starts[kWidthAxis]; - if ((h_offset != 0 && h_offset != 1) || (w_offset != 0 && w_offset != 1)) { - return false; - } +struct SliceChain { + // Ordered from the Concat backwards, so index 0 feeds the Concat. + InlinedVector nodes; + const NodeArg* root_input; + SlicePhase phase; +}; - phase = {h_offset, w_offset}; - return true; -} +// Walks back from one Concat input through up to kMaxSliceChainLength Slice +// nodes, accumulating which spatial axis each one strides. Succeeds only when +// height and width are each strided exactly once and every other axis passes +// through untouched. +bool TryMatchSliceChain(Graph& graph, + Node& concat, + size_t input_index, + const std::string& provider_type, + const logging::Logger& logger, + SliceChain& chain) { + chain.nodes.clear(); + chain.root_input = nullptr; + + std::array axis_strided{}; + std::array axis_offset{}; + + Node* node = GetMutableInputProducerNode(graph, concat, input_index); + const Node* consumer = &concat; + + while (chain.nodes.size() < kMaxSliceChainLength) { + if (node == nullptr || node == &concat || + node->GetExecutionProviderType() != provider_type || + graph.NodeProducesGraphOutput(*node)) { + return false; + } + + // The node feeding the Concat must be ours alone. A node further upstream + // may legitimately be shared with a sibling branch, which is how a focus + // layer that splits height once and width twice is usually exported. + if (chain.nodes.empty()) { + if (!HasSingleOutputEdgeToNode(*node, *consumer)) { + return false; + } + } else if (node->GetOutputEdgesCount() == 0) { + return false; + } + + IntValues starts; + IntValues ends; + IntValues axes; + IntValues steps; + if (!GetSliceInfo(graph, *node, logger, starts, ends, axes, steps)) { + return false; + } + + NormalizedSliceParams params{}; + if (!NormalizeSliceParams(starts, ends, axes, steps, params)) { + return false; + } + + if (node->InputDefs().empty()) { + return false; + } + + const NodeArg* node_input = node->InputDefs()[0]; + if (node_input == nullptr || !node_input->Exists()) { + return false; + } + + // Batch and channel must pass through untouched. + for (const int64_t axis : {int64_t{0}, kChannelAxis}) { + const size_t axis_index = onnxruntime::narrow(axis); + if (params.starts[axis_index] != 0 || params.steps[axis_index] != 1 || + !IsFullExtentEnd(*node_input, axis, params.ends[axis_index])) { + return false; + } + } -bool IsSingleConsumerOfConcat(const Node& slice, const Node& concat) { - return HasSingleOutputEdgeToNode(slice, concat); + bool strided_here = false; + + for (const int64_t axis : {kHeightAxis, kWidthAxis}) { + const size_t axis_index = onnxruntime::narrow(axis); + + if (!IsFullExtentEnd(*node_input, axis, params.ends[axis_index])) { + return false; + } + + if (params.steps[axis_index] == kBlockSize) { + if (axis_strided[axis_index] || + (params.starts[axis_index] != 0 && params.starts[axis_index] != 1)) { + return false; + } + + axis_strided[axis_index] = true; + axis_offset[axis_index] = params.starts[axis_index]; + strided_here = true; + } else if (params.steps[axis_index] != 1 || params.starts[axis_index] != 0) { + return false; + } + } + + if (!strided_here) { + return false; + } + + chain.nodes.push_back(node); + + const size_t height_index = onnxruntime::narrow(kHeightAxis); + const size_t width_index = onnxruntime::narrow(kWidthAxis); + + if (axis_strided[height_index] && axis_strided[width_index]) { + chain.root_input = node_input; + chain.phase = {axis_offset[height_index], axis_offset[width_index]}; + return true; + } + + consumer = node; + node = GetMutableInputProducerNode(graph, *node, 0); + } + + return false; } bool TryGetPhasePermutation(const std::array& actual_phases, @@ -399,6 +489,11 @@ NodeArg* CreateInt64Initializer(Graph& graph, } bool FuseSliceConcatToSpaceToDepth(Node& concat, Graph& graph, const logging::Logger& logger) { + // Fusion runs after function expansion, and SpaceToDepth-28 has no native CPU kernel. + if (graph.DomainToVersionMap().at(kOnnxDomain) >= 28) { + return false; + } + if (!graph_utils::IsSupportedOptypeVersionAndDomain(concat, "Concat", {4, 11, 13}, kOnnxDomain) || concat.InputDefs().size() != 4) { return false; @@ -414,10 +509,9 @@ bool FuseSliceConcatToSpaceToDepth(Node& concat, Graph& graph, const logging::Lo return false; } - Node* slice_nodes[4]{}; const NodeArg* common_input = nullptr; const auto& provider_type = concat.GetExecutionProviderType(); - NormalizedSliceParams reference_params{}; + std::array chains{}; std::array actual_phases{}; for (size_t i = 0; i < concat.InputDefs().size(); ++i) { @@ -426,45 +520,25 @@ bool FuseSliceConcatToSpaceToDepth(Node& concat, Graph& graph, const logging::Lo return false; } - Node* slice = GetMutableInputProducerNode(graph, concat, i); - if (slice == nullptr || slice == &concat || slice->GetExecutionProviderType() != provider_type || - !IsSingleConsumerOfConcat(*slice, concat)) { + if (!TryMatchSliceChain(graph, concat, i, provider_type, logger, chains[i])) { return false; } if (i == 0) { - common_input = slice->InputDefs()[0]; + common_input = chains[i].root_input; if (common_input == nullptr || !IsSupportedSpaceToDepthInputType(*common_input)) { return false; } - } - - ORT_ENFORCE(common_input != nullptr); - - NormalizedSliceParams current_params{}; - SlicePhase phase{}; - if (!TryMatchSlicePhase(graph, *slice, *common_input, logger, current_params, phase)) { - return false; - } - - actual_phases[i] = phase; - - if (i == 0) { - reference_params = current_params; - } else if (current_params.ends != reference_params.ends || - current_params.steps != reference_params.steps || - current_params.starts[0] != reference_params.starts[0] || - current_params.starts[1] != reference_params.starts[1]) { - return false; - } - - if (graph.NodeProducesGraphOutput(*slice)) { + } else if (chains[i].root_input != common_input) { + // Every branch has to strip the same tensor. return false; } - slice_nodes[i] = slice; + actual_phases[i] = chains[i].phase; } + ORT_ENFORCE(common_input != nullptr); + std::array phase_permutation{}; if (!TryGetPhasePermutation(actual_phases, phase_permutation)) { return false; @@ -527,11 +601,13 @@ bool FuseSliceConcatToSpaceToDepth(Node& concat, Graph& graph, const logging::Lo replacement_end = &gather; } - // Explicitly transfer the shared data-input edge from the first Slice to - // SpaceToDepth. This avoids graph_utils::MoveAllNodeInputEdges(), which is - // not defined in extended minimal builds. + // Explicitly transfer the shared data-input edge from the topmost Slice of + // the first branch to SpaceToDepth. This avoids + // graph_utils::MoveAllNodeInputEdges(), which is not defined in extended + // minimal builds. { - const auto data_input_edges = graph_utils::GraphEdge::GetNodeInputEdges(*slice_nodes[0], 0); + Node& chain_root = *chains[0].nodes.back(); + const auto data_input_edges = graph_utils::GraphEdge::GetNodeInputEdges(chain_root, 0); if (!data_input_edges.empty()) { ORT_ENFORCE(data_input_edges.size() == 1, "Expected a single data input edge for Slice node."); const auto& data_input_edge = data_input_edges[0]; @@ -546,9 +622,33 @@ bool FuseSliceConcatToSpaceToDepth(Node& concat, Graph& graph, const logging::Lo graph.AddEdge(replacement_end->Index(), edge.dst_node, 0, edge.dst_arg_index); } - for (Node* node : {slice_nodes[0], slice_nodes[1], slice_nodes[2], slice_nodes[3], &concat}) { - graph_utils::RemoveNodeOutputEdges(graph, *node); - graph.RemoveNode(node->Index()); + // Drop the Concat and the Slice nodes that fed it directly, then any upstream + // Slice left without consumers. Upstream nodes can be shared between two + // branches, so they are deduplicated and only removed once orphaned. + InlinedVector upstream_nodes; + + graph_utils::RemoveNodeOutputEdges(graph, concat); + graph.RemoveNode(concat.Index()); + + for (const SliceChain& chain : chains) { + for (size_t i = 1; i < chain.nodes.size(); ++i) { + if (std::find(upstream_nodes.begin(), upstream_nodes.end(), chain.nodes[i]) == upstream_nodes.end()) { + upstream_nodes.push_back(chain.nodes[i]); + } + } + } + + for (const SliceChain& chain : chains) { + Node& node = *chain.nodes[0]; + graph_utils::RemoveNodeOutputEdges(graph, node); + graph.RemoveNode(node.Index()); + } + + for (Node* node : upstream_nodes) { + if (node->GetOutputEdgesCount() == 0 && !graph.NodeProducesGraphOutput(*node)) { + graph_utils::RemoveNodeOutputEdges(graph, *node); + graph.RemoveNode(node->Index()); + } } LOGS(logger, INFO) << "Fused Slice+Concat downsample pattern into " diff --git a/onnxruntime/core/optimizer/transpose_optimization/optimizer_api.h b/onnxruntime/core/optimizer/transpose_optimization/optimizer_api.h index cd0094877657d..ea3af9abd3bba 100644 --- a/onnxruntime/core/optimizer/transpose_optimization/optimizer_api.h +++ b/onnxruntime/core/optimizer/transpose_optimization/optimizer_api.h @@ -478,7 +478,7 @@ class GraphRef { } // namespace api constexpr int64_t kMinSupportedOpset = 7; -constexpr int64_t kMaxSupportedOpset = 27; +constexpr int64_t kMaxSupportedOpset = 28; // enum of results that a CostCheckFn can return. enum class CostCheckResult { diff --git a/onnxruntime/core/platform/telemetry_1ds.cc b/onnxruntime/core/platform/telemetry_1ds.cc index 3d9a9d04cbbd8..fa11a2f1bde68 100644 --- a/onnxruntime/core/platform/telemetry_1ds.cc +++ b/onnxruntime/core/platform/telemetry_1ds.cc @@ -342,6 +342,14 @@ OneDsTelemetry::OneDsTelemetry(const Telemetry& local_telemetry) : OneDsTelemetr } OneDsTelemetry::~OneDsTelemetry() { +#if defined(_WIN32) + // ExitProcess terminates the other threads before running DLL destructors. + // The uploader cannot drain WinHTTP callbacks here, and a terminated thread + // may still own one of our locks. Normal destruction/FreeLibrary still cleans up. + if (telemetry_internal::IsDllShutdownInProgress()) { + return; + } +#endif std::lock_guard lock(global_mutex_); global_register_count_--; diff --git a/onnxruntime/core/platform/telemetry_1ds.h b/onnxruntime/core/platform/telemetry_1ds.h index 2d0eabfd48dc2..665d38acf2ffe 100644 --- a/onnxruntime/core/platform/telemetry_1ds.h +++ b/onnxruntime/core/platform/telemetry_1ds.h @@ -140,6 +140,8 @@ class OneDsTelemetry : public Telemetry { const std::string& lib_path) const override; private: + friend class OneDsTelemetryTestPeer; + static void ReportFailure(const char* operation_name, const char* message); template diff --git a/onnxruntime/core/platform/telemetry_1ds_platform.h b/onnxruntime/core/platform/telemetry_1ds_platform.h index 2cfdbd99dd0a8..311b80b19c553 100644 --- a/onnxruntime/core/platform/telemetry_1ds_platform.h +++ b/onnxruntime/core/platform/telemetry_1ds_platform.h @@ -13,4 +13,8 @@ const char* GetDefaultEncodedToken(); int32_t GetProcessorCount(); std::string GetCachePath(const std::string& directory); +#if defined(_WIN32) +bool IsDllShutdownInProgress(); +#endif + } // namespace onnxruntime::telemetry_internal diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index bc3aa37ceeb16..ec4d81aa7f5bb 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -309,14 +309,18 @@ common::Status D3D12FileBufferLoader::LoadInternal( std::future allocation_future = std::async( std::launch::async, [this, sizes = std::move(sizes)]() noexcept { AllocationResult allocation; - try { + ORT_TRY { allocation.status = AllocateDestinationGpuBuffers(sizes, allocation.batch); - } catch (const std::exception& ex) { - allocation.status = ORT_MAKE_STATUS( - ONNXRUNTIME, RUNTIME_EXCEPTION, - "Destination allocation failed: ", ex.what()); - } catch (...) { + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + allocation.status = ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "Destination allocation failed: ", ex.what()); + }); + } + ORT_CATCH(...) { allocation.status = ORT_MAKE_STATUS( ONNXRUNTIME, RUNTIME_EXCEPTION, "Destination allocation failed with an unknown exception."); @@ -327,6 +331,7 @@ common::Status D3D12FileBufferLoader::LoadInternal( // CreateFileReadPlan groups buffer sources by file and coalesces overlapping // or adjacent sector-aligned ranges into read regions to avoid redundant reads. std::vector files; + BufferCollection loaded_batch; uint64_t last_submitted_fence = 0; bool load_succeeded = false; // Resources referenced by in-flight I/O or GPU work must remain alive on @@ -337,10 +342,7 @@ common::Status D3D12FileBufferLoader::LoadInternal( WaitForFenceUncancelled(last_submitted_fence); } if (allocation_future.valid()) { - try { - (void)allocation_future.get(); - } catch (...) { - } + allocation_future.wait(); } }); @@ -352,7 +354,6 @@ common::Status D3D12FileBufferLoader::LoadInternal( } ORT_RETURN_IF_ERROR(PrepareUploadStagingSlots(maximum_alignment, cancellation)); - BufferCollection loaded_batch; bool allocation_ready = false; auto ensure_allocation = [&]() -> common::Status { if (allocation_ready) { @@ -1111,47 +1112,71 @@ void D3D12FileBufferLoader::DrainActiveReads() noexcept { } } +common::Status D3D12FileBufferLoader::Create( + ID3D12Device* device, + std::unique_ptr& loader) noexcept { + return Create(device, loader, Config{}); +} + common::Status D3D12FileBufferLoader::Create( ID3D12Device* device, std::unique_ptr& loader, const Config& config) noexcept { loader.reset(); - try { + common::Status status; + ORT_TRY { auto instance = std::unique_ptr( new D3D12FileBufferLoader(device, config)); ORT_RETURN_IF_ERROR(instance->Initialize()); loader = std::move(instance); return common::Status::OK(); - } catch (const std::exception& ex) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, RUNTIME_EXCEPTION, - "D3D12FileBufferLoader::Create failed: ", - ex.what()); - } catch (...) { - return ORT_MAKE_STATUS( + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + status = ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "D3D12FileBufferLoader::Create failed: ", + ex.what()); + }); + } + ORT_CATCH(...) { + status = ORT_MAKE_STATUS( ONNXRUNTIME, RUNTIME_EXCEPTION, "D3D12FileBufferLoader::Create failed with an unknown exception."); } + return status; +} + +common::Status D3D12FileBufferLoader::Load( + const std::vector& ranges, + BufferCollection& result) noexcept { + return Load(ranges, result, CancellationToken{}); } common::Status D3D12FileBufferLoader::Load( const std::vector& ranges, BufferCollection& result, const CancellationToken& cancellation) noexcept { - try { + common::Status status; + ORT_TRY { return LoadInternal(ranges, result, cancellation); - } catch (const std::exception& ex) { + } + ORT_CATCH(const std::exception& ex) { result.Clear(); - return ORT_MAKE_STATUS( - ONNXRUNTIME, RUNTIME_EXCEPTION, - "D3D12FileBufferLoader::Load failed: ", - ex.what()); - } catch (...) { + ORT_HANDLE_EXCEPTION([&]() { + status = ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "D3D12FileBufferLoader::Load failed: ", + ex.what()); + }); + } + ORT_CATCH(...) { result.Clear(); - return ORT_MAKE_STATUS( + status = ORT_MAKE_STATUS( ONNXRUNTIME, RUNTIME_EXCEPTION, "D3D12FileBufferLoader::Load failed with an unknown exception."); } + return status; } } // namespace d3d12 diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h index 961a5ca0c2bf8..f1d782e46a536 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h @@ -55,10 +55,13 @@ class D3D12FileBufferLoader { } }; + static common::Status Create( + ID3D12Device* device, + std::unique_ptr& loader) noexcept; static common::Status Create( ID3D12Device* device, std::unique_ptr& loader, - const Config& config = {}) noexcept; + const Config& config) noexcept; ~D3D12FileBufferLoader(); @@ -69,10 +72,13 @@ class D3D12FileBufferLoader { // The returned buffers correspond one-to-one with ranges and preserve input order. // Each buffer is in D3D12_RESOURCE_STATE_COMMON when this method succeeds. // Resource widths may include padding; ranges specify the loaded data lengths. + common::Status Load( + const std::vector& ranges, + BufferCollection& result) noexcept; common::Status Load( const std::vector& ranges, BufferCollection& result, - const CancellationToken& cancellation = {}) noexcept; + const CancellationToken& cancellation) noexcept; private: struct FileReadRegion; diff --git a/onnxruntime/core/platform/windows/telemetry_1ds.cc b/onnxruntime/core/platform/windows/telemetry_1ds.cc index 88ad276a796a0..5fb8bd45bacb5 100644 --- a/onnxruntime/core/platform/windows/telemetry_1ds.cc +++ b/onnxruntime/core/platform/windows/telemetry_1ds.cc @@ -54,6 +54,15 @@ bool PrepareSampledProcessEvent(EventProperties& event) { namespace telemetry_internal { +bool IsDllShutdownInProgress() { + using RtlDllShutdownInProgressFn = BOOLEAN(NTAPI*)(); + const HMODULE ntdll = ::GetModuleHandleW(L"ntdll.dll"); + const auto is_shutting_down = + ntdll == nullptr ? nullptr + : reinterpret_cast(::GetProcAddress(ntdll, "RtlDllShutdownInProgress")); + return is_shutting_down != nullptr && is_shutting_down(); +} + const char* GetDefaultEncodedToken() { return "fllXSmJHWBYMX1cqWldIYU0KTA0IVnpWD0lkRloSWwhILF5WGmBDDUdEC1csC0NMZhZYWVEPBH9DWEhlR15MCl1UfgwPVWVDX0U="; } diff --git a/onnxruntime/core/providers/cann/cann_graph.cc b/onnxruntime/core/providers/cann/cann_graph.cc index fc2c6fcbf361b..d431c89d4d219 100644 --- a/onnxruntime/core/providers/cann/cann_graph.cc +++ b/onnxruntime/core/providers/cann/cann_graph.cc @@ -70,7 +70,8 @@ std::vector SupportONNXModel(const GraphViewer& graph_viewer) { const auto& node = graph_viewer.GetNode(index); if (node->Domain() != kOnnxDomain || domain_version < lower_bound || - !cann_supported_ops.count(node->OpType())) { + !cann_supported_ops.count(node->OpType()) || + ((node->OpType() == "Mod" || node->OpType() == "SpaceToDepth") && node->SinceVersion() >= 28)) { unsupported_nodes.push_back(index); continue; } diff --git a/onnxruntime/core/providers/coreml/builders/impl/base_op_builder.h b/onnxruntime/core/providers/coreml/builders/impl/base_op_builder.h index 609a965fa5198..0e49f488f8684 100644 --- a/onnxruntime/core/providers/coreml/builders/impl/base_op_builder.h +++ b/onnxruntime/core/providers/coreml/builders/impl/base_op_builder.h @@ -46,7 +46,7 @@ class BaseOpBuilder : public IOpBuilder { const logging::Logger& logger) const; virtual int GetMinSupportedOpSet(const Node& /*node*/) const { return 1; } - virtual int GetMaxSupportedOpSet(const Node& /*node*/) const { return 27; } + virtual int GetMaxSupportedOpSet(const Node& /*node*/) const { return 28; } bool HasSupportedOpSet(const Node& node, const logging::Logger& logger) const; bool HasSupportedInputs(const Node& node, const OpBuilderInputParams& input_params, diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index c3e062611d5cf..701424a532d07 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -599,9 +599,10 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOn class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, MLFloat16, Gemm); #endif class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, GatherElements); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint8_t, BitShift); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint32_t, BitShift); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint64_t, BitShift); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27, uint8_t, BitShift); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27, uint16_t, BitShift); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27, uint32_t, BitShift); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27, uint64_t, BitShift); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, Pad); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 11, GatherND); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 26, Range); @@ -837,7 +838,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int16_t, Neg); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int32_t, Neg); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int64_t, Neg); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Mod); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 27, Mod); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, Mod); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, float, Abs); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double, Abs); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int8_t, Abs); @@ -871,7 +873,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double, Log); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 14, Pow); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, DepthToSpace); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, SpaceToDepth); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 27, SpaceToDepth); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Slice); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 17, Split); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 20, Unsqueeze); @@ -1499,7 +1501,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 24, uint8_t, TopK); // Opset 25 -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 25, Cast); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 25, 27, Cast); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 25, ConstantOfShape); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 25, int32_t, DequantizeLinear); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 25, uint8_t, DequantizeLinear); @@ -1555,6 +1557,17 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, // Opset 27 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 27, Range); +// Opset 28 +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int8_t, BitShift); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int16_t, BitShift); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int32_t, BitShift); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int64_t, BitShift); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint8_t, BitShift); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint16_t, BitShift); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint32_t, BitShift); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint64_t, BitShift); + // !!PLEASE READ BELOW!! Following that, add new entries above this comment /* *** IMPORTANT! *** @@ -2321,12 +2334,14 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { double, Gemm)>, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -2770,7 +2785,8 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -2814,7 +2830,7 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, // opset 25 - BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -3758,6 +3774,25 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { // opset 27 BuildKernelCreateInfo, + + // opset 28 + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, }; static_assert(registration_internal::IsKernelRegistrationTableValid( function_table, BuildKernelCreateInfo)); diff --git a/onnxruntime/core/providers/cpu/math/element_wise_ops.cc b/onnxruntime/core/providers/cpu/math/element_wise_ops.cc index 4a9c4646a2c34..9f20c2ab02387 100644 --- a/onnxruntime/core/providers/cpu/math/element_wise_ops.cc +++ b/onnxruntime/core/providers/cpu/math/element_wise_ops.cc @@ -12,6 +12,9 @@ #include "core/mlas/inc/mlas.h" #include +#include +#include +#include namespace onnxruntime { // Supported types for operators that have type reduction enabled @@ -489,10 +492,18 @@ REG_ELEMENTWISE_VERSIONED_TYPED_KERNEL(Mean, 8, 12, float, Mean_8); // Supposed to add BFloat16 but we are not supporting now, however, separate registration REG_ELEMENTWISE_TYPED_KERNEL(Mean, 13, float, Mean_8); -REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 11, uint8_t, BitShift); -// REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 11, uint16_t, BitShift); -REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 11, uint32_t, BitShift); -REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 11, uint64_t, BitShift); +REG_ELEMENTWISE_VERSIONED_TYPED_KERNEL(BitShift, 11, 27, uint8_t, BitShift); +REG_ELEMENTWISE_VERSIONED_TYPED_KERNEL(BitShift, 11, 27, uint16_t, BitShift); +REG_ELEMENTWISE_VERSIONED_TYPED_KERNEL(BitShift, 11, 27, uint32_t, BitShift); +REG_ELEMENTWISE_VERSIONED_TYPED_KERNEL(BitShift, 11, 27, uint64_t, BitShift); +REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, int8_t, BitShift); +REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, int16_t, BitShift); +REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, int32_t, BitShift); +REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, int64_t, BitShift); +REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, uint8_t, BitShift); +REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, uint16_t, BitShift); +REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, uint32_t, BitShift); +REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, uint64_t, BitShift); REG_ELEMENTWISE_TYPED_KERNEL(BitwiseAnd, 18, int8_t, BitwiseAnd); REG_ELEMENTWISE_TYPED_KERNEL(BitwiseAnd, 18, int16_t, BitwiseAnd); @@ -1334,17 +1345,54 @@ BitShift::BitShift(const OpKernelInfo& info) : OpKernel(info) { ORT_THROW("Invalid direction value of '", direction, "'. Valid values are 'LEFT' or 'RIGHT'."); } -// Shifting by >= the bit width of an unsigned type is undefined behavior in C++. -// On x86, 64-bit shifts mask the shift amount to 6 bits, so shift by 64 acts like shift by 0. -// Guard against this by returning 0 when the shift amount >= the bit width. +template +T FromUnsignedBits(std::make_unsigned_t bits) { + static_assert(std::is_integral_v); + T value; + std::memcpy(&value, &bits, sizeof(value)); + return value; +} + +template +inline bool IsInvalidShift(T shift) { + using UnsignedT = std::make_unsigned_t; + constexpr UnsignedT bit_width = std::numeric_limits::digits; + return (std::is_signed_v && shift < 0) || static_cast(shift) >= bit_width; +} + template inline T SafeShiftLeft(T value, T shift) { - return shift >= sizeof(T) * 8 ? T{0} : value << shift; + if (IsInvalidShift(shift)) { + return T{0}; + } + + using UnsignedT = std::make_unsigned_t; + const auto result = static_cast( + static_cast(value) << static_cast(shift)); + return FromUnsignedBits(result); } template inline T SafeShiftRight(T value, T shift) { - return shift >= sizeof(T) * 8 ? T{0} : value >> shift; + if (IsInvalidShift(shift)) { + if constexpr (std::is_signed_v) { + return value < 0 ? T{-1} : T{0}; + } else { + return T{0}; + } + } + + using UnsignedT = std::make_unsigned_t; + constexpr UnsignedT bit_width = std::numeric_limits::digits; + const auto shift_amount = static_cast(shift); + auto result = static_cast(static_cast(value) >> shift_amount); + if constexpr (std::is_signed_v) { + if (value < 0 && shift_amount != 0) { + result |= std::numeric_limits::max() << (bit_width - shift_amount); + } + } + + return FromUnsignedBits(result); } template @@ -2097,6 +2145,7 @@ class Mod final : public OpKernel { private: bool fmod_{false}; + bool supports_float_floor_mod_{false}; bool divisor_is_validated_constant_{false}; }; @@ -2110,9 +2159,19 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( BuildKernelDefConstraintsFromTypeList()), Mod); -ONNX_CPU_OPERATOR_KERNEL( +ONNX_CPU_OPERATOR_VERSIONED_KERNEL( Mod, 13, + 27, + KernelDefBuilder() + .TypeConstraint( + "T", + BuildKernelDefConstraintsFromTypeList()), + Mod); + +ONNX_CPU_OPERATOR_KERNEL( + Mod, + 28, KernelDefBuilder() .TypeConstraint( "T", @@ -2121,6 +2180,26 @@ ONNX_CPU_OPERATOR_KERNEL( namespace mod_internal { +template +inline T TruncatingRemainder(T x, T y) { + if constexpr (std::is_signed_v) { + if (x == std::numeric_limits::min() && y == T{-1}) { + return T{0}; + } + } + + return x % y; +} + +template +inline T FMod(T x, T y) { + if constexpr (std::is_integral_v) { + return TruncatingRemainder(x, y); + } else { + return static_cast(std::fmod(x, y)); + } +} + template void BroadCastFMod(OpKernelContext* context) { ProcessBroadcastSpanFuncs funcs{ @@ -2131,7 +2210,7 @@ void BroadCastFMod(OpKernelContext* context) { std::transform(Y.begin(), Y.end(), output.begin(), [X](T y) { - return static_cast(std::fmod(X, y)); + return FMod(X, y); }); }, [](BroadcastHelper& per_iter_bh) { @@ -2141,7 +2220,7 @@ void BroadCastFMod(OpKernelContext* context) { std::transform(X.begin(), X.end(), output.begin(), [Y](T x) { - return static_cast(std::fmod(x, Y)); + return FMod(x, Y); }); }, [](BroadcastHelper& per_iter_bh) { @@ -2151,7 +2230,7 @@ void BroadCastFMod(OpKernelContext* context) { std::transform(X.begin(), X.end(), Y.begin(), output.begin(), [](T x, T y) { - return static_cast(std::fmod(x, y)); + return FMod(x, y); }); }}; @@ -2160,11 +2239,11 @@ void BroadCastFMod(OpKernelContext* context) { template inline T Modulus(T x, T y) { - auto res = x % y; + auto res = TruncatingRemainder(x, y); if ((res < 0 && y > 0) || (res > 0 && y < 0)) { res += y; } - return static_cast(res); + return res; } template @@ -2204,6 +2283,56 @@ void BroadCastMod(OpKernelContext* context) { UntypedBroadcastTwo(*context, funcs); } +template +inline T FloorMod(T x, T y) { + auto res = std::fmod(x, y); + if (res == T{0}) { + return std::copysign(T{0}, y); + } + + if ((res < T{0} && y > T{0}) || (res > T{0} && y < T{0})) { + res += y; + } + return res; +} + +template +void BroadCastFloorMod(OpKernelContext* context) { + ProcessBroadcastSpanFuncs funcs{ + [](BroadcastHelper& per_iter_bh) { + const T& X = per_iter_bh.ScalarInput0(); + auto Y = per_iter_bh.SpanInput1(); + auto output = per_iter_bh.OutputSpan(); + + std::transform(Y.begin(), Y.end(), output.begin(), + [X](T y) { + return FloorMod(X, y); + }); + }, + [](BroadcastHelper& per_iter_bh) { + auto X = per_iter_bh.SpanInput0(); + const T& Y = per_iter_bh.ScalarInput1(); + auto output = per_iter_bh.OutputSpan(); + + std::transform(X.begin(), X.end(), output.begin(), + [Y](T x) { + return FloorMod(x, Y); + }); + }, + [](BroadcastHelper& per_iter_bh) { + auto X = per_iter_bh.SpanInput0(); + auto Y = per_iter_bh.SpanInput1(); + auto output = per_iter_bh.OutputSpan(); + + std::transform(X.begin(), X.end(), Y.begin(), output.begin(), + [](T x, T y) { + return FloorMod(x, y); + }); + }}; + + UntypedBroadcastTwo(*context, funcs); +} + void BroadCastMLFloat16FMod(OpKernelContext* context) { ProcessBroadcastSpanFuncs funcs{ [](BroadcastHelper& per_iter_bh) { @@ -2242,6 +2371,42 @@ void BroadCastMLFloat16FMod(OpKernelContext* context) { UntypedBroadcastTwo(*context, funcs); } +void BroadCastMLFloat16FloorMod(OpKernelContext* context) { + ProcessBroadcastSpanFuncs funcs{ + [](BroadcastHelper& per_iter_bh) { + const auto X = per_iter_bh.ScalarInput0(); + auto Y = per_iter_bh.SpanInput1(); + auto output = per_iter_bh.OutputSpan(); + + std::transform(Y.begin(), Y.end(), output.begin(), + [X_fl = X.ToFloat()](const MLFloat16& y) { + return MLFloat16(FloorMod(X_fl, y.ToFloat())); + }); + }, + [](BroadcastHelper& per_iter_bh) { + auto X = per_iter_bh.SpanInput0(); + const MLFloat16 Y = per_iter_bh.ScalarInput1(); + auto output = per_iter_bh.OutputSpan(); + + std::transform(X.begin(), X.end(), output.begin(), + [Y_fl = Y.ToFloat()](const MLFloat16& x) { + return MLFloat16(FloorMod(x.ToFloat(), Y_fl)); + }); + }, + [](BroadcastHelper& per_iter_bh) { + auto X = per_iter_bh.SpanInput0(); + auto Y = per_iter_bh.SpanInput1(); + auto output = per_iter_bh.OutputSpan(); + + std::transform(X.begin(), X.end(), Y.begin(), output.begin(), + [](const MLFloat16& x, const MLFloat16& y) { + return MLFloat16(FloorMod(x.ToFloat(), y.ToFloat())); + }); + }}; + + UntypedBroadcastTwo(*context, funcs); +} + template struct CallModImpl; @@ -2263,7 +2428,7 @@ struct CheckZeroDivisorImpl { // Generic implementation of Mod kernel, non-floating point types template struct CallModImpl::value>::type> { - void operator()(bool fmod, OpKernelContext* ctx) const { + void operator()(bool fmod, bool /*supports_float_floor_mod*/, OpKernelContext* ctx) const { if (fmod) { BroadCastFMod(ctx); } else { @@ -2275,24 +2440,35 @@ struct CallModImpl::value> // Generic implementation of Mod kernel, floating point types template struct CallModImpl::value, void>::type> { - void operator()(bool fmod, OpKernelContext* ctx) const { - ORT_ENFORCE(fmod, "fmod attribute must be true for floating point types"); - BroadCastFMod(ctx); + void operator()(bool fmod, bool supports_float_floor_mod, OpKernelContext* ctx) const { + ORT_ENFORCE(fmod || supports_float_floor_mod, + "fmod attribute must be true for floating point types before opset 28"); + if (fmod) { + BroadCastFMod(ctx); + } else { + BroadCastFloorMod(ctx); + } } }; // MLFloat16 implementation of Mod kernel template <> struct CallModImpl { - void operator()(bool fmod, OpKernelContext* ctx) const { - ORT_ENFORCE(fmod, "fmod attribute must be true for floating point types"); - BroadCastMLFloat16FMod(ctx); + void operator()(bool fmod, bool supports_float_floor_mod, OpKernelContext* ctx) const { + ORT_ENFORCE(fmod || supports_float_floor_mod, + "fmod attribute must be true for floating point types before opset 28"); + if (fmod) { + BroadCastMLFloat16FMod(ctx); + } else { + BroadCastMLFloat16FloorMod(ctx); + } } }; } // namespace mod_internal Mod::Mod(const OpKernelInfo& info) : OpKernel(info) { + supports_float_floor_mod_ = info.node().SinceVersion() >= 28; int64_t fmod = 0; Status s = info.GetAttr("fmod", &fmod); if (s.IsOK()) { @@ -2327,7 +2503,7 @@ Status Mod::Compute(OpKernelContext* context) const { } utils::MLTypeCallDispatcherFromTypeList t_disp(dt_type); - t_disp.Invoke(fmod_, context); + t_disp.Invoke(fmod_, supports_float_floor_mod_, context); return Status::OK(); } diff --git a/onnxruntime/core/providers/cpu/math/gemm.cc b/onnxruntime/core/providers/cpu/math/gemm.cc index 8ca374fe543e5..c0db108f485a8 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.cc +++ b/onnxruntime/core/providers/cpu/math/gemm.cc @@ -219,18 +219,25 @@ void Gemm_MLFloat16(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE trans_b, } const bool use_mlas_no_bias = beta == onnxruntime::MLFloat16::Zero; const bool use_mlas_bias = beta == onnxruntime::MLFloat16::One && support_mlas_bias; - if (has_accelerated_half_gemm && trans_a == CblasNoTrans && trans_b == CblasNoTrans && + const bool half_gemm_supports_trans_b = + trans_b == CblasNoTrans || + (trans_b == CblasTrans && + MlasHalfGemmNativePackBSize(CblasNoTrans, CblasTrans, + static_cast(N), static_cast(K), + mlas_backend_kernel_selector_config) != 0); + if (has_accelerated_half_gemm && half_gemm_supports_trans_b && trans_a == CblasNoTrans && alpha == onnxruntime::MLFloat16::One && (use_mlas_no_bias || use_mlas_bias)) { MLAS_HALF_GEMM_DATA_PARAMS data{}; data.A = a_data; data.lda = K; data.B = b_data; - data.ldb = N; + data.ldb = trans_b == CblasNoTrans ? N : K; data.C = y_data; data.ldc = N; if (use_mlas_bias && c_shape != nullptr) { data.Bias = c_data; } + data.BIsTransposed = trans_b == CblasTrans; data.BackendKernelSelectorConfig = mlas_backend_kernel_selector_config; MlasHalfGemmBatch(M, N, K, 1, &data, thread_pool); return; diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc index d317250267689..cb3a4f5d982f8 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.cc +++ b/onnxruntime/core/providers/cpu/math/matmul.cc @@ -179,9 +179,8 @@ Status MatMul::Compute(OpKernelContext* ctx) const { return Status::OK(); if (helper.K() == 0) { - EigenMatrixMapRowMajor dest(y->MutableData(), - narrow(helper.M()), narrow(helper.N())); - dest.setZero(); + auto output_span = gsl::make_span(y->MutableData(), narrow(y->Shape().Size())); + std::fill(output_span.begin(), output_span.end(), double{}); return Status::OK(); } @@ -476,9 +475,8 @@ Status MatMul::Compute(OpKernelContext* ctx) const { if (helper.K() == 0) { // When we have (M, 0, N) then the inputs are empty, but the output should // be filled out with zeros. - EigenMatrixMapRowMajor dest(y->MutableData(), - narrow(helper.M()), narrow(helper.N())); - dest.setZero(); + auto output_span = gsl::make_span(y->MutableData(), narrow(y->Shape().Size())); + std::fill(output_span.begin(), output_span.end(), float{}); return Status::OK(); } diff --git a/onnxruntime/core/providers/cpu/math/matmul_helper.h b/onnxruntime/core/providers/cpu/math/matmul_helper.h index 7ff77d18d1738..65a9a54e99759 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_helper.h +++ b/onnxruntime/core/providers/cpu/math/matmul_helper.h @@ -47,7 +47,7 @@ class MatMulComputeHelper { // A: [M1, M2, ... K], B: [1, ..., 1, K, N] // A: [M1, M2, ... K], B: [1, ..., 1, N, K]^T if (!transa && !trans_batch_a && !trans_batch_b && left_num_dims >= 2 && right_num_dims >= 2 && left_num_dims >= right_num_dims && - orig_right_shape.SizeToDimension(right_num_dims - 1) == orig_right_shape[right_num_dims - 2]) { + orig_right_shape.SizeToDimension(right_num_dims - 2) == 1) { M_ = static_cast(orig_left_shape.SizeToDimension(left_num_dims - 1)); K_ = static_cast(orig_left_shape[left_num_dims - 1]); N_ = static_cast(transb ? orig_right_shape[right_num_dims - 2] : orig_right_shape[right_num_dims - 1]); @@ -137,7 +137,10 @@ class MatMulComputeHelper { // broadcasting for all output dims except last two for (size_t idx_dim = 0; idx_dim < num_dims_with_pad - 2; ++idx_dim) { - output_dims[idx_dim] = std::max(left_padded_dims_[idx_dim], right_padded_dims_[idx_dim]); + const auto left_dim = left_padded_dims_[idx_dim]; + const auto right_dim = right_padded_dims_[idx_dim]; + const auto min_dim = std::min(left_dim, right_dim); + output_dims[idx_dim] = min_dim == 0 ? 0 : std::max(left_dim, right_dim); if (left_padded_dims_[idx_dim] != output_dims[idx_dim]) ORT_RETURN_IF_NOT(left_padded_dims_[idx_dim] == 1, "left operand cannot broadcast on dim ", idx_dim); if (right_padded_dims_[idx_dim] != output_dims[idx_dim]) @@ -231,7 +234,7 @@ class MatMulComputeHelper { output_broadcast_dims_.resize(num_broadcasted_dims_); for (size_t i = num_broadcasted_dims_; i > 0; --i) { size_t idx = i - 1; - output_broadcast_dims_[idx] = std::max(left_padded_dims_[idx], right_padded_dims_[idx]); + output_broadcast_dims_[idx] = static_cast(output_shape_[idx]); output_broadcast_strides_[idx] = ((i == num_broadcasted_dims_) ? 1 : output_broadcast_strides_[idx + 1] * output_broadcast_dims_[idx + 1]); left_padded_strides_[idx] = ((i == num_broadcasted_dims_) ? 1 : left_padded_strides_[idx + 1] * left_padded_dims_[idx + 1]); right_padded_strides_[idx] = ((i == num_broadcasted_dims_) ? 1 : right_padded_strides_[idx + 1] * right_padded_dims_[idx + 1]); diff --git a/onnxruntime/core/providers/cpu/nn/conv.cc b/onnxruntime/core/providers/cpu/nn/conv.cc index bccf67fd68ab2..2205f8cf705e5 100644 --- a/onnxruntime/core/providers/cpu/nn/conv.cc +++ b/onnxruntime/core/providers/cpu/nn/conv.cc @@ -15,6 +15,7 @@ */ /* Modifications Copyright (c) Microsoft. */ +#include #include #include "core/providers/cpu/nn/conv.h" @@ -102,6 +103,19 @@ Status Conv::Compute(OpKernelContext* context) const { return Status::OK(); } + if (C == 0) { + T* output = Y->MutableData(); + const size_t spatial_size = narrow(output_shape.Size()); + const T* bias = B == nullptr ? nullptr : B->Data(); + for (int64_t n = 0; n < N; ++n) { + for (int64_t c = 0; c < M; ++c) { + std::fill_n(output, spatial_size, bias == nullptr ? T{} : bias[c]); + output += spatial_size; + } + } + return Status::OK(); + } + const int64_t input_image_size = input_shape.Size(); const int64_t output_image_size = output_shape.Size(); const int64_t kernel_size = TensorShape(kernel_shape).Size(); @@ -295,6 +309,33 @@ Status Conv::Compute(OpKernelContext* context) const { return Status::OK(); } + if (C == 0) { + if (Sum != nullptr) { + ORT_RETURN_IF_NOT(Sum->Shape() == Y->Shape(), "output and sum shape must match"); + } + const size_t output_size = narrow(Y->Shape().Size()); + const size_t spatial_size = narrow(output_shape.Size()); + const size_t batch_count = narrow(N); + const size_t channel_count = narrow(M); + const float* bias = B == nullptr ? nullptr : B->Data(); + const float* sum = Sum == nullptr ? nullptr : Sum->Data(); + float* output = Y->MutableData(); + + for (size_t n = 0; n < batch_count; ++n) { + for (size_t c = 0; c < channel_count; ++c) { + for (size_t i = 0; i < spatial_size; ++i) { + const size_t index = channels_last_ + ? (n * spatial_size + i) * channel_count + c + : (n * channel_count + c) * spatial_size + i; + output[index] = (bias == nullptr ? 0.0f : bias[c]) + (sum == nullptr ? 0.0f : sum[index]); + } + } + } + + MlasActivation(&activation_, output, nullptr, 1, output_size, output_size); + return Status::OK(); + } + AllocatorPtr alloc; ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc)); diff --git a/onnxruntime/core/providers/cpu/tensor/cast_op.cc b/onnxruntime/core/providers/cpu/tensor/cast_op.cc index 1d5a7c63228b3..22a3e7b62c458 100644 --- a/onnxruntime/core/providers/cpu/tensor/cast_op.cc +++ b/onnxruntime/core/providers/cpu/tensor/cast_op.cc @@ -1212,9 +1212,20 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( .MayInplace(0, 0), // allocation planner will check input and output sizes match before inplacing Cast); -ONNX_CPU_OPERATOR_KERNEL( +ONNX_CPU_OPERATOR_VERSIONED_KERNEL( Cast, 25, + 27, + KernelDefBuilder() + .TypeConstraint("T1", BuildKernelDefConstraintsFromTypeList()) + .TypeConstraint("T2", BuildKernelDefConstraintsFromTypeList()) + .MayInplace(0, 0), + Cast); + +// FP6 conversions are deferred; preserve the existing types for CPU fallback casts. +ONNX_CPU_OPERATOR_KERNEL( + Cast, + 28, KernelDefBuilder() .TypeConstraint("T1", BuildKernelDefConstraintsFromTypeList()) .TypeConstraint("T2", BuildKernelDefConstraintsFromTypeList()) diff --git a/onnxruntime/core/providers/cpu/tensor/onehot.cc b/onnxruntime/core/providers/cpu/tensor/onehot.cc index c6b42c2f560a9..25ee45dc9f794 100644 --- a/onnxruntime/core/providers/cpu/tensor/onehot.cc +++ b/onnxruntime/core/providers/cpu/tensor/onehot.cc @@ -16,7 +16,9 @@ limitations under the License. #include "core/providers/cpu/tensor/onehot.h" +#include #include +#include #include "core/common/eigen_common_wrapper.h" #include "core/common/safeint.h" @@ -151,11 +153,11 @@ struct EigenTensorTypes { }; namespace generator { -template +template class OneGenerator { public: EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE - OneGenerator(const typename EigenTensorTypes::ConstMatrix& indices, + OneGenerator(const typename EigenTensorTypes::ConstMatrix& indices, const typename EigenTensorTypes::ConstScalar& on_value, const typename EigenTensorTypes::ConstScalar& off_value) : indices_(indices), on_value_(on_value), off_value_(off_value) {} @@ -168,7 +170,7 @@ class OneGenerator { } private: - const typename EigenTensorTypes::ConstMatrix indices_; + const typename EigenTensorTypes::ConstMatrix indices_; const typename EigenTensorTypes::ConstScalar on_value_; const typename EigenTensorTypes::ConstScalar off_value_; }; @@ -207,21 +209,35 @@ Status OneHotOp::Compute(OpKernelContext* p_op_ke Eigen::array indices_dims_e = { {static_cast(prefix_dim_size), static_cast(suffix_dim_size)}}; - // Handle negative indices. It's faster to create a new indices instead of comparing in generator - // since generator has much larger loops. + // Cast floating-point indices to int64 before handling negative values, as required by ONNX. + // Keep invalid values out of range without converting them to int64. const auto* indices_data = indices->Data(); const auto indices_size = indices->Shape().Size(); - std::vector adjusted_indices; + using index_type = std::conditional_t, int64_t, in_type>; + std::vector adjusted_indices; adjusted_indices.reserve(onnxruntime::narrow(indices_size)); for (int64_t i = 0; i < indices_size; ++i) { - if (indices_data[i] < 0) - adjusted_indices.push_back(indices_data[i] + static_cast(depth_val)); - else - adjusted_indices.push_back(indices_data[i]); + index_type index; + if constexpr (std::is_floating_point_v) { + const double value = static_cast(indices_data[i]); + const double int64_min = static_cast(std::numeric_limits::min()); + const double int64_max_exclusive = -int64_min; + if (!std::isfinite(value) || value < int64_min || value >= int64_max_exclusive) { + adjusted_indices.push_back(static_cast(depth_val)); + continue; + } + index = static_cast(value); + } else { + index = indices_data[i]; + } + + if (index < 0) { + index += static_cast(depth_val); + } + adjusted_indices.push_back(index); } - indices_data = adjusted_indices.data(); - typename EigenTensorTypes::ConstEigenTensorMap indices_tensor_e(indices_data, indices_dims_e); + typename EigenTensorTypes::ConstEigenTensorMap indices_tensor_e(adjusted_indices.data(), indices_dims_e); // Split output into 3-Tensor of size: // prefix_dim_size x depth x suffix_dim_size. @@ -233,7 +249,7 @@ Status OneHotOp::Compute(OpKernelContext* p_op_ke typename EigenTensorTypes::ConstScalar on_value_e(values_data + 1); typename EigenTensorTypes::ConstScalar off_value_e(values_data); - generator::OneGenerator generator(indices_tensor_e, on_value_e, off_value_e); + generator::OneGenerator generator(indices_tensor_e, on_value_e, off_value_e); // TODO potential optimization opportunity // TODO figure out the eigen threadpool stuff for use here diff --git a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc index f8b78e80e7cac..f59cd88508b71 100644 --- a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc +++ b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc @@ -18,9 +18,10 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( DataTypeImpl::GetTensorType()}), SpaceToDepth); -ONNX_CPU_OPERATOR_KERNEL( +ONNX_CPU_OPERATOR_VERSIONED_KERNEL( SpaceToDepth, 13, + 27, KernelDefBuilder() .TypeConstraint("T", {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), diff --git a/onnxruntime/core/providers/cuda/activation/activations_impl.cu b/onnxruntime/core/providers/cuda/activation/activations_impl.cu index 811f49a0cca19..02c32591ee373 100644 --- a/onnxruntime/core/providers/cuda/activation/activations_impl.cu +++ b/onnxruntime/core/providers/cuda/activation/activations_impl.cu @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include +#include #include "activations_impl.h" #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cu_inc/unary_elementwise_impl.cuh" @@ -47,28 +48,22 @@ struct OP_Selu : public CtxSelu { template struct OP_Sigmoid : public CtxSigmoid { __device__ __inline__ T operator()(const T& a) const { - return a > T(0) ? (T)1 / ((T)1. + _Exp(-a)) : (T)1 - (T)1 / ((T)1 + _Exp(a)); - } -}; - -template <> -struct OP_Sigmoid : public CtxSigmoid { - __device__ __inline__ half operator()(const half& a) const { - // For some small negative values, it will loss precision if using half for compute. - // The cast between half and float below won't cause perf issue since there is cast inside _Exp if input is half. - float af = static_cast(a); - float res = af > 0.0f ? 1.0f / (1.0f + _Exp(-af)) : 1.0f - 1.0f / (1.0f + _Exp(af)); - return static_cast(res); + using ComputeT = std::conditional_t, double, float>; + const ComputeT value = static_cast(a); + const ComputeT exponent = _Exp(value > ComputeT(0) ? -value : value); + // Subtracting from one for negative logits erases small, representable probabilities. + return static_cast(value > ComputeT(0) ? ComputeT(1) / (ComputeT(1) + exponent) + : exponent / (ComputeT(1) + exponent)); } }; template struct OP_Softplus : public CtxSoftplus { __device__ __inline__ T operator()(const T& a) const { - if (a > (T)0) - return a + _Log(_Exp(-a) + (T)1); - else - return _Log(_Exp(a) + (T)1); + using ComputeT = std::conditional_t, double, float>; + const ComputeT value = static_cast(a); + // Adding one before taking log erases small positive tails, especially in FP16/BF16. + return static_cast(value > ComputeT(0) ? value + log1p(_Exp(-value)) : log1p(_Exp(value))); } }; diff --git a/onnxruntime/core/providers/cuda/cu_inc/common.cuh b/onnxruntime/core/providers/cuda/cu_inc/common.cuh index f53da982538bd..93257da22e269 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/common.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/common.cuh @@ -501,6 +501,10 @@ __device__ __inline__ half _Gelu(half a) { template __device__ __inline__ T _Mod(T a, T b) { + if (a == std::numeric_limits::min() && b == T(-1)) { + return T(0); + } + T r = a % b; T zero = T(0); if ((r > zero && b < zero) || (r < zero && b > zero)) { @@ -509,8 +513,46 @@ __device__ __inline__ T _Mod(T a, T b) { return r; } +template <> +__device__ __inline__ float _Mod(float a, float b) { + float r = fmodf(a, b); + if (r == 0.0f) { + return copysignf(0.0f, b); + } + if ((r > 0.0f && b < 0.0f) || (r < 0.0f && b > 0.0f)) { + r += b; + } + return r; +} + +template <> +__device__ __inline__ double _Mod(double a, double b) { + double r = fmod(a, b); + if (r == 0.0) { + return copysign(0.0, b); + } + if ((r > 0.0 && b < 0.0) || (r < 0.0 && b > 0.0)) { + r += b; + } + return r; +} + +template <> +__device__ __inline__ half _Mod(half a, half b) { + return half(_Mod(static_cast(a), static_cast(b))); +} + +template <> +__device__ __inline__ BFloat16 _Mod(BFloat16 a, BFloat16 b) { + return BFloat16(_Mod(static_cast(a), static_cast(b))); +} + template __device__ __inline__ T _Fmod(T a, T b) { + if (a == std::numeric_limits::min() && b == T(-1)) { + return T(0); + } + return a % b; } diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 5fa508ee6448d..6c1a9a3195b44 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -46,6 +46,10 @@ #include "core/providers/cuda/cuda_stream_handle.h" +// Persistent fpA_intB tactic cache flush at EP teardown (no-op unless onnxruntime_USE_FPA_INTB_GEMM). +// The header is self-contained (no contrib/CUTLASS dependencies) and resolves to an inline no-op +// when the fpA_intB path is not compiled in. +#include "contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h" #if !defined(DISABLE_CONTRIB_OPS) && USE_FPA_INTB_GEMM #include "contrib_ops/cuda/quantization/matmul_nbits_workspace_estimate.h" #endif @@ -413,6 +417,10 @@ std::optional CUDAExecutionProvider::ShouldConvertDataLayoutForOp([[maybe_ } CUDAExecutionProvider::~CUDAExecutionProvider() { + // Persist fpA_intB MatMulNBits tactics staged in memory to disk (best-effort, dirty-guarded). + // No-op in builds without onnxruntime_USE_FPA_INTB_GEMM. The plugin EP flushes from ~CudaEp. + onnxruntime::contrib::cuda::FlushMatMulNBitsTacticCaches(); + // clean up thread local context caches { std::lock_guard lock(context_state_.mutex); @@ -570,6 +578,23 @@ void CUDAExecutionProvider::RetainBufferForGraphCapture(std::shared_ptr bu captured_host_buffers_.push_back(std::move(buffer)); } +void CUDAExecutionProvider::SetSessionModelPath(std::filesystem::path path) const { + if (!path.empty() && path.is_relative()) { + std::error_code ec; + std::filesystem::path abs_path = std::filesystem::absolute(path, ec); + if (!ec && !abs_path.empty()) { + path = std::move(abs_path); + } + } + std::lock_guard lock(model_path_mutex_); + model_path_ = std::move(path); +} + +std::filesystem::path CUDAExecutionProvider::GetSessionModelPath() const { + std::lock_guard lock(model_path_mutex_); + return model_path_; +} + namespace cuda { template <> @@ -1327,7 +1352,7 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kO class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 17, double, Pad); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 17, MLFloat16, Pad); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 17, bool, Pad); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, SpaceToDepth); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 27, SpaceToDepth); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, DepthToSpace); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int8_t, Sign); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int16_t, Sign); @@ -1354,7 +1379,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, Mod); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 27, Mod); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, Mod); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 18, int8_t, QuantizeLinear); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 18, uint8_t, QuantizeLinear); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 18, int8_t, DequantizeLinear); @@ -1842,25 +1868,25 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 24, T class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 24, 24, Unsqueeze); // Opset 25. -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, float, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, double, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, MLFloat16, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, BFloat16, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, int8_t, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, int16_t, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, int32_t, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, int64_t, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, uint8_t, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, uint16_t, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, uint32_t, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, uint64_t, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, bool, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, float, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, double, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, MLFloat16, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, BFloat16, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, int8_t, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, int16_t, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, int32_t, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, int64_t, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, uint8_t, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, uint16_t, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, uint32_t, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, uint64_t, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, bool, Cast); #if !defined(DISABLE_FLOAT8_TYPES) -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, Float8E4M3FN, Cast); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, Float8E5M2, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, Float8E4M3FN, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, Float8E5M2, Cast); #endif #if !defined(DISABLE_FLOAT4_TYPES) -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, Float4E2M1x2, Cast); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, 27, Float4E2M1x2, Cast); #endif class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, ConstantOfShape); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, uint8_t, float, DequantizeLinear); @@ -1909,6 +1935,28 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 25, U // Opset 27. class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 27, Range); + +// Opset 28. +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, float, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, double, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, MLFloat16, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, BFloat16, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, int8_t, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, int16_t, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, int32_t, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, int64_t, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, uint8_t, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, uint16_t, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, uint32_t, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, uint64_t, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, bool, Cast); +#if !defined(DISABLE_FLOAT8_TYPES) +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, Float8E4M3FN, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, Float8E5M2, Cast); +#endif +#if !defined(DISABLE_FLOAT4_TYPES) +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, Float4E2M1x2, Cast); +#endif #endif static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { @@ -2620,7 +2668,7 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -2647,7 +2695,8 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -3136,25 +3185,25 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, // Opset 25 - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, #if !defined(DISABLE_FLOAT8_TYPES) - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, #endif #if !defined(DISABLE_FLOAT4_TYPES) - BuildKernelCreateInfo, + BuildKernelCreateInfo, #endif BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -3202,6 +3251,27 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, // Opset 27 BuildKernelCreateInfo, + // Opset 28 + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, +#if !defined(DISABLE_FLOAT8_TYPES) + BuildKernelCreateInfo, + BuildKernelCreateInfo, +#endif +#if !defined(DISABLE_FLOAT4_TYPES) + BuildKernelCreateInfo, +#endif #endif }; @@ -3245,6 +3315,9 @@ void InitializeRegistry() { } void DeleteRegistry() { + // Save any tactic caches of sessions still alive when the provider is shut down, before worker + // threads are stopped or the provider DLL is unloaded. + onnxruntime::contrib::cuda::FlushMatMulNBitsTacticCaches(); s_kernel_registry.reset(); } @@ -3449,6 +3522,10 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, std::vector> result; const logging::Logger& logger = *GetLogger(); + if (!graph.ModelPath().empty()) { + SetSessionModelPath(graph.ModelPath()); + } + // Figure out the memory limit if accountant is available size_t memory_threshold = std::numeric_limits::max(); SafeInt consumed_memory = 0; @@ -3636,7 +3713,7 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, input_defs.size() > 11 && input_defs[11] != nullptr && input_defs[11]->Exists() && graph.IsConstantInitializer(input_defs[11]->Name(), true); - // The factory validates the envelope once. It remains reader-only here. + // The factory validates the caller-declared envelope once. const auto gqa_workspace_max_total_sequence_length = resource_accountant->GetWorkspaceEstimatorConfig() .cuda_gqa_workspace_max_total_sequence_length; @@ -3651,6 +3728,13 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, LOGS(logger, VERBOSE) << "Level-1 memory estimate for " << node->Name() << ": runtime workspace=" << ws->total_workspace_bytes << " bytes"; + } else { + // The generic safety margin is an ad-hoc fallback, not a proven GQA + // workspace bound. Capacity-aware placement therefore fails closed. + LOGS(logger, WARNING) + << "CUDA_EP declining capacity-aware assignment for GroupQueryAttention node " + << node->Name() << " because no proven route-aware workspace estimate is available."; + continue; } } #endif diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.h b/onnxruntime/core/providers/cuda/cuda_execution_provider.h index c61ef9f2c2109..841c9ef7aeb85 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.h @@ -4,6 +4,7 @@ #pragma once +#include // NOLINT(build/c++17) #include #include @@ -131,10 +132,18 @@ class CUDAExecutionProvider : public IExecutionProvider { // by a later run. Held until the provider is destroyed. void RetainBufferForGraphCapture(std::shared_ptr buffer) const; + // Path of the model this provider partitioned, recorded in GetCapability. Empty for models loaded + // from bytes, or before partitioning. + void SetSessionModelPath(std::filesystem::path path) const; + std::filesystem::path GetSessionModelPath() const; + private: mutable std::mutex captured_host_buffers_mutex_; mutable std::vector> captured_host_buffers_; + mutable std::mutex model_path_mutex_; + mutable std::filesystem::path model_path_; + CUDAExecutionProviderInfo info_; cudaDeviceProp device_prop_; bool external_stream_ = false; diff --git a/onnxruntime/core/providers/cuda/cuda_kernel.h b/onnxruntime/core/providers/cuda/cuda_kernel.h index 22e39f434ece5..0df29f53acd0c 100644 --- a/onnxruntime/core/providers/cuda/cuda_kernel.h +++ b/onnxruntime/core/providers/cuda/cuda_kernel.h @@ -121,6 +121,9 @@ class CudaKernel : public OpKernel { } const cudaDeviceProp& GetDeviceProp() const { return provider_->GetDeviceProp(); } + std::filesystem::path GetSessionModelPath() const { return provider_->GetSessionModelPath(); } + // Identifies the session (execution provider instance) that owns this kernel. + const void* GetSessionKey() const { return provider_; } int GetCudnnConvAlgo() const { return provider_->GetCudnnConvAlgo(); } bool GetCudnnConvUseMaxWorkspace() const { return provider_->GetCudnnConvUseMaxWorkspace(); } bool GetCudnnConv1dPadToNc1d() const { return provider_->GetCudnnConv1dPadToNc1d(); } diff --git a/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc b/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc index 388a30d4a289f..22f14b6123eb8 100755 --- a/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc +++ b/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc @@ -85,7 +85,7 @@ class CUDA_NHWC_OP_VERSIONED_CLASS_NAME(1, 10, DepthToSpace); class CUDA_NHWC_OP_VERSIONED_CLASS_NAME(11, 12, DepthToSpace); class CUDA_NHWC_OP_CLASS_NAME(13, DepthToSpace); class CUDA_NHWC_OP_VERSIONED_CLASS_NAME(1, 12, SpaceToDepth); -class CUDA_NHWC_OP_CLASS_NAME(13, SpaceToDepth); +class CUDA_NHWC_OP_VERSIONED_CLASS_NAME(13, 27, SpaceToDepth); class CUDA_NHWC_OP_VERSIONED_TYPED_CLASS_NAME(1, 12, float, LRN); class CUDA_NHWC_OP_VERSIONED_TYPED_CLASS_NAME(1, 12, double, LRN); class CUDA_NHWC_OP_VERSIONED_TYPED_CLASS_NAME(1, 12, MLFloat16, LRN); @@ -154,7 +154,7 @@ Status RegisterCudaNhwcKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index b1a818e10c101..1e16328ef59fe 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -694,10 +694,12 @@ struct CudaSyncStreamImpl : OrtSyncStreamImpl { struct CudaEpFactory : OrtEpFactory { using MemoryInfoUniquePtr = std::unique_ptr>; - CudaEpFactory(const OrtApi& ort_api_in, const OrtLogger& default_logger_in) : ort_api{ort_api_in}, + CudaEpFactory(const OrtApi& ort_api_in, const OrtLogger& default_logger_in) : OrtEpFactory{}, + ort_api{ort_api_in}, ep_api{*ort_api_in.GetEpApi()}, default_logger{default_logger_in}, data_transfer_impl{ort_api_in} { + ort_version_supported = ORT_API_VERSION; GetName = GetNameImpl; GetVendor = GetVendorImpl; GetVendorId = GetVendorIdImpl; diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc index babbb4b3ba672..0340f3b4c7e8b 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc @@ -467,7 +467,14 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( BuildKernelDefConstraints()), Mod); -ONNX_OPERATOR_KERNEL_EX(Mod, kOnnxDomain, 13, kCudaExecutionProvider, +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Mod, kOnnxDomain, 13, 27, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", BuildKernelDefConstraints()), + Mod); + +ONNX_OPERATOR_KERNEL_EX(Mod, kOnnxDomain, 28, kCudaExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", BuildKernelDefConstraints()), @@ -478,10 +485,10 @@ Status Mod::ComputeInternal(OpKernelContext* context) const { BinaryElementwisePreparation prepare; ORT_RETURN_IF_ERROR(Prepare(context, &prepare)); auto element_type = prepare.lhs_tensor->GetElementType(); - ORT_ENFORCE(fmod_ || element_type == on::TensorProto_DataType_INT32 || + ORT_ENFORCE(fmod_ || supports_float_floor_mod_ || element_type == on::TensorProto_DataType_INT32 || element_type == on::TensorProto_DataType_INT64 || element_type == on::TensorProto_DataType_UINT32 || element_type == on::TensorProto_DataType_UINT64, - "Non-fmod can support integer types only."); + "Non-fmod Mod supports floating-point types only starting in opset 28 (set fmod=1 for older opsets)."); #define CASE_MOD_ELEMENT_TYPE(name, onnx_type, data_type) \ case onnx_type: { \ Impl_##name::MappedType>( \ @@ -514,6 +521,10 @@ Status Mod::ComputeInternal(OpKernelContext* context) const { CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_INT64, int64_t); CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_UINT32, uint32_t); CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_UINT64, uint64_t); + CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_FLOAT, float); + CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_DOUBLE, double); + CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_FLOAT16, MLFloat16); + CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_BFLOAT16, BFloat16); default: return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported element type: ", DataTypeImpl::ToString(prepare.lhs_tensor->DataType())); diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h index 8bc40ae7d2ce7..337c7909f46e9 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h @@ -221,11 +221,13 @@ class Mod final : public BinaryElementwise { Mod(const OpKernelInfo& info) : BinaryElementwise(info) { int64_t fmod = info.GetAttrOrDefault("fmod", 0LL); fmod_ = fmod != 0; + supports_float_floor_mod_ = info.node().SinceVersion() >= 28; } Status ComputeInternal(OpKernelContext* context) const override; private: bool fmod_{false}; + bool supports_float_floor_mod_{false}; }; template diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu index 4e71a9f7c5090..2b08614b1e9c4 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu @@ -162,7 +162,7 @@ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(Xor, bool) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(PRelu) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Max) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Min) -SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZIL(Mod) +SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Mod) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Fmod) // create declarations for impl for Pow diff --git a/onnxruntime/core/providers/cuda/nn/layer_norm_impl.cu b/onnxruntime/core/providers/cuda/nn/layer_norm_impl.cu index 0e6f4396086df..aab461529143c 100644 --- a/onnxruntime/core/providers/cuda/nn/layer_norm_impl.cu +++ b/onnxruntime/core/providers/cuda/nn/layer_norm_impl.cu @@ -240,7 +240,7 @@ struct SharedMemory { }; } // namespace -template +template __global__ void cuApplyLayerNorm( V* __restrict__ output_vals, U* __restrict__ mean, @@ -292,8 +292,16 @@ __global__ void cuApplyLayerNorm( for (int i = thrx; i < n2 / 2; i += numx) { const float2 input_pair = __half22float2(input2[i]); const float2 gamma_pair = __half22float2(gamma2[i]); - output2[i] = __floats2half2_rn(gamma_pair.x * c_inv_std_dev * input_pair.x, - gamma_pair.y * c_inv_std_dev * input_pair.y); + if constexpr (stash_normalized) { + const half2 normalized = __floats2half2_rn(c_inv_std_dev * input_pair.x, + c_inv_std_dev * input_pair.y); + const float2 normalized_pair = __half22float2(normalized); + output2[i] = __floats2half2_rn(gamma_pair.x * normalized_pair.x, + gamma_pair.y * normalized_pair.y); + } else { + output2[i] = __floats2half2_rn(gamma_pair.x * c_inv_std_dev * input_pair.x, + gamma_pair.y * c_inv_std_dev * input_pair.y); + } } use_scalar = false; } @@ -315,7 +323,15 @@ __global__ void cuApplyLayerNorm( U beta_i = (beta != nullptr) ? (U)beta[index] : (U)0; if (simplified) { - ovals[i] = static_cast(gamma_i * c_inv_std_dev * curr); + if constexpr (stash_normalized) { + // ONNX RMSNormalization rounds standardization to X's type before scaling. + const T normalized = static_cast(c_inv_std_dev * curr); + using ScaleType = std::conditional_t, double, U>; + const ScaleType scale = gamma != nullptr ? static_cast(gamma[index]) : ScaleType(1); + ovals[i] = static_cast(scale * static_cast(normalized)); + } else { + ovals[i] = static_cast(gamma_i * c_inv_std_dev * curr); + } } else { ovals[i] = static_cast(gamma_i * c_inv_std_dev * (curr - mu) + beta_i); } @@ -343,7 +359,7 @@ int32_t round_up_power_of_2(int32_t v) { return v; } -template +template void HostApplyLayerNorm( const cudaDeviceProp& prop, cudaStream_t stream, @@ -379,7 +395,7 @@ void HostApplyLayerNorm( const dim3 blocks(1, std::min(n1, maxGridY), 1); int nshared = threads.y > 1 ? threads.y * sizeof(U) + (threads.y / 2) * sizeof(U) : 0; - cuApplyLayerNorm<<>>( + cuApplyLayerNorm<<>>( output, mean, inv_std_dev, @@ -411,5 +427,20 @@ LAYERNORM_LINEAR_IMPL(half, float, float, false) LAYERNORM_LINEAR_IMPL(BFloat16, float, BFloat16, true) LAYERNORM_LINEAR_IMPL(BFloat16, float, BFloat16, false) +#define RMSNORM_LINEAR_IMPL(T, U, V) \ + template void HostApplyLayerNorm(const cudaDeviceProp& prop, cudaStream_t stream, V* output, \ + U* mean, U* inv_std_dev, const T* input, int n1, int n2, \ + double epsilon, const V* gamma, const V* beta, \ + int broadcast_param, \ + const T* skip, const T* bias, T* skip_input_bias_add_output); + +RMSNORM_LINEAR_IMPL(float, float, float) +RMSNORM_LINEAR_IMPL(half, float, half) +RMSNORM_LINEAR_IMPL(double, double, double) +RMSNORM_LINEAR_IMPL(double, float, double) +RMSNORM_LINEAR_IMPL(float, float, half) +RMSNORM_LINEAR_IMPL(half, float, float) +RMSNORM_LINEAR_IMPL(BFloat16, float, BFloat16) + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/nn/layer_norm_impl.h b/onnxruntime/core/providers/cuda/nn/layer_norm_impl.h index 4e74aa9ab67b4..c7e5e9ee2329d 100644 --- a/onnxruntime/core/providers/cuda/nn/layer_norm_impl.h +++ b/onnxruntime/core/providers/cuda/nn/layer_norm_impl.h @@ -28,7 +28,7 @@ namespace onnxruntime { namespace cuda { -template +template void HostApplyLayerNorm( const cudaDeviceProp& prop, cudaStream_t stream, diff --git a/onnxruntime/core/providers/cuda/nn/rms_norm.cc b/onnxruntime/core/providers/cuda/nn/rms_norm.cc index 1f0eb1d60ed08..1d20c75224114 100644 --- a/onnxruntime/core/providers/cuda/nn/rms_norm.cc +++ b/onnxruntime/core/providers/cuda/nn/rms_norm.cc @@ -12,21 +12,13 @@ namespace onnxruntime { namespace cuda { -// For T and V both are double, we need U to be double -template -using MeanVarType = typename std::conditional< - std::is_same::value && std::is_same::value, - double, - float>::type; - -// RMSNorm uses LayerNorm kernel, which only supports X and scale both -// being the same data type. +// The default stage-one stash is FP32 independently of input and scale types. #define REGISTER_KERNEL_TYPED(T, V) \ ONNX_OPERATOR_TYPED_KERNEL_EX(RMSNormalization, kOnnxDomain, 23, T##_##V, kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ .TypeConstraint("V", DataTypeImpl::GetTensorType()), \ - RMSNorm, V>); + RMSNorm); REGISTER_KERNEL_TYPED(float, float) REGISTER_KERNEL_TYPED(float, MLFloat16) @@ -44,6 +36,11 @@ RMSNorm::RMSNorm(const OpKernelInfo& op_kernel_info) float tmp_epsilon; ORT_ENFORCE(op_kernel_info.GetAttr("epsilon", &tmp_epsilon).IsOK()); epsilon_ = tmp_epsilon; + stash_type_ = op_kernel_info.GetAttrOrDefault("stash_type", 1); + ORT_ENFORCE(stash_type_ == 1 || + (stash_type_ == 11 && std::is_same_v && std::is_same_v), + "CUDA RMSNormalization supports stash_type=1 (FLOAT), or stash_type=11 (DOUBLE) " + "with double input and scale."); } template @@ -89,16 +86,28 @@ Status RMSNorm::ComputeInternal(OpKernelContext* ctx) const { params.norm_size > std::numeric_limits::max() / params.num_rows), "RMSNormalization input is too large for CUDA kernel indexing: norm_size exceeds safe limits or num_rows * norm_size exceeds INT_MAX"); - // For RMSNorm, we don't need mean and inv_var data, so we can pass nullptr. - CudaU* mean_data = nullptr; - CudaU* inv_var_data = nullptr; - - // simplified should always be true for RMSNorm - HostApplyLayerNorm( - GetDeviceProp(), Stream(ctx), Y_data, mean_data, inv_var_data, X_data, - onnxruntime::narrow(params.num_rows), onnxruntime::narrow(params.norm_size), epsilon_, - scale_data, bias_data, - onnxruntime::narrow(params.broadcast_param)); + // Unlike historical SimplifiedLayerNormalization, ONNX RMSNorm rounds before scaling. + if constexpr (std::is_same_v && std::is_same_v) { + if (stash_type_ == 11) { + HostApplyLayerNorm( + GetDeviceProp(), Stream(ctx), Y_data, nullptr, nullptr, X_data, + onnxruntime::narrow(params.num_rows), onnxruntime::narrow(params.norm_size), epsilon_, + scale_data, bias_data, + onnxruntime::narrow(params.broadcast_param)); + } else { + HostApplyLayerNorm( + GetDeviceProp(), Stream(ctx), Y_data, nullptr, nullptr, X_data, + onnxruntime::narrow(params.num_rows), onnxruntime::narrow(params.norm_size), epsilon_, + scale_data, bias_data, + onnxruntime::narrow(params.broadcast_param)); + } + } else { + HostApplyLayerNorm( + GetDeviceProp(), Stream(ctx), Y_data, nullptr, nullptr, X_data, + onnxruntime::narrow(params.num_rows), onnxruntime::narrow(params.norm_size), epsilon_, + scale_data, bias_data, + onnxruntime::narrow(params.broadcast_param)); + } CUDA_RETURN_IF_ERROR(cudaGetLastError()); return Status::OK(); } diff --git a/onnxruntime/core/providers/cuda/nn/rms_norm.h b/onnxruntime/core/providers/cuda/nn/rms_norm.h index 9f86fcc4470ee..8019669efb227 100644 --- a/onnxruntime/core/providers/cuda/nn/rms_norm.h +++ b/onnxruntime/core/providers/cuda/nn/rms_norm.h @@ -17,8 +17,7 @@ class RMSNorm final : public CudaKernel { private: int64_t axis_; double epsilon_; - // The stash_type is not supported in the current implementation. - // int64_t stash_type; + int64_t stash_type_; }; } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc b/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc index 9442c843f1197..cc1ab3a85def7 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc @@ -24,6 +24,7 @@ #include "core/graph/constants.h" #include "core/providers/cuda/cuda_nhwc_ops.h" +#include "contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h" namespace onnxruntime { namespace cuda_plugin { @@ -225,6 +226,9 @@ CudaEp::CudaEp(CudaEpFactory& factory, const Config& config, const OrtLogger& lo } CudaEp::~CudaEp() { + // The plugin build excludes cuda_execution_provider.cc, so persist staged fpA_intB tactics here. + onnxruntime::contrib::cuda::FlushMatMulNBitsTacticCaches(); + std::lock_guard lock(per_thread_contexts_mutex_); for (const auto& cache_weak : per_thread_context_caches_) { auto cache = cache_weak.lock(); @@ -257,6 +261,14 @@ OrtStatus* ORT_API_CALL CudaEp::GetCapabilityImpl( return nullptr; } + if (auto model_path = graph.GetModelPath(); !model_path.empty()) { + std::error_code ec; + std::filesystem::path path(model_path); + std::filesystem::path abs_path = std::filesystem::absolute(path, ec); + onnxruntime::cuda::detail::GetCudaKernelAdapterRuntimeConfigForProvider(static_cast(ep->EpImpl())) + ->SetModelPath((!ec && !abs_path.empty()) ? std::move(abs_path) : std::move(path)); + } + // Three-phase filtering determines which graph nodes run on this EP: // Phase 1: Collect tentative nodes that have a registered CUDA kernel. // Phase 2: Filter out CPU-preferred nodes (cheap ops where device-to-host diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc index c8db242e12e75..a498cd33ec949 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc @@ -5,6 +5,7 @@ #include "cuda_device_mapping.h" #include "cuda_ep.h" #include "cuda_plugin_kernels.h" +#include "contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h" #include "core/common/string_utils.h" #include "core/session/onnxruntime_c_api.h" #include @@ -54,6 +55,10 @@ CudaEpFactory::CudaEpFactory(const OrtApi& ort_api, const OrtEpApi& ep_api, } CudaEpFactory::~CudaEpFactory() { + // Save any tactic caches of sessions still alive when the factory is released, before worker + // threads are stopped or the plugin DLL is unloaded. + onnxruntime::contrib::cuda::FlushMatMulNBitsTacticCaches(); + for (auto& [key, entry] : device_cache_) { static_cast(key); for (auto& arena : entry.device_arenas) { diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index 8d9ea3725b8ee..cf322fb3615c7 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -505,6 +505,7 @@ class PluginKernelCollector { // which is initialized in CudaEpFactory::CudaEpFactory. #include +#include // NOLINT(build/c++17) #include #include #include @@ -548,6 +549,27 @@ struct CudaKernelAdapterRuntimeConfig { std::lock_guard lock(captured_host_buffers_mutex); captured_host_buffers.push_back(std::move(buffer)); } + + // Path of the model the EP partitioned (recorded in GetCapability); empty if loaded from bytes. + void SetModelPath(std::filesystem::path path) { + if (!path.empty() && path.is_relative()) { + std::error_code ec; + std::filesystem::path abs_path = std::filesystem::absolute(path, ec); + if (!ec && !abs_path.empty()) { + path = std::move(abs_path); + } + } + std::lock_guard lock(model_path_mutex); + model_path = std::move(path); + } + std::filesystem::path GetModelPath() const { + std::lock_guard lock(model_path_mutex); + return model_path; + } + + private: + mutable std::mutex model_path_mutex; + std::filesystem::path model_path; }; template struct SizeOf { @@ -1159,7 +1181,9 @@ class CudaKernel : public OpKernel { return device_prop_; } - inline int GetDeviceId() const { return device_id_; } + std::filesystem::path GetSessionModelPath() const { return runtime_config_->GetModelPath(); } + // Identifies the session (execution provider instance) that owns this kernel. + const void* GetSessionKey() const { return runtime_config_.get(); } int GetCudnnConvAlgo() const { return runtime_config_->cudnn_conv_algo; } bool GetCudnnConvUseMaxWorkspace() const { return runtime_config_->cudnn_conv_use_max_workspace; } bool GetCudnnConv1dPadToNc1d() const { return runtime_config_->cudnn_conv1d_pad_to_nc1d; } diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index 60ffd772c97f1..75024a420d3a9 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -275,7 +275,7 @@ __device__ void reduce_all( template __global__ void reduce_matrix_columns_kernel( const int num_rows, const int num_cols, const TIn* const input, TOut* const output, - TBuf* const block_reductions_buffer, int* const block_done_counts_buffer) { + TBuf* const block_reductions_buffer, int* const block_done_counts_buffer, const int input_row_stride) { const int num_blocks_in_grid_row = gridDim.x; const int row_id_in_grid = blockIdx.y; const int num_grid_rows = gridDim.y; @@ -283,7 +283,7 @@ __global__ void reduce_matrix_columns_kernel( // one row per iteration // row_id is int64_t to avoid int overflow in offset calculations for (int64_t row_id = row_id_in_grid; row_id < num_rows; row_id += num_grid_rows) { - const TIn* const row_data = input + row_id * num_cols; + const TIn* const row_data = input + row_id * input_row_stride; TOut* const row_output = output + row_id; TBuf* const row_block_reductions_buffer = block_reductions_buffer + row_id * num_blocks_in_grid_row; int* const row_block_done_counts_buffer = block_done_counts_buffer + row_id; @@ -296,8 +296,10 @@ __global__ void reduce_matrix_columns_kernel( template Status call_reduce_matrix_columns( - cudaStream_t stream, const TIn* input, TOut* output, const int num_rows, const int num_cols, void* buffer, size_t buffer_size) { + cudaStream_t stream, const TIn* input, TOut* output, const int num_rows, const int num_cols, + void* buffer, size_t buffer_size, int input_row_stride = 0) { ORT_ENFORCE(num_rows >= 0 && num_cols >= 0); + ORT_ENFORCE(input_row_stride == 0 || input_row_stride >= num_cols); using TBuf = AccumulationType_t; @@ -319,7 +321,8 @@ Status call_reduce_matrix_columns( const int shared_mem_size = sizeof(TBuf) * block_dim.x * block_dim.y / GPU_WARP_SIZE_HOST; reduce_matrix_columns_kernel <<>>( - num_rows, num_cols, input, output, block_reductions_buffer, block_done_counts_buffer); + num_rows, num_cols, input, output, block_reductions_buffer, block_done_counts_buffer, + input_row_stride == 0 ? num_cols : input_row_stride); return Status::OK(); } @@ -528,6 +531,26 @@ INSTANTIATE_REDUCE_MATRIX_COLUMNS(double); INSTANTIATE_REDUCE_MATRIX_COLUMNS(BFloat16); #undef INSTANTIATE_REDUCE_MATRIX_COLUMNS +struct RoundedSquare { + __device__ __forceinline__ float operator()(float value) const { + return __fmul_rn(value, value); + } +}; + +template +Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, + int m, int n, void* buffer, size_t buffer_size, int input_row_stride) { + return detail::call_reduce_matrix_columns( + stream, input, output, m, n, buffer, buffer_size, input_row_stride); +} + +#define INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(T) \ + template Status reduce_square_matrix_columns(cudaStream_t, const T*, float*, int, int, void*, size_t, int) +INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(float); +INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(half); +INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS(BFloat16); +#undef INSTANTIATE_REDUCE_SQUARE_MATRIX_COLUMNS + namespace detail { constexpr int kMaxReduceRank = 16; diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index b772876db6395..9c877c0743b31 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -106,6 +106,11 @@ 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); +/** Sum FP32-rounded squares across columns, preserving reduce_matrix_columns' summation order. */ +template +Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, + int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0); + /** * 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 diff --git a/onnxruntime/core/providers/cuda/tensor/cast_op.cc b/onnxruntime/core/providers/cuda/tensor/cast_op.cc index 2ed08e25d02d2..6146b2ca934d1 100644 --- a/onnxruntime/core/providers/cuda/tensor/cast_op.cc +++ b/onnxruntime/core/providers/cuda/tensor/cast_op.cc @@ -100,10 +100,20 @@ const std::vector& CastOpTypeConstraints() { .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ .TypeConstraint("T2", CastOpTypeConstraints()), \ Cast); \ + ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ + Cast, \ + kOnnxDomain, \ + 25, 27, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", CastOpTypeConstraints()), \ + Cast); \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ Cast, \ kOnnxDomain, \ - 25, \ + 28, \ T, \ kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ @@ -411,11 +421,21 @@ SPECIALIZE_IMPL(BFloat16) .TypeConstraint("T2", OutputTypeConstraints), \ Cast); -#define REGISTER_KERNEL_TYPED_25(T, OutputTypeConstraints) \ +#define REGISTER_KERNEL_TYPED_25_TO_28(T, OutputTypeConstraints) \ + ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ + Cast, \ + kOnnxDomain, \ + 25, 27, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", OutputTypeConstraints), \ + Cast); \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ Cast, \ kOnnxDomain, \ - 25, \ + 28, \ T, \ kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ @@ -425,20 +445,20 @@ SPECIALIZE_IMPL(BFloat16) #if !defined(DISABLE_FLOAT8_TYPES) -#define SPECIALIZE_IMPL_19_TO_25(T) \ +#define SPECIALIZE_IMPL_19_TO_28(T) \ REGISTER_KERNEL_TYPED_19_TO_22(T) \ REGISTER_KERNEL_TYPED_23_TO_24(T, CastOpTypeConstraints()) \ - REGISTER_KERNEL_TYPED_25(T, CastOpTypeConstraints()) \ + REGISTER_KERNEL_TYPED_25_TO_28(T, CastOpTypeConstraints()) \ template Status Cast::ComputeInternal(OpKernelContext* context) const; -SPECIALIZE_IMPL_19_TO_25(Float8E4M3FN) -SPECIALIZE_IMPL_19_TO_25(Float8E5M2) +SPECIALIZE_IMPL_19_TO_28(Float8E4M3FN) +SPECIALIZE_IMPL_19_TO_28(Float8E5M2) #endif #if !defined(DISABLE_FLOAT4_TYPES) REGISTER_KERNEL_TYPED_23_TO_24(Float4E2M1x2, {DataTypeImpl::GetTensorType()}) -REGISTER_KERNEL_TYPED_25(Float4E2M1x2, {DataTypeImpl::GetTensorType()}) +REGISTER_KERNEL_TYPED_25_TO_28(Float4E2M1x2, {DataTypeImpl::GetTensorType()}) template Status Cast::ComputeInternal(OpKernelContext* context) const; #endif diff --git a/onnxruntime/core/providers/cuda/tensor/space_depth_ops.cc b/onnxruntime/core/providers/cuda/tensor/space_depth_ops.cc index cbfc62494fde4..3c5fe31f35309 100644 --- a/onnxruntime/core/providers/cuda/tensor/space_depth_ops.cc +++ b/onnxruntime/core/providers/cuda/tensor/space_depth_ops.cc @@ -37,10 +37,11 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( SpaceToDepth); #endif -ONNX_OPERATOR_KERNEL_EX( +ONNX_OPERATOR_VERSIONED_KERNEL_EX( SpaceToDepth, kOnnxDomain, 13, + 27, kCudaExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", @@ -50,10 +51,11 @@ ONNX_OPERATOR_KERNEL_EX( SpaceToDepth); #ifdef ENABLE_CUDA_NHWC_OPS -ONNX_OPERATOR_KERNEL_EX( +ONNX_OPERATOR_VERSIONED_KERNEL_EX( SpaceToDepth, kMSInternalNHWCDomain, 13, + 27, kCudaExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp index 076027dd3672f..3edc8d216f7c1 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp @@ -391,6 +391,16 @@ HRESULT STDMETHODCALLTYPE AbiCustomRegistry::RegisterOperatorKernel( builder.SetDomain(opKernel->domain) .SinceVersion(opKernel->minimumOperatorSetVersion) .Provider(providerType); + if (isInternalOperator && + opKernel->executionType == MLOperatorExecutionType::D3D12 && + opKernel->domain != nullptr && + std::string_view(opKernel->domain).empty() && + std::string_view(opKernel->name) == "BitShift" && + opKernel->minimumOperatorSetVersion == 11) + { + // DirectML masks out-of-range shift counts, which does not satisfy the opset-28 semantics. + builder.SinceVersion(11, 27); + } std::string_view name(opKernel->name); if (name == "MemcpyToHost") diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp index b64a5265f56e3..94129987905e5 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp @@ -731,6 +731,134 @@ class DmlOperatorElementwiseMod : public DmlOperator auto fmod = kernelInfo.GetOptionalAttribute(AttrName::Fmod, 0); + const DML_TENSOR_DATA_TYPE dataType = m_inputTensorDescs[0].GetDmlDataType(); + if (!fmod && (dataType == DML_TENSOR_DATA_TYPE_FLOAT16 || dataType == DML_TENSOR_DATA_TYPE_FLOAT32)) + { + // Intermediate tensors must be dense, even when the graph inputs have broadcast strides. + TensorDesc valueTensorDesc(dataType, m_outputTensorDescs[0].GetSizes()); + TensorDesc conditionTensorDesc(DML_TENSOR_DATA_TYPE_UINT8, m_outputTensorDescs[0].GetSizes()); + DML_TENSOR_DESC valueDesc = valueTensorDesc.GetDmlDesc(); + DML_TENSOR_DESC conditionDesc = conditionTensorDesc.GetDmlDesc(); + + struct Source + { + uint32_t index; + bool isInput = false; + }; + const Source x{0, true}; + const Source y{1, true}; + std::vector operators; + std::vector inputEdges; + std::vector intermediateEdges; + auto addNode = [&](const auto& desc, std::initializer_list sources) -> Source + { + const uint32_t nodeIndex = gsl::narrow_cast(operators.size()); + using Desc = std::decay_t; + operators.push_back({ApiTraits::OperatorDescTraits::Type, &desc}); + uint32_t inputIndex = 0; + for (const Source& source : sources) + { + if (source.isInput) + { + DML_INPUT_GRAPH_EDGE_DESC edge{}; + edge.GraphInputIndex = source.index; + edge.ToNodeIndex = nodeIndex; + edge.ToNodeInputIndex = inputIndex; + inputEdges.push_back(edge); + } + else + { + DML_INTERMEDIATE_GRAPH_EDGE_DESC edge{}; + edge.FromNodeIndex = source.index; + edge.ToNodeIndex = nodeIndex; + edge.ToNodeInputIndex = inputIndex; + intermediateEdges.push_back(edge); + } + ++inputIndex; + } + return {nodeIndex}; + }; + + DML_FILL_VALUE_CONSTANT_OPERATOR_DESC zeroDesc{}; + zeroDesc.OutputTensor = &valueDesc; + zeroDesc.ValueDataType = dataType; + const Source zero = addNode(zeroDesc, {}); + DML_FILL_VALUE_CONSTANT_OPERATOR_DESC negativeZeroDesc = zeroDesc; + DML_FILL_VALUE_CONSTANT_OPERATOR_DESC infinityDesc = zeroDesc; + if (dataType == DML_TENSOR_DATA_TYPE_FLOAT16) + { + negativeZeroDesc.Value.UInt16 = 0x8000; + infinityDesc.Value.UInt16 = 0x7c00; + } + else + { + negativeZeroDesc.Value.Float32 = -0.0f; + infinityDesc.Value.Float32 = std::numeric_limits::infinity(); + } + const Source negativeZero = addNode(negativeZeroDesc, {}); + const Source infinity = addNode(infinityDesc, {}); + + DML_ELEMENT_WISE_ABS_OPERATOR_DESC absXDesc{&inputDescs[0], &valueDesc}; + DML_ELEMENT_WISE_ABS_OPERATOR_DESC absYDesc{&inputDescs[1], &valueDesc}; + const Source absX = addNode(absXDesc, {x}); + const Source absY = addNode(absYDesc, {y}); + DML_ELEMENT_WISE_LOGICAL_LESS_THAN_OPERATOR_DESC finiteXDesc{&valueDesc, &valueDesc, &conditionDesc}; + const Source finiteX = addNode(finiteXDesc, {absX, infinity}); + DML_ELEMENT_WISE_LOGICAL_EQUALS_OPERATOR_DESC infiniteYDesc{&valueDesc, &valueDesc, &conditionDesc}; + const Source infiniteY = addNode(infiniteYDesc, {absY, infinity}); + DML_ELEMENT_WISE_LOGICAL_AND_OPERATOR_DESC infiniteDivisorDesc{&conditionDesc, &conditionDesc, &conditionDesc}; + const Source infiniteDivisor = addNode(infiniteDivisorDesc, {finiteX, infiniteY}); + + DML_ELEMENT_WISE_LOGICAL_LESS_THAN_OPERATOR_DESC negativeXDesc{&inputDescs[0], &valueDesc, &conditionDesc}; + DML_ELEMENT_WISE_LOGICAL_LESS_THAN_OPERATOR_DESC negativeYDesc{&inputDescs[1], &valueDesc, &conditionDesc}; + const Source negativeX = addNode(negativeXDesc, {x, zero}); + const Source negativeY = addNode(negativeYDesc, {y, zero}); + DML_ELEMENT_WISE_LOGICAL_EQUALS_OPERATOR_DESC sameSignDesc{&conditionDesc, &conditionDesc, &conditionDesc}; + const Source sameSign = addNode(sameSignDesc, {negativeX, negativeY}); + DML_ELEMENT_WISE_MODULUS_FLOOR_OPERATOR_DESC floorDesc{&inputDescs[0], &inputDescs[1], &valueDesc}; + const Source floor = addNode(floorDesc, {x, y}); + DML_ELEMENT_WISE_IF_OPERATOR_DESC infiniteResultDesc{&conditionDesc, &inputDescs[0], &inputDescs[1], &valueDesc}; + const Source infiniteResult = addNode(infiniteResultDesc, {sameSign, x, y}); + DML_ELEMENT_WISE_IF_OPERATOR_DESC resultDesc{&conditionDesc, &valueDesc, &valueDesc, &valueDesc}; + const Source result = addNode(resultDesc, {infiniteDivisor, infiniteResult, floor}); + + DML_ELEMENT_WISE_LOGICAL_EQUALS_OPERATOR_DESC zeroResultDesc{&valueDesc, &valueDesc, &conditionDesc}; + const Source zeroResult = addNode(zeroResultDesc, {result, zero}); + DML_ELEMENT_WISE_LOGICAL_EQUALS_OPERATOR_DESC zeroXDesc{&inputDescs[0], &valueDesc, &conditionDesc}; + const Source zeroX = addNode(zeroXDesc, {x, zero}); + DML_ELEMENT_WISE_LOGICAL_OR_OPERATOR_DESC anyZeroDesc{&conditionDesc, &conditionDesc, &conditionDesc}; + const Source anyZero = addNode(anyZeroDesc, {zeroX, zeroResult}); + DML_ELEMENT_WISE_LOGICAL_GREATER_THAN_OPERATOR_DESC validYDesc{&valueDesc, &valueDesc, &conditionDesc}; + const Source validY = addNode(validYDesc, {absY, zero}); + DML_ELEMENT_WISE_LOGICAL_AND_OPERATOR_DESC validZeroDesc{&conditionDesc, &conditionDesc, &conditionDesc}; + const Source validZero = addNode(validZeroDesc, {anyZero, validY}); + // Select literal signed zeros rather than relying on arithmetic to preserve the sign bit. + DML_ELEMENT_WISE_IF_OPERATOR_DESC signedZeroDesc{&conditionDesc, &valueDesc, &valueDesc, &valueDesc}; + const Source signedZero = addNode(signedZeroDesc, {negativeY, negativeZero, zero}); + DML_ELEMENT_WISE_IF_OPERATOR_DESC outputDesc{&conditionDesc, &valueDesc, &valueDesc, &outputDescs[0]}; + const Source output = addNode(outputDesc, {validZero, signedZero, result}); + + std::vector nodes; + nodes.reserve(operators.size()); + for (const DML_OPERATOR_DESC& op : operators) + { + nodes.push_back(&op); + } + DML_OUTPUT_GRAPH_EDGE_DESC outputEdge{}; + outputEdge.FromNodeIndex = output.index; + MLOperatorGraphDesc graphDesc{}; + graphDesc.nodeCount = gsl::narrow_cast(nodes.size()); + graphDesc.nodes = nodes.data(); + graphDesc.inputEdgeCount = gsl::narrow_cast(inputEdges.size()); + graphDesc.inputEdges = inputEdges.data(); + graphDesc.intermediateEdgeCount = gsl::narrow_cast(intermediateEdges.size()); + graphDesc.intermediateEdges = intermediateEdges.data(); + graphDesc.outputEdgeCount = 1; + graphDesc.outputEdges = &outputEdge; + SetDmlOperatorGraphDesc(std::move(graphDesc), kernelInfo); + return; + } + // Note TRUNCATE and FLOOR modulus operator descriptions are identical. static_assert(sizeof(DML_ELEMENT_WISE_MODULUS_TRUNCATE_OPERATOR_DESC) == sizeof(DML_ELEMENT_WISE_MODULUS_FLOOR_OPERATOR_DESC)); DML_ELEMENT_WISE_MODULUS_TRUNCATE_OPERATOR_DESC opDesc = {}; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorSpaceToDepth.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorSpaceToDepth.cpp index 9d1faef3fd2a6..81eff170f1312 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorSpaceToDepth.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorSpaceToDepth.cpp @@ -30,6 +30,12 @@ class DmlOperatorSpaceToDepth : public DmlOperator, public SpaceToDepthHelper } }; +void CALLBACK QuerySpaceToDepth(IMLOperatorSupportQueryContextPrivate* context, bool* isSupported) +{ + MLOperatorAttributes attributes(context); + *isSupported = attributes.GetOptionalAttribute(AttrName::Mode, "DCR") == "DCR"; +} + DML_OP_DEFINE_CREATION_FUNCTION(SpaceToDepth, DmlOperatorSpaceToDepth); } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp index 26f998c7521a2..edfe419028090 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -543,6 +543,7 @@ DML_OP_EXTERN_CREATION_FUNCTION(MatMulNBits); DML_OP_EXTERN_QUERY_FUNCTION(MaxPool); DML_OP_EXTERN_QUERY_FUNCTION(Slice); DML_OP_EXTERN_QUERY_FUNCTION(Resize); +DML_OP_EXTERN_QUERY_FUNCTION(SpaceToDepth); DML_OP_EXTERN_QUERY_FUNCTION(Scatter); DML_OP_EXTERN_QUERY_FUNCTION(EinSum); DML_OP_EXTERN_QUERY_FUNCTION(RecurrentNeuralNetwork); @@ -813,8 +814,8 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO_VER( 19, Pad, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1, 2, 3) /*pads, value, axes*/)}, #endif - {REG_INFO( 7, SpaceToDepth, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO( 13, SpaceToDepth, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO( 7, SpaceToDepth, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QuerySpaceToDepth)}, + {REG_INFO( 13, SpaceToDepth, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QuerySpaceToDepth)}, {REG_INFO( 7, DepthToSpace, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, {REG_INFO( 11, DepthToSpace, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, {REG_INFO( 13, DepthToSpace, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, @@ -1092,6 +1093,7 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO( 13, Cast, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, {REG_INFO( 19, Cast, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, {REG_INFO( 21, Cast, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, + {REG_INFO( 28, Cast, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, {REG_INFO_VER( 15, CastLike, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, {REG_INFO_VER( 19, CastLike, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, {REG_INFO_VER( 21, CastLike, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, @@ -1143,6 +1145,7 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO( 20, IsInf, typeNameListTwo, supportedTypeListIsInf, DmlGraphSupport::Supported)}, {REG_INFO( 10, Mod, typeNameListDefault, supportedTypeListNumericDefault, DmlGraphSupport::Supported)}, {REG_INFO( 13, Mod, typeNameListDefault, supportedTypeListNumericDefault, DmlGraphSupport::Supported)}, + {REG_INFO( 28, Mod, typeNameListDefault, supportedTypeListNumericDefault, DmlGraphSupport::Supported)}, {REG_INFO( 11, BitShift, typeNameListDefault, supportedTypeListUInt8to64, DmlGraphSupport::Supported)}, {REG_INFO( 11, Round, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO( 10, ReverseSequence, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h index b4d402a1d9e77..c78439c9ebff3 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h @@ -457,6 +457,12 @@ namespace OperatorHelper static const int sc_sinceVer_GroupNorm = 21; } + namespace OnnxOperatorSet28 + { + static const int sc_sinceVer_Cast = 28; + static const int sc_sinceVer_Mod = 28; + } + namespace MsftOperatorSet1 { static const int sc_sinceVer_DmlFusedConv = 1; diff --git a/onnxruntime/core/providers/js/js_execution_provider.cc b/onnxruntime/core/providers/js/js_execution_provider.cc index 135aba2e32b59..c3e0072fd0dab 100644 --- a/onnxruntime/core/providers/js/js_execution_provider.cc +++ b/onnxruntime/core/providers/js/js_execution_provider.cc @@ -123,7 +123,8 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomai class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, 12, Cast); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, 18, Cast); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 19, 20, Cast); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 21, Cast); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 21, 27, Cast); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 28, Cast); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 6, 10, Clip); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, 11, Clip); @@ -466,7 +467,8 @@ std::unique_ptr RegisterKernels() { KERNEL_CREATE_INFO_VERSIONED(9, 12, Cast), KERNEL_CREATE_INFO_VERSIONED(13, 18, Cast), KERNEL_CREATE_INFO_VERSIONED(19, 20, Cast), - KERNEL_CREATE_INFO(21, Cast), + KERNEL_CREATE_INFO_VERSIONED(21, 27, Cast), + KERNEL_CREATE_INFO(28, Cast), // activations KERNEL_CREATE_INFO_VERSIONED(6, 10, Clip), diff --git a/onnxruntime/core/providers/js/operators/cast.cc b/onnxruntime/core/providers/js/operators/cast.cc index f499d0627e032..f03a59f0bd922 100644 --- a/onnxruntime/core/providers/js/operators/cast.cc +++ b/onnxruntime/core/providers/js/operators/cast.cc @@ -58,10 +58,19 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .TypeConstraint("T1", CastOpTypeConstraints()) .TypeConstraint("T2", CastOpTypeConstraints()), Cast); +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Cast, + kOnnxDomain, + 21, 27, + kJsExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T1", CastOpTypeConstraints()) + .TypeConstraint("T2", CastOpTypeConstraints()), + Cast); ONNX_OPERATOR_KERNEL_EX( Cast, kOnnxDomain, - 21, + 28, kJsExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T1", CastOpTypeConstraints()) diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc index 8908013a52291..72c5676c20fa0 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc @@ -377,6 +377,11 @@ std::vector toVector(const ONNX_NAMESPACE::int64s& nums) { static bool IsUnsupportedOpMode(const onnxruntime::GraphViewer& graph_viewer, const Node* node) { std::vector input_nodes; const auto& optype = node->OpType(); + // These opset-28 semantics require CPU kernels or ONNX function expansion. + if ((optype == "Mod" || optype == "SpaceToDepth") && node->SinceVersion() >= 28) { + return true; + } + if (optype == "ArgMax" || optype == "ArgMin") { const auto& attributes = node->GetAttributes(); // we do not support select_last_index = 1 for now diff --git a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc index 37c76c9f46846..9be6b9939404d 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc @@ -204,7 +204,11 @@ struct MigraphXEpFactory : OrtEpFactory { const char* ep_name, OrtHardwareDeviceType hw_type, const OrtLogger& default_logger_in) - : ort_api{ort_api_in}, default_logger{default_logger_in}, ep_name{ep_name}, ort_hw_device_type{hw_type} { + : OrtEpFactory{}, + ort_api{ort_api_in}, + default_logger{default_logger_in}, + ep_name{ep_name}, + ort_hw_device_type{hw_type} { ort_version_supported = ORT_API_VERSION; // set to the ORT version we were compiled with GetName = GetNameImpl; GetVendor = GetVendorImpl; diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/base_op_builder.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/base_op_builder.h index a573b73d2bbb5..c52458b7046c5 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/base_op_builder.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/base_op_builder.h @@ -72,7 +72,7 @@ class BaseOpBuilder : public IOpBuilder { const OpSupportCheckParams& params) const; virtual int GetMinSupportedOpSet(const NodeUnit& /* node_unit */) const { return 1; } - virtual int GetMaxSupportedOpSet(const NodeUnit& /* node_unit */) const { return 27; } + virtual int GetMaxSupportedOpSet(const NodeUnit& /* node_unit */) const { return 28; } // Check if this node_unit's type is supported // SingleNode type NodeUnit is supported diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/dequantizelinear_op_builder.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/dequantizelinear_op_builder.cc index ed9062f894660..72544d87413cd 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/dequantizelinear_op_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/dequantizelinear_op_builder.cc @@ -37,6 +37,8 @@ class DequantizeLinearOpBuilder : public BaseOpBuilder { return ANEURALNETWORKS_FEATURE_LEVEL_1; } + int GetMaxSupportedOpSet(const NodeUnit& /* node_unit */) const override { return 27; } + bool HasSupportedInputOutputsImpl( const GraphViewer& graph_viewer, const NodeUnit& node_unit, const OpSupportCheckParams& params) const override { diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/quantizelinear_op_builder.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/quantizelinear_op_builder.cc index eb81f5e3f59ee..9a94f37b68ae6 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/quantizelinear_op_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/quantizelinear_op_builder.cc @@ -37,6 +37,8 @@ class QuantizeLinearOpBuilder : public BaseOpBuilder { return ANEURALNETWORKS_FEATURE_LEVEL_3; } + int GetMaxSupportedOpSet(const NodeUnit& /* node_unit */) const override { return 27; } + bool HasSupportedInputOutputsImpl( const GraphViewer& graph_viewer, const NodeUnit& node_unit, const OpSupportCheckParams& params) const override { diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc index 23a53788585d6..af979ce89d7d8 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc @@ -1042,7 +1042,8 @@ struct NvTensorRtRtxEpFactory : OrtEpFactory { using MemoryInfoUniquePtr = std::unique_ptr>; NvTensorRtRtxEpFactory(const OrtApi& ort_api_in, - const OrtLogger& default_logger_in) : ort_api{ort_api_in}, + const OrtLogger& default_logger_in) : OrtEpFactory{}, + ort_api{ort_api_in}, ep_api{*ort_api_in.GetEpApi()}, default_logger{default_logger_in}, data_transfer_impl{ort_api_in} { diff --git a/onnxruntime/core/providers/openvino/ov_factory.cc b/onnxruntime/core/providers/openvino/ov_factory.cc index 5119c611d3f3d..23a10971a489f 100644 --- a/onnxruntime/core/providers/openvino/ov_factory.cc +++ b/onnxruntime/core/providers/openvino/ov_factory.cc @@ -30,7 +30,8 @@ static void InitCxxApi(const OrtApiBase& ort_api_base) { } OpenVINOEpPluginFactory::OpenVINOEpPluginFactory(ApiPtrs apis, const std::string& ov_metadevice_name, std::shared_ptr core) - : ApiPtrs{apis}, + : OrtEpFactory{}, + ApiPtrs{apis}, ep_name_(ov_metadevice_name.empty() ? provider_name_ : std::string(provider_name_) + "." + ov_metadevice_name), device_type_(ov_metadevice_name), ov_core_(std::move(core)) { diff --git a/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc b/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc index 7aae81b987c0d..24c5388fe25c7 100644 --- a/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc +++ b/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc @@ -727,6 +727,10 @@ bool DataOps::dimension_unsupported(const Node* node) { bool DataOps::node_is_supported(const NodeIndex node_idx, bool& has_external_weights_) { const auto& node = graph_viewer_.GetNode(node_idx); const auto& optype = node->OpType(); + // These opset-28 semantics require CPU kernels or ONNX function expansion. + if ((optype == "Mod" || optype == "BitShift" || optype == "SpaceToDepth") && node->SinceVersion() >= 28) { + return false; + } #ifndef NDEBUG if (openvino_ep::backend_utils::IsDebugEnabled()) { diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/mod_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/mod_op_builder.cc index 1d61ac9b9599d..5db8e638e83c1 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/mod_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/mod_op_builder.cc @@ -47,12 +47,21 @@ Status ModOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper, bool do_op_validation) const { ORT_UNUSED_PARAMETER(do_op_validation); NodeAttrHelper node_helper(node_unit); - int64_t fmod = node_helper.Get("fmod", static_cast(0)); // 0=integer mod. 1=float mod. + int64_t fmod = node_helper.Get("fmod", static_cast(0)); if (1 == fmod) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "QNN Mod Op only support fmod == 0 for now."); } const auto& inputs = node_unit.Inputs(); + int32_t input_type = 0; + ORT_RETURN_IF_ERROR(utils::GetOnnxTensorElemDataType(inputs[0].node_arg, input_type)); + // a - b * floor(a / b) does not preserve the ONNX floating-point edge-case semantics. + ORT_RETURN_IF(input_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT || + input_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 || + input_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE || + input_type == ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16, + "QNN EP does not support floating-point Mod with fmod == 0."); + const auto input_count = GetInputCountQnnRequired(node_unit); for (size_t input_i = 0; input_i < input_count; ++input_i) { ORT_RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, inputs[input_i], logger, input_names)); diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/simple_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/simple_op_builder.cc index 5537a388f598c..9d59ac1ea5417 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/simple_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/simple_op_builder.cc @@ -48,6 +48,12 @@ Status SimpleOpBuilder::ExplicitOpCheck(QnnModelWrapper& qnn_model_wrapper, const NodeUnit& node_unit) const { const std::string& op_type = node_unit.OpType(); + if (op_type == "SpaceToDepth") { + NodeAttrHelper node_helper(node_unit); + ORT_RETURN_IF_NOT(node_helper.Get("mode", "DCR") == "DCR", + "QNN EP only supports SpaceToDepth with DCR mode."); + } + if (op_type == "GridSample") { NodeAttrHelper node_helper(node_unit); std::string mode = node_helper.Get("mode", "linear"); diff --git a/onnxruntime/core/providers/qnn/qnn_provider_factory.cc b/onnxruntime/core/providers/qnn/qnn_provider_factory.cc index a51a1dc785340..96dfffdc20ede 100644 --- a/onnxruntime/core/providers/qnn/qnn_provider_factory.cc +++ b/onnxruntime/core/providers/qnn/qnn_provider_factory.cc @@ -208,7 +208,8 @@ struct QnnEpFactory : OrtEpFactory { QnnEpFactory(const OrtApi& ort_api_in, const OrtLogger& default_logger_in, const char* ep_name) - : ort_api{ort_api_in}, + : OrtEpFactory{}, + ort_api{ort_api_in}, default_logger{default_logger_in}, ep_name{ep_name} { ort_version_supported = ORT_API_VERSION; diff --git a/onnxruntime/core/providers/vitisai/vitisai_provider_factory.cc b/onnxruntime/core/providers/vitisai/vitisai_provider_factory.cc index 89daef43baeb7..0c35ed1c393b4 100644 --- a/onnxruntime/core/providers/vitisai/vitisai_provider_factory.cc +++ b/onnxruntime/core/providers/vitisai/vitisai_provider_factory.cc @@ -105,7 +105,7 @@ struct VitisAI_Provider : Provider { struct VitisAIEpFactory : OrtEpFactory { VitisAIEpFactory(const OrtApi& ort_api_in, const OrtLogger& default_logger_in) - : ort_api{ort_api_in}, default_logger{default_logger_in} { + : OrtEpFactory{}, ort_api{ort_api_in}, default_logger{default_logger_in} { ort_version_supported = ORT_API_VERSION; GetName = GetNameImpl; GetVendor = GetVendorImpl; diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.h index 488581dbd0e97..8a9f342871f3d 100644 --- a/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.h +++ b/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.h @@ -49,7 +49,7 @@ class BaseOpBuilder : public IOpBuilder { virtual bool IsQuantizedOp(const NodeUnit& /* node_unit */) const { return false; } virtual int GetMinSupportedOpSet(const NodeUnit& /* node_unit */) const { return 1; } - virtual int GetMaxSupportedOpSet(const NodeUnit& /* node_unit */) const { return 27; } + virtual int GetMaxSupportedOpSet(const NodeUnit& /* node_unit */) const { return 28; } virtual bool HasSupportedInputOutputsImpl( const InitializedTensorSet& initializers, const NodeUnit& node_unit) const; diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.cc b/onnxruntime/core/providers/webgpu/buffer_manager.cc index 247b2ee922ea2..220ec4942f42e 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.cc +++ b/onnxruntime/core/providers/webgpu/buffer_manager.cc @@ -545,10 +545,7 @@ void BufferManager::MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGP "Source and destination buffers must have enough space for the copy operation. src_size=", src_size, ", dst_size=", dst_size, ", copy_size=", copy_size, "."); - ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording)); - auto& command_encoder = context_.GetCommandEncoder(recording); - context_.EndComputePass(recording); - command_encoder.CopyBufferToBuffer(src, 0, dst, 0, copy_size); + context_.CopyBuffer(src, dst, copy_size, recording); } WGPUBuffer BufferManager::Create(CommandRecordingState& recording, size_t size, wgpu::BufferUsage usage, diff --git a/onnxruntime/core/providers/webgpu/compute_context.h b/onnxruntime/core/providers/webgpu/compute_context.h index 03f5ff9ac5869..7f6e10a166e2d 100644 --- a/onnxruntime/core/providers/webgpu/compute_context.h +++ b/onnxruntime/core/providers/webgpu/compute_context.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include @@ -77,6 +78,9 @@ class ComputeContextBase { inline bool HasFeature(wgpu::FeatureName feature) const { return webgpu_context_.DeviceHasFeature(feature); } + inline bool HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName feature) const { + return webgpu_context_.Instance().HasWGSLLanguageFeature(feature); + } inline const wgpu::AdapterPropertiesSubgroupMatrixConfigs& SubgroupMatrixConfigs() const { return webgpu_context_.SubgroupMatrixConfigs(); } @@ -123,18 +127,20 @@ class ComputeContextBase { return ep_.EnableMatmulFp32Accumulation(); } + inline std::optional ForcedMatMulAlgorithm() const { + return ep_.ForcedMatMulAlgorithm(); + } + // // Get the logger. // + inline const auto& Logger() const { #if defined(ORT_USE_EP_API_ADAPTERS) - inline const onnxruntime::ep::adapter::Logger& Logger() const { return ep_.GetEpLogger(); - } #else - inline const logging::Logger& Logger() const { return *ep_.GetLogger(); - } #endif + } // // Run a compute shader program. @@ -143,6 +149,14 @@ class ComputeContextBase { return webgpu_context_.Run(*this, program); } + inline Status FlushAndWaitChecked() { + return webgpu_context_.FlushAndWaitChecked(ep_.BufferManager(), ep_.Recording()); + } + + inline Status RunAndWaitChecked(const std::function& enqueue) { + return webgpu_context_.RunAndWaitChecked(ep_.BufferManager(), ep_.Recording(), enqueue); + } + protected: WebGpuContext& webgpu_context_; const WebGpuExecutionProvider& ep_; diff --git a/onnxruntime/core/providers/webgpu/ep/README.md b/onnxruntime/core/providers/webgpu/ep/README.md index 641dd4dc77a1d..e2d52a67e9dbe 100644 --- a/onnxruntime/core/providers/webgpu/ep/README.md +++ b/onnxruntime/core/providers/webgpu/ep/README.md @@ -151,6 +151,17 @@ error for checked completion; ordinary Flush/Run behavior is unchanged. For externally supplied devices, earlier errors captured by the caller's scopes or callbacks remain the caller's responsibility. +`WebGpuContext::RunAndWaitChecked(buffer_manager, recording, enqueue)` extends +that boundary to the supplied Status-returning callback. All three scopes are +active before the callback creates, validates or encodes programs and remain +active through flush/wait. The callback may enqueue multiple programs on the +owning recording; it must stay on the calling thread. An enqueue failure or +exception is reported alongside completion errors, and already-enqueued work +is drained before all scopes are popped. Scope cleanup also runs on exception +unwinding. Caller-owned scopes and external-device callbacks are not replaced. +The ordinary `FlushAndWaitChecked` entry point uses the same implementation +with an empty enqueue callback and retains its existing boundary. + The helper polls `Device::GetLostFuture` after waiting, rather than treating a successful queue callback as proof of device health. This also detects loss of an external device without replacing its callbacks. A failed wait or @@ -162,6 +173,7 @@ this is not a new stream or cross-provider API. Focused context tests cover empty, pending and submitted work, deferred dispatch, validation at Flush, scope cleanup after a failed wait, +enqueue failure/exception cleanup, early shader errors on an external device, validation/OOM injection and external device loss. Dawn's existing InjectError API accepts validation and OOM only. Deterministic internal-error and non-success work-done callback injection is not available here; those branches @@ -202,6 +214,75 @@ are also checked before that external readback, including after a cancelled Run. Concurrent profiling, cross-device transfer, and successful arbitrary foreign stream overrides are not established by these tests. Performance must be measured separately. +### GatedDeltaNet state replay + +`com.microsoft::GatedDeltaNetStateReplay` version 1 is a WebGPU-only, FP32 +operator for one request and one recurrent-state descriptor. Inputs 0-2 are +full, rank-1 source-state, compact-capsule and inactive destination backing +tensors. Input 3 is CPU `int64[11]`, in this order: +`source_offset, destination_offset, decay_offset, key_offset, delta_offset, +Hv, Dv, Dk, Hk, capacity, kept_count`. Offsets count FP32 elements. +Dimensions and capacity must be positive, and `1 <= kept_count <= capacity`. +The caller must provide valid captured transitions for the entire kept prefix. + +The `source_state` selected by `source_offset` must be the state snapshot from +immediately before the captured GatedDeltaNet transitions began. Using the +producer's final state or another same-shaped snapshot is invalid because +captured delta is state-dependent. Capsule geometry (`Hv, Dv, Dk, Hk`) and +capacity must match the producer that generated the capsule. `capacity` is +the original capsule capture/layout capacity. `kept_count` only selects the +leading transitions to replay; `kept_count` must NOT be substituted for +`capacity` when deriving the decay/key/delta section offsets. + +The output must be explicitly I/O-bound to the same allocation as the +destination input; `.MayInplace` alone does not guarantee aliasing of a graph +input. Writable destination backing must not alias source or capsule backing. +Only the selected destination row is written. Source, capsule and neighboring +destination slots are preserved, including unused capsule tails. + +State layout is `[Hv,Dv,Dk]`; capsule views are `[capacity,Hv]` decay, +`[capacity,Hk,Dk]` key and `[capacity,Hv,Dv]` delta. Each invocation owns one +state element, maps `hk = floor(h*Hk/Hv)`, and applies the captured factors in +token order: `s *= decay[t,h]; s += key[t,hk,i] * delta[t,h,v]`. +Workgroups contain 128 invocations; large dispatches use a flattened 2D grid. +`BufferView` accessors add element-offset uniforms while binding the backing +allocation (or its aligned framework segments), so a logical capsule row does +not require a storage-binding-aligned offset. Source and capsule may share +read-only backing. Storage binding counts use `ProgramBase::InputBufferOwner` and +`OutputBufferOwner`, counting each owner's full backing segments once, plus +one uniform binding. Distinct Tensor wrappers are not deduplicated merely +because their raw buffer handles match. Extents, uint32 indexing and device +binding limits are checked before dispatch. + +The kernel calls `ComputeContextBase::RunAndWaitChecked()` on the owning EP +recording, with its enqueue helper inside the checked callback. This captures +errors from shader/pipeline creation and binding/encoding as well as the +subsequent flush/wait, including on externally supplied devices. The enqueue +helper remains separate so future internal callers can enqueue multiple +descriptors within one checked region and check completion once. A failing +completion returns failure rather than publishing success; destination contents +on failure are unspecified and must not be committed. As with the context helper, +errors from before the checked region remain the external caller's responsibility. + +This milestone has no graph-capture support, convolution replay, batching, +or external session integration. Tests cover all kept prefixes at capacities +1, 4, 7 and 8, non-divisible head mappings, Qwen's 230,720-byte capsule-row +offset (64 modulo 256), segmented and shared source/capsule backing, poisoned +tails, preservation, alias requirements and invalid metadata. GPU/CPU replay comparisons use +absolute tolerance `1e-5` plus relative tolerance `2e-5`, not bitwise equality. +Artificially small segmentation limits use the existing native-only test factory. +`CapturedStateReplayMatchesIndependentPrefixes` uses the actual WebGPU producer's +capsule, replays the original incoming state and compares with independent prefix +runs with capture disabled. `RunCompletesBeforeReadback` checks the owning recording +and polls queue completion with timeout zero immediately after Session::Run, before +any tensor readback. The completion observation uses native Dawn context access and +is not compiled for plugin EP, WASM or external-Dawn builds. +`CheckedRegionReportsRunProgramBindingError` uses a validation-enabled external +device without session/kernel scopes. A source buffer lacking Storage usage +must report the original CreateBindGroup error through the replay call, not just +the subsequent invalid-encoder error, while leaving the caller's outer scope +empty. A valid replay before and after the rejected run verifies recovery. + ### Missing parts This section describes what is missing. diff --git a/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc b/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc index 3cf422e9243eb..8e0d20031571e 100644 --- a/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc +++ b/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc @@ -273,12 +273,13 @@ Status RunBinaryProgram(ComputeContext& context, program.AddInput({rhs_tensor, ProgramTensorMetadataDependency::Type}); } // Mode Vectorize broadcast - // cache hint: "V{a_rank};{b_rank};{output_rank}" + // Boolx4 storage does not distinguish vector reads from lane broadcasts. program .AddIndices(std::move(reshaped_output_shape)) .AddIndices(std::move(reshaped_lhs_shape)) .AddIndices(std::move(reshaped_rhs_shape)) - .CacheHint("V"); + .CacheHint("V", is_lhs_bool && (shared_dimension_divisible_by_4 || a_last_dim_divisible_by_4), + is_rhs_bool && (shared_dimension_divisible_by_4 || b_last_dim_divisible_by_4)); } else { // Mode Broadcast // cache hint: "B" @@ -467,7 +468,7 @@ std::string GetPowImpl(int lhs_element_type, int /* rhs_element_type */) { return SS_GET(s); } -WEBGPU_BINARY_IMPL(Pow, "pow_v(a, b)", GetPowImpl) +WEBGPU_BINARY_IMPL(Pow, "pow_v(vec4(a), vec4(b))", GetPowImpl) WEBGPU_BINARY_IMPL(PRelu, "select(b * a, a, a >= vec4(0))") WEBGPU_BINARY_IMPL(Equal, "vec4(vec4(a) == vec4(b))") WEBGPU_BINARY_IMPL(Greater, "vec4(vec4(a) > vec4(b))") diff --git a/onnxruntime/core/providers/webgpu/math/matmul.cc b/onnxruntime/core/providers/webgpu/math/matmul.cc index cd23d7224f1a9..d5eda34a9a67a 100644 --- a/onnxruntime/core/providers/webgpu/math/matmul.cc +++ b/onnxruntime/core/providers/webgpu/math/matmul.cc @@ -12,19 +12,25 @@ #include "core/providers/webgpu/nn/fuse_utils.h" #include "core/providers/webgpu/data_transfer.h" #include "core/providers/webgpu/tensor/cast.h" +#include "core/providers/webgpu/vendor/intel/math/matmul_execution_planner.h" #include "core/providers/webgpu/vendor/intel/math/matmul.h" #include "core/providers/webgpu/webgpu_utils.h" namespace onnxruntime { namespace webgpu { -std::unique_ptr CreateSubgroupMatrixMatMulImpl(const ComputeContextBase& context); +std::unique_ptr CreateSubgroupMatrixMatMulImpl(const ComputeContextBase& context); -MatMulOptImpl* MatMulOptImplCache::GetOrCreate(const ComputeContextBase& context) { - std::call_once(subgroup_impl_init_flag_, [&]() { - subgroup_impl_ = CreateSubgroupMatrixMatMulImpl(context); +void MatMulComputeDispatcher::Initialize(const ComputeContextBase& context) { + std::call_once(init_flag_, [&]() { + subgroup_matrix_impl_ = CreateSubgroupMatrixMatMulImpl(context); + if (context.AdapterInfo().vendor == std::string_view{"intel"}) { + planner_ = std::make_unique( + context.GetSplitKConfig()); + } else { + planner_ = std::make_unique(context.GetSplitKConfig()); + } }); - return subgroup_impl_.get(); } ONNX_OPERATOR_VERSIONED_KERNEL_EX( @@ -45,30 +51,61 @@ ONNX_OPERATOR_KERNEL_EX( .TypeConstraint("T", WebGpuSupportedFloatTypes()), MatMul); +Status MatMulGemvProgram::GenerateShaderCode(ShaderHelper& shader) const { + const auto& a = shader.AddInput("a", ShaderUsage::None); + const auto& b = shader.AddInput("b", ShaderUsage::None); + const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias); + return WGSL_TEMPLATE_APPLY(shader, "math/matmul_gemv.wgsl.template", + WGSL_TEMPLATE_VARIABLE(a, a), + WGSL_TEMPLATE_VARIABLE(b, b), + WGSL_TEMPLATE_VARIABLE(output, output)); +} + namespace { -class MatMulGemvProgram final : public Program { +class MatMulZeroKProgram final : public Program { public: - MatMulGemvProgram() : Program{"MatMulGemv"} {} - - static constexpr uint32_t kWorkgroupSizeX = 1; - static constexpr uint32_t kWorkgroupSizeY = 128; - static_assert(kWorkgroupSizeY != 0 && (kWorkgroupSizeY & (kWorkgroupSizeY - 1)) == 0, - "MatMulGemvProgram requires kWorkgroupSizeY to be a non-zero power of two " - "because the WGSL reduction loop halves workgroup_size_y each iteration."); + MatMulZeroKProgram(const Activation& activation, bool has_bias, + bool is_channels_last, bool bias_is_scalar) + : Program{"MatMulZeroK"}, + activation_(activation), + has_bias_(has_bias), + is_channels_last_(is_channels_last), + bias_is_scalar_(bias_is_scalar) {} Status GenerateShaderCode(ShaderHelper& shader) const override { - const auto& a = shader.AddInput("a", ShaderUsage::None); - const auto& b = shader.AddInput("b", ShaderUsage::None); - const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias); - return WGSL_TEMPLATE_APPLY(shader, "math/matmul_gemv.wgsl.template", - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output)); + const auto& output = shader.AddOutput( + "output", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + shader.AdditionalImplementation() + << GetActivationDeclaration(activation_, "output_value_t", "output_element_t"); + if (has_bias_) { + shader.AddInput("bias", ShaderUsage::UseValueTypeAlias); + } + shader.MainFunctionBody() + << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size") + << "var value = output_value_t(0);\n"; + if (has_bias_) { + const char* index = bias_is_scalar_ ? "0" : is_channels_last_ ? "global_idx % uniforms.N" + : "(global_idx / uniforms.N) % uniforms.M"; + shader.MainFunctionBody() << "value = output_value_t(bias[" << index << "]);\n"; + } + shader.MainFunctionBody() + << GetActivationSnippet(activation_, "output_value_t", "output_element_t") + << output.SetByOffset("global_idx", "value"); + return Status::OK(); } - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"K", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}); + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"output_size", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + WEBGPU_PROGRAM_ACTIVATION_UNIFORM_VARIABLES); + + private: + const Activation activation_; + const bool has_bias_; + const bool is_channels_last_; + const bool bias_is_scalar_; }; } // namespace @@ -90,13 +127,17 @@ static std::string CalcResult(int64_t components, int64_t a_components, int64_t } Status MatMulNaiveProgram::GenerateShaderCode(ShaderHelper& shader) const { - const auto& a = shader.AddInput("a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | - ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); - const auto& b = shader.AddInput("b", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | - ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); - - const int a_components = a.NumComponents(); - const int components = b.NumComponents(); // components of N + const ShaderVariableHelper* a = nullptr; + const ShaderVariableHelper* b = nullptr; + int a_components = 1; + if (!is_zero_k_) { + a = &shader.AddInput("a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | + ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + b = &shader.AddInput("b", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | + ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + a_components = a->NumComponents(); + } + const int components = NumberOfComponents(Outputs()[0].var_type); // components of N std::string process_bias; if (has_bias_) { @@ -110,32 +151,34 @@ Status MatMulNaiveProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); shader.AdditionalImplementation() << GetActivationDeclaration(activation_, "output_value_t", "output_element_t"); - const auto& batch_dims = shader.AddIndices("batch_dims"); shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size") << "let col = (global_idx % (uniforms.N / " << components << ")) * " << components << ";\n" << "var index1 = global_idx / (uniforms.N / " << components << ");\n" << "let stride1 = uniforms.M / " << output_number_ << ";\n" << "let row = (index1 % stride1) * " << output_number_ << ";\n" - << "let batch = index1 / stride1;\n"; - if (output_rank_ != 2) { - shader.MainFunctionBody() << "let batch_indices = " << batch_dims.OffsetToIndices("batch") << ";\n"; + << "let batch = index1 / stride1;\n" + << "var values: array;\n"; + if (!is_zero_k_) { + const auto& batch_dims = shader.AddIndices("batch_dims"); + if (output_rank_ != 2) { + shader.MainFunctionBody() << "let batch_indices = " << batch_dims.OffsetToIndices("batch") << ";\n"; + } + shader.MainFunctionBody() << "var a_indices: a_indices_t;\n" + << ConvertOutputBatchIndicesToInputBatchIndices("a", *a, a->Rank() - 2, batch_dims.Rank(), "batch_indices") + << a->IndicesSet("a_indices", a->Rank() - 2, 0) << "\n" + << a->IndicesSet("a_indices", a->Rank() - 1, 0) << "\n" + << "let a_offset = " << a->IndicesToOffset("a_indices") << "*" << a_components << ";\n" + << "var b_indices: b_indices_t;\n" + << ConvertOutputBatchIndicesToInputBatchIndices("b", *b, b->Rank() - 2, batch_dims.Rank(), "batch_indices") + << b->IndicesSet("b_indices", b->Rank() - 2, 0) << "\n" + << b->IndicesSet("b_indices", b->Rank() - 1, 0) << "\n" + << "let b_offset = " << b->IndicesToOffset("b_indices") << " * " << components << ";\n" + << "for (var k: u32 = 0u; k < uniforms.K; k = k + " << a_components << ") {\n" + << CalcResult(components, a_components, output_number_) << "\n" + << "}\n"; } - shader.MainFunctionBody() << "var a_indices: a_indices_t;\n" - << ConvertOutputBatchIndicesToInputBatchIndices("a", a, a.Rank() - 2, batch_dims.Rank(), "batch_indices") - << a.IndicesSet("a_indices", a.Rank() - 2, 0) << "\n" - << a.IndicesSet("a_indices", a.Rank() - 1, 0) << "\n" - << "let a_offset = " << a.IndicesToOffset("a_indices") << "*" << a_components << ";\n" - << "var b_indices: b_indices_t;\n" - << ConvertOutputBatchIndicesToInputBatchIndices("b", b, b.Rank() - 2, batch_dims.Rank(), "batch_indices") - << b.IndicesSet("b_indices", b.Rank() - 2, 0) << "\n" - << b.IndicesSet("b_indices", b.Rank() - 1, 0) << "\n" - << "let b_offset = " << b.IndicesToOffset("b_indices") << " * " << components << ";\n" - << "var values: array;\n" - << "for (var k: u32 = 0u; k < uniforms.K; k = k + " << a_components << ") {\n" - << CalcResult(components, a_components, output_number_) << "\n" - << "}\n" - << "for (var i = 0u; i < " << output_number_ << "u; i++) {\n" + shader.MainFunctionBody() << "for (var i = 0u; i < " << output_number_ << "u; i++) {\n" << " var value = values[i];\n" << process_bias << "\n" << apply_activation << "\n" @@ -181,236 +224,380 @@ Status MatMul::ComputeInternal(ComputeContext& context) const { inputs[1] = &promoted_b; } - return ComputeMatMul(&context, Activation(), inputs, output_tensor, - /*is_channels_last=*/true, compute_cache_, b_is_constant_); + return compute_dispatcher_.Compute(context, Activation(), inputs, output_tensor, + /*is_channels_last=*/true, b_is_constant_); } -Status ComputeMatMul(ComputeContext* context, - const Activation& activation, std::vector& inputs, Tensor* output_tensor, bool is_channels_last, - MatMulOptImplCache& cache, - bool b_is_constant) { +static Status ApplyMatMulNaive(ComputeContext& context, + const Activation& activation, + const std::vector& inputs, + Tensor* output_tensor, + bool is_channels_last, + const MatMulComputeHelper& helper) { const auto* a = inputs[0]; const auto* b = inputs[1]; - bool has_bias = inputs.size() > 2; - const TensorShape& logical_a_shape = a->Shape(); - const TensorShape& logical_b_shape = b->Shape(); - ORT_RETURN_IF_NOT(logical_a_shape.NumDimensions() >= 2 && logical_b_shape.NumDimensions() >= 2, - "ComputeMatMul expects matrix or batched-matrix inputs."); - - MatMulComputeHelper helper; - ORT_THROW_IF_ERROR(helper.Compute(logical_a_shape, logical_b_shape)); - - const bool requires_fp32_accumulation = context->EnableMatmulFp32Accumulation() && a->IsDataType(); - MatMulOptImpl* subgroup_impl = requires_fp32_accumulation ? nullptr : cache.GetOrCreate(*context); - if (subgroup_impl != nullptr) { - bool handled = false; - ORT_RETURN_IF_ERROR(subgroup_impl->Compute( - *context, inputs, output_tensor, - activation, is_channels_last, b_is_constant, handled)); - if (handled) { - return Status::OK(); - } - } - if (helper.N() < 8 && helper.K() < 8) { - const uint32_t m = narrow(helper.M()); - const uint32_t n = narrow(helper.N()); - const uint32_t k = narrow(helper.K()); - const int components = GetMaxComponents(n); - const int a_components = GetMaxComponents(k); - const int64_t output_number = GetMaxComponents(m); - const TensorShape& logical_output_shape = helper.OutputShape(); - const size_t output_rank = logical_output_shape.NumDimensions(); - const TensorShape outer_dims = - output_rank > 2 ? logical_output_shape.Slice(0, output_rank - 2) : TensorShape({}); - const int64_t output_rows = logical_a_shape[logical_a_shape.NumDimensions() - 2]; - const TensorShape output_program_shape{ - outer_dims.Size(), output_rows, n / components}; - const uint32_t output_size = - narrow(logical_output_shape.Size() / components / output_number); - - MatMulNaiveProgram program{activation, output_rank, output_number, has_bias, is_channels_last}; + const bool has_bias = inputs.size() > 2; + const uint32_t m = narrow(helper.M()); + const uint32_t n = narrow(helper.N()); + const uint32_t k = narrow(helper.K()); + const bool is_zero_k = k == 0; + const int components = GetMaxComponents(n); + const int a_components = is_zero_k ? 1 : GetMaxComponents(k); + const int64_t output_number = GetMaxComponents(m); + const TensorShape& logical_output_shape = helper.OutputShape(); + const size_t output_rank = logical_output_shape.NumDimensions(); + const TensorShape outer_dims = + output_rank > 2 ? logical_output_shape.Slice(0, output_rank - 2) : TensorShape({}); + const int64_t output_rows = a->Shape()[a->Shape().NumDimensions() - 2]; + const TensorShape output_program_shape{ + outer_dims.Size(), output_rows, n / components}; + const uint32_t output_size = + narrow(logical_output_shape.Size() / components / output_number); + + MatMulNaiveProgram program{activation, output_rank, output_number, has_bias, + is_channels_last, is_zero_k}; + program + .CacheHint(activation.CacheKey(), std::to_string(components), + std::to_string(a_components), std::to_string(output_number), + std::to_string(is_channels_last), std::to_string(is_zero_k)); + if (!is_zero_k) { program - .CacheHint(activation.CacheKey(), std::to_string(components), - std::to_string(a_components), std::to_string(output_number), - std::to_string(is_channels_last)) .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, a_components}, - {b, ProgramTensorMetadataDependency::TypeAndRank, components}}); - if (has_bias) { - const int bias_components = is_channels_last ? components : 1; - program.AddInput({inputs[2], ProgramTensorMetadataDependency::Rank, bias_components}); - } - program - .AddOutputs({{output_tensor, ProgramTensorMetadataDependency::None, - output_program_shape, components}}) - .SetDispatchGroupSize(CeilDiv(output_size, 64u)) - .AddIndices(outer_dims) - .AddUniformVariables({{output_size}, {m}, {n}, {k}}); - AppendActivationUniformsData(activation, program); - return context->RunProgram(program); + {b, ProgramTensorMetadataDependency::TypeAndRank, components}}) + .AddIndices(outer_dims); } - - if (intel::CanApplyMatMulIntel(*context, helper.M(), helper.N(), helper.K())) { - return intel::ApplyMatMulIntel(*context, activation, inputs, output_tensor, is_channels_last); + if (has_bias) { + const int bias_components = is_channels_last ? components : 1; + program.AddInput({inputs[2], ProgramTensorMetadataDependency::Rank, bias_components}); } + program + .AddOutputs({{output_tensor, ProgramTensorMetadataDependency::None, + output_program_shape, components}}) + .SetDispatchGroupSize(CeilDiv(output_size, 64u)) + .AddUniformVariables({{output_size}, {m}, {n}, {k}}); + AppendActivationUniformsData(activation, program); + return context.RunProgram(program); +} - TensorShape a_shape = logical_a_shape; - TensorShape b_shape = logical_b_shape; +static Status ApplyMatMulGemv(ComputeContext& context, + const std::vector& inputs, + Tensor* output_tensor, + const MatMulComputeHelper& helper) { + const auto* a = inputs[0]; + const auto* b = inputs[1]; + const uint32_t m = narrow(a->Shape().SizeToDimension(a->Shape().NumDimensions() - 1)); + const uint32_t n = narrow(helper.N()); + const uint32_t k = narrow(helper.K()); + const uint32_t n_vec_count = n / 4; + + MatMulGemvProgram program; + program.AddInputs({{a, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, + {b, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten, 4}) + .AddUniformVariables({{k}, {n}, {m}}) + .SetWorkgroupSize(MatMulGemvProgram::kWorkgroupSizeX, MatMulGemvProgram::kWorkgroupSizeY) + .SetDispatchGroupSize(CeilDiv(n_vec_count, MatMulGemvProgram::kWorkgroupSizeX) * m); + return context.RunProgram(program); +} + +static Status ApplyMatMulPacked(ComputeContext& context, + const Activation& activation, + const std::vector& inputs, + Tensor* output_tensor, + bool is_channels_last, + const MatMulComputeHelper& helper, + const MatMulPackedConfiguration& configuration, + bool use_split_k) { + const auto* a = inputs[0]; + const auto* b = inputs[1]; + const bool has_bias = inputs.size() > 2; + const bool requires_fp32_accumulation = + context.EnableMatmulFp32Accumulation() && a->IsDataType(); + TensorShape a_shape = a->Shape(); + TensorShape b_shape = b->Shape(); TensorShape output_shape = helper.OutputShape(); - const int64_t batchA = + const int64_t batch_a = a_shape.NumDimensions() > 2 ? a_shape.SizeToDimension(a_shape.NumDimensions() - 2) : 1; - const int64_t batchB = + const int64_t batch_b = b_shape.NumDimensions() > 2 ? b_shape.SizeToDimension(b_shape.NumDimensions() - 2) : 1; - // The generic path benefits from folding A's batch dimensions into M when B - // is shared. The subgroup and Intel paths derive their own dispatch shapes - // directly from the tensor views and have already declined above. - if (batchA != 1 && batchB == 1) { - const int64_t batchAndM = a_shape.SizeToDimension(a_shape.NumDimensions() - 1); - a_shape = TensorShape({batchAndM, helper.K()}); + if (batch_a != 1 && batch_b == 1) { + const int64_t batch_and_m = a_shape.SizeToDimension(a_shape.NumDimensions() - 1); + a_shape = TensorShape({batch_and_m, helper.K()}); b_shape = TensorShape({helper.K(), helper.N()}); - output_shape = TensorShape({batchAndM, helper.N()}); + output_shape = TensorShape({batch_and_m, helper.N()}); } - // helpful dimension variables - TensorShape outer_dims_a = a_shape.NumDimensions() > 2 - ? a_shape.Slice(0, a_shape.NumDimensions() - 2) - : TensorShape({}); - - TensorShape outer_dims_b = b_shape.NumDimensions() > 2 - ? b_shape.Slice(0, b_shape.NumDimensions() - 2) - : TensorShape({}); - - TensorShape outer_dims = output_shape.NumDimensions() > 2 - ? output_shape.Slice(0, output_shape.NumDimensions() - 2) - : TensorShape({}); - + const TensorShape outer_dims_a = a_shape.NumDimensions() > 2 + ? a_shape.Slice(0, a_shape.NumDimensions() - 2) + : TensorShape({}); + const TensorShape outer_dims_b = b_shape.NumDimensions() > 2 + ? b_shape.Slice(0, b_shape.NumDimensions() - 2) + : TensorShape({}); + const TensorShape outer_dims = output_shape.NumDimensions() > 2 + ? output_shape.Slice(0, output_shape.NumDimensions() - 2) + : TensorShape({}); const int64_t batch_size = outer_dims.Size(); - - // Get dimensions for matrix multiplication from TensorShape - const uint32_t dim_a_outer = narrow(a_shape[a_shape.NumDimensions() - 2]); // left matrix second dimension - const uint32_t dim_inner = narrow(a_shape[a_shape.NumDimensions() - 1]); // left matrix first dimension - const uint32_t dim_b_outer = narrow(b_shape[b_shape.NumDimensions() - 1]); // right matrix first dimension - - // A single row has no cross-row weight reuse. Stream B directly and reduce K - // within each workgroup in f32, retaining the existing paths outside this regime. - if (batch_size == 1 && dim_a_outer == 1 && - dim_b_outer >= 16 && dim_b_outer <= 64 && dim_b_outer % 4 == 0 && - dim_inner >= 2048 && dim_inner <= 8192 && - !has_bias && activation.activation_kind_ == ActivationKind::None && - a->IsDataType() && b->IsDataType() && output_tensor->IsDataType()) { - const uint32_t n_vec_count = dim_b_outer / 4; - MatMulGemvProgram program; - program.AddInputs({{a, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, - {b, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}}) - .AddOutput({output_tensor, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten, 4}) - .AddUniformVariables({{dim_inner}, {dim_b_outer}}) - .SetWorkgroupSize(MatMulGemvProgram::kWorkgroupSizeX, MatMulGemvProgram::kWorkgroupSizeY) - .SetDispatchGroupSize(CeilDiv(n_vec_count, MatMulGemvProgram::kWorkgroupSizeX)); - return context->RunProgram(program); - } - + const uint32_t dim_a_outer = narrow(a_shape[a_shape.NumDimensions() - 2]); + const uint32_t dim_inner = narrow(a_shape[a_shape.NumDimensions() - 1]); + const uint32_t dim_b_outer = narrow(b_shape[b_shape.NumDimensions() - 1]); const bool is_vec4 = dim_inner % 4 == 0 && dim_b_outer % 4 == 0; - InlinedVector elements_per_thread = dim_a_outer <= 8 - ? InlinedVector({4, 1, 1}) - : InlinedVector({4, 4, 1}); - - const uint32_t dispatch_x = narrow((dim_b_outer + MatMul::MATMUL_PACKED_WORKGROUP_SIZE_X * elements_per_thread[0] - 1) / - (MatMul::MATMUL_PACKED_WORKGROUP_SIZE_X * elements_per_thread[0])); - const uint32_t dispatch_y = narrow((dim_a_outer + MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Y * elements_per_thread[1] - 1) / - (MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Y * elements_per_thread[1])); - uint32_t dispatch_z = narrow(CeilDiv( - batch_size, MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Z * elements_per_thread[2])); + ORT_RETURN_IF_NOT(IsMatMulPackedConfigurationValid(configuration, is_vec4, use_split_k), + "MatMul packed configuration is invalid for ", + use_split_k ? "Split-K." : "the packed algorithm."); + InlinedVector elements_per_thread{ + configuration.elements_per_thread[0], + configuration.elements_per_thread[1], + configuration.elements_per_thread[2]}; + const auto dispatch_x_value = TryGetMatMulPackedDispatchGroupCount( + dim_b_outer, configuration.workgroup_size[0], configuration.elements_per_thread[0]); + const auto dispatch_y_value = TryGetMatMulPackedDispatchGroupCount( + dim_a_outer, configuration.workgroup_size[1], configuration.elements_per_thread[1]); + const auto dispatch_z_value = TryGetMatMulPackedDispatchGroupCount( + batch_size, configuration.workgroup_size[2], configuration.elements_per_thread[2]); + ORT_RETURN_IF_NOT(dispatch_x_value.has_value() && + dispatch_y_value.has_value() && + dispatch_z_value.has_value(), + "MatMul packed dispatch dimensions exceed uint32 limits."); + const uint32_t dispatch_x = *dispatch_x_value; + const uint32_t dispatch_y = *dispatch_y_value; + uint32_t dispatch_z = *dispatch_z_value; const int components = is_vec4 ? 4 : 1; - const TensorShape a_shape_temp = CreateMatMulIntermediateShape(outer_dims_a, dim_a_outer, dim_inner, components); - const TensorShape b_shape_temp = CreateMatMulIntermediateShape(outer_dims_b, dim_inner, dim_b_outer, components); - const TensorShape output_shape_temp = TensorShape({batch_size, dim_a_outer, dim_b_outer / components}); - + const TensorShape a_shape_temp = + CreateMatMulIntermediateShape(outer_dims_a, dim_a_outer, dim_inner, components); + const TensorShape b_shape_temp = + CreateMatMulIntermediateShape(outer_dims_b, dim_inner, dim_b_outer, components); + const TensorShape output_shape_temp{batch_size, dim_a_outer, dim_b_outer / components}; Tensor split_k_output; - ProgramOutput output(output_tensor, ProgramTensorMetadataDependency::TypeAndRank, output_shape_temp, components); + ProgramOutput output(output_tensor, ProgramTensorMetadataDependency::TypeAndRank, + output_shape_temp, components); const Tensor* bias = has_bias ? inputs[2] : nullptr; bool use_bias_in_matmul = has_bias; uint32_t split_dim_inner = 1; uint32_t splits_per_batch = 1; - // Current Split-K implementation relies on atomic operations, which are not deterministic. - if (!context->KernelContext().GetUseDeterministicCompute()) { - const SplitKConfig& split_k_config = context->GetSplitKConfig(); - const bool need_split_k = split_k_config.UseSplitK( - is_vec4, activation.activation_kind_, batch_size, dim_a_outer, dim_b_outer, - dim_inner, is_channels_last); - if (need_split_k) { - ORT_ENFORCE(is_vec4, "Split-K MatMul requires vec4 packing."); - - if (has_bias) { - ORT_ENFORCE(is_channels_last, "Split-K MatMul only supports channels-last format."); - } - - Tensor* reduction_output = output_tensor; - if (requires_fp32_accumulation) { - split_k_output = context->CreateGPUTensor(DataTypeImpl::GetType(), output_tensor->Shape()); - reduction_output = &split_k_output; - output = ProgramOutput(reduction_output, ProgramTensorMetadataDependency::TypeAndRank, output_shape_temp, components); - } - - const auto fill_bias_program = CreateMatMulFillBiasOrZeroBeforeSplitKProgram(bias, reduction_output, /*is_gemm*/ false, /*beta*/ 1.0f, /*bias_components*/ 4, output_shape_temp, narrow(batch_size)); - ORT_RETURN_IF_ERROR(context->RunProgram(fill_bias_program)); - - // `bias` has been handled in the execution of `fill_bias_program` so we don't need to set - // `bias` again in `MatMulProgram`. - use_bias_in_matmul = false; - - // With Split-K, `dim_inner` will be split into multiple parts. `dispatch_z` encodes - // both the split-k index and the batch index: dispatch_z = splits_per_batch * batch_size. - split_dim_inner = split_k_config.GetSplitDimInner(); - splits_per_batch = (dim_inner + split_dim_inner - 1) / split_dim_inner; - const uint64_t dispatch_z_u64 = static_cast(batch_size) * static_cast(splits_per_batch); - ORT_ENFORCE(dispatch_z_u64 <= static_cast(std::numeric_limits::max()), - "dispatch_z exceeds uint32_t range: ", dispatch_z_u64); - dispatch_z = narrow(dispatch_z_u64); - - // The output should be declared in atomic types in `MatMulProgram` for the use of atomic - // built-in functions. - output.is_atomic = true; + if (use_split_k) { + ORT_RETURN_IF(context.KernelContext().GetUseDeterministicCompute(), + "MatMul algorithm packed_split_k does not support deterministic compute."); + ORT_RETURN_IF_NOT(configuration.split_dim_inner > 1, + "MatMul algorithm packed_split_k is not configured for this adapter."); + ORT_RETURN_IF_NOT(is_vec4, + "MatMul algorithm packed_split_k requires vec4 packing."); + ORT_RETURN_IF_NOT(activation.activation_kind_ == ActivationKind::None, + "MatMul algorithm packed_split_k does not support a fused activation."); + ORT_RETURN_IF_NOT(!has_bias || is_channels_last, + "MatMul algorithm packed_split_k requires channels-last bias layout."); + + Tensor* reduction_output = output_tensor; + if (requires_fp32_accumulation) { + split_k_output = context.CreateGPUTensor(DataTypeImpl::GetType(), output_tensor->Shape()); + reduction_output = &split_k_output; + output = ProgramOutput(reduction_output, ProgramTensorMetadataDependency::TypeAndRank, + output_shape_temp, components); } + + const auto fill_bias_program = CreateMatMulFillBiasOrZeroBeforeSplitKProgram( + bias, reduction_output, /*is_gemm=*/false, /*beta=*/1.0f, + /*output_components=*/4, output_shape_temp, narrow(batch_size)); + ORT_RETURN_IF_ERROR(context.RunProgram(fill_bias_program)); + use_bias_in_matmul = false; + split_dim_inner = configuration.split_dim_inner; + const auto splits_per_batch_value = TryGetMatMulPackedDispatchGroupCount( + dim_inner, split_dim_inner, 1); + ORT_RETURN_IF_NOT(splits_per_batch_value.has_value(), + "MatMul packed Split-K count exceeds uint32 limits."); + splits_per_batch = *splits_per_batch_value; + const uint64_t dispatch_z_u64 = + static_cast(batch_size) * static_cast(splits_per_batch); + ORT_RETURN_IF_NOT(dispatch_z_u64 <= static_cast(std::numeric_limits::max()), + "MatMul algorithm packed_split_k dispatch_z exceeds uint32_t range: ", dispatch_z_u64); + dispatch_z = narrow(dispatch_z_u64); + output.is_atomic = true; } - MatMulProgram matmul_program{activation, use_bias_in_matmul, is_vec4, elements_per_thread, is_channels_last, split_dim_inner, requires_fp32_accumulation}; - matmul_program - .CacheHint(activation.CacheKey(), absl::StrJoin(elements_per_thread, "-"), std::to_string(is_vec4), components, is_channels_last, split_dim_inner, requires_fp32_accumulation) + MatMulProgram program{activation, use_bias_in_matmul, is_vec4, elements_per_thread, + is_channels_last, split_dim_inner, configuration.tile_inner, + requires_fp32_accumulation}; + program + .CacheHint(activation.CacheKey(), absl::StrJoin(elements_per_thread, "-"), + absl::StrJoin(configuration.workgroup_size, "-"), + std::to_string(is_vec4), components, is_channels_last, + split_dim_inner, configuration.tile_inner, requires_fp32_accumulation) .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, a_shape_temp, components}, {b, ProgramTensorMetadataDependency::TypeAndRank, b_shape_temp, components}}) .AddUniformVariables({{dim_a_outer}, {dim_b_outer}, {dim_inner}, {dispatch_x}, {dispatch_y}, {dispatch_z}, {splits_per_batch}}) .AddIndices(outer_dims) .SetDispatchGroupSize(dispatch_x, dispatch_y, dispatch_z) - .SetWorkgroupSize(MatMul::MATMUL_PACKED_WORKGROUP_SIZE_X, MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Y, MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Z) + .SetWorkgroupSize(configuration.workgroup_size[0], + configuration.workgroup_size[1], + configuration.workgroup_size[2]) .AddOutput(std::move(output)); - // Activation uniforms must remain last because definitions and values are matched by index. - AppendActivationUniformsData(activation, matmul_program); + AppendActivationUniformsData(activation, program); if (use_bias_in_matmul) { - auto bias_components = is_channels_last ? components : 1; - TensorShape reduced_bias_shape = ReduceShapeByComponents(bias->Shape(), bias_components); - matmul_program.AddInput({bias, ProgramTensorMetadataDependency::Rank, reduced_bias_shape, bias_components}); + const int bias_components = is_channels_last ? components : 1; + const TensorShape reduced_bias_shape = ReduceShapeByComponents(bias->Shape(), bias_components); + program.AddInput({bias, ProgramTensorMetadataDependency::Rank, reduced_bias_shape, bias_components}); } - ORT_RETURN_IF_ERROR(context->RunProgram(matmul_program)); - if (split_dim_inner > 1 && requires_fp32_accumulation) { + ORT_RETURN_IF_ERROR(context.RunProgram(program)); + if (use_split_k && requires_fp32_accumulation) { const uint32_t output_size = narrow(output_tensor->Shape().Size()); const uint32_t vec_size = output_size / 4; - CastProgram cast_program{ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, false, true, false, false}; + CastProgram cast_program{ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, + false, true, false, false}; cast_program .AddInput({&split_k_output, ProgramTensorMetadataDependency::Type, {vec_size}, 4}) .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {vec_size}, 4}) .SetDispatchGroupSize(CeilDiv(vec_size, static_cast(WORKGROUP_SIZE))) .AddUniformVariables({{vec_size}, {output_size}}) .CacheHint(std::to_string(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16)); - return context->RunProgram(cast_program); + return context.RunProgram(cast_program); } return Status::OK(); } +Status MatMulComputeDispatcher::Compute(ComputeContext& context, + const Activation& activation, + const std::vector& inputs, + Tensor* output_tensor, + bool is_channels_last, + bool b_is_constant) { + const auto* a = inputs[0]; + const auto* b = inputs[1]; + const bool has_bias = inputs.size() > 2; + const TensorShape& logical_a_shape = a->Shape(); + const TensorShape& logical_b_shape = b->Shape(); + ORT_RETURN_IF_NOT(logical_a_shape.NumDimensions() >= 2 && logical_b_shape.NumDimensions() >= 2, + "ComputeMatMul expects matrix or batched-matrix inputs."); + + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(logical_a_shape, logical_b_shape)); + + if (helper.OutputShape().Size() == 0) { + return Status::OK(); + } + if (helper.K() == 0) { + if (const auto forced_algorithm = context.ForcedMatMulAlgorithm()) { + ORT_RETURN_IF_NOT(*forced_algorithm == MatMulAlgorithm::Naive, + "MatMul algorithm ", MatMulAlgorithmName(*forced_algorithm), + " does not support these inputs or this device."); + } + // Empty inputs have null GPU buffers and must not be bound to a MatMul shader. + const Tensor* bias = has_bias ? inputs[2] : nullptr; + const bool bias_is_scalar = bias && bias->Shape().Size() == 1; + if (bias) { + const int64_t expected_bias_size = is_channels_last ? helper.N() : helper.M(); + ORT_RETURN_IF_NOT(bias_is_scalar || bias->Shape().Size() == expected_bias_size, + "Zero-K MatMul bias must be scalar or match the output channel dimension."); + } + const uint32_t output_size = narrow(helper.OutputShape().Size()); + MatMulZeroKProgram program{activation, bias != nullptr, is_channels_last, bias_is_scalar}; + program.CacheHint(activation.CacheKey(), bias != nullptr, is_channels_last, bias_is_scalar) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::TypeAndRank, 1}) + .AddUniformVariables( + {{output_size}, {narrow(helper.M())}, {narrow(helper.N())}}) + .SetDispatchGroupSize(CeilDiv(output_size, static_cast(WORKGROUP_SIZE))); + AppendActivationUniformsData(activation, program); + if (bias) { + program.AddInput({bias, ProgramTensorMetadataDependency::TypeAndRank, 1}); + } + return context.RunProgram(program); + } + + Initialize(context); + const bool requires_fp32_accumulation = + context.EnableMatmulFp32Accumulation() && a->IsDataType(); + SubgroupMatrixMatMulImpl* subgroup_impl = subgroup_matrix_impl_.get(); + const bool can_use_subgroup_matrix = + !requires_fp32_accumulation && subgroup_impl != nullptr && + subgroup_impl->CanApply(context, inputs, is_channels_last, b_is_constant); + const std::optional subgroup_size = intel::SelectMatMulSubgroupSize(context); + const bool has_subgroup_capability = subgroup_size.has_value(); + const int64_t batch_a = + logical_a_shape.NumDimensions() > 2 + ? logical_a_shape.SizeToDimension(logical_a_shape.NumDimensions() - 2) + : 1; + const int64_t batch_b = + logical_b_shape.NumDimensions() > 2 + ? logical_b_shape.SizeToDimension(logical_b_shape.NumDimensions() - 2) + : 1; + const TensorShape& logical_output_shape = helper.OutputShape(); + const uint64_t batch_size = logical_output_shape.NumDimensions() > 2 + ? narrow(logical_output_shape.SizeToDimension( + logical_output_shape.NumDimensions() - 2)) + : 1; + const bool folds_batch_into_m = batch_a != 1 && batch_b == 1; + const int64_t packed_m = folds_batch_into_m + ? logical_a_shape.SizeToDimension(logical_a_shape.NumDimensions() - 1) + : helper.M(); + const uint64_t packed_batch_size = folds_batch_into_m ? 1 : batch_size; + const bool can_run_gemv = + packed_batch_size == 1 && packed_m >= 1 && packed_m <= 8 && + helper.N() > 0 && helper.N() % 4 == 0 && helper.K() > 0 && + !has_bias && activation.activation_kind_ == ActivationKind::None && + a->IsDataType() && b->IsDataType() && + output_tensor->IsDataType(); + + MatMulAlgorithmSelectionParams selection_params{}; + selection_params.m = helper.M(); + selection_params.n = helper.N(); + selection_params.k = helper.K(); + selection_params.packed_m = packed_m; + selection_params.batch_size = batch_size; + selection_params.packed_batch_size = packed_batch_size; + selection_params.adapter_architecture = context.AdapterInfo().architecture; + selection_params.a_data_type = a->GetElementType(); + selection_params.b_data_type = b->GetElementType(); + selection_params.can_use_subgroup_matrix = can_use_subgroup_matrix; + selection_params.can_run_gemv = can_run_gemv; + selection_params.has_subgroup_capability = has_subgroup_capability; + selection_params.subgroup_size = subgroup_size.value_or(0); + selection_params.is_vec4 = helper.K() % 4 == 0 && helper.N() % 4 == 0; + selection_params.deterministic_compute = context.KernelContext().GetUseDeterministicCompute(); + selection_params.has_fused_activation = activation.activation_kind_ != ActivationKind::None; + selection_params.has_bias = has_bias; + selection_params.is_channels_last = is_channels_last; + + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + ORT_RETURN_IF_ERROR(planner_->CreateExecutionPlan( + selection_params, context.ForcedMatMulAlgorithm(), plan)); + const MatMulAlgorithm algorithm = plan.algorithm; + const auto* packed_configuration = std::get_if(&plan.configuration); + const auto* subgroup_configuration = std::get_if(&plan.configuration); + + switch (algorithm) { + case MatMulAlgorithm::SubgroupMatrix: + ORT_RETURN_IF_NOT(subgroup_impl != nullptr, + "MatMul algorithm subgroup_matrix is unavailable."); + return subgroup_impl->Compute( + context, inputs, output_tensor, activation, is_channels_last, b_is_constant); + case MatMulAlgorithm::Gemv: + return ApplyMatMulGemv(context, inputs, output_tensor, helper); + case MatMulAlgorithm::Naive: + return ApplyMatMulNaive( + context, activation, inputs, output_tensor, is_channels_last, helper); + case MatMulAlgorithm::Subgroup: + return intel::ApplyMatMulSubgroup( + context, activation, inputs, output_tensor, is_channels_last, + subgroup_configuration->subgroup_size); + case MatMulAlgorithm::Packed: + return ApplyMatMulPacked( + context, activation, inputs, output_tensor, is_channels_last, helper, + *packed_configuration, + /*use_split_k=*/false); + case MatMulAlgorithm::PackedSplitK: + return ApplyMatMulPacked( + context, activation, inputs, output_tensor, is_channels_last, helper, + *packed_configuration, + /*use_split_k=*/true); + } + + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unknown MatMul algorithm."); +} + MatMulFillBiasOrZeroBeforeSplitKProgram CreateMatMulFillBiasOrZeroBeforeSplitKProgram( const Tensor* bias, Tensor* output, diff --git a/onnxruntime/core/providers/webgpu/math/matmul.h b/onnxruntime/core/providers/webgpu/math/matmul.h index 4aaca49f0edbe..7c2af1c6ec294 100644 --- a/onnxruntime/core/providers/webgpu/math/matmul.h +++ b/onnxruntime/core/providers/webgpu/math/matmul.h @@ -3,12 +3,10 @@ #pragma once -#include -#include - #include "core/providers/webgpu/webgpu_kernel.h" #include "core/providers/webgpu/program.h" #include "core/providers/cpu/math/matmul_helper.h" +#include "core/providers/webgpu/math/matmul_compute_dispatcher.h" #include "core/providers/webgpu/math/matmul_utils.h" #include "core/providers/webgpu/math/matmul_packed.h" #include "core/providers/webgpu/webgpu_utils.h" @@ -17,35 +15,6 @@ namespace onnxruntime { namespace webgpu { -class MatMulOptImpl { - public: - virtual ~MatMulOptImpl() = default; - - virtual Status Compute(ComputeContext& context, - const std::vector& inputs, - Tensor* output, - const Activation& activation, - bool is_channels_last, - bool b_is_constant, - /*out*/ bool& handled) = 0; -}; - -class MatMulOptImplCache { - public: - MatMulOptImplCache() = default; - ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MatMulOptImplCache); - - MatMulOptImpl* GetOrCreate(const ComputeContextBase& context); - - private: - std::once_flag subgroup_impl_init_flag_; - std::unique_ptr subgroup_impl_; -}; - -Status ComputeMatMul(ComputeContext* context, const Activation& activation, std::vector& inputs, Tensor* output, - bool is_channels_last, MatMulOptImplCache& cache, - bool b_is_constant = false); - MatMulFillBiasOrZeroBeforeSplitKProgram CreateMatMulFillBiasOrZeroBeforeSplitKProgram( const Tensor* bias, Tensor* output, @@ -72,14 +41,15 @@ class MatMul final : public WebGpuKernel { constexpr static uint32_t MATMUL_PACKED_WORKGROUP_SIZE_Z = 1; private: - mutable MatMulOptImplCache compute_cache_; + mutable MatMulComputeDispatcher compute_dispatcher_; bool b_is_constant_ = false; }; class MatMulNaiveProgram final : public Program { public: - MatMulNaiveProgram(const Activation& activation, const size_t output_rank, int64_t output_number, bool has_bias, bool is_channels_last = false) - : Program{"MatMulNaive"}, activation_(activation), output_rank_(output_rank), output_number_(output_number), has_bias_{has_bias}, is_channels_last_(is_channels_last) { + MatMulNaiveProgram(const Activation& activation, const size_t output_rank, int64_t output_number, + bool has_bias, bool is_channels_last = false, bool is_zero_k = false) + : Program{"MatMulNaive"}, activation_(activation), output_rank_(output_rank), output_number_(output_number), has_bias_{has_bias}, is_channels_last_(is_channels_last), is_zero_k_(is_zero_k) { } Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -96,6 +66,24 @@ class MatMulNaiveProgram final : public Program { const int64_t output_number_; const bool has_bias_; const bool is_channels_last_; + const bool is_zero_k_; +}; + +class MatMulGemvProgram final : public Program { + public: + MatMulGemvProgram() : Program{"MatMulGemv"} {} + + static constexpr uint32_t kWorkgroupSizeX = 1; + static constexpr uint32_t kWorkgroupSizeY = 128; + static_assert(kWorkgroupSizeY != 0 && (kWorkgroupSizeY & (kWorkgroupSizeY - 1)) == 0, + "MatMulGemvProgram requires kWorkgroupSizeY to be a non-zero power of two " + "because the WGSL reduction loop halves workgroup_size_y each iteration."); + + Status GenerateShaderCode(ShaderHelper& shader) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"K", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}); }; } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/math/matmul_algorithm.cc b/onnxruntime/core/providers/webgpu/math/matmul_algorithm.cc new file mode 100644 index 0000000000000..2e8ef72d255f7 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/matmul_algorithm.cc @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/math/matmul_algorithm.h" + +namespace onnxruntime { +namespace webgpu { + +std::string_view MatMulAlgorithmName(MatMulAlgorithm algorithm) { + switch (algorithm) { + case MatMulAlgorithm::SubgroupMatrix: + return "subgroup_matrix"; + case MatMulAlgorithm::Gemv: + return "gemv"; + case MatMulAlgorithm::Naive: + return "naive"; + case MatMulAlgorithm::Subgroup: + return "subgroup"; + case MatMulAlgorithm::Packed: + return "packed"; + case MatMulAlgorithm::PackedSplitK: + return "packed_split_k"; + } + return "unknown"; +} + +} // namespace webgpu +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/webgpu/math/matmul_algorithm.h b/onnxruntime/core/providers/webgpu/math/matmul_algorithm.h new file mode 100644 index 0000000000000..566dd884aa837 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/matmul_algorithm.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +namespace onnxruntime { +namespace webgpu { + +enum class MatMulAlgorithm { + SubgroupMatrix, + Gemv, + Naive, + Subgroup, + Packed, + PackedSplitK, +}; + +std::string_view MatMulAlgorithmName(MatMulAlgorithm algorithm); + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/math/matmul_compute_dispatcher.h b/onnxruntime/core/providers/webgpu/math/matmul_compute_dispatcher.h new file mode 100644 index 0000000000000..1ec5f3ff7667a --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/matmul_compute_dispatcher.h @@ -0,0 +1,66 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "core/common/common.h" +#include "core/common/status.h" +#include "core/providers/webgpu/math/matmul_execution_planner.h" + +namespace onnxruntime { + +class Tensor; + +namespace webgpu { + +struct Activation; +class ComputeContext; +class ComputeContextBase; + +// Stateful subgroup-matrix implementation owned independently from selection policy. +class SubgroupMatrixMatMulImpl { + public: + SubgroupMatrixMatMulImpl() = default; + virtual ~SubgroupMatrixMatMulImpl() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(SubgroupMatrixMatMulImpl); + + virtual bool CanApply(const ComputeContext& context, + const std::vector& inputs, + bool is_channels_last, + bool b_is_constant) const = 0; + + virtual Status Compute(ComputeContext& context, + const std::vector& inputs, + Tensor* output, + const Activation& activation, + bool is_channels_last, + bool b_is_constant) = 0; +}; + +// Selects, validates, and dispatches one concrete implementation per invocation. +class MatMulComputeDispatcher { + public: + MatMulComputeDispatcher() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MatMulComputeDispatcher); + + Status Compute(ComputeContext& context, + const Activation& activation, + const std::vector& inputs, + Tensor* output, + bool is_channels_last, + bool b_is_constant = false); + + private: + void Initialize(const ComputeContextBase& context); + + std::once_flag init_flag_; + std::unique_ptr subgroup_matrix_impl_; + std::unique_ptr planner_; +}; + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/math/matmul_execution_planner.cc b/onnxruntime/core/providers/webgpu/math/matmul_execution_planner.cc new file mode 100644 index 0000000000000..dc3350f0dcc83 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/matmul_execution_planner.cc @@ -0,0 +1,345 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/math/matmul_execution_planner.h" + +#include +#include + +#include "core/common/safeint.h" +#include "core/providers/webgpu/vendor/intel/math/gemm_subgroup.h" + +namespace onnxruntime { +namespace webgpu { + +namespace { + +constexpr uint64_t kMaxUint32 = std::numeric_limits::max(); + +bool IsUint32Dimension(int64_t dimension) { + return dimension >= 0 && static_cast(dimension) <= kMaxUint32; +} + +bool IsMatMulNaiveOutputSizeValid(const MatMulAlgorithmSelectionParams& params) { + if (!IsUint32Dimension(params.m) || + !IsUint32Dimension(params.n) || + !IsUint32Dimension(params.k)) { + return false; + } + + const uint64_t output_rows = static_cast(params.m) / GetMaxComponents(params.m); + const uint64_t output_columns = static_cast(params.n) / GetMaxComponents(params.n); + uint64_t output_size = params.batch_size; + for (const uint64_t factor : {output_rows, output_columns}) { + uint64_t product{}; + if (!SafeMultiply(output_size, factor, product) || product > kMaxUint32) { + return false; + } + output_size = product; + } + return true; +} + +bool IsDispatchGroupCountProductValid(uint32_t x, uint32_t y, uint32_t z) { + // ProgramManager normalizes device-specific per-axis limits. Plans only need to + // guarantee that ShaderHelper's flattened logical workgroup ID fits in uint32. + uint32_t xy{}; + uint32_t xyz{}; + return SafeMultiply(x, y, xy) && SafeMultiply(xy, z, xyz); +} + +bool IsMatMulAlgorithmConfigurationCompatible(const MatMulExecutionPlan& plan) { + switch (plan.algorithm) { + case MatMulAlgorithm::SubgroupMatrix: + return std::holds_alternative(plan.configuration); + case MatMulAlgorithm::Gemv: + return std::holds_alternative(plan.configuration); + case MatMulAlgorithm::Naive: + return std::holds_alternative(plan.configuration); + case MatMulAlgorithm::Subgroup: + return std::holds_alternative(plan.configuration); + case MatMulAlgorithm::Packed: + case MatMulAlgorithm::PackedSplitK: + return std::holds_alternative(plan.configuration); + } + return false; +} + +} // namespace + +bool IsMatMulPackedConfigurationValid( + const MatMulPackedConfiguration& configuration, + bool is_vec4, + bool use_split_k) { + if (configuration.workgroup_size[0] == 0 || + configuration.workgroup_size[1] == 0 || + configuration.workgroup_size[2] != 1 || + configuration.elements_per_thread[0] == 0 || + configuration.elements_per_thread[1] == 0 || + configuration.elements_per_thread[2] != 1 || + configuration.tile_inner == 0) { + return false; + } + + if (configuration.tile_inner % configuration.workgroup_size[0] != 0 || + configuration.tile_inner % configuration.workgroup_size[1] != 0) { + return false; + } + + if (is_vec4) { + const uint32_t a_vector_width = configuration.tile_inner / configuration.workgroup_size[0]; + if (configuration.elements_per_thread[0] != 4 || + a_vector_width != 4) { + return false; + } + } + + return !use_split_k || + (configuration.split_dim_inner > 1 && + configuration.split_dim_inner % configuration.tile_inner == 0); +} + +std::optional TryGetMatMulPackedDispatchGroupCount( + uint64_t dimension, + uint32_t workgroup_size, + uint32_t elements_per_thread) { + if (workgroup_size == 0 || elements_per_thread == 0) { + return std::nullopt; + } + + const uint64_t elements_per_group = + static_cast(workgroup_size) * elements_per_thread; + const uint64_t group_count = + dimension == 0 ? 0 : 1 + (dimension - 1) / elements_per_group; + if (group_count > std::numeric_limits::max()) { + return std::nullopt; + } + + return static_cast(group_count); +} + +MatMulExecutionPlanner::MatMulExecutionPlanner(SplitKConfig split_k_config) + : split_k_config_{std::move(split_k_config)} {} + +MatMulExecutionPlanner::~MatMulExecutionPlanner() = default; + +MatMulAlgorithm MatMulExecutionPlanner::SelectAlgorithm( + const MatMulAlgorithmSelectionParams& params, + std::optional forced_algorithm) const { + if (forced_algorithm.has_value()) { + return *forced_algorithm; + } + if (params.k == 0) { + return MatMulAlgorithm::Naive; + } + if (const auto vendor_algorithm = SelectVendorAlgorithm(params); vendor_algorithm.has_value()) { + return *vendor_algorithm; + } + return SelectCommonAlgorithm(params); +} + +Status MatMulExecutionPlanner::CreateExecutionPlan( + const MatMulAlgorithmSelectionParams& params, + std::optional forced_algorithm, + MatMulExecutionPlan& plan) const { + const MatMulAlgorithm algorithm = SelectAlgorithm(params, forced_algorithm); + std::optional configuration = + SelectVendorConfiguration(algorithm, params); + if (!configuration.has_value()) { + configuration = SelectCommonConfiguration(algorithm, params); + } + MatMulExecutionPlan candidate{algorithm, std::move(*configuration)}; + ORT_RETURN_IF_ERROR(ValidateExecutionPlan(params, candidate)); + plan = std::move(candidate); + return Status::OK(); +} + +std::optional MatMulExecutionPlanner::SelectVendorAlgorithm( + const MatMulAlgorithmSelectionParams& /*params*/) const { + return std::nullopt; +} + +std::optional MatMulExecutionPlanner::SelectVendorConfiguration( + MatMulAlgorithm /*algorithm*/, + const MatMulAlgorithmSelectionParams& /*params*/) const { + return std::nullopt; +} + +MatMulAlgorithm MatMulExecutionPlanner::SelectCommonAlgorithm( + const MatMulAlgorithmSelectionParams& params) const { + if (params.can_use_subgroup_matrix) { + return MatMulAlgorithm::SubgroupMatrix; + } + if (ShouldUseGemv(params)) { + return MatMulAlgorithm::Gemv; + } + if (params.n < 8 && params.k < 8) { + return MatMulAlgorithm::Naive; + } + if (ShouldUseSplitK(params)) { + return MatMulAlgorithm::PackedSplitK; + } + return MatMulAlgorithm::Packed; +} + +MatMulAlgorithmConfiguration MatMulExecutionPlanner::SelectCommonConfiguration( + MatMulAlgorithm algorithm, + const MatMulAlgorithmSelectionParams& params) const { + switch (algorithm) { + case MatMulAlgorithm::SubgroupMatrix: + return MatMulSubgroupMatrixConfiguration{}; + case MatMulAlgorithm::Gemv: + return MatMulGemvConfiguration{}; + case MatMulAlgorithm::Naive: + return MatMulNaiveConfiguration{}; + case MatMulAlgorithm::Subgroup: + return MatMulSubgroupConfiguration{params.subgroup_size}; + case MatMulAlgorithm::Packed: + case MatMulAlgorithm::PackedSplitK: { + MatMulPackedConfiguration configuration{}; + configuration.elements_per_thread = + params.packed_m <= 8 ? std::array{4, 1, 1} + : std::array{4, 4, 1}; + configuration.split_dim_inner = + algorithm == MatMulAlgorithm::PackedSplitK + ? split_k_config_.GetSplitDimInner() + : 1; + return configuration; + } + } + return MatMulNaiveConfiguration{}; +} + +Status MatMulExecutionPlanner::ValidateExecutionPlan( + const MatMulAlgorithmSelectionParams& params, + const MatMulExecutionPlan& plan) const { + ORT_RETURN_IF_NOT(IsMatMulAlgorithmConfigurationCompatible(plan), + "MatMul algorithm ", MatMulAlgorithmName(plan.algorithm), + " received an incompatible vendor configuration."); + + switch (plan.algorithm) { + case MatMulAlgorithm::SubgroupMatrix: + ORT_RETURN_IF_NOT(params.can_use_subgroup_matrix, + "MatMul algorithm subgroup_matrix does not support these inputs or this device."); + break; + case MatMulAlgorithm::Gemv: + ORT_RETURN_IF_NOT(params.can_run_gemv && + IsUint32Dimension(params.n) && + IsUint32Dimension(params.k), + "MatMul algorithm gemv does not support these inputs or this device."); + break; + case MatMulAlgorithm::Naive: + ORT_RETURN_IF_NOT(IsMatMulNaiveOutputSizeValid(params), + "MatMul algorithm naive dimensions or vectorized output size exceed uint32 limits."); + break; + case MatMulAlgorithm::Subgroup: { + const auto& configuration = std::get(plan.configuration); + const bool supported_size = configuration.subgroup_size == 8 || + configuration.subgroup_size == 16 || + configuration.subgroup_size == 32; + ORT_RETURN_IF_NOT(params.has_subgroup_capability && params.k > 0 && supported_size && + configuration.subgroup_size == params.subgroup_size, + "MatMul algorithm subgroup does not support these inputs, configuration, or device."); + + const bool can_fold_batch_into_m = params.packed_batch_size != params.batch_size; + const bool fold_batch_into_m = intel::ShouldFoldMatMulSubgroupBatch( + params.adapter_architecture, params.m, can_fold_batch_into_m); + const int64_t effective_m = fold_batch_into_m ? params.packed_m : params.m; + const uint64_t effective_batch_size = fold_batch_into_m ? params.packed_batch_size : params.batch_size; + ORT_RETURN_IF_NOT(effective_m > 0 && IsUint32Dimension(effective_m) && + params.n > 0 && IsUint32Dimension(params.n) && + IsUint32Dimension(params.k) && + effective_batch_size > 0 && effective_batch_size <= kMaxUint32, + "MatMul subgroup dimensions must fit in uint32."); + + const uint32_t elements_per_thread_y = static_cast( + intel::MatMulSubgroupElementsPerThreadY(params.adapter_architecture, + static_cast(effective_m))); + const auto dispatch_x = TryGetMatMulPackedDispatchGroupCount( + static_cast(params.n), intel::kSubgroupLogicalWorkGroupSizeX, 4); + const auto dispatch_y = TryGetMatMulPackedDispatchGroupCount( + static_cast(effective_m), intel::kSubgroupLogicalWorkGroupSizeY, + elements_per_thread_y); + const auto dispatch_z = TryGetMatMulPackedDispatchGroupCount( + effective_batch_size, intel::kSubgroupLogicalWorkGroupSizeZ, 1); + ORT_RETURN_IF_NOT(dispatch_x.has_value() && dispatch_y.has_value() && dispatch_z.has_value() && + IsDispatchGroupCountProductValid(*dispatch_x, *dispatch_y, *dispatch_z), + "MatMul subgroup dispatch group count exceeds uint32 limits."); + break; + } + case MatMulAlgorithm::Packed: + case MatMulAlgorithm::PackedSplitK: { + const bool use_split_k = plan.algorithm == MatMulAlgorithm::PackedSplitK; + const auto& configuration = std::get(plan.configuration); + ORT_RETURN_IF_NOT(params.k > 0 && + IsMatMulPackedConfigurationValid(configuration, params.is_vec4, use_split_k), + "MatMul algorithm ", MatMulAlgorithmName(plan.algorithm), + " has invalid inputs or packed configuration."); + if (use_split_k) { + ORT_RETURN_IF_NOT(!params.deterministic_compute && params.is_vec4 && + !params.has_fused_activation && + (!params.has_bias || params.is_channels_last), + "MatMul algorithm packed_split_k does not support these inputs."); + } + + ORT_RETURN_IF_NOT(IsUint32Dimension(params.packed_m) && + IsUint32Dimension(params.n) && + IsUint32Dimension(params.k), + "MatMul packed dimensions must fit in uint32."); + const auto dispatch_x = TryGetMatMulPackedDispatchGroupCount( + static_cast(params.n), configuration.workgroup_size[0], + configuration.elements_per_thread[0]); + const auto dispatch_y = TryGetMatMulPackedDispatchGroupCount( + static_cast(params.packed_m), configuration.workgroup_size[1], + configuration.elements_per_thread[1]); + const auto dispatch_z = TryGetMatMulPackedDispatchGroupCount( + params.packed_batch_size, configuration.workgroup_size[2], + configuration.elements_per_thread[2]); + ORT_RETURN_IF_NOT(dispatch_x.has_value() && dispatch_y.has_value() && dispatch_z.has_value(), + "MatMul packed dispatch dimensions exceed uint32 limits."); + + uint32_t final_dispatch_z = *dispatch_z; + if (use_split_k) { + const auto splits_per_batch = TryGetMatMulPackedDispatchGroupCount( + static_cast(params.k), configuration.split_dim_inner, 1); + uint32_t split_dispatch_z{}; + ORT_RETURN_IF_NOT(params.packed_batch_size > 0 && splits_per_batch.has_value() && + SafeMultiply(final_dispatch_z, *splits_per_batch, split_dispatch_z), + "MatMul packed Split-K dispatch dimension exceeds uint32 limits."); + final_dispatch_z = split_dispatch_z; + } + ORT_RETURN_IF_NOT(IsDispatchGroupCountProductValid(*dispatch_x, *dispatch_y, final_dispatch_z), + "MatMul packed combined dispatch group count exceeds uint32 limits."); + break; + } + } + + return Status::OK(); +} + +bool MatMulExecutionPlanner::ShouldUseSplitK( + const MatMulAlgorithmSelectionParams& params) const { + if (params.deterministic_compute || params.has_fused_activation || + params.packed_m < 0 || params.n < 0 || params.k < 0) { + return false; + } + + return split_k_config_.UseSplitK( + params.is_vec4, + ActivationKind::None, + params.packed_batch_size, + static_cast(params.packed_m), + static_cast(params.n), + static_cast(params.k), + params.is_channels_last); +} + +bool MatMulExecutionPlanner::ShouldUseGemv( + const MatMulAlgorithmSelectionParams& params) const { + return params.can_run_gemv && + params.n >= 16 && params.n <= 64 && + params.k >= 2048 && params.k <= 8192; +} + +} // namespace webgpu +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/webgpu/math/matmul_execution_planner.h b/onnxruntime/core/providers/webgpu/math/matmul_execution_planner.h new file mode 100644 index 0000000000000..e176d214a32af --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/matmul_execution_planner.h @@ -0,0 +1,125 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include + +#include "core/common/common.h" +#include "core/providers/webgpu/math/matmul_algorithm.h" +#include "core/providers/webgpu/webgpu_utils.h" + +namespace onnxruntime { +namespace webgpu { + +// Immutable problem facts used by automatic selection and execution tuning. +struct MatMulAlgorithmSelectionParams { + int64_t m = 0; + int64_t n = 0; + int64_t k = 0; + int64_t packed_m = 0; + uint64_t batch_size = 1; + uint64_t packed_batch_size = 1; + std::string_view adapter_architecture; + int32_t a_data_type = 0; + int32_t b_data_type = 0; + bool can_use_subgroup_matrix = false; + bool can_run_gemv = false; + bool has_subgroup_capability = false; + uint32_t subgroup_size = 0; + bool is_vec4 = false; + bool deterministic_compute = false; + bool has_fused_activation = false; + bool has_bias = false; + bool is_channels_last = true; +}; + +// Algorithm-specific tuning selected together with the implementation. +struct MatMulSubgroupMatrixConfiguration {}; +struct MatMulGemvConfiguration {}; +struct MatMulNaiveConfiguration {}; +struct MatMulSubgroupConfiguration { + uint32_t subgroup_size = 0; +}; + +struct MatMulPackedConfiguration { + std::array workgroup_size{8, 8, 1}; + std::array elements_per_thread{4, 4, 1}; + uint32_t tile_inner = 32; + uint32_t split_dim_inner = 1; +}; + +bool IsMatMulPackedConfigurationValid( + const MatMulPackedConfiguration& configuration, + bool is_vec4, + bool use_split_k); + +std::optional TryGetMatMulPackedDispatchGroupCount( + uint64_t dimension, + uint32_t workgroup_size, + uint32_t elements_per_thread); + +using MatMulAlgorithmConfiguration = + std::variant; + +// Complete per-invocation decision consumed by the compute dispatcher. +struct MatMulExecutionPlan { + MatMulAlgorithm algorithm; + MatMulAlgorithmConfiguration configuration; +}; + +// Pure policy object that selects an algorithm and creates a host-validated execution plan. +// Device resource validation and execution remain in MatMulComputeDispatcher. +class MatMulExecutionPlanner { + public: + explicit MatMulExecutionPlanner(SplitKConfig split_k_config = {}); + virtual ~MatMulExecutionPlanner(); + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MatMulExecutionPlanner); + + // Selects one algorithm from immutable invocation facts without validating executability. + MatMulAlgorithm SelectAlgorithm( + const MatMulAlgorithmSelectionParams& params, + std::optional forced_algorithm = std::nullopt) const; + + // Creates a host-validated plan containing the selected algorithm and its typed tuning + // configuration. Unsupported forced requests return an error without falling back. + Status CreateExecutionPlan( + const MatMulAlgorithmSelectionParams& params, + std::optional forced_algorithm, + MatMulExecutionPlan& plan) const; + + protected: + virtual std::optional SelectVendorAlgorithm( + const MatMulAlgorithmSelectionParams& params) const; + + virtual std::optional SelectVendorConfiguration( + MatMulAlgorithm algorithm, + const MatMulAlgorithmSelectionParams& params) const; + + private: + MatMulAlgorithm SelectCommonAlgorithm(const MatMulAlgorithmSelectionParams& params) const; + + MatMulAlgorithmConfiguration SelectCommonConfiguration( + MatMulAlgorithm algorithm, + const MatMulAlgorithmSelectionParams& params) const; + + Status ValidateExecutionPlan( + const MatMulAlgorithmSelectionParams& params, + const MatMulExecutionPlan& plan) const; + + bool ShouldUseSplitK(const MatMulAlgorithmSelectionParams& params) const; + bool ShouldUseGemv(const MatMulAlgorithmSelectionParams& params) const; + + SplitKConfig split_k_config_; +}; + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/math/matmul_gemv.wgsl.template b/onnxruntime/core/providers/webgpu/math/matmul_gemv.wgsl.template index bce589743af44..5b1a5314498f6 100644 --- a/onnxruntime/core/providers/webgpu/math/matmul_gemv.wgsl.template +++ b/onnxruntime/core/providers/webgpu/math/matmul_gemv.wgsl.template @@ -6,12 +6,19 @@ var partials: array, workgroup_size_x * workgroup_size_y>; $MAIN { - let n_vec = workgroup_idx * workgroup_size_x + local_id.x; let n_vec_count = uniforms.N / 4u; + let groups_per_row = (n_vec_count + workgroup_size_x - 1u) / workgroup_size_x; + // Use the flattened index because dispatch normalization can reshape the workgroup grid. + let row = workgroup_idx / groups_per_row; + // Reject padded rows uniformly before buffer accesses or workgroup barriers. + if (row >= uniforms.M) { + return; + } + let n_vec = (workgroup_idx % groups_per_row) * workgroup_size_x + local_id.x; var acc = vec4(0); if (n_vec < n_vec_count) { for (var k = local_id.y; k < uniforms.K; k += workgroup_size_y) { - let av = f32(a.getByOffset(k)); + let av = f32(a.getByOffset(row * uniforms.K + k)); let bv = vec4(b.getByOffset(k * n_vec_count + n_vec)); acc = fma(vec4(av), bv, acc); } @@ -27,6 +34,6 @@ $MAIN { workgroupBarrier(); } if (local_id.y == 0u && n_vec < n_vec_count) { - output.setByOffset(n_vec, output_value_t(partials[local_idx])); + output.setByOffset(row * n_vec_count + n_vec, output_value_t(partials[local_idx])); } } diff --git a/onnxruntime/core/providers/webgpu/math/matmul_packed.cc b/onnxruntime/core/providers/webgpu/math/matmul_packed.cc index ce7c663bf0d9f..20a753d60fd37 100644 --- a/onnxruntime/core/providers/webgpu/math/matmul_packed.cc +++ b/onnxruntime/core/providers/webgpu/math/matmul_packed.cc @@ -47,12 +47,14 @@ Status MatMulProgram::GenerateShaderCode(ShaderHelper& shader) const { ORT_RETURN_IF_ERROR(MakeMatMulPackedVec4Source( shader, elements_per_thread_, WorkgroupSizeX(), WorkgroupSizeY(), data_type, &batch_dims, /*transA = */ false, /*transB = */ false, /*alpha = */ 1.f, /*need_handle_matmul = */ true, - /*output_components = */ 4, /*tile_inner = */ 32, need_split_k, split_dim_inner_, + /*output_components = */ 4, tile_inner_, need_split_k, split_dim_inner_, use_f32_accumulation_)); } else { ORT_RETURN_IF_ERROR(MakeMatMulPackedSource( shader, elements_per_thread_, WorkgroupSizeX(), WorkgroupSizeY(), data_type, &batch_dims, - false, false, 1.f, true, 32, need_split_k, split_dim_inner_, use_f32_accumulation_)); + /*transpose_a = */ false, /*transpose_b = */ false, /*alpha = */ 1.f, + /*need_handle_matmul = */ true, tile_inner_, need_split_k, split_dim_inner_, + use_f32_accumulation_)); } return Status::OK(); } diff --git a/onnxruntime/core/providers/webgpu/math/matmul_packed.h b/onnxruntime/core/providers/webgpu/math/matmul_packed.h index d1903b1d65905..1c8cd5d1f32a2 100644 --- a/onnxruntime/core/providers/webgpu/math/matmul_packed.h +++ b/onnxruntime/core/providers/webgpu/math/matmul_packed.h @@ -13,9 +13,14 @@ namespace onnxruntime { namespace webgpu { class MatMulProgram final : public Program { public: - MatMulProgram(const Activation& activation, bool bias, bool is_vec4, - const gsl::span& elements_per_thread, bool is_channels_last = false, - uint32_t split_dim_inner = 1, bool use_f32_accumulation = false) + MatMulProgram(const Activation& activation, + bool bias, + bool is_vec4, + const gsl::span& elements_per_thread, + bool is_channels_last = false, + uint32_t split_dim_inner = 1, + uint32_t tile_inner = 32, + bool use_f32_accumulation = false) : Program{"MatMul"}, activation_(activation), has_bias_{bias}, @@ -23,6 +28,7 @@ class MatMulProgram final : public Program { elements_per_thread_(elements_per_thread.begin(), elements_per_thread.end()), is_channels_last_(is_channels_last), split_dim_inner_(split_dim_inner), + tile_inner_(tile_inner), use_f32_accumulation_(use_f32_accumulation) {} Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -44,6 +50,7 @@ class MatMulProgram final : public Program { const InlinedVector elements_per_thread_; bool is_channels_last_ = false; uint32_t split_dim_inner_ = 1; + uint32_t tile_inner_ = 32; bool use_f32_accumulation_ = false; }; diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc index 3b2c75a9f1391..aa17baeb2bae1 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc @@ -54,8 +54,8 @@ class SubgroupMatrixMatMulProgram final : public Program= N), zero-filling columns [N, N_b). Gives B an even row stride so the -// subgroup-matrix f16 load's 4-byte row-start alignment holds for odd N. +// (N_b >= N), zero-filling columns [N, N_b). Gives B an aligned row stride so the +// subgroup-matrix f16 load's row-start alignment holds for any N. class SubgroupMatrixMatMulPadBProgram final : public Program { public: SubgroupMatrixMatMulPadBProgram() : Program{"SubgroupMatrixMatMulPadB"} {} @@ -73,24 +73,19 @@ class SubgroupMatrixMatMulPadBProgram final : public Program& inputs, - Tensor* output, - const Activation& activation, - bool is_channels_last, - bool b_is_constant, - /*out*/ bool& handled) override { - handled = false; - + bool CanApply(const ComputeContext& context, + const std::vector& inputs, + bool is_channels_last, + bool b_is_constant) const override { const auto* a = inputs[0]; const auto* b = inputs[1]; const auto& a_shape = a->Shape(); @@ -100,72 +95,129 @@ class SubgroupMatrixMatMulImpl final : public MatMulOptImpl { const size_t b_rank = b_shape.NumDimensions(); if ((!is_channels_last && has_bias) || a_rank < 2 || b_rank < 2 || !a->IsDataType() || !b->IsDataType()) { - return Status::OK(); + return false; } const uint32_t K = narrow(a_shape[a_rank - 1]); if (K == 0) { - return Status::OK(); + return false; } uint32_t M = 0; uint32_t N = 0; uint32_t batch = 1; if (b_rank == 2) { - ORT_ENFORCE(narrow(b_shape[0]) == K, - "MatMul contraction dim mismatch: A K=", K, " vs B rows=", b_shape[0]); + if (narrow(b_shape[0]) != K) { + return false; + } M = narrow(a_shape.Size() / static_cast(K)); N = narrow(b_shape[1]); } else { - if (a_rank != b_rank) { - return Status::OK(); + if (a_rank != b_rank || narrow(b_shape[b_rank - 2]) != K) { + return false; } - ORT_ENFORCE(narrow(b_shape[b_rank - 2]) == K, - "MatMul contraction dim mismatch: A K=", K, - " vs B rows=", b_shape[b_rank - 2]); M = narrow(a_shape[a_rank - 2]); N = narrow(b_shape[b_rank - 1]); for (size_t i = 0; i + 2 < a_rank; ++i) { if (a_shape[i] != b_shape[i]) { - return Status::OK(); + return false; } } batch = narrow(a_shape.SizeToDimension(a_rank - 2)); } if (M == 0 || N == 0) { - return Status::OK(); + return false; } const std::optional tiling = tiling_selector_(context, M, N, K, batch); if (!tiling) { - return Status::OK(); + return false; } const auto& config = config_; - const bool needs_padded_b = N % 2 != 0; - // Require whole subgroup-matrix K blocks. An odd-width B must be constant - // because its padded copy is cached by this implementation. - if (config.K == 0 || K % config.K != 0 || - (needs_padded_b && !b_is_constant)) { - return Status::OK(); + const uint32_t n_align = config.Is(16, 16, 16) ? 8 : 2; + const bool needs_padded_b = N % n_align != 0; + if (needs_padded_b && (!b_is_constant || N > std::numeric_limits::max() - n_align)) { + return false; + } + const uint32_t n_b = needs_padded_b ? (N / n_align + 1) * n_align : N; + return config.K != 0 && K % config.K == 0 && + M >= tiling->tile_m && n_b >= tiling->tile_n; + } + + Status Compute(ComputeContext& context, + const std::vector& inputs, + Tensor* output, + const Activation& activation, + bool is_channels_last, + bool b_is_constant) override { + ORT_RETURN_IF_NOT(inputs.size() >= 2, + "MatMul algorithm subgroup_matrix requires two inputs."); + const auto* a = inputs[0]; + const auto* b = inputs[1]; + const auto& a_shape = a->Shape(); + const auto& b_shape = b->Shape(); + const bool has_bias = inputs.size() > 2; + const size_t a_rank = a_shape.NumDimensions(); + const size_t b_rank = b_shape.NumDimensions(); + ORT_RETURN_IF_NOT((is_channels_last || !has_bias) && a_rank >= 2 && b_rank >= 2 && + a->IsDataType() && b->IsDataType(), + "MatMul algorithm subgroup_matrix does not support these inputs."); + const uint32_t K = narrow(a_shape[a_rank - 1]); + ORT_RETURN_IF_NOT(K != 0, + "MatMul algorithm subgroup_matrix requires a nonzero contraction dimension."); + + uint32_t M = 0; + uint32_t N = 0; + uint32_t batch = 1; + if (b_rank == 2) { + ORT_ENFORCE(narrow(b_shape[0]) == K, + "MatMul contraction dim mismatch: A K=", K, " vs B rows=", b_shape[0]); + M = narrow(a_shape.Size() / static_cast(K)); + N = narrow(b_shape[1]); + } else { + ORT_ENFORCE(narrow(b_shape[b_rank - 2]) == K, + "MatMul contraction dim mismatch: A K=", K, + " vs B rows=", b_shape[b_rank - 2]); + M = narrow(a_shape[a_rank - 2]); + N = narrow(b_shape[b_rank - 1]); + for (size_t i = 0; i + 2 < a_rank; ++i) { + ORT_ENFORCE(a_shape[i] == b_shape[i]); + } + batch = narrow(a_shape.SizeToDimension(a_rank - 2)); } + ORT_RETURN_IF_NOT(M != 0 && N != 0, + "MatMul algorithm subgroup_matrix requires nonzero output dimensions."); - // N_b is just N rounded up to even - compute it before doing any padding work so + const std::optional tiling = tiling_selector_(context, M, N, K, batch); + ORT_RETURN_IF_NOT(tiling.has_value(), + "MatMul algorithm subgroup_matrix has no tiling for these inputs."); + + const auto& config = config_; + // Vulkan cooperative matrix loads need 16 byte aligned rows so the + // 16x16x16 config pads B to a multiple of 8 columns. + const uint32_t n_align = config.Is(16, 16, 16) ? 8 : 2; + const bool needs_padded_b = N % n_align != 0; + // Require whole subgroup-matrix K blocks. An unaligned B must be constant + // because its padded copy is cached by this implementation. + ORT_RETURN_IF_NOT(config.K != 0 && K % config.K == 0 && + (!needs_padded_b || b_is_constant), + "MatMul algorithm subgroup_matrix requires aligned K and a constant unaligned-N weight."); + + // N_b is just N rounded up to a multiple of n_align. Compute it before padding so // the tile-fit check below can bail out without a wasted pad dispatch. uint32_t N_b = N; if (needs_padded_b) { - ORT_RETURN_IF_NOT(N < std::numeric_limits::max(), - "Cannot pad odd-N B because N+1 exceeds uint32_t range."); - N_b = N + 1; + ORT_RETURN_IF_NOT(N <= std::numeric_limits::max() - n_align, + "Cannot pad B because its padded width exceeds uint32_t range."); + N_b = (N / n_align + 1) * n_align; } // The kernel keeps its operand loads in bounds by shifting a trailing partial // tile back, which is only possible when the tile fits within M and N. - if (M < tiling->tile_m || N_b < tiling->tile_n) { - return Status::OK(); - } + ORT_ENFORCE(M >= tiling->tile_m && N_b >= tiling->tile_n); - // The optimized path will run: now materialize the even-strided B for odd N. + // The optimized path will run: now materialize the aligned B for unaligned N. const Tensor* b_used = b; if (needs_padded_b) { ORT_RETURN_IF_ERROR(EnsurePaddedB(context, *b, b_shape, N, N_b)); @@ -208,13 +260,12 @@ class SubgroupMatrixMatMulImpl final : public MatMulOptImpl { } ORT_RETURN_IF_ERROR(context.RunProgram(program)); - handled = true; return Status::OK(); } private: - // Lazily builds an even-strided copy of a constant weight B [..., K, N] with odd N - // by widening its last dim to N_b = N + 1 (zero-filling the extra column) and + // Lazily builds an aligned copy of a constant weight B [..., K, N] with unaligned N + // by widening its last dim to an aligned N_b (zero-filling the extra columns) and // caches it, so the per-run pad cost is paid once. Works for a 2D weight [K, N] // and a batched weight [batch, K, N] alike: the pad pass treats B as a flat // [rows, N] -> [rows, N_b] copy over rows = numel / N (= K, or batch*K), which is @@ -232,7 +283,7 @@ class SubgroupMatrixMatMulImpl final : public MatMulOptImpl { const TensorShape padded_shape{padded_dims}; const int64_t output_size_i64 = padded_shape.Size(); ORT_RETURN_IF_NOT(output_size_i64 <= static_cast(std::numeric_limits::max()), - "Cannot pad odd-N B because the padded tensor has ", output_size_i64, + "Cannot pad B because the padded tensor has ", output_size_i64, " elements, exceeding uint32_t shader indexing range."); const uint32_t output_size = narrow(output_size_i64); @@ -258,35 +309,39 @@ class SubgroupMatrixMatMulImpl final : public MatMulOptImpl { // padded_b_ persists the outcome across calls: call_once runs the body only on // the first call, so a failed pad stays failed (and null) on later calls. return padded_b_ ? Status::OK() - : ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to pad odd-N B for subgroup-matrix MatMul."); + : ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to pad B for subgroup-matrix MatMul."); } const SubgroupMatrixConfig config_; SubgroupMatrixTilingSelector tiling_selector_; - // Cached even-strided B for odd N; built once by EnsurePaddedB. + // Cached aligned B for unaligned N; built once by EnsurePaddedB. mutable std::once_flag pad_once_; mutable std::unique_ptr padded_b_; }; -Status GenerateShaderCode8x16x16(ShaderHelper& shader, - uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, - uint32_t split_k) { - return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul_8x16x16.wgsl.template", +Status GenerateSubgroupMatrixShaderCode(ShaderHelper& shader, const SubgroupMatrixConfig& config, + uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, + uint32_t split_k) { + return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul.wgsl.template", WGSL_TEMPLATE_PARAMETER(sg_mat_count_m, sg_mat_count_m), WGSL_TEMPLATE_PARAMETER(sg_mat_count_n, sg_mat_count_n), - WGSL_TEMPLATE_PARAMETER(sg_mat_k, 16), - WGSL_TEMPLATE_PARAMETER(sg_mat_m, 8), - WGSL_TEMPLATE_PARAMETER(sg_mat_n, 16), + WGSL_TEMPLATE_PARAMETER(sg_mat_k, config.K), + WGSL_TEMPLATE_PARAMETER(sg_mat_m, config.M), + WGSL_TEMPLATE_PARAMETER(sg_mat_n, config.N), WGSL_TEMPLATE_PARAMETER(split_k, split_k)); } -// Default tiling used on any vendor without a specialized policy: a fixed 32x32 -// output tile with no split-K. +// Default tiling used on any vendor without a specialized policy: a 64x64 output +// tile split two ways along K when M and N are at least 64 and otherwise a 32x32 +// output tile split four ways along K. SubgroupMatrixTilingSelector MakeDefaultTilingSelector() { - return [](const ComputeContext&, uint32_t /*M*/, uint32_t /*N*/, + return [](const ComputeContext&, uint32_t M, uint32_t N, uint32_t /*K*/, uint32_t /*batch*/) -> std::optional { - return SubgroupMatrixTiling{32, 32, 1}; + if (M >= 64 && N >= 64) { + return SubgroupMatrixTiling{64, 64, 2}; + } + return SubgroupMatrixTiling{32, 32, 4}; }; } @@ -313,30 +368,31 @@ Status SubgroupMatrixMatMulProgram::GenerateShaderCode(ShaderHelper& shader) con << " " << output.SetByOffset("output_offset", "value") << "\n" << "}\n"; - if (config_.Is(8, 16, 16)) { - return GenerateShaderCode8x16x16(shader, sg_mat_count_m_, sg_mat_count_n_, split_k_); + if (config_.Is(8, 16, 16) || config_.Is(16, 16, 16)) { + return GenerateSubgroupMatrixShaderCode(shader, config_, sg_mat_count_m_, sg_mat_count_n_, split_k_); } return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::NOT_IMPLEMENTED, "Unsupported subgroup matrix config dimensions."); } -std::unique_ptr CreateSubgroupMatrixMatMulImpl(const ComputeContextBase& context) { - // Only run on devices that report the 8x16x16 F16 subgroup-matrix config this - // kernel is implemented for and can provide its required subgroup size. +std::unique_ptr CreateSubgroupMatrixMatMulImpl(const ComputeContextBase& context) { + // Only run on devices that report an F16 subgroup-matrix config this kernel + // implements (8x16x16 or 16x16x16) and can provide its required subgroup size. constexpr auto kF16 = wgpu::SubgroupMatrixComponentType::F16; - const auto config = SelectSubgroupMatrixConfig(context, {{kF16, kF16, 8, 16, 16, 32, false}}); + const auto config = SelectSubgroupMatrixConfig(context, {{kF16, kF16, 8, 16, 16, 32, false}, + {kF16, kF16, 16, 16, 16, 32, false}}); if (!config) { return nullptr; } // Intel GPUs use a tuned/heuristic tiling policy; every other vendor falls back - // to a fixed default tiling. + // to a default tiling. const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"}; SubgroupMatrixTilingSelector tiling_selector = is_intel ? intel::CreateSubgroupMatrixTilingSelector(context) : MakeDefaultTilingSelector(); if (!tiling_selector) { return nullptr; } - return std::make_unique(*config, std::move(tiling_selector)); + return std::make_unique(*config, std::move(tiling_selector)); } } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_8x16x16.wgsl.template b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template similarity index 97% rename from onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_8x16x16.wgsl.template rename to onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template index 8d31e9e742482..a41c869cefd9a 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_8x16x16.wgsl.template +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.wgsl.template @@ -1,10 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. // -// Subgroup-matrix MatMul kernel (F16) for the 8x16x16 (sg_mat_m x sg_mat_n x -// sg_mat_k) subgroup-matrix shape. The sg_mat_m/n/k values are still supplied as -// params, but this kernel is currently only validated for 8/16/16; other -// subgroup-matrix shapes are TBD and should get their own template variant. +// Subgroup-matrix MatMul kernel (F16) for the 8x16x16 and 16x16x16 (sg_mat_m x +// sg_mat_n x sg_mat_k) subgroup-matrix shapes. The sg_mat_m/n/k values are still +// supplied as params, but this kernel is only validated for 8/16/16 and 16/16/16. +// Other subgroup-matrix shapes are TBD and should get their own template variant. // // Computes Y[MxN] = A[MxK] @ B[KxN] using subgroupMatrixMultiplyAccumulate. // @@ -13,8 +13,8 @@ // B: weight, loaded directly from global memory in plain row-major KxN layout as a // right operand (subgroup_matrix_right) with row stride // N_b. N_b is the padded row stride of the (possibly column-padded) B buffer and -// equals N unless the host padded B to an even stride to satisfy the load's -// 4-byte row-start alignment (see the odd-N handling in subgroup_matrix_matmul.cc). +// equals N unless the host padded B to an aligned stride to satisfy the load's +// row-start alignment (see the B padding in subgroup_matrix_matmul.cc). // // Workgroup: split_k subgroups x 32 lanes (split_k in {1,2,4,8}). // Tile: kTileM x kTileN, chosen adaptively by the host config provider. Each @@ -24,7 +24,7 @@ // tile in shared memory, and the partials are summed at write-out. // // Preconditions enforced by the host: K % sg_mat_k == 0; B's row stride N_b is -// even (the load requires 4-byte-aligned row starts); and the tile fits the +// padded to the row-start alignment the load requires and the tile fits the // matrix (kTileM <= M, kTileN <= N - the host falls back otherwise). M and N may // otherwise be any size: a trailing partial tile is shifted back to end exactly // at M / N_b so no operand load runs past the end of A or B, and the rows/columns @@ -84,9 +84,9 @@ $MAIN { let n_tile_start = n_tile * kTileN; let m_tile_start = m_tile * kTileM; // Shift a trailing partial tile back so its operand loads stay inside A and B. - // The N shift uses B's padded stride N_b so the shifted origin stays even (the - // load needs 4-byte-aligned row starts); the extra column is clipped at - // write-out. The max() only guards the u32 subtraction; the host guarantees + // The N shift uses B's padded stride N_b so the shifted origin stays aligned (the + // load needs aligned row starts). The extra columns are clipped at write-out. + // The max() only guards the u32 subtraction; the host guarantees // the tile fits. let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM); diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template index f5d0d6996bced..8932768357e2b 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template @@ -3,8 +3,8 @@ // Copies a row-major f16 weight B [K, N] into a column-padded [K, N_b] buffer // (N_b >= N), zero-filling columns [N, N_b). One thread per padded-output element. -// Gives B an even row stride so the subgroup-matrix f16 load's 4-byte row-start -// alignment holds for odd N. See EnsurePaddedB in subgroup_matrix_matmul.cc. +// Gives B an aligned row stride so the subgroup-matrix f16 load's row-start +// alignment holds for any N. See EnsurePaddedB in subgroup_matrix_matmul.cc. #use guardAgainstOutOfBoundsWorkgroupSizes #use .getByOffset .setByOffset diff --git a/onnxruntime/core/providers/webgpu/nn/conv.cc b/onnxruntime/core/providers/webgpu/nn/conv.cc index 6ffb91783381c..6e7fb6f80a4d1 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv.cc @@ -9,7 +9,6 @@ #include "core/providers/webgpu/tensor/transpose.h" #include "core/providers/webgpu/nn/grouped_conv.h" #include "core/providers/webgpu/webgpu_utils.h" -#include "core/providers/webgpu/math/matmul.h" namespace onnxruntime { namespace webgpu { @@ -282,8 +281,8 @@ Status Conv::ComputeInternal(ComputeContext& context if (has_bias) { matmul_inputs.push_back(bias); } - return ComputeMatMul(&context, activation_, matmul_inputs, output, is_channels_last, - matmul_compute_cache_, matmul_b_is_constant); + return matmul_compute_dispatcher_.Compute(context, activation_, matmul_inputs, output, is_channels_last, + matmul_b_is_constant); } // Transpose weights when necessary Tensor transposed_kernel; diff --git a/onnxruntime/core/providers/webgpu/nn/conv.h b/onnxruntime/core/providers/webgpu/nn/conv.h index a64a206f1ef34..af98fa76e30ca 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.h +++ b/onnxruntime/core/providers/webgpu/nn/conv.h @@ -8,7 +8,7 @@ #include "core/providers/cpu/nn/conv_attributes.h" #include "core/providers/webgpu/program.h" #include "core/providers/webgpu/shader_helper.h" -#include "core/providers/webgpu/math/matmul.h" +#include "core/providers/webgpu/math/matmul_compute_dispatcher.h" #include "core/providers/webgpu/nn/fuse_utils.h" namespace onnxruntime { @@ -29,7 +29,8 @@ class Conv : public WebGpuKernel { public: Conv(const OpKernelInfo& info) : WebGpuKernel(info), conv_attrs_(info) { if (is_fused) { - ORT_ENFORCE(GetFusedActivationAttr(info, activation_).IsOK()); + // Throw the Status, not the expression: the message names the activation ORT_ENFORCE would drop. + ORT_THROW_IF_ERROR(GetFusedActivationAttr(info, activation_)); } } Status ComputeInternal(ComputeContext& context) const override; @@ -49,7 +50,7 @@ class Conv : public WebGpuKernel { // Layout of the tensor ComputeInternal ends up consuming -- `prepacked_kernel_` when it // is set, otherwise input 1. Stays `OIHW` while `prepacked_kernel_` is null. KernelLayout kernel_layout_{KernelLayout::OIHW}; - mutable MatMulOptImplCache matmul_compute_cache_; + mutable MatMulComputeDispatcher matmul_compute_dispatcher_; }; Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const TensorShape& kernel_shape, Tensor* transposed_kernel, const InlinedVector& perm); diff --git a/onnxruntime/core/providers/webgpu/nn/conv2d_mm.cc b/onnxruntime/core/providers/webgpu/nn/conv2d_mm.cc index fd6a2627cea66..b5aa00dc3d225 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv2d_mm.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv2d_mm.cc @@ -40,7 +40,8 @@ std::string Conv2dMMProgram::Conv2dCommonSnippet(const ShaderVariableHelper& x, ORT_THROW("inner_element_size ", inner_element_size, " is not supported."); } }; - const std::string coord_a_snippet = is_channels_last_ ? "let coord = vec4(batch, xRow, xCol, xCh / " + std::to_string(inner_element_size_x == 3 ? 4 : inner_element_size_x) + ");" : "let coord = vec4(batch, xCh, xRow, xCol);"; + // The vec3 path assembles three scalar loads, so its channel coordinate is not packed. + const std::string coord_a_snippet = is_channels_last_ ? "let coord = vec4(batch, xRow, xCol, xCh / " + std::to_string(inner_element_size_x == 3 ? 1 : inner_element_size_x) + ");" : "let coord = vec4(batch, xCh, xRow, xCol);"; const std::string coord_res_snippet = is_channels_last_ ? "let coords = vec4(batch, row / outWidth, row % outWidth, col / " + std::to_string(inner_element_size) + ");" : "let coords = vec4(batch, row, col / outWidth, col % outWidth);"; const std::string xHeight = is_channels_last_ ? "i32(uniforms.x_shape[1])" : "i32(uniforms.x_shape[2])"; diff --git a/onnxruntime/core/providers/webgpu/nn/conv_backprop.cc b/onnxruntime/core/providers/webgpu/nn/conv_backprop.cc index 779e21125e3f2..ee61079fd8c80 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv_backprop.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv_backprop.cc @@ -67,7 +67,8 @@ Status ConvTranspose2DProgram::GenerateShaderCode(ShaderHelper& shader) const { ss << "let wValue = " << w.GetByIndices("w_indices_t(u32(wRPerm), u32(wCPerm), inputChannel, wOutChannel / " + std::to_string(b_components_) + ")") << ";\n" << "dotProd = dotProd + xValue * wValue;\n"; } else if (a_components_ == b_components_ && components_ == 1) { - ss << "let wValue = " << w.GetByIndices("w_indices_t(u32(wRPerm), u32(wCPerm), inputChannel, wOutChannel)") << ";\n" + // With one output channel per group, weight vectors span input channels. + ss << "let wValue = " << w.GetByIndices("w_indices_t(u32(wRPerm), u32(wCPerm), inputChannel / " + std::to_string(b_components_) + ", wOutChannel)") << ";\n" << "dotProd = dotProd + dot(xValue, wValue);\n"; } else { for (uint32_t i = 0; i < a_components_; ++i) { @@ -146,7 +147,7 @@ Status ConvTranspose2DProgram::GenerateShaderCode(ShaderHelper& shader) const { << " var inputChannel = groupId * uniforms.input_channels_per_group;\n"; if (pack_input_as4_) { shader.MainFunctionBody() << " let dy_indices = dy_indices_t(batch, idyR, idyC, inputChannel / " << a_components_ << ");\n" - << " let w_indices = w_indices_t(u32(wRPerm), u32(wCPerm), inputChannel, wOutChannel / " << b_components_ << ");\n" + << " let w_indices = w_indices_t(u32(wRPerm), u32(wCPerm), inputChannel / " << b_components_ << ", wOutChannel);\n" << " var x_offset = " << dy.IndicesToOffset("dy_indices") << ";\n" << " var w_offset = " << w.IndicesToOffset("w_indices") << ";\n"; } diff --git a/onnxruntime/core/providers/webgpu/nn/conv_transpose.cc b/onnxruntime/core/providers/webgpu/nn/conv_transpose.cc index afa56d9f5263d..26ba673890038 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv_transpose.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv_transpose.cc @@ -19,12 +19,8 @@ Status ConvTranspose::ComputeInternal(ComputeContext& context) const auto* input = context.Input(0); const auto* filter = prepacked_filter_ ? prepacked_filter_.get() : context.Input(1); TensorShape input_shape = input->Shape(); - TensorShape filter_shape = filter->Shape(); - const bool is_prepacked = weight_layout_ == WeightLayout::DHWOI; - if (is_prepacked) { - // Recover the logical ONNX shape before validating channels or inferring output dimensions. - filter_shape = TensorShape{filter_shape[4], filter_shape[3], filter_shape[0], filter_shape[1], filter_shape[2]}; - } + const bool is_prepacked = weight_layout_ != WeightLayout::ONNX; + TensorShape filter_shape = is_prepacked ? original_filter_shape_ : filter->Shape(); const auto rank = input_shape.NumDimensions(); if (rank < 3) { @@ -56,7 +52,6 @@ Status ConvTranspose::ComputeInternal(ComputeContext& context) return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input channels is not divisible by group."); } - const InlinedVector perm = {2, 3, 0, 1}; TensorShapeVector local_output_padding(conv_transpose_attrs_.output_padding.begin(), conv_transpose_attrs_.output_padding.end()); ConvAttributes::ConvPadVector local_pads(conv_transpose_attrs_.pads.begin(), conv_transpose_attrs_.pads.end()); TensorShapeVector local_dilations(conv_transpose_attrs_.dilations.begin(), conv_transpose_attrs_.dilations.end()); @@ -159,12 +154,15 @@ Status ConvTranspose::ComputeInternal(ComputeContext& context) strides.insert(strides.begin(), 1); dilations.insert(dilations.begin(), 1); } - // Transpose weights Tensor transposed_filter; - ORT_RETURN_IF_ERROR(TransposeKernel(context, filter, filter_shape, &transposed_filter, perm)); + if (!is_prepacked) { + const InlinedVector perm{2, 3, 0, 1}; + ORT_RETURN_IF_ERROR(TransposeKernel(context, filter, filter_shape, &transposed_filter, perm)); + filter = &transposed_filter; + } TensorShape output_shape(output_shape_vector); - TensorShape transposed_filter_shape = transposed_filter.Shape(); - std::vector inputs = {input, &transposed_filter}; + TensorShape transposed_filter_shape = filter->Shape(); + std::vector inputs = {input, filter}; std::vector input_output_shapes = {input_shape, transposed_filter_shape}; if (has_bias) { inputs.push_back(bias); @@ -185,7 +183,8 @@ Status ConvTranspose::PrePackInternal(ComputeContextBase& cont AllocatorPtr alloc, /*out*/ bool& is_packed) { is_packed = false; - if (input_idx != 1 || tensor.Shape().NumDimensions() != 5 || tensor.Shape().Size() == 0) { + const auto rank = tensor.Shape().NumDimensions(); + if (input_idx != 1 || rank < 3 || rank > 5 || tensor.Shape().Size() == 0) { return Status::OK(); } @@ -196,13 +195,27 @@ Status ConvTranspose::PrePackInternal(ComputeContextBase& cont return Status::OK(); } - const InlinedVector perm{2, 3, 4, 1, 0}; - const auto& shape = tensor.Shape(); - const TensorShape packed_shape{shape[2], shape[3], shape[4], shape[1], shape[0]}; + TensorShapeVector unpacked_shape = tensor.Shape().AsShapeVector(); + InlinedVector perm{2, 3, 0, 1}; + auto packed_layout = WeightLayout::HWIO; + if (rank == 3) { + // The 1D compute path uses the 2D shader with a synthetic height of one. + unpacked_shape.insert(unpacked_shape.begin() + 2, 1); + } else if (rank == 5) { + perm = {2, 3, 4, 1, 0}; + packed_layout = WeightLayout::DHWOI; + } + TensorShapeVector packed_shape; + packed_shape.reserve(perm.size()); + for (const auto axis : perm) { + packed_shape.push_back(unpacked_shape[axis]); + } auto packed = std::make_unique(tensor.DataType(), packed_shape, alloc); - ORT_RETURN_IF_ERROR(Transpose::DoTranspose(context, perm, tensor, *packed)); + const Tensor unpacked = CreateTensorView(tensor, unpacked_shape); + ORT_RETURN_IF_ERROR(Transpose::DoTranspose(context, perm, unpacked, *packed)); + original_filter_shape_ = tensor.Shape(); prepacked_filter_ = std::move(packed); - weight_layout_ = WeightLayout::DHWOI; + weight_layout_ = packed_layout; is_packed = true; return Status::OK(); } diff --git a/onnxruntime/core/providers/webgpu/nn/conv_transpose.h b/onnxruntime/core/providers/webgpu/nn/conv_transpose.h index a4e1901ea395d..554c624747751 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv_transpose.h +++ b/onnxruntime/core/providers/webgpu/nn/conv_transpose.h @@ -29,12 +29,14 @@ class ConvTranspose final : public WebGpuKernel { private: enum class WeightLayout { - IODHW, // ONNX weights: [C_in, C_out/group, kD, kH, kW]. - DHWOI, // Prepacked weights: [kD, kH, kW, C_out/group, C_in]. + ONNX, // Unpacked weights: [C_in, C_out/group, spatial kernel dimensions...]. + HWIO, // Prepacked 1D/2D weights: [kH, kW, C_in, C_out/group]; kH is 1 for 1D. + DHWOI, // Prepacked 3D weights: [kD, kH, kW, C_out/group, C_in]. }; std::unique_ptr prepacked_filter_; - WeightLayout weight_layout_{WeightLayout::IODHW}; + TensorShape original_filter_shape_; + WeightLayout weight_layout_{WeightLayout::ONNX}; }; } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/nn/fuse_utils.cc b/onnxruntime/core/providers/webgpu/nn/fuse_utils.cc index b8137aa1f495a..19c2a4c403d14 100644 --- a/onnxruntime/core/providers/webgpu/nn/fuse_utils.cc +++ b/onnxruntime/core/providers/webgpu/nn/fuse_utils.cc @@ -45,6 +45,19 @@ size_t GetActivationUsedUniformCount(const Activation& activation) { } } +// ConvActivationFusion writes the attribute from the core library while the table above lives in +// the EP, so separately versioned binaries can disagree on the activation set. Name that cause: +// the bare activation name alone sends people looking for a bug in their model. +Status UnsupportedActivationError(const std::string& activation_type) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "The WebGPU EP cannot fuse activation '", activation_type, + "'. The attribute comes from ConvActivationFusion in the onnxruntime core " + "library; if the core library and the WebGPU EP binary are from different " + "builds, core can emit an activation this EP does not implement. Either " + "match the two versions, or create the session with graph_optimization_level " + "ORT_ENABLE_BASIC so the activation stays a separate node."); +} + } // namespace Status GetFusedActivationAttr(const OpKernelInfo& info, Activation& activation) { @@ -97,7 +110,7 @@ Status GetFusedActivationAttr(const OpKernelInfo& info, Activation& activation) activation.activation_kind_ = ActivationKind::ThresholdedRelu; activation_params_count = 1; } else { - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "unimplemented activation: " + activation_type); + return UnsupportedActivationError(activation_type); } std::vector activation_params; diff --git a/onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc b/onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc index c068928b526be..563804b4d7135 100644 --- a/onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc +++ b/onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc @@ -39,6 +39,23 @@ std::pair ChooseTileSize(uint32_t im2col_m, uint32_t im2col_ return kTileSizes.back(); } +uint32_t GetFixedSubgroupSize(const ComputeContextBase& context) { + if (!context.HasFeature(wgpu::FeatureName::Subgroups) || + !context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { + return 0; + } + + // The subgroup path requires a fixed subgroup size >= tile_m (16 or 32). 32 is the best for both + // tilings. + constexpr uint32_t kFixedSubgroupSize = 32; + const auto& adapter_info = context.AdapterInfo(); + if (adapter_info.subgroupMinSize <= kFixedSubgroupSize && kFixedSubgroupSize <= adapter_info.subgroupMaxSize) { + return kFixedSubgroupSize; + } + + return 0; +} + // Add support for more devices. bool IsDeviceSupported(const ComputeContextBase& context) { const wgpu::AdapterInfo& adapter_info = context.AdapterInfo(); @@ -173,15 +190,14 @@ Status ApplyIm2ColMatMulProgram(ComputeContext& context, const auto [tile_m, tile_n] = ChooseTileSize(im2col_m, im2col_n); const uint32_t workgroup_size = tile_n; - // Check the device's subgroup size before shader compilation to avoid potential performance penalties - // associated with conditional checks in the shader runtime. - // - // Ensure the subgroup size must be greater than or equal to `tile_m` to safely enable `use_subgroup`. - // If the status of this condition is uncertain, the feature must be disabled. - const bool use_subgroup = false; + const uint32_t subgroup_size = GetFixedSubgroupSize(context); + const bool use_subgroup = subgroup_size > 0; const uint32_t vec_size = channel_input % 4 == 0 ? 4 : (channel_input % 2 == 0 ? 2 : 1); Im2ColMatMulProgram im2col_mm_program{has_bias, tile_m, tile_n, vec_size, use_subgroup, activation}; im2col_mm_program.SetWorkgroupSize(workgroup_size); + if (use_subgroup) { + im2col_mm_program.SetSubgroupSize(subgroup_size); + } const uint32_t M_tiles = CeilDiv(im2col_m, tile_m); const uint32_t N_tiles = CeilDiv(im2col_n, tile_n); diff --git a/onnxruntime/core/providers/webgpu/nn/im2col_matmul.wgsl.template b/onnxruntime/core/providers/webgpu/nn/im2col_matmul.wgsl.template index 132f0cc0b4ea0..df71ed6e7e185 100644 --- a/onnxruntime/core/providers/webgpu/nn/im2col_matmul.wgsl.template +++ b/onnxruntime/core/providers/webgpu/nn/im2col_matmul.wgsl.template @@ -125,9 +125,14 @@ $MAIN { for (var inner_k_idx = 0u; inner_k_idx < TILE_K_VEC_SIZE; inner_k_idx++) { let weight_data = weight_tile[inner_k_idx][local_idx]; #if use_subgroup - let src_data = src_tile[inner_k_idx][sg_id]; + // The subgroup size is fixed at 32; wrap the index so the 16x64 tiling does not read out of bounds. + let src_data = src_tile[inner_k_idx][sg_id % TILE_M_SIZE]; for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { +#if vec_size == 1 + results[m_idx] += output_element_t(weight_data * subgroupShuffle(src_data, m_idx)); +#else results[m_idx] += output_element_t(dot(weight_data, subgroupShuffle(src_data, m_idx))); +#endif } #else for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { diff --git a/onnxruntime/core/providers/webgpu/nn/instance_norm.cc b/onnxruntime/core/providers/webgpu/nn/instance_norm.cc index bc3fe538346ca..47651f7d7b790 100644 --- a/onnxruntime/core/providers/webgpu/nn/instance_norm.cc +++ b/onnxruntime/core/providers/webgpu/nn/instance_norm.cc @@ -73,7 +73,8 @@ Status ComputeChannelScaleAndShift(ComputeContext& context, const Tensor* input, TensorShape reduced_output_shape(reduced_output_shape_vector); *output = context.CreateGPUTensor(input->DataType(), output_shape); ComputeChannelScaleShiftProgram program = ComputeChannelScaleShiftProgram(components, epsilon, workgroup_size); - program.CacheHint(components, units_of_work) + // Match the epsilon literal embedded in GenerateShaderCode. + program.CacheHint(components, units_of_work, std::to_string(epsilon)) .AddInputs({{input, ProgramTensorMetadataDependency::TypeAndRank, reduced_input_shape, components}, {scale, ProgramTensorMetadataDependency::TypeAndRank}, {bias, ProgramTensorMetadataDependency::TypeAndRank}}) diff --git a/onnxruntime/core/providers/webgpu/nn/layer_norm.cc b/onnxruntime/core/providers/webgpu/nn/layer_norm.cc index fbad41cc14491..a010c22121445 100644 --- a/onnxruntime/core/providers/webgpu/nn/layer_norm.cc +++ b/onnxruntime/core/providers/webgpu/nn/layer_norm.cc @@ -228,7 +228,8 @@ Status RunLayerNormProgram(ComputeContext& context, LayerNormProgram program{bias != nullptr, simplified, mean != nullptr, inv_std_dev != nullptr, split_norm_dim, fp32_normalization}; - program.CacheHint(components, simplified, split_norm_dim, fp32_normalization) + // Output count alone cannot distinguish mean-only from inverse-standard-deviation-only. + program.CacheHint(components, simplified, split_norm_dim, fp32_normalization, mean != nullptr) .AddInputs({{x, ProgramTensorMetadataDependency::Type, GetOverrideShape(x->Shape(), components), components}}) .AddInputs( {{scale, ProgramTensorMetadataDependency::Type, GetOverrideShape(scale->Shape(), components), components}}) diff --git a/onnxruntime/core/providers/webgpu/program.h b/onnxruntime/core/providers/webgpu/program.h index 67e94d3035616..8b62c69da3115 100644 --- a/onnxruntime/core/providers/webgpu/program.h +++ b/onnxruntime/core/providers/webgpu/program.h @@ -353,6 +353,12 @@ class ProgramBase { // requires the device to support wgpu::FeatureName::SubgroupSizeControl; see ShaderHelper::GenerateSourceCode. ProgramBase& SetSubgroupSize(uint32_t size); + // Opt in only when the shader uses subgroup_index and subgroup_id is supported by the instance. + ProgramBase& SetUseSubgroupId(bool value) { + use_subgroup_id_ = value; + return *this; + } + // add a uniform variable. // // the specified uniform variable should match the uniform definition in the class, @@ -398,6 +404,7 @@ class ProgramBase { inline uint32_t WorkgroupSizeY() const { return workgroup_size_y_; } inline uint32_t WorkgroupSizeZ() const { return workgroup_size_z_; } inline uint32_t SubgroupSize() const { return subgroup_size_; } + inline bool UseSubgroupId() const { return use_subgroup_id_; } inline const std::vector& UniformVariables() const { return variables_; } inline const std::vector& OverridableConstants() const { return overridable_constants_; } @@ -427,6 +434,7 @@ class ProgramBase { uint32_t workgroup_size_y_; uint32_t workgroup_size_z_; uint32_t subgroup_size_; + bool use_subgroup_id_{false}; std::vector variables_; std::vector overridable_constants_; diff --git a/onnxruntime/core/providers/webgpu/program_cache_key.cc b/onnxruntime/core/providers/webgpu/program_cache_key.cc index ae83918a77376..029da45cdcf03 100644 --- a/onnxruntime/core/providers/webgpu/program_cache_key.cc +++ b/onnxruntime/core/providers/webgpu/program_cache_key.cc @@ -71,6 +71,9 @@ std::string CalculateProgramCacheKey(const ProgramBase& program, // = ||... // = ;; ss << program.Name(); + if (program.UseSubgroupId()) { + ss << "[SubgroupId]"; + } // append custom cache hint if any if (auto& hint = program.CacheHint(); !hint.empty()) { diff --git a/onnxruntime/core/providers/webgpu/quantization/quantize_linear.cc b/onnxruntime/core/providers/webgpu/quantization/quantize_linear.cc index 2ce9c14e3b406..99e3bb0b67fa1 100644 --- a/onnxruntime/core/providers/webgpu/quantization/quantize_linear.cc +++ b/onnxruntime/core/providers/webgpu/quantization/quantize_linear.cc @@ -170,15 +170,18 @@ Status DequantizeLinear::ComputeInternal(ComputeContext& context) const { : PackingMode::None; bool packed = packing != PackingMode::None; bool is_packed_signed = x_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8 || x_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4; - int64_t axis = (axis_ >= 0) ? axis_ : axis_ + x_shape.NumDimensions(); - - int max_components = GetMaxComponents(x_size); + int64_t axis = (axis_ >= 0) ? axis_ : axis_ + static_cast(x_shape.NumDimensions()); // scaler - single scaler for all elements - bool per_layer = x_scale_rank == 0 || (x_scale_rank == 1 && x_scale->Shape()[0] == 1); + bool per_layer = block_size_ == 0 && + (x_scale_rank == 0 || (x_scale_rank == 1 && x_scale->Shape()[0] == 1)); // 1D tensor - 1 scaler for per axis - bool per_axis = per_layer == false && x_scale_rank == 1; + bool per_axis = block_size_ == 0 && per_layer == false && x_scale_rank == 1; + ORT_RETURN_IF(!per_layer && (axis < 0 || axis >= static_cast(x_shape.NumDimensions())), + "axis must be within the input rank"); + ORT_RETURN_IF(!per_layer && !per_axis && x_scale_rank != static_cast(x_shape.NumDimensions()), + "x_scale and x must have the same rank for blocked quantization"); // Compute effective block_size. When block_size_ is 0 (default) but scale is 1D with // fewer elements than the input dimension on the axis, infer block_size from the ratio. @@ -187,6 +190,7 @@ Status DequantizeLinear::ComputeInternal(ComputeContext& context) const { int64_t input_dim = x_shape[onnxruntime::narrow(axis)]; int64_t scale_dim = x_scale->Shape()[0]; if (scale_dim < input_dim) { + ORT_RETURN_IF(scale_dim == 0, "x_scale must be nonempty on a nonempty quantize axis"); block_size = input_dim / scale_dim; per_axis = false; // treat as block quantization } @@ -199,6 +203,7 @@ Status DequantizeLinear::ComputeInternal(ComputeContext& context) const { const auto& scale_shape = x_scale->Shape(); for (size_t i = 0; i < x_shape.NumDimensions(); i++) { if (scale_shape[i] < x_shape[i]) { + ORT_RETURN_IF(scale_shape[i] == 0, "x_scale must be nonempty on a nonempty quantize axis"); axis = static_cast(i); block_size = x_shape[i] / scale_shape[i]; break; @@ -224,6 +229,8 @@ Status DequantizeLinear::ComputeInternal(ComputeContext& context) const { } } if (x_zeropoint != nullptr) { + ORT_RETURN_IF(x_zeropoint->Shape().NumDimensions() != scale_shape.NumDimensions(), + "x_zero_point and x_scale must have the same rank for blocked quantization"); for (size_t i = 0; i < x_shape.NumDimensions(); i++) { ORT_RETURN_IF(x_zeropoint->Shape()[i] != scale_shape[i], "x_zero_point and x_scale must have the same shape for blocked quantization"); @@ -231,6 +238,25 @@ Status DequantizeLinear::ComputeInternal(ComputeContext& context) const { } } + if (per_layer) { + ORT_RETURN_IF(x_zeropoint != nullptr && + !(x_zeropoint->Shape().NumDimensions() == 0 || + (x_zeropoint->Shape().NumDimensions() == 1 && x_zeropoint->Shape()[0] == 1)), + "x_zero_point must be a scalar or size-1 vector for per-tensor quantization"); + } else if (per_axis) { + const int64_t input_dim = x_shape[onnxruntime::narrow(axis)]; + ORT_RETURN_IF(x_scale->Shape()[0] != input_dim, + "x_scale must match the quantize axis for per-axis quantization"); + ORT_RETURN_IF(x_zeropoint != nullptr && + (x_zeropoint->Shape().NumDimensions() != 1 || x_zeropoint->Shape()[0] != input_dim), + "x_zero_point must match x_scale for per-axis quantization"); + } + + if (x_size == 0) { + return Status::OK(); + } + + int max_components = GetMaxComponents(x_size); bool use_components = per_layer && packing != PackingMode::Packed4 && (!packed || max_components == 4); int components = use_components ? max_components : 1; int input_component = use_components ? max_components : 1; diff --git a/onnxruntime/core/providers/webgpu/shader_helper.cc b/onnxruntime/core/providers/webgpu/shader_helper.cc index edfc8227b198f..f1e7b78db8f5f 100644 --- a/onnxruntime/core/providers/webgpu/shader_helper.cc +++ b/onnxruntime/core/providers/webgpu/shader_helper.cc @@ -39,6 +39,10 @@ ShaderHelper::ShaderHelper(const ProgramBase& program, Status ShaderHelper::Init() { // dispatch group size is normalized so no need to validate it here + ORT_RETURN_IF_NOT(!program_.UseSubgroupId() || + (webgpu_context_.DeviceHasFeature(wgpu::FeatureName::Subgroups) && + webgpu_context_.Instance().HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId)), + "Program ", program_.Name(), " requires subgroups and the subgroup_id WGSL language feature."); // validate workgroup size // normalize 0 (meaning "use default") the same way GenerateSourceCode() does, so validation @@ -87,9 +91,11 @@ Status ShaderHelper::Init() { " @builtin(local_invocation_id) local_id : vec3"; if (webgpu_context_.DeviceHasFeature(wgpu::FeatureName::Subgroups)) { body_ss_ << ",\n" - " @builtin(subgroup_invocation_id) sg_id : u32,\n" - " @builtin(subgroup_id) subgroup_index : u32,\n" - " @builtin(subgroup_size) sg_size : u32"; + " @builtin(subgroup_invocation_id) sg_id : u32,\n"; + if (program_.UseSubgroupId()) { + body_ss_ << " @builtin(subgroup_id) subgroup_index : u32,\n"; + } + body_ss_ << " @builtin(subgroup_size) sg_size : u32"; } // When using indirect dispatch, avoid @builtin(num_workgroups) to skip Dawn's validation // and duplication overhead in TransformIndirectDispatchBuffer. diff --git a/onnxruntime/core/providers/webgpu/tensor/transpose.cc b/onnxruntime/core/providers/webgpu/tensor/transpose.cc index 7bb85ad74bb8f..a8160100e4180 100644 --- a/onnxruntime/core/providers/webgpu/tensor/transpose.cc +++ b/onnxruntime/core/providers/webgpu/tensor/transpose.cc @@ -150,6 +150,10 @@ Status Transpose::DoTranspose(onnxruntime::webgpu::ComputeContextBase& context, } TensorShape output_shape(output_dims); + if (input_shape.Size() == 0) { + return Status::OK(); + } + // Check if `OIHW2OHWIProgram` can be applied. // // `OIHW2OHWIProgram` was originally designed to transpose 4D weights from OIHW diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.cc b/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.cc index f43c53362f6d8..1e32fefde8d6a 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.cc +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.cc @@ -164,12 +164,39 @@ bool CanApplySubgroup(const ComputeContext& context, int64_t M, int64_t N, int64 return false; } +bool ShouldFoldMatMulSubgroupBatch(std::string_view architecture, + int64_t m, + bool can_fold_batch_into_m) { + if (!can_fold_batch_into_m) { + return false; + } + + const bool is_xe_lpg_or_xe_3lpg = architecture == gpu_arch::kXeLpg || + architecture == gpu_arch::kXe3Lpg; + const int64_t m_mod_32 = m % 32; + return !is_xe_lpg_or_xe_3lpg || (m_mod_32 > 0 && m_mod_32 <= 24); +} + +int64_t MatMulSubgroupElementsPerThreadY(std::string_view architecture, uint32_t m) { + const bool is_xe_lpg_or_xe_3lpg = architecture == gpu_arch::kXeLpg || + architecture == gpu_arch::kXe3Lpg; + return m <= 8 ? 1 : (m <= 16 ? 2 : (m <= 32 ? 4 : (is_xe_lpg_or_xe_3lpg ? 4 : 8))); +} + int64_t ElementsPerThreadY(ComputeContext& context, uint32_t M) { - // For Xe-LPG and Xe-3LPG, we have observed that 4 elements per thread is optimal when M is large. - const auto& arch = context.AdapterInfo().architecture; - const bool is_xe_lpg_or_xe_3lpg = arch == gpu_arch::kXeLpg || - arch == gpu_arch::kXe3Lpg; - return M <= 8 ? 1 : (M <= 16 ? 2 : (M <= 32 ? 4 : (is_xe_lpg_or_xe_3lpg ? 4 : 8))); + return MatMulSubgroupElementsPerThreadY(context.AdapterInfo().architecture, M); +} + +bool CanUseAVec4CooperativeLoad(std::string_view architecture, + uint32_t dim_inner, + int64_t elements_per_thread_y) { + // A 32-wide subgroup has four 8-lane cooperative-load groups. Every generated + // subgroup-size branch must distribute rows evenly across its lane groups. + constexpr int64_t max_lane_groups = 4; + return architecture == gpu_arch::kXe3Lpg && + dim_inner % 4 == 0 && + elements_per_thread_y > 0 && + elements_per_thread_y % max_lane_groups == 0; } Status MakeMatMulSubgroupSource(ShaderHelper& shader, diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.h b/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.h index 403a9b7f56595..3f6599cb59c8c 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.h +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup.h @@ -4,16 +4,12 @@ #pragma once #include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/vendor/intel/math/gemm_subgroup_utils.h" namespace onnxruntime { namespace webgpu { namespace intel { -namespace gpu_arch { -inline constexpr std::string_view kXeLpg = "xe-lpg"; -inline constexpr std::string_view kXe3Lpg = "xe-3lpg"; -} // namespace gpu_arch - const uint32_t kSubgroupLogicalWorkGroupSizeX = 32; const uint32_t kSubgroupLogicalWorkGroupSizeY = 8; const uint32_t kSubgroupLogicalWorkGroupSizeZ = 1; diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup_utils.h b/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup_utils.h new file mode 100644 index 0000000000000..5d78f9feff4cb --- /dev/null +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/gemm_subgroup_utils.h @@ -0,0 +1,35 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +namespace onnxruntime { +namespace webgpu { +namespace intel { + +namespace gpu_arch { +inline constexpr std::string_view kXeLpg = "xe-lpg"; +inline constexpr std::string_view kXe3Lpg = "xe-3lpg"; +} // namespace gpu_arch + +bool ShouldFoldMatMulSubgroupBatch(std::string_view architecture, + int64_t m, + bool can_fold_batch_into_m); + +int64_t MatMulSubgroupElementsPerThreadY(std::string_view architecture, uint32_t m); + +bool CanUseAVec4CooperativeLoad(std::string_view architecture, + uint32_t dim_inner, + int64_t elements_per_thread_y); + +std::optional SelectMatMulSubgroupSize(uint32_t adapter_min_subgroup_size, + uint32_t adapter_max_subgroup_size, + bool has_subgroup_size_control); + +} // namespace intel +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.cc b/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.cc index 8ea1022363a19..6a69bb46e1592 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.cc +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.cc @@ -12,6 +12,12 @@ namespace onnxruntime { namespace webgpu { namespace intel { +namespace { + +constexpr std::array kSupportedMatMulSubgroupSizes{32, 16, 8}; + +} // namespace + Status MatMulSubgroupProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& a = shader.AddInput("a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); @@ -37,15 +43,46 @@ Status MatMulSubgroupProgram::GenerateShaderCode(ShaderHelper& shader) const { return Status::OK(); } -bool CanApplyMatMulIntel(const ComputeContext& context, int64_t M, int64_t N, int64_t K) { - return CanApplySubgroup(context, M, N, K); +std::optional SelectMatMulSubgroupSize(uint32_t adapter_min_subgroup_size, + uint32_t adapter_max_subgroup_size, + bool has_subgroup_size_control) { + if (adapter_min_subgroup_size == adapter_max_subgroup_size) { + for (const uint32_t size : kSupportedMatMulSubgroupSizes) { + if (adapter_min_subgroup_size == size) { + return size; + } + } + return std::nullopt; + } + + if (has_subgroup_size_control) { + for (const uint32_t size : kSupportedMatMulSubgroupSizes) { + if (adapter_min_subgroup_size <= size && size <= adapter_max_subgroup_size) { + return size; + } + } + } + return std::nullopt; } -Status ApplyMatMulIntel(ComputeContext& context, - const Activation& activation, - const std::vector& inputs, - Tensor* output, - bool is_channels_last) { +std::optional SelectMatMulSubgroupSize(const ComputeContext& context) { + if (!context.HasFeature(wgpu::FeatureName::Subgroups)) { + return std::nullopt; + } + + const auto& adapter_info = context.AdapterInfo(); + return SelectMatMulSubgroupSize( + adapter_info.subgroupMinSize, + adapter_info.subgroupMaxSize, + context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)); +} + +Status ApplyMatMulSubgroup(ComputeContext& context, + const Activation& activation, + const std::vector& inputs, + Tensor* output, + bool is_channels_last, + uint32_t subgroup_size) { const auto* a = inputs[0]; const auto* b = inputs[1]; bool has_bias = inputs.size() > 2; @@ -66,12 +103,7 @@ Status ApplyMatMulIntel(ComputeContext& context, // fold actually claws that waste back. Otherwise the Z-dispatch path wins. const int64_t M = output_shape[output_shape.NumDimensions() - 2]; const auto& arch = context.AdapterInfo().architecture; - const bool is_xe_lpg_or_xe_3lpg = arch == gpu_arch::kXeLpg || - arch == gpu_arch::kXe3Lpg; - // 32 = kSubgroupLogicalWorkGroupSizeY * ElementsPerThreadY(M > 32) on Xe-LPG/3LPG - const int64_t m_mod_32 = M % 32; - const bool xe_lpg_or_xe_3lpg_fold_ok = (m_mod_32 > 0 && m_mod_32 <= 24); - if (batchA != 1 && batchB == 1 && (!is_xe_lpg_or_xe_3lpg || xe_lpg_or_xe_3lpg_fold_ok)) { + if (ShouldFoldMatMulSubgroupBatch(arch, M, batchA != 1 && batchB == 1)) { // dimensions of A: [`batchA` * M, K] int64_t batchAndM = a_shape.SizeToDimension(a_shape.NumDimensions() - 1); TensorShapeVector dims_a = {batchAndM, helper.K()}; @@ -107,12 +139,13 @@ Status ApplyMatMulIntel(ComputeContext& context, const bool is_vec4 = dim_b_outer % 4 == 0; // vec4 A loads and double-buffering of the B tile are only enabled on Xe-3LPG. const bool is_xe_3lpg = arch == gpu_arch::kXe3Lpg; - // Load A from global memory as vec4 when K is a multiple of 4; otherwise fall back to scalar load. - const bool a_vec4 = is_xe_3lpg && dim_inner % 4 == 0; // Double-buffering of the B tile (held in workgroup memory) is only enabled for float16 B inputs. const bool b_is_fp16 = is_xe_3lpg && b->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; const bool use_f32_accumulation = context.EnableMatmulFp32Accumulation() && a->IsDataType(); InlinedVector elements_per_thread = InlinedVector({4, ElementsPerThreadY(context, dim_a_outer), 1}); + // Fall back to scalar A loads when rows cannot be distributed evenly among + // the cooperative vec4 lane groups (for example, forced low-M execution). + const bool a_vec4 = CanUseAVec4CooperativeLoad(arch, dim_inner, elements_per_thread[1]); const uint32_t dispatch_x = narrow((dim_b_outer + kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0] - 1) / (kSubgroupLogicalWorkGroupSizeX * elements_per_thread[0])); @@ -130,6 +163,9 @@ Status ApplyMatMulIntel(ComputeContext& context, MatMulSubgroupProgram program{activation, has_bias, is_vec4, a_vec4, b_is_fp16, use_f32_accumulation, is_channels_last, elements_per_thread}; + if (context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) { + program.SetSubgroupSize(subgroup_size); + } program .CacheHint(activation.CacheKey(), absl::StrJoin(elements_per_thread, "-"), a_vec4, b_is_fp16, use_f32_accumulation, is_channels_last) diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.h b/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.h index ba6309fca5bf9..586bf209da0e0 100644 --- a/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.h +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul.h @@ -3,6 +3,9 @@ #pragma once +#include +#include + #include "core/providers/webgpu/webgpu_kernel.h" #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/program.h" @@ -12,6 +15,8 @@ namespace onnxruntime { namespace webgpu { namespace intel { +// TODO: Move this common subgroup implementation out of vendor/intel. Keep only +// Intel-specific automatic-selection thresholds and tuning under the vendor directory. class MatMulSubgroupProgram final : public Program { public: MatMulSubgroupProgram(const Activation& activation, @@ -52,13 +57,14 @@ class MatMulSubgroupProgram final : public Program { const InlinedVector elements_per_thread_; }; -bool CanApplyMatMulIntel(const ComputeContext& context, int64_t M, int64_t N, int64_t K); +std::optional SelectMatMulSubgroupSize(const ComputeContext& context); -Status ApplyMatMulIntel(ComputeContext& context, - const Activation& activation, - const std::vector& inputs, - Tensor* output, - bool is_channels_last); +Status ApplyMatMulSubgroup(ComputeContext& context, + const Activation& activation, + const std::vector& inputs, + Tensor* output, + bool is_channels_last, + uint32_t subgroup_size); } // namespace intel } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul_execution_planner.cc b/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul_execution_planner.cc new file mode 100644 index 0000000000000..30f9757948847 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul_execution_planner.cc @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/vendor/intel/math/matmul_execution_planner.h" + +#include + +namespace onnxruntime { +namespace webgpu { +namespace intel { + +IntelMatMulExecutionPlanner::IntelMatMulExecutionPlanner(SplitKConfig split_k_config) + : MatMulExecutionPlanner{std::move(split_k_config)} {} + +std::optional IntelMatMulExecutionPlanner::SelectVendorAlgorithm( + const MatMulAlgorithmSelectionParams& params) const { + if (params.can_use_subgroup_matrix) { + return MatMulAlgorithm::SubgroupMatrix; + } + if (params.has_subgroup_capability && + params.m >= 64 && params.n >= 512 && params.k >= 32) { + return MatMulAlgorithm::Subgroup; + } + return std::nullopt; +} + +} // namespace intel +} // namespace webgpu +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul_execution_planner.h b/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul_execution_planner.h new file mode 100644 index 0000000000000..c2254bc9e5c05 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/matmul_execution_planner.h @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/math/matmul_execution_planner.h" + +namespace onnxruntime { +namespace webgpu { +namespace intel { + +class IntelMatMulExecutionPlanner final : public MatMulExecutionPlanner { + public: + IntelMatMulExecutionPlanner() = default; + explicit IntelMatMulExecutionPlanner(SplitKConfig split_k_config); + + protected: + std::optional SelectVendorAlgorithm( + const MatMulAlgorithmSelectionParams& params) const override; +}; + +} // namespace intel +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/split_k_config.cc b/onnxruntime/core/providers/webgpu/vendor/intel/math/split_k_config.cc new file mode 100644 index 0000000000000..4c5f8f02e45a5 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/split_k_config.cc @@ -0,0 +1,41 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/vendor/intel/math/split_k_config.h" + +namespace onnxruntime { +namespace webgpu { +namespace intel { + +SplitKConfig CreateSplitKConfig(std::string_view architecture) { + // Disable Split-K on old Intel GPUs. + if (architecture == "gen-7" || architecture == "gen-8" || + architecture == "gen-9" || architecture == "gen-11") { + return {}; + } + + constexpr uint32_t max_batch_size = 8; + constexpr uint32_t split_dim_inner = 256; + constexpr uint32_t min_dim_inner_with_split_k = split_dim_inner * 2; + + if (architecture == "xe-2lpg" || architecture == "xe-2hpg" || + architecture == "gen-12hp") { + // These thresholds are verified on Intel discrete GPUs and Lunar Lake iGPUs. + return SplitKConfig{ + max_batch_size, split_dim_inner, min_dim_inner_with_split_k, {{768, 52.0}, {2304, 35.0}, {3072, 21.5}, {4096, 16.0}}}; + } + + if (architecture == "xe-3lpg") { + // These thresholds are verified on Intel Panther Lake iGPUs (12Xe). + return SplitKConfig{ + max_batch_size, split_dim_inner, min_dim_inner_with_split_k, {{768, 40.0}, {1792, 22.0}, {3072, 18.0}, {4096, 10.0}}}; + } + + // Default thresholds for newer Intel GPUs, chosen on a gen-12lp GPU with 32 EUs. + return SplitKConfig{ + max_batch_size, split_dim_inner, min_dim_inner_with_split_k, {{768, 20.0}, {1792, 13.0}, {3072, 8.0}, {4096, 6.0}}}; +} + +} // namespace intel +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/vendor/intel/math/split_k_config.h b/onnxruntime/core/providers/webgpu/vendor/intel/math/split_k_config.h new file mode 100644 index 0000000000000..c2b826126222c --- /dev/null +++ b/onnxruntime/core/providers/webgpu/vendor/intel/math/split_k_config.h @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "core/providers/webgpu/webgpu_utils.h" + +namespace onnxruntime { +namespace webgpu { +namespace intel { + +SplitKConfig CreateSplitKConfig(std::string_view architecture); + +} // namespace intel +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 08693c6528787..587b430e180d1 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -9,6 +9,8 @@ #include #include +#include + #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) #include #include @@ -454,7 +456,7 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { program_mgr_ = std::make_unique(*this); // create split-k config - split_k_config_ = std::make_unique(adapter_info_); + split_k_config_ = std::make_unique(CreateSplitKConfig(adapter_info_)); // set query type #if !defined(__wasm__) @@ -625,7 +627,7 @@ Status WebGpuContext::EncodeDeferredDispatches(CommandRecordingState& recording) DispatchCommand(dispatch, recording); if (recording.graph_capture_state == GraphCaptureState::Capturing) { ORT_ENFORCE(recording.external_captured_commands != nullptr); - recording.external_captured_commands->push_back(std::move(dispatch)); + recording.external_captured_commands->push_back({std::move(dispatch)}); } } @@ -957,7 +959,7 @@ Status WebGpuContext::Run(ComputeContextBase& context, const ProgramBase& progra // Record the ready bind group and return. The deferred drain only needs to wait for the pipeline // and encode the dispatch commands. - webgpu::CapturedCommandInfo command; + webgpu::DispatchCommandInfo command; command.program_key = key; if (program_artifact != nullptr) { command.compute_pipeline = program_artifact->compute_pipeline; @@ -1342,6 +1344,12 @@ Status WebGpuContext::CheckDeviceStatus() const { Status WebGpuContext::FlushAndWaitChecked(const webgpu::BufferManager& buffer_mgr, CommandRecordingState& recording) { + return RunAndWaitChecked(buffer_mgr, recording, []() { return Status::OK(); }); +} + +Status WebGpuContext::RunAndWaitChecked(const webgpu::BufferManager& buffer_mgr, + CommandRecordingState& recording, + const std::function& enqueue) { ORT_RETURN_IF_NOT(HasDevice(), "WebGPU checked completion requires a device."); PushErrorScope(wgpu::ErrorFilter::Internal); @@ -1360,6 +1368,27 @@ Status WebGpuContext::FlushAndWaitChecked(const webgpu::BufferManager& buffer_mg result = Status(result.Category(), result.Code(), result.ErrorMessage() + "; " + status.ErrorMessage()); }; + size_t scopes_to_pop = 3; + const auto pop_scopes = [&]() { + while (scopes_to_pop > 0) { + --scopes_to_pop; + record_failure(PopErrorScope()); + } + }; + const auto cleanup = gsl::finally(pop_scopes); + + ORT_TRY { + record_failure(enqueue()); + } + ORT_CATCH(const std::exception& ex) { + ORT_HANDLE_EXCEPTION([&]() { + record_failure(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "WebGPU checked enqueue failed: ", ex.what())); + }); + } + ORT_CATCH(...) { + record_failure(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "WebGPU checked enqueue failed: unknown exception.")); + } + ORT_TRY { record_failure(Flush(buffer_mgr, recording)); } @@ -1392,9 +1421,7 @@ Status WebGpuContext::FlushAndWaitChecked(const webgpu::BufferManager& buffer_mg } // Do not return early: every scope must be popped even if flushing or waiting failed. - for (int i = 0; i < 3; ++i) { - record_failure(PopErrorScope()); - } + pop_scopes(); record_failure(CheckDeviceStatus()); return result; } @@ -1521,7 +1548,7 @@ wgpu::BindGroup WebGpuContext::CreateBindGroup(const std::vector& bi return wgpu::BindGroup::Acquire(bind_group); } -void WebGpuContext::DispatchCommand(const webgpu::CapturedCommandInfo& command, +void WebGpuContext::DispatchCommand(const webgpu::DispatchCommandInfo& command, CommandRecordingState& recording) { ORT_ENFORCE(command.compute_pipeline.has_value()); ORT_ENFORCE(command.bind_group != nullptr); @@ -1545,7 +1572,20 @@ void WebGpuContext::DispatchCommand(const webgpu::CapturedCommandInfo& command, } } -void WebGpuContext::CaptureBegin(std::vector* captured_commands, +void WebGpuContext::CopyBuffer(WGPUBuffer source, WGPUBuffer destination, uint64_t size, + CommandRecordingState& recording) { + // Materialize earlier deferred dispatches before inserting the copy into either command sequence. + ORT_THROW_IF_ERROR(EncodeDeferredDispatches(recording)); + EndComputePass(recording); + GetCommandEncoder(recording).CopyBufferToBuffer(source, 0, destination, 0, size); + if (recording.graph_capture_state == GraphCaptureState::Capturing) { + ORT_ENFORCE(recording.external_captured_commands != nullptr); + recording.external_captured_commands->push_back( + {CapturedBufferCopy{wgpu::Buffer{source}, wgpu::Buffer{destination}, size}}); + } +} + +void WebGpuContext::CaptureBegin(std::vector* captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording) { LOGS_DEFAULT(VERBOSE) << "CaptureBegin with external storage"; @@ -1561,7 +1601,7 @@ void WebGpuContext::CaptureBegin(std::vector* captu recording.graph_capture_state = GraphCaptureState::Capturing; } -void WebGpuContext::Replay(const std::vector& captured_commands, +void WebGpuContext::Replay(const std::vector& captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording) { LOGS_DEFAULT(VERBOSE) << "Replay with external storage"; @@ -1569,7 +1609,12 @@ void WebGpuContext::Replay(const std::vector& captu // Replay all captured commands from the provided vector const size_t command_count = captured_commands.size(); for (size_t i = 0; i < command_count; ++i) { - auto& command = captured_commands[i]; + const auto& operation = captured_commands[i].operation; + if (const auto* copy = std::get_if(&operation)) { + CopyBuffer(copy->source.Get(), copy->destination.Get(), copy->size, recording); + continue; + } + const auto& command = std::get(operation); // Restore profiling info when profiling is enabled. All commands are expected // to have profiling data in this mode to keep pending_kernels_ consistent // with num_pending_dispatches_. @@ -1600,12 +1645,10 @@ void WebGpuContext::CaptureEnd(CommandRecordingState& recording) { recording.external_captured_commands = nullptr; } -void WebGpuContext::ReleaseGraphResources(std::vector& captured_commands) { +void WebGpuContext::ReleaseGraphResources(std::vector& captured_commands) { LOGS_DEFAULT(VERBOSE) << "ReleaseGraphResources: Releasing " << captured_commands.size() << " captured command resources"; - for (auto& command : captured_commands) { - command.bind_group = nullptr; - } + captured_commands.clear(); } std::mutex WebGpuContextFactory::mutex_; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index cc3b9c7b9466d..a1ba426a713b9 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -11,6 +11,7 @@ #include #include #include +#include #include #include "core/providers/webgpu/webgpu_external_header.h" @@ -104,7 +105,7 @@ struct PendingPipelineBuild { // // A command with both fields empty is therefore valid only in state 3 before deferred pipeline // resolution. DispatchCommand() requires compute_pipeline to be set. -struct CapturedCommandInfo { +struct DispatchCommandInfo { std::string program_key; std::optional pending_build; std::optional compute_pipeline; @@ -116,6 +117,17 @@ struct CapturedCommandInfo { std::optional pending_kernel_info; }; +struct CapturedBufferCopy { + wgpu::Buffer source; + wgpu::Buffer destination; + uint64_t size; +}; + +// Copies retain both buffers and share dispatch ordering, but do not consume timestamp query slots. +struct CapturedCommand { + std::variant operation; +}; + // A command recording timeline has one caller at a time. ORT serializes Run on a Session; // plugin streamless operations use independent encoders. BufferManager synchronizes buffer // releases from concurrent allocator callers separately from command recording. @@ -133,10 +145,10 @@ struct CommandRecordingState { std::mutex pending_release_callbacks_mutex; std::vector> pending_release_callbacks; #endif - std::vector deferred_dispatches; + std::vector deferred_dispatches; std::vector pending_kernels; GraphCaptureState graph_capture_state{GraphCaptureState::Default}; - std::vector* external_captured_commands = nullptr; + std::vector* external_captured_commands = nullptr; }; struct WebGpuBufferCacheConfig { @@ -311,14 +323,15 @@ class WebGpuContext final { recording.compute_pass_encoder = nullptr; } } - void CaptureBegin(std::vector* captured_commands, + void CopyBuffer(WGPUBuffer source, WGPUBuffer destination, uint64_t size, CommandRecordingState& recording); + void CaptureBegin(std::vector* captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording); void CaptureEnd(CommandRecordingState& recording); - void Replay(const std::vector& captured_commands, + void Replay(const std::vector& captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording); - void ReleaseGraphResources(std::vector& captured_commands); + void ReleaseGraphResources(std::vector& captured_commands); Status Flush(const webgpu::BufferManager& buffer_mgr, CommandRecordingState& recording); @@ -326,6 +339,11 @@ class WebGpuContext final { // Serialize with other operations on this recording. Scopes cover flushing, not earlier encoding. Status FlushAndWaitChecked(const webgpu::BufferManager& buffer_mgr, CommandRecordingState& recording); + // Scopes cover enqueue (which may run multiple programs), flushing and completion. + // Drain and resolve all scopes even if enqueue fails or throws. Do not switch threads in enqueue. + Status RunAndWaitChecked(const webgpu::BufferManager& buffer_mgr, CommandRecordingState& recording, + const std::function& enqueue); + // Context-level managers are shared by sessions and synchronize their buffer caches internally. webgpu::BufferManager& BufferManager() const { return *buffer_mgr_; } webgpu::BufferManager& InitializerBufferManager() const { return *initializer_buffer_mgr_; } @@ -438,7 +456,7 @@ class WebGpuContext final { const std::vector& bind_buffers_segments, const wgpu::BindGroupLayout& bind_group_layout, std::string_view label) const; - void DispatchCommand(const webgpu::CapturedCommandInfo& command, CommandRecordingState& recording); + void DispatchCommand(const webgpu::DispatchCommandInfo& command, CommandRecordingState& recording); Status EncodeDeferredDispatches(CommandRecordingState& recording); std::vector GetEnabledAdapterToggles() const; @@ -536,7 +554,7 @@ class WebGpuContext final { GraphCaptureState graph_capture_state_{GraphCaptureState::Default}; // External vector to store captured commands, owned by EP - std::vector* external_captured_commands_ = nullptr; + std::vector* external_captured_commands_ = nullptr; }; #if !defined(__wasm__) diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 57d2e578f8477..493b5eec03c6c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -473,7 +473,8 @@ std::unique_ptr RegisterKernels(bool enable_graph_capture, bool ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastVersionedKernelInfo(21, 22, enable_int64))); ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastVersionedKernelInfo(23, 23, enable_int64))); ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastVersionedKernelInfo(24, 24, enable_int64))); - ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo(25, enable_int64))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastVersionedKernelInfo(25, 27, enable_int64))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo(28, enable_int64))); // Register int64 Clip kernels with conditional int64 support RegisterClipInt64Kernels(*kernel_registry, enable_int64); @@ -625,6 +626,7 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id, multi_rotary_cache_concat_offset_{config.multi_rotary_cache_concat_offset}, kv_cache_quantization_bits_{config.kv_cache_quantization_bits}, enable_matmul_fp32_accumulation_{config.enable_matmul_fp32_accumulation}, + forced_matmul_algorithm_{config.forced_matmul_algorithm}, #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) weight_load_acceleration_mode_{config.weight_load_acceleration_mode}, #endif @@ -974,7 +976,7 @@ Status WebGpuExecutionProvider::OnRunEnd(bool /* sync_stream */, const onnxrunti } #endif }); - // When capturing, flushing creates the replay-ready CapturedCommandInfo entries before + // When capturing, flushing creates the replay-ready dispatch entries before // CaptureEnd() detaches their external storage. Status flush_status = context_.Flush(BufferManager(), *recording_); diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h index 20294f59a7003..3103d1ffcd548 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -17,6 +18,7 @@ #include "core/graph/constants.h" #include "core/providers/providers.h" #include "core/providers/webgpu/buffer_manager.h" +#include "core/providers/webgpu/math/matmul_algorithm.h" #include "core/providers/webgpu/session_buffer_pool.h" #include "core/providers/webgpu/webgpu_provider_options.h" @@ -45,8 +47,7 @@ class GpuBufferAllocator; class D3D12ImportedBufferRegistry; #endif -// Forward declare CapturedCommandInfo which is now defined in webgpu_context.h -struct CapturedCommandInfo; +struct CapturedCommand; // The actual implementation of kernel registration. std::shared_ptr GetKernelRegistry(bool enable_graph_capture, bool enable_int64); @@ -69,6 +70,7 @@ struct WebGpuExecutionProviderConfig { // This is the single line that decides the shipped default for the // "enableMatmulFp32Accumulation" provider option. bool enable_matmul_fp32_accumulation{false}; + std::optional forced_matmul_algorithm; std::vector force_cpu_node_names{}; }; @@ -135,6 +137,7 @@ class WebGpuExecutionProvider : public IExecutionProvider { uint32_t KvCacheQuantizationBits() const { return kv_cache_quantization_bits_; } bool KvCacheQuantizationEnabled() const { return kv_cache_quantization_bits_ != 0; } bool EnableMatmulFp32Accumulation() const { return enable_matmul_fp32_accumulation_; } + std::optional ForcedMatMulAlgorithm() const { return forced_matmul_algorithm_; } #if defined(ORT_USE_EP_API_ADAPTERS) onnxruntime::ep::adapter::Logger& GetEpLogger() const; @@ -157,6 +160,7 @@ class WebGpuExecutionProvider : public IExecutionProvider { uint32_t multi_rotary_cache_concat_offset_ = 0; uint32_t kv_cache_quantization_bits_ = 0; bool enable_matmul_fp32_accumulation_ = false; + std::optional forced_matmul_algorithm_; #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) webgpu::WeightLoadAccelerationMode weight_load_acceleration_mode_{ webgpu::WeightLoadAccelerationMode::Off}; @@ -186,7 +190,7 @@ class WebGpuExecutionProvider : public IExecutionProvider { std::unique_ptr session_buffer_pool_; // Store captured commands per graph annotation ID - std::unordered_map> captured_graphs_; + std::unordered_map> captured_graphs_; // Track which graph annotation IDs have completed capture std::unordered_set captured_graph_ids_; diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index ea5f529a9b21b..520103e687241 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -468,18 +468,19 @@ WebGpuContextConfig ParseWebGpuContextConfig( } // namespace static std::shared_ptr CreateWebGpuProviderFactory( - const ConfigOptions& config_options, uint64_t test_only_max_storage_buffer_binding_size) { + const ConfigOptions& config_options, const WebGpuExecutionProviderTestOptions& test_options) { const auto weight_load_acceleration_mode = ResolveWeightLoadAccelerationMode(config_options); // prepare WebGpuExecutionProviderConfig WebGpuExecutionProviderConfig webgpu_ep_config = ParseEpConfig(config_options, weight_load_acceleration_mode); + webgpu_ep_config.forced_matmul_algorithm = test_options.forced_matmul_algorithm; // prepare WebGpuContextConfig WebGpuContextConfig config = ParseWebGpuContextConfig(config_options, weight_load_acceleration_mode); - config.test_only_max_storage_buffer_binding_size = test_only_max_storage_buffer_binding_size; + config.test_only_max_storage_buffer_binding_size = test_options.max_storage_buffer_binding_size; // Load the Dawn library and create the WebGPU instance. auto& context = WebGpuContextFactory::CreateContext(config); @@ -489,12 +490,12 @@ static std::shared_ptr CreateWebGpuProviderFactory( } std::shared_ptr WebGpuProviderFactoryCreator::Create(const ConfigOptions& config_options) { - return CreateWebGpuProviderFactory(config_options, 0); + return CreateWebGpuProviderFactory(config_options, {}); } std::shared_ptr WebGpuProviderFactoryCreator::CreateForTesting( - const ConfigOptions& config_options, uint64_t max_storage_buffer_binding_size) { - return CreateWebGpuProviderFactory(config_options, max_storage_buffer_binding_size); + const ConfigOptions& config_options, const WebGpuExecutionProviderTestOptions& test_options) { + return CreateWebGpuProviderFactory(config_options, test_options); } // WebGPU DataTransfer implementation wrapper for the C API diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory_creator.h b/onnxruntime/core/providers/webgpu/webgpu_provider_factory_creator.h index 7e4ed3fe3ed89..49b027f207ec9 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory_creator.h +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory_creator.h @@ -3,11 +3,14 @@ #pragma once +#include #include +#include #include "core/framework/provider_options.h" #include "core/providers/providers.h" +#include "core/providers/webgpu/math/matmul_algorithm.h" #include "core/providers/webgpu/webgpu_provider_options.h" struct OrtDataTransferImpl; @@ -15,10 +18,15 @@ struct OrtDataTransferImpl; namespace onnxruntime { struct ConfigOptions; +struct WebGpuExecutionProviderTestOptions { + uint64_t max_storage_buffer_binding_size{0}; + std::optional forced_matmul_algorithm; +}; + struct WebGpuProviderFactoryCreator { static std::shared_ptr Create(const ConfigOptions& config_options); static std::shared_ptr CreateForTesting( - const ConfigOptions& config_options, uint64_t max_storage_buffer_binding_size); + const ConfigOptions& config_options, const WebGpuExecutionProviderTestOptions& test_options); }; // Each copy uses its session stream or the GPU tensor's registered context. diff --git a/onnxruntime/core/providers/webgpu/webgpu_utils.cc b/onnxruntime/core/providers/webgpu/webgpu_utils.cc index 62d7b9934e164..a22b93a0b3cab 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_utils.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_utils.cc @@ -3,7 +3,10 @@ #include "core/providers/webgpu/webgpu_utils.h" #include +#include + #include "core/providers/webgpu/shader_variable.h" +#include "core/providers/webgpu/vendor/intel/math/split_k_config.h" namespace onnxruntime { namespace webgpu { @@ -25,55 +28,32 @@ TensorShape ReduceShapeByComponents(const TensorShape& shape, int64_t components return TensorShape(shape_vector); } -SplitKConfig::SplitKConfig(const wgpu::AdapterInfo& adapter_info) { - if (adapter_info.vendor == std::string_view{"intel"}) { - // Disable Split-K on old Intel GPUs. - if (adapter_info.architecture == std::string_view{"gen-7"} || - adapter_info.architecture == std::string_view{"gen-8"} || - adapter_info.architecture == std::string_view{"gen-9"} || - adapter_info.architecture == std::string_view{"gen-11"}) { - enable_split_k_ = false; - } else if (adapter_info.architecture == std::string_view{"xe-2lpg"} || - adapter_info.architecture == std::string_view{"xe-2hpg"} || - adapter_info.architecture == std::string_view{"gen-12hp"}) { - // Below thresholds are only verified on Intel discrete GPUs and Lunar Lake iGPUs. - enable_split_k_ = true; - - max_batch_size_ = 8; - split_dim_inner_ = 256; - min_dim_inner_with_split_k_ = split_dim_inner_ * 2; - - configs_per_dim_inner_range_.emplace_back(768, 52.0); - configs_per_dim_inner_range_.emplace_back(2304, 35.0); - configs_per_dim_inner_range_.emplace_back(3072, 21.5); - configs_per_dim_inner_range_.emplace_back(4096, 16.0); - } else if (adapter_info.architecture == std::string_view{"xe-3lpg"}) { - // Below thresholds are only verified on Intel Panther Lake iGPUs (12Xe). - enable_split_k_ = true; - - max_batch_size_ = 8; - split_dim_inner_ = 256; - min_dim_inner_with_split_k_ = split_dim_inner_ * 2; - - configs_per_dim_inner_range_.emplace_back(768, 40.0); - configs_per_dim_inner_range_.emplace_back(1792, 22.0); - configs_per_dim_inner_range_.emplace_back(3072, 18.0); - configs_per_dim_inner_range_.emplace_back(4096, 10.0); - } else { - // Below are the default thresholds on newer Intel GPUs. These values are chosen on - // Intel "gen-12lp" GPU with 32EUs. - enable_split_k_ = true; - - max_batch_size_ = 8; - split_dim_inner_ = 256; - min_dim_inner_with_split_k_ = split_dim_inner_ * 2; - - configs_per_dim_inner_range_.emplace_back(768, 20.0); - configs_per_dim_inner_range_.emplace_back(1792, 13.0); - configs_per_dim_inner_range_.emplace_back(3072, 8.0); - configs_per_dim_inner_range_.emplace_back(4096, 6.0); - } +SplitKConfig::SplitKConfig( + uint32_t max_batch_size, + uint32_t split_dim_inner, + uint32_t min_dim_inner_with_split_k, + std::initializer_list> configs_per_dim_inner_range) + : enable_split_k_{true}, + split_dim_inner_{split_dim_inner}, + min_dim_inner_with_split_k_{min_dim_inner_with_split_k}, + max_batch_size_{max_batch_size} { + configs_per_dim_inner_range_.reserve(configs_per_dim_inner_range.size()); + for (const auto& [max_dim_inner, rate] : configs_per_dim_inner_range) { + configs_per_dim_inner_range_.emplace_back(max_dim_inner, rate); + } +} + +SplitKConfig CreateSplitKConfig(const wgpu::AdapterInfo& adapter_info) { + return CreateSplitKConfig( + std::string_view{adapter_info.vendor}, + std::string_view{adapter_info.architecture}); +} + +SplitKConfig CreateSplitKConfig(std::string_view vendor, std::string_view architecture) { + if (vendor == "intel") { + return intel::CreateSplitKConfig(architecture); } + return {}; } SplitKConfig::ConfigAtRange::ConfigAtRange(uint32_t max_dim_inner, double rate) @@ -88,9 +68,9 @@ bool SplitKConfig::UseSplitK( bool is_vec4, ActivationKind activation_kind, uint64_t batch_size, - uint32_t dim_a_outer, - uint32_t dim_b_outer, - uint32_t dim_inner, + uint64_t dim_a_outer, + uint64_t dim_b_outer, + uint64_t dim_inner, bool is_channels_last) const { if (!enable_split_k_) { return false; diff --git a/onnxruntime/core/providers/webgpu/webgpu_utils.h b/onnxruntime/core/providers/webgpu/webgpu_utils.h index 9d3dd833315e5..7f3f96a1f12d7 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_utils.h +++ b/onnxruntime/core/providers/webgpu/webgpu_utils.h @@ -4,7 +4,12 @@ #pragma once #include +#include +#include #include +#include +#include + #include "core/common/common.h" #include "core/framework/tensor.h" #include "core/framework/tensor_shape.h" @@ -109,11 +114,17 @@ inline Tensor CreateTensorView(const Tensor& tensor, MLDataType new_data_type, c */ class SplitKConfig { public: - explicit SplitKConfig(const wgpu::AdapterInfo& adapter_info); + SplitKConfig() = default; + SplitKConfig( + uint32_t max_batch_size, + uint32_t split_dim_inner, + uint32_t min_dim_inner_with_split_k, + std::initializer_list> configs_per_dim_inner_range); bool UseSplitK( bool is_vec4, ActivationKind activation_kind, uint64_t batch_size, - uint32_t dim_a_outer, uint32_t dim_b_outer, uint32_t dim_inner, bool is_channels_last = true) const; + uint64_t dim_a_outer, uint64_t dim_b_outer, uint64_t dim_inner, + bool is_channels_last = true) const; uint32_t GetSplitDimInner() const; @@ -133,6 +144,9 @@ class SplitKConfig { std::vector configs_per_dim_inner_range_; }; +SplitKConfig CreateSplitKConfig(const wgpu::AdapterInfo& adapter_info); +SplitKConfig CreateSplitKConfig(std::string_view vendor, std::string_view architecture); + /** * Generates WGSL (WebGPU Shading Language) code for performing an atomic add operation * on a non-integer value (e.g., floating-point) in a shader. diff --git a/onnxruntime/core/providers/webnn/builders/impl/base_op_builder.h b/onnxruntime/core/providers/webnn/builders/impl/base_op_builder.h index f4429fbb06457..d1334ba9a5420 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/base_op_builder.h +++ b/onnxruntime/core/providers/webnn/builders/impl/base_op_builder.h @@ -52,7 +52,7 @@ class BaseOpBuilder : public IOpBuilder { // We still set the minimal supported opset to 1 as we couldn't // get the model opset version at this stage. virtual int GetMinSupportedOpSet(const Node& /* node */) const { return 1; } - virtual int GetMaxSupportedOpSet(const Node& /* node */) const { return 27; } + virtual int GetMaxSupportedOpSet(const Node& /* node */) const { return 28; } private: bool HasSupportedOpSet(const Node& node, const logging::Logger& logger) const; diff --git a/onnxruntime/core/providers/webnn/builders/impl/qdq_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/qdq_op_builder.cc index e4e3a09391c3b..5b6ab6420ca45 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/qdq_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/qdq_op_builder.cc @@ -25,6 +25,8 @@ class QDQOpBuilder : public BaseOpBuilder { const WebnnDeviceType /* device_type */, const logging::Logger& logger) const override; bool HasSupportedInputsImpl(const GraphViewer&, const Node& node, const emscripten::val& wnn_limits, const logging::Logger& logger) const override; + + int GetMaxSupportedOpSet(const Node& /* node */) const override { return 27; } }; Status QDQOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, diff --git a/onnxruntime/core/session/compile_api.cc b/onnxruntime/core/session/compile_api.cc index b2bf8ed6e3691..b924e35d9edae 100644 --- a/onnxruntime/core/session/compile_api.cc +++ b/onnxruntime/core/session/compile_api.cc @@ -485,6 +485,7 @@ static constexpr OrtCompileApi ort_compile_api = { &OrtCompileAPI::ModelCompilationOptions_SetEpContextDataWriteFunc, &OrtCompileAPI::ModelCompilationOptions_SetOutputModelExternalInitializersBuffer, &OrtCompileAPI::ModelCompilationOptions_SetOutputModelExternalInitializersAlignment, + // End of Version 31 - DO NOT MODIFY ABOVE }; // checks that we don't violate the rule that the functions must remain in the slots they were originally assigned @@ -496,6 +497,8 @@ static_assert(offsetof(OrtCompileApi, ModelCompilationOptions_SetInputModel) / s "Size of version 24 of Api cannot change"); static_assert(offsetof(OrtCompileApi, ModelCompilationOptions_SetWeightlessEnabled) / sizeof(void*) == 15, "Size of version 29 of Api cannot change"); +static_assert(offsetof(OrtCompileApi, ModelCompilationOptions_SetOutputModelExternalInitializersAlignment) / sizeof(void*) == 18, + "Size of version 31 of Api cannot change"); // So that nobody forgets to finish an API version, this check will serve as a reminder: static_assert(std::string_view(ORT_VERSION) == "1.32.0", diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index fbadf6da73678..52edede9208e0 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -296,7 +296,7 @@ Status Environment::Initialize(std::unique_ptr logging_ auto& domainToVersionRangeInstance = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance(); if (domainToVersionRangeInstance.Map().find(onnxruntime::kMSDomain) == domainToVersionRangeInstance.Map().end()) { // External shared providers may have already added kMSDomain - domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 1); + domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 2); } domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSExperimentalDomain, 1, 1); domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSNchwcDomain, 1, 1); diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 9227930f312dd..e9015b3e40b09 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -1249,7 +1249,9 @@ common::Status InferenceSession::LoadOnnxModel(const PathString& model_uri) { const bool strict_shape_type_inference = session_options_.config_options.GetConfigOrDefault( kOrtSessionOptionsConfigStrictShapeTypeInference, "0") == "1"; - ModelOptions model_opts{true, strict_shape_type_inference, check_load_cancellation_fn_}; + const bool allow_released_opsets_only = session_options_.config_options.GetConfigOrDefault( + kOrtSessionOptionsConfigStrictAllowReleasedOpsetsOnly, "1") == "1"; + ModelOptions model_opts{allow_released_opsets_only, strict_shape_type_inference, check_load_cancellation_fn_}; // When set, the external initializers folder overrides the model's own directory as the // base for resolving external data. The model bytes are still read from model_uri. @@ -1345,6 +1347,8 @@ common::Status InferenceSession::Load(const void* model_data, int model_data_len const bool strict_shape_type_inference = session_options_.config_options.GetConfigOrDefault( kOrtSessionOptionsConfigStrictShapeTypeInference, "0") == "1"; + const bool allow_released_opsets_only = session_options_.config_options.GetConfigOrDefault( + kOrtSessionOptionsConfigStrictAllowReleasedOpsetsOnly, "1") == "1"; PathString external_data_model_path = GetExternalInitializersFolderModelPath(session_options_.config_options); if (!external_data_model_path.empty()) { @@ -1353,7 +1357,8 @@ common::Status InferenceSession::Load(const void* model_data, int model_data_len return onnxruntime::Model::Load(std::move(model_proto), model_location_, model, HasLocalSchema() ? &custom_schema_registries_ : nullptr, *session_logger_, - ModelOptions{true, strict_shape_type_inference, check_load_cancellation_fn_}); + ModelOptions{allow_released_opsets_only, strict_shape_type_inference, + check_load_cancellation_fn_}); }; return LoadWithLoader(loader, "model_loading_array"); @@ -1409,6 +1414,8 @@ common::Status InferenceSession::LoadOnnxModel(ModelProto model_proto) { #endif const bool strict_shape_type_inference = session_options_.config_options.GetConfigOrDefault( kOrtSessionOptionsConfigStrictShapeTypeInference, "0") == "1"; + const bool allow_released_opsets_only = session_options_.config_options.GetConfigOrDefault( + kOrtSessionOptionsConfigStrictAllowReleasedOpsetsOnly, "1") == "1"; PathString external_data_model_path = GetExternalInitializersFolderModelPath(session_options_.config_options); if (!external_data_model_path.empty()) { @@ -1418,7 +1425,8 @@ common::Status InferenceSession::LoadOnnxModel(ModelProto model_proto) { // This call will move model_proto to the constructed model instance return onnxruntime::Model::Load(std::move(model_proto), model_location_, model, HasLocalSchema() ? &custom_schema_registries_ : nullptr, *session_logger_, - ModelOptions{true, strict_shape_type_inference, check_load_cancellation_fn_}); + ModelOptions{allow_released_opsets_only, strict_shape_type_inference, + check_load_cancellation_fn_}); }; return LoadWithLoader(loader, "model_loading_proto"); @@ -1516,12 +1524,15 @@ common::Status InferenceSession::Load(const OrtModel& model_editor_api_model) { const bool strict_shape_type_inference = session_options_.config_options.GetConfigOrDefault( kOrtSessionOptionsConfigStrictShapeTypeInference, "0") == "1"; + const bool allow_released_opsets_only = session_options_.config_options.GetConfigOrDefault( + kOrtSessionOptionsConfigStrictAllowReleasedOpsetsOnly, "1") == "1"; // need to go from unique_ptr to shared_ptr when moving into model_ std::unique_ptr tmp_model; ORT_RETURN_IF_ERROR(Model::LoadFromModelEditorApiModel(model_editor_api_model, HasLocalSchema() ? &custom_schema_registries_ : nullptr, - ModelOptions{true, strict_shape_type_inference, + ModelOptions{allow_released_opsets_only, + strict_shape_type_inference, check_load_cancellation_fn_}, *session_logger_, tmp_model)); diff --git a/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc b/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc index 33652534004b3..c1bbf1a0dbb21 100644 --- a/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc +++ b/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc @@ -3,6 +3,7 @@ #include "core/session/plugin_ep/ep_plugin_provider_interfaces.h" +#include #include #include #include @@ -17,6 +18,7 @@ #include "core/common/inlined_containers.h" #include +#include "core/graph/constants.h" #include "core/graph/ep_api_types.h" #include "core/graph/model_editor_api_types.h" #include "core/session/abi_devices.h" @@ -337,6 +339,24 @@ PluginExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_vie continue; } + // Plugin accounting cannot receive Level-1 estimates yet. Reject the whole + // group rather than use generic or profiled costs for an unproven GQA bound. + if (resource_accountant != nullptr && Type() == kCudaExecutionProvider) { + const auto unproven_gqa = std::find_if( + node_grouping.nodes.begin(), node_grouping.nodes.end(), [](const EpNode* ep_node) { + const Node& node = ep_node->GetInternalNode(); + return node.GetExecutionProviderType().empty() && + node.OpType() == "GroupQueryAttention" && node.Domain() == kMSDomain; + }); + if (unproven_gqa != node_grouping.nodes.end()) { + LOGS(logger, WARNING) + << Type() << " declining capacity-aware assignment for GroupQueryAttention node " + << (*unproven_gqa)->GetInternalNode().Name() + << " because the plugin accounting path has no proven route-aware workspace estimate."; + continue; + } + } + if (node_grouping.kind == OrtEpGraphSupportInfo::NodeGroupingKind::kSingleAssignedNode) { if (node_grouping.nodes.size() != 1) { LOGS(logger, ERROR) << "OrtEp::GetCapability() for " << Type() << " did not specify exactly one valid node " diff --git a/onnxruntime/core/session/provider_bridge_ort.h b/onnxruntime/core/session/provider_bridge_ort.h index d15de1076b2c7..b04ceec474968 100644 --- a/onnxruntime/core/session/provider_bridge_ort.h +++ b/onnxruntime/core/session/provider_bridge_ort.h @@ -3,9 +3,10 @@ #pragma once +struct ProviderInfo_OpenVINO; + namespace onnxruntime { -struct ProviderInfo_OpenVINO; ProviderInfo_OpenVINO* TryGetProviderInfo_OpenVINO(); bool InitProvidersSharedLibrary(); diff --git a/onnxruntime/python/backend/backend.py b/onnxruntime/python/backend/backend.py index 69be7a7657adf..5edac6df88fc0 100644 --- a/onnxruntime/python/backend/backend.py +++ b/onnxruntime/python/backend/backend.py @@ -135,6 +135,9 @@ def prepare(cls, model, device=None, **kwargs): return OnnxRuntimeBackendRep(model) elif isinstance(model, (str, bytes)): options = SessionOptions() + options.add_session_config_entry( + "session.allow_released_opsets_only", "1" if cls.allowReleasedOpsetsOnly else "0" + ) for k, v in kwargs.items(): if k in _ALLOWED_SESSION_OPTIONS: setattr(options, k, v) diff --git a/onnxruntime/python/backend/backend_rep.py b/onnxruntime/python/backend/backend_rep.py index 950ce417c6c2d..75b682fc9002e 100644 --- a/onnxruntime/python/backend/backend_rep.py +++ b/onnxruntime/python/backend/backend_rep.py @@ -6,6 +6,7 @@ Implements ONNX's backend API. """ +import numpy as np from onnx.backend.base import BackendRep from onnxruntime import RunOptions @@ -23,6 +24,10 @@ ) +def _normalize_numpy_scalar(value): + return np.asarray(value) if isinstance(value, np.generic) else value + + class OnnxRuntimeBackendRep(BackendRep): """ Wraps an :class:`onnxruntime.InferenceSession` to implement ONNX's @@ -61,7 +66,7 @@ def run(self, inputs, **kwargs): # type: (Any, **Any) -> Tuple[Any, ...] if isinstance(inputs, list): inps = {} for i, inp in enumerate(self._session.get_inputs()): - inps[inp.name] = inputs[i] + inps[inp.name] = _normalize_numpy_scalar(inputs[i]) outs = self._session.run(None, inps, options) if isinstance(outs, list): return outs @@ -72,5 +77,5 @@ def run(self, inputs, **kwargs): # type: (Any, **Any) -> Tuple[Any, ...] inp = self._session.get_inputs() if len(inp) != 1: raise RuntimeError(f"Model expect {len(inp)} inputs") - inps = {inp[0].name: inputs} + inps = {inp[0].name: _normalize_numpy_scalar(inputs)} return self._session.run(None, inps, options) diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index f8a569f26e122..1211899ceb103 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -1474,6 +1474,11 @@ static std::shared_ptr CreateExecutionProviderFactory cit == provider_options_map.end() ? ProviderOptions{} : cit->second, &session_options); #endif } else if (type == kWebGpuExecutionProvider) { +#if !defined(ORT_MINIMAL_BUILD) + if (auto ep_factory = try_create_registered_plugin_factory(); ep_factory) { + return ep_factory; + } +#endif #if defined(USE_WEBGPU) return onnxruntime::WebGpuProviderFactoryCreator::Create(session_options.config_options); #endif diff --git a/onnxruntime/python/tools/fpa_intb_tune.py b/onnxruntime/python/tools/fpa_intb_tune.py new file mode 100644 index 0000000000000..77f9688bb58e3 --- /dev/null +++ b/onnxruntime/python/tools/fpa_intb_tune.py @@ -0,0 +1,346 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Offline tactic tuning tool for the MatMulNBits fpA_intB CUDA path. + +This tool tunes the weight-only GEMM tactics for a model's ``MatMulNBits`` nodes and +writes the results to a persistent tactic cache file that ONNX Runtime can reuse on +later runs (see docs/contrib_ops/cuda/gemm_profiler_cache.md). + +How it works: + * It creates a CUDA execution-provider session for the model with the fpA_intB path and the + cache prefix enabled through session config entries. Kernel construction profiles the + configured M buckets that are not already in ``.matmulnbits_fpa_intb.tsv`` + (all of them with ``--retune``) and stages their successful tactics in memory. + * It then (best-effort) runs dummy inferences at each requested M value so that any + additional buckets are profiled lazily. + * All staged tactics are written to ``.matmulnbits_fpa_intb.tsv`` when the session + (and its CUDA execution provider) is released. + +The generated cache is hardware/build specific: it is only reused on the same GPU +model + SM + CUDA runtime + ORT version. + +Example: + python -m onnxruntime.tools.fpa_intb_tune \ + --model model.onnx \ + --output-prefix /path/to/cache/mymodel \ + --m-values 1,8,16,32,64,128,256,512,1024,2048 +""" + +from __future__ import annotations + +import argparse +import ctypes +import gc +import os +import sys +import tempfile +from urllib.parse import unquote + +import numpy as np + +import onnxruntime as ort + +_CACHE_TABLE_SUFFIX = ".matmulnbits_fpa_intb.tsv" +_CACHE_KEY_COLUMNS = ( + "n_16b", + "k", + "activation_dtype", + "weight_type", + "bits", + "block_size", + "has_zero_points", + "zero_point_dtype", + "gemv_enabled", + "has_bias", + "packing_sm", + "paired_gemv_mode", + "wave_aware_gemv", +) +_CACHE_CONFIG_COLUMNS = ( + "valid_config", + "sm_version", + "tile80", + "tile90", + "tile100", + "tile120", + "split_k_style", + "split_k", + "stages", + "cluster", + "mainloop", + "epilogue", + "tma", + "enable_cuda_kernel", + "cuda_kernel_variant", +) + + +def _parse_m_values(text: str) -> list[int]: + values = [] + for t in text.split(","): + token = t.strip() + if not token: + continue + m = int(token) + if m > 0: + values.append(m) + # De-duplicate while keeping ascending order. + return sorted(set(values)) + + +def _make_session_options(output_prefix: str, m_values: list[int], retune: bool = False) -> ort.SessionOptions: + """Enables the fpA_intB path, its M-bucket sweep, and the persistent cache for one session.""" + sess_options = ort.SessionOptions() + sess_options.add_session_config_entry("ep.cuda.gemm_tactic_cache_prefix", output_prefix) + sess_options.add_session_config_entry("ep.cuda.gemm_tactic_cache_mode", "save" if retune else "load_save") + sess_options.add_session_config_entry("ep.cuda.fpa_intb_gemm", "1") + if m_values: + sess_options.add_session_config_entry("ep.cuda.fpa_intb_profile_m", ",".join(str(m) for m in m_values)) + return sess_options + + +def _numpy_dtype_for(ort_type: str): + mapping = { + "tensor(float16)": np.float16, + "tensor(float)": np.float32, + "tensor(double)": np.float64, + "tensor(int64)": np.int64, + "tensor(int32)": np.int32, + "tensor(int16)": np.int16, + "tensor(int8)": np.int8, + "tensor(uint8)": np.uint8, + "tensor(bool)": np.bool_, + } + if ort_type not in mapping: + raise ValueError(f"Unsupported dummy input type: {ort_type}") + return mapping[ort_type] + + +def _make_dummy_inputs(session, m: int) -> dict: + """Best-effort dummy inputs. The first symbolic/dynamic dim of each input is set to + ``m`` (a heuristic to drive the GEMM M dimension); remaining dynamic dims become 1. + """ + + feeds = {} + for inp in session.get_inputs(): + shape = [] + replaced = False + for dim in inp.shape: + if isinstance(dim, int) and dim > 0: + shape.append(dim) + else: + shape.append(m if not replaced else 1) + replaced = True + if inp.type == "tensor(bfloat16)": + # BF16 zero has the same bits as uint16 zero; NumPy extension dtypes are not accepted by run(). + feeds[inp.name] = ort.OrtValue.ortvalue_from_numpy_with_onnx_type(np.zeros(shape, dtype=np.uint16), 16) + else: + feeds[inp.name] = ort.OrtValue.ortvalue_from_numpy(np.zeros(shape, dtype=_numpy_dtype_for(inp.type))) + return feeds + + +def _current_signature(device_id: int) -> dict[str, str]: + # Use the CUDA driver for device identity and the wheel metadata for the toolkit it was built with. + try: + from onnxruntime.capi.build_and_package_info import cuda_version # noqa: PLC0415 + except ImportError as exc: + raise RuntimeError( + "The tuning tool requires an ONNX Runtime package built with CUDA. " + "CPU packages with a registered CUDA plugin are not supported because their " + "metadata does not identify the plugin's CUDA build. Use a built-in CUDA package instead." + ) from exc + if not cuda_version: + raise RuntimeError("The tuning tool requires a package with built-in CUDA version metadata.") + + driver = ctypes.WinDLL("nvcuda.dll") if sys.platform == "win32" else ctypes.CDLL("libcuda.so.1") + + def check(code): + if code != 0: + raise RuntimeError(f"Cannot query the CUDA cache signature (CUDA driver error {code})") + + check(driver.cuInit(0)) + device = ctypes.c_int() + check(driver.cuDeviceGet(ctypes.byref(device), device_id)) + name = ctypes.create_string_buffer(256) + check(driver.cuDeviceGetName(name, len(name), device)) + major, minor = ctypes.c_int(), ctypes.c_int() + check(driver.cuDeviceComputeCapability(ctypes.byref(major), ctypes.byref(minor), device)) + toolkit = cuda_version.split(".") + return { + "device_name": name.value.decode("utf-8"), + "sm": str(major.value * 10 + minor.value), + "cuda_runtime": str(int(toolkit[0]) * 1000 + int(toolkit[1]) * 10), + # Match the C++ cache header, which writes the base ORT_VERSION without wheel dev/local suffixes. + "ort_version": ort.get_version_string(), + } + + +def _summarize_cache(cache_path: str, signature: dict[str, str], *, display: bool = True) -> set[tuple[str, ...]]: + if not os.path.exists(cache_path): + raise RuntimeError( + f"No tactic cache was produced at {cache_path}; check that the model uses fpA_intB MatMulNBits." + ) + + header = {} + columns = None + n_key_col = None + rows = 0 + unique_keys = set() + entries = set() + with open(cache_path, encoding="utf-8") as f: + for raw_line in f: + line = raw_line.rstrip("\n") + if not line: + continue + if line.startswith("#"): + parts = line[1:].strip().split("\t") + if len(parts) >= 2: + header[parts[0]] = unquote(parts[1]) + continue + fields = line.split("\t") + if columns is None: + columns = fields + n_key_col = {name: i for i, name in enumerate(columns)} + required = (*_CACHE_KEY_COLUMNS, "m_bucket", *_CACHE_CONFIG_COLUMNS) + if not set(required).issubset(n_key_col): + raise RuntimeError(f"Missing problem-key or tactic-setting columns in {cache_path}") + continue + if len(fields) != len(columns) or fields[n_key_col["valid_config"]] != "1": + raise RuntimeError(f"Invalid or unsuccessful tactic in {cache_path}") + rows += 1 + # Build a key tuple from the problem-key columns for a unique-shape count. + key = tuple(fields[n_key_col[c]] for c in _CACHE_KEY_COLUMNS) + unique_keys.add(key) + config = tuple(fields[n_key_col[c]] for c in _CACHE_CONFIG_COLUMNS) + entries.add((*key, fields[n_key_col["m_bucket"]], *config)) + + expected = { + **signature, + "ort_cuda_gemm_tactic_cache": "v1", + "table": "matmulnbits_fpa_intb", + "tactic_selection_version": "4", + } + for key, value in expected.items(): + if header.get(key) != value: + raise RuntimeError(f"Inapplicable tactic cache: {key}={header.get(key)!r}, expected {value!r}") + if rows == 0: + raise RuntimeError(f"No successful tactics in {cache_path}") + if not display: + return entries + print(f"Cache ready: {cache_path}") + print(f" device_name : {header.get('device_name', '?')}") + print(f" sm : {header.get('sm', '?')}") + print(f" cuda_runtime : {header.get('cuda_runtime', '?')}") + print(f" ort_version : {header.get('ort_version', '?')}") + print(f" ort_git_commit : {header.get('ort_git_commit', '?')}") + print(f" unique shapes : {len(unique_keys)}") + print(f" tuned (shape, M) : {rows}") + return entries + + +def tune(model: str, output_prefix: str, m_values: list[int], run_inference: bool, retune: bool = False) -> str: + available = ort.get_available_providers() + if "CUDAExecutionProvider" not in available: + raise RuntimeError( + f"CUDAExecutionProvider is not available in this onnxruntime build. Available providers: {available}" + ) + + print(f"Creating CUDA session for {model} (this profiles the M buckets)...") + sess_options = _make_session_options(output_prefix, m_values, retune) + with tempfile.TemporaryDirectory(prefix="ort_fpa_intb_tune_") as directory: + results_prefix = os.path.join(directory, "selected") + sess_options.add_session_config_entry("ep.cuda.gemm_tactic_cache_tuning_results_prefix", results_prefix) + session = ort.InferenceSession(model, sess_options, providers=["CUDAExecutionProvider"], enable_fallback=False) + try: + if "CUDAExecutionProvider" not in session.get_providers(): + raise RuntimeError("Tuning session did not activate CUDAExecutionProvider") + session.disable_fallback() + device_id = int(session.get_provider_options().get("CUDAExecutionProvider", {}).get("device_id", "0")) + signature = _current_signature(device_id) + + if run_inference: + for m in m_values: + try: + feeds = _make_dummy_inputs(session, m) + session.run_with_ort_values(None, feeds) + print(f" ran dummy inference for M={m}") + except Exception as exc: + print(f" skipped dummy inference for M={m}: {exc}") + finally: + # Both initial and lazy tactics reach disk at CUDA EP teardown. + del session + gc.collect() + + results_path = results_prefix + _CACHE_TABLE_SUFFIX + if not os.path.exists(results_path): + raise RuntimeError("This session did not select any successful fpA_intB MatMulNBits tactics for the model.") + selected = _summarize_cache(results_path, signature, display=False) + cache_path = output_prefix + _CACHE_TABLE_SUFFIX + persisted = _summarize_cache(cache_path, signature, display=False) + if not selected.issubset(persisted): + raise RuntimeError(f"Selected tactics for this model were not persisted in {cache_path}") + _summarize_cache(cache_path, signature) + print(f" selected this run: {len(selected)}") + return cache_path + + +def _build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Offline tactic tuning for MatMulNBits fpA_intB CUDA kernels.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--model", required=True, help="Path to the ONNX model to tune.") + parser.add_argument( + "--output-prefix", + required=True, + help="Cache file prefix. Writes '.matmulnbits_fpa_intb.tsv'.", + ) + parser.add_argument( + "--m-values", + default="1,2,4,8,16,32,64,128,256,512,1024,2048", + help="Comma-separated M buckets to profile (sets ep.cuda.fpa_intb_profile_m).", + ) + parser.add_argument( + "--no-inference", + action="store_true", + help="Only trigger construction-time profiling; skip the dummy inference loop.", + ) + parser.add_argument( + "--retune", + action="store_true", + help="Ignore tactics already in the cache file and profile every bucket again.", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = _build_arg_parser().parse_args(argv) + + if not os.path.exists(args.model): + print(f"ERROR: model not found: {args.model}", file=sys.stderr) + return 2 + + m_values = _parse_m_values(args.m_values) + if not m_values: + print("ERROR: --m-values did not contain any positive integers.", file=sys.stderr) + return 2 + + output_prefix = os.path.abspath(args.output_prefix) + out_dir = os.path.dirname(output_prefix) + if out_dir: + os.makedirs(out_dir, exist_ok=True) + + tune( + model=args.model, + output_prefix=output_prefix, + m_values=m_values, + run_inference=not args.no_inference, + retune=args.retune, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/onnxruntime/python/tools/transformers/float16.py b/onnxruntime/python/tools/transformers/float16.py index b55fdef09d693..ffe094ea0b7c2 100644 --- a/onnxruntime/python/tools/transformers/float16.py +++ b/onnxruntime/python/tools/transformers/float16.py @@ -210,13 +210,16 @@ def convert_float_to_float16( if isinstance(model, str): model_path = model if version.parse(onnx.__version__) >= version.parse("1.8.0") and not disable_shape_infer: - # shape_infer_model_path should be in the same folder of model_path - with tempfile.NamedTemporaryFile(dir=os.path.dirname(model_path)) as tmpfile: + # Keep external data paths relative to the original model and close the handle before ONNX reopens it. + with tempfile.NamedTemporaryFile(dir=os.path.dirname(model_path), delete=False) as tmpfile: shape_infer_model_path = tmpfile.name + try: # infer_shapes_path can be used for model >2GB, and infer_shapes cannot. infer_shapes_path(model_path, shape_infer_model_path) model = onnx.load(shape_infer_model_path) disable_shape_infer = True + finally: + os.remove(shape_infer_model_path) else: model = onnx.load(model_path) @@ -430,6 +433,8 @@ def convert_float_to_float16( f"initializer is used by both fp32 and fp16 nodes. Consider add these nodes to block list:{value.fp16_nodes}" ) + cast_output_names: dict[tuple[str, int], str] = {} + # Some operators have data type fixed as float for some input. Add a float16 to float cast for those inputs. for node in mixed_float_type_node_list: for i, input_name in enumerate(node.input): @@ -437,18 +442,21 @@ def convert_float_to_float16( continue for value_info in value_info_list: if input_name == value_info.name: - # create new value_info for current node's new input name - new_value_info = model.graph.value_info.add() - new_value_info.CopyFrom(value_info) - output_name = input_name + "_cast_to_fp32" - new_value_info.name = output_name - new_value_info.type.tensor_type.elem_type = TensorProto.FLOAT - # add Cast node (from tensor(float16) to tensor(float) before current node - node_name = input_name + "_cast_to_fp32_node" - new_node = [helper.make_node("Cast", [input_name], [output_name], to=1, name=node_name)] - model.graph.node.extend(new_node) + cast_key = (input_name, TensorProto.FLOAT) + if cast_key not in cast_output_names: + # create new value_info for the shared Cast output + new_value_info = model.graph.value_info.add() + new_value_info.CopyFrom(value_info) + output_name = input_name + "_cast_to_fp32" + new_value_info.name = output_name + new_value_info.type.tensor_type.elem_type = TensorProto.FLOAT + # add Cast node (from tensor(float16) to tensor(float) before current node + node_name = input_name + "_cast_to_fp32_node" + new_node = [helper.make_node("Cast", [input_name], [output_name], to=1, name=node_name)] + model.graph.node.extend(new_node) + cast_output_names[cast_key] = output_name # change current node's input name - node.input[i] = output_name + node.input[i] = cast_output_names[cast_key] break accuracy_type = TensorProto.BFLOAT16 if use_bfloat16_as_blocked_nodes_dtype else TensorProto.FLOAT @@ -461,18 +469,24 @@ def convert_float_to_float16( input_name = node.input[i] for value_info in value_info_list: if input_name == value_info.name: - # create new value_info for current node's new input name - new_value_info = model.graph.value_info.add() - new_value_info.CopyFrom(value_info) - output_name = input_name + "_cast_to_fp32" - new_value_info.name = output_name - new_value_info.type.tensor_type.elem_type = accuracy_type - # add Cast node (from tensor(float16) to tensor(float) before current node - node_name = input_name + "_cast_to_fp32_node" - new_node = [helper.make_node("Cast", [input_name], [output_name], to=accuracy_type, name=node_name)] - model.graph.node.extend(new_node) + cast_key = (input_name, accuracy_type) + if cast_key not in cast_output_names: + # create new value_info for the shared Cast output + new_value_info = model.graph.value_info.add() + new_value_info.CopyFrom(value_info) + cast_type_name = "fp32" if accuracy_type == TensorProto.FLOAT else "bfloat16" + output_name = input_name + f"_cast_to_{cast_type_name}" + new_value_info.name = output_name + new_value_info.type.tensor_type.elem_type = accuracy_type + # add Cast node (from tensor(float16) to tensor(float) before current node + node_name = input_name + f"_cast_to_{cast_type_name}_node" + new_node = [ + helper.make_node("Cast", [input_name], [output_name], to=accuracy_type, name=node_name) + ] + model.graph.node.extend(new_node) + cast_output_names[cast_key] = output_name # change current node's input name - node.input[i] = output_name + node.input[i] = cast_output_names[cast_key] break # if output's name is in the value_info_list meaning output is tensor(float16) type, insert a float to # float16 Cast node after the node, change current node's output name and create new value_info for the new name diff --git a/onnxruntime/test/contrib_ops/activation_op_test.cc b/onnxruntime/test/contrib_ops/activation_op_test.cc index f4223206cdc1d..9f0a4199c3749 100644 --- a/onnxruntime/test/contrib_ops/activation_op_test.cc +++ b/onnxruntime/test/contrib_ops/activation_op_test.cc @@ -6,7 +6,9 @@ #include "core/common/cpuid_info.h" #include "test/providers/provider_test_utils.h" #include "test/providers/cpu/activation/activation_op_test.h" +#include #include +#include #include #if defined(USE_DNNL) @@ -168,5 +170,62 @@ TEST_F(ActivationOpTest, QuickGelu) { } } +#if defined(USE_CUDA) +template +static void CheckCudaQuickGeluPrecision() { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP unavailable"; + for (float alpha : {1.0f, 1.702f, -1.702f, 0.0f}) { + SCOPED_TRACE(alpha); + std::vector input; + std::vector expected; + // Include float exp's subnormal/underflow region while the fused product is representable. + for (float value : {-128.0f, -108.0f, -104.0f, -100.0f, -96.0f, -88.0f, -87.5f, -87.0f, + -80.0f, -64.0f, -60.0f, -52.0f, -32.0f, -24.0f, -20.0f, -16.0f, -12.0f, -8.0f, -4.0f, + -1.0f, -0.125f, 0.0f, 0.125f, 1.0f, 4.0f, 8.0f, + 12.0f, 16.0f, 20.0f, 24.0f, 32.0f, 52.0f, 60.0f, 64.0f, 80.0f, + 87.0f, 87.5f, 88.0f, 96.0f, 100.0f, 104.0f, 108.0f, 128.0f}) { + input.emplace_back(value); + const double scaled = static_cast(value) * static_cast(alpha); + const double exponent = std::exp(-std::abs(scaled)); + const double sigmoid = scaled >= 0 ? 1.0 / (1.0 + exponent) : exponent / (1.0 + exponent); + const double result = static_cast(value) * sigmoid; + if constexpr (std::is_same_v) { + expected.emplace_back(result); + } else { + expected.emplace_back(static_cast(result)); + } + } + OpTester tester("QuickGelu", 1, kMSDomain); + tester.AddAttribute("alpha", alpha); + const std::vector dimensions{static_cast(input.size())}; + tester.AddInput("X", dimensions, input); + tester.AddOutput("Y", dimensions, expected); + tester.SetOutputAbsErr("Y", 0.0f); + tester.SetOutputRelErr("Y", std::is_same_v ? 2e-14f : 2e-6f); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + +TEST_F(ActivationOpTest, CudaQuickGeluFloatPrecision) { + CheckCudaQuickGeluPrecision(); +} + +TEST_F(ActivationOpTest, CudaQuickGeluDoublePrecision) { + CheckCudaQuickGeluPrecision(); +} + +TEST_F(ActivationOpTest, CudaQuickGeluFloat16Precision) { + CheckCudaQuickGeluPrecision(); +} + +TEST_F(ActivationOpTest, CudaQuickGeluBFloat16Precision) { + CheckCudaQuickGeluPrecision(); +} +#endif + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc index eeb48a0579b58..f857d895570bb 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -1528,7 +1528,9 @@ std::unique_ptr TryGetEpWithVarlenCausalConvWithState() { } std::unique_ptr GetWebGpuEpWithTestStorageBufferBindingSize(uint64_t max_size) { - return WebGpuExecutionProviderWithTestStorageBufferBindingSize(max_size); + WebGpuExecutionProviderTestOptions test_options{}; + test_options.max_storage_buffer_binding_size = max_size; + return WebGpuExecutionProviderWithTestOptions(test_options); } // Transpose a single request's (channels, length) reference block to the token-major diff --git a/onnxruntime/test/contrib_ops/conv_unknown_activation_test.cc b/onnxruntime/test/contrib_ops/conv_unknown_activation_test.cc new file mode 100644 index 0000000000000..1285a1291fd42 --- /dev/null +++ b/onnxruntime/test/contrib_ops/conv_unknown_activation_test.cc @@ -0,0 +1,92 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include + +#include "gtest/gtest.h" + +#include "core/graph/model.h" +#include "core/session/inference_session.h" +#include "test/test_environment.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime { +namespace test { + +namespace { + +// Builds a FusedConv with the given activation attributes and returns the Initialize() status. +Status InitializeFusedConv(const std::string& activation, const std::vector& activation_params, + std::unique_ptr webgpu_ep) { + constexpr int64_t kChannels = 4; + auto& logger = DefaultLoggingManager().DefaultLogger(); + Model model("ConvUnknownActivation", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 17}, {kMSDomain, 1}}, {}, logger); + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + + auto* input = builder.MakeInput({1, kChannels, 8}, -1.0f, 1.0f); + auto* weight = builder.MakeInitializer({kChannels, kChannels, 3}, -1.0f, 1.0f); + auto* output = builder.MakeOutput(); + Node& conv = builder.AddNode("FusedConv", {input, weight}, {output}, kMSDomain); + conv.AddAttribute("kernel_shape", std::vector{3}); + conv.AddAttribute("pads", std::vector{1, 1}); + conv.AddAttribute("strides", std::vector{1}); + conv.AddAttribute("dilations", std::vector{1}); + conv.AddAttribute("group", static_cast(1)); + conv.AddAttribute("activation", activation); + if (!activation_params.empty()) { + conv.AddAttribute("activation_params", activation_params); + } + + builder.SetGraphOutputs(); + ORT_RETURN_IF_ERROR(graph.Resolve()); + + std::string model_data; + model.ToProto().SerializeToString(&model_data); + + SessionOptions so; + so.session_logid = "Conv_WebGPU.FusedActivation"; + InferenceSession session{so, GetEnvironment()}; + ORT_RETURN_IF_ERROR(session.RegisterExecutionProvider(std::move(webgpu_ep))); + ORT_RETURN_IF_ERROR(session.Load(model_data.data(), static_cast(model_data.size()))); + return session.Initialize(); +} + +} // namespace + +// ConvActivationFusion writes the attribute from the core library while the activation set lives +// in the EP, so a core library newer than the EP binary can fuse an activation the kernel has +// never heard of. That has to report the activation by name, not a bare failed assertion. +TEST(Conv_WebGPU, UnknownFusedActivationIsReportedByName) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + + const Status status = InitializeFusedConv("NotAnActivation", {}, std::move(webgpu_ep)); + ASSERT_FALSE(status.IsOK()) << "the WebGPU EP has no shader for 'NotAnActivation'"; + EXPECT_NE(status.ErrorMessage().find("NotAnActivation"), std::string::npos) + << "the error must name the activation it could not fuse, got: " << status.ErrorMessage(); +} + +// Every parse failure keeps its Status message, not only the unknown-name one. +TEST(Conv_WebGPU, FusedActivationParamCountMismatchIsReported) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + + // LeakyRelu takes exactly one parameter. + const Status status = InitializeFusedConv("LeakyRelu", {0.1f, 0.2f}, std::move(webgpu_ep)); + ASSERT_FALSE(status.IsOK()); + EXPECT_NE(status.ErrorMessage().find("activation_params"), std::string::npos) + << "the error must name the attribute that is wrong, got: " << status.ErrorMessage(); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc index e7077e8073cd8..7f743a14faf4d 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc @@ -675,33 +675,30 @@ TEST(FpAIntBGemvTest, SupportUsesDeviceAndKernelArchitectures) { TEST(FpAIntBGemvTest, TacticCacheSeparatesDeviceAndQuantization) { using TacticCache = std::unordered_map; const wo_profile::GemmIdCore base(384, 1536, onnxruntime::llm::nvinfer::DataType::kHALF, - 80, 80, false, 0, 0, 4, 32, false, false); + 80, 4, 32, false, true); std::vector ids{base}; - for (int field = 0; field < 8; ++field) { + for (int field = 0; field < 7; ++field) { auto id = base; switch (field) { case 0: id.device_id = 1; break; case 1: - id.packing_sm = 90; + id.sm = 90; break; case 2: - id.device_sm = 90; - break; - case 3: id.dtype = onnxruntime::llm::nvinfer::DataType::kBF16; break; - case 4: + case 3: id.quant_bits = 8; break; - case 5: + case 4: id.group_size = 64; break; - case 6: + case 5: id.has_bias = true; break; - case 7: + case 6: id.has_zeros = true; break; } @@ -769,7 +766,7 @@ TEST(Int4DecodeTileTest, ProfilerEligibilityAndFallback) { profiler.setDecodeInterleave(4); for (auto type : {wo::KernelType::FP16Int4Groupwise, wo::KernelType::BF16Int4Groupwise}) { for (int arch : {75, 80, 86, 89, 90, 100, 120}) { - profiler.setCudaKernelType(type, 80, arch); + profiler.setCudaKernelType(type, arch); std::set variants; for (const auto& tactic : profiler.getTactics(1, 384, 1536)) { if (tactic.enableCudaKernel) variants.insert(tactic.cudaKernelVariant); @@ -811,7 +808,7 @@ TEST(Int4DecodeTileTest, ProfilerEligibilityAndFallback) { } // Upstream opt-in modes must not suppress M=1 candidates or admit them at M=5..8. - profiler.setCudaKernelType(wo::KernelType::FP16Int4Groupwise, 80, 80); + profiler.setCudaKernelType(wo::KernelType::FP16Int4Groupwise, 80); for (bool wave_aware : {false, true}) { profiler.setWaveAwareGemv(wave_aware); for (int mode : {0, 1, 2}) { @@ -861,8 +858,10 @@ TEST(FpAIntBGemvTest, WaveAwareDispatchUsesSyntheticSmCount) { // Keep upstream wave-aware and default tactic caches distinct after extending the key. TEST(FpAIntBGemvTest, TacticCacheSeparatesWaveAwareMode) { using TacticCache = std::unordered_map; - wo_profile::GemmIdCore const default_id(10240, 4096, onnxruntime::llm::nvinfer::DataType::kHALF, 80, false); - wo_profile::GemmIdCore const wave_aware_id(10240, 4096, onnxruntime::llm::nvinfer::DataType::kHALF, 80, true); + wo_profile::GemmIdCore const default_id(10240, 4096, onnxruntime::llm::nvinfer::DataType::kHALF, + 80, 4, 128, false, true, false, "", false); + wo_profile::GemmIdCore const wave_aware_id(10240, 4096, onnxruntime::llm::nvinfer::DataType::kHALF, + 80, 4, 128, false, true, false, "", true); TacticCache cache{{default_id, 4}, {wave_aware_id, 2}}; ASSERT_EQ(cache.size(), 2u); @@ -880,7 +879,9 @@ TEST(FpAIntBGemvTest, TacticCacheSeparatesPairedAndWaveAwareModes) { for (bool wave_aware : {false, true}) { for (int mode : {0, 1, 2}) { wo_profile::GemmIdCore const id(512, 1024, onnxruntime::llm::nvinfer::DataType::kHALF, - 80, wave_aware, mode); + 80, 4, 128, false, true, false, "", wave_aware, mode); + EXPECT_EQ(id.wave_aware, wave_aware); + EXPECT_EQ(id.tag, mode); cache->createMProfileMap(id); CutlassGemmConfig tactic; tactic.enableCudaKernel = true; @@ -894,7 +895,7 @@ TEST(FpAIntBGemvTest, TacticCacheSeparatesPairedAndWaveAwareModes) { for (bool wave_aware : {false, true}) { for (int mode : {0, 1, 2}) { wo_profile::GemmIdCore const id(512, 1024, onnxruntime::llm::nvinfer::DataType::kHALF, - 80, wave_aware, mode); + 80, 4, 128, false, true, false, "", wave_aware, mode); for (int m : {5, 6, 7, 8}) { auto const tactic = profiler.getBestConfig(m, id); ASSERT_TRUE(tactic.has_value()); diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc new file mode 100644 index 0000000000000..9f807df4eefa5 --- /dev/null +++ b/onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc @@ -0,0 +1,1365 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Unit tests for the fpA_intB / MatMulNBits GEMM tactic cache utilities. +// GemmTacticCacheTest uses synthetic signatures and needs no GPU. +// GemmTacticCacheCudaTest additionally exercises the profiler's CUDA events. +// +// Built into onnxruntime_providers_cuda_ut (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON); run like: +// ./onnxruntime_provider_test --gtest_filter=CUDA_EP_Unittest.All +// Plugin builds register these directly: --gtest_filter=GemmTacticCache*.* +#if USE_FPA_INTB_GEMM +#include + +#include +#include +#include // NOLINT(build/c++17) +#include +#include +#include +#include +#include +#include +#include + +#ifdef _WIN32 +#include +#endif + +#include "contrib_ops/cuda/llm/gemm_tactic_cache.h" +#include "contrib_ops/cuda/llm/gemm_profiler.h" +#if !defined(ORT_UNIT_TEST_HAS_CUDA_PLUGIN_EP) +#include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" +#include "contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h" +#include "core/providers/cuda/cuda_execution_provider.h" +#include "core/providers/cuda/cuda_execution_provider_info.h" +#endif +#include "test/util/include/scoped_env_vars.h" + +namespace onnxruntime { +namespace test { + +namespace gc = onnxruntime::llm::gemm_cache; +using onnxruntime::llm::cutlass_extensions::ClusterShape; +using onnxruntime::llm::cutlass_extensions::CutlassGemmConfig; +using onnxruntime::llm::cutlass_extensions::CutlassTileConfig; +using onnxruntime::llm::cutlass_extensions::CutlassTileConfigSM90; +using onnxruntime::llm::cutlass_extensions::MainloopScheduleType; +using onnxruntime::llm::cutlass_extensions::SplitKStyle; + +namespace { + +#if !defined(ORT_UNIT_TEST_HAS_CUDA_PLUGIN_EP) +class CacheTestProfiler : public onnxruntime::llm::kernels::weight_only::WeightOnlyGroupwiseQuantGemmPluginProfiler { + public: + size_t launches = 0; + + protected: + bool hasRunnerContext() const override { return true; } + void runTactic(int, int, int, const CutlassGemmConfig&, char*, const cudaStream_t&) override { ++launches; } + size_t computeTmpSize(size_t, size_t, size_t) override { return 0; } + std::vector getTactics(int, int, int) const override { return {CutlassGemmConfig{}}; } + bool checkTactic(int, int, int, const CutlassGemmConfig&) const override { return true; } + float getSelectionTime(int, int, int, const CutlassGemmConfig&, float time) const override { return time; } +}; +#endif // !defined(ORT_UNIT_TEST_HAS_CUDA_PLUGIN_EP) + +gc::HardwareSignature MakeSignature(const std::string& device_name = "TEST GPU 4090", int sm = 80) { + gc::HardwareSignature sig; + sig.device_name = device_name; + sig.sm = sm; + sig.multiprocessor_count = 108; + sig.cuda_runtime = 13000; + sig.cuda_driver = 13000; + sig.ort_version = "1.28.0"; + sig.ort_git_commit = "deadbeef"; + sig.ort_build_config = "Release"; + return sig; +} + +gc::MatMulNBitsKey MakeKey(int n_16b = 3072, int k = 4096) { + gc::MatMulNBitsKey key; + key.n_16b = n_16b; + key.k = k; + key.activation_dtype = "half"; + key.weight_type = "uint4b_t"; + key.bits = 4; + key.block_size = 64; + key.has_zero_points = true; + key.zero_point_dtype = "uint4b_t"; + key.gemv_enabled = true; + key.packing_sm = 80; + return key; +} + +CutlassGemmConfig MakeSm80Config() { + CutlassGemmConfig c; + c.sm_version = 80; + c.tile_config_sm80 = CutlassTileConfig::CtaShape64x128x64_WarpShape32x64x64; + c.split_k_style = SplitKStyle::SPLIT_K_SERIAL; + c.split_k_factor = 4; + c.stages = 3; + c.is_tma_warp_specialized = false; + c.enableCudaKernel = false; + return c; +} + +CutlassGemmConfig MakeSm90Config() { + CutlassGemmConfig c; + c.sm_version = 90; + c.tile_config_sm90 = CutlassTileConfigSM90::CtaShape128x128x128B; + c.mainloop_schedule = MainloopScheduleType::COOPERATIVE; + c.cluster_shape = ClusterShape::ClusterShape_2x1x1; + c.is_tma_warp_specialized = true; + c.enableCudaKernel = false; + return c; +} + +void ExpectConfigEqual(const CutlassGemmConfig& a, const CutlassGemmConfig& b) { + EXPECT_EQ(a.sm_version, b.sm_version); + EXPECT_EQ(static_cast(a.tile_config_sm80), static_cast(b.tile_config_sm80)); + EXPECT_EQ(static_cast(a.tile_config_sm90), static_cast(b.tile_config_sm90)); + EXPECT_EQ(static_cast(a.tile_config_sm100), static_cast(b.tile_config_sm100)); + EXPECT_EQ(static_cast(a.tile_config_sm120), static_cast(b.tile_config_sm120)); + EXPECT_EQ(static_cast(a.split_k_style), static_cast(b.split_k_style)); + EXPECT_EQ(a.split_k_factor, b.split_k_factor); + EXPECT_EQ(a.stages, b.stages); + EXPECT_EQ(static_cast(a.cluster_shape), static_cast(b.cluster_shape)); + EXPECT_EQ(static_cast(a.mainloop_schedule), static_cast(b.mainloop_schedule)); + EXPECT_EQ(static_cast(a.epilogue_schedule), static_cast(b.epilogue_schedule)); + EXPECT_EQ(a.is_tma_warp_specialized, b.is_tma_warp_specialized); + EXPECT_EQ(a.enableCudaKernel, b.enableCudaKernel); + EXPECT_EQ(a.cudaKernelVariant, b.cudaKernelVariant); +} + +// Returns a unique temp file prefix and removes any leftover files from a prior run. +std::string UniqueTempPrefix(const std::string& tag) { + auto dir = std::filesystem::temp_directory_path(); + std::string prefix = (dir / ("ort_gemm_tactic_cache_test_" + tag)).string(); + std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + std::error_code ec; + std::filesystem::remove(file, ec); + std::filesystem::remove(file + ".lock", ec); + std::filesystem::remove(file + ".tmp", ec); + return prefix; +} + +void CleanUp(const std::string& file) { + std::error_code ec; + std::filesystem::remove(file, ec); + std::filesystem::remove(file + ".lock", ec); + std::filesystem::remove(file + ".tmp", ec); +} + +} // namespace + +TEST(GemmTacticCacheTest, TsvEncodeDecodeRoundTrip) { + const std::vector samples = { + "NVIDIA A100-SXM4-80GB", + "has\ttab", + "has\nnewline", + "has%percent", + "mix\t%\n\rall", + "", + }; + for (const auto& s : samples) { + std::string encoded = gc::TsvEncode(s); + EXPECT_EQ(encoded.find('\t'), std::string::npos) << "encoded still has tab: " << s; + EXPECT_EQ(encoded.find('\n'), std::string::npos) << "encoded still has newline: " << s; + EXPECT_EQ(gc::TsvDecode(encoded), s) << "round-trip failed for: " << s; + } +} + +TEST(GemmTacticCacheTest, TsvDecodeKeepsMalformedEscapes) { + EXPECT_EQ(gc::TsvDecode("%-1"), "%-1"); + EXPECT_EQ(gc::TsvDecode("% 9"), "% 9"); + EXPECT_EQ(gc::TsvDecode("%4"), "%4"); + EXPECT_EQ(gc::TsvDecode("a%41b"), "aAb"); +} + +TEST(GemmTacticCacheTest, ResolveFilePathSessionConfig) { + ScopedEnvironmentVariables scoped_env_vars{ + {{gc::kEnvCachePrefix, "/tmp/env_model"}, {gc::kEnvCacheDir, "/tmp/env_dir"}}}; + const auto sig = MakeSignature("NVIDIA H200", 90); + const std::string suffix = ".matmulnbits_fpa_intb.tsv"; + EXPECT_EQ(gc::MatMulNBitsTacticCache::ResolveFilePath("/tmp/dir", "", sig), + "/tmp/dir/NVIDIA_H200_sm90" + suffix); + EXPECT_EQ(gc::MatMulNBitsTacticCache::ResolveFilePath("", "/tmp/model", sig), "/tmp/model" + suffix); + // A prefix wins over a directory. + EXPECT_EQ(gc::MatMulNBitsTacticCache::ResolveFilePath("/tmp/dir", "/tmp/model", sig), "/tmp/model" + suffix); +} + +TEST(GemmTacticCacheTest, ResolveFilePathEnvironmentDisabled) { + const auto sig = MakeSignature(); + ScopedEnvironmentVariables scoped_env_vars{ + {{gc::kEnvCachePrefix, std::nullopt}, {gc::kEnvCacheDir, std::nullopt}}}; + EXPECT_TRUE(gc::MatMulNBitsTacticCache::ResolveFilePath("", "", sig).empty()); + + ScopedEnvironmentVariables empty_env_vars{ + {{gc::kEnvCachePrefix, ""}, {gc::kEnvCacheDir, ""}}}; + EXPECT_TRUE(gc::MatMulNBitsTacticCache::ResolveFilePath("", "", sig).empty()); +} + +TEST(GemmTacticCacheTest, ResolveFilePathEnvironmentFallback) { + const auto sig = MakeSignature("NVIDIA H200", 90); + const std::string suffix = ".matmulnbits_fpa_intb.tsv"; + const std::string dir = "/tmp/cache \xE4\xB8\xAD"; + const std::string prefix = "/tmp/model \xE4\xB8\xAD"; + ScopedEnvironmentVariables scoped_env_vars{ + {{gc::kEnvCachePrefix, std::nullopt}, {gc::kEnvCacheDir, dir}}}; + EXPECT_EQ(gc::MatMulNBitsTacticCache::ResolveFilePath("", "", sig), + dir + "/NVIDIA_H200_sm90" + suffix); + + { + ScopedEnvironmentVariables prefix_env_var{{{gc::kEnvCachePrefix, prefix}}}; + EXPECT_EQ(gc::MatMulNBitsTacticCache::ResolveFilePath("", "", sig), prefix + suffix); + } + + { + ScopedEnvironmentVariables prefix_only_env_vars{ + {{gc::kEnvCachePrefix, prefix}, {gc::kEnvCacheDir, std::nullopt}}}; + EXPECT_EQ(gc::MatMulNBitsTacticCache::ResolveFilePath("", "", sig), prefix + suffix); + } +} + +TEST(GemmTacticCacheTest, ConfigColumnsRoundTripSm80) { + const CutlassGemmConfig original = MakeSm80Config(); + std::vector row; + gc::AppendConfigColumns(row, original); + ASSERT_EQ(row.size(), static_cast(gc::kNumConfigColumns)); + + auto parsed = gc::ParseConfigColumns(row, 0); + ASSERT_TRUE(parsed.has_value()); + ASSERT_TRUE(parsed->has_value()); + ExpectConfigEqual(original, **parsed); +} + +TEST(GemmTacticCacheTest, ConfigColumnsRoundTripSm90) { + const CutlassGemmConfig original = MakeSm90Config(); + std::vector row; + gc::AppendConfigColumns(row, original); + auto parsed = gc::ParseConfigColumns(row, 0); + ASSERT_TRUE(parsed.has_value()); + ASSERT_TRUE(parsed->has_value()); + ExpectConfigEqual(original, **parsed); +} + +TEST(GemmTacticCacheTest, ConfigColumnsNullTacticRoundTrip) { + std::vector row; + gc::AppendConfigColumns(row, std::nullopt); // profiled bucket with no valid tactic + auto parsed = gc::ParseConfigColumns(row, 0); + ASSERT_TRUE(parsed.has_value()); // outer: the columns parsed + EXPECT_FALSE(parsed->has_value()); // inner: no valid tactic +} + +TEST(GemmTacticCacheTest, ConfigColumnsRoundTripPairedGemv) { + CutlassGemmConfig original; + original.enableCudaKernel = true; + original.cudaKernelVariant = 1; + std::vector row; + gc::AppendConfigColumns(row, original); + auto parsed = gc::ParseConfigColumns(row, 0); + ASSERT_TRUE(parsed.has_value() && parsed->has_value()); + ExpectConfigEqual(original, **parsed); + + row.back() = "8"; + EXPECT_FALSE(gc::ParseConfigColumns(row, 0).has_value()); + row.back() = "-1"; + EXPECT_FALSE(gc::ParseConfigColumns(row, 0).has_value()); +} + +TEST(GemmTacticCacheTest, ConfigColumnsRoundTripDecodeVariants) { + for (int variant = 2; variant <= 7; ++variant) { + auto config = MakeSm80Config(); + config.enableCudaKernel = true; + config.cudaKernelVariant = variant; + std::vector row; + gc::AppendConfigColumns(row, config); + const auto parsed = gc::ParseConfigColumns(row, 0); + ASSERT_TRUE(parsed.has_value() && parsed->has_value()); + ExpectConfigEqual(config, **parsed); + row[row.size() - 2] = "0"; + EXPECT_FALSE(gc::ParseConfigColumns(row, 0).has_value()); + } +} + +TEST(GemmTacticCacheTest, DecodeUpgradeRejectsOldSelectionVersion) { + const auto file = UniqueTempPrefix("decode_version") + ".matmulnbits_fpa_intb.tsv"; + gc::MatMulNBitsTacticCache writer(file, MakeSignature()); + writer.Put(MakeKey(), 1, MakeSm80Config()); + ASSERT_TRUE(writer.Flush().IsOK()); + std::ifstream input(file); + std::string contents((std::istreambuf_iterator(input)), std::istreambuf_iterator()); + const auto position = contents.find("# tactic_selection_version\t4"); + ASSERT_NE(position, std::string::npos); + contents.replace(position, std::string("# tactic_selection_version\t4").size(), "# tactic_selection_version\t3"); + std::istringstream old(contents); + gc::MatMulNBitsTacticCache reader(file, MakeSignature()); + ASSERT_TRUE(reader.Load(old).IsOK()); + EXPECT_TRUE(reader.GetAll(MakeKey()).empty()); + CleanUp(file); +} + +TEST(GemmTacticCacheTest, StoreLoadRoundTrip) { + const std::string prefix = UniqueTempPrefix("roundtrip"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::HardwareSignature sig = MakeSignature(); + const gc::MatMulNBitsKey key = MakeKey(); + + { + gc::MatMulNBitsTacticCache cache(file, sig); + cache.Put(key, 1, MakeSm80Config()); + cache.Put(key, 64, MakeSm90Config()); + cache.Put(key, 128, std::nullopt); // negative result + ASSERT_TRUE(cache.Flush().IsOK()); + } + + gc::MatMulNBitsTacticCache reloaded(file, sig); + ASSERT_TRUE(reloaded.Load().IsOK()); + + auto c1 = reloaded.Get(key, 1); + ASSERT_TRUE(c1.has_value()); + ASSERT_TRUE(c1->has_value()); + ExpectConfigEqual(MakeSm80Config(), **c1); + + auto c64 = reloaded.Get(key, 64); + ASSERT_TRUE(c64.has_value()); + ASSERT_TRUE(c64->has_value()); + ExpectConfigEqual(MakeSm90Config(), **c64); + + auto c128 = reloaded.Get(key, 128); + EXPECT_FALSE(c128.has_value()); // Temporary profiling failures are never persisted. + + EXPECT_FALSE(reloaded.Get(key, 999).has_value()); // never profiled + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, SignatureMismatchRejected) { + const std::string prefix = UniqueTempPrefix("sigmismatch"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::MatMulNBitsKey key = MakeKey(); + + { + gc::MatMulNBitsTacticCache cache(file, MakeSignature("NVIDIA RTX 4090", 89)); + cache.Put(key, 1, MakeSm80Config()); + ASSERT_TRUE(cache.Flush().IsOK()); + } + + // Different device name -> file must be rejected on load. + gc::MatMulNBitsTacticCache other(file, MakeSignature("NVIDIA RTX 4060", 89)); + ASSERT_TRUE(other.Load().IsOK()); + EXPECT_FALSE(other.Get(key, 1).has_value()); + + // Same signature -> accepted. + gc::MatMulNBitsTacticCache same(file, MakeSignature("NVIDIA RTX 4090", 89)); + ASSERT_TRUE(same.Load().IsOK()); + EXPECT_TRUE(same.Get(key, 1).has_value()); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, GitCommitAndBuildConfigNotStrict) { + const std::string prefix = UniqueTempPrefix("nonstrict"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::MatMulNBitsKey key = MakeKey(); + + // Write with one git commit / build config. + gc::HardwareSignature writer = MakeSignature(); + writer.ort_git_commit = "commit_aaaa"; + writer.ort_build_config = "Release"; + { + gc::MatMulNBitsTacticCache cache(file, writer); + cache.Put(key, 1, MakeSm80Config()); + ASSERT_TRUE(cache.Flush().IsOK()); + } + + // A reader with a DIFFERENT git commit and build config, but the same device/sm/cuda_runtime/ + // ort_version, must still accept the cache: those two fields are diagnostic-only. + gc::HardwareSignature reader = MakeSignature(); + reader.ort_git_commit = "commit_bbbb"; + reader.ort_build_config = "Debug"; + gc::MatMulNBitsTacticCache reloaded(file, reader); + ASSERT_TRUE(reloaded.Load().IsOK()); + EXPECT_TRUE(reloaded.Get(key, 1).has_value()); + + // A different ort_version is still rejected (it remains part of the strict guard). + gc::HardwareSignature other_version = MakeSignature(); + other_version.ort_version = "9.9.9"; + gc::MatMulNBitsTacticCache rejected(file, other_version); + ASSERT_TRUE(rejected.Load().IsOK()); + EXPECT_FALSE(rejected.Get(key, 1).has_value()); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, AppendedColumnTolerated) { + const std::string prefix = UniqueTempPrefix("appendcol"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::HardwareSignature sig = MakeSignature(); + const gc::MatMulNBitsKey key = MakeKey(); + + { + gc::MatMulNBitsTacticCache cache(file, sig); + cache.Put(key, 8, MakeSm80Config()); + ASSERT_TRUE(cache.Flush().IsOK()); + } + + // Simulate a future writer that appended an extra trailing column to the header and rows. + std::vector lines; + { + std::ifstream in(file); + std::string line; + bool header_seen = false; + while (std::getline(in, line)) { + if (!line.empty() && line[0] != '#') { + line += (header_seen ? "\t999" : "\tfuture_col"); + header_seen = true; + } + lines.push_back(line); + } + } + { + std::ofstream out(file, std::ios::trunc); + for (const auto& l : lines) { + out << l << '\n'; + } + } + + // Reader must map by column name and ignore the unknown appended column. + gc::MatMulNBitsTacticCache reloaded(file, sig); + ASSERT_TRUE(reloaded.Load().IsOK()); + auto c8 = reloaded.Get(key, 8); + ASSERT_TRUE(c8.has_value()); + ASSERT_TRUE(c8->has_value()); + ExpectConfigEqual(MakeSm80Config(), **c8); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, FlushMergesSequentialRows) { + const std::string prefix = UniqueTempPrefix("merge"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::HardwareSignature sig = MakeSignature(); + const gc::MatMulNBitsKey key = MakeKey(); + + // First writer persists bucket 1. + { + gc::MatMulNBitsTacticCache a(file, sig); + a.Put(key, 1, MakeSm80Config()); + ASSERT_TRUE(a.Flush().IsOK()); + } + + // Second writer (independent instance) persists bucket 64; Flush must merge, not clobber. + { + gc::MatMulNBitsTacticCache b(file, sig); + b.Put(key, 64, MakeSm90Config()); + ASSERT_TRUE(b.Flush().IsOK()); + } + + gc::MatMulNBitsTacticCache reloaded(file, sig); + ASSERT_TRUE(reloaded.Load().IsOK()); + EXPECT_TRUE(reloaded.Get(key, 1).has_value()); + EXPECT_TRUE(reloaded.Get(key, 64).has_value()); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, SelectionVersionMismatchRejected) { + const std::string prefix = UniqueTempPrefix("selversion"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::HardwareSignature sig = MakeSignature(); + const gc::MatMulNBitsKey key = MakeKey(); + { + gc::MatMulNBitsTacticCache cache(file, sig); + cache.Put(key, 1, MakeSm80Config()); + ASSERT_TRUE(cache.Flush().IsOK()); + } + + std::vector lines; + { + std::ifstream in(file); + std::string line; + while (std::getline(in, line)) { + if (line.rfind("# tactic_selection_version\t", 0) == 0) { + line = "# tactic_selection_version\tstale"; + } + lines.push_back(line); + } + } + { + std::ofstream out(file, std::ios::trunc); + for (const auto& l : lines) { + out << l << '\n'; + } + } + + gc::MatMulNBitsTacticCache reloaded(file, sig); + ASSERT_TRUE(reloaded.Load().IsOK()); + EXPECT_FALSE(reloaded.Get(key, 1).has_value()); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, FlushCreatesMissingDirectory) { + const auto dir = std::filesystem::temp_directory_path() / "ort_gemm_tactic_cache_test_newdir"; + std::error_code ec; + std::filesystem::remove_all(dir, ec); + const std::string file = (dir / "nested" / "cache.matmulnbits_fpa_intb.tsv").string(); + const gc::HardwareSignature sig = MakeSignature(); + const gc::MatMulNBitsKey key = MakeKey(); + { + gc::MatMulNBitsTacticCache cache(file, sig); + cache.Put(key, 1, MakeSm80Config()); + ASSERT_TRUE(cache.Flush().IsOK()); + } + + gc::MatMulNBitsTacticCache reloaded(file, sig); + ASSERT_TRUE(reloaded.Load().IsOK()); + EXPECT_TRUE(reloaded.Get(key, 1).has_value()); + + // Only the cache and its lock file remain; per-writer temp files are renamed away. + size_t entries = 0; + for (const auto& entry : std::filesystem::directory_iterator(dir / "nested")) { + static_cast(entry); + ++entries; + } + EXPECT_EQ(entries, 2u); + std::filesystem::remove_all(dir, ec); +} + +// The process-global profile map must not share tactics across variants that the persistent key separates. +TEST(GemmTacticCacheTest, GemmIdCoreSeparatesQuantVariants) { + using onnxruntime::llm::kernels::weight_only::GemmIdCore; + using onnxruntime::llm::kernels::weight_only::GemmIdCoreHash; + const auto dtype = onnxruntime::llm::nvinfer::DataType::kHALF; + const GemmIdCore base(1024, 4096, dtype, 80, 4, 64, false, true); + + EXPECT_EQ(base, GemmIdCore(1024, 4096, dtype, 80, 4, 64, false, true)); + EXPECT_EQ(GemmIdCoreHash{}(base), GemmIdCoreHash{}(GemmIdCore(1024, 4096, dtype, 80, 4, 64, false, true))); + EXPECT_FALSE(base == GemmIdCore(1024, 4096, dtype, 80, 8, 64, false, true)); + EXPECT_FALSE(base == GemmIdCore(1024, 4096, dtype, 80, 4, 128, false, true)); + EXPECT_FALSE(base == GemmIdCore(1024, 4096, dtype, 80, 4, 64, true, true)); + EXPECT_FALSE(base == GemmIdCore(1024, 4096, dtype, 80, 4, 64, false, false)); + EXPECT_FALSE(base == GemmIdCore(1024, 4096, dtype, 80, 4, 64, false, true, true)); + EXPECT_FALSE(base == GemmIdCore(1024, 4096, dtype, 80, 4, 64, false, true, false, "NVIDIA RTX 4090")); + + // Same packing SM, different GPU: the process-global map must not mix their tactics. + const GemmIdCore rtx4090(1024, 4096, dtype, 80, 4, 64, false, true, false, "NVIDIA RTX 4090"); + const GemmIdCore rtx4060(1024, 4096, dtype, 80, 4, 64, false, true, false, "NVIDIA RTX 4060"); + EXPECT_FALSE(rtx4090 == rtx4060); + EXPECT_EQ(rtx4090, GemmIdCore(1024, 4096, dtype, 80, 4, 64, false, true, false, "NVIDIA RTX 4090")); + EXPECT_FALSE(base == GemmIdCore(1024, 4096, dtype, 80, 4, 64, false, true, false, "", true, 0)); + EXPECT_FALSE(base == GemmIdCore(1024, 4096, dtype, 80, 4, 64, false, true, false, "", false, 1)); + EXPECT_FALSE(base == GemmIdCore(1024, 4096, dtype, 80, 4, 64, false, true, false, "", false, 2)); +} + +TEST(GemmTacticCacheTest, GemvOptionsArePartOfPersistentKey) { + const std::string file = UniqueTempPrefix("gemv_options") + ".matmulnbits_fpa_intb.tsv"; + const auto sig = MakeSignature(); + const auto base = MakeKey(); + auto paired = base; + paired.paired_gemv_mode = 1; + auto forced = base; + forced.paired_gemv_mode = 2; + auto wave_aware = base; + wave_aware.wave_aware_gemv = true; + CutlassGemmConfig paired_config; + paired_config.enableCudaKernel = true; + paired_config.cudaKernelVariant = 1; + { + gc::MatMulNBitsTacticCache cache(file, sig); + cache.Put(base, 8, MakeSm80Config()); + cache.Put(paired, 8, paired_config); + cache.Put(forced, 8, paired_config); + cache.Put(wave_aware, 8, MakeSm90Config()); + ASSERT_TRUE(cache.Flush().IsOK()); + } + gc::MatMulNBitsTacticCache reloaded(file, sig); + ASSERT_TRUE(reloaded.Load().IsOK()); + for (const auto& key : {base, paired, forced, wave_aware}) { + auto config = reloaded.Get(key, 8); + ASSERT_TRUE(config.has_value() && config->has_value()); + ExpectConfigEqual(key.paired_gemv_mode != 0 ? paired_config + : key.wave_aware_gemv ? MakeSm90Config() + : MakeSm80Config(), + **config); + } + CleanUp(file); +} + +#if !defined(ORT_UNIT_TEST_HAS_CUDA_PLUGIN_EP) +TEST(GemmTacticCacheTest, CachedGemvRespectsPairedMode) { + using namespace onnxruntime::llm::kernels::weight_only; + class Profiler : public WeightOnlyGroupwiseQuantGemmPluginProfiler { + public: + using WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic; + } profiler; + CutlassGemmConfig paired; + paired.enableCudaKernel = true; + paired.cudaKernelVariant = 1; + CutlassGemmConfig sequential = paired; + sequential.cudaKernelVariant = 0; + for (int mode : {0, 1, 2}) { + profiler.setPairedGemvMode(mode); + EXPECT_EQ(profiler.checkTactic(8, 16, 128, paired), mode != 0); + EXPECT_EQ(profiler.checkTactic(8, 16, 128, sequential), mode != 2); + EXPECT_EQ(profiler.checkTactic(8, 16, 128, MakeSm80Config()), mode != 2); + EXPECT_FALSE(profiler.checkTactic(4, 16, 128, paired)); + EXPECT_TRUE(profiler.checkTactic(4, 16, 128, sequential)); + EXPECT_FALSE(profiler.checkTactic(16, 16, 128, paired)); + EXPECT_FALSE(profiler.checkTactic(16, 16, 128, sequential)); + } +} + +TEST(GemmTacticCacheTest, CachedDecodeVariantRequiresMatchingGeometry) { + using namespace onnxruntime::llm::kernels::weight_only; + class Profiler : public WeightOnlyGroupwiseQuantGemmPluginProfiler { + public: + using WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic; + } profiler; + profiler.setQuant(4, false, false); + profiler.setGroupSize(32); + profiler.setDecodeInterleave(4); + profiler.setCudaKernelType(KernelType::FP16Int4Groupwise, 80); + CutlassGemmConfig config; + config.enableCudaKernel = true; + for (int variant = 2; variant <= 7; ++variant) { + config.cudaKernelVariant = variant; + EXPECT_TRUE(profiler.checkTactic(1, 384, 1536, config)); + EXPECT_FALSE(profiler.checkTactic(8, 384, 1536, config)); + EXPECT_FALSE(profiler.checkTactic(1, 384, 1500, config)); + } + profiler.setDecodeInterleave(1); + EXPECT_FALSE(profiler.checkTactic(1, 384, 1536, config)); + profiler.setDecodeInterleave(4); + profiler.setQuant(4, true, false); + EXPECT_FALSE(profiler.checkTactic(1, 384, 1536, config)); +} +#endif + +TEST(GemmTacticCacheTest, BiasIsPartOfPersistentKey) { + const std::string prefix = UniqueTempPrefix("bias"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::HardwareSignature sig = MakeSignature(); + gc::MatMulNBitsKey no_bias = MakeKey(); + gc::MatMulNBitsKey with_bias = MakeKey(); + with_bias.has_bias = true; + + { + gc::MatMulNBitsTacticCache cache(file, sig); + cache.Put(no_bias, 1, MakeSm80Config()); + cache.Put(with_bias, 1, MakeSm90Config()); + ASSERT_TRUE(cache.Flush().IsOK()); + } + + gc::MatMulNBitsTacticCache reloaded(file, sig); + ASSERT_TRUE(reloaded.Load().IsOK()); + auto biasless = reloaded.Get(no_bias, 1); + auto biasful = reloaded.Get(with_bias, 1); + ASSERT_TRUE(biasless.has_value() && biasless->has_value()); + ASSERT_TRUE(biasful.has_value() && biasful->has_value()); + ExpectConfigEqual(MakeSm80Config(), **biasless); + ExpectConfigEqual(MakeSm90Config(), **biasful); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, FlushFailsWhenFileLockUnavailable) { + const std::string prefix = UniqueTempPrefix("lockfail"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::HardwareSignature sig = MakeSignature(); + + // A directory at the lock path makes lock acquisition fail on every platform. + std::error_code ec; + ASSERT_TRUE(std::filesystem::create_directory(file + ".lock", ec)) << ec.message(); + + gc::MatMulNBitsTacticCache cache(file, sig); + cache.Put(MakeKey(), 1, MakeSm80Config()); + EXPECT_FALSE(cache.Flush().IsOK()); + EXPECT_FALSE(std::filesystem::exists(file)); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, StaleReaderDoesNotOverwriteRetunedBucket) { + const auto file = UniqueTempPrefix("stale") + ".matmulnbits_fpa_intb.tsv"; + const auto key = MakeKey(); + gc::MatMulNBitsTacticCache first(file, MakeSignature()); + first.Put(key, 1, MakeSm80Config()); + ASSERT_TRUE(first.Flush().IsOK()); + gc::MatMulNBitsTacticCache stale(file, MakeSignature()); + ASSERT_TRUE(stale.Load().IsOK()); + first.Put(key, 1, MakeSm90Config()); + ASSERT_TRUE(first.Flush().IsOK()); + stale.Put(key, 64, MakeSm80Config()); + ASSERT_TRUE(stale.Flush().IsOK()); + ASSERT_TRUE(stale.Get(key, 1).has_value() && stale.Get(key, 1)->has_value()); + ExpectConfigEqual(**stale.Get(key, 1), MakeSm90Config()); + ASSERT_TRUE(stale.ReloadIfChanged().IsOK()); + ExpectConfigEqual(**stale.Get(key, 1), MakeSm90Config()); + gc::MatMulNBitsTacticCache result(file, MakeSignature()); + ASSERT_TRUE(result.Load().IsOK()); + ASSERT_TRUE(result.Get(key, 1).has_value()); + ExpectConfigEqual(**result.Get(key, 1), MakeSm90Config()); + EXPECT_TRUE(result.Get(key, 64).has_value()); + CleanUp(file); +} + +TEST(GemmTacticCacheTest, ConcurrentWritersPreserveAllBuckets) { + const auto file = UniqueTempPrefix("concurrent") + ".matmulnbits_fpa_intb.tsv"; + constexpr int writers = 8; + std::barrier start(writers); + std::vector threads; + for (int i = 1; i <= writers; ++i) { + threads.emplace_back([&, i] { + gc::MatMulNBitsTacticCache cache(file, MakeSignature()); + cache.Put(MakeKey(), i, MakeSm80Config()); + start.arrive_and_wait(); + EXPECT_TRUE(cache.Flush().IsOK()); + }); + } + for (auto& thread : threads) thread.join(); + gc::MatMulNBitsTacticCache result(file, MakeSignature()); + ASSERT_TRUE(result.Load().IsOK()); + EXPECT_EQ(result.GetAll(MakeKey()).size(), writers); + CleanUp(file); +} + +TEST(GemmTacticCacheTest, ProfilingFailureIsRetriedAfterReload) { + const auto file = UniqueTempPrefix("retry") + ".matmulnbits_fpa_intb.tsv"; + gc::MatMulNBitsTacticCache failed(file, MakeSignature()); + failed.Put(MakeKey(), 1, std::nullopt); + ASSERT_TRUE(failed.Flush().IsOK()); + gc::MatMulNBitsTacticCache healthy(file, MakeSignature()); + ASSERT_TRUE(healthy.Load().IsOK()); + ASSERT_FALSE(healthy.Get(MakeKey(), 1).has_value()); + healthy.Put(MakeKey(), 1, MakeSm80Config()); + ASSERT_TRUE(healthy.Flush().IsOK()); + gc::MatMulNBitsTacticCache result(file, MakeSignature()); + ASSERT_TRUE(result.Load().IsOK()); + ASSERT_TRUE(result.Get(MakeKey(), 1).has_value()); + CleanUp(file); +} + +TEST(GemmTacticCacheTest, PartialReadDoesNotAcceptRows) { + const auto file = UniqueTempPrefix("partial") + ".matmulnbits_fpa_intb.tsv"; + gc::MatMulNBitsTacticCache writer(file, MakeSignature()); + writer.Put(MakeKey(), 1, MakeSm80Config()); + ASSERT_TRUE(writer.Flush().IsOK()); + std::ifstream input(file); + std::string contents((std::istreambuf_iterator(input)), std::istreambuf_iterator()); + input.close(); + class FailingBuffer : public std::streambuf { + public: + explicit FailingBuffer(std::string& data) { setg(data.data(), data.data(), data.data() + data.size()); } + int_type underflow() override { throw std::runtime_error("injected read error after a complete row"); } + } buffer(contents); + std::istream broken(&buffer); + gc::MatMulNBitsTacticCache reader(file, MakeSignature()); + EXPECT_FALSE(reader.Load(broken).IsOK()); + EXPECT_TRUE(reader.GetAll(MakeKey()).empty()); + CleanUp(file); +} + +TEST(GemmTacticCacheTest, LegacyNegativeRowsAreIgnored) { + const auto file = UniqueTempPrefix("legacy_negative") + ".matmulnbits_fpa_intb.tsv"; + gc::MatMulNBitsTacticCache writer(file, MakeSignature()); + writer.Put(MakeKey(), 1, MakeSm80Config()); + ASSERT_TRUE(writer.Flush().IsOK()); + std::ifstream input(file); + std::string contents((std::istreambuf_iterator(input)), std::istreambuf_iterator()); + input.close(); + auto columns_text = [](const std::optional& config) { + std::vector columns; + gc::AppendConfigColumns(columns, config); + std::string text; + for (const auto& column : columns) text += "\t" + column; + return text; + }; + const auto successful = columns_text(MakeSm80Config()); + const auto position = contents.find(successful); + ASSERT_NE(position, std::string::npos); + contents.replace(position, successful.size(), columns_text(std::nullopt)); + std::istringstream legacy(contents); + gc::MatMulNBitsTacticCache reader(file, MakeSignature()); + ASSERT_TRUE(reader.Load(legacy).IsOK()); + EXPECT_FALSE(reader.Get(MakeKey(), 1).has_value()); + CleanUp(file); +} + +#ifdef _WIN32 +TEST(GemmTacticCacheTest, UnreadableSnapshotIsNotReplaced) { + const auto file = UniqueTempPrefix("unreadable") + ".matmulnbits_fpa_intb.tsv"; + gc::MatMulNBitsTacticCache cache(file, MakeSignature()); + cache.Put(MakeKey(), 1, MakeSm80Config()); + ASSERT_TRUE(cache.Flush().IsOK()); + // Deny reads while still permitting replacement. Flush must fail at the read, not the rename. + HANDLE handle = CreateFileA(file.c_str(), GENERIC_WRITE, FILE_SHARE_WRITE | FILE_SHARE_DELETE, + nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + ASSERT_NE(handle, INVALID_HANDLE_VALUE); + cache.Put(MakeKey(), 64, MakeSm90Config()); + EXPECT_FALSE(cache.Load().IsOK()); + EXPECT_FALSE(cache.Flush().IsOK()); + CloseHandle(handle); + gc::MatMulNBitsTacticCache reader(file, MakeSignature()); + ASSERT_TRUE(reader.Load().IsOK()); + EXPECT_TRUE(reader.Get(MakeKey(), 1).has_value()); + EXPECT_FALSE(reader.Get(MakeKey(), 64).has_value()); + EXPECT_TRUE(cache.Flush().IsOK()); // Failed writes retain their dirty rows. + CleanUp(file); +} + +TEST(GemmTacticCacheTest, FailedReplacementRemovesTemporaryFile) { + const auto file = UniqueTempPrefix("replacefail") + ".matmulnbits_fpa_intb.tsv"; + gc::MatMulNBitsTacticCache cache(file, MakeSignature()); + cache.Put(MakeKey(), 1, MakeSm80Config()); + ASSERT_TRUE(cache.Flush().IsOK()); + HANDLE handle = CreateFileA(file.c_str(), GENERIC_READ, FILE_SHARE_READ, + nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + ASSERT_NE(handle, INVALID_HANDLE_VALUE); + cache.Put(MakeKey(), 64, MakeSm90Config()); + EXPECT_FALSE(cache.Flush().IsOK()); + CloseHandle(handle); + const auto path = std::filesystem::path(file); + for (const auto& entry : std::filesystem::directory_iterator(path.parent_path())) { + EXPECT_NE(entry.path().filename().string().find(path.filename().string() + ".tmp."), 0u); + } + gc::MatMulNBitsTacticCache reader(file, MakeSignature()); + ASSERT_TRUE(reader.Load().IsOK()); + EXPECT_FALSE(reader.Get(MakeKey(), 64).has_value()); + EXPECT_TRUE(cache.Flush().IsOK()); + CleanUp(file); +} +#endif + +TEST(GemmTacticCacheTest, SharedExactAndRoundedHitsReachSeparatePrefixes) { + using namespace onnxruntime::llm::kernels::weight_only; + class Profiler : public GemmPluginProfiler, GemmIdCore, GemmIdCoreHash> { + public: + std::shared_ptr cache; + int stages = 0; + + protected: + void runTactic(int, int, int, const CutlassGemmConfig&, char*, const cudaStream_t&) override {} + size_t computeTmpSize(size_t, size_t, size_t) override { return 0; } + std::vector getTactics(int, int, int) const override { return {}; } + void stagePersistentCache(const GemmIdCore&, const MProfileMap& map, bool) override { + ++stages; + if (cache) { + for (const auto& [m, config] : map) cache->Put(MakeKey(), m, config); + } + } + }; + const auto file_a = UniqueTempPrefix("shared_a") + ".matmulnbits_fpa_intb.tsv"; + const auto file_b = UniqueTempPrefix("shared_b") + ".matmulnbits_fpa_intb.tsv"; + for (bool persist_a : {false, true}) { + auto shared = std::make_shared(); + GemmIdCore id(3072, 4096, onnxruntime::llm::nvinfer::DataType::kHALF); + shared->createMProfileMap(id); + Profiler a, b; + a.setSelectionTactics(shared); + b.setSelectionTactics(shared); + if (persist_a) a.cache = std::make_shared(file_a, MakeSignature()); + b.cache = std::make_shared(file_b, MakeSignature()); + (*shared->getMProfileMap(id))[64] = MakeSm80Config(); + ASSERT_TRUE(a.getBestConfigOrProfile(64, id).has_value()); + ASSERT_TRUE(b.getBestConfigOrProfile(64, id).has_value()); + (*shared->getMProfileMap(id))[128] = MakeSm90Config(); + ASSERT_TRUE(a.getBestConfigOrProfile(128, id).has_value()); + ASSERT_TRUE(b.getBestConfig(65, id).has_value()); // Capture-time rounded hit. + ASSERT_TRUE(b.getBestConfigOrProfile(65, id).has_value()); + EXPECT_EQ(b.stages, 2); // Repeated inference performs no staging work. + ASSERT_TRUE(b.cache->Flush().IsOK()); + gc::MatMulNBitsTacticCache result(file_b, MakeSignature()); + ASSERT_TRUE(result.Load().IsOK()); + EXPECT_EQ(result.GetAll(MakeKey()).size(), 2u); + CleanUp(file_a); + CleanUp(file_b); + } +} + +#if !defined(ORT_UNIT_TEST_HAS_CUDA_PLUGIN_EP) +TEST(GemmTacticCacheTest, CachedValidationRespectsSessionLaunchLimit) { + using namespace onnxruntime::llm::kernels::weight_only; + class Profiler : public WeightOnlyGroupwiseQuantGemmPluginProfiler { + public: + std::vector validated; + + void load(const GemmIdCore& id, MProfileMap& map) { + loadPersistentCache(id, map, true); + } + + protected: + std::vector getTactics(int, int, int) const override { return {MakeSm80Config()}; } + bool validatePersistentTactic(int m, int, int, const Config&) override { + validated.push_back(m); + return true; + } + }; + const auto file = UniqueTempPrefix("launch_limit") + ".matmulnbits_fpa_intb.tsv"; + auto cache = std::make_shared(file, MakeSignature()); + const auto key = MakeKey(); + cache->Put(key, 64, MakeSm80Config()); + cache->Put(key, 8192, MakeSm80Config()); + ASSERT_TRUE(cache->Flush().IsOK()); + + Profiler profiler; + profiler.setQuant(4, false, true); + profiler.setGroupSize(64); + profiler.setCudaKernelType(KernelType::FP16Int4Groupwise, 80); + profiler.setMaxProfileM(64); + profiler.setPersistentCache(cache); + Profiler::MProfileMap map; + profiler.load(GemmIdCore(key.n_16b, key.k, onnxruntime::llm::nvinfer::DataType::kHALF), map); + EXPECT_EQ(profiler.validated, std::vector({64})); + EXPECT_EQ(map.size(), 1u); + EXPECT_EQ(profiler.getMaxProfileM(), 64); + EXPECT_TRUE(cache->Get(key, 8192).has_value()); + CleanUp(file); +} +#endif + +TEST(GemmTacticCacheCudaTest, TemporaryRunnerFailureDoesNotPoisonNextProfiler) { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA device required for profiler events"; + } + using namespace onnxruntime::llm::kernels::weight_only; + class Profiler : public GemmPluginProfiler, GemmIdCore, GemmIdCoreHash> { + public: + std::shared_ptr cache; + bool fail = false; + int launches = 0; + + protected: + void runTactic(int, int, int, const CutlassGemmConfig&, char*, const cudaStream_t&) override { + ++launches; + if (fail) throw std::runtime_error("temporary runner failure"); + } + size_t computeTmpSize(size_t, size_t, size_t) override { return 1; } + std::vector getTactics(int, int, int) const override { return {MakeSm80Config()}; } + void loadPersistentCache(const GemmIdCore&, MProfileMap& map, bool) override { + for (const auto& entry : cache->GetAll(MakeKey())) map.insert(entry); + } + void stagePersistentCache(const GemmIdCore&, const MProfileMap& map, bool) override { + for (const auto& [m, config] : map) cache->Put(MakeKey(), m, config); + } + }; + const auto file = UniqueTempPrefix("runner_failure") + ".matmulnbits_fpa_intb.tsv"; + const auto dtype = onnxruntime::llm::nvinfer::DataType::kHALF; + const GemmIdCore id(3072, 4096, dtype); + for (bool fail : {true, false}) { + Profiler profiler; + profiler.fail = fail; + profiler.cache = std::make_shared(file, MakeSignature()); + ASSERT_TRUE(profiler.cache->Load().IsOK()); + profiler.setAllocator(std::make_shared()); // Fake runner never dereferences workspace. + profiler.profileTactics(std::make_shared(0), dtype, GemmDims(1, 1, id.n, id.k), id); + EXPECT_GT(profiler.launches, 0); + EXPECT_EQ(profiler.getBestConfig(1, id).has_value(), !fail); + ASSERT_TRUE(profiler.cache->Flush().IsOK()); + } + gc::MatMulNBitsTacticCache result(file, MakeSignature()); + ASSERT_TRUE(result.Load().IsOK()); + EXPECT_TRUE(result.Get(MakeKey(), 1).has_value()); + CleanUp(file); +} + +TEST(GemmTacticCacheTest, ParseCacheAccessAndFlag) { + auto both = gc::ParseCacheAccess(""); + ASSERT_TRUE(both.has_value()); + EXPECT_TRUE(both->load && both->save); + auto load_save = gc::ParseCacheAccess("LOAD_SAVE"); + ASSERT_TRUE(load_save.has_value()); + EXPECT_TRUE(load_save->load && load_save->save); + auto load = gc::ParseCacheAccess("load"); + ASSERT_TRUE(load.has_value()); + EXPECT_TRUE(load->load && !load->save); + auto save = gc::ParseCacheAccess("save"); + ASSERT_TRUE(save.has_value()); + EXPECT_TRUE(!save->load && save->save); + EXPECT_FALSE(gc::ParseCacheAccess("readonly").has_value()); + + EXPECT_EQ(gc::ParseCacheFlag(""), std::optional(false)); + EXPECT_EQ(gc::ParseCacheFlag("0"), std::optional(false)); + EXPECT_EQ(gc::ParseCacheFlag("On"), std::optional(true)); + EXPECT_EQ(gc::ParseCacheFlag("1"), std::optional(true)); + EXPECT_FALSE(gc::ParseCacheFlag("2").has_value()); +} + +TEST(GemmTacticCacheTest, ResolveFilePathSidecar) { + const auto sig = MakeSignature("NVIDIA H200", 90); + const std::string suffix = ".matmulnbits_fpa_intb.tsv"; + const std::filesystem::path abs_model_path = + std::filesystem::temp_directory_path() / "models" / "m.onnx"; + ASSERT_TRUE(abs_model_path.is_absolute()); + const auto abs_model_u8 = abs_model_path.u8string(); + const std::string abs_model(abs_model_u8.begin(), abs_model_u8.end()); + EXPECT_EQ(gc::MatMulNBitsTacticCache::ResolveFilePath("", "", sig, abs_model), abs_model + suffix); + const auto rel_resolved = gc::MatMulNBitsTacticCache::ResolveFilePath("", "", sig, "models/m.onnx"); + EXPECT_TRUE(std::filesystem::path(rel_resolved).is_absolute()); + // Explicit session locations win over the sidecar. + EXPECT_EQ(gc::MatMulNBitsTacticCache::ResolveFilePath("", "/tmp/p", sig, abs_model), "/tmp/p" + suffix); + EXPECT_EQ(gc::MatMulNBitsTacticCache::ResolveFilePath("/tmp/dir", "", sig, abs_model), + "/tmp/dir/NVIDIA_H200_sm90" + suffix); +} + +TEST(GemmTacticCacheTest, ReloadIfChangedPicksUpOtherWriters) { + const std::string prefix = UniqueTempPrefix("reload"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::HardwareSignature sig = MakeSignature(); + const gc::MatMulNBitsKey key = MakeKey(); + + gc::MatMulNBitsTacticCache reader(file, sig); + ASSERT_TRUE(reader.Load().IsOK()); // file does not exist yet + + { + gc::MatMulNBitsTacticCache writer(file, sig); + writer.Put(key, 1, MakeSm80Config()); + ASSERT_TRUE(writer.Flush().IsOK()); + } + ASSERT_TRUE(reader.ReloadIfChanged().IsOK()); + ASSERT_TRUE(reader.Get(key, 1).has_value() && reader.Get(key, 1)->has_value()); + ExpectConfigEqual(**reader.Get(key, 1), MakeSm80Config()); + + { + gc::MatMulNBitsTacticCache writer(file, sig); + writer.Put(key, 64, MakeSm90Config()); + ASSERT_TRUE(writer.Flush().IsOK()); + } + ASSERT_TRUE(reader.ReloadIfChanged().IsOK()); + EXPECT_TRUE(reader.Get(key, 64).has_value()); + + // When another writer replaces an existing clean bucket on disk, ReloadIfChanged() must refresh + // the clean entry while preserving any locally dirty bucket in memory. + reader.Put(key, 64, MakeSm80Config()); + { + gc::MatMulNBitsTacticCache writer(file, sig); + writer.Put(key, 1, MakeSm90Config()); + ASSERT_TRUE(writer.Flush().IsOK()); + } + ASSERT_TRUE(reader.ReloadIfChanged().IsOK()); + ASSERT_TRUE(reader.Get(key, 1).has_value() && reader.Get(key, 1)->has_value()); + ExpectConfigEqual(**reader.Get(key, 1), MakeSm90Config()); + ASSERT_TRUE(reader.Get(key, 64).has_value() && reader.Get(key, 64)->has_value()); + ExpectConfigEqual(**reader.Get(key, 64), MakeSm80Config()); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, FlushAdoptsRowsFromOtherWriters) { + const std::string prefix = UniqueTempPrefix("adopt"); + const std::string file = prefix + ".matmulnbits_fpa_intb.tsv"; + const gc::HardwareSignature sig = MakeSignature(); + const gc::MatMulNBitsKey key = MakeKey(); + + gc::MatMulNBitsTacticCache self(file, sig); + ASSERT_TRUE(self.Load().IsOK()); + { + gc::MatMulNBitsTacticCache other(file, sig); + other.Put(key, 1, MakeSm80Config()); + ASSERT_TRUE(other.Flush().IsOK()); + } + self.Put(key, 64, MakeSm90Config()); + ASSERT_TRUE(self.Flush().IsOK()); + + // The other writer's row is now known in memory, and no reload is needed to see it. + EXPECT_TRUE(self.Get(key, 1).has_value()); + ASSERT_TRUE(self.ReloadIfChanged().IsOK()); + EXPECT_TRUE(self.Get(key, 64).has_value()); + + CleanUp(file); +} + +#if !defined(ORT_UNIT_TEST_HAS_CUDA_PLUGIN_EP) +TEST(GemmTacticCacheTest, SessionProfileMapsDeduplicateOnlyWithinOneLiveSession) { + using onnxruntime::llm::kernels::weight_only::GemmPluginProfilerManager; + GemmPluginProfilerManager manager; + const int first_session = 1; + const int second_session = 2; + auto first_map = manager.getSessionProfileMap(&first_session); + EXPECT_EQ(first_map, manager.getSessionProfileMap(&first_session)); + EXPECT_NE(first_map, manager.getSessionProfileMap(&second_session)); + using onnxruntime::llm::kernels::weight_only::GemmIdCore; + const GemmIdCore key(1024, 4096, onnxruntime::llm::nvinfer::DataType::kHALF); + first_map->createMProfileMap(key); + first_map->getMProfileMap(key)->emplace(1, MakeSm80Config()); + auto second_map = manager.getSessionProfileMap(&second_session); + EXPECT_FALSE(second_map->existsMProfileMap(key)); + first_map.reset(); + EXPECT_FALSE(manager.getSessionProfileMap(&first_session)->existsMProfileMap(key)); +} + +TEST(GemmTacticCacheTest, RetuneProfilesAgainAcrossConcurrentAndSequentialSessions) { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA events are required for the profiling-loop smoke test."; + } + using onnxruntime::llm::kernels::weight_only::GemmDims; + using onnxruntime::llm::kernels::weight_only::GemmIdCore; + using onnxruntime::llm::kernels::weight_only::GemmPluginProfilerManager; + GemmPluginProfilerManager manager; + const int first_session = 1; + const int retune_session = 2; + const auto dtype = onnxruntime::llm::nvinfer::DataType::kHALF; + const GemmIdCore key(1024, 4096, dtype); + const GemmDims dims(1, 1, 1024, 4096); + const auto allocator = std::make_shared(); + const auto access = gc::ParseCacheAccess("save"); + ASSERT_TRUE(access.has_value()); + auto first = manager.createGemmPluginProfiler(false); + first->setPersistentCache(nullptr, *access, manager.getSessionProfileMap(&first_session)); + first->setAllocator(allocator); + first->profileTactics(nullptr, dtype, dims, key); + EXPECT_GT(first->launches, 0u); + + auto same_session = manager.createGemmPluginProfiler(false); + same_session->setPersistentCache(nullptr, *access, manager.getSessionProfileMap(&first_session)); + same_session->setAllocator(allocator); + same_session->profileTactics(nullptr, dtype, dims, key); + EXPECT_EQ(same_session->launches, 0u); + + auto retune = manager.createGemmPluginProfiler(false); + retune->setPersistentCache(nullptr, *access, manager.getSessionProfileMap(&retune_session)); + retune->setAllocator(allocator); + retune->profileTactics(nullptr, dtype, dims, key); + EXPECT_GT(retune->launches, 0u); + EXPECT_TRUE(first->getBestConfig(1, key).has_value()); + + first.reset(); + same_session.reset(); + auto sequential = manager.createGemmPluginProfiler(false); + sequential->setPersistentCache(nullptr, *access, manager.getSessionProfileMap(&first_session)); + sequential->setAllocator(allocator); + sequential->profileTactics(nullptr, dtype, dims, key); + EXPECT_GT(sequential->launches, 0u); + EXPECT_TRUE(retune->getBestConfig(1, key).has_value()); +} + +TEST(GemmTacticCacheTest, NormalRetuneNormalKeepsRetunedTacticWithoutChangingLiveSession) { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA events are required for the profiling-loop test."; + } + using namespace onnxruntime::llm::kernels::weight_only; + class ConfigurableProfiler : public CacheTestProfiler { + public: + int preferred_sm = 80; + + protected: + std::vector getTactics(int, int, int) const override { + return {MakeSm80Config(), MakeSm90Config()}; + } + float getSelectionTime(int, int, int, const CutlassGemmConfig& tactic, float) const override { + return tactic.sm_version == preferred_sm ? 0.1f : 10.0f; + } + bool validatePersistentTactic(int, int, int, const CutlassGemmConfig&) override { return true; } + }; + + const std::string file = UniqueTempPrefix("normal_retune_normal") + ".matmulnbits_fpa_intb.tsv"; + const auto sig = MakeSignature(); + auto cache = std::make_shared(file, sig); + ASSERT_TRUE(cache->Load().IsOK()); + + GemmPluginProfilerManager manager; + const int session_1 = 1; + const int session_2_retune = 2; + const int session_3 = 3; + const auto dtype = onnxruntime::llm::nvinfer::DataType::kHALF; + const GemmIdCore id(3072, 4096, dtype); + const GemmDims dims(1, 1, id.n, id.k); + const auto allocator = std::make_shared(); + + const auto load_save_access = gc::ParseCacheAccess("load_save"); + const auto save_access = gc::ParseCacheAccess("save"); + ASSERT_TRUE(load_save_access.has_value() && save_access.has_value()); + + auto configure_profiler = [&](ConfigurableProfiler& p, int preferred_sm, + const gc::CacheAccess& access, const void* session_key) { + p.preferred_sm = preferred_sm; + p.setQuant(4, false, true); + p.setGroupSize(64); + p.setCudaKernelType(KernelType::FP16Int4Groupwise, 80); + p.setProfileMOverride({1}); + p.setPersistentCache(cache, access, manager.getSessionProfileMap(session_key)); + p.setAllocator(allocator); + }; + + // 1. First normal session selects tactic A (MakeSm80Config) and deduplicates within the session. + auto first_normal = manager.createGemmPluginProfiler(false); + configure_profiler(*first_normal, 80, *load_save_access, &session_1); + first_normal->profileTactics(nullptr, dtype, dims, id, true); + EXPECT_GT(first_normal->launches, 0u); + ASSERT_TRUE(first_normal->getBestConfig(1, id).has_value()); + ExpectConfigEqual(*first_normal->getBestConfig(1, id), MakeSm80Config()); + + auto first_normal_dup = manager.createGemmPluginProfiler(false); + configure_profiler(*first_normal_dup, 80, *load_save_access, &session_1); + first_normal_dup->profileTactics(nullptr, dtype, dims, id, true); + EXPECT_EQ(first_normal_dup->launches, 0u); + ASSERT_TRUE(cache->Flush().IsOK()); + + // 2. A save-mode session retunes the same shape to tactic B (MakeSm90Config) and saves B. + auto retune = manager.createGemmPluginProfiler(false); + configure_profiler(*retune, 90, *save_access, &session_2_retune); + retune->profileTactics(nullptr, dtype, dims, id, true); + EXPECT_GT(retune->launches, 0u); + ASSERT_TRUE(retune->getBestConfig(1, id).has_value()); + ExpectConfigEqual(*retune->getBestConfig(1, id), MakeSm90Config()); + ASSERT_TRUE(cache->Flush().IsOK()); + + // The still-running first normal session must keep its original choice (A). + ASSERT_TRUE(first_normal->getBestConfig(1, id).has_value()); + ExpectConfigEqual(*first_normal->getBestConfig(1, id), MakeSm80Config()); + + // Lazily growing the first session's map with a new M bucket must publish only the newly + // selected bucket (64 -> A), preserving within-session sharing while keeping the retuned + // tactic (1 -> B) persisted. + const size_t launches_before_lazy = first_normal->launches; + const auto lazy_config = first_normal->getBestConfigOrProfile(64, id); + ASSERT_TRUE(lazy_config.has_value()); + ExpectConfigEqual(*lazy_config, MakeSm80Config()); + EXPECT_GT(first_normal->launches, launches_before_lazy); + + const auto shared_lazy_config = first_normal_dup->getBestConfigOrProfile(64, id); + ASSERT_TRUE(shared_lazy_config.has_value()); + ExpectConfigEqual(*shared_lazy_config, MakeSm80Config()); + EXPECT_EQ(first_normal_dup->launches, 0u); + ASSERT_TRUE(cache->Flush().IsOK()); + ASSERT_TRUE(cache->Get(MakeKey(), 1).has_value() && cache->Get(MakeKey(), 1)->has_value()); + ExpectConfigEqual(**cache->Get(MakeKey(), 1), MakeSm90Config()); + ASSERT_TRUE(cache->Get(MakeKey(), 64).has_value() && cache->Get(MakeKey(), 64)->has_value()); + ExpectConfigEqual(**cache->Get(MakeKey(), 64), MakeSm80Config()); + + // 3. The next normal session must load tactic B from the cache without re-profiling or overwriting B with A. + auto second_normal = manager.createGemmPluginProfiler(false); + configure_profiler(*second_normal, 80, *load_save_access, &session_3); + second_normal->profileTactics(nullptr, dtype, dims, id, true); + EXPECT_EQ(second_normal->launches, 0u); + ASSERT_TRUE(second_normal->getBestConfig(1, id).has_value()); + ExpectConfigEqual(*second_normal->getBestConfig(1, id), MakeSm90Config()); + ASSERT_TRUE(cache->Flush().IsOK()); + + gc::MatMulNBitsTacticCache verify(file, sig); + ASSERT_TRUE(verify.Load().IsOK()); + ASSERT_TRUE(verify.Get(MakeKey(), 1).has_value() && verify.Get(MakeKey(), 1)->has_value()); + ExpectConfigEqual(**verify.Get(MakeKey(), 1), MakeSm90Config()); + ASSERT_TRUE(verify.Get(MakeKey(), 64).has_value() && verify.Get(MakeKey(), 64)->has_value()); + ExpectConfigEqual(**verify.Get(MakeKey(), 64), MakeSm80Config()); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, ShutdownFlushAvoidsLockedMutexesAndSavesBeforeStop) { + const std::string file = UniqueTempPrefix("shutdown_flush") + ".matmulnbits_fpa_intb.tsv"; + const auto sig = MakeSignature(); + const auto key = MakeKey(); + + auto cache = onnxruntime::contrib::cuda::GetGlobalMatMulNBitsTacticCacheForTesting(file, sig, /*load*/ false); + ASSERT_NE(cache, nullptr); + cache->Put(key, 1, MakeSm80Config()); + + // Simulate a worker thread stopped while holding the cache mutex: shutdown flush must not block. + { + std::barrier locked(2); + std::barrier release(2); + std::thread stopped_worker([&]() { + auto lock = cache->LockMutexForTesting(); + locked.arrive_and_wait(); + release.arrive_and_wait(); + }); + locked.arrive_and_wait(); + EXPECT_FALSE(cache->TryFlush().IsOK()); + onnxruntime::contrib::cuda::FlushMatMulNBitsTacticCachesForShutdown(); + EXPECT_FALSE(std::filesystem::exists(file)); + release.arrive_and_wait(); + stopped_worker.join(); + } + + // Simulate a worker thread stopped while holding the global registry mutex: shutdown flush must not block. + { + std::barrier locked(2); + std::barrier release(2); + std::thread stopped_worker([&]() { + auto lock = onnxruntime::contrib::cuda::LockMatMulNBitsTacticCacheRegistryForTesting(); + locked.arrive_and_wait(); + release.arrive_and_wait(); + }); + locked.arrive_and_wait(); + onnxruntime::contrib::cuda::FlushMatMulNBitsTacticCachesForShutdown(); + EXPECT_FALSE(std::filesystem::exists(file)); + release.arrive_and_wait(); + stopped_worker.join(); + } + + // Saving before worker threads are stopped persists the staged tactics. + onnxruntime::contrib::cuda::FlushMatMulNBitsTacticCaches(); + gc::MatMulNBitsTacticCache verify(file, sig); + ASSERT_TRUE(verify.Load().IsOK()); + ASSERT_TRUE(verify.Get(key, 1).has_value() && verify.Get(key, 1)->has_value()); + ExpectConfigEqual(**verify.Get(key, 1), MakeSm80Config()); + + CleanUp(file); +} + +TEST(GemmTacticCacheTest, RelativeModelPathSavedNextToModelAfterDirectoryChange) { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA device required for CUDAExecutionProvider construction."; + } + + const auto base_dir = std::filesystem::temp_directory_path() / "ort_gemm_tactic_cache_rel_model_test"; + const auto models_dir = base_dir / "models"; + const auto other_dir = base_dir / "other_cwd"; + std::error_code ec; + std::filesystem::remove_all(base_dir, ec); + ASSERT_TRUE(std::filesystem::create_directories(models_dir, ec)) << ec.message(); + ASSERT_TRUE(std::filesystem::create_directories(other_dir, ec)) << ec.message(); + + const auto original_cwd = std::filesystem::current_path(ec); + ASSERT_FALSE(ec) << ec.message(); + struct CwdGuard { + std::filesystem::path saved; + std::filesystem::path cleanup_dir; + ~CwdGuard() { + std::error_code ignore; + std::filesystem::current_path(saved, ignore); + std::filesystem::remove_all(cleanup_dir, ignore); + } + } cwd_guard{original_cwd, base_dir}; + + std::filesystem::current_path(base_dir, ec); + ASSERT_FALSE(ec) << ec.message(); + + const auto sig = MakeSignature(); + const auto key = MakeKey(); + { + CUDAExecutionProvider ep(CUDAExecutionProviderInfo{}); + // Record a relative model path ("models/model.onnx") as GetCapability does at session creation. + ep.SetSessionModelPath(std::filesystem::path("models") / "model.onnx"); + EXPECT_TRUE(ep.GetSessionModelPath().is_absolute()); + + const auto model_u8 = ep.GetSessionModelPath().u8string(); + const std::string sidecar_prefix(model_u8.begin(), model_u8.end()); + const std::string cache_file = gc::MatMulNBitsTacticCache::ResolveFilePath("", "", sig, sidecar_prefix); + auto cache = onnxruntime::contrib::cuda::GetGlobalMatMulNBitsTacticCacheForTesting(cache_file, sig, /*load*/ false); + ASSERT_NE(cache, nullptr); + cache->Put(key, 1, MakeSm80Config()); + + // Change working directory before closing the session (~CUDAExecutionProvider flushes the cache). + std::filesystem::current_path(other_dir, ec); + ASSERT_FALSE(ec) << ec.message(); + } + + const auto expected_cache = models_dir / "model.onnx.matmulnbits_fpa_intb.tsv"; + const auto wrong_cache = other_dir / "models" / "model.onnx.matmulnbits_fpa_intb.tsv"; + EXPECT_TRUE(std::filesystem::exists(expected_cache)); + EXPECT_FALSE(std::filesystem::exists(wrong_cache)); +} +#endif // !defined(ORT_UNIT_TEST_HAS_CUDA_PLUGIN_EP) + +} // namespace test +} // namespace onnxruntime + +#endif // USE_FPA_INTB_GEMM diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc index 00cbba0e19f61..c47c029a58168 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc @@ -1,10 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include #include -#include #include #include @@ -12,6 +12,11 @@ #include "contrib_ops/cuda/moe/qmoe_kernels.h" #include "core/providers/cuda/cuda_common.h" +#include "core/platform/env_var_utils.h" + +#if defined(ENABLE_FP4) && defined(USE_FP4_QMOE) +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h" +#endif namespace onnxruntime { namespace test { @@ -163,259 +168,280 @@ TEST(CUDA_EP_Unittest, QMoEFp8DequantizesOnlySelectedExpertsAndGathersBias) { } } -template -void TestAllFp8Codes() { - std::array codes{}; - for (int i = 0; i < 256; ++i) { - codes[i] = static_cast(i); - } - const float scale = 1.0f; - CudaBuffer weights(codes.size()); - CudaBuffer scales(sizeof(scale)); - CudaBuffer output(codes.size() * sizeof(T)); - weights.Upload(codes.data()); - scales.Upload(&scale); - qmoe::LaunchQMoEDequantizeFp8Weights(weights.As(), scales.As(), output.As(), - 1, 1, 256, nullptr); - std::array actual{}; - output.Download(actual.data()); - for (int code = 0; code < 256; ++code) { - SCOPED_TRACE(code); - const float value = static_cast(actual[code]); - if ((code & 127) == 127) { - EXPECT_TRUE(std::isnan(value)); - continue; - } - const int exponent = (code >> 3) & 15; - const int mantissa = code & 7; - const float magnitude = exponent == 0 ? std::ldexp(static_cast(mantissa), -9) - : std::ldexp(1.0f + static_cast(mantissa) / 8.0f, exponent - 7); - const float expected = code & 128 ? -magnitude : magnitude; - EXPECT_EQ(value, expected); - EXPECT_EQ(std::signbit(value), std::signbit(expected)); - } +[[maybe_unused]] float FiniteFp8Reference(int code) { + const int exponent = (code >> 3) & 15; + const int mantissa = code & 7; + const float magnitude = exponent == 0 ? std::ldexp(static_cast(mantissa), -9) + : std::ldexp(1.0f + static_cast(mantissa) / 8.0f, exponent - 7); + return code & 128 ? -magnitude : magnitude; } -TEST(CUDA_EP_Unittest, QMoEFp8AllCodesPreserveFiniteValuesSignedZerosAndNaNs) { - TestAllFp8Codes(); - TestAllFp8Codes<__nv_bfloat16>(); +#if defined(ENABLE_FP4) && defined(USE_FP4_QMOE) +bool SupportsFp8Gemv() { + int device = 0; + cudaDeviceProp properties{}; + CUDA_CALL_THROW(cudaGetDevice(&device)); + CUDA_CALL_THROW(cudaGetDeviceProperties(&properties, device)); + return properties.major >= 8; } +#endif -template -void TestFp8ProjectionScaleBlocks(int fusion, int block_size, int k, bool graph, bool gemm) { - constexpr int n = 6; - constexpr int rows = 3; - constexpr int num_experts = 2; - const bool split = fusion == 1; - const int source_n = split ? n / 2 : n; - const int scale_n = (source_n + block_size - 1) / block_size; - const int scale_k = (k + block_size - 1) / block_size; - const std::array experts{0, 0, 1}; - const std::array row_map{5, 0, 4}; - const std::array expert_offsets{0, 2, 3}; - std::vector weights(num_experts * source_n * k); +template +void TestFp8GemvCodeDecoding(int scale_type, int block_size, int k, int fusion, bool graph) { + constexpr int n = 256; + const int source_n = fusion == 1 ? n / 2 : n; + const int scales_per_expert = + ((source_n + block_size - 1) / block_size) * ((k + block_size - 1) / block_size); + std::vector weights(2 * source_n * k); std::vector up_weights(weights.size()); - std::vector scales(num_experts * scale_n * scale_k); - std::vector up_scales(scales.size()); - std::vector input(rows * k); - std::vector output(rows * n); - for (size_t i = 0; i < weights.size(); ++i) { - weights[i] = static_cast(i % 3 == 0 ? 0xb8 : (i % 3 == 1 ? 0x38 : 0x40)); - up_weights[i] = 0x30; + for (int expert = 0; expert < 2; ++expert) { + for (int row = 0; row < source_n; ++row) { + const int code = fusion == 1 ? row * 2 : row; + std::fill_n(weights.begin() + (expert * source_n + row) * k, k, static_cast(code)); + std::fill_n(up_weights.begin() + (expert * source_n + row) * k, k, static_cast(code + 1)); + } } - for (size_t i = 0; i < scales.size(); ++i) { - scales[i] = ScaleT(static_cast(i % 5 + 1) / 32.0f); - up_scales[i] = ScaleT(static_cast(i % 5 + 2) / 32.0f); + const int k_blocks = (k + block_size - 1) / block_size; + std::vector scales(2 * scales_per_expert); + for (int expert = 0; expert < 2; ++expert) { + for (int block = 0; block < scales_per_expert; ++block) { + scales[expert * scales_per_expert + block] = + S((expert + 1) * 0.25f + ((block / k_blocks) % 4) * 0.125f + (block % k_blocks) * 0.25f); + } } - for (size_t i = 0; i < input.size(); ++i) { - input[i] = T(static_cast(static_cast(i % 7) - 3) / 16.0f); + std::vector input(2 * k); + for (int col = 0; col < k; ++col) { + input[col] = T(col == k - 1 ? 2.0f : (1 + col % 4) * 0.03125f); + input[k + col] = T(col == k - 1 ? -4.0f : -(1 + col % 3) * 0.0625f); } - CudaBuffer d_weights(weights.size()), d_up_weights(up_weights.size()); - CudaBuffer d_scales(scales.size() * sizeof(ScaleT)), d_up_scales(up_scales.size() * sizeof(ScaleT)); - CudaBuffer d_input(input.size() * sizeof(T)), d_output(output.size() * sizeof(T)); - CudaBuffer d_experts(sizeof(experts)), d_rows(sizeof(row_map)), d_offsets(sizeof(expert_offsets)); - CudaBuffer d_tiles((num_experts + 1) * sizeof(int)); - d_weights.Upload(weights.data()); - d_up_weights.Upload(up_weights.data()); - d_scales.Upload(scales.data()); - d_up_scales.Upload(up_scales.data()); - d_input.Upload(input.data()); - d_experts.Upload(experts.data()); - d_rows.Upload(row_map.data()); - d_offsets.Upload(expert_offsets.data()); + const std::array experts{1, 0}; + const std::array source_rows{1, 0}; + CudaBuffer device_weights(weights.size()); + CudaBuffer device_up_weights(up_weights.size()); + CudaBuffer device_scales(scales.size() * sizeof(S)); + CudaBuffer device_input(input.size() * sizeof(T)); + CudaBuffer device_experts(sizeof(experts)); + CudaBuffer device_source_rows(sizeof(source_rows)); + CudaBuffer device_output(2 * n * sizeof(T)); + device_weights.Upload(weights.data()); + device_up_weights.Upload(up_weights.data()); + device_scales.Upload(scales.data()); + device_input.Upload(input.data()); + device_experts.Upload(experts.data()); + device_source_rows.Upload(source_rows.data()); qmoe::QMoEFp8ProjectionParams params; - params.weights = d_weights.As(); - params.scales = d_scales.data; - params.scale_type = std::is_same_v ? 0 : (std::is_same_v ? 1 : 2); - params.up_scale_type = params.scale_type; - params.up_weights = split ? d_up_weights.As() : nullptr; - params.up_scales = split ? d_up_scales.data : nullptr; - params.experts = d_experts.As(); - params.row_to_unpermuted = d_rows.As(); - params.expert_offsets = d_offsets.As(); - params.num_experts = num_experts; - params.num_rows = rows; - params.expanded_rows = rows; + params.weights = device_weights.As(); + params.scales = device_scales.data; + params.scale_type = scale_type; + if (fusion == 1) { + params.up_weights = device_up_weights.As(); + params.up_scales = device_scales.data; + params.up_scale_type = scale_type; + } + params.fusion = fusion; + params.experts = device_experts.As(); + params.row_to_unpermuted = device_source_rows.As(); + params.num_experts = 2; + params.num_rows = 2; + params.expanded_rows = 2; params.n = n; params.k = k; params.block_size = block_size; - params.fusion = fusion; - if (gemm) { - qmoe::LaunchQMoEFp8ExpertTiles(params.expert_offsets, d_tiles.As(), num_experts, nullptr); - params.tile_offsets = d_tiles.As(); - CUDA_CALL_THROW(cudaDeviceSynchronize()); - } cudaStream_t stream; - CUDA_CALL_THROW(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); - if (graph) { - CUDA_CALL_THROW(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal)); - } - qmoe::LaunchQMoEFp8Projection(params, d_input.As(), d_output.As(), stream); + CUDA_CALL_THROW(cudaStreamCreate(&stream)); if (graph) { cudaGraph_t captured; cudaGraphExec_t executable; + CUDA_CALL_THROW(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + qmoe::LaunchQMoEFp8Projection(params, device_input.As(), device_output.As(), stream); CUDA_CALL_THROW(cudaStreamEndCapture(stream, &captured)); - CUDA_CALL_THROW(cudaGraphInstantiate(&executable, captured, nullptr, nullptr, 0)); - for (int replay = 0; replay < 3; ++replay) { - CUDA_CALL_THROW(cudaGraphLaunch(executable, stream)); + size_t node_count = 0; + CUDA_CALL_THROW(cudaGraphGetNodes(captured, nullptr, &node_count)); + std::vector nodes(node_count); + CUDA_CALL_THROW(cudaGraphGetNodes(captured, nodes.data(), &node_count)); + EXPECT_EQ(node_count, 1u); + if (node_count == 1) { + cudaKernelNodeParams node{}; + CUDA_CALL_THROW(cudaGraphKernelNodeGetParams(nodes[0], &node)); + EXPECT_EQ(qmoe::QMoEFp8GemvKernelUsesNativeDecode(node.func), + ParseEnvironmentVariableWithDefault("ORT_QMOE_FP8_GEMV_NATIVE_DECODE", false)); } + CUDA_CALL_THROW(cudaGraphInstantiate(&executable, captured, nullptr, nullptr, 0)); + CUDA_CALL_THROW(cudaGraphLaunch(executable, stream)); CUDA_CALL_THROW(cudaStreamSynchronize(stream)); CUDA_CALL_THROW(cudaGraphExecDestroy(executable)); CUDA_CALL_THROW(cudaGraphDestroy(captured)); + } else { + qmoe::LaunchQMoEFp8Projection(params, device_input.As(), device_output.As(), stream); + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); } - CUDA_CALL_THROW(cudaStreamSynchronize(stream)); CUDA_CALL_THROW(cudaStreamDestroy(stream)); - d_output.Download(output.data()); - for (int row = 0; row < rows; ++row) { + std::array actual{}; + device_output.Download(actual.data()); + for (int row = 0; row < 2; ++row) { for (int col = 0; col < n; ++col) { - const bool up = split && col % 2; - const int source_col = split ? col / 2 : (fusion == 2 ? col / 2 + (col % 2) * (n / 2) : col); - float expected = 0.0f; + const int code = fusion == 2 ? col / 2 + (col % 2) * (n / 2) : col; + SCOPED_TRACE(row); + SCOPED_TRACE(code); + const float value = static_cast(actual[row * n + col]); + if ((code & 127) == 127) { + EXPECT_TRUE(std::isnan(value)); + continue; + } + const int source_col = fusion == 1 ? col / 2 : code; + double sum = 0.0; for (int inner = 0; inner < k; ++inner) { - const int weight_index = (experts[row] * source_n + source_col) * k + inner; - const int scale_index = (experts[row] * scale_n + source_col / block_size) * scale_k + - inner / block_size; - const uint8_t code = up ? up_weights[weight_index] : weights[weight_index]; - const float value = code == 0xb8 ? -1.0f : (code == 0x38 ? 1.0f : (code == 0x40 ? 2.0f : 0.5f)); - const T weight(value * static_cast(up ? up_scales[scale_index] : scales[scale_index])); - expected += static_cast(input[(row_map[row] % rows) * k + inner]) * static_cast(weight); + const float scale = static_cast(scales[experts[row] * scales_per_expert + + (source_col / block_size) * k_blocks + inner / block_size]); + const T scaled_weight(FiniteFp8Reference(code) * scale); + sum += static_cast(scaled_weight) * static_cast(input[source_rows[row] * k + inner]); } - ASSERT_EQ(static_cast(output[row * n + col]), static_cast(T(expected))) - << "row=" << row << " col=" << col; + const float expected = static_cast(T(static_cast(sum))); + const float tolerance = std::is_same_v ? 0.001f : 0.008f; + EXPECT_NEAR(value, expected, std::max(1e-7f, std::abs(expected) * tolerance)); } } } -TEST(CUDA_EP_Unittest, QMoEFp8ProjectionPreservesScaleBlocksLayoutsAndGraphReplay) { - int device = 0; - cudaDeviceProp properties{}; - ASSERT_EQ(cudaGetDevice(&device), cudaSuccess); - ASSERT_EQ(cudaGetDeviceProperties(&properties, device), cudaSuccess); - for (int fusion : {0, 1, 2}) { - for (int block_size : {64, 128}) { - for (int k : {127, 128, 257}) { - for (bool graph : {false, true}) { - for (bool gemm : {false, true}) { - if (gemm && properties.major < 8) { - continue; - } - SCOPED_TRACE(::testing::Message() << "fusion=" << fusion << " block=" << block_size - << " k=" << k << " graph=" << graph << " gemm=" << gemm); - TestFp8ProjectionScaleBlocks(fusion, block_size, k, graph, gemm); - TestFp8ProjectionScaleBlocks(fusion, block_size, k, graph, gemm); - TestFp8ProjectionScaleBlocks(fusion, block_size, k, graph, gemm); - TestFp8ProjectionScaleBlocks<__nv_bfloat16, float>(fusion, block_size, k, graph, gemm); - TestFp8ProjectionScaleBlocks<__nv_bfloat16, half>(fusion, block_size, k, graph, gemm); - TestFp8ProjectionScaleBlocks<__nv_bfloat16, __nv_bfloat16>(fusion, block_size, k, graph, gemm); - } - } +void TestFp8GemvCodeConfigurations() { + for (int block_size : {16, 128}) { + for (int fusion : {0, 1, 2}) { + for (bool graph : {false, true}) { + SCOPED_TRACE(block_size); + SCOPED_TRACE(fusion); + SCOPED_TRACE(graph); + const int k = block_size * 2 + 1; + TestFp8GemvCodeDecoding(0, block_size, k, fusion, graph); + TestFp8GemvCodeDecoding(1, block_size, k, fusion, graph); + TestFp8GemvCodeDecoding(2, block_size, k, fusion, graph); + TestFp8GemvCodeDecoding<__nv_bfloat16, float>(0, block_size, k, fusion, graph); + TestFp8GemvCodeDecoding<__nv_bfloat16, half>(1, block_size, k, fusion, graph); + TestFp8GemvCodeDecoding<__nv_bfloat16, __nv_bfloat16>(2, block_size, k, fusion, graph); } } } } +TEST(CUDA_EP_Unittest, QMoEFp8GemvAllCodesWithMappedRowsPartialKBlocksAndGraphReplay) { + TestFp8GemvCodeConfigurations(); +} + +TEST(CUDA_EP_Unittest, QMoEFp8GemvNativeDecode) { + if (!ParseEnvironmentVariableWithDefault("ORT_QMOE_FP8_GEMV_NATIVE_DECODE", false)) { + GTEST_SKIP() << "Native decoding coverage runs in its flag-enabled CTest process."; + } + if (!qmoe::QMoEFp8GemvSupportsNativeDecode()) { + GTEST_SKIP() << "Native E4M3 decoding requires FP8-enabled kernels compiled for SM90+."; + } + TestFp8GemvCodeConfigurations(); +} + +#if defined(ENABLE_FP4) && defined(USE_FP4_QMOE) template -void TestNvfp4RowMajorDequantization(int k, bool compact, bool with_bias) { - constexpr int source_experts = 5; - constexpr int n = 3; - const std::array selected{4, 1, -1}; - const int capacity = compact ? static_cast(selected.size()) : source_experts; - constexpr std::array magnitudes{0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; - constexpr std::array scale_codes{0x00, 0x01, 0x20, 0x38, 0x41, 0x7e, 0x80, 0xb8}; - constexpr std::array scale_values{0.0f, 1.0f / 512.0f, 0.125f, 1.0f, 2.25f, 448.0f, -0.0f, -1.0f}; - std::vector weights(source_experts * n * k / 2); - std::vector scales(source_experts * n * k / 16); - std::vector globals(source_experts); - std::vector bias(source_experts * n); - for (int expert = 0; expert < source_experts; ++expert) { - globals[expert] = 0.37f * (expert + 1); - for (int row = 0; row < n; ++row) { - const int source_row = expert * n + row; - bias[source_row] = T(static_cast(source_row)); - for (int col = 0; col < k; col += 2) { - const int low = (col + row + expert) % 16; - const int high = (low + 1) % 16; - weights[source_row * (k / 2) + col / 2] = static_cast(low | (high << 4)); - } - for (int group = 0; group < k / 16; ++group) { - scales[source_row * (k / 16) + group] = scale_codes[(group + source_row) % scale_codes.size()]; - } +void TestNvfp4GemvScaleCodes(int k) { + constexpr int n = 256; + std::vector weights(n * k / 2, 0); + std::vector scales(n * k / 16); + const int blocks = k / 16; + for (int row = 0; row < n; ++row) { + for (int block = 0; block < blocks; ++block) { + weights[row * k / 2 + block * 8] = 0x22; + scales[row * blocks + block] = + block == 0 ? static_cast(row) : static_cast(0x30 + (row + block) % 16); } } - std::vector output(capacity * n * k, T(-123.0f)); - std::vector output_bias(capacity * n, T(-123.0f)); + std::vector input(k, T(0.0f)); + for (int block = 0; block < blocks; ++block) { + input[block * 16] = T(1.0f + (block % 3) * 0.5f); + input[block * 16 + 1] = T(0.5f + (block % 2) * 0.25f); + } + const float global = 0.37f; + const int expert = 0; CudaBuffer device_weights(weights.size()); CudaBuffer device_scales(scales.size()); - CudaBuffer device_globals(globals.size() * sizeof(float)); - CudaBuffer device_bias(bias.size() * sizeof(T)); - CudaBuffer device_output(output.size() * sizeof(T)); - CudaBuffer device_output_bias(output_bias.size() * sizeof(T)); - CudaBuffer device_selected(selected.size() * sizeof(int)); + CudaBuffer device_input(input.size() * sizeof(T)); + CudaBuffer device_global(sizeof(global)); + CudaBuffer device_expert(sizeof(expert)); + CudaBuffer device_output(n * sizeof(T)); device_weights.Upload(weights.data()); device_scales.Upload(scales.data()); - device_globals.Upload(globals.data()); - device_bias.Upload(bias.data()); - device_output.Upload(output.data()); - device_output_bias.Upload(output_bias.data()); - device_selected.Upload(selected.data()); - qmoe::LaunchQMoEDequantizeNvfp4Weights( - device_weights.As(), device_scales.As(), device_globals.As(), - device_output.As(), capacity, n, k, nullptr, - compact ? device_selected.As() : nullptr, - with_bias ? device_bias.As() : nullptr, - with_bias ? device_output_bias.As() : nullptr, true); - device_output.Download(output.data()); - device_output_bias.Download(output_bias.data()); - for (int slot = 0; slot < capacity; ++slot) { - const int expert = compact ? selected[slot] : slot; - for (int row = 0; row < n; ++row) { - const T expected_bias = with_bias && expert >= 0 ? bias[expert * n + row] : T(-123.0f); - EXPECT_EQ(std::memcmp(&output_bias[slot * n + row], &expected_bias, sizeof(T)), 0); - for (int col = 0; col < k; ++col) { - T expected(-123.0f); - if (expert >= 0) { - const int code = (col + row + expert) % 16; - const float value = code < 8 ? magnitudes[code] : -magnitudes[code - 8]; - expected = T(value * scale_values[(col / 16 + expert * n + row) % scale_values.size()] * globals[expert]); - } - ASSERT_EQ(std::memcmp(&output[(slot * n + row) * k + col], &expected, sizeof(T)), 0) - << "slot=" << slot << " row=" << row << " col=" << col; + device_input.Upload(input.data()); + device_global.Upload(&global); + device_expert.Upload(&expert); + using namespace onnxruntime::llm::kernels::moe_gemv; + cudaStream_t stream; + cudaGraph_t graph; + cudaGraphExec_t executable; + CUDA_CALL_THROW(cudaStreamCreate(&stream)); + CUDA_CALL_THROW(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + launch_moe_gemv_fp4_symmetric( + device_input.As(), device_weights.As(), nullptr, device_scales.As(), + device_global.As(), nullptr, device_output.As(), nullptr, device_expert.As(), + 1, 1, n, k, 16, 90, MoeGemvConfig::kDefault, false, true, stream, true); + CUDA_CALL_THROW(cudaStreamEndCapture(stream, &graph)); + size_t node_count = 0; + CUDA_CALL_THROW(cudaGraphGetNodes(graph, nullptr, &node_count)); + std::vector nodes(node_count); + CUDA_CALL_THROW(cudaGraphGetNodes(graph, nodes.data(), &node_count)); + EXPECT_EQ(node_count, 1u); + if (node_count == 1) { + cudaKernelNodeParams node{}; + CUDA_CALL_THROW(cudaGraphKernelNodeGetParams(nodes[0], &node)); + EXPECT_EQ(Fp4MoeGemvKernelUsesNativeScale(node.func), + ParseEnvironmentVariableWithDefault("ORT_NVFP4_GEMV_NATIVE_SCALE", false)); + } + CUDA_CALL_THROW(cudaGraphInstantiate(&executable, graph, nullptr, nullptr, 0)); + CUDA_CALL_THROW(cudaGraphLaunch(executable, stream)); + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + CUDA_CALL_THROW(cudaGraphExecDestroy(executable)); + CUDA_CALL_THROW(cudaGraphDestroy(graph)); + CUDA_CALL_THROW(cudaStreamDestroy(stream)); + std::array actual{}; + device_output.Download(actual.data()); + for (int code = 0; code < 256; ++code) { + SCOPED_TRACE(code); + const float value = static_cast(actual[code]); + if ((code & 127) == 127) { + EXPECT_TRUE(std::isnan(value)); + } else { + double sum = 0.0; + for (int block = 0; block < blocks; ++block) { + const float scale = static_cast(T(FiniteFp8Reference(scales[code * blocks + block]) * global)); + sum += scale * static_cast(input[block * 16]); + sum += scale * static_cast(input[block * 16 + 1]); } + const float expected = static_cast(T(static_cast(sum))); + const float tolerance = std::is_same_v ? 0.001f : 0.008f; + EXPECT_NEAR(value, expected, std::max(1e-7f, std::abs(expected) * tolerance)); } } } -TEST(CUDA_EP_Unittest, QMoENvfp4RowMajorVectorDequantizationPreservesValuesAndBias) { - for (int k : {16, 48, 1280, 2560}) { - for (bool compact : {false, true}) { - for (bool with_bias : {false, true}) { - SCOPED_TRACE(::testing::Message() << "k=" << k << " compact=" << compact << " bias=" << with_bias); - TestNvfp4RowMajorDequantization(k, compact, with_bias); - TestNvfp4RowMajorDequantization<__nv_bfloat16>(k, compact, with_bias); - } - } +void TestNvfp4GemvScaleCodeConfigurations() { + for (int k : {16, 80, 1040}) { + SCOPED_TRACE(k); + TestNvfp4GemvScaleCodes(k); + TestNvfp4GemvScaleCodes<__nv_bfloat16>(k); + } +} + +TEST(CUDA_EP_Unittest, QMoENvfp4GemvAllScaleCodesPreserveScaledValuesAndNaNs) { + if (!SupportsFp8Gemv()) { + GTEST_SKIP() << "NVFP4 GEMV requires SM80+."; + } + TestNvfp4GemvScaleCodeConfigurations(); +} + +TEST(CUDA_EP_Unittest, QMoENvfp4GemvNativeScale) { + if (!ParseEnvironmentVariableWithDefault("ORT_NVFP4_GEMV_NATIVE_SCALE", false)) { + GTEST_SKIP() << "Native-scale coverage runs in its flag-enabled CTest process."; + } + if (!SupportsFp8Gemv() || !onnxruntime::llm::kernels::moe_gemv::Fp4MoeGemvSupportsNativeScale()) { + GTEST_SKIP() << "Native E4M3 scale conversion requires FP8-enabled kernels compiled for SM90+."; } + TestNvfp4GemvScaleCodeConfigurations(); } +#endif } // namespace } // namespace test diff --git a/onnxruntime/test/contrib_ops/engram_ops_test.cc b/onnxruntime/test/contrib_ops/engram_ops_test.cc index d45282428e11f..f1fc4660e9485 100644 --- a/onnxruntime/test/contrib_ops/engram_ops_test.cc +++ b/onnxruntime/test/contrib_ops/engram_ops_test.cc @@ -204,9 +204,9 @@ void RunEngramGateNormedTest(float tolerance, bool packed = false) { OpTester test("EngramGate", 1, kMSDomain); test.AddAttribute("epsilon", kEpsilon); const std::vector key_shape = packed ? std::vector{1, hc_mult, hidden_size} - : std::vector{1, 1, hc_mult, hidden_size}; + : std::vector{1, 1, hc_mult, hidden_size}; const std::vector value_shape = packed ? std::vector{1, hidden_size} - : std::vector{1, 1, hidden_size}; + : std::vector{1, 1, hidden_size}; test.AddInput("key", key_shape, ToTensorType(key)); test.AddInput("query", key_shape, ToTensorType(query)); test.AddInput("value", value_shape, ToTensorType(value)); @@ -273,9 +273,9 @@ void RunEngramGateVectorizedTest(float tolerance, int64_t hidden, bool packed = OpTester test("EngramGate", 1, kMSDomain); test.AddAttribute("epsilon", kEpsilon); const std::vector key_shape = packed ? std::vector{kBatch * kSequence, kHcMult, hidden} - : std::vector{kBatch, kSequence, kHcMult, hidden}; + : std::vector{kBatch, kSequence, kHcMult, hidden}; const std::vector value_shape = packed ? std::vector{kBatch * kSequence, hidden} - : std::vector{kBatch, kSequence, hidden}; + : std::vector{kBatch, kSequence, hidden}; test.AddInput("key", key_shape, ToTensorType(key)); test.AddInput("query", key_shape, ToTensorType(query)); test.AddInput("value", value_shape, ToTensorType(value)); diff --git a/onnxruntime/test/contrib_ops/fused_conv_test.cc b/onnxruntime/test/contrib_ops/fused_conv_test.cc index aa13e2f6d56f0..4e885a3770a24 100644 --- a/onnxruntime/test/contrib_ops/fused_conv_test.cc +++ b/onnxruntime/test/contrib_ops/fused_conv_test.cc @@ -670,6 +670,41 @@ TEST(FusedConvTest, Cpu_Conv3D_Batched_Relu) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +TEST(FusedConvTest, Cpu_ZeroChannel_SumShapeValidation) { + auto run_zero_channel_with_sum = [](const vector& sum_shape, + const vector& sum_data, + OpTester::ExpectResult expect_result, + const string& expected_error) { + OpTester test("FusedConv", 1, onnxruntime::kMSDomain); + test.AddAttribute("group", static_cast(1)); + test.AddAttribute("kernel_shape", vector{1, 1}); + test.AddAttribute("pads", vector{0, 0, 0, 0}); + test.AddAttribute("strides", vector{1, 1}); + test.AddAttribute("dilations", vector{1, 1}); + test.AddAttribute("activation", string("Relu")); + + test.AddInput("X", {1, 0, 2, 2}, {}); + test.AddInput("W", {1, 0, 1, 1}, {}); + test.AddInput("B", {1}, {1.0f}); + test.AddInput("Z", sum_shape, sum_data); + test.AddOutput("Y", {1, 1, 2, 2}, {3.0f, 4.0f, 5.0f, 6.0f}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(expect_result, expected_error, {}, nullptr, &execution_providers); + }; + + // Matching residual shape succeeds. + run_zero_channel_with_sum({1, 1, 2, 2}, {2.0f, 3.0f, 4.0f, 5.0f}, + OpTester::ExpectResult::kExpectSuccess, ""); + // Smaller residual shape must be rejected before reading Sum data. + run_zero_channel_with_sum({1, 1, 1, 1}, {2.0f}, + OpTester::ExpectResult::kExpectFailure, "output and sum shape must match"); + // Empty residual shape must also be rejected. + run_zero_channel_with_sum({1, 0, 2, 2}, {}, + OpTester::ExpectResult::kExpectFailure, "output and sum shape must match"); +} + #endif } // namespace test diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index 4db622e45ad2b..541a4ecf6502e 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -348,11 +348,11 @@ void RunTypedCase(const Geometry& g, const Options& o, const Inputs& in_raw, flo if (use_webgpu) { #ifdef USE_WEBGPU if (test_max_storage_buffer_binding_size != 0) { + WebGpuExecutionProviderTestOptions test_options{}; + test_options.max_storage_buffer_binding_size = test_max_storage_buffer_binding_size; webgpu_ep = webgpu_config != nullptr - ? WebGpuExecutionProviderWithTestStorageBufferBindingSize( - *webgpu_config, test_max_storage_buffer_binding_size) - : WebGpuExecutionProviderWithTestStorageBufferBindingSize( - test_max_storage_buffer_binding_size); + ? WebGpuExecutionProviderWithTestOptions(*webgpu_config, test_options) + : WebGpuExecutionProviderWithTestOptions(test_options); } else { webgpu_ep = webgpu_config != nullptr ? WebGpuExecutionProviderWithOptions(*webgpu_config) : DefaultWebGpuExecutionProvider(); @@ -932,6 +932,140 @@ TEST(GatedDeltaNetWebGpuTest, CompactStateUpdatesFp16QwenGeometry) { /*rank4=*/false, /*fetches=*/nullptr, /*use_webgpu=*/true); } +#ifdef USE_WEBGPU +TEST(GatedDeltaNetWebGpuTest, CapturedStateReplayMatchesIndependentPrefixes) { + auto replay_ep = DefaultWebGpuExecutionProvider(); + if (replay_ep == nullptr) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + Geometry geometry{5, 1, 2, 6, 16, 5}; + Inputs inputs = MakeInputs(geometry, 233); + for (auto& value : inputs.state0) { + value *= 8.0f; + } + inputs.capture_count = {geometry.total_tokens}; + Options options; + options.gate_activation = "qwen"; + options.beta_activation = "sigmoid"; + options.qk_l2_norm = 1; + // Capture capacity intentionally exceeds both the full run length and every partial prefix. + options.state_update_capacity = 7; + std::vector producer_fetches; + ASSERT_NO_FATAL_FAILURE(RunTypedCase(geometry, options, inputs, 4e-4f, 4e-4f, + /*rank4=*/false, &producer_fetches, /*use_webgpu=*/true)); + ASSERT_EQ(producer_fetches.size(), 3u); + const auto& captured = producer_fetches[2].Get(); + const int64_t capacity = options.state_update_capacity; + const int64_t decay_elements = capacity * geometry.hv; + const int64_t key_elements = capacity * geometry.hq * geometry.dk; + const int64_t capsule_elements = decay_elements + key_elements + capacity * geometry.hv * geometry.dv; + ASSERT_EQ(captured.Shape(), TensorShape({1, capsule_elements})); + const int64_t state_elements = static_cast(inputs.state0.size()); + constexpr int64_t kDestinationOffset = 5; + std::vector destination(static_cast(state_elements + 11), -931.25f); + std::array descriptor{0, kDestinationOffset, 0, decay_elements, decay_elements + key_elements, + geometry.hv, geometry.dv, geometry.dk, geometry.hq, capacity, 1}; + + Model model("gated_delta_net_captured_replay", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kMSDomain, 1}}, {}, + DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + auto tensor_arg = [&](const char* name, int element_type, int64_t extent) -> NodeArg& { + ONNX_NAMESPACE::TypeProto type; + auto* tensor_type = type.mutable_tensor_type(); + tensor_type->set_elem_type(element_type); + tensor_type->mutable_shape()->add_dim()->set_dim_value(extent); + return graph.GetOrCreateNodeArg(name, &type); + }; + auto& source_arg = tensor_arg("source", ONNX_NAMESPACE::TensorProto_DataType_FLOAT, state_elements); + auto& capsule_arg = tensor_arg("capsule", ONNX_NAMESPACE::TensorProto_DataType_FLOAT, capsule_elements); + auto& destination_arg = tensor_arg("destination", ONNX_NAMESPACE::TensorProto_DataType_FLOAT, + static_cast(destination.size())); + auto& metadata_arg = tensor_arg("metadata", ONNX_NAMESPACE::TensorProto_DataType_INT64, 11); + auto& output_arg = graph.GetOrCreateNodeArg("output", nullptr); + auto& node = graph.AddNode("replay", "GatedDeltaNetStateReplay", "", + {&source_arg, &capsule_arg, &destination_arg, &metadata_arg}, + {&output_arg}, nullptr, kMSDomain); + node.SetExecutionProviderType(kWebGpuExecutionProvider); + graph.SetOutputs({&output_arg}); + ASSERT_STATUS_OK(graph.Resolve()); + std::stringstream stream; + ASSERT_TRUE(model.ToProto().SerializeToOstream(&stream)); + auto* ep = replay_ep.get(); + const auto allocators = ep->CreatePreferredAllocators(); + ASSERT_FALSE(allocators.empty()); + const auto memory_info = allocators[0]->Info(); + SessionOptions session_options; + session_options.graph_optimization_level = TransformerLevel::Default; + InferenceSession session(session_options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(replay_ep))); + ASSERT_STATUS_OK(session.Load(stream)); + ASSERT_STATUS_OK(session.Initialize()); + const auto allocator = session.GetAllocator(memory_info); + ASSERT_NE(allocator, nullptr); + CPUAllocator cpu_allocator; + auto upload = [&](const Tensor& cpu, OrtValue& gpu) -> Status { + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape{cpu.Shape().Size()}, allocator, gpu); + return ep->GetDataTransfer()->CopyTensor(cpu, *gpu.GetMutable()); + }; + Tensor incoming(DataTypeImpl::GetType(), TensorShape{state_elements}, + inputs.state0.data(), cpu_allocator.Info()); + Tensor cpu_destination(DataTypeImpl::GetType(), TensorShape{static_cast(destination.size())}, + destination.data(), cpu_allocator.Info()); + OrtValue source_value, capsule_value, destination_value, metadata_value; + ASSERT_STATUS_OK(upload(incoming, source_value)); + ASSERT_STATUS_OK(upload(captured, capsule_value)); + ASSERT_STATUS_OK(upload(cpu_destination, destination_value)); + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape{11}, descriptor.data(), + cpu_allocator.Info(), metadata_value); + std::unique_ptr binding; + ASSERT_STATUS_OK(session.NewIOBinding(&binding)); + ASSERT_STATUS_OK(binding->BindInput("source", source_value)); + ASSERT_STATUS_OK(binding->BindInput("capsule", capsule_value)); + ASSERT_STATUS_OK(binding->BindInput("destination", destination_value)); + ASSERT_STATUS_OK(binding->BindInput("metadata", metadata_value)); + ASSERT_STATUS_OK(binding->BindOutput("output", destination_value)); + + for (int kept = 1; kept <= geometry.total_tokens; ++kept) { + SCOPED_TRACE(kept); + Geometry prefix_geometry = geometry; + prefix_geometry.total_tokens = kept; + Inputs prefix_inputs = inputs; + prefix_inputs.q.resize(static_cast(kept * geometry.hq * geometry.dk)); + prefix_inputs.k.resize(static_cast(kept * geometry.hq * geometry.dk)); + prefix_inputs.v.resize(static_cast(kept * geometry.hv * geometry.dv)); + prefix_inputs.decay.resize(static_cast(kept * geometry.hv)); + prefix_inputs.beta.resize(static_cast(kept * geometry.hv)); + prefix_inputs.capture_count.clear(); + Options prefix_options = options; + prefix_options.state_update_capacity = 0; + std::vector prefix_fetches; + ASSERT_NO_FATAL_FAILURE(RunTypedCase(prefix_geometry, prefix_options, prefix_inputs, 4e-4f, 4e-4f, + /*rank4=*/false, &prefix_fetches, /*use_webgpu=*/true)); + ASSERT_EQ(prefix_fetches.size(), 3u); + const auto expected = prefix_fetches[1].Get().DataAsSpan(); + ASSERT_EQ(expected.size(), static_cast(state_elements)); + + descriptor[10] = kept; + ASSERT_STATUS_OK(session.Run(RunOptions{}, *binding)); + std::vector actual(destination.size()); + Tensor cpu_actual(DataTypeImpl::GetType(), cpu_destination.Shape(), actual.data(), cpu_allocator.Info()); + ASSERT_STATUS_OK(ep->GetDataTransfer()->CopyTensor(destination_value.Get(), cpu_actual)); + for (size_t index = 0; index < expected.size(); ++index) { + const float value = actual[static_cast(kDestinationOffset) + index]; + ASSERT_TRUE(std::isfinite(value)) << index; + ASSERT_LE(std::abs(value - expected[index]), 1e-5f + 2e-5f * std::abs(expected[index])) + << "state element " << index; + } + for (size_t index = 0; index < actual.size(); ++index) { + if (index < kDestinationOffset || index >= static_cast(kDestinationOffset + state_elements)) { + EXPECT_EQ(actual[index], destination[index]) << "neighboring destination element " << index; + } + } + } +} +#endif // USE_WEBGPU + TEST(GatedDeltaNetWebGpuTest, InactiveCompactStateUpdatesAreZero) { if (NeedSkipGatedDeltaNetWebGpuTest()) { GTEST_SKIP() << "WebGPU execution provider is not available"; @@ -1570,7 +1704,9 @@ TEST(GatedDeltaNetTest, AliasedStateIoBindingRecurrentAndChunked) { } TEST(GatedDeltaNetWebGpuTest, AliasedStateIoBinding) { - auto webgpu_ep = WebGpuExecutionProviderWithTestStorageBufferBindingSize(64 * 1024); + WebGpuExecutionProviderTestOptions test_options{}; + test_options.max_storage_buffer_binding_size = 64 * 1024; + auto webgpu_ep = WebGpuExecutionProviderWithTestOptions(test_options); if (webgpu_ep == nullptr) { GTEST_SKIP() << "WebGPU execution provider is not available"; } diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_state_replay_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_state_replay_op_test.cc new file mode 100644 index 0000000000000..f7d5abdb3c4b9 --- /dev/null +++ b/onnxruntime/test/contrib_ops/gated_delta_net_state_replay_op_test.cc @@ -0,0 +1,510 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/graph/model.h" +#include "gtest/gtest.h" +#include "test/unittest_util/framework_test_utils.h" +#include "test/util/include/asserts.h" +#include "test/util/include/default_providers.h" +#include "test/util/include/test_environment.h" + +#ifdef USE_WEBGPU +#include "core/providers/webgpu/webgpu_provider_options.h" +#include "core/session/IOBinding.h" +#include "core/session/inference_session.h" +#if !defined(ORT_USE_EP_API_ADAPTERS) && !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) +#include "core/providers/webgpu/webgpu_context.h" +#include "core/providers/webgpu/webgpu_execution_provider.h" +#include "test/providers/webgpu/webgpu_context_test_helper.h" +#endif +#endif + +namespace onnxruntime::test { +namespace { + +using Shapes = std::array, 4>; + +std::unique_ptr MakeReplayModel(const Shapes& shapes) { + auto model = std::make_unique( + "gated_delta_net_state_replay", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), std::unordered_map{{kMSDomain, 1}}, + std::vector{}, DefaultLoggingManager().DefaultLogger()); + auto& graph = model->MainGraph(); + constexpr std::array names{ + "source_backing", "capsule_backing", "destination_backing", "metadata"}; + std::vector inputs; + for (size_t index = 0; index < names.size(); ++index) { + ONNX_NAMESPACE::TypeProto type; + auto* tensor_type = type.mutable_tensor_type(); + tensor_type->set_elem_type(index == 3 ? ONNX_NAMESPACE::TensorProto_DataType_INT64 + : ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + for (const auto dim : shapes[index]) { + auto* dimension = tensor_type->mutable_shape()->add_dim(); + if (dim >= 0) { + dimension->set_dim_value(dim); + } else { + dimension->set_dim_param("backing_extent"); + } + } + inputs.push_back(&graph.GetOrCreateNodeArg(names[index], &type)); + } + auto& output = graph.GetOrCreateNodeArg("destination_out", nullptr); + auto& node = graph.AddNode("replay", "GatedDeltaNetStateReplay", "", inputs, {&output}, nullptr, kMSDomain); + node.SetExecutionProviderType(kWebGpuExecutionProvider); + graph.SetOutputs({&output}); + return model; +} + +TEST(GatedDeltaNetStateReplayShapeInferenceTest, DestinationShapeAndType) { + auto model = MakeReplayModel(Shapes{{{20}, {30}, {40}, {11}}}); + ASSERT_STATUS_OK(model->MainGraph().Resolve()); + const auto* output = model->MainGraph().GetOutputs()[0]; + ASSERT_EQ(output->TypeAsProto()->tensor_type().elem_type(), ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + ASSERT_EQ(output->Shape()->dim_size(), 1); + EXPECT_EQ(output->Shape()->dim(0).dim_value(), 40); + + model = MakeReplayModel(Shapes{{{-1}, {-1}, {-1}, {11}}}); + ASSERT_STATUS_OK(model->MainGraph().Resolve()); + EXPECT_EQ(model->MainGraph().GetOutputs()[0]->Shape()->dim(0).dim_param(), "backing_extent"); +} + +TEST(GatedDeltaNetStateReplayShapeInferenceTest, RejectsRanksAndMetadataLength) { + for (size_t index = 0; index < 4; ++index) { + auto shapes = Shapes{{{20}, {30}, {40}, {11}}}; + shapes[index] = {1, shapes[index][0]}; + auto model = MakeReplayModel(shapes); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(model->MainGraph().Resolve(), "inputs must have rank 1"); + } + for (const int64_t size : {0, 10, 12}) { + auto model = MakeReplayModel(Shapes{{{20}, {30}, {40}, {size}}}); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(model->MainGraph().Resolve(), "metadata must have shape [11]"); + } +} + +#ifdef USE_WEBGPU + +struct Geometry { + int64_t hv, dv, dk, hk, capacity; +}; + +class ReplayHarness { + public: + Status Initialize(const Geometry& geometry, + std::unique_ptr provider = DefaultWebGpuExecutionProvider(), + bool shared_source_capsule = false) { + ORT_RETURN_IF_NOT(provider, "WebGPU EP is unavailable."); + provider_ = provider.get(); + const auto allocators = provider_->CreatePreferredAllocators(); + ORT_RETURN_IF_NOT(!allocators.empty(), "WebGPU EP has no preferred allocator."); + const OrtMemoryInfo memory_info = allocators[0]->Info(); + + const auto [hv, dv, dk, hk, capacity] = geometry; + state_elements_ = static_cast(hv * dv * dk); + const int64_t row_elements = capacity * (hv + hk * dk + hv * dv); + desc_ = {static_cast(state_elements_) + 3, static_cast(state_elements_) + 5, + row_elements, row_elements + capacity * hv, row_elements + capacity * (hv + hk * dk), + hv, dv, dk, hk, capacity, 1}; + source_.resize(2 * state_elements_ + 11); + capsule_.resize(static_cast(2 * row_elements + 7)); + shared_source_capsule_ = shared_source_capsule; + if (shared_source_capsule_) { + const auto source_extent = source_.size(); + const auto shared_extent = source_extent + capsule_.size(); + source_.resize(shared_extent); + capsule_.resize(shared_extent); + for (size_t field = 2; field < 5; ++field) { + desc_[field] += static_cast(source_extent); + } + } + destination_.assign(3 * state_elements_ + 17, -931.25f); + for (size_t index = 0; index < source_.size(); ++index) { + source_[index] = static_cast(static_cast(index % 113) - 56) / 97.0f; + } + + auto model = MakeReplayModel(Shapes{{{static_cast(source_.size())}, + {static_cast(capsule_.size())}, + {static_cast(destination_.size())}, + {11}}}); + ORT_RETURN_IF_ERROR(model->MainGraph().Resolve()); + std::stringstream stream; + model->ToProto().SerializeToOstream(&stream); + SessionOptions session_options; + session_options.graph_optimization_level = TransformerLevel::Default; + session_ = std::make_unique(session_options, GetEnvironment()); + ORT_RETURN_IF_ERROR(session_->RegisterExecutionProvider(std::move(provider))); + ORT_RETURN_IF_ERROR(session_->Load(stream)); + ORT_RETURN_IF_ERROR(session_->Initialize()); + allocator_ = session_->GetAllocator(memory_info); + ORT_RETURN_IF_NOT(allocator_, "WebGPU allocator is unavailable."); + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape{static_cast(source_.size())}, + allocator_, source_value_); + if (shared_source_capsule_) { + capsule_value_ = source_value_; + } else { + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape{static_cast(capsule_.size())}, + allocator_, capsule_value_); + } + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape{static_cast(destination_.size())}, + allocator_, destination_value_); + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape{11}, desc_.data(), + cpu_allocator_.Info(), metadata_value_); + return Status::OK(); + } + + Status Prepare(int64_t kept_count) { + desc_[10] = kept_count; + std::fill(capsule_.begin(), capsule_.end(), std::numeric_limits::quiet_NaN()); + const auto hv = desc_[5], dv = desc_[6], dk = desc_[7], hk = desc_[8]; + for (int64_t t = 0; t < kept_count; ++t) { + for (int64_t h = 0; h < hv; ++h) { + capsule_[static_cast(desc_[2] + t * hv + h)] = 0.71f + 0.013f * static_cast((t + h) % 12); + for (int64_t v = 0; v < dv; ++v) { + capsule_[static_cast(desc_[4] + (t * hv + h) * dv + v)] = + static_cast((t * 11 + h * 7 + v) % 23 - 11) / 31.0f; + } + } + for (int64_t h = 0; h < hk; ++h) { + for (int64_t i = 0; i < dk; ++i) { + capsule_[static_cast(desc_[3] + (t * hk + h) * dk + i)] = + static_cast((t * 13 + h * 5 + i) % 29 - 14) / 37.0f; + } + } + } + if (shared_source_capsule_) { + std::copy_n(source_.begin() + desc_[0], state_elements_, capsule_.begin() + desc_[0]); + source_ = capsule_; + } + ORT_RETURN_IF_ERROR(Upload(source_, source_value_)); + ORT_RETURN_IF_ERROR(Upload(capsule_, capsule_value_)); + return Upload(destination_, destination_value_); + } + + Status Run(bool alias_output = true, const OrtValue* source_override = nullptr, + const OrtValue* capsule_override = nullptr) { + std::unique_ptr binding; + ORT_RETURN_IF_ERROR(session_->NewIOBinding(&binding)); + ORT_RETURN_IF_ERROR(binding->BindInput("source_backing", source_override ? *source_override : source_value_)); + ORT_RETURN_IF_ERROR(binding->BindInput("capsule_backing", capsule_override ? *capsule_override : capsule_value_)); + ORT_RETURN_IF_ERROR(binding->BindInput("destination_backing", destination_value_)); + ORT_RETURN_IF_ERROR(binding->BindInput("metadata", metadata_value_)); + OrtValue separate_output; + if (!alias_output) { + Tensor::InitOrtValue(DataTypeImpl::GetType(), destination_value_.Get().Shape(), + allocator_, separate_output); + } + ORT_RETURN_IF_ERROR(binding->BindOutput("destination_out", alias_output ? destination_value_ : separate_output)); + return session_->Run(RunOptions{}, *binding); + } + + std::vector Reference() const { + const auto hv = desc_[5], dv = desc_[6], dk = desc_[7], hk = desc_[8]; + std::vector state(source_.begin() + desc_[0], source_.begin() + desc_[0] + state_elements_); + // Token-major CPU oracle, independent of the shader's flattened per-element loop. + for (int64_t t = 0; t < desc_[10]; ++t) { + for (int64_t h = 0; h < hv; ++h) { + const float decay = capsule_[static_cast(desc_[2] + t * hv + h)]; + const int64_t key_head = h * hk / hv; + for (int64_t v = 0; v < dv; ++v) { + const float delta = capsule_[static_cast(desc_[4] + (t * hv + h) * dv + v)]; + for (int64_t i = 0; i < dk; ++i) { + auto& value = state[static_cast((h * dv + v) * dk + i)]; + value *= decay; + value += capsule_[static_cast(desc_[3] + (t * hk + key_head) * dk + i)] * delta; + } + } + } + } + return state; + } + + void Verify() { + const auto expected = Reference(); + std::vector actual(destination_.size()); + ASSERT_STATUS_OK(Download(destination_value_, actual)); + const auto offset = static_cast(desc_[1]); + for (size_t index = 0; index < expected.size(); ++index) { + ASSERT_TRUE(std::isfinite(actual[offset + index])) << index; + constexpr float kAbsoluteTolerance = 1e-5f, kRelativeTolerance = 2e-5f; + ASSERT_LE(std::abs(actual[offset + index] - expected[index]), + kAbsoluteTolerance + kRelativeTolerance * std::abs(expected[index])) + << "state element " << index; + } + EXPECT_EQ(std::memcmp(actual.data(), destination_.data(), offset * sizeof(float)), 0); + const auto tail = offset + expected.size(); + EXPECT_EQ(std::memcmp(actual.data() + tail, destination_.data() + tail, + (actual.size() - tail) * sizeof(float)), + 0); + std::vector source_copy(source_.size()), capsule_copy(capsule_.size()); + ASSERT_STATUS_OK(Download(source_value_, source_copy)); + ASSERT_STATUS_OK(Download(capsule_value_, capsule_copy)); + EXPECT_EQ(std::memcmp(source_copy.data(), source_.data(), source_.size() * sizeof(float)), 0); + // Byte comparison is intentional for preservation checks, including poisoned NaN tails. + EXPECT_EQ(std::memcmp(capsule_copy.data(), capsule_.data(), capsule_.size() * sizeof(float)), 0); + } + + void VerifyDestinationUnchanged() { + std::vector actual(destination_.size()); + ASSERT_STATUS_OK(Download(destination_value_, actual)); + EXPECT_EQ(std::memcmp(actual.data(), destination_.data(), destination_.size() * sizeof(float)), 0); + } + + std::array& Descriptor() { return desc_; } + + const Tensor& SourceBacking() const { return source_value_.Get(); } + + Status RunWithAliasedDestination(size_t input_index) { + OrtValue aliased_input; + const auto& input = (input_index == 0 ? source_value_ : capsule_value_).Get(); + Tensor::InitOrtValue(DataTypeImpl::GetType(), input.Shape(), + destination_value_.GetMutable()->MutableDataRaw(), + input.Location(), aliased_input); + return Run(true, input_index == 0 ? &aliased_input : nullptr, + input_index == 1 ? &aliased_input : nullptr); + } + + Status RunWithSeparateCapsuleWrapper() { + OrtValue capsule_wrapper; + const auto& backing = capsule_value_.Get(); + Tensor::InitOrtValue(DataTypeImpl::GetType(), backing.Shape(), + capsule_value_.GetMutable()->MutableDataRaw(), + backing.Location(), capsule_wrapper); + return Run(true, nullptr, &capsule_wrapper); + } + + private: + Status Upload(std::vector& data, OrtValue& gpu_value) { + Tensor cpu_tensor(DataTypeImpl::GetType(), TensorShape{static_cast(data.size())}, + data.data(), cpu_allocator_.Info()); + return provider_->GetDataTransfer()->CopyTensor(cpu_tensor, *gpu_value.GetMutable()); + } + + Status Download(const OrtValue& gpu_value, std::vector& data) { + Tensor cpu_tensor(DataTypeImpl::GetType(), TensorShape{static_cast(data.size())}, + data.data(), cpu_allocator_.Info()); + return provider_->GetDataTransfer()->CopyTensor(gpu_value.Get(), cpu_tensor); + } + + std::unique_ptr session_; + IExecutionProvider* provider_ = nullptr; + AllocatorPtr allocator_; + CPUAllocator cpu_allocator_; + std::vector source_, capsule_, destination_; + size_t state_elements_ = 0; + bool shared_source_capsule_ = false; + std::array desc_{}; + OrtValue source_value_, capsule_value_, destination_value_, metadata_value_; +}; + +class GatedDeltaNetStateReplayTest : public ::testing::Test { + protected: + void SetUp() override { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU EP is unavailable."; + } + } +}; + +TEST_F(GatedDeltaNetStateReplayTest, EveryPrefixWithNonDivisibleHeadsAndPoisonedTails) { + for (const int64_t capacity : {1, 4, 8}) { + SCOPED_TRACE(capacity); + ReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({5, 7, 11, 3, capacity})); + for (int64_t kept = 1; kept <= capacity; ++kept) { + SCOPED_TRACE(kept); + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } + } +} + +TEST_F(GatedDeltaNetStateReplayTest, QwenGeometryAndUnalignedRealCapsuleRow) { + ReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({48, 128, 128, 16, 7})); + EXPECT_EQ(harness.Descriptor()[2], 57680); + EXPECT_EQ(harness.Descriptor()[2] * static_cast(sizeof(float)) % 256, 64); + for (int64_t kept = 1; kept <= 7; ++kept) { + SCOPED_TRACE(kept); + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } +} + +#if !defined(ORT_USE_EP_API_ADAPTERS) && !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) +TEST_F(GatedDeltaNetStateReplayTest, RunCompletesBeforeReadback) { + auto provider = DefaultWebGpuExecutionProvider(); + ASSERT_NE(provider, nullptr); + auto& ep = static_cast(*provider); + auto& context = webgpu::WebGpuContextFactory::GetContext(ep.GetDeviceId()); + ReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({48, 128, 128, 16, 7}, std::move(provider))); + for (const int64_t kept : {1, 4, 7}) { + SCOPED_TRACE(kept); + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + + // No readback, Flush or blocking Wait may mask the operator's completion boundary. + const auto& recording = ep.Recording(); + ASSERT_TRUE(recording.deferred_dispatches.empty()); + ASSERT_EQ(recording.command_encoder, nullptr); + ASSERT_FALSE(recording.has_unsubmitted_work.load(std::memory_order_relaxed)); + auto result = std::make_shared(wgpu::QueueWorkDoneStatus::CallbackCancelled); + const auto future = context.Device().GetQueue().OnSubmittedWorkDone( + wgpu::CallbackMode::WaitAnyOnly, + [result](wgpu::QueueWorkDoneStatus status, wgpu::StringView /*message*/) noexcept { + *result = status; + }); + ASSERT_EQ(context.Instance().WaitAny(future, 0), wgpu::WaitStatus::Success); + ASSERT_EQ(*result, wgpu::QueueWorkDoneStatus::Success); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } +} + +TEST_F(GatedDeltaNetStateReplayTest, CheckedRegionReportsRunProgramBindingError) { + RunWithExternalCompletionContext([](auto& context, auto&, const auto&, auto& provider) { + ReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({5, 7, 11, 3, 4}, std::move(provider))); + ASSERT_STATUS_OK(harness.Prepare(2)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_STATUS_OK(harness.Prepare(2)); + + const auto& source = harness.SourceBacking(); + wgpu::BufferDescriptor desc{}; + desc.size = source.SizeInBytes(); + desc.usage = wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + auto buffer = context.Device().CreateBuffer(&desc); + OrtValue invalid_source; + Tensor::InitOrtValue(DataTypeImpl::GetType(), source.Shape(), buffer.Get(), + source.Location(), invalid_source); + + context.PushErrorScope(); + const auto status = harness.Run(true, &invalid_source); + const auto outer_status = context.PopErrorScope(); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("expected usage")); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("CreateBindGroup")); + EXPECT_TRUE(outer_status.IsOK()) << outer_status.ErrorMessage(); + harness.VerifyDestinationUnchanged(); + ASSERT_STATUS_OK(harness.Run()); + harness.Verify(); + }); +} +#endif + +TEST_F(GatedDeltaNetStateReplayTest, SegmentedBackingViews) { + WebGpuExecutionProviderTestOptions test_options{}; + test_options.max_storage_buffer_binding_size = 65536; + auto provider = WebGpuExecutionProviderWithTestOptions(test_options); + if (!provider) { + GTEST_SKIP() << "Artificially small storage binding limits require the native test factory."; + } + ReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({3, 64, 64, 16, 8}, std::move(provider))); + // Source/capsule each have two segments; destination has three. Logical views cross segment boundaries. + ASSERT_STATUS_OK(harness.Prepare(8)); + ASSERT_STATUS_OK(harness.Run()); + harness.Verify(); +} + +TEST_F(GatedDeltaNetStateReplayTest, SharedSourceCapsuleBindingOwner) { + constexpr uint64_t kSegmentBytes = 65536; + ConfigOptions options; + ASSERT_STATUS_OK(options.AddConfigEntry(webgpu::options::kMaxStorageBuffersPerShaderStage, "8")); + WebGpuExecutionProviderTestOptions test_options{}; + test_options.max_storage_buffer_binding_size = kSegmentBytes; + auto provider = WebGpuExecutionProviderWithTestOptions(options, test_options); + if (!provider) { + GTEST_SKIP() << "Artificially small storage binding limits require the native test factory."; + } + ReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({3, 64, 64, 16, 8}, std::move(provider), true)); + // Shared read-only backing: 44,050 floats (3 segments); destination: 36,881 floats (3 segments). + // Distinct Tensor wrappers do not share Program binding ownership, even with the same raw buffer. + ASSERT_STATUS_OK(harness.Prepare(3)); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.RunWithSeparateCapsuleWrapper(), "exceed WebGPU binding limits"); + harness.VerifyDestinationUnchanged(); + for (int64_t kept = 1; kept <= 8; ++kept) { + SCOPED_TRACE(kept); + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } +} + +TEST_F(GatedDeltaNetStateReplayTest, RejectsInvalidMetadataBeforeWriting) { + ReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({5, 7, 11, 3, 4})); + ASSERT_STATUS_OK(harness.Prepare(2)); + const auto valid = harness.Descriptor(); + for (size_t field = 0; field < valid.size(); ++field) { + SCOPED_TRACE(field); + for (const int64_t invalid : + std::array{-1, static_cast(std::numeric_limits::max()) + 1}) { + harness.Descriptor() = valid; + harness.Descriptor()[field] = invalid; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(), "outside the uint32 indexing range"); + } + } + for (size_t field = 5; field < 10; ++field) { + harness.Descriptor() = valid; + harness.Descriptor()[field] = 0; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(), "dimensions and capacity must be positive"); + } + for (const int64_t kept : {0, 5}) { + harness.Descriptor() = valid; + harness.Descriptor()[10] = kept; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(), "kept_count must be in [1, capacity]"); + } + for (size_t field = 0; field < 5; ++field) { + harness.Descriptor() = valid; + harness.Descriptor()[field] = std::numeric_limits::max(); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(), "view exceeds its backing tensor"); + } + for (size_t field = 5; field < 10; ++field) { + harness.Descriptor() = valid; + harness.Descriptor()[field] = std::numeric_limits::max(); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(), "geometry exceeds the uint32 indexing range"); + } + harness.VerifyDestinationUnchanged(); + harness.Descriptor() = valid; + ASSERT_STATUS_OK(harness.Run()); + harness.Verify(); +} + +TEST_F(GatedDeltaNetStateReplayTest, RequiresPreboundDestinationAlias) { + ReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({5, 7, 11, 3, 4})); + ASSERT_STATUS_OK(harness.Prepare(2)); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(false), "output must be prebound to destination_backing"); + harness.VerifyDestinationUnchanged(); + ASSERT_STATUS_OK(harness.Run()); + harness.Verify(); +} + +TEST_F(GatedDeltaNetStateReplayTest, RejectsWritableInputAliasing) { + ReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({5, 7, 11, 3, 4})); + ASSERT_STATUS_OK(harness.Prepare(2)); + for (size_t input_index = 0; input_index < 2; ++input_index) { + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.RunWithAliasedDestination(input_index), + "destination backing must be distinct"); + } + harness.VerifyDestinationUnchanged(); +} + +#endif // USE_WEBGPU + +} // namespace +} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/group_query_attention_growing_cache_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_growing_cache_test.cc new file mode 100644 index 0000000000000..ad87f741361ae --- /dev/null +++ b/onnxruntime/test/contrib_ops/group_query_attention_growing_cache_test.cc @@ -0,0 +1,329 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// GroupQueryAttention with a growing KV cache whose past_key/past_value are planner-owned values. +// Shape inference used to copy the past length onto present, symbolic or fixed, although present +// grows with the runtime total_sequence_length. The allocation planner then gave present a +// past-sized buffer (MayInplace alias or free list) and every growing step failed with +// "Shape mismatch attempting to re-use buffer". + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "core/graph/model.h" +#include "core/session/inference_session.h" +#include "test/test_environment.h" +#include "test/unittest_util/framework_test_utils.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" +#include "test/util/include/default_providers.h" +#include "test/util/include/test_utils.h" + +namespace onnxruntime { +namespace test { + +namespace { + +constexpr int64_t kBatch = 1; +constexpr int64_t kSeqLen = 1; +constexpr int64_t kNumHeads = 4; +constexpr int64_t kKvNumHeads = 2; +constexpr int64_t kHeadSize = 32; +constexpr int64_t kLayers = 2; + +using Dim = std::variant; + +// How the planner gets hold of a past-sized buffer to give to present. +enum class PastSizedBuffer { + // Gather(past_keys, layer) -> GroupQueryAttention. past_key dies at the node, so a + // MayInplace(past_key -> present_key) alias on the kernel def is honoured. + kPastItself, + // Gather -> Mul -> GroupQueryAttention. The Gather output is dead and on the planner's free list + // by the time the present outputs are planned. No kernel def alias is involved. + kFreeList, +}; + +struct ModelOptions { + PastSizedBuffer kind = PastSizedBuffer::kPastItself; + Dim past_len = std::string("past_len"); + // total_sequence_length as an initializer instead of a graph input. + std::optional constant_total_seq_len; + bool sliding_window_cache = false; +}; + +// Aggregate initialization with fewer members than the struct has trips -Wmissing-field-initializers. +ModelOptions Options(PastSizedBuffer kind, const Dim& past_len = std::string("past_len")) { + ModelOptions options; + options.kind = kind; + options.past_len = past_len; + return options; +} + +struct KvCacheModel { + std::string data; + // query, key, value, past_keys, past_values, seqlens_k, total_seq_len (empty when constant). + std::array input_names; + std::vector output_names; +}; + +// A KV cache held as one tensor for all layers and sliced per layer. +KvCacheModel BuildKvCacheModel(const ModelOptions& options) { + auto& logger = DefaultLoggingManager().DefaultLogger(); + Model model("GqaGrowingKvCache", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, logger); + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + + auto* query = builder.MakeSymbolicInput({kBatch, kSeqLen, kNumHeads * kHeadSize}); + auto* key = builder.MakeSymbolicInput({kBatch, kSeqLen, kKvNumHeads * kHeadSize}); + auto* value = builder.MakeSymbolicInput({kBatch, kSeqLen, kKvNumHeads * kHeadSize}); + auto* past_keys = builder.MakeSymbolicInput({kLayers, kBatch, kKvNumHeads, options.past_len, kHeadSize}); + auto* past_values = builder.MakeSymbolicInput({kLayers, kBatch, kKvNumHeads, options.past_len, kHeadSize}); + auto* seqlens_k = builder.MakeSymbolicInput({kBatch}); + auto* total_seq_len = options.constant_total_seq_len + ? builder.MakeScalarInitializer(*options.constant_total_seq_len) + : builder.MakeSymbolicInput({}); + + auto* layer_index = builder.MakeScalarInitializer(0); + auto* layer_axis = builder.Make1DInitializer({0}); + + auto slice_layer = [&](NodeArg* cache) { + auto* gathered = builder.MakeIntermediate(); + builder.AddNode("Gather", {cache, layer_index}, {gathered}).AddAttribute("axis", static_cast(0)); + if (options.kind == PastSizedBuffer::kPastItself) { + return gathered; + } + // Both Mul inputs are the same value, so Mul cannot run in place and `gathered` is freed. + auto* squared = builder.MakeIntermediate(); + builder.AddNode("Mul", {gathered, gathered}, {squared}); + return squared; + }; + auto* past_key = slice_layer(past_keys); + auto* past_value = slice_layer(past_values); + + auto* attn_out = builder.MakeOutput(); + auto* present_key = builder.MakeIntermediate(); + auto* present_value = builder.MakeIntermediate(); + Node& gqa = builder.AddNode("GroupQueryAttention", + {query, key, value, past_key, past_value, seqlens_k, total_seq_len}, + {attn_out, present_key, present_value}, kMSDomain); + gqa.AddAttribute("num_heads", kNumHeads); + gqa.AddAttribute("kv_num_heads", kKvNumHeads); + if (options.sliding_window_cache) { + gqa.AddAttribute("sliding_window_cache", static_cast(1)); + gqa.AddAttribute("local_window_size", static_cast(8)); + } + + auto* new_keys = builder.MakeOutput(); + auto* new_values = builder.MakeOutput(); + builder.AddNode("Unsqueeze", {present_key, layer_axis}, {new_keys}); + builder.AddNode("Unsqueeze", {present_value, layer_axis}, {new_values}); + + builder.SetGraphOutputs(); + EXPECT_STATUS_OK(graph.Resolve()); + + KvCacheModel result; + model.ToProto().SerializeToString(&result.data); + result.input_names = {query->Name(), key->Name(), value->Name(), past_keys->Name(), past_values->Name(), + seqlens_k->Name(), options.constant_total_seq_len ? "" : total_seq_len->Name()}; + result.output_names = {attn_out->Name(), new_keys->Name(), new_values->Name()}; + return result; +} + +// Small, exactly representable values keep the cross-EP comparison well conditioned. +std::vector SmallValues(size_t count, int seed) { + std::vector values(count); + for (size_t i = 0; i < count; ++i) { + values[i] = static_cast((static_cast(i) * 5 + seed) % 17 - 8) / 16.0f; + } + return values; +} + +template +void AddFeed(NameMLValMap& feeds, const std::string& name, + const std::vector& shape, const std::vector& data) { + OrtValue value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], shape, data, &value); + feeds.insert({name, std::move(value)}); +} + +struct Step { + // Length of the KV buffer that reaches the node. + int64_t past_len; + int32_t seqlens_k; + // How long the present cache has to be once this step's key and value are appended. + int32_t total_seq_len; +}; + +// Fed to a model with a symbolic past length. The first three steps grow present beyond the past +// buffer; the last two do not. +constexpr Step kSymbolicSteps[] = {{0, 0, 1}, {1, 1, 2}, {2, 2, 3}, {1, 0, 1}, {4, 3, 4}}; + +// Each fed to a model whose declared past length equals the step's past length. A fixed declared +// length does not make the cache static: total_sequence_length is still a runtime input. +constexpr Step kFixedSteps[] = {{1, 1, 2}, {2, 2, 3}, {16, 16, 17}}; + +std::string Describe(const Step& step) { + return "past_len=" + std::to_string(step.past_len) + " total_seq_len=" + std::to_string(step.total_seq_len); +} + +NameMLValMap MakeFeeds(const KvCacheModel& model, const Step& step) { + const size_t q_size = static_cast(kBatch * kSeqLen * kNumHeads * kHeadSize); + const size_t kv_size = static_cast(kBatch * kSeqLen * kKvNumHeads * kHeadSize); + const size_t past_size = static_cast(kLayers * kBatch * kKvNumHeads * step.past_len * kHeadSize); + const std::vector past_shape{kLayers, kBatch, kKvNumHeads, step.past_len, kHeadSize}; + + NameMLValMap feeds; + AddFeed(feeds, model.input_names[0], {kBatch, kSeqLen, kNumHeads * kHeadSize}, SmallValues(q_size, 0)); + AddFeed(feeds, model.input_names[1], {kBatch, kSeqLen, kKvNumHeads * kHeadSize}, SmallValues(kv_size, 3)); + AddFeed(feeds, model.input_names[2], {kBatch, kSeqLen, kKvNumHeads * kHeadSize}, SmallValues(kv_size, 7)); + AddFeed(feeds, model.input_names[3], past_shape, SmallValues(past_size, 11)); + AddFeed(feeds, model.input_names[4], past_shape, SmallValues(past_size, 13)); + AddFeed(feeds, model.input_names[5], {kBatch}, {step.seqlens_k}); + AddFeed(feeds, model.input_names[6], {}, {step.total_seq_len}); + return feeds; +} + +// The present_key sequence dimension inferred for the model's GroupQueryAttention node, read back +// from a fresh load with the shapes recorded at build time dropped. +ONNX_NAMESPACE::TensorShapeProto_Dimension InferredPresentLength(const KvCacheModel& model) { + ONNX_NAMESPACE::ModelProto proto; + EXPECT_TRUE(proto.ParseFromString(model.data)); + proto.mutable_graph()->clear_value_info(); + std::shared_ptr loaded; + EXPECT_STATUS_OK(Model::Load(std::move(proto), loaded, nullptr, DefaultLoggingManager().DefaultLogger())); + for (const Node& node : loaded->MainGraph().Nodes()) { + if (node.OpType() == "GroupQueryAttention") { + const auto* shape = node.OutputDefs()[1]->Shape(); + EXPECT_NE(shape, nullptr); + EXPECT_EQ(shape->dim_size(), 4); + EXPECT_EQ(shape->dim(1).dim_value(), kKvNumHeads); + EXPECT_EQ(shape->dim(3).dim_value(), kHeadSize); + return shape->dim(2); + } + } + ADD_FAILURE() << "no GroupQueryAttention node"; + return {}; +} + +void ExpectDynamic(const ONNX_NAMESPACE::TensorShapeProto_Dimension& dim) { + EXPECT_FALSE(dim.has_dim_value()) << dim.dim_value(); + EXPECT_FALSE(dim.has_dim_param()) << dim.dim_param(); +} + +void RunOnCpu(const KvCacheModel& model, const Step& step, bool enable_mem_reuse, std::vector& fetches) { + SessionOptions so; + so.session_logid = "GroupQueryAttentionGrowingCache"; + so.enable_mem_reuse = enable_mem_reuse; + InferenceSession session{so, GetEnvironment()}; + ASSERT_STATUS_OK(session.Load(model.data.data(), static_cast(model.data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + ASSERT_STATUS_OK(session.Run(RunOptions{}, MakeFeeds(model, step), model.output_names, &fetches)); +} + +// The CPU kernel declares no alias, so this is the free-list route alone; a session with memory +// reuse disabled is the reference. +void ExpectCpuReuseMatchesNoReuse(const KvCacheModel& model, const Step& step) { + std::vector expected; + std::vector actual; + ASSERT_NO_FATAL_FAILURE(RunOnCpu(model, step, /*enable_mem_reuse*/ false, expected)); + ASSERT_NO_FATAL_FAILURE(RunOnCpu(model, step, /*enable_mem_reuse*/ true, actual)); + + ASSERT_EQ(actual.size(), expected.size()); + const TensorShape new_keys_shape{1, kBatch, kKvNumHeads, std::max(step.past_len, step.total_seq_len), + kHeadSize}; + EXPECT_EQ(actual[1].Get().Shape(), new_keys_shape); + for (size_t i = 0; i < actual.size(); ++i) { + VerifyOutput(model.output_names[i], expected[i].Get(), actual[i].Get(), 1e-5f); + } +} + +void ExpectWebGpuMatchesCpu(const KvCacheModel& model, const Step& step, const char* log_id) { + const gsl::span model_bytes{reinterpret_cast(model.data.data()), + model.data.size()}; + // Gather alone would satisfy ExpectedEPNodeAssignment::Some. + const std::function gqa_runs_on_webgpu = [](const Graph& graph) { + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "GroupQueryAttention") { + EXPECT_EQ(node.GetExecutionProviderType(), kWebGpuExecutionProvider); + } + } + }; + EPVerificationParams params; + params.ep_node_assignment = ExpectedEPNodeAssignment::Some; + params.fp32_abs_err = 1e-4f; + params.graph_verifier = &gqa_runs_on_webgpu; + RunAndVerifyOutputsWithEP(model_bytes, log_id, DefaultWebGpuExecutionProvider(), MakeFeeds(model, step), params); +} + +void RunOnWebGpu(PastSizedBuffer kind, const char* log_id) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + const KvCacheModel symbolic = BuildKvCacheModel(Options(kind)); + for (const Step& step : kSymbolicSteps) { + SCOPED_TRACE("symbolic " + Describe(step)); + ExpectWebGpuMatchesCpu(symbolic, step, log_id); + } + for (const Step& step : kFixedSteps) { + SCOPED_TRACE("fixed " + Describe(step)); + ExpectWebGpuMatchesCpu(BuildKvCacheModel(Options(kind, Dim{step.past_len})), step, log_id); + } +} + +} // namespace + +TEST(GroupQueryAttentionGrowingCacheTest, SymbolicPastLengthLeavesPresentDynamic) { + ExpectDynamic(InferredPresentLength(BuildKvCacheModel(ModelOptions{}))); +} + +TEST(GroupQueryAttentionGrowingCacheTest, FixedPastLengthLeavesPresentDynamic) { + ExpectDynamic(InferredPresentLength(BuildKvCacheModel(Options(PastSizedBuffer::kPastItself, Dim{16})))); +} + +TEST(GroupQueryAttentionGrowingCacheTest, ConstantTotalLengthGivesExactPresentLength) { + ModelOptions grows = Options(PastSizedBuffer::kPastItself, Dim{16}); + grows.constant_total_seq_len = 17; + EXPECT_EQ(InferredPresentLength(BuildKvCacheModel(grows)).dim_value(), 17); + ModelOptions fits = Options(PastSizedBuffer::kPastItself, Dim{16}); + fits.constant_total_seq_len = 8; + EXPECT_EQ(InferredPresentLength(BuildKvCacheModel(fits)).dim_value(), 16); +} + +// A windowed cache is capacity-sized and evicts internally, so present keeps the past length. +TEST(GroupQueryAttentionGrowingCacheTest, SlidingWindowCacheKeepsPastLength) { + ModelOptions windowed = Options(PastSizedBuffer::kPastItself, Dim{16}); + windowed.sliding_window_cache = true; + EXPECT_EQ(InferredPresentLength(BuildKvCacheModel(windowed)).dim_value(), 16); +} + +TEST(GroupQueryAttentionGrowingCacheTest, PresentKvMayGrowBeyondFreedPastSizedBuffer) { + const KvCacheModel symbolic = BuildKvCacheModel(Options(PastSizedBuffer::kFreeList)); + for (const Step& step : kSymbolicSteps) { + SCOPED_TRACE("symbolic " + Describe(step)); + ExpectCpuReuseMatchesNoReuse(symbolic, step); + } + for (const Step& step : kFixedSteps) { + SCOPED_TRACE("fixed " + Describe(step)); + ExpectCpuReuseMatchesNoReuse(BuildKvCacheModel(Options(PastSizedBuffer::kFreeList, Dim{step.past_len})), step); + } +} + +TEST(GroupQueryAttentionGrowingCacheTest, PresentKvMayGrowBeyondPlannerOwnedPast_WebGPU) { + RunOnWebGpu(PastSizedBuffer::kPastItself, "GroupQueryAttentionGrowingCacheTest.PlannerOwnedPast_WebGPU"); +} + +TEST(GroupQueryAttentionGrowingCacheTest, PresentKvMayGrowBeyondFreedPastSizedBuffer_WebGPU) { + RunOnWebGpu(PastSizedBuffer::kFreeList, "GroupQueryAttentionGrowingCacheTest.FreedPastSizedBuffer_WebGPU"); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index c38b3788bc1dc..8eff8420fe527 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -1879,9 +1879,9 @@ TEST(GroupQueryAttentionTest, FlashDecoding_LocalWindow_Reference_CPU) { constexpr int kv_block_size = 16; constexpr int kv_chunk_count = (total_length + kv_block_size - 1) / kv_block_size; constexpr int thread_count = 12; + // thread_count fixes the MLAS task count, not the pool's DegreeOfParallelism (which varies by machine). concurrency::ThreadPool thread_pool(&Env::Default(), ThreadOptions{}, nullptr, thread_count, concurrency::kSpinDurationDefault, false); - ASSERT_EQ(concurrency::ThreadPool::DegreeOfParallelism(&thread_pool), thread_count); for (int head_size : {8, 24, 256}) { for (int window_size : {-1, 1, 24, 40}) { SCOPED_TRACE(testing::Message() << "head_size=" << head_size << " window=" << window_size); diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 451d3d591cc0e..73d2c27567ed5 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -2187,6 +2187,33 @@ TEST(MatMulNBits, PrePack_InvalidBShape_RejectsAtSessionInit) { {}, nullptr, &execution_providers); } +TEST(MatMulNBits, PrePack_Fp16InvalidBShape_RejectsAtSessionInit) { + constexpr int64_t M = 1, N = 4, K = 32, block_size = 32; + constexpr int64_t k_blocks = (K + block_size - 1) / block_size; + constexpr int64_t blob_size = block_size * QBits / 8; + constexpr int64_t bad_N = N - 1; + + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("block_size", block_size); + test.AddAttribute("bits", QBits); + test.AddAttribute("accuracy_level", int64_t{0}); + + test.AddInput("A", {M, K}, std::vector(M * K, MLFloat16(1.0f)), false); + test.AddInput("B", {bad_N, k_blocks, blob_size}, + std::vector(bad_N * k_blocks * blob_size, 0), true); + test.AddInput("scales", {N, k_blocks}, + std::vector(N * k_blocks, MLFloat16(1.0f)), true); + test.AddOutput("Y", {M, N}, std::vector(M * N, MLFloat16(0.0f))); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, + "MatMulNBits PrePack: B initializer shape", + {}, nullptr, &execution_providers); +} + // B shape has the wrong rank (2D instead of (N, k_blocks, blob_size)). TEST(MatMulNBits, PrePack_InvalidBRank_RejectsAtSessionInit) { constexpr int64_t M = 1, N = 4, K = 32, block_size = 32; @@ -2257,6 +2284,33 @@ TEST(MatMulNBits, PrePack_InvalidScalesShape_RejectsAtSessionInit) { {}, nullptr, &execution_providers); } +TEST(MatMulNBits, PrePack_Fp16InvalidScalesShape_RejectsAtSessionInit) { + constexpr int64_t M = 1, N = 4, K = 32, block_size = 32; + constexpr int64_t k_blocks = (K + block_size - 1) / block_size; + constexpr int64_t blob_size = block_size * QBits / 8; + constexpr int64_t bad_N = N - 1; + + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("block_size", block_size); + test.AddAttribute("bits", QBits); + test.AddAttribute("accuracy_level", int64_t{0}); + + test.AddInput("A", {M, K}, std::vector(M * K, MLFloat16(1.0f)), false); + test.AddInput("B", {N, k_blocks, blob_size}, + std::vector(N * k_blocks * blob_size, 0), true); + test.AddInput("scales", {bad_N, k_blocks}, + std::vector(bad_N * k_blocks, MLFloat16(1.0f)), true); + test.AddOutput("Y", {M, N}, std::vector(M * N, MLFloat16(0.0f))); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, + "MatMulNBits PrePack: scales initializer shape", + {}, nullptr, &execution_providers); +} + // uint8 (packed) zero_points shape does not match the // [N * zp_blob_size] / [N, zp_blob_size] layout derived from attributes. TEST(MatMulNBits, PrePack_InvalidUInt8ZeroPointsShape_RejectsAtSessionInit) { diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 12148390f63aa..4ed7e9e4a3968 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -147,6 +147,115 @@ static void RunMoETest(const std::vector& input, const std::vector // TODO(wy): Add python parity tests that can serve as examples. Need cutlass upgrade to build cutlass extensions to // add weights preprocesser to onnxruntime_pybind_quant.cc +template +static void RunCudaReluSquaredMoETest(int minimum_cuda_arch, int expert_weight_bits = 0) { + if (!HasCudaEnvironment(minimum_cuda_arch)) { + GTEST_SKIP() << "CUDA architecture does not support the test element type"; + } + + constexpr int num_rows = 2; + constexpr int num_experts = 2; + constexpr int hidden_size = 128; + constexpr int inter_size = 128; + std::vector input(num_rows * hidden_size); + std::vector expected(input.size()); + std::vector fc1(num_experts * inter_size * hidden_size, ElementType(0.0f)); + std::vector fc2(num_experts * hidden_size * inter_size, ElementType(0.0f)); + std::vector fc1_bias(num_experts * inter_size); + std::vector fc2_bias(num_experts * hidden_size); + + for (int expert = 0; expert < num_experts; ++expert) { + for (int channel = 0; channel < hidden_size; ++channel) { + const int diagonal = expert * hidden_size * inter_size + channel * hidden_size + channel; + fc1[diagonal] = ElementType(static_cast(expert + 1)); + fc2[diagonal] = ElementType(1.0f); + fc1_bias[expert * inter_size + channel] = ElementType(expert == 0 ? 0.5f : -0.5f); + fc2_bias[expert * hidden_size + channel] = ElementType(expert == 0 ? 0.25f : -0.25f); + } + } + for (int row = 0; row < num_rows; ++row) { + for (int channel = 0; channel < hidden_size; ++channel) { + const float value = static_cast((channel + row) % 8 - 3) * 0.5f; + input[row * hidden_size + channel] = ElementType(value); + const float first = std::max(0.0f, value + 0.5f); + const float second = std::max(0.0f, 2.0f * value - 0.5f); + expected[row * hidden_size + channel] = ElementType(0.5f * (first * first + second * second)); + } + } + + OpTester tester(expert_weight_bits == 0 ? "MoE" : "QMoE", 1, onnxruntime::kMSDomain); + tester.AddAttribute("k", 2); + tester.AddAttribute("activation_type", "relu2"); + tester.AddAttribute("normalize_routing_weights", 1); + tester.AddInput("input", {num_rows, hidden_size}, input); + tester.AddInput("router_probs", {num_rows, num_experts}, + std::vector(num_rows * num_experts, ElementType(0.5f))); + if (expert_weight_bits == 0) { + tester.AddInput("fc1_experts_weights", {num_experts, inter_size, hidden_size}, fc1); + tester.AddInput("fc1_experts_bias", {num_experts, inter_size}, fc1_bias); + tester.AddInput("fc2_experts_weights", {num_experts, hidden_size, inter_size}, fc2); + tester.AddInput("fc2_experts_bias", {num_experts, hidden_size}, fc2_bias); + } else { + ASSERT_TRUE(expert_weight_bits == 4 || expert_weight_bits == 8); + const int pack_size = 8 / expert_weight_bits; + const int zero_point = 1 << (expert_weight_bits - 1); + const uint8_t packed_zero = expert_weight_bits == 4 ? 0x88 : 0x80; + std::vector quantized_fc1(num_experts * hidden_size * inter_size / pack_size, packed_zero); + std::vector quantized_fc2(quantized_fc1.size(), packed_zero); + for (int expert = 0; expert < num_experts; ++expert) { + for (int channel = 0; channel < hidden_size; ++channel) { + const int diagonal = expert * hidden_size * inter_size / pack_size + + channel * inter_size / pack_size + channel / pack_size; + const int shift = (channel % pack_size) * expert_weight_bits; + const uint8_t mask = static_cast(((1 << expert_weight_bits) - 1) << shift); + quantized_fc1[diagonal] = static_cast( + (quantized_fc1[diagonal] & ~mask) | ((zero_point + expert + 1) << shift)); + quantized_fc2[diagonal] = static_cast( + (quantized_fc2[diagonal] & ~mask) | ((zero_point + 1) << shift)); + } + } + tester.AddAttribute("expert_weight_bits", expert_weight_bits); + tester.AddAttribute("weights_prepacked", 0); + tester.AddInput("fc1_experts_weights", + {num_experts, inter_size, hidden_size / pack_size}, quantized_fc1, true); + tester.AddInput("fc1_scales", {num_experts, inter_size}, + std::vector(num_experts * inter_size, ElementType(1.0f)), true); + tester.AddInput("fc1_experts_bias", {num_experts, inter_size}, fc1_bias); + tester.AddInput("fc2_experts_weights", + {num_experts, hidden_size, inter_size / pack_size}, quantized_fc2, true); + tester.AddInput("fc2_scales", {num_experts, hidden_size}, + std::vector(num_experts * hidden_size, ElementType(1.0f)), true); + tester.AddInput("fc2_experts_bias", {num_experts, hidden_size}, fc2_bias); + } + tester.AddOutput("output", {num_rows, hidden_size}, expected); + tester.SetOutputAbsErr("output", 0.001f); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(MoETest, CudaReluSquaredFloat32) { + RunCudaReluSquaredMoETest(700); +} + +TEST(MoETest, CudaReluSquaredFloat16) { + RunCudaReluSquaredMoETest(700); +} + +TEST(MoETest, CudaReluSquaredBFloat16) { + RunCudaReluSquaredMoETest(800); +} + +TEST(MoETest, CudaReluSquaredQuantizedInt4) { + RunCudaReluSquaredMoETest(750, 4); +} + +TEST(MoETest, CudaReluSquaredQuantizedInt8) { + RunCudaReluSquaredMoETest(750, 8); +} + static void RunQMoETest(const std::vector& input, const std::vector& router_probs, const std::vector& fc1_experts_weights, const std::vector& fc2_experts_weights, diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_merge_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_merge_op_test.cc index 87eedc3b39dc7..e4b88bb789ec5 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_merge_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_merge_op_test.cc @@ -4,6 +4,7 @@ #include #include "gtest/gtest.h" +#include "contrib_ops/cpu/sparse/packed_sparse_attention_indexer_merge_common.h" #include "core/graph/constants.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" @@ -53,6 +54,43 @@ TEST(PackedSparseAttentionIndexerMerge, StableOverlapAndSharedRows) { {4, 5, 5}, {0, 0, 0}); } +TEST(PackedSparseAttentionIndexerMerge, HashCapacityUsesInsertionBound) { + using contrib::indexer_merge::HashCapacity; + EXPECT_EQ(HashCapacity(int64_t{1} << 30, 3), 8); + EXPECT_EQ(HashCapacity(8, 8), 16); + EXPECT_EQ(HashCapacity(4, INT32_MAX), 16); + EXPECT_EQ(HashCapacity(int64_t{1} << 30, INT32_MAX), int64_t{1} << 32); + EXPECT_EQ(HashCapacity(8, 0), 1); +} + +TEST(PackedSparseAttentionIndexerMerge, PowerOfTwoCapacityAndDuplicatePrefixes) { + RunMerge(1, 4, 4, {0, 1, 2, 3}, {4}, {0}, {0}, {4}, {0, 1, 2, 3}, {4}, {0}); + RunMerge(1, 8, 2, {1, 1, 1, 1, 2, 2, 2, 2}, {8}, {0}, {1}, {3}, {1, 2}, {2}, {0}); + RunMerge(1, 8, 2, {1, 1, 2, 2, 3, 3, 4, 4}, {8}, {0}, {0}, {0}, {-1, -1}, {0}, {2}); + RunMerge(1, 513, 2, std::vector(513, 7), {513}, {0}, {7}, {9}, {7, 8}, {2}, {0}); +} + +TEST(PackedSparseAttentionIndexerMerge, AppendIndicesDuplicateChunks) { + auto provider = DefaultCudaExecutionProvider(); + if (!provider) GTEST_SKIP() << "Requires CUDA"; + OpTester test("PackedSparseAttentionIndexerMerge", 1, kMSDomain); + test.AddAttribute("policy_mode", std::string("append_indices")); + test.AddAttribute("max_output_entries", int64_t{2}); + test.AddInput("base", {1, 257}, std::vector(257, 7)); + test.AddInput("counts", {1}, {257}); + test.AddInput("rows", {1}, {0}); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddInput("additional", {1, 256}, std::vector(256, 8)); + test.AddInput("additional_counts", {1}, {256}); + test.AddOutput("selected", {1, 2}, {7, 8}); + test.AddOutput("selected_counts", {1}, {2}); + test.AddOutput("status", {1}, {0}); + std::vector> providers; + providers.push_back(std::move(provider)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + TEST(PackedSparseAttentionIndexerMerge, InvalidMetadataAndOverflow) { RunMerge(2, 3, 3, {1, 2, 3, -1, 2, 3}, {3, 3}, {-1, 2, 0, 0, 0, 1, 0}, {0, 0, -1, 5, 3, 0, 0}, @@ -104,7 +142,7 @@ TEST(PackedSparseAttentionIndexerMerge, QwenCapacity) { RunMerge(1, 2051, 2058, base, {2051}, {0}, {2051}, {2058}, expected, {2058}, {0}); } -TEST(PackedSparseAttentionIndexerMerge, LargeCapacityFallback) { +TEST(PackedSparseAttentionIndexerMerge, LargeOutputCapacityWithSmallBase) { std::vector expected(4097, -1); expected[0] = 7; expected[1] = 5; @@ -112,6 +150,22 @@ TEST(PackedSparseAttentionIndexerMerge, LargeCapacityFallback) { RunMerge(1, 5, 4097, {7, 7, 5, 5, -1}, {4}, {0}, {5}, {8}, expected, {3}, {0}); } +TEST(PackedSparseAttentionIndexerMerge, StableCompactionAcrossWorkgroupChunks) { + std::vector base(513); + for (int32_t column = 0; column < 513; ++column) base[column] = (column * 37) % 257; + std::vector expected(base.begin(), base.begin() + 257); + for (int32_t value = 257; value < 265; ++value) expected.push_back(value); + expected.resize(270, -1); + expected.insert(expected.end(), base.begin(), base.begin() + 257); + expected.resize(540, -1); + expected.resize(810, -1); + RunMerge(1, 513, 270, base, {513}, {0, 0, 0}, {250, 0, 0}, {265, 0, 300}, + expected, {265, 257, 0}, {0, 0, 2}); + + base.back() = -1; + RunMerge(1, 513, 2, base, {513}, {0}, {0}, {0}, {-1, -1}, {0}, {1}); +} + #ifndef ORT_NO_EXCEPTIONS TEST(PackedSparseAttentionIndexerMerge, RejectUnsupportedPolicy) { OpTester test("PackedSparseAttentionIndexerMerge", 1, kMSDomain); diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 6baeb9fd32725..27406a7a76cac 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -107,6 +107,14 @@ struct GraphOptions { bool omit_csa_gate_output = false; int output_count = psai::kFixedOutputCount; std::string policy_mode = psai::kPolicyModeQsa; + bool indexshare = false; + std::string indexshare_state_policy = "default"; + std::vector mode_shape{1}; + int64_t merged_rows = 0; + int64_t merged_width = 8; + int64_t merged_counts_rows = -1; + int64_t merged_status_rows = -1; + int64_t projection_tokens = -1; }; int64_t BufferCapacity(int64_t compress_ratio) { return 2 * compress_ratio - 1; } @@ -117,6 +125,8 @@ void AddNode(ModelTestBuilder& builder, const GraphOptions& options) { const bool is_csa = options.policy_mode == psai::kPolicyModeCsa; const int64_t width = is_csa ? 2 * options.head_size : options.head_size; const int64_t buffer_capacity = BufferCapacity(options.compress_ratio); + const int64_t projection_tokens = + options.projection_tokens >= 0 ? options.projection_tokens : options.total_tokens; NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); std::vector inputs{ @@ -135,7 +145,7 @@ void AddNode(ModelTestBuilder& builder, const GraphOptions& options) { builder.MakeInput(std::vector{options.batch_size}), }; if (is_csa || options.add_csa_inputs) { - inputs.push_back(builder.MakeInput(std::vector{options.total_tokens, width})); + inputs.push_back(builder.MakeInput(std::vector{projection_tokens, width})); inputs.push_back(builder.MakeInput(std::vector{options.compress_ratio, width})); inputs.push_back(builder.MakeInput(std::vector{options.total_tokens, options.num_heads})); } else { @@ -144,7 +154,7 @@ void AddNode(ModelTestBuilder& builder, const GraphOptions& options) { inputs.push_back(&empty); } if (is_csa || options.add_position_ids) { - inputs.push_back(builder.MakeInput(std::vector{options.total_tokens})); + inputs.push_back(builder.MakeInput(std::vector{projection_tokens})); } else { inputs.push_back(&empty); } @@ -171,16 +181,30 @@ void AddNode(ModelTestBuilder& builder, const GraphOptions& options) { if (options.add_state_update_active) { inputs.push_back(builder.MakeInput(std::vector{1})); } + if (options.indexshare) { + inputs.resize(psai::kIndexShareMode, &empty); + inputs.push_back(builder.MakeInput(options.mode_shape)); + inputs.push_back(builder.MakeInput(std::vector{options.merged_rows, options.merged_width})); + inputs.push_back(builder.MakeInput(std::vector{ + options.merged_counts_rows >= 0 ? options.merged_counts_rows : options.merged_rows})); + inputs.push_back(builder.MakeInput(std::vector{ + options.merged_status_rows >= 0 ? options.merged_status_rows : options.merged_rows})); + } std::vector outputs; - for (int i = 0; i < options.output_count; ++i) { - const bool omit_gate = (!is_csa && !options.add_qsa_gate_output) || (is_csa && options.omit_csa_gate_output); - outputs.push_back(i == psai::kPresentGateBuffer && omit_gate ? &empty : builder.MakeOutput()); + for (int i = 0; i < (options.indexshare ? psai::kIndexShareStatus + 1 : options.output_count); ++i) { + const bool omit_gate_output = + i == psai::kPresentGateBuffer && + ((is_csa && options.omit_csa_gate_output) || (!is_csa && !options.add_qsa_gate_output)); + const bool omit_capture_output = options.indexshare && i == psai::kStateUpdate && options.output_count <= i; + outputs.push_back(omit_gate_output || omit_capture_output ? &empty : builder.MakeOutput()); } Node& node = builder.AddNode("PackedSparseAttentionIndexer", inputs, outputs, kMSDomain); node.AddAttribute("policy_mode", options.policy_mode); node.AddAttribute("compress_ratio", options.compress_ratio); node.AddAttribute("state_capacity", options.state_capacity); + node.AddAttribute("indexshare_state_policy", options.indexshare_state_policy); + if (options.indexshare) node.AddAttribute("max_output_entries", int64_t{8}); if (options.state_update_capacity != 0) { node.AddAttribute("state_update_capacity", options.state_update_capacity); } @@ -193,6 +217,93 @@ void AddNode(ModelTestBuilder& builder, const GraphOptions& options) { } // namespace +TEST(PackedSparseAttentionIndexerShapeInferenceTest, ValidatesIndexShareInterface) { + for (int invalid = 0; invalid < 8; ++invalid) { + GraphOptions options; + options.indexshare = true; + switch (invalid) { + case 0: + options.mode_shape = {2}; + break; + case 1: + options.mode_shape = {1, 1}; + break; + case 2: + options.merged_width = 7; + break; + case 3: + options.merged_counts_rows = 1; + break; + case 4: + options.merged_status_rows = 1; + break; + case 5: + options.merged_rows = 3; + break; + case 6: + options.indexshare_state_policy = "invalid"; + break; + case 7: + options.indexshare_state_policy = "advance"; + break; + } + std::unique_ptr model; + SCOPED_TRACE(invalid); + EXPECT_FALSE(BuildAndResolve([&](ModelTestBuilder& builder) { AddNode(builder, options); }, model).IsOK()); + } + GraphOptions options; + options.indexshare = true; + std::unique_ptr model; + ASSERT_STATUS_OK(BuildAndResolve([&](ModelTestBuilder& builder) { AddNode(builder, options); }, model)); +} + +TEST(PackedSparseAttentionIndexerShapeInferenceTest, IndexShareCsaValidatesProjectionRows) { + for (bool mismatch : {false, true}) { + GraphOptions options; + options.indexshare = true; + options.policy_mode = psai::kPolicyModeCsa; + options.total_tokens = 0; + options.key_total_tokens = 2; + options.projection_tokens = mismatch ? 1 : 2; + options.merged_rows = 2; + std::unique_ptr model; + const auto status = BuildAndResolve([&](ModelTestBuilder& builder) { AddNode(builder, options); }, model); + EXPECT_EQ(status.IsOK(), !mismatch) << status.ErrorMessage(); + } +} + +TEST(PackedSparseAttentionIndexerShapeInferenceTest, MergeValidatesPolicyInputsAndShapes) { + for (bool append_indices : {false, true}) { + for (int invalid = -1; invalid < 8; ++invalid) { + std::unique_ptr model; + const auto status = BuildAndResolve([&](ModelTestBuilder& builder) { + NodeArg* empty = &builder.graph_.GetOrCreateNodeArg("", nullptr); + std::vector inputs{ + builder.MakeInput(std::vector{2, 4}), + builder.MakeInput(invalid == 0 ? std::vector{2, 1} : invalid == 1 ? std::vector{3} + : std::vector{2}), + builder.MakeInput(std::vector{3}), empty, empty, empty, empty}; + const int first = append_indices ? 5 : 3; + inputs[first] = builder.MakeInput(invalid == 2 ? std::vector{3, 1, 1} : append_indices ? std::vector{3, 2} + : std::vector{3}); + inputs[first + 1] = builder.MakeInput(invalid == 3 ? std::vector{2} : invalid == 7 ? std::vector{3, 1} + : std::vector{3}); + if (invalid == 4) inputs[first] = empty; + if (invalid == 5) inputs[append_indices ? 3 : 5] = builder.MakeInput(std::vector{3}); + if (invalid == 6) inputs[first] = builder.MakeInput( + append_indices ? std::vector{2, 2} : std::vector{2}); + Node& node = builder.AddNode("PackedSparseAttentionIndexerMerge", inputs, + {builder.MakeOutput(), builder.MakeOutput(), builder.MakeOutput()}, kMSDomain); + node.AddAttribute("policy_mode", std::string(append_indices ? "append_indices" : "append_range")); + node.AddAttribute("max_output_entries", int64_t{8}); + }, + model); + SCOPED_TRACE(invalid); + EXPECT_EQ(status.IsOK(), invalid == -1) << status.ErrorMessage(); + } + } +} + TEST(PackedSparseAttentionIndexerShapeInferenceTest, QsaInfersFixedCapacityAndState) { GraphOptions options; std::unique_ptr model; @@ -797,7 +908,7 @@ QsaPackedProblem MakeQsaPackedProblem(QsaPackedProblem problem = {}) { template void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedProblem(), ProviderKind provider_kind = ProviderKind::Cuda, QsaPackedResult* actual = nullptr, - bool packed_qk = false, int padding = 0, const ConfigOptions* webgpu_options = nullptr) { + bool packed_qk = false, const ConfigOptions* webgpu_options = nullptr, int padding = 0) { auto provider = CreateProvider(provider_kind, webgpu_options); if (provider == nullptr) { GTEST_SKIP() << (provider_kind == ProviderKind::Cuda ? "CUDA" : "WebGPU") @@ -932,47 +1043,60 @@ TEST(PackedSparseAttentionIndexerTest, QsaPackedQueryKey) { } TEST(PackedSparseAttentionIndexerTest, IndexShareRefreshPreservesBudgetAndPadsRows) { - RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(), ProviderKind::Cuda, nullptr, true, 6); + RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(), ProviderKind::Cuda, nullptr, true, nullptr, 6); } TEST(PackedSparseAttentionIndexerTest, IndexShareReusePackedQkAndStatus) { - auto provider = DefaultCudaExecutionProvider(); - if (!provider) GTEST_SKIP() << "Requires CUDA"; - OpTester test("PackedSparseAttentionIndexer", 1, kMSDomain); - test.AddAttribute("policy_mode", std::string("qsa")); - test.AddAttribute("compress_ratio", int64_t{2}); - test.AddAttribute("state_capacity", int64_t{4}); - test.AddAttribute("token_budget", int64_t{2}); - test.AddAttribute("max_output_entries", int64_t{5}); - test.AddInput("query_key", {0, 4}, {}); - test.AddOptionalInputEdge(); - test.AddInput("q_norm", {2}, {1, 1}); - test.AddInput("k_norm", {2}, {1, 1}); - test.AddInput("cos", {8, 2}, std::vector(16, 1)); - test.AddInput("sin", {8, 2}, std::vector(16, 0)); - test.AddInput("cumulative", {3}, {0, 1, 2}); - test.AddInput("past_lengths", {2}, {5, 5}); - for (int index = 8; index <= 10; ++index) test.AddOptionalInputEdge(); - test.AddOptionalInputEdge(); - test.AddInput("past_keys", {2, 4, 2}, std::vector(16, 1)); - test.AddInput("past_buffer", {2, 3, 2}, std::vector(12, 2)); - test.AddOptionalInputEdge(); - test.AddInput("state_lengths", {2, 2}, {2, 1, 2, 1}); - test.AddOptionalInputEdge(); - test.AddOptionalInputEdge(); - test.AddInput("mode", {1}, {1}); - test.AddInput("merged_indices", {2, 5}, {0, 2, 5, -1, -1, -1, -1, -1, -1, -1}); - test.AddInput("merged_counts", {2}, {3, 0}); - test.AddInput("merged_status", {2}, {0, 2}); - test.AddOutput("selected", {2, 5}, {0, 2, 5, -1, -1, -1, -1, -1, -1, -1}); - test.AddOutput("selected_counts", {2}, {3, 0}); - test.AddOutput("present_keys", {2, 4, 2}, std::vector(16, 1)); - test.AddOutput("present_buffer", {2, 3, 2}, std::vector(12, 2)); - test.AddOptionalOutputEdge(); - test.AddOutput("present_lengths", {2, 2}, {2, 1, 2, 1}); - test.AddOptionalOutputEdge(); - test.AddOutput("status", {2}, {0, 2}); - RunOnProvider(test, std::move(provider)); + for (int invalid = 0; invalid < 7; ++invalid) { + SCOPED_TRACE(invalid); + std::vector cumulative{0, 1, 2}; + std::vector past_lengths{5, 5}; + std::vector state_lengths{2, 1, 2, 1}; + std::vector merged_counts{3, 3}; + if (invalid == 1) merged_counts[0] = -1; + if (invalid == 2) merged_counts[0] = 6; + if (invalid == 3) cumulative[2] = 5; + if (invalid == 4) past_lengths[0] = -1; + if (invalid == 5) state_lengths[0] = 5; + if (invalid == 6) state_lengths[1] = 0; + auto provider = DefaultCudaExecutionProvider(); + if (!provider) GTEST_SKIP() << "Requires CUDA"; + OpTester test("PackedSparseAttentionIndexer", 1, kMSDomain); + test.AddAttribute("policy_mode", std::string("qsa")); + test.AddAttribute("compress_ratio", int64_t{2}); + test.AddAttribute("state_capacity", int64_t{4}); + test.AddAttribute("token_budget", int64_t{2}); + test.AddAttribute("max_output_entries", int64_t{5}); + test.AddInput("query_key", {0, 4}, {}); + test.AddOptionalInputEdge(); + test.AddInput("q_norm", {2}, {1, 1}); + test.AddInput("k_norm", {2}, {1, 1}); + test.AddInput("cos", {8, 2}, std::vector(16, 1)); + test.AddInput("sin", {8, 2}, std::vector(16, 0)); + test.AddInput("cumulative", {3}, cumulative); + test.AddInput("past_lengths", {2}, past_lengths); + for (int index = 8; index <= 10; ++index) test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddInput("past_keys", {2, 4, 2}, std::vector(16, 1)); + test.AddInput("past_buffer", {2, 3, 2}, std::vector(12, 2)); + test.AddOptionalInputEdge(); + test.AddInput("state_lengths", {2, 2}, state_lengths); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddInput("mode", {1}, {1}); + test.AddInput("merged_indices", {2, 5}, {0, 2, 5, -1, -1, 9, 8, 7, 6, 5}); + test.AddInput("merged_counts", {2}, merged_counts); + test.AddInput("merged_status", {2}, {0, 2}); + test.AddOutput("selected", {2, 5}, invalid == 0 ? std::vector{0, 2, 5, -1, -1, -1, -1, -1, -1, -1} : std::vector(10, -1)); + test.AddOutput("selected_counts", {2}, {invalid == 0 ? 3 : 0, 0}); + test.AddOutput("present_keys", {2, 4, 2}, std::vector(16, 1)); + test.AddOutput("present_buffer", {2, 3, 2}, std::vector(12, 2)); + test.AddOptionalOutputEdge(); + test.AddOutput("present_lengths", {2, 2}, state_lengths); + test.AddOptionalOutputEdge(); + test.AddOutput("status", {2}, {invalid == 0 ? 0 : 1, 2}); + RunOnProvider(test, std::move(provider)); + } } TEST(PackedSparseAttentionIndexerTest, QsaPrefillThenDecodeIndependentState) { @@ -1089,6 +1213,63 @@ TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalTopK) { RunQsaPackedTest(1.0e-5f, std::move(problem)); } +TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalPackedVerification) { + QsaPackedProblem problem; + problem.batch_size = 3; + problem.cumulative_sequence_lengths = {0, 0, 7, 8}; + problem.past_sequence_lengths = {0, 4101, 16388}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.compress_ratio = 4; + problem.token_budget = 2048; + problem.state_capacity = 65536; + problem = MakeQsaPackedProblem(std::move(problem)); + std::fill(problem.query.begin(), problem.query.end(), 1.0f); + std::fill(problem.query_norm_weight.begin(), problem.query_norm_weight.end(), 1.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + for (int batch = 0; batch < problem.batch_size; ++batch) { + for (int block = 0; block < problem.state_capacity; ++block) { + const float value = static_cast(block + 1) / problem.state_capacity; + const size_t offset = (static_cast(batch) * problem.state_capacity + block) * problem.head_size; + std::fill_n(problem.past_key_state.begin() + offset, problem.head_size, value); + } + } + RunQsaPackedTest(1.0e-5f, std::move(problem), ProviderKind::Cuda); +} + +TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalPairedVerification) { + for (int tokens : {5, 8}) { + QsaPackedProblem problem; + problem.batch_size = 1; + problem.cumulative_sequence_lengths = {0, tokens}; + problem.past_sequence_lengths = {4098}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.compress_ratio = 4; + problem.token_budget = 2048; + problem.state_capacity = 4099; + problem = MakeQsaPackedProblem(std::move(problem)); + std::fill(problem.query.begin(), problem.query.end(), 1.0f); + std::fill(problem.key.begin(), problem.key.end(), 1.0f); + std::fill(problem.query_norm_weight.begin(), problem.query_norm_weight.end(), 1.0f); + std::fill(problem.key_norm_weight.begin(), problem.key_norm_weight.end(), 1.0f); + std::fill(problem.past_kv_buffer.begin(), problem.past_kv_buffer.end(), 1.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + for (int block = 0; block < problem.state_capacity; ++block) { + const float value = static_cast(block + 1) / problem.state_capacity; + std::fill_n(problem.past_key_state.begin() + static_cast(block) * problem.head_size, + problem.head_size, value); + } + RunQsaPackedTest(1.0e-5f, problem, ProviderKind::Cuda); + RunQsaPackedTest(2.0e-3f, problem, ProviderKind::Cuda); + RunQsaPackedTest(2.0e-2f, std::move(problem), ProviderKind::Cuda); + } +} + TEST(PackedSparseAttentionIndexerTest, QsaRaggedPrefill) { QsaPackedProblem problem; problem.batch_size = 1; @@ -1475,7 +1656,8 @@ CsaPackedProblem MakeCsaPackedProblem(CsaPackedProblem problem = {}) { template void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, ProviderKind provider_kind = ProviderKind::Cuda, CsaPackedResult* actual = nullptr, - bool reuse = false, bool reuse_failure = false, const ConfigOptions* webgpu_options = nullptr) { + const ConfigOptions* webgpu_options = nullptr, bool reuse = false, bool reuse_failure = false, + int32_t merged_status = 0, int32_t merged_count = 1) { auto provider = CreateProvider(provider_kind, webgpu_options); if (provider == nullptr) { GTEST_SKIP() << (provider_kind == ProviderKind::Cuda ? "CUDA" : "WebGPU") @@ -1557,8 +1739,8 @@ void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, test.AddOptionalInputEdge(); test.AddInput("mode", {1}, {1}); test.AddInput("merged", {batch_size, output_capacity}, merged); - test.AddInput("merged_counts", {batch_size}, std::vector(static_cast(batch_size), 1)); - test.AddInput("merged_status", {batch_size}, std::vector(static_cast(batch_size), 0)); + test.AddInput("merged_counts", {batch_size}, std::vector(static_cast(batch_size), merged_count)); + test.AddInput("merged_status", {batch_size}, std::vector(static_cast(batch_size), merged_status)); } test.AddOutput("selected_indices", {total_tokens, output_capacity}, expected.selected_indices); test.AddOutput("selected_counts", {total_tokens}, expected.selected_counts); @@ -1571,7 +1753,8 @@ void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, test.AddOutput("present_state_lengths", {batch_size, 2}, expected.present_state_lengths); if (reuse) { test.AddOptionalOutputEdge(); - test.AddOutput("status", {batch_size}, std::vector(static_cast(batch_size), reuse_failure ? 1 : 0)); + test.AddOutput("status", {batch_size}, std::vector(static_cast(batch_size), merged_status != 0 ? merged_status : reuse_failure ? 1 + : 0)); } RunOnProvider(test, std::move(provider)); if (actual != nullptr) { @@ -1594,8 +1777,8 @@ TEST(PackedSparseAttentionIndexerTest, IndexShareCsaAdvanceWithoutScoring) { problem.past_sequence_lengths = {3, 1}; problem.past_state_lengths = {1, 3, 1, 1}; problem = MakeCsaPackedProblem(std::move(problem)); - RunCsaPackedTest(problem, 1.0e-5f, ProviderKind::Cuda, nullptr, true); - RunCsaPackedTest(problem, 4.0e-3f, ProviderKind::Cuda, nullptr, true); + RunCsaPackedTest(problem, 1.0e-5f, ProviderKind::Cuda, nullptr, nullptr, true); + RunCsaPackedTest(problem, 4.0e-3f, ProviderKind::Cuda, nullptr, nullptr, true); } TEST(PackedSparseAttentionIndexerTest, IndexShareCsaRejectsOffsetsAndStateOverflow) { @@ -1603,10 +1786,21 @@ TEST(PackedSparseAttentionIndexerTest, IndexShareCsaRejectsOffsetsAndStateOverfl problem.cumulative_sequence_lengths = {0, 0, 2}; problem.past_sequence_lengths = {3, 3}; problem.past_state_lengths = {1, 3, 1, 3}; - RunCsaPackedTest(MakeCsaPackedProblem(problem), 1.0e-5f, ProviderKind::Cuda, nullptr, true, true); + RunCsaPackedTest(MakeCsaPackedProblem(problem), 1.0e-5f, ProviderKind::Cuda, nullptr, nullptr, true, true); problem.cumulative_sequence_lengths = {0, 1, 2}; problem.state_capacity = 1; - RunCsaPackedTest(MakeCsaPackedProblem(problem), 1.0e-5f, ProviderKind::Cuda, nullptr, true, true); + RunCsaPackedTest(MakeCsaPackedProblem(problem), 1.0e-5f, ProviderKind::Cuda, nullptr, nullptr, true, true); +} + +TEST(PackedSparseAttentionIndexerTest, IndexShareCsaRejectsInvalidSelectionsBeforeAdvancing) { + CsaPackedProblem problem; + problem.cumulative_sequence_lengths = {0, 1, 2}; + problem.past_sequence_lengths = {3, 1}; + problem.past_state_lengths = {1, 3, 1, 1}; + problem = MakeCsaPackedProblem(std::move(problem)); + RunCsaPackedTest(problem, 1.0e-5f, ProviderKind::Cuda, nullptr, nullptr, true, true, 2, 1); + RunCsaPackedTest(problem, 1.0e-5f, ProviderKind::Cuda, nullptr, nullptr, true, true, 0, -1); + RunCsaPackedTest(problem, 1.0e-5f, ProviderKind::Cuda, nullptr, nullptr, true, true, 0, INT32_MAX); } TEST(PackedSparseAttentionIndexerTest, CsaFloat16) { RunCsaPackedTest(MakeCsaPackedProblem(), 4.0e-3f); } diff --git a/onnxruntime/test/contrib_ops/paged_attention_op_test.cc b/onnxruntime/test/contrib_ops/paged_attention_op_test.cc index bfe64593c2ad1..2a698262a90ca 100644 --- a/onnxruntime/test/contrib_ops/paged_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/paged_attention_op_test.cc @@ -2530,7 +2530,7 @@ TEST(PagedAttention, EndToEnd_Prefill_MultiBatch_Varlen_Fused) { // Fused paged prefill across a block boundary. block_size == max_k_step (32 // for fp16 head_size <= 128) is the tightest configuration -// ShouldRunFusedPagedPrefill accepts. token_count = 64 spans two paged +// CanApplyFlashAttentionPagedPrefill accepts. token_count = 64 spans two paged // blocks, and block_table = {3, 1} makes the two logical pages map to // non-adjacent physical blocks — any indexing bug that assumes physical // adjacency of consecutive logical pages will fail here (see the alignment @@ -2572,12 +2572,12 @@ TEST(PagedAttention, EndToEnd_Decode_BlockBoundaryCrossing) { RunEndToEndCaseOnAvailableProviders(c); } -// Force the ShouldRunFusedPagedPrefill *reject* branch. With fp16 head_size +// Force the CanApplyFlashAttentionPagedPrefill *reject* branch. With fp16 head_size // 128, max_k_step is 32, so block_size 16 fails the block_size >= max_k_step // guard. PagedAttention must then take the gather-then-flash cascade (the // merged #31611 code path). Verifies output parity against the reference in // the fallback path, locking down the reject-boundary so a future refactor -// of ShouldRunFusedPagedPrefill can't silently change it. +// of CanApplyFlashAttentionPagedPrefill can't silently change it. TEST(PagedAttention, EndToEnd_Prefill_ForcedFallback_BlockSizeBelowMaxKStep) { EndToEndCase c{}; c.batch_size = 1; diff --git a/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc b/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc index 9186acf6e348a..5a5561e4eb460 100644 --- a/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc +++ b/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc @@ -4,6 +4,8 @@ #include "test/util/include/default_providers.h" #include +#include +#include namespace onnxruntime { namespace test { @@ -324,5 +326,172 @@ TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_Int8_Saturate) { run_test(false); } +template +void TestQLinearSoftmaxEmptyInput() { + for (int64_t opset : {12, 13}) { + for (const TensorShapeVector& dims : {TensorShapeVector{0, 2, 37}, + TensorShapeVector{2, 0, 37}, + TensorShapeVector{2, 37, 0}}) { + for (bool add_shape : {false, true}) { + for (bool fixed_scales : {false, true}) { + SCOPED_TRACE(testing::Message() << "opset=" << opset << " dims=" << testing::PrintToString(dims) + << " add_shape=" << add_shape << " fixed_scales=" << fixed_scales); + OpTester test("QLinearSoftmax", 1, onnxruntime::kMSDomain); + test.ConfigEp(DefaultCpuExecutionProvider()); + test.AddAttribute("axis", -1); + test.AddAttribute("opset", opset); + test.AddShapeToTensorData(add_shape); + test.AddInput("X", dims, {}); + test.AddInput("X_scale", {}, {0.105f}, fixed_scales); + test.AddInput("X_zero_point", {}, {0}, fixed_scales); + test.AddInput("Y_scale", {}, {1.0f / 256.0f}, fixed_scales); + test.AddInput("Y_zero_point", {}, {0}, fixed_scales); + test.AddOutput("Y", dims, {}); + test.RunWithConfig(); + } + } + } + } +} + +TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_UInt8_EmptyInput) { + TestQLinearSoftmaxEmptyInput(); +} + +TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_Int8_EmptyInput) { + TestQLinearSoftmaxEmptyInput(); +} + +template +void TestQLinearSoftmaxSmallOutputScale() { + const T high = std::numeric_limits::max(); + const T low = std::numeric_limits::lowest(); + for (int64_t opset : {12, 13}) { + for (bool nonlast_axis : {false, true}) { + const std::vector dims = nonlast_axis ? std::vector{2, 1} : std::vector{1, 2}; + for (bool add_shape : {false, true}) { + for (bool fixed_scales : {false, true}) { + for (float output_scale : {std::ldexp(1.0f, -25), std::ldexp(1.0f, -40), + std::numeric_limits::min(), std::numeric_limits::denorm_min()}) { + volatile float scale = output_scale; + volatile double promoted_scale = scale; + if (promoted_scale == 0.0) { + continue; // Some platforms flush subnormal inputs to zero. + } + for (T zero_point : {low, static_cast(23)}) { + for (bool equal_inputs : {false, true}) { + SCOPED_TRACE(testing::Message() << "opset=" << opset << " nonlast_axis=" << nonlast_axis + << " add_shape=" << add_shape << " fixed_scales=" << fixed_scales + << " output_scale=" << output_scale << " zero_point=" << int(zero_point) + << " equal_inputs=" << equal_inputs); + OpTester test("QLinearSoftmax", 1, onnxruntime::kMSDomain); + test.ConfigEp(DefaultCpuExecutionProvider()); + test.AddAttribute("axis", nonlast_axis ? 0 : 1); + test.AddAttribute("opset", opset); + test.AddShapeToTensorData(add_shape); + test.AddInput("X", dims, equal_inputs ? std::vector{0, 0} : std::vector{low, high}); + test.AddInput("X_scale", {}, {1.0f}, fixed_scales); + test.AddInput("X_zero_point", {}, {0}, fixed_scales); + test.AddInput("Y_scale", {}, {output_scale}, fixed_scales); + test.AddInput("Y_zero_point", {}, {zero_point}, fixed_scales); + // Equal logits saturate both outputs. A difference of 255 leaves the smaller probability + // below half of even the smallest positive float scale, so it quantizes to the zero point. + test.AddOutput("Y", dims, {equal_inputs ? high : zero_point, high}); + test.RunWithConfig(); + } + } + } + } + } + } + } +} + +TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_UInt8_SmallOutputScale) { + TestQLinearSoftmaxSmallOutputScale(); +} + +TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_Int8_SmallOutputScale) { + TestQLinearSoftmaxSmallOutputScale(); +} + +template +void TestQLinearSoftmaxSingleton() { + for (int64_t opset : {12, 13}) { + for (bool add_shape : {false, true}) { + for (bool fixed_scales : {false, true}) { + for (T zero_point : {std::numeric_limits::lowest(), static_cast(23)}) { + SCOPED_TRACE(testing::Message() << "opset=" << opset << " add_shape=" << add_shape + << " fixed_scales=" << fixed_scales << " zero_point=" << int(zero_point)); + OpTester test("QLinearSoftmax", 1, onnxruntime::kMSDomain); + test.ConfigEp(DefaultCpuExecutionProvider()); + test.AddAttribute("axis", -1); + test.AddAttribute("opset", opset); + test.AddShapeToTensorData(add_shape); + test.AddInput("X", {2, 1}, {std::numeric_limits::lowest(), std::numeric_limits::max()}); + test.AddInput("X_scale", {}, {0.105f}, fixed_scales); + test.AddInput("X_zero_point", {}, {0}, fixed_scales); + test.AddInput("Y_scale", {}, {1.0f / 256.0f}, fixed_scales); + test.AddInput("Y_zero_point", {}, {zero_point}, fixed_scales); + test.AddOutput("Y", {2, 1}, {std::numeric_limits::max(), std::numeric_limits::max()}); + test.RunWithConfig(); + } + } + } + } +} + +TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_UInt8_Singleton) { + TestQLinearSoftmaxSingleton(); +} + +TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_Int8_Singleton) { + TestQLinearSoftmaxSingleton(); +} + +template +void TestQLinearSoftmaxRounding() { + for (int64_t opset : {12, 13}) { + for (bool fixed_scales : {false, true}) { + for (const auto& values : {std::pair{1.0f / 3.0f, 2}, std::pair{1.0f / 5.0f, 2}, + std::pair{1.0f / 11.0f, 6}, std::pair{1.0f / 9.0f, 4}, + std::pair{1.0f / 59.0f, 30}, std::pair{1.0f / 249.0f, 124}}) { + SCOPED_TRACE(testing::Message() << "opset=" << opset << " fixed_scales=" << fixed_scales + << " output_scale=" << values.first); + OpTester test("QLinearSoftmax", 1, onnxruntime::kMSDomain); + test.ConfigEp(DefaultCpuExecutionProvider()); + test.AddAttribute("axis", -1); + test.AddAttribute("opset", opset); + test.AddInput("X", {1, 2}, {0, 0}); + test.AddInput("X_scale", {}, {0.105f}, fixed_scales); + test.AddInput("X_zero_point", {}, {0}, fixed_scales); + test.AddInput("Y_scale", {}, {values.first}, fixed_scales); + test.AddInput("Y_zero_point", {}, {0}, fixed_scales); + test.AddOutput("Y", {1, 2}, {static_cast(values.second), static_cast(values.second)}); + test.RunWithConfig(); + } + OpTester test("QLinearSoftmax", 1, onnxruntime::kMSDomain); + test.ConfigEp(DefaultCpuExecutionProvider()); + test.AddAttribute("axis", -1); + test.AddAttribute("opset", opset); + test.AddInput("X", {1, 2}, {0, 16}); + test.AddInput("X_scale", {}, {1.0f}, fixed_scales); + test.AddInput("X_zero_point", {}, {0}, fixed_scales); + test.AddInput("Y_scale", {}, {std::ldexp(1.0f, -25)}, fixed_scales); + test.AddInput("Y_zero_point", {}, {0}, fixed_scales); + test.AddOutput("Y", {1, 2}, {4, std::numeric_limits::max()}); + test.RunWithConfig(); + } + } +} + +TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_UInt8_Rounding) { + TestQLinearSoftmaxRounding(); +} + +TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_Int8_Rounding) { + TestQLinearSoftmaxRounding(); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc index 658c8695643e6..a6957e18bc194 100644 --- a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc +++ b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc @@ -2,6 +2,8 @@ // Licensed under the MIT License. #include +#include +#include #include #include "gtest/gtest.h" #include "contrib_ops/cpu/bert/rotary_embedding.h" @@ -21,6 +23,149 @@ enum class TensorType { }; } // anonymous namespace +class FloatRotaryChainTester : public OpTester { + public: + FloatRotaryChainTester(int type, int heads, int rotary_dim, bool interleaved) + : OpTester("RotaryEmbedding", 1, kMSDomain, false), + type_(type), + heads_(heads), + rotary_dim_(rotary_dim), + interleaved_(interleaved) {} + + void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, + std::vector>&) override { + auto value = [&](const std::string& name) { return &graph.GetOrCreateNodeArg(name, nullptr); }; + std::vector widened = inputs; + for (size_t index : {size_t{0}, size_t{2}, size_t{3}}) { + const auto name = "float_" + std::to_string(index); + widened[index] = value(name); + graph.AddNode(name, "Cast", "", {inputs[index]}, {widened[index]}) + .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + } + auto* result = value("float_result"); + auto& rotary = graph.AddNode("rotary", "RotaryEmbedding", "", widened, {result}, nullptr, kMSDomain); + rotary.AddAttribute("num_heads", static_cast(heads_)); + rotary.AddAttribute("rotary_embedding_dim", static_cast(rotary_dim_)); + rotary.AddAttribute("interleaved", static_cast(interleaved_)); + graph.AddNode("rounded", "Cast", "", {result}, {outputs[0]}) + .AddAttribute("to", static_cast(type_)); + graph.SetOutputs(outputs); + } + + private: + void CreateModelToTest(const ModelOptions& options, Model*& model) override { + model = &BuildModel({{kOnnxDomain, 17}}, options); + ORT_THROW_IF_ERROR(model->MainGraph().Resolve()); + } + + int type_, heads_, rotary_dim_; + bool interleaved_; +}; + +template +void TestFloatRotaryArithmetic(int type) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA SM80 is required"; + } + for (int sequence : {1, 2, 9}) { + for (int width : {64, 256}) { + for (bool transposed : {false, true}) { + for (bool interleaved : {false, true}) { + SCOPED_TRACE("sequence=" + std::to_string(sequence) + " width=" + std::to_string(width) + + " transposed=" + std::to_string(transposed) + " interleaved=" + std::to_string(interleaved)); + constexpr int batch = 2, heads = 3; + const int rotary_dim = 64; + const int cache_rows = batch * sequence + 3; + std::vector shape = transposed ? std::vector{batch, heads, sequence, width} + : std::vector{batch, sequence, heads * width}; + std::vector input(batch * sequence * heads * width), cosine(cache_rows * rotary_dim / 2), + sine(cosine.size()); + const bool scalar_position = !transposed; + std::vector positions(scalar_position ? 1 : batch * sequence); + for (size_t i = 0; i < input.size(); ++i) input[i] = T(std::sin(i * 0.71f) * (1 + i % 17)); + for (size_t i = 0; i < cosine.size(); ++i) { + cosine[i] = T(std::cos(i * 0.13f)); + sine[i] = T(std::sin(i * 0.13f)); + } + for (size_t i = 0; i < positions.size(); ++i) positions[i] = scalar_position ? 1 : (i * 7) % cache_rows; + FloatRotaryChainTester reference(type, heads, rotary_dim, interleaved); + OpTester fused("RotaryEmbedding", 2, kMSDomain, false); + fused.AddAttribute("num_heads", heads); + fused.AddAttribute("rotary_embedding_dim", rotary_dim); + fused.AddAttribute("interleaved", interleaved); + fused.AddAttribute("use_fp32_arithmetic", 1); + std::vector expected; + for (OpTester* tester : {static_cast(&reference), &fused}) { + tester->AddInput("input", shape, input); + tester->AddInput("positions", scalar_position ? std::vector{1} : std::vector{batch, sequence}, positions); + tester->AddInput("cosine", {cache_rows, rotary_dim / 2}, cosine); + tester->AddInput("sine", {cache_rows, rotary_dim / 2}, sine); + tester->AddOutput("output", shape, std::vector(input.size())); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester->Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + auto fetches = tester->GetFetches(); + ASSERT_EQ(fetches.size(), 1); + if (expected.empty()) { + expected = std::move(fetches); + } else { + const auto& before = expected[0].Get(); + const auto& after = fetches[0].Get(); + ASSERT_EQ(before.SizeInBytes(), after.SizeInBytes()); + ASSERT_EQ(std::memcmp(before.DataRaw(), after.DataRaw(), before.SizeInBytes()), 0); + } + } + } + } + } + } +} + +TEST(RotaryEmbeddingTest, FloatArithmeticMatchesWidenRotateRoundFloat16) { + TestFloatRotaryArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); +} + +TEST(RotaryEmbeddingTest, FloatArithmeticMatchesWidenRotateRoundBFloat16) { + TestFloatRotaryArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); +} + +TEST(RotaryEmbeddingTest, Version2LegacyDefaultAndInvalidMode) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA SM80 is required"; + } + for (int mode : {0, 2}) { + OpTester test("RotaryEmbedding", 2, kMSDomain); + if (mode != 0) test.AddAttribute("use_fp32_arithmetic", mode); + test.AddAttribute("num_heads", 1); + test.AddInput("input", {1, 1, 4}, ToFloat16({1, 2, 3, 4})); + test.AddInput("positions", {1}, {0}); + test.AddInput("cos", {1, 2}, ToFloat16({1, 1})); + test.AddInput("sin", {1, 2}, ToFloat16({0, 0})); + test.AddOutput("output", {1, 1, 4}, ToFloat16({1, 2, 3, 4})); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, + mode == 0 ? "" : "use_fp32_arithmetic must be 0 or 1", {}, nullptr, &providers); + } +} + +TEST(RotaryEmbeddingTest, Version2DoesNotFallBackToLegacyCpuArithmetic) { + OpTester test("RotaryEmbedding", 2, kMSDomain); + test.AddAttribute("use_fp32_arithmetic", 1); + test.AddAttribute("num_heads", 1); + test.AddInput("input", {1, 1, 4}, {1, 2, 3, 4}); + test.AddInput("positions", {1}, {0}); + test.AddInput("cos", {1, 2}, {1, 1}); + test.AddInput("sin", {1, 2}, {0, 0}); + test.AddOutput("output", {1, 1, 4}, {1, 2, 3, 4}); + std::vector> providers; + providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "Could not find an implementation for RotaryEmbedding(2)", + {}, nullptr, &providers); +} + static void RunTest( const std::vector& input_data, const std::vector& position_ids, diff --git a/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc index 3cf1bb9e2ae93..0561c46571879 100644 --- a/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc @@ -386,7 +386,9 @@ void QsaReference(const QsaProblem& problem, std::vector& selected, std } else { pooled[static_cast(d)] += problem.key[(static_cast(b) * problem.sequence_length + - position - problem.past_sequence_length) * head_size + d]; + position - problem.past_sequence_length) * + head_size + + d]; } } } @@ -492,7 +494,6 @@ void QsaReference(const QsaProblem& problem, std::vector& selected, std } } } - } struct CsaProblem { diff --git a/onnxruntime/test/framework/ep_plugin_provider_test.cc b/onnxruntime/test/framework/ep_plugin_provider_test.cc index 8b29a518e77f2..792b8b7b5ff07 100644 --- a/onnxruntime/test/framework/ep_plugin_provider_test.cc +++ b/onnxruntime/test/framework/ep_plugin_provider_test.cc @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -31,6 +32,7 @@ #include "core/session/plugin_ep/ep_allocator_utils.h" #include "test/util/include/api_asserts.h" #include "test/util/include/asserts.h" +#include "test/util/include/file_util.h" #include "test/util/include/test_environment.h" namespace onnxruntime::test { @@ -1945,6 +1947,137 @@ static IResourceAccountant* CreateNoThresholdAccountant(std::optionalend() ? it->second.get() : nullptr; } +#if !defined(DISABLE_CONTRIB_OPS) +TEST(PluginExecutionProviderTest, GetCapability_CudaGqaRequiresProvenEstimate) { + enum class AccountingMode { + None, + AdHoc, + Profile, + }; + const auto run_test = [](bool fuse_nodes, AccountingMode accounting_mode, bool previously_assigned, + const char* ep_name, const char* envelope) { + const bool use_accountant = accounting_mode != AccountingMode::None; + const bool use_profile = accounting_mode == AccountingMode::Profile; + SCOPED_TRACE(::testing::Message() << "fuse_nodes=" << fuse_nodes << ", use_accountant=" << use_accountant + << ", use_profile=" << use_profile + << ", previously_assigned=" << previously_assigned << ", ep_name=" << ep_name + << ", envelope=" << envelope); + auto [ep, ort_ep] = test_plugin_ep::MakeTestOrtEp({}, [ep_name](test_plugin_ep::TestOrtEp& test_ep) { + if (std::strcmp(ep_name, kCudaExecutionProvider) == 0) { + test_ep.GetName = [](const OrtEp*) noexcept { return kCudaExecutionProvider; }; + } + }); + ort_ep->GetCapability = fuse_nodes ? GetCapabilityTakeAllNodesOneGroup + : [](OrtEp* this_ptr, const OrtGraph* graph, + OrtEpGraphSupportInfo* graph_support_info) noexcept -> OrtStatus* { + auto* test_ep = static_cast(this_ptr); + for (const auto& node : Ort::ConstGraph{graph}.GetNodes()) { + if (auto* status = test_ep->ep_api->EpGraphSupportInfo_AddSingleNode(graph_support_info, node)) { + return status; + } + } + return nullptr; + }; + + const auto& logger = DefaultLoggingManager().DefaultLogger(); + Model model("gqa_plugin_capacity", false, logger); + Graph& graph = model.MainGraph(); + const auto make_arg = [&graph](const char* name, int32_t element_type, + std::initializer_list dimensions) { + ONNX_NAMESPACE::TypeProto type; + auto* tensor = type.mutable_tensor_type(); + tensor->set_elem_type(element_type); + auto* shape = tensor->mutable_shape(); + for (int64_t dimension : dimensions) { + shape->add_dim()->set_dim_value(dimension); + } + return &graph.GetOrCreateNodeArg(name, &type); + }; + constexpr int32_t kFloat16 = ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; + constexpr int32_t kInt32 = ONNX_NAMESPACE::TensorProto_DataType_INT32; + auto* query = make_arg("query", kFloat16, {1, 1, 64}); + auto* key = make_arg("key", kFloat16, {1, 1, 32}); + auto* value = make_arg("value", kFloat16, {1, 1, 32}); + auto* past_key = make_arg("past_key", kFloat16, {1, 1, 128, 32}); + auto* past_value = make_arg("past_value", kFloat16, {1, 1, 128, 32}); + auto* seqlens = make_arg("seqlens_k", kInt32, {1}); + auto* total_length = make_arg("total_sequence_length", kInt32, {}); + auto* output = make_arg("output", kFloat16, {1, 1, 64}); + auto* present_key = make_arg("present_key", kFloat16, {1, 1, 128, 32}); + auto* present_value = make_arg("present_value", kFloat16, {1, 1, 128, 32}); + auto& gqa = graph.AddNode("gqa", "GroupQueryAttention", "", + {query, key, value, past_key, past_value, seqlens, total_length}, + {output, present_key, present_value}, nullptr, kMSDomain); + gqa.AddAttribute("num_heads", int64_t{2}); + gqa.AddAttribute("kv_num_heads", int64_t{1}); + auto* identity_output = make_arg("identity_output", kFloat16, {1, 1, 64}); + graph.AddNode("identity", "Identity", "", {output}, {identity_output}); + ASSERT_STATUS_OK(graph.Resolve()); + if (previously_assigned) { + gqa.SetExecutionProviderType(ep_name); + } + + PathString stats_path = ORT_TSTR("gqa_plugin_stats_XXXXXX"); + if (use_profile) { + FILE* stats_file = nullptr; + ASSERT_NO_FATAL_FAILURE(CreateTestFile(stats_file, stats_path)); + ASSERT_NE(stats_file, nullptr); + ASSERT_EQ(std::fclose(stats_file), 0); + } + ScopedFileDeleter stats_cleanup(use_profile ? stats_path : PathString{}); + if (use_profile) { + std::ofstream stats{std::filesystem::path{stats_path}}; + ASSERT_TRUE(stats); + stats << "#name,input_sizes,initializers_sizes,total_dynamic_sizes,total_temp_allocations\n" + << IResourceAccountant::MakeUniqueNodeName(gqa) << ",0,0,0,0\n"; + stats.close(); + ASSERT_TRUE(stats); + } + + ConfigOptions config; + const std::string settings = (std::strcmp(envelope, "0") == 0 ? "," : "1024,") + + (use_profile ? ToUTF8String(stats_path) : std::string{}); + ASSERT_STATUS_OK(config.AddConfigEntry(kOrtSessionOptionsResourceCudaPartitioningSettings, settings.c_str())); + ASSERT_STATUS_OK(config.AddConfigEntry(kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength, envelope)); + std::optional acc_map; + ASSERT_STATUS_OK(CreateAccountants(config, /*model_path=*/{}, acc_map)); + auto* accountant = acc_map->at(kCudaExecutionProvider).get(); + if (use_profile) { + EXPECT_EQ(std::get(accountant->ComputeResourceCount(gqa)), 0u); + EXPECT_EQ(accountant->GetPendingWorkspaceEstimateSelection(gqa.Index()).source, + WorkspaceEstimateSource::kProfile); + accountant->ResetForNewPass(); + } + GraphViewer graph_viewer(graph); + auto capabilities = ep->GetCapability(graph_viewer, test_plugin_ep::MockKernelLookup(), + GraphOptimizerRegistry(nullptr, nullptr, &logger), + use_accountant ? accountant : nullptr); + + const bool decline_gqa = use_accountant && !previously_assigned && + std::strcmp(ep_name, kCudaExecutionProvider) == 0; + ASSERT_EQ(capabilities.size(), fuse_nodes ? (decline_gqa ? 0u : 1u) : (decline_gqa ? 1u : 2u)); + if (decline_gqa && !fuse_nodes) { + ASSERT_EQ(capabilities[0]->sub_graph->nodes.size(), 1u); + EXPECT_EQ(graph.GetNode(capabilities[0]->sub_graph->nodes[0])->OpType(), "Identity"); + EXPECT_TRUE(capabilities[0]->sub_graph->IsAccountingEnabled()); + } + EXPECT_FALSE(accountant->IsStopIssued()); + }; + + for (bool fuse_nodes : {false, true}) { + for (AccountingMode accounting_mode : {AccountingMode::None, AccountingMode::AdHoc, AccountingMode::Profile}) { + for (bool previously_assigned : {false, true}) { + for (const char* ep_name : {kCudaExecutionProvider, "TestOrtEp"}) { + for (const char* envelope : {"0", "512"}) { + run_test(fuse_nodes, accounting_mode, previously_assigned, ep_name, envelope); + } + } + } + } + } +} +#endif // !defined(DISABLE_CONTRIB_OPS) + // Helper: call GetCapability on a mock EP with a no-threshold accountant, returning the accountant for inspection. static IResourceAccountant* CallGetCapabilityWithAccountant( IExecutionProvider& ep, diff --git a/onnxruntime/test/framework/inference_session_test.cc b/onnxruntime/test/framework/inference_session_test.cc index f7fab9b320c4f..dd4a9c8965504 100644 --- a/onnxruntime/test/framework/inference_session_test.cc +++ b/onnxruntime/test/framework/inference_session_test.cc @@ -5245,7 +5245,7 @@ TEST(InferenceSessionTests, CompileApiExternalBufferKeepsSubgraphInitializersEmb if (buffer_destination && has_main_initializer) continue; #endif #if defined(__EMSCRIPTEN__) - // The WASM test module does not provide Module.MountedFiles for file-backed external data. + // File-backed subgraph tensors require a Module.MountedFiles mapping in WASM. if (!buffer_destination) continue; #endif { @@ -5295,7 +5295,7 @@ TEST(InferenceSessionTests, CompileApiExternalBufferKeepsSubgraphInitializersEmb } } -// The WASM test module cannot load host filesystem external data without Module.MountedFiles. +// Subgraph and attribute external tensors in this test load from the host filesystem. #if !defined(DISABLE_EXTERNAL_INITIALIZERS) && !defined(__EMSCRIPTEN__) TEST(InferenceSessionTests, ExternalInitializerInjectionIsMainGraphOnly) { const std::filesystem::path model_dir = ORT_TSTR("external_buffers_attribute_model_dir"); diff --git a/onnxruntime/test/framework/insert_cast_transformer_test.cc b/onnxruntime/test/framework/insert_cast_transformer_test.cc index ae9fccf8fedd3..21bd7ba8ff622 100644 --- a/onnxruntime/test/framework/insert_cast_transformer_test.cc +++ b/onnxruntime/test/framework/insert_cast_transformer_test.cc @@ -20,6 +20,36 @@ namespace test { #define MODEL_FOLDER ORT_TSTR("testdata/transform/") typedef std::vector ArgMap; + +TEST(TransformerTest, Fp16OnlyContribOpIsNotCastToUnsupportedFloat32) { +#if defined(DISABLE_CONTRIB_OPS) || defined(DISABLE_FLOAT8_TYPES) + GTEST_SKIP() << "Block-quantized FP8 MatMul is unavailable in this build."; +#else + auto model = std::make_shared( + "test", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + std::unordered_map{{onnxruntime::kOnnxDomain, 22}, {onnxruntime::kMSDomain, 1}}, + std::vector(), DefaultLoggingManager().DefaultLogger()); + onnxruntime::Graph& graph = model->MainGraph(); + + TypeProto fp16, fp8, float32; + fp16.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT16); + fp8.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT8E4M3FN); + float32.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); + onnxruntime::NodeArg a("A", &fp16), b("B", &fp8), scale("b_scale", &float32), y("Y", &fp16); + auto& node = graph.AddNode("fp8_matmul", "MatMulBlockQuantizedFp8Weight", "", + ArgMap{&a, &b, &scale}, ArgMap{&y}, nullptr, kMSDomain); + graph.SetOutputs({node.OutputDefs()[0]}); + + ASSERT_STATUS_OK(graph.Resolve()); + InsertCastTransformer transformer("Test", DefaultCpuExecutionProvider()->GetKernelRegistry().get()); + bool modified = false; + ASSERT_STATUS_OK(transformer.Apply(graph, modified, DefaultLoggingManager().DefaultLogger())); + ASSERT_STATUS_OK(graph.Resolve()); + EXPECT_FALSE(modified); + EXPECT_EQ(graph.NumberOfNodes(), 1); +#endif +} + TEST(TransformerTest, InsertCastGPUTest) { auto model = std::make_shared("test", false, DefaultLoggingManager().DefaultLogger()); onnxruntime::Graph& graph = model->MainGraph(); diff --git a/onnxruntime/test/framework/resource_accountant_test.cc b/onnxruntime/test/framework/resource_accountant_test.cc index a693771a1419c..7b4ebec73571d 100644 --- a/onnxruntime/test/framework/resource_accountant_test.cc +++ b/onnxruntime/test/framework/resource_accountant_test.cc @@ -722,9 +722,9 @@ TEST(RealAccountantTest, StatsPath_Level1EstimateUsesMaximumWorkspace) { EXPECT_EQ(graph_reservations.at(h.node_b->Index()).bytes, size_t{800}); } -// A stats file may have incomplete coverage. Preserve the historical behavior -// where unknown nodes have zero cost and do not enter ad-hoc bookkeeping. -TEST(RealAccountantTest, StatsPath_UnknownNodeHasZeroCost) { +// A stats file may have incomplete coverage. Preserve zero cost for an unknown +// node without an estimator, but do not discard an explicit Level-1 estimate. +TEST(RealAccountantTest, StatsPath_UnknownNodeHonorsLevel1Estimate) { SharedWeightGraph h; ASSERT_NO_FATAL_FAILURE(SharedWeightGraph::Create(h)); @@ -756,17 +756,27 @@ TEST(RealAccountantTest, StatsPath_UnknownNodeHasZeroCost) { auto known_cost = accountant->ComputeResourceCount(*h.node_a); EXPECT_EQ(std::get(known_cost), size_t{1000}); - auto unknown_cost = accountant->ComputeResourceCount( + const auto unknown_without_estimate = + accountant->ComputeResourceCount(*h.node_b); + EXPECT_EQ(std::get(unknown_without_estimate), size_t{0}); + auto pending_workspace = + accountant->GetPendingWorkspaceEstimateSelection(h.node_b->Index()); + EXPECT_EQ(pending_workspace.bytes, size_t{0}); + EXPECT_EQ(pending_workspace.source, WorkspaceEstimateSource::kNone); + + const auto unknown_cost = accountant->ComputeResourceCount( *h.node_b, Level1MemoryEstimate{ /*runtime_workspace_bytes=*/250, /*persistent_prepack_bytes=*/300, /*initialization_scratch_bytes=*/400}); - EXPECT_EQ(std::get(unknown_cost), size_t{0}); - const auto pending_workspace = + EXPECT_EQ(std::get(unknown_cost), size_t{2550}); + pending_workspace = accountant->GetPendingWorkspaceEstimateSelection(h.node_b->Index()); - EXPECT_EQ(pending_workspace.bytes, size_t{0}); - EXPECT_EQ(pending_workspace.source, WorkspaceEstimateSource::kNone); + EXPECT_EQ(pending_workspace.bytes, size_t{250}); + EXPECT_EQ(pending_workspace.source, WorkspaceEstimateSource::kEstimator); + EXPECT_EQ(pending_workspace.persistent_prepack_bytes, size_t{300}); + EXPECT_EQ(pending_workspace.initialization_scratch_bytes, size_t{400}); IndexedSubGraph sub_graph; sub_graph.nodes.push_back(h.node_a->Index()); @@ -776,9 +786,12 @@ TEST(RealAccountantTest, StatsPath_UnknownNodeHasZeroCost) { sub_graph.AppendNodeCost(unknown_cost); sub_graph.AccountForAllNodes(); - EXPECT_EQ(GetSizeT(accountant->GetConsumedAmount()), size_t{1000}); - EXPECT_EQ(accountant->GetCommittedWorkspaceEstimate(), size_t{400}); + EXPECT_EQ(GetSizeT(accountant->GetConsumedAmount()), size_t{3550}); + EXPECT_EQ(accountant->GetCommittedWorkspaceEstimate(), size_t{650}); + EXPECT_EQ(accountant->GetCommittedPersistentPrepackEstimate(), size_t{300}); + EXPECT_EQ(accountant->GetCommittedInitializationScratchEstimate(), size_t{400}); EXPECT_EQ(accountant->GetWorkspaceEstimateSourceCounts().profile, size_t{1}); + EXPECT_EQ(accountant->GetWorkspaceEstimateSourceCounts().estimator, size_t{1}); } // Factory with no limit and no stats file creates accountant with no threshold. @@ -833,7 +846,7 @@ TEST(RealAccountantTest, FactoryParsesGqaWorkspaceEnvelope) { }; for (const auto& test_case : { Case{nullptr, 0}, Case{"0", 0}, Case{"4096", 4096}, - Case{"9223372036854775807", std::numeric_limits::max()}}) { + Case{"2147483647", std::numeric_limits::max()}}) { SCOPED_TRACE(test_case.value != nullptr ? test_case.value : ""); ConfigOptions config; ASSERT_STATUS_OK(config.AddConfigEntry( @@ -856,7 +869,8 @@ TEST(RealAccountantTest, FactoryRejectsInvalidGqaWorkspaceEnvelope) { for (bool enable_partitioning : {false, true}) { SCOPED_TRACE(enable_partitioning); for (const char* value : {"-1", "", "malformed", "4096x", "+4096", - " 4096", "4096 ", "9223372036854775808"}) { + " 4096", "4096 ", "2147483648", + "9223372036854775807", "9223372036854775808"}) { SCOPED_TRACE(value); ConfigOptions config; if (enable_partitioning) { diff --git a/onnxruntime/test/framework/sparse_kernels_test.cc b/onnxruntime/test/framework/sparse_kernels_test.cc index 57b785df18494..b37ef94c0315c 100644 --- a/onnxruntime/test/framework/sparse_kernels_test.cc +++ b/onnxruntime/test/framework/sparse_kernels_test.cc @@ -2431,6 +2431,74 @@ TEST(SparseTensorConversionTests, SparseTensorProtoToDense_ValuesSizeMismatch_Ra EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("raw_data size")); } +TEST(SparseTensorConversionTests, SparseTensorProtoToDense_IndicesSizeMismatch) { + ONNX_NAMESPACE::SparseTensorProto sparse; + sparse.mutable_values()->set_name("test_tensor_idx_size_mismatch"); + sparse.add_dims(utils::kMaxEmbeddedInitializerSizeInBytes / sizeof(float) + 1); + + auto* val = sparse.mutable_values(); + val->add_dims(2); + val->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + val->add_float_data(1.0f); + val->add_float_data(2.0f); + + auto* ind = sparse.mutable_indices(); + ind->set_name("test_tensor_idx_size_mismatch_indices"); + ind->add_dims(2); + ind->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + ind->add_int64_data(0); + + ONNX_NAMESPACE::TensorProto dense; + auto status = utils::SparseTensorProtoToDenseTensorProto(sparse, {}, dense); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("data field count")); +} + +TEST(SparseTensorConversionTests, SparseTensorProtoToDense_IndicesSizeMismatch_RawData) { + ONNX_NAMESPACE::SparseTensorProto sparse; + sparse.mutable_values()->set_name("test_tensor_idx_size_mismatch_raw"); + sparse.add_dims(utils::kMaxEmbeddedInitializerSizeInBytes / sizeof(float) + 1); + + auto* val = sparse.mutable_values(); + val->add_dims(2); + val->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + val->add_float_data(1.0f); + val->add_float_data(2.0f); + + auto* ind = sparse.mutable_indices(); + ind->set_name("test_tensor_idx_size_mismatch_raw_indices"); + ind->add_dims(2); + ind->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + + int64_t raw_index = 0; + onnxruntime::utils::SetRawDataInTensorProto(*ind, &raw_index, sizeof(raw_index)); + + ONNX_NAMESPACE::TensorProto dense; + auto status = utils::SparseTensorProtoToDenseTensorProto(sparse, {}, dense); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("raw_data size")); +} + +TEST(SparseTensorConversionTests, SparseTensorProtoToDense_RejectsUnsupportedIndexTypeForEmptyTensor) { + ONNX_NAMESPACE::SparseTensorProto sparse; + sparse.mutable_values()->set_name("unsupported_index_type"); + sparse.add_dims(0); + + auto* values = sparse.mutable_values(); + values->add_dims(0); + values->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + + auto* indices = sparse.mutable_indices(); + indices->add_dims(0); + indices->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + + ONNX_NAMESPACE::TensorProto dense; + const auto status = utils::SparseTensorProtoToDenseTensorProto(sparse, {}, dense); + ASSERT_FALSE(status.IsOK()); + EXPECT_EQ(status.Code(), common::INVALID_GRAPH); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("indices. Should be one of the following types")); +} + // Tests for SparseTensorProtoToDenseTensorProto with negative indices (model-loading path) TEST(SparseTensorConversionTests, SparseTensorProtoToDense_NegativeIndex_Rank1) { // Dense size 4 diff --git a/onnxruntime/test/framework/tensorutils_test.cc b/onnxruntime/test/framework/tensorutils_test.cc index 6f941cd95d741..779c6cfff5366 100644 --- a/onnxruntime/test/framework/tensorutils_test.cc +++ b/onnxruntime/test/framework/tensorutils_test.cc @@ -61,6 +61,8 @@ constexpr bool TensorProtoElementSizesAreConstexpr() { sizeof(uint8_t), // FLOAT8E8M0 sizeof(uint8_t), // UINT2 sizeof(uint8_t), // INT2 + sizeof(uint8_t), // FLOAT6E2M3 + sizeof(uint8_t), // FLOAT6E3M2 }; for (size_t index = 0; index < expected_sizes.size(); ++index) { diff --git a/onnxruntime/test/framework/type_info_test.cc b/onnxruntime/test/framework/type_info_test.cc index 9368f5da889b0..7631b5eaaf58b 100644 --- a/onnxruntime/test/framework/type_info_test.cc +++ b/onnxruntime/test/framework/type_info_test.cc @@ -11,6 +11,7 @@ #include "core/framework/onnxruntime_map_type_info.h" #include "core/framework/onnxruntime_sequence_type_info.h" #include "core/framework/tensor_type_and_shape.h" +#include "core/framework/tensorprotoutils.h" #include "core/framework/onnxruntime_typeinfo.h" #include "core/session/onnxruntime_type_conversion.h" @@ -49,6 +50,8 @@ constexpr bool TensorElementTypeConversionIsConstexpr() { ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E8M0, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT2, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT2, + ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E2M3, + ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E3M2, }; for (size_t index = 0; index < expected_types.size(); ++index) { @@ -72,6 +75,13 @@ constexpr bool TensorElementTypeConversionIsConstexpr() { static_assert(TensorElementTypeConversionIsConstexpr()); +TEST(TypeInfoTests, CApiElementTypeFromProtoTypeFloat6) { + EXPECT_EQ(utils::CApiElementTypeFromProtoType(ONNX_NAMESPACE::TensorProto_DataType_FLOAT6E2M3), + ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E2M3); + EXPECT_EQ(utils::CApiElementTypeFromProtoType(ONNX_NAMESPACE::TensorProto_DataType_FLOAT6E3M2), + ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E3M2); +} + TEST(TypeInfoTests, TensorElementTypeConversions) { struct TypePair { ONNXTensorElementDataType api_type; @@ -105,6 +115,8 @@ TEST(TypeInfoTests, TensorElementTypeConversions) { {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E8M0, ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E8M0}, {ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT2, ONNX_NAMESPACE::TensorProto_DataType_UINT2}, {ONNX_TENSOR_ELEMENT_DATA_TYPE_INT2, ONNX_NAMESPACE::TensorProto_DataType_INT2}, + {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E2M3, ONNX_NAMESPACE::TensorProto_DataType_FLOAT6E2M3}, + {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT6E3M2, ONNX_NAMESPACE::TensorProto_DataType_FLOAT6E3M2}, }; static_assert(std::size(types) == ONNX_NAMESPACE::TensorProto_DataType_DataType_ARRAYSIZE); for (const auto& [api_type, proto_type] : types) { diff --git a/onnxruntime/test/mlas/unittest/test_eltwise.cpp b/onnxruntime/test/mlas/unittest/test_eltwise.cpp index 136d3a9a756b4..d3e530fb5ca74 100644 --- a/onnxruntime/test/mlas/unittest/test_eltwise.cpp +++ b/onnxruntime/test/mlas/unittest/test_eltwise.cpp @@ -48,6 +48,9 @@ class MlasEltwiseAddTest : public MlasTestBase { #if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64) void TestFp16(size_t N, float MinimumValue, float MaximumValue, const std::optional& ScalarValue = std::nullopt) { + if (!MlasFp16AccelerationSupported()) { + return; + } MLAS_FP16* InputLeft = BufferInputLeftFp16.GetBuffer(N); MLAS_FP16* InputRight = BufferInputRightFp16.GetBuffer(N); MLAS_FP16* Output = BufferOutputFp16.GetBuffer(N); diff --git a/onnxruntime/test/mlas/unittest/test_fp16_activation.cpp b/onnxruntime/test/mlas/unittest/test_fp16_activation.cpp index 969997d2b84ec..3460a502c3d7f 100644 --- a/onnxruntime/test/mlas/unittest/test_fp16_activation.cpp +++ b/onnxruntime/test/mlas/unittest/test_fp16_activation.cpp @@ -149,7 +149,10 @@ class MlasFp16ActivationTest : public MlasTestBase { } }; -static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_execute) { +static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_execute) -> size_t { + if (!MlasFp16AccelerationSupported()) { + return 0; + } return is_short_execute ? MlasDirectShortExecuteTests::RegisterShortExecute() : 0; }); diff --git a/onnxruntime/test/mlas/unittest/test_gelu_erf_neon_fp16.cpp b/onnxruntime/test/mlas/unittest/test_gelu_erf_neon_fp16.cpp index bd6ecb69339c3..04a5cf7e41d2e 100644 --- a/onnxruntime/test/mlas/unittest/test_gelu_erf_neon_fp16.cpp +++ b/onnxruntime/test/mlas/unittest/test_gelu_erf_neon_fp16.cpp @@ -270,7 +270,7 @@ class MlasNeonGeluFP16Test : public MlasTestBase { static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_execute) { size_t count = 0; - if (is_short_execute) { + if (is_short_execute && MlasFp16AccelerationSupported()) { count += MlasDirectShortExecuteTests::RegisterShortExecute(); count += MlasDirectShortExecuteTests::RegisterShortExecute(); } diff --git a/onnxruntime/test/mlas/unittest/test_halfgemm.cpp b/onnxruntime/test/mlas/unittest/test_halfgemm.cpp index 87278cb25a88a..782d30e22910a 100644 --- a/onnxruntime/test/mlas/unittest/test_halfgemm.cpp +++ b/onnxruntime/test/mlas/unittest/test_halfgemm.cpp @@ -869,6 +869,40 @@ TEST(HalfGemmKleidiAISVE2p1, RejectsKAboveKernelLimit) { #endif } +TEST(HalfGemmKleidiAISVE2p1, RejectsTransposedB) { +#if defined(MLAS_TARGET_ARM64) && defined(USE_KLEIDIAI) && defined(MLAS_ENABLE_TEST_HOOKS) + const auto& cpuid = MLAS_CPUIDINFO::GetCPUIDInfo(); + if (!cpuid.HasArmSVE2p1() || cpuid.HasArm_SME() || cpuid.HasArm_SME2()) { + GTEST_SKIP() << "SVE2.1 must be selected without an SME-priority backend."; + } + ASSERT_TRUE(IsSve2p1HalfGemmSelected()); + + constexpr size_t M = 2; + constexpr size_t N = 3; + constexpr size_t K = 4; + std::vector a(M * K, MLFp16(1.0f)); + std::vector b(N * K, MLFp16(1.0f)); + std::vector c(M * N, MLFp16(-1.0f)); + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.A = a.data(); + data.B = b.data(); + data.C = reinterpret_cast(c.data()); + data.lda = K; + data.ldb = K; + data.ldc = N; + data.BIsTransposed = true; + + EXPECT_EQ(ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasTrans, N, K), size_t{0}); + EXPECT_FALSE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr)); +#if !defined(ORT_NO_EXCEPTIONS) + EXPECT_THROW(MlasHalfGemmBatch(M, N, K, 1, &data, nullptr), std::runtime_error); +#endif +#else + GTEST_SKIP() << "SVE2.1 HalfGemm requires an ARM64 KleidiAI test-hook build."; +#endif +} + TEST(HalfGemmKleidiAISVE2p1, RuntimePackedTailCasesSingleThread) { #if defined(MLAS_TARGET_ARM64) && defined(USE_KLEIDIAI) && defined(MLAS_ENABLE_TEST_HOOKS) const auto& cpuid = MLAS_CPUIDINFO::GetCPUIDInfo(); @@ -1124,6 +1158,131 @@ TEST(HalfGemmKleidiAIPath, PackedBFloatSingleThreadVariedShapesAndBiasWithoutOut // KleidiAI-specific packed-B uses a separate direct-consumption contract from // generic halfgemm PackB. Unsupported combinations fail at the public API // boundary because generic MLAS cannot consume this backend-native layout. +TEST(HalfGemmKleidiAIPath, TransposedBRuntimePacking) { +#if defined(MLAS_TARGET_ARM64) + const auto& cpuid = MLAS_CPUIDINFO::GetCPUIDInfo(); + if (!cpuid.HasArm_SME() && !cpuid.HasArm_SME2()) { + GTEST_SKIP() << "Transposed-B HalfGemm requires an SME or SME2 backend."; + } + + constexpr size_t M = 2; + constexpr size_t N = 3; + constexpr size_t K = 4; + constexpr size_t ldb = K + 1; + std::vector a(M * K); + std::vector b(N * ldb, MLFp16(-1.0f)); + std::vector bias{MLFp16(1.0f), MLFp16(2.0f), MLFp16(3.0f)}; + std::vector c(M * N, MLFp16(0.0f)); + + for (size_t m = 0; m < M; ++m) { + for (size_t k = 0; k < K; ++k) { + a[m * K + k] = MLFp16(static_cast(m * K + k + 1)); + } + } + for (size_t n = 0; n < N; ++n) { + for (size_t k = 0; k < K; ++k) { + b[n * ldb + k] = MLFp16(static_cast((n + 1) * (k + 1))); + } + } + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.A = a.data(); + data.B = b.data(); + data.Bias = reinterpret_cast(bias.data()); + data.C = reinterpret_cast(c.data()); + data.lda = K; + data.ldb = ldb; + data.ldc = N; + data.BIsTransposed = true; + + ASSERT_NE(ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasTrans, N, K), size_t{0}); + ASSERT_TRUE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr)); + for (size_t m = 0; m < M; ++m) { + for (size_t n = 0; n < N; ++n) { + float expected = float(bias[n]); + for (size_t k = 0; k < K; ++k) { + expected += float(a[m * K + k]) * float(b[n * ldb + k]); + } + EXPECT_EQ(c[m * N + n], MLFp16(expected)); + } + } + + std::fill(c.begin(), c.end(), MLFp16(-1.0f)); + data.ldb = K - 1; + EXPECT_FALSE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr)); + for (const auto value : c) { + EXPECT_EQ(value, MLFp16(-1.0f)); + } +#else + GTEST_SKIP() << "Transposed-B HalfGemm requires an ARM64 KleidiAI build."; +#endif +} + +TEST(HalfGemmKleidiAIPath, TransposedBBackendNativePacking) { +#if defined(MLAS_TARGET_ARM64) + const auto& cpuid = MLAS_CPUIDINFO::GetCPUIDInfo(); + if (!cpuid.HasArm_SME() && !cpuid.HasArm_SME2()) { + GTEST_SKIP() << "Transposed-B HalfGemm requires an SME or SME2 backend."; + } + + constexpr size_t M = 2; + constexpr size_t N = 3; + constexpr size_t K = 4; + constexpr size_t ldb = K + 1; + std::vector a(M * K); + std::vector b(N * ldb, MLFp16(-1.0f)); + std::vector c(M * N, MLFp16(0.0f)); + + for (size_t m = 0; m < M; ++m) { + for (size_t k = 0; k < K; ++k) { + a[m * K + k] = MLFp16(static_cast(m * K + k + 1)); + } + } + for (size_t n = 0; n < N; ++n) { + for (size_t k = 0; k < K; ++k) { + b[n * ldb + k] = MLFp16(static_cast((n + 1) * (k + 1))); + } + } + + const size_t packed_b_size = + ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasTrans, N, K); + ASSERT_NE(packed_b_size, size_t{0}); + std::vector packed_b(packed_b_size); + ASSERT_TRUE(ArmKleidiAI::MlasHalfGemmKleidiAIPackB( + CblasNoTrans, CblasTrans, N, K, + reinterpret_cast(b.data()), ldb, packed_b.data())); + + MLAS_HALF_GEMM_DATA_PARAMS data{}; + data.A = a.data(); + data.B = packed_b.data(); + data.C = reinterpret_cast(c.data()); + data.lda = K; + data.ldb = 0; + data.ldc = N; + data.BIsBackendNativePacked = true; + + ASSERT_TRUE(ArmKleidiAI::MlasHalfGemmBatch(M, N, K, 1, &data, nullptr)); + for (size_t m = 0; m < M; ++m) { + for (size_t n = 0; n < N; ++n) { + float expected = 0.0f; + for (size_t k = 0; k < K; ++k) { + expected += float(a[m * K + k]) * float(b[n * ldb + k]); + } + EXPECT_EQ(c[m * N + n], MLFp16(expected)); + } + } + + constexpr std::byte packed_b_initial_value{0x7f}; + std::fill(packed_b.begin(), packed_b.end(), packed_b_initial_value); + EXPECT_FALSE(ArmKleidiAI::MlasHalfGemmKleidiAIPackB( + CblasNoTrans, CblasTrans, N, K, + reinterpret_cast(b.data()), K - 1, packed_b.data())); + ExpectBufferFilledWith(packed_b, packed_b_initial_value); +#else + GTEST_SKIP() << "Transposed-B HalfGemm requires an ARM64 KleidiAI build."; +#endif +} + TEST(HalfGemmKleidiAIPath, KleidiAIPackedBWithBiasIsRejected) { if (!MlasFp16AccelerationSupported()) { GTEST_SKIP(); @@ -1247,7 +1406,7 @@ TEST(HalfGemmKleidiAIPath, KleidiAIPackedBSizeRejectsUnsupportedTranspose) { constexpr size_t K = 9; EXPECT_EQ(ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasTrans, CblasNoTrans, N, K), size_t{0}); - EXPECT_EQ(ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasTrans, N, K), size_t{0}); + EXPECT_EQ(ArmKleidiAI::MlasHalfGemmKleidiAIPackBSize(CblasNoTrans, CblasConjTrans, N, K), size_t{0}); } TEST(HalfGemmKleidiAIPath, KleidiAIPackedBRejectsUnsupportedTranspose) { @@ -1268,7 +1427,7 @@ TEST(HalfGemmKleidiAIPath, KleidiAIPackedBRejectsUnsupportedTranspose) { EXPECT_FALSE(ArmKleidiAI::MlasHalfGemmKleidiAIPackB( CblasTrans, CblasNoTrans, N, K, reinterpret_cast(B.data()), N, packed_b.data())); EXPECT_FALSE(ArmKleidiAI::MlasHalfGemmKleidiAIPackB( - CblasNoTrans, CblasTrans, N, K, reinterpret_cast(B.data()), N, packed_b.data())); + CblasNoTrans, CblasConjTrans, N, K, reinterpret_cast(B.data()), N, packed_b.data())); } TEST(HalfGemmKleidiAIPath, KleidiAIPackedBRejectsInvalidLeadingDimension) { diff --git a/onnxruntime/test/mlas/unittest/test_hgemm_neon.cpp b/onnxruntime/test/mlas/unittest/test_hgemm_neon.cpp index 01f1086a505e2..2ae7ed7057fad 100644 --- a/onnxruntime/test/mlas/unittest/test_hgemm_neon.cpp +++ b/onnxruntime/test/mlas/unittest/test_hgemm_neon.cpp @@ -823,10 +823,12 @@ static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_exe size_t count = 0; if (is_short_execute) { count += MlasDirectShortExecuteTests::RegisterShortExecute(); - count += MlasDirectShortExecuteTests::RegisterShortExecute(); - count += MlasDirectShortExecuteTests::RegisterShortExecute(); - count += MlasDirectShortExecuteTests::RegisterShortExecute(); - count += MlasDirectShortExecuteTests::RegisterShortExecute(); + if (MlasFp16AccelerationSupported()) { + count += MlasDirectShortExecuteTests::RegisterShortExecute(); + count += MlasDirectShortExecuteTests::RegisterShortExecute(); + count += MlasDirectShortExecuteTests::RegisterShortExecute(); + count += MlasDirectShortExecuteTests::RegisterShortExecute(); + } } return count; }); diff --git a/onnxruntime/test/mlas/unittest/test_hqnbitgemm_neon.cpp b/onnxruntime/test/mlas/unittest/test_hqnbitgemm_neon.cpp index 3762e30af352d..c20642ff8c2f7 100644 --- a/onnxruntime/test/mlas/unittest/test_hqnbitgemm_neon.cpp +++ b/onnxruntime/test/mlas/unittest/test_hqnbitgemm_neon.cpp @@ -743,16 +743,16 @@ static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_exe if (GetMlasPlatform().QNBitGemmDispatch->HQ4BitGemmPackQuantBData) { count += MlasDirectShortExecuteTests::RegisterShortExecute(); } - if (GetMlasPlatform().QNBitGemmDispatch->HQ4BitBlkDequantBForHgemm_CompFp16) { + if (GetMlasPlatform().QNBitGemmDispatch->HQ4BitBlkDequantBForHgemm_CompFp16 && MlasFp16AccelerationSupported()) { count += MlasDirectShortExecuteTests::RegisterShortExecute(); } - if (GetMlasPlatform().QNBitGemmDispatch->HQ4BitGemmKernel_CompFp16) { + if (GetMlasPlatform().QNBitGemmDispatch->HQ4BitGemmKernel_CompFp16 && MlasFp16AccelerationSupported()) { count += MlasDirectShortExecuteTests::RegisterShortExecute(); } if (GetMlasPlatform().QNBitGemmDispatch->HQ8BitGemmPackQuantBData) { count += MlasDirectShortExecuteTests::RegisterShortExecute(); } - if (GetMlasPlatform().QNBitGemmDispatch->HQ8BitBlkDequantBForHgemm_CompFp16) { + if (GetMlasPlatform().QNBitGemmDispatch->HQ8BitBlkDequantBForHgemm_CompFp16 && MlasFp16AccelerationSupported()) { count += MlasDirectShortExecuteTests::RegisterShortExecute(); } } diff --git a/onnxruntime/test/mlas/unittest/test_snchwc_partition.cpp b/onnxruntime/test/mlas/unittest/test_snchwc_partition.cpp index 56b8f4f2e80a3..fc251f62f6728 100644 --- a/onnxruntime/test/mlas/unittest/test_snchwc_partition.cpp +++ b/onnxruntime/test/mlas/unittest/test_snchwc_partition.cpp @@ -89,20 +89,27 @@ struct PartitionShape { // // The shape grid. Output channel counts are chosen to cover every ragged tail // the mapping has to handle: 16 and 32 are shorter than one filter set, 64 and -// 256 divide evenly, and 48/96/144/208 leave last sets of 3, 2, 1 and 1 blocks. +// 256 divide evenly, and 48/96/144/208 leave last sets of 3, 2, 1 and 1 blocks +// at FilterSetSize=4. The same channel counts are replayed at FilterSetSize=2 +// and 1 -- the sizes ChooseFilterSetSize falls back to under MLAS_NCHWC_FSS_TARGET +// -- which lands on a different, independent set of ragged remainders (see the +// comment on ChooseFilterSetSizeDispatch below for the worked-out remainders). // std::vector AllShapes() { const size_t OutputChannelValues[] = {16, 32, 48, 64, 96, 144, 208, 256}; const size_t OutputHeightValues[] = {1, 2, 7, 13, 52}; const size_t BatchGroupValues[] = {1, 2, 8}; + const size_t FilterSetSizeValues[] = {4, 2, 1}; std::vector Shapes; - for (size_t OutputChannels : OutputChannelValues) { - for (size_t OutputHeight : OutputHeightValues) { - for (size_t BatchGroupCount : BatchGroupValues) { - Shapes.push_back(PartitionShape{16, 4, OutputChannels, OutputHeight, BatchGroupCount}); + for (size_t FilterSetSize : FilterSetSizeValues) { + for (size_t OutputChannels : OutputChannelValues) { + for (size_t OutputHeight : OutputHeightValues) { + for (size_t BatchGroupCount : BatchGroupValues) { + Shapes.push_back(PartitionShape{16, FilterSetSize, OutputChannels, OutputHeight, BatchGroupCount}); + } } } } @@ -293,3 +300,81 @@ TEST(SnchwcPartition, WorkedExampleFromDocumentation) { ASSERT_LE(Worst, Ideal + Shape.FilterSetSize) << "worst thread took " << Worst << " block-rows against ideal " << Ideal; } + +// +// MlasNchwcChooseFilterSetSize halves MaximumFilterSetSize (4) down to 2 or 1 +// whenever the larger size would not produce Target work units per thread. +// With MLAS_NCHWC_FSS_TARGET unset in every other test in this binary, Target +// is always 0 and the function always takes its first-line early return -- +// the halving loop below is otherwise completely unexercised. Calling the +// pure helper directly, rather than going through the environment variable, +// makes every rung of that loop reachable regardless of what the process-wide +// cached env var read elsewhere resolved to. +// +// Shape: BlockSize=16, OutputChannels=80 -> Blocks=5 (ragged at every size: +// FilterSetCount/LastSetFilterCount are 2/1 at Size=4, 3/1 at Size=2, 5/1 at +// Size=1), BatchCount=1, GroupCount=1, OutputHeight=10. WorkAt(Size) = +// BatchCount * GroupCount * ceil(Blocks/Size) * OutputHeight: +// WorkAt(4) = 1 * 1 * ceil(5/4) * 10 = 20 +// WorkAt(2) = 1 * 1 * ceil(5/2) * 10 = 30 +// WorkAt(1) = 1 * 1 * ceil(5/1) * 10 = 50 +// + +TEST(SnchwcPartition, ChooseFilterSetSizeDispatch) { + constexpr size_t MaximumFilterSetSize = 4; + constexpr size_t BlockSize = 16; + constexpr size_t OutputChannels = 80; // Blocks = 5, ragged at every size. + constexpr size_t BatchCount = 1; + constexpr size_t GroupCount = 1; + constexpr size_t OutputHeight = 10; + + // Target == 0 (MLAS_NCHWC_FSS_TARGET unset) always short-circuits to the + // maximum, regardless of thread count or shape -- the one path every other + // test in this binary exercises. + EXPECT_EQ(MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, /*Target=*/0, /*tids=*/8, + BlockSize, OutputChannels, BatchCount, GroupCount, OutputHeight), + MaximumFilterSetSize); + + // Single-threaded also always short-circuits, even with a Target set: there + // is no second thread to starve. + EXPECT_EQ(MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, /*Target=*/100, /*tids=*/1, + BlockSize, OutputChannels, BatchCount, GroupCount, OutputHeight), + MaximumFilterSetSize); + + // WorkAt(4) = 20 >= Wanted: stays at the maximum. + EXPECT_EQ(MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, /*Target=*/5, /*tids=*/4, + BlockSize, OutputChannels, BatchCount, GroupCount, OutputHeight), + size_t{4}); + + // Boundary: Wanted == WorkAt(4) exactly still satisfies ">=" and keeps 4. + EXPECT_EQ(MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, /*Target=*/5, /*tids=*/4, + BlockSize, OutputChannels, BatchCount, GroupCount, OutputHeight), + size_t{4}); + + // Boundary: Wanted == WorkAt(4) + 1 falls one short, forcing exactly one + // halving; WorkAt(2) = 30 comfortably covers it. + EXPECT_EQ(MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, /*Target=*/7, /*tids=*/3, + BlockSize, OutputChannels, BatchCount, GroupCount, OutputHeight), + size_t{2}); + + // WorkAt(4) = 20 and WorkAt(2) = 30 both fall short of Wanted = 35; only + // WorkAt(1) = 50 covers it, forcing the halving loop all the way to 1. + EXPECT_EQ(MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, /*Target=*/7, /*tids=*/5, + BlockSize, OutputChannels, BatchCount, GroupCount, OutputHeight), + size_t{1}); + + // Even an unreasonably large Target cannot go below 1: the loop condition + // is "while (Size > 1)", so it always terminates at the floor. + EXPECT_EQ(MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, /*Target=*/1000000, /*tids=*/32, + BlockSize, OutputChannels, BatchCount, GroupCount, OutputHeight), + size_t{1}); + + // BatchCount and GroupCount multiply into the same work-unit count as + // OutputHeight; scaling them instead of OutputHeight must dispatch + // identically to the WorkAt(2)=30-vs-Wanted=21 case above (2 * 1 * ceil(5/2) * + // 5 = 30, same as 1 * 1 * ceil(5/2) * 10). + EXPECT_EQ(MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, /*Target=*/7, /*tids=*/3, + BlockSize, OutputChannels, /*BatchCount=*/2, /*GroupCount=*/1, + /*OutputHeight=*/5), + size_t{2}); +} diff --git a/onnxruntime/test/mlas/unittest/test_softcap.cpp b/onnxruntime/test/mlas/unittest/test_softcap.cpp index c3aa3fed6c011..7b10b170bc6f1 100644 --- a/onnxruntime/test/mlas/unittest/test_softcap.cpp +++ b/onnxruntime/test/mlas/unittest/test_softcap.cpp @@ -102,11 +102,15 @@ class MlasComputeSoftcapTest : public MlasTestBase { } }; +#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64) +// Both suites only have FP16 test paths, so register them only when FP16 +// arithmetic is available; otherwise they would pass without running any check. static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_execute) { size_t count = 0; - if (is_short_execute) { + if (is_short_execute && MlasFp16AccelerationSupported()) { count += MlasDirectShortExecuteTests::RegisterShortExecute(); count += MlasDirectShortExecuteTests::RegisterShortExecute(); } return count; }); +#endif // defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64) diff --git a/onnxruntime/test/mlas/unittest/test_softmax.cpp b/onnxruntime/test/mlas/unittest/test_softmax.cpp index 4d7a45143b311..8bf1438ee0fc3 100644 --- a/onnxruntime/test/mlas/unittest/test_softmax.cpp +++ b/onnxruntime/test/mlas/unittest/test_softmax.cpp @@ -44,6 +44,9 @@ class MlasComputeExpTest : public MlasTestBase { #if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64) void TestFp16(size_t N, float MinimumValue, float MaximumValue) { + if (!MlasFp16AccelerationSupported()) { + return; + } MLAS_FP16* Input = BufferInputFp16.GetBuffer(N); MLAS_FP16* Output = BufferOutputFp16.GetBuffer(N); @@ -70,6 +73,9 @@ class MlasComputeExpTest : public MlasTestBase { } void TestSumFp16(size_t N, float MinimumValue, float MaximumValue) { + if (!MlasFp16AccelerationSupported()) { + return; + } MLAS_FP16* Input = BufferInputFp16.GetBuffer(N); MLAS_FP16* Output = BufferOutputFp16.GetBuffer(N); @@ -168,6 +174,9 @@ class MlasSoftmaxTest : public MlasTestBase { #if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64) void TestReduceMaxFp16(size_t N, float MinimumValue, float MaximumValue) { + if (!MlasFp16AccelerationSupported()) { + return; + } MLAS_FP16* Input = BufferInputFp16.GetBuffer(N); std::default_random_engine generator(static_cast(N)); @@ -193,6 +202,9 @@ class MlasSoftmaxTest : public MlasTestBase { } void TestFp16(size_t N, size_t D, float MinimumValue, float MaximumValue, bool LogSoftmax, bool SmoothSoftmax) { + if (!MlasFp16AccelerationSupported()) { + return; + } MLAS_FP16* Input = BufferInputFp16.GetBuffer(N * D); MLAS_FP16* Output = BufferOutputFp16.GetBuffer(N * D); float* InputReference = BufferInput.GetBuffer(N * D); diff --git a/onnxruntime/test/onnx/TestCase.cc b/onnxruntime/test/onnx/TestCase.cc index 86eb9623b4905..a2c54a6904ea8 100644 --- a/onnxruntime/test/onnx/TestCase.cc +++ b/onnxruntime/test/onnx/TestCase.cc @@ -958,6 +958,8 @@ std::unique_ptr> GetBrokenTests(const std::string& provider {"cast_UINT4_to_FLOAT", "Skipped until onnxruntime/cmake/external/onnx points to onnx 1.19 which should include @onnx/onnx/pull/7074"}, {"cast_UINT4_to_FLOAT16", "Skipped until onnxruntime/cmake/external/onnx points to onnx 1.19 which should include @onnx/onnx/pull/7074"}, {"cast_UINT4_to_UINT8", "Skipped until onnxruntime/cmake/external/onnx points to onnx 1.19 which should include @onnx/onnx/pull/7074"}, + {"reduce_max_empty_set_bool", + "ReduceMax does not support bool; mirrors onnx_backend_test_series_filters.jsonc global exclusion"}, // onnx#7959: the materialized node corpus now runs on the CPU EP via the // onnx_test_node_materialized ctest (previously the C++ node runner was QNN-only). ORT has // no per-channel-quantization ConvInteger support on any EP -- the CPU kernel throws and @@ -1091,8 +1093,6 @@ std::unique_ptr> GetBrokenTests(const std::string& provider {"resize_upsample_sizes_nearest_ceil_half_pixel", "Bad onnx test output. Needs test fix."}, {"resize_upsample_sizes_nearest_floor_align_corners", "Bad onnx test output. Needs test fix."}, {"resize_upsample_sizes_nearest_round_prefer_ceil_asymmetric", "Bad onnx test output. Needs test fix."}, - {"bitshift_right_uint16", "BitShift(11) uint16 support not enabled currently"}, - {"bitshift_left_uint16", "BitShift(11) uint16 support not enabled currently"}, {"maxunpool_export_with_output_shape", "Invalid output in ONNX test. See https://github.com/onnx/onnx/issues/2398"}, {"cntk_simple_seg", "Bad onnx test output caused by wrong SAME_UPPER/SAME_LOWER for ConvTranspose"}, diff --git a/onnxruntime/test/onnx/main.cc b/onnxruntime/test/onnx/main.cc index 5740a407c7ef6..1817a92c91e7f 100644 --- a/onnxruntime/test/onnx/main.cc +++ b/onnxruntime/test/onnx/main.cc @@ -908,6 +908,9 @@ select from 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. \n)"); ORT_TSTR("rotary_embedding_3d_input_expanded"), ORT_TSTR("rotary_embedding_expanded"), ORT_TSTR("rotary_embedding_interleaved_expanded"), + // QNN CPU returns uniform Softmax probabilities for the expanded function's all-negative-infinity rows, + // so its zero-row guard is not preserved. Keep the native Attention case enabled. + ORT_TSTR("attention_4d_causal_nonpad_negative_offset_structural_empty_expanded"), // QNN don't support fmod = 1 ORT_TSTR("mod_mixed_sign_float64"), ORT_TSTR("mod_mixed_sign_float32"), diff --git a/onnxruntime/test/onnx/tensorprotoutils.cc b/onnxruntime/test/onnx/tensorprotoutils.cc index 2af930ff75bcf..4ec48f558741b 100644 --- a/onnxruntime/test/onnx/tensorprotoutils.cc +++ b/onnxruntime/test/onnx/tensorprotoutils.cc @@ -19,6 +19,7 @@ #include "core/framework/endian_utils.h" #include "core/graph/onnx_protobuf.h" #include "core/session/onnxruntime_cxx_api.h" +#include "core/session/onnxruntime_type_conversion.h" #include "mem_buffer.h" struct OrtStatus { @@ -621,41 +622,8 @@ ORT_API(void, OrtUninitializeBuffer, _In_opt_ void* input, size_t input_len, enu (Y*)preallocated, static_cast(tensor_size)); \ break; -#define CASE_TYPE(X) \ - case onnx::TensorProto_DataType_##X: \ - return ONNX_TENSOR_ELEMENT_DATA_TYPE_##X; - ONNXTensorElementDataType CApiElementTypeFromProtoType(int type) { - switch (type) { - CASE_TYPE(FLOAT) - CASE_TYPE(UINT8) - CASE_TYPE(INT8) - CASE_TYPE(UINT16) - CASE_TYPE(INT16) - CASE_TYPE(INT32) - CASE_TYPE(INT64) - CASE_TYPE(STRING) - CASE_TYPE(BOOL) - CASE_TYPE(FLOAT16) - CASE_TYPE(DOUBLE) - CASE_TYPE(UINT32) - CASE_TYPE(UINT64) - CASE_TYPE(COMPLEX64) - CASE_TYPE(COMPLEX128) - CASE_TYPE(BFLOAT16) - CASE_TYPE(FLOAT8E4M3FN) - CASE_TYPE(FLOAT8E4M3FNUZ) - CASE_TYPE(FLOAT8E5M2) - CASE_TYPE(FLOAT8E5M2FNUZ) - CASE_TYPE(FLOAT8E8M0) - CASE_TYPE(FLOAT4E2M1) - CASE_TYPE(UINT4) - CASE_TYPE(INT4) - CASE_TYPE(UINT2) - CASE_TYPE(INT2) - default: - return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; - } + return utils::ToOrtTensorElementDataType(type); } ONNXTensorElementDataType GetTensorElementType(const onnx::TensorProto& tensor_proto) { diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 933167791554b..3faced9bacb23 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -5,6 +5,7 @@ #pragma warning(disable : 4244) #endif +#include #include #include #include @@ -3080,6 +3081,7 @@ TEST_F(GraphTransformationTests, DoNotApplyFuseMatmulBNDirectly) { ASSERT_EQ(op_to_count["Gemm"], 0); } +// should not fuse: the Reshape to [1, 1, 1, 5] moves the MatMul columns off the channel axis TEST_F(GraphTransformationTests, FuseMatmulBNWithOnlyReshape) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/fuse-matmul-bn-only-reshape.onnx"; @@ -3087,15 +3089,6 @@ TEST_F(GraphTransformationTests, FuseMatmulBNWithOnlyReshape) { ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); Graph& graph = p_model->MainGraph(); - std::string expected_output_name; - GraphViewer graphViewer(graph); - for (auto& node_index : graphViewer.GetNodesInTopologicalOrder()) { - auto& node = *graph.GetNode(node_index); - if (node.OpType() == "MatMul") { - expected_output_name = node.OutputDefs()[0]->Name(); - } - } - onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; auto rule_transformer_L1 = std::make_unique("RuleTransformerL1"); ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); @@ -3104,18 +3097,12 @@ TEST_F(GraphTransformationTests, FuseMatmulBNWithOnlyReshape) { ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); std::map op_to_count = CountOpsInGraph(graph); - ASSERT_EQ(op_to_count["BatchNormalization"], 0); - ASSERT_EQ(op_to_count["MatMul"], 0); - ASSERT_EQ(op_to_count["Gemm"], 1); - - for (auto& node : graph.Nodes()) { - if (node.OpType() == "Gemm") { - ASSERT_EQ(node.OutputDefs()[0]->Name(), expected_output_name) - << "fusion should produce the same output name as the MatMul node"; - } - } + ASSERT_EQ(op_to_count["BatchNormalization"], 1); + ASSERT_EQ(op_to_count["MatMul"], 1); + ASSERT_EQ(op_to_count["Gemm"], 0); } +// should not fuse: the Transpose moves the MatMul columns off the channel axis TEST_F(GraphTransformationTests, FuseMatmulBNWithOnlyTranspose) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/fuse-matmul-bn-only-transpose.onnx"; @@ -3123,15 +3110,6 @@ TEST_F(GraphTransformationTests, FuseMatmulBNWithOnlyTranspose) { ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); Graph& graph = p_model->MainGraph(); - std::string expected_output_name; - GraphViewer graphViewer(graph); - for (auto& node_index : graphViewer.GetNodesInTopologicalOrder()) { - auto& node = *graph.GetNode(node_index); - if (node.OpType() == "MatMul") { - expected_output_name = node.OutputDefs()[0]->Name(); - } - } - onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; auto rule_transformer_L1 = std::make_unique("RuleTransformerL1"); ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); @@ -3140,16 +3118,9 @@ TEST_F(GraphTransformationTests, FuseMatmulBNWithOnlyTranspose) { ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); std::map op_to_count = CountOpsInGraph(graph); - ASSERT_EQ(op_to_count["BatchNormalization"], 0); - ASSERT_EQ(op_to_count["MatMul"], 0); - ASSERT_EQ(op_to_count["Gemm"], 1); - - for (auto& node : graph.Nodes()) { - if (node.OpType() == "Gemm") { - ASSERT_EQ(node.OutputDefs()[0]->Name(), expected_output_name) - << "fusion should produce the same output name as the MatMul node"; - } - } + ASSERT_EQ(op_to_count["BatchNormalization"], 1); + ASSERT_EQ(op_to_count["MatMul"], 1); + ASSERT_EQ(op_to_count["Gemm"], 0); } TEST_F(GraphTransformationTests, FuseMatmulBNWithoutBatchNormalization) { @@ -3199,6 +3170,48 @@ TEST_F(GraphTransformationTests, FuseMatmulBNWithNonIgnorableNode) { ASSERT_EQ(op_to_count["Gemm"], 0); } +// The fusion is only valid if the Reshape keeps the MatMul columns as the BatchNormalization channels. +TEST_F(GraphTransformationTests, FuseMatmulBNWithChannelChangingReshape) { + for (const auto& [reshape_shape, should_fuse] : std::vector, bool>>{ + {{2, 4, 1, 1}, true}, + {{1, 4, 2, 1}, false}}) { + auto build_test_case = [&reshape_shape](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({2, 3}, -1.0f, 1.0f); + auto* matmul_b = builder.MakeInitializer({3, 4}, -1.0f, 1.0f); + auto* shape = builder.MakeInitializer({4}, reshape_shape); + auto* bn_scale = builder.MakeInitializer({4}, 0.5f, 1.5f); + auto* bn_bias = builder.MakeInitializer({4}, -1.0f, 1.0f); + auto* bn_mean = builder.MakeInitializer({4}, -1.0f, 1.0f); + auto* bn_var = builder.MakeInitializer({4}, 0.5f, 1.5f); + auto* matmul_output = builder.MakeIntermediate(); + auto* reshape_output = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("MatMul", {input, matmul_b}, {matmul_output}); + builder.AddNode("Reshape", {matmul_output, shape}, {reshape_output}); + builder.AddNode("BatchNormalization", {reshape_output, bn_scale, bn_bias, bn_mean, bn_var}, {output}) + .AddAttribute("epsilon", 1e-5f); + }; + + auto check_transformed_graph = [should_fuse](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + ASSERT_EQ(op_to_count["BatchNormalization"], should_fuse ? 0 : 1); + ASSERT_EQ(op_to_count["Gemm"], should_fuse ? 1 : 0); + }; + + auto rule_transformer_L1 = std::make_unique("RuleTransformerL1"); + ASSERT_STATUS_OK(rule_transformer_L1->Register(std::make_unique())); + TransformerTester(build_test_case, + check_transformed_graph, + TransformerLevel::Default, + TransformerLevel::Level1, + 15, + 1e-4, + 1e-4, + std::move(rule_transformer_L1)); + } +} + TEST_F(GraphTransformationTests, DontFuseConvWithBNWithOptionalOutputs) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/fuse-conv-bn-no-bias.onnx"; @@ -4561,13 +4574,14 @@ ConvActivationBuilder SimpleActivation(const std::string& op_type, // Use fp16, group 1, NHWC, and a non-1x1 kernel to select Im2ColMatMulProgram. void RunWebGpuIm2ColActivationParity(const ConvActivationBuilder& add_activation, const std::string& expected_activation, - int opset_version) { + int opset_version, + int64_t input_channels = 4) { if (!DefaultWebGpuExecutionProvider()) { GTEST_SKIP() << "WebGPU EP unavailable in this build."; } - const std::vector input_shape{1, 4, 14, 14}; - const std::vector weight_shape{8, 4, 3, 3}; + const std::vector input_shape{1, input_channels, 14, 14}; + const std::vector weight_shape{8, input_channels, 3, 3}; auto to_fp16 = [](const std::vector& values) { std::vector converted; @@ -4588,8 +4602,10 @@ void RunWebGpuIm2ColActivationParity(const ConvActivationBuilder& add_activation }; auto build_test_case = [&](ModelTestBuilder& builder) { - auto* input = builder.MakeInput(input_shape, to_fp16(ramp(1 * 4 * 14 * 14, -3.0f, 3.0f))); - auto* weight = builder.MakeInitializer(weight_shape, to_fp16(ramp(8 * 4 * 3 * 3, -0.5f, 0.5f))); + auto* input = builder.MakeInput( + input_shape, to_fp16(ramp(static_cast(input_channels * 14 * 14), -3.0f, 3.0f))); + auto* weight = builder.MakeInitializer( + weight_shape, to_fp16(ramp(static_cast(8 * input_channels * 3 * 3), -0.5f, 0.5f))); auto* bias = builder.MakeInitializer({weight_shape[0]}, to_fp16(ramp(8, -0.5f, 0.5f))); auto* conv_out = builder.MakeIntermediate(); auto* output = builder.MakeOutput(); @@ -4644,10 +4660,11 @@ void RunWebGpuIm2ColActivationParity(const ConvActivationBuilder& add_activation "whether the im2col program was selected"; }; - RunWebGpuFusionTransformerTest(build_test_case, check_transformed_graph, TransformerLevel::Level1, TransformerLevel::Level2, opset_version, - /*per_sample_tolerance=*/2e-2, - /*relative_per_sample_tolerance=*/2e-2, - /*transformer=*/nullptr, []() { return DefaultWebGpuExecutionProvider(); }, [](SessionOptions& session_options) { + RunWebGpuFusionTransformerTest( + build_test_case, check_transformed_graph, TransformerLevel::Level1, TransformerLevel::Level2, opset_version, + /*per_sample_tolerance=*/2e-2, + /*relative_per_sample_tolerance=*/2e-2, + /*transformer=*/nullptr, []() { return DefaultWebGpuExecutionProvider(); }, [](SessionOptions& session_options) { session_options.enable_profiling = true; session_options.profile_file_prefix = ORT_TSTR("webgpu_im2col_activation"); }); @@ -4759,6 +4776,16 @@ TEST_F(GraphTransformationTests, WebGpuIm2ColConvReluFusionMatchesUnfusedResults RunWebGpuIm2ColActivationParity(SimpleActivation("Relu"), "Relu", 17); } +// Odd input channels give vec_size == 1, which takes the scalar path of the shader (including the +// subgroup variant when subgroup-size control is available). 4 channels give vec_size == 4 and 2 give 2. +TEST_F(GraphTransformationTests, WebGpuIm2ColConvOddInputChannelsMatchesUnfusedResults) { + RunWebGpuIm2ColActivationParity(SimpleActivation("Relu"), "Relu", 17, /*input_channels=*/3); +} + +TEST_F(GraphTransformationTests, WebGpuIm2ColConvTwoInputChannelsMatchesUnfusedResults) { + RunWebGpuIm2ColActivationParity(SimpleActivation("Relu"), "Relu", 17, /*input_channels=*/2); +} + TEST_F(GraphTransformationTests, WebGpuIm2ColConvLeakyReluFusionMatchesUnfusedResults) { RunWebGpuIm2ColActivationParity( SimpleActivation("LeakyRelu", kOnnxDomain, [](Node& node) { node.AddAttribute("alpha", 0.25f); }), @@ -7307,7 +7334,8 @@ TEST_F(GraphTransformationTests, ReshapeFusionContiguousReshapesWithSharedShapeP }; std::unique_ptr transformer = std::make_unique(); - ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 18, *logger_, std::move(transformer), TransformerLevel::Level1, 1, [](Graph&) { return Status::OK(); }, post_graph_checker)); + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 18, *logger_, std::move(transformer), TransformerLevel::Level1, 1, [](Graph&) { return Status::OK(); }, post_graph_checker)); } // Full-session test of a Reshape chain whose first shape input is produced by a Where node. @@ -7802,6 +7830,12 @@ TEST_F(GraphTransformationTests, SliceConcatToSpaceToDepthFusionTest) { auto check_transformed_graph = [get_op_count](InferenceSessionWrapper& session) { const Graph& graph = session.GetGraph(); const auto op_to_count = CountOpsInGraph(graph); + if (graph.DomainToVersionMap().at(kOnnxDomain) >= 28) { + ASSERT_EQ(get_op_count(op_to_count, "Slice"), 4); + ASSERT_EQ(get_op_count(op_to_count, "Concat"), 1); + ASSERT_EQ(get_op_count(op_to_count, "SpaceToDepth"), 0); + return; + } ASSERT_TRUE(op_to_count.count("Slice") == 0 || op_to_count.at("Slice") == 0); ASSERT_TRUE(op_to_count.count("Concat") == 0 || op_to_count.at("Concat") == 0); ASSERT_EQ(get_op_count(op_to_count, "SpaceToDepth"), 1); @@ -7814,14 +7848,17 @@ TEST_F(GraphTransformationTests, SliceConcatToSpaceToDepthFusionTest) { } }; - TransformerTester(build_test_case, - check_transformed_graph, - TransformerLevel::Default, - TransformerLevel::Level1, - 13, - 0.0, - 0.0, - std::make_unique()); + for (int opset : {13, 27, 28}) { + SCOPED_TRACE(opset); + TransformerTester(build_test_case, + check_transformed_graph, + TransformerLevel::Default, + TransformerLevel::Level1, + opset, + 0.0, + 0.0, + std::make_unique()); + } } TEST_F(GraphTransformationTests, SliceConcatToSpaceToDepthFusionWithConstantNodesTest) { @@ -8222,6 +8259,166 @@ TEST_F(GraphTransformationTests, SliceConcatToSpaceToDepthFusionNotTriggeredForR 1, pre_graph_checker, post_graph_checker)); } +namespace { + +// Builds the initializers shared by the chained-Slice focus tests below. +struct ChainedSliceInitializers { + NodeArg* axis_h; + NodeArg* axis_w; + NodeArg* step2; + NodeArg* start0; + NodeArg* start1; + NodeArg* end_full; +}; + +ChainedSliceInitializers MakeChainedSliceInitializers(ModelTestBuilder& builder, int64_t extent) { + return {builder.Make1DInitializer({2}), + builder.Make1DInitializer({3}), + builder.Make1DInitializer({2}), + builder.Make1DInitializer({0}), + builder.Make1DInitializer({1}), + builder.Make1DInitializer({extent})}; +} + +int GetOpCountOrZero(const OpCountMap& op_to_count, std::string_view op_type) { + const auto it = op_to_count.find(std::string(op_type)); + return it == op_to_count.end() ? 0 : it->second; +} + +} // namespace + +// A focus layer is commonly exported as two chained Slice nodes per branch, one +// striding height and one striding width, rather than a single Slice striding +// both axes. The composed phase is still a blocksize=2 space-to-depth, so the +// chain must fuse. +TEST_F(GraphTransformationTests, SliceConcatToSpaceToDepthFusionChainedSlicesTest) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({1, 3, 8, 8}, -1.0f, 1.0f); + const auto init = MakeChainedSliceInitializers(builder, 8); + + auto slice_height_then_width = [&](NodeArg* h_start, NodeArg* w_start) { + auto* h_out = builder.MakeIntermediate(); + auto* w_out = builder.MakeIntermediate(); + builder.AddNode("Slice", {input, h_start, init.end_full, init.axis_h, init.step2}, {h_out}); + builder.AddNode("Slice", {h_out, w_start, init.end_full, init.axis_w, init.step2}, {w_out}); + return w_out; + }; + + // Canonical block order, so no channel permutation is required. + auto* branch00 = slice_height_then_width(init.start0, init.start0); + auto* branch01 = slice_height_then_width(init.start0, init.start1); + auto* branch10 = slice_height_then_width(init.start1, init.start0); + auto* branch11 = slice_height_then_width(init.start1, init.start1); + + auto* concat_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Concat", {branch00, branch01, branch10, branch11}, {concat_out}) + .AddAttribute("axis", static_cast(1)); + builder.AddNode("Identity", {concat_out}, {output}); + }; + + auto check_transformed_graph = [](InferenceSessionWrapper& session) { + const auto op_to_count = CountOpsInGraph(session.GetGraph()); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "Slice"), 0); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "Concat"), 0); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "SpaceToDepth"), 1); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "Gather"), 0); + }; + + TransformerTester(build_test_case, check_transformed_graph, TransformerLevel::Default, + TransformerLevel::Level1, 13, 0.0, 0.0, + std::make_unique()); +} + +// The height Slice is normally shared by two width branches, so the graph holds +// six Slice nodes rather than eight and the upstream node has two consumers. +// This is the shape a YOLO focus layer actually exports, including its +// non-canonical (0,0) (1,0) (0,1) (1,1) block order, which needs a Gather to +// restore the channel order that Concat produced. +TEST_F(GraphTransformationTests, SliceConcatToSpaceToDepthFusionSharedHeightSliceTest) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({1, 3, 8, 8}, -1.0f, 1.0f); + const auto init = MakeChainedSliceInitializers(builder, 8); + + auto slice_height = [&](NodeArg* h_start) { + auto* h_out = builder.MakeIntermediate(); + builder.AddNode("Slice", {input, h_start, init.end_full, init.axis_h, init.step2}, {h_out}); + return h_out; + }; + auto slice_width = [&](NodeArg* h_out, NodeArg* w_start) { + auto* w_out = builder.MakeIntermediate(); + builder.AddNode("Slice", {h_out, w_start, init.end_full, init.axis_w, init.step2}, {w_out}); + return w_out; + }; + + auto* top = slice_height(init.start0); + auto* bottom = slice_height(init.start1); + + auto* top_left = slice_width(top, init.start0); + auto* top_right = slice_width(top, init.start1); + auto* bottom_left = slice_width(bottom, init.start0); + auto* bottom_right = slice_width(bottom, init.start1); + + auto* concat_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Concat", {top_left, bottom_left, top_right, bottom_right}, {concat_out}) + .AddAttribute("axis", static_cast(1)); + builder.AddNode("Identity", {concat_out}, {output}); + }; + + auto check_transformed_graph = [](InferenceSessionWrapper& session) { + const auto op_to_count = CountOpsInGraph(session.GetGraph()); + // Both the width Slices and the shared height Slices must be gone. + ASSERT_EQ(GetOpCountOrZero(op_to_count, "Slice"), 0); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "Concat"), 0); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "SpaceToDepth"), 1); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "Gather"), 1); + }; + + TransformerTester(build_test_case, check_transformed_graph, TransformerLevel::Default, + TransformerLevel::Level1, 13, 0.0, 0.0, + std::make_unique()); +} + +// A chain that strides the same axis twice is a 4x height decimation, not a +// space-to-depth, and must be left alone. +TEST_F(GraphTransformationTests, SliceConcatToSpaceToDepthFusionNotTriggeredForRepeatedAxisChainTest) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({1, 3, 8, 8}, -1.0f, 1.0f); + const auto init = MakeChainedSliceInitializers(builder, 8); + + auto slice_height_twice = [&](NodeArg* first_start, NodeArg* second_start) { + auto* first_out = builder.MakeIntermediate(); + auto* second_out = builder.MakeIntermediate(); + builder.AddNode("Slice", {input, first_start, init.end_full, init.axis_h, init.step2}, {first_out}); + builder.AddNode("Slice", {first_out, second_start, init.end_full, init.axis_h, init.step2}, {second_out}); + return second_out; + }; + + auto* branch00 = slice_height_twice(init.start0, init.start0); + auto* branch01 = slice_height_twice(init.start0, init.start1); + auto* branch10 = slice_height_twice(init.start1, init.start0); + auto* branch11 = slice_height_twice(init.start1, init.start1); + + auto* concat_out = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Concat", {branch00, branch01, branch10, branch11}, {concat_out}) + .AddAttribute("axis", static_cast(1)); + builder.AddNode("Identity", {concat_out}, {output}); + }; + + auto check_transformed_graph = [](InferenceSessionWrapper& session) { + const auto op_to_count = CountOpsInGraph(session.GetGraph()); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "Slice"), 8); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "Concat"), 1); + ASSERT_EQ(GetOpCountOrZero(op_to_count, "SpaceToDepth"), 0); + }; + + TransformerTester(build_test_case, check_transformed_graph, TransformerLevel::Default, + TransformerLevel::Level1, 13, 0.0, 0.0, + std::make_unique()); +} + TEST_F(GraphTransformationTests, ExpandElimination) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "expand_elimination.onnx"; std::shared_ptr model; @@ -11664,6 +11861,133 @@ TEST_F(GraphTransformationTests, MatMulScaleFusionDoubleType) { {kCpuExecutionProvider}); } +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyDoubleScale) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2, 2}}); + auto* input_b = builder.MakeInput({{2, 2}}); + auto* scale = builder.MakeScalarInitializer(std::nextafter(1.0, 2.0)); + auto* scaled_a = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Mul", {input_a, scale}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Mul") == 1); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyDoubleDivScale) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2, 2}}); + auto* input_b = builder.MakeInput({{2, 2}}); + auto* divisor = builder.MakeScalarInitializer(3.0); + auto* scaled_a = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Div", {input_a, divisor}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Div") == 1); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyCombinedDoubleScale) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2, 2}}); + auto* input_b = builder.MakeInput({{2, 2}}); + const double scale_value = static_cast(std::nextafter(1.0f, 2.0f)); + auto* scale1 = builder.MakeScalarInitializer(scale_value); + auto* scale2 = builder.MakeScalarInitializer(scale_value); + auto* scaled_a = builder.MakeIntermediate(); + auto* matmul_output = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Mul", {input_a, scale1}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {matmul_output}); + builder.AddNode("Mul", {matmul_output, scale2}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Mul") == 2); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsCancellingLossyDoubleScales) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{1, 1}}); + auto* input_b = builder.MakeInput({{1, 1}}); + auto* scale1 = builder.MakeScalarInitializer(1e300); + auto* scale2 = builder.MakeScalarInitializer(1e-300); + auto* scaled_a = builder.MakeIntermediate(); + auto* matmul_output = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Mul", {input_a, scale1}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {matmul_output}); + builder.AddNode("Mul", {matmul_output, scale2}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Mul") == 2); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, MatMulScaleFusionRejectsLossyIntegerScale) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2, 2}}); + auto* input_b = builder.MakeInput({{2, 2}}); + auto* scale = builder.MakeScalarInitializer((int64_t{1} << 24) + 1); + auto* scaled_a = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Mul", {input_a, scale}, {scaled_a}); + builder.AddNode("MatMul", {scaled_a, input_b}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("Mul") == 1); + TEST_RETURN_IF_NOT(op_to_count.at("MatMul") == 1); + TEST_RETURN_IF(op_to_count.count("com.microsoft.FusedMatMul") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + TEST_F(GraphTransformationTests, MatMulScaleFusionWithScaleInput) { TestMatMulScaleFusion( MODEL_FOLDER "fusion/matmul_scale_with_scale_input.onnx", *logger_, diff --git a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc index f1e20844c26d8..6696790410603 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc +++ b/onnxruntime/test/optimizer/graph_transform_test_layernorm.cc @@ -243,11 +243,12 @@ TEST_F(GraphTransformationTests, LayerNormWithSubDupFusionTest) { } void BuildLayerNorm(ModelTestBuilder& builder, std::vector reduce1_axes = {-1}, - std::vector reduce2_axes = {-1}) { + std::vector reduce2_axes = {-1}, bool epsilon_input_first = false, + float epsilon = 1e-5f) { std::vector input_shape = {2, 3, 3, 3}; auto* data_arg = builder.MakeInput(input_shape); auto* pow_initializer = builder.MakeInitializer({}, {2.0f}); - auto* add_initializer = builder.MakeInitializer({}, {1e-5f}); + auto* add_initializer = builder.MakeInitializer({}, {epsilon}); std::vector normalized_shape = {}; int64_t normalized_shape_size = 1; auto raxes = reduce1_axes; @@ -295,7 +296,11 @@ void BuildLayerNorm(ModelTestBuilder& builder, std::vector reduce1_axes } else { builder.AddNode("ReduceMean", {pow_out}, {reduce_mean_out_2}).AddAttribute("axes", reduce2_axes); } - builder.AddNode("Add", {reduce_mean_out_2, add_initializer}, {add_out_1}); + if (epsilon_input_first) { + builder.AddNode("Add", {add_initializer, reduce_mean_out_2}, {add_out_1}); + } else { + builder.AddNode("Add", {reduce_mean_out_2, add_initializer}, {add_out_1}); + } builder.AddNode("Sqrt", {add_out_1}, {sqrt_out}); builder.AddNode("Div", {cast_out_1, sqrt_out}, {div_out}); builder.AddNode("Cast", {div_out}, {cast_out_2}) @@ -344,6 +349,31 @@ TEST_F(GraphTransformationTests, LayerNormWithCastFusionTest_5) { TransformerLevel::Level2, 1, pre_graph_checker, post_graph_checker)); } +TEST_F(GraphTransformationTests, LayerNormFusionAddEpsilonInput0) { + auto build_test_case = [](ModelTestBuilder& builder) { + BuildLayerNorm(builder, {-1}, {-1}, true, 1e-4f); + }; + + auto post_graph_checker = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count.at("LayerNormalization") == 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "LayerNormalization") { + const auto epsilon_it = node.GetAttributes().find("epsilon"); + TEST_RETURN_IF_NOT(epsilon_it != node.GetAttributes().end()); + TEST_RETURN_IF_NOT(epsilon_it->second.f() == 1e-4f); + } + } + return Status::OK(); + }; + + const InlinedHashSet compatible_execution_providers; + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 14, *logger_, + std::make_unique(compatible_execution_providers, TransformerLevel::Level2), + TransformerLevel::Level2, 1, nullptr, post_graph_checker)); +} + TEST_F(GraphTransformationTests, LayerNormWithCastFusionTest_6) { auto build_test_case = [](ModelTestBuilder& builder) { BuildLayerNorm(builder, {-2}, {-1}); @@ -498,6 +528,187 @@ TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionTest) { } } +TEST_F(GraphTransformationTests, LayerNormFusionsRejectLossyDoubleEpsilon) { + auto build_layer_norm = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0, 1.0, 1.0, 1.0}); + auto* bias = builder.MakeInitializer({4}, {0.0, 0.0, 0.0, 0.0}); + auto* exponent = builder.MakeScalarInitializer(2.0); + auto* epsilon = builder.MakeScalarInitializer(1e-50); + auto* mean1 = builder.MakeIntermediate(); + auto* sub = builder.MakeIntermediate(); + auto* squared = builder.MakeIntermediate(); + auto* mean2 = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* scaled = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("ReduceMean", {input}, {mean1}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Sub", {input, mean1}, {sub}); + builder.AddNode("Pow", {sub, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean2}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean2, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {sub, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {scaled}); + builder.AddNode("Add", {scaled, bias}, {output}); + }; + + auto build_simplified_layer_norm = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0, 1.0, 1.0, 1.0}); + auto* exponent = builder.MakeScalarInitializer(2.0); + auto* epsilon = builder.MakeScalarInitializer(1e-50); + auto* squared = builder.MakeIntermediate(); + auto* mean = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Pow", {input, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {input, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF(op_to_count.count("LayerNormalization") != 0); + TEST_RETURN_IF(op_to_count.count("SimplifiedLayerNormalization") != 0); + TEST_RETURN_IF_NOT(op_to_count.at("Pow") == 1); + return Status::OK(); + }; + + const InlinedHashSet compatible_execution_providers; + ASSERT_STATUS_OK(TestGraphTransformer( + build_layer_norm, 14, *logger_, + std::make_unique(compatible_execution_providers, TransformerLevel::Level2), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); + ASSERT_STATUS_OK(TestGraphTransformer( + build_simplified_layer_norm, 14, *logger_, + std::make_unique(), + TransformerLevel::Level2, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, LayerNormFusionsAllowExactDoubleEpsilon) { + auto build_layer_norm = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0, 1.0, 1.0, 1.0}); + auto* bias = builder.MakeInitializer({4}, {0.0, 0.0, 0.0, 0.0}); + auto* exponent = builder.MakeScalarInitializer(2.0); + auto* epsilon = builder.MakeScalarInitializer(0.125); + auto* mean1 = builder.MakeIntermediate(); + auto* sub = builder.MakeIntermediate(); + auto* squared = builder.MakeIntermediate(); + auto* mean2 = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* scaled = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("ReduceMean", {input}, {mean1}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Sub", {input, mean1}, {sub}); + builder.AddNode("Pow", {sub, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean2}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean2, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {sub, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {scaled}); + builder.AddNode("Add", {scaled, bias}, {output}); + }; + + auto build_simplified_layer_norm = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0, 1.0, 1.0, 1.0}); + auto* exponent = builder.MakeScalarInitializer(2.0); + auto* epsilon = builder.MakeScalarInitializer(0.125); + auto* squared = builder.MakeIntermediate(); + auto* mean = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Pow", {input, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {input, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {output}); + }; + + auto make_checker = [](std::string_view fused_op) { + return [fused_op](Graph& graph) { + TEST_RETURN_IF_NOT(OpCount(CountOpsInGraph(graph), std::string{fused_op}) == 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == fused_op) { + const auto epsilon_it = node.GetAttributes().find("epsilon"); + TEST_RETURN_IF_NOT(epsilon_it != node.GetAttributes().end()); + TEST_RETURN_IF_NOT(epsilon_it->second.f() == 0.125f); + } + } + return Status::OK(); + }; + }; + + const InlinedHashSet compatible_execution_providers; + ASSERT_STATUS_OK(TestGraphTransformer( + build_layer_norm, 14, *logger_, + std::make_unique(compatible_execution_providers, TransformerLevel::Level2), + TransformerLevel::Level2, 1, nullptr, make_checker("LayerNormalization"))); + ASSERT_STATUS_OK(TestGraphTransformer( + build_simplified_layer_norm, 14, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, nullptr, make_checker("SimplifiedLayerNormalization"))); +} + +TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionPreservesIntegerToFloatCastEpsilon) { + constexpr uint64_t kEpsilonSource = (uint64_t{1} << 63) + (uint64_t{1} << 39) + 1; + auto build_test_case = [=](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 4}}); + auto* scale = builder.MakeInitializer({4}, {1.0f, 1.0f, 1.0f, 1.0f}); + auto* exponent = builder.MakeScalarInitializer(2.0f); + auto* epsilon_source = builder.MakeScalarInitializer(kEpsilonSource); + auto* epsilon = builder.MakeIntermediate(); + auto* squared = builder.MakeIntermediate(); + auto* mean = builder.MakeIntermediate(); + auto* variance = builder.MakeIntermediate(); + auto* stddev = builder.MakeIntermediate(); + auto* normalized = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + + builder.AddNode("Cast", {epsilon_source}, {epsilon}) + .AddAttribute("to", static_cast(TensorProto_DataType_FLOAT)); + builder.AddNode("Pow", {input, exponent}, {squared}); + builder.AddNode("ReduceMean", {squared}, {mean}).AddAttribute("axes", std::vector{-1}); + builder.AddNode("Add", {mean, epsilon}, {variance}); + builder.AddNode("Sqrt", {variance}, {stddev}); + builder.AddNode("Div", {input, stddev}, {normalized}); + builder.AddNode("Mul", {normalized, scale}, {output}); + }; + + auto check_after = [=](Graph& graph) { + TEST_RETURN_IF_NOT(OpCount(CountOpsInGraph(graph), "SimplifiedLayerNormalization") == 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "SimplifiedLayerNormalization") { + const auto epsilon_it = node.GetAttributes().find("epsilon"); + TEST_RETURN_IF_NOT(epsilon_it != node.GetAttributes().end()); + TEST_RETURN_IF_NOT(epsilon_it->second.f() == static_cast(kEpsilonSource)); + } + } + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 17, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, nullptr, check_after)); +} + TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionSharedCastPowExponent) { for (const bool has_leading_cast : {false, true}) { auto build_test_case = [has_leading_cast](ModelTestBuilder& builder) { @@ -2664,26 +2875,16 @@ TEST_F(GraphTransformationTests, LayerNormFusion_ZeroElementEpsilon) { // Must handle zero-element epsilon gracefully. ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); - // Verify the fusion fired and produced a LayerNormalization with the - // default epsilon (1e-5f), per the zero-element fallback in - // layer_norm_fusion.cc. + // A non-scalar epsilon cannot be represented by the fused float attribute, + // so the original graph must remain unchanged. std::map op_to_count = CountOpsInGraph(graph); - EXPECT_EQ(op_to_count["LayerNormalization"], 1); - EXPECT_EQ(op_to_count["ReduceMean"], 0); - EXPECT_EQ(op_to_count["Sub"], 0); - EXPECT_EQ(op_to_count["Pow"], 0); - EXPECT_EQ(op_to_count["Sqrt"], 0); - EXPECT_EQ(op_to_count["Div"], 0); - EXPECT_EQ(op_to_count["Mul"], 0); - - for (const auto& node : graph.Nodes()) { - if (node.OpType() == "LayerNormalization") { - const auto& attrs = node.GetAttributes(); - auto it = attrs.find("epsilon"); - ASSERT_NE(it, attrs.end()); - EXPECT_FLOAT_EQ(it->second.f(), 1e-5f); - } - } + EXPECT_EQ(op_to_count["LayerNormalization"], 0); + EXPECT_EQ(op_to_count["ReduceMean"], 2); + EXPECT_EQ(op_to_count["Sub"], 1); + EXPECT_EQ(op_to_count["Pow"], 1); + EXPECT_EQ(op_to_count["Sqrt"], 1); + EXPECT_EQ(op_to_count["Div"], 1); + EXPECT_EQ(op_to_count["Mul"], 1); } // Test: SimplifiedLayerNormFusion gracefully handles a zero-element epsilon initializer. @@ -2773,25 +2974,15 @@ TEST_F(GraphTransformationTests, SimplifiedLayerNormFusion_ZeroElementEpsilon) { ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level1)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); - // Verify the fusion fired and produced a SimplifiedLayerNormalization with - // the default epsilon (1e-5f), per the zero-element fallback in - // layer_norm_fusion.cc. + // A non-scalar epsilon cannot be represented by the fused float attribute, + // so the original graph must remain unchanged. std::map op_to_count = CountOpsInGraph(graph); - EXPECT_EQ(op_to_count["SimplifiedLayerNormalization"], 1); - EXPECT_EQ(op_to_count["Pow"], 0); - EXPECT_EQ(op_to_count["ReduceMean"], 0); - EXPECT_EQ(op_to_count["Sqrt"], 0); - EXPECT_EQ(op_to_count["Div"], 0); - EXPECT_EQ(op_to_count["Mul"], 0); - - for (const auto& node : graph.Nodes()) { - if (node.OpType() == "SimplifiedLayerNormalization") { - const auto& attrs = node.GetAttributes(); - auto it = attrs.find("epsilon"); - ASSERT_NE(it, attrs.end()); - EXPECT_FLOAT_EQ(it->second.f(), 1e-5f); - } - } + EXPECT_EQ(op_to_count["SimplifiedLayerNormalization"], 0); + EXPECT_EQ(op_to_count["Pow"], 1); + EXPECT_EQ(op_to_count["ReduceMean"], 1); + EXPECT_EQ(op_to_count["Sqrt"], 1); + EXPECT_EQ(op_to_count["Div"], 1); + EXPECT_EQ(op_to_count["Mul"], 1); } #endif diff --git a/onnxruntime/test/optimizer/layout_transformation_potentially_added_ops_test.cc b/onnxruntime/test/optimizer/layout_transformation_potentially_added_ops_test.cc index c092c660633b1..c9fa32d569cf7 100644 --- a/onnxruntime/test/optimizer/layout_transformation_potentially_added_ops_test.cc +++ b/onnxruntime/test/optimizer/layout_transformation_potentially_added_ops_test.cc @@ -27,9 +27,15 @@ TEST(LayoutTransformationPotentiallyAddedOpsTests, OpsHaveLatestVersions) { if (prev_domain != it->domain || prev_op_type != it->op_type) { const auto* schema = schema_registry->GetSchema(std::string{it->op_type}, INT_MAX, std::string{it->domain}); ASSERT_NE(schema, nullptr); - EXPECT_EQ(schema->SinceVersion(), it->since_version) - << "A new version for op " << it->op_type << " (" << schema->SinceVersion() - << ") is available. Please update kLayoutTransformationPotentiallyAddedOps to include it."; + const bool is_staged_opset_28_qdq = + it->domain == kOnnxDomain && schema->SinceVersion() == 28 && + (it->op_type == "DequantizeLinear" || it->op_type == "QuantizeLinear"); + if (!is_staged_opset_28_qdq) { + EXPECT_EQ(schema->SinceVersion(), it->since_version) + << "A new version for op " << it->op_type << " (" << schema->SinceVersion() + << ") is available. Please update kLayoutTransformationPotentiallyAddedOps to include it."; + } + // TODO(#32372): Remove this exception when the opset-28 Q/DQ kernels land. prev_domain = it->domain; prev_op_type = it->op_type; } diff --git a/onnxruntime/test/optimizer/nchwc_optimizer_test.cc b/onnxruntime/test/optimizer/nchwc_optimizer_test.cc index 6cd73d5fc1d9e..ee686fad94ed7 100644 --- a/onnxruntime/test/optimizer/nchwc_optimizer_test.cc +++ b/onnxruntime/test/optimizer/nchwc_optimizer_test.cc @@ -238,6 +238,63 @@ void NchwcOptimizerTester(const std::function& bu #ifndef DISABLE_CONTRIB_OPS +TEST(NchwcOptimizerTests, ConvWithZeroOutputChannelsIsNotTransformed) { + for (bool no_bias : {false, true}) { + SCOPED_TRACE(no_bias ? "without bias" : "with bias"); + auto build_test_case = [&](NchwcTestHelper& helper) { + auto* input_arg = helper.MakeInput({1, 8, 5, 5}); + helper.AddConvNode(input_arg, helper.MakeOutput(), {0, 8, 1, 1}, no_bias); + }; + + auto check_nchwc_graph = [](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + EXPECT_EQ(op_to_count["Conv"], 1); + EXPECT_EQ(op_to_count["com.microsoft.nchwc.Conv"], 0); + EXPECT_EQ(op_to_count["com.microsoft.nchwc.ReorderInput"], 0); + EXPECT_EQ(op_to_count["com.microsoft.nchwc.ReorderOutput"], 0); + }; + + NchwcOptimizerTester(build_test_case, check_nchwc_graph, 18); + } +} + +TEST(NchwcOptimizerTests, ConvWithZeroInputChannelsIsNotTransformed) { + if (MlasNchwcGetBlockSize() <= 1) { + GTEST_SKIP() << "NCHWc is not supported on this platform."; + } + + for (bool no_bias : {false, true}) { + SCOPED_TRACE(no_bias ? "without bias" : "with bias"); + std::unordered_map domain_to_version{ + {kOnnxDomain, 18}, {kMSDomain, 1}, {kMSNchwcDomain, 1}}; + Model model("nchwc_zero_input_channels", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, {}, DefaultLoggingManager().DefaultLogger()); + Graph& graph = model.MainGraph(); + NchwcTestHelper helper(graph); + + auto* input_arg = helper.MakeInput({1, 0, 5, 5}); + auto* weights_arg = helper.MakeInitializer({8, 0, 1, 1}, {}); + std::vector input_args{input_arg, weights_arg}; + if (!no_bias) { + input_args.push_back(helper.Make1DInitializer({1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f})); + } + Node& conv = helper.AddNode("Conv", input_args, {helper.MakeOutput()}); + ASSERT_STATUS_OK(graph.Resolve()); + conv.SetExecutionProviderType(kCpuExecutionProvider); + + NchwcTransformer transformer; + bool modified = false; + ASSERT_STATUS_OK(transformer.Apply(graph, modified, DefaultLoggingManager().DefaultLogger())); + EXPECT_FALSE(modified); + + auto op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["Conv"], 1); + EXPECT_EQ(op_to_count["com.microsoft.nchwc.Conv"], 0); + EXPECT_EQ(op_to_count["com.microsoft.nchwc.ReorderInput"], 0); + EXPECT_EQ(op_to_count["com.microsoft.nchwc.ReorderOutput"], 0); + } +} + TEST(NchwcOptimizerTests, BinaryInputsWithUnexpectedRankAreNotTransformed) { if (MlasNchwcGetBlockSize() <= 1) { return; @@ -435,6 +492,47 @@ TEST(NchwcOptimizerTests, ConvNchwcHardSigmoidNotHardSwish) { NchwcOptimizerTester(build_test_case, check_nchwc_graph); } + // Keep the ragged filter sets test as well. + { + // The NCHWc convolution algorithm splits the output channel blocks into filter + // sets that are processed one kernel invocation at a time, and those sets are + // also the unit of work handed to each thread. When the block count is not a + // multiple of the filter set size the sets are unevenly sized, so the per-set + // filter count and the offset of each set into the output, filter and bias + // buffers all have to be derived rather than assumed. Get either wrong and + // channels are written to the wrong place. + // + // The channel counts below leave a remainder for both the 8 and 16 wide NCHWc + // block sizes, so the ragged path is exercised whichever one the platform + // selects: at block size 16 the block counts are 3, 5, 6 and 7, and at block + // size 8 they are 6, 10, 12 and 14. + auto test_case = [&](int64_t output_channels, int64_t kernel_size) { + auto build_test_case = [&](NchwcTestHelper& helper) { + auto* input_arg = helper.MakeInput({1, 64, 28, 28}); + auto* output_arg = helper.MakeOutput(); + + helper.AddConvNode(input_arg, output_arg, {output_channels, 64, kernel_size, kernel_size}); + }; + + auto check_nchwc_graph = [&](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + EXPECT_EQ(op_to_count["com.microsoft.nchwc.Conv"], 1); + }; + + NchwcOptimizerTester(build_test_case, check_nchwc_graph); + }; + + // Ragged (non-multiple) channel counts. + test_case(103, 3); + test_case(173, 3); + + // Both the general and the pointwise algorithms share this filter set logic. + for (int64_t output_channels : {48, 80, 96, 112}) { + test_case(output_channels, 3); + test_case(output_channels, 1); + } + } + // Sub-case 2: Mul(x, HardSigmoid(x)) diamond but with mismatched alpha (0.2, // not 1/6). TryFuseNchwcHardSwish must reject it at the alpha/beta guard so // the HardSigmoid is NOT rewritten as a HardSwish activation. diff --git a/onnxruntime/test/optimizer/qdq_transformer_test.cc b/onnxruntime/test/optimizer/qdq_transformer_test.cc index de9090d905ca8..e6139363ade79 100644 --- a/onnxruntime/test/optimizer/qdq_transformer_test.cc +++ b/onnxruntime/test/optimizer/qdq_transformer_test.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include "core/common/inlined_containers_fwd.h" #include "core/common/span_utils.h" @@ -4107,6 +4108,36 @@ TEST(QDQTransformerTests, Softmax_U8U8) { QDQTransformerSoftmaxTests(); } +TEST(QDQTransformerTests, SoftmaxSaturatesOutputWithoutOverflow) { + for (int opset : {12, 13}) { + for (int64_t width : {1, 2}) { + for (float output_scale : {1.0f / 256.0f, std::ldexp(1.0f, -25)}) { + SCOPED_TRACE(testing::Message() << "opset=" << opset << " width=" << width + << " output_scale=" << output_scale); + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({1, width}, std::vector(static_cast(width), 127)); + auto* dq_output = builder.MakeIntermediate(); + builder.AddDequantizeLinearNode(input, 0.105f, 127, dq_output); + auto* softmax_output = builder.MakeIntermediate(); + builder.AddNode("Softmax", {dq_output}, {softmax_output}).AddAttribute("axis", int64_t{1}); + auto* q_output = builder.MakeIntermediate(); + builder.AddQuantizeLinearNode(softmax_output, output_scale, 0, q_output); + builder.AddDequantizeLinearNode(q_output, 1.0f / 256.0f, 0, builder.MakeOutput()); + }; + auto check_graph = [](InferenceSessionWrapper& session) { + const auto counts = CountOpsInGraph(session.GetGraph()); + EXPECT_EQ(counts.at("com.microsoft.QLinearSoftmax"), 1); + EXPECT_EQ(counts.count("Softmax"), 0); + }; + TransformerTester(build_test_case, check_graph, + TransformerLevel::Level1, TransformerLevel::Level2, + opset, 0.0f, 0.0f, + std::make_unique(QDQIsInt8Allowed())); + } + } + } +} + #endif // !defined(DISABLE_CONTRIB_OPS) TEST(QDQTransformerTests, QDQPropagation_QBackward) { @@ -5690,7 +5721,10 @@ TEST(QDQTransformerTests, QDQPropagation_GH11605_Opset13) { // test removal of Q->DQ pairs by QDQFinalCleanupTransformer TEST(QDQTransformerTests, QDQFinalCleanupTransformerReportsIntentionalRemoval) { auto& logger = DefaultLoggingManager().DefaultLogger(); - Model model("QDQFinalCleanupRemovalTester", false, logger); + std::unordered_map domain_to_version; + domain_to_version[kOnnxDomain] = 25; + Model model("QDQFinalCleanupRemovalTester", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {}, logger); Graph& graph = model.MainGraph(); ModelTestBuilder builder(graph); diff --git a/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc new file mode 100644 index 0000000000000..e04fd6357fc22 --- /dev/null +++ b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc @@ -0,0 +1,110 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct CloseWin32Handle { + void operator()(HANDLE handle) const { CloseHandle(handle); } +}; +using UniqueHandle = std::unique_ptr; + +int Fail(const char* message) { + std::fprintf(stderr, "%s (Win32 error: %lu)\n", message, GetLastError()); + return EXIT_FAILURE; +} + +int RunChild(const wchar_t* library_path, std::wstring_view mode, HANDLE verified) { + if (!SetEnvironmentVariableW(L"ORT_DISABLE_TELEMETRY", L"1")) { + return Fail("Could not disable network telemetry"); + } + const HMODULE library = LoadLibraryW(library_path); + if (!library) { + return Fail("LoadLibraryW failed"); + } + using RunTest = void (*)(bool, bool, HANDLE); + const auto run_test = reinterpret_cast(GetProcAddress(library, "OrtTestTelemetryShutdown")); + if (!run_test) { + return Fail("GetProcAddress failed"); + } + const bool process_exit = mode == L"process_exit"; + run_test(process_exit, mode == L"normal_destruction", verified); + const DWORD expected_signal = mode == L"normal_destruction" ? WAIT_OBJECT_0 : WAIT_TIMEOUT; + if (WaitForSingleObject(verified, 0) != expected_signal) { + return Fail("Telemetry destruction happened at the wrong lifecycle stage"); + } + if (process_exit) { + ExitProcess(EXIT_SUCCESS); + } + if (!FreeLibrary(library)) { + return Fail("FreeLibrary failed"); + } + return EXIT_SUCCESS; +} + +} // namespace + +int wmain(int argc, wchar_t** argv) { + if (argc == 5 && std::wstring_view(argv[1]) == L"--child") { + const auto verified = reinterpret_cast(static_cast(std::wcstoull(argv[4], nullptr, 10))); + return RunChild(argv[2], argv[3], verified); + } + if (argc != 3) { + return Fail("Expected test DLL path and lifecycle mode"); + } + const std::wstring_view mode(argv[2]); + if (mode != L"process_exit" && mode != L"normal_destruction" && mode != L"explicit_unload") { + return Fail("Unknown lifecycle mode"); + } + + // The parent checks a signal from the DLL destructor, so an early child exit + // cannot pass without actually exercising and verifying telemetry teardown. + SECURITY_ATTRIBUTES attributes{sizeof(SECURITY_ATTRIBUTES), nullptr, TRUE}; + UniqueHandle verified(CreateEventW(&attributes, TRUE, FALSE, nullptr)); + if (!verified) { + return Fail("CreateEventW failed"); + } + std::wstring executable(32768, L'\0'); + const DWORD length = GetModuleFileNameW(nullptr, executable.data(), static_cast(executable.size())); + if (length == 0 || length >= executable.size()) { + return Fail("GetModuleFileNameW failed"); + } + executable.resize(length); + std::wstring command = L"\"" + executable + L"\" --child \"" + argv[1] + L"\" " + argv[2] + L" " + + std::to_wstring(reinterpret_cast(verified.get())); + STARTUPINFOW startup{}; + startup.cb = sizeof(startup); + PROCESS_INFORMATION info{}; + if (!CreateProcessW(executable.c_str(), command.data(), nullptr, nullptr, TRUE, CREATE_NO_WINDOW, + nullptr, nullptr, &startup, &info)) { + return Fail("CreateProcessW failed"); + } + UniqueHandle process(info.hProcess); + UniqueHandle thread(info.hThread); + if (WaitForSingleObject(process.get(), 10000) != WAIT_OBJECT_0) { + TerminateProcess(process.get(), EXIT_FAILURE); + WaitForSingleObject(process.get(), 2000); + return Fail("Child process did not exit within 10 seconds"); + } + DWORD exit_code = EXIT_FAILURE; + if (!GetExitCodeProcess(process.get(), &exit_code)) { + return Fail("GetExitCodeProcess failed"); + } + if (exit_code != EXIT_SUCCESS) { + std::fprintf(stderr, "Child process failed with exit code %lu\n", exit_code); + return EXIT_FAILURE; + } + if (WaitForSingleObject(verified.get(), 0) != WAIT_OBJECT_0) { + return Fail("Child did not verify telemetry destruction"); + } + return EXIT_SUCCESS; +} diff --git a/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc new file mode 100644 index 0000000000000..0b7b15235ff48 --- /dev/null +++ b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc @@ -0,0 +1,123 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include +#include +#include +#include + +#include +#include + +#include "core/platform/telemetry_1ds.h" +#include "core/platform/telemetry_1ds_platform.h" + +namespace onnxruntime { +namespace { + +void Check(bool condition, const char* message) { + if (!condition) { + std::fprintf(stderr, "%s (Win32 error: %lu)\n", message, GetLastError()); + std::fflush(stderr); + // Also fail reliably from a DLL process-detach destructor. + TerminateProcess(GetCurrentProcess(), EXIT_FAILURE); + std::abort(); + } +} + +} // namespace + +class OneDsTelemetryTestPeer { + public: + static void Install(Microsoft::Applications::Events::ILogManager& manager) { + Check(OneDsTelemetry::global_register_count_ == 1, "Expected exactly one telemetry instance"); + Check(!OneDsTelemetry::log_manager_ && !OneDsTelemetry::config_, "Real uploader must be disabled"); + OneDsTelemetry::config_ = std::make_unique(); + OneDsTelemetry::log_manager_ = &manager; + } + + static bool IsClean() { + return OneDsTelemetry::global_register_count_ == 0 && + !OneDsTelemetry::log_manager_ && !OneDsTelemetry::config_; + } +}; + +namespace { + +class MockLogManager : public Microsoft::Applications::Events::NullLogManager { + public: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MockLogManager); + + explicit MockLogManager(bool pending) : callback_(CreateEventW(nullptr, TRUE, !pending, nullptr)) { + Check(callback_ != nullptr, "CreateEventW failed"); + } + + ~MockLogManager() override { CloseHandle(callback_); } + + Microsoft::Applications::Events::status_t Flush() override { + ++flush_count; + return Microsoft::Applications::Events::STATUS_SUCCESS; + } + + void FlushAndTeardown() override { + Check(flush_count == 1, "Expected Flush before FlushAndTeardown"); + ++teardown_count; + // A pending callback cannot arrive after ExitProcess has killed the other + // threads. Bound the mock's wait so a regression does not leave a stuck DLL. + Check(WaitForSingleObject(callback_, 2000) == WAIT_OBJECT_0, "Timed out draining the pending uploader"); + } + + int flush_count = 0; + int teardown_count = 0; + + private: + HANDLE callback_; +}; + +class TelemetryLifetime { + public: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TelemetryLifetime); + + TelemetryLifetime(bool process_exit, HANDLE verified) + : manager_(process_exit), process_exit_(process_exit), verified_(verified) { + telemetry_.emplace(); + OneDsTelemetryTestPeer::Install(manager_); + } + + ~TelemetryLifetime() { + if (telemetry_) { + Destroy(); + } + } + + void Destroy() { + Check(telemetry_internal::IsDllShutdownInProgress() == process_exit_, "Unexpected DLL shutdown state"); + telemetry_.reset(); + const int expected_calls = process_exit_ ? 0 : 1; + Check(manager_.flush_count == expected_calls, "Unexpected Flush count"); + Check(manager_.teardown_count == expected_calls, "Unexpected FlushAndTeardown count"); + if (!process_exit_) { + Check(OneDsTelemetryTestPeer::IsClean(), "Telemetry state was not cleaned up"); + } + Check(SetEvent(verified_) != FALSE, "Could not signal verified destruction"); + } + + private: + MockLogManager manager_; + std::optional telemetry_; + bool process_exit_; + HANDLE verified_; +}; + +} // namespace +} // namespace onnxruntime + +extern "C" void OrtTestTelemetryShutdown(bool process_exit, bool destroy_now, HANDLE verified) { + // A DLL-owned static exercises CRT destruction during real Windows shutdown. + static onnxruntime::TelemetryLifetime lifetime(process_exit, verified); + if (destroy_now) { + lifetime.Destroy(); + } +} diff --git a/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def new file mode 100644 index 0000000000000..517b8cf037088 --- /dev/null +++ b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def @@ -0,0 +1,3 @@ +LIBRARY onnxruntime_telemetry_shutdown_test_library +EXPORTS + OrtTestTelemetryShutdown diff --git a/onnxruntime/test/providers/cpu/activation/activation_op_test.cc b/onnxruntime/test/providers/cpu/activation/activation_op_test.cc index c438cb00b401b..5284517a05717 100644 --- a/onnxruntime/test/providers/cpu/activation/activation_op_test.cc +++ b/onnxruntime/test/providers/cpu/activation/activation_op_test.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include +#include #include "activation_op_test.h" #include @@ -707,6 +708,82 @@ TEST_F(ActivationOpTest, Softplus_Opset22) { } #if defined(USE_CUDA) +template +static void CheckCudaActivationSmallPositiveTails(const char* op_name, const Reference& reference) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP unavailable"; + std::vector input; + std::vector expected; + for (float value : {-100.0f, -80.0f, -40.0f, -30.0f, -20.0f, -16.0f, -14.0f, + -12.0f, -8.0f, -5.0f, -1.0f, -0.125f, 0.0f, 0.125f, + 1.0f, 5.0f, 8.0f, 14.0f, 20.0f, 40.0f, 100.0f}) { + input.emplace_back(value); + const double result = reference(static_cast(value)); + if constexpr (std::is_same_v) { + expected.emplace_back(result); + } else { + expected.emplace_back(static_cast(result)); + } + } + OpTester tester(op_name, 22); + const std::vector dimensions{static_cast(input.size())}; + tester.AddInput("X", dimensions, input); + tester.AddOutput("Y", dimensions, expected); + tester.SetOutputAbsErr("Y", 0.0f); + tester.SetOutputRelErr("Y", std::is_same_v ? 2e-14f : 2e-6f); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +template +static void CheckCudaSoftplusSmallPositiveTails() { + CheckCudaActivationSmallPositiveTails("Softplus", [](double value) { + return std::max(value, 0.0) + std::log1p(std::exp(-std::abs(value))); + }); +} + +template +static void CheckCudaSigmoidSmallPositiveTails() { + CheckCudaActivationSmallPositiveTails("Sigmoid", [](double value) { + const double exponent = std::exp(-std::abs(value)); + return value > 0 ? 1.0 / (1.0 + exponent) : exponent / (1.0 + exponent); + }); +} + +TEST_F(ActivationOpTest, CudaSigmoidFloatSmallPositiveTails) { + CheckCudaSigmoidSmallPositiveTails(); +} + +TEST_F(ActivationOpTest, CudaSigmoidDoubleSmallPositiveTails) { + CheckCudaSigmoidSmallPositiveTails(); +} + +TEST_F(ActivationOpTest, CudaSigmoidFloat16SmallPositiveTails) { + CheckCudaSigmoidSmallPositiveTails(); +} + +TEST_F(ActivationOpTest, CudaSigmoidBFloat16SmallPositiveTails) { + CheckCudaSigmoidSmallPositiveTails(); +} + +TEST_F(ActivationOpTest, CudaSoftplusFloatSmallPositiveTails) { + CheckCudaSoftplusSmallPositiveTails(); +} + +TEST_F(ActivationOpTest, CudaSoftplusDoubleSmallPositiveTails) { + CheckCudaSoftplusSmallPositiveTails(); +} + +TEST_F(ActivationOpTest, CudaSoftplusFloat16SmallPositiveTails) { + CheckCudaSoftplusSmallPositiveTails(); +} + +TEST_F(ActivationOpTest, CudaSoftplusBFloat16SmallPositiveTails) { + CheckCudaSoftplusSmallPositiveTails(); +} + TEST_F(ActivationOpTest, Softplus_bfloat16_Opset22) { if (!HasCudaEnvironment(530)) { LOGS_DEFAULT(WARNING) << "Hardware does NOT support BF16"; diff --git a/onnxruntime/test/providers/cpu/controlflow/scan_test.cc b/onnxruntime/test/providers/cpu/controlflow/scan_test.cc index 4583b772bb30c..c35a45d590eda 100644 --- a/onnxruntime/test/providers/cpu/controlflow/scan_test.cc +++ b/onnxruntime/test/providers/cpu/controlflow/scan_test.cc @@ -1061,7 +1061,8 @@ TEST(Scan8, NumScanInputsExceedsVariadicInputs) { test.AddOutput("scan_output_2", {1, 2, 1}, {0.f, 0.f}); test.AddOutput("scan_output_3", {1, 2, 1}, {0.f, 0.f}); - test.Run(OpTester::ExpectResult::kExpectFailure, "Invalid 'num_scan_inputs' of 10. Value must be between 1 and 3", + test.Run(OpTester::ExpectResult::kExpectFailure, + "num_scan_inputs (10) plus the sequence_lens input cannot exceed the number of Scan inputs (4).", options.excluded_provider_types); } @@ -1165,7 +1166,7 @@ TEST(Scan9, NumScanInputsExceedsVariadicInputs) { // inferencing error wrapping a shape-inference error, rather than the kernel-level message produced // by the opset 8 test above. test.Run(OpTester::ExpectResult::kExpectFailure, - "[ShapeInferenceError] Mismatch between number of inferred and declared dimensions", + "num_scan_inputs (10) cannot exceed the number of Scan inputs (3).", options.excluded_provider_types); } #endif // !defined(ORT_NO_EXCEPTIONS) diff --git a/onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc b/onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc index f3cc72249b764..703c4fc04b74b 100644 --- a/onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc +++ b/onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc @@ -15,6 +15,12 @@ #include #include #include +#include + +#ifdef USE_DML +#include "core/providers/dml/dml_session_options_config_keys.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#endif namespace onnxruntime { namespace test { @@ -4315,6 +4321,7 @@ TEST(MathOpTest, ErfCheckMultiThreadDataChunking) { } constexpr int ModOp_ver = 10; +constexpr int ModOp_ver28 = 28; TEST(ModOpTest, Fmod_float_mixed_sign) { OpTester test("Mod", ModOp_ver); @@ -4351,6 +4358,92 @@ TEST(ModOpTest, Fmod_float16_mixed_sign) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kQnnExecutionProvider}); } +template +void TestFloorModFloatingPointOpset28(const std::vector& x, + const std::vector& y, + const std::vector& expected) { + OpTester test("Mod", 28); + test.AddAttribute("fmod", 0); + const std::vector shape{static_cast(x.size())}; + test.AddInput("X", shape, x); + test.AddInput("Y", shape, y); + test.AddOutput("Z", shape, expected); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kQnnExecutionProvider, kOpenVINOExecutionProvider}); +} + +TEST(ModOpTest, FloorMod_float_mixed_sign_opset28) { + TestFloorModFloatingPointOpset28( + {-4.3f, 7.2f, 5.0f, 4.3f, -7.2f, 8.0f}, + {2.1f, -3.4f, 8.0f, -2.1f, 3.4f, 5.0f}, + {2.0f, -3.0f, 5.0f, -2.0f, 3.0f, 3.0f}); +} + +TEST(ModOpTest, FloorMod_double_mixed_sign_opset28) { + TestFloorModFloatingPointOpset28( + {-4.3, 7.2, 5.0, 4.3, -7.2, 8.0}, + {2.1, -3.4, 8.0, -2.1, 3.4, 5.0}, + {2.0, -3.0, 5.0, -2.0, 3.0, 3.0}); +} + +TEST(ModOpTest, FloorMod_float16_mixed_sign_opset28) { + TestFloorModFloatingPointOpset28( + MakeMLFloat16({-4.3f, 7.2f, 5.0f, 4.3f, -7.2f, 8.0f}), + MakeMLFloat16({2.1f, -3.4f, 8.0f, -2.1f, 3.4f, 5.0f}), + MakeMLFloat16({2.0f, -3.0f, 5.0f, -2.0f, 3.0f, 3.0f})); +} + +TEST(ModOpTest, FloorMod_float_special_values_opset28) { + constexpr auto infinity = std::numeric_limits::infinity(); + constexpr auto nan = std::numeric_limits::quiet_NaN(); + TestFloorModFloatingPointOpset28( + {6.0f, -6.0f, 5.0f, -5.0f, infinity, 1.0f, nan, 1.0f}, + {-3.0f, 3.0f, infinity, -infinity, 2.0f, 0.0f, 2.0f, nan}, + {-0.0f, 0.0f, 5.0f, -5.0f, nan, nan, nan, nan}); +} + +TEST(ModOpTest, FloorMod_opset28_provider_fallback) { + InlinedVector> providers; + providers.push_back(DefaultCpuExecutionProvider()); +#ifdef USE_CANN + OrtCANNProviderOptions cann_options{}; + cann_options.enable_cann_graph = 1; + auto cann_factory = CannProviderFactoryCreator::Create(&cann_options); + ASSERT_NE(cann_factory, nullptr); + providers.push_back(cann_factory->CreateProvider()); +#endif + providers.push_back(DefaultMIGraphXExecutionProvider()); + providers.push_back(DefaultOpenVINOExecutionProvider()); + for (auto& provider : providers) { + if (!provider) { + continue; + } + SCOPED_TRACE(provider->Type()); + const std::vector shape{4}; + const float infinity = std::numeric_limits::infinity(); + const std::vector x{-5.0f, 5.0f, -5.0f, 5.0f}; + const std::vector y{3.0f, -3.0f, infinity, -infinity}; + OpTester test("Mod", 28); + test.AddAttribute("fmod", 0); + test.AddInput("X", shape, x); + test.AddInput("Y", shape, y); + test.AddOutput("Z", shape, {1.0f, -1.0f, infinity, -infinity}); + test.ConfigEp(DefaultCpuExecutionProvider()).RunWithConfig(); + std::string model_data; + auto& model = test.BuildModel(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + model.ToProto().SerializeToString(&model_data); + EPVerificationParams params; + params.ep_node_assignment = provider->Type() == kCpuExecutionProvider + ? ExpectedEPNodeAssignment::All + : ExpectedEPNodeAssignment::None; + RunAndVerifyOutputsWithEP( + AsByteSpan(model_data.data(), model_data.size()), CurrentTestName(), std::move(provider), + {{"X", CreateInputOrtValueOnCPU(shape, x)}, {"Y", CreateInputOrtValueOnCPU(shape, y)}}, + params); + } +} + #if defined(USE_CUDA) TEST(ModOpTest, Fmod_bfloat16_mixed_sign) { OpTester test("Mod", 13); @@ -4367,6 +4460,176 @@ TEST(ModOpTest, Fmod_bfloat16_mixed_sign) { } #endif +TEST(ModOpTest, FloorMod_float_mixed_sign) { + OpTester test("Mod", ModOp_ver28); + test.AddInput("X", {6}, {-4.3f, 7.2f, 5.0f, 4.3f, -7.2f, 8.0f}); + test.AddInput("Y", {6}, {2.1f, -3.4f, 8.0f, -2.1f, 3.4f, 5.0f}); + test.AddOutput("Z", {6}, {2.0f, -3.0f, 5.0f, -2.0f, 3.0f, 3.0f}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider}); +} + +TEST(ModOpTest, FloorMod_double_mixed_sign) { + OpTester test("Mod", ModOp_ver28); + test.AddInput("X", {6}, {-4.3, 7.2, 5.0, 4.3, -7.2, 8.0}); + test.AddInput("Y", {6}, {2.1, -3.4, 8.0, -2.1, 3.4, 5.0}); + test.AddOutput("Z", {6}, {2.0, -3.0, 5.0, -2.0, 3.0, 3.0}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider}); +} + +TEST(ModOpTest, FloorMod_float16_mixed_sign) { + OpTester test("Mod", ModOp_ver28); + test.AddInput("X", {6}, MakeMLFloat16({-4.3f, 7.2f, 5.0f, 4.3f, -7.2f, 8.0f})); + test.AddInput("Y", {6}, MakeMLFloat16({2.1f, -3.4f, 8.0f, -2.1f, 3.4f, 5.0f})); + test.AddOutput("Z", {6}, MakeMLFloat16({2.0f, -3.0f, 5.0f, -2.0f, 3.0f, 3.0f})); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider}); +} + +TEST(ModOpTest, FloorMod_float_edge_cases) { + const float nan = std::numeric_limits::quiet_NaN(); + const float inf = std::numeric_limits::infinity(); + OpTester test("Mod", ModOp_ver28); + test.AddInput("X", {14}, + {0.0f, -0.0f, 0.0f, -0.0f, -3.0f, 3.0f, -1.0f, 1.0f, + inf, -inf, 1.0f, 1.0f, nan, 1.0f}); + test.AddInput("Y", {14}, + {-2.0f, 2.0f, 2.0f, -2.0f, inf, inf, -inf, -inf, + 2.0f, 2.0f, 0.0f, -0.0f, 2.0f, nan}); + test.AddOutput("Z", {14}, + {-0.0f, 0.0f, 0.0f, -0.0f, inf, 3.0f, -1.0f, -inf, + nan, nan, nan, nan, nan, nan}); + test.SetCustomOutputVerifier([](const std::vector& fetches, + const std::string& /*provider_type*/) { + ASSERT_EQ(fetches.size(), 1u); + ASSERT_TRUE(fetches[0].IsTensor()); + const float* output = fetches[0].Get().Data(); + + EXPECT_EQ(output[0], 0.0f); + EXPECT_TRUE(std::signbit(output[0])); + EXPECT_EQ(output[1], 0.0f); + EXPECT_FALSE(std::signbit(output[1])); + EXPECT_EQ(output[2], 0.0f); + EXPECT_FALSE(std::signbit(output[2])); + EXPECT_EQ(output[3], 0.0f); + EXPECT_TRUE(std::signbit(output[3])); + EXPECT_EQ(output[4], std::numeric_limits::infinity()); + EXPECT_EQ(output[5], 3.0f); + EXPECT_EQ(output[6], -1.0f); + EXPECT_EQ(output[7], -std::numeric_limits::infinity()); + for (size_t i = 8; i < 14; ++i) { + EXPECT_TRUE(std::isnan(output[i])) << "output index " << i; + } + }); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider}); +} + +#ifdef USE_DML +template +void TestDirectMLFloorModBroadcast() { + const float nan = std::numeric_limits::quiet_NaN(); + const float inf = std::numeric_limits::infinity(); + const std::vector x{T(0.0f), T(-0.0f), T(6.0f), T(-6.0f), T(3.0f), + T(-3.0f), T(inf), T(-inf), T(nan), T(1.0f)}; + const std::vector y{T(2.0f), T(-2.0f), T(inf), T(-inf), T(0.0f), T(-0.0f), T(nan)}; + const std::vector expected{ + 0.0f, -0.0f, 0.0f, -0.0f, nan, nan, nan, + 0.0f, -0.0f, 0.0f, -0.0f, nan, nan, nan, + 0.0f, -0.0f, 6.0f, -inf, nan, nan, nan, + 0.0f, -0.0f, inf, -6.0f, nan, nan, nan, + 1.0f, -1.0f, 3.0f, -inf, nan, nan, nan, + 1.0f, -1.0f, inf, -3.0f, nan, nan, nan, + nan, nan, nan, nan, nan, nan, nan, + nan, nan, nan, nan, nan, nan, nan, + nan, nan, nan, nan, nan, nan, nan, + 1.0f, -1.0f, 1.0f, -inf, nan, nan, nan}; + std::vector typed_expected; + typed_expected.reserve(expected.size()); + for (float value : expected) { + typed_expected.emplace_back(value); + } + for (bool disable_fusion : {false, true}) { + SCOPED_TRACE(disable_fusion ? "standalone kernel" : "fused graph"); + OpTester test("Mod", 28); + test.AddAttribute("fmod", 0); + test.AddInput("X", {10, 1}, x); + test.AddInput("Y", {1, 7}, y); + test.AddOutput("Z", {10, 7}, typed_expected); + test.SetCustomOutputVerifier([&expected](const std::vector& fetches, + const std::string& provider_type) { + ASSERT_EQ(provider_type, kDmlExecutionProvider); + ASSERT_EQ(fetches.size(), 1u); + const auto& tensor = fetches[0].Get(); + ASSERT_EQ(tensor.Shape().Size(), static_cast(expected.size())); + const T* output = tensor.Data(); + for (size_t i = 0; i < expected.size(); ++i) { + const float actual = static_cast(output[i]); + if (std::isnan(expected[i])) { + EXPECT_TRUE(std::isnan(actual)) << "output index " << i; + } else { + EXPECT_EQ(actual, expected[i]) << "output index " << i; + EXPECT_EQ(std::signbit(actual), std::signbit(expected[i])) << "output index " << i; + } + } + }); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry( + kOrtSessionOptionsConfigDisableDmlGraphFusion, disable_fusion ? "1" : "0")); + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(options).ConfigEp(DefaultDmlExecutionProvider()).RunWithConfig(); + } +} + +TEST(ModOpTest, FloorMod_DirectML_float_broadcast) { + TestDirectMLFloorModBroadcast(); +} + +TEST(ModOpTest, FloorMod_DirectML_float16_broadcast) { + TestDirectMLFloorModBroadcast(); +} +#endif + +TEST(ModOpTest, Signed_integer_overflow_case) { + OpTester test("Mod", ModOp_ver28); + test.AddInput("X", {2}, {std::numeric_limits::min(), 7}); + test.AddInput("Y", {2}, {-1, -3}); + test.AddOutput("Z", {2}, {0, -2}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider}); +} + +TEST(ModOpTest, Signed_integer_overflow_case_fmod) { + OpTester test("Mod", ModOp_ver28); + test.AddAttribute("fmod", 1); + test.AddInput("X", {2}, {std::numeric_limits::min(), 7}); + test.AddInput("Y", {2}, {-1, -3}); + test.AddOutput("Z", {2}, {0, 1}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider}); +} + +TEST(ModOpTest, Int64_fmod_precision) { + OpTester test("Mod", ModOp_ver28); + test.AddAttribute("fmod", 1); + test.AddInput("X", {1}, {9007199254740993LL}); + test.AddInput("Y", {1}, {10}); + test.AddOutput("Z", {1}, {3}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider}); +} + +TEST(ModOpTest, UInt64_fmod_precision) { + OpTester test("Mod", ModOp_ver28); + test.AddAttribute("fmod", 1); + test.AddInput("X", {1}, {18446744073709551613ULL}); + test.AddInput("Y", {1}, {10}); + test.AddOutput("Z", {1}, {3}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kQnnExecutionProvider}); +} + TEST(ModOpTest, Int8_mixed_sign) { OpTester test("Mod", ModOp_ver); test.AddInput("X", {6}, {-4, 7, 5, 4, -7, 8}); @@ -4652,6 +4915,125 @@ TEST(BitShiftOpTest, BroadcastXRight_Uint8) { test.Run(); } +template +void TestSignedBitShiftOpset28() { + constexpr T bit_width = static_cast(std::numeric_limits>::digits); + + { + OpTester test("BitShift", 28); + test.AddAttribute("direction", "LEFT"); + test.AddInput("X", {6}, {T{-8}, T{-3}, T{1}, std::numeric_limits::max(), T{-2}, T{2}}); + test.AddInput("Y", {6}, {T{1}, T{2}, T{3}, T{1}, T{-1}, bit_width}); + test.AddOutput("Z", {6}, {T{-16}, T{-12}, T{8}, T{-2}, T{0}, T{0}}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider}); + } + + { + OpTester test("BitShift", 28); + test.AddAttribute("direction", "RIGHT"); + test.AddInput("X", {6}, {T{-8}, T{-3}, T{1}, T{-2}, T{-2}, T{2}}); + test.AddInput("Y", {6}, {T{1}, T{2}, T{3}, T{-1}, bit_width, bit_width}); + test.AddOutput("Z", {6}, {T{-4}, T{-1}, T{0}, T{-1}, T{-1}, T{0}}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider}); + } + + { + OpTester test("BitShift", 28); + test.AddAttribute("direction", "RIGHT"); + test.AddInput("X", {1}, {T{-8}}); + test.AddInput("Y", {4}, {T{1}, T{2}, bit_width, T{-1}}); + test.AddOutput("Z", {4}, {T{-4}, T{-2}, T{-1}, T{-1}}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider}); + } + + { + OpTester test("BitShift", 28); + test.AddAttribute("direction", "LEFT"); + test.AddInput("X", {2}, {std::numeric_limits::max(), T{-8}}); + test.AddInput("Y", {1}, {T{1}}); + test.AddOutput("Z", {2}, {T{-2}, T{-16}}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider}); + } +} + +TEST(BitShiftOpTest, SignedInt8Opset28) { + TestSignedBitShiftOpset28(); +} + +TEST(BitShiftOpTest, SignedInt16Opset28) { + TestSignedBitShiftOpset28(); +} + +TEST(BitShiftOpTest, SignedInt32Opset28) { + TestSignedBitShiftOpset28(); +} + +TEST(BitShiftOpTest, SignedInt64Opset28) { + TestSignedBitShiftOpset28(); +} + +template +void TestBitShiftOpset28ProviderFallback(const std::vector& x, const std::vector& y, + const std::vector& expected) { + const std::vector shape{static_cast(x.size())}; + OpTester test("BitShift", 28); + test.AddAttribute("direction", "RIGHT"); + test.AddInput("X", shape, x); + test.AddInput("Y", shape, y); + test.AddOutput("Z", shape, expected); + test.ConfigEp(DefaultCpuExecutionProvider()).RunWithConfig(); + auto provider = DefaultOpenVINOExecutionProvider(); + if (!provider) { + return; + } + + auto& model = test.BuildModel(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + std::string model_data; + ASSERT_TRUE(model.ToProto().SerializeToString(&model_data)); + EPVerificationParams params; + params.ep_node_assignment = ExpectedEPNodeAssignment::None; + RunAndVerifyOutputsWithEP( + AsByteSpan(model_data.data(), model_data.size()), CurrentTestName(), std::move(provider), + {{"X", CreateInputOrtValueOnCPU(shape, x)}, {"Y", CreateInputOrtValueOnCPU(shape, y)}}, + params); +} + +TEST(BitShiftOpTest, Opset28ProviderFallback) { + TestBitShiftOpset28ProviderFallback({-8, -3, -2, -2}, {1, 2, -1, 32}, {-4, -1, -1, -1}); + TestBitShiftOpset28ProviderFallback({8, 255, 1, 42}, {1, 32, 33, 64}, {4, 0, 0, 0}); +} + +template +void TestUnsignedBitShiftRegistration(int opset) { + OpTester test("BitShift", opset); + test.AddAttribute("direction", "LEFT"); + test.AddInput("X", {2}, {T{1}, T{3}}); + test.AddInput("Y", {2}, {T{1}, T{2}}); + test.AddOutput("Z", {2}, {T{2}, T{12}}); + test.Run(); +} + +TEST(BitShiftOpTest, Uint16Opset11) { + TestUnsignedBitShiftRegistration(11); +} + +TEST(BitShiftOpTest, UnsignedTypesOpset28) { + TestUnsignedBitShiftRegistration(28); + TestUnsignedBitShiftRegistration(28); + TestUnsignedBitShiftRegistration(28); + TestUnsignedBitShiftRegistration(28); +} + +TEST(BitShiftOpTest, RightShiftByBitWidth_Uint64Opset28) { + OpTester test("BitShift", 28); + test.AddAttribute("direction", "RIGHT"); + test.AddInput("X", {4}, {1000, 255, 1, 42}); + test.AddInput("Y", {4}, {64, 64, 64, 64}); + test.AddOutput("Z", {4}, {0, 0, 0, 0}); + test.Run(); +} + // Test that shift amounts >= bit width produce 0 (not undefined behavior). // DirectML EP has the same hardware-level shift masking behavior, so skip these tests for DML. TEST(BitShiftOpTest, RightShiftByBitWidth_Uint64) { diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc index 87005c72b4c45..8b28addc18fc8 100644 --- a/onnxruntime/test/providers/cpu/math/gemm_test.cc +++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc @@ -343,7 +343,6 @@ TEST(GemmOpTest, GemmNoTrans_f16) { } } -// Only CUDA and CoreML kernels have float 16 support TEST(GemmOpTest, GemmTransB_f16) { #ifdef USE_CUDA int min_cuda_architecture = 530; @@ -365,6 +364,26 @@ TEST(GemmOpTest, GemmTransB_f16) { std::vector f_B(12); ConvertFloatToMLFloat16(A.data(), f_A.data(), 8); ConvertFloatToMLFloat16(B.data(), f_B.data(), 12); + { + // Missing C makes beta zero, which is eligible for the MLAS path. + std::vector f_Y(6); + std::vector Y{7.1f, -6.2f, -19.0f, -7.1f, 6.2f, 19.0f}; + ConvertFloatToMLFloat16(Y.data(), f_Y.data(), 6); + + OpTester test("Gemm", 13); + + test.AddAttribute("transA", (int64_t)0); + test.AddAttribute("transB", (int64_t)1); + test.AddAttribute("alpha", 1.0f); + test.AddAttribute("beta", 0.0f); + test.AddInput("A", {2, 4}, f_A); + test.AddInput("B", {3, 4}, f_B, true); + test.AddOutput("Y", {2, 3}, f_Y); + test.SetOutputTolerance(0.005f); + test.ConfigExcludeEps({kTensorrtExecutionProvider}) // TensorRT: fp16 is not supported + .Config(run_with_tunable_op) + .RunWithConfig(); + } { // bias is a scalar and transB is True std::vector f_Y(6); diff --git a/onnxruntime/test/providers/cpu/math/matmul_test.cc b/onnxruntime/test/providers/cpu/math/matmul_test.cc index c3393c622cb8a..48ac5da24ff1f 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc @@ -3,10 +3,13 @@ #include "gtest/gtest.h" +#include "core/providers/cpu/math/matmul_helper.h" #include "test/providers/provider_test_utils.h" #include "test/common/dnnl_op_test_utils.h" #include "test/common/cuda_op_test_utils.h" #include "test/common/tensor_op_test_utils.h" +#include "test/test_environment.h" +#include "test/unittest_util/graph_transform_test_builder.h" #include "default_providers.h" namespace onnxruntime { @@ -500,6 +503,124 @@ TEST(MathOpTest, MatMulZeroKDoubleType) { RunMatMulZeroKTest(); } +template +void RunMatMulZeroKBroadcastCpuTest(bool transpose_a_with_vector_b = false) { + const std::vector a_dims = transpose_a_with_vector_b ? std::vector{2, 0, 3} + : std::vector{2, 1, 3, 0}; + const std::vector b_dims = transpose_a_with_vector_b ? std::vector{0} + : std::vector{1, 4, 0, 5}; + const std::vector output_dims = transpose_a_with_vector_b ? std::vector{2, 3} + : std::vector{2, 4, 3, 5}; + const size_t output_size = static_cast(TensorShape(output_dims).Size()); + const auto* domain = transpose_a_with_vector_b ? kMSDomain : kOnnxDomain; + const int opset = transpose_a_with_vector_b ? 1 : 13; + Model model("zero_k_cpu_broadcast", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{domain, opset}}, {}, + DefaultLoggingManager().DefaultLogger()); + ModelTestBuilder builder(model.MainGraph()); + auto* a = builder.MakeInput(a_dims, std::vector{}); + auto* b = builder.MakeInput(b_dims, std::vector{}); + auto* y = builder.MakeOutput(output_dims); + auto& node = builder.AddNode(transpose_a_with_vector_b ? "FusedMatMul" : "MatMul", {a, b}, {y}, domain); + if (transpose_a_with_vector_b) { + node.AddAttribute("transA", int64_t{1}); + } + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + const std::string bytes = model.ToProto().SerializeAsString(); + SessionOptions options; + options.intra_op_param.thread_pool_size = 1; + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultCpuExecutionProvider())); + ASSERT_STATUS_OK(session.Load(bytes.data(), static_cast(bytes.size()))); + ASSERT_STATUS_OK(session.Initialize()); + const auto allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + OrtValue output; + CreateMLValue(allocator, output_dims, std::vector(output_size, T{-99}), &output); + std::vector outputs{output}; + for (int run = 0; run < 2; ++run) { + // Poison every batch so untouched memory cannot appear correctly zeroed. + std::fill_n(output.GetMutable()->MutableData(), output_size, T{-99}); + ASSERT_STATUS_OK(session.Run(RunOptions{}, builder.feeds_, builder.output_names_, &outputs)); + const auto& tensor = outputs[0].Get(); + ASSERT_EQ(tensor.Shape(), TensorShape(output_dims)); + ASSERT_EQ(tensor.DataRaw(), output.Get().DataRaw()); + for (T value : tensor.DataAsSpan()) { + EXPECT_EQ(value, T{}) << "run=" << run; + } + } +} + +TEST(MathOpTest, MatMulZeroKBroadcastCpuFloat) { + RunMatMulZeroKBroadcastCpuTest(); +} + +TEST(MathOpTest, MatMulZeroKBroadcastCpuDouble) { + RunMatMulZeroKBroadcastCpuTest(); +} + +#ifndef DISABLE_CONTRIB_OPS +TEST(MathOpTest, FusedMatMulZeroKTransposeARightVectorCpuFloat) { + RunMatMulZeroKBroadcastCpuTest(true); +} + +TEST(MathOpTest, FusedMatMulZeroKTransposeARightVectorCpuDouble) { + RunMatMulZeroKBroadcastCpuTest(true); +} +#endif + +TEST(MathOpTest, MatMulHelperZeroBatchBroadcast) { + const auto check = [](const TensorShape& a, const TensorShape& b, const TensorShape& expected) { + for (bool fill_offsets : {false, true}) { + MatMulComputeHelper helper; + ASSERT_STATUS_OK(helper.Compute(a, b, false, false, false, false, fill_offsets)); + ASSERT_EQ(helper.OutputShape(), expected); + ASSERT_EQ(helper.OutputShape().Size(), 0); + if (!fill_offsets) { + helper.FillOffsets(); + } + EXPECT_TRUE(helper.LeftOffsets().empty()); + EXPECT_TRUE(helper.RightOffsets().empty()); + EXPECT_TRUE(helper.OutputOffsets().empty()); + } + }; + ASSERT_NO_FATAL_FAILURE(check(TensorShape({1, 3, 0}), TensorShape({0, 0, 5}), TensorShape({0, 3, 5}))); + ASSERT_NO_FATAL_FAILURE(check(TensorShape({1, 3, 4}), TensorShape({0, 4, 5}), TensorShape({0, 3, 5}))); + ASSERT_NO_FATAL_FAILURE(check(TensorShape({2, 1, 3, 0}), TensorShape({1, 0, 0, 5}), TensorShape({2, 0, 3, 5}))); + ASSERT_NO_FATAL_FAILURE(check(TensorShape({0, 1, 3, 0}), TensorShape({1, 2, 0, 5}), TensorShape({0, 2, 3, 5}))); +} + +TEST(MathOpTest, MatMulHelperRejectsIncompatibleZeroBatch) { + MatMulComputeHelper helper; + EXPECT_FALSE(helper.Compute(TensorShape({2, 3, 0}), TensorShape({0, 0, 5})).IsOK()); + EXPECT_FALSE(helper.Compute(TensorShape({0, 1, 3, 0}), TensorShape({2, 2, 0, 5})).IsOK()); +} + +template +void RunMatMulZeroBatchCpuTest(std::initializer_list a_dims, + std::initializer_list b_dims, + std::initializer_list y_dims) { + OpTester test("MatMul", 13); + test.AddInput("A", a_dims, std::vector(static_cast(TensorShape(a_dims).Size()), T{1})); + test.AddInput("B", b_dims, std::vector(static_cast(TensorShape(b_dims).Size()), T{1})); + test.AddOutput("Y", y_dims, std::vector{}); + test.ConfigEp(DefaultCpuExecutionProvider()).RunWithConfig(); +} + +TEST(MathOpTest, MatMulZeroBatchCpuFloat) { + RunMatMulZeroBatchCpuTest({1, 3, 0}, {0, 0, 5}, {0, 3, 5}); + RunMatMulZeroBatchCpuTest({1, 3, 4}, {0, 4, 5}, {0, 3, 5}); + RunMatMulZeroBatchCpuTest({2, 1, 3, 0}, {1, 0, 0, 5}, {2, 0, 3, 5}); + RunMatMulZeroBatchCpuTest({0, 1, 3, 0}, {1, 2, 0, 5}, {0, 2, 3, 5}); +} + +TEST(MathOpTest, MatMulZeroBatchCpuDouble) { + RunMatMulZeroBatchCpuTest({1, 3, 0}, {0, 0, 5}, {0, 3, 5}); + RunMatMulZeroBatchCpuTest({1, 3, 4}, {0, 4, 5}, {0, 3, 5}); + RunMatMulZeroBatchCpuTest({2, 1, 3, 0}, {1, 0, 0, 5}, {2, 0, 3, 5}); + RunMatMulZeroBatchCpuTest({0, 1, 3, 0}, {1, 2, 0, 5}, {0, 2, 3, 5}); +} + #if defined(USE_CUDA) || defined(USE_COREML) || defined(USE_XNNPACK) TEST(MathOpTest, MatMul_Float16) { #ifdef USE_CUDA diff --git a/onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc b/onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc index 24a41983207b3..778ea25317c5b 100644 --- a/onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc @@ -13,6 +13,7 @@ namespace test { template static void CheckCudaRmsNormTypes() { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP unavailable"; + using ScaleType = std::conditional_t, double, float>; const auto as_double = [](auto value) { if constexpr (std::is_same_v || std::is_same_v) { return static_cast(value.ToFloat()); @@ -26,17 +27,19 @@ static void CheckCudaRmsNormTypes() { std::vector scale(input.size()); std::vector expected(input.size()); for (int row = 0; row < 3; ++row) { - double sum = 0.0; + float sum = 0.0f; for (int channel = 0; channel < width; ++channel) { const int index = row * width + channel; input[index] = InputType(static_cast((channel % 7) - 3) * 0.25f + row); scale[index] = OutputType(1.0f + row * 0.5f); - sum += as_double(input[index]) * as_double(input[index]); + const float x = static_cast(as_double(input[index])); + sum += x * x; } for (int channel = 0; channel < width; ++channel) { const int index = row * width + channel; - expected[index] = OutputType(static_cast(as_double(input[index]) * as_double(scale[index]) / - std::sqrt(sum / width + 1e-5))); + const InputType normalized(static_cast(as_double(input[index])) / std::sqrt(sum / width + 1e-5f)); + expected[index] = OutputType(static_cast(as_double(normalized)) * + static_cast(as_double(scale[index]))); } } OpTester tester("RMSNormalization", 23); @@ -68,6 +71,123 @@ TEST(RMSNormalizationOpTest, CudaRegisteredTypesAndBroadcast) { CheckCudaRmsNormTypes(); } +TEST(RMSNormalizationOpTest, CudaDoubleUsesFloatStashAndDoubleScale) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP unavailable"; + for (int width : {7, 32, 513, 2688}) { + SCOPED_TRACE(testing::Message() << "width=" << width); + std::vector input(width); + std::vector scale(width); + for (int channel = 0; channel < width; ++channel) { + const double sign = channel % 2 == 0 ? 1.0 : -1.0; + input[channel] = 1.0 + sign * std::ldexp(1.0, -28); + scale[channel] = 1.0 + sign * std::ldexp(1.0, -40); + } + // Every input rounds to 1 in the default FP32 stash; the scale remains double. + OpTester tester("RMSNormalization", 23); + tester.AddAttribute("epsilon", 0.0f); + tester.AddInput("X", {1, width}, input); + tester.AddInput("scale", {width}, scale); + tester.AddOutput("Y", {1, width}, scale); + tester.SetOutputAbsErr("Y", 0.0f); + tester.SetOutputRelErr("Y", 0.0f); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + +TEST(RMSNormalizationOpTest, CudaDoubleExplicitDoubleStash) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP unavailable"; + const std::vector input{1.0 + std::ldexp(1.0, -28), 1.0 - std::ldexp(1.0, -28)}; + const std::vector scale{1.0 + std::ldexp(1.0, -40), 1.0 - std::ldexp(1.0, -40)}; + OpTester tester("RMSNormalization", 23); + tester.AddAttribute("epsilon", 0.0f); + tester.AddAttribute("stash_type", 11); + tester.AddInput("X", {1, 2}, input); + tester.AddInput("scale", {2}, scale); + tester.AddOutput("Y", {1, 2}, {input[0] * scale[0], input[1] * scale[1]}); + tester.SetOutputAbsErr("Y", 0.0f); + tester.SetOutputRelErr("Y", 0.0f); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(RMSNormalizationOpTest, CudaRejectsUnsupportedStashTypes) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP unavailable"; + for (int64_t stash_type : {10, 11, 16}) { + SCOPED_TRACE(testing::Message() << "stash_type=" << stash_type); + OpTester tester("RMSNormalization", 23); + tester.AddAttribute("stash_type", stash_type); + tester.AddInput("X", {1, 2}, {1.0f, 1.0f}); + tester.AddInput("scale", {2}, {1.0f, 1.0f}); + tester.AddOutput("Y", {1, 2}, {1.0f, 1.0f}); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectFailure, + "CUDA RMSNormalization supports stash_type=1 (FLOAT)", {}, nullptr, &providers); + } +} + +template +static void CheckCudaRmsNormStandardizationRounding() { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP unavailable"; + for (int width : {7, 8, 31, 32, 2688}) { + SCOPED_TRACE(testing::Message() << "width=" << width); + std::vector input(2 * width); + std::vector scale(input.size()); + std::vector expected(input.size()); + for (int row = 0; row < 2; ++row) { + float sum = 0.0f; + for (int channel = 0; channel < width; ++channel) { + const int index = row * width + channel; + input[index] = InputType((channel % 2 == 0 ? 1.25f : 3.5f) + row * 0.25f); + scale[index] = OutputType(1.015625f + (channel % 5) * 0.125f); + const float x = input[index].ToFloat(); + sum += x * x; + } + const float rms = std::sqrt(sum / width + 1e-5f); + for (int channel = 0; channel < width; ++channel) { + const int index = row * width + channel; + // ONNX stash_type=1 rounds standardization to X's type before scaling. + const InputType normalized(input[index].ToFloat() / rms); + expected[index] = OutputType(normalized.ToFloat() * (1.015625f + (channel % 5) * 0.125f)); + } + } + OpTester tester("RMSNormalization", 23); + tester.AddAttribute("epsilon", 1e-5f); + tester.AddAttribute("stash_type", 1); + tester.AddInput("X", {1, 2, width}, input); + tester.AddInput("scale", {1, 2, width}, scale); + tester.AddOutput("Y", {1, 2, width}, expected); + tester.SetOutputAbsErr("Y", 0.0f); + tester.SetOutputRelErr("Y", 0.0f); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + +TEST(RMSNormalizationOpTest, CudaBFloat16RoundsStandardizationBeforeScale) { + CheckCudaRmsNormStandardizationRounding(); +} + +TEST(RMSNormalizationOpTest, CudaFloat16RoundsStandardizationBeforeScale) { + CheckCudaRmsNormStandardizationRounding(); +} + +TEST(RMSNormalizationOpTest, CudaFloat16ToFloatRoundsStandardizationBeforeScale) { + CheckCudaRmsNormStandardizationRounding(); +} + // Large, tiny, and zero FP16 inputs must stay finite in scalar and half2 output paths. TEST(RMSNormalizationOpTest, CudaFloat16FiniteRange) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP unavailable"; diff --git a/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc b/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc index 05bd7aef1e440..07eca67baed4a 100644 --- a/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc @@ -155,6 +155,9 @@ std::vector CastedValues(gsl::span src) { } struct CastNonStringTester { + int opset{21}; + bool cuda_only{false}; + template void operator()(const std::pair&) { SCOPED_TRACE( @@ -177,7 +180,8 @@ struct CastNonStringTester { auto output_span = gsl::make_span(output_buffer.get(), size); CastSpan(input_span, output_span); - TestCastOp(input_span, output_span, shape.AsShapeVector()); + TestCastOp(input_span, output_span, shape.AsShapeVector(), + OpTester::ExpectResult::kExpectSuccess, "", opset, Saturate::None, cuda_only); } }; @@ -194,6 +198,25 @@ TEST(CastOpTest, NonStringTypes) { CastNonStringTester{}); } +TEST(CastOpTest, NonStringTypes_Opset25To28) { + for (int opset : {25, 27, 28}) { + SCOPED_TRACE(opset); + boost::mp11::mp_for_each>( + CastNonStringTester{opset}); + } +} + +#if defined(USE_CUDA) +TEST(CastOpTest, NonStringTypes_Opset28_Cuda) { + if (!HasCudaEnvironment(530)) { + GTEST_SKIP() << "CUDA Cast tests require a CUDA device with compute capability >= 5.3."; + } + + boost::mp11::mp_for_each>( + CastNonStringTester{28, true}); +} +#endif + TEST(CastOpTest, FromString) { const std::vector shape{2, 2, 2}; const std::vector string_data = {"-inf", "+INF", "0.9767611", "0.28280696", diff --git a/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc b/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc index 495fea5735b32..36c3f61a59630 100644 --- a/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc +++ b/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc @@ -1,6 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include + #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" @@ -166,8 +169,13 @@ TEST(CompressTest, Compress_3dims_neg_axis) { // inputs to {0, 1}) and asserts the output is sized by truthiness. TEST(CompressTest, Compress_cuda_non_canonical_bool_condition) { // Build: output = Compress(input, condition, axis=0) - auto model = std::make_unique("compress_non_canonical_bool", false, - DefaultLoggingManager().DefaultLogger()); + const std::unordered_map domain_to_version{{kOnnxDomain, 27}}; + auto model = std::make_unique( + "compress_non_canonical_bool", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, std::vector{}, DefaultLoggingManager().DefaultLogger()); + ASSERT_EQ(model->ToProto().opset_import_size(), 1); + EXPECT_EQ(model->ToProto().opset_import(0).domain(), kOnnxDomain); + EXPECT_EQ(model->ToProto().opset_import(0).version(), 27); onnxruntime::Graph& graph = model->MainGraph(); ONNX_NAMESPACE::TypeProto tensor_float; diff --git a/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc index ed449b96eacca..07f9addbfd83b 100644 --- a/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc @@ -3,6 +3,7 @@ #include +#include "default_providers.h" #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" #include "test/common/trt_op_test_utils.h" @@ -43,6 +44,40 @@ TEST(OneHotOpTest, DefaultAxis_float_float_float /*indices, output, depth*/) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", ExcludeTrtOnA100()); } +TEST(OneHotOpTest, FractionalFloatIndicesAreCastToInt64) { + for (int opset : {9, 11}) { + OpTester test("OneHot", opset); + test.AddInput("indices", {6}, {0.9f, 1.8f, -1.2f, -4.3f, -5.3f, 2.0f}); + test.AddInput("depth", {}, {4.0f}); + test.AddInput("values", {2}, {2.0f, 3.0f}); + test.AddOutput("output", {6, 4}, + {3.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 3.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 3.0f, + 3.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 3.0f, 2.0f}); + test.ConfigEp(DefaultCpuExecutionProvider()).RunWithConfig(); + } +} + +TEST(OneHotOpTest, NonFiniteAndUnrepresentableFloatIndicesAreOutOfRange) { + OpTester test("OneHot", 11); + test.AddInput("indices", {4}, + {std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN(), + std::numeric_limits::max()}); + test.AddInput("depth", {}, {4.0f}); + test.AddInput("values", {2}, {2.0f, 3.0f}); + test.AddOutput("output", {4, 4}, + {2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f}); + test.ConfigEp(DefaultCpuExecutionProvider()).RunWithConfig(); +} + TEST(OneHotOpTest, DefaultAxis_int64_int32_float /*indices, output, depth*/) { OpTester test("OneHot", 11); test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); diff --git a/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc b/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc index 6404418e75b09..c981c1c70b8a2 100644 --- a/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc @@ -2267,6 +2267,15 @@ void DequantizeLinearOp21BlockedTest_Float8_Succeed(std::vector&& dims, test.Run(BaseTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); } +TEST(DequantizeLinearOp21BlockedTest, SignedInt_UseZeroPoint_OneDimension) { + std::vector x_scale{0.5f, 2.0f}; + std::vector zero_point{0, 10}; + std::vector x{10, 20, 30, 40, 50, 60, 70, 80}; + std::vector y{5.0f, 10.0f, 15.0f, 20.0f, 80.0f, 100.0f, 120.0f, 140.0f}; + + DequantizeLinearOp21BlockedTest_Int_Succeed({8}, 0, 4, x, x_scale, zero_point, y); +} + TEST(DequantizeLinearOp21BlockedTest, SignedInt_NoZeroPoint_FirstAxis) { std::vector x_scale{-2.0, -4.0, 3.5, 1.0, 2.0, 4.0, -3.5, -1.0}; std::vector zero_point; diff --git a/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc b/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc index ac7d896258603..7df530e0ff078 100644 --- a/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc @@ -9,6 +9,9 @@ #include "core/framework/allocator.h" #include "core/providers/cpu/tensor/space_depth_ops.h" #include "core/mlas/inc/mlas.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/default_providers.h" +#include "test/util/include/inference_session_wrapper.h" namespace onnxruntime { namespace test { @@ -121,6 +124,88 @@ TEST(TensorOpTest, SpaceToDepthTest_1_double) { test.AddOutput("output", {N, C * blocksize * blocksize, H / blocksize, W / blocksize}, result); test.Run(); } + +TEST(TensorOpTest, SpaceToDepthOpset28ModesAndProviderRouting) { + const std::vector input{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; + const std::vector dcr{0, 2, 8, 10, 1, 3, 9, 11, 4, 6, 12, 14, 5, 7, 13, 15}; + const std::vector crd{0, 2, 1, 3, 4, 6, 5, 7, 8, 10, 9, 11, 12, 14, 13, 15}; + for (int opset : {13, 27, 28}) { + for (const std::string mode : {"", "DCR", "CRD"}) { + if (opset < 28 && !mode.empty()) { + continue; + } + SCOPED_TRACE(opset); + SCOPED_TRACE(mode); + Model model("SpaceToDepthRouting", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, opset}}, {}, + DefaultLoggingManager().DefaultLogger()); + ModelTestBuilder builder(model.MainGraph()); + auto* x = builder.MakeInput({1, 2, 2, 4}, input); + auto* z = builder.MakeOutput(); + auto& node = builder.AddNode("SpaceToDepth", {x}, {z}); + node.AddAttribute("blocksize", int64_t{2}); + if (!mode.empty()) { + node.AddAttribute("mode", mode); + } + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + std::string model_data; + ASSERT_TRUE(model.ToProto().SerializeToString(&model_data)); + + InlinedVector> providers; + providers.push_back(DefaultCpuExecutionProvider()); + providers.push_back(DefaultCudaExecutionProvider()); +#ifdef ENABLE_CUDA_NHWC_OPS + providers.push_back(DefaultCudaNHWCExecutionProvider()); +#endif + providers.push_back(DefaultDmlExecutionProvider()); + providers.push_back(DefaultOpenVINOExecutionProvider()); + providers.push_back(DefaultMIGraphXExecutionProvider()); +#ifdef USE_CANN + OrtCANNProviderOptions cann_options{}; + cann_options.enable_cann_graph = 1; + auto cann_factory = CannProviderFactoryCreator::Create(&cann_options); + ASSERT_NE(cann_factory, nullptr); + providers.push_back(cann_factory->CreateProvider()); +#endif + for (auto& provider : providers) { + if (!provider) { + continue; + } + const std::string provider_type = provider->Type(); + SCOPED_TRACE(provider_type); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + options.enable_mem_pattern = false; + options.execution_mode = ExecutionMode::ORT_SEQUENTIAL; + InferenceSessionWrapper session{options, GetEnvironment()}; + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(provider))); + ASSERT_STATUS_OK(session.Load(model_data.data(), static_cast(model_data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + bool has_space_to_depth = false; + for (const auto& resolved_node : session.GetGraph().Nodes()) { + has_space_to_depth |= resolved_node.OpType() == "SpaceToDepth"; + } + if (opset == 28 && (mode == "CRD" || provider_type == kCpuExecutionProvider || + provider_type == kCudaExecutionProvider)) { + EXPECT_FALSE(has_space_to_depth); + } else if (opset < 28 && provider_type == kCpuExecutionProvider) { + EXPECT_TRUE(has_space_to_depth); + } + + std::vector fetches; + ASSERT_STATUS_OK(session.Run(RunOptions{}, builder.feeds_, builder.output_names_, &fetches)); + ASSERT_EQ(fetches.size(), 1u); + const auto& tensor = fetches[0].Get(); + EXPECT_EQ(tensor.Shape(), TensorShape({1, 8, 1, 2})); + const auto& expected = mode == "CRD" ? crd : dcr; + const auto* actual = tensor.Data(); + EXPECT_EQ(std::vector(actual, actual + tensor.Shape().Size()), expected); + } + } + } +} + TEST(TensorOpTest, SpaceToDepthTest_2) { OpTester test("SpaceToDepth"); constexpr int64_t blocksize = 3; diff --git a/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc b/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc index 622039c20f5a2..ab2c385685ea5 100644 --- a/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc @@ -57,6 +57,7 @@ using contrib::cuda::SetGroupQueryAttentionWorkspaceRequirements; constexpr int kMath = static_cast(AttentionBackend::MATH); constexpr int kFlash = static_cast(AttentionBackend::FLASH_ATTENTION); +constexpr int kCudnn = static_cast(AttentionBackend::CUDNN_FLASH_ATTENTION); WorkspaceInputShape Known(std::initializer_list dims) { return WorkspaceInputShape::PresentWithShape(TensorShape{TensorShapeVector{dims}}); @@ -174,7 +175,8 @@ void SetValueInfo(ONNX_NAMESPACE::ValueInfoProto& value_info, } } -std::string BuildGroupQueryAttentionKernelModel() { +std::string BuildGroupQueryAttentionKernelModel(bool sliding_window_cache = true, + bool include_head_sink = true) { ONNX_NAMESPACE::ModelProto model; model.set_ir_version(ONNX_NAMESPACE::IR_VERSION); auto* onnx_opset = model.add_opset_import(); @@ -192,9 +194,10 @@ std::string BuildGroupQueryAttentionKernelModel() { node->set_op_type("GroupQueryAttention"); for (const char* input_name : {"query", "key", "value", "past_key", "past_value", "seqlens_k", - "total_sequence_length", "", "", "", "", "head_sink"}) { + "total_sequence_length", "", "", "", ""}) { node->add_input(input_name); } + node->add_input(include_head_sink ? "head_sink" : ""); node->add_output("output"); node->add_output("present_key"); node->add_output("present_value"); @@ -207,8 +210,10 @@ std::string BuildGroupQueryAttentionKernelModel() { }; add_int_attribute("num_heads", 8); add_int_attribute("kv_num_heads", 2); - add_int_attribute("local_window_size", 256); - add_int_attribute("sliding_window_cache", 1); + if (sliding_window_cache) { + add_int_attribute("local_window_size", 256); + add_int_attribute("sliding_window_cache", 1); + } constexpr int32_t kFloat16 = ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; constexpr int32_t kInt32 = ONNX_NAMESPACE::TensorProto_DataType_INT32; @@ -223,11 +228,13 @@ std::string BuildGroupQueryAttentionKernelModel() { SetValueInfo(*graph->add_output(), "present_key", kFloat16, {2, 2, 256, 64}); SetValueInfo(*graph->add_output(), "present_value", kFloat16, {2, 2, 256, 64}); - auto* head_sink = graph->add_initializer(); - head_sink->set_name("head_sink"); - head_sink->set_data_type(kFloat16); - head_sink->add_dims(8); - head_sink->mutable_raw_data()->assign(8 * sizeof(uint16_t), '\0'); + if (include_head_sink) { + auto* head_sink = graph->add_initializer(); + head_sink->set_name("head_sink"); + head_sink->set_data_type(kFloat16); + head_sink->add_dims(8); + head_sink->mutable_raw_data()->assign(8 * sizeof(uint16_t), '\0'); + } std::string bytes; model.SerializeToString(&bytes); @@ -248,6 +255,58 @@ bool HasCudaDevice() { return cudaGetDeviceCount(&device_count) == cudaSuccess && device_count > 0; } +void CheckNonWindowedStaticCacheCapacityBound(int kernel_selection, GQABackend backend) { + AttentionKernelOptions options; + options.InitializeOnce(kernel_selection, true); + auto config = Config(); + config.sliding_window_cache = false; + config.local_window_size = -1; + config.enable_xqa = false; + config.max_total_sequence_length = 512; + constexpr int64_t kPastCapacity = 1024; + const auto estimate = EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(/*sequence=*/1, /*head=*/64, kPastCapacity), + Device(), options); + ASSERT_TRUE(estimate.has_value()); + const bool use_mea = backend == GQABackend::MemoryEfficient; + ASSERT_TRUE(HasGQAReachableBackend( + estimate->sized_backends, + use_mea ? GQAReachableBackend::MemoryEfficient : GQAReachableBackend::Unfused)); + + GQAWorkspaceProblem problem; + problem.qkv_element_size = 2; + problem.cache_element_size = 2; + problem.batch_size = 2; + problem.sequence_length = 1; + problem.num_heads = 8; + problem.kv_num_heads = 2; + problem.head_size = 64; + problem.present_kv_cache_capacity = kPastCapacity; + problem.requires_separate_past_buffer = true; + problem.past_kv_cache_capacity = kPastCapacity; + GQAConcreteRoute route; + route.backend = backend; + route.preparation.preprocess_mode = + use_mea ? GQAPreprocessMode::MemoryEfficient : GQAPreprocessMode::Unfused; + route.unfused.total_sequence_length = config.max_total_sequence_length; + const auto runtime = GetGQACompleteWorkspaceRecipe(problem, route); + ASSERT_TRUE(runtime.status.IsOK()); + constexpr size_t kPastBytes = 2 * 2 * kPastCapacity * 64 * 2; + EXPECT_EQ(runtime.recipe.preparation.separate_past_bytes, kPastBytes); + if (use_mea) { + constexpr size_t kExpandedCacheBytes = 2 * 8 * kPastCapacity * 64 * 2; + EXPECT_EQ(runtime.recipe.memory_efficient.expanded_key_bytes, kExpandedCacheBytes); + EXPECT_EQ(runtime.recipe.memory_efficient.expanded_value_bytes, kExpandedCacheBytes); + } + EXPECT_GE(estimate->total_workspace_bytes, runtime.recipe.total_workspace_bytes); + + const auto smaller_cache = EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(/*sequence=*/1, /*head=*/64, /*capacity=*/512), + Device(), options); + ASSERT_TRUE(smaller_cache.has_value()); + EXPECT_GT(estimate->total_workspace_bytes, smaller_cache->total_workspace_bytes); +} + TEST(GroupQueryAttentionWorkspaceEstimateTest, ParsesPackedAndSeparateLayouts) { AttentionKernelOptions options; options.InitializeOnce(kMath, true); @@ -372,9 +431,10 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, ParsedEnvelopeReachesNodeConfigWi auto non_windowed_config = *gqa_config; non_windowed_config.sliding_window_cache = false; - EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( - non_windowed_config, SeparateShapes(), Device(), options) - .has_value()); + EXPECT_EQ(EstimateGroupQueryAttentionWorkspace( + non_windowed_config, SeparateShapes(), Device(), options) + .has_value(), + test_case.expected > 0); } }); } @@ -487,18 +547,151 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti } } -TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedTotalKvAndAliasingAreUnavailable) { +TEST(GroupQueryAttentionWorkspaceEstimateTest, CapacityAwarePartitioningRequiresProvenEstimate) { + if (!HasCudaDevice()) { + GTEST_SKIP() << "A CUDA device is required for the budget integration test."; + } + + ScopedEnvironmentVariables scoped_env_vars{{{"ORT_ENABLE_XQA", "0"}}}; + const std::string model_bytes = BuildGroupQueryAttentionKernelModel( + /*sliding_window_cache=*/false, /*include_head_sink=*/false); + + struct Case { + int kernel_selection; + const char* envelope; + bool expect_cuda; + }; + for (const auto& test_case : { + Case{kMath, nullptr, false}, + Case{kCudnn, "512", false}, + Case{kMath, "512", true}, + }) { + SCOPED_TRACE(testing::Message() + << "kernel=" << test_case.kernel_selection + << " envelope=" << (test_case.envelope == nullptr ? "" : test_case.envelope)); + SessionOptions session_options; + session_options.graph_optimization_level = TransformerLevel::Default; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsResourceCudaPartitioningSettings, "1048576,")); + if (test_case.envelope != nullptr) { + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength, test_case.envelope)); + } + + CUDAExecutionProviderInfo provider_info; + provider_info.sdpa_kernel = test_case.kernel_selection; + InferenceSessionWrapper session(session_options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider( + std::make_shared(provider_info))); + ASSERT_STATUS_OK(session.Load(model_bytes.data(), static_cast(model_bytes.size()))); + ASSERT_STATUS_OK(session.Initialize()); + + const Node* node = FindNodeByOpType(session.GetGraph(), "GroupQueryAttention"); + ASSERT_NE(node, nullptr); + EXPECT_EQ(node->GetExecutionProviderType() == kCudaExecutionProvider, + test_case.expect_cuda); + } +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedWithoutEnvelopeIsUnavailable) { AttentionKernelOptions options; options.InitializeOnce(kMath, true); auto config = Config(); config.sliding_window_cache = false; config.local_window_size = -1; + // total_sequence_length is a runtime scalar. Without the declared KV-length + // envelope the non-windowed estimate cannot bound the present length, so it + // declines -- the behavior before the envelope knob existed. EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( config, SeparateShapes(/*sequence=*/1, /*head=*/64, /*capacity=*/128), Device(), options) .has_value()); } +TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedEnvelopeBoundsWorkspace) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto config = Config(); + config.sliding_window_cache = false; + config.local_window_size = -1; + config.max_total_sequence_length = 512; + const auto estimate = EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(/*sequence=*/4, /*head=*/64, /*capacity=*/256), + Device(), options); + ASSERT_TRUE(estimate.has_value()); + EXPECT_GT(estimate->total_workspace_bytes, 0u); + + // The envelope is the KV length the estimate is taken over, so raising it + // cannot lower the bound. + auto larger = config; + larger.max_total_sequence_length = 1024; + const auto larger_estimate = EstimateGroupQueryAttentionWorkspace( + larger, SeparateShapes(/*sequence=*/4, /*head=*/64, /*capacity=*/256), + Device(), options); + ASSERT_TRUE(larger_estimate.has_value()); + EXPECT_GE(larger_estimate->total_workspace_bytes, estimate->total_workspace_bytes); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedStaticCacheCapacityBoundsPreservation) { + CheckNonWindowedStaticCacheCapacityBound(kMath, GQABackend::Unfused); +} + +#if USE_MEMORY_EFFICIENT_ATTENTION +TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedStaticCacheCapacityBoundsMeaExpansion) { + CheckNonWindowedStaticCacheCapacityBound( + static_cast(AttentionBackend::EFFICIENT_ATTENTION), GQABackend::MemoryEfficient); +} +#endif + +TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedEnvelopeCannotBeSmallerThanQueryBound) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto config = Config(); + config.sliding_window_cache = false; + config.local_window_size = -1; + config.max_total_sequence_length = 3; + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(/*sequence=*/4, /*head=*/64, /*capacity=*/1024), + Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedStaticCacheCapacityMustFitInt32) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto config = Config(); + config.sliding_window_cache = false; + config.local_window_size = -1; + config.max_total_sequence_length = 512; + const int64_t capacity = static_cast(std::numeric_limits::max()) + 1; + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(/*sequence=*/1, /*head=*/64, capacity), + Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, PartialAliasPreservationAddsFullPastCopy) { + auto bounds = Bounds(); + bounds.reachable_backends = GQAReachableBackend::Unfused; + const auto without = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(without.status.IsOK()); + EXPECT_GT(without.total_workspace_bytes, 0u); + + // Non-windowed partial aliasing copies the full past cache, which stacks on + // top of the selected route's workspace rather than replacing it. + bounds.account_partial_alias_preservation = true; + const auto with_copy = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(with_copy.status.IsOK()); + + const size_t expected_copy = + static_cast(bounds.batch_size_bound) * + static_cast(bounds.kv_num_heads) * + static_cast(bounds.present_kv_cache_capacity_bound) * + static_cast(bounds.head_size_bound) * bounds.cache_element_size; + EXPECT_EQ(with_copy.total_workspace_bytes, + without.total_workspace_bytes + expected_copy); +} + TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsCacheCapacityDifferentFromWindow) { AttentionKernelOptions options; options.InitializeOnce(kMath, true); diff --git a/onnxruntime/test/providers/dml_onnx_tensor_wrapper_test.cc b/onnxruntime/test/providers/dml_onnx_tensor_wrapper_test.cc index c9d0d190e6ff9..6855b309f3e30 100644 --- a/onnxruntime/test/providers/dml_onnx_tensor_wrapper_test.cc +++ b/onnxruntime/test/providers/dml_onnx_tensor_wrapper_test.cc @@ -59,10 +59,9 @@ enum class ValueAttributePayload { // IMLOperatorAttributes1::GetTensorAttribute, which wraps the raw AttributeProto tensor in an // OnnxTensorWrapper. // -// Attribute tensors are not graph initializers, so none of the framework-level guards apply to them: -// Graph::ConvertInitializersIntoOrtValues never sees them, and ONNX shape inference for -// ConstantOfShape only reads the attribute's data_type (to type the output), never its data. That -// leaves OnnxTensorWrapper as the only place the payload can be validated. +// Attribute tensors are not graph initializers, so Graph::ConvertInitializersIntoOrtValues never +// sees them. Older ONNX versions also did not validate byte-aligned raw_data sizes in check_tensor, +// leaving OnnxTensorWrapper as the only place the payload could be validated. std::string BuildConstantOfShapeModel(ValueAttributePayload payload) { ONNX_NAMESPACE::ModelProto model; model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); @@ -113,8 +112,8 @@ std::string BuildConstantOfShapeModel(ValueAttributePayload payload) { return serialized; } -// Loads |serializedModel| into a session that has only the DML EP registered, and returns the status -// of Initialize(). |skipped| is set when no DirectML device is available. +// Loads |serializedModel| into a session that has only the DML EP registered, and returns the first +// failure from Load() or Initialize(). |skipped| is set when no DirectML device is available. // // Constant folding is disabled so the ConstantOfShape node survives graph optimization and is handed // to the DML EP. Without this, the CPU EP would evaluate the node during Level1 optimization and the @@ -166,15 +165,19 @@ TEST(DmlOnnxTensorWrapperTest, TruncatedRawDataAttributeTensorIsRejected) { GTEST_SKIP() << "DirectML execution provider is not available on this machine."; } - ASSERT_FALSE(status.IsOK()) << "Session initialization should reject a tensor whose declared shape " - "exceeds the data it carries."; - - // The detailed message raised inside the constructor does not survive: OnnxTensorWrapper is - // constructed behind the IMLOperatorAttributes1 COM boundary, whose ORT_CATCH_RETURN reduces the - // exception to a bare HRESULT before it is rethrown. 0x80070057 is E_INVALIDARG, so this asserts the - // model was rejected as malformed rather than failing for an unrelated reason. - EXPECT_NE(status.ErrorMessage().find("80070057"), std::string::npos) - << "Unexpected failure reason: " << status.ErrorMessage(); + ASSERT_FALSE(status.IsOK()) << "Model loading or session initialization should reject a tensor whose declared " + "shape exceeds the data it carries."; + + // ONNX 1.23 rejects the malformed attribute during Graph::Resolve. With older ONNX versions it + // reaches OnnxTensorWrapper, where the detailed constructor exception is reduced to E_INVALIDARG + // (0x80070057) at the IMLOperatorAttributes1 COM boundary. + const std::string& errorMessage = status.ErrorMessage(); + EXPECT_TRUE( + errorMessage.find("80070057") != std::string::npos || + errorMessage.find( + "raw_data size (1 bytes) is too small for the declared shape and type (8 bytes required)") != + std::string::npos) + << "Unexpected failure reason: " << errorMessage; } // The constructor's typed-field branch has the same defect in principle: UnpackTensor sizes the buffer diff --git a/onnxruntime/test/providers/graph_capture_test.cc b/onnxruntime/test/providers/graph_capture_test.cc index 127efae7c5b0c..3b8dbde43e4eb 100644 --- a/onnxruntime/test/providers/graph_capture_test.cc +++ b/onnxruntime/test/providers/graph_capture_test.cc @@ -13,6 +13,7 @@ #include "core/graph/constants.h" #include "core/session/onnxruntime_cxx_api.h" #include "core/session/onnxruntime_run_options_config_keys.h" +#include "core/session/onnxruntime_session_options_config_keys.h" using namespace Ort; @@ -163,6 +164,154 @@ TEST(WebGpuDispatchBatchingTests, ProfilingDispatchCountStaysBoundedAcrossAlloca } } +class WebGpuScatterGraphCaptureTests : public testing::TestWithParam {}; + +void TestScatterReplayCopiesInput(const char* op_type, const std::vector& indices_shape, + bool enable_graph_capture, bool surrounding_dispatches = false) { + Env& env = *ort_env; + SessionOptions session_options; + session_options.DisableMemPattern(); + session_options.SetGraphOptimizationLevel(ORT_DISABLE_ALL); + session_options.SetLogSeverityLevel(1); + session_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"); + if (surrounding_dispatches) { + session_options.EnableProfiling(ORT_TSTR("webgpu_scatter_capture_test")); + } + ASSERT_NO_FATAL_FAILURE(AppendWebGpuEp( + env, session_options, {{"enableGraphCapture", enable_graph_capture ? "1" : "0"}})); + + const std::vector data_shape{3}; + const std::vector updates_shape{1}; + TensorTypeAndShapeInfo data_info(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, data_shape); + TensorTypeAndShapeInfo indices_info(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, indices_shape); + TensorTypeAndShapeInfo updates_info(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, updates_shape); + auto data_type = TypeInfo::CreateTensorInfo(data_info.GetConst()); + auto indices_type = TypeInfo::CreateTensorInfo(indices_info.GetConst()); + auto updates_type = TypeInfo::CreateTensorInfo(updates_info.GetConst()); + std::vector inputs; + inputs.emplace_back("X", data_type.GetConst()); + inputs.emplace_back("indices", indices_type.GetConst()); + inputs.emplace_back("updates", updates_type.GetConst()); + std::vector outputs; + outputs.emplace_back("Y", data_type.GetConst()); + if (surrounding_dispatches) { + outputs.emplace_back("T", data_type.GetConst()); + outputs.emplace_back("Z", data_type.GetConst()); + } + Graph graph; + graph.SetInputs(inputs); + graph.SetOutputs(outputs); + if (surrounding_dispatches) { + Node before("Relu", onnxruntime::kOnnxDomain, "before", {"X"}, {"T"}); + graph.AddNode(before); + } + Node scatter(op_type, onnxruntime::kOnnxDomain, "scatter", + {surrounding_dispatches ? "T" : "X", "indices", "updates"}, {"Y"}); + graph.AddNode(scatter); + if (surrounding_dispatches) { + Node after("Relu", onnxruntime::kOnnxDomain, "after", {"Y"}, {"Z"}); + graph.AddNode(after); + } + Model model(std::vector{{onnxruntime::kOnnxDomain, 18}}); + model.AddGraph(graph); + Session session(env, model, session_options); + + MemoryInfo gpu_mem_info("WebGPU_Buffer", OrtDeviceAllocator, 0, OrtMemTypeDefault); + Allocator gpu_allocator(session, gpu_mem_info); + auto gpu_x = Value::CreateTensor(gpu_allocator, data_shape.data(), data_shape.size()); + auto gpu_indices = Value::CreateTensor(gpu_allocator, indices_shape.data(), indices_shape.size()); + auto gpu_updates = Value::CreateTensor(gpu_allocator, updates_shape.data(), updates_shape.size()); + auto gpu_y = Value::CreateTensor(gpu_allocator, data_shape.data(), data_shape.size()); + Value gpu_t{nullptr}; + Value gpu_z{nullptr}; + // Separate bound output storage forces the copy path despite the kernel's MayInplace registration. + ASSERT_NE(gpu_x.GetTensorRawData(), gpu_y.GetTensorRawData()); + + IoBinding binding(session); + binding.BindInput("X", gpu_x); + binding.BindInput("indices", gpu_indices); + binding.BindInput("updates", gpu_updates); + binding.BindOutput("Y", gpu_y); + if (surrounding_dispatches) { + gpu_t = Value::CreateTensor(gpu_allocator, data_shape.data(), data_shape.size()); + gpu_z = Value::CreateTensor(gpu_allocator, data_shape.data(), data_shape.size()); + ASSERT_NE(gpu_t.GetTensorRawData(), gpu_y.GetTensorRawData()); + binding.BindOutput("T", gpu_t); + binding.BindOutput("Z", gpu_z); + } + + MemoryInfo cpu_mem_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + std::vector x{1.0f, 2.0f, 3.0f}; + std::vector indices{1}; + std::vector updates{9.0f}; + std::vector result(3); + std::vector input_after_run(3); + auto cpu_x = Value::CreateTensor(cpu_mem_info, x.data(), x.size(), data_shape.data(), data_shape.size()); + auto cpu_indices = Value::CreateTensor( + cpu_mem_info, indices.data(), indices.size(), indices_shape.data(), indices_shape.size()); + auto cpu_updates = Value::CreateTensor( + cpu_mem_info, updates.data(), updates.size(), updates_shape.data(), updates_shape.size()); + auto cpu_result = Value::CreateTensor( + cpu_mem_info, result.data(), result.size(), data_shape.data(), data_shape.size()); + auto cpu_input_after_run = Value::CreateTensor( + cpu_mem_info, input_after_run.data(), input_after_run.size(), data_shape.data(), data_shape.size()); + auto copy = [&](const Value& src, Value& dst) { + auto status = env.CopyTensor(src, dst, nullptr); + ASSERT_TRUE(status.IsOK()) << status.GetErrorMessage(); + }; + + ASSERT_NO_FATAL_FAILURE(copy(cpu_indices, gpu_indices)); + RunOptions run_options; + run_options.AddConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation, "1"); + for (int run = 0; run < 4; ++run) { + SCOPED_TRACE(testing::Message() << op_type << ", capture=" << enable_graph_capture << ", run=" << run); + if (enable_graph_capture && run == 3) { + session.ReleaseCapturedGraph(1); + } + x[0] = 1.0f + 3.0f * run; + x[1] = 2.0f + 3.0f * run; + x[2] = 3.0f + 3.0f * run; + updates[0] = 9.0f - run; + ASSERT_NO_FATAL_FAILURE(copy(cpu_x, gpu_x)); + ASSERT_NO_FATAL_FAILURE(copy(cpu_updates, gpu_updates)); + binding.SynchronizeInputs(); + session.Run(run_options, binding); + binding.SynchronizeOutputs(); + ASSERT_NO_FATAL_FAILURE(copy(gpu_y, cpu_result)); + ASSERT_NO_FATAL_FAILURE(copy(gpu_x, cpu_input_after_run)); + EXPECT_EQ(input_after_run, x) << "Scatter must not overwrite the separately bound input"; + EXPECT_EQ(result, (std::vector{x[0], updates[0], x[2]})) + << "Replay must copy the new input values at positions not written by scatter"; + if (surrounding_dispatches) { + ASSERT_NO_FATAL_FAILURE(copy(gpu_z, cpu_result)); + EXPECT_EQ(result, (std::vector{x[0], updates[0], x[2]})) + << "The dispatch following scatter must observe the copy and the scatter update"; + } + } + if (enable_graph_capture) { + session.ReleaseCapturedGraph(1); + } + if (surrounding_dispatches) { + AllocatorWithDefaultOptions allocator; + auto profile_path = session.EndProfilingAllocated(allocator); + std::filesystem::remove(profile_path.get()); + } +} + +TEST_P(WebGpuScatterGraphCaptureTests, ScatterNDNonInplaceCopiesUpdatedInput) { + TestScatterReplayCopiesInput("ScatterND", {1, 1}, GetParam()); +} + +TEST_P(WebGpuScatterGraphCaptureTests, ScatterElementsNonInplaceCopiesUpdatedInput) { + TestScatterReplayCopiesInput("ScatterElements", {1}, GetParam()); +} + +TEST_P(WebGpuScatterGraphCaptureTests, ScatterNDCopyBetweenProfiledDispatches) { + TestScatterReplayCopiesInput("ScatterND", {1, 1}, GetParam(), true); +} + +INSTANTIATE_TEST_SUITE_P(CaptureDisabledAndEnabled, WebGpuScatterGraphCaptureTests, testing::Bool()); + TEST(GraphCaptureTests, TestReleaseCapturedGraph) { Env& env = *ort_env; diff --git a/onnxruntime/test/providers/qnn/mod_op_test.cc b/onnxruntime/test/providers/qnn/mod_op_test.cc index 271fa76a6d639..06793f72cda51 100644 --- a/onnxruntime/test/providers/qnn/mod_op_test.cc +++ b/onnxruntime/test/providers/qnn/mod_op_test.cc @@ -4,6 +4,7 @@ #if !defined(ORT_MINIMAL_BUILD) #include +#include #include #include "test/providers/qnn/qnn_test_utils.h" @@ -38,6 +39,25 @@ static void RunModTest(const std::vector>& input_defs, // CPU tests: // +TEST_F(QnnCPUBackendTests, Mod_float_Floor_Opset28_Fallback) { + const float inf = std::numeric_limits::infinity(); + RunModTest({TestInputDef({8}, false, {-5.0f, 5.0f, -5.0f, 5.0f, -5.0f, 5.0f, -5.0f, 5.0f}), + TestInputDef({8}, false, {3.0f, 3.0f, -3.0f, -3.0f, inf, inf, -inf, -inf})}, + {utils::MakeAttribute("fmod", static_cast(0))}, + ExpectedEPNodeAssignment::None, + "cpu", + 28); +} + +TEST_F(QnnCPUBackendTests, Mod_integer_Floor_Opset28) { + RunModTest({TestInputDef({4}, false, {-5, 5, -5, 5}), + TestInputDef({4}, false, {3, 3, -3, -3})}, + {utils::MakeAttribute("fmod", static_cast(0))}, + ExpectedEPNodeAssignment::All, + "cpu", + 28); +} + // Test that Mod with dynamic divisor. TEST_F(QnnCPUBackendTests, Mod_dynamic_Divisor) { RandomValueGenerator rand_gen{optional{2345}}; diff --git a/onnxruntime/test/providers/qnn/qnn_model_wrapper_test.cc b/onnxruntime/test/providers/qnn/qnn_model_wrapper_test.cc index 54f966928e577..505d7018cccc6 100644 --- a/onnxruntime/test/providers/qnn/qnn_model_wrapper_test.cc +++ b/onnxruntime/test/providers/qnn/qnn_model_wrapper_test.cc @@ -10,9 +10,11 @@ #if !defined(ORT_MINIMAL_BUILD) && BUILD_QNN_EP_STATIC_LIB #include "core/graph/model.h" +#include "core/providers/qnn/builder/op_builder_factory.h" #include "core/providers/qnn/builder/qnn_model_wrapper.h" #include "core/providers/qnn/builder/qnn_def.h" #include "test/util/include/default_providers.h" +#include "test/util/include/asserts.h" #include "test/util/include/test_environment.h" using namespace onnxruntime; @@ -56,6 +58,46 @@ struct QnnModelWrapperTestContext { } // namespace +TEST(QnnModelWrapperTest, SpaceToDepth_RejectsCRDBeforeLayoutTransformation) { + for (const std::string mode : {"DCR", "CRD"}) { + QnnModelWrapperTestContext ctx; + Graph& graph = ctx.model->MainGraph(); + ONNX_NAMESPACE::TypeProto input_type; + auto* input_tensor = input_type.mutable_tensor_type(); + input_tensor->set_elem_type(ONNX_NAMESPACE::TensorProto::FLOAT); + for (int64_t dim : {1, 2, 4, 4}) { + input_tensor->mutable_shape()->add_dim()->set_dim_value(dim); + } + ONNX_NAMESPACE::TypeProto output_type; + auto* output_tensor = output_type.mutable_tensor_type(); + output_tensor->set_elem_type(ONNX_NAMESPACE::TensorProto::FLOAT); + for (int64_t dim : {1, 8, 2, 2}) { + output_tensor->mutable_shape()->add_dim()->set_dim_value(dim); + } + auto& input = graph.GetOrCreateNodeArg("input", &input_type); + auto& output = graph.GetOrCreateNodeArg("output", &output_type); + Node& node = graph.AddNode("SpaceToDepth", "SpaceToDepth", "", {&input}, {&output}); + node.AddAttribute("blocksize", static_cast(2)); + node.AddAttribute("mode", mode); + ASSERT_STATUS_OK(graph.Resolve()); + ctx.graph_viewer = std::make_unique(graph); + ctx.input_index_map = {{"input", 0}}; + ctx.output_index_map = {{"output", 0}}; + + auto wrapper = ctx.CreateWrapper(ModelSettings{}); + const auto* builder = GetOpBuilder("SpaceToDepth"); + ASSERT_NE(builder, nullptr); + const NodeUnit node_unit(node); + auto status = builder->IsOpSupported(*wrapper, node_unit, DefaultLoggingManager().DefaultLogger()); + if (mode == "DCR") { + ASSERT_STATUS_OK(status); + } else { + EXPECT_FALSE(status.IsOK()); + EXPECT_NE(status.ErrorMessage().find("only supports SpaceToDepth with DCR mode"), std::string::npos); + } + } +} + // Verifies that when htp_shared_memory is disabled (default), the mem type of a // graph input tensor remains QNN_TENSORMEMTYPE_RAW. TEST(QnnModelWrapperTest, AddTensorWrapper_SharedMemoryDisabled_GraphInput_MemTypeIsRaw) { diff --git a/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc new file mode 100644 index 0000000000000..7cb3986c71d04 --- /dev/null +++ b/onnxruntime/test/providers/webgpu/contrib_ops/matmul_block_quantized_fp8_weight_test.cc @@ -0,0 +1,352 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "core/common/float8.h" +#include "core/providers/webgpu/math/subgroup_matrix_config.h" +#include "core/providers/webgpu/program_cache_key.h" +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_context.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime { +namespace test { + +#if defined(USE_WEBGPU) && !defined(DISABLE_FLOAT8_TYPES) +namespace { +class SubgroupIdTestProgram final : public webgpu::Program { + public: + SubgroupIdTestProgram() : Program{"SubgroupIdTest"} {} + Status GenerateShaderCode(webgpu::ShaderHelper&) const override { return Status::OK(); } +}; + +void RunFp8MatMul(const std::vector& a_dims, int64_t n, int64_t block_size, + bool with_a_scale, bool with_bias, bool distinct_rows = false) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + + const auto k = a_dims.back(); + int64_t m = 1; + for (size_t i = 0; i + 1 < a_dims.size(); ++i) { + m *= a_dims[i]; + } + const auto blocks = (k + block_size - 1) / block_size; + std::vector a; + std::vector b; + std::vector scales; + std::vector biases; + for (int64_t i = 0; i < m * k; ++i) { + a.emplace_back(0.125f * static_cast((i % 11) - 5) + + (distinct_rows ? 0.015625f * static_cast(i / k) : 0.0f)); + } + for (int64_t i = 0; i < n; ++i) { + biases.emplace_back(static_cast((i % 5) - 2) * 0.25f); + for (int64_t j = 0; j < blocks; ++j) { + scales.push_back(0.125f * static_cast(1 + (i + j) % 4)); + } + for (int64_t j = 0; j < k; ++j) { + b.emplace_back(0.25f * static_cast(((i * 3 + j) % 9) - 4), true); + } + } + + std::vector expected; + expected.reserve(m * n); + for (int64_t row = 0; row < m; ++row) { + for (int64_t col = 0; col < n; ++col) { + float sum = 0.0f; + for (int64_t j = 0; j < k; ++j) { + float activation = a[row * k + j].ToFloat(); + if (with_a_scale) { + activation = MLFloat16(Float8E4M3FN(activation / 0.375f, true).ToFloat() * 0.375f).ToFloat(); + } + const auto weight = MLFloat16(b[col * k + j].ToFloat() * + scales[col * blocks + j / block_size]) + .ToFloat(); + sum += activation * weight; + } + expected.emplace_back(MLFloat16(sum + (with_bias ? biases[col].ToFloat() : 0.0f))); + } + } + + std::vector output_dims(a_dims.begin(), a_dims.end()); + output_dims.back() = n; + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", block_size); + test.AddInput("A", a_dims, a); + test.AddInput("B", {n, k}, b, /*is_initializer=*/true); + test.AddInput("b_scale", {n, blocks}, scales, /*is_initializer=*/true); + if (with_a_scale) { + test.AddInput("a_scale", {}, {0.375f}, /*is_initializer=*/true); + } else { + test.AddOptionalInputEdge(); + } + if (with_bias) { + test.AddInput("bias", {n}, biases); + } + test.AddOutput("Y", output_dims, expected); + test.SetOutputTolerance(0.03f); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} + +bool SupportsFp8MatrixPath() { + const auto& context = webgpu::WebGpuContextFactory::DefaultContext(); + bool supported = false; + if (context.Instance().HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) && + context.DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { + const auto& configs = context.SubgroupMatrixConfigs(); + const auto& adapter = context.AdapterInfo(); + supported = webgpu::detail::SelectSubgroupMatrixConfigFromAdapterConfigs( + {configs.configs, configs.configCount}, adapter.subgroupMinSize, adapter.subgroupMaxSize, + context.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl), + {{wgpu::SubgroupMatrixComponentType::F16, wgpu::SubgroupMatrixComponentType::F32, + 16, 16, 16, 32, false}}) + .has_value(); + } + webgpu::WebGpuContextFactory::ReleaseContext(0); + return supported; +} +} // namespace + +TEST(MatMulBlockQuantizedFp8WeightTest, SubgroupIdIsOptInAndCacheKeyed) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + const auto& context = webgpu::WebGpuContextFactory::DefaultContext(); + SubgroupIdTestProgram program; + const auto original_key = webgpu::CalculateProgramCacheKey(program, {}, {}); + webgpu::ShaderHelper ordinary_shader(program, program.Metadata(), context, {}, {}, 1, 1, 1); + ASSERT_STATUS_OK(ordinary_shader.Init()); + EXPECT_EQ(std::move(ordinary_shader.MainFunctionBody()).str().find("@builtin(subgroup_id)"), std::string::npos); + + program.SetUseSubgroupId(true); + EXPECT_NE(original_key, webgpu::CalculateProgramCacheKey(program, {}, {})); + webgpu::ShaderHelper matrix_shader(program, program.Metadata(), context, {}, {}, 1, 1, 1); + const bool supported = context.DeviceHasFeature(wgpu::FeatureName::Subgroups) && + context.Instance().HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId); + const auto status = matrix_shader.Init(); + EXPECT_EQ(status.IsOK(), supported); + if (supported) { + EXPECT_NE(std::move(matrix_shader.MainFunctionBody()).str().find("@builtin(subgroup_id)"), std::string::npos); + } + webgpu::WebGpuContextFactory::ReleaseContext(0); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, DecodeRankOneOddBytes) { + RunFp8MatMul({3}, 1, 2, false, false); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, ScalarDispatchBeyondWebGpuDimensionLimit) { + RunFp8MatMul({1, 1}, 65537, 1, false, false); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, DecodeWithBlockScalesAndBias) { + RunFp8MatMul({1, 33}, 19, 16, false, true); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqBatched) { + RunFp8MatMul({2, 3, 33}, 19, 16, true, true); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, MatrixTailWithBlockScales) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + if (!SupportsFp8MatrixPath()) { + GTEST_SKIP() << "16x16x16 FP16/F32 subgroup matrix with subgroup size 32 is unavailable"; + } + RunFp8MatMul({9, 33}, 19, 16, false, true); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, MatrixWithActivationQdq) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + if (!SupportsFp8MatrixPath()) { + GTEST_SKIP() << "16x16x16 FP16/F32 subgroup matrix with subgroup size 32 is unavailable"; + } + RunFp8MatMul({8, 32}, 16, 32, true, false); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, MatrixSubgroupTilesAndSecondWorkgroup) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + if (!SupportsFp8MatrixPath()) { + GTEST_SKIP() << "16x16x16 FP16/F32 subgroup matrix with subgroup size 32 is unavailable"; + } + RunFp8MatMul({64, 16}, 16, 16, false, false, true); + RunFp8MatMul({65, 16}, 16, 16, false, false, true); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, SubnormalSignedWeightsAndBias) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + const std::vector b{ + Float8E4M3FN(uint8_t{0x01}, Float8E4M3FN::FromBits()), + Float8E4M3FN(uint8_t{0x87}, Float8E4M3FN::FromBits()), + Float8E4M3FN(0.5625f, true), + Float8E4M3FN(-1.0f, true)}; + float expected = 0.25f; + for (const auto weight : b) { + expected += weight.ToFloat(); + } + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{4}); + test.AddInput("A", {1, 4}, std::vector(4, MLFloat16(1.0f))); + test.AddInput("B", {1, 4}, b, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddOptionalInputEdge(); + test.AddInput("bias", {1}, {MLFloat16(0.25f)}); + test.AddOutput("Y", {1, 1}, {MLFloat16(expected)}); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqPreservesNaN) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + const float nan = std::numeric_limits::quiet_NaN(); + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{1}); + test.AddInput("A", {2, 1}, ToFloat16({nan, 1.0f})); + test.AddInput("B", {1, 1}, {Float8E4M3FN(1.0f, true)}, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddInput("a_scale", {}, {0.375f}); + test.AddOutput("Y", {2, 1}, ToFloat16({nan, 1.03125f})); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqTieAndSaturation) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{1}); + test.AddInput("A", {7, 1}, + ToFloat16({0.3984375f, -0.3984375f, 0.4453125f, -0.4453125f, + 200.0f, -200.0f, 0.0f})); + test.AddInput("B", {1, 1}, {Float8E4M3FN(1.0f, true)}, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddInput("a_scale", {}, {0.375f}); + test.AddOutput("Y", {7, 1}, + ToFloat16({0.375f, -0.375f, 0.46875f, -0.46875f, + 168.0f, -168.0f, 0.0f})); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqAdjacentEvenAndOddTies) { + for (float scale : {0.375f, 0.75f, 1.0f}) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + const float factor = scale / 0.375f; + const float delta = factor * 0.000244140625f; + const float even_tie = factor * 0.3984375f; + const float odd_tie = factor * 0.4453125f; + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{1}); + test.AddInput("A", {6, 1}, + ToFloat16({even_tie - delta, even_tie, even_tie + delta, + odd_tie - delta, odd_tie, odd_tie + delta})); + test.AddInput("B", {1, 1}, {Float8E4M3FN(1.0f, true)}, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddInput("a_scale", {}, {scale}); + test.AddOutput("Y", {6, 1}, + ToFloat16({factor * 0.375f, factor * 0.375f, factor * 0.421875f, + factor * 0.421875f, factor * 0.46875f, factor * 0.46875f})); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); + } +} + +TEST(MatMulBlockQuantizedFp8WeightTest, ActivationQdqReciprocalMidpoint) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + constexpr float scale = 1.21603262424469f; + const auto input = ToFloat16({0.0272979736328125f, 0.027313232421875f, 0.0273284912109375f}); + ASSERT_LT(input[1].ToFloat() * (1.0f / scale), 0.0224609375f); + + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{1}); + test.AddInput("A", {3, 1}, input); + test.AddInput("B", {1, 1}, {Float8E4M3FN(1.0f, true)}, true); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddInput("a_scale", {}, {scale}); + test.AddOutput("Y", {3, 1}, + ToFloat16({0.021484375f * scale, 0.021484375f * scale, 0.0234375f * scale})); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectSuccess, {}, {}, nullptr, &providers); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, EmptyReductionWithBias) { + RunFp8MatMul({2, 0}, 3, 16, false, true); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, EmptyOutput) { + RunFp8MatMul({0, 16}, 3, 16, false, false); +} + +TEST(MatMulBlockQuantizedFp8WeightTest, InvalidScaleShape) { + auto ep = DefaultWebGpuExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", int64_t{16}); + test.AddInput("A", {1, 17}, std::vector(17, MLFloat16(1.0f))); + test.AddInput("B", {1, 17}, std::vector(17, Float8E4M3FN(1.0f, true))); + test.AddInput("b_scale", {1, 1}, {1.0f}); + test.AddOutput("Y", {1, 1}, {MLFloat16(17.0f)}); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + std::vector> providers; + providers.push_back(std::move(ep)); + test.Run(options, OpTester::ExpectResult::kExpectFailure, "b_scale must have shape", {}, nullptr, &providers); +} + +#endif + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/webgpu/conv2d_mm_test.cc b/onnxruntime/test/providers/webgpu/conv2d_mm_test.cc new file mode 100644 index 0000000000000..c271ed6fc3134 --- /dev/null +++ b/onnxruntime/test/providers/webgpu/conv2d_mm_test.cc @@ -0,0 +1,105 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "gtest/gtest.h" + +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "default_providers.h" +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime { +namespace test { +namespace { + +template +void RunConv2dMMInputChannelsTest(int64_t input_channels, bool weights_are_initializers) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + + constexpr int64_t batch = 2; + constexpr int64_t height = 8; + constexpr int64_t width = 8; + constexpr int64_t output_channels = 16; + constexpr int64_t kernel_size = 3; + std::vector input(static_cast(batch * input_channels * height * width)); + std::vector weights(static_cast(output_channels * input_channels * kernel_size * kernel_size)); + std::vector bias(static_cast(output_channels)); + std::vector expected(static_cast(batch * output_channels * height * width)); + // Distinct channel values expose reads from the wrong input-channel triplet. + // Binary fractions keep both the inputs and the reference accumulation exact. + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast(static_cast(i % 37) - 18) / 8.0f; + } + for (size_t i = 0; i < weights.size(); ++i) { + weights[i] = static_cast(static_cast(i % 11) - 5) / 16.0f; + } + for (int64_t c = 0; c < output_channels; ++c) { + bias[static_cast(c)] = static_cast(c - 8) / 8.0f; + } + for (int64_t n = 0; n < batch; ++n) { + for (int64_t oc = 0; oc < output_channels; ++oc) { + for (int64_t h = 0; h < height; ++h) { + for (int64_t w = 0; w < width; ++w) { + float sum = bias[static_cast(oc)]; + for (int64_t ic = 0; ic < input_channels; ++ic) { + for (int64_t kh = 0; kh < kernel_size; ++kh) { + for (int64_t kw = 0; kw < kernel_size; ++kw) { + const int64_t ih = h + kh - 1; + const int64_t iw = w + kw - 1; + if (ih < 0 || ih >= height || iw < 0 || iw >= width) { + continue; + } + const size_t input_index = static_cast(((n * input_channels + ic) * height + ih) * width + iw); + const size_t weight_index = + static_cast(((oc * input_channels + ic) * kernel_size + kh) * kernel_size + kw); + sum += input[input_index] * weights[weight_index]; + } + } + } + expected[static_cast(((n * output_channels + oc) * height + h) * width + w)] = sum; + } + } + } + } + + OpTester test("Conv", 11); + test.AddAttribute("kernel_shape", std::vector{kernel_size, kernel_size}); + test.AddAttribute("pads", std::vector{1, 1, 1, 1}); + test.AddInput("X", {batch, input_channels, height, width}, std::vector(input.begin(), input.end())); + test.AddInput("W", {output_channels, input_channels, kernel_size, kernel_size}, + std::vector(weights.begin(), weights.end()), weights_are_initializers); + test.AddInput("B", {output_channels}, std::vector(bias.begin(), bias.end()), weights_are_initializers); + test.AddOutput("Y", {batch, output_channels, height, width}, std::vector(expected.begin(), expected.end())); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(options).ConfigEp(std::move(webgpu_ep)).RunWithConfig(); +} + +class Conv2dMMInputChannelsTest : public ::testing::TestWithParam> {}; + +TEST_P(Conv2dMMInputChannelsTest, Float) { + const auto [input_channels, weights_are_initializers] = GetParam(); + RunConv2dMMInputChannelsTest(input_channels, weights_are_initializers); +} + +TEST_P(Conv2dMMInputChannelsTest, Float16) { + const auto [input_channels, weights_are_initializers] = GetParam(); + RunConv2dMMInputChannelsTest(input_channels, weights_are_initializers); +} + +// Three channels only reads the first triplet. Larger multiples of three must +// also read subsequent triplets. Include scalar and vec4 dispatch controls. +INSTANTIATE_TEST_SUITE_P( + Conv_WebGPU, Conv2dMMInputChannelsTest, + ::testing::Combine(::testing::Values(int64_t{3}, int64_t{6}, int64_t{9}, int64_t{15}, + int64_t{5}, int64_t{4}, int64_t{12}), + ::testing::Bool())); + +} // namespace +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc b/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc index 59c91c5a930e0..a197551015f3e 100644 --- a/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc +++ b/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc @@ -17,6 +17,9 @@ #include "core/providers/webgpu/buffer_manager.h" #include "core/providers/webgpu/compute_context.h" #include "core/providers/webgpu/math/gemm_utils.h" +#include "core/providers/webgpu/math/matmul.h" +#include "core/providers/webgpu/math/matmul_packed.h" +#include "core/providers/webgpu/math/matmul_utils.h" #include "core/providers/webgpu/nn/conv2d_mm.h" #include "core/providers/webgpu/program_manager.h" #include "core/providers/webgpu/webgpu_provider_factory_creator.h" @@ -105,8 +108,9 @@ class DispatchTestKernel final : public OpKernel { Status Compute(OpKernelContext*) const override { return Status::OK(); } }; +template void RunAndReadProgram(WebGpuExecutionProvider& ep, const ProgramBase& program, - WGPUBuffer output_buffer, std::vector& output_data) { + WGPUBuffer output_buffer, std::vector& output_data) { auto& context = WebGpuContextFactory::GetContext(0); ConfigOptions options; Model model("dispatch_padding", false, DefaultLoggingManager().DefaultLogger()); @@ -120,11 +124,12 @@ void RunAndReadProgram(WebGpuExecutionProvider& ep, const ProgramBase& program, DispatchTestKernel kernel(info); ComputeContextBase compute_context(context, ep, kernel); ASSERT_STATUS_OK(compute_context.RunProgram(program)); - ep.BufferManager().Download(ep.Recording(), output_buffer, output_data.data(), output_data.size() * sizeof(float)); + ep.BufferManager().Download(ep.Recording(), output_buffer, output_data.data(), output_data.size() * sizeof(T)); } -wgpu::Buffer CreateStorageBuffer(WebGpuExecutionProvider& ep, gsl::span data) { - const size_t size = data.size() * sizeof(float); +template +wgpu::Buffer CreateStorageBuffer(WebGpuExecutionProvider& ep, std::vector& data) { + const size_t size = data.size() * sizeof(T); auto buffer = wgpu::Buffer::Acquire(ep.BufferManager().Create( ep.Recording(), size, wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst)); ep.BufferManager().Upload(ep.Recording(), data.data(), buffer.Get(), size); @@ -147,6 +152,61 @@ class WebGpuDispatchExecutionTest : public ::testing::Test { std::unique_ptr ep_; }; +TEST_F(WebGpuDispatchExecutionTest, PackedMatMulSupportsNonDefaultValidTuning) { + constexpr uint32_t M = 4; + constexpr uint32_t N = 64; + constexpr uint32_t K = 32; + std::vector a_data(M * K); + std::vector b_data(K * N); + std::vector output_data(M * N, 0.0f); + for (size_t i = 0; i < a_data.size(); ++i) { + a_data[i] = static_cast(static_cast(i % 7) - 3) * 0.125f; + } + for (size_t i = 0; i < b_data.size(); ++i) { + b_data[i] = static_cast(static_cast(i % 11) - 5) * 0.0625f; + } + + auto a_buffer = CreateStorageBuffer(Ep(), a_data); + auto b_buffer = CreateStorageBuffer(Ep(), b_data); + auto output_buffer = CreateStorageBuffer(Ep(), output_data); + const OrtMemoryInfo memory_info( + WEBGPU_BUFFER, OrtDeviceAllocator, webgpu::WebGpuDevice(0), OrtMemTypeDefault); + Tensor a(DataTypeImpl::GetType(), TensorShape{M, K}, a_buffer.Get(), memory_info); + Tensor b(DataTypeImpl::GetType(), TensorShape{K, N}, b_buffer.Get(), memory_info); + Tensor output(DataTypeImpl::GetType(), TensorShape{M, N}, output_buffer.Get(), memory_info); + + const TensorShape outer_dims{}; + const TensorShape a_program_shape = CreateMatMulIntermediateShape(outer_dims, M, K, 4); + const TensorShape b_program_shape = CreateMatMulIntermediateShape(outer_dims, K, N, 4); + const TensorShape output_program_shape{1, M, N / 4}; + InlinedVector elements_per_thread{4, 2, 1}; + const Activation activation; + MatMulProgram program{activation, /*bias=*/false, /*is_vec4=*/true, + elements_per_thread, /*is_channels_last=*/true, + /*split_dim_inner=*/1, /*tile_inner=*/64}; + program.AddInputs({{&a, ProgramTensorMetadataDependency::TypeAndRank, a_program_shape, 4}, + {&b, ProgramTensorMetadataDependency::TypeAndRank, b_program_shape, 4}}) + .AddUniformVariables({{M}, {N}, {K}, {uint32_t{1}}, {uint32_t{1}}, {uint32_t{1}}, {uint32_t{1}}}) + .AddIndices(outer_dims) + .SetDispatchGroupSize(1, 1, 1) + .SetWorkgroupSize(16, 4, 1) + .AddOutput(ProgramOutput(&output, ProgramTensorMetadataDependency::Rank, + output_program_shape, 4)); + AppendActivationUniformsData(activation, program); + + RunAndReadProgram(Ep(), program, output_buffer.Get(), output_data); + for (uint32_t m = 0; m < M; ++m) { + for (uint32_t n = 0; n < N; ++n) { + float expected = 0.0f; + for (uint32_t k = 0; k < K; ++k) { + expected += a_data[m * K + k] * b_data[k * N + n]; + } + EXPECT_NEAR(output_data[m * N + n], expected, 1e-5f) + << "at output (" << m << ", " << n << ")"; + } + } +} + constexpr float kOutputPaddingCanary = -12345.0f; void CheckResultsAndPadding(gsl::span output_data, size_t output_size, float expected) { @@ -300,6 +360,54 @@ TEST_F(WebGpuDispatchExecutionTest, Conv2dMMPreservesPaddingInBothLayouts) { } } +TEST_F(WebGpuDispatchExecutionTest, MatMulGemvPreservesPaddingAfterNormalization) { + auto& context = WebGpuContextFactory::GetContext(Ep().GetDeviceId()); + if (!context.DeviceHasFeature(wgpu::FeatureName::ShaderF16)) { + GTEST_SKIP() << "FP16 is not supported by the WebGPU adapter."; + } + constexpr uint32_t N = 28; + constexpr uint32_t K = 2048; + const MLFloat16 canary{-42.0f}; + const OrtMemoryInfo memory_info(WEBGPU_BUFFER, OrtDeviceAllocator, webgpu::WebGpuDevice(0), OrtMemTypeDefault); + for (const uint32_t M : {1u, 2u}) { + SCOPED_TRACE(M); + std::array physical{M * (N / 4), 1, 1}; + ASSERT_STATUS_OK(webgpu::detail::NormalizeDispatchGroupSize(physical[0], physical[1], physical[2], 5)); + const uint32_t padded_rows = CeilDiv(physical[0] * physical[1] * physical[2], N / 4); + ASSERT_GT(padded_rows, M); + + // Bind padded buffers so a missing row guard visibly overwrites the canary. + std::vector a_data(padded_rows * K, MLFloat16{0.5f}); + for (uint32_t m = 0; m < M; ++m) { + std::fill_n(a_data.begin() + m * K, K, MLFloat16{0.125f * static_cast(m + 1)}); + } + std::vector b_data(K * N, MLFloat16{0.125f}); + std::vector output_data(padded_rows * N, canary); + auto a_buffer = CreateStorageBuffer(Ep(), a_data); + auto b_buffer = CreateStorageBuffer(Ep(), b_data); + auto output_buffer = CreateStorageBuffer(Ep(), output_data); + Tensor a(DataTypeImpl::GetType(), {M, K}, a_buffer.Get(), memory_info); + Tensor b(DataTypeImpl::GetType(), {K, N}, b_buffer.Get(), memory_info); + Tensor output(DataTypeImpl::GetType(), {M, N}, output_buffer.Get(), memory_info); + MatMulGemvProgram program; + program.AddInputs({{&a, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, + {&b, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}}) + .AddOutput({&output, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten, 4}) + .AddUniformVariables({{K}, {N}, {M}}) + .SetWorkgroupSize(MatMulGemvProgram::kWorkgroupSizeX, MatMulGemvProgram::kWorkgroupSizeY) + .SetDispatchGroupSize(physical[0], physical[1], physical[2]); + ASSERT_NO_FATAL_FAILURE(RunAndReadProgram(Ep(), program, output_buffer.Get(), output_data)); + for (uint32_t m = 0; m < M; ++m) { + for (uint32_t n = 0; n < N; ++n) { + EXPECT_EQ(output_data[m * N + n].ToFloat(), 32.0f * static_cast(m + 1)); + } + } + for (size_t i = M * N; i < output_data.size(); ++i) { + EXPECT_EQ(output_data[i].ToFloat(), canary.ToFloat()) << "at padding index " << i; + } + } +} + } // namespace } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/webgpu/graph_capture_test_utils.h b/onnxruntime/test/providers/webgpu/graph_capture_test_utils.h new file mode 100644 index 0000000000000..d3b3bc143dbac --- /dev/null +++ b/onnxruntime/test/providers/webgpu/graph_capture_test_utils.h @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include + +#include "core/framework/allocator.h" +#include "core/framework/tensor.h" +#include "core/session/IOBinding.h" +#include "core/session/inference_session.h" +#include "test/providers/provider_test_utils.h" +#include "test/test_environment.h" + +namespace onnxruntime::test { + +class WebGpuGraphCaptureTester final : public OpTester { + public: + using OpTester::OpTester; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WebGpuGraphCaptureTester); + + template + void RunWithBoundCapture(std::unique_ptr provider, + const SessionOptions& options, const TensorShape& shape, + gsl::span expected) { + SetTestFunctionCalled(); + auto& capture_model = BuildModel(); + ASSERT_STATUS_OK(capture_model.MainGraph().Resolve(GetRunContext().resolve_options)); + const std::string model = capture_model.ToProto().SerializeAsString(); + auto* capture_provider = provider.get(); + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(provider))); + ASSERT_STATUS_OK(session.Load(model.data(), static_cast(model.size()))); + ASSERT_STATUS_OK(session.Initialize()); + const auto allocators = capture_provider->CreatePreferredAllocators(); + ASSERT_FALSE(allocators.empty()); + const auto allocator = allocators.back(); + ASSERT_NE(allocator->Info().device.Type(), OrtDevice::CPU); + std::unique_ptr binding; + ASSERT_STATUS_OK(session.NewIOBinding(&binding)); + std::unordered_map feeds; + std::vector output_names; + FillFeedsAndOutputNames(feeds, output_names); + ASSERT_EQ(output_names.size(), 1u); + for (const auto& [name, value] : feeds) { + const auto& source = value.template Get(); + OrtValue input; + Tensor::InitOrtValue(source.DataType(), source.Shape(), allocator, input); + if (source.Shape().Size() != 0) { + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(source, *input.GetMutable())); + } + ASSERT_STATUS_OK(binding->BindInput(name, input)); + } + OrtValue output; + Tensor::InitOrtValue(DataTypeImpl::GetType(), shape, allocator, output); + ASSERT_STATUS_OK(binding->BindOutput(output_names[0], output)); + ASSERT_STATUS_OK(binding->SynchronizeInputs()); + Tensor cpu_output(DataTypeImpl::GetType(), shape, CPUAllocator::DefaultInstance()); + ASSERT_EQ(static_cast(shape.Size()), expected.size()); + for (int run = 0; run < 4; ++run) { + if (!expected.empty()) { + std::fill_n(cpu_output.MutableData(), expected.size(), T{-99.0f}); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(cpu_output, *output.GetMutable())); + } + ASSERT_STATUS_OK(session.Run(RunOptions{}, *binding)); + ASSERT_STATUS_OK(binding->SynchronizeOutputs()); + ASSERT_EQ(binding->GetOutputs()[0].Get().Shape(), shape); + if (!expected.empty()) { + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(output.Get(), cpu_output)); + for (size_t index = 0; index < expected.size(); ++index) { + EXPECT_EQ(cpu_output.Data()[index], expected[index]) << "run=" << run << " index=" << index; + } + } + } + EXPECT_TRUE(capture_provider->IsGraphCaptured(0)); + } +}; + +} // namespace onnxruntime::test diff --git a/onnxruntime/test/providers/webgpu/math/element_wise_ops_test.cc b/onnxruntime/test/providers/webgpu/math/element_wise_ops_test.cc index afc120e4500d2..acbd2246c19fc 100644 --- a/onnxruntime/test/providers/webgpu/math/element_wise_ops_test.cc +++ b/onnxruntime/test/providers/webgpu/math/element_wise_ops_test.cc @@ -149,6 +149,68 @@ TEST(MathOpTest, Add_Broadcast_WebGpu_UnequalRank_ShorterLhs) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +namespace { +template +void TestPowBroadcast(const std::vector& base_shape, const std::vector& base, + const std::vector& exponent_shape, const std::vector& exponent, + const std::vector& output_shape, const std::vector& expected) { + for (bool exponent_is_initializer : {false, true}) { + SCOPED_TRACE(exponent_is_initializer); + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU execution provider is not enabled in this build."; + } + OpTester test("Pow", 15); + test.AddInput("X", base_shape, base); + test.AddInput("Y", exponent_shape, exponent, exponent_is_initializer); + test.AddOutput("Z", output_shape, expected); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(options).ConfigEp(std::move(webgpu_ep)).RunWithConfig(); + } +} +} // namespace + +TEST(MathOpTest, Pow_BroadcastExponent_WebGpu) { + // Issue #32842: a rank-4 singleton uses vectorized broadcasting rather than the scalar path. + const std::vector shape{3, 2, 64, 64}; + std::vector base(3 * 2 * 64 * 64); + std::vector expected(base.size()); + for (size_t i = 0; i < base.size(); ++i) { + base[i] = static_cast(static_cast(i % 5) - 2); + expected[i] = base[i] * base[i] * base[i]; + } + TestPowBroadcast(shape, base, {1, 1, 1, 1}, {3.0f}, shape, expected); +} + +TEST(MathOpTest, Pow_BroadcastBase_WebGpu) { + TestPowBroadcast({1, 1, 1, 1}, {2.0f}, {1, 1, 2, 4}, + {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 3.0f, 2.0f, 1.0f}, + {1, 1, 2, 4}, {1.0f, 2.0f, 4.0f, 8.0f, 16.0f, 8.0f, 4.0f, 2.0f}); +} + +TEST(MathOpTest, Pow_BroadcastExponent_MixedTypes_WebGpu) { + TestPowBroadcast({2, 4}, {-2.0f, -1.0f, 0.0f, 2.0f, -3.0f, -2.0f, 1.0f, 3.0f}, + {2, 1}, {3, 2}, {2, 4}, {-8.0f, -1.0f, 0.0f, 8.0f, 9.0f, 4.0f, 1.0f, 9.0f}); +} + +TEST(MathOpTest, Pow_BroadcastBase_MixedTypes_WebGpu) { + TestPowBroadcast({2, 1}, {-2, 3}, {2, 4}, {0.0f, 1.0f, 2.0f, 3.0f, 0.0f, 1.0f, 2.0f, 3.0f}, + {2, 4}, {1, -2, 4, -8, 1, 3, 9, 27}); +} + +TEST(MathOpTest, Pow_BroadcastExponent_Float16_WebGpu) { + TestPowBroadcast({2, 4}, MakeMLFloat16({-2.0f, -1.0f, 0.0f, 2.0f, -3.0f, -2.0f, 1.0f, 3.0f}), + {2, 1}, MakeMLFloat16({3.0f, 2.0f}), {2, 4}, + MakeMLFloat16({-8.0f, -1.0f, 0.0f, 8.0f, 9.0f, 4.0f, 1.0f, 9.0f})); +} + +TEST(MathOpTest, Pow_BroadcastBase_Float16_WebGpu) { + TestPowBroadcast({2, 1}, MakeMLFloat16({-2.0f, 3.0f}), + {2, 4}, MakeMLFloat16({0.0f, 1.0f, 2.0f, 3.0f, 0.0f, 1.0f, 2.0f, 3.0f}), + {2, 4}, MakeMLFloat16({1.0f, -2.0f, 4.0f, -8.0f, 1.0f, 3.0f, 9.0f, 27.0f})); +} + // WebGPU EP currently handles a special case for supporting Pow op: // A Pow followed by a Cast to int64 type. TEST(MathOpTest, Pow_float_sqrt) { diff --git a/onnxruntime/test/providers/webgpu/math/matmul_test.cc b/onnxruntime/test/providers/webgpu/math/matmul_test.cc index cd42da8f1e083..2ff6ad9efca94 100644 --- a/onnxruntime/test/providers/webgpu/math/matmul_test.cc +++ b/onnxruntime/test/providers/webgpu/math/matmul_test.cc @@ -3,11 +3,15 @@ #include #include +#include #include #include #include "gtest/gtest.h" #include "core/providers/webgpu/webgpu_provider_options.h" +#if !defined(ORT_USE_EP_API_ADAPTERS) +#include "core/providers/webgpu/webgpu_context.h" +#endif #include "core/session/onnxruntime_session_options_config_keys.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" @@ -180,14 +184,15 @@ TEST(MathOpTest, MatMulPackedScalarFp16LongReductionUsesFloat32Accumulator) { } } -// f16 MatMul cases that exercise the Intel 8x16x16 subgroup-matrix impl. +// f16 MatMul cases that exercise the 8x16x16 and 16x16x16 subgroup-matrix impl. // The host picks the tile shape adaptively (TileM in {8,16,32,64}, TileN in // {16,32,64}); M and N may be any size and K must be a multiple of 16. When the // output has few tiles and K is large, the host also splits K across multiple // cooperating subgroups (split_k in {1,2,4,8}) that reduce in shared memory. B -// is a constant initializer. On non-Intel hardware these fall back to the -// default impl but still validate correctness. -static void RunSubgroupMatrixMatMulTest(const std::vector& a_dims, int64_t K, int64_t N) { +// is a constant initializer except in the last case. On hardware that reports +// neither config these fall back to the default impl but still validate correctness. +static void RunSubgroupMatrixMatMulTest(const std::vector& a_dims, int64_t K, int64_t N, + bool b_is_initializer = true) { int64_t M = 1; for (size_t i = 0; i + 1 < a_dims.size(); ++i) { M *= a_dims[i]; @@ -225,7 +230,7 @@ static void RunSubgroupMatrixMatMulTest(const std::vector& a_dims, int6 OpTester test("MatMul", 14); test.AddInput("A", a_dims, f_A); - test.AddInput("B", {K, N}, f_B, /*is_initializer=*/true); + test.AddInput("B", {K, N}, f_B, b_is_initializer); test.AddOutput("Y", y_dims, f_Y); test.SetOutputTolerance(0.02f); test.ConfigExcludeEps({kTensorrtExecutionProvider}) @@ -277,7 +282,7 @@ TEST(MathOpTest, MatMulSubgroupMatrix) { // Batched A folds to M=8; one tile + K=256 -> split_k=8. {"SplitKBatched (2*4 -> M=8)", {2, 4, 256}, 256, 16}, // Odd N: the subgroup f16 load needs an even B row stride, so a constant odd-N - // weight is padded once to N+1 (even) and cached; output is still written at + // weight is padded once to an even width and cached. Output is still written at // the real, odd N. Covers small/large odd N, min K, partial M, batched-A fold, // and split-K, all with odd N. {"OddN15 (N=15)", {32, 64}, 64, 15}, @@ -286,12 +291,171 @@ TEST(MathOpTest, MatMulSubgroupMatrix) { {"OddN PartialM (M=40,N=31)", {40, 64}, 64, 31}, {"OddN BatchedA (2*32 -> M=64,N=63)", {2, 32, 64}, 64, 63}, {"OddN SplitK (K=256,N=17)", {8, 256}, 256, 17}, + // N is even but not a multiple of 8 so only the 16x16x16 config pads B. + {"UnalignedN36 (N=36)", {64, 128}, 128, 36}, + {"UnalignedN44 PartialM (M=40,N=44)", {40, 128}, 128, 44}, + // Odd N and partial M past 64 so the default tiling takes its 64x64 tile on padded B. + {"OddN LargeTile (M=72,N=67)", {72, 96}, 96, 67}, }; for (const auto& c : cases) { SCOPED_TRACE(c.name); RunSubgroupMatrixMatMulTest(c.a_dims, c.K, c.N); } + + // B is a runtime input here so its padded copy could not be cached and the + // 16x16x16 config falls back to the default impl for this width. + SCOPED_TRACE("RuntimeB UnalignedN36 (N=36)"); + RunSubgroupMatrixMatMulTest({64, 128}, 128, 36, /*b_is_initializer=*/false); +} + +class WebGpuMatMulGemvTest : public ::testing::Test { + protected: + void SetUp() override { + WebGpuExecutionProviderTestOptions options{}; + options.forced_matmul_algorithm = webgpu::MatMulAlgorithm::Gemv; + try { + auto ep = WebGpuExecutionProviderWithTestOptions(options); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider test options are not available."; + } +#if !defined(ORT_USE_EP_API_ADAPTERS) + auto& context = webgpu::WebGpuContextFactory::GetContext(ep->GetDeviceId()); + if (!context.HasDevice() || !context.DeviceHasFeature(wgpu::FeatureName::ShaderF16)) { + GTEST_SKIP() << "A usable FP16 WebGPU adapter is not available."; + } +#endif + } catch (const OnnxRuntimeException& error) { + if (std::string_view{error.what()}.find("Failed to get a WebGPU adapter") != std::string_view::npos) { + GTEST_SKIP() << error.what(); + } + throw; + } + } + + void Run(const std::vector& a_dims, int64_t K, int64_t N, bool expect_gemv, + bool force_gemv = false, bool batched_weights = false) { + SCOPED_TRACE(::testing::PrintToString(a_dims)); + SCOPED_TRACE(K); + SCOPED_TRACE(N); + SCOPED_TRACE(batched_weights); + const int64_t M = TensorShape{a_dims}.SizeToDimension(a_dims.size() - 1); + const int64_t rows_per_batch = a_dims[a_dims.size() - 2]; + const int64_t weight_batches = batched_weights ? M / rows_per_batch : 1; + std::vector a(M * K), b(weight_batches * K * N), y(M * N); + for (int64_t i = 0; i < M * K; ++i) { + a[i] = MLFloat16(static_cast((i % 7) - 3) * 0.125f); + } + for (int64_t i = 0; i < weight_batches * K * N; ++i) { + const int64_t batch = i / (K * N); + b[i] = MLFloat16(static_cast((i % 5) - 2) * 0.125f + static_cast(batch) * 0.25f); + } + for (int64_t m = 0; m < M; ++m) { + const int64_t b_offset = batched_weights ? (m / rows_per_batch) * K * N : 0; + for (int64_t n = 0; n < N; ++n) { + float sum = 0.0f; + for (int64_t k = 0; k < K; ++k) { + sum += a[m * K + k].ToFloat() * b[b_offset + k * N + n].ToFloat(); + } + y[m * N + n] = MLFloat16(sum); + } + } + + struct DispatchCounts { + size_t total = 0; + size_t gemv = 0; + } dispatches; + SessionOptions session_options; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + session_options.session_log_severity_level = ORT_LOGGING_LEVEL_INFO; + session_options.user_logging_param = &dispatches; + session_options.user_logging_function = [](void* param, OrtLoggingLevel, const char*, const char*, + const char*, const char* message) { + auto& counts = *static_cast(param); + const std::string_view text{message}; + if (text.find("Starting program \"") == 0) { + ++counts.total; + if (text.find("Starting program \"MatMulGemv") == 0) { + ++counts.gemv; + } + } + }; + WebGpuExecutionProviderTestOptions options{}; + options.forced_matmul_algorithm = webgpu::MatMulAlgorithm::Gemv; + auto ep = force_gemv ? WebGpuExecutionProviderWithTestOptions(options) : DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + + auto y_dims = a_dims; + y_dims.back() = N; + auto b_dims = batched_weights ? a_dims : std::vector{K, N}; + b_dims[b_dims.size() - 2] = K; + b_dims.back() = N; + OpTester test("MatMul", 14); + test.AddInput("A", a_dims, a); + test.AddInput("B", b_dims, b); + test.AddOutput("Y", y_dims, y); + test.SetOutputAbsErr("Y", 0.02f); + test.Config(session_options).ConfigEp(std::move(ep)).RunWithConfig(); + ASSERT_GT(dispatches.total, 0u); + if (expect_gemv) { + EXPECT_EQ(dispatches.gemv, dispatches.total); + } else { + EXPECT_EQ(dispatches.gemv, 0u); + } + } +}; + +TEST_F(WebGpuMatMulGemvTest, SmallRows) { + for (int64_t m = 1; m <= 8; ++m) { + Run({m, 5120}, 5120, 48, true, true); + } + Run({4, 2048}, 2048, 16, true, true); + Run({4, 8192}, 8192, 64, true, true); + Run({3, 2049}, 2049, 20, true, true); + Run({2, 2, 5120}, 5120, 48, true, true); +} + +TEST_F(WebGpuMatMulGemvTest, AutomaticSelection) { + // Non-aligned K excludes subgroup-matrix kernels without forcing GEMV. + for (int64_t m = 1; m <= 8; ++m) { + Run({m, 2049}, 2049, 48, true); + } +} + +TEST_F(WebGpuMatMulGemvTest, GuardFallbacks) { + Run({9, 2049}, 2049, 48, false); + Run({4, 2047}, 2047, 48, false); + Run({4, 8193}, 8193, 48, false); + Run({4, 5120}, 5120, 12, false); + Run({4, 5120}, 5120, 68, false); + Run({4, 5120}, 5120, 49, false); +} + +TEST_F(WebGpuMatMulGemvTest, BatchedWeightsFallback) { + Run({2, 2, 2049}, 2049, 48, false, /*force_gemv=*/false, /*batched_weights=*/true); +} + +TEST_F(WebGpuMatMulGemvTest, ForcedGemvRejectsTooManyRows) { + constexpr int64_t M = 9; + constexpr int64_t K = 2049; + constexpr int64_t N = 48; + WebGpuExecutionProviderTestOptions options{}; + options.forced_matmul_algorithm = webgpu::MatMulAlgorithm::Gemv; + auto ep = WebGpuExecutionProviderWithTestOptions(options); + ASSERT_NE(ep, nullptr); + + SessionOptions session_options; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + + OpTester test("MatMul", 14); + test.AddInput("A", {M, K}, std::vector(M * K, MLFloat16{1.0f})); + test.AddInput("B", {K, N}, std::vector(K * N, MLFloat16{1.0f})); + test.AddOutput("Y", {M, N}, std::vector(M * N, MLFloat16{static_cast(K)})); + test.Config(session_options) + .ConfigEp(std::move(ep)) + .Config(OpTester::ExpectResult::kExpectFailure, + "MatMul algorithm gemv does not support these inputs or this device.") + .RunWithConfig(); } #endif // defined(USE_WEBGPU) #endif diff --git a/onnxruntime/test/providers/webgpu/matmul_execution_planner_test.cc b/onnxruntime/test/providers/webgpu/matmul_execution_planner_test.cc new file mode 100644 index 0000000000000..5f82fe53ede1f --- /dev/null +++ b/onnxruntime/test/providers/webgpu/matmul_execution_planner_test.cc @@ -0,0 +1,903 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "gtest/gtest.h" + +#include "core/providers/webgpu/math/matmul_algorithm.h" +#include "core/providers/webgpu/math/matmul_execution_planner.h" +#include "core/providers/webgpu/math/matmul_compute_dispatcher.h" +#include "core/providers/webgpu/vendor/intel/math/gemm_subgroup_utils.h" +#include "core/providers/webgpu/vendor/intel/math/matmul_execution_planner.h" +#include "core/providers/webgpu/vendor/intel/math/split_k_config.h" + +namespace onnxruntime { +namespace webgpu { +namespace test { + +namespace { + +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); + +class AlwaysPackedVendorPlanner final : public MatMulExecutionPlanner { + protected: + std::optional SelectVendorAlgorithm( + const MatMulAlgorithmSelectionParams& /*params*/) const override { + return MatMulAlgorithm::Packed; + } +}; + +class TunedPackedVendorPlanner final : public MatMulExecutionPlanner { + protected: + std::optional SelectVendorAlgorithm( + const MatMulAlgorithmSelectionParams& /*params*/) const override { + return MatMulAlgorithm::Naive; + } + + std::optional SelectVendorConfiguration( + MatMulAlgorithm algorithm, + const MatMulAlgorithmSelectionParams& params) const override { + const bool is_packed_algorithm = algorithm == MatMulAlgorithm::Packed || + algorithm == MatMulAlgorithm::PackedSplitK; + if (!is_packed_algorithm || + params.adapter_architecture != "test-architecture" || + params.a_data_type != 10 || params.b_data_type != 10) { + return std::nullopt; + } + + MatMulPackedConfiguration configuration{}; + configuration.workgroup_size = {16, 4, 1}; + configuration.elements_per_thread = {4, 2, 1}; + configuration.tile_inner = 64; + configuration.split_dim_inner = 128; + return configuration; + } +}; + +class VendorSplitKThresholdPlanner final : public MatMulExecutionPlanner { + public: + VendorSplitKThresholdPlanner() + : MatMulExecutionPlanner{SplitKConfig{ + /*max_batch_size=*/8, + /*split_dim_inner=*/128, + /*min_dim_inner_with_split_k=*/256, + {{4096, 16.0}}}} {} + + protected: + std::optional SelectVendorAlgorithm( + const MatMulAlgorithmSelectionParams& params) const override { + if (params.batch_size <= 16 && params.is_vec4 && + !params.deterministic_compute && !params.has_fused_activation && + (!params.has_bias || params.is_channels_last)) { + return MatMulAlgorithm::PackedSplitK; + } + return std::nullopt; + } +}; + +class IncompatibleConfigurationPlanner final : public MatMulExecutionPlanner { + protected: + std::optional SelectVendorConfiguration( + MatMulAlgorithm algorithm, + const MatMulAlgorithmSelectionParams& /*params*/) const override { + return algorithm == MatMulAlgorithm::Packed + ? std::optional{MatMulGemvConfiguration{}} + : std::nullopt; + } +}; + +class InvalidPackedConfigurationPlanner final : public MatMulExecutionPlanner { + protected: + std::optional SelectVendorConfiguration( + MatMulAlgorithm algorithm, + const MatMulAlgorithmSelectionParams& /*params*/) const override { + if (algorithm != MatMulAlgorithm::Packed) { + return std::nullopt; + } + MatMulPackedConfiguration configuration{}; + configuration.workgroup_size[0] = 0; + return configuration; + } +}; + +class MismatchedSubgroupConfigurationPlanner final : public MatMulExecutionPlanner { + protected: + std::optional SelectVendorConfiguration( + MatMulAlgorithm algorithm, + const MatMulAlgorithmSelectionParams& /*params*/) const override { + return algorithm == MatMulAlgorithm::Subgroup + ? std::optional{MatMulSubgroupConfiguration{8}} + : std::nullopt; + } +}; + +MatMulExecutionPlan CreatePlan( + const MatMulExecutionPlanner& planner, + const MatMulAlgorithmSelectionParams& params, + std::optional forced_algorithm = std::nullopt) { + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + const Status status = planner.CreateExecutionPlan(params, forced_algorithm, plan); + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); + return plan; +} + +} // namespace + +TEST(MatMulAlgorithmNameTest, ReturnsEveryAlgorithmName) { + struct TestCase { + std::string_view name; + MatMulAlgorithm algorithm; + }; + + constexpr TestCase test_cases[] = { + {"subgroup_matrix", MatMulAlgorithm::SubgroupMatrix}, + {"gemv", MatMulAlgorithm::Gemv}, + {"naive", MatMulAlgorithm::Naive}, + {"subgroup", MatMulAlgorithm::Subgroup}, + {"packed", MatMulAlgorithm::Packed}, + {"packed_split_k", MatMulAlgorithm::PackedSplitK}, + }; + + for (const auto& test_case : test_cases) { + SCOPED_TRACE(test_case.name); + EXPECT_EQ(MatMulAlgorithmName(test_case.algorithm), test_case.name); + } +} + +TEST(SplitKConfigTest, IntelArchitectureProfilesPreserveCurrentBoundaries) { + const SplitKConfig discrete_config = intel::CreateSplitKConfig("xe-2lpg"); + EXPECT_EQ(discrete_config.GetSplitDimInner(), 256u); + EXPECT_TRUE(discrete_config.UseSplitK( + /*is_vec4=*/true, ActivationKind::None, /*batch_size=*/1, + /*dim_a_outer=*/192, /*dim_b_outer=*/160, /*dim_inner=*/1024)); + + const SplitKConfig xe3_config = intel::CreateSplitKConfig("xe-3lpg"); + EXPECT_FALSE(xe3_config.UseSplitK( + /*is_vec4=*/true, ActivationKind::None, /*batch_size=*/1, + /*dim_a_outer=*/192, /*dim_b_outer=*/160, /*dim_inner=*/1024)); + EXPECT_TRUE(xe3_config.UseSplitK( + /*is_vec4=*/true, ActivationKind::None, /*batch_size=*/1, + /*dim_a_outer=*/128, /*dim_b_outer=*/128, /*dim_inner=*/1024)); + + const SplitKConfig default_config = intel::CreateSplitKConfig("gen-12lp"); + EXPECT_FALSE(default_config.UseSplitK( + /*is_vec4=*/true, ActivationKind::None, /*batch_size=*/1, + /*dim_a_outer=*/128, /*dim_b_outer=*/128, /*dim_inner=*/1024)); + + const SplitKConfig legacy_config = intel::CreateSplitKConfig("gen-9"); + EXPECT_EQ(legacy_config.GetSplitDimInner(), 0u); + EXPECT_FALSE(legacy_config.UseSplitK( + /*is_vec4=*/true, ActivationKind::None, /*batch_size=*/1, + /*dim_a_outer=*/1, /*dim_b_outer=*/1, /*dim_inner=*/1024)); +} + +TEST(SplitKConfigTest, FactoryRoutesOnlySupportedVendorProfiles) { + EXPECT_EQ(CreateSplitKConfig("intel", "xe-2lpg").GetSplitDimInner(), 256u); + EXPECT_EQ(CreateSplitKConfig("nvidia", "pascal").GetSplitDimInner(), 0u); +} + +TEST(MatMulExecutionPlannerTest, ForcedAlgorithmTakesPrecedence) { + AlwaysPackedVendorPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.can_use_subgroup_matrix = true; + + EXPECT_EQ(planner.SelectAlgorithm(params, MatMulAlgorithm::Naive), MatMulAlgorithm::Naive); +} + +TEST(MatMulExecutionPlannerTest, ReevaluatesSelectionForEachRuntimeShape) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 4; + params.packed_m = 4; + params.n = 7; + params.k = 7; + + EXPECT_EQ(CreatePlan(planner, params).algorithm, MatMulAlgorithm::Naive); + + params.m = 64; + params.packed_m = 64; + params.n = 64; + params.k = 64; + EXPECT_EQ(CreatePlan(planner, params).algorithm, MatMulAlgorithm::Packed); +} + +TEST(MatMulExecutionPlannerTest, VendorCanTuneForcedAlgorithmConfiguration) { + TunedPackedVendorPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.adapter_architecture = "test-architecture"; + params.a_data_type = 10; + params.b_data_type = 10; + params.k = 64; + + const MatMulExecutionPlan plan = + CreatePlan(planner, params, MatMulAlgorithm::Packed); + + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Packed); + const auto& configuration = + std::get(plan.configuration); + EXPECT_EQ(configuration.workgroup_size, + (std::array{16, 4, 1})); + EXPECT_EQ(configuration.elements_per_thread, + (std::array{4, 2, 1})); + EXPECT_EQ(configuration.tile_inner, 64u); + EXPECT_EQ(configuration.split_dim_inner, 128u); +} + +TEST(MatMulExecutionPlannerTest, VendorCanSetIndependentSplitKThresholds) { + VendorSplitKThresholdPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 32; + params.n = 64; + params.k = 1024; + params.batch_size = 16; + params.packed_batch_size = 16; + params.is_vec4 = true; + params.is_channels_last = true; + + const MatMulExecutionPlan plan = CreatePlan(planner, params); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::PackedSplitK); + EXPECT_EQ(std::get(plan.configuration).split_dim_inner, 128u); + + params.batch_size = 17; + params.packed_batch_size = 17; + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Packed); +} + +TEST(MatMulExecutionPlannerTest, InjectedSplitKPolicyCreatesPackedSplitKPlan) { + intel::IntelMatMulExecutionPlanner planner{intel::CreateSplitKConfig("xe-2lpg")}; + MatMulAlgorithmSelectionParams params{}; + params.m = 192; + params.packed_m = 192; + params.n = 160; + params.k = 1024; + params.is_vec4 = true; + + const MatMulExecutionPlan plan = CreatePlan(planner, params); + + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::PackedSplitK); + const auto& configuration = + std::get(plan.configuration); + EXPECT_EQ(configuration.split_dim_inner, 256u); +} + +TEST(MatMulExecutionPlannerTest, CommonPackedConfigurationPreservesCurrentTuning) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 8; + params.packed_m = 8; + params.n = 64; + params.k = 64; + + const MatMulExecutionPlan small_m_plan = CreatePlan(planner, params); + const auto& small_m_configuration = + std::get(small_m_plan.configuration); + EXPECT_EQ(small_m_configuration.workgroup_size, + (std::array{8, 8, 1})); + EXPECT_EQ(small_m_configuration.elements_per_thread, + (std::array{4, 1, 1})); + EXPECT_EQ(small_m_configuration.tile_inner, 32u); + + params.m = 9; + params.packed_m = 9; + const MatMulExecutionPlan large_m_plan = CreatePlan(planner, params); + const auto& large_m_configuration = + std::get(large_m_plan.configuration); + EXPECT_EQ(large_m_configuration.elements_per_thread, + (std::array{4, 4, 1})); +} + +TEST(MatMulAlgorithmConfigurationTest, PackedConfigurationKeepsBatchAxesUntiled) { + MatMulPackedConfiguration configuration{}; + EXPECT_TRUE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/false, + /*use_split_k=*/false)); + EXPECT_TRUE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/true, + /*use_split_k=*/false)); + + configuration.workgroup_size[2] = 2; + EXPECT_FALSE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/false, + /*use_split_k=*/false)); + + configuration.workgroup_size[2] = 1; + configuration.elements_per_thread[2] = 2; + EXPECT_FALSE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/false, + /*use_split_k=*/false)); +} + +TEST(MatMulAlgorithmConfigurationTest, PackedVec4RequiresSupportedCooperativeLoads) { + MatMulPackedConfiguration configuration{}; + + configuration.elements_per_thread[0] = 8; + EXPECT_FALSE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/true, + /*use_split_k=*/false)); + + configuration.elements_per_thread[0] = 4; + configuration.workgroup_size = {16, 4, 1}; + configuration.tile_inner = 16; + EXPECT_FALSE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/true, + /*use_split_k=*/false)); + + configuration.tile_inner = 48; + EXPECT_FALSE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/true, + /*use_split_k=*/false)); + + configuration.tile_inner = 60; + EXPECT_FALSE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/true, + /*use_split_k=*/false)); +} + +TEST(MatMulAlgorithmConfigurationTest, SubgroupConfigurationCarriesSelectedSize) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 1; + params.n = 1; + params.k = 1; + params.has_subgroup_capability = true; + params.subgroup_size = 16; + + const auto plan = CreatePlan(planner, params, MatMulAlgorithm::Subgroup); + const auto* configuration = std::get_if(&plan.configuration); + ASSERT_NE(configuration, nullptr); + EXPECT_EQ(configuration->subgroup_size, 16u); +} + +TEST(MatMulAlgorithmConfigurationTest, SubgroupSizeSelectionRequiresSupportedSize) { + EXPECT_EQ(intel::SelectMatMulSubgroupSize(8, 8, false), 8u); + EXPECT_EQ(intel::SelectMatMulSubgroupSize(16, 16, false), 16u); + EXPECT_EQ(intel::SelectMatMulSubgroupSize(32, 32, false), 32u); + EXPECT_FALSE(intel::SelectMatMulSubgroupSize(64, 64, false).has_value()); + EXPECT_FALSE(intel::SelectMatMulSubgroupSize(8, 32, false).has_value()); + EXPECT_EQ(intel::SelectMatMulSubgroupSize(8, 32, true), 32u); + EXPECT_EQ(intel::SelectMatMulSubgroupSize(8, 16, true), 16u); + EXPECT_FALSE(intel::SelectMatMulSubgroupSize(4, 4, true).has_value()); +} + +TEST(MatMulAlgorithmConfigurationTest, SplitKConfigurationRequiresTileAlignedSplits) { + MatMulPackedConfiguration configuration{}; + configuration.tile_inner = 32; + configuration.split_dim_inner = 64; + EXPECT_TRUE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/true, + /*use_split_k=*/true)); + + configuration.split_dim_inner = 48; + EXPECT_FALSE(IsMatMulPackedConfigurationValid(configuration, /*is_vec4=*/true, + /*use_split_k=*/true)); +} + +TEST(MatMulAlgorithmConfigurationTest, PackedDispatchArithmeticIsOverflowSafe) { + const auto maximum_tuning_dispatch = TryGetMatMulPackedDispatchGroupCount( + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max()); + ASSERT_TRUE(maximum_tuning_dispatch.has_value()); + EXPECT_EQ(*maximum_tuning_dispatch, 1u); + + EXPECT_EQ(TryGetMatMulPackedDispatchGroupCount( + std::numeric_limits::max(), 1, 1), + std::nullopt); +} + +TEST(MatMulExecutionPlannerTest, CommonFallbackPrefersSubgroupMatrix) { + MatMulExecutionPlanner planner{intel::CreateSplitKConfig("xe-2lpg")}; + MatMulAlgorithmSelectionParams params{}; + params.m = 64; + params.packed_m = 64; + params.n = 64; + params.k = 1024; + params.can_use_subgroup_matrix = true; + params.is_vec4 = true; + + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::SubgroupMatrix); +} + +TEST(MatMulExecutionPlannerTest, GemvPrecedesGenericFallbacks) { + MatMulExecutionPlanner planner{intel::CreateSplitKConfig("xe-2lpg")}; + MatMulAlgorithmSelectionParams params{}; + params.m = 1; + params.packed_m = 1; + params.n = 48; + params.k = 5120; + params.can_run_gemv = true; + params.is_vec4 = true; + + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Gemv); + + params.can_use_subgroup_matrix = true; + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::SubgroupMatrix); +} + +TEST(MatMulExecutionPlannerTest, GemvAutomaticSelectionUsesPerformanceThresholds) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 1; + params.n = 16; + params.k = 2048; + params.can_run_gemv = true; + + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Gemv); + + params.n = 12; + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Packed); + params.n = 68; + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Packed); + params.n = 16; + + params.k = 2047; + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Packed); + params.k = 8193; + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Packed); +} + +TEST(MatMulAlgorithmConfigurationTest, GemvUsesTypedEmptyConfiguration) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 1; + params.n = 48; + params.k = 5120; + params.can_run_gemv = true; + + const MatMulExecutionPlan plan = CreatePlan(planner, params); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Gemv); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); +} + +TEST(MatMulExecutionPlannerTest, VendorPolicyPrecedesCommonHeuristics) { + AlwaysPackedVendorPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.n = 4; + params.k = 4; + params.can_use_subgroup_matrix = true; + + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Packed); +} + +TEST(MatMulExecutionPlannerTest, ZeroContractionDimensionPrecedesVendorPolicy) { + AlwaysPackedVendorPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.n = 8; + params.k = 0; + + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Naive); +} + +TEST(MatMulExecutionPlannerTest, NaiveUsesStrictSmallDimensionBoundaries) { + MatMulExecutionPlanner planner; + + MatMulAlgorithmSelectionParams small_params{}; + small_params.n = 7; + small_params.k = 7; + EXPECT_EQ(planner.SelectAlgorithm(small_params), MatMulAlgorithm::Naive); + + MatMulAlgorithmSelectionParams n_boundary = small_params; + n_boundary.n = 8; + EXPECT_EQ(planner.SelectAlgorithm(n_boundary), MatMulAlgorithm::Packed); + + MatMulAlgorithmSelectionParams k_boundary = small_params; + k_boundary.k = 8; + EXPECT_EQ(planner.SelectAlgorithm(k_boundary), MatMulAlgorithm::Packed); +} + +TEST(MatMulExecutionPlannerTest, ZeroContractionDimensionUsesNaive) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.n = 8; + params.k = 0; + + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Naive); +} + +TEST(MatMulExecutionPlannerTest, IntelPlannerAppliesCurrentVendorRule) { + intel::IntelMatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 64; + params.n = 512; + params.k = 32; + params.has_subgroup_capability = true; + + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Subgroup); + + params.n = 511; + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::Packed); +} + +TEST(MatMulExecutionPlannerTest, IntelPlannerPreservesSubgroupMatrixPrecedence) { + intel::IntelMatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 64; + params.n = 512; + params.k = 32; + params.can_use_subgroup_matrix = true; + params.has_subgroup_capability = true; + + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::SubgroupMatrix); +} + +TEST(MatMulExecutionPlannerTest, SplitKPrecedesPackedFallback) { + MatMulExecutionPlanner planner{intel::CreateSplitKConfig("xe-2lpg")}; + MatMulAlgorithmSelectionParams params{}; + params.m = 64; + params.packed_m = 64; + params.n = 64; + params.k = 1024; + params.is_vec4 = true; + EXPECT_EQ(planner.SelectAlgorithm(params), MatMulAlgorithm::PackedSplitK); + + MatMulExecutionPlanner disabled_planner; + EXPECT_EQ(disabled_planner.SelectAlgorithm(params), MatMulAlgorithm::Packed); +} + +TEST(MatMulExecutionPlanValidationTest, SplitKRejectsEachHardConstraint) { + MatMulExecutionPlanner planner{SplitKConfig{ + /*max_batch_size=*/8, + /*split_dim_inner=*/128, + /*min_dim_inner_with_split_k=*/256, + {{4096, 16.0}}}}; + MatMulAlgorithmSelectionParams params{}; + params.m = 32; + params.packed_m = 32; + params.n = 64; + params.k = 256; + params.is_vec4 = true; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::PackedSplitK, plan).IsOK()); + + params.deterministic_compute = true; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::PackedSplitK, plan).IsOK()); + params.deterministic_compute = false; + + params.is_vec4 = false; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::PackedSplitK, plan).IsOK()); + params.is_vec4 = true; + + params.has_fused_activation = true; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::PackedSplitK, plan).IsOK()); + params.has_fused_activation = false; + + params.has_bias = true; + params.is_channels_last = false; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::PackedSplitK, plan).IsOK()); +} + +TEST(MatMulExecutionPlanValidationTest, ForcedFailureDoesNotFallBackOrModifyOutput) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 1; + params.n = 48; + params.k = 5120; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + const Status status = planner.CreateExecutionPlan(params, MatMulAlgorithm::Gemv, plan); + + EXPECT_FALSE(status.IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Naive); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); +} + +TEST(MatMulExecutionPlanValidationTest, GemvRequiresCompatibleInputs) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 1; + params.n = 48; + params.k = 5120; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Gemv, plan).IsOK()); + params.can_run_gemv = true; + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Gemv, plan).IsOK()); +} + +TEST(MatMulExecutionPlanValidationTest, GemvRejectsDimensionsOutsideUint32Range) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 1; + params.n = std::numeric_limits::max() - 3; + params.k = std::numeric_limits::max(); + params.can_run_gemv = true; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Gemv, plan).IsOK()); + + plan = {MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + params.n = static_cast(std::numeric_limits::max()) + 1; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Gemv, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Naive); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); + + params.n = 4; + params.k = static_cast(std::numeric_limits::max()) + 1; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Gemv, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Naive); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); +} + +TEST(MatMulExecutionPlanValidationTest, NaiveRejectsDimensionsAndVectorizedOutputOutsideUint32Range) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = std::numeric_limits::max(); + params.n = 1; + MatMulExecutionPlan plan{MatMulAlgorithm::Packed, MatMulPackedConfiguration{}}; + + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Naive, plan).IsOK()); + + plan = {MatMulAlgorithm::Packed, MatMulPackedConfiguration{}}; + params.m = static_cast(std::numeric_limits::max()) + 1; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Naive, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Packed); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); + + params.m = 1; + params.n = static_cast(std::numeric_limits::max()) + 1; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Naive, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Packed); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); + + params.n = 1; + params.k = static_cast(std::numeric_limits::max()) + 1; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Naive, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Packed); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); + + params.k = 0; + params.m = std::numeric_limits::max(); + params.batch_size = 2; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Naive, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Packed); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); +} + +TEST(MatMulExecutionPlanValidationTest, PackedAlgorithmsRejectZeroContractionDimension) { + MatMulExecutionPlanner planner{SplitKConfig{ + /*max_batch_size=*/8, + /*split_dim_inner=*/128, + /*min_dim_inner_with_split_k=*/256, + {{4096, 16.0}}}}; + MatMulAlgorithmSelectionParams params{}; + params.m = 32; + params.packed_m = 32; + params.n = 64; + params.is_vec4 = true; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::PackedSplitK, plan).IsOK()); + + params.k = 256; + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::PackedSplitK, plan).IsOK()); +} + +TEST(MatMulExecutionPlanValidationTest, RejectsInvalidVendorConfigurations) { + MatMulAlgorithmSelectionParams params{}; + params.k = 64; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + IncompatibleConfigurationPlanner incompatible_planner; + EXPECT_FALSE(incompatible_planner.CreateExecutionPlan( + params, MatMulAlgorithm::Packed, plan) + .IsOK()); + + InvalidPackedConfigurationPlanner invalid_packed_planner; + EXPECT_FALSE(invalid_packed_planner.CreateExecutionPlan( + params, MatMulAlgorithm::Packed, plan) + .IsOK()); + + MismatchedSubgroupConfigurationPlanner mismatched_subgroup_planner; + params.has_subgroup_capability = true; + params.subgroup_size = 16; + EXPECT_FALSE(mismatched_subgroup_planner.CreateExecutionPlan( + params, MatMulAlgorithm::Subgroup, plan) + .IsOK()); +} + +TEST(MatMulExecutionPlanValidationTest, RejectsPackedDimensionsOutsideUint32Range) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.packed_m = 32; + params.n = 64; + params.k = 64; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + constexpr int64_t max_shader_dimension = std::numeric_limits::max(); + + params.packed_m = max_shader_dimension; + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); + params.packed_m = max_shader_dimension + 1; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); + + params.packed_m = 32; + params.n = max_shader_dimension; + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); + params.n = max_shader_dimension + 1; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); + + params.n = 64; + params.k = max_shader_dimension; + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); + params.k = max_shader_dimension + 1; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); +} + +TEST(MatMulExecutionPlanValidationTest, RejectsPackedDispatchOverflow) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.k = 64; + params.packed_batch_size = static_cast(std::numeric_limits::max()) + 1; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); +} + +TEST(MatMulExecutionPlanValidationTest, RejectsPackedCombinedDispatchOverflow) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.packed_m = 2097152; + params.n = 2097152; + params.k = 64; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Packed, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Naive); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); +} + +TEST(MatMulExecutionPlanValidationTest, RejectsSplitKDispatchOverflow) { + MatMulExecutionPlanner planner{SplitKConfig{ + /*max_batch_size=*/std::numeric_limits::max(), + /*split_dim_inner=*/128, + /*min_dim_inner_with_split_k=*/256, + {{4096, 16.0}}}}; + MatMulAlgorithmSelectionParams params{}; + params.m = 32; + params.packed_m = 32; + params.n = 64; + params.k = 256; + params.packed_batch_size = std::numeric_limits::max(); + params.is_vec4 = true; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan( + params, MatMulAlgorithm::PackedSplitK, plan) + .IsOK()); +} + +TEST(MatMulExecutionPlanValidationTest, RejectsSplitKCombinedDispatchOverflow) { + MatMulExecutionPlanner planner{SplitKConfig{ + /*max_batch_size=*/8, + /*split_dim_inner=*/128, + /*min_dim_inner_with_split_k=*/256, + {{4096, 16.0}}}}; + MatMulAlgorithmSelectionParams params{}; + params.packed_m = 1048576; + params.n = 2097152; + params.k = 256; + params.is_vec4 = true; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan( + params, MatMulAlgorithm::PackedSplitK, plan) + .IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Naive); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); +} + +TEST(MatMulAlgorithmPrerequisiteTest, IntelAVec4RequiresCompatibleRowsPerThread) { + EXPECT_FALSE(intel::CanUseAVec4CooperativeLoad(intel::gpu_arch::kXe3Lpg, 32, 1)); + EXPECT_FALSE(intel::CanUseAVec4CooperativeLoad(intel::gpu_arch::kXe3Lpg, 32, 2)); + EXPECT_TRUE(intel::CanUseAVec4CooperativeLoad(intel::gpu_arch::kXe3Lpg, 32, 4)); + EXPECT_FALSE(intel::CanUseAVec4CooperativeLoad(intel::gpu_arch::kXe3Lpg, 31, 4)); + EXPECT_FALSE(intel::CanUseAVec4CooperativeLoad(intel::gpu_arch::kXeLpg, 32, 4)); +} + +TEST(MatMulExecutionPlanValidationTest, IntelCapabilityDoesNotIncludeAutomaticThresholds) { + MatMulAlgorithmSelectionParams below_heuristic_threshold{}; + below_heuristic_threshold.m = 1; + below_heuristic_threshold.n = 1; + below_heuristic_threshold.k = 1; + below_heuristic_threshold.has_subgroup_capability = true; + below_heuristic_threshold.subgroup_size = 16; + intel::IntelMatMulExecutionPlanner planner; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_TRUE(planner.CreateExecutionPlan( + below_heuristic_threshold, MatMulAlgorithm::Subgroup, plan) + .IsOK()); + EXPECT_NE(planner.SelectAlgorithm(below_heuristic_threshold), MatMulAlgorithm::Subgroup); +} + +TEST(MatMulExecutionPlanValidationTest, SubgroupRejectsZeroContractionDimension) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 1; + params.n = 1; + params.has_subgroup_capability = true; + params.subgroup_size = 16; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Subgroup, plan).IsOK()); + params.k = 1; + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Subgroup, plan).IsOK()); +} + +TEST(MatMulExecutionPlanValidationTest, SubgroupRejectsUnsupportedConfiguredSize) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 1; + params.n = 1; + params.has_subgroup_capability = true; + params.k = 1; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + for (uint32_t subgroup_size : {8u, 16u, 32u}) { + params.subgroup_size = subgroup_size; + EXPECT_TRUE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Subgroup, plan).IsOK()); + } + + params.subgroup_size = 64; + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Subgroup, plan).IsOK()); +} + +TEST(MatMulExecutionPlanValidationTest, SubgroupRejectsDimensionsOutsideUint32Range) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = static_cast(std::numeric_limits::max()) + 1; + params.packed_m = params.m; + params.n = 1; + params.k = 1; + params.has_subgroup_capability = true; + params.subgroup_size = 16; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Subgroup, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Naive); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); +} + +TEST(MatMulExecutionPlanValidationTest, SubgroupRejectsCombinedDispatchOverflow) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 4194304; + params.packed_m = params.m; + params.n = 8388608; + params.k = 1; + params.has_subgroup_capability = true; + params.subgroup_size = 16; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Subgroup, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Naive); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); +} + +TEST(MatMulExecutionPlanValidationTest, SubgroupRejectsFoldedBatchDispatchOverflow) { + MatMulExecutionPlanner planner; + MatMulAlgorithmSelectionParams params{}; + params.m = 2097152; + params.packed_m = 4194304; + params.n = 8388608; + params.k = 1; + params.batch_size = 2; + params.packed_batch_size = 1; + params.has_subgroup_capability = true; + params.subgroup_size = 16; + MatMulExecutionPlan plan{MatMulAlgorithm::Naive, MatMulNaiveConfiguration{}}; + + EXPECT_FALSE(planner.CreateExecutionPlan(params, MatMulAlgorithm::Subgroup, plan).IsOK()); + EXPECT_EQ(plan.algorithm, MatMulAlgorithm::Naive); + EXPECT_TRUE(std::holds_alternative(plan.configuration)); +} + +} // namespace test +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/webgpu/matmul_large_test.cc b/onnxruntime/test/providers/webgpu/matmul_large_test.cc index 3375fc5e6b53b..baa0697599d41 100644 --- a/onnxruntime/test/providers/webgpu/matmul_large_test.cc +++ b/onnxruntime/test/providers/webgpu/matmul_large_test.cc @@ -3,12 +3,33 @@ #include #include +#include #include +#include +#include +#include +#include #include #include "gtest/gtest.h" -#include "test/providers/provider_test_utils.h" + +#include "core/graph/onnx_protobuf.h" +#include "core/providers/cpu/math/matmul_helper.h" +#include "core/providers/webgpu/math/matmul_algorithm.h" +#include "core/providers/webgpu/webgpu_provider_options.h" +#if !defined(ORT_USE_EP_API_ADAPTERS) +#include "core/providers/webgpu/math/subgroup_matrix_config.h" +#include "core/providers/webgpu/vendor/intel/math/gemm_subgroup_utils.h" +#include "core/providers/webgpu/webgpu_context.h" +#endif +#include "core/session/onnxruntime_session_options_config_keys.h" #include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/providers/webgpu/graph_capture_test_utils.h" +#include "test/test_environment.h" +#include "test/unittest_util/framework_test_utils.h" +#include "test/util/include/asserts.h" +#include "test/util/include/inference_session_wrapper.h" #include "default_providers.h" namespace onnxruntime { @@ -72,15 +93,98 @@ static void ComputeExpectedResult(std::initializer_list a_dims, if (!b_is_vector) output_dims.push_back(N); } +static std::optional GetForcedAlgorithmUnsupportedReason( + const IExecutionProvider& ep, + webgpu::MatMulAlgorithm algorithm) { +#if defined(ORT_USE_EP_API_ADAPTERS) + ORT_UNUSED_PARAMETER(ep); + switch (algorithm) { + case webgpu::MatMulAlgorithm::Subgroup: + case webgpu::MatMulAlgorithm::PackedSplitK: + case webgpu::MatMulAlgorithm::SubgroupMatrix: + return "hardware-specific forced MatMul tests require direct adapter capability inspection."; + case webgpu::MatMulAlgorithm::Gemv: + case webgpu::MatMulAlgorithm::Naive: + case webgpu::MatMulAlgorithm::Packed: + return std::nullopt; + } + return std::nullopt; +#else + auto& context = webgpu::WebGpuContextFactory::GetContext(ep.GetDeviceId()); + + switch (algorithm) { + case webgpu::MatMulAlgorithm::Subgroup: { + if (!context.DeviceHasFeature(wgpu::FeatureName::Subgroups)) { + return "subgroup requires the WebGPU Subgroups feature."; + } + const auto& adapter_info = context.AdapterInfo(); + if (!webgpu::intel::SelectMatMulSubgroupSize( + adapter_info.subgroupMinSize, + adapter_info.subgroupMaxSize, + context.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl)) + .has_value()) { + return "subgroup requires a supported subgroup size (8, 16, or 32)."; + } + break; + } + case webgpu::MatMulAlgorithm::PackedSplitK: + if (context.GetSplitKConfig().GetSplitDimInner() == 0) { + return "packed_split_k is not configured for the selected adapter."; + } + break; + case webgpu::MatMulAlgorithm::SubgroupMatrix: { + if (!context.DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { + return "subgroup_matrix requires the WebGPU subgroup-matrix feature."; + } + + const auto& adapter_info = context.AdapterInfo(); + const auto& device_configs = context.SubgroupMatrixConfigs(); + constexpr auto kF16 = wgpu::SubgroupMatrixComponentType::F16; + if (!webgpu::detail::SelectSubgroupMatrixConfigFromAdapterConfigs( + {device_configs.configs, device_configs.configCount}, + adapter_info.subgroupMinSize, adapter_info.subgroupMaxSize, + context.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl), + {{kF16, kF16, 8, 16, 16, 32, false}, + {kF16, kF16, 16, 16, 16, 32, false}})) { + return "subgroup_matrix requires an 8x16x16 or 16x16x16 F16 configuration with subgroup size 32."; + } + break; + } + case webgpu::MatMulAlgorithm::Gemv: + case webgpu::MatMulAlgorithm::Naive: + case webgpu::MatMulAlgorithm::Packed: + break; + } + + return std::nullopt; +#endif +} + template void RunTestTyped(std::initializer_list a_dims, std::initializer_list b_dims, - bool b_is_constant = false) { + bool b_is_constant = false, + std::optional forced_algorithm = std::nullopt, + OpTester::ExpectResult expected_result = OpTester::ExpectResult::kExpectSuccess, + const char* expected_error = nullptr) { static_assert(std::is_same_v || std::is_same_v, "unexpected type for T"); - auto webgpu_ep = DefaultWebGpuExecutionProvider(); + std::unique_ptr webgpu_ep; + if (forced_algorithm.has_value()) { + WebGpuExecutionProviderTestOptions test_options{}; + test_options.forced_matmul_algorithm = forced_algorithm; + webgpu_ep = WebGpuExecutionProviderWithTestOptions(test_options); + } else { + webgpu_ep = DefaultWebGpuExecutionProvider(); + } if (!webgpu_ep) { GTEST_SKIP() << "WebGPU execution provider is not available."; } + if (forced_algorithm.has_value()) { + if (const auto reason = GetForcedAlgorithmUnsupportedReason(*webgpu_ep, *forced_algorithm); + reason.has_value()) { + GTEST_SKIP() << *reason; + } + } RandomValueGenerator random{1234}; std::vector a_vals(random.Gaussian(AsSpan(a_dims), 0.0f, 0.25f)); @@ -103,7 +207,12 @@ void RunTestTyped(std::initializer_list a_dims, std::initializer_list @@ -112,6 +221,149 @@ void RunBothTypes(std::initializer_list a_dims, std::initializer_list(a_dims, b_dims); } +template +void RunZeroKTest(std::initializer_list a_dims, + std::initializer_list b_dims, bool graph_capture = false) { + ConfigOptions config; + ASSERT_STATUS_OK(config.AddConfigEntry( + webgpu::options::kEnableGraphCapture, + graph_capture ? webgpu::options::kEnableGraphCapture_ON : webgpu::options::kEnableGraphCapture_OFF)); + ASSERT_STATUS_OK(config.AddConfigEntry( + webgpu::options::kValidationMode, webgpu::options::kValidationMode_full)); + auto ep = WebGpuExecutionProviderWithOptions(config); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + MatMulComputeHelper helper; + ASSERT_STATUS_OK(helper.Compute(TensorShape(a_dims), TensorShape(b_dims))); + ASSERT_EQ(helper.K(), 0); + TensorShapeVector output_dims; + std::vector expected; + ComputeExpectedResult(a_dims, b_dims, {}, {}, output_dims, expected); + ASSERT_EQ(helper.OutputShape(), TensorShape(output_dims)); + WebGpuGraphCaptureTester test("MatMul", 13); + test.AddInput("A", a_dims, std::vector{}); + test.AddInput("B", b_dims, std::vector{}); + const std::vector expected_output(expected.size(), T{0.0f}); + test.AddOutput("Y", output_dims, expected_output); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + if (graph_capture) { + test.RunWithBoundCapture(std::move(ep), options, TensorShape(output_dims), expected_output); + } else { + test.Config(options).ConfigEp(std::move(ep)).RunWithConfig(); + } +} + +TEST(MatMulZeroKTest, Float32MatrixAndBatch) { + RunZeroKTest({2, 0}, {0, 3}); + RunZeroKTest({1, 35, 0}, {0, 1024}); + RunZeroKTest({2, 1, 3, 0}, {1, 4, 0, 5}); +} + +TEST(MatMulZeroKTest, HelperPreservesBroadcastBatchesWithZeroK) { + MatMulComputeHelper helper; + ASSERT_STATUS_OK(helper.Compute(TensorShape({2, 1, 3, 0}), TensorShape({1, 4, 0, 5}))); + EXPECT_EQ(helper.OutputShape(), TensorShape({2, 4, 3, 5})); + ASSERT_STATUS_OK(helper.Compute(TensorShape({2, 1, 3, 4}), TensorShape({1, 4, 4, 5}))); + EXPECT_EQ(helper.OutputShape(), TensorShape({2, 4, 3, 5})); +} + +TEST(MatMulZeroKTest, Float16MatrixAndBatch) { + RunZeroKTest({2, 0}, {0, 3}); + RunZeroKTest({1, 35, 0}, {0, 1024}); + RunZeroKTest({2, 1, 3, 0}, {1, 4, 0, 5}); +} + +TEST(MatMulZeroKTest, VectorAndScalarOutputs) { + RunZeroKTest({0}, {0, 5}); + RunZeroKTest({3, 0}, {0}); + RunZeroKTest({0}, {0}); + RunZeroKTest({0}, {0, 5}); + RunZeroKTest({0}, {0}); +} + +TEST(MatMulZeroKTest, EmptyOutput) { + RunZeroKTest({0, 0}, {0, 3}); + RunZeroKTest({1, 3, 0}, {0, 0, 5}); + RunZeroKTest({2, 1, 3, 0}, {1, 0, 0, 5}); + RunZeroKTest({0, 1, 3, 0}, {1, 2, 0, 5}); + RunZeroKTest({0, 0}, {0, 3}); + RunZeroKTest({1, 3, 0}, {0, 0, 5}); + RunZeroKTest({2, 1, 3, 0}, {1, 0, 0, 5}); + RunZeroKTest({0, 1, 3, 0}, {1, 2, 0, 5}); +} + +TEST(MatMulZeroKTest, GraphCaptureReplay) { + RunZeroKTest({1, 35, 0}, {0, 1024}, true); + RunZeroKTest({1, 35, 0}, {0, 1024}, true); +} + +template +void RunConvMatMulTest(bool channels_last, bool initializer, bool empty_output = false) { + ConfigOptions config; + ASSERT_STATUS_OK(config.AddConfigEntry(webgpu::options::kValidationMode, + webgpu::options::kValidationMode_full)); + ASSERT_STATUS_OK(config.AddConfigEntry(webgpu::options::kEnableGraphCapture, + webgpu::options::kEnableGraphCapture_ON)); + auto ep = WebGpuExecutionProviderWithOptions(config); + if (!ep) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + WebGpuGraphCaptureTester test("Conv", 11, channels_last ? kMSInternalNHWCDomain : kOnnxDomain); + test.AddAttribute("kernel_shape", std::vector{1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0}); + test.AddAttribute("strides", std::vector{1, 1}); + if (channels_last) { + test.AddAttribute("activation", std::string{"Clip"}); + test.AddAttribute("activation_params", std::vector{-0.5f, 2.0f}); + } + const int64_t batch = empty_output ? 0 : 2; + const int64_t channels = empty_output ? 4 : 0; + const std::vector input_shape = channels_last ? std::vector{batch, 2, 3, channels} + : std::vector{batch, channels, 2, 3}; + const std::vector output_shape = channels_last ? std::vector{batch, 2, 3, 3} + : std::vector{batch, 3, 2, 3}; + test.AddInput("X", input_shape, std::vector{}); + test.AddInput("W", {3, channels, 1, 1}, std::vector(3 * channels, T{1.0f}), initializer); + test.AddInput("B", {3}, std::vector{T{-2.0f}, T{-0.25f}, T{3.5f}}); + std::vector expected; + for (int64_t index = 0; index < batch * 18; ++index) { + const int64_t channel = channels_last ? index % 3 : (index / 6) % 3; + const float value = channel == 0 ? -2.0f : channel == 1 ? -0.25f + : 3.5f; + expected.push_back(T{channels_last ? std::clamp(value, -0.5f, 2.0f) : value}); + } + test.AddOutput("Y", output_shape, expected); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.RunWithBoundCapture(std::move(ep), options, TensorShape(output_shape), expected); +} + +TEST(MatMulZeroKTest, ChannelsFirstConvBias) { + for (bool initializer : {false, true}) { + RunConvMatMulTest(false, initializer); + RunConvMatMulTest(false, initializer); + } +} + +#ifndef DISABLE_CONTRIB_OPS +TEST(MatMulZeroKTest, ChannelsLastConvBiasAndActivation) { + for (bool initializer : {false, true}) { + RunConvMatMulTest(true, initializer); + RunConvMatMulTest(true, initializer); + } +} + +TEST(MatMulZeroKTest, SharedHelperEmptyOutputChannelsLastConv) { + for (bool initializer : {false, true}) { + RunConvMatMulTest(true, initializer, true); + RunConvMatMulTest(true, initializer, true); + } +} +#endif + TEST(MatMulNaiveProgramTest, Broadcast4DExecution) { RunTestTyped({3, 1, 1, 2}, {2, 2, 2}); RunTestTyped({2, 2, 3, 2}, {2, 1}); @@ -123,6 +375,10 @@ TEST(MatMulNaiveProgramTest, VectorExecution) { RunTestTyped({3}, {3}); } +TEST(MatMulNaiveProgramTest, ZeroContractionDimensionExecution) { + RunTestTyped({1, 0}, {0, 8}); +} + TEST(MatMulProgramTest, VectorFallbackExecution) { RunTestTyped({8}, {8, 3}); RunTestTyped({2, 8}, {8}); @@ -131,6 +387,161 @@ TEST(MatMulProgramTest, VectorFallbackExecution) { RunTestTyped({2, 2, 8}, {8}); } +static std::string BuildDynamicMatMulModelBytes() { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(ONNX_NAMESPACE::IR_VERSION); + auto* opset = model.add_opset_import(); + opset->set_domain(""); + opset->set_version(13); + + auto* graph = model.mutable_graph(); + graph->set_name("dynamic_matmul"); + + auto set_float_shape = [](ONNX_NAMESPACE::ValueInfoProto* value_info, + const char* first_dimension, + const char* second_dimension) { + auto* tensor_type = value_info->mutable_type()->mutable_tensor_type(); + tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + auto* shape = tensor_type->mutable_shape(); + shape->add_dim()->set_dim_param(first_dimension); + if (second_dimension != nullptr) { + shape->add_dim()->set_dim_param(second_dimension); + } else { + shape->add_dim()->set_dim_value(8); + } + }; + + auto* a = graph->add_input(); + a->set_name("A"); + set_float_shape(a, "M", "K"); + + auto* b = graph->add_input(); + b->set_name("B"); + set_float_shape(b, "K", nullptr); + + auto* y = graph->add_output(); + y->set_name("Y"); + set_float_shape(y, "M", nullptr); + + auto* node = graph->add_node(); + node->set_name("MatMul"); + node->set_op_type("MatMul"); + node->add_input("A"); + node->add_input("B"); + node->add_output("Y"); + + std::string bytes; + model.SerializeToString(&bytes); + return bytes; +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedNaive) { + RunTestTyped({8, 8}, {8, 8}, false, webgpu::MatMulAlgorithm::Naive); + RunTestTyped({1, 0}, {0, 8}, false, webgpu::MatMulAlgorithm::Naive); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedGemv) { + RunTestTyped({1, 128}, {128, 4}, false, webgpu::MatMulAlgorithm::Gemv); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedGemvRejectsFloatInputs) { + RunTestTyped({1, 2048}, {2048, 16}, false, webgpu::MatMulAlgorithm::Gemv, + OpTester::ExpectResult::kExpectFailure, + "MatMul algorithm gemv"); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedPacked) { + RunTestTyped({2, 2}, {2, 2}, false, webgpu::MatMulAlgorithm::Packed); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedPackedRejectsZeroContractionDimension) { + RunTestTyped({1, 0}, {0, 1}, false, webgpu::MatMulAlgorithm::Packed, + OpTester::ExpectResult::kExpectFailure, + "MatMul algorithm packed"); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedSubgroup) { + RunTestTyped({8, 32}, {32, 64}, false, webgpu::MatMulAlgorithm::Subgroup); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedSubgroupRejectsZeroContractionDimension) { + RunTestTyped({1, 0}, {0, 1}, false, webgpu::MatMulAlgorithm::Subgroup, + OpTester::ExpectResult::kExpectFailure, + "MatMul algorithm subgroup"); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedPackedSplitK) { + RunTestTyped({1, 1024}, {1024, 16}, false, webgpu::MatMulAlgorithm::PackedSplitK); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedPackedSplitKHandlesBatchedPartialSplit) { + RunTestTyped({2, 1, 1028}, {2, 1028, 16}, false, + webgpu::MatMulAlgorithm::PackedSplitK); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedSubgroupMatrix) { + RunTestTyped({32, 16}, {16, 32}, false, webgpu::MatMulAlgorithm::SubgroupMatrix); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedSubgroupMatrixUnalignedConstantB) { + RunTestTyped({32, 16}, {16, 36}, true, webgpu::MatMulAlgorithm::SubgroupMatrix); +} + +TEST(WebGpuMatMulAlgorithmTest, ForcedSubgroupMatrixRejectsFloatInputs) { + RunTestTyped({32, 16}, {16, 32}, false, webgpu::MatMulAlgorithm::SubgroupMatrix, + OpTester::ExpectResult::kExpectFailure, + "MatMul algorithm subgroup_matrix"); +} + +TEST(WebGpuMatMulAlgorithmTest, ReselectsAlgorithmForEachDynamicShape) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + + SessionOptions session_options; + session_options.session_logid = "WebGpuDynamicMatMulAlgorithmSelection"; + InferenceSessionWrapper session(session_options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(webgpu_ep))); + + const std::string model_bytes = BuildDynamicMatMulModelBytes(); + ASSERT_STATUS_OK(session.Load(model_bytes.data(), static_cast(model_bytes.size()))); + ASSERT_STATUS_OK(session.Initialize()); + + const std::vector output_names{"Y"}; + auto run = [&](const TensorShape& a_shape, std::vector a_data, + const TensorShape& b_shape, std::vector b_data, + const TensorShape& expected_shape, float expected_value) { + OrtValue a_value; + OrtValue b_value; + CreateMLValue(a_shape.GetDims(), a_data.data(), OrtMemoryInfo(), &a_value); + CreateMLValue(b_shape.GetDims(), b_data.data(), OrtMemoryInfo(), &b_value); + + NameMLValMap feeds; + feeds.emplace("A", std::move(a_value)); + feeds.emplace("B", std::move(b_value)); + std::vector fetches; + ASSERT_STATUS_OK(session.Run(feeds, output_names, &fetches)); + ASSERT_EQ(fetches.size(), 1u); + const Tensor& output = fetches[0].Get(); + ASSERT_EQ(output.Shape(), expected_shape); + for (float value : output.DataAsSpan()) { + EXPECT_EQ(value, expected_value); + } + }; + + // The first invocation uses Packed. The second has K=0 and must reselect Naive on the same + // session and kernel; reusing the first execution plan would fail Packed's nonzero-K prerequisite. + run(TensorShape({8, 8}), std::vector(64, 1.0f), + TensorShape({8, 8}), std::vector(64, 1.0f), + TensorShape({8, 8}), 8.0f); + + const float unused_storage = 0.0f; + run(TensorShape({1, 0}), std::vector{unused_storage}, + TensorShape({0, 8}), std::vector{unused_storage}, + TensorShape({1, 8}), 0.0f); +} + // 2D aligned baseline shapes. TEST(MatMul_Large, DISABLED_Aligned) { RunBothTypes({128, 64}, {64, 1024}); @@ -159,7 +570,7 @@ TEST(MatMul_Large, DISABLED_Broadcast4D) { } // Batched B (true bmm): A [..., M, K] x B [..., K, N] with matching batch. On the -// Intel subgroup path each (A, B) slice is dispatched on z. Covers small and +// Subgroup path each (A, B) slice is dispatched on z. Covers small and // larger batch counts with tile-aligned per-slice shapes. TEST(MatMul_Large, DISABLED_BatchedB) { RunBothTypes({2, 128, 64}, {2, 64, 1024}); @@ -212,7 +623,7 @@ TEST(MatMul_Large, DISABLED_ConstantWeightOddN) { // Broadcasted batch dims that are NOT identical but share the same batch // *product* (A=[2,1,...], B=[1,2,...] -> [2,2,...]; A=[1,4,...], B=[4,1,...] -> // [4,4,...]). A product-only batch check would wrongly route these onto the -// Intel subgroup path, which pairs slice i of A with slice i of B and copies A's +// Subgroup path, which pairs slice i of A with slice i of B and copies A's // shape to the output - producing the wrong output shape and mismatched pairing. // N is even and every per-slice shape is tile-aligned, so only the // identical-batch-dims guard (not the odd-N or partial-tile fallbacks) keeps them diff --git a/onnxruntime/test/providers/webgpu/nn/conv_transpose_test.cc b/onnxruntime/test/providers/webgpu/nn/conv_transpose_test.cc index ba7f980528fb4..d7f9ea2817e47 100644 --- a/onnxruntime/test/providers/webgpu/nn/conv_transpose_test.cc +++ b/onnxruntime/test/providers/webgpu/nn/conv_transpose_test.cc @@ -17,7 +17,7 @@ namespace onnxruntime { namespace test { namespace { -struct ConvTranspose3DAttributes { +struct ConvTransposeTestAttributes { std::vector strides; std::vector dilations; std::vector pads; @@ -28,12 +28,13 @@ struct ConvTranspose3DAttributes { }; // Layout, float16, initializer weights/bias, opset. -class ConvTranspose3DWebGpuTest : public testing::TestWithParam> { +class ConvTransposeWebGpuTestBase : public testing::TestWithParam> { protected: - void Run(const std::vector& input_shape, - const std::vector& weight_shape, - const std::vector& output_shape, - const ConvTranspose3DAttributes& attrs = {}, bool has_bias = false) { + void Run(const TensorShapeVector& input_shape, + const TensorShapeVector& weight_shape, + const TensorShapeVector& output_shape, + const ConvTransposeTestAttributes& attrs = {}, bool has_bias = false, + bool disable_prepacking = false) { const auto [is_nhwc, is_fp16, initializer, opset] = GetParam(); auto ep = DefaultWebGpuExecutionProvider(is_nhwc); if (!ep) { @@ -95,16 +96,148 @@ class ConvTranspose3DWebGpuTest : public testing::TestWithParam("X", input_shape, {1.0f, 2.0f}); + test.AddInput("W", weight_shape, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}, initializer); + test.AddOutput("Y", {1}, {0.0f}); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + size_t prepacked_weights = 0; + ASSERT_NO_FATAL_FAILURE(test.Config(options) + .Config(OpTester::ExpectResult::kExpectFailure, + "X num_dims does not match W num_dims.") + .ConfigEp(std::move(ep)) + .RunWithConfig(&prepacked_weights)); + EXPECT_EQ(prepacked_weights, initializer ? 1U : 0U); + } + } + } +} + TEST(ConvTransposeWebGpuTest, UnsupportedSpatialRank) { auto ep = DefaultWebGpuExecutionProvider(); if (!ep) { @@ -125,33 +258,33 @@ TEST_P(ConvTranspose3DWebGpuTest, DefaultAttributes) { } TEST_P(ConvTranspose3DWebGpuTest, GroupsAndBias) { - ConvTranspose3DAttributes attrs; + ConvTransposeTestAttributes attrs; attrs.group = 2; Run({2, 4, 2, 2, 3}, {4, 3, 2, 3, 1}, {2, 6, 3, 4, 3}, attrs, true); } TEST_P(ConvTranspose3DWebGpuTest, Depthwise) { - ConvTranspose3DAttributes attrs; + ConvTransposeTestAttributes attrs; attrs.group = 3; Run({1, 3, 2, 2, 2}, {3, 1, 2, 2, 2}, {1, 3, 3, 3, 3}, attrs, true); } TEST_P(ConvTranspose3DWebGpuTest, Vec4GroupsAndBias) { - ConvTranspose3DAttributes attrs; + ConvTransposeTestAttributes attrs; attrs.group = 2; // Eight channels per group exercises multiple vec4 loads and a nonzero group offset. Run({2, 16, 2, 2, 2}, {16, 3, 2, 1, 3}, {2, 6, 3, 2, 4}, attrs, true); } TEST_P(ConvTranspose3DWebGpuTest, Vec2GroupAlignment) { - ConvTranspose3DAttributes attrs; + ConvTransposeTestAttributes attrs; attrs.group = 2; // Total channels are divisible by four, but each group's six channels require vec2. Run({2, 12, 2, 2, 2}, {12, 3, 2, 1, 3}, {2, 6, 3, 2, 4}, attrs, true); } TEST_P(ConvTranspose3DWebGpuTest, StridesDilationsAndAsymmetricPadding) { - ConvTranspose3DAttributes attrs; + ConvTransposeTestAttributes attrs; attrs.strides = {2, 3, 2}; attrs.dilations = {2, 1, 2}; attrs.pads = {1, 0, 3, 0, 1, 1}; @@ -160,28 +293,28 @@ TEST_P(ConvTranspose3DWebGpuTest, StridesDilationsAndAsymmetricPadding) { } TEST_P(ConvTranspose3DWebGpuTest, SameUpper) { - ConvTranspose3DAttributes attrs; + ConvTransposeTestAttributes attrs; attrs.strides = {2, 2, 2}; attrs.auto_pad = "SAME_UPPER"; Run({1, 2, 2, 2, 2}, {2, 3, 3, 3, 3}, {1, 3, 4, 4, 4}, attrs); } TEST_P(ConvTranspose3DWebGpuTest, SameLower) { - ConvTranspose3DAttributes attrs; + ConvTransposeTestAttributes attrs; attrs.strides = {2, 2, 2}; attrs.auto_pad = "SAME_LOWER"; Run({1, 2, 2, 2, 2}, {2, 3, 3, 3, 3}, {1, 3, 4, 4, 4}, attrs, true); } TEST_P(ConvTranspose3DWebGpuTest, ExplicitOutputShape) { - ConvTranspose3DAttributes attrs; + ConvTransposeTestAttributes attrs; attrs.strides = {2, 2, 2}; attrs.output_shape = {3, 3, 3}; Run({1, 2, 2, 2, 2}, {2, 3, 3, 2, 2}, {1, 3, 3, 3, 3}, attrs); } TEST_P(ConvTranspose3DWebGpuTest, PaddingLargerThanKernel) { - ConvTranspose3DAttributes attrs; + ConvTransposeTestAttributes attrs; attrs.pads = {2, 1, 0, 0, 0, 0}; Run({1, 1, 4, 3, 2}, {1, 1, 2, 2, 2}, {1, 1, 3, 3, 3}, attrs); } diff --git a/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc b/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc new file mode 100644 index 0000000000000..99050394fd8a5 --- /dev/null +++ b/onnxruntime/test/providers/webgpu/quantization/matmul_nbits_test.cc @@ -0,0 +1,152 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_CONTRIB_OPS) + +#include +#include + +#include "gtest/gtest.h" + +#include "core/common/inlined_containers.h" +#include "core/common/span_utils.h" +#include "core/providers/webgpu/webgpu_context.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime { +namespace test { +namespace { + +void RunMatMulNBitsTileTest(int64_t bits, int64_t M, int64_t N, int64_t K, + bool has_zero_points, bool has_bias) { + SCOPED_TRACE(testing::Message() << "bits=" << bits << ", M=" << M << ", N=" << N << ", K=" << K + << ", zero_points=" << has_zero_points << ", bias=" << has_bias); + + auto provider = DefaultWebGpuExecutionProvider(); + ASSERT_NE(provider, nullptr); + const auto& context = webgpu::WebGpuContextFactory::GetContext(provider->GetDeviceId()); + if (!context.DeviceHasFeature(wgpu::FeatureName::ShaderF16)) { + GTEST_SKIP() << "WebGPU adapter does not support shader-f16"; + } + + constexpr int64_t block_size = 32; + const int64_t blocks = K / block_size; + const int64_t values_per_byte = 8 / bits; + const int64_t blob_size = block_size / values_per_byte; + const int64_t zero_point_bytes = (blocks + values_per_byte - 1) / values_per_byte; + const int32_t default_zero_point = 1 << (bits - 1); + + RandomValueGenerator random{1234}; + const auto a = FloatsToMLFloat16s(random.Gaussian(AsSpan({M, K}), 0.0f, 0.25f)); + const auto weights = random.Uniform(AsSpan({N, K}), 0, (1 << bits) - 1); + const auto bias = FloatsToMLFloat16s(random.Uniform(AsSpan({N}), -0.5f, 0.5f)); + InlinedVector packed_weights(N * blocks * blob_size, 0); + InlinedVector packed_zero_points(N * zero_point_bytes, 0); + InlinedVector scales(N * blocks); + InlinedVector dequantized_weights(N * K); + + // Construct the quantized inputs locally; the reference uses the unpacked values + // so packing/layout errors are not shared between the kernel and the reference. + for (int64_t n = 0; n < N; ++n) { + for (int64_t block = 0; block < blocks; ++block) { + const int64_t scale_index = n * blocks + block; + scales[scale_index] = MLFloat16(static_cast(1 + scale_index % 4) / + (bits == 4 ? 128.0f : 2048.0f)); + const int32_t zero_point = has_zero_points + ? default_zero_point + static_cast(scale_index % 5) - 2 + : default_zero_point; + packed_zero_points[n * zero_point_bytes + block / values_per_byte] |= + static_cast(zero_point << ((block % values_per_byte) * bits)); + for (int64_t k = block * block_size; k < (block + 1) * block_size; ++k) { + const int64_t index = n * K + k; + packed_weights[index / values_per_byte] |= + static_cast(weights[index] << ((k % values_per_byte) * bits)); + dequantized_weights[index] = static_cast(weights[index] - zero_point) * scales[scale_index].ToFloat(); + } + } + } + + InlinedVector expected(M * N); + for (int64_t m = 0; m < M; ++m) { + for (int64_t n = 0; n < N; ++n) { + float sum = 0.0f; + for (int64_t k = 0; k < K; ++k) { + sum += a[m * K + k].ToFloat() * dequantized_weights[n * K + k]; + } + expected[m * N + n] = MLFloat16(sum + (has_bias ? bias[n].ToFloat() : 0.0f)); + } + } + + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("bits", bits); + test.AddAttribute("block_size", block_size); + test.AddAttribute("accuracy_level", 4); + test.AddInput("A", {M, K}, a); + test.AddInput("B", {N, blocks, blob_size}, packed_weights.data(), packed_weights.size(), true); + test.AddInput("scales", {N, blocks}, scales.data(), scales.size(), true); + if (has_zero_points) { + test.AddInput("zero_points", {N, zero_point_bytes}, packed_zero_points.data(), packed_zero_points.size(), true); + } else { + test.AddOptionalInputEdge(); + } + test.AddOptionalInputEdge(); // g_idx + if (has_bias) { + test.AddInput("bias", {N}, bias, true); + } + test.AddOutput("Y", {M, N}, expected.data(), expected.size()); + test.SetOutputAbsErr("Y", 0.1f); + test.SetOutputRelErr("Y", 0.02f); + + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(options).ConfigEp(std::move(provider)).RunWithConfig(); +} + +void RunMatMulNBitsTileTests(int64_t bits) { + for (bool has_zero_points : {false, true}) { + for (bool has_bias : {false, true}) { + // AMD D3D12 wave64: small-M/BK=64, four-wave/BK=128 with an M tail, + // and two-wave/BK=64. Other adapters validate their selected WebGPU path. + // N=128 checks valid output from both column subgroups of the wide tile. + RunMatMulNBitsTileTest(bits, 32, 128, 64, has_zero_points, has_bias); + RunMatMulNBitsTileTest(bits, 400, 64, 128, has_zero_points, has_bias); + RunMatMulNBitsTileTest(bits, 640, 64, 128, has_zero_points, has_bias); + + // N=64 needs only one column wave; N=192 retains the wide tile's partial N edge. + RunMatMulNBitsTileTest(bits, 32, 64, 64, has_zero_points, has_bias); + RunMatMulNBitsTileTest(bits, 32, 192, 64, has_zero_points, has_bias); + + // Initial BK=128 does not divide K: reduce it to 64 for K=192 and + // to 32 for K=96 so no iteration reads beyond the prepacked K dimension. + RunMatMulNBitsTileTest(bits, 400, 64, 192, has_zero_points, has_bias); + RunMatMulNBitsTileTest(bits, 400, 64, 96, has_zero_points, has_bias); + + // Two-wave BK=128 and wide BK=32, both with partial 16-row prepack blocks. + RunMatMulNBitsTileTest(bits, 257, 64, 128, has_zero_points, has_bias); + RunMatMulNBitsTileTest(bits, 33, 128, 32, has_zero_points, has_bias); + + // The single-wave tile must also handle reduced BK with a partial prepack block. + RunMatMulNBitsTileTest(bits, 33, 64, 32, has_zero_points, has_bias); + } + } +} + +} // namespace + +TEST(WebGpuMatMulNBitsTest, Float16_4Bits_SubgroupMatrixTiles) { + RunMatMulNBitsTileTests(4); +} + +TEST(WebGpuMatMulNBitsTest, Float16_8Bits_SubgroupMatrixTiles) { + RunMatMulNBitsTileTests(8); +} + +} // namespace test +} // namespace onnxruntime + +#endif // !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_CONTRIB_OPS) diff --git a/onnxruntime/test/providers/webgpu/shader_cache_regression_test.cc b/onnxruntime/test/providers/webgpu/shader_cache_regression_test.cc new file mode 100644 index 0000000000000..2c8625b1b0d60 --- /dev/null +++ b/onnxruntime/test/providers/webgpu/shader_cache_regression_test.cc @@ -0,0 +1,198 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "core/graph/graph.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +namespace { + +// Keep both variants in one session so they share the ready/pending shader cache. +class CachePairTester final : public OpTester { + public: + CachePairTester(const char* op, const char* domain, int version, size_t input_count, bool reverse) + : OpTester(op, version, domain), op_{op}, domain_{domain}, input_count_{input_count}, reverse_{reverse} {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CachePairTester); + + void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, + std::vector>&) override { + for (size_t step = 0; step < 2; ++step) { + const size_t index = reverse_ ? 1 - step : step; + std::vector node_inputs(inputs.begin() + index * input_count_, + inputs.begin() + (index + 1) * input_count_); + auto& node = graph.AddNode(MakeString("variant", index), op_, "", node_inputs, {outputs[index]}, nullptr, domain_); + if (std::string_view{op_} == "InstanceNormalization") { + node.AddAttribute("epsilon", index == 0 ? 1e-5f : 1.0f); + } + } + } + + private: + const char* op_; + const char* domain_; + const size_t input_count_; + const bool reverse_; +}; + +std::unique_ptr CreateCacheTestProvider() { + ConfigOptions provider_options; + ORT_THROW_IF_ERROR(provider_options.AddConfigEntry("ep.webgpuexecutionprovider.validationMode", "full")); + ORT_THROW_IF_ERROR(provider_options.AddConfigEntry("ep.webgpuexecutionprovider.preferredLayout", "NCHW")); + return WebGpuExecutionProviderWithOptions(provider_options); +} + +void RunOnWebGpu(OpTester& test, std::unique_ptr provider) { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.SetNumRunCalls(2); + test.Config(options).ConfigEp(std::move(provider)).RunWithConfig(); +} + +#ifndef DISABLE_CONTRIB_OPS +template +void AddBiasVariant(CachePairTester& test, const char* op, int index, int64_t channels) { + const bool add = std::string_view{op} == "BiasAdd"; + const int64_t output_channels = add ? channels : channels / 2; + test.AddInput(MakeString("x", index).c_str(), {1, 1, channels}, std::vector(channels, T{1.0f})); + test.AddInput(MakeString("b", index).c_str(), {channels}, std::vector(channels, T{2.0f})); + if (add) { + test.AddInput(MakeString("r", index).c_str(), {1, 1, channels}, std::vector(channels, T{3.0f})); + } + const float expected = add ? 6.0f : 4.5f * (1.0f + std::erf(3.0f / std::sqrt(2.0f))); + test.AddOutput(MakeString("y", index).c_str(), {1, 1, output_channels}, + std::vector(output_channels, T{expected})); +} + +TEST(WebGpuShaderCacheTest, BiasAddVectorWidths) { + for (bool reverse : {false, true}) { + // Scalar and vec2 must each coexist with vec4 at the same rank and dtype. + for (int64_t channels : {3, 6}) { + SCOPED_TRACE(MakeString("reverse=", reverse, " channels=", channels)); + auto provider = CreateCacheTestProvider(); + if (!provider) { + GTEST_SKIP() << "WebGPU EP is not available"; + } + CachePairTester test{"BiasAdd", kMSDomain, 1, 3, reverse}; + AddBiasVariant(test, "BiasAdd", 0, channels); + AddBiasVariant(test, "BiasAdd", 1, 8); + RunOnWebGpu(test, std::move(provider)); + } + } +} + +void TestBiasElementTypes(const char* op) { + for (bool reverse : {false, true}) { + SCOPED_TRACE(MakeString(op, " reverse=", reverse)); + auto provider = CreateCacheTestProvider(); + if (!provider) { + GTEST_SKIP() << "WebGPU EP is not available"; + } + CachePairTester test{op, kMSDomain, 1, std::string_view{op} == "BiasAdd" ? 3u : 2u, reverse}; + AddBiasVariant(test, op, 0, 8); + AddBiasVariant(test, op, 1, 8); + RunOnWebGpu(test, std::move(provider)); + } +} + +TEST(WebGpuShaderCacheTest, BiasAddElementTypes) { + TestBiasElementTypes("BiasAdd"); +} + +TEST(WebGpuShaderCacheTest, BiasSplitGeluElementTypes) { + TestBiasElementTypes("BiasSplitGelu"); +} +#endif + +class LayerNormOutputsTester final : public OpTester { + public: + explicit LayerNormOutputsTester(bool reverse) : OpTester("LayerNormalization", 17), reverse_{reverse} {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(LayerNormOutputsTester); + + void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, + std::vector>&) override { + auto& absent = graph.GetOrCreateNodeArg("", nullptr); + for (int i : {0, 1}) { + const bool mean = (i == 0) != reverse_; + graph.AddNode(mean ? "mean" : "inverse_std", "LayerNormalization", "", inputs, + mean ? std::vector{outputs[0], outputs[1]} + : std::vector{outputs[2], &absent, outputs[3]}); + } + } + + private: + const bool reverse_; +}; + +TEST(WebGpuShaderCacheTest, LayerNormOptionalOutputRoles) { + for (bool reverse : {false, true}) { + SCOPED_TRACE(reverse); + auto provider = CreateCacheTestProvider(); + if (!provider) { + GTEST_SKIP() << "WebGPU EP is not available"; + } + LayerNormOutputsTester test{reverse}; + test.AddInput("x", {1, 4}, {1, 2, 3, 4}); + test.AddInput("scale", {4}, {1, 1, 1, 1}); + const float inv_std = 1.0f / std::sqrt(1.25f + 1e-5f); + const std::vector normalized{-1.5f * inv_std, -0.5f * inv_std, 0.5f * inv_std, 1.5f * inv_std}; + test.AddOutput("mean_y", {1, 4}, normalized); + test.AddOutput("mean", {1, 1}, {2.5f}); + test.AddOutput("inverse_y", {1, 4}, normalized); + test.AddOutput("inverse_std", {1, 1}, {inv_std}); + RunOnWebGpu(test, std::move(provider)); + } +} + +TEST(WebGpuShaderCacheTest, InstanceNormEpsilon) { + for (bool reverse : {false, true}) { + SCOPED_TRACE(reverse); + auto provider = CreateCacheTestProvider(); + if (!provider) { + GTEST_SKIP() << "WebGPU EP is not available"; + } + CachePairTester test{"InstanceNormalization", kOnnxDomain, 17, 3, reverse}; + for (int index : {0, 1}) { + test.AddInput(MakeString("x", index).c_str(), {1, 1, 4}, {1, 2, 3, 4}); + test.AddInput(MakeString("s", index).c_str(), {1}, {1}); + test.AddInput(MakeString("b", index).c_str(), {1}, {0}); + } + const float inv = 1.0f / std::sqrt(1.25f + 1e-5f); + test.AddOutput("y0", {1, 1, 4}, {-1.5f * inv, -0.5f * inv, 0.5f * inv, 1.5f * inv}); + test.AddOutput("y1", {1, 1, 4}, {-1, -1.0f / 3, 1.0f / 3, 1}); + RunOnWebGpu(test, std::move(provider)); + } +} + +TEST(WebGpuShaderCacheTest, BooleanBroadcastReadModes) { + for (bool reverse : {false, true}) { + SCOPED_TRACE(reverse); + auto provider = CreateCacheTestProvider(); + if (!provider) { + GTEST_SKIP() << "WebGPU EP is not available"; + } + CachePairTester test{"And", kOnnxDomain, 17, 2, reverse}; + test.AddInput("a0", {2, 1}, {true, false}); + test.AddInput("b0", {1, 4}, {true, false, true, false}); + test.AddInput("a1", {1, 4}, {true, false, true, false}); + test.AddInput("b1", {2, 1}, {true, false}); + const std::array expected{true, false, true, false, false, false, false, false}; + test.AddOutput("y0", {2, 4}, expected.data(), expected.size()); + test.AddOutput("y1", {2, 4}, expected.data(), expected.size()); + RunOnWebGpu(test, std::move(provider)); + } +} + +} // namespace +} // namespace onnxruntime::test diff --git a/onnxruntime/test/providers/webgpu/tensor/quantize_linear_test.cc b/onnxruntime/test/providers/webgpu/tensor/quantize_linear_test.cc index 5bb1b5416cac7..abe1a222fa587 100644 --- a/onnxruntime/test/providers/webgpu/tensor/quantize_linear_test.cc +++ b/onnxruntime/test/providers/webgpu/tensor/quantize_linear_test.cc @@ -11,8 +11,10 @@ #include "gtest/gtest.h" #include "core/graph/graph.h" +#include "core/providers/webgpu/webgpu_provider_options.h" #include "core/session/onnxruntime_session_options_config_keys.h" #include "test/providers/provider_test_utils.h" +#include "test/providers/webgpu/graph_capture_test_utils.h" #include "test/util/include/default_providers.h" namespace onnxruntime::test { @@ -94,5 +96,92 @@ TEST(WebGpuDequantizeLinearTest, Int8Float16ScalarAndVectorCache) { TestPackedDequantizeLinearCache(); } +template +void TestEmptyDequantization(const std::vector& input_shape, + const std::vector& scale_shape, + const std::vector* zero_point_shape, + int64_t block_size, const char* error = "") { + ConfigOptions config; + ASSERT_STATUS_OK(config.AddConfigEntry(webgpu::options::kValidationMode, + webgpu::options::kValidationMode_full)); + ASSERT_STATUS_OK(config.AddConfigEntry(webgpu::options::kEnableGraphCapture, + webgpu::options::kEnableGraphCapture_ON)); + auto provider = WebGpuExecutionProviderWithOptions(config); + if (!provider) { + GTEST_SKIP() << "WebGPU EP is not available"; + } + + WebGpuGraphCaptureTester test("DequantizeLinear", 21); + // Exercise kernel validation rather than model-load shape inference. + test.AddShapeToTensorData(false); + test.AddAttribute("axis", int64_t{0}); + test.AddAttribute("block_size", block_size); + test.AddInput("x", input_shape, {}); + test.AddInput("x_scale", scale_shape, + std::vector(TensorShape(scale_shape).Size(), T{0.5f})); + if (zero_point_shape) { + test.AddInput("x_zero_point", *zero_point_shape, + std::vector(TensorShape(*zero_point_shape).Size(), 0)); + } + test.AddOutput("y", input_shape, {}); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + options.graph_optimization_level = TransformerLevel::Default; + const bool expect_failure = error[0] != '\0'; + if (expect_failure) { + test.Config(options) + .Config(OpTester::ExpectResult::kExpectFailure, error) + .ConfigEp(std::move(provider)) + .RunWithConfig(); + } else { + test.RunWithBoundCapture(std::move(provider), options, TensorShape(input_shape), {}); + } +} + +TEST(WebGpuDequantizeLinearTest, EmptyBlockedShapesAreValidatedBeforeNoDispatch) { + for (const auto& shape : {std::vector{35, 0}, std::vector{0, 16}}) { + const std::vector scale_shape{(shape[0] + 31) / 32, shape[1]}; + TestEmptyDequantization(shape, scale_shape, nullptr, 32); + TestEmptyDequantization(shape, scale_shape, &scale_shape, 32); + } + TestEmptyDequantization({35, 0}, {3, 0}, nullptr, 32, + "x_scale must be ceil(Di/block_size)"); + TestEmptyDequantization({0, 16}, {0, 15}, nullptr, 32, + "same shape on non-quantize axes"); + TestEmptyDequantization({0, 16}, {0, 16, 1}, nullptr, 32, + "same rank for blocked quantization"); + TestEmptyDequantization({35, 0}, {0}, nullptr, 32, + "same rank for blocked quantization"); + const std::vector wrong_zero_shape{2, 1}; + TestEmptyDequantization({35, 0}, {2, 0}, &wrong_zero_shape, 32, + "x_zero_point and x_scale must have the same shape"); + const std::vector wrong_zero_rank{0}; + TestEmptyDequantization({35, 0}, {2, 0}, &wrong_zero_rank, 32, + "x_zero_point and x_scale must have the same rank"); +} + +TEST(WebGpuDequantizeLinearTest, EmptyShapeBlockInferenceIsSafe) { + TestEmptyDequantization({32, 0}, {2, 0}, nullptr, 0); + TestEmptyDequantization({0, 16}, {0, 16}, nullptr, 0); + TestEmptyDequantization({35, 0}, {0, 0}, nullptr, 0, + "x_scale must be nonempty on a nonempty quantize axis"); + TestEmptyDequantization({35, 0}, {0}, nullptr, 0, + "x_scale must be nonempty on a nonempty quantize axis"); +} + +TEST(WebGpuDequantizeLinearTest, EmptyPerTensorAndPerAxisShapesAreValidated) { + const std::vector scalar{}; + const std::vector empty_axis{0}; + TestEmptyDequantization({0, 16}, {}, &scalar, 0); + TestEmptyDequantization({0, 16}, {0}, &empty_axis, 0); + TestEmptyDequantization({0, 16}, {}, &empty_axis, 0, + "x_zero_point must be a scalar or size-1 vector"); + TestEmptyDequantization({0, 16}, {2}, nullptr, 0, + "x_scale must match the quantize axis"); + const std::vector wrong_zero_shape{2}; + TestEmptyDequantization({0, 16}, {0}, &wrong_zero_shape, 0, + "x_zero_point must match x_scale"); +} + } // namespace } // namespace onnxruntime::test diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index ad92ee33b7430..d0ba4ead1b8ca 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -27,6 +27,7 @@ #include "core/graph/onnx_protobuf.h" #include "core/providers/webgpu/allocator.h" #include "core/providers/webgpu/buffer_manager.h" +#include "core/providers/webgpu/program_cache_key.h" #include "core/providers/webgpu/webgpu_context.h" #include "core/providers/webgpu/webgpu_execution_provider.h" #include "core/providers/webgpu/webgpu_provider_factory_creator.h" @@ -46,6 +47,7 @@ #if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) #include "dawn/native/DawnNative.h" +#include "test/providers/webgpu/webgpu_context_test_helper.h" #endif namespace onnxruntime { @@ -230,7 +232,7 @@ void TestCopyAfterDeferredDispatch(bool upload, bool checked_completion = false) bind_group_desc.layout = pipeline.GetBindGroupLayout(0); bind_group_desc.entryCount = entries.size(); bind_group_desc.entries = entries.data(); - webgpu::CapturedCommandInfo dispatch; + webgpu::DispatchCommandInfo dispatch; dispatch.compute_pipeline = pipeline; dispatch.bind_group = context.Device().CreateBindGroup(&bind_group_desc); recording.deferred_dispatches.push_back(std::move(dispatch)); @@ -343,7 +345,7 @@ TEST(WebGpuContextTest, CheckedCompletionDrainsAfterDeferredDispatchFailure) { auto& webgpu_ep = static_cast(*ep); auto& context = webgpu::WebGpuContextFactory::GetContext(0); auto& recording = webgpu_ep.Recording(); - webgpu::CapturedCommandInfo invalid_dispatch; + webgpu::DispatchCommandInfo invalid_dispatch; invalid_dispatch.program_key = "checked-completion-missing-pipeline"; recording.deferred_dispatches.push_back(std::move(invalid_dispatch)); recording.has_unsubmitted_work.store(true, std::memory_order_relaxed); @@ -371,63 +373,27 @@ Status StartPipelineCallbackTestBuild(webgpu::WebGpuContext& context, webgpu::Pe bool invalid_shader) { PipelineCallbackTestProgram program; program.SetWorkgroupSize(1); + program.CacheHint(invalid_shader); auto metadata = program.Metadata(); const webgpu::ProgramConstant invalid_constant{"invalid identifier", uint32_t{0}}; if (invalid_shader) { metadata.constants = {&invalid_constant, 1}; } webgpu::ProgramManager manager{context}; + const auto program_key = webgpu::CalculateProgramCacheKey(program, {}, {}); build.name = program.Name(); build.callback_context = std::make_shared(); - return manager.Build(program, metadata, {}, {}, program.Name(), 1, 1, 1, + return manager.Build(program, metadata, {}, {}, program_key, 1, 1, 1, build.bind_group_layout, build.shape_uniform_ranks, build.future, build.callback_context); } -template -void RunWithExternalCompletionContext(TestBody test_body, bool timed_wait = true) { - ConfigOptions options; - auto owned_ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); - ASSERT_NE(owned_ep, nullptr); - - const auto instance_feature = wgpu::InstanceFeatureName::TimedWaitAny; - wgpu::InstanceDescriptor instance_desc{}; - instance_desc.requiredFeatureCount = timed_wait ? 1 : 0; - instance_desc.requiredFeatures = timed_wait ? &instance_feature : nullptr; - dawn::native::Instance instance{&instance_desc}; - wgpu::RequestAdapterOptions adapter_options{}; - adapter_options.backendType = static_cast(webgpu::WebGpuContextConfig{}.backend_type); - auto adapters = instance.EnumerateAdapters(&adapter_options); - ASSERT_FALSE(adapters.empty()); - - const auto device_feature = wgpu::FeatureName::ImplicitDeviceSynchronization; - auto lost = std::make_shared(false); - wgpu::DeviceDescriptor device_desc{}; - device_desc.requiredFeatureCount = 1; - device_desc.requiredFeatures = &device_feature; - device_desc.SetDeviceLostCallback( - wgpu::CallbackMode::WaitAnyOnly, - [lost](const wgpu::Device&, wgpu::DeviceLostReason, wgpu::StringView) noexcept { *lost = true; }); - auto device = wgpu::Device::Acquire(adapters.front().CreateDevice(&device_desc)); - ASSERT_NE(device, nullptr); - - ConfigOptions external_options; - ASSERT_STATUS_OK(external_options.AddConfigEntry(kDeviceId, "1")); - ASSERT_STATUS_OK(external_options.AddConfigEntry( - kWebGpuInstance, std::to_string(reinterpret_cast(instance.Get())).c_str())); - ASSERT_STATUS_OK(external_options.AddConfigEntry( - kWebGpuDevice, std::to_string(reinterpret_cast(device.Get())).c_str())); - auto external_ep = WebGpuProviderFactoryCreator::Create(external_options)->CreateProvider(); - auto& webgpu_ep = static_cast(*external_ep); - auto& context = webgpu::WebGpuContextFactory::GetContext(1); - test_body(context, webgpu_ep, lost); -} #endif TEST(WebGpuContextTest, SuccessfulPipelineBuildReleasesCallbackOwnership) { #if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) GTEST_SKIP() << "Dawn native device creation is unavailable."; #else - RunWithExternalCompletionContext([](auto& context, auto&, const auto&) { + RunWithExternalCompletionContext([](auto& context, auto&, const auto&, auto&) { webgpu::PendingPipelineBuild build; ASSERT_STATUS_OK(StartPipelineCallbackTestBuild(context, build, false)); const std::weak_ptr weak_context = build.callback_context; @@ -445,10 +411,10 @@ TEST(WebGpuContextTest, FailedPipelineWaitRetainsCallbackUntilDelivery) { #if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) GTEST_SKIP() << "Dawn native device creation is unavailable."; #else - RunWithExternalCompletionContext([](auto& context, auto& ep, const auto&) { + RunWithExternalCompletionContext([](auto& context, auto& ep, const auto&, auto&) { // Invalid WGSL makes the pipeline event ready synchronously, but WaitAnyOnly defers delivery. context.PushErrorScope(); - webgpu::CapturedCommandInfo dispatch; + webgpu::DispatchCommandInfo dispatch; dispatch.program_key = "failed-wait-pipeline-callback"; auto& build = dispatch.pending_build.emplace(); ASSERT_STATUS_OK(StartPipelineCallbackTestBuild(context, build, true)); @@ -493,9 +459,9 @@ TEST(WebGpuContextTest, FailedPipelineWaitRetainsCallbackAfterContextDestruction wgpu::Instance callback_instance; wgpu::Future future{}; std::weak_ptr weak_context; - RunWithExternalCompletionContext([&](auto& context, auto& ep, const auto&) { + RunWithExternalCompletionContext([&](auto& context, auto& ep, const auto&, auto&) { context.PushErrorScope(); - webgpu::CapturedCommandInfo dispatch; + webgpu::DispatchCommandInfo dispatch; dispatch.program_key = "post-context-pipeline-callback"; auto& build = dispatch.pending_build.emplace(); ASSERT_STATUS_OK(StartPipelineCallbackTestBuild(context, build, true)); @@ -528,7 +494,7 @@ TEST(WebGpuContextTest, CheckedCompletionPropagatesValidationAndPreservesOuterSc #if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) GTEST_SKIP() << "Dawn native device creation is unavailable."; #else - RunWithExternalCompletionContext([](auto& context, auto& ep, const auto&) { + RunWithExternalCompletionContext([](auto& context, auto& ep, const auto&, auto&) { auto& recording = ep.Recording(); wgpu::BufferDescriptor desc{}; desc.size = 64; @@ -555,7 +521,7 @@ TEST(WebGpuContextTest, ErrorScopesPropagateValidationAndOutOfMemory) { #if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) GTEST_SKIP() << "Dawn native error injection is unavailable."; #else - RunWithExternalCompletionContext([](auto& context, auto& ep, const auto&) { + RunWithExternalCompletionContext([](auto& context, auto& ep, const auto&, auto&) { for (const auto filter : {wgpu::ErrorFilter::Validation, wgpu::ErrorFilter::OutOfMemory}) { SCOPED_TRACE(static_cast(filter)); const auto type = filter == wgpu::ErrorFilter::Validation ? wgpu::ErrorType::Validation @@ -575,7 +541,7 @@ TEST(WebGpuContextTest, CheckedCompletionCleansScopesAfterWaitFailure) { #if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) GTEST_SKIP() << "Dawn native device creation is unavailable."; #else - RunWithExternalCompletionContext([](auto& context, auto& ep, const auto&) { + RunWithExternalCompletionContext([](auto& context, auto& ep, const auto&, auto&) { const auto status = context.FlushAndWaitChecked(ep.BufferManager(), ep.Recording()); EXPECT_FALSE(status.IsOK()); EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("Failed to wait for the operation")); @@ -594,6 +560,77 @@ TEST(WebGpuContextTest, CheckedCompletionCleansScopesAfterWaitFailure) { #endif } +TEST(WebGpuContextTest, CheckedRegionReportsEarlyShaderErrorOnExternalDevice) { +#if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) + GTEST_SKIP() << "Dawn native device creation is unavailable."; +#else + RunWithExternalCompletionContext([](auto& context, auto& ep, const auto&, auto&) { + wgpu::ShaderSourceWGSL source{}; + source.code = "checked_region_invalid_shader"; + wgpu::ShaderModuleDescriptor desc{}; + desc.nextInChain = &source; + const auto status = context.RunAndWaitChecked(ep.BufferManager(), ep.Recording(), [&]() { + context.Device().CreateShaderModule(&desc); + return Status::OK(); + }); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("CreateShaderModule")); + // No invalid command is recorded: a successful queue wait alone cannot detect this error. + EXPECT_EQ(ep.Recording().command_encoder, nullptr); + ExpectQueueCompleted(context); + ASSERT_STATUS_OK(context.FlushAndWaitChecked(ep.BufferManager(), ep.Recording())); + }); +#endif +} + +void TestCheckedRegionEnqueueFailure(bool throw_exception) { +#if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) + ORT_UNUSED_PARAMETER(throw_exception); + GTEST_SKIP() << "Dawn native error injection is unavailable."; +#else + RunWithExternalCompletionContext([throw_exception](auto& context, auto& ep, const auto&, auto&) { + context.PushErrorScope(); + context.Device().InjectError(wgpu::ErrorType::Validation, "checked-region-outer-error"); + size_t enqueue_count = 0; + const auto status = context.RunAndWaitChecked(ep.BufferManager(), ep.Recording(), [&]() -> Status { + ++enqueue_count; + context.GetCommandEncoder(ep.Recording()); + context.Device().InjectError(wgpu::ErrorType::Validation, "checked-region-validation-error"); + context.Device().InjectError(wgpu::ErrorType::OutOfMemory, "checked-region-oom-error"); + if (throw_exception) { + ORT_THROW("checked-region-enqueue-failure"); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "checked-region-enqueue-failure"); + }); + const auto outer_status = context.PopErrorScope(); + EXPECT_EQ(enqueue_count, 1u); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("checked-region-enqueue-failure")); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("checked-region-validation-error")); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("checked-region-oom-error")); + EXPECT_FALSE(outer_status.IsOK()); + EXPECT_THAT(outer_status.ErrorMessage(), testing::HasSubstr("checked-region-outer-error")); + EXPECT_TRUE(ep.Recording().deferred_dispatches.empty()); + EXPECT_EQ(ep.Recording().command_encoder, nullptr); + EXPECT_FALSE(ep.Recording().has_unsubmitted_work.load(std::memory_order_relaxed)); + ExpectQueueCompleted(context); + ASSERT_STATUS_OK(context.FlushAndWaitChecked(ep.BufferManager(), ep.Recording())); + }); +#endif +} + +TEST(WebGpuContextTest, CheckedRegionCleansScopesAfterEnqueueFailure) { + TestCheckedRegionEnqueueFailure(false); +} + +TEST(WebGpuContextTest, CheckedRegionCleansScopesAfterEnqueueException) { +#ifdef ORT_NO_EXCEPTIONS + GTEST_SKIP() << "Exceptions are disabled."; +#else + TestCheckedRegionEnqueueFailure(true); +#endif +} + void TestCheckedCompletionUncapturedError(wgpu::ErrorType type) { #if !GTEST_HAS_DEATH_TEST || defined(__wasm__) || defined(USE_EXTERNAL_DAWN) ORT_UNUSED_PARAMETER(type); @@ -628,7 +665,7 @@ TEST(WebGpuContextTest, CheckedCompletionDetectsLostExternalDevice) { GTEST_SKIP() << "Dawn native device loss injection is unavailable."; #else for (bool pending_encoder : {false, true}) { - RunWithExternalCompletionContext([pending_encoder](auto& context, auto& ep, const auto& lost) { + RunWithExternalCompletionContext([pending_encoder](auto& context, auto& ep, const auto& lost, auto&) { auto& recording = ep.Recording(); ASSERT_STATUS_OK(context.FlushAndWaitChecked(ep.BufferManager(), recording)); if (pending_encoder) { @@ -758,7 +795,7 @@ TEST(WebGpuContextTest, IndependentClearPreservesPendingAndCapturedUniformBuffer constexpr auto storage_usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; - std::vector captured_commands; + std::vector captured_commands; context.CaptureBegin(&captured_commands, manager, session_recording); auto uniform_buffer = manager.Create(session_recording, 64, uniform_usage); wgpu::Buffer retained_uniform{uniform_buffer}; @@ -886,7 +923,7 @@ TEST(WebGpuContextTest, FailedFlushDiscardsPendingBuffersFromBothManagers) { auto shared_released = shared_manager.Create(recording, 64, usage); wgpu::Buffer retained{released}; wgpu::Buffer shared_retained{shared_released}; - webgpu::CapturedCommandInfo dispatch; + webgpu::DispatchCommandInfo dispatch; dispatch.program_key = "missing_pipeline_for_failed_flush_test"; recording.deferred_dispatches.push_back(std::move(dispatch)); recording.has_unsubmitted_work.store(true, std::memory_order_relaxed); @@ -970,6 +1007,42 @@ TEST(WebGpuContextTest, SessionAllocatorSubmitsReusedBufferClearOutsideRun) { allocator.Free(allocation); } +TEST(WebGpuContextTest, ReplaysCopyOnlyGraphAndRetainsBuffersUntilRelease) { + ConfigOptions options; + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + webgpu::BufferManager manager(context, webgpu::BufferCacheMode::Disabled, + webgpu::BufferCacheMode::Disabled, webgpu::BufferCacheMode::Disabled, + webgpu::BufferCacheMode::Disabled); + webgpu::CommandRecordingState recording; + constexpr auto usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + std::array values{1, 2, 3}; + auto source = wgpu::Buffer::Acquire(manager.Create(recording, sizeof(values), usage)); + auto intermediate = wgpu::Buffer::Acquire(manager.Create(recording, sizeof(values), usage)); + auto destination = wgpu::Buffer::Acquire(manager.Create(recording, sizeof(values), usage)); + manager.Upload(recording, values.data(), source.Get(), sizeof(values)); + + std::vector commands; + context.CaptureBegin(&commands, manager, recording); + manager.MemCpy(recording, source.Get(), intermediate.Get(), sizeof(values)); + manager.MemCpy(recording, intermediate.Get(), destination.Get(), sizeof(values)); + ASSERT_STATUS_OK(context.Flush(manager, recording)); + context.CaptureEnd(recording); + ASSERT_EQ(commands.size(), 2u); + EXPECT_EQ(recording.num_pending_dispatches, 0u); + EXPECT_EQ(ReadBufferWithExternalCommandEncoder<3>(context, destination.Get()), values); + + // Only the captured copies now own the intermediate; replay must keep it alive. + intermediate = nullptr; + values = {4, 5, 6}; + manager.Upload(recording, values.data(), source.Get(), sizeof(values)); + context.Replay(commands, manager, recording); + EXPECT_EQ(ReadBufferWithExternalCommandEncoder<3>(context, destination.Get()), values); + EXPECT_EQ(recording.num_pending_dispatches, 0u); + context.ReleaseGraphResources(commands); + EXPECT_TRUE(commands.empty()); +} + TEST(WebGpuContextTest, DoesNotCaptureDeviceAllocatorBufferClear) { ConfigOptions options; auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); @@ -982,7 +1055,7 @@ TEST(WebGpuContextTest, DoesNotCaptureDeviceAllocatorBufferClear) { webgpu::BufferCacheMode::Disabled, webgpu::BufferCacheMode::Disabled, webgpu::BufferCacheMode::Disabled); - std::vector captured_commands; + std::vector captured_commands; webgpu::GpuBufferAllocator allocator( webgpu_ep->GetDeviceId(), [&buffer_manager]() -> const webgpu::BufferManager& { return buffer_manager; }, diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test_helper.h b/onnxruntime/test/providers/webgpu/webgpu_context_test_helper.h new file mode 100644 index 0000000000000..1b167d9b1cde2 --- /dev/null +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test_helper.h @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "dawn/native/DawnNative.h" +#include "core/providers/webgpu/webgpu_context.h" +#include "core/providers/webgpu/webgpu_execution_provider.h" +#include "core/providers/webgpu/webgpu_provider_factory_creator.h" +#include "core/providers/webgpu/webgpu_provider_options.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime::test { + +template +void RunWithExternalCompletionContext(TestBody test_body, bool timed_wait = true) { + ConfigOptions options; + auto owned_ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + ASSERT_NE(owned_ep, nullptr); + + const auto instance_feature = wgpu::InstanceFeatureName::TimedWaitAny; + wgpu::InstanceDescriptor instance_desc{}; + instance_desc.requiredFeatureCount = timed_wait ? 1 : 0; + instance_desc.requiredFeatures = timed_wait ? &instance_feature : nullptr; + std::atomic uncaptured_errors{0}; + dawn::native::Instance instance{&instance_desc}; + wgpu::RequestAdapterOptions adapter_options{}; + adapter_options.backendType = static_cast(webgpu::WebGpuContextConfig{}.backend_type); + auto adapters = instance.EnumerateAdapters(&adapter_options); + ASSERT_FALSE(adapters.empty()); + + const auto device_feature = wgpu::FeatureName::ImplicitDeviceSynchronization; + auto lost = std::make_shared(false); + wgpu::DeviceDescriptor device_desc{}; + device_desc.requiredFeatureCount = 1; + device_desc.requiredFeatures = &device_feature; + device_desc.SetDeviceLostCallback( + wgpu::CallbackMode::WaitAnyOnly, + [lost](const wgpu::Device&, wgpu::DeviceLostReason, wgpu::StringView) noexcept { *lost = true; }); + device_desc.SetUncapturedErrorCallback( + [](const wgpu::Device&, wgpu::ErrorType, wgpu::StringView, std::atomic* count) noexcept { + count->fetch_add(1, std::memory_order_relaxed); + }, + &uncaptured_errors); + auto device = wgpu::Device::Acquire(adapters.front().CreateDevice(&device_desc)); + ASSERT_NE(device, nullptr); + + ConfigOptions external_options; + ASSERT_STATUS_OK(external_options.AddConfigEntry(webgpu::options::kDeviceId, "1")); + // Keep Dawn validation enabled, without session/kernel scopes masking the checked boundary. + ASSERT_STATUS_OK(external_options.AddConfigEntry(webgpu::options::kValidationMode, + webgpu::options::kValidationMode_wgpuOnly)); + ASSERT_STATUS_OK(external_options.AddConfigEntry( + webgpu::options::kWebGpuInstance, std::to_string(reinterpret_cast(instance.Get())).c_str())); + ASSERT_STATUS_OK(external_options.AddConfigEntry( + webgpu::options::kWebGpuDevice, std::to_string(reinterpret_cast(device.Get())).c_str())); + auto external_ep = WebGpuProviderFactoryCreator::Create(external_options)->CreateProvider(); + ASSERT_NE(external_ep, nullptr); + auto& webgpu_ep = static_cast(*external_ep); + auto& context = webgpu::WebGpuContextFactory::GetContext(1); + test_body(context, webgpu_ep, lost, external_ep); + // ForceLoss itself reports an uncaptured internal error before the checked call. + if (!*lost) { + EXPECT_EQ(uncaptured_errors.load(std::memory_order_relaxed), 0u); + } +} + +} // namespace onnxruntime::test diff --git a/onnxruntime/test/python/onnx_backend_test_series.py b/onnxruntime/test/python/onnx_backend_test_series.py index 62159721a2108..f12af65a6946f 100644 --- a/onnxruntime/test/python/onnx_backend_test_series.py +++ b/onnxruntime/test/python/onnx_backend_test_series.py @@ -17,6 +17,12 @@ import onnx.defs import onnxruntime.backend as backend # pylint: disable=consider-using-from-import +from onnxruntime.backend.backend import OnnxRuntimeBackend + +if os.getenv("ORT_BACKEND_TEST_ALLOW_UNRELEASED_OPSETS") == "1" and OnnxRuntimeBackend.allowReleasedOpsetsOnly: + raise RuntimeError( + "onnx_backend_test_series requires ALLOW_RELEASED_ONNX_OPSET_ONLY=0 before importing onnxruntime.backend." + ) pytest_plugins = ("onnx.backend.test.report",) diff --git a/onnxruntime/test/python/onnx_node_test_equivalence_test.py b/onnxruntime/test/python/onnx_node_test_equivalence_test.py index 1daa2f4d1eac7..0516f2c2c28a8 100644 --- a/onnxruntime/test/python/onnx_node_test_equivalence_test.py +++ b/onnxruntime/test/python/onnx_node_test_equivalence_test.py @@ -30,9 +30,8 @@ backstop. Version parity: the materialized model.onnx opset/IR is stamped from the -compiled onnx schema registry, so this test asserts the installed onnx equals -the ``cmake/deps.txt`` pin (compared on the release base, so a pre-release rcN/.dev -wheel of the pinned tag is accepted). numpy parity is NOT asserted -- the core +compiled onnx schema registry, so this test asserts the installed onnx exactly +equals the ``cmake/deps.txt`` pin. numpy parity is NOT asserted -- the core treats numpy as an informational soft-warning, and recomputed float outputs that differ by within a small ULP band are accepted as Class A numpy-gen skew. """ @@ -44,6 +43,7 @@ import sys import tempfile import unittest +from unittest import mock # Make the sibling materializer importable regardless of CWD. _TOOLS_PYTHON = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "tools", "python")) @@ -62,14 +62,14 @@ def _repo_root() -> str: def _expected_onnx_version() -> str | None: - """Parse the onnx pin (archive tag) from cmake/deps.txt.""" + """Derive the expected wheel version from the archive pin in cmake/deps.txt.""" deps = os.path.join(_repo_root(), "cmake", "deps.txt") if not os.path.isfile(deps): return None with open(deps) as f: for line in f: if line.startswith("onnx;"): - m = re.search(r"/tags/v([0-9]+\.[0-9]+\.[0-9]+)\.zip", line) + m = re.search(r"/v([0-9]+\.[0-9]+\.[0-9]+)\.zip;", line) if m: return m.group(1) return None @@ -123,6 +123,14 @@ def _find_oracle_node_dir() -> str | None: class NodeTestEquivalence(unittest.TestCase): + def test_expected_version_comes_from_archive_pin(self) -> None: + deps = "# ONNX_VERSION=1.23.0\nonnx;https://github.com/onnx/onnx/archive/refs/tags/v1.23.1.zip;archive_hash\n" + with ( + mock.patch("os.path.isfile", return_value=True), + mock.patch("builtins.open", mock.mock_open(read_data=deps)), + ): + self.assertEqual(_expected_onnx_version(), "1.23.1") + def test_materialized_equals_on_disk_oracle(self) -> None: oracle = _find_oracle_node_dir() if oracle is None: @@ -138,12 +146,11 @@ def test_materialized_equals_on_disk_oracle(self) -> None: except ImportError as exc: # pragma: no cover self.skipTest(f"onnx not importable: {exc}") - if expected_onnx and mat._release_base(onnx.__version__) != mat._release_base(expected_onnx): + if expected_onnx and onnx.__version__ != expected_onnx: self.skipTest( f"installed onnx=={onnx.__version__} != cmake/deps.txt pin " f"{expected_onnx}; cannot certify equivalence under version skew. " - f"Install the pinned onnx to run this check. (A pre-release rcN/.dev " - f"of the same release base is accepted.)" + "Install the pinned onnx to run this check." ) with tempfile.TemporaryDirectory(prefix="ort_node_mat_") as tmp: diff --git a/onnxruntime/test/python/onnxruntime_test_python_backend.py b/onnxruntime/test/python/onnxruntime_test_python_backend.py index bb83f6d36011f..2cfa152b94a17 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_backend.py +++ b/onnxruntime/test/python/onnxruntime_test_python_backend.py @@ -5,13 +5,18 @@ import os import tempfile import unittest +from types import SimpleNamespace +from unittest.mock import patch import numpy as np from helper import get_name from numpy.testing import assert_allclose +from onnx import TensorProto, helper import onnxruntime as onnxrt import onnxruntime.backend as backend +from onnxruntime.backend.backend import OnnxRuntimeBackend +from onnxruntime.backend.backend_rep import OnnxRuntimeBackendRep class TestBackend(unittest.TestCase): @@ -23,6 +28,84 @@ def test_run_model(self): output_expected = np.array([[1.0, 4.0], [9.0, 16.0], [25.0, 36.0]], dtype=np.float32) np.testing.assert_allclose(res[0], output_expected, rtol=1e-05, atol=1e-08) + @patch.dict(os.environ, {"ALLOW_RELEASED_ONNX_OPSET_ONLY": "1"}) + def test_prepare_released_opset28_model_when_policy_enabled(self): + model = self._make_opset28_model() + original_policy = OnnxRuntimeBackend.allowReleasedOpsetsOnly + try: + OnnxRuntimeBackend.allowReleasedOpsetsOnly = True + rep = backend.prepare(model) + finally: + OnnxRuntimeBackend.allowReleasedOpsetsOnly = original_policy + + result = rep.run(np.array([1.0], dtype=np.float32)) + np.testing.assert_array_equal(result[0], np.array([1.0], dtype=np.float32)) + + @patch.dict(os.environ, {"ALLOW_RELEASED_ONNX_OPSET_ONLY": "1"}) + def test_prepare_model_bytes_rejects_unreleased_opset29_when_policy_enabled(self): + model = self._make_opset29_model() + original_policy = OnnxRuntimeBackend.allowReleasedOpsetsOnly + try: + OnnxRuntimeBackend.allowReleasedOpsetsOnly = True + with self.assertRaisesRegex(onnxrt.capi.onnxruntime_pybind11_state.Fail, "Opset 29 is under development"): + backend.prepare(model.SerializeToString()) + finally: + OnnxRuntimeBackend.allowReleasedOpsetsOnly = original_policy + + @patch.dict(os.environ, {"ALLOW_RELEASED_ONNX_OPSET_ONLY": "1"}) + def test_prepare_model_bytes_allows_unreleased_opset29_when_policy_disabled(self): + model = self._make_opset29_model() + + original_policy = OnnxRuntimeBackend.allowReleasedOpsetsOnly + try: + OnnxRuntimeBackend.allowReleasedOpsetsOnly = False + rep = backend.prepare(model.SerializeToString()) + finally: + OnnxRuntimeBackend.allowReleasedOpsetsOnly = original_policy + + result = rep.run(np.array([1.0], dtype=np.float32)) + np.testing.assert_array_equal(result[0], np.array([1.0], dtype=np.float32)) + + @patch.dict(os.environ, {"ALLOW_RELEASED_ONNX_OPSET_ONLY": "1"}) + def test_prepare_model_path_allows_unreleased_opset29_when_policy_disabled(self): + model = self._make_opset29_model() + with tempfile.TemporaryDirectory() as tmpdir: + model_path = os.path.join(tmpdir, "opset29.onnx") + with open(model_path, "wb") as model_file: + model_file.write(model.SerializeToString()) + + original_policy = OnnxRuntimeBackend.allowReleasedOpsetsOnly + try: + OnnxRuntimeBackend.allowReleasedOpsetsOnly = False + rep = backend.prepare(model_path) + finally: + OnnxRuntimeBackend.allowReleasedOpsetsOnly = original_policy + + result = rep.run(np.array([1.0], dtype=np.float32)) + np.testing.assert_array_equal(result[0], np.array([1.0], dtype=np.float32)) + + @staticmethod + def _make_opset29_model(): + # A passthrough graph tests load policy without requiring unregistered opset-29 schemas. + value_info = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1]) + return helper.make_model( + helper.make_graph([], "opset29", [value_info], [value_info]), + opset_imports=[helper.make_opsetid("", 29)], + ) + + @staticmethod + def _make_opset28_model(): + model = helper.make_model( + helper.make_graph( + [helper.make_node("Identity", ["X"], ["Y"])], + "opset28", + [helper.make_tensor_value_info("X", TensorProto.FLOAT, [1])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1])], + ), + opset_imports=[helper.make_opsetid("", 28)], + ) + return model + def test_allocation_plan_works_with_only_execute_path_to_fetches_option(self): """ (inp0) (inp1) @@ -66,6 +149,54 @@ def test_allocation_plan_works_with_only_execute_path_to_fetches_option(self): assert_allclose(session_run_results[0], -(inp0 - inp1)) +class TestBackendRepInputNormalization(unittest.TestCase): + class _Session: + def __init__(self, input_names): + self._inputs = [SimpleNamespace(name=name) for name in input_names] + self.received_inputs = None + + def get_inputs(self): + return self._inputs + + def run(self, _, inputs, __): + self.received_inputs = inputs + return [inputs] + + def test_run_normalizes_numpy_scalar_inputs_to_zero_dimensional_arrays(self): + session = self._Session(["first", "second"]) + rep = OnnxRuntimeBackendRep(session) + + rep.run([np.int32(3), np.float32(4.5)]) + + self.assertEqual(session.received_inputs["first"].shape, ()) + self.assertEqual(session.received_inputs["second"].shape, ()) + self.assertEqual(session.received_inputs["first"].dtype, np.dtype(np.int32)) + self.assertEqual(session.received_inputs["second"].dtype, np.dtype(np.float32)) + + def test_run_normalizes_single_numpy_scalar_to_zero_dimensional_array(self): + session = self._Session(["input"]) + rep = OnnxRuntimeBackendRep(session) + + rep.run(np.int64(7)) + + self.assertEqual(session.received_inputs["input"].shape, ()) + self.assertEqual(session.received_inputs["input"].dtype, np.dtype(np.int64)) + + def test_run_preserves_non_scalar_inputs(self): + values = ( + (np.array([1], dtype=np.int32),), + {"key": np.array([2], dtype=np.int32)}, + ) + for value in values: + with self.subTest(value=value): + session = self._Session(["input"]) + rep = OnnxRuntimeBackendRep(session) + + rep.run(value) + + self.assertIs(session.received_inputs["input"], value) + + class TestBackendKwargsAllowlist(unittest.TestCase): """Tests that the SessionOptions/RunOptions kwargs allowlist correctly blocks dangerous attributes and allows safe ones, preventing arbitrary file writes diff --git a/onnxruntime/test/python/onnxruntime_test_python_cuda_kernel.py b/onnxruntime/test/python/onnxruntime_test_python_cuda_kernel.py new file mode 100644 index 0000000000000..fdd2f90dff90a --- /dev/null +++ b/onnxruntime/test/python/onnxruntime_test_python_cuda_kernel.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import json +import tempfile +import unittest +from pathlib import Path + +import numpy as np +from onnx import TensorProto, helper + +import onnxruntime as onnxrt + + +def create_mul_model(): + graph = helper.make_graph( + [helper.make_node("Mul", ["x", "y"], ["z"])], + "mul_kernel", + [helper.make_tensor_value_info(name, TensorProto.FLOAT, ["size"]) for name in ("x", "y")], + [helper.make_tensor_value_info("z", TensorProto.FLOAT, ["size"])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + model.ir_version = 7 + return model + + +class TestCudaKernel(unittest.TestCase): + def _run_mul(self, provider): + self.assertIn(provider, onnxrt.get_available_providers()) + with tempfile.TemporaryDirectory() as directory: + options = onnxrt.SessionOptions() + options.graph_optimization_level = onnxrt.GraphOptimizationLevel.ORT_DISABLE_ALL + options.enable_profiling = True + options.profile_file_prefix = str(Path(directory) / "mul_profile") + if provider == "CUDAExecutionProvider": + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + + session = onnxrt.InferenceSession( + create_mul_model().SerializeToString(), sess_options=options, providers=[provider] + ) + try: + session.disable_fallback() + self.assertIn(provider, session.get_providers()) + for size in (17, 257): + with self.subTest(provider=provider, size=size): + x = np.arange(size, dtype=np.float32) - np.float32(size // 2) + y = (np.arange(size, dtype=np.float32) % np.float32(7) - np.float32(3)) / np.float32(4) + actual = session.run(["z"], {"x": x, "y": y})[0] + np.testing.assert_array_equal(actual, x * y) + finally: + profile_path = session.end_profiling() + + with open(profile_path) as profile: + events = json.load(profile) + mul_events = [ + event + for event in events + if event.get("cat") == "Node" + and event.get("name", "").endswith("_kernel_time") + and event.get("args", {}).get("op_name") == "Mul" + ] + self.assertEqual(len(mul_events), 2, "Expected a profiled Mul kernel for each inference") + self.assertEqual({event["args"].get("provider") for event in mul_events}, {provider}) + + def test_mul_cpu(self): + self._run_mul("CPUExecutionProvider") + + def test_mul_cuda(self): + self._run_mul("CUDAExecutionProvider") + + +if __name__ == "__main__": + unittest.main() diff --git a/onnxruntime/test/python/onnxruntime_test_python_fpa_intb_tune.py b/onnxruntime/test/python/onnxruntime_test_python_fpa_intb_tune.py new file mode 100644 index 0000000000000..48f5c343ea976 --- /dev/null +++ b/onnxruntime/test/python/onnxruntime_test_python_fpa_intb_tune.py @@ -0,0 +1,591 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import contextlib +import csv +import gc +import io +import os +import subprocess +import sys +import tempfile +import time +import types +import unittest +from pathlib import Path +from typing import ClassVar +from unittest.mock import MagicMock, patch + +import numpy as np +from onnx import TensorProto, helper, numpy_helper, save + +import onnxruntime as ort +from onnxruntime.tools import fpa_intb_tune as tune + + +class TestFpAIntBTune(unittest.TestCase): + signature: ClassVar[dict[str, str]] = { + "device_name": "Test GPU", + "sm": "80", + "cuda_runtime": "13040", + "ort_version": "1.28.0", + } + + def write_cache(self, prefix, *, tactic_values=None, **overrides): + header = { + **self.signature, + "ort_cuda_gemm_tactic_cache": "v1", + "table": "matmulnbits_fpa_intb", + "tactic_selection_version": "4", + **overrides, + } + columns = (*tune._CACHE_KEY_COLUMNS, "m_bucket", *tune._CACHE_CONFIG_COLUMNS) + row = dict.fromkeys(columns, "0") + row.update( + n_16b="16", + k="128", + activation_dtype="half", + weight_type="uint4b_t", + bits="4", + block_size="32", + zero_point_dtype="none", + gemv_enabled="1", + packing_sm="80", + m_bucket="1", + valid_config="1", + sm_version="80", + split_k="1", + stages="3", + ) + row.update(tactic_values or {}) + path = Path(str(prefix) + tune._CACHE_TABLE_SUFFIX) + path.write_text( + "".join(f"# {k}\t{v}\n" for k, v in header.items()) + + "\t".join(columns) + + "\n" + + "\t".join(row[c] for c in columns) + + "\n", + encoding="utf-8", + ) + return path + + def test_cpu_fallback_rejects_stale_output(self): + with tempfile.TemporaryDirectory() as directory: + prefix = str(Path(directory) / "cache") + path = self.write_cache(prefix, device_name="Stale GPU") + original = path.read_bytes() + session = MagicMock() + session.get_providers.return_value = ["CPUExecutionProvider"] + with ( + patch.object(ort, "get_available_providers", return_value=["CUDAExecutionProvider"]), + patch.object(ort, "InferenceSession", return_value=session) as create, + self.assertRaisesRegex(RuntimeError, "did not activate"), + ): + tune.tune("model.onnx", prefix, [1], False) + self.assertFalse(create.call_args.kwargs["enable_fallback"]) + self.assertEqual(original, path.read_bytes()) + + def test_valid_existing_hit_succeeds_without_rewrite(self): + with tempfile.TemporaryDirectory() as directory: + prefix = str(Path(directory) / "cache") + path = self.write_cache(prefix) + timestamp = path.stat().st_mtime_ns + session = MagicMock() + session.get_providers.return_value = ["CUDAExecutionProvider"] + session.get_provider_options.return_value = {"CUDAExecutionProvider": {"device_id": "0"}} + + def create_session(_model, options, **_kwargs): + self.write_cache(options.get_session_config_entry("ep.cuda.gemm_tactic_cache_tuning_results_prefix")) + return session + + with ( + patch.object(ort, "get_available_providers", return_value=["CUDAExecutionProvider"]), + patch.object(ort, "InferenceSession", side_effect=create_session), + patch.object(tune, "_current_signature", return_value=self.signature), + ): + self.assertEqual(tune.tune("model.onnx", prefix, [1], False), str(path)) + session.disable_fallback.assert_called_once() + self.assertEqual(timestamp, path.stat().st_mtime_ns) + + def test_unrelated_existing_cache_does_not_report_success(self): + with tempfile.TemporaryDirectory() as directory: + prefix = str(Path(directory) / "cache") + path = self.write_cache(prefix) + original = path.read_bytes() + session = MagicMock() + session.get_providers.return_value = ["CUDAExecutionProvider"] + session.get_provider_options.return_value = {"CUDAExecutionProvider": {"device_id": "0"}} + with ( + patch.object(ort, "get_available_providers", return_value=["CUDAExecutionProvider"]), + patch.object(ort, "InferenceSession", return_value=session), + patch.object(tune, "_current_signature", return_value=self.signature), + self.assertRaisesRegex(RuntimeError, "did not select any successful"), + ): + tune.tune("unrelated.onnx", prefix, [1], False) + self.assertEqual(original, path.read_bytes()) + + def test_cpu_package_with_cuda_plugin_reports_metadata_limitation(self): + with ( + patch.dict(sys.modules, {"onnxruntime.capi.build_and_package_info": types.SimpleNamespace()}), + self.assertRaisesRegex(RuntimeError, "CPU packages with a registered CUDA plugin are not supported"), + ): + tune._current_signature(0) + + def test_selected_model_entries_must_reach_output_cache(self): + with tempfile.TemporaryDirectory() as directory: + prefix = str(Path(directory) / "cache") + path = self.write_cache(prefix) + session = MagicMock() + session.get_providers.return_value = ["CUDAExecutionProvider"] + session.get_provider_options.return_value = {"CUDAExecutionProvider": {"device_id": "0"}} + + def create_session(_model, options, **_kwargs): + report = self.write_cache( + options.get_session_config_entry("ep.cuda.gemm_tactic_cache_tuning_results_prefix") + ) + report.write_text(report.read_text().replace("\n16\t", "\n32\t"), encoding="utf-8") + return session + + with ( + patch.object(ort, "get_available_providers", return_value=["CUDAExecutionProvider"]), + patch.object(ort, "InferenceSession", side_effect=create_session), + patch.object(tune, "_current_signature", return_value=self.signature), + self.assertRaisesRegex(RuntimeError, "not persisted"), + ): + tune.tune("model.onnx", prefix, [1], False) + self.assertTrue(path.exists()) + + def test_corrected_settings_not_saved_does_not_report_success(self): + for column in tune._CACHE_CONFIG_COLUMNS[1:]: + with self.subTest(column=column), tempfile.TemporaryDirectory() as directory: + prefix = str(Path(directory) / "cache") + path = self.write_cache(prefix, tactic_values={column: "3" if column == "split_k" else "99"}) + original = path.read_bytes() + session = MagicMock() + session.get_providers.return_value = ["CUDAExecutionProvider"] + session.get_provider_options.return_value = {"CUDAExecutionProvider": {"device_id": "0"}} + + def create_session(_model, options, _session=session, **_kwargs): + # Simulate a failed main-cache flush while the separate report is writable. + self.write_cache( + options.get_session_config_entry("ep.cuda.gemm_tactic_cache_tuning_results_prefix") + ) + return _session + + output = io.StringIO() + with ( + patch.object(ort, "get_available_providers", return_value=["CUDAExecutionProvider"]), + patch.object(ort, "InferenceSession", side_effect=create_session), + patch.object(tune, "_current_signature", return_value=self.signature), + contextlib.redirect_stdout(output), + self.assertRaisesRegex(RuntimeError, "not persisted"), + ): + tune.tune("model.onnx", prefix, [1], False) + self.assertNotIn("Cache ready", output.getvalue()) + self.assertEqual(original, path.read_bytes()) + + def test_missing_tactic_settings_are_rejected(self): + with tempfile.TemporaryDirectory() as directory: + path = self.write_cache(Path(directory) / "cache") + text = path.read_text(encoding="utf-8").replace("split_k\t", "missing_split_k\t") + path.write_text(text, encoding="utf-8") + with self.assertRaisesRegex(RuntimeError, "Missing problem-key or tactic-setting columns"): + tune._summarize_cache(str(path), self.signature) + + def test_selected_tactics_report(self): + with tempfile.TemporaryDirectory() as directory: + prefix = Path(directory) / "selected" + path = self.write_cache(prefix) + selected = _read_selected_tactics(prefix) + self.assertEqual(set(selected), {1}) + self.assertEqual(selected[1]["split_k"], "1") + self.assertEqual(selected[1]["activation_dtype"], "half") + path.write_text(path.read_text().replace("\t1\t80\t", "\t0\t80\t"), encoding="utf-8") + with self.assertRaisesRegex(RuntimeError, "No successful selected tactics"): + _read_selected_tactics(prefix) + path.write_text("\t".join((*tune._CACHE_KEY_COLUMNS, "m_bucket", *tune._CACHE_CONFIG_COLUMNS)) + "\n") + with self.assertRaisesRegex(RuntimeError, "No successful selected tactics"): + _read_selected_tactics(prefix) + + def test_incompatible_or_missing_cache_is_an_error(self): + with tempfile.TemporaryDirectory() as directory: + prefix = str(Path(directory) / "cache") + for field in ("device_name", "sm", "cuda_runtime", "ort_version", "tactic_selection_version"): + with self.subTest(field=field): + path = self.write_cache(prefix, **{field: "stale"}) + with self.assertRaisesRegex(RuntimeError, "Inapplicable"): + tune._summarize_cache(str(path), self.signature) + path.unlink() + with self.assertRaisesRegex(RuntimeError, "No tactic cache"): + tune._summarize_cache(str(path), self.signature) + + def test_signature_uses_native_ort_version(self): + class FakeDriver: + def __getattr__(self, _): + return lambda *args: 0 + + fake_build_info = types.SimpleNamespace(cuda_version="13.0") + with ( + patch.dict(sys.modules, {"onnxruntime.capi.build_and_package_info": fake_build_info}), + patch.object(tune.ctypes, "CDLL", return_value=FakeDriver(), create=True), + patch.object(tune.ctypes, "WinDLL", return_value=FakeDriver(), create=True), + patch.object(tune.ort, "get_version_string", return_value="1.28.0"), + patch.object(tune.ort, "__version__", "1.28.0.dev20260101"), + ): + signature = tune._current_signature(0) + + self.assertEqual(signature["ort_version"], "1.28.0") + + def test_bf16_dummy_feed_uses_typed_ortvalue(self): + model = helper.make_model( + helper.make_graph( + [helper.make_node("Identity", ["x"], ["y"])], + "bf16", + [helper.make_tensor_value_info("x", TensorProto.BFLOAT16, ["m", 4])], + [helper.make_tensor_value_info("y", TensorProto.BFLOAT16, ["m", 4])], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=10, + ) + session = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + feeds = tune._make_dummy_inputs(session, 3) + self.assertEqual(feeds["x"].element_type(), TensorProto.BFLOAT16) + outputs = session.run_with_ort_values(None, feeds) + self.assertEqual(outputs[0].shape(), [3, 4]) + self.assertEqual(outputs[0].element_type(), TensorProto.BFLOAT16) + + +def _lock_file(file): + if sys.platform == "win32": + import msvcrt # noqa: PLC0415 + + file.seek(0) + msvcrt.locking(file.fileno(), msvcrt.LK_LOCK, 1) + else: + import fcntl # noqa: PLC0415 + + fcntl.flock(file, fcntl.LOCK_EX) + + +def _unlock_file(file): + if sys.platform == "win32": + import msvcrt # noqa: PLC0415 + + file.seek(0) + msvcrt.locking(file.fileno(), msvcrt.LK_UNLCK, 1) + else: + import fcntl # noqa: PLC0415 + + fcntl.flock(file, fcntl.LOCK_UN) + + +def _read_selected_tactics(prefix): + with open(str(prefix) + tune._CACHE_TABLE_SUFFIX, encoding="utf-8") as file: + rows = list(csv.DictReader((line for line in file if not line.startswith("#")), delimiter="\t")) + if not rows or any(row["valid_config"] != "1" for row in rows): + raise RuntimeError("No successful selected tactics") + return {int(row["m_bucket"]): row for row in rows} + + +def _gpu_worker(model, prefix, output, chunk="0"): + plugin = os.environ.get("ORT_CUDA_PLUGIN_PATH") + # Plugin wheels may have already registered CUDA during import. + if plugin and not any(device.ep_name == "CUDAExecutionProvider" for device in ort.get_ep_devices()): + ort.register_execution_provider_library("CUDAExecutionProvider", plugin) + ort.set_default_logger_severity(1) + options = tune._make_session_options(prefix, [1, 64]) + options.add_session_config_entry("ep.cuda.gemm_tactic_cache_tuning_results_prefix", output + ".selected") + if chunk != "0": + os.environ["ORT_MATMULNBITS_FORCE_CHUNKED"] = "1" + options.add_session_config_entry("ep.cuda.matmul_nbits_m_chunk_size", chunk) + options.log_severity_level = 1 + session = ort.InferenceSession(model, options, providers=["CUDAExecutionProvider"], enable_fallback=False) + if "CUDAExecutionProvider" not in session.get_providers(): + raise RuntimeError("CUDA is not active") + session.disable_fallback() + outputs = [] + for m in (1, 64, 65): # 65 exercises lazy bucket 128 and its teardown flush. + inputs = np.arange(m * 128, dtype=np.float32).reshape(m, 128) % 17 / 17 + outputs.append(session.run(None, {"A": inputs.astype(np.float16)})[0]) + np.savez(output, *outputs) + Path(output + ".ready").touch() + del session + gc.collect() + + +class TestFpAIntBWorker(unittest.TestCase): + def test_cuda_plugin_registration(self): + cuda_device = types.SimpleNamespace(ep_name="CUDAExecutionProvider") + cpu_device = types.SimpleNamespace(ep_name="CPUExecutionProvider") + cases = ( + ("no_plugin", "", [], False), + ("unregistered", "/cuda.so", [], True), + ("other_provider", "/cuda.so", [cpu_device], True), + ("already_registered", "/cuda.so", [cpu_device, cuda_device], False), + ) + for name, plugin, devices, should_register in cases: + with self.subTest(name=name): + session = MagicMock() + session.get_providers.return_value = ["CUDAExecutionProvider"] + with ( + patch.dict(os.environ, {"ORT_CUDA_PLUGIN_PATH": plugin}), + patch.object(ort, "get_ep_devices", return_value=devices) as get_devices, + patch.object(ort, "register_execution_provider_library") as register, + patch.object(ort, "set_default_logger_severity"), + patch.object(tune, "_make_session_options") as options, + patch.object(ort, "InferenceSession", return_value=session) as create, + patch.object(np, "savez"), + patch.object(Path, "touch"), + ): + _gpu_worker("model.onnx", "cache", "output.npz") + if should_register: + register.assert_called_once_with("CUDAExecutionProvider", plugin) + else: + register.assert_not_called() + if plugin: + get_devices.assert_called_once_with() + else: + get_devices.assert_not_called() + self.assertFalse(create.call_args.kwargs["enable_fallback"]) + options.return_value.add_session_config_entry.assert_called_once_with( + "ep.cuda.gemm_tactic_cache_tuning_results_prefix", "output.npz.selected" + ) + session.disable_fallback.assert_called_once_with() + self.assertEqual(session.run.call_count, 3) + + def test_cuda_plugin_registration_failure_is_not_suppressed(self): + with ( + patch.dict(os.environ, {"ORT_CUDA_PLUGIN_PATH": "/cuda.so"}), + patch.object(ort, "get_ep_devices", return_value=[]), + patch.object(ort, "register_execution_provider_library", side_effect=RuntimeError("Plugin load failed")), + patch.object(ort, "InferenceSession") as create, + self.assertRaisesRegex(RuntimeError, "Plugin load failed"), + ): + _gpu_worker("model.onnx", "cache", "output.npz") + create.assert_not_called() + + +@unittest.skipUnless(os.environ.get("ORT_RUN_CUDA_TACTIC_CACHE_TESTS") == "1", "opt-in CUDA cache integration") +class TestFpAIntBCacheCuda(unittest.TestCase): + def test_fresh_process_reuse_repair_and_lock_contention(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + rng = np.random.default_rng(42) + weights = rng.integers(0, 256, size=(64, 4, 16), dtype=np.uint8) + scales = np.full((64, 4), 0.125, dtype=np.float16) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "MatMulNBits", + ["A", "B", "scales"], + ["Y"], + domain="com.microsoft", + K=128, + N=64, + bits=4, + block_size=32, + ) + ], + "cache", + [helper.make_tensor_value_info("A", TensorProto.FLOAT16, ["m", 128])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT16, ["m", 64])], + [numpy_helper.from_array(weights, "B"), numpy_helper.from_array(scales, "scales")], + ), + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ir_version=10, + ) + model_path = root / "model.onnx" + save(model, model_path) + prefix = str(root / "cache") + + def command(output, chunk=0): + return [sys.executable, __file__, "--gpu-worker", str(model_path), prefix, str(output), str(chunk)] + + def run(name, chunk=0): + output = root / f"{name}.npz" + result = subprocess.run( + command(output, chunk), check=False, capture_output=True, text=True, timeout=180 + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + selected = _read_selected_tactics(str(output) + ".selected") + with np.load(output) as arrays: + return {key: arrays[key] for key in arrays.files}, result.stderr, selected + + cold, _, cold_tactics = run("cold") + cache_path = Path(prefix + tune._CACHE_TABLE_SUFFIX) + self.assertTrue(cache_path.exists()) + original = cache_path.read_bytes() + warm, log, warm_tactics = run("warm") + self.assertEqual(cold_tactics, warm_tactics) + # Plugin loggers cache their severity during automatic import-time registration. + if not os.environ.get("ORT_CUDA_PLUGIN_PATH"): + self.assertIn("validated fpA_intB tactics from", log) + self.assertEqual(original, cache_path.read_bytes()) + for key in cold: + np.testing.assert_allclose(warm[key], cold[key], rtol=1e-3, atol=1e-3) + + # Exercise the actual tool's session-specific evidence, not just its file parser. + args = [ + sys.executable, + "-m", + "onnxruntime.tools.fpa_intb_tune", + "--model", + str(model_path), + "--output-prefix", + prefix, + "--m-values", + "1,64", + "--no-inference", + ] + if not os.environ.get("ORT_CUDA_PLUGIN_PATH"): + result = subprocess.run(args, check=False, capture_output=True, text=True, timeout=180) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertIn("selected this run:", result.stdout) + self.assertEqual(original, cache_path.read_bytes()) + + unrelated = helper.make_model( + helper.make_graph( + [helper.make_node("Identity", ["A"], ["Y"])], + "unrelated", + [helper.make_tensor_value_info("A", TensorProto.FLOAT, [1])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1])], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=10, + ) + unrelated_path = root / "unrelated.onnx" + save(unrelated, unrelated_path) + args[args.index("--model") + 1] = str(unrelated_path) + result = subprocess.run(args, check=False, capture_output=True, text=True, timeout=180) + self.assertNotEqual(result.returncode, 0) + self.assertIn("did not select any successful", result.stderr) + self.assertNotIn("Cache ready", result.stdout) + self.assertEqual(original, cache_path.read_bytes()) + + lines = cache_path.read_text(encoding="utf-8").splitlines() + header_index = next(i for i, line in enumerate(lines) if not line.startswith("#")) + columns = lines[header_index].split("\t") + large = next( + line.split("\t") + for line in lines[header_index + 1 :] + if line.split("\t")[columns.index("m_bucket")] == "64" + ) + large[columns.index("m_bucket")] = "8192" + cache_path.write_text("\n".join([*lines, "\t".join(large)]) + "\n", encoding="utf-8") + chunked, log, chunked_tactics = run("chunked", chunk=64) + self.assertEqual(set(chunked_tactics), {1, 64}) + self.assertEqual(chunked_tactics, {m: cold_tactics[m] for m in (1, 64)}) + if not os.environ.get("ORT_CUDA_PLUGIN_PATH"): + self.assertIn("Loaded 2 validated fpA_intB tactics", log) + unpacked = np.empty((64, 128), dtype=np.float32) + packed = weights.reshape(64, 64) + unpacked[:, ::2] = packed & 15 + unpacked[:, 1::2] = packed >> 4 + dense = (unpacked - 8) * 0.125 + for key, m in zip(cold, (1, 64, 65), strict=True): + inputs = (np.arange(m * 128, dtype=np.float32).reshape(m, 128) % 17 / 17).astype(np.float16) + reference = inputs.astype(np.float32) @ dense.T + # A final one-row chunk uses GEMV instead of GEMM, with different FP16 rounding. + np.testing.assert_allclose(chunked[key], reference, rtol=1e-3, atol=1e-2) + np.testing.assert_allclose(cold[key], reference, rtol=1e-3, atol=1e-2) + cache_path.write_bytes(original) + + for i in range(header_index + 1, len(lines)): + row = lines[i].split("\t") + if row[columns.index("m_bucket")] == "64": + row[columns.index("split_k")] = "3" # K=128 cannot be split into three aligned K tiles. + lines[i] = "\t".join(row) + cache_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + if not os.environ.get("ORT_CUDA_PLUGIN_PATH"): + args[args.index("--model") + 1] = str(model_path) + damaged = cache_path.read_bytes() + lock_path = Path(str(cache_path) + ".lock") + lock_path.unlink() + lock_path.mkdir() # Make the main cache's file-lock acquisition fail. + try: + result = subprocess.run(args, check=False, capture_output=True, text=True, timeout=180) + self.assertNotEqual(result.returncode, 0) + self.assertIn("not persisted", result.stderr) + self.assertNotIn("Cache ready", result.stdout) + self.assertEqual(damaged, cache_path.read_bytes()) + finally: + lock_path.rmdir() + repaired, log, repaired_tactics = run("repaired") + self.assertNotEqual(repaired_tactics[64]["split_k"], "3") + self.assertTrue("Rejecting cached GEMM tactic" in log or "Dropping incompatible cached" in log, log) + for key in cold: + np.testing.assert_allclose(repaired[key], cold[key], rtol=1e-3, atol=1e-3) + + # Holding the cross-process file lock must not prevent construction or cached inference. + # Remove an initial bucket to force construction to produce dirty rows; a pure cache hit + # would not exercise the old construction-time Flush() path. + lines = cache_path.read_text(encoding="utf-8").splitlines() + lines = lines[: header_index + 1] + [ + line for line in lines[header_index + 1 :] if line.split("\t")[columns.index("m_bucket")] != "64" + ] + cache_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + with open(str(cache_path) + ".lock", "a+b") as lock, tempfile.TemporaryFile(mode="w+t") as log_file: + _lock_file(lock) + output = root / "locked.npz" + process = subprocess.Popen(command(output), stdout=log_file, stderr=subprocess.STDOUT, text=True) + try: + deadline = time.monotonic() + 60 + while not Path(str(output) + ".ready").exists() and process.poll() is None: + if time.monotonic() >= deadline: + self.fail("File lock blocked session construction or inference") + time.sleep(0.1) + self.assertTrue(Path(str(output) + ".ready").exists()) + finally: + _unlock_file(lock) + try: + process.wait(timeout=60) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + raise + log_file.seek(0) + self.assertEqual(process.returncode, 0, log_file.read()) + + # Loading a relative model path and changing cwd before closing the session must still write + # the sidecar cache next to the model rather than under the new working directory. + models_dir = root / "rel_models" + other_cwd = root / "other_cwd" + models_dir.mkdir() + other_cwd.mkdir() + rel_model_abs = models_dir / "model.onnx" + save(model, rel_model_abs) + saved_cwd = Path.cwd() + try: + os.chdir(root) + plugin = os.environ.get("ORT_CUDA_PLUGIN_PATH") + if plugin and not any(device.ep_name == "CUDAExecutionProvider" for device in ort.get_ep_devices()): + ort.register_execution_provider_library("CUDAExecutionProvider", plugin) + sidecar_options = ort.SessionOptions() + sidecar_options.add_session_config_entry("ep.cuda.fpa_intb_gemm", "1") + sidecar_options.add_session_config_entry("ep.cuda.fpa_intb_profile_m", "1") + sidecar_options.add_session_config_entry("ep.cuda.gemm_tactic_cache_to_model", "1") + sidecar_options.add_session_config_entry("ep.cuda.gemm_tactic_cache_mode", "save") + sidecar_session = ort.InferenceSession( + str(Path("rel_models") / "model.onnx"), + sidecar_options, + providers=["CUDAExecutionProvider"], + enable_fallback=False, + ) + os.chdir(other_cwd) + del sidecar_session + gc.collect() + finally: + os.chdir(saved_cwd) + expected_sidecar = Path(str(rel_model_abs) + tune._CACHE_TABLE_SUFFIX) + wrong_sidecar = other_cwd / "rel_models" / ("model.onnx" + tune._CACHE_TABLE_SUFFIX) + self.assertTrue(expected_sidecar.exists()) + self.assertFalse(wrong_sidecar.exists()) + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "--gpu-worker": + _gpu_worker(*sys.argv[2:]) + else: + with contextlib.redirect_stdout(io.StringIO()): + unittest.main() diff --git a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py index 96bb35672a6cf..6ac581d72cab8 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py +++ b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py @@ -882,9 +882,6 @@ def infer(query_width, packed): self.assertEqual(separate_graph.output[0].type.tensor_type.shape.dim[2].dim_value, 33) self.assertFalse(separate_graph.output[1].type.tensor_type.HasField("shape")) - - - def test_unsqueeze_opset_11(self): graph = helper.make_graph( [ 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 23920aae2d539..5ebe8b49433e2 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -12,6 +12,7 @@ import sys import unittest from contextlib import contextmanager +from unittest.mock import MagicMock, patch import numpy as np from onnx import ModelProto, TensorProto, helper, numpy_helper @@ -140,7 +141,7 @@ def _check_prepacked_parity( prepacked_output = self._run_model(prepacked_model, a) except Exception as exc: outside_compact_contract = block_size != 32 or has_bias or weight_prepacked != 1 - if outside_compact_contract and "compact fpA_intB build supports prepacked weights only" in str(exc): + if outside_compact_contract and "compact fpA_intB build supports prepacked weights for" in str(exc): self.skipTest("case is outside the compact fpA_intB build contract") raise @@ -191,6 +192,30 @@ def test_int8_sm90_prepacked_weight_matches_runtime_prepack(self): self._check_sm90_parity(bits=8, block_size=128, m=32) +class TestMatMulNBitsCompactContract(unittest.TestCase): + def test_skips_only_unsupported_compact_prepacked_cases(self): + message = "This compact fpA_intB build supports prepacked weights for FP16/BF16 activations" + cases = ( + (64, message, unittest.SkipTest), + (32, message, RuntimeError), + (64, "Unexpected execution failure", RuntimeError), + ) + for block_size, error, expected_exception in cases: + with self.subTest(block_size=block_size, error=error): + case = TestMatMulNBitsPrepackedCuda() + weights = np.zeros((512, 256 // block_size, block_size // 2), dtype=np.uint8) + scales = np.ones(weights.shape[:2], dtype=np.float16) + packer = MagicMock() + packer.pack_weights_for_cuda_mixed_gemm.return_value = weights.reshape(-1).view(np.int8) + with ( + patch.object(sys.modules[__name__], "_cuda_quant", packer), + patch.object(case, "_quantize_weight", return_value=(weights, scales)), + patch.object(case, "_run_model", side_effect=[np.zeros((1, 512)), RuntimeError(error)]), + self.assertRaises(expected_exception), + ): + case._check_prepacked_parity(bits=4, block_size=block_size, m=1) + + @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): @@ -224,6 +249,16 @@ def test_torch_packer_matches_cuda_oracle(self): with self.subTest(bits=bits, force_arch=force_arch, n=n, k=k): self._check(bits, force_arch, n, k) + def test_explicit_sm90_layout_does_not_fall_back_to_sm80(self): + # Layout selection must not depend on whether SM90 compute kernels were compiled. + n = k = 128 + for bits in (4, 8): + with self.subTest(bits=bits): + q = np.random.default_rng(42).integers(0, 256, size=(n, k // (8 // bits)), dtype=np.uint8) + sm80 = np.asarray(_cuda_quant.pack_weights_for_cuda_mixed_gemm(q, n, k, bits, 80)) + sm90 = np.asarray(_cuda_quant.pack_weights_for_cuda_mixed_gemm(q, n, k, bits, 90)) + self.assertFalse(np.array_equal(sm80, sm90), "force_arch=90 must preserve the SM90 weight layout") + @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") diff --git a/onnxruntime/test/python/requirements.txt b/onnxruntime/test/python/requirements.txt index 296ac83f3b6b5..fc11cff5648c0 100644 --- a/onnxruntime/test/python/requirements.txt +++ b/onnxruntime/test/python/requirements.txt @@ -1,3 +1,3 @@ -onnx==1.22.0 +onnx==1.23.2 pytest onnx-ir diff --git a/onnxruntime/test/python/transformers/benchmark_qwen38_mtp_kernels.py b/onnxruntime/test/python/transformers/benchmark_qwen38_mtp_kernels.py new file mode 100644 index 0000000000000..e338337beee87 --- /dev/null +++ b/onnxruntime/test/python/transformers/benchmark_qwen38_mtp_kernels.py @@ -0,0 +1,858 @@ +import argparse +import functools +import json +import os +import re +import sys +from pathlib import Path + +import cupy +import numpy + +ROOT = Path(__file__).resolve().parents[3] +QSA_SOURCE = ROOT / "contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu" + + +@functools.cache +def compile_device_code(code, name): + cuda_path = os.environ.get("CUDA_PATH") or cupy.cuda.get_cuda_path() + options = ("--std=c++17", f"-I{cuda_path}/include", f"-I{cuda_path}/include/cccl") + for arch in filter(None, os.environ.get("ORT_KERNEL_PROBE_ARCHS", "").split(",")): + cupy.cuda.compiler.compile_using_nvrtc(code, options=options, arch=arch, name_expressions=(name,)) + module = cupy.RawModule(code=code, options=options, name_expressions=(name,)) + return module.get_function(name) + + +def device_function(source, name, qualifier="__device__"): + match = re.search(qualifier + r"[^\n]*\b" + name + r"\(", source) + if match is None: + raise ValueError(f"Missing device function: {name}") + body = source.index("{", match.end()) + depth = 1 + end = body + 1 + while depth: + depth += (source[end] == "{") - (source[end] == "}") + end += 1 + template = re.search(r"template\s*<[^;{}]*>\s*$", source[: match.start()]) + start = template.start() if template is not None else match.start() + return source[start:end] + + +def qsa_prefix(source): + header = (QSA_SOURCE.parent / "packed_sparse_attention_indexer_impl.h").read_text() + math_source = (QSA_SOURCE.parent / "sparse_attention_indexer_device_math.cuh").read_text() + topk_source = (ROOT / "core/providers/cuda/cu_inc/topk_warp_sort.cuh").read_text() + params = header[header.index("struct PackedSparseAttentionIndexerParams") : header.index("};") + 2] + fields = re.findall(r"\b(int|bool|float) (\w+) = [^;]+;", params) + types = {"int": numpy.int32, "bool": numpy.bool_, "float": numpy.float32} + params_dtype = numpy.dtype([(name, types[kind]) for kind, name in fields], align=True) + helpers = source[ + source.index("template \nstruct alignas") : source.index( + "template \n__global__ void QsaUpdateStateKernel" + ) + ] + prefix = """ +#include +#include +#include +namespace std { using cuda::std::numeric_limits; } +#ifndef UINT_MAX +#define UINT_MAX 0xffffffffu +#endif +typedef unsigned long long uint64_t; +typedef long long int64_t; +typedef int int32_t; +typedef unsigned int uint32_t; +template __device__ float to_float(T value) { return static_cast(value); } +""" + prefix += params + f"\nstatic_assert(sizeof(PackedSparseAttentionIndexerParams) == {params_dtype.itemsize});\n" + prefix += """ +constexpr int kHierarchicalTileBlocks = 8, kQwenNumHeads = 4, kQwenHeadSize = 128; +constexpr int kWarpSize = 32, kQwenCompressRatio = 4; +namespace psai { constexpr int kKeyStateLength = 0; } +namespace topk { constexpr uint64_t kPaddingSortKey = 0; +""" + prefix += device_function(topk_source, "WarpReduceSum") + "\n" + prefix += device_function(topk_source, "PackStableSortKey") + "\n}\n" + prefix += helpers + "\n" + device_function(math_source, "SaiCausalThreshold") + "\n" + return prefix, params_dtype + + +def compile_qsa(source, element_type, paired=False): + prefix, params_dtype = qsa_prefix(source) + body = source[ + source.index("template \n__global__ void QsaScoreTileTopKKernel") : source.index( + "__device__ __forceinline__ uint64_t QsaMergeRank" + ) + ] + name = f"{'QsaPairedScoreTileTopKKernel' if paired else 'QsaScoreTileTopKKernel'}<{element_type}>" + return compile_device_code(prefix + body, name), params_dtype + + +def qsa_score_grid(source, rows, tiles): + if "const int token = static_cast(blockIdx.y)" in source: + return (min(tiles, 128), rows) + return (rows * tiles, 1) + + +def compile_qsa_emit(source): + prefix, _ = qsa_prefix(source) + topk_source = (ROOT / "core/providers/cuda/cu_inc/topk_warp_sort.cuh").read_text() + prefix += "\nnamespace topk {\n" + device_function(topk_source, "UnpackStableSortIndex") + "\n}\n" + start = source.index("__global__ void QsaEmitHierarchicalTopKKernel") + end = source.index("// One block per query token.", start) + return compile_device_code(prefix + source[start:end], "QsaEmitHierarchicalTopKKernel") + + +def graph_latency(kernel, grid, block, args, iterations=100): + cupy.cuda.Stream.null.synchronize() + stream = cupy.cuda.Stream(non_blocking=True) + with stream: + for _ in range(10): + kernel(grid, block, args) + stream.synchronize() + stream.begin_capture() + with stream: + for _ in range(iterations): + kernel(grid, block, args) + graph = stream.end_capture() + start = cupy.cuda.Event() + stop = cupy.cuda.Event() + samples = [] + for _ in range(5): + start.record(stream) + graph.launch(stream) + stop.record(stream) + stop.synchronize() + samples.append(cupy.cuda.get_elapsed_time(start, stop) / iterations) + return float(numpy.median(samples)) + + +def key_data(values, element_type): + if element_type == "float": + return cupy.asarray(values) + if element_type == "half": + return cupy.asarray(values.astype(numpy.float16)) + bits = values.view(numpy.uint32) + rounded = bits + numpy.uint32(0x7FFF) + ((bits >> 16) & 1) + return cupy.asarray((rounded >> 16).astype(numpy.uint16)) + + +def compile_qsa_merge(source): + prefix, _ = qsa_prefix(source) + body = source[ + source.index("__device__ __forceinline__ uint64_t QsaMergeRank") : source.index( + "__global__ void QsaEmitHierarchicalTopKKernel" + ) + ] + code = prefix + "\nconstexpr int kBoundedTopKMax = 512;\n" + body + merge = compile_device_code(code, "QsaMergeTileTopKKernel") + compact = compile_device_code(code, "QsaCompactTileTopKKernel") if "QsaCompactTileTopKKernel" in body else None + return compact, merge + + +def merge_operations(kernels, data, rows, stride): + compact, merge = kernels + scratch = cupy.full(data.size, 0xBEEF, dtype=cupy.uint64) + alternate = cupy.full_like(scratch, 0xDEAD) + operations = [] + current = data + width = 8 + if compact is not None and rows >= 4: + chunks = (stride + 2047) // 2048 + compact_stride = chunks * 512 + operations.append( + (compact, (chunks, rows), (256,), (current, scratch, numpy.int32(stride), numpy.int32(compact_stride))) + ) + current, scratch = scratch, alternate + stride = compact_stride + width = 512 + while width < stride: + pairs = (stride + 2 * width - 1) // (2 * width) + operations.append( + ( + merge, + (rows * pairs,), + (256,), + (current, scratch, numpy.int32(stride), numpy.int32(width), numpy.int32(rows)), + ) + ) + current, scratch = scratch, alternate if current is data else current + width *= 2 + return operations, current[: rows * stride].reshape(rows, stride)[:, :512] + + +def launch_operations(operations): + for kernel, grid, block, arguments in operations: + kernel(grid, block, arguments) + + +def sequence_latency(operations): + def launch(_grid, _block, _arguments): + launch_operations(operations) + + return graph_latency(launch, (), (), ()) + + +def run_qsa_merge(baseline, check_only, sanitizer_case=False): + old_kernels = compile_qsa_merge(baseline) + new_kernels = compile_qsa_merge(QSA_SOURCE.read_text()) + random = numpy.random.default_rng(17) + checks = 0 + for rows in (8,) if sanitizer_case else (1, 2, 8): + for stride in (4104,) if sanitizer_case else (2048, 4104, 65536): + for visible in (9, stride) if sanitizer_case else (0, 1, 9, 1025, stride): + for tied in (False, True): + values = numpy.zeros((rows, stride), dtype=numpy.uint64) + if tied: + values[:, :visible] = numpy.uint64(1 << 63) + numpy.arange( + visible, 0, -1, dtype=numpy.int64 + ).astype(numpy.uint64) + else: + values[:, :visible] = random.integers(1, 2**63, (rows, visible), dtype=numpy.uint64) + values = numpy.sort(values.reshape(rows, -1, 8), axis=2)[:, :, ::-1].copy().reshape(rows, stride) + data = cupy.asarray(values) + expected = cupy.sort(data, axis=1)[:, ::-1][:, :512] + old_ops, old_result = merge_operations(old_kernels, data, rows, stride) + new_ops, new_result = merge_operations(new_kernels, data, rows, stride) + launch_operations(old_ops) + launch_operations(new_ops) + cupy.testing.assert_array_equal(expected, old_result) + cupy.testing.assert_array_equal(expected, new_result) + checks += 1 + if not check_only and rows in (2, 8) and stride == 65536 and visible in (1025, stride) and not tied: + old_ms = sequence_latency(old_ops) + new_ms = sequence_latency(new_ops) + print( + json.dumps( + { + "kernel": "qsa_merge_hierarchy", + "rows": rows, + "visible_blocks": visible, + "baseline_ms": old_ms, + "candidate_ms": new_ms, + "speedup": old_ms / new_ms, + } + ), + flush=True, + ) + print(f"QSA compact-merge oracle checks passed: {checks}", flush=True) + + +def run_qsa(baseline, check_only, sanitizer_case=False): + candidate = QSA_SOURCE.read_text() + random = numpy.random.default_rng(42) + checks = 0 + old_mergers = compile_qsa_merge(baseline) + new_mergers = compile_qsa_merge(candidate) + old_emit = compile_qsa_emit(baseline) + new_emit = compile_qsa_emit(candidate) + for element_type in ("float", "half", "__nv_bfloat16"): + old_kernel, params_dtype = compile_qsa(baseline, element_type) + new_kernel, _ = compile_qsa(candidate, element_type) + paired_kernel, _ = compile_qsa(candidate, element_type, paired=True) + for rows in (8,) if sanitizer_case else (1, 2, 5, 7, 8): + for capacity in (4099,) if sanitizer_case else (2048, 4099, 65536): + keys = key_data(random.standard_normal((capacity, 128), dtype=numpy.float32), element_type) + queries = cupy.asarray(random.standard_normal((rows, 4, 128), dtype=numpy.float32)) + for visible in (9,) if sanitizer_case else (0, 1, 9, 1024, capacity): + params = numpy.zeros((), dtype=params_dtype) + for name, value in { + "batch_size": 1, + "total_tokens": rows, + "num_heads": 4, + "head_size": 128, + "compress_ratio": 4, + "state_capacity": capacity, + "capacity": 2051, + "block_topk": 512, + "scale": 0.0883883461356163, + }.items(): + params[name] = value + cumulative = cupy.asarray([0, rows], dtype=cupy.int32) + past = cupy.asarray([visible * 4], dtype=cupy.int32) + positions = cupy.zeros(rows, dtype=cupy.int64) + lengths = cupy.asarray([[visible, 0]], dtype=cupy.int32) + overflow = cupy.zeros(1, dtype=cupy.int32) + tiles = (capacity + 7) // 8 + expected = cupy.full((rows, tiles * 8), 0xDEAD, dtype=cupy.uint64) + actual = cupy.full_like(expected, 0xBEEF) + common = (keys, queries, cumulative, past, positions, lengths, overflow) + old_args = (*common, expected, numpy.int32(tiles), params[()]) + new_args = (*common, actual, numpy.int32(tiles), params[()]) + old_grid = qsa_score_grid(baseline, rows, tiles) + use_paired = rows >= 4 + candidate_kernel = paired_kernel if use_paired else new_kernel + new_grid = (min(tiles, 128), (rows + 1) // 2 if use_paired else rows) + new_block = (256,) if use_paired else (1024,) + old_kernel(old_grid, (1024,), old_args) + candidate_kernel(new_grid, new_block, new_args) + cupy.testing.assert_array_equal(expected, actual) + old_merge_ops, old_result = merge_operations(old_mergers, expected, rows, tiles * 8) + new_merge_ops, new_result = merge_operations(new_mergers, actual, rows, tiles * 8) + expected_indices = cupy.full((rows, 2051), -77, dtype=cupy.int32) + actual_indices = cupy.full_like(expected_indices, -88) + expected_counts = cupy.full(rows, -77, dtype=cupy.int32) + actual_counts = cupy.full_like(expected_counts, -88) + emit_common = (cumulative, past, positions, lengths, overflow) + old_emit_args = ( + old_result, + *emit_common, + numpy.int32(old_result.strides[0] // 8), + expected_indices, + expected_counts, + params[()], + ) + new_emit_args = ( + new_result, + *emit_common, + numpy.int32(new_result.strides[0] // 8), + actual_indices, + actual_counts, + params[()], + ) + old_ops = [ + (old_kernel, old_grid, (1024,), old_args), + *old_merge_ops, + (old_emit, (rows,), (128,), old_emit_args), + ] + new_ops = [ + (candidate_kernel, new_grid, new_block, new_args), + *new_merge_ops, + (new_emit, (rows,), (128,), new_emit_args), + ] + launch_operations(old_ops) + launch_operations(new_ops) + cupy.testing.assert_array_equal(expected_indices, actual_indices) + cupy.testing.assert_array_equal(expected_counts, actual_counts) + checks += 1 + if not check_only and rows in (2, 8) and capacity == 65536 and visible in (1024, capacity): + old_ms = graph_latency(old_kernel, old_grid, (1024,), old_args) + new_ms = graph_latency(candidate_kernel, new_grid, new_block, new_args) + print( + json.dumps( + { + "kernel": "qsa_score", + "dtype": element_type, + "rows": rows, + "visible_blocks": visible, + "baseline_ms": old_ms, + "candidate_ms": new_ms, + "speedup": old_ms / new_ms, + } + ), + flush=True, + ) + old_pipeline_ms = sequence_latency(old_ops) + new_pipeline_ms = sequence_latency(new_ops) + print( + json.dumps( + { + "kernel": "qsa_score_merge_emit", + "dtype": element_type, + "rows": rows, + "visible_blocks": visible, + "baseline_ms": old_pipeline_ms, + "candidate_ms": new_pipeline_ms, + "speedup": old_pipeline_ms / new_pipeline_ms, + } + ), + flush=True, + ) + checks += qsa_ragged_replay(old_kernel, new_kernel, params_dtype, element_type, random, baseline) + checks += qsa_ragged_replay( + old_kernel, paired_kernel, params_dtype, element_type, random, baseline, paired=True + ) + print(f"QSA exact-key checks passed: {checks}", flush=True) + + +def qsa_ragged_replay(old_kernel, new_kernel, params_dtype, element_type, random, baseline, paired=False): + rows, capacity, batches = 8, 4099, 1 if paired else 3 + params = numpy.zeros((), dtype=params_dtype) + for name, value in { + "batch_size": batches, + "total_tokens": rows, + "num_heads": 4, + "head_size": 128, + "compress_ratio": 4, + "state_capacity": capacity, + "capacity": 2051, + "block_topk": 512, + "scale": 0.0883883461356163, + }.items(): + params[name] = value + keys = key_data(random.standard_normal((batches, capacity, 128), dtype=numpy.float32), element_type) + query = cupy.asarray(random.standard_normal((rows, 4, 128), dtype=numpy.float32)) + cumulative = cupy.asarray([0, rows] if paired else [0, 0, 7, 8], dtype=cupy.int32) + past = cupy.zeros(batches, dtype=cupy.int32) + positions = cupy.zeros(rows, dtype=cupy.int64) + lengths = cupy.zeros((batches, 2), dtype=cupy.int32) + overflow = cupy.zeros(batches, dtype=cupy.int32) + tiles = (capacity + 7) // 8 + expected = cupy.empty((rows, tiles * 8), dtype=cupy.uint64) + actual = cupy.empty_like(expected) + common = (keys, query, cumulative, past, positions, lengths, overflow) + old_args = (*common, expected, numpy.int32(tiles), params[()]) + new_args = (*common, actual, numpy.int32(tiles), params[()]) + old_ops = [(old_kernel, qsa_score_grid(baseline, rows, tiles), (1024,), old_args)] + new_ops = [ + (new_kernel, (min(tiles, 128), (rows + 1) // 2 if paired else rows), (256,) if paired else (1024,), new_args) + ] + if paired: + old_merge_ops, old_result = merge_operations(compile_qsa_merge(baseline), expected, rows, tiles * 8) + new_merge_ops, new_result = merge_operations(compile_qsa_merge(QSA_SOURCE.read_text()), actual, rows, tiles * 8) + expected_indices = cupy.empty((rows, 2051), dtype=cupy.int32) + actual_indices = cupy.empty_like(expected_indices) + expected_counts = cupy.empty(rows, dtype=cupy.int32) + actual_counts = cupy.empty_like(expected_counts) + emit_common = (cumulative, past, positions, lengths, overflow) + old_ops += [ + *old_merge_ops, + ( + compile_qsa_emit(baseline), + (rows,), + (128,), + ( + old_result, + *emit_common, + numpy.int32(old_result.strides[0] // 8), + expected_indices, + expected_counts, + params[()], + ), + ), + ] + new_ops += [ + *new_merge_ops, + ( + compile_qsa_emit(QSA_SOURCE.read_text()), + (rows,), + (128,), + ( + new_result, + *emit_common, + numpy.int32(new_result.strides[0] // 8), + actual_indices, + actual_counts, + params[()], + ), + ), + ] + cupy.cuda.Stream.null.synchronize() + stream = cupy.cuda.Stream(non_blocking=True) + stream.begin_capture() + with stream: + launch_operations(old_ops) + launch_operations(new_ops) + graph = stream.end_capture() + checks = 0 + overflow_cases = ((0,), (1,)) if paired else ((0, 0, 0), (0, 1, 0), (0, 0, 1), (0, 1, 1)) + visibility_cases = [(visible, visible * 4) for visible in (0, 1, 9, 1025, capacity)] + if paired: + visibility_cases.append((1026, 4098)) + for visible, past_length in visibility_cases: + for overflow_values in overflow_cases: + lengths.set( + numpy.array([[visible, 0]] if paired else [[0, 0], [visible, 0], [visible // 2, 0]], dtype=numpy.int32) + ) + past.set(numpy.array([past_length] if paired else [0, past_length, (visible // 2) * 4], dtype=numpy.int32)) + overflow.set(numpy.array(overflow_values, dtype=numpy.int32)) + expected.fill(0xDEAD) + actual.fill(0xBEEF) + cupy.cuda.Stream.null.synchronize() + graph.launch(stream) + stream.synchronize() + cupy.testing.assert_array_equal(expected, actual) + if paired: + cupy.testing.assert_array_equal(expected_indices, actual_indices) + cupy.testing.assert_array_equal(expected_counts, actual_counts) + checks += 1 + return checks + + +def main(): + parser = argparse.ArgumentParser(description="Standalone NVRTC parity/timing probe; does not rebuild ORT.") + parser.add_argument("--baseline-stdin", action="store_true", required=True) + parser.add_argument("--check-only", action="store_true") + parser.add_argument("--sanitizer-case", action="store_true") + parser.add_argument("--kernel", choices=("qsa", "qsa_merge", "nvfp4", "nvfp4_reuse", "nvfp4_decode"), default="qsa") + args = parser.parse_args() + baseline = sys.stdin.read() + if args.kernel == "qsa": + run_qsa(baseline, args.check_only, args.sanitizer_case) + elif args.kernel == "qsa_merge": + run_qsa_merge(baseline, args.check_only, args.sanitizer_case) + elif args.kernel == "nvfp4_reuse": + run_nvfp4_reuse(baseline, args.check_only, args.sanitizer_case) + run_nvfp4_reuse(baseline, args.check_only, args.sanitizer_case, map_sources=False) + elif args.kernel == "nvfp4_decode": + run_nvfp4_decode(baseline, args.check_only, args.sanitizer_case) + else: + run_nvfp4(baseline, args.check_only, args.sanitizer_case) + + +def compile_nvfp4(source, element_type, fused, bias, wide, reuse=False, decode=False): + common = (ROOT / "contrib_ops/cuda/llm/moe_gemm/common.h").read_text() + details = (ROOT / "contrib_ops/cuda/llm/fpA_intB_gemv/details.h").read_text() + interleaved = (ROOT / "contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h").read_text() + uninterleave = device_function(interleaved.replace("CUTLASS_DEVICE", "__device__"), "fp4_e2m1x8_uninterleave") + activation = common[common.index("struct ActivationParams") : common.index(" ActivationParams()")] + activation += "};" + enum = common[common.index("enum class ActivationType") : common.index("};") + 2] + converter_start = details.index("template ") + converter_end = details.index("\n};", converter_start) + 3 + pair_start = source.index("template \n__device__") + pair_end = source.index("// NVFP4 schema weights", pair_start) + prefix = ( + """ +#include +#include +#include +#include +namespace std { +using cuda::std::is_same_v; +using cuda::std::conditional_t; +using cuda::std::numeric_limits; +} +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef long long int64_t; +namespace cutlass_kernels { +""" + + enum + + activation + + "\n}\nnamespace cutlass { namespace detail {\n" + + uninterleave + ) + prefix += "\n}}\nnamespace fiv {\n" + details[converter_start:converter_end] + "\n}\n" + if decode: + decode_start = source.index("template \n__device__ __forceinline__ float PackedLowToFloat") + decode_kernel = device_function(source, "MoeGemvFp4DecodeKernel", "__global__") + decode_end = source.index(decode_kernel, decode_start) + len(decode_kernel) + kernel_body = source[decode_start:decode_end] + else: + kernel_body = device_function(source, "MoeGemvFp4RawKPackedKernel", "__global__") + body = device_function(source, "DecodeE4M3Fn") + "\n" + source[pair_start:pair_end] + kernel_body + flags = f"{str(fused).lower()},{str(bias).lower()}" + if decode: + lanes = 4 if wide else 8 + min_blocks = 4 if wide else 2 + name = f"MoeGemvFp4DecodeKernel<{element_type},{flags},{lanes},2,2,{min_blocks}>" + else: + lanes = ",4" if wide else "" + if reuse: + lanes = ",8,true" + name = f"MoeGemvFp4RawKPackedKernel<{element_type},{flags}{lanes}>" + activation_dtype = numpy.dtype( + [ + ("activation_type", numpy.int32), + ("swiglu_alpha", numpy.uint64), + ("swiglu_beta", numpy.uint64), + ("swiglu_limit", numpy.uint64), + ("alpha", numpy.float32), + ("beta", numpy.float32), + ("limit", numpy.float32), + ("swiglu_fusion", numpy.int32), + ], + align=True, + ) + return compile_device_code(prefix + body, name), activation_dtype + + +def run_nvfp4_reuse(baseline, check_only, sanitizer_case=False, map_sources=True): + candidate = (ROOT / "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu").read_text() + random = numpy.random.default_rng(51) + checks = 0 + for element_type in ("half", "__nv_bfloat16"): + for rows in (70,) if sanitizer_case else (70, 80): + for group in (5,) if sanitizer_case else (1, 2, 3, 4, 5, 8): + for with_bias in (False, True): + old_kernel, activation_dtype = compile_nvfp4(baseline, element_type, True, with_bias, False) + new_kernel, _ = compile_nvfp4(candidate, element_type, True, with_bias, False, reuse=True) + experts, columns, reduction = 512, 1280, 2560 + tokens = rows // 10 if map_sources else rows + identifiers = numpy.arange(rows, dtype=numpy.int32) // group + counts = numpy.bincount(identifiers, minlength=experts) + offsets = cupy.asarray(numpy.concatenate(([0], numpy.cumsum(counts))).astype(numpy.int64)) + mapped_experts = cupy.asarray(identifiers) + mapped_rows = cupy.arange(rows, dtype=cupy.int32) % tokens if map_sources else numpy.uint64(0) + act = key_data( + random.standard_normal((tokens, reduction), dtype=numpy.float32) * 0.05, element_type + ) + weight = cupy.asarray( + random.integers(0, 256, (experts, columns, reduction // 2), dtype=numpy.uint8) + ) + scales = cupy.asarray( + random.integers(16, 64, (experts, columns, reduction // 16), dtype=numpy.uint8) + ) + globals_ = cupy.full(experts, 0.125, dtype=cupy.float32) + biases = key_data( + random.standard_normal((experts, columns), dtype=numpy.float32) * 0.01, element_type + ) + output_dtype = cupy.float16 if element_type == "half" else cupy.uint16 + expected = cupy.full((rows, columns // 2), 0xDEAD, dtype=output_dtype) + actual = cupy.full_like(expected, 0xBEEF) + activation = numpy.zeros((), dtype=activation_dtype) + activation["alpha"] = 1.0 + activation["limit"] = numpy.inf + head = (act, weight, scales, globals_, biases if with_bias else numpy.uint64(0)) + tail = ( + offsets, + mapped_experts, + numpy.int32(experts), + numpy.int64(columns * reduction // 2), + numpy.int64(columns * reduction // 16), + numpy.int32(columns), + numpy.int32(reduction), + activation[()], + mapped_rows, + numpy.int32(tokens), + ) + old_args = (*head, expected, *tail) + new_args = (*head, actual, *tail) + grid = (rows, columns // 16) + old_kernel(grid, (128,), old_args) + new_kernel(grid, (128,), new_args) + cupy.testing.assert_array_equal(expected, actual) + checks += 1 + if not check_only and not with_bias: + old_ms = graph_latency(old_kernel, grid, (128,), old_args) + new_ms = graph_latency(new_kernel, grid, (128,), new_args) + print( + json.dumps( + { + "kernel": "nvfp4_fc1_expert_reuse", + "dtype": element_type, + "expanded_rows": rows, + "rows_per_expert": group, + "mapped_sources": map_sources, + "baseline_ms": old_ms, + "candidate_ms": new_ms, + "speedup": old_ms / new_ms, + } + ), + flush=True, + ) + print(f"FC1 expert-reuse exact-output checks passed: {checks}, mapped sources: {map_sources}", flush=True) + + +def output_float(values, element_type): + result = cupy.asnumpy(values) + if element_type == "half": + return result.astype(numpy.float32) + return (result.astype(numpy.uint32) << 16).view(numpy.float32) + + +def run_nvfp4(baseline, check_only, sanitizer_case=False): + source_path = ROOT / "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu" + candidate = source_path.read_text() + random = numpy.random.default_rng(37) + wide = "int KLanes" in candidate + checks = 0 + for element_type in ("half", "__nv_bfloat16"): + for rows in (80,) if sanitizer_case else (10, 20, 70, 80): + shapes = ( + ((514, 528, True),) if sanitizer_case else ((1280, 2560, True), (2560, 640, False), (514, 528, True)) + ) + for columns, reduction, fused in shapes: + for with_bias in (False, True): + old_kernel, activation_dtype = compile_nvfp4(baseline, element_type, fused, with_bias, False) + use_wide = wide and reduction <= 1024 + new_kernel, _ = compile_nvfp4(candidate, element_type, fused, with_bias, use_wide) + experts = 512 + tokens = rows // 10 + act = key_data( + random.standard_normal((tokens, reduction), dtype=numpy.float32) * 0.05, element_type + ) + weight = cupy.asarray( + random.integers(0, 256, (experts, columns, reduction // 2), dtype=numpy.uint8) + ) + scales = cupy.asarray( + random.integers(16, 64, (experts, columns, reduction // 16), dtype=numpy.uint8) + ) + globals_ = cupy.full(experts, 0.125, dtype=cupy.float32) + biases = key_data( + random.standard_normal((experts, columns), dtype=numpy.float32) * 0.01, element_type + ) + mapped_experts = cupy.asarray(random.integers(0, experts, rows, dtype=numpy.int32)) + mapped_rows = cupy.asarray(random.integers(0, tokens, rows, dtype=numpy.int32)) + offsets = cupy.zeros(experts + 1, dtype=cupy.int64) + output_columns = columns // 2 if fused else columns + output_dtype = cupy.float16 if element_type == "half" else cupy.uint16 + expected = cupy.empty((rows, output_columns), dtype=output_dtype) + actual = cupy.empty_like(expected) + activation = numpy.zeros((), dtype=activation_dtype) + activation["alpha"] = 1.0 + activation["limit"] = numpy.inf + common = (act, weight, scales, globals_, biases if with_bias else numpy.uint64(0)) + tail = ( + offsets, + mapped_experts, + numpy.int32(experts), + numpy.int64(columns * reduction // 2), + numpy.int64(columns * reduction // 16), + numpy.int32(columns), + numpy.int32(reduction), + activation[()], + mapped_rows, + numpy.int32(tokens), + ) + old_args = (*common, expected, *tail) + new_args = (*common, actual, *tail) + old_grid = (rows, (columns + 15) // 16) + tile_columns = 32 if use_wide else 16 + new_grid = (rows, (columns + tile_columns - 1) // tile_columns) + old_kernel(old_grid, (128,), old_args) + new_kernel(new_grid, (128,), new_args) + numpy.testing.assert_allclose( + output_float(actual, element_type), + output_float(expected, element_type), + rtol=0.002 if element_type == "half" else 0.01, + atol=0.0005, + ) + checks += 1 + if not check_only and not with_bias and columns != 514: + old_ms = graph_latency(old_kernel, old_grid, (128,), old_args) + new_ms = graph_latency(new_kernel, new_grid, (128,), new_args) + print( + json.dumps( + { + "kernel": "nvfp4_fc1" if fused else "nvfp4_fc2", + "dtype": element_type, + "expanded_rows": rows, + "baseline_ms": old_ms, + "candidate_ms": new_ms, + "speedup": old_ms / new_ms, + } + ), + flush=True, + ) + print(f"NVFP4 parity checks passed: {checks}", flush=True) + + +def graph_latency_rotating(kernel, grid, block, step_args_list, shared_mem=0, iterations=300, replays=10): + cupy.cuda.Stream.null.synchronize() + stream = cupy.cuda.Stream(non_blocking=True) + num_steps = len(step_args_list) + with stream: + for i in range(min(10, num_steps)): + kernel(grid, block, step_args_list[i], shared_mem=shared_mem) + stream.synchronize() + stream.begin_capture() + with stream: + for i in range(iterations): + kernel(grid, block, step_args_list[i % num_steps], shared_mem=shared_mem) + graph = stream.end_capture() + start = cupy.cuda.Event() + stop = cupy.cuda.Event() + samples_us = [] + for _ in range(replays): + start.record(stream) + graph.launch(stream) + stop.record(stream) + stop.synchronize() + samples_us.append(cupy.cuda.get_elapsed_time(start, stop) * 1000.0 / iterations) + return ( + float(numpy.median(samples_us)), + float(numpy.min(samples_us)), + float(numpy.max(samples_us)), + float(numpy.std(samples_us)), + ) + + +def run_nvfp4_decode(baseline, check_only, sanitizer_case=False): + source_path = ROOT / "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu" + candidate = source_path.read_text() + random = numpy.random.default_rng(73) + checks = 0 + experts = 512 + routing_steps = 64 + for element_type in ("half", "__nv_bfloat16"): + for rows in (10,) if sanitizer_case else (10, 20): + for columns, reduction, fused in ((1280, 2560, True), (2560, 640, False)): + for with_bias in (False, True): + use_wide = reduction <= 1024 + old_kernel, activation_dtype = compile_nvfp4(baseline, element_type, fused, with_bias, use_wide) + new_kernel, _ = compile_nvfp4(candidate, element_type, fused, with_bias, use_wide, decode=True) + tokens = rows // 10 + act = key_data( + random.standard_normal((tokens, reduction), dtype=numpy.float32) * 0.05, element_type + ) + weight = cupy.asarray( + random.integers(0, 256, (experts, columns, reduction // 2), dtype=numpy.uint8) + ) + scales = cupy.asarray( + random.integers(16, 64, (experts, columns, reduction // 16), dtype=numpy.uint8) + ) + globals_ = cupy.full(experts, 0.09 / 448.0, dtype=cupy.float32) + biases = key_data( + random.standard_normal((experts, columns), dtype=numpy.float32) * 0.01, element_type + ) + # 64 disjoint routing steps striding across all 512 experts so consecutive launches in the + # 300-launch CUDA graph stream the full 800 MiB (FC1) / 400 MiB (FC2) expert bank rather than + # reusing L2-resident expert weights. + step_experts = [ + cupy.asarray((numpy.arange(rows, dtype=numpy.int32) + step * rows) % experts) + for step in range(routing_steps) + ] + mapped_rows = cupy.asarray(numpy.arange(rows, dtype=numpy.int32) % tokens) + offsets = cupy.zeros(experts + 1, dtype=cupy.int64) + output_columns = columns // 2 if fused else columns + output_dtype = cupy.float16 if element_type == "half" else cupy.uint16 + expected = cupy.full((rows, output_columns), 0xDEAD, dtype=output_dtype) + actual = cupy.full_like(expected, 0xBEEF) + activation = numpy.zeros((), dtype=activation_dtype) + activation["alpha"] = 1.0 + activation["limit"] = numpy.inf + common = (act, weight, scales, globals_, biases if with_bias else numpy.uint64(0)) + tail_suffix = ( + numpy.int32(experts), + numpy.int64(columns * reduction // 2), + numpy.int64(columns * reduction // 16), + numpy.int32(columns), + numpy.int32(reduction), + activation[()], + mapped_rows, + numpy.int32(tokens), + ) + old_step_args = [(*common, expected, offsets, exp, *tail_suffix) for exp in step_experts] + new_step_args = [(*common, actual, offsets, exp, *tail_suffix) for exp in step_experts] + old_tile_cols = 32 if use_wide else 16 + new_tile_cols = 64 if use_wide else 32 + old_grid = (rows, (columns + old_tile_cols - 1) // old_tile_cols) + new_grid = (rows, (columns + new_tile_cols - 1) // new_tile_cols) + decode_smem = (reduction // 32) * 36 * 4 + old_kernel(old_grid, (128,), old_step_args[0]) + new_kernel(new_grid, (64,), new_step_args[0], shared_mem=decode_smem) + cupy.testing.assert_array_equal(expected, actual) + checks += 1 + if not check_only and not with_bias: + old_med, old_min, old_max, old_std = graph_latency_rotating( + old_kernel, old_grid, (128,), old_step_args + ) + new_med, new_min, new_max, new_std = graph_latency_rotating( + new_kernel, new_grid, (64,), new_step_args, shared_mem=decode_smem + ) + print( + json.dumps( + { + "kernel": "nvfp4_decode_fc1" if fused else "nvfp4_decode_fc2", + "dtype": element_type, + "expanded_rows": rows, + "baseline_us": old_med, + "baseline_range_us": [old_min, old_max], + "baseline_std_us": old_std, + "candidate_us": new_med, + "candidate_range_us": [new_min, new_max], + "candidate_std_us": new_std, + "speedup": old_med / new_med, + } + ), + flush=True, + ) + print(f"NVFP4 decode bit-exact checks passed: {checks}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/test/python/transformers/test_benchmark_qwen38_mtp_kernels.py b/onnxruntime/test/python/transformers/test_benchmark_qwen38_mtp_kernels.py new file mode 100644 index 0000000000000..aa1984843c549 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_benchmark_qwen38_mtp_kernels.py @@ -0,0 +1,97 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import importlib.util +import sys +import unittest.mock +from pathlib import Path + + +def load_benchmark(): + path = Path(__file__).with_name("benchmark_qwen38_mtp_kernels.py") + spec = importlib.util.spec_from_file_location("benchmark_qwen38_mtp_kernels", path) + module = importlib.util.module_from_spec(spec) + with unittest.mock.patch.dict(sys.modules, {"cupy": unittest.mock.Mock()}): + spec.loader.exec_module(module) + return module + + +class TestBenchmarkCudaSource(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.benchmark = load_benchmark() + + def test_templated_device_function_preserves_declaration(self): + source = "template \n__device__ float decode(int x) { return x; }\nvoid host() {}" + actual = self.benchmark.device_function(source, "decode") + self.assertTrue(actual.startswith("template ")) + self.assertNotIn("void host", actual) + + def test_non_template_device_function_stays_non_template(self): + source = "template \nvoid unrelated() {}\n__device__ float decode(int x) { return x; }" + actual = self.benchmark.device_function(source, "decode") + self.assertTrue(actual.startswith("__device__")) + + def test_global_function_stops_before_host_helpers(self): + source = "template \n__global__ void kernel(T x) { if (x) { x++; } }\nbool host() { return true; }" + actual = self.benchmark.device_function(source, "kernel", "__global__") + self.assertTrue(actual.startswith("template ")) + self.assertNotIn("bool host", actual) + + def test_missing_function_is_an_error(self): + with self.assertRaisesRegex(ValueError, "Missing device function"): + self.benchmark.device_function("", "missing") + + def test_nvfp4_generated_source_excludes_host_inspection(self): + source = """ +template +__device__ float DecodeE4M3Fn(unsigned char code) { return code; } +template +__device__ void StorePair() {} +// NVFP4 schema weights +template +__global__ void MoeGemvFp4RawKPackedKernel() { if (NativeScale) { return; } } +bool Fp4MoeGemvSupportsNativeScale() { return cudaFuncGetAttributes(); } +bool Fp4MoeGemvKernelUsesNativeScale() { return false; } +template +__device__ __forceinline__ float PackedLowToFloat(uint32_t word) { return word; } +template +__global__ void MoeGemvFp4DecodeKernel() {} +template +void LaunchMoeGemvFp4RawNPacked() {} +""" + headers = { + "common.h": "enum class ActivationType { Identity };\nstruct ActivationParams {\n ActivationParams() {} };", + "details.h": "template \nstruct Converter {\n};", + "interleaved_numeric_conversion.h": "CUTLASS_DEVICE int fp4_e2m1x8_uninterleave(int value) { return value; }", + } + for element in ("half", "__nv_bfloat16"): + for fused, bias, wide, reuse, decode in ( + (False, False, True, False, False), + (True, True, False, False, False), + (True, False, False, True, False), + (True, False, False, False, True), + (False, True, True, False, True), + ): + with self.subTest(element=element, fused=fused, bias=bias, wide=wide, reuse=reuse, decode=decode): + with ( + unittest.mock.patch.object(Path, "read_text", lambda path: headers[path.name]), + unittest.mock.patch.object( + self.benchmark, "compile_device_code", return_value=object() + ) as compile_code, + ): + self.benchmark.compile_nvfp4(source, element, fused, bias, wide, reuse, decode=decode) + generated, _ = compile_code.call_args.args + self.assertIn("template ", generated) + if decode: + self.assertIn("__global__ void MoeGemvFp4DecodeKernel", generated) + else: + self.assertIn("__global__ void MoeGemvFp4RawKPackedKernel", generated) + self.assertNotIn("bool Fp4MoeGemvSupportsNativeScale", generated) + self.assertNotIn("bool Fp4MoeGemvKernelUsesNativeScale", generated) + self.assertNotIn("cudaFuncGetAttributes", generated) + + +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 95df52c3f6e61..670766f1c8fb0 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -30,8 +30,8 @@ EP_GRAPH_ASSIGNMENT_CONFIG_KEY = "session.record_ep_graph_assignment_info" NO_CUDNN_PLUGIN_TEST = os.getenv("ORT_TEST_CUDA_PLUGIN_NO_CUDNN", "").upper() in {"1", "ON", "TRUE", "YES"} requires_cudnn = unittest.skipIf(NO_CUDNN_PLUGIN_TEST, "test requires cuDNN-backed CUDA plugin kernels") -# Use the latest released ai.onnx opset so the model builders stay current as ONNX releases new opsets. -DEFAULT_ONNX_OPSET = max(v for (d, v) in helper.OP_SET_ID_VERSION_MAP if d == "ai.onnx") +# CUDA plugin runtime currently supports released ai.onnx opsets through 27. +DEFAULT_ONNX_OPSET = 27 def _make_released_opset_model(graph, producer_name="onnx-example"): diff --git a/onnxruntime/test/python/transformers/test_float16.py b/onnxruntime/test/python/transformers/test_float16.py index 6ba9eefa3d046..219348837f4ce 100644 --- a/onnxruntime/test/python/transformers/test_float16.py +++ b/onnxruntime/test/python/transformers/test_float16.py @@ -7,11 +7,15 @@ """Tests for float16 conversion (convert_float_to_float16).""" +import tempfile import unittest +from pathlib import Path +from unittest.mock import patch import numpy as np import onnx from onnx import TensorProto, helper, numpy_helper +from parameterized import parameterized from parity_utilities import find_transformers_source if find_transformers_source(): @@ -20,7 +24,7 @@ from onnxruntime.transformers.float16 import convert_float_to_float16 -def _make_resize_model_opset11(num_resize_nodes=2, use_empty_names=True): +def _make_resize_model_opset11(num_resize_nodes=2, use_empty_names=True, share_scales=False): """Create a minimal ONNX model with multiple Resize nodes (opset 11+). Resize opset 11+: inputs are [X, roi, scales, sizes]. @@ -33,7 +37,7 @@ def _make_resize_model_opset11(num_resize_nodes=2, use_empty_names=True): prev_output = "input" for idx in range(num_resize_nodes): roi_name = f"roi_{idx}" - scales_name = f"scales_{idx}" + scales_name = "scales" if share_scales else f"scales_{idx}" output_name = f"resize_out_{idx}" if idx < num_resize_nodes - 1 else "output" node = helper.make_node( @@ -49,10 +53,17 @@ def _make_resize_model_opset11(num_resize_nodes=2, use_empty_names=True): initializers = [] for idx in range(num_resize_nodes): roi = numpy_helper.from_array(np.array([], dtype=np.float32), name=f"roi_{idx}") - scales = numpy_helper.from_array(np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), name=f"scales_{idx}") - initializers.extend([roi, scales]) - - graph = helper.make_graph(nodes, "resize_test", [graph_input], [graph_output], initializer=initializers) + initializers.append(roi) + if not share_scales or idx == 0: + scales = numpy_helper.from_array( + np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), name="scales" if share_scales else f"scales_{idx}" + ) + initializers.append(scales) + + value_info = [helper.make_tensor_value_info("scales", TensorProto.FLOAT, [4])] if share_scales else None + graph = helper.make_graph( + nodes, "resize_test", [graph_input], [graph_output], initializer=initializers, value_info=value_info + ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 11)]) model = onnx.shape_inference.infer_shapes(model) return model @@ -128,6 +139,27 @@ def _make_blocked_node_model(num_nodes=2, use_empty_names=True): return model +def _make_mixed_float_bfloat16_model(): + """Create a model where one value feeds both a FLOAT and a BFLOAT16 cast.""" + data_input = helper.make_tensor_value_info("data", TensorProto.FLOAT, [1, 1, 2, 2]) + shared_input = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [4]) + resize_output = helper.make_tensor_value_info("resize_output", TensorProto.FLOAT, [1, 1, 4, 4]) + upsample_output = helper.make_tensor_value_info("upsample_output", TensorProto.FLOAT, [1, 1, 4, 4]) + nodes = [ + helper.make_node("Resize", ["data", "roi", "shared"], ["resize_output"], mode="nearest"), + helper.make_node("Upsample", ["data", "shared"], ["upsample_output"], mode="nearest"), + ] + roi = numpy_helper.from_array(np.array([], dtype=np.float32), name="roi") + graph = helper.make_graph( + nodes, + "mixed_float_bfloat16_test", + [data_input, shared_input], + [resize_output, upsample_output], + initializer=[roi], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 11)]) + + def _make_constant_of_shape_model(): """Create a model with a ConstantOfShape node that omits its optional value attribute.""" graph_input = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 4]) @@ -178,6 +210,17 @@ def test_resize_opset11_cast_naming_unique(self): len(cast_outputs), len(set(cast_outputs)), f"Duplicate Cast output names found: {cast_outputs}" ) + def test_resize_opset11_shared_scales_reuses_cast(self): + """A shared Resize scales input should use one Cast node for all consumers.""" + model = _make_resize_model_opset11(num_resize_nodes=2, use_empty_names=False, share_scales=True) + converted = convert_float_to_float16(model, keep_io_types=True) + + scales_casts = [n for n in converted.graph.node if n.op_type == "Cast" and n.input[0] == "scales"] + self.assertEqual(len(scales_casts), 1) + + resize_nodes = [n for n in converted.graph.node if n.op_type == "Resize"] + self.assertEqual([n.input[2] for n in resize_nodes], [scales_casts[0].output[0]] * len(resize_nodes)) + def test_resize_opset11_scales_initializer_stays_fp32(self): """Resize scales initializer (input index 2) should stay float32 after conversion. @@ -249,6 +292,19 @@ def test_blocked_node_cast_naming_unique(self): len(cast_outputs), len(set(cast_outputs)), f"Duplicate Cast output names found: {cast_outputs}" ) + def test_mixed_float_bfloat16_cast_naming_unique(self): + """A shared input for FLOAT and BFLOAT16 consumers should get distinct Cast names.""" + model = _make_mixed_float_bfloat16_model() + converted = convert_float_to_float16( + model, keep_io_types=False, use_bfloat16_as_blocked_nodes_dtype=True, disable_shape_infer=True + ) + + shared_casts = [n for n in converted.graph.node if n.op_type == "Cast" and n.input[0] == "shared"] + cast_types = {next(a.i for a in n.attribute if a.name == "to") for n in shared_casts} + self.assertEqual(cast_types, {TensorProto.FLOAT, TensorProto.BFLOAT16}) + self.assertEqual(len({n.name for n in shared_casts}), 2) + self.assertEqual(len({n.output[0] for n in shared_casts}), 2) + def test_resize_with_op_block_list(self): """When Resize is in op_block_list, Cast nodes should have unique names.""" model = _make_resize_model_opset11(num_resize_nodes=2, use_empty_names=True) @@ -295,5 +351,59 @@ def test_constant_of_shape_without_value_gets_fp16_value(self): np.testing.assert_array_equal(numpy_helper.to_array(value), np.array([0], dtype=np.float16)) +class TestFloat16ModelPathConversion(unittest.TestCase): + def _save_model(self, directory, external_data=False): + graph = helper.make_graph( + [ + helper.make_node("Mul", ["input", "weight"], ["product"]), + helper.make_node("Identity", ["product"], ["output"]), + ], + "model_path_test", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, [2])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [2])], + [numpy_helper.from_array(np.array([0.5, 2.0], dtype=np.float32), "weight")], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)], ir_version=8) + model_path = directory / "model with spaces.onnx" + onnx.save_model( + model, + model_path, + save_as_external_data=external_data, + all_tensors_to_one_file=True, + location=f"{directory.name}.weights.bin", + size_threshold=0, + ) + return model_path + + @parameterized.expand([(False,), (True,)]) + def test_model_path_matches_model_proto(self, external_data): + with tempfile.TemporaryDirectory() as temporary_directory: + directory = Path(temporary_directory) + model_path = self._save_model(directory, external_data) + original_files = {path.name: path.read_bytes() for path in directory.iterdir()} + for keep_io_types in (False, True, ["input"]): + with self.subTest(keep_io_types=keep_io_types): + expected = convert_float_to_float16(onnx.load(model_path), keep_io_types=keep_io_types) + for _ in range(2): + converted = convert_float_to_float16(str(model_path), keep_io_types=keep_io_types) + self.assertEqual(converted, expected) + self.assertEqual({path.name: path.read_bytes() for path in directory.iterdir()}, original_files) + + @parameterized.expand([("infer_shapes_path",), ("onnx.load",)]) + def test_model_path_removes_temporary_file_on_error(self, operation): + with tempfile.TemporaryDirectory() as temporary_directory: + directory = Path(temporary_directory) + model_path = self._save_model(directory) + original_files = {path.name: path.read_bytes() for path in directory.iterdir()} + error = RuntimeError("model path conversion failed") + with ( + patch(f"{convert_float_to_float16.__module__}.{operation}", side_effect=error), + self.assertRaises(RuntimeError) as raised, + ): + convert_float_to_float16(str(model_path)) + self.assertIs(raised.exception, error) + self.assertEqual({path.name: path.read_bytes() for path in directory.iterdir()}, original_files) + + if __name__ == "__main__": unittest.main() diff --git a/onnxruntime/test/python/transformers/test_matmul_nbits_unaligned4.py b/onnxruntime/test/python/transformers/test_matmul_nbits_unaligned4.py new file mode 100644 index 0000000000000..9f5e2ec08ad28 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_matmul_nbits_unaligned4.py @@ -0,0 +1,193 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""CUDA tests for the opt-in ORT_MATMULNBITS_UNALIGNED4_FAST_PATH (FP16 INT4, N % 8 == 4, M <= 4). + +The selector must be bit-identical to the existing aligned kernels evaluated on the same weights padded to the +next multiple of 8 output columns, and must not change any ineligible configuration. +""" + +import os +import subprocess +import sys +import unittest +from pathlib import Path + +import numpy as np +from cuda_plugin_ep_helper import CUDA_PLUGIN_EP_NAME, ensure_cuda_plugin_ep_registered +from onnx import TensorProto, helper, numpy_helper + +import onnxruntime as ort + +ENV_NAME = "ORT_MATMULNBITS_UNALIGNED4_FAST_PATH" +LOG_PREFIX = "[MatMulNBitsUnaligned4]" +BLOCK_SIZE = 32 + + +def _cuda_session_options(): + """Returns a plugin EP device, the CUDA provider name, or None if CUDA is unavailable.""" + if ensure_cuda_plugin_ep_registered(): + devices = [d for d in ort.get_ep_devices() if d.ep_name == CUDA_PLUGIN_EP_NAME] + if devices: + return devices[0] + if "CUDAExecutionProvider" in ort.get_available_providers(): + return "CUDAExecutionProvider" + return None + + +def _create_session(target, n, k, weights, scales, enabled, bias=False, zero_points=False, g_idx=False): + os.environ[ENV_NAME] = "1" if enabled else "0" + inputs = ["A", "B", "scales"] + initializers = [numpy_helper.from_array(weights, "B"), numpy_helper.from_array(scales, "scales")] + if zero_points or g_idx or bias: + inputs += ["zp" if zero_points else "", "g_idx" if g_idx else "", "bias" if bias else ""] + if zero_points: + initializers.append(numpy_helper.from_array(np.full((n, (k // BLOCK_SIZE + 1) // 2), 0x35, np.uint8), "zp")) + if g_idx: + num_blocks = k // BLOCK_SIZE + reversed_g_idx = (num_blocks - 1) - (np.arange(k, dtype=np.int32) // BLOCK_SIZE) + initializers.append(numpy_helper.from_array(reversed_g_idx, "g_idx")) + if bias: + initializers.append(numpy_helper.from_array(np.linspace(0.125, 1.5, n, dtype=np.float16), "bias")) + node = helper.make_node( + "MatMulNBits", + inputs, + ["Y"], + name=f"tail_N{n}_K{k}", + domain="com.microsoft", + N=n, + K=k, + bits=4, + block_size=BLOCK_SIZE, + ) + graph = helper.make_graph( + [node], + "unaligned4", + [helper.make_tensor_value_info("A", TensorProto.FLOAT16, ["M", k])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT16, ["M", n])], + initializers, + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = 10 + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + if isinstance(target, str): + return ort.InferenceSession( + model.SerializeToString(), options, providers=[(target, {"enable_cuda_graph": "0"})] + ) + options.add_provider_for_devices([target], {"enable_cuda_graph": "0"}) + return ort.InferenceSession(model.SerializeToString(), options) + + +def _pattern(name, rng, n, k): + if name == "constant": + return ( + np.full((n, k // BLOCK_SIZE, 16), 0x99, np.uint8), + np.full((n, k // BLOCK_SIZE), 0.03125, np.float16), + np.full((8, k), 0.125, np.float16), + ) + if name == "signed_zero": + return ( + rng.integers(0, 256, (n, k // BLOCK_SIZE, 16), dtype=np.uint8), + np.resize(np.array([0.0, -0.0, 0.0625, -0.0625], np.float16), (n, k // BLOCK_SIZE)), + np.resize(np.array([0.0, -0.0], np.float16), (8, k)), + ) + if name == "adversarial": + return ( + np.resize(np.array([0x00, 0x0F, 0xF0, 0xFF, 0x88, 0x17, 0xE9], np.uint8), (n, k // BLOCK_SIZE, 16)), + np.resize(np.array([0.0625, -0.0625, 2**-14, 2**-24, 0.0, -0.0], np.float16), (n, k // BLOCK_SIZE)), + np.resize(np.array([1, -1, 8, -8, 2**-14, -(2**-14), 2**-24, -0.0], np.float16), (8, k)), + ) + return ( + rng.integers(0, 256, (n, k // BLOCK_SIZE, 16), dtype=np.uint8), + rng.uniform(0.001, 0.02, (n, k // BLOCK_SIZE)).astype(np.float16), + rng.normal(0, 0.02, (8, k)).astype(np.float16), + ) + + +def _worker(): + target = _cuda_session_options() + if target is None: + print("SKIP") + return + for name in ("ORT_FPA_INTB_GEMM", "ORT_MATMULNBITS_DEQUANT_CACHE", "ORT_MATMULNBITS_FORCE_CHUNKED"): + os.environ[name] = "0" + + rng = np.random.default_rng(34812) + comparisons = 0 + for n in (4, 12, 324): + for k in (32, 320, 10240): + for pattern in ("constant", "signed_zero", "adversarial", "realish"): + weights, scales, a = _pattern(pattern, rng, n, k) + # Padded output rows are independent poison; the valid outputs must not depend on them. + padded_weights = np.concatenate([weights, np.full((4, k // BLOCK_SIZE, 16), 0xFE, np.uint8)]) + padded_scales = np.concatenate([scales, np.full((4, k // BLOCK_SIZE), -0.375, np.float16)]) + tail = _create_session(target, n, k, weights, scales, True) + reference = _create_session(target, n + 4, k, padded_weights, padded_scales, False) + for m in (1, 2, 3, 4): + feed = {"A": a[:m].copy()} + actual = tail.run(["Y"], feed)[0] + expected = reference.run(["Y"], feed)[0][:, :n].copy() + if actual.tobytes() != expected.tobytes(): + raise AssertionError(f"bit mismatch M={m} N={n} K={k} pattern={pattern}") + comparisons += 1 + del tail, reference + + n, k = 12, 320 + weights, scales, a_excl = _pattern("realish", rng, n, k) + scale_only_session = _create_session(target, n, k, weights, scales, False) + scale_only_m2 = scale_only_session.run(["Y"], {"A": a_excl[:2].copy()})[0].tobytes() + del scale_only_session + for label, kwargs, m in ( + ("M5", {}, 5), + ("M8", {}, 8), + ("M9", {}, 9), + ("bias", {"bias": True}, 2), + ("zero_points", {"zero_points": True}, 2), + ("g_idx", {"g_idx": True}, 2), + ): + feed_a = np.resize(a_excl, (m, k)).copy() + outputs = [] + for enabled in (False, True): + session = _create_session(target, n, k, weights, scales, enabled, **kwargs) + outputs.append(session.run(["Y"], {"A": feed_a})[0].tobytes()) + del session + if outputs[0] != outputs[1]: + raise AssertionError(f"{label}: selector changed an ineligible configuration") + if kwargs and outputs[0] == scale_only_m2: + raise AssertionError(f"{label}: exclusion input did not change output vs scale-only baseline") + print(f"COMPARISONS={comparisons}") + + +class TestMatMulNBitsUnaligned4(unittest.TestCase): + def test_tail_dispatch_matches_padded_aligned_kernels(self): + # Run in a subprocess so the dispatch receipts written to stderr can be inspected. + env = os.environ.copy() + package_root = str(Path(ort.__file__).resolve().parent.parent) + pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = os.pathsep.join(filter(None, [package_root, pythonpath])) + result = subprocess.run( + [sys.executable, str(Path(__file__).resolve()), "--worker"], + capture_output=True, + text=True, + check=False, + env=env, + ) + self.assertEqual(result.returncode, 0, result.stdout[-2000:] + result.stderr[-6000:]) + if "SKIP" in result.stdout: + self.skipTest("CUDA execution provider is not available") + self.assertIn("COMPARISONS=144", result.stdout) + lines = [line for line in result.stderr.splitlines() if line.startswith(LOG_PREFIX)] + self.assertTrue(any("branch=M1_SHARED_TAIL" in line for line in lines)) + self.assertTrue(any("branch=BATCHED_REGISTER_TAIL" in line for line in lines)) + # Only the 36 tail-enabled sessions may report a receipt, at most once per branch each. + self.assertLessEqual(len(lines), 72) + + +if __name__ == "__main__": + if sys.argv[1:] == ["--worker"]: + _worker() + else: + unittest.main() diff --git a/onnxruntime/test/python/transformers/test_moe_cuda.py b/onnxruntime/test/python/transformers/test_moe_cuda.py index 5c92b5e197e4f..11731f295fad6 100644 --- a/onnxruntime/test/python/transformers/test_moe_cuda.py +++ b/onnxruntime/test/python/transformers/test_moe_cuda.py @@ -1734,7 +1734,11 @@ def create_sparse_mixer_onnx_graph( ], ) - return helper.make_model(graph, producer_name="MoE_Model") + return helper.make_model( + graph, + producer_name="MoE_Model", + opset_imports=[helper.make_opsetid("", 27), helper.make_opsetid("com.microsoft", 1)], + ) class TestSparseMixer(unittest.TestCase): @@ -1777,6 +1781,8 @@ def test_sparse_mixer_functional(self, onnx_dtype): fc2_bias, onnx_dtype, ) + opset_imports = {opset_import.domain: opset_import.version for opset_import in onnx_model.opset_import} + self.assertEqual(opset_imports, {"": 27, "com.microsoft": 1}) sess_options = onnxruntime.SessionOptions() sess = onnxruntime.InferenceSession(onnx_model.SerializeToString(), sess_options, providers=get_ort_provider()) diff --git a/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py index 751ef131ef0b9..b9dd703f3b6af 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py @@ -412,7 +412,9 @@ def weight(name, expert): for expert in experts.unique().tolist(): rows, slots = torch.where(experts == expert) - projection = (x[rows].float() @ weight("fc1", expert).T).to(x.dtype).float() + projection = x[rows].float() @ weight("fc1", expert).T + if fusion or "fc3_weights" in tensors: + projection = projection.to(x.dtype).float() if "fc1_bias" in tensors: projection += tensors["fc1_bias"][expert].float() if fusion: @@ -649,6 +651,44 @@ def test_block_fp8_fused_without_weight_scratch(capfd, monkeypatch, tokens, path assert "DequantWeightBytes=0" in log +@pytest.mark.skipif( + not torch.cuda.is_available() or not has_fp8_qmoe or torch.cuda.get_device_capability()[0] < 8, + reason="SM80+ CUDA FP8 QMoE required", +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("tokens,path", [(1, "fp8_fused_gemv"), (32, "fp8_fused_gemm")]) +def test_block_fp8_fused_silu_accumulator_precision(capfd, monkeypatch, dtype, tokens, path): + tensors = TestQMoEBlockFP8._inputs(hidden=16, inter=16, experts=1, tokens=tokens, block=16, dtype=dtype) + del tensors["fc3_weights"], tensors["fc3_scales"] + tensors["input"].zero_() + tensors["fc1_scales"].fill_(1) + tensors["fc2_scales"].fill_(1) + weights = torch.zeros(1, 16, 16, device=device) + tensors["fc1_bias"] = torch.zeros(1, 16, device=device, dtype=dtype) + if dtype == torch.float16: + tensors["input"][:, 0] = -4096 + weights[:, :, 0] = 32 + else: + tensors["input"][:, :3] = torch.tensor([128, 127, 0.5], device=device, dtype=dtype) + weights[:, :, :3] = 1 + tensors["fc1_bias"].fill_(-255) + tensors["fc1_weights"] = weights.to(torch.float8_e4m3fn).view(torch.uint8) + tensors["fc2_weights"] = torch.eye(16, device=device).unsqueeze(0).to(torch.float8_e4m3fn).view(torch.uint8) + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + monkeypatch.setenv("ORT_ENABLE_FP8_FUSED", "0") + dense = TestQMoEBlockFP8._execute(tensors, block=16, top_k=1) + capfd.readouterr() + monkeypatch.setenv("ORT_ENABLE_FP8_FUSED", "1") + actual = TestQMoEBlockFP8._execute(tensors, block=16, top_k=1) + log = capfd.readouterr().out + assert path in log + assert "DequantWeightBytes=0" in log + expected = TestQMoEBlockFP8._reference(tensors, 16, 0, top_k=1) + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(actual, dense, atol=0, rtol=0) + + @pytest.mark.skipif( not torch.cuda.is_available() or not has_fp8_qmoe or torch.cuda.get_device_capability()[0] < 8, reason="SM80+ CUDA FP8 QMoE required", diff --git a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py index 9551359248e50..cc80c4306f92a 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py @@ -18,6 +18,8 @@ # SM-agnostic; the tests require SM80+ / CUDA / an ENABLE_FP4 + USE_FP4_QMOE build. # -------------------------------------------------------------------------- +import base64 +import json import os import re import subprocess @@ -83,7 +85,7 @@ def fp4_e2m1_quantize(values): return quantized.reshape(values.shape), codes.reshape(values.shape) -def quantize_weight_to_nvfp4(weight, block_size=NVFP4_BLOCK_SIZE): +def quantize_weight_to_nvfp4(weight, block_size=NVFP4_BLOCK_SIZE, global_scale=None): """ Quantize a per-expert weight matrix [N, K] to NVFP4. @@ -103,19 +105,23 @@ def quantize_weight_to_nvfp4(weight, block_size=NVFP4_BLOCK_SIZE): w = weight.float() num_blocks = k // block_size - blocks = w.reshape(n, num_blocks, block_size) + if global_scale is not None: + w = w.clone() + w[0, :block_size] *= (abs(global_scale) * FP4_MAX * E4M3_MAX) / w[0, :block_size].abs().amax() + global_scale = float(global_scale) + else: + global_amax = w.abs().amax().item() + # Per-expert global scale (weight_scale_2). Chosen so that per-block E4M3 + # scales land within [0, E4M3_MAX]. + global_scale = global_amax / (FP4_MAX * E4M3_MAX) + if global_scale <= 0: + global_scale = 1.0 + blocks = w.reshape(n, num_blocks, block_size) block_amax = blocks.abs().amax(dim=-1) # [N, num_blocks] - global_amax = w.abs().amax().item() - - # Per-expert global scale (weight_scale_2). Chosen so that per-block E4M3 - # scales land within [0, E4M3_MAX]. - global_scale = global_amax / (FP4_MAX * E4M3_MAX) - if global_scale <= 0: - global_scale = 1.0 # Ideal per-block scale before E4M3 quantization. - block_scale_f = block_amax / FP4_MAX / global_scale # [N, num_blocks], all <= E4M3_MAX + block_scale_f = (block_amax / FP4_MAX / abs(global_scale)).clamp(max=E4M3_MAX) # [N, num_blocks], all <= E4M3_MAX block_scale_e4m3 = block_scale_f.to(torch.float8_e4m3fn) block_scale_deq = block_scale_e4m3.float() # DecodeE4M3(byte) @@ -334,6 +340,9 @@ def _run_nvfp4_moe_test( offline_prepacked=False, enable_cuda_graph=False, weights_prepacked=None, + global_scale=None, + fc1_global_scale=None, + fc2_global_scale=None, swiglu_params=(1.702, 1.0, 7.0), ): self._skip_if_no_fp4() @@ -351,17 +360,19 @@ def _run_nvfp4_moe_test( fc1_packed, fc1_bs, fc1_gs, fc1_deq = [], [], [], [] fc2_packed, fc2_bs, fc2_gs, fc2_deq = [], [], [], [] + fc1_scale_override = fc1_global_scale if fc1_global_scale is not None else global_scale + fc2_scale_override = fc2_global_scale if fc2_global_scale is not None else global_scale for _ in range(num_experts): w1 = torch.randn(fc1_n, fc1_k, device=device) * 0.1 - p1, b1, g1, d1 = quantize_weight_to_nvfp4(w1, block_size) + p1, b1, g1, d1 = quantize_weight_to_nvfp4(w1, block_size, global_scale=fc1_scale_override) fc1_packed.append(p1) fc1_bs.append(b1) fc1_gs.append(torch.tensor(g1, dtype=torch.float32)) fc1_deq.append(d1) w2 = torch.randn(fc2_n, fc2_k, device=device) * 0.1 - p2, b2, g2, d2 = quantize_weight_to_nvfp4(w2, block_size) + p2, b2, g2, d2 = quantize_weight_to_nvfp4(w2, block_size, global_scale=fc2_scale_override) fc2_packed.append(p2) fc2_bs.append(b2) fc2_gs.append(torch.tensor(g2, dtype=torch.float32)) @@ -929,18 +940,6 @@ def test_nvfp4_fp16_gemv_disabled_swiglu(self): gemv_mode="0", ) - def test_nvfp4_fp16_prepacking_disabled_uses_raw_fallback(self): - self._run_nvfp4_moe_test( - hidden_size=64, - inter_size=64, - num_experts=4, - top_k=2, - num_tokens=4, - onnx_dtype=TensorProto.FLOAT16, - gemv_mode="1", - disable_prepacking=True, - ) - @parameterized.expand( [ (onnx_dtype, disable_prepacking) @@ -1015,6 +1014,32 @@ def test_nvfp4_offline_row_major_gemv_partial_k_blocks(self, onnx_dtype, num_tok torch.testing.assert_close(packed, legacy, atol=tolerance, rtol=tolerance) torch.testing.assert_close(packed, fallback, atol=tolerance, rtol=tolerance) + @parameterized.expand( + [ + (onnx_dtype, disable_prepacking) + for onnx_dtype in (TensorProto.FLOAT16, TensorProto.BFLOAT16) + for disable_prepacking in (False, True) + ] + ) + def test_nvfp4_offline_row_major_above_gemv_limit(self, onnx_dtype, disable_prepacking): + shape = dict( + hidden_size=512, + inter_size=512, + num_experts=16, + top_k=10, + num_tokens=9, + onnx_dtype=onnx_dtype, + use_swiglu=True, + gemv_mode="1", + disable_prepacking=disable_prepacking, + use_bias=True, + input_scale=0.1, + ) + packed = self._run_nvfp4_moe_test(**shape, offline_prepacked=True) + legacy = self._run_nvfp4_moe_test(**shape) + tolerance = 0.03 if onnx_dtype == TensorProto.BFLOAT16 else 0.005 + torch.testing.assert_close(packed, legacy, atol=tolerance, rtol=tolerance) + @parameterized.expand( [ (onnx_dtype, disable_prepacking) @@ -1052,6 +1077,307 @@ def test_nvfp4_offline_row_major_cuda_graph(self, onnx_dtype): enable_cuda_graph=True, ) + # Decode GEMV parity cases (row-major weights_prepacked=1, fused SwiGLU FC1) as + # (default_fc1_decode, default_fc2_decode, optin_fc1_decode, optin_fc2_decode, fc1_scale_ok, fc2_scale_ok, shape). + # - Cases 1-2: Qwen3.8 shapes (FC1 N=1280 K=2560 -> 8 K lanes, FC2 N=2560 K=640 -> 4 K lanes) at 10/20 expanded + # rows, inside the default SM90 envelope. + # - Case 3: Qwen3.8 shape at 60 expanded rows; declined by default (ORT_FP4_GEMV_DECODE_KERNEL=1 caps expanded + # rows at 20) and enabled by opt-in ORT_FP4_GEMV_DECODE_KERNEL=2 (up to 80 expanded rows). + # - Case 4: Partial last K tile (98 and 34 blocks) at 6 expanded rows. + # - Case 5: Controlled global_scale just below 3.98 / 448 (table-free decode branch, eligible by default). + # - Case 6: Mixed per-projection scale eligibility (FC1 below 3.98 / 448, FC2 above); default mode runs + # MoeGemvFp4DecodeKernel for FC1 and keeps MoeGemvFp4RawKPackedKernel for FC2. + # - Case 7: Controlled global_scale above 3.98 / 448; default mode keeps MoeGemvFp4RawKPackedKernel for both + # projections to avoid the slower generic decoder, while opt-in mode (ORT_FP4_GEMV_DECODE_KERNEL=2) exercises + # the generic Fp4I2FConverter branch inside MoeGemvFp4DecodeKernel. + # - Case 8: K % 32 != 0, declined in all modes in favor of MoeGemvFp4RawKPackedKernel. + _DECODE_GEMV_CASES = ( + ( + True, + True, + True, + True, + True, + True, + dict(hidden_size=2560, inter_size=640, num_experts=16, top_k=10, num_tokens=1, use_bias=False), + ), + ( + True, + True, + True, + True, + True, + True, + dict(hidden_size=2560, inter_size=640, num_experts=16, top_k=10, num_tokens=2, use_bias=True), + ), + ( + False, + False, + True, + True, + True, + True, + dict(hidden_size=2560, inter_size=640, num_experts=16, top_k=10, num_tokens=6, use_bias=False), + ), + ( + True, + True, + True, + True, + True, + True, + dict(hidden_size=1568, inter_size=544, num_experts=12, top_k=6, num_tokens=1, use_bias=True), + ), + ( + True, + True, + True, + True, + True, + True, + dict( + hidden_size=1568, + inter_size=544, + num_experts=12, + top_k=6, + num_tokens=1, + use_bias=True, + global_scale=3.90 / E4M3_MAX, + ), + ), + ( + True, + False, + True, + True, + True, + False, + dict( + hidden_size=1568, + inter_size=544, + num_experts=12, + top_k=6, + num_tokens=1, + use_bias=True, + fc1_global_scale=3.90 / E4M3_MAX, + fc2_global_scale=4.10 / E4M3_MAX, + ), + ), + ( + False, + False, + True, + True, + False, + False, + dict( + hidden_size=1568, + inter_size=544, + num_experts=12, + top_k=6, + num_tokens=1, + use_bias=True, + global_scale=4.10 / E4M3_MAX, + ), + ), + ( + False, + False, + False, + False, + True, + True, + dict(hidden_size=528, inter_size=656, num_experts=16, top_k=10, num_tokens=6, use_bias=True), + ), + ) + + _DECODE_GEMV_RUNNER_SCRIPT = """ +import base64, json, sys +import torch +from torch.autograd import DeviceType +from torch.profiler import ProfilerActivity, profile +from test_qmoe_nvfp4_cuda import TestQMoENVFP4, device + +test = TestQMoENVFP4() +onnx_dtype = int(sys.argv[1]) +with profile(activities=[ProfilerActivity.CUDA]) as probe_prof: + _ = torch.zeros(16, device=device) + 1 + torch.cuda.synchronize() +cupti_available = any(e.device_type == DeviceType.CUDA for e in probe_prof.events()) +print("DECODE_GEMV_CUPTI", "1" if cupti_available else "0", sep="\\t", flush=True) +for idx, case in enumerate(json.loads(sys.argv[2])): + print(f"DECODE_GEMV_CASE_BEGIN\\t{idx}", flush=True) + with profile(activities=[ProfilerActivity.CUDA]) as prof: + out = test._run_nvfp4_moe_test(onnx_dtype=onnx_dtype, gemv_mode="1", offline_prepacked=True, **case) + names = sorted({e.name for e in prof.events() if "MoeGemvFp4" in e.name}) + payload = base64.b64encode(out.view(torch.int16).cpu().numpy().tobytes()).decode() + print("DECODE_GEMV_RESULT", str(idx), json.dumps(names), payload, sep="\\t", flush=True) +""" + + def _run_decode_gemv_subprocess(self, onnx_dtype, cases, decode_mode, native_scale=0): + env = dict(os.environ) + env["ORT_FP4_GEMV_DECODE_KERNEL"] = str(decode_mode) + env["ORT_NVFP4_GEMV_NATIVE_SCALE"] = str(native_scale) + env["ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO"] = "1" + env.pop("ORT_NVFP4_FC1_EXPERT_REUSE", None) + proc = subprocess.run( + [sys.executable, "-c", self._DECODE_GEMV_RUNNER_SCRIPT, str(int(onnx_dtype)), json.dumps(cases)], + cwd=os.path.dirname(os.path.abspath(__file__)), + env=env, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual( + proc.returncode, + 0, + f"ORT_FP4_GEMV_DECODE_KERNEL={decode_mode} ORT_NVFP4_GEMV_NATIVE_SCALE={native_scale} run failed:\n" + f"{proc.stdout}\n{proc.stderr}", + ) + cupti_available = False + current_idx = None + case_diags = {} + results = {} + for line in proc.stdout.splitlines(): + if line.startswith("DECODE_GEMV_CUPTI\t"): + cupti_available = line.split("\t")[1] == "1" + elif line.startswith("DECODE_GEMV_CASE_BEGIN\t"): + current_idx = int(line.split("\t")[1]) + elif line.startswith("QMoE NVFP4 DecodeImageAvailable=") and current_idx is not None: + fields = dict(item.split("=", 1) for item in line.split()[2:]) + case_diags[current_idx] = fields + elif line.startswith("DECODE_GEMV_RESULT\t"): + _, idx_str, names_json, payload = line.split("\t") + results[int(idx_str)] = (json.loads(names_json), payload) + self.assertEqual(len(results), len(cases), proc.stdout) + self.assertEqual(len(case_diags), len(cases), proc.stdout) + return cupti_available, [(case_diags[i], results[i][0], results[i][1]) for i in range(len(cases))] + + def _assert_fc1_fc2_gemv_dispatch( + self, + case, + diag, + kernels, + cupti_available, + expected_fc1_scale_ok, + expected_fc2_scale_ok, + expected_fc1_decode, + expected_fc2_decode, + ): + self.assertEqual(diag.get("FC1ScaleEligible"), "1" if expected_fc1_scale_ok else "0", (case, diag)) + self.assertEqual(diag.get("FC2ScaleEligible"), "1" if expected_fc2_scale_ok else "0", (case, diag)) + self.assertEqual(diag.get("FC1Kernel"), "decode" if expected_fc1_decode else "kpacked", (case, diag)) + self.assertEqual(diag.get("FC2Kernel"), "decode" if expected_fc2_decode else "kpacked", (case, diag)) + if not cupti_available: + return + self.assertTrue(kernels, f"CUPTI is available but no MoeGemvFp4 kernel events were recorded for {case}") + fc1_decode = any(re.search(r"MoeGemvFp4DecodeKernel<[^,>]+,\s*true\b", name) for name in kernels) + fc1_kpacked = any(re.search(r"MoeGemvFp4RawKPackedKernel<[^,>]+,\s*true\b", name) for name in kernels) + fc2_decode = any(re.search(r"MoeGemvFp4DecodeKernel<[^,>]+,\s*false\b", name) for name in kernels) + fc2_kpacked = any(re.search(r"MoeGemvFp4RawKPackedKernel<[^,>]+,\s*false\b", name) for name in kernels) + self.assertEqual(fc1_decode, expected_fc1_decode, (case, kernels)) + self.assertEqual(fc1_kpacked, not expected_fc1_decode, (case, kernels)) + self.assertEqual(fc2_decode, expected_fc2_decode, (case, kernels)) + self.assertEqual(fc2_kpacked, not expected_fc2_decode, (case, kernels)) + + @parameterized.expand([(TensorProto.FLOAT16,), (TensorProto.BFLOAT16,)]) + def test_nvfp4_offline_row_major_decode_gemv_bit_exact(self, onnx_dtype): + # The decode GEMV reproduces MoeGemvFp4RawKPackedKernel's summation order, so the QMoE output must be + # bit-identical across ORT_FP4_GEMV_DECODE_KERNEL=0 (opt-out), 1 (default measured SM90 envelope with + # per-projection scale eligibility), and 2 (opt-in full SM90+ envelope including generic scale decode). + self._skip_if_no_fp4() + cases = [{"use_swiglu": True, "input_scale": 0.1, **spec[-1]} for spec in self._DECODE_GEMV_CASES] + optout_cupti, optout_runs = self._run_decode_gemv_subprocess(onnx_dtype, cases, decode_mode=0) + default_cupti, default_runs = self._run_decode_gemv_subprocess(onnx_dtype, cases, decode_mode=1) + native_cupti, native_runs = self._run_decode_gemv_subprocess(onnx_dtype, cases, decode_mode=1, native_scale=1) + optin_cupti, optin_runs = self._run_decode_gemv_subprocess(onnx_dtype, cases, decode_mode=2) + + sm = _cuda_sm() + for ( + (def_fc1, def_fc2, optin_fc1, optin_fc2, fc1_scale_ok, fc2_scale_ok, _), + case, + (optout_diag, optout_kernels, optout_payload), + (default_diag, default_kernels, default_payload), + (native_diag, native_kernels, native_payload), + (optin_diag, optin_kernels, optin_payload), + ) in zip(self._DECODE_GEMV_CASES, cases, optout_runs, default_runs, native_runs, optin_runs, strict=True): + decode_image_available = default_diag.get("DecodeImageAvailable") == "1" + self.assertEqual(optout_diag.get("DecodeImageAvailable") == "1", decode_image_available, case) + self.assertEqual(native_diag.get("DecodeImageAvailable") == "1", decode_image_available, case) + self.assertEqual(optin_diag.get("DecodeImageAvailable") == "1", decode_image_available, case) + if sm < 90: + self.assertFalse(decode_image_available, (case, default_diag)) + + self._assert_fc1_fc2_gemv_dispatch( + case, optout_diag, optout_kernels, optout_cupti, fc1_scale_ok, fc2_scale_ok, False, False + ) + can_decode_default = sm == 90 and decode_image_available + self._assert_fc1_fc2_gemv_dispatch( + case, + default_diag, + default_kernels, + default_cupti, + fc1_scale_ok, + fc2_scale_ok, + def_fc1 and can_decode_default, + def_fc2 and can_decode_default, + ) + native_fc2_decode = ( + def_fc2 + and can_decode_default + and not (onnx_dtype == TensorProto.FLOAT16 and case["inter_size"] <= 1024) + ) + self._assert_fc1_fc2_gemv_dispatch( + case, + native_diag, + native_kernels, + native_cupti, + fc1_scale_ok, + fc2_scale_ok, + def_fc1 and can_decode_default, + native_fc2_decode, + ) + can_decode_optin = sm >= 90 and decode_image_available + self._assert_fc1_fc2_gemv_dispatch( + case, + optin_diag, + optin_kernels, + optin_cupti, + fc1_scale_ok, + fc2_scale_ok, + optin_fc1 and can_decode_optin, + optin_fc2 and can_decode_optin, + ) + + expected = torch.frombuffer(bytearray(base64.b64decode(optout_payload)), dtype=torch.int16) + default_out = torch.frombuffer(bytearray(base64.b64decode(default_payload)), dtype=torch.int16) + native_out = torch.frombuffer(bytearray(base64.b64decode(native_payload)), dtype=torch.int16) + optin_out = torch.frombuffer(bytearray(base64.b64decode(optin_payload)), dtype=torch.int16) + self.assertTrue( + torch.equal(default_out, expected), + f"default decode GEMV output differs bitwise from the KPacked kernel for {case}", + ) + self.assertTrue( + torch.equal(native_out, expected), + f"combined native-scale decode GEMV output differs bitwise from the KPacked kernel for {case}", + ) + self.assertTrue( + torch.equal(optin_out, expected), + f"opt-in decode GEMV output differs bitwise from the KPacked kernel for {case}", + ) + + def test_nvfp4_fp16_prepacking_disabled_uses_raw_fallback(self): + self._run_nvfp4_moe_test( + hidden_size=64, + inter_size=64, + num_experts=4, + top_k=2, + num_tokens=4, + onnx_dtype=TensorProto.FLOAT16, + gemv_mode="1", + disable_prepacking=True, + ) + @parameterized.expand( [ (TensorProto.FLOAT16, 2), @@ -1379,292 +1705,5 @@ def test_nvfp4_kpacked_mtp_qwen38_official_shape(onnx_dtype, num_tokens, capfd, assert "DequantWeightBytes=" not in diagnostics -def _enable_nvfp4_packed_prefill(monkeypatch): - monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "1") - monkeypatch.setenv("ORT_ENABLE_NVFP4_CUTLASS_GEMM", "0") - - -@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") -@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) -@pytest.mark.parametrize("num_tokens", [17, 65, 257]) -@pytest.mark.parametrize("hidden_size,inter_size", [(128, 64), (64, 128)]) -def test_nvfp4_packed_prefill_reference(onnx_dtype, num_tokens, hidden_size, inter_size, capfd, monkeypatch): - _enable_nvfp4_packed_prefill(monkeypatch) - monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") - shape = dict( - hidden_size=hidden_size, - inter_size=inter_size, - num_experts=16, - top_k=10, - num_tokens=num_tokens, - onnx_dtype=onnx_dtype, - use_swiglu=True, - use_bias=True, - gemv_mode="0", - offline_prepacked=True, - input_scale=0.1, - enable_cuda_graph=True, - ) - packed = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) - assert "Route=packed_nvfp4_prefill" in capfd.readouterr().out - monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "0") - dense = TestQMoENVFP4()._run_nvfp4_moe_test(**{**shape, "enable_cuda_graph": False}) - tolerance = 0.03 if onnx_dtype == TensorProto.BFLOAT16 else 0.005 - torch.testing.assert_close(packed, dense, atol=tolerance, rtol=tolerance) - - -@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") -@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) -def test_nvfp4_packed_prefill_qwen_dimensions(onnx_dtype, capfd, monkeypatch): - _enable_nvfp4_packed_prefill(monkeypatch) - monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") - shape = dict( - hidden_size=2560, - inter_size=640, - num_experts=16, - top_k=10, - num_tokens=65, - onnx_dtype=onnx_dtype, - use_swiglu=True, - use_bias=True, - gemv_mode="0", - offline_prepacked=True, - input_scale=0.05, - ) - packed = TestQMoENVFP4()._run_nvfp4_moe_test(**shape, enable_cuda_graph=True) - assert "Route=packed_nvfp4_prefill" in capfd.readouterr().out - monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "0") - dense = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) - tolerance = 0.03 if onnx_dtype == TensorProto.BFLOAT16 else 0.005 - torch.testing.assert_close(packed, dense, atol=tolerance, rtol=tolerance) - - -@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") -@pytest.mark.parametrize("fused", ["0", "1"]) -@pytest.mark.parametrize("reuse_scratch", ["0", "1"]) -@pytest.mark.parametrize("tiled_down", ["0", "1"]) -def test_nvfp4_packed_prefill_empty_experts(fused, reuse_scratch, tiled_down, capfd, monkeypatch): - _enable_nvfp4_packed_prefill(monkeypatch) - monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", fused) - monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", reuse_scratch) - monkeypatch.setenv("ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT", tiled_down) - monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") - logits = numpy.tile(numpy.arange(16, dtype=numpy.float32), (65, 1)) - shape = dict( - hidden_size=64, - inter_size=64, - num_experts=16, - top_k=2, - num_tokens=65, - onnx_dtype=TensorProto.FLOAT16, - use_swiglu=True, - use_bias=True, - gemv_mode="0", - offline_prepacked=True, - input_scale=0.1, - router_logits_override=logits, - row_tile_size=17, - enable_cuda_graph=True, - ) - packed = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) - assert "Route=packed_nvfp4_prefill" in capfd.readouterr().out - monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "0") - dense = TestQMoENVFP4()._run_nvfp4_moe_test(**{**shape, "enable_cuda_graph": False}) - torch.testing.assert_close(packed, dense, atol=0.005, rtol=0.005) - - -@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") -@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) -@pytest.mark.parametrize("num_tokens", [17, 65]) -@pytest.mark.parametrize("hidden_size,inter_size", [(128, 64), (2560, 640)]) -@pytest.mark.parametrize("use_bias", [False, True]) -@pytest.mark.parametrize( - "swiglu_params", - [ - (1.702, 1.0, 7.0), - (1.0, 0.0, float("inf")), - (0.5, 0.2, 0.1), - ], -) -def test_nvfp4_packed_prefill_fused_swiglu( - onnx_dtype, num_tokens, hidden_size, inter_size, use_bias, swiglu_params, capfd, monkeypatch -): - _enable_nvfp4_packed_prefill(monkeypatch) - monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") - shape = dict( - hidden_size=hidden_size, - inter_size=inter_size, - num_experts=16, - top_k=10, - num_tokens=num_tokens, - onnx_dtype=onnx_dtype, - use_swiglu=True, - use_bias=use_bias, - gemv_mode="0", - offline_prepacked=True, - input_scale=0.05, - enable_cuda_graph=True, - swiglu_params=swiglu_params, - ) - monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", "0") - unfused = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) - unfused_diagnostics = capfd.readouterr().out - assert "Route=packed_nvfp4_prefill" in unfused_diagnostics - monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", "1") - fused = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) - fused_diagnostics = capfd.readouterr().out - assert "Route=packed_nvfp4_prefill" in fused_diagnostics - unfused_workspace = int(re.search(r"RunnerWorkspaceBytes=(\d+)", unfused_diagnostics).group(1)) - fused_workspace = int(re.search(r"RunnerWorkspaceBytes=(\d+)", fused_diagnostics).group(1)) - gate_up_bytes = num_tokens * 10 * inter_size * 2 * 2 - assert unfused_workspace - fused_workspace == (gate_up_bytes + 255) // 256 * 256 - torch.testing.assert_close(fused, unfused, atol=0, rtol=0) - - -@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") -@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) -@pytest.mark.parametrize("fused", ["0", "1"]) -@pytest.mark.parametrize("row_tile_size", [0, 17]) -@pytest.mark.parametrize( - "hidden_size,inter_size,num_experts,top_k", - [ - (128, 64, 16, 10), - (2560, 640, 16, 10), - (64, 64, 512, 1), - (2560, 640, 512, 10), - ], -) -def test_nvfp4_packed_prefill_reuse_routing_scratch( - onnx_dtype, fused, row_tile_size, hidden_size, inter_size, num_experts, top_k, capfd, monkeypatch -): - _enable_nvfp4_packed_prefill(monkeypatch) - monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", fused) - monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") - shape = dict( - hidden_size=hidden_size, - inter_size=inter_size, - num_experts=num_experts, - top_k=top_k, - num_tokens=65, - onnx_dtype=onnx_dtype, - use_swiglu=True, - use_bias=True, - gemv_mode="0", - offline_prepacked=True, - input_scale=0.05, - enable_cuda_graph=True, - row_tile_size=row_tile_size, - ) - monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", "0") - separate = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) - separate_diagnostics = capfd.readouterr().out - monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", "1") - reused = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) - reused_diagnostics = capfd.readouterr().out - assert "Route=packed_nvfp4_prefill" in separate_diagnostics - assert "Route=packed_nvfp4_prefill" in reused_diagnostics - separate_bytes = int(re.search(r"RunnerWorkspaceBytes=(\d+)", separate_diagnostics).group(1)) - reused_bytes = int(re.search(r"RunnerWorkspaceBytes=(\d+)", reused_diagnostics).group(1)) - rows = row_tile_size or 65 - block_rows = 32 - while block_rows < 1024 and (rows + block_rows - 1) // block_rows * num_experts > block_rows: - block_rows *= 2 - blocks = (rows + block_rows - 1) // block_rows - - def align(value): - return (value + 255) // 256 * 256 - - routing_bytes = 2 * align(num_experts * blocks * 4) + align(rows * num_experts * 4) - input_bytes = align(rows * top_k * hidden_size * 2) - assert separate_bytes - reused_bytes == min(routing_bytes, input_bytes) - torch.testing.assert_close(reused, separate, atol=0, rtol=0) - - -@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") -@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) -@pytest.mark.parametrize( - "hidden_size,inter_size,num_experts,top_k", - [ - (64, 64, 16, 2), - (1152, 128, 16, 10), - (2560, 640, 16, 10), - (1152, 64, 512, 10), - ], -) -@pytest.mark.parametrize("row_tile_size", [0, 17]) -@pytest.mark.parametrize("reuse_scratch", ["0", "1"]) -def test_nvfp4_packed_prefill_tiled_down( - onnx_dtype, hidden_size, inter_size, num_experts, top_k, row_tile_size, reuse_scratch, capfd, monkeypatch -): - _enable_nvfp4_packed_prefill(monkeypatch) - monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", "1") - monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", reuse_scratch) - monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") - shape = dict( - hidden_size=hidden_size, - inter_size=inter_size, - num_experts=num_experts, - top_k=top_k, - num_tokens=65, - onnx_dtype=onnx_dtype, - use_swiglu=True, - use_bias=row_tile_size == 0, - gemv_mode="0", - offline_prepacked=True, - input_scale=0.05, - enable_cuda_graph=True, - row_tile_size=row_tile_size, - ) - monkeypatch.setenv("ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT", "0") - baseline = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) - baseline_bytes = int(re.search(r"RunnerWorkspaceBytes=(\d+)", capfd.readouterr().out).group(1)) - monkeypatch.setenv("ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT", "1") - tiled = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) - tiled_bytes = int(re.search(r"RunnerWorkspaceBytes=(\d+)", capfd.readouterr().out).group(1)) - rows = row_tile_size or 65 - - def aligned(value): - return (value + 255) // 256 * 256 - - input_bytes = aligned(rows * top_k * hidden_size * 2) - tiled_bytes_reservation = aligned(rows * top_k * min(hidden_size, 1024) * 2) - if reuse_scratch == "1": - block_rows = 32 - while block_rows < 1024 and (rows + block_rows - 1) // block_rows * num_experts > block_rows: - block_rows *= 2 - blocks = (rows + block_rows - 1) // block_rows - sorting_bytes = 2 * aligned(num_experts * blocks * 4) + aligned(rows * num_experts * 4) - input_bytes = max(input_bytes, sorting_bytes) - tiled_bytes_reservation = max(tiled_bytes_reservation, sorting_bytes) - assert baseline_bytes - tiled_bytes == input_bytes - tiled_bytes_reservation - torch.testing.assert_close(tiled, baseline, atol=0, rtol=0) - - -@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") -@pytest.mark.parametrize("mode", ["raw_layout", "unsupported_shape", "decode"]) -def test_nvfp4_packed_prefill_fallback(mode, capfd, monkeypatch): - _enable_nvfp4_packed_prefill(monkeypatch) - monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", "1") - monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", "1") - monkeypatch.setenv("ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT", "1") - monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") - decode = mode == "decode" - TestQMoENVFP4()._run_nvfp4_moe_test( - hidden_size=512 if decode else (80 if mode == "unsupported_shape" else 64), - inter_size=512 if decode else 64, - num_experts=16, - top_k=10, - num_tokens=1 if decode else 17, - onnx_dtype=TensorProto.FLOAT16, - use_swiglu=True, - gemv_mode="1" if decode else "0", - offline_prepacked=mode != "raw_layout", - input_scale=0.1, - ) - diagnostics = capfd.readouterr().out - assert "Route=packed_nvfp4_prefill" not in diagnostics - assert f"Route={'fp4_gemv_raw' if decode else 'grouped_moe'}" in diagnostics - - if __name__ == "__main__": unittest.main() diff --git a/onnxruntime/test/python/transformers/test_webgpu_matmul.py b/onnxruntime/test/python/transformers/test_webgpu_matmul.py index a9f8b9e6568bc..632190de11edb 100644 --- a/onnxruntime/test/python/transformers/test_webgpu_matmul.py +++ b/onnxruntime/test/python/transformers/test_webgpu_matmul.py @@ -80,13 +80,15 @@ def pointwise_conv_model(b, bias=None, relu=False): ).SerializeToString() -def matmul_session(model, device, *, capture=False, robustness=True): +def matmul_session(model, device, *, capture=False, robustness=True, fp32_accumulation=False): options = ort.SessionOptions() options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") options.use_deterministic_compute = True options.log_severity_level = 1 provider_options = {"validationMode": "full", "enableRobustness": str(int(robustness))} + if fp32_accumulation: + provider_options["enableMatmulFp32Accumulation"] = "1" if capture: provider_options["enableGraphCapture"] = "1" options.add_provider_for_devices([device], provider_options) @@ -152,6 +154,7 @@ def run_case( capture=False, updates=False, robustness=True, + fp32_accumulation=False, model=None, bias=None, relu=False, @@ -162,6 +165,7 @@ def run_case( self.device, capture=capture, robustness=robustness, + fp32_accumulation=fp32_accumulation, ) log.seek(0) initialization_log = log.read().decode("utf-8", errors="replace") @@ -211,9 +215,15 @@ def run_case( def test_exact_qwen_shapes(self): rng = np.random.default_rng(40) b = rng.normal(0, 0.1, (5120, 48)).astype(np.float16) - for m in (1, 8, 66): + for m in (1, 8, 9, 66): with self.subTest(m=m, output_shape=(m, 48)): - self.run_case(rng.normal(0, 0.1, (m, 5120)).astype(np.float16), b, gemv=m == 1) + # Exclude higher-priority f16 subgroup-matrix kernels without forcing GEMV. + self.run_case( + rng.normal(0, 0.1, (m, 5120)).astype(np.float16), + b, + gemv=m <= 8, + fp32_accumulation=True, + ) def test_target_and_nearby_shapes(self): rng = np.random.default_rng(42) @@ -250,8 +260,8 @@ def test_selector_boundaries(self): (1, 5120, n, np.float16, accepted) for n, accepted in [(15, False), (16, True), (17, False), (63, False), (64, True), (65, False)] ), - (1, 5120, 48, np.float16, True), - (2, 5120, 48, np.float16, False), + # Unaligned K excludes subgroup-matrix kernels so they cannot mask the M limit. + *((m, 2049, 48, np.float16, m <= 8) for m in range(1, 10)), (1, 5120, 48, np.float32, False), ] for m, k, n, dtype, gemv in cases: @@ -262,13 +272,22 @@ def test_selector_boundaries(self): gemv=gemv, ) - def test_batch_broadcast_stays_outside_gate(self): + def test_folded_a_batch_uses_gemv(self): rng = np.random.default_rng(48) - for a_shape, b_shape in [((2, 1, 5120), (5120, 48)), ((1, 5120), (2, 5120, 48))]: + self.run_case( + rng.normal(0, 0.1, (2, 1, 5120)).astype(np.float16), + rng.normal(0, 0.1, (5120, 48)).astype(np.float16), + ) + + def test_non_foldable_batched_weights_stay_outside_gate(self): + rng = np.random.default_rng(51) + for a_shape, b_shape in [((1, 2049), (2, 2049, 48)), ((2, 2, 2049), (2, 2049, 48))]: with self.subTest(a_shape=a_shape, b_shape=b_shape): + b = rng.normal(0, 0.1, b_shape).astype(np.float16) + b[1] = -b[0] self.run_case( rng.normal(0, 0.1, a_shape).astype(np.float16), - rng.normal(0, 0.1, b_shape).astype(np.float16), + b, gemv=False, ) @@ -348,7 +367,6 @@ def test_dynamic_weights_and_capture_replay(self): def test_outside_gate(self): rng = np.random.default_rng(45) cases = [ - ((2, 5120), (5120, 48), np.float16), ((2, 1, 2048), (2, 2048, 48), np.float16), ((1, 1024), (1024, 48), np.float16), ((1, 8193), (8193, 48), np.float16), diff --git a/onnxruntime/test/shared_lib/optional_provider_probe_test.cc b/onnxruntime/test/shared_lib/optional_provider_probe_test.cc index 3f64968484730..d611e61e239fd 100644 --- a/onnxruntime/test/shared_lib/optional_provider_probe_test.cc +++ b/onnxruntime/test/shared_lib/optional_provider_probe_test.cc @@ -91,7 +91,7 @@ TEST_F(OptionalProviderProbeTest, MissingExportUnloadsAndAllowsReplacement) { TEST_F(OptionalProviderProbeTest, SuccessfulProbeInitializesOnceAndCanReload) { ASSERT_NO_FATAL_FAILURE(Install(ORT_OPTIONAL_PROBE_shared_LIBRARY, shared_path_)); ASSERT_NO_FATAL_FAILURE(Install(ORT_OPTIONAL_PROBE_valid_LIBRARY, provider_path_)); - const auto* info = TryGetProviderInfo_OpenVINO(); + const ::ProviderInfo_OpenVINO* info = TryGetProviderInfo_OpenVINO(); ASSERT_NE(info, nullptr); EXPECT_EQ(TryGetProviderInfo_OpenVINO(), info); ASSERT_NO_FATAL_FAILURE(CheckInitializedOnce()); diff --git a/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc b/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc index a10eb97496186..325fa58772f76 100644 --- a/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc +++ b/onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc @@ -55,6 +55,18 @@ "^test_attention_(23|24)_.*qk_matmul_output_mode3", // #28994 CUDA NOT_IMPLEMENTED mode>kQK (attention.cc:1477); #27712 to support // TODO: support qk_matmul_output modes beyond kQK in Attention-cuda (see issue #27712) // Tests combining qk_matmul with softcap need unfused-path qk_matmul support (deferred). + "^test_attention_3d_local_window_cuda$", // exits with 0xDEAD in the CUDA backend test + "^test_attention_bidirectional_window_cuda$", // exits with 0xDEAD in the CUDA backend test + "^test_attention_local_window_cuda$", // exits with 0xDEAD in the CUDA backend test + "^test_attention_local_window_ext_cache_float16_mask_cuda$", // exits with 0xDEAD in the CUDA backend test + "^test_attention_local_window_ext_cache_rank2_mask_cuda$", // exits with 0xDEAD in the CUDA backend test + "^test_attention_local_window_ext_cache_rank3_head_mask_cuda$", // exits with 0xDEAD in the CUDA backend test + "^test_attention_local_window_ext_cache_rank4_batch_mask_cuda$", // exits with 0xDEAD in the CUDA backend test + "^test_attention_local_window_gqa_rank4_mask_cuda$", // exits with an access violation in the CUDA backend test + "^test_attention_local_window_gqa_rank4_mask_expanded_cuda$", // expanded Attention output mismatch + "^test_attention_local_window_rank1_boolean_mask_cuda$", // exits with an access violation in the CUDA backend test + "^test_attention_local_window_with_past_cuda$", // exits with an access violation in the CUDA backend test + "^test_attention_4d_causal_nonpad_negative_offset_structural_empty_expanded_cuda$", // expanded Attention output mismatch "^test_attention_3d_with_past_and_present_qk_matmul_softcap_cuda", // qk_matmul modes beyond kQK not supported "^test_attention_4d_with_qk_matmul_softcap_cuda", // qk_matmul modes beyond kQK not supported "^test_attention_3d_with_past_and_present_qk_matmul_bias_cuda", // QK matmul + bias not supported in Attention-cuda @@ -84,6 +96,179 @@ "^test_castlike_UINT2_to*", // ORT does not support uint2 "^test_cast_INT2_to*", // ORT does not support int2 "^test_cast_UINT2_to*", // ORT does not support uint2 + // TODO(#32378): Remove the exact BF16 Attention failures when the function dependencies support BF16. + "^test_attention_3d_causal_bf16_cpu$", + "^test_attention_3d_causal_bf16_expanded_cpu$", + "^test_attention_4d_attn_mask_causal_bf16_cpu$", + "^test_attention_4d_attn_mask_causal_bf16_expanded_cpu$", + "^test_attention_4d_causal_bf16_cpu$", + "^test_attention_4d_causal_bf16_expanded_cpu$", + "^test_attention_4d_causal_padded_kv_bf16_cpu$", + "^test_attention_4d_causal_padded_kv_bf16_expanded_cpu$", + "^test_attention_4d_padded_kv_bf16_cpu$", + "^test_attention_4d_padded_kv_bf16_expanded_cpu$", + "^test_attention_3d_causal_bf16_cuda$", + "^test_attention_3d_causal_bf16_expanded_cuda$", + "^test_attention_4d_attn_mask_causal_bf16_cuda$", + "^test_attention_4d_attn_mask_causal_bf16_expanded_cuda$", + "^test_attention_4d_causal_bf16_cuda$", + "^test_attention_4d_causal_bf16_expanded_cuda$", + "^test_attention_4d_causal_padded_kv_bf16_cuda$", + "^test_attention_4d_causal_padded_kv_bf16_expanded_cuda$", + "^test_attention_4d_padded_kv_bf16_cuda$", + "^test_attention_4d_padded_kv_bf16_expanded_cuda$", + "^test_attention_local_window_ext_cache_float16_mask_cpu$", + "^test_celu_bfloat16_cpu$", + "^test_celu_bfloat16_expanded_cpu$", + "^test_celu_cpu$", + "^test_celu_expanded_cpu$", + "^test_celu_float16_cpu$", + "^test_celu_float16_expanded_cpu$", + // TODO(#32381): Remove the exact log-reduction failures when their Cast(28) dependency is available. + "^test_reduce_log_sum_exp_default_axes_keepdims_example_cpu$", + "^test_reduce_log_sum_exp_default_axes_keepdims_example_expanded_cpu$", + "^test_reduce_log_sum_exp_default_axes_keepdims_random_cpu$", + "^test_reduce_log_sum_exp_default_axes_keepdims_random_expanded_cpu$", + "^test_reduce_log_sum_exp_do_not_keepdims_example_cpu$", + "^test_reduce_log_sum_exp_do_not_keepdims_example_expanded_cpu$", + "^test_reduce_log_sum_exp_do_not_keepdims_random_cpu$", + "^test_reduce_log_sum_exp_do_not_keepdims_random_expanded_cpu$", + "^test_reduce_log_sum_exp_empty_set_cpu$", + "^test_reduce_log_sum_exp_empty_set_expanded_cpu$", + "^test_reduce_log_sum_exp_keepdims_example_cpu$", + "^test_reduce_log_sum_exp_keepdims_example_expanded_cpu$", + "^test_reduce_log_sum_exp_keepdims_random_cpu$", + "^test_reduce_log_sum_exp_keepdims_random_expanded_cpu$", + "^test_reduce_log_sum_exp_negative_axes_keepdims_example_cpu$", + "^test_reduce_log_sum_exp_negative_axes_keepdims_example_expanded_cpu$", + "^test_reduce_log_sum_exp_negative_axes_keepdims_random_cpu$", + "^test_reduce_log_sum_exp_negative_axes_keepdims_random_expanded_cpu$", + // TODO: Remove the exact Compress(28) failures when the existing kernel supports opset 28. + "^test_compress_0_cpu$", + "^test_compress_1_cpu$", + "^test_compress_bfloat16_cpu$", + "^test_compress_default_axis_cpu$", + "^test_compress_negative_axis_cpu$", + // TODO: Remove the exact Einsum(28) failures when the existing kernel supports opset 28. + "^test_einsum_batch_diagonal_cpu$", + "^test_einsum_batch_matmul_bfloat16_cpu$", + "^test_einsum_batch_matmul_cpu$", + "^test_einsum_inner_prod_cpu$", + "^test_einsum_scalar_cpu$", + "^test_einsum_sum_bfloat16_cpu$", + "^test_einsum_sum_cpu$", + "^test_einsum_transpose_bfloat16_cpu$", + "^test_einsum_transpose_cpu$", + // TODO: Remove the exact OneHot(28) failures when the existing kernel supports opset 28. + "^test_onehot_negative_indices_cpu$", + "^test_onehot_out_of_range_indices_cpu$", + "^test_onehot_with_axis_cpu$", + "^test_onehot_with_bfloat16_values_cpu$", + "^test_onehot_with_negative_axis_cpu$", + "^test_onehot_without_axis_cpu$", + // TODO: Remove this test when ReduceMax(28) supports an empty bool input. + "^test_reduce_max_empty_set_bool_cpu$", + "^test_reduce_max_empty_set_bool_cuda$", + // TODO(#32381): Remove the exact Swish(24) function-dependency failures when SwiGLU is supported. + "^test_swiglu_alpha_cpu$", + "^test_swiglu_alpha_expanded_cpu$", + "^test_swiglu_cpu$", + "^test_swiglu_expanded_cpu$", + "^test_swiglu_float16_cpu$", + "^test_swiglu_float16_expanded_cpu$", + // TODO: Remove the exact Unique(28) failures when the existing kernel supports opset 28. + "^test_unique_bfloat16_sorted_without_axis_cpu$", + "^test_unique_length_1_cpu$", + "^test_unique_not_sorted_without_axis_cpu$", + "^test_unique_sorted_with_axis_3d_cpu$", + "^test_unique_sorted_with_axis_cpu$", + "^test_unique_sorted_with_negative_axis_cpu$", + "^test_unique_sorted_without_axis_cpu$", + // TODO: Remove these exact CUDA skips when fallback nodes have providers and Swish(24) has a CPU kernel. + "^test_celu_bfloat16_cuda$", + "^test_celu_bfloat16_expanded_cuda$", + "^test_celu_cuda$", + "^test_celu_expanded_cuda$", + "^test_celu_float16_cuda$", + "^test_celu_float16_expanded_cuda$", + "^test_compress_0_cuda$", + "^test_compress_1_cuda$", + "^test_compress_bfloat16_cuda$", + "^test_compress_default_axis_cuda$", + "^test_compress_negative_axis_cuda$", + "^test_einsum_batch_diagonal_cuda$", + "^test_einsum_batch_matmul_bfloat16_cuda$", + "^test_einsum_batch_matmul_cuda$", + "^test_einsum_inner_prod_cuda$", + "^test_einsum_scalar_cuda$", + "^test_einsum_sum_bfloat16_cuda$", + "^test_einsum_sum_cuda$", + "^test_einsum_transpose_bfloat16_cuda$", + "^test_einsum_transpose_cuda$", + "^test_onehot_negative_indices_cuda$", + "^test_onehot_out_of_range_indices_cuda$", + "^test_onehot_with_axis_cuda$", + "^test_onehot_with_bfloat16_values_cuda$", + "^test_onehot_with_negative_axis_cuda$", + "^test_onehot_without_axis_cuda$", + "^test_reduce_log_sum_exp_default_axes_keepdims_example_cuda$", + "^test_reduce_log_sum_exp_default_axes_keepdims_example_expanded_cuda$", + "^test_reduce_log_sum_exp_default_axes_keepdims_random_cuda$", + "^test_reduce_log_sum_exp_default_axes_keepdims_random_expanded_cuda$", + "^test_reduce_log_sum_exp_do_not_keepdims_example_cuda$", + "^test_reduce_log_sum_exp_do_not_keepdims_example_expanded_cuda$", + "^test_reduce_log_sum_exp_do_not_keepdims_random_cuda$", + "^test_reduce_log_sum_exp_do_not_keepdims_random_expanded_cuda$", + "^test_reduce_log_sum_exp_keepdims_example_cuda$", + "^test_reduce_log_sum_exp_keepdims_example_expanded_cuda$", + "^test_reduce_log_sum_exp_keepdims_random_cuda$", + "^test_reduce_log_sum_exp_keepdims_random_expanded_cuda$", + "^test_reduce_log_sum_exp_negative_axes_keepdims_example_cuda$", + "^test_reduce_log_sum_exp_negative_axes_keepdims_example_expanded_cuda$", + "^test_reduce_log_sum_exp_negative_axes_keepdims_random_cuda$", + "^test_reduce_log_sum_exp_negative_axes_keepdims_random_expanded_cuda$", + "^test_swiglu_alpha_cuda$", + "^test_swiglu_alpha_expanded_cuda$", + "^test_swiglu_cuda$", + "^test_swiglu_expanded_cuda$", + "^test_swiglu_float16_cuda$", + "^test_swiglu_float16_expanded_cuda$", + "^test_unique_bfloat16_sorted_without_axis_cuda$", + "^test_unique_length_1_cuda$", + "^test_unique_not_sorted_without_axis_cuda$", + "^test_unique_sorted_with_axis_3d_cuda$", + "^test_unique_sorted_with_axis_cuda$", + "^test_unique_sorted_with_negative_axis_cuda$", + "^test_unique_sorted_without_axis_cuda$", + "^test_mod_mixed_sign_int16_cuda$", + "^test_mod_mixed_sign_int8_cuda$", + "^test_mod_uint16_cuda$", + "^test_mod_uint8_cuda$", + // TODO(#32372): Remove these exact skips when the opset-28 Q/DQ kernels land. + "^test_dequantizelinear(_(cpu|cuda))?$", + "^test_dequantizelinear_axis(_(cpu|cuda))?$", + "^test_dequantizelinear_e4m3fn(_(cpu|cuda))?$", + "^test_dequantizelinear_e4m3fn_float16(_(cpu|cuda))?$", + "^test_dequantizelinear_e4m3fn_zero_point(_(cpu|cuda))?$", + "^test_dequantizelinear_e5m2(_(cpu|cuda))?$", + "^test_dequantizelinear_float4e2m1(_(cpu|cuda))?$", + "^test_dequantizelinear_int16(_(cpu|cuda))?$", + "^test_dequantizelinear_int2(_(cpu|cuda))?$", + "^test_dequantizelinear_uint16(_(cpu|cuda))?$", + "^test_dequantizelinear_uint2(_(cpu|cuda))?$", + "^test_quantizelinear(_(cpu|cuda))?$", + "^test_quantizelinear_axis(_(cpu|cuda))?$", + "^test_quantizelinear_e4m3fn(_(cpu|cuda))?$", + "^test_quantizelinear_e5m2(_(cpu|cuda))?$", + "^test_quantizelinear_float4e2m1(_(cpu|cuda))?$", + "^test_quantizelinear_int16(_(cpu|cuda))?$", + "^test_quantizelinear_int2(_(cpu|cuda))?$", + "^test_quantizelinear_uint16(_(cpu|cuda))?$", + "^test_quantizelinear_uint2(_(cpu|cuda))?$", + // TODO(#32378): Remove these exact skips when opset-28 ReverseSequence support lands. + ["opset28", "^test_reversesequence_batch(_(cpu|cuda))?$"], + ["opset28", "^test_reversesequence_bfloat16(_(cpu|cuda))?$"], + ["opset28", "^test_reversesequence_time(_(cpu|cuda))?$"], "^test_convinteger_with_padding*", // ORT does not support per-channel quantization for ConvInteger "^test_cast_e8m0_*", // ORT does not support float8e8m0 "^test_batchnorm_epsilon_training_mode", diff --git a/onnxruntime/test/testdata/onnx_backend_test_series_overrides.jsonc b/onnxruntime/test/testdata/onnx_backend_test_series_overrides.jsonc index f4da6b0d37a47..b56b0e5a97720 100644 --- a/onnxruntime/test/testdata/onnx_backend_test_series_overrides.jsonc +++ b/onnxruntime/test/testdata/onnx_backend_test_series_overrides.jsonc @@ -21,5 +21,8 @@ }, // Key: str, the name of the test as defined by ONNX without any device suffix. // Val: float, max relative difference between expected and actual. - "rtol_overrides": {} + "rtol_overrides": { + "test_attention_4d_causal_fp16": 0.002, + "test_attention_4d_causal_fp16_expanded": 0.002 + } } diff --git a/onnxruntime/test/unittest_util/graph_transform_test_builder.cc b/onnxruntime/test/unittest_util/graph_transform_test_builder.cc index 325f820343091..4e67113a4d76a 100644 --- a/onnxruntime/test/unittest_util/graph_transform_test_builder.cc +++ b/onnxruntime/test/unittest_util/graph_transform_test_builder.cc @@ -169,9 +169,7 @@ void TransformerTester(const std::function& buil SessionOptions session_options; session_options.graph_optimization_level = transformer ? baseline_level : level; // Mirror the model build's shape/type-inference strictness onto the session so that - // loading the serialized model applies the same options. (allow_released_opsets_only is - // passed explicitly to Load below, since the byte/proto Load overloads don't read it - // from the config options.) + // loading the serialized model applies the same options. ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry( kOrtSessionOptionsConfigStrictShapeTypeInference, model_options.strict_shape_type_inference ? "1" : "0")); @@ -187,8 +185,7 @@ void TransformerTester(const std::function& buil ASSERT_STATUS_OK(session.RegisterExecutionProvider(ep_shared)); } - // Load via an istream so the under-development opset flag can be honored: the istream Load - // overload takes allow_released_opsets_only explicitly (the byte/proto overloads hardcode it). + // Load via an istream so the under-development opset flag can be passed explicitly. // This lets models built at an unreleased ONNX opset (e.g. opset 27 in ONNX 1.22) load on // strict (ALLOW_RELEASED_ONNX_OPSET_ONLY default) CI legs. std::istringstream model_istream(model_data); diff --git a/onnxruntime/test/unittest_util/op_tester.cc b/onnxruntime/test/unittest_util/op_tester.cc index b633989866277..2b1f4d87df2e1 100644 --- a/onnxruntime/test/unittest_util/op_tester.cc +++ b/onnxruntime/test/unittest_util/op_tester.cc @@ -40,6 +40,12 @@ onnxruntime::Model& OpTester::BuildModel(const std::unordered_map domain_to_version(extra_domain_to_version); const auto& domain = Domain(); + if (domain != kOnnxDomain && domain_to_version.count(kOnnxDomain) == 0) { + const auto& last_released_versions = + ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance().LastReleaseVersionMap(); + domain_to_version.insert({kOnnxDomain, last_released_versions.at(kOnnxDomain)}); + } + if (domain_to_version.count(domain) == 0) { domain_to_version.insert({domain, Opset()}); } else { diff --git a/onnxruntime/test/util/default_providers.cc b/onnxruntime/test/util/default_providers.cc index b730060441eec..f2453df76a106 100644 --- a/onnxruntime/test/util/default_providers.cc +++ b/onnxruntime/test/util/default_providers.cc @@ -419,17 +419,18 @@ std::unique_ptr WebGpuExecutionProviderWithOptions(const Con #endif } -std::unique_ptr WebGpuExecutionProviderWithTestStorageBufferBindingSize(uint64_t max_size) { - return WebGpuExecutionProviderWithTestStorageBufferBindingSize(ConfigOptions{}, max_size); +std::unique_ptr WebGpuExecutionProviderWithTestOptions( + const WebGpuExecutionProviderTestOptions& test_options) { + return WebGpuExecutionProviderWithTestOptions(ConfigOptions{}, test_options); } -std::unique_ptr WebGpuExecutionProviderWithTestStorageBufferBindingSize( - const ConfigOptions& config_options, uint64_t max_size) { +std::unique_ptr WebGpuExecutionProviderWithTestOptions( + const ConfigOptions& config_options, const WebGpuExecutionProviderTestOptions& test_options) { #if defined(USE_WEBGPU) && !defined(ORT_USE_EP_API_ADAPTERS) - return WebGpuProviderFactoryCreator::CreateForTesting(config_options, max_size)->CreateProvider(); + return WebGpuProviderFactoryCreator::CreateForTesting(config_options, test_options)->CreateProvider(); #else ORT_UNUSED_PARAMETER(config_options); - ORT_UNUSED_PARAMETER(max_size); + ORT_UNUSED_PARAMETER(test_options); return nullptr; #endif } diff --git a/onnxruntime/test/util/include/default_providers.h b/onnxruntime/test/util/include/default_providers.h index c0ceee6da60a2..9e424573525d7 100644 --- a/onnxruntime/test/util/include/default_providers.h +++ b/onnxruntime/test/util/include/default_providers.h @@ -5,6 +5,7 @@ #include "core/common/optional.h" #include "core/providers/providers.h" #include "core/providers/provider_factory_creators.h" +#include "core/providers/webgpu/webgpu_provider_factory_creator.h" #include "core/framework/config_options.h" #include "core/framework/execution_provider.h" @@ -63,9 +64,10 @@ std::unique_ptr QnnExecutionProviderWithOptions(const Provid std::unique_ptr DefaultXnnpackExecutionProvider(); std::unique_ptr DefaultWebGpuExecutionProvider(bool is_nhwc = true); std::unique_ptr WebGpuExecutionProviderWithOptions(const ConfigOptions& config_options); -std::unique_ptr WebGpuExecutionProviderWithTestStorageBufferBindingSize(uint64_t max_size); -std::unique_ptr WebGpuExecutionProviderWithTestStorageBufferBindingSize( - const ConfigOptions& config_options, uint64_t max_size); +std::unique_ptr WebGpuExecutionProviderWithTestOptions( + const WebGpuExecutionProviderTestOptions& test_options); +std::unique_ptr WebGpuExecutionProviderWithTestOptions( + const ConfigOptions& config_options, const WebGpuExecutionProviderTestOptions& test_options); std::unique_ptr DefaultCannExecutionProvider(); std::unique_ptr DefaultDmlExecutionProvider(); diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 92f803030ada4..207038805f624 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -5082,7 +5082,7 @@ void RegisterOrtOpSchemas() { auto& domainToVersionRangeInstance = ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance(); if (domainToVersionRangeInstance.Map().find(onnxruntime::kMSDomain) == domainToVersionRangeInstance.Map().end()) { // External shared providers may have already added kMSDomain - domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 1); + domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSDomain, 1, 2); } domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSExperimentalDomain, 1, 1); domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSNchwcDomain, 1, 1); diff --git a/orttraining/orttraining/core/optimizer/cast_sce_loss_fusion.cc b/orttraining/orttraining/core/optimizer/cast_sce_loss_fusion.cc index cb692d347d706..a54af274724b8 100644 --- a/orttraining/orttraining/core/optimizer/cast_sce_loss_fusion.cc +++ b/orttraining/orttraining/core/optimizer/cast_sce_loss_fusion.cc @@ -30,7 +30,7 @@ Status CastSceLossFusion::ApplyImpl(Graph& graph, bool& modified, int graph_leve Node* input_node = graph.GetMutableProducerNode(node.MutableInputDefs()[0]->Name()); - if (!(graph_utils::IsSupportedOptypeVersionAndDomain(*input_node, "Cast", {9, 13, 19}))) { + if (!input_node || !(graph_utils::IsSupportedOptypeVersionAndDomain(*input_node, "Cast", {9, 13, 19}))) { continue; } diff --git a/orttraining/orttraining/core/optimizer/loss_rewriter.cc b/orttraining/orttraining/core/optimizer/loss_rewriter.cc index f2fc0df42924f..60f95707fa104 100644 --- a/orttraining/orttraining/core/optimizer/loss_rewriter.cc +++ b/orttraining/orttraining/core/optimizer/loss_rewriter.cc @@ -25,6 +25,10 @@ Status SoftmaxCrossEntropyLossInternalFusion::ApplyImpl(Graph& graph, bool& modi } if (p_first_input->OpType() == "Cast") { + if (p_first_input->GetOutputEdgesCount() != 1 || graph.NodeProducesGraphOutput(*p_first_input)) { + continue; + } + p_first_input = graph_utils::GetInputNode(*p_first_input, 0); if (!p_first_input || p_first_input->OpType() != "LogSoftmax") { continue; @@ -32,6 +36,30 @@ Status SoftmaxCrossEntropyLossInternalFusion::ApplyImpl(Graph& graph, bool& modi } Node& log_softmax_node = *graph.GetNode(p_first_input->Index()); + if (log_softmax_node.GetOutputEdgesCount() != 1 || graph.NodeProducesGraphOutput(log_softmax_node)) { + continue; + } + + const auto* input_shape = log_softmax_node.InputDefs()[0]->Shape(); + if (log_softmax_node.SinceVersion() < 13 && + (input_shape == nullptr || input_shape->dim_size() != 2)) { + continue; + } + + const auto axis_attr = log_softmax_node.GetAttributes().find("axis"); + int64_t axis = axis_attr == log_softmax_node.GetAttributes().end() + ? (log_softmax_node.SinceVersion() < 13 ? 1 : -1) + : axis_attr->second.i(); + if (axis < 0) { + if (input_shape == nullptr) { + continue; + } + axis += input_shape->dim_size(); + } + if (axis != 1) { + continue; + } + graph_utils::RemoveNode(graph, log_softmax_node); const auto& loss_inputs = loss_node.MutableInputDefs(); auto& loss_outputs = loss_node.MutableOutputDefs(); diff --git a/orttraining/orttraining/test/optimizer/graph_transform_test.cc b/orttraining/orttraining/test/optimizer/graph_transform_test.cc index 26220f781d91f..1278f6e30dd81 100644 --- a/orttraining/orttraining/test/optimizer/graph_transform_test.cc +++ b/orttraining/orttraining/test/optimizer/graph_transform_test.cc @@ -793,6 +793,27 @@ TEST_F(GraphTransformationTests, CastSceLossFusion) { ASSERT_EQ(op_to_count["Cast"], 9); } +TEST_F(GraphTransformationTests, CastSceLossFusionWithGraphInput) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* scores = builder.MakeInput({{2, 3}}); + auto* labels = builder.MakeInput({{2}}); + auto* loss = builder.MakeOutput(std::vector{}); + auto* log_prob = builder.MakeOutput(std::vector{2, 3}); + builder.AddNode("SoftmaxCrossEntropyLossInternal", {scores, labels}, {loss, log_prob}, kMSDomain) + .AddAttribute("reduction", "mean"); + }; + + auto check_unchanged = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.at("com.microsoft.SoftmaxCrossEntropyLossInternal") == 1); + TEST_RETURN_IF(op_count.count("Cast") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level2, 1, check_unchanged, check_unchanged)); +} + Node* GetNodeByName(Graph& graph, std::string node_name) { GraphViewer graph_viewer(graph); const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder(); @@ -1207,10 +1228,13 @@ TEST_F(GraphTransformationTests, SoftmaxCrossEntropyLossInternalFusionWithoutCas TypeProto tensor_float; tensor_float.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); + TypeProto tensor_float_matrix = tensor_float; + tensor_float_matrix.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(2); + tensor_float_matrix.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(3); TypeProto tensor_int; tensor_int.mutable_tensor_type()->set_elem_type(TensorProto_DataType_INT64); - onnxruntime::NodeArg x_def("X", &tensor_float); - onnxruntime::NodeArg ls_out_def("ls_out", &tensor_float); + onnxruntime::NodeArg x_def("X", &tensor_float_matrix); + onnxruntime::NodeArg ls_out_def("ls_out", &tensor_float_matrix); onnxruntime::NodeArg target_def("target", &tensor_int); onnxruntime::NodeArg weight_def("weight", &tensor_float); onnxruntime::NodeArg ignore_index_def("ignore_index", &tensor_int); @@ -1243,13 +1267,19 @@ TEST_F(GraphTransformationTests, SoftmaxCrossEntropyLossInternalFusionWithCast) TypeProto tensor_half; tensor_half.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT16); + TypeProto tensor_half_matrix = tensor_half; + tensor_half_matrix.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(2); + tensor_half_matrix.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(3); TypeProto tensor_float; tensor_float.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); + TypeProto tensor_float_matrix = tensor_float; + tensor_float_matrix.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(2); + tensor_float_matrix.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(3); TypeProto tensor_int; tensor_int.mutable_tensor_type()->set_elem_type(TensorProto_DataType_INT64); - onnxruntime::NodeArg x_def("X", &tensor_half); - onnxruntime::NodeArg ls_out_def("ls_out", &tensor_half); - onnxruntime::NodeArg ct_out_def("ct_out", &tensor_float); + onnxruntime::NodeArg x_def("X", &tensor_half_matrix); + onnxruntime::NodeArg ls_out_def("ls_out", &tensor_half_matrix); + onnxruntime::NodeArg ct_out_def("ct_out", &tensor_float_matrix); onnxruntime::NodeArg target_def("target", &tensor_int); onnxruntime::NodeArg weight_def("weight", &tensor_float); onnxruntime::NodeArg ignore_index_def("ignore_index", &tensor_int); @@ -1278,6 +1308,161 @@ TEST_F(GraphTransformationTests, SoftmaxCrossEntropyLossInternalFusionWithCast) ASSERT_TRUE(op_to_count["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); } +TEST_F(GraphTransformationTests, SoftmaxCrossEntropyLossInternalFusionGuards) { + struct TestCase { + int64_t axis; + bool public_log_prob; + bool should_fuse; + }; + + for (const auto test_case : + {TestCase{1, false, true}, TestCase{-1, false, true}, TestCase{0, false, false}, + TestCase{1, true, false}}) { + auto build_test_case = [test_case](ModelTestBuilder& builder) { + auto* scores = builder.MakeInput({{2, 3}}); + auto* labels = builder.MakeInput({{2}}); + NodeArg* log_prob = test_case.public_log_prob + ? builder.MakeOutput(std::vector{2, 3}) + : builder.MakeIntermediate({{2, 3}}); + auto* loss = builder.MakeOutput(std::vector{}); + builder.AddNode("LogSoftmax", {scores}, {log_prob}).AddAttribute("axis", test_case.axis); + builder.AddNode("NegativeLogLikelihoodLossInternal", {log_prob, labels}, {loss}, kMSDomain) + .AddAttribute("reduction", "mean"); + }; + + auto check_before = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.at("LogSoftmax") == 1); + TEST_RETURN_IF_NOT(op_count.at("com.microsoft.NegativeLogLikelihoodLossInternal") == 1); + return Status::OK(); + }; + auto check_after = [test_case](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.count("LogSoftmax") == (test_case.should_fuse ? 0U : 1U)); + TEST_RETURN_IF_NOT(op_count.count("com.microsoft.NegativeLogLikelihoodLossInternal") == + (test_case.should_fuse ? 0U : 1U)); + TEST_RETURN_IF_NOT(op_count.count("com.microsoft.SoftmaxCrossEntropyLossInternal") == + (test_case.should_fuse ? 1U : 0U)); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, + std::make_unique(), + TransformerLevel::Level1, 1, check_before, check_after)); + } +} + +TEST_F(GraphTransformationTests, SoftmaxCrossEntropyLossInternalFusionPre13RequiresRankTwo) { + for (const auto& scores_shape : {std::vector{2, 3}, std::vector{2, 3, 4}}) { + const bool should_fuse = scores_shape.size() == 2; + auto build_test_case = [scores_shape](ModelTestBuilder& builder) { + auto labels_shape = scores_shape; + labels_shape.erase(labels_shape.begin() + 1); + auto* scores = builder.MakeInput(scores_shape); + auto* labels = builder.MakeInput(labels_shape); + auto* log_prob = builder.MakeIntermediate(scores_shape); + auto* loss = builder.MakeOutput(std::vector{}); + builder.AddNode("LogSoftmax", {scores}, {log_prob}).AddAttribute("axis", int64_t{1}); + builder.AddNode("NegativeLogLikelihoodLossInternal", {log_prob, labels}, {loss}, kMSDomain) + .AddAttribute("reduction", "mean"); + }; + + auto check_after = [should_fuse](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(OpCount(op_count, "LogSoftmax") == (should_fuse ? 0 : 1)); + TEST_RETURN_IF_NOT(OpCount(op_count, "com.microsoft.NegativeLogLikelihoodLossInternal") == + (should_fuse ? 0 : 1)); + TEST_RETURN_IF_NOT(OpCount(op_count, "com.microsoft.SoftmaxCrossEntropyLossInternal") == + (should_fuse ? 1 : 0)); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 12, *logger_, std::make_unique(), + TransformerLevel::Level1, 1, nullptr, check_after)); + } +} + +TEST_F(GraphTransformationTests, SoftmaxCrossEntropyLossInternalFusionPre13RejectsUnknownRank) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* scores = builder.MakeInput(std::nullopt); + auto* labels = builder.MakeInput(std::nullopt); + auto* log_prob = builder.MakeIntermediate(std::nullopt); + auto* loss = builder.MakeOutput(std::vector{}); + builder.AddNode("LogSoftmax", {scores}, {log_prob}).AddAttribute("axis", int64_t{1}); + builder.AddNode("NegativeLogLikelihoodLossInternal", {log_prob, labels}, {loss}, kMSDomain) + .AddAttribute("reduction", "mean"); + }; + + auto check_unchanged = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(OpCount(op_count, "LogSoftmax") == 1); + TEST_RETURN_IF_NOT(OpCount(op_count, "com.microsoft.NegativeLogLikelihoodLossInternal") == 1); + TEST_RETURN_IF_NOT(OpCount(op_count, "com.microsoft.SoftmaxCrossEntropyLossInternal") == 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 12, *logger_, std::make_unique(), + TransformerLevel::Level1, 1, check_unchanged, check_unchanged)); +} + +TEST_F(GraphTransformationTests, SoftmaxCrossEntropyLossInternalFusionRejectsSharedIntermediates) { + enum class SharedValue { + DirectLogSoftmaxOutput, + LogSoftmaxOutput, + CastOutput, + PublicCastOutput, + }; + + for (const auto shared_value : {SharedValue::DirectLogSoftmaxOutput, SharedValue::LogSoftmaxOutput, + SharedValue::CastOutput, SharedValue::PublicCastOutput}) { + auto build_test_case = [shared_value](ModelTestBuilder& builder) { + const bool use_cast = shared_value != SharedValue::DirectLogSoftmaxOutput; + const bool public_cast_output = shared_value == SharedValue::PublicCastOutput; + auto* scores = use_cast ? builder.MakeInput({{2, 3}}) : builder.MakeInput({{2, 3}}); + auto* labels = builder.MakeInput({{2}}); + auto* log_prob = use_cast ? builder.MakeIntermediate({{2, 3}}) + : builder.MakeIntermediate({{2, 3}}); + NodeArg* nll_input = log_prob; + NodeArg* cast_log_prob = nullptr; + auto* loss = builder.MakeOutput(std::vector{}); + + builder.AddNode("LogSoftmax", {scores}, {log_prob}).AddAttribute("axis", int64_t{-1}); + if (use_cast) { + cast_log_prob = public_cast_output + ? builder.MakeOutput(std::vector{2, 3}) + : builder.MakeIntermediate(std::vector{2, 3}); + builder.AddNode("Cast", {log_prob}, {cast_log_prob}) + .AddAttribute("to", static_cast(TensorProto_DataType_FLOAT)); + nll_input = cast_log_prob; + } + builder.AddNode("NegativeLogLikelihoodLossInternal", {nll_input, labels}, {loss}, kMSDomain) + .AddAttribute("reduction", "mean"); + if (!public_cast_output) { + auto* shared_output = builder.MakeOutput(); + builder.AddNode("Identity", + {shared_value == SharedValue::CastOutput ? cast_log_prob : log_prob}, + {shared_output}); + } + }; + + auto check_unchanged = [shared_value](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(OpCount(op_count, "LogSoftmax") == 1); + TEST_RETURN_IF_NOT(OpCount(op_count, "Cast") == + (shared_value == SharedValue::DirectLogSoftmaxOutput ? 0 : 1)); + TEST_RETURN_IF_NOT(OpCount(op_count, "com.microsoft.NegativeLogLikelihoodLossInternal") == 1); + TEST_RETURN_IF_NOT(OpCount(op_count, "com.microsoft.SoftmaxCrossEntropyLossInternal") == 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer( + build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level1, 1, check_unchanged, check_unchanged)); + } +} + class LSTMReplacementTestsParameterized : public GraphTransformationTests, public ::testing::WithParamInterface> { }; diff --git a/plugin-ep-cuda/VERSION_NUMBER b/plugin-ep-cuda/VERSION_NUMBER index 0d91a54c7d439..1d0ba9ea182b0 100644 --- a/plugin-ep-cuda/VERSION_NUMBER +++ b/plugin-ep-cuda/VERSION_NUMBER @@ -1 +1 @@ -0.3.0 +0.4.0 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 0d2f339ca25ce..0e64e5205b785 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 @@ -18,7 +18,8 @@ ONNX;ONNX Runtime;Machine Learning;AI;Deep Learning;CUDA;GPU - MIT + LICENSE + https://onnxruntime.ai https://github.com/microsoft/onnxruntime git © Microsoft Corporation. All rights reserved. @@ -31,6 +32,9 @@ + + +