Language Modeling Implementations

Sequence Modeling

RNN

 1import torch
 2from torch.nn.utils.rnn import pad_sequence
 3from collections import OrderedDict
 4
 5class MyRNNCell(torch.nn.Module):
 6    def __init__(self, in_dim, out_dim):
 7        super().__init__()
 8        self.h = torch.zeros(out_dim)
 9        self.ih = torch.nn.Linear(in_features=in_dim, out_features=out_dim, bias=True)
10        self.hh = torch.nn.Linear(in_features=out_dim, out_features=out_dim, bias=True)
11        self.norm = torch.nn.LayerNorm(normalized_shape=out_dim)
12        self.relu = torch.nn.ReLU()
13
14    def detach_hidden(self):
15        self.h.detach_()
16
17    def forward(self, x):
18        # recurrence here
19        a = self.hh(self.h) + self.ih(x)
20        a = self.norm(a + x)
21        self.h = self.relu(a)
22        return self.h
23
24class MyDeepRNN(torch.nn.Module):
25
26    def __init__(self, in_dim, hidden_dim, out_dim):
27        super().__init__()
28        self.embedding = torch.nn.Embedding(num_embeddings=in_dim, embedding_dim=hidden_dim, max_norm=True)
29        self.rnn = torch.nn.Sequential(OrderedDict([
30                        ('rnn1', MyRNNCell(in_dim=hidden_dim, out_dim=hidden_dim)),
31                        ('rnn2', MyRNNCell(in_dim=hidden_dim, out_dim=hidden_dim)),
32                        ('rnn3', MyRNNCell(in_dim=hidden_dim, out_dim=hidden_dim)),
33                        ('rnn4', MyRNNCell(in_dim=hidden_dim, out_dim=hidden_dim))
34                    ]))
35        self.ho = torch.nn.Linear(in_features=hidden_dim, out_features=out_dim, bias=True)
36        self.norm = torch.nn.LayerNorm(normalized_shape=out_dim)
37
38    def detach_hidden(self):
39        for unit in self.rnn:
40            unit.detach_hidden()
41
42    def display_params(self):
43        for name, param in self.named_parameters():
44            print(f'{name}={param.shape}')
45
46    def forward(self, x):
47        h = self.embedding(x)
48        h = self.rnn(h) + h #residual connection
49        o = self.ho(h)
50        y = self.norm(o)
51        return y
52
53if __name__ == '__main__':
54    inputs = ['hello, world!\n', 'this is my first rnn demo.\n', 'excited? nope!\n'] # \n=EOS
55    x_seq = [torch.LongTensor(list(input.encode('utf-8'))) for input in inputs]
56    x_padded = pad_sequence(x_seq, batch_first=False, padding_value=0) #.permute(1,0)
57
58    # since we're only predicting with the last, and we have sequences of different lengths
59    # we create multiple minimatches where we pass first K-1 tokens in the sequence and ask it to predict K
60    # and we vary K to be a sequence of size 2, to all the way N
61    batches = [x_padded[:n+2,:] for n in range(x_padded.shape[0]-1)] # create N(N-1) batches from size N
62
63    vocab_size = 256
64    hidden_dim = 8
65    num_epochs = 100
66
67    model = MyDeepRNN(in_dim=vocab_size, hidden_dim=hidden_dim, out_dim=vocab_size)
68    optim = torch.optim.Adam(model.parameters())
69    criterion = torch.nn.CrossEntropyLoss(ignore_index=0)
70
71    for epoch in range(num_epochs):
72        for batch in batches:
73            model.detach_hidden()
74
75            for input, target in zip(batch[:-1], batch[1:]):
76                pred = model(input)
77
78            loss = criterion(pred, target)
79            optim.zero_grad()
80            loss.backward()
81            optim.step()
82
83        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}")

LSTM

  1import torch
  2from torch.nn.utils.rnn import pad_sequence
  3from collections import OrderedDict
  4
  5class MyLSTMCell(torch.nn.Module):
  6    def __init__(self, in_dim, out_dim):
  7        super().__init__()        
  8        self.c = torch.zeros(out_dim) # long term memory
  9        self.h = torch.zeros(out_dim) # short term memory
 10
 11        # process input gates
 12        self._ii = torch.nn.Linear(in_features=in_dim, out_features=out_dim, bias=True)
 13        self._hi = torch.nn.Linear(in_features=out_dim, out_features=out_dim, bias=True)
 14        self._inorm = torch.nn.LayerNorm(normalized_shape=out_dim)
 15        self._ai = torch.nn.Sigmoid()
 16
 17        # process forget gates
 18        self._if = torch.nn.Linear(in_features=in_dim, out_features=out_dim, bias=True)
 19        self._hf = torch.nn.Linear(in_features=out_dim, out_features=out_dim, bias=True)
 20        self._fnorm = torch.nn.LayerNorm(normalized_shape=out_dim)
 21        self._af = torch.nn.Sigmoid()
 22
 23        # process gating gates
 24        self._ig = torch.nn.Linear(in_features=in_dim, out_features=out_dim, bias=True)
 25        self._hg = torch.nn.Linear(in_features=out_dim, out_features=out_dim, bias=True)
 26        self._gnorm = torch.nn.LayerNorm(normalized_shape=out_dim)
 27        self._ag = torch.nn.Tanh()
 28
 29        # process output gates
 30        self._io = torch.nn.Linear(in_features=in_dim, out_features=out_dim, bias=True)
 31        self._ho = torch.nn.Linear(in_features=out_dim, out_features=out_dim, bias=True)
 32        self._onorm = torch.nn.LayerNorm(normalized_shape=out_dim)
 33        self._ao = torch.nn.Sigmoid()
 34
 35        self._ch = torch.nn.Tanh()
 36
 37    def detach_hidden(self):
 38        self.h.detach_()
 39        self.c.detach_()
 40
 41    def forward(self, x):
 42        f = self._hf(self.h) + self._if(x)
 43        f = self._fnorm(f)
 44        f = self._af(f)
 45
 46        i = self._hi(self.h) + self._ii(x)
 47        i = self._inorm(f)
 48        i = self._ai(i)
 49
 50        g = self._hg(self.h) + self._ig(x)
 51        g = self._gnorm(g)
 52        g = self._ag(g)
 53
 54        o = self._ho(self.h) + self._io(x)
 55        o = self._onorm(o)
 56        o = self._ao(o)
 57
 58        # long term memory: forget old stuff a little bit, add new stuff a little bit
 59        # note that c is just multiplied by the forget gate and added by the input - no weights being multiplied here
 60        self.c = torch.mul(self.c, f) + torch.mul(i, g)
 61
 62        # short term memory: remember the output primarily, also add a little bit from the new long-term memory that formed
 63        self.h = torch.mul(o, self._ch(self.c))
 64
 65        return self.h
 66
 67class MyDeepLSTM(torch.nn.Module):
 68
 69    def __init__(self, in_dim, hidden_dim, out_dim):
 70        super().__init__()
 71        self.embedding = torch.nn.Embedding(num_embeddings=in_dim, embedding_dim=hidden_dim, max_norm=True)
 72        self.rnn = torch.nn.Sequential(OrderedDict([
 73                        ('rnn1', MyLSTMCell(in_dim=hidden_dim, out_dim=hidden_dim)),
 74                        ('rnn2', MyLSTMCell(in_dim=hidden_dim, out_dim=hidden_dim)),
 75                        ('rnn3', MyLSTMCell(in_dim=hidden_dim, out_dim=hidden_dim)),
 76                        ('rnn4', MyLSTMCell(in_dim=hidden_dim, out_dim=hidden_dim))
 77                    ]))
 78        self.ho = torch.nn.Linear(in_features=hidden_dim, out_features=out_dim, bias=True)
 79        self.norm = torch.nn.LayerNorm(normalized_shape=out_dim)
 80
 81    def detach_hidden(self):
 82        for unit in self.rnn:
 83            unit.detach_hidden()
 84
 85    def display_params(self):
 86        for name, param in self.named_parameters():
 87            print(f'{name}={param.shape}')
 88
 89    def forward(self, x):
 90        h = self.embedding(x)
 91        h = self.rnn(h)
 92        o = self.ho(h)
 93        y = self.norm(o)
 94        return y
 95
 96if __name__ == '__main__':
 97    inputs = ['hello, world!\n', 'this is my first rnn demo.\n', 'excited? nope!\n'] # \n=EOS
 98    x_seq = [torch.LongTensor(list(input.encode('utf-8'))) for input in inputs]
 99    x_padded = pad_sequence(x_seq, batch_first=False, padding_value=0) #.permute(1,0)
