## Coverage Report

#### Created: 2026-05-04 15:30

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/liu/actions-runner/_work/ccv/ccv/lib/nnc/cmd/scaled_dot_product_attention/ccv_nnc_scaled_dot_product_attention_cpu_ref.c

    Line
    Count
    Source
    1
    #include "ccv.h"
    2
    #include "ccv_internal.h"
    3
    #include "nnc/ccv_nnc.h"
    4
    #include "nnc/ccv_nnc_easy.h"
    5
    #include "nnc/ccv_nnc_internal.h"
    6
    #ifdef USE_OPENMP
    7
    #include <omp.h>
    8
    #endif
    9
    #ifdef USE_DISPATCH
    10
    #include <dispatch/dispatch.h>
    11
    #endif
    12
    13
    // Shared methods.
    14
    #include "../_ccv_nnc_cpu_ref.h"
    15
    16
    static int _ccv_nnc_scaled_dot_product_attention_forw(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)
    17
    11
    {
    18
    11
      assert(input_size >= 3);
    19
    11
      assert(output_size >= 1);
    20
    11
      const int is_varlen = cmd.info.scaled_dot_product_attention.is_varlen;
    21
    11
      ccv_nnc_tensor_view_t* const q = (ccv_nnc_tensor_view_t*)inputs[0];
    22
    11
      ccv_nnc_tensor_view_t* const k = (ccv_nnc_tensor_view_t*)inputs[1];
    23
    11
      ccv_nnc_tensor_view_t* const v = (ccv_nnc_tensor_view_t*)inputs[2];
    24
    11
      ccv_nnc_tensor_view_t* const attn_mask = input_size > 3 ? (ccv_nnc_tensor_view_t*)inputs[3]8 : 03;
    25
    11
      ccv_nnc_tensor_view_t* const w = input_size > 4 ? (ccv_nnc_tensor_view_t*)inputs[4]7 : 04;
    26
    11
      ccv_nnc_tensor_view_t* const bias = input_size > 5 ? (ccv_nnc_tensor_view_t*)inputs[5]7 : 04;
    27
    11
      ccv_nnc_tensor_view_t* const q_seq_offsets = is_varlen && input_size > 62 ? (ccv_nnc_tensor_view_t*)inputs[6]2 : 09;
    28
    11
      ccv_nnc_tensor_view_t* const kv_seq_offsets = is_varlen && input_size > 72 ? (ccv_nnc_tensor_view_t*)inputs[7]2 : 09;
    29
    11
      if (bias) // bias always requires a weight matrix.
    30
    3
        { assert(w); }
    31
    11
      if (is_varlen && (2attn_mask2 || w2 || bias2 || !q_seq_offsets2 || !kv_seq_offsets2))
    32
    0
        return CCV_NNC_EXEC_INVALID;
    33
    11
      ccv_nnc_tensor_view_t* const c = (w) ? (ccv_nnc_tensor_view_t*)outputs[2]3 : (ccv_nnc_tensor_view_t*)outputs[0]8;
    34
    11
      const int q_nd = ccv_nnc_tensor_nd(q->info.dim);
    35
    11
      assert(q_nd == 3 || q_nd == 4);
    36
    11
      const int k_nd = ccv_nnc_tensor_nd(k->info.dim);
    37
    11
      assert(k_nd == 3 || k_nd == 4);
    38
    11
      const int v_nd = ccv_nnc_tensor_nd(v->info.dim);
    39
    11
      assert(v_nd == 3 || v_nd == 4);
    40
    11
      const int c_nd = ccv_nnc_tensor_nd(c->info.dim);
    41
    11
      assert(c_nd == 3 || c_nd == 4);
    42
    11
      assert(q_nd == k_nd && k_nd == v_nd && v_nd == c_nd);
    43
    11
      if (is_varlen && q_nd != 42)
    44
    0
        return CCV_NNC_EXEC_INVALID;
    45
      // Assuming this is float 32.
    46
    11
      int qdim[CCV_NNC_MAX_DIM_ALLOC];
    47
    11
      int kdim[CCV_NNC_MAX_DIM_ALLOC];
    48
    11
      int vdim[CCV_NNC_MAX_DIM_ALLOC];
    49
    11
      int cdim[CCV_NNC_MAX_DIM_ALLOC];
    50
    11
      int amdim[CCV_NNC_MAX_DIM_ALLOC];
    51
    11
      ccv_nnc_tensor_view_get_dim(q, qdim);
    52
    11
      ccv_nnc_tensor_view_get_dim(k, kdim);
    53
    11
      ccv_nnc_tensor_view_get_dim(v, vdim);
    54
    11
      ccv_nnc_tensor_view_get_dim(c, cdim);
    55
    11
      if (is_varlen)
    56
    2
      {
    57
    2
        assert(q_seq_offsets->info.datatype == CCV_32S);
    58
    2
        assert(kv_seq_offsets->info.datatype == CCV_32S);
    59
    2
        assert(CCV_IS_TENSOR_CONTIGUOUS(q_seq_offsets));
    60
    2
        assert(CCV_IS_TENSOR_CONTIGUOUS(kv_seq_offsets));
    61
    2
      }
    62
    11
      if (q_nd == 3)
    63
    0
      {
    64
    0
        qdim[0] = qdim[1], qdim[1] = qdim[2], qdim[2] = 1;
    65
    0
        kdim[0] = kdim[1], kdim[1] = kdim[2], kdim[2] = 1;
    66
    0
        vdim[0] = vdim[1], vdim[1] = vdim[2], vdim[2] = 1;
    67
    0
        cdim[0] = cdim[1], cdim[1] = cdim[2], cdim[2] = 1;
    68
    0
      }
    69
    11
      assert(qdim[0] == kdim[0] && kdim[0] == vdim[0] && vdim[0] == cdim[0]);
    70
    11
      assert(qdim[2] == cdim[2]);
    71
    11
      assert(kdim[2] == vdim[2]);
    72
    11
      assert(qdim[2] % kdim[2] == 0);
    73
    11
      assert(qdim[2] >= kdim[2]);
    74
    11
      assert(qdim[3] == kdim[3]);
    75
    11
      assert(kdim[1] == vdim[1]);
    76
    11
      assert(cdim[1] == qdim[1]);
    77
    11
      assert(cdim[3] == vdim[3]);
    78
    11
      assert(CCV_NNC_MAX_DIM == 2); // Need to change this logic for CCV_NNC_MAX_DIM == other number.
    79
    11
      int qstride[CCV_NNC_MAX_DIM_ALLOC];
    80
    11
      int kstride[CCV_NNC_MAX_DIM_ALLOC];
    81
    11
      int vstride[CCV_NNC_MAX_DIM_ALLOC];
    82
    11
      int cstride[CCV_NNC_MAX_DIM_ALLOC];
    83
    11
      int amstride[CCV_NNC_MAX_DIM_ALLOC];
    84
    11
      ccv_nnc_tensor_view_get_stride(q, qstride);
    85
    11
      ccv_nnc_tensor_view_get_stride(k, kstride);
    86
    11
      ccv_nnc_tensor_view_get_stride(v, vstride);
    87
    11
      ccv_nnc_tensor_view_get_stride(c, cstride);
    88
    11
      if (q_nd == 3)
    89
    0
      {
    90
    0
        qstride[0] = qstride[1], qstride[1] = qstride[2], qstride[2] = qstride[3];
    91
    0
        kstride[0] = kstride[1], kstride[1] = kstride[2], kstride[2] = kstride[3];
    92
    0
        vstride[0] = vstride[1], vstride[1] = vstride[2], vstride[2] = vstride[3];
    93
    0
        cstride[0] = cstride[1], cstride[1] = cstride[2], cstride[2] = cstride[3];
    94
    0
      }
    95
    11
      if (attn_mask)
    96
    2
      {
    97
    2
        ccv_nnc_tensor_view_get_dim(attn_mask, amdim);
    98
    2
        ccv_nnc_tensor_view_get_stride(attn_mask, amstride);
    99
    2
        assert(amdim[0] == qdim[0] || amdim[0] == 1);
    100
    2
        assert(amdim[1] == qdim[2] || amdim[1] == 1);
    101
    2
        assert(amdim[2] == qdim[1]);
    102
    2
        assert(amdim[3] == kdim[1]);
    103
    2
      }
    104
    11
      int i[CCV_NNC_MAX_DIM + 2];
    105
    11
      float* qk = ccv_nnc_stream_context_get_workspace(stream_context, sizeof(float) * qdim[1] * kdim[1], CCV_TENSOR_CPU_MEMORY);
    106
    11
      const float* const qp = q->data.f32;
    107
    11
      const float* const kp = k->data.f32;
    108
    11
      const float* const vp = v->data.f32;
    109
    11
      const float* const amp = attn_mask ? attn_mask->data.f322 : 09;
    110
    11
      float* const cp = c->data.f32;
    111
    11
      const float scale = cmd.info.scaled_dot_product_attention.scale;
    112
    11
      const int is_causal = cmd.info.scaled_dot_product_attention.is_causal;
    113
    11
      const int h_h_k_ratio = qdim[2] / kdim[2];
    114
    11
      assert(kdim[2] == vdim[2]);
    115
    11
      assert(qdim[2] >= kdim[2]);
    116
    11
      assert(qdim[2] % kdim[2] == 0);
    117
    11
      if (is_varlen)
    118
    2
      {
    119
    2
        const int batch_size = ccv_nnc_tensor_count(q_seq_offsets->info) - 1;
    120
    2
        assert(batch_size > 0);
    121
    2
        assert(ccv_nnc_tensor_count(kv_seq_offsets->info) == batch_size + 1);
    122
    2
        assert(qdim[0] == 1);
    123
    2
        assert(kdim[0] == 1);
    124
    2
        assert(vdim[0] == 1);
    125
    2
        assert(cdim[0] == 1);
    126
    2
        assert(cdim[1] == qdim[1]);
    127
    2
        assert(cdim[2] == qdim[2]);
    128
    2
        assert(cdim[3] == vdim[3]);
    129
    2
        assert(kdim[1] == vdim[1]);
    130
    2
        const int* const q_offset = q_seq_offsets->data.i32;
    131
    2
        const int* const kv_offset = kv_seq_offsets->data.i32;
    132
    2
        assert(q_offset[0] == 0);
    133
    2
        assert(kv_offset[0] == 0);
    134
    2
        assert(q_offset[batch_size] == qdim[1]);
    135
    2
        assert(kv_offset[batch_size] == kdim[1]);
    136
    10
        for (i[0] = 0; 2i[0] < batch_size; i[0]++8)
    137
    8
        {
    138
    8
          const int q_start = q_offset[i[0]];
    139
    8
          const int q_end = q_offset[i[0] + 1];
    140
    8
          const int k_start = kv_offset[i[0]];
    141
    8
          const int k_end = kv_offset[i[0] + 1];
    142
    8
          assert(q_start <= q_end);
    143
    8
          assert(k_start <= k_end);
    144
    8
          const int R = q_end - q_start;
    145
    8
          const int K = k_end - k_start;
    146
    8
          assert(R > 0);
    147
    8
          assert(K > 0);
    148
    8
          assert(R <= cmd.info.scaled_dot_product_attention.max_seqlen_q);
    149
    8
          assert(K <= cmd.info.scaled_dot_product_attention.max_seqlen_kv);
    150
    8
          const float* const qp0 = qp + q_start * qstride[1];
    151
    8
          const float* const kp0 = kp + k_start * kstride[1];
    152
    8
          const float* const vp0 = vp + k_start * vstride[1];
    153
    8
          float* const cp0 = cp + q_start * cstride[1];
    154
    40
          for (i[1] = 0; i[1] < qdim[2]; i[1]++32)
    155
    32
          {
    156
    32
            const float* const qp1 = qp0 + i[1] * qstride[2];
    157
    32
            const float* const kp1 = kp0 + (i[1] / h_h_k_ratio) * kstride[2];
    158
    32
            const float* const vp1 = vp0 + (i[1] / h_h_k_ratio) * vstride[2];
    159
    32
            float* const cp1 = cp0 + i[1] * cstride[2];
    160
    136
            parallel_for32(x, R) {
    161
    136
              int y, k;
    162
    136
              const float* const qp2 = qp1 + x * qstride[1];
    163
    136
              float* const cp2 = cp1 + x * cstride[1];
    164
    136
              float* const qk0 = qk + x * K;
    165
    808
              for (y = 0; y < K; y++672)
    166
    672
              {
    167
    672
                const float* const kp2 = kp1 + y * kstride[1];
    168
    672
                float v = 0;
    169
    43.6k
                for (k = 0; k < qdim[3]; k++43.0k)
    170
    43.0k
                  v += qp2[k * qstride[3]] * kp2[k * kstride[3]];
    171
    672
                qk0[y] = scale * v;
    172
    672
              }
    173
    136
              if (is_causal)
    174
    136
              {
    175
    136
                const int x_end = ccv_max(x - R + K + 1, 0);
    176
    368
                for (y = x_end; y < K; y++232)
    177
    232
                  qk0[y] = 0;
    178
    136
                double maxval = qk0[0];
    179
    472
                for (y = 1; y < x_end; y++336)
    180
    336
                  if (qk0[y] > maxval)
    181
    118
                    maxval = qk0[y];
    182
    136
                double sumval = 0;
    183
    576
                for (y = 0; y < x_end; y++440)
    184
    440
                  sumval += (qk0[y] = expf(qk0[y] - maxval));
    185
    136
                sumval = 1.0 / sumval;
    186
    576
                for (y = 0; y < x_end; y++440)
    187
    440
                  qk0[y] *= sumval;
    188
    136
              } else {
    189
    0
                double maxval = qk0[0];
    190
    0
                for (y = 1; y < K; y++)
    191
    0
                  if (qk0[y] > maxval)
    192
    0
                    maxval = qk0[y];
    193
    0
                double sumval = 0;
    194
    0
                for (y = 0; y < K; y++)
    195
    0
                  sumval += (qk0[y] = expf(qk0[y] - maxval));
    196
    0
                sumval = 1.0 / sumval;
    197
    0
                for (y = 0; y < K; y++)
    198
    0
                  qk0[y] *= sumval;
    199
    0
              }
    200
    8.84k
              for (k = 0; k < vdim[3]; k++8.70k)
    201
    8.70k
                cp2[k * cstride[3]] = 0;
    202
    808
              for (y = 0; y < K; y++672)
    203
    672
              {
    204
    672
                const float* const vp2 = vp1 + y * vstride[1];
    205
    672
                const float v = qk0[y];
    206
    43.6k
                for (k = 0; k < vdim[3]; k++43.0k)
    207
    43.0k
                  cp2[k * cstride[3]] += v * vp2[k * vstride[3]];
    208
    672
              }
    209
    136
            } parallel_endfor
    210
    32
          }
    211
    8
        }
    212
    2
        return CCV_NNC_EXEC_SUCCESS;
    213
    2
      }
    214
    297
      for (i[0] = 0; 9i[0] < qdim[0]; i[0]++288)
    215
    288
      {
    216
    288
        const float* const qp0 = qp + i[0] * qstride[0];
    217
    288
        const float* const kp0 = kp + i[0] * kstride[0];
    218
    288
        const float* const vp0 = vp + i[0] * vstride[0];
    219
    288
        const float* const amp0 = amp && amdim[0] > 164 ? amp + i[0] * amstride[0]0 : amp;
    220
    288
        float* const cp0 = cp + i[0] * cstride[0];
    221
    2.59k
        for (i[1] = 0; i[1] < qdim[2]; i[1]++2.30k)
    222
    2.30k
        {
    223
    2.30k
          const float* const qp1 = qp0 + i[1] * qstride[2];
    224
    2.30k
          const float* const kp1 = kp0 + (i[1] / h_h_k_ratio) * kstride[2];
    225
    2.30k
          const float* const vp1 = vp0 + (i[1] / h_h_k_ratio) * vstride[2];
    226
    2.30k
          const float* const amp1 = amp && amdim[1] > 1512 ? amp0 + i[1] * amstride[1]0 : amp0;
    227
    2.30k
          float* const cp1 = cp0 + i[1] * cstride[2];
    228
          // Compute Q @ K^T
    229
    294k
          parallel_for2.30k(x, qdim[1]) {
    230
    294k
            int y, k;
    231
    294k
            const float* const qp2 = qp1 + x * qstride[1];
    232
    294k
            float* const cp2 = cp1 + x * cstride[1];
    233
    294k
            float* const qk0 = qk + x * kdim[1];
    234
    294k
            const float* const amp2 = amp1 ? amp1 + x * amstride[2]65.5k : 0229k;
    235
    294k
            if (attn_mask)
    236
    65.5k
            {
    237
    8.45M
              for (y = 0; y < kdim[1]; y++8.38M)
    238
    8.38M
              {
    239
    8.38M
                const float* const kp2 = kp1 + y * kstride[1];
    240
    8.38M
                float v = 0;
    241
    545M
                for (k = 0; k < qdim[3]; k++536M)
    242
    536M
                  v += qp2[k * qstride[3]] * kp2[k * kstride[3]];
    243
    8.38M
                qk0[y] = scale * v + amp2[y * amstride[3]];
    244
    8.38M
              }
    245
    229k
            } else {
    246
    29.5M
              for (y = 0; y < kdim[1]; y++29.3M)
    247
    29.3M
              {
    248
    29.3M
                const float* const kp2 = kp1 + y * kstride[1];
    249
    29.3M
                float v = 0;
    250
    1.90G
                for (k = 0; k < qdim[3]; k++1.87G)
    251
    1.87G
                  v += qp2[k * qstride[3]] * kp2[k * kstride[3]];
    252
    29.3M
                qk0[y] = scale * v;
    253
    29.3M
              }
    254
    229k
            }
    255
            // Compute softmax on qk.
    256
    294k
            if (is_causal)
    257
    0
            {
    258
    0
              const int x_end = ccv_max(x - qdim[1] + kdim[1] + 1, 0);
    259
    0
              for (y = x_end; y < kdim[1]; y++)
    260
    0
                qk0[y] = 0;
    261
    0
              double maxval = qk0[0];
    262
    0
              for (y = 1; y < x_end; y++)
    263
    0
                if (qk0[y] > maxval)
    264
    0
                  maxval = qk0[y];
    265
    0
              double sumval = 0;
    266
    0
              for (y = 0; y < x_end; y++)
    267
    0
                sumval += (qk0[y] = expf(qk0[y] - maxval));
    268
    0
              sumval = 1.0 / sumval;
    269
    0
              for (y = 0; y < x_end; y++)
    270
    0
                qk0[y] *= sumval;
    271
    294k
            } else {
    272
    294k
              double maxval = qk0[0];
    273
    37.7M
              for (y = 1; y < kdim[1]; y++37.4M)
    274
    37.4M
                if (qk0[y] > maxval)
    275
    1.25M
                  maxval = qk0[y];
    276
    294k
              double sumval = 0;
    277
    38.0M
              for (y = 0; y < kdim[1]; y++37.7M)
    278
    37.7M
                sumval += (qk0[y] = expf(qk0[y] - maxval));
    279
    294k
              sumval = 1.0 / sumval;
    280
    38.0M
              for (y = 0; y < kdim[1]; y++37.7M)
    281
    37.7M
                qk0[y] *= sumval;
    282
    294k
            }
    283
    28.6M
            for (k = 0; k < vdim[3]; k++28.3M)
    284
    28.3M
              cp2[k * cstride[3]] = 0;
    285
    38.0M
            for (y = 0; y < kdim[1]; y++37.7M)
    286
    37.7M
            {
    287
    37.7M
              const float* const vp2 = vp1 + y * vstride[1];
    288
    37.7M
              const float v = qk0[y];
    289
    3.66G
              for (k = 0; k < vdim[3]; k++3.62G)
    290
    3.62G
                cp2[k * cstride[3]] += v * vp2[k * vstride[3]];
    291
    37.7M
            }
    292
    294k
          } parallel_endfor
    293
    2.30k
        }
    294
    288
      }
    295
    9
      if (w)
    296
    3
      {
    297
    3
        const int num_heads = cdim[2];
    298
    3
        ccv_nnc_tensor_view_t* const d = (ccv_nnc_tensor_view_t*)outputs[0];
    299
    3
        const int w_nd = ccv_nnc_tensor_nd(w->info.dim);
    300
    3
        assert(w_nd == 2);
    301
    3
        assert(CCV_IS_TENSOR_CONTIGUOUS(w));
    302
    3
        const int d_nd = ccv_nnc_tensor_nd(d->info.dim);
    303
    3
        assert(d_nd == 3);
    304
    3
        int ddim[CCV_NNC_MAX_DIM_ALLOC];
    305
    3
        int dstride[CCV_NNC_MAX_DIM_ALLOC];
    306
    3
        ccv_nnc_tensor_view_get_dim(d, ddim);
    307
    3
        ccv_nnc_tensor_view_get_stride(d, dstride);
    308
    3
        assert(ddim[2] == cdim[1]);
    309
    3
        assert(ddim[3] == num_heads * cdim[3]);
    310
    3
        assert(w->info.dim[1] == ddim[3]);
    311
    3
        assert(w->info.dim[0] == ddim[3]);
    312
    3
        float* const dp = d->data.f32;
    313
    3
        const float* const wp = w->data.f32;
    314
    3
        const float* const cp = c->data.f32;
    315
    3
        if (bias)
    316
    3
        {
    317
    3
          assert(ccv_nnc_tensor_count(bias->info) == ddim[3]);
    318
    3
          assert(CCV_IS_TENSOR_CONTIGUOUS(bias));
    319
    3
          const float* const biasp = bias->data.f32;
    320
    99
          for (i[0] = 0; i[0] < ddim[1]; i[0]++96)
    321
    96
          {
    322
    96
            const float* const cp0 = cp + i[0] * cstride[0];
    323
    96
            float* const dp0 = dp + i[0] * dstride[1];
    324
    12.2k
            parallel_for96(y, ddim[2]) {
    325
    12.2k
              int x, j, k;
    326
    12.2k
              const float* const cp1 = cp0 + y * cstride[1];
    327
    12.2k
              float* const dp1 = dp0 + y * dstride[2];
    328
    9.44M
              for (x = 0; x < ddim[3]; x++9.43M)
    329
    9.43M
              {
    330
    9.43M
                const float* const wp0 = wp + x * ddim[3];
    331
    9.43M
                float v = biasp[x];
    332
    84.9M
                for (j = 0; j < num_heads; j++75.4M)
    333
    75.4M
                {
    334
    75.4M
                  const float* const cp2 = cp1 + j * cstride[2];
    335
    7.32G
                  for (k = 0; k < cdim[3]; k++7.24G)
    336
    7.24G
                    v += wp0[j * cdim[3] + k] * cp2[k * cstride[3]];
    337
    75.4M
                }
    338
    9.43M
                dp1[x * dstride[3]] = v;
    339
    9.43M
              }
    340
    12.2k
            } parallel_endfor
    341
    96
          }
    342
    3
        } else {
    343
    0
          for (i[0] = 0; i[0] < ddim[1]; i[0]++)
    344
    0
          {
    345
    0
            const float* const cp0 = cp + i[0] * cstride[0];
    346
    0
            float* const dp0 = dp + i[0] * dstride[1];
    347
    0
            parallel_for(y, ddim[2]) {
    348
    0
              int x, j, k;
    349
    0
              const float* const cp1 = cp0 + y * cstride[1];
    350
    0
              float* const dp1 = dp0 + y * dstride[2];
    351
    0
              for (x = 0; x < ddim[3]; x++)
    352
    0
              {
    353
    0
                const float* const wp0 = wp + x * ddim[3];
    354
    0
                float v = 0;
    355
    0
                for (j = 0; j < num_heads; j++)
    356
    0
                {
    357
    0
                  const float* const cp2 = cp1 + j * cstride[2];
    358
    0
                  for (k = 0; k < cdim[3]; k++)
    359
    0
                    v += wp0[j * cdim[3] + k] * cp2[k * cstride[3]];
    360
    0
                }
    361
    0
                dp1[x * dstride[3]] = v;
    362
    0
              }
    363
    0
            } parallel_endfor
    364
    0
          }
    365
    0
        }
    366
    3
      }
    367
    9
      return CCV_NNC_EXEC_SUCCESS;
    368
    9
    }
    369
    370
    static int _ccv_nnc_scaled_dot_product_attention_back(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)
    371
    1
    {
    372
      // Assuming no saved_softmax, we need to recompute from q, k, v.
    373
      // We cannot do this with masks (yet).
    374
    1
      assert(input_size >= 6);
    375
    1
      if (cmd.info.scaled_dot_product_attention.is_varlen)
    376
    0
        return CCV_NNC_EXEC_INVALID;
    377
    1
      ccv_nnc_tensor_view_t* const g = (ccv_nnc_tensor_view_t*)inputs[0];
    378
    1
      ccv_nnc_tensor_view_t* const q = (ccv_nnc_tensor_view_t*)inputs[3];
    379
    1
      ccv_nnc_tensor_view_t* const k = (ccv_nnc_tensor_view_t*)inputs[4];
    380
    1
      ccv_nnc_tensor_view_t* const v = (ccv_nnc_tensor_view_t*)inputs[5];
    381
    1
      ccv_nnc_tensor_view_t* const dq = (ccv_nnc_tensor_view_t*)outputs[0];
    382
    1
      ccv_nnc_tensor_view_t* const dk = (ccv_nnc_tensor_view_t*)outputs[1];
    383
    1
      ccv_nnc_tensor_view_t* const dv = (ccv_nnc_tensor_view_t*)outputs[2];
    384
    1
      const int q_nd = ccv_nnc_tensor_nd(q->info.dim);
    385
    1
      assert(q_nd == 3 || q_nd == 4);
    386
    1
      const int k_nd = ccv_nnc_tensor_nd(k->info.dim);
    387
    1
      assert(k_nd == 3 || k_nd == 4);
    388
    1
      const int v_nd = ccv_nnc_tensor_nd(v->info.dim);
    389
    1
      assert(v_nd == 3 || v_nd == 4);
    390
    1
      const int g_nd = ccv_nnc_tensor_nd(g->info.dim);
    391
    1
      assert(g_nd == 3 || g_nd == 4);
    392
    1
      const int dq_nd = ccv_nnc_tensor_nd(dq->info.dim);
    393
    1
      assert(dq_nd == 3 || dq_nd == 4);
    394
    1
      assert(dq_nd == q_nd);
    395
    1
      const int dk_nd = ccv_nnc_tensor_nd(dk->info.dim);
    396
    1
      assert(dk_nd == 3 || dk_nd == 4);
    397
    1
      assert(dk_nd == k_nd);
    398
    1
      const int dv_nd = ccv_nnc_tensor_nd(dv->info.dim);
    399
    1
      assert(dv_nd == 3 || dv_nd == 4);
    400
    1
      assert(dv_nd == v_nd);
    401
    1
      assert(q_nd == k_nd && k_nd == v_nd && v_nd == g_nd);
    402
      // Assuming this is float 32.
    403
    1
      int qdim[CCV_NNC_MAX_DIM_ALLOC];
    404
    1
      int kdim[CCV_NNC_MAX_DIM_ALLOC];
    405
    1
      int vdim[CCV_NNC_MAX_DIM_ALLOC];
    406
    1
      int gdim[CCV_NNC_MAX_DIM_ALLOC];
    407
    1
      int dqdim[CCV_NNC_MAX_DIM_ALLOC];
    408
    1
      int dkdim[CCV_NNC_MAX_DIM_ALLOC];
    409
    1
      int dvdim[CCV_NNC_MAX_DIM_ALLOC];
    410
    1
      ccv_nnc_tensor_view_get_dim(q, qdim);
    411
    1
      ccv_nnc_tensor_view_get_dim(k, kdim);
    412
    1
      ccv_nnc_tensor_view_get_dim(v, vdim);
    413
    1
      ccv_nnc_tensor_view_get_dim(g, gdim);
    414
    1
      ccv_nnc_tensor_view_get_dim(dq, dqdim);
    415
    1
      ccv_nnc_tensor_view_get_dim(dk, dkdim);
    416
    1
      ccv_nnc_tensor_view_get_dim(dv, dvdim);
    417
    1
      if (q_nd == 3)
    418
    0
      {
    419
    0
        qdim[0] = qdim[1], qdim[1] = qdim[2], qdim[2] = 1;
    420
    0
        kdim[0] = kdim[1], kdim[1] = kdim[2], kdim[2] = 1;
    421
    0
        vdim[0] = vdim[1], vdim[1] = vdim[2], vdim[2] = 1;
    422
    0
        gdim[0] = gdim[1], gdim[1] = gdim[2], gdim[2] = 1;
    423
    0
        dqdim[0] = dqdim[1], dqdim[1] = dqdim[2], dqdim[2] = 1;
    424
    0
        dkdim[0] = dkdim[1], dkdim[1] = dkdim[2], dkdim[2] = 1;
    425
    0
        dvdim[0] = dvdim[1], dvdim[1] = dvdim[2], dvdim[2] = 1;
    426
    0
      }
    427
    1
      assert(qdim[0] == kdim[0] && kdim[0] == vdim[0] && vdim[0] == gdim[0]);
    428
    1
      assert(qdim[2] == gdim[2]);
    429
    1
      assert(kdim[2] == vdim[2]);
    430
    1
      assert(qdim[2] % kdim[2] == 0);
    431
    1
      assert(qdim[2] >= kdim[2]);
    432
    1
      assert(qdim[3] == kdim[3]);
    433
    1
      assert(kdim[1] == vdim[1]);
    434
    1
      assert(gdim[1] == qdim[1]);
    435
    1
      assert(gdim[3] == vdim[3]);
    436
    1
      assert(CCV_NNC_MAX_DIM == 2); // Need to change this logic for CCV_NNC_MAX_DIM == other number.
    437
    1
      int qstride[CCV_NNC_MAX_DIM_ALLOC];
    438
    1
      int kstride[CCV_NNC_MAX_DIM_ALLOC];
    439
    1
      int vstride[CCV_NNC_MAX_DIM_ALLOC];
    440
    1
      int gstride[CCV_NNC_MAX_DIM_ALLOC];
    441
    1
      int dqstride[CCV_NNC_MAX_DIM_ALLOC];
    442
    1
      int dkstride[CCV_NNC_MAX_DIM_ALLOC];
    443
    1
      int dvstride[CCV_NNC_MAX_DIM_ALLOC];
    444
    1
      ccv_nnc_tensor_view_get_stride(q, qstride);
    445
    1
      ccv_nnc_tensor_view_get_stride(k, kstride);
    446
    1
      ccv_nnc_tensor_view_get_stride(v, vstride);
    447
    1
      ccv_nnc_tensor_view_get_stride(g, gstride);
    448
    1
      ccv_nnc_tensor_view_get_stride(dq, dqstride);
    449
    1
      ccv_nnc_tensor_view_get_stride(dk, dkstride);
    450
    1
      ccv_nnc_tensor_view_get_stride(dv, dvstride);
    451
    1
      if (q_nd == 3)
    452
    0
      {
    453
    0
        qstride[0] = qstride[1], qstride[1] = qstride[2], qstride[2] = qstride[3];
    454
    0
        kstride[0] = kstride[1], kstride[1] = kstride[2], kstride[2] = kstride[3];
    455
    0
        vstride[0] = vstride[1], vstride[1] = vstride[2], vstride[2] = vstride[3];
    456
    0
        gstride[0] = gstride[1], gstride[1] = gstride[2], gstride[2] = gstride[3];
    457
    0
        dqstride[0] = dqstride[1], dqstride[1] = dqstride[2], dqstride[2] = dqstride[3];
    458
    0
        dkstride[0] = dkstride[1], dkstride[1] = dkstride[2], dkstride[2] = dkstride[3];
    459
    0
        dvstride[0] = dvstride[1], dvstride[1] = dvstride[2], dvstride[2] = dvstride[3];
    460
    0
      }
    461
    1
      int i[CCV_NNC_MAX_DIM + 2];
    462
    1
      float* qk = ccv_nnc_stream_context_get_workspace(stream_context, sizeof(float) * 2 * kdim[1], CCV_TENSOR_CPU_MEMORY);
    463
    1
      const float* const qp = q->data.f32;
    464
    1
      const float* const kp = k->data.f32;
    465
    1
      const float* const vp = v->data.f32;
    466
    1
      const float* const gp = g->data.f32;
    467
    1
      float* const dqp = dq->data.f32;
    468
    1
      float* const dkp = dk->data.f32;
    469
    1
      float* const dvp = dv->data.f32;
    470
    1
      const float scale = cmd.info.scaled_dot_product_attention.scale;
    471
    1
      const int is_causal = cmd.info.scaled_dot_product_attention.is_causal;
    472
    1
      const int h_h_k_ratio = qdim[2] / kdim[2];
    473
    33
      for (i[0] = 0; i[0] < qdim[0]; i[0]++32)
    474
    32
      {
    475
    32
        const float* const qp0 = qp + i[0] * qstride[0];
    476
    32
        const float* const kp0 = kp + i[0] * kstride[0];
    477
    32
        const float* const vp0 = vp + i[0] * vstride[0];
    478
    32
        const float* const gp0 = gp + i[0] * gstride[0];
    479
    32
        float* const dqp0 = dqp + i[0] * dqstride[0];
    480
    32
        float* const dkp0 = dkp + i[0] * dkstride[0];
    481
    32
        float* const dvp0 = dvp + i[0] * dvstride[0];
    482
    288
        for (i[1] = 0; i[1] < qdim[2]; i[1]++256)
    483
    256
        {
    484
    256
          const float* const qp1 = qp0 + i[1] * qstride[2];
    485
    256
          const float* const kp1 = kp0 + (i[1] / h_h_k_ratio) * kstride[2];
    486
    256
          const float* const vp1 = vp0 + (i[1] / h_h_k_ratio) * vstride[2];
    487
    256
          const float* const gp1 = gp0 + i[1] * gstride[2];
    488
    256
          float* const dqp1 = dqp0 + i[1] * dqstride[2];
    489
    256
          float* const dkp1 = dkp0 + (i[1] / h_h_k_ratio) * dkstride[2];
    490
    256
          float* const dvp1 = dvp0 + (i[1] / h_h_k_ratio) * dvstride[2];
    491
          // Compute Q @ K^T
    492
    256
          int x, y, k;
    493
    33.0k
          for (x = 0; x < qdim[1]; x++32.7k)
    494
    32.7k
          {
    495
    32.7k
            float* const dqp2 = dqp1 + x * dqstride[1];
    496
    2.12M
            for (k = 0; k < qdim[3]; k++2.09M)
    497
    2.09M
              dqp2[k * dqstride[3]] = 0;
    498
    32.7k
          }
    499
          // Only zero out when it is at 0-index.
    500
    256
          if (i[1] % h_h_k_ratio == 0)
    501
    33.0k
            for (y = 0; 256y < kdim[1]; y++32.7k)
    502
    32.7k
            {
    503
    32.7k
              float* const dkp2 = dkp1 + y * dkstride[1];
    504
    2.12M
              for (k = 0; k < qdim[3]; k++2.09M)
    505
    2.09M
                dkp2[k * dkstride[3]] = 0;
    506
    32.7k
            }
    507
          // Only zero out when it is at 0-index.
    508
    256
          if (i[1] % h_h_k_ratio == 0)
    509
    33.0k
            for (y = 0; 256y < kdim[1]; y++32.7k)
    510
    32.7k
            {
    511
    32.7k
              float* const dvp2 = dvp1 + y * dvstride[1];
    512
    3.17M
              for (k = 0; k < vdim[3]; k++3.14M)
    513
    3.14M
                dvp2[k * dvstride[3]] = 0;
    514
    32.7k
            }
    515
    33.0k
          for (x = 0; x < qdim[1]; x++32.7k)
    516
    32.7k
          {
    517
    32.7k
            const float* const qp2 = qp1 + x * qstride[1];
    518
    32.7k
            const float* const gp2 = gp1 + x * gstride[1];
    519
    32.7k
            float* const qk0 = qk;
    520
    32.7k
            float* const qks0 = qk + kdim[1];
    521
    4.22M
            for (y = 0; y < kdim[1]; y++4.19M)
    522
    4.19M
            {
    523
    4.19M
              const float* const kp2 = kp1 + y * kstride[1];
    524
    4.19M
              float v = 0;
    525
    272M
              for (k = 0; k < qdim[3]; k++268M)
    526
    268M
                v += qp2[k * qstride[3]] * kp2[k * kstride[3]];
    527
    4.19M
              qk0[y] = scale * v;
    528
    4.19M
            }
    529
            // Compute softmax on qk.
    530
    32.7k
            if (is_causal)
    531
    0
            {
    532
    0
              const int x_end = ccv_max(x - qdim[1] + kdim[1] + 1, 0);
    533
    0
              for (y = x_end; y < kdim[1]; y++)
    534
    0
                qk0[y] = 0;
    535
    0
              double maxval = qk0[0];
    536
    0
              for (y = 1; y < x_end; y++)
    537
    0
                if (qk0[y] > maxval)
    538
    0
                  maxval = qk0[y];
    539
    0
              double sumval = 0;
    540
    0
              for (y = 0; y < x_end; y++)
    541
    0
                sumval += (qk0[y] = expf(qk0[y] - maxval));
    542
    0
              sumval = 1.0 / sumval;
    543
    0
              for (y = 0; y < x_end; y++)
    544
    0
                qk0[y] *= sumval;
    545
    32.7k
            } else {
    546
    32.7k
              double maxval = qk0[0];
    547
    4.19M
              for (y = 1; y < kdim[1]; y++4.16M)
    548
    4.16M
                if (qk0[y] > maxval)
    549
    146k
                  maxval = qk0[y];
    550
    32.7k
              double sumval = 0;
    551
    4.22M
              for (y = 0; y < kdim[1]; y++4.19M)
    552
    4.19M
                sumval += (qk0[y] = expf(qk0[y] - maxval));
    553
    32.7k
              sumval = 1.0 / sumval;
    554
    4.22M
              for (y = 0; y < kdim[1]; y++4.19M)
    555
    4.19M
                qk0[y] *= sumval;
    556
    32.7k
            }
    557
    4.22M
            for (y = 0; y < kdim[1]; y++4.19M)
    558
    4.19M
            {
    559
    4.19M
              float* const dvp2 = dvp1 + y * dvstride[1];
    560
    4.19M
              const float v = qk0[y];
    561
    406M
              for (k = 0; k < vdim[3]; k++402M)
    562
    402M
                dvp2[k * dvstride[3]] += v * gp2[k * gstride[3]];
    563
    4.19M
            }
    564
    32.7k
            double sumval = 0;
    565
    4.22M
            for (y = 0; y < kdim[1]; y++4.19M)
    566
    4.19M
            {
    567
    4.19M
              const float* const vp2 = vp1 + y * vstride[1];
    568
    4.19M
              float v = 0;
    569
    406M
              for (k = 0; k < vdim[3]; k++402M)
    570
    402M
                v += gp2[k * gstride[3]] * vp2[k * vstride[3]];
    571
    4.19M
              qks0[y] = v;
    572
    4.19M
              sumval += v * qk0[y];
    573
    4.19M
            }
    574
    4.22M
            for (y = 0; y < kdim[1]; y++4.19M)
    575
    4.19M
              qk0[y] = (qks0[y] - sumval) * qk0[y];
    576
    32.7k
            float* const dqp2 = dqp1 + x * dqstride[1];
    577
    4.22M
            for (y = 0; y < kdim[1]; y++4.19M)
    578
    4.19M
            {
    579
    4.19M
              const float* const kp2 = kp1 + y * kstride[1];
    580
    4.19M
              float* const dkp2 = dkp1 + y * dkstride[1];
    581
    4.19M
              const float v = scale * qk0[y];
    582
    272M
              for (k = 0; k < qdim[3]; k++268M)
    583
    268M
              {
    584
    268M
                dqp2[k * dqstride[3]] += v * kp2[k * kstride[3]];
    585
    268M
                dkp2[k * dkstride[3]] += v * qp2[k * qstride[3]];
    586
    268M
              }
    587
    4.19M
            }
    588
    32.7k
          }
    589
    256
        }
    590
    32
      }
    591
    1
      return CCV_NNC_EXEC_SUCCESS;
    592
    1
    }
    593
    594
    REGISTER_COMMAND_BACKEND(CCV_NNC_SCALED_DOT_PRODUCT_ATTENTION_FORWARD, CCV_NNC_BACKEND_CPU_REF)(ccv_nnc_cmd_backend_registry_t* const registry)
    595
    1
    {
    596
    1
      registry->tensor_formats = CCV_TENSOR_FORMAT_NHWC;
    597
    1
      registry->tensor_datatypes = CCV_32F | CCV_32S;
    598
    1
      registry->tensor_memory = CCV_TENSOR_CPU_MEMORY;
    599
    1
      registry->algorithms = 1;
    600
    1
      registry->exec = _ccv_nnc_scaled_dot_product_attention_forw;
    601
    1
    }
    602
    603
    REGISTER_COMMAND_BACKEND(CCV_NNC_SCALED_DOT_PRODUCT_ATTENTION_BACKWARD, CCV_NNC_BACKEND_CPU_REF)(ccv_nnc_cmd_backend_registry_t* const registry)
    604
    1
    {
    605
    1
      registry->tensor_formats = CCV_TENSOR_FORMAT_NHWC;
    606
    1
      registry->tensor_datatypes = CCV_32F;
    607
    1
      registry->tensor_memory = CCV_TENSOR_CPU_MEMORY;
    608
    1
      registry->algorithms = 1;
    609
    1
      registry->exec = _ccv_nnc_scaled_dot_product_attention_back;
    610
    1
    }