mha.py 33.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# Copyright (c) 2022, Tri Dao.

import math
from functools import partial

import torch
import torch.nn as nn
import torch.nn.functional as F

from einops import rearrange

try:
    from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_func
    from flash_attn.flash_attn_interface import flash_attn_unpadded_kvpacked_func
except ImportError:
    flash_attn_unpadded_qkvpacked_func, flash_attn_unpadded_kvpacked_func = None, None

try:
    from flash_attn.ops.flash_attn_triton import flash_attn_qkvpacked_func, flash_attn_kvpacked_func
except ImportError:
    flash_attn_qkvpacked_func, flash_attn_kvpacked_func = None, None

try:
Tri Dao's avatar
Tri Dao committed
24
    from flash_attn.ops.fused_dense import FusedDense, ColumnParallelLinear, RowParallelLinear
25
except ImportError:
Tri Dao's avatar
Tri Dao committed
26
    FusedDense, ColumnParallelLinear, RowParallelLinear = None, None, None
27
28
29
30
31
32

try:
    from flash_attn.layers.rotary import RotaryEmbedding
except ImportError:
    RotaryEmbedding = None

33
34
35
36
37
try:
    import ft_attention
except ImportError:
    ft_attention = None

38
39
40
41
42
43
44
45
46
47
48
49

class FlashSelfAttention(nn.Module):
    """Implement the scaled dot product attention with softmax.
    Arguments
    ---------
        softmax_scale: The temperature to use for the softmax attention.
                      (default: 1/sqrt(d_keys) where d_keys is computed at
                      runtime)
        attention_dropout: The dropout rate to apply to the attention
                           (default: 0.0)
    """
    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0,
Tri Dao's avatar
Tri Dao committed
50
                 triton=False):
51
52
53
54
55
56
57
        super().__init__()
        if attention_dropout != 0.0 or not triton:
            assert flash_attn_unpadded_qkvpacked_func is not None, 'FlashAttention is not installed'
        if attention_dropout == 0.0 and triton:
            assert flash_attn_qkvpacked_func is not None, 'FlashAttention Triton is not installed'
        self.causal = causal
        self.softmax_scale = softmax_scale
58
        self.drop = nn.Dropout(attention_dropout)
59
60
        self.triton = triton

Tri Dao's avatar
Tri Dao committed
61
    def forward(self, qkv, causal=None, cu_seqlens=None, max_seqlen=None):
62
63
64
        """Implements the multihead softmax attention.
        Arguments
        ---------
Tri Dao's avatar
Tri Dao committed
65
66
67
68
            qkv: The tensor containing the query, key, and value.
                If cu_seqlens is None and max_seqlen is None, then qkv has shape (B, S, 3, H, D).
                If cu_seqlens is not None and max_seqlen is not None, then qkv has shape
                (total, 3, H, D), where total is the sum of the sequence lengths in the batch.
Tri Dao's avatar
Tri Dao committed
69
            causal: if passed, will override self.causal
Tri Dao's avatar
Tri Dao committed
70
71
72
73
74
75
76
            cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
                of the sequences in the batch, used to index into qkv.
            max_seqlen: int. Maximum sequence length in the batch.
        Returns:
        --------
            out: (total, H, D) if cu_seqlens is not None and max_seqlen is not None,
                else (B, S, H, D).
77
78
79
        """
        assert qkv.dtype in [torch.float16, torch.bfloat16]
        assert qkv.is_cuda
Tri Dao's avatar
Tri Dao committed
80
        causal = self.causal if causal is None else causal
Tri Dao's avatar
Tri Dao committed
81
82
83
84
85
86
        unpadded = cu_seqlens is not None
        if unpadded:
            assert cu_seqlens.dtype == torch.int32
            assert max_seqlen is not None
            assert isinstance(max_seqlen, int)
            return flash_attn_unpadded_qkvpacked_func(
87
                qkv, cu_seqlens, max_seqlen, self.drop.p if self.training else 0.0,
Tri Dao's avatar
Tri Dao committed
88
                softmax_scale=self.softmax_scale, causal=causal
89
            )
Tri Dao's avatar
Tri Dao committed
90
91
92
        else:
            batch_size, seqlen = qkv.shape[0], qkv.shape[1]
            # Triton version doesn't support dropout
93
            if self.triton and (self.drop.p == 0 or not self.training):
Tri Dao's avatar
Tri Dao committed
94
                output = flash_attn_qkvpacked_func(qkv, None, causal, self.softmax_scale)