100
101    # since we're only predicting with the last, and we have sequences of different lengths
102    # we create multiple minimatches where we pass first K-1 tokens in the sequence and ask it to predict K
103    # and we vary K to be a sequence of size 2, to all the way N
104    batches = [x_padded[:n+2,:] for n in range(x_padded.shape[0]-1)] # create N(N-1) batches from size N
105
106    vocab_size = 256
107    hidden_dim = 8
108    num_epochs = 100
109
110    model = MyDeepLSTM(in_dim=vocab_size, hidden_dim=hidden_dim, out_dim=vocab_size)
111    optim = torch.optim.Adam(model.parameters())
112    criterion = torch.nn.CrossEntropyLoss(ignore_index=0)
113
114    for epoch in range(num_epochs):
115        for batch in batches:
116            model.detach_hidden()
117
118            for input, target in zip(batch[:-1], batch[1:]):
119                pred = model(input)
120
121            loss = criterion(pred, target)
122            optim.zero_grad()
123            loss.backward()
124            optim.step()
125
126        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}")

Attention

Understanding Einsum

  1import torch
  2import torch.nn.functional as F
  3
  4"""
  5Rule of thumb:
  6------------------------------------------------------------------------
  7(a) dimensions that appear in the output would appear in the outer-loop.
  8    we'll fill in for these dimensions element-wise.
  9(b) dimensions that appear in both the inputs (common dimensions) 
 10    are multiplied element-wise.
 11(c) dimensions that appear in the inputs but not on the output are summed over.
 12    this means for dimensions that satisfy both (b) and (c) are first multiplied
 13    and then summer over.
 14"""
 15
 16def test_tensorprod():
 17    torch.manual_seed(42)
 18    X = torch.randn((4,5))
 19    Y = torch.randn((5,3))
 20    
 21    Expected = torch.einsum('ij,jk->ijk', X, Y)
 22
 23    # einsum impl
 24    Actual = torch.zeros_like(Expected)
 25
 26    # output dimension
 27    for i in torch.arange(X.shape[-2]):
 28        # output dimension
 29        for j in torch.arange(X.shape[-1]):
 30        # output dimension
 31            for k in torch.arange(Y.shape[-1]):
 32                # loop through common dimension j for element-wise product
 33                Actual[i,j,k] = X[i,j] * Y[j,k]
 34
 35    assert(torch.all(torch.isclose(Expected, Actual)))
 36
 37def test_matmul():
 38    torch.manual_seed(42)
 39    X = torch.randn((4,5))
 40    Y = torch.randn((5,3))
 41    
 42    Expected = torch.einsum('ij,jk->ik', X, Y)
 43
 44    # einsum impl
 45    Actual = torch.zeros_like(Expected)
 46
 47    # output dimension
 48    for i in torch.arange(X.shape[-2]):
 49        # output dimension
 50        for k in torch.arange(Y.shape[-1]):
 51            # loop through common dimension j for element-wise product
 52            Actual[i,k] = torch.dot(X[i,:], Y[:,k])
 53
 54    assert(torch.all(torch.isclose(Expected, Actual)))
 55
 56def test_matmul_reduce_k():
 57    torch.manual_seed(42)
 58    X = torch.randn((4,5))
 59    Y = torch.randn((5,3))
 60    
 61    Expected = torch.einsum('ij,jk->i', X, Y)
 62
 63    # einsum impl
 64    Actual = torch.zeros_like(Expected)
 65
 66    # since k is free for Y, we sum over k and cache it
 67    sum_k_Y = torch.sum(Y,dim=-1)
 68
 69    # output dimension
 70    for i in torch.arange(X.shape[-2]):        
 71        # loop through common dimension j for element-wise product
 72        Actual[i] = torch.dot(X[i,:], sum_k_Y)
 73
 74    assert(torch.all(torch.isclose(Expected, Actual)))
 75
 76def test_matmul_reduce_i():
 77    torch.manual_seed(42)
 78    X = torch.randn((4,5))
 79    Y = torch.randn((5,3))
 80    
 81    Expected = torch.einsum('ij,jk->k', X, Y)
 82
 83    # einsum impl
 84    Actual = torch.zeros_like(Expected)
 85
 86    # since i is free for X, we sum over i and cache it
 87    sum_i_X = torch.sum(X,dim=-2)
 88
 89    # output dimension
 90    for k in torch.arange(Y.shape[-1]):        
 91        # loop through common dimension j for element-wise product
 92        Actual[k] = torch.dot(sum_i_X, Y[:,k])
 93
 94    assert(torch.all(torch.isclose(Expected, Actual)))
 95
 96def test_matmul_reduce_j():
 97    torch.manual_seed(42)
 98    X = torch.randn((4,5))
 99    Y = torch.randn((5,3))
