code: def _rotary_kernel(
      q,
      input_scale,
      output_scale,
      Cos,
      Sin,
      q_bs_stride,
      q_h_stride,
      q_d_stride,
      cos_bs_stride,
      cos_d_stride,
      total_len,
      HEAD_NUM: tl.constexpr,
      BLOCK_HEAD: tl.constexpr,
      BLOCK_SEQ: tl.constexpr,
      HEAD_DIM: tl.constexpr,
  ):
      current_head_index = tl.program_id(0)
      current_seq_index = tl.program_id(1)
  
      dim_range0 = tl.arange(0, HEAD_DIM // 2)
      dim_range1 = tl.arange(HEAD_DIM // 2, HEAD_DIM)
  
      current_head_range = current_head_index * BLOCK_HEAD + tl.arange(
          0, BLOCK_HEAD
      )
      current_seq_range = current_seq_index * BLOCK_SEQ + tl.arange(
          0, BLOCK_SEQ
      )
  
      off_q0 = (
          current_seq_range[:, None, None] * q_bs_stride
          + current_head_range[None, :, None] * q_h_stride
          + dim_range0[None, None, :] * q_d_stride
      )
      off_q1 = (
          current_seq_range[:, None, None] * q_bs_stride
          + current_head_range[None, :, None] * q_h_stride
          + dim_range1[None, None, :] * q_d_stride
      )
  
      off_dimcos_sin = (
          current_seq_range[:, None, None] * cos_bs_stride
          + dim_range0[None, None, :] * cos_d_stride
      )
  
      q0 = tl.load(
          q + off_q0,
          mask=(current_seq_range[:, None, None] < total_len)
          & (current_head_range[None, :, None] < HEAD_NUM),
          other=0.0,
      )
      q1 = tl.load(
          q + off_q1,
          mask=(current_seq_range[:, None, None] < total_len)
          & (current_head_range[None, :, None] < HEAD_NUM),
          other=0.0,
      )
  
      cos = tl.load(
          Cos + off_dimcos_sin,
          mask=current_seq_range[:, None, None] < total_len,
          other=0.0,
      )
      sin = tl.load(
          Sin + off_dimcos_sin,
          mask=current_seq_range[:, None, None] < total_len,
          other=0.0,
      )
  
      q0 = q0.to(tl.float32) * input_scale
      q1 = q1.to(tl.float32) * input_scale
  
      out0 = (q0 * cos - q1 * sin) / output_scale
      out1 = (q0 * sin + q1 * cos) / output_scale
  
      out0 = out0.to(tl.int8)
      out1 = out1.to(tl.int8)
  
      tl.store(
          q + off_q0,
          out0,
          mask=(current_seq_range[:, None, None] < total_len)
          & (current_head_range[None, :, None] < HEAD_NUM),
      )
      tl.store(
          q + off_q1,
          out1,
          mask=(current_seq_range[:, None, None] < total_len)
          & (current_head_range[None, :, None] < HEAD_NUM),
      )
  
      return
  
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 1 for calling a global function with example usage: _rotary_kernel(q, input_scale, output_scale, Cos, Sin, q_bs_stride, q_h_stride, q_d_stride, cos_bs_stride, cos_d_stride, total_len, HEAD_NUM, BLOCK_HEAD, BLOCK_SEQ, HEAD_DIM) without return types