# Copyright (c) 2024 Microsoft
# Licensed under The MIT License [see LICENSE for details]

import inspect
import json
import os
from importlib import import_module

from transformers.models.llama.modeling_llama import *
from transformers.utils.import_utils import _is_package_available

if _is_package_available("vllm"):
    from vllm.attention.backends.flash_attn import *

from ..ops.block_sparse_flash_attention import block_sparse_attention
from ..ops.pit_sparse_flash_attention_v2 import vertical_slash_sparse_attention
from ..ops.streaming_kernel import streaming_forward, streaming_forward2
from .snap_kv import *

last_q = 64
arange = torch.arange(last_q, device="cuda")
LAST_Q_MASK = arange[None, None, :, None] >= arange[None, None, None, :]
ROPE_TYPE = None
SEARCH_MASK = None

def init_minference_parameters(self):
    config = self.config.to_dict()
    self.starting_layer = config.get("starting_layer", 0)
    self.is_search = config.get("is_search", False)

    # self.n_init = config.get("n_init", 128)
    # self.n_local = config.get("n_local", 3968)

    self.ne_inf = None
    self.config_path = config.get("config_path", "")
    if os.path.exists(self.config_path) and self.layer_idx < len(json.load(open(self.config_path))):
        self.best_pattern = {int(ii): jj for ii, jj in json.load(open(self.config_path))[self.layer_idx].items()}
    else:
        self.best_pattern = {}
    self.vertical, self.slash = None, None

    # import apply_rotary_pos_emb
    if "apply_rotary_pos_emb" not in self.__dict__:
        global apply_rotary_pos_emb
        model_path = self.rotary_emb.__class__.__module__
        apply_rotary_pos_emb = getattr(import_module(model_path), "apply_rotary_pos_emb")
        self.apply_rotary_pos_emb = True

def sum_all_diagonal_matrix(mat: torch.tensor):
    b, h, n, m = mat.shape
    zero_mat = torch.zeros((b, h, n, n)).to(mat.device) # Zero matrix used for padding
    mat_padded =  torch.cat((zero_mat, mat, zero_mat), -1) # pads the matrix on left and right
    mat_strided = mat_padded.as_strided((1, 1, n, n + m), (1, n * (2 * n + m), 2 * n + m + 1, 1)) # Change the strides
    sum_diags = torch.sum(mat_strided, 2) # Sums the resulting matrix's columns
    return sum_diags[:,:,1:]

def gather(t, dim, i):
    """A broadcasting version of torch.gather."""
    dim += (dim < 0) * t.ndim
    return t.gather(dim, i.expand(*t.shape[:dim], i.shape[dim], *t.shape[dim + 1 :]))

def gather_qkv(q, k, v, attention_mask):
    attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(q.size(-1)) + attention_mask
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
    attn_output = torch.matmul(attn_weights, v)
    return attn_output

