前一篇是: 0-使能 CANN 后端, 但目前这一篇跟前一篇并没有什么太大的关系, 且本篇原来打算的题目是 1-给 Ascend310B 补全算子 , 主要是打算实现:

1.BatchMatmulCustom: n_dim=3 时, 对应的 CANN 官方算子为 aclnnBatchMatMul;

2.MatmulCustom: n_dim=2 时, 对应的 CANN 官方算子为 aclnnMatmul;

3.MmCustom: 其他情形(由于 GGML 中张量最多 4 维, 那么也就是 n_dim=2 时), 对应的 CANN 官方算子为 aclnnMm

但发现 Ascend310B 这边特性缺的实在是太多(算子库中甚至连 Cast 和 Pad 都不支持), 在小尺寸矩阵乘法上(好像是 k<32 时)会出问题(当然应该是我写得太烂了).

但最麻烦的一点是 CANN 这一套工具链部署算子(可能是我使用的方法有问题)比较费劲, 想要算子在 llama.cpp 中可用, 得编译打包然后安装, 整个耗时差不多 1 分钟. 且在算子这边做完了算子直调的测试后, 放 llama.cpp 上仍然会有各种奇奇怪怪的问题, 也比较难 Debug.

遂放弃的原来补全算子的打算, 决定曲线救国. 先考虑使用 aclblas 来做矩阵乘法, 但这个 CBLAS 接口跟我想象中有点不太一样, 我希望是类似 cuBLAS 这种, 但这个好像还得手动设置模型之类的, 无法比较方便地去使用, 遂放弃.

然后发现了 CATLASS, 发现这个框架的 examples/00_basic_matmul 是可以在 Ascend310B 上跑的, 其他的示例改改也能用, 同时比较方便去封装成静态库与动态库, 于是就利用其写了个 AscendBLAS.

AscendBLAS 的代码开源在: https://github.com/xvyv99/ascend-blas , 适配了 Ascend310B 的 llama.cpp 代码开源在: https://github.com/xvyv99/llama.cpp .

CATLASS

CATLASS(CANN Templates for Linear Algebra Subroutines), 定位类似于 Nvidia 的 CUTLASS. 相比于直接去写 AscendC 算子, 使用 CATLASS 框架能够很大程度上降低心智负担, 同时框架也比较现代, 构建与测试都比较方便.

后续的开发基于 CATLASS 的提交 108c123335b08ba0dc1a72258bf764bce295935d.

Ascend310B 适配

由于 CATLASS 并不原生支持 Ascend310B, 所以 examples 下的许多例子都需要做相应修改才能正常运行.

比如说 examples/00_basic_matmul/CMakeLists.txt 中得把 catlass_example_add_executable 中的 mix 改成 cube, 这应该是因为 Ascend310B 并不支持 Mix 模式, 否则运行时会有问题 Compare failed., 且这个改动对其他例子也适用.

然后就是 examples/18_gemv_aic/, 这个也会有问题(此时出现的是运行时报错 507015), 猜测可能是因为不支持 AIC 跟 AIV 核间同步的相关配置, 而这部分在 include/catlass/gemv/kernel/kernel_gemv_aic.hpp 中, 删掉跟 CrossCore 相关的配置即可.

AscendBLAS

最开始时提到, 由于 CANN 的 CBLAS 接口使用起来有点麻烦, 遂利用 CATLASS 写了个类似 CBLAS 接口的库 AscendBLAS. 为了方便起见, API 并不与 BLAS 对齐, 且目前仅实现了 AscendBLAS::Gemm(来源于 CATLASS 的 examples/00_basic_matmul) 与 AscendBLAS::Gemv(来源于 CATLASS 的 examples/18_gemv_aic/).

其中 AscendBLAS::Gemm 的函数签名如下:

/**
 * @brief General matrix-matrix multiplication: C = alpha * op(A) * op(B) + beta * C
 *
 * @param transA  Transpose operation for matrix A
 * @param transB  Transpose operation for matrix B
 * @param m       Number of rows in A and C
 * @param n       Number of columns in B and C
 * @param k       Number of columns in A and rows in B
 * @param alpha   Scalar multiplier for A * B
 * @param a       Pointer to matrix A
 * @param lda     Leading dimension of A
 * @param b       Pointer to matrix B
 * @param ldb     Leading dimension of B
 * @param beta    Scalar multiplier for C
 * @param c       Pointer to matrix C
 * @param ldc     Leading dimension of C
 * @param dtype   Data type (FP32 or FP16)
 * @param stream  Ascend stream for asynchronous execution
 */
