import math
import random
import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from itertools import product

def gen_args(batch, head_size, is_pagedattn, has_alibi, kv_data_type, max_seqlen, seq_q, head_num, num_kv_heads):
    input_qkv = torch.randn((batch, 3 * seq_q, 3 * head_num, head_size)).mlu().half()
    input_q = input_qkv[:,0:seq_q, 0:head_num,:].view(batch, seq_q, head_num, head_size)
    context_lens = torch.randint(seq_q, max_seqlen + 1, (batch, ), dtype=torch.int32).mlu()
    max_context_len = int(max(context_lens))
    if is_pagedattn:
        block_size = 16
    else:
        block_size = max_seqlen
    num_blocks = batch * ((max_seqlen + block_size - 1) // block_size)
    max_num_blocks_per_seq = (max_context_len + block_size - 1) // block_size
    cache_shape = (num_blocks, num_kv_heads, block_size, head_size)
    scale_shape = (num_blocks, num_kv_heads, block_size)
    block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
    block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
    if kv_data_type is not torch.int8:
        key_cache = torch.randn(size=cache_shape, dtype=torch.float16).mlu()
        value_cache = torch.randn(size=cache_shape, dtype=torch.float16).mlu()
        key_cache_scale = None
        value_cache_scale = None
    else:
        key_cache = torch.zeros(cache_shape).uniform_(-128, 128).to(kv_data_type).mlu()
        value_cache = torch.zeros(cache_shape).uniform_(-128, 128).to(kv_data_type).mlu()
        key_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
        value_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
    alibi_slopes = None
    if has_alibi:
        alibi_slopes = torch.zeros((batch, head_num), dtype=torch.float32).mlu()
        alibi_slopes.uniform_(0, 0.125)
    softmax_scale = 1 / math.sqrt(head_size)
    return (input_q.contiguous(), key_cache, value_cache, torch.empty_like(input_q), block_tables, context_lens, None,
             key_cache_scale, value_cache_scale, alibi_slopes, max_context_len, -1, -1, softmax_scale, False, -1)

def gen_params(batch, head_size, head_size_v, is_pagedattn, has_alibi, dtype, kv_dtype, max_seqlen, seq_q, head_num, num_kv_heads):
    input_qkv = torch.randn((batch, 3 * seq_q, 3 * head_num, head_size), dtype = dtype).mlu()
    input_q = input_qkv[:,0:seq_q, 0:head_num,:].view(batch, seq_q, head_num, head_size)
    context_lens = torch.randint(seq_q, max_seqlen + 1, (batch, ), dtype=torch.int32).mlu()
    max_context_len = int(max(context_lens))
    if is_pagedattn:
        block_size = 16
    else:
        block_size = max_seqlen
    num_blocks = batch * ((max_seqlen + block_size - 1) // block_size)
    max_num_blocks_per_seq = (max_context_len + block_size - 1) // block_size
    cache_shape_k = (num_blocks, num_kv_heads, block_size, head_size)
    cache_shape_v = (num_blocks, num_kv_heads, block_size, head_size_v)
    scale_shape = (num_blocks, num_kv_heads, block_size)
    block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
    block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
    if kv_dtype is torch.int8:
        key_cache = torch.zeros(cache_shape_k).uniform_(-128, 127).to(torch.int8).mlu()
        value_cache = torch.zeros(cache_shape_v).uniform_(-128, 127).to(torch.int8).mlu()
        key_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
        value_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
    else:
        key_cache = torch.randn(size=cache_shape_k, dtype=dtype).mlu()
        value_cache = torch.randn(size=cache_shape_v, dtype=dtype).mlu()
        key_cache_scale = None
        value_cache_scale = None
    alibi_slopes = None
    if has_alibi:
        alibi_slopes = torch.zeros((batch, head_num), dtype=torch.float32).mlu()
        alibi_slopes.uniform_(0, 0.125)
    return input_q, key_cache, value_cache, block_tables, context_lens, key_cache_scale, value_cache_scale, alibi_slopes, max_context_len

class TestSingleQueryAttnOp(BtTestCase):
    def op_impl_base(self, *args):
        q, k_cache, v_cache, out, block_tables, context_lens, k_cache_quant_scale, v_cache_quant_scale, \
            alibi_slopes, max_contxt_len, windows_size_left, windows_size_right, softmax_scale, return_lse, \
                kv_cache_quant_bit_size = args
        base_output = single_query_cached_kv_attn(q, k_cache, v_cache, block_tables, context_lens, k_cache_quant_scale,
           v_cache_quant_scale, alibi_slopes, windows_size_left, windows_size_right, softmax_scale, return_lse)
        return base_output

    def test_single_query_attention(self):
        head_num = 16
        batch_list = [5, 12]
        num_kv_heads = 4
        head_size_list = [(128, 128), (192, 384)]
        seq_len_list = [512]
        is_pagedattn_list = [False, True]
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name:
           print("pagedattn is not implement on mlu370, skip it")
           is_pagedattn_list = [False]
        has_alibi_list = [True, False]
        seq_q_list = [1, 5]
        window_size_list = [(-1, -1), (10, -1)]
        data_type_list = [torch.float16, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            data_type_list.append(torch.bfloat16)

        args = product(batch_list, head_size_list, is_pagedattn_list, has_alibi_list, seq_len_list, seq_q_list, window_size_list, data_type_list)
        for batch, (head_size, head_size_v), is_pagedattn, has_alibi, max_seqlen, seq_q, (window_size_left, window_size_right), dtype in args:
            print("batch: {}, max_seqlen: {}, head_size: {}, head_size_v: {}, is_pagedattn: {}, has_alibi {}, seq_q {}, window_size_left {},\
window_size_right {}, dtype {}, testing...".format(
              batch, max_seqlen, head_size, head_size_v, is_pagedattn, has_alibi, seq_q, window_size_left, window_size_right, dtype))
            # prepare input
            params = gen_params(batch, head_size, head_size_v, is_pagedattn, has_alibi, dtype, dtype, max_seqlen, seq_q, head_num, num_kv_heads)
            input_q, key_cache, value_cache, block_tables, context_lens, key_cache_scale, value_cache_scale, alibi_slopes, max_context_len = params
            softmax_scale = 1 / math.sqrt(head_size)
            torch_output = self.op_impl_base(input_q, key_cache, value_cache,
                                          None, block_tables, context_lens, key_cache_scale,
                                          value_cache_scale, alibi_slopes, max_context_len,
                                          window_size_left, window_size_right, softmax_scale, False, -1)
            tmo_output1 = ops.single_query_cached_kv_attn(input_q, key_cache, value_cache,
                                          None, block_tables, context_lens, key_cache_scale,
                                          value_cache_scale, alibi_slopes, max_context_len,
                                          window_size_left, window_size_right, softmax_scale)
            self.assertTensorsEqual(torch_output.cpu().float(), tmo_output1.cpu().float(),
                                    0.003, use_MSE=True)
            if seq_q == 1:
                torch_output, torch_lse = self.op_impl_base(input_q, key_cache, value_cache,
                                            None, block_tables, context_lens, key_cache_scale,
                                            value_cache_scale, alibi_slopes, max_context_len,
                                            window_size_left, window_size_right, softmax_scale, True, -1)
                tmo_output, tmo_lse = ops.single_query_cached_kv_attn(input_q, key_cache, value_cache,
                                            None, block_tables, context_lens, key_cache_scale,
                                            value_cache_scale, alibi_slopes, max_context_len,
                                            window_size_left, window_size_right, softmax_scale, True)
                self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                        0.003, use_MSE=True)
                self.assertTensorsEqual(torch_lse.cpu().float(), tmo_lse.cpu().float(),
                                        0.0001, use_MSE=True)

    # @unittest.skip("not test")
    def test_single_query_attention_quantize_kv(self):
        head_num = 16
        batch_list = [5, 12]
        num_kv_heads = 4
        head_size_list = [(128, 128), (16, 384)]
        seq_len_list = [512]
        is_pagedattn_list = [False, True]
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name:
           print("pagedattn is not implement on mlu370, skip it")
           is_pagedattn_list = [False]
        has_alibi_list = [True, False]
        quant_mode_list = ['per_token', 'per_channel']
        kv_data_type_list = [torch.int8]
        seq_q_list = [1, 5]
        window_size_list = [(-1, -1), (10, -1)]
        args = product(batch_list, head_size_list, is_pagedattn_list, has_alibi_list, kv_data_type_list, quant_mode_list, seq_len_list, seq_q_list, window_size_list)
        for batch, (head_size, head_size_v), is_pagedattn, has_alibi, kv_data_type, quant_mode, max_seqlen, seq_q, (window_size_left, window_size_right) in args:
            print("batch: {}, max_seqlen: {}, head_size: {}, head_size_v: {}, is_pagedattn: {}, has_alibi {}, kv_datatype {}, \
quant_mode {}, seq_q {}, window_size_left {}, window_size_right {}, testing...".format(
              batch, max_seqlen, head_size, head_size_v, is_pagedattn, has_alibi, kv_data_type, quant_mode, seq_q, window_size_left, window_size_right))
            # prepare input
            params = gen_params(batch, head_size, head_size_v, is_pagedattn, has_alibi, torch.float16, kv_data_type, max_seqlen, seq_q, head_num, num_kv_heads)
            input_q, key_cache, value_cache, block_tables, context_lens, key_cache_scale, value_cache_scale, alibi_slopes, max_context_len = params
            softmax_scale = 1 / math.sqrt(head_size)
            torch_output_contiguous = self.op_impl_base(input_q.contiguous(),
                                   key_cache, value_cache, None, block_tables, context_lens,
                                   key_cache_scale, value_cache_scale, alibi_slopes,
                                   max_context_len, window_size_left, window_size_right, softmax_scale, False, -1)
            tmo_output_contiguous = ops.single_query_cached_kv_attn(input_q.contiguous(),
                                   key_cache, value_cache, None, block_tables, context_lens,
                                   key_cache_scale, value_cache_scale, alibi_slopes,
                                   max_context_len, window_size_left, window_size_right, softmax_scale)
            self.assertTensorsEqual(torch_output_contiguous.cpu().float(), tmo_output_contiguous.cpu().float(),
                                    0.003, use_MSE=True)
            if seq_q == 1:
                torch_output, torch_lse = self.op_impl_base(input_q, key_cache, value_cache,
                                    None, block_tables, context_lens, key_cache_scale, value_cache_scale,
                                    alibi_slopes, max_context_len, window_size_left, window_size_right, softmax_scale, True, -1)
                tmo_output, tmo_lse = ops.single_query_cached_kv_attn(input_q, key_cache, value_cache,
                                    None, block_tables, context_lens, key_cache_scale, value_cache_scale,
                                    alibi_slopes, max_context_len, window_size_left, window_size_right, softmax_scale, True)
                self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                        0.005, use_MSE=True)
                self.assertTensorsEqual(torch_lse.cpu().float(), tmo_lse.cpu().float(), 0.0003, use_MSE=True)

    def test_single_query_attention_int4_kv(self):
        head_num = 16
        batch_list = [5, 12]
        num_kv_heads = 4
        head_size_list = [(64, 128), (256, 128), (64, 384)]
        seq_len_list = [512]
        is_pagedattn_list = [False, True]
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name:
           print("pagedattn is not implement on mlu370, skip it")
           is_pagedattn_list = [False]
        has_alibi_list = [True, False]
        quant_mode_list = ['per_token', 'per_channel', 'per_token_group']
        data_type_list = [torch.float, torch.half]
        if torch_mlu.mlu.is_bf16_supported():
            data_type_list.append(torch.bfloat16)
        kv_data_type = torch.int8
        seq_q_list = [1, 5]
        #int4 range
        quant_bit = 4
        int_max = float(2 ** (quant_bit - 1) - 1)
        int_min = -float(2 ** (quant_bit - 1))
        window_size_list = [(-1, -1), (20, -1)]
        args = product(batch_list, head_size_list, is_pagedattn_list, has_alibi_list, quant_mode_list, seq_len_list, seq_q_list, window_size_list, data_type_list)
        for batch, (head_size, head_size_v), is_pagedattn, has_alibi, quant_mode, max_seqlen, seq_q, (window_size_left, window_size_right), data_type in args:
            print("kv4: batch: {},  head_size: {}, head_size_v: {}, is_pagedattn: {}, has_alibi {}, quant_mode {}, max_seqlen: {}, seq_q {}, \
window_size_left {}, window_size_right {}, data_type {}, testing...".format(
              batch, head_size, head_size_v, is_pagedattn, has_alibi, quant_mode, max_seqlen, seq_q, window_size_left, window_size_right, data_type))
            # prepare input
            input_qkv = torch.randn((batch, seq_q, 3 * head_num, head_size), dtype=data_type).mlu()
            input_q = input_qkv[..., 0:head_num,:]

            context_lens = torch.randint(seq_q, max_seqlen + 1, (batch, ), dtype=torch.int32).mlu()
            max_context_len = context_lens.max().item()
            if max_seqlen % 2 == 1:
                max_seqlen += 1
            if is_pagedattn:
                block_size = 16
            else:
                block_size = max_seqlen
            num_blocks = batch * ((max_seqlen + block_size - 1) // block_size)
            max_num_blocks_per_seq = (max_context_len + block_size - 1) // block_size
            cache_shape_k = (num_blocks, num_kv_heads, block_size, head_size)
            cache_shape_v = (num_blocks, num_kv_heads, block_size, head_size_v)
            cache_shape_k_int4 = (num_blocks, num_kv_heads, block_size, int(head_size/2))
            cache_shape_v_int4 = (num_blocks, num_kv_heads, int(block_size/2), head_size_v)
            cache_shape_v_int4_tmp = (num_blocks, num_kv_heads, head_size_v, int(block_size/2))

            if quant_mode == "per_channel":
                scale_shape_k = (num_kv_heads, head_size)
                scale_shape_v = (num_kv_heads, head_size_v)
            elif quant_mode == "per_token":
                scale_shape_k = (num_blocks, num_kv_heads, block_size)
                scale_shape_v = (num_blocks, num_kv_heads, block_size)
            elif quant_mode == "per_token_group":
                scale_shape_k = (num_blocks, num_kv_heads, block_size, 1) #group_size = head_size
                scale_shape_v = (num_blocks, num_kv_heads, block_size, 1)
            block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
            block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)

            key_cache = torch.zeros(cache_shape_k).uniform_(int_min, int_max).to(kv_data_type).mlu()
            value_cache = torch.zeros(cache_shape_v).uniform_(int_min, int_max).to(kv_data_type).mlu()
            key_cache_scale = torch.randn(size=scale_shape_k, dtype=torch.float32).mlu()
            value_cache_scale = torch.randn(size=scale_shape_v, dtype=torch.float32).mlu()
            key_cache_view = key_cache.reshape(-1, head_size)
            value_cache_view = value_cache.transpose(2, 3).reshape(-1, block_size)
            key_cache_int4 = PairlyPackInt8(key_cache_view).view(cache_shape_k_int4)
            value_cache_int4 = PairlyPackInt8(value_cache_view).view(cache_shape_v_int4_tmp).transpose(2,3).contiguous()

            alibi_slopes = None
            if has_alibi:
                alibi_slopes = torch.zeros((batch, head_num), dtype=torch.float32).mlu()
                alibi_slopes.uniform_(0, 0.125)
            softmax_scale = 1 / math.sqrt(head_size)

            torch_output = self.op_impl_base(input_q.contiguous(),
                                          key_cache, value_cache,
                                          None, block_tables, context_lens,
                                          key_cache_scale, value_cache_scale,
                                          alibi_slopes, max_context_len,
                                          window_size_left, window_size_right, softmax_scale, False, 4)
            tmo_output_contigous = ops.single_query_cached_kv_attn(input_q.contiguous(),
                                          key_cache_int4, value_cache_int4,
                                          None, block_tables, context_lens,
                                          key_cache_scale, value_cache_scale,
                                          alibi_slopes, max_context_len,
                                          window_size_left, window_size_right, softmax_scale, False, 4)
            self.assertTensorsEqual(torch_output.cpu().float(), tmo_output_contigous.cpu().float(),
                                    0.003, use_MSE=True)
            tmo_output_inplace = torch.empty((batch, seq_q, head_num, head_size_v), dtype=data_type, device="mlu")
            ops.single_query_cached_kv_attn(input_q, key_cache_int4, value_cache_int4,
                                        tmo_output_inplace, block_tables, context_lens,
                                        key_cache_scale, value_cache_scale,
                                        alibi_slopes, max_context_len,
                                        window_size_left, window_size_right, softmax_scale, False, 4)
            self.assertTensorsEqual(torch_output.cpu().float(), tmo_output_inplace.cpu().float(),
                                    0.003, use_MSE=True)
            if seq_q == 1:
                torch_output, torch_lse = self.op_impl_base(input_q,
                                            key_cache, value_cache,
                                            None, block_tables, context_lens,
                                            key_cache_scale, value_cache_scale,
                                            alibi_slopes, max_context_len,
                                            window_size_left, window_size_right, softmax_scale, True, 4)
                tmo_output1, tmo_lse = ops.single_query_cached_kv_attn(input_q,
                                            key_cache_int4, value_cache_int4,
                                            None, block_tables, context_lens,
                                            key_cache_scale, value_cache_scale,
                                            alibi_slopes, max_context_len,
                                            window_size_left, window_size_right, softmax_scale, True, 4)
                self.assertTensorsEqual(torch_lse.cpu().float(), tmo_lse.cpu().float(),
                                        0.003, use_MSE=True)
                self.assertTensorsEqual(torch_output.cpu().float(), tmo_output1.cpu().float(),
                                        0.003, use_MSE=True)

    def test_single_query_attention_inplace(self):
        head_num = 16
        batch_list = [5]
        num_kv_heads = 4
        head_size_list = [64]
        seq_len_list = [512]
        is_pagedattn_list = [False, True]
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name:
           print("pagedattn is not implement on mlu370, skip it")
           is_pagedattn_list = [False]
        has_alibi_list = [True, False]
        kv_data_type_list = [torch.int8, torch.float16]
        seq_q_list = [1, 5]
        window_size_list = [(-1, -1), (20, -1)]
        args = product(batch_list, head_size_list, is_pagedattn_list, has_alibi_list, kv_data_type_list, seq_len_list, seq_q_list, window_size_list)
        for batch, head_size, is_pagedattn, has_alibi, kv_data_type, max_seqlen, seq_q, (window_size_left, window_size_right) in args:
            print("batch: {}, max_seqlen: {}, head_size: {}, is_pagedattn: {}, has_alibi {}, kv_datatype {}, seq_q {}, window_size_left {}, window_size_right {}, testing...".format(
              batch, max_seqlen,  head_size, is_pagedattn, has_alibi, kv_data_type, seq_q, window_size_left, window_size_right))
            # prepare input
            params = gen_params(batch, head_size, head_size, is_pagedattn, has_alibi, torch.float16, kv_data_type, max_seqlen, seq_q, head_num, num_kv_heads)
            input_q, key_cache, value_cache, block_tables, context_lens, key_cache_scale, value_cache_scale, alibi_slopes, max_context_len = params
            softmax_scale = 1 / math.sqrt(head_size)
            tmo_output = torch.empty_like(input_q)
            tmo_output_contigous = torch.empty_like(input_q)
            torch_output = self.op_impl_base(input_q, key_cache, value_cache, None,
                                          block_tables, context_lens, key_cache_scale,
                                          value_cache_scale, alibi_slopes, max_context_len,
                                          window_size_left, window_size_right, softmax_scale, False, -1)
            ops.single_query_cached_kv_attn(input_q, key_cache, value_cache, tmo_output,
                                          block_tables, context_lens, key_cache_scale,
                                          value_cache_scale, alibi_slopes, max_context_len,
                                          window_size_left, window_size_right, softmax_scale)
            ops.single_query_cached_kv_attn(input_q.contiguous(), key_cache, value_cache,
                                          tmo_output_contigous, block_tables, context_lens,
                                          key_cache_scale, value_cache_scale, alibi_slopes,
                                          max_context_len, window_size_left, window_size_right, softmax_scale)
            self.assertTensorsEqual(tmo_output.cpu().float(), tmo_output_contigous.cpu().float(),
                                    0.000, use_MSE=True)
            self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                    0.003, use_MSE=True)

    # 防呆测试
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
        func = ops.single_query_cached_kv_attn
        batch, seq_q, head_num, num_kv_heads, head_size_qk, head_size_v, max_seqlen, softmax_scale = 5, 1, 8, 3, 64, 128, 512, 0.625
        dtype = torch.float16
        input = torch.randn((batch, seq_q, head_num, head_size_qk), dtype = dtype).mlu()
        context_lens = torch.randint(seq_q, max_seqlen + 1, (batch, ), dtype=torch.int32).mlu()
        max_context_len = int(max(context_lens))
        block_size = 16
        num_blocks = batch * ((max_seqlen + block_size - 1) // block_size)
        max_num_blocks_per_seq = (max_context_len + block_size - 1) // block_size
        cache_shape_k = (num_blocks, num_kv_heads, block_size, head_size_qk)
        cache_shape_v = (num_blocks, num_kv_heads, block_size, head_size_v)
        block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
        block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
        key_cache = torch.randn(size=cache_shape_k, dtype=dtype).mlu()
        value_cache = torch.randn(size=cache_shape_v, dtype=dtype).mlu()
        key_cache_scale = None
        value_cache_scale = None
        window_size_left, window_size_right = 10, 10
        self.assertException("only support windows_size_right < 0 currently.",
                                func, input, key_cache, value_cache,
                                None, block_tables, context_lens,
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, True, -1)
        window_size_right = -1
        self.assertException("num_heads need be mutiple of num_kv_heads.",
                                func, input, key_cache, value_cache,
                                None, block_tables, context_lens,
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, True, -1)
        num_kv_heads = 1
        cache_shape_k = (num_blocks, num_kv_heads, block_size, head_size_qk)
        cache_shape_v = (num_blocks, num_kv_heads, block_size, head_size_v)
        key_cache = torch.randn(size=cache_shape_k, dtype=dtype).mlu()
        value_cache = torch.randn(size=cache_shape_v, dtype=dtype).mlu()
        self.assertException("illegal quant bit size, only support 4, 8 or -1.",
                                func, input, key_cache, value_cache,
                                None, block_tables, context_lens,
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, True, 16)
        input1 = torch.randn((batch, seq_q, head_num, head_size_qk * 2), dtype = dtype).mlu()
        input = input1[..., 0::2]
        self.assertException("q last two dim need be contiguous.",
                                func, input, key_cache, value_cache,
                                None, block_tables, context_lens,
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, True, -1)
        input = torch.randn((batch, seq_q, head_num, head_size_qk), dtype = dtype).mlu()
        key_cache1 = key_cache[..., :8, :]
        value_cache1 = value_cache[..., :8, :]
        self.assertException("k_cache and v_cache need be contiguous.",
                                func, input, key_cache1, value_cache1,
                                None, block_tables, context_lens,
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, True, -1)
        self.assertException("q_ori need be mlu tensor.",
                                func, input.cpu(), key_cache, value_cache,
                                None, block_tables, context_lens,
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, True, -1)
        input = torch.randn((batch, 5, head_num, head_size_qk), dtype = dtype).mlu()
        self.assertException("return lse only support seq_q = 1 currently.",
                                func, input, key_cache, value_cache,
                                None, block_tables, context_lens,
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, True, -1)
        self.assertException("block_tables type need be torch::kInt32 or torch::kLong.",
                                func, input, key_cache, value_cache,
                                None, block_tables.float(), context_lens,
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, False, -1)
        self.assertException("context_lens type need be torch::kInt32.",
                                func, input, key_cache, value_cache,
                                None, block_tables, context_lens.to(torch.int64),
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, False, -1)
        self.assertException("context_lens need be contiguous.",
                                func, input, key_cache, value_cache,
                                None, block_tables, context_lens[0::2],
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, False, -1)
        scale_shape = (num_blocks)
        key_cache = torch.zeros(cache_shape_k).uniform_(-128, 127).to(torch.int8).mlu()
        value_cache = torch.zeros(cache_shape_v).uniform_(-128, 127).to(torch.int8).mlu()
        key_cache_scale = torch.randn(scale_shape, dtype=torch.float32).mlu()
        value_cache_scale = torch.randn(scale_shape, dtype=torch.float32).mlu()
        self.assertException("k_cache_quant_scale must be 2d or 3d or 4d.",
                                func, input, key_cache, value_cache,
                                None, block_tables, context_lens,
                                key_cache_scale, value_cache_scale,
                                None, max_context_len,
                                window_size_left, window_size_right, softmax_scale, True, -1)

    def test_inductor(self):
        batch, seq_q, head_num, num_kv_heads, head_size, max_seqlen = 1, 5, 16, 16, 128, 512
        is_pagedattn_list = [False, True] if "MLU3" not in torch.mlu.get_device_name() else [False]
        has_alibi_list = [True, False]
        test_flags = product(is_pagedattn_list, has_alibi_list)
        for is_pagedattn, has_alibi in test_flags:
            args = gen_args(batch, head_size, is_pagedattn, has_alibi, torch.int8, max_seqlen, seq_q, head_num, num_kv_heads)
            self.base_opcheck(torch.ops.torch_mlu_ops.single_query_cached_kv_attn, args)

if __name__ == '__main__':
    exit(run_unittest(TestSingleQueryAttnOp))