Tri Dao's avatar
Tri Dao committed
95
96
97
98
99
100
            else:
                qkv = rearrange(qkv, 'b s ... -> (b s) ...')
                max_seqlen = seqlen
                cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32,
                                        device=qkv.device)
                output = flash_attn_unpadded_qkvpacked_func(
101
                    qkv, cu_seqlens, max_seqlen, self.drop.p if self.training else 0.0,
Tri Dao's avatar
Tri Dao committed
102
                    softmax_scale=self.softmax_scale, causal=causal
Tri Dao's avatar
Tri Dao committed
103
104
105
                )
                output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
            return output
106
107
108
109
110
111
112
113
114
115
116
117
118


class FlashCrossAttention(nn.Module):
    """Implement the scaled dot product attention with softmax.
    Arguments
    ---------
        softmax_scale: The temperature to use for the softmax attention.
                      (default: 1/sqrt(d_keys) where d_keys is computed at
                      runtime)
        attention_dropout: The dropout rate to apply to the attention
                           (default: 0.0)
    """
    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0,
Tri Dao's avatar
Tri Dao committed
119
                 triton=False):
120
121
122
123
124
125
126
        super().__init__()
        if attention_dropout != 0.0 or not triton:
            assert flash_attn_unpadded_kvpacked_func is not None, 'FlashAttention is not installed'
        if attention_dropout == 0.0 and triton:
            assert flash_attn_kvpacked_func is not None, 'FlashAttention Triton is not installed'
        self.causal = causal
        self.softmax_scale = softmax_scale
127
        self.drop = nn.Dropout(attention_dropout)
128
129
        self.triton = triton

Tri Dao's avatar
Tri Dao committed
130
131
    def forward(self, q, kv, causal=None, cu_seqlens=None, max_seqlen=None,
                cu_seqlens_k=None, max_seqlen_k=None):
132
133
134
135
136
        """Implements the multihead softmax attention.
        Arguments
        ---------
            q: The tensor containing the query. (B, Sq, H, D)
            kv: The tensor containing the key and value. (B, Sk, 2, H, D)
Tri Dao's avatar
Tri Dao committed
137
            causal: if passed, will override self.causal
138
139
140
141
142
143
            cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
                of the sequences in the batch, used to index into q.
            max_seqlen: int. Maximum sequence length in the batch of q.
            cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
                of the sequences in the batch, used to index into kv.
            max_seqlen_k: int. Maximum sequence length in the batch of k and v.
144
145
146
        """
        assert q.dtype in [torch.float16, torch.bfloat16]
        assert q.is_cuda and kv.is_cuda
Tri Dao's avatar
Tri Dao committed
147
        causal = self.causal if causal is None else causal
148
149
150
151
152
153
154
155
156
157
158
        unpadded = cu_seqlens is not None
        if unpadded:
            assert cu_seqlens.dtype == torch.int32
            assert max_seqlen is not None
            assert isinstance(max_seqlen, int)
            assert cu_seqlens_k is not None
            assert cu_seqlens_k.dtype == torch.int32
            assert max_seqlen_k is not None
            assert isinstance(max_seqlen, int)
            return flash_attn_unpadded_kvpacked_func(
                q, kv, cu_seqlens, cu_seqlens_k, max_seqlen, max_seqlen_k,
159
                self.drop.p if self.training else 0.0,
Tri Dao's avatar
Tri Dao committed
160
                softmax_scale=self.softmax_scale, causal=causal
161
            )
162
163
164
165
        else:
            batch_size, seqlen_q = q.shape[0], q.shape[1]
            seqlen_k = kv.shape[1]
            assert kv.shape[0] == batch_size and kv.shape[3] == q.shape[2] and kv.shape[4] == q.shape[3]
166
            if self.triton and (self.drop.p == 0.0 or not self.training):  # Triton version doesn't support dropout
Tri Dao's avatar
Tri Dao committed
167
                output = flash_attn_kvpacked_func(q, kv, None, causal, self.softmax_scale)
168
169
170
171
172
173
174
175
176
            else:
                q = rearrange(q, 'b s ... -> (b s) ...')
                kv = rearrange(kv, 'b s ... -> (b s) ...')
                cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q,
                                            dtype=torch.int32, device=q.device)
                cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, step=seqlen_k,
                                            dtype=torch.int32, device=kv.device)
                output = flash_attn_unpadded_kvpacked_func(
                    q, kv, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k,
177
                    self.drop.p if self.training else 0.0,
Tri Dao's avatar
Tri Dao committed
178
                    softmax_scale=self.softmax_scale, causal=causal
179
180
181
                )
                output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
            return output
182
183
184
185
186
187
188
189
190
191
192
193


class SelfAttention(nn.Module):
    """Implement the scaled dot product attention with softmax.
    Arguments
    ---------
        softmax_scale: The temperature to use for the softmax attention.
                      (default: 1/sqrt(d_keys) where d_keys is computed at
                      runtime)
        attention_dropout: The dropout rate to apply to the attention
                           (default: 0.0)
    """
