Merge pull request #30405 from Anemptyship:leader/attention-memory

dnn: avoid unused quadratic workspace for MLAS attention - #30405

`AttentionLayerImpl` allocates a dense `[batch * heads, sequence, sequence]` probability buffer before calling MLAS FlashAttention, even though that path never uses it. With 12 heads and 8,192 tokens, this is an unused 3 GiB allocation that can make inference fail before the tiled kernel runs.

This patch shares the existing fast-path eligibility check between memory planning and execution. It keeps one-element placeholders for the unused probability and output workspaces, preserving the internal buffer indices used by RoPE. If MLAS declines the call, the generic path allocates the dense buffers it needs. Masked attention and builds without MLAS keep their existing workspace plans. The numerical kernels are unchanged.

I reproduced the allocation with a generated FP32 ONNX Attention model loaded through `readNetFromONNX`, using the 5.x graph engine. For batch 1, 12 heads, head width 16 and 8,192 tokens:

| Measurement | Before | After |
| --- | ---: | ---: |
| Largest `cv::Mat` allocation | 3 GiB | 18 MiB |
| Peak live `cv::Mat` storage | 3,279,180,288 bytes | 51,663,368 bytes |
| Forward with a 512 MiB address-space limit | Allocation failure | Completes |
| Peak RSS without the limit | 65,664 KiB | 65,616 KiB |

The allocation trace goes through `allocateLayerOutputs` and `Mat::fit` before the layer's `forward`. The same failure is resolved at 4,096 and 16,384 tokens; at 16,384 tokens, the largest allocation drops from 12 GiB to 36 MiB. Outputs from the constrained runs match the unrestricted runs byte for byte.

The limit above is Linux `RLIMIT_AS`, which limits virtual address space. RSS barely changes because the unused buffer is not touched. The allocation measurements cover `cv::Mat` storage only.

Validation on Linux x86-64:

- 40 targeted tests pass with MLAS in Release, and the same 40 pass under ASan/UBSan with MLAS and SIMD disabled. The workspace regression test fails against the baseline library.
- 56 ONNX baseline/patch comparisons are byte-identical, covering sequence lengths, 1/6 threads, static and dynamic weights, masks, and RoPE. Four of these force `mlasFlashAttention` to return false after dispatch to exercise fallback from a compact workspace plan.
- Reusing a graph across changing sequence lengths passes 40 comparisons against fresh graphs on each version. A separate layer-level matrix adds 144 byte-identical comparisons.
- The existing Attention perf fixture passes all five cases at 1 and 6 threads with `--perf_verify_sanity=true` and no regression data files. This includes new constant-weight cases at 257, 1,024, and 4,096 tokens. Separate longer pinned timing runs measured approximately +0.5% at 257 tokens and -1.1% at 1,024 tokens; timings vary on this shared host, so this PR makes no speedup claim.

The half-input audit passed six configurations that the graph promotes to FP32. Static, unmasked FP16 constants exposed a separate existing weight-packing issue: the baseline already produces invalid outputs, and ASan reports an out-of-bounds read in unchanged packing code. Those cases are excluded from passing validation. ARM and other platforms still need CI.

This follows the MLAS integration in #29126. #30126 addresses a different source of memory use (weight ownership), but touches the same class. `AttentionOnnxAiLayerImpl` is outside this change.

### Pull Request Readiness Checklist

