# -------------------------------------------------------------------------
#  This file is part of the MindStudio project.
# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
#
# MindStudio is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
# -------------------------------------------------------------------------

from collections import namedtuple
import torch
import torch.nn as nn
import numpy as np

from einops import rearrange


from msprobe.pytorch.common.utils import logger

GTYPE = torch.float64  # arm host必须选择float64，x86环境选择float32即可，64也行。arm计算很慢，s=8k的场景建议使用x86
SOFTMAX_BUILD_MODE = "QKV"  # "MAX_SUM"

FaForwardParams = namedtuple("FaForwardParams",
                             ["q", "k", "v", "drop_mask", "attn_mask", "pse", "scalar_value", "keep_prob"])
FaBackwardParams = namedtuple("FaBackwardParams",
                              ["dx", "q", "k", "v", "softmax_res", "drop_mask", "pse", "scalar_value", "keep_prob"])
RebuildSoftmaxParams = namedtuple("RebuildSoftmaxParams",
                                  ["q", "k", "attn_mask", "pse", "scalar_value", "softmax_max", "softmax_sum"])


def softmax_forward(x):
    x_max = torch.max(x, dim=-1, keepdims=True)[0]
    x_sub = x.sub(x_max)
    y = torch.exp(x_sub)
    x_sum = y.sum(dim=-1, keepdims=True)
    res = y.div(x_sum)
    return res, x_max, x_sum


def softmax_grad(dp, softmax_res):
    muls = dp * softmax_res
    muls_r = muls.sum(dim=-1, keepdims=True)
    sub_r = dp - muls_r
    res = sub_r * softmax_res
    return res


def broadcast_kv(num_heads, num_kv_heads, kv_tensor, dtype):
    # 检查维度
    if kv_tensor.dim() != 4:
        raise ValueError(f"broadcast_kv: kv_tensor 必须是 4 维 (B, N_kv, S, D)，但得到 {kv_tensor.shape}")
    if num_kv_heads == 0 or num_kv_heads > num_heads:
        raise ValueError("broadcast_kv: num_kv_heads 必须大于 0 且不超过 num_heads。")
    if num_heads % num_kv_heads != 0:
        raise ValueError(f"broadcast_kv: num_heads({num_heads}) 必须能被 num_kv_heads({num_kv_heads}) 整除。")


    factor = num_heads // num_kv_heads
    kv_shape = kv_tensor.shape
    b = kv_shape[0]
    s = kv_shape[2]
    d = kv_shape[3]
    kv_res = torch.zeros([b, num_heads, s, d]).to(dtype)
    for i in range(num_heads):
        j = i // factor
        kv_res[:, i:i + 1, :, :] = kv_tensor[:, j:j + 1, :, :]
    return kv_res


def calculate_qk(q, k, attn_mask, pse, scalar_value):
    # 基本形状检查
    if q.dim() < 4 or k.dim() < 4:
        raise ValueError(f"calculate_qk: q,k 必须至少 4 维，q={q.dim()}，k={k.dim()}")
    # 检查 head_dim 一致性
    if q.size(-1) != k.size(-1):
        raise ValueError(f"calculate_qk: q.head_dim({q.size(-1)}) != k.head_dim({k.size(-1)})")

    if k.dim() != 4:
        raise ValueError(f"k tensor dimension must be 4, but got {k.dim()} dimensions (shape: {k.shape})")

    if k.dim() == 3:
        k = k.unsqueeze(1)  # 在head维度扩展

    if pse is None or len(pse.shape) == 0:
        qk = torch.matmul(q, k.permute(0, 1, 3, 2)).mul(scalar_value)
    else:
        qk = (torch.matmul(q, k.permute(0, 1, 3, 2)) + pse).mul(scalar_value)
    if attn_mask is None or len(attn_mask.shape) == 0:
        return qk
    else:
        qk = qk + attn_mask.bool() * (-40000.0)  # -10000
    return qk


def fusion_attention_forward(forward_params):
    q = forward_params.q
    k = forward_params.k
    v = forward_params.v
    drop_mask = forward_params.drop_mask
    attn_mask = forward_params.attn_mask
    pse = forward_params.pse
    scalar_value = forward_params.scalar_value
    keep_prob = forward_params.keep_prob

    # 拦截 keep_prob 为 0 的情况，防止除零
    if keep_prob == 0:
        raise ValueError("fusion_attention_forward: keep_prob 不能为 0，避免除零错误。")

    qk = calculate_qk(q, k, attn_mask, pse, scalar_value)
    softmax_res, softmax_max, softmax_sum = softmax_forward(qk)
    if drop_mask is None or len(drop_mask.shape) == 0:
        drop_res = softmax_res
    else:
        drop_res = softmax_res * drop_mask * (1.0 / keep_prob)
    y = torch.matmul(drop_res, v)
    return y, softmax_max, softmax_sum


