// Copyright 2025 Tencent
// SPDX-License-Identifier: BSD-3-Clause

#include "testutil.h"

#include <float.h>

static int test_multiheadattention_cross_kvcache(const ncnn::Mat& q, const ncnn::Mat& k, const ncnn::Mat& v, int embed_dim, int num_heads, int attn_mask, int input_kvcache)
{
    const int qdim = q.w;
    const int kdim = k.w;
    const int vdim = v.w;

    ncnn::ParamDict pd;
    pd.set(0, embed_dim);
    pd.set(1, num_heads);
    pd.set(2, embed_dim * qdim);
    pd.set(3, kdim);
    pd.set(4, vdim);
    pd.set(5, attn_mask);
    pd.set(6, 1.f / sqrtf(embed_dim / num_heads));
    pd.set(7, 1); // kv_cache

    std::vector<ncnn::Mat> weights(8);
    weights[0] = RandomMat(embed_dim * qdim);
    weights[1] = RandomMat(embed_dim);
    weights[2] = RandomMat(embed_dim * kdim);
    weights[3] = RandomMat(embed_dim);
    weights[4] = RandomMat(embed_dim * vdim);
    weights[5] = RandomMat(embed_dim);
    weights[6] = RandomMat(qdim * embed_dim);
    weights[7] = RandomMat(qdim);

    std::vector<ncnn::Mat> as(3);
    as[0] = q;
    as[1] = k;
    as[2] = v;

    if (attn_mask)
    {
        as.push_back(RandomMat(k.h, q.h));
    }

    if (input_kvcache)
    {
        as.push_back(RandomMat(embed_dim / num_heads, k.h, num_heads));
        as.push_back(RandomMat(embed_dim / num_heads, k.h, num_heads));
    }
    else
    {
        as.push_back(ncnn::Mat());
        as.push_back(ncnn::Mat());
    }

    float epsilon = 0.005;

    int ret = test_layer("MultiHeadAttention", pd, weights, as, 3, epsilon);
    if (ret != 0)
    {
        fprintf(stderr, "test_multiheadattention_cross_kvcache failed q=(%d %d) k=(%d %d) v=(%d %d) embed_dim=%d num_heads=%d kdim=%d vdim=%d attn_mask=%d input_kvcache=%d\n", q.w, q.h, k.w, k.h, v.w, v.h, embed_dim, num_heads, kdim, vdim, attn_mask, input_kvcache);
    }

    return ret;
}

static int test_multiheadattention_cross_kvcache(const ncnn::Mat& q, const ncnn::Mat& k, const ncnn::Mat& v, int embed_dim, int num_heads, int attn_mask)
{
    return test_multiheadattention_cross_kvcache(q, k, v, embed_dim, num_heads, attn_mask, 0) || test_multiheadattention_cross_kvcache(q, k, v, embed_dim, num_heads, attn_mask, 1);
}

static int test_multiheadattention_self_kvcache_prefill(const ncnn::Mat& q, int embed_dim, int num_heads, int flag = 0)
{
    const int qdim = q.w;
    const int dst_seqlen = q.h;

    ncnn::ParamDict pd;
    pd.set(0, embed_dim);
    pd.set(1, num_heads);
    pd.set(2, embed_dim * qdim);
    pd.set(3, qdim);
    pd.set(4, qdim);
    pd.set(6, 0.7f / sqrtf(embed_dim / num_heads));
    pd.set(5, 1); // attn_mask
    pd.set(7, 1); // kv_cache

    std::vector<ncnn::Mat> weights(8);
    weights[0] = RandomMat(embed_dim * qdim);
    weights[1] = RandomMat(embed_dim);
    weights[2] = RandomMat(embed_dim * qdim);
    weights[3] = RandomMat(embed_dim);
    weights[4] = RandomMat(embed_dim * qdim);
    weights[5] = RandomMat(embed_dim);
    weights[6] = RandomMat(qdim * embed_dim);
    weights[7] = RandomMat(qdim);

    ncnn::Mat attn_mask(dst_seqlen, dst_seqlen);
    attn_mask.fill(0.f);
    for (int i = 0; i < dst_seqlen; i++)
    {
        for (int j = i + 1; j < dst_seqlen; j++)
        {
            attn_mask.row(i)[j] = -60000.f;
        }
    }

    std::vector<ncnn::Mat> as(4);
    as[0] = q;
    as[1] = attn_mask;
    as[2] = ncnn::Mat();
    as[3] = ncnn::Mat();

    float epsilon = 0.005;

    int ret = test_layer("MultiHeadAttention", pd, weights, as, 3, epsilon, flag);
    if (ret != 0)
    {
        fprintf(stderr, "test_multiheadattention_self_kvcache_prefill failed q=(%d %d) embed_dim=%d num_heads=%d\n", q.w, q.h, embed_dim, num_heads);
    }

    return ret;
}

