class_docstr: Region where selective activation recomputation is applied.
  See Figure 3. in Reducing Activation Recomputation in Large Transformer Models
  https://arxiv.org/pdf/2205.05198.pdf for more details.
class_docstr_tok: ["Region","where","selective","activation","recomputation","is","applied",".","See","Figure","3.","in","Reducing","Activation","Recomputation","in","Large","Transformer","Models","https",":","\/\/arxiv.org\/pdf\/2205.05198.pdf","for","more","details","."]
code: def flash_attention_cuda(
      self,
      query_layer,
      key_layer,
      value_layer,
      attention_mask,
      is_causal,
  ):
      batch_size, seqlen, nheads, _ = query_layer.shape
  
      # True: attend / False: not attend
      if attention_mask is None:
          attention_mask_q = torch.ones(
              batch_size,
              query_layer.shape[1],
              device=query_layer.device,
          ).bool()
          attention_mask_kv = torch.ones(
              batch_size, key_layer.shape[1], device=query_layer.device
          ).bool()
      elif len(attention_mask.shape) == 4:
          # [b, 1, sq, sk] -> [b, sq] / [b, sk]
          attention_mask_q = torch.any(
              torch.eq(attention_mask, False), dim=3
          ).squeeze(1)
          attention_mask_kv = torch.any(
              torch.eq(attention_mask, False), dim=2
          ).squeeze(1)
      else:
          assert len(attention_mask.shape) == 2
          attention_mask_q = ~attention_mask
          attention_mask_kv = ~attention_mask
  
      seqlens_q_in_batch = len(
          attention_mask_q.sum(dim=-1, dtype=torch.int32).unique()
      )
      seqlens_kv_in_batch = len(
          attention_mask_kv.sum(dim=-1, dtype=torch.int32).unique()
      )
  
      if (
          seqlens_q_in_batch == 1
          and seqlens_kv_in_batch == 1
          and flash_attn_func is not None
      ):
          # [b, sq, np, hn]
          context_layer = flash_attn_func(
              query_layer,
              key_layer,
              value_layer,
              dropout_p=self.attention_dropout_p
              if self.training
              else 0.0,
              causal=is_causal,
          )
      else:
          q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(
              query_layer, attention_mask_q
          )
          k, _, cu_seqlens_k, max_seqlen_k = unpad_input(
              key_layer, attention_mask_kv
          )
          v, _, _, _ = unpad_input(value_layer, attention_mask_kv)
  
          context_layer = flash_attn_unpadded_func(
              q,
              k,
              v,
              cu_seqlens_q,
              cu_seqlens_k,
              max_seqlen_q,
              max_seqlen_k,
              dropout_p=self.attention_dropout_p
              if self.training
              else 0.0,
              causal=is_causal,
          )
  
          # [b, sq, np, hn]
          context_layer = pad_input(
              context_layer, indices_q, batch_size, seqlen
          )
  
      # [b, sq, np, hn] -> [b, np, sq, hn]
      context_layer = context_layer.permute(0, 2, 1, 3)
      return context_layer
  
imports: import math
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, InfusedAdapterConfig, LoraKQVAdapterConfig, LoraKQVAdapterWeightTyingConfig, LoraKVAdapterConfig, LoraQAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.fused_softmax import MatchedScaleMaskSoftmax
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.position_embedding import XPOSPositionEmbedding
  from nemo.collections.nlp.modules.common.megatron.position_embedding.rotary_position_embedding import apply_rotary_pos_emb
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, _cast_if_autocast_enabled, attention_mask_func
  from nemo.core import adapter_mixins
label_desc: Use image node_id 5 for calling the CoreAttention obj's underlying member method code with example usage: obj.flash_attention_cuda(query_layer, key_layer, value_layer, attention_mask, is_causal) and returns: context_layer