mirror of
https://github.com/opencv/opencv.git
synced 2026-07-29 07:13:02 +04:00
Merge pull request #29104 from abhishek-gola:sdpa
Added SDPA layer (Scaled Dot Product Attention) #29104 Merge with: https://github.com/opencv/opencv_extra/pull/1374 ### Pull Request Readiness Checklist See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request - [x] I agree to contribute to the project under Apache 2 License. - [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV - [x] The PR is proposed to the proper branch - [x] There is a reference to the original bug report and related work - [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable Patch to opencv_extra has the same branch name. - [x] The feature is well documented and sample code can be built with the project CMake
This commit is contained in:
Abhishek Gola
GitHub
@@ -1872,6 +1872,12 @@ CV__DNN_INLINE_NS_BEGIN
|
||||
static Ptr<AttentionLayer> 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<SDPALayer> create(const LayerParams ¶ms);
|
||||
};
|
||||
|
||||
class CV_EXPORTS RotaryEmbeddingLayer : public Layer {
|
||||
public:
|
||||
static Ptr<RotaryEmbeddingLayer> create(const LayerParams ¶ms);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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<float>("scale", 1.f);
|
||||
}
|
||||
|
||||
virtual bool supportBackend(int backendId) CV_OVERRIDE {
|
||||
return backendId == DNN_BACKEND_OPENCV;
|
||||
}
|
||||
|
||||
virtual bool getMemoryShapes(const std::vector<MatShape> &inputs,
|
||||
const int requiredOutputs,
|
||||
std::vector<MatShape> &outputs,
|
||||
std::vector<MatShape> &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<MatShape> &inputs,
|
||||
const std::vector<MatShape> &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<Mat> 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<const float>();
|
||||
float* k_data = K.ptr<float>();
|
||||
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<float>(), K.ptr<float>(), V.ptr<float>(),
|
||||
Y.ptr<float>(),
|
||||
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>();
|
||||
const float* k_p = K.ptr<const float>();
|
||||
const float* v_p = V.ptr<const float>();
|
||||
float* y_p = Y.ptr<float>();
|
||||
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<float> 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<unsigned char> flash_scratch;
|
||||
};
|
||||
|
||||
Ptr<SDPALayer> SDPALayer::create(const LayerParams ¶ms) {
|
||||
return makePtr<SDPALayerImpl>(params);
|
||||
}
|
||||
|
||||
}} // cv::dnn
|
||||
@@ -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"] =
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<float>();
|
||||
const float* kt = KT.ptr<float>();
|
||||
const float* v = V.ptr<float>();
|
||||
|
||||
int szY[3] = {B, S_q, H * D_v};
|
||||
Mat Y(3, szY, CV_32F, Scalar(0));
|
||||
float* y = Y.ptr<float>();
|
||||
|
||||
const float inv_scale = (scale > 0.f) ? 1.f / scale : 1.f;
|
||||
std::vector<float> 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::dnn::EngineType>(
|
||||
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");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user