static int test_multiheadattention_self_kvcache_decode(const ncnn::Mat& q, int embed_dim, int num_heads, int cur_seqlen)
{
    const int qdim = q.w;
    const int past_seqlen = q.h;

    ncnn::ParamDict pd;
    pd.set(0, embed_dim);
    pd.set(1, num_heads);
    pd.set(2, embed_dim * qdim);
    pd.set(3, qdim);
    pd.set(4, qdim);
    pd.set(6, 0.7f / sqrtf(embed_dim / num_heads));
    pd.set(5, 0); // attn_mask
    pd.set(7, 1); // kv_cache

    std::vector<ncnn::Mat> weights(8);
    weights[0] = RandomMat(embed_dim * qdim);
    weights[1] = RandomMat(embed_dim);
    weights[2] = RandomMat(embed_dim * qdim);
    weights[3] = RandomMat(embed_dim);
    weights[4] = RandomMat(embed_dim * qdim);
    weights[5] = RandomMat(embed_dim);
    weights[6] = RandomMat(qdim * embed_dim);
    weights[7] = RandomMat(qdim);

    std::vector<ncnn::Mat> as(3);
    as[0] = RandomMat(qdim, cur_seqlen);
    as[1] = RandomMat(embed_dim / num_heads, past_seqlen, num_heads);
    as[2] = RandomMat(embed_dim / num_heads, past_seqlen, num_heads);

    float epsilon = 0.005;

    int ret = test_layer("MultiHeadAttention", pd, weights, as, 3, epsilon);
    if (ret != 0)
    {
        fprintf(stderr, "test_multiheadattention_self_kvcache_decode failed q=(%d %d) embed_dim=%d num_heads=%d cur_seqlen=%d\n", q.w, q.h, embed_dim, num_heads, cur_seqlen);
    }

    return ret;
}

static int test_multiheadattention_singleton_mask(int embed_dim, int num_heads, int src_seqlen, int kv_cache)
{
    const int past_seqlen = kv_cache ? 5 : 0;
    const int dst_seqlen = past_seqlen + src_seqlen;

    ncnn::ParamDict pd;
    pd.set(0, embed_dim);
    pd.set(1, num_heads);
    pd.set(2, embed_dim * embed_dim);
    pd.set(3, embed_dim);
    pd.set(4, embed_dim);
    pd.set(5, 1); // attn_mask
    pd.set(7, kv_cache);

    std::vector<ncnn::Mat> weights(8);
    weights[0] = RandomMat(embed_dim * embed_dim);
    weights[1] = RandomMat(embed_dim);
    weights[2] = RandomMat(embed_dim * embed_dim);
    weights[3] = RandomMat(embed_dim);
    weights[4] = RandomMat(embed_dim * embed_dim);
    weights[5] = RandomMat(embed_dim);
    weights[6] = RandomMat(embed_dim * embed_dim);
    weights[7] = RandomMat(embed_dim);

    std::vector<ncnn::Mat> as(kv_cache ? 4 : 2);
    as[0] = RandomMat(embed_dim, src_seqlen);
    as[1] = RandomMat(dst_seqlen, src_seqlen, -2.f, 0.f);
    if (kv_cache)
    {
        as[2] = RandomMat(embed_dim / num_heads, past_seqlen, num_heads);
        as[3] = RandomMat(embed_dim / num_heads, past_seqlen, num_heads);
    }

    const int top_blob_count = kv_cache ? 3 : 1;
    const int typeindex = ncnn::layer_to_index("MultiHeadAttention");
    const float epsilon = 0.005f;

    std::vector<ncnn::Mat> reference;
    int ret = test_layer_naive(typeindex, pd, weights, as, top_blob_count, reference, 0);

    as[1] = as[1].reshape(dst_seqlen, src_seqlen, 1);
    std::vector<ncnn::Mat> reference_singleton;
    if (ret == 0)
        ret = test_layer_naive(typeindex, pd, weights, as, top_blob_count, reference_singleton, 0);
    if (ret == 0)
        ret = CompareMat(reference, reference_singleton, epsilon);
    if (ret == 0)
        ret = test_layer("MultiHeadAttention", pd, weights, as, top_blob_count, epsilon);

    if (ret != 0)
        fprintf(stderr, "test_multiheadattention_singleton_mask failed embed_dim=%d num_heads=%d src_seqlen=%d kv_cache=%d\n", embed_dim, num_heads, src_seqlen, kv_cache);

    return ret;
}