def fusion_attention_backward(backward_params):
    dx = backward_params.dx
    q = backward_params.q
    k = backward_params.k
    v = backward_params.v
    softmax_res = backward_params.softmax_res
    drop_mask = backward_params.drop_mask
    pse = backward_params.pse
    scalar_value = backward_params.scalar_value
    keep_prob = backward_params.keep_prob

    # 拦截 keep_prob 为 0 的情况，防止除零
    if keep_prob == 0:
        raise ValueError("fusion_attention_backward: keep_prob 不能为 0，避免除零错误。")

    dp = torch.matmul(dx, v.permute(0, 1, 3, 2))
    if drop_mask is None or len(drop_mask.shape) == 0:
        drop_res = softmax_res.permute(0, 1, 3, 2)
        dp_drop = dp
    else:
        drop_res = softmax_res.mul(drop_mask).mul(1.0 / keep_prob).permute(0, 1, 3, 2)
        dp_drop = dp * drop_mask * (1.0 / keep_prob)
    dv = torch.matmul(drop_res, dx)
    softmax_grad_res = (softmax_grad(dp_drop, softmax_res) * scalar_value)
    dq = torch.matmul(softmax_grad_res, k)
    dk = torch.matmul(softmax_grad_res.permute(0, 1, 3, 2), q)
    return dq, dk, dv


def parse_bsnd_args(query, key, head_num, input_layout):
    supported_input_layout = ["BSH", "SBH", "BSND", "BNSD", "TND"]
    b, s1, s2, n1, n2, d, h1, h2 = None, None, None, head_num, None, None, None, None

    if not isinstance(input_layout, str) or input_layout not in supported_input_layout:
        raise ValueError(f"Invalid input_layout arg which must be one of {supported_input_layout}.")

    if input_layout == "TND":
        raise ValueError(f"input_layout {input_layout} does not supported for now.")

    # 防止 head_num 为 0
    if n1 == 0:
        raise ValueError("parse_bsnd_args: head_num (n1) 不能为 0，避免除零错误。")

    try:
        if input_layout == "BSH":
            b, s1, h1 = query.shape
            _, s2, h2 = key.shape
            d = h1 // n1
            # 拦截 d 为 0 的情况
            if d == 0:
                raise ValueError("parse_bsnd_args: 计算得到的 head_dim d 不能为 0。")
            n2 = h2 // d
        elif input_layout == "SBH":
            s1, b, h1 = query.shape
            s2, _, h2 = key.shape
            d = h1 // n1
            if d == 0:
                raise ValueError("parse_bsnd_args: 计算得到的 head_dim d 不能为 0。")
            n2 = h2 // d
        elif input_layout == "BSND":
            b, s1, n1, d = query.shape
            _, s2, n2, _ = key.shape
            if d == 0:
                raise ValueError("parse_bsnd_args: head_dim d 不能为 0。")
            h1 = n1 * d
            h2 = n2 * d
        elif input_layout == "BNSD":
            b, n1, s1, d = query.shape
            _, n2, s2, _ = key.shape
            if d == 0:
                raise ValueError("parse_bsnd_args: head_dim d 不能为 0。")
            h1 = n1 * d
            h2 = n2 * d
    except Exception as e:
        raise ValueError(f"query.shape: {query.shape}, key.shape: {key.shape}, parse_bsnd_args error: {e}") from e

    ret = (b, s1, s2, n1, n2, d, h1, h2, query.dtype)
    return ret


def convert_from_bnsd(_input, input_layout):
    """
    transform qkv from bnsd to input_layout.
    B: batch_size
    S: sequence_length
    N: num_heads
    D: head_dim
    Args:
       _input (torch.Tensor): tensor of shape (B,N,S,D)
        input_layout (str): "BSH" or "SBH" or "BSND" or "BNSD" or "TND"
    Returns:
        tensor of shape (B,N,S,D) or (B,S,N,D) or (S,B,H) or (B,S,H)
    """
    if input_layout == "BSH":
        # (B,N,S,D)=>(B,S,N*D)
        out = rearrange(_input, 'b n s d -> b s (n d)').contiguous()
    elif input_layout == "SBH":
        # (B,N,S,D)=>(S,B,N*D)
        out = rearrange(_input, 'b n s d -> s b (n d)').contiguous()
    elif input_layout == "BSND":
        # (B,N,S,D)=>(B,S,N,D)
        out = rearrange(_input, 'b n s d -> b s n d').contiguous()
    elif input_layout == "TND":
        raise ValueError(f"input_layout {input_layout} does not supported for now.")
    else:
        out = _input
    return out


