From 5c8d3b60f4a94aa7d4fb74939dada9d39a2e46a1 Mon Sep 17 00:00:00 2001 From: nklskyoy <31723634+nklskyoy@users.noreply.github.com> Date: Fri, 13 Mar 2026 08:23:30 +0100 Subject: [PATCH] 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 --- .../dnn/src/layers/cpu_kernels/fast_gemm.cpp | 298 ++++++++++++++++++ .../dnn/src/layers/cpu_kernels/fast_gemm.hpp | 24 +- .../cpu_kernels/fast_gemm_kernels.default.hpp | 225 ++++++++++++- .../cpu_kernels/fast_gemm_kernels.simd.hpp | 239 +++++++++++++- 4 files changed, 778 insertions(+), 8 deletions(-) diff --git a/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp b/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp index 5aab38dd2e..81d9390100 100644 --- a/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp +++ b/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp @@ -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 &A_offsets, const std::vector &B_offsets, const std::vector &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 char *b = (const char *)B.ptr(); + char *c = (char *)C.ptr(); + +#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 &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 packed_K; + for (size_t s = 0; s < K.size(); s++){ + packed_K.push_back(K[s].ptr()); + } + + char*a = A.ptr(); + bool isQ3D = shape_q.size() == 3; +#if CV_TRY_NEON + if (opt.use_neon) + opt_NEON::pagedAttnQKGemmKernel( + Q.ptr(), 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(), 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(), 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(), packed_K, a, + B, T_q, Nq, N_k, T_s, D, + esz, isQ3D + ); + else +#endif + cpu_baseline::pagedAttnQKGemmKernel( + Q.ptr(), packed_K, a, + B, T_q, Nq, N_k, T_s, D, + esz, isQ3D + ); + + +} + + +void pagedAttnAVGemm( + const Mat& A,const std::vector &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 packed_V; + for (size_t s = 0; s < V.size(); s++){ + packed_V.push_back(V[s].ptr()); + } + + bool canonical_layout = shape(Out).size() == 3; +#if CV_TRY_NEON + if (opt.use_neon) + opt_NEON::pagedAttnAVGemmKernel( + A.ptr(), packed_V, Out.ptr(), + 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(), packed_V, Out.ptr(), + 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(), packed_V, Out.ptr(), + 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(), packed_V, Out.ptr(), + B, T_q, Nq, N_k, T_s, D, + esz, canonical_layout, fastGemmPackBSize(D, T_s, opt) + ); + } + else +#endif + { + cpu_baseline::pagedAttnAVGemmKernel( + A.ptr(), packed_V, Out.ptr(), + B, T_q, Nq, N_k, T_s, D, + esz, canonical_layout, fastGemmPackBSize(D, T_s, opt) + ); + + } + + +} + }} // cv::dnn diff --git a/modules/dnn/src/layers/cpu_kernels/fast_gemm.hpp b/modules/dnn/src/layers/cpu_kernels/fast_gemm.hpp index a207c63c3c..c9bdfa9ad4 100644 --- a/modules/dnn/src/layers/cpu_kernels/fast_gemm.hpp +++ b/modules/dnn/src/layers/cpu_kernels/fast_gemm.hpp @@ -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 &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 &A_offsets, const std::vector &B_offsets, const std::vector &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 &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 &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 \ No newline at end of file diff --git a/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.default.hpp b/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.default.hpp index 980796cf12..5ef81f1d50 100644 --- a/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.default.hpp +++ b/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.default.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 &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 &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((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 &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(FAST_GEMM_F32_MC), + GEMM_NC = static_cast(FAST_GEMM_F32_NC), + GEMM_MR = static_cast(FAST_GEMM_F32_MR), + GEMM_NR = static_cast(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(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 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((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((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 &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(FAST_GEMM_F32_MC), + GEMM_NC = static_cast(FAST_GEMM_F32_NC), + GEMM_MR = static_cast(FAST_GEMM_F32_MR), + GEMM_NR = static_cast(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(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 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((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((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 diff --git a/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.simd.hpp b/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.simd.hpp index 98a3095bf6..b9ce5a9113 100644 --- a/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.simd.hpp +++ b/modules/dnn/src/layers/cpu_kernels/fast_gemm_kernels.simd.hpp @@ -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 &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 &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((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(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 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 &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(FAST_GEMM_F32_MC), + GEMM_NC = static_cast(FAST_GEMM_F32_NC), + GEMM_MR = static_cast(FAST_GEMM_F32_MR), + GEMM_NR = static_cast(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(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 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((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((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 &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(FAST_GEMM_F32_MC), + GEMM_NC = static_cast(FAST_GEMM_F32_NC), + GEMM_MR = static_cast(FAST_GEMM_F32_MR), + GEMM_NR = static_cast(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(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 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((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((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