llm.cpp/include/attention.h at master · LMGNU/llm.cpp · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
#pragma once
// ============================================================
//  include/attention.h  –  Causal self-attention
//  Mirrors: class Head and class MultiHeadAttention in Python
// ============================================================

#include "config/config.h"
#include "linear.h"
#include "tensor.h"

#include <fstream>
#include <vector>

// ------------------------------------------------------------------
// Single causal attention head
// ------------------------------------------------------------------
struct Head
{
      int head_size;
      Linear key, query, value; // each [n_embd head_size], no bias

      Head() = default;

      Head(int n_embd, int hs, std::mt19937 &rng)
            : head_size(hs), key(n_embd, hs, false, rng), query(n_embd, hs, false, rng),
              value(n_embd, hs, false, rng)
      {
      }

      // x: [B, T, n_embd]   [B, T, head_size]
      Tensor forward(const Tensor &x, bool training, std::mt19937 &rng) const
      {
            int B = x.shape[0], T = x.shape[1];

            Tensor k = key.forward(x);   // [B, T, hs]
            Tensor q = query.forward(x); // [B, T, hs]
            Tensor v = value.forward(x); // [B, T, hs]

            // scaled dot-product  wei = q @ k^T * scale
            float scale = 1.0f / std::sqrt((float)head_size);
            Tensor kt = transpose23(k); // [B, hs, T]
            Tensor wei = bmm(q, kt);    // [B, T, T]

            // mask upper triangle -inf
            for (int b = 0; b < B; ++b)
                  for (int i = 0; i < T; ++i)
                        for (int j = i + 1; j < T; ++j)
                              wei.at(b, i, j) = -1e30f;

            // scale then softmax
            wei = scale3d_inplace(wei, scale);
            wei = softmax3d(wei);

            // dropout on attention weights
            wei = dropout(wei, DROPOUT, training, rng);

            // wei @ v   [B, T, hs]
            return bmm(wei, v);
      }

      int num_params() const
      {
            return key.num_params() + query.num_params() + value.num_params();
      }

      void save(std::ofstream &f) const
      {
            key.save(f);
            query.save(f);
            value.save(f);
      }
      void load(std::ifstream &f)
      {
            key.load(f);
            query.load(f);
            value.load(f);
      }

    private:
      static Tensor scale3d_inplace(Tensor t, float s)
      {
            for (auto &v : t.data)
                  v *= s;
            return t;
      }
};

// Multi-head causal attention
struct MultiHeadAttention
{
      int num_heads, head_size, n_embd;
      std::vector<Head> heads;
      Linear proj;

      MultiHeadAttention() = default;

      MultiHeadAttention(int n_embd_, int num_h, int hs, std::mt19937 &rng)
            : num_heads(num_h), head_size(hs), n_embd(n_embd_), proj(num_h * hs, n_embd_, true, rng)
      {
            for (int i = 0; i < num_h; ++i)
                  heads.emplace_back(n_embd_, hs, rng);
      }

      // x: [B, T, n_embd] to  [B, T, n_embd]
      Tensor forward(const Tensor &x, bool training, std::mt19937 &rng) const
      {
            std::vector<Tensor> head_outs;
            head_outs.reserve(num_heads);
            for (auto &h : heads)
                  head_outs.push_back(h.forward(x, training, rng));

            Tensor concat = cat_last(head_outs); // [B, T, num_heads*hs]
            Tensor out = proj.forward(concat);   // [B, T, n_embd]
            return dropout(out, DROPOUT, training, rng);
      }

      int num_params() const
      {
            int n = proj.num_params();
            for (auto &h : heads)
                  n += h.num_params();
            return n;
      }

      void save(std::ofstream &f) const
      {
            for (auto &h : heads)
                  h.save(f);
            proj.save(f);
      }
      void load(std::ifstream &f)
      {
            for (auto &h : heads)
                  h.load(f);
            proj.load(f);
      }
};