code: def flashattn_attn(
      self,
      hidden_states: torch.FloatTensor,
      attention_mask: torch.FloatTensor,
      position_ids: torch.LongTensor,
      past_key_value: Optional[Tuple[torch.Tensor]] = None,
      output_attentions: Optional[
          bool
      ] = False,  # pylint: disable=unused-argument
      use_cache: Optional[bool] = False,
      cu_seqlens: Optional[torch.Tensor] = None,
      max_seqlen: Optional[torch.Tensor] = None,
  ) -> Tuple[
      torch.Tensor,
      Optional[torch.Tensor],
      Optional[Tuple[torch.Tensor]],
  ]:
      bsz, q_len, _ = hidden_states.size()
  
      query_states = self.q_proj(hidden_states)
      key_states = self.k_proj(hidden_states)
      value_states = self.v_proj(hidden_states)
  
      query_states = query_states.view(
          bsz, q_len, self.num_heads, self.head_dim
      ).transpose(1, 2)
      key_states = key_states.view(
          bsz, q_len, self.num_key_value_heads, self.head_dim
      ).transpose(1, 2)
      value_states = value_states.view(
          bsz, q_len, self.num_key_value_heads, self.head_dim
      ).transpose(1, 2)
  
      query_rot = query_states[..., : self.rotary_ndims]
      query_pass = query_states[..., self.rotary_ndims :]
      key_rot = key_states[..., : self.rotary_ndims]
      key_pass = key_states[..., self.rotary_ndims :]
  
      kv_seq_len = key_states.shape[-2]
      if past_key_value is not None:
          kv_seq_len += past_key_value[0].shape[-2]
      cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
      query_states, key_states = apply_rotary_pos_emb(
          query_rot, key_rot, cos, sin, position_ids
      )
  
      # [batch_size, num_heads, seq_len, head_dim]
      query_states = torch.cat((query_states, query_pass), dim=-1)
      key_states = torch.cat((key_states, key_pass), dim=-1)
  
      if past_key_value is not None:
          # Reuse k, v, self_attention
          key_states = torch.cat((past_key_value[0], key_states), dim=2)
          value_states = torch.cat(
              (past_key_value[1], value_states), dim=2
          )
  
      past_key_value = (key_states, value_states) if use_cache else None
  
      # Repeat k/v heads if n_kv_heads < n_heads
      key_states = repeat_kv(key_states, self.num_key_value_groups)
      value_states = repeat_kv(value_states, self.num_key_value_groups)
  
      if (
          cu_seqlens is not None
          and max_seqlen is not None
          and cu_seqlens.dim() == 1
      ):
          # special handling using sample packing
          qkv = torch.stack(
              [query_states, key_states, value_states], dim=2
          )  # [bsz, nh, 3, q_len, hd]
          qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
          qkv = rearrange(qkv, "b s ... -> (b s) ...")
          softmax_scale = None
  
          output = flash_attn_varlen_qkvpacked_func(
              qkv,
              cu_seqlens,
              max_seqlen,
              0.0,
              softmax_scale=softmax_scale,
              causal=True,
          )
  
          attn_output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
          attn_output = rearrange(attn_output, "b s h d -> b s (h d)")
      else:
          attn_weights = torch.matmul(
              query_states, key_states.transpose(2, 3)
          ) / math.sqrt(self.head_dim)
  
          if attn_weights.size() != (
              bsz,
              self.num_heads,
              q_len,
              kv_seq_len,
          ):
              raise ValueError(
                  f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
                  f" {attn_weights.size()}"
              )
  
          if attention_mask is not None:
              if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                  raise ValueError(
                      f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
                  )
              attn_weights = attn_weights + attention_mask
  
          # Upcast attention to fp32
          attn_weights = nn.functional.softmax(
              attn_weights, dim=-1, dtype=torch.float32
          ).to(query_states.dtype)
          attn_output = torch.matmul(attn_weights, value_states)
  
          if attn_output.size() != (
              bsz,
              self.num_heads,
              q_len,
              self.head_dim,
          ):
              raise ValueError(
                  f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
                  f" {attn_output.size()}"
              )
  
          # Merge heads
          attn_output = attn_output.transpose(1, 2).contiguous()
          attn_output = attn_output.reshape(
              bsz, q_len, self.hidden_size
          )
  
      # Final linear projection
      attn_output = self.o_proj(attn_output)
  
      return attn_output, None, past_key_value
  
file_path: axolotl/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
imports: import importlib
  import math
  from typing import Optional, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from accelerate import init_empty_weights
  from einops import rearrange
  from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
  from torch import nn
  from transformers import AutoConfig, AutoModelForCausalLM
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.utils import logging
  from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
label_desc: Use image node_id 5 for calling a global function with example usage: flashattn_attn(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cu_seqlens, max_seqlen) and returns: attn_output,  None,  past_key_value