import torch
import unittest
from torch.nn import functional as F
import torch_mlu_ops as ops
from common_utils import *
import random

def genSlotMapping(batch, block_size):
  output = []
  for i in range(batch):
    idx = random.randint(i * block_size, (i + 1) * block_size - 1)
    output.append(idx)

  return output

def rotate(x: torch.Tensor):
  x1, x2 = x.chunk(2, dim=-1)
  return torch.cat((-x2, x1), dim=-1)

def quant(input: torch.Tensor):
  input_fp32 = input.to(torch.float32)
  max_value, _ = torch.max(input_fp32.abs(), dim=-1, keepdim=True)
  scale = max_value / 7
  scaled_input = torch.round(input_fp32 / scale)
  return scaled_input.to(torch.int8), scale[..., 0], input_fp32 / scale

class TestFusedRopeOp(BtTestCase):
  def op_impl_base(self, *args):
    qkv, k_cache_hp, v_cache_hp, sin_cache, cos_cache, position_id, gamma, beta, \
       k_cache_lp, v_cache_lp, cache_bs_id_hp, cache_seq_offsets_hp, cache_bs_id_lp, \
          cache_seq_offsets_lp, k_scale_hp, v_scale_hp, k_scale_lp, v_scale_lp, slot_mapping_hp, \
             slot_mapping_lp, eps = args
    mixed_cache = k_cache_lp is not None and v_cache_lp is not None
    qkv = qkv.to(torch.float32)
    sin_cache = sin_cache.to(torch.float32)
    cos_cache = cos_cache.to(torch.float32)
    gamma = gamma.to(torch.float32)
    beta = beta.to(torch.float32)
    if k_scale_hp is not None:
      k_scale_hp = 1 / k_scale_hp
    if v_scale_hp is not None:
      v_scale_hp = 1 / v_scale_hp
    if not mixed_cache and k_scale_hp is None:
      k_cache_hp = k_cache_hp.to(torch.float32)
      v_cache_hp = v_cache_hp.to(torch.float32)
    discrete_batch_hp = cache_bs_id_hp is not None
    discrete_batch_lp = cache_bs_id_lp is not None
    paged_cache_hp = slot_mapping_hp is not None
    paged_cache_lp = slot_mapping_lp is not None
    head_size = qkv.shape[-1]
    rope_dim = head_size
    batch_size = qkv.shape[0]
    head_qkv = qkv.shape[-2]
    kv_heads = k_cache_hp.shape[1]
    q_heads = head_qkv - 2 * kv_heads
    head_qk = q_heads + kv_heads
    group_num = 1
    group_size = head_size
    if mixed_cache:
      group_num = k_scale_lp.shape[-1]
      group_size = int(head_size / group_num)
      k_cache_lp_shape = list(k_cache_lp.shape)
      k_cache_lp_shape[-1] *= 2
      k_cache_lp_int8 = UnpackInt4(k_cache_lp).reshape(k_cache_lp_shape)
      v_cache_lp_shape = list(v_cache_lp.shape)
      v_cache_lp_shape.append(2)
      v_cache_lp_int8 = UnpackInt4(v_cache_lp).reshape(v_cache_lp_shape)

    qk = qkv[:, :, 0:q_heads + kv_heads].clone()

    for i in range(batch_size):
      qk_i = qk[i]
      sin_cache_i = sin_cache[position_id[i]:position_id[i] + 1]
      cos_cache_i = cos_cache[position_id[i]:position_id[i] + 1]
      sin_cache_i = sin_cache_i[:1]
      cos_cache_i = cos_cache_i[:1]
      rot = rotate(qk_i)

      qk_i[:] = rot * sin_cache_i.unsqueeze(1) + qk_i * cos_cache_i.unsqueeze(1)

    q = qk[:, :, 0:q_heads]
    k = qk[:, :, q_heads:head_qk].contiguous().reshape(batch_size, kv_heads, head_size)

    qkv_q = qkv[:, :, 0:q_heads]
    qkv_q[...] = q

    shape_k = k.shape
    k = torch.reshape(k, (-1, shape_k[-1]))
    k_norm = F.layer_norm(k, (head_size,), gamma, beta, eps)
    k_norm = torch.reshape(k_norm, shape_k)

    k_out = k_norm.reshape(batch_size, 1, kv_heads, head_size)
    v_out = qkv[:, :, head_qk:head_qkv].contiguous()
    k_out_hp = k_out.clone()
    v_out_hp = v_out.clone()
    k_out_lp = None
    v_out_lp = None
    if k_scale_hp is not None and v_scale_hp is not None:
      k_scale_hp = k_scale_hp.reshape(kv_heads, head_size)
      v_scale_hp = v_scale_hp.reshape(kv_heads, head_size)
      k_out_hp = (k_out * k_scale_hp).round().clamp(-128, 127).to(torch.int8)
      v_out_hp = (v_out * v_scale_hp).round().clamp(-128, 127).to(torch.int8)

    if paged_cache_hp:
      block_size = k_cache_hp.shape[2]
      for i in range(batch_size):
          if slot_mapping_hp[i] >= 0:
              block_id = torch.div(slot_mapping_hp[i], block_size, rounding_mode='floor')
              block_offset = slot_mapping_hp[i] % block_size
              k_cache_hp[block_id, :, block_offset, :] = k_out_hp[i]
              v_cache_hp[block_id, :, block_offset, :] = v_out_hp[i]
    else:
      for i in range(batch_size):
        key_i = k_out_hp[i].transpose(1, 0)
        value_i = v_out_hp[i].transpose(1, 0)

        cache_bs_id_hp_i = cache_bs_id_hp[i] if discrete_batch_hp else i
        cache_seqlen_offset_hp_i = cache_seq_offsets_hp[i]
        if cache_seqlen_offset_hp_i < 0 or cache_bs_id_hp_i < 0:
          continue

        key_cache_hp_i = \
          k_cache_hp[cache_bs_id_hp_i, :, cache_seqlen_offset_hp_i:cache_seqlen_offset_hp_i + 1]
        key_cache_hp_i[...] = key_i[...]

        value_cache_hp_i = \
          v_cache_hp[cache_bs_id_hp_i, :, cache_seqlen_offset_hp_i:cache_seqlen_offset_hp_i + 1]
        value_cache_hp_i[...] = value_i[...]

    if mixed_cache:
      for i in range(batch_size):
        key_i = k_out[i].reshape(kv_heads, group_num, group_size)
        value_i = v_out[i].reshape(kv_heads, group_num, group_size)
        key_i_lp, key_scale_i_lp, _ = quant(key_i)
        value_i_lp, value_scale_i_lp, scaled_input = quant(value_i)

        if paged_cache_lp:
          block_size = k_cache_lp_int8.shape[2]
          if slot_mapping_lp[i] >= 0:
            block_id = torch.div(slot_mapping_lp[i], block_size, rounding_mode='floor')
            block_offset_k = slot_mapping_lp[i] % block_size
            block_offset_v = torch.div(block_offset_k, 2, rounding_mode='floor')
            odd_even_offset = block_offset_k % 2
            k_cache_lp_int8[block_id, :, block_offset_k, :] = key_i_lp.reshape(kv_heads, head_size)
            v_cache_lp_int8[block_id, :, block_offset_v, :, odd_even_offset] = value_i_lp.reshape(kv_heads, head_size)
            k_scale_lp[block_id, :, block_offset_k, :] = key_scale_i_lp[...]
            v_scale_lp[block_id, :, block_offset_k, :] = value_scale_i_lp[...]
        else:
            cache_bs_id_lp_i = cache_bs_id_lp[i] if discrete_batch_lp else i
            cache_seqlen_offset_lp_i = cache_seq_offsets_lp[i]
            v_seq_offset = torch.div(cache_seqlen_offset_lp_i, 2, rounding_mode='floor')
            odd_even_offset = cache_seqlen_offset_lp_i % 2
            if cache_seqlen_offset_lp_i < 0 or cache_bs_id_lp_i < 0:
              continue
            key_cache_lp_i = \
              k_cache_lp_int8[cache_bs_id_lp_i, :, cache_seqlen_offset_lp_i]
            key_cache_lp_i[...] = key_i_lp.reshape(kv_heads, head_size)
            key_scale_lp_i = \
              k_scale_lp[cache_bs_id_lp_i, :, cache_seqlen_offset_lp_i]
            key_scale_lp_i[...] = key_scale_i_lp[...]

            value_cache_lp_i = \
              v_cache_lp_int8[cache_bs_id_lp_i, :, v_seq_offset, :, odd_even_offset]
            value_cache_lp_i[...] = value_i_lp.reshape(kv_heads, head_size)
            value_scale_lp_i = \
              v_scale_lp[cache_bs_id_lp_i, :, cache_seqlen_offset_lp_i]
            value_scale_lp_i[...] = value_scale_i_lp[...]
    out = (qkv, k_cache_hp, v_cache_hp)
    if mixed_cache:
      k_cache_lp = PairlyPackInt8(k_cache_lp_int8.view(-1, head_size)).reshape(k_cache_lp.shape)
      v_cache_lp_int8 = v_cache_lp_int8.transpose(2, 3)
      s0,s1,s2,s3,s4 = v_cache_lp_int8.shape
      v_cache_lp = PairlyPackInt8(v_cache_lp_int8.reshape(-1, s3 * s4)).reshape(s0, s1, s2, s3 * s4 // 2).transpose(2,3)
      out += (k_cache_lp, v_cache_lp)
      if k_scale_lp is not None:
        out += (k_scale_lp, v_scale_lp)
    return out

  @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "fused_rope not support MLU3XX device")
  def test_fused_rope(self):
    random_cases = 100
    random.seed(355)
    eps = 1e-6
    for _ in range(random_cases):
      if random_cases == 1:
        bs = 512
        seq_len = 1
        q_heads = 8
        kv_heads = 1
        head_size = 128
        rope_dim = 128
        dtype = torch.float16
        need_quant_kv = False
        mixed_cache = False
        max_decode_len_hp = 128
        max_decode_len_lp = 128
        discrete_batch_hp = True
        discrete_batch_lp = False
        paged_cache_hp = False
        paged_cache_lp = True
        block_size_hp = 16
        block_size_lp = 128
        num_blocks_hp = int((bs * max_decode_len_hp + block_size_hp - 1) / block_size_hp)
        num_blocks_lp = int((bs * max_decode_len_lp + block_size_lp - 1) / block_size_lp)
        group_num = 1
      else:
        bs = random.randint(1, 512)
        seq_len = 1
        q_heads = random.randint(1, 32)
        kv_heads = random.randint(1, 32)
        head_size_list = [32, 64, 96, 128]
        head_size = random.choice(head_size_list)
        rope_dim = head_size
        dtype_list = [torch.half, torch.bfloat16]
        bool_list = [True, False]
        block_size_list = [16, 32, 64, 128]
        dtype = random.choice(dtype_list)
        need_quant_kv = random.choice(bool_list)
        mixed_cache = random.choice(bool_list)
        max_decode_len_hp = random.randint(128, 1024)
        max_decode_len_lp = random.randint(128, 1024)
        discrete_batch_hp = random.choice(bool_list)
        discrete_batch_lp = random.choice(bool_list)
        paged_cache_hp = random.choice(bool_list)
        paged_cache_lp = random.choice(bool_list)
        block_size_hp = random.choice(block_size_list)
        num_blocks_hp = int((bs * max_decode_len_hp + block_size_hp - 1) / block_size_hp)
        block_size_lp = random.choice(block_size_list)
        num_blocks_lp = int((bs * max_decode_len_lp + block_size_lp - 1) / block_size_lp)
        group_num_list = [1, 2, 4, 8]
        group_num = random.choice(group_num_list)

      if mixed_cache:
        need_quant_kv = True
      group_size = int(head_size / group_num)
      if max_decode_len_lp % 2 != 0:
        max_decode_len_lp = max_decode_len_lp - 1
      print("bs: {}, seq_len: {}, q_heads: {}, kv_heads: {}, head_size: {}, rope_dim: {}, "
            "dtype: {}, mixed_cache: {}, quant_kv: {}, paged_cache_hp: {}, paged_cache_lp: {}, "
            "discrete_batch_hp: {}, discrete_batch_lp: {}."
            .format(bs, seq_len, q_heads, kv_heads, head_size, rope_dim, dtype, mixed_cache,
            need_quant_kv, paged_cache_hp, paged_cache_lp, discrete_batch_hp, discrete_batch_lp))

      if mixed_cache:
        print("max_decode_len_hp: {}, max_decode_len_lp: {}, num_blocks_hp: {}, num_blocks_lp: {}, "
              "block_size_hp: {}, block_size_lp: {}, group_num: {}, testing..."
              .format(max_decode_len_hp, max_decode_len_lp, num_blocks_hp, num_blocks_lp,
              block_size_hp, block_size_lp, group_num))

      max_bs_hp = bs + 1 if discrete_batch_hp else bs
      max_bs_lp = bs + 1 if discrete_batch_lp else bs
      cache_size = 1 if need_quant_kv else 2
      cache_bytes_hp = num_blocks_hp * kv_heads * block_size_hp * head_size if paged_cache_hp else \
                       max_bs_hp * kv_heads * max_decode_len_hp * head_size
      cache_bytes_hp = cache_bytes_hp * cache_size
      cache_bytes_lp = 0
      if mixed_cache:
        cache_bytes_lp = num_blocks_lp * kv_heads * block_size_lp * head_size if paged_cache_lp else \
                         max_bs_lp * kv_heads * max_decode_len_lp * head_size

      if cache_bytes_hp > 2**31 or cache_bytes_lp > 2**31:
        print("cache bytes can not be larger than int32max. ")
        continue

      input_shape = (bs, seq_len, q_heads + 2 * kv_heads, head_size)
      input = torch.randn(size=input_shape, dtype=dtype).mlu()

      if paged_cache_hp:
        cache_shape_hp = (num_blocks_hp, kv_heads, block_size_hp, head_size)
      else:
        cache_shape_hp = (max_bs_hp, kv_heads, max_decode_len_hp, head_size)
      if need_quant_kv:
        k_cache_hp = torch.randn(cache_shape_hp, dtype=dtype, device='mlu')
        k_cache_hp = (k_cache_hp - 0.5) * 256
        k_cache_hp = k_cache_hp.to(torch.int8)
        v_cache_hp = torch.randn(cache_shape_hp, dtype=dtype, device='mlu')
        v_cache_hp = (k_cache_hp - 0.5) * 256
        v_cache_hp = k_cache_hp.to(torch.int8)
        k_scale_ops_hp = 1 / (torch.randn(size=(kv_heads, head_size), dtype=torch.float).abs().mlu() + 0.01)
        v_scale_ops_hp = 1 / (torch.randn(size=(kv_heads, head_size), dtype=torch.float).abs().mlu() + 0.01)
      else:
        k_cache_hp = torch.randn(cache_shape_hp, dtype=dtype, device='mlu')
        v_cache_hp = torch.randn(cache_shape_hp, dtype=dtype, device='mlu')
        k_scale_ops_hp = None
        v_scale_ops_hp = None

      k_cache_lp = None
      v_cache_lp = None
      k_scale_lp = None
      v_scale_lp = None
      if mixed_cache:
        if paged_cache_lp:
          k_scale_lp = torch.randn(size=(num_blocks_lp, kv_heads, block_size_lp, group_num), dtype=torch.float).mlu()
          v_scale_lp = torch.randn(size=(num_blocks_lp, kv_heads, block_size_lp, group_num), dtype=torch.float).mlu()
          cache_raw = torch.randn((num_blocks_lp * kv_heads * block_size_lp, head_size), dtype=dtype, device='mlu')
          max_value = torch.amax(torch.abs(cache_raw))
          cache_raw = cache_raw * (7 / max_value)
          cache_raw = cache_raw.to(torch.int8)
          k_cache_lp = PairlyPackInt8(cache_raw).reshape(num_blocks_lp, kv_heads, block_size_lp, int(head_size / 2))
          cache_raw = torch.randn((int(num_blocks_lp * kv_heads * block_size_lp / 2), head_size * 2), dtype=dtype, device='mlu')
          max_value = torch.amax(torch.abs(cache_raw))
          cache_raw = cache_raw * (7 / max_value)
          cache_raw = cache_raw.to(torch.int8)
          v_cache_lp = PairlyPackInt8(cache_raw).reshape(num_blocks_lp, kv_heads, int(block_size_lp / 2), head_size)
          k_cache_lp_ref_shape = (num_blocks_lp, kv_heads, block_size_lp, head_size)
          v_cache_lp_ref_shape = (num_blocks_lp, kv_heads, int(block_size_lp / 2), head_size, 2)
        else:
          k_scale_lp = torch.randn(size=(max_bs_lp, kv_heads, max_decode_len_lp, group_num), dtype=torch.float).mlu()
          v_scale_lp = torch.randn(size=(max_bs_lp, kv_heads, max_decode_len_lp, group_num), dtype=torch.float).mlu()
          cache_raw = torch.randn((max_bs_lp * kv_heads * max_decode_len_lp, head_size), dtype=dtype, device='mlu')
          max_value = torch.amax(torch.abs(cache_raw))
          cache_raw = cache_raw * (7 / max_value)
          cache_raw = cache_raw.to(torch.int8)
          k_cache_lp = PairlyPackInt8(cache_raw).reshape(max_bs_lp, kv_heads, max_decode_len_lp, int(head_size / 2))
          cache_raw = torch.randn((int(max_bs_lp * kv_heads * max_decode_len_lp / 2), head_size * 2), dtype=dtype, device='mlu')
          max_value = torch.amax(torch.abs(cache_raw))
          cache_raw = cache_raw * (7 / max_value)
          cache_raw = cache_raw.to(torch.int8)
          v_cache_lp = PairlyPackInt8(cache_raw).reshape(max_bs_lp,kv_heads, int(max_decode_len_lp / 2), head_size)
          k_cache_lp_ref_shape = (max_bs_lp, kv_heads, max_decode_len_lp, head_size)
          v_cache_lp_ref_shape = (max_bs_lp, kv_heads, int(max_decode_len_lp / 2), head_size, 2)
        del cache_raw

      cache_bs_id_hp = None
      cache_bs_id_lp = None
      cache_seq_offsets_hp = None
      cache_seq_offsets_lp = None
      slot_mapping_hp = None
      slot_mapping_lp = None
      if not paged_cache_hp:
        if discrete_batch_hp:
          cache_bs_id_hp = random.sample([*range(0, max_bs_hp)], bs)
          cache_bs_id_hp = torch.IntTensor(cache_bs_id_hp).mlu()
        cache_seq_offsets_hp = torch.randint(size=(bs, ), low=-1, high=max_decode_len_hp - 2,
                                             dtype=torch.int32, device='mlu')
      else:
        slot_mapping_hp = random.sample([*range(-1, block_size_hp * num_blocks_hp)], bs)
        slot_mapping_hp = torch.IntTensor(slot_mapping_hp).mlu()

      input_ref = input.clone()
      k_cache_hp_ref = k_cache_hp.clone()
      v_cache_hp_ref = v_cache_hp.clone()
      k_cache_lp_ref = None
      v_cache_lp_ref = None
      k_scale_lp_ref = None
      v_scale_lp_ref = None
      if mixed_cache:
        if not paged_cache_lp:
          if discrete_batch_lp:
            cache_bs_id_lp = random.sample([*range(0, max_bs_lp)], bs)
            cache_bs_id_lp = torch.IntTensor(cache_bs_id_lp).mlu()
          cache_seq_offsets_lp = torch.randint(size=(bs, ), low=-1, high=max_decode_len_lp - 2,
                                             dtype=torch.int32, device='mlu')
        else:
          slot_mapping_lp = genSlotMapping(bs, block_size_lp)
          slot_mapping_lp = torch.IntTensor(slot_mapping_lp).mlu()
        k_cache_lp_ref = k_cache_lp.clone()
        v_cache_lp_ref = v_cache_lp.clone()
        k_scale_lp_ref = k_scale_lp.clone()
        v_scale_lp_ref = v_scale_lp.clone()

      cos_table = torch.randn(size=(bs, rope_dim), dtype=dtype).mlu()
      sin_table = torch.randn(size=(bs, rope_dim), dtype=dtype).mlu()
      gamma = torch.randn(size=(head_size, ), dtype=dtype).mlu()
      beta = torch.randn(size=(head_size, ), dtype=dtype).mlu()
      position_id = torch.randint(size=(bs, ), low=0, high=bs, dtype=torch.int32, device='mlu')

      notify_start = torch.mlu.Event(enable_timing=True)
      notify_end = torch.mlu.Event(enable_timing=True)
      notify_start.record()
      base_output = self.op_impl_base(input_ref, k_cache_hp_ref, v_cache_hp_ref, sin_table,      \
                        cos_table, position_id, gamma, beta,                                     \
                        k_cache_lp_ref, v_cache_lp_ref, cache_bs_id_hp, cache_seq_offsets_hp,    \
                        cache_bs_id_lp, cache_seq_offsets_lp, k_scale_ops_hp, v_scale_ops_hp,    \
                        k_scale_lp_ref, v_scale_lp_ref, slot_mapping_hp, slot_mapping_lp, eps)
      notify_end.record()
      notify_end.synchronize()
      time = notify_start.hardware_time(notify_end)
      print("baseline hw_time is: ", time, "us")
      del input_ref, k_cache_hp_ref, v_cache_hp_ref, k_cache_lp_ref, v_cache_lp_ref, k_scale_lp_ref, v_scale_lp_ref

      notify_start.record()
      loop = 1
      for _ in range(loop):
        ops.fused_rope(input, k_cache_hp, v_cache_hp, sin_table, cos_table, position_id,          \
                       gamma, beta, k_cache_lp, v_cache_lp, cache_bs_id_hp, cache_seq_offsets_hp, \
                       cache_bs_id_lp, cache_seq_offsets_lp, k_scale_ops_hp, v_scale_ops_hp,      \
                       k_scale_lp, v_scale_lp, slot_mapping_hp, slot_mapping_lp, eps)

      notify_end.record()
      notify_end.synchronize()
      time = notify_start.hardware_time(notify_end) / loop
      print("hw time is: ", time, "us")

      print("check input diff \n")
      self.assertTensorsEqual(input.cpu().float(), base_output[0].cpu().float(), 0.003, use_MSE=True)
      print("pass \n")
      print("check key cache hp diff \n")
      self.assertTensorsEqual(k_cache_hp.cpu().float(), base_output[1].cpu().float(), 0.003, use_MSE=True)
      print("pass \n")
      print("check value cache hp diff \n")
      self.assertTensorsEqual(v_cache_hp.cpu().float(), base_output[2].cpu().float(), 0.003, use_MSE=True)
      print("pass \n")
      if mixed_cache:
        k_cache_lp_int8 = UnpackInt4(k_cache_lp).reshape(k_cache_lp_ref_shape)
        v_cache_lp_int8 = UnpackInt4(v_cache_lp).reshape(v_cache_lp_ref_shape)
        k_cache_lp_ref_int8 = UnpackInt4(base_output[3]).reshape(k_cache_lp_ref_shape)
        v_cache_lp_ref_int8 = UnpackInt4(base_output[4]).reshape(v_cache_lp_ref_shape)

        print("check key cache lp diff \n")
        self.assertTensorsEqual(k_cache_lp_int8.cpu().float(), k_cache_lp_ref_int8.cpu().float(), 1)
        print("pass \n")
        print("check key scale lp diff \n")
        self.assertTensorsEqual(k_scale_lp.cpu().float(), base_output[5].cpu().float(), 0.003, use_MSE=True)
        print("pass \n")
        print("check value cache lp diff \n")
        self.assertTensorsEqual(v_cache_lp_int8.cpu().float(), v_cache_lp_ref_int8.cpu().float(), 1)
        print("pass \n")
        print("check value scale lp diff \n")
        self.assertTensorsEqual(v_scale_lp.cpu().float(), base_output[6].cpu().float(), 0.003, use_MSE=True)
        print("pass \n")

  @unittest.skipIf('MLU3' in torch.mlu.get_device_name(), "fused_rope not support MLU3XX device")
  def test_inductor(self):
    bs, seq_len, q_heads, kv_heads, head_size, rope_dim, max_decode_len, dtype= 40, 1, 8, 1, 128, 128, 2048, torch.bfloat16
    input_shape = (bs, seq_len, q_heads + 2 * kv_heads, head_size)
    input = torch.randn(size=input_shape, dtype=dtype).mlu()
    max_bs = bs + 1
    cos_table = torch.randn(size=(bs, rope_dim), dtype=dtype).mlu()
    sin_table = torch.randn(size=(bs, rope_dim), dtype=dtype).mlu()
    gamma = torch.randn(size=(head_size, ), dtype=dtype).mlu()
    beta = torch.randn(size=(head_size, ), dtype=dtype).mlu()

    cache = torch.randn((2, max_bs, kv_heads, max_decode_len, head_size), dtype=dtype, device='mlu')
    k_cache = cache[0]
    v_cache = cache[1]
    cache_bs_id = random.sample([*range(0, max_bs)], bs)
    cache_bs_id = torch.IntTensor(cache_bs_id).mlu()
    cache_seq_offsets = torch.randint(size=(bs, ), low=-1, high=max_decode_len - 2,
                                        dtype=torch.int32, device='mlu')
    position_id = torch.randint(size=(bs, ), low=0, high=bs, dtype=torch.int32, device='mlu')
    args = (input, k_cache, v_cache, None, None, sin_table, cos_table, position_id, gamma, beta, None,
            None, None, None, cache_bs_id, cache_seq_offsets, None, None, None, None, 1e-5)
    self.base_opcheck(torch.ops.torch_mlu_ops.fused_rope, args)

if __name__ == "__main__":
  exit(run_unittest(TestFusedRopeOp))