Tri Dao's avatar
Tri Dao committed
194
    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0):
195
196
197
        super().__init__()
        self.causal = causal
        self.softmax_scale = softmax_scale
198
        self.drop = nn.Dropout(attention_dropout)
199

Tri Dao's avatar
Tri Dao committed
200
    def forward(self, qkv, causal=None, key_padding_mask=None):
201
202
203
204
        """Implements the multihead softmax attention.
        Arguments
        ---------
            qkv: The tensor containing the query, key, and value. (B, S, 3, H, D)
Tri Dao's avatar
Tri Dao committed
205
            causal: if passed, will override self.causal
Tri Dao's avatar
Tri Dao committed
206
207
            key_padding_mask: boolean mask to apply to the attention weights. True means to keep,
                False means to mask out. (B, S)
208
209
        """
        batch_size, seqlen = qkv.shape[0], qkv.shape[1]
Tri Dao's avatar
Tri Dao committed
210
        causal = self.causal if causal is None else causal
211
212
213
        q, k, v = qkv.unbind(dim=2)
        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
        scores = torch.einsum('bthd,bshd->bhts', q, k * softmax_scale)
Tri Dao's avatar
Tri Dao committed
214
215
216
217
218
219
        if key_padding_mask is not None:
            padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype,
                                      device=scores.device)
            padding_mask.masked_fill_(key_padding_mask, 0.0)
            # TD [2022-09-30]: Adding is faster than masked_fill_ (idk why, just better kernel I guess)
            scores = scores + rearrange(padding_mask, 'b s -> b 1 1 s')
Tri Dao's avatar
Tri Dao committed
220
        if causal:
221
222
223
224
225
226
            # "triu_tril_cuda_template" not implemented for 'BFloat16'
            # So we have to construct the mask in float
            causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)
            # TD [2022-09-30]: Adding is faster than masked_fill_ (idk why, just better kernel I guess)
            scores = scores + causal_mask.to(dtype=scores.dtype)
        attention = torch.softmax(scores, dim=-1, dtype=v.dtype)
227
        attention_drop = self.drop(attention)
228
229
230
231
232
233
234
235
236
237
238
239
240
241
        output = torch.einsum('bhts,bshd->bthd', attention_drop, v)
        return output


class CrossAttention(nn.Module):
    """Implement the scaled dot product attention with softmax.
    Arguments
    ---------
        softmax_scale: The temperature to use for the softmax attention.
                      (default: 1/sqrt(d_keys) where d_keys is computed at
                      runtime)
        attention_dropout: The dropout rate to apply to the attention
                           (default: 0.0)
    """
Tri Dao's avatar
Tri Dao committed
242
    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0):
243
244
245
        super().__init__()
        self.causal = causal
        self.softmax_scale = softmax_scale
246
        self.drop = nn.Dropout(attention_dropout)
247

Tri Dao's avatar
Tri Dao committed
248
    def forward(self, q, kv, causal=None, key_padding_mask=None):
249
250
251
252
253
        """Implements the multihead softmax attention.
        Arguments
        ---------
            q: The tensor containing the query. (B, Sq, H, D)
            kv: The tensor containing the key and value. (B, Sk, 2, H, D)
Tri Dao's avatar
Tri Dao committed
254
            causal: if passed, will override self.causal
255
256
            key_padding_mask: boolean mask to apply to the attention weights. True means to keep,
                False means to mask out. (B, Sk)
257
258
        """
        batch_size, seqlen_q = q.shape[0], q.shape[1]
Tri Dao's avatar
Tri Dao committed
259
        causal = self.causal if causal is None else causal
260
261
262
263
264
        seqlen_k = kv.shape[1]
        assert kv.shape[0] == batch_size and kv.shape[3] == q.shape[2] and kv.shape[4] == q.shape[3]
        k, v = kv.unbind(dim=2)
        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
        scores = torch.einsum('bthd,bshd->bhts', q, k * softmax_scale)
265
266
267
268
269
270
        if key_padding_mask is not None:
            padding_mask = torch.full((batch_size, seqlen_k), -10000.0, dtype=scores.dtype,
                                      device=scores.device)
            padding_mask.masked_fill_(key_padding_mask, 0.0)
            # TD [2022-09-30]: Adding is faster than masked_fill_ (idk why, just better kernel I guess)
            scores = scores + rearrange(padding_mask, 'b s -> b 1 1 s')
Tri Dao's avatar
Tri Dao committed
271
        if causal:
272
273
274
275
276
277
278
            # "triu_tril_cuda_template" not implemented for 'BFloat16'
            # So we have to construct the mask in float
            causal_mask = torch.triu(torch.full((seqlen_q, seqlen_k), -10000.0,
                                                device=scores.device), 1)
            # TD [2022-09-30]: Adding is faster than masked_fill_ (idk why, just better kernel I guess)
            scores = scores + causal_mask.to(dtype=scores.dtype)
        attention = torch.softmax(scores, dim=-1, dtype=v.dtype)
279
        attention_drop = self.drop(attention)
280
281
282
283
284
        output = torch.einsum('bhts,bshd->bthd', attention_drop, v)
        return output


class LinearResidual(nn.Linear):
Tri Dao's avatar
Tri Dao committed
285
    """Wrap nn.Linear to return the residual as well. For compatibility with FusedDense.
286
287
288
289
290
291
    """

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return super().forward(input), input


292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
def _update_kv_cache(kv, inference_params, layer_idx):
    """kv: (batch_size, seqlen, 2, nheads, head_dim) or (batch_size, 1, 2, nheads, head_dim)
    """
    # Pre-allocate memory for key-values for inference.
    num_heads, head_dim = kv.shape[-2:]
    if layer_idx not in inference_params.key_value_memory_dict:
        kv_cache = torch.empty(
            inference_params.max_batch_size, inference_params.max_sequence_len, 2,
            num_heads, head_dim, dtype=kv.dtype, device=kv.device
        )
        inference_params.key_value_memory_dict[layer_idx] = kv_cache
    else:
        if not inference_params.fused_ft_kernel:
            kv_cache = inference_params.key_value_memory_dict[layer_idx]
        else:
            # For FT, k_cache has shape (b, h, headdim / packsize, s, packsize)
            # where packsize = 4 if fp32, 8 if fp16 or bf16.
            # v_cache has shape (b, h, s, headdim)
            k_cache, v_cache = inference_params.key_value_memory_dict[layer_idx]
            kv_cache = None
    # Adjust key and value for inference
    batch_start = inference_params.batch_size_offset
    batch_end = batch_start + kv.shape[0]
    sequence_start = inference_params.sequence_len_offset
    sequence_end = sequence_start + kv.shape[1]
    assert batch_end <= (kv_cache.shape[0] if kv_cache is not None else v_cache.shape[0])
    assert sequence_end <= (kv_cache.shape[1] if kv_cache is not None else v_cache.shape[2])
    # Copy key and values.
    if not inference_params.fused_ft_kernel:
        assert kv_cache is not None
        kv_cache[batch_start:batch_end, sequence_start:sequence_end, ...] = kv
        kv = kv_cache[batch_start:batch_end, :sequence_end, ...]
        return kv
    else:
        assert inference_params.sequence_len_offset == 0
        # FT kernel requires different layouts for the k_cache and v_cache.
        assert kv.dtype in [torch.float16, torch.bfloat16, torch.float32]
        packsize = 4 if kv.dtype == torch.float32 else 8
        if kv_cache is not None:
            kv_cache[batch_start:batch_end, sequence_start:sequence_end, ...] = kv
            k_cache = rearrange(kv_cache[:, :, 0], 'b s h (d packsize) -> b h d s packsize',
                                packsize=packsize).contiguous()
            v_cache = rearrange(kv_cache[:, :, 1], 'b s h d -> b h s d').contiguous()
            inference_params.key_value_memory_dict[layer_idx] = (k_cache, v_cache)
        else:
            k_cache[batch_start:batch_end, :, :, :sequence_end, :] = rearrange(
                kv[:, :, 0], 'b s h (d packsize) -> b h d s packsize', packsize=packsize
            )
            v_cache[batch_start:batch_end, :, :sequence_end, :] = rearrange(
                kv[:, :, 1], 'b s h d -> b h s d'
            )
        return kv


346
347
348
349
class MHA(nn.Module):
    """Multi-head self-attention and cross-attention
    """

Tri Dao's avatar
Tri Dao committed
350
351
352
353
    def __init__(self, embed_dim, num_heads, cross_attn=False,
                 qkv_proj_bias=True, out_proj_bias=True,
                 dropout=0.0, softmax_scale=None, causal=False, layer_idx=None, dwconv=False,
                 rotary_emb_dim=0, rotary_emb_scale_base=None, rotary_emb_interleaved=False,
354
355
356
357
358
359
360
361
362
363
364
365
                 fused_bias_fc=False, use_flash_attn=False, return_residual=False,
                 checkpointing=False, device=None, dtype=None) -> None:
        """
            return_residual: whether to return the input x along with the output. This is for
                performance reason: for post-norm architecture, returning the input allows us
                to fuse the backward of nn.Linear with the residual connection.
        """
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        self.embed_dim = embed_dim
        self.cross_attn = cross_attn
        self.causal = causal
