From ccc2d33b71d6bd8973665aa912b1651bf3f1156d Mon Sep 17 00:00:00 2001
From: yaoguany <89233842+yaoguany@users.noreply.github.com>
Date: Mon, 18 Sep 2023 16:37:08 +0800
Subject: [PATCH 1/4] update llama flash attention

---
 .../flash_attention/llama_flash_attention.py  | 158 +++++++++++-------
 1 file changed, 96 insertions(+), 62 deletions(-)

diff --git a/src/lmflow/utils/flash_attention/llama_flash_attention.py b/src/lmflow/utils/flash_attention/llama_flash_attention.py
index 4159629c6..790bd403e 100644
--- a/src/lmflow/utils/flash_attention/llama_flash_attention.py
+++ b/src/lmflow/utils/flash_attention/llama_flash_attention.py
@@ -2,6 +2,7 @@
 
 import torch
 from torch import nn
+import math
 
 import transformers
 from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
@@ -10,110 +11,143 @@
 
 #try to import flash_attn 2.x.x, if not, import flash_attn 1.x.x
 try:
-    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func
+    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
 except:
-    from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_func
+    from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_func as flash_attn_varlen_qkvpacked_func
 
 from flash_attn.bert_padding import unpad_input, pad_input
 
 
 def forward(
-    self,
-    hidden_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.Tensor] = None,
-    past_key_value: Optional[Tuple[torch.Tensor]] = None,
-    output_attentions: bool = False,
-    use_cache: bool = False,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor],
+           Optional[Tuple[torch.Tensor]]]:
     """Input shape: Batch x Time x Channel
 
     attention_mask: [bsz, q_len]
     """
     bsz, q_len, _ = hidden_states.size()
 