def convert_to_bnsd(_input, n, input_layout):
    """
    transform qkv from input_layout to bnsd.
    B: batch_size
    S: sequence_length
    N: num_heads
    D: head_dim
    Args:
        _input (torch.Tensor): tensor of shape (B,N,S,D) or (B,S,N,D) or (S,B,H) or (B,S,H)
        n (int): num_heads
        input_layout (str):"BSH" or "SBH" or "BSND" or "BNSD" or "TND"
    Returns:
        tensor of shape (B,N,S,D)
    """
    if input_layout == "BSH":
        # (B,S,N*D)=>(B,N,S,D)
        out = rearrange(_input, 'b s (n d) -> b n s d', n=n)
    elif input_layout == "SBH":
        # (S,B,N*D)=>(B,N,S,D)
        out = rearrange(_input, 's b (n d) -> b n s d', n=n)
    elif input_layout == "BSND":
        # (B,S,N,D)=>(B,N,S,D)
        out = rearrange(_input, 'b s n d -> b n s d', n=n)
    elif input_layout == "TND":
        raise ValueError(f"input_layout {input_layout} does not supported for now.")
    else:
        out = _input
    if out.dim() != 4:
        raise ValueError(f"convert qkv format failed with input_layout {input_layout}.")
    return out.to(GTYPE)


def generate_attn_mask(*args):
    """
    # 当sparse_mode=2、3、4时小算子到融合算子会走这个优化，反过来看就要拆解回原来的基本实现
    ===> attn_mask = torch.from_numpy(np.triu(np.ones([2048, 2048]), k=1)).to(dtype)
    """

    sparse_mode, attn_mask, b, n1, s1, s2, pre_tocken, next_tocken, dtype = args
    shape = [s1, s2]

    if attn_mask is not None:
        # 当FA的输入已经包含attn_mask时，可以认为已经是转换之后的mask矩阵了，有三种特殊场景，即稀疏矩阵场景，需要进行逆向还原
        if sparse_mode == 2 or sparse_mode == 3 or sparse_mode == 4:
            logger.info(f"s1: {s1}, s2:{s2}, attn_mask.shape:{attn_mask.shape}, attn_mask.dtype:{attn_mask.dtype}")

            if attn_mask.dim() == 2 and attn_mask.shape[0] == 2048 and attn_mask.shape[1] == 2048:
                if attn_mask.equal(torch.from_numpy(np.triu(np.ones([2048, 2048]), k=1)).to(attn_mask.dtype)):
                    if sparse_mode == 2:
                        attn_mask = torch.from_numpy(np.triu(np.ones(shape), k=1))
                    elif sparse_mode == 3:
                        attn_mask = torch.from_numpy(np.triu(np.ones(shape), k=s2 - s1 + 1))
                    elif sparse_mode == 4:
                        attn_mask_u = torch.from_numpy(np.triu(np.ones(shape), k=next_tocken + 1))
                        attn_mask_l = torch.from_numpy(np.tril(np.ones(shape), k=-pre_tocken - 1))
                        attn_mask = attn_mask_u + attn_mask_l
                    logger.debug(f"反向转换attn_mask {attn_mask.shape}")
                    return attn_mask.to(dtype)

        return attn_mask.to(dtype)

    if attn_mask is not None:
        if attn_mask.dim() == 2:
            if attn_mask.shape[0] != s1 or attn_mask.shape[1] != s2:
                raise ValueError(f"Invalid attn_mask shape `SS` {attn_mask.shape}")
            shape = [s1, s2]
        elif attn_mask.dim() == 4:
            if attn_mask.shape[1] == 1:
                shape = [b, 1, s1, s2] if b != 1 else [1, 1, s1, s2]
            else:
                shape = [b, n1, s1, s2] if b != 1 else [1, n1, s1, s2]

    if sparse_mode == 0:
        attn_mask_u = torch.from_numpy(np.triu(np.ones(shape), k=next_tocken + 1))
        attn_mask_l = torch.from_numpy(np.tril(np.ones(shape), k=-pre_tocken - 1))
        attn_mask = attn_mask_u + attn_mask_l
    elif sparse_mode == 1:  # no sparse
        attn_mask = torch.from_numpy(np.zeros(shape))
    elif sparse_mode == 2:
        attn_mask = torch.from_numpy(np.triu(np.ones(shape), k=1))
    elif sparse_mode == 3:
        attn_mask = torch.from_numpy(np.triu(np.ones(shape), k=s2 - s1 + 1))
    elif sparse_mode == 4:
        attn_mask_u = torch.from_numpy(np.triu(np.ones(shape), k=next_tocken + 1))
        attn_mask_l = torch.from_numpy(np.tril(np.ones(shape), k=-pre_tocken - 1))
        attn_mask = attn_mask_u + attn_mask_l
    # 注:不会出现sparse_mode=5的情况，该情况要求必须要传入attn_mask，且attn_mask矩阵数据格式须为BNSS或B1SS，
    # 因此可以认为FA的输入已经是正确的attn_mask了
    return attn_mask.to(dtype)