def search_pattern(q, k, head):
    q_len = q.shape[2]
    head_dim = q.shape[-1]

    def vertical_and_slash(vertical_size, slash_size):
        last_q = 64
        q_len = q.shape[2]
        qk_idxs = [ii + q_len for ii in list(range(-last_q, 0, 1))]
        qk = torch.matmul(q[:,:,qk_idxs,:], k.transpose(2, 3))/ math.sqrt(head_dim) + attention_mask[:,:,qk_idxs]
        qk = torch.nn.functional.softmax(qk, dim=-1, dtype=torch.float32)
        vertical = qk.sum(-2, keepdim=True)
        vertical[...,:30] = 10000
        vertical_topk = torch.topk(-vertical, q_len - vertical_size, -1).indices

        slash = sum_all_diagonal_matrix(qk)[...,:-last_q + 1]
        slash[...,-30:] = 10000
        slash_topk = slash
        slash = torch.topk(slash, slash_size, -1).indices - (q_len - 1)
        slash = torch.stack([torch.sparse.spdiags(torch.ones(slash_size, q_len), slash.cpu()[0][_], (q_len, q_len)).to_dense() for _ in range(1)]).to(q.device)

        est_attn = torch.ones_like(attn_weights)
        dim = 3
        est_attn = est_attn.scatter(3, vertical_topk.expand(*est_attn.shape[:dim], vertical_topk.shape[dim], *est_attn.shape[dim + 1 :]), 0)
        est_attn = est_attn + slash

        est_attn = (est_attn > 0).float()
        est_attn = torch.tril(est_attn)
        attn_weights_x = attn_weights * est_attn
        res3 = attn_weights_x[:,:,2500:].sum(-1).mean(-1).squeeze().float().detach().cpu().numpy()
        return res3

    def stream_llm(vertical_size, slash_size):
        q_len = q.shape[2]

        mask = torch.triu(torch.tril(torch.ones(q_len, q_len), 0), -slash_size).to(q)
        mask[:,:vertical_size] = 1
        mask = mask.unsqueeze(0).unsqueeze(1)

        est_attn = torch.tril(mask)
        attn_weights_x = attn_weights * est_attn
        res3 = attn_weights_x[:,:,2500:].sum(-1).mean(-1).squeeze().float().detach().cpu().numpy()
        return res3

    def block_sparse(topk_ratio, slash_size=None):
        block_num = (q_len -1) // 32 + 1
        block_q = torch.zeros(1,1,block_num * 32,head_dim).to(q)
        block_q[:,:,:q_len] = q
        block_q = block_q.reshape(1,1,block_num,32,-1).mean(-2)
        block_k = torch.zeros(1,1,block_num * 32,head_dim).to(k)
        block_k[:,:,:q_len] = k
        block_k = block_k.reshape(1,1,block_num,32,-1).mean(-2)

        qk = torch.matmul(block_q, block_k.transpose(2, 3)) + attention_mask[:,:,:block_num,:block_num]
        est_attn = torch.ones_like(qk)
        block_topk = torch.topk(-qk, block_num - block_num//topk_ratio, -1).indices

        dim = 3
        est_attn = est_attn.scatter(3, block_topk.expand(*est_attn.shape[:dim], block_topk.shape[dim], *est_attn.shape[dim + 1 :]), 0)
        est_attn = est_attn.unsqueeze(3).unsqueeze(-1).repeat(1,1,1,32,1,32).reshape(1,1,block_num * 32, block_num * 32)[...,:q_len,:q_len]
        est_attn = torch.tril(est_attn)

        attn_weights_x = attn_weights * est_attn
        res2 = attn_weights_x[:,:,2500:].sum(-1).mean(-1).squeeze().float().detach().cpu().numpy()
        return res2

    global SEARCH_MASK
    if SEARCH_MASK is None:
        attention_mask = torch.full((q_len, q_len), torch.finfo(q.dtype).min, device="cuda")
        mask_cond = torch.arange(attention_mask.size(-1), device="cuda")
        attention_mask.masked_fill_(mask_cond < (mask_cond + 1).view(attention_mask.size(-1), 1), 0)
        attention_mask = attention_mask[None, None, :]
        SEARCH_MASK = attention_mask
    else:
        attention_mask = SEARCH_MASK
    attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(head_dim) + attention_mask
    attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
    best_s, best_v, best_score, best_ty = 0, 0, 0, ""
    all_info = []
    for ty, fc in [("stream_llm", stream_llm), ("vertical_and_slash", vertical_and_slash), ("block_sparse", block_sparse)]:
        if ty == "stream_llm":
            vs_list = [(100, 800)]
        elif ty == "vertical_and_slash":
            vs_list = [(30, 800), (100, 750), (500, 700), (3500, 100)]
        else:
            vs_list = [(8, 1)]
        for v_size, s_size in vs_list:
            score = fc(v_size, s_size)
            score = score.item()
            all_info.append([ty, v_size, s_size, score])
            if score > best_score:
                best_score = score
                best_s, best_v = s_size, v_size
                best_ty = ty
    if best_ty == "stream_llm":
        best_ty = "vertical_and_slash"
    if best_ty == "block_sparse":
        best_ty, best_v, best_s = "vertical_and_slash", 1000, 6096
    print(head, best_ty, best_v, best_s, best_score)
    return (best_ty, best_v, best_s, best_score)

def search_pattern_v2(q, k, v, head):
    q_len = q.shape[2]
    head_dim = q.shape[-1]
    def vertical_and_slash_kernel(q, k, v, vertical_size, slash_size):
        vertical_size, slash_size  = min(q_len, max(vertical_size, 30)), min(q_len, max(slash_size, 50))
        last_q = 64
        qk = torch.einsum(f'bhmk, bhnk -> bhmn', q[:,:,-last_q:,:], k)
        qk[:, :, :, -last_q:] = torch.where(LAST_Q_MASK, qk[:, :, :, -last_q:], -torch.inf)
        qk = torch.nn.functional.softmax(qk, dim=-1, dtype=torch.float32)
        vertical = qk.sum(-2, keepdim=True)
        vertical[...,:30] = torch.inf
        vertical_topk = torch.topk(vertical, vertical_size, -1).indices

        slash = sum_all_diagonal_matrix(qk)[...,:-last_q + 1]
        slash[...,-30:] = torch.inf
        slash_topk = slash
        slash = (q_len - 1) - torch.topk(slash, slash_size, -1).indices

        return vertical_slash_sparse_attention(q, k, v, vertical_topk, slash)
    def dense(q, k, v, vertical_size=None, slash_size=None):
        return flash_attn_func(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1,2), 0.0, softmax_scale=None, causal=q_len != 1).view(bsz, 1, q_len, head_dim)
    def block_sparse_kernel(q, k, v, vertical_size=None, slash_size=None):
        topk = 100
        return block_sparse_attention(q, k, v, topk)

    best_s, best_v, best_score, best_ty = 0, 0, float("inf"), ""
    bsz = q.shape[0]
    all_info = []
    ref = dense(q, k, v)
    for ty, fc in [("stream_llm", streaming_forward), ("vertical_and_slash", vertical_and_slash_kernel), ("block_sparse", block_sparse_kernel)]:
        if ty == "stream_llm":
            vs_list = [(100, 800)]
        elif ty == "vertical_and_slash":
            vs_list = [(30, 800), (100, 800), (100, 750), (500, 700), (3500, 100), (1000, 4096)]
        else:
            vs_list = [(10, 1)]
        for v_size, s_size in vs_list:
            score = fc(q, k, v, v_size, s_size)
            # delta = (ref - score).abs().sum()
            delta = ((ref - score).abs() > 5e-3).sum()
            score = delta.item()
            all_info.append([ty, v_size, s_size, score])
            if score < best_score:
                best_score = score
                best_s, best_v = s_size, v_size
                best_ty = ty
    print(head, best_ty, best_v, best_s, best_score)
    return all_info