static int test_multiheadattention_singleton_mask()
{
    return 0
           || test_multiheadattention_singleton_mask(12, 4, 3, 0)
           || test_multiheadattention_singleton_mask(12, 4, 4, 0)
           || test_multiheadattention_singleton_mask(16, 4, 3, 0)
           || test_multiheadattention_singleton_mask(16, 4, 4, 0)
           || test_multiheadattention_singleton_mask(16, 4, 1, 1)
           || test_multiheadattention_singleton_mask(16, 4, 3, 1)
           || test_multiheadattention_singleton_mask(16, 4, 4, 1);
}

static int test_multiheadattention_0()
{
    return 0
           || test_multiheadattention_cross_kvcache(RandomMat(62, 66), RandomMat(32, 66), RandomMat(20, 66), 62, 2, 0)
           || test_multiheadattention_cross_kvcache(RandomMat(26, 64), RandomMat(32, 64), RandomMat(18, 64), 26, 2, 1)
           || test_multiheadattention_cross_kvcache(RandomMat(64, 128), RandomMat(64, 128), RandomMat(64, 128), 64, 4, 0)
           || test_multiheadattention_cross_kvcache(RandomMat(48, 127), RandomMat(64, 127), RandomMat(64, 127), 64, 16, 1)
           || test_multiheadattention_cross_kvcache(RandomMat(16, 128), RandomMat(44, 128), RandomMat(55, 128), 16, 2, 0)
           || test_multiheadattention_cross_kvcache(RandomMat(12, 128), RandomMat(44, 127), RandomMat(55, 127), 16, 4, 1)
           || test_multiheadattention_cross_kvcache(RandomMat(12, 17), RandomMat(28, 127), RandomMat(32, 127), 12, 3, 0)
           || test_multiheadattention_cross_kvcache(RandomMat(12, 17), RandomMat(28, 32), RandomMat(11, 32), 12, 3, 1);
}

static int test_multiheadattention_1()
{
    return 0
           || test_multiheadattention_self_kvcache_prefill(RandomMat(64, 128), 64, 4)
           || test_multiheadattention_self_kvcache_prefill(RandomMat(48, 127), 64, 8)
           || test_multiheadattention_self_kvcache_prefill(RandomMat(48, 127), 64, 8, TEST_LAYER_DISABLE_GPU_TESTING);
}

static int test_multiheadattention_2()
{
    return 0
           || test_multiheadattention_self_kvcache_decode(RandomMat(64, 128), 64, 4, 1)
           || test_multiheadattention_self_kvcache_decode(RandomMat(48, 127), 64, 8, 1)
           || test_multiheadattention_self_kvcache_decode(RandomMat(64, 128), 64, 4, 32);
}

#if NCNN_INT8
static int test_multiheadattention_int8_cross_kvcache(const ncnn::Mat& q, const ncnn::Mat& k, const ncnn::Mat& v, int embed_dim, int num_heads, int attn_mask, int input_kvcache)
{
    const int qdim = q.w;
    const int kdim = k.w;
    const int vdim = v.w;

    ncnn::ParamDict pd;
    pd.set(0, embed_dim);
    pd.set(1, num_heads);
    pd.set(2, embed_dim * qdim);
    pd.set(3, kdim);
    pd.set(4, vdim);
    pd.set(5, attn_mask);
    pd.set(6, 1.f / sqrtf(embed_dim / num_heads));
    pd.set(7, 1);  // kv_cache
    pd.set(18, 2); // int8_scale_term

    std::vector<ncnn::Mat> weights(12);
    weights[0] = RandomS8Mat(embed_dim * qdim);
    weights[1] = RandomMat(embed_dim);
    weights[2] = RandomS8Mat(embed_dim * kdim);
    weights[3] = RandomMat(embed_dim);
    weights[4] = RandomS8Mat(embed_dim * vdim);
    weights[5] = RandomMat(embed_dim);
    weights[6] = RandomS8Mat(qdim * embed_dim);
    weights[7] = RandomMat(qdim);
    weights[8] = RandomMat(embed_dim, 160.f, 200.f);
    weights[9] = RandomMat(embed_dim, 160.f, 200.f);
    weights[10] = RandomMat(embed_dim, 160.f, 200.f);
    weights[11] = RandomMat(1, 160.f, 200.f);

    std::vector<ncnn::Mat> as(3);
    as[0] = q;
    as[1] = k;
    as[2] = v;

    if (attn_mask)
    {
        as.push_back(RandomMat(k.h, q.h));
    }

    if (input_kvcache)
    {
        as.push_back(RandomMat(embed_dim / num_heads, k.h, num_heads));
        as.push_back(RandomMat(embed_dim / num_heads, k.h, num_heads));
    }
    else
    {
        as.push_back(ncnn::Mat());
        as.push_back(ncnn::Mat());
    }

    float epsilon = 0.1;

    int ret = test_layer("MultiHeadAttention", pd, weights, as, 3, epsilon);
    if (ret != 0)
    {
        fprintf(stderr, "test_multiheadattention_int8_cross_kvcache failed q=(%d %d) k=(%d %d) v=(%d %d) embed_dim=%d num_heads=%d kdim=%d vdim=%d attn_mask=%d input_kvcache=%d\n", q.w, q.h, k.w, k.h, v.w, v.h, embed_dim, num_heads, kdim, vdim, attn_mask, input_kvcache);
    }

    return ret;
}

