diff --git a/.github/workflows/android.yml b/.github/workflows/android.yml index e1e79a9402fe5..0a79b089d4192 100644 --- a/.github/workflows/android.yml +++ b/.github/workflows/android.yml @@ -24,11 +24,7 @@ permissions: jobs: AndroidBinarySizeCheckJob_MinimalBaseline: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=AndroidBinarySizeCheckJob_MinimalBaseline-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 env: CCACHE_DIR: ~/.cache/ccache # explicitly set to prevent any fallback to `~/.ccache` @@ -119,11 +115,7 @@ jobs: shell: bash android_nnapi_ep: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=android_nnapi_ep-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 env: CCACHE_DIR: ~/.cache/ccache # explicitly set to prevent any fallback to `~/.ccache` @@ -234,11 +226,7 @@ jobs: android_cpu_ep: name: Android CI Pipeline - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=android_cpu_ep-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 env: CCACHE_DIR: ~/.cache/ccache # explicitly set to prevent any fallback to `~/.ccache` diff --git a/.github/workflows/cffconvert.yml b/.github/workflows/cffconvert.yml index ae53c32fefe2a..7e7b8d03ac315 100644 --- a/.github/workflows/cffconvert.yml +++ b/.github/workflows/cffconvert.yml @@ -9,11 +9,7 @@ on: jobs: validate: name: "validate" - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=cffconvert-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 steps: - name: Check out a copy of the repository diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 8de673d769677..28bc77e9ec705 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -23,11 +23,7 @@ concurrency: jobs: analyze: name: Analyze - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=codeql-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}-${{ matrix.language}}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: actions: read diff --git a/.github/workflows/gradle-wrapper-validation.yml b/.github/workflows/gradle-wrapper-validation.yml index f00dfa87632b0..73c0d31fd0c9c 100644 --- a/.github/workflows/gradle-wrapper-validation.yml +++ b/.github/workflows/gradle-wrapper-validation.yml @@ -13,11 +13,7 @@ on: jobs: validation: name: "Validation" - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=gradle-wrapper-validation-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 diff --git a/.github/workflows/labeler.yml b/.github/workflows/labeler.yml index 08a7573a06552..89c91b080923e 100644 --- a/.github/workflows/labeler.yml +++ b/.github/workflows/labeler.yml @@ -8,11 +8,7 @@ permissions: jobs: triage: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=labeler-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 steps: - uses: github/issue-labeler@ccca1e16fcb66762200e616fc29b57b940c1cb67 # v3.5 diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index a74a78bced5df..4c90725baa0e8 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -39,11 +39,7 @@ jobs: lint-python-format: # Required workflow name: Python format - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=lint-python-format-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: contents: read diff --git a/.github/workflows/linux-wasm-ci-build-and-test-workflow.yml b/.github/workflows/linux-wasm-ci-build-and-test-workflow.yml index 5d05a92892aba..2d6a4527e35ef 100644 --- a/.github/workflows/linux-wasm-ci-build-and-test-workflow.yml +++ b/.github/workflows/linux-wasm-ci-build-and-test-workflow.yml @@ -41,11 +41,7 @@ on: jobs: build-wasm: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-linux-large", - "JobId=build-wasm-${{inputs.job_name}}-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-32-amd-scus-ss] timeout-minutes: 120 env: buildArch: x64 diff --git a/.github/workflows/linux_ci.yml b/.github/workflows/linux_ci.yml index e94c71f146ed7..d82f0ccd61feb 100644 --- a/.github/workflows/linux_ci.yml +++ b/.github/workflows/linux_ci.yml @@ -43,7 +43,7 @@ jobs: name: Build Linux x64 Debug (ASan) uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + pool_name: "ort-github-linux-16-amd-wcus-ss" build_config: Debug architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/inference/x86_64/default/cpu/Dockerfile @@ -58,7 +58,7 @@ jobs: name: Build Linux x64 Release uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + pool_name: "ort-github-linux-16-amd-wcus-ss" build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cpu @@ -74,7 +74,7 @@ jobs: name: Build Linux x64 Release (Python 3.14) uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + pool_name: "ort-github-linux-16-amd-wcus-ss" build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cpu @@ -93,7 +93,7 @@ jobs: name: Build Linux x64 Release with training uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + pool_name: "ort-github-linux-16-amd-wcus-ss" build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cpu @@ -108,7 +108,7 @@ jobs: name: Build Linux x64 MinSizeRel (Exceptions Disabled) uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + pool_name: "ort-github-linux-16-amd-wcus-ss" build_config: MinSizeRel architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/inference/x86_64/default/cpu/Dockerfile @@ -125,7 +125,7 @@ jobs: name: Build Linux arm64 Debug uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-ARM-CPU" + pool_name: "ort-github-linux-16-arm64-eus2-ss" build_config: Debug architecture: arm64 dockerfile_path: tools/ci_build/github/linux/docker/inference/aarch64/default/cpu/Dockerfile @@ -140,7 +140,7 @@ jobs: name: Build Linux arm64 Release uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-ARM-CPU" + pool_name: "ort-github-linux-16-arm64-eus2-ss" build_config: Release architecture: arm64 dockerfile_path: tools/ci_build/github/linux/docker/inference/aarch64/python/cpu/Dockerfile diff --git a/.github/workflows/linux_cuda_ci.yml b/.github/workflows/linux_cuda_ci.yml index d0d37b2c94bab..695b1b741b802 100644 --- a/.github/workflows/linux_cuda_ci.yml +++ b/.github/workflows/linux_cuda_ci.yml @@ -23,7 +23,7 @@ jobs: # This job runs on a CPU node using the reusable build workflow uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" # Build pool + pool_name: "ort-github-linux-16-amd-wcus-ss" # Build pool build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cuda @@ -55,10 +55,7 @@ jobs: test-linux-cuda-x64-release: name: Test Linux CUDA x64 Release needs: build-linux-cuda-x64-release - runs-on: - - self-hosted - - "1ES.Pool=onnxruntime-github-linux-a10" - - "JobId=test-linux-cuda-x64-release-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" + runs-on: [ort-github-linux-6-a10-scus-ss] timeout-minutes: 210 permissions: contents: read diff --git a/.github/workflows/linux_cuda_no_cudnn.yml b/.github/workflows/linux_cuda_no_cudnn.yml index 505ef4b646572..780ae7cc1ea29 100644 --- a/.github/workflows/linux_cuda_no_cudnn.yml +++ b/.github/workflows/linux_cuda_no_cudnn.yml @@ -27,7 +27,7 @@ jobs: name: Build Linux CUDA x64 Release without cuDNN link uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + pool_name: "ort-github-linux-16-amd-wcus-ss" build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cuda @@ -57,10 +57,7 @@ jobs: smoke-linux-cuda-no-cudnn-x64-release: name: Smoke Linux CUDA x64 Release without cuDNN runtime use needs: build-linux-cuda-no-cudnn-x64-release - runs-on: - - self-hosted - - "1ES.Pool=onnxruntime-github-linux-a10" - - "JobId=smoke-linux-cuda-no-cudnn-x64-release-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" + runs-on: [ort-github-linux-6-a10-scus-ss] timeout-minutes: 210 permissions: contents: read diff --git a/.github/workflows/linux_cuda_plugin_ci.yml b/.github/workflows/linux_cuda_plugin_ci.yml index eda64f6b98476..2c394e7868da3 100644 --- a/.github/workflows/linux_cuda_plugin_ci.yml +++ b/.github/workflows/linux_cuda_plugin_ci.yml @@ -22,7 +22,7 @@ jobs: name: Build Linux CUDA Plugin EP x64 Release uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" # Build pool + pool_name: "ort-github-linux-16-amd-wcus-ss" # Build pool build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cuda @@ -54,10 +54,7 @@ jobs: test-linux-cuda-plugin-x64-release: name: Test Linux CUDA Plugin EP x64 Release needs: build-linux-cuda-plugin-x64-release - runs-on: - - self-hosted - - "1ES.Pool=onnxruntime-github-linux-a10" - - "JobId=test-linux-cuda-plugin-x64-release-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" + runs-on: [ort-github-linux-6-a10-scus-ss] timeout-minutes: 210 permissions: contents: read diff --git a/.github/workflows/linux_minimal_build.yml b/.github/workflows/linux_minimal_build.yml index 4b68bb4c19ebc..6f2657f42f096 100644 --- a/.github/workflows/linux_minimal_build.yml +++ b/.github/workflows/linux_minimal_build.yml @@ -22,11 +22,7 @@ jobs: # Job 1: Build full onnxruntime and generate ORT format test files build_full_ort: name: 1. Build Full ORT and Generate ORT Files - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_full_ort-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: contents: read @@ -86,11 +82,7 @@ jobs: # Job 2: Build minimal onnxruntime [exceptions DISABLED, type reduction DISABLED, training ops ENABLED] build_minimal_exceptions_disabled: name: 2. Build Minimal (Exceptions Disabled) - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_minimal_exceptions_disabled-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: # Permissions needed for build-docker-image contents: read @@ -163,11 +155,7 @@ jobs: build_minimal_custom_ops: name: 3a. Build Minimal (Custom Ops) needs: build_full_ort # Depends on Job 1 for test data - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_minimal_custom_ops-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: # Permissions needed for build-docker-image contents: read @@ -219,11 +207,7 @@ jobs: build_minimal_type_reduction: name: 3b. Build Minimal (Type Reduction) needs: build_full_ort # Depends on Job 1 for test data - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_minimal_type_reduction-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: # Permissions needed for build-docker-image contents: read @@ -274,11 +258,7 @@ jobs: # Job 4: Build minimal onnxruntime [exceptions ENABLED, type reduction ENABLED (globally allowed types)] and run tests build_minimal_globally_allowed_types: name: 4. Build Minimal (Globally Allowed Types) - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_minimal_globally_allowed_types-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: # Permissions needed for build-docker-image contents: read @@ -330,11 +310,7 @@ jobs: # Job 5: Build extended minimal onnxruntime and run tests build_extended_minimal: name: 5. Build Extended Minimal - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_extended_minimal-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: # Permissions needed for build-docker-image contents: read @@ -417,11 +393,7 @@ jobs: # Job 6a: Regular build with python and all optional features disabled. build_regular_no_optional: name: 6a. Build Regular (No Optional Features) - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_regular_no_optional-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: # Permissions needed for build-docker-image contents: read @@ -515,11 +487,7 @@ jobs: # Job 6b: Minimal build with all optional features disabled. build_minimal_no_optional: name: 6b. Build Minimal (No Optional Features) - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_minimal_no_optional-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: # Permissions needed for build-docker-image contents: read @@ -602,11 +570,7 @@ jobs: # Job 6c: Extended minimal build with all optional features disabled. build_extended_minimal_no_optional: name: 6c. Build Extended Minimal (No Optional Features) - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_extended_minimal_no_optional-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: # Permissions needed for build-docker-image contents: read @@ -697,11 +661,7 @@ jobs: build_extended_minimal_android: name: 7. Build Extended Minimal (Android NNAPI) needs: build_full_ort # Depends on Job 1 for test data - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=build_extended_minimal_android-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: # Permissions needed for build-docker-image contents: read diff --git a/.github/workflows/linux_openvino_ci.yml b/.github/workflows/linux_openvino_ci.yml index e04bd6fd7d7af..8823126682571 100644 --- a/.github/workflows/linux_openvino_ci.yml +++ b/.github/workflows/linux_openvino_ci.yml @@ -23,7 +23,7 @@ jobs: # Use the reusable workflow as the other Linux CI pipelines uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + pool_name: "ort-github-linux-16-amd-wcus-ss" build_config: Release # Architecture: OpenVino only supports Intel X64 architecture: x64 diff --git a/.github/workflows/linux_tensorrt_ci.yml b/.github/workflows/linux_tensorrt_ci.yml index 3e881f35796bf..d4f9862c9cb21 100644 --- a/.github/workflows/linux_tensorrt_ci.yml +++ b/.github/workflows/linux_tensorrt_ci.yml @@ -23,7 +23,7 @@ jobs: # This job runs on a CPU node using the reusable build workflow uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" # Build pool + pool_name: "ort-github-linux-16-amd-wcus-ss" # Build pool build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cuda @@ -43,7 +43,7 @@ jobs: # Build-only job for CUDA minimal build (no tests, no unit tests) uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" + pool_name: "ort-github-linux-16-amd-wcus-ss" build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_cuda @@ -61,10 +61,7 @@ jobs: test-linux-TensorRT-x64-release: name: Test Linux TensorRT x64 Release needs: build-linux-TensorRT-x64-release - runs-on: - - self-hosted - - "1ES.Pool=onnxruntime-github-linux-a10" - - "JobId=test-linux-TensorRT-x64-release-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" + runs-on: [ort-github-linux-6-a10-scus-ss] timeout-minutes: 210 permissions: contents: read diff --git a/.github/workflows/linux_webgpu.yml b/.github/workflows/linux_webgpu.yml index 44aad2400af75..378474304b37e 100644 --- a/.github/workflows/linux_webgpu.yml +++ b/.github/workflows/linux_webgpu.yml @@ -23,7 +23,7 @@ jobs: # This job runs on a CPU node using the reusable build workflow uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" # Build pool + pool_name: "ort-github-linux-16-amd-wcus-ss" # Build pool build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_webgpu @@ -45,7 +45,7 @@ jobs: name: Build Linux WebGPU Static Plugin x64 Release uses: ./.github/workflows/reusable_linux_build.yml with: - pool_name: "onnxruntime-github-Ubuntu2204-AMD-CPU" # Build pool + pool_name: "ort-github-linux-16-amd-wcus-ss" # Build pool build_config: Release architecture: x64 dockerfile_path: tools/ci_build/github/linux/docker/Dockerfile.manylinux2_28_webgpu diff --git a/.github/workflows/pr_checks.yml b/.github/workflows/pr_checks.yml index c7b755150233d..c2955323055b4 100644 --- a/.github/workflows/pr_checks.yml +++ b/.github/workflows/pr_checks.yml @@ -19,11 +19,7 @@ concurrency: jobs: auto-apply-fixes: name: Suggest fixes - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=pr-auto-apply-fixes-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 permissions: contents: read diff --git a/.github/workflows/publish-c-apidocs.yml b/.github/workflows/publish-c-apidocs.yml index b05ef2010fc49..984ba12b4c3f6 100644 --- a/.github/workflows/publish-c-apidocs.yml +++ b/.github/workflows/publish-c-apidocs.yml @@ -30,11 +30,7 @@ permissions: jobs: build: name: Generate C/C++ API docs - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=apidocs-c-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 diff --git a/.github/workflows/publish-csharp-apidocs.yml b/.github/workflows/publish-csharp-apidocs.yml index 0206a4e3eec0e..9405bbb55eece 100644 --- a/.github/workflows/publish-csharp-apidocs.yml +++ b/.github/workflows/publish-csharp-apidocs.yml @@ -25,11 +25,7 @@ permissions: jobs: build: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=publish-csharp-apidocs-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 210 env: DOCFXVERSION: 2.62.2 diff --git a/.github/workflows/publish-gh-pages.yml b/.github/workflows/publish-gh-pages.yml index fe2419d25e4a8..cff5956820ff8 100644 --- a/.github/workflows/publish-gh-pages.yml +++ b/.github/workflows/publish-gh-pages.yml @@ -8,11 +8,7 @@ on: jobs: placeholder: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=publish-gh-pages-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 steps: - name: Placeholder step to have workflow included in the GitHub web UI diff --git a/.github/workflows/publish-java-apidocs.yml b/.github/workflows/publish-java-apidocs.yml index c9c5eb8cf3dd8..48ba4b6061fc4 100644 --- a/.github/workflows/publish-java-apidocs.yml +++ b/.github/workflows/publish-java-apidocs.yml @@ -26,11 +26,7 @@ permissions: jobs: build: name: Generate Java docs - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=apidocs-java-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 diff --git a/.github/workflows/publish-js-apidocs.yml b/.github/workflows/publish-js-apidocs.yml index a62c635b1f337..46f9a238e0e35 100644 --- a/.github/workflows/publish-js-apidocs.yml +++ b/.github/workflows/publish-js-apidocs.yml @@ -26,11 +26,7 @@ permissions: jobs: build: name: Generate JS API docs - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=apidocs-js-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 diff --git a/.github/workflows/publish-python-apidocs.yml b/.github/workflows/publish-python-apidocs.yml index cfc68f70bae16..642d4b8ea0851 100644 --- a/.github/workflows/publish-python-apidocs.yml +++ b/.github/workflows/publish-python-apidocs.yml @@ -28,11 +28,7 @@ permissions: jobs: build: name: Generate Python API docs - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=apidocs-python-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 diff --git a/.github/workflows/react_native.yml b/.github/workflows/react_native.yml index c67ab5bbb72a9..16b1196c810aa 100644 --- a/.github/workflows/react_native.yml +++ b/.github/workflows/react_native.yml @@ -14,11 +14,7 @@ concurrency: jobs: build_android_packages: name: Build Android AAR Packages - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=react_build_android_packages-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 outputs: aar_path: ${{ runner.temp }}/.artifacts @@ -76,11 +72,7 @@ jobs: react_native_ci_android: name: React Native CI Android needs: build_android_packages - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=react_native_ci_android-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 90 steps: - name: Set ANDROID_AVD_HOME environment variable diff --git a/.github/workflows/reusable_linux_build.yml b/.github/workflows/reusable_linux_build.yml index eada4a6149d8a..c057ac129ae43 100644 --- a/.github/workflows/reusable_linux_build.yml +++ b/.github/workflows/reusable_linux_build.yml @@ -4,7 +4,7 @@ on: workflow_call: inputs: pool_name: - description: 'The specific 1ES pool name (e.g., onnxruntime-github-Ubuntu2204-AMD-CPU)' + description: 'The 1ES scale-set name (e.g., ort-github-linux-16-amd-wcus-ss)' required: true type: string build_config: @@ -74,10 +74,7 @@ on: jobs: build_test_pipeline: - runs-on: - - self-hosted - - "1ES.Pool=${{ inputs.pool_name }}" - - "JobId=${{ inputs.job_identifier }}-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" + runs-on: ["${{ inputs.pool_name }}"] timeout-minutes: 120 permissions: contents: read diff --git a/.github/workflows/title-only-labeler.yml b/.github/workflows/title-only-labeler.yml index d3a3d91f6bd40..bf178e330dbc7 100644 --- a/.github/workflows/title-only-labeler.yml +++ b/.github/workflows/title-only-labeler.yml @@ -8,11 +8,7 @@ permissions: jobs: triage: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Ubuntu2204-AMD-CPU", - "JobId=title-only-labeler-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-linux-16-amd-wcus-ss] timeout-minutes: 120 steps: - uses: github/issue-labeler@ccca1e16fcb66762200e616fc29b57b940c1cb67 # v3.5 diff --git a/.github/workflows/windows-web-ci-workflow.yml b/.github/workflows/windows-web-ci-workflow.yml index d6bf97c7a3eda..d25140d2b03c9 100644 --- a/.github/workflows/windows-web-ci-workflow.yml +++ b/.github/workflows/windows-web-ci-workflow.yml @@ -22,11 +22,7 @@ on: jobs: build_onnxruntime_web: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-WEBGPU-A10", - "JobId=build_onnxruntime_web-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-webgpu-6-a10-scus-ss] env: webgpu_commandline_extra_flags: "--chromium-flags=--ignore-gpu-blocklist --chromium-flags=--gpu-vendor-id=0x10de" diff --git a/.github/workflows/windows_arm64_python.yml b/.github/workflows/windows_arm64_python.yml index a0be7cfc66327..430e48c69238c 100644 --- a/.github/workflows/windows_arm64_python.yml +++ b/.github/workflows/windows_arm64_python.yml @@ -21,11 +21,7 @@ concurrency: jobs: build_windows_arm64_python_wheel: name: Build Windows ARM64 Python 3.12 wheel - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-medium-win-gh-arm64-scus", - "JobId=windows-arm64-python-3.12-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-arm64-scus-ss] timeout-minutes: 180 strategy: fail-fast: false diff --git a/.github/workflows/windows_build_x64_asan.yml b/.github/workflows/windows_build_x64_asan.yml index d312981afb34c..72422c982a8e8 100644 --- a/.github/workflows/windows_build_x64_asan.yml +++ b/.github/workflows/windows_build_x64_asan.yml @@ -14,11 +14,7 @@ concurrency: jobs: build_x64: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-mms", - "JobId=windows-build-x64-asan-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-32-amd-eus2-ss] timeout-minutes: 300 env: diff --git a/.github/workflows/windows_cuda.yml b/.github/workflows/windows_cuda.yml index 360e445df5211..e197d8b2fc200 100644 --- a/.github/workflows/windows_cuda.yml +++ b/.github/workflows/windows_cuda.yml @@ -19,11 +19,7 @@ concurrency: jobs: build: name: Windows GPU CUDA CI Pipeline - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=windows-cuda-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 210 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 @@ -199,11 +195,7 @@ jobs: name: Windows GPU CUDA CI Pipeline Test Job needs: build timeout-minutes: 300 - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=windows-cuda-test-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 with: diff --git a/.github/workflows/windows_cuda_no_cudnn.yml b/.github/workflows/windows_cuda_no_cudnn.yml index d02d40319af3b..55fafc3de4c29 100644 --- a/.github/workflows/windows_cuda_no_cudnn.yml +++ b/.github/workflows/windows_cuda_no_cudnn.yml @@ -23,11 +23,7 @@ concurrency: jobs: build: name: Windows CUDA Plugin EP Build without cuDNN - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=windows-cuda-plugin-no-cudnn-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 210 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 @@ -168,11 +164,7 @@ jobs: name: Windows CUDA Plugin EP Test without cuDNN needs: build timeout-minutes: 120 - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=windows-cuda-plugin-no-cudnn-test-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 with: diff --git a/.github/workflows/windows_cuda_plugin.yml b/.github/workflows/windows_cuda_plugin.yml index fc0c2ff92787c..804e8c38ed4da 100644 --- a/.github/workflows/windows_cuda_plugin.yml +++ b/.github/workflows/windows_cuda_plugin.yml @@ -18,11 +18,7 @@ concurrency: jobs: build: name: Windows CUDA Plugin EP Build - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=windows-cuda-plugin-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 210 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 @@ -131,11 +127,7 @@ jobs: name: Windows CUDA Plugin EP Test needs: build timeout-minutes: 120 - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=windows-cuda-plugin-test-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 with: diff --git a/.github/workflows/windows_dml.yml b/.github/workflows/windows_dml.yml index 13db252b65b05..b6be18dd0941a 100644 --- a/.github/workflows/windows_dml.yml +++ b/.github/workflows/windows_dml.yml @@ -25,11 +25,7 @@ jobs: ALLOW_RELEASED_ONNX_OPSET_ONLY: '0' DocUpdateNeeded: false ONNXRUNTIME_TEST_GPU_DEVICE_ID: '0' - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=windows-dml-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] timeout-minutes: 210 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 diff --git a/.github/workflows/windows_gpu_doc_gen.yml b/.github/workflows/windows_gpu_doc_gen.yml index 4666d2f2ec50b..44dbf54e4c943 100644 --- a/.github/workflows/windows_gpu_doc_gen.yml +++ b/.github/workflows/windows_gpu_doc_gen.yml @@ -79,11 +79,7 @@ jobs: ALLOW_RELEASED_ONNX_OPSET_ONLY: '0' AZCOPY_AUTO_LOGIN_TYPE: MSI AZCOPY_MSI_CLIENT_ID: d712a4c7-a0cf-4e87-af75-31510eba0a8e - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=windows-gpu-doc-gen-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] timeout-minutes: 300 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 diff --git a/.github/workflows/windows_openvino.yml b/.github/workflows/windows_openvino.yml index 0521867cafb6d..d0fada58be699 100644 --- a/.github/workflows/windows_openvino.yml +++ b/.github/workflows/windows_openvino.yml @@ -18,11 +18,7 @@ concurrency: jobs: BUILD_OPENVINO_EP: name: Windows OpenVINO CI Pipeline - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=BUILD_OPENVINO_EP-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 240 env: AZCOPY_AUTO_LOGIN_TYPE: MSI diff --git a/.github/workflows/windows_qnn_x64.yml b/.github/workflows/windows_qnn_x64.yml index cdf556749e8e5..0601cf890ddd2 100644 --- a/.github/workflows/windows_qnn_x64.yml +++ b/.github/workflows/windows_qnn_x64.yml @@ -18,11 +18,7 @@ concurrency: jobs: build_test_qnn_ep: name: Windows x64 QNN CI Pipeline (${{ matrix.QnnLibKind }}) - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=build_test_qnn_ep-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}-${{matrix.QnnLibKind}}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 120 strategy: matrix: diff --git a/.github/workflows/windows_tensorrt.yml b/.github/workflows/windows_tensorrt.yml index 18882550f1231..c6b111948fe60 100644 --- a/.github/workflows/windows_tensorrt.yml +++ b/.github/workflows/windows_tensorrt.yml @@ -19,11 +19,7 @@ concurrency: jobs: build: name: Windows GPU TensorRT CI Pipeline - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=windows-tensorrt-build-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 210 steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 @@ -167,11 +163,7 @@ jobs: name: Windows GPU TensorRT CI Pipeline Test Job needs: build timeout-minutes: 300 - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=windows-tensorrt-test-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] steps: - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0 with: diff --git a/.github/workflows/windows_webgpu.yml b/.github/workflows/windows_webgpu.yml index e2457ad51e2c0..58f95a8c2a3e3 100644 --- a/.github/workflows/windows_webgpu.yml +++ b/.github/workflows/windows_webgpu.yml @@ -25,11 +25,7 @@ env: jobs: webgpu_compile_x64_RelWithDebInfo: name: webgpu_compile_x64_RelWithDebInfo (${{ matrix.vcpkg_option }}) - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=webgpu_compile_x64_RelWithDebInfo-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}-${{ matrix.vcpkg_option }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 300 strategy: matrix: &webgpu_matrix @@ -99,6 +95,13 @@ jobs: with: dotnet-version: "8.x" + - name: Test WebGPU NuGet packaging + run: python -m unittest plugin-ep-webgpu\csharp\test\test_pack_nuget.py -v + shell: cmd + working-directory: ${{ github.workspace }} + env: + ORT_WEBGPU_TEST_REQUIRE_DOTNET_SDK: "1" + - name: Use Nuget 6.x uses: nuget/setup-nuget@d105a947828025cd7a980103c35ba2bfae586d0f # v2.0.2 with: @@ -167,11 +170,7 @@ jobs: webgpu_test_x64_RelWithDebInfo: name: webgpu_test_x64_RelWithDebInfo (${{ matrix.vcpkg_option }}) needs: webgpu_compile_x64_RelWithDebInfo - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=webgpu_test_x64_RelWithDebInfo-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}-${{ matrix.vcpkg_option }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] timeout-minutes: 300 strategy: matrix: *webgpu_matrix @@ -289,11 +288,7 @@ jobs: webgpu_plugin_compile_x64_RelWithDebInfo: name: webgpu_plugin_compile_x64_RelWithDebInfo - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=webgpu_plugin_compile_x64_RelWithDebInfo-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 300 env: setVcvars: true @@ -381,11 +376,7 @@ jobs: # Preserve the legacy required-check name so branch protection waits for GPU tests. name: webgpu_plugin_build_x64_RelWithDebInfo needs: webgpu_plugin_compile_x64_RelWithDebInfo - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=webgpu_plugin_test_x64_RelWithDebInfo-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] timeout-minutes: 300 env: setVcvars: true @@ -446,11 +437,7 @@ jobs: # 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: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=webgpu_static_plugin_build_x64_RelWithDebInfo-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] timeout-minutes: 300 env: OnnxRuntimeBuildDirectory: ${{ github.workspace }} @@ -515,11 +502,7 @@ jobs: } webgpu_external_dawn_build_x64_RelWithDebInfo: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=webgpu_external_dawn_build_x64_RelWithDebInfo-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] timeout-minutes: 300 steps: - name: Checkout @@ -600,11 +583,7 @@ jobs: webgpu_dawn_package_build_x64_RelWithDebInfo: needs: webgpu_external_dawn_build_x64_RelWithDebInfo - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-Win2022-GPU-A10", - "JobId=webgpu_dawn_package_build_x64_RelWithDebInfo-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-6-a10-scus-ss] timeout-minutes: 300 steps: - name: Checkout @@ -720,11 +699,7 @@ jobs: if ($LASTEXITCODE -ne 0) { throw "Plugin-backed cache exception containment failed" } webgpu_minimal_build_edge_build_x64_RelWithDebInfo: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=webgpu_minimal_build_edge_build_x64_RelWithDebInfo-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 300 env: OrtPackageId: Microsoft.ML.OnnxRuntime diff --git a/.github/workflows/windows_x64_debug_build_x64_debug.yml b/.github/workflows/windows_x64_debug_build_x64_debug.yml index 52a7a213496b8..e26b7961dd12f 100644 --- a/.github/workflows/windows_x64_debug_build_x64_debug.yml +++ b/.github/workflows/windows_x64_debug_build_x64_debug.yml @@ -13,11 +13,7 @@ concurrency: jobs: build_x64_debug: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=windows-x64-debug-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 300 steps: diff --git a/.github/workflows/windows_x64_release_build_x64_release.yml b/.github/workflows/windows_x64_release_build_x64_release.yml index 2b0d54031f775..35d1ab2275aae 100644 --- a/.github/workflows/windows_x64_release_build_x64_release.yml +++ b/.github/workflows/windows_x64_release_build_x64_release.yml @@ -13,11 +13,7 @@ concurrency: jobs: build_x64_release: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=windows-x64-release-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 300 steps: diff --git a/.github/workflows/windows_x64_release_ep_generic_interface_build_x64_release_ep_generic_interface.yml b/.github/workflows/windows_x64_release_ep_generic_interface_build_x64_release_ep_generic_interface.yml index 0ca8e5e408228..834f5f1303d7a 100644 --- a/.github/workflows/windows_x64_release_ep_generic_interface_build_x64_release_ep_generic_interface.yml +++ b/.github/workflows/windows_x64_release_ep_generic_interface_build_x64_release_ep_generic_interface.yml @@ -13,11 +13,7 @@ concurrency: jobs: build_x64_release_ep_generic_interface: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=build_x64_release_ep_generic_interface-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 300 steps: diff --git a/.github/workflows/windows_x64_release_vitisai_build_x64_release.yml b/.github/workflows/windows_x64_release_vitisai_build_x64_release.yml index 8c31abcd2d4aa..530ee5af49cb1 100644 --- a/.github/workflows/windows_x64_release_vitisai_build_x64_release.yml +++ b/.github/workflows/windows_x64_release_vitisai_build_x64_release.yml @@ -13,11 +13,7 @@ concurrency: jobs: build_x64_release_vitisai: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=build_x64_release_vitisai-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 300 steps: diff --git a/.github/workflows/windows_x64_release_xnnpack.yml b/.github/workflows/windows_x64_release_xnnpack.yml index f315de01b959f..b6ee0f7e29163 100644 --- a/.github/workflows/windows_x64_release_xnnpack.yml +++ b/.github/workflows/windows_x64_release_xnnpack.yml @@ -13,11 +13,7 @@ concurrency: jobs: build_x64_release_xnnpack: - runs-on: [ - "self-hosted", - "1ES.Pool=onnxruntime-github-vs2022-latest", - "JobId=build_x64_release_xnnpack-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" - ] + runs-on: [ort-github-win-16-amd-scus-ss] timeout-minutes: 300 steps: diff --git a/cgmanifests/webgpu/README.md b/cgmanifests/webgpu/README.md index cf03477ea6bbe..a10677d5c75b2 100644 --- a/cgmanifests/webgpu/README.md +++ b/cgmanifests/webgpu/README.md @@ -25,16 +25,48 @@ registration and explain the packaging path in `comments` and `detectedComponent When rolling Dawn or changing WebGPU packaging: -1. Update the Dawn registration to match the `dawn` entry in `cmake/deps.txt`. -2. Re-audit the Dawn dependency graph for the pinned Dawn commit: +For a complete Dawn update, run the [Dawn update tool](../../plugin-ep-webgpu/tools/dawn/README.md): + +```powershell +python plugin-ep-webgpu/tools/dawn/update_dawn.py +``` + +It selects the latest dated release by default, updates the archive pin/checksum, this manifest, SwiftShader, +and active patches, then validates the result automatically. Use `--revision TAG_OR_COMMIT` to choose a version. +The commands below are for manifest-only maintenance; they are not additional steps after a complete update. + +1. After updating the `dawn` entry in `cmake/deps.txt`, refresh the registered Git commit hashes: + + ```powershell + python cgmanifests/webgpu/update_webgpu_cgmanifest.py + ``` + + The updater resolves Dawn's commit or tag, downloads `DEPS` at that exact commit, and updates the Dawn root, + registered Dawn dependency commits, and `dependencyRoots`. It recognizes Chromium mirrors of public repositories + and updates renamed DEPS paths in comments. It preserves the reviewed classifications, other comments, public + repository URLs, and detected locations, including the separate redistributed DXC release registration. + + This command needs network access and Git on `PATH` when Dawn is pinned to a tag. It parses DEPS without executing + it and does not clone dependency repositories or evaluate platform conditions. A missing, changed, or ambiguous + registered repository stops the update before writing the manifest; review that entry using the inventory below. + +2. Re-audit the Dawn dependency graph for the pinned Dawn commit. To generate a JSON inventory of every Git entry in + Dawn's DEPS (path, repository URL, commit hash, and condition), run: + + ```powershell + python cgmanifests/webgpu/update_webgpu_cgmanifest.py --list-deps + ``` + - Start from the Dawn commit in `cmake/deps.txt`; do not audit Dawn `main` or a different roll. - Inspect Dawn's `tools/fetch_dawn_dependencies.py` at that commit. For ORT's normal source-fetch path, `cmake/external/onnxruntime_external_deps.cmake` enables `DAWN_FETCH_DEPENDENCIES`, so the script's `required_submodules` list is the primary set of Dawn source dependencies fetched for the build. - - Cross-reference each fetched submodule path with Dawn's `DEPS` file to get the public upstream repository URL, - commit, and condition. Use public upstream identities in this manifest, not internal mirrors. - - Compare that fetched set against this manifest. Add new fetched components, update changed commits or repository - URLs, and remove entries that are no longer fetched or relevant unless CG/legal guidance requires keeping them. + - Cross-reference each fetched submodule path with the generated inventory. Its URLs are Dawn's checkout URLs; + use public upstream identities in this manifest, not internal mirrors. + - Compare that fetched set against this manifest. The updater refreshes existing registrations; it does not decide + which components belong in the manifest. Add new fetched components, review changed repository URLs, and remove + entries that are no longer fetched or relevant unless CG/legal guidance requires keeping them. Rerun the updater + after editing registrations. - Cross-check ORT's Dawn CMake options in `cmake/external/onnxruntime_external_deps.cmake` and Dawn's `third_party/CMakeLists.txt` before classifying a component. Components that are redistributed, statically linked, or otherwise part of the WebGPU package/runtime closure should not be marked as development dependencies; build @@ -50,10 +82,14 @@ When rolling Dawn or changing WebGPU packaging: ```powershell python cgmanifests/webgpu/validate_webgpu_cgmanifest.py + python cgmanifests/webgpu/update_webgpu_cgmanifest.py --check ``` -The validator checks for stale Dawn and DXC pins, but it does not replace the manual dependency classification review -in step 2. +The validator checks the local Dawn pin, DXC release metadata, and Dawn dependency roots without network access. +The updater's `--check` mode also resolves the pinned Dawn revision and compares every registered Dawn dependency +against its DEPS, returning a nonzero exit code for stale hashes or paths without writing. Neither replaces the +dependency selection and classification review in step 2. Use `--deps-file PATH` and `--manifest PATH` to update or +check alternate input files. Non-git Dawn toolchain packages from CIPD/GCS, such as GN, Ninja, CMake, Go, Siso, reclient, and sysroots, are intentionally not registered here unless they become redistributed or CG/legal guidance requires build input coverage. diff --git a/cgmanifests/webgpu/test_update_webgpu_cgmanifest.py b/cgmanifests/webgpu/test_update_webgpu_cgmanifest.py new file mode 100644 index 0000000000000..3dc5fc71b469b --- /dev/null +++ b/cgmanifests/webgpu/test_update_webgpu_cgmanifest.py @@ -0,0 +1,245 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import copy +import io +import json +import subprocess +import tempfile +import unittest +from contextlib import ExitStack, redirect_stderr, redirect_stdout +from pathlib import Path +from unittest.mock import patch + +import update_webgpu_cgmanifest as updater +import validate_webgpu_cgmanifest as validator + +OLD = "1" * 40 +DAWN = "2" * 40 +CHILD = "3" * 40 +TAG_OBJECT = "4" * 40 +REPO = "https://github.com/microsoft/DirectXShaderCompiler.git" +PATH = "third_party/directx-shader-compiler/src" +SOURCE = f""" +vars = {{'chromium_git': 'https://chromium.googlesource.com', 'revision': '{CHILD}', 'version': Str('3.0')}} +deps = {{ + '{PATH}': {{ + 'url': '{{chromium_git}}/external/github.com/microsoft/DirectXShaderCompiler.git@{{revision}}', + 'condition': 'host_os == "win"', + }}, + 'tools/win': Var('chromium_git') + '/chromium/src/tools/win@{CHILD}', + 'tools/cipd': {{'dep_type': 'cipd', 'packages': [{{'version': Var('version')}}]}}, + 'tools/gcs': {{'dep_type': 'gcs', 'objects': []}}, + 'removed': None, +}} +""" + + +def make_manifest(): + return { + "version": 1, + "registrations": [ + { + "component": { + "type": "git", + "git": {"repositoryUrl": validator.DAWN_REPOSITORY_URL, "commitHash": OLD}, + }, + "comments": "runtime; Dawn root", + }, + { + "component": {"type": "git", "git": {"repositoryUrl": REPO, "commitHash": OLD, "tag": "v1.8.2502"}}, + "comments": "runtime; separately packaged DXC release; SHA256: keep-me", + }, + { + "component": {"type": "git", "git": {"repositoryUrl": REPO, "commitHash": OLD}}, + "comments": "build-tool; Dawn DEPS third_party/dxc. Reviewed classification.", + "developmentDependency": True, + "detectedComponentLocations": ["{SourceFileRoot}/cmake/deps.txt"], + "dependencyRoots": [ + {"type": "git", "git": {"repositoryUrl": validator.DAWN_REPOSITORY_URL, "commitHash": OLD}} + ], + }, + ], + } + + +class DawnDepsTests(unittest.TestCase): + def test_parse_all_git_pins_including_conditional_entries(self): + deps = updater.parse_dawn_deps(SOURCE) + self.assertEqual(set(deps), {PATH, "tools/win"}) + self.assertEqual(deps[PATH]["commitHash"], CHILD) + self.assertEqual(deps[PATH]["condition"], 'host_os == "win"') + self.assertEqual(deps["tools/win"]["repositoryUrl"], "https://chromium.googlesource.com/chromium/src/tools/win") + + def test_does_not_execute_deps_or_hooks(self): + with tempfile.TemporaryDirectory() as directory: + marker = Path(directory) / "executed" + code = f"__import__('pathlib').Path({str(marker)!r}).touch()" + self.assertEqual(updater.parse_dawn_deps(SOURCE + "\n" + code), updater.parse_dawn_deps(SOURCE)) + with self.assertRaises(ValueError): + updater.parse_dawn_deps(f"vars = {{}}\ndeps = {{'malicious': {code}}}") + self.assertFalse(marker.exists()) + + def test_rejects_unsupported_or_ambiguous_metadata(self): + sources = [ + SOURCE.replace(CHILD, "main"), + SOURCE.replace("{revision}", "{missing}"), + SOURCE + "\ndeps = {}", + "vars = {}\ndeps = {}", + "vars = {}\ndeps = dict()", + "vars = {}\ndeps = {'a': None, 'a': None}", + "vars = {}\ndeps = {", + ] + for source in sources: + with self.subTest(source=source), self.assertRaises((ValueError, KeyError, SyntaxError)): + updater.parse_dawn_deps(source) + + def test_upstream_mirrors_and_split_llvm_repository_identities(self): + aliases = [ + ( + "https://chromium.googlesource.com/external/github.com/Google/googletest", + "https://github.com/google/googletest.git", + ), + ( + "https://chromium.googlesource.com/external/gitlab.freedesktop.org/mesa/mesa/", + "https://gitlab.freedesktop.org/mesa/mesa.git", + ), + ] + for mirror, upstream in aliases: + with self.subTest(mirror=mirror): + self.assertEqual(updater._repository_identity(mirror), updater._repository_identity(upstream)) + llvm = "https://chromium.googlesource.com/external/github.com/llvm/llvm-project/libcxx" + self.assertEqual(updater._repository_identity(llvm + ".git"), llvm) + + def test_resolves_lightweight_and_annotated_tags(self): + for refs in [f"{DAWN}\trefs/tags/v1\n", f"{TAG_OBJECT}\trefs/tags/v1\n{DAWN}\trefs/tags/v1^{{}}\n"]: + with self.subTest(refs=refs), patch.object(updater.subprocess, "run") as run: + run.return_value.stdout = refs + self.assertEqual(updater._resolve_dawn_commit("tag", "v1"), DAWN) + self.assertEqual(run.call_args.args[0][-2:], ["refs/tags/v1", "refs/tags/v1^{}"]) + with patch.object(updater.subprocess, "run") as run: + self.assertEqual(updater._resolve_dawn_commit("commitHash", DAWN), DAWN) + run.assert_not_called() + with patch.object(updater.subprocess, "run") as run: + run.return_value.stdout = "" + with self.assertRaisesRegex(ValueError, "could not resolve"): + updater._resolve_dawn_commit("tag", "missing") + + def test_download_uses_resolved_commit(self): + with patch.object(updater.urllib.request, "urlopen") as urlopen: + urlopen.return_value.__enter__.return_value.read.return_value = SOURCE.encode() + self.assertEqual(updater._download_deps(DAWN), SOURCE) + urlopen.assert_called_once_with(f"https://raw.githubusercontent.com/google/dawn/{DAWN}/DEPS", timeout=30) + + +class ManifestTests(unittest.TestCase): + def test_updates_hashes_and_renamed_path_preserving_reviewed_metadata(self): + original = make_manifest() + snapshot = copy.deepcopy(original) + updated = updater.update_manifest(original, DAWN, "v1", updater.parse_dawn_deps(SOURCE)) + expected = copy.deepcopy(original) + root, _, child = expected["registrations"] + root["component"]["git"].update(commitHash=DAWN, tag="v1") + child["component"]["git"]["commitHash"] = CHILD + child["comments"] = child["comments"].replace("third_party/dxc", PATH) + child["dependencyRoots"][0]["git"]["commitHash"] = DAWN + self.assertEqual(updated, expected) + self.assertEqual(original, snapshot) + self.assertEqual(updater.update_manifest(updated, DAWN, "v1", updater.parse_dawn_deps(SOURCE)), updated) + + def test_commit_pin_removes_old_tag(self): + manifest = make_manifest() + manifest["registrations"][0]["component"]["git"]["tag"] = "old-tag" + updated = updater.update_manifest(manifest, DAWN, None, updater.parse_dawn_deps(SOURCE)) + self.assertNotIn("tag", updated["registrations"][0]["component"]["git"]) + + def test_rejects_removed_changed_or_ambiguous_repositories(self): + deps = updater.parse_dawn_deps(SOURCE) + variants = [ + {"tools/win": deps["tools/win"]}, + {"third_party/dxc": deps["tools/win"], PATH: deps[PATH]}, + {PATH: deps[PATH], "another/path": deps[PATH]}, + ] + for dependencies in variants: + with self.subTest(dependencies=dependencies), self.assertRaisesRegex(ValueError, "No manifest written"): + updater.update_manifest(make_manifest(), DAWN, None, dependencies) + + def test_duplicate_repository_is_disambiguated_by_path(self): + deps = updater.parse_dawn_deps(SOURCE) + deps["third_party/dxc"] = dict(deps[PATH], commitHash=TAG_OBJECT) + updated = updater.update_manifest(make_manifest(), DAWN, None, deps) + self.assertEqual(updated["registrations"][2]["component"]["git"]["commitHash"], TAG_OBJECT) + + +class CommandTests(unittest.TestCase): + def setUp(self): + contexts = ExitStack() + self.addCleanup(contexts.close) + directory = contexts.enter_context(tempfile.TemporaryDirectory()) + self.manifest = Path(directory) / "manifest.json" + self.manifest.write_text(json.dumps(make_manifest(), indent=2) + "\n", encoding="utf-8") + self.deps_file = Path(directory) / "deps.txt" + self.deps_file.write_text( + f"dawn;https://github.com/google/dawn/archive/{DAWN}.zip;archive-hash\n", encoding="utf-8" + ) + self.download = contexts.enter_context(patch.object(updater, "_download_deps", return_value=SOURCE)) + self.stdout = contexts.enter_context(redirect_stdout(io.StringIO())) + self.stderr = contexts.enter_context(redirect_stderr(io.StringIO())) + self.args = ["--manifest", str(self.manifest), "--deps-file", str(self.deps_file)] + + def test_check_update_check_and_idempotence(self): + before = self.manifest.read_bytes() + self.assertEqual(updater.main([*self.args, "--check"]), 1) + self.assertEqual(self.manifest.read_bytes(), before) + self.assertEqual(updater.main(self.args), 0) + after = self.manifest.read_bytes() + self.assertNotEqual(after, before) + self.assertEqual(updater.main([*self.args, "--check"]), 0) + self.assertEqual(updater.main(self.args), 0) + self.assertEqual(self.manifest.read_bytes(), after) + + def test_failures_do_not_write_partial_manifest(self): + before = self.manifest.read_bytes() + for error in [OSError("network failed"), subprocess.TimeoutExpired("git", 30)]: + with self.subTest(error=error): + self.download.side_effect = error + self.assertEqual(updater.main(self.args), 1) + self.assertEqual(self.manifest.read_bytes(), before) + self.download.side_effect = None + for source in ["invalid DEPS!", SOURCE.replace("DirectXShaderCompiler", "different-repo")]: + with self.subTest(source=source): + self.download.return_value = source + self.assertEqual(updater.main(self.args), 1) + self.assertEqual(self.manifest.read_bytes(), before) + + def test_list_all_deps_without_a_manifest(self): + self.manifest.unlink() + self.assertEqual(updater.main([*self.args, "--list-deps"]), 0) + self.assertEqual(json.loads(self.stdout.getvalue()), updater.parse_dawn_deps(SOURCE)) + self.assertFalse(self.manifest.exists()) + + def test_tag_pin_in_deps_txt(self): + self.deps_file.write_text( + "dawn;https://github.com/google/dawn/archive/refs/tags/v1.zip;hash\n", encoding="utf-8" + ) + with patch.object(updater.subprocess, "run") as run: + run.return_value.stdout = f"{DAWN}\trefs/tags/v1\n" + self.assertEqual(updater.main(self.args), 0) + self.download.assert_called_once_with(DAWN) + manifest = json.loads(self.manifest.read_text(encoding="utf-8")) + self.assertEqual(manifest["registrations"][0]["component"]["git"]["tag"], "v1") + with patch.object(validator, "_dawn_pin_from_deps_txt", return_value=("tag", "v1")): + validator._validate_dawn_root(manifest) + validator._validate_dawn_dependency_roots(manifest) + + def test_rejects_missing_or_duplicate_dawn_pins(self): + pin = self.deps_file.read_text(encoding="utf-8") + for text in ["", pin + pin, pin.replace(DAWN, "main")]: + with self.subTest(text=text): + self.deps_file.write_text(text, encoding="utf-8") + self.assertEqual(updater.main(self.args), 1) + self.download.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/cgmanifests/webgpu/update_webgpu_cgmanifest.py b/cgmanifests/webgpu/update_webgpu_cgmanifest.py new file mode 100644 index 0000000000000..e8565b2b012d3 --- /dev/null +++ b/cgmanifests/webgpu/update_webgpu_cgmanifest.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 + +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +"""Update WebGPU manifest Git commits from the Dawn version pinned in cmake/deps.txt.""" + +from __future__ import annotations + +import argparse +import ast +import copy +import json +import re +import subprocess +import sys +import urllib.request +from pathlib import Path +from typing import Any + +import validate_webgpu_cgmanifest as validator + + +def _assignment(tree: ast.Module, name: str) -> ast.expr: + values = [ + node.value + for node in tree.body + if isinstance(node, ast.Assign) and any(isinstance(t, ast.Name) and t.id == name for t in node.targets) + ] + if len(values) != 1: + raise ValueError(f"expected one top-level {name} assignment in Dawn DEPS") + return values[0] + + +def _dictionary(node: ast.expr) -> dict[str, ast.expr]: + if not isinstance(node, ast.Dict): + raise ValueError("expected a literal dictionary in Dawn DEPS") + result = {} + for key, value in zip(node.keys, node.values, strict=True): + if not isinstance(key, ast.Constant) or not isinstance(key.value, str) or key.value in result: + raise ValueError("expected unique literal string keys in Dawn DEPS") + result[key.value] = value + return result + + +def _string(node: ast.expr, variables: dict[str, Any]) -> str: + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add): + return _string(node.left, variables) + _string(node.right, variables) + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "Var" + and len(node.args) == 1 + and not node.keywords + ): + value = variables[_string(node.args[0], {})] + if isinstance(value, str): + return value + raise ValueError(f"unsupported Dawn DEPS string expression: {ast.dump(node)}") + + +def parse_dawn_deps(source: str) -> dict[str, dict[str, str]]: + """Read Git pins without executing DEPS, hooks, or dependency conditions.""" + tree = ast.parse(source) + variables = {} + for name, value in _dictionary(_assignment(tree, "vars")).items(): + if ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Name) + and value.func.id == "Str" + and len(value.args) == 1 + and not value.keywords + ): + variables[name] = str(ast.literal_eval(value.args[0])) + else: + variables[name] = ast.literal_eval(value) + + dependencies = {} + for path, entry in _dictionary(_assignment(tree, "deps")).items(): + node = entry + condition = "" + if isinstance(node, ast.Dict): + fields = _dictionary(node) + if "url" not in fields: + continue # CIPD/GCS packages do not identify Git source commits. + node = fields["url"] + if "condition" in fields: + condition = _string(fields["condition"], variables) + if isinstance(node, ast.Constant) and node.value is None: + continue + url = _string(node, variables).format(**variables) + repository, separator, commit = url.rpartition("@") + if not separator or not repository.startswith("https://") or not re.fullmatch(r"[0-9a-fA-F]{40}", commit): + raise ValueError(f"expected an HTTPS Git URL and full commit hash for {path}: {url}") + dependencies[path] = {"repositoryUrl": repository, "commitHash": commit.lower(), "condition": condition} + if not dependencies: + raise ValueError("no Git dependencies found in Dawn DEPS") + return dependencies + + +def _resolve_dawn_commit(pin_field: str, pin: str) -> str: + if pin_field == "commitHash": + return pin + ref = f"refs/tags/{pin}" + result = subprocess.run( + ["git", "ls-remote", validator.DAWN_REPOSITORY_URL, ref, ref + "^{}"], + check=True, + capture_output=True, + text=True, + timeout=30, + ) + refs = dict(line.split()[::-1] for line in result.stdout.splitlines()) + # Annotated tags point at tag objects; use the peeled commit when present. + commit = refs.get(ref + "^{}", refs.get(ref, "")) + if not re.fullmatch(r"[0-9a-f]{40}", commit): + raise ValueError(f"could not resolve Dawn tag {pin} to a commit") + return commit + + +def _download_deps(commit: str) -> str: + url = f"https://raw.githubusercontent.com/google/dawn/{commit}/DEPS" + with urllib.request.urlopen(url, timeout=30) as response: + return response.read().decode("utf-8") + + +def _repository_identity(url: str) -> str: + url = url.rstrip("/").removesuffix(".git") + mirror = "https://chromium.googlesource.com/external/github.com/" + # LLVM split repositories have no equivalent GitHub owner/repo identity. + if url.startswith(mirror) and len(url.removeprefix(mirror).split("/")) == 2: + url = "https://github.com/" + url.removeprefix(mirror) + url = url.replace( + "https://chromium.googlesource.com/external/gitlab.freedesktop.org/", "https://gitlab.freedesktop.org/" + ) + if url.startswith("https://github.com/"): + url = url.lower() + return url + + +def update_manifest( + manifest: dict[str, Any], commit: str, tag: str | None, dependencies: dict[str, dict[str, str]] +) -> dict[str, Any]: + """Refresh registered Dawn pins, preserving classifications and unrelated entries.""" + updated = copy.deepcopy(manifest) + dawn = validator._find_git_registration(updated, validator.DAWN_REPOSITORY_URL) + dawn_git = dawn["component"]["git"] + dawn_git["commitHash"] = commit + if tag is None: + dawn_git.pop("tag", None) + else: + dawn_git["tag"] = tag + + errors = [] + for registration in updated["registrations"]: + comments = registration.get("comments", "") + if "Dawn DEPS" not in comments: + continue + git = validator._git_component(registration) + path_match = re.search(r"\bDawn DEPS ([^\s,;]+)", comments) + if git is None or path_match is None or "tag" in git: + errors.append(f"expected an untagged Git registration with a Dawn DEPS path: {comments}") + continue + old_path = path_match.group(1).rstrip(".") + identity = _repository_identity(git["repositoryUrl"]) + candidates = [p for p, dep in dependencies.items() if _repository_identity(dep["repositoryUrl"]) == identity] + if old_path in dependencies: + candidates = [old_path] if old_path in candidates else [] + if len(candidates) != 1: + errors.append( + f"{old_path}: cannot uniquely match {git['repositoryUrl']}; review removed/changed repositories" + ) + continue + path = candidates[0] + registration["component"]["git"]["commitHash"] = dependencies[path]["commitHash"] + registration["comments"] = comments.replace(f"Dawn DEPS {old_path}", f"Dawn DEPS {path}", 1) + roots = registration.setdefault( + "dependencyRoots", [{"type": "git", "git": {"repositoryUrl": validator.DAWN_REPOSITORY_URL}}] + ) + if len(roots) != 1 or roots[0].get("git", {}).get("repositoryUrl") != validator.DAWN_REPOSITORY_URL: + errors.append(f"{old_path}: expected one Dawn dependency root") + continue + roots[0]["git"]["commitHash"] = commit + if "tag" in roots[0]["git"]: + if tag is None: + roots[0]["git"].pop("tag") + else: + roots[0]["git"]["tag"] = tag + + if errors: + raise ValueError( + "\n".join(errors) + "\nNo manifest written. Use --list-deps to inspect the pinned dependencies." + ) + return updated + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, default=validator.WEBGPU_CGMANIFEST) + parser.add_argument("--deps-file", type=Path, default=validator.DEPS_TXT) + modes = parser.add_mutually_exclusive_group() + modes.add_argument("--check", action="store_true", help="check registered Dawn hashes without writing") + modes.add_argument("--list-deps", action="store_true", help="print all pinned Dawn Git dependencies as JSON") + args = parser.parse_args(argv) + try: + pin_field, pin = validator._dawn_pin_from_deps_txt(args.deps_file) + commit = _resolve_dawn_commit(pin_field, pin) + dependencies = parse_dawn_deps(_download_deps(commit)) + print(f"Dawn {pin} -> {commit}; {len(dependencies)} Git dependencies", file=sys.stderr) + if args.list_deps: + print(json.dumps(dependencies, indent=2) + "\n", end="") + return 0 + manifest = json.loads(args.manifest.read_text(encoding="utf-8")) + updated = update_manifest(manifest, commit, pin if pin_field == "tag" else None, dependencies) + if updated == manifest: + print(f"Registered Dawn hashes are up to date in {args.manifest}") + return 0 + if args.check: + print( + f"ERROR: stale Dawn hashes or paths in {args.manifest}; run without --check to update", file=sys.stderr + ) + return 1 + # Resolve and validate every entry before touching the maintained manifest. + args.manifest.write_text(json.dumps(updated, indent=2) + "\n", encoding="utf-8") + print(f"Updated {args.manifest}. Review dependency additions/removals and classifications with --list-deps.") + except (OSError, ValueError, KeyError, SyntaxError, subprocess.SubprocessError) as ex: + print(f"ERROR: {ex}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/cgmanifests/webgpu/validate_webgpu_cgmanifest.py b/cgmanifests/webgpu/validate_webgpu_cgmanifest.py index 393f441060115..ea8a7732f84f6 100644 --- a/cgmanifests/webgpu/validate_webgpu_cgmanifest.py +++ b/cgmanifests/webgpu/validate_webgpu_cgmanifest.py @@ -76,29 +76,27 @@ def _find_git_registration(manifest: dict[str, Any], repository_url: str, *, tag return matches[0] -def _dawn_pin_from_deps_txt() -> tuple[str, str]: - deps_text = DEPS_TXT.read_text(encoding="utf-8") - commit_match = re.search( - r"^dawn;https://github\.com/google/dawn/archive/([0-9a-f]{40})\.zip;", deps_text, re.MULTILINE +def _dawn_pin_from_deps_txt(deps_file: Path = DEPS_TXT) -> tuple[str, str]: + deps_text = deps_file.read_text(encoding="utf-8") + matches = re.findall( + r"^dawn;https://github\.com/google/dawn/archive/([0-9a-f]{40}|refs/tags/[^;]+)\.zip;", + deps_text, + re.MULTILINE, ) - if commit_match: - return "commitHash", commit_match.group(1) - - tag_match = re.search( - r"^dawn;https://github\.com/google/dawn/archive/refs/tags/([^;]+)\.zip;", deps_text, re.MULTILINE - ) - if tag_match: - return "tag", tag_match.group(1) - - raise ValueError(f"could not find Dawn commit or tag in {DEPS_TXT}") + if len(matches) != 1: + raise ValueError(f"expected one Dawn commit or tag in {deps_file}") + pin = matches[0] + if pin.startswith("refs/tags/"): + return "tag", pin.removeprefix("refs/tags/") + return "commitHash", pin -def _dxc_release_from_pipeline() -> tuple[str, str, str]: - pipeline_text = PLUGIN_WIN_WEBGPU_STAGE.read_text(encoding="utf-8") +def _dxc_release_from_pipeline(pipeline_file: Path = PLUGIN_WIN_WEBGPU_STAGE) -> tuple[str, str, str]: + pipeline_text = pipeline_file.read_text(encoding="utf-8") url_match = re.search(r'\$dxcZipUrl = "([^"]+)"', pipeline_text) hash_match = re.search(r'\$expectedHash = "([0-9A-Fa-f]+)"', pipeline_text) if not url_match or not hash_match: - raise ValueError(f"could not find DXC release URL/hash in {PLUGIN_WIN_WEBGPU_STAGE}") + raise ValueError(f"could not find DXC release URL/hash in {pipeline_file}") tag_match = re.search(r"/download/(v[^/]+)/", url_match.group(1)) if not tag_match: @@ -107,21 +105,21 @@ def _dxc_release_from_pipeline() -> tuple[str, str, str]: return tag_match.group(1), url_match.group(1), hash_match.group(1).upper() -def _validate_dawn_root(manifest: dict[str, Any]) -> None: +def _validate_dawn_root(manifest: dict[str, Any], deps_file: Path = DEPS_TXT) -> None: registration = _find_git_registration(manifest, DAWN_REPOSITORY_URL) git = _git_component(registration) if git is None: raise ValueError("Dawn registration must be a git component") - pin_field, expected_pin = _dawn_pin_from_deps_txt() + pin_field, expected_pin = _dawn_pin_from_deps_txt(deps_file) if git.get(pin_field) != expected_pin: raise ValueError( - f"Dawn manifest {pin_field} {git.get(pin_field)} does not match {DEPS_TXT} {pin_field} {expected_pin}" + f"Dawn manifest {pin_field} {git.get(pin_field)} does not match {deps_file} {pin_field} {expected_pin}" ) -def _validate_dxc_release(manifest: dict[str, Any]) -> None: - expected_tag, expected_url, expected_hash = _dxc_release_from_pipeline() +def _validate_dxc_release(manifest: dict[str, Any], pipeline_file: Path = PLUGIN_WIN_WEBGPU_STAGE) -> None: + expected_tag, expected_url, expected_hash = _dxc_release_from_pipeline(pipeline_file) registration = _find_git_registration(manifest, DXC_REPOSITORY_URL, tag=expected_tag) git = _git_component(registration) if git is None: @@ -162,12 +160,22 @@ def _validate_dawn_dependency_roots(manifest: dict[str, Any]) -> None: raise ValueError(f"Dawn dependency root does not match {DAWN_REPOSITORY_URL}@{dawn_commit}: {comments}") +def validate_manifest( + manifest: dict[str, Any], + *, + deps_file: Path = DEPS_TXT, + pipeline_file: Path = PLUGIN_WIN_WEBGPU_STAGE, +) -> None: + """Validate reviewed metadata against the selected checkout's pins.""" + _validate_dawn_root(manifest, deps_file) + _validate_dxc_release(manifest, pipeline_file) + _validate_dawn_dependency_roots(manifest) + + def main() -> int: try: manifest = _load_manifest() - _validate_dawn_root(manifest) - _validate_dxc_release(manifest) - _validate_dawn_dependency_roots(manifest) + validate_manifest(manifest) except (OSError, ValueError) as ex: print(f"ERROR: {ex}", file=sys.stderr) return 1 diff --git a/cmake/external/onnxruntime_external_deps.cmake b/cmake/external/onnxruntime_external_deps.cmake index b7bc9aff89c08..946d1c4d685de 100644 --- a/cmake/external/onnxruntime_external_deps.cmake +++ b/cmake/external/onnxruntime_external_deps.cmake @@ -754,13 +754,6 @@ if (onnxruntime_USE_WEBGPU) if (NOT onnxruntime_ENABLE_DAWN_BACKEND_D3D12) message(FATAL_ERROR "DAWN_USE_AGILITY_SDK requires the Dawn D3D12 backend.") endif() - if (onnxruntime_USE_EP_API_ADAPTERS) - # Plugin EP packages cannot guarantee that the Agility SDK runtime DLLs are deployed - # next to the host executable. - message(FATAL_ERROR - "DAWN_USE_AGILITY_SDK is not supported with onnxruntime_USE_EP_API_ADAPTERS=ON (plugin EP build). " - "It is intended for local development builds only.") - endif() endif() # TODO: the following code is used to disable building Dawn using vcpkg temporarily diff --git a/cmake/onnxruntime_providers_webgpu.cmake b/cmake/onnxruntime_providers_webgpu.cmake index 9428befacd09e..e7dd0a7055a98 100644 --- a/cmake/onnxruntime_providers_webgpu.cmake +++ b/cmake/onnxruntime_providers_webgpu.cmake @@ -312,6 +312,10 @@ endif() endif() + if (DAWN_USE_AGILITY_SDK) + add_dependencies(onnxruntime_providers_webgpu copy_agility_sdk_dlls) + endif() + # In the plugin EP shared library build the WebGPU EP is its own DLL # (onnxruntime_providers_webgpu.dll) rather than being linked into onnxruntime.dll, so the # delay-load flags accumulated above (notably /DELAYLOAD:user32.dll) must be applied to THIS target. diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index fb976e6ffbc92..55653aaabb438 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -1267,6 +1267,19 @@ if (onnxruntime_USE_WEBGPU) $/onnxruntime/capi/ ) endif() + + if (DAWN_USE_AGILITY_SDK) + add_dependencies(onnxruntime_pybind11_state copy_agility_sdk_dlls) + add_custom_command( + TARGET onnxruntime_pybind11_state POST_BUILD + COMMAND ${CMAKE_COMMAND} -E make_directory + $/onnxruntime/capi/D3D12 + COMMAND ${CMAKE_COMMAND} -E copy_if_different + $/D3D12/D3D12Core.dll + $/D3D12/d3d12SDKLayers.dll + $/onnxruntime/capi/D3D12 + ) + endif() endif() if (onnxruntime_BUILD_DAWN_SHARED_LIBRARY) add_custom_command( diff --git a/cmake/patches/flatbuffers/flatbuffers.patch b/cmake/patches/flatbuffers/flatbuffers.patch index 9fb58e301bba8..afa75e16d3766 100644 --- a/cmake/patches/flatbuffers/flatbuffers.patch +++ b/cmake/patches/flatbuffers/flatbuffers.patch @@ -28,3 +28,24 @@ index bc828a31..3d3effe8 100644 // ElementaryType unsigned short base_type : 4; // Either vector (in table) or array (in struct) +diff --git a/grpc/examples/go/greeter/client/go.mod b/grpc/examples/go/greeter/client/go.mod +index 6f151292..e2107728 100644 +--- a/grpc/examples/go/greeter/client/go.mod ++++ b/grpc/examples/go/greeter/client/go.mod +@@ -10 +10 @@ require ( +- google.golang.org/grpc v1.35.0 ++ google.golang.org/grpc v1.84.0 +diff --git a/grpc/examples/go/greeter/models/go.mod b/grpc/examples/go/greeter/models/go.mod +index c62fb381..b25c1662 100644 +--- a/grpc/examples/go/greeter/models/go.mod ++++ b/grpc/examples/go/greeter/models/go.mod +@@ -7 +7 @@ require ( +- google.golang.org/grpc v1.35.0 ++ google.golang.org/grpc v1.84.0 +diff --git a/grpc/examples/go/greeter/server/go.mod b/grpc/examples/go/greeter/server/go.mod +index 3e10997e..8a1b7208 100644 +--- a/grpc/examples/go/greeter/server/go.mod ++++ b/grpc/examples/go/greeter/server/go.mod +@@ -10 +10 @@ require ( +- google.golang.org/grpc v1.39.0-dev ++ google.golang.org/grpc v1.84.0 diff --git a/docs/BuildWithDawnAgilitySDK.md b/docs/BuildWithDawnAgilitySDK.md index 71337e6dd4a52..fa75e7abe64cf 100644 --- a/docs/BuildWithDawnAgilitySDK.md +++ b/docs/BuildWithDawnAgilitySDK.md @@ -1,6 +1,28 @@ # Build WebGPU with the D3D12 Agility SDK -Use `--use_dawn_agility_sdk` to build Dawn's D3D12 backend with ONNX Runtime's pinned preview D3D12 Agility SDK: +Use `--use_dawn_agility_sdk` to build Dawn's D3D12 backend with ONNX Runtime's pinned preview D3D12 Agility SDK. + +For the WebGPU Plugin EP wheel (`onnxruntime-ep-webgpu`), first build the shared plugin library and then package it: + +```powershell +python tools\ci_build\build.py ` + --build_dir build\WebGpuPlugin ` + --no_telemetry ` + --config Release ` + --update --build ` + --use_webgpu shared_lib ` + --use_dawn_agility_sdk ` + --skip_tests ` + --parallel + +python plugin-ep-webgpu\python\build_wheel.py ` + --binary_dir build\WebGpuPlugin\Release\Release ` + --version 0.1.0.devYYYYMMDD ` + --use-dawn-agility-sdk ` + --output_dir dist +``` + +For a local ONNX Runtime Python wheel with the built-in WebGPU EP: ```powershell python tools\ci_build\build.py ` @@ -10,14 +32,19 @@ python tools\ci_build\build.py ` --update --build ` --use_webgpu ` --use_dawn_agility_sdk ` - --build_shared_lib ` - --skip_tests + --build_wheel ` + --skip_tests ` + --parallel ``` -This option is intended for local development and supports Windows desktop x86, x64, and ARM64 targets. Windows ARM32, -ARM64EC, and WindowsStore/UWP targets are not supported. WebGPU Plugin EP, Python wheels, C#, NuGet, Java, and Node.js -packages are also not supported because they do not deploy the required D3D12 runtime DLLs. Custom Dawn checkouts -selected with `onnxruntime_CUSTOM_DAWN_SRC_PATH` are not supported. +Wheels with the built-in WebGPU EP are no longer officially published. Local builds retain Agility SDK support +for now; the built-in WebGPU EP is planned for removal. Use the plugin EP for the published package workflow. + +This option supports Windows desktop x86, x64, and ARM64 targets. Windows ARM32, ARM64EC, and WindowsStore/UWP targets +are not supported. Both local built-in WebGPU EP Python wheels and WebGPU Plugin EP Python wheels are supported. +C#, NuGet, Java, and Node.js packaging through `build.py` is not supported because those paths do not deploy the +required DLLs; plugin NuGet packaging uses the separate `plugin-ep-webgpu/csharp/pack_nuget.py` script. +Custom Dawn checkouts selected with `onnxruntime_CUSTOM_DAWN_SRC_PATH` are not supported. The pinned SDK requires Windows 10 version 1909 or newer. For versions 1909, 2004, and 20H2, the minimum OS build revisions are: @@ -33,8 +60,9 @@ Because the pinned SDK is a preview release, Windows Developer Mode must be enab and consult Microsoft's [AgilitySDK 721 Preview announcement](https://devblogs.microsoft.com/directx/announcing-agilitysdk-721-preview-and-more-shader-model-6-10-features/) for the current per-vendor driver links and feature-support table. Feature availability depends on the GPU and driver. -At runtime, Dawn loads the SDK DLLs from the `D3D12` directory next to the executable. Look for `[AgilitySDK] active` -in the Dawn log to confirm successful activation. +At runtime, Dawn loads the SDK DLLs from the `D3D12` directory next to the Dawn module. In a local wheel with the +built-in WebGPU EP this is `onnxruntime/capi/D3D12`; in the WebGPU plugin wheel it is `onnxruntime_ep_webgpu/D3D12`. +Look for `[AgilitySDK] active` in the Dawn log to confirm successful activation. For direct CMake configuration, enable `onnxruntime_USE_WEBGPU`, `onnxruntime_ENABLE_DAWN_BACKEND_D3D12`, and `DAWN_USE_AGILITY_SDK`. The SDK is downloaded to `third_party/agility-sdk` in the Dawn source tree managed by ONNX diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 2703de7cc1fae..517f11af63dd1 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -19,6 +19,7 @@ Do not modify directly.* * com.microsoft.CausalConvWithState * com.microsoft.ComplexMul * com.microsoft.ComplexMulConj + * com.microsoft.ConvStateReplay * com.microsoft.ConvTransposeWithDynamicPads * com.microsoft.CropAndResize * com.microsoft.DecoderAttention @@ -1116,6 +1117,69 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.ConvStateReplay** + + Reconstruct rolling convolution history after retaining a prefix of captured inputs, without + performing convolution arithmetic. The WebGPU-only implementation copies FP16 or FP32 payload + bits exactly, including signed zero and NaN payloads, and waits for checked GPU completion. + + Inputs 0-2 are full rank-1 backing allocations of the same element type. + CPU int64 metadata [8] is ordered as: + source_offset_bytes, destination_offset_bytes, capture_offset_bytes, + batch_size, channels, history_width, capacity, kept_count. + Offsets are logical byte offsets into the full backing tensors, aligned to the element size + (two bytes for FP16, four for FP32), not WebGPU storage-binding offsets. + All fields are nonnegative; channels must be positive; 0 <= kept_count <= capacity. + Zero batch size, history width, or capture capacity is valid. Zero state size is a no-op. + Geometry, backing byte extents and metadata must fit the implementation's uint32 indexing range. + + The source view is the incoming snapshot immediately before capture, with layout + [batch_size, channels, history_width]. Captured inputs have layout [batch_size, capacity, channels]. + capacity is the original capture stride, not kept_count. All retained inputs must have been captured. + For each batch b, channel c and position p, with W = history_width and k = kept_count: + destination[b,c,p] = source[b,c,p+k] if p+k < W + = capture[b,p+k-W,c] otherwise. + For k >= W only the last W inputs of the retained prefix are used; unused capture tails are not read. + + The output must be explicitly prebound to the same full allocation as destination_backing. + Destination backing must be distinct from source and capture backing, even for disjoint logical views. + Only the designated destination view is modified; neighboring elements and all source/capture data + are preserved. One descriptor is supported per invocation; graph capture is not supported. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Inputs + +
+
source_backing : T
+
Full backing tensor containing the incoming convolution history.
+
capture_backing : T
+
Full backing tensor containing captured token-major inputs.
+
destination_backing : T
+
Full inactive destination-state backing tensor.
+
metadata : TI
+
CPU int64 [8]: byte offsets, batch, channels, history, capacity, kept_count.
+
+ +#### Outputs + +
+
destination_out : T
+
Same backing allocation and shape as destination_backing.
+
+ +#### Type Constraints + +
+
T : tensor(float16), tensor(float)
+
State and captured inputs are FP16 or FP32.
+
TI : tensor(int64)
+
CPU replay metadata is int64.
+
+ + ### **com.microsoft.ConvTransposeWithDynamicPads** #### Version @@ -5464,9 +5528,9 @@ This version of the operator has been available since version 1 of the 'com.micr
cumulative_sequence_length : S
A tensor with shape (batch_size + 1). It specifies the cumulative sequence lengths between the packed entries in Q/K/V.
past_seqlens : S
-
A tensor with shape (batch_size). It specifies the past lengths of cached sequence in the KV cache.
+
A tensor with shape (batch_size). It specifies the past lengths of cached sequences. CUDA sanitizes malformed values on device: an invalid value suppresses that sequence's cache writes and produces zero output. A backend requiring exact host lengths may instead return INVALID_ARGUMENT when attention_metadata is absent.
block_table : S
-
2D tensor with shape (batch_size, max_blocks_per_sequence) that maps each sequence in the batch to itscorresponding blocks in the KV cache.
+
2D tensor with shape (batch_size, max_blocks_per_sequence) that maps each sequence in the batch to itscorresponding blocks in the KV cache. -1 denotes an unmapped page. CUDA converts values below -1 or greater than or equal to num_blocks to the unmapped sentinel before resolving cache addresses; mapped pages remain attendable.
cos_cache (optional) : T
2D tensor with shape (max total seqlen, head_size / 2).
sin_cache (optional) : T
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 5617394bfa8df..cd5ed43763c43 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -119,7 +119,8 @@ The **OpSet Version** column uses the following notation: |||[17, 19]|**T1** = tensor(double), tensor(float)
**T2** = tensor(int32), tensor(int64)| |DeformConv|*in* X:**T**
*in* W:**T**
*in* offset:**T**
*in* B:**T**
*in* mask:**T**
*out* Y:**T**|22+|**T** = tensor(double), tensor(float)| |||[19, 21]|**T** = tensor(double), tensor(float)| -|DepthToSpace|*in* input:**T**
*out* output:**T**|13+|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| +|DepthToSpace|*in* input:**T**
*out* output:**T**|28+|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| +|||[13, 27]|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| |||[11, 12]|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| |||[1, 10]|**T** = tensor(double), tensor(float)| |DequantizeLinear|*in* x:**T**
*in* x_scale:**tensor(float)**
*in* x_zero_point:**T**
*out* y:**tensor(float)**

or

*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T2**

or

*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T3**|25+|**T1** = tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int2), tensor(int32), tensor(int4), tensor(int8), tensor(uint16), tensor(uint2), tensor(uint4), tensor(uint8)
**T2** = tensor(float), tensor(float16)
**T3** = tensor(float), tensor(float16)| @@ -485,7 +486,8 @@ The **OpSet Version** column uses the following notation: |||[1, 21]|**T** = tensor(float)| |Softsign|*in* input:**T**
*out* output:**T**|22+|**T** = tensor(float)| |||[1, 21]|**T** = tensor(float)| -|SpaceToDepth|*in* input:**T**
*out* output:**T**|[13, 27]|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| +|SpaceToDepth|*in* input:**T**
*out* output:**T**|28+|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| +|||[13, 27]|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| |||[1, 12]|**T** = tensor(double), tensor(float), tensor(int8), tensor(uint8)| |Split|*in* input:**T**
*in* split:**T**
*out* outputs...:**T**

or

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

or

*in* input:**T**
*out* outputs:**T**|18+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[13, 17]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| @@ -739,7 +741,8 @@ The **OpSet Version** column uses the following notation: |||[11, 13]|**T** = tensor(double), tensor(float), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)
**T2** = tensor(int32), tensor(int64)| |DeformConv|*in* X:**T**
*in* W:**T**
*in* offset:**T**
*in* B:**T**
*in* mask:**T**
*out* Y:**T**|22+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |||[19, 21]|**T** = tensor(double), tensor(float), tensor(float16)| -|DepthToSpace|*in* input:**T**
*out* output:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16)| +|DepthToSpace|*in* input:**T**
*out* output:**T**|28+|**T** = tensor(double), tensor(float), tensor(float16)| +|||[13, 27]|**T** = tensor(double), tensor(float), tensor(float16)| |||[11, 12]|**T** = tensor(double), tensor(float), tensor(float16)| |||[1, 10]|**T** = tensor(double), tensor(float), tensor(float16)| |DequantizeLinear|*in* x:**T**
*in* x_scale:**tensor(float)**
*in* x_zero_point:**T**
*out* y:**tensor(float)**

or

*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T2**

or

*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T3**|25+|**T1** = tensor(float8e4m3fn), tensor(float8e5m2), tensor(int4), tensor(int8), tensor(uint4), tensor(uint8)
**T2** = tensor(float), tensor(float16)
**T3** = tensor(float), tensor(float16)| @@ -1031,7 +1034,8 @@ The **OpSet Version** column uses the following notation: |||[1, 21]|**T** = tensor(double), tensor(float), tensor(float16)| |Softsign|*in* input:**T**
*out* output:**T**|22+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |||[1, 21]|**T** = tensor(double), tensor(float), tensor(float16)| -|SpaceToDepth|*in* input:**T**
*out* output:**T**|[13, 27]|**T** = tensor(double), tensor(float), tensor(float16)| +|SpaceToDepth|*in* input:**T**
*out* output:**T**|28+|**T** = tensor(double), tensor(float), tensor(float16)| +|||[13, 27]|**T** = tensor(double), tensor(float), tensor(float16)| |||[1, 12]|**T** = tensor(double), tensor(float), tensor(float16)| |Split|*in* input:**T**
*in* split:**T**
*out* outputs...:**T**

or

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

or

*in* input:**T**
*out* outputs:**T**|18+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[13, 17]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| @@ -1196,7 +1200,8 @@ The **OpSet Version** column uses the following notation: |ConvTranspose|*in* X:**T**
*in* W:**T**
*in* B:**T**
*out* Y:**T**|22+|**T** = tensor(float), tensor(float16)| |||[11, 21]|**T** = tensor(float), tensor(float16)| |||[1, 10]|**T** = tensor(float), tensor(float16)| -|DepthToSpace|*in* input:**T**
*out* output:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16)| +|DepthToSpace|*in* input:**T**
*out* output:**T**|28+|**T** = tensor(double), tensor(float), tensor(float16)| +|||[13, 27]|**T** = tensor(double), tensor(float), tensor(float16)| |||[11, 12]|**T** = tensor(double), tensor(float), tensor(float16)| |||[1, 10]|**T** = tensor(double), tensor(float), tensor(float16)| |GlobalAveragePool|*in* X:**T**
*out* Y:**T**|22+|**T** = tensor(float), tensor(float16)| @@ -1213,7 +1218,8 @@ The **OpSet Version** column uses the following notation: |||10|**I** = tensor(int64)
**T** = tensor(float), tensor(float16)| |||[8, 9]|**I** = tensor(int64)
**T** = tensor(float), tensor(float16)| |||[1, 7]|**T** = tensor(float), tensor(float16)| -|SpaceToDepth|*in* input:**T**
*out* output:**T**|[13, 27]|**T** = tensor(double), tensor(float), tensor(float16)| +|SpaceToDepth|*in* input:**T**
*out* output:**T**|28+|**T** = tensor(double), tensor(float), tensor(float16)| +|||[13, 27]|**T** = tensor(double), tensor(float), tensor(float16)| |||[1, 12]|**T** = tensor(double), tensor(float), tensor(float16)| | | | | @@ -1295,7 +1301,8 @@ The **OpSet Version** column uses the following notation: |||11+|**T** = tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| |DFT|*in* input:**T1**
*in* dft_length:**T2**
*in* axis:**tensor(int64)**
*out* output:**T1**

or

*in* input:**T1**
*in* dft_length:**T2**
*out* output:**T1**|20+|**T1** = tensor(double), tensor(float), tensor(float16)
**T2** = tensor(int32), tensor(int64)| |||17+|**T1** = tensor(double), tensor(float), tensor(float16)
**T2** = tensor(int32), tensor(int64)| -|DepthToSpace|*in* input:**T**
*out* output:**T**|13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|DepthToSpace|*in* input:**T**
*out* output:**T**|28+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|||13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||11+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||1+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |DequantizeLinear|*in* x:**T**
*in* x_scale:**tensor(float)**
*in* x_zero_point:**T**
*out* y:**tensor(float)**

or

*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T2**

or

*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T3**|21+|**T1** = tensor(int4), tensor(int8), tensor(uint4), tensor(uint8)
**T2** = tensor(float), tensor(float16)| @@ -1571,7 +1578,8 @@ The **OpSet Version** column uses the following notation: |||1+|**T** = tensor(float), tensor(float16)| |Softplus|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |Softsign|*in* input:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| -|SpaceToDepth|*in* input:**T**
*out* output:**T**|13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|SpaceToDepth|*in* input:**T**
*out* output:**T**|28+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|||13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||1+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |Split|*in* input:**T**
*in* split:**T**
*out* outputs...:**T**

or

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

or

*in* input:**T**
*out* outputs:**T**|18+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md index 73682f6a237c6..4adb08e5080dd 100644 --- a/docs/contrib_ops/cuda/moe_qmoe.md +++ b/docs/contrib_ops/cuda/moe_qmoe.md @@ -831,6 +831,9 @@ debug switches. | `ORT_FP4_GEMV_AUTOTUNE` | `0` | Opt-in per-shape autotune of the GEMV CtaN/Threads tiling. Enabling it synchronizes the first uncached inference for each shape. | | `ORT_FP4_GEMV_AUTOTUNE_LOG` | `0` | Set to `1` to log the chosen GEMV configs per shape. | | `ORT_FP4_GEMV_DECODE_KERNEL` | `1` | NVFP4 offline K-packed (`weights_prepacked=1`) decode GEMV mode: `1` (default) enables `MoeGemvFp4DecodeKernel` in the measured SM90 decode envelope (`expanded_num_rows <= 20`, `N, K <= 2560`, and per-projection `|global_scale| * 448 < 3.98`); `0` forces `MoeGemvFp4RawKPackedKernel`; `2` opts into the full kernel envelope (`SM90+`, `expanded_num_rows <= 80`, `K <= 8192`, including the generic scale decoder). Outputs are bit-identical across all modes. Read once per process. See [§9b.2](#9b2-fused-gemv-decode-group-16). | +| `ORT_FP4_GEMV_FUSED_ROUTE` | `1` | Single-CTA fused softmax + top-k + expert-map decode routing kernel for offline K-packed NVFP4 on SM90+ (`num_rows <= 4`, `num_experts == 512`, `k in {1, 2, 4, 6, 8, 10}`). Bit-identical to `LaunchSoftmaxTopK` + `fusedBuildExpertMapsSortFirstToken`; set to `0` to disable. | +| `ORT_FP4_GEMV_PDL` | `1` | Programmatic dependent launch (`cudaLaunchAttributeProgrammaticStreamSerialization`) overlapping FC1 -> FC2 -> finalize on the offline K-packed NVFP4 decode path on SM90+. Set to `0` to disable. | +| `ORT_FP4_GEMV_FC1_PDL` | `0` | Opt-in programmatic dependent launch of FC1 behind the fused routing kernel. Off by default. | | `ORT_FP4_GEMV_INTERLEAVED` | `0` | **Experimental, opt-in.** Routes the MXFP4 decode GEMV through the `ColumnMajorInterleaved` weight layout (`kInterleave=4`, `kStepK=32`) with dtype-conditional accumulation. fp16 gets faster decode; bf16 stays accuracy-safe. Default off keeps the shipping `ColumnMajor` path byte-for-byte unchanged. See [§9.10](#910-interleaved-gemv-layout--dtype-conditional-accumulation). | | `ORT_FP4_GEMV_INTERLEAVED_HALFACC` | `0` | **Override.** When `ORT_FP4_GEMV_INTERLEAVED=1`, forces 16-bit accumulation for *both* fp16 and bf16, overriding the dtype-conditional policy; regresses bf16 accuracy, so it is off by default. | | `ORT_FP4_SM80_GEMM` | `1` | Routes SM80+ (including SM120/SM121) FP4 prefill through the fused-dequant grouped GEMM. Set to `0` to force dense fallback for debugging or comparison. Decode routes through the fused MXFP4 GEMV, reading the *same* pre-packed buffer as prefill. In this regime the raw e2m1 initializers are released after `PrePack`. See [§9.11](#911-single-copy-sm80-weights-prefill--decode-share-one-buffer). | diff --git a/docs/contrib_ops/cuda/paged_attention.md b/docs/contrib_ops/cuda/paged_attention.md index 7bcd49109abc0..b45d817e786a6 100644 --- a/docs/contrib_ops/cuda/paged_attention.md +++ b/docs/contrib_ops/cuda/paged_attention.md @@ -83,8 +83,8 @@ never populated or consumed), `head_sink` / smooth softmax, QK-Norm, quantized c Structural limits in the current implementation: -- `batch_size <= 256` — `LaunchGetCumulativeSeqlensKV` uses a per-block `cub::BlockScan` with 256 - threads and independent blocks. +- Sequence metadata is canonicalized with parallel device kernels and grid-wide scans; this path + has no fixed batch-size limit. - `block_size % 256 == 0` — see [§18](#18-known-defects-to-fix-first); this is almost certainly a bug. *(Partly true. It is a genuine FlashAttention tiling constraint, but a head-size-dependent one. See the implementation note in §18.1: validation now accepts any power-of-two `block_size >= 16` and @@ -181,6 +181,24 @@ matches the landing order in [§19](#19-phasing), so the schema grows monotonica For `k_cache_dtype=v_cache_dtype="int4"`, the cache tensors use `uint8` storage and their last dimension is `(head_size + 1) / 2`, not `head_size`. +`cumulative_sequence_length`, `past_seqlens`, and `block_table` are sanitized on the CUDA compute +stream before any cache address is resolved: + +- cumulative query offsets are clamped to `[0, token_count]`, made nondecreasing, and fixed to start + at zero and end at `token_count`; +- a negative `past_seqlens` value, or one that leaves insufficient capacity for that sequence's + query tokens, suppresses that sequence's cache writes and produces zero output; +- `block_table == -1` remains the supported unmapped-page sentinel; values `< -1` or + `>= num_blocks` are converted to that sentinel. Sentinel pages are excluded from attention while + mapped pages remain attendable. + +When `attention_metadata` supplies replay-wide bounds, or while a CUDA Graph is being captured or +replayed, sanitization is entirely device-side and malformed values follow the behavior above. +Without metadata, a backend that needs exact host lengths for workspace sizing or XQA eligibility +performs the existing device-to-host length readback. On that diagnostic path, malformed +`cumulative_sequence_length` or `past_seqlens` returns `INVALID_ARGUMENT` instead of being silently +sanitized. Block-table safety remains device-side on every path. + `max_context_len` is the largest per-sequence total KV length in the batch, bounded above by `block_table.shape[1] * block_size`. @@ -520,9 +538,8 @@ measurable win for large batches and eliminates the OOB failure mode. The kernel - Rank 1, `shape[0] == token_count`. - Element type `int32`. -- Range checking on device is a debug-build assertion only; host-side range checking would require a - D→H copy per step. The op documents that out-of-range values are undefined behavior, consistent - with `block_table` today. +- Range checking is performed on device without a host synchronization. `-1` suppresses the write; + values below `-1` or at/above `num_blocks * block_size` are also treated as suppressed writes. ### 5.5 Why not derive-only @@ -1451,7 +1468,6 @@ Consolidated, to be implemented in `paged_attention_helper::CheckInputs`. Every - `block_table` rank 2 with `dim0 == batch_size`. - `cos_cache`/`sin_cache` both present or both absent; required when `do_rotary == 1`. - `key_cache_out` must alias `key_cache`; same for value. -- `batch_size <= 256` (BlockScan limitation — to be lifted, see §18). **Corrected:** - `block_size` must be a power of two in `{16, 32, 64, 128, 256}` (**replaces** `block_size % 256 == 0`; see §18). @@ -1552,7 +1568,7 @@ reference oracle for every shared feature and catches drift automatically. Run i - `block_table` entries of `-1` (unmapped) are masked out. - Ragged batches including sequences with **zero** new tokens. - `token_count == 0` early-out. -- `batch_size` at and just above the BlockScan limit (must error, not corrupt). +- `batch_size` at and above the former 256-sequence BlockScan limit. - Sliding-window block pruning: pruned run bit-matches the unpruned run. - Rolling sliding-window cache (§9.3): a run holding only `ceil(W / block_size) + 1` blocks per sequence, recycled through `slot_mapping` with the evicted `block_table` entries set to `-1`, @@ -1658,9 +1674,9 @@ These block the feature work and should land ahead of it. `block_table`. Guard with an early `return` when `token_id >= cumulative_seqlens_q[batch_size]`. (`slot_mapping` removes the search entirely on the write path, but the gather path still needs the fix.) -3. **`batch_size <= 256`.** Replace the per-block `cub::BlockScan` with a grid-wide scan (or a single - 256-thread block doing a strided serial scan) so continuous batching is not capped at 256 - concurrent sequences — a real limit for a serving op. +3. **`batch_size <= 256`.** **Fixed.** Sequence metadata canonicalization now uses grid-wide CUB + scans and parallel per-sequence bounds checks, so continuous batching is no longer capped at + 256 concurrent sequences. 4. **Per-step D→H synchronization — and it blocks CUDA graph capture.** ~~`max_query_len` (and `total_kv_tokens` for MEA) are obtained via `cudaStreamSynchronize` every step, once per layer. This is not only a throughput bug for the op's primary use case: `cudaStreamSynchronize` on a diff --git a/include/onnxruntime/core/framework/kernel_dispatch_policy.h b/include/onnxruntime/core/framework/kernel_dispatch_policy.h new file mode 100644 index 0000000000000..7cfd063ea70ee --- /dev/null +++ b/include/onnxruntime/core/framework/kernel_dispatch_policy.h @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +namespace onnxruntime { + +// Hardware-neutral intent for future operator/EP policy resolution. +// Reader-only: no value currently changes kernel dispatch or workspace sizing. +enum class KernelDispatchPolicy { + Auto, + Latency, + Memory, + Safe, +}; + +// Matching is case-sensitive. Empty and "auto" select Auto; failure leaves policy unchanged. +inline bool TryParseKernelDispatchPolicy( + std::string_view value, + KernelDispatchPolicy& policy) noexcept { + if (value.empty() || value == "auto") { + policy = KernelDispatchPolicy::Auto; + return true; + } + if (value == "latency") { + policy = KernelDispatchPolicy::Latency; + return true; + } + if (value == "memory") { + policy = KernelDispatchPolicy::Memory; + return true; + } + if (value == "safe") { + policy = KernelDispatchPolicy::Safe; + return true; + } + return false; +} + +} // namespace onnxruntime diff --git a/include/onnxruntime/core/framework/resource_accountant.h b/include/onnxruntime/core/framework/resource_accountant.h index c7a80839b7221..669bdf72edaa3 100644 --- a/include/onnxruntime/core/framework/resource_accountant.h +++ b/include/onnxruntime/core/framework/resource_accountant.h @@ -69,6 +69,8 @@ struct WorkspaceEstimatorConfig { std::optional cuda_fpa_intb_profile_m; // Validated nonnegative, int32-compatible envelope; zero means unspecified. int64_t cuda_gqa_workspace_max_total_sequence_length = 0; + // Raw session.kernel_dispatch_policy value; consumers share the framework parser. + std::optional kernel_dispatch_policy = std::nullopt; }; // Reads the caller-declared CUDA GQA envelope and validates its runtime ABI range. 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 1251fae691b7c..39200df71ad80 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -503,6 +503,15 @@ static const char* const kOrtSessionOptionsCudaFpAIntBGemvWaveAware = "ep.cuda.f static const char* const kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength = "ep.cuda.gqa_workspace_max_total_sequence_length"; +/// Reader-only, hardware-neutral kernel dispatch intent shared across operators and EPs. +/// Valid values are "auto", "latency", "memory", and "safe"; matching is case-sensitive. +/// Unset or empty values map to "auto". Other values cause INVALID_ARGUMENT during +/// session initialization. +/// CUDA GroupQueryAttention Level-1 estimation is the first consumer; sizing does not use it. +/// No value changes runtime dispatch, workspace estimates, or partitioning, and this option +/// does not enforce a workload envelope or memory budget or provide a no-OOM guarantee. +static const char* const kOrtSessionOptionsKernelDispatchPolicy = "session.kernel_dispatch_policy"; + /// Maximum number of rows of input A per CUDA MatMulNBits fpA_intB GEMM launch. Values below 8192 are /// rounded down to a supported tactic-profiler M bucket. Chunking requires M to exceed this limit /// and the 256 MiB estimated A/C row-size gate; ORT_MATMULNBITS_FORCE_CHUNKED=1 bypasses that gate. diff --git a/onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h index 6855e348fcfd6..214a26be6df0b 100644 --- a/onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h @@ -258,6 +258,10 @@ Status CheckBlockTable(const T* block_table, const int batch_size, int& max_num_ return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "block_table dimension 0 should be batch_size, got ", block_table_dims[0]); + } else if (block_table_dims[1] <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "block_table dimension 1 must be positive, got ", + block_table_dims[1]); } max_num_blocks_per_seq = static_cast(block_table_dims[1]); return Status::OK(); @@ -396,6 +400,129 @@ inline Status CheckKVCacheDataType(const KVCacheDataType cache_dtype, const KVCa "'. Leave the attribute at '' to use the tensor's element type."); } +inline Status CheckSequenceLengthValues(const int32_t* cumulative_seqlens_q, + const int32_t* cumulative_seqlens_kv, + int batch_size, + int max_num_blocks_per_seq, + int block_size, + int token_count) { + if (batch_size <= 0 || max_num_blocks_per_seq <= 0 || block_size <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "PagedAttention validation requires positive " + "batch_size/max_num_blocks_per_seq/block_size."); + } + + if (cumulative_seqlens_q[0] != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "cumulative_seqlens_q must start with 0, got ", + cumulative_seqlens_q[0]); + } + + if (cumulative_seqlens_q[batch_size] != token_count) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "cumulative_seqlens_q must end with token_count (", token_count, + "), got ", cumulative_seqlens_q[batch_size]); + } + + if (cumulative_seqlens_kv[0] != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "cumulative_seqlens_kv must start with 0, got ", + cumulative_seqlens_kv[0]); + } + + const int64_t max_cache_sequence_length = static_cast(max_num_blocks_per_seq) * block_size; + for (int b = 0; b < batch_size; ++b) { + const int32_t q_start = cumulative_seqlens_q[b]; + const int32_t q_end = cumulative_seqlens_q[b + 1]; + + if (q_start < 0 || q_end < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "cumulative_seqlens_q values must be non-negative. Invalid value at index ", + (q_start < 0 ? b : b + 1), ": ", (q_start < 0 ? q_start : q_end)); + } + + if (q_end < q_start) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "cumulative_seqlens_q must be non-decreasing."); + } + + const int32_t q_len = q_end - q_start; + const int32_t kv_start = cumulative_seqlens_kv[b]; + const int32_t kv_end = cumulative_seqlens_kv[b + 1]; + if (kv_start < 0 || kv_end < kv_start) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "cumulative_seqlens_kv must be non-negative and non-decreasing."); + } + + const int64_t past_length = static_cast(kv_end) - kv_start - q_len; + if (past_length < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "past_seqlens values must be non-negative. Invalid value: ", + past_length); + } + + if (q_len == 0) { + if (past_length > max_cache_sequence_length) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "past_seqlens exceeds max_num_blocks_per_seq * block_size for zero-token sequence. Invalid value: ", + past_length); + } + } else { + if (past_length >= max_cache_sequence_length) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "past_seqlens must be less than max_num_blocks_per_seq * block_size when q_len > 0. Invalid value: ", + past_length); + } + + const int64_t last_position = past_length + q_len - 1; + if (last_position >= max_cache_sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "past_seqlens + query_length exceeds block_table capacity for a sequence."); + } + } + } + + return Status::OK(); +} + +inline Status CheckSequenceLengthAndPastValues(const int32_t* cumulative_seqlens_q, + const int32_t* past_seqlens, + int batch_size, + int max_num_blocks_per_seq, + int block_size, + int token_count) { + if (cumulative_seqlens_q[0] != 0 || + cumulative_seqlens_q[batch_size] != token_count) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "cumulative_seqlens_q must start with 0 and end with token_count."); + } + + const int64_t capacity = static_cast(max_num_blocks_per_seq) * block_size; + for (int b = 0; b < batch_size; ++b) { + const int32_t q_start = cumulative_seqlens_q[b]; + const int32_t q_end = cumulative_seqlens_q[b + 1]; + const int32_t past = past_seqlens[b]; + if (q_start < 0 || q_end < q_start || past < 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "cumulative_seqlens_q must be non-negative and non-decreasing, and past_seqlens must be non-negative."); + } + + const int64_t q_len = static_cast(q_end) - q_start; + if ((q_len == 0 && past > capacity) || + (q_len > 0 && (past >= capacity || static_cast(past) + q_len > capacity))) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "past_seqlens and query length exceed block_table capacity."); + } + } + + return Status::OK(); +} + template Status CheckInputs(const T* query, const T* key, diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index a5c69221347c8..79659e6b0ef50 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_data.h @@ -232,7 +232,9 @@ struct PagedAttentionData { const float* v_scale = nullptr; const int* cumulative_seqlens_q = nullptr; const int* past_seqlens = nullptr; + const int* sequence_validity = nullptr; const int* block_table = nullptr; + const int* cudnn_block_table = nullptr; // Optional explicit write slots, one per query token, into the cache viewed as // [num_blocks * block_size, kv_num_heads, head_size]. A value of -1 suppresses the K/V store // for that token (prefix cache hit / rejected speculative token). nullptr keeps the legacy diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index a6b97286d1733..55ef9970d669d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -124,7 +124,7 @@ __global__ void EmbedLayerNormKernel( int hidden_size, const int* input_ids, const int* segment_ids, const T* beta, const T* gamma, const T* word_embedding, const T* position_embedding, const T* segment_embedding, float epsilon, T* output, T* embedding_sum, const int* position_ids, const bool broadcast_position_ids) { - KeyValuePairSum pair_sum; + if (hidden_size == 0) return; // 1. lookup word and segment of the block // blockIdx.x = position in the sequence // blockIdx.y = batch @@ -133,8 +133,8 @@ __global__ void EmbedLayerNormKernel( __shared__ int word_id; __shared__ int segment_id; __shared__ int position_id; + __shared__ float origin; - const float rld = 1.f / hidden_size; const int sequence_position = blockIdx.y * gridDim.x + blockIdx.x; if (threadIdx.x == 0) { word_id = input_ids[sequence_position]; @@ -150,6 +150,10 @@ __global__ void EmbedLayerNormKernel( } else { position_id = position_ids[sequence_position]; } + const T first_word = word_embedding[word_id * hidden_size]; + const T first_segment = segment_embedding == nullptr ? T(0) : segment_embedding[segment_id * hidden_size]; + const T first_position = position_embedding[position_id * hidden_size]; + origin = static_cast(first_word + first_segment + first_position); } __syncthreads(); @@ -162,7 +166,7 @@ __global__ void EmbedLayerNormKernel( // the output offset is given by b * (sequence_length * hidden_size) + s * hidden_size const int output_offset = sequence_position * hidden_size; - cub::KeyValuePair thread_data(0.f, 0.f); + WelfordData thread_data{}; for (int it = threadIdx.x; it < hidden_size; it += TPB) { const T w(word_embedding[word_offset + it]); @@ -178,12 +182,11 @@ __global__ void EmbedLayerNormKernel( } const float val_f = static_cast(val); - const float rldval = rld * val_f; - thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val_f)); + thread_data.Update(val_f - origin); } // 3. layer norm on the sum - LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, epsilon, output); + LayerNorm(thread_data, origin, hidden_size, output_offset, beta, gamma, epsilon, output); } template diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h index a17fb92bd300e..dc566eecc1106 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h @@ -541,14 +541,26 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons // We move K and V to the last block. const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; const int* block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; + const auto page_id = [block_table](int index) { + return block_table[index] < 0 ? 0 : block_table[index]; + }; + const auto mask_unmapped_page = [block_table, ¶ms](auto& scores, int block) { + if (block_table != nullptr && + block_table[block * kBlockN / params.page_block_size] < 0) { +#pragma unroll + for (int i = 0; i < size(scores); ++i) { + scores(i) = -kInfinity; + } + } + }; const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; const index_t row_offset_k = block_table == nullptr ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride - : block_table[block_table_idx] * params.k_batch_stride + block_table_offset * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride; + : page_id(block_table_idx) * params.k_batch_stride + block_table_offset * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride; const index_t row_offset_v = block_table == nullptr ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride - : block_table[block_table_idx] * params.v_batch_stride + block_table_offset * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride; + : page_id(block_table_idx) * params.v_batch_stride + block_table_offset * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride; Tensor mQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb)), make_shape(binfo.actual_seqlen_q, params.h, params.d), @@ -579,6 +591,13 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons Tensor tKsK = gmem_thr_copy_QKV.partition_D(sK); Tensor tVgV = gmem_thr_copy_QKV.partition_S(gV); // (VCPY, VCPY_N, VCPY_K) Tensor tVsV = gmem_thr_copy_QKV.partition_D(sV); + const auto clear_unmapped_v = [block_table, ¶ms, &tVsV](int block) { + if (block_table != nullptr && + block_table[block * kBlockN / params.page_block_size] < 0) { + clear(tVsV); + __syncthreads(); + } + }; typename Kernel_traits::TiledMma tiled_mma; auto thr_mma = tiled_mma.get_thread_slice(tidx); @@ -721,7 +740,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons const int block_table_offset_cur = n_block * kBlockN - block_table_idx_cur * params.page_block_size; const int block_table_idx_next = (n_block - 1) * kBlockN / params.page_block_size; const int block_table_offset_next = (n_block - 1) * kBlockN - block_table_idx_next * params.page_block_size; - const int table_diff = block_table[block_table_idx_next] - block_table[block_table_idx_cur]; + const int table_diff = page_id(block_table_idx_next) - page_id(block_table_idx_cur); const int offset_diff = block_table_offset_next - block_table_offset_cur; tVgV.data() = tVgV.data() + table_diff * params.v_batch_stride + offset_diff * params.v_row_stride; tKgK.data() = tKgK.data() + table_diff * params.k_batch_stride + offset_diff * params.k_row_stride; @@ -815,7 +834,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons const int block_table_offset_cur = (n_block + 1) * kBlockN - block_table_idx_cur * params.page_block_size; const int block_table_idx_next = n_block * kBlockN / params.page_block_size; const int block_table_offset_next = n_block * kBlockN - block_table_idx_next * params.page_block_size; - tVgV.data() = tVgV.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.v_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.v_row_stride; + tVgV.data() = tVgV.data() + (page_id(block_table_idx_next) - page_id(block_table_idx_cur)) * params.v_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.v_row_stride; } FLASH_NAMESPACE::copy(gmem_tiled_copy_QKV, tVgV, tVsV, tKVcKV, tKVpKV); } else { @@ -835,9 +854,11 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons mask.template apply_mask( acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, kNWarps * 16); + mask_unmapped_page(acc_s, n_block); FLASH_NAMESPACE::cp_async_wait<0>(); __syncthreads(); + clear_unmapped_v(n_block); // if (tidx == 0 && blockIdx.y == 0 && blockIdx.z == 0) { print(tVsV); } // __syncthreads(); @@ -850,7 +871,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons const int block_table_offset_cur = n_block * kBlockN - block_table_idx_cur * params.page_block_size; const int block_table_idx_next = (n_block - 1) * kBlockN / params.page_block_size; const int block_table_offset_next = (n_block - 1) * kBlockN - block_table_idx_next * params.page_block_size; - tKgK.data() = tKgK.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.k_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.k_row_stride; + tKgK.data() = tKgK.data() + (page_id(block_table_idx_next) - page_id(block_table_idx_cur)) * params.k_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.k_row_stride; } FLASH_NAMESPACE::copy(gmem_tiled_copy_QKV, tKgK, tKsK, tKVcKV, tKVpKV); // This cp_async_fence needs to be in the if block, otherwise the synchronization @@ -859,9 +880,19 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons } // We have key_padding_mask so we'll need to Check_inf - masking_step == 0 - ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) - : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + if (block_table != nullptr) { + masking_step == 0 + ? softmax.template softmax_rescale_o( + acc_s, acc_o, params.scale_softmax_log2) + : softmax.template softmax_rescale_o( + acc_s, acc_o, params.scale_softmax_log2); + } else { + masking_step == 0 + ? softmax.template softmax_rescale_o( + acc_s, acc_o, params.scale_softmax_log2) + : softmax.template softmax_rescale_o( + acc_s, acc_o, params.scale_softmax_log2); + } // if (cute::thread0()) { print(scores_max); print(scores_sum); print(scores); } // Convert acc_s from fp32 to fp16/bf16 @@ -893,7 +924,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons const int block_table_offset_cur = (n_block + 1) * kBlockN - block_table_idx_cur * params.page_block_size; const int block_table_idx_next = n_block * kBlockN / params.page_block_size; const int block_table_offset_next = n_block * kBlockN - block_table_idx_next * params.page_block_size; - tVgV.data() = tVgV.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.v_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.v_row_stride; + tVgV.data() = tVgV.data() + (page_id(block_table_idx_next) - page_id(block_table_idx_cur)) * params.v_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.v_row_stride; } FLASH_NAMESPACE::copy(gmem_tiled_copy_QKV, tVgV, tVsV, tKVcKV, tKVpKV); cute::cp_async_fence(); @@ -907,6 +938,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons FLASH_NAMESPACE::cp_async_wait<0>(); __syncthreads(); + clear_unmapped_v(n_block); if (n_block > n_block_min) { // Advance gK if (block_table == nullptr) { @@ -916,7 +948,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons const int block_table_offset_cur = n_block * kBlockN - block_table_idx_cur * params.page_block_size; const int block_table_idx_next = (n_block - 1) * kBlockN / params.page_block_size; const int block_table_offset_next = (n_block - 1) * kBlockN - block_table_idx_next * params.page_block_size; - tKgK.data() = tKgK.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.k_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.k_row_stride; + tKgK.data() = tKgK.data() + (page_id(block_table_idx_next) - page_id(block_table_idx_cur)) * params.k_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.k_row_stride; } FLASH_NAMESPACE::copy(gmem_tiled_copy_QKV, tKgK, tKsK, tKVcKV, tKVpKV); // This cp_async_fence needs to be in the if block, otherwise the synchronization @@ -926,7 +958,14 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons mask.template apply_mask( acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, kNWarps * 16); - softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + mask_unmapped_page(acc_s, n_block); + if (block_table != nullptr) { + softmax.template softmax_rescale_o( + acc_s, acc_o, params.scale_softmax_log2); + } else { + softmax.template softmax_rescale_o( + acc_s, acc_o, params.scale_softmax_log2); + } Tensor rP = FLASH_NAMESPACE::convert_type(acc_s); // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) 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 68ad24bbf7f71..27e67b93d7a71 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 @@ -410,7 +410,8 @@ std::optional BuildBounds( std::optional GetGroupQueryAttentionWorkspaceEstimateConfig( const Node& node, bool head_sink_is_constant_initializer, - int64_t max_total_sequence_length) { + int64_t max_total_sequence_length, + KernelDispatchPolicy dispatch_policy) { if (node.OpType() != "GroupQueryAttention") return std::nullopt; GQAWorkspaceEstimateConfig config; bool found_heads = false; @@ -481,6 +482,7 @@ std::optional GetGroupQueryAttentionWorkspaceEstimat ParseEnvironmentVariableWithDefault("ORT_DISABLE_FLASH_DECODE", false); config.head_sink_may_be_prepacked = head_sink_is_constant_initializer; config.max_total_sequence_length = max_total_sequence_length; + config.dispatch_policy = dispatch_policy; return config; } @@ -503,9 +505,11 @@ std::optional EstimateGroupQueryAttentionWorkspace( const cudaDeviceProp& device_prop, const AttentionKernelOptions& kernel_options, bool head_sink_is_constant_initializer, - int64_t max_total_sequence_length) { - const auto config = GetGroupQueryAttentionWorkspaceEstimateConfig( - node, head_sink_is_constant_initializer, max_total_sequence_length); + int64_t max_total_sequence_length, + KernelDispatchPolicy dispatch_policy) { + auto config = GetGroupQueryAttentionWorkspaceEstimateConfig( + node, head_sink_is_constant_initializer, max_total_sequence_length, + dispatch_policy); if (!config.has_value()) return std::nullopt; return EstimateGroupQueryAttentionWorkspace( *config, input_shapes, device_prop, kernel_options); 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 7af799f27694e..05de1e1f11471 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 @@ -13,6 +13,7 @@ #include #include "core/common/inlined_containers.h" +#include "core/framework/kernel_dispatch_policy.h" #include "core/framework/level1_memory_estimate.h" #include "core/framework/workspace_input_shape.h" #include "core/framework/workspace_requirement.h" @@ -51,12 +52,15 @@ struct GQAWorkspaceEstimateConfig { // 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; + // Parsed intent only; all values currently leave sizing and dispatch unchanged. + KernelDispatchPolicy dispatch_policy = KernelDispatchPolicy::Auto; }; std::optional GetGroupQueryAttentionWorkspaceEstimateConfig( const Node& node, bool head_sink_is_constant_initializer = false, - int64_t max_total_sequence_length = 0); + int64_t max_total_sequence_length = 0, + KernelDispatchPolicy dispatch_policy = KernelDispatchPolicy::Auto); std::optional EstimateGroupQueryAttentionWorkspace( const GQAWorkspaceEstimateConfig& config, @@ -70,7 +74,8 @@ std::optional EstimateGroupQueryAttentionWorkspace( const cudaDeviceProp& device_prop, const AttentionKernelOptions& kernel_options, bool head_sink_is_constant_initializer = false, - int64_t max_total_sequence_length = 0); + int64_t max_total_sequence_length = 0, + KernelDispatchPolicy dispatch_policy = KernelDispatchPolicy::Auto); void SetGroupQueryAttentionWorkspaceRequirements( const GQAWorkspaceAggregate& estimate, diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index 0ed2c125405e3..9ce0807420e27 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -69,31 +69,49 @@ __device__ inline nv_bfloat162 AddHalf2(const nv_bfloat162 a, const nv_bfloat162 return __hadd2(a, b); } -struct KeyValuePairSum { - __device__ inline cub::KeyValuePair operator()(const cub::KeyValuePair& a, - const cub::KeyValuePair& b) { - return cub::KeyValuePair(a.key + b.key, a.value + b.value); +struct WelfordData { + float mean; + float m2; + int count; + + __device__ inline void Update(float value) { + ++count; + const float delta = value - mean; + mean += delta / static_cast(count); + m2 += delta * (value - mean); + } +}; + +struct WelfordCombine { + __device__ inline WelfordData operator()(const WelfordData& a, const WelfordData& b) const { + if (a.count == 0) return b; + if (b.count == 0) return a; + + const int count = a.count + b.count; + const float delta = b.mean - a.mean; + const float b_weight = static_cast(b.count) / static_cast(count); + return {a.mean + delta * b_weight, + a.m2 + b.m2 + delta * delta * static_cast(a.count) * b_weight, + count}; } }; template __device__ inline void LayerNorm( - const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, + const WelfordData& thread_data, const float origin, const int ld, const int offset, const T* beta, const T* gamma, const float epsilon, T* output) { - // Assuming thread_data is already divided by ld - // Uses fp32 accumulation for mean/variance to avoid overflow in fp16/bf16. + // Statistics and centering stay relative to origin to retain small FP32 differences. - using BlockReduce = cub::BlockReduce, TPB>; + using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage temp_storage; __shared__ float mu; // mean __shared__ float rsigma; // 1 / std.dev. - KeyValuePairSum pair_sum; - const auto sum_kv = BlockReduce(temp_storage).Reduce(thread_data, pair_sum); + const auto stats = BlockReduce(temp_storage).Reduce(thread_data, WelfordCombine{}); if (threadIdx.x == 0) { - mu = sum_kv.key; - rsigma = rsqrtf(sum_kv.value - mu * mu + epsilon); + mu = stats.mean; + rsigma = rsqrtf(stats.m2 / static_cast(stats.count) + epsilon); } __syncthreads(); @@ -102,7 +120,7 @@ __device__ inline void LayerNorm( const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); const float b = (nullptr == beta) ? 0.f : static_cast(beta[i]); - output[idx] = static_cast(g * (val - mu) * rsigma + b); + output[idx] = static_cast(g * ((val - origin) - mu) * rsigma + b); } } @@ -132,15 +150,12 @@ __device__ inline void SimplifiedLayerNorm( } template -__device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair& thread_data, +__device__ inline void LayerNormSmall(const T* input_v, const WelfordData& thread_data, const float origin, const int ld, const int idx, const T* beta, const T* gamma, const float epsilon, T* output) { - // Assuming thread_data is already divided by ld - // Small settings: the block covers the leading dimension TPB >= ld. The input - // value is available in a register - // Uses fp32 accumulation for mean/variance to avoid overflow in fp16/bf16. + // Input values are in registers; statistics and centering are relative to origin. using VecT = aligned_vector; - using BlockReduce = cub::BlockReduce, TPB>; + using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage temp_storage; __shared__ float mu; // mean __shared__ float rsigma; // 1 / std.dev. @@ -158,12 +173,11 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair< *gamma_val = *reinterpret_cast(&gamma[threadIdx.x * ILP]); } - KeyValuePairSum pair_sum; - const cub::KeyValuePair sum_kv = BlockReduce(temp_storage).Reduce(thread_data, pair_sum); + const auto stats = BlockReduce(temp_storage).Reduce(thread_data, WelfordCombine{}); if (threadIdx.x == 0) { - mu = sum_kv.key; - rsigma = rsqrtf(sum_kv.value - mu * mu + epsilon); + mu = stats.mean; + rsigma = rsqrtf(stats.m2 / static_cast(stats.count) + epsilon); } __syncthreads(); @@ -173,7 +187,7 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair< const float in_f = static_cast(input_v[i]); const float g_f = static_cast(gamma_v[i]); const float b_f = (beta != nullptr) ? static_cast(beta_v[i]) : 0.f; - output_v[i] = static_cast(g_f * (in_f - mu) * rsigma + b_f); + output_v[i] = static_cast(g_f * ((in_f - origin) - mu) * rsigma + b_f); } VecT* output_val = reinterpret_cast(&output_v); diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc b/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc index 54aee560c52f9..0433fd9959858 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc @@ -267,6 +267,16 @@ Status PackedAttention::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_ERROR(PackedAttentionWorkspaceStatusToStatus(workspace_result.status)); const PackedAttentionWorkspaceRecipe& workspace_recipe = workspace_result.recipe; + if (problem.backend == PackedAttentionBackend::Unfused) { + auto validation_flag = this->template GetScratchBuffer( + 1, this->GetComputeStream(context)); + ORT_RETURN_IF_ERROR(ValidatePackedAttentionTokenOffset( + token_offset->Data(), + parameters.batch_size * parameters.sequence_length, + validation_flag.get(), + this->Stream(context))); + } + auto gemm_buffer = this->template GetScratchBuffer( workspace_recipe.projection_bytes, this->GetComputeStream(context)); const int m = workspace_recipe.projection_m; diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu index ce05de4c9e45e..297f8b2febaaa 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu @@ -28,6 +28,60 @@ namespace cuda { constexpr int32_t kMAX_THREADS_PER_BLOCK = 256; +__global__ void ValidatePackedAttentionTokenOffsetKernel( + const int32_t* token_offset, + int32_t token_offset_count, + int32_t* validation_flag) { + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + i < token_offset_count; + i += static_cast(gridDim.x) * blockDim.x) { + const int32_t offset = token_offset[i]; + if (offset < 0 || offset >= token_offset_count) { + atomicExch(validation_flag, 1); + return; + } + } +} + +Status ValidatePackedAttentionTokenOffset( + const int32_t* token_offset, + int32_t token_offset_count, + int32_t* validation_flag, + cudaStream_t stream) { + if (token_offset_count == 0) { + return Status::OK(); + } + + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(stream, &capture_status)); + if (capture_status != cudaStreamCaptureStatusNone) { + return Status::OK(); + } + + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(validation_flag, 0, sizeof(int32_t), stream)); + constexpr int32_t threads_per_block = 256; + const int32_t required_blocks = + static_cast((static_cast(token_offset_count) + threads_per_block - 1) / + threads_per_block); + const int32_t blocks = std::min(required_blocks, 1024); + ValidatePackedAttentionTokenOffsetKernel<<>>( + token_offset, token_offset_count, validation_flag); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + + int32_t host_validation_flag = 0; + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync( + &host_validation_flag, validation_flag, sizeof(int32_t), + cudaMemcpyDeviceToHost, stream)); + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); + if (host_validation_flag != 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "PackedAttention token_offset values must be in [0, B * S)."); + } + + return Status::OK(); +} + // Grid: (S, B) // Block: 256 // For unfused PackedAttention @@ -58,6 +112,9 @@ __global__ void AddBiasTransposeQKVPacked( const int packing_token_idx = b * S + s; const int padding_token_idx = token_offset[packing_token_idx]; + if (padding_token_idx < 0 || padding_token_idx >= gridDim.x * gridDim.y) { + return; + } b = padding_token_idx / S; s = padding_token_idx - b * S; @@ -311,6 +368,10 @@ __global__ void __launch_bounds__(kMAX_THREADS_PER_BLOCK) int token_idx = blockIdx.x; int source_idx = token_offset[token_idx]; if (source_idx < 0 || source_idx >= B * S) { + target += token_idx * N * H; + for (int i = threadIdx.x; i < N * H; i += blockDim.x) { + target[i] = T{}; + } return; } diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h index c1224bf9ef57a..958f0425481fb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h @@ -13,6 +13,12 @@ namespace onnxruntime { namespace contrib { namespace cuda { +Status ValidatePackedAttentionTokenOffset( + const int32_t* token_offset, + int32_t token_offset_count, + int32_t* validation_flag, + cudaStream_t stream); + template T* PackedAttentionWorkspaceAt(T* workspace, size_t offset_bytes) { if (offset_bytes == 0) { diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc b/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc index 2e2542b948a3b..64746d835b270 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc @@ -6,6 +6,7 @@ #include #include +#include "core/common/safeint.h" #include "core/providers/cuda/cuda_common.h" #include "core/platform/env_var_utils.h" #include "contrib_ops/cpu/utils/dump_tensor.h" @@ -168,6 +169,7 @@ PagedAttention::PagedAttention(const OpKernelInfo& info) template Status PagedAttention::ComputeInternal(OpKernelContext* context) const { auto ort_stream = GetOrtStream(context); + cudaStream_t cuda_stream = static_cast(ort_stream.get()->GetHandle()); const Tensor* query = context->Input(0); const Tensor* key = context->Input(1); @@ -299,6 +301,32 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons if (parameters.token_count == 0) { return Status::OK(); } + if (parameters.num_blocks <= 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "PagedAttention requires at least one physical cache block for non-empty input."); + } + + const SafeInt safe_batch_size(parameters.batch_size); + const size_t block_table_element_count = + safe_batch_size * SafeInt(parameters.max_num_blocks_per_seq); + if (block_table_element_count > static_cast(std::numeric_limits::max())) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "block_table element count exceeds the CUDA kernel indexing limit."); + } + SafeInt safe_max_kv_len_capacity(parameters.max_num_blocks_per_seq); + safe_max_kv_len_capacity *= parameters.block_size; + if (safe_max_kv_len_capacity > std::numeric_limits::max()) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "block_table sequence capacity exceeds the CUDA int32 indexing limit."); + } + const int max_kv_len_capacity = static_cast(safe_max_kv_len_capacity); + auto sanitized_block_table = GetScratchBuffer(block_table_element_count, GetComputeStream(context)); + ORT_RETURN_IF_ERROR(LaunchSanitizeBlockTable( + reinterpret_cast(block_table->Data()), + sanitized_block_table.get(), block_table_element_count, parameters.num_blocks, + cuda_stream)); // Kernel backend selection. The choice depends only on static shapes and on the optional // 'attention_metadata' bounds, never on a device-to-host readback, so it is identical on every @@ -372,9 +400,26 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons parameters.token_count <= device_prop.maxGridSize[1] && GetPagedDecodeSharedMemoryBytes(parameters.head_size) <= static_cast(device_prop.sharedMemPerBlock); - size_t cumulative_seqlens_kv_bytes = sizeof(int) * (parameters.batch_size + 1); + const size_t cumulative_sequence_length_count = safe_batch_size + 1; + const size_t sanitized_sequence_length_count = SafeInt(2) * safe_batch_size + 1; + if (sanitized_sequence_length_count > static_cast(std::numeric_limits::max())) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "sequence length element count exceeds the CUDA kernel indexing limit."); + } + const size_t cumulative_seqlens_kv_bytes = + SafeInt(sizeof(int)) * cumulative_sequence_length_count; auto cumulative_seqlens_kv_buffer = GetScratchBuffer(cumulative_seqlens_kv_bytes, GetComputeStream(context)); int* cumulative_seqlens_kv_ptr = reinterpret_cast(cumulative_seqlens_kv_buffer.get()); + auto sanitized_sequence_lengths = + GetScratchBuffer(sanitized_sequence_length_count, GetComputeStream(context)); + int* sanitized_cumulative_seqlens_q = sanitized_sequence_lengths.get(); + int* sanitized_past_seqlens = sanitized_cumulative_seqlens_q + cumulative_sequence_length_count; + auto sequence_validity = GetScratchBuffer(parameters.batch_size, GetComputeStream(context)); + size_t sequence_sanitizer_workspace_bytes = 0; + ORT_RETURN_IF_ERROR(GetSanitizeSequenceLengthsWorkspaceSize( + parameters.batch_size, sequence_sanitizer_workspace_bytes, cuda_stream)); + auto sequence_sanitizer_workspace = + GetScratchBuffer(sequence_sanitizer_workspace_bytes, GetComputeStream(context)); // The fused prologue (QK-Norm and/or rotary) writes densified Q and K into the workspace, so it // needs room for both. Plain packed-QKV only needs to densify Q. @@ -387,15 +432,23 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons } auto workspace_buffer = GetScratchBuffer(workspace_buffer_bytes, GetComputeStream(context)); - // Populate cumulative_seqlens_kv for all backends. Every kernel that needs a per-sequence KV - // length reads it from here on device; the host only ever uses upper bounds. - cudaStream_t cuda_stream = static_cast(ort_stream.get()->GetHandle()); - ORT_RETURN_IF_ERROR(LaunchGetCumulativeSeqlensKV( + // Canonicalize untrusted sequence metadata and populate cumulative_seqlens_kv for every backend. + // This stays entirely on the compute stream, so CUDA Graph replay sees current values without a + // host synchronization. + ORT_RETURN_IF_ERROR(LaunchSanitizeSequenceLengths( + sanitized_cumulative_seqlens_q, + sanitized_past_seqlens, + sequence_validity.get(), cumulative_seqlens_kv_ptr, reinterpret_cast(cumulative_seqlens_q->Data()), reinterpret_cast(past_seqlens->Data()), - parameters.batch_size, cuda_stream)); - + sequence_sanitizer_workspace.get(), + sequence_sanitizer_workspace_bytes, + parameters.batch_size, + parameters.max_num_blocks_per_seq, + parameters.block_size, + parameters.token_count, + cuda_stream)); int total_kv_tokens = 0; int max_query_len = 0; int max_kv_len = 0; @@ -414,7 +467,6 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons // query tokens can belong to a single sequence, and a sequence can address at most // block_table.shape[1] * block_size cached tokens. Those are valid but loose, so they cost some // empty thread blocks rather than correctness. - const int max_kv_len_capacity = parameters.max_num_blocks_per_seq * parameters.block_size; const bool has_metadata_bounds = attention_metadata != nullptr; int max_query_len_bound = parameters.token_count; int max_kv_len_bound = max_kv_len_capacity; @@ -520,15 +572,6 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons const bool portable_spec_dec_candidate = has_metadata_bounds && max_query_len_bound > 1 && max_query_len_bound <= 8 && (std::is_same_v || (kIsQuantizedCache && per_channel_k && !enable_per_channel_xqa_)); - // cuDNN paged SDPA (decode-only, unquantized cache). Preferred over FlashAttention when eligible; - // XQA still wins its target case (fp16, group_size 6, head_size 256, native page size). The - // eligibility is intentionally metadata-gated so the selection never triggers a new D->H readback - // -- max_query_len_bound == 1 with has_metadata_bounds is the same signal XQA uses. The - // xqa_spec_dec_candidate / portable_spec_dec_candidate paths are orthogonal because they require - // max_query_len_bound > 1. - // Enablement mirrors GroupQueryAttention's cuDNN tier: explicit opt-in via the sdpa_kernel bit / - // ORT_ENABLE_CUDNN_FLASH_ATTENTION=1, plus auto-on for sm>=90. The env var set to 0 kills every - // cuDNN attention path uniformly. const bool cudnn_paged_enabled = enable_cudnn_paged_ || (auto_enable_cudnn_paged_ && device_prop.major >= 9); const bool cudnn_paged_eligible = @@ -550,21 +593,8 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons /*sequence_length_q=*/1, parameters.block_size); bool use_cudnn_paged = cudnn_paged_eligible && !fp16_xqa_eligible; - - // Pre-dispatch buildability probe for cuDNN paged. is_supported_paged only checks static shapes; - // the planner may still reject the compiled graph, and we cannot fall back once dispatch has - // committed. The graph cache in run_paged is thread_local and keyed on the full PagedGraphParams - // (shape + handle), so a node-wide scalar latch would be wrong on two axes: (a) after a - // successful probe for one shape, a later Compute with a different batch_size / - // max_num_blocks_per_seq / num_blocks would skip the check and hit an unbuildable dispatch, and - // (b) a second worker thread has its own empty thread_local cache. So probe the *current* shape - // and handle every Run using the same key run_paged will use. try_build_paged_graph is a - // cache-first read (returns true on hit without touching cuDNN) and folds isCapturing() - // internally: on a cache miss during graph capture it returns false rather than attempting a - // non-capturable build. A false result clears use_cudnn_paged for this Run only, so the cascade - // drops to FlashAttention / MemoryEfficientAttention. if (use_cudnn_paged) { - const bool ok = onnxruntime::cudnn_sdpa::try_build_paged_graph( + use_cudnn_paged = onnxruntime::cudnn_sdpa::try_build_paged_graph( parameters.batch_size, parameters.num_heads, parameters.kv_num_heads, parameters.head_size, parameters.head_size, @@ -575,9 +605,6 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons std::is_same::value, GetCudnnHandle(context), ort_stream.get()); - if (!ok) { - use_cudnn_paged = false; - } } // cuDNN paged SDPA is causal-only. FlashAttention, paged decode, and CUTLASS all consume the @@ -591,9 +618,10 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons const bool use_memory_efficient_attention = mea_eligible && !use_paged_decode && !use_cudnn_paged; - // Both gather-based backends need a dense KV staging buffer when the cache is quantized - // (FlashAttention cannot read a quantized page, and the CUTLASS kernel is not paged at all). - const bool needs_dense_kv = use_memory_efficient_attention || (use_flash_attention && kIsQuantizedCache); + // Quantized Flash and CUTLASS need dense staging. Native Flash reads the paged cache directly; + // its page loader masks the documented -1 sentinel without materializing capacity-sized K/V. + const bool needs_dense_kv = + use_memory_efficient_attention || (use_flash_attention && kIsQuantizedCache); // The dense buffer keeps the grouped layout for FlashAttention (it does GQA internally) and is // GQA-expanded for the CUTLASS kernel. const int gathered_num_heads = use_memory_efficient_attention ? parameters.num_heads : parameters.kv_num_heads; @@ -629,13 +657,13 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons // a large over-allocation for a short prefill. Read the exact value back instead. // 2. XQA needs the one-token-per-sequence proof described above. // - // Neither case can occur on a capturable step: a captured step is decode-shaped on a paged cache - // (so no gather runs), and a producer that captures must supply 'attention_metadata' anyway -- - // its bounds are the only replay-safe source of per-step information. The synchronization is - // therefore gone for every configuration CUDA Graphs can reach, including an unquantized cache. + // During graph capture, static capacity bounds replace an otherwise-required readback. Optional + // attention_metadata can tighten those bounds, but correctness and capture never require it. const bool needs_readback = !has_metadata_bounds && (needs_dense_kv || xqa_candidate); + const bool perform_readback = + needs_readback && !onnxruntime::llm::common::isCapturing(cuda_stream); - if (!needs_readback) { + if (!perform_readback) { max_query_len = max_query_len_bound; max_kv_len = max_kv_len_bound; // Upper bound: no sequence holds more than max_kv_len_bound cached tokens. Read only by the @@ -648,12 +676,36 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons const int kCumulativeCount = parameters.batch_size + 1; auto cum_q_pinned = this->AllocateBufferOnCPUPinned(kCumulativeCount); auto cum_kv_pinned = this->AllocateBufferOnCPUPinned(kCumulativeCount); + auto original_cum_q_pinned = this->AllocateBufferOnCPUPinned(kCumulativeCount); + auto original_past_pinned = this->AllocateBufferOnCPUPinned(parameters.batch_size); CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cum_q_pinned.get(), - reinterpret_cast(cumulative_seqlens_q->Data()), + sanitized_cumulative_seqlens_q, sizeof(int) * kCumulativeCount, cudaMemcpyDeviceToHost, cuda_stream)); CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cum_kv_pinned.get(), cumulative_seqlens_kv_ptr, sizeof(int) * kCumulativeCount, cudaMemcpyDeviceToHost, cuda_stream)); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync( + original_cum_q_pinned.get(), cumulative_seqlens_q->Data(), + sizeof(int) * kCumulativeCount, cudaMemcpyDeviceToHost, cuda_stream)); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync( + original_past_pinned.get(), past_seqlens->Data(), + sizeof(int) * parameters.batch_size, cudaMemcpyDeviceToHost, cuda_stream)); CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream)); + + ORT_RETURN_IF_ERROR(paged_attention_helper::CheckSequenceLengthAndPastValues( + original_cum_q_pinned.get(), + original_past_pinned.get(), + parameters.batch_size, + parameters.max_num_blocks_per_seq, + parameters.block_size, + parameters.token_count)); + ORT_RETURN_IF_ERROR(paged_attention_helper::CheckSequenceLengthValues( + cum_q_pinned.get(), + cum_kv_pinned.get(), + parameters.batch_size, + parameters.max_num_blocks_per_seq, + parameters.block_size, + parameters.token_count)); + for (int i = 0; i < parameters.batch_size; ++i) { const int q_len_i = cum_q_pinned.get()[i + 1] - cum_q_pinned.get()[i]; if (q_len_i > max_query_len) { @@ -707,37 +759,28 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons use_xqa_decode = (xqa_smem_ok != 0); use_xqa_spec_dec = use_xqa_decode && use_xqa_spec_dec; } - // XQA was statically preferred over cuDNN paged (see the `!fp16_xqa_eligible` term above), so - // when the XQA shared-memory / symbol probe fails at runtime the documented cascade drops to - // cuDNN paged before FlashAttention -- not straight to FlashAttention. Retry the shape-keyed - // buildability probe now; it is cache-first, so on steady-state Runs this is one hash lookup, - // and it only pays for a real cuDNN graph build when cuDNN is about to serve the Run. Guarding - // on `fp16_xqa_eligible` scopes the retry to the specific "XQA-preferred-but-runtime-failed" - // case, so we don't re-run a probe that already failed at line ~527 for a non-XQA path (that - // failure is a definitive per-Run decision). - if (fp16_xqa_eligible && cudnn_paged_eligible && !use_cudnn_paged && !use_xqa_decode) { - if (onnxruntime::cudnn_sdpa::try_build_paged_graph( - parameters.batch_size, - parameters.num_heads, parameters.kv_num_heads, - parameters.head_size, parameters.head_size, - parameters.num_blocks, - parameters.block_size, - parameters.max_num_blocks_per_seq, - cudnn_scale, - std::is_same::value, - GetCudnnHandle(context), - ort_stream.get())) { - use_cudnn_paged = true; + if (fp16_xqa_eligible && cudnn_paged_eligible && + !use_cudnn_paged && !use_xqa_decode) { + use_cudnn_paged = onnxruntime::cudnn_sdpa::try_build_paged_graph( + parameters.batch_size, + parameters.num_heads, parameters.kv_num_heads, + parameters.head_size, parameters.head_size, + parameters.num_blocks, + parameters.block_size, + parameters.max_num_blocks_per_seq, + cudnn_scale, + std::is_same::value, + GetCudnnHandle(context), + ort_stream.get()); + if (use_cudnn_paged) { use_paged_decode = false; use_flash_attention = false; - // MemoryEfficientAttention is ineligible here: this branch requires fp16_xqa_eligible, - // which implies a FlashAttention-supported shape, while MEA requires !flash_eligible. } } // Native-cache XQA promotion is speculative until the one-token-per-sequence and shared-memory // checks pass. Restore Flash for ragged decode steps and unsupported devices instead of leaving - // them on the portable scalar paged-decode fallback. This is safe without dense KV staging - // because the native FP16/BF16 cache is already a Flash-supported dtype. Skip when cuDNN paged + // them on the portable scalar paged-decode fallback. Dense KV staging was reserved above for + // this fallback so unmapped pages remain masked. Skip when cuDNN paged // already won this step; the two selections would otherwise coexist for the corner case where // native XQA is opted in on a bf16 cache with head_size=256 / group_size=6 and max_query_len=1 // (native_spec_xqa_eligible is a shape gate and stays true regardless of query length). @@ -905,12 +948,21 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons } #endif - // cuDNN paged SDPA needs a small [batch_size] int32 buffer holding per-sequence KV lengths, - // filled on device by LaunchGetSeqlensKVDecode right before dispatch. + // cuDNN paged SDPA needs compacted page tables and their mapped per-sequence KV lengths. IAllocatorUniquePtr cudnn_seqlens_kv_buffer; + IAllocatorUniquePtr cudnn_block_table_buffer; if (use_cudnn_paged) { cudnn_seqlens_kv_buffer = GetScratchBuffer( sizeof(int) * static_cast(parameters.batch_size), GetComputeStream(context)); + cudnn_block_table_buffer = + GetScratchBuffer(block_table_element_count, GetComputeStream(context)); + ORT_RETURN_IF_ERROR(LaunchPrepareCudnnBlockTable( + sanitized_block_table.get(), cudnn_block_table_buffer.get(), + cumulative_seqlens_kv_ptr, + reinterpret_cast(cudnn_seqlens_kv_buffer.get()), + sequence_validity.get(), block_table_element_count, + parameters.batch_size, parameters.max_num_blocks_per_seq, + parameters.block_size, cuda_stream)); } // Print debug info @@ -951,10 +1003,13 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons : reinterpret_cast(const_cast(value_cache->Data())); data.k_scale = k_scale == nullptr ? nullptr : k_scale->Data(); data.v_scale = v_scale == nullptr ? nullptr : v_scale->Data(); - data.cumulative_seqlens_q = reinterpret_cast(cumulative_seqlens_q->Data()); - data.past_seqlens = reinterpret_cast(past_seqlens->Data()); + data.cumulative_seqlens_q = sanitized_cumulative_seqlens_q; + data.past_seqlens = sanitized_past_seqlens; + data.sequence_validity = sequence_validity.get(); data.cumulative_seqlens_kv = cumulative_seqlens_kv_ptr; - data.block_table = reinterpret_cast(block_table->Data()); + data.block_table = sanitized_block_table.get(); + data.cudnn_block_table = + use_cudnn_paged ? cudnn_block_table_buffer.get() : nullptr; data.slot_mapping = slot_mapping == nullptr ? nullptr : reinterpret_cast(slot_mapping->Data()); data.head_sink = head_sink == nullptr ? nullptr : reinterpret_cast(head_sink->Data()); data.q_norm_weight = q_norm_weight == nullptr ? nullptr : reinterpret_cast(q_norm_weight->Data()); diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu index eb936079c7c51..4734c87ebfd91 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu @@ -4,6 +4,7 @@ #include #include // FLT_MAX #include +#include #include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" @@ -26,6 +27,121 @@ namespace onnxruntime { namespace contrib { namespace cuda { +__global__ void SanitizeBlockTableKernel(const int32_t* block_table, int32_t* sanitized_block_table, + size_t element_count, int num_blocks) { + const size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < element_count) { + const int32_t block_id = block_table[index]; + sanitized_block_table[index] = block_id >= -1 && block_id < num_blocks ? block_id : -1; + } +} + +Status LaunchSanitizeBlockTable(const int32_t* block_table, int32_t* sanitized_block_table, + size_t element_count, int num_blocks, cudaStream_t stream) { + if (element_count == 0) { + return Status::OK(); + } + + if (element_count > static_cast(std::numeric_limits::max())) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "block_table element count exceeds the CUDA kernel indexing limit."); + } + constexpr int kThreadsPerBlock = 256; + const size_t blocks = (element_count + kThreadsPerBlock - 1) / kThreadsPerBlock; + SanitizeBlockTableKernel<<(blocks), kThreadsPerBlock, 0, stream>>>( + block_table, sanitized_block_table, element_count, num_blocks); + return CUDA_CALL(cudaGetLastError()); +} + +__global__ void PrepareCudnnBlockTableKernel( + const int32_t* block_table, int32_t* cudnn_block_table, + const int32_t* cumulative_seqlens_kv, int32_t* cudnn_seqlens_kv, + int32_t* sequence_validity, int batch_size, + int max_num_blocks_per_seq, int block_size) { + for (int batch = blockIdx.x; batch < batch_size; batch += gridDim.x) { + const int row_offset = batch * max_num_blocks_per_seq; + if (threadIdx.x == 0) { + const int sequence_length = + cumulative_seqlens_kv[batch + 1] - cumulative_seqlens_kv[batch]; + const int live_blocks = (sequence_length + block_size - 1) / block_size; + const int last_block_tokens = sequence_length % block_size; + int mapped_blocks = 0; + int mapped_length = 0; + for (int block = 0; block < live_blocks; ++block) { + const int block_id = block_table[row_offset + block]; + if (block_id >= 0) { + cudnn_block_table[row_offset + mapped_blocks++] = block_id; + mapped_length += + block + 1 == live_blocks && last_block_tokens != 0 + ? last_block_tokens + : block_size; + } + } + for (int block = mapped_blocks; block < max_num_blocks_per_seq; ++block) { + cudnn_block_table[row_offset + block] = 0; + } + cudnn_seqlens_kv[batch] = mapped_length; + if (mapped_length == 0) { + sequence_validity[batch] = 0; + } + } + } +} + +Status LaunchPrepareCudnnBlockTable( + const int32_t* block_table, int32_t* cudnn_block_table, + const int32_t* cumulative_seqlens_kv, int32_t* cudnn_seqlens_kv, + int32_t* sequence_validity, size_t element_count, int batch_size, + int max_num_blocks_per_seq, int block_size, + cudaStream_t stream) { + ORT_RETURN_IF_NOT( + element_count <= static_cast(std::numeric_limits::max()), + "block_table element count exceeds the CUDA kernel indexing limit."); + if (element_count == 0) { + return Status::OK(); + } + constexpr int kMaxBlocks = 65535; + const int blocks = std::min(batch_size, kMaxBlocks); + PrepareCudnnBlockTableKernel<<>>( + block_table, cudnn_block_table, cumulative_seqlens_kv, + cudnn_seqlens_kv, sequence_validity, batch_size, + max_num_blocks_per_seq, block_size); + return CUDA_CALL(cudaGetLastError()); +} + +template +__global__ void MaskInvalidSequenceOutputs( + T* output, const int32_t* sequence_validity, + const int32_t* cumulative_seqlens_q, int batch_size, int v_hidden_size) { + for (int batch = blockIdx.x; batch < batch_size; batch += gridDim.x) { + if (sequence_validity[batch] != 0) { + continue; + } + + const int64_t query_start = cumulative_seqlens_q[batch]; + const int64_t query_length = cumulative_seqlens_q[batch + 1] - query_start; + const int64_t element_count = query_length * v_hidden_size; + for (int64_t offset = threadIdx.x; offset < element_count; offset += blockDim.x) { + output[query_start * v_hidden_size + offset] = static_cast(0.0f); + } + } +} + +template +Status LaunchMaskInvalidSequenceOutputs( + T* output, const int32_t* sequence_validity, + const int32_t* cumulative_seqlens_q, int batch_size, + int v_hidden_size, cudaStream_t stream) { + constexpr int kThreadsPerBlock = 256; + constexpr int kMaxBlocks = 65535; + if (batch_size > 0 && v_hidden_size > 0) { + const int blocks = std::min(batch_size, kMaxBlocks); + MaskInvalidSequenceOutputs<<>>( + output, sequence_validity, cumulative_seqlens_q, batch_size, v_hidden_size); + } + return CUDA_CALL(cudaGetLastError()); +} + ////////// Quantized paged KV cache helpers // // Symmetric, zero-point-free quantization with the same numerics GroupQueryAttention uses @@ -337,76 +453,136 @@ Status LaunchQkNormRotaryKernel(cudaStream_t stream, T* output, const T* input, return CUDA_CALL(cudaGetLastError()); } -// Single-block inclusive scan over the per-sequence KV lengths. One block loops over the batch in -// kBlockSize-sized tiles carrying a running total, so there is no cap on batch_size (the previous -// implementation launched independent blocks whose cub::BlockScan did not compose, which silently -// produced wrong offsets past 256 concurrent sequences). -template -__global__ void GetCumulativeSeqlensKV(int32_t* cumulative_seqlens_kv, const int32_t* cumulative_seqlens_q, - const int32_t* past_seqlens, const int batch_size) { - typedef cub::BlockScan BlockScan; - __shared__ typename BlockScan::TempStorage temp_storage; - __shared__ int running_total; - - if (threadIdx.x == 0) { - cumulative_seqlens_kv[0] = 0; - running_total = 0; +struct MaxInt32 { + __host__ __device__ __forceinline__ int32_t operator()(int32_t lhs, int32_t rhs) const { + return lhs > rhs ? lhs : rhs; } - __syncthreads(); +}; - for (int base = 0; base < batch_size; base += kBlockSize) { - const int id = base + static_cast(threadIdx.x); - // Sum past_seqlens to the new sequence length (which we get by subtracting cumulative_seqlens_q), - // then inclusive-scan across present sequence lengths. - const int length = (id < batch_size) - ? past_seqlens[id] + cumulative_seqlens_q[id + 1] - cumulative_seqlens_q[id] - : 0; - int prefix = 0; - int aggregate = 0; - BlockScan(temp_storage).InclusiveSum(length, prefix, aggregate); - if (id < batch_size) { - cumulative_seqlens_kv[id + 1] = running_total + prefix; - } - __syncthreads(); // all reads of running_total and of temp_storage are done - if (threadIdx.x == 0) { - running_total += aggregate; - } - __syncthreads(); // running_total visible, temp_storage safe to reuse +struct SaturatingAddInt32 { + __host__ __device__ __forceinline__ int32_t operator()(int32_t lhs, int32_t rhs) const { + const int64_t sum = static_cast(lhs) + rhs; + return sum < INT32_MAX ? static_cast(sum) : INT32_MAX; } -} +}; -Status LaunchGetCumulativeSeqlensKV(int32_t* cumulative_seqlens_kv, const int32_t* cumulative_seqlens_q, - const int32_t* past_seqlens, const int batch_size, cudaStream_t stream) { - constexpr int kThreads = 256; - GetCumulativeSeqlensKV<<<1, kThreads, 0, stream>>>(cumulative_seqlens_kv, cumulative_seqlens_q, - past_seqlens, batch_size); - return CUDA_CALL(cudaGetLastError()); +__global__ void ClampCumulativeSequenceLengths(int32_t* sanitized_cumulative_seqlens_q, + const int32_t* cumulative_seqlens_q, + int batch_size, + int token_count) { + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index > batch_size) { + return; + } + + if (index == 0) { + sanitized_cumulative_seqlens_q[index] = 0; + } else if (index == batch_size) { + sanitized_cumulative_seqlens_q[index] = token_count; + } else { + sanitized_cumulative_seqlens_q[index] = min(token_count, max(0, cumulative_seqlens_q[index])); + } } -// Fills seqlens_kv[i] = past_seqlens[i] + (cumulative_seqlens_q[i+1] - cumulative_seqlens_q[i]) -// for the cuDNN paged SDPA graph's per-batch KV padding-mask input. Deriving the query count from -// cumulative_seqlens_q rather than hard-coding +1 avoids baking the "decode-only" invariant into -// the kernel: on every currently reachable call the difference is 1 (the cuDNN paged tier is -// gated to max_query_len_bound == 1), and any future extension to multi-token steps stays -// numerically correct here without a second edit. -__global__ void GetSeqlensKV(int32_t* seqlens_kv, const int32_t* past_seqlens, - const int32_t* cumulative_seqlens_q, - const int batch_size) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < batch_size) { - const int q_len = cumulative_seqlens_q[i + 1] - cumulative_seqlens_q[i]; - seqlens_kv[i] = past_seqlens[i] + q_len; +__global__ void SanitizePastSequenceLengths(int32_t* sanitized_past_seqlens, + int32_t* sequence_validity, + int32_t* cumulative_seqlens_kv, + const int32_t* sanitized_cumulative_seqlens_q, + const int32_t* past_seqlens, + int batch_size, + int max_num_blocks_per_seq, + int block_size) { + const int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b < batch_size) { + const int32_t query_length = + sanitized_cumulative_seqlens_q[b + 1] - sanitized_cumulative_seqlens_q[b]; + const int64_t mapped_capacity = static_cast(max_num_blocks_per_seq) * block_size; + const int64_t available_past_capacity = mapped_capacity - query_length; + const int64_t max_past_length = available_past_capacity > 0 ? available_past_capacity : 0; + const int64_t input_past_length = past_seqlens[b]; + if (query_length > mapped_capacity || + input_past_length < 0 || input_past_length > max_past_length) { + sanitized_past_seqlens[b] = 0; + sequence_validity[b] = 0; + cumulative_seqlens_kv[b + 1] = 0; + if (b == 0) { + cumulative_seqlens_kv[0] = 0; + } + return; + } + sequence_validity[b] = 1; + const int64_t bounded_past_length = + input_past_length; + sanitized_past_seqlens[b] = static_cast(bounded_past_length); + + const int64_t uncapped_kv_length = bounded_past_length + query_length; + const int64_t kv_length = uncapped_kv_length < mapped_capacity ? uncapped_kv_length : mapped_capacity; + cumulative_seqlens_kv[b + 1] = + static_cast(kv_length < INT32_MAX ? kv_length : INT32_MAX); + + if (b == 0) { + cumulative_seqlens_kv[0] = 0; + } } } -Status LaunchGetSeqlensKV(int32_t* seqlens_kv, const int32_t* past_seqlens, - const int32_t* cumulative_seqlens_q, - const int batch_size, cudaStream_t stream) { - constexpr int kThreads = 128; - const int blocks = (batch_size + kThreads - 1) / kThreads; - GetSeqlensKV<<>>(seqlens_kv, past_seqlens, cumulative_seqlens_q, - batch_size); - return CUDA_CALL(cudaGetLastError()); +Status GetSanitizeSequenceLengthsWorkspaceSize(int batch_size, + size_t& workspace_bytes, + cudaStream_t stream) { + if (batch_size <= 0 || batch_size > (INT32_MAX - 1) / 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "batch_size exceeds the sequence sanitizer indexing limit."); + } + const int element_count = batch_size + 1; + int32_t* placeholder = nullptr; + size_t max_scan_bytes = 0; + size_t sum_scan_bytes = 0; + CUDA_RETURN_IF_ERROR(cub::DeviceScan::InclusiveScan( + nullptr, max_scan_bytes, placeholder, placeholder, MaxInt32{}, element_count, stream)); + CUDA_RETURN_IF_ERROR(cub::DeviceScan::InclusiveScan( + nullptr, sum_scan_bytes, placeholder, placeholder, SaturatingAddInt32{}, element_count, stream)); + workspace_bytes = std::max(max_scan_bytes, sum_scan_bytes); + return Status::OK(); +} + +Status LaunchSanitizeSequenceLengths(int32_t* sanitized_cumulative_seqlens_q, + int32_t* sanitized_past_seqlens, + int32_t* sequence_validity, + int32_t* cumulative_seqlens_kv, + const int32_t* cumulative_seqlens_q, + const int32_t* past_seqlens, + void* workspace, + size_t workspace_bytes, + int batch_size, + int max_num_blocks_per_seq, + int block_size, + int token_count, + cudaStream_t stream) { + if (batch_size <= 0 || batch_size > (INT32_MAX - 1) / 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "batch_size exceeds the sequence sanitizer indexing limit."); + } + constexpr int kThreadsPerBlock = 256; + const int cumulative_element_count = batch_size + 1; + const int cumulative_blocks = (cumulative_element_count + kThreadsPerBlock - 1) / kThreadsPerBlock; + ClampCumulativeSequenceLengths<<>>( + sanitized_cumulative_seqlens_q, cumulative_seqlens_q, batch_size, token_count); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + CUDA_RETURN_IF_ERROR(cub::DeviceScan::InclusiveScan( + workspace, workspace_bytes, + sanitized_cumulative_seqlens_q, sanitized_cumulative_seqlens_q, + MaxInt32{}, cumulative_element_count, stream)); + + const int batch_blocks = (batch_size + kThreadsPerBlock - 1) / kThreadsPerBlock; + SanitizePastSequenceLengths<<>>( + sanitized_past_seqlens, sequence_validity, cumulative_seqlens_kv, sanitized_cumulative_seqlens_q, + past_seqlens, batch_size, max_num_blocks_per_seq, block_size); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + CUDA_RETURN_IF_ERROR(cub::DeviceScan::InclusiveScan( + workspace, workspace_bytes, + cumulative_seqlens_kv, cumulative_seqlens_kv, + SaturatingAddInt32{}, cumulative_element_count, stream)); + return Status::OK(); } // Resolves the flat cache slot that a query token's K/V is written to, in the cache viewed as @@ -416,34 +592,41 @@ Status LaunchGetSeqlensKV(int32_t* seqlens_kv, const int32_t* past_seqlens, // past_seqlens[b] + (token_id - cumulative_seqlens_q[b]) of its own sequence. The binary search is // guarded against token_id >= cumulative_seqlens_q[batch_size], which previously walked off the end // of past_seqlens / block_table. +__device__ __forceinline__ int FindTokenBatch( + int token_id, const int* cumulative_seqlens_q, int batch_size) { + if (token_id < 0 || token_id >= cumulative_seqlens_q[batch_size]) { + return -1; + } + int left = 0; + int right = batch_size - 1; + while (left < right) { + const int mid = left + (right - left) / 2; + if (token_id < cumulative_seqlens_q[mid + 1]) { + right = mid; + } else { + left = mid + 1; + } + } + return left; +} + struct DerivedSlotResolver { const int* __restrict__ block_table; const int* __restrict__ past_seqlens; const int* __restrict__ cumulative_seqlens_q; + const int* __restrict__ sequence_validity; int batch_size; int max_num_blocks_per_seq; int block_size; __device__ __forceinline__ int operator()(int token_id) const { - if (token_id < 0 || token_id >= cumulative_seqlens_q[batch_size]) { + const int batch_id = FindTokenBatch(token_id, cumulative_seqlens_q, batch_size); + if (batch_id < 0 || sequence_validity[batch_id] == 0) { return -1; } - // cumulative_seqlens_q is a non-decreasing prefix sum, so binary search finds the owning - // sequence in log2(batch_size) steps instead of the previous O(batch_size) scan. - int left = 0; - int right = batch_size - 1; - while (left < right) { - const int mid = left + (right - left) / 2; - if (token_id < cumulative_seqlens_q[mid + 1]) { - right = mid; - } else { - left = mid + 1; - } - } - const int batch_id = left; const int position = past_seqlens[batch_id] + (token_id - cumulative_seqlens_q[batch_id]); const int block_idx_in_seq = position / block_size; - if (block_idx_in_seq >= max_num_blocks_per_seq) { + if (block_idx_in_seq < 0 || block_idx_in_seq >= max_num_blocks_per_seq) { return -1; } const int block_id = block_table[batch_id * max_num_blocks_per_seq + block_idx_in_seq]; @@ -459,8 +642,17 @@ struct DerivedSlotResolver { // kernel, owns block placement. It also removes the per-thread binary search entirely. struct ExplicitSlotResolver { const int* __restrict__ slot_mapping; + const int* __restrict__ cumulative_seqlens_q = nullptr; + const int* __restrict__ sequence_validity = nullptr; + int batch_size = 0; __device__ __forceinline__ int operator()(int token_id) const { + if (sequence_validity != nullptr) { + const int batch_id = FindTokenBatch(token_id, cumulative_seqlens_q, batch_size); + if (batch_id < 0 || sequence_validity[batch_id] == 0) { + return -1; + } + } return slot_mapping[token_id]; } }; @@ -576,20 +768,23 @@ template Status LaunchReshapeAndCache(const T* key, const T* value, TCACHE* key_cache, TCACHE* value_cache, const float* k_scale, const float* v_scale, const bool k_per_channel, const bool v_per_channel, const int* block_table, - const int* past_seqlens, const int* cumulative_seqlens_q, const int* slot_mapping, + const int* past_seqlens, const int* cumulative_seqlens_q, + const int* sequence_validity, const int* slot_mapping, const int batch_size, const int max_num_blocks_per_seq, const int token_count, const int kv_hidden_size, const int block_size, const int num_blocks, const int key_stride, const int value_stride, cudaStream_t stream, const int max_threads_per_block) { const int64_t num_slots = static_cast(num_blocks) * block_size; if (slot_mapping != nullptr) { - ExplicitSlotResolver resolver{slot_mapping}; + ExplicitSlotResolver resolver{ + slot_mapping, cumulative_seqlens_q, sequence_validity, batch_size}; return LaunchReshapeAndCacheImpl( key, value, key_cache, value_cache, k_scale, v_scale, k_per_channel, v_per_channel, resolver, token_count, kv_hidden_size, key_stride, value_stride, num_slots, stream, max_threads_per_block); } - DerivedSlotResolver resolver{block_table, past_seqlens, cumulative_seqlens_q, batch_size, - max_num_blocks_per_seq, block_size}; + DerivedSlotResolver resolver{ + block_table, past_seqlens, cumulative_seqlens_q, sequence_validity, + batch_size, max_num_blocks_per_seq, block_size}; return LaunchReshapeAndCacheImpl( key, value, key_cache, value_cache, k_scale, v_scale, k_per_channel, v_per_channel, resolver, token_count, kv_hidden_size, key_stride, value_stride, num_slots, stream, max_threads_per_block); @@ -1475,17 +1670,24 @@ Status PrepareQueryAndCache(cudaStream_t stream, contrib::PagedAttentionParamete const bool v_per_channel = parameters.v_quant_type == KVQuantizationType::PER_CHANNEL; if constexpr (std::is_same_v) { if (data.slot_mapping != nullptr) { - ORT_RETURN_IF_ERROR(LaunchCacheHeads(key, value, data, parameters, ExplicitSlotResolver{data.slot_mapping}, - key_stride, value_stride, stream)); + ORT_RETURN_IF_ERROR(LaunchCacheHeads( + key, value, data, parameters, + ExplicitSlotResolver{ + data.slot_mapping, cumulative_seqlens_q, + data.sequence_validity, batch_size}, + key_stride, value_stride, stream)); } else { - DerivedSlotResolver resolver{data.block_table, past_seqlens, cumulative_seqlens_q, batch_size, - parameters.max_num_blocks_per_seq, parameters.block_size}; + DerivedSlotResolver resolver{ + data.block_table, past_seqlens, cumulative_seqlens_q, + data.sequence_validity, batch_size, + parameters.max_num_blocks_per_seq, parameters.block_size}; ORT_RETURN_IF_ERROR(LaunchCacheHeads(key, value, data, parameters, resolver, key_stride, value_stride, stream)); } } else { ORT_RETURN_IF_ERROR((LaunchReshapeAndCache( key, value, data.key_cache, data.value_cache, data.k_scale, data.v_scale, k_per_channel, v_per_channel, - const_cast(data.block_table), past_seqlens, cumulative_seqlens_q, data.slot_mapping, batch_size, + const_cast(data.block_table), past_seqlens, cumulative_seqlens_q, + data.sequence_validity, data.slot_mapping, batch_size, parameters.max_num_blocks_per_seq, token_count, kv_hidden_size, parameters.block_size, parameters.num_blocks, key_stride, value_stride, stream, max_threads_per_block))); } @@ -1877,12 +2079,6 @@ Status CudnnPagedAttention( ORT_RETURN_IF_ERROR((PrepareQueryAndCache(stream, parameters, data, max_threads_per_block, &query))); - // Per-batch KV lengths for the padding-mask input. Derived from cumulative_seqlens_q so this - // stays correct if the cuDNN paged tier is ever relaxed beyond decode-only. - ORT_RETURN_IF_ERROR(LaunchGetSeqlensKV( - data.cudnn_seqlens_kv, data.past_seqlens, data.cumulative_seqlens_q, - parameters.batch_size, stream)); - cudnnHandle_t cudnn_handle = static_cast(data.cudnn_handle); bool cache_hit = false; const bool ok = onnxruntime::cudnn_sdpa::run_paged( @@ -1890,7 +2086,7 @@ Status CudnnPagedAttention( /*q=*/reinterpret_cast(query), /*k_cache=*/reinterpret_cast(data.key_cache), /*v_cache=*/reinterpret_cast(data.value_cache), - /*block_table=*/const_cast(data.block_table), + /*block_table=*/const_cast(data.cudnn_block_table), /*mask_sequence_lengths_kv=*/data.cudnn_seqlens_kv, parameters.batch_size, parameters.num_heads, @@ -1972,9 +2168,6 @@ Status FlashAttention( void* softmax_lse = reinterpret_cast(data.softmax_lse); if constexpr (IsQuantizedCache::value) { - // FlashAttention cannot read a quantized page, so dequantize the live context into a dense - // packed-varlen [total_kv_tokens, kv_num_heads, head_size] buffer (no GQA expansion — Flash - // does the grouping itself) and use the non-paged varlen entry point. ORT_RETURN_IF_ERROR((LaunchGatherAndExpandPagedKVCache( data.key_cache, data.value_cache, data.gathered_key, data.gathered_value, data.k_scale, data.v_scale, k_per_channel, v_per_channel, @@ -1989,14 +2182,12 @@ Status FlashAttention( local_window_size - 1, /*max_num_blocks_per_seq*/ 0, /*page_block_size*/ 1, data.flash_num_splits, data.flash_softmax_lse_accum, data.flash_out_accum)); } else { - void* key_cache = reinterpret_cast(data.key_cache); - void* value_cache = reinterpret_cast(data.value_cache); ORT_RETURN_IF_ERROR(onnxruntime::flash::mha_varlen_fwd( - device_prop, stream, q, key_cache, value_cache, output, cumulative_seqlens_q, cumulative_seqlens_kv, + device_prop, stream, q, reinterpret_cast(data.key_cache), + reinterpret_cast(data.value_cache), output, cumulative_seqlens_q, cumulative_seqlens_kv, /*seqused_k*/ nullptr, block_table, softmax_lse, batch_size, num_heads, kv_num_heads, head_size, max_query_len, data.max_kv_len, token_count, scale, softcap, parameters.is_causal, is_bf16, - local_window_size - 1, - max_num_blocks_per_seq, block_size, + local_window_size - 1, max_num_blocks_per_seq, block_size, data.flash_num_splits, data.flash_softmax_lse_accum, data.flash_out_accum)); } @@ -2125,39 +2316,35 @@ Status QkvToContext( PagedAttentionData& data) { auto stream = static_cast(ort_stream->GetHandle()); const float scale = parameters.scale == 0.0f ? 1.f / sqrt(static_cast(parameters.head_size)) : parameters.scale; + Status attention_status; // LATENT (MLA) has its own backend: no other kernel can serve v_head_size != head_size over a // single aliased cache. Validation guarantees an explicit scale here, so the default above is // never the one used. if (parameters.is_latent_kv) { - return LatentAttention(device_prop, stream, parameters, data, scale); - } - - if (data.use_xqa_decode) { - return PagedXqaDecodeAttention(device_prop, stream, parameters, data, scale); - } - - if (data.use_cudnn_paged) { - return CudnnPagedAttention(device_prop, ort_stream, parameters, data); - } - - if (data.use_paged_decode) { - return PagedDecodeAttention(device_prop, stream, parameters, data, scale); - } - + attention_status = LatentAttention(device_prop, stream, parameters, data, scale); + } else if (data.use_xqa_decode) { + attention_status = PagedXqaDecodeAttention(device_prop, stream, parameters, data, scale); + } else if (data.use_cudnn_paged) { + attention_status = CudnnPagedAttention(device_prop, ort_stream, parameters, data); + } else if (data.use_paged_decode) { + attention_status = PagedDecodeAttention(device_prop, stream, parameters, data, scale); #if USE_FLASH_ATTENTION - if (data.use_flash_attention) { - return FlashAttention(device_prop, stream, parameters, data, scale); - } + } else if (data.use_flash_attention) { + attention_status = FlashAttention(device_prop, stream, parameters, data, scale); #endif - #if USE_MEMORY_EFFICIENT_ATTENTION - if (data.use_memory_efficient_attention) { - return EfficientAttention(device_prop, stream, parameters, data, scale); - } + } else if (data.use_memory_efficient_attention) { + attention_status = EfficientAttention(device_prop, stream, parameters, data, scale); #endif + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "No PagedAttention kernel available for the current configuration."); + } - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "No PagedAttention kernel available for the current configuration."); + ORT_RETURN_IF_ERROR(attention_status); + return LaunchMaskInvalidSequenceOutputs( + data.output, data.sequence_validity, data.cumulative_seqlens_q, + parameters.batch_size, parameters.v_hidden_size, stream); } #define INSTANTIATE_PAGED_ATTENTION(T, TCACHE) \ diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h index e310d53ba0b79..8bb7cd558bc9a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h @@ -39,21 +39,32 @@ Status LaunchUnpackQKVCumulative(const T* packed_qkv, T* unpacked_q, T* unpacked const int kv_num_heads, const int head_size, const int token_count, cudaStream_t stream, const int max_threads_per_block); -// Exposed so paged_attention.cc can populate cumulative_seqlens_kv on both the FA and MEA -// dispatch paths (producer hoisted out of FlashAttention/UnfusedAttention in impl.cu). -Status LaunchGetCumulativeSeqlensKV(int32_t* cumulative_seqlens_kv, const int32_t* cumulative_seqlens_q, - const int32_t* past_seqlens, const int batch_size, cudaStream_t stream); +Status GetSanitizeSequenceLengthsWorkspaceSize(int batch_size, + size_t& workspace_bytes, + cudaStream_t stream); -// Produces per-batch KV lengths seqlens_kv[i] = past_seqlens[i] + (cumulative_seqlens_q[i+1] - -// cumulative_seqlens_q[i]) for the cuDNN paged SDPA backend's padding-mask input. Deriving the -// query count from cumulative_seqlens_q keeps the kernel correct if a caller ever routes a -// multi-token step through this path -- the cuDNN paged tier is decode-only today, so -// cumulative_seqlens_q[i+1] - cumulative_seqlens_q[i] is 1 for every eligible step, but the -// derivation avoids baking that invariant into a magic constant. -Status LaunchGetSeqlensKV(int32_t* seqlens_kv, const int32_t* past_seqlens, - const int32_t* cumulative_seqlens_q, - const int batch_size, cudaStream_t stream); +Status LaunchSanitizeSequenceLengths(int32_t* sanitized_cumulative_seqlens_q, + int32_t* sanitized_past_seqlens, + int32_t* sequence_validity, + int32_t* cumulative_seqlens_kv, + const int32_t* cumulative_seqlens_q, + const int32_t* past_seqlens, + void* workspace, + size_t workspace_bytes, + int batch_size, + int max_num_blocks_per_seq, + int block_size, + int token_count, + cudaStream_t stream); +Status LaunchSanitizeBlockTable(const int32_t* block_table, int32_t* sanitized_block_table, + size_t element_count, int num_blocks, cudaStream_t stream); +Status LaunchPrepareCudnnBlockTable(const int32_t* block_table, int32_t* cudnn_block_table, + const int32_t* cumulative_seqlens_kv, + int32_t* cudnn_seqlens_kv, int32_t* sequence_validity, + size_t element_count, int batch_size, + int max_num_blocks_per_seq, int block_size, + cudaStream_t stream); // Paged decode backend sizing helpers, used by paged_attention.cc to test eligibility (the kernel // needs more dynamic shared memory than the device provides for very wide heads) and to size the // split-KV workspaces. diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index 7f5169639ff8d..d4aeca0e657bd 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -31,6 +31,7 @@ limitations under the License. #include "contrib_ops/cuda/bert/skip_layer_norm_impl.h" #include #include +#include namespace onnxruntime { namespace contrib { @@ -68,20 +69,34 @@ bool CanVectorized(void* output, void* sum_output, const void* input, const void next_size / num_unroll >= kMinBlockSize && next_size / num_unroll <= kMaxBlockSize; } + +template +__device__ inline float SkipLayerNormOrigin( + const T* input, const T* skip, const T* bias, int offset, int skip_size) { + __shared__ float origin; + if (threadIdx.x == 0) { + T value = input[offset]; + if (bias != nullptr) value += bias[0]; + value += skip[offset % skip_size]; + origin = static_cast(value); + } + // Read the origin before any thread writes the fused residual sum. + __syncthreads(); + return origin; +} } // namespace template __global__ void SkipLayerNormKernel( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, const int ld, int skip_size) { + if (ld == 0) return; const float reverse_ld = 1.f / ld; const int offset = blockIdx.x * ld; const bool has_bias = (bias != nullptr); - // Reduce sum of x and x^2, and the results are divided by ld. - // Uses fp32 accumulation to avoid overflow in fp16/bf16. - KeyValuePairSum pair_sum; - cub::KeyValuePair thread_data(0.f, 0.f); + const float origin = Simplified ? 0.f : SkipLayerNormOrigin(input, skip, bias, offset, skip_size); + std::conditional_t thread_data{}; for (int i = threadIdx.x; i < ld; i += TPB) { const int idx = offset + i; @@ -93,8 +108,12 @@ __global__ void SkipLayerNormKernel( val += skip[idx % skip_size]; const float val_f = static_cast(val); - const float rldval = reverse_ld * val_f; - thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val_f)); + if constexpr (Simplified) { + const float rldval = reverse_ld * val_f; + thread_data += rldval * val_f; + } else { + thread_data.Update(val_f - origin); + } if (sum_output != nullptr) { sum_output[idx] = val; @@ -103,11 +122,11 @@ __global__ void SkipLayerNormKernel( output[idx] = val; } - if (Simplified) { - SimplifiedLayerNorm(thread_data.value, ld, offset, gamma, epsilon, output); - return; + if constexpr (Simplified) { + SimplifiedLayerNorm(thread_data, ld, offset, gamma, epsilon, output); + } else { + LayerNorm(thread_data, origin, ld, offset, beta, gamma, epsilon, output); } - LayerNorm(thread_data, ld, offset, beta, gamma, epsilon, output); } // Vectorized kernel @@ -115,13 +134,15 @@ template __global__ void SkipLayerNormKernelSmall( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, int ld, int skip_size) { + if (ld == 0) return; const float rld = 1.f / ld; const int idx = blockIdx.x * ld + threadIdx.x * ILP; using VecT = aligned_vector; T sum_v[ILP]; - cub::KeyValuePair thread_data(0.f, 0.f); + const float origin = Simplified ? 0.f : SkipLayerNormOrigin(input, skip, bias, blockIdx.x * ld, skip_size); + std::conditional_t thread_data{}; if (ILP * threadIdx.x < ld) { // load data under this guard to avoid reading out-of-bounds T skip_v[ILP], bias_v[ILP]; @@ -139,8 +160,6 @@ __global__ void SkipLayerNormKernelSmall( *bias_val = *reinterpret_cast(&bias[threadIdx.x * ILP]); } - float rldval_sum = 0.f; - float rldvalsq_sum = 0.f; const bool has_sum_output = (sum_output != nullptr); #pragma unroll @@ -151,23 +170,24 @@ __global__ void SkipLayerNormKernelSmall( sum_v[i] += skip_v[i]; const float val_f = static_cast(sum_v[i]); - const float rldval = rld * val_f; - rldval_sum += rldval; - rldvalsq_sum += rldval * val_f; + if constexpr (Simplified) { + const float rldval = rld * val_f; + thread_data += rldval * val_f; + } else { + thread_data.Update(val_f - origin); + } } if (has_sum_output) { *(reinterpret_cast(&sum_output[idx])) = *reinterpret_cast(&sum_v); } - - thread_data = cub::KeyValuePair(rldval_sum, rldvalsq_sum); } - if (Simplified) { - SimplifiedLayerNormSmall(sum_v, thread_data.value, ld, idx, gamma, epsilon, output); - return; + if constexpr (Simplified) { + SimplifiedLayerNormSmall(sum_v, thread_data, ld, idx, gamma, epsilon, output); + } else { + LayerNormSmall(sum_v, thread_data, origin, ld, idx, beta, gamma, epsilon, output); } - LayerNormSmall(sum_v, thread_data, ld, idx, beta, gamma, epsilon, output); } template 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 459372618401e..6d930d8c9f606 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc @@ -218,6 +218,25 @@ std::vector WeightOnlyGroupw return tactics; } +std::optional +WeightOnlyGroupwiseQuantGemmPluginProfiler::getDeterministicConfig(int m) const { + const auto configs = mRunner->getConfigs(); + if (m < 16) { + for (const auto& config : configs) { + if (config.enableCudaKernel) { + return config; + } + } + } + for (const auto& config : configs) { + if (!config.enableCudaKernel && config.split_k_style == cutlass_extensions::SplitKStyle::NO_SPLIT_K && + (mArch < 89 || config.stages != 2)) { + return config; + } + } + return std::nullopt; +} + bool WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic(int m, int n, int k, Config const& tactic) const { // stop to profile Cuda kernel for m >= 16 if (tactic.enableCudaKernel) { 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 fd610aa5174ff..47691d7b492df 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h @@ -80,6 +80,8 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler public: using Config = onnxruntime::llm::cutlass_extensions::CutlassGemmConfig; + std::optional getDeterministicConfig(int m) const; + // Parses a comma-separated list of M buckets (e.g. "1,8,64,512") into a sorted, de-duplicated, // positive list (empty when the string is empty/blank). Used for the ep.cuda.fpa_intb_profile_m // session-config key and the ORT_FPA_INTB_PROFILE_M env var, both resolved by the kernel. 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 6951bf9a5b458..03ba285f83e19 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 @@ -481,9 +481,22 @@ __global__ void MoeGemvFp4RawKPackedKernel( bool Fp4MoeGemvSupportsNativeScale() { #if defined(ORT_CUDA_TYPE_HELPER_HAS_FP8) + int device = -1; + if (cudaGetDevice(&device) != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } + int sm_major = 0; + if (cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess || sm_major < 9) { + (void)cudaGetLastError(); + return false; + } cudaFuncAttributes attributes{}; - CUDA_CALL_THROW(cudaFuncGetAttributes( - &attributes, MoeGemvFp4RawKPackedKernel)); + if (cudaFuncGetAttributes( + &attributes, MoeGemvFp4RawKPackedKernel) != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } return attributes.binaryVersion >= 90 && attributes.ptxVersion >= 90; #else return false; @@ -747,6 +760,18 @@ __global__ void __launch_bounds__(Warps * 32, MinBlocks) MoeGemvFp4DecodeKernel( static_assert(KLanes == 4 || KLanes == 8); static_assert(Batch >= 1); + // Programmatic dependent launch (only active when launched with the PDL attribute, otherwise no-ops): let the next + // kernel of the MoE chain start its prologue while this one runs. FC2 reads only data of kernels that completed + // before FC1 passed its wait (expert maps, weights) before its own wait and the FC1 output after it. + // FC1 triggers only after its own wait (the routing kernel may be its primary grid), so FC2 and the finalize + // kernel, which read routing metadata before their waits, never launch before that metadata is complete. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (FusedSwiGlu) { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + asm volatile("griddepcontrol.launch_dependents;"); +#endif + const int row = static_cast(blockIdx.x); int expert = permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : 0; #pragma unroll 1 @@ -797,6 +822,11 @@ __global__ void __launch_bounds__(Warps * 32, MinBlocks) MoeGemvFp4DecodeKernel( uint32_t scale_codes[Batch][kDecodeColsPerThread]; load_batch(0, packed, scale_codes); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FusedSwiGlu) { + asm volatile("griddepcontrol.wait;" ::: "memory"); // the activation row is the FC1 output + } +#endif // Stage the activation row as FP32, one padded 32-element unit per lane pair. extern __shared__ __align__(16) float act_smem[]; { @@ -867,6 +897,47 @@ __global__ void __launch_bounds__(Warps * 32, MinBlocks) MoeGemvFp4DecodeKernel( #endif } +// Opt-out for the programmatic dependent launch of FC2 (env ORT_FP4_GEMV_PDL=0). +static bool Fp4MoeGemvUsePdl() { + static const bool enabled = onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_PDL", 1) != 0; + return enabled; +} + +template +static bool QuerySm90DecodeKernelCaps(KernelPtr kernel, bool* pdl_supported = nullptr) { + if (pdl_supported != nullptr) { + *pdl_supported = false; + } + int device = -1; + if (cudaGetDevice(&device) != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } + int sm_major = 0; + if (cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess || sm_major < 9) { + (void)cudaGetLastError(); + return false; + } + cudaFuncAttributes attributes{}; + if (cudaFuncGetAttributes(&attributes, kernel) != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } + const bool image_available = attributes.ptxVersion >= 90 && attributes.binaryVersion >= 90; + if (!image_available) { + return false; + } + if (pdl_supported != nullptr) { + int driver_version = 0; + if (cudaDriverGetVersion(&driver_version) == cudaSuccess) { + *pdl_supported = driver_version >= 12000; + } else { + (void)cudaGetLastError(); + } + } + return true; +} + template void LaunchMoeGemvFp4RawNPacked( const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales, @@ -874,7 +945,8 @@ void LaunchMoeGemvFp4RawNPacked( const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, int64_t n, int64_t k, cutlass_kernels::ActivationParams activation_params, const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, bool weights_row_major, int sm, - bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr) { + bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr, + bool programmatic_dependent_launch = false) { if (decode_kernel_launched != nullptr) { *decode_kernel_launched = false; } @@ -927,26 +999,49 @@ void LaunchMoeGemvFp4RawNPacked( static_cast((n + kBlockCols - 1) / kBlockCols)); const auto launch_bias = [&](auto bias_tag) { constexpr bool kEnableBias = decltype(bias_tag)::value; + auto kernel = MoeGemvFp4DecodeKernel; // PTX compiled below SM90 has an empty decode body, even when JIT-loaded on SM90+. // Image selection depends on the device SM; refresh when a thread switches architectures. static thread_local int cached_sm = -1; static thread_local bool available = false; + static thread_local bool pdl_supported = false; if (cached_sm != sm) { - cudaFuncAttributes attributes{}; - CUDA_CALL_THROW(cudaFuncGetAttributes( - &attributes, MoeGemvFp4DecodeKernel)); - available = attributes.ptxVersion >= 90; + available = QuerySm90DecodeKernelCaps(kernel, &pdl_supported); cached_sm = sm; } if (!available) { return false; } - MoeGemvFp4DecodeKernel - <<>>( - act, weight, block_scales, global_scales, bias, out, - expert_first_token_offset, permuted_row_to_expert, num_experts, - weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k), - activation_params, permuted_row_to_source_row, static_cast(num_rows)); + if (programmatic_dependent_launch && Fp4MoeGemvUsePdl() && pdl_supported) { + cudaLaunchConfig_t config{}; + config.gridDim = decode_grid; + config.blockDim = dim3(kWarps * 32); + config.dynamicSmemBytes = decode_smem; + config.stream = stream; + cudaLaunchAttribute attribute[1]{}; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attribute; + config.numAttrs = 1; + const cudaError_t status = cudaLaunchKernelEx( + &config, kernel, 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 (status == cudaSuccess) { + return true; + } + if (status != cudaErrorInvalidValue && status != cudaErrorNotSupported) { + ORT_ENFORCE(status == cudaSuccess, "FP4 GEMV decode launch failed: ", cudaGetErrorString(status)); + } + (void)cudaGetLastError(); + pdl_supported = false; + } + kernel<<>>( + 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)); + CUDA_CALL_THROW(cudaGetLastError()); return true; }; if (bias != nullptr) { @@ -1108,10 +1203,8 @@ bool is_moe_gemv_fp4_decode_image_available(int sm, bool is_fp16) { } const auto check_type = [&](auto type_tag) { using T = typename decltype(type_tag)::type; - cudaFuncAttributes attributes{}; - CUDA_CALL_THROW(cudaFuncGetAttributes( - &attributes, MoeGemvFp4DecodeKernel)); - return attributes.ptxVersion >= 90; + return QuerySm90DecodeKernelCaps( + MoeGemvFp4DecodeKernel); }; if (is_fp16) { return check_type(TypeTag{}); @@ -1131,7 +1224,7 @@ void launch_moe_gemv_fp4_symmetric(const T* act, const uint8_t* weight, const T* int num_experts, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, int sm, MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed, cudaStream_t stream, bool weights_row_major, bool decode_scale_eligible, - bool* decode_kernel_launched) { + bool* decode_kernel_launched, bool programmatic_dependent_launch) { if (decode_kernel_launched != nullptr) { *decode_kernel_launched = false; } @@ -1143,7 +1236,7 @@ void launch_moe_gemv_fp4_symmetric(const T* act, const uint8_t* weight, const T* act, weight, raw_block_scales, raw_global_scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, n, k, cutlass_kernels::ActivationParams{}, nullptr, expanded_num_rows, stream, - weights_row_major, sm, decode_scale_eligible, decode_kernel_launched); + weights_row_major, sm, decode_scale_eligible, decode_kernel_launched, programmatic_dependent_launch); return; } // Interleaved path: ColumnMajorInterleaved layout + dtype-conditional accumulation + smaller @@ -1203,7 +1296,7 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( cutlass_kernels::ActivationParams activation_params, MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed, const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, bool weights_row_major, - bool decode_scale_eligible, bool* decode_kernel_launched) { + bool decode_scale_eligible, bool* decode_kernel_launched, bool programmatic_dependent_launch) { if (decode_kernel_launched != nullptr) { *decode_kernel_launched = false; } @@ -1215,7 +1308,7 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( act, weight, raw_block_scales, raw_global_scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, inter_size * 2, k, activation_params, permuted_row_to_source_row, num_rows, stream, - weights_row_major, sm, decode_scale_eligible, decode_kernel_launched); + weights_row_major, sm, decode_scale_eligible, decode_kernel_launched, programmatic_dependent_launch); return; } // Interleaved path: ColumnMajorInterleaved layout + dtype-conditional accumulation + smaller @@ -1271,23 +1364,23 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( template void launch_moe_gemv_fp4_symmetric( const half*, const uint8_t*, const half*, const uint8_t*, const float*, const half*, half*, const int64_t*, const int*, int, - int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, bool, cudaStream_t, bool, bool, bool*); + int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, bool, cudaStream_t, bool, bool, bool*, bool); template void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( const half*, const uint8_t*, const half*, const uint8_t*, const float*, const half*, half*, const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, MoeGemvConfig, bool, bool, - const int*, int64_t, cudaStream_t, bool, bool, bool*); + const int*, int64_t, cudaStream_t, bool, bool, bool*, bool); #ifdef ENABLE_BF16 template void launch_moe_gemv_fp4_symmetric<__nv_bfloat16>( const __nv_bfloat16*, const uint8_t*, const __nv_bfloat16*, const uint8_t*, const float*, const __nv_bfloat16*, __nv_bfloat16*, const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, bool, cudaStream_t, - bool, bool, bool*); + bool, bool, bool*, bool); template void launch_moe_gemv_fp4_symmetric_interleaved_swiglu<__nv_bfloat16>( const __nv_bfloat16*, const uint8_t*, const __nv_bfloat16*, const uint8_t*, const float*, const __nv_bfloat16*, __nv_bfloat16*, const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, - MoeGemvConfig, bool, bool, const int*, int64_t, cudaStream_t, bool, bool, bool*); + MoeGemvConfig, bool, bool, const int*, int64_t, cudaStream_t, bool, bool, bool*, bool); #endif } // namespace moe_gemv diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h index 4081168fd07cc..47dfe69f1da58 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 @@ -122,7 +122,12 @@ void launch_moe_gemv_fp4_symmetric( const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, int sm, MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed, cudaStream_t stream, bool weights_row_major = false, - bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr); + bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr, + bool programmatic_dependent_launch = false); + +// `programmatic_dependent_launch` (raw row-major NVFP4 decode kernel only) launches the grid with the PDL attribute: it +// must directly follow the fused-SwiGLU launch below on the same stream and may only read that kernel's output after +// its own griddepcontrol.wait (see MoeGemvFp4DecodeKernel). // Launches the MXFP4 MoE GEMV and fuses interleaved SwiGLU activation. // weight/scales/bias use raw FC1 output width n = 2 * inter_size @@ -135,7 +140,8 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( int64_t inter_size, int64_t k, int group_size, int sm, cutlass_kernels::ActivationParams activation_params, MoeGemvConfig config, bool sm80_pair_interleaved, bool raw_n_packed, const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, - bool weights_row_major = false, bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr); + bool weights_row_major = false, bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr, + bool programmatic_dependent_launch = false); } // namespace moe_gemv } // namespace kernels diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu index 43bd9170cc772..9bac149cebe96 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu @@ -1592,13 +1592,17 @@ constexpr static int FINALIZE_THREADS_PER_BLOCK = 256; // accesses that can all be in flight at the same time. constexpr static int FINALIZE_MAX_STAGED_TOPK = FINALIZE_THREADS_PER_BLOCK; +// Expert rows whose loads are issued together by the staged finalize kernel (top-k <= 16 takes one round). +constexpr static int FINALIZE_LOAD_BATCH = 16; + // Final kernel to unpermute and scale // This kernel unpermutes the original data, does the k-way reduction and performs the final skip connection. template __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted_rows, OutputType* reduced_unpermuted_output, const ScaleBiasType* bias, const float* scales, const int* unpermuted_row_to_permuted_row, const int* token_selected_experts, const int64_t orig_cols, - const int64_t experts_per_token, const int num_experts_per_node, const int start_expert_id) { + const int64_t experts_per_token, const int num_experts_per_node, const int start_expert_id, + const bool early_routing_loads) { const int64_t original_row = blockIdx.x; const int64_t num_rows = gridDim.x; const auto offset = original_row * orig_cols; @@ -1609,7 +1613,7 @@ __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted assert(orig_cols % FINALIZE_ELEM_PER_THREAD == 0); const int64_t start_offset = threadIdx.x; - const int64_t stride = FINALIZE_THREADS_PER_BLOCK; + const int64_t stride = blockDim.x; const int64_t num_elems_in_col = orig_cols / FINALIZE_ELEM_PER_THREAD; using BiasElem = cutlass::Array; @@ -1621,14 +1625,16 @@ __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted auto* reduced_row_ptr_v = reinterpret_cast(reduced_row_ptr); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + if (!(STAGE_ROUTING && early_routing_loads)) { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } #endif __shared__ int s_permuted_row[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1]; __shared__ int s_expert_id[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1]; __shared__ float s_row_scale[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1]; if constexpr (STAGE_ROUTING) { - for (int64_t k_idx = threadIdx.x; k_idx < experts_per_token; k_idx += FINALIZE_THREADS_PER_BLOCK) { + for (int64_t k_idx = threadIdx.x; k_idx < experts_per_token; k_idx += blockDim.x) { int64_t const k_offset = original_row * experts_per_token + k_idx; int const expert_id = token_selected_experts[k_offset] - start_expert_id; bool const is_local_expert = expert_id >= 0 && expert_id < num_experts_per_node; @@ -1639,42 +1645,66 @@ __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted } __syncthreads(); } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (STAGE_ROUTING && early_routing_loads) { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif #pragma unroll for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { ComputeElem thread_output; thread_output.fill(0); -#pragma unroll 4 - for (int k_idx = 0; k_idx < experts_per_token; ++k_idx) { - int64_t expanded_permuted_row; - int64_t expert_id; - float row_scale; - if constexpr (STAGE_ROUTING) { - expanded_permuted_row = s_permuted_row[k_idx]; - expert_id = s_expert_id[k_idx]; - row_scale = s_row_scale[k_idx]; - if (expanded_permuted_row < 0) { - continue; + if constexpr (STAGE_ROUTING) { + // Issue the loads of up to FINALIZE_LOAD_BATCH expert rows together (they are independent), then + // accumulate them in top-k order: the sums are the same as with one load per iteration. + for (int k_base = 0; k_base < experts_per_token; k_base += FINALIZE_LOAD_BATCH) { + InputElem batch[FINALIZE_LOAD_BATCH]; +#pragma unroll + for (int i = 0; i < FINALIZE_LOAD_BATCH; ++i) { + const int k_idx = k_base + i; + if (k_idx < experts_per_token && s_permuted_row[k_idx] >= 0) { + batch[i] = expanded_permuted_rows_v[static_cast(s_permuted_row[k_idx]) * num_elems_in_col + elem_index]; + } } - } else { +#pragma unroll + for (int i = 0; i < FINALIZE_LOAD_BATCH; ++i) { + const int k_idx = k_base + i; + if (k_idx >= experts_per_token || s_permuted_row[k_idx] < 0) { + continue; + } + const int64_t expert_id = s_expert_id[k_idx]; + const float row_scale = s_row_scale[k_idx]; + ComputeElem expert_result = arrayConvert(batch[i]); + if (bias) { + const auto* bias_ptr = bias_v + expert_id * num_elems_in_col; + expert_result = expert_result + arrayConvert(bias_ptr[elem_index]); + } + + thread_output = thread_output + row_scale * expert_result; + } + } + } else { +#pragma unroll 4 + for (int k_idx = 0; k_idx < experts_per_token; ++k_idx) { const int64_t k_offset = original_row * experts_per_token + k_idx; - expert_id = token_selected_experts[k_offset] - start_expert_id; + const int64_t expert_id = token_selected_experts[k_offset] - start_expert_id; if (expert_id < 0 || expert_id >= num_experts_per_node) { continue; } - expanded_permuted_row = unpermuted_row_to_permuted_row[original_row + k_idx * num_rows]; - row_scale = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset]; - } + const int64_t expanded_permuted_row = unpermuted_row_to_permuted_row[original_row + k_idx * num_rows]; + const float row_scale = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset]; - const auto* expanded_permuted_rows_row_ptr = expanded_permuted_rows_v + expanded_permuted_row * num_elems_in_col; + const auto* expanded_permuted_rows_row_ptr = expanded_permuted_rows_v + expanded_permuted_row * num_elems_in_col; - ComputeElem expert_result = arrayConvert(expanded_permuted_rows_row_ptr[elem_index]); - if (bias) { - const auto* bias_ptr = bias_v + expert_id * num_elems_in_col; - expert_result = expert_result + arrayConvert(bias_ptr[elem_index]); - } + ComputeElem expert_result = arrayConvert(expanded_permuted_rows_row_ptr[elem_index]); + if (bias) { + const auto* bias_ptr = bias_v + expert_id * num_elems_in_col; + expert_result = expert_result + arrayConvert(bias_ptr[elem_index]); + } - thread_output = thread_output + row_scale * expert_result; + thread_output = thread_output + row_scale * expert_result; + } } OutputElem output_elem = arrayConvert(thread_output); @@ -1843,13 +1873,47 @@ __global__ void finalizeMoeRoutingNoFillingKernel(const GemmOutputType* expanded #endif } +template +struct FinalizeKernelLaunchCaps { + int max_threads_per_block; + bool sm90_pdl_available; +}; + +template +FinalizeKernelLaunchCaps getFinalizeKernelLaunchCaps(KernelPtr func) { + static thread_local int cached_device = -1; + static thread_local KernelPtr cached_func = nullptr; + static thread_local FinalizeKernelLaunchCaps cached_caps{FINALIZE_THREADS_PER_BLOCK, false}; + int device = -1; + CUDA_CALL_THROW(cudaGetDevice(&device)); + if (device != cached_device || func != cached_func) { + cudaFuncAttributes attr{}; + CUDA_CALL_THROW(cudaFuncGetAttributes(&attr, func)); + int driver_version = 0; + int sm_major = 0; + const bool driver_and_arch_ok = + cudaDriverGetVersion(&driver_version) == cudaSuccess && driver_version >= 12000 && + cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device) == cudaSuccess && sm_major >= 9; + if (!driver_and_arch_ok) { + (void)cudaGetLastError(); + } + const int max_threads = std::max(32, std::min(1024, (attr.maxThreadsPerBlock / 32) * 32)); + cached_caps = {max_threads, + driver_and_arch_ok && attr.ptxVersion >= 90 && attr.binaryVersion >= 90}; + cached_device = device; + cached_func = func; + } + return cached_caps; +} + template void finalizeMoeRoutingKernelLauncher(const GemmOutputType* expanded_permuted_rows, OutputType* reduced_unpermuted_output, const ScaleBiasType* bias, const float* final_scales, const int* unpermuted_row_to_permuted_row, const int* permuted_row_to_unpermuted_row, const int* token_selected_experts, const int64_t* expert_first_token_offset, const int64_t num_rows, const int64_t cols, const int64_t experts_per_token, const int64_t num_experts_per_node, - MOEParallelismConfig parallelism_config, const bool enable_alltoall, cudaStream_t stream) { + MOEParallelismConfig parallelism_config, const bool enable_alltoall, cudaStream_t stream, + const bool early_routing_loads) { // Only add bias on rank 0 for tensor parallelism const bool is_rank_0 = parallelism_config.tp_rank == 0; const ScaleBiasType* bias_ptr = is_rank_0 ? bias : nullptr; @@ -1860,12 +1924,13 @@ void finalizeMoeRoutingKernelLauncher(const GemmOutputType* expanded_permuted_ro int num_experts_per_node_int = SafeInt(num_experts_per_node); const int start_expert_id = num_experts_per_node_int * parallelism_config.ep_rank; - cudaLaunchConfig_t config; + const bool env_pdl = onnxruntime::llm::common::getEnvEnablePDL(); + cudaLaunchConfig_t config{}; config.dynamicSmemBytes = 0; config.stream = stream; - cudaLaunchAttribute attrs[1]; + cudaLaunchAttribute attrs[1]{}; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = onnxruntime::llm::common::getEnvEnablePDL(); + attrs[0].val.programmaticStreamSerializationAllowed = env_pdl ? 1 : 0; config.numAttrs = 1; config.attrs = attrs; @@ -1880,23 +1945,30 @@ void finalizeMoeRoutingKernelLauncher(const GemmOutputType* expanded_permuted_ro auto func = final_scales ? &finalizeMoeRoutingNoFillingKernel : &finalizeMoeRoutingNoFillingKernel; - cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, - unpermuted_row_to_permuted_row, permuted_row_to_unpermuted_row, token_selected_experts, - expert_first_token_offset, num_rows, cols, experts_per_token, num_experts_per_node_int, start_expert_id); + CUDA_CALL_THROW(cudaLaunchKernelEx( + &config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, + unpermuted_row_to_permuted_row, permuted_row_to_unpermuted_row, token_selected_experts, + expert_first_token_offset, num_rows, cols, experts_per_token, num_experts_per_node_int, start_expert_id)); } else { + if (num_rows <= 0) { + return; + } // If all-gather reduce-scatter is used, finalizeMoeRouting must fill invalid output tokens with zeros. const int64_t blocks = num_rows; - const int64_t threads = FINALIZE_THREADS_PER_BLOCK; + // One thread per 128-bit column vector (clamped to the kernel's register-limited maxThreadsPerBlock) so + // each block reduces its row in as few passes as possible; the one-row kernel below keeps the fixed block size. + const int64_t column_vectors = cols / kFinalizeElemPerThread; config.gridDim = blocks; - config.blockDim = threads; if (num_rows == 1 && experts_per_token > 0 && experts_per_token <= 4) { + config.blockDim = FINALIZE_THREADS_PER_BLOCK; #define LAUNCH_FINALIZE_ONE_ROW(EXPERTS_PER_TOKEN) \ do { \ auto func = final_scales \ ? &finalizeMoeRoutingOneRowKernel \ : &finalizeMoeRoutingOneRowKernel; \ - cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, \ - unpermuted_row_to_permuted_row, token_selected_experts, cols, num_experts_per_node_int, start_expert_id); \ + CUDA_CALL_THROW(cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, \ + unpermuted_row_to_permuted_row, token_selected_experts, cols, num_experts_per_node_int, \ + start_expert_id)); \ } while (0) switch (experts_per_token) { @@ -1918,16 +1990,42 @@ void finalizeMoeRoutingKernelLauncher(const GemmOutputType* expanded_permuted_ro auto func = final_scales ? &finalizeMoeRoutingKernel : &finalizeMoeRoutingKernel; - cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, - unpermuted_row_to_permuted_row, token_selected_experts, cols, experts_per_token, num_experts_per_node_int, - start_expert_id); + const auto caps = getFinalizeKernelLaunchCaps(func); + config.blockDim = std::min( + caps.max_threads_per_block, + std::max(FINALIZE_THREADS_PER_BLOCK / 2, (column_vectors + 31) / 32 * 32)); + const bool use_early_routing_loads = early_routing_loads && caps.sm90_pdl_available; + if (use_early_routing_loads) { + attrs[0].val.programmaticStreamSerializationAllowed = 1; + const cudaError_t status = cudaLaunchKernelEx( + &config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, + unpermuted_row_to_permuted_row, token_selected_experts, cols, experts_per_token, num_experts_per_node_int, + start_expert_id, true); + if (status == cudaSuccess) { + return; + } + if (status != cudaErrorInvalidValue && status != cudaErrorNotSupported) { + CUDA_CALL_THROW(status); + } + (void)cudaGetLastError(); + attrs[0].val.programmaticStreamSerializationAllowed = 0; + } + CUDA_CALL_THROW(cudaLaunchKernelEx( + &config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, + unpermuted_row_to_permuted_row, token_selected_experts, cols, experts_per_token, num_experts_per_node_int, + start_expert_id, false)); } else { auto func = final_scales ? &finalizeMoeRoutingKernel : &finalizeMoeRoutingKernel; - cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, - unpermuted_row_to_permuted_row, token_selected_experts, cols, experts_per_token, num_experts_per_node_int, - start_expert_id); + const auto caps = getFinalizeKernelLaunchCaps(func); + config.blockDim = std::min( + caps.max_threads_per_block, + std::max(FINALIZE_THREADS_PER_BLOCK / 2, (column_vectors + 31) / 32 * 32)); + CUDA_CALL_THROW(cudaLaunchKernelEx( + &config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, + unpermuted_row_to_permuted_row, token_selected_experts, cols, experts_per_token, num_experts_per_node_int, + start_expert_id, false)); } } } @@ -1939,7 +2037,7 @@ void finalizeMoeRoutingKernelLauncher(const GemmOutputType* expanded_permuted_ro const int* expanded_dest_row_to_expanded_source_row, const int* expert_for_source_row, \ const int64_t* expert_first_token_offset, const int64_t num_rows, const int64_t cols, \ const int64_t experts_per_token, const int64_t num_experts_per_node, MOEParallelismConfig parallelism_config, \ - const bool enable_alltoall, cudaStream_t stream); + const bool enable_alltoall, cudaStream_t stream, const bool early_routing_loads); // Instantiate the data types that are used by the external pytorch op INSTANTIATE_FINALIZE_MOE_ROUTING(half, half, half); diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h index 296bebdac494d..46e86b6b9a524 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h @@ -60,13 +60,17 @@ void expandInputRowsKernelLauncher(const InputActivationsType* unpermuted_input, int64_t* expert_first_token_offset, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, const TmaWarpSpecializedGroupedGemmInput::ElementSF* input_sf, const void* prequant_scales, cudaStream_t stream); +// early_routing_loads: launch with the PDL attribute and read the routing metadata (written by kernels that completed +// before the producer of `expanded_permuted_rows` started) before waiting for that producer, which only the expanded +// rows depend on. Only valid for the staged-routing kernel (experts_per_token <= 256). template void finalizeMoeRoutingKernelLauncher(const GemmOutputType* expanded_permuted_rows, OutputType* reduced_unpermuted_output, const ScaleBiasType* bias, const float* final_scales, const int* unpermuted_row_to_permuted_row, const int* permuted_row_to_unpermuted_row, const int* token_selected_experts, const int64_t* expert_first_token_offset, const int64_t num_rows, const int64_t cols, const int64_t experts_per_token, const int64_t num_experts_per_node, - MOEParallelismConfig parallelism_config, const bool enable_alltoall, cudaStream_t stream); + MOEParallelismConfig parallelism_config, const bool enable_alltoall, cudaStream_t stream, + const bool early_routing_loads = false); } // namespace cutlass_kernels } // namespace onnxruntime::llm::kernels diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index a48b2b9ca8e3c..990d8c3996cc1 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -2268,32 +2268,62 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { act_params.swiglu_fusion = swiglu_fusion; act_params.limit = swiglu_limit_; - const auto fused_routing = route_tile(0, num_rows); - ORT_ENFORCE(fused_routing.router_logits == nullptr, - "QMoE FP4 GEMV does not support fused routing."); + // Decode routing (softmax + top-k + expert maps) in one single-CTA launch for the raw row-major NVFP4 GEMV route; + // bit-identical to LaunchSoftmaxTopK + fusedBuildExpertMapsSortFirstToken (env ORT_FP4_GEMV_FUSED_ROUTE=0 disables). + static const bool fp4_gemv_pdl = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_PDL", 1) != 0; + static const bool fp4_gemv_fused_route = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_FUSED_ROUTE", 1) != 0; + static const bool fp4_gemv_fc1_pdl = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_FC1_PDL", 0) != 0; + bool fused_route_launched = false; #if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) - if (routing_snapshot_) { - ORT_RETURN_IF_ERROR(routing_snapshot_->Capture( - expert_indices, SafeInt(moe_params.num_rows) * SafeInt(k_), stream)); - } + const bool routing_snapshot_active = routing_snapshot_.has_value(); +#else + constexpr bool routing_snapshot_active = false; #endif + if (fp4_gemv_fused_route && use_raw_nvfp4_gemv && nvfp4_weights_row_major_ && sm_ >= 90 && + !routing_snapshot_active && (is_fp16 || is_bf16)) { + if (is_fp16) { + fused_route_launched = LaunchFusedRouteMaps( + static_cast(router_probs->DataRaw()), expert_scales, expert_indices, p_r2u, + unpermuted_row_to_permuted_row, p_exp, p_efto, static_cast(num_rows), num_experts, + static_cast(k_), normalize_routing_weights_, fp4_gemv_pdl, stream); + } else { + fused_route_launched = LaunchFusedRouteMaps( + static_cast(router_probs->DataRaw()), expert_scales, expert_indices, p_r2u, + unpermuted_row_to_permuted_row, p_exp, p_efto, static_cast(num_rows), num_experts, + static_cast(k_), normalize_routing_weights_, fp4_gemv_pdl, stream); + } + } IAllocatorUniquePtr blocked_expert_counts; IAllocatorUniquePtr blocked_expert_counts_cumsum; IAllocatorUniquePtr blocked_row_to_unpermuted_row; - if (!ck::fusedBuildExpertMapsSortFirstToken( - expert_indices, p_r2u, unpermuted_row_to_permuted_row, p_exp, p_efto, - num_rows, num_experts, static_cast(k_), 0, num_experts, stream)) { - const int64_t num_tokens_per_block = ck::computeNumTokensPerBlock(num_rows, num_experts); - const int64_t num_blocks_per_seq = onnxruntime::llm::common::ceilDiv(num_rows, num_tokens_per_block); - const size_t num_counts = SafeInt(num_experts) * num_blocks_per_seq; - blocked_expert_counts = GetScratchBuffer(num_counts, GetComputeStream(context)); - blocked_expert_counts_cumsum = GetScratchBuffer(num_counts, GetComputeStream(context)); - blocked_row_to_unpermuted_row = GetScratchBuffer( - SafeInt(num_experts) * num_rows, GetComputeStream(context)); - ck::threeStepBuildExpertMapsSortFirstToken( - expert_indices, p_exp, p_r2u, unpermuted_row_to_permuted_row, p_efto, - blocked_expert_counts.get(), blocked_expert_counts_cumsum.get(), - blocked_row_to_unpermuted_row.get(), num_rows, num_experts, k_, 0, stream); + if (!fused_route_launched) { + const auto fused_routing = route_tile(0, num_rows); + ORT_ENFORCE(fused_routing.router_logits == nullptr, + "QMoE FP4 GEMV does not support fused routing."); +#if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) + if (routing_snapshot_) { + ORT_RETURN_IF_ERROR(routing_snapshot_->Capture( + expert_indices, SafeInt(moe_params.num_rows) * SafeInt(k_), stream)); + } +#endif + if (!ck::fusedBuildExpertMapsSortFirstToken( + expert_indices, p_r2u, unpermuted_row_to_permuted_row, p_exp, p_efto, + num_rows, num_experts, static_cast(k_), 0, num_experts, stream)) { + const int64_t num_tokens_per_block = ck::computeNumTokensPerBlock(num_rows, num_experts); + const int64_t num_blocks_per_seq = onnxruntime::llm::common::ceilDiv(num_rows, num_tokens_per_block); + const size_t num_counts = SafeInt(num_experts) * num_blocks_per_seq; + blocked_expert_counts = GetScratchBuffer(num_counts, GetComputeStream(context)); + blocked_expert_counts_cumsum = GetScratchBuffer(num_counts, GetComputeStream(context)); + blocked_row_to_unpermuted_row = GetScratchBuffer( + SafeInt(num_experts) * num_rows, GetComputeStream(context)); + ck::threeStepBuildExpertMapsSortFirstToken( + expert_indices, p_exp, p_r2u, unpermuted_row_to_permuted_row, p_efto, + blocked_expert_counts.get(), blocked_expert_counts_cumsum.get(), + blocked_row_to_unpermuted_row.get(), num_rows, num_experts, k_, 0, stream); + } } const void* fc1_bias = fc1_experts_bias_optional ? fc1_experts_bias_optional->DataRaw() : nullptr; @@ -2336,6 +2366,10 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { bool fc1_used_decode = false; bool fc2_used_decode = false; + // FC1 -> FC2 -> finalize overlap through programmatic dependent launch (row-major NVFP4 decode kernels only): + // FC2 prefetches its weights and finalize its routing metadata while the previous kernel finishes. + const bool fc2_finalize_pdl = + fp4_gemv_pdl && use_raw_nvfp4_gemv && nvfp4_weights_row_major_ && sm_ >= 90; auto run_fused = [&](auto* t_ptr) { using T = std::remove_pointer_t; const bool skip_expand = fp4_gemv_skip_expand_; @@ -2346,7 +2380,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { quant_params, false, p_efto, nullptr, nullptr, nullptr, stream); } - auto launch_fc1 = [&](MoeGemvConfig cfg) { + auto launch_fc1 = [&](MoeGemvConfig cfg, bool pdl = false) { gemv::launch_moe_gemv_fp4_symmetric_interleaved_swiglu( skip_expand ? static_cast(input->DataRaw()) : static_cast(p_act_buf.get()), gemv_fc1_weight, @@ -2355,9 +2389,9 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { static_cast(fc1_bias), static_cast(p_fc1_buf.get()), p_efto, p_exp, num_experts, expanded, inter, hidden, gemv_group_size, sm_, act_params, cfg, fc1_gemv_sm80_layout, use_raw_nvfp4_gemv, skip_expand ? p_r2u : nullptr, num_rows, stream, - nvfp4_weights_row_major_, nvfp4_fc1_decode_scale_eligible_, &fc1_used_decode); + nvfp4_weights_row_major_, nvfp4_fc1_decode_scale_eligible_, &fc1_used_decode, pdl); }; - auto launch_fc2 = [&](MoeGemvConfig cfg) { + auto launch_fc2 = [&](MoeGemvConfig cfg, bool pdl = false) { gemv::launch_moe_gemv_fp4_symmetric( static_cast(p_fc1_buf.get()), gemv_fc2_weight, @@ -2366,7 +2400,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { static_cast(fc2_bias), static_cast(p_fc2_buf.get()), p_efto, p_exp, num_experts, expanded, hidden, inter, gemv_group_size, sm_, cfg, fc2_gemv_sm80_layout, use_raw_nvfp4_gemv, stream, nvfp4_weights_row_major_, - nvfp4_fc2_decode_scale_eligible_, &fc2_used_decode); + nvfp4_fc2_decode_scale_eligible_, &fc2_used_decode, pdl); }; if (do_tune) { @@ -2451,13 +2485,13 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } } - launch_fc1(fc1_config); - launch_fc2(fc2_config); + launch_fc1(fc1_config, !do_tune && fused_route_launched && skip_expand && fp4_gemv_fc1_pdl); + launch_fc2(fc2_config, fc2_finalize_pdl && fc1_used_decode); ck::finalizeMoeRoutingKernelLauncher( static_cast(p_fc2_buf.get()), static_cast(output->MutableDataRaw()), nullptr, expert_scales, unpermuted_row_to_permuted_row, p_r2u, expert_indices, p_efto, num_rows, hidden, static_cast(k_), num_experts, - parallelism_config, false, stream); + parallelism_config, false, stream, fc2_finalize_pdl && fc2_used_decode); }; if (is_fp16_) { diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuh b/onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuh new file mode 100644 index 0000000000000..a615a5b922f6a --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuh @@ -0,0 +1,223 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "core/providers/cuda/cu_inc/topk_warp_sort.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// Decode routing in one launch: softmax + top-k (bit-identical to SoftmaxTopKMergeKernel) followed by the +// (expert, token)-sorted expert maps of fusedBuildExpertMapsSortFirstTokenKernel. +// +// One warp per row. Lane l emulates the four "threads" v = l + 32 w (w = 0..3) of the 128-thread merge kernel, each of +// which owns logits [4v, 4v + 4): the max / exp-sum reduction trees are reproduced operation for operation +// (cub::BlockReduce: ascending warp shuffle tree per 32 threads, the four warp aggregates summed in order), +// but entirely in registers. The top-k selection is exact on unique keys (score, lower index first), so the selection +// algorithm cannot change the result: logits are fp16 / bf16 values, so the sortable 32-bit score has 13 zero low bits +// and a 32-bit key (score | inverted index) orders exactly like the 64-bit key of the merge kernel. +constexpr int kFusedRouteMaxRows = 4; +constexpr int kFusedRouteMaxK = 16; +constexpr int kFusedRouteExperts = 512; // the kernel is specialised for exactly 512 experts +constexpr int kFusedRouteThreads = 128; +static_assert(kFusedRouteMaxRows * kFusedRouteMaxK <= 64); + +__device__ __forceinline__ uint32_t FusedRouteKey(float logit, int index, bool valid) { + const uint32_t bits = logit == 0.0f ? 0u : __float_as_uint(logit); + const uint32_t sortable = (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); + return valid ? ((sortable & 0xFFFFE000u) | (0x1FFFu - static_cast(index))) : 0u; +} + +__device__ __forceinline__ float FusedRouteScore(uint32_t key) { + const uint32_t sortable = key & 0xFFFFE000u; + return __uint_as_float((sortable & 0x80000000u) ? (sortable & 0x7FFFFFFFu) : (~sortable & 0xFFFFE000u)); +} + +__device__ __forceinline__ void FusedRouteLoad4(const half* p, float (&v)[4]) { + const uint2 raw = *reinterpret_cast(p); + const half2 a = *reinterpret_cast(&raw.x); + const half2 b = *reinterpret_cast(&raw.y); + v[0] = __low2float(a); + v[1] = __high2float(a); + v[2] = __low2float(b); + v[3] = __high2float(b); +} + +__device__ __forceinline__ void FusedRouteLoad4(const __nv_bfloat16* p, float (&v)[4]) { + const uint2 raw = *reinterpret_cast(p); + v[0] = __uint_as_float(raw.x << 16); + v[1] = __uint_as_float(raw.x & 0xFFFF0000u); + v[2] = __uint_as_float(raw.y << 16); + v[3] = __uint_as_float(raw.y & 0xFFFF0000u); +} + +// K > 0: top-k count known at compile time (fully unrolled selection); K == 0: runtime k. +template +__global__ void __launch_bounds__(kFusedRouteThreads) FusedSoftmaxTopKMapsKernel( + const T* logits, float* topk_scales, int* topk_indices, int* permuted_row_to_unpermuted_row, + int* unpermuted_row_to_permuted_row, int* permuted_expert, int64_t* expert_first_token_offset, int num_rows, + int k_runtime, bool normalize_scales) { + constexpr int num_experts = kFusedRouteExperts; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Programmatic dependent launch: the router logits come from the previous kernel; let the expert GEMV kernel + // start its prologue right away (it waits for this grid to complete before reading the maps). + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + constexpr int kMaxSel = K > 0 ? K : kFusedRouteMaxK; + const int k = K > 0 ? K : k_runtime; + __shared__ int s_key[kFusedRouteMaxRows * kFusedRouteMaxK]; + __shared__ int s_sorted[128]; + + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + s_sorted[threadIdx.x] = INT_MAX; // padding for the branchless searches below + + if (warp < num_rows) { + const T* row_logits = logits + static_cast(warp) * num_experts; // 8-byte aligned (checked by the launcher) + float scores[16]; + uint32_t keys[16]; + float local_max = onnxruntime::cuda::topk::kNegativeInfinity; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const int base = (lane + 32 * w) * 4; + float v[4]; + FusedRouteLoad4(row_logits + base, v); +#pragma unroll + for (int j = 0; j < 4; ++j) { + scores[w * 4 + j] = v[j] == 0.0f ? 0.0f : v[j]; + keys[w * 4 + j] = FusedRouteKey(v[j], base + j, true); + local_max = fmaxf(local_max, v[j]); + } + } + // Max over all experts (order independent: the per-lane values are never NaN). +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const float other = __shfl_xor_sync(0xFFFFFFFFu, local_max, offset); + local_max = (local_max < other) ? other : local_max; + } + const float max_val = local_max; + + // k rounds of an exact warp-wide arg-max (independent of the softmax sums below, so the scheduler overlaps them). + uint32_t winners[kMaxSel]; +#pragma unroll + for (int r = 0; r < kMaxSel; ++r) { + if (K > 0 || r < k) { + uint32_t best = keys[0]; +#pragma unroll + for (int j = 1; j < 16; ++j) { + best = max(best, keys[j]); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + winners[r] = __reduce_max_sync(0xFFFFFFFFu, best); +#else +#pragma unroll + for (int offset = 16; offset >= 1; offset >>= 1) { + best = max(best, __shfl_xor_sync(0xFFFFFFFFu, best, offset)); + } + winners[r] = best; +#endif +#pragma unroll + for (int j = 0; j < 16; ++j) { + keys[j] = keys[j] == winners[r] ? 0u : keys[j]; + } + } else { + winners[r] = 0u; + } + } + + float part[4]; +#pragma unroll + for (int w = 0; w < 4; ++w) { + float sum = 0.0f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + sum += expf(scores[w * 4 + j] - max_val); + } + part[w] = sum; + } +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { +#pragma unroll + for (int w = 0; w < 4; ++w) { + part[w] += __shfl_xor_sync(0xFFFFFFFFu, part[w], offset); + } + } + // lane 0 holds the tree result that cub's shfl_down reduction leaves in lane 0 + const float block_sum = ((__shfl_sync(0xFFFFFFFFu, part[0], 0) + __shfl_sync(0xFFFFFFFFu, part[1], 0)) + + __shfl_sync(0xFFFFFFFFu, part[2], 0)) + + __shfl_sync(0xFFFFFFFFu, part[3], 0); + const float inv_sum = onnxruntime::cuda::topk::SafeInvSum(block_sum); + + uint32_t selected = 0; + float prob = 0.0f; + float scale_sum = 0.0f; +#pragma unroll + for (int r = 0; r < kMaxSel; ++r) { + if (K > 0 || r < k) { + const float p = onnxruntime::cuda::topk::SoftmaxScale(FusedRouteScore(winners[r]), max_val, inv_sum); + scale_sum += p; // sequential rank order, as thread 0 of the merge kernel + if (lane == r) { + selected = winners[r]; + prob = p; + } + } + } + const float denom = onnxruntime::cuda::topk::TopKNormalizeDenom(normalize_scales, scale_sum); + if (lane < k) { + const int expert = 0x1FFF - static_cast(selected & 0x1FFFu); + topk_scales[static_cast(warp) * k + lane] = normalize_scales ? (prob / denom) : prob; + topk_indices[static_cast(warp) * k + lane] = expert; + s_key[warp * k + lane] = (expert << 2) | warp; // unique per (expert, token); sorts by expert, then token + } + } + const int num_pairs = num_rows * k; + __syncthreads(); + + // Expert maps: pair p = (token p / k, rank p % k) is ranked by (expert, token), the order of the radix rank kernel. + if (static_cast(threadIdx.x) < num_pairs) { + const int p = static_cast(threadIdx.x); + const int key = s_key[p]; + int rank = 0; + for (int q = 0; q < num_pairs; ++q) { + rank += s_key[q] < key ? 1 : 0; + } + const int token = key & 3; + const int unpermuted_row = (p - token * k) * num_rows + token; + const int expert = key >> 2; + permuted_row_to_unpermuted_row[rank] = unpermuted_row; + unpermuted_row_to_permuted_row[unpermuted_row] = rank; + permuted_expert[rank] = expert; + s_sorted[rank] = expert; + } + __syncthreads(); + // expert_first_token_offset[e] = number of pairs routed to an expert < e (e = 0..num_experts) + int count[5]; +#pragma unroll + for (int i = 0; i < 5; ++i) { + const int e = static_cast(threadIdx.x) + i * kFusedRouteThreads; + count[i] = 0; +#pragma unroll + for (int step = 64; step >= 1; step >>= 1) { + count[i] += s_sorted[count[i] + step - 1] < e ? step : 0; + } + } +#pragma unroll + for (int i = 0; i < 5; ++i) { + const int e = static_cast(threadIdx.x) + i * kFusedRouteThreads; + if (e <= num_experts) { + expert_first_token_offset[e] = count[i]; + } + } +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu index 331351a6ac33b..e3210c7df02cb 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu @@ -3,6 +3,7 @@ // Licensed under the MIT License. #include "contrib_ops/cuda/moe/qmoe_kernels.h" +#include "contrib_ops/cuda/moe/qmoe_fused_route.cuh" #include #include @@ -393,6 +394,93 @@ __global__ void SoftmaxTopKWarpMergeKernel(const T* logits, float* topk_scales, } } +template +bool IsFusedRouteKernelSm90Available(KernelPtr kernel) { + static thread_local int cached_device = -1; + static thread_local KernelPtr cached_kernel = nullptr; + static thread_local bool available = false; + int device = -1; + if (cudaGetDevice(&device) != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } + if (device != cached_device || kernel != cached_kernel) { + int driver_version = 0; + int sm_major = 0; + cudaFuncAttributes attributes{}; + if (cudaDriverGetVersion(&driver_version) != cudaSuccess || + cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess || + cudaFuncGetAttributes(&attributes, kernel) != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } + available = driver_version >= 12000 && sm_major >= 9 && + attributes.ptxVersion >= 90 && attributes.binaryVersion >= 90; + cached_device = device; + cached_kernel = kernel; + } + return available; +} + +template +bool LaunchFusedRouteMapsT(const T* logits, float* topk_scales, int* topk_indices, + int* permuted_row_to_unpermuted_row, int* unpermuted_row_to_permuted_row, + int* permuted_expert, int64_t* expert_first_token_offset, int num_rows, + int num_experts, int k, bool normalize_scales, bool pdl, cudaStream_t stream) { + // k is limited to the values the single-CTA radix-rank map builder supports (1, 2, 4, 6, 8, 10): the fused path then + // replaces exactly the launches it is bit-identical to; other k keep the generic route. + const bool supported_k = k == 1 || k == 2 || k == 4 || k == 6 || k == 8 || k == 10; + if (num_rows < 1 || num_rows > kFusedRouteMaxRows || num_experts != kFusedRouteExperts || !supported_k || + reinterpret_cast(logits) % 8 != 0) { + return false; + } + const auto kernel = k == 10 ? FusedSoftmaxTopKMapsKernel : FusedSoftmaxTopKMapsKernel; + if (!IsFusedRouteKernelSm90Available(kernel)) { + return false; + } + if (pdl) { + cudaLaunchConfig_t config{}; + config.gridDim = dim3(1); + config.blockDim = dim3(kFusedRouteThreads); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]{}; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attribute; + config.numAttrs = 1; + const cudaError_t status = cudaLaunchKernelEx( + &config, kernel, logits, topk_scales, topk_indices, permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, permuted_expert, expert_first_token_offset, num_rows, k, normalize_scales); + if (status == cudaSuccess) { + return true; + } + if (status != cudaErrorInvalidValue && status != cudaErrorNotSupported) { + ORT_ENFORCE(status == cudaSuccess, "fused QMoE routing launch failed: ", cudaGetErrorString(status)); + } + (void)cudaGetLastError(); + } + kernel<<<1, kFusedRouteThreads, 0, stream>>>( + logits, topk_scales, topk_indices, permuted_row_to_unpermuted_row, + unpermuted_row_to_permuted_row, permuted_expert, expert_first_token_offset, num_rows, k, normalize_scales); + CUDA_CALL_THROW(cudaGetLastError()); + return true; +} + +bool LaunchFusedRouteMaps(const half* logits, float* topk_scales, int* topk_indices, int* r2u, int* u2r, int* expert, + int64_t* offsets, int num_rows, int num_experts, int k, bool normalize_scales, bool pdl, + cudaStream_t stream) { + return LaunchFusedRouteMapsT(logits, topk_scales, topk_indices, r2u, u2r, expert, offsets, num_rows, + num_experts, k, normalize_scales, pdl, stream); +} + +bool LaunchFusedRouteMaps(const __nv_bfloat16* logits, float* topk_scales, int* topk_indices, int* r2u, int* u2r, + int* expert, int64_t* offsets, int num_rows, int num_experts, int k, bool normalize_scales, + bool pdl, cudaStream_t stream) { + return LaunchFusedRouteMapsT<__nv_bfloat16>(logits, topk_scales, topk_indices, r2u, u2r, expert, offsets, num_rows, + num_experts, k, normalize_scales, pdl, stream); +} + template void DispatchSoftmaxTopK(const T* logits, float* topk_scales, int* topk_indices, int num_rows, int num_experts, int k, bool normalize_scales, diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h index 00a49bddbe72d..cb911b9bdc2ed 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h @@ -42,6 +42,17 @@ void LaunchSoftmaxTopK( bool normalize_scales, cudaStream_t stream); +// Decode-only fused routing: softmax + top-k (bit-identical to LaunchSoftmaxTopK) and the (expert, token)-sorted expert +// maps of fusedBuildExpertMapsSortFirstToken in one single-CTA launch. Returns false (nothing launched) outside its +// envelope: 1..4 rows, exactly 512 experts, k in {1, 2, 4, 6, 8, 10}, 8-byte aligned logits. +bool LaunchFusedRouteMaps(const half* logits, float* topk_scales, int* topk_indices, int* permuted_row_to_unpermuted_row, + int* unpermuted_row_to_permuted_row, int* permuted_expert, int64_t* expert_first_token_offset, + int num_rows, int num_experts, int k, bool normalize_scales, bool pdl, cudaStream_t stream); +bool LaunchFusedRouteMaps(const __nv_bfloat16* logits, float* topk_scales, int* topk_indices, + int* permuted_row_to_unpermuted_row, int* unpermuted_row_to_permuted_row, + int* permuted_expert, int64_t* expert_first_token_offset, int num_rows, int num_experts, + int k, bool normalize_scales, bool pdl, cudaStream_t stream); + void LaunchSparseMixerTop2( const float* input, float* output, diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu index 152a8c5914894..052449ecaca45 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu @@ -26,40 +26,16 @@ namespace onnxruntime { namespace contrib { namespace cuda { -// The specialized router GEMV kernel only handles M=1, or batch size 1. -constexpr int kRouterM = 1; - -// MoE router shape (N = number of experts, K = hidden size). The specialization is exact-shape -// gated to the GPT-OSS-20B router projection to keep the dispatch change conservative. -constexpr int kGptOssRouterN = 32; -constexpr int kGptOssRouterK = 2880; - -static bool IsRouterGemvSpecializationDisabled() { - // Use ORT's cross-platform env var helper instead of std::getenv, which is unsafe on Windows. +static bool IsSingleRowGemvSpecializationDisabled() { + // Keep the existing cross-platform control for both MoE router and language-model projections. return ParseEnvironmentVariableWithDefault("ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", false); } -// The router GEMV kernel handles any symmetric (no zero point) M=1 shape with an int4 group size of -// 32 or 64 (whichever quantizes best) and N divisible by kColsPerThreadBlock. We gate on the exact -// GPT-OSS-20B router shape to avoid changing the dispatch for general MatMulNBits cases. K must be a -// multiple of the group size (always true for a router) and N a multiple of kColsPerThreadBlock -// (checked in TryMatMul4Bits). Note kPerIter (256) is divisible by both 32 and 64, so the scale -// stride is exact. -static bool IsSupportedRouterGemvShape(const uint8_t* zero_points, int m, int n, int k, int block_size) { - if (zero_points != nullptr || m != kRouterM || (block_size != 32 && block_size != 64)) { - return false; - } - if (k % block_size != 0) { - return false; - } - return n == kGptOssRouterN && k == kGptOssRouterK; // gpt-oss-20b -} - // Reduces kUnroll groups of kPerIter elements per step, advancing the packed-weight pointer, scale -// index and k position in lockstep. Factored out of MatMulFloatInt4RouterKernel so the three unroll +// index and k position in lockstep. Factored out of MatMulFloatInt4SingleRowKernel so the three unroll // factors (16, 4, 1) share one implementation instead of a macro. template -__device__ __forceinline__ void RouterUnrollReduction( +__device__ __forceinline__ void SingleRowUnrollReduction( const uint8_t*& b_data_quant, const T* scales_data, const T* a_data, @@ -82,14 +58,15 @@ __device__ __forceinline__ void RouterUnrollReduction( } } -// GEMV specialization for MoE routers: output(1, N) = a(1, K) x dequant(B(N, K)) [+ bias(N)]. +// Single-row projection GEMV: output(1, N) = a(1, K) x dequant(B(N, K)) [+ bias(N)]. // B is 4-bit block-quantized (symmetric, no zero point) with group size BlockSize (32 or 64). One warp -// computes one expert column; the thread block holds kColsPerThreadBlock warps. N is passed via the -// grid and K at runtime, so a single instantiation per (T, BlockSize) serves every router shape. +// computes one output column; the thread block holds kColsPerThreadBlock warps. Dispatch is exact-shape +// gated to the GPT-OSS-20B router (N=32, K=2880, groups 32/64) and Qwen draft head +// (N=248320, K=2560, group 32). A single instantiation per (T, BlockSize) serves both workloads. // Requirements (satisfied by the dispatch in TryMatMul4Bits): N % kColsPerThreadBlock == 0, // K % BlockSize == 0, and kPerIter % BlockSize == 0 so the per-iteration scale stride is exact. template -__global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt4RouterKernel( +__global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt4SingleRowKernel( T* output, const T* a_data, const uint8_t* b_data_quant, @@ -113,9 +90,9 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt int k_id = 0; int scale_id = lane_id * 8 / BlockSize; - RouterUnrollReduction(b_data_quant, scales_data, a_data, k, k_id, scale_id, sums); - RouterUnrollReduction(b_data_quant, scales_data, a_data, k, k_id, scale_id, sums); - RouterUnrollReduction(b_data_quant, scales_data, a_data, k, k_id, scale_id, sums); + SingleRowUnrollReduction(b_data_quant, scales_data, a_data, k, k_id, scale_id, sums); + SingleRowUnrollReduction(b_data_quant, scales_data, a_data, k, k_id, scale_id, sums); + SingleRowUnrollReduction(b_data_quant, scales_data, a_data, k, k_id, scale_id, sums); if (k_id + lane_id * 8 < k) { uint32_t value = *(reinterpret_cast(b_data_quant)); @@ -707,15 +684,15 @@ bool TryMatMul4Bits( return false; } - if (IsSupportedRouterGemvShape(zero_points, m, n, k, block_size) && - !IsRouterGemvSpecializationDisabled()) { + if (IsSupportedSingleRowGemvShape(zero_points != nullptr, m, n, k, block_size) && + !IsSingleRowGemvSpecializationDisabled()) { const dim3 blocks(n / kColsPerThreadBlock, 1); const dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); if (block_size == 32) { - MatMulFloatInt4RouterKernel<<>>( + MatMulFloatInt4SingleRowKernel<<>>( output, a_data, b_data_quant, scales_data, bias_data, n, k); } else { - MatMulFloatInt4RouterKernel<<>>( + MatMulFloatInt4SingleRowKernel<<>>( output, a_data, b_data_quant, scales_data, bias_data, n, k); } return true; diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index e458dd438f8fd..f866641b34146 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -1208,8 +1208,10 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { const int rows = std::min(chunk_m, m - row_start); // Only a trailing partial chunk can change rows, so this looks up at most two tactics. if (rows != tactic_m) { - bestTactic = stream_is_capturing ? gemmProfiler_->getBestConfig(rows, gemmId_) - : gemmProfiler_->getBestConfigOrProfile(rows, gemmId_); + bestTactic = ctx->GetUseDeterministicCompute() + ? gemmProfiler_->getDeterministicConfig(rows) + : (stream_is_capturing ? gemmProfiler_->getBestConfig(rows, gemmId_) + : gemmProfiler_->getBestConfigOrProfile(rows, gemmId_)); if (!bestTactic.has_value()) { return ORT_MAKE_STATUS( ONNXRUNTIME, FAIL, diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh index eab444df8a2d2..9293163fafde8 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh @@ -8,6 +8,14 @@ namespace onnxruntime { namespace contrib { namespace cuda { +constexpr bool IsSupportedSingleRowGemvShape(bool has_zero_points, int m, int n, int k, int block_size) { + if (has_zero_points || m != 1 || (block_size != 32 && block_size != 64)) { + return false; + } + return (n == 32 && k == 2880) || + (n == 248320 && k == 2560 && block_size == 32); +} + inline constexpr int kMatMul8BitsDefaultMaxRows = 5; inline constexpr int kMatMul8BitsSm121Fp16MaxRows = 8; diff --git a/onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.cc b/onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.cc new file mode 100644 index 0000000000000..0da84e807aa8a --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.cc @@ -0,0 +1,205 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/webgpu/bert/conv_state_replay.h" + +#include +#include +#include + +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" + +namespace onnxruntime::contrib::webgpu { + +using namespace onnxruntime::webgpu; + +ONNX_OPERATOR_KERNEL_EX( + ConvStateReplay, + kMSDomain, + 1, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedFloatTypes()) + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) + .InputMemoryType(OrtMemTypeCPUInput, 3) + .MayInplace(2, 0), + ConvStateReplay); + +namespace { + +constexpr uint32_t kWorkgroupSize = 128; +constexpr uint64_t kMaxIndex = std::numeric_limits::max(); + +class ConvStateReplayProgram final : public Program { + public: + ConvStateReplayProgram(bool half, bool has_capture) + : Program{"ConvStateReplay"}, half_{half}, has_capture_{has_capture} {} + + Status GenerateShaderCode(ShaderHelper& shader) const override { + const auto& source = shader.AddInput("source", ShaderUsage::UseUniform); + const ShaderVariableHelper* capture = nullptr; + if (has_capture_) { + capture = &shader.AddInput("capture", ShaderUsage::UseUniform); + } + const auto& destination = shader.AddOutput("destination", ShaderUsage::UseUniform); + return WGSL_TEMPLATE_APPLY(shader, "bert/conv_state_replay.wgsl.template", + WGSL_TEMPLATE_PARAMETER(half, half_), + WGSL_TEMPLATE_PARAMETER(has_capture, has_capture_), + WGSL_TEMPLATE_OPTIONAL_VARIABLE(capture, capture), + WGSL_TEMPLATE_VARIABLE(destination, destination), + WGSL_TEMPLATE_VARIABLE(source, source)); + } + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"state_elements", ProgramUniformVariableDataType::Uint32}, + {"output_words", ProgramUniformVariableDataType::Uint32}, + {"channels", ProgramUniformVariableDataType::Uint32}, + {"history_width", ProgramUniformVariableDataType::Uint32}, + {"capacity", ProgramUniformVariableDataType::Uint32}, + {"kept_count", ProgramUniformVariableDataType::Uint32}, + {"source_lane", ProgramUniformVariableDataType::Uint32}, + {"destination_lane", ProgramUniformVariableDataType::Uint32}, + {"capture_lane", ProgramUniformVariableDataType::Uint32}); + + private: + const bool half_; + const bool has_capture_; +}; + +Status EnqueueReplay(onnxruntime::webgpu::ComputeContext& context, const Tensor& source, const Tensor& capture, + Tensor& destination, const std::array& desc, + uint32_t state_elements, uint32_t capture_elements, + uint32_t output_words, uint32_t dispatch_x, uint32_t dispatch_y) { + const bool half = source.IsDataType(); + const uint32_t element_bytes = half ? 2 : 4; + const uint32_t elements_per_word = 4 / element_bytes; + const bool has_capture = desc[7] > 0; + // WebGPU allocations are rounded to four bytes. Reinterpret the entire allocation, never + // do pointer arithmetic on an opaque buffer handle or load/store floating-point values. + const auto word_shape = [](const Tensor& tensor) { + return TensorShape{static_cast((tensor.SizeInBytes() + 3) / 4)}; + }; + Tensor source_words(DataTypeImpl::GetType(), word_shape(source), + const_cast(source.DataRaw()), source.Location()); + Tensor capture_words(DataTypeImpl::GetType(), word_shape(capture), + const_cast(capture.DataRaw()), capture.Location()); + Tensor destination_words(DataTypeImpl::GetType(), word_shape(destination), + destination.MutableDataRaw(), destination.Location()); + const auto view_words = [elements_per_word, element_bytes](uint32_t offset, uint32_t count) { + return TensorShape{static_cast( + (static_cast(offset % 4 / element_bytes) + count + elements_per_word - 1) / elements_per_word)}; + }; + ConvStateReplayProgram program{half, has_capture}; + program.AddInput(ProgramInput::BufferView(&source_words, ProgramTensorMetadataDependency::Type, + view_words(desc[0], state_elements), desc[0] / 4)); + if (has_capture) { + program.AddInput(ProgramInput::BufferView(&source == &capture ? &source_words : &capture_words, + ProgramTensorMetadataDependency::Type, + view_words(desc[2], capture_elements), desc[2] / 4)); + } + program.AddOutput(ProgramOutput::BufferView(&destination_words, ProgramTensorMetadataDependency::Type, + TensorShape{output_words}, desc[1] / 4)) + .CacheHint(half, has_capture) + .SetDispatchGroupSize(dispatch_x, dispatch_y) + .SetWorkgroupSize(kWorkgroupSize) + .AddUniformVariables({{state_elements}, + {output_words}, + {desc[4]}, + {desc[5]}, + {desc[6]}, + {desc[7]}, + {desc[0] % 4 / element_bytes}, + {desc[1] % 4 / element_bytes}, + {desc[2] % 4 / element_bytes}}); + + const auto& limits = context.DeviceLimits(); + const auto segments = [&limits](const Tensor& tensor) { + return (static_cast(tensor.SizeInBytes()) - 1) / limits.maxStorageBufferBindingSize + 1; + }; + uint64_t bindings = segments(destination_words); + for (size_t index = 0; index < program.Inputs().size(); ++index) { + if (program.InputBufferOwner(index) == index) { + bindings += segments(*program.Inputs()[index].tensor); + } + } + ORT_RETURN_IF_NOT(bindings <= limits.maxStorageBuffersPerShaderStage && + bindings + 1 <= limits.maxBindingsPerBindGroup, + "ConvStateReplay backing tensors exceed WebGPU binding limits."); + return context.RunProgram(program); +} + +} // namespace + +Status ConvStateReplay::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + ORT_RETURN_IF(context.IsGraphCaptureEnabled(), "ConvStateReplay does not support graph capture."); + const auto* source = context.Input(0); + const auto* capture = context.Input(1); + const auto* destination = context.Input(2); + const auto* metadata = context.Input(3); + ORT_RETURN_IF_NOT(source && capture && destination && metadata, "ConvStateReplay requires all four inputs."); + for (const auto* backing : {source, capture, destination}) { + ORT_RETURN_IF_NOT((backing->IsDataType() || backing->IsDataType()) && + backing->GetElementType() == source->GetElementType() && + backing->Shape().NumDimensions() == 1, + "ConvStateReplay backing tensors must be rank-1 with the same FP16 or FP32 type."); + ORT_RETURN_IF_NOT(static_cast(backing->SizeInBytes()) <= kMaxIndex, + "ConvStateReplay backing extent exceeds the uint32 indexing range."); + } + ORT_RETURN_IF_NOT(metadata->IsDataType() && metadata->Shape() == TensorShape{8}, + "ConvStateReplay metadata must be CPU int64 with shape [8]."); + std::array desc{}; + for (size_t index = 0; index < desc.size(); ++index) { + const auto value = metadata->Data()[index]; + ORT_RETURN_IF_NOT(value >= 0 && static_cast(value) <= kMaxIndex, + "ConvStateReplay metadata field ", index, " is outside the uint32 indexing range."); + desc[index] = static_cast(value); + } + const uint64_t batch = desc[3], channels = desc[4], width = desc[5], capacity = desc[6]; + ORT_RETURN_IF_NOT(channels > 0, "ConvStateReplay channels must be positive."); + ORT_RETURN_IF_NOT(desc[7] <= capacity, "ConvStateReplay kept_count must be in [0, capacity]."); + const uint32_t element_bytes = source->IsDataType() ? 2 : 4; + ORT_RETURN_IF_NOT(batch * channels <= kMaxIndex, + "ConvStateReplay geometry exceeds the uint32 indexing range."); + const uint64_t state_elements = batch * channels * width; + const uint64_t capture_elements = batch * channels * capacity; + ORT_RETURN_IF_NOT(state_elements <= kMaxIndex / element_bytes && + capture_elements <= kMaxIndex / element_bytes, + "ConvStateReplay geometry exceeds the uint32 indexing range."); + const auto check_view = [element_bytes](const Tensor& backing, uint32_t offset, uint64_t count) -> Status { + ORT_RETURN_IF_NOT(offset % element_bytes == 0, "ConvStateReplay byte offsets must be element-aligned."); + ORT_RETURN_IF_NOT(static_cast(offset) + count * element_bytes <= backing.SizeInBytes(), + "ConvStateReplay 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(*capture, desc[2], capture_elements)); + ORT_RETURN_IF(destination->DataRaw() && + (destination->DataRaw() == source->DataRaw() || destination->DataRaw() == capture->DataRaw()), + "ConvStateReplay destination backing must be distinct from source and capture."); + auto* output = context.Output(0, destination->Shape()); + ORT_RETURN_IF_NOT(output && output->MutableDataRaw() == destination->DataRaw(), + "ConvStateReplay output must be prebound to destination_backing."); + if (state_elements == 0) { + return Status::OK(); + } + + const uint64_t output_words = (desc[1] % 4 + state_elements * element_bytes + 3) / 4; + const auto& limits = context.DeviceLimits(); + const uint64_t groups = (output_words + kWorkgroupSize - 1) / kWorkgroupSize; + const uint64_t dispatch_x = std::min(groups, static_cast(limits.maxComputeWorkgroupsPerDimension)); + const uint64_t dispatch_y = (groups + dispatch_x - 1) / dispatch_x; + ORT_RETURN_IF_NOT(dispatch_y <= limits.maxComputeWorkgroupsPerDimension && + dispatch_x * dispatch_y * kWorkgroupSize <= kMaxIndex + 1, + "ConvStateReplay dispatch exceeds the supported indexing range."); + return context.RunAndWaitChecked([&]() { + return EnqueueReplay(context, *source, *capture, *output, desc, + static_cast(state_elements), static_cast(capture_elements), + static_cast(output_words), static_cast(dispatch_x), + static_cast(dispatch_y)); + }); +} + +} // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.h b/onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.h new file mode 100644 index 0000000000000..b337e436b6fc6 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/conv_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 ConvStateReplay final : public onnxruntime::webgpu::WebGpuKernel { + public: + explicit ConvStateReplay(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/conv_state_replay.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.wgsl.template new file mode 100644 index 0000000000000..4c425f0a079fe --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.wgsl.template @@ -0,0 +1,57 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#use .getByOffset .setByOffset +#param half +#param has_capture + +fn replay_element(index: u32) -> u32 { + let position = index % uniforms.history_width; + let channel = (index / uniforms.history_width) % uniforms.channels; + let batch = index / (uniforms.history_width * uniforms.channels); + if (uniforms.kept_count < uniforms.history_width && + position < uniforms.history_width - uniforms.kept_count) { + let source_index = index + uniforms.kept_count + uniforms.source_lane; +#if half + return (source.getByOffset(source_index / 2u) >> ((source_index % 2u) * 16u)) & 0xffffu; +#else + return source.getByOffset(source_index); +#endif + } +#if has_capture + var token: u32; + if (uniforms.kept_count >= uniforms.history_width) { + token = uniforms.kept_count - uniforms.history_width + position; + } else { + token = position - (uniforms.history_width - uniforms.kept_count); + } + let capture_index = (batch * uniforms.capacity + token) * uniforms.channels + channel + uniforms.capture_lane; +#if half + return (capture.getByOffset(capture_index / 2u) >> ((capture_index % 2u) * 16u)) & 0xffffu; +#else + return capture.getByOffset(capture_index); +#endif +#else + return 0u; +#endif +} + +$MAIN { + if (global_idx >= uniforms.output_words) { + return; + } +#if half + // One invocation owns a physical word, including both halves across logical row boundaries. + var bits = destination.getByOffset(global_idx); + for (var lane = 0u; lane < 2u; lane++) { + let element = global_idx * 2u + lane; + if (element >= uniforms.destination_lane && element - uniforms.destination_lane < uniforms.state_elements) { + let shift = lane * 16u; + bits = (bits & ~(0xffffu << shift)) | (replay_element(element - uniforms.destination_lane) << shift); + } + } + destination.setByOffset(global_idx, bits); +#else + destination.setByOffset(global_idx, replay_element(global_idx)); +#endif +} diff --git a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index 426d4057be6f4..d1cf54166e658 100644 --- a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc @@ -3,6 +3,7 @@ #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" #include "contrib_ops/webgpu/bert/causal_conv_with_state.h" +#include "contrib_ops/webgpu/bert/conv_state_replay.h" #include "contrib_ops/webgpu/bert/dynamic_sparse_attention.h" #include "contrib_ops/webgpu/bert/varlen_causal_conv_with_state.h" #include "contrib_ops/webgpu/bert/engram_gate.h" @@ -40,6 +41,7 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 829d99683d7bb..c4a83a9a5fdd8 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -2,15 +2,22 @@ // Licensed under the MIT License. #include "core/framework/allocation_planner.h" -#include +#if !defined(ORT_MINIMAL_BUILD) +#include "core/framework/allocation_planner_reuse_index.h" +#else #include +#endif +#include +#include #include #include #include #include #include #include +#include #include "core/common/exceptions.h" +#include "core/common/hash_combine.h" #include "core/common/inlined_containers.h" #include "core/common/safeint.h" #include "core/platform/env.h" @@ -208,6 +215,9 @@ class PlannerImpl { // Removing initializers is a temporary measure needed to limit the number of copies of // tensors in GPU memory. OrtValueIndex reused_buffer_index = -1; // index of original buffer to reuse +#if !defined(ORT_MINIMAL_BUILD) + OrtValueIndex next_free_buffer = -1; +#endif bool is_inplace_reuse = false; #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) OrtValueIndex inplace_reused_buffer_index = -1; // index of original buffer to reuse inplace @@ -217,6 +227,11 @@ class PlannerImpl { // ort_value_info_ is indexed by an OrtValueIndex std::vector ort_value_info_; + // Full builds use the reuse index. All minimal variants retain the baseline LIFO linear + // algorithm and do not receive the indexed planner's startup-time improvement. +#if !defined(ORT_MINIMAL_BUILD) + allocation_planner_internal::FreeBufferIndex freelists_; +#else // FreeBufferInfo is used to track information about ml-values whose buffers are // free to be reused. struct FreeBufferInfo { @@ -231,6 +246,7 @@ class PlannerImpl { // freelist_ : a list of ml-values whose buffers are free to be reused, sorted by when // they became free (more recently freed earlier in the list). std::list freelist_; +#endif OrtValueIndex Index(const OrtValueName& name) { OrtValueIndex result; @@ -499,33 +515,28 @@ class PlannerImpl { /*! \brief Given a tensor-type, return the primitive element type of the tensor. */ - static MLDataType GetPrimitiveElementType(const DataType& tensor_type) { + static MLDataType GetPrimitiveElementType(DataType tensor_type) { MLDataType ml_data_type = DataTypeImpl::GetDataType(*tensor_type); const TensorTypeBase* tensor_type_base = ml_data_type->AsTensorType(); ORT_ENFORCE(nullptr != tensor_type_base); return tensor_type_base->GetElementType(); } - /*! \brief Given a tensor-type, return the size in bytes of the C++ carrier used for an element. - */ - static size_t GetElementSize(const DataType& tensor_type) { - return GetPrimitiveElementType(tensor_type)->Size(); - } + struct ElementStorage { + size_t element_size; + int32_t sub_element_count; + }; - /*! \brief Given a tensor-type, return how many logical (sub-byte) elements are packed into one - * carrier element. Returns 1 for regular types and >1 for packed sub-byte types (e.g. 2 for int4/uint4). - */ - static int32_t GetSubElemCount(const DataType& tensor_type) { - const auto* prim_type = GetPrimitiveElementType(tensor_type)->AsPrimitiveDataType(); - return prim_type != nullptr ? prim_type->GetNumSubElems() : 1; + // Share physical storage metadata between the reuse key and the final compatibility check. + static ElementStorage GetElementStorage(MLDataType element_type) { + const auto* prim_type = element_type->AsPrimitiveDataType(); + return {element_type->Size(), prim_type != nullptr ? prim_type->GetNumSubElems() : 1}; } static bool SameSize(const TensorShapeProto& shape1, const onnxruntime::NodeArg& arg1, const TensorShapeProto& shape2, const onnxruntime::NodeArg& arg2) { - const auto& ptype1 = arg1.Type(); - const auto& ptype2 = arg2.Type(); - auto type1_size = GetElementSize(ptype1); - auto type2_size = GetElementSize(ptype2); + const auto storage1 = GetElementStorage(GetPrimitiveElementType(arg1.Type())); + const auto storage2 = GetElementStorage(GetPrimitiveElementType(arg2.Type())); bool is_type1_string = arg1.TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING; bool is_type2_string = arg2.TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING; @@ -535,7 +546,7 @@ class PlannerImpl { // (e.g. uint4[1024] needs 512 bytes while uint8[1024] needs 1024 bytes). Reusing the smaller buffer for the // larger tensor produces a heap buffer overflow when the tensor is later written. Only treat the tensors as // the same size when the sub-element packing density also matches, which guarantees identical storage bytes. - if (GetSubElemCount(ptype1) != GetSubElemCount(ptype2)) { + if (storage1.sub_element_count != storage2.sub_element_count) { return false; } @@ -544,7 +555,7 @@ class PlannerImpl { // If either of the tensors is a string, don't treat them the same. Moreover, reusing a string tensor for a string // tensor without releasing the previous memory can cause memory leaks; hence we don't allow reuse across string // tensors as well. - return !(is_type1_string || is_type2_string) && (type1_size == type2_size) && SameShape(shape1, shape2); + return !(is_type1_string || is_type2_string) && (storage1.element_size == storage2.element_size) && SameShape(shape1, shape2); /* TODO: we can improve this if the concrete shapes are known for both as below. Unclear whether this is worthwhile though. @@ -578,17 +589,71 @@ class PlannerImpl { return SameSize(*p_shape1, arg1, *p_shape2, arg2); } - // Find if freelist contains a buffer of the same size as output_arg +#if !defined(ORT_MINIMAL_BUILD) + std::optional GetReuseKey(const NodeArg& arg, const TensorShapeProto& shape) { + if (shape.dim_size() == 0) return std::nullopt; + const auto* element_type = GetPrimitiveElementType(arg.Type()); + if (element_type == DataTypeImpl::GetType()) return std::nullopt; + + size_t key = AllocPlan(arg.Name()).location.Hash(); + const auto storage = GetElementStorage(element_type); + HashCombine(storage.element_size, key); + HashCombine(storage.sub_element_count, key); + HashCombine(shape.dim_size(), key); + for (const auto& dim : shape.dim()) { + if (utils::HasDimValue(dim)) { + HashCombine(false, key); + HashCombine(dim.dim_value(), key); + } else if (utils::HasDimParam(dim) && !dim.dim_param().empty()) { + HashCombine(true, key); + HashCombine(dim.dim_param(), key); + } else { + // SameShape cannot match an anonymous or empty symbolic dimension. + return std::nullopt; + } + } + return key; + } + + void AddFreeBuffer(OrtValueIndex value) { + if (!context_->GetEnableMemoryReuse()) return; + const auto* arg = ort_value_info_[value].p_def_site; + if (!arg || IsNonTensor(*arg)) return; + const auto* shape = context_->GetShape(*arg); + if (!shape) return; + if (auto key = GetReuseKey(*arg, *shape)) { + auto& head = freelists_.GetOrCreate(*key); + ort_value_info_[value].next_free_buffer = + head ? static_cast(*head) : -1; + head = static_cast(value); + } + } +#endif + + // Find a compatible buffer; full builds restrict the candidates with the reuse index. bool FindReusableTensor(const onnxruntime::NodeArg& output_arg, OrtValueIndex* reusable_tensor) { if (!context_->GetEnableMemoryReuse()) { return false; } auto p_required_buffer_shape = context_->GetShape(output_arg); if (nullptr == p_required_buffer_shape || p_required_buffer_shape->dim_size() == 0) return false; +#if !defined(ORT_MINIMAL_BUILD) + auto key = GetReuseKey(output_arg, *p_required_buffer_shape); + if (!key) return false; + auto* head = freelists_.Find(*key); + if (head == nullptr || !*head) return false; +#endif auto& required_memory_info = AllocPlan(output_arg.Name()).location; +#if !defined(ORT_MINIMAL_BUILD) + OrtValueIndex previous = -1; + for (auto value = static_cast(**head); value != -1; + previous = value, value = ort_value_info_[value].next_free_buffer) { + size_t reusable = static_cast(value); +#else for (auto it = freelist_.begin(); it != freelist_.end(); ++it) { size_t reusable = static_cast(it->ml_value); +#endif const onnxruntime::NodeArg* p_node_arg = ort_value_info_.at(reusable).p_def_site; if (!p_node_arg) { // TODO this should be an error case, needs more investigation @@ -610,8 +675,20 @@ class PlannerImpl { if (nullptr != p_available_buffer_shape) { if (SameSize(*p_available_buffer_shape, *p_node_arg, *p_required_buffer_shape, output_arg)) { +#if !defined(ORT_MINIMAL_BUILD) + *reusable_tensor = value; + auto next = ort_value_info_[reusable].next_free_buffer; + if (previous != -1) { + ort_value_info_[previous].next_free_buffer = next; + } else if (next != -1) { + *head = static_cast(next); + } else { + head->reset(); + } +#else *reusable_tensor = it->ml_value; freelist_.erase(it); +#endif return true; } } @@ -1540,7 +1617,11 @@ class PlannerImpl { #endif ORT_RETURN_IF_ERROR(ComputeSingleStreamReusePlan(i)); ClearUseCount(); +#if !defined(ORT_MINIMAL_BUILD) + freelists_.clear(); // DONOT share free buffers across streams +#else freelist_.clear(); // DONOT share freelist across streams +#endif } #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) CalculateLifetime(ort_value_usecount); @@ -1674,7 +1755,11 @@ class PlannerImpl { // The index will be -1 if it's an initializer that was removed as part of a temporary workaround. // See comments in the OrtValueInfo definition. if ((original != -1) && (0 == DecrementUseCount(original))) { +#if !defined(ORT_MINIMAL_BUILD) + AddFreeBuffer(original); +#else freelist_.push_front(FreeBufferInfo(original, program_counter)); +#endif } } } @@ -1694,7 +1779,11 @@ class PlannerImpl { // The index will be -1 if it's an initializer that was removed as part of a temporary workaround. // See comments in the OrtValueInfo definition. if ((original != -1) && (0 == DecrementUseCount(original))) { +#if !defined(ORT_MINIMAL_BUILD) + AddFreeBuffer(original); +#else freelist_.push_front(FreeBufferInfo(original, program_counter)); +#endif } } } @@ -1708,7 +1797,11 @@ class PlannerImpl { // The index will be -1 if it's an initializer that was removed as part of a temporary workaround. // See comments in the OrtValueInfo definition. if (0 == DecrementUseCount(original)) { +#if !defined(ORT_MINIMAL_BUILD) + AddFreeBuffer(original); +#else freelist_.push_front(FreeBufferInfo(original, program_counter)); +#endif } } } @@ -2427,7 +2520,7 @@ Status PlannerImpl::CreatePlan( ORT_RETURN_IF_ERROR(ComputeAllocationOrder()); #endif - // convert information in the freelist_ into a deallocation plan in required format + // convert buffer lifetime information into a deallocation plan in required format ORT_RETURN_IF_ERROR(GenerateDeallocationPlan()); // generate program counter diff --git a/onnxruntime/core/framework/allocation_planner_reuse_index.h b/onnxruntime/core/framework/allocation_planner_reuse_index.h new file mode 100644 index 0000000000000..33a66305d0366 --- /dev/null +++ b/onnxruntime/core/framework/allocation_planner_reuse_index.h @@ -0,0 +1,100 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include +#include +#include +#include +#include + +#include + +#include "core/common/common.h" +#include "core/common/safeint.h" + +namespace onnxruntime::allocation_planner_internal { + +// Planner-only full-word index. Head absence does not remove a key from a probe chain. +class FreeBufferIndex { + public: + FreeBufferIndex() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(FreeBufferIndex); + + std::optional* Find(size_t key) { + if (slots_.empty()) return nullptr; + auto& slot = slots_[FindSlot(slots_, key)]; + return slot.occupied ? &slot.head : nullptr; + } + + std::optional& GetOrCreate(size_t key) { + if (slots_.empty()) Grow(); + size_t index = FindSlot(slots_, key); + if (!slots_[index].occupied) { + // Leave at least one unused slot, including when occupied heads are consumed. + if (size_ == slots_.size() - slots_.size() / 4) { + Grow(); + index = FindSlot(slots_, key); + } + auto& slot = slots_[index]; + slot.key = key; + slot.head.reset(); + slot.occupied = true; + ++size_; + } + return slots_[index].head; + } + + void clear() { + for (auto& slot : slots_) slot.occupied = false; + size_ = 0; + } + + private: + friend struct FreeBufferIndexTestPeer; + + struct Slot { + size_t key = 0; + std::optional head; + bool occupied = false; + }; + + static size_t Hash(size_t key) { + static_assert(std::has_unique_object_representations_v); + // Hash every byte to disperse aligned dimensions; equality still compares the full word. + return std::hash{}( + std::string_view(reinterpret_cast(&key), sizeof(key))); + } + + static size_t FindSlot(gsl::span slots, size_t key) { + const size_t mask = slots.size() - 1; + size_t index = Hash(key) & mask; + while (slots[index].occupied && slots[index].key != key) index = (index + 1) & mask; + return index; + } + + static size_t NextCapacity(size_t capacity) { + if (capacity == 0) return 8; + return SafeInt(capacity) * 2; + } + + void Grow() { + std::vector next(NextCapacity(slots_.size())); + for (const auto& slot : slots_) { + if (slot.occupied) next[FindSlot(next, slot.key)] = slot; + } + slots_.swap(next); + } + + // A nonempty table exceeds inline-vector storage even at its initial capacity. + std::vector slots_; + size_t size_ = 0; +}; + +} // namespace onnxruntime::allocation_planner_internal + +#endif // !defined(ORT_MINIMAL_BUILD) diff --git a/onnxruntime/core/framework/resource_accountant.cc b/onnxruntime/core/framework/resource_accountant.cc index 39510f1202895..2777647b7096e 100644 --- a/onnxruntime/core/framework/resource_accountant.cc +++ b/onnxruntime/core/framework/resource_accountant.cc @@ -10,6 +10,7 @@ #include "core/common/string_utils.h" #include "core/framework/config_options.h" +#include "core/framework/kernel_dispatch_policy.h" #include "core/framework/max_shape_override.h" #include "core/framework/murmurhash3.h" #include "core/framework/tensorprotoutils.h" @@ -571,6 +572,20 @@ Status CreateAccountants( WorkspaceEstimatorConfig estimator_config{ config_options.GetConfigEntry(kOrtSessionOptionsCudaFpAIntBGemm), config_options.GetConfigEntry(kOrtSessionOptionsCudaFpAIntBProfileM)}; + estimator_config.kernel_dispatch_policy = + config_options.GetConfigEntry(kOrtSessionOptionsKernelDispatchPolicy); + if (estimator_config.kernel_dispatch_policy.has_value()) { + KernelDispatchPolicy parsed_policy = KernelDispatchPolicy::Auto; + if (!TryParseKernelDispatchPolicy( + *estimator_config.kernel_dispatch_policy, parsed_policy)) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + kOrtSessionOptionsKernelDispatchPolicy, + " must be empty or one of 'auto', 'latency', 'memory', or 'safe' " + "(case-sensitive), but got: '", + *estimator_config.kernel_dispatch_policy, "'"); + } + } ORT_RETURN_IF_ERROR(ParseCudaGqaWorkspaceMaxTotalSequenceLength( config_options, estimator_config.cuda_gqa_workspace_max_total_sequence_length)); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index fecac2f17f523..1fb0dfb3304c0 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -1902,12 +1902,17 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "S") .Input(6, "past_seqlens", - "A tensor with shape (batch_size). It specifies the past lengths of cached sequence in the KV cache.", + "A tensor with shape (batch_size). It specifies the past lengths of cached sequences. " + "CUDA sanitizes malformed values on device: an invalid value suppresses that sequence's cache writes " + "and produces zero output. A backend requiring exact host lengths may instead return INVALID_ARGUMENT " + "when attention_metadata is absent.", "S") .Input(7, "block_table", "2D tensor with shape (batch_size, max_blocks_per_sequence) that maps each sequence in the batch to its" - "corresponding blocks in the KV cache.", + "corresponding blocks in the KV cache. -1 denotes an unmapped page. CUDA converts values below -1 or " + "greater than or equal to num_blocks to the unmapped sentinel before resolving cache addresses; mapped " + "pages remain attendable.", "S") .Input(8, "cos_cache", @@ -5285,6 +5290,63 @@ their BufferViews have the same Program binding owner, not merely equal raw buff } })); +ONNX_MS_OPERATOR_SET_SCHEMA( + ConvStateReplay, 1, + OpSchema() + .SetDoc(R"DOC( +Reconstruct rolling convolution history after retaining a prefix of captured inputs, without +performing convolution arithmetic. The WebGPU-only implementation copies FP16 or FP32 payload +bits exactly, including signed zero and NaN payloads, and waits for checked GPU completion. + +Inputs 0-2 are full rank-1 backing allocations of the same element type. +CPU int64 metadata [8] is ordered as: + source_offset_bytes, destination_offset_bytes, capture_offset_bytes, + batch_size, channels, history_width, capacity, kept_count. +Offsets are logical byte offsets into the full backing tensors, aligned to the element size +(two bytes for FP16, four for FP32), not WebGPU storage-binding offsets. +All fields are nonnegative; channels must be positive; 0 <= kept_count <= capacity. +Zero batch size, history width, or capture capacity is valid. Zero state size is a no-op. +Geometry, backing byte extents and metadata must fit the implementation's uint32 indexing range. + +The source view is the incoming snapshot immediately before capture, with layout +[batch_size, channels, history_width]. Captured inputs have layout [batch_size, capacity, channels]. +capacity is the original capture stride, not kept_count. All retained inputs must have been captured. +For each batch b, channel c and position p, with W = history_width and k = kept_count: + destination[b,c,p] = source[b,c,p+k] if p+k < W + = capture[b,p+k-W,c] otherwise. +For k >= W only the last W inputs of the retained prefix are used; unused capture tails are not read. + +The output must be explicitly prebound to the same full allocation as destination_backing. +Destination backing must be distinct from source and capture backing, even for disjoint logical views. +Only the designated destination view is modified; neighboring elements and all source/capture data +are preserved. One descriptor is supported per invocation; graph capture is not supported. +)DOC") + .Input(0, "source_backing", "Full backing tensor containing the incoming convolution history.", "T") + .Input(1, "capture_backing", "Full backing tensor containing captured token-major inputs.", "T") + .Input(2, "destination_backing", "Full inactive destination-state backing tensor.", "T") + .Input(3, "metadata", "CPU int64 [8]: byte offsets, batch, channels, history, capacity, kept_count.", "TI") + .Output(0, "destination_out", "Same backing allocation and shape as destination_backing.", "T") + .TypeConstraint("T", {"tensor(float16)", "tensor(float)"}, "State and captured inputs are FP16 or 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("ConvStateReplay inputs must have rank 1."); + } + if (index == 3 && shape.dim(0).has_dim_value() && shape.dim(0).dim_value() != 8) { + fail_shape_inference("ConvStateReplay metadata must have shape [8]."); + } + } + 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/internal_nhwc_onnx_schemas.cc b/onnxruntime/core/graph/contrib_ops/internal_nhwc_onnx_schemas.cc index 6bf19654a3ce9..2ae696dedd82a 100644 --- a/onnxruntime/core/graph/contrib_ops/internal_nhwc_onnx_schemas.cc +++ b/onnxruntime/core/graph/contrib_ops/internal_nhwc_onnx_schemas.cc @@ -109,6 +109,7 @@ void OpSet_Internal_NHWC_ONNX::ForEachSchema(const std::function()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index bdf53208969b0..ac48926d34392 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -912,28 +912,55 @@ Status Node::LoadEdgesFromOrtFormat(const onnxruntime::fbs::NodeEdge& fbs_node_e ORT_RETURN_IF(fbs_node_edges.node_index() != index_, "input index: ", fbs_node_edges.node_index(), " is not the same as this node's index:", index_); - auto add_edges = [&graph](const flatbuffers::Vector* fbs_edges, - EdgeSet& edge_set, const std::string& dst_name) -> Status { + auto add_edges = [this, &graph]( + const flatbuffers::Vector* fbs_edges, + bool input_edges) -> Status { if (fbs_edges) { for (const auto* fbs_edge : *fbs_edges) { - ORT_RETURN_IF(nullptr == fbs_edge, "Node::LoadEdgesFromOrtFormat, edge is missing for ", dst_name); + ORT_RETURN_IF(nullptr == fbs_edge, "Node::LoadEdgesFromOrtFormat, edge is missing."); const auto edge_node_index = fbs_edge->node_index(); const size_t node_slot_count = static_cast(graph.MaxNodeIndex()); ORT_RETURN_IF(static_cast(edge_node_index) >= node_slot_count, - "Node::LoadEdgesFromOrtFormat, ", dst_name, " has out-of-range node index ", + "Node::LoadEdgesFromOrtFormat, edge has out-of-range node index ", edge_node_index, ". Invalid ORT format model."); const auto* edge_node = graph.GetNode(edge_node_index); ORT_RETURN_IF(edge_node == nullptr, - "Node::LoadEdgesFromOrtFormat, ", dst_name, " references missing node ", + "Node::LoadEdgesFromOrtFormat, edge references missing node ", edge_node_index, ". Invalid ORT format model."); - edge_set.emplace(*edge_node, fbs_edge->src_arg_index(), fbs_edge->dst_arg_index()); + + const int src_arg_index = fbs_edge->src_arg_index(); + const int dst_arg_index = fbs_edge->dst_arg_index(); + ORT_RETURN_IF(src_arg_index == INT_MAX || dst_arg_index == INT_MAX, + "Node::LoadEdgesFromOrtFormat, control edges are not supported."); + + const Node& src_node = input_edges ? *edge_node : *this; + const Node& dst_node = input_edges ? *this : *edge_node; + ORT_RETURN_IF(src_arg_index < 0 || + static_cast(src_arg_index) >= src_node.OutputDefs().size(), + "Node::LoadEdgesFromOrtFormat, edge has out-of-range src_arg_index."); + const size_t explicit_input_count = dst_node.InputDefs().size(); + const size_t input_count = explicit_input_count + dst_node.ImplicitInputDefs().size(); + ORT_RETURN_IF(dst_arg_index < 0 || static_cast(dst_arg_index) >= input_count, + "Node::LoadEdgesFromOrtFormat, edge has out-of-range dst_arg_index."); + + const NodeArg* src_arg = src_node.OutputDefs()[src_arg_index]; + const NodeArg* dst_arg = static_cast(dst_arg_index) < explicit_input_count + ? dst_node.InputDefs()[dst_arg_index] + : dst_node.ImplicitInputDefs()[dst_arg_index - explicit_input_count]; + ORT_RETURN_IF(!src_arg->Exists() || !dst_arg->Exists(), + "Node::LoadEdgesFromOrtFormat, edge references a missing optional NodeArg."); + ORT_RETURN_IF(src_arg != dst_arg, + "Node::LoadEdgesFromOrtFormat, edge connects mismatched NodeArgs."); + + auto& edge_set = input_edges ? relationships_.input_edges : relationships_.output_edges; + edge_set.emplace(*edge_node, src_arg_index, dst_arg_index); } } return Status::OK(); }; - ORT_RETURN_IF_ERROR(add_edges(fbs_node_edges.input_edges(), relationships_.input_edges, "input edges")); - ORT_RETURN_IF_ERROR(add_edges(fbs_node_edges.output_edges(), relationships_.output_edges, "output edges")); + ORT_RETURN_IF_ERROR(add_edges(fbs_node_edges.input_edges(), true)); + ORT_RETURN_IF_ERROR(add_edges(fbs_node_edges.output_edges(), false)); return Status::OK(); } diff --git a/onnxruntime/core/optimizer/gemm_transpose_fusion.cc b/onnxruntime/core/optimizer/gemm_transpose_fusion.cc index ed7675739b175..bced8655f11e3 100644 --- a/onnxruntime/core/optimizer/gemm_transpose_fusion.cc +++ b/onnxruntime/core/optimizer/gemm_transpose_fusion.cc @@ -36,7 +36,8 @@ Status GemmTransposeFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& m auto new_gemm_input_defs = gemm_node.MutableInputDefs(); // check if input A is a Transpose - if (A_node_ptr != nullptr && A_node_ptr->OpType() == "Transpose" && IsMatrixTranspose(*A_node_ptr)) { + if (A_node_ptr != nullptr && A_node_ptr->OpType() == "Transpose" && IsMatrixTranspose(*A_node_ptr) && + !graph.NodeProducesGraphOutput(*A_node_ptr)) { // make sure all consumers are gemm nodes to avoid possible double transpose std::vector gemm_nodes = graph_utils::FindChildrenByType(*A_node_ptr, "Gemm"); if (gemm_nodes.size() == A_node_ptr->GetOutputEdgesCount()) { @@ -55,7 +56,8 @@ Status GemmTransposeFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& m } } // check if input B is a Transpose - if (B_node_ptr != nullptr && B_node_ptr->OpType() == "Transpose" && IsMatrixTranspose(*B_node_ptr)) { + if (B_node_ptr != nullptr && B_node_ptr->OpType() == "Transpose" && IsMatrixTranspose(*B_node_ptr) && + !graph.NodeProducesGraphOutput(*B_node_ptr)) { std::vector gemm_nodes = graph_utils::FindChildrenByType(*B_node_ptr, "Gemm"); if (gemm_nodes.size() == B_node_ptr->GetOutputEdgesCount()) { Node& B_node = *graph.GetNode(B_node_ptr->Index()); diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index 9348ffbbf5bc7..adc0671c67d8f 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -872,8 +872,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, float, Log); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double, Log); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 14, Pow); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, DepthToSpace); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 27, DepthToSpace); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 27, SpaceToDepth); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, DepthToSpace); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, SpaceToDepth); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Slice); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 17, Split); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, 20, Unsqueeze); @@ -2852,8 +2854,12 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo dims(x_rank, 1); + const size_t sr = s.NumDimensions(); + const size_t count = std::min(sr, x_rank); + for (size_t i = 0; i < count; ++i) { + dims[x_rank - 1 - i] = s.GetDims()[sr - 1 - i]; + } + return dims; + }; + // X shape is (B, S, ...) - if (axis == 2 && - x_shape.NumDimensions() >= 3 && - x_shape.NumDimensions() == scale_shape.NumDimensions() && - (bias_shape == nullptr || *bias_shape == scale_shape)) { - for (size_t i = 2; i < x_shape.NumDimensions(); ++i) { - if (x_shape.GetDims()[i] != scale_shape.GetDims()[i]) { + if (axis == 2 && x_rank >= 3) { + const InlinedVector scale_dims = right_align(scale_shape); + if (bias_shape != nullptr && right_align(*bias_shape) != scale_dims) { + // scale and bias must share the same (right-aligned) broadcast layout. + return kLayerNormInvalidInput; + } + + for (size_t i = 2; i < x_rank; ++i) { + if (x_shape.GetDims()[i] != scale_dims[i]) { // scale cannot be broadcasted to X. It is invalid input. return kLayerNormInvalidInput; } } - if (x_shape.GetDims()[0] == scale_shape.GetDims()[0]) { + if (x_shape.GetDims()[0] == scale_dims[0]) { // scale and bias shape is (B, S, ...). - if (x_shape.GetDims()[1] == scale_shape.GetDims()[1]) { + if (x_shape.GetDims()[1] == scale_dims[1]) { return 1; } // scale and bias shape is (B, 1, ...), returns S - if (scale_shape.GetDims()[1] == 1) { + if (scale_dims[1] == 1) { return x_shape.GetDims()[1]; } - } else if (scale_shape.GetDims()[0] == 1) { + } else if (scale_dims[0] == 1) { // scale and bias shape is (1, S, ...), returns -S - if (x_shape.GetDims()[1] == scale_shape.GetDims()[1]) { + if (x_shape.GetDims()[1] == scale_dims[1]) { return -(x_shape.GetDims()[1]); } } diff --git a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc index f59cd88508b71..c934e4f08e2a4 100644 --- a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc +++ b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc @@ -29,6 +29,16 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( DataTypeImpl::GetTensorType()}), SpaceToDepth); +ONNX_CPU_OPERATOR_KERNEL( + SpaceToDepth, + 28, + KernelDefBuilder() + .TypeConstraint("T", {DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + SpaceToDepth); + ONNX_CPU_OPERATOR_VERSIONED_KERNEL( DepthToSpace, 1, 10, @@ -48,9 +58,20 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( DataTypeImpl::GetTensorType()}), DepthToSpace); -ONNX_CPU_OPERATOR_KERNEL( +ONNX_CPU_OPERATOR_VERSIONED_KERNEL( DepthToSpace, 13, + 27, + KernelDefBuilder() + .TypeConstraint("T", {DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + DepthToSpace); + +ONNX_CPU_OPERATOR_KERNEL( + DepthToSpace, + 28, KernelDefBuilder() .TypeConstraint("T", {DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), @@ -108,7 +129,11 @@ Status SpaceToDepth::Compute(OpKernelContext* context) const { Tensor& output = *context->Output(0, {batch, output_depth, output_height, output_width}); - std::array permutation{{0, 3, 5, 1, 2, 4}}; + const auto permutation = is_dcr_ ? std::array{{0, 3, 5, 1, 2, 4}} + : std::array{{0, 1, 3, 5, 2, 4}}; + const auto out_dim1 = onnxruntime::narrow(is_dcr_ ? blocksize_ : input_depth); + const auto out_dim2 = onnxruntime::narrow(blocksize_); + const auto out_dim3 = onnxruntime::narrow(is_dcr_ ? input_depth : blocksize_); if (input.IsDataType()) { SpaceDepthOpCpuImpl(input.Data(), output.MutableData(), permutation, @@ -118,9 +143,7 @@ Status SpaceToDepth::Compute(OpKernelContext* context) const { onnxruntime::narrow(blocksize_), onnxruntime::narrow(input_width / blocksize_), onnxruntime::narrow(blocksize_), - onnxruntime::narrow(blocksize_), - onnxruntime::narrow(blocksize_), - onnxruntime::narrow(input_depth), + out_dim1, out_dim2, out_dim3, onnxruntime::narrow(input_height / blocksize_), onnxruntime::narrow(input_width / blocksize_)); } else if (input.IsDataType()) { @@ -131,9 +154,7 @@ Status SpaceToDepth::Compute(OpKernelContext* context) const { onnxruntime::narrow(blocksize_), onnxruntime::narrow(input_width / blocksize_), onnxruntime::narrow(blocksize_), - onnxruntime::narrow(blocksize_), - onnxruntime::narrow(blocksize_), - onnxruntime::narrow(input_depth), + out_dim1, out_dim2, out_dim3, onnxruntime::narrow(input_height / blocksize_), onnxruntime::narrow(input_width / blocksize_)); } else if (input.IsDataType() || input.IsDataType()) { @@ -147,9 +168,7 @@ Status SpaceToDepth::Compute(OpKernelContext* context) const { onnxruntime::narrow(blocksize_), onnxruntime::narrow(input_width / blocksize_), onnxruntime::narrow(blocksize_), - onnxruntime::narrow(blocksize_), - onnxruntime::narrow(blocksize_), - onnxruntime::narrow(input_depth), + out_dim1, out_dim2, out_dim3, onnxruntime::narrow(input_height / blocksize_), onnxruntime::narrow(input_width / blocksize_)); } else { diff --git a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.h b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.h index 3db5f2565226b..3e64bdef076fd 100644 --- a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.h +++ b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.h @@ -26,13 +26,12 @@ template inline bool ReadIsDCR(const KernelInfoType& info) { bool is_dcr = true; std::string mode; - // If mode doesn't exist, then it is the default "DCR" mode - // (or) it is an opset < 11 model for which the only mode is "DCR" mode. + // A missing mode is the default DCR mode, including pre-mode schemas. if (info.GetAttr("mode", &mode).IsOK()) { if (mode == "CRD") { is_dcr = false; } else if (mode != "DCR") { - ORT_THROW("DepthToSpace op: only 'DCR' and 'CRD' modes are supported"); + ORT_THROW("SpaceDepth op: only 'DCR' and 'CRD' modes are supported"); } } @@ -130,10 +129,13 @@ class SpaceDepthBase { class SpaceToDepth final : public OpKernel, SpaceDepthBase { public: - explicit SpaceToDepth(const OpKernelInfo& info) : OpKernel(info), SpaceDepthBase(info) { - } + explicit SpaceToDepth(const OpKernelInfo& info) + : OpKernel(info), SpaceDepthBase(info), is_dcr_(space_depth_internal::ReadIsDCR(info)) {} Status Compute(OpKernelContext* context) const override; + + private: + bool is_dcr_ = true; }; class DepthToSpace final : public OpKernel, SpaceDepthBase { diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 0697ff2148fa5..deb03dbd8bd0a 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -8,6 +8,7 @@ #include "core/common/inlined_containers.h" #include "core/common/parse_string.h" #include "core/framework/int4.h" +#include "core/framework/kernel_dispatch_policy.h" #if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) #include "core/framework/node_shape_resolver.h" #endif @@ -1353,7 +1354,9 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kO class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 17, MLFloat16, Pad); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 17, bool, Pad); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 27, SpaceToDepth); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, DepthToSpace); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 27, DepthToSpace); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, SpaceToDepth); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 28, DepthToSpace); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int8_t, Sign); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int16_t, Sign); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int32_t, Sign); @@ -2678,8 +2681,12 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -3736,13 +3743,22 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, input_defs[11]->Exists() && graph.IsConstantInitializer(input_defs[11]->Name(), true); // The factory validates the caller-declared envelope once. + const auto& estimator_config = resource_accountant->GetWorkspaceEstimatorConfig(); const auto gqa_workspace_max_total_sequence_length = - resource_accountant->GetWorkspaceEstimatorConfig() - .cuda_gqa_workspace_max_total_sequence_length; + estimator_config.cuda_gqa_workspace_max_total_sequence_length; + // Reader-only intent; no value changes route selection or sizing. + KernelDispatchPolicy kernel_dispatch_policy = KernelDispatchPolicy::Auto; + if (const auto& policy = estimator_config.kernel_dispatch_policy; + policy.has_value()) { + ORT_ENFORCE( + TryParseKernelDispatchPolicy(*policy, kernel_dispatch_policy), + "Invalid ", kOrtSessionOptionsKernelDispatchPolicy, + " in WorkspaceEstimatorConfig: '", *policy, "'"); + } const auto ws = contrib::cuda::EstimateGroupQueryAttentionWorkspace( *node, gsl::make_span(input_shapes), GetDeviceProp(), *GetAttentionKernelOptions(), head_sink_is_constant_initializer, - gqa_workspace_max_total_sequence_length); + gqa_workspace_max_total_sequence_length, kernel_dispatch_policy); if (ws.has_value()) { Level1MemoryEstimate estimate; contrib::cuda::SetGroupQueryAttentionLevel1MemoryEstimate(*ws, estimate); diff --git a/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc b/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc index 22f14b6123eb8..cd004138aed5c 100755 --- a/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc +++ b/onnxruntime/core/providers/cuda/cuda_nhwc_kernels.cc @@ -83,9 +83,11 @@ class CUDA_NHWC_OP_TYPED_CLASS_NAME(15, double, BatchNormalization); class CUDA_NHWC_OP_TYPED_CLASS_NAME(15, MLFloat16, BatchNormalization); class CUDA_NHWC_OP_VERSIONED_CLASS_NAME(1, 10, DepthToSpace); class CUDA_NHWC_OP_VERSIONED_CLASS_NAME(11, 12, DepthToSpace); -class CUDA_NHWC_OP_CLASS_NAME(13, DepthToSpace); +class CUDA_NHWC_OP_VERSIONED_CLASS_NAME(13, 27, DepthToSpace); +class CUDA_NHWC_OP_CLASS_NAME(28, DepthToSpace); class CUDA_NHWC_OP_VERSIONED_CLASS_NAME(1, 12, SpaceToDepth); class CUDA_NHWC_OP_VERSIONED_CLASS_NAME(13, 27, SpaceToDepth); +class CUDA_NHWC_OP_CLASS_NAME(28, SpaceToDepth); class CUDA_NHWC_OP_VERSIONED_TYPED_CLASS_NAME(1, 12, float, LRN); class CUDA_NHWC_OP_VERSIONED_TYPED_CLASS_NAME(1, 12, double, LRN); class CUDA_NHWC_OP_VERSIONED_TYPED_CLASS_NAME(1, 12, MLFloat16, LRN); @@ -152,9 +154,11 @@ Status RegisterCudaNhwcKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/providers/cuda/tensor/space_depth_ops.cc b/onnxruntime/core/providers/cuda/tensor/space_depth_ops.cc index 3c5fe31f35309..dfffad54f5308 100644 --- a/onnxruntime/core/providers/cuda/tensor/space_depth_ops.cc +++ b/onnxruntime/core/providers/cuda/tensor/space_depth_ops.cc @@ -50,6 +50,18 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( DataTypeImpl::GetTensorType()}), SpaceToDepth); +ONNX_OPERATOR_KERNEL_EX( + SpaceToDepth, + kOnnxDomain, + 28, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", + {DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + SpaceToDepth); + #ifdef ENABLE_CUDA_NHWC_OPS ONNX_OPERATOR_VERSIONED_KERNEL_EX( SpaceToDepth, @@ -63,6 +75,18 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), SpaceToDepth); + +ONNX_OPERATOR_KERNEL_EX( + SpaceToDepth, + kMSInternalNHWCDomain, + 28, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", + {DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + SpaceToDepth); #endif ONNX_OPERATOR_VERSIONED_KERNEL_EX( @@ -121,10 +145,11 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( DepthToSpace); #endif -ONNX_OPERATOR_KERNEL_EX( +ONNX_OPERATOR_VERSIONED_KERNEL_EX( DepthToSpace, kOnnxDomain, 13, + 27, kCudaExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", @@ -133,11 +158,36 @@ ONNX_OPERATOR_KERNEL_EX( DataTypeImpl::GetTensorType()}), DepthToSpace); -#ifdef ENABLE_CUDA_NHWC_OPS ONNX_OPERATOR_KERNEL_EX( + DepthToSpace, + kOnnxDomain, + 28, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", + {DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + DepthToSpace); + +#ifdef ENABLE_CUDA_NHWC_OPS +ONNX_OPERATOR_VERSIONED_KERNEL_EX( DepthToSpace, kMSInternalNHWCDomain, 13, + 27, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", + {DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + DepthToSpace); + +ONNX_OPERATOR_KERNEL_EX( + DepthToSpace, + kMSInternalNHWCDomain, + 28, kCudaExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", @@ -193,15 +243,25 @@ Status SpaceToDepth::ComputeInternal(OpKernelContext* context) const { input_width / blocksize_, blocksize_, input_depth}; // We will pass in the "virtual" output shape to be used by DoTranspose() in SpaceDepthOpCudaImpl(...) - TensorShape virtual_output_shape = (Layout == LAYOUT_NCHW) - ? TensorShape{batch, blocksize_, blocksize_, input_depth, - input_height / blocksize_, input_width / blocksize_} - : TensorShape{batch, input_height / blocksize_, input_width / blocksize_, - blocksize_, blocksize_, input_depth}; - - std::vector permutation = (Layout == LAYOUT_NCHW) - ? std::vector{0, 3, 5, 1, 2, 4} - : std::vector{0, 1, 3, 2, 4, 5}; + TensorShape virtual_output_shape; + std::vector permutation; + if (is_dcr_) { + virtual_output_shape = (Layout == LAYOUT_NCHW) + ? TensorShape{batch, blocksize_, blocksize_, input_depth, + input_height / blocksize_, input_width / blocksize_} + : TensorShape{batch, input_height / blocksize_, input_width / blocksize_, + blocksize_, blocksize_, input_depth}; + permutation = (Layout == LAYOUT_NCHW) + ? std::vector{0, 3, 5, 1, 2, 4} + : std::vector{0, 1, 3, 2, 4, 5}; + } else { + virtual_output_shape = (Layout == LAYOUT_NCHW) + ? TensorShape{batch, input_depth, blocksize_, blocksize_, + input_height / blocksize_, input_width / blocksize_} + : TensorShape{batch, input_height / blocksize_, input_width / blocksize_, + input_depth, blocksize_, blocksize_}; + permutation = {0, 1, 3, 5, 2, 4}; + } ORT_RETURN_IF_ERROR( SpaceDepthOpCudaImpl(GetDeviceProp(), Stream(context), GetCublasHandle(context), input, output, permutation, diff --git a/onnxruntime/core/providers/cuda/tensor/space_depth_ops.h b/onnxruntime/core/providers/cuda/tensor/space_depth_ops.h index 5cb4c54aeef3c..28fe9b6c3248d 100644 --- a/onnxruntime/core/providers/cuda/tensor/space_depth_ops.h +++ b/onnxruntime/core/providers/cuda/tensor/space_depth_ops.h @@ -13,9 +13,12 @@ template class SpaceToDepth final : public CudaKernel, SpaceDepthBase { public: explicit SpaceToDepth(const OpKernelInfo& info) - : CudaKernel(info), SpaceDepthBase(info) {} + : CudaKernel(info), SpaceDepthBase(info), is_dcr_(space_depth_internal::ReadIsDCR(info)) {} Status ComputeInternal(OpKernelContext* context) const override; + + private: + bool is_dcr_ = true; }; template diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp index edfe419028090..c9d1cfad2a646 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -816,9 +816,11 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO( 7, SpaceToDepth, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QuerySpaceToDepth)}, {REG_INFO( 13, SpaceToDepth, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QuerySpaceToDepth)}, + {REG_INFO( 28, SpaceToDepth, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QuerySpaceToDepth)}, {REG_INFO( 7, DepthToSpace, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, {REG_INFO( 11, DepthToSpace, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, {REG_INFO( 13, DepthToSpace, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO( 28, DepthToSpace, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, {REG_INFO( 7, Tile, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, {REG_INFO( 13, Tile, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, {REG_INFO( 8, Expand, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h index c78439c9ebff3..e8b22d2d83d15 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h @@ -459,6 +459,8 @@ namespace OperatorHelper namespace OnnxOperatorSet28 { + static const int sc_sinceVer_DepthToSpace = 28; + static const int sc_sinceVer_SpaceToDepth = 28; static const int sc_sinceVer_Cast = 28; static const int sc_sinceVer_Mod = 28; } diff --git a/onnxruntime/core/providers/webgpu/ep/README.md b/onnxruntime/core/providers/webgpu/ep/README.md index e2d52a67e9dbe..77d67027589b9 100644 --- a/onnxruntime/core/providers/webgpu/ep/README.md +++ b/onnxruntime/core/providers/webgpu/ep/README.md @@ -214,6 +214,52 @@ 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. +### Convolution state replay + +`com.microsoft::ConvStateReplay` version 1 reconstructs convolution history without +convolution arithmetic. It is WebGPU-only and accepts FP16 or FP32 tensors. Inputs +0-2 are full rank-1 source, capture, and inactive destination backing tensors of +the same type. Input 3 is CPU `int64[8]`, ordered as: +`source_offset_bytes, destination_offset_bytes, capture_offset_bytes, +batch_size, channels, history_width, capacity, kept_count`. +Offsets are bytes from the allocation base, aligned to the element size; FP16 +offsets need not be four-byte aligned. State views have layout +`[batch_size, channels, history_width]`, and captured inputs have layout +`[batch_size, capacity, channels]`. `capacity` is the original batch stride, +not the retained count. All fields are nonnegative, channels must be positive, +and `0 <= kept_count <= capacity`. Zero batch/history is a no-op; zero capacity +is valid with zero retained inputs. Full view extents, including capture capacity, +must fit their backings even when only part of a view is read. + +The caller must supply the incoming state immediately before capture and valid +captured inputs for the retained prefix. With width `W` and retained count `k`, +the output history is the final `W` entries of the concatenation of the incoming +history and the first `k` captured inputs. Thus `k=0` copies the source, while +`k>=W` selects only the last `W` inputs of the retained prefix. + +The output must be explicitly I/O-bound to the same full allocation as +`destination_backing`; `.MayInplace` alone does not guarantee this. Destination +backing must not alias source or capture backing, even for disjoint views. No +other destination elements or source/capture payloads are modified. + +The kernel wraps the full opaque buffer handles as non-owning `uint32` tensors. +`ProgramInput::BufferView` / `ProgramOutput::BufferView` supply word offsets via +uniforms, while binding full allocations (or framework-aligned segments). +FP16's remaining two-byte offset is a lane uniform. One invocation owns each +destination word and merges selected 16-bit payloads with the original word at +view boundaries, preserving adjacent halves without races. No floating-point +loads, stores, conversions, or arithmetic are used, including for FP32, so +signed zeros, subnormals and NaN payloads are preserved exactly. Backings must +have WebGPU's four-byte allocation padding, including odd-sized FP16 tensors. + +The operation uses checked GPU completion before returning, like the separate +GatedDeltaNet replay operator. Failure must not be treated as a committed state +update. There is no graph-capture support or GenAI integration in this milestone. +Metadata and backing byte extents are limited to uint32; geometry, alignment, +bounds, output aliasing, dispatch, and device binding limits are checked. +Integration should convert descriptor offsets to bytes, retain the original +capture capacity, prebind the inactive bank, and publish it only after success. + ### GatedDeltaNet state replay `com.microsoft::GatedDeltaNetStateReplay` version 1 is a WebGPU-only, FP32 @@ -264,8 +310,9 @@ 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 +This operator has no graph-capture support, batching, or external session +integration; convolution replay uses the separate `ConvStateReplay` operator. +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 diff --git a/onnxruntime/core/providers/webgpu/tensor/depth_to_space.cc b/onnxruntime/core/providers/webgpu/tensor/depth_to_space.cc index 9b54eca641dd7..d1e73045f1773 100644 --- a/onnxruntime/core/providers/webgpu/tensor/depth_to_space.cc +++ b/onnxruntime/core/providers/webgpu/tensor/depth_to_space.cc @@ -31,10 +31,12 @@ namespace webgpu { DepthToSpace); WEBGPU_DEPTH_TO_SPACE_VERSIONED_KERNEL(11, 12, kOnnxDomain, false) -WEBGPU_DEPTH_TO_SPACE_KERNEL(13, kOnnxDomain, false) +WEBGPU_DEPTH_TO_SPACE_VERSIONED_KERNEL(13, 27, kOnnxDomain, false) +WEBGPU_DEPTH_TO_SPACE_KERNEL(28, kOnnxDomain, false) WEBGPU_DEPTH_TO_SPACE_VERSIONED_KERNEL(11, 12, kMSInternalNHWCDomain, true) -WEBGPU_DEPTH_TO_SPACE_KERNEL(13, kMSInternalNHWCDomain, true) +WEBGPU_DEPTH_TO_SPACE_VERSIONED_KERNEL(13, 27, kMSInternalNHWCDomain, true) +WEBGPU_DEPTH_TO_SPACE_KERNEL(28, kMSInternalNHWCDomain, true) void AppendPermFunction(OStringStream& os, const ShaderVariableHelper& input, const int64_t* perm) { os << "fn perm(i: input_indices_t) -> input_indices_t {\n" diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 742b469edb6ae..2af2fd6b49e6c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -278,9 +278,14 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index fbadf6da73678..bc8009b05c9e7 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -303,10 +303,14 @@ Status Environment::Initialize(std::unique_ptr logging_ // we have static registrations for NHWC versions of ONNX operators so this domain needs to extend to the // latest ONNX version - auto onnx_version = domainToVersionRangeInstance.LastReleaseVersionMap() - .find(ONNX_NAMESPACE::ONNX_DOMAIN) - ->second; - domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSInternalNHWCDomain, 1, onnx_version); + const auto onnx_version = domainToVersionRangeInstance.Map() + .find(ONNX_NAMESPACE::ONNX_DOMAIN) + ->second.second; + const auto last_released_onnx_version = domainToVersionRangeInstance.LastReleaseVersionMap() + .find(ONNX_NAMESPACE::ONNX_DOMAIN) + ->second; + domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kMSInternalNHWCDomain, 1, onnx_version, + last_released_onnx_version); domainToVersionRangeInstance.AddDomainToVersion(onnxruntime::kPytorchAtenDomain, 1, 1); #ifdef USE_DML diff --git a/onnxruntime/python/tools/transformers/models/llama/requirements.txt b/onnxruntime/python/tools/transformers/models/llama/requirements.txt index 46a4ed1ca615e..25529827157de 100644 --- a/onnxruntime/python/tools/transformers/models/llama/requirements.txt +++ b/onnxruntime/python/tools/transformers/models/llama/requirements.txt @@ -1,7 +1,7 @@ onnxscript>=0.2.3 optimum>=1.14.1 optree -transformers==5.3.0 +transformers==5.10.1 torch>=2.7.0 onnx==1.22.0 datasets>=2.8.0 diff --git a/onnxruntime/python/tools/transformers/models/phi2/requirements.txt b/onnxruntime/python/tools/transformers/models/phi2/requirements.txt index 4d81d8e6ea6b4..b4a785d303513 100644 --- a/onnxruntime/python/tools/transformers/models/phi2/requirements.txt +++ b/onnxruntime/python/tools/transformers/models/phi2/requirements.txt @@ -1,3 +1,3 @@ onnx==1.22.0 -transformers>=5.3.0 +transformers>=5.10.1 onnxscript>=0.1.0.dev20240126 diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements/requirements.txt b/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements/requirements.txt index ff57fd6232463..f6fd689eec9ba 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements/requirements.txt +++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements/requirements.txt @@ -1,6 +1,6 @@ huggingface_hub>=1.23.0,<2.0 diffusers==0.40.0 -transformers==5.3.0 +transformers==5.10.1 numpy>=1.24.1 accelerate onnx==1.18.0 diff --git a/onnxruntime/python/tools/transformers/models/whisper/requirements.txt b/onnxruntime/python/tools/transformers/models/whisper/requirements.txt index 2f65d30816d33..6b08503b2289b 100644 --- a/onnxruntime/python/tools/transformers/models/whisper/requirements.txt +++ b/onnxruntime/python/tools/transformers/models/whisper/requirements.txt @@ -1,5 +1,5 @@ torch==2.7.0 -transformers==5.3.0 +transformers==5.10.1 openai-whisper==20240927 ffmpeg-python datasets diff --git a/onnxruntime/python/tools/transformers/requirements.txt b/onnxruntime/python/tools/transformers/requirements.txt index c9e46d8fa57f0..dd8c91ff6df6e 100644 --- a/onnxruntime/python/tools/transformers/requirements.txt +++ b/onnxruntime/python/tools/transformers/requirements.txt @@ -4,7 +4,7 @@ psutil py-cpuinfo py3nvml packaging -transformers >= 5.3.0 +transformers >= 5.10.1 scipy sentencepiece pillow diff --git a/onnxruntime/test/contrib_ops/conv_state_replay_op_test.cc b/onnxruntime/test/contrib_ops/conv_state_replay_op_test.cc new file mode 100644 index 0000000000000..8da29ff420656 --- /dev/null +++ b/onnxruntime/test/contrib_ops/conv_state_replay_op_test.cc @@ -0,0 +1,412 @@ +// 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/util/include/asserts.h" +#include "test/util/include/default_providers.h" +#include "test/util/include/test_environment.h" + +#ifdef USE_WEBGPU +#include "core/session/IOBinding.h" +#include "core/session/inference_session.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#endif + +namespace onnxruntime::test { +namespace { + +using Shapes = std::array, 4>; + +std::unique_ptr MakeConvReplayModel(const Shapes& shapes, bool half) { + auto model = std::make_unique( + "conv_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", "capture_backing", "destination_backing", "metadata"}; + std::vector inputs; + for (size_t index = 0; index < names.size(); ++index) { + ONNX_NAMESPACE::TypeProto type; + auto* tensor = type.mutable_tensor_type(); + tensor->set_elem_type(index == 3 ? ONNX_NAMESPACE::TensorProto_DataType_INT64 + : half ? ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 + : ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + for (const auto dim : shapes[index]) { + auto* dimension = tensor->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", "ConvStateReplay", "", inputs, {&output}, nullptr, kMSDomain); + node.SetExecutionProviderType(kWebGpuExecutionProvider); + graph.SetOutputs({&output}); + return model; +} + +TEST(ConvStateReplayShapeInferenceTest, DestinationShapeAndType) { + for (const bool half : {false, true}) { + auto model = MakeConvReplayModel(Shapes{{{20}, {30}, {40}, {8}}}, half); + ASSERT_STATUS_OK(model->MainGraph().Resolve()); + const auto* output = model->MainGraph().GetOutputs()[0]; + EXPECT_EQ(output->TypeAsProto()->tensor_type().elem_type(), + half ? ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 : ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + ASSERT_EQ(output->Shape()->dim_size(), 1); + EXPECT_EQ(output->Shape()->dim(0).dim_value(), 40); + model = MakeConvReplayModel(Shapes{{{-1}, {-1}, {-1}, {8}}}, half); + ASSERT_STATUS_OK(model->MainGraph().Resolve()); + EXPECT_EQ(model->MainGraph().GetOutputs()[0]->Shape()->dim(0).dim_param(), "backing_extent"); + } +} + +TEST(ConvStateReplayShapeInferenceTest, RejectsRanksAndMetadataLength) { + for (size_t index = 0; index < 4; ++index) { + auto shapes = Shapes{{{20}, {30}, {40}, {8}}}; + shapes[index] = {1, shapes[index][0]}; + auto model = MakeConvReplayModel(shapes, false); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(model->MainGraph().Resolve(), "inputs must have rank 1"); + } + for (const int64_t size : {0, 7, 9}) { + auto model = MakeConvReplayModel(Shapes{{{20}, {30}, {40}, {size}}}, true); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(model->MainGraph().Resolve(), "metadata must have shape [8]"); + } +} + +#ifdef USE_WEBGPU + +struct Geometry { + int64_t batch, channels, width, capacity; +}; + +// Store the CPU oracle and transfers as integers so neither can canonicalize NaNs. +template +class ConvReplayHarness { + public: + Status Initialize(const Geometry& geometry, const std::array& offsets = {5, 7, 9}, + const std::array& tails = {13, 11, 15}, + std::unique_ptr provider = DefaultWebGpuExecutionProvider()) { + 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 [batch, channels, width, capacity] = geometry; + desc_ = {offsets[0] * static_cast(sizeof(Bits)), offsets[1] * static_cast(sizeof(Bits)), + offsets[2] * static_cast(sizeof(Bits)), + batch, channels, width, capacity, 0}; + const int64_t state_elements = batch * channels * width; + source_.resize(static_cast(offsets[0] + state_elements + tails[0])); + destination_.resize(static_cast(offsets[1] + state_elements + tails[1])); + capture_.resize(static_cast(offsets[2] + batch * channels * capacity + tails[2])); + Fill(source_, 17); + Fill(capture_, 113); + Fill(destination_, 971); + auto model = MakeConvReplayModel(Shapes{{{static_cast(source_.size())}, + {static_cast(capture_.size())}, + {static_cast(destination_.size())}, + {8}}}, + sizeof(Bits) == 2); + ORT_RETURN_IF_ERROR(model->MainGraph().Resolve()); + std::stringstream stream; + model->ToProto().SerializeToOstream(&stream); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + ORT_RETURN_IF_ERROR(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + session_ = std::make_unique(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."); + const std::array backing_values{std::make_pair(&source_, &source_value_), std::make_pair(&capture_, &capture_value_), + std::make_pair(&destination_, &destination_value_)}; + for (const auto& item : backing_values) { + Tensor::InitOrtValue(ElementType(), TensorShape{static_cast(item.first->size())}, + allocator_, *item.second); + } + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape{8}, desc_.data(), + cpu_allocator_.Info(), metadata_value_); + return Status::OK(); + } + + Status Prepare(int64_t kept) { + desc_[7] = kept; + ORT_RETURN_IF_ERROR(Transfer(source_, source_value_, true)); + ORT_RETURN_IF_ERROR(Transfer(capture_, capture_value_, true)); + return Transfer(destination_, destination_value_, true); + } + + Status Run(bool alias_output = true, int aliased_input = -1) { + std::unique_ptr binding; + ORT_RETURN_IF_ERROR(session_->NewIOBinding(&binding)); + OrtValue input_alias; + if (aliased_input >= 0) { + const auto& input = (aliased_input == 0 ? source_value_ : capture_value_).template Get(); + Tensor::InitOrtValue(ElementType(), input.Shape(), destination_value_.GetMutable()->MutableDataRaw(), + input.Location(), input_alias); + } + ORT_RETURN_IF_ERROR(binding->BindInput("source_backing", aliased_input == 0 ? input_alias : source_value_)); + ORT_RETURN_IF_ERROR(binding->BindInput("capture_backing", aliased_input == 1 ? input_alias : capture_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(ElementType(), 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); + } + + void Verify() { + auto expected = destination_; + const auto batch = desc_[3], channels = desc_[4], width = desc_[5], capacity = desc_[6]; + for (int64_t b = 0; b < batch; ++b) { + for (int64_t c = 0; c < channels; ++c) { + const auto row = static_cast((b * channels + c) * width); + const auto source_start = source_.begin() + desc_[0] / sizeof(Bits) + row; + std::deque history(source_start, source_start + width); + // Independent token-by-token rolling queue, not the shader's direct selection formula. + for (int64_t t = 0; t < desc_[7] && width > 0; ++t) { + history.pop_front(); + history.push_back(capture_[static_cast(desc_[2] / sizeof(Bits) + (b * capacity + t) * channels + c)]); + } + std::copy(history.begin(), history.end(), expected.begin() + desc_[1] / sizeof(Bits) + row); + } + } + VerifyBacking(expected, destination_value_); + VerifyBacking(source_, source_value_); + VerifyBacking(capture_, capture_value_); + } + + void VerifyDestinationUnchanged() { VerifyBacking(destination_, destination_value_); } + std::array& Descriptor() { return desc_; } + size_t SourceBytes() const { return source_.size() * sizeof(Bits); } + size_t CaptureBytes() const { return capture_.size() * sizeof(Bits); } + size_t DestinationBytes() const { return destination_.size() * sizeof(Bits); } + + private: + static MLDataType ElementType() { + return sizeof(Bits) == 2 ? DataTypeImpl::GetType() : DataTypeImpl::GetType(); + } + + static void Fill(std::vector& values, uint32_t seed) { + constexpr std::array special{ + 0x00000000, 0x80000000, 0x7f800000, 0xff800000, 0x7fc01234, 0x7f801234, 0xffc54321, 0xff800001, + 0x00000001, 0x80000001, 0x007fffff, 0x807fffff, 0x3f800000, 0xbf800000, 0x7fffffff, 0xffffffff}; + for (size_t index = 0; index < values.size(); ++index) { + if constexpr (sizeof(Bits) == 2) { + // Odd multiplier visits all 65,536 half payloads in a sufficiently large backing. + values[index] = static_cast(index * 40503 + seed); + } else { + values[index] = static_cast(index % 2 == 0 ? special[(index / 2 + seed) % special.size()] + : static_cast(index * 2654435761u + seed)); + } + } + } + + Status Transfer(std::vector& data, OrtValue& value, bool upload) { + if (data.empty()) { + return Status::OK(); + } + Tensor cpu_tensor(ElementType(), TensorShape{static_cast(data.size())}, + data.data(), cpu_allocator_.Info()); + return upload ? provider_->GetDataTransfer()->CopyTensor(cpu_tensor, *value.GetMutable()) + : provider_->GetDataTransfer()->CopyTensor(value.Get(), cpu_tensor); + } + + void VerifyBacking(const std::vector& expected, OrtValue& value) { + std::vector actual(expected.size()); + ASSERT_STATUS_OK(Transfer(actual, value, false)); + ASSERT_EQ(actual, expected); + } + + std::unique_ptr session_; + IExecutionProvider* provider_ = nullptr; + AllocatorPtr allocator_; + CPUAllocator cpu_allocator_; + std::vector source_, capture_, destination_; + std::array desc_{}; + OrtValue source_value_, capture_value_, destination_value_, metadata_value_; +}; + +template +class ConvStateReplayTest : public ::testing::Test { + protected: + void SetUp() override { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU EP is unavailable."; + } + } +}; + +using ReplayPayloadTypes = ::testing::Types; +TYPED_TEST_SUITE(ConvStateReplayTest, ReplayPayloadTypes); + +TYPED_TEST(ConvStateReplayTest, EveryPrefixBatchChannelsAndHistoryWidths) { + for (const int64_t capacity : {1, 4, 8}) { + for (const int64_t width : {1, 3, 11}) { + SCOPED_TRACE(::testing::Message() << "capacity " << capacity << " width " << width); + ConvReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({2, 5, width, capacity})); + for (int64_t kept = 0; kept <= capacity; ++kept) { + SCOPED_TRACE(kept); + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } + } + } +} + +TYPED_TEST(ConvStateReplayTest, ArbitraryOffsetsPreserveAdjacentElements) { + for (const auto offsets : {std::array{0, 0, 0}, {1, 1, 1}, {2, 3, 4}, {31, 128, 257}}) { + SCOPED_TRACE(::testing::Message() << offsets[0] << "," << offsets[1] << "," << offsets[2]); + ConvReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({3, 3, 5, 7}, offsets)); + for (int64_t kept = 0; kept <= 7; ++kept) { + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } + } +} + +TYPED_TEST(ConvStateReplayTest, ExactEndAndSingleElementBoundaries) { + for (const int64_t width : {1, 2, 5}) { + for (const int64_t offset : {0, 1}) { + ConvReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({1, 1, width, 7}, {offset, offset, offset}, {0, 0, 0})); + for (int64_t kept = 0; kept <= 7; ++kept) { + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } + } + } +} + +TYPED_TEST(ConvStateReplayTest, ZeroCapacityHistoryAndBatch) { + for (const auto geometry : {Geometry{2, 3, 5, 0}, {2, 3, 0, 4}, {0, 3, 5, 4}, {0, 3, 0, 0}}) { + for (const int64_t padding : {0, 1}) { + ConvReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize(geometry, {padding, padding, padding}, {0, 0, 0})); + for (int64_t kept = 0; kept <= geometry.capacity; ++kept) { + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } + } + } +} + +TYPED_TEST(ConvStateReplayTest, ExactSpecialPayloadBits) { + ConvReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({1, 1, 65536, 65537}, {1, 3, 5})); + // Source-only and capture-only selections cover every FP16 bit pattern, including + // signaling NaNs, both zero signs, infinities and subnormals, at odd half offsets. + for (const int64_t kept : {0, 1, 65536, 65537}) { + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } +} + +TYPED_TEST(ConvStateReplayTest, SegmentedBackingViews) { + WebGpuExecutionProviderTestOptions options{}; + options.max_storage_buffer_binding_size = 16384 * sizeof(TypeParam); + auto provider = WebGpuExecutionProviderWithTestOptions(options); + if (!provider) { + GTEST_SKIP() << "Artificial storage limits require the native test factory."; + } + ConvReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({2, 1000, 7, 8}, {11001, 13001, 15001}, {1, 1, 1}, std::move(provider))); + for (const int64_t kept : {0, 3, 7, 8}) { + ASSERT_STATUS_OK(harness.Prepare(kept)); + ASSERT_STATUS_OK(harness.Run()); + ASSERT_NO_FATAL_FAILURE(harness.Verify()); + } +} + +TYPED_TEST(ConvStateReplayTest, RejectsInvalidMetadataBeforeWriting) { + ConvReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({2, 3, 5, 7})); + ASSERT_STATUS_OK(harness.Prepare(3)); + 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"); + } + } + harness.Descriptor() = valid; + harness.Descriptor()[4] = 0; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(), "channels must be positive"); + harness.Descriptor() = valid; + harness.Descriptor()[7] = 8; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(), "kept_count must be in [0, capacity]"); + for (size_t field = 0; field < 3; ++field) { + harness.Descriptor() = valid; + harness.Descriptor()[field] += 1; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(), "byte offsets must be element-aligned"); + harness.Descriptor() = valid; + const auto bytes = field == 0 ? harness.SourceBytes() + : field == 1 ? harness.DestinationBytes() + : harness.CaptureBytes(); + harness.Descriptor()[field] = static_cast(bytes); + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(), "view exceeds its backing tensor"); + } + for (size_t field = 3; field < 7; ++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(); +} + +TYPED_TEST(ConvStateReplayTest, RequiresPreboundDestinationAlias) { + ConvReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({2, 3, 5, 7})); + ASSERT_STATUS_OK(harness.Prepare(4)); + 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(); +} + +TYPED_TEST(ConvStateReplayTest, RejectsWritableInputAliasing) { + ConvReplayHarness harness; + ASSERT_STATUS_OK(harness.Initialize({2, 3, 5, 7})); + ASSERT_STATUS_OK(harness.Prepare(4)); + for (int index = 0; index < 2; ++index) { + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(harness.Run(true, index), "destination backing must be distinct"); + } + harness.VerifyDestinationUnchanged(); +} + +#endif // USE_WEBGPU + +} // namespace +} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/paged_attention_helper_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/paged_attention_helper_test.cc new file mode 100644 index 0000000000000..9c057107273c6 --- /dev/null +++ b/onnxruntime/test/contrib_ops/cuda_kernels/paged_attention_helper_test.cc @@ -0,0 +1,356 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#include "core/framework/tensor_shape.h" +#include "contrib_ops/cpu/bert/paged_attention_helper.h" +#include "contrib_ops/cuda/bert/paged_attention_impl.h" + +namespace onnxruntime { +namespace test { + +namespace { + +struct FakeTensor { + explicit FakeTensor(std::initializer_list dims) : shape_(dims) {} + + const TensorShape& Shape() const { return shape_; } + + private: + TensorShape shape_; +}; + +} // namespace + +TEST(PagedAttentionHelperTest, CheckSequenceLengthTensorsRejectsWrongSeqlensLength) { + FakeTensor cumulative_sequence_length({65}); + FakeTensor seqlens({1}); + + int batch_size = 0; + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthTensors( + &cumulative_sequence_length, &seqlens, batch_size); + + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("seqlens must be shape (batch_size).")); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthTensorsAcceptsMatchingSeqlensLength) { + FakeTensor cumulative_sequence_length({65}); + FakeTensor seqlens({64}); + + int batch_size = 0; + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthTensors( + &cumulative_sequence_length, &seqlens, batch_size); + + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); + EXPECT_EQ(batch_size, 64); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthAndPastValuesRejectsOriginalMalformedMetadata) { + const int32_t shifted_cumulative[] = {1, 2}; + const int32_t valid_past[] = {0}; + EXPECT_FALSE( + onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthAndPastValues( + shifted_cumulative, valid_past, 1, 1, 16, 1) + .IsOK()); + + const int32_t valid_cumulative[] = {0, 1}; + const int32_t negative_past[] = {-1}; + EXPECT_FALSE( + onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthAndPastValues( + valid_cumulative, negative_past, 1, 1, 16, 1) + .IsOK()); +} + +TEST(PagedAttentionHelperTest, CheckBlockTableRejectsZeroWidth) { + FakeTensor block_table({2, 0}); + int max_num_blocks_per_seq = 0; + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckBlockTable( + &block_table, 2, max_num_blocks_per_seq); + + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("dimension 1 must be positive")); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthValuesRejectsPastSeqlensOverflow) { + const int32_t cumulative_seqlens_q[] = {0, 2}; + const int32_t cumulative_seqlens_kv[] = {0, 17}; + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthValues( + cumulative_seqlens_q, cumulative_seqlens_kv, 1, 1, 16, 2); + + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("past_seqlens + query_length exceeds block_table capacity")); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthValuesAcceptsValidInputs) { + const int32_t cumulative_seqlens_q[] = {0, 2, 5}; + const int32_t cumulative_seqlens_kv[] = {0, 12, 18}; + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthValues( + cumulative_seqlens_q, cumulative_seqlens_kv, 2, 2, 16, 5); + + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthValuesRejectsNegativeCumulativeSeqLen) { + const int32_t cumulative_seqlens_q[] = {-1, 2}; + const int32_t cumulative_seqlens_kv[] = {0, 3}; + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthValues( + cumulative_seqlens_q, cumulative_seqlens_kv, 1, 1, 16, 2); + + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("cumulative_seqlens_q must start with 0")); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthValuesRejectsCumulativeNotStartingAtZero) { + const int32_t cumulative_seqlens_q[] = {1, 2}; + const int32_t cumulative_seqlens_kv[] = {0, 1}; + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthValues( + cumulative_seqlens_q, cumulative_seqlens_kv, 1, 1, 16, 2); + + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("cumulative_seqlens_q must start with 0")); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthValuesRejectsKvCumulativeNotStartingAtZero) { + const int32_t cumulative_seqlens_q[] = {0, 1}; + const int32_t cumulative_seqlens_kv[] = {1, 2}; + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthValues( + cumulative_seqlens_q, cumulative_seqlens_kv, 1, 1, 16, 1); + + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("cumulative_seqlens_kv must start with 0")); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthValuesRejectsNegativePastSeqlens) { + const int32_t cumulative_seqlens_q[] = {0, 2}; + const int32_t cumulative_seqlens_kv[] = {0, 1}; + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthValues( + cumulative_seqlens_q, cumulative_seqlens_kv, 1, 1, 16, 2); + + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("past_seqlens values must be non-negative")); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthValuesAllowsFullCacheWithZeroTokens) { + // q_len == 0 (no new tokens) and past_length == max_cache_sequence_length (full cache) + // This should be allowed because there's no write to cache + const int32_t cumulative_seqlens_q[] = {0, 0}; // q_len = 0 + const int32_t cumulative_seqlens_kv[] = {0, 256}; // max_cache_sequence_length = 1 * 256 = 256 + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthValues( + cumulative_seqlens_q, cumulative_seqlens_kv, 1, 1, 256, 0); + + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthValuesRejectsFullCacheWithNewTokens) { + // q_len > 0 (has new tokens) and past_length == max_cache_sequence_length (full cache) + // This should be rejected because we need space for new tokens + const int32_t cumulative_seqlens_q[] = {0, 1}; // q_len = 1 + const int32_t cumulative_seqlens_kv[] = {0, 257}; // past_length = 256 + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthValues( + cumulative_seqlens_q, cumulative_seqlens_kv, 1, 1, 256, 1); + + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("past_seqlens must be less than max_num_blocks_per_seq * block_size when q_len > 0")); +} + +TEST(PagedAttentionHelperTest, CheckSequenceLengthValuesRejectsTokenCountMismatch) { + const int32_t cumulative_seqlens_q[] = {0, 1}; + const int32_t cumulative_seqlens_kv[] = {0, 1}; + + const auto status = onnxruntime::contrib::paged_attention_helper::CheckSequenceLengthValues( + cumulative_seqlens_q, cumulative_seqlens_kv, 1, 1, 16, 2); + + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("must end with token_count")); +} + +TEST(PagedAttentionHelperTest, SanitizeBlockTablePreservesSentinelAndBoundsInvalidIds) { + const std::vector input{-2, -1, 0, 3, 4}; + const std::vector expected{-1, -1, 0, 3, -1}; + int32_t* input_device = nullptr; + int32_t* output_device = nullptr; + ASSERT_EQ(cudaSuccess, cudaMalloc(&input_device, input.size() * sizeof(int32_t))); + ASSERT_EQ(cudaSuccess, cudaMalloc(&output_device, input.size() * sizeof(int32_t))); + auto cleanup = gsl::finally([&]() { + cudaFree(input_device); + cudaFree(output_device); + }); + + ASSERT_EQ(cudaSuccess, + cudaMemcpy(input_device, input.data(), input.size() * sizeof(int32_t), cudaMemcpyHostToDevice)); + const auto status = onnxruntime::contrib::cuda::LaunchSanitizeBlockTable( + input_device, output_device, static_cast(input.size()), 4, nullptr); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + + std::vector actual(input.size()); + ASSERT_EQ(cudaSuccess, + cudaMemcpy(actual.data(), output_device, actual.size() * sizeof(int32_t), cudaMemcpyDeviceToHost)); + EXPECT_EQ(actual, expected); +} + +TEST(PagedAttentionHelperTest, SanitizeBlockTableAllowsZeroElements) { + const auto status = onnxruntime::contrib::cuda::LaunchSanitizeBlockTable( + nullptr, nullptr, 0, 0, nullptr); + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); +} + +TEST(PagedAttentionHelperTest, SanitizeBlockTableRejectsUnsupportedElementCount) { + const auto status = onnxruntime::contrib::cuda::LaunchSanitizeBlockTable( + nullptr, nullptr, static_cast(std::numeric_limits::max()) + 1, + 0, nullptr); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("exceeds the CUDA kernel indexing limit")); +} + +TEST(PagedAttentionHelperTest, SanitizeSequenceLengthsCanonicalizesUnsafeInputs) { + const std::vector cumulative_seqlens_q{5, -2, 1}; + const std::vector past_seqlens{-4, 100}; + constexpr int batch_size = 2; + + int32_t* input_device = nullptr; + int32_t* output_device = nullptr; + void* workspace_device = nullptr; + ASSERT_EQ(cudaSuccess, cudaMalloc(&input_device, 5 * sizeof(int32_t))); + ASSERT_EQ(cudaSuccess, cudaMalloc(&output_device, 10 * sizeof(int32_t))); + size_t workspace_bytes = 0; + const auto workspace_status = onnxruntime::contrib::cuda::GetSanitizeSequenceLengthsWorkspaceSize( + batch_size, workspace_bytes, nullptr); + ASSERT_TRUE(workspace_status.IsOK()) << workspace_status.ErrorMessage(); + ASSERT_EQ(cudaSuccess, cudaMalloc(&workspace_device, workspace_bytes)); + auto cleanup = gsl::finally([&]() { + cudaFree(input_device); + cudaFree(output_device); + cudaFree(workspace_device); + }); + + ASSERT_EQ(cudaSuccess, + cudaMemcpy(input_device, cumulative_seqlens_q.data(), 3 * sizeof(int32_t), cudaMemcpyHostToDevice)); + ASSERT_EQ(cudaSuccess, + cudaMemcpy(input_device + 3, past_seqlens.data(), 2 * sizeof(int32_t), cudaMemcpyHostToDevice)); + + const auto status = onnxruntime::contrib::cuda::LaunchSanitizeSequenceLengths( + output_device, output_device + 3, output_device + 5, output_device + 7, + input_device, input_device + 3, workspace_device, workspace_bytes, + batch_size, 2, 16, 3, nullptr); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + + std::vector actual(10); + ASSERT_EQ(cudaSuccess, + cudaMemcpy(actual.data(), output_device, actual.size() * sizeof(int32_t), cudaMemcpyDeviceToHost)); + EXPECT_EQ(actual, (std::vector{0, 0, 3, 0, 0, 0, 0, 0, 0, 0})); +} + +TEST(PagedAttentionHelperTest, SanitizeSequenceLengthsPreservesPastWithLeadingEvictedBlock) { + const std::vector cumulative_seqlens_q{0, 1}; + const std::vector past_seqlens{31}; + constexpr int batch_size = 1; + + int32_t* input_device = nullptr; + int32_t* output_device = nullptr; + void* workspace_device = nullptr; + ASSERT_EQ(cudaSuccess, cudaMalloc(&input_device, 3 * sizeof(int32_t))); + ASSERT_EQ(cudaSuccess, cudaMalloc(&output_device, 6 * sizeof(int32_t))); + size_t workspace_bytes = 0; + const auto workspace_status = onnxruntime::contrib::cuda::GetSanitizeSequenceLengthsWorkspaceSize( + batch_size, workspace_bytes, nullptr); + ASSERT_TRUE(workspace_status.IsOK()) << workspace_status.ErrorMessage(); + ASSERT_EQ(cudaSuccess, cudaMalloc(&workspace_device, workspace_bytes)); + auto cleanup = gsl::finally([&]() { + cudaFree(input_device); + cudaFree(output_device); + cudaFree(workspace_device); + }); + + ASSERT_EQ(cudaSuccess, + cudaMemcpy(input_device, cumulative_seqlens_q.data(), 2 * sizeof(int32_t), cudaMemcpyHostToDevice)); + ASSERT_EQ(cudaSuccess, + cudaMemcpy(input_device + 2, past_seqlens.data(), sizeof(int32_t), cudaMemcpyHostToDevice)); + + const auto status = onnxruntime::contrib::cuda::LaunchSanitizeSequenceLengths( + output_device, output_device + 2, output_device + 3, output_device + 4, + input_device, input_device + 2, workspace_device, workspace_bytes, + batch_size, 2, 16, 1, nullptr); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + + std::vector actual(6); + ASSERT_EQ(cudaSuccess, + cudaMemcpy(actual.data(), output_device, actual.size() * sizeof(int32_t), cudaMemcpyDeviceToHost)); + EXPECT_EQ(actual, (std::vector{0, 1, 31, 1, 0, 32})); +} + +TEST(PagedAttentionHelperTest, SanitizeSequenceLengthsSupportsBatchAboveFormerBlockScanLimit) { + constexpr int batch_size = 257; + std::vector cumulative_seqlens_q(batch_size + 1); + std::iota(cumulative_seqlens_q.begin(), cumulative_seqlens_q.end(), 0); + const std::vector past_seqlens(batch_size, 0); + + int32_t* input_device = nullptr; + int32_t* output_device = nullptr; + void* workspace_device = nullptr; + const size_t input_count = cumulative_seqlens_q.size() + past_seqlens.size(); + const size_t output_count = + cumulative_seqlens_q.size() * 2 + past_seqlens.size() * 2; + ASSERT_EQ(cudaSuccess, cudaMalloc(&input_device, input_count * sizeof(int32_t))); + ASSERT_EQ(cudaSuccess, cudaMalloc(&output_device, output_count * sizeof(int32_t))); + size_t workspace_bytes = 0; + const auto workspace_status = onnxruntime::contrib::cuda::GetSanitizeSequenceLengthsWorkspaceSize( + batch_size, workspace_bytes, nullptr); + ASSERT_TRUE(workspace_status.IsOK()) << workspace_status.ErrorMessage(); + ASSERT_EQ(cudaSuccess, cudaMalloc(&workspace_device, workspace_bytes)); + auto cleanup = gsl::finally([&]() { + cudaFree(input_device); + cudaFree(output_device); + cudaFree(workspace_device); + }); + + ASSERT_EQ(cudaSuccess, + cudaMemcpy(input_device, cumulative_seqlens_q.data(), + cumulative_seqlens_q.size() * sizeof(int32_t), cudaMemcpyHostToDevice)); + ASSERT_EQ(cudaSuccess, + cudaMemcpy(input_device + cumulative_seqlens_q.size(), past_seqlens.data(), + past_seqlens.size() * sizeof(int32_t), cudaMemcpyHostToDevice)); + + const auto status = onnxruntime::contrib::cuda::LaunchSanitizeSequenceLengths( + output_device, + output_device + cumulative_seqlens_q.size(), + output_device + cumulative_seqlens_q.size() + past_seqlens.size(), + output_device + cumulative_seqlens_q.size() + past_seqlens.size() * 2, + input_device, + input_device + cumulative_seqlens_q.size(), + workspace_device, workspace_bytes, + batch_size, 1, 16, batch_size, nullptr); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + + std::vector actual(output_count); + ASSERT_EQ(cudaSuccess, + cudaMemcpy(actual.data(), output_device, actual.size() * sizeof(int32_t), cudaMemcpyDeviceToHost)); + EXPECT_TRUE(std::equal(cumulative_seqlens_q.begin(), cumulative_seqlens_q.end(), actual.begin())); + EXPECT_TRUE(std::all_of(actual.begin() + cumulative_seqlens_q.size(), + actual.begin() + cumulative_seqlens_q.size() + past_seqlens.size(), + [](int32_t value) { return value == 0; })); + EXPECT_TRUE(std::all_of( + actual.begin() + cumulative_seqlens_q.size() + past_seqlens.size(), + actual.begin() + cumulative_seqlens_q.size() + past_seqlens.size() * 2, + [](int32_t value) { return value == 1; })); + EXPECT_TRUE(std::equal(cumulative_seqlens_q.begin(), cumulative_seqlens_q.end(), + actual.begin() + cumulative_seqlens_q.size() + past_seqlens.size() * 2)); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc index d32c7fc224fd8..eca0a33410af0 100644 --- a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc +++ b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc @@ -3,6 +3,8 @@ #include "embed_layer_norm_test_vectors.h" #include "gtest/gtest.h" +#include +#include #include "test/common/tensor_op_test_utils.h" #include "test/common/cuda_op_test_utils.h" #include "test/providers/provider_test_utils.h" @@ -10,6 +12,118 @@ namespace onnxruntime { namespace test { +#if defined(USE_CUDA) +template +static void CheckCudaHighOffsetVariance() { + if (!HasCudaEnvironment(0)) GTEST_SKIP() << "CUDA EP unavailable"; + + constexpr float epsilon = 1e-5f; + const auto native_value = [](float value) { + if constexpr (std::is_same_v) { + return value; + } else { + return T(value).ToFloat(); + } + }; + for (int hidden_size : {1, 3, 4, 5, 128, 129, 320, 768, 769, 8192, 8193}) { + SCOPED_TRACE(hidden_size); + for (bool has_segment : {false, true}) { + SCOPED_TRACE(has_segment); + const std::vector shape{1, 2, hidden_size}; + const std::vector table_shape{2, hidden_size}; + const std::vector parameter_shape{hidden_size}; + std::vector pattern{1000.01f, 1000.02f, 999.99f, 1000.0f}; + float position_value = 128.0f; + if constexpr (std::is_same_v) { + pattern = {60000.0f, 60032.0f, 59968.0f, 60000.0f}; + position_value = -59872.0f; + } + std::vector word(2 * hidden_size); + std::vector position(2 * hidden_size, has_segment ? position_value : 0.0f); + std::vector segment(hidden_size, 512.0f); + std::vector gamma(hidden_size); + std::vector beta(hidden_size); + std::vector sum(2 * hidden_size); + std::vector expected(2 * hidden_size); + for (int i = 0; i < hidden_size; ++i) { + word[i] = pattern[i % 4]; + word[hidden_size + i] = std::is_same_v ? -1000.0f : pattern[0]; + if constexpr (std::is_same_v) { + segment[i] = i % 2 == 0 ? 0.25f : 0.5f; + } + gamma[i] = i % 2 == 0 ? 0.5f : -1.25f; + beta[i] = i % 2 == 0 ? 0.25f : -0.75f; + } + for (int row = 0; row < 2; ++row) { + double mean = 0.0; + for (int i = 0; i < hidden_size; ++i) { + const int index = row * hidden_size + i; + const float value = native_value(native_value(word[index] + (has_segment ? segment[i] : 0.0f)) + + position[index]); + sum[index] = value; + mean += static_cast(value); + } + mean /= hidden_size; + double variance = 0.0; + for (int i = 0; i < hidden_size; ++i) { + const double delta = static_cast(sum[row * hidden_size + i]) - mean; + variance += delta * delta; + } + variance /= hidden_size; + const double inv_std = 1.0 / std::sqrt(variance + static_cast(epsilon)); + for (int i = 0; i < hidden_size; ++i) { + const int index = row * hidden_size + i; + expected[index] = static_cast((static_cast(sum[index]) - mean) * + inv_std * gamma[i] + + beta[i]); + } + } + + OpTester tester("EmbedLayerNormalization", 1, kMSDomain); + tester.AddAttribute("epsilon", epsilon); + tester.AddInput("input_ids", {1, 2}, {0, 1}); + if (has_segment) { + tester.AddInput("segment_ids", {1, 2}, {0, 0}); + } else { + tester.AddOptionalInputEdge(); + } + tester.AddInput("word_embedding", table_shape, std::vector(word.begin(), word.end())); + tester.AddInput("position_embedding", table_shape, std::vector(position.begin(), position.end())); + if (has_segment) { + tester.AddInput("segment_embedding", {1, hidden_size}, std::vector(segment.begin(), segment.end())); + } else { + tester.AddOptionalInputEdge(); + } + tester.AddInput("gamma", parameter_shape, std::vector(gamma.begin(), gamma.end())); + tester.AddInput("beta", parameter_shape, std::vector(beta.begin(), beta.end())); + tester.AddInput("mask", {1, 2}, {1, 0}); + tester.AddOutput("output", shape, std::vector(expected.begin(), expected.end())); + tester.AddOutput("mask_index", {1}, {1}); + tester.AddOutput("embedding_sum", shape, std::vector(sum.begin(), sum.end())); + if constexpr (std::is_same_v) { + tester.SetOutputAbsErr("output", 1e-4f); + tester.SetOutputRelErr("output", 1e-4f); + } + tester.SetOutputAbsErr("embedding_sum", 0.0f); + tester.SetOutputRelErr("embedding_sum", 0.0f); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + } +} + +TEST(EmbedLayerNormTest, CudaHighOffsetVariance) { + CheckCudaHighOffsetVariance(); +} + +TEST(EmbedLayerNormTest, CudaFloat16HighOffsetVariance) { + CheckCudaHighOffsetVariance(); +} +#endif + static void RunTest(const embedlayernorm::OpData& data, bool use_float16 = false, bool sum_output = false, diff --git a/onnxruntime/test/contrib_ops/function_ops_test.cc b/onnxruntime/test/contrib_ops/function_ops_test.cc index 2036328053d17..5df6136568174 100644 --- a/onnxruntime/test/contrib_ops/function_ops_test.cc +++ b/onnxruntime/test/contrib_ops/function_ops_test.cc @@ -113,5 +113,25 @@ TEST_F(ContribFunExpansionTest, FastGeluWithoutBias) { CheckFastGelu(false); } +void CheckSpaceDepthFunction(const char* op_type, std::vector input_shape, + int64_t blocksize, const char* mode) { + FunctionTestCase test_case(op_type, kOnnxDomain); + test_case.AddOpset(kOnnxDomain, 28); + test_case.AddInput("input", input_shape); + test_case.AddOutput("output"); + test_case.AddAttribute("blocksize", blocksize); + test_case.AddAttribute("mode", mode); + test_case.RunTest(/*allow_released_opsets_only*/ false); +} + +TEST_F(ContribFunExpansionTest, SpaceDepthOpset28FunctionParity) { + for (const auto* mode : {"DCR", "CRD"}) { + CheckSpaceDepthFunction("SpaceToDepth", {1, 2, 4, 6}, 2, mode); + CheckSpaceDepthFunction("SpaceToDepth", {1, 2, 6, 9}, 3, mode); + CheckSpaceDepthFunction("DepthToSpace", {1, 8, 2, 3}, 2, mode); + CheckSpaceDepthFunction("DepthToSpace", {1, 18, 2, 3}, 3, mode); + } +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_fpa_intb_layout_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_fpa_intb_layout_test.cc index 25f0947be58bb..1e78d94a8e61b 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_fpa_intb_layout_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_fpa_intb_layout_test.cc @@ -30,7 +30,7 @@ namespace { // the first tile. N is a multiple of 128 so 2-bit weights are fpA_intB-eligible. constexpr int64_t kN = 128, kK = 256; -void ExpectExactDequantizedWeights(int64_t bits) { +void ExpectExactDequantizedWeights(int64_t bits, bool deterministic = false) { SCOPED_TRACE("bits=" + std::to_string(bits)); ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_FPA_INTB_GEMM", "1"}}}; @@ -83,6 +83,7 @@ void ExpectExactDequantizedWeights(int64_t bits) { eps.emplace_back(DefaultCudaExecutionProvider()); SessionOptions session_options; session_options.use_per_session_threads = false; + session_options.use_deterministic_compute = deterministic; ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); test.Config(session_options); test.ConfigEps(std::move(eps)); @@ -101,6 +102,13 @@ TEST(MatMulNBitsFpAIntBLayout, Int2WeightsRoundTripExactly) { ExpectExactDequantizedWeights(2); } +TEST(MatMulNBitsFpAIntBLayout, DeterministicWeightsRoundTripExactly) { + if (!HasCudaEnvironment(800)) GTEST_SKIP() << "fpA_intB layout test requires SM80+"; + for (int64_t bits : {2, 4, 8}) { + ExpectExactDequantizedWeights(bits, true); + } +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/packed_attention_op_test.cc b/onnxruntime/test/contrib_ops/packed_attention_op_test.cc index 1ae779c264fed..3664391110e8c 100644 --- a/onnxruntime/test/contrib_ops/packed_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_attention_op_test.cc @@ -1,13 +1,23 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include +#include + #include "core/common/narrow.h" +#include "core/graph/model.h" #include "core/platform/env_var_utils.h" +#include "core/providers/cuda/cuda_provider_options.h" +#include "core/session/IOBinding.h" +#include "core/session/inference_session.h" #include "gtest/gtest.h" #include "test/common/tensor_op_test_utils.h" #include "test/common/cuda_op_test_utils.h" #include "test/providers/provider_test_utils.h" +#include "test/unittest_util/framework_test_utils.h" #include "test/util/include/scoped_env_vars.h" +#include "test/util/include/test_environment.h" #include "contrib_ops/cpu/bert/attention_common.h" #include "test/contrib_ops/attention_op_test_helper.h" @@ -595,6 +605,204 @@ TEST(PackedAttentionTest, PackedBatchWithQKV) { qkv_sizes); } +TEST(PackedAttentionTest, InvalidTokenOffset_Unfused_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "PackedAttention token offset validation requires a CUDA device."; + } + + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableTrtFlashAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedSelfAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}}}; + + constexpr int kTokenCount = 1; + constexpr int kBatchSize = 1; + constexpr int kSequenceLength = 2; + constexpr int kHiddenSize = 8; + + for (const std::vector& token_offset : { + std::vector{-1, 1}, + std::vector{0, kBatchSize * kSequenceLength}}) { + SCOPED_TRACE(token_offset[0]); + SCOPED_TRACE(token_offset[1]); + + OpTester tester("PackedAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", 1); + tester.AddInput("input", {kTokenCount, kHiddenSize}, + std::vector(kTokenCount * kHiddenSize)); + tester.AddInput("weight", {kHiddenSize, 3 * kHiddenSize}, + std::vector(3 * kHiddenSize * kHiddenSize)); + tester.AddInput("bias", {3 * kHiddenSize}, + std::vector(3 * kHiddenSize)); + tester.AddInput("token_offset", {kBatchSize, kSequenceLength}, token_offset); + tester.AddInput("cumulative_sequence_length", {kBatchSize + 1}, {0, kTokenCount}); + tester.AddOutput("output", {kTokenCount, kHiddenSize}, + std::vector(kTokenCount * kHiddenSize)); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run( + OpTester::ExpectResult::kExpectFailure, + "PackedAttention token_offset values must be in [0, B * S).", + {}, nullptr, &execution_providers); + } +} + +TEST(PackedAttentionTest, InvalidTokenOffset_UnfusedCudaGraphReplay) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "PackedAttention CUDA graph test requires a CUDA device."; + } + + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableTrtFlashAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedSelfAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}}}; + + constexpr int64_t kTokenCount = 2; + constexpr int64_t kBatchSize = 1; + constexpr int64_t kSequenceLength = 2; + constexpr int64_t kHiddenSize = 8; + + auto make_tensor_type = [](int32_t element_type, std::initializer_list shape) { + ONNX_NAMESPACE::TypeProto type; + auto* tensor_type = type.mutable_tensor_type(); + tensor_type->set_elem_type(element_type); + for (int64_t dimension : shape) { + tensor_type->mutable_shape()->add_dim()->set_dim_value(dimension); + } + return type; + }; + + auto input_type = make_tensor_type( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT, {kTokenCount, kHiddenSize}); + auto weights_type = make_tensor_type( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT, {kHiddenSize, 3 * kHiddenSize}); + auto bias_type = make_tensor_type( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT, {3 * kHiddenSize}); + auto token_offset_type = make_tensor_type( + ONNX_NAMESPACE::TensorProto_DataType_INT32, {kBatchSize, kSequenceLength}); + auto cumulative_sequence_length_type = make_tensor_type( + ONNX_NAMESPACE::TensorProto_DataType_INT32, {kBatchSize + 1}); + auto output_type = make_tensor_type( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT, {kTokenCount, kHiddenSize}); + + Model model("packed_attention_cuda_graph", true, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 17}, {kMSDomain, 1}}, + {}, DefaultLoggingManager().DefaultLogger(), ModelOptions(true, true)); + auto& graph = model.MainGraph(); + std::vector inputs{ + &graph.GetOrCreateNodeArg("input", &input_type), + &graph.GetOrCreateNodeArg("weights", &weights_type), + &graph.GetOrCreateNodeArg("bias", &bias_type), + &graph.GetOrCreateNodeArg("token_offset", &token_offset_type), + &graph.GetOrCreateNodeArg("cumulative_sequence_length", &cumulative_sequence_length_type)}; + std::vector outputs{ + &graph.GetOrCreateNodeArg("output", &output_type)}; + auto& node = graph.AddNode( + "packed_attention", "PackedAttention", "", inputs, outputs, nullptr, kMSDomain); + node.AddAttribute("num_heads", int64_t{1}); + ASSERT_STATUS_OK(graph.Resolve()); + + std::string model_data; + ASSERT_TRUE(model.ToProto().SerializeToString(&model_data)); + + OrtCUDAProviderOptionsV2 provider_options{}; + provider_options.do_copy_in_default_stream = true; + provider_options.enable_cuda_graph = true; + auto cuda_ep = CudaExecutionProviderWithOptions(&provider_options); + ASSERT_NE(cuda_ep, nullptr); + IExecutionProvider* cuda_ep_ptr = cuda_ep.get(); + + SessionOptions session_options; + InferenceSession session(session_options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(cuda_ep))); + std::istringstream model_stream(model_data); + ASSERT_STATUS_OK(session.Load(model_stream)); + ASSERT_STATUS_OK(session.Initialize()); + + auto gpu_allocators = cuda_ep_ptr->CreatePreferredAllocators(); + auto gpu_allocator = std::find_if(gpu_allocators.begin(), gpu_allocators.end(), [](const auto& allocator) { + return allocator->Info().device.Type() == OrtDevice::GPU && + allocator->Info().alloc_type != OrtAllocatorType::OrtReadOnlyAllocator && + allocator->Info().mem_type == OrtMemTypeDefault; + }); + ASSERT_NE(gpu_allocator, gpu_allocators.end()); + auto allocator = session.GetAllocator((*gpu_allocator)->Info()); + ASSERT_NE(allocator, nullptr); + auto cpu_allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + + auto make_gpu_value = [&](const auto& values, const TensorShape& shape) { + using Element = typename std::decay_t::value_type; + Tensor cpu_tensor( + DataTypeImpl::GetType(), shape, + const_cast(values.data()), cpu_allocator->Info()); + Tensor gpu_tensor(DataTypeImpl::GetType(), shape, allocator); + ORT_THROW_IF_ERROR(cuda_ep_ptr->GetDataTransfer()->CopyTensor(cpu_tensor, gpu_tensor)); + OrtValue result; + Tensor::InitOrtValue(std::move(gpu_tensor), result); + return result; + }; + + auto input = make_gpu_value( + std::vector(kTokenCount * kHiddenSize), + TensorShape{kTokenCount, kHiddenSize}); + auto weights = make_gpu_value( + std::vector(kHiddenSize * 3 * kHiddenSize), + TensorShape{kHiddenSize, 3 * kHiddenSize}); + std::vector bias_data(3 * kHiddenSize); + std::fill(bias_data.begin() + 2 * kHiddenSize, bias_data.end(), 1.0f); + auto bias = make_gpu_value(bias_data, TensorShape{3 * kHiddenSize}); + auto token_offset = make_gpu_value( + std::vector{0, 1}, TensorShape{kBatchSize, kSequenceLength}); + auto cumulative_sequence_length = make_gpu_value( + std::vector{0, 2}, TensorShape{kBatchSize + 1}); + Tensor output_tensor( + DataTypeImpl::GetType(), + TensorShape{kTokenCount, kHiddenSize}, allocator); + OrtValue output; + Tensor::InitOrtValue(std::move(output_tensor), output); + + std::unique_ptr binding; + ASSERT_STATUS_OK(session.NewIOBinding(&binding)); + ASSERT_STATUS_OK(binding->BindInput("input", input)); + ASSERT_STATUS_OK(binding->BindInput("weights", weights)); + ASSERT_STATUS_OK(binding->BindInput("bias", bias)); + ASSERT_STATUS_OK(binding->BindInput("token_offset", token_offset)); + ASSERT_STATUS_OK(binding->BindInput("cumulative_sequence_length", cumulative_sequence_length)); + ASSERT_STATUS_OK(binding->BindOutput("output", output)); + ASSERT_STATUS_OK(binding->SynchronizeInputs()); + + RunOptions run_options; + ASSERT_STATUS_OK(run_options.config_options.AddConfigEntry("gpu_graph_id", "1")); + for (int i = 0; i < 3; ++i) { + ASSERT_STATUS_OK(session.Run(run_options, *binding)); + } + ASSERT_TRUE(cuda_ep_ptr->IsGraphCaptured(1)); + + const std::vector invalid_offsets{-1, kBatchSize * kSequenceLength}; + Tensor invalid_offsets_tensor( + DataTypeImpl::GetType(), TensorShape{kBatchSize, kSequenceLength}, + const_cast(invalid_offsets.data()), cpu_allocator->Info()); + ASSERT_STATUS_OK(cuda_ep_ptr->GetDataTransfer()->CopyTensor( + invalid_offsets_tensor, *token_offset.GetMutable())); + ASSERT_STATUS_OK(binding->SynchronizeInputs()); + ASSERT_STATUS_OK(session.Run(run_options, *binding)); + ASSERT_STATUS_OK(binding->SynchronizeOutputs()); + + Tensor cpu_output( + DataTypeImpl::GetType(), + TensorShape{kTokenCount, kHiddenSize}, cpu_allocator); + ASSERT_STATUS_OK(cuda_ep_ptr->GetDataTransfer()->CopyTensor( + output.Get(), cpu_output)); + for (float value : cpu_output.DataAsSpan()) { + EXPECT_EQ(value, 0.0f); + } +} + static void RunModelWithRandomInput( int64_t batch_size, int64_t sequence_length, diff --git a/onnxruntime/test/contrib_ops/paged_attention_op_test.cc b/onnxruntime/test/contrib_ops/paged_attention_op_test.cc index 2a698262a90ca..258ef5b0a3ff3 100644 --- a/onnxruntime/test/contrib_ops/paged_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/paged_attention_op_test.cc @@ -84,12 +84,18 @@ struct IoBindingCase { bool enable_cuda_graph = false; bool irregular_layout = false; bool discriminating_attention = false; + bool poison_first_value_cache_page = false; std::vector past_seqlens; std::vector> replay_past_seqlens; std::vector cumulative_seqlens_q; std::vector block_table; std::vector attention_metadata; std::string expected_error; + bool allow_malformed_sequence_metadata = false; + bool skip_reference_check = false; + bool verify_malformed_cache_unchanged = false; + bool verify_malformed_output_finite = false; + bool verify_malformed_output_zero = false; }; // Masked positions get zero probability. Uses fp32 throughout to establish a @@ -318,29 +324,36 @@ void RunIoBindingCase(std::unique_ptr execution_provider, ASSERT_TRUE(c.past_seqlens.empty() || c.past_seqlens.size() == static_cast(batch_size)); ASSERT_TRUE(c.past_seqlens.empty() || c.replay_past_seqlens.empty()); - for (int32_t initial_length : c.past_seqlens) { - ASSERT_GE(initial_length, 0); - ASSERT_GT(max_num_blocks_per_seq, initial_length / block_size); - } - for (const auto& replay_lengths : c.replay_past_seqlens) { - ASSERT_EQ(replay_lengths.size(), static_cast(batch_size)); - for (int32_t replay_length : replay_lengths) { - ASSERT_GE(replay_length, 0); - ASSERT_GT(max_num_blocks_per_seq, replay_length / block_size); + if (!c.allow_malformed_sequence_metadata) { + for (int32_t initial_length : c.past_seqlens) { + ASSERT_GE(initial_length, 0); + ASSERT_GT(max_num_blocks_per_seq, initial_length / block_size); + } + for (const auto& replay_lengths : c.replay_past_seqlens) { + ASSERT_EQ(replay_lengths.size(), static_cast(batch_size)); + for (int32_t replay_length : replay_lengths) { + ASSERT_GE(replay_length, 0); + ASSERT_GT(max_num_blocks_per_seq, replay_length / block_size); + } } } ASSERT_TRUE(!c.replay_past_seqlens.empty() || !c.past_seqlens.empty() || max_num_blocks_per_seq > past_seqlen / block_size); - ASSERT_LE(batch_size * max_num_blocks_per_seq, num_blocks); + if (!c.allow_malformed_sequence_metadata) { + ASSERT_LE(batch_size * max_num_blocks_per_seq, num_blocks); + } ASSERT_EQ(num_heads % kv_num_heads, 0); ASSERT_TRUE(c.block_table.empty() || c.block_table.size() == static_cast(batch_size * max_num_blocks_per_seq)); ASSERT_TRUE(c.cumulative_seqlens_q.empty() || - (c.cumulative_seqlens_q.size() == static_cast(batch_size + 1) && - c.cumulative_seqlens_q.front() == 0 && c.cumulative_seqlens_q.back() == token_count)); - for (int32_t block_id : c.block_table) { - ASSERT_GE(block_id, 0); - ASSERT_LT(block_id, num_blocks); + c.cumulative_seqlens_q.size() == static_cast(batch_size + 1)); + if (!c.allow_malformed_sequence_metadata) { + ASSERT_TRUE(c.cumulative_seqlens_q.empty() || + (c.cumulative_seqlens_q.front() == 0 && c.cumulative_seqlens_q.back() == token_count)); + for (int32_t block_id : c.block_table) { + ASSERT_GE(block_id, -1); + ASSERT_LT(block_id, num_blocks); + } } std::unordered_map domain_to_version = {{onnxruntime::kMSDomain, 1}}; @@ -622,6 +635,10 @@ void RunIoBindingCase(std::unique_ptr execution_provider, } } } + if (c.poison_first_value_cache_page) { + std::fill_n(value_cache_data.begin(), block_size * kv_hidden_size, + MLFloat16(std::numeric_limits::quiet_NaN())); + } std::vector query_data_bf16; std::vector key_data_bf16; std::vector value_data_bf16; @@ -785,6 +802,48 @@ void RunIoBindingCase(std::unique_ptr execution_provider, return; } ASSERT_STATUS_OK(run_status); + if (c.skip_reference_check) { + if (c.verify_malformed_cache_unchanged && run_index + 1 == run_count) { + ASSERT_FALSE(c.bf16_query); + ASSERT_FALSE(quantized_cache); + + Tensor cpu_key_cache( + DataTypeImpl::GetType(), + TensorShape({num_blocks, block_size, kv_num_heads, head_size}), cpu_alloc); + Tensor cpu_value_cache( + DataTypeImpl::GetType(), + TensorShape({num_blocks, block_size, kv_num_heads, head_size}), cpu_alloc); + ORT_THROW_IF_ERROR( + execution_provider_ptr->GetDataTransfer()->CopyTensor( + key_cache_value.Get(), cpu_key_cache)); + ORT_THROW_IF_ERROR( + execution_provider_ptr->GetDataTransfer()->CopyTensor( + value_cache_value.Get(), cpu_value_cache)); + const auto actual_key_cache = cpu_key_cache.DataAsSpan(); + const auto actual_value_cache = cpu_value_cache.DataAsSpan(); + ASSERT_EQ(actual_key_cache.size(), key_cache_data.size()); + ASSERT_EQ(actual_value_cache.size(), value_cache_data.size()); + for (size_t i = 0; i < key_cache_data.size(); ++i) { + EXPECT_EQ(actual_key_cache[i].ToFloat(), key_cache_data[i].ToFloat()); + EXPECT_EQ(actual_value_cache[i].ToFloat(), value_cache_data[i].ToFloat()); + } + } + if (c.verify_malformed_output_finite && run_index + 1 == run_count) { + Tensor cpu_output( + DataTypeImpl::GetType(), + TensorShape({token_count, hidden_size}), cpu_alloc); + ORT_THROW_IF_ERROR( + execution_provider_ptr->GetDataTransfer()->CopyTensor( + output_value.Get(), cpu_output)); + for (const MLFloat16 value : cpu_output.DataAsSpan()) { + EXPECT_TRUE(std::isfinite(value.ToFloat())); + if (c.verify_malformed_output_zero) { + EXPECT_EQ(value.ToFloat(), 0.0f); + } + } + } + continue; + } Tensor cpu_output(c.bf16_query ? DataTypeImpl::GetType() : DataTypeImpl::GetType(), TensorShape({token_count, hidden_size}), cpu_alloc); @@ -833,6 +892,10 @@ void RunIoBindingCase(std::unique_ptr execution_provider, for (int slot = first_visible_slot; slot <= last_visible_slot; ++slot) { const int block_id = block_table_data[b * max_num_blocks_per_seq + slot / block_size]; + if (block_id < 0) { + scores[slot] = -std::numeric_limits::infinity(); + continue; + } float dot = 0.0f; for (int dim = 0; dim < head_size; ++dim) { const int query_index = (token * num_heads + q_head) * head_size + dim; @@ -857,6 +920,9 @@ void RunIoBindingCase(std::unique_ptr execution_provider, for (int slot = first_visible_slot; slot <= last_visible_slot; ++slot) { const int block_id = block_table_data[b * max_num_blocks_per_seq + slot / block_size]; + if (block_id < 0) { + continue; + } const int cache_index = CacheIndex(block_id, slot % block_size, kv_head, dim, block_size, kv_num_heads, head_size); const float value_element = native_bf16_cache ? value_cache_bf16[cache_index].ToFloat() @@ -884,6 +950,10 @@ void RunIoBindingCase(std::unique_ptr execution_provider, EXPECT_TRUE(execution_provider_ptr->IsGraphCaptured(1)); } + if (c.skip_reference_check) { + return; + } + const auto& outputs = io_binding->GetOutputs(); MLDataType cache_data_type = c.int8_cache ? DataTypeImpl::GetType() : c.bf16_query ? DataTypeImpl::GetType() @@ -1831,6 +1901,46 @@ TEST(PagedAttention, Cuda_FlashSplitKvLongContext) { #endif } +TEST(PagedAttention, Cuda_FlashUnmappedPageIgnoresPoisonedValueCache) { +#if defined(USE_FLASH_ATTENTION) + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableFlashAttention, "0"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}, + {onnxruntime::contrib::attention::kDisableDecoderAttention, "1"}, + {onnxruntime::contrib::attention::kEnableCudnnFlashAttention, "0"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + if (GetCudaArchitecture() < 800) { + GTEST_SKIP() << "Flash Attention requires compute capability 8.0 or later."; + } + + IoBindingCase c; + c.num_heads = 2; + c.kv_num_heads = 1; + c.head_size = 128; + c.num_blocks = 2; + c.max_num_blocks_per_seq = 2; + c.past_seqlens = {257}; + c.block_table = {-1, 1}; + c.attention_metadata = {1, 512, 258}; + c.poison_first_value_cache_page = true; + c.skip_reference_check = true; + c.verify_malformed_output_finite = true; + + testing::internal::CaptureStdout(); + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); + const std::string debug_output = testing::internal::GetCapturedStdout(); + + EXPECT_NE(debug_output.find("SdpaKernel=FLASH_ATTENTION"), std::string::npos) << debug_output; +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + TEST(PagedAttention, Cuda_FlashSplitKvCudaGraphReplay) { #if defined(USE_FLASH_ATTENTION) ScopedEnvironmentVariables scoped_env_vars{ @@ -1885,6 +1995,180 @@ TEST(PagedAttention, Cuda_FlashSplitKvCudaGraphReplay) { #endif } +TEST(PagedAttention, CudaGraphWithoutAttentionMetadata) { + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + + OrtCUDAProviderOptionsV2 provider_options{}; + provider_options.enable_cuda_graph = true; + + IoBindingCase c; + c.enable_cuda_graph = true; + c.replay_past_seqlens = {{4}, {5}, {6}, {7}}; + RunIoBindingCase(CudaExecutionProviderWithOptions(&provider_options), + kCudaExecutionProvider, true, false, c); +} + +TEST(PagedAttention, CudaMalformedSequenceMetadataIsSanitizedWithoutReadback) { + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + + IoBindingCase c; + c.cumulative_seqlens_q = {0, 1}; + c.past_seqlens = {-4}; + c.block_table = {-2}; + c.allow_malformed_sequence_metadata = true; + c.skip_reference_check = true; + c.verify_malformed_cache_unchanged = true; + c.verify_malformed_output_finite = true; + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); +} + +TEST(PagedAttention, CudaMalformedSequenceMetadataIsSanitizedWithMetadata) { + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + + IoBindingCase c; + c.cumulative_seqlens_q = {0, 1}; + c.past_seqlens = {-4}; + c.block_table = {-2}; + c.attention_metadata = {1, 1}; + c.allow_malformed_sequence_metadata = true; + c.skip_reference_check = true; + c.verify_malformed_cache_unchanged = true; + c.verify_malformed_output_finite = true; + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); +} + +TEST(PagedAttention, CudaOversizedQueryIsSuppressedWithoutReadback) { +#if defined(USE_FLASH_ATTENTION) + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + if (GetCudaArchitecture() < 800) { + GTEST_SKIP() << "Native FlashAttention requires compute capability 8.0 or later."; + } + + IoBindingCase c; + c.token_count = 257; + c.head_size = 128; + c.block_size = 256; + c.num_blocks = 1; + c.max_num_blocks_per_seq = 1; + c.cumulative_seqlens_q = {0, 257}; + c.past_seqlens = {0}; + c.block_table = {0}; + c.allow_malformed_sequence_metadata = true; + c.skip_reference_check = true; + c.verify_malformed_cache_unchanged = true; + c.verify_malformed_output_finite = true; + c.verify_malformed_output_zero = true; + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); +#else + GTEST_SKIP() << "FlashAttention is not enabled in this build."; +#endif +} + +TEST(PagedAttention, CudaGraphOversizedQueryIsSuppressed) { +#if defined(USE_FLASH_ATTENTION) + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + if (GetCudaArchitecture() < 800) { + GTEST_SKIP() << "Native FlashAttention requires compute capability 8.0 or later."; + } + + OrtCUDAProviderOptionsV2 provider_options{}; + provider_options.enable_cuda_graph = true; + + IoBindingCase c; + c.token_count = 257; + c.head_size = 128; + c.block_size = 256; + c.num_blocks = 1; + c.max_num_blocks_per_seq = 1; + c.cumulative_seqlens_q = {0, 257}; + c.replay_past_seqlens = {{0}, {0}, {0}, {0}}; + c.block_table = {0}; + c.attention_metadata = {257, 256}; + c.enable_cuda_graph = true; + c.allow_malformed_sequence_metadata = true; + c.skip_reference_check = true; + c.verify_malformed_cache_unchanged = true; + c.verify_malformed_output_finite = true; + c.verify_malformed_output_zero = true; + RunIoBindingCase( + CudaExecutionProviderWithOptions(&provider_options), + kCudaExecutionProvider, true, false, c); +#else + GTEST_SKIP() << "FlashAttention is not enabled in this build."; +#endif +} + +TEST(PagedAttention, CudaRejectsZeroPhysicalCacheBlocks) { + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + + IoBindingCase c; + c.num_blocks = 0; + c.max_num_blocks_per_seq = 1; + c.block_table = {-1}; + c.allow_malformed_sequence_metadata = true; + c.expected_error = "requires at least one physical cache block"; + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); +} + +TEST(PagedAttention, CudaGraphFlashMalformedPageIsMasked) { +#if defined(USE_FLASH_ATTENTION) +#if defined(ORT_UNIT_TEST_ENABLE_DYNAMIC_PLUGIN_EP_USAGE) || defined(ORT_QUICK_BUILD) || defined(EXCLUDE_SM_80) + GTEST_SKIP() << "This build does not expose the native FlashAttention geometry used by this test."; +#endif + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "FlashAttention requires a CUDA device with compute capability 8.0 or newer."; + } + + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableFlashAttention, "0"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}, + {onnxruntime::contrib::attention::kDisableDecoderAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}, + {"ORT_ENABLE_XQA", "0"}, + {"ORT_ENABLE_CUDNN_FLASH_ATTENTION", "0"}}}; + + OrtCUDAProviderOptionsV2 provider_options{}; + provider_options.enable_cuda_graph = true; + + IoBindingCase malformed; + malformed.batch_size = 1; + malformed.token_count = 1; + malformed.num_heads = 2; + malformed.kv_num_heads = 1; + malformed.head_size = 64; + malformed.block_size = 256; + malformed.num_blocks = 4; + malformed.max_num_blocks_per_seq = 4; + malformed.cumulative_seqlens_q = {0, 1}; + malformed.replay_past_seqlens = {{768}, {768}, {768}, {768}}; + malformed.block_table = {0, malformed.num_blocks, malformed.num_blocks, malformed.num_blocks}; + malformed.attention_metadata = {1, 769}; + malformed.allow_malformed_sequence_metadata = true; + malformed.skip_reference_check = true; + malformed.verify_malformed_cache_unchanged = true; + malformed.verify_malformed_output_finite = true; + malformed.enable_cuda_graph = true; + RunIoBindingCase( + CudaExecutionProviderWithOptions(&provider_options), + kCudaExecutionProvider, true, false, malformed); +#else + GTEST_SKIP() << "FlashAttention is not enabled in this build."; +#endif +} + TEST(PagedAttention, Cuda_FlashSplitKvInt8Cache) { #if defined(USE_FLASH_ATTENTION) ScopedEnvironmentVariables scoped_env_vars{ @@ -2023,12 +2307,10 @@ TEST(PagedAttention, Cuda_CudnnPagedDispatchWhenEnabled) { } TEST(PagedAttention, Cuda_CudnnPagedRunsWhenPreferredXqaIsRejected) { - ScopedEnvironmentVariables scoped_env_vars{ + ScopedEnvironmentVariables common_env_vars{ EnvVarMap{ {onnxruntime::contrib::attention::kEnableCudnnFlashAttention, "1"}, - {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}, - {"ORT_ENABLE_XQA_NATIVE_KV", "1"}, - {"ORT_TEST_ONLY_PAGED_ATTENTION_XQA_SHARED_MEMORY_LIMIT", "0"}}}; + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; if (DefaultCudaExecutionProvider() == nullptr) { GTEST_SKIP() << "CUDA EP not available."; @@ -2041,14 +2323,27 @@ TEST(PagedAttention, Cuda_CudnnPagedRunsWhenPreferredXqaIsRejected) { c.num_heads = 6; c.head_size = 256; + { + ScopedEnvironmentVariables disable_xqa{ + EnvVarMap{{"ORT_ENABLE_XQA_NATIVE_KV", "0"}}}; + testing::internal::CaptureStdout(); + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); + const std::string probe_output = testing::internal::GetCapturedStdout(); + if (probe_output.find("SdpaKernel=CUDNN_FLASH_ATTENTION") == std::string::npos) { + GTEST_SKIP() << "cuDNN paged SDPA is not runnable for the target XQA shape.\n" + << probe_output; + } + } + + ScopedEnvironmentVariables reject_xqa{ + EnvVarMap{ + {"ORT_ENABLE_XQA_NATIVE_KV", "1"}, + {"ORT_TEST_ONLY_PAGED_ATTENTION_XQA_SHARED_MEMORY_LIMIT", "0"}}}; testing::internal::CaptureStdout(); RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); const std::string debug_output = testing::internal::GetCapturedStdout(); - if (debug_output.find("SdpaKernel=CUDNN_FLASH_ATTENTION") == std::string::npos) { - GTEST_SKIP() << "cuDNN paged SDPA is not runnable in this build/device configuration.\n" - << debug_output; - } + EXPECT_NE(debug_output.find("SdpaKernel=CUDNN_FLASH_ATTENTION"), std::string::npos) << debug_output; EXPECT_EQ(debug_output.find("SdpaKernel=XQA"), std::string::npos) << debug_output; EXPECT_EQ(debug_output.find("SdpaKernel=FLASH_ATTENTION"), std::string::npos) << debug_output; } @@ -2084,17 +2379,19 @@ TEST(PagedAttention, Cuda_CudnnPagedCudaGraphReplay) { IoBindingCase c = MakeCudnnPagedDecodeCase(); c.enable_cuda_graph = true; c.irregular_layout = true; + c.past_seqlen = 511; + c.block_table = {-1, 3}; // Five Runs: Runs 1-2 are warm-ups (populate the thread_local plan cache), Run 3 begins capture, // Runs 4-5 replay. All Runs share the same PagedGraphParams key (batch / heads / head_size / // blocks / block_size / max_num_blocks_per_seq / scale / dtype / handle), so every call to // try_build_paged_graph -- warm-up, capturing, replaying -- hits the cached plan. past_seqlen // varies across Runs but is not part of the cache key. c.replay_past_seqlens = { - {256}, - {260}, - {300}, - {340}, - {380}, + {511}, + {510}, + {509}, + {508}, + {507}, }; testing::internal::CaptureStdout(); diff --git a/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc b/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc index ea9ff00c08f8c..5ebe805553bd7 100644 --- a/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc +++ b/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc @@ -2,6 +2,8 @@ // Licensed under the MIT License. #include "gtest/gtest.h" +#include +#include #include "core/session/onnxruntime_cxx_api.h" #include "test/common/tensor_op_test_utils.h" #include "test/common/cuda_op_test_utils.h" @@ -11,6 +13,137 @@ namespace onnxruntime { namespace test { constexpr float epsilon_ = 1e-12f; +#if defined(USE_CUDA) +TEST(SkipLayerNormTest, CudaEmptyHiddenDimension) { + if (!HasCudaEnvironment(0)) GTEST_SKIP() << "CUDA EP unavailable"; + + OpTester tester("SkipLayerNormalization", 1, kMSDomain); + tester.AddAttribute("epsilon", 1e-5f); + tester.AddInput("input", {1, 2, 0}, {}); + tester.AddInput("skip", {1, 2, 0}, {}); + tester.AddInput("gamma", {0}, {}); + tester.AddInput("beta", {0}, {}); + tester.AddOutput("output", {1, 2, 0}, {}); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +template +static void CheckCudaHighOffsetVariance() { + if (!HasCudaEnvironment(0)) GTEST_SKIP() << "CUDA EP unavailable"; + + constexpr float epsilon = 1e-5f; + const auto native_value = [](float value) { + if constexpr (std::is_same_v) { + return value; + } else { + return T(value).ToFloat(); + } + }; + // Cover scalar/vectorized dispatch, partial blocks, and multiple values per thread. + for (int hidden_size : {1, 3, 4, 5, 128, 129, 320, 768, 769, 8192, 8193}) { + SCOPED_TRACE(hidden_size); + for (bool has_bias : {false, true}) { + SCOPED_TRACE(has_bias); + const std::vector shape{1, 2, hidden_size}; + const std::vector parameter_shape{hidden_size}; + std::vector pattern{1000.01f, 1000.02f, 999.99f, 1000.0f}; + float skip_value = 256.0f; + if constexpr (std::is_same_v) { + pattern = {60000.0f, 60032.0f, 59968.0f, 60000.0f}; + skip_value = -59872.0f; + } else if constexpr (std::is_same_v) { + pattern = {1024.0f, 1032.0f, 1016.0f, 1024.0f}; + skip_value = -896.0f; + } + std::vector input(2 * hidden_size); + std::vector skip(2 * hidden_size, has_bias ? skip_value : 0.0f); + std::vector bias(hidden_size, 512.0f); + std::vector gamma(hidden_size); + std::vector beta(hidden_size); + std::vector sum(2 * hidden_size); + std::vector expected(2 * hidden_size); + for (int i = 0; i < hidden_size; ++i) { + input[i] = pattern[i % 4]; + input[hidden_size + i] = std::is_same_v ? -1000.0f : pattern[0]; + // Small biases round before the low-precision residual cancellation. + if constexpr (std::is_same_v) { + bias[i] = i % 2 == 0 ? 0.25f : 0.5f; + } else if constexpr (std::is_same_v) { + bias[i] = i % 2 == 0 ? 0.5f : 1.5f; + } + gamma[i] = i % 2 == 0 ? 0.5f : -1.25f; + beta[i] = i % 2 == 0 ? 0.25f : -0.75f; + } + for (int row = 0; row < 2; ++row) { + double mean = 0.0; + for (int i = 0; i < hidden_size; ++i) { + const int index = row * hidden_size + i; + float value = input[index]; + if (has_bias) value = native_value(value + bias[i]); + value = native_value(value + skip[index]); + sum[index] = value; + mean += static_cast(value); + } + mean /= hidden_size; + double variance = 0.0; + for (int i = 0; i < hidden_size; ++i) { + const double delta = static_cast(sum[row * hidden_size + i]) - mean; + variance += delta * delta; + } + variance /= hidden_size; + const double inv_std = 1.0 / std::sqrt(variance + static_cast(epsilon)); + for (int i = 0; i < hidden_size; ++i) { + const int index = row * hidden_size + i; + expected[index] = static_cast((static_cast(sum[index]) - mean) * + inv_std * gamma[i] + + beta[i]); + } + } + + OpTester tester("SkipLayerNormalization", 1, kMSDomain); + tester.AddAttribute("epsilon", epsilon); + tester.AddInput("input", shape, std::vector(input.begin(), input.end())); + tester.AddInput("skip", shape, std::vector(skip.begin(), skip.end())); + tester.AddInput("gamma", parameter_shape, std::vector(gamma.begin(), gamma.end())); + tester.AddInput("beta", parameter_shape, std::vector(beta.begin(), beta.end())); + if (has_bias) tester.AddInput("bias", parameter_shape, std::vector(bias.begin(), bias.end())); + tester.AddOutput("output", shape, std::vector(expected.begin(), expected.end())); + tester.AddOptionalOutputEdge(); + tester.AddOptionalOutputEdge(); + tester.AddOutput("sum", shape, std::vector(sum.begin(), sum.end())); + if constexpr (std::is_same_v) { + tester.SetOutputAbsErr("output", 1e-4f); + tester.SetOutputRelErr("output", 1e-4f); + } + tester.SetOutputAbsErr("sum", 0.0f); + tester.SetOutputRelErr("sum", 0.0f); + SessionOptions options; + ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + } +} + +TEST(SkipLayerNormTest, CudaHighOffsetVariance) { + CheckCudaHighOffsetVariance(); +} + +TEST(SkipLayerNormTest, CudaFloat16HighOffsetVariance) { + CheckCudaHighOffsetVariance(); +} + +TEST(SkipLayerNormTest, CudaBFloat16HighOffsetVariance) { + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA BF16 unavailable"; + CheckCudaHighOffsetVariance(); +} +#endif + static void RunOneTest( const std::vector& input_data, const std::vector& skip_data, diff --git a/onnxruntime/test/framework/allocation_planner_reuse_index_test.cc b/onnxruntime/test/framework/allocation_planner_reuse_index_test.cc new file mode 100644 index 0000000000000..6637b11e35ccb --- /dev/null +++ b/onnxruntime/test/framework/allocation_planner_reuse_index_test.cc @@ -0,0 +1,162 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "core/framework/allocation_planner_reuse_index.h" + +namespace onnxruntime::allocation_planner_internal { + +struct FreeBufferIndexTestPeer { + static size_t Capacity(const FreeBufferIndex& index) { return index.slots_.size(); } + static size_t Size(const FreeBufferIndex& index) { return index.size_; } + static size_t Hash(size_t key) { return FreeBufferIndex::Hash(key); } + static size_t NextCapacity(size_t capacity) { return FreeBufferIndex::NextCapacity(capacity); } +}; + +namespace { + +TEST(FreeBufferIndexTest, EmptyZeroAndMaxWords) { + FreeBufferIndex index; + const auto max = std::numeric_limits::max(); + EXPECT_EQ(index.Find(0), nullptr); + EXPECT_EQ(index.Find(max), nullptr); + EXPECT_FALSE(index.GetOrCreate(0).has_value()); + index.GetOrCreate(0) = max; + EXPECT_FALSE(index.GetOrCreate(max).has_value()); + index.GetOrCreate(max) = 0; + ASSERT_NE(index.Find(0), nullptr); + ASSERT_NE(index.Find(max), nullptr); + EXPECT_EQ(*index.Find(0), max); + EXPECT_EQ(*index.Find(max), 0u); + EXPECT_EQ(FreeBufferIndexTestPeer::Size(index), 2u); +} + +TEST(FreeBufferIndexTest, FullWordCollisionsKeepConsumedProbeChains) { + // 32 words in eight initial buckets guarantee at least four colliding words. + // Equal low halves also detect accidentally narrowed key comparisons. + constexpr int shift = std::numeric_limits::digits / 2; + std::array, 8> groups; + for (size_t i = 0; i < 32; ++i) { + size_t key = (i << shift) | 7; + groups[FreeBufferIndexTestPeer::Hash(key) & 7].push_back(key); + } + const auto group = std::find_if(groups.begin(), groups.end(), + [](const auto& keys) { return keys.size() >= 4; }); + ASSERT_NE(group, groups.end()); + FreeBufferIndex index; + for (size_t i = 0; i < 4; ++i) index.GetOrCreate((*group)[i]) = i; + EXPECT_EQ(FreeBufferIndexTestPeer::Capacity(index), 8u); + index.Find((*group)[0])->reset(); + index.Find((*group)[1])->reset(); + for (size_t i = 2; i < 4; ++i) { + ASSERT_NE(index.Find((*group)[i]), nullptr); + EXPECT_EQ(*index.Find((*group)[i]), i); + } + EXPECT_FALSE(index.GetOrCreate((*group)[0]).has_value()); + index.GetOrCreate((*group)[0]) = std::numeric_limits::max(); + EXPECT_EQ(*index.Find((*group)[0]), std::numeric_limits::max()); + EXPECT_EQ(FreeBufferIndexTestPeer::Size(index), 4u); +} + +TEST(FreeBufferIndexTest, DuplicateAndReaddedHeadsDoNotGrowAtLoadLimit) { + FreeBufferIndex index; + for (size_t i = 0; i < 6; ++i) index.GetOrCreate(i) = i + 10; + EXPECT_EQ(FreeBufferIndexTestPeer::Capacity(index), 8u); + EXPECT_EQ(index.GetOrCreate(2), 12u); + index.Find(2)->reset(); + EXPECT_FALSE(index.GetOrCreate(2).has_value()); + index.GetOrCreate(2) = 99; + EXPECT_EQ(FreeBufferIndexTestPeer::Capacity(index), 8u); + EXPECT_EQ(FreeBufferIndexTestPeer::Size(index), 6u); + index.Find(4)->reset(); + index.GetOrCreate(6) = 16; + EXPECT_EQ(FreeBufferIndexTestPeer::Capacity(index), 16u); + EXPECT_EQ(FreeBufferIndexTestPeer::Size(index), 7u); + EXPECT_EQ(*index.Find(2), 99u); + ASSERT_NE(index.Find(4), nullptr); + EXPECT_FALSE(index.Find(4)->has_value()); +} + +TEST(FreeBufferIndexTest, WrappingProbeChainRetainsConsumedEntries) { + std::vector keys; + for (size_t key = 0; key < 4096 && keys.size() < 4; ++key) { + if ((FreeBufferIndexTestPeer::Hash(key) & 7) == 7) keys.push_back(key); + } + ASSERT_EQ(keys.size(), 4u); + FreeBufferIndex index; + for (size_t i = 0; i < 3; ++i) index.GetOrCreate(keys[i]) = i; + index.Find(keys[0])->reset(); + EXPECT_EQ(*index.Find(keys[2]), 2u); + EXPECT_EQ(index.Find(keys[3]), nullptr); + index.GetOrCreate(keys[3]) = 3; + EXPECT_EQ(*index.Find(keys[3]), 3u); + EXPECT_FALSE(index.GetOrCreate(keys[0]).has_value()); + EXPECT_EQ(FreeBufferIndexTestPeer::Capacity(index), 8u); +} + +TEST(FreeBufferIndexTest, GrowthPreservesEveryWordAndConsumedBucket) { + FreeBufferIndex index; + std::map> expected; + const auto max = std::numeric_limits::max(); + for (size_t i = 0; i < 2048; ++i) { + const size_t key = i == 0 ? max : i - 1; + auto& head = index.GetOrCreate(key); + if (i % 5 == 0) { + head.reset(); + expected[key] = std::nullopt; + } else { + const size_t value = i == 1 ? max : i == 2 ? 0 + : max - i; + head = value; + expected[key] = value; + } + const auto capacity = FreeBufferIndexTestPeer::Capacity(index); + EXPECT_EQ(capacity & (capacity - 1), 0u); + EXPECT_LE(FreeBufferIndexTestPeer::Size(index), capacity - capacity / 4); + } + for (const auto& [key, head] : expected) { + ASSERT_NE(index.Find(key), nullptr); + EXPECT_EQ(*index.Find(key), head); + } + EXPECT_EQ(FreeBufferIndexTestPeer::Size(index), expected.size()); + EXPECT_EQ(index.Find(4096), nullptr); +} + +TEST(FreeBufferIndexTest, ClearRemovesAllKeysAndStaleHeads) { + FreeBufferIndex index; + for (size_t i = 0; i < 30; ++i) index.GetOrCreate(i) = i; + index.Find(3)->reset(); + const auto capacity = FreeBufferIndexTestPeer::Capacity(index); + index.clear(); + EXPECT_EQ(FreeBufferIndexTestPeer::Size(index), 0u); + EXPECT_EQ(FreeBufferIndexTestPeer::Capacity(index), capacity); + for (size_t i = 0; i < 30; ++i) { + EXPECT_EQ(index.Find(i), nullptr); + EXPECT_FALSE(index.GetOrCreate(i).has_value()); + index.GetOrCreate(i) = i + 50; + } + for (size_t i = 0; i < 30; ++i) EXPECT_EQ(*index.Find(i), i + 50); +} + +#if GTEST_HAS_EXCEPTIONS && !defined(ORT_NO_EXCEPTIONS) +TEST(FreeBufferIndexTest, CheckedCapacityGrowth) { + EXPECT_EQ(FreeBufferIndexTestPeer::NextCapacity(0), 8u); + EXPECT_EQ(FreeBufferIndexTestPeer::NextCapacity(8), 16u); + EXPECT_THROW(FreeBufferIndexTestPeer::NextCapacity(std::numeric_limits::max()), + OnnxRuntimeException); +} +#endif + +} // namespace +} // namespace onnxruntime::allocation_planner_internal + +#endif // !defined(ORT_MINIMAL_BUILD) diff --git a/onnxruntime/test/framework/allocation_planner_test.cc b/onnxruntime/test/framework/allocation_planner_test.cc index daff2cd55a59e..dec0ec5e748c0 100644 --- a/onnxruntime/test/framework/allocation_planner_test.cc +++ b/onnxruntime/test/framework/allocation_planner_test.cc @@ -115,12 +115,16 @@ class SequentialPlannerTestContext : public ISequentialPlannerContext { SequentialPlannerTestContext(ShapeMap* shape_map) : shape_map_(shape_map) {} TensorShapeProto* GetShape(const onnxruntime::NodeArg& arg) const override { + ++shape_queries_; auto iter = shape_map_->find(&arg); return (shape_map_->end() != iter) ? iter->second : nullptr; } + size_t GetShapeQueryCount() const { return shape_queries_; } + private: ShapeMap* shape_map_; + mutable size_t shape_queries_ = 0; }; class ParallelPlannerTestContext : public SequentialPlannerTestContext { @@ -165,6 +169,7 @@ class PlannerTest : public ::testing::Test { std::unique_ptr state_; ShapeMap shape_map_; std::optional plan_; + size_t shape_queries_ = 0; public: PlannerTest() @@ -208,6 +213,11 @@ class PlannerTest : public ::testing::Test { return (name_to_arg_[name] = &graph_.GetOrCreateNodeArg(name, &float_type_.value)); } + template + void SetTensorType(const std::string& name) { + name_to_arg_[name] = &graph_.GetOrCreateNodeArg(name, DataTypeImpl::GetTensorType()->GetTypeProto()); + } + onnxruntime::Node* AddNode(::onnxruntime::KernelDef& kernel_def, std::string& input, std::string& output) { auto node = std::make_unique(graph_, kernel_def.OpName(), Arg(input), Arg(output)); auto* p_node = node->p_node; @@ -349,6 +359,7 @@ class PlannerTest : public ::testing::Test { plan_); EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); + shape_queries_ = test_context.GetShapeQueryCount(); // AllocationPlanTestUtility::BasicIntegrityCheck(*plan_, name_to_arg_.size()); } } @@ -359,6 +370,13 @@ class PlannerTest : public ::testing::Test { EXPECT_EQ(plan_->allocation_plan[id].alloc_kind, kind) << "Error in allocation kind for " << name; } + void CheckReusedBuffer(const std::string& name, const std::string& expected) { + int id, expected_id; + index(name, id); + index(expected, expected_id); + EXPECT_EQ(plan_->allocation_plan[id].reused_buffer, expected_id); + } + void CheckFreed(int step_number, std::initializer_list freed_items) { // TODO: add the checker for new implementation of release plan //// create set and check equality @@ -382,6 +400,26 @@ class PlannerTest : public ::testing::Test { protected: Graph& GetGraph() { return graph_; } + size_t GetShapeQueryCount() const { return shape_queries_; } + template + void CheckStorageSizedReuse(AllocKind expected) { + std::string input = "input", a = "a", b = "b", c = "c", output = "output"; + SetTensorType(b); + SetTensorType(c); + SetTensorType(output); + AddNormalNode(input, a); + auto cast_kernel = KernelDefBuilder().SetName("Cast").Provider(kCpuExecutionProvider).SinceVersion(9).Build(); + AddNode(*cast_kernel, a, b)->AddAttribute("to", static_cast(utils::ToTensorProtoElementType())); + AddNormalNode(b, c); + AddNormalNode(c, output); + Shape compatible{"batch", "width"}, different{"batch", "other"}; + SetShape({{a, &compatible.value}, {b, &different.value}, {c, &compatible.value}}); + + CreatePlan(); + + CheckAllocKind(c, expected); + if (expected == AllocKind::kReuse) CheckReusedBuffer(c, a); + } const SequentialExecutionPlan& GetPlan() const { return *plan_; } const SessionState& GetState() const { return *state_; } ExecutionProviders& GetExecutionProviders() { return execution_providers_; } @@ -716,6 +754,123 @@ TEST_F(PlannerTest, InPlaceSizeMismatchTest) { CheckFreed(3, {X2}); } +TEST_F(PlannerTest, ReuseIndexPreservesMostRecentlyFreedCompatibleBuffer) { + std::string input = "input", a = "a", b = "b", c = "c"; + std::string joined = "joined", reused = "reused", output = "output"; + AddNormalNode(input, a); + AddNormalNode(input, b); + AddNormalNode(input, c); + auto sum_kernel = KernelDefBuilder().SetName("Sum").Provider(kCpuExecutionProvider).SinceVersion(8).Build(); + std::string sum_name = "sum"; + std::vector inputs{Arg(a), Arg(b), Arg(c)}, outputs{Arg(joined)}; + AddNode(*sum_kernel, sum_name, inputs, outputs); + AddNormalNode(joined, reused); + AddNormalNode(reused, output); + Shape compatible{"batch", "width"}, incompatible{"batch", "other"}, join_shape{"batch", "joined"}; + SetShape({{a, &compatible.value}, {b, &compatible.value}, {c, &incompatible.value}, {joined, &join_shape.value}, {reused, &compatible.value}, {output, &compatible.value}}); + + CreatePlan(); + + // Sum frees a, b, then c. Skip the newer incompatible c, but prefer b to a. + CheckAllocKind(reused, AllocKind::kReuse); + CheckReusedBuffer(reused, b); + CheckAllocKind(output, AllocKind::kAllocateOutput); +} + +TEST_F(PlannerTest, ReuseIndexPreservesEmptyAndUnknownShapeEligibility) { + std::string input = "input", a = "a", b = "b", c = "c", output = "output"; + AddNormalNode(input, a); + AddNormalNode(a, b); + AddNormalNode(b, c); + AddNormalNode(c, output); + Shape empty_symbol{""}, different{"other"}; + SetShape({{a, &empty_symbol.value}, {b, &different.value}, {c, &empty_symbol.value}}); + + CreatePlan(); + + CheckAllocKind(a, AllocKind::kAllocate); + CheckAllocKind(b, AllocKind::kAllocate); + CheckAllocKind(c, AllocKind::kAllocate); + CheckAllocKind(output, AllocKind::kAllocateOutput); +} + +TEST_F(PlannerTest, ReuseIndexAllowsZeroSizedTensors) { + std::string input = "input", a = "a", b = "b", c = "c", output = "output"; + AddNormalNode(input, a); + AddNormalNode(a, b); + AddNormalNode(b, c); + AddNormalNode(c, output); + Shape empty{0, 8}, different{1, 8}; + SetShape({{a, &empty.value}, {b, &different.value}, {c, &empty.value}}); + + CreatePlan(); + + CheckAllocKind(c, AllocKind::kReuse); + CheckReusedBuffer(c, a); +} + +TEST_F(PlannerTest, ReuseIndexDoesNotReuseScalarBuffers) { + std::string input = "input", a = "a", b = "b", c = "c", output = "output"; + AddNormalNode(input, a); + AddNormalNode(a, b); + AddNormalNode(b, c); + AddNormalNode(c, output); + Shape scalar(std::initializer_list{}), different{1}; + SetShape({{a, &scalar.value}, {b, &different.value}, {c, &scalar.value}}); + + CreatePlan(); + + CheckAllocKind(c, AllocKind::kAllocate); +} + +TEST_F(PlannerTest, ReuseIndexAllowsEqualWidthDifferentElementTypes) { + CheckStorageSizedReuse(AllocKind::kReuse); +} + +TEST_F(PlannerTest, ReuseIndexRejectsDifferentElementWidths) { + CheckStorageSizedReuse(AllocKind::kAllocate); +} + +TEST_F(PlannerTest, ReuseIndexDoesNotReuseStringBuffers) { + std::string input = "input", a = "a", b = "b", c = "c", output = "output"; + for (const auto& name : {input, a, b, c, output}) SetTensorType(name); + AddNormalNode(input, a); + AddNormalNode(a, b); + AddNormalNode(b, c); + AddNormalNode(c, output); + Shape shape{"batch", "width"}; + SetShape({{a, &shape.value}, {b, &shape.value}, {c, &shape.value}}); + + CreatePlan(); + + CheckAllocKind(c, AllocKind::kAllocate); +} + +// The indexed complexity contract applies to full builds, not the minimal linear fallback. +#if !defined(ORT_MINIMAL_BUILD) +TEST_F(PlannerTest, ReuseIndexAvoidsQuadraticShapeQueries) { + constexpr int node_count = 512; + InlinedVector names; + InlinedVector shapes; + names.reserve(node_count + 1); + shapes.reserve(node_count + 1); + for (int i = 0; i <= node_count; ++i) { + names.push_back("value_" + std::to_string(i)); + shapes.emplace_back(std::initializer_list{i + 1}); + SetShape(names.back(), &shapes.back().value); + if (i > 0) AddNormalNode(names[i - 1], names[i]); + } + + CreatePlan(); + + // Count existing shape-context queries instead of relying on wall-clock time. + // Every freed buffer has a distinct shape, so no full-list scans are needed. + RecordProperty("shape_queries", std::to_string(GetShapeQueryCount())); + EXPECT_LE(GetShapeQueryCount(), 3U * node_count); + for (int i = 1; i < node_count; ++i) CheckAllocKind(names[i], AllocKind::kAllocate); +} +#endif + // Test operator<< to output details of an allocation & execution plan. TEST_F(PlannerTest, PlanOutputTest) { // tensor variables: @@ -2211,29 +2366,29 @@ TEST(AllocationPlannerTest, AvoidReuseOfPackedSubByteBufferForFullByteTensor) { return t; }; - auto create_model = [&]() -> Model { + auto create_model = [&](TensorProto_DataType source_elem_type, TensorProto_DataType dest_elem_type) -> Model { Model model("packed_subbyte_reuse", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}}, {}, DefaultLoggingManager().DefaultLogger()); Graph& graph = model.MainGraph(); TypeProto float_type = make_tensor_type(TensorProto_DataType_FLOAT); - TypeProto uint4_type = make_tensor_type(TensorProto_DataType_UINT4); - TypeProto uint8_type = make_tensor_type(TensorProto_DataType_UINT8); + TypeProto source_type = make_tensor_type(source_elem_type); + TypeProto dest_type = make_tensor_type(dest_elem_type); auto& X = graph.GetOrCreateNodeArg("X", &float_type); - auto& A = graph.GetOrCreateNodeArg("A", &uint4_type); // packed sub-byte buffer: ceil(1023/2) = 512 bytes + auto& A = graph.GetOrCreateNodeArg("A", &source_type); auto& B = graph.GetOrCreateNodeArg("B", &float_type); // consumes A -> A becomes dead here - auto& C = graph.GetOrCreateNodeArg("C", &uint8_type); // full-byte buffer: 1023 bytes + auto& C = graph.GetOrCreateNodeArg("C", &dest_type); auto& Y = graph.GetOrCreateNodeArg("Y", &float_type); auto add_cast = [&graph](const std::string& name, NodeArg& in, NodeArg& out, TensorProto_DataType to) { auto& node = graph.AddNode(name, "Cast", name, {&in}, {&out}); node.AddAttribute("to", static_cast(to)); }; - add_cast("cast_to_uint4", X, A, TensorProto_DataType_UINT4); + add_cast("cast_to_source", X, A, source_elem_type); add_cast("cast_a_to_float", A, B, TensorProto_DataType_FLOAT); - add_cast("cast_to_uint8", B, C, TensorProto_DataType_UINT8); + add_cast("cast_to_dest", B, C, dest_elem_type); add_cast("cast_c_to_float", C, Y, TensorProto_DataType_FLOAT); graph.SetInputs({&X}); @@ -2242,48 +2397,64 @@ TEST(AllocationPlannerTest, AvoidReuseOfPackedSubByteBufferForFullByteTensor) { return model; }; - SessionOptions so; - // Keep memory reuse on (default) and avoid graph optimizations that could rewrite the Cast chain. - so.graph_optimization_level = TransformerLevel::Default; - InferenceSession sess{so, GetEnvironment()}; - - std::string serialized; - ASSERT_TRUE(create_model().ToProto().SerializeToString(&serialized)); - std::stringstream sstr(serialized); - ASSERT_STATUS_OK(sess.Load(sstr)); - ASSERT_STATUS_OK(sess.Initialize()); - - const auto& session_state = sess.GetSessionState(); - const auto& ort_value_index_map = session_state.GetOrtValueNameIdxMap(); - const SequentialExecutionPlan* plan = session_state.GetExecutionPlan(); - - OrtValueIndex a_index, c_index; - ASSERT_STATUS_OK(ort_value_index_map.GetIdx("A", a_index)); - ASSERT_STATUS_OK(ort_value_index_map.GetIdx("C", c_index)); - - // The uint8 tensor C must never be planned to reuse the (smaller) uint4 tensor A's buffer. - const auto& c_plan = plan->allocation_plan[c_index]; - const bool reuses_a = (c_plan.alloc_kind == AllocKind::kReuse) && (c_plan.reused_buffer == a_index); - EXPECT_FALSE(reuses_a) << "uint8[" << kDim << "] output must not reuse the uint4[" << kDim << "] buffer"; - - // The model must still run and produce correct results after round-tripping through uint4 and uint8. - constexpr float kValue = 3.0f; - std::vector dims{kDim}; - std::vector input_data(static_cast(kDim), kValue); - OrtValue input_value; - CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], dims, input_data, &input_value); - - NameMLValMap feeds{{"X", input_value}}; - std::vector output_names{"Y"}; - std::vector fetches; - ASSERT_STATUS_OK(sess.Run(feeds, output_names, &fetches)); - - ASSERT_EQ(fetches.size(), 1u); - const Tensor& out = fetches[0].Get(); - ASSERT_EQ(out.Shape().Size(), kDim); - const float* out_data = out.Data(); - for (int64_t i = 0; i < kDim; ++i) { - ASSERT_EQ(out_data[i], kValue) << "mismatch at index " << i; + struct StorageCase { + TensorProto_DataType source; + TensorProto_DataType destination; + bool reuse; + }; + // Exercise both packing-density rejection and compatible packed/dense reuse, plus string exclusion. + for (const auto& storage_case : { + StorageCase{TensorProto_DataType_UINT4, TensorProto_DataType_UINT8, false}, + StorageCase{TensorProto_DataType_UINT8, TensorProto_DataType_UINT4, false}, + StorageCase{TensorProto_DataType_UINT4, TensorProto_DataType_UINT4, true}, + StorageCase{TensorProto_DataType_UINT4, TensorProto_DataType_INT4, true}, + StorageCase{TensorProto_DataType_UINT8, TensorProto_DataType_INT8, true}, + StorageCase{TensorProto_DataType_STRING, TensorProto_DataType_STRING, false}, + }) { + SCOPED_TRACE(::testing::Message() << "source=" << storage_case.source + << ", destination=" << storage_case.destination); + SessionOptions so; + // Keep memory reuse on (default) and avoid graph optimizations that could rewrite the Cast chain. + so.graph_optimization_level = TransformerLevel::Default; + InferenceSession sess{so, GetEnvironment()}; + + std::string serialized; + ASSERT_TRUE(create_model(storage_case.source, storage_case.destination).ToProto().SerializeToString(&serialized)); + std::stringstream sstr(serialized); + ASSERT_STATUS_OK(sess.Load(sstr)); + ASSERT_STATUS_OK(sess.Initialize()); + + const auto& session_state = sess.GetSessionState(); + const auto& ort_value_index_map = session_state.GetOrtValueNameIdxMap(); + const SequentialExecutionPlan* plan = session_state.GetExecutionPlan(); + + OrtValueIndex a_index, c_index; + ASSERT_STATUS_OK(ort_value_index_map.GetIdx("A", a_index)); + ASSERT_STATUS_OK(ort_value_index_map.GetIdx("C", c_index)); + + const auto& c_plan = plan->allocation_plan[c_index]; + const bool reuses_a = (c_plan.alloc_kind == AllocKind::kReuse) && (c_plan.reused_buffer == a_index); + EXPECT_EQ(reuses_a, storage_case.reuse); + + // The model must still run and produce correct results after round-tripping through both storage types. + constexpr float kValue = 3.0f; + std::vector dims{kDim}; + std::vector input_data(static_cast(kDim), kValue); + OrtValue input_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], dims, input_data, &input_value); + + NameMLValMap feeds{{"X", input_value}}; + std::vector output_names{"Y"}; + std::vector fetches; + ASSERT_STATUS_OK(sess.Run(feeds, output_names, &fetches)); + + ASSERT_EQ(fetches.size(), 1u); + const Tensor& out = fetches[0].Get(); + ASSERT_EQ(out.Shape().Size(), kDim); + const float* out_data = out.Data(); + for (int64_t i = 0; i < kDim; ++i) { + ASSERT_EQ(out_data[i], kValue) << "mismatch at index " << i; + } } } diff --git a/onnxruntime/test/framework/ort_model_only_test.cc b/onnxruntime/test/framework/ort_model_only_test.cc index 8b6b262b4cc04..2192ea56db714 100644 --- a/onnxruntime/test/framework/ort_model_only_test.cc +++ b/onnxruntime/test/framework/ort_model_only_test.cc @@ -88,6 +88,40 @@ Status LoadOrtBuffer(const std::vector& buffer, bool use_buffer_for_ini return session_object.Load(buffer.data(), static_cast(buffer.size())); } +std::vector BuildOrtModelWithEdge(int32_t src_arg_index, int32_t dst_arg_index, + bool input_edge = true, bool mismatched_args = false) { + return BuildOrtModelBuffer([&](flatbuffers::FlatBufferBuilder& builder) { + std::vector> node_args{ + fbs::CreateValueInfoDirect(builder, "input", "", CreateFloatTensorTypeInfo(builder, 1)), + fbs::CreateValueInfoDirect(builder, "x", "", CreateFloatTensorTypeInfo(builder, 1)), + fbs::CreateValueInfoDirect(builder, "y", "", CreateFloatTensorTypeInfo(builder, 1)), + fbs::CreateValueInfoDirect(builder, "other", "", CreateFloatTensorTypeInfo(builder, 1))}; + std::vector> empty_args; + std::vector> source_inputs{builder.CreateSharedString("input")}; + std::vector> source_outputs{builder.CreateSharedString("x")}; + std::vector> destination_inputs{ + builder.CreateSharedString(mismatched_args ? "other" : "x")}; + std::vector> destination_outputs{builder.CreateSharedString("y")}; + std::vector arg_counts{1}; + std::vector> nodes{ + fbs::CreateNodeDirect(builder, "source", "", "", 1, 0, "Identity", + fbs::NodeType::Primitive, nullptr, &source_inputs, &source_outputs, + nullptr, &arg_counts, &empty_args), + fbs::CreateNodeDirect(builder, "destination", "", "", 1, 1, "Identity", + fbs::NodeType::Primitive, nullptr, &destination_inputs, &destination_outputs, + nullptr, &arg_counts, &empty_args)}; + std::vector input_edges{fbs::EdgeEnd(0, src_arg_index, dst_arg_index)}; + std::vector output_edges{fbs::EdgeEnd(1, src_arg_index, dst_arg_index)}; + std::vector> node_edges{ + input_edge ? fbs::CreateNodeEdgeDirect(builder, 1, &input_edges) + : fbs::CreateNodeEdgeDirect(builder, 0, nullptr, &output_edges)}; + std::vector> graph_inputs{builder.CreateSharedString("input")}; + std::vector> graph_outputs{builder.CreateSharedString("y")}; + return fbs::CreateGraphDirect( + builder, nullptr, &node_args, &nodes, 2, &node_edges, &graph_inputs, &graph_outputs); + }); +} + } // namespace static void RunOrtModel(const OrtModelTestInfo& test_info) { @@ -222,6 +256,24 @@ TEST(OrtModelTest, RejectsInitializerRawDataSizeMismatch) { EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("raw data size mismatch")); } +TEST(OrtModelTest, RejectsInvalidSerializedEdges) { + for (const bool input_edge : {false, true}) { + for (const auto& [src_arg_index, dst_arg_index, error] : { + std::tuple{1, 0, "out-of-range src_arg_index"}, + std::tuple{0, 1, "out-of-range dst_arg_index"}, + std::tuple{INT_MAX, INT_MAX, "control edges are not supported"}}) { + const auto status = LoadOrtBuffer( + BuildOrtModelWithEdge(src_arg_index, dst_arg_index, input_edge)); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr(error)); + } + + const auto status = LoadOrtBuffer(BuildOrtModelWithEdge(0, 0, input_edge, true)); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("mismatched NodeArgs")); + } +} + TEST(OrtModelTest, RejectsDanglingNodeEdge) { const auto buffer = BuildOrtModelBuffer([](flatbuffers::FlatBufferBuilder& builder) { std::vector> node_edges{ diff --git a/onnxruntime/test/framework/resource_accountant_test.cc b/onnxruntime/test/framework/resource_accountant_test.cc index 7b4ebec73571d..b72b4c2bf98f2 100644 --- a/onnxruntime/test/framework/resource_accountant_test.cc +++ b/onnxruntime/test/framework/resource_accountant_test.cc @@ -3,6 +3,7 @@ #include "core/framework/resource_accountant.h" #include "core/framework/config_options.h" +#include "core/framework/kernel_dispatch_policy.h" #include "core/graph/indexed_sub_graph.h" #include "core/graph/constants.h" #include "core/graph/model.h" @@ -815,6 +816,39 @@ TEST(RealAccountantTest, Factory_NoLimitNoStats) { EXPECT_FALSE(accountant->GetThreshold().has_value()); } +TEST(KernelDispatchPolicyTest, PartialEstimatorConfigDefaultsPolicyToUnset) { + const WorkspaceEstimatorConfig config{std::nullopt, std::nullopt}; + EXPECT_FALSE(config.kernel_dispatch_policy.has_value()); + EXPECT_EQ(config.cuda_gqa_workspace_max_total_sequence_length, int64_t{0}); +} + +TEST(KernelDispatchPolicyTest, ParsesValidTokens) { + struct Case { + const char* value; + KernelDispatchPolicy expected; + }; + for (const auto& test_case : { + Case{"", KernelDispatchPolicy::Auto}, + Case{"auto", KernelDispatchPolicy::Auto}, + Case{"latency", KernelDispatchPolicy::Latency}, + Case{"memory", KernelDispatchPolicy::Memory}, + Case{"safe", KernelDispatchPolicy::Safe}}) { + SCOPED_TRACE(test_case.value); + KernelDispatchPolicy parsed = KernelDispatchPolicy::Safe; + EXPECT_TRUE(TryParseKernelDispatchPolicy(test_case.value, parsed)); + EXPECT_EQ(parsed, test_case.expected); + } +} + +TEST(KernelDispatchPolicyTest, RejectsInvalidTokens) { + for (const char* value : {"Latency", "Memory", "SAFE", "memroy", " memory", "memory ", " memory ", " ", "unknown"}) { + SCOPED_TRACE(value); + KernelDispatchPolicy parsed = KernelDispatchPolicy::Memory; + EXPECT_FALSE(TryParseKernelDispatchPolicy(value, parsed)); + EXPECT_EQ(parsed, KernelDispatchPolicy::Memory); + } +} + TEST(RealAccountantTest, FactoryRetainsNarrowWorkspaceEstimatorConfig) { ConfigOptions config; ASSERT_STATUS_OK(config.AddConfigEntry( @@ -825,6 +859,8 @@ TEST(RealAccountantTest, FactoryRetainsNarrowWorkspaceEstimatorConfig) { kOrtSessionOptionsCudaFpAIntBProfileM, "1,16")); ASSERT_STATUS_OK(config.AddConfigEntry( kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength, "4096")); + ASSERT_STATUS_OK(config.AddConfigEntry( + kOrtSessionOptionsKernelDispatchPolicy, "memory")); ASSERT_STATUS_OK(config.AddConfigEntry("unrelated.config", "not copied")); std::optional acc_map; @@ -837,6 +873,72 @@ TEST(RealAccountantTest, FactoryRetainsNarrowWorkspaceEstimatorConfig) { EXPECT_EQ(estimator_config.cuda_fpa_intb_profile_m, std::optional{"1,16"}); EXPECT_EQ(estimator_config.cuda_gqa_workspace_max_total_sequence_length, int64_t{4096}); + EXPECT_EQ(estimator_config.kernel_dispatch_policy, + std::optional{"memory"}); +} + +TEST(RealAccountantTest, FactoryRetainsKernelDispatchPolicyReader) { + struct Case { + const char* value; + }; + for (const auto& test_case : { + Case{nullptr}, Case{""}, Case{"auto"}, Case{"latency"}, Case{"memory"}, + Case{"safe"}}) { + SCOPED_TRACE(test_case.value != nullptr ? test_case.value : ""); + ConfigOptions config; + ASSERT_STATUS_OK(config.AddConfigEntry( + kOrtSessionOptionsResourceCudaPartitioningSettings, "1000,")); + if (test_case.value != nullptr) { + ASSERT_STATUS_OK(config.AddConfigEntry( + kOrtSessionOptionsKernelDispatchPolicy, test_case.value)); + } + std::optional acc_map; + ASSERT_STATUS_OK(CreateAccountants(config, PathString(), acc_map)); + ASSERT_TRUE(acc_map.has_value()); + const auto& estimator_config = + acc_map->at(kCudaExecutionProvider)->GetWorkspaceEstimatorConfig(); + if (test_case.value == nullptr) { + EXPECT_FALSE(estimator_config.kernel_dispatch_policy.has_value()); + } else { + EXPECT_EQ(estimator_config.kernel_dispatch_policy, + std::optional{test_case.value}); + } + EXPECT_EQ(estimator_config.cuda_gqa_workspace_max_total_sequence_length, int64_t{0}); + } +} + +TEST(RealAccountantTest, KernelDispatchPolicyDoesNotEnablePartitioning) { + for (const char* value : {"", "auto", "latency", "memory", "safe"}) { + SCOPED_TRACE(value); + ConfigOptions config; + ASSERT_STATUS_OK(config.AddConfigEntry(kOrtSessionOptionsKernelDispatchPolicy, value)); + std::optional acc_map; + ASSERT_STATUS_OK(CreateAccountants(config, PathString(), acc_map)); + EXPECT_FALSE(acc_map.has_value()); + } +} + +TEST(RealAccountantTest, FactoryRejectsInvalidKernelDispatchPolicy) { + for (bool enable_partitioning : {false, true}) { + SCOPED_TRACE(enable_partitioning); + for (const char* value : {"Latency", "Memory", "SAFE", "memroy", " memory", "memory ", " memory ", " ", "unknown"}) { + SCOPED_TRACE(value); + ConfigOptions config; + if (enable_partitioning) { + ASSERT_STATUS_OK(config.AddConfigEntry( + kOrtSessionOptionsResourceCudaPartitioningSettings, "1000,")); + } + ASSERT_STATUS_OK(config.AddConfigEntry( + kOrtSessionOptionsKernelDispatchPolicy, value)); + std::optional acc_map; + const auto status = CreateAccountants(config, PathString(), acc_map); + ASSERT_FALSE(status.IsOK()); + EXPECT_EQ(status.Code(), common::StatusCode::INVALID_ARGUMENT); + EXPECT_NE(status.ErrorMessage().find(kOrtSessionOptionsKernelDispatchPolicy), + std::string::npos); + EXPECT_FALSE(acc_map.has_value()); + } + } } TEST(RealAccountantTest, FactoryParsesGqaWorkspaceEnvelope) { diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 3faced9bacb23..b241caf60289d 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -5985,6 +5985,163 @@ TEST_F(GraphTransformationTests, GemmTransposeFusionPreservesIdentityOutputWhenF } } +TEST_F(GraphTransformationTests, GemmTransposeFusionPreservesPublicInputTransposeWhenFusingOtherInput) { + for (bool public_input_a : {false, true}) { + for (bool initial_trans_a : {false, true}) { + for (bool initial_trans_b : {false, true}) { + SCOPED_TRACE(testing::Message() << "public_input_a=" << public_input_a << ", transA=" << initial_trans_a + << ", transB=" << initial_trans_b); + + std::string public_transpose_output_name; + auto build_test_case = [public_input_a, initial_trans_a, initial_trans_b, + &public_transpose_output_name](ModelTestBuilder& builder) { + const std::vector input_a_shape = initial_trans_a ? std::vector{3, 4} + : std::vector{4, 3}; + const std::vector input_b_shape = initial_trans_b ? std::vector{4, 5} + : std::vector{5, 4}; + const std::vector transpose_a_shape = initial_trans_a ? std::vector{4, 3} + : std::vector{3, 4}; + const std::vector transpose_b_shape = initial_trans_b ? std::vector{5, 4} + : std::vector{4, 5}; + auto* input_a = builder.MakeInput(input_a_shape, "A"); + auto* input_b = builder.MakeInput(input_b_shape, "B"); + auto* transpose_a_output = public_input_a ? builder.MakeOutput(transpose_a_shape) + : builder.MakeIntermediate(transpose_a_shape); + auto* transpose_b_output = public_input_a ? builder.MakeIntermediate(transpose_b_shape) + : builder.MakeOutput(transpose_b_shape); + + builder.AddNode("Transpose", {input_a}, {transpose_a_output}) + .AddAttribute("perm", std::vector{1, 0}); + builder.AddNode("Transpose", {input_b}, {transpose_b_output}) + .AddAttribute("perm", std::vector{1, 0}); + public_transpose_output_name = (public_input_a ? transpose_a_output : transpose_b_output)->Name(); + + auto* output = builder.MakeOutput(std::vector{3, 5}); + auto& gemm = builder.AddNode("Gemm", {transpose_a_output, transpose_b_output}, {output}); + gemm.AddAttribute("transA", static_cast(initial_trans_a)); + gemm.AddAttribute("transB", static_cast(initial_trans_b)); + gemm.AddAttribute("alpha", 1.0f); + gemm.AddAttribute("beta", 1.0f); + }; + + auto check_graph = [public_input_a, initial_trans_a, initial_trans_b, + &public_transpose_output_name](Graph& graph) { + auto op_counts = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_counts["Transpose"] == 1); + TEST_RETURN_IF_NOT(op_counts["Gemm"] == 1); + + const Node* gemm_node = nullptr; + const Node* public_transpose = nullptr; + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "Gemm") { + gemm_node = &node; + } else if (node.OpType() == "Transpose") { + public_transpose = &node; + TEST_RETURN_IF_NOT(graph.NodeProducesGraphOutput(node)); + } + } + TEST_RETURN_IF_NOT(gemm_node != nullptr); + TEST_RETURN_IF_NOT(public_transpose != nullptr); + + const size_t public_input_index = public_input_a ? 0 : 1; + const size_t fused_input_index = 1 - public_input_index; + const std::string expected_public_input_name = public_input_a ? "A" : "B"; + const std::string expected_fused_input_name = public_input_a ? "B" : "A"; + const bool expected_trans_a = public_input_a ? initial_trans_a : !initial_trans_a; + const bool expected_trans_b = public_input_a ? !initial_trans_b : initial_trans_b; + TEST_RETURN_IF_NOT(gemm_node->InputDefs().size() == 2); + TEST_RETURN_IF_NOT(public_transpose->InputDefs()[0]->Name() == expected_public_input_name); + TEST_RETURN_IF_NOT(public_transpose->OutputDefs()[0]->Name() == public_transpose_output_name); + TEST_RETURN_IF_NOT(gemm_node->InputDefs()[public_input_index]->Name() == public_transpose_output_name); + TEST_RETURN_IF_NOT(gemm_node->InputDefs()[fused_input_index]->Name() == expected_fused_input_name); + TEST_RETURN_IF_NOT(gemm_node->GetAttributes().at("transA").i() == + static_cast(expected_trans_a)); + TEST_RETURN_IF_NOT(gemm_node->GetAttributes().at("transB").i() == + static_cast(expected_trans_b)); + return Status::OK(); + }; + + auto rule_transformer = std::make_unique("RuleTransformer"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::move(rule_transformer), + TransformerLevel::Level1, 1, nullptr, check_graph)); + } + } + } +} + +TEST_F(GraphTransformationTests, GemmTransposeFusionPreservesPublicInputTransposesWhenFusingOutput) { + std::string transpose_a_output_name; + std::string transpose_b_output_name; + std::string final_output_name; + auto build_test_case = [&transpose_a_output_name, &transpose_b_output_name, + &final_output_name](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput(std::vector{4, 3}, "A"); + auto* input_b = builder.MakeInput(std::vector{5, 4}, "B"); + auto* transpose_a_output = builder.MakeOutput(std::vector{3, 4}); + auto* transpose_b_output = builder.MakeOutput(std::vector{4, 5}); + transpose_a_output_name = transpose_a_output->Name(); + transpose_b_output_name = transpose_b_output->Name(); + + builder.AddNode("Transpose", {input_a}, {transpose_a_output}).AddAttribute("perm", std::vector{1, 0}); + builder.AddNode("Transpose", {input_b}, {transpose_b_output}).AddAttribute("perm", std::vector{1, 0}); + + auto* gemm_output = builder.MakeIntermediate(std::vector{3, 5}); + auto& gemm = builder.AddNode("Gemm", {transpose_a_output, transpose_b_output}, {gemm_output}); + gemm.AddAttribute("transA", int64_t{0}); + gemm.AddAttribute("transB", int64_t{0}); + gemm.AddAttribute("alpha", 1.0f); + gemm.AddAttribute("beta", 1.0f); + auto* output = builder.MakeOutput(std::vector{5, 3}); + final_output_name = output->Name(); + builder.AddNode("Transpose", {gemm_output}, {output}).AddAttribute("perm", std::vector{1, 0}); + }; + + auto check_graph = [&transpose_a_output_name, &transpose_b_output_name, &final_output_name](Graph& graph) { + auto op_counts = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_counts["Transpose"] == 2); + TEST_RETURN_IF_NOT(op_counts["Gemm"] == 1); + + const Node* gemm_node = nullptr; + const Node* transpose_a = nullptr; + const Node* transpose_b = nullptr; + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "Gemm") { + gemm_node = &node; + } else if (node.OpType() == "Transpose") { + TEST_RETURN_IF_NOT(graph.NodeProducesGraphOutput(node)); + TEST_RETURN_IF_NOT(node.InputDefs().size() == 1); + if (node.InputDefs()[0]->Name() == "A") { + transpose_a = &node; + TEST_RETURN_IF_NOT(node.OutputDefs()[0]->Name() == transpose_a_output_name); + } else if (node.InputDefs()[0]->Name() == "B") { + transpose_b = &node; + TEST_RETURN_IF_NOT(node.OutputDefs()[0]->Name() == transpose_b_output_name); + } else { + TEST_RETURN_IF_NOT(false); + } + } + } + TEST_RETURN_IF_NOT(gemm_node != nullptr); + TEST_RETURN_IF_NOT(transpose_a != nullptr); + TEST_RETURN_IF_NOT(transpose_b != nullptr); + TEST_RETURN_IF_NOT(gemm_node->InputDefs().size() == 2); + TEST_RETURN_IF_NOT(gemm_node->InputDefs()[0]->Name() == transpose_b_output_name); + TEST_RETURN_IF_NOT(gemm_node->InputDefs()[1]->Name() == transpose_a_output_name); + TEST_RETURN_IF_NOT(gemm_node->GetAttributes().at("transA").i() == 1); + TEST_RETURN_IF_NOT(gemm_node->GetAttributes().at("transB").i() == 1); + TEST_RETURN_IF_NOT(graph.NodeProducesGraphOutput(*gemm_node)); + TEST_RETURN_IF_NOT(gemm_node->OutputDefs().size() == 1); + TEST_RETURN_IF_NOT(gemm_node->OutputDefs()[0]->Name() == final_output_name); + return Status::OK(); + }; + + auto rule_transformer = std::make_unique("RuleTransformer"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::move(rule_transformer), + TransformerLevel::Level1, 1, nullptr, check_graph)); +} + // (A')'B' = AB' where transpose has multiple consumers TEST_F(GraphTransformationTests, GemmTransposeFusion2OutputsFromTranspose) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/gemm_transpose_2outputs_from_transpose.onnx"; diff --git a/onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc b/onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc index 778ea25317c5b..95b563dde1f3e 100644 --- a/onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc @@ -543,6 +543,79 @@ TEST(RMSNormalizationOpTest, RMSNorm_Scale_1xSx3_Axis2) { test.ConfigEp(std::move(cpu)).RunWithConfig(); } +// Rank-reduced scale (S, W) for an (B, S, W) input with axis=2. This is +// NumPy-broadcast-equivalent to a (1, S, W) scale and must produce the same +// result as RMSNorm_Scale_1xSx3_Axis2. It is a regression test for a CUDA +// LayerNorm/RMSNorm kernel bug where a scale whose rank was smaller than the +// input rank fell back to a generic broadcasting path that the CUDA kernel +// does not implement, silently applying scale[0:W] to every row. Runs on all +// available EPs (including CUDA) so the GPU path is covered. +TEST(RMSNormalizationOpTest, RMSNorm_Scale_RankReduced_Sx3_Axis2) { + OpTester test("RMSNormalization", 23); + test.AddAttribute("epsilon", 1e-5f); + test.AddAttribute("axis", 2); + + std::vector x(2 * 4 * 3); + for (int i = 0; i < static_cast(x.size()); ++i) { + x[i] = static_cast(i); + } + test.AddInput("X", {2, 4, 3}, x); + + // Rank-2 scale (S, W) broadcasts over the batch dimension. + test.AddInput("scale", + {4, 3}, + { + 1.1f, + 1.1f, + 1.1f, + 1.2f, + 1.2f, + 1.2f, + 1.3f, + 1.3f, + 1.3f, + 1.4f, + 1.4f, + 1.4f, + }, + true); + + test.AddOutput( + "Y", {2, 4, 3}, + { + 0.0000f, + 0.8521f, + 1.7041f, + 0.8818f, + 1.1758f, + 1.4697f, + 1.1068f, + 1.2912f, + 1.4757f, + 1.2558f, + 1.3954f, + 1.5349f, + + 1.0134f, + 1.0978f, + 1.1823f, + 1.1235f, + 1.1984f, + 1.2733f, + 1.2304f, + 1.2988f, + 1.3672f, + 1.3354f, + 1.3990f, + 1.4626f, + }); + + // TRT, DNNL, OpenVINO, NNAPI, QNN, and WebGPU don't support this combination. + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kDnnlExecutionProvider, kOpenVINOExecutionProvider, + kNnapiExecutionProvider, kQnnExecutionProvider, kWebGpuExecutionProvider}); +} + TEST(RMSNormalizationOpTest, RMSNorm_Scale_NoBroadcast_BxSx3_Axis2) { OpTester test("RMSNormalization", 23); test.AddAttribute("epsilon", 1e-05f); diff --git a/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc b/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc index 7df530e0ff078..5ca69614847e6 100644 --- a/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc @@ -6,7 +6,9 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" #include "core/framework/allocator.h" +#include "core/framework/utils.h" #include "core/providers/cpu/tensor/space_depth_ops.h" #include "core/mlas/inc/mlas.h" #include "test/unittest_util/graph_transform_test_builder.h" @@ -63,6 +65,142 @@ TEST(TensorOpTest, RejectsBlocksizeWhoseSquareExceedsInt64) { EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("multiple of (block_size * block_size)")); } +TEST(TensorOpTest, SpaceToDepthOpset28Modes) { + constexpr int64_t blocksize = 2; + const std::vector input = {0.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, + 8.f, 9.f, 10.f, 11.f, 12.f, 13.f, 14.f, 15.f}; + + for (const auto& [mode, expected] : std::vector>>{ + {"DCR", {0.f, 2.f, 8.f, 10.f, 1.f, 3.f, 9.f, 11.f, 4.f, 6.f, 12.f, 14.f, 5.f, 7.f, 13.f, 15.f}}, + {"CRD", {0.f, 2.f, 1.f, 3.f, 4.f, 6.f, 5.f, 7.f, 8.f, 10.f, 9.f, 11.f, 12.f, 14.f, 13.f, 15.f}}}) { + OpTester test("SpaceToDepth", 28); + test.SetAllowUnreleasedOnnxOpset(); + test.AddAttribute("blocksize", blocksize); + test.AddAttribute("mode", mode.c_str()); + test.AddInput("input", {1, 2, 2, 4}, input); + test.AddOutput("output", {1, 8, 1, 2}, expected); + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + +TEST(TensorOpTest, SpaceToDepthOpset28BlocksizeThree) { + OpTester test("SpaceToDepth", 28); + test.SetAllowUnreleasedOnnxOpset(); + test.AddAttribute("blocksize", int64_t{3}); + test.AddAttribute("mode", "CRD"); + test.AddInput("input", {1, 1, 3, 6}, + {0.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, + 9.f, 10.f, 11.f, 12.f, 13.f, 14.f, 15.f, 16.f, 17.f}); + test.AddOutput("output", {1, 9, 1, 2}, + {0.f, 3.f, 1.f, 4.f, 2.f, 5.f, 6.f, 9.f, 7.f, + 10.f, 8.f, 11.f, 12.f, 15.f, 13.f, 16.f, 14.f, 17.f}); + test.Run(); +} + +#if defined(USE_CUDA) && defined(ENABLE_CUDA_NHWC_OPS) +TEST(TensorOpTest, SpaceToDepthOpset28NhwcCrd) { + auto cuda_provider = DefaultCudaExecutionProvider(); + if (cuda_provider == nullptr) { + GTEST_SKIP() << "CUDA execution provider is not available."; + } + + OpTester test("SpaceToDepth", 28, kMSInternalNHWCDomain); + test.SetAllowUnreleasedOnnxOpset(); + test.AddAttribute("blocksize", int64_t{2}); + test.AddAttribute("mode", "CRD"); + test.AddInput("input", {1, 2, 4, 2}, + {0.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, + 8.f, 9.f, 10.f, 11.f, 12.f, 13.f, 14.f, 15.f}); + test.AddOutput("output", {1, 1, 2, 8}, + {0.f, 2.f, 8.f, 10.f, 1.f, 3.f, 9.f, 11.f, + 4.f, 6.f, 12.f, 14.f, 5.f, 7.f, 13.f, 15.f}); + + std::vector> execution_providers; + execution_providers.push_back(std::move(cuda_provider)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + +TEST(TensorOpTest, DepthToSpaceOpset28Modes) { + constexpr int64_t blocksize = 2; + const std::vector expected = {0.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, + 8.f, 9.f, 10.f, 11.f, 12.f, 13.f, 14.f, 15.f}; + + for (const auto& [mode, input] : std::vector>>{ + {"DCR", {0.f, 2.f, 8.f, 10.f, 1.f, 3.f, 9.f, 11.f, 4.f, 6.f, 12.f, 14.f, 5.f, 7.f, 13.f, 15.f}}, + {"CRD", {0.f, 2.f, 1.f, 3.f, 4.f, 6.f, 5.f, 7.f, 8.f, 10.f, 9.f, 11.f, 12.f, 14.f, 13.f, 15.f}}}) { + OpTester test("DepthToSpace", 28); + test.SetAllowUnreleasedOnnxOpset(); + test.AddAttribute("blocksize", blocksize); + test.AddAttribute("mode", mode.c_str()); + test.AddInput("input", {1, 8, 1, 2}, input); + test.AddOutput("output", {1, 2, 2, 4}, expected); + test.Run(); + } +} + +TEST(TensorOpTest, DepthToSpaceOpset28BlocksizeThree) { + OpTester test("DepthToSpace", 28); + test.SetAllowUnreleasedOnnxOpset(); + test.AddAttribute("blocksize", int64_t{3}); + test.AddAttribute("mode", "CRD"); + test.AddInput("input", {1, 9, 1, 2}, + {0.f, 3.f, 1.f, 4.f, 2.f, 5.f, 6.f, 9.f, 7.f, + 10.f, 8.f, 11.f, 12.f, 15.f, 13.f, 16.f, 14.f, 17.f}); + test.AddOutput("output", {1, 1, 3, 6}, + {0.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, + 9.f, 10.f, 11.f, 12.f, 13.f, 14.f, 15.f, 16.f, 17.f}); + test.Run(); +} + +TEST(TensorOpTest, SpaceDepthOpset28RejectsInvalidAttributesAndShapes) { + const auto run_on_cpu = [](OpTester& test, const char* expected_failure) { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, expected_failure, {}, nullptr, &execution_providers); + }; + + { + OpTester test("SpaceToDepth", 28); + test.SetAllowUnreleasedOnnxOpset(); + test.AddAttribute("blocksize", int64_t{2}); + test.AddAttribute("mode", "invalid"); + test.AddInput("input", {1, 1, 2, 2}, {0.f, 1.f, 2.f, 3.f}); + test.AddOutput("output", {1, 4, 1, 1}, {0.f, 1.f, 2.f, 3.f}); + run_on_cpu(test, "only 'DCR' and 'CRD' modes are supported"); + } + + { + OpTester test("DepthToSpace", 28); + test.SetAllowUnreleasedOnnxOpset(); + test.AddAttribute("blocksize", int64_t{2}); + test.AddAttribute("mode", "invalid"); + test.AddInput("input", {1, 4, 1, 1}, {0.f, 1.f, 2.f, 3.f}); + test.AddOutput("output", {1, 1, 2, 2}, {0.f, 1.f, 2.f, 3.f}); + run_on_cpu(test, "only 'DCR' and 'CRD' modes are supported"); + } + + { + OpTester test("SpaceToDepth", 28); + test.SetAllowUnreleasedOnnxOpset(); + test.AddAttribute("blocksize", int64_t{2}); + test.AddInput("input", {1, 1, 3, 4}, std::vector(12)); + test.AddOutput("output", {1, 4, 1, 2}, std::vector(8)); + run_on_cpu(test, "input height to be a multiple"); + } + + { + OpTester test("DepthToSpace", 28); + test.SetAllowUnreleasedOnnxOpset(); + test.AddAttribute("blocksize", int64_t{2}); + test.AddInput("input", {1, 3, 1, 1}, std::vector(3)); + test.AddOutput("output", {1, 1, 2, 2}, std::vector(4)); + run_on_cpu(test, "Can't merge shape info"); + } +} + TEST(TensorOpTest, SpaceToDepthTest_1) { OpTester test("SpaceToDepth"); constexpr int64_t blocksize = 2; @@ -155,8 +293,10 @@ TEST(TensorOpTest, SpaceToDepthOpset28ModesAndProviderRouting) { InlinedVector> providers; providers.push_back(DefaultCpuExecutionProvider()); providers.push_back(DefaultCudaExecutionProvider()); + const IExecutionProvider* cuda_nhwc_provider = nullptr; #ifdef ENABLE_CUDA_NHWC_OPS providers.push_back(DefaultCudaNHWCExecutionProvider()); + cuda_nhwc_provider = providers.back().get(); #endif providers.push_back(DefaultDmlExecutionProvider()); providers.push_back(DefaultOpenVINOExecutionProvider()); @@ -173,7 +313,9 @@ TEST(TensorOpTest, SpaceToDepthOpset28ModesAndProviderRouting) { continue; } const std::string provider_type = provider->Type(); + const bool is_cuda_nhwc = provider.get() == cuda_nhwc_provider; SCOPED_TRACE(provider_type); + SCOPED_TRACE(is_cuda_nhwc); SessionOptions options; options.graph_optimization_level = TransformerLevel::Default; options.enable_mem_pattern = false; @@ -182,15 +324,36 @@ TEST(TensorOpTest, SpaceToDepthOpset28ModesAndProviderRouting) { ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(provider))); ASSERT_STATUS_OK(session.Load(model_data.data(), static_cast(model_data.size()))); ASSERT_STATUS_OK(session.Initialize()); - bool has_space_to_depth = false; - for (const auto& resolved_node : session.GetGraph().Nodes()) { - has_space_to_depth |= resolved_node.OpType() == "SpaceToDepth"; - } - if (opset == 28 && (mode == "CRD" || provider_type == kCpuExecutionProvider || - provider_type == kCudaExecutionProvider)) { - EXPECT_FALSE(has_space_to_depth); - } else if (opset < 28 && provider_type == kCpuExecutionProvider) { - EXPECT_TRUE(has_space_to_depth); + const bool falls_back_to_cpu = + opset == 28 && ((provider_type == kDmlExecutionProvider && mode == "CRD") || + provider_type == kOpenVINOExecutionProvider || + provider_type == kMIGraphXExecutionProvider || + provider_type == kCannExecutionProvider); + const bool uses_native_kernel = falls_back_to_cpu || provider_type == kCpuExecutionProvider || + provider_type == kCudaExecutionProvider; + if (uses_native_kernel || provider_type == kDmlExecutionProvider || + provider_type == kOpenVINOExecutionProvider) { + const std::string expected_provider = falls_back_to_cpu ? kCpuExecutionProvider : provider_type; + size_t compute_node_count = 0; + size_t layout_transpose_count = 0; + for (const auto& resolved_node : session.GetGraph().Nodes()) { + if (utils::IsMemcpyNode(resolved_node)) { + continue; + } + EXPECT_EQ(resolved_node.GetExecutionProviderType(), expected_provider); + // NHWC CUDA wraps the native kernel with NCHW/NHWC layout conversions. + if (is_cuda_nhwc && resolved_node.OpType() == "Transpose") { + ++layout_transpose_count; + continue; + } + ++compute_node_count; + // Graph EPs may replace supported operators with fused nodes. + if (uses_native_kernel) { + EXPECT_EQ(resolved_node.OpType(), "SpaceToDepth"); + } + } + EXPECT_EQ(compute_node_count, 1u); + EXPECT_EQ(layout_transpose_count, is_cuda_nhwc ? 2u : 0u); } std::vector fetches; diff --git a/onnxruntime/test/providers/cuda/test_cases/cuda_test_bridge.cc b/onnxruntime/test/providers/cuda/test_cases/cuda_test_bridge.cc index 4dff2a64f94a9..cb0d0a5f87ee7 100644 --- a/onnxruntime/test/providers/cuda/test_cases/cuda_test_bridge.cc +++ b/onnxruntime/test/providers/cuda/test_cases/cuda_test_bridge.cc @@ -10,20 +10,22 @@ namespace onnxruntime::test { #if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) std::optional GetGroupQueryAttentionWorkspaceEstimateConfigForTest( const void* node, bool head_sink_is_constant_initializer, - int64_t max_total_sequence_length) { + int64_t max_total_sequence_length, + KernelDispatchPolicy dispatch_policy) { return contrib::cuda::GetGroupQueryAttentionWorkspaceEstimateConfig( *static_cast(node), head_sink_is_constant_initializer, - max_total_sequence_length); + max_total_sequence_length, dispatch_policy); } std::optional EstimateGroupQueryAttentionWorkspaceForTest( const void* node, gsl::span input_shapes, const cudaDeviceProp& device_prop, const AttentionKernelOptions& kernel_options, bool head_sink_is_constant_initializer, - int64_t max_total_sequence_length) { + int64_t max_total_sequence_length, + KernelDispatchPolicy dispatch_policy) { return contrib::cuda::EstimateGroupQueryAttentionWorkspace( *static_cast(node), input_shapes, device_prop, kernel_options, - head_sink_is_constant_initializer, max_total_sequence_length); + head_sink_is_constant_initializer, max_total_sequence_length, dispatch_policy); } std::optional EstimatePackedAttentionWorkspaceForTest( diff --git a/onnxruntime/test/providers/cuda/test_cases/cuda_test_bridge.h b/onnxruntime/test/providers/cuda/test_cases/cuda_test_bridge.h index 2bd13c9334724..69e399692a579 100644 --- a/onnxruntime/test/providers/cuda/test_cases/cuda_test_bridge.h +++ b/onnxruntime/test/providers/cuda/test_cases/cuda_test_bridge.h @@ -21,13 +21,15 @@ namespace onnxruntime::test { #if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) std::optional GetGroupQueryAttentionWorkspaceEstimateConfigForTest( const void* node, bool head_sink_is_constant_initializer = false, - int64_t max_total_sequence_length = 0); + int64_t max_total_sequence_length = 0, + KernelDispatchPolicy dispatch_policy = KernelDispatchPolicy::Auto); std::optional EstimateGroupQueryAttentionWorkspaceForTest( const void* node, gsl::span input_shapes, const cudaDeviceProp& device_prop, const AttentionKernelOptions& kernel_options, bool head_sink_is_constant_initializer = false, - int64_t max_total_sequence_length = 0); + int64_t max_total_sequence_length = 0, + KernelDispatchPolicy dispatch_policy = KernelDispatchPolicy::Auto); std::optional EstimatePackedAttentionWorkspaceForTest( const void* node, gsl::span input_shapes, 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 ab2c385685ea5..98109ff3211eb 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 @@ -391,7 +391,36 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, NodeAdapterParsesAttributesAndTyp EXPECT_GT(estimate->total_workspace_bytes, 0u); } -TEST(GroupQueryAttentionWorkspaceEstimateTest, ParsedEnvelopeReachesNodeConfigWithoutChangingEstimate) { +TEST(GroupQueryAttentionWorkspaceEstimateTest, NodeAdapterForwardsDispatchPolicy) { + WithGroupQueryAttentionNode([](const Node& node) { + const auto default_config = GetGroupQueryAttentionWorkspaceEstimateConfigForTest(&node); + ASSERT_TRUE(default_config.has_value()); + EXPECT_EQ(default_config->dispatch_policy, KernelDispatchPolicy::Auto); + + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + const auto baseline = EstimateGroupQueryAttentionWorkspaceForTest( + &node, SeparateShapes(), Device(), options); + ASSERT_TRUE(baseline.has_value()); + + for (const auto policy : {KernelDispatchPolicy::Auto, KernelDispatchPolicy::Latency, + KernelDispatchPolicy::Memory, KernelDispatchPolicy::Safe}) { + SCOPED_TRACE(static_cast(policy)); + const auto config = GetGroupQueryAttentionWorkspaceEstimateConfigForTest( + &node, false, 4096, policy); + ASSERT_TRUE(config.has_value()); + EXPECT_EQ(config->max_total_sequence_length, int64_t{4096}); + EXPECT_EQ(config->dispatch_policy, policy); + const auto estimate = EstimateGroupQueryAttentionWorkspaceForTest( + &node, SeparateShapes(), Device(), options, false, 4096, policy); + ASSERT_TRUE(estimate.has_value()); + EXPECT_EQ(estimate->total_workspace_bytes, baseline->total_workspace_bytes); + EXPECT_EQ(estimate->sized_backends, baseline->sized_backends); + } + }); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, ParsedEnvelopePreservesWindowedEstimateAndEnablesNonWindowed) { WithGroupQueryAttentionNode([](const Node& node) { AttentionKernelOptions options; options.InitializeOnce(kMath, true); @@ -692,6 +721,46 @@ TEST(GroupQueryAttentionWorkspaceBoundsTest, PartialAliasPreservationAddsFullPas without.total_workspace_bytes + expected_copy); } +TEST(GroupQueryAttentionWorkspaceEstimateTest, DispatchPolicyIsReaderOnlyForBothCacheModes) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + for (bool windowed : {false, true}) { + SCOPED_TRACE(windowed); + auto config = Config(); + EXPECT_EQ(config.dispatch_policy, KernelDispatchPolicy::Auto); + config.sliding_window_cache = windowed; + config.local_window_size = windowed ? 256 : -1; + config.max_total_sequence_length = windowed ? 0 : 512; + const auto baseline = EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(), Device(), options); + ASSERT_TRUE(baseline.has_value()); + for (const auto policy : {KernelDispatchPolicy::Auto, KernelDispatchPolicy::Latency, + KernelDispatchPolicy::Memory, KernelDispatchPolicy::Safe}) { + config.dispatch_policy = policy; + const auto estimate = EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(), Device(), options); + ASSERT_TRUE(estimate.has_value()); + EXPECT_EQ(estimate->total_workspace_bytes, baseline->total_workspace_bytes); + EXPECT_EQ(estimate->sized_backends, baseline->sized_backends); + } + } +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, DispatchPolicyDoesNotEnableUnspecifiedEnvelope) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto config = Config(); + config.sliding_window_cache = false; + config.local_window_size = -1; + for (const auto policy : {KernelDispatchPolicy::Latency, KernelDispatchPolicy::Memory, + KernelDispatchPolicy::Safe, KernelDispatchPolicy::Auto}) { + config.dispatch_policy = policy; + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(), Device(), options) + .has_value()); + } +} + TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsCacheCapacityDifferentFromWindow) { AttentionKernelOptions options; options.InitializeOnce(kMath, true); diff --git a/onnxruntime/test/providers/cuda/test_cases/matmul_nbits_workspace_test.cc b/onnxruntime/test/providers/cuda/test_cases/matmul_nbits_workspace_test.cc index 1dd5c240435a4..424a4fc83a38c 100644 --- a/onnxruntime/test/providers/cuda/test_cases/matmul_nbits_workspace_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/matmul_nbits_workspace_test.cc @@ -37,6 +37,7 @@ // below (TensorProto_DataType_FLOAT16 etc.) come from provider_api.h via matmul_nbits.h. #include "contrib_ops/cuda/llm/fpA_intB_gemm/fpA_intB_gemm.h" #include "contrib_ops/cuda/quantization/matmul_nbits.h" +#include "contrib_ops/cuda/quantization/matmul_nbits.cuh" #include "test/providers/cuda/test_cases/matmul_nbits_workspace_test_probe.h" namespace onnxruntime { @@ -79,6 +80,112 @@ bool CheckDefault(int32_t elem_type = kFp16, int64_t N = 256, int64_t K = 1024, } } // namespace +TEST(MatMulNBitsWorkspace, DeterministicTacticsIgnoreProfiledCache) { + using namespace onnxruntime::llm::cutlass_extensions; + using namespace onnxruntime::llm::kernels::weight_only; + using onnxruntime::llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunner; + using Runner = CutlassFpAIntBGemmRunner; + + class TestProfiler : public WeightOnlyGroupwiseQuantGemmPluginProfiler { + public: + TestProfiler() { + mRunner = std::make_shared(); + mRunner->setArch(80); + setCudaKernelType({}, 80); + } + } profiler; + + const GemmIdCore gemm_id(640, 2560, onnxruntime::llm::nvinfer::DataType::kHALF, 80); + auto cache = std::make_shared(); + cache->createMProfileMap(gemm_id); + profiler.setSelectionTactics(cache); + + const auto decode = profiler.getDeterministicConfig(1); + const auto prefill = profiler.getDeterministicConfig(77); + ASSERT_TRUE(decode.has_value()); + ASSERT_TRUE(prefill.has_value()); + EXPECT_TRUE(decode->enableCudaKernel); + EXPECT_FALSE(prefill->enableCudaKernel); + EXPECT_EQ(prefill->split_k_style, SplitKStyle::NO_SPLIT_K); + EXPECT_EQ(prefill->split_k_factor, 1); + + for (int split_k : {2, 5}) { + const CutlassGemmConfig timed(CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64, + SplitKStyle::SPLIT_K_SERIAL, split_k, 3); + for (int rows : {1, 77}) { + (*cache->getMProfileMap(gemm_id))[rows] = timed; + ASSERT_TRUE(profiler.getBestConfig(rows, gemm_id).has_value()); + EXPECT_EQ(profiler.getBestConfig(rows, gemm_id)->split_k_factor, split_k); + const auto selected = profiler.getDeterministicConfig(rows); + ASSERT_TRUE(selected.has_value()); + EXPECT_EQ(selected->toString(), (rows == 1 ? decode : prefill)->toString()); + } + } +} + +TEST(MatMulNBitsWorkspace, DeterministicTacticsUseSupportedStages) { + using namespace onnxruntime::llm::cutlass_extensions; + using onnxruntime::llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunner; + using Runner = CutlassFpAIntBGemmRunner; + + class TestProfiler : public WeightOnlyGroupwiseQuantGemmPluginProfiler { + public: + explicit TestProfiler(int arch) { + mRunner = std::make_shared(); + mRunner->setArch(arch); + setCudaKernelType({}, arch); + } + }; + + for (int arch : {75, 80, 86, 89, 90}) { + SCOPED_TRACE(arch); + TestProfiler profiler(arch); + for (int rows : {16, 77}) { + const auto config = profiler.getDeterministicConfig(rows); + ASSERT_TRUE(config.has_value()); + EXPECT_FALSE(config->enableCudaKernel); + EXPECT_EQ(config->split_k_style, SplitKStyle::NO_SPLIT_K); + if (config->is_tma_warp_specialized) { + EXPECT_EQ(arch, 90); + EXPECT_EQ(config->sm_version, 90); + EXPECT_NE(config->tile_config_sm90, CutlassTileConfigSM90::Undefined); + EXPECT_NE(config->tile_config_sm90, CutlassTileConfigSM90::ChooseWithHeuristic); + // Native SM90 tactics do not use the legacy split-K factor or pipeline-stage count. + EXPECT_EQ(config->split_k_factor, -1); + EXPECT_EQ(config->stages, -1); + } else { + EXPECT_LT(config->sm_version, 90); + EXPECT_EQ(config->split_k_factor, 1); + EXPECT_GE(config->stages, arch < 89 ? 2 : 3); + EXPECT_LE(config->stages, arch < 80 ? 2 : 4); + } + } + } +} + +TEST(MatMulNBitsWorkspace, SingleRowInt4GemvShapeEligibility) { + using onnxruntime::contrib::cuda::IsSupportedSingleRowGemvShape; + EXPECT_TRUE(IsSupportedSingleRowGemvShape(false, 1, 248320, 2560, 32)); + EXPECT_TRUE(IsSupportedSingleRowGemvShape(false, 1, 32, 2880, 32)); + EXPECT_TRUE(IsSupportedSingleRowGemvShape(false, 1, 32, 2880, 64)); + EXPECT_FALSE(IsSupportedSingleRowGemvShape(true, 1, 248320, 2560, 32)); + EXPECT_FALSE(IsSupportedSingleRowGemvShape(true, 1, 32, 2880, 32)); + for (int rows : {0, 2, 8}) { + EXPECT_FALSE(IsSupportedSingleRowGemvShape(false, rows, 248320, 2560, 32)); + } + for (int columns : {248288, 248319, 248321, 248352}) { + EXPECT_FALSE(IsSupportedSingleRowGemvShape(false, 1, columns, 2560, 32)); + } + for (int reduction : {2528, 2559, 2561, 2592}) { + EXPECT_FALSE(IsSupportedSingleRowGemvShape(false, 1, 248320, reduction, 32)); + } + for (int block_size : {16, 64, 128}) { + EXPECT_FALSE(IsSupportedSingleRowGemvShape(false, 1, 248320, 2560, block_size)); + } +} + // --------------------------------------------------------------------------- // Shared workspace-size formula (ComputeFpAIntBGemmWorkspaceSize). // The non-SM90 branch is device-independent so it can be asserted with exact values. diff --git a/onnxruntime/test/providers/qnn/qnn_test_utils.cc b/onnxruntime/test/providers/qnn/qnn_test_utils.cc index a6d43a3d3a9d9..5187474ed0761 100644 --- a/onnxruntime/test/providers/qnn/qnn_test_utils.cc +++ b/onnxruntime/test/providers/qnn/qnn_test_utils.cc @@ -101,7 +101,8 @@ void TryEnableQNNSaver(ProviderOptions& qnn_options) { void RunQnnModelTest(const GetTestModelFn& build_test_case, ProviderOptions provider_options, int opset_version, ExpectedEPNodeAssignment expected_ep_assignment, float fp32_abs_err, logging::Severity log_severity, bool verify_outputs, - std::function* ep_graph_checker) { + std::function* ep_graph_checker, + bool allow_released_opsets_only) { std::filesystem::path output_dir; if (QNNTestEnvironment::GetInstance().dump_onnx() || QNNTestEnvironment::GetInstance().dump_json() || @@ -124,7 +125,8 @@ void RunQnnModelTest(const GetTestModelFn& build_test_case, ProviderOptions prov onnxruntime::Model model("QNN_EP_TestModel", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {}, - logging_manager.DefaultLogger()); + logging_manager.DefaultLogger(), + ModelOptions{allow_released_opsets_only, false}); Graph& graph = model.MainGraph(); ModelTestBuilder helper(graph); build_test_case(helper); @@ -158,7 +160,7 @@ void RunQnnModelTest(const GetTestModelFn& build_test_case, ProviderOptions prov RunAndVerifyOutputsWithEP(AsByteSpan(model_data.data(), model_data.size()), "QNN_EP_TestLogID", QnnExecutionProviderWithOptions(provider_options), helper.feeds_, verification_params, - {}, verify_outputs); + {}, verify_outputs, allow_released_opsets_only); } void RunQnnModelTestHTPNoVerify(const GetTestModelFn& build_test_case, ProviderOptions provider_options, diff --git a/onnxruntime/test/providers/qnn/qnn_test_utils.h b/onnxruntime/test/providers/qnn/qnn_test_utils.h index 88780d81e5d1a..2865895faf018 100644 --- a/onnxruntime/test/providers/qnn/qnn_test_utils.h +++ b/onnxruntime/test/providers/qnn/qnn_test_utils.h @@ -1322,7 +1322,8 @@ void RunQnnModelTest(const GetTestModelFn& build_test_case, ProviderOptions prov float fp32_abs_err = 1e-5f, logging::Severity log_severity = logging::Severity::kERROR, bool verify_outputs = true, - std::function* ep_graph_checker = nullptr); + std::function* ep_graph_checker = nullptr, + bool allow_released_opsets_only = true); /** * Runs a test model on the QNN HTP backend and verifies node assignment to QNN EP without comparing outputs to ORT CPU. diff --git a/onnxruntime/test/providers/qnn/simple_op_htp_test.cc b/onnxruntime/test/providers/qnn/simple_op_htp_test.cc index b84b361b61367..2d0d0907e3c54 100644 --- a/onnxruntime/test/providers/qnn/simple_op_htp_test.cc +++ b/onnxruntime/test/providers/qnn/simple_op_htp_test.cc @@ -25,7 +25,8 @@ static void RunOpTestOnCPU(const std::string& op_type, const std::vector& attrs, int opset_version, ExpectedEPNodeAssignment expected_ep_assignment, - const std::string& op_domain = kOnnxDomain) { + const std::string& op_domain = kOnnxDomain, + bool allow_released_opsets_only = true) { ProviderOptions provider_options; provider_options["backend_type"] = "cpu"; provider_options["offload_graph_io_quantization"] = "0"; @@ -33,7 +34,8 @@ static void RunOpTestOnCPU(const std::string& op_type, RunQnnModelTest(BuildOpTestCase(op_type, input_defs, {}, attrs, op_domain), provider_options, opset_version, - expected_ep_assignment); + expected_ep_assignment, + 1e-5f, logging::Severity::kERROR, true, nullptr, allow_released_opsets_only); } template @@ -44,7 +46,8 @@ static void RunOpTestOnCPU(const std::string& op_type, const std::vector& attrs, int opset_version, ExpectedEPNodeAssignment expected_ep_assignment, - const std::string& op_domain = kOnnxDomain) { + const std::string& op_domain = kOnnxDomain, + bool allow_released_opsets_only = true) { ProviderOptions provider_options; provider_options["backend_type"] = "cpu"; provider_options["offload_graph_io_quantization"] = "0"; @@ -52,7 +55,8 @@ static void RunOpTestOnCPU(const std::string& op_type, RunQnnModelTest(BuildOpTestCase(op_type, input_defs_1, input_defs_2, input_defs_3, attrs, op_domain), provider_options, opset_version, - expected_ep_assignment); + expected_ep_assignment, + 1e-5f, logging::Severity::kERROR, true, nullptr, allow_released_opsets_only); } // Test float DepthToSpace on the QNN CPU backend. @@ -112,6 +116,17 @@ TEST_F(QnnCPUBackendTests, SpaceToDepth_Flaky2) { } } +TEST_F(QnnCPUBackendTests, SpaceToDepth_Opset28CrdFallsBackToCpu) { + RunOpTestOnCPU("SpaceToDepth", + {TestInputDef({1, 1, 2, 2}, false, {0.f, 1.f, 2.f, 3.f})}, + {utils::MakeAttribute("blocksize", int64_t{2}), + utils::MakeAttribute("mode", "CRD")}, + 28, + ExpectedEPNodeAssignment::None, + kOnnxDomain, + false); +} + // Test f32 Relu on the CPU backend. // TODO: When this is fixed, enable ActivationOpTest.Relu test in cpu/activation/activation_op_test tests. // Disabled because QNN SDK 2.17 Relu treats inf as FLT_MAX. diff --git a/onnxruntime/test/python/transformers/benchmark_qwen38_mtp_kernels.py b/onnxruntime/test/python/transformers/benchmark_qwen38_mtp_kernels.py index e338337beee87..a6ed0c4481769 100644 --- a/onnxruntime/test/python/transformers/benchmark_qwen38_mtp_kernels.py +++ b/onnxruntime/test/python/transformers/benchmark_qwen38_mtp_kernels.py @@ -9,7 +9,26 @@ import cupy import numpy -ROOT = Path(__file__).resolve().parents[3] + +def resolve_ort_root(script_path=None): + raw_script = Path(script_path or __file__) + marker = Path("contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu") + roots = [*raw_script.resolve().parents, *raw_script.absolute().parents] + for env_name in ("GITHUB_WORKSPACE", "BUILD_SOURCESDIRECTORY"): + env_value = os.environ.get(env_name) + if env_value: + env_path = Path(env_value).resolve() + roots.extend((env_path, *env_path.parents)) + cwd = Path.cwd().resolve() + roots.extend((cwd, *cwd.parents)) + for base in roots: + for candidate in (base, base / "onnxruntime", base / "onnxruntime" / "onnxruntime"): + if (candidate / marker).is_file(): + return candidate + return raw_script.resolve().parents[3] + + +ROOT = resolve_ort_root() QSA_SOURCE = ROOT / "contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu" diff --git a/onnxruntime/test/python/transformers/test_benchmark_qwen38_mtp_kernels.py b/onnxruntime/test/python/transformers/test_benchmark_qwen38_mtp_kernels.py index aa1984843c549..d6aa16d0d830d 100644 --- a/onnxruntime/test/python/transformers/test_benchmark_qwen38_mtp_kernels.py +++ b/onnxruntime/test/python/transformers/test_benchmark_qwen38_mtp_kernels.py @@ -2,7 +2,9 @@ # Licensed under the MIT License. import importlib.util +import os import sys +import tempfile import unittest.mock from pathlib import Path @@ -58,6 +60,7 @@ def test_nvfp4_generated_source_excludes_host_inspection(self): __device__ __forceinline__ float PackedLowToFloat(uint32_t word) { return word; } template __global__ void MoeGemvFp4DecodeKernel() {} +static bool Fp4MoeGemvUsePdl() { return onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_PDL", 1) != 0; } template void LaunchMoeGemvFp4RawNPacked() {} """ @@ -90,8 +93,53 @@ def test_nvfp4_generated_source_excludes_host_inspection(self): self.assertIn("__global__ void MoeGemvFp4RawKPackedKernel", generated) self.assertNotIn("bool Fp4MoeGemvSupportsNativeScale", generated) self.assertNotIn("bool Fp4MoeGemvKernelUsesNativeScale", generated) + self.assertNotIn("Fp4MoeGemvUsePdl", generated) + self.assertNotIn("ParseEnvironmentVariableWithDefault", generated) self.assertNotIn("cudaFuncGetAttributes", generated) + def test_real_moe_gemv_fp4_source_excludes_host_helpers(self): + real_source = (self.benchmark.ROOT / "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu").read_text() + for decode in (False, True): + with ( + self.subTest(decode=decode), + unittest.mock.patch.object( + self.benchmark, "compile_device_code", return_value=object() + ) as compile_code, + ): + self.benchmark.compile_nvfp4(real_source, "half", True, False, False, False, decode=decode) + generated, _ = compile_code.call_args.args + self.assertNotIn("Fp4MoeGemvSupportsNativeScale", generated) + self.assertNotIn("Fp4MoeGemvKernelUsesNativeScale", generated) + self.assertNotIn("Fp4MoeGemvUsePdl", generated) + self.assertNotIn("QuerySm90DecodeKernelCaps", generated) + self.assertNotIn("ParseEnvironmentVariableWithDefault", generated) + self.assertNotIn("cudaFuncGetAttributes", generated) + + def test_resolve_ort_root_from_copied_build_tree_and_workspace_env(self): + marker = Path("contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu") + with tempfile.TemporaryDirectory() as temp_dir: + temp_root = Path(temp_dir).resolve() + ort_root = temp_root / "repo" / "onnxruntime" + (ort_root / marker).parent.mkdir(parents=True, exist_ok=True) + (ort_root / marker).write_text("// stub\n") + + copied_in_repo = ( + temp_root + / "repo" + / "build" + / "Release" + / "Release" + / "transformers" + / "benchmark_qwen38_mtp_kernels.py" + ) + self.assertEqual(self.benchmark.resolve_ort_root(copied_in_repo), ort_root) + + copied_in_temp = ( + temp_root / "_temp" / "build" / "RelWithDebInfo" / "RelWithDebInfo" / "transformers" / "benchmark.py" + ) + with unittest.mock.patch.dict(os.environ, {"GITHUB_WORKSPACE": str(temp_root / "repo")}, clear=False): + self.assertEqual(self.benchmark.resolve_ort_root(copied_in_temp), ort_root) + 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 2b95ffa455fe4..1b41ce9109f45 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py @@ -632,7 +632,7 @@ def _compute_reference( output = torch.zeros(num_tokens, hidden_size, device=x.device, dtype=torch.float32) expert_mask = F.one_hot(topk_idx, num_classes=num_experts).permute(2, 1, 0) - for e in range(num_experts): + for e in topk_idx.unique().tolist(): idx, top_x = torch.where(expert_mask[e]) if top_x.shape[0] == 0: continue @@ -1110,6 +1110,24 @@ def test_nvfp4_offline_row_major_cuda_graph(self, onnx_dtype): True, dict(hidden_size=2560, inter_size=640, num_experts=16, top_k=10, num_tokens=6, use_bias=False), ), + ( + False, + False, + True, + True, + True, + True, + dict(hidden_size=5120, inter_size=640, num_experts=16, top_k=10, num_tokens=2, use_bias=True), + ), + ( + True, + True, + True, + True, + True, + True, + dict(hidden_size=512, inter_size=512, num_experts=512, top_k=10, num_tokens=1, use_bias=True), + ), ( True, True, @@ -1200,17 +1218,26 @@ def test_nvfp4_offline_row_major_cuda_graph(self, onnx_dtype): print(f"DECODE_GEMV_CASE_BEGIN\\t{idx}", flush=True) with profile(activities=[ProfilerActivity.CUDA]) as prof: out = test._run_nvfp4_moe_test(onnx_dtype=onnx_dtype, gemv_mode="1", offline_prepacked=True, **case) - names = sorted({e.name for e in prof.events() if "MoeGemvFp4" in e.name}) + names = sorted({ + e.name for e in prof.events() + if "MoeGemvFp4" in e.name or "FusedSoftmaxTopKMapsKernel" in e.name + }) payload = base64.b64encode(out.view(torch.int16).cpu().numpy().tobytes()).decode() print("DECODE_GEMV_RESULT", str(idx), json.dumps(names), payload, sep="\\t", flush=True) """ - def _run_decode_gemv_subprocess(self, onnx_dtype, cases, decode_mode, native_scale=0): + def _run_decode_gemv_subprocess(self, onnx_dtype, cases, decode_mode, native_scale=0, extra_env=None): env = dict(os.environ) env["ORT_FP4_GEMV_DECODE_KERNEL"] = str(decode_mode) + if decode_mode == 0: + env["ORT_FP4_GEMV_FUSED_ROUTE"] = "0" + env["ORT_FP4_GEMV_PDL"] = "0" + env["ORT_FP4_GEMV_FC1_PDL"] = "0" env["ORT_NVFP4_GEMV_NATIVE_SCALE"] = str(native_scale) env["ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO"] = "1" env.pop("ORT_NVFP4_FC1_EXPERT_REUSE", None) + if extra_env: + env.update(extra_env) proc = subprocess.run( [sys.executable, "-c", self._DECODE_GEMV_RUNNER_SCRIPT, str(int(onnx_dtype)), json.dumps(cases)], cwd=os.path.dirname(os.path.abspath(__file__)), @@ -1313,6 +1340,11 @@ def test_nvfp4_offline_row_major_decode_gemv_bit_exact(self, onnx_dtype): def_fc1 and can_decode_default, def_fc2 and can_decode_default, ) + native_fc1_decode = ( + def_fc1 + and can_decode_default + and not (onnx_dtype == TensorProto.FLOAT16 and case["hidden_size"] <= 1024) + ) native_fc2_decode = ( def_fc2 and can_decode_default @@ -1325,7 +1357,7 @@ def test_nvfp4_offline_row_major_decode_gemv_bit_exact(self, onnx_dtype): native_cupti, fc1_scale_ok, fc2_scale_ok, - def_fc1 and can_decode_default, + native_fc1_decode, native_fc2_decode, ) can_decode_optin = sm >= 90 and decode_image_available @@ -1357,6 +1389,63 @@ def test_nvfp4_offline_row_major_decode_gemv_bit_exact(self, onnx_dtype): f"opt-in decode GEMV output differs bitwise from the KPacked kernel for {case}", ) + @parameterized.expand([(TensorProto.FLOAT16,), (TensorProto.BFLOAT16,)]) + def test_nvfp4_offline_row_major_decode_gemv_fc1_pdl(self, onnx_dtype): + self._skip_if_no_fp4() + base_case = dict( + hidden_size=512, + inter_size=512, + num_experts=512, + top_k=2, + num_tokens=2, + use_bias=False, + use_swiglu=True, + input_scale=0.1, + ) + cases = [ + {**base_case, "enable_cuda_graph": False}, + {**base_case, "enable_cuda_graph": True}, + ] + _, baseline_runs = self._run_decode_gemv_subprocess(onnx_dtype, cases, decode_mode=0) + pdl_cupti, pdl_runs = self._run_decode_gemv_subprocess( + onnx_dtype, + cases, + decode_mode=1, + extra_env={ + "ORT_FP4_GEMV_FUSED_ROUTE": "1", + "ORT_FP4_GEMV_PDL": "1", + "ORT_FP4_GEMV_FC1_PDL": "1", + "ORT_DISABLE_FP4_GEMV_SKIP_EXPAND": "0", + }, + ) + sm = _cuda_sm() + for case, (_, _, baseline_payload), (pdl_diag, pdl_kernels, pdl_payload) in zip( + cases, baseline_runs, pdl_runs, strict=True + ): + decode_image_available = pdl_diag.get("DecodeImageAvailable") == "1" + can_decode_default = sm == 90 and decode_image_available + self._assert_fc1_fc2_gemv_dispatch( + case, + pdl_diag, + pdl_kernels, + pdl_cupti, + True, + True, + can_decode_default, + can_decode_default, + ) + if pdl_cupti and sm >= 90 and decode_image_available and not case["enable_cuda_graph"]: + self.assertTrue( + any("FusedSoftmaxTopKMapsKernel" in name for name in pdl_kernels), + (case, pdl_kernels), + ) + expected = torch.frombuffer(bytearray(base64.b64decode(baseline_payload)), dtype=torch.int16) + actual = torch.frombuffer(bytearray(base64.b64decode(pdl_payload)), dtype=torch.int16) + self.assertTrue( + torch.equal(actual, expected), + f"FC1 PDL decode GEMV output differs bitwise from serialized baseline for {case}", + ) + def test_nvfp4_fp16_prepacking_disabled_uses_raw_fallback(self): self._run_nvfp4_moe_test( hidden_size=64, diff --git a/onnxruntime/test/unittest_util/function_test_util.cc b/onnxruntime/test/unittest_util/function_test_util.cc index 5edcd1db4541b..50411b8b779a6 100644 --- a/onnxruntime/test/unittest_util/function_test_util.cc +++ b/onnxruntime/test/unittest_util/function_test_util.cc @@ -43,7 +43,9 @@ ONNX_NAMESPACE::TypeProto FunctionTestCase::TensorType(int32_t elem_type, std::v return typeProto; } -std::vector FunctionTestCase::Run(onnxruntime::Model& model, NameMLValMap& feeds, std::vector output_names) { +std::vector FunctionTestCase::Run(onnxruntime::Model& model, NameMLValMap& feeds, + std::vector output_names, + bool allow_released_opsets_only) { SessionOptions session_options; InferenceSession session_object{session_options, GetEnvironment()}; @@ -51,7 +53,7 @@ std::vector FunctionTestCase::Run(onnxruntime::Model& model, NameMLVal const bool serialization_status = model.ToProto().SerializeToString(&serialized_model); EXPECT_TRUE(serialization_status) << "Failed to serialize proto to string"; std::stringstream sstr(serialized_model); - auto status = session_object.Load(sstr); + auto status = session_object.Load(sstr, allow_released_opsets_only); EXPECT_TRUE(status.IsOK()); status = session_object.Initialize(); EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); @@ -128,7 +130,7 @@ onnxruntime::Node& FunctionTestCase::AddCallNodeTo(onnxruntime::Graph& graph) { return graph.AddNode("fncallnode", opname, "function call node", input_arg_ptrs, output_arg_ptrs, &attributes, domain); } -std::unique_ptr FunctionTestCase::CreateModel(bool inline_call) { +std::unique_ptr FunctionTestCase::CreateModel(bool inline_call, bool allow_released_opsets_only) { if (opsets.size() == 0) { // Default opsets opsets[kOnnxDomain] = 13; @@ -136,7 +138,8 @@ std::unique_ptr FunctionTestCase::CreateModel(bool inline_call) { } std::unique_ptr model(new Model("test", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), - opsets, {}, DefaultLoggingManager().DefaultLogger())); + opsets, {}, DefaultLoggingManager().DefaultLogger(), + ModelOptions{allow_released_opsets_only, false})); onnxruntime::Graph& graph = model->MainGraph(); auto& call_node = AddCallNodeTo(graph); @@ -154,12 +157,12 @@ std::unique_ptr FunctionTestCase::CreateModel(bool inline_call) { return model; } -void FunctionTestCase::RunTest() { - auto model1 = CreateModel(false); - auto results1 = Run(*model1, input_value_map, output_names); +void FunctionTestCase::RunTest(bool allow_released_opsets_only) { + auto model1 = CreateModel(false, allow_released_opsets_only); + auto results1 = Run(*model1, input_value_map, output_names, allow_released_opsets_only); - auto model2 = CreateModel(true); - auto results2 = Run(*model2, input_value_map, output_names); + auto model2 = CreateModel(true, allow_released_opsets_only); + auto results2 = Run(*model2, input_value_map, output_names, allow_released_opsets_only); AssertEqual(results1, results2); } diff --git a/onnxruntime/test/unittest_util/function_test_util.h b/onnxruntime/test/unittest_util/function_test_util.h index a7c9dd9f2d15a..675db4c851a7b 100644 --- a/onnxruntime/test/unittest_util/function_test_util.h +++ b/onnxruntime/test/unittest_util/function_test_util.h @@ -130,11 +130,13 @@ struct FunctionTestCase { onnxruntime::Node& AddCallNodeTo(onnxruntime::Graph& graph); - std::unique_ptr CreateModel(bool inline_call = false); + std::unique_ptr CreateModel(bool inline_call = false, bool allow_released_opsets_only = true); - static std::vector Run(onnxruntime::Model& model, NameMLValMap& feeds, std::vector output_names); + static std::vector Run(onnxruntime::Model& model, NameMLValMap& feeds, + std::vector output_names, + bool allow_released_opsets_only = true); - void RunTest(); + void RunTest(bool allow_released_opsets_only = true); static void AssertEqual(const std::vector& results1, const std::vector& results2); diff --git a/onnxruntime/test/util/include/test_utils.h b/onnxruntime/test/util/include/test_utils.h index f55295ac8aec4..9af8eb92ec5c4 100644 --- a/onnxruntime/test/util/include/test_utils.h +++ b/onnxruntime/test/util/include/test_utils.h @@ -70,7 +70,8 @@ void RunAndVerifyOutputsWithEP(ModelPathOrBytes model_path_or_bytes, const NameMLValMap& feeds, const EPVerificationParams& params = EPVerificationParams(), const std::function& session_options_updater = {}, - bool verify_outputs = true); + bool verify_outputs = true, + bool allow_released_opsets_only = true); // Tests model loading only. // This can be used to test EPs in builds where only loading (and not running) of a model is supported. diff --git a/onnxruntime/test/util/test_utils.cc b/onnxruntime/test/util/test_utils.cc index f6d5d133262c4..132d9f810ec79 100644 --- a/onnxruntime/test/util/test_utils.cc +++ b/onnxruntime/test/util/test_utils.cc @@ -3,6 +3,8 @@ #include "test/util/include/test_utils.h" +#include + #include "core/common/narrow.h" #include "core/common/span_utils.h" #include "core/framework/ort_value.h" @@ -143,7 +145,8 @@ void RunAndVerifyOutputsWithEP(ModelPathOrBytes model_path_or_bytes, std::string const NameMLValMap& feeds, const EPVerificationParams& params, const std::function& session_options_updater, - bool verify_outputs) { + bool verify_outputs, + bool allow_released_opsets_only) { std::vector model_data_buffer{}; const auto model_data = GetModelBytes(model_path_or_bytes, model_data_buffer); @@ -160,7 +163,21 @@ void RunAndVerifyOutputsWithEP(ModelPathOrBytes model_path_or_bytes, std::string // get expected output from CPU EP // InferenceSessionWrapper session_object{so, GetEnvironment()}; - ASSERT_STATUS_OK(session_object.Load(model_data.data(), static_cast(model_data.size()))); + const auto load_model = [&](InferenceSessionWrapper& session) { +#if defined(ORT_MINIMAL_BUILD) + ORT_ENFORCE(allow_released_opsets_only, "Unreleased ONNX opsets require a non-minimal build."); + return session.Load(model_data.data(), static_cast(model_data.size())); +#else + if (allow_released_opsets_only) { + return session.Load(model_data.data(), static_cast(model_data.size())); + } + + const std::string model_data_string(reinterpret_cast(model_data.data()), model_data.size()); + std::istringstream model_stream(model_data_string); + return session.Load(model_stream, false); +#endif + }; + ASSERT_STATUS_OK(load_model(session_object)); ASSERT_STATUS_OK(session_object.Initialize()); const auto& graph = session_object.GetGraph(); @@ -185,7 +202,7 @@ void RunAndVerifyOutputsWithEP(ModelPathOrBytes model_path_or_bytes, std::string // InferenceSessionWrapper session_object2{so, GetEnvironment()}; ASSERT_STATUS_OK(session_object2.RegisterExecutionProvider(std::move(execution_provider))); - ASSERT_STATUS_OK(session_object2.Load(model_data.data(), static_cast(model_data.size()))); + ASSERT_STATUS_OK(load_model(session_object2)); ASSERT_STATUS_OK(session_object2.Initialize()); const auto& graph2 = session_object2.GetGraph(); diff --git a/orttraining/orttraining/core/optimizer/qdq_fusion.cc b/orttraining/orttraining/core/optimizer/qdq_fusion.cc index 6219d1cfafae1..145714e8bca4f 100644 --- a/orttraining/orttraining/core/optimizer/qdq_fusion.cc +++ b/orttraining/orttraining/core/optimizer/qdq_fusion.cc @@ -5,6 +5,7 @@ #include "core/graph/graph_utils.h" #include "core/optimizer/initializer.h" +#include "core/optimizer/utils.h" namespace onnxruntime { @@ -17,7 +18,7 @@ int ReplaceOrCreateZeroPointInitializer(Graph& graph, Node& quantize_node) { // not previously exist. const auto quant_node_input_defs = quantize_node.MutableInputDefs(); - int zero_point_type = ONNX_NAMESPACE::TensorProto_DataType_INT8; + int zero_point_type = ONNX_NAMESPACE::TensorProto_DataType_UINT8; ONNX_NAMESPACE::TensorProto zero_point_tensor_float; if (quant_node_input_defs.size() >= 3) { // The quantize node has the zero point input @@ -41,8 +42,6 @@ int ReplaceOrCreateZeroPointInitializer(Graph& graph, Node& quantize_node) { for (const auto dim : zero_point_tensor_int.dims()) { zero_point_tensor_float.add_dims(dim); } - graph.RemoveInitializedTensor(zero_point_tensor_int.name()); - // Since the quantize node has the zero point initializer input, replace it graph_utils::ReplaceNodeInput(quantize_node, 2, graph_utils::AddInitializerWithOrtValue(graph, zero_point_tensor_float)); @@ -94,7 +93,8 @@ std::pair CheckForQDQPatternMatch(Graph& graph, Node& quantize_node // Try to match the current node with QuantizeLinear in the effort of searching for the pattern // QuantizeLinear -> DequantizeLinear. if (!graph_utils::IsSupportedOptypeVersionAndDomain(quantize_node, "QuantizeLinear", {10, 13, 19}) || - !graph_utils::IsSupportedProvider(quantize_node, compatible_execution_providers)) { + !graph_utils::IsSupportedProvider(quantize_node, compatible_execution_providers) || + graph.NodeProducesGraphOutput(quantize_node)) { return {false, nullptr}; } @@ -117,6 +117,26 @@ std::pair CheckForQDQPatternMatch(Graph& graph, Node& quantize_node "Actual: QuantizeLinear (", quantize_node.Name(), ") is followed by ", dequantize_node_ptr->OpType(), "(", dequantize_node_ptr->Name(), ")."); + const auto& quantize_inputs = quantize_node.InputDefs(); + const auto& dequantize_inputs = dequantize_node_ptr->InputDefs(); + const auto* quantize_axis = graph_utils::GetNodeAttribute(quantize_node, "axis"); + const auto* dequantize_axis = graph_utils::GetNodeAttribute(*dequantize_node_ptr, "axis"); + if (quantize_inputs.size() != dequantize_inputs.size() || + (quantize_axis ? quantize_axis->i() : 1) != (dequantize_axis ? dequantize_axis->i() : 1)) { + return {false, nullptr}; + } + + if (!optimizer_utils::IsScalar(*quantize_inputs[1]) || + (quantize_inputs.size() > 2 && !optimizer_utils::IsScalar(*quantize_inputs[2]))) { + return {false, nullptr}; + } + + for (size_t i = 1; i < quantize_inputs.size(); ++i) { + if (quantize_inputs[i]->Name() != dequantize_inputs[i]->Name()) { + return {false, nullptr}; + } + } + return {true, dequantize_node_ptr}; } @@ -141,8 +161,8 @@ Status QDQFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const } // QuantizeLinear zero_point is INT8 or UINT8. FakeQuant uses quant_zero_point as FLOAT. - // So, remove the old initializers and update the zero point to be of FLOAT type if it exists. - // If the initializer does not exist, create a new zero point initializer with the correct type. + // Add a FLOAT replacement for FakeQuant and retain the original initializer until graph cleanup + // so any other consumers remain valid. Create a zero point initializer when none exists. const auto quant_type = ReplaceOrCreateZeroPointInitializer(graph, quantize_node); // Fuse the QDQ pattern into FakeQuant and move the inputs and outputs. diff --git a/orttraining/orttraining/core/optimizer/scaled_sum_fusion.cc b/orttraining/orttraining/core/optimizer/scaled_sum_fusion.cc index e6319952dfae7..0b48ab2fd7113 100644 --- a/orttraining/orttraining/core/optimizer/scaled_sum_fusion.cc +++ b/orttraining/orttraining/core/optimizer/scaled_sum_fusion.cc @@ -205,7 +205,7 @@ Status ScaledSumFusion::ApplyImpl(Graph& graph, bool& modified, int /*graph_leve } else { // The other input is 1). a constant initializer or graph input, OR 2). it is not a scale operator: // then we only add node arg into data input args, NOT need add any mode into nodes_to_remove. - data_input_args.push_back(mutable_the_other_input_node->MutableInputDefs()[0]); + data_input_args.push_back(the_other_input_arg); scales.push_back(scale_value); } } diff --git a/orttraining/orttraining/core/optimizer/transpose_replacement..cc b/orttraining/orttraining/core/optimizer/transpose_replacement..cc index 48e9c4d6e6a07..d4ea00375ad8e 100644 --- a/orttraining/orttraining/core/optimizer/transpose_replacement..cc +++ b/orttraining/orttraining/core/optimizer/transpose_replacement..cc @@ -61,8 +61,8 @@ Status TransposeReplacement::Apply(Graph& graph, return Status::OK(); } -bool TransposeReplacement::SatisfyCondition(const Graph&, const Node&, const logging::Logger&) const { - return true; +bool TransposeReplacement::SatisfyCondition(const Graph&, const Node& node, const logging::Logger&) const { + return node.GetAttributes().find("perm") != node.GetAttributes().end(); } } // namespace onnxruntime diff --git a/orttraining/orttraining/test/optimizer/graph_transform_test.cc b/orttraining/orttraining/test/optimizer/graph_transform_test.cc index 1278f6e30dd81..191746b48f365 100644 --- a/orttraining/orttraining/test/optimizer/graph_transform_test.cc +++ b/orttraining/orttraining/test/optimizer/graph_transform_test.cc @@ -889,6 +889,27 @@ TEST_F(GraphTransformationTests, TransposeReplacement) { } } +TEST_F(GraphTransformationTests, TransposeReplacementWithoutPerm) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2, 3}}); + auto* output = builder.MakeOutput(); + builder.AddNode("Transpose", {input}, {output}); + }; + + auto check_not_replaced = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.at("Transpose") == 1); + TEST_RETURN_IF(op_count.count("Reshape") != 0); + return Status::OK(); + }; + + auto transformer = std::make_unique("TransposeReplacement"); + ASSERT_STATUS_OK(transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::move(transformer), + TransformerLevel::Level1, 1, + check_not_replaced, check_not_replaced)); +} + TEST_F(GraphTransformationTests, MegatronMLPPartitionRank0) { auto model_uri = MODEL_FOLDER "model_parallel/mlp_megatron_basic_test.onnx"; std::shared_ptr p_model; @@ -1641,7 +1662,7 @@ INSTANTIATE_TEST_SUITE_P( std::make_tuple(2, true, true))); class QDQFusionTestsParameterized : public GraphTransformationTests, - public ::testing::WithParamInterface> { + public ::testing::WithParamInterface> { }; TEST_P(QDQFusionTestsParameterized, CheckModelComposition) { @@ -1662,6 +1683,135 @@ TEST_P(QDQFusionTestsParameterized, CheckModelComposition) { ASSERT_EQ(op_to_count_post_fusion["QuantizeLinear"], 0); ASSERT_EQ(op_to_count_post_fusion["DequantizeLinear"], 0); ASSERT_EQ(op_to_count_post_fusion["com.microsoft.FakeQuant"], 1); + + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "FakeQuant") { + EXPECT_EQ(node.GetAttributes().at("quant_min").i(), std::get<1>(GetParam())); + EXPECT_EQ(node.GetAttributes().at("quant_max").i(), std::get<2>(GetParam())); + } + } +} + +TEST_F(GraphTransformationTests, QDQFusionPreservesPublicQuantizeOutput) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2}}); + auto* scale = builder.MakeScalarInitializer(0.1f); + auto* zero_point = builder.MakeScalarInitializer(0); + auto* quantized = builder.MakeOutput(std::vector{2}); + auto* output = builder.MakeOutput(std::vector{2}); + builder.AddNode("QuantizeLinear", {input, scale, zero_point}, {quantized}); + builder.AddNode("DequantizeLinear", {quantized, scale, zero_point}, {output}); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.at("QuantizeLinear") == 1); + TEST_RETURN_IF_NOT(op_count.at("DequantizeLinear") == 1); + TEST_RETURN_IF(op_count.count("com.microsoft.FakeQuant") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level1, 1, check_not_fused, check_not_fused)); +} + +TEST_F(GraphTransformationTests, QDQFusionPreservesSharedZeroPoint) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2}}); + auto* input_b = builder.MakeInput({{2}}); + auto* scale = builder.MakeScalarInitializer(0.1f); + auto* zero_point = builder.MakeScalarInitializer(0); + auto* quantized_a = builder.MakeIntermediate(); + auto* quantized_b = builder.MakeIntermediate(); + auto* output_a = builder.MakeOutput(std::vector{2}); + auto* output_b = builder.MakeOutput(std::vector{2}); + builder.AddNode("QuantizeLinear", {input_a, scale, zero_point}, {quantized_a}); + builder.AddNode("DequantizeLinear", {quantized_a, scale, zero_point}, {output_a}); + builder.AddNode("QuantizeLinear", {input_b, scale, zero_point}, {quantized_b}); + builder.AddNode("DequantizeLinear", {quantized_b, scale, zero_point}, {output_b}); + }; + + auto check_before = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.at("QuantizeLinear") == 2); + TEST_RETURN_IF_NOT(op_count.at("DequantizeLinear") == 2); + return Status::OK(); + }; + auto check_after = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF(op_count.count("QuantizeLinear") != 0); + TEST_RETURN_IF(op_count.count("DequantizeLinear") != 0); + TEST_RETURN_IF_NOT(op_count.at("com.microsoft.FakeQuant") == 2); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level1, 1, check_before, check_after)); +} + +TEST_F(GraphTransformationTests, QDQFusionRequiresMatchingQuantizationParameters) { + enum class Mismatch { Scale, + ZeroPoint, + Axis }; + + auto check_not_fused = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.at("QuantizeLinear") == 1); + TEST_RETURN_IF_NOT(op_count.at("DequantizeLinear") == 1); + TEST_RETURN_IF(op_count.count("com.microsoft.FakeQuant") != 0); + return Status::OK(); + }; + + for (const auto mismatch : {Mismatch::Scale, Mismatch::ZeroPoint, Mismatch::Axis}) { + auto build_test_case = [mismatch](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2}}); + auto* quantize_scale = builder.MakeScalarInitializer(0.1f); + auto* dequantize_scale = mismatch == Mismatch::Scale + ? builder.MakeScalarInitializer(0.2f) + : quantize_scale; + auto* quantize_zero_point = builder.MakeScalarInitializer(0); + auto* dequantize_zero_point = mismatch == Mismatch::ZeroPoint + ? builder.MakeScalarInitializer(1) + : quantize_zero_point; + auto* quantized = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(std::vector{2}); + auto& quantize = builder.AddNode( + "QuantizeLinear", {input, quantize_scale, quantize_zero_point}, {quantized}); + builder.AddNode( + "DequantizeLinear", {quantized, dequantize_scale, dequantize_zero_point}, {output}); + if (mismatch == Mismatch::Axis) { + quantize.AddAttribute("axis", static_cast(0)); + } + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level1, 1, check_not_fused, check_not_fused)); + } +} + +TEST_F(GraphTransformationTests, QDQFusionRejectsPerAxisParameters) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({{2}}); + auto* scale = builder.MakeInitializer({2}, {0.1f, 0.2f}); + auto* zero_point = builder.MakeInitializer({2}, {0, 0}); + auto* quantized = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(std::vector{2}); + auto& quantize = builder.AddNode("QuantizeLinear", {input, scale, zero_point}, {quantized}); + auto& dequantize = builder.AddNode("DequantizeLinear", {quantized, scale, zero_point}, {output}); + quantize.AddAttribute("axis", static_cast(0)); + dequantize.AddAttribute("axis", static_cast(0)); + }; + + auto check_not_fused = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.at("QuantizeLinear") == 1); + TEST_RETURN_IF_NOT(op_count.at("DequantizeLinear") == 1); + TEST_RETURN_IF(op_count.count("com.microsoft.FakeQuant") != 0); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level1, 1, check_not_fused, check_not_fused)); } TEST_F(GraphTransformationTests, Conv1dReplacement_TakeEffect) { @@ -1813,9 +1963,9 @@ INSTANTIATE_TEST_SUITE_P( QDQFusionTests, QDQFusionTestsParameterized, ::testing::Values( - std::make_tuple(MODEL_FOLDER "fusion/qdq_fusion_int8.onnx"), - std::make_tuple(MODEL_FOLDER "fusion/qdq_fusion_uint8.onnx"), - std::make_tuple(MODEL_FOLDER "fusion/qdq_fusion_zp_not_provided.onnx"))); + std::make_tuple(MODEL_FOLDER "fusion/qdq_fusion_int8.onnx", -128, 127), + std::make_tuple(MODEL_FOLDER "fusion/qdq_fusion_uint8.onnx", 0, 255), + std::make_tuple(MODEL_FOLDER "fusion/qdq_fusion_zp_not_provided.onnx", 0, 255))); // We only tested on CUDA run. #if defined(USE_CUDA) @@ -2092,6 +2242,9 @@ TEST_F(GraphTransformationTests, ScaledSumFusionThreeInputs_LastAddNotHaveScaleI TEST_RETURN_IF_NOT(1.0f / 0.5f == attrs.at("scale_0").f()); TEST_RETURN_IF_NOT(1.0f / 0.3f == attrs.at("scale_1").f()); TEST_RETURN_IF_NOT(1.0f == attrs.at("scale_2").f()); + const Node* third_input_producer = graph.GetProducerNode(node.InputDefs()[2]->Name()); + TEST_RETURN_IF_NOT(third_input_producer != nullptr); + TEST_RETURN_IF_NOT(third_input_producer->OpType() == "Sub"); } } @@ -2138,6 +2291,47 @@ TEST_F(GraphTransformationTests, ScaledSumFusionThreeInputs_LastAddNotHaveScaleI } } +TEST_F(GraphTransformationTests, ScaledSumFusionWithGraphInputThirdTerm) { + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input_a = builder.MakeInput({{2, 3}}); + auto* input_b = builder.MakeInput({{2, 3}}); + auto* input_c = builder.MakeInput({{2, 3}}); + auto* scale_a = builder.MakeScalarInitializer(0.5f); + auto* scale_b = builder.MakeScalarInitializer(0.25f); + auto* scaled_a = builder.MakeIntermediate(); + auto* scaled_b = builder.MakeIntermediate(); + auto* sum_ab = builder.MakeIntermediate(); + auto* sum_abc = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("Div", {input_a, scale_a}, {scaled_a}); + builder.AddNode("Div", {input_b, scale_b}, {scaled_b}); + builder.AddNode("Add", {scaled_a, scaled_b}, {sum_ab}); + builder.AddNode("Add", {sum_ab, input_c}, {sum_abc}); + builder.AddNode("Identity", {sum_abc}, {output}); + }; + + auto check_before = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.at("Div") == 2); + TEST_RETURN_IF_NOT(op_count.at("Add") == 2); + return Status::OK(); + }; + auto check_after = [](Graph& graph) { + const auto op_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count.at("com.microsoft.ScaledSum") == 1); + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "ScaledSum") { + TEST_RETURN_IF_NOT(node.InputDefs().size() == 3U); + TEST_RETURN_IF(graph.GetProducerNode(node.InputDefs()[2]->Name()) != nullptr); + } + } + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 13, *logger_, std::make_unique(), + TransformerLevel::Level1, 1, check_before, check_after)); +} + /* Test graph as below. graph input [1, 1, 256, 256] (float) scalar_0 graph input [1, 1, 256, 256] (float) diff --git a/plugin-ep-webgpu/README.md b/plugin-ep-webgpu/README.md index c26d9207d5df0..833054bc399aa 100644 --- a/plugin-ep-webgpu/README.md +++ b/plugin-ep-webgpu/README.md @@ -32,7 +32,10 @@ build (`--use_webgpu shared_lib`). The resulting binaries are then packaged into - Per-platform zip packages for Foundry Local consumption. On Windows, the packages also bundle the DirectX Shader Compiler runtime (`dxcompiler.dll`) from the -[DXC GitHub releases](https://github.com/microsoft/DirectXShaderCompiler/releases); CI fetches these automatically. +[DXC GitHub releases](https://github.com/microsoft/DirectXShaderCompiler/releases); CI fetches it automatically. +The Windows Python wheels, NuGet package, and Foundry Local zip packages produced by CI also bundle the pinned preview +D3D12 Agility SDK runtime. See the [Agility SDK guidance](../docs/BuildWithDawnAgilitySDK.md) for Windows, driver, and +Developer Mode prerequisites. Packaging is driven by the *WebGPU Plugin EP Packaging Pipeline* ([`plugin-webgpu-pipeline.yml`](../tools/ci_build/github/azure-pipelines/plugin-webgpu-pipeline.yml)), diff --git a/plugin-ep-webgpu/_packaging_utils.py b/plugin-ep-webgpu/_packaging_utils.py index 84850e4dee5fe..ceeb48ab10087 100644 --- a/plugin-ep-webgpu/_packaging_utils.py +++ b/plugin-ep-webgpu/_packaging_utils.py @@ -11,6 +11,31 @@ # Matches "@var@" template variables. _TEMPLATE_VARIABLE_PATTERN = re.compile(r"@(\w+)@") +PLUGIN_BINARIES = { + "Windows": (Path("onnxruntime_providers_webgpu.dll"), Path("dxcompiler.dll")), + "Linux": (Path("libonnxruntime_providers_webgpu.so"),), + "Darwin": (Path("libonnxruntime_providers_webgpu.dylib"),), +} + +WINDOWS_AGILITY_SDK_BINARIES = (Path("D3D12/D3D12Core.dll"), Path("D3D12/d3d12SDKLayers.dll")) + + +def get_plugin_binaries(target_platform: str, *, include_agility_sdk_files: bool = False) -> tuple[Path, ...]: + """Return the required plugin binaries for a target platform. + + Returned paths are relative to the binary directory provided by the caller. + `target_platform` uses `platform.system()` names such as "Windows", "Linux", and "Darwin". + """ + try: + binaries = list(PLUGIN_BINARIES[target_platform]) + except KeyError as e: + raise ValueError(f"Unsupported WebGPU packaging platform: {target_platform}") from e + + if target_platform == "Windows" and include_agility_sdk_files: + binaries.extend(WINDOWS_AGILITY_SDK_BINARIES) + + return tuple(binaries) + def gen_file_from_template( template_file: Path, output_file: Path, variable_substitutions: dict[str, str], strict: bool = True diff --git a/plugin-ep-webgpu/csharp/Microsoft.ML.OnnxRuntime.EP.WebGpu/Microsoft.ML.OnnxRuntime.EP.WebGpu.csproj b/plugin-ep-webgpu/csharp/Microsoft.ML.OnnxRuntime.EP.WebGpu/Microsoft.ML.OnnxRuntime.EP.WebGpu.csproj index eadda6cb8482f..8ff6222d336ea 100644 --- a/plugin-ep-webgpu/csharp/Microsoft.ML.OnnxRuntime.EP.WebGpu/Microsoft.ML.OnnxRuntime.EP.WebGpu.csproj +++ b/plugin-ep-webgpu/csharp/Microsoft.ML.OnnxRuntime.EP.WebGpu/Microsoft.ML.OnnxRuntime.EP.WebGpu.csproj @@ -33,6 +33,9 @@ +