def shift_matrix(mat):
    b, h, _, n = mat.shape
    zero_mat = torch.zeros((b, h, n, n)).to(mat.device) # Zero matrix used for padding
    mat_padded =  torch.cat((zero_mat, mat, zero_mat), -1) # pads the matrix on left and right
    mat_strided = mat_padded.as_strided((1, 1, n, n + 2 * n), (1, n * (2 * n + n), 2 * n + n - 1, 1)) # Change the strides
    return mat_strided[...,2 * n-1:-1]

def repeat(self, q, k, v, attention_mask):
    q_len = q.shape[2]
    if q_len == 1:
        return gather_qkv(q, k, v, attention_mask)
    qk = torch.matmul(q[:,:,-1:,:], k.transpose(2, 3)) / math.sqrt(self.head_dim)
    qk = qk.repeat(1,1,q_len, 1)
    qk = shift_matrix(qk) + attention_mask
    attn_weights = nn.functional.softmax(qk, dim=-1, dtype=torch.float32).to(q.dtype)
    attn_output = torch.matmul(attn_weights, v)
    return attn_output

def gather_last_q_vertical_slash_topk_v4(self, q, k, v, head_id):
    kv_seq_len = k.size(2)

    def vertical_and_slash(attn_weights, vertical_size, slash_size):
        last_q = 64
        q_len = q.shape[2]
        vertical_size, slash_size  = min(q_len, max(vertical_size, 30)), min(q_len, max(slash_size, 50))
        qk_idxs = [ii + q_len for ii in list(range(-last_q, 0, 1))]
        qk = torch.matmul(q[:,:,qk_idxs,:], k.transpose(2, 3))/ math.sqrt(self.head_dim) + attention_mask[:,:,qk_idxs]
        qk = torch.nn.functional.softmax(qk, dim=-1, dtype=torch.float32)
        vertical = qk.sum(-2, keepdim=True)
        vertical[...,:30] = -self.ne_inf
        vertical_topk = torch.topk(-vertical, q_len - vertical_size, -1).indices

        slash = sum_all_diagonal_matrix(qk)[...,:-last_q + 1]
        slash[...,-30:] = -self.ne_inf
        slash_topk = slash
        slash = torch.topk(slash, slash_size, -1).indices - (q_len - 1)
        slash = torch.stack([torch.sparse.spdiags(torch.ones(slash_size, q_len), slash.cpu()[0][_], (q_len, q_len)).to_dense() for _ in range(1)]).to(q.device)

        est_attn = torch.ones_like(attn_weights)
        dim = 3
        est_attn = est_attn.scatter(3, vertical_topk.expand(*est_attn.shape[:dim], vertical_topk.shape[dim], *est_attn.shape[dim + 1 :]), 0)
        est_attn = est_attn + slash

        est_attn = (est_attn > 0).float()
        est_attn = torch.tril(est_attn)
        est_attn = (est_attn == 0).int() * self.ne_inf
        attn_weights = attn_weights + est_attn
        if self.kv_cache_compressed_v4:
            self.vertical = torch.topk(vertical, vertical_size * 4, -1).indices
            self.slash = (torch.topk(slash_topk, slash_size * 4, -1).indices - (q_len - 1)).unsqueeze(2)
        return attn_weights

    def stream_llm(attn_weights, vertical_size, slash_size):
        q_len = q.shape[2]
        vertical_size, slash_size = min(q_len, max(vertical_size, 30)), min(q_len, max(slash_size, 50))
        mask = torch.triu(torch.tril(torch.ones(q_len, q_len), 0), -slash_size).to(q)
        mask[:,:vertical_size] = 1
        mask = mask.unsqueeze(0).unsqueeze(1)

        est_attn = torch.tril(mask)
        est_attn = (est_attn == 0).int() * self.ne_inf
        attn_weights = attn_weights + est_attn
        if self.kv_cache_compressed_v4:
            self.vertical = torch.Tensor(list(range(vertical_size * 4))).long().to(q.device).unsqueeze(0).unsqueeze(0).unsqueeze(0)
            self.slash = torch.Tensor(list(range(-slash_size * 4, 1))).long().to(q.device).unsqueeze(0).unsqueeze(0).unsqueeze(0)
        return attn_weights

    def block_sparse(attn_weights, topk_ratio, slash_size=None, block_size=8):
        block_num = (q_len -1) // block_size + 1
        block_q = torch.zeros(1,1,block_num * block_size,head_dim).to(q)
        block_q[:,:,:q_len] = q
        block_q = block_q.reshape(1,1,block_num,block_size,-1).mean(-2)
        block_k = torch.zeros(1,1,block_num * block_size,head_dim).to(k)
        block_k[:,:,:q_len] = k
        block_k = block_k.reshape(1,1,block_num,block_size,-1).mean(-2)

        qk = torch.matmul(block_q, block_k.transpose(2, 3)) + attention_mask[:,:,:block_num,:block_num]
        est_attn = torch.ones_like(qk)
        block_topk = torch.topk(-qk, block_num - block_num//topk_ratio, -1).indices

        dim = 3
        est_attn = est_attn.scatter(3, block_topk.expand(*est_attn.shape[:dim], block_topk.shape[dim], *est_attn.shape[dim + 1 :]), 0)
        est_attn = est_attn.unsqueeze(3).unsqueeze(-1).repeat(1,1,1,block_size,1,block_size).reshape(1,1,block_num * block_size, block_num * block_size)[...,:q_len,:q_len]
        est_attn = torch.tril(est_attn)
        est_attn = (est_attn == 0).int()
        attn_weights = attn_weights + est_attn
        return attn_weights

    def dialted(q,k,v, type):
        q_len = q.shape[2]
        n_init = min(1024, q_len)
        vertical_topk = torch.arange(0, n_init, device=q.device)[None, None, None, :]

        slash = torch.arange(0, q_len, device=q.device)
        if type == 'dilated1':
            # 8k local with 1 interval
            slash = slash[-8192::2][None, None, :]
        elif type == 'dilated2':
            # 2k dense local + 4k local with 1 interval
            slash = torch.cat([slash[-2048:], slash[-6144:-2048:2]], 0)[None, None, :]

        slash = (q_len - 1) - slash
        return vertical_slash_sparse_attention(q, k, v, vertical_topk, slash)

    def vertical_and_slash_kernel(q, k, v, vertical_size, slash_size):
        vertical_size, slash_size  = min(q_len, max(vertical_size, 30)), min(q_len, max(slash_size, 50))
        last_q = min(64, q_len)
        qk = torch.einsum(f'bhmk, bhnk -> bhmn', q[:,:,-last_q:,:], k)
        qk[:, :, :, -last_q:] = torch.where(LAST_Q_MASK[...,-last_q:,-last_q:].to(q.device), qk[:, :, :, -last_q:], -torch.inf)
        qk = torch.nn.functional.softmax(qk, dim=-1, dtype=torch.float32)
        vertical = qk.sum(-2, keepdim=True)
        vertical[...,:30] = torch.inf
        vertical_topk = torch.topk(vertical, vertical_size, -1).indices

        slash = sum_all_diagonal_matrix(qk)[...,:-last_q + 1]
        slash[...,-100:] = torch.inf
        slash_topk = slash
        slash = (q_len - 1) - torch.topk(slash, slash_size, -1).indices

        return vertical_slash_sparse_attention(q, k, v, vertical_topk, slash)

    def vertical_and_slash_kernel_static(q, k, v, vertical_size, slash_size):
        if "vs" in self.__dict__:
            vertical_topk, slash = self.vs
        else:
            vertical_size, slash_size  = min(q_len, max(vertical_size, 30)), min(q_len, max(slash_size, 50))
            last_q = 64
            qk = torch.einsum(f'bhmk, bhnk -> bhmn', q[:,:,-last_q:,:], k)
            qk[:, :, :, -last_q:] = torch.where(LAST_Q_MASK, qk[:, :, :, -last_q:], -torch.inf)
            qk = torch.nn.functional.softmax(qk, dim=-1, dtype=torch.float32)
            vertical = qk.sum(-2, keepdim=True)
            vertical[...,:30] = torch.inf
            vertical_topk = torch.topk(vertical, vertical_size, -1).indices

            slash = sum_all_diagonal_matrix(qk)[...,:-last_q + 1]
            slash[...,-30:] = torch.inf
            slash_topk = slash
            slash = (q_len - 1) - torch.topk(slash, slash_size, -1).indices
            self.vs = vertical_topk, slash

        return vertical_slash_sparse_attention(q, k, v, vertical_topk, slash)
    def dense(q, k, v, vertical_size=None, slash_size=None):
        return flash_attn_func(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1,2), 0.0, softmax_scale=None, causal=q_len != 1).view(bsz, 1, q_len, self.head_dim)
    def block_sparse_kernel(q, k, v, vertical_size=None, slash_size=None):
        topk = 100
        return block_sparse_attention(q, k, v, topk)

    q_len = q.shape[2]
    bsz = q.shape[0]

    if self.config.to_dict().get("dilated1", False):
        return dialted(q, k, v, 'dilated1')
    if self.config.to_dict().get("dilated2", False):
        return dialted(q, k, v, 'dilated2')
    if self.config.to_dict().get("dense", False):
        return dense(q, k, v)
    if self.config.to_dict().get("streaming", False):
        return streaming_forward(q, k, v, self.config.streaming_kwargs["n_init"], self.config.streaming_kwargs["n_local"])

    ty, vertical_size, slash_size, _ = self.best_pattern.get(head_id, ("vertical_and_slash", 1000, 6096, 1))

    if self.config.to_dict().get("static_pattern", False):
        return vertical_and_slash_kernel_static(q, k, v, vertical_size, slash_size)
    if self.config.to_dict().get("vs_only", False):
        return vertical_and_slash_kernel(q, k, v, vertical_size, slash_size)

    if q_len == 1:
        return dense(q, k, v)

    fc = {
        "stream_llm": streaming_forward,
        "vertical_and_slash": vertical_and_slash_kernel,
        "block_sparse": block_sparse_kernel,
    }[ty]
    return fc(q, k, v, vertical_size, slash_size)

