code: def test_flash_self_attention(self, cfg, model_parallel_config):
      device = cfg["device"]
      layer_number = cfg["layer_number"]
      bz, sl, np, h = cfg["bz"], cfg["sq"], cfg["head"], cfg["hidden"]
      hn = h // np
  
      q = torch.rand(sl, bz, np, hn, device=device).half()
      k = torch.rand(sl, bz, np, hn, device=device).half()
      v = torch.rand(sl, bz, np, hn, device=device).half()
  
      attention_mask_2d = torch.arange(sl, device=device).unsqueeze(
          0
      ) < torch.randint(1, sl, (bz,), device=device).unsqueeze(1)
  
      attention_mask_padding_3d = build_attention_mask_3d(
          source_mask=attention_mask_2d,
          target_mask=attention_mask_2d,
          attn_mask_type=AttnMaskType.padding,
      ).unsqueeze(1)
  
      attention_mask_causal_3d = build_attention_mask_3d(
          source_mask=attention_mask_2d,
          target_mask=attention_mask_2d,
          attn_mask_type=AttnMaskType.causal,
      ).unsqueeze(1)
  
      # Non-causal
      attention = CoreAttention(
          config=model_parallel_config,
          layer_number=layer_number,
          num_attention_heads=np,
          hidden_size=h,
          attn_mask_type=AttnMaskType.padding,
          attention_dropout=0.0,
      )
  
      attention_fa = CoreAttention(
          config=model_parallel_config,
          layer_number=layer_number,
          num_attention_heads=np,
          hidden_size=h,
          attn_mask_type=AttnMaskType.padding,
          attention_dropout=0.0,
          use_flash_attention=True,
      )
  
      out = attention(q, k, v, attention_mask_padding_3d)
      out_fa = attention_fa(q, k, v, attention_mask_padding_3d)
      torch.testing.assert_close(out, out_fa)
      out_fa = attention_fa(q, k, v, ~attention_mask_2d)
      torch.testing.assert_close(out, out_fa)
  
      # Causal
      attention = CoreAttention(
          config=model_parallel_config,
          layer_number=layer_number,
          num_attention_heads=np,
          hidden_size=h,
          attn_mask_type=AttnMaskType.causal,
          attention_dropout=0.0,
          apply_query_key_layer_scaling=False,
      )
  
      attention_fa = CoreAttention(
          config=model_parallel_config,
          layer_number=layer_number,
          num_attention_heads=np,
          hidden_size=h,
          attn_mask_type=AttnMaskType.causal,
          attention_dropout=0.0,
          use_flash_attention=True,
      )
  
      out = attention(q, k, v, attention_mask_causal_3d)
      out_fa = attention_fa(q, k, v, attention_mask_causal_3d)
      torch.testing.assert_close(out, out_fa)
      out_fa = attention_fa(q, k, v, ~attention_mask_2d)
      torch.testing.assert_close(out, out_fa)
  
file_path: NeMo/tests/collections/nlp/test_flash_attention.py
imports: import random
  import pytest
  import torch
  from megatron.core import ModelParallelConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.modules.common.megatron.attention import CoreAttention
  from nemo.collections.nlp.modules.common.megatron.megatron_init import initialize_model_parallel_for_nemo
  from nemo.collections.nlp.modules.common.megatron.utils import build_attention_mask_3d
  from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
label_desc: Use image node_id 4 for calling the TestFlashAttention obj's underlying member method code with example usage: obj.test_flash_self_attention(cfg, model_parallel_config) without return types