attn.py 10.1 KB
Newer Older
1
import os
Casper Hansen's avatar
Casper Hansen committed
2
import math
Haotian Tang's avatar
Haotian Tang committed
3
4
import torch
import torch.nn as nn
Casper Hansen's avatar
Casper Hansen committed
5
from torch.nn import functional as F
Casper Hansen's avatar
Casper Hansen committed
6
from awq.modules.fused.cache import WindowedCache
7
from awq.utils.fused_utils import get_attention_shapes
Casper Hansen's avatar
Casper Hansen committed
8

9

Casper's avatar
Casper committed
10
try:
11
    import awq_ft_ext
Casper's avatar
Casper committed
12
13
14
    FT_INSTALLED = True
except:
    FT_INSTALLED = False
qwopqwop200's avatar
qwopqwop200 committed
15

16
17
18
19
20
21
22
23
24
HF_NEW_CACHE_FORMAT = False

import transformers
# https://github.com/huggingface/transformers/pull/26681 introduced a new cache format
HF_NEW_CACHE_FORMAT = hasattr(transformers, "cache_utils")
if HF_NEW_CACHE_FORMAT:
    from transformers.cache_utils import DynamicCache


Casper Hansen's avatar
Casper Hansen committed
25
class RoPE(nn.Module):
Casper's avatar
Casper committed
26
    def __init__(self, hidden_size, n_heads, max_seq_len, device, rope_theta):
Casper Hansen's avatar
Casper Hansen committed
27
28
29
        super(RoPE, self).__init__()
        
        self.freqs_cis = nn.Parameter(
Casper's avatar
Casper committed
30
31
32
            self.precompute_freqs_cis(
                hidden_size // n_heads, max_seq_len * 2, rope_theta
            ).to(device),
Casper Hansen's avatar
Casper Hansen committed
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
            requires_grad=False
        )

    @staticmethod
    def precompute_freqs_cis(dim: int, end: int, theta=10000.0):
        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
        t = torch.arange(end)
        freqs = torch.outer(t, freqs).float()
        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
        return freqs_cis

    @staticmethod
    def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
        ndim = x.ndim
        assert 0 <= 1 < ndim
        assert freqs_cis.shape == (x.shape[1], x.shape[-1])
        shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
        return freqs_cis.view(*shape)

    def forward(self, xq: torch.Tensor, xk: torch.Tensor, start_pos: int, seqlen: int):
        xq_ = torch.view_as_complex(
            xq.float().reshape(*xq.shape[:-1], 2, -1).transpose(-2, -1).contiguous()
        )
        xk_ = torch.view_as_complex(
            xk.float().reshape(*xk.shape[:-1], 2, -1).transpose(-2, -1).contiguous()
        )
        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
        freqs_cis = self.reshape_for_broadcast(freqs_cis, xq_).to(xq_.device)
        
        xq_out = torch.view_as_real(xq_ * freqs_cis).transpose(-2, -1).flatten(3)
        xk_out = torch.view_as_real(xk_ * freqs_cis).transpose(-2, -1).flatten(3)
        
        return xq_out.type_as(xq), xk_out.type_as(xk)
Casper Hansen's avatar
Casper Hansen committed
66

Casper Hansen's avatar
Casper Hansen committed
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
class ALiBi(nn.Module):
    def __init__(self, n_heads, max_seq_len, device, alibi_bias_max=8):
        super(ALiBi, self).__init__()
        
        # Initialize ALiBi slopes and bias
        slopes, bias = self.build_alibi_bias(n_heads, max_seq_len, alibi_bias_max=alibi_bias_max)
        self.slopes = nn.Parameter(slopes.float().to(device), requires_grad=False)
        self.bias = nn.Parameter(bias.float().to(device), requires_grad=False)

    @staticmethod
    def gen_slopes(n_heads, alibi_bias_max=8):
        _n_heads = 2 ** math.ceil(math.log2(n_heads))
        m = torch.arange(1, _n_heads + 1, dtype=torch.float32)
        m = m.mul(alibi_bias_max / _n_heads)
        slopes = 1.0 / torch.pow(2, m)
        
        if _n_heads != n_heads:
            slopes = torch.cat([slopes[1::2], slopes[::2]])[:n_heads]
            
        return slopes.view(1, n_heads, 1, 1)

    @staticmethod
    def build_alibi_bias(n_heads, seq_len, alibi_bias_max=8, dtype=torch.float32):
        alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32).view(1, 1, 1, seq_len)
        slopes = ALiBi.gen_slopes(n_heads, alibi_bias_max)
        alibi_bias = alibi_bias * slopes
        slopes = slopes.squeeze(0).squeeze(-1).squeeze(-1)
        return slopes.to(dtype=dtype), alibi_bias.to(dtype=dtype)
    
    def forward(self, scores, seqlen):
        scores += self.bias[..., :seqlen]
        return scores
Casper Hansen's avatar
Casper Hansen committed
99
100