100    
101    Expected = torch.einsum('ij,jk->j', X, Y)
102
103    # einsum impl
104    Actual = torch.zeros_like(Expected)
105
106    # since i is free for X, we sum over i and cache it
107    sum_i_X = torch.sum(X,dim=-2)
108    # since k is free for Y, we sum over k and cache it
109    sum_k_Y = torch.sum(Y,dim=-1)
110
111    # output dimension
112    for j in torch.arange(X.shape[-1]):        
113        # loop through common dimension j for element-wise product
114        Actual[j] = sum_i_X[j] * sum_k_Y[j]
115
116    assert(torch.all(torch.isclose(Expected, Actual)))
117
118def test_matmul_reduce_all():
119    torch.manual_seed(42)
120    X = torch.randn((4,5))
121    Y = torch.randn((5,3))
122    
123    Expected = torch.einsum('ij,jk->', X, Y)
124
125    # einsum impl
126    Actual = torch.zeros_like(Expected)
127
128    # since i is free for X, we sum over i and cache it
129    sum_i_X = torch.sum(X,dim=-2)
130    # since k is free for Y, we sum over k and cache it
131    sum_k_Y = torch.sum(Y,dim=-1)
132
133    # output dimension
134    Actual = torch.dot(sum_i_X, sum_k_Y)
135
136    assert(torch.all(torch.isclose(Expected, Actual)))
137
138###################################################
139# Test with multi-head attention code
140###################################################
141def mha_par_batched(Q,K,V):
142    """
143    args:
144        Q: [b,h,n,k]
145        K: [b,h,m,k]
146        V: [b,h,m,v]
147    returns:
148        O: [b,h,n,v]
149    """
150    logits = torch.einsum('bhnk,bhmk->bhnm', Q, K)
151    weights = F.softmax(logits, dim=-1)
152    O = torch.einsum('bhnm,bhmv->bhnv', weights, V)
153    return logits, O
154
155def mha_par_batched_impl(Q,K,V):
156    _b = Q.shape[0]
157    _h = Q.shape[1]
158    _n = Q.shape[2]
159    _m = K.shape[2]
160    _v = V.shape[-1]
161    logits = torch.zeros(_b,_h,_n,_m)
162    for b in torch.arange(_b):
163        for h in torch.arange(_h):
164            for n in torch.arange(_n):
165                for m in torch.arange(_m):
166                    logits[b,h,n,m] = torch.dot(Q[b,h,n,:], K[b,h,m,:])
167    weights = F.softmax(logits, dim=-1)
168    O = torch.zeros(_b,_h,_n,_v)
169    for b in torch.arange(_b):
170        for h in torch.arange(_h):
171            for n in torch.arange(_n):
172                for v in torch.arange(_v):
173                    O[b,h,n,v] = torch.dot(weights[b,h,n,:], V[b,h,:,v])
174    return logits, O
175
176def test_mha_par_batched():
177    torch.manual_seed(42)
178    b = 5
179    h = 2
180    n = m = 4
181    k = v = 4
182    Q = torch.randn((b,h,n,k))
183    K = torch.randn((b,h,m,k))
184    V = torch.randn((b,h,m,v))
185    logits_expected, O_expected = mha_par_batched(Q,K,V)
186    logits_actual, O_actual = mha_par_batched_impl(Q,K,V)
187
188    assert(torch.all(torch.isclose(logits_expected, logits_actual)))
189    assert(torch.all(torch.isclose(O_expected, O_actual)))
190
191if __name__ == '__main__':
192    test_tensorprod()
193    test_matmul()
194    test_matmul_reduce_k()
195    test_matmul_reduce_i()
196    test_matmul_reduce_j()
197    test_matmul_reduce_all()
198    test_mha_par_batched()

Dot product Attention (single query)

 1def attn(q,K,V):
 2    """
 3    args:
 4        q: [k]
 5        K: [m,k]
 6        V: [m,v]
 7    returns:
 8        y: [v]
 9    """
10    logits = torch.einsum('k,mk->m',q,K)
11    weights = F.softmax(logits, dim=0)
12    y = torch.einsum('m,mv->v',weights, V)
13    return y
14
15class Attention(torch.nn.Module):
16    def __init__(self, d, k, v):
17        super(Attention, self).__init__()
18        self.Wq = nn.Parameter(torch.randn(d,k))
19        self.Wk = nn.Parameter(torch.randn(d,k))
20        self.Wv = nn.Parameter(torch.randn(d,v))
21
22    def forward(self, x, M):
23        q = torch.einsum('d,dk->k', x, self.Wq)
24        K = torch.einsum('md,dk->mk', M, self.Wk)
25        V = torch.einsum('md,dv->mv', M, self.Wv)
26        return attn(q,K,V)

Multi-head Attention (single query)

 1def mha(q,K,V):
 2    """
 3    args:
 4        q: [h,k]
 5        K: [h,m,k]
 6        V: [h,m,v]
 7    returns:
 8        o: [h,v]
 9    """
10    logits = torch.einsum('hk,hmk->hm',q,K)
11    weights = F.softmax(logits, dim=1)
12    o = torch.einsum('hm,hmv->hv',weights, V)
13    return o
14
15class MultiHeadAttention(torch.nn.Module):
16    def __init__(self, h, d, k, v):
17        super(MultiHeadAttention, self).__init__()
18        self.Wq = nn.Parameter(torch.randn(h,d,k))
19        self.Wk = nn.Parameter(torch.randn(h,d,k))
20        self.Wv = nn.Parameter(torch.randn(h,d,v))
21        self.Wo = nn.Parameter(torch.randn(h,v,d))
22
23    def forward(self, x, M):
24        q = torch.einsum('d,hdk->hk', x, self.Wq)
25        K = torch.einsum('md,hdk->hmk', M, self.Wk)
26        V = torch.einsum('md,hdv->hmv', M, self.Wv)
27        o = mha(q, K, V)
28        y = torch.einsum('hv,hvd->d', o, self.Wo)
29        return y

Multi-head Attention (sequential query)

 1def mha(q,K,V):
 2    """
 3    args:
 4        q: [h,k]
 5        K: [h,m,k]
 6        V: [h,m,v]
 7    returns:
 8        o: [h,v]
 9    """
