Merge pull request #29127 from Prasadayus:KV-cache

Add KV cache with paged attention and prefetch #29127

Closes: https://github.com/opencv/opencv/issues/27159

### 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

modules/dnn/src/kv_cache_manager.cpp

@@ -48,6 +48,7 @@ void setKVCacheManager(Ptr<Net::Impl> netimpl)
     manager.netimpl = netimpl;
     manager.opt.init();
     initKVDataRecursively(netimpl->mainGraph, manager.kData, manager.vData, manager.opt);
+    manager.buildRoutes();
 
     manager.isInitialized = true;
     netimpl->useKVCache = true;
@@ -61,6 +62,78 @@ void KVCacheManager::init()
     CV_Assert(isInitialized);
 }
 
+void KVCacheManager::buildRoutes()
+{
+    presentToPastRoutes.clear();
+    hasRoutes = false;
+
+    if (!netimpl || !netimpl->mainGraph)
+        return;
+
+    const std::vector<Arg>& gr_outputs = netimpl->mainGraph->outputs();
+    for (const Arg& out_arg : gr_outputs)
+    {
+        const ArgData& out_adata = netimpl->args.at(out_arg.idx);
+        const std::string& out_name = out_adata.name;
+
+        if (out_name.compare(0, 8, "present.") != 0)
+            continue;
+
+        std::string past_name = "past_key_values." + out_name.substr(8);
+        auto it = netimpl->argnames.find(past_name);
+        if (it == netimpl->argnames.end())
+            continue;
+
+        Arg past_arg((int)it->second);
+        if (netimpl->args.at(past_arg.idx).kind != DNN_ARG_INPUT)
+            continue;
+
+        presentToPastRoutes.emplace_back(out_arg.idx, past_arg.idx);
+    }
+
+    hasRoutes = !presentToPastRoutes.empty();
+    if (hasRoutes)
+        initPastTensors();
+}
+
+void KVCacheManager::initPastTensors()
+{
+    for (const auto& route : presentToPastRoutes)
+    {
+        const ArgData& past_adata = netimpl->args.at(route.second);
+        const MatShape& decl_shape = past_adata.shape;
+        if (decl_shape.dims <= 0)
+            continue;
+
+        // Replace symbolic dims (stored as <=0): batch (dim 0) -> 1, all others -> 0 for empty-sequence state.
+        std::vector<int> shape_vec(decl_shape.dims);
+        for (int d = 0; d < decl_shape.dims; d++)
+            shape_vec[d] = (decl_shape[d] > 0) ? decl_shape[d] : (d == 0 ? 1 : 0);
+
+        int dtype = past_adata.type;
+        if (dtype < 0)
+            dtype = CV_32F;
+
+        Mat& past_t = netimpl->__tensors__.at(route.second);
+        past_t = Mat(shape_vec, dtype, Scalar(0));
+        netimpl->finalizeLayers = true;
+    }
+}
+
+void KVCacheManager::applyRoutes()
+{
+    for (const auto& route : presentToPastRoutes)
+    {
+        const Mat& present_t = netimpl->argTensor(Arg(route.first));
+        Mat& past_t = netimpl->__tensors__.at(route.second);
+        if (present_t.empty())
+            continue;
+        if (past_t.shape() != present_t.shape() || past_t.type() != present_t.type())
+            netimpl->finalizeLayers = true;
+        present_t.copyTo(past_t);
+    }
+}
+
 void KVCache::grow(const Mat& newData) {
     CV_Assert(newData.dims == 4 || newData.dims == 3);
 
@@ -191,21 +264,16 @@ void KCache::growGenerate(const Mat& newData){
     auto* page = pages[cur_page].ptr<float>();
     const auto* data = newData.ptr<float>();
 
-    for (int b = 0; b < batch_size; b++){
-        for (int h = 0; h < nHeads; h++){
-            for(int j = 0; j < headDim; j++) {
-                int step =
-                    b * nHeads * headDim +
-                    h * headDim +
-                    j;
-                page[
-                    b * nHeads * Ps +
-                    h * Ps +
-                    t0 + pageSize * j
-                ] = *(data + step);
-            }
+    const int nstripes = batch_size * nHeads;
+    parallel_for_(Range(0, nstripes), [&](const Range& r) {
+        for (int i = r.start; i < r.end; i++) {
+            int b = i / nHeads, h = i % nHeads;
+            const float* src = data + (b * nHeads + h) * headDim;
+            float* dst = page + b * nHeads * Ps + h * Ps + t0;
+            for (int j = 0; j < headDim; j++)
+                dst[j * pageSize] = src[j];
         }
-    }
+    });
 
     nTokens += 1;
 }
@@ -225,8 +293,10 @@ void VCache::growGenerate(const Mat& newData){
     auto* page = pages[cur_page].ptr<float>();
     const auto* data = newData.ptr<float>();
 
-    for (int b = 0; b < batch_size; b++){
-        for (int h = 0; h < nHeads; h++){
+    const int nstripes = batch_size * nHeads;
+    parallel_for_(Range(0, nstripes), [&](const Range& r) {
+        for (int i = r.start; i < r.end; i++) {
+            int b = i / nHeads, h = i % nHeads;
             for (int j = 0; j <= (headDim - 1) / Nr; j++) {
                 int step = b * nHeads * headDim + h * headDim + j * Nr;
                 int copy_size = std::min(Nr, headDim - j * Nr);
@@ -242,11 +312,10 @@ void VCache::growGenerate(const Mat& newData){
                 }
             }
         }
-    }
+    });
 
     nTokens += 1;
 }
 
-
 CV__DNN_INLINE_NS_END
 }}

modules/dnn/src/kv_cache_manager.hpp

@@ -89,7 +89,14 @@ struct KVCacheManager
     FastGemmOpt opt;
     bool isInitialized = false;
 
+    // present.* output arg idx -> past_key_values.* input arg idx
+    std::vector<std::pair<int, int>> presentToPastRoutes;
+    bool hasRoutes = false;
+
     void init();
+    void buildRoutes();
+    void applyRoutes();
+    void initPastTensors();
 };
 
 void setKVCacheManager(Ptr<Net::Impl> netimpl);

