code: def _flash_attn_ssa(
      query, key, value, attention_mask=None, head_mask=None
  ):
      # transform the data into the qkv packed form
      qkv = torch.stack(
          [query, key, value], dim=2
      )  # [bsz, nh, 3, q_len, hd]
      qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
      bsz, q_len = qkv.shape[:2]
  
      qkv = rearrange(qkv, "b s ... -> (b s) ...")
      cu_q_lens = torch.arange(
          0,
          (bsz + 1) * q_len,
          step=q_len,
          dtype=torch.int32,
          device=qkv.device,
      )
      output = flash_attn_varlen_qkvpacked_func(
          qkv, cu_q_lens, q_len, 0.0, softmax_scale=None, causal=True
      )
      output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
  
      # disable attn weights by returning None when using flash attention
      return output, None
  
file_path: long-lora/gptneox_attn_replace.py
imports: from typing import Optional, Tuple
  import warnings
  import torch
  import transformers
  from einops import rearrange
  from flash_attn import flash_attn_varlen_qkvpacked_func, flash_attn_varlen_func
  from flash_attn.bert_padding import unpad_input, pad_input
label_desc: Use image node_id 3 for calling a global function with example usage: _flash_attn_ssa(query, key, value, attention_mask, head_mask) and returns: output,  None