static int test_multiheadattention_int8_cross_kvcache(const ncnn::Mat& q, const ncnn::Mat& k, const ncnn::Mat& v, int embed_dim, int num_heads, int attn_mask)
{
    return test_multiheadattention_int8_cross_kvcache(q, k, v, embed_dim, num_heads, attn_mask, 0) || test_multiheadattention_int8_cross_kvcache(q, k, v, embed_dim, num_heads, attn_mask, 1);
}

static int test_multiheadattention_int8_self_kvcache_prefill(const ncnn::Mat& q, int embed_dim, int num_heads)
{
    const int qdim = q.w;
    const int dst_seqlen = q.h;

    ncnn::ParamDict pd;
    pd.set(0, embed_dim);
    pd.set(1, num_heads);
    pd.set(2, embed_dim * qdim);
    pd.set(3, qdim);
    pd.set(4, qdim);
    pd.set(6, 0.7f / sqrtf(embed_dim / num_heads));
    pd.set(5, 1);  // attn_mask
    pd.set(7, 1);  // kv_cache
    pd.set(18, 2); // int8_scale_term

    std::vector<ncnn::Mat> weights(12);
    weights[0] = RandomS8Mat(embed_dim * qdim);
    weights[1] = RandomMat(embed_dim);
    weights[2] = RandomS8Mat(embed_dim * qdim);
    weights[3] = RandomMat(embed_dim);
    weights[4] = RandomS8Mat(embed_dim * qdim);
    weights[5] = RandomMat(embed_dim);
    weights[6] = RandomS8Mat(qdim * embed_dim);
    weights[7] = RandomMat(qdim);
    weights[8] = RandomMat(embed_dim, 160.f, 200.f);
    weights[9] = RandomMat(embed_dim, 160.f, 200.f);
    weights[10] = RandomMat(embed_dim, 160.f, 200.f);
    weights[11] = RandomMat(1, 160.f, 200.f);

    ncnn::Mat attn_mask(dst_seqlen, dst_seqlen);
    attn_mask.fill(0.f);
    for (int i = 0; i < dst_seqlen; i++)
    {
        for (int j = i + 1; j < dst_seqlen; j++)
        {
            attn_mask.row(i)[j] = -60000.f;
        }
    }

    std::vector<ncnn::Mat> as(4);
    as[0] = q;
    as[1] = attn_mask;
    as[2] = ncnn::Mat();
    as[3] = ncnn::Mat();

    float epsilon = 0.1;

    int ret = test_layer("MultiHeadAttention", pd, weights, as, 3, epsilon);
    if (ret != 0)
    {
        fprintf(stderr, "test_multiheadattention_int8_self_kvcache_prefill failed q=(%d %d) embed_dim=%d num_heads=%d\n", q.w, q.h, embed_dim, num_heads);
    }

    return ret;
}