def generate_kv(key, value, n1, n2):
    # N不等长适配by cdy
    if not (n1 == n2):
        k_new = broadcast_kv(n1, n2, key, key.dtype)
        v_new = broadcast_kv(n1, n2, value, value.dtype)
    else:
        k_new = key
        v_new = value
    return k_new, v_new


def rebuid_softmax_by_qkv(q, k, attn_mask, pse, scalar_value):
    """
    attention = softmax(QK^T/sqrt(d))V
    softmax(x_i) = e^(x_i - x_max) / sum(e^(x_i - x_max))
    """
    logger.info("Using QKV to rebuild original softmax")
    qk = calculate_qk(q, k, attn_mask, pse, scalar_value)
    softmax_res, _, _ = softmax_forward(qk)
    return softmax_res


def rebuild_softmax_by_max_sum(softmax_params):
    """
    attention = softmax(QK^T/sqrt(d))V
    softmax(x_i) = e^(x_i - x_max_i) / x_sum_i)
    """
    q = softmax_params.q
    k = softmax_params.k
    attn_mask = softmax_params.attn_mask
    pse = softmax_params.pse
    scalar_value = softmax_params.scalar_value
    softmax_max = softmax_params.softmax_max
    softmax_sum = softmax_params.softmax_sum
    logger.info("Using softmax_max and softmax_sum to rebuild original softmax")

    qk = calculate_qk(q, k, attn_mask, pse, scalar_value)
    if softmax_max.shape[-1] == 0:
        raise ValueError(f"softmax_max.shape[-1] must be non-zero, softmax_max.shape: {softmax_max.shape}")
    repeat_dim = qk.shape[-1] // softmax_max.shape[-1]
    softmax_res = torch.exp(qk.sub(softmax_max.repeat(1, 1, 1, repeat_dim))).div(
        softmax_sum.repeat(1, 1, 1, repeat_dim))
    return softmax_res


def get_head_num(*args, **kwargs):
    if kwargs.get("head_num", None):
        head_num = kwargs.get("head_num")
    elif len(args) >= 4:
        head_num = args[3]
    else:
        raise ValueError(f"Unsupported npu_fusion_attention args {args}.")
    return head_num


def get_input_layout(*args, **kwargs):
    if kwargs.get("input_layout", None):
        input_layout = kwargs.get("input_layout")
    elif len(args) >= 5:
        input_layout = args[4]
    else:
        raise ValueError(f"Unsupported npu_fusion_attention args {args}.")
    return input_layout


