code: def get_llama_flash_attention_forward(shard_config: ShardConfig):
      from transformers.models.llama.modeling_llama import (
          LlamaAttention,
          apply_rotary_pos_emb,
      )
  
      from colossalai.kernel.cuda_native import (
          AttnMaskType,
          ColoAttention,
      )
  
      llama_version = 2
      try:
          from transformers.models.llama.modeling_llama import repeat_kv
      except:
          warnings.warn(
              "using llamav1, llamav1 hasn't repeat_kv function"
          )
          llama_version = 1
  
      def forward(
          self: LlamaAttention,
          hidden_states: torch.Tensor,
          attention_mask: Optional[torch.Tensor] = None,
          position_ids: Optional[torch.LongTensor] = None,
          past_key_value: Optional[Tuple[torch.Tensor]] = None,
          output_attentions: bool = False,
          use_cache: bool = False,
          **kwargs,
      ) -> Tuple[
          torch.Tensor,
          Optional[torch.Tensor],
          Optional[Tuple[torch.Tensor]],
      ]:
          bsz, q_len, _ = hidden_states.size()
          assert (
              q_len % 4 == 0
          ), "Flash Attention Error: The sequence length should be a multiple of 4."
  
          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)
          )
  
          kv_seq_len = key_states.shape[-2]
          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
          )
  
          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
          )
  
          # repeat k/v heads if n_kv_heads < n_heads
          if llama_version == 2:
              key_states = repeat_kv(
                  key_states, self.num_key_value_groups
              )
              value_states = repeat_kv(
                  value_states, self.num_key_value_groups
              )
  
          me_input_shape = (bsz, q_len, self.num_heads, self.head_dim)
          query_states = (
              query_states.transpose(1, 2)
              .contiguous()
              .view(*me_input_shape)
          )
          key_states = (
              key_states.transpose(1, 2)
              .contiguous()
              .view(*me_input_shape)
          )
          value_states = (
              value_states.transpose(1, 2)
              .contiguous()
              .view(*me_input_shape)
          )
  
          flash_attention_mask = None
          attn_mask_type = AttnMaskType.causal
          if (
              not getattr(shard_config, "causal_lm", False)
              and attention_mask != 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()}"
                  )
              flash_attention_mask = ~(
                  attention_mask[:, :, -1].squeeze(1).to(torch.bool)
              ).contiguous()
              attn_mask_type = AttnMaskType.paddedcausal
  
          attention = ColoAttention(
              embed_dim=self.hidden_size, num_heads=self.num_heads
          )
          attn_output = attention(
              query_states,
              key_states,
              value_states,
              attn_mask=flash_attention_mask,
              attn_mask_type=attn_mask_type,
          )
  
          attn_output = self.o_proj(attn_output)
  
          return attn_output, None, past_key_value
  
      return forward
  
file_path: ColossalAI/colossalai/shardformer/modeling/llama.py
imports: import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import torch.distributed
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
  from transformers.models.llama.modeling_llama import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer.shard import ShardConfig
  from ..layer import cross_entropy_1d
label_desc: Use image node_id 1 for calling a global function with example usage: get_llama_flash_attention_forward(shard_config) and returns: forward, attn_output,  None,  past_key_value