static int test_multiheadattention_int8_self_kvcache_decode(const ncnn::Mat& q, int embed_dim, int num_heads, int cur_seqlen)
{
    const int qdim = q.w;
    const int past_seqlen = q.h;

    ncnn::ParamDict pd;
    pd.set(0, embed_dim);
    pd.set(1, num_heads);
    pd.set(2, embed_dim * qdim);
    pd.set(3, qdim);
    pd.set(4, qdim);
    pd.set(6, 0.7f / sqrtf(embed_dim / num_heads));
    pd.set(5, 0);  // attn_mask
    pd.set(7, 1);  // kv_cache
    pd.set(18, 2); // int8_scale_term

    std::vector<ncnn::Mat> weights(12);
    weights[0] = RandomS8Mat(embed_dim * qdim);
    weights[1] = RandomMat(embed_dim);
    weights[2] = RandomS8Mat(embed_dim * qdim);
    weights[3] = RandomMat(embed_dim);
    weights[4] = RandomS8Mat(embed_dim * qdim);
    weights[5] = RandomMat(embed_dim);
    weights[6] = RandomS8Mat(qdim * embed_dim);
    weights[7] = RandomMat(qdim);
    weights[8] = RandomMat(embed_dim, 160.f, 200.f);
    weights[9] = RandomMat(embed_dim, 160.f, 200.f);
    weights[10] = RandomMat(embed_dim, 160.f, 200.f);
    weights[11] = RandomMat(1, 160.f, 200.f);

    std::vector<ncnn::Mat> as(3);
    as[0] = RandomMat(qdim, cur_seqlen);
    as[1] = RandomMat(embed_dim / num_heads, past_seqlen, num_heads);
    as[2] = RandomMat(embed_dim / num_heads, past_seqlen, num_heads);

    float epsilon = 0.1;

    int ret = test_layer("MultiHeadAttention", pd, weights, as, 3, epsilon);
    if (ret != 0)
    {
        fprintf(stderr, "test_multiheadattention_int8_self_kvcache_decode failed q=(%d %d) embed_dim=%d num_heads=%d cur_seqlen=%d\n", q.w, q.h, embed_dim, num_heads, cur_seqlen);
    }

    return ret;
}

static int test_multiheadattention_3()
{
    return 0
           || test_multiheadattention_int8_cross_kvcache(RandomMat(62, 66), RandomMat(32, 66), RandomMat(20, 66), 62, 2, 0)
           || test_multiheadattention_int8_cross_kvcache(RandomMat(26, 64), RandomMat(32, 64), RandomMat(18, 64), 26, 2, 1)
           || test_multiheadattention_int8_cross_kvcache(RandomMat(64, 128), RandomMat(64, 128), RandomMat(64, 128), 64, 4, 0)
           || test_multiheadattention_int8_cross_kvcache(RandomMat(48, 127), RandomMat(64, 127), RandomMat(64, 127), 64, 16, 1)
           || test_multiheadattention_int8_cross_kvcache(RandomMat(16, 128), RandomMat(44, 128), RandomMat(55, 128), 16, 2, 0)
           || test_multiheadattention_int8_cross_kvcache(RandomMat(12, 128), RandomMat(44, 127), RandomMat(55, 127), 16, 4, 1)
           || test_multiheadattention_int8_cross_kvcache(RandomMat(12, 17), RandomMat(28, 127), RandomMat(32, 127), 12, 3, 0)
           || test_multiheadattention_int8_cross_kvcache(RandomMat(12, 17), RandomMat(28, 32), RandomMat(11, 32), 12, 3, 1);
}

static int test_multiheadattention_4()
{
    return 0
           || test_multiheadattention_int8_self_kvcache_prefill(RandomMat(64, 128), 64, 4)
           || test_multiheadattention_int8_self_kvcache_prefill(RandomMat(48, 127), 64, 8);
}

static int test_multiheadattention_5()
{
    return 0
           || test_multiheadattention_int8_self_kvcache_decode(RandomMat(64, 128), 64, 4, 1)
           || test_multiheadattention_int8_self_kvcache_decode(RandomMat(48, 127), 64, 8, 1)
           || test_multiheadattention_int8_self_kvcache_decode(RandomMat(64, 128), 64, 4, 32);
}
#endif

