code: def get_flash_core_attention_forward():
      from colossalai.kernel.cuda_native import (
          AttnMaskType,
          ColoAttention,
      )
  
      from .chatglm2_6b.modeling_chatglm import CoreAttention
  
      def forward(
          self: CoreAttention,
          query_layer,
          key_layer,
          value_layer,
          attention_mask,
      ):
          pytorch_major_version = int(torch.__version__.split(".")[0])
          if pytorch_major_version >= 2:
              query_layer, key_layer, value_layer = [
                  k.permute(1, 2, 0, 3)
                  for k in [query_layer, key_layer, value_layer]
              ]
              if (
                  attention_mask is None
                  and query_layer.shape[2] == key_layer.shape[2]
              ):
                  context_layer = (
                      torch.nn.functional.scaled_dot_product_attention(
                          query_layer,
                          key_layer,
                          value_layer,
                          is_causal=True,
                      )
                  )
              else:
                  if attention_mask is not None:
                      attention_mask = ~attention_mask
                  context_layer = (
                      torch.nn.functional.scaled_dot_product_attention(
                          query_layer,
                          key_layer,
                          value_layer,
                          attention_mask,
                      )
                  )
              context_layer = context_layer.permute(2, 0, 1, 3)
              new_context_layer_shape = context_layer.size()[:-2] + (
                  self.hidden_size_per_partition,
              )
              context_layer = context_layer.reshape(
                  *new_context_layer_shape
              )
          else:
              # Raw attention scores
              query_layer = query_layer.permute(1, 0, 2, 3).contiguous()
              key_layer = key_layer.permute(1, 0, 2, 3).contiguous()
              value_layer = value_layer.permute(1, 0, 2, 3).contiguous()
  
              scale = 1.0 / self.norm_factor
              if self.coeff is not None:
                  scale = scale * self.coeff
  
              flash_attention_mask = None
              attn_mask_type = None
              if attention_mask is None:
                  attn_mask_type = AttnMaskType.causal
              else:
                  flash_attention_mask = ~(
                      attention_mask[:, :, -1].squeeze(1).to(torch.bool)
                  ).contiguous()
                  if not torch.all(flash_attention_mask):
                      attn_mask_type = AttnMaskType.paddedcausal
  
              attention = ColoAttention(
                  embed_dim=self.hidden_size_per_partition,
                  num_heads=self.num_attention_heads_per_partition,
                  dropout=self.attention_dropout.p,
                  scale=scale,
              )
              context_layer = attention(
                  query_layer,
                  key_layer,
                  value_layer,
                  attn_mask=flash_attention_mask,
                  attn_mask_type=attn_mask_type,
              )
  
              context_layer = context_layer.permute(
                  1, 0, -1
              ).contiguous()
  
          return context_layer
  
      return forward
  
file_path: ColossalAI/colossalai/shardformer/modeling/chatglm2.py
imports: from typing import List, Optional, Tuple
  import torch
  import torch.utils.checkpoint
  from torch.nn import CrossEntropyLoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer import ShardConfig
  from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel
label_desc: Use image node_id 1 for calling a global function with example usage: get_flash_core_attention_forward() and returns: forward, context_layer