modules/dnn/src/layers/gemm_layer.cpp

@@ -255,7 +255,7 @@ public:
         LayerGemmOpMode mode = getOpMode(inputs.size(), blobs.size());
 
         // pack B if it is const
-        if (constB(mode)) {
+        if (constB(mode) && blobs[0].data != last_packed_blob_data) {
             fastGemmPackB(blobs[0], packed_B, trans_b, opt);
 
             // Pre-pack B in the "thin" layout when the gemm shape has a
@@ -304,6 +304,7 @@ public:
                 }
             }
 #endif
+            last_packed_blob_data = blobs[0].data;
         }
 
         if (constC(mode) && flatten_a) {
@@ -578,6 +579,7 @@ private:
     std::vector<float> broadcast_C;
     int real_ndims_C;
     FastGemmOpt opt;
+    const uchar* last_packed_blob_data = nullptr;
 };
 
 Ptr<GemmLayer> GemmLayer::create(const LayerParams& params) {

modules/dnn/src/layers/matmul_layer.cpp

@@ -154,20 +154,24 @@ class MatMulLayerImpl CV_FINAL : public MatMulLayer {
                    C_shape = shape(outputs[0]);
         helper.compute(trans_a, trans_b, A_shape, B_shape, C_shape);
 
-        if (!blobs.empty()) {
-            fastGemmPackB(blobs[0], packed_input_B, trans_b, opt);
+        // Pack only 2D weight matrices; skip higher-dim tensors (e.g. Q@K^T in attention).
+        const Mat* B_mat = !blobs.empty() ? &blobs[0] :
+                           (inputs.size() >= 2 && inputs[1].dims == 2 ? &inputs[1] : nullptr);
+        if (B_mat && B_mat->data != last_packed_input_B_data) {
+            fastGemmPackB(*B_mat, packed_input_B, trans_b, opt);
             helper.updatePackedBOffsets(packed_input_B.size());
 
-            if (helper.batch == 1 && blobs[0].type() == CV_32F &&
+            if (helper.batch == 1 && B_mat->type() == CV_32F &&
                 fastGemmThinEligible(helper.M, helper.N, helper.K)) {
                 thin_packed_B.resize(fastGemmThinPackBSize(helper.N, helper.K));
                 fastGemmThinPackB(helper.N, helper.K,
-                                  blobs[0].ptr<const float>(),
+                                  B_mat->ptr<const float>(),
                                   (size_t)helper.ldb0, (size_t)helper.ldb1,
                                   thin_packed_B.data());
             } else {
                 thin_packed_B.clear();
             }
+            last_packed_input_B_data = B_mat->data;
         }
 
         // broadcast bias if needed
@@ -289,10 +293,17 @@ class MatMulLayerImpl CV_FINAL : public MatMulLayer {
                          a, helper.lda0, helper.lda1,
                          thin_packed_B.data(), beta,
                          y, helper.ldc, opt.multi_thread);
-        } else {
+        } else if (!packed_input_B.empty()) {
             fastGemmBatch(helper.batch, helper.A_offsets.data(), helper.packed_B_offsets.data(), helper.C_offsets.data(),
                           helper.M, helper.N, helper.K, alpha, a, helper.lda0, helper.lda1,
                           packed_input_B.data(), beta, y, helper.ldc, opt);
+        } else {
+            // truly dynamic B (changes every call — no packing cache available)
+            const auto &B = inputs[1];
+            const auto *b = B.ptr<const float>();
+            fastGemmBatch(helper.batch, helper.A_offsets.data(), helper.B_offsets.data(), helper.C_offsets.data(),
+                          helper.M, helper.N, helper.K, alpha, a, helper.lda0, helper.lda1,
+                          b, helper.ldb0, helper.ldb1, beta, y, helper.ldc, opt);
         }
     }
 
@@ -523,6 +534,8 @@ class MatMulLayerImpl CV_FINAL : public MatMulLayer {
     std::vector<float> thin_packed_B;
     Mat broadcast_bias;
 
+    const uchar* last_packed_input_B_data = nullptr;
+
     FastGemmOpt opt;
     MatMulHelper helper;
 };

modules/dnn/src/net_impl2.cpp

@@ -684,6 +684,10 @@ void Net::Impl::forwardMainGraph(InputArrayOfArrays inputs, OutputArrayOfArrays
     // [TODO] if a target or backend change or there are some other important
     // global changes in configuration, finalizeLayers should be set to 'true' again
     finalizeLayers = false;
+
+    // Feed present.* outputs back as past_key_values.* inputs for the next step (causal-lm-with-past).
+    if (useKVCache && kvCacheManager.hasRoutes)
+        kvCacheManager.applyRoutes();
 }
 
 void Net::Impl::forwardWithSingleOutput(const std::string& outname, OutputArrayOfArrays outputBlobs)

modules/dnn/test/test_layers.cpp

@@ -3188,4 +3188,4 @@ TEST(Layer_Test_Softmax, NoNaN_AllNegInf)
     }
 }
 
-}} // namespace
+}} // namespace