def npu_fusion_attention_forward_patch(*args, **kwargs):
    if len(args) < 2:
        raise RuntimeError("npu_fusion_attention_forward_patch: length of args should be greater than or equal to 2.")

    # query, key, value, head_num, input_layout
    head_num = get_head_num(*args, **kwargs)
    input_layout = get_input_layout(*args, **kwargs)

    b, s1, s2, n1, n2, d, h1, h2, dtype = parse_bsnd_args(args[0], args[1], head_num, input_layout)
    # 此处 d 已在 parse_bsnd_args 中检查为非零
    if n1 == n2 and s1 == s2:
        logger.debug(f"running case : BNSD = {b}_{n1}_{s1}_{d}, sparse = {kwargs.get('sparse_mode', 0)}")
    else:
        logger.debug(f"running case: BNSD = {b}_{n1}({n2})_{s1}({s2})_{d}, sparse = {kwargs.get('sparse_mode', 0)}")
    if n2 == 0:
        raise ValueError("n2 不能为 0，避免除零错误。")
    if not (n1 % n2 == 0 and n1 >= n2):
        raise ValueError(f"N1与N2不匹配,请检查: n1 = {n1}, n2 = {n2}.")

    dims_kwargs = {
        "b": b, "s1": s1, "s2": s2, "n1": n1, "n2": n2,
        "d": d, "h1": h1, "h2": h2, "dtype": dtype
    }
    new_kwargs = {
        "keep_prob": 1,  # 注意：如果外部传入 keep_prob 为 0，也会在 fusion_attention_forward 中捕获
        "scalar_value": kwargs.get("scalar_value", 1 / (d ** 0.5)),
        "sparse_mode": kwargs.get("sparse_mode", 0),
        "prefix": kwargs.get("prefix"),
        "pre_tockens": kwargs.get("pre_tockens", 2147483647),
        "next_tockens": kwargs.get("next_tockens", 2147483647),
        "pse": kwargs.get("pse"),
        "padding_mask": kwargs.get("padding_mask"),
        "attn_mask": kwargs.get("attn_mask")
    }

    return args, dims_kwargs, new_kwargs


def npu_fusion_attention_backward_patch(*args, **kwargs):
    if len(args) != 6:
        raise ValueError(f"Unsupported npu_fusion_attention_grad args {args}.")

    b, s1, s2, n1, n2, d, h1, h2, dtype = parse_bsnd_args(args[0], args[1], args[4], args[5])
    # 此处 d 已在 parse_bsnd_args 中检查为非零
    if n1 == n2 and s1 == s2:
        logger.info(f"running case : bnsd = {b}_{n1}_{s1}_{d}, sparse = {kwargs.get('sparse_mode', 0)}")
    else:
        logger.info(f"running case: bnsd = {b}_{n1}({n2})_{s1}({s2})_{d}, sparse = {kwargs.get('sparse_mode', 0)}")
    if n2 == 0:
        raise ValueError("n2 不能为 0，避免除零错误。")
    if not (n1 % n2 == 0 and n1 >= n2):
        raise ValueError(f"N1与N2不匹配,请检查: n1 = {n1}, n2 = {n2}.")

    dims_kwargs = {
        "b": b, "s1": s1, "s2": s2, "n1": n1, "n2": n2,
        "d": d, "h1": h1, "h2": h2, "dtype": dtype
    }

    new_kwargs = {
        "keep_prob": 1,  # 同上，fusion_attention_backward 内会拦截 keep_prob 为 0 的情况
        "scalar_value_value": kwargs.get("scalar_value_value", 1 / (d ** 0.5)),
        "sparse_mode": kwargs.get("sparse_mode", 0),
        "prefix": kwargs.get("prefix"),
        "pre_tockens": kwargs.get("pre_tockens", 2147483647),
        "next_tockens": kwargs.get("next_tockens", 2147483647),
        "pse": kwargs.get("pse"),
        "padding_mask": kwargs.get("padding_mask"),
        "softmax_max": kwargs.get("softmax_max"),
        "softmax_sum": kwargs.get("softmax_sum"),
        "softmax_in": kwargs.get("softmax_in"),
        "attention_in": kwargs.get("attention_in"),
        "seed": kwargs.get("seed", 0),
        "offset": kwargs.get("offset", 0),
        "numels": kwargs.get("numels", 0),
        "attn_mask": kwargs.get("attn_mask")
    }

    return args, dims_kwargs, new_kwargs