10    logits = torch.einsum('hk,hmk->hm',q,K)
11    weights = F.softmax(logits, dim=1)
12    o = torch.einsum('hm,hmv->hv', weights, V)
13    return o
14
15class MultiHeadAttentionSequential(torch.nn.Module):
16    def __init__(self, h, d, k, v):
17        super(MultiHeadAttentionSequential, self).__init__()
18        self.Wq = nn.Parameter(torch.randn(h,d,k))
19        self.Wk = nn.Parameter(torch.randn(h,d,k))
20        self.Wv = nn.Parameter(torch.randn(h,d,v))
21        self.Wo = nn.Parameter(torch.randn(h,v,d))
22
23    '''
24    args:
25        x: [d]
26        prev_K: [h,m,k]
27        prev_V: [h,m,v]
28    returns:
29        y: [d]
30        K: [h,m+1,k]
31        V: [h,m+1,v]
32    '''
33    def forward(self, x, prev_K, prev_V):
34        q = torch.einsum('d,hdk->hk', x, self.Wq)
35        K = torch.cat((prev_K, torch.einsum('d,hdk->hk', x, self.Wk).unsqueeze(-2)), dim=-2)
36        V = torch.cat((prev_V, torch.einsum('d,hdv->hv', x, self.Wv).unsqueeze(-2)), dim=-2)
37        o = mha(q, K, V)
38        y = torch.einsum('hv,hvd->d', o, self.Wo)
39        return y, K, V

Masked Multi-head Attention (parallel query)

 1def mha_par(Q,K,V,mask):
 2    """
 3    args:
 4        Q: [h,n,k]
 5        K: [h,m,k]
 6        V: [h,m,v]
 7        mask: [n,m]
 8    returns:
 9        O: [h,n,v]
10    """
11    logits = torch.einsum('hnk,hmk->hnm',Q,K)
12    attn_mask = mask == 0
13    logits.masked_fill_(attn_mask, float('-inf'))
14    weights = F.softmax(logits, dim=-1)
15    O = torch.einsum('hnm,hmv->hnv', weights, V)
16    return O
17
18class MaskedMultiHeadAttentionParallel(torch.nn.Module):
19    def __init__(self, h, d, k, v):
20        super(MaskedMultiHeadAttentionParallel, self).__init__()
21        self.Wq = nn.Parameter(torch.randn(h,d,k))
22        self.Wk = nn.Parameter(torch.randn(h,d,k))
23        self.Wv = nn.Parameter(torch.randn(h,d,v))
24        self.Wo = nn.Parameter(torch.randn(h,v,d))
25
26    def forward(self, X, M, mask):
27        Q = torch.einsum('nd,hdk->hnk', X, self.Wq)
28        K = torch.einsum('md,hdk->hmk', M, self.Wk)
29        V = torch.einsum('md,hdv->hmv', M, self.Wv)
30        O = mha_par(Q, K, V, mask)
31        Y = torch.einsum('hnv,hvd->nd', O, self.Wo)
32        return Y

Masked Multi-head Attention Batched (parallel query)

 1def mha_par_batched(Q,K,V,mask):
 2    """
 3    args:
 4        Q: [b,h,n,k]
 5        K: [b,h,m,k]
 6        V: [b,h,m,v]
 7        mask: [n,m]
 8    returns:
 9        O: [b,h,n,v]
10    """
11    logits = torch.einsum('bhnk,bhmk->bhnm', Q, K)
12    attn_mask = mask == 0
13    logits.masked_fill_(attn_mask, float('-inf'))
14    weights = F.softmax(logits, dim=-1)
15    O = torch.einsum('bhnm,bhmv->bhnv', weights, V)
16    return O
17
18class MaskedMultiHeadAttentionParallelBatched(torch.nn.Module):
19    def __init__(self, h, d, k, v):
20        super(MaskedMultiHeadAttentionParallelBatched, self).__init__()
21        self.Wq = nn.Parameter(torch.randn(h,d,k))
22        self.Wk = nn.Parameter(torch.randn(h,d,k))
23        self.Wv = nn.Parameter(torch.randn(h,d,v))
24        self.Wo = nn.Parameter(torch.randn(h,v,d))
25
26    def forward(self, X, M, mask):
27        Q = torch.einsum('bnd,hdk->bhnk', X, self.Wq)
28        K = torch.einsum('bmd,hdk->bhmk', M, self.Wk)
29        V = torch.einsum('bmd,hdv->bhmv', M, self.Wv)
30        O = mha_par_batched(Q,K,V,mask)
31        Y = torch.einsum('bhnv,hvd->bnd', O, self.Wo)
32        return Y

Multi-head Attention Batched (sequential query)

 1def mha_batched(q,K,V):
 2    """
 3    args:
 4        q: [b,h,k]
 5        K: [b,h,m,k]
 6        V: [b,h,m,v]
 7    returns:
 8        O: [b,h,v]
 9    """
10    logits = torch.einsum('bhk,bhmk->bhm',q,K)
11    weights = F.softmax(logits, dim=-1)
12    O = torch.einsum('bhm,bhmv->bhv', weights, V)
13    return O
14
15class MultiHeadAttentionSequentialBatched(torch.nn.Module):
16    def __init__(self, h, d, k, v):
17        super(MultiHeadAttentionSequentialBatched, self).__init__()
18        self.Wq = nn.Parameter(torch.randn(h,d,k))
19        self.Wk = nn.Parameter(torch.randn(h,d,k))
20        self.Wv = nn.Parameter(torch.randn(h,d,v))
21        self.Wo = nn.Parameter(torch.randn(h,v,d))
22
23    '''
24    args:
25        x: [b,d]
26        prev_K: [b,h,m,k]
27        prev_V: [b,h,m,v]
28    returns:
29        y: [b,d]
30        K: [b,h,m+1,k]
31        V: [b,h,m+1,v]
32    '''
33    def forward(self, x, prev_K, prev_V):
34        Q = torch.einsum('bd,hdk->bhk', x, self.Wq)
35        K = torch.cat((prev_K, torch.einsum('bd,hdk->bhk', x, self.Wk).unsqueeze(-2)), dim=-2)
36        V = torch.cat((prev_V, torch.einsum('bd,hdv->bhv', x, self.Wv).unsqueeze(-2)), dim=-2)
37        O = mha_batched(Q, K, V)
38        Y = torch.einsum('bhv,hvd->bd', O, self.Wo)
39        return Y, K, V