class QuantAttentionFused(nn.Module):
Casper Hansen's avatar
Casper Hansen committed
101
    def __init__(self, hidden_size, n_heads, n_kv_heads, qkv_layer, o_proj, dev, max_seq_len, 
Casper's avatar
Casper committed
102
                       use_alibi=False, attention_shapes=None, rope_theta=10000):
Casper Hansen's avatar
Casper Hansen committed
103
104
        super().__init__()
        self.hidden_size = hidden_size
Casper Hansen's avatar
Casper Hansen committed
105
106
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
107
        self.n_kv_groups = n_heads // n_kv_heads if n_kv_heads != 0 else 0
Casper Hansen's avatar
Casper Hansen committed
108
        self.head_dim = self.hidden_size // n_heads
Casper Hansen's avatar
Casper Hansen committed
109
110
111
        self.qkv_proj = qkv_layer
        self.o_proj = o_proj
        self.start_pos = 0
Casper Hansen's avatar
Casper Hansen committed
112
        self.use_alibi = use_alibi
113
        self.cache_batch_size = int(os.getenv("AWQ_BATCH_SIZE", "1"))
114
        self.max_seq_len = max_seq_len
115
        self.is_hf_transformers = False
Casper's avatar
Casper committed
116
        self.rope_theta = rope_theta
Casper Hansen's avatar
Casper Hansen committed
117
118
119
120
121
122
123

        # attention shapes for self attention
        self.attention_shapes = get_attention_shapes(
            attention_shapes, max_seq_len, self.cache_batch_size, n_heads, n_kv_heads, self.head_dim
        )
        # cache store that rolls cache
        self.cache = WindowedCache(
Casper's avatar
Casper committed
124
            self.attention_shapes["cache_v"], self.attention_shapes["cache_k"], self.max_seq_len, dev
Casper Hansen's avatar
Casper Hansen committed
125
        )
Casper Hansen's avatar
Casper Hansen committed
126

127
        if use_alibi:
Casper Hansen's avatar
Casper Hansen committed
128
            self.alibi = ALiBi(n_heads, max_seq_len, dev)
129
130
131
            self.rotary_dim = 0
            self.is_neox = False
        else:
Casper Hansen's avatar
Casper Hansen committed
132
            self.alibi = None
Casper's avatar
Casper committed
133
            self.rope = RoPE(hidden_size, n_heads, max_seq_len, dev, rope_theta)
134
135
136
            self.rotary_dim = self.head_dim
            self.is_neox = True
    
Casper Hansen's avatar
Casper Hansen committed
137
    def forward(self, hidden_states:torch.Tensor, attention_mask=None, *args, **kwargs):
Casper Hansen's avatar
Casper Hansen committed
138
        bsz, seqlen, _ = hidden_states.shape
139

Casper's avatar
Casper committed
140
        # Reallocate cache if batch size changes
141
        if bsz != self.cache_batch_size:
Casper's avatar
Casper committed
142
143
144
145
146
147
            if bsz > self.cache_batch_size:
                self.cache.increase_batch_size(bsz)
                self.cache_batch_size = bsz
            elif bsz < self.cache_batch_size:
                self.cache.decrease_batch_size(bsz)
                self.cache_batch_size = bsz
148
149
150
151
152
153
154

            # Always reset to 0
            self.start_pos = 0 

        # In case we re-generate, we need to refresh the starting position 
        # to 0. We detect it by checking if `past_key_values` is set to None, 
        # which indicates that we are on the first step of `generate()`.
155
156
        # This is only applicable for `transformers` integration
        if self.is_hf_transformers and "past_key_value" in kwargs and kwargs["past_key_value"] is None:
157
158
            self.start_pos = 0

Casper Hansen's avatar
Casper Hansen committed
159
        xqkv = self.qkv_proj(hidden_states)
160
        xqkv = xqkv.view((bsz, seqlen) + self.attention_shapes["xqkv_view"])
Casper Hansen's avatar
Casper Hansen committed
161
        
162
163
164
        xq = self.attention_shapes["xq_slice"](xqkv)
        xk = self.attention_shapes["xk_slice"](xqkv)
        xv = self.attention_shapes["xv_slice"](xqkv)
Haotian Tang's avatar
Haotian Tang committed
165

Casper's avatar
Casper committed
166
        if seqlen > 1 or not FT_INSTALLED:
Casper Hansen's avatar
Casper Hansen committed
167
            xq = xq.view((bsz, seqlen) + self.attention_shapes["xq_view"])
168
169
            xk = xk.view((bsz, seqlen) + self.attention_shapes["xk_view"])
            xv = xv.view((bsz, seqlen) + self.attention_shapes["xv_view"])
Haotian Tang's avatar
Haotian Tang committed
170

171
            if not self.use_alibi:
Casper Hansen's avatar
Casper Hansen committed
172
                xq, xk = self.rope.forward(xq, xk, self.start_pos, seqlen)
