From 4c184ac48295b9e64bf83a344c2d46bd7030f6fa Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 20 Oct 2026 15:57:38 -0700 Subject: [PATCH 01/17] Add opset 28 SpaceToDepth support (#32362) ### Summary - Add native ONNX opset-28 `SpaceToDepth` support for both `DCR` and `CRD` modes on CPU and CUDA, including internal NHWC kernels. - Version-split `SpaceToDepth` and `DepthToSpace` registrations at opset 28 across CPU, CUDA, DML, JS, and WebGPU; add DML ordered `SpaceToDepth` support and make QNN decline unsupported CRD nodes for fallback. - Add opset-28 CPU mode, shape-validation, non-square/blocksize, and native-vs-context-dependent-function parity coverage. ### Validation - `git diff --check` - Source-level DCR/CRD layout checks against the pinned ONNX 1.23 function bodies. - `build.py --config Debug --build --build_dir build\\windows-spacedepth --cmake_generator "Visual Studio 17 2022" --target onnxruntime_provider_test` is blocked by the base branch's unrelated FP6 `TensorProto` datatype-map static assertion in `tensorprotoutils.h`, before producing a test binary. --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Co-authored-by: Ganesan Ramalingam --- docs/OperatorKernels.md | 24 ++- .../contrib_ops/internal_nhwc_onnx_schemas.cc | 2 + .../providers/cpu/cpu_execution_provider.cc | 12 +- .../providers/cpu/tensor/space_depth_ops.cc | 41 ++-- .../providers/cpu/tensor/space_depth_ops.h | 12 +- .../providers/cuda/cuda_execution_provider.cc | 12 +- .../core/providers/cuda/cuda_nhwc_kernels.cc | 8 +- .../providers/cuda/tensor/space_depth_ops.cc | 82 ++++++-- .../providers/cuda/tensor/space_depth_ops.h | 5 +- .../src/Operators/OperatorRegistration.cpp | 2 + .../OperatorAuthorHelper/OperatorVersions.h | 2 + .../providers/webgpu/tensor/depth_to_space.cc | 6 +- .../webgpu/webgpu_execution_provider.cc | 9 +- onnxruntime/core/session/environment.cc | 12 +- .../test/contrib_ops/function_ops_test.cc | 20 ++ .../cpu/tensor/space_depth_ops_test.cc | 181 +++++++++++++++++- .../test/providers/qnn/qnn_test_utils.cc | 8 +- .../test/providers/qnn/qnn_test_utils.h | 3 +- .../test/providers/qnn/simple_op_htp_test.cc | 23 ++- .../test/unittest_util/function_test_util.cc | 21 +- .../test/unittest_util/function_test_util.h | 8 +- onnxruntime/test/util/include/test_utils.h | 3 +- onnxruntime/test/util/test_utils.cc | 23 ++- 23 files changed, 434 insertions(+), 85 deletions(-) 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/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, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo()}), 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..d341152f18930 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -1353,7 +1353,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 +2680,12 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, 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/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/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/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/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/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(); From 786177ab6f07721f928a54d656afd5bcb6d9a0c7 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane <111780983+apsonawane@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:25:21 -0700 Subject: [PATCH 02/17] cuda pagedattention seqlens shape fix (#31651) This pull request introduces improved validation for input tensors in the Paged Attention CUDA operator, ensuring that sequence lengths and block table values are properly checked before computation. It adds a new comprehensive validation function, integrates it into the main compute path, and provides unit tests to verify correctness. Additionally, a minor bug in shape checking is fixed. **Input validation improvements:** * Added a new function `CheckBlockTableAndPastSeqLensValues` in `paged_attention_helper.h` to validate that `cumulative_sequence_length`, `past_seqlens`, and `block_table` values are within expected ranges, preventing invalid memory accesses and logical errors. * Integrated this validation into the main compute function in `paged_attention.cc`, copying relevant device data to CPU and running the checks before computation proceeds. * Fixed a bug in `CheckSequenceLengthTensors` where the logical condition for validating the shape of `seqlens` was incorrect (`&&` replaced with `||`). **Testing:** * Added a new test file `paged_attention_helper_test.cc` with unit tests covering valid and invalid cases for the new validation logic, ensuring robustness against edge cases. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 4 +- docs/contrib_ops/cuda/paged_attention.md | 36 +- .../cpu/bert/paged_attention_helper.h | 127 +++++ .../contrib_ops/cuda/bert/attention_data.h | 2 + .../bert/flash_attention/flash_fwd_kernel.h | 61 ++- .../contrib_ops/cuda/bert/paged_attention.cc | 201 +++++--- .../cuda/bert/paged_attention_impl.cu | 433 +++++++++++++----- .../cuda/bert/paged_attention_impl.h | 37 +- .../core/graph/contrib_ops/bert_defs.cc | 9 +- .../paged_attention_helper_test.cc | 356 ++++++++++++++ .../contrib_ops/paged_attention_op_test.cc | 353 ++++++++++++-- 11 files changed, 1357 insertions(+), 262 deletions(-) create mode 100644 onnxruntime/test/contrib_ops/cuda_kernels/paged_attention_helper_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 2703de7cc1fae..603118adeeed8 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -5464,9 +5464,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/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/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/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/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/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index fecac2f17f523..ffdb88d239ec3 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", 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/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(); From a243f8cb323ca90b929d61ae820a78b70422bd54 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane <111780983+apsonawane@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:32:12 -0700 Subject: [PATCH 03/17] Validate PackedAttention token offsets (#33330) This pull request introduces validation for the `token_offset` input in the CUDA implementation of the PackedAttention operator, specifically for the unfused backend. The main goal is to ensure that invalid `token_offset` values are detected early and reported with a clear error message, improving robustness and debuggability. The change includes both the implementation of the validation logic and new tests to verify its correctness. **Validation improvements for PackedAttention:** * Added a CUDA kernel and a host function (`ValidatePackedAttentionTokenOffset`) to check that all `token_offset` values are within the valid range `[0, B * S)` for the unfused backend, and integrated this validation into `PackedAttention::ComputeInternal`. [[1]](diffhunk://#diff-4e4c1980e908e7c273f52041c3dc123cb7a8c4ee20779fb28b4975911556c9dbR31-R84) [[2]](diffhunk://#diff-450dced704ee2fc03ee3611224808eb8a00671f7587e48599ce17c21d5e7c9e8R16-R21) [[3]](diffhunk://#diff-fd7731431d1e80901e297b2ce7f1018fc7305faf58145cb36733bd19b0aeb9d2R270-R279) * Updated relevant CUDA kernels to handle out-of-range `token_offset` values gracefully by skipping computation or zeroing out invalid output regions. [[1]](diffhunk://#diff-4e4c1980e908e7c273f52041c3dc123cb7a8c4ee20779fb28b4975911556c9dbR115-R117) [[2]](diffhunk://#diff-4e4c1980e908e7c273f52041c3dc123cb7a8c4ee20779fb28b4975911556c9dbR371-R374) **Testing enhancements:** * Added a new unit test (`InvalidTokenOffset_Unfused_CUDA`) to verify that invalid `token_offset` values for the unfused backend are correctly detected and result in a descriptive error message. --- .../contrib_ops/cuda/bert/packed_attention.cc | 10 + .../cuda/bert/packed_attention_impl.cu | 61 +++++ .../cuda/bert/packed_attention_impl.h | 6 + .../contrib_ops/packed_attention_op_test.cc | 208 ++++++++++++++++++ 4 files changed, 285 insertions(+) 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/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, From 942356fce6bdc78c016e4b902e49492fdf695453 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 20 Oct 2026 16:59:48 -0700 Subject: [PATCH 04/17] [CUDA] Add fused decode routing, single-pass finalize, and PDL overlap for NVFP4 QMoE (#33366) > **Stacked on #33354 (and #33094).** The base branch is `tlwu/20261013/nvfp4-qmoe-decode-gemv`, so the diff below is this change only. Once #33094 and #33354 merge, this PR will be retargeted to `main`. Opened as a draft until then. ## Description Optimizes the non-GEMV kernels and inter-kernel launch serialization along the offline K-packed (`weights_prepacked=1`) NVFP4 `QMoE` decode path on SM90+: 1. **Single-CTA fused decode routing (`FusedSoftmaxTopKMapsKernel`)**: Replaces `LaunchSoftmaxTopK` + `fusedBuildExpertMapsSortFirstToken` in a single CTA launch for `num_rows <= 4`, `num_experts == 512`, and `k in {1, 2, 4, 6, 8, 10}`. One warp per row reproduces the `SoftmaxTopKMergeKernel` reduction trees operation for operation and emits the `(expert, token)`-sorted expert maps bit-identically (`ORT_FP4_GEMV_FUSED_ROUTE=0` disables). 2. **Single-pass `finalizeMoeRoutingKernel`**: Sizes the block to one thread per 128-bit column vector (up to 1024 threads instead of a fixed 256) so hidden dimensions up to 4096 (such as Qwen3.8's `hidden=2560`) reduce in a single pass, and issues batched loads of up to 16 expert rows together while preserving top-k accumulation order (`7.7 us -> 4.0 us` per layer on H200). 3. **Programmatic dependent launch (`PDL`) overlap of `FC1 -> FC2 -> finalize`**: Launches FC2 and `finalizeMoeRoutingKernel` with `cudaLaunchAttributeProgrammaticStreamSerialization` (`ORT_FP4_GEMV_PDL=0` disables; `ORT_FP4_GEMV_FC1_PDL=1` optionally enables PDL for FC1 behind fused routing). FC1 triggers dependents at start, FC2 prefetches its first batch of weights/scales before waiting on FC1's output, and `finalizeMoeRoutingKernel` stages routing metadata into shared memory before waiting on FC2's output (`~2.2 us` saved per layer on H200, preserved under CUDA Graph capture). All three optimizations preserve exact floating-point evaluation order and produce **bit-identical** outputs. ## Summary of Changes | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuh` | New `FusedSoftmaxTopKMapsKernel` single-CTA fused softmax + top-k + expert-map kernel (`E=512`, `rows <= 4`, `k in {1,2,4,6,8,10}`) | | `onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h` / `.cu` | `LaunchFusedRouteMaps` declarations and launcher | | `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu` / `moe_util_kernels.h` | Single-pass `finalizeMoeRoutingKernel` (`blockDim.x` up to 1024, `FINALIZE_LOAD_BATCH=16`) and `early_routing_loads` PDL support | | `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h` / `.cu` | `griddepcontrol.wait` / `griddepcontrol.launch_dependents` in `MoeGemvFp4DecodeKernel` and `programmatic_dependent_launch` attribute in `LaunchMoeGemvFp4RawNPacked` (`ORT_FP4_GEMV_PDL`) | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` | Wire `LaunchFusedRouteMaps` and `fc2_finalize_pdl` into `QMoE::ComputeInternal` | | `docs/contrib_ops/cuda/moe_qmoe.md` | Document `ORT_FP4_GEMV_FUSED_ROUTE`, `ORT_FP4_GEMV_PDL`, and `ORT_FP4_GEMV_FC1_PDL` | ## Testing Validated on NVIDIA H200 (`sm_90a`, CUDA 13.0): - `python -m pytest onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py`: **178 passed, 8 skipped, 0 failed** (including `test_nvfp4_offline_row_major_decode_gemv_bit_exact` for FP16 and BF16). - End-to-end Qwen3.8-Flash-Next NVFP4 (`v12b` 3-session CUDA-graph + MTP engine, 48 MoE layers + 1 MTP layer): - **Bit-exactness**: `engine_gate.py --prefill steps` passes **6/6** with 100% identical token sequences (`256/256` tokens on `text2048` and `prose2048`) compared against the baseline without fused routing / PDL. - **End-to-end decode speed (`bench_engine_mtp.py`, 256 generated tokens, median of 3 runs)**: | Context | Before (`#33354` only) `ms/cycle` | With this PR `ms/cycle` | Cycle speedup | Decode throughput (`tok/s`) | |---:|---:|---:|---:|---:| | 2,048 | 8.274 ms | **7.657 ms** | **-0.617 ms (-7.5%)** | **244.98 tok/s** (was 227.64) | | 8,192 | 8.164 ms | **7.836 ms** | **-0.328 ms (-4.0%)** | **232.49 tok/s** (was 225.81) | | 32,768 | 8.352 ms | **8.027 ms** | **-0.325 ms (-3.9%)** | **237.25 tok/s** (was 229.97) | - `lintrunner` on all touched files: no issues. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/moe_qmoe.md | 3 + .../cuda/llm/moe_gemm/moe_gemv_fp4.cu | 143 +++++++++-- .../cuda/llm/moe_gemm/moe_gemv_fp4.h | 10 +- .../cuda/llm/moe_gemm/moe_kernels.cu | 188 +++++++++++---- .../cuda/llm/moe_gemm/moe_util_kernels.h | 6 +- .../contrib_ops/cuda/moe/moe_quantization.cc | 90 ++++--- .../contrib_ops/cuda/moe/qmoe_fused_route.cuh | 223 ++++++++++++++++++ .../contrib_ops/cuda/moe/qmoe_kernels.cu | 88 +++++++ .../contrib_ops/cuda/moe/qmoe_kernels.h | 11 + .../benchmark_qwen38_mtp_kernels.py | 21 +- .../test_benchmark_qwen38_mtp_kernels.py | 48 ++++ .../transformers/test_qmoe_nvfp4_cuda.py | 97 +++++++- 12 files changed, 822 insertions(+), 106 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuh 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/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/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, From 53e61f3e4ff833be5bdfe9993b3acda43f2fceea Mon Sep 17 00:00:00 2001 From: Mustapha Jaber Date: Tue, 20 Oct 2026 17:53:35 -0700 Subject: [PATCH 05/17] Update transformers and FlatBuffers Go example grpc dependency pins (#33501) ## Description - Update `transformers` requirement pins from `5.3.0` to `5.10.1` across the Python transformer tooling and test `requirements.txt` manifests. - Update `google.golang.org/grpc` to `v1.84.0` in the FlatBuffers Go greeter example `go.mod` files (`grpc/examples/go/greeter/{client,models,server}/go.mod`) via `cmake/patches/flatbuffers/flatbuffers.patch`. ## Motivation and Context Keep Python `transformers` and third-party Go module dependency manifests up to date with recent patched releases. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- cmake/patches/flatbuffers/flatbuffers.patch | 21 +++++++++++++++++++ .../models/llama/requirements.txt | 2 +- .../transformers/models/phi2/requirements.txt | 2 +- .../requirements/requirements.txt | 2 +- .../models/whisper/requirements.txt | 2 +- .../tools/transformers/requirements.txt | 2 +- .../docker/scripts/lort/requirements.txt | 2 +- .../transformers-test/requirements.txt | 2 +- 8 files changed, 28 insertions(+), 7 deletions(-) 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/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/tools/ci_build/github/linux/docker/scripts/lort/requirements.txt b/tools/ci_build/github/linux/docker/scripts/lort/requirements.txt index 314aaa3f6a541..fcf19ded1e9a9 100644 --- a/tools/ci_build/github/linux/docker/scripts/lort/requirements.txt +++ b/tools/ci_build/github/linux/docker/scripts/lort/requirements.txt @@ -1,4 +1,4 @@ -transformers>=5.3.0 +transformers>=5.10.1 beartype==0.15.0 flatbuffers cerberus diff --git a/tools/ci_build/requirements/transformers-test/requirements.txt b/tools/ci_build/requirements/transformers-test/requirements.txt index cc6fcc1d923c4..9d096c235e592 100644 --- a/tools/ci_build/requirements/transformers-test/requirements.txt +++ b/tools/ci_build/requirements/transformers-test/requirements.txt @@ -6,7 +6,7 @@ numpy==2.2.6; python_version < "3.11" numpy==2.4.2; python_version >= "3.11" torch==2.10.0 torchvision==0.25.0 -transformers==5.3.0 +transformers==5.10.1 parameterized>=0.8.1 sentencepiece psutil From fb8bfaaade9c517be99194ab24dafa38e8dab16c Mon Sep 17 00:00:00 2001 From: May Mei Date: Wed, 21 Oct 2026 09:29:38 +0800 Subject: [PATCH 06/17] [WebGPU] Package D3D12 Agility SDK for Python wheels (#32795) This commit packages the D3D12 Agility SDK runtime DLLs in both the standard ONNX Runtime Python wheel with the built-in WebGPU EP and the standalone WebGPU Plugin EP Python wheel. --- .github/workflows/windows_webgpu.yml | 7 + .../external/onnxruntime_external_deps.cmake | 7 - cmake/onnxruntime_providers_webgpu.cmake | 4 + cmake/onnxruntime_python.cmake | 13 + docs/BuildWithDawnAgilitySDK.md | 46 ++- plugin-ep-webgpu/README.md | 5 +- plugin-ep-webgpu/_packaging_utils.py | 25 ++ .../Microsoft.ML.OnnxRuntime.EP.WebGpu.csproj | 3 + ...Microsoft.ML.OnnxRuntime.EP.WebGpu.targets | 33 ++ .../README.md | 12 + plugin-ep-webgpu/csharp/README.md | 27 ++ plugin-ep-webgpu/csharp/pack_nuget.py | 101 ++++-- .../csharp/test/test_pack_nuget.py | 312 ++++++++++++++++++ plugin-ep-webgpu/python/README.md | 35 ++ plugin-ep-webgpu/python/build_wheel.py | 59 ++-- .../python/onnxruntime_ep_webgpu/README.md | 12 + plugin-ep-webgpu/python/pyproject.toml.in | 2 +- .../python/test/test_webgpu_plugin_ep.py | 16 +- setup.py | 2 +- tools/ci_build/build.py | 18 +- tools/ci_build/build_args.py | 7 +- .../plugin-webgpu-nuget-packaging-stage.yml | 7 + .../stages/plugin-webgpu-packaging-stage.yml | 12 +- .../stages/plugin-win-webgpu-stage.yml | 3 + .../stages/plugin-win-webgpu-test-stage.yml | 1 + 25 files changed, 669 insertions(+), 100 deletions(-) create mode 100644 plugin-ep-webgpu/csharp/Microsoft.ML.OnnxRuntime.EP.WebGpu/Microsoft.ML.OnnxRuntime.EP.WebGpu.targets create mode 100644 plugin-ep-webgpu/csharp/test/test_pack_nuget.py diff --git a/.github/workflows/windows_webgpu.yml b/.github/workflows/windows_webgpu.yml index e2457ad51e2c0..7fb1ae7c60f15 100644 --- a/.github/workflows/windows_webgpu.yml +++ b/.github/workflows/windows_webgpu.yml @@ -99,6 +99,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: 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/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/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 @@ + ### Motivation and Context --- .github/workflows/android.yml | 18 +----- .github/workflows/cffconvert.yml | 6 +- .github/workflows/codeql.yml | 6 +- .../workflows/gradle-wrapper-validation.yml | 6 +- .github/workflows/labeler.yml | 6 +- .github/workflows/lint.yml | 6 +- .../linux-wasm-ci-build-and-test-workflow.yml | 6 +- .github/workflows/linux_ci.yml | 14 ++--- .github/workflows/linux_cuda_ci.yml | 7 +-- .github/workflows/linux_cuda_no_cudnn.yml | 7 +-- .github/workflows/linux_cuda_plugin_ci.yml | 7 +-- .github/workflows/linux_minimal_build.yml | 60 ++++--------------- .github/workflows/linux_openvino_ci.yml | 2 +- .github/workflows/linux_tensorrt_ci.yml | 9 +-- .github/workflows/linux_webgpu.yml | 4 +- .github/workflows/pr_checks.yml | 6 +- .github/workflows/publish-c-apidocs.yml | 6 +- .github/workflows/publish-csharp-apidocs.yml | 6 +- .github/workflows/publish-gh-pages.yml | 6 +- .github/workflows/publish-java-apidocs.yml | 6 +- .github/workflows/publish-js-apidocs.yml | 6 +- .github/workflows/publish-python-apidocs.yml | 6 +- .github/workflows/react_native.yml | 12 +--- .github/workflows/reusable_linux_build.yml | 7 +-- .github/workflows/title-only-labeler.yml | 6 +- .github/workflows/windows-web-ci-workflow.yml | 6 +- .github/workflows/windows_arm64_python.yml | 6 +- .github/workflows/windows_build_x64_asan.yml | 6 +- .github/workflows/windows_cuda.yml | 12 +--- .github/workflows/windows_cuda_no_cudnn.yml | 12 +--- .github/workflows/windows_cuda_plugin.yml | 12 +--- .github/workflows/windows_dml.yml | 6 +- .github/workflows/windows_gpu_doc_gen.yml | 6 +- .github/workflows/windows_openvino.yml | 6 +- .github/workflows/windows_qnn_x64.yml | 6 +- .github/workflows/windows_tensorrt.yml | 12 +--- .github/workflows/windows_webgpu.yml | 48 +++------------ .../windows_x64_debug_build_x64_debug.yml | 6 +- .../windows_x64_release_build_x64_release.yml | 6 +- ...build_x64_release_ep_generic_interface.yml | 6 +- ..._x64_release_vitisai_build_x64_release.yml | 6 +- .../workflows/windows_x64_release_xnnpack.yml | 6 +- 42 files changed, 78 insertions(+), 321 deletions(-) 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 7fb1ae7c60f15..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 @@ -174,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 @@ -296,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 @@ -388,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 @@ -453,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 }} @@ -522,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 @@ -607,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 @@ -727,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: From b923407cf40455dfca0b72bae63e369c004eb92e Mon Sep 17 00:00:00 2001 From: Akshay Sonawane <111780983+apsonawane@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:39:13 -0700 Subject: [PATCH 09/17] Validate ORT format edge slots (#32200) This pull request enhances the robustness and correctness of ONNX Runtime's ORT format model loading, particularly around edge validation in computational graphs. It introduces stricter checks for edge consistency, improves error handling, and significantly expands test coverage to prevent invalid graph structures from being loaded. **Edge validation and canonicalization improvements:** * The `Node::LoadEdgesFromOrtFormat` method was refactored to enforce stricter validation of edge indices, ensure control edges are handled correctly, and guarantee that edges only connect matching `NodeArgs`. This includes checks for out-of-range indices, mismatched arguments, and correct handling of control edges. (`onnxruntime/core/graph/graph.cc`, `include/onnxruntime/core/graph/graph.h`) [[1]](diffhunk://#diff-e231a92b40d89409cc8e82436be0a15bc87ef95c93b303b9feaeab6e50c8835cL892-R955) [[2]](diffhunk://#diff-aaea1507ec81a94c72a1fa72ce320df712156b665f7798573be3f7e439bb4c37L523-R523) * After edge loading, the graph now checks that each destination argument slot has at most one producer, preventing invalid graphs with multiple producers for the same input. (`onnxruntime/core/graph/graph.cc`) **Test coverage expansion:** * New helper functions and tests were added to rigorously verify edge validation logic. These tests cover scenarios such as out-of-range source/destination indices, mismatched node arguments, canonicalization of reciprocal edges, control edge handling, and detection of multiple producers for a single destination slot. (`onnxruntime/test/framework/ort_model_only_test.cc`) [[1]](diffhunk://#diff-091729c3958c3945dfe649e7a1c3313b941a3da9a27662f047cc35c948196e49R90-R160) [[2]](diffhunk://#diff-091729c3958c3945dfe649e7a1c3313b941a3da9a27662f047cc35c948196e49R302-R369) These changes collectively ensure that only valid and well-formed graphs are accepted from ORT format models, reducing the risk of subtle bugs or crashes due to malformed edge definitions. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/graph/graph.cc | 43 ++++++++++++--- .../test/framework/ort_model_only_test.cc | 52 +++++++++++++++++++ 2 files changed, 87 insertions(+), 8 deletions(-) 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/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{ From c55286e574e728117542787ab674db1f52aab290 Mon Sep 17 00:00:00 2001 From: Yuhx141 <155034285+Yuhx141@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:45:56 +0800 Subject: [PATCH 10/17] Fix correctness guards in training graph optimizers (#32982) ### Description This change adds missing semantic guards in three training graph optimizers: - QDQFusion uses the ONNX UINT8 default when zero point is omitted, leaves shared initializers for graph cleanup, skips a QuantizeLinear whose output is public, and only fuses when the Q/DQ scale, zero-point inputs, and effective axis match. - TransposeReplacement skips the optional default-permutation form instead of indexing a missing attribute. - ScaledSumFusion keeps the original unscaled third input, including graph inputs and outputs of non-scale producers. Regression tests cover each case and retain the existing successful fusion paths. ### Motivation and Context The affected rewrites could change valid model semantics or fail while optimizing valid graphs. Fixes #32966 Fixes #32967 Fixes #32968 Fixes #32969 Fixes #32977 Fixes #32980 ### Testing Built `onnxruntime_test_all` with `--enable_training` in Release mode. After the Q/DQ parameter guard was added, 10 focused tests from two suites passed, covering the new regressions and neighboring successful fusion paths. --------- Signed-off-by: Yuhx141 <155034285+Yuhx141@users.noreply.github.com> Co-authored-by: xadupre Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../orttraining/core/optimizer/qdq_fusion.cc | 32 ++- .../core/optimizer/scaled_sum_fusion.cc | 2 +- .../core/optimizer/transpose_replacement..cc | 4 +- .../test/optimizer/graph_transform_test.cc | 202 +++++++++++++++++- 4 files changed, 227 insertions(+), 13 deletions(-) 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) From 314ff8d4e87816bafa1b8e0fefa476c896c4dedd Mon Sep 17 00:00:00 2001 From: Chi Lo <54722500+chilo-ms@users.noreply.github.com> Date: Wed, 21 Oct 2026 08:54:41 -0700 Subject: [PATCH 11/17] [Framework/CUDA] Add generic kernel dispatch policy reader (#33042) ### Description Add reader-only plumbing for the hardware-neutral session option `session.kernel_dispatch_policy`. The intent is generic across operators and EPs. `KernelDispatchPolicy` and the shared, header-only `TryParseKernelDispatchPolicy` live in `core/framework/kernel_dispatch_policy.h`, with no CUDA or graph dependencies. CUDA GQA is the first consumer, not the owner of the policy. For valid values, **no policy changes runtime dispatch, workspace estimates, or partitioning** in this PR. `latency`, `memory`, and `safe` remain intent metadata for future policy resolution. In particular, `memory` does not reduce the reservation, and `safe` does not select a bounded fallback or provide a no-OOM guarantee. #### Reader contract - Unset, empty, and explicit `auto` resolve to `Auto`. The other valid tokens are exactly `latency`, `memory`, and `safe`. - Unknown, differently cased, whitespace-only, and whitespace-padded non-empty values fail session initialization with `INVALID_ARGUMENT`; they never silently resolve to `Auto`. Validation runs even when partitioning is disabled. - The accountant factory retains the validated raw optional string without changing other settings. The CUDA consumer enforces successful parsing rather than silently accepting an invalid configuration. - `GetGroupQueryAttentionWorkspaceEstimateConfig` accepts `dispatch_policy = KernelDispatchPolicy::Auto` alongside the envelope and initializes both fields in one place. The node-estimator overload and CUDA test bridge forward the policy through that helper. - Setting the policy alone does not enable partitioning or non-windowed GQA estimation without a KV envelope. - This replaces the earlier, unreleased attention-specific option name; there is no compatibility alias. #### Integration and regression coverage Preserve #33014/#33016's envelope validation against `[0, INT32_MAX]`, invalid-value diagnostics, and fail-closed GQA placement when no proven estimate is available. Reuse the exported graph-config helper instead of duplicating node parsing. Keep envelopes, eligibility, and sizing local to GQA; other operators and EPs can adopt the shared intent independently. Framework regressions cover valid-token parsing, rejected tokens (including `Memory`, `SAFE`, `memroy`, and ` memory `), unchanged parser output on failure, validation with partitioning enabled and disabled, raw-token retention, unset defaults, and policy-only configuration not enabling partitioning. Existing KV-envelope boundary tests remain intact. CUDA metadata regressions cover default and explicit policy forwarding through the node-config helper/test bridge, identical bytes/backend sets for windowed and non-windowed estimation, and unchanged decline without a non-windowed envelope. #### Validation for this review update - Changed-file lintrunner/clang-format check: passed for all 10 edited C++ files. - `git diff --check`: passed. - No full ORT build or C++ regression suites were run locally for this update, as requested. Earlier-head build evidence does not validate this head; full framework/CUDA validation is left to CI. ### Motivation and Context Provide a generic policy representation and reader without coupling other consumers to GQA or CUDA. Rejecting typos now prevents future enforcement (including #33316's committed policies) from silently falling back to `Auto` and bypassing the requested restrictions. This PR builds on #33016, which is now merged, and targets `main`. Committed policies, operator/EP-specific resolution, route restrictions, runtime coverage/envelope enforcement, bounded fallback, and reduced policy-bound estimates remain separate follow-ups. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Chi Lo Copilot-Session: b95559c4-5cbe-4787-8dd7-ffa45a2157f9 Copilot-Session: 24b20465-fd4d-4a64-8174-aa0622c5164d --- .../core/framework/kernel_dispatch_policy.h | 42 ++++++++ .../core/framework/resource_accountant.h | 2 + .../onnxruntime_session_options_config_keys.h | 9 ++ ...roup_query_attention_workspace_estimate.cc | 12 ++- ...group_query_attention_workspace_estimate.h | 9 +- .../core/framework/resource_accountant.cc | 15 +++ .../providers/cuda/cuda_execution_provider.cc | 16 ++- .../framework/resource_accountant_test.cc | 102 ++++++++++++++++++ .../cuda/test_cases/cuda_test_bridge.cc | 10 +- .../cuda/test_cases/cuda_test_bridge.h | 6 +- ...query_attention_workspace_estimate_test.cc | 71 +++++++++++- 11 files changed, 278 insertions(+), 16 deletions(-) create mode 100644 include/onnxruntime/core/framework/kernel_dispatch_policy.h 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/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/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/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index d341152f18930..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 @@ -3742,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/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/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); From 8b17e671f4b6f3011712ca1c4d8651ba220994da Mon Sep 17 00:00:00 2001 From: aciddelgado <139922440+aciddelgado@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:29:12 -0700 Subject: [PATCH 12/17] [WebGPU] Add bit-exact convolution state replay for speculative decoding (#33493) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Description: ### Description Add `com.microsoft::ConvStateReplay` v1, a dedicated WebGPU contrib operator that reconstructs rolling convolution history from an incoming state snapshot and a retained prefix of captured inputs. It performs copy/selection only, without convolution arithmetic. - Supports FP16 and FP32, batch/channel dimensions, configurable history width and capture capacity, and all retained-prefix counts from zero through capacity, including counts greater than the history width. - Accepts full source, capture, and destination backing tensors plus CPU `int64[8]` metadata describing byte offsets, geometry, capacity, and retained count. - Uses `ProgramInput::BufferView` / `ProgramOutput::BufferView` with integer-word access and logical offsets, including two-byte-aligned FP16 subviews. Integer-only copies preserve signed zero, subnormals, and NaN payloads exactly; masked boundary writes preserve neighboring elements. - Validates geometry, alignment, bounds, binding limits, and destination aliasing, and returns after checked GPU completion. The output must be explicitly prebound to the inactive destination backing. - Adds schema/kernel registration, operator documentation, and focused tests against an independent CPU rolling-history reference, with CPU fallback disabled. Validation: built the native WebGPU provider and ran the following from `build/webgpu-pipeline-lifetime/Release` using Mesa lavapipe: XDG_RUNTIME_DIR=/tmp VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/lvp_icd.x86_64.json ./onnxruntime_provider_test --gtest_filter='*ConvStateReplay*:*GatedDeltaNetStateReplay*' Result: **31 tests passed, zero failures or skips**—20 new convolution replay tests and 11 existing GatedDeltaNet replay tests. Coverage includes every retained prefix, batch/channel indexing, segmented buffers, arbitrary offsets, adjacent-state preservation, every FP16 payload pattern, FP32 special values, zero/boundary cases, and invalid metadata/bounds/aliasing. Repository lint passed and generated contrib documentation was verified. ### Motivation and Context During speculative-decoding partial acceptance, the current GenAI WebGPU fallback downloads convolution state and captured inputs, reconstructs the accepted-prefix history on CPU, then uploads it into the inactive state bank. This fallback accounted for approximately 7.31 seconds of overhead in the reported long Qwen 3.8 DFlash2 width-1 benchmark. Convolution replay only needs to select the final history window from the incoming state followed by the retained captured inputs. This operator provides the GPU primitive needed to eliminate that round trip in a subsequent GenAI integration. This PR is limited to the ORT operator, registration, documentation, and tests. It does not integrate GenAI, change GatedDeltaNet replay behavior, modify the tensor-transfer API, or change production kernel-selection policies. End-to-end performance improvement has not been measured; graph capture is not supported. --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- docs/ContribOperators.md | 64 +++ .../webgpu/bert/conv_state_replay.cc | 205 +++++++++ .../webgpu/bert/conv_state_replay.h | 17 + .../bert/conv_state_replay.wgsl.template | 57 +++ .../webgpu/webgpu_contrib_kernels.cc | 2 + .../core/graph/contrib_ops/bert_defs.cc | 57 +++ onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 + .../core/providers/webgpu/ep/README.md | 51 ++- .../contrib_ops/conv_state_replay_op_test.cc | 412 ++++++++++++++++++ 9 files changed, 865 insertions(+), 2 deletions(-) create mode 100644 onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.cc create mode 100644 onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.h create mode 100644 onnxruntime/contrib_ops/webgpu/bert/conv_state_replay.wgsl.template create mode 100644 onnxruntime/test/contrib_ops/conv_state_replay_op_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 603118adeeed8..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 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/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index ffdb88d239ec3..1fb0dfb3304c0 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5290,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/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index d68cdd392666b..e51a77568683b 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -96,6 +96,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPreMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPostMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNet); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNetStateReplay); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ConvStateReplay); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedAdd); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, NGramHashMapping); @@ -230,6 +231,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); 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/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 From db5f5eaff7c32e09e90520374ac19b78d89094ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B0=8F=E5=93=A5=E4=B8=8D=E5=A4=AA=E9=80=8D=E9=81=A5?= <115146404+xgbtxy@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:36:24 +0800 Subject: [PATCH 13/17] Fix Gemm transpose fusion with public input outputs (#33484) ### Description Check whether each input Transpose produces a graph output before folding it into Gemm. This preserves the model's public outputs while allowing the other, internal input Transpose to be fused. The regression tests cover either input being public with all existing `transA`/`transB` combinations, and both inputs being public when an output Transpose triggers the rewrite. They check the surviving nodes, input/output names, and transpose attributes. ### Motivation and Context Fixes #33441. `SatisfyCondition` can accept Gemm because one input Transpose is safe to fuse. `Apply` then independently processes both inputs, but previously did not repeat the graph-output check. A public input Transpose could therefore be removed, leaving a declared model output without a producer. I reproduced the report with ONNX Runtime 1.31.0 on Windows x64. The unoptimized graph and the pass-disabled control match the reference result; the optimized session fails to resolve the public output. ### Validation - Built current main (`a7a1172c`) and the patch with MSVC on Windows x64, Release configuration. - With the new regressions and the original implementation: 9 existing GemmTransposeFusion tests passed, and the 2 new tests failed. After the fix: all 11 passed. - `GraphTransformationTests.*`: 442 passed, 17 skipped because their GPU execution providers are unavailable in this CPU build. This includes the 11 focused tests. - The original issue's model now runs with optimization enabled and both outputs exactly match `ReferenceEvaluator`. - A local 64-case model matrix covers public-output placement, existing transpose flags, bias, and output transposition. All outputs from the newly built Python module exactly match NumPy. The released 1.31.0 build fails 36 of these cases, while its unoptimized and pass-disabled controls match. - `lintrunner` for both changed files and `git diff --check` pass. AI assistance was used for implementation, regression design, and validation. --- .../core/optimizer/gemm_transpose_fusion.cc | 6 +- .../test/optimizer/graph_transform_test.cc | 157 ++++++++++++++++++ 2 files changed, 161 insertions(+), 2 deletions(-) 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/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"; From a926b837fb8a39b0e4840fcb997ff1336cdf5db1 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 21 Oct 2026 12:56:31 -0700 Subject: [PATCH 14/17] Fix CUDA LayerNorm/RMSNorm silent wrong output for rank-reduced scale (#28971) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description The CUDA `LayerNormalization` / `RMSNormalization` kernels only support broadcasting scale/bias through the encoded `broadcast_param` fast-path; they do **not** implement the generic NumPy-broadcasting path the CPU kernels use. `LayerNormHelper::GetBroadcastParam` required `scale.rank == X.rank` to emit a fast-path value, so a rank-reduced but NumPy-broadcastable scale — e.g. scale `(S, W)` for input `(B, S, W)` with `axis=2` — returned `kLayerNormInvalidInput`. On CPU this falls back to the correct generic path, but the CUDA kernel silently used `broadcast_param = -1`, computing offset `norm_size * (x_row % 1) = 0` and applying `scale[0:W]` to **every** row → wrong results (no error raised). ### Impact This breaks Mamba2 / hybrid SSM models (e.g. **NemotronH**), whose gated `RMSNorm` uses a per-group scale `(n_groups, head_dim)` broadcast over the batch dimension. Such models produce correct logits on CPU but a wrong argmax on CUDA. Minimal repro (fp16): ```python # X:(5,8,960), scale:(8,960), axis=-1 # CPU matches numpy reference; CUDA was off by ~13 (max-abs) before this fix. ``` ### Fix Right-align the scale/bias shapes to `X`'s rank inside `GetBroadcastParam` (NumPy-style, prepending 1s) before matching the `(B,S,...)`, `(B,1,...)`, `(1,S,...)` fast-path layouts. A rank-reduced `(S, W)` scale now maps onto the existing `(1, S, ...)` fast-path and yields the correct per-row offset on CUDA — identical to an equal-rank `(1, S, W)` scale. - Equal-rank shapes are unchanged (right-alignment is the identity). - The CPU generic path is unaffected. ### Testing Adds a CPU+CUDA regression test `RMSNorm_Scale_RankReduced_Sx3_Axis2`. Verified end-to-end: the NemotronH prefill argmax now matches between CPU and CUDA, and a broadcast sweep (RMSNorm + LayerNorm, multiple scale shapes, CPU vs CUDA vs numpy) passes at float epsilon with no regressions. --------- Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/providers/cpu/nn/layer_norm_helper.h | 44 ++++++++--- .../test/providers/cpu/nn/rms_norm_op_test.cc | 73 +++++++++++++++++++ 2 files changed, 106 insertions(+), 11 deletions(-) diff --git a/onnxruntime/core/providers/cpu/nn/layer_norm_helper.h b/onnxruntime/core/providers/cpu/nn/layer_norm_helper.h index e544c4a819f83..d1c47dafb968f 100644 --- a/onnxruntime/core/providers/cpu/nn/layer_norm_helper.h +++ b/onnxruntime/core/providers/cpu/nn/layer_norm_helper.h @@ -215,31 +215,53 @@ class LayerNormHelper { int64_t axis) { // Note that when size of scale and bias is norm_size, it won't enter this function (see CheckInputs). + const size_t x_rank = x_shape.NumDimensions(); + + // Right-align scale/bias dims to X's rank, prepending 1s (NumPy-style + // broadcasting). This lets a rank-reduced scale such as (S, ...) for an + // (B, S, ...) input map onto the (1, S, ...) fast-path encoded in + // broadcast_param, instead of falling back to the generic broadcasting + // path. The fast-path is the only broadcasting mechanism the CUDA + // LayerNorm/RMSNorm kernels support, so without this alignment such a + // scale would be silently mishandled on CUDA. + auto right_align = [x_rank](const TensorShape& s) { + InlinedVector 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/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); From f7b31b70308167e2f277b419f1931aed837c83ad Mon Sep 17 00:00:00 2001 From: Rui Ren Date: Wed, 21 Oct 2026 13:07:31 -0700 Subject: [PATCH 15/17] [CUDA] Cosmos3 - Fix variance cancellation in fused BERT layer normalization (#33357) ### Description Fix variance cancellation in the fused CUDA BERT `SkipLayerNormalization` and `EmbedLayerNormalization` kernels. This is a standalone change against `main`, independent of #33312. The shared regular LayerNorm helpers currently compute variance as `E[x*x] - E[x]*E[x]`. FP32 accumulation does not make that subtraction stable for high-offset, low-variance inputs. The result can be substantially wrong or lead to NaN outputs. Ordinary CUDA `LayerNormalization` already uses Welford/Chan; this change addresses the separate fused helpers. - Replace raw moments with FP32 Welford updates and Chan/CUB block reduction, including count-zero identities for inactive lanes. - Accumulate and normalize relative to a row origin. Keep the mean in that shifted coordinate system instead of reconstructing a rounded high-offset FP32 mean. - Apply the fix to scalar and vectorized skip kernels and the embedding kernel. - Preserve native-dtype input/bias/skip and word/segment/position addition order and rounding, optional residual/embedding-sum outputs, masks, and simplified RMSNorm arithmetic. - Avoid reading a row origin for an empty hidden dimension; cover the existing CUDA SkipLayerNormalization empty-row behavior. No operator schema, public launch interface, workspace, provider option, or ordinary LayerNormalization changes. ### Motivation and evidence For the actual FP32 values obtained from `[1000.01, 1000.02, 999.99, 1000.0]`, centered FP64 population variance is approximately `0.00012524426`. In the preceding isolated CUDA reproduction, with epsilon `1e-5`, the raw-moments skip kernel returned NaNs at widths 4 and 768. The embedding kernel returned NaNs at width 4 and substantially incorrect finite values at width 768. The new FP32 regression tests fail on the unchanged upstream kernels and pass with this fix. No variance clamp, larger epsilon, or relaxed FP32 tolerance is used. ### Validation Tested on H100 (SM90), CUDA 13, using the compiled CUDA provider and C++ operator test executable. ```bash cmake --build build/Cuda13/Release \ --target onnxruntime_providers_cuda onnxruntime_provider_test -j12 cd build/Cuda13/Release CUDA_VISIBLE_DEVICES=0 ./onnxruntime_provider_test \ --gtest_filter='SkipLayerNormTest.*:EmbedLayerNormTest.*' ``` **44 passed, zero skips.** All six new CUDA regression tests executed. The new tests explicitly disable CPU EP fallback. Coverage includes FP32/FP16/BF16 skip normalization and FP32/FP16 embedding normalization; hidden widths 1, 3, 4, 5, 128, 129, 320, 768, 769, 8192 and 8193; constant rows; padded blocks; scalar/vectorized dispatch; multi-element-per-thread reduction; optional bias/segment inputs; masks; and optional sum outputs. Low-precision residual-cancellation cases verify that native addition rounding is retained. The oracle uses centered FP64 arithmetic over the actual native-rounded fused sums. FP32 output bounds remain absolute/relative `1e-4`; low-precision outputs use the repository's existing dtype-aware tolerances. Optional sum outputs require exact equality. `clang-format --dry-run --Werror` and `git diff --check` pass. Validated source commit: `5ff00337b71aa939e023f8342899575e5dc6ba56`, based on `07daf405c788c300aa9fd2988104b29bcbdede23`. CUDA provider SHA-256: `01826247f90f3b6a64f5fde8b9dcf6bf6e62728af34614494b9949029c38c460`. ### Scope and related work Related: #32439 proposes Welford as part of stacked CUDA statistics-output and addition-contract work. This PR fixed the onboarding precision issue of Cosmos3 model, https://github.com/onnxruntime/mobius/pull/466 Signed-off-by: rui-ren <15321482+rui-ren@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/embed_layer_norm_impl.cu | 15 +- .../contrib_ops/cuda/bert/layer_norm.cuh | 62 ++++---- .../cuda/bert/skip_layer_norm_impl.cu | 64 ++++++--- .../contrib_ops/embed_layer_norm_op_test.cc | 114 +++++++++++++++ .../test/contrib_ops/skiplayernorm_op_test.cc | 133 ++++++++++++++++++ 5 files changed, 336 insertions(+), 52 deletions(-) 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/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/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/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/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, From c33b4eb542f360ba077518b3966ca00b09bdd884 Mon Sep 17 00:00:00 2001 From: Rui Ren Date: Wed, 21 Oct 2026 13:10:01 -0700 Subject: [PATCH 16/17] Avoid quadratic allocation-planner free-buffer scans (#33407) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Summary Optimize ONNX Runtime's allocation planner to avoid quadratic free-buffer scans during session initialization. The existing planner maintains a linear list of reusable buffers and searches it when assigning memory to intermediate tensors. For large graphs containing many incompatible buffer candidates, repeated linear scans can result in **O(N²) allocation-planning overhead**. This PR introduces a lightweight, compatibility-keyed free-buffer index to reduce unnecessary candidate scans while preserving existing memory reuse semantics. **Key results:** - **283.38× faster CPU allocation planning** on the synthetic stress workload. - **297.03× faster CUDA allocation planning** on the same workload. - **2.91× faster CPU session initialization**. - **2.75× faster CUDA session initialization**. - No material peak host-memory regression observed. These improvements apply to the targeted pathological workload and should not be generalized to all models. ## Motivation ONNX Runtime reuses memory from intermediate tensors whose lifetimes have ended to reduce memory consumption and avoid unnecessary allocations. Previously, the allocation planner searched a linear free-buffer list to identify compatible reuse candidates. For a graph with many intermediate tensors and incompatible buffer shapes, this can result in repeated examination of buffers that cannot be reused. A motivating large decoder initialization exhibited: | Metric | Observed | |---|---:| | Free-list candidate visits | 150,994,944 | | Successful buffer reuse matches | 6 | This indicates substantial unnecessary search work during allocation-plan construction. The objective is to eliminate this pathological scanning behavior without changing which buffers are eligible for reuse or how execution memory is assigned. ## Implementation Introduce a private `FreeBufferIndex` for non-minimal builds. The index uses open addressing with linear probing to map compatibility keys to free-buffer chain heads. Instead of repeatedly traversing the entire free-buffer list, the planner first locates the relevant compatibility group. Key implementation properties: - **Indexed lookup:** Expected O(1) compatibility-group lookup under normal hash-table conditions, rather than repeatedly scanning unrelated free buffers. - **Full-word keys:** Hashes the complete `size_t` key representation to improve distribution for aligned values. - **Collision safety:** Uses full-key equality checks and linear probing to resolve hash collisions. - **Preserved reuse behavior:** Existing compatibility checks and buffer-selection ordering remain authoritative. - **Explicit occupancy tracking:** An empty chain head does not remove a registered key or break a probe chain. - **Bounded load factor:** The table grows before exceeding approximately 75% occupancy. - **Checked capacity growth:** Uses `SafeInt` to prevent silent capacity arithmetic overflow. - **Minimal-build compatibility:** `ORT_MINIMAL_BUILD` variants retain the original linear planner to avoid additional binary-size overhead. The optimization changes how candidate buffers are located, not the memory reuse rules themselves. No new runtime configuration flags or execution-provider-specific allocation semantics are introduced. ## Performance Evaluation Benchmarked baseline (`main`) against the optimized implementation. All reported measurements are **medians of 10 runs per build**. ### 1. Synthetic Stress Workload The workload consists of: - 8,192 independent four-Cast chains. - 32,768 total nodes. - Different tensor shapes across chains. - Graph optimizations disabled to preserve the stress workload. This workload exercises the original planner's pathological behavior when many free buffers are incompatible with subsequent allocation requests. | Metric | CPU main | CPU optimized | CUDA main | CUDA optimized | |---|---:|---:|---:|---:| | Allocation/execution-plan construction | 26,723.73 ms | 94.30 ms | 25,533.78 ms | 85.96 ms | | Complete session initialization | 40,914.74 ms | 14,079.38 ms | 39,827.44 ms | 14,494.00 ms | | Peak host RSS | 233.81 MiB | 233.79 MiB | 512.63 MiB | 513.02 MiB | | **Planner speedup** | — | **283.38×** | — | **297.03×** | | **Session initialization speedup** | — | **2.91×** | — | **2.75×** | The results demonstrate that reducing incompatible free-buffer scans substantially improves allocation-planning scalability. More importantly, the planner improvement translates into measurable end-to-end session initialization savings. ### 2. Control Workloads and Decoder Model Additional benchmarks were conducted to determine whether the optimization introduces regressions or provides benefits outside the pathological workload. The following measurements are complete session initialization medians. | Workload | CPU main → optimized | CUDA main → optimized | |---|---|---| | Same-shape control: 8,192 chains | 13,916.98 → 14,083.31 ms | 14,527.28 → 14,565.04 ms | | Random-weight 32-layer Llama decoder | 349.77 → 349.42 ms | 803.88 → 798.69 ms | **Observations:** - The same-shape control does not exhibit the large improvement seen in the different-shape stress workload. - The CPU same-shape control is approximately 1.2% slower, while the CUDA difference is approximately 0.3%. Additional variance analysis would be needed to determine whether these differences are statistically meaningful. - The random-weight Llama decoder shows effectively unchanged initialization performance. - These results indicate that the optimization specifically addresses pathological free-buffer scanning rather than universally accelerating model initialization. ### 3. Peak Host-Memory Usage | Configuration | Main | Optimized | Difference | |---|---:|---:|---:| | CPU | 233.81 MiB | 233.79 MiB | −0.02 MiB | | CUDA | 512.63 MiB | 513.02 MiB | +0.39 MiB | No material peak host RSS regression was observed in the synthetic stress workload. The additional compatibility index has a memory cost, but its contribution did not produce a meaningful increase in measured peak process RSS. Peak host RSS does not directly measure the index's standalone memory footprint or GPU device-memory consumption. ## Correctness and Regression Validation The implementation has been validated using focused correctness, allocation-plan equivalence, and complexity regression tests. | Validation | Result | |---|---| | Full-build test harness | 28 passed, 0 skipped | | Independent test rerun | 28 passed, 0 skipped | | Complete allocation/execution plan comparison | 15 plans identical to baseline | | Linear-algorithm semantic control | 20 passed | | Deterministic complexity regression | Passed | | CUDA CI planner cases | 12 passed, 0 skipped/failures | | Android minimal cross-compile, link, and size gate | Passed | The Android minimal-size gate passed at 1,607,070 bytes against the unchanged 1,607,680-byte limit. Android emulator runtime tests were skipped by the workflow and are not claimed as validated. ### Algorithmic Complexity Validation A deterministic complexity test compared the number of shape compatibility queries performed by the original and indexed planners. | Metric | Original planner | Indexed planner | |---|---:|---:| | Shape compatibility queries | 130,306 | 1,022 | This represents approximately **127× fewer shape compatibility queries** in the focused test. The result validates the intended reduction in unnecessary search work. It is not an end-to-end performance speedup measurement. The allocation-plan equivalence tests additionally verify that the optimized planner preserves the tested allocation and execution decisions. ## Scope and Limitations This PR focuses specifically on allocation-planning efficiency. - It reduces CPU-side planning overhead during session initialization. - It does not change operator implementations, numerical precision, or GPU kernels. - It does not directly optimize steady-state inference latency. - It preserves the existing buffer reuse eligibility rules and is intended to maintain equivalent allocation/execution plans. - Minimal builds retain the original linear implementation. - Hash-table lookups have expected O(1) complexity, but worst-case collisions and candidate-chain traversal can still require additional work. - The synthetic stress-test speedups are not representative of every model. - The random-weight 32-layer Llama decoder benchmark shows no substantial initialization improvement. ## Conclusion This PR eliminates a pathological source of quadratic free-buffer scanning in ONNX Runtime's allocation planner. The benchmark results demonstrate **283–297× faster allocation-plan construction** and **2.75–2.91× faster complete session initialization** on the targeted synthetic workload, without a material increase in measured peak host RSS. Control workloads and a 32-layer Llama decoder show little or no performance change, indicating that the optimization primarily benefits graphs with large numbers of incompatible reusable-buffer candidates. The change improves allocation-planning scalability while preserving the existing memory reuse behavior and keeping minimal builds on the original implementation. --------- Signed-off-by: rui-ren <15321482+rui-ren@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/framework/allocation_planner.cc | 135 +++++++-- .../allocation_planner_reuse_index.h | 100 +++++++ .../allocation_planner_reuse_index_test.cc | 162 +++++++++++ .../test/framework/allocation_planner_test.cc | 269 ++++++++++++++---- 4 files changed, 596 insertions(+), 70 deletions(-) create mode 100644 onnxruntime/core/framework/allocation_planner_reuse_index.h create mode 100644 onnxruntime/test/framework/allocation_planner_reuse_index_test.cc 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/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; + } } } From ed8f79bcef8b073d83270a45d68046a72a5166bd Mon Sep 17 00:00:00 2001 From: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:15:46 -0700 Subject: [PATCH 17/17] [CUDA] Honor deterministic compute for MatMulNBits tactic selection (#32898) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description When deterministic compute is enabled, CUDA MatMulNBits now selects a stable fpA-intB tactic instead of using the timing-profiled choice. For small row counts it selects the available CUDA GEMV tactic; otherwise it selects a fixed non-split-K GEMM tactic. This path bypasses the profiled tactic cache. With deterministic compute disabled, the existing profiling behavior is unchanged. The PR adds tests for cache independence and exact 2-, 4-, and 8-bit weight reconstruction in deterministic mode. ### Motivation and Context Fresh processes running the same quantized model with greedy decoding sometimes selected different MatMulNBits tile and split-K tactics. Those tactics produced slightly different FP16 outputs from identical inputs, eventually changing the ninth generated token from `680` to `264`. ORT’s existing deterministic-compute flag did not previously govern this tactic selection. On SM80, the isolated MatMul output matched bitwise in 12/12 fresh processes after the change. In a downstream full-model validation, 36/36 deterministic runs produced the same token stream and tactic fingerprint; with the flag disabled, both original branches still appeared. Deterministic selection may be slower than a profiled tactic. Validation does not establish identical results across hardware or batch shapes. ### Qwen INT4 decode specialization This PR now also routes the Qwen Flash-Next language-model projection through the existing M=1 symmetric INT4 CUDA GEMV specialization for the exact shape: - `M = 1` - `N = 248320` - `K = 2560` - block size `32` - no zero points The dispatch remains shape- and format-gated, so other `MatMulNBits` workloads retain their existing selection behavior. The specialization reduces single-token draft-head projection overhead used by Qwen speculative decoding. --------- Co-authored-by: Akshay Sonawane Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- .../cuda/llm/fpA_intB_gemm_profiler.cc | 19 ++++ .../cuda/llm/fpA_intB_gemm_profiler.h | 2 + .../cuda/quantization/matmul_4bits.cu | 55 +++------ .../cuda/quantization/matmul_nbits.cc | 6 +- .../cuda/quantization/matmul_nbits.cuh | 8 ++ .../matmul_nbits_fpa_intb_layout_test.cc | 10 +- .../test_cases/matmul_nbits_workspace_test.cc | 107 ++++++++++++++++++ 7 files changed, 165 insertions(+), 42 deletions(-) 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/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/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/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.