diff --git a/modules/dnn/include/opencv2/dnn/all_layers.hpp b/modules/dnn/include/opencv2/dnn/all_layers.hpp index d11af929e6..2e8f56add1 100644 --- a/modules/dnn/include/opencv2/dnn/all_layers.hpp +++ b/modules/dnn/include/opencv2/dnn/all_layers.hpp @@ -1872,6 +1872,12 @@ CV__DNN_INLINE_NS_BEGIN static Ptr create(const LayerParams ¶ms); }; + // Scaled-dot-product attention on pre-projected, pre-reshaped Q / K^T / V. + class CV_EXPORTS SDPALayer : public Layer { + public: + static Ptr create(const LayerParams ¶ms); + }; + class CV_EXPORTS RotaryEmbeddingLayer : public Layer { public: static Ptr create(const LayerParams ¶ms); diff --git a/modules/dnn/src/init.cpp b/modules/dnn/src/init.cpp index ca78e3b9a1..cc696a2db6 100644 --- a/modules/dnn/src/init.cpp +++ b/modules/dnn/src/init.cpp @@ -212,6 +212,7 @@ void initializeLayerFactory() CV_DNN_REGISTER_LAYER_CLASS(Expand, ExpandLayer); CV_DNN_REGISTER_LAYER_CLASS(InstanceNormalization, InstanceNormLayer); CV_DNN_REGISTER_LAYER_CLASS(Attention, AttentionLayer); + CV_DNN_REGISTER_LAYER_CLASS(SDPA, SDPALayer); CV_DNN_REGISTER_LAYER_CLASS(AttentionOnnxAi, AttentionOnnxAiLayer); CV_DNN_REGISTER_LAYER_CLASS(RotaryEmbedding, RotaryEmbeddingLayer); CV_DNN_REGISTER_LAYER_CLASS(GroupNormalization, GroupNormLayer); diff --git a/modules/dnn/src/layers/sdpa_layer.cpp b/modules/dnn/src/layers/sdpa_layer.cpp new file mode 100644 index 0000000000..3e4ea19f76 --- /dev/null +++ b/modules/dnn/src/layers/sdpa_layer.cpp @@ -0,0 +1,193 @@ +// This file is part of OpenCV project. +// It is subject to the license terms in the LICENSE file found in the top-level directory +// of this distribution and at http://opencv.org/license.html. +// Copyright (C) 2026, BigVision LLC, all rights reserved. +// Third party copyrights are property of their respective owners. + +#include "../precomp.hpp" +#include "cpu_kernels/mlas_gemm.hpp" + +namespace cv { namespace dnn { + +// Scaled-Dot-Product Attention on pre-projected Q, K^T, V. Fuses +// QK^T -> [scale] -> softmax -> ·V -> Transpose -> Reshape into one +// MlasFlashAttention call. Used by the SR-attention fusion pass. +class SDPALayerImpl CV_FINAL : public SDPALayer { +public: + SDPALayerImpl(const LayerParams ¶ms) { + setParamsFrom(params); + scale = params.get("scale", 1.f); + } + + virtual bool supportBackend(int backendId) CV_OVERRIDE { + return backendId == DNN_BACKEND_OPENCV; + } + + virtual bool getMemoryShapes(const std::vector &inputs, + const int requiredOutputs, + std::vector &outputs, + std::vector &internals) const CV_OVERRIDE { + CV_CheckEQ(inputs.size(), 3u, "DNN/SDPA: expects Q, K^T, V"); + const auto &q = inputs[0]; // (B, H, S_q, D) + const auto &kT = inputs[1]; // (B, H, D, S_kv) + const auto &v = inputs[2]; // (B, H, S_kv, D_v) + CV_CheckEQ(q.dims, 4, "DNN/SDPA: Q must be 4D"); + CV_CheckEQ(kT.dims, 4, "DNN/SDPA: K must be 4D"); + CV_CheckEQ(v.dims, 4, "DNN/SDPA: V must be 4D"); + const int B = q[0], H = q[1], S_q = q[2], D = q[3]; + const int S_kv = v[2], D_v = v[3]; + CV_CheckEQ(kT[0], B, "DNN/SDPA: K batch"); + CV_CheckEQ(kT[1], H, "DNN/SDPA: K heads"); + CV_CheckEQ(kT[2], D, "DNN/SDPA: K^T inner dim must equal Q's head_dim"); + CV_CheckEQ(kT[3], S_kv, "DNN/SDPA: K^T outer dim must equal V's seq len"); + CV_CheckEQ(v[0], B, "DNN/SDPA: V batch"); + CV_CheckEQ(v[1], H, "DNN/SDPA: V heads"); + + outputs.assign(1, MatShape{B, S_q, H * D_v}); + // Scratch for K^T re-transposed to (B, H, S_kv, D). + internals.assign(1, MatShape{B * H * S_kv * D}); + return false; + } + + virtual int64 getFLOPS(const std::vector &inputs, + const std::vector &outputs) const CV_OVERRIDE { + const auto &q = inputs[0]; + const auto &v = inputs[2]; + const int64 B = q[0], H = q[1], S_q = q[2], D = q[3]; + const int64 S_kv = v[2], D_v = v[3]; + return B * H * (CV_BIG_INT(2) * S_q * S_kv * D + 4 * S_q * S_kv + + CV_BIG_INT(2) * S_q * S_kv * D_v); + } + + void forward(InputArrayOfArrays inputs_arr, + OutputArrayOfArrays outputs_arr, + OutputArrayOfArrays internals_arr) CV_OVERRIDE + { + CV_TRACE_FUNCTION(); + CV_TRACE_ARG_VALUE(name, "name", name.c_str()); + + if (inputs_arr.depth() == CV_16F) { + forward_fallback(inputs_arr, outputs_arr, internals_arr); + return; + } + + std::vector inputs, outputs, internals; + inputs_arr.getMatVector(inputs); + outputs_arr.getMatVector(outputs); + internals_arr.getMatVector(internals); + + const Mat& Q = inputs[0]; // (B, H, S_q, D) + const Mat& KT = inputs[1]; // (B, H, D, S_kv) + const Mat& V = inputs[2]; // (B, H, S_kv, D_v) + Mat& Y = outputs[0]; // (B, S_q, H*D_v) + + const int B = Q.size[0]; + const int H = Q.size[1]; + const int S_q = Q.size[2]; + const int D = Q.size[3]; + const int S_kv = V.size[2]; + const int D_v = V.size[3]; + + // Re-transpose K^T (B,H,D,S_kv) -> K (B,H,S_kv,D) for MlasFlashAttention. + Mat &K = internals[0]; + const float* kT_data = KT.ptr(); + float* k_data = K.ptr(); + const size_t kT_bhd_stride = (size_t)D * S_kv; + const size_t k_bhd_stride = (size_t)S_kv * D; + cv::parallel_for_(cv::Range(0, B * H), [&](const cv::Range& r) { + for (int bh = r.start; bh < r.end; bh++) { + const float* src = kT_data + (size_t)bh * kT_bhd_stride; + float* dst = k_data + (size_t)bh * k_bhd_stride; + for (int s = 0; s < S_kv; s++) { + for (int d = 0; d < D; d++) { + dst[(size_t)s * D + d] = src[(size_t)d * S_kv + s]; + } + } + } + }); + + const int q_block = std::min(256, S_q); + const int kv_block = std::min(256, S_kv); + const int threads = std::max(1, cv::getNumThreads()); + + const size_t per_thread = + mlasFlashAttentionBufferBytesPerThread(q_block, kv_block, D_v); + if (mlasAvailable() && per_thread > 0) { + flash_scratch.resize((size_t)threads * per_thread); + const float effective_scale = (scale > 0.f) ? (1.f / scale) : 1.f; + if (mlasFlashAttention(Q.ptr(), K.ptr(), V.ptr(), + Y.ptr(), + B, H, S_q, S_kv, D, D_v, effective_scale, + q_block, kv_block, + flash_scratch.data(), threads)) + { + return; + } + } + + forward_fallback_impl(Q, K, V, Y, B, H, S_q, S_kv, D, D_v); + } + +private: + // Unfused SDPA, used when MLAS isn't available. Not perf-critical. + void forward_fallback_impl(const Mat& Q, const Mat& K, const Mat& V, Mat& Y, + int B, int H, int S_q, int S_kv, int D, int D_v) const + { + const float* q_p = Q.ptr(); + const float* k_p = K.ptr(); + const float* v_p = V.ptr(); + float* y_p = Y.ptr(); + const float inv_scale = (scale > 0.f) ? (1.f / scale) : 1.f; + + cv::parallel_for_(cv::Range(0, B * H), [&](const cv::Range& r) { + std::vector attn((size_t)S_q * S_kv); + for (int bh = r.start; bh < r.end; bh++) { + const int b = bh / H, h = bh % H; + const float* Q_bh = q_p + (size_t)bh * S_q * D; + const float* K_bh = k_p + (size_t)bh * S_kv * D; + const float* V_bh = v_p + (size_t)bh * S_kv * D_v; + float* Y_b = y_p + (size_t)b * S_q * H * D_v; + + for (int i = 0; i < S_q; i++) { + for (int j = 0; j < S_kv; j++) { + float s = 0.f; + for (int d = 0; d < D; d++) + s += Q_bh[(size_t)i * D + d] * K_bh[(size_t)j * D + d]; + attn[(size_t)i * S_kv + j] = s * inv_scale; + } + float mx = attn[(size_t)i * S_kv]; + for (int j = 1; j < S_kv; j++) + if (attn[(size_t)i * S_kv + j] > mx) mx = attn[(size_t)i * S_kv + j]; + float sum = 0.f; + for (int j = 0; j < S_kv; j++) { + float e = std::exp(attn[(size_t)i * S_kv + j] - mx); + attn[(size_t)i * S_kv + j] = e; + sum += e; + } + float inv_sum = 1.f / sum; + for (int j = 0; j < S_kv; j++) + attn[(size_t)i * S_kv + j] *= inv_sum; + + // out[b, i, h, :] = attn[i, :] @ V_bh + float* out_row = Y_b + ((size_t)i * H + h) * D_v; + for (int d = 0; d < D_v; d++) out_row[d] = 0.f; + for (int j = 0; j < S_kv; j++) { + float a = attn[(size_t)i * S_kv + j]; + const float* V_row = V_bh + (size_t)j * D_v; + for (int d = 0; d < D_v; d++) + out_row[d] += a * V_row[d]; + } + } + } + }); + } + + float scale; + std::vector flash_scratch; +}; + +Ptr SDPALayer::create(const LayerParams ¶ms) { + return makePtr(params); +} + +}} // cv::dnn diff --git a/modules/dnn/src/onnx/onnx_importer2.cpp b/modules/dnn/src/onnx/onnx_importer2.cpp index 438370f5fd..a8212fb781 100644 --- a/modules/dnn/src/onnx/onnx_importer2.cpp +++ b/modules/dnn/src/onnx/onnx_importer2.cpp @@ -284,6 +284,7 @@ protected: // URL: https://github.com/microsoft/onnxruntime/blob/master/docs/ContribOperators.md void parseAttention (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto); void parseAttentionOnnxAi (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto); + void parseSDPA (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto); void parseDequantizeLinear (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto); void parseQuantizeLinear (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto); void parseCustomLayer (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto); @@ -2608,6 +2609,11 @@ void ONNXImporter2::parseAttentionOnnxAi(LayerParams& params, const opencv_onnx: addLayer(params, node_proto, n_inputs); } +void ONNXImporter2::parseSDPA(LayerParams& params, const opencv_onnx::NodeProto& node_proto) { + CV_CheckEQ(node_proto.input_size(), 3, "ONNXImporter2/parseSDPA: SDPA expects 3 inputs (Q, K^T, V)"); + addLayer(params, node_proto, 3); +} + void ONNXImporter2::parseRoiAlign(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto) { layerParams.type = "RoiAlign"; @@ -2708,7 +2714,7 @@ void ONNXImporter2::buildDispatchMap_ONNX_AI() dispatch["RotaryEmbedding"] = &ONNXImporter2::parseRotaryEmbedding; dispatch["NegativeLogLikelihoodLoss"] = &ONNXImporter2::parseNegativeLogLikelihoodLoss; dispatch["SoftmaxCrossEntropyLoss"] = &ONNXImporter2::parseSoftmaxCrossEntropyLoss; - // @TODO@ONNX: Add support for SDPA + dispatch["SDPA"] = &ONNXImporter2::parseSDPA; dispatch["Equal"] = dispatch["Greater"] = dispatch["Less"] = dispatch["Pow"] = dispatch["Add"] = dispatch["Sub"] = dispatch["Mul"] = dispatch["Div"] = dispatch["GreaterOrEqual"] = diff --git a/modules/dnn/test/test_layers.cpp b/modules/dnn/test/test_layers.cpp index 94af725bce..6616aeda1b 100644 --- a/modules/dnn/test/test_layers.cpp +++ b/modules/dnn/test/test_layers.cpp @@ -768,8 +768,6 @@ TEST(Layer_MHARoPe_Test_Accuracy_with_, Pytorch) normAssert(h_t_reference, outputs[0]); } - - TEST_P(Test_Caffe_layers, Accum) { #ifdef OPENCV_DNN_EXTERNAL_PROTOBUF diff --git a/modules/dnn/test/test_onnx_importer.cpp b/modules/dnn/test/test_onnx_importer.cpp index 27ede08b53..036f0eee4a 100644 --- a/modules/dnn/test/test_onnx_importer.cpp +++ b/modules/dnn/test/test_onnx_importer.cpp @@ -28,6 +28,63 @@ static std::string _tf(TString filename, bool required = true) return findDataFile(std::string("dnn/onnx/") + filename, required); } +static Mat sdpaReference(const Mat& Q, const Mat& KT, const Mat& V, float scale) +{ + CV_Assert(Q.dims == 4 && KT.dims == 4 && V.dims == 4); + const int B = Q.size[0], H = Q.size[1], S_q = Q.size[2], D = Q.size[3]; + const int S_kv = V.size[2], D_v = V.size[3]; + + const float* q = Q.ptr(); + const float* kt = KT.ptr(); + const float* v = V.ptr(); + + int szY[3] = {B, S_q, H * D_v}; + Mat Y(3, szY, CV_32F, Scalar(0)); + float* y = Y.ptr(); + + const float inv_scale = (scale > 0.f) ? 1.f / scale : 1.f; + std::vector attn(S_kv); + + for (int b = 0; b < B; b++) + for (int h = 0; h < H; h++) + { + const float* Q_bh = q + ((size_t)(b * H + h) * S_q ) * D; + const float* KT_bh = kt + ((size_t)(b * H + h) * D ) * S_kv; + const float* V_bh = v + ((size_t)(b * H + h) * S_kv) * D_v; + + for (int i = 0; i < S_q; i++) + { + for (int j = 0; j < S_kv; j++) + { + float s = 0.f; + for (int d = 0; d < D; d++) + s += Q_bh[i * D + d] * KT_bh[d * S_kv + j]; + attn[j] = s * inv_scale; + } + float mx = attn[0]; + for (int j = 1; j < S_kv; j++) + mx = std::max(mx, attn[j]); + float sum = 0.f; + for (int j = 0; j < S_kv; j++) + { + attn[j] = std::exp(attn[j] - mx); + sum += attn[j]; + } + const float inv_sum = 1.f / sum; + + float* out_row = y + ((size_t)(b * S_q + i) * H + h) * D_v; + for (int d = 0; d < D_v; d++) + { + float acc = 0.f; + for (int j = 0; j < S_kv; j++) + acc += attn[j] * V_bh[j * D_v + d]; + out_row[d] = acc * inv_sum; + } + } + } + return Y; +} + class Test_ONNX_layers : public DNNTestLayer { public: @@ -155,6 +212,40 @@ public: if (checkNoFallbacks) expectNoFallbacksFromIE(net); } + + // Runs an SDPA ONNX model through the importer + SDPALayer and checks the + // output against an in-test attention reference computed from the same inputs. + void testSDPAModel(const String& basename, double l1, double lInf) + { + // SDPA is only handled by the new-engine ONNX importer. + auto engine_forced = static_cast( + cv::utils::getConfigurationParameterSizeT("OPENCV_FORCE_DNN_ENGINE", cv::dnn::ENGINE_AUTO)); + if (engine_forced == cv::dnn::ENGINE_CLASSIC) { + applyTestTag(CV_TEST_TAG_DNN_SKIP_PARSER); + return; + } + + Mat Q = blobFromNPY(_tf("data/input_" + basename + "_0.npy")); + Mat KT = blobFromNPY(_tf("data/input_" + basename + "_1.npy")); + Mat V = blobFromNPY(_tf("data/input_" + basename + "_2.npy")); + + Net net = readNetFromONNX(_tf("models/" + basename + ".onnx", required)); + ASSERT_FALSE(net.empty()); + net.setPreferableBackend(backend); + net.setPreferableTarget(target); + + net.setInputsNames({"0", "1", "2"}); + net.setInput(Q, "0"); + net.setInput(KT, "1"); + net.setInput(V, "2"); + Mat out = net.forward(""); + + // The models are generated with scale = sqrt(head_dim). + Mat ref = sdpaReference(Q, KT, V, std::sqrt((float)Q.size[3])); + + EXPECT_EQ(ref.shape(), out.shape()); + normAssert(ref, out, basename.c_str(), l1 ? l1 : default_l1, lInf ? lInf : default_lInf); + } }; TEST_P(Test_ONNX_layers, InstanceNorm) @@ -3438,6 +3529,16 @@ TEST_P(Test_ONNX_layers, PyTorchAttentionSingleHead) { testONNXModels("pytorch_attention_single_head"); } +// Batch + multi-head, square attention (B=2, H=4, S_q=S_kv=16, D=D_v=32). +TEST_P(Test_ONNX_layers, SDPA_MultiHead) { + testSDPAModel("sdpa_multi_head", 1e-4, 5e-4); +} +// Cross-attention with asymmetric shapes (S_q=12 != S_kv=20, D=16 != D_v=24) +// to exercise the K^T re-transpose and head-merge indexing. +TEST_P(Test_ONNX_layers, SDPA_CrossAttention) { + testSDPAModel("sdpa_cross_attention", 1e-4, 5e-4); +} + TEST_P(Test_ONNX_layers, PyTorchUnflatten){ testONNXModels("unflatten"); }