#if NCNN_BATCH
static int test_multiheadattention_kvcache_batch_rejected()
{
    ncnn::ParamDict pd;
    pd.set(0, 16);
    pd.set(1, 4);
    pd.set(2, 16 * 12);
    pd.set(3, 12);
    pd.set(4, 12);
    pd.set(7, 1); // kv_cache

    ncnn::Layer* op = ncnn::create_layer_cpu("MultiHeadAttention");
    if (!op)
        return -1;

    int ret = op->load_param(pd);
    if (ret == 0 && !op->support_batch)
        ret = -1;

    ncnn::Mat query;
    query.create(12, 1, 4u, 1, 2);

    std::vector<ncnn::Mat> bottoms(3);
    bottoms[0] = query;
    std::vector<ncnn::Mat> tops(3);

    if (ret == 0 && op->forward(bottoms, tops, ncnn::Option()) != -1)
        ret = -1;

    delete op;

    if (ret != 0)
        fprintf(stderr, "test_multiheadattention_kvcache_batch_rejected failed ret=%d\n", ret);

    return ret;
}
#endif // NCNN_BATCH

static int test_multiheadattention_shared_kv_kvcache(const std::vector<ncnn::Mat>& as, int attn_mask)
{
    const int embed_dim = 16;
    const int num_heads = 4;

    ncnn::ParamDict pd;
    pd.set(0, embed_dim);
    pd.set(1, num_heads);
    pd.set(2, embed_dim * embed_dim);
    pd.set(3, embed_dim);
    pd.set(4, embed_dim);
    pd.set(5, attn_mask);
    pd.set(6, 1.f / sqrtf(embed_dim / num_heads));
    pd.set(7, 1);

    // identity projections keep shared-input and cache indexing numerically observable
    std::vector<ncnn::Mat> weights(8);
    for (int projection = 0; projection < 4; projection++)
    {
        weights[projection * 2] = ncnn::Mat(embed_dim * embed_dim);
        weights[projection * 2].fill(0.f);
        for (int i = 0; i < embed_dim; i++)
            weights[projection * 2][i * embed_dim + i] = 1.f;
        weights[projection * 2 + 1] = ncnn::Mat(embed_dim);
        weights[projection * 2 + 1].fill(0.f);
    }

    const float epsilon = 0.005f;
    int ret = test_layer("MultiHeadAttention", pd, weights, as, 3, epsilon);
    if (ret != 0)
        fprintf(stderr, "test_multiheadattention_shared_kv_kvcache failed inputs=%d attn_mask=%d\n", (int)as.size(), attn_mask);
    return ret;
}

static int test_multiheadattention_shared_kv_kvcache_boundaries()
{
    ncnn::Mat q(16, 3);
    ncnn::Mat kv(16, 4);
    ncnn::Mat cached_k(4, 2, 4);
    ncnn::Mat cached_v(4, 2, 4);
    for (size_t i = 0; i < q.total(); i++)
        q[i] = ((int)(i % 11) - 5) * 0.0625f;
    for (size_t i = 0; i < kv.total(); i++)
        kv[i] = ((int)(i % 7) - 3) * 0.125f;
    for (size_t i = 0; i < cached_k.total(); i++)
        cached_k[i] = ((int)(i % 9) - 4) * 0.125f;
    for (size_t i = 0; i < cached_v.total(); i++)
        cached_v[i] = ((int)(i % 13) - 6) * 0.0625f;

    // cross attention reuses the nonempty two-token cache instead of appending shared kv
    std::vector<ncnn::Mat> unmasked(4);
    unmasked[0] = q;
    unmasked[1] = kv;
    unmasked[2] = cached_k;
    unmasked[3] = cached_v;

    // per-head masks include channel padding and distinct finite key penalties
    ncnn::Mat mask(2, 3, 4);
    mask.fill(0.f);
    for (int head = 0; head < 4; head++)
    {
        ncnn::Mat channel = mask.channel(head);
        for (int row = 0; row < 3; row++)
            channel.row(row)[head % 2] = -0.5f * (head + 1);
    }
    std::vector<ncnn::Mat> masked(5);
    masked[0] = q;
    masked[1] = kv;
    masked[2] = mask;
    masked[3] = cached_k;
    masked[4] = cached_v;

    return test_multiheadattention_shared_kv_kvcache(unmasked, 0)
           || test_multiheadattention_shared_kv_kvcache(masked, 1);
}

int main()
{
    SRAND(7767517);

    return 0
           || test_multiheadattention_0()
           || test_multiheadattention_1()
           || test_multiheadattention_2()
           || test_multiheadattention_singleton_mask()
#if NCNN_INT8
           || test_multiheadattention_3()
           || test_multiheadattention_4()
           || test_multiheadattention_5()
#endif
#if NCNN_BATCH
           || test_multiheadattention_kvcache_batch_rejected()
#endif // NCNN_BATCH
           || test_multiheadattention_shared_kv_kvcache_boundaries();
}