- [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 code under GPL or another license incompatible with OpenCV.
- [x] The PR targets the proper branch (`5.x`).
- [x] Related work is referenced above; the reproducer and measurements describe the bug.
- [x] Accuracy and performance tests are included. They generate their inputs and require no new `opencv_extra` data.
- [x] No public API changes or new samples requiring documentation.

modules/dnn/perf/perf_layer.cpp

@@ -799,6 +799,18 @@ PERF_TEST_P_(Layer_Attention, BERT_ConstWeights) {
     test_layer({1, 128, 768}, {768, 768, 768}, 12, /*const_weights=*/true);
 }
 
+PERF_TEST_P_(Layer_Attention, LongSequence_257) {
+    test_layer({1, 257, 192}, {192, 192, 192}, 12, /*const_weights=*/true);
+}
+
+PERF_TEST_P_(Layer_Attention, LongSequence_1024) {
+    test_layer({1, 1024, 192}, {192, 192, 192}, 12, /*const_weights=*/true);
+}
+
+PERF_TEST_P_(Layer_Attention, LongSequence_4096) {
+    test_layer({1, 4096, 192}, {192, 192, 192}, 12, /*const_weights=*/true);
+}
+
 // GroupQueryAttention runs on AttentionOnnxAi; these flags are what the ONNX importer sets.
 struct Layer_GroupQueryAttention : public TestBaseWithParam<tuple<Backend, Target>> {
     void test_layer(int B, int S, int Sp, int num_heads, int kv_num_heads, int D) {

modules/dnn/src/layers/attention_layer.cpp

@@ -218,6 +218,13 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
         MatShape gemm_buffer_shape{batch_size_, seq_len_, hidden_size_},
                  attention_prob_shape{batch_size_ * num_heads_, seq_len_, seq_len_},
                  output_buffer_shape{batch_size_ * num_heads_, seq_len_, v_head_size_};
+        if (useFlashAttention(inputs.size(), batch_size_, seq_len_, v_head_size_))
+        {
+            // FlashAttention writes directly to the output and never uses either
+            // dense workspace. Keep their slots so the RoPE table index is stable.
+            attention_prob_shape = MatShape{1};
+            output_buffer_shape = MatShape{1};
+        }
         internals.assign(1, gemm_buffer_shape);
         internals.push_back(attention_prob_shape);
         internals.push_back(output_buffer_shape);
@@ -499,13 +506,7 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
             }
         }
 
-        const int num_non_blob_inputs = (int)inputs.size();
-        const bool has_mask = (!blobs.empty() && num_non_blob_inputs >= 2) ||
-                                ( blobs.empty() && num_non_blob_inputs >= 4);
-
-        if (mlasAvailable() && !has_mask &&
-            batch_size > 0 && num_heads > 0 && seq_len > 0 &&
-            qkv_head_sizes[0] > 0 && qkv_head_sizes[2] > 0)
+        if (useFlashAttention(inputs.size(), (int)batch_size, (int)seq_len, (int)qkv_head_sizes[2]))
         {
             const int B   = (int)batch_size;
             const int H   = (int)num_heads;
@@ -533,6 +534,8 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
 
         // Compute Softmax(scale * MatMul(Q, K))
         auto &attention_prob = internals[1];
+        // getMemoryShapes() plans dense buffers whenever FlashAttention is not used
+        CV_CheckEQ(attention_prob.total(), (size_t)(batch_size * num_heads * seq_len * seq_len), "");
         {
             auto *output = attention_prob.ptr<float>();
 
@@ -631,6 +634,7 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
 
         // output = attention_prob @ V, then transpose back to [B, S, H*D]
         auto &output_buffer = internals[2];
+        CV_CheckEQ(output_buffer.total(), (size_t)(batch_size * num_heads * seq_len * qkv_head_sizes[2]), "");
         {
             auto *output = outputs[0].ptr<float>();
             auto *output_buff = output_buffer.ptr<float>();
@@ -674,6 +678,14 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
     }
 
  private:
+    bool useFlashAttention(size_t inputCount, int batch, int sequence, int valueHeadSize) const
+    {
+        const bool hasMask = (!blobs.empty() && inputCount >= 2) ||
+                             (blobs.empty() && inputCount >= 4);
+        return mlasAvailable() && !hasMask && batch > 0 && num_heads > 0 && sequence > 0 &&
+               qkv_head_sizes[0] > 0 && valueHeadSize > 0;
+    }
+
     size_t num_heads;
     std::vector<size_t> qkv_hidden_sizes; // order: {qk_hidden_size, qk_hidden_size, v_hidden_size}
     float scale;

modules/dnn/test/test_attention_workspace.cpp

@@ -0,0 +1,105 @@
+// 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.
+
+#include "test_precomp.hpp"
+#include <opencv2/dnn/all_layers.hpp>
+#include <opencv2/dnn/shape_utils.hpp>
+#include <cstring>
+
+namespace opencv_test { namespace {
+
+static LayerParams attentionWorkspaceParams(int valueHidden, bool rotary, Mat& weights, Mat& bias)
+{
+    const int channels = 32;
+    int hidden[] = {channels, channels, valueHidden};
+    LayerParams params;
+    params.set("num_heads", 2);
+    params.set("qkv_hidden_sizes", DictValue::arrayInt(hidden, 3));
+    params.set("do_rotary", rotary);
+    weights.create(channels, 2 * channels + valueHidden, CV_32F);
+    bias.create(MatShape{2 * channels + valueHidden}, CV_32F);
+    theRNG().fill(weights, RNG::UNIFORM, -.1f, .1f);
+    theRNG().fill(bias, RNG::UNIFORM, -.1f, .1f);
+    return params;
+}
+
+typedef testing::TestWithParam<tuple<bool, bool, bool> > Layer_AttentionWorkspace;
+
+TEST_P(Layer_AttentionWorkspace, memoryPlan)
+{
+    const bool dynamicWeights = get<0>(GetParam()), hasMask = get<1>(GetParam()), rotary = get<2>(GetParam());
+    Mat weights, bias;
+    LayerParams params = attentionWorkspaceParams(48, rotary, weights, bias);
+    if (!dynamicWeights)
+        params.blobs = {weights, bias};
+    Ptr<Layer> layer = AttentionLayer::create(params);
+#ifdef HAVE_MLAS
+    size_t previous = 0;
+#endif
+    for (int seq : {1024, 2048})
+    {
+        std::vector<MatShape> inputs{MatShape{2, seq, 32}}, outputs, internals;
+        if (dynamicWeights)
+        {
+            inputs.push_back(weights.shape());
+            inputs.push_back(bias.shape());
+        }
+        if (hasMask)
+            inputs.push_back(MatShape{seq, seq});
+        layer->getMemoryShapes(inputs, 1, outputs, internals);
+        ASSERT_EQ(size_t(rotary ? 4 : 3), internals.size());
+        size_t elements = 0;
+        for (const MatShape& shape : internals)
+            elements += total(shape);
+#ifdef HAVE_MLAS
+        if (!hasMask)
+        {
+            // Linear workspace growth, including QKV and optional RoPE.
+            EXPECT_LT(elements, size_t(2 * seq * 128));
+            if (previous)
+            {
+                EXPECT_LE(elements, 2 * previous);
+            }
+            previous = elements;
+        }
+        else
+#endif
+        {
+            EXPECT_EQ(MatShape({4, seq, seq}), internals[1]);
+            EXPECT_EQ(MatShape({4, seq, 24}), internals[2]);
+        }
+    }
+}
+
+typedef testing::TestWithParam<bool> Layer_AttentionWorkspaceShapes;
+
+// the same net runs inputs of different lengths one after another
+TEST_P(Layer_AttentionWorkspaceShapes, changingInputShape)
+{
+    const bool rotary = GetParam();
+    Mat weights, bias;
+    LayerParams params = attentionWorkspaceParams(48, rotary, weights, bias);
+    params.blobs = {weights, bias};
+    Net net;
+    net.addLayerToPrev("attention", "Attention", CV_32F, params);
+    for (int seq : {17, 257, 1, 17})
+    {
+        SCOPED_TRACE(testing::Message() << "sequence=" << seq);
+        Mat input(MatShape{2, seq, 32}, CV_32F);
+        theRNG().fill(input, RNG::UNIFORM, -1.f, 1.f);
+        std::vector<Mat> inputs{input};
+        std::vector<Mat> expected;
+        runLayer(AttentionLayer::create(params), inputs, expected);
+        net.setInput(input);
+        Mat actual = net.forward();
+        ASSERT_EQ(expected[0].shape(), actual.shape());
+        EXPECT_EQ(0, std::memcmp(expected[0].data, actual.data, expected[0].total() * expected[0].elemSize()));
+    }
+}
+
+INSTANTIATE_TEST_CASE_P(/**/, Layer_AttentionWorkspace, testing::Combine(
+    testing::Bool(), testing::Bool(), testing::Bool()));
+INSTANTIATE_TEST_CASE_P(/**/, Layer_AttentionWorkspaceShapes, testing::Bool());
+
+}} // namespace