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","."]
docstr: Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
  first unpad the input, then computes the attention scores and pad the final attention scores.
  
  Args:
      query_states (`torch.Tensor`):
          Input query states to be passed to Flash Attention API
      key_states (`torch.Tensor`):
          Input key states to be passed to Flash Attention API
      value_states (`torch.Tensor`):
          Input value states to be passed to Flash Attention API
      attention_mask (`torch.Tensor`):
          The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
          position of padding tokens and 1 for the position of non-padding tokens.
      dropout (`int`, *optional*):
          Attention dropout
      softmax_scale (`float`, *optional*):
          The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
docstr_tok: ["Calls","the","forward","method","of","Flash","Attention","-","if","the","input","hidden","states","contain","at","least","one","padding","token","first","unpad","the","input",",","then","computes","the","attention","scores","and","pad","the","final","attention","scores",".","Args",":","query_states","(","`","torch.Tensor","`",")",":","Input","query","states","to","be","passed","to","Flash","Attention","API","key_states","(","`","torch.Tensor","`",")",":","Input","key","states","to","be","passed","to","Flash","Attention","API","value_states","(","`","torch.Tensor","`",")",":","Input","value","states","to","be","passed","to","Flash","Attention","API","attention_mask","(","`","torch.Tensor","`",")",":","The","padding","mask","-","corresponds","to","a","tensor","of","size","`","(","batch_size",",","seq_len",")","`","where","0","stands","for","the","position","of","padding","tokens","and","1","for","the","position","of","non-padding","tokens",".","dropout","(","`","int","`",",","*","optional","*",")",":","Attention","dropout","softmax_scale","(","`","float","`",",","*","optional","*",")",":","The","scaling","of","QK^T","before","applying","softmax",".","Default","to","1","\/","sqrt","(","head_dim",")"]
code: def _flash_attention_forward(
      self,
      query_states,
      key_states,
      value_states,
      attention_mask,
      query_length,
      dropout=0.0,
      softmax_scale=None,
  ):
      """
      Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
      first unpad the input, then computes the attention scores and pad the final attention scores.
  
      Args:
          query_states (`torch.Tensor`):
              Input query states to be passed to Flash Attention API
          key_states (`torch.Tensor`):
              Input key states to be passed to Flash Attention API
          value_states (`torch.Tensor`):
              Input value states to be passed to Flash Attention API
          attention_mask (`torch.Tensor`):
              The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
              position of padding tokens and 1 for the position of non-padding tokens.
          dropout (`int`, *optional*):
              Attention dropout
          softmax_scale (`float`, *optional*):
              The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
      """
      if not self._flash_attn_uses_top_left_mask:
          causal = self.is_causal
      else:
          # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
          causal = self.is_causal and query_length != 1
  
      # Contains at least one padding token in the sequence
      if attention_mask is not None:
          batch_size = query_states.shape[0]
          (
              query_states,
              key_states,
              value_states,
              indices_q,
              cu_seq_lens,
              max_seq_lens,
          ) = self._upad_input(
              query_states,
              key_states,
              value_states,
              attention_mask,
              query_length,
          )
  
          cu_seqlens_q, cu_seqlens_k = cu_seq_lens
          max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
  
          attn_output_unpad = flash_attn_varlen_func(
              query_states,
              key_states,
              value_states,
              cu_seqlens_q=cu_seqlens_q,
              cu_seqlens_k=cu_seqlens_k,
              max_seqlen_q=max_seqlen_in_batch_q,
              max_seqlen_k=max_seqlen_in_batch_k,
              dropout_p=dropout,
              softmax_scale=softmax_scale,
              causal=causal,
          )
  
          attn_output = pad_input(
              attn_output_unpad, indices_q, batch_size, query_length
          )
      else:
          attn_output = flash_attn_func(
              query_states,
              key_states,
              value_states,
              dropout,
              softmax_scale=softmax_scale,
              causal=causal,
          )
  
      return attn_output
  
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 3 for calling the GPTBigCodeFlashAttention2 obj's underlying member method code with example usage: obj._flash_attention_forward(query_states, key_states, value_states, attention_mask, query_length, dropout, softmax_scale) and returns: attn_output