Tri Dao's avatar
Tri Dao committed
366
        self.layer_idx = layer_idx
367
368
        self.dwconv = dwconv
        self.rotary_emb_dim = rotary_emb_dim
Tri Dao's avatar
Tri Dao committed
369
        self.use_flash_attn = use_flash_attn
370
371
372
373
374
375
376
377
378
379
        self.return_residual = return_residual
        self.checkpointing = checkpointing

        self.num_heads = num_heads
        assert self.embed_dim % num_heads == 0, "self.kdim must be divisible by num_heads"
        self.head_dim = self.embed_dim // num_heads

        if self.rotary_emb_dim > 0:
            assert not cross_attn, 'MHA with rotary embedding does not support cross-attention yet'
            assert RotaryEmbedding is not None, 'rotary_emb is not installed'
Tri Dao's avatar
Tri Dao committed
380
            self.rotary_emb = RotaryEmbedding(self.rotary_emb_dim, scale_base=rotary_emb_scale_base,
Tri Dao's avatar
Tri Dao committed
381
                                              interleaved=rotary_emb_interleaved, device=device)
382

Tri Dao's avatar
Tri Dao committed
383
        if fused_bias_fc and FusedDense is None:
384
            raise ImportError('fused_dense is not installed')
Tri Dao's avatar
Tri Dao committed
385
386
387
        linear_cls = nn.Linear if not fused_bias_fc else FusedDense
        linear_resid_cls = (LinearResidual if not fused_bias_fc
                            else partial(FusedDense, return_residual=True))
Tri Dao's avatar
Tri Dao committed
388
389
        inner_attn_cls = FlashSelfAttention if use_flash_attn else SelfAttention
        inner_cross_attn_cls = FlashCrossAttention if use_flash_attn else CrossAttention
390
391
        if not self.cross_attn:
            if not self.return_residual:
Tri Dao's avatar
Tri Dao committed
392
393
                self.Wqkv = linear_cls(embed_dim, 3 * embed_dim, bias=qkv_proj_bias,
                                       **factory_kwargs)
394
            else:
Tri Dao's avatar
Tri Dao committed
395
396
                self.Wqkv = linear_resid_cls(embed_dim, 3 * embed_dim, bias=qkv_proj_bias,
                                             **factory_kwargs)
397
398
399
400
            if self.dwconv:
                self.dwconv_qkv = nn.Conv1d(3 * embed_dim, 3 * embed_dim, kernel_size=3, padding=2,
                                            groups=3 * embed_dim)
        else:
Tri Dao's avatar
Tri Dao committed
401
            self.Wq = linear_cls(embed_dim, embed_dim, bias=qkv_proj_bias, **factory_kwargs)
402
            if not self.return_residual:
Tri Dao's avatar
Tri Dao committed
403
404
                self.Wkv = linear_cls(embed_dim, 2 * embed_dim, bias=qkv_proj_bias,
                                      **factory_kwargs)
405
            else:
Tri Dao's avatar
Tri Dao committed
406
407
                self.Wkv = linear_resid_cls(embed_dim, 2 * embed_dim, bias=qkv_proj_bias,
                                            **factory_kwargs)
408
409
            if self.dwconv:
                self.dwconv_q = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, padding=2,
Tri Dao's avatar
Tri Dao committed
410
                                          groups=embed_dim)
411
                self.dwconv_kv = nn.Conv1d(2 * embed_dim, 2 * embed_dim, kernel_size=3, padding=2,
Tri Dao's avatar
Tri Dao committed
412
                                          groups=2 * embed_dim)
413
        self.inner_attn = inner_attn_cls(causal=causal, softmax_scale=softmax_scale,
Tri Dao's avatar
Tri Dao committed
414
                                         attention_dropout=dropout)
Tri Dao's avatar
Tri Dao committed
415
416
        self.inner_cross_attn = inner_cross_attn_cls(causal=causal, softmax_scale=softmax_scale,
                                                     attention_dropout=dropout)
Tri Dao's avatar
Tri Dao committed
417
        self.out_proj = linear_cls(embed_dim, embed_dim, bias=out_proj_bias, **factory_kwargs)
418

419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, fused_ft_kernel=True):
        dtype = self.out_proj.weight.dtype if dtype is None else dtype
        device = self.out_proj.weight.device
        if not fused_ft_kernel:
            return torch.empty(batch_size, max_seqlen, 2, self.num_heads, self.head_dim,
                               dtype=dtype, device=device)
        else:
            assert dtype in [torch.float16, torch.bfloat16, torch.float32]
            packsize = 4 if dtype == torch.float32 else 8
            assert self.head_dim % packsize == 0
            k_cache = torch.empty(batch_size, self.num_heads, self.head_dim // packsize, max_seqlen,
                                  packsize, dtype=dtype, device=device)
            v_cache = torch.empty(batch_size, self.num_heads, max_seqlen, self.head_dim,
                                  dtype=dtype, device=device)
            return k_cache, v_cache