Masked Multi-query Attention Batched (parallel query)

 1def mqa_par_batched(Q,K,V,mask):
 2    """
 3    Args:
 4        Q: [b,h,n,k]
 5        K: [b,m,k]
 6        V: [b,m,v]
 7        mask: [n,m]
 8    Returns:
 9        O: [b,h,n,v]
10    """
11    logits = torch.einsum('bhnk,bmk->bhnm', Q, K)
12    attn_mask = mask == 0
13    logits.masked_fill_(attn_mask, float('-inf'))
14    weights = F.softmax(logits, dim=-1)
15    O = torch.einsum('bhnm,bmv->bhnv', weights, V)
16    return O
17
18class MaskedMultiQueryAttentionParallelBatched(torch.nn.Module):
19    def __init__(self, h, d, k, v):
20        super(MaskedMultiQueryAttentionParallelBatched, self).__init__()
21        self.Wq = nn.Parameter(torch.randn(h,d,k))
22        self.Wk = nn.Parameter(torch.randn(d,k))
23        self.Wv = nn.Parameter(torch.randn(d,v))
24        self.Wo = nn.Parameter(torch.randn(h,v,d))
25
26    def forward(self, X, M, mask):
27        """
28        Args:
29            X: [b,n,d]
30            M: [b,m,d]
31            mask: [n,m]
32        Returns:
33            Y: [b,n,d]
34        """
35        Q = torch.einsum('bnd,hdk->bhnk', X, self.Wq)
36        K = torch.einsum('bmd,dk->bmk', M, self.Wk)
37        V = torch.einsum('bmd,dv->bmv', M, self.Wv)
38        O = mqa_par_batched(Q,K,V,mask)
39        Y = torch.einsum('bhnv,hvd->bnd', O, self.Wo)
40        return Y

Multi-query Attention Batched (sequential query)

 1def mqa_batched(q,K,V):
 2    """
 3    Args:
 4        q: [b,h,k]
 5        k: [b,m,k]
 6        v: [b,m,v]
 7    Returns:
 8        o: [b,h,n,v]
 9    """
10    logits = torch.einsum('bhk,bmk->bhm', q, K)
11    weights = F.softmax(logits, dim=-1)
12    o = torch.einsum('bhm,bmv->bhv', weights, V)
13    return o
14
15class MultiQueryAttentionSequentialBatched(torch.nn.Module):
16    def __init__(self, h, d, k, v):
17        super(MultiQueryAttentionSequentialBatched, self).__init__()
18        self.Wq = nn.Parameter(torch.randn(h,d,k))
19        self.Wk = nn.Parameter(torch.randn(d,k))
20        self.Wv = nn.Parameter(torch.randn(d,v))
21        self.Wo = nn.Parameter(torch.randn(h,v,d))
22
23    def forward(self, x, prev_K, prev_V):
24        """
25        Args:
26            x: [b,d]
27            prev_K: [b,m,k]
28            prev_V: [b,m,v]
29        Returns:
30            y: [b,d]
31            K: [b,m+1,k]
32            V: [b,m+1,v]
33        """
34        q = torch.einsum('bd,hdk->bhk', x, self.Wq)
35        K = torch.cat((prev_K, torch.einsum('bd,dk->bk', x, self.Wk).unsqueeze(-2)), dim=-2)
36        V = torch.cat((prev_V, torch.einsum('bd,dv->bv', x, self.Wv).unsqueeze(-2)), dim=-2)
37        o = mqa_batched(q, K, V)
38        y = torch.einsum('bhv,hvd->bd', o, self.Wo)
39        return y, K, V

