class_docstr: GPTBigCode flash attention module. This module inherits from `GPTBigCodeAttention` as the weights of the module
  stays untouched. The only required change would be on the forward pass where it needs to correctly call the public
  API of flash attention and deal with padding tokens in case the input contains any of them.
class_docstr_tok: ["GPTBigCode","flash","attention","module",".","This","module","inherits","from","`","GPTBigCodeAttention","`","as","the","weights","of","the","module","stays","untouched",".","The","only","required","change","would","be","on","the","forward","pass","where","it","needs","to","correctly","call","the","public","API","of","flash","attention","and","deal","with","padding","tokens","in","case","the","input","contains","any","of","them","."]
code: def forward(
      self,
      hidden_states: torch.Tensor,
      layer_past: Optional[torch.Tensor] = None,
      attention_mask: Optional[torch.Tensor] = None,
      head_mask: Optional[torch.Tensor] = None,
      encoder_hidden_states: Optional[torch.Tensor] = None,
      encoder_attention_mask: Optional[torch.Tensor] = None,
      use_cache: Optional[bool] = False,
      output_attentions: Optional[bool] = False,
  ) -> Union[
      Tuple[torch.Tensor, Optional[torch.Tensor]],
      Tuple[
          torch.Tensor, Optional[torch.Tensor], Tuple[torch.Tensor, ...]
      ],
  ]:
      if encoder_hidden_states is not None:
          if not hasattr(self, "q_attn") or not self.is_cross_attention:
              raise ValueError(
                  "If class is used as cross attention, the weights `q_attn` have to be defined. "
                  "Please make sure to instantiate class with `GPTBigCodeAttention(..., is_cross_attention=True)`."
              )
  
          query = self.q_attn(hidden_states)
          key_value = self.c_attn(encoder_hidden_states)
          attention_mask = encoder_attention_mask
      elif self.multi_query:
          query, key_value = self.c_attn(hidden_states).split(
              (self.embed_dim, 2 * self.kv_dim), dim=2
          )
      else:
          # Note: We split as (self.num_heads, 3, self.head_dim) instead of (3, self.num_heads, self.head_dim),
          # i.e., the memory layout is not the same as GPT2.
          # This makes the concatenation with past_key_value more efficient.
          query, key_value = (
              self.c_attn(hidden_states)
              .view(
                  *hidden_states.shape[:2],
                  self.num_heads,
                  3 * self.head_dim,
              )
              .transpose(1, 2)
              .split((self.head_dim, 2 * self.head_dim), dim=3)
          )
  
      if layer_past is not None:
          key_value = torch.cat((layer_past, key_value), dim=-2)
      present = key_value if use_cache else None
  
      key, value = key_value.split(
          (self.head_dim, self.head_dim), dim=-1
      )
  
      # Flash attention requires the input to have the shape
      # batch_size x seq_length x head_dim x hidden_dim
      if self.multi_query:
          batch_size, query_length, _ = query.shape
          query = query.reshape(
              batch_size, query_length, self.num_heads, self.head_dim
          )
          key = key.unsqueeze(2)
          value = value.unsqueeze(2)
      else:
          query_length = query.shape[2]
          batch_size, _, tgt, _ = key.shape
          query = query.transpose(1, 2).reshape(
              batch_size, query_length, self.num_heads, self.head_dim
          )
          key = key.transpose(1, 2).reshape(
              batch_size, tgt, self.num_heads, self.head_dim
          )
          value = value.transpose(1, 2).reshape(
              batch_size, tgt, self.num_heads, self.head_dim
          )
  
      attn_dropout = self.attn_pdrop if self.training else 0.0
  
      softmax_dtype = (
          torch.float32
          if self.attention_softmax_in_fp32
          else query.dtype
      )
      upcast = query.dtype != softmax_dtype
      softmax_scale = (
          self.layer_idx + 1
          if self.scale_attention_softmax_in_fp32 and upcast
          else 1
      )
      softmax_scale = softmax_scale**-1
      if self.scale_attn_weights:
          softmax_scale /= self.head_dim**0.5
  
      # In PEFT, usually we cast the layer norms in float32 for training stability reasons
      # therefore the input hidden states gets silently casted in float32. Hence, we need
      # cast them back in float16 just to be sure everything works as expected.
      input_dtype = query.dtype
      if input_dtype == torch.float32:
          if torch.is_autocast_enabled():
              target_dtype = torch.get_autocast_gpu_dtype()
          # Handle the case where the model is quantized
          elif hasattr(self.config, "_pre_quantization_dtype"):
              target_dtype = self.config._pre_quantization_dtype
          else:
              target_dtype = self.c_attn.weight.dtype
  
          logger.warning_once(
              f"The input hidden states seems to be silently casted in float32, this might be related to"
              f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
              f" {target_dtype}."
          )
          query = query.to(target_dtype)
          key = key.to(target_dtype)
          value = value.to(target_dtype)
  
      attn_output = self._flash_attention_forward(
          query,
          key,
          value,
          attention_mask,
          query_length,
          dropout=attn_dropout,
          softmax_scale=softmax_scale,
      )
  
      attn_weights_reshaped = attn_output.reshape(
          batch_size, query_length, self.num_heads * self.head_dim
      )
      attn_output = self.c_proj(attn_weights_reshaped)
      attn_output = self.resid_dropout(attn_output)
  
      outputs = (attn_output, present)
  
      if output_attentions:
          if self.multi_query:
              # Transpose to return weights in the usual format (batch_size, num_heads, query_length, key_length)
              attn_weights_reshaped = attn_weights_reshaped.transpose(
                  1, 2
              )
      else:
          attn_weights_reshaped = None
  
      outputs += (attn_weights_reshaped,)
  
      return outputs
  
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from ...activations import ACT2FN
  from ...modeling_attn_mask_utils import AttentionMaskConverter
  from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions, SequenceClassifierOutputWithPast, TokenClassifierOutput
  from ...modeling_utils import PreTrainedModel
  from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10, logging
  from .configuration_gpt_bigcode import GPTBigCodeConfig
label_desc: Use image node_id 2 for calling the GPTBigCodeFlashAttention2 obj's underlying member method code with example usage: obj.forward(hidden_states, layer_past, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions) and returns: outputs