前言

最近计划研究下llama.cpp是如何实现混合精度量化矩阵乘的。粗看了下llama.cpp的代码，应该是比fastertransformer要友好得多的。

首先量化了一个Q5_0的llama2，然后用build/bin/llama_simple跑了一下，并用nsys抓了一下时间线。最新的llama.cpp引入了CUDA graph技术，大概是把一堆算子封装成一个graph同一进行调度，这导致直接用nsys抓到的只有一个graph，不清楚里面具体调了哪些kernel。想要知道调用了哪些kernel，需要增加参数 --cuda-graph-trace=node。最终命令为

1	nsys profile --cuda-graph-trace=node -o output ./llama-simple -m /home/data/models/llama2-7b_q5_0/Llama-2-7B-chat-hf_q5_0.gguf

在时间线里可以明确看到时长占比达到92%的kernelmul_mat_vec_q，混合精度量化的矩阵乘估计就是靠这个了。

一下子就抓到关节，比fastertransformer友好多了。

参数

分析参数可以知道，函数实现的是$X * Y$，其中X是矩阵，Y是向量，Y的列数的取值被外层函数的switch限定在1-8之间。
const void * vx：X的数据指针
const void * vy：Y的数据指针
float * dst：存放结果的数据指针
const int ncols_x：X的列数，即K
const int nrows_x, ：X的行数，即M
const int nrows_y：Y的行数，即K
const int nrows_dst：结果矩阵的行数

此外函数还接受两个模板参数，
ggml_type type：量化方法
int ncols_y：Y的列数，即N

另外由于是CUDA Kernel，还有一些Kernel参数。函数的block num和block dim由函数calc_launch_params决定：

static std::pair<dim3, dim3> calc_launch_params(const int ncols_y, const int nrows_x, const int warp_size, const mmvq_parameter_table_id table_id) {
    const int64_t nblocks = (nrows_x + calc_rows_per_block(ncols_y, table_id) - 1) / calc_rows_per_block(ncols_y, table_id);
    const dim3 block_nums(nblocks, 1, 1);
    const dim3 block_dims(warp_size, calc_nwarps(ncols_y, table_id), 1);
    return {block_nums, block_dims};
}

llama.cpp没有封装div_ceil，这个nblocks乍一看还有点懵，其实就是div_ceil(nrows_x, rows_per_block)。rows_per_block在W5A8的量化场景下依ncols_y的不同取1或2。
block的第一维取warp_size即32，第二维在ncols_y小于等于4时取4,5~8取2，其他情况取1。

代码分析

进到函数先计算了一大堆要用到的量，逐一看一下。

constexpr int qk  = ggml_cuda_type_traits<type>::qk;
constexpr int qi  = ggml_cuda_type_traits<type>::qi;
constexpr int vdr = get_vdr_mmvq(type);
constexpr mmvq_parameter_table_id table_id = get_device_table_id();
constexpr int nwarps = calc_nwarps(ncols_y, table_id);
constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_y, table_id);
constexpr int warp_size = ggml_cuda_get_physical_warp_size();

constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);

const     int tid = warp_size*threadIdx.y + threadIdx.x;
const     int row0 = rows_per_cuda_block*blockIdx.x;
const     int blocks_per_row_x = ncols_x / qk;
const     int blocks_per_col_y = nrows_y / QK8_1;
constexpr int blocks_per_iter = vdr * nwarps*warp_size / qi;

qk和qi都是从ggml_type中提取出的，其中qk表示量化块大小，即对于输入矩阵会分成qk*qk大小的小块处理，而qi表示每个线程处理的量化元素数。在量化方式取Q5_0时，两者分别取32和4。
vdr（vector dot ratio）通过查找表获得，表示每个线程在调用向量点积内核时处理的连续整数数量。

table_id大概是一个配置表之类的东西，正常情况下就取MMVQ_PARAMETERS_GENERIC。

nwarps为warp个数，根据ncols_y和table_id取， rows_per_cuda_block同理，内部都是一些查找表。

warp_size通常都取32，一些特殊的硬件除外。

vec_dot_q_cuda则为不同的量化模式匹配不同的计算核。

tid就是算下线程的id。
row0应该是该块负责的线程的第一行的索引。

Liu Yao's attic

llama.cpp的混合精度量化计算

前言

参数

代码分析