void Gemm(
    Transpose transA, Transpose transB,
    int m, int n, int k,
    float alpha,
    const void* a, int lda,
    const void* b, int ldb,
    float beta,
    void* c, int ldc,
    DataType dtype,
    aclrtStream stream
);

且支持的精度有:

1.fp32xfp32->fp32;

2.fp16xfp16->fp16;

3.int8xint8->int32.

ggml-blas 后端

看起来 ggml-blas 后端并没有什么人去用, 因为我运行时遇到了很严重的问题: Misc. bug: out-of-range access during model loading with BLAS backend. 具体原因是 ggml-blas 后端在获取可用内存和总内存时(通过 ggml_backend_blas_device_get_memory)均得到了 0, 然后 llama_model::load_tensors 方法中就出现了除零异常.

Llama.cpp 的 docs/build.md 中提到使用 BLAS 只会在 Prefill 阶段可能带来一些性能提升(且是在 batch sizes 大于 32 的情形下, 虽然默认是 512), 不会对 Decode 阶段带来性能提升.

Building the program with BLAS support may lead to some performance improvements in prompt processing using batch sizes higher than 32 (the default is 512). Using BLAS doesn't affect the generation performance.

在测试时发现一个比较有意思的点, 就算 llama-cli 处设置 --device none, 此时还是会使用 ggml-blas 后端进行加速. 这是因为 ggml-blas 后端是一个 GGML_BACKEND_DEVICE_TYPE_ACCEL 类型的后端(在 ggml_backend_blas_device_get_type 函数处).

这个类型的后端, 不管是否指定 --device, 都会跟 ggml-cpu 一起初始化(在 src/llama-context.cpp 的 llama_context::llama_context 方法中有体现). 除了 ggml-blas 外, 还有 ggml-zdnn 跟 ggml-zendnn 也是这个类型的后端.

代码概览

ggml-blas 后端只实现了两个算子 GGML_OP_MUL_MAT 与 `GGML_OP_OUT_PROD.

对于 GGML_OP_MUL_MAT 算子, 其内部会先将张量 src0 的(如果能转换到 fp32 的话)类型都转换成 fp32 后再计算, 转换这步可以使用 OpenMP 或者 std::async 去进行并行. 接着通过 for-loop 循环的方式去调用 cblas_sgemm 来计算 Batch Matmul.

为什么选 ggml-blas

选用的 ggml-blas 进行魔改的原因是其比较纯净, 可以当成脚手架来使用. 我们可以直接封装一个包含算子的类 BLAS API 动态链接库, 然后替换原 BLAS 调用就可以用了, 而不用处理比较麻烦的 aclTensor 类型转换等比较繁琐的事情.

创建 ggml-ascendrc 后端

为了不与 ggml-blas 后端混淆, 所以基于 ggml-blas 创建了一个新后端, 称为 ggml-ascendrc(取这个名称主要是因为目前 OrangePI AI Pro 的 NPU Ascend310B1 是在 Ascend RC 形态下工作的).

其中 ggml_backend_ascendrc_context 参考(其实是照搬)了 ggml-cann 后端 ggml_backend_cann_context 的写法.

然后就是把 GGML_OP_MUL_MAT 算子实现中的 cblas_sgemm 改成了调用 AscendBLAS::Gemm.

不过目前只支持 fp32 和 fp16, 对于 src0 为 fp16 的情形, ggml-ascendrc 会先将 src1 从 fp32 转换为 fp16, 然后调用 AscendBLAS 的函数得到 fp16 的结果, 最后将结果从 fp16 转换为 fp32.

性能分析

跑 test-backend-ops 的 test 结果如下.

而 Prefill 阶段的 perf 结果如下: Decode 阶段的 perf 结果如下:

可以看出, 比 CPU 后端(其 Prefill 阶段最高应该只有不到 50 GFLOPS)要好一点, 但不多.

因为根据我这边的测试(测试基于 AscendBLAS 和 MindSpore 的 mindspore.ops.matmul), Ascend310B1 在 fp16 精度下的 GEMM 应该是能到 4TOPS, 所以应该有很大的优化空间.

目前我认为主要的可以优化的地方在于:

1.Decode 阶段采用 GEMV;

2.将 for-loop 去做 Batch Matmul 的实现替换为原生的 Batch Matmul 实现;

3.去做随路量化;

其他的得 Profile 后才能确定了.

经过目前的改动, llama.cpp 已经可以用到 NPU Ascend310B 去进行加速了, 虽然速度比较捉急, 不过至少比 OrangePI AI Pro 孱弱的 CPU 要好很多了( .

以下是使用 ggml-ascendrc 后端跑的效果:

以下是 OrangePI AI Pro CPU 跑的效果: