mirror of
https://github.com/opencv/opencv.git
synced 2026-07-29 15:23:05 +04:00
Merge pull request #28524 from nklskyoy:attn-kv-cache
CPU Kernels for fp32 KV Cache #28524 The kernels are: - pagedAttnQKGemmKernel - pagedAttnAVGemmKernel ### 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 - [ ] The PR is proposed to the proper branch - [ ] There is a reference to the original bug report and related work - [ ] There is accuracy test, performance test and test data in opencv_extra repository, if applicable Patch to opencv_extra has the same branch name. - [ ] The feature is well documented and sample code can be built with the project CMake
This commit is contained in:
nklskyoy
GitHub
@@ -20,6 +20,84 @@
|
||||
|
||||
namespace cv { namespace dnn {
|
||||
|
||||
int fastGemmMC(const FastGemmOpt &opt) {
|
||||
#if CV_TRY_NEON
|
||||
if (opt.use_neon) {
|
||||
return opt_NEON::fastGemmMC();
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_AVX2
|
||||
if (opt.use_avx2) {
|
||||
return opt_AVX2::fastGemmMC();
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_AVX
|
||||
if (opt.use_avx) {
|
||||
return opt_AVX::fastGemmMC();
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_LASX
|
||||
if (opt.use_lasx) {
|
||||
return opt_LASX::fastGemmMC();
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
return cpu_baseline::fastGemmMC();
|
||||
}
|
||||
}
|
||||
|
||||
int fastGemmNC(const FastGemmOpt &opt) {
|
||||
#if CV_TRY_NEON
|
||||
if (opt.use_neon) {
|
||||
return opt_NEON::fastGemmNC();
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_AVX2
|
||||
if (opt.use_avx2) {
|
||||
return opt_AVX2::fastGemmNC();
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_AVX
|
||||
if (opt.use_avx) {
|
||||
return opt_AVX::fastGemmNC();
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_LASX
|
||||
if (opt.use_lasx) {
|
||||
return opt_LASX::fastGemmNC();
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
return cpu_baseline::fastGemmNC();
|
||||
}
|
||||
}
|
||||
|
||||
int fastGemmKC(const FastGemmOpt &opt) {
|
||||
#if CV_TRY_NEON
|
||||
if (opt.use_neon) {
|
||||
return opt_NEON::fastGemmKC();
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_AVX2
|
||||
if (opt.use_avx2) {
|
||||
return opt_AVX2::fastGemmKC();
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_AVX
|
||||
if (opt.use_avx) {
|
||||
return opt_AVX::fastGemmKC();
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_LASX
|
||||
if (opt.use_lasx) {
|
||||
return opt_LASX::fastGemmKC();
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
return cpu_baseline::fastGemmKC();
|
||||
}
|
||||
}
|
||||
|
||||
size_t fastGemmPackBSize(size_t N, size_t K, const FastGemmOpt &opt) {
|
||||
#if CV_TRY_NEON
|
||||
if (opt.use_neon) {
|
||||
@@ -399,4 +477,224 @@ void fastGemmBatch(bool trans_a, bool trans_b,
|
||||
helper.ldb1, beta, c, helper.ldc, opt);
|
||||
}
|
||||
|
||||
|
||||
void fastGemmBatch(size_t batch,
|
||||
const std::vector<size_t> &A_offsets, const std::vector<size_t> &B_offsets, const std::vector<size_t> &C_offsets,
|
||||
int M, int N, int K, float alpha, const Mat&A, int lda0, int lda1,
|
||||
const Mat&B, int ldb0, int ldb1, float beta, Mat&C, int ldc, FastGemmOpt &opt){
|
||||
const char *a = (const char *)A.ptr<const float>();
|
||||
const char *b = (const char *)B.ptr<const float>();
|
||||
char *c = (char *)C.ptr<float>();
|
||||
|
||||
#if CV_TRY_NEON
|
||||
if (opt.use_neon) {
|
||||
opt_NEON::fastGemmBatchKernel(batch, A_offsets.data(), B_offsets.data(), C_offsets.data(), M, N, K, alpha, a, lda0, lda1, b, ldb0, ldb1, beta, c, ldc, sizeof(float));
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_AVX2
|
||||
if (opt.use_avx2) {
|
||||
opt_AVX2::fastGemmBatchKernel(batch, A_offsets.data(), B_offsets.data(), C_offsets.data(), M, N, K, alpha, a, lda0, lda1, b, ldb0, ldb1, beta, c, ldc, sizeof(float));
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_AVX
|
||||
if (opt.use_avx) {
|
||||
opt_AVX::fastGemmBatchKernel(batch, A_offsets.data(), B_offsets.data(), C_offsets.data(), M, N, K, alpha, a, lda0, lda1, b, ldb0, ldb1, beta, c, ldc, sizeof(float));
|
||||
} else
|
||||
#endif
|
||||
#if CV_TRY_LASX
|
||||
if (opt.use_lasx) {
|
||||
opt_LASX::fastGemmBatchKernel(batch, A_offsets.data(), B_offsets.data(), C_offsets.data(), M, N, K, alpha, a, lda0, lda1, b, ldb0, ldb1, beta, c, ldc, sizeof(float));
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
cpu_baseline::fastGemmBatchKernel(batch, A_offsets.data(), B_offsets.data(), C_offsets.data(), M, N, K, alpha, a, lda0, lda1, b, ldb0, ldb1, beta, c, ldc, sizeof(float));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void pagedAttnQKGemm(
|
||||
const Mat& Q,const std::vector<Mat> &K, Mat& A,
|
||||
int T_q, int Nq, int N_k, int T_s, int D,
|
||||
const FastGemmOpt &opt
|
||||
) {
|
||||
size_t esz = Q.elemSize();
|
||||
|
||||
for (size_t s = 0; s < K.size(); s++){
|
||||
CV_CheckTypeEQ(Q.type(), K[s].type(), "pagedAttnQKGemmKernel: Q and K should have the same type");
|
||||
CV_CheckTrue(esz == K[s].elemSize(), "pagedAttnQKGemmKernel: Q and K should have the same element size");
|
||||
}
|
||||
CV_CheckTypeEQ(Q.type(), A.type(), "pagedAttnQKGemmKernel: Q and A should have the same type");
|
||||
|
||||
CV_CheckTrue(
|
||||
T_s % fastGemmNC(opt) == 0,
|
||||
"pagedAttnQKGemmKernel: T_s should be divisible by the macro tile size"
|
||||
);
|
||||
|
||||
const auto shape_q = shape(Q);
|
||||
CV_CheckTrue(
|
||||
((shape_q.size() == 3) || (shape_q.size() == 4)),
|
||||
"pagedAttnQKGemmKernel: Q must be 3D or 4D (T_q x Nq x D) or (kq_groups x T_q x Nq x D)"
|
||||
);
|
||||
const auto shape_a = shape(A);
|
||||
CV_CheckTrue( shape_a.size() == 4, "pagedAttnQKGemmKernel: A must be 4D (B x N_q x T_q x D)" );
|
||||
|
||||
const int B = shape_a[0];
|
||||
for (size_t s = 0; s < K.size(); s++) {
|
||||
const auto shape_k = shape(K[s]);
|
||||
CV_CheckTrue(
|
||||
shape_k.size() == 3,
|
||||
"pagedAttnQKGemmKernel: each K must be 4D (B x N_k x T_s x D)"
|
||||
);
|
||||
CV_CheckEQ(shape_k[0], B, "pagedAttnQKGemmKernel: the batch size of K should be the same as A");
|
||||
CV_CheckEQ(shape_k[1], N_k, "pagedAttnQKGemmKernel: the number of heads in K should match that of Q");
|
||||
CV_CheckEQ(shape_k[2], D * T_s, "pagedAttnQKGemmKernel: the head dimension of K should match that of Q and A");
|
||||
}
|
||||
|
||||
std::vector<const char*> packed_K;
|
||||
for (size_t s = 0; s < K.size(); s++){
|
||||
packed_K.push_back(K[s].ptr<const char>());
|
||||
}
|
||||
|
||||
char*a = A.ptr<char>();
|
||||
bool isQ3D = shape_q.size() == 3;
|
||||
#if CV_TRY_NEON
|
||||
if (opt.use_neon)
|
||||
opt_NEON::pagedAttnQKGemmKernel(
|
||||
Q.ptr<const char>(), packed_K, a,
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, isQ3D
|
||||
);
|
||||
else
|
||||
#endif
|
||||
#if CV_TRY_AVX2
|
||||
if (opt.use_avx2)
|
||||
opt_AVX2::pagedAttnQKGemmKernel(
|
||||
Q.ptr<const char>(), packed_K, a,
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, isQ3D
|
||||
);
|
||||
else
|
||||
#endif
|
||||
#if CV_TRY_AVX
|
||||
if (opt.use_avx)
|
||||
opt_AVX::pagedAttnQKGemmKernel(
|
||||
Q.ptr<const char>(), packed_K, a,
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, isQ3D
|
||||
);
|
||||
else
|
||||
#endif
|
||||
#if CV_TRY_LASX
|
||||
if (opt.use_lasx)
|
||||
opt_LASX::pagedAttnQKGemmKernel(
|
||||
Q.ptr<const char>(), packed_K, a,
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, isQ3D
|
||||
);
|
||||
else
|
||||
#endif
|
||||
cpu_baseline::pagedAttnQKGemmKernel(
|
||||
Q.ptr<const char>(), packed_K, a,
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, isQ3D
|
||||
);
|
||||
|
||||
|
||||
}
|
||||
|
||||
|
||||
void pagedAttnAVGemm(
|
||||
const Mat& A,const std::vector<Mat> &V, Mat& Out,
|
||||
int T_q, int Nq, int N_k, int T_s, int D,
|
||||
const FastGemmOpt &opt
|
||||
) {
|
||||
size_t esz = A.elemSize();
|
||||
|
||||
for (size_t s = 0; s < V.size(); s++) {
|
||||
CV_CheckTypeEQ(A.type(), V[s].type(), "pagedAttnAVGemmKernel: A and V should have the same type");
|
||||
CV_CheckTrue(esz == V[s].elemSize(), "pagedAttnAVGemmKernel: A and V should have the same element size");
|
||||
}
|
||||
CV_CheckTypeEQ(A.type(), Out.type(), "pagedAttnAVGemmKernel: A and Out should have the same type");
|
||||
|
||||
CV_CheckTrue(
|
||||
T_s % fastGemmKC(opt) == 0,
|
||||
"pagedAttnAVGemmKernel: T_s should be divisible by the macro tile size"
|
||||
);
|
||||
|
||||
const auto shape_a = shape(A);
|
||||
CV_CheckTrue(
|
||||
shape_a.size() == 4,
|
||||
"pagedAttnAVGemmKernel: A must be 4D (B x n_q x T_q x D)"
|
||||
);
|
||||
|
||||
const int B = shape_a[0];
|
||||
for (size_t s = 0; s < V.size(); s++) {
|
||||
const auto shape_v = shape(V[s]);
|
||||
CV_CheckTrue(
|
||||
shape_v.size() == 3,
|
||||
"pagedAttnAVGemmKernel: each V must be 3D (B x N_k x (T_s * D))"
|
||||
);
|
||||
CV_CheckEQ(shape_v[0], B, "pagedAttnAVGemmKernel: the batch size of V should be the same as A");
|
||||
CV_CheckEQ(shape_v[1], N_k, "pagedAttnAVGemmKernel: the number of heads in V should match that of A");
|
||||
CV_CheckEQ(shape_v[2], (int)fastGemmPackBSize(D, T_s, opt),
|
||||
"pagedAttnAVGemmKernel: the last dimension of V-pages should be be equal to packed D x T_s Mat size");
|
||||
}
|
||||
|
||||
std::vector<const char*> packed_V;
|
||||
for (size_t s = 0; s < V.size(); s++){
|
||||
packed_V.push_back(V[s].ptr<const char>());
|
||||
}
|
||||
|
||||
bool canonical_layout = shape(Out).size() == 3;
|
||||
#if CV_TRY_NEON
|
||||
if (opt.use_neon)
|
||||
opt_NEON::pagedAttnAVGemmKernel(
|
||||
A.ptr<const char>(), packed_V, Out.ptr<char>(),
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, canonical_layout, fastGemmPackBSize(D, T_s, opt)
|
||||
);
|
||||
else
|
||||
#endif
|
||||
#if CV_TRY_AVX2
|
||||
if (opt.use_avx2) {
|
||||
opt_AVX2::pagedAttnAVGemmKernel(
|
||||
A.ptr<const char>(), packed_V, Out.ptr<char>(),
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, canonical_layout, fastGemmPackBSize(D, T_s, opt)
|
||||
);
|
||||
}
|
||||
else
|
||||
#endif
|
||||
#if CV_TRY_AVX
|
||||
if (opt.use_avx){
|
||||
opt_AVX::pagedAttnAVGemmKernel(
|
||||
A.ptr<const char>(), packed_V, Out.ptr<char>(),
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, canonical_layout, fastGemmPackBSize(D, T_s, opt)
|
||||
);
|
||||
}
|
||||
else
|
||||
#endif
|
||||
#if CV_TRY_LASX
|
||||
if (opt.use_lasx){
|
||||
opt_LASX::pagedAttnAVGemmKernel(
|
||||
A.ptr<const char>(), packed_V, Out.ptr<char>(),
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, canonical_layout, fastGemmPackBSize(D, T_s, opt)
|
||||
);
|
||||
}
|
||||
else
|
||||
#endif
|
||||
{
|
||||
cpu_baseline::pagedAttnAVGemmKernel(
|
||||
A.ptr<const char>(), packed_V, Out.ptr<char>(),
|
||||
B, T_q, Nq, N_k, T_s, D,
|
||||
esz, canonical_layout, fastGemmPackBSize(D, T_s, opt)
|
||||
);
|
||||
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
|
||||
}} // cv::dnn
|
||||
|
||||
@@ -43,6 +43,7 @@ struct FastGemmOpt {
|
||||
bool all() {
|
||||
return use_avx || use_avx2 || use_neon || use_lasx;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
struct MatMulHelper {
|
||||
@@ -154,6 +155,11 @@ struct MatMulHelper {
|
||||
size_t fastGemmPackBSize(size_t N, size_t K, const FastGemmOpt &opt);
|
||||
|
||||
void fastGemmPackB(const Mat &m, std::vector<float> &packed_B, bool trans, FastGemmOpt &opt);
|
||||
|
||||
int fastGemmMC(const FastGemmOpt &opt);
|
||||
int fastGemmNC(const FastGemmOpt &opt);
|
||||
int fastGemmKC(const FastGemmOpt &opt);
|
||||
|
||||
void fastGemmPackB(bool trans, size_t N, size_t K, const float *B, size_t ldb, float *packed_B, const FastGemmOpt &opt);
|
||||
|
||||
void fastGemm(bool trans_a, int M, int N, int K,
|
||||
@@ -175,7 +181,23 @@ void fastGemmBatch(size_t batch, const size_t *A_offsets, const size_t *B_offset
|
||||
const float *packed_B, float beta, float *C, int ldc, FastGemmOpt &opt);
|
||||
void fastGemmBatch(bool trans_a, bool trans_b, float alpha, const Mat &A,
|
||||
const Mat &B, float beta, Mat &C, FastGemmOpt &opt);
|
||||
void fastGemmBatch(size_t batch,
|
||||
const std::vector<size_t> &A_offsets, const std::vector<size_t> &B_offsets, const std::vector<size_t> &C_offsets,
|
||||
int M, int N, int K, float alpha, const Mat&A, int lda0, int lda1,
|
||||
const Mat&B, int ldb0, int ldb1, float beta, Mat&C, int ldc, FastGemmOpt &opt);
|
||||
|
||||
void pagedAttnQKGemm(
|
||||
const Mat& Q, const std::vector<Mat> &K, Mat& A,
|
||||
int T_q, int Nq, int N_k, int T_s, int D,
|
||||
const FastGemmOpt &opts
|
||||
);
|
||||
void pagedAttnAVGemm(
|
||||
const Mat& A,const std::vector<Mat> &V, Mat& Out,
|
||||
int T_q, int Nq, int N_k, int T_s, int D,
|
||||
const FastGemmOpt &opt
|
||||
);
|
||||
|
||||
|
||||
}} // cv::dnn
|
||||
|
||||
#endif // OPENCV_DNN_FAST_GEMM_HPP
|
||||
#endif // OPENCV_DNN_FAST_GEMM_HPP
|
||||
@@ -78,6 +78,10 @@ namespace cv { namespace dnn { namespace cpu_baseline {
|
||||
|
||||
size_t fastGemmPackBSize(int N, int K);
|
||||
|
||||
int fastGemmMC();
|
||||
int fastGemmNC();
|
||||
int fastGemmKC();
|
||||
|
||||
void fastGemmPackBKernel(const char *B, char *packed_B, size_t N, size_t K, size_t ldb0, size_t ldb1, size_t esz);
|
||||
|
||||
void fastGemmKernel(size_t M, size_t N, size_t K,
|
||||
@@ -94,7 +98,16 @@ void fastGemmBatchKernel(size_t batch, const size_t *A_offsets, const size_t *B_
|
||||
void fastGemmBatchKernel(size_t batch, const size_t *A_offsets, const size_t *B_offsets, const size_t *C_offsets,
|
||||
size_t M, size_t N, size_t K, float alpha, const char *A, size_t lda0, size_t lda1,
|
||||
const char *packed_B, float beta, char *C, size_t ldc, size_t esz);
|
||||
|
||||
void pagedAttnAVGemmKernel(
|
||||
const char* A, const std::vector<const char *> &V, char*Out,
|
||||
size_t B, size_t T_a, size_t Nq, size_t N_k, size_t T_s, size_t D,
|
||||
size_t esz, bool canonical_layout, size_t packed_stride
|
||||
);
|
||||
void pagedAttnQKGemmKernel(
|
||||
const char *Q, const std::vector<const char *> &K, char *A,
|
||||
size_t B, size_t T_q, size_t Nq, size_t N_k, size_t T_s, size_t D,
|
||||
size_t esz, bool isQ3d
|
||||
);
|
||||
FAST_GEMM_IMPLEMENT_PACK(8, _f32, float, float)
|
||||
FAST_GEMM_IMPLEMENT_PACK(12, _f32, float, float)
|
||||
|
||||
@@ -105,6 +118,18 @@ size_t fastGemmPackBSize(int N, int K) {
|
||||
return static_cast<size_t>((N + NC - 1) / NC) * NC * K;
|
||||
}
|
||||
|
||||
int fastGemmMC() {
|
||||
return FAST_GEMM_F32_MC;
|
||||
}
|
||||
|
||||
int fastGemmNC() {
|
||||
return FAST_GEMM_F32_NC;
|
||||
}
|
||||
|
||||
int fastGemmKC() {
|
||||
return FAST_GEMM_F32_PACKED_STRIDE_K;
|
||||
}
|
||||
|
||||
void fastGemmPackBKernel(const char *B, char *packed_B, size_t N, size_t K, size_t ldb0, size_t ldb1, size_t esz) {
|
||||
size_t GEMM_NC = FAST_GEMM_F32_NC, GEMM_NR = FAST_GEMM_F32_NR;
|
||||
size_t NC = (((GEMM_NC < N ? GEMM_NC : N) + GEMM_NR - 1) / GEMM_NR) * GEMM_NR;
|
||||
@@ -462,6 +487,204 @@ void fastGemmBatchKernel(size_t batch, const size_t *A_offsets, const size_t *B_
|
||||
parallel_for_(Range(0, total), fn, nstripes);
|
||||
}
|
||||
|
||||
|
||||
|
||||
// specially for attention
|
||||
// compute QK, supports KV Cache
|
||||
// pages in K are always of shape [B, N_kv, T_s, D]
|
||||
// tot. seq len T = T_s * S
|
||||
void pagedAttnQKGemmKernel(
|
||||
const char *Q, const std::vector<const char *> &K, char *A,
|
||||
size_t B, size_t T_q, size_t Nq, size_t N_k, size_t T_s, size_t D,
|
||||
size_t esz, bool isQ3d
|
||||
) {
|
||||
size_t GEMM_MC = static_cast<size_t>(FAST_GEMM_F32_MC),
|
||||
GEMM_NC = static_cast<size_t>(FAST_GEMM_F32_NC),
|
||||
GEMM_MR = static_cast<size_t>(FAST_GEMM_F32_MR),
|
||||
GEMM_NR = static_cast<size_t>(FAST_GEMM_F32_NR);
|
||||
|
||||
size_t MC = (((GEMM_MC < T_q ? GEMM_MC : T_q) + GEMM_MR - 1) / GEMM_MR) * GEMM_MR;
|
||||
size_t NC = (((GEMM_NC < T_s ? GEMM_NC : T_s) + GEMM_NR - 1) / GEMM_NR) * GEMM_NR;
|
||||
size_t KC = std::min( static_cast<size_t>(FAST_GEMM_F32_PACKED_STRIDE_K), D);
|
||||
|
||||
size_t buff_size = KC * MC * esz;
|
||||
|
||||
const size_t m_tiles = (T_q + MC - 1) / MC;
|
||||
const size_t n_tiles = (T_s + NC - 1) / NC;
|
||||
const size_t tiles_per_mat = m_tiles * n_tiles;
|
||||
const size_t S = K.size();
|
||||
const size_t T = S * T_s;
|
||||
const int n_kq_groups = Nq / N_k;
|
||||
size_t batch = Nq * B;
|
||||
// batch = B x N_q
|
||||
int total = S * batch * tiles_per_mat;
|
||||
|
||||
auto fn = [&](const Range &r) {
|
||||
cv::AutoBuffer<char, FAST_GEMM_MAX_STACKBUF> packed_q_buff;
|
||||
packed_q_buff.allocate(buff_size);
|
||||
char*packed_q = packed_q_buff.data();
|
||||
|
||||
size_t start = r.start;
|
||||
size_t end = r.end;
|
||||
size_t ldq0 = isQ3d ? Nq * D : D;
|
||||
|
||||
for (size_t tile_idx = start; tile_idx < end; tile_idx++) {
|
||||
size_t idx = tile_idx / tiles_per_mat;
|
||||
size_t b = idx / (Nq * S);
|
||||
size_t nq = (idx - b * Nq * S) / S;
|
||||
size_t n_k = nq / n_kq_groups;
|
||||
size_t s = (idx - b * Nq * S) % S;
|
||||
|
||||
const size_t m_tiles_index = (tile_idx - idx * tiles_per_mat) / n_tiles;
|
||||
const size_t n_tiles_index = tile_idx % n_tiles;
|
||||
|
||||
size_t i0 = m_tiles_index * MC;
|
||||
size_t j0 = n_tiles_index * NC;
|
||||
size_t mc = T_q - i0 < MC ? T_q - i0 : MC;
|
||||
size_t nc = T_s - j0 < NC ? T_s - j0 : NC;
|
||||
|
||||
size_t q_offset = b * Nq * T_q * D +
|
||||
(isQ3d ? nq * D : nq * T_q * D);
|
||||
const char *q_block = Q + q_offset * esz;
|
||||
|
||||
size_t k_offset = b * N_k * T_s * D +
|
||||
n_k * T_s * D +
|
||||
j0 * D;
|
||||
const char *k_block = (const char *)K[s] + k_offset * esz;
|
||||
|
||||
// save result to A[b, n_q, : , T_s * s : T_s * (s + 1)]
|
||||
const int a_offset = b * Nq * T_q * T +
|
||||
nq * T_q * T +
|
||||
T_s * s +
|
||||
i0 * T + j0;
|
||||
char* a_block = A + a_offset * esz;
|
||||
|
||||
int _nc = static_cast<int>((nc + GEMM_NR - 1) / GEMM_NR) * GEMM_NR * esz;
|
||||
|
||||
for(int k0 = 0; k0 < D; k0 += KC)
|
||||
{
|
||||
int kc = D - k0 < KC ? D - k0 : KC;
|
||||
// pack q
|
||||
size_t step_q = (i0 * ldq0 + k0) * esz;
|
||||
fast_gemm_pack8_f32(mc, kc, q_block + step_q, ldq0, 1, packed_q);
|
||||
// run kernel
|
||||
fast_gemm_macro_kernel(mc, nc, kc, packed_q, k_block, 1.f, a_block, T, esz);
|
||||
k_block += _nc * kc;
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
int cost_per_thread = static_cast<int>((D / KC) * (MC / GEMM_MR) * (NC / GEMM_NR));
|
||||
double nstripes = (size_t)total * cost_per_thread * (1 / 1024.0);
|
||||
parallel_for_(Range(0, total), fn, nstripes);
|
||||
}
|
||||
|
||||
|
||||
|
||||
// specially for attention
|
||||
// compute QK, supports KV Cache
|
||||
// A is of shape [B, N_q, T_q, T_s * S]
|
||||
// pages in K are always of shape [B, N_kv, T_s, D]
|
||||
// tot. seq len T = T_s * S
|
||||
void pagedAttnAVGemmKernel(
|
||||
const char* A, const std::vector<const char *> &V, char*Out,
|
||||
size_t B, size_t T_a, size_t Nq, size_t N_k, size_t T_s, size_t D,
|
||||
size_t esz, bool canonical_layout, size_t packed_stride
|
||||
) {
|
||||
|
||||
size_t GEMM_MC = static_cast<size_t>(FAST_GEMM_F32_MC),
|
||||
GEMM_NC = static_cast<size_t>(FAST_GEMM_F32_NC),
|
||||
GEMM_MR = static_cast<size_t>(FAST_GEMM_F32_MR),
|
||||
GEMM_NR = static_cast<size_t>(FAST_GEMM_F32_NR);
|
||||
|
||||
const size_t S = V.size();
|
||||
const size_t T = S * T_s;
|
||||
|
||||
size_t MC = (((GEMM_MC < T_a ? GEMM_MC : T_a) + GEMM_MR - 1) / GEMM_MR) * GEMM_MR;
|
||||
size_t NC = (((GEMM_NC < D ? GEMM_NC : D) + GEMM_NR - 1) / GEMM_NR) * GEMM_NR;
|
||||
size_t KC = std::min( static_cast<size_t>(FAST_GEMM_F32_PACKED_STRIDE_K), T);
|
||||
|
||||
size_t buff_size = KC * MC * esz;
|
||||
|
||||
const size_t m_tiles = (T_a + MC - 1) / MC;
|
||||
const size_t n_tiles = (D + NC - 1) / NC;
|
||||
const size_t tiles_per_mat = m_tiles * n_tiles;
|
||||
const int n_kq_groups = Nq / N_k;
|
||||
size_t batch = Nq * B;
|
||||
int total = batch * tiles_per_mat;
|
||||
|
||||
auto fn = [&](const Range &r) {
|
||||
cv::AutoBuffer<char, FAST_GEMM_MAX_STACKBUF> packed_a_buff;
|
||||
packed_a_buff.allocate(buff_size);
|
||||
char* packed_a = packed_a_buff.data();
|
||||
|
||||
size_t start = r.start;
|
||||
size_t end = r.end;
|
||||
size_t ldc0 = canonical_layout ? D : Nq * D;
|
||||
|
||||
for (size_t tile_idx = start; tile_idx < end; tile_idx++) {
|
||||
size_t idx = tile_idx / tiles_per_mat;
|
||||
size_t b = idx / Nq;
|
||||
size_t nq = idx % Nq;
|
||||
size_t n_k = nq / n_kq_groups;
|
||||
|
||||
size_t m_tiles_index = (tile_idx - idx * tiles_per_mat) / n_tiles;
|
||||
size_t n_tiles_index = tile_idx % n_tiles;
|
||||
|
||||
size_t i0 = m_tiles_index * MC;
|
||||
size_t j0 = n_tiles_index * NC;
|
||||
|
||||
size_t mc = T_a - i0 < MC ? T_a - i0 : MC;
|
||||
size_t nc = D - j0 < NC ? D - j0 : NC;
|
||||
|
||||
const size_t a_offset = b * Nq * T_a * T +
|
||||
nq * T_a * T +
|
||||
i0 * T;
|
||||
const char*a_block = A + a_offset * esz;
|
||||
|
||||
// start at the 0th row
|
||||
// column offset is determined according to packed layout
|
||||
size_t v_offset_base = b * N_k * packed_stride +
|
||||
n_k * packed_stride +
|
||||
n_tiles_index * T_s * NC;
|
||||
|
||||
size_t o_offset = b * Nq * T_a * D;
|
||||
if (canonical_layout)
|
||||
o_offset += nq * T_a * D +
|
||||
i0 * D + j0;
|
||||
else
|
||||
o_offset += i0 * Nq * D +
|
||||
nq * D + j0;
|
||||
char* out_block = Out + o_offset * esz;
|
||||
|
||||
int _nc = static_cast<int>((nc + GEMM_NR - 1) / GEMM_NR) * GEMM_NR * esz;
|
||||
|
||||
for(int k0 = 0; k0 < T; k0 += KC)
|
||||
{
|
||||
// T_s should be divisible by KC
|
||||
// pack a
|
||||
fast_gemm_pack8_f32(mc, KC, a_block + k0 * esz, T, 1, packed_a);
|
||||
|
||||
// get the k-block
|
||||
size_t sk = k0 / T_s; // which page
|
||||
size_t k = k0 % T_s;
|
||||
size_t v_offset = v_offset_base * esz + k * _nc;
|
||||
const char *v_block = V[sk] + v_offset;
|
||||
|
||||
// run kernel
|
||||
fast_gemm_macro_kernel(mc, nc, KC, packed_a, v_block, 1.f, out_block, ldc0, esz);
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
int cost_per_thread = static_cast<int>((T / KC) * (MC / GEMM_MR) * (NC / GEMM_NR));
|
||||
double nstripes = (size_t)total * cost_per_thread * (1 / 1024.0);
|
||||
parallel_for_(Range(0, total), fn, nstripes);
|
||||
}
|
||||
|
||||
|
||||
}}} // cv::dnn::cpu_baseline
|
||||
|
||||
#undef FAST_GEMM_STORAGE
|
||||
|
||||
@@ -117,6 +117,10 @@ CV_CPU_OPTIMIZATION_NAMESPACE_BEGIN
|
||||
|
||||
size_t fastGemmPackBSize(int N, int K);
|
||||
|
||||
int fastGemmMC();
|
||||
int fastGemmNC();
|
||||
int fastGemmKC();
|
||||
|
||||
void fastGemmPackBKernel(const char *B, char *packed_B, size_t N, size_t K, size_t ldb0, size_t ldb1, size_t esz);
|
||||
|
||||
void fastGemmKernel(size_t M, size_t N, size_t K,
|
||||
@@ -134,6 +138,18 @@ void fastGemmBatchKernel(size_t batch, const size_t *A_offsets, const size_t *B_
|
||||
size_t M, size_t N, size_t K, float alpha, const char *A, size_t lda0, size_t lda1,
|
||||
const char *packed_B, float beta, char *C, size_t ldc, size_t esz);
|
||||
|
||||
void pagedAttnQKGemmKernel(
|
||||
const char *Q, const std::vector<const char *> &K, char *A,
|
||||
size_t B, size_t T_q, size_t Nq, size_t N_k, size_t T_s, size_t D,
|
||||
size_t esz, bool isQ3d
|
||||
);
|
||||
|
||||
void pagedAttnAVGemmKernel(
|
||||
const char* A, const std::vector<const char *> &K, char*Out,
|
||||
size_t B, size_t T_a, size_t Nq, size_t N_k, size_t T_s, size_t D,
|
||||
size_t esz, bool canonical_layout, size_t packed_stride
|
||||
);
|
||||
|
||||
#ifndef CV_CPU_OPTIMIZATION_DECLARATIONS_ONLY
|
||||
|
||||
/*
|
||||
@@ -538,6 +554,10 @@ size_t fastGemmPackBSize(int N, int K) {
|
||||
return static_cast<size_t>((N + NC - 1) / NC) * NC * K;
|
||||
}
|
||||
|
||||
int fastGemmMC() {return FAST_GEMM_F32_MC;}
|
||||
int fastGemmNC() {return FAST_GEMM_F32_NC;}
|
||||
int fastGemmKC() {return FAST_GEMM_F32_PACKED_STRIDE_K;}
|
||||
|
||||
void fastGemmPackBKernel(const char *B, char *packed_B, size_t N, size_t K, size_t ldb0, size_t ldb1, size_t esz) {
|
||||
size_t GEMM_NC = FAST_GEMM_F32_NC, GEMM_NR = FAST_GEMM_F32_NR;
|
||||
size_t NC = (((GEMM_NC < N ? GEMM_NC : N) + GEMM_NR - 1) / GEMM_NR) * GEMM_NR;
|
||||
@@ -844,13 +864,15 @@ void fastGemmBatchKernel(size_t batch, const size_t *A_offsets, const size_t *B_
|
||||
size_t KC = std::min( static_cast<size_t>(FAST_GEMM_F32_PACKED_STRIDE_K), K);
|
||||
|
||||
size_t buff_size = KC * MC * esz;
|
||||
bool use_stackbuff = buff_size <= FAST_GEMM_MAX_STACKBUF;
|
||||
size_t m_tiles = (M + MC - 1) / MC;
|
||||
size_t n_tiles = (N + NC - 1) / NC;
|
||||
size_t total_tiles = m_tiles * n_tiles;
|
||||
|
||||
auto fn = [&](const Range &r) {
|
||||
char* packed_a = (char*)(use_stackbuff ? alloca(buff_size) : malloc(buff_size));
|
||||
cv::AutoBuffer<char, FAST_GEMM_MAX_STACKBUF> packed_a_buff;
|
||||
packed_a_buff.allocate(buff_size);
|
||||
char*packed_a = packed_a_buff.data();
|
||||
|
||||
const char *packed_b = packed_B;
|
||||
size_t start = r.start;
|
||||
size_t end = r.end;
|
||||
@@ -907,10 +929,6 @@ void fastGemmBatchKernel(size_t batch, const size_t *A_offsets, const size_t *B_
|
||||
packed_b += _nc * kc;
|
||||
}
|
||||
}
|
||||
|
||||
if (!use_stackbuff) {
|
||||
free(packed_a);
|
||||
}
|
||||
};
|
||||
|
||||
int total = batch * total_tiles;
|
||||
@@ -919,6 +937,215 @@ void fastGemmBatchKernel(size_t batch, const size_t *A_offsets, const size_t *B_
|
||||
parallel_for_(Range(0, total), fn, nstripes);
|
||||
}
|
||||
|
||||
|
||||
// specially for attention
|
||||
// compute QK, supports KV Cache
|
||||
// pages in K are always of shape [B, N_kv, T_s, D]
|
||||
// tot. seq len T = T_s * S
|
||||
void pagedAttnQKGemmKernel(
|
||||
const char *Q, const std::vector<const char *> &K, char *A,
|
||||
size_t B, size_t T_q, size_t Nq, size_t N_k, size_t T_s, size_t D,
|
||||
size_t esz, bool isQ3d
|
||||
) {
|
||||
size_t GEMM_MC = static_cast<size_t>(FAST_GEMM_F32_MC),
|
||||
GEMM_NC = static_cast<size_t>(FAST_GEMM_F32_NC),
|
||||
GEMM_MR = static_cast<size_t>(FAST_GEMM_F32_MR),
|
||||
GEMM_NR = static_cast<size_t>(FAST_GEMM_F32_NR);
|
||||
|
||||
size_t MC = (((GEMM_MC < T_q ? GEMM_MC : T_q) + GEMM_MR - 1) / GEMM_MR) * GEMM_MR;
|
||||
size_t NC = (((GEMM_NC < T_s ? GEMM_NC : T_s) + GEMM_NR - 1) / GEMM_NR) * GEMM_NR;
|
||||
size_t KC = std::min( static_cast<size_t>(FAST_GEMM_F32_PACKED_STRIDE_K), D);
|
||||
|
||||
size_t buff_size = KC * MC * esz;
|
||||
|
||||
const size_t m_tiles = (T_q + MC - 1) / MC;
|
||||
const size_t n_tiles = (T_s + NC - 1) / NC;
|
||||
const size_t tiles_per_mat = m_tiles * n_tiles;
|
||||
const size_t S = K.size();
|
||||
const size_t T = S * T_s;
|
||||
const int n_kq_groups = Nq / N_k;
|
||||
size_t batch = Nq * B;
|
||||
// batch = B x N_q
|
||||
int total = S * batch * tiles_per_mat;
|
||||
|
||||
auto fn = [&](const Range &r) {
|
||||
cv::AutoBuffer<char, FAST_GEMM_MAX_STACKBUF> packed_q_buff;
|
||||
packed_q_buff.allocate(buff_size);
|
||||
char*packed_q = packed_q_buff.data();
|
||||
|
||||
size_t start = r.start;
|
||||
size_t end = r.end;
|
||||
size_t ldq0 = isQ3d ? Nq * D : D;
|
||||
|
||||
for (size_t tile_idx = start; tile_idx < end; tile_idx++) {
|
||||
size_t idx = tile_idx / tiles_per_mat;
|
||||
size_t b = idx / (Nq * S);
|
||||
size_t nq = (idx - b * Nq * S) / S;
|
||||
size_t n_k = nq / n_kq_groups;
|
||||
size_t s = (idx - b * Nq * S) % S;
|
||||
|
||||
const size_t m_tiles_index = (tile_idx - idx * tiles_per_mat) / n_tiles;
|
||||
const size_t n_tiles_index = tile_idx % n_tiles;
|
||||
|
||||
size_t i0 = m_tiles_index * MC;
|
||||
size_t j0 = n_tiles_index * NC;
|
||||
size_t mc = T_q - i0 < MC ? T_q - i0 : MC;
|
||||
size_t nc = T_s - j0 < NC ? T_s - j0 : NC;
|
||||
|
||||
size_t q_offset = b * Nq * T_q * D +
|
||||
(isQ3d ? nq * D : nq * T_q * D);
|
||||
const char *q_block = Q + q_offset * esz;
|
||||
|
||||
size_t k_offset = b * N_k * T_s * D +
|
||||
n_k * T_s * D +
|
||||
j0 * D;
|
||||
const char *k_block = (const char *)K[s] + k_offset * esz;
|
||||
|
||||
// save result to A[b, n_q, : , T_s * s : T_s * (s + 1)]
|
||||
const int a_offset = b * Nq * T_q * T +
|
||||
nq * T_q * T +
|
||||
T_s * s +
|
||||
i0 * T + j0;
|
||||
char* a_block = A + a_offset * esz;
|
||||
|
||||
int _nc = static_cast<int>((nc + GEMM_NR - 1) / GEMM_NR) * GEMM_NR * esz;
|
||||
for(int k0 = 0; k0 < D; k0 += KC)
|
||||
{
|
||||
int kc = D - k0 < KC ? D - k0 : KC;
|
||||
// pack q
|
||||
size_t step_q = (i0 * ldq0 + k0) * esz;
|
||||
#if CV_NEON && CV_NEON_AARCH64
|
||||
fast_gemm_pack8_f32(mc, kc, q_block + step_q, ldq0, 1, packed_q);
|
||||
#elif CV_AVX
|
||||
fast_gemm_pack12_f32(mc, kc, q_block + step_q, ldq0, 1, packed_q);
|
||||
#elif CV_LASX
|
||||
fast_gemm_pack12_f32(mc, kc, q_block + step_q, ldq0, 1, packed_q);
|
||||
#elif CV_SIMD128
|
||||
fast_gemm_pack8_f32(mc, kc, q_block + step_q, ldq0, 1, packed_q);
|
||||
#endif
|
||||
// run kernel
|
||||
fast_gemm_macro_kernel(mc, nc, kc, packed_q, k_block, 1.f, a_block, T, esz);
|
||||
k_block += _nc * kc;
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
int cost_per_thread = static_cast<int>((D / KC) * (MC / GEMM_MR) * (NC / GEMM_NR));
|
||||
double nstripes = (size_t)total * cost_per_thread * (1 / 1024.0);
|
||||
parallel_for_(Range(0, total), fn, nstripes);
|
||||
}
|
||||
|
||||
|
||||
|
||||
// specially for attention
|
||||
// compute QK, supports KV Cache
|
||||
// A is of shape [B, N_q, T_q, T_s * S]
|
||||
// pages in K are always of shape [B, N_kv, T_s, D]
|
||||
// tot. seq len T = T_s * S
|
||||
void pagedAttnAVGemmKernel(
|
||||
const char* A, const std::vector<const char *> &V, char*Out,
|
||||
size_t B, size_t T_a, size_t Nq, size_t N_k, size_t T_s, size_t D,
|
||||
size_t esz, bool canonical_layout, size_t packed_stride
|
||||
) {
|
||||
size_t GEMM_MC = static_cast<size_t>(FAST_GEMM_F32_MC),
|
||||
GEMM_NC = static_cast<size_t>(FAST_GEMM_F32_NC),
|
||||
GEMM_MR = static_cast<size_t>(FAST_GEMM_F32_MR),
|
||||
GEMM_NR = static_cast<size_t>(FAST_GEMM_F32_NR);
|
||||
|
||||
const size_t S = V.size();
|
||||
const size_t T = S * T_s;
|
||||
|
||||
size_t MC = (((GEMM_MC < T_a ? GEMM_MC : T_a) + GEMM_MR - 1) / GEMM_MR) * GEMM_MR;
|
||||
size_t NC = (((GEMM_NC < D ? GEMM_NC : D) + GEMM_NR - 1) / GEMM_NR) * GEMM_NR;
|
||||
size_t KC = std::min( static_cast<size_t>(FAST_GEMM_F32_PACKED_STRIDE_K), T);
|
||||
|
||||
size_t buff_size = KC * MC * esz;
|
||||
|
||||
const size_t m_tiles = (T_a + MC - 1) / MC;
|
||||
const size_t n_tiles = (D + NC - 1) / NC;
|
||||
const size_t tiles_per_mat = m_tiles * n_tiles;
|
||||
const int n_kq_groups = Nq / N_k;
|
||||
size_t batch = Nq * B;
|
||||
int total = batch * tiles_per_mat;
|
||||
|
||||
auto fn = [&](const Range &r) {
|
||||
cv::AutoBuffer<char, FAST_GEMM_MAX_STACKBUF> packed_a_buff;
|
||||
packed_a_buff.allocate(buff_size);
|
||||
char* packed_a = packed_a_buff.data();
|
||||
|
||||
size_t start = r.start;
|
||||
size_t end = r.end;
|
||||
size_t ldc0 = canonical_layout ? D : Nq * D;
|
||||
|
||||
for (size_t tile_idx = start; tile_idx < end; tile_idx++) {
|
||||
size_t idx = tile_idx / tiles_per_mat;
|
||||
size_t b = idx / Nq;
|
||||
size_t nq = idx % Nq;
|
||||
size_t n_k = nq / n_kq_groups;
|
||||
|
||||
size_t m_tiles_index = (tile_idx - idx * tiles_per_mat) / n_tiles;
|
||||
size_t n_tiles_index = tile_idx % n_tiles;
|
||||
|
||||
size_t i0 = m_tiles_index * MC;
|
||||
size_t j0 = n_tiles_index * NC;
|
||||
|
||||
size_t mc = T_a - i0 < MC ? T_a - i0 : MC;
|
||||
size_t nc = D - j0 < NC ? D - j0 : NC;
|
||||
|
||||
const size_t a_offset = b * Nq * T_a * T +
|
||||
nq * T_a * T +
|
||||
i0 * T;
|
||||
const char*a_block = A + a_offset * esz;
|
||||
|
||||
// start at the 0th row
|
||||
// column offset is determined according to packed layout
|
||||
size_t v_offset_base = b * N_k * packed_stride +
|
||||
n_k * packed_stride +
|
||||
n_tiles_index * T_s * NC;
|
||||
|
||||
size_t o_offset = b * Nq * T_a * D;
|
||||
if (canonical_layout)
|
||||
o_offset += nq * T_a * D +
|
||||
i0 * D + j0;
|
||||
else
|
||||
o_offset += i0 * Nq * D +
|
||||
nq * D + j0;
|
||||
char* out_block = Out + o_offset * esz;
|
||||
|
||||
int _nc = static_cast<int>((nc + GEMM_NR - 1) / GEMM_NR) * GEMM_NR * esz;
|
||||
|
||||
for(int k0 = 0; k0 < T; k0 += KC)
|
||||
{
|
||||
// get the k-block
|
||||
size_t sk = k0 / T_s; // which page
|
||||
size_t k = k0 % T_s;
|
||||
size_t v_offset = v_offset_base * esz + k * _nc;
|
||||
const char *v_block = V[sk] + v_offset;
|
||||
|
||||
// T_s should be divisible by KC
|
||||
// pack
|
||||
#if CV_NEON && CV_NEON_AARCH64
|
||||
fast_gemm_pack8_f32(mc, KC, a_block + k0 * esz, T, 1, packed_a);
|
||||
#elif CV_AVX
|
||||
fast_gemm_pack12_f32(mc, KC, a_block + k0 * esz, T, 1, packed_a);
|
||||
#elif CV_LASX
|
||||
fast_gemm_pack12_f32(mc, KC, a_block + k0 * esz, T, 1, packed_a);
|
||||
#elif CV_SIMD128
|
||||
fast_gemm_pack8_f32(mc, KC, a_block + k0 * esz, T, 1, packed_a);
|
||||
#endif
|
||||
// run kernel
|
||||
fast_gemm_macro_kernel(mc, nc, KC, packed_a, v_block, 1.f, out_block, ldc0, esz);
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
int cost_per_thread = static_cast<int>((T / KC) * (MC / GEMM_MR) * (NC / GEMM_NR));
|
||||
double nstripes = (size_t)total * cost_per_thread * (1 / 1024.0);
|
||||
parallel_for_(Range(0, total), fn, nstripes);
|
||||
}
|
||||
|
||||
#endif // CV_CPU_OPTIMIZATION_DECLARATIONS_ONLY
|
||||
|
||||
CV_CPU_OPTIMIZATION_NAMESPACE_END
|
||||
|
||||
Reference in New Issue
Block a user