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)
      qkv = qkv.view(
          -1, self.num_groups, self.num_heads + 2, self.head_size
      )
  
      # Split on group dimension
      query, kv = qkv.split(
          [self.num_heads, 2],
          dim=2,
      )
      # Merge groups and heads
      query = query.reshape(
          -1, self.num_groups * self.num_heads, self.head_size
      )
  
      # Inplace rotary
      self.rotary_emb(query, torch.select(kv, dim=2, index=0), cos, sin)
  
      paged_attention.reshape_and_cache(
          kv[:, :, 0].contiguous(),
          kv[:, :, 1].contiguous(),
          kv_cache[0],
          kv_cache[1],
          slots,
      )
  
      # output
      attn_output = torch.empty_like(query)
  
      # Prefill
      if cu_seqlen_prefill is not None:
          # flash attention
          flash_attn.attention(
              query,
              torch.select(kv, dim=2, index=0),
              torch.select(kv, dim=2, 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.dense(
          attn_output.view(
              -1, self.num_groups * self.num_heads * self.head_size
          )
      )
  
imports: import torch
  import torch.distributed
  from torch import nn
  from transformers.modeling_utils import PreTrainedModel
  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.flash_attn import attention
  from text_generation_server.utils.layers import TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelHead, FastLayerNorm, PositionRotaryEmbedding, get_linear
label_desc: Use image node_id 2 for calling the FlashRWLargeAttention 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