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_cuda_ci.yml b/.github/workflows/linux_cuda_ci.yml index d5fceea751d14..af9257757a0d1 100644 --- a/.github/workflows/linux_cuda_ci.yml +++ b/.github/workflows/linux_cuda_ci.yml @@ -138,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: 'ORT_RUN_CUDA_TACTIC_CACHE_TESTS=1 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/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/.github/workflows/windows_webgpu.yml b/.github/workflows/windows_webgpu.yml index a54d9d7f12590..e2457ad51e2c0 100644 --- a/.github/workflows/windows_webgpu.yml +++ b/.github/workflows/windows_webgpu.yml @@ -343,9 +343,9 @@ jobs: --cmake_generator "Visual Studio 17 2022" ` --enable_onnx_tests ` --use_webgpu shared_lib ` - --build_wheel ` --use_vcpkg --use_vcpkg_ms_internal_asset_cache ` --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=ON onnxruntime_ENABLE_DAWN_BACKEND_D3D12=1 onnxruntime_ENABLE_DAWN_BACKEND_VULKAN=1 ` + --disable_rtti ` --enable_lto if ($lastExitCode -ne 0) { @@ -410,7 +410,7 @@ jobs: architecture: x64 - name: Install python modules - run: python -m pip install -r tools\ci_build\github\windows\python\requirements.txt -r plugin-ep-webgpu\python\requirements-build-wheel.txt + run: python -m pip install -r tools\ci_build\github\windows\python\requirements.txt shell: cmd working-directory: ${{ github.workspace }} @@ -426,7 +426,6 @@ jobs: - name: Run tests shell: pwsh - working-directory: ${{ runner.temp }} run: | python.exe ${{ github.workspace }}\tools\python\update_ctest_path.py ` "${{ github.workspace }}\RelWithDebInfo\CTestTestfile.cmake" ` @@ -444,27 +443,6 @@ jobs: exit $lastExitCode } - $binaryDir = "${{ github.workspace }}\RelWithDebInfo\RelWithDebInfo" - $pluginWheelDir = "${{ runner.temp }}\webgpu-plugin-wheels" - python "${{ github.workspace }}\plugin-ep-webgpu\python\build_wheel.py" ` - --binary_dir $binaryDir --version 0.0.0 --output_dir $pluginWheelDir - if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } - - # Install the host from this PR, not the minimum-runtime wheel used by package compatibility CI. - $hostWheels = @(Get-ChildItem "$binaryDir\dist\onnxruntime_webgpu-*.whl") - $pluginWheels = @(Get-ChildItem "$pluginWheelDir\onnxruntime_ep_webgpu-*.whl") - if ($hostWheels.Count -ne 1 -or $pluginWheels.Count -ne 1) { - throw "Expected exactly one source-built host wheel and one plugin wheel" - } - # Dependencies were installed once for this job in Install python modules. - python -m pip install --no-deps --force-reinstall $hostWheels[0].FullName $pluginWheels[0].FullName - if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } - - $env:ORT_WEBGPU_TEST_REQUIRE_EP_DEVICE = '1' - $env:ORT_WEBGPU_TEST_PROVIDER_NAME = '1' - python -u "${{ github.workspace }}\plugin-ep-webgpu\python\test\test_webgpu_plugin_ep.py" - if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } - # Same plugin EP sources as webgpu_plugin_build_x64_RelWithDebInfo, but linked into the ORT binary # instead of built as a separate shared library. webgpu_static_plugin_build_x64_RelWithDebInfo: 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/deps.txt b/cmake/deps.txt index 2be571e06fbe5..b31831f7eef7e 100644 --- a/cmake/deps.txt +++ b/cmake/deps.txt @@ -58,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/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 94ce9f307c893..aec4bb0cb2274 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -3095,8 +3095,7 @@ if (onnxruntime_USE_WEBGPU AND onnxruntime_USE_EXTERNAL_DAWN AND TARGET dawn::da target_compile_features(onnxruntime_webgpu_external_dawn_test PRIVATE cxx_std_20) endif() -# This test uses the built-in WebGPU append API and does not register a plugin. -if (onnxruntime_USE_WEBGPU AND NOT onnxruntime_USE_EP_API_ADAPTERS AND WIN32 AND onnxruntime_BUILD_SHARED_LIB AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten" AND NOT onnxruntime_MINIMAL_BUILD) +if (onnxruntime_USE_WEBGPU AND WIN32 AND onnxruntime_BUILD_SHARED_LIB AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten" AND NOT onnxruntime_MINIMAL_BUILD) AddTest(DYN TARGET onnxruntime_webgpu_delay_load_test SOURCES ${onnxruntime_webgpu_delay_load_test_SRC} @@ -3135,4 +3134,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/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 f9e8e0013b41c..de716517c700d 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 @@ -84,6 +85,7 @@ Do not modify directly.* * com.microsoft.PackedAttention * com.microsoft.PackedMultiHeadAttention * com.microsoft.PackedSparseAttentionIndexer + * com.microsoft.PackedSparseAttentionIndexerMerge * com.microsoft.Pad * com.microsoft.PagedAttention * com.microsoft.QAttention @@ -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: @@ -5098,6 +5172,10 @@ This version of the operator has been available since version 1 of the 'com.micr
Only for policy_mode 'csa': scale applied to head_weights. Default is 1/sqrt(num_heads). Must be omitted when policy_mode is 'qsa'.
index_topk : int
Only for policy_mode 'csa': number of compressed entries selected per query. Must be > 0. Must be omitted when policy_mode is 'qsa'.
+
indexshare_state_policy : string
+
Reuse state policy: default holds QSA and advances CSA; hold or advance overrides it.
+
max_output_entries : int
+
Fixed padded output width for the optional IndexShare interface.
policy_mode : string (required)
Indexer policy. Must be exactly 'qsa' (token indexer) or 'csa' (compressed block indexer).
scale : float
@@ -5110,7 +5188,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Only for policy_mode 'qsa': maximum number of tokens selected from complete blocks. Must be > 0 and divisible by compress_ratio. Must be omitted when policy_mode is 'csa'.
-#### Inputs (16 - 18) +#### Inputs (16 - 22)
query : T
@@ -5149,9 +5227,17 @@ This version of the operator has been available since version 1 of the 'com.micr
Only for policy_mode 'qsa': number of leading token transitions to capture for each request, with shape (batch_size). Values are clamped to the request length and state_update_capacity. Required when state_update_capacity is positive.
state_update_active (optional) : M
Only for policy_mode 'qsa': optional capture gate with shape (1). A zero value disables capture.
+
indexshare_mode (optional) : M
+
CPU scalar [1]: 0 refresh, 1 reuse.
+
merged_indices (optional) : M
+
Packed merged selection [N,max_output_entries].
+
merged_counts (optional) : M
+
Merged counts [N].
+
merged_status (optional) : M
+
Merged device status [N].
-#### Outputs (6 - 7) +#### Outputs (6 - 8)
selected_indices : M
@@ -5168,6 +5254,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Updated generic per-request state length, with the same fixed shape as past_state_lengths.
state_update (optional) : T
Only for policy_mode 'qsa': compact transition payloads with shape (batch_size, state_update_capacity, head_size). A token that completes a compression block stores the prepared block representative; any other captured token stores its raw key. Inactive and unused slots are zero.
+
indexshare_status (optional) : M
+
Selection status [N], forwarded on reuse and zero on refresh.
#### Type Constraints @@ -5182,6 +5270,61 @@ 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 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 + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
max_output_entries : int (required)
+
Fixed output width, in [1, 2^30].
+
policy_mode : string (required)
+
append_range or append_indices.
+
+ +#### Inputs (3 - 7) + +
+
base_indices : T
+
Cached selections, shape (R, C).
+
base_counts : T
+
Valid base-prefix lengths, shape (R).
+
base_row_indices : T
+
Output-query to cached-row mapping, shape (N).
+
range_starts (optional) : T
+
Inclusive nonnegative starts, shape (N).
+
range_ends (optional) : T
+
Exclusive ends not less than starts, shape (N).
+
additional_indices (optional) : T
+
Required for append_indices: selections of shape (N, A); absent for append_range.
+
additional_counts (optional) : T
+
Required for append_indices: valid prefix lengths of shape (N); absent for append_range.
+
+ +#### Outputs + +
+
selected_indices : T
+
Merged indices, shape (N, max_output_entries).
+
selected_counts : T
+
Merged prefix lengths, shape (N).
+
status : T
+
Per-query device validation status, shape (N).
+
+ +#### Type Constraints + +
+
T : tensor(int32)
+
Integer selection tensors.
+
+ + ### **com.microsoft.Pad** Given `data` tensor, pads, mode, and value. @@ -6121,6 +6264,10 @@ This version of the operator has been available since version 1 of the 'com.micr 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. @@ -6182,7 +6329,7 @@ This version of the operator has been available since version 1 of the 'com.micr
use_sparse_mixer : int
Whether to use sparse mixer
weights_prepacked : int
-
Only meaningful when quant_type='int'. Tri-state control over the layout of the int4/int8 fc1/fc2 weight initializers. The concrete prepacked layouts selected by -1 and 1 are determined by the execution provider. 0: the initializers are raw, un-prepacked [E, N, K/pack] tensors as produced by quantize_matmul_{4,8}bits. Defaults to -1.
+
Expert-weight layout selector with quantization-specific meanings. For quant_type='int', -1 and 1 select execution-provider-specific prepacked int4/int8 fc1/fc2 layouts. 0 selects raw [E,N,K/pack] tensors as produced by quantize_matmul_{4,8}bits. For quant_type='nvfp4' on CUDA, -1 (default) and 0 select legacy N-packed [E,K,N/2] bytes; 1 selects K-packed row-major [E,N,K/2] bytes while retaining logical [E,K,N/2] dimensions. NVFP4 block scales and global scales retain their original layouts and values. Ignored for other quantization types. Supported values are -1, 0, and 1; defaults to -1.
#### Inputs (6 - 21) @@ -8624,3 +8771,5 @@ 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 6cbca2e452b8a..669e3083f6b15 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1136,7 +1136,8 @@ The **OpSet Version** column uses the following notation: |NhwcConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |PackedAttention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* token_offset:**M**
*in* cumulative_sequence_length:**M**
*in* attention_bias:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |PackedMultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* token_offset:**M**
*in* cumulative_sequence_length:**M**
*in* attention_bias:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| -|PackedSparseAttentionIndexer|*in* query:**T**
*in* key:**T**
*in* query_norm_weight:**T**
*in* key_norm_weight:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* cumulative_sequence_lengths:**M**
*in* past_sequence_lengths:**M**
*in* gate:**T**
*in* position_bias:**T**
*in* head_weights:**T**
*in* position_ids:**I**
*in* past_key_state:**T**
*in* past_kv_buffer:**T**
*in* past_gate_buffer:**T**
*in* past_state_lengths:**M**
*in* state_update_capture_count:**M**
*in* state_update_active:**M**
*out* selected_indices:**M**
*out* selected_counts:**M**
*out* present_key_state:**T**
*out* present_kv_buffer:**T**
*out* present_gate_buffer:**T**
*out* present_state_lengths:**M**
*out* state_update:**T**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|PackedSparseAttentionIndexer|*in* query:**T**
*in* key:**T**
*in* query_norm_weight:**T**
*in* key_norm_weight:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* cumulative_sequence_lengths:**M**
*in* past_sequence_lengths:**M**
*in* gate:**T**
*in* position_bias:**T**
*in* head_weights:**T**
*in* position_ids:**I**
*in* past_key_state:**T**
*in* past_kv_buffer:**T**
*in* past_gate_buffer:**T**
*in* past_state_lengths:**M**
*in* state_update_capture_count:**M**
*in* state_update_active:**M**
*in* indexshare_mode:**M**
*in* merged_indices:**M**
*in* merged_counts:**M**
*in* merged_status:**M**
*out* selected_indices:**M**
*out* selected_counts:**M**
*out* present_key_state:**T**
*out* present_kv_buffer:**T**
*out* present_gate_buffer:**T**
*out* present_state_lengths:**M**
*out* state_update:**T**
*out* indexshare_status:**M**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|PackedSparseAttentionIndexerMerge|*in* base_indices:**T**
*in* base_counts:**T**
*in* base_row_indices:**T**
*in* range_starts:**T**
*in* range_ends:**T**
*in* additional_indices:**T**
*in* additional_counts:**T**
*out* selected_indices:**T**
*out* selected_counts:**T**
*out* status:**T**|1+|**T** = tensor(int32)| |PagedAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* key_cache:**T_CACHE**
*in* value_cache:**T_CACHE**
*in* cumulative_sequence_length:**S**
*in* past_seqlens:**S**
*in* block_table:**S**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* slot_mapping:**S**
*in* head_sink:**T**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* attention_metadata:**S**
*out* output:**T**
*out* key_cache_out:**T_CACHE**
*out* value_cache_out:**T_CACHE**|1+|**S** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(int8), tensor(uint8)
**T_KV_SCALE** = tensor(float)| |QAttention|*in* input:**T1**
*in* weight:**T2**
*in* bias:**T3**
*in* input_scale:**T3**
*in* weight_scale:**T3**
*in* mask_index:**T4**
*in* input_zero_point:**T1**
*in* weight_zero_point:**T2**
*in* past:**T3**
*out* output:**T3**
*out* present:**T3**|1+|**T1** = tensor(int8)
**T2** = tensor(int8)
**T3** = tensor(float), tensor(float16)
**T4** = tensor(int32)| |QMoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T1**
*in* fc1_scales:**T2**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T1**
*in* fc2_scales:**T2**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T1**
*in* fc3_scales:**T2**
*in* fc3_experts_bias:**T**
*in* fc1_zero_points:**T1**
*in* fc2_zero_points:**T1**
*in* fc3_zero_points:**T1**
*in* router_weights:**T**
*in* fc1_global_scale:**T4**
*in* fc2_global_scale:**T4**
*in* fc1_act_scale:**T4**
*in* fc2_act_scale:**T4**
*in* fc1_act_block_scale:**T2**
*in* fc2_act_block_scale:**T2**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(float8e4m3fn), tensor(uint8)
**T2** = tensor(bfloat16), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e8m0)
**T4** = tensor(float)| 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 new file mode 100644 index 0000000000000..dd4cf123b6b32 --- /dev/null +++ b/docs/contrib_ops/PackedSparseAttentionIndexerMerge.md @@ -0,0 +1,45 @@ +# PackedSparseAttentionIndexerMerge + +`com.microsoft::PackedSparseAttentionIndexerMerge`, version 1, forms a stable union +of a mapped base selection and additional indices. CUDA and WebGPU kernels +support `policy_mode="append_range"`; CUDA also supports `append_indices`. +The operator owns no persistent state. +It performs no attention, scoring, sorting, or TopK computation. + +All tensors use `int32`. Common inputs are `base_indices [R,C]`, +`base_counts [R]`, and `base_row_indices [N]`. For `append_range`, slots 3 and 4 +are `range_starts [N]` and `range_ends [N]`; slots 5 and 6 must be absent. +For `append_indices`, slots 3 and 4 must be absent, and slots 5 and 6 are +`additional_indices [N,A]` and `additional_counts [N]`. The required +`max_output_entries` attribute is positive and at most `2^30`. + +For query `i`, preserve the first occurrence of each valid base entry from +row `base_row_indices[i]`, then append unseen values from +`[range_starts[i],range_ends[i])` in ascending order. Padding outside the valid +base prefix is ignored. For `append_indices`, append unseen values from the +additional valid prefix in their original order. +Outputs are `selected_indices [N,max_output_entries]`, +`selected_counts [N]`, and `status [N]`. Unused output columns are `-1`. + +Status is `0` on success, `1` for invalid row/count/range metadata or a negative +valid-prefix index, and `2` when the deduplicated union exceeds capacity. On +failure, the entire row is `-1` and its count is zero. Shape, type, and attribute +errors are ordinary session/operator validation errors. WebGPU rejects +`append_indices`. + +Callers must supply one index namespace, ensure causal visibility, and check +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. 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 +`onnxruntime_provider_test`. \ No newline at end of file diff --git a/docs/contrib_ops/cuda/gemm_profiler_cache.md b/docs/contrib_ops/cuda/gemm_profiler_cache.md index 155266287da35..db72ee973c5d0 100644 --- a/docs/contrib_ops/cuda/gemm_profiler_cache.md +++ b/docs/contrib_ops/cuda/gemm_profiler_cache.md @@ -258,9 +258,23 @@ unit-testable. - Sidecars use the same cache prefix scheme: - `".matmulnbits_fpa_intb.tsv"` - `".qmoe_gemm.tsv"` -- Writing is opt-in via session config or env var. Because C++ kernels do not always have a reliable - model path (for example, models loaded from bytes), the core implementation should prefer an - explicit cache prefix and let the offline Python tool derive `.*.tsv` when possible. +- 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`, @@ -289,9 +303,13 @@ unit-testable. |---|---| | `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) diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index cdc8ad9f46a72..fd81528b5564f 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -544,17 +544,7 @@ sessions, and the first-time tuning cost is reduced. Design details: (`{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 buckets are intentionally **not** written to disk -on the hot path (that would put file I/O on inference latency); instead each -lazily profiled bucket is **staged** into the process-global in-memory cache -right after it is profiled, and the cache is written to disk at CUDA EP -teardown (`~CUDAExecutionProvider` for the built-in EP, `~CudaEp` for the CUDA -plugin EP). Disk -persistence covers both the construction-time sweep and lazily discovered buckets, -batched into one flush at teardown. Release the session before inspecting the file. -Only successful tactics are saved; temporary profiling failures are retried in a -new process. Different cache prefixes also receive new buckets from shared-map hits. -Lazy profiling is also skipped +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 @@ -562,17 +552,33 @@ needs. Set `ORT_FPA_INTB_PROFILE_M` (comma-separated) to override the bucket set its maximum also sets the initial profile range. -**Enabling persistence.** Persistence is opt-in. Provide either a directory or an -explicit file prefix, via environment variable or session-config option: +**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`. | - -Resolution order: session-config prefix → session-config dir → env prefix → env -dir. A non-empty prefix wins over a directory, and any session-config value -overrides both env vars. When none is set, nothing is +| `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 @@ -621,6 +627,8 @@ session config entries, creates a CUDA-EP session (which profiles the bucket set 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 @@ -669,9 +677,11 @@ 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_FPA_INTB_PROFILE_M` | comma list, unset | Override the M buckets profiled for the fpA_intB tactic cache (§6.3). 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.3). 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.3); writes `.matmulnbits_fpa_intb.tsv`. Session-config equivalent: `ep.cuda.gemm_tactic_cache_prefix`. | +| `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. | diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md index 353b7a7751e0a..d41eb6620e73f 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). @@ -1060,18 +1081,20 @@ 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. #### Raw-layout memory option @@ -1105,6 +1128,19 @@ the raw kernel, so profiling them only adds synchronization and repeated work. W ### 9b.3 Fast E2M1 → half/bf16 decode +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, @@ -1210,7 +1246,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. @@ -1227,9 +1263,13 @@ 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. + 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`); @@ -1238,7 +1278,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. @@ -1255,7 +1295,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. diff --git a/docs/design/webgpu_matmul_execution_planner.md b/docs/design/webgpu_matmul_execution_planner.md index 38c2b1b085beb..23935af9f5c32 100644 --- a/docs/design/webgpu_matmul_execution_planner.md +++ b/docs/design/webgpu_matmul_execution_planner.md @@ -9,7 +9,7 @@ Make every `MatMulComputeDispatcher` implementation path explicit and independen Introduce `MatMulAlgorithm` with six concrete values: - `SubgroupMatrix`: the common subgroup-matrix implementation in `subgroup_matrix_matmul.cc`. -- `Gemv`: the vendor-neutral FP16 single-row specialization in `MatMulGemvProgram`. +- `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. @@ -23,9 +23,9 @@ 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 single-row FP16 shape, `Naive` when `N < 8 && K < 8`, `PackedSplitK` when the existing `SplitKConfig::UseSplitK` rule succeeds, and otherwise `Packed`. +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 batch size 1, effective `M == 1`, nonzero `K`, FP16 inputs and output, `N` divisible by four, no bias, and no fused activation. 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. +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`. @@ -53,7 +53,7 @@ Configuration is represented by an algorithm-specific variant rather than a bag 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 single-row, 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. +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 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/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 572dbe2aa0a22..c38ecd3a6794f 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -7585,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/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc b/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc index 79aecaf6a572a..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]; diff --git a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h index a5b2be971c8aa..8ed3e7adb01f0 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h +++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h @@ -53,7 +53,11 @@ enum InputIndex : int { kPastStateLengths = 15, // generic: [batch_size, 2], int32 kStateUpdateCaptureCount = 16, // qsa only, optional: [batch_size], int32 kStateUpdateActive = 17, // qsa only, optional: [1], int32 - kInputCount = 18, + kIndexShareMode = 18, + kMergedIndices = 19, + kMergedCounts = 20, + kMergedStatus = 21, + kInputCount = 22, }; // Fixed output slots. present_gate_buffer is declared (with an empty name) but not produced for @@ -66,7 +70,8 @@ enum OutputIndex : int { kPresentGateBuffer = 4, // csa only: same shape as past_gate_buffer kPresentStateLengths = 5, // [batch_size, 2], int32 kStateUpdate = 6, // qsa only, optional: [batch_size, state_update_capacity, head_size] - kOutputCount = 7, + kIndexShareStatus = 7, + kOutputCount = 8, }; // Every PackedSparseAttentionIndexer node declares all 6 fixed outputs; present_gate_buffer is an 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 new file mode 100644 index 0000000000000..f30c837563882 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_merge_common.h @@ -0,0 +1,57 @@ +#pragma once + +#include +#include +#include +#ifndef SHARED_PROVIDER +#include "core/framework/op_kernel.h" +#endif + +namespace onnxruntime::contrib::indexer_merge { + +template +inline int64_t ReadCapacity(const KernelInfo& info) { + std::string policy; + int64_t capacity; + ORT_THROW_IF_ERROR(info.GetAttr("policy_mode", &policy)); + ORT_ENFORCE(policy == "append_range" || policy == "append_indices", "Unknown PackedSparseAttentionIndexerMerge policy"); + ORT_THROW_IF_ERROR(info.GetAttr("max_output_entries", &capacity)); + ORT_ENFORCE(capacity > 0 && capacity <= (int64_t{1} << 30), + "max_output_entries must be in [1, 2^30]"); + return capacity; +} + +struct Dimensions { + int32_t rows; + int32_t capacity; + int32_t queries; + 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) { + ORT_RETURN_IF_NOT(base && counts && rows && starts && ends, "All append_range inputs are required"); + ORT_RETURN_IF_NOT(base->Shape().NumDimensions() == 2 && rows->Shape().NumDimensions() == 1, + "base_indices must be rank 2 and base_row_indices rank 1"); + const int64_t cached_rows = base->Shape()[0]; + const int64_t base_capacity = base->Shape()[1]; + const int64_t queries = rows->Shape()[0]; + ORT_RETURN_IF(cached_rows > INT32_MAX || base_capacity > INT32_MAX || queries > INT32_MAX, + "Selection dimensions must fit int32"); + 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"); + dimensions = {static_cast(cached_rows), static_cast(base_capacity), + static_cast(queries), HashCapacity(output_capacity, base_capacity)}; + return Status::OK(); +} + +} // namespace onnxruntime::contrib::indexer_merge \ No newline at end of file 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/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 2825ca8927721..546b5b12e178f 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -280,6 +280,7 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, SparseAttentionIndexer); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, PackedSparseAttentionIndexer); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, PackedSparseAttentionIndexer); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, PackedSparseAttentionIndexer); +class CUDA_MS_OP_CLASS_NAME(1, PackedSparseAttentionIndexerMerge); class CUDA_MS_OP_THREE_TYPED_CLASS_NAME(1, uint8_t, float, int32_t, GatherBlockQuantized); class CUDA_MS_OP_THREE_TYPED_CLASS_NAME(1, uint8_t, MLFloat16, int32_t, GatherBlockQuantized); @@ -649,6 +650,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { 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 b1f9da27c09d5..459372618401e 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc @@ -337,7 +337,7 @@ onnxruntime::llm::gemm_cache::MatMulNBitsKey WeightOnlyGroupwiseQuantGemmPluginP void WeightOnlyGroupwiseQuantGemmPluginProfiler::loadPersistentCache( GemmIdCore const& gemmId, MProfileMap& map, bool hasWeightOnlyCudaKernel) { - if (mCache == nullptr) { + if (mCache == nullptr || !mCacheAccess.load) { return; } auto key = makeCacheKey(gemmId, hasWeightOnlyCudaKernel); @@ -388,7 +388,9 @@ void WeightOnlyGroupwiseQuantGemmPluginProfiler::loadPersistentCache( continue; } // Do not clobber tactics already selected in-process this session. - map.emplace(m, config); + if (map.emplace(m, config).second) { + markCacheBucketStaged(mCache.get(), gemmId, m); + } ++accepted; } if (accepted != 0) { @@ -407,10 +409,14 @@ bool WeightOnlyGroupwiseQuantGemmPluginProfiler::stageProfiledTactics( if (!config.has_value()) { continue; } - if (mTuningResultsCache != nullptr) { + if (mTuningResultsCache != nullptr && + markCacheBucketStaged(mTuningResultsCache.get(), gemmId, m)) { mTuningResultsCache->Put(key, m, config); } - if (mCache == nullptr) { + 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 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 4498e1eeac91e..fd610aa5174ff 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 @@ -117,9 +119,17 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler } // Attaches the process-global persistent tactic cache. A nullptr keeps the - // in-process-only behavior (no disk reads/writes). - void setPersistentCache(std::shared_ptr cache) { + // 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) { @@ -163,6 +173,7 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler 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; @@ -186,6 +197,7 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler 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 8cc2b150d6c66..9e4df525c4b24 100644 --- a/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h @@ -19,9 +19,12 @@ #include #include +#include #include +#include #include #include +#include #include #include #include @@ -29,6 +32,7 @@ #include #include #include +#include #include #include @@ -222,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; @@ -240,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) { @@ -323,27 +342,52 @@ class GemmPluginProfiler { // 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 persist the - // profiled tactics afterwards. Both are invoked while holding the profile-map writer lock. + // 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*/) {} - // Called after a lazily profiled bucket is inserted: stage the tactics into the in-memory cache - // WITHOUT writing to disk, so any later flush point (e.g. CUDA EP teardown) persists them. 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. - if (mStagedBucketCount.load(std::memory_order_relaxed) != map.size()) { - stagePersistentCache(gemmId, map, mHasWeightOnlyCudaKernel); - mStagedBucketCount.store(map.size(), std::memory_order_relaxed); + // 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); @@ -395,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 @@ -580,7 +644,7 @@ std::optional GemmPluginProfiler #include #include -#include +#include // NOLINT(build/c++17) #include #include +#include #include +#include +#include #include +#include #include "onnxruntime_config.h" @@ -146,19 +151,21 @@ int CurrentProcessId() { // RAII cross-process advisory lock on ".lock". class ScopedFileLock { public: - explicit ScopedFileLock(const std::string& path) : lock_path_(Utf8Path(path + ".lock")) { + 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{}; - locked_ = ::LockFileEx(handle_, LOCKFILE_EXCLUSIVE_LOCK, 0, MAXDWORD, MAXDWORD, &ov) != 0; + 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) { - locked_ = ::flock(fd_, LOCK_EX) == 0; + const int operation = LOCK_EX | (try_only ? LOCK_NB : 0); + locked_ = ::flock(fd_, operation) == 0; } #endif } @@ -360,6 +367,38 @@ std::optional> ParseConfigColumns( 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 && @@ -409,10 +448,23 @@ MatMulNBitsTacticCache::MatMulNBitsTacticCache(std::string file_path, HardwareSi : 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 HardwareSignature& signature, + const std::string& sidecar_prefix) { std::string prefix = config_prefix; std::string dir = config_dir; - // Any session-config value overrides both env vars, so an env prefix cannot shadow a session 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); @@ -465,6 +517,11 @@ MatMulNBitsTacticCache::BucketMap MatMulNBitsTacticCache::GetAll(const MatMulNBi 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()) { @@ -472,20 +529,29 @@ void MatMulNBitsTacticCache::Put(const MatMulNBitsKey& key, int m_bucket, } std::lock_guard guard(mutex_); table_[key][m_bucket] = config; - dirty_rows_[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_); } - return Load(in); + 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) { @@ -620,12 +686,39 @@ onnxruntime::common::Status MatMulNBitsTacticCache::Load(std::istream& in) { for (auto& [key, buckets] : loaded) { auto& dest = table_[key]; for (auto& [m, cfg] : buckets) { - dest.emplace(m, cfg); // Do not clobber entries already set in-process. + 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. @@ -711,8 +804,21 @@ onnxruntime::common::Status MatMulNBitsTacticCache::WriteAllLocked( } 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::lock_guard guard(mutex_); + 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(); } @@ -730,7 +836,7 @@ onnxruntime::common::Status MatMulNBitsTacticCache::Flush() { // 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_); + 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"); } @@ -747,37 +853,65 @@ onnxruntime::common::Status MatMulNBitsTacticCache::Flush() { } decltype(dirty_rows_) pending; { - std::lock_guard guard(mutex_); + 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, cfg] : buckets) { - dest[m] = cfg; + for (const auto& [m, dirty] : buckets) { + dest[m] = dirty.config; } } - dirty_rows_.clear(); } - auto restore_pending = [&]() { - std::lock_guard guard(mutex_); - for (const auto& [key, buckets] : pending) { - for (const auto& [m, cfg] : buckets) { - dirty_rows_[key].emplace(m, cfg); // Keep newer changes made while the write was in flight. + 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); } } - }; - onnxruntime::common::Status status; - try { - status = WriteAllLocked(merged); - } catch (...) { - restore_pending(); - throw; + if (dirty_key_it->second.empty()) { + dirty_rows_.erase(dirty_key_it); + } } - if (!status.IsOK()) { - restore_pending(); + 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; } diff --git a/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.h b/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.h index dcc09d4fdac07..82a79d19e11b5 100644 --- a/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.h +++ b/onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.h @@ -24,6 +24,7 @@ #pragma once #include +#include // NOLINT(build/c++17) #include #include #include @@ -31,6 +32,7 @@ #include #include #include +#include #include #include "core/common/status.h" @@ -52,12 +54,33 @@ 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"; @@ -146,9 +169,10 @@ struct MatMulNBitsKeyHash { 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, env prefix, env dir. + // 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 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 @@ -180,11 +204,38 @@ class MatMulNBitsTacticCache { // 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; @@ -194,7 +245,9 @@ class MatMulNBitsTacticCache { 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. - std::unordered_map dirty_rows_; + 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_gemv_fp4.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu index a23e3635f2f78..7b23d9c5c5cc0 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 @@ -177,6 +177,34 @@ __device__ __forceinline__ float DecodeE4M3Fn(uint8_t code) { return sign ? -value : value; } +template +__device__ __forceinline__ void StoreMoeGemvFp4Pair( + float acc0, float acc1, const T* expert_bias, T* out, int row, int column, int output_columns, + int expert, cutlass_kernels::ActivationParams activation_params) { + if constexpr (EnableBias) { + acc0 += static_cast(expert_bias[column]); + acc1 += static_cast(expert_bias[column + 1]); + } + if constexpr (FusedSwiGlu) { + const float* alpha = activation_params.swiglu_alpha; + const float* beta = activation_params.swiglu_beta; + const float* limit = activation_params.swiglu_limit; + const float activation_alpha = alpha ? alpha[expert] : activation_params.alpha; + const float activation_beta = beta ? beta[expert] : activation_params.beta; + const float activation_limit = limit ? limit[expert] : activation_params.limit; + if (isfinite(activation_limit)) { + acc0 = fminf(acc0, activation_limit); + acc1 = fminf(fmaxf(acc1, -activation_limit), activation_limit); + } + acc1 += activation_beta; + const float sigmoid = 1.0f / (1.0f + expf(-activation_alpha * acc0)); + out[static_cast(row) * (output_columns / 2) + column / 2] = static_cast(acc0 * sigmoid * acc1); + } else { + out[static_cast(row) * output_columns + column] = static_cast(acc0); + out[static_cast(row) * output_columns + column + 1] = static_cast(acc1); + } +} + // NVFP4 schema weights are [E, K, N/2], with adjacent output columns in the two nibbles of // each byte. Transpose an N16/K1024 tile in shared memory so 64 independent K groups can // reuse each block scale for 16 values. Padding distributes compute-time reads across banks. @@ -230,7 +258,8 @@ __global__ void MoeGemvFp4RawNPackedKernel( const int tile_n = static_cast(blockIdx.y) * kNLanes * kColsPerThread; uint2 packed = {}; if (k_tile + tile_row < k) { - const uint8_t* weights = expert_weight + static_cast(k_tile + tile_row) * (n / 2); + const int reduction_column = k_tile + tile_row; + const uint8_t* weights = expert_weight + static_cast(reduction_column) * (n / 2); if (n % (kNLanes * kColsPerThread) == 0) { packed = *reinterpret_cast(weights + tile_n / 2); } else { @@ -318,27 +347,126 @@ __global__ void MoeGemvFp4RawNPackedKernel( acc0 += partials[partial][col][n_lane]; acc1 += partials[partial][col + 1][n_lane]; } - if constexpr (EnableBias) { - acc0 += static_cast(expert_bias[n0]); - acc1 += static_cast(expert_bias[n0 + 1]); + StoreMoeGemvFp4Pair( + acc0, acc1, expert_bias, out, row, n0, n, expert, activation_params); + } +#endif +} + +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, + 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 output_columns, int reduction_size, + 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 = KLanes; + constexpr int kBlockSize = 16; + constexpr int kColsPerBlock = 128 / kKLanes; + static_assert(kKLanes == 4 || kKLanes == 8); + constexpr unsigned kFullMask = 0xffffffffu; + using PackedT = std::conditional_t, half2, __nv_bfloat162>; + + 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; } - if constexpr (FusedSwiGlu) { - const float* alpha = activation_params.swiglu_alpha; - const float* beta = activation_params.swiglu_beta; - const float* limit = activation_params.swiglu_limit; - const float activation_alpha = alpha ? alpha[expert] : activation_params.alpha; - const float activation_beta = beta ? beta[expert] : activation_params.beta; - const float activation_limit = limit ? limit[expert] : activation_params.limit; - if (isfinite(activation_limit)) { - acc0 = fminf(acc0, activation_limit); - acc1 = fminf(fmaxf(acc1, -activation_limit), activation_limit); + break; + } + if (expert < 0 || expert >= num_experts) { + return; + } + + 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[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( + expert_scales[static_cast(column) * (reduction_size / kBlockSize) + + reduction_column / kBlockSize]) * + global_scale); + alignas(4) T scale_pair[2] = {scale, scale}; + 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][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][0] + accumulators[0][1]; +#pragma unroll + for (int offset = kKLanes / 2; offset > 0; offset /= 2) { + accumulator += __shfl_xor_sync(kFullMask, accumulator, offset, kKLanes); + } + const float paired_accumulator = __shfl_xor_sync(kFullMask, accumulator, kKLanes); + if (k_lane == 0 && column < output_columns && column % 2 == 0) { + 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); } - acc1 += activation_beta; - const float sigmoid = 1.0f / (1.0f + expf(-activation_alpha * acc0)); - out[static_cast(row) * (n / 2) + n0 / 2] = static_cast(acc0 * sigmoid * acc1); - } else { - out[static_cast(row) * n + n0] = static_cast(acc0); - out[static_cast(row) * n + n0 + 1] = static_cast(acc1); } } #endif @@ -350,12 +478,47 @@ 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) { + const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, bool weights_row_major) { 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) { + static const bool enable_expert_reuse = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_NVFP4_FC1_EXPERT_REUSE", 0) == 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; + const auto launch_k_packed = [&](auto k_lanes, auto reuse_rows) { + constexpr int kKLanes = decltype(k_lanes)::value; + constexpr bool kReuseRows = decltype(reuse_rows)::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)); + } + }; + if (k <= 1024) { + launch_k_packed(std::integral_constant{}, std::false_type{}); + } else if (reuse_experts) { + launch_k_packed(std::integral_constant{}, std::true_type{}); + } else { + launch_k_packed(std::integral_constant{}, std::false_type{}); + } + return; + } if (bias != nullptr) { MoeGemvFp4RawNPackedKernel<<>>( act, weight, block_scales, global_scales, bias, out, @@ -376,7 +539,8 @@ void LaunchMoeGemvFp4RawNPacked( // selected by `config`, and the per-thread step is StepK = 128 / activation_bits = 8 // (not 128 / weight_bits). bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, - MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed) { + MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed, + bool weights_row_major) { if (sm < 80) { return false; } @@ -386,7 +550,9 @@ bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int if (k % group_size != 0) { return false; } - if (expanded_num_rows <= 0 || expanded_num_rows > kMaxProfiledExpandedRowsFp4) { + const int64_t max_expanded_rows = + raw_n_packed && weights_row_major ? kMaxProfiledExpandedRowsNvfp4KPacked : kMaxProfiledExpandedRowsFp4; + if (expanded_num_rows <= 0 || expanded_num_rows > max_expanded_rows) { return false; } if (n < kMinProfiledProblemDim || k < kMinProfiledProblemDim) { @@ -455,7 +621,7 @@ 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) { + cudaStream_t stream, bool weights_row_major) { ORT_UNUSED_PARAMETER(sm); if (raw_n_packed) { ORT_ENFORCE(group_size == 16, "Raw N-packed FP4 GEMV is NVFP4-only."); @@ -464,7 +630,8 @@ void launch_moe_gemv_fp4_symmetric(const T* act, const uint8_t* weight, const T* LaunchMoeGemvFp4RawNPacked( 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); + expanded_num_rows, n, k, cutlass_kernels::ActivationParams{}, nullptr, expanded_num_rows, stream, + weights_row_major); return; } // Interleaved path: ColumnMajorInterleaved layout + dtype-conditional accumulation + smaller @@ -523,7 +690,7 @@ 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) { + const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, bool weights_row_major) { ORT_UNUSED_PARAMETER(sm); if (raw_n_packed) { ORT_ENFORCE(group_size == 16, "Raw N-packed FP4 GEMV is NVFP4-only."); @@ -532,7 +699,8 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( LaunchMoeGemvFp4RawNPacked( 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); + expanded_num_rows, inter_size * 2, k, activation_params, permuted_row_to_source_row, num_rows, stream, + weights_row_major); return; } // Interleaved path: ColumnMajorInterleaved layout + dtype-conditional accumulation + smaller @@ -588,22 +756,22 @@ 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); + int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, bool, cudaStream_t, 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); + const int*, int64_t, cudaStream_t, 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); + const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, bool, cudaStream_t, 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); + MoeGemvConfig, bool, bool, const int*, int64_t, cudaStream_t, 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 b803c1023f078..ed6d6d72f3a24 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 @@ -30,6 +30,7 @@ enum class MoeGemvConfig { // Cover Qwen-style top_k=8 MTP decode. An (N+1)-token verification for // num_speculative_tokens=N expands to (N+1)*8 rows, up to 64 for N=7. inline constexpr int64_t kMaxProfiledExpandedRowsFp4 = 64; +inline constexpr int64_t kMaxProfiledExpandedRowsNvfp4KPacked = 80; // True when the opt-in interleaved MXFP4 GEMV path is enabled (env ORT_FP4_GEMV_INTERLEAVED=1). // It combines three changes over the default path: (a) the INT4-style ColumnMajorInterleaved FP4 @@ -51,8 +52,9 @@ MoeGemvConfig Fp4MoeGemvDefaultConfig(int64_t expanded_num_rows, int64_t n, int6 int multi_processor_count); // FP4 GEMV shape support. MXFP4 uses the prepacked ColumnMajor or SM80 pair-interleaved layout; -// NVFP4 sets raw_n_packed and directly consumes the schema [E,K,N/2] layout. Requires sm >= 80 -// and the profiled small-decode row/dim bounds. The opt-in interleaved layout is MXFP4-only. +// NVFP4 sets raw_n_packed and directly consumes the schema [E,K,N/2] layout. weights_row_major +// selects native K-packed NVFP4 with an 80-row bound instead of 64. Requires sm >= 80 and the +// profiled small-decode row/dim bounds. The opt-in interleaved layout is MXFP4-only. // See launch_moe_gemv_fp4_symmetric. bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size); bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, @@ -62,7 +64,8 @@ bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int // un-permute it in-register. It implies the interleaved shape rules (MXFP4 group_size 32, // n % 16 == 0, k % 64 == 0) regardless of ORT_FP4_GEMV_INTERLEAVED. bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, - MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed = false); + MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed = false, + bool weights_row_major = false); // Layout-only (row-count independent) form of the rules above, for PrePack: true when a GEMV can // decode an [n, k] MXFP4 problem straight out of the SM80 grouped-GEMM pair-interleaved buffer, so @@ -91,7 +94,7 @@ 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 raw_n_packed, cudaStream_t stream, bool weights_row_major = false); // Launches the MXFP4 MoE GEMV and fuses interleaved SwiGLU activation. // weight/scales/bias use raw FC1 output width n = 2 * inter_size @@ -103,7 +106,8 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, 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); + const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, + bool weights_row_major = false); } // namespace moe_gemv } // namespace kernels diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index 277a182ae2e87..872893f086fc4 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -379,6 +379,8 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE ORT_ENFORCE(quant_type_ != "fp8", "QMoE quant_type='fp8' requires USE_FP8_QMOE with CUDA 11.8 or newer."); ORT_ENFORCE(quant_type_ != "wfp4afp8", "QMoE quant_type='wfp4afp8' requires USE_FP8_QMOE with CUDA 11.8 or newer."); #endif + enable_fp8_fused_ = quant_type_ == "fp8" && sm_ >= 80 && + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_ENABLE_FP8_FUSED", 1) != 0; using namespace onnxruntime::llm::kernels::cutlass_kernels; @@ -675,6 +677,12 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE #endif } // end integer quantization + nvfp4_weights_row_major_ = quant_type_ == "nvfp4" && weights_prepacked_mode == 1; + if (nvfp4_weights_row_major_) { + nvfp4_gemv_raw_layout_ = true; + enable_fp4_gemv_autotune_ = false; + } + ORT_ENFORCE(m_moe_runner != nullptr, "QMoE: failed to construct MoE runner for quant_type='", quant_type_, "', expert_weight_bits=", expert_weight_bits_, @@ -1087,6 +1095,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { : static_cast(moe_params.num_rows) * static_cast(k_); const bool route_native_fp4 = fp4_native_available && + (!nvfp4_weights_row_major_ || (packed_fp4_fc1_block_scales_ != nullptr && + packed_fp4_fc2_block_scales_ != nullptr)) && (fp4_native_max_tokens_per_expert_ <= 0 || avg_tokens_per_expert <= fp4_native_max_tokens_per_expert_) && // Native FP4xFP4 is underfilled for decode (small M) and quantizes activations to NVFP4; // keep NVFP4 decode shapes (num_rows < prefill threshold) off the native runner so they route @@ -1155,8 +1165,11 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { (fc2_gemv_sm80_layout ? gemv_fp4_fc2_weights_ != nullptr : gemv_fp4_fc2_weights_decode_ != nullptr)) : (gemv_fp4_fc1_weights_ != nullptr && gemv_fp4_fc2_weights_ != nullptr)) && (use_raw_nvfp4_gemv - ? (packed_fp4_fc1_block_scales_ != nullptr && packed_fp4_fc2_block_scales_ != nullptr && - packed_fc1_global_scale_ != nullptr && packed_fc2_global_scale_ != nullptr) + ? (nvfp4_weights_row_major_ + ? (context->Input(3) != nullptr && context->Input(6) != nullptr && + context->Input(15) != nullptr && context->Input(16) != nullptr) + : (packed_fp4_fc1_block_scales_ != nullptr && packed_fp4_fc2_block_scales_ != nullptr && + packed_fc1_global_scale_ != nullptr && packed_fc2_global_scale_ != nullptr)) : (gemv_fp4_fc1_scales_ != nullptr && gemv_fp4_fc2_scales_ != nullptr)); bool use_fp4_gemv = false; const bool fp4_gemv_prologue_supported = @@ -1171,10 +1184,11 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { use_fp4_gemv = moe_params.num_rows > 0 && moe_params.num_rows <= 256 && expanded > 0 && gemv::is_moe_gemv_fp4_supported(sm_, expanded, fc1_n, moe_params.hidden_size, gemv_group_size, - gemv::MoeGemvConfig::kDefault, fc1_gemv_sm80_layout, use_raw_nvfp4_gemv) && + gemv::MoeGemvConfig::kDefault, fc1_gemv_sm80_layout, use_raw_nvfp4_gemv, + nvfp4_weights_row_major_) && gemv::is_moe_gemv_fp4_supported(sm_, expanded, moe_params.hidden_size, moe_params.inter_size, gemv_group_size, gemv::MoeGemvConfig::kDefault, - fc2_gemv_sm80_layout, use_raw_nvfp4_gemv); + fc2_gemv_sm80_layout, use_raw_nvfp4_gemv, nvfp4_weights_row_major_); } bool use_packed_int_gemv = false; @@ -1256,6 +1270,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_), @@ -1299,7 +1317,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { // 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) { + if (!use_packed_int && !use_fp8_fused) { std::lock_guard profiler_lock(mGemmProfilerMutex); // Profiling launches grouped-GEMM kernels, records/synchronizes CUDA events, and @@ -1505,6 +1523,123 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { return fused_routing; }; + if (use_fp8_fused) { + Tensor* output = context->Output(0, input->Shape()); + namespace ck = onnxruntime::llm::kernels::cutlass_kernels; + const int num_experts = narrow(moe_params.num_experts); + const int hidden = narrow(moe_params.hidden_size); + const int inter = narrow(moe_params.inter_size); + const int fusion = split_fp8_fc1 ? 1 : (is_fused_swiglu ? static_cast(std::max(1, swiglu_fusion_)) : 0); + 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(); + 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)); + 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) { + return tensor->IsDataType() ? 0 : (tensor->IsDataType() ? 1 : 2); + }; + const auto* fc1_scale_tensor = context->Input(3); + const auto* fc2_scale_tensor = context->Input(6); + const auto* fc3_scale_tensor = split_fp8_fc1 ? context->Input(9) : nullptr; + for (int64_t tile = 0; tile < row_tile_plan.TileCount(); ++tile) { + const int rows = narrow(row_tile_plan.RowsInTile(tile)); + const int routes = rows * narrow(k_); + const int64_t row_offset = row_tile_plan.RowOffset(tile); + 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, routes, stream)); + } +#endif + IAllocatorUniquePtr counts; + IAllocatorUniquePtr cumulative_counts; + IAllocatorUniquePtr blocked_rows; + if (!ck::fusedBuildExpertMapsSortFirstToken( + expert_indices, row_map.get(), unpermuted_row_to_permuted_row, sorted_experts.get(), expert_offsets.get(), + rows, num_experts, narrow(k_), 0, num_experts, stream)) { + const int64_t tokens_per_block = ck::computeNumTokensPerBlock(rows, num_experts); + const int64_t blocks = onnxruntime::llm::common::ceilDiv(rows, tokens_per_block); + counts = GetScratchBuffer(SafeInt(num_experts) * blocks, GetComputeStream(context)); + cumulative_counts = GetScratchBuffer(SafeInt(num_experts) * blocks, GetComputeStream(context)); + blocked_rows = GetScratchBuffer(SafeInt(num_experts) * rows, GetComputeStream(context)); + ck::threeStepBuildExpertMapsSortFirstToken( + expert_indices, sorted_experts.get(), row_map.get(), unpermuted_row_to_permuted_row, expert_offsets.get(), + counts.get(), cumulative_counts.get(), blocked_rows.get(), rows, num_experts, k_, 0, stream); + } + const bool gemm = rows > 8; + if (gemm) { + LaunchQMoEFp8ExpertTiles(expert_offsets.get(), tile_offsets.get(), num_experts, stream); + } + auto execute = [&](auto* type_ptr) { + using T = std::remove_pointer_t; + QMoEFp8ProjectionParams params; + params.weights = static_cast(fc1_experts_weights->DataRaw()); + params.scales = fc1_scale_tensor->DataRaw(); + params.scale_type = scale_type(fc1_scale_tensor); + params.up_weights = split_fp8_fc1 ? static_cast(fc3_experts_weights->DataRaw()) : nullptr; + params.up_scales = fc3_scale_tensor ? fc3_scale_tensor->DataRaw() : nullptr; + params.up_scale_type = fc3_scale_tensor ? scale_type(fc3_scale_tensor) : 0; + params.row_to_unpermuted = row_map.get(); + params.experts = sorted_experts.get(); + params.expert_offsets = expert_offsets.get(); + params.tile_offsets = gemm ? tile_offsets.get() : nullptr; + params.num_experts = num_experts; + params.num_rows = rows; + params.expanded_rows = routes; + params.n = fc1_n; + params.k = hidden; + 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); + params.weights = static_cast(fc2_experts_weights->DataRaw()); + params.scales = fc2_scale_tensor->DataRaw(); + params.scale_type = scale_type(fc2_scale_tensor); + params.up_weights = nullptr; + params.up_scales = nullptr; + params.row_to_unpermuted = nullptr; + params.n = hidden; + params.k = inter; + params.fusion = 0; + LaunchQMoEFp8Projection(params, static_cast(activation.get()), static_cast(result.get()), stream); + const auto* fc2_bias = fc2_experts_bias_optional ? static_cast(fc2_experts_bias_optional->DataRaw()) : nullptr; + ck::finalizeMoeRoutingKernelLauncher( + static_cast(result.get()), static_cast(output->MutableDataRaw()) + row_offset * hidden, + fc2_bias, expert_scales, unpermuted_row_to_permuted_row, row_map.get(), expert_indices, expert_offsets.get(), + rows, hidden, k_, num_experts, parallelism_config, false, stream); + }; + if (is_fp16_) { + execute(static_cast(nullptr)); + } else { + execute(static_cast<__nv_bfloat16*>(nullptr)); + } +#if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) + if (routing_snapshot_) { + ORT_RETURN_IF_ERROR(routing_snapshot_->Consume()); + } +#endif + } + if (enable_kernel_debug_info_) { + const size_t activation_bytes = SafeInt(max_routes) * (fc1_n + 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); + std::cout << "QMoE FP8 ExpertCapacity=" << num_experts << " DequantWeightBytes=0" << std::endl; + } + return Status::OK(); + } + // Holders for packed tensors (if packing is needed for SwiGLU) IAllocatorUniquePtr packed_fc1_scales_holder; IAllocatorUniquePtr packed_fc1_zp_holder; @@ -2088,19 +2223,23 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { : gemv_fp4_fc2_weights_) .get()); const uint8_t* gemv_fc1_raw_block_scales = - is_nvfp4 ? static_cast((gemv_fp4_fc1_block_raw_ ? gemv_fp4_fc1_block_raw_ - : packed_fp4_fc1_block_scales_) - .get()) - : nullptr; + nvfp4_weights_row_major_ ? static_cast(context->Input(3)->DataRaw()) + : is_nvfp4 ? static_cast((gemv_fp4_fc1_block_raw_ ? gemv_fp4_fc1_block_raw_ + : packed_fp4_fc1_block_scales_) + .get()) + : nullptr; const uint8_t* gemv_fc2_raw_block_scales = - is_nvfp4 ? static_cast((gemv_fp4_fc2_block_raw_ ? gemv_fp4_fc2_block_raw_ - : packed_fp4_fc2_block_scales_) - .get()) - : nullptr; + nvfp4_weights_row_major_ ? static_cast(context->Input(6)->DataRaw()) + : is_nvfp4 ? static_cast((gemv_fp4_fc2_block_raw_ ? gemv_fp4_fc2_block_raw_ + : packed_fp4_fc2_block_scales_) + .get()) + : nullptr; const float* gemv_fc1_raw_global_scales = - is_nvfp4 ? static_cast(packed_fc1_global_scale_.get()) : nullptr; + nvfp4_weights_row_major_ ? context->Input(15)->Data() + : (is_nvfp4 ? static_cast(packed_fc1_global_scale_.get()) : nullptr); const float* gemv_fc2_raw_global_scales = - is_nvfp4 ? static_cast(packed_fc2_global_scale_.get()) : nullptr; + nvfp4_weights_row_major_ ? context->Input(16)->Data() + : (is_nvfp4 ? static_cast(packed_fc2_global_scale_.get()) : nullptr); auto p_r2u_buf = GetScratchBuffer(expanded, GetComputeStream(context)); auto p_exp_buf = GetScratchBuffer(expanded, GetComputeStream(context)); auto p_efto_buf = GetScratchBuffer(num_experts + 1, GetComputeStream(context)); @@ -2202,7 +2341,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { gemv_fc1_raw_global_scales, 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); + fc1_gemv_sm80_layout, use_raw_nvfp4_gemv, skip_expand ? p_r2u : nullptr, num_rows, stream, + nvfp4_weights_row_major_); }; auto launch_fc2 = [&](MoeGemvConfig cfg) { gemv::launch_moe_gemv_fp4_symmetric( @@ -2212,7 +2352,7 @@ 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); + fc2_gemv_sm80_layout, use_raw_nvfp4_gemv, stream, nvfp4_weights_row_major_); }; if (do_tune) { @@ -2250,7 +2390,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { float best_fc1 = std::numeric_limits::max(); for (MoeGemvConfig cfg : kCandidates) { if (!gemv::is_moe_gemv_fp4_supported(sm_, expanded, fc1_n, hidden, gemv_group_size, cfg, - fc1_gemv_sm80_layout, is_nvfp4)) { + fc1_gemv_sm80_layout, is_nvfp4, nvfp4_weights_row_major_)) { continue; } const float ms = time_launch([&] { launch_fc1(cfg); }); @@ -2269,7 +2409,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { float best_fc2 = std::numeric_limits::max(); for (MoeGemvConfig cfg : kCandidates) { if (!gemv::is_moe_gemv_fp4_supported(sm_, expanded, hidden, inter, gemv_group_size, cfg, - fc2_gemv_sm80_layout, is_nvfp4)) { + fc2_gemv_sm80_layout, is_nvfp4, nvfp4_weights_row_major_)) { continue; } const float ms = time_launch([&] { launch_fc2(cfg); }); @@ -2342,14 +2482,16 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { // ``session.disable_prepacking`` is set) the repacked buffers stay null and // falling through to the raw initializer bytes would feed a non-CUTLASS // layout to the runner, producing silently wrong output. Fail loudly. - if (packed_fp4_fc1_weights_ == nullptr || packed_fp4_fc2_weights_ == nullptr) { + if (!nvfp4_weights_row_major_ && (packed_fp4_fc1_weights_ == nullptr || packed_fp4_fc2_weights_ == nullptr)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "QMoE native FP4 requires PrePack to run, but the repacked FP4 weight " "buffers were not produced (is session.disable_prepacking set?). " "Enable prepacking to use the native FP4 path."); } - fc1_weight_data = packed_fp4_fc1_weights_.get(); - fc2_weight_data = packed_fp4_fc2_weights_.get(); + if (!nvfp4_weights_row_major_) { + fc1_weight_data = packed_fp4_fc1_weights_.get(); + fc2_weight_data = packed_fp4_fc2_weights_.get(); + } } else if (int_weights_consumed_by_prepack) { // PrePack converted the raw int4/int8 weights to the CUTLASS fpA_intB // layout that the runner consumes and freed the source initializer @@ -2419,8 +2561,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 && @@ -2428,16 +2570,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, @@ -2467,7 +2611,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { if (is_nvfp4) { LaunchQMoEDequantizeNvfp4Weights(weights, block_scales, global_scale, out_h, num_experts, n, k, stream, compact_to_expert.get(), bias ? static_cast(bias->DataRaw()) : nullptr, - static_cast(output_bias)); + static_cast(output_bias), nvfp4_weights_row_major_); } else { LaunchQMoEDequantizeFp4Weights(weights, block_scales, global_scale, out_h, num_experts, n, k, stream); } @@ -2476,7 +2620,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { if (is_nvfp4) { LaunchQMoEDequantizeNvfp4Weights(weights, block_scales, global_scale, out_b, num_experts, n, k, stream, compact_to_expert.get(), bias ? static_cast(bias->DataRaw()) : nullptr, - static_cast<__nv_bfloat16*>(output_bias)); + static_cast<__nv_bfloat16*>(output_bias), nvfp4_weights_row_major_); } else { LaunchQMoEDequantizeFp4Weights(weights, block_scales, global_scale, out_b, num_experts, n, k, stream); } @@ -2714,6 +2858,10 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, bool& is_packed, PrePackedWeights* prepacked_weights) { ORT_UNUSED_PARAMETER(prepacked_weights); is_packed = false; + if (nvfp4_weights_row_major_ && (input_idx == 2 || input_idx == 5 || + (use_fp4_dequant_fallback_ && (input_idx == 3 || input_idx == 6)))) { + return Status::OK(); + } if (quant_type_ == "int" && (input_idx == 3 || input_idx == 6)) { ORT_RETURN_IF_NOT(is_fp16_ ? tensor.IsDataType() : tensor.IsDataType(), "QMoE integer fc", input_idx == 3 ? 1 : 2, "_scales must match the activation type."); @@ -2941,7 +3089,7 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, // SM100 WFP4AFP8 path; byte-level swizzle is format-agnostic). Keep a raw E4M3 copy in // gemv_fp4_fc1_block_raw_ for the GEMV decode and dequant fallback paths. PrePackSwizzleBlockScales(tensor, stream, alloc, packed_fp4_fc1_block_scales_, is_packed); - if (enable_fp4_gemv_ && tensor.Shape().NumDimensions() == 3) { + if (enable_fp4_gemv_ && !nvfp4_weights_row_major_ && tensor.Shape().NumDimensions() == 3) { bool raw_packed = false; PrePackCopyToGpu(tensor, stream, alloc, gemv_fp4_fc1_block_raw_, raw_packed); gemv_fp4_fc1_scale_e_ = tensor.Shape()[0]; @@ -2994,7 +3142,7 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, } if (quant_type_ == "nvfp4" && !use_fp4_dequant_fallback_) { PrePackSwizzleBlockScales(tensor, stream, alloc, packed_fp4_fc2_block_scales_, is_packed); - if (enable_fp4_gemv_ && tensor.Shape().NumDimensions() == 3) { + if (enable_fp4_gemv_ && !nvfp4_weights_row_major_ && tensor.Shape().NumDimensions() == 3) { bool raw_packed = false; PrePackCopyToGpu(tensor, stream, alloc, gemv_fp4_fc2_block_raw_, raw_packed); gemv_fp4_fc2_scale_e_ = tensor.Shape()[0]; diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h index 2fa6e24aa0df5..a6bd4d13154b1 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h @@ -92,6 +92,7 @@ class QMoE final : public CudaKernel, public MoEBase { // layout on every GPU, so -1 and 1 are currently equivalent for the CUDA EP; // 1 is reserved for a possible future Hopper-specific layout (e.g. W4A8). bool weights_prepacked_ = true; + bool nvfp4_weights_row_major_ = 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)`` @@ -104,6 +105,7 @@ class QMoE final : public CudaKernel, public MoEBase { bool use_fp4_dequant_fallback_ = false; // Dequantizes FP8 weights to FP16/BF16 scratch buffers before invoking the A16 MoE runner. bool use_fp8_dequant_fallback_ = false; + bool enable_fp8_fused_ = false; // WFP4AFP8 (W4A8) requires SM100+ (Blackwell) block-scaled tensor ops. On older GPUs we // dequantize MXFP4 weights to FP16/BF16 and run the dense A16 MoE runner. bool use_wfp4afp8_dequant_fallback_ = false; diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu index 558e8985fd4b0..f8b394dae4bfe 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu @@ -6,6 +6,7 @@ #include #include +#include #include #include @@ -1545,6 +1546,195 @@ void LaunchQMoECompactExperts( CUDA_CALL_THROW(cudaGetLastError()); } +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; + int source_n = params.n; + if (params.up_weights) { + if (row % 2) { + weights = params.up_weights; + scales = params.up_scales; + scale_type = params.up_scale_type; + } + source_n /= 2; + row /= 2; + } else if (params.fusion == 2) { + row = row / 2 + (row % 2) * (source_n / 2); + } + 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_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) { + 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; + if (output_col >= params.n) { + return; + } + const int input_row = params.row_to_unpermuted ? params.row_to_unpermuted[row] % params.num_rows : row; + float sum = 0.0f; + 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); + } +} + +__global__ void QMoEFp8ExpertTilesKernel(const int64_t* expert_offsets, int* tile_offsets, int num_experts) { + int tiles = 0; + tile_offsets[0] = 0; + for (int expert = 0; expert < num_experts; ++expert) { + tiles += static_cast((expert_offsets[expert + 1] - expert_offsets[expert] + 15) / 16); + tile_offsets[expert + 1] = tiles; + } +} + +void LaunchQMoEFp8ExpertTiles(const int64_t* expert_offsets, int* tile_offsets, + int num_experts, cudaStream_t stream) { + QMoEFp8ExpertTilesKernel<<<1, 1, 0, stream>>>(expert_offsets, tile_offsets, num_experts); + CUDA_CALL_THROW(cudaGetLastError()); +} + +template +__global__ void QMoEFp8GemmKernel(QMoEFp8ProjectionParams params, const T* input, T* output) { +#if __CUDA_ARCH__ >= 800 + namespace wmma = nvcuda::wmma; + const int tile = static_cast(blockIdx.y); + if (tile >= params.tile_offsets[params.num_experts]) { + return; + } + int first = 0; + int last = params.num_experts; + while (first < last) { + const int middle = (first + last) / 2; + if (params.tile_offsets[middle + 1] <= tile) { + first = middle + 1; + } else { + last = middle; + } + } + const int expert = first; + const int64_t row_start = params.expert_offsets[expert] + (tile - params.tile_offsets[expert]) * 16; + const int64_t row_end = params.expert_offsets[expert + 1]; + const int column_start = static_cast(blockIdx.x) * 64; + __shared__ __align__(32) T activation_tile[16 * 32]; + __shared__ __align__(32) T weight_tile[64 * 32]; + __shared__ __align__(32) float output_tile[16 * 64]; + wmma::fragment accumulator; + wmma::fill_fragment(accumulator, 0.0f); + const int warp = threadIdx.x / 32; + for (int k_start = 0; k_start < params.k; k_start += 32) { + for (int index = threadIdx.x; index < 16 * 32; index += 128) { + const int64_t row = row_start + index / 32; + const int col = k_start + index % 32; + T value = static_cast(0.0f); + if (row < row_end && col < params.k) { + const int input_row = params.row_to_unpermuted ? params.row_to_unpermuted[row] % params.num_rows + : static_cast(row); + value = input[static_cast(input_row) * params.k + col]; + } + activation_tile[index] = value; + } + for (int index = threadIdx.x; index < 64 * 32; index += 128) { + const int row = column_start + index / 32; + const int col = k_start + index % 32; + weight_tile[index] = row < params.n && col < params.k + ? QMoEFp8Weight(params, expert, row, col) + : static_cast(0.0f); + } + __syncthreads(); + for (int inner = 0; inner < 32; inner += 16) { + wmma::fragment activation; + wmma::fragment weight; + wmma::load_matrix_sync(activation, activation_tile + inner, 32); + wmma::load_matrix_sync(weight, weight_tile + warp * 16 * 32 + inner, 32); + wmma::mma_sync(accumulator, activation, weight, accumulator); + } + __syncthreads(); + } + wmma::store_matrix_sync(output_tile + warp * 16, accumulator, 64, wmma::mem_row_major); + __syncthreads(); + for (int index = threadIdx.x; index < 16 * 64; index += 128) { + 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]); + } + } +#endif +} + +template +void LaunchQMoEFp8Projection(const QMoEFp8ProjectionParams& params, + const T* input, T* 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); + } else { + QMoEFp8GemvKernel<<>>(params, input, output); + } + CUDA_CALL_THROW(cudaGetLastError()); +} + +template +__global__ void QMoEFp8ActivationKernel(const T* 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; + if (index >= static_cast(rows) * inter_size) { + return; + } + const int row = static_cast(index / inter_size); + const int col = static_cast(index % inter_size); + const int n = inter_size * (fusion ? 2 : 1); + const int gate_col = fusion ? 2 * col : col; + float gate = static_cast(input[static_cast(row) * n + gate_col]); + float up = fusion ? static_cast(input[static_cast(row) * n + gate_col + 1]) : 1.0f; + if (bias) { + const int bias_col = fusion == 2 ? col : gate_col; + gate += static_cast(bias[static_cast(experts[row]) * n + bias_col]); + if (fusion) { + up += static_cast(bias[static_cast(experts[row]) * n + + (fusion == 2 ? col + inter_size : gate_col + 1)]); + } + } + if (fusion) { + gate = fminf(gate, limit); + up = fminf(fmaxf(up, -limit), limit) + beta; + } + output[index] = static_cast(gate / (1.0f + expf(-alpha * gate)) * up); +} + +template +void LaunchQMoEFp8Activation(const T* 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; + QMoEFp8ActivationKernel<<((count + 255) / 256), 256, 0, stream>>>( + input, bias, output, experts, rows, inter_size, fusion, alpha, beta, limit); + CUDA_CALL_THROW(cudaGetLastError()); +} + +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 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 __global__ void QMoEDequantizeFp8WeightsKernel( const uint8_t* weights, @@ -1674,7 +1864,8 @@ __global__ void QMoEDequantizeNvfp4WeightsKernel( int k, const int* compact_to_expert, const T* bias, - T* output_bias) { + T* output_bias, + bool weights_row_major) { int64_t total = static_cast(num_experts) * n * k; int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (index >= total) { @@ -1692,8 +1883,12 @@ __global__ void QMoEDequantizeNvfp4WeightsKernel( int col = static_cast(offset - static_cast(row) * k); int packed_n = n / 2; - uint8_t packed = packed_weights[(static_cast(expert) * k + col) * packed_n + row / 2]; - uint8_t fp4_code = (row & 1) == 0 ? (packed & 0x0F) : (packed >> 4); + const int64_t packed_index = weights_row_major + ? (static_cast(expert) * n + row) * (k / 2) + col / 2 + : (static_cast(expert) * k + col) * packed_n + row / 2; + const int nibble = weights_row_major ? col : row; + uint8_t packed = packed_weights[packed_index]; + uint8_t fp4_code = (nibble & 1) == 0 ? (packed & 0x0F) : (packed >> 4); constexpr int kNvfp4BlockSize = 16; int scale_k = k / kNvfp4BlockSize; @@ -1717,10 +1912,11 @@ void LaunchQMoEDequantizeNvfp4WeightsImpl( cudaStream_t stream, const int* compact_to_expert, const T* bias, - T* output_bias) { + T* output_bias, + bool weights_row_major) { ORT_ENFORCE(bias == nullptr || output_bias != nullptr, "QMoE NVFP4 bias gathering requires an output buffer."); constexpr int block = 256; - if (!compact_to_expert && !bias && QMoEDequantizeFp4VecApplies<16>(num_experts, n, k)) { + 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, k / kQMoEDequantizeFp4TileK, num_experts); @@ -1733,7 +1929,7 @@ void LaunchQMoEDequantizeNvfp4WeightsImpl( int grid = onnxruntime::narrow((total + block - 1) / block); QMoEDequantizeNvfp4WeightsKernel<<>>( packed_weights, block_scales, global_scales, output, num_experts, n, k, - compact_to_expert, bias, output_bias); + compact_to_expert, bias, output_bias, weights_row_major); CUDA_CALL_THROW(cudaGetLastError()); } @@ -1748,9 +1944,10 @@ void LaunchQMoEDequantizeNvfp4Weights( cudaStream_t stream, const int* compact_to_expert, const half* bias, - half* output_bias) { + half* output_bias, + bool weights_row_major) { LaunchQMoEDequantizeNvfp4WeightsImpl(packed_weights, block_scales, global_scales, output, num_experts, n, k, - stream, compact_to_expert, bias, output_bias); + stream, compact_to_expert, bias, output_bias, weights_row_major); } void LaunchQMoEDequantizeNvfp4Weights( @@ -1764,9 +1961,10 @@ void LaunchQMoEDequantizeNvfp4Weights( cudaStream_t stream, const int* compact_to_expert, const __nv_bfloat16* bias, - __nv_bfloat16* output_bias) { + __nv_bfloat16* output_bias, + bool weights_row_major) { LaunchQMoEDequantizeNvfp4WeightsImpl(packed_weights, block_scales, global_scales, output, num_experts, n, k, - stream, compact_to_expert, bias, output_bias); + stream, compact_to_expert, bias, output_bias, weights_row_major); } // NVFP4 counterpart of QMoECombineFp4ScalesForGemvKernel. Identical [E, n, k_blocks] -> diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h index a0e965a36cb64..7bb9b771a92d2 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h @@ -270,6 +270,38 @@ void LaunchQMoECompactExperts( int64_t num_routes, cudaStream_t stream); +struct QMoEFp8ProjectionParams { + const uint8_t* weights = nullptr; + const uint8_t* up_weights = nullptr; + const void* scales = nullptr; + const void* up_scales = nullptr; + int scale_type = 0; + int up_scale_type = 0; + const int* row_to_unpermuted = nullptr; + const int* experts = nullptr; + const int64_t* expert_offsets = nullptr; + const int* tile_offsets = nullptr; + int num_experts = 0; + int num_rows = 0; + int expanded_rows = 0; + int n = 0; + int k = 0; + int block_size = 0; + int fusion = 0; +}; + +void LaunchQMoEFp8ExpertTiles(const int64_t* expert_offsets, int* tile_offsets, + int num_experts, cudaStream_t stream); + +template +void LaunchQMoEFp8Projection(const QMoEFp8ProjectionParams& params, + const T* input, T* output, cudaStream_t stream); + +template +void LaunchQMoEFp8Activation(const T* 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); + void LaunchQMoEDequantizeFp8Weights( const uint8_t* weights, const float* scales, @@ -314,7 +346,8 @@ void LaunchQMoEDequantizeNvfp4Weights( cudaStream_t stream, const int* compact_to_expert = nullptr, const half* bias = nullptr, - half* output_bias = nullptr); + half* output_bias = nullptr, + bool weights_row_major = false); void LaunchQMoEDequantizeNvfp4Weights( const uint8_t* packed_weights, @@ -327,7 +360,8 @@ void LaunchQMoEDequantizeNvfp4Weights( cudaStream_t stream, const int* compact_to_expert = nullptr, const __nv_bfloat16* bias = nullptr, - __nv_bfloat16* output_bias = nullptr); + __nv_bfloat16* output_bias = nullptr, + bool weights_row_major = false); // Repack column-major FP4 packed weights to row-major layout on GPU. // Input shape interpretation: [experts, k, n/2] (col-major packed), diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index 1678b8698405d..660f73ad0fa10 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -5,12 +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/providers/cpu/math/matmul_helper.h" @@ -129,28 +140,102 @@ struct GlobalTacticCacheRegistry { }; GlobalTacticCacheRegistry& GetGlobalTacticCacheRegistry() { - static GlobalTacticCacheRegistry registry; - return registry; + // 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 the resolved location (creating, loading, and registering it on -// first use). Sessions configured with different cache directories/prefixes each get their own cache. -// Returns nullptr when persistence is not configured, or when the location is already bound to a -// different GPU's signature (e.g. one explicit prefix shared by heterogeneous devices). During a -// session, lazily profiled tactics are only staged into these in-memory caches; the disk write -// happens in FlushMatMulNBitsTacticCaches() at CUDA EP teardown. +// 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( - const std::string& config_dir, const std::string& config_prefix, const cudaDeviceProp& device_prop) { - using onnxruntime::llm::gemm_cache::HardwareSignature; + std::string file_path, onnxruntime::llm::gemm_cache::HardwareSignature signature, bool load) { using onnxruntime::llm::gemm_cache::MatMulNBitsTacticCache; - HardwareSignature signature = HardwareSignature::FromDevice( - device_prop.name, device_prop.major * 10 + device_prop.minor, device_prop.multiProcessorCount); - std::string file_path = MatMulNBitsTacticCache::ResolveFilePath(config_dir, config_prefix, signature); - if (file_path.empty()) { - return nullptr; - } - auto& registry = GetGlobalTacticCacheRegistry(); std::lock_guard lock(registry.mutex); auto it = registry.caches.find(file_path); @@ -161,9 +246,13 @@ static std::shared_ptr Get 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. @@ -171,37 +260,26 @@ static std::shared_ptr Get return cache; } -// Flushes every registered tactic cache to disk. This is the single place lazily-discovered tactics -// reach disk, which keeps file I/O off the inference path. 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. +// 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() { - try { - auto& registry = GetGlobalTacticCacheRegistry(); - std::vector> caches; - { - std::lock_guard lock(registry.mutex); - caches.reserve(registry.caches.size()); - for (const auto& entry : registry.caches) { - caches.push_back(entry.second); - } - } - for (const auto& cache : caches) { - if (cache == nullptr) { - continue; - } - try { - auto status = cache->Flush(); - if (!status.IsOK()) { - ORT_LLM_LOG_WARNING("Failed to flush MatMulNBits gemm tactic cache: " + status.ErrorMessage()); - } - } catch (...) { - // Swallow: cache persistence is best-effort and must not escape EP teardown. - } - } - } catch (...) { - // Registry snapshot allocation must not escape EP teardown either. - } + 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; @@ -610,6 +688,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; @@ -720,18 +833,21 @@ void MatMulNBits::InitGemmProfiler(int sm) { gemmProfiler_->setGroupSize(static_cast(block_size_)); gemmProfiler_->setDecodeInterleave(weight_prepacked_ == kMatMulNBitsWeightPrepackedSm90 ? 1 : 4); - // Resolve the persistent tactic cache location from session config (falls back to env vars). - const auto& config_options = this->Info().GetConfigOptions(); - const std::string cache_dir = - config_options.GetConfigOrDefault(onnxruntime::llm::gemm_cache::kSessionConfigCacheDir, ""); - const std::string cache_prefix = - config_options.GetConfigOrDefault(onnxruntime::llm::gemm_cache::kSessionConfigCachePrefix, ""); - gemmProfiler_->setPersistentCache(GetGlobalMatMulNBitsTacticCache(cache_dir, cache_prefix, this->GetDeviceProp())); + 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 = - config_options.GetConfigOrDefault(onnxruntime::llm::gemm_cache::kSessionConfigTuningResultsPrefix, ""); + info.GetConfigOptions().GetConfigOrDefault(gc::kSessionConfigTuningResultsPrefix, ""); if (!tuning_results_prefix.empty()) { - gemmProfiler_->setTuningResultsCache( - GetGlobalMatMulNBitsTacticCache("", tuning_results_prefix, this->GetDeviceProp())); + gemmProfiler_->setTuningResultsCache(GetGlobalMatMulNBitsTacticCache( + gc::MatMulNBitsTacticCache::ResolveFilePath("", tuning_results_prefix, cache_signature), cache_signature, + /*load*/ false)); } auto allocator = this->Info().GetAllocator(OrtMemType::OrtMemTypeDefault); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h index 0a3e54a55539e..e59deb39341e3 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h @@ -3,6 +3,14 @@ #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 { @@ -11,12 +19,18 @@ namespace cuda { // // 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. In builds without onnxruntime_USE_FPA_INTB_GEMM -// there is no tactic cache and this is a no-op. +// 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 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 4e2bed76767d6..01bbefdf685af 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -31,6 +31,10 @@ namespace psai = onnxruntime::contrib::packed_sparse_attention_indexer; .TypeConstraint("I", DataTypeImpl::GetTensorType()) \ .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ .InputMemoryType(OrtMemTypeCPUInput, psai::kStateUpdateActive) \ + .InputMemoryType(OrtMemTypeCPUInput, psai::kIndexShareMode) \ + .MayInplace(psai::kMergedIndices, psai::kSelectedIndices) \ + .MayInplace(psai::kMergedCounts, psai::kSelectedCounts) \ + .MayInplace(psai::kMergedStatus, psai::kIndexShareStatus) \ .MayInplace(psai::kPastKeyState, psai::kPresentKeyState) \ .MayInplace(psai::kPastKvBuffer, psai::kPresentKvBuffer) \ .MayInplace(psai::kPastGateBuffer, psai::kPresentGateBuffer) \ @@ -147,10 +151,57 @@ PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo ORT_ENFORCE(epsilon_ >= 0.0f, "PackedSparseAttentionIndexer: epsilon must be >= 0, got ", epsilon_); has_scale_ = info.GetAttr("scale", &scale_).IsOK(); head_weight_scale_ = head_weight_scale; + output_capacity_ = info.GetAttrOrDefault("max_output_entries", 0); + const std::string state_policy = info.GetAttrOrDefault("indexshare_state_policy", "default"); + ORT_ENFORCE(state_policy == "default" || state_policy == "hold" || state_policy == "advance", + "indexshare_state_policy must be default, hold, or advance"); + advance_reuse_state_ = state_policy == "advance" || (state_policy == "default" && policy_ == psai::Policy::kCsa); + ORT_ENFORCE(!advance_reuse_state_ || policy_ == psai::Policy::kCsa, "State-only advance requires CSA"); + ORT_ENFORCE(output_capacity_ == 0 || + (output_capacity_ >= psai::SelectedCapacity(policy_, token_budget_, index_topk_, compress_ratio_) && + output_capacity_ <= std::numeric_limits::max()), + "PackedSparseAttentionIndexer: max_output_entries must cover the original selection capacity"); } template Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context) const { + bool reuse = false; + const Tensor* mode = context->InputCount() > psai::kIndexShareMode + ? context->Input(psai::kIndexShareMode) + : nullptr; + if (mode != nullptr) { + ORT_RETURN_IF_ERROR(CheckShape(mode, "indexshare_mode", {1})); + ORT_RETURN_IF(output_capacity_ == 0, "IndexShare requires max_output_entries"); + const int32_t value = mode->Data()[0]; + ORT_RETURN_IF(value != 0 && value != 1, "indexshare_mode must be 0 (refresh) or 1 (reuse)"); + for (int index : {psai::kMergedIndices, psai::kMergedCounts, psai::kMergedStatus}) { + ORT_RETURN_IF(index >= context->InputCount() || context->Input(index) == nullptr, + "All IndexShare merge inputs are required"); + } + const Tensor* merged_indices = context->Input(psai::kMergedIndices); + ORT_RETURN_IF(merged_indices->Shape().NumDimensions() != 2 || merged_indices->Shape()[1] != output_capacity_, + "Merged indices must match max_output_entries"); + 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})); + 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"); + for (int index : {psai::kMergedIndices, psai::kMergedCounts, psai::kMergedStatus}) { + ORT_RETURN_IF(index < context->InputCount() && context->Input(index) != nullptr, + "Merge inputs require indexshare_mode"); + } + } const bool is_qsa = policy_ == psai::Policy::kQsa; constexpr int kCsaOnlyInputs[] = {psai::kGate, psai::kPositionBias, psai::kHeadWeights}; for (int index : kCsaOnlyInputs) { @@ -178,6 +229,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); } @@ -289,7 +341,8 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con ORT_RETURN_IF(present_gate_buffer != nullptr, "PackedSparseAttentionIndexer: output ", psai::kPresentGateBuffer, " must be omitted for policy_mode 'qsa'"); - Tensor* selected_indices = context->Output(psai::kSelectedIndices, TensorShape({total_tokens, capacity})); + const int64_t output_capacity = output_capacity_ == 0 ? capacity : output_capacity_; + Tensor* selected_indices = context->Output(psai::kSelectedIndices, TensorShape({total_tokens, output_capacity})); Tensor* selected_counts = context->Output(psai::kSelectedCounts, TensorShape({total_tokens})); Tensor* present_key_state = context->Output(psai::kPresentKeyState, key_state_shape); Tensor* present_kv_buffer = @@ -324,6 +377,7 @@ 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.output_capacity = static_cast(output_capacity); params.epsilon = epsilon_; params.scale = has_scale_ ? scale_ : 1.0f / std::sqrt(static_cast(head_size)); params.block_topk = static_cast(token_budget_ / compress_ratio_); @@ -364,7 +418,7 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con } template -Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context) const { +Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, const bool state_only) const { using CudaT = typename OrtToCudaType::type; const Tensor* query = context->Input(psai::kQuery); @@ -390,7 +444,9 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context) con "PackedSparseAttentionIndexer: query must have shape (total_tokens, num_heads * head_size), " "got ", query_shape.ToString()); - const int64_t total_tokens = query_shape[0]; + ORT_RETURN_IF(state_only && (key == nullptr || key->Shape().NumDimensions() != 2), + "CSA state-only advance requires rank-2 keys"); + const int64_t total_tokens = state_only ? key->Shape()[0] : query_shape[0]; ORT_RETURN_IF(query_norm_weight == nullptr, "PackedSparseAttentionIndexer: query_norm_weight is required"); const auto& query_norm_shape = query_norm_weight->Shape(); ORT_RETURN_IF_NOT(query_norm_shape.NumDimensions() == 1, @@ -425,7 +481,7 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context) con ORT_RETURN_IF_ERROR(CheckShape(key_norm_weight, "key_norm_weight", {head_size})); ORT_RETURN_IF_ERROR(CheckShape(gate, "gate", {total_tokens, width})); ORT_RETURN_IF_ERROR(CheckShape(position_bias, "position_bias", {compress_ratio_, width})); - ORT_RETURN_IF_ERROR(CheckShape(head_weights, "head_weights", {total_tokens, num_heads})); + ORT_RETURN_IF_ERROR(CheckShape(head_weights, "head_weights", {state_only ? 0 : total_tokens, num_heads})); ORT_RETURN_IF_ERROR(CheckShape(position_ids, "position_ids", {total_tokens})); RotaryCacheShape rotary; @@ -455,7 +511,8 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context) con const int64_t capacity = psai::SelectedCapacity(psai::Policy::kCsa, token_budget_, index_topk_, compress_ratio_); - Tensor* selected_indices = context->Output(psai::kSelectedIndices, TensorShape({total_tokens, capacity})); + const int64_t output_capacity = output_capacity_ == 0 ? capacity : output_capacity_; + Tensor* selected_indices = context->Output(psai::kSelectedIndices, TensorShape({total_tokens, output_capacity})); Tensor* selected_counts = context->Output(psai::kSelectedCounts, TensorShape({total_tokens})); Tensor* present_key_state = context->Output(psai::kPresentKeyState, key_state_shape); Tensor* present_kv_buffer = @@ -484,17 +541,42 @@ 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.output_capacity = static_cast(output_capacity); params.epsilon = epsilon_; params.scale = has_scale_ ? scale_ : 1.0f / std::sqrt(static_cast(head_size)); params.index_topk = static_cast(index_topk_); params.head_weight_scale = has_head_weight_scale_ ? head_weight_scale_ : 1.0f / std::sqrt(static_cast(num_heads)); + params.state_only = state_only; auto float_workspace = GetScratchBuffer(GetCsaPackedWorkspaceFloatCount(params), GetComputeStream(context)); auto overflow_flags = GetScratchBuffer(static_cast(std::max(batch_size, 1)), GetComputeStream(context)); - return LaunchCsaPackedSparseAttentionIndexer( + 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()), reinterpret_cast(key->Data()), @@ -519,7 +601,117 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context) con reinterpret_cast(present_gate_buffer->MutableData()), present_state_lengths->MutableData(), float_workspace.get(), - overflow_flags.get()); + overflow_flags.get())); + if (!state_only) return Status::OK(); + 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), static_cast(state_capacity), + static_cast(compress_ratio_), false, overflow_flags.get()); +} + +template +Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) const { + const Tensor* query = context->Input(psai::kQuery); + const Tensor* lengths = context->Input(psai::kPastSequenceLengths); + const Tensor* indices = context->Input(psai::kMergedIndices); + const Tensor* counts = context->Input(psai::kMergedCounts); + const Tensor* status = context->Input(psai::kMergedStatus); + ORT_RETURN_IF(query == nullptr || query->Shape().NumDimensions() != 2 || query->Shape()[0] != 0, + "Reuse requires zero projected query rows"); + ORT_RETURN_IF(lengths == nullptr || lengths->Shape().NumDimensions() != 1, "Reuse requires request lengths"); + const int64_t batch = lengths->Shape()[0]; + ORT_RETURN_IF_ERROR(CheckIntDimension("batch_size", batch)); + ORT_RETURN_IF_ERROR(CheckShape(indices, "merged_indices", {batch, output_capacity_})); + ORT_RETURN_IF_ERROR(CheckShape(counts, "merged_counts", {batch})); + ORT_RETURN_IF_ERROR(CheckShape(status, "merged_status", {batch})); + ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kCumulativeSequenceLengths), + "cumulative_sequence_lengths", {batch + 1})); + const bool is_qsa = policy_ == psai::Policy::kQsa; + const bool advance = !is_qsa && advance_reuse_state_; + const Tensor* norm = context->Input(psai::kQueryNormWeight); + ORT_RETURN_IF(norm == nullptr || norm->Shape().NumDimensions() != 1 || norm->Shape()[0] <= 0, + "Reuse requires query norm weights"); + const int64_t head_size = norm->Shape()[0]; + ORT_RETURN_IF_ERROR(CheckIntDimension("head_size", head_size, false)); + ORT_RETURN_IF(!is_qsa && head_size > std::numeric_limits::max() / 2, "CSA reuse width exceeds INT_MAX"); + 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})); + ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kPastKvBuffer), "past_kv_buffer", + {batch, psai::GenericBufferCapacity(compress_ratio_), width})); + ORT_RETURN_IF_ERROR(CheckShape(context->Input(psai::kPastStateLengths), "past_state_lengths", {batch, 2})); + 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 { + 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})); + } + const Tensor* active = context->Input(psai::kStateUpdateActive); + if (active != nullptr) { + ORT_RETURN_IF_ERROR(CheckShape(active, "state_update_active", {1})); + } + } + 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}, + std::pair{psai::kPastKeyState, psai::kPresentKeyState}, + std::pair{psai::kPastKvBuffer, psai::kPresentKvBuffer}, + std::pair{psai::kPastStateLengths, psai::kPresentStateLengths}, + std::pair{psai::kPastGateBuffer, psai::kPresentGateBuffer}}) { + if (is_qsa && pair.first == psai::kPastGateBuffer) 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()); + ORT_RETURN_IF(output == nullptr, "Reuse state and selection outputs are required"); + if (input->SizeInBytes() != 0 && output->MutableDataRaw() != input->DataRaw()) { + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output->MutableDataRaw(), input->DataRaw(), input->SizeInBytes(), + cudaMemcpyDeviceToDevice, Stream(context))); + } + } + 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"); + Tensor* snapshot = context->Output(psai::kStateUpdate, {batch, state_update_capacity_, qsa_norm->Shape()[0]}); + if (snapshot != nullptr && snapshot->SizeInBytes() != 0) { + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(snapshot->MutableDataRaw(), 0, snapshot->SizeInBytes(), Stream(context))); + } + } + 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(state_capacity_), static_cast(compress_ratio_), is_qsa); } template class PackedSparseAttentionIndexer; diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h index 112be6e3f46af..9330edc01b0a8 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h @@ -19,7 +19,8 @@ class PackedSparseAttentionIndexer final : public onnxruntime::cuda::CudaKernel private: Status ComputeQsa(OpKernelContext* context) const; - Status ComputeCsa(OpKernelContext* context) const; + Status ComputeCsa(OpKernelContext* context, bool state_only = false) const; + Status ComputeReuse(OpKernelContext* context) const; packed_sparse_attention_indexer::Policy policy_; int64_t compress_ratio_; @@ -27,6 +28,8 @@ class PackedSparseAttentionIndexer final : public onnxruntime::cuda::CudaKernel int64_t state_update_capacity_; int64_t token_budget_; int64_t index_topk_; + int64_t output_capacity_; + bool advance_reuse_state_; float epsilon_; float scale_; float head_weight_scale_; 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 a5ed4610319e9..3304a7839dd3d 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 @@ -467,24 +467,29 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* __shared__ uint64_t keys[kHierarchicalTileBlocks]; __shared__ float head_scores[kHierarchicalTileBlocks][kQwenNumHeads]; - const int64_t total_tiles = static_cast(params.total_tokens) * tile_count; const int warp = static_cast(threadIdx.x) / kWarpSize; const int lane = static_cast(threadIdx.x) % kWarpSize; const int block_slot = warp / kQwenNumHeads; const int head = warp % kQwenNumHeads; - for (int64_t work = blockIdx.x; work < total_tiles; work += gridDim.x) { - const int tile = static_cast(work % tile_count); - const int token = static_cast(work / tile_count); - const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token); - const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength]; - const int64_t abs_position = - params.has_position_ids - ? position_ids[token] - : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]); - const int64_t causal_count = SaiCausalThreshold(abs_position, kQwenCompressRatio); - const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after); - const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0; + const int token = static_cast(blockIdx.y); + const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token); + const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength]; + const int64_t abs_position = + params.has_position_ids + ? position_ids[token] + : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]); + const int64_t causal_count = SaiCausalThreshold(abs_position, kQwenCompressRatio); + const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after); + const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0; + const int active_tiles = (block_count + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks; + const int key_stride = tile_count * kHierarchicalTileBlocks; + const int64_t row_base = static_cast(token) * key_stride; + for (int index = active_tiles * kHierarchicalTileBlocks + blockIdx.x * blockDim.x + threadIdx.x; + index < key_stride; index += gridDim.x * blockDim.x) { + tile_keys[row_base + index] = topk::kPaddingSortKey; + } + for (int tile = static_cast(blockIdx.x); tile < active_tiles; tile += static_cast(gridDim.x)) { const int block_index = tile * kHierarchicalTileBlocks + block_slot; float score = 0.0f; if (block_index < block_count) { @@ -530,7 +535,7 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* } } __syncthreads(); - const int64_t output_base = work * kHierarchicalTileBlocks; + const int64_t output_base = row_base + tile * kHierarchicalTileBlocks; if (threadIdx.x < kHierarchicalTileBlocks) { tile_keys[output_base + threadIdx.x] = keys[threadIdx.x]; } @@ -538,6 +543,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); @@ -558,6 +665,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); @@ -588,7 +736,7 @@ __global__ void QsaEmitHierarchicalTopKKernel(const uint64_t* tile_keys, int32_t* selected_indices, int32_t* selected_counts, PackedSparseAttentionIndexerParams params) { for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) { - int32_t* output = selected_indices + static_cast(token) * params.capacity; + int32_t* output = selected_indices + static_cast(token) * params.output_capacity; const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token); const int64_t abs_position = params.has_position_ids @@ -614,7 +762,7 @@ __global__ void QsaEmitHierarchicalTopKKernel(const uint64_t* tile_keys, output[selected * params.compress_ratio + t] = static_cast(block_start + t); } for (int position = selected * params.compress_ratio + tail_count + threadIdx.x; - position < params.capacity; position += blockDim.x) { + position < params.output_capacity; position += blockDim.x) { output[position] = -1; } if (threadIdx.x == 0) { @@ -641,8 +789,8 @@ __global__ void QsaSelectKernel(const float* block_scores, const int32_t* topk_i int* shared_index = reinterpret_cast(shared + shared_entries); for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) { - int32_t* out_row = selected_indices + static_cast(token) * params.capacity; - for (int p = static_cast(threadIdx.x); p < params.capacity; p += static_cast(blockDim.x)) { + int32_t* out_row = selected_indices + static_cast(token) * params.output_capacity; + for (int p = static_cast(threadIdx.x); p < params.output_capacity; p += static_cast(blockDim.x)) { out_row[p] = -1; } __syncthreads(); @@ -906,6 +1054,8 @@ __global__ void CsaUpdateStateKernel(const T* key, const T* gate, const T* key_n cumulative_sequence_lengths[0] != 0 || 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; @@ -1153,8 +1303,8 @@ __global__ void CsaSelectKernel(const float* scores, const int32_t* cumulative_s int* shared_index = reinterpret_cast(shared + blockDim.x); for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) { - int32_t* out_row = selected_indices + static_cast(token) * params.capacity; - for (int p = static_cast(threadIdx.x); p < params.capacity; p += static_cast(blockDim.x)) { + int32_t* out_row = selected_indices + static_cast(token) * params.output_capacity; + for (int p = static_cast(threadIdx.x); p < params.output_capacity; p += static_cast(blockDim.x)) { out_row[p] = -1; } __syncthreads(); @@ -1207,6 +1357,51 @@ __global__ void CsaSelectKernel(const float* scores, const int32_t* cumulative_s } // namespace +namespace { +__global__ void ValidatePackedReuseKernel(const int32_t* cumulative_lengths, int32_t* indices, + 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; + __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, 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, past_sequence_lengths, + past_state_lengths, batch_size, state_capacity, compress_ratio, is_qsa); + return CUDA_CALL(cudaGetLastError()); +} + size_t GetQsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& params) { const size_t rows = static_cast(params.total_tokens); const size_t query_count = rows * params.num_heads * params.head_size; @@ -1218,6 +1413,7 @@ size_t GetQsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& } size_t GetCsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& params) { + if (params.state_only) return 0; const size_t rows = static_cast(params.total_tokens); return rows * params.num_heads * params.head_size + rows * static_cast(std::max(params.state_capacity, 1)); } @@ -1292,22 +1488,41 @@ Status LaunchQsaPackedSparseAttentionIndexer( const size_t key_count = static_cast(params.total_tokens) * key_stride; uint64_t* merge_input = reinterpret_cast(selection_workspace); uint64_t* merge_output = merge_input + key_count; - const int64_t tile_work = static_cast(params.total_tokens) * tile_count; - QsaScoreTileTopKKernel<<(std::min(tile_work, kSaiMaxGridDimX)), - kHierarchicalScoreThreads, 0, stream>>>( - present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, - present_state_lengths, overflow_flags, merge_input, tile_count, params); - for (int list_width = kHierarchicalTileBlocks; list_width < key_stride; list_width *= 2) { - const int pairs_per_row = (key_stride + 2 * list_width - 1) / (2 * list_width); + if (params.batch_size == 1 && params.total_tokens >= 4) { + const dim3 score_grid(static_cast(std::min(tile_count, 128)), + static_cast((params.total_tokens + 1) / 2)); + QsaPairedScoreTileTopKKernel<<>>( + present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, merge_input, tile_count, params); + } else { + const dim3 score_grid(static_cast(std::min(tile_count, 128)), + static_cast(params.total_tokens)); + QsaScoreTileTopKKernel<<>>( + present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, merge_input, tile_count, params); + } + int compact_stride = key_stride; + int first_list_width = kHierarchicalTileBlocks; + if (params.total_tokens >= 4) { + const int chunks = (key_stride + 2047) / 2048; + compact_stride = chunks * kBoundedTopKMax; + QsaCompactTileTopKKernel<<(chunks), static_cast(params.total_tokens)), + kHierarchicalMergeThreads, 0, stream>>>( + merge_input, merge_output, key_stride, compact_stride); + std::swap(merge_input, merge_output); + first_list_width = kBoundedTopKMax; + } + for (int list_width = first_list_width; list_width < compact_stride; list_width *= 2) { + const int pairs_per_row = (compact_stride + 2 * list_width - 1) / (2 * list_width); const int64_t merge_work = static_cast(params.total_tokens) * pairs_per_row; QsaMergeTileTopKKernel<<(std::min(merge_work, kSaiMaxGridDimX)), kHierarchicalMergeThreads, 0, stream>>>( - merge_input, merge_output, key_stride, list_width, params.total_tokens); + merge_input, merge_output, compact_stride, list_width, params.total_tokens); std::swap(merge_input, merge_output); } QsaEmitHierarchicalTopKKernel<<>>( merge_input, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, - overflow_flags, key_stride, selected_indices, selected_counts, params); + overflow_flags, compact_stride, selected_indices, selected_counts, params); } else if (use_ragged_prefill) { QsaRaggedPrefillScoreKernel<<>>( present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, @@ -1390,7 +1605,7 @@ Status LaunchCsaPackedSparseAttentionIndexer( past_sequence_lengths, past_kv_buffer, past_gate_buffer, past_state_lengths, present_key_state, present_kv_buffer, present_gate_buffer, present_state_lengths, overflow_flags, params); - if (params.total_tokens == 0) { + if (params.total_tokens == 0 || params.state_only) { return CUDA_CALL(cudaGetLastError()); } 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 aa8b346d044c6..81f15253ffd22 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 @@ -31,6 +31,8 @@ struct PackedSparseAttentionIndexerParams { int state_capacity = 0; // past_key_state.shape[1] int buffer_capacity = 0; // past_kv_buffer.shape[1] == 2 * compress_ratio - 1 int capacity = 0; // selected_indices.shape[1] + int output_capacity = 0; + bool state_only = false; float epsilon = 1e-6f; float scale = 0.0f; @@ -46,6 +48,11 @@ struct PackedSparseAttentionIndexerParams { // Scratch requirements, in float elements. 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, 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 new file mode 100644 index 0000000000000..00dc75288b395 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge.cc @@ -0,0 +1,74 @@ +#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" + +namespace onnxruntime::contrib::cuda { + +class PackedSparseAttentionIndexerMerge final : public onnxruntime::cuda::CudaKernel { + public: + explicit PackedSparseAttentionIndexerMerge(const OpKernelInfo& info) + : CudaKernel(info), capacity_(indexer_merge::ReadCapacity(info)) { + std::string policy; + ORT_THROW_IF_ERROR(info.GetAttr("policy_mode", &policy)); + append_indices_ = policy == "append_indices"; + } + + Status ComputeInternal(OpKernelContext* context) const override { + const auto* base = context->Input(0); + const auto* counts = context->Input(1); + const auto* rows = context->Input(2); + const auto* starts = context->Input(3); + const auto* ends = context->Input(4); + const auto* additional = context->Input(5); + const auto* additional_counts = context->Input(6); + indexer_merge::Dimensions dimensions; + if (append_indices_) { + ORT_RETURN_IF(starts || ends, "append_indices does not accept ranges"); + ORT_RETURN_IF_ERROR(indexer_merge::Validate(base, counts, rows, rows, rows, capacity_, dimensions)); + ORT_RETURN_IF(additional == nullptr || additional_counts == nullptr || + additional->Shape().NumDimensions() != 2 || + additional->Shape()[0] != dimensions.queries || + 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)); + } + auto* output = context->Output(0, {dimensions.queries, capacity_}); + auto* output_counts = context->Output(1, {dimensions.queries}); + auto* status = context->Output(2, {dimensions.queries}); + if (dimensions.queries == 0) return Status::OK(); + auto workspace = GetScratchBuffer( + 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(), + additional->Data(), additional_counts->Data(), output->MutableData(), + output_counts->MutableData(), status->MutableData(), workspace.get(), + dimensions.rows, dimensions.capacity, static_cast(additional->Shape()[1]), + dimensions.queries, static_cast(capacity_), dimensions.hash_capacity)); + return Status::OK(); + } + CUDA_RETURN_IF_ERROR(LaunchPackedSparseAttentionIndexerMerge( + Stream(context), base->Data(), counts->Data(), rows->Data(), + starts->Data(), ends->Data(), output->MutableData(), + output_counts->MutableData(), status->MutableData(), workspace.get(), + dimensions.rows, dimensions.capacity, dimensions.queries, static_cast(capacity_), + dimensions.hash_capacity)); + return Status::OK(); + } + + private: + int64_t capacity_; + bool append_indices_; +}; + +ONNX_OPERATOR_KERNEL_EX(PackedSparseAttentionIndexerMerge, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), + PackedSparseAttentionIndexerMerge); + +} // namespace onnxruntime::contrib::cuda \ No newline at end of file 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 new file mode 100644 index 0000000000000..5d5fa04c9517c --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.cu @@ -0,0 +1,269 @@ +#include "packed_sparse_attention_indexer_merge_impl.h" + +#include + +namespace onnxruntime::contrib::cuda { +namespace { + +constexpr int kThreads = 256; + +__device__ int64_t HashSlot(int value, int64_t capacity) { + return (static_cast(value) * 2654435761u) & (capacity - 1); +} + +__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); + } + return -1; +} + +__device__ int InclusiveScan(int* flags, int value) { + const int lane = threadIdx.x; + flags[lane] = value; + __syncthreads(); + for (int stride = 1; stride < kThreads; stride *= 2) { + const int previous = lane >= stride ? flags[lane - stride] : 0; + __syncthreads(); + flags[lane] += previous; + __syncthreads(); + } + return flags[lane]; +} + +__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, int64_t hash_capacity) { + const int query = blockIdx.x; + const int lane = threadIdx.x; + int* output = indices + static_cast(query) * output_capacity; + int* keys = workspace + static_cast(query) * hash_capacity * 2; + int* first = keys + hash_capacity; + __shared__ int flags[kThreads]; + __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; + } + for (int64_t column = lane; column < output_capacity; column += kThreads) output[column] = -1; + __syncthreads(); + + const int row = base_rows[query]; + const int start = starts[query]; + const int end = ends[query]; + if (row < 0 || row >= cached_rows || start < 0 || end < start) { + if (lane == 0) status[query] = 1; + return; + } + const int count = base_counts[row]; + if (count < 0 || count > base_capacity) { + if (lane == 0) status[query] = 1; + return; + } + const int* base = base_indices + static_cast(row) * base_capacity; + for (int64_t column = lane; column < count; column += kThreads) { + if (base[column] < 0) atomicExch(&failed, 1); + } + __syncthreads(); + if (failed) { + if (lane == 0) status[query] = 1; + return; + } + + for (int64_t slot = lane; slot < hash_capacity; slot += kThreads) { + keys[slot] = -1; + first[slot] = INT_MAX; + } + __syncthreads(); + 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); + } + 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 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) { + if (static_cast(emitted) + rank <= output_capacity) output[emitted + rank - 1] = value; + if (value >= start && value < end) atomicAdd(&overlap, 1); + } + __syncthreads(); + if (lane == 0) emitted += flags[kThreads - 1]; + __syncthreads(); + } + + const int64_t total = static_cast(emitted) + static_cast(end) - start - overlap; + if (total > output_capacity) { + for (int64_t column = lane; column < output_capacity; column += kThreads) output[column] = -1; + if (lane == 0) status[query] = 2; + return; + } + for (int64_t chunk = start; chunk < end; chunk += kThreads) { + const int64_t value = chunk + lane; + 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); + __syncthreads(); + if (lane == 0) emitted += flags[kThreads - 1]; + __syncthreads(); + } + if (lane == 0) counts[query] = emitted; +} + +__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, int64_t hash_capacity) { + const int query = blockIdx.x; + const int lane = threadIdx.x; + int* output = indices + static_cast(query) * output_capacity; + int* keys = workspace + static_cast(query) * hash_capacity * 2; + int* first = keys + hash_capacity; + __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 (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]; + if (row < 0 || row >= cached_rows || added_count < 0 || added_count > additional_capacity) { + if (lane == 0) status[query] = 1; + return; + } + const int base_count = base_counts[row]; + if (base_count < 0 || base_count > base_capacity) { + if (lane == 0) status[query] = 1; + return; + } + 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 (int64_t column = lane; column < total; column += kThreads) { + if ((column < base_count ? base[column] : added[column - base_count]) < 0) atomicExch(&failed, 1); + } + __syncthreads(); + if (failed) { + if (lane == 0) status[query] = 1; + return; + } + for (int64_t slot = lane; slot < hash_capacity; slot += kThreads) { + keys[slot] = -1; + first[slot] = INT_MAX; + } + __syncthreads(); + 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); + } + if (!inserted) atomicExch(&failed, 2); + } + __syncthreads(); + if (failed) { + if (lane == 0) status[query] = 2; + return; + } + } + 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 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; + __syncthreads(); + if (lane == 0) emitted += flags[kThreads - 1]; + __syncthreads(); + } + if (emitted > output_capacity) { + 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; + } +} + +} // namespace + +cudaError_t LaunchPackedSparseAttentionIndexerMergeIndices( + cudaStream_t stream, const int32_t* base_indices, const int32_t* base_counts, + 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, int64_t hash_capacity) { + if (query_count == 0) return cudaSuccess; + MergeIndicesKernel<<>>( + base_indices, base_counts, base_rows, additional_indices, additional_counts, + indices, counts, status, workspace, cached_rows, base_capacity, additional_capacity, + output_capacity, hash_capacity); + return cudaGetLastError(); +} + +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, int64_t hash_capacity) { + if (query_count == 0) return cudaSuccess; + MergeRangeKernel<<>>( + base_indices, base_counts, base_rows, starts, ends, indices, counts, status, + workspace, cached_rows, base_capacity, output_capacity, hash_capacity); + return cudaGetLastError(); +} + +} // namespace onnxruntime::contrib::cuda \ No newline at end of file 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 new file mode 100644 index 0000000000000..7a80b33916678 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_merge_impl.h @@ -0,0 +1,22 @@ +#pragma once + +#include +#include + +namespace onnxruntime::contrib::cuda { + +cudaError_t LaunchPackedSparseAttentionIndexerMergeIndices( + cudaStream_t stream, const int32_t* base_indices, const int32_t* base_counts, + 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, 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, int64_t hash_capacity); + +} // namespace onnxruntime::contrib::cuda \ No newline at end of file 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 e147d806b2666..accf0241c1d2e 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc @@ -893,6 +893,8 @@ PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo& i } Status PackedSparseAttentionIndexer::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + ORT_RETURN_IF(context.InputCount() > psai::kIndexShareMode && context.Input(psai::kIndexShareMode), + "PackedSparseAttentionIndexer IndexShare requires the CUDA provider"); const bool is_qsa = policy_ == psai::Policy::kQsa; constexpr int kCsaOnlyInputs[] = {psai::kGate, psai::kPositionBias, psai::kHeadWeights}; for (int index : kCsaOnlyInputs) { 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 new file mode 100644 index 0000000000000..f1fc78f849699 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.cc @@ -0,0 +1,158 @@ +#include + +#include "contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.h" +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "core/providers/webgpu/shader_helper.h" + +namespace onnxruntime::contrib::webgpu { + +ONNX_OPERATOR_KERNEL_EX(PackedSparseAttentionIndexerMerge, kMSDomain, 1, kWebGpuExecutionProvider, + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), + PackedSparseAttentionIndexerMerge); + +Status PackedSparseAttentionIndexerMergeProgram::GenerateShaderCode(ShaderHelper& shader) const { + const auto& base = shader.AddInput("base", ShaderUsage::UseValueTypeAlias); + const auto& counts = shader.AddInput("counts", ShaderUsage::UseValueTypeAlias); + const auto& rows = shader.AddInput("rows", ShaderUsage::UseValueTypeAlias); + const auto& starts = shader.AddInput("starts", ShaderUsage::UseValueTypeAlias); + const auto& ends = shader.AddInput("ends", ShaderUsage::UseValueTypeAlias); + 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); + 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 << "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" + << "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" + << "result_shared = 1u;\n" + << "} else {\n" + << "let count = " << counts.GetByOffset("u32(row)") << ";\n" + << "if (count < 0 || u32(count) > uniforms.base_capacity) {\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" + << "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(); +} + +Status PackedSparseAttentionIndexerMerge::ComputeInternal(ComputeContext& context) const { + const auto* base = context.Input(0); + const auto* counts = context.Input(1); + const auto* rows = context.Input(2); + const auto* starts = context.Input(3); + const auto* ends = context.Input(4); + ORT_RETURN_IF(context.Input(5) || context.Input(6), "append_range does not accept additional index inputs"); + indexer_merge::Dimensions dimensions; + ORT_RETURN_IF_ERROR(indexer_merge::Validate(base, counts, rows, starts, ends, capacity_, dimensions)); + ORT_RETURN_IF(base->Shape().Size() > UINT32_MAX || static_cast(dimensions.queries) * capacity_ > UINT32_MAX, + "WebGPU selection tensor offsets must fit uint32"); + auto* output = context.Output(0, TensorShape({dimensions.queries, capacity_})); + auto* output_counts = context.Output(1, TensorShape({dimensions.queries})); + auto* status = context.Output(2, TensorShape({dimensions.queries})); + if (dimensions.queries == 0) return Status::OK(); + std::optional empty_binding; + if (base->Shape().Size() == 0 || counts->Shape().Size() == 0) { + empty_binding.emplace(context.CreateGPUTensor(DataTypeImpl::GetType(), TensorShape({1}))); + } + 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; + program.SetWorkgroupSize(256) + .SetDispatchGroupSize(dimensions.queries) + .AddInputs({{base_binding, ProgramTensorMetadataDependency::Type}, + {counts_binding, ProgramTensorMetadataDependency::Type}, + {rows, ProgramTensorMetadataDependency::Type}, + {starts, ProgramTensorMetadataDependency::Type}, + {ends, ProgramTensorMetadataDependency::Type}}) + .AddOutputs({{output, ProgramTensorMetadataDependency::Type}, + {output_counts, ProgramTensorMetadataDependency::Type}, + {status, ProgramTensorMetadataDependency::Type}}) + .AddUniformVariables({{static_cast(dimensions.rows)}, + {static_cast(dimensions.capacity)}, + {static_cast(dimensions.queries)}, + {static_cast(capacity_)}}); + return context.RunProgram(program); +} + +} // namespace onnxruntime::contrib::webgpu \ No newline at end of file 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 new file mode 100644 index 0000000000000..404ddef014e2b --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer_merge.h @@ -0,0 +1,35 @@ +#pragma once + +#include "core/providers/webgpu/program.h" +#include "core/providers/webgpu/webgpu_kernel.h" +#include "contrib_ops/cpu/sparse/packed_sparse_attention_indexer_merge_common.h" + +namespace onnxruntime::contrib::webgpu { + +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + +class PackedSparseAttentionIndexerMergeProgram final : public Program { + public: + 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}); +}; + +class PackedSparseAttentionIndexerMerge final : public WebGpuKernel { + public: + explicit PackedSparseAttentionIndexerMerge(const OpKernelInfo& info) + : WebGpuKernel(info), capacity_(indexer_merge::ReadCapacity(info)) { + ORT_ENFORCE(info.GetAttrOrDefault("policy_mode", "append_range") == "append_range", + "WebGPU PackedSparseAttentionIndexerMerge supports only append_range"); + } + Status ComputeInternal(ComputeContext& context) const override; + + private: + int64_t capacity_; +}; + +} // namespace onnxruntime::contrib::webgpu \ No newline at end of file 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/quantization/subgroup_matrix_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc index 2fb980fa5737b..8f6b9f6763108 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc @@ -61,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), @@ -73,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)); @@ -141,7 +146,10 @@ Status SubgroupMatrixMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader } 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_, 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."); @@ -187,6 +195,7 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te // Determine tile sizes first (needed for prepack padding). 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)) { // Default: 64x64 tile size, 256 thread. @@ -208,10 +217,38 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te } } } else if (config.Is(16, 16, 16)) { - // 16x16x16 config: 4 subgroups, 128 threads, 128x128 tiles - tile_m = 128; - tile_n = 128; - workgroup_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, @@ -256,7 +293,9 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te // 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_m != 0); - SubgroupMatrixMatMulNBitsProgram matmul_program{nbits, config, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, tile_m, tile_n, has_tail_buffer}; + 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. @@ -278,8 +317,8 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te 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, - tile_m, tile_n, 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) { matmul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -378,7 +417,17 @@ bool CanApplySubgroupMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& cont 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 2f3c49171b80c..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,16 +20,20 @@ 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, uint32_t tile_m, uint32_t tile_n, 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}, has_weight_idx_indirect_{has_weight_idx_indirect}, - tile_m_(tile_m), - tile_n_(tile_n), has_tail_buffer_{has_tail_buffer} {}; Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( @@ -44,12 +48,13 @@ 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/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index 1574ee7d0bcb3..6a24368bfc61c 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" @@ -49,6 +50,7 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -71,6 +73,7 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, 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/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 0b304dd05d262..fecac2f17f523 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2007,6 +2007,83 @@ ONNX_MS_OPERATOR_SET_SCHEMA( PagedAttentionTypeAndShapeInference(ctx); })); +static OpSchema PackedSparseAttentionIndexerMergeSchema() { + return OpSchema() + .SetDoc( + "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.") + .Attr("policy_mode", "append_range or append_indices.", AttributeProto::STRING) + .Attr("max_output_entries", "Fixed output width, in [1, 2^30].", AttributeProto::INT) + .Input(0, "base_indices", "Cached selections, shape (R, C).", "T") + .Input(1, "base_counts", "Valid base-prefix lengths, shape (R).", "T") + .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", "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") + .TypeConstraint("T", {"tensor(int32)"}, "Integer selection tensors.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + const auto* policy = ctx.getAttribute("policy_mode"); + const auto* capacity = ctx.getAttribute("max_output_entries"); + if (!policy || (policy->s() != "append_range" && policy->s() != "append_indices") || !capacity || capacity->i() <= 0 || + capacity->i() > (int64_t{1} << 30)) { + fail_shape_inference("Invalid PackedSparseAttentionIndexerMerge policy or capacity"); + } + for (size_t output = 0; output < 3; ++output) propagateElemTypeFromInputToOutput(ctx, 0, output); + 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"); + auto* shape = getOutputShape(ctx, 0); + *shape->add_dim() = rows.dim(0); + shape->add_dim()->set_dim_value(capacity->i()); + propagateShapeFromInputToOutput(ctx, 2, 1); + propagateShapeFromInputToOutput(ctx, 2, 2); + } + }); +} + +ONNX_MS_OPERATOR_SET_SCHEMA(PackedSparseAttentionIndexerMerge, 1, PackedSparseAttentionIndexerMergeSchema()); + constexpr const char* SparsePagedAttention_ver1_doc = R"DOC( Selected-index attention over the PagedAttention main K/V cache. @@ -2690,6 +2767,26 @@ 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)); + for (int index : {psai::kMergedIndices, psai::kMergedCounts, psai::kMergedStatus}) { + if (PackedSparseAttentionIndexerHasInput(ctx, index) != has_indexshare) { + fail_shape_inference("PackedSparseAttentionIndexer: all IndexShare inputs must be provided together"); + } + } + 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)) { + fail_shape_inference("PackedSparseAttentionIndexer: IndexShare requires its status output"); + } const int64_t compress_ratio = getAttribute(ctx, "compress_ratio", static_cast(0)); if (compress_ratio <= 0 || @@ -2763,8 +2860,7 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference } const bool has_capture_count = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateCaptureCount); const bool has_state_update_active = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateActive); - const bool has_state_update_output = - ctx.getNumOutputs() > static_cast(psai::kStateUpdate) && ctx.getOutputType(psai::kStateUpdate) != nullptr; + const bool has_state_update_output = ctx.hasOutput(psai::kStateUpdate); if (!is_qsa && (state_update_capacity > 0 || has_capture_count || has_state_update_active || has_state_update_output)) { fail_shape_inference("PackedSparseAttentionIndexer: state update capture is only valid for policy_mode 'qsa'"); @@ -2776,7 +2872,7 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference } if (ctx.getNumOutputs() < static_cast(psai::kFixedOutputCount) || - ctx.getNumOutputs() > static_cast(psai::kOutputCount)) { + ctx.getNumOutputs() > static_cast(has_indexshare ? psai::kOutputCount : psai::kStateUpdate + 1)) { fail_shape_inference("PackedSparseAttentionIndexer: expected ", psai::kFixedOutputCount, " or ", psai::kOutputCount, " declared outputs, got ", ctx.getNumOutputs()); } @@ -2787,6 +2883,7 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference } updateOutputElemType(ctx, psai::kSelectedIndices, ONNX_NAMESPACE::TensorProto_DataType_INT32); updateOutputElemType(ctx, psai::kSelectedCounts, ONNX_NAMESPACE::TensorProto_DataType_INT32); + if (has_indexshare) updateOutputElemType(ctx, psai::kIndexShareStatus, ONNX_NAMESPACE::TensorProto_DataType_INT32); updateOutputElemType(ctx, psai::kPresentStateLengths, ONNX_NAMESPACE::TensorProto_DataType_INT32); propagateElemTypeFromInputToOutput(ctx, psai::kPastKeyState, psai::kPresentKeyState); propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kPresentKvBuffer); @@ -2842,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, @@ -2861,7 +2962,16 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference }; const int64_t buffer_capacity = psai::GenericBufferCapacity(compress_ratio); - require_equal_dims(key_shape, 0, query_shape, 0, "key dimension 0 must equal query dimension 0"); + 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"); require_equal_dims(past_sequence_shape, 0, key_state_shape, 0, @@ -2906,7 +3016,9 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference } } if (!is_qsa) { - require_equal_dims(gate_shape, 0, query_shape, 0, "gate dimension 0 must equal query dimension 0"); + 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"); if (head_weights_shape != nullptr && query_shape != nullptr && @@ -2926,8 +3038,8 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference require_equal_dims(gate_buffer_shape, 2, kv_buffer_shape, 2, "past_gate_buffer and past_kv_buffer widths must match"); } - require_equal_dims(position_ids_shape, 0, query_shape, 0, - "position_ids dimension 0 must equal query dimension 0"); + if (!has_indexshare) require_equal_dims(position_ids_shape, 0, query_shape, 0, + "position_ids dimension 0 must equal query dimension 0"); if (cos_shape != nullptr && sin_shape != nullptr) { if (cos_shape->dim_size() != sin_shape->dim_size()) { fail_shape_inference("PackedSparseAttentionIndexer: cos_cache and sin_cache ranks must match"); @@ -2981,15 +3093,26 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference fail_shape_inference("PackedSparseAttentionIndexer: query width must be divisible by head_size"); } - const int64_t capacity = psai::SelectedCapacity(policy, token_budget, index_topk, compress_ratio); + const int64_t base_capacity = psai::SelectedCapacity(policy, token_budget, index_topk, compress_ratio); + if (has_indexshare && output_capacity < base_capacity) { + fail_shape_inference("PackedSparseAttentionIndexer: max_output_entries is smaller than the selection budget"); + } + const int64_t capacity = has_indexshare ? output_capacity : base_capacity; ONNX_NAMESPACE::TensorShapeProto selected_shape; - SparseAttentionIndexerAppendDim(selected_shape, total_tokens_dim); + if (has_indexshare) + selected_shape.add_dim(); + else + SparseAttentionIndexerAppendDim(selected_shape, total_tokens_dim); selected_shape.add_dim()->set_dim_value(capacity); updateOutputShape(ctx, psai::kSelectedIndices, selected_shape); ONNX_NAMESPACE::TensorShapeProto counts_shape; - SparseAttentionIndexerAppendDim(counts_shape, total_tokens_dim); + if (has_indexshare) + counts_shape.add_dim(); + else + SparseAttentionIndexerAppendDim(counts_shape, total_tokens_dim); updateOutputShape(ctx, psai::kSelectedCounts, counts_shape); + if (has_indexshare) updateOutputShape(ctx, psai::kIndexShareStatus, counts_shape); } // State never grows: present_* always has exactly the same fixed shape as past_*. @@ -3227,6 +3350,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Only for policy_mode 'qsa': optional capture gate with shape (1). A zero value disables capture.", "M", OpSchema::Optional) + .Input(18, "indexshare_mode", "CPU scalar [1]: 0 refresh, 1 reuse.", "M", OpSchema::Optional) + .Input(19, "merged_indices", "Packed merged selection [N,max_output_entries].", "M", OpSchema::Optional) + .Input(20, "merged_counts", "Merged counts [N].", "M", OpSchema::Optional) + .Input(21, "merged_status", "Merged device status [N].", "M", OpSchema::Optional) + .Attr("max_output_entries", "Fixed padded output width for the optional IndexShare interface.", + AttributeProto::INT, static_cast(0)) + .Attr("indexshare_state_policy", "Reuse state policy: default holds QSA and advances CSA; hold or advance overrides it.", + AttributeProto::STRING, std::string("default")) .Output(0, "selected_indices", "Selected entries with shape (total_tokens, capacity). capacity is " @@ -3264,6 +3395,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Inactive and unused slots are zero.", "T", OpSchema::Optional) + .Output(7, "indexshare_status", "Selection status [N], forwarded on reuse and zero on refresh.", + "M", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain floating point tensors to float, float16 and bfloat16.") @@ -5085,6 +5218,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 diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index eed58ac863515..c38c4d1e1936e 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1491,6 +1491,10 @@ constexpr const char* qMoE_ver1_doc = R"DOC( 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. @@ -1597,10 +1601,13 @@ ONNX_MS_OPERATOR_SET_SCHEMA( AttributeProto::STRING, std::string("int")) .Attr("weights_prepacked", - "Only meaningful when quant_type='int'. Tri-state control over the layout of the " - "int4/int8 fc1/fc2 weight initializers. The concrete prepacked layouts selected by " - "-1 and 1 are determined by the execution provider. 0: the initializers are raw, " - "un-prepacked [E, N, K/pack] tensors as produced by quantize_matmul_{4,8}bits. Defaults to -1.", + "Expert-weight layout selector with quantization-specific meanings. For quant_type='int', " + "-1 and 1 select execution-provider-specific prepacked int4/int8 fc1/fc2 layouts. " + "0 selects raw [E,N,K/pack] tensors as produced by quantize_matmul_{4,8}bits. " + "For quant_type='nvfp4' on CUDA, -1 (default) and 0 select legacy N-packed [E,K,N/2] bytes; " + "1 selects K-packed row-major [E,N,K/2] bytes while retaining logical [E,K,N/2] dimensions. " + "NVFP4 block scales and global scales retain their original layouts and values. " + "Ignored for other quantization types. Supported values are -1, 0, and 1; defaults to -1.", AttributeProto::INT, static_cast(-1)) .Input(0, diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index e3ffa851d5d0c..d68cdd392666b 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -95,6 +95,7 @@ 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); @@ -123,6 +124,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipSimplifiedLayerNorma class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SparseAttention); 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); @@ -227,6 +229,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); @@ -258,6 +261,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); 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/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/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 873cb34cd0b78..6c1a9a3195b44 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -578,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 <> @@ -3298,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(); } @@ -3502,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; @@ -3689,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; @@ -3704,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 00fd7007703f1..a9a28070ebf11 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 @@ -130,10 +131,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 e71d65679d8a1..7b15a7b56279c 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_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index 00519584950b8..9bbc212787f22 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -688,10 +688,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/plugin/cuda_ep.cc b/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc index 45ea3abbb52c5..1b6d09ed83df5 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc @@ -259,6 +259,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 3e4a7bb3c0332..bf9f31c8634d1 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 7ea0a4d427955..e00ccaab64d95 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 @@ -547,6 +548,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,6 +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/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/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/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/webgpu/compute_context.h b/onnxruntime/core/providers/webgpu/compute_context.h index bdcbafccace57..576a0c5789964 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 @@ -145,6 +146,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/matmul.cc b/onnxruntime/core/providers/webgpu/math/matmul.cc index 42d76f3aa8331..d5eda34a9a67a 100644 --- a/onnxruntime/core/providers/webgpu/math/matmul.cc +++ b/onnxruntime/core/providers/webgpu/math/matmul.cc @@ -51,6 +51,16 @@ 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 MatMulZeroKProgram final : public Program { @@ -98,30 +108,6 @@ class MatMulZeroKProgram final : public Program { const bool bias_is_scalar_; }; -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 { - 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)); - } - - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"K", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}); -}; - } // namespace static std::string CalcResult(int64_t components, int64_t a_components, int64_t output_number) { @@ -299,6 +285,7 @@ static Status ApplyMatMulGemv(ComputeContext& context, 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; @@ -307,9 +294,9 @@ static Status ApplyMatMulGemv(ComputeContext& context, program.AddInputs({{a, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, {b, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}}) .AddOutput({output_tensor, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten, 4}) - .AddUniformVariables({{k}, {n}}) + .AddUniformVariables({{k}, {n}, {m}}) .SetWorkgroupSize(MatMulGemvProgram::kWorkgroupSizeX, MatMulGemvProgram::kWorkgroupSizeY) - .SetDispatchGroupSize(CeilDiv(n_vec_count, MatMulGemvProgram::kWorkgroupSizeX)); + .SetDispatchGroupSize(CeilDiv(n_vec_count, MatMulGemvProgram::kWorkgroupSizeX) * m); return context.RunProgram(program); } @@ -543,8 +530,12 @@ Status MatMulComputeDispatcher::Compute(ComputeContext& context, 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 = - batch_size == 1 && helper.M() == 1 && + 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() && @@ -554,11 +545,9 @@ Status MatMulComputeDispatcher::Compute(ComputeContext& context, selection_params.m = helper.M(); selection_params.n = helper.N(); selection_params.k = helper.K(); - selection_params.packed_m = folds_batch_into_m - ? logical_a_shape.SizeToDimension(logical_a_shape.NumDimensions() - 1) - : helper.M(); + selection_params.packed_m = packed_m; selection_params.batch_size = batch_size; - selection_params.packed_batch_size = folds_batch_into_m ? 1 : 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(); diff --git a/onnxruntime/core/providers/webgpu/math/matmul.h b/onnxruntime/core/providers/webgpu/math/matmul.h index 86e8c3df0a915..7c2af1c6ec294 100644 --- a/onnxruntime/core/providers/webgpu/math/matmul.h +++ b/onnxruntime/core/providers/webgpu/math/matmul.h @@ -69,5 +69,22 @@ class MatMulNaiveProgram final : public Program { 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 } // 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/nn/conv.h b/onnxruntime/core/providers/webgpu/nn/conv.h index 6684b7ca5259c..af98fa76e30ca 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.h +++ b/onnxruntime/core/providers/webgpu/nn/conv.h @@ -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; 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/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 433ea0d32b886..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 @@ -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; } diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index fc7426537b296..a1ba426a713b9 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -339,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_; } 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/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index a1adfcd9d7c1a..f8a569f26e122 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -1474,10 +1474,8 @@ static std::shared_ptr CreateExecutionProviderFactory cit == provider_options_map.end() ? ProviderOptions{} : cit->second, &session_options); #endif } else if (type == kWebGpuExecutionProvider) { -#if defined(USE_WEBGPU) && !defined(ORT_USE_EP_API_ADAPTERS) +#if defined(USE_WEBGPU) return onnxruntime::WebGpuProviderFactoryCreator::Create(session_options.config_options); -#elif !defined(ORT_MINIMAL_BUILD) - return try_create_registered_plugin_factory(); #endif } else if (type == kCannExecutionProvider) { #ifdef USE_CANN diff --git a/onnxruntime/python/tools/fpa_intb_tune.py b/onnxruntime/python/tools/fpa_intb_tune.py index d86401fa9d5c8..77f9688bb58e3 100644 --- a/onnxruntime/python/tools/fpa_intb_tune.py +++ b/onnxruntime/python/tools/fpa_intb_tune.py @@ -11,10 +11,12 @@ 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 and stages their successful tactics in memory. + 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 is released. + 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. @@ -88,10 +90,11 @@ def _parse_m_values(text: str) -> list[int]: return sorted(set(values)) -def _make_session_options(output_prefix: str, m_values: list[int]) -> ort.SessionOptions: +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)) @@ -237,7 +240,7 @@ def _summarize_cache(cache_path: str, signature: dict[str, str], *, display: boo return entries -def tune(model: str, output_prefix: str, m_values: list[int], run_inference: bool) -> str: +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( @@ -245,7 +248,7 @@ def tune(model: str, output_prefix: str, m_values: list[int], run_inference: boo ) print(f"Creating CUDA session for {model} (this profiles the M buckets)...") - sess_options = _make_session_options(output_prefix, m_values) + 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) @@ -304,6 +307,11 @@ def _build_arg_parser() -> argparse.ArgumentParser: 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 @@ -329,6 +337,7 @@ def main(argv: list[str] | None = None) -> int: output_prefix=output_prefix, m_values=m_values, run_inference=not args.no_inference, + retune=args.retune, ) return 0 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/gemm_tactic_cache_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc index 717f9e183d704..9f807df4eefa5 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc @@ -13,8 +13,9 @@ #include #include -#include +#include // NOLINT(build/c++17) #include +#include #include #include #include @@ -29,6 +30,9 @@ #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" @@ -45,6 +49,21 @@ 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; @@ -675,6 +694,10 @@ TEST(GemmTacticCacheTest, StaleReaderDoesNotOverwriteRetunedBucket) { 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()); @@ -943,6 +966,399 @@ TEST(GemmTacticCacheCudaTest, TemporaryRunnerFailureDoesNotPoisonNextProfiler) { 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 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 ce49b29e76761..a4778180005f3 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -1013,6 +1013,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"; 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/packed_sparse_attention_indexer_merge_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_merge_op_test.cc new file mode 100644 index 0000000000000..e4b88bb789ec5 --- /dev/null +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_merge_op_test.cc @@ -0,0 +1,186 @@ +#include +#include +#include +#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" + +namespace onnxruntime::test { +namespace { + +void RunMerge(int64_t cached_rows, int64_t base_capacity, int64_t output_capacity, + const std::vector& base, const std::vector& counts, + const std::vector& rows, const std::vector& starts, + const std::vector& ends, const std::vector& expected, + const std::vector& expected_counts, const std::vector& status) { + bool tested = false; + for (int provider_kind = 0; provider_kind < 2; ++provider_kind) { + std::unique_ptr provider; + if (provider_kind == 0) provider = DefaultCudaExecutionProvider(); +#ifdef USE_WEBGPU + if (provider_kind == 1) provider = DefaultWebGpuExecutionProvider(); +#endif + if (!provider) continue; + tested = true; + OpTester test("PackedSparseAttentionIndexerMerge", 1, kMSDomain); + test.AddAttribute("policy_mode", std::string("append_range")); + test.AddAttribute("max_output_entries", output_capacity); + const int64_t queries = static_cast(rows.size()); + test.AddInput("base_indices", {cached_rows, base_capacity}, base); + test.AddInput("base_counts", {cached_rows}, counts); + test.AddInput("base_row_indices", {queries}, rows); + test.AddInput("range_starts", {queries}, starts); + test.AddInput("range_ends", {queries}, ends); + test.AddOutput("selected_indices", {queries, output_capacity}, expected); + test.AddOutput("selected_counts", {queries}, expected_counts); + test.AddOutput("status", {queries}, status); + std::vector> providers; + providers.push_back(std::move(provider)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + if (!tested) GTEST_SKIP() << "Requires CUDA or WebGPU"; +} + +} // namespace + +TEST(PackedSparseAttentionIndexerMerge, StableOverlapAndSharedRows) { + RunMerge(2, 6, 7, {9, 1, 9, 3, 1, -1, 6, 6, 2, 7, -1, -1}, {5, 4}, + {1, 0, 0}, {5, 1, 10}, {8, 5, 12}, + {6, 2, 7, 5, -1, -1, -1, 9, 1, 3, 2, 4, -1, -1, 9, 1, 3, 10, 11, -1, -1}, + {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}, + {0, 0, 1, 4, 5, 0, std::numeric_limits::max()}, + std::vector(21, -1), {0, 0, 0, 0, 0, 0, 0}, {1, 1, 1, 1, 2, 1, 2}); +} + +TEST(PackedSparseAttentionIndexerMerge, InvalidCounts) { + RunMerge(2, 2, 3, {1, 2, 3, 4}, {-1, 3}, {0, 1}, {0, 0}, {0, 0}, + std::vector(6, -1), {0, 0}, {1, 1}); +} + +TEST(PackedSparseAttentionIndexerMerge, EmptyBaseAndQueries) { + RunMerge(1, 0, 4, {}, {0}, {0}, {2}, {5}, {2, 3, 4, -1}, {3}, {0}); + RunMerge(1, 2, 4, {1, 2}, {2}, {}, {}, {}, {}, {}, {}); + RunMerge(0, 2, 4, {}, {}, {0}, {0}, {0}, {-1, -1, -1, -1}, {0}, {1}); +} + +TEST(PackedSparseAttentionIndexerMerge, PackedName) { + RunMerge(1, 2, 4, {1, 2}, {2}, {0}, {2}, {4}, {1, 2, 3, -1}, {3}, {0}); +} + +TEST(PackedSparseAttentionIndexerMerge, AppendIndicesStableUnionAndErrors) { + 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{4}); + test.AddInput("base", {2, 3}, {3, 1, 3, 5, 6, 7}); + test.AddInput("counts", {2}, {3, 3}); + test.AddInput("rows", {4}, {0, 1, 0, -1}); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddInput("additional", {4, 3}, {1, 8, 8, 8, 9, 10, -1, 0, 0, 0, 0, 0}); + test.AddInput("additional_counts", {4}, {3, 3, 1, 0}); + test.AddOutput("selected", {4, 4}, {3, 1, 8, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}); + test.AddOutput("selected_counts", {4}, {3, 0, 0, 0}); + test.AddOutput("status", {4}, {0, 2, 1, 1}); + std::vector> providers; + providers.push_back(std::move(provider)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(PackedSparseAttentionIndexerMerge, QwenCapacity) { + std::vector base(2051); + for (int32_t column = 0; column < 2051; ++column) base[column] = column; + auto expected = base; + for (int32_t column = 2051; column < 2058; ++column) expected.push_back(column); + RunMerge(1, 2051, 2058, base, {2051}, {0}, {2051}, {2058}, expected, {2058}, {0}); +} + +TEST(PackedSparseAttentionIndexerMerge, LargeOutputCapacityWithSmallBase) { + std::vector expected(4097, -1); + expected[0] = 7; + expected[1] = 5; + expected[2] = 6; + 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); + test.AddAttribute("policy_mode", std::string("invalid_policy")); + test.AddAttribute("max_output_entries", int64_t{1}); + test.AddInput("base", {1, 1}, {0}); + test.AddInput("counts", {1}, {1}); + test.AddInput("rows", {1}, {0}); + test.AddInput("starts", {1}, {0}); + test.AddInput("ends", {1}, {0}); + test.AddOutput("selected", {1, 1}, {0}); + test.AddOutput("selected_counts", {1}, {1}); + test.AddOutput("status", {1}, {0}); + test.Run(OpTester::ExpectResult::kExpectFailure, "Invalid PackedSparseAttentionIndexerMerge policy or capacity"); +} +#endif + +} // namespace onnxruntime::test \ No newline at end of file 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 6c8b21529e4ca..3ccf50716367c 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 @@ -111,6 +111,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; } @@ -121,6 +129,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{ @@ -139,7 +149,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 { @@ -148,7 +158,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); } @@ -175,18 +185,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) { + 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)); - outputs.push_back(omit_gate_output ? &empty : builder.MakeOutput()); + 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); } @@ -199,6 +221,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; @@ -237,7 +346,7 @@ TEST(PackedSparseAttentionIndexerShapeInferenceTest, QsaInfersStateUpdate) { options.state_update_capacity = 4; options.add_capture_count = true; options.add_state_update_active = true; - options.output_count = psai::kOutputCount; + options.output_count = psai::kStateUpdate + 1; std::unique_ptr model; ASSERT_STATUS_OK(BuildAndResolve([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, model)); @@ -425,13 +534,13 @@ TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsWrongOutputCount) { GraphOptions options; options.output_count = 4; ExpectResolveFailure([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, - "output size 4 not in range [min=6, max=7]"); + "output size 4 not in range [min=6, max=8]"); } TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsStateUpdateCapacityWithoutCaptureCount) { GraphOptions options; options.state_update_capacity = 4; - options.output_count = psai::kOutputCount; + options.output_count = psai::kStateUpdate + 1; ExpectResolveFailure([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, "state_update_capture_count is required"); } @@ -441,7 +550,7 @@ TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsCsaStateUpdateCaptur options.policy_mode = psai::kPolicyModeCsa; options.state_update_capacity = 4; options.add_capture_count = true; - options.output_count = psai::kOutputCount; + options.output_count = psai::kStateUpdate + 1; ExpectResolveFailure([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, "only valid for policy_mode 'qsa'"); } @@ -841,7 +950,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, 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") @@ -859,6 +968,15 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP QsaPackedResult expected; QsaPackedReference(problem, expected); + const int64_t output_capacity = problem.Capacity() + padding; + if (padding > 0) { + std::vector padded(static_cast(problem.TotalTokens() * output_capacity), -1); + for (int token = 0; token < problem.TotalTokens(); ++token) { + std::copy_n(expected.selected_indices.begin() + token * problem.Capacity(), problem.Capacity(), + padded.begin() + token * output_capacity); + } + expected.selected_indices = std::move(padded); + } const int64_t total_tokens = problem.TotalTokens(); const int64_t batch_size = problem.batch_size; @@ -919,7 +1037,16 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP } } - test.AddOutput("selected_indices", {total_tokens, problem.Capacity()}, expected.selected_indices); + if (padding > 0) { + test.AddAttribute("max_output_entries", output_capacity); + if (problem.state_update_capacity == 0) test.AddOptionalInputEdge(); + if (!problem.state_update_active.has_value()) test.AddOptionalInputEdge(); + test.AddInput("mode", {1}, {0}); + test.AddInput("merged", {0, output_capacity}, {}); + test.AddInput("merged_counts", {0}, {}); + test.AddInput("merged_status", {0}, {}); + } + test.AddOutput("selected_indices", {total_tokens, output_capacity}, expected.selected_indices); test.AddOutput("selected_counts", {total_tokens}, expected.selected_counts); test.AddOutput("present_key_state", {batch_size, problem.state_capacity, head_size}, ToElementType(expected.present_key_state), false, 0.0f, tolerance); @@ -931,6 +1058,10 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP test.AddOutput("state_update", {batch_size, problem.state_update_capacity, head_size}, ToElementType(expected.state_update), false, 0.0f, tolerance); } + if (padding > 0) { + if (problem.state_update_capacity == 0) test.AddOptionalOutputEdge(); + test.AddOutput("status", {total_tokens}, std::vector(static_cast(total_tokens), 0)); + } RunOnProvider(test, std::move(provider)); if (actual != nullptr) { const auto& fetches = test.GetFetches(); @@ -953,6 +1084,63 @@ TEST(PackedSparseAttentionIndexerTest, QsaPackedQueryKey) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(), ProviderKind::Cuda, nullptr, true); } +TEST(PackedSparseAttentionIndexerTest, IndexShareRefreshPreservesBudgetAndPadsRows) { + RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(), ProviderKind::Cuda, nullptr, true, nullptr, 6); +} + +TEST(PackedSparseAttentionIndexerTest, IndexShareReusePackedQkAndStatus) { + 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) { if (DefaultCudaExecutionProvider() == nullptr) { GTEST_SKIP() << "CUDA execution provider is not available"; @@ -1067,6 +1255,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; @@ -1427,7 +1672,8 @@ CsaPackedProblem MakeCsaPackedProblem(CsaPackedProblem problem = {}) { template void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, ProviderKind provider_kind = ProviderKind::Cuda, CsaPackedResult* actual = nullptr, - 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") @@ -1456,6 +1702,23 @@ void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, const int64_t head_size = problem.head_size; const int64_t width = problem.Width(); const int64_t buffer_capacity = problem.BufferCapacity(); + const int64_t output_capacity = problem.index_topk + (reuse ? 6 : 0); + std::vector merged; + if (reuse) { + ASSERT_EQ(total_tokens, batch_size); + merged.assign(static_cast(batch_size * output_capacity), -1); + for (int64_t row = 0; row < batch_size; ++row) merged[static_cast(row * output_capacity)] = 0; + expected.selected_indices = merged; + expected.selected_counts.assign(static_cast(batch_size), 1); + if (reuse_failure) { + expected.selected_indices.assign(static_cast(batch_size * output_capacity), -1); + expected.selected_counts.assign(static_cast(batch_size), 0); + expected.present_key_state = problem.past_key_state; + expected.present_kv_buffer = problem.past_kv_buffer; + expected.present_gate_buffer = problem.past_gate_buffer; + expected.present_state_lengths = problem.past_state_lengths; + } + } OpTester test("PackedSparseAttentionIndexer", 1, onnxruntime::kMSDomain); test.AddAttribute("policy_mode", std::string(psai::kPolicyModeCsa)); @@ -1465,7 +1728,8 @@ void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, if (problem.scale.has_value()) test.AddAttribute("scale", *problem.scale); if (problem.head_weight_scale.has_value()) test.AddAttribute("head_weight_scale", *problem.head_weight_scale); - test.AddInput("query", {total_tokens, problem.num_heads * head_size}, ToElementType(problem.query)); + test.AddInput("query", {reuse ? 0 : total_tokens, problem.num_heads * head_size}, + reuse ? std::vector{} : ToElementType(problem.query)); test.AddInput("key", {total_tokens, width}, ToElementType(problem.key)); test.AddInput("query_norm_weight", {head_size}, ToElementType(problem.query_norm_weight)); test.AddInput("key_norm_weight", {head_size}, ToElementType(problem.key_norm_weight)); @@ -1475,7 +1739,8 @@ void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, test.AddInput("past_sequence_lengths", {batch_size}, problem.past_sequence_lengths); test.AddInput("gate", {total_tokens, width}, ToElementType(problem.gate)); test.AddInput("position_bias", {problem.compress_ratio, width}, ToElementType(problem.position_bias)); - test.AddInput("head_weights", {total_tokens, problem.num_heads}, ToElementType(problem.head_weights)); + test.AddInput("head_weights", {reuse ? 0 : total_tokens, problem.num_heads}, + reuse ? std::vector{} : ToElementType(problem.head_weights)); test.AddInput("position_ids", {total_tokens}, problem.position_ids); test.AddInput("past_key_state", {batch_size, problem.state_capacity, head_size}, ToElementType(problem.past_key_state)); @@ -1484,7 +1749,16 @@ void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, ToElementType(problem.past_gate_buffer)); test.AddInput("past_state_lengths", {batch_size, 2}, problem.past_state_lengths); - test.AddOutput("selected_indices", {total_tokens, problem.index_topk}, expected.selected_indices); + if (reuse) { + test.AddAttribute("max_output_entries", output_capacity); + test.AddOptionalInputEdge(); + 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), 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); test.AddOutput("present_key_state", {batch_size, problem.state_capacity, head_size}, ToElementType(expected.present_key_state), false, 0.0f, tolerance); @@ -1493,6 +1767,11 @@ void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, test.AddOutput("present_gate_buffer", {batch_size, buffer_capacity, width}, ToElementType(expected.present_gate_buffer), false, 0.0f, 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), merged_status != 0 ? merged_status : reuse_failure ? 1 + : 0)); + } RunOnProvider(test, std::move(provider)); if (actual != nullptr) { const auto& fetches = test.GetFetches(); @@ -1508,6 +1787,38 @@ void RunCsaPackedTest(const CsaPackedProblem& base, float tolerance, TEST(PackedSparseAttentionIndexerTest, CsaFloat) { RunCsaPackedTest(MakeCsaPackedProblem(), 1.0e-5f); } +TEST(PackedSparseAttentionIndexerTest, IndexShareCsaAdvanceWithoutScoring) { + 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); + RunCsaPackedTest(problem, 4.0e-3f, ProviderKind::Cuda, nullptr, nullptr, true); +} + +TEST(PackedSparseAttentionIndexerTest, IndexShareCsaRejectsOffsetsAndStateOverflow) { + CsaPackedProblem problem; + 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, nullptr, true, true); + problem.cumulative_sequence_lengths = {0, 1, 2}; + problem.state_capacity = 1; + 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); } TEST(PackedSparseAttentionIndexerTest, CsaBFloat16) { RunCsaPackedTest(MakeCsaPackedProblem(), 3.0e-2f); } diff --git a/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc b/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc index 2d202dfdf62a4..5a5561e4eb460 100644 --- a/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc +++ b/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc @@ -326,6 +326,42 @@ 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(); 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/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/optimizer/nchwc_optimizer_test.cc b/onnxruntime/test/optimizer/nchwc_optimizer_test.cc index f9bdc00b03f8e..ee686fad94ed7 100644 --- a/onnxruntime/test/optimizer/nchwc_optimizer_test.cc +++ b/onnxruntime/test/optimizer/nchwc_optimizer_test.cc @@ -259,27 +259,39 @@ TEST(NchwcOptimizerTests, ConvWithZeroOutputChannelsIsNotTransformed) { } 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"); - auto build_test_case = [&](NchwcTestHelper& helper) { - 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})); - } - helper.AddNode("Conv", input_args, {helper.MakeOutput()}); - }; + 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 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); - }; + 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); - NchwcOptimizerTester(build_test_case, check_nchwc_graph, 18); + 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); } } 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/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/webgpu/dispatch_normalization_test.cc b/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc index 6cfadb642c0a5..a197551015f3e 100644 --- a/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc +++ b/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc @@ -17,6 +17,7 @@ #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" @@ -107,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()); @@ -122,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); @@ -357,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/math/matmul_test.cc b/onnxruntime/test/providers/webgpu/math/matmul_test.cc index 89cc941cb6ca4..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" @@ -304,6 +308,155 @@ TEST(MathOpTest, MatMulSubgroupMatrix) { 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/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/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index dd6888c18de3b..d0ba4ead1b8ca 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -47,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 { @@ -386,51 +387,13 @@ Status StartPipelineCallbackTestBuild(webgpu::WebGpuContext& context, webgpu::Pe 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; @@ -448,7 +411,7 @@ 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::DispatchCommandInfo dispatch; @@ -496,7 +459,7 @@ 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::DispatchCommandInfo dispatch; dispatch.program_key = "post-context-pipeline-callback"; @@ -531,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; @@ -558,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 @@ -578,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")); @@ -597,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); @@ -631,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) { 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/onnxruntime_test_python_fpa_intb_tune.py b/onnxruntime/test/python/onnxruntime_test_python_fpa_intb_tune.py index 3bce450c8f74d..48f5c343ea976 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_fpa_intb_tune.py +++ b/onnxruntime/test/python/onnxruntime_test_python_fpa_intb_tune.py @@ -547,6 +547,41 @@ def run(name, chunk=0): 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": 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..80862cbf88b6b --- /dev/null +++ b/onnxruntime/test/python/transformers/benchmark_qwen38_mtp_kernels.py @@ -0,0 +1,724 @@ +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): + match = re.search(r"__device__[^\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 + return source[match.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"), 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) + else: + run_nvfp4(baseline, args.check_only, args.sanitizer_case) + + +def compile_nvfp4(source, element_type, fused, bias, wide, reuse=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) + kernel_start = source.index("__global__ void MoeGemvFp4RawKPackedKernel") + kernel_start = source.rindex("template <", 0, kernel_start) + kernel_end = source.index("template \nvoid LaunchMoeGemvFp4RawNPacked", kernel_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" + body = ( + device_function(source, "DecodeE4M3Fn") + "\n" + source[pair_start:pair_end] + source[kernel_start:kernel_end] + ) + flags = f"{str(fused).lower()},{str(bias).lower()}" + 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) + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py index e314b1d61b45f..751ef131ef0b9 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py @@ -356,16 +356,27 @@ def _inputs( return tensors @staticmethod - def _session(tensors, block=128, fusion=0, top_k=10, normalize=1, initializers=False, row_tile_size=0): + def _session( + tensors, + block=128, + fusion=0, + top_k=10, + normalize=1, + initializers=False, + row_tile_size=0, + enable_cuda_graph=False, + ): dtype = TensorProto.BFLOAT16 if tensors["input"].dtype == torch.bfloat16 else TensorProto.FLOAT16 model, input_types = create_block_fp8_moe_graph(tensors, top_k, dtype, block, fusion, normalize, initializers) options = onnxruntime.SessionOptions() options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") options.add_session_config_entry("ep.cuda.qmoe_row_tile_size", str(row_tile_size)) - return onnxruntime.InferenceSession( - model, options, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] - ), input_types + provider = resolve_cuda_plugin_ep("CUDAExecutionProvider") + if enable_cuda_graph: + provider_name, provider_options = provider if isinstance(provider, tuple) else (provider, {}) + provider = (provider_name, {**provider_options, "enable_cuda_graph": "1"}) + return onnxruntime.InferenceSession(model, options, providers=[provider]), input_types @classmethod def _execute(cls, tensors, block=128, fusion=0, top_k=10, normalize=1, session=None, initializers=False): @@ -606,6 +617,7 @@ def test_block_fp8_qwen38_flash_next_official_shape(self): @pytest.mark.skipif(not torch.cuda.is_available() or not has_fp8_qmoe, reason="CUDA FP8 QMoE required") def test_block_fp8_compact_decode_scratch(capfd, monkeypatch): monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + monkeypatch.setenv("ORT_ENABLE_FP8_FUSED", "0") tensors = TestQMoEBlockFP8._inputs(experts=512, tokens=1) tensors["router_probs"] = (torch.randperm(512, device=device).float() / 128 - 2).bfloat16().unsqueeze(0) actual = TestQMoEBlockFP8._execute(tensors) @@ -614,5 +626,81 @@ def test_block_fp8_compact_decode_scratch(capfd, monkeypatch): assert "QMoE FP8 ExpertCapacity=10 DequantWeightBytes=768000" in capfd.readouterr().out +@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("tokens,path", [(1, "fp8_fused_gemv"), (32, "fp8_fused_gemm")]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("scale_dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_block_fp8_fused_without_weight_scratch(capfd, monkeypatch, tokens, path, dtype, scale_dtype): + tensors = TestQMoEBlockFP8._inputs( + experts=512 if tokens == 1 else 16, tokens=tokens, dtype=dtype, scale_dtype=scale_dtype + ) + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + monkeypatch.setenv("ORT_ENABLE_FP8_FUSED", "0") + reference = TestQMoEBlockFP8._execute(tensors) + capfd.readouterr() + monkeypatch.setenv("ORT_ENABLE_FP8_FUSED", "1") + actual = TestQMoEBlockFP8._execute(tensors) + torch.testing.assert_close(actual.float(), reference.float(), atol=0.003, rtol=0.04) + log = capfd.readouterr().out + assert path in log + 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("fusion", [0, 1, 2]) +@pytest.mark.parametrize("row_tile_size", [0, 17]) +def test_block_fp8_fused_gemm_partial_tiles(monkeypatch, dtype, fusion, row_tile_size): + monkeypatch.setenv("ORT_ENABLE_FP8_FUSED", "1") + tensors = TestQMoEBlockFP8._inputs(hidden=72, inter=40, experts=32, tokens=35, block=32, fusion=fusion, dtype=dtype) + tensors["router_probs"].fill_(-10) + selected = [31, 7, 2, 19, 25, 0, 17, 11, 9, 23] + tensors["router_probs"][:, selected] = torch.linspace(2, 1, 10, device=device, dtype=dtype) + tensors["fc2_bias"] = torch.randn(32, 72, device=device, dtype=dtype) * 0.01 + if fusion == 1: + tensors["fc1_bias"] = torch.randn(32, 80, device=device, dtype=dtype) * 0.01 + session = TestQMoEBlockFP8._session(tensors, block=32, fusion=fusion, row_tile_size=row_tile_size) + actual = TestQMoEBlockFP8._execute(tensors, block=32, fusion=fusion, session=session) + expected = TestQMoEBlockFP8._reference(tensors, 32, fusion) + torch.testing.assert_close(actual.float(), expected.float(), atol=0.003, rtol=0.04) + + +@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("tokens", [1, 32]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_block_fp8_fused_cuda_graph_replay(capfd, monkeypatch, tokens, dtype): + monkeypatch.setenv("ORT_ENABLE_FP8_FUSED", "1") + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + tensors = TestQMoEBlockFP8._inputs(experts=32, tokens=tokens, dtype=dtype) + session, input_types = TestQMoEBlockFP8._session(tensors, enable_cuda_graph=True) + output = torch.empty_like(tensors["input"]) + binding = session.io_binding() + for graph_input in session.get_inputs(): + tensor = tensors[graph_input.name] + binding.bind_input(graph_input.name, "cuda", 0, input_types[graph_input.name], tensor.shape, tensor.data_ptr()) + binding.bind_output("output", "cuda", 0, input_types["input"], output.shape, output.data_ptr()) + for selected in ([31, 7, 2, 19, 25, 0, 17, 11, 9, 23], [1, 4, 6, 8, 10, 12, 14, 16, 18, 20]): + tensors["router_probs"].fill_(-10) + tensors["router_probs"][:, selected] = torch.linspace(2, 1, 10, device=device, dtype=dtype) + tensors["fc2_scales"].mul_(0.5) + for _replay in range(2): + output.zero_() + torch.cuda.synchronize() + session.run_with_iobinding(binding) + binding.synchronize_outputs() + expected = TestQMoEBlockFP8._reference(tensors, 128, 0) + torch.testing.assert_close(output.float(), expected.float(), atol=0.003, rtol=0.04) + assert "DequantWeightBytes=0" in capfd.readouterr().out + + if __name__ == "__main__": unittest.main() diff --git a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py index dc20d461c8814..cf7f0c762e5f4 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py @@ -179,6 +179,7 @@ def create_nvfp4_moe_onnx_graph( use_swiglu=False, fc1_bias=None, fc2_bias=None, + weights_prepacked=None, ): """Build ONNX model with QMoE operator in NVFP4 mode.""" inputs = [ @@ -223,6 +224,9 @@ def create_nvfp4_moe_onnx_graph( ), ] + if weights_prepacked is not None: + nodes[0].attribute.append(helper.make_attribute("weights_prepacked", weights_prepacked)) + initializers = [] # FC1 / FC2 packed weights [E, K, N/2] uint8 @@ -325,6 +329,9 @@ def _run_nvfp4_moe_test( row_tile_size=0, use_bias=False, rtol_override=0.0, + offline_prepacked=False, + enable_cuda_graph=False, + weights_prepacked=None, ): self._skip_if_no_fp4() @@ -359,6 +366,13 @@ def _run_nvfp4_moe_test( fc1_weights = torch.stack(fc1_packed, dim=0) # [E, K, N/2] fc2_weights = torch.stack(fc2_packed, dim=0) # [E, K, N/2] + if offline_prepacked: + packed_weights = [] + for weights in (fc1_weights, fc2_weights): + codes = torch.stack((weights & 15, weights >> 4), dim=-1) + codes = codes.reshape(weights.shape[0], weights.shape[1], -1).transpose(1, 2) + packed_weights.append(((codes[..., 0::2] & 15) | (codes[..., 1::2] << 4)).reshape(weights.shape)) + fc1_weights, fc2_weights = packed_weights fc1_block_scales = torch.stack(fc1_bs, dim=0) # [E, N, K//16] fc2_block_scales = torch.stack(fc2_bs, dim=0) # [E, N, K//16] fc1_global_scale = torch.stack(fc1_gs) # [E] @@ -386,6 +400,7 @@ def _run_nvfp4_moe_test( use_swiglu=use_swiglu, fc1_bias=fc1_bias, fc2_bias=fc2_bias, + weights_prepacked=1 if offline_prepacked else weights_prepacked, ) opts = onnxruntime.SessionOptions() @@ -399,10 +414,12 @@ def _run_nvfp4_moe_test( prev_gemv_env = os.environ.get("ORT_ENABLE_FP4_GEMV") if gemv_mode is not None: os.environ["ORT_ENABLE_FP4_GEMV"] = gemv_mode + provider = resolve_cuda_plugin_ep("CUDAExecutionProvider") + if enable_cuda_graph: + provider_name, provider_options = provider if isinstance(provider, tuple) else (provider, {}) + provider = (provider_name, {**provider_options, "enable_cuda_graph": "1"}) try: - session = onnxruntime.InferenceSession( - onnx_model, opts, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] - ) + session = onnxruntime.InferenceSession(onnx_model, opts, providers=[provider]) except Exception as e: if "ENABLE_FP4" in str(e) or "requires USE_FP4_QMOE" in str(e): self.skipTest(f"NVFP4 not supported in this build: {e}") @@ -437,6 +454,14 @@ def _run_nvfp4_moe_test( raise iobinding.synchronize_outputs() + if enable_cuda_graph: + captured_output = output_tensor.clone() + output_tensor.zero_() + torch.cuda.synchronize() + session.run_with_iobinding(iobinding) + iobinding.synchronize_outputs() + self.assertTrue(torch.equal(captured_output, output_tensor), "NVFP4 CUDA graph replay differs from capture") + ort_output = output_tensor.clone() self.assertTrue(torch.isfinite(ort_output).all().item(), "NVFP4 MoE output contains NaN or infinity") @@ -474,7 +499,7 @@ def _run_nvfp4_moe_test( # not. Allow the corresponding headroom only for shapes that actually route native; # every other shape keeps the strict bound. A genuinely broken kernel still produces # error far above this (order 1.0+), so gross regressions are still caught. - if _routes_native_fp4_prefill(num_tokens): + if _routes_native_fp4_prefill(num_tokens) and not (offline_prepacked and disable_prepacking): atol = 0.25 if torch_dtype == torch.float16 else 0.28 tolerance = atol + rtol_override * ref_output.float().abs().max().item() self.assertLess( @@ -493,6 +518,7 @@ def _assert_invalid_nvfp4_model( truncate_fc1_global_scale=False, hidden_size=64, inter_size=64, + weights_prepacked=None, ): self._skip_if_no_fp4() num_experts = 2 @@ -516,10 +542,11 @@ def _assert_invalid_nvfp4_model( fc2_block_scales=fc2_scales, fc2_global_scale=torch.ones(num_experts, dtype=torch.float32), block_size=block_size, + weights_prepacked=weights_prepacked, ) opts = onnxruntime.SessionOptions() opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL - with self.assertRaisesRegex(Exception, "block_size|fc1_scales|fc1_global_scale|multiple of"): + with self.assertRaisesRegex(Exception, "block_size|fc1_scales|fc1_global_scale|multiple of|weights_prepacked"): session = onnxruntime.InferenceSession( model, opts, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] ) @@ -534,6 +561,24 @@ def _assert_invalid_nvfp4_model( def test_nvfp4_rejects_wrong_block_size(self): self._assert_invalid_nvfp4_model(block_size=32) + @parameterized.expand([(2,), (-2,)]) + def test_nvfp4_rejects_unknown_weight_layout(self, weights_prepacked): + self._assert_invalid_nvfp4_model(weights_prepacked=weights_prepacked) + + @parameterized.expand([(None,), (-1,), (0,)]) + def test_nvfp4_legacy_weights_prepacked(self, weights_prepacked): + self._run_nvfp4_moe_test( + hidden_size=64, + inter_size=64, + num_experts=4, + top_k=2, + num_tokens=2, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + gemv_mode="0", + weights_prepacked=weights_prepacked, + ) + def test_nvfp4_rejects_malformed_prepacked_block_scales(self): self._assert_invalid_nvfp4_model(truncate_fc1_scales=True) @@ -874,6 +919,143 @@ def test_nvfp4_fp16_gemv_disabled_swiglu(self): gemv_mode="0", ) + @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_fallback(self, onnx_dtype, disable_prepacking): + self._run_nvfp4_moe_test( + hidden_size=64, + inter_size=64, + num_experts=4, + top_k=2, + num_tokens=4, + onnx_dtype=onnx_dtype, + gemv_mode="0", + disable_prepacking=disable_prepacking, + use_bias=True, + row_tile_size=2, + offline_prepacked=True, + ) + + @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_gemv(self, onnx_dtype, disable_prepacking): + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=16, + top_k=10, + num_tokens=2, + onnx_dtype=onnx_dtype, + use_swiglu=True, + gemv_mode="1", + disable_prepacking=disable_prepacking, + use_bias=True, + offline_prepacked=True, + ) + + @parameterized.expand( + [ + (onnx_dtype, num_tokens, use_bias) + for onnx_dtype in (TensorProto.FLOAT16, TensorProto.BFLOAT16) + for num_tokens in (1, 6) + for use_bias in (False, True) + ] + ) + def test_nvfp4_offline_row_major_gemv_partial_k_blocks(self, onnx_dtype, num_tokens, use_bias): + shape = dict( + hidden_size=528, + inter_size=656, + num_experts=16, + top_k=10, + num_tokens=num_tokens, + onnx_dtype=onnx_dtype, + use_swiglu=True, + use_bias=use_bias, + row_tile_size=2 if num_tokens > 1 else 0, + input_scale=0.1, + ) + packed = self._run_nvfp4_moe_test( + **shape, gemv_mode="1", offline_prepacked=True, enable_cuda_graph=num_tokens == 1 + ) + legacy = self._run_nvfp4_moe_test(**shape, gemv_mode="1") + fallback = self._run_nvfp4_moe_test(**shape, gemv_mode="0", offline_prepacked=True) + tolerance = 0.03 if onnx_dtype == TensorProto.BFLOAT16 else 0.005 + 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) + for onnx_dtype in (TensorProto.FLOAT16, TensorProto.BFLOAT16) + for disable_prepacking in (False, True) + ] + ) + def test_nvfp4_offline_row_major_prefill(self, onnx_dtype, disable_prepacking): + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=4, + top_k=2, + num_tokens=64, + onnx_dtype=onnx_dtype, + use_swiglu=True, + gemv_mode="0", + disable_prepacking=disable_prepacking, + offline_prepacked=True, + input_scale=0.1, + ) + + @parameterized.expand([(TensorProto.FLOAT16,), (TensorProto.BFLOAT16,)]) + def test_nvfp4_offline_row_major_cuda_graph(self, onnx_dtype): + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=16, + top_k=10, + num_tokens=2, + onnx_dtype=onnx_dtype, + use_swiglu=True, + gemv_mode="1", + offline_prepacked=True, + enable_cuda_graph=True, + ) + def test_nvfp4_fp16_prepacking_disabled_uses_raw_fallback(self): self._run_nvfp4_moe_test( hidden_size=64, @@ -1079,6 +1261,40 @@ def test_nvfp4_fp16_gemv_vs_fallback_parity(self): ) +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="CUDA NVFP4 QMoE required") +@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) +@pytest.mark.parametrize("num_tokens", list(range(1, 10))) +@pytest.mark.parametrize("use_bias", [False, True]) +@pytest.mark.parametrize("row_tile_size", [0, 2]) +def test_nvfp4_kpacked_mtp_window(onnx_dtype, num_tokens, use_bias, row_tile_size, capfd, monkeypatch): + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + monkeypatch.setenv("ORT_FP4_PREFILL_MIN_TOKENS", "32") + shape = dict( + hidden_size=528, + inter_size=656, + num_experts=16, + top_k=10, + num_tokens=num_tokens, + onnx_dtype=onnx_dtype, + use_swiglu=True, + use_bias=use_bias, + row_tile_size=row_tile_size, + offline_prepacked=True, + input_scale=0.1, + ) + packed = TestQMoENVFP4()._run_nvfp4_moe_test(**shape, gemv_mode="1", enable_cuda_graph=num_tokens == 8) + diagnostics = capfd.readouterr().out + if num_tokens <= 8: + assert "Route=fp4_gemv_raw" in diagnostics + assert "DequantWeightBytes=" not in diagnostics + else: + assert "Route=grouped_moe" in diagnostics + assert "DequantWeightBytes=" in diagnostics + fallback = TestQMoENVFP4()._run_nvfp4_moe_test(**shape, gemv_mode="0") + tolerance = 0.03 if onnx_dtype == TensorProto.BFLOAT16 else 0.005 + torch.testing.assert_close(packed, fallback, atol=tolerance, rtol=tolerance) + + @pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="CUDA NVFP4 QMoE required") @pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) @pytest.mark.parametrize("use_swiglu", [False, True]) @@ -1153,5 +1369,31 @@ def test_nvfp4_compact_qwen38_official_shape(num_tokens, capfd, monkeypatch): assert f"QMoE NVFP4 ExpertCapacity={capacity} DequantWeightBytes={weight_bytes}" in capfd.readouterr().out +@pytest.mark.skipif( + os.getenv("ORT_RUN_LARGE_NVFP4_QMOE_TEST") != "1" or not torch.cuda.is_available() or not has_fp4_qmoe, + reason="Opt-in official 512-expert NVFP4 memory test", +) +@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) +@pytest.mark.parametrize("num_tokens", [7, 8]) +def test_nvfp4_kpacked_mtp_qwen38_official_shape(onnx_dtype, num_tokens, capfd, monkeypatch): + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + monkeypatch.setenv("ORT_FP4_PREFILL_MIN_TOKENS", "32") + TestQMoENVFP4()._run_nvfp4_moe_test( + hidden_size=2560, + inter_size=640, + num_experts=512, + top_k=10, + num_tokens=num_tokens, + onnx_dtype=onnx_dtype, + use_swiglu=True, + gemv_mode="1", + offline_prepacked=True, + input_scale=0.05, + ) + diagnostics = capfd.readouterr().out + assert "Route=fp4_gemv_raw" in diagnostics + assert "DequantWeightBytes=" not 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/plugin-ep-webgpu/python/README.md b/plugin-ep-webgpu/python/README.md index e77236766c710..3b878b65b779e 100644 --- a/plugin-ep-webgpu/python/README.md +++ b/plugin-ep-webgpu/python/README.md @@ -54,14 +54,6 @@ The test validates import, EP registration, device discovery, and inference (req inference portion). Set the environment variable `ORT_TEST_VERBOSE=1` to print additional diagnostic information (environment, available providers, discovered devices, etc.). -By default, inference uses `SessionOptions.add_provider_for_devices`, which is supported by the minimum core runtime. -To also test plugin selection through `InferenceSession(..., providers=[("WebGpuExecutionProvider", options)])`, install -a host wheel built from this checkout with `--use_webgpu shared_lib --build_wheel` and set -`ORT_WEBGPU_TEST_PROVIDER_NAME=1`. The Windows WebGPU plugin PR job runs this check in its existing `Run tests` step, -using the locally built host and plugin wheels and reusing the dependencies installed for the job. -The minimum-runtime package tests leave it disabled. -Set `ORT_WEBGPU_TEST_REQUIRE_EP_DEVICE=1` to fail instead of skipping inference when no WebGPU device is available. - ## Versioning The package version is derived from [`../VERSION_NUMBER`](../VERSION_NUMBER) by the packaging pipeline (see diff --git a/plugin-ep-webgpu/python/test/test_webgpu_plugin_ep.py b/plugin-ep-webgpu/python/test/test_webgpu_plugin_ep.py index 94ada9ac558d5..3d62a6a09e467 100644 --- a/plugin-ep-webgpu/python/test/test_webgpu_plugin_ep.py +++ b/plugin-ep-webgpu/python/test/test_webgpu_plugin_ep.py @@ -13,10 +13,6 @@ The inference test is skipped gracefully if no WebGPU device is available (e.g., on CPU-only build agents), unless ORT_WEBGPU_TEST_REQUIRE_EP_DEVICE is enabled. - -Set ORT_WEBGPU_TEST_PROVIDER_NAME=1 when testing a host built from this checkout -with --use_webgpu shared_lib to also exercise provider-name selection. The minimum -supported runtime only needs to support the device-selection API. """ import os @@ -34,7 +30,6 @@ VERBOSE = os.environ.get("ORT_TEST_VERBOSE", "").strip().lower() in ("1", "true", "yes") REQUIRE_EP_DEVICE = os.environ.get("ORT_WEBGPU_TEST_REQUIRE_EP_DEVICE", "").strip().lower() in ("1", "true", "yes") -TEST_PROVIDER_NAME = os.environ.get("ORT_WEBGPU_TEST_PROVIDER_NAME", "").strip().lower() in ("1", "true", "yes") def debug_print(*args, **kwargs): @@ -134,14 +129,8 @@ def test_registration_and_inference(): print("\n--- ORT EP devices ---") print_ep_devices(all_ep_devices) - webgpu_ep_devices = [ - d - for d in all_ep_devices - if d.ep_name == ep_name - and d.ep_metadata.get("library_path") - and Path(d.ep_metadata["library_path"]).resolve() == Path(lib_path).resolve() - ] - print(f"Found {len(webgpu_ep_devices)} WebGPU plugin EP device(s)") + webgpu_ep_devices = [d for d in all_ep_devices if d.ep_name == ep_name] + print(f"Found {len(webgpu_ep_devices)} WebGPU EP device(s)") if not webgpu_ep_devices: if REQUIRE_EP_DEVICE: @@ -153,10 +142,9 @@ def test_registration_and_inference(): webgpu_ep_device = webgpu_ep_devices[0] # Create session with WebGPU EP - provider_options = {"validationMode": "disabled"} sess_options = ort.SessionOptions() sess_options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") - sess_options.add_provider_for_devices([webgpu_ep_device], provider_options) + sess_options.add_provider_for_devices([webgpu_ep_device], {}) assert sess_options.has_providers(), "SessionOptions should have providers after add_provider_for_devices" print("OK: Session options configured with WebGPU EP") @@ -172,36 +160,14 @@ def test_registration_and_inference(): y = np.array([[2.0, 3.0, 4.0], [5.0, 6.0, 7.0]], dtype=np.float32) expected = x * y - try: - assert sess.get_provider_options()[ep_name] == provider_options - outputs = sess.run(None, {"x": x, "y": y}) - result = outputs[0] + outputs = sess.run(None, {"x": x, "y": y}) + result = outputs[0] - np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-5) - print("OK: Inference result matches expected output") - finally: - del sess - print("OK: Session released") + np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-5) + print("OK: Inference result matches expected output") - if TEST_PROVIDER_NAME: - # This binding path is not available in the minimum supported host runtime. - named_options = ort.SessionOptions() - named_options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") - named_session = ort.InferenceSession( - model_path, sess_options=named_options, providers=[(ep_name, provider_options)] - ) - try: - assert ep_name in named_session.get_providers(), "Registered WebGPU provider was not selected" - # Verify that the plugin receives the options supplied through the provider-name API. - assert named_session.get_provider_options()[ep_name] == provider_options, ( - "Provider-name selection did not use the WebGPU plugin with its requested options" - ) - np.testing.assert_allclose( - named_session.run(None, {"x": x, "y": y})[0], expected, rtol=1e-5, atol=1e-5 - ) - finally: - del named_session - print("OK: Registered plugin inference via provider name (CPU fallback disabled)") + del sess + print("OK: Session released") finally: ort.unregister_execution_provider_library(registration_name) diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/matmul_gemv.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/matmul_gemv.h index 01d1b5a4edc1e..1e063b4985b89 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/matmul_gemv.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/math/matmul_gemv.h @@ -26,61 +26,75 @@ ss << "\n"; // 8 | $MAIN { MainFunctionStart(); ss << "\n"; -// 9 | let n_vec = workgroup_idx * workgroup_size_x + local_id.x; -ss << " let n_vec = workgroup_idx * workgroup_size_x + local_id.x;\n"; -// 10 | let n_vec_count = uniforms.N / 4u; +// 9 | let n_vec_count = uniforms.N / 4u; ss << " let n_vec_count = uniforms.N / 4u;\n"; -// 11 | var acc = vec4(0); +// 10 | let groups_per_row = (n_vec_count + workgroup_size_x - 1u) / workgroup_size_x; +ss << " let groups_per_row = (n_vec_count + workgroup_size_x - 1u) / workgroup_size_x;\n"; +// 11 | // Use the flattened index because dispatch normalization can reshape the workgroup grid. +ss << "\n"; +// 12 | let row = workgroup_idx / groups_per_row; +ss << " let row = workgroup_idx / groups_per_row;\n"; +// 13 | // Reject padded rows uniformly before buffer accesses or workgroup barriers. +ss << "\n"; +// 14 | if (row >= uniforms.M) { +ss << " if (row >= uniforms.M) {\n"; +// 15 | return; +ss << " return;\n"; +// 16 | } +ss << " }\n"; +// 17 | let n_vec = (workgroup_idx % groups_per_row) * workgroup_size_x + local_id.x; +ss << " let n_vec = (workgroup_idx % groups_per_row) * workgroup_size_x + local_id.x;\n"; +// 18 | var acc = vec4(0); ss << " var acc = vec4(0);\n"; -// 12 | if (n_vec < n_vec_count) { +// 19 | if (n_vec < n_vec_count) { ss << " if (n_vec < n_vec_count) {\n"; -// 13 | for (var k = local_id.y; k < uniforms.K; k += workgroup_size_y) { +// 20 | for (var k = local_id.y; k < uniforms.K; k += workgroup_size_y) { ss << " for (var k = local_id.y; k < uniforms.K; k += workgroup_size_y) {\n"; -// 14 | let av = f32(a.getByOffset(k)); +// 21 | let av = f32(a.getByOffset(row * uniforms.K + k)); ss << " let av = f32("; -ss << __var_a->GetByOffset("k"); +ss << __var_a->GetByOffset("row * uniforms.K + k"); ss << ");\n"; -// 15 | let bv = vec4(b.getByOffset(k * n_vec_count + n_vec)); +// 22 | let bv = vec4(b.getByOffset(k * n_vec_count + n_vec)); ss << " let bv = vec4("; ss << __var_b->GetByOffset("k * n_vec_count + n_vec"); ss << ");\n"; -// 16 | acc = fma(vec4(av), bv, acc); +// 23 | acc = fma(vec4(av), bv, acc); ss << " acc = fma(vec4(av), bv, acc);\n"; -// 17 | } +// 24 | } ss << " }\n"; -// 18 | } +// 25 | } ss << " }\n"; -// 19 | partials[local_idx] = acc; +// 26 | partials[local_idx] = acc; ss << " partials[local_idx] = acc;\n"; -// 20 | workgroupBarrier(); +// 27 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 21 | +// 28 | ss << "\n"; -// 22 | // All lanes participate in the barriers, including a partial N tile. -// 23 | for (var stride = workgroup_size_y / 2u; stride > 0u; stride /= 2u) { +// 29 | // All lanes participate in the barriers, including a partial N tile. +// 30 | for (var stride = workgroup_size_y / 2u; stride > 0u; stride /= 2u) { ss << " for (var stride = workgroup_size_y / 2u; stride > 0u; stride /= 2u) {\n"; -// 24 | if (local_id.y < stride) { +// 31 | if (local_id.y < stride) { ss << " if (local_id.y < stride) {\n"; -// 25 | partials[local_idx] += partials[local_idx + stride * workgroup_size_x]; +// 32 | partials[local_idx] += partials[local_idx + stride * workgroup_size_x]; ss << " partials[local_idx] += partials[local_idx + stride * workgroup_size_x];\n"; -// 26 | } +// 33 | } ss << " }\n"; -// 27 | workgroupBarrier(); +// 34 | workgroupBarrier(); ss << " workgroupBarrier();\n"; -// 28 | } +// 35 | } ss << " }\n"; -// 29 | if (local_id.y == 0u && n_vec < n_vec_count) { +// 36 | if (local_id.y == 0u && n_vec < n_vec_count) { ss << " if (local_id.y == 0u && n_vec < n_vec_count) {\n"; -// 30 | output.setByOffset(n_vec, output_value_t(partials[local_idx])); +// 37 | output.setByOffset(row * n_vec_count + n_vec, output_value_t(partials[local_idx])); ss << " "; -ss << __var_output->SetByOffset("n_vec", "output_value_t(partials[local_idx])"); +ss << __var_output->SetByOffset("row * n_vec_count + n_vec", "output_value_t(partials[local_idx])"); ss << ";\n"; -// 31 | } +// 38 | } ss << " }\n"; -// 32 | } +// 39 | } MainFunctionEnd(); ss << "\n"; -// 33 | +// 40 | return Status::OK(); diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h index 9306021ec2b09..b85499e73881b 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h @@ -34,7 +34,7 @@ std::string pass_as_string(T&& v) { // Include template implementations -#include "wgsl_template_gen/generated/math/matmul_gemv.h" // 5ae837e52fbaa856b637c6f8801781389406334a2359f8efc49c49226d8e2983 +#include "wgsl_template_gen/generated/math/matmul_gemv.h" // ad945fab49f626209e741b76ac3900250cb43159d4ba8a3f66de12fad97f5c43 #include "wgsl_template_gen/generated/math/subgroup_matrix_gemm_8x16x16.h" // 1aaff98897bad5e6e46888a2948ca3ac245eda97bdbbe7463c7cec4b6da6d0a9 #include "wgsl_template_gen/generated/math/subgroup_matrix_matmul.h" // 5eb470fbe35d55b1ccfbf62f8d780a6f2798853a2628f81e19622f0b24a287f7 #include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // c01fae062d8059604670635900f3dd036b692cc3d16ebb8a1e54491e83439e66 diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/matmul_gemv.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/matmul_gemv.h index 4ecd3eabbef1f..b342a4321eaa2 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/matmul_gemv.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/matmul_gemv.h @@ -26,61 +26,75 @@ ss << __str_5; // 8 | $MAIN { MainFunctionStart(); ss << __str_5; -// 9 | let n_vec = workgroup_idx * workgroup_size_x + local_id.x; +// 9 | let n_vec_count = uniforms.N / 4u; ss << __str_6; -// 10 | let n_vec_count = uniforms.N / 4u; +// 10 | let groups_per_row = (n_vec_count + workgroup_size_x - 1u) / workgroup_size_x; ss << __str_7; -// 11 | var acc = vec4(0); +// 11 | // Use the flattened index because dispatch normalization can reshape the workgroup grid. +ss << __str_5; +// 12 | let row = workgroup_idx / groups_per_row; ss << __str_8; -// 12 | if (n_vec < n_vec_count) { +// 13 | // Reject padded rows uniformly before buffer accesses or workgroup barriers. +ss << __str_5; +// 14 | if (row >= uniforms.M) { ss << __str_9; -// 13 | for (var k = local_id.y; k < uniforms.K; k += workgroup_size_y) { +// 15 | return; ss << __str_10; -// 14 | let av = f32(a.getByOffset(k)); +// 16 | } ss << __str_11; -ss << __var_a->GetByOffset(__str_0); +// 17 | let n_vec = (workgroup_idx % groups_per_row) * workgroup_size_x + local_id.x; ss << __str_12; -// 15 | let bv = vec4(b.getByOffset(k * n_vec_count + n_vec)); +// 18 | var acc = vec4(0); ss << __str_13; -ss << __var_b->GetByOffset(__str_1); -ss << __str_12; -// 16 | acc = fma(vec4(av), bv, acc); +// 19 | if (n_vec < n_vec_count) { ss << __str_14; -// 17 | } +// 20 | for (var k = local_id.y; k < uniforms.K; k += workgroup_size_y) { ss << __str_15; -// 18 | } +// 21 | let av = f32(a.getByOffset(row * uniforms.K + k)); ss << __str_16; -// 19 | partials[local_idx] = acc; +ss << __var_a->GetByOffset(__str_0); ss << __str_17; -// 20 | workgroupBarrier(); +// 22 | let bv = vec4(b.getByOffset(k * n_vec_count + n_vec)); ss << __str_18; -// 21 | -ss << __str_5; -// 22 | // All lanes participate in the barriers, including a partial N tile. -// 23 | for (var stride = workgroup_size_y / 2u; stride > 0u; stride /= 2u) { +ss << __var_b->GetByOffset(__str_1); +ss << __str_17; +// 23 | acc = fma(vec4(av), bv, acc); ss << __str_19; -// 24 | if (local_id.y < stride) { +// 24 | } ss << __str_20; -// 25 | partials[local_idx] += partials[local_idx + stride * workgroup_size_x]; +// 25 | } +ss << __str_11; +// 26 | partials[local_idx] = acc; ss << __str_21; -// 26 | } -ss << __str_15; -// 27 | workgroupBarrier(); +// 27 | workgroupBarrier(); ss << __str_22; -// 28 | } -ss << __str_16; -// 29 | if (local_id.y == 0u && n_vec < n_vec_count) { +// 28 | +ss << __str_5; +// 29 | // All lanes participate in the barriers, including a partial N tile. +// 30 | for (var stride = workgroup_size_y / 2u; stride > 0u; stride /= 2u) { ss << __str_23; -// 30 | output.setByOffset(n_vec, output_value_t(partials[local_idx])); +// 31 | if (local_id.y < stride) { ss << __str_24; -ss << __var_output->SetByOffset(__str_2, __str_3); +// 32 | partials[local_idx] += partials[local_idx + stride * workgroup_size_x]; ss << __str_25; -// 31 | } -ss << __str_16; -// 32 | } +// 33 | } +ss << __str_20; +// 34 | workgroupBarrier(); +ss << __str_26; +// 35 | } +ss << __str_11; +// 36 | if (local_id.y == 0u && n_vec < n_vec_count) { +ss << __str_27; +// 37 | output.setByOffset(row * n_vec_count + n_vec, output_value_t(partials[local_idx])); +ss << __str_28; +ss << __var_output->SetByOffset(__str_2, __str_3); +ss << __str_29; +// 38 | } +ss << __str_11; +// 39 | } MainFunctionEnd(); ss << __str_5; -// 33 | +// 40 | return Status::OK(); diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_gemm_8x16x16.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_gemm_8x16x16.h index cfd85efc0507a..6127c98acc23e 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_gemm_8x16x16.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_gemm_8x16x16.h @@ -82,35 +82,35 @@ Status ApplyTemplate<"math/subgroup_matrix_gemm_8x16x16.wgsl.template">(ShaderHe // 57 | #use .setByOffset // 58 | // 59 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) -ss << __str_28; +ss << __str_32; ss << __param_sg_mat_m; -ss << __str_12; +ss << __str_17; // 60 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) -ss << __str_29; +ss << __str_33; ss << __param_sg_mat_n; -ss << __str_12; +ss << __str_17; // 61 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) -ss << __str_30; +ss << __str_34; ss << __param_sg_mat_k; -ss << __str_12; +ss << __str_17; // 62 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M -ss << __str_31; +ss << __str_35; ss << __param_sg_mat_count_m; -ss << __str_12; +ss << __str_17; // 63 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N -ss << __str_32; +ss << __str_36; ss << __param_sg_mat_count_n; -ss << __str_12; +ss << __str_17; // 64 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) -ss << __str_33; +ss << __str_37; // 65 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) -ss << __str_34; +ss << __str_38; // 66 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K -ss << __str_35; +ss << __str_39; ss << __param_split_k; -ss << __str_12; +ss << __str_17; // 67 | const kSubgroupSize: u32 = 32; // Lanes per subgroup -ss << __str_36; +ss << __str_40; // 68 | ss << __str_5; // 69 | // Scratch space for kSplitK partial [kTileM, kTileN] tiles, row-major. Each @@ -118,7 +118,7 @@ ss << __str_5; // 71 | // summed during write-back. Storing here also lets us write to global memory // 72 | // with M/N bounds checks for dimensions that are not multiples of the tile size. // 73 | var scratch: array; -ss << __str_37; +ss << __str_41; // 74 | ss << __str_5; // 75 | $MAIN { @@ -132,105 +132,105 @@ ss << __str_5; // 80 | // true linear index, invariant under normalization, laid out as // 81 | // m_tile * num_n_tile + n_tile. // 82 | let num_n_tile = (uniforms.N + kTileN - 1u) / kTileN; -ss << __str_38; +ss << __str_42; // 83 | let n_tile = workgroup_idx % num_n_tile; // which kTileN-wide column tile -ss << __str_39; +ss << __str_43; // 84 | let m_tile = workgroup_idx / num_n_tile; // which kTileM-tall row tile -ss << __str_40; +ss << __str_44; // 85 | let n_tile_start = n_tile * kTileN; -ss << __str_41; +ss << __str_45; // 86 | let m_tile_start = m_tile * kTileM; -ss << __str_42; +ss << __str_46; // 87 | // Shift a trailing partial tile back so its operand loads stay inside A and B. ss << __str_5; // 88 | // The max() only guards the u32 subtraction; the host guarantees the tile fits. // 89 | let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN); -ss << __str_43; +ss << __str_47; // 90 | let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM); -ss << __str_44; +ss << __str_48; // 91 | let k_blocks = uniforms.K / kSgMatK; -ss << __str_45; +ss << __str_49; // 92 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) -ss << __str_46; +ss << __str_50; // 93 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) -ss << __str_47; +ss << __str_51; // 94 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's scratch slot -ss << __str_48; +ss << __str_52; // 95 | ss << __str_5; // 96 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). // 97 | var sg_mat_c0_0: subgroup_matrix_result; -ss << __str_49; +ss << __str_53; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 98 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 99 | var sg_mat_c0_1: subgroup_matrix_result; -ss << __str_52; +ss << __str_56; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 100 | #endif } // 101 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 102 | var sg_mat_c0_2: subgroup_matrix_result; -ss << __str_53; +ss << __str_57; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 103 | #endif } // 104 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 105 | var sg_mat_c0_3: subgroup_matrix_result; -ss << __str_54; +ss << __str_58; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 106 | #endif } // 107 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { // 108 | var sg_mat_c1_0: subgroup_matrix_result; -ss << __str_55; +ss << __str_59; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 109 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 110 | var sg_mat_c1_1: subgroup_matrix_result; -ss << __str_56; +ss << __str_60; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 111 | #endif } // 112 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 113 | var sg_mat_c1_2: subgroup_matrix_result; -ss << __str_57; +ss << __str_61; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 114 | #endif } // 115 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 116 | var sg_mat_c1_3: subgroup_matrix_result; -ss << __str_58; +ss << __str_62; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 117 | #endif } // 118 | #endif @@ -238,39 +238,39 @@ ss << __str_51; // 119 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { // 120 | var sg_mat_c2_0: subgroup_matrix_result; -ss << __str_59; +ss << __str_63; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 121 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 122 | var sg_mat_c2_1: subgroup_matrix_result; -ss << __str_60; +ss << __str_64; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 123 | #endif } // 124 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 125 | var sg_mat_c2_2: subgroup_matrix_result; -ss << __str_61; +ss << __str_65; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 126 | #endif } // 127 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 128 | var sg_mat_c2_3: subgroup_matrix_result; -ss << __str_62; +ss << __str_66; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 129 | #endif } // 130 | #endif @@ -278,39 +278,39 @@ ss << __str_51; // 131 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { // 132 | var sg_mat_c3_0: subgroup_matrix_result; -ss << __str_63; +ss << __str_67; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 133 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 134 | var sg_mat_c3_1: subgroup_matrix_result; -ss << __str_64; +ss << __str_68; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 135 | #endif } // 136 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 137 | var sg_mat_c3_2: subgroup_matrix_result; -ss << __str_65; +ss << __str_69; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 138 | #endif } // 139 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 140 | var sg_mat_c3_3: subgroup_matrix_result; -ss << __str_66; +ss << __str_70; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 141 | #endif } // 142 | #endif @@ -318,39 +318,39 @@ ss << __str_51; // 143 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { // 144 | var sg_mat_c4_0: subgroup_matrix_result; -ss << __str_67; +ss << __str_71; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 145 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 146 | var sg_mat_c4_1: subgroup_matrix_result; -ss << __str_68; +ss << __str_72; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 147 | #endif } // 148 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 149 | var sg_mat_c4_2: subgroup_matrix_result; -ss << __str_69; +ss << __str_73; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 150 | #endif } // 151 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 152 | var sg_mat_c4_3: subgroup_matrix_result; -ss << __str_70; +ss << __str_74; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 153 | #endif } // 154 | #endif @@ -358,39 +358,39 @@ ss << __str_51; // 155 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { // 156 | var sg_mat_c5_0: subgroup_matrix_result; -ss << __str_71; +ss << __str_75; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 157 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 158 | var sg_mat_c5_1: subgroup_matrix_result; -ss << __str_72; +ss << __str_76; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 159 | #endif } // 160 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 161 | var sg_mat_c5_2: subgroup_matrix_result; -ss << __str_73; +ss << __str_77; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 162 | #endif } // 163 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 164 | var sg_mat_c5_3: subgroup_matrix_result; -ss << __str_74; +ss << __str_78; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 165 | #endif } // 166 | #endif @@ -398,39 +398,39 @@ ss << __str_51; // 167 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { // 168 | var sg_mat_c6_0: subgroup_matrix_result; -ss << __str_75; +ss << __str_79; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 169 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 170 | var sg_mat_c6_1: subgroup_matrix_result; -ss << __str_76; +ss << __str_80; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 171 | #endif } // 172 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 173 | var sg_mat_c6_2: subgroup_matrix_result; -ss << __str_77; +ss << __str_81; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 174 | #endif } // 175 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 176 | var sg_mat_c6_3: subgroup_matrix_result; -ss << __str_78; +ss << __str_82; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 177 | #endif } // 178 | #endif @@ -438,48 +438,48 @@ ss << __str_51; // 179 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { // 180 | var sg_mat_c7_0: subgroup_matrix_result; -ss << __str_79; +ss << __str_83; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 181 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 182 | var sg_mat_c7_1: subgroup_matrix_result; -ss << __str_80; +ss << __str_84; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 183 | #endif } // 184 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 185 | var sg_mat_c7_2: subgroup_matrix_result; -ss << __str_81; +ss << __str_85; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 186 | #endif } // 187 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 188 | var sg_mat_c7_3: subgroup_matrix_result; -ss << __str_82; +ss << __str_86; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 189 | #endif } // 190 | #endif } // 191 | // 192 | for (var kb: u32 = sg_index; kb < k_blocks; kb = kb + kSplitK) { -ss << __str_83; +ss << __str_87; // 193 | let kb_k = kb * kSgMatK; -ss << __str_84; +ss << __str_88; // 194 | // Load the B right tiles for this K block. ss << __str_5; // 195 | #if trans_b @@ -488,144 +488,144 @@ if (__param_trans_b) { // 197 | // N block ni starts at logical element (n = global_base_n + ni*kSgMatN, // 198 | // k = kb_k), i.e. flat offset (global_base_n + ni*kSgMatN) * K + kb_k. // 199 | var sg_mat_b0: subgroup_matrix_right = -ss << __str_85; +ss << __str_89; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 200 | subgroupMatrixLoad, col_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_88; +ss << __str_92; // 201 | &input_b, (global_base_n + 0 * kSgMatN) * uniforms.K + kb_k, uniforms.K); -ss << __str_89; +ss << __str_93; // 202 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 203 | var sg_mat_b1: subgroup_matrix_right = -ss << __str_90; +ss << __str_94; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 204 | subgroupMatrixLoad, col_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_88; +ss << __str_92; // 205 | &input_b, (global_base_n + 1 * kSgMatN) * uniforms.K + kb_k, uniforms.K); -ss << __str_91; +ss << __str_95; // 206 | #endif } // 207 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 208 | var sg_mat_b2: subgroup_matrix_right = -ss << __str_92; +ss << __str_96; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 209 | subgroupMatrixLoad, col_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_88; +ss << __str_92; // 210 | &input_b, (global_base_n + 2 * kSgMatN) * uniforms.K + kb_k, uniforms.K); -ss << __str_93; +ss << __str_97; // 211 | #endif } // 212 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 213 | var sg_mat_b3: subgroup_matrix_right = -ss << __str_94; +ss << __str_98; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 214 | subgroupMatrixLoad, col_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_88; +ss << __str_92; // 215 | &input_b, (global_base_n + 3 * kSgMatN) * uniforms.K + kb_k, uniforms.K); -ss << __str_95; +ss << __str_99; // 216 | #endif } // 217 | #else } else { // 218 | // B stored [K, N] row-major: load row-major with stride N. // 219 | let b_base = kb_k * uniforms.N + global_base_n; -ss << __str_96; +ss << __str_100; // 220 | var sg_mat_b0: subgroup_matrix_right = -ss << __str_85; +ss << __str_89; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 221 | subgroupMatrixLoad, row_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_97; +ss << __str_101; // 222 | &input_b, b_base + 0 * kSgMatN, uniforms.N); -ss << __str_98; +ss << __str_102; // 223 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 224 | var sg_mat_b1: subgroup_matrix_right = -ss << __str_90; +ss << __str_94; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 225 | subgroupMatrixLoad, row_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_97; +ss << __str_101; // 226 | &input_b, b_base + 1 * kSgMatN, uniforms.N); -ss << __str_99; +ss << __str_103; // 227 | #endif } // 228 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 229 | var sg_mat_b2: subgroup_matrix_right = -ss << __str_92; +ss << __str_96; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 230 | subgroupMatrixLoad, row_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_97; +ss << __str_101; // 231 | &input_b, b_base + 2 * kSgMatN, uniforms.N); -ss << __str_100; +ss << __str_104; // 232 | #endif } // 233 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 234 | var sg_mat_b3: subgroup_matrix_right = -ss << __str_94; +ss << __str_98; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 235 | subgroupMatrixLoad, row_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_97; -// 236 | &input_b, b_base + 3 * kSgMatN, uniforms.N); ss << __str_101; +// 236 | &input_b, b_base + 3 * kSgMatN, uniforms.N); +ss << __str_105; // 237 | #endif } // 238 | #endif @@ -638,290 +638,290 @@ if (__param_trans_a) { // 243 | // M block mi starts at logical element (m = m_base + mi*kSgMatM, k = kb_k), // 244 | // i.e. flat offset kb_k * M + (m_base + mi*kSgMatM). // 245 | let a_row = kb_k * uniforms.M; -ss << __str_102; +ss << __str_106; // 246 | var sg_mat_a0: subgroup_matrix_left = -ss << __str_103; +ss << __str_107; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 247 | subgroupMatrixLoad, col_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_88; +ss << __str_92; // 248 | &input_a, a_row + (m_base + 0 * kSgMatM), uniforms.M); -ss << __str_105; +ss << __str_109; // 249 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { // 250 | var sg_mat_a1: subgroup_matrix_left = -ss << __str_106; +ss << __str_110; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 251 | subgroupMatrixLoad, col_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_88; +ss << __str_92; // 252 | &input_a, a_row + (m_base + 1 * kSgMatM), uniforms.M); -ss << __str_107; +ss << __str_111; // 253 | #endif } // 254 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { // 255 | var sg_mat_a2: subgroup_matrix_left = -ss << __str_108; +ss << __str_112; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 256 | subgroupMatrixLoad, col_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_88; +ss << __str_92; // 257 | &input_a, a_row + (m_base + 2 * kSgMatM), uniforms.M); -ss << __str_109; +ss << __str_113; // 258 | #endif } // 259 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { // 260 | var sg_mat_a3: subgroup_matrix_left = -ss << __str_110; +ss << __str_114; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 261 | subgroupMatrixLoad, col_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_88; +ss << __str_92; // 262 | &input_a, a_row + (m_base + 3 * kSgMatM), uniforms.M); -ss << __str_111; +ss << __str_115; // 263 | #endif } // 264 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { // 265 | var sg_mat_a4: subgroup_matrix_left = -ss << __str_112; +ss << __str_116; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 266 | subgroupMatrixLoad, col_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_88; +ss << __str_92; // 267 | &input_a, a_row + (m_base + 4 * kSgMatM), uniforms.M); -ss << __str_113; +ss << __str_117; // 268 | #endif } // 269 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { // 270 | var sg_mat_a5: subgroup_matrix_left = -ss << __str_114; +ss << __str_118; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 271 | subgroupMatrixLoad, col_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_88; +ss << __str_92; // 272 | &input_a, a_row + (m_base + 5 * kSgMatM), uniforms.M); -ss << __str_115; +ss << __str_119; // 273 | #endif } // 274 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { // 275 | var sg_mat_a6: subgroup_matrix_left = -ss << __str_116; +ss << __str_120; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 276 | subgroupMatrixLoad, col_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_88; +ss << __str_92; // 277 | &input_a, a_row + (m_base + 6 * kSgMatM), uniforms.M); -ss << __str_117; +ss << __str_121; // 278 | #endif } // 279 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { // 280 | var sg_mat_a7: subgroup_matrix_left = -ss << __str_118; +ss << __str_122; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 281 | subgroupMatrixLoad, col_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_88; +ss << __str_92; // 282 | &input_a, a_row + (m_base + 7 * kSgMatM), uniforms.M); -ss << __str_119; +ss << __str_123; // 283 | #endif } // 284 | #else } else { // 285 | // A stored [M, K] row-major: load row-major with stride K. // 286 | let a_col = kb_k; -ss << __str_120; +ss << __str_124; // 287 | var sg_mat_a0: subgroup_matrix_left = -ss << __str_103; +ss << __str_107; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 288 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 289 | &input_a, (m_base + 0 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_121; +ss << __str_125; // 290 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { // 291 | var sg_mat_a1: subgroup_matrix_left = -ss << __str_106; +ss << __str_110; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 292 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 293 | &input_a, (m_base + 1 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_122; +ss << __str_126; // 294 | #endif } // 295 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { // 296 | var sg_mat_a2: subgroup_matrix_left = -ss << __str_108; +ss << __str_112; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 297 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 298 | &input_a, (m_base + 2 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_123; +ss << __str_127; // 299 | #endif } // 300 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { // 301 | var sg_mat_a3: subgroup_matrix_left = -ss << __str_110; +ss << __str_114; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 302 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 303 | &input_a, (m_base + 3 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_124; +ss << __str_128; // 304 | #endif } // 305 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { // 306 | var sg_mat_a4: subgroup_matrix_left = -ss << __str_112; +ss << __str_116; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 307 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 308 | &input_a, (m_base + 4 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_125; +ss << __str_129; // 309 | #endif } // 310 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { // 311 | var sg_mat_a5: subgroup_matrix_left = -ss << __str_114; +ss << __str_118; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 312 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 313 | &input_a, (m_base + 5 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_126; +ss << __str_130; // 314 | #endif } // 315 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { // 316 | var sg_mat_a6: subgroup_matrix_left = -ss << __str_116; +ss << __str_120; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 317 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 318 | &input_a, (m_base + 6 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_127; +ss << __str_131; // 319 | #endif } // 320 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { // 321 | var sg_mat_a7: subgroup_matrix_left = -ss << __str_118; +ss << __str_122; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 322 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 323 | &input_a, (m_base + 7 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_128; +ss << __str_132; // 324 | #endif } // 325 | #endif @@ -929,45 +929,45 @@ ss << __str_128; // 326 | // 327 | // Accumulate every (M block, N block) pair. // 328 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); -ss << __str_129; +ss << __str_133; // 329 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 330 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); -ss << __str_130; +ss << __str_134; // 331 | #endif } // 332 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 333 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); -ss << __str_131; +ss << __str_135; // 334 | #endif } // 335 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 336 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); -ss << __str_132; +ss << __str_136; // 337 | #endif } // 338 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { // 339 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); -ss << __str_133; +ss << __str_137; // 340 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 341 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); -ss << __str_134; +ss << __str_138; // 342 | #endif } // 343 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 344 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); -ss << __str_135; +ss << __str_139; // 345 | #endif } // 346 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 347 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); -ss << __str_136; +ss << __str_140; // 348 | #endif } // 349 | #endif @@ -975,23 +975,23 @@ ss << __str_136; // 350 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { // 351 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); -ss << __str_137; +ss << __str_141; // 352 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 353 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); -ss << __str_138; +ss << __str_142; // 354 | #endif } // 355 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 356 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); -ss << __str_139; +ss << __str_143; // 357 | #endif } // 358 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 359 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); -ss << __str_140; +ss << __str_144; // 360 | #endif } // 361 | #endif @@ -999,23 +999,23 @@ ss << __str_140; // 362 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { // 363 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); -ss << __str_141; +ss << __str_145; // 364 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 365 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); -ss << __str_142; +ss << __str_146; // 366 | #endif } // 367 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 368 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); -ss << __str_143; +ss << __str_147; // 369 | #endif } // 370 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 371 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); -ss << __str_144; +ss << __str_148; // 372 | #endif } // 373 | #endif @@ -1023,23 +1023,23 @@ ss << __str_144; // 374 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { // 375 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); -ss << __str_145; +ss << __str_149; // 376 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 377 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); -ss << __str_146; +ss << __str_150; // 378 | #endif } // 379 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 380 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); -ss << __str_147; +ss << __str_151; // 381 | #endif } // 382 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 383 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); -ss << __str_148; +ss << __str_152; // 384 | #endif } // 385 | #endif @@ -1047,23 +1047,23 @@ ss << __str_148; // 386 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { // 387 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); -ss << __str_149; +ss << __str_153; // 388 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 389 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); -ss << __str_150; +ss << __str_154; // 390 | #endif } // 391 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 392 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); -ss << __str_151; +ss << __str_155; // 393 | #endif } // 394 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 395 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); -ss << __str_152; +ss << __str_156; // 396 | #endif } // 397 | #endif @@ -1071,23 +1071,23 @@ ss << __str_152; // 398 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { // 399 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); -ss << __str_153; +ss << __str_157; // 400 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 401 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); -ss << __str_154; +ss << __str_158; // 402 | #endif } // 403 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 404 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); -ss << __str_155; +ss << __str_159; // 405 | #endif } // 406 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 407 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); -ss << __str_156; +ss << __str_160; // 408 | #endif } // 409 | #endif @@ -1095,74 +1095,74 @@ ss << __str_156; // 410 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { // 411 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); -ss << __str_157; +ss << __str_161; // 412 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 413 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); -ss << __str_158; +ss << __str_162; // 414 | #endif } // 415 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 416 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); -ss << __str_159; +ss << __str_163; // 417 | #endif } // 418 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 419 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); -ss << __str_160; +ss << __str_164; // 420 | #endif } // 421 | #endif } // 422 | } -ss << __str_15; +ss << __str_20; // 423 | ss << __str_5; // 424 | // Store this subgroup's partial results into its split-K slot of scratch, // 425 | // laid out as kSplitK x [kTileM, kTileN] row-major. Offset of block (mi, ni) // 426 | // within a slot = mi * kSgMatM * kTileN + ni * kSgMatN, stride kTileN. // 427 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); -ss << __str_161; +ss << __str_165; // 428 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 429 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); -ss << __str_162; +ss << __str_166; // 430 | #endif } // 431 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 432 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); -ss << __str_163; +ss << __str_167; // 433 | #endif } // 434 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 435 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); -ss << __str_164; +ss << __str_168; // 436 | #endif } // 437 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { // 438 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); -ss << __str_165; +ss << __str_169; // 439 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 440 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); -ss << __str_166; +ss << __str_170; // 441 | #endif } // 442 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 443 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); -ss << __str_167; +ss << __str_171; // 444 | #endif } // 445 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 446 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); -ss << __str_168; +ss << __str_172; // 447 | #endif } // 448 | #endif @@ -1170,23 +1170,23 @@ ss << __str_168; // 449 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { // 450 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); -ss << __str_169; +ss << __str_173; // 451 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); -ss << __str_170; +ss << __str_174; // 453 | #endif } // 454 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 455 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); -ss << __str_171; +ss << __str_175; // 456 | #endif } // 457 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 458 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); -ss << __str_172; +ss << __str_176; // 459 | #endif } // 460 | #endif @@ -1194,23 +1194,23 @@ ss << __str_172; // 461 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { // 462 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); -ss << __str_173; +ss << __str_177; // 463 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 464 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); -ss << __str_174; +ss << __str_178; // 465 | #endif } // 466 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 467 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); -ss << __str_175; +ss << __str_179; // 468 | #endif } // 469 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 470 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); -ss << __str_176; +ss << __str_180; // 471 | #endif } // 472 | #endif @@ -1218,23 +1218,23 @@ ss << __str_176; // 473 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { // 474 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); -ss << __str_177; +ss << __str_181; // 475 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 476 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); -ss << __str_178; +ss << __str_182; // 477 | #endif } // 478 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 479 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); -ss << __str_179; +ss << __str_183; // 480 | #endif } // 481 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 482 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); -ss << __str_180; +ss << __str_184; // 483 | #endif } // 484 | #endif @@ -1242,23 +1242,23 @@ ss << __str_180; // 485 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { // 486 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); -ss << __str_181; +ss << __str_185; // 487 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 488 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); -ss << __str_182; +ss << __str_186; // 489 | #endif } // 490 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 491 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); -ss << __str_183; +ss << __str_187; // 492 | #endif } // 493 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 494 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); -ss << __str_184; +ss << __str_188; // 495 | #endif } // 496 | #endif @@ -1266,23 +1266,23 @@ ss << __str_184; // 497 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { // 498 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); -ss << __str_185; +ss << __str_189; // 499 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 500 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); -ss << __str_186; +ss << __str_190; // 501 | #endif } // 502 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 503 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); -ss << __str_187; +ss << __str_191; // 504 | #endif } // 505 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 506 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); -ss << __str_188; +ss << __str_192; // 507 | #endif } // 508 | #endif @@ -1290,23 +1290,23 @@ ss << __str_188; // 509 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { // 510 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); -ss << __str_189; +ss << __str_193; // 511 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 512 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); -ss << __str_190; +ss << __str_194; // 513 | #endif } // 514 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 515 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); -ss << __str_191; +ss << __str_195; // 516 | #endif } // 517 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 518 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); -ss << __str_192; +ss << __str_196; // 519 | #endif } // 520 | #endif @@ -1316,7 +1316,7 @@ ss << __str_192; if (__param_split_k >= 2) { // 523 | // Make all subgroups' partial tiles visible before the reduction. // 524 | workgroupBarrier(); -ss << __str_22; +ss << __str_26; // 525 | ss << __str_5; // 526 | // Sum pass: a single subgroup reduces the kSplitK partial [kTileM, kTileN] @@ -1325,27 +1325,27 @@ ss << __str_5; // 529 | // is safe: a position is only overwritten by the same lane that already read // 530 | // it as the slot-0 term of its own sum. // 531 | if (sg_index == 0u) { -ss << __str_193; +ss << __str_197; // 532 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { -ss << __str_194; +ss << __str_198; // 533 | var acc: f16 = f16(0); -ss << __str_195; +ss << __str_199; // 534 | for (var s: u32 = 0; s < kSplitK; s++) { -ss << __str_196; +ss << __str_200; // 535 | acc += scratch[s * kTileM * kTileN + idx]; -ss << __str_197; +ss << __str_201; // 536 | } -ss << __str_198; +ss << __str_202; // 537 | scratch[idx] = acc; -ss << __str_199; +ss << __str_203; // 538 | } -ss << __str_200; +ss << __str_204; // 539 | } -ss << __str_15; +ss << __str_20; // 540 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. ss << __str_5; // 541 | workgroupBarrier(); -ss << __str_22; +ss << __str_26; // 542 | #endif } // 543 | @@ -1359,43 +1359,43 @@ ss << __str_22; // 551 | // is scaled by alpha and, when has_c, has beta * C[m, n] added (C broadcast via // 552 | // strides). // 553 | let n_skip = n_tile_start - global_base_n; -ss << __str_201; +ss << __str_205; // 554 | let n_count = min(kTileN, uniforms.N - global_base_n); -ss << __str_202; +ss << __str_206; // 555 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { -ss << __str_203; +ss << __str_207; // 556 | let global_m = m_base + r; -ss << __str_204; +ss << __str_208; // 557 | if (global_m >= m_tile_start && global_m < uniforms.M) { -ss << __str_205; +ss << __str_209; // 558 | let scratch_base = r * kTileN; -ss << __str_206; +ss << __str_210; // 559 | let out_base = global_m * uniforms.N + global_base_n; -ss << __str_207; +ss << __str_211; // 560 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { -ss << __str_208; +ss << __str_212; // 561 | var val = output_element_t(scratch[scratch_base + i]) * output_element_t(uniforms.alpha); -ss << __str_209; +ss << __str_213; // 562 | #if has_c if (__param_has_c) { // 563 | let global_n = global_base_n + i; -ss << __str_210; +ss << __str_214; // 564 | let c_offset = global_m * uniforms.c_stride_m + global_n * uniforms.c_stride_n; -ss << __str_211; +ss << __str_215; // 565 | val += output_element_t(uniforms.beta) * input_c[c_offset]; -ss << __str_212; +ss << __str_216; // 566 | #endif } // 567 | output.setByOffset(out_base + i, val); -ss << __str_213; -ss << __var_output->SetByOffset(__str_26, __str_27); -ss << __str_25; +ss << __str_217; +ss << __var_output->SetByOffset(__str_30, __str_31); +ss << __str_29; // 568 | } -ss << __str_198; +ss << __str_202; // 569 | } -ss << __str_200; +ss << __str_204; // 570 | } -ss << __str_15; +ss << __str_20; // 571 | } // MAIN MainFunctionEnd(); ss << __str_5; diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul.h index eeae32076ebe6..d44023edb406f 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul.h @@ -72,35 +72,35 @@ Status ApplyTemplate<"math/subgroup_matrix_matmul.wgsl.template">(ShaderHelper& // 53 | #use .setByOffset // 54 | // // 55 | const kSgMatM: u32 = u32(sg_mat_m); // Subgroup matrix M dimension (rows) -ss << __str_28; +ss << __str_32; ss << __param_sg_mat_m; -ss << __str_12; +ss << __str_17; // 56 | const kSgMatN: u32 = u32(sg_mat_n); // Subgroup matrix N dimension (columns) -ss << __str_29; +ss << __str_33; ss << __param_sg_mat_n; -ss << __str_12; +ss << __str_17; // 57 | const kSgMatK: u32 = u32(sg_mat_k); // Subgroup matrix K dimension (reduction) -ss << __str_30; +ss << __str_34; ss << __param_sg_mat_k; -ss << __str_12; +ss << __str_17; // 58 | const kSgMatCountM: u32 = u32(sg_mat_count_m); // Number of subgroup matrices per tile along M -ss << __str_31; +ss << __str_35; ss << __param_sg_mat_count_m; -ss << __str_12; +ss << __str_17; // 59 | const kSgMatCountN: u32 = u32(sg_mat_count_n); // Number of subgroup matrices per tile along N -ss << __str_32; +ss << __str_36; ss << __param_sg_mat_count_n; -ss << __str_12; +ss << __str_17; // 60 | const kTileM: u32 = kSgMatM * kSgMatCountM; // Tile size along M (output rows per workgroup) -ss << __str_33; +ss << __str_37; // 61 | const kTileN: u32 = kSgMatN * kSgMatCountN; // Tile size along N (output columns per workgroup) -ss << __str_34; +ss << __str_38; // 62 | const kSplitK: u32 = u32(split_k); // Number of subgroups cooperating along K -ss << __str_35; +ss << __str_39; ss << __param_split_k; -ss << __str_12; +ss << __str_17; // 63 | const kSubgroupSize: u32 = 32; // Lanes per subgroup -ss << __str_36; +ss << __str_40; // 64 | // ss << __str_5; // 65 | // Scratch space for kSplitK partial [kTileM, kTileN] tiles, row-major. Each @@ -108,7 +108,7 @@ ss << __str_5; // 67 | // summed during write-back. Storing here also lets us write to global memory // 68 | // with M/N bounds checks for dimensions that are not multiples of the tile size. // 69 | var scratch: array; -ss << __str_37; +ss << __str_41; // 70 | // ss << __str_5; // 71 | $MAIN { @@ -122,21 +122,21 @@ ss << __str_5; // 76 | // used directly. workgroup_idx is the true linear index, invariant under // 77 | // normalization; batch slices are contiguous blocks of tiles_per_slice tiles. // 78 | let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM; -ss << __str_214; +ss << __str_218; // 79 | let tiles_per_slice = uniforms.num_n_tile * num_m_tile; -ss << __str_215; +ss << __str_219; // 80 | let batch_id = workgroup_idx / tiles_per_slice; // which batch slice (0 for a shared 2D weight) -ss << __str_216; +ss << __str_220; // 81 | let slice_idx = workgroup_idx % tiles_per_slice; // tile index within the slice -ss << __str_217; +ss << __str_221; // 82 | let n_tile = slice_idx % uniforms.num_n_tile; // which kTileN-wide column tile -ss << __str_218; +ss << __str_222; // 83 | let m_tile = slice_idx / uniforms.num_n_tile; // which kTileM-tall row tile -ss << __str_219; +ss << __str_223; // 84 | let n_tile_start = n_tile * kTileN; -ss << __str_41; +ss << __str_45; // 85 | let m_tile_start = m_tile * kTileM; -ss << __str_42; +ss << __str_46; // 86 | // Shift a trailing partial tile back so its operand loads stay inside A and B. ss << __str_5; // 87 | // The N shift uses B's padded stride N_b so the shifted origin stays aligned (the @@ -144,101 +144,101 @@ ss << __str_5; // 89 | // The max() only guards the u32 subtraction; the host guarantees // 90 | // the tile fits. // 91 | let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN); -ss << __str_220; +ss << __str_224; // 92 | let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM); -ss << __str_44; +ss << __str_48; // 93 | // Flat-element offsets into A/B/output for this batch slice, derived from ss << __str_5; // 94 | // M/N/K. For a shared 2D weight batch_id is 0, so B collapses to its base. // 95 | let a_batch_offset = batch_id * uniforms.M * uniforms.K; -ss << __str_221; +ss << __str_225; // 96 | let b_batch_offset = batch_id * uniforms.K * uniforms.N_b; -ss << __str_222; +ss << __str_226; // 97 | let out_batch_offset = batch_id * uniforms.M * uniforms.N; -ss << __str_223; +ss << __str_227; // 98 | let k_blocks = uniforms.K / kSgMatK; -ss << __str_45; +ss << __str_49; // 99 | let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) -ss << __str_46; +ss << __str_50; // 100 | let sg_lane = local_idx % kSubgroupSize; // lane within the subgroup (0..kSubgroupSize-1) -ss << __str_47; +ss << __str_51; // 101 | let sg_scratch_base = sg_index * kTileM * kTileN; // this subgroup's scratch slot -ss << __str_48; +ss << __str_52; // 102 | // ss << __str_5; // 103 | // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c_). // 104 | var sg_mat_c0_0: subgroup_matrix_result; -ss << __str_49; +ss << __str_53; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 105 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 106 | var sg_mat_c0_1: subgroup_matrix_result; -ss << __str_52; +ss << __str_56; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 107 | #endif } // 108 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 109 | var sg_mat_c0_2: subgroup_matrix_result; -ss << __str_53; +ss << __str_57; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 110 | #endif } // 111 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 112 | var sg_mat_c0_3: subgroup_matrix_result; -ss << __str_54; +ss << __str_58; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 113 | #endif } // 114 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { // 115 | var sg_mat_c1_0: subgroup_matrix_result; -ss << __str_55; +ss << __str_59; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 116 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 117 | var sg_mat_c1_1: subgroup_matrix_result; -ss << __str_56; +ss << __str_60; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 118 | #endif } // 119 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 120 | var sg_mat_c1_2: subgroup_matrix_result; -ss << __str_57; +ss << __str_61; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 121 | #endif } // 122 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 123 | var sg_mat_c1_3: subgroup_matrix_result; -ss << __str_58; +ss << __str_62; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 124 | #endif } // 125 | #endif @@ -246,39 +246,39 @@ ss << __str_51; // 126 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { // 127 | var sg_mat_c2_0: subgroup_matrix_result; -ss << __str_59; +ss << __str_63; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 128 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 129 | var sg_mat_c2_1: subgroup_matrix_result; -ss << __str_60; +ss << __str_64; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 130 | #endif } // 131 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 132 | var sg_mat_c2_2: subgroup_matrix_result; -ss << __str_61; +ss << __str_65; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 133 | #endif } // 134 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 135 | var sg_mat_c2_3: subgroup_matrix_result; -ss << __str_62; +ss << __str_66; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 136 | #endif } // 137 | #endif @@ -286,39 +286,39 @@ ss << __str_51; // 138 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { // 139 | var sg_mat_c3_0: subgroup_matrix_result; -ss << __str_63; +ss << __str_67; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 140 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 141 | var sg_mat_c3_1: subgroup_matrix_result; -ss << __str_64; +ss << __str_68; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 142 | #endif } // 143 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 144 | var sg_mat_c3_2: subgroup_matrix_result; -ss << __str_65; +ss << __str_69; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 145 | #endif } // 146 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 147 | var sg_mat_c3_3: subgroup_matrix_result; -ss << __str_66; +ss << __str_70; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 148 | #endif } // 149 | #endif @@ -326,39 +326,39 @@ ss << __str_51; // 150 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { // 151 | var sg_mat_c4_0: subgroup_matrix_result; -ss << __str_67; +ss << __str_71; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 152 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 153 | var sg_mat_c4_1: subgroup_matrix_result; -ss << __str_68; +ss << __str_72; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 154 | #endif } // 155 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 156 | var sg_mat_c4_2: subgroup_matrix_result; -ss << __str_69; +ss << __str_73; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 157 | #endif } // 158 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 159 | var sg_mat_c4_3: subgroup_matrix_result; -ss << __str_70; +ss << __str_74; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 160 | #endif } // 161 | #endif @@ -366,39 +366,39 @@ ss << __str_51; // 162 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { // 163 | var sg_mat_c5_0: subgroup_matrix_result; -ss << __str_71; +ss << __str_75; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 164 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 165 | var sg_mat_c5_1: subgroup_matrix_result; -ss << __str_72; +ss << __str_76; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 166 | #endif } // 167 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 168 | var sg_mat_c5_2: subgroup_matrix_result; -ss << __str_73; +ss << __str_77; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 169 | #endif } // 170 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 171 | var sg_mat_c5_3: subgroup_matrix_result; -ss << __str_74; +ss << __str_78; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 172 | #endif } // 173 | #endif @@ -406,39 +406,39 @@ ss << __str_51; // 174 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { // 175 | var sg_mat_c6_0: subgroup_matrix_result; -ss << __str_75; +ss << __str_79; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 176 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 177 | var sg_mat_c6_1: subgroup_matrix_result; -ss << __str_76; +ss << __str_80; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 178 | #endif } // 179 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 180 | var sg_mat_c6_2: subgroup_matrix_result; -ss << __str_77; +ss << __str_81; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 181 | #endif } // 182 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 183 | var sg_mat_c6_3: subgroup_matrix_result; -ss << __str_78; +ss << __str_82; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 184 | #endif } // 185 | #endif @@ -446,116 +446,116 @@ ss << __str_51; // 186 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { // 187 | var sg_mat_c7_0: subgroup_matrix_result; -ss << __str_79; +ss << __str_83; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 188 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 189 | var sg_mat_c7_1: subgroup_matrix_result; -ss << __str_80; +ss << __str_84; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 190 | #endif } // 191 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 192 | var sg_mat_c7_2: subgroup_matrix_result; -ss << __str_81; +ss << __str_85; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 193 | #endif } // 194 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 195 | var sg_mat_c7_3: subgroup_matrix_result; -ss << __str_82; +ss << __str_86; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_51; +ss << __str_55; // 196 | #endif } // 197 | #endif } // 198 | // // 199 | for (var kb: u32 = sg_index; kb < k_blocks; kb = kb + kSplitK) { -ss << __str_83; +ss << __str_87; // 200 | // Load the B right tiles for this K block (KxN_b row-major, stride N_b). ss << __str_5; // 201 | let b_base = b_batch_offset + kb * kSgMatK * uniforms.N_b + global_base_n; -ss << __str_224; +ss << __str_228; // 202 | var sg_mat_b0: subgroup_matrix_right = -ss << __str_85; +ss << __str_89; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 203 | subgroupMatrixLoad, row_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_97; +ss << __str_101; // 204 | &input_b, b_base + 0 * kSgMatN, uniforms.N_b); -ss << __str_225; +ss << __str_229; // 205 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 206 | var sg_mat_b1: subgroup_matrix_right = -ss << __str_90; +ss << __str_94; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 207 | subgroupMatrixLoad, row_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_97; +ss << __str_101; // 208 | &input_b, b_base + 1 * kSgMatN, uniforms.N_b); -ss << __str_226; +ss << __str_230; // 209 | #endif } // 210 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 211 | var sg_mat_b2: subgroup_matrix_right = -ss << __str_92; +ss << __str_96; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 212 | subgroupMatrixLoad, row_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_97; +ss << __str_101; // 213 | &input_b, b_base + 2 * kSgMatN, uniforms.N_b); -ss << __str_227; +ss << __str_231; // 214 | #endif } // 215 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 216 | var sg_mat_b3: subgroup_matrix_right = -ss << __str_94; +ss << __str_98; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_86; +ss << __str_90; // 217 | subgroupMatrixLoad, row_major>( -ss << __str_87; +ss << __str_91; ss << __param_sg_mat_n; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_k; -ss << __str_97; +ss << __str_101; // 218 | &input_b, b_base + 3 * kSgMatN, uniforms.N_b); -ss << __str_228; +ss << __str_232; // 219 | #endif } // 220 | // @@ -563,189 +563,189 @@ ss << __str_228; // 222 | // a_batch_offset folds this z-slice's flat offset into the column term so // 223 | // every A load below reads from the correct batch slice. // 224 | let a_col = a_batch_offset + kb * kSgMatK; -ss << __str_229; +ss << __str_233; // 225 | var sg_mat_a0: subgroup_matrix_left = -ss << __str_103; +ss << __str_107; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 226 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 227 | &input_a, (m_base + 0 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_121; +ss << __str_125; // 228 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { // 229 | var sg_mat_a1: subgroup_matrix_left = -ss << __str_106; +ss << __str_110; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 230 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 231 | &input_a, (m_base + 1 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_122; +ss << __str_126; // 232 | #endif } // 233 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { // 234 | var sg_mat_a2: subgroup_matrix_left = -ss << __str_108; +ss << __str_112; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 235 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 236 | &input_a, (m_base + 2 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_123; +ss << __str_127; // 237 | #endif } // 238 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { // 239 | var sg_mat_a3: subgroup_matrix_left = -ss << __str_110; +ss << __str_114; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 240 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 241 | &input_a, (m_base + 3 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_124; +ss << __str_128; // 242 | #endif } // 243 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { // 244 | var sg_mat_a4: subgroup_matrix_left = -ss << __str_112; +ss << __str_116; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 245 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 246 | &input_a, (m_base + 4 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_125; +ss << __str_129; // 247 | #endif } // 248 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { // 249 | var sg_mat_a5: subgroup_matrix_left = -ss << __str_114; +ss << __str_118; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 250 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 251 | &input_a, (m_base + 5 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_126; +ss << __str_130; // 252 | #endif } // 253 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { // 254 | var sg_mat_a6: subgroup_matrix_left = -ss << __str_116; +ss << __str_120; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 255 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 256 | &input_a, (m_base + 6 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_127; +ss << __str_131; // 257 | #endif } // 258 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { // 259 | var sg_mat_a7: subgroup_matrix_left = -ss << __str_118; +ss << __str_122; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_86; +ss << __str_90; // 260 | subgroupMatrixLoad, row_major>( -ss << __str_104; +ss << __str_108; ss << __param_sg_mat_k; -ss << __str_50; +ss << __str_54; ss << __param_sg_mat_m; -ss << __str_97; +ss << __str_101; // 261 | &input_a, (m_base + 7 * kSgMatM) * uniforms.K + a_col, uniforms.K); -ss << __str_128; +ss << __str_132; // 262 | #endif } // 263 | // // 264 | // Accumulate every (M block, N block) pair. // 265 | sg_mat_c0_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b0, sg_mat_c0_0); -ss << __str_129; +ss << __str_133; // 266 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 267 | sg_mat_c0_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b1, sg_mat_c0_1); -ss << __str_130; +ss << __str_134; // 268 | #endif } // 269 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 270 | sg_mat_c0_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b2, sg_mat_c0_2); -ss << __str_131; +ss << __str_135; // 271 | #endif } // 272 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 273 | sg_mat_c0_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a0, sg_mat_b3, sg_mat_c0_3); -ss << __str_132; +ss << __str_136; // 274 | #endif } // 275 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { // 276 | sg_mat_c1_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b0, sg_mat_c1_0); -ss << __str_133; +ss << __str_137; // 277 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 278 | sg_mat_c1_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b1, sg_mat_c1_1); -ss << __str_134; +ss << __str_138; // 279 | #endif } // 280 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 281 | sg_mat_c1_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b2, sg_mat_c1_2); -ss << __str_135; +ss << __str_139; // 282 | #endif } // 283 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 284 | sg_mat_c1_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a1, sg_mat_b3, sg_mat_c1_3); -ss << __str_136; +ss << __str_140; // 285 | #endif } // 286 | #endif @@ -753,23 +753,23 @@ ss << __str_136; // 287 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { // 288 | sg_mat_c2_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b0, sg_mat_c2_0); -ss << __str_137; +ss << __str_141; // 289 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 290 | sg_mat_c2_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b1, sg_mat_c2_1); -ss << __str_138; +ss << __str_142; // 291 | #endif } // 292 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 293 | sg_mat_c2_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b2, sg_mat_c2_2); -ss << __str_139; +ss << __str_143; // 294 | #endif } // 295 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 296 | sg_mat_c2_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a2, sg_mat_b3, sg_mat_c2_3); -ss << __str_140; +ss << __str_144; // 297 | #endif } // 298 | #endif @@ -777,23 +777,23 @@ ss << __str_140; // 299 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { // 300 | sg_mat_c3_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b0, sg_mat_c3_0); -ss << __str_141; +ss << __str_145; // 301 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 302 | sg_mat_c3_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b1, sg_mat_c3_1); -ss << __str_142; +ss << __str_146; // 303 | #endif } // 304 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 305 | sg_mat_c3_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b2, sg_mat_c3_2); -ss << __str_143; +ss << __str_147; // 306 | #endif } // 307 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 308 | sg_mat_c3_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a3, sg_mat_b3, sg_mat_c3_3); -ss << __str_144; +ss << __str_148; // 309 | #endif } // 310 | #endif @@ -801,23 +801,23 @@ ss << __str_144; // 311 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { // 312 | sg_mat_c4_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b0, sg_mat_c4_0); -ss << __str_145; +ss << __str_149; // 313 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 314 | sg_mat_c4_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b1, sg_mat_c4_1); -ss << __str_146; +ss << __str_150; // 315 | #endif } // 316 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 317 | sg_mat_c4_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b2, sg_mat_c4_2); -ss << __str_147; +ss << __str_151; // 318 | #endif } // 319 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 320 | sg_mat_c4_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a4, sg_mat_b3, sg_mat_c4_3); -ss << __str_148; +ss << __str_152; // 321 | #endif } // 322 | #endif @@ -825,23 +825,23 @@ ss << __str_148; // 323 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { // 324 | sg_mat_c5_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b0, sg_mat_c5_0); -ss << __str_149; +ss << __str_153; // 325 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 326 | sg_mat_c5_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b1, sg_mat_c5_1); -ss << __str_150; +ss << __str_154; // 327 | #endif } // 328 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 329 | sg_mat_c5_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b2, sg_mat_c5_2); -ss << __str_151; +ss << __str_155; // 330 | #endif } // 331 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 332 | sg_mat_c5_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a5, sg_mat_b3, sg_mat_c5_3); -ss << __str_152; +ss << __str_156; // 333 | #endif } // 334 | #endif @@ -849,23 +849,23 @@ ss << __str_152; // 335 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { // 336 | sg_mat_c6_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b0, sg_mat_c6_0); -ss << __str_153; +ss << __str_157; // 337 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 338 | sg_mat_c6_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b1, sg_mat_c6_1); -ss << __str_154; +ss << __str_158; // 339 | #endif } // 340 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 341 | sg_mat_c6_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b2, sg_mat_c6_2); -ss << __str_155; +ss << __str_159; // 342 | #endif } // 343 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 344 | sg_mat_c6_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a6, sg_mat_b3, sg_mat_c6_3); -ss << __str_156; +ss << __str_160; // 345 | #endif } // 346 | #endif @@ -873,74 +873,74 @@ ss << __str_156; // 347 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { // 348 | sg_mat_c7_0 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b0, sg_mat_c7_0); -ss << __str_157; +ss << __str_161; // 349 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 350 | sg_mat_c7_1 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b1, sg_mat_c7_1); -ss << __str_158; +ss << __str_162; // 351 | #endif } // 352 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 353 | sg_mat_c7_2 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b2, sg_mat_c7_2); -ss << __str_159; +ss << __str_163; // 354 | #endif } // 355 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 356 | sg_mat_c7_3 = subgroupMatrixMultiplyAccumulate(sg_mat_a7, sg_mat_b3, sg_mat_c7_3); -ss << __str_160; +ss << __str_164; // 357 | #endif } // 358 | #endif } // 359 | } -ss << __str_15; +ss << __str_20; // 360 | // ss << __str_5; // 361 | // Store this subgroup's partial results into its split-K slot of scratch, // 362 | // laid out as kSplitK x [kTileM, kTileN] row-major. Offset of block (mi, ni) // 363 | // within a slot = mi * kSgMatM * kTileN + ni * kSgMatN, stride kTileN. // 364 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c0_0, kTileN); -ss << __str_161; +ss << __str_165; // 365 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 366 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c0_1, kTileN); -ss << __str_162; +ss << __str_166; // 367 | #endif } // 368 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 369 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c0_2, kTileN); -ss << __str_163; +ss << __str_167; // 370 | #endif } // 371 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 372 | subgroupMatrixStore(&scratch, sg_scratch_base + 0 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c0_3, kTileN); -ss << __str_164; +ss << __str_168; // 373 | #endif } // 374 | #if sg_mat_count_m >= 2 if (__param_sg_mat_count_m >= 2) { // 375 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c1_0, kTileN); -ss << __str_165; +ss << __str_169; // 376 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 377 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c1_1, kTileN); -ss << __str_166; +ss << __str_170; // 378 | #endif } // 379 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 380 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c1_2, kTileN); -ss << __str_167; +ss << __str_171; // 381 | #endif } // 382 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 383 | subgroupMatrixStore(&scratch, sg_scratch_base + 1 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c1_3, kTileN); -ss << __str_168; +ss << __str_172; // 384 | #endif } // 385 | #endif @@ -948,23 +948,23 @@ ss << __str_168; // 386 | #if sg_mat_count_m >= 3 if (__param_sg_mat_count_m >= 3) { // 387 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c2_0, kTileN); -ss << __str_169; +ss << __str_173; // 388 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 389 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c2_1, kTileN); -ss << __str_170; +ss << __str_174; // 390 | #endif } // 391 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 392 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c2_2, kTileN); -ss << __str_171; +ss << __str_175; // 393 | #endif } // 394 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 395 | subgroupMatrixStore(&scratch, sg_scratch_base + 2 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c2_3, kTileN); -ss << __str_172; +ss << __str_176; // 396 | #endif } // 397 | #endif @@ -972,23 +972,23 @@ ss << __str_172; // 398 | #if sg_mat_count_m >= 4 if (__param_sg_mat_count_m >= 4) { // 399 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c3_0, kTileN); -ss << __str_173; +ss << __str_177; // 400 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 401 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c3_1, kTileN); -ss << __str_174; +ss << __str_178; // 402 | #endif } // 403 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 404 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c3_2, kTileN); -ss << __str_175; +ss << __str_179; // 405 | #endif } // 406 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 407 | subgroupMatrixStore(&scratch, sg_scratch_base + 3 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c3_3, kTileN); -ss << __str_176; +ss << __str_180; // 408 | #endif } // 409 | #endif @@ -996,23 +996,23 @@ ss << __str_176; // 410 | #if sg_mat_count_m >= 5 if (__param_sg_mat_count_m >= 5) { // 411 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c4_0, kTileN); -ss << __str_177; +ss << __str_181; // 412 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 413 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c4_1, kTileN); -ss << __str_178; +ss << __str_182; // 414 | #endif } // 415 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 416 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c4_2, kTileN); -ss << __str_179; +ss << __str_183; // 417 | #endif } // 418 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 419 | subgroupMatrixStore(&scratch, sg_scratch_base + 4 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c4_3, kTileN); -ss << __str_180; +ss << __str_184; // 420 | #endif } // 421 | #endif @@ -1020,23 +1020,23 @@ ss << __str_180; // 422 | #if sg_mat_count_m >= 6 if (__param_sg_mat_count_m >= 6) { // 423 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c5_0, kTileN); -ss << __str_181; +ss << __str_185; // 424 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 425 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c5_1, kTileN); -ss << __str_182; +ss << __str_186; // 426 | #endif } // 427 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 428 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c5_2, kTileN); -ss << __str_183; +ss << __str_187; // 429 | #endif } // 430 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 431 | subgroupMatrixStore(&scratch, sg_scratch_base + 5 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c5_3, kTileN); -ss << __str_184; +ss << __str_188; // 432 | #endif } // 433 | #endif @@ -1044,23 +1044,23 @@ ss << __str_184; // 434 | #if sg_mat_count_m >= 7 if (__param_sg_mat_count_m >= 7) { // 435 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c6_0, kTileN); -ss << __str_185; +ss << __str_189; // 436 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 437 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c6_1, kTileN); -ss << __str_186; +ss << __str_190; // 438 | #endif } // 439 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 440 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c6_2, kTileN); -ss << __str_187; +ss << __str_191; // 441 | #endif } // 442 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 443 | subgroupMatrixStore(&scratch, sg_scratch_base + 6 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c6_3, kTileN); -ss << __str_188; +ss << __str_192; // 444 | #endif } // 445 | #endif @@ -1068,23 +1068,23 @@ ss << __str_188; // 446 | #if sg_mat_count_m >= 8 if (__param_sg_mat_count_m >= 8) { // 447 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 0 * kSgMatN, sg_mat_c7_0, kTileN); -ss << __str_189; +ss << __str_193; // 448 | #if sg_mat_count_n >= 2 if (__param_sg_mat_count_n >= 2) { // 449 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 1 * kSgMatN, sg_mat_c7_1, kTileN); -ss << __str_190; +ss << __str_194; // 450 | #endif } // 451 | #if sg_mat_count_n >= 3 if (__param_sg_mat_count_n >= 3) { // 452 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 2 * kSgMatN, sg_mat_c7_2, kTileN); -ss << __str_191; +ss << __str_195; // 453 | #endif } // 454 | #if sg_mat_count_n >= 4 if (__param_sg_mat_count_n >= 4) { // 455 | subgroupMatrixStore(&scratch, sg_scratch_base + 7 * kSgMatM * kTileN + 3 * kSgMatN, sg_mat_c7_3, kTileN); -ss << __str_192; +ss << __str_196; // 456 | #endif } // 457 | #endif @@ -1094,7 +1094,7 @@ ss << __str_192; if (__param_split_k >= 2) { // 460 | // Make all subgroups' partial tiles visible before the reduction. // 461 | workgroupBarrier(); -ss << __str_22; +ss << __str_26; // 462 | // ss << __str_5; // 463 | // Sum pass: a single subgroup reduces the kSplitK partial [kTileM, kTileN] @@ -1103,27 +1103,27 @@ ss << __str_5; // 466 | // is safe: a position is only overwritten by the same lane that already read // 467 | // it as the slot-0 term of its own sum. // 468 | if (sg_index == 0u) { -ss << __str_193; +ss << __str_197; // 469 | for (var idx: u32 = sg_lane; idx < kTileM * kTileN; idx = idx + kSubgroupSize) { -ss << __str_194; +ss << __str_198; // 470 | var acc: f16 = f16(0); -ss << __str_195; +ss << __str_199; // 471 | for (var s: u32 = 0; s < kSplitK; s++) { -ss << __str_196; +ss << __str_200; // 472 | acc += scratch[s * kTileM * kTileN + idx]; -ss << __str_197; +ss << __str_201; // 473 | } -ss << __str_198; +ss << __str_202; // 474 | scratch[idx] = acc; -ss << __str_199; +ss << __str_203; // 475 | } -ss << __str_200; +ss << __str_204; // 476 | } -ss << __str_15; +ss << __str_20; // 477 | // Publish the summed tile (slot 0) to every subgroup for the write-out pass. ss << __str_5; // 478 | workgroupBarrier(); -ss << __str_22; +ss << __str_26; // 479 | #endif } // 480 | // @@ -1135,31 +1135,31 @@ ss << __str_22; // 486 | // shifted tile never writes the same output element twice; M and N are // 487 | // bounds-checked so non-multiple dimensions only write valid data. // 488 | let n_skip = n_tile_start - global_base_n; -ss << __str_201; +ss << __str_205; // 489 | let n_count = min(kTileN, uniforms.N - global_base_n); -ss << __str_202; +ss << __str_206; // 490 | for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) { -ss << __str_203; +ss << __str_207; // 491 | let global_m = m_base + r; -ss << __str_204; +ss << __str_208; // 492 | if (global_m >= m_tile_start && global_m < uniforms.M) { -ss << __str_205; +ss << __str_209; // 493 | let scratch_base = r * kTileN; -ss << __str_206; +ss << __str_210; // 494 | let out_base = out_batch_offset + global_m * uniforms.N + global_base_n; -ss << __str_230; +ss << __str_234; // 495 | for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) { -ss << __str_208; +ss << __str_212; // 496 | write_output(out_base + i, global_base_n + i, -ss << __str_231; +ss << __str_235; // 497 | output_value_t(scratch[scratch_base + i])); -ss << __str_232; +ss << __str_236; // 498 | } -ss << __str_198; +ss << __str_202; // 499 | } -ss << __str_200; +ss << __str_204; // 500 | } -ss << __str_15; +ss << __str_20; // 501 | } // MAIN MainFunctionEnd(); ss << __str_5; diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_pad_b.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_pad_b.h index 37bed16c4d318..72c280617cd01 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_pad_b.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_pad_b.h @@ -28,27 +28,27 @@ Status ApplyTemplate<"math/subgroup_matrix_matmul_pad_b.wgsl.template">(ShaderHe MainFunctionStart(); ss << __str_5; // 13 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); -ss << __str_237; -ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_233); -ss << __str_25; +ss << __str_241; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_237); +ss << __str_29; // 14 | let r = global_idx / uniforms.N_b; // padded-output row -ss << __str_238; +ss << __str_242; // 15 | let c = global_idx % uniforms.N_b; // padded-output column -ss << __str_239; +ss << __str_243; // 16 | var v = output_value_t(0); -ss << __str_240; +ss << __str_244; // 17 | if (c < uniforms.N) { // real column -> copy; else zero pad -ss << __str_241; +ss << __str_245; // 18 | v = output_value_t(input_b.getByOffset(r * uniforms.N + c)); -ss << __str_242; -ss << __var_input_b->GetByOffset(__str_234); -ss << __str_12; +ss << __str_246; +ss << __var_input_b->GetByOffset(__str_238); +ss << __str_17; // 19 | } -ss << __str_16; +ss << __str_11; // 20 | output.setByOffset(global_idx, v); -ss << __str_237; -ss << __var_output->SetByOffset(__str_235, __str_236); -ss << __str_25; +ss << __str_241; +ss << __var_output->SetByOffset(__str_239, __str_240); +ss << __str_29; // 21 | } // MAIN MainFunctionEnd(); ss << __str_5; diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/conv_transpose3d.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/conv_transpose3d.h index ace6f78799dba..f86290a3a1457 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/conv_transpose3d.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/conv_transpose3d.h @@ -35,15 +35,15 @@ Status ApplyTemplate<"nn/conv_transpose3d.wgsl.template">(ShaderHelper& shader_h // 12 | #if components == 4 if (__param_components == 4) { // 13 | alias Accumulator = vec4; -ss << __str_258; +ss << __str_262; // 14 | #elif components == 2 } else if (__param_components == 2) { // 15 | alias Accumulator = vec2; -ss << __str_259; +ss << __str_263; // 16 | #else } else { // 17 | alias Accumulator = f32; -ss << __str_260; +ss << __str_264; // 18 | #endif } // 19 | @@ -51,199 +51,199 @@ ss << __str_260; MainFunctionStart(); ss << __str_5; // 21 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); -ss << __str_237; -ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_233); -ss << __str_25; +ss << __str_241; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_237); +ss << __str_29; // 22 | let output_indices = output.offsetToIndices(global_idx); -ss << __str_261; -ss << __var_output->OffsetToIndices(__str_235); -ss << __str_25; +ss << __str_265; +ss << __var_output->OffsetToIndices(__str_239); +ss << __str_29; // 23 | let batch = output.indicesGet(output_indices, 0); -ss << __str_262; -ss << __var_output->IndicesGet(__str_243, __str_244); -ss << __str_25; +ss << __str_266; +ss << __var_output->IndicesGet(__str_247, __str_248); +ss << __str_29; // 24 | #if is_channels_last if (__param_is_channels_last) { // 25 | let output_channel = output.indicesGet(output_indices, 4); -ss << __str_263; -ss << __var_output->IndicesGet(__str_243, __str_245); -ss << __str_25; +ss << __str_267; +ss << __var_output->IndicesGet(__str_247, __str_249); +ss << __str_29; // 26 | let spatial_dim = 1u; -ss << __str_264; +ss << __str_268; // 27 | #else } else { // 28 | let output_channel = output.indicesGet(output_indices, 1); -ss << __str_263; -ss << __var_output->IndicesGet(__str_243, __str_246); -ss << __str_25; +ss << __str_267; +ss << __var_output->IndicesGet(__str_247, __str_250); +ss << __str_29; // 29 | let spatial_dim = 2u; -ss << __str_265; +ss << __str_269; // 30 | #endif } // 31 | #if is_prepacked if (__param_is_prepacked) { // 32 | let output_channels_per_group = getElementAt(uniforms.w_shape, 3, w.rank); -ss << __str_266; -ss << GetElementAt(__str_247, __str_248, __var_w->Rank()); -ss << __str_25; +ss << __str_270; +ss << GetElementAt(__str_251, __str_252, __var_w->Rank()); +ss << __str_29; // 33 | let kernel_shape = vec3(getElementAt(uniforms.w_shape, 0, w.rank), -ss << __str_267; -ss << GetElementAt(__str_247, __str_244, __var_w->Rank()); -ss << __str_268; +ss << __str_271; +ss << GetElementAt(__str_251, __str_248, __var_w->Rank()); +ss << __str_272; // 34 | getElementAt(uniforms.w_shape, 1, w.rank), -ss << __str_269; -ss << GetElementAt(__str_247, __str_246, __var_w->Rank()); -ss << __str_268; +ss << __str_273; +ss << GetElementAt(__str_251, __str_250, __var_w->Rank()); +ss << __str_272; // 35 | getElementAt(uniforms.w_shape, 2, w.rank)); -ss << __str_269; -ss << GetElementAt(__str_247, __str_249, __var_w->Rank()); -ss << __str_12; +ss << __str_273; +ss << GetElementAt(__str_251, __str_253, __var_w->Rank()); +ss << __str_17; // 36 | #else } else { // 37 | let output_channels_per_group = getElementAt(uniforms.w_shape, 1, w.rank); -ss << __str_266; -ss << GetElementAt(__str_247, __str_246, __var_w->Rank()); -ss << __str_25; +ss << __str_270; +ss << GetElementAt(__str_251, __str_250, __var_w->Rank()); +ss << __str_29; // 38 | let kernel_shape = vec3(getElementAt(uniforms.w_shape, 2, w.rank), -ss << __str_267; -ss << GetElementAt(__str_247, __str_249, __var_w->Rank()); -ss << __str_268; +ss << __str_271; +ss << GetElementAt(__str_251, __str_253, __var_w->Rank()); +ss << __str_272; // 39 | getElementAt(uniforms.w_shape, 3, w.rank), -ss << __str_269; -ss << GetElementAt(__str_247, __str_248, __var_w->Rank()); -ss << __str_268; +ss << __str_273; +ss << GetElementAt(__str_251, __str_252, __var_w->Rank()); +ss << __str_272; // 40 | getElementAt(uniforms.w_shape, 4, w.rank)); -ss << __str_269; -ss << GetElementAt(__str_247, __str_245, __var_w->Rank()); -ss << __str_12; +ss << __str_273; +ss << GetElementAt(__str_251, __str_249, __var_w->Rank()); +ss << __str_17; // 41 | #endif } // 42 | let group = output_channel / output_channels_per_group; -ss << __str_270; +ss << __str_274; // 43 | let w_channel = output_channel % output_channels_per_group; -ss << __str_271; +ss << __str_275; // 44 | let corner = vec3( -ss << __str_272; +ss << __str_276; // 45 | i32(output.indicesGet(output_indices, spatial_dim)), -ss << __str_273; -ss << __var_output->IndicesGet(__str_243, __str_250); -ss << __str_274; +ss << __str_277; +ss << __var_output->IndicesGet(__str_247, __str_254); +ss << __str_278; // 46 | i32(output.indicesGet(output_indices, spatial_dim + 1u)), -ss << __str_273; -ss << __var_output->IndicesGet(__str_243, __str_251); -ss << __str_274; +ss << __str_277; +ss << __var_output->IndicesGet(__str_247, __str_255); +ss << __str_278; // 47 | i32(output.indicesGet(output_indices, spatial_dim + 2u))) + vec3(uniforms.pads); -ss << __str_273; -ss << __var_output->IndicesGet(__str_243, __str_252); -ss << __str_275; +ss << __str_277; +ss << __var_output->IndicesGet(__str_247, __str_256); +ss << __str_279; // 48 | var value = Accumulator(0); -ss << __str_276; +ss << __str_280; // 49 | ss << __str_5; // 50 | // Gather input positions using signed integer coordinates so padding and stride // 51 | // alignment remain exact even for float16 tensors with large spatial dimensions. // 52 | for (var kd = 0u; kd < kernel_shape.x; kd++) { -ss << __str_277; +ss << __str_281; // 53 | let pos_d = corner.x - i32(kd) * i32(uniforms.dilations.x); -ss << __str_278; +ss << __str_282; // 54 | if (pos_d < 0 || pos_d % i32(uniforms.strides.x) != 0) { -ss << __str_279; +ss << __str_283; // 55 | continue; -ss << __str_280; +ss << __str_284; // 56 | } -ss << __str_15; +ss << __str_20; // 57 | let id = u32(pos_d) / uniforms.strides.x; -ss << __str_281; +ss << __str_285; // 58 | if (id >= getElementAt(uniforms.x_shape, spatial_dim, x.rank)) { -ss << __str_282; -ss << GetElementAt(__str_253, __str_250, __var_x->Rank()); -ss << __str_283; +ss << __str_286; +ss << GetElementAt(__str_257, __str_254, __var_x->Rank()); +ss << __str_287; // 59 | continue; -ss << __str_280; +ss << __str_284; // 60 | } -ss << __str_15; +ss << __str_20; // 61 | for (var kh = 0u; kh < kernel_shape.y; kh++) { -ss << __str_284; +ss << __str_288; // 62 | let pos_h = corner.y - i32(kh) * i32(uniforms.dilations.y); -ss << __str_285; +ss << __str_289; // 63 | if (pos_h < 0 || pos_h % i32(uniforms.strides.y) != 0) { -ss << __str_286; +ss << __str_290; // 64 | continue; -ss << __str_287; +ss << __str_291; // 65 | } -ss << __str_288; +ss << __str_292; // 66 | let ih = u32(pos_h) / uniforms.strides.y; -ss << __str_289; +ss << __str_293; // 67 | if (ih >= getElementAt(uniforms.x_shape, spatial_dim + 1u, x.rank)) { -ss << __str_290; -ss << GetElementAt(__str_253, __str_251, __var_x->Rank()); -ss << __str_283; -// 68 | continue; +ss << __str_294; +ss << GetElementAt(__str_257, __str_255, __var_x->Rank()); ss << __str_287; +// 68 | continue; +ss << __str_291; // 69 | } -ss << __str_288; +ss << __str_292; // 70 | for (var kw = 0u; kw < kernel_shape.z; kw++) { -ss << __str_291; +ss << __str_295; // 71 | let pos_w = corner.z - i32(kw) * i32(uniforms.dilations.z); -ss << __str_292; +ss << __str_296; // 72 | if (pos_w < 0 || pos_w % i32(uniforms.strides.z) != 0) { -ss << __str_293; +ss << __str_297; // 73 | continue; -ss << __str_294; +ss << __str_298; // 74 | } -ss << __str_200; +ss << __str_204; // 75 | let iw = u32(pos_w) / uniforms.strides.z; -ss << __str_295; +ss << __str_299; // 76 | if (iw >= getElementAt(uniforms.x_shape, spatial_dim + 2u, x.rank)) { -ss << __str_296; -ss << GetElementAt(__str_253, __str_252, __var_x->Rank()); -ss << __str_283; +ss << __str_300; +ss << GetElementAt(__str_257, __str_256, __var_x->Rank()); +ss << __str_287; // 77 | continue; -ss << __str_294; +ss << __str_298; // 78 | } -ss << __str_200; +ss << __str_204; // 79 | for (var c = 0u; c < uniforms.input_channels_per_group; c += components) { -ss << __str_297; +ss << __str_301; ss << __param_components; -ss << __str_283; +ss << __str_287; // 80 | let input_channel = group * uniforms.input_channels_per_group + c; -ss << __str_298; +ss << __str_302; // 81 | #if is_channels_last if (__param_is_channels_last) { // 82 | let x_indices = x_indices_t(batch, id, ih, iw, input_channel / components); -ss << __str_299; +ss << __str_303; ss << __param_components; -ss << __str_12; +ss << __str_17; // 83 | let x_value = Accumulator(x.getByIndices(x_indices)); -ss << __str_300; -ss << __var_x->GetByIndices(__str_254); -ss << __str_12; +ss << __str_304; +ss << __var_x->GetByIndices(__str_258); +ss << __str_17; // 84 | #else } else { // 85 | #if components > 1 if (__param_components > 1) { // 86 | var x_value = Accumulator(0); -ss << __str_301; +ss << __str_305; // 87 | for (var lane = 0u; lane < components; lane++) { -ss << __str_302; +ss << __str_306; ss << __param_components; -ss << __str_303; +ss << __str_307; // 88 | let x_indices = x_indices_t(batch, input_channel + lane, id, ih, iw); -ss << __str_304; +ss << __str_308; // 89 | x_value[lane] = f32(x.getByIndices(x_indices)); -ss << __str_305; -ss << __var_x->GetByIndices(__str_254); -ss << __str_12; +ss << __str_309; +ss << __var_x->GetByIndices(__str_258); +ss << __str_17; // 90 | } -ss << __str_306; +ss << __str_310; // 91 | #else } else { // 92 | let x_indices = x_indices_t(batch, input_channel, id, ih, iw); -ss << __str_307; +ss << __str_311; // 93 | let x_value = f32(x.getByIndices(x_indices)); -ss << __str_308; -ss << __var_x->GetByIndices(__str_254); -ss << __str_12; +ss << __str_312; +ss << __var_x->GetByIndices(__str_258); +ss << __str_17; // 94 | #endif } // 95 | #endif @@ -251,49 +251,49 @@ ss << __str_12; // 96 | #if is_prepacked if (__param_is_prepacked) { // 97 | let w_indices = w_indices_t(kd, kh, kw, w_channel, input_channel / components); -ss << __str_309; +ss << __str_313; ss << __param_components; -ss << __str_12; +ss << __str_17; // 98 | #else } else { // 99 | let w_indices = w_indices_t(input_channel, w_channel, kd, kh, kw); -ss << __str_310; +ss << __str_314; // 100 | #endif } // 101 | value += x_value * Accumulator(w.getByIndices(w_indices)); -ss << __str_311; -ss << __var_w->GetByIndices(__str_255); -ss << __str_12; +ss << __str_315; +ss << __var_w->GetByIndices(__str_259); +ss << __str_17; // 102 | } -ss << __str_200; +ss << __str_204; // 103 | } -ss << __str_288; +ss << __str_292; // 104 | } -ss << __str_15; +ss << __str_20; // 105 | } -ss << __str_16; +ss << __str_11; // 106 | #if components > 1 if (__param_components > 1) { // 107 | var result = dot(value, Accumulator(1)); -ss << __str_312; +ss << __str_316; // 108 | #else } else { // 109 | var result = value; -ss << __str_313; +ss << __str_317; // 110 | #endif } // 111 | #if has_bias if (__param_has_bias) { // 112 | result += f32(bias.getByOffset(output_channel)); -ss << __str_314; -ss << __var_bias->GetByOffset(__str_256); -ss << __str_12; +ss << __str_318; +ss << __var_bias->GetByOffset(__str_260); +ss << __str_17; // 113 | #endif } // 114 | output.setByOffset(global_idx, output_value_t(result)); -ss << __str_237; -ss << __var_output->SetByOffset(__str_235, __str_257); -ss << __str_25; +ss << __str_241; +ss << __var_output->SetByOffset(__str_239, __str_261); +ss << __str_29; // 115 | } MainFunctionEnd(); ss << __str_5; diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h index 2ec07b7140896..84976d4a91fa1 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h @@ -43,331 +43,331 @@ Status ApplyTemplate<"nn/im2col_matmul.wgsl.template">(ShaderHelper& shader_help // 18 | // im2col access for src: [N, H_i, W_i, C_i / vec_size] // 19 | // Conceptual Matrix Shape: N * (H_o * W_o) x (K_h * K_w * C_i / vec_size) // 20 | fn load_src(batch : u32, m : u32, k_packed_idx : u32) -> src_value_t { -ss << __str_319; +ss << __str_323; // 21 | if (batch >= uniforms.batch || m >= uniforms.im2col_m || k_packed_idx * vec_size >= uniforms.im2col_k) { -ss << __str_320; +ss << __str_324; ss << __param_vec_size; -ss << __str_321; +ss << __str_325; // 22 | return src_value_t(); -ss << __str_322; +ss << __str_326; // 23 | } -ss << __str_16; +ss << __str_11; // 24 | ss << __str_5; // 25 | let channel_i_vec = uniforms.channel_i / vec_size; -ss << __str_323; +ss << __str_327; ss << __param_vec_size; -ss << __str_25; +ss << __str_29; // 26 | ss << __str_5; // 27 | // 1. Decompose M index (H_o * W_o) into (h_idx, w_idx) // 28 | let h_idx = m / uniforms.output_w; // Output H index (H_o) -ss << __str_324; +ss << __str_328; // 29 | let w_idx = m % uniforms.output_w; // Output W index (W_o) -ss << __str_325; +ss << __str_329; // 30 | ss << __str_5; // 31 | // 2. Decompose K index into (k_h, k_w, c_i_vec_idx) // 32 | let c_i_vec_idx = k_packed_idx % channel_i_vec; -ss << __str_326; +ss << __str_330; // 33 | let k_h_w_idx = k_packed_idx / channel_i_vec; -ss << __str_327; +ss << __str_331; // 34 | let k_h = k_h_w_idx / uniforms.kernel_w; // Kernel Row -ss << __str_328; +ss << __str_332; // 35 | let k_w = k_h_w_idx % uniforms.kernel_w; // Kernel Column -ss << __str_329; +ss << __str_333; // 36 | ss << __str_5; // 37 | // 3. Calculate the coordinate in the padded input tensor // 38 | let src_h_coord_padded = h_idx * uniforms.strides.x + k_h * uniforms.dilations.x; -ss << __str_330; +ss << __str_334; // 39 | let src_w_coord_padded = w_idx * uniforms.strides.y + k_w * uniforms.dilations.y; -ss << __str_331; +ss << __str_335; // 40 | ss << __str_5; // 41 | // 4. Calculate the coordinate in the original input tensor // 42 | let src_h_coord : i32 = i32(src_h_coord_padded) - i32(uniforms.pads.x); -ss << __str_332; +ss << __str_336; // 43 | let src_w_coord : i32 = i32(src_w_coord_padded) - i32(uniforms.pads.y); -ss << __str_333; +ss << __str_337; // 44 | ss << __str_5; // 45 | // 5. Check for padding/out-of-bounds // 46 | if (src_h_coord < 0 || src_h_coord >= i32(uniforms.src_h) || -ss << __str_334; +ss << __str_338; // 47 | src_w_coord < 0 || src_w_coord >= i32(uniforms.src_w)) { -ss << __str_335; +ss << __str_339; // 48 | return src_value_t(); -ss << __str_322; +ss << __str_326; // 49 | } -ss << __str_16; +ss << __str_11; // 50 | ss << __str_5; // 51 | // 6. Calculate final NHWC index // 52 | let src_idx = batch * uniforms.src_h * uniforms.src_w * channel_i_vec + -ss << __str_336; +ss << __str_340; // 53 | u32(src_h_coord) * uniforms.src_w * channel_i_vec + -ss << __str_337; +ss << __str_341; // 54 | u32(src_w_coord) * channel_i_vec + -ss << __str_338; +ss << __str_342; // 55 | c_i_vec_idx; -ss << __str_339; +ss << __str_343; // 56 | return src.getByOffset(src_idx); -ss << __str_340; -ss << __var_src->GetByOffset(__str_315); -ss << __str_25; +ss << __str_344; +ss << __var_src->GetByOffset(__str_319); +ss << __str_29; // 57 | } -ss << __str_341; +ss << __str_345; // 58 | ss << __str_5; // 59 | // weight shape: [Co, K_h, K_w, C_i / vec_size] (CoHWCi) // 60 | fn load_weight(n : u32, k_packed_idx : u32) -> weight_value_t { -ss << __str_342; +ss << __str_346; // 61 | if (n < uniforms.im2col_n && k_packed_idx < uniforms.im2col_k / vec_size) { -ss << __str_343; +ss << __str_347; ss << __param_vec_size; -ss << __str_283; +ss << __str_287; // 62 | let weight_idx = n * uniforms.im2col_k / vec_size + -ss << __str_344; +ss << __str_348; ss << __param_vec_size; -ss << __str_345; +ss << __str_349; // 63 | k_packed_idx; -ss << __str_346; +ss << __str_350; // 64 | return weight.getByOffset(weight_idx); -ss << __str_347; -ss << __var_weight->GetByOffset(__str_316); -ss << __str_25; +ss << __str_351; +ss << __var_weight->GetByOffset(__str_320); +ss << __str_29; // 65 | } -ss << __str_16; +ss << __str_11; // 66 | return weight_value_t(); -ss << __str_348; +ss << __str_352; // 67 | } -ss << __str_341; +ss << __str_345; // 68 | ss << __str_5; // 69 | fn load_bias(n : u32) -> output_element_t { -ss << __str_349; +ss << __str_353; // 70 | #if has_bias if (__param_has_bias) { // 71 | if (n < uniforms.im2col_n) { -ss << __str_350; +ss << __str_354; // 72 | return output_element_t(bias[n]); -ss << __str_351; +ss << __str_355; // 73 | } -ss << __str_16; +ss << __str_11; // 74 | #endif } // 75 | return output_element_t(); -ss << __str_352; +ss << __str_356; // 76 | } -ss << __str_341; +ss << __str_345; // 77 | ss << __str_5; // 78 | // output shape: [N, H_o, W_o, C_o] (NHWC) // 79 | fn write_output(batch : u32, m : u32, n : u32, value : output_element_t) { -ss << __str_353; +ss << __str_357; // 80 | if (batch < uniforms.batch && m < uniforms.im2col_m && n < uniforms.im2col_n) { -ss << __str_354; +ss << __str_358; // 81 | let output_idx = batch * uniforms.im2col_m * uniforms.im2col_n + -ss << __str_355; +ss << __str_359; // 82 | m * uniforms.im2col_n + -ss << __str_356; +ss << __str_360; // 83 | n; -ss << __str_357; +ss << __str_361; // 84 | output.setByOffset(output_idx, value); -ss << __str_24; -ss << __var_output->SetByOffset(__str_317, __str_318); -ss << __str_25; +ss << __str_28; +ss << __var_output->SetByOffset(__str_321, __str_322); +ss << __str_29; // 85 | } -ss << __str_16; +ss << __str_11; // 86 | } -ss << __str_341; +ss << __str_345; // 87 | ss << __str_5; // 88 | const TILE_M_SIZE : u32 = tile_m; -ss << __str_358; +ss << __str_362; ss << __param_tile_m; -ss << __str_25; +ss << __str_29; // 89 | const TILE_N_SIZE : u32 = tile_n; -ss << __str_359; +ss << __str_363; ss << __param_tile_n; -ss << __str_25; +ss << __str_29; // 90 | // In dimension K, the tile consists of 16 scalars, requiring `16 / vec_size` vector loads. ss << __str_5; // 91 | const TILE_K_VEC_SIZE : u32 = 16 / vec_size; -ss << __str_360; +ss << __str_364; ss << __param_vec_size; -ss << __str_25; +ss << __str_29; // 92 | // In dimensions M and N, since a workgroup has 64 threads, it advances by `64 / TILE_K_VEC_SIZE`. ss << __str_5; // 93 | const ADVANCE_DIM = 64 / TILE_K_VEC_SIZE; -ss << __str_361; +ss << __str_365; // 94 | ss << __str_5; // 95 | var src_tile : array, TILE_K_VEC_SIZE>; -ss << __str_362; +ss << __str_366; // 96 | var weight_tile : array, TILE_K_VEC_SIZE>; -ss << __str_363; +ss << __str_367; // 97 | ss << __str_5; // 98 | $MAIN { MainFunctionStart(); ss << __str_5; // 99 | let batch = workgroup_idx / (uniforms.M_tiles * uniforms.N_tiles); -ss << __str_364; +ss << __str_368; // 100 | let m_global_base = ((workgroup_idx / uniforms.N_tiles) % uniforms.M_tiles) * TILE_M_SIZE; -ss << __str_365; +ss << __str_369; // 101 | let n_global_base = (workgroup_idx % uniforms.N_tiles) * TILE_N_SIZE; -ss << __str_366; +ss << __str_370; // 102 | ss << __str_5; // 103 | var results : array; -ss << __str_367; +ss << __str_371; // 104 | for (var k_idx = 0u; k_idx < uniforms.K_tiles; k_idx++) { -ss << __str_368; +ss << __str_372; // 105 | for (var src_m = 0u; src_m < TILE_M_SIZE; src_m += ADVANCE_DIM) { -ss << __str_369; +ss << __str_373; // 106 | // Loads a 64 vec of src into the workgroup memory. ss << __str_5; // 107 | let load_src_m = src_m + local_idx / TILE_K_VEC_SIZE; -ss << __str_370; +ss << __str_374; // 108 | let load_src_k = local_idx % TILE_K_VEC_SIZE; -ss << __str_371; +ss << __str_375; // 109 | ss << __str_5; // 110 | src_tile[load_src_k][load_src_m] = load_src(batch, -ss << __str_372; +ss << __str_376; // 111 | m_global_base + load_src_m, -ss << __str_373; +ss << __str_377; // 112 | k_idx * TILE_K_VEC_SIZE + load_src_k); -ss << __str_374; +ss << __str_378; // 113 | } -ss << __str_15; +ss << __str_20; // 114 | ss << __str_5; // 115 | for (var weight_n = 0u; weight_n < TILE_N_SIZE; weight_n += ADVANCE_DIM) { -ss << __str_375; +ss << __str_379; // 116 | // Loads a 64 vec of weight into the workgroup memory. ss << __str_5; // 117 | let load_weight_n = weight_n + local_idx / TILE_K_VEC_SIZE; -ss << __str_376; +ss << __str_380; // 118 | let load_weight_k = local_idx % TILE_K_VEC_SIZE; -ss << __str_377; +ss << __str_381; // 119 | ss << __str_5; // 120 | weight_tile[load_weight_k][load_weight_n] = load_weight(n_global_base + load_weight_n, -ss << __str_378; +ss << __str_382; // 121 | k_idx * TILE_K_VEC_SIZE + load_weight_k); -ss << __str_379; +ss << __str_383; // 122 | } -ss << __str_15; +ss << __str_20; // 123 | workgroupBarrier(); -ss << __str_22; +ss << __str_26; // 124 | ss << __str_5; // 125 | for (var inner_k_idx = 0u; inner_k_idx < TILE_K_VEC_SIZE; inner_k_idx++) { -ss << __str_380; +ss << __str_384; // 126 | let weight_data = weight_tile[inner_k_idx][local_idx]; -ss << __str_381; +ss << __str_385; // 127 | #if use_subgroup if (__param_use_subgroup) { // 128 | // The subgroup size is fixed at 32; wrap the index so the 16x64 tiling does not read out of bounds. // 129 | let src_data = src_tile[inner_k_idx][sg_id % TILE_M_SIZE]; -ss << __str_382; +ss << __str_386; // 130 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { -ss << __str_383; +ss << __str_387; // 131 | #if vec_size == 1 if (__param_vec_size == 1) { // 132 | results[m_idx] += output_element_t(weight_data * subgroupShuffle(src_data, m_idx)); -ss << __str_384; +ss << __str_388; // 133 | #else } else { // 134 | results[m_idx] += output_element_t(dot(weight_data, subgroupShuffle(src_data, m_idx))); -ss << __str_385; +ss << __str_389; // 135 | #endif } // 136 | } -ss << __str_288; +ss << __str_292; // 137 | #else } else { // 138 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { -ss << __str_383; +ss << __str_387; // 139 | #if vec_size == 1 if (__param_vec_size == 1) { // 140 | results[m_idx] += output_element_t(weight_data * src_tile[inner_k_idx][m_idx]); -ss << __str_386; +ss << __str_390; // 141 | #else } else { // 142 | results[m_idx] += output_element_t(dot(weight_data, src_tile[inner_k_idx][m_idx])); -ss << __str_387; +ss << __str_391; // 143 | #endif } // 144 | } -ss << __str_288; +ss << __str_292; // 145 | #endif } // 146 | } -ss << __str_15; +ss << __str_20; // 147 | workgroupBarrier(); -ss << __str_22; +ss << __str_26; // 148 | } -ss << __str_16; +ss << __str_11; // 149 | ss << __str_5; // 150 | let m_base = m_global_base; -ss << __str_388; +ss << __str_392; // 151 | let n_base = n_global_base + local_idx; -ss << __str_389; +ss << __str_393; // 152 | ss << __str_5; // 153 | let bias = load_bias(n_base); -ss << __str_390; +ss << __str_394; // 154 | for (var m_idx = 0u; m_idx < TILE_M_SIZE; m_idx++) { -ss << __str_391; +ss << __str_395; // 155 | var output_data = results[m_idx] + bias; -ss << __str_392; +ss << __str_396; // 156 | #if activation_kind == 1 if (__param_activation_kind == 1) { // 157 | output_data = max(output_data, output_element_t(0)); -ss << __str_393; +ss << __str_397; // 158 | #elif activation_kind == 2 } else if (__param_activation_kind == 2) { // 159 | output_data = output_element_t(1) / (output_element_t(1) + exp(-output_data)); -ss << __str_394; +ss << __str_398; // 160 | #elif activation_kind == 3 } else if (__param_activation_kind == 3) { // 161 | output_data = clamp(output_data, output_element_t(uniforms.activation_param_0), output_element_t(uniforms.activation_param_1)); -ss << __str_395; +ss << __str_399; // 162 | #elif activation_kind == 4 } else if (__param_activation_kind == 4) { // 163 | output_data = clamp(output_element_t(uniforms.activation_param_0) * output_data + output_element_t(uniforms.activation_param_1), output_element_t(0), output_element_t(1)); -ss << __str_396; +ss << __str_400; // 164 | #elif activation_kind == 5 } else if (__param_activation_kind == 5) { // 165 | output_data = select(output_element_t(uniforms.activation_param_0) * output_data, output_data, output_data >= output_element_t(0)); -ss << __str_397; +ss << __str_401; // 166 | #elif activation_kind == 6 } else if (__param_activation_kind == 6) { // 167 | output_data = tanh(output_data); -ss << __str_398; +ss << __str_402; // 168 | #elif activation_kind == 7 } else if (__param_activation_kind == 7) { // 169 | #if quick_gelu_unit_alpha if (__param_quick_gelu_unit_alpha) { // 170 | output_data = output_data * (output_element_t(1) / (output_element_t(1) + exp(-output_data))); -ss << __str_399; +ss << __str_403; // 171 | #else } else { // 172 | output_data = output_data * (output_element_t(1) / (output_element_t(1) + exp(-(output_element_t(uniforms.activation_param_0) * output_data)))); -ss << __str_400; +ss << __str_404; // 173 | #endif } // 174 | #endif } // 175 | write_output(batch, m_base + m_idx, n_base, output_data); -ss << __str_401; +ss << __str_405; // 176 | } -ss << __str_16; +ss << __str_11; // 177 | } // MAIN MainFunctionEnd(); ss << __str_5; diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h index e934c4592df46..06e403c405f07 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h @@ -19,105 +19,105 @@ Status ApplyTemplate<"tensor/oihw_to_ohwi.wgsl.template">(ShaderHelper& shader_h // 4 | #use .getByOffset .setByOffset // 5 | // 6 | fn load_src(co : u32, ci : u32, h_w : u32) -> src_element_t { -ss << __str_403; +ss << __str_407; // 7 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_404; +ss << __str_408; // 8 | let offset = co * uniforms.I * uniforms.H * uniforms.W + -ss << __str_405; +ss << __str_409; // 9 | ci * uniforms.H * uniforms.W + -ss << __str_406; +ss << __str_410; // 10 | h_w; -ss << __str_407; +ss << __str_411; // 11 | return src.getByOffset(offset); -ss << __str_347; -ss << __var_src->GetByOffset(__str_402); -ss << __str_25; +ss << __str_351; +ss << __var_src->GetByOffset(__str_406); +ss << __str_29; // 12 | } -ss << __str_16; +ss << __str_11; // 13 | return src_element_t(); -ss << __str_408; +ss << __str_412; // 14 | } -ss << __str_341; +ss << __str_345; // 15 | ss << __str_5; // 16 | fn write_output(co : u32, h_w : u32, ci : u32, value : output_element_t) { -ss << __str_409; +ss << __str_413; // 17 | if (co < uniforms.O && ci < uniforms.I && h_w < uniforms.H * uniforms.W) { -ss << __str_404; +ss << __str_408; // 18 | let offset = co * uniforms.H * uniforms.W * uniforms.I + -ss << __str_410; +ss << __str_414; // 19 | h_w * uniforms.I + -ss << __str_411; +ss << __str_415; // 20 | ci; -ss << __str_412; +ss << __str_416; // 21 | output.setByOffset(offset, value); -ss << __str_24; -ss << __var_output->SetByOffset(__str_402, __str_318); -ss << __str_25; +ss << __str_28; +ss << __var_output->SetByOffset(__str_406, __str_322); +ss << __str_29; // 22 | } -ss << __str_16; +ss << __str_11; // 23 | } -ss << __str_341; +ss << __str_345; // 24 | ss << __str_5; // 25 | var data_cache : array, 4>; -ss << __str_413; +ss << __str_417; // 26 | ss << __str_5; // 27 | $MAIN { MainFunctionStart(); ss << __str_5; // 28 | let group_co : u32 = workgroup_idx / uniforms.Ci_tiles; -ss << __str_414; +ss << __str_418; // 29 | let group_ci : u32 = (workgroup_idx % uniforms.Ci_tiles) * 64; -ss << __str_415; +ss << __str_419; // 30 | ss << __str_5; // 31 | if (group_co >= uniforms.O || group_ci >= uniforms.I) { -ss << __str_416; +ss << __str_420; // 32 | return; -ss << __str_417; +ss << __str_10; // 33 | } -ss << __str_16; +ss << __str_11; // 34 | ss << __str_5; // 35 | for (var h_w_idx = 0u; h_w_idx < uniforms.H_W_tiles; h_w_idx++) { -ss << __str_418; +ss << __str_421; // 36 | // load ss << __str_5; // 37 | for (var ci_idx = 0u; ci_idx < 64u; ci_idx += 16u) { -ss << __str_419; +ss << __str_422; // 38 | let load_ci_idx = ci_idx + local_idx / 4; -ss << __str_420; +ss << __str_423; // 39 | let load_h_w_idx = local_idx % 4; -ss << __str_421; +ss << __str_424; // 40 | ss << __str_5; // 41 | data_cache[load_h_w_idx][load_ci_idx] = load_src(group_co, -ss << __str_422; +ss << __str_425; // 42 | group_ci + load_ci_idx, -ss << __str_423; +ss << __str_426; // 43 | h_w_idx * 4 + load_h_w_idx); -ss << __str_424; +ss << __str_427; // 44 | } -ss << __str_15; +ss << __str_20; // 45 | workgroupBarrier(); -ss << __str_22; +ss << __str_26; // 46 | ss << __str_5; // 47 | // store // 48 | for (var local_h_w_idx = 0u; local_h_w_idx < 4u; local_h_w_idx++) { -ss << __str_425; +ss << __str_428; // 49 | let output_data = data_cache[local_h_w_idx][local_idx]; -ss << __str_426; +ss << __str_429; // 50 | write_output(group_co, h_w_idx * 4 + local_h_w_idx, group_ci + local_idx, output_data); -ss << __str_427; +ss << __str_430; // 51 | } -ss << __str_15; +ss << __str_20; // 52 | workgroupBarrier(); -ss << __str_22; +ss << __str_26; // 53 | } -ss << __str_16; +ss << __str_11; // 54 | } // MAIN MainFunctionEnd(); ss << __str_5; diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h index 90d556155c4b3..a40daa675052e 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h @@ -34,145 +34,145 @@ Status ApplyTemplate<"tensor/pad.wgsl.template">(ShaderHelper& shader_helper, Te MainFunctionStart(); ss << __str_5; // 15 | guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); -ss << __str_237; -ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_233); -ss << __str_25; +ss << __str_241; +ss << shader_helper.GuardAgainstOutOfBoundsWorkgroupSizes(__str_237); +ss << __str_29; // 16 | ss << __str_5; // 17 | let constant_value = -ss << __str_433; +ss << __str_436; // 18 | #if is_float16 if (__param_is_float16) { // 19 | bitcast>(uniforms.constant_value)[0]; -ss << __str_434; +ss << __str_437; // 20 | #else } else { // 21 | bitcast(uniforms.constant_value); -ss << __str_435; +ss << __str_438; // 22 | #endif } // 23 | // 24 | #if dim_value_zero if (__param_dim_value_zero) { // 25 | output[global_idx] = constant_value; -ss << __str_436; +ss << __str_439; // 26 | #else } else { // 27 | let output_indices = output.offsetToIndices(global_idx); -ss << __str_261; -ss << __var_output->OffsetToIndices(__str_235); -ss << __str_25; +ss << __str_265; +ss << __var_output->OffsetToIndices(__str_239); +ss << __str_29; // 28 | var input_index = u32(0); -ss << __str_437; +ss << __str_440; // 29 | var use_pad_value = false; -ss << __str_438; +ss << __str_441; // 30 | var in_coord = i32(0); -ss << __str_439; +ss << __str_442; // 31 | ss << __str_5; // 32 | for (var dim = 0; dim < output.rank && !use_pad_value; dim++) { -ss << __str_440; +ss << __str_443; ss << __var_output->Rank(); -ss << __str_441; +ss << __str_444; // 33 | let output_index = i32(getElementAt(output_indices, dim, output.rank)); -ss << __str_442; -ss << GetElementAt(__str_243, __str_428, __var_output->Rank()); -ss << __str_12; +ss << __str_445; +ss << GetElementAt(__str_247, __str_431, __var_output->Rank()); +ss << __str_17; // 34 | let lower_pads = getElementAt(uniforms.lower_pads, dim, output.rank); -ss << __str_443; -ss << GetElementAt(__str_429, __str_428, __var_output->Rank()); -ss << __str_25; +ss << __str_446; +ss << GetElementAt(__str_432, __str_431, __var_output->Rank()); +ss << __str_29; // 35 | let data_shape = i32(getElementAt(uniforms.data_shape, dim, output.rank)); -ss << __str_444; -ss << GetElementAt(__str_430, __str_428, __var_output->Rank()); -ss << __str_12; +ss << __str_447; +ss << GetElementAt(__str_433, __str_431, __var_output->Rank()); +ss << __str_17; // 36 | #if pad_mode == PAD_MODE_CONSTANT if (__param_pad_mode == 0) { // 37 | if (output_index < lower_pads || output_index >= data_shape + lower_pads) { -ss << __str_445; +ss << __str_448; // 38 | use_pad_value = true; -ss << __str_446; +ss << __str_449; // 39 | #elif pad_mode == PAD_MODE_EDGE } else if (__param_pad_mode == 2) { // 40 | if (output_index < lower_pads) { -ss << __str_447; +ss << __str_450; // 41 | in_coord = 0; -ss << __str_448; +ss << __str_451; // 42 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_449; +ss << __str_452; // 43 | in_coord = data_shape - 1; -ss << __str_450; +ss << __str_453; // 44 | #elif pad_mode == PAD_MODE_REFLECT } else if (__param_pad_mode == 1) { // 45 | if (output_index < lower_pads || output_index >= data_shape + lower_pads) { -ss << __str_445; +ss << __str_448; // 46 | in_coord = output_index - lower_pads; -ss << __str_451; +ss << __str_454; // 47 | if (in_coord < 0) { -ss << __str_452; +ss << __str_455; // 48 | in_coord = -in_coord; -ss << __str_453; +ss << __str_456; // 49 | } -ss << __str_288; +ss << __str_292; // 50 | let _2n_1 = 2 * (data_shape - 1); -ss << __str_454; +ss << __str_457; // 51 | in_coord = in_coord % _2n_1; -ss << __str_455; +ss << __str_458; // 52 | if (in_coord >= data_shape) { -ss << __str_456; +ss << __str_459; // 53 | in_coord = _2n_1 - in_coord; -ss << __str_457; +ss << __str_460; // 54 | } -ss << __str_288; +ss << __str_292; // 55 | #else // PAD_MODE_WRAP } else { // 56 | if (output_index < lower_pads) { -ss << __str_447; +ss << __str_450; // 57 | in_coord = data_shape + output_index - lower_pads; -ss << __str_458; +ss << __str_461; // 58 | } else if (output_index >= data_shape + lower_pads) { -ss << __str_449; +ss << __str_452; // 59 | in_coord = output_index - data_shape - lower_pads; -ss << __str_459; +ss << __str_462; // 60 | #endif // pad_mode } // 61 | } else { -ss << __str_460; +ss << __str_463; // 62 | in_coord = output_index - lower_pads; -ss << __str_461; +ss << __str_464; // 63 | } -ss << __str_15; +ss << __str_20; // 64 | ss << __str_5; // 65 | #if pad_mode == PAD_MODE_WRAP if (__param_pad_mode == 3) { // 66 | in_coord = ((in_coord % data_shape) + data_shape) % data_shape; -ss << __str_462; +ss << __str_465; // 67 | #endif } // 68 | // 69 | input_index += select(u32(in_coord) -ss << __str_463; +ss << __str_466; // 70 | #if output.rank > 1 if (__var_output->Rank() > 1) { // 71 | * getElementAt(uniforms.data_stride, dim, output.rank - 1) -ss << __str_464; -ss << GetElementAt(__str_431, __str_428, __var_output->Rank() - 1); +ss << __str_467; +ss << GetElementAt(__str_434, __str_431, __var_output->Rank() - 1); ss << __str_5; // 72 | #endif } // 73 | , u32(in_coord), dim == output.rank - 1); -ss << __str_465; +ss << __str_468; ss << __var_output->Rank(); -ss << __str_466; +ss << __str_469; // 74 | } -ss << __str_16; +ss << __str_11; // 75 | ss << __str_5; // 76 | output.setByOffset(global_idx, select(data[input_index], constant_value, use_pad_value)); -ss << __str_237; -ss << __var_output->SetByOffset(__str_235, __str_432); -ss << __str_25; +ss << __str_241; +ss << __var_output->SetByOffset(__str_239, __str_435); +ss << __str_29; // 77 | #endif } // 78 | } // MAIN diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h index 1a7809deb37b9..3bb8396be1a50 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h @@ -31,18 +31,18 @@ std::string pass_as_string(T&& v) { } } // namespace wgsl_detail -#include "wgsl_template_gen/string_table.h" // 05862fc625af56779334dcecc6e032931ea9864f12f60bb4d05a797c56cda2e6 +#include "wgsl_template_gen/string_table.h" // 0164a7040addea01a1f6e9ff048e816229c3b38c2d8eee7625c241df23ef3072 // Include template implementations -#include "wgsl_template_gen/generated/math/matmul_gemv.h" // 59e4865f8dfb09777dff88b1b5983d18704a0713a43e26a48ee0479420b5dba5 -#include "wgsl_template_gen/generated/math/subgroup_matrix_gemm_8x16x16.h" // c1d802ca557f329f392fcd71ac806d73af886bcc1a4e0c19f4c39affb0bfd99e -#include "wgsl_template_gen/generated/math/subgroup_matrix_matmul.h" // 34ecf1392f9e66b1e8573f76daf878b219e991b50c55301d54cdcc8ee91b6e53 -#include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // 4417a0eb6c04ce936a4be2d7a3080254b26ef44cccae021de9a5f2c2c40ab902 -#include "wgsl_template_gen/generated/nn/conv_transpose3d.h" // 6ac2b5204cde9e0f987817c38e684629d1aa3aa2105f858e9d1a044864029a66 -#include "wgsl_template_gen/generated/nn/im2col_matmul.h" // dacf36da447a95df88adaca6763f006767a0ecf4d234d14b883ff0aa27a2b96f -#include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // b4979ab8ce0d5a0ff32d75b6bf2cbc3919807cb8f870d95537f2ff88bcd23155 -#include "wgsl_template_gen/generated/tensor/pad.h" // e20a8f93e1f0622e0d355fe97380e7dd61a95b09d9736b44430c6bba9df63cbb +#include "wgsl_template_gen/generated/math/matmul_gemv.h" // 512f5afa8421c4630912b0d87d46061d43e1b3226fbbc6c993241b285adf2fc7 +#include "wgsl_template_gen/generated/math/subgroup_matrix_gemm_8x16x16.h" // 4251accac15bcce4e64a1ec7582de07ee5626d6cfe05527a9187586ade097870 +#include "wgsl_template_gen/generated/math/subgroup_matrix_matmul.h" // 82baad53c5eb895073a12f45a2bb5ff6e073bbe089448e0d5ba52255c4ee7ce7 +#include "wgsl_template_gen/generated/math/subgroup_matrix_matmul_pad_b.h" // 4b28d0cbe9ca50020bbcb75121e787982c268992efdac401b821dee7c59c604e +#include "wgsl_template_gen/generated/nn/conv_transpose3d.h" // 33bd52e16ae1e748ef45b492ae4730f520eeb9706cab055391d296ec54265bab +#include "wgsl_template_gen/generated/nn/im2col_matmul.h" // 7185c8aeed8391c916822b35d16fcbfdd342e4c3cc5ea9a4d09667a2651a3c08 +#include "wgsl_template_gen/generated/tensor/oihw_to_ohwi.h" // 473c95130b9ad534439ed746bf37df45be692792142d57c48db3cea0e864f352 +#include "wgsl_template_gen/generated/tensor/pad.h" // b60f1767797524bcf834535f8e4f8d214c0c8a6f6d7c4e12ceae44c39abec554 #pragma pop_macro("MainFunctionStart") #pragma pop_macro("MainFunctionEnd") \ No newline at end of file diff --git a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h index 8ea5a5f2cd151..9764cea3df588 100644 --- a/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h +++ b/tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h @@ -6,470 +6,473 @@ #endif // String table constants -constexpr const char* __str_0 = "k"; +constexpr const char* __str_0 = "row * uniforms.K + k"; constexpr const char* __str_1 = "k * n_vec_count + n_vec"; -constexpr const char* __str_2 = "n_vec"; +constexpr const char* __str_2 = "row * n_vec_count + n_vec"; constexpr const char* __str_3 = "output_value_t(partials[local_idx])"; constexpr const char* __str_4 = "var partials: array, workgroup_size_x * workgroup_size_y>;\n"; constexpr const char* __str_5 = "\n"; -constexpr const char* __str_6 = " let n_vec = workgroup_idx * workgroup_size_x + local_id.x;\n"; -constexpr const char* __str_7 = " let n_vec_count = uniforms.N / 4u;\n"; -constexpr const char* __str_8 = " var acc = vec4(0);\n"; -constexpr const char* __str_9 = " if (n_vec < n_vec_count) {\n"; -constexpr const char* __str_10 = " for (var k = local_id.y; k < uniforms.K; k += workgroup_size_y) {\n"; -constexpr const char* __str_11 = " let av = f32("; -constexpr const char* __str_12 = ");\n"; -constexpr const char* __str_13 = " let bv = vec4("; -constexpr const char* __str_14 = " acc = fma(vec4(av), bv, acc);\n"; -constexpr const char* __str_15 = " }\n"; -constexpr const char* __str_16 = " }\n"; -constexpr const char* __str_17 = " partials[local_idx] = acc;\n"; -constexpr const char* __str_18 = " workgroupBarrier();\n"; -constexpr const char* __str_19 = " for (var stride = workgroup_size_y / 2u; stride > 0u; stride /= 2u) {\n"; -constexpr const char* __str_20 = " if (local_id.y < stride) {\n"; -constexpr const char* __str_21 = " partials[local_idx] += partials[local_idx + stride * workgroup_size_x];\n"; -constexpr const char* __str_22 = " workgroupBarrier();\n"; -constexpr const char* __str_23 = " if (local_id.y == 0u && n_vec < n_vec_count) {\n"; -constexpr const char* __str_24 = " "; -constexpr const char* __str_25 = ";\n"; -constexpr const char* __str_26 = "out_base + i"; -constexpr const char* __str_27 = "val"; -constexpr const char* __str_28 = "const kSgMatM: u32 = u32("; -constexpr const char* __str_29 = "const kSgMatN: u32 = u32("; -constexpr const char* __str_30 = "const kSgMatK: u32 = u32("; -constexpr const char* __str_31 = "const kSgMatCountM: u32 = u32("; -constexpr const char* __str_32 = "const kSgMatCountN: u32 = u32("; -constexpr const char* __str_33 = "const kTileM: u32 = kSgMatM * kSgMatCountM;\n"; -constexpr const char* __str_34 = "const kTileN: u32 = kSgMatN * kSgMatCountN;\n"; -constexpr const char* __str_35 = "const kSplitK: u32 = u32("; -constexpr const char* __str_36 = "const kSubgroupSize: u32 = 32;\n"; -constexpr const char* __str_37 = "var scratch: array;\n"; -constexpr const char* __str_38 = " let num_n_tile = (uniforms.N + kTileN - 1u) / kTileN;\n"; -constexpr const char* __str_39 = " let n_tile = workgroup_idx % num_n_tile;\n"; -constexpr const char* __str_40 = " let m_tile = workgroup_idx / num_n_tile;\n"; -constexpr const char* __str_41 = " let n_tile_start = n_tile * kTileN;\n"; -constexpr const char* __str_42 = " let m_tile_start = m_tile * kTileM;\n"; -constexpr const char* __str_43 = " let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN);\n"; -constexpr const char* __str_44 = " let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM);\n"; -constexpr const char* __str_45 = " let k_blocks = uniforms.K / kSgMatK;\n"; -constexpr const char* __str_46 = " let sg_index = local_idx / kSubgroupSize;\n"; -constexpr const char* __str_47 = " let sg_lane = local_idx % kSubgroupSize;\n"; -constexpr const char* __str_48 = " let sg_scratch_base = sg_index * kTileM * kTileN;\n"; -constexpr const char* __str_49 = " var sg_mat_c0_0: subgroup_matrix_result