Tri Dao's avatar
Tri Dao committed
435
    def _update_kv_cache(self, kv, inference_params):
436
        """kv: (batch_size, seqlen, 2, nheads, head_dim) or (batch_size, 1, 2, nheads, head_dim)
Tri Dao's avatar
Tri Dao committed
437
438
439
        """
        assert not self.dwconv, 'Generation does not support dwconv yet'
        assert self.layer_idx is not None, 'Generation requires layer_idx in the constructor'
440
        return _update_kv_cache(kv, inference_params, self.layer_idx)
Tri Dao's avatar
Tri Dao committed
441

442
    def forward(self, x, x_kv=None, key_padding_mask=None, cu_seqlens=None, max_seqlen=None,
443
                mixer_subset=None, inference_params=None, **kwargs):
444
445
        """
        Arguments:
Tri Dao's avatar
Tri Dao committed
446
447
448
            x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim) if
                cu_seqlens is None and max_seqlen is None, else (total, hidden_dim) where total
                is the is the sum of the sequence lengths in the batch.
449
            x_kv: (batch, seqlen, hidden_dim), only applicable for cross-attention. If None, use x.
Tri Dao's avatar
Tri Dao committed
450
451
452
453
454
455
            cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
                of the sequences in the batch, used to index into x. Only applicable when using
                FlashAttention.
            max_seqlen: int. Maximum sequence length in the batch.
            key_padding_mask: boolean mask, True means to keep, False means to mask out.
                (batch, seqlen). Only applicable when not using FlashAttention.
456
457
458
            mixer_subset: for cross-attention only. If not None, will take a subset of x
                before applying the query projection. Useful for e.g., ViT where we only care
                about the CLS token in the last layer.
Tri Dao's avatar
Tri Dao committed
459
460
            inference_params: for generation. Adapted from Megatron-LM (and Apex)
            https://github.com/NVIDIA/apex/blob/3ff1a10f72ec07067c4e44759442329804ac5162/apex/transformer/testing/standalone_transformer_lm.py#L470
461
        """
Tri Dao's avatar
Tri Dao committed
462
463
464
465
466
467
468
469
470
471
        if cu_seqlens is not None:
            assert max_seqlen is not None
            assert key_padding_mask is None
            assert self.use_flash_attn
            assert not self.dwconv
            assert self.rotary_emb_dim == 0
        if key_padding_mask is not None:
            assert cu_seqlens is None
            assert max_seqlen is None
            assert not self.use_flash_attn
Tri Dao's avatar
Tri Dao committed
472
473
474
475
        if inference_params is not None:
            assert key_padding_mask is None
            assert cu_seqlens is None and max_seqlen is None
            assert not self.dwconv
Tri Dao's avatar
Tri Dao committed
476

477
478
        kwargs = ({'cu_seqlens': cu_seqlens, 'max_seqlen': max_seqlen, **kwargs}
                  if self.use_flash_attn else {'key_padding_mask': key_padding_mask, **kwargs})
479
        if not self.cross_attn:
480
            assert x_kv is None and mixer_subset is None
481
482
483
484
485
486
487
            if not self.return_residual:
                qkv = self.Wqkv(x)
            else:
                qkv, x = self.Wqkv(x)
            if self.dwconv:
                qkv = rearrange(self.dwconv_qkv(rearrange(qkv, 'b s d -> b d s'))[..., :-2],
                                'b d s -> b s d').contiguous()
Tri Dao's avatar
Tri Dao committed
488
            qkv = rearrange(qkv, '... (three h d) -> ... three h d', three=3, d=self.head_dim)
Tri Dao's avatar
Tri Dao committed
489
490
491
492
493
494
495
            if inference_params is None:
                if self.rotary_emb_dim > 0:
                    qkv = self.rotary_emb(qkv)
                if not self.checkpointing:
                    context = self.inner_attn(qkv, **kwargs)
                else:
                    context = torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, **kwargs)
496
            else:
497
498
499
500
501
502
503
504
505
506
                if (not inference_params.fused_ft_kernel) or inference_params.sequence_len_offset == 0:
                    if self.rotary_emb_dim > 0:
                        qkv = self.rotary_emb(qkv, seqlen_offset=inference_params.sequence_len_offset)
                    q = qkv[:, :, 0]
                    kv = self._update_kv_cache(qkv[:, :, 1:], inference_params)
                    # If we're processing the prompt, causal=None (use self.causal).
                    # If we're decoding, then causal=False.
                    causal = None if inference_params.sequence_len_offset == 0 else False
                    context = self.inner_cross_attn(q, kv, causal=causal)
                else:
507
                    assert inference_params.fused_ft_kernel
508
                    assert ft_attention is not None
509
510
511
512
513
                    batch_start = inference_params.batch_size_offset
                    batch_end = batch_start + qkv.shape[0]
                    k_cache, v_cache = inference_params.key_value_memory_dict[self.layer_idx]
                    lengths_per_sample = (inference_params.lengths_per_sample[batch_start:batch_end]
                                          if inference_params.lengths_per_sample is not None else None)
514
515
                    context = ft_attention.single_query_attention(
                        *rearrange(qkv, 'b 1 three h d -> b three h d').unbind(dim=1),
516
517
                        k_cache[batch_start:batch_end], v_cache[batch_start:batch_end],
                        lengths_per_sample, inference_params.sequence_len_offset,
518
                        self.rotary_emb_dim,
Tri Dao's avatar
Tri Dao committed
519
520
                        # neox_rotary_style
                        (not self.rotary_emb.interleaved) if self.rotary_emb_dim > 0 else True
521
522
                    )
                    context = rearrange(context, 'b h d -> b 1 h d')
523
        else:
524
            if not self.return_residual:
525
                q = self.Wq(x if mixer_subset is None else x[:, mixer_subset])
526
527
528
529
530
531
                kv = self.Wkv(x_kv if x_kv is not None else x)
            else:
                if x_kv is not None:
                    kv, x_kv = self.Wkv(x_kv)
                else:
                    kv, x = self.Wkv(x)
532
                q = self.Wq(x if mixer_subset is None else x[:, mixer_subset])
Tri Dao's avatar
Tri Dao committed
533
534
            q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim)
            kv = rearrange(kv, '... (two h d) -> ... two h d', two=2, d=self.head_dim)
535
536
537
538
539
            if self.dwconv:
                q = rearrange(self.dwconv_q(rearrange(q, 'b s d -> b d s'))[..., :-2],
                              'b d s -> b s d').contiguous()
                kv = rearrange(self.dwconv_kv(rearrange(kv, 'b s d -> b d s'))[..., :-2],
                               'b d s -> b s d').contiguous()
Tri Dao's avatar
Tri Dao committed
540
541
            if inference_params is None:
                if not self.checkpointing:
542
                    context = self.inner_cross_attn(q, kv, **kwargs)
Tri Dao's avatar
Tri Dao committed
543
                else:
544
                    context = torch.utils.checkpoint.checkpoint(self.inner_cross_attn, q, kv, **kwargs)
545
            else:
Tri Dao's avatar
Tri Dao committed
546
547
                kv = self._update_kv_cache(kv)
                context = self.inner_cross_attn(q, kv, causal=False)
Tri Dao's avatar
Tri Dao committed
548
        out = self.out_proj(rearrange(context, '... h d -> ... (h d)'))
549
        return out if not self.return_residual else (out, x)
Tri Dao's avatar
Tri Dao committed
550
551
552
553
554
555


class ParallelMHA(nn.Module):
    """Multi-head self-attention and cross-attention
    """

Tri Dao's avatar
Tri Dao committed
556
557
558
559
    def __init__(self, embed_dim, num_heads, process_group, qkv_proj_bias=True, out_proj_bias=True,
                 dropout=0.0, softmax_scale=None, causal=False, layer_idx=None,
                 rotary_emb_dim=0, rotary_emb_scale_base=None, rotary_emb_interleaved=False,
                 use_flash_attn=False, checkpointing=False,
560
                 sequence_parallel=True, device=None, dtype=None) -> None:
Tri Dao's avatar
Tri Dao committed
561
562
563
564
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        self.embed_dim = embed_dim
        self.causal = causal
565
        self.layer_idx = layer_idx
Tri Dao's avatar
Tri Dao committed
566
567
568
569
570
571
572
573
574
575
576
        self.rotary_emb_dim = rotary_emb_dim
        self.use_flash_attn = use_flash_attn
        self.checkpointing = checkpointing

        self.num_heads = num_heads
        assert self.embed_dim % num_heads == 0, "self.kdim must be divisible by num_heads"
        self.head_dim = self.embed_dim // num_heads

        if self.rotary_emb_dim > 0:
            assert RotaryEmbedding is not None, 'rotary_emb is not installed'
            self.rotary_emb = RotaryEmbedding(self.rotary_emb_dim, scale_base=rotary_emb_scale_base,
Tri Dao's avatar
Tri Dao committed
577
                                              interleaved=rotary_emb_interleaved, device=device)