def apply_rotary_pos_emb_single(q, cos, sin, position_ids, unsqueeze_dim=1):
    # cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    # sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    return (q * cos) + (rotate_half(q) * sin)

def minference_forward():
    def forward(
        self,
        hidden_states,
        attention_mask,
        position_ids,
        past_key_value,
        output_attentions,
        use_cache,
        **kwargs,
    ):
        self.init_minference_parameters()
        self.ne_inf = torch.finfo(hidden_states.dtype).min

        bsz, q_len, _ = hidden_states.size()

        if "q_proj" in self.__dict__["_modules"]:
            query_states = self.q_proj(hidden_states)
            key_states = self.k_proj(hidden_states)
            value_states = self.v_proj(hidden_states)
        else:
            qkv = self.qkv_proj(hidden_states)
            query_pos = self.num_heads * self.head_dim
            query_states, key_states, value_states = torch.split(qkv, query_pos, -1)

        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        kv_seq_len = key_states.shape[-2]
        if past_key_value is not None:
            if self.layer_idx is None:
                raise ValueError(
                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
                    "with a layer index."
                )
            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
        global ROPE_TYPE
        if ROPE_TYPE is None:
            ROPE_TYPE = "seq_len" in inspect.signature(self.rotary_emb.forward).parameters
        if ROPE_TYPE:
            cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
        else:
            cos, sin = self.rotary_emb(value_states, position_ids)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

        if past_key_value is not None:
            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)
        if self.is_search:
            if os.path.exists(self.config_path):
                config_list = json.load(open(self.config_path))
                if self.layer_idx < len(config_list):
                    assert False
            else:
                config_list = []
            config = {}
            print("Layer", self.layer_idx)
        if q_len != 1:
            output = torch.empty_like(query_states)
            for head in range(query_states.size(1)):
                q = query_states[:, head, :, :].unsqueeze(1)
                k = key_states[:, head, :, :].unsqueeze(1)
                v = value_states[:, head, :, :].unsqueeze(1)
                if self.is_search:
                    config[head] = search_pattern(q, k, head)
                if self.layer_idx >= self.starting_layer and not self.is_search:
                    attn_output = self.gather_last_q_vertical_slash_topk_v4(q, k, v, head)
                elif is_flash_attn_2_available():
                    attn_output = flash_attn_func(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1,2), 0.0, softmax_scale=None, causal=q_len != 1).view(bsz, 1, q_len, self.head_dim)
                else:
                    attn_output = gather_qkv(q, k, v, attention_mask)
                output[:, head:head + 1] = attn_output
            if self.is_search:
                config_list.append(config)
                with open(self.config_path, 'w') as json_file:
                    json.dump(config_list, json_file)
        else:
            output =  flash_attn_func(query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1,2), 0.0, softmax_scale=None, causal=q_len != 1).view(bsz, query_states.size(1), q_len, self.head_dim)
        attn_output = output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
        attn_output = self.o_proj(attn_output)

        return attn_output, None, past_key_value

    return forward

