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

class TestSessionCacheAttnOp(BtTestCase):
    def op_impl_base(self, *args):
        def dequant_from_cache(key_cache, value_cache, key_cache_scale, value_cache_scale, cache_bs_id, cache_lens,
                               cache_seq_offset,  quant_bit, quant_layout, max_cache_len):
            batch = cache_bs_id.size(0)
            head_num = key_cache.size(1)
            head_size = value_cache.size(-1)
            cache_len = key_cache.size(-2)
            cache_shape = (batch, max_cache_len, head_num, head_size)
            key_cache_mem = torch.zeros(cache_shape, dtype=torch.float, device='mlu')
            value_cache_mem = torch.zeros_like(key_cache_mem)
            for i in range(batch):
                batch_id = cache_bs_id[i]
                cache_offset = 0 if cache_seq_offset is None else cache_seq_offset[i]
                key_cache_data = key_cache[batch_id] # head_num_kv, cache_mem_len, head_size
                value_cache_data = value_cache[batch_id]
                if quant_bit == 4:
                    key_quant_data_int8 = torch.zeros(head_num, cache_len, head_size, dtype=torch.int8, device="mlu")
                    value_quant_data_int8 = torch.zeros(head_num, cache_len, head_size, dtype=torch.int8, device="mlu")
                    key_quant_data_int8[...,::2] = key_cache_data << 4 >> 4
                    key_quant_data_int8[...,1::2] = key_cache_data >> 4
                    key_quant_data_fp32 = key_quant_data_int8.clone().to(torch.float)
                    value_quant_data_int8[:,0::2,:] = value_cache_data << 4 >> 4
                    value_quant_data_int8[:,1::2,:] = value_cache_data >> 4
                    value_quant_data_fp32 = value_quant_data_int8.clone().to(torch.float)
                else:
                    key_quant_data_fp32 = key_cache_data.clone().to(torch.float)
                    value_quant_data_fp32 = value_cache_data.clone().to(torch.float)

                if quant_layout == 'per_channel':
                    key_quant_data_fp32 = key_quant_data_fp32 * key_cache_scale[..., None, :]
                    value_quant_data_fp32 = value_quant_data_fp32 * value_cache_scale[..., None, :]
                else:  # per token
                    key_cache_scale_data = key_cache_scale[batch_id]
                    value_cache_scale_data = value_cache_scale[batch_id]
                    key_quant_data_fp32 = key_quant_data_fp32 * key_cache_scale_data[..., None]
                    value_quant_data_fp32 = value_quant_data_fp32 * value_cache_scale_data[..., None]
                key_quant_data_fp32 = key_quant_data_fp32[:, cache_offset:cache_offset+max_cache_len, :] #cut vaild cache
                value_quant_data_fp32 = value_quant_data_fp32[:, cache_offset:cache_offset+max_cache_len, :]
                # head_num_kv, max_cache_len, headsize-> max_cache_len, head_num_kv, headsize
                key_cache_mem[i] =  key_quant_data_fp32.transpose(0,1)
                value_cache_mem[i] = value_quant_data_fp32.transpose(0,1)
            return  key_cache_mem, value_cache_mem
        def scale_dot_attn(q_sess, key_cache, value_cache, cu_seq_lens_q, cu_seq_lens_cache, alibi_slope, is_causal, softmax_scale):
            # cache_shape: [batch,  cache_len, head_kv,head_size]
            batch = cu_seq_lens_q.size(0) - 1
            head_num_q = q_sess.size(-2)
            head_num_kv = key_cache.size(-2)
            assert head_num_q >= head_num_kv and head_num_q % head_num_kv == 0
            group = head_num_q // head_num_kv
            inf = 1e6
            device= 'mlu'
            out_list = []
            for i in range(batch):
                q = q_sess[cu_seq_lens_q[i]:cu_seq_lens_q[i+1], ...]
                k = key_cache[cu_seq_lens_cache[i]:cu_seq_lens_cache[i+1], ...]  # [cache_len, head_num_kv, head_size]
                v = value_cache[cu_seq_lens_cache[i]:cu_seq_lens_cache[i+1], ...]
                k = torch.repeat_interleave(k, group, dim=1) #[cache_len, head_num_q, head_size]
                v = torch.repeat_interleave(v, group, dim=1)
                qk = torch.einsum('qhd,khd->hqk', q, k) * softmax_scale
                seq_q, seq_k = q.size(0), k.size(0)
                if alibi_slope is not None:
                    slope = alibi_slope.reshape(1, head_num_q, 1, 1)
                    slope_bias = torch.zeros(1, head_num_q, seq_q, seq_k).to(device=device)
                    if is_causal:
                        relative_pos = torch.arange(-seq_k + 1, 1, dtype=torch.float32).to(device=device)
                        slope_bias = relative_pos * slope
                    else:
                        row_idx = torch.arange(seq_q, dtype=torch.long).reshape(-1, 1)
                        col_idx = torch.arange(seq_k, dtype=torch.long)
                        relative_pos = torch.abs(row_idx + seq_k - seq_q - col_idx).to(device=device)
                        slope_bias = -slope * relative_pos.to(dtype=slope.dtype)
                    qk += (slope_bias.squeeze(0))
                if is_causal:
                    assert seq_q <= seq_k, "seq_q <= seq_k if causal=True"
                    zeros = torch.zeros(seq_q, seq_k-seq_q, dtype=torch.float, device="mlu")
                    tri = torch.full((seq_q, seq_q), -inf, dtype=torch.float, device="mlu").triu(diagonal=1)
                    mask = torch.cat([zeros, tri], dim=1) # (q, k-q) + (q, q) => (q, k)
                    qk += mask
                attn = torch.softmax(qk, dim=-1, dtype=torch.float).to(q.dtype)
                qkv = torch.einsum('hqk,khd->qhd', attn, v)
                out_list.append(qkv)
            output = torch.cat(out_list, dim=0)
            return output

        q_sess, k_sess, v_sess, key_cache1, value_cache1, key_cache_scale1, value_cache_scale1, cache_lens1,\
                                cache_seq_offset1, quant_bit1, quant_layout1, max_cache_len1,\
                                key_cache2, value_cache2, key_cache_scale2, value_cache_scale2, cache_lens2,\
                                cache_seq_offset2, quant_bit2, quant_layout2, max_cache_len2,\
                                sess_lens, cache_bs_id, cu_seq_lens_q, is_causal, softmax_scale = args
        #1. dequant cache
        #input cache shape [max_batch, head_num_kv, cache_mem_len, head_size]
        #output cache shape [batch, max_cache_len, head_num_kv, head_size]
        key_cache_mem1, value_cache_mem1 = dequant_from_cache(key_cache1, value_cache1, key_cache_scale1,
                                        value_cache_scale1, cache_bs_id, cache_lens1, cache_seq_offset1,
                                        quant_bit1, quant_layout1, max_cache_len1)
        key_cache_mem2, value_cache_mem2 = None, None
        if key_cache2 is not None:
            key_cache_mem2, value_cache_mem2 = dequant_from_cache(key_cache2, value_cache2, key_cache_scale2,
                                        value_cache_scale2, cache_bs_id, cache_lens2, cache_seq_offset2,
                                        quant_bit2, quant_layout2, max_cache_len2)
        #2. concat cache
        batch = cache_bs_id.size(0)
        head_num_kv = k_sess.size(-2)
        head_size = k_sess.size(-1)
        # concat cache1 and cache2
        if key_cache2 is not None:
            cache_lens = cache_lens1 + cache_lens2
            max_cache_len = max_cache_len1 + max_cache_len2
            key_cache_mem = torch.zeros(batch, max_cache_len, head_num_kv, head_size, dtype=torch.float, device='mlu')
            value_cache_mem = torch.zeros(batch, max_cache_len, head_num_kv, head_size, dtype=torch.float, device='mlu')
            for i in range(batch): #concat
                len1 = cache_lens1[i]
                len2 = cache_lens2[i]
                key_cache_mem[i, :len1, ...] = key_cache_mem1[i, :len1, ...]
                key_cache_mem[i, len1:len1+len2, ...] = key_cache_mem2[i, :len2, ...]
                value_cache_mem[i, :len1, ...] = value_cache_mem1[i, :len1, ...]
                value_cache_mem[i, len1:len1+len2, ...] = value_cache_mem2[i, :len2, ...]
        else:
            key_cache_mem, value_cache_mem = key_cache_mem1, value_cache_mem1
            cache_lens = cache_lens1
        #concat mem cache and sess cache
        concat_cache_lens = cache_lens + sess_lens
        cu_seq_lens_cache = torch.zeros((batch+1), dtype=torch.int32)
        cu_seq_lens_cache[1:] = torch.cumsum(concat_cache_lens, dim=-1)
        total_cache_len = cu_seq_lens_cache[-1]
        concate_key_cache = torch.zeros((total_cache_len, head_num_kv, head_size), dtype=torch.float, device='mlu')
        concate_value_cache = torch.zeros((total_cache_len, head_num_kv, head_size), dtype=torch.float, device='mlu')
        for i in range(batch): #concat
            mem_len = cache_lens[i]
            sess_len = sess_lens[i]
            off = cu_seq_lens_cache[i]
            concate_key_cache[off:off+mem_len, ...] = key_cache_mem[i, :mem_len, ...]
            concate_key_cache[off+mem_len:off+mem_len+sess_len, ...] = k_sess[i, :sess_len, ...]
            concate_value_cache[off:off+mem_len, ...] = value_cache_mem[i, :mem_len, ...]
            concate_value_cache[off+mem_len:off+mem_len+sess_len, ...] = v_sess[i, :sess_len, ...]
        #3. attn
        attn_out = scale_dot_attn(q_sess, concate_key_cache, concate_value_cache, cu_seq_lens_q, cu_seq_lens_cache, None, is_causal, softmax_scale)
        return attn_out, concate_key_cache, concate_value_cache

    def test_session_cache_attn_quant_kv(self):
        max_batch = 64
        cache_mem_len = 4096
        head_num_kv = 1
        head_num_q = 8
        head_size = 128
        batch = 32
        max_sess_len = 100
        max_cache_len = 3072
        dtype_list=[torch.float16]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        quant_bit_list = [4, 8]
        quant_layout_list = ['per_token', 'per_channel']
        is_causal_list = [True, False]
        arg = product(is_causal_list, dtype_list, quant_bit_list, quant_layout_list)
        softmax_scale = 1 / math.sqrt(head_size)
        for is_causal, dtype, quant_bit, quant_layout in arg:
            print(f"is_causal: {is_causal}, dtype: {dtype}, quant_bit: {quant_bit}, quant_layout: {quant_layout}")
            if quant_bit == 8:
                cache_shape_key = (max_batch, head_num_kv, cache_mem_len, head_size)
                cache_shape_value = cache_shape_key
            else:
                cache_shape_key = (max_batch, head_num_kv, cache_mem_len, head_size//2)
                cache_shape_value = (max_batch, head_num_kv, cache_mem_len//2, head_size)
            if quant_layout == "per_token":
                scale_shape = (max_batch, head_num_kv, cache_mem_len)
                quant_mode = 1
            else: # per channel
                scale_shape = (head_num_kv, head_size)
                quant_mode = 0
            key_cache = torch.randint(-127, 128, cache_shape_key, device='mlu').to(torch.int8)
            value_cache = torch.randint(-127, 128, cache_shape_value, device='mlu').to(torch.int8)
            key_cache_scale = torch.randn(scale_shape, device='mlu', dtype=torch.float)
            value_cache_scale = torch.randn(scale_shape, device='mlu', dtype=torch.float)
            key_cache_scale =  torch.fill(key_cache_scale, 0.01)
            value_cache_scale = torch.fill(value_cache_scale, 0.01)
            cache_lens = torch.randint(1, max_cache_len + 1, (batch,), dtype=torch.int32, device='mlu')
            sess_lens = torch.randint(1, max_sess_len + 1, (batch,), dtype=torch.int32, device='mlu')
            context_lens = cache_lens + sess_lens
            max_cache_len_new = torch.max(context_lens)
            cu_seq_lens_q = torch.zeros(batch+1, dtype=torch.int32)
            cu_seq_lens_q[1:] = torch.cumsum(sess_lens, dim=-1)
            total_sess_len = cu_seq_lens_q[-1]
            cu_seq_lens_q=cu_seq_lens_q.mlu()
            cache_bs_id = random.sample([*range(0, max_batch)], batch)
            cache_bs_id = torch.IntTensor(cache_bs_id).mlu() #block_tables
            mem_cache_seq_offset = torch.randint(0, cache_mem_len-max_cache_len, (batch,),  dtype=torch.int32, device='mlu')

            q_sess = torch.randn(total_sess_len, head_num_q, head_size, dtype=dtype, device='mlu')
            k_sess = torch.randn(batch, max_sess_len, head_num_kv, head_size, dtype=dtype, device='mlu')
            v_sess = torch.randn(batch, max_sess_len, head_num_kv, head_size, dtype=dtype, device='mlu')
            #fill sess cache
            cache_seq_offset = torch.zeros((batch + 1), dtype=torch.int32)
            cache_seq_offset[1:]= torch.cumsum(context_lens, dim=-1)
            cache_seq_offset = cache_seq_offset.mlu()
            context_seq_offset = cache_seq_offset[:-1]
            sess_seq_offset = context_seq_offset + cache_lens
            total_mem_len = cache_seq_offset[-1]
            cache_shape_mem = (total_mem_len, head_num_kv, head_size)#NT,H,C
            key_cache_mem = torch.zeros(cache_shape_mem, dtype= dtype, device='mlu')
            value_cache_mem = torch.zeros(cache_shape_mem, dtype= dtype, device='mlu')
            for i in range(batch):
                offset = sess_seq_offset[i]
                len = sess_lens[i]
                key_cache_mem[offset:offset+len, ...] = k_sess[i, :len, ...]
                value_cache_mem[offset:offset+len, ...] = v_sess[i, :len, ...]

            # baseline
            base_output, base_key_cache, base_value_cache = self.op_impl_base(q_sess.float(), k_sess.float(), v_sess.float(),
                                key_cache, value_cache, key_cache_scale, value_cache_scale, cache_lens, mem_cache_seq_offset,
                                quant_bit, quant_layout, max_cache_len,
                                None, None, None, None, None, None, -1, None, -1,
                                sess_lens, cache_bs_id,  cu_seq_lens_q, is_causal, softmax_scale)

            #tmo
            #1. dequant cache
            ops.dequant_from_linear_cache(key_cache_mem, value_cache_mem, key_cache, value_cache, key_cache_scale,
                                          value_cache_scale, cache_lens, max_cache_len, context_seq_offset,
                                          cache_bs_id, mem_cache_seq_offset, quant_mode, quant_bit)
            self.assertTensorsEqual(base_key_cache.cpu().float(), key_cache_mem.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(base_value_cache.cpu().float(), value_cache_mem.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)
            #2. flash_attn
            tmo_output = ops.flash_attention(q_sess, key_cache_mem, value_cache_mem, None, cu_seq_lens_q, cache_seq_offset,
                                None, None, max_sess_len, max_cache_len_new, softmax_scale,
                                is_causal, -1, -1, torch.float, False, None, None, None)
            self.assertTensorsEqual(base_output.cpu().float(), tmo_output.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)

    def test_session_cache_attn_kv_mixquant(self):
        max_batch = 64
        cache_mem_len_int4 = 4064
        cache_mem_len_int8 = 32
        head_num_kv = 1
        head_num_q = 8
        head_size = 128
        batch = 32
        max_sess_len = 100
        max_cache_len_int4 = 3072
        max_cache_len_int8 = 32
        dtype_list=[torch.float16]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        is_causal_list = [True, False]
        arg = product(is_causal_list, dtype_list)
        softmax_scale = 1 / math.sqrt(head_size)
        for is_causal, dtype in arg:
            print(f"is_causal: {is_causal}, dtype: {dtype}")
            cache_shape_int8 = (max_batch, head_num_kv, cache_mem_len_int8, head_size)
            cache_shape_int4_key = (max_batch, head_num_kv, cache_mem_len_int4, head_size//2)
            cache_shape_int4_value = (max_batch, head_num_kv, cache_mem_len_int4//2, head_size)
            scale_shape_int4 = (max_batch, head_num_kv, cache_mem_len_int4) #per_token
            quant_mode_int4 = 1
            scale_shape_int8 = (head_num_kv, head_size) # per channel
            quant_mode_int8 = 0
            key_cache_int8 = torch.randint(-127, 128, cache_shape_int8, device='mlu').to(torch.int8)
            value_cache_int8 = torch.randint(-127, 128, cache_shape_int8, device='mlu').to(torch.int8)
            key_cache_int4 = torch.randint(-127, 128, cache_shape_int4_key, device='mlu').to(torch.int8)
            value_cache_int4 = torch.randint(-127, 128, cache_shape_int4_value, device='mlu').to(torch.int8)
            key_cache_scale_int4 = torch.randn(scale_shape_int4, device='mlu', dtype=torch.float)
            value_cache_scale_int4 = torch.randn(scale_shape_int4, device='mlu', dtype=torch.float)
            key_cache_scale_int8 = torch.randn(scale_shape_int8, device='mlu', dtype=torch.float)
            value_cache_scale_int8 = torch.randn(scale_shape_int8, device='mlu', dtype=torch.float)
            key_cache_scale_int4 =  torch.fill(key_cache_scale_int4, 0.01)
            value_cache_scale_int4 = torch.fill(value_cache_scale_int4, 0.01)
            key_cache_scale_int8 =  torch.fill(key_cache_scale_int8, 0.01)
            value_cache_scale_int8 = torch.fill(value_cache_scale_int8, 0.01)
            cache_lens_int4 = torch.randint(1, max_cache_len_int4 + 1, (batch,), dtype=torch.int32, device='mlu')
            cache_lens_int8 = torch.randint(1, max_cache_len_int8 + 1, (batch,), dtype=torch.int32, device='mlu')
            sess_lens = torch.randint(max_sess_len-1, max_sess_len, (batch,), dtype=torch.int32, device='mlu')
            cu_seq_lens_q = torch.zeros(batch+1, dtype=torch.int32)
            cu_seq_lens_q[1:] = torch.cumsum(sess_lens, dim=-1)
            total_sess_len = cu_seq_lens_q[-1]
            cu_seq_lens_q=cu_seq_lens_q.mlu()
            context_lens = cache_lens_int4 + cache_lens_int8 + sess_lens
            max_cache_len_new = torch.max(context_lens)
            cache_bs_id = random.sample([*range(0, max_batch)], batch)
            cache_bs_id = torch.IntTensor(cache_bs_id).mlu() #block_tables
            cache_seq_offset_int4 = torch.randint(0, cache_mem_len_int4-max_cache_len_int4, (batch,),  dtype=torch.int32, device='mlu')
            cache_seq_offset_int8 = None

            q_sess = torch.randn(total_sess_len, head_num_q, head_size, dtype=dtype, device='mlu')
            k_sess = torch.randn(batch, max_sess_len, head_num_kv, head_size, dtype=dtype, device='mlu')
            v_sess = torch.randn(batch, max_sess_len, head_num_kv, head_size, dtype=dtype, device='mlu')
            #fill sess cache
            cache_seq_offset1 = torch.zeros((batch + 1), dtype=torch.int32)
            cache_seq_offset2 = torch.zeros((batch + 1), dtype=torch.int32)
            cache_seq_offset1[1:]= torch.cumsum(context_lens, dim=-1)
            cache_seq_offset2[:-1] = cache_seq_offset1[:-1] + cache_lens_int4.cpu()
            cache_seq_offset1 = cache_seq_offset1.mlu()
            cache_seq_offset2 = cache_seq_offset2.mlu()

            context_seq_offset1 = cache_seq_offset1[:-1]
            context_seq_offset2 = cache_seq_offset2[:-1]
            sess_seq_offset = context_seq_offset2 + cache_lens_int8
            total_mem_len = cache_seq_offset1[-1]
            cache_shape_mem = (total_mem_len, head_num_kv, head_size)#NT,H,C
            key_cache_mem = torch.zeros(cache_shape_mem, dtype= dtype, device='mlu')
            value_cache_mem = torch.zeros(cache_shape_mem, dtype= dtype, device='mlu')
            for i in range(batch):
                offset = sess_seq_offset[i]
                len = sess_lens[i]
                key_cache_mem[offset:offset+len, ...] = k_sess[i, :len, ...]
                value_cache_mem[offset:offset+len, ...] = v_sess[i, :len, ...]

            # baseline
            base_output, base_key_cache, base_value_cache = self.op_impl_base(q_sess.float(), k_sess.float(), v_sess.float(),
                                key_cache_int4, value_cache_int4, key_cache_scale_int4, value_cache_scale_int4, cache_lens_int4,
                                cache_seq_offset_int4, 4, 'per_token', max_cache_len_int4,
                                key_cache_int8, value_cache_int8, key_cache_scale_int8, value_cache_scale_int8, cache_lens_int8,
                                cache_seq_offset_int8, 8, 'per_channel', max_cache_len_int8,
                                sess_lens, cache_bs_id, cu_seq_lens_q, is_causal, softmax_scale)
            #tmo
            #1. dequant cache
            ops.dequant_from_linear_cache(key_cache_mem, value_cache_mem, key_cache_int4, value_cache_int4,
                                          key_cache_scale_int4, value_cache_scale_int4, cache_lens_int4, max_cache_len_int4,
                                          context_seq_offset1, cache_bs_id, cache_seq_offset_int4, quant_mode_int4, 4)
            ops.dequant_from_linear_cache(key_cache_mem, value_cache_mem, key_cache_int8, value_cache_int8,
                                          key_cache_scale_int8, value_cache_scale_int8, cache_lens_int8, max_cache_len_int8,
                                          context_seq_offset2, cache_bs_id, cache_seq_offset_int8, quant_mode_int8, 8)

            self.assertTensorsEqual(base_key_cache.cpu().float(), key_cache_mem.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(base_value_cache.cpu().float(), value_cache_mem.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)
            #2. flash_attn
            tmo_output = ops.flash_attention(q_sess, key_cache_mem, value_cache_mem, None, cu_seq_lens_q, cache_seq_offset1,
                                None, None, max_sess_len, max_cache_len_new, softmax_scale,
                                is_causal, -1, -1, torch.float, False, None, None, None)
            self.assertTensorsEqual(base_output.cpu().float(), tmo_output.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        return super().test_inductor()

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