本文分享自华为云社区《昇腾CANN 7.0 黑科技：大模型推理部署技术解密》，作者：昇腾 CANN。

近期，随着生成式 AI、大模型进入公众视野，越来越多的人意识到抓住 AI 的爆发就是抓住未来智能化变革的契机。AI 基础设施如何快速部署使用，以及如何提升推理性能，逐渐成为众多企业关注的焦点。

CANN 作为最接近昇腾 AI 系列硬件产品的一层，通过软硬件联合设计，打造出适合昇腾 AI 处理器的软件架构，充分使能和释放昇腾硬件的澎湃算力。针对大模型推理场景，CANN 最新发布的 CANN 7.0 版本有机整合各内部组件，支持大模型的量化压缩、分布式切分编译、分布式加载部署，并在基础加速库、图编译优化、模型执行调度等方面针对大模型进行极致性能优化。

自动并行切分实现大模型分布式部署：

针对 LLM 模型巨大的计算和内存开销，CANN 提供自动并行切分能力，实现大模型在昇腾集群的分布式部署。自动并行切分过程可以分为 5 个步骤：

自动切分的策略以物理集群信息和模型结构为输入，进行负载切分优化的空间建模，通过策略生成-策略应用-性能模拟的多轮迭代，进而搜索得到优化的切分部署策略。

KV Cache 机制减少重复推理计算：

LLM 模型推理计算的过程可以分为 prompt 处理和后续输出 token 的自回归计算。前者有大量数据的矩阵乘，是典型的计算密集型处理，而后者随着 LLM 的执行，会积累越来越多的对话内容，基于历史输出计算得到新的 token 输出。以“盘古是一个语言模型”为例，输入内容后，每一个 token 都会生成对应的 Q、K 和 V 向量，在 attention 部分进行矩阵乘和 softmax 等计算。在这个过程中，用户 prompt 加上已经输出的 token 都要作为下一次迭代的输入，都要重新计算相应的 QKV，这造成了大量的重复计算。

为此，业界提出了 KV Cache 方法，将已经出现的 token 所计算得出的 K 和 V 向量保存在内存，仅计算最新一个 token 的 QKV，再进行矩阵乘和 softmax 计算，本质上是以空间换时间。

目前，CANN 已经全面支持 KV Cache，并实现了 KV Cache 的分布式存储、更新和复位，有效加速自回归阶段计算。

量化技术有效降低内存占用：

量化是 AI 领域的常见技术，在大模型时代，量化还有不同的特点和要求。LLM 的权重分布相对均匀，而 FM 数据存在很多离群点。传统量化算法中，直接抛弃离群点或将所有离群点纳入量化范围，均会导致精度损失，为此 CANN 支持仅 Weight 量化，INT8 量化场景相比 FP16 可降低 50%权重内存空间占用。

同时支持 KV Cache 量化，KV Cache 本质上是空间换时间，随模型层数、sequence length 的线性增长，KV Cache 量化可降低一半存储。

FlashAttention 融合算子降低访存开销：

LLM 模型中大量使用了 Multi-Head Atten-tion 结构，这不仅带来了巨大的计算量，保存数据所需的内存容量也是计算系统的关键瓶颈。对此，业界提出了 FlashAttention 融合算子，其原理是对 attention 处理过程进行切分和计算等价，使得 attention 的多个步骤可以在一个算子中完成，并且通过多重循环、每次处理一小部分数据，以近似流式的方式访问 HBM，减少了 HBM 访问的总数据量，并能够将计算和数据搬运更好的重叠隐藏。

来源：https://arxiv.org/pdf/2205.14135.pdf

CANN 针对昇腾 AI 处理器的 HBM 和缓存大小，以及数据搬运通路，优化实现 FlashAttention 融合算子，充分利用片上缓存，提升 Attention 处理性能可达 50%。

Auto Batching 调度提升算力利用率:

面对 input 阶段 compute-bound、output 阶段 memory-bound 的计算特征，以及 LLM 业务的时延需求，CANN 支持多个 input 和 output 计算集群的异构部署，并支持 LLM 计算任务的 auto batching 调度，提升 AI 算力利用率。它的原理是将不同的服务请求尽可能地聚合处理：在 input 阶段通过单 batch 和预置的多种 sequence length 模型推理，尽量降低每个请求的启动开销；在 output 阶段以 iteration 粒度调度多个服务，尽可能拼成 batch 处理，以提升计算密度，平衡计算和访存。

支持 Torch.Compile 计算图提高编程效率：

为了使开发者能够更简单的将 LLM 在昇腾平台运行推理，CANN 实现了 PyTorch 的计算图支持。开发者只需要使用 PyTorch 原生的 torch.-compile 接口，CANN 使能的 NPU 后端就会对 PyTorch 生成的 FX Graph 进行接管，基于 trace 逻辑将 AtenIR 转换为 AIR，再进行端到端的图编译深度优化，从而降低推理阶段的内存需求、提升计算性能，同时最大程度的减少开发者的修改工作。

来源：https://pytorch.org/get-started/pytorch-2.0/

这里有一个 CANN 大模型推理上手的示例。在编译阶段使用 ATC 工具对 pb 或 onnx 模型进行编译，命令参数与 CV 等经典 AI 模型类似，只是增加了集群信息和切分信息的输入。打开集群开关以及并行切分开关，同时传入集群配置文件和切分方式的配置文件，ATC 就会在编译过程中自动实现模型的切分和通信算子插入。

atc --model=./matmul2.pb
--soc_version=Ascend910
--output=test910_parallel
--distributed_cluster_build=1
--cluster_config=./numa_config_910_2p.json
--enable_graph_parallel="1"
--graph_parallel_option_path=./parallel_option.json

在执行阶段，通过 LoadGraph 接口载入 om 离线模型，CANN 会将各个模型切片载入到相应的昇腾 AI 处理器 device 上，然后再使用既有的 RunGraph 接口即可执行推理。

经过计算/通信并行、图优化、算子调优等优化，LLAMA 65B 推理性能可较优化前提升一倍以上，端到端耗时可以达到 8s 左右，仍有提升空间。

总而言之，在大模型技术日新月异不断迭代的时代背景下，昇腾 CANN 将会持续深耕大模型优化 &加速技术，比如继续探索面向在线服务的调度优化，缩短服务时延；基于计算图的 weight 预取与 Cache 驻留优化，提升访存性能；亲和 FlashAttention 业界最新融合算子，提升计算性能；支持更丰富的量化计算组合、模型稀疏，降低内存占用...随着大模型规模化商业落地，以昇腾 CANN 为核心的昇腾 AI 基础软硬件平台，将持续提升大模型推理部署场景的核心竞争力，为客户提供最优选择！

点击关注，第一时间了解华为云新鲜技术~