mirror of
https://github.com/Tencent/ncnn.git
synced 2026-10-19 06:21:39 +08:00
nihui
458 lines
16 KiB
C++
458 lines
16 KiB
C++
// Copyright 2026 Tencent
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
#include "testutil.h"
|
|
|
|
#include "layer_type.h"
|
|
#include "multiheadattention.h"
|
|
|
|
#include <string.h>
|
|
|
|
static void pack_signed_weight(unsigned char* ptr, int k, int bits, int q)
|
|
{
|
|
const unsigned int mask = (1u << bits) - 1u;
|
|
const unsigned int v = (unsigned int)q & mask;
|
|
const int bit_offset = k * bits;
|
|
|
|
for (int b = 0; b < bits; b++)
|
|
{
|
|
if (v & (1u << b))
|
|
{
|
|
const int out_bit = bit_offset + b;
|
|
ptr[out_bit / 8] |= (unsigned char)(1u << (out_bit % 8));
|
|
}
|
|
}
|
|
}
|
|
|
|
static int float2int_weight(float v, int bits)
|
|
{
|
|
const int qmax = (1 << (bits - 1)) - 1;
|
|
int q = (int)roundf(v);
|
|
if (q > qmax) q = qmax;
|
|
if (q < -qmax) q = -qmax;
|
|
return q;
|
|
}
|
|
|
|
static int weight_block_quantize_term(int bits, int block_size, int input_scale = 0)
|
|
{
|
|
const int block_size_code = block_size == 32 ? 0 : block_size == 64 ? 1 : block_size == 128 ? 2 : -1;
|
|
if ((bits != 4 && bits != 6 && bits != 8) || block_size_code < 0)
|
|
return 0;
|
|
|
|
return bits * 100 + (input_scale ? 10 : 0) + block_size_code;
|
|
}
|
|
|
|
static int weight_quantize_packed_k_bytes(int K, int bits)
|
|
{
|
|
return (K * bits + 7) / 8;
|
|
}
|
|
|
|
static ncnn::Mat make_input_scales(int K)
|
|
{
|
|
ncnn::Mat input_scales(K);
|
|
float* ptr = input_scales;
|
|
for (int k = 0; k < K; k++)
|
|
ptr[k] = 0.75f + (k % 5) * 0.15f;
|
|
|
|
return input_scales;
|
|
}
|
|
|
|
static ncnn::Mat scale_weight_by_input_scales(const ncnn::Mat& weight_data, const ncnn::Mat& input_scales, int inverse)
|
|
{
|
|
const int K = weight_data.w;
|
|
const int N = weight_data.h;
|
|
|
|
ncnn::Mat weight_data1(K, N);
|
|
const float* input_scale_ptr = input_scales;
|
|
|
|
for (int n = 0; n < N; n++)
|
|
{
|
|
const float* ptr = weight_data.row(n);
|
|
float* outptr = weight_data1.row(n);
|
|
|
|
for (int k = 0; k < K; k++)
|
|
outptr[k] = inverse ? ptr[k] / input_scale_ptr[k] : ptr[k] * input_scale_ptr[k];
|
|
}
|
|
|
|
return weight_data1;
|
|
}
|
|
|
|
static int quantize_weight(const ncnn::Mat& weight_data, int bits, int block_size, ncnn::Mat& weight_data_quantized, ncnn::Mat& weight_data_quantize_scales, ncnn::Mat& weight_data_dequantized)
|
|
{
|
|
const int K = weight_data.w;
|
|
const int N = weight_data.h;
|
|
const int packed_k_bytes = weight_quantize_packed_k_bytes(K, bits);
|
|
const int block_count = (K + block_size - 1) / block_size;
|
|
|
|
weight_data_quantized.create(packed_k_bytes, N, (size_t)1u);
|
|
weight_data_quantize_scales.create(block_count, N);
|
|
weight_data_dequantized.create(K, N);
|
|
if (weight_data_quantized.empty() || weight_data_quantize_scales.empty() || weight_data_dequantized.empty())
|
|
return -100;
|
|
|
|
memset(weight_data_quantized.data, 0, weight_data_quantized.total() * weight_data_quantized.elemsize);
|
|
|
|
const int qmax = (1 << (bits - 1)) - 1;
|
|
for (int n = 0; n < N; n++)
|
|
{
|
|
const float* ptr = weight_data.row(n);
|
|
float* scale_ptr = weight_data_quantize_scales.row(n);
|
|
unsigned char* qptr = weight_data_quantized.row<unsigned char>(n);
|
|
float* deqptr = weight_data_dequantized.row(n);
|
|
|
|
for (int b = 0; b < block_count; b++)
|
|
{
|
|
const int k0 = b * block_size;
|
|
const int max_kk = block_size < K - k0 ? block_size : K - k0;
|
|
|
|
float absmax = 0.f;
|
|
for (int k = 0; k < max_kk; k++)
|
|
{
|
|
const float v = (float)fabs(ptr[k0 + k]);
|
|
if (v > absmax)
|
|
absmax = v;
|
|
}
|
|
|
|
const float scale = absmax == 0.f ? 1.f : (float)qmax / absmax;
|
|
scale_ptr[b] = scale;
|
|
|
|
for (int k = 0; k < max_kk; k++)
|
|
{
|
|
const int q = float2int_weight(ptr[k0 + k] * scale, bits);
|
|
pack_signed_weight(qptr, k0 + k, bits, q);
|
|
deqptr[k0 + k] = q / scale;
|
|
}
|
|
}
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int make_mha_weights(int qdim, int kdim, int vdim, int embed_dim, int bits, int block_size, std::vector<ncnn::Mat>& weights, std::vector<ncnn::Mat>& ref_weights, int input_scale = 0)
|
|
{
|
|
ncnn::Mat q_weight_data = RandomMat(qdim, embed_dim, -1.f, 1.f);
|
|
ncnn::Mat k_weight_data = RandomMat(kdim, embed_dim, -1.f, 1.f);
|
|
ncnn::Mat v_weight_data = RandomMat(vdim, embed_dim, -1.f, 1.f);
|
|
ncnn::Mat out_weight_data = RandomMat(embed_dim, qdim, -1.f, 1.f);
|
|
|
|
ncnn::Mat q_input_scales;
|
|
ncnn::Mat k_input_scales;
|
|
ncnn::Mat v_input_scales;
|
|
ncnn::Mat out_input_scales;
|
|
|
|
if (input_scale)
|
|
{
|
|
q_input_scales = make_input_scales(qdim);
|
|
k_input_scales = make_input_scales(kdim);
|
|
v_input_scales = make_input_scales(vdim);
|
|
out_input_scales = make_input_scales(embed_dim);
|
|
|
|
q_weight_data = scale_weight_by_input_scales(q_weight_data, q_input_scales, 1);
|
|
k_weight_data = scale_weight_by_input_scales(k_weight_data, k_input_scales, 1);
|
|
v_weight_data = scale_weight_by_input_scales(v_weight_data, v_input_scales, 1);
|
|
out_weight_data = scale_weight_by_input_scales(out_weight_data, out_input_scales, 1);
|
|
}
|
|
|
|
ncnn::Mat q_weight_data_quantized;
|
|
ncnn::Mat k_weight_data_quantized;
|
|
ncnn::Mat v_weight_data_quantized;
|
|
ncnn::Mat out_weight_data_quantized;
|
|
ncnn::Mat q_weight_data_scales;
|
|
ncnn::Mat k_weight_data_scales;
|
|
ncnn::Mat v_weight_data_scales;
|
|
ncnn::Mat out_weight_data_scales;
|
|
ncnn::Mat q_weight_data_dequantized;
|
|
ncnn::Mat k_weight_data_dequantized;
|
|
ncnn::Mat v_weight_data_dequantized;
|
|
ncnn::Mat out_weight_data_dequantized;
|
|
|
|
int ret = quantize_weight(q_weight_data, bits, block_size, q_weight_data_quantized, q_weight_data_scales, q_weight_data_dequantized);
|
|
if (ret != 0)
|
|
return ret;
|
|
ret = quantize_weight(k_weight_data, bits, block_size, k_weight_data_quantized, k_weight_data_scales, k_weight_data_dequantized);
|
|
if (ret != 0)
|
|
return ret;
|
|
ret = quantize_weight(v_weight_data, bits, block_size, v_weight_data_quantized, v_weight_data_scales, v_weight_data_dequantized);
|
|
if (ret != 0)
|
|
return ret;
|
|
ret = quantize_weight(out_weight_data, bits, block_size, out_weight_data_quantized, out_weight_data_scales, out_weight_data_dequantized);
|
|
if (ret != 0)
|
|
return ret;
|
|
|
|
if (input_scale)
|
|
{
|
|
q_weight_data_dequantized = scale_weight_by_input_scales(q_weight_data_dequantized, q_input_scales, 0);
|
|
k_weight_data_dequantized = scale_weight_by_input_scales(k_weight_data_dequantized, k_input_scales, 0);
|
|
v_weight_data_dequantized = scale_weight_by_input_scales(v_weight_data_dequantized, v_input_scales, 0);
|
|
out_weight_data_dequantized = scale_weight_by_input_scales(out_weight_data_dequantized, out_input_scales, 0);
|
|
}
|
|
|
|
ncnn::Mat q_bias_data = RandomMat(embed_dim, -1.f, 1.f);
|
|
ncnn::Mat k_bias_data = RandomMat(embed_dim, -1.f, 1.f);
|
|
ncnn::Mat v_bias_data = RandomMat(embed_dim, -1.f, 1.f);
|
|
ncnn::Mat out_bias_data = RandomMat(qdim, -1.f, 1.f);
|
|
|
|
weights.resize(input_scale ? 16 : 12);
|
|
weights[0] = q_weight_data_quantized;
|
|
weights[1] = q_bias_data;
|
|
weights[2] = k_weight_data_quantized;
|
|
weights[3] = k_bias_data;
|
|
weights[4] = v_weight_data_quantized;
|
|
weights[5] = v_bias_data;
|
|
weights[6] = out_weight_data_quantized;
|
|
weights[7] = out_bias_data;
|
|
weights[8] = q_weight_data_scales;
|
|
weights[9] = k_weight_data_scales;
|
|
weights[10] = v_weight_data_scales;
|
|
weights[11] = out_weight_data_scales;
|
|
|
|
if (input_scale)
|
|
{
|
|
weights[12] = q_input_scales;
|
|
weights[13] = k_input_scales;
|
|
weights[14] = v_input_scales;
|
|
weights[15] = out_input_scales;
|
|
}
|
|
|
|
ref_weights.resize(8);
|
|
ref_weights[0] = q_weight_data_dequantized.reshape(embed_dim * qdim);
|
|
ref_weights[1] = q_bias_data;
|
|
ref_weights[2] = k_weight_data_dequantized.reshape(embed_dim * kdim);
|
|
ref_weights[3] = k_bias_data;
|
|
ref_weights[4] = v_weight_data_dequantized.reshape(embed_dim * vdim);
|
|
ref_weights[5] = v_bias_data;
|
|
ref_weights[6] = out_weight_data_dequantized.reshape(qdim * embed_dim);
|
|
ref_weights[7] = out_bias_data;
|
|
|
|
return 0;
|
|
}
|
|
|
|
static ncnn::ParamDict make_mha_param(int qdim, int kdim, int vdim, int embed_dim, int num_heads, int attn_mask, int kv_cache, int quantize_term)
|
|
{
|
|
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, 0.7f / sqrtf(embed_dim / num_heads));
|
|
pd.set(7, kv_cache);
|
|
if (quantize_term)
|
|
pd.set(18, quantize_term);
|
|
|
|
return pd;
|
|
}
|
|
|
|
static int run_mha_layer(const ncnn::ParamDict& pd, const std::vector<ncnn::Mat>& weights, const std::vector<ncnn::Mat>& inputs, int top_blob_count, std::vector<ncnn::Mat>& outputs)
|
|
{
|
|
ncnn::Option opt;
|
|
opt.num_threads = 2;
|
|
opt.use_packing_layout = false;
|
|
opt.use_fp16_packed = false;
|
|
opt.use_fp16_storage = false;
|
|
opt.use_fp16_arithmetic = false;
|
|
opt.use_bf16_storage = false;
|
|
|
|
return test_layer_cpu(ncnn::LayerType::MultiHeadAttention, pd, weights, opt, inputs, top_blob_count, outputs, std::vector<ncnn::Mat>(), TEST_LAYER_DISABLE_GPU_TESTING);
|
|
}
|
|
|
|
static int test_multiheadattention_block_quant(int qdim, int kdim, int vdim, int embed_dim, int num_heads, int bits, int block_size, int attn_mask, int input_scale = 0)
|
|
{
|
|
std::vector<ncnn::Mat> weights;
|
|
std::vector<ncnn::Mat> ref_weights;
|
|
int ret = make_mha_weights(qdim, kdim, vdim, embed_dim, bits, block_size, weights, ref_weights, input_scale);
|
|
if (ret != 0)
|
|
return ret;
|
|
|
|
std::vector<ncnn::Mat> inputs(3);
|
|
inputs[0] = RandomMat(qdim, 5, -1.f, 1.f);
|
|
inputs[1] = RandomMat(kdim, 6, -1.f, 1.f);
|
|
inputs[2] = RandomMat(vdim, 6, -1.f, 1.f);
|
|
|
|
if (attn_mask)
|
|
inputs.push_back(RandomMat(6, 5, -1.f, 0.f));
|
|
|
|
const int quantize_term = weight_block_quantize_term(bits, block_size, input_scale);
|
|
const ncnn::ParamDict pd = make_mha_param(qdim, kdim, vdim, embed_dim, num_heads, attn_mask, 0, quantize_term);
|
|
const ncnn::ParamDict ref_pd = make_mha_param(qdim, kdim, vdim, embed_dim, num_heads, attn_mask, 0, 0);
|
|
|
|
std::vector<ncnn::Mat> outputs;
|
|
std::vector<ncnn::Mat> refs;
|
|
ret = run_mha_layer(pd, weights, inputs, 1, outputs);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant failed ret=%d qdim=%d kdim=%d vdim=%d embed_dim=%d bits=%d block_size=%d attn_mask=%d input_scale=%d\n", ret, qdim, kdim, vdim, embed_dim, bits, block_size, attn_mask, input_scale);
|
|
return ret;
|
|
}
|
|
|
|
ret = run_mha_layer(ref_pd, ref_weights, inputs, 1, refs);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant reference failed ret=%d qdim=%d kdim=%d vdim=%d embed_dim=%d\n", ret, qdim, kdim, vdim, embed_dim);
|
|
return ret;
|
|
}
|
|
|
|
ret = CompareMat(outputs, refs, 0.001f);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant compare failed qdim=%d kdim=%d vdim=%d embed_dim=%d bits=%d block_size=%d attn_mask=%d input_scale=%d\n", qdim, kdim, vdim, embed_dim, bits, block_size, attn_mask, input_scale);
|
|
return ret;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int test_multiheadattention_block_quant_kvcache(int attn_mask = 0, int input_scale = 0)
|
|
{
|
|
const int qdim = 10;
|
|
const int embed_dim = 8;
|
|
const int num_heads = 2;
|
|
const int bits = 4;
|
|
const int block_size = 64;
|
|
|
|
std::vector<ncnn::Mat> weights;
|
|
std::vector<ncnn::Mat> ref_weights;
|
|
int ret = make_mha_weights(qdim, qdim, qdim, embed_dim, bits, block_size, weights, ref_weights, input_scale);
|
|
if (ret != 0)
|
|
return ret;
|
|
|
|
std::vector<ncnn::Mat> inputs(attn_mask ? 4 : 3);
|
|
inputs[0] = RandomMat(qdim, 3, -1.f, 1.f);
|
|
if (attn_mask)
|
|
{
|
|
inputs[1] = RandomMat(8, 3, -1.f, 0.f);
|
|
inputs[2] = RandomMat(5, embed_dim, -1.f, 1.f);
|
|
inputs[3] = RandomMat(5, embed_dim, -1.f, 1.f);
|
|
}
|
|
else
|
|
{
|
|
inputs[1] = RandomMat(5, embed_dim, -1.f, 1.f);
|
|
inputs[2] = RandomMat(5, embed_dim, -1.f, 1.f);
|
|
}
|
|
|
|
const int quantize_term = weight_block_quantize_term(bits, block_size, input_scale);
|
|
const ncnn::ParamDict pd = make_mha_param(qdim, qdim, qdim, embed_dim, num_heads, attn_mask, 1, quantize_term);
|
|
const ncnn::ParamDict ref_pd = make_mha_param(qdim, qdim, qdim, embed_dim, num_heads, attn_mask, 1, 0);
|
|
|
|
std::vector<ncnn::Mat> outputs;
|
|
std::vector<ncnn::Mat> refs;
|
|
ret = run_mha_layer(pd, weights, inputs, 3, outputs);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant_kvcache failed ret=%d attn_mask=%d input_scale=%d\n", ret, attn_mask, input_scale);
|
|
return ret;
|
|
}
|
|
|
|
ret = run_mha_layer(ref_pd, ref_weights, inputs, 3, refs);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant_kvcache reference failed ret=%d attn_mask=%d input_scale=%d\n", ret, attn_mask, input_scale);
|
|
return ret;
|
|
}
|
|
|
|
ret = CompareMat(outputs, refs, 0.001f);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant_kvcache compare failed attn_mask=%d input_scale=%d\n", attn_mask, input_scale);
|
|
return ret;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int test_multiheadattention_block_quant_cross_kvcache(int attn_mask = 0, int input_scale = 0)
|
|
{
|
|
const int qdim = 65;
|
|
const int kdim = 33;
|
|
const int vdim = 49;
|
|
const int embed_dim = 64;
|
|
const int num_heads = 4;
|
|
const int bits = 6;
|
|
const int block_size = 32;
|
|
|
|
std::vector<ncnn::Mat> weights;
|
|
std::vector<ncnn::Mat> ref_weights;
|
|
int ret = make_mha_weights(qdim, kdim, vdim, embed_dim, bits, block_size, weights, ref_weights, input_scale);
|
|
if (ret != 0)
|
|
return ret;
|
|
|
|
std::vector<ncnn::Mat> inputs(attn_mask ? 6 : 5);
|
|
inputs[0] = RandomMat(qdim, 3, -1.f, 1.f);
|
|
inputs[1] = RandomMat(kdim, 2, -1.f, 1.f);
|
|
inputs[2] = RandomMat(vdim, 2, -1.f, 1.f);
|
|
if (attn_mask)
|
|
{
|
|
inputs[3] = RandomMat(5, 3, -1.f, 0.f);
|
|
inputs[4] = RandomMat(5, embed_dim, -1.f, 1.f);
|
|
inputs[5] = RandomMat(5, embed_dim, -1.f, 1.f);
|
|
}
|
|
else
|
|
{
|
|
inputs[3] = RandomMat(5, embed_dim, -1.f, 1.f);
|
|
inputs[4] = RandomMat(5, embed_dim, -1.f, 1.f);
|
|
}
|
|
|
|
const int quantize_term = weight_block_quantize_term(bits, block_size, input_scale);
|
|
const ncnn::ParamDict pd = make_mha_param(qdim, kdim, vdim, embed_dim, num_heads, attn_mask, 1, quantize_term);
|
|
const ncnn::ParamDict ref_pd = make_mha_param(qdim, kdim, vdim, embed_dim, num_heads, attn_mask, 1, 0);
|
|
|
|
std::vector<ncnn::Mat> outputs;
|
|
std::vector<ncnn::Mat> refs;
|
|
ret = run_mha_layer(pd, weights, inputs, 3, outputs);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant_cross_kvcache failed ret=%d attn_mask=%d input_scale=%d\n", ret, attn_mask, input_scale);
|
|
return ret;
|
|
}
|
|
|
|
ret = run_mha_layer(ref_pd, ref_weights, inputs, 3, refs);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant_cross_kvcache reference failed ret=%d attn_mask=%d input_scale=%d\n", ret, attn_mask, input_scale);
|
|
return ret;
|
|
}
|
|
|
|
ret = CompareMat(outputs, refs, 0.001f);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant_cross_kvcache compare failed attn_mask=%d input_scale=%d\n", attn_mask, input_scale);
|
|
return ret;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
int main()
|
|
{
|
|
SRAND(7767517);
|
|
|
|
#if !NCNN_WEIGHT_QUANT
|
|
ncnn::ParamDict pd = make_mha_param(5, 5, 5, 4, 2, 0, 0, 410);
|
|
|
|
ncnn::MultiHeadAttention mha;
|
|
if (mha.load_param(pd) == 0)
|
|
{
|
|
fprintf(stderr, "test_multiheadattention_block_quant failed NCNN_WEIGHT_QUANT=OFF accepted weight block quantization\n");
|
|
return -1;
|
|
}
|
|
|
|
return 0;
|
|
#else
|
|
return 0
|
|
|| test_multiheadattention_block_quant(13, 9, 11, 8, 2, 4, 32, 0)
|
|
|| test_multiheadattention_block_quant(10, 10, 10, 8, 2, 6, 64, 1)
|
|
|| test_multiheadattention_block_quant(12, 7, 9, 8, 2, 8, 128, 0)
|
|
|| test_multiheadattention_block_quant(13, 9, 11, 8, 2, 4, 64, 1, 1)
|
|
|| test_multiheadattention_block_quant(65, 65, 65, 64, 4, 6, 64, 1)
|
|
|| test_multiheadattention_block_quant(65, 33, 49, 64, 4, 4, 32, 0, 1)
|
|
|| test_multiheadattention_block_quant_kvcache()
|
|
|| test_multiheadattention_block_quant_kvcache(0, 1)
|
|
|| test_multiheadattention_block_quant_kvcache(1)
|
|
|| test_multiheadattention_block_quant_kvcache(1, 1)
|
|
|| test_multiheadattention_block_quant_cross_kvcache()
|
|
|| test_multiheadattention_block_quant_cross_kvcache(1)
|
|
|| test_multiheadattention_block_quant_cross_kvcache(0, 1);
|
|
#endif
|
|
}
|