code: def decoder_layer_forward(
      self,
      hidden_states: Optional[torch.FloatTensor],
      attention_mask: Optional[torch.FloatTensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      past_key_value: Optional[Tuple[torch.Tensor]] = None,
      output_attentions: Optional[bool] = False,
      use_cache: Optional[bool] = False,
      cu_seqlens: Optional[torch.Tensor] = None,
      max_seqlen: Optional[torch.Tensor] = None,
  ) -> Union[
      Tuple[torch.Tensor],
      Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]],
  ]:
      # pylint: disable=duplicate-code
      residual = hidden_states
  
      hidden_states = self.input_layernorm(hidden_states)
  
      # Self Attention
      (
          hidden_states,
          self_attn_weights,
          present_key_value,
      ) = self.self_attn(
          hidden_states=hidden_states,
          attention_mask=attention_mask,
          position_ids=position_ids,
          past_key_value=past_key_value,
          output_attentions=output_attentions,
          use_cache=use_cache,
          cu_seqlens=cu_seqlens,
          max_seqlen=max_seqlen,
      )
      hidden_states = residual + hidden_states
  
      # Fully Connected
      residual = hidden_states
      hidden_states = self.post_attention_layernorm(hidden_states)
      hidden_states = self.mlp(hidden_states)
      hidden_states = residual + hidden_states
  
      outputs = (hidden_states,)
  
      if output_attentions:
          outputs += (self_attn_weights,)
  
      if use_cache:
          outputs += (present_key_value,)
  
      return outputs
  
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 6 for calling a global function with example usage: decoder_layer_forward(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cu_seqlens, max_seqlen) and returns: outputs