Haotian Tang's avatar
Haotian Tang committed
173

Casper Hansen's avatar
Casper Hansen committed
174
            self.cache.to(xq)
Haotian Tang's avatar
Haotian Tang committed
175

Casper Hansen's avatar
Casper Hansen committed
176
177
            values_store = xv.transpose(2, 1)
            keys_store = (
Casper Hansen's avatar
Casper Hansen committed
178
                xk.reshape((bsz, seqlen) + self.attention_shapes["xk_reshape"])
Casper Hansen's avatar
Casper Hansen committed
179
180
181
                .permute(0, 2, 3, 1, 4)
                .contiguous()
            )
Casper Hansen's avatar
Casper Hansen committed
182
            
Casper Hansen's avatar
Casper Hansen committed
183
            self.cache.update_kv(values_store, keys_store, bsz, self.start_pos, seqlen)
Casper Hansen's avatar
Casper Hansen committed
184

Casper's avatar
Casper committed
185
            # Only necessary to retrieve from cache when we are not processing context
qwopqwop200's avatar
fix bug  
qwopqwop200 committed
186
            if seqlen == 1:
Casper Hansen's avatar
Casper Hansen committed
187
                xv, xk = self.cache.get_kv(bsz, self.start_pos, seqlen, self.head_dim)
188

Casper's avatar
Casper committed
189
            
Casper Hansen's avatar
Casper Hansen committed
190
191
            keys = xk
            values = xv
192
193
194
195
196

            if self.n_kv_groups != 0:
                keys = torch.repeat_interleave(keys, dim=2, repeats=self.n_kv_groups)
                values = torch.repeat_interleave(values, dim=2, repeats=self.n_kv_groups)
            
Casper Hansen's avatar
Casper Hansen committed
197
198
199
200
201
202
            xq = xq.transpose(1, 2)
            keys = keys.transpose(1, 2)
            values = values.transpose(1, 2)
            scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)

            if self.use_alibi:
Casper Hansen's avatar
Casper Hansen committed
203
                scores = self.alibi.forward(scores, seqlen)
Casper Hansen's avatar
Casper Hansen committed
204

205
206
            # When seqlen is 1, there is nothing else to attend to
            if attention_mask is not None and seqlen > 1:
Casper Hansen's avatar
Casper Hansen committed
207
208
209
210
                scores = scores + attention_mask  # (bs, n_local_heads, slen, cache_len + slen)
            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
            output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
            attention_weight = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
Casper Hansen's avatar
Casper Hansen committed
211
        else:
212
213
214
215
            xq = xq.view((bsz,) + self.attention_shapes["single_xq_view"])
            xk = xk.view((bsz,) + self.attention_shapes["single_xk_view"])
            xv = xv.view((bsz,) + self.attention_shapes["single_xv_view"])

Casper Hansen's avatar
Casper Hansen committed
216
            alibi_slopes = self.alibi.slopes if self.alibi is not None else None
217
            attention_weight = awq_ft_ext.single_query_attention(
Casper Hansen's avatar
Casper Hansen committed
218
219
220
                xq, # query
                xk, # key
                xv, # value
Casper Hansen's avatar
Casper Hansen committed
221
222
                self.cache.k, # key cache
                self.cache.v, # value cache
Casper Hansen's avatar
Casper Hansen committed
223
                None, # length per sample
Casper Hansen's avatar
Casper Hansen committed
224
                alibi_slopes, # alibi slopes
Casper Hansen's avatar
Casper Hansen committed
225
226
                self.start_pos, # timestep
                self.rotary_dim, # rotary embedding dimension
Casper's avatar
Casper committed
227
                self.rope_theta, # rotary embedding base
228
                self.is_neox, # is neox
Casper Hansen's avatar
Casper Hansen committed
229
            )
Casper Hansen's avatar
Casper Hansen committed
230
            attention_weight = attention_weight.reshape(bsz, 1, -1)
Casper Hansen's avatar
Casper Hansen committed
231
        
Casper Hansen's avatar
Casper Hansen committed
232
        attn_output = self.o_proj(attention_weight)
Casper Hansen's avatar
Casper Hansen committed
233
        self.start_pos += seqlen
Haotian Tang's avatar
Haotian Tang committed
234

Casper Hansen's avatar
Casper Hansen committed
235
        # past_key_value is replaced with cache_v, cache_k, returning empty data
236
237
238
        # we pass a dummy past kv cache for transformers to be able to retrieve the correct info 
        # about past key length
        past_key_value = [torch.zeros(1, 1, self.start_pos, 1)]
239
240
241
242
243
244

        if HF_NEW_CACHE_FORMAT and self.is_hf_transformers:
            new_cache = DynamicCache()
            new_cache.update(past_key_value[0], past_key_value[0], layer_idx=0)
            past_key_value = new_cache

245
        return attn_output, attention_weight, past_key_value