def minference_kv_cache_cpu_forward():
    def forward(
        self,
        hidden_states,
        attention_mask,
        position_ids,
        past_key_value,
        output_attentions,
        use_cache,
        **kwargs,
    ):
        self.init_minference_parameters()
        self.ne_inf = torch.finfo(hidden_states.dtype).min

        bsz, q_len, hidden_dim = hidden_states.size()
        kv_seq_len = q_len
        if use_cache and past_key_value is not None:
            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)

        global ROPE_TYPE
        if ROPE_TYPE is None:
            ROPE_TYPE = "seq_len" in inspect.signature(self.rotary_emb.forward).parameters
        if ROPE_TYPE:
            cos, sin = self.rotary_emb(hidden_states, seq_len=kv_seq_len)
        else:
            cos, sin = self.rotary_emb(hidden_states, position_ids)
        cache_kwargs = {"sin": sin, "cos": cos}

        attn_out = torch.empty_like(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
        act_num_heads = self.num_heads // self.num_key_value_groups
        if use_cache:
            k = torch.zeros(bsz, act_num_heads, q_len, self.head_dim).to(hidden_states.dtype).cpu()
            v = torch.zeros(bsz, act_num_heads, q_len, self.head_dim).to(hidden_states.dtype).cpu()
        part_k, part_v = None, None
        for head in range(self.num_heads):
            if "q_proj" in self.__dict__["_modules"]:
                part_q = F.linear(hidden_states, self.q_proj.weight.view(self.num_heads, self.head_dim, hidden_dim)[head]).unsqueeze(2)
            else:
                part_q = F.linear(hidden_states, self.qkv_proj.weight.view(3, self.num_heads, self.head_dim, hidden_dim)[0][head]).unsqueeze(2)
            part_q = apply_rotary_pos_emb_single(part_q.transpose(1, 2), cos, sin, position_ids)

            if head % self.num_key_value_groups == 0:
                if "q_proj" in self.__dict__["_modules"]:
                    part_k = F.linear(hidden_states, self.k_proj.weight.view(act_num_heads, self.head_dim, hidden_dim)[head // self.num_key_value_groups]).unsqueeze(2)
                    part_v = F.linear(hidden_states, self.v_proj.weight.view(act_num_heads, self.head_dim, hidden_dim)[head // self.num_key_value_groups]).unsqueeze(2).transpose(1, 2)
                else:
                    part_k = F.linear(hidden_states, self.qkv_proj.weight.view(3, act_num_heads, self.head_dim, hidden_dim)[1][head // self.num_key_value_groups]).unsqueeze(2)
                    part_v = F.linear(hidden_states, self.qkv_proj.weight.view(3, act_num_heads, self.head_dim, hidden_dim)[2][head // self.num_key_value_groups]).unsqueeze(2).transpose(1, 2)

                part_k = apply_rotary_pos_emb_single(part_k.transpose(1, 2), cos, sin, position_ids)
                if use_cache and past_key_value is not None:
                    k[:,head // self.num_key_value_groups] = part_k.cpu()
                    v[:,head // self.num_key_value_groups] = part_v.cpu()
                    part_k, part_v = past_key_value.get(part_k, part_v, self.layer_idx, head // self.num_key_value_groups, cache_kwargs)

            if self.layer_idx >= self.starting_layer:
                part_o = self.gather_last_q_vertical_slash_topk_v4(part_q, part_k, part_v, head)
            else:
                part_o = flash_attn_func(part_q, part_k, part_v.transpose(1, 2), 0.0, softmax_scale=None, causal=True).view(bsz, part_q.shape[1], self.head_dim)
            attn_out[:, :, head, :] = part_o

        if use_cache and past_key_value is not None:
            past_key_value.update(k, v, self.layer_idx, cache_kwargs)
        torch.matmul(attn_out.view(bsz, q_len, hidden_dim), self.o_proj.weight.T, out=hidden_states)
        torch.cuda.empty_cache()
        return (hidden_states, None, past_key_value)

    return forward

def minference_with_snapkv_forward():
    def forward(
        self,
        hidden_states,
        attention_mask,
        position_ids,
        past_key_value,
        output_attentions,
        use_cache,
        **kwargs,
    ):
        self.init_minference_parameters()
        self.ne_inf = torch.finfo(hidden_states.dtype).min

        init_snapkv(self)

        bsz, q_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        kv_seq_len = key_states.shape[-2]
        if past_key_value is not None:
            if self.layer_idx is None:
                raise ValueError(
                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
                    "with a layer index."
                )

            if hasattr(self, "kv_seq_len"): #[SnapKV] add kv_seq_len
                if self.kv_seq_len != 0:
                    kv_seq_len += self.kv_seq_len
                else:
                    kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
            else:
                kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
        global ROPE_TYPE
        if ROPE_TYPE is None:
            ROPE_TYPE = "seq_len" in inspect.signature(self.rotary_emb.forward).parameters
        if ROPE_TYPE:
            cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
        else:
            cos, sin = self.rotary_emb(value_states, position_ids)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        if past_key_value is not None:
            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
            if key_states.shape[-2] == kv_seq_len: # [SnapKV] add kv_cluster
                self.kv_seq_len = kv_seq_len # [SnapKV] register kv_seq_len
                key_states_compress, value_states_compress = self.kv_cluster.update_kv(key_states, query_states, value_states, attention_mask, self.num_key_value_groups)
                past_key_value.update(key_states_compress, value_states_compress, self.layer_idx, cache_kwargs)
            else:
                self.kv_seq_len += q_len
                key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

        if self.layer_idx >= self.starting_layer:
            assert query_states.size(1) == key_states.size(1) == value_states.size(1)
            output = torch.empty_like(query_states)
            for head in range(query_states.size(1)):
                q = query_states[:, head, :, :].unsqueeze(1)
                k = key_states[:, head, :, :].unsqueeze(1)
                v = value_states[:, head, :, :].unsqueeze(1)
                output[:, head:head + 1] = self.gather_last_q_vertical_slash_topk_v4(q, k, v, head)

            attn_output = output.transpose(1, 2).contiguous()
            attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
            attn_output = self.o_proj(attn_output)
            return attn_output, None, past_key_value

        else:
            output = torch.empty_like(query_states)
            for head in range(query_states.size(1)):
                q = query_states[:, head, :, :].unsqueeze(1)
                k = key_states[:, head, :, :].unsqueeze(1)
                v = value_states[:, head, :, :].unsqueeze(1)
                if is_flash_attn_2_available():
                    attn_output = flash_attn_func(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1,2), 0.0, softmax_scale=None, causal=q_len != 1).view(bsz, 1, q.shape[2], self.head_dim)
                else:
                    attn_output = gather_qkv(q, k, v, attention_mask)
                output[:, head:head + 1] = attn_output
            attn_output = output.transpose(1, 2).contiguous()
            attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
            attn_output = self.o_proj(attn_output)

            return attn_output, None, past_key_value

    return forward

def gather_last_q_vertical_slash_topk_vllm(self, q, k, v, head_id):
    kv_seq_len = k.size(2)
    head_dim = q.size(-1)

    def vertical_and_slash_kernel(q, k, v, vertical_size, slash_size):
        vertical_size, slash_size  = min(q_len, max(vertical_size, 30)), min(q_len, max(slash_size, 50))
        last_q = min(64, q_len)
        qk = torch.einsum(f'bhmk, bhnk -> bhmn', q[:,:,-last_q:,:], k)

        qk[:, :, :, -last_q:] = torch.where(LAST_Q_MASK[...,-last_q:,-last_q:], qk[:, :, :, -last_q:], -torch.inf)
        qk = torch.nn.functional.softmax(qk, dim=-1, dtype=torch.float32)
        vertical = qk.sum(-2, keepdim=True)
        vertical[...,:30] = torch.inf
        vertical_topk = torch.topk(vertical, vertical_size, -1).indices

        slash = sum_all_diagonal_matrix(qk)[...,:-last_q + 1]
        slash[...,-100:] = torch.inf
        slash_topk = slash
        slash = (q_len - 1) - torch.topk(slash, slash_size, -1).indices

        return vertical_slash_sparse_attention(q, k, v, vertical_topk, slash)

    def block_sparse_kernel(q, k, v, vertical_size=None, slash_size=None):
        topk = 100
        return block_sparse_attention(q, k, v, topk)

    def dense(q, k, v, vertical_size=None, slash_size=None):
        return flash_attn_func(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1,2), 0.0, softmax_scale=None, causal=q_len != 1).view(bsz, 1, q_len, head_dim)

    q_len = q.shape[2]
    bsz = q.shape[0]

    ty, vertical_size, slash_size, _ = self.best_pattern[head_id]

    if q_len == 1:
        return dense(q, k, v)

    fc = {
        "stream_llm": streaming_forward,
        "vertical_and_slash": vertical_and_slash_kernel,
        "block_sparse": block_sparse_kernel,
    }[ty]
    return fc(q, k, v, vertical_size, slash_size)

def minference_vllm_forward(
    pattern_config
):
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata[FlashAttentionMetadata],
        kv_scale: float,
        layer_idx: int,
    ) -> torch.Tensor:
        """Forward pass with FlashAttention and PagedAttention.

        Args:
            query: shape = [num_tokens, num_heads * head_size]
            key: shape = [num_tokens, num_kv_heads * head_size]
            value: shape = [num_tokens, num_kv_heads * head_size]
            kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]
            attn_metadata: Metadata for attention.
        Returns:
            shape = [num_tokens, num_heads * head_size]
        """
        self.best_pattern = {int(ii): jj for ii, jj in pattern_config[layer_idx].items()}
        def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
            """
            This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
            num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
            """
            slen, num_key_value_heads, head_dim = hidden_states.shape
            if n_rep == 1:
                return hidden_states
            hidden_states = hidden_states[:, None, :, :].expand(slen, n_rep, num_key_value_heads, head_dim)
            return hidden_states.reshape(slen, num_key_value_heads * n_rep, head_dim)

        def minference_prefill_func(
            q, k, v,

        ):
            # (seq_len, num_heads, head_size)
            if q.size(-2) != k.size(-2):
                k = repeat_kv(k, q.size(-2) // k.size(-2))
                v = repeat_kv(v, q.size(-2) // v.size(-2))

            output = torch.empty_like(q)
            for head in range(q.size(-2)):
                q_head = q[:, head, :].unsqueeze(1)
                k_head = k[:, head, :].unsqueeze(1)
                v_head = v[:, head, :].unsqueeze(1)

                # (1, seq_len, num_heads, head_size)
                q_head = q_head[None, ...]
                k_head = k_head[None, ...]
                v_head = v_head[None, ...]

                q_head = q_head.transpose(1, 2)
                k_head = k_head.transpose(1, 2)
                v_head = v_head.transpose(1, 2)

                out = self.gather_last_q_vertical_slash_topk_vllm(q_head, k_head, v_head, head)

                out = out.transpose(1, 2).squeeze(0).contiguous()
                output[:, head:head+1, :] = out
            return output

        num_tokens, hidden_size = query.shape
        # Reshape the query, key, and value tensors.
        query = query.view(-1, self.num_heads, self.head_size)
        key = key.view(-1, self.num_kv_heads, self.head_size)
        value = value.view(-1, self.num_kv_heads, self.head_size)

        if kv_cache is not None:
            key_cache, value_cache = PagedAttention.split_kv_cache(
                kv_cache, self.num_kv_heads, self.head_size)

            # Reshape the input keys and values and store them in the cache.
            # If kv_cache is not provided, the new key and value tensors are
            # not cached. This happens during the initial memory profiling run.
            PagedAttention.write_to_paged_cache(key, value, key_cache,
                                                value_cache,
                                                attn_metadata.slot_mapping,
                                                attn_metadata.kv_cache_dtype,
                                                kv_scale)

        num_prefill_tokens = attn_metadata.num_prefill_tokens
        num_decode_tokens = attn_metadata.num_decode_tokens
        assert key.shape[0] == num_prefill_tokens + num_decode_tokens
        assert value.shape[0] == num_prefill_tokens + num_decode_tokens

        output = torch.empty_like(query)
        # Query for decode. KV is not needed because it is already cached.
        decode_query = query[num_prefill_tokens:]
        # QKV for prefill.
        query = query[:num_prefill_tokens]
        key = key[:num_prefill_tokens]
        value = value[:num_prefill_tokens]

        assert query.shape[0] == num_prefill_tokens
        assert decode_query.shape[0] == num_decode_tokens

        if prefill_meta := attn_metadata.prefill_metadata:
            # Prompt run.
            if kv_cache is None or prefill_meta.block_tables.numel() == 0:
                # normal attention
                # When block_tables are not filled, it means q and k are the
                # prompt, and they have the same length.
                # (seq_len, num_heads, head_size)
                # out = flash_attn_varlen_func(
                #     q=query,
                #     k=key,
                #     v=value,
                #     cu_seqlens_q=prefill_meta.seq_start_loc,
                #     cu_seqlens_k=prefill_meta.seq_start_loc,
                #     max_seqlen_q=prefill_meta.max_prompt_len,
                #     max_seqlen_k=prefill_meta.max_prompt_len,
                #     softmax_scale=self.scale,
                #     causal=True,
                #     window_size=self.sliding_window,
                #     alibi_slopes=self.alibi_slopes,
                # )
                out = minference_prefill_func(query, key, value)
                assert output[:num_prefill_tokens].shape == out.shape
                output[:num_prefill_tokens] = out
            else:
                # prefix-enabled attention
                # TODO(Hai) this triton kernel has regression issue (broke) to
                # deal with different data types between KV and FP8 KV cache,
                # to be addressed separately.
                output[:num_prefill_tokens] = PagedAttention.forward_prefix(
                    query,
                    key,
                    value,
                    key_cache,
                    value_cache,
                    prefill_meta.block_tables,
                    prefill_meta.subquery_start_loc,
                    prefill_meta.prompt_lens_tensor,
                    prefill_meta.context_lens,
                    prefill_meta.max_subquery_len,
                    self.alibi_slopes,
                )
        if decode_meta := attn_metadata.decode_metadata:
            # Decoding run.
            output[num_prefill_tokens:] = PagedAttention.forward_decode(
                decode_query,
                key_cache,
                value_cache,
                decode_meta.block_tables,
                decode_meta.context_lens,
                decode_meta.max_context_len,
                attn_metadata.kv_cache_dtype,
                self.num_kv_heads,
                self.scale,
                self.alibi_slopes,
                kv_scale,
            )

        # Reshape the output tensor.
        return output.view(num_tokens, hidden_size)

    return forward