Tri Dao's avatar
Tri Dao committed
578
579
580

        if ColumnParallelLinear is None or RowParallelLinear is None:
            raise ImportError('fused_dense is not installed')
Tri Dao's avatar
Tri Dao committed
581
582
        self.Wqkv = ColumnParallelLinear(embed_dim, 3 * embed_dim, process_group,
                                         bias=qkv_proj_bias,
583
                                         sequence_parallel=sequence_parallel, **factory_kwargs)
Tri Dao's avatar
Tri Dao committed
584
        inner_attn_cls = FlashSelfAttention if use_flash_attn else SelfAttention
585
        inner_cross_attn_cls = FlashCrossAttention if use_flash_attn else CrossAttention
Tri Dao's avatar
Tri Dao committed
586
587
        self.inner_attn = inner_attn_cls(causal=causal, softmax_scale=softmax_scale,
                                         attention_dropout=dropout)
588
589
        self.inner_cross_attn = inner_cross_attn_cls(causal=causal, softmax_scale=softmax_scale,
                                                     attention_dropout=dropout)
590
        self.out_proj = RowParallelLinear(embed_dim, embed_dim, process_group,
Tri Dao's avatar
Tri Dao committed
591
                                          bias=out_proj_bias,
592
                                          sequence_parallel=sequence_parallel, **factory_kwargs)
Tri Dao's avatar
Tri Dao committed
593

594
    def forward(self, x, seqlen=None, inference_params=None, **kwargs):
Tri Dao's avatar
Tri Dao committed
595
596
597
598
599
600
601
602
603
604
605
606
607
        """
        Arguments:
            x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim) if seqlen=None.
                If seqlen is not None, x is (batch * seqlen, hidden_dim). This is so that when we
                split x during sequence parallel, we split the batch * seqlen dimension
                (in case batch is small).
        """
        qkv = self.Wqkv(x)
        if seqlen is None:
            qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, d=self.head_dim)
        else:
            qkv = rearrange(qkv, '(b s) (three h d) -> b s three h d', s=seqlen, three=3,
                            d=self.head_dim)
608
609
610
611
612
613
614
        if inference_params is None:
            if self.rotary_emb_dim > 0:
                qkv = self.rotary_emb(qkv)
            if not self.checkpointing:
                context = self.inner_attn(qkv, **kwargs)
            else:
                context = torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, **kwargs)
Tri Dao's avatar
Tri Dao committed
615
        else:
616
617
618
619
620
621
622
623
624
625
626
627
628
            if (not inference_params.fused_ft_kernel) or inference_params.sequence_len_offset == 0:
                if self.rotary_emb_dim > 0:
                    qkv = self.rotary_emb(qkv, seqlen_offset=inference_params.sequence_len_offset)
                q = qkv[:, :, 0]
                assert self.layer_idx is not None, 'Generation requires layer_idx in the constructor'
                kv = _update_kv_cache(qkv[:, :, 1:], inference_params, self.layer_idx)
                # If we're processing the prompt, causal=None (use self.causal).
                # If we're decoding, then causal=False.
                causal = None if inference_params.sequence_len_offset == 0 else False
                context = self.inner_cross_attn(q, kv, causal=causal)
            else:
                assert inference_params.fused_ft_kernel
                assert ft_attention is not None
629
630
631
632
633
                batch_start = inference_params.batch_size_offset
                batch_end = batch_start + qkv.shape[0]
                k_cache, v_cache = inference_params.key_value_memory_dict[self.layer_idx]
                lengths_per_sample = (inference_params.lengths_per_sample[batch_start:batch_end]
                                      if inference_params.lengths_per_sample is not None else None)
634
635
                context = ft_attention.single_query_attention(
                    *rearrange(qkv, 'b 1 three h d -> b three h d').unbind(dim=1),
636
637
638
                    k_cache[batch_start:batch_end], v_cache[batch_start:batch_end],
                    lengths_per_sample, inference_params.sequence_len_offset,
                    self.rotary_emb_dim, inference_params.sequence_len_offset,
Tri Dao's avatar
Tri Dao committed
639
640
                    # neox_rotary_style
                    (not self.rotary_emb.interleaved) if self.rotary_emb_dim > 0 else True
641
642
                )
                context = rearrange(context, 'b h d -> b 1 h d')
Tri Dao's avatar
Tri Dao committed
643
644
645
646
647
648
        if seqlen is None:
            context = rearrange(context, 'b s h d -> b s (h d)')
        else:
            context = rearrange(context, 'b s h d -> (b s) (h d)')
        out = self.out_proj(context)
        return out