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(
      self,
      query_layer,
      key_layer,
      value_layer,
      attention_mask,
      attention_bias,
      inference_mode,
  ):
      query_layer = rearrange(query_layer, "sq b np hn -> b sq np hn")
      key_layer = rearrange(key_layer, "sk b np hn -> b sk np hn")
      value_layer = rearrange(value_layer, "sv b np hn -> b sv np hn")
  
      # Use to ensure dtype cast to fp16 or bf16
      query_layer = _cast_if_autocast_enabled(query_layer)
      key_layer = _cast_if_autocast_enabled(key_layer)
      value_layer = _cast_if_autocast_enabled(value_layer)
      attention_bias = _cast_if_autocast_enabled(attention_bias)
  
      is_causal = (
          self.attn_mask_type == AttnMaskType.causal
          and not inference_mode
      )
  
      if attention_bias is not None:
          return self.flash_attention_triton(
              query_layer,
              key_layer,
              value_layer,
              attention_mask,
              attention_bias,
              is_causal,
          )
      else:
          return self.flash_attention_cuda(
              query_layer,
              key_layer,
              value_layer,
              attention_mask,
              is_causal,
          )
  
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 4 for calling the CoreAttention obj's underlying member method code with example usage: obj.flash_attention(query_layer, key_layer, value_layer, attention_mask, attention_bias, inference_mode) and returns: self, self