mirror of
https://github.com/Tencent/ncnn.git
synced 2026-10-19 23:03:17 +08:00
nihui
113 lines
3.8 KiB
C++
113 lines
3.8 KiB
C++
// Copyright 2025 Tencent
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
#include "testutil.h"
|
|
|
|
static int test_sdpa(const ncnn::Mat& q, const ncnn::Mat& k, const ncnn::Mat& v, int attn_mask, float scale = 0.f)
|
|
{
|
|
const int src_seqlen = q.h;
|
|
const int dst_seqlen = k.h;
|
|
|
|
ncnn::ParamDict pd;
|
|
pd.set(5, attn_mask);
|
|
pd.set(6, scale);
|
|
|
|
std::vector<ncnn::Mat> weights(0);
|
|
|
|
std::vector<ncnn::Mat> as(3);
|
|
as[0] = q;
|
|
as[1] = k;
|
|
as[2] = v;
|
|
|
|
if (attn_mask)
|
|
{
|
|
as.push_back(RandomMat(dst_seqlen, src_seqlen));
|
|
}
|
|
|
|
float epsilon = 0.001;
|
|
|
|
int ret = test_layer("SDPA", pd, weights, as, 1, epsilon);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_sdpa failed q=(%d %d %d) k=(%d %d %d) v=(%d %d %d) attn_mask=%d scale=%f\n", q.w, q.h, q.c, k.w, k.h, k.c, v.w, v.h, v.c, attn_mask, scale);
|
|
}
|
|
|
|
return ret;
|
|
}
|
|
|
|
static int test_sdpa_0()
|
|
{
|
|
return 0
|
|
|| test_sdpa(RandomMat(32, 66, 8), RandomMat(32, 66, 8), RandomMat(20, 66, 8), 0)
|
|
|| test_sdpa(RandomMat(26, 64, 8), RandomMat(26, 61, 8), RandomMat(18, 61, 8), 1)
|
|
|| test_sdpa(RandomMat(40, 62, 7), RandomMat(40, 61, 7), RandomMat(24, 61, 7), 0)
|
|
|| test_sdpa(RandomMat(24, 22, 6), RandomMat(24, 19, 6), RandomMat(16, 19, 6), 1)
|
|
|| test_sdpa(RandomMat(64, 128, 12), RandomMat(64, 128, 2), RandomMat(64, 128, 2), 0)
|
|
|| test_sdpa(RandomMat(64, 122, 12), RandomMat(64, 127, 2), RandomMat(48, 127, 2), 1)
|
|
|| test_sdpa(RandomMat(44, 128, 4), RandomMat(44, 123, 4), RandomMat(55, 123, 4), 0, 1.f)
|
|
|| test_sdpa(RandomMat(12, 127, 4), RandomMat(12, 127, 4), RandomMat(55, 127, 4), 1, 1.f)
|
|
|| test_sdpa(RandomMat(28, 17, 15), RandomMat(28, 127, 5), RandomMat(32, 127, 5), 0, 0.1f)
|
|
|| test_sdpa(RandomMat(28, 17, 15), RandomMat(28, 32, 5), RandomMat(11, 32, 5), 1, -0.4f);
|
|
}
|
|
|
|
#if NCNN_INT8
|
|
static int test_sdpa_int8(const ncnn::Mat& q, const ncnn::Mat& k, const ncnn::Mat& v, int attn_mask, float scale = 0.f)
|
|
{
|
|
const int src_seqlen = q.h;
|
|
const int dst_seqlen = k.h;
|
|
|
|
ncnn::ParamDict pd;
|
|
pd.set(5, attn_mask);
|
|
pd.set(6, scale);
|
|
pd.set(18, 2); // int8_scale_term
|
|
|
|
std::vector<ncnn::Mat> weights(0);
|
|
|
|
std::vector<ncnn::Mat> as(3);
|
|
as[0] = q;
|
|
as[1] = k;
|
|
as[2] = v;
|
|
|
|
if (attn_mask)
|
|
{
|
|
as.push_back(RandomMat(dst_seqlen, src_seqlen));
|
|
}
|
|
|
|
float epsilon = 0.01;
|
|
|
|
int ret = test_layer("SDPA", pd, weights, as, 1, epsilon);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_sdpa_int8 failed q=(%d %d %d) k=(%d %d %d) v=(%d %d %d) attn_mask=%d scale=%f\n", q.w, q.h, q.c, k.w, k.h, k.c, v.w, v.h, v.c, attn_mask, scale);
|
|
}
|
|
|
|
return ret;
|
|
}
|
|
|
|
static int test_sdpa_1()
|
|
{
|
|
return 0
|
|
|| test_sdpa_int8(RandomMat(32, 66, 8), RandomMat(32, 66, 8), RandomMat(20, 66, 8), 0)
|
|
|| test_sdpa_int8(RandomMat(26, 64, 8), RandomMat(26, 61, 8), RandomMat(18, 61, 8), 1)
|
|
|| test_sdpa_int8(RandomMat(40, 62, 7), RandomMat(40, 61, 7), RandomMat(24, 61, 7), 0)
|
|
|| test_sdpa_int8(RandomMat(24, 22, 6), RandomMat(24, 19, 6), RandomMat(16, 19, 6), 1)
|
|
|| test_sdpa_int8(RandomMat(64, 128, 12), RandomMat(64, 128, 2), RandomMat(64, 128, 2), 0)
|
|
|| test_sdpa_int8(RandomMat(48, 122, 12), RandomMat(64, 127, 2), RandomMat(64, 127, 2), 1)
|
|
|| test_sdpa_int8(RandomMat(44, 128, 4), RandomMat(44, 123, 4), RandomMat(55, 123, 4), 0, 1.f)
|
|
|| test_sdpa_int8(RandomMat(12, 127, 4), RandomMat(12, 127, 4), RandomMat(55, 127, 4), 1, 1.f)
|
|
|| test_sdpa_int8(RandomMat(28, 17, 15), RandomMat(28, 127, 5), RandomMat(32, 127, 5), 0, 0.1f)
|
|
|| test_sdpa_int8(RandomMat(28, 17, 15), RandomMat(28, 32, 5), RandomMat(11, 32, 5), 1, -0.4f);
|
|
}
|
|
#endif
|
|
|
|
int main()
|
|
{
|
|
SRAND(7767517);
|
|
|
|
#if NCNN_INT8
|
|
return test_sdpa_0() || test_sdpa_1();
|
|
#else
|
|
return test_sdpa_0();
|
|
#endif
|
|
}
|