class FlashAttentionScore(nn.Module):
    def __init__(self):
        super(FlashAttentionScore, self).__init__()
        # You can initialize any parameters here if necessary

    def forward(self, *inputs, **kwargs):
        # Extract the inputs for the attention calculation
        new_args, dims_kwargs, new_kwargs = npu_fusion_attention_forward_patch(*inputs, **kwargs)
        query, key, value = new_args[0], new_args[1], new_args[2]

        input_layout = get_input_layout(*inputs, **kwargs)

        n1 = dims_kwargs.get("n1")
        n2 = dims_kwargs.get("n2")
        s1 = dims_kwargs.get("s1")
        s2 = dims_kwargs.get("s2")
        b = dims_kwargs.get("b")
        dtype = dims_kwargs.get("dtype")
        attn_mask = new_kwargs.get("attn_mask")
        keep_prob = new_kwargs.get("keep_prob")
        sparse_mode = new_kwargs.get("sparse_mode")
        pre_tockens = new_kwargs.get("pre_tockens")
        next_tockens = new_kwargs.get("next_tokens")
        pse = new_kwargs.get("real_shift")
        scalar_value = new_kwargs.get("scalar_value")

        args_temp = [sparse_mode, attn_mask, b, n1, s1, s2, pre_tockens, next_tockens, dtype]

        attn_mask = generate_attn_mask(*args_temp)
        query = convert_to_bnsd(query, n1, input_layout)
        key = convert_to_bnsd(key, n2, input_layout)
        value = convert_to_bnsd(value, n2, input_layout)

        forward_params = FaForwardParams(
            q=query,
            k=key,
            v=value,
            drop_mask=None,
            attn_mask=attn_mask,
            pse=pse,
            scalar_value=scalar_value,
            keep_prob=keep_prob
        )

        out_golden, softmax_max, softmax_sum = fusion_attention_forward(forward_params)

        # If output dimension is 5, reshape accordingly
        if out_golden.dim() == 5:
            out_golden = out_golden.reshape(out_golden.size(0),
                                            out_golden.size(1) * out_golden.size(2),
                                            out_golden.size(3), out_golden.size(4))

        out_golden = convert_from_bnsd(out_golden, input_layout)

        # Ensure the output matches the desired layout
        out_golden = out_golden.cpu(), softmax_max.repeat(1, 1, 1, 8).cpu(), softmax_sum.repeat(1, 1, 1, 8).cpu()

        return out_golden

    def backward(self, *inputs, **kwargs):
        # The backward pass will be similar to what was described for the gradient computation
        new_args, dims_kwargs, new_kwargs = npu_fusion_attention_backward_patch(*inputs, **kwargs)
        query, key, value, dx, input_layout = new_args[0], new_args[1], new_args[2], new_args[3], new_args[5]
        n1 = dims_kwargs.get("n1")
        n2 = dims_kwargs.get("n2")
        s1 = dims_kwargs.get("s1")
        s2 = dims_kwargs.get("s2")
        b = dims_kwargs.get("b")
        dtype = dims_kwargs.get("dtype")
        attn_mask = new_kwargs.get("attn_mask")
        keep_prob = new_kwargs.get("keep_prob")
        sparse_mode = new_kwargs.get("sparse_mode")
        pre_tockens = new_kwargs.get("pre_tockens")
        next_tockens = new_kwargs.get("next_tockens")
        pse = new_kwargs.get("pse")
        softmax_max = new_kwargs.get("softmax_max")
        softmax_sum = new_kwargs.get("softmax_sum")
        scalar_value = new_kwargs.get("scalar_value")

        args_temp = [sparse_mode, attn_mask, b, n1, s1, s2, pre_tockens, next_tockens, dtype]
        attn_mask = generate_attn_mask(*args_temp)

        query = convert_to_bnsd(query, n1, input_layout)
        dx = convert_to_bnsd(dx, n1, input_layout)
        key = convert_to_bnsd(key, n2, input_layout)
        value = convert_to_bnsd(value, n2, input_layout)

        k_new, v_new = generate_kv(key, value, n1, n2)

        if SOFTMAX_BUILD_MODE == "QKV":
            softmax_res = rebuid_softmax_by_qkv(query, k_new, attn_mask, pse, scalar_value)
        else:
            softmax_params = RebuildSoftmaxParams(query, k_new, attn_mask, pse, scalar_value, softmax_max, softmax_sum)
            softmax_res = rebuild_softmax_by_max_sum(softmax_params)

        backward_params = FaBackwardParams(dx, query, k_new, v_new, softmax_res, None, pse, scalar_value, keep_prob)
        dq, dk, dv = fusion_attention_backward(backward_params)

        # Reshape as needed
        if dq.dim() == 5:
            dq = dq.reshape(dq.size(0), dq.size(1) * dq.size(2), dq.size(3), dq.size(4))
        if dk.dim() == 5:
            dk = dk.reshape(dk.size(0), dk.size(1) * dk.size(2), dk.size(3), dk.size(4))
        if dv.dim() == 5:
            dv = dv.reshape(dv.size(0), dv.size(1) * dv.size(2), dv.size(3), dv.size(4))

        dq = convert_from_bnsd(dq, input_layout)
        dk = convert_from_bnsd(dk, input_layout)
        dv = convert_from_bnsd(dv, input_layout)

        return dq.cpu(), dk.cpu(), dv.cpu()
