code: def attention(
      q,
      k,
      v,
      out,
      cu_seqlens,
      max_s,
      softmax_scale,
      window_size_left=-1,
  ):
      if window_size_left <= 0 and window_size_left != -1:
          raise ValueError("`window_size_left` must be > 0 or -1")
  
      if HAS_FLASH_ATTN_V2_CUDA:
          return flash_attn_2_cuda.varlen_fwd(
              q,
              k,
              v,
              out,
              cu_seqlens,
              cu_seqlens,
              max_s,
              max_s,
              0.0,
              softmax_scale,
              False,
              True,
              window_size_left,
              0,
              False,
              None,
          )
      elif HAS_FLASH_ATTN_V2_ROCM:
          if window_size_left != -1:
              raise ValueError(
                  f"RoCm version of Flash Attention v2 does not support window attention (window_size_left != -1, got window_size_left={window_size_left})."
              )
  
          # RoCm flash API does not take the window_size_left and window_size_right arguments.
          return flash_attn_2_cuda.varlen_fwd(
              q,
              k,
              v,
              out,
              cu_seqlens,
              cu_seqlens,
              max_s,
              max_s,
              0.0,
              softmax_scale,
              False,
              True,
              False,
              None,
          )
      elif HAS_FLASH_ATTN:
          if window_size_left != -1:
              raise NotImplementedError(
                  "window_size_left is only available with flash attn v2"
              )
  
          # Flash attention v1 requires q, k and v to have the same number of heads
          if k.shape[1] != q.shape[1]:
              # MQA expand
              if k.shape[1] == 1:
                  k = k.expand(-1, q.shape[1], -1)
              # Grouped attention reshape
              else:
                  original_shape = k.shape
                  k = (
                      k.unsqueeze(2)
                      .expand(-1, -1, q.shape[1] // k.shape[1], -1)
                      .reshape(original_shape[0], -1, original_shape[2])
                  )
          if v.shape[1] != q.shape[1]:
              # MQA expand
              if v.shape[1] == 1:
                  v = v.expand(-1, q.shape[1], -1)
              # Grouped attention reshape
              else:
                  original_shape = v.shape
                  v = (
                      v.unsqueeze(2)
                      .expand(-1, -1, q.shape[1] // v.shape[1], -1)
                      .reshape(original_shape[0], -1, original_shape[2])
                  )
  
          return flash_attn_cuda.fwd(
              q,
              k,
              v,
              out,
              cu_seqlens,
              cu_seqlens,
              max_s,
              max_s,
              0.0,
              softmax_scale,
              False,
              True,
              False,
              0,
              None,
          )
  
      raise NotImplementedError("flash attention is not installed")
  
imports: import os
  import torch
  from loguru import logger
  from text_generation_server.utils.import_utils import IS_CUDA_SYSTEM, IS_ROCM_SYSTEM
label_desc: Use image node_id 1 for calling a global function with example usage: attention(q, k, v, out, cu_seqlens, max_s, softmax_scale, window_size_left) and returns: flash_attn_2_cuda, flash_attn_2_cuda, flash_attn_cuda