-    query_states = (
-        self.q_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    key_states = (
-        self.k_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    value_states = (
-        self.v_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
+    query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
     # [bsz, q_len, nh, hd]
     # [bsz, nh, q_len, hd]
 
     kv_seq_len = key_states.shape[-2]
-    assert past_key_value is None, "past_key_value is not supported"
+    # assert past_key_value is None, "past_key_value is not supported"
+    if past_key_value is not None:
+        kv_seq_len += past_key_value[0].shape[-2]
 
     cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-    query_states, key_states = apply_rotary_pos_emb(
-        query_states, key_states, cos, sin, position_ids
-    )
+    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
     # [bsz, nh, t, hd]
-    assert not output_attentions, "output_attentions is not supported"
-    assert not use_cache, "use_cache is not supported"
+    # assert not output_attentions, "output_attentions is not supported"
+    # assert not use_cache, "use_cache is not supported"
+    if past_key_value is not None:
+        # reuse k, v, self_attention
+        key_states = torch.cat([past_key_value[0], key_states], dim=2)
+        value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+    past_key_value = (key_states, value_states) if use_cache else None
 
     # Flash attention codes from
     # https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attention.py
 
     # transform the data into the format required by flash attention
-    qkv = torch.stack(
-        [query_states, key_states, value_states], dim=2
-    )  # [bsz, nh, 3, q_len, hd]
+    # import pdb; pdb.set_trace()
+
+    if query_states.shape[-2] == 1 or query_states.shape[-2] != key_states.shape[-2]:
+        # decode token-by-token, do not use flash attention
+        # in incremental state, do not use flash attention
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+            attn_weights = torch.max(attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min))
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    qkv = torch.stack([query_states, key_states, value_states], dim=2)  # [bsz, nh, 3, q_len, hd]
     qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
     # We have disabled _prepare_decoder_attention_mask in LlamaModel
     # the attention_mask should be the same as the key_padding_mask
     key_padding_mask = attention_mask
 
     if key_padding_mask is None:
-        qkv = rearrange(qkv, "b s ... -> (b s) ...")
+        qkv = rearrange(qkv, 'b s ... -> (b s) ...')
         max_s = q_len
-        cu_q_lens = torch.arange(
-            0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32, device=qkv.device
-        )
-        output = flash_attn_unpadded_qkvpacked_func(
-            qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
+        cu_q_lens = torch.arange(0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32,
+                                 device=qkv.device)
+        output = flash_attn_varlen_qkvpacked_func(
+            qkv, cu_q_lens, max_s, 0.0,
+            softmax_scale=None, causal=True
         )
-        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
+        output = rearrange(output, '(b s) ... -> b s ...', b=bsz)
     else:
         nheads = qkv.shape[-2]
-        x = rearrange(qkv, "b s three h d -> b s (three h d)")
+        x = rearrange(qkv, 'b s three h d -> b s (three h d)')
         x_unpad, indices, cu_q_lens, max_s = unpad_input(x, key_padding_mask)
-        x_unpad = rearrange(
-            x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads
+        x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads)
+        output_unpad = flash_attn_varlen_qkvpacked_func(
+            x_unpad, cu_q_lens, max_s, 0.0,
+            softmax_scale=None, causal=True
         )
-        output_unpad = flash_attn_unpadded_qkvpacked_func(
-            x_unpad, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
-        )
-        output = rearrange(
-            pad_input(
-                rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, bsz, q_len
-            ),
-            "b s (h d) -> b s h d",
-            h=nheads,
-        )
-    return self.o_proj(rearrange(output, "b s h d -> b s (h d)")), None, None
+        output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'),
+                                     indices, bsz, q_len),
+                           'b s (h d) -> b s h d', h=nheads)
+    attn_output = self.o_proj(rearrange(output, 'b s h d -> b s (h d)'))
+
+    return attn_output, None, past_key_value
 
 
 # Disable the transformation of the attention mask in LlamaModel as the flash attention
 # requires the attention mask to be the same as the key_padding_mask
-def _prepare_decoder_attention_mask(
-    self, attention_mask, input_shape, inputs_embeds, past_key_values_length
-):
+def _prepare_decoder_attention_mask(self, attention_mask, input_shape,
+                                    inputs_embeds, past_key_values_length):
     # [bsz, seq_len]
-    return attention_mask
+    if input_shape[-1] > 1 and past_key_values_length == 0:  # encode
+        return attention_mask
+    return transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask(self, attention_mask,
+                                                                                              input_shape,
+                                                                                              inputs_embeds,
+                                                                                              past_key_values_length)
 
 
 def replace_llama_attn_with_flash_attn():
-    transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (
-        _prepare_decoder_attention_mask
-    )
+    transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = _prepare_decoder_attention_mask
     transformers.models.llama.modeling_llama.LlamaAttention.forward = forward
\ No newline at end of file

From bb9b1c0414eeec924688a039bf8f9c761dbbda40 Mon Sep 17 00:00:00 2001
From: yaoguany <89233842+yaoguany@users.noreply.github.com>
Date: Mon, 18 Sep 2023 20:24:36 +0800
Subject: [PATCH 2/4] update decoder attention mask code

---
 .../flash_attention/llama_flash_attention.py  | 44 +++++++++++++------
 1 file changed, 30 insertions(+), 14 deletions(-)

diff --git a/src/lmflow/utils/flash_attention/llama_flash_attention.py b/src/lmflow/utils/flash_attention/llama_flash_attention.py
index 790bd403e..76f99fa6b 100644
--- a/src/lmflow/utils/flash_attention/llama_flash_attention.py
+++ b/src/lmflow/utils/flash_attention/llama_flash_attention.py
@@ -5,7 +5,7 @@
 import math
 
 import transformers
-from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb,_make_causal_mask,_expand_mask
 
 from einops import rearrange
 
@@ -19,15 +19,14 @@
 
 
 def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        past_key_value: Optional[Tuple[torch.Tensor]] = None,
-        output_attentions: bool = False,
-        use_cache: bool = False,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor],
-           Optional[Tuple[torch.Tensor]]]:
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.Tensor] = None,
+    past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor],Optional[Tuple[torch.Tensor]]]:
     """Input shape: Batch x Time x Channel
 
     attention_mask: [bsz, q_len]
@@ -142,10 +141,27 @@ def _prepare_decoder_attention_mask(self, attention_mask, input_shape,
     # [bsz, seq_len]
     if input_shape[-1] > 1 and past_key_values_length == 0:  # encode
         return attention_mask
-    return transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask(self, attention_mask,
-                                                                                              input_shape,
-                                                                                              inputs_embeds,
-                                                                                              past_key_values_length)
+    # create causal mask
+    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+    combined_attention_mask = None
+    if input_shape[-1] > 1:
+        combined_attention_mask = _make_causal_mask(
+            input_shape,
+            inputs_embeds.dtype,
+            device=inputs_embeds.device,
+            past_key_values_length=past_key_values_length,
+        )
+
+    if attention_mask is not None:
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
+            inputs_embeds.device
+        )
+        combined_attention_mask = (
+            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+        )
+
+    return combined_attention_mask
 
 
 def replace_llama_attn_with_flash_attn():

From 0a639642dadac15687a1fa4e13ae13ff68347e30 Mon Sep 17 00:00:00 2001
From: yaoguany <89233842+yaoguany@users.noreply.github.com>
Date: Mon, 18 Sep 2023 20:32:45 +0800
Subject: [PATCH 3/4] update codes

---
 .../flash_attention/llama_flash_attention.py  | 59 ++++++++++++-------
 1 file changed, 38 insertions(+), 21 deletions(-)

diff --git a/src/lmflow/utils/flash_attention/llama_flash_attention.py b/src/lmflow/utils/flash_attention/llama_flash_attention.py
index 76f99fa6b..23ffa219d 100644
--- a/src/lmflow/utils/flash_attention/llama_flash_attention.py
+++ b/src/lmflow/utils/flash_attention/llama_flash_attention.py
@@ -26,16 +26,28 @@ def forward(
     past_key_value: Optional[Tuple[torch.Tensor]] = None,
     output_attentions: bool = False,
     use_cache: bool = False,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor],Optional[Tuple[torch.Tensor]]]:
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
     """Input shape: Batch x Time x Channel
 
     attention_mask: [bsz, q_len]
     """
     bsz, q_len, _ = hidden_states.size()
 
-    query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-    key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-    value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    query_states = (
+        self.q_proj(hidden_states)
+        .view(bsz, q_len, self.num_heads, self.head_dim)
+        .transpose(1, 2)
+    )
+    key_states = (
+        self.k_proj(hidden_states)
+        .view(bsz, q_len, self.num_heads, self.head_dim)
+        .transpose(1, 2)
+    )
+    value_states = (
+        self.v_proj(hidden_states)
+        .view(bsz, q_len, self.num_heads, self.head_dim)
+        .transpose(1, 2)
+    )
     # [bsz, q_len, nh, hd]
     # [bsz, nh, q_len, hd]
 
@@ -45,7 +57,9 @@ def forward(
         kv_seq_len += past_key_value[0].shape[-2]
 
     cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+    query_states, key_states = apply_rotary_pos_emb(
+        query_states, key_states, cos, sin, position_ids
+    )
     # [bsz, nh, t, hd]
     # assert not output_attentions, "output_attentions is not supported"
     # assert not use_cache, "use_cache is not supported"
@@ -108,30 +122,33 @@ def forward(
     key_padding_mask = attention_mask
 
     if key_padding_mask is None:
-        qkv = rearrange(qkv, 'b s ... -> (b s) ...')
+        qkv = rearrange(qkv, "b s ... -> (b s) ...")
         max_s = q_len
-        cu_q_lens = torch.arange(0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32,
-                                 device=qkv.device)
+        cu_q_lens = torch.arange(
+            0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32, device=qkv.device
+        )
         output = flash_attn_varlen_qkvpacked_func(
-            qkv, cu_q_lens, max_s, 0.0,
-            softmax_scale=None, causal=True
+            qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
         )
-        output = rearrange(output, '(b s) ... -> b s ...', b=bsz)
+        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
     else:
         nheads = qkv.shape[-2]
-        x = rearrange(qkv, 'b s three h d -> b s (three h d)')
+        x = rearrange(qkv, "b s three h d -> b s (three h d)")
         x_unpad, indices, cu_q_lens, max_s = unpad_input(x, key_padding_mask)
-        x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads)
+        x_unpad = rearrange(
+            x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads
+        )
         output_unpad = flash_attn_varlen_qkvpacked_func(
-            x_unpad, cu_q_lens, max_s, 0.0,
-            softmax_scale=None, causal=True
+            x_unpad, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
         )
-        output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'),
-                                     indices, bsz, q_len),
-                           'b s (h d) -> b s h d', h=nheads)
-    attn_output = self.o_proj(rearrange(output, 'b s h d -> b s (h d)'))
-
-    return attn_output, None, past_key_value
+        output = rearrange(
+            pad_input(
+                rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, bsz, q_len
+            ),
+            "b s (h d) -> b s h d",
+            h=nheads,
+        )
+    return self.o_proj(rearrange(output, 'b s h d -> b s (h d)')), None, past_key_value
 
 
 # Disable the transformation of the attention mask in LlamaModel as the flash attention

From bf87e6d5519a65310a5cdcffb0abfcba605a2202 Mon Sep 17 00:00:00 2001
From: yaoguany <89233842+yaoguany@users.noreply.github.com>
Date: Mon, 18 Sep 2023 20:35:36 +0800
Subject: [PATCH 4/4] update code to make minimum change

---
 .../utils/flash_attention/llama_flash_attention.py   | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/lmflow/utils/flash_attention/llama_flash_attention.py b/src/lmflow/utils/flash_attention/llama_flash_attention.py
index 23ffa219d..5855c9146 100644
--- a/src/lmflow/utils/flash_attention/llama_flash_attention.py
+++ b/src/lmflow/utils/flash_attention/llama_flash_attention.py
@@ -115,7 +115,10 @@ def forward(
 
         return attn_output, attn_weights, past_key_value
 
-    qkv = torch.stack([query_states, key_states, value_states], dim=2)  # [bsz, nh, 3, q_len, hd]
+    # transform the data into the format required by flash attention
+    qkv = torch.stack(
+        [query_states, key_states, value_states], dim=2
+    )  # [bsz, nh, 3, q_len, hd]
     qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
     # We have disabled _prepare_decoder_attention_mask in LlamaModel
     # the attention_mask should be the same as the key_padding_mask
@@ -148,13 +151,14 @@ def forward(
             "b s (h d) -> b s h d",
             h=nheads,
         )
-    return self.o_proj(rearrange(output, 'b s h d -> b s (h d)')), None, past_key_value
+    return self.o_proj(rearrange(output, "b s h d -> b s (h d)")), None, past_key_value
 
 
 # Disable the transformation of the attention mask in LlamaModel as the flash attention
 # requires the attention mask to be the same as the key_padding_mask
-def _prepare_decoder_attention_mask(self, attention_mask, input_shape,
-                                    inputs_embeds, past_key_values_length):
+def _prepare_decoder_attention_mask(
+    self, attention_mask, input_shape, inputs_embeds, past_key_values_length
+):
     # [bsz, seq_len]
     if input_shape[-1] > 1 and past_key_values_length == 0:  # encode
         return attention_mask