code: def test_split_rotary_position_qkv():
      batch_size, offset_seq_len = kt.bs_sl()
      nhead = kt.nhead
      head_dim = 128
      seq_len = 1
      seq_len = random.randint(1, 2048)
      offset_seq_len = random.randint(0, 2048 - seq_len)
      outshape = kt.rand((batch_size, nhead, seq_len, head_dim))
  
      cachek = kt.rand((batch_size, nhead, offset_seq_len, head_dim))
      cachev = kt.rand((batch_size, nhead, offset_seq_len, head_dim))
      q_tensor = kt.rand((batch_size, seq_len, nhead, head_dim))
      k_tensor = kt.rand((batch_size, seq_len, nhead, head_dim))
      v_tensor = kt.rand((batch_size, seq_len, nhead, head_dim))
      qkv_tensor = torch.cat((q_tensor, k_tensor, v_tensor), dim=2)
  
      out_cachek = torch.cat((cachek, outshape), dim=2)
      out_cachev = torch.cat((cachev, outshape), dim=2)
  
      func = None
      if kt.dtype == torch.float:
          func = cuda_module.torch_launch_split_rotary_position_fp32
      elif kt.dtype == torch.half:
          func = cuda_module.torch_launch_split_rotary_position_fp16
  
      custom_q = torch.empty_like(q_tensor)
  
      def custom():
          func(
              qkv_tensor,
              custom_q,
              out_cachek,
              out_cachev,
              batch_size,
              nhead,
              offset_seq_len,
              seq_len,
              head_dim,
          )
          return [
              custom_q.contiguous(),
              out_cachek.contiguous(),
              out_cachev.contiguous(),
          ]
  
      inv_freq = 1.0 / (
          10000
          ** (
              torch.arange(0, head_dim, 2).float().to(device="cuda:0")
              / head_dim
          )
      )
      t = torch.arange(2048, device="cuda:0", dtype=inv_freq.dtype)
      freqs = torch.einsum("i,j->ij", t, inv_freq)
      # Different from paper, but it uses a different permutation in order to obtain the same calculation
      emb = torch.cat((freqs, freqs), dim=-1)
      cos_cached = emb.cos()[None, None, :, :].to(
          device="cuda:0", dtype=kt.dtype
      )
      sin_cached = emb.sin()[None, None, :, :].to(
          device="cuda:0", dtype=kt.dtype
      )
  
      def baseline():
          trans_q = q_tensor.transpose(1, 2)
          trans_k = k_tensor.transpose(1, 2)
          trans_v = v_tensor.transpose(1, 2)
          kv_seq_len = offset_seq_len + seq_len
          cos = cos_cached[:, :, :kv_seq_len, ...]
          sin = sin_cached[:, :, :kv_seq_len, ...]
          gather_indices = (
              (torch.arange(seq_len) + offset_seq_len)[
                  None, None, :, None
              ]
              .repeat(batch_size, cos.shape[1], 1, cos.shape[3])
              .to("cuda:0")
          )
          cos = torch.gather(
              cos.repeat(gather_indices.shape[0], 1, 1, 1),
              2,
              gather_indices,
          )
          sin = torch.gather(
              sin.repeat(gather_indices.shape[0], 1, 1, 1),
              2,
              gather_indices,
          )
  
          def rotate_half(x):
              """Rotates half the hidden dims of the input."""
              x1 = x[..., : x.shape[-1] // 2]
              x2 = x[..., x.shape[-1] // 2 :]
              return torch.cat((-x2, x1), dim=-1)
  
          q_out = (trans_q * cos) + (rotate_half(trans_q) * sin)
          k_out = (trans_k * cos) + (rotate_half(trans_k) * sin)
          k_out = torch.cat((cachek, k_out), dim=2)
          v_out = torch.cat((cachev, trans_v), dim=2)
          return [
              q_out.contiguous(),
              k_out.contiguous(),
              v_out.contiguous(),
          ]
  
      return custom, baseline
  
imports: import sys
  from __init__ import lightseq_dir
  import random
  import torch
  from torch._C import ScriptModule, dtype
  from torch.nn import functional
  from csrc.pytorch.builder import CudaKernelBuilder
  from csrc.tests.util import TestDecorator
  from torch_crf import CRF
  from transformers import LlamaModel
  from transformers.activations import SiLUActivation
label_desc: Use image node_id 31 for calling a global function with example usage: test_split_rotary_position_qkv() and returns: custom,  baseline, list, list, torch