mirror of
https://github.com/Tencent/ncnn.git
synced 2026-10-18 21:54:39 +08:00
風的彷徨
ab3790faca
Keep persistent KV caches on the blob allocator, broadcast single-channel attention masks across heads on CPU and Vulkan paths, and add long multi-token cache regressions.
594 lines
20 KiB
C++
594 lines
20 KiB
C++
// 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();
|
|
}
|