Unit Test

  1"""
  2Sample output:
  3--------------------------------------
  4Attention
  5tensor([ -3.0437,  -5.3354,  -2.7996,   4.7690,   6.1953,  -3.1872,   2.4339,
  6         -4.9126,  -0.5149,  -3.6056,   1.6128, -14.4580,  -2.2639,  -2.7896,
  7         -0.7055,   6.9216], grad_fn=<ViewBackward0>)
  8MultiHeadAttention
  9tensor([-18.4005,  11.4970,   5.9840,   9.8845, -15.3181,   1.3615,  -2.5959,
 10         17.3029, -11.3590,  25.8750, -14.3187,  -3.3374,   2.2135, -13.3058,
 11         -1.9368, -15.8990], grad_fn=<ViewBackward0>)
 12MultiHeadAttentionSequential
 13tensor([-18.4005,  11.4970,   5.9840,   9.8845, -15.3181,   1.3615,  -2.5959,
 14         17.3029, -11.3590,  25.8750, -14.3187,  -3.3374,   2.2135, -13.3058,
 15         -1.9368, -15.8990], grad_fn=<ViewBackward0>)
 16MaskedMultiHeadAttentionParallel
 17tensor([-18.4005,  11.4970,   5.9840,   9.8845, -15.3181,   1.3615,  -2.5959,
 18         17.3029, -11.3590,  25.8750, -14.3187,  -3.3374,   2.2135, -13.3058,
 19         -1.9368, -15.8990], requires_grad=True)
 20MultiHeadAttentionSequentialBatched
 21tensor([-18.4005,  11.4970,   5.9840,   9.8845, -15.3181,   1.3615,  -2.5959,
 22         17.3029, -11.3590,  25.8750, -14.3187,  -3.3374,   2.2135, -13.3058,
 23         -1.9368, -15.8990], requires_grad=True)
 24MaskedMultiHeadAttentionParallelBatched
 25tensor([-18.4005,  11.4970,   5.9840,   9.8845, -15.3181,   1.3615,  -2.5959,
 26         17.3029, -11.3590,  25.8750, -14.3187,  -3.3374,   2.2135, -13.3058,
 27         -1.9368, -15.8990], requires_grad=True)
 28MultiQueryAttentionSequentialBatched
 29tensor([ -6.9532,  16.7457, -11.2753,  -6.5196,  22.5019,  13.0422,  -6.6303,
 30         -6.5617,   2.0252,   2.2950,   1.2324,  15.0855, -20.8858,  17.2391,
 31         -8.0939,  -1.4088], requires_grad=True)
 32MaskedMultiQueryAttentionParallelBatched
 33tensor([ -6.9532,  16.7457, -11.2753,  -6.5196,  22.5018,  13.0422,  -6.6303,
 34         -6.5617,   2.0252,   2.2950,   1.2324,  15.0855, -20.8858,  17.2391,
 35         -8.0938,  -1.4088], requires_grad=True)
 36GroupedQueryAttentionSequentialBatched(g=1)
 37tensor([ -6.9532,  16.7457, -11.2753,  -6.5196,  22.5019,  13.0422,  -6.6303,
 38         -6.5617,   2.0252,   2.2950,   1.2324,  15.0855, -20.8858,  17.2391,
 39         -8.0939,  -1.4088], requires_grad=True)
 40GroupedQueryAttentionSequentialBatched(g=2)
 41tensor([-18.4005,  11.4970,   5.9840,   9.8845, -15.3181,   1.3615,  -2.5959,
 42         17.3029, -11.3590,  25.8750, -14.3187,  -3.3374,   2.2135, -13.3058,
 43         -1.9368, -15.8990], requires_grad=True)
 44"""
 45import torch
 46import torch.nn as nn
 47import torch.nn.functional as F
 48
 49def attn(q,K,V):
 50    """
 51    Args:
 52        q: [k]
 53        K: [m,k]
 54        V: [m,v]
 55    Returns:
 56        y: [v]
 57    """
 58    logits = torch.einsum('k,mk->m',q,K)
 59    weights = F.softmax(logits, dim=0)
 60    y = torch.einsum('m,mv->v', weights, V)
 61    return y
 62
 63def mha(q,K,V):
 64    """
 65    Args:
 66        q: [h,k]
 67        K: [h,m,k]
 68        V: [h,m,v]
 69    Returns:
 70        o: [h,v]
 71    """
 72    logits = torch.einsum('hk,hmk->hm',q,K)
 73    weights = F.softmax(logits, dim=-1)
 74    o = torch.einsum('hm,hmv->hv', weights, V)
 75    return o
 76
 77def mha_par(Q,K,V,mask):
 78    """
 79    Args:
 80        Q: [h,n,k]
 81        K: [h,m,k]
 82        V: [h,m,v]
 83        mask: [n,m]
 84    Returns:
 85        O: [h,n,v]
 86    """
 87    logits = torch.einsum('hnk,hmk->hnm',Q,K)
 88    attn_mask = mask == 0
 89    logits.masked_fill_(attn_mask, float('-inf'))
 90    weights = F.softmax(logits, dim=-1)
 91    O = torch.einsum('hnm,hmv->hnv', weights, V)
 92    return O
 93
 94def mha_batched(q,K,V):
 95    """
 96    Args:
 97        q: [b,h,k]
 98        K: [b,h,m,k]
 99        V: [b,h,m,v]
100    Returns:
101        O: [b,h,v]
102    """
103    logits = torch.einsum('bhk,bhmk->bhm',q,K)
104    weights = F.softmax(logits, dim=-1)
105    o = torch.einsum('bhm,bhmv->bhv', weights, V)
106    return o
107
108def gqa_batched(q,K,V):
109    """
110    Args:
111        q: [b,h,k]
112        K: [b,g,m,k]
113        V: [b,g,m,v]
114    Returns:
115        O: [b,h,v]
116    """
117    # TODO: fixit
118    h = q.shape[-2]
119    g = K.shape[-3]
120    r = int(h/g)
121    K = torch.cat([K]*r,dim=-3) # back to dim
122    V = torch.cat([V]*r,dim=-3) # back to dim
123    o = mha_batched(q,K,V)
124    return o
125
126def mqa_batched(q,K,V):
127    """
128    Args:
129        q: [b,h,k]
130        k: [b,m,k]
131        v: [b,m,v]
132    Returns:
133        o: [b,h,n,v]
134    """
135    logits = torch.einsum('bhk,bmk->bhm', q, K)
136    weights = F.softmax(logits, dim=-1)
137    o = torch.einsum('bhm,bmv->bhv', weights, V)
138    return o
139
140def mha_par_batched(Q,K,V,mask):
141    """
142    Args:
143        Q: [b,h,n,k]
144        K: [b,h,m,k]
145        V: [b,h,m,v]
146        mask: [n,m]
147    Returns:
148        O: [b,h,n,v]
149    """
150    logits = torch.einsum('bhnk,bhmk->bhnm', Q, K)
151    attn_mask = mask == 0
152    logits.masked_fill_(attn_mask, float('-inf'))
153    weights = F.softmax(logits, dim=-1)
154    O = torch.einsum('bhnm,bhmv->bhnv', weights, V)
155    return O
156
157def mqa_par_batched(Q,K,V,mask):
158    """
159    Args:
160        Q: [b,h,n,k]
161        K: [b,m,k]
162        V: [b,m,v]
163        mask: [n,m]
164    Returns:
165        O: [b,h,n,v]
166    """
167    logits = torch.einsum('bhnk,bmk->bhnm', Q, K)
168    attn_mask = mask == 0
169    logits.masked_fill_(attn_mask, float('-inf'))
170    weights = F.softmax(logits, dim=-1)
171    O = torch.einsum('bhnm,bmv->bhnv', weights, V)
172    return O
173
174class Attention(torch.nn.Module):
175    def __init__(self, d, k, v):
176        super(Attention, self).__init__()
177        self.Wq = nn.Parameter(torch.randn(d,k))
178        self.Wk = nn.Parameter(torch.randn(d,k))
179        self.Wv = nn.Parameter(torch.randn(d,v))
180
181    def forward(self, x, M):
182        q = torch.einsum('d,dk->k', x, self.Wq)
183        K = torch.einsum('md,dk->mk', M, self.Wk)
184        V = torch.einsum('md,dv->mv', M, self.Wv)
185        y = attn(q,K,V)
186        return y
187
188class MultiHeadAttention(torch.nn.Module):
189    def __init__(self, h, d, k, v):
190        super(MultiHeadAttention, self).__init__()
191        self.Wq = nn.Parameter(torch.randn(h,d,k))
192        self.Wk = nn.Parameter(torch.randn(h,d,k))
193        self.Wv = nn.Parameter(torch.randn(h,d,v))
194        self.Wo = nn.Parameter(torch.randn(h,v,d))
195
196    def forward(self, x, M):
197        q = torch.einsum('d,hdk->hk', x, self.Wq)
198        K = torch.einsum('md,hdk->hmk', M, self.Wk)
199        V = torch.einsum('md,hdv->hmv', M, self.Wv)
200        o = mha(q, K, V)
201        y = torch.einsum('hv,hvd->d', o, self.Wo)
202        return y
203
204class MaskedMultiHeadAttentionParallel(torch.nn.Module):
205    def __init__(self, h, d, k, v):
206        super(MaskedMultiHeadAttentionParallel, self).__init__()
207        self.Wq = nn.Parameter(torch.randn(h,d,k))
208        self.Wk = nn.Parameter(torch.randn(h,d,k))
209        self.Wv = nn.Parameter(torch.randn(h,d,v))
210        self.Wo = nn.Parameter(torch.randn(h,v,d))
211
212    def forward(self, X, M, mask):
213        Q = torch.einsum('nd,hdk->hnk', X, self.Wq)
214        K = torch.einsum('md,hdk->hmk', M, self.Wk)
215        V = torch.einsum('md,hdv->hmv', M, self.Wv)
216        O = mha_par(Q, K, V, mask)
217        Y = torch.einsum('hnv,hvd->nd', O, self.Wo)
218        return Y
219
220class MaskedMultiHeadAttentionParallelBatched(torch.nn.Module):
221    def __init__(self, h, d, k, v):
222        super(MaskedMultiHeadAttentionParallelBatched, self).__init__()
223        self.Wq = nn.Parameter(torch.randn(h,d,k))
224        self.Wk = nn.Parameter(torch.randn(h,d,k))
225        self.Wv = nn.Parameter(torch.randn(h,d,v))
226        self.Wo = nn.Parameter(torch.randn(h,v,d))
227
228    def forward(self, X, M, mask):
229        Q = torch.einsum('bnd,hdk->bhnk', X, self.Wq)
230        K = torch.einsum('bmd,hdk->bhmk', M, self.Wk)
231        V = torch.einsum('bmd,hdv->bhmv', M, self.Wv)
232        O = mha_par_batched(Q,K,V,mask)
233        Y = torch.einsum('bhnv,hvd->bnd', O, self.Wo)
234        return Y
235
236class MultiHeadAttentionSequential(torch.nn.Module):
237    def __init__(self, h, d, k, v):
238        super(MultiHeadAttentionSequential, self).__init__()
239        self.Wq = nn.Parameter(torch.randn(h,d,k))
240        self.Wk = nn.Parameter(torch.randn(h,d,k))
241        self.Wv = nn.Parameter(torch.randn(h,d,v))
242        self.Wo = nn.Parameter(torch.randn(h,v,d))
243
244    def forward(self, x, prev_K, prev_V):
245        """
246        Args:
247            x: [d]
248            prev_K: [h,m,k]
249            prev_V: [h,m,v]
250        Returns:
251            y: [d]
252            K: [h,m+1,k]
253            V: [h,m+1,v]
254        """
255        q = torch.einsum('d,hdk->hk', x, self.Wq)
256        K = torch.cat((prev_K, torch.einsum('d,hdk->hk', x, self.Wk).unsqueeze(-2)), dim=-2)
257        V = torch.cat((prev_V, torch.einsum('d,hdv->hv', x, self.Wv).unsqueeze(-2)), dim=-2)
258        o = mha(q, K, V)
259        y = torch.einsum('hv,hvd->d', o, self.Wo)
260        return y, K, V
261
262class MultiHeadAttentionSequentialBatched(torch.nn.Module):
263    def __init__(self, h, d, k, v):
264        super(MultiHeadAttentionSequentialBatched, self).__init__()
265        self.Wq = nn.Parameter(torch.randn(h,d,k))
266        self.Wk = nn.Parameter(torch.randn(h,d,k))
267        self.Wv = nn.Parameter(torch.randn(h,d,v))
268        self.Wo = nn.Parameter(torch.randn(h,v,d))
269
270    def forward(self, x, prev_K, prev_V):
271        """
272        Args:
273            x: [b,d]
274            prev_K: [b,h,m,k]
275            prev_V: [b,h,m,v]
276        Returns:
277            y: [b,d]
278            K: [b,h,m+1,k]
279            V: [b,h,m+1,v]
280        """
281        q = torch.einsum('bd,hdk->bhk', x, self.Wq)
282        K = torch.cat((prev_K, torch.einsum('bd,hdk->bhk', x, self.Wk).unsqueeze(-2)), dim=-2)
283        V = torch.cat((prev_V, torch.einsum('bd,hdv->bhv', x, self.Wv).unsqueeze(-2)), dim=-2)
284        o = mha_batched(q, K, V)
285        y = torch.einsum('bhv,hvd->bd', o, self.Wo)
286        return y, K, V
287
288class GroupedQueryAttentionSequentialBatched(torch.nn.Module):
289    def __init__(self, h, g, d, k, v):
290        super(GroupedQueryAttentionSequentialBatched, self).__init__()
291        self.Wq = nn.Parameter(torch.randn(h,d,k))
292        self.Wk = nn.Parameter(torch.randn(g,d,k))
293        self.Wv = nn.Parameter(torch.randn(g,d,v))
294        self.Wo = nn.Parameter(torch.randn(h,v,d))
295
296    def forward(self, x, prev_K, prev_V):
297        """
298        Args:
299            x: [b,d]
300            prev_K: [b,g,m,k]
301            prev_V: [b,g,m,v]
302        Returns:
303            y: [b,d]
304            K: [b,m+1,k]
305            V: [b,m+1,v]
306        """
307        q = torch.einsum('bd,hdk->bhk', x, self.Wq)
308        K = torch.cat((prev_K, torch.einsum('bd,gdk->bgk', x, self.Wk).unsqueeze(-2)), dim=-2)
309        V = torch.cat((prev_V, torch.einsum('bd,gdv->bgv', x, self.Wv).unsqueeze(-2)), dim=-2)
310        o = gqa_batched(q, K, V)
311        y = torch.einsum('bhv,hvd->bd', o, self.Wo)
312        return y, K, V
313
314class MaskedMultiQueryAttentionParallelBatched(torch.nn.Module):
315    def __init__(self, h, d, k, v):
316        super(MaskedMultiQueryAttentionParallelBatched, self).__init__()
317        self.Wq = nn.Parameter(torch.randn(h,d,k))
318        self.Wk = nn.Parameter(torch.randn(d,k))
319        self.Wv = nn.Parameter(torch.randn(d,v))
320        self.Wo = nn.Parameter(torch.randn(h,v,d))
321
322    def forward(self, X, M, mask):
323        """
324        Args:
325            X: [b,n,d]
326            M: [b,m,d]
327            mask: [n,m]
328        Returns:
329            Y: [b,n,d]
330        """
331        Q = torch.einsum('bnd,hdk->bhnk', X, self.Wq)
332        K = torch.einsum('bmd,dk->bmk', M, self.Wk)
333        V = torch.einsum('bmd,dv->bmv', M, self.Wv)
334        O = mqa_par_batched(Q,K,V,mask)
335        Y = torch.einsum('bhnv,hvd->bnd', O, self.Wo)
336        return Y
337
338class MultiQueryAttentionSequentialBatched(torch.nn.Module):
339    def __init__(self, h, d, k, v):
340        super(MultiQueryAttentionSequentialBatched, self).__init__()
341        self.Wq = nn.Parameter(torch.randn(h,d,k))
342        self.Wk = nn.Parameter(torch.randn(d,k))
343        self.Wv = nn.Parameter(torch.randn(d,v))
344        self.Wo = nn.Parameter(torch.randn(h,v,d))
345
346    def forward(self, x, prev_K, prev_V):
347        """
348        Args:
349            x: [b,d]
350            prev_K: [b,m,k]
351            prev_V: [b,m,v]
352        Returns:
353            y: [b,d]
354            K: [b,m+1,k]
355            V: [b,m+1,v]
356        """
357        q = torch.einsum('bd,hdk->bhk', x, self.Wq)
358        K = torch.cat((prev_K, torch.einsum('bd,dk->bk', x, self.Wk).unsqueeze(-2)), dim=-2)
359        V = torch.cat((prev_V, torch.einsum('bd,dv->bv', x, self.Wv).unsqueeze(-2)), dim=-2)
360        o = mqa_batched(q, K, V)
361        y = torch.einsum('bhv,hvd->bd', o, self.Wo)
362        return y, K, V
363
364def test_attn(Attention, d, k, M):
365    model = Attention(d,k,d)
366    y = model(M[0],M)
367    print(f'Attention\n{y}')
368
369def test_mha(MultiHeadAttention, d, k, v, h, M):
370    torch.manual_seed(42)
371    model = MultiHeadAttention(h,d,k,v)
372    y = model(M[-1],M)
373    with torch.no_grad():
374        print(f'MultiHeadAttention\n{y}')
375
376def test_mha_seq(MultiHeadAttentionSequential, d, k, v, h, M):
377    torch.manual_seed(42)
378    model = MultiHeadAttentionSequential(h,d,k,v)
379    prev_K = torch.FloatTensor(h,0,k)
380    prev_V = torch.FloatTensor(h,0,v)
381    y = None
382
383    for x in M:
384        y, prev_K, prev_V = model(x, prev_K, prev_V)
385    
386    with torch.no_grad():
387        print(f'MultiHeadAttentionSequential\n{y}')
388
389def test_mha_par(MaskedMultiHeadAttentionParallel, d, k, v, h, M, mask):
390    torch.manual_seed(42)
391    model = MaskedMultiHeadAttentionParallel(h,d,k,v)
392    Y = model(M,M,mask)
393    with torch.no_grad():
394        print(f'MaskedMultiHeadAttentionParallel\n{Y[-1]}')
395
396def test_mha_par_batched(MaskedMultiHeadAttentionParallelBatched, d, k, v, h, M, mask):
397    torch.manual_seed(42)
398    model = MaskedMultiHeadAttentionParallelBatched(h,d,k,v)
399    X = M.unsqueeze(0)
400    Y = model(X,X,mask)
401    with torch.no_grad():
402        Y = Y.squeeze(0)
403        print(f'MaskedMultiHeadAttentionParallelBatched\n{Y[-1]}')
404
405def test_mha_seq_batched(MultiHeadAttentionSequentialBatched, d, k, v, h, M):
406    torch.manual_seed(42)
407    model = MultiHeadAttentionSequentialBatched(h,d,k,v)
408    prev_K = torch.FloatTensor(1,h,0,k)
409    prev_V = torch.FloatTensor(1,h,0,v)
410    y = None
411
412    for x in M:
413        y, prev_K, prev_V = model(x.unsqueeze(0), prev_K, prev_V)
414
415    with torch.no_grad():
416        y = y.squeeze(0)
417        print(f'MultiHeadAttentionSequentialBatched\n{y}')
418
419def test_mqa_par_batched(MaskedMultiQueryAttentionParallelBatched, d, k, v, h, M, mask):
420    torch.manual_seed(42)
421    model = MaskedMultiQueryAttentionParallelBatched(h,d,k,v)
422    X = M.unsqueeze(0)
423    Y = model(X,X,mask)
424    with torch.no_grad():
425        Y = Y.squeeze(0)
426        print(f'MaskedMultiQueryAttentionParallelBatched\n{Y[-1]}')
427
428def test_mqa_seq_batched(MultiQueryAttentionSequentialBatched, d, k, v, h, M):
429    torch.manual_seed(42)
430    model = MultiQueryAttentionSequentialBatched(h,d,k,v)
431    prev_K = torch.FloatTensor(1,0,k)
432    prev_V = torch.FloatTensor(1,0,v)
433    y = None
434
435    for x in M:
436        y, prev_K, prev_V = model(x.unsqueeze(0), prev_K, prev_V)
437
438    with torch.no_grad():
439        y = y.squeeze(0)
440        print(f'MultiQueryAttentionSequentialBatched\n{y}')
441
442def test_gqa1_seq_batched(GroupedQueryAttentionSequentialBatched, d, k, v, h, M):
443    torch.manual_seed(42)
444    g = 1
445    model = GroupedQueryAttentionSequentialBatched(h,g,d,k,v)
446    prev_K = torch.FloatTensor(1,g,0,k)
447    prev_V = torch.FloatTensor(1,g,0,v)
448    y = None
449
450    for x in M:
451        y, prev_K, prev_V = model(x.unsqueeze(0), prev_K, prev_V)
452
453    with torch.no_grad():
454        y = y.squeeze(0)
455        print(f'GroupedQueryAttentionSequentialBatched(g={g})\n{y}')
456
457def test_gqah_seq_batched(GroupedQueryAttentionSequentialBatched, d, k, v, h, M):
458    torch.manual_seed(42)
459    g = h
460    model = GroupedQueryAttentionSequentialBatched(h,g,d,k,v)
461    prev_K = torch.FloatTensor(1,g,0,k)
462    prev_V = torch.FloatTensor(1,g,0,v)
463    y = None
464
465    for x in M:
466        y, prev_K, prev_V = model(x.unsqueeze(0), prev_K, prev_V)
467
468    with torch.no_grad():
469        y = y.squeeze(0)
470        print(f'GroupedQueryAttentionSequentialBatched(g={g})\n{y}')
471
472if __name__ == '__main__':    
473    m = 10
474    d = 16
475    k = 8
476    v = 8
477    h = 2
478
479    M = torch.randn((m,d))
480    # triangular mask mimicing the decoder
481    # same code can be reused for encoder with all bits on
482    mask = torch.tril(torch.ones(M.shape[0],M.shape[0]))
483
484    #########################################
485    # SHA
486    #########################################
487    test_attn(Attention, d, k, M)
488
489    #########################################
490    # MHA
491    #########################################
492    test_mha(MultiHeadAttention, d, k, v, h, M)
493    test_mha_seq(MultiHeadAttentionSequential, d, k, v, h, M)
494    test_mha_par(MaskedMultiHeadAttentionParallel, d, k, v, h, M, mask)    
495    test_mha_seq_batched(MultiHeadAttentionSequentialBatched, d, k, v, h, M)
496    test_mha_par_batched(MaskedMultiHeadAttentionParallelBatched, d, k, v, h, M, mask)
497
498    #########################################
499    # MQA
500    #########################################    
501    test_mqa_seq_batched(MultiQueryAttentionSequentialBatched, d, k, v, h, M)
502    test_mqa_par_batched(MaskedMultiQueryAttentionParallelBatched, d, k, v, h, M, mask)
503
504    #########################################
505    # GQA
506    #########################################    
507    test_gqa1_seq_batched(GroupedQueryAttentionSequentialBatched, d, k, v, h, M)
508    test_gqah_seq_batched(GroupedQueryAttentionSequentialBatched, d, k, v, h, M)