diff --git a/CMakeLists.txt b/CMakeLists.txt index 43bd9427a36b..407c8d98e243 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -90,6 +90,7 @@ option(NCNN_ASAN "build for address sanitizer" OFF) option(NCNN_BUILD_BENCHMARK "build benchmark" ON) option(NCNN_PYTHON "build python api" OFF) option(NCNN_INT8 "int8 inference" ON) +option(NCNN_WEIGHT_QUANT "weight quantized inference" ON) option(NCNN_BF16 "bf16 inference" ON) option(NCNN_FORCE_INLINE "force inline some function" ON) diff --git a/docs/developer-guide/operation-param-weight-table.md b/docs/developer-guide/operation-param-weight-table.md index aa5c99adf5e6..5a035ed9a576 100644 --- a/docs/developer-guide/operation-param-weight-table.md +++ b/docs/developer-guide/operation-param-weight-table.md @@ -131,6 +131,25 @@ ||2|expand_c|0| ||3|axes|[ ]| |Flatten||| +|Gemm|0|alpha|1.f|A_data B_data C_data A_data_int8_scales B_data_int8_scales/B_data_quantize_scales B_data_input_scales| +||1|beta|1.f| +||2|transA|0| +||3|transB|0| +||4|constantA|0| +||5|constantB|0| +||6|constantC|0| +||7|constantM|0| +||8|constantN|0| +||9|constantK|0| +||10|constant_broadcast_type_C|0| +||11|output_N1M|0| +||12|output_elempack|0| +||13|output_elemtype|0| +||14|output_transpose|0| +||18|quantize_term|0| +||20|constant_TILE_M|0| +||21|constant_TILE_N|0| +||22|constant_TILE_K|0| |HardSigmoid|0|alpha|0.2f|| ||1|beta|0.5f| |HardSwish|0|alpha|0.2f|| diff --git a/docs/developer-guide/operators.md b/docs/developer-guide/operators.md index c83b5c2a226c..6e153feefd74 100644 --- a/docs/developer-guide/operators.md +++ b/docs/developer-guide/operators.md @@ -1151,7 +1151,7 @@ y = (gemm(a, b) + c * beta) * alpha | 12 | output_elempack | int | 0 | | | 13 | output_elemtype | int | 0 | | | 14 | output_transpose | int| 0 | | -| 18 | int8_scale_term | int | 0 | | +| 18 | quantize_term | int | 0 | 0=no quant, nonzero non-block value below 400 is legacy Gemm int8 except obsolete 4/5/6, 400/401/402=int4 block32/64/128, 410/411/412=int4 block32/64/128 with input scale, 600/601/602=int6 block32/64/128, 610/611/612=int6 block32/64/128 with input scale, 800/801/802=int8 block32/64/128, 810/811/812=int8 block32/64/128 with input scale | | 20 | constant_TILE_M | int | 0 | | | 21 | constant_TILE_N | int | 0 | | | 22 | constant_TILE_K | int | 0 | | @@ -1163,6 +1163,16 @@ y = (gemm(a, b) + c * beta) * alpha | C_data | float | [1], [M] or [N] or [1, M] or [N,1] or [N, M] | | A_data_int8_scales| float | [M] | | B_data_int8_scales| float | [1] | +| B_data_quantize_scales| float | [ceil(K / block_size), N] for block quantized constant B | +| B_data_input_scales| float | [K] for block quantized constant B with input scale | + +For weight-only block quantized Gemm: + +* `constantA=0`, `constantB=1`, `transA=0`, `transB=1` +* output is fp32 pack1 with `output_N1M=0`, `output_elempack=0`, `output_transpose=0` +* `B_data` is tagged int8 bytes with shape `[ceil(K * weight_bits / 8), N]` +* scales and optional input scales are raw fp32 data +* packing is signed symmetric scale-only, no zero point # GridSample ``` @@ -1575,7 +1585,7 @@ y = affine(out) | 5 | attn_mask | int | 0 | | | 6 | scale | float | 1.f / sqrt(embed_dim / num_heads) | | | 7 | kv_cache | int | 0 | | -| 18 | int8_scale_term | int | 0 | | +| 18 | quantize_term | int | 0 | 0=no quant, nonzero non-block value below 400 is legacy MultiHeadAttention int8 except obsolete 4/5/6, 400/401/402=int4 block32/64/128, 410/411/412=int4 block32/64/128 with input scale, 600/601/602=int6 block32/64/128, 610/611/612=int6 block32/64/128 with input scale, 800/801/802=int8 block32/64/128, 810/811/812=int8 block32/64/128 with input scale | | weight | type | shape | | ------------- | ----- | --------------------- | @@ -1591,6 +1601,16 @@ y = affine(out) | k_weight_data_int8_scales| float | [embed_dim] | | v_weight_data_int8_scales| float | [embed_dim] | | out_weight_data_int8_scales| float | [1] | +| q_weight_data_quantize_scales| float | [ceil(qdim / block_size), embed_dim] for block quantized weight | +| k_weight_data_quantize_scales| float | [ceil(kdim / block_size), embed_dim] for block quantized weight | +| v_weight_data_quantize_scales| float | [ceil(vdim / block_size), embed_dim] for block quantized weight | +| out_weight_data_quantize_scales| float | [ceil(embed_dim / block_size), qdim] for block quantized weight | +| q_weight_data_input_scales| float | [qdim] for block quantized weight with input scale | +| k_weight_data_input_scales| float | [kdim] for block quantized weight with input scale | +| v_weight_data_input_scales| float | [vdim] for block quantized weight with input scale | +| out_weight_data_input_scales| float | [embed_dim] for block quantized weight with input scale | + +Weight-only block quantized MultiHeadAttention stores q/k/v/out weights as tagged int8 bytes with signed symmetric int4/int6/int8 packing. Activations and output are fp32. Input-scale terms add per-input-channel multipliers for q/k/v/out. # MVN ``` diff --git a/docs/how-to-use-and-FAQ/quantized-int8-inference.md b/docs/how-to-use-and-FAQ/quantized-int8-inference.md index 10ab96843c92..002810dcee3d 100644 --- a/docs/how-to-use-and-FAQ/quantized-int8-inference.md +++ b/docs/how-to-use-and-FAQ/quantized-int8-inference.md @@ -101,6 +101,120 @@ ncnn2table can generate static weight scales without a calibration dataset for R ./ncnn2int8 mobilenet-opt.param mobilenet-opt.bin mobilenet-int8.param mobilenet-int8.bin mobilenet.table ``` +## Weight-only block quantized Gemm and MultiHeadAttention + +LLM-oriented `Gemm` and `MultiHeadAttention` weight-only block quantization is separate from the post training int8 activation/weight inference flow above. It stores weights as signed symmetric int4/int6/int8 blocks with one fp32 scale per K block and output remains fp32. + +The recommended workflow mirrors `ncnn2table` and `ncnn2int8`: + +```shell +./ncnnllm2table in.param in.bin model.llm.table method=minmax bits=6 block=64 +./ncnnllm2int468 in.param in.bin out.param out.bin model.llm.table +``` + +`ncnnllm2table` options follow the same trailing key-value style as `ncnn2table`: + +| key | values | default | description | +| --- | ------ | ------- | ----------- | +| `method` | `minmax`, `mseclip`, `awq`, `gptq` | `minmax` | offline quantization method | +| `bits` | `4`, `6`, `8` | `6` | signed weight bit width | +| `block` | `32`, `64`, `128` | `64` | K block size | +| `thread` | positive integer | `1` | worker threads | +| `type` | `1` | `1` | calibration input file type, npy only | +| `shape` | `[w,h,...]` | none | calibration input shape in ncnn order | +| `awq_steps` | non-negative integer | `20` | AWQ input-scale search steps | +| `awq_samples` | positive integer | `128` | max activation rows used by AWQ search | +| `awq_max_scale` | float > 1 | `16` | AWQ input scale clamp | +| `awq_inner` | `minmax`, `mseclip` | `minmax` | weight scale method used inside AWQ | +| `gptq_samples` | positive integer | `128` | max activation rows used by GPTQ | +| `gptq_damp` | non-negative float | `0.01` | GPTQ Hessian damping ratio | + +`method=minmax` uses per-block absmax scaling. + +`method=mseclip` searches clipped absmax candidates and picks the scale with the smallest block weight reconstruction error. It is calibration-free and still exports the same symmetric scale-only runtime format. It is not AWQ and does not add activation rescaling metadata. + +`method=awq` requires calibration data. It exports block scale rows plus `_input_scale` rows. Runtime still computes signed symmetric block dequantization; `method=awq` is only table provenance. + +```shell +./ncnnllm2table in.param in.bin calib.list awq.llm.table method=awq bits=4 block=64 type=1 shape=[...] +./ncnnllm2int468 in.param in.bin awq.param awq.bin awq.llm.table +``` + +`method=gptq` requires calibration data. It uses Fixed-Scale GPTQ: the runtime block scale is chosen from the original weight block, then block-local Hessian error compensation writes exact qweight sidecar files. The converter imports those qvalues directly instead of requantizing fp32 weights. + +```shell +./ncnnllm2table in.param in.bin calib.list gptq.llm.table method=gptq bits=4 block=128 type=1 shape=[...] +./ncnnllm2int468 in.param in.bin gptq.param gptq.bin gptq.llm.table +``` + +The calibration list format follows `ncnn2table`: one npy path per line for one input, or comma-separated list files for multiple inputs. All input lists must have the same sample count. + +The llm table uses one line for each quantized weight: + +```text +gemm_name_param_1 bits=4 block=64 method=mseclip scale0 scale1 ... +mha_name_param_0 bits=4 block=64 method=mseclip scale0 scale1 ... +mha_name_param_1 bits=4 block=64 method=mseclip scale0 scale1 ... +mha_name_param_2 bits=4 block=64 method=mseclip scale0 scale1 ... +mha_name_param_3 bits=4 block=64 method=mseclip scale0 scale1 ... +``` + +For `MultiHeadAttention`, `_param_0/_param_1/_param_2/_param_3` are q/k/v/out weights. All four rows must use the same bits and block size. + +`bits` and `block` are row-local, so one table can mix int4/int6/int8 and block32/block64/block128 for different layers. `method` is written for readability and logging. Unknown `key=value` fields are rejected. + +`out.llm.table` is a text file and may be edited before conversion. A missing Gemm row skips that Gemm. For `MultiHeadAttention`, all four q/k/v/out rows must exist together; deleting all four skips the layer. + +The scale-only table format is: + +```text +layer_param_x bits=4/6/8 block=32/64/128 method=minmax/mseclip/awq scale0 scale1 ... +``` + +Exact qweight import uses: + +```text +gemm_name_param_1 bits=4 block=128 method=gptq qweight=gemm.qweight scale0 scale1 ... +``` + +The `qweight` path is relative to the table directory unless absolute. qweight data is packed in runtime layout: output rows `N`, input columns `K`, signed int4/int6/int8. The converter validates qweight byte count, tail padding bits, unsupported negative sentinel values, and scale coefficient count. + +The runtime quantize term is stored in param id `18`. + +```text +400/401/402 int4 block=32/64/128 +600/601/602 int6 block=32/64/128 +800/801/802 int8 block=32/64/128 +``` + +Optional per-input-channel multipliers can be stored by adding separate `_input_scale` rows. They are applied to the current input channel inside the dot product: + +```text +sum += (x[k] * input_scale[k]) * (qweight[n,k] / block_scale[n,k/block]) +``` + +```text +gemm_name_param_1_input_scale method=awq coeff0 coeff1 ... +mha_name_param_0_input_scale method=awq coeff0 coeff1 ... +mha_name_param_1_input_scale method=awq coeff0 coeff1 ... +mha_name_param_2_input_scale method=awq coeff0 coeff1 ... +mha_name_param_3_input_scale method=awq coeff0 coeff1 ... +``` + +For `Gemm`, the input scale count is `constantK`. For `MultiHeadAttention`, q/k/v/out input scale counts are qdim/kdim/vdim/embed_dim, and either all four rows exist or none exist. `ncnnllm2table method=minmax/mseclip` does not write input scale rows; `method=awq` writes them. + +Input-scale models use the same bits/block terms with tens digit `1`: `410/411/412`, `610/611/612`, and `810/811/812`. + +For quick conversion without saving a table, `ncnnllm2int468` can still compute scales directly: + +```shell +./ncnnllm2int468 in.param in.bin out.param out.bin method=minmax bits=6 block=64 +``` + +AWQ and GPTQ are offline tool methods and must not become runtime `quantize_term` values. SmoothQuant is not implemented here; it needs separate offline graph rewrite or scale folding rules. + +This format does not use zero point or asymmetric dequantization metadata. + ## use ncnn int8 inference the ncnn library would use int8 inference automatically, nothing changed in your code diff --git a/src/layer/arm/gemm_arm.cpp b/src/layer/arm/gemm_arm.cpp index 0c57613bbaa9..38528293aa7f 100644 --- a/src/layer/arm/gemm_arm.cpp +++ b/src/layer/arm/gemm_arm.cpp @@ -4584,8 +4584,13 @@ static int gemm_AT_BT_arm(const Mat& AT, const Mat& BT, const Mat& C, Mat& top_b int Gemm_arm::create_pipeline(const Option& opt) { + if (weight_block_quantize) + { + return 0; + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return create_pipeline_int8(opt); } @@ -4742,8 +4747,13 @@ int Gemm_arm::create_pipeline(const Option& opt) int Gemm_arm::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const { + if (weight_block_quantize) + { + return Gemm::forward(bottom_blobs, top_blobs, opt); + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return forward_int8(bottom_blobs, top_blobs, opt); } diff --git a/src/layer/arm/multiheadattention_arm.cpp b/src/layer/arm/multiheadattention_arm.cpp index 46cefd4822e7..f76e1a54e862 100644 --- a/src/layer/arm/multiheadattention_arm.cpp +++ b/src/layer/arm/multiheadattention_arm.cpp @@ -4,6 +4,7 @@ #include "multiheadattention_arm.h" #include "cpu.h" +#include "gemm.h" #include "layer_type.h" namespace ncnn { @@ -34,6 +35,9 @@ MultiHeadAttention_arm::MultiHeadAttention_arm() int MultiHeadAttention_arm::create_pipeline(const Option& _opt) { + if (weight_block_quantize) + return 0; + Option opt = _opt; if (int8_scale_term) { @@ -265,6 +269,9 @@ int MultiHeadAttention_arm::create_pipeline(const Option& _opt) int MultiHeadAttention_arm::destroy_pipeline(const Option& _opt) { + if (weight_block_quantize) + return 0; + Option opt = _opt; if (int8_scale_term) { @@ -330,6 +337,9 @@ int MultiHeadAttention_arm::destroy_pipeline(const Option& _opt) int MultiHeadAttention_arm::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& _opt) const { + if (weight_block_quantize) + return MultiHeadAttention::forward(bottom_blobs, top_blobs, _opt); + int q_blob_i = 0; int k_blob_i = 0; int v_blob_i = 0; diff --git a/src/layer/gemm.cpp b/src/layer/gemm.cpp index 442472f0f19a..f22a04f1ed7e 100644 --- a/src/layer/gemm.cpp +++ b/src/layer/gemm.cpp @@ -3,12 +3,93 @@ #include "gemm.h" +#include +#include + namespace ncnn { +static bool gemm_is_weight_block_quantize(int quantize_term) +{ + const int weight_bits = quantize_term / 100; + const int format_code = quantize_term % 100 / 10; + const int block_size_code = quantize_term % 10; + + if (weight_bits != 4 && weight_bits != 6 && weight_bits != 8) + return false; + + if (format_code != 0 && format_code != 1) + return false; + + if (block_size_code < 0 || block_size_code > 2) + return false; + + return true; +} + +static bool gemm_weight_quantize_has_input_scale(int quantize_term) +{ + if (!gemm_is_weight_block_quantize(quantize_term)) + return false; + + return quantize_term % 100 / 10 == 1; +} + +static int gemm_weight_quantize_bits(int quantize_term) +{ + return gemm_is_weight_block_quantize(quantize_term) ? quantize_term / 100 : 0; +} + +static int gemm_weight_quantize_block_size(int quantize_term) +{ + if (!gemm_is_weight_block_quantize(quantize_term)) + return 0; + + const int block_size_code = quantize_term % 10; + return block_size_code == 0 ? 32 : block_size_code == 1 ? 64 : 128; +} + +static int gemm_weight_quantize_packed_k_bytes(int constantK, int weight_bits) +{ + if (constantK <= 0 || weight_bits <= 0) + return -1; + + const size_t packed_k_bytes = ((size_t)constantK * weight_bits + 7) / 8; + if (packed_k_bytes > (size_t)INT_MAX) + return -1; + + return (int)packed_k_bytes; +} + +#if NCNN_WEIGHT_QUANT +static inline int gemm_weight_block_quantize_sign_extend(int v, int bits) +{ + const int sign_bit = 1 << (bits - 1); + return (v ^ sign_bit) - sign_bit; +} + +static inline int gemm_weight_block_quantize_unpack(const unsigned char* ptr, int k, int bits, int packed_k_bytes) +{ + const int bit_offset = k * bits; + const int byte_offset = bit_offset / 8; + const int bit_shift = bit_offset % 8; + + unsigned int v = ptr[byte_offset]; + if (byte_offset + 1 < packed_k_bytes) + v |= (unsigned int)ptr[byte_offset + 1] << 8; + if (byte_offset + 2 < packed_k_bytes) + v |= (unsigned int)ptr[byte_offset + 2] << 16; + + const int mask = (1 << bits) - 1; + return gemm_weight_block_quantize_sign_extend((v >> bit_shift) & mask, bits); +} +#endif // NCNN_WEIGHT_QUANT + Gemm::Gemm() { one_blob_only = false; support_inplace = false; + + weight_block_quantize = 0; } int Gemm::load_param(const ParamDict& pd) @@ -28,12 +109,50 @@ int Gemm::load_param(const ParamDict& pd) output_elempack = pd.get(12, 0); output_elemtype = pd.get(13, 0); output_transpose = pd.get(14, 0); - int8_scale_term = pd.get(18, 0); + quantize_term = pd.get(18, 0); + weight_block_quantize = gemm_is_weight_block_quantize(quantize_term); constant_TILE_M = pd.get(20, 0); constant_TILE_N = pd.get(21, 0); constant_TILE_K = pd.get(22, 0); - if (int8_scale_term) + if (quantize_term == 4 || quantize_term == 5 || quantize_term == 6) + { + NCNN_LOGE("Gemm quantize_term %d is an obsolete block quantization prototype value", quantize_term); + return -1; + } + + if (quantize_term >= 400 && !weight_block_quantize) + { + NCNN_LOGE("Gemm unsupported quantize_term %d", quantize_term); + return -1; + } + + if (weight_block_quantize) + { +#if NCNN_WEIGHT_QUANT + if (constantA != 0 || constantB != 1 || transA != 0 || transB != 1) + { + NCNN_LOGE("Gemm weight block quantization only supports constantA=0 constantB=1 transA=0 transB=1"); + return -1; + } + + if (output_N1M != 0 || output_elempack != 0 || (output_elemtype != 0 && output_elemtype != 1) || output_transpose != 0) + { + NCNN_LOGE("Gemm weight block quantization only supports fp32 pack1 non-transposed output"); + return -1; + } + + support_packing = false; + support_bf16_storage = false; + support_fp16_storage = false; + support_vulkan = false; + support_vulkan_packing = false; +#else + NCNN_LOGE("please build ncnn with NCNN_WEIGHT_QUANT enabled for weight quantized inference"); + return -1; +#endif + } + else if (quantize_term) { #if !NCNN_INT8 NCNN_LOGE("please build ncnn with NCNN_INT8 enabled for int8 inference"); @@ -87,6 +206,16 @@ int Gemm::load_model(const ModelBin& mb) { if (transB == 0) B_data = mb.load(constantN, constantK, 0); +#if NCNN_WEIGHT_QUANT + else if (weight_block_quantize) + { + const int weight_bits = gemm_weight_quantize_bits(quantize_term); + const int packed_k_bytes = gemm_weight_quantize_packed_k_bytes(constantK, weight_bits); + if (packed_k_bytes <= 0) + return -100; + B_data = mb.load(packed_k_bytes, constantN, 0); + } +#endif // NCNN_WEIGHT_QUANT else B_data = mb.load(constantK, constantN, 0); if (B_data.empty()) @@ -109,8 +238,49 @@ int Gemm::load_model(const ModelBin& mb) return -100; } +#if NCNN_WEIGHT_QUANT + if (weight_block_quantize) + { + const int weight_bits = gemm_weight_quantize_bits(quantize_term); + const int block_size = gemm_weight_quantize_block_size(quantize_term); + const int packed_k_bytes = gemm_weight_quantize_packed_k_bytes(constantK, weight_bits); + if (packed_k_bytes <= 0) + return -100; + const int block_count = (constantK + block_size - 1) / block_size; + + if (B_data.elemsize != 1u || B_data.w != packed_k_bytes || B_data.h != constantN) + { + NCNN_LOGE("Gemm weight block quantized B_data shape mismatch"); + return -100; + } + + B_data_quantize_scales = mb.load(block_count, constantN, 1); + if (B_data_quantize_scales.empty()) + return -100; + + if (B_data_quantize_scales.elemsize != 4u || B_data_quantize_scales.w != block_count || B_data_quantize_scales.h != constantN) + { + NCNN_LOGE("Gemm weight block quantize scale shape mismatch"); + return -100; + } + + if (gemm_weight_quantize_has_input_scale(quantize_term)) + { + B_data_input_scales = mb.load(constantK, 1); + if (B_data_input_scales.empty()) + return -100; + + if (B_data_input_scales.elemsize != 4u || B_data_input_scales.w != constantK) + { + NCNN_LOGE("Gemm weight block quantize input scale shape mismatch"); + return -100; + } + } + } +#endif // NCNN_WEIGHT_QUANT + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term && !weight_block_quantize) { if (constantA == 1) { @@ -194,6 +364,157 @@ static void gemm_transB(const Mat& A, const Mat& BT, const Mat& C, Mat& top_blob } } +#if NCNN_WEIGHT_QUANT +int Gemm::forward_weight_block_quantize(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const +{ + const Mat& A = bottom_blobs[0]; + if (A.elemsize != 4u || A.elempack != 1) + { + NCNN_LOGE("Gemm weight block quantize only supports fp32 input"); + return -1; + } + + const int K = A.w; + if (K != constantK) + { + NCNN_LOGE("Gemm weight block quantize K mismatch"); + return -1; + } + + const int weight_bits = gemm_weight_quantize_bits(quantize_term); + const int block_size = gemm_weight_quantize_block_size(quantize_term); + const int packed_k_bytes = gemm_weight_quantize_packed_k_bytes(constantK, weight_bits); + if (packed_k_bytes <= 0) + return -1; + + const bool has_input_scale = gemm_weight_quantize_has_input_scale(quantize_term); + const float* input_scale_ptr = has_input_scale ? (const float*)B_data_input_scales : 0; + + const int M = A.dims == 3 ? A.c : A.h; + const int N = constantN; + + Mat C; + int broadcast_type_C = -1; + if (constantC) + { + C = C_data; + broadcast_type_C = constant_broadcast_type_C; + } + else + { + if (bottom_blobs.size() == 2) + { + C = bottom_blobs[1]; + } + + if (!C.empty()) + { + bool matched = false; + if (C.dims == 1 && C.w == 1) + { + broadcast_type_C = 0; + matched = true; + } + if (C.dims == 1 && C.w == M) + { + broadcast_type_C = 1; + matched = true; + } + if (C.dims == 1 && C.w == N) + { + broadcast_type_C = 4; + matched = true; + } + if (C.dims == 2 && C.w == 1 && C.h == M) + { + broadcast_type_C = 2; + matched = true; + } + if (C.dims == 2 && C.w == N && C.h == M) + { + broadcast_type_C = 3; + matched = true; + } + if (C.dims == 2 && C.w == N && C.h == 1) + { + broadcast_type_C = 4; + matched = true; + } + + if (!matched || C.elemsize != 4u || C.elempack != 1) + { + NCNN_LOGE("Gemm weight block quantize unsupported C"); + return -1; + } + } + } + + if (!C.empty() && (C.elemsize != 4u || C.elempack != 1)) + { + NCNN_LOGE("Gemm weight block quantize only supports fp32 C"); + return -1; + } + + Mat& top_blob = top_blobs[0]; + top_blob.create(N, M, (size_t)4u, opt.blob_allocator); + if (top_blob.empty()) + return -100; + + const size_t A_hstep = A.dims == 3 ? A.cstep : (size_t)A.w; + const float* ptrC = C; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int i = 0; i < M; i++) + { + const float* ptrA = (const float*)A + i * A_hstep; + float* outptr = top_blob.row(i); + + for (int j = 0; j < N; j++) + { + float sum = 0.f; + if (ptrC) + { + if (broadcast_type_C == 0) + sum = ptrC[0]; + if (broadcast_type_C == 1) + sum = ptrC[i]; + if (broadcast_type_C == 2) + sum = ptrC[i]; + if (broadcast_type_C == 3) + sum = ptrC[i * N + j]; + if (broadcast_type_C == 4) + sum = ptrC[j]; + + sum *= beta; + } + + const unsigned char* ptrB = B_data.row(j); + const float* scale_ptr = B_data_quantize_scales.row(j); + + for (int k0 = 0; k0 < K; k0 += block_size) + { + const int max_kk = block_size < K - k0 ? block_size : K - k0; + const float descale = 1.f / scale_ptr[k0 / block_size]; + + for (int kk = 0; kk < max_kk; kk++) + { + const int k = k0 + kk; + const int q = gemm_weight_block_quantize_unpack(ptrB, k, weight_bits, packed_k_bytes); + float v = ptrA[k]; + if (input_scale_ptr) + v *= input_scale_ptr[k]; + sum += v * (q * descale); + } + } + + outptr[j] = sum * alpha; + } + } + + return 0; +} +#endif // NCNN_WEIGHT_QUANT + #if NCNN_INT8 static inline signed char float2int8(float v) { @@ -228,8 +549,7 @@ static void gemm_transB_int8(const Mat& A_int8, const Mat& BT_int8, const Mat& A { sum += ptrA[k] * ptrBT[k]; #if __mips_loongson_mmi && !__mips_msa - // GCC may mis-vectorize this int8 dot loop with -mloongson-mmi. - // Keep this loop scalar without disabling tree-vectorize globally. + // workaround gcc mis-vectorization on loongson-mmi asm volatile("" :: : "memory"); #endif @@ -290,11 +610,30 @@ int Gemm::forward(const Mat& bottom_blob, Mat& top_blob, const Option& opt) cons int Gemm::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const { +#if NCNN_WEIGHT_QUANT + if (weight_block_quantize) + { + return forward_weight_block_quantize(bottom_blobs, top_blobs, opt); + } +#else + if (weight_block_quantize) + { + NCNN_LOGE("please build ncnn with NCNN_WEIGHT_QUANT enabled for weight quantized inference"); + return -1; + } +#endif + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return forward_int8(bottom_blobs, top_blobs, opt); } +#else + if (quantize_term) + { + NCNN_LOGE("please build ncnn with NCNN_INT8 enabled for int8 inference"); + return -1; + } #endif // NCNN_INT8 const Mat& A0 = constantA ? A_data : bottom_blobs[0]; diff --git a/src/layer/gemm.h b/src/layer/gemm.h index 120d9ee91259..6f974210033e 100644 --- a/src/layer/gemm.h +++ b/src/layer/gemm.h @@ -22,6 +22,10 @@ class Gemm : public Layer virtual int forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const; protected: +#if NCNN_WEIGHT_QUANT + int forward_weight_block_quantize(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const; +#endif + #if NCNN_INT8 int forward_int8(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const; #endif @@ -44,7 +48,12 @@ class Gemm : public Layer int output_elemtype; // 0=auto 1=fp32 int output_transpose; - int int8_scale_term; + union + { + int quantize_term; + int int8_scale_term; + }; + int weight_block_quantize; int constant_TILE_M; int constant_TILE_N; @@ -59,6 +68,8 @@ class Gemm : public Layer Mat A_data_int8_scales; float B_data_int8_scale; #endif + Mat B_data_quantize_scales; + Mat B_data_input_scales; }; } // namespace ncnn diff --git a/src/layer/loongarch/gemm_loongarch.cpp b/src/layer/loongarch/gemm_loongarch.cpp index 4f6f21ff0bfe..72832080bcf7 100644 --- a/src/layer/loongarch/gemm_loongarch.cpp +++ b/src/layer/loongarch/gemm_loongarch.cpp @@ -7372,8 +7372,13 @@ int Gemm_loongarch::create_pipeline(const Option& opt) CT_data.release(); nT = 0; + if (weight_block_quantize) + { + return 0; + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return create_pipeline_int8(opt); } @@ -7522,13 +7527,16 @@ int Gemm_loongarch::create_pipeline(const Option& opt) int Gemm_loongarch::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const { + if (weight_block_quantize) + { + return Gemm::forward(bottom_blobs, top_blobs, opt); + } + #if NCNN_INT8 -#if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return forward_int8(bottom_blobs, top_blobs, opt); } -#endif #endif const Mat& bottom_blob = bottom_blobs.empty() ? AT_data : bottom_blobs[0]; diff --git a/src/layer/loongarch/multiheadattention_loongarch.cpp b/src/layer/loongarch/multiheadattention_loongarch.cpp index 20af084272ab..395f191f3b84 100644 --- a/src/layer/loongarch/multiheadattention_loongarch.cpp +++ b/src/layer/loongarch/multiheadattention_loongarch.cpp @@ -3,6 +3,7 @@ #include "multiheadattention_loongarch.h" +#include "gemm.h" #include "layer_type.h" namespace ncnn { @@ -30,6 +31,9 @@ MultiHeadAttention_loongarch::MultiHeadAttention_loongarch() int MultiHeadAttention_loongarch::create_pipeline(const Option& _opt) { + if (weight_block_quantize) + return 0; + Option opt = _opt; if (int8_scale_term) { @@ -256,6 +260,9 @@ int MultiHeadAttention_loongarch::create_pipeline(const Option& _opt) int MultiHeadAttention_loongarch::destroy_pipeline(const Option& _opt) { + if (weight_block_quantize) + return 0; + Option opt = _opt; if (int8_scale_term) { @@ -315,6 +322,9 @@ int MultiHeadAttention_loongarch::destroy_pipeline(const Option& _opt) int MultiHeadAttention_loongarch::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& _opt) const { + if (weight_block_quantize) + return MultiHeadAttention::forward(bottom_blobs, top_blobs, _opt); + int q_blob_i = 0; int k_blob_i = 0; int v_blob_i = 0; diff --git a/src/layer/mips/gemm_mips.cpp b/src/layer/mips/gemm_mips.cpp index 77d80676bdb0..81f0cb108b5d 100644 --- a/src/layer/mips/gemm_mips.cpp +++ b/src/layer/mips/gemm_mips.cpp @@ -4477,8 +4477,13 @@ int Gemm_mips::create_pipeline(const Option& opt) CT_data.release(); nT = 0; + if (weight_block_quantize) + { + return 0; + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return create_pipeline_int8(opt); } @@ -4613,13 +4618,16 @@ int Gemm_mips::create_pipeline(const Option& opt) int Gemm_mips::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const { + if (weight_block_quantize) + { + return Gemm::forward(bottom_blobs, top_blobs, opt); + } + #if NCNN_INT8 -#if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return forward_int8(bottom_blobs, top_blobs, opt); } -#endif #endif const Mat& bottom_blob = bottom_blobs.empty() ? AT_data : bottom_blobs[0]; diff --git a/src/layer/mips/multiheadattention_mips.cpp b/src/layer/mips/multiheadattention_mips.cpp index da51e1f0b4a5..177e06acc711 100644 --- a/src/layer/mips/multiheadattention_mips.cpp +++ b/src/layer/mips/multiheadattention_mips.cpp @@ -3,6 +3,7 @@ #include "multiheadattention_mips.h" +#include "gemm.h" #include "layer_type.h" namespace ncnn { @@ -30,6 +31,9 @@ MultiHeadAttention_mips::MultiHeadAttention_mips() int MultiHeadAttention_mips::create_pipeline(const Option& _opt) { + if (weight_block_quantize) + return 0; + Option opt = _opt; if (int8_scale_term) { @@ -256,6 +260,9 @@ int MultiHeadAttention_mips::create_pipeline(const Option& _opt) int MultiHeadAttention_mips::destroy_pipeline(const Option& _opt) { + if (weight_block_quantize) + return 0; + Option opt = _opt; if (int8_scale_term) { @@ -315,6 +322,9 @@ int MultiHeadAttention_mips::destroy_pipeline(const Option& _opt) int MultiHeadAttention_mips::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& _opt) const { + if (weight_block_quantize) + return MultiHeadAttention::forward(bottom_blobs, top_blobs, _opt); + int q_blob_i = 0; int k_blob_i = 0; int v_blob_i = 0; diff --git a/src/layer/multiheadattention.cpp b/src/layer/multiheadattention.cpp index ce637b7d13ef..c7cf2d24cc54 100644 --- a/src/layer/multiheadattention.cpp +++ b/src/layer/multiheadattention.cpp @@ -4,11 +4,122 @@ #include "multiheadattention.h" #include +#include namespace ncnn { +static bool mha_is_weight_block_quantize(int quantize_term) +{ + const int weight_bits = quantize_term / 100; + const int format_code = quantize_term % 100 / 10; + const int block_size_code = quantize_term % 10; + + if (weight_bits != 4 && weight_bits != 6 && weight_bits != 8) + return false; + + if (format_code != 0 && format_code != 1) + return false; + + if (block_size_code < 0 || block_size_code > 2) + return false; + + return true; +} + +static bool mha_weight_quantize_has_input_scale(int quantize_term) +{ + if (!mha_is_weight_block_quantize(quantize_term)) + return false; + + return quantize_term % 100 / 10 == 1; +} + +static int mha_weight_quantize_bits(int quantize_term) +{ + return mha_is_weight_block_quantize(quantize_term) ? quantize_term / 100 : 0; +} + +static int mha_weight_quantize_block_size(int quantize_term) +{ + if (!mha_is_weight_block_quantize(quantize_term)) + return 0; + + const int block_size_code = quantize_term % 10; + return block_size_code == 0 ? 32 : block_size_code == 1 ? 64 : 128; +} + +static int mha_weight_quantize_packed_k_bytes(int constantK, int weight_bits) +{ + if (constantK <= 0 || weight_bits <= 0) + return -1; + + const size_t packed_k_bytes = ((size_t)constantK * weight_bits + 7) / 8; + if (packed_k_bytes > (size_t)INT_MAX) + return -1; + + return (int)packed_k_bytes; +} + +#if NCNN_WEIGHT_QUANT +static inline int mha_weight_block_quantize_sign_extend(int v, int bits) +{ + const int sign_bit = 1 << (bits - 1); + return (v ^ sign_bit) - sign_bit; +} + +static inline int mha_weight_block_quantize_unpack(const unsigned char* ptr, int k, int bits, int packed_k_bytes) +{ + const int bit_offset = k * bits; + const int byte_offset = bit_offset / 8; + const int bit_shift = bit_offset % 8; + + unsigned int v = ptr[byte_offset]; + if (byte_offset + 1 < packed_k_bytes) + v |= (unsigned int)ptr[byte_offset + 1] << 8; + if (byte_offset + 2 < packed_k_bytes) + v |= (unsigned int)ptr[byte_offset + 2] << 16; + + const int mask = (1 << bits) - 1; + return mha_weight_block_quantize_sign_extend((v >> bit_shift) & mask, bits); +} + +static int mha_check_weight_block_quantize_shape(const Mat& weight_data, const Mat& weight_data_quantize_scales, int K, int N, int weight_bits, int block_size, const char* name) +{ + const int packed_k_bytes = mha_weight_quantize_packed_k_bytes(K, weight_bits); + if (packed_k_bytes <= 0) + return -100; + const int block_count = (K + block_size - 1) / block_size; + + if (weight_data.elemsize != 1u || weight_data.w != packed_k_bytes || weight_data.h != N) + { + NCNN_LOGE("MultiHeadAttention weight block quantized %s_data shape mismatch", name); + return -100; + } + + if (weight_data_quantize_scales.elemsize != 4u || weight_data_quantize_scales.w != block_count || weight_data_quantize_scales.h != N) + { + NCNN_LOGE("MultiHeadAttention weight block quantize %s_scale shape mismatch", name); + return -100; + } + + return 0; +} + +static int mha_check_weight_block_quantize_input_scales(const Mat& input_scales, int K, const char* name) +{ + if (input_scales.elemsize != 4u || input_scales.w != K) + { + NCNN_LOGE("MultiHeadAttention weight block quantize %s_input_scale shape mismatch", name); + return -100; + } + + return 0; +} +#endif // NCNN_WEIGHT_QUANT + MultiHeadAttention::MultiHeadAttention() { + weight_block_quantize = 0; } int MultiHeadAttention::load_param(const ParamDict& pd) @@ -21,7 +132,47 @@ int MultiHeadAttention::load_param(const ParamDict& pd) attn_mask = pd.get(5, 0); scale = pd.get(6, 1.f / sqrtf(embed_dim / num_heads)); kv_cache = pd.get(7, 0); - int8_scale_term = pd.get(18, 0); + quantize_term = pd.get(18, 0); + weight_block_quantize = mha_is_weight_block_quantize(quantize_term); + + if (quantize_term == 4 || quantize_term == 5 || quantize_term == 6) + { + NCNN_LOGE("MultiHeadAttention quantize_term %d is an obsolete block quantization prototype value", quantize_term); + return -1; + } + + if (quantize_term >= 400 && !weight_block_quantize) + { + NCNN_LOGE("MultiHeadAttention unsupported quantize_term %d", quantize_term); + return -1; + } + + if (weight_block_quantize) + { +#if NCNN_WEIGHT_QUANT + if (embed_dim <= 0 || num_heads <= 0 || embed_dim % num_heads != 0 || weight_data_size <= 0 || weight_data_size % embed_dim != 0 || kdim <= 0 || vdim <= 0) + { + NCNN_LOGE("MultiHeadAttention weight block quantization requires valid embed_dim/num_heads/weight_data_size/kdim/vdim"); + return -1; + } + + support_packing = false; + support_bf16_storage = false; + support_fp16_storage = false; + support_vulkan = false; + support_vulkan_packing = false; +#else + NCNN_LOGE("please build ncnn with NCNN_WEIGHT_QUANT enabled for weight quantized inference"); + return -1; +#endif + } + else if (quantize_term) + { +#if !NCNN_INT8 + NCNN_LOGE("please build ncnn with NCNN_INT8 enabled for int8 inference"); + return -1; +#endif + } return 0; } @@ -30,6 +181,98 @@ int MultiHeadAttention::load_model(const ModelBin& mb) { const int qdim = weight_data_size / embed_dim; +#if NCNN_WEIGHT_QUANT + if (weight_block_quantize) + { + const int weight_bits = mha_weight_quantize_bits(quantize_term); + const int block_size = mha_weight_quantize_block_size(quantize_term); + + const int q_packed_k_bytes = mha_weight_quantize_packed_k_bytes(qdim, weight_bits); + const int k_packed_k_bytes = mha_weight_quantize_packed_k_bytes(kdim, weight_bits); + const int v_packed_k_bytes = mha_weight_quantize_packed_k_bytes(vdim, weight_bits); + const int out_packed_k_bytes = mha_weight_quantize_packed_k_bytes(embed_dim, weight_bits); + if (q_packed_k_bytes <= 0 || k_packed_k_bytes <= 0 || v_packed_k_bytes <= 0 || out_packed_k_bytes <= 0) + return -100; + + q_weight_data = mb.load(q_packed_k_bytes, embed_dim, 0); + if (q_weight_data.empty()) + return -100; + + q_bias_data = mb.load(embed_dim, 1); + if (q_bias_data.empty()) + return -100; + + k_weight_data = mb.load(k_packed_k_bytes, embed_dim, 0); + if (k_weight_data.empty()) + return -100; + + k_bias_data = mb.load(embed_dim, 1); + if (k_bias_data.empty()) + return -100; + + v_weight_data = mb.load(v_packed_k_bytes, embed_dim, 0); + if (v_weight_data.empty()) + return -100; + + v_bias_data = mb.load(embed_dim, 1); + if (v_bias_data.empty()) + return -100; + + out_weight_data = mb.load(out_packed_k_bytes, qdim, 0); + if (out_weight_data.empty()) + return -100; + + out_bias_data = mb.load(qdim, 1); + if (out_bias_data.empty()) + return -100; + + q_weight_data_quantize_scales = mb.load((qdim + block_size - 1) / block_size, embed_dim, 1); + k_weight_data_quantize_scales = mb.load((kdim + block_size - 1) / block_size, embed_dim, 1); + v_weight_data_quantize_scales = mb.load((vdim + block_size - 1) / block_size, embed_dim, 1); + out_weight_data_quantize_scales = mb.load((embed_dim + block_size - 1) / block_size, qdim, 1); + if (q_weight_data_quantize_scales.empty() || k_weight_data_quantize_scales.empty() || v_weight_data_quantize_scales.empty() || out_weight_data_quantize_scales.empty()) + return -100; + + int ret = mha_check_weight_block_quantize_shape(q_weight_data, q_weight_data_quantize_scales, qdim, embed_dim, weight_bits, block_size, "q_weight"); + if (ret != 0) + return ret; + ret = mha_check_weight_block_quantize_shape(k_weight_data, k_weight_data_quantize_scales, kdim, embed_dim, weight_bits, block_size, "k_weight"); + if (ret != 0) + return ret; + ret = mha_check_weight_block_quantize_shape(v_weight_data, v_weight_data_quantize_scales, vdim, embed_dim, weight_bits, block_size, "v_weight"); + if (ret != 0) + return ret; + ret = mha_check_weight_block_quantize_shape(out_weight_data, out_weight_data_quantize_scales, embed_dim, qdim, weight_bits, block_size, "out_weight"); + if (ret != 0) + return ret; + + if (mha_weight_quantize_has_input_scale(quantize_term)) + { + q_weight_data_input_scales = mb.load(qdim, 1); + k_weight_data_input_scales = mb.load(kdim, 1); + v_weight_data_input_scales = mb.load(vdim, 1); + out_weight_data_input_scales = mb.load(embed_dim, 1); + if (q_weight_data_input_scales.empty() || k_weight_data_input_scales.empty() || v_weight_data_input_scales.empty() || out_weight_data_input_scales.empty()) + return -100; + + ret = mha_check_weight_block_quantize_input_scales(q_weight_data_input_scales, qdim, "q_weight"); + if (ret != 0) + return ret; + ret = mha_check_weight_block_quantize_input_scales(k_weight_data_input_scales, kdim, "k_weight"); + if (ret != 0) + return ret; + ret = mha_check_weight_block_quantize_input_scales(v_weight_data_input_scales, vdim, "v_weight"); + if (ret != 0) + return ret; + ret = mha_check_weight_block_quantize_input_scales(out_weight_data_input_scales, embed_dim, "out_weight"); + if (ret != 0) + return ret; + } + + return 0; + } +#endif // NCNN_WEIGHT_QUANT + q_weight_data = mb.load(embed_dim * qdim, 0); if (q_weight_data.empty()) return -100; @@ -63,7 +306,7 @@ int MultiHeadAttention::load_model(const ModelBin& mb) return -100; #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { q_weight_data_int8_scales = mb.load(embed_dim, 1); k_weight_data_int8_scales = mb.load(embed_dim, 1); @@ -78,8 +321,15 @@ int MultiHeadAttention::load_model(const ModelBin& mb) // refers to https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html int MultiHeadAttention::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const { +#if NCNN_WEIGHT_QUANT + if (weight_block_quantize) + { + return forward_weight_block_quantize(bottom_blobs, top_blobs, opt); + } +#endif + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return forward_int8(bottom_blobs, top_blobs, opt); } @@ -347,15 +597,369 @@ int MultiHeadAttention::forward(const std::vector& bottom_blobs, std::vecto #pragma omp parallel for num_threads(opt.num_threads) for (int i = 0; i < src_seqlen; i++) { - const float* kptr = (const float*)out_weight_data; float* outptr = top_blob.row(i); for (int j = 0; j < qdim; j++) { + const float* kptr = (const float*)out_weight_data + j * embed_dim; + float sum = out_bias_data[j]; for (int k = 0; k < embed_dim; k++) { - sum += qkv_cross.row(k)[i] * *kptr++; + sum += qkv_cross.row(k)[i] * kptr[k]; + } + + outptr[j] = sum; + } + } + } + + if (kv_cache) + { + // assert top_blobs.size() == 3 + top_blobs[1] = k_affine; + top_blobs[2] = v_affine; + } + + return 0; +} + +#if NCNN_WEIGHT_QUANT +int MultiHeadAttention::forward_weight_block_quantize(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const +{ + int q_blob_i = 0; + int k_blob_i = 0; + int v_blob_i = 0; + int attn_mask_i = 0; + int cached_xk_i = 0; + int cached_xv_i = 0; + resolve_bottom_blob_index((int)bottom_blobs.size(), q_blob_i, k_blob_i, v_blob_i, attn_mask_i, cached_xk_i, cached_xv_i); + + const Mat& q_blob = bottom_blobs[q_blob_i]; + const Mat& k_blob = bottom_blobs[k_blob_i]; + const Mat& v_blob = bottom_blobs[v_blob_i]; + const Mat& attn_mask_blob = attn_mask ? bottom_blobs[attn_mask_i] : Mat(); + const Mat& cached_xk_blob = kv_cache ? bottom_blobs[cached_xk_i] : Mat(); + const Mat& cached_xv_blob = kv_cache ? bottom_blobs[cached_xv_i] : Mat(); + + // | self-attention cross-attention + // w/o kvcache | past(0) + cur cur + // with kvcache | past + cur past + + const int src_seqlen = q_blob.h; + const int cur_seqlen = k_blob.h; + const int past_seqlen = kv_cache && !cached_xk_blob.empty() ? cached_xk_blob.w : 0; + const int dst_seqlen = past_seqlen > 0 ? (q_blob_i == k_blob_i ? (past_seqlen + cur_seqlen) : past_seqlen) : cur_seqlen; + + const int embed_dim_per_head = embed_dim / num_heads; + const int qdim = weight_data_size / embed_dim; + const int weight_bits = mha_weight_quantize_bits(quantize_term); + const int block_size = mha_weight_quantize_block_size(quantize_term); + const bool has_input_scale = mha_weight_quantize_has_input_scale(quantize_term); + + const float* q_input_scale_ptr = has_input_scale ? (const float*)q_weight_data_input_scales : 0; + const float* k_input_scale_ptr = has_input_scale ? (const float*)k_weight_data_input_scales : 0; + const float* v_input_scale_ptr = has_input_scale ? (const float*)v_weight_data_input_scales : 0; + const float* out_input_scale_ptr = has_input_scale ? (const float*)out_weight_data_input_scales : 0; + + // assert k_blob.h == v_blob.h + + Mat q_affine; + { + q_affine.create(src_seqlen, embed_dim, 4u, opt.workspace_allocator); + if (q_affine.empty()) + return -100; + + const int packed_k_bytes = mha_weight_quantize_packed_k_bytes(qdim, weight_bits); + + #pragma omp parallel for num_threads(opt.num_threads) + for (int i = 0; i < src_seqlen; i++) + { + for (int j = 0; j < embed_dim; j++) + { + const float* ptr = q_blob.row(i); + const unsigned char* kptr = q_weight_data.row(j); + const float* scale_ptr = q_weight_data_quantize_scales.row(j); + + float sum = q_bias_data[j]; + for (int k0 = 0; k0 < qdim; k0 += block_size) + { + const int max_kk = block_size < qdim - k0 ? block_size : qdim - k0; + const float descale = 1.f / scale_ptr[k0 / block_size]; + + for (int kk = 0; kk < max_kk; kk++) + { + const int k = k0 + kk; + const int q = mha_weight_block_quantize_unpack(kptr, k, weight_bits, packed_k_bytes); + float v = ptr[k]; + if (q_input_scale_ptr) + v *= q_input_scale_ptr[k]; + sum += v * (q * descale); + } + } + + q_affine.row(j)[i] = sum * scale; + } + } + } + + Mat k_affine; + if (past_seqlen > 0 && q_blob_i != k_blob_i) + { + k_affine = cached_xk_blob; + } + else + { + k_affine.create(dst_seqlen, embed_dim, 4u, opt.workspace_allocator); + if (k_affine.empty()) + return -100; + + if (past_seqlen > 0) + { + // reuse cached_xk + #pragma omp parallel for num_threads(opt.num_threads) + for (int i = 0; i < embed_dim; i++) + { + memcpy(k_affine.row(i), cached_xk_blob.row(i), past_seqlen * sizeof(float)); + } + } + + const int packed_k_bytes = mha_weight_quantize_packed_k_bytes(kdim, weight_bits); + + #pragma omp parallel for num_threads(opt.num_threads) + for (int i = 0; i < cur_seqlen; i++) + { + for (int j = 0; j < embed_dim; j++) + { + const float* ptr = k_blob.row(i); + const unsigned char* kptr = k_weight_data.row(j); + const float* scale_ptr = k_weight_data_quantize_scales.row(j); + + float sum = k_bias_data[j]; + for (int k0 = 0; k0 < kdim; k0 += block_size) + { + const int max_kk = block_size < kdim - k0 ? block_size : kdim - k0; + const float descale = 1.f / scale_ptr[k0 / block_size]; + + for (int kk = 0; kk < max_kk; kk++) + { + const int k = k0 + kk; + const int q = mha_weight_block_quantize_unpack(kptr, k, weight_bits, packed_k_bytes); + float v = ptr[k]; + if (k_input_scale_ptr) + v *= k_input_scale_ptr[k]; + sum += v * (q * descale); + } + } + + k_affine.row(j)[past_seqlen + i] = sum; + } + } + } + + Mat v_affine; + if (past_seqlen > 0 && q_blob_i != v_blob_i) + { + v_affine = cached_xv_blob; + } + else + { + v_affine.create(dst_seqlen, embed_dim, 4u, opt.workspace_allocator); + if (v_affine.empty()) + return -100; + + if (past_seqlen > 0) + { + // reuse cached_xv + #pragma omp parallel for num_threads(opt.num_threads) + for (int i = 0; i < embed_dim; i++) + { + memcpy(v_affine.row(i), cached_xv_blob.row(i), past_seqlen * sizeof(float)); + } + } + + const int packed_k_bytes = mha_weight_quantize_packed_k_bytes(vdim, weight_bits); + + #pragma omp parallel for num_threads(opt.num_threads) + for (int i = 0; i < cur_seqlen; i++) + { + for (int j = 0; j < embed_dim; j++) + { + const float* ptr = v_blob.row(i); + const unsigned char* kptr = v_weight_data.row(j); + const float* scale_ptr = v_weight_data_quantize_scales.row(j); + + float sum = v_bias_data[j]; + for (int k0 = 0; k0 < vdim; k0 += block_size) + { + const int max_kk = block_size < vdim - k0 ? block_size : vdim - k0; + const float descale = 1.f / scale_ptr[k0 / block_size]; + + for (int kk = 0; kk < max_kk; kk++) + { + const int k = k0 + kk; + const int q = mha_weight_block_quantize_unpack(kptr, k, weight_bits, packed_k_bytes); + float v = ptr[k]; + if (v_input_scale_ptr) + v *= v_input_scale_ptr[k]; + sum += v * (q * descale); + } + } + + v_affine.row(j)[past_seqlen + i] = sum; + } + } + } + + Mat qk_cross; + { + qk_cross.create(dst_seqlen, src_seqlen, num_heads, 4u, opt.workspace_allocator); + if (qk_cross.empty()) + return -100; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = 0; q < num_heads; q++) + { + const Mat q_affine_head = q_affine.row_range(q * embed_dim_per_head, embed_dim_per_head); + const Mat k_affine_head = k_affine.row_range(q * embed_dim_per_head, embed_dim_per_head); + Mat qk_cross_head = qk_cross.channel(q); + + for (int i = 0; i < src_seqlen; i++) + { + float* outptr = qk_cross_head.row(i); + + for (int j = 0; j < dst_seqlen; j++) + { + float sum = 0.f; + for (int l = 0; l < embed_dim_per_head; l++) + { + sum += q_affine_head.row(l)[i] * k_affine_head.row(l)[j]; + } + + outptr[j] = sum; + } + } + } + } + + if (attn_mask) + { + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = 0; q < num_heads; q++) + { + const Mat& maskm = attn_mask_blob.dims == 3 ? attn_mask_blob.channel(q) : attn_mask_blob; + Mat qk_cross_head = qk_cross.channel(q); + + for (int i = 0; i < src_seqlen; i++) + { + const float* mptr = maskm.row(i); + float* outptr = qk_cross_head.row(i); + + for (int j = 0; j < dst_seqlen; j++) + { + outptr[j] += mptr[j]; + } + } + } + } + + // softmax(qk_cross) + { + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = 0; q < num_heads; q++) + { + Mat qk_cross_head = qk_cross.channel(q); + + for (int i = 0; i < src_seqlen; i++) + { + float* ptr = qk_cross_head.row(i); + + float max = -FLT_MAX; + for (int j = 0; j < dst_seqlen; j++) + { + max = std::max(max, ptr[j]); + } + + float sum = 0.f; + for (int j = 0; j < dst_seqlen; j++) + { + ptr[j] = (float)expf(ptr[j] - max); + sum += ptr[j]; + } + + for (int j = 0; j < dst_seqlen; j++) + { + ptr[j] /= sum; + } + } + } + } + + Mat qkv_cross; + { + qkv_cross.create(src_seqlen, embed_dim, 4u, opt.workspace_allocator); + if (qkv_cross.empty()) + return -100; + + #pragma omp parallel for num_threads(opt.num_threads) + for (int q = 0; q < num_heads; q++) + { + const Mat qk_cross_head = qk_cross.channel(q); + const Mat v_affine_head = v_affine.row_range(q * embed_dim_per_head, embed_dim_per_head); + Mat qkv_cross_head = qkv_cross.row_range(q * embed_dim_per_head, embed_dim_per_head); + + for (int i = 0; i < src_seqlen; i++) + { + for (int j = 0; j < embed_dim_per_head; j++) + { + const float* qkptr = qk_cross_head.row(i); + const float* vptr = v_affine_head.row(j); + + float sum = 0.f; + for (int k = 0; k < dst_seqlen; k++) + { + sum += *qkptr++ * *vptr++; + } + + qkv_cross_head.row(j)[i] = sum; + } + } + } + } + + Mat& top_blob = top_blobs[0]; + { + top_blob.create(qdim, src_seqlen, 4u, opt.blob_allocator); + if (top_blob.empty()) + return -100; + + const int packed_k_bytes = mha_weight_quantize_packed_k_bytes(embed_dim, weight_bits); + + #pragma omp parallel for num_threads(opt.num_threads) + for (int i = 0; i < src_seqlen; i++) + { + float* outptr = top_blob.row(i); + + for (int j = 0; j < qdim; j++) + { + const unsigned char* kptr = out_weight_data.row(j); + const float* scale_ptr = out_weight_data_quantize_scales.row(j); + + float sum = out_bias_data[j]; + for (int k0 = 0; k0 < embed_dim; k0 += block_size) + { + const int max_kk = block_size < embed_dim - k0 ? block_size : embed_dim - k0; + const float descale = 1.f / scale_ptr[k0 / block_size]; + + for (int kk = 0; kk < max_kk; kk++) + { + const int k = k0 + kk; + const int q = mha_weight_block_quantize_unpack(kptr, k, weight_bits, packed_k_bytes); + float v = qkv_cross.row(k)[i]; + if (out_input_scale_ptr) + v *= out_input_scale_ptr[k]; + sum += v * (q * descale); + } } outptr[j] = sum; @@ -372,6 +976,7 @@ int MultiHeadAttention::forward(const std::vector& bottom_blobs, std::vecto return 0; } +#endif // NCNN_WEIGHT_QUANT void MultiHeadAttention::resolve_bottom_blob_index(int bottom_blob_count, int& q_blob_i, int& k_blob_i, int& v_blob_i, int& attn_mask_i, int& cached_xk_i, int& cached_xv_i) const { diff --git a/src/layer/multiheadattention.h b/src/layer/multiheadattention.h index 10f67c1a481f..30bb88347a40 100644 --- a/src/layer/multiheadattention.h +++ b/src/layer/multiheadattention.h @@ -22,6 +22,10 @@ class MultiHeadAttention : public Layer protected: void resolve_bottom_blob_index(int bottom_blob_count, int& q_blob_i, int& k_blob_i, int& v_blob_i, int& attn_mask_i, int& cached_xk_i, int& cached_xv_i) const; +#if NCNN_WEIGHT_QUANT + int forward_weight_block_quantize(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const; +#endif + #if NCNN_INT8 int forward_int8(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const; #endif @@ -36,7 +40,12 @@ class MultiHeadAttention : public Layer float scale; int kv_cache; - int int8_scale_term; + union + { + int quantize_term; + int int8_scale_term; + }; + int weight_block_quantize; Mat q_weight_data; Mat q_bias_data; @@ -53,6 +62,14 @@ class MultiHeadAttention : public Layer Mat v_weight_data_int8_scales; float out_weight_data_int8_scale; #endif + Mat q_weight_data_quantize_scales; + Mat k_weight_data_quantize_scales; + Mat v_weight_data_quantize_scales; + Mat out_weight_data_quantize_scales; + Mat q_weight_data_input_scales; + Mat k_weight_data_input_scales; + Mat v_weight_data_input_scales; + Mat out_weight_data_input_scales; }; } // namespace ncnn diff --git a/src/layer/riscv/gemm_riscv.cpp b/src/layer/riscv/gemm_riscv.cpp index b8578e704aa1..9e03af872798 100644 --- a/src/layer/riscv/gemm_riscv.cpp +++ b/src/layer/riscv/gemm_riscv.cpp @@ -1871,8 +1871,13 @@ static int gemm_AT_BT_riscv(const Mat& AT, const Mat& BT, const Mat& C, Mat& top int Gemm_riscv::create_pipeline(const Option& opt) { + if (weight_block_quantize) + { + return 0; + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { support_packing = false; support_fp16_storage = false; @@ -2016,8 +2021,13 @@ int Gemm_riscv::create_pipeline(const Option& opt) int Gemm_riscv::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const { + if (weight_block_quantize) + { + return Gemm::forward(bottom_blobs, top_blobs, opt); + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return Gemm::forward_int8(bottom_blobs, top_blobs, opt); } diff --git a/src/layer/vulkan/gemm_vulkan.cpp b/src/layer/vulkan/gemm_vulkan.cpp index e8b59736f794..b2e57c2de6ea 100644 --- a/src/layer/vulkan/gemm_vulkan.cpp +++ b/src/layer/vulkan/gemm_vulkan.cpp @@ -49,8 +49,14 @@ Gemm_vulkan::Gemm_vulkan() int Gemm_vulkan::create_pipeline(const Option& opt) { + if (weight_block_quantize) + { + NCNN_LOGE("Gemm weight block quantization is not supported by Vulkan"); + return -1; + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return create_pipeline_int8(opt); } @@ -651,8 +657,14 @@ int Gemm_vulkan::destroy_pipeline(const Option& /*opt*/) int Gemm_vulkan::upload_model(VkTransfer& cmd, const Option& opt) { + if (weight_block_quantize) + { + NCNN_LOGE("Gemm weight block quantization is not supported by Vulkan"); + return -1; + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return upload_model_int8(cmd, opt); } @@ -684,8 +696,14 @@ int Gemm_vulkan::upload_model(VkTransfer& cmd, const Option& opt) int Gemm_vulkan::forward(const std::vector& bottom_blobs, std::vector& top_blobs, VkCompute& cmd, const Option& opt) const { + if (weight_block_quantize) + { + NCNN_LOGE("Gemm weight block quantization is not supported by Vulkan"); + return -1; + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { return forward_int8(bottom_blobs, top_blobs, cmd, opt); } diff --git a/src/layer/vulkan/multiheadattention_vulkan.cpp b/src/layer/vulkan/multiheadattention_vulkan.cpp index 675eeb687ef9..e027755957e5 100644 --- a/src/layer/vulkan/multiheadattention_vulkan.cpp +++ b/src/layer/vulkan/multiheadattention_vulkan.cpp @@ -3,6 +3,7 @@ #include "multiheadattention_vulkan.h" +#include "gemm.h" #include "layer_shader_type.h" #include "layer_type.h" @@ -48,6 +49,12 @@ int MultiHeadAttention_vulkan::load_param(const ParamDict& pd) int MultiHeadAttention_vulkan::create_pipeline(const Option& opt) { + if (weight_block_quantize) + { + NCNN_LOGE("MultiHeadAttention weight block quantization is not supported by Vulkan"); + return -1; + } + // const int embed_dim_per_head = embed_dim / num_heads; const int qdim = weight_data_size / embed_dim; { @@ -255,6 +262,9 @@ int MultiHeadAttention_vulkan::create_pipeline(const Option& opt) int MultiHeadAttention_vulkan::destroy_pipeline(const Option& opt) { + if (weight_block_quantize) + return 0; + if (q_gemm) { q_gemm->destroy_pipeline(opt); @@ -326,6 +336,12 @@ int MultiHeadAttention_vulkan::destroy_pipeline(const Option& opt) int MultiHeadAttention_vulkan::upload_model(VkTransfer& cmd, const Option& opt) { + if (weight_block_quantize) + { + NCNN_LOGE("MultiHeadAttention weight block quantization is not supported by Vulkan"); + return -1; + } + if (q_gemm) { q_gemm->upload_model(cmd, opt); @@ -351,6 +367,12 @@ int MultiHeadAttention_vulkan::upload_model(VkTransfer& cmd, const Option& opt) int MultiHeadAttention_vulkan::forward(const std::vector& bottom_blobs, std::vector& top_blobs, VkCompute& cmd, const Option& opt) const { + if (weight_block_quantize) + { + NCNN_LOGE("MultiHeadAttention weight block quantization is not supported by Vulkan"); + return -1; + } + int q_blob_i = 0; int k_blob_i = 0; int v_blob_i = 0; diff --git a/src/layer/x86/gemm_x86.cpp b/src/layer/x86/gemm_x86.cpp index b8dad80ff07a..0d176d4778a5 100644 --- a/src/layer/x86/gemm_x86.cpp +++ b/src/layer/x86/gemm_x86.cpp @@ -7434,8 +7434,13 @@ static int gemm_AT_BT_x86(const Mat& AT, const Mat& BT, const Mat& C, Mat& top_b int Gemm_x86::create_pipeline(const Option& opt) { + if (weight_block_quantize) + { + return 0; + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { support_bf16_storage = false; return create_pipeline_int8(opt); @@ -7588,8 +7593,13 @@ int Gemm_x86::create_pipeline(const Option& opt) int Gemm_x86::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const { + if (weight_block_quantize) + { + return Gemm::forward(bottom_blobs, top_blobs, opt); + } + #if NCNN_INT8 - if (int8_scale_term) + if (quantize_term) { // return Gemm::forward_int8(bottom_blobs, top_blobs, opt); return forward_int8(bottom_blobs, top_blobs, opt); diff --git a/src/layer/x86/multiheadattention_x86.cpp b/src/layer/x86/multiheadattention_x86.cpp index 745d4aeb19fc..c63e47cbb7a6 100644 --- a/src/layer/x86/multiheadattention_x86.cpp +++ b/src/layer/x86/multiheadattention_x86.cpp @@ -3,6 +3,7 @@ #include "multiheadattention_x86.h" +#include "gemm.h" #include "layer_type.h" namespace ncnn { @@ -31,6 +32,9 @@ MultiHeadAttention_x86::MultiHeadAttention_x86() int MultiHeadAttention_x86::create_pipeline(const Option& _opt) { + if (weight_block_quantize) + return 0; + Option opt = _opt; if (int8_scale_term) { @@ -257,6 +261,9 @@ int MultiHeadAttention_x86::create_pipeline(const Option& _opt) int MultiHeadAttention_x86::destroy_pipeline(const Option& _opt) { + if (weight_block_quantize) + return 0; + Option opt = _opt; if (int8_scale_term) { @@ -316,6 +323,9 @@ int MultiHeadAttention_x86::destroy_pipeline(const Option& _opt) int MultiHeadAttention_x86::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& _opt) const { + if (weight_block_quantize) + return MultiHeadAttention::forward(bottom_blobs, top_blobs, _opt); + int q_blob_i = 0; int k_blob_i = 0; int v_blob_i = 0; diff --git a/src/platform.h.in b/src/platform.h.in index 5473925816eb..821470c53114 100644 --- a/src/platform.h.in +++ b/src/platform.h.in @@ -58,6 +58,7 @@ #cmakedefine01 NCNN_ZVFH #cmakedefine01 NCNN_XTHEADVECTOR #cmakedefine01 NCNN_INT8 +#cmakedefine01 NCNN_WEIGHT_QUANT #cmakedefine01 NCNN_BF16 #cmakedefine01 NCNN_FORCE_INLINE diff --git a/tests/test_gemm_block_quant.cpp b/tests/test_gemm_block_quant.cpp new file mode 100644 index 000000000000..39bea99dfc1f --- /dev/null +++ b/tests/test_gemm_block_quant.cpp @@ -0,0 +1,251 @@ +// Copyright 2026 Tencent +// SPDX-License-Identifier: BSD-3-Clause + +#include "testutil.h" + +#include "gemm.h" +#include "layer_type.h" + +#include + +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(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 ncnn::ParamDict make_gemm_param(int M, int N, int K, int quantize_term) +{ + ncnn::ParamDict pd; + pd.set(2, 0); // transA + pd.set(3, 1); // transB + pd.set(4, 0); // constantA + pd.set(5, 1); // constantB + pd.set(6, 0); // constantC + pd.set(7, M); + pd.set(8, N); + pd.set(9, K); + if (quantize_term) + pd.set(18, quantize_term); + + return pd; +} + +static int test_gemm_block_quant(int M, int N, int K, int bits, int block_size, int broadcast_type_C, int input_scale = 0) +{ + const int quantize_term = weight_block_quantize_term(bits, block_size, input_scale); + + ncnn::Mat A = RandomMat(K, M, -2.f, 2.f); + ncnn::Mat B = RandomMat(K, N, -2.f, 2.f); + + ncnn::Mat B_input_scales; + if (input_scale) + { + B_input_scales = make_input_scales(K); + B = scale_weight_by_input_scales(B, B_input_scales, 1); + } + + ncnn::Mat B_quantized; + ncnn::Mat B_quantize_scales; + ncnn::Mat B_dequantized; + int ret = quantize_weight(B, bits, block_size, B_quantized, B_quantize_scales, B_dequantized); + if (ret != 0) + return ret; + + if (input_scale) + B_dequantized = scale_weight_by_input_scales(B_dequantized, B_input_scales, 0); + + std::vector weights(2); + weights[0] = B_quantized; + weights[1] = B_quantize_scales; + if (input_scale) + weights.push_back(B_input_scales); + + std::vector ref_weights(1); + ref_weights[0] = B_dequantized; + + std::vector inputs; + inputs.push_back(A); + + if (broadcast_type_C == 0) inputs.push_back(RandomMat(1)); + if (broadcast_type_C == 1) inputs.push_back(RandomMat(M)); + if (broadcast_type_C == 2) inputs.push_back(RandomMat(1, M)); + if (broadcast_type_C == 3) inputs.push_back(RandomMat(N, M)); + if (broadcast_type_C == 4) inputs.push_back(RandomMat(N)); + + 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; + + std::vector outputs; + std::vector refs; + ret = test_layer_cpu(ncnn::LayerType::Gemm, make_gemm_param(M, N, K, quantize_term), weights, opt, inputs, 1, outputs, std::vector(), TEST_LAYER_DISABLE_GPU_TESTING); + if (ret != 0) + { + fprintf(stderr, "test_gemm_block_quant failed ret=%d M=%d N=%d K=%d bits=%d block_size=%d broadcast_type_C=%d input_scale=%d\n", ret, M, N, K, bits, block_size, broadcast_type_C, input_scale); + return ret; + } + + ret = test_layer_cpu(ncnn::LayerType::Gemm, make_gemm_param(M, N, K, 0), ref_weights, opt, inputs, 1, refs, std::vector(), TEST_LAYER_DISABLE_GPU_TESTING); + if (ret != 0) + { + fprintf(stderr, "test_gemm_block_quant reference failed ret=%d M=%d N=%d K=%d\n", ret, M, N, K); + return ret; + } + + ret = CompareMat(outputs, refs, 0.001f); + if (ret != 0) + { + fprintf(stderr, "test_gemm_block_quant compare failed M=%d N=%d K=%d bits=%d block_size=%d broadcast_type_C=%d input_scale=%d\n", M, N, K, bits, block_size, broadcast_type_C, input_scale); + return ret; + } + + return 0; +} + +int main() +{ + SRAND(7767517); + +#if !NCNN_WEIGHT_QUANT + ncnn::ParamDict pd = make_gemm_param(3, 4, 5, 410); + + ncnn::Gemm gemm; + if (gemm.load_param(pd) == 0) + { + fprintf(stderr, "test_gemm_block_quant failed NCNN_WEIGHT_QUANT=OFF accepted weight block quantization\n"); + return -1; + } + + return 0; +#else + return 0 + || test_gemm_block_quant(3, 5, 65, 4, 64, -1) + || test_gemm_block_quant(3, 4, 33, 4, 32, 0) + || test_gemm_block_quant(4, 7, 67, 6, 64, 3) + || test_gemm_block_quant(3, 4, 33, 6, 32, 2) + || test_gemm_block_quant(2, 3, 5, 8, 32, 4) + || test_gemm_block_quant(2, 5, 65, 8, 32, -1) + || test_gemm_block_quant(3, 4, 129, 6, 128, -1) + || test_gemm_block_quant(3, 5, 65, 4, 64, -1, 1) + || test_gemm_block_quant(2, 4, 31, 6, 32, 1, 1); +#endif +} diff --git a/tests/test_multiheadattention_block_quant.cpp b/tests/test_multiheadattention_block_quant.cpp new file mode 100644 index 000000000000..cf1927b5d257 --- /dev/null +++ b/tests/test_multiheadattention_block_quant.cpp @@ -0,0 +1,390 @@ +// Copyright 2026 Tencent +// SPDX-License-Identifier: BSD-3-Clause + +#include "testutil.h" + +#include "layer_type.h" +#include "multiheadattention.h" + +#include + +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(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& weights, std::vector& 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& weights, const std::vector& inputs, int top_blob_count, std::vector& 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(), 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 weights; + std::vector 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 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 outputs; + std::vector 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 weights; + std::vector 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 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 outputs; + std::vector 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; +} + +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_kvcache() + || test_multiheadattention_block_quant_kvcache(0, 1) + || test_multiheadattention_block_quant_kvcache(1) + || test_multiheadattention_block_quant_kvcache(1, 1); +#endif +} diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 284d8dac16fb..c35d1b2b0a8b 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -16,11 +16,10 @@ endif() add_subdirectory(caffe) add_subdirectory(mxnet) add_subdirectory(darknet) -if(NCNN_INT8) - add_subdirectory(quantize) -else() - message(WARNING "NCNN_INT8 disabled, quantize tools won't be built") +if(NOT NCNN_INT8) + message(WARNING "NCNN_INT8 disabled, post-training int8 quantize tools won't be built") endif() +add_subdirectory(quantize) add_executable(ncnn2mem ncnn2mem.cpp) target_link_libraries(ncnn2mem PRIVATE ncnn) diff --git a/tools/modelwriter.h b/tools/modelwriter.h index 111310e90528..3a7359d03551 100644 --- a/tools/modelwriter.h +++ b/tools/modelwriter.h @@ -108,6 +108,23 @@ #include "layer/yolodetectionoutput.h" #include "layer/yolov3detectionoutput.h" +static bool modelwriter_is_weight_block_quantize(int quantize_term) +{ + const int weight_bits = quantize_term / 100; + const int format_code = quantize_term % 100 / 10; + const int block_size_code = quantize_term % 10; + + return (weight_bits == 4 || weight_bits == 6 || weight_bits == 8) && (format_code == 0 || format_code == 1) && block_size_code >= 0 && block_size_code <= 2; +} + +static bool modelwriter_weight_block_quantize_has_input_scale(int quantize_term) +{ + if (!modelwriter_is_weight_block_quantize(quantize_term)) + return false; + + return quantize_term % 100 / 10 == 1; +} + // for gen_random_weight #include "../tests/prng.h" @@ -239,6 +256,7 @@ class ModelWriter : public ncnn::Net int fprintf_param_float_array(int id, const ncnn::Mat& m, FILE* pp); int fwrite_weight_tag_data(const ncnn::Mat& data, FILE* bp, float a = -1.2f, float b = 1.2f); + int fwrite_weight_tag_data_no_random(const ncnn::Mat& data, FILE* bp); int fwrite_weight_data(const ncnn::Mat& data, FILE* bp, float a = -1.2f, float b = 1.2f); int save(const char* parampath, const char* binpath); @@ -665,6 +683,16 @@ int ModelWriter::fwrite_weight_tag_data(const ncnn::Mat& data, FILE* bp, float a return 0; } +int ModelWriter::fwrite_weight_tag_data_no_random(const ncnn::Mat& data, FILE* bp) +{ + const int old_gen_random_weight = gen_random_weight; + gen_random_weight = false; + int ret = fwrite_weight_tag_data(data, bp); + gen_random_weight = old_gen_random_weight; + + return ret; +} + int ModelWriter::fwrite_weight_data(const ncnn::Mat& data, FILE* bp, float a, float b) { int p0 = ftell(bp); @@ -1840,27 +1868,48 @@ int ModelWriter::save(const char* parampath, const char* binpath) fprintf_param_value(" 12=%d", output_elempack) fprintf_param_value(" 13=%d", output_elemtype) fprintf_param_value(" 14=%d", output_transpose) - fprintf_param_value(" 18=%d", int8_scale_term) + fprintf_param_value(" 18=%d", quantize_term) fprintf_param_value(" 20=%d", constant_TILE_M) fprintf_param_value(" 21=%d", constant_TILE_N) fprintf_param_value(" 22=%d", constant_TILE_K) + const bool weight_block_quantize + = modelwriter_is_weight_block_quantize(op->quantize_term); + if (op->constantA == 1) { fwrite_weight_tag_data(op->A_data, bp); } if (op->constantB == 1) { - fwrite_weight_tag_data(op->B_data, bp); + if (weight_block_quantize) + { + fwrite_weight_tag_data_no_random(op->B_data, bp); + } + else + { + fwrite_weight_tag_data(op->B_data, bp); + } } if (op->constantC == 1 && op->constant_broadcast_type_C != -1) { fwrite_weight_tag_data(op->C_data, bp); } + if (weight_block_quantize) + { + if (op->constantB == 1) + { + fwrite_weight_data(op->B_data_quantize_scales, bp, 0.001, 1); + if (modelwriter_weight_block_quantize_has_input_scale(op->quantize_term)) + { + fwrite_weight_data(op->B_data_input_scales, bp, 0.001, 1); + } + } + } #if NCNN_INT8 // write int8_scale data - if (op->int8_scale_term) + if (op->quantize_term && !weight_block_quantize) { if (op->constantA == 1) { @@ -2125,20 +2174,65 @@ int ModelWriter::save(const char* parampath, const char* binpath) fprintf_param_value(" 5=%d", attn_mask) fprintf_param_value(" 6=%e", scale) fprintf_param_value(" 7=%d", kv_cache) - fprintf_param_value(" 18=%d", int8_scale_term) + fprintf_param_value(" 18=%d", quantize_term) + + const bool weight_block_quantize + = modelwriter_is_weight_block_quantize(op->quantize_term); - fwrite_weight_tag_data(op->q_weight_data, bp); + if (weight_block_quantize) + { + fwrite_weight_tag_data_no_random(op->q_weight_data, bp); + } + else + { + fwrite_weight_tag_data(op->q_weight_data, bp); + } fwrite_weight_data(op->q_bias_data, bp); - fwrite_weight_tag_data(op->k_weight_data, bp); + if (weight_block_quantize) + { + fwrite_weight_tag_data_no_random(op->k_weight_data, bp); + } + else + { + fwrite_weight_tag_data(op->k_weight_data, bp); + } fwrite_weight_data(op->k_bias_data, bp); - fwrite_weight_tag_data(op->v_weight_data, bp); + if (weight_block_quantize) + { + fwrite_weight_tag_data_no_random(op->v_weight_data, bp); + } + else + { + fwrite_weight_tag_data(op->v_weight_data, bp); + } fwrite_weight_data(op->v_bias_data, bp); - fwrite_weight_tag_data(op->out_weight_data, bp); + if (weight_block_quantize) + { + fwrite_weight_tag_data_no_random(op->out_weight_data, bp); + } + else + { + fwrite_weight_tag_data(op->out_weight_data, bp); + } fwrite_weight_data(op->out_bias_data, bp); + if (weight_block_quantize) + { + fwrite_weight_data(op->q_weight_data_quantize_scales, bp, 0.001, 1); + fwrite_weight_data(op->k_weight_data_quantize_scales, bp, 0.001, 1); + fwrite_weight_data(op->v_weight_data_quantize_scales, bp, 0.001, 1); + fwrite_weight_data(op->out_weight_data_quantize_scales, bp, 0.001, 1); + if (modelwriter_weight_block_quantize_has_input_scale(op->quantize_term)) + { + fwrite_weight_data(op->q_weight_data_input_scales, bp, 0.001, 1); + fwrite_weight_data(op->k_weight_data_input_scales, bp, 0.001, 1); + fwrite_weight_data(op->v_weight_data_input_scales, bp, 0.001, 1); + fwrite_weight_data(op->out_weight_data_input_scales, bp, 0.001, 1); + } + } #if NCNN_INT8 // write int8_scale data - if (op->int8_scale_term) + else if (op->quantize_term) { fwrite_weight_data(op->q_weight_data_int8_scales, bp, 90, 100); fwrite_weight_data(op->k_weight_data_int8_scales, bp, 90, 100); diff --git a/tools/quantize/CMakeLists.txt b/tools/quantize/CMakeLists.txt index 72c76d135015..553b6f1f04b6 100644 --- a/tools/quantize/CMakeLists.txt +++ b/tools/quantize/CMakeLists.txt @@ -1,43 +1,54 @@ -if(NCNN_PIXEL) - if(NOT NCNN_SIMPLEOCV) - find_package(OpenCV QUIET COMPONENTS opencv_world) - # for opencv 2.4 on ubuntu 16.04, there is no opencv_world but OpenCV_FOUND will be TRUE - if("${OpenCV_LIBS}" STREQUAL "") - set(OpenCV_FOUND FALSE) - endif() - if(NOT OpenCV_FOUND) - find_package(OpenCV QUIET COMPONENTS core highgui imgproc imgcodecs) +add_executable(ncnnllm2int468 ncnnllm2int468.cpp) +target_link_libraries(ncnnllm2int468 PRIVATE ncnn) + +add_executable(ncnnllm2table ncnnllm2table.cpp) +target_link_libraries(ncnnllm2table PRIVATE ncnn) + +set_property(TARGET ncnnllm2int468 PROPERTY FOLDER "tools/optimization") +set_property(TARGET ncnnllm2table PROPERTY FOLDER "tools/optimization") +ncnn_install_tool(ncnnllm2int468) +ncnn_install_tool(ncnnllm2table) + +if(NCNN_INT8) + if(NCNN_PIXEL) + if(NOT NCNN_SIMPLEOCV) + find_package(OpenCV QUIET COMPONENTS opencv_world) + # for opencv 2.4 on ubuntu 16.04, there is no opencv_world but OpenCV_FOUND will be TRUE + if("${OpenCV_LIBS}" STREQUAL "") + set(OpenCV_FOUND FALSE) + endif() + if(NOT OpenCV_FOUND) + find_package(OpenCV QUIET COMPONENTS core highgui imgproc imgcodecs) + endif() + if(NOT OpenCV_FOUND) + find_package(OpenCV QUIET COMPONENTS core highgui imgproc) + endif() endif() - if(NOT OpenCV_FOUND) - find_package(OpenCV QUIET COMPONENTS core highgui imgproc) + + if(OpenCV_FOUND) + add_executable(ncnn2table ncnn2table.cpp) + target_include_directories(ncnn2table PRIVATE ${OpenCV_INCLUDE_DIRS}) + target_link_libraries(ncnn2table PRIVATE ncnn ${OpenCV_LIBS}) + elseif(NCNN_SIMPLEOCV) + add_executable(ncnn2table ncnn2table.cpp) + target_compile_definitions(ncnn2table PUBLIC USE_NCNN_SIMPLEOCV) + target_link_libraries(ncnn2table PRIVATE ncnn) + else() + add_executable(ncnn2table ncnn2table.cpp imreadwrite.cpp) + target_compile_definitions(ncnn2table PUBLIC USE_LOCAL_IMREADWRITE) + target_link_libraries(ncnn2table PRIVATE ncnn) endif() - endif() - set(OpenCV_FOUND FALSE) - - if(OpenCV_FOUND) - add_executable(ncnn2table ncnn2table.cpp) - target_include_directories(ncnn2table PRIVATE ${OpenCV_INCLUDE_DIRS}) - target_link_libraries(ncnn2table PRIVATE ncnn ${OpenCV_LIBS}) - elseif(NCNN_SIMPLEOCV) - add_executable(ncnn2table ncnn2table.cpp) - target_compile_definitions(ncnn2table PUBLIC USE_NCNN_SIMPLEOCV) - target_link_libraries(ncnn2table PRIVATE ncnn) - else() - add_executable(ncnn2table ncnn2table.cpp imreadwrite.cpp) - target_compile_definitions(ncnn2table PUBLIC USE_LOCAL_IMREADWRITE) - target_link_libraries(ncnn2table PRIVATE ncnn) + # add ncnn2table tool to a virtual project group + set_property(TARGET ncnn2table PROPERTY FOLDER "tools/optimization") + ncnn_install_tool(ncnn2table) endif() - # add ncnn2table tool to a virtual project group - set_property(TARGET ncnn2table PROPERTY FOLDER "tools/optimization") -endif() - -add_executable(ncnn2int8 ncnn2int8.cpp) -target_link_libraries(ncnn2int8 PRIVATE ncnn) + add_executable(ncnn2int8 ncnn2int8.cpp) + target_link_libraries(ncnn2int8 PRIVATE ncnn) -# add ncnn2int8 tool to a virtual project group -set_property(TARGET ncnn2int8 PROPERTY FOLDER "tools/optimization") -ncnn_install_tool(ncnn2table) -ncnn_install_tool(ncnn2int8) + # add ncnn2int8 tool to a virtual project group + set_property(TARGET ncnn2int8 PROPERTY FOLDER "tools/optimization") + ncnn_install_tool(ncnn2int8) +endif() diff --git a/tools/quantize/ncnn2int8.cpp b/tools/quantize/ncnn2int8.cpp index 41a2115b094a..2e7706bd25f8 100644 --- a/tools/quantize/ncnn2int8.cpp +++ b/tools/quantize/ncnn2int8.cpp @@ -697,7 +697,7 @@ int NetQuantize::quantize_gemm() gemm->B_data = B_data_int8; } - gemm->int8_scale_term = 2; + gemm->quantize_term = 2; } return 0; diff --git a/tools/quantize/ncnnllm2int468.cpp b/tools/quantize/ncnnllm2int468.cpp new file mode 100644 index 000000000000..936895b90f9f --- /dev/null +++ b/tools/quantize/ncnnllm2int468.cpp @@ -0,0 +1,1110 @@ +// Copyright 2019 BUG1989 +// Copyright 2026 Tencent +// SPDX-License-Identifier: BSD-3-Clause + +#ifdef _MSC_VER +#define _CRT_SECURE_NO_DEPRECATE +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +// ncnn public header +#include "datareader.h" +#include "layer.h" +#include "layer_type.h" +#include "net.h" + +// ncnn private header +#include "../modelwriter.h" +#include "ncnnllm_quant.h" + +class DataReaderFromEmpty : public ncnn::DataReader +{ +public: + virtual int scan(const char* /*format*/, void* /*p*/) const + { + return 0; + } + virtual size_t read(void* buf, size_t size) const + { + memset(buf, 0, size); + return size; + } +}; + +class LLMWeightScale +{ +public: + LLMWeightScale() + { + bits = 0; + block_size = 0; + used = false; + } + +public: + int bits; + int block_size; + std::string qweight; + std::string method; + ncnn::Mat scales; + bool used; +}; + +static bool qweight_path_is_absolute(const char* path) +{ + if (path[0] == '/') + return true; + +#ifdef _WIN32 + if (strlen(path) >= 3 && path[1] == ':' && (path[2] == '\\' || path[2] == '/')) + return true; +#endif + + return false; +} + +static std::string resolve_qweight_path(const char* tablepath, const char* path) +{ + if (qweight_path_is_absolute(path)) + return std::string(path); + + const char* slash = strrchr(tablepath, '/'); +#ifdef _WIN32 + const char* backslash = strrchr(tablepath, '\\'); + if (backslash && (!slash || backslash > slash)) + slash = backslash; +#endif + + if (!slash) + return std::string(path); + + return std::string(tablepath, slash - tablepath + 1) + path; +} + +static inline int sign_extend(int v, int bits) +{ + const int sign_bit = 1 << (bits - 1); + return (v ^ sign_bit) - sign_bit; +} + +static inline int unpack_signed_weight(const unsigned char* ptr, int k, int bits, int packed_k_bytes) +{ + const int bit_offset = k * bits; + const int byte_offset = bit_offset / 8; + const int bit_shift = bit_offset % 8; + + unsigned int v = ptr[byte_offset]; + if (byte_offset + 1 < packed_k_bytes) + v |= (unsigned int)ptr[byte_offset + 1] << 8; + if (byte_offset + 2 < packed_k_bytes) + v |= (unsigned int)ptr[byte_offset + 2] << 16; + + const int mask = (1 << bits) - 1; + return sign_extend((v >> bit_shift) & mask, bits); +} + +static bool parse_int_string(const char* s, int& v) +{ + int nconsumed = 0; + if (sscanf(s, "%d%n", &v, &nconsumed) != 1 || s[nconsumed] != '\0') + return false; + + return true; +} + +static bool parse_float_string(const char* s, float& v) +{ + int nconsumed = 0; + if (sscanf(s, "%f%n", &v, &nconsumed) != 1 || s[nconsumed] != '\0') + return false; + + return true; +} + +static bool read_line(FILE* fp, std::vector& line) +{ + line.clear(); + + char buf[4096]; + while (fgets(buf, sizeof(buf), fp)) + { + const size_t len = strlen(buf); + line.insert(line.end(), buf, buf + len); + if (len > 0 && buf[len - 1] == '\n') + break; + } + + if (line.empty()) + return false; + + line.push_back('\0'); + return true; +} + +static bool read_llm_scale_table(const char* filepath, std::map& llm_scale_table) +{ + llm_scale_table.clear(); + + FILE* fp = fopen(filepath, "rb"); + if (!fp) + { + fprintf(stderr, "Open %s failed.\n", filepath); + return false; + } + + std::string key_str; + std::vector scales; + + std::vector line; + char* pch = NULL; + int lineno = 0; + + while (read_line(fp, line)) + { + lineno++; + line[strcspn(line.data(), "\r\n")] = 0; + + pch = strtok(line.data(), " \t"); + if (pch == NULL) + continue; + if (pch[0] == '#') + continue; + + char key[256]; + sscanf(pch, "%255s", key); + key_str = key; + + if (llm_scale_table.find(key_str) != llm_scale_table.end()) + { + fprintf(stderr, "%s:%d duplicate key %s\n", filepath, lineno, key); + fclose(fp); + return false; + } + + LLMWeightScale scale; + scales.clear(); + + bool coeff_started = false; + bool has_bits = false; + bool has_block = false; + bool has_qweight = false; + bool has_method = false; + + while ((pch = strtok(NULL, " \t")) != NULL) + { + char* eqs = strchr(pch, '='); + if (eqs && !coeff_started) + { + eqs[0] = '\0'; + const char* k = pch; + const char* v = eqs + 1; + + if (v[0] == '\0') + { + fprintf(stderr, "%s:%d malformed metadata %s=\n", filepath, lineno, k); + fclose(fp); + return false; + } + + if (strcmp(k, "bits") == 0) + { + if (has_bits) + { + fprintf(stderr, "%s:%d duplicate metadata %s\n", filepath, lineno, k); + fclose(fp); + return false; + } + + has_bits = true; + if (!parse_int_string(v, scale.bits)) + { + fprintf(stderr, "%s:%d malformed bits=%s\n", filepath, lineno, v); + fclose(fp); + return false; + } + } + else if (strcmp(k, "block") == 0) + { + if (has_block) + { + fprintf(stderr, "%s:%d duplicate metadata %s\n", filepath, lineno, k); + fclose(fp); + return false; + } + + has_block = true; + if (!parse_int_string(v, scale.block_size)) + { + fprintf(stderr, "%s:%d malformed block=%s\n", filepath, lineno, v); + fclose(fp); + return false; + } + } + else if (strcmp(k, "qweight") == 0) + { + if (has_qweight) + { + fprintf(stderr, "%s:%d duplicate metadata %s\n", filepath, lineno, k); + fclose(fp); + return false; + } + + has_qweight = true; + scale.qweight = resolve_qweight_path(filepath, v); + } + else if (strcmp(k, "method") == 0) + { + if (has_method) + { + fprintf(stderr, "%s:%d duplicate metadata %s\n", filepath, lineno, k); + fclose(fp); + return false; + } + + has_method = true; + scale.method = v; + } + else if (strcmp(k, "zero_point") == 0 || strcmp(k, "zp") == 0 || strcmp(k, "g_idx") == 0) + { + fprintf(stderr, "%s:%d unsupported metadata %s\n", filepath, lineno, k); + fclose(fp); + return false; + } + else + { + fprintf(stderr, "%s:%d unsupported metadata %s\n", filepath, lineno, k); + fclose(fp); + return false; + } + } + else + { + if (eqs && coeff_started) + { + fprintf(stderr, "%s:%d key=value token after coefficients started\n", filepath, lineno); + fclose(fp); + return false; + } + + coeff_started = true; + + float coeff = 0.f; + if (!parse_float_string(pch, coeff)) + { + fprintf(stderr, "%s:%d malformed coefficient %s\n", filepath, lineno, pch); + fclose(fp); + return false; + } + + scales.push_back(coeff); + } + } + + scale.scales = ncnn::Mat((int)scales.size(), (void*)scales.data()).clone(); + llm_scale_table[key_str] = scale; + } + + fclose(fp); + + return true; +} + +static int make_input_scaled_weight(const ncnn::Mat& weight_data, const ncnn::Mat& input_scales, ncnn::Mat& weight_data_scaled) +{ + const int K = weight_data.w; + const int N = weight_data.h; + + if (input_scales.w != K) + { + fprintf(stderr, "input scale count mismatch expected=%d got=%d\n", K, input_scales.w); + return -1; + } + + weight_data_scaled.create(K, N); + if (weight_data_scaled.empty()) + return -100; + + const float* input_scale_ptr = input_scales; + for (int n = 0; n < N; n++) + { + const float* ptr = weight_data.row(n); + float* outptr = weight_data_scaled.row(n); + + for (int k = 0; k < K; k++) + outptr[k] = ptr[k] / input_scale_ptr[k]; + } + + return 0; +} + +class NetQuantize : public ModelWriter +{ +public: + NetQuantize(); + +public: + int quantize_gemm(int block_size, int weight_bits, int method); + int quantize_gemm_from_table(std::map& llm_scale_table); + int quantize_multiheadattention(int block_size, int weight_bits, int method); + int quantize_multiheadattention_from_table(std::map& llm_scale_table); +}; + +NetQuantize::NetQuantize() + : ModelWriter() +{ +} + +static void print_usage(const char* argv0) +{ + fprintf(stderr, "Usage: %s [inparam] [inbin] [outparam] [outbin] [(key=value)...]\n", argv0); + fprintf(stderr, " %s [inparam] [inbin] [outparam] [outbin] [llm.table] [(key=value)...]\n", argv0); + fprintf(stderr, " method=minmax/mseclip\n"); + fprintf(stderr, " bits=4/6/8\n"); + fprintf(stderr, " block=32/64/128\n"); + fprintf(stderr, "Sample usage:\n"); + fprintf(stderr, " %s model.param model.bin model-int4.param model-int4.bin method=mseclip bits=4 block=64\n", argv0); + fprintf(stderr, " %s model.param model.bin model-int4.param model-int4.bin model.llm.table\n", argv0); +} + +static int llm_table_row_to_scales(const char* key, const LLMWeightScale& scale, int K, int N, int& weight_bits, int& block_size, ncnn::Mat& weight_data_quantize_scales) +{ + if (scale.bits == 0 || scale.block_size == 0) + { + fprintf(stderr, "%s missing mandatory metadata\n", key); + return -1; + } + + weight_bits = scale.bits; + block_size = scale.block_size; + + if (llm_weight_block_quantize_term(weight_bits, block_size) == 0) + { + fprintf(stderr, "%s unsupported bits=%d block=%d\n", key, weight_bits, block_size); + return -1; + } + + const int block_count = (K + block_size - 1) / block_size; + const size_t weight_scale_count = (size_t)N * block_count; + + if ((size_t)scale.scales.w != weight_scale_count) + { + fprintf(stderr, "%s coefficient count mismatch expected=%zu got=%d\n", key, weight_scale_count, scale.scales.w); + return -1; + } + + weight_data_quantize_scales.create(block_count, N); + if (weight_data_quantize_scales.empty()) + return -100; + + memcpy(weight_data_quantize_scales.data, scale.scales.data, weight_scale_count * sizeof(float)); + + return 0; +} + +static int read_qweight_file(const char* key, const LLMWeightScale& scale, int K, int N, int weight_bits, ncnn::Mat& weight_data_quantized) +{ + const int packed_k_bytes = llm_weight_quantize_packed_k_bytes(K, weight_bits); + if (packed_k_bytes <= 0) + { + fprintf(stderr, "%s qweight packed byte count is invalid K=%d bits=%d\n", key, K, weight_bits); + return -1; + } + const size_t qweight_bytes = (size_t)N * packed_k_bytes; + + FILE* fp = fopen(scale.qweight.c_str(), "rb"); + if (!fp) + { + fprintf(stderr, "%s fopen qweight %s failed\n", key, scale.qweight.c_str()); + return -1; + } + + fseek(fp, 0, SEEK_END); + const long file_size = ftell(fp); + fseek(fp, 0, SEEK_SET); + + if (file_size != (long)qweight_bytes) + { + fprintf(stderr, "%s qweight byte count mismatch expected=%zu got=%ld\n", key, qweight_bytes, file_size); + fclose(fp); + return -1; + } + + weight_data_quantized.create(packed_k_bytes, N, (size_t)1u); + if (weight_data_quantized.empty()) + { + fclose(fp); + return -100; + } + + if (fread(weight_data_quantized.data, 1, qweight_bytes, fp) != qweight_bytes) + { + fprintf(stderr, "%s fread qweight %s failed\n", key, scale.qweight.c_str()); + fclose(fp); + return -1; + } + + fclose(fp); + + const int used_bits = (int)(((size_t)K * weight_bits) % 8); + const unsigned char tail_mask = used_bits == 0 ? 0 : (unsigned char)(0xffu << used_bits); + const int qmin = -(1 << (weight_bits - 1)); + + for (int n = 0; n < N; n++) + { + const unsigned char* qptr = weight_data_quantized.row(n); + + if (tail_mask && (qptr[packed_k_bytes - 1] & tail_mask)) + { + fprintf(stderr, "%s qweight tail padding bits are not zero row=%d\n", key, n); + return -1; + } + + for (int k = 0; k < K; k++) + { + const int q = unpack_signed_weight(qptr, k, weight_bits, packed_k_bytes); + if (q == qmin) + { + fprintf(stderr, "%s qweight contains unsupported value row=%d k=%d q=%d\n", key, n, k, q); + return -1; + } + } + } + + return 0; +} + +static int llm_table_row_to_qweight(const char* key, const LLMWeightScale& scale, int K, int N, int& weight_bits, int& block_size, ncnn::Mat& weight_data_quantize_scales, ncnn::Mat& weight_data_quantized) +{ + if (scale.bits == 0 || scale.block_size == 0 || scale.qweight.empty()) + { + fprintf(stderr, "%s missing mandatory metadata\n", key); + return -1; + } + + weight_bits = scale.bits; + block_size = scale.block_size; + + if (llm_weight_block_quantize_term(weight_bits, block_size) == 0) + { + fprintf(stderr, "%s unsupported bits=%d block=%d\n", key, weight_bits, block_size); + return -1; + } + + const int block_count = (K + block_size - 1) / block_size; + const size_t weight_scale_count = (size_t)N * block_count; + + if ((size_t)scale.scales.w != weight_scale_count) + { + fprintf(stderr, "%s coefficient count mismatch expected=%zu got=%d\n", key, weight_scale_count, scale.scales.w); + return -1; + } + + weight_data_quantize_scales.create(block_count, N); + if (weight_data_quantize_scales.empty()) + return -100; + + memcpy(weight_data_quantize_scales.data, scale.scales.data, weight_scale_count * sizeof(float)); + const float* scale_ptr = weight_data_quantize_scales; + for (size_t i = 0; i < weight_scale_count; i++) + { + if (!(scale_ptr[i] > 0.f) || scale_ptr[i] > FLT_MAX) + { + fprintf(stderr, "%s invalid weight scale index=%zu coeff=%f\n", key, i, scale_ptr[i]); + return -1; + } + } + + return read_qweight_file(key, scale, K, N, weight_bits, weight_data_quantized); +} + +static int llm_table_row_to_input_scales(const char* key, const LLMWeightScale& scale, int K, ncnn::Mat& input_scales) +{ + if (scale.scales.w != K) + { + fprintf(stderr, "%s coefficient count mismatch expected=%d got=%d\n", key, K, scale.scales.w); + return -1; + } + + const float* ptr = scale.scales; + for (int k = 0; k < K; k++) + { + if (!(ptr[k] > 0.f) || ptr[k] > FLT_MAX) + { + fprintf(stderr, "%s invalid input scale index=%d coeff=%f\n", key, k, ptr[k]); + return -1; + } + } + + input_scales = scale.scales.clone(); + if (input_scales.empty()) + return -100; + + return 0; +} + +static void print_unused_llm_table_rows(const std::map& llm_scale_table) +{ + for (std::map::const_iterator it = llm_scale_table.begin(); it != llm_scale_table.end(); ++it) + { + if (!it->second.used) + fprintf(stderr, "warning: unused llm table row %s\n", it->first.c_str()); + } +} + +int NetQuantize::quantize_gemm(int block_size, int weight_bits, int method) +{ + const int quantize_term = llm_weight_block_quantize_term(weight_bits, block_size); + if (quantize_term == 0) + { + fprintf(stderr, "unsupported bits=%d or block=%d\n", weight_bits, block_size); + return -1; + } + + int quantized_count = 0; + + for (size_t i = 0; i < layers.size(); i++) + { + if (layers[i]->type != "Gemm") + continue; + + ncnn::Gemm* gemm = (ncnn::Gemm*)layers[i]; + + const char* reason = 0; + if (!is_supported_llm_gemm(gemm, &reason)) + { + print_skip_gemm(gemm, reason); + continue; + } + + fprintf(stderr, "quantize_gemm bits=%d block_size=%d term=%d %s\n", weight_bits, block_size, quantize_term, gemm_name(gemm)); + + ncnn::Mat B_data_quantized; + ncnn::Mat B_data_quantize_scales; + const int ret = make_and_pack_gemm_B(gemm->B_data, block_size, weight_bits, method, B_data_quantized, B_data_quantize_scales); + if (ret != 0) + return ret; + + gemm->B_data = B_data_quantized; + gemm->B_data_quantize_scales = B_data_quantize_scales; + gemm->quantize_term = quantize_term; + + quantized_count++; + } + + fprintf(stderr, "quantized %d Gemm layers\n", quantized_count); + + return 0; +} + +int NetQuantize::quantize_gemm_from_table(std::map& llm_scale_table) +{ + int quantized_count = 0; + + for (size_t i = 0; i < layers.size(); i++) + { + if (layers[i]->type != "Gemm") + continue; + + ncnn::Gemm* gemm = (ncnn::Gemm*)layers[i]; + char key[256]; + snprintf(key, 256, "%s_param_1", gemm_name(gemm)); + char input_scale_key[256]; + snprintf(input_scale_key, 256, "%s_param_1_input_scale", gemm_name(gemm)); + + const char* reason = 0; + if (!is_supported_llm_gemm(gemm, &reason)) + { + std::map::iterator iter = llm_scale_table.find(key); + if (iter != llm_scale_table.end()) + { + iter->second.used = true; + fprintf(stderr, "table row %s targets unsupported Gemm %s\n", key, reason); + return -1; + } + + print_skip_gemm(gemm, reason); + continue; + } + + std::map::iterator iter = llm_scale_table.find(key); + if (iter == llm_scale_table.end()) + { + fprintf(stderr, "skip_gemm %s missing table row %s\n", gemm_name(gemm), key); + continue; + } + + LLMWeightScale& scale = iter->second; + scale.used = true; + + const int constantN = gemm->constantN; + const int constantK = gemm->constantK; + int weight_bits = 0; + int block_size = 0; + ncnn::Mat B_data_quantize_scales; + ncnn::Mat B_data_quantized; + const bool qweight_row = !scale.qweight.empty(); + int ret = 0; + if (qweight_row) + ret = llm_table_row_to_qweight(key, scale, constantK, constantN, weight_bits, block_size, B_data_quantize_scales, B_data_quantized); + else + ret = llm_table_row_to_scales(key, scale, constantK, constantN, weight_bits, block_size, B_data_quantize_scales); + if (ret != 0) + return ret; + + std::map::iterator input_scale_iter = llm_scale_table.find(input_scale_key); + const bool has_input_scale = input_scale_iter != llm_scale_table.end(); + + if (qweight_row && has_input_scale) + { + input_scale_iter->second.used = true; + fprintf(stderr, "%s does not support input_scale with qweight yet\n", key); + return -1; + } + + ncnn::Mat B_data_input_scales; + if (has_input_scale) + { + input_scale_iter->second.used = true; + + ret = llm_table_row_to_input_scales(input_scale_key, input_scale_iter->second, constantK, B_data_input_scales); + if (ret != 0) + return ret; + } + + const int quantize_term = llm_weight_block_quantize_term(weight_bits, block_size, has_input_scale); + fprintf(stderr, "quantize_gemm table bits=%d block_size=%d term=%d method=%s %s\n", weight_bits, block_size, quantize_term, scale.method.c_str(), gemm_name(gemm)); + + if (!qweight_row) + { + if (has_input_scale) + { + ncnn::Mat B_data_scaled; + ret = make_input_scaled_weight(gemm->B_data, B_data_input_scales, B_data_scaled); + if (ret != 0) + return ret; + + ret = pack_gemm_B_from_scales(B_data_scaled, B_data_quantize_scales, block_size, weight_bits, B_data_quantized); + } + else + { + ret = pack_gemm_B_from_scales(gemm->B_data, B_data_quantize_scales, block_size, weight_bits, B_data_quantized); + } + if (ret != 0) + return ret; + } + + gemm->B_data = B_data_quantized; + gemm->B_data_quantize_scales = B_data_quantize_scales; + gemm->B_data_input_scales = B_data_input_scales; + gemm->quantize_term = quantize_term; + + quantized_count++; + } + + fprintf(stderr, "quantized %d Gemm layers\n", quantized_count); + + return 0; +} + +int NetQuantize::quantize_multiheadattention(int block_size, int weight_bits, int method) +{ + const int quantize_term = llm_weight_block_quantize_term(weight_bits, block_size); + if (quantize_term == 0) + { + fprintf(stderr, "unsupported bits=%d or block=%d\n", weight_bits, block_size); + return -1; + } + + int quantized_count = 0; + + for (size_t i = 0; i < layers.size(); i++) + { + if (layers[i]->type != "MultiHeadAttention") + continue; + + ncnn::MultiHeadAttention* mha = (ncnn::MultiHeadAttention*)layers[i]; + + const char* reason = 0; + if (!is_supported_llm_multiheadattention(mha, &reason)) + { + print_skip_multiheadattention(mha, reason); + continue; + } + + fprintf(stderr, "quantize_multiheadattention bits=%d block_size=%d term=%d %s\n", weight_bits, block_size, quantize_term, multiheadattention_name(mha)); + + const int qdim = mha->weight_data_size / mha->embed_dim; + + const ncnn::Mat q_weight_data = mha->q_weight_data.reshape(qdim, mha->embed_dim); + const ncnn::Mat k_weight_data = mha->k_weight_data.reshape(mha->kdim, mha->embed_dim); + const ncnn::Mat v_weight_data = mha->v_weight_data.reshape(mha->vdim, mha->embed_dim); + const ncnn::Mat out_weight_data = mha->out_weight_data.reshape(mha->embed_dim, qdim); + + 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_quantize_scales; + ncnn::Mat k_weight_data_quantize_scales; + ncnn::Mat v_weight_data_quantize_scales; + ncnn::Mat out_weight_data_quantize_scales; + + int ret = make_and_pack_gemm_B(q_weight_data, block_size, weight_bits, method, q_weight_data_quantized, q_weight_data_quantize_scales); + if (ret != 0) + return ret; + ret = make_and_pack_gemm_B(k_weight_data, block_size, weight_bits, method, k_weight_data_quantized, k_weight_data_quantize_scales); + if (ret != 0) + return ret; + ret = make_and_pack_gemm_B(v_weight_data, block_size, weight_bits, method, v_weight_data_quantized, v_weight_data_quantize_scales); + if (ret != 0) + return ret; + ret = make_and_pack_gemm_B(out_weight_data, block_size, weight_bits, method, out_weight_data_quantized, out_weight_data_quantize_scales); + if (ret != 0) + return ret; + + mha->q_weight_data = q_weight_data_quantized; + mha->k_weight_data = k_weight_data_quantized; + mha->v_weight_data = v_weight_data_quantized; + mha->out_weight_data = out_weight_data_quantized; + mha->q_weight_data_quantize_scales = q_weight_data_quantize_scales; + mha->k_weight_data_quantize_scales = k_weight_data_quantize_scales; + mha->v_weight_data_quantize_scales = v_weight_data_quantize_scales; + mha->out_weight_data_quantize_scales = out_weight_data_quantize_scales; + mha->quantize_term = quantize_term; + + quantized_count++; + } + + fprintf(stderr, "quantized %d MultiHeadAttention layers\n", quantized_count); + + return 0; +} + +int NetQuantize::quantize_multiheadattention_from_table(std::map& llm_scale_table) +{ + int quantized_count = 0; + + for (size_t i = 0; i < layers.size(); i++) + { + if (layers[i]->type != "MultiHeadAttention") + continue; + + ncnn::MultiHeadAttention* mha = (ncnn::MultiHeadAttention*)layers[i]; + + char keys[4][256]; + char input_scale_keys[4][256]; + for (int j = 0; j < 4; j++) + { + snprintf(keys[j], 256, "%s_param_%d", multiheadattention_name(mha), j); + snprintf(input_scale_keys[j], 256, "%s_param_%d_input_scale", multiheadattention_name(mha), j); + } + + std::map::iterator iters[4]; + int present_count = 0; + for (int j = 0; j < 4; j++) + { + iters[j] = llm_scale_table.find(keys[j]); + if (iters[j] != llm_scale_table.end()) + present_count++; + } + + if (present_count == 0) + continue; + + for (int j = 0; j < 4; j++) + { + if (iters[j] != llm_scale_table.end()) + iters[j]->second.used = true; + } + + if (present_count != 4) + { + fprintf(stderr, "MultiHeadAttention %s requires all table rows %s %s %s %s\n", multiheadattention_name(mha), keys[0], keys[1], keys[2], keys[3]); + return -1; + } + + const char* reason = 0; + if (!is_supported_llm_multiheadattention(mha, &reason)) + { + fprintf(stderr, "table rows target unsupported MultiHeadAttention %s %s\n", multiheadattention_name(mha), reason); + return -1; + } + + const int qdim = mha->weight_data_size / mha->embed_dim; + + const int Ks[4] = {qdim, mha->kdim, mha->vdim, mha->embed_dim}; + const int Ns[4] = {mha->embed_dim, mha->embed_dim, mha->embed_dim, qdim}; + + std::map::iterator input_scale_iters[4]; + int input_scale_present_count = 0; + for (int j = 0; j < 4; j++) + { + input_scale_iters[j] = llm_scale_table.find(input_scale_keys[j]); + if (input_scale_iters[j] != llm_scale_table.end()) + input_scale_present_count++; + } + + if (input_scale_present_count != 0) + { + for (int j = 0; j < 4; j++) + { + if (input_scale_iters[j] != llm_scale_table.end()) + input_scale_iters[j]->second.used = true; + } + + if (input_scale_present_count != 4) + { + fprintf(stderr, "MultiHeadAttention %s requires all input scale table rows %s %s %s %s\n", multiheadattention_name(mha), input_scale_keys[0], input_scale_keys[1], input_scale_keys[2], input_scale_keys[3]); + return -1; + } + } + + int weight_bits[4]; + int block_size[4]; + bool qweight_rows[4]; + ncnn::Mat weight_data_quantize_scales[4]; + ncnn::Mat weight_data_quantized[4]; + for (int j = 0; j < 4; j++) + { + qweight_rows[j] = !iters[j]->second.qweight.empty(); + + int ret = 0; + if (qweight_rows[j]) + ret = llm_table_row_to_qweight(keys[j], iters[j]->second, Ks[j], Ns[j], weight_bits[j], block_size[j], weight_data_quantize_scales[j], weight_data_quantized[j]); + else + ret = llm_table_row_to_scales(keys[j], iters[j]->second, Ks[j], Ns[j], weight_bits[j], block_size[j], weight_data_quantize_scales[j]); + if (ret != 0) + return ret; + } + + for (int j = 1; j < 4; j++) + { + if (weight_bits[j] != weight_bits[0] || block_size[j] != block_size[0]) + { + fprintf(stderr, "MultiHeadAttention %s table rows require same bits/block for q/k/v/out\n", multiheadattention_name(mha)); + return -1; + } + + if (qweight_rows[j] != qweight_rows[0]) + { + fprintf(stderr, "MultiHeadAttention %s table rows require same format for q/k/v/out\n", multiheadattention_name(mha)); + return -1; + } + } + + ncnn::Mat input_scales[4]; + const bool has_input_scale = input_scale_present_count == 4; + if (has_input_scale) + { + for (int j = 0; j < 4; j++) + { + if (qweight_rows[j]) + { + fprintf(stderr, "MultiHeadAttention %s does not support input_scale with qweight yet\n", multiheadattention_name(mha)); + return -1; + } + } + } + if (has_input_scale) + { + for (int j = 0; j < 4; j++) + { + int ret = llm_table_row_to_input_scales(input_scale_keys[j], input_scale_iters[j]->second, Ks[j], input_scales[j]); + if (ret != 0) + return ret; + } + } + + const int quantize_term = llm_weight_block_quantize_term(weight_bits[0], block_size[0], has_input_scale); + fprintf(stderr, "quantize_multiheadattention table bits=%d block_size=%d term=%d %s\n", weight_bits[0], block_size[0], quantize_term, multiheadattention_name(mha)); + + const ncnn::Mat q_weight_data = mha->q_weight_data.reshape(qdim, mha->embed_dim); + const ncnn::Mat k_weight_data = mha->k_weight_data.reshape(mha->kdim, mha->embed_dim); + const ncnn::Mat v_weight_data = mha->v_weight_data.reshape(mha->vdim, mha->embed_dim); + const ncnn::Mat out_weight_data = mha->out_weight_data.reshape(mha->embed_dim, qdim); + + const ncnn::Mat weights[4] = {q_weight_data, k_weight_data, v_weight_data, out_weight_data}; + for (int j = 0; j < 4; j++) + { + if (qweight_rows[j]) + continue; + + int ret = 0; + if (has_input_scale) + { + ncnn::Mat weight_data_scaled; + ret = make_input_scaled_weight(weights[j], input_scales[j], weight_data_scaled); + if (ret != 0) + return ret; + + ret = pack_gemm_B_from_scales(weight_data_scaled, weight_data_quantize_scales[j], block_size[0], weight_bits[0], weight_data_quantized[j]); + } + else + { + ret = pack_gemm_B_from_scales(weights[j], weight_data_quantize_scales[j], block_size[0], weight_bits[0], weight_data_quantized[j]); + } + if (ret != 0) + return ret; + } + + mha->q_weight_data = weight_data_quantized[0]; + mha->k_weight_data = weight_data_quantized[1]; + mha->v_weight_data = weight_data_quantized[2]; + mha->out_weight_data = weight_data_quantized[3]; + mha->q_weight_data_quantize_scales = weight_data_quantize_scales[0]; + mha->k_weight_data_quantize_scales = weight_data_quantize_scales[1]; + mha->v_weight_data_quantize_scales = weight_data_quantize_scales[2]; + mha->out_weight_data_quantize_scales = weight_data_quantize_scales[3]; + mha->q_weight_data_input_scales = input_scales[0]; + mha->k_weight_data_input_scales = input_scales[1]; + mha->v_weight_data_input_scales = input_scales[2]; + mha->out_weight_data_input_scales = input_scales[3]; + mha->quantize_term = quantize_term; + + quantized_count++; + } + + fprintf(stderr, "quantized %d MultiHeadAttention layers\n", quantized_count); + + return 0; +} + +int main(int argc, char** argv) +{ + if (argc < 5) + { + print_usage(argv[0]); + return -1; + } + + for (int i = 1; i < argc; i++) + { + if (argv[i][0] == '-') + { + print_usage(argv[0]); + return -1; + } + } + + const char* inparam = argv[1]; + const char* inbin = argv[2]; + const char* outparam = argv[3]; + const char* outbin = argv[4]; + const char* tablepath = 0; + + int kv_start = 5; + if (argc >= 6 && strchr(argv[5], '=') == NULL) + { + tablepath = argv[5]; + kv_start = 6; + } + + int weight_bits = 6; + int block_size = 64; + const char* method = "minmax"; + + for (int i = kv_start; i < argc; i++) + { + // key=value + char* kv = argv[i]; + char* eqs = strchr(kv, '='); + if (eqs == NULL) + { + fprintf(stderr, "unrecognized arg %s\n", kv); + return -1; + } + + eqs[0] = '\0'; + const char* key = kv; + const char* value = eqs + 1; + + if (strcmp(key, "method") == 0) + method = value; + else if (strcmp(key, "bits") == 0) + weight_bits = atoi(value); + else if (strcmp(key, "block") == 0) + block_size = atoi(value); + else + { + fprintf(stderr, "unrecognized arg %s\n", key); + return -1; + } + } + + int quantize_method = LLM_QUANT_METHOD_MINMAX; + if (!tablepath) + { + quantize_method = llm_quant_method_from_string(method); + if (quantize_method < 0) + { + fprintf(stderr, "unsupported method=%s\n", method); + return -1; + } + + if (quantize_method == LLM_QUANT_METHOD_AWQ || quantize_method == LLM_QUANT_METHOD_GPTQ) + { + fprintf(stderr, "method=%s requires llm.table\n", method); + return -1; + } + + if (llm_weight_block_quantize_term(weight_bits, block_size) == 0) + { + print_usage(argv[0]); + return -1; + } + } + + std::map llm_scale_table; + if (tablepath) + { + if (!read_llm_scale_table(tablepath, llm_scale_table)) + return -1; + } + + NetQuantize quantizer; + quantizer.storage_type = 1; // keep existing prototype behavior for unrelated fp32 weights + + if (quantizer.load_param(inparam) != 0) + return -1; + + if (strcmp(inbin, "null") == 0) + { + DataReaderFromEmpty dr; + if (quantizer.load_model(dr) != 0) + return -1; + } + else + { + if (quantizer.load_model(inbin) != 0) + return -1; + } + + if (tablepath) + { + if (quantizer.quantize_gemm_from_table(llm_scale_table) != 0) + return -1; + if (quantizer.quantize_multiheadattention_from_table(llm_scale_table) != 0) + return -1; + print_unused_llm_table_rows(llm_scale_table); + } + else if (quantizer.quantize_gemm(block_size, weight_bits, quantize_method) != 0) + { + return -1; + } + else if (quantizer.quantize_multiheadattention(block_size, weight_bits, quantize_method) != 0) + { + return -1; + } + + quantizer.save(outparam, outbin); + + return 0; +} diff --git a/tools/quantize/ncnnllm2table.cpp b/tools/quantize/ncnnllm2table.cpp new file mode 100644 index 000000000000..ef3fdcdefe32 --- /dev/null +++ b/tools/quantize/ncnnllm2table.cpp @@ -0,0 +1,1813 @@ +// Copyright 2026 Tencent +// SPDX-License-Identifier: BSD-3-Clause + +#ifdef _MSC_VER +#define _CRT_SECURE_NO_DEPRECATE +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// npy format header +#include "npy.hpp" + +// ncnn public header +#include "datareader.h" +#include "layer.h" +#include "layer_type.h" +#include "net.h" + +// ncnn private header +#include "../modelwriter.h" +#include "ncnnllm_quant.h" + +class DataReaderFromEmpty : public ncnn::DataReader +{ +public: + virtual int scan(const char* /*format*/, void* /*p*/) const + { + return 0; + } + virtual size_t read(void* buf, size_t size) const + { + memset(buf, 0, size); + return size; + } +}; + +class QuantActStat +{ +public: + QuantActStat() + { + width = 0; + count = 0; + max_sample_count = 0; + } + +public: + int width; + int count; + int max_sample_count; + ncnn::Mat sum_abs; + ncnn::Mat max_abs; + std::vector samples; +}; + +class QuantNet : public ModelWriter +{ +public: + QuantNet(); + +public: + std::vector > listspaths; + std::vector > shapes; + std::vector input_blobs; + std::vector gemm_act_stats; + std::vector mha_act_stats; + int quantize_num_threads; + int file_type; + int awq_steps; + int awq_samples; + int awq_inner_method; + float awq_max_scale; + int gptq_samples; + float gptq_damp; + bool use_calibration_dataset; + +public: + int init(); + int collect_activation_stats(); + int save_table(const char* tablepath, int block_size, int weight_bits, int method) const; +}; + +QuantNet::QuantNet() + : ModelWriter() +{ + quantize_num_threads = 1; + file_type = 1; + awq_steps = 20; + awq_samples = 128; + awq_inner_method = LLM_QUANT_METHOD_MINMAX; + awq_max_scale = 16.f; + gptq_samples = 128; + gptq_damp = 0.01f; + use_calibration_dataset = false; +} + +static void show_usage(const char* argv0) +{ + fprintf(stderr, "Usage: %s [inparam] [inbin] [outtable] [(key=value)...]\n", argv0); + fprintf(stderr, " %s [inparam] [inbin] [caliblist] [outtable] [(key=value)...]\n", argv0); + fprintf(stderr, " method=minmax/mseclip/awq/gptq\n"); + fprintf(stderr, " bits=4/6/8\n"); + fprintf(stderr, " block=32/64/128\n"); + fprintf(stderr, " thread=8\n"); + fprintf(stderr, " type=1\n"); + fprintf(stderr, " shape=[w,h,...]\n"); + fprintf(stderr, " awq_steps=20\n"); + fprintf(stderr, " awq_samples=128\n"); + fprintf(stderr, " awq_max_scale=16\n"); + fprintf(stderr, " awq_inner=minmax/mseclip\n"); + fprintf(stderr, " gptq_samples=128\n"); + fprintf(stderr, " gptq_damp=0.01\n"); + fprintf(stderr, "Sample usage:\n"); + fprintf(stderr, " %s model.param model.bin model.llm.table method=mseclip bits=4 block=64\n", argv0); + fprintf(stderr, " %s model.param model.bin calib.list model.llm.table method=awq bits=4 block=64 shape=[4096]\n", argv0); + fprintf(stderr, " %s model.param model.bin calib.list model.llm.table method=gptq bits=4 block=64 shape=[4096]\n", argv0); +} + +static ncnn::Mat read_npy(const std::vector& shape, const std::string& npypath) +{ + npy::npy_data d; + try + { + d = npy::read_npy(npypath); + } + catch (const std::exception& e) + { + fprintf(stderr, "npy::read_npy %s exception: %s\n", npypath.c_str(), e.what()); + return ncnn::Mat(); + } + + const std::vector& npy_shape = d.shape; + const size_t dims = shape.size(); + + if (dims != npy_shape.size()) + { + fprintf(stderr, "%s expect %d dims, but got %d\n", npypath.c_str(), (int)dims, (int)npy_shape.size()); + return ncnn::Mat(); + } + + for (size_t i = 0; i < dims; i++) + { + if ((unsigned long)shape[i] != npy_shape[dims - 1 - i]) + { + fprintf(stderr, "%s shape mismatch\n", npypath.c_str()); + return ncnn::Mat(); + } + } + + if (dims == 1) + return ncnn::Mat(shape[0], (void*)d.data.data()).reshape(shape[0]).clone(); + if (dims == 2) + return ncnn::Mat(shape[0] * shape[1], (void*)d.data.data()).reshape(shape[0], shape[1]).clone(); + if (dims == 3) + return ncnn::Mat(shape[0] * shape[1] * shape[2], (void*)d.data.data()).reshape(shape[0], shape[1], shape[2]).clone(); + if (dims == 4) + return ncnn::Mat(shape[0] * shape[1] * shape[2] * shape[3], (void*)d.data.data()).reshape(shape[0], shape[1], shape[2], shape[3]).clone(); + + fprintf(stderr, "%s dims %d is unsupported\n", npypath.c_str(), (int)dims); + return ncnn::Mat(); +} + +static std::vector > parse_comma_path_list(char* s) +{ + std::vector > aps; + + char* pch = strtok(s, ","); + while (pch != NULL) + { + FILE* fp = fopen(pch, "rb"); + if (!fp) + { + fprintf(stderr, "fopen %s failed\n", pch); + break; + } + + std::vector paths; + + char line[1024]; + while (!feof(fp)) + { + char* ss = fgets(line, 1024, fp); + if (!ss) + break; + + char filepath[256]; + int nscan = sscanf(line, "%255s", filepath); + if (nscan != 1) + continue; + + paths.push_back(std::string(filepath)); + } + + fclose(fp); + + aps.push_back(paths); + + pch = strtok(NULL, ","); + } + + return aps; +} + +static std::vector > parse_comma_int_array_list(char* s) +{ + std::vector > aai; + + char* pch = strtok(s, "[]"); + while (pch != NULL) + { + int v; + int nconsumed = 0; + int nscan = sscanf(pch, "%d%n", &v, &nconsumed); + if (nscan == 1) + { + pch += nconsumed; + + std::vector ai; + ai.push_back(v); + + nscan = sscanf(pch, ",%d%n", &v, &nconsumed); + while (nscan == 1) + { + pch += nconsumed; + + ai.push_back(v); + + nscan = sscanf(pch, ",%d%n", &v, &nconsumed); + } + + aai.push_back(ai); + } + + pch = strtok(NULL, "[]"); + } + + return aai; +} + +static std::string make_qweight_filename(const char* tablepath, const char* key, std::string& qweight_path) +{ + const char* slash = strrchr(tablepath, '/'); +#ifdef _WIN32 + const char* backslash = strrchr(tablepath, '\\'); + if (backslash && (!slash || backslash > slash)) + slash = backslash; +#endif + + const char* table_base = slash ? slash + 1 : tablepath; + std::string qweight_name = std::string(table_base) + "." + key + ".qweight"; + for (size_t i = 0; i < qweight_name.size(); i++) + { + if (qweight_name[i] == '/' || qweight_name[i] == '\\' || qweight_name[i] == ':' || qweight_name[i] == ' ') + qweight_name[i] = '_'; + } + + qweight_path = qweight_name; + if (slash) + qweight_path = std::string(tablepath, slash - tablepath + 1) + qweight_name; + + return qweight_name; +} + +static int write_raw_mat_file(const char* path, const ncnn::Mat& m) +{ + FILE* fp = fopen(path, "wb"); + if (!fp) + { + fprintf(stderr, "fopen %s failed\n", path); + return -1; + } + + const size_t logical_count = (size_t)m.w * m.h * m.d * m.c; + if (fwrite(m.data, m.elemsize, logical_count, fp) != logical_count) + { + fprintf(stderr, "fwrite %s failed\n", path); + fclose(fp); + return -1; + } + + fclose(fp); + + return 0; +} + +static int init_act_stat(QuantActStat& stat, int width, int max_sample_count) +{ + stat.width = width; + stat.count = 0; + stat.max_sample_count = max_sample_count; + + stat.sum_abs.create(width); + stat.max_abs.create(width); + if (stat.sum_abs.empty() || stat.max_abs.empty()) + return -100; + + stat.sum_abs.fill(0.f); + stat.max_abs.fill(0.f); + stat.samples.clear(); + stat.samples.reserve((size_t)width * max_sample_count); + + return 0; +} + +static int collect_act_row(QuantActStat& stat, const float* ptr) +{ + const int width = stat.width; + float* sum_abs_ptr = stat.sum_abs; + float* max_abs_ptr = stat.max_abs; + + for (int k = 0; k < width; k++) + { + const float v = (float)fabs(ptr[k]); + sum_abs_ptr[k] += v; + max_abs_ptr[k] = std::max(max_abs_ptr[k], v); + } + + if ((int)(stat.samples.size() / width) < stat.max_sample_count) + { + for (int k = 0; k < width; k++) + stat.samples.push_back(ptr[k]); + } + + stat.count++; + + return 0; +} + +static int collect_act_rows(QuantActStat& stat, const ncnn::Mat& m, int width, const char* key) +{ + if (m.empty()) + { + fprintf(stderr, "%s activation blob is empty\n", key); + return -1; + } + + if (m.elemsize != 4u) + { + fprintf(stderr, "%s activation elemsize %d is unsupported\n", key, (int)m.elemsize); + return -1; + } + + if (stat.width != width) + { + fprintf(stderr, "%s activation width mismatch expected=%d got=%d\n", key, stat.width, width); + return -1; + } + + if (m.dims == 2 && m.w == width) + { + for (int i = 0; i < m.h; i++) + collect_act_row(stat, m.row(i)); + return 0; + } + + if (m.dims == 1 && m.w == width) + { + collect_act_row(stat, (const float*)m); + return 0; + } + + const int total = (int)m.total(); + if (total % width != 0) + { + fprintf(stderr, "%s activation size %d is not divisible by width %d\n", key, total, width); + return -1; + } + + const float* ptr = m; + const int sample_count = total / width; + for (int i = 0; i < sample_count; i++) + collect_act_row(stat, ptr + i * width); + + return 0; +} + +static int resolve_mha_bottom_blob_index(const ncnn::MultiHeadAttention* mha, int bottom_blob_count, int& q_blob_i, int& k_blob_i, int& v_blob_i, int& attn_mask_i, int& cached_xk_i, int& cached_xv_i) +{ + q_blob_i = 0; + k_blob_i = 0; + v_blob_i = 0; + attn_mask_i = 0; + cached_xk_i = 0; + cached_xv_i = 0; + + if (mha->kv_cache) + { + if (mha->attn_mask) + { + if (bottom_blob_count == 4) + { + q_blob_i = 0; + k_blob_i = 0; + v_blob_i = 0; + attn_mask_i = 1; + cached_xk_i = 2; + cached_xv_i = 3; + return 0; + } + if (bottom_blob_count == 5) + { + q_blob_i = 0; + k_blob_i = 1; + v_blob_i = 1; + attn_mask_i = 2; + cached_xk_i = 3; + cached_xv_i = 4; + return 0; + } + if (bottom_blob_count == 6) + { + q_blob_i = 0; + k_blob_i = 1; + v_blob_i = 2; + attn_mask_i = 3; + cached_xk_i = 4; + cached_xv_i = 5; + return 0; + } + } + else + { + if (bottom_blob_count == 3) + { + q_blob_i = 0; + k_blob_i = 0; + v_blob_i = 0; + cached_xk_i = 1; + cached_xv_i = 2; + return 0; + } + if (bottom_blob_count == 4) + { + q_blob_i = 0; + k_blob_i = 1; + v_blob_i = 1; + cached_xk_i = 2; + cached_xv_i = 3; + return 0; + } + if (bottom_blob_count == 5) + { + q_blob_i = 0; + k_blob_i = 1; + v_blob_i = 2; + cached_xk_i = 3; + cached_xv_i = 4; + return 0; + } + } + } + else + { + if (mha->attn_mask) + { + if (bottom_blob_count == 2) + { + q_blob_i = 0; + k_blob_i = 0; + v_blob_i = 0; + attn_mask_i = 1; + return 0; + } + if (bottom_blob_count == 3) + { + q_blob_i = 0; + k_blob_i = 1; + v_blob_i = 1; + attn_mask_i = 2; + return 0; + } + if (bottom_blob_count == 4) + { + q_blob_i = 0; + k_blob_i = 1; + v_blob_i = 2; + attn_mask_i = 3; + return 0; + } + } + else + { + if (bottom_blob_count == 1) + { + q_blob_i = 0; + k_blob_i = 0; + v_blob_i = 0; + return 0; + } + if (bottom_blob_count == 2) + { + q_blob_i = 0; + k_blob_i = 1; + v_blob_i = 1; + return 0; + } + if (bottom_blob_count == 3) + { + q_blob_i = 0; + k_blob_i = 1; + v_blob_i = 2; + return 0; + } + } + } + + fprintf(stderr, "MultiHeadAttention %s unsupported bottom blob count %d\n", multiheadattention_name(mha), bottom_blob_count); + return -1; +} + +static int collect_mha_out_act_rows(const ncnn::MultiHeadAttention* mha, const ncnn::Mat& q_blob, const ncnn::Mat& k_blob, const ncnn::Mat& v_blob, const ncnn::Mat& attn_mask_blob, const ncnn::Mat& cached_xk_blob, const ncnn::Mat& cached_xv_blob, int q_blob_i, int k_blob_i, int v_blob_i, QuantActStat& stat, const char* key, int num_threads) +{ + const int src_seqlen = q_blob.h; + const int cur_seqlen = k_blob.h; + const int past_seqlen = mha->kv_cache && !cached_xk_blob.empty() ? cached_xk_blob.w : 0; + const int dst_seqlen = past_seqlen > 0 ? (q_blob_i == k_blob_i ? (past_seqlen + cur_seqlen) : past_seqlen) : cur_seqlen; + + const int embed_dim = mha->embed_dim; + const int num_heads = mha->num_heads; + const int embed_dim_per_head = embed_dim / num_heads; + const int qdim = mha->weight_data_size / embed_dim; + + ncnn::Mat q_affine; + q_affine.create(src_seqlen, embed_dim); + if (q_affine.empty()) + return -100; + + #pragma omp parallel for num_threads(num_threads) + for (int i = 0; i < src_seqlen; i++) + { + const float* kptr = (const float*)mha->q_weight_data; + + for (int j = 0; j < embed_dim; j++) + { + const float* ptr = q_blob.row(i); + + float sum = mha->q_bias_data[j]; + for (int k = 0; k < qdim; k++) + sum += *ptr++ * *kptr++; + + q_affine.row(j)[i] = sum * mha->scale; + } + } + + ncnn::Mat k_affine; + if (past_seqlen > 0 && q_blob_i != k_blob_i) + { + k_affine = cached_xk_blob; + } + else + { + k_affine.create(dst_seqlen, embed_dim); + if (k_affine.empty()) + return -100; + + if (past_seqlen > 0) + { + #pragma omp parallel for num_threads(num_threads) + for (int i = 0; i < embed_dim; i++) + memcpy(k_affine.row(i), cached_xk_blob.row(i), past_seqlen * sizeof(float)); + } + + #pragma omp parallel for num_threads(num_threads) + for (int i = 0; i < cur_seqlen; i++) + { + const float* kptr = (const float*)mha->k_weight_data; + + for (int j = 0; j < embed_dim; j++) + { + const float* ptr = k_blob.row(i); + + float sum = mha->k_bias_data[j]; + for (int k = 0; k < mha->kdim; k++) + sum += *ptr++ * *kptr++; + + k_affine.row(j)[past_seqlen + i] = sum; + } + } + } + + ncnn::Mat v_affine; + if (past_seqlen > 0 && q_blob_i != v_blob_i) + { + v_affine = cached_xv_blob; + } + else + { + v_affine.create(dst_seqlen, embed_dim); + if (v_affine.empty()) + return -100; + + if (past_seqlen > 0) + { + #pragma omp parallel for num_threads(num_threads) + for (int i = 0; i < embed_dim; i++) + memcpy(v_affine.row(i), cached_xv_blob.row(i), past_seqlen * sizeof(float)); + } + + #pragma omp parallel for num_threads(num_threads) + for (int i = 0; i < cur_seqlen; i++) + { + const float* kptr = (const float*)mha->v_weight_data; + + for (int j = 0; j < embed_dim; j++) + { + const float* ptr = v_blob.row(i); + + float sum = mha->v_bias_data[j]; + for (int k = 0; k < mha->vdim; k++) + sum += *ptr++ * *kptr++; + + v_affine.row(j)[past_seqlen + i] = sum; + } + } + } + + ncnn::Mat qk_cross; + qk_cross.create(dst_seqlen, src_seqlen, num_heads); + if (qk_cross.empty()) + return -100; + + #pragma omp parallel for num_threads(num_threads) + for (int q = 0; q < num_heads; q++) + { + const ncnn::Mat q_affine_head = q_affine.row_range(q * embed_dim_per_head, embed_dim_per_head); + const ncnn::Mat k_affine_head = k_affine.row_range(q * embed_dim_per_head, embed_dim_per_head); + ncnn::Mat qk_cross_head = qk_cross.channel(q); + + for (int i = 0; i < src_seqlen; i++) + { + float* outptr = qk_cross_head.row(i); + + for (int j = 0; j < dst_seqlen; j++) + { + float sum = 0.f; + for (int l = 0; l < embed_dim_per_head; l++) + sum += q_affine_head.row(l)[i] * k_affine_head.row(l)[j]; + + outptr[j] = sum; + } + } + } + + if (mha->attn_mask) + { + #pragma omp parallel for num_threads(num_threads) + for (int q = 0; q < num_heads; q++) + { + const ncnn::Mat& maskm = attn_mask_blob.dims == 3 ? attn_mask_blob.channel(q) : attn_mask_blob; + ncnn::Mat qk_cross_head = qk_cross.channel(q); + + for (int i = 0; i < src_seqlen; i++) + { + const float* mptr = maskm.row(i); + float* outptr = qk_cross_head.row(i); + + for (int j = 0; j < dst_seqlen; j++) + outptr[j] += mptr[j]; + } + } + } + + #pragma omp parallel for num_threads(num_threads) + for (int q = 0; q < num_heads; q++) + { + ncnn::Mat qk_cross_head = qk_cross.channel(q); + + for (int i = 0; i < src_seqlen; i++) + { + float* ptr = qk_cross_head.row(i); + + float max = -FLT_MAX; + for (int j = 0; j < dst_seqlen; j++) + max = std::max(max, ptr[j]); + + float sum = 0.f; + for (int j = 0; j < dst_seqlen; j++) + { + ptr[j] = (float)expf(ptr[j] - max); + sum += ptr[j]; + } + + for (int j = 0; j < dst_seqlen; j++) + ptr[j] /= sum; + } + } + + ncnn::Mat qkv_cross; + qkv_cross.create(src_seqlen, embed_dim); + if (qkv_cross.empty()) + return -100; + + #pragma omp parallel for num_threads(num_threads) + for (int q = 0; q < num_heads; q++) + { + const ncnn::Mat qk_cross_head = qk_cross.channel(q); + const ncnn::Mat v_affine_head = v_affine.row_range(q * embed_dim_per_head, embed_dim_per_head); + ncnn::Mat qkv_cross_head = qkv_cross.row_range(q * embed_dim_per_head, embed_dim_per_head); + + for (int i = 0; i < src_seqlen; i++) + { + for (int j = 0; j < embed_dim_per_head; j++) + { + const float* qkptr = qk_cross_head.row(i); + const float* vptr = v_affine_head.row(j); + + float sum = 0.f; + for (int k = 0; k < dst_seqlen; k++) + sum += *qkptr++ * *vptr++; + + qkv_cross_head.row(j)[i] = sum; + } + } + } + + if (embed_dim != stat.width) + { + fprintf(stderr, "%s out activation width mismatch expected=%d got=%d\n", key, stat.width, embed_dim); + return -1; + } + + std::vector row(embed_dim); + for (int i = 0; i < src_seqlen; i++) + { + for (int k = 0; k < embed_dim; k++) + row[k] = qkv_cross.row(k)[i]; + + collect_act_row(stat, row.data()); + } + + return 0; +} + +static int make_input_scaled_weight(const ncnn::Mat& weight_data, const ncnn::Mat& input_scales, ncnn::Mat& weight_data_scaled) +{ + const int K = weight_data.w; + const int N = weight_data.h; + + if (input_scales.w != K) + { + fprintf(stderr, "input scale count mismatch expected=%d got=%d\n", K, input_scales.w); + return -1; + } + + weight_data_scaled.create(K, N); + if (weight_data_scaled.empty()) + return -100; + + const float* input_scale_ptr = input_scales; + for (int n = 0; n < N; n++) + { + const float* ptr = weight_data.row(n); + float* outptr = weight_data_scaled.row(n); + + for (int k = 0; k < K; k++) + { + const float s = input_scale_ptr[k]; + if (!(s > 0.f) || s > FLT_MAX) + { + fprintf(stderr, "invalid input scale k=%d scale=%f\n", k, s); + return -1; + } + + outptr[k] = ptr[k] / s; + } + } + + return 0; +} + +static int make_dequant_weight(const ncnn::Mat& weight_data, const ncnn::Mat& weight_data_quantize_scales, int block_size, int weight_bits, ncnn::Mat& weight_data_dequantized) +{ + const int K = weight_data.w; + const int N = weight_data.h; + const int block_count = (K + block_size - 1) / block_size; + + if (weight_data_quantize_scales.w != block_count || weight_data_quantize_scales.h != N) + { + fprintf(stderr, "weight scale shape mismatch expected=%d,%d got=%d,%d\n", block_count, N, weight_data_quantize_scales.w, weight_data_quantize_scales.h); + return -1; + } + + weight_data_dequantized.create(K, N); + if (weight_data_dequantized.empty()) + return -100; + + for (int n = 0; n < N; n++) + { + const float* ptr = weight_data.row(n); + const float* scale_ptr = weight_data_quantize_scales.row(n); + float* outptr = weight_data_dequantized.row(n); + + for (int b = 0; b < block_count; b++) + { + const int k0 = b * block_size; + const int max_kk = std::min(block_size, K - k0); + const float scale = scale_ptr[b]; + + for (int k = 0; k < max_kk; k++) + { + const int q = float2int_weight(ptr[k0 + k] * scale, weight_bits); + outptr[k0 + k] = q / scale; + } + } + } + + return 0; +} + +static double calc_awq_reconstruction_error(const ncnn::Mat& weight_data, const ncnn::Mat& input_scales, const ncnn::Mat& weight_data_dequantized, const QuantActStat& stat) +{ + const int K = weight_data.w; + const int N = weight_data.h; + const int sample_count = (int)(stat.samples.size() / K); + + const float* input_scale_ptr = input_scales; + double error = 0.0; + + for (int i = 0; i < sample_count; i++) + { + const float* sample_ptr = &stat.samples[(size_t)i * K]; + + for (int n = 0; n < N; n++) + { + const float* wptr = weight_data.row(n); + const float* deqptr = weight_data_dequantized.row(n); + + double ref = 0.0; + double out = 0.0; + for (int k = 0; k < K; k++) + { + const float x = sample_ptr[k]; + ref += (double)x * wptr[k]; + out += (double)(x * input_scale_ptr[k]) * deqptr[k]; + } + + const double diff = ref - out; + error += diff * diff; + } + } + + return error; +} + +static int make_awq_input_scales(const ncnn::Mat& weight_data, const QuantActStat& stat, int block_size, int weight_bits, int inner_method, int awq_steps, float awq_max_scale, ncnn::Mat& input_scales, ncnn::Mat& weight_data_quantize_scales, int num_threads) +{ + const int K = weight_data.w; + const int N = weight_data.h; + + if (stat.width != K || stat.count == 0 || stat.samples.empty()) + { + fprintf(stderr, "awq activation stat is empty or shape mismatch K=%d stat_width=%d count=%d samples=%d\n", K, stat.width, stat.count, (int)(stat.samples.size() / std::max(K, 1))); + return -1; + } + + ncnn::Mat act_mean(K); + ncnn::Mat weight_mean(K); + if (act_mean.empty() || weight_mean.empty()) + return -100; + + const float* sum_abs_ptr = stat.sum_abs; + float* act_mean_ptr = act_mean; + for (int k = 0; k < K; k++) + act_mean_ptr[k] = sum_abs_ptr[k] / stat.count; + + weight_mean.fill(0.f); + float* weight_mean_ptr = weight_mean; + for (int n = 0; n < N; n++) + { + const float* ptr = weight_data.row(n); + for (int k = 0; k < K; k++) + weight_mean_ptr[k] += (float)fabs(ptr[k]); + } + for (int k = 0; k < K; k++) + weight_mean_ptr[k] /= N; + + ncnn::Mat best_input_scales; + ncnn::Mat best_quantize_scales; + double best_error = DBL_MAX; + + const int search_steps = std::max(awq_steps, 0); + for (int s = 0; s <= search_steps; s++) + { + ncnn::Mat input_scales1(K); + if (input_scales1.empty()) + return -100; + + float* input_scale_ptr = input_scales1; + if (s == 0 || search_steps == 0) + { + for (int k = 0; k < K; k++) + input_scale_ptr[k] = 1.f; + } + else + { + const float alpha = (float)s / search_steps; + double logsum = 0.0; + for (int k = 0; k < K; k++) + { + float raw = powf((weight_mean_ptr[k] + 1e-6f) / (act_mean_ptr[k] + 1e-6f), alpha); + if (raw < 1e-12f) + raw = 1e-12f; + input_scale_ptr[k] = raw; + logsum += ::log((double)raw); + } + + const float geomean = (float)::exp(logsum / K); + for (int k = 0; k < K; k++) + { + float v = input_scale_ptr[k] / geomean; + if (v < 1.f / awq_max_scale) + v = 1.f / awq_max_scale; + if (v > awq_max_scale) + v = awq_max_scale; + + input_scale_ptr[k] = v; + } + } + + ncnn::Mat weight_data_scaled; + int ret = make_input_scaled_weight(weight_data, input_scales1, weight_data_scaled); + if (ret != 0) + return ret; + + ncnn::Mat quantize_scales1; + ret = make_gemm_B_scales(weight_data_scaled, block_size, weight_bits, inner_method, quantize_scales1, num_threads); + if (ret != 0) + return ret; + + ncnn::Mat weight_data_dequantized; + ret = make_dequant_weight(weight_data_scaled, quantize_scales1, block_size, weight_bits, weight_data_dequantized); + if (ret != 0) + return ret; + + const double error = calc_awq_reconstruction_error(weight_data, input_scales1, weight_data_dequantized, stat); + if (error < best_error) + { + best_error = error; + best_input_scales = input_scales1; + best_quantize_scales = quantize_scales1; + } + } + + input_scales = best_input_scales.clone(); + weight_data_quantize_scales = best_quantize_scales.clone(); + if (input_scales.empty() || weight_data_quantize_scales.empty()) + return -100; + + return 0; +} + +static int invert_gptq_hessian(std::vector& m, int n) +{ + std::vector inv((size_t)n * n, 0.0); + for (int i = 0; i < n; i++) + inv[(size_t)i * n + i] = 1.0; + + for (int i = 0; i < n; i++) + { + int pivot = i; + double pivot_abs = fabs(m[(size_t)i * n + i]); + for (int j = i + 1; j < n; j++) + { + const double v = fabs(m[(size_t)j * n + i]); + if (v > pivot_abs) + { + pivot_abs = v; + pivot = j; + } + } + + if (pivot_abs < 1e-12) + return -1; + + if (pivot != i) + { + for (int j = 0; j < n; j++) + { + std::swap(m[(size_t)i * n + j], m[(size_t)pivot * n + j]); + std::swap(inv[(size_t)i * n + j], inv[(size_t)pivot * n + j]); + } + } + + const double diag = m[(size_t)i * n + i]; + for (int j = 0; j < n; j++) + { + m[(size_t)i * n + j] /= diag; + inv[(size_t)i * n + j] /= diag; + } + + for (int r = 0; r < n; r++) + { + if (r == i) + continue; + + const double f = m[(size_t)r * n + i]; + if (f == 0.0) + continue; + + for (int c = 0; c < n; c++) + { + m[(size_t)r * n + c] -= f * m[(size_t)i * n + c]; + inv[(size_t)r * n + c] -= f * inv[(size_t)i * n + c]; + } + } + } + + m.swap(inv); + + return 0; +} + +static int make_gptq_qweight(const ncnn::Mat& weight_data, const QuantActStat& stat, int block_size, int weight_bits, float damp, ncnn::Mat& weight_data_quantize_scales, ncnn::Mat& weight_data_quantized) +{ + const int K = weight_data.w; + const int N = weight_data.h; + const int block_count = (K + block_size - 1) / block_size; + const int packed_k_bytes = llm_weight_quantize_packed_k_bytes(K, weight_bits); + if (packed_k_bytes <= 0) + return -1; + const int sample_count = (int)(stat.samples.size() / K); + + if (stat.width != K || sample_count == 0) + { + fprintf(stderr, "gptq activation stat is empty or shape mismatch K=%d stat_width=%d samples=%d\n", K, stat.width, sample_count); + return -1; + } + + weight_data_quantize_scales.create(block_count, N); + weight_data_quantized.create(packed_k_bytes, N, (size_t)1u); + if (weight_data_quantize_scales.empty() || weight_data_quantized.empty()) + return -100; + + memset(weight_data_quantized.data, 0, weight_data_quantized.total() * weight_data_quantized.elemsize); + + for (int b = 0; b < block_count; b++) + { + const int k0 = b * block_size; + const int B = std::min(block_size, K - k0); + + std::vector H((size_t)B * B, 0.0); + for (int s = 0; s < sample_count; s++) + { + const float* x = &stat.samples[(size_t)s * K + k0]; + for (int i = 0; i < B; i++) + { + for (int j = 0; j < B; j++) + H[(size_t)i * B + j] += (double)x[i] * x[j]; + } + } + + double diag_mean = 0.0; + for (int i = 0; i < B; i++) + diag_mean += H[(size_t)i * B + i]; + diag_mean /= B; + if (!(diag_mean > 0.0) || diag_mean > DBL_MAX) + diag_mean = 1.0; + + for (int i = 0; i < B; i++) + H[(size_t)i * B + i] += damp * diag_mean; + + if (invert_gptq_hessian(H, B) != 0) + { + fprintf(stderr, "gptq hessian inverse failed block=%d, use identity hessian inverse\n", b); + for (int i = 0; i < B * B; i++) + H[i] = 0.0; + for (int i = 0; i < B; i++) + H[(size_t)i * B + i] = 1.0; + } + + std::vector work((size_t)N * B); + for (int n = 0; n < N; n++) + { + const float* ptr = weight_data.row(n) + k0; + for (int k = 0; k < B; k++) + work[(size_t)n * B + k] = ptr[k]; + } + + for (int n = 0; n < N; n++) + { + const float* ptr = weight_data.row(n) + k0; + float* scale_ptr = weight_data_quantize_scales.row(n); + // fixed-scale gptq keeps the runtime block scale tied to the original weight block + scale_ptr[b] = choose_weight_scale(ptr, B, weight_bits, LLM_QUANT_METHOD_MINMAX); + } + + for (int k = 0; k < B; k++) + { + double hdiag = H[(size_t)k * B + k]; + if (fabs(hdiag) < 1e-12) + hdiag = hdiag < 0.0 ? -1e-12 : 1e-12; + + for (int n = 0; n < N; n++) + { + const float scale = weight_data_quantize_scales.row(n)[b]; + const double w = work[(size_t)n * B + k]; + const int q = float2int_weight((float)(w * scale), weight_bits); + const double deq = (double)q / scale; + const double err = (w - deq) / hdiag; + + unsigned char* qptr = weight_data_quantized.row(n); + pack_signed_weight(qptr, k0 + k, weight_bits, q); + + for (int j = k + 1; j < B; j++) + work[(size_t)n * B + j] -= err * H[(size_t)k * B + j]; + } + } + } + + return 0; +} + +int QuantNet::init() +{ + const int max_sample_count = std::max(awq_samples, gptq_samples); + + for (size_t i = 0; i < layers.size(); i++) + { + const ncnn::Layer* layer = layers[i]; + if (layer->type == "Input") + input_blobs.push_back(layer->tops[0]); + } + + for (size_t i = 0; i < layers.size(); i++) + { + if (layers[i]->type == "Gemm") + { + const ncnn::Gemm* gemm = (const ncnn::Gemm*)layers[i]; + + const char* reason = 0; + if (!is_supported_llm_gemm(gemm, &reason)) + continue; + + QuantActStat stat; + int ret = init_act_stat(stat, gemm->constantK, max_sample_count); + if (ret != 0) + return ret; + + gemm_act_stats.push_back(stat); + } + else if (layers[i]->type == "MultiHeadAttention") + { + const ncnn::MultiHeadAttention* mha = (const ncnn::MultiHeadAttention*)layers[i]; + + const char* reason = 0; + if (!is_supported_llm_multiheadattention(mha, &reason)) + continue; + + const size_t mha_act_offset = mha_act_stats.size(); + mha_act_stats.resize(mha_act_offset + 4); + + const int qdim = mha->weight_data_size / mha->embed_dim; + const int Ks[4] = {qdim, mha->kdim, mha->vdim, mha->embed_dim}; + + for (int j = 0; j < 4; j++) + { + int ret = init_act_stat(mha_act_stats[mha_act_offset + j], Ks[j], max_sample_count); + if (ret != 0) + return ret; + } + } + } + + return 0; +} + +static int check_calibration_input(const QuantNet& table) +{ + const int input_blob_count = (int)table.input_blobs.size(); + if (input_blob_count == 0) + { + fprintf(stderr, "no input blob found\n"); + return -1; + } + + if ((int)table.listspaths.size() != input_blob_count) + { + fprintf(stderr, "calibration list count mismatch expected=%d got=%d\n", input_blob_count, (int)table.listspaths.size()); + return -1; + } + + if ((int)table.shapes.size() != input_blob_count) + { + fprintf(stderr, "shape count mismatch expected=%d got=%d\n", input_blob_count, (int)table.shapes.size()); + return -1; + } + + if (table.file_type != 1) + { + fprintf(stderr, "ncnnllm2table calibration supports type=1 only\n"); + return -1; + } + + if (table.listspaths[0].empty()) + { + fprintf(stderr, "calibration list is empty\n"); + return -1; + } + + const int file_count = (int)table.listspaths[0].size(); + for (int i = 1; i < input_blob_count; i++) + { + if ((int)table.listspaths[i].size() != file_count) + { + fprintf(stderr, "calibration list size mismatch input=%d expected=%d got=%d\n", i, file_count, (int)table.listspaths[i].size()); + return -1; + } + } + + return 0; +} + +int QuantNet::collect_activation_stats() +{ + int ret = check_calibration_input(*this); + if (ret != 0) + return ret; + + const int input_blob_count = (int)input_blobs.size(); + const int file_count = (int)listspaths[0].size(); + + for (int i = 0; i < file_count; i++) + { + if (i % 100 == 0) + fprintf(stderr, "collect_activation_stats %.2f%% [ %d / %d ]\n", i * 100.f / file_count, i, file_count); + + ncnn::Extractor ex = create_extractor(); + ex.set_light_mode(false); + + for (int j = 0; j < input_blob_count; j++) + { + ncnn::Mat in = read_npy(shapes[j], listspaths[j][i]); + if (in.empty()) + return -1; + + ret = ex.input(input_blobs[j], in); + if (ret != 0) + { + fprintf(stderr, "input blob %d failed ret=%d\n", input_blobs[j], ret); + return ret; + } + } + + size_t gemm_act_index = 0; + size_t mha_act_index = 0; + + for (size_t j = 0; j < layers.size(); j++) + { + if (layers[j]->type == "Gemm") + { + const ncnn::Gemm* gemm = (const ncnn::Gemm*)layers[j]; + + const char* reason = 0; + if (!is_supported_llm_gemm(gemm, &reason)) + continue; + + char key[256]; + snprintf(key, 256, "%s_param_1", gemm_name(gemm)); + + if (gemm_act_index >= gemm_act_stats.size()) + { + fprintf(stderr, "%s activation stat is missing\n", key); + return -1; + } + + ncnn::Mat bottom_blob; + ret = ex.extract(gemm->bottoms[0], bottom_blob); + if (ret != 0) + { + fprintf(stderr, "extract %s bottom failed ret=%d\n", gemm_name(gemm), ret); + return ret; + } + + ret = collect_act_rows(gemm_act_stats[gemm_act_index], bottom_blob, gemm->constantK, key); + if (ret != 0) + return ret; + + gemm_act_index++; + } + else if (layers[j]->type == "MultiHeadAttention") + { + const ncnn::MultiHeadAttention* mha = (const ncnn::MultiHeadAttention*)layers[j]; + + const char* reason = 0; + if (!is_supported_llm_multiheadattention(mha, &reason)) + continue; + + int q_blob_i = 0; + int k_blob_i = 0; + int v_blob_i = 0; + int attn_mask_i = 0; + int cached_xk_i = 0; + int cached_xv_i = 0; + ret = resolve_mha_bottom_blob_index(mha, (int)mha->bottoms.size(), q_blob_i, k_blob_i, v_blob_i, attn_mask_i, cached_xk_i, cached_xv_i); + if (ret != 0) + return ret; + + ncnn::Mat q_blob; + ncnn::Mat k_blob; + ncnn::Mat v_blob; + ncnn::Mat attn_mask_blob; + ncnn::Mat cached_xk_blob; + ncnn::Mat cached_xv_blob; + + ret = ex.extract(mha->bottoms[q_blob_i], q_blob); + if (ret != 0) + { + fprintf(stderr, "extract %s q bottom failed ret=%d\n", multiheadattention_name(mha), ret); + return ret; + } + ret = ex.extract(mha->bottoms[k_blob_i], k_blob); + if (ret != 0) + { + fprintf(stderr, "extract %s k bottom failed ret=%d\n", multiheadattention_name(mha), ret); + return ret; + } + ret = ex.extract(mha->bottoms[v_blob_i], v_blob); + if (ret != 0) + { + fprintf(stderr, "extract %s v bottom failed ret=%d\n", multiheadattention_name(mha), ret); + return ret; + } + if (mha->attn_mask) + { + ret = ex.extract(mha->bottoms[attn_mask_i], attn_mask_blob); + if (ret != 0) + { + fprintf(stderr, "extract %s attn_mask bottom failed ret=%d\n", multiheadattention_name(mha), ret); + return ret; + } + } + if (mha->kv_cache) + { + ret = ex.extract(mha->bottoms[cached_xk_i], cached_xk_blob); + if (ret != 0) + { + fprintf(stderr, "extract %s cached_xk bottom failed ret=%d\n", multiheadattention_name(mha), ret); + return ret; + } + ret = ex.extract(mha->bottoms[cached_xv_i], cached_xv_blob); + if (ret != 0) + { + fprintf(stderr, "extract %s cached_xv bottom failed ret=%d\n", multiheadattention_name(mha), ret); + return ret; + } + } + + const int qdim = mha->weight_data_size / mha->embed_dim; + + char q_key[256]; + char k_key[256]; + char v_key[256]; + char out_key[256]; + snprintf(q_key, 256, "%s_param_0", multiheadattention_name(mha)); + snprintf(k_key, 256, "%s_param_1", multiheadattention_name(mha)); + snprintf(v_key, 256, "%s_param_2", multiheadattention_name(mha)); + snprintf(out_key, 256, "%s_param_3", multiheadattention_name(mha)); + + const size_t mha_act_offset = mha_act_index * 4; + if (mha_act_offset + 3 >= mha_act_stats.size()) + { + fprintf(stderr, "%s activation stat is missing\n", multiheadattention_name(mha)); + return -1; + } + + ret = collect_act_rows(mha_act_stats[mha_act_offset], q_blob, qdim, q_key); + if (ret != 0) + return ret; + ret = collect_act_rows(mha_act_stats[mha_act_offset + 1], k_blob, mha->kdim, k_key); + if (ret != 0) + return ret; + ret = collect_act_rows(mha_act_stats[mha_act_offset + 2], v_blob, mha->vdim, v_key); + if (ret != 0) + return ret; + + ret = collect_mha_out_act_rows(mha, q_blob, k_blob, v_blob, attn_mask_blob, cached_xk_blob, cached_xv_blob, q_blob_i, k_blob_i, v_blob_i, mha_act_stats[mha_act_offset + 3], out_key, quantize_num_threads); + if (ret != 0) + return ret; + + mha_act_index++; + } + } + } + + fprintf(stderr, "collect_activation_stats 100.00%% [ %d / %d ]\n", file_count, file_count); + + return 0; +} + +int QuantNet::save_table(const char* tablepath, int block_size, int weight_bits, int method) const +{ + FILE* fp = fopen(tablepath, "wb"); + if (!fp) + { + fprintf(stderr, "fopen %s failed\n", tablepath); + return -1; + } + + int table_count = 0; + size_t gemm_act_index = 0; + size_t mha_act_index = 0; + + for (size_t i = 0; i < layers.size(); i++) + { + if (layers[i]->type == "Gemm") + { + const ncnn::Gemm* gemm = (const ncnn::Gemm*)layers[i]; + + const char* reason = 0; + if (!is_supported_llm_gemm(gemm, &reason)) + { + print_skip_gemm(gemm, reason); + continue; + } + + char key[256]; + snprintf(key, 256, "%s_param_1", gemm_name(gemm)); + + const QuantActStat* act_stat = 0; + if (method == LLM_QUANT_METHOD_AWQ || method == LLM_QUANT_METHOD_GPTQ) + { + if (gemm_act_index >= gemm_act_stats.size()) + { + fprintf(stderr, "%s activation stat is missing\n", key); + fclose(fp); + return -1; + } + + act_stat = &gemm_act_stats[gemm_act_index]; + } + + ncnn::Mat B_data_quantize_scales; + ncnn::Mat B_data_input_scales; + ncnn::Mat B_data_quantized; + int ret = 0; + if (method == LLM_QUANT_METHOD_AWQ) + { + ret = make_awq_input_scales(gemm->B_data, *act_stat, block_size, weight_bits, awq_inner_method, awq_steps, awq_max_scale, B_data_input_scales, B_data_quantize_scales, quantize_num_threads); + } + else if (method == LLM_QUANT_METHOD_GPTQ) + { + ret = make_gptq_qweight(gemm->B_data, *act_stat, block_size, weight_bits, gptq_damp, B_data_quantize_scales, B_data_quantized); + } + else + { + ret = make_gemm_B_scales(gemm->B_data, block_size, weight_bits, method, B_data_quantize_scales, quantize_num_threads); + } + if (ret != 0) + { + fclose(fp); + return ret; + } + + if (method == LLM_QUANT_METHOD_GPTQ) + { + std::string qweight_path; + const std::string qweight_name = make_qweight_filename(tablepath, key, qweight_path); + ret = write_raw_mat_file(qweight_path.c_str(), B_data_quantized); + if (ret != 0) + { + fclose(fp); + return ret; + } + + if (write_llm_qweight_table_row(fp, key, weight_bits, block_size, method, qweight_name.c_str(), B_data_quantize_scales) != 0) + { + fclose(fp); + return -1; + } + } + else if (write_llm_table_row(fp, key, weight_bits, block_size, method, B_data_quantize_scales) != 0) + { + fclose(fp); + return -1; + } + + fprintf(stderr, "write_llm_table %s bits=%d block=%d method=%s\n", key, weight_bits, block_size, llm_quant_method_to_string(method)); + table_count++; + + if (method == LLM_QUANT_METHOD_AWQ) + { + char input_scale_key[512]; + snprintf(input_scale_key, 512, "%s_input_scale", key); + + if (write_llm_input_scale_row(fp, input_scale_key, method, B_data_input_scales) != 0) + { + fclose(fp); + return -1; + } + + fprintf(stderr, "write_llm_table %s method=%s\n", input_scale_key, llm_quant_method_to_string(method)); + table_count++; + } + + gemm_act_index++; + } + else if (layers[i]->type == "MultiHeadAttention") + { + const ncnn::MultiHeadAttention* mha = (const ncnn::MultiHeadAttention*)layers[i]; + + const char* reason = 0; + if (!is_supported_llm_multiheadattention(mha, &reason)) + { + print_skip_multiheadattention(mha, reason); + continue; + } + + const int qdim = mha->weight_data_size / mha->embed_dim; + + const ncnn::Mat q_weight_data = mha->q_weight_data.reshape(qdim, mha->embed_dim); + const ncnn::Mat k_weight_data = mha->k_weight_data.reshape(mha->kdim, mha->embed_dim); + const ncnn::Mat v_weight_data = mha->v_weight_data.reshape(mha->vdim, mha->embed_dim); + const ncnn::Mat out_weight_data = mha->out_weight_data.reshape(mha->embed_dim, qdim); + + const ncnn::Mat weights[4] = {q_weight_data, k_weight_data, v_weight_data, out_weight_data}; + + const size_t mha_act_offset = mha_act_index * 4; + if (method == LLM_QUANT_METHOD_AWQ || method == LLM_QUANT_METHOD_GPTQ) + { + if (mha_act_offset + 3 >= mha_act_stats.size()) + { + fprintf(stderr, "%s activation stat is missing\n", multiheadattention_name(mha)); + fclose(fp); + return -1; + } + } + + for (int j = 0; j < 4; j++) + { + char key[256]; + snprintf(key, 256, "%s_param_%d", multiheadattention_name(mha), j); + + ncnn::Mat weight_data_quantize_scales; + ncnn::Mat weight_data_input_scales; + ncnn::Mat weight_data_quantized; + int ret = 0; + if (method == LLM_QUANT_METHOD_AWQ) + { + ret = make_awq_input_scales(weights[j], mha_act_stats[mha_act_offset + j], block_size, weight_bits, awq_inner_method, awq_steps, awq_max_scale, weight_data_input_scales, weight_data_quantize_scales, quantize_num_threads); + } + else if (method == LLM_QUANT_METHOD_GPTQ) + { + ret = make_gptq_qweight(weights[j], mha_act_stats[mha_act_offset + j], block_size, weight_bits, gptq_damp, weight_data_quantize_scales, weight_data_quantized); + } + else + { + ret = make_gemm_B_scales(weights[j], block_size, weight_bits, method, weight_data_quantize_scales, quantize_num_threads); + } + if (ret != 0) + { + fclose(fp); + return ret; + } + + if (method == LLM_QUANT_METHOD_GPTQ) + { + std::string qweight_path; + const std::string qweight_name = make_qweight_filename(tablepath, key, qweight_path); + ret = write_raw_mat_file(qweight_path.c_str(), weight_data_quantized); + if (ret != 0) + { + fclose(fp); + return ret; + } + + if (write_llm_qweight_table_row(fp, key, weight_bits, block_size, method, qweight_name.c_str(), weight_data_quantize_scales) != 0) + { + fclose(fp); + return -1; + } + } + else if (write_llm_table_row(fp, key, weight_bits, block_size, method, weight_data_quantize_scales) != 0) + { + fclose(fp); + return -1; + } + + fprintf(stderr, "write_llm_table %s bits=%d block=%d method=%s\n", key, weight_bits, block_size, llm_quant_method_to_string(method)); + table_count++; + + if (method == LLM_QUANT_METHOD_AWQ) + { + char input_scale_key[512]; + snprintf(input_scale_key, 512, "%s_input_scale", key); + + if (write_llm_input_scale_row(fp, input_scale_key, method, weight_data_input_scales) != 0) + { + fclose(fp); + return -1; + } + + fprintf(stderr, "write_llm_table %s method=%s\n", input_scale_key, llm_quant_method_to_string(method)); + table_count++; + } + } + + mha_act_index++; + } + } + + fclose(fp); + + fprintf(stderr, "ncnn llm quant table created %d rows\n", table_count); + + return 0; +} + +int main(int argc, char** argv) +{ + if (argc < 4) + { + show_usage(argv[0]); + return -1; + } + + for (int i = 1; i < argc; i++) + { + if (argv[i][0] == '-') + { + show_usage(argv[0]); + return -1; + } + } + + char* inparam = argv[1]; + char* inbin = argv[2]; + char* caliblist = 0; + char* outtable = 0; + int kv_start = 4; + + if (argc >= 5 && strchr(argv[3], '=') == NULL && strchr(argv[4], '=') == NULL) + { + caliblist = argv[3]; + outtable = argv[4]; + kv_start = 5; + } + else + { + outtable = argv[3]; + kv_start = 4; + } + + int weight_bits = 6; + int block_size = 64; + const char* method = "minmax"; + const char* awq_inner = "minmax"; + int thread = 1; + + QuantNet table; + table.use_calibration_dataset = caliblist != 0; + if (caliblist) + table.listspaths = parse_comma_path_list(caliblist); + + for (int i = kv_start; i < argc; i++) + { + char* kv = argv[i]; + char* eqs = strchr(kv, '='); + if (eqs == NULL) + { + fprintf(stderr, "unrecognized arg %s\n", kv); + return -1; + } + + eqs[0] = '\0'; + const char* key = kv; + char* value = eqs + 1; + + if (strcmp(key, "method") == 0) + method = value; + else if (strcmp(key, "bits") == 0) + weight_bits = atoi(value); + else if (strcmp(key, "block") == 0) + block_size = atoi(value); + else if (strcmp(key, "thread") == 0) + thread = atoi(value); + else if (strcmp(key, "type") == 0) + table.file_type = atoi(value); + else if (strcmp(key, "shape") == 0) + table.shapes = parse_comma_int_array_list(value); + else if (strcmp(key, "awq_steps") == 0) + table.awq_steps = atoi(value); + else if (strcmp(key, "awq_samples") == 0) + table.awq_samples = atoi(value); + else if (strcmp(key, "awq_max_scale") == 0) + table.awq_max_scale = (float)atof(value); + else if (strcmp(key, "awq_inner") == 0) + awq_inner = value; + else if (strcmp(key, "gptq_samples") == 0) + table.gptq_samples = atoi(value); + else if (strcmp(key, "gptq_damp") == 0) + table.gptq_damp = (float)atof(value); + else + { + fprintf(stderr, "unrecognized arg %s\n", key); + return -1; + } + } + + const int quantize_method = llm_quant_method_from_string(method); + if (quantize_method < 0) + { + fprintf(stderr, "unsupported method=%s\n", method); + return -1; + } + + if (quantize_method == LLM_QUANT_METHOD_AWQ && !table.use_calibration_dataset) + { + fprintf(stderr, "method=awq requires calibration list\n"); + return -1; + } + + if (quantize_method == LLM_QUANT_METHOD_GPTQ && !table.use_calibration_dataset) + { + fprintf(stderr, "method=gptq requires calibration list\n"); + return -1; + } + + if (quantize_method != LLM_QUANT_METHOD_AWQ && quantize_method != LLM_QUANT_METHOD_GPTQ && table.use_calibration_dataset) + { + fprintf(stderr, "calibration list is only used by method=awq/gptq\n"); + return -1; + } + + table.awq_inner_method = llm_quant_method_from_string(awq_inner); + if (table.awq_inner_method != LLM_QUANT_METHOD_MINMAX && table.awq_inner_method != LLM_QUANT_METHOD_MSECLIP) + { + fprintf(stderr, "unsupported awq_inner=%s\n", awq_inner); + return -1; + } + + if (llm_weight_block_quantize_term(weight_bits, block_size) == 0) + { + show_usage(argv[0]); + return -1; + } + + if (thread < 1) + { + fprintf(stderr, "malformed thread %d\n", thread); + return -1; + } + + if (table.awq_steps < 0) + { + fprintf(stderr, "malformed awq_steps %d\n", table.awq_steps); + return -1; + } + + if (table.awq_samples < 1) + { + fprintf(stderr, "malformed awq_samples %d\n", table.awq_samples); + return -1; + } + + if (table.gptq_samples < 1) + { + fprintf(stderr, "malformed gptq_samples %d\n", table.gptq_samples); + return -1; + } + + if (!(table.awq_max_scale > 1.f) || table.awq_max_scale > FLT_MAX) + { + fprintf(stderr, "malformed awq_max_scale %f\n", table.awq_max_scale); + return -1; + } + + if (!(table.gptq_damp >= 0.f) || table.gptq_damp > FLT_MAX) + { + fprintf(stderr, "malformed gptq_damp %f\n", table.gptq_damp); + return -1; + } + + table.storage_type = 1; + table.quantize_num_threads = thread; + table.opt.num_threads = thread; + + if (table.load_param(inparam) != 0) + return -1; + + if (strcmp(inbin, "null") == 0) + { + DataReaderFromEmpty dr; + if (table.load_model(dr) != 0) + return -1; + } + else + { + if (table.load_model(inbin) != 0) + return -1; + } + + if (table.init() != 0) + return -1; + + if (quantize_method == LLM_QUANT_METHOD_AWQ || quantize_method == LLM_QUANT_METHOD_GPTQ) + { + if (table.collect_activation_stats() != 0) + return -1; + } + + if (table.save_table(outtable, block_size, weight_bits, quantize_method) != 0) + return -1; + + return 0; +} diff --git a/tools/quantize/ncnnllm_quant.h b/tools/quantize/ncnnllm_quant.h new file mode 100644 index 000000000000..cd8736449a32 --- /dev/null +++ b/tools/quantize/ncnnllm_quant.h @@ -0,0 +1,429 @@ +// Copyright 2026 Tencent +// SPDX-License-Identifier: BSD-3-Clause + +#ifndef NCNNLLM_QUANT_H +#define NCNNLLM_QUANT_H + +#include +#include +#include +#include +#include +#include + +enum +{ + LLM_QUANT_METHOD_MINMAX = 0, + LLM_QUANT_METHOD_MSECLIP = 1, + LLM_QUANT_METHOD_AWQ = 2, + LLM_QUANT_METHOD_GPTQ = 3 +}; + +static inline const char* llm_quant_method_to_string(int method) +{ + if (method == LLM_QUANT_METHOD_MINMAX) + return "minmax"; + if (method == LLM_QUANT_METHOD_MSECLIP) + return "mseclip"; + if (method == LLM_QUANT_METHOD_AWQ) + return "awq"; + if (method == LLM_QUANT_METHOD_GPTQ) + return "gptq"; + + return ""; +} + +static inline int llm_quant_method_from_string(const char* method) +{ + if (strcmp(method, "minmax") == 0) + return LLM_QUANT_METHOD_MINMAX; + if (strcmp(method, "mseclip") == 0) + return LLM_QUANT_METHOD_MSECLIP; + if (strcmp(method, "awq") == 0) + return LLM_QUANT_METHOD_AWQ; + if (strcmp(method, "gptq") == 0) + return LLM_QUANT_METHOD_GPTQ; + + return -1; +} + +static inline int llm_weight_block_quantize_term(int weight_bits, int block_size, bool input_scale = false) +{ + const int block_size_code = block_size == 32 ? 0 : block_size == 64 ? 1 : block_size == 128 ? 2 : -1; + if ((weight_bits != 4 && weight_bits != 6 && weight_bits != 8) || block_size_code < 0) + return 0; + + return weight_bits * 100 + (input_scale ? 10 : 0) + block_size_code; +} + +static inline int llm_weight_quantize_packed_k_bytes(int constantK, int weight_bits) +{ + if (constantK <= 0 || weight_bits <= 0) + return -1; + + const size_t packed_k_bytes = ((size_t)constantK * weight_bits + 7) / 8; + if (packed_k_bytes > (size_t)INT_MAX) + return -1; + + return (int)packed_k_bytes; +} + +static inline int float2int_weight(float v, int weight_bits) +{ + const int qmax = (1 << (weight_bits - 1)) - 1; + int q = static_cast(roundf(v)); + if (q > qmax) + q = qmax; + if (q < -qmax) + q = -qmax; + + return q; +} + +static inline void pack_signed_weight(unsigned char* ptr, int k, int weight_bits, int q) +{ + const unsigned int mask = (1u << weight_bits) - 1u; + const unsigned int v = (unsigned int)q & mask; + const int bit_offset = k * weight_bits; + + for (int b = 0; b < weight_bits; b++) + { + if (v & (1u << b)) + { + const int out_bit = bit_offset + b; + ptr[out_bit / 8] |= (unsigned char)(1u << (out_bit % 8)); + } + } +} + +static inline const char* gemm_name(const ncnn::Gemm* gemm) +{ +#if NCNN_STRING + return gemm->name.c_str(); +#else + (void)gemm; + return ""; +#endif +} + +static inline void print_skip_gemm(const ncnn::Gemm* gemm, const char* reason) +{ + fprintf(stderr, "skip_gemm %s %s\n", gemm_name(gemm), reason); +} + +static inline const char* multiheadattention_name(const ncnn::MultiHeadAttention* mha) +{ +#if NCNN_STRING + return mha->name.c_str(); +#else + (void)mha; + return ""; +#endif +} + +static inline void print_skip_multiheadattention(const ncnn::MultiHeadAttention* mha, const char* reason) +{ + fprintf(stderr, "skip_multiheadattention %s %s\n", multiheadattention_name(mha), reason); +} + +static inline bool is_supported_llm_gemm(const ncnn::Gemm* gemm, const char** reason) +{ + if (gemm->alpha != 1.f || gemm->beta != 1.f) + { + *reason = "alpha/beta is not 1"; + return false; + } + + if (gemm->transA != 0 || gemm->transB != 1) + { + *reason = "requires transA=0 transB=1"; + return false; + } + + if (gemm->constantA != 0 || gemm->constantB != 1 || gemm->constantC != 1 || gemm->constant_broadcast_type_C != -1) + { + *reason = "requires constantA=0 constantB=1 constantC=1 broadcastC=-1"; + return false; + } + + if (gemm->constantM != 0) + { + *reason = "requires dynamic M"; + return false; + } + + if (gemm->output_N1M != 0 || gemm->output_elempack != 0 || gemm->output_elemtype != 0 || gemm->output_transpose != 0) + { + *reason = "unsupported output layout"; + return false; + } + + if (gemm->quantize_term != 0) + { + *reason = "already quantized"; + return false; + } + + if (gemm->constant_TILE_M != 0 || gemm->constant_TILE_N != 0 || gemm->constant_TILE_K != 0) + { + *reason = "tiled Gemm is not supported"; + return false; + } + + const int constantN = gemm->constantN; + const int constantK = gemm->constantK; + + if (constantN <= 0 || constantK <= 0 || gemm->B_data.w != constantK || gemm->B_data.h != constantN || gemm->B_data.elemsize != 4u) + { + *reason = "B weight shape or storage is unsupported"; + return false; + } + + return true; +} + +static inline bool is_supported_llm_multiheadattention(const ncnn::MultiHeadAttention* mha, const char** reason) +{ + if (mha->quantize_term != 0) + { + *reason = "already quantized"; + return false; + } + + if (mha->embed_dim <= 0 || mha->num_heads <= 0 || mha->embed_dim % mha->num_heads != 0 || mha->weight_data_size <= 0 || mha->weight_data_size % mha->embed_dim != 0 || mha->kdim <= 0 || mha->vdim <= 0) + { + *reason = "requires valid embed_dim/num_heads/weight_data_size/kdim/vdim"; + return false; + } + + const int qdim = mha->weight_data_size / mha->embed_dim; + if (mha->q_weight_data.elemsize != 4u || mha->q_weight_data.w != mha->embed_dim * qdim) + { + *reason = "q weight shape or storage is unsupported"; + return false; + } + + if (mha->k_weight_data.elemsize != 4u || mha->k_weight_data.w != mha->embed_dim * mha->kdim) + { + *reason = "k weight shape or storage is unsupported"; + return false; + } + + if (mha->v_weight_data.elemsize != 4u || mha->v_weight_data.w != mha->embed_dim * mha->vdim) + { + *reason = "v weight shape or storage is unsupported"; + return false; + } + + if (mha->out_weight_data.elemsize != 4u || mha->out_weight_data.w != qdim * mha->embed_dim) + { + *reason = "out weight shape or storage is unsupported"; + return false; + } + + if (mha->q_bias_data.elemsize != 4u || mha->q_bias_data.w != mha->embed_dim || mha->k_bias_data.elemsize != 4u || mha->k_bias_data.w != mha->embed_dim || mha->v_bias_data.elemsize != 4u || mha->v_bias_data.w != mha->embed_dim || mha->out_bias_data.elemsize != 4u || mha->out_bias_data.w != qdim) + { + *reason = "bias shape or storage is unsupported"; + return false; + } + + return true; +} + +static inline float choose_weight_scale(const float* ptr, int size, int weight_bits, int method) +{ + float absmax = 0.f; + for (int i = 0; i < size; i++) + { + absmax = std::max(absmax, (float)fabs(ptr[i])); + } + + if (absmax == 0.f) + return 1.f; + + const int qmax = (1 << (weight_bits - 1)) - 1; + if (method != LLM_QUANT_METHOD_MSECLIP) + return (float)qmax / absmax; + + float best_scale = (float)qmax / absmax; + float best_error = 0.f; + for (int i = 0; i < size; i++) + { + const int q = float2int_weight(ptr[i] * best_scale, weight_bits); + const float deq = q / best_scale; + const float diff = ptr[i] - deq; + best_error += diff * diff; + } + + const int search_steps = 20; + for (int s = 1; s <= search_steps; s++) + { + const float shrink = 1.f - 0.5f * s / search_steps; + const float scale = (float)qmax / (absmax * shrink); + + float error = 0.f; + for (int i = 0; i < size; i++) + { + const int q = float2int_weight(ptr[i] * scale, weight_bits); + const float deq = q / scale; + const float diff = ptr[i] - deq; + error += diff * diff; + } + + if (error < best_error) + { + best_error = error; + best_scale = scale; + } + } + + return best_scale; +} + +static inline int make_gemm_B_scales(const ncnn::Mat& B_data, int block_size, int weight_bits, int method, ncnn::Mat& B_data_quantize_scales, int num_threads = 1) +{ + const int constantN = B_data.h; + const int constantK = B_data.w; + const int block_count = (constantK + block_size - 1) / block_size; + + B_data_quantize_scales.create(block_count, constantN); + if (B_data_quantize_scales.empty()) + return -100; + + #pragma omp parallel for num_threads(num_threads) + for (int n = 0; n < constantN; n++) + { + const float* ptr = B_data.row(n); + float* scale_ptr = B_data_quantize_scales.row(n); + + for (int b = 0; b < block_count; b++) + { + const int k0 = b * block_size; + const int max_kk = std::min(block_size, constantK - k0); + scale_ptr[b] = choose_weight_scale(ptr + k0, max_kk, weight_bits, method); + } + } + + return 0; +} + +static inline int pack_gemm_B_from_scales(const ncnn::Mat& B_data, const ncnn::Mat& B_data_quantize_scales, int block_size, int weight_bits, ncnn::Mat& B_data_quantized) +{ + const int constantN = B_data.h; + const int constantK = B_data.w; + const int block_count = (constantK + block_size - 1) / block_size; + + if (B_data_quantize_scales.w != block_count || B_data_quantize_scales.h != constantN) + { + fprintf(stderr, "Gemm B scale shape mismatch expected=%d,%d got=%d,%d\n", block_count, constantN, B_data_quantize_scales.w, B_data_quantize_scales.h); + return -1; + } + + for (int n = 0; n < constantN; n++) + { + const float* scale_ptr = B_data_quantize_scales.row(n); + for (int b = 0; b < block_count; b++) + { + const float scale = scale_ptr[b]; + if (!(scale > 0.f) || scale > FLT_MAX) + { + fprintf(stderr, "Gemm B scale is invalid n=%d block=%d scale=%f\n", n, b, scale); + return -1; + } + } + } + + const int packed_k_bytes = llm_weight_quantize_packed_k_bytes(constantK, weight_bits); + if (packed_k_bytes <= 0) + return -1; + + B_data_quantized.create(packed_k_bytes, constantN, (size_t)1u); + if (B_data_quantized.empty()) + return -100; + + memset(B_data_quantized.data, 0, B_data_quantized.total() * B_data_quantized.elemsize); + + for (int n = 0; n < constantN; n++) + { + const float* ptr = B_data.row(n); + const float* scale_ptr = B_data_quantize_scales.row(n); + unsigned char* qptr = B_data_quantized.row(n); + + for (int b = 0; b < block_count; b++) + { + const int k0 = b * block_size; + const int max_kk = std::min(block_size, constantK - k0); + const float scale = scale_ptr[b]; + + for (int k = 0; k < max_kk; k++) + { + const int q = float2int_weight(ptr[k0 + k] * scale, weight_bits); + pack_signed_weight(qptr, k0 + k, weight_bits, q); + } + } + } + + return 0; +} + +static inline int make_and_pack_gemm_B(const ncnn::Mat& B_data, int block_size, int weight_bits, int method, ncnn::Mat& B_data_quantized, ncnn::Mat& B_data_quantize_scales) +{ + int ret = make_gemm_B_scales(B_data, block_size, weight_bits, method, B_data_quantize_scales); + if (ret != 0) + return ret; + + return pack_gemm_B_from_scales(B_data, B_data_quantize_scales, block_size, weight_bits, B_data_quantized); +} + +static inline int write_llm_table_row(FILE* fp, const char* key, int weight_bits, int block_size, int method, const ncnn::Mat& scales) +{ + if (weight_bits != 4 && weight_bits != 6 && weight_bits != 8) + return -1; + + fprintf(fp, "%s bits=%d block=%d method=%s ", key, weight_bits, block_size, llm_quant_method_to_string(method)); + + const float* ptr = scales; + const size_t size = (size_t)scales.w * scales.h * scales.d * scales.c; + for (size_t i = 0; i < size; i++) + { + fprintf(fp, "%.9g ", ptr[i]); + } + fprintf(fp, "\n"); + + return 0; +} + +static inline int write_llm_input_scale_row(FILE* fp, const char* key, int method, const ncnn::Mat& scales) +{ + fprintf(fp, "%s method=%s ", key, llm_quant_method_to_string(method)); + + const float* ptr = scales; + const size_t size = (size_t)scales.w * scales.h * scales.d * scales.c; + for (size_t i = 0; i < size; i++) + { + fprintf(fp, "%.9g ", ptr[i]); + } + fprintf(fp, "\n"); + + return 0; +} + +static inline int write_llm_qweight_table_row(FILE* fp, const char* key, int weight_bits, int block_size, int method, const char* qweight, const ncnn::Mat& scales) +{ + if (weight_bits != 4 && weight_bits != 6 && weight_bits != 8) + return -1; + + fprintf(fp, "%s bits=%d block=%d method=%s qweight=%s ", key, weight_bits, block_size, llm_quant_method_to_string(method), qweight); + + const float* ptr = scales; + const size_t size = (size_t)scales.w * scales.h * scales.d * scales.c; + for (size_t i = 0; i < size; i++) + { + fprintf(fp, "%.9g ", ptr[i]); + } + fprintf(fp, "\n"); + + return 0; +} + +#endif // NCNNLLM_QUANT_H