class_docstr: Region where selective activation recomputation is applied.
  This region is memory intensive but less compute intensive which
  makes activation checkpointing more efficient for LLMs (20B+).
  See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
  
  We use the following notation:
   h: hidden size
   n: number of attention heads
   p: number of tensor model parallel partitions
   b: batch size
   s: sequence length
class_docstr_tok: ["Region","where","selective","activation","recomputation","is","applied",".","This","region","is","memory","intensive","but","less","compute","intensive","which","makes","activation","checkpointing","more","efficient","for","LLMs","(","20B+",")",".","See","Reducing","Activation","Recomputation","in","Large","Transformer","Models",":","https",":","\/\/arxiv.org\/abs\/2205.05198","for","more","details",".","We","use","the","following","notation",":","h",":","hidden","size","n",":","number","of","attention","heads","p",":","number","of","tensor","model","parallel","partitions","b",":","batch","size","s",":","sequence","length"]
code: def forward(
      self,
      query: Tensor,
      key: Tensor,
      value: Tensor,
      attention_mask: Tensor,
      attn_mask_type: AttnMaskType = None,
  ):
      # ===================================
      # Raw attention scores. [b, n/p, s, s]
      # ===================================
  
      # expand the key and value [sk, b, ng, hn] -> [sk, b, np, hn]
      # This is a noop for normal attention where ng == np. When using group query attention this
      # creates a view that has the keys and values virtually repeated along their dimension to
      # match the number of queries.
  
      # attn_mask_type is not used.
      if (
          self.num_attention_heads_per_partition
          // self.num_query_groups_per_partition
          > 1
      ):
          key = key.repeat_interleave(
              self.num_attention_heads_per_partition
              // self.num_query_groups_per_partition,
              dim=2,
          )
          value = value.repeat_interleave(
              self.num_attention_heads_per_partition
              // self.num_query_groups_per_partition,
              dim=2,
          )
  
      # [b, np, sq, sk]
      output_size = (
          query.size(1),
          query.size(2),
          query.size(0),
          key.size(0),
      )
  
      # [sq, b, np, hn] -> [sq, b * np, hn]
      # This will be a simple view when doing normal attention, but in group query attention
      # the key and value tensors are repeated to match the queries so you can't use simple strides
      # to extract the queries.
      query = query.reshape(
          output_size[2], output_size[0] * output_size[1], -1
      )
      # [sk, b, np, hn] -> [sk, b * np, hn]
      key = key.view(
          output_size[3], output_size[0] * output_size[1], -1
      )
  
      # preallocting input tensor: [b * np, sq, sk]
      matmul_input_buffer = (
          parallel_state.get_global_memory_buffer().get_tensor(
              (
                  output_size[0] * output_size[1],
                  output_size[2],
                  output_size[3],
              ),
              query.dtype,
              "mpu",
          )
      )
  
      # Raw attention scores. [b * np, sq, sk]
      matmul_result = torch.baddbmm(
          matmul_input_buffer,
          query.transpose(0, 1),  # [b * np, sq, hn]
          key.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
          beta=0.0,
          alpha=(1.0 / self.norm_factor),
      )
  
      # change view to [b, np, sq, sk]
      attention_scores = matmul_result.view(*output_size)
  
      # ===========================
      # Attention probs and dropout
      # ===========================
  
      # attention scores and attention mask [b, np, sq, sk]
      attention_probs: Tensor = self.scale_mask_softmax(
          attention_scores, attention_mask
      )
  
      # This is actually dropping out entire tokens to attend to, which might
      # seem a bit unusual, but is taken from the original Transformer paper.
  
      if not self.config.sequence_parallel:
          with tensor_parallel.get_cuda_rng_tracker().fork():
              attention_probs = self.attention_dropout(attention_probs)
      else:
          attention_probs = self.attention_dropout(attention_probs)
  
      # =========================
      # Context layer. [sq, b, hp]
      # =========================
  
      # value -> context layer.
      # [sk, b, np, hn] --> [b, np, sq, hn]
  
      # context layer shape: [b, np, sq, hn]
      output_size = (
          value.size(1),
          value.size(2),
          query.size(0),
          value.size(3),
      )
  
      # change view [sk, b * np, hn]
      value = value.view(
          value.size(0), output_size[0] * output_size[1], -1
      )
  
      # change view [b * np, sq, sk]
      attention_probs = attention_probs.view(
          output_size[0] * output_size[1], output_size[2], -1
      )
  
      # matmul: [b * np, sq, hn]
      context = torch.bmm(attention_probs, value.transpose(0, 1))
  
      # change view [b, np, sq, hn]
      context = context.view(*output_size)
  
      # [b, np, sq, hn] --> [sq, b, np, hn]
      context = context.permute(2, 0, 1, 3).contiguous()
  
      # [sq, b, np, hn] --> [sq, b, hp]
      new_context_shape = context.size()[:-2] + (
          self.hidden_size_per_partition,
      )
      context = context.view(*new_context_shape)
  
      return context
  
file_path: megatron-lm/megatron/core/transformer/dot_product_attention.py
imports: import math
  import torch
  from torch import Tensor
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.fusions.fused_softmax import FusedScaleMaskSoftmax
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.transformer.utils import attention_mask_func
  from megatron.core.utils import divide
label_desc: Use image node_id 2 for calling the DotProductAttention obj's underlying member method code with example usage: obj.forward(query, key, value, attention_mask, attn_mask_type) and returns: context