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_triton(
      self,
      query_layer,
      key_layer,
      value_layer,
      attention_mask,
      attention_bias,
      is_causal,
  ):
      if self.attention_dropout_p > 0.0:
          raise NotImplementedError(
              f"attention_dropout not implemented for flash_attention with attention bias"
          )
  
      if attention_mask is not None:
          if len(attention_mask.shape) == 4:
              # [b, 1, sq, sk] -> [b, 1, sq, 1] / [b, 1, 1, sk]
              attention_mask_q = torch.any(
                  torch.eq(attention_mask, False), dim=3
              ).unsqueeze(3)
              attention_mask_kv = torch.any(
                  torch.eq(attention_mask, False), dim=2
              ).unsqueeze(2)
          else:
              # [b, s] -> [b, 1, s, 1] / [b, 1, 1, s]
              assert len(attention_mask.shape) == 2
              attention_mask_q = (
                  (~attention_mask).unsqueeze(1).unsqueeze(3)
              )
              attention_mask_kv = (
                  (~attention_mask).unsqueeze(1).unsqueeze(2)
              )
  
          if attention_bias.shape[2] == attention_mask_q.shape[2]:
              attention_bias = attention_bias.masked_fill(
                  ~attention_mask_q, torch.finfo(query_layer.dtype).min
              )
          if attention_bias.shape[3] == attention_mask_kv.shape[3]:
              attention_bias = attention_bias.masked_fill(
                  ~attention_mask_kv, torch.finfo(query_layer.dtype).min
              )
  
      context_layer = flash_attn_func_triton(
          query_layer,
          key_layer,
          value_layer,
          attention_bias,
          is_causal,
      )
  
      # [b, sq, np, hn] -> [b, np, sq, hn]
      context_layer = context_layer.permute(0, 2, 1, 3)
  
      if attention_mask is not None:
          context_layer = context_layer * attention_mask_q
  
      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 6 for calling the CoreAttention obj's underlying member method code with example usage: obj.flash_attention_triton(query_layer, key_layer, value_layer, attention_mask, attention_bias, is_causal) and returns: context_layer