code: def forward(
      self,
      hidden_states,
      cos,
      sin,
      cu_seqlen_prefill,
      kv_cache,
      block_tables,
      slots,
      input_lengths,
      max_s,
  ):
      qkv = self.query_key_value(hidden_states)
      query, kv = qkv.split(
          [
              self.head_size * self.num_heads,
              2 * self.head_size * self.num_key_value_heads,
          ],
          dim=1,
      )
      query = query.view(-1, self.num_heads, self.head_size)
      kv = kv.view(-1, 2, self.num_key_value_heads, self.head_size)
  
      self.rotary_emb(query, torch.select(kv, dim=1, index=0), cos, sin)
  
      paged_attention.reshape_and_cache(
          kv[:, 0], kv[:, 1], kv_cache[0], kv_cache[1], slots
      )
  
      # output tensor
      attn_output = torch.empty_like(query)
  
      # Prefill
      if cu_seqlen_prefill is not None:
          # flash attention
          flash_attn.attention(
              query,
              torch.select(kv, dim=1, index=0),
              torch.select(kv, dim=1, index=1),
              attn_output,
              cu_seqlen_prefill,
              max_s,
              self.softmax_scale,
          )
      # Decode
      else:
          paged_attention.attention(
              attn_output,
              query,
              kv_cache[0],
              kv_cache[1],
              self.kv_head_mapping,
              self.softmax_scale,
              block_tables,
              input_lengths,
              max_s,
          )
  
      return self.o_proj(
          attn_output.view(-1, self.num_heads * self.head_size)
      )
  
imports: import torch
  import torch.distributed
  from torch import nn
  from transformers.activations import ACT2FN
  from transformers.configuration_utils import PretrainedConfig
  from typing import Optional, List, Tuple
  from text_generation_server.utils import paged_attention, flash_attn
  from text_generation_server.utils.layers import TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, PositionRotaryEmbedding, TensorParallelHead, get_linear, FastRMSNorm
label_desc: Use image node_id 2 for calling the FlashLlamaAttention obj's underlying member method code with example usage: obj.forward(hidden_states, cos, sin, cu_seqlen_prefill, kv_cache, block_tables, slots, input_lengths, max_s) and returns: self