code: def test_cross_attn_inference(self, model_parallel_config):
      num_layers = 1
      init_method_std = 0.02
      batch = 2
      neighbors = 2
      # rotary pos emb dim
      dim = 128
      pad_id = 19999
      num_attention_heads = 8
      chunks = 32
      text_chunk_size = 64
      context_chunk_size = 2 * text_chunk_size
      input_length = chunks * text_chunk_size
      vocab_size = 20000
  
      rot_dim = dim // num_attention_heads
      rotary_pos_emb = RotaryEmbedding(rot_dim).cuda().half()
  
      hidden = torch.randint(
          0, vocab_size, (input_length, batch)
      ).cuda()  # (seq, batch, dim)
      hidden_mask = (hidden != pad_id).cuda()
      hidden_emb = (
          torch.rand(input_length, batch, dim).cuda().half()
      )  # (seq, batch, dim)
  
      retrieved = torch.randint(
          0, vocab_size, (chunks, neighbors, context_chunk_size, batch)
      ).cuda()
      # retrieved tokens - (num chunks, num retrieved neighbors, retrieved chunk with continuation, batch)
  
      # context attention mask [b, np, sq, sk]
      context_mask = (retrieved != pad_id).cuda()
      retrieved_emb = (
          torch.rand(chunks, neighbors, context_chunk_size, batch, dim)
          .cuda()
          .half()
      )
      # retrieved tokens - (num chunks, num retrieved neighbors, retrieved chunk with continuation, batch, hidden)
  
      # need to add extra chunk size, since it will be shifted
      cross_attn_q_pos_emb = rotary_pos_emb(
          text_chunk_size + text_chunk_size - 1,
          offset=-text_chunk_size + 1,
      )
      cross_attn_k_pos_emb = rotary_pos_emb(context_chunk_size)
      cross_attn_pos_emb = (cross_attn_q_pos_emb, cross_attn_k_pos_emb)
  
      def get_attn_mask_3d(hidden_mask, context_mask, chunks):
          causal_padding = text_chunk_size - 1
          reminder = (
              text_chunk_size - (hidden_mask.shape[0] + 1)
          ) % text_chunk_size
          hidden_mask = F.pad(
              hidden_mask,
              (0, 0, -causal_padding, reminder),
              value=False,
          )
  
          dec_attn_mask = rearrange(
              hidden_mask, "(k n) b -> (b k) n", k=chunks
          )
          context_attn_mask = rearrange(
              context_mask, "k r n b -> (b k) (r n)"
          )
          enc_dec_attn_mask_3d = build_attention_mask_3d(
              source_mask=dec_attn_mask,
              target_mask=context_attn_mask,
              attn_mask_type=AttnMaskType.padding,
          )
          enc_dec_attn_mask_3d = enc_dec_attn_mask_3d[:, None, :, :]
          return enc_dec_attn_mask_3d
  
      enc_dec_attn_mask_3d = get_attn_mask_3d(
          hidden_mask, context_mask, chunks
      )
  
      init_method = init_method_normal(init_method_std)
  
      scaled_init_method = scaled_init_method_normal(
          init_method_std, num_layers
      )
      cross_attn = (
          ParallelChunkedCrossAttention(
              config=model_parallel_config,
              init_method=init_method,
              output_layer_init_method=scaled_init_method,
              layer_number=1,
              num_attention_heads=num_attention_heads,
              hidden_size=dim,
              precision=16,
              chunk_size=text_chunk_size,
              masked_softmax_fusion=False,
          )
          .cuda()
          .half()
      )
  
      out, bias = cross_attn(
          hidden_emb,
          enc_dec_attn_mask_3d,
          encoder_output=retrieved_emb,
          rotary_pos_emb=cross_attn_pos_emb,
      )
      assert out.shape == torch.Size([input_length, batch, dim])
      assert bias.shape == torch.Size([dim])
  
      attn_mask_3d = None
  
      out_1, b = cross_attn(
          hidden_emb[:62],
          attn_mask_3d,
          encoder_output=None,
          rotary_pos_emb=cross_attn_pos_emb,
          set_inference_key_value_memory=True,
          inference_max_sequence_len=input_length,
      )
      assert (
          out_1 - torch.zeros_like(hidden_emb[:62])
      ).abs().max() == 0
      out_1, b = cross_attn(
          hidden_emb[62:63],
          attn_mask_3d,
          encoder_output=None,
          rotary_pos_emb=cross_attn_pos_emb,
          set_inference_key_value_memory=False,
          inference_max_sequence_len=input_length,
      )
      assert (
          out_1 - torch.zeros_like(hidden_emb[62:63])
      ).abs().max() == 0
  
      attn_mask_3d = get_attn_mask_3d(
          hidden_mask[:64], context_mask[:1], 1
      )
      out_2, b = cross_attn(
          hidden_emb[63:64],
          attn_mask_3d,
          encoder_output=retrieved_emb[:1],
          rotary_pos_emb=cross_attn_pos_emb,
          set_inference_key_value_memory=False,
          inference_max_sequence_len=input_length,
      )
      assert (out[63] - out_2[0]).abs().max().item() < 1e-2
  
      for i in range(64, 127):
          attn_mask_3d = get_attn_mask_3d(
              hidden_mask[: i + 1], context_mask[:1], 1
          )
          out_2, b = cross_attn(
              hidden_emb[i : i + 1],
              attn_mask_3d,
              encoder_output=retrieved_emb[:1],
              rotary_pos_emb=cross_attn_pos_emb,
              set_inference_key_value_memory=False,
              inference_max_sequence_len=input_length,
          )
      i = 127
      attn_mask_3d = get_attn_mask_3d(
          hidden_mask[: i + 1], context_mask[:2], 2
      )
      out_3, b = cross_attn(
          hidden_emb[i : i + 1],
          attn_mask_3d,
          encoder_output=retrieved_emb[:2],
          rotary_pos_emb=cross_attn_pos_emb,
          set_inference_key_value_memory=False,
          inference_max_sequence_len=input_length,
      )
      assert (out[i] - out_3[0]).abs().max().item() < 1e-2
  
      attn_mask_3d = get_attn_mask_3d(
          hidden_mask[:130], context_mask[:2], 2
      )
  
      out_1, b = cross_attn(
          hidden_emb[:130],
          attn_mask_3d,
          encoder_output=retrieved_emb[:2],
          rotary_pos_emb=cross_attn_pos_emb,
          set_inference_key_value_memory=True,
          inference_max_sequence_len=input_length,
      )
  
      assert (out[:130] - out_1[:130]).abs().max().item() < 1e-2
  
      for i in range(130, 191):
          attn_mask_3d = get_attn_mask_3d(
              hidden_mask[: i + 1], context_mask[:2], 2
          )
          out_2, b = cross_attn(
              hidden_emb[i : i + 1],
              attn_mask_3d,
              encoder_output=retrieved_emb[:2],
              rotary_pos_emb=cross_attn_pos_emb,
              set_inference_key_value_memory=False,
              inference_max_sequence_len=input_length,
          )
      i = 191
      attn_mask_3d = get_attn_mask_3d(
          hidden_mask[: i + 1], context_mask[:3], 3
      )
      out_4, b = cross_attn(
          hidden_emb[i : i + 1],
          attn_mask_3d,
          encoder_output=retrieved_emb[:3],
          rotary_pos_emb=cross_attn_pos_emb,
          set_inference_key_value_memory=False,
          inference_max_sequence_len=input_length,
      )
      assert (out[i] - out_4[0]).abs().max().item() < 1e-2
  
file_path: NeMo/tests/collections/nlp/test_retrieval_module_inference.py
imports: import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.modules.common.megatron.attention import ParallelChunkedCrossAttention
  from nemo.collections.nlp.modules.common.megatron.layer_type import LayerType
  from nemo.collections.nlp.modules.common.megatron.megatron_init import initialize_model_parallel_for_nemo
  from nemo.collections.nlp.modules.common.megatron.position_embedding import RotaryEmbedding
  from nemo.collections.nlp.modules.common.megatron.retrieval_token_level_encoder_decoder import MegatronRetrievalTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.retrieval_transformer import MegatronRetrievalTransformerDecoderModule, MegatronRetrievalTransformerEncoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import build_attention_mask_3d, init_method_normal, scaled_init_method_normal
  from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
label_desc: Use image node_id 3 for calling the TestRetrievalModuleInference obj's underlying member method code with example usage: obj.test_cross_attn_inference(model_parallel_config) and returns: enc_dec_attn_mask_3d