想用大模型赚钱？这个实力强劲的新面孔决定先把推理利润打下来。
大模型业务到底多烧钱？前段时间，华尔街日报的一则报道给出了参考答案。
报道显示，微软的 GitHub Copilot 业务（背后由 OpenAI 的 GPT 大模型支撑）虽然每月收费 10 美元，但平均还是要为每个用户倒贴 20 美元。可见当前 AI 服务提供商们正面临着严峻的经济账挑战 —— 这些服务不仅构建利润高昂，运营利润也非常高。
有人比喻说：「使用 AI 总结电子邮件，就像是让兰博基尼送披萨外卖。」
对此，OpenAI 算过一笔更详细的账：当上下文长度为 8K 时，每 1K 输出 token 的利润为 3 美分，输出的利润为 6 美分。目前，OpenAI 拥有 1.8 亿用户，每天收到的查询数量超过 1000 万次。这样算来，为了运营 ChatGPT 这样的模型，OpenAI 每天都需要在必要的盘算硬件上投入大约 700 万美元，可以说是贵得吓人。
降低 LLM 的推理利润势在必行，而提升推理速度成为一条行之有效的关键路径。
实际上，研究社区已经提出了不少用于减速 LLM 推理任务的技术，包括 DeepSpeed、FlexGen、vLLM、OpenPPL、FlashDecoding 和 TensorRT-LLM 等。这些技术自然也各有优势和短板。其中，FlashDecoding 是 FlashAttention 作者、斯坦福大学团队的 Tri Dao 等人在上个月提出的一种 state-of-the-art 方法，它通过并行加载数据，大幅提升了 LLM 的推理速度，被认为极具潜力。但与此同时，它也引入了一些不必要的盘算开销，因此依然存在很大的优化空间。
为了进一步解决问题，近日，来自无问芯穹（Infinigence-AI）、清华大学和上海交通大学的联合团队提出了一种新方法 FlashDecoding++，不仅能带来比之前方法更强的减速能力（可以将 GPU 推理提速 2-4 倍），更重要的是还同时支持 NVIDIA 和 AMD 的 GPU！它的核心思想是通过异步方法完成注意力盘算的真正并行，并针对「矮胖」矩阵乘优化减速 Decode 阶段的盘算。
论文地址：https://arxiv.org/pdf/2311.01282.pdf
将 GPU 推理提速 2-4 倍，
FlashDecoding++ 是怎么做到的？
LLM 推理任务一般为输出一段文字（token），通过 LLM 模型盘算继续生成文字或其他形式的内容。
LLM 的推理盘算可被分为 Prefill 和 Decode 两个阶段，其中 Prefill 阶段通过理解输出文字，生成第一个 token；Decode 阶段则顺序输出后续 token。在两个阶段，LLM 推理的盘算可被分为注意力盘算和矩阵乘盘算两个主要部份。
对于注意力盘算，现有工作如 FlashDecoding 切分注意力盘算中的 softmax 算子完成并行加载数据。这一方法由于需要在不同部份 softmax 同步最大值，在注意力盘算中引入了 20% 的盘算开销。而对于矩阵乘盘算，在 Decode 阶段，左乘矩阵多表现为「矮胖」矩阵，即其行数一般不大（如 <=8），现有 LLM 推理引擎通过补 0 将行数扩充到 64 从而利用 Tensor Core 等架构减速，从而导致大量的无效盘算（乘 0）。
为解决上述问题，「FlashDecoding++」的核心思想在于，通过异步方法完成注意力盘算的真正并行，并针对「矮胖」矩阵乘优化减速 Decode 阶段的盘算。
异步并行部份 softmax 盘算
图 1 异步并行部份 softmax 盘算
先前工作对每个部份 softmax 盘算求输出最大值作为缩放系数，避免 softmax 盘算中 e 指数的溢出，这就导致了不同部份 softmax 盘算的同步开销（图 1 (a)(b)）。
图 2 softmax 输出值统计分布
「FlashDecoding++」指出，对于大部份 LLM，其 softmax 的输出分布较为集中。如图 2 所示，Llama2-7B 的 softmax 输出 99.99% 以上集中在 [-16.8, 6.5] 这个区间。因此，「FlashDecoding++」提出在部份 softmax 盘算时使用一个固定的最大值（图 1 (c)），从而避免了不同部份 softmax 盘算间的频繁同步。而当小概率发生的输出超出给定范围时，「FlashDecoding++」对这一部份的 softmax 盘算退化为原先的盘算方法。
「矮胖」矩阵乘的优化
图 3 「矮胖」矩阵乘切分与双缓存机制
由于 Decode 阶段的输出为一个或几个 token 向量，因此该阶段的矩阵乘表现为「矮胖」形状。以矩阵 A×B=C 为例，A 与 B 矩阵的形状为 M×K 与 K×N，「矮胖」矩阵乘即 M 较小的情况。「FlashDecoding++」指出「矮胖」矩阵乘一般缓存受限，并提出双缓存机制等优化手段进行减速（图 3）。
图 4 自适应矩阵乘完成
此外，「FlashDecoding++」进一步指出，在 LLM 推理阶段，针对特定模型，N 和 K 的取值固定。因此，「FlashDecoding++」会根据 M 的大小，自适应选取矩阵乘的最优完成。
将 GPU 推理提速 2-4 倍
图 5 「FlashDecoding++」NVIDIA 与 AMD 平台 LLM 推理（Llama2-7B模型，batchsize=1）
目前，「FlashDecoding++」可以完成 NVIDIA 与 AMD 等多款 GPU 后端的 LLM 推理减速（图 5）。通过减速 Prefill 阶段的首 token 生成速度，以及 Decode 阶段每个 token 的生成速度，「FlashDecoding++」可以在长、短文本的生成上均取得减速效果。相较于 FlashDecoding，「FlashDecoding++」在 NVIDIA A100 上的推理平均减速 37%，并在 NVIDIA 和 AMD 的多 GPU 后端上相较于 Hugging Face 完成减速多达 2-4 倍。
AI 大模型创业新秀：无问芯穹
该研究的三位共同一作分别是无问芯穹首席科学家、上海交通大学副教授戴国浩博士，无问芯穹研究实习生、清华大学硕士生洪可，无问芯穹研究实习生、上海交通大学博士生许珈铭。通讯作者为上海交通大学戴国浩教授和清华大学电子工程系主任汪玉教授。
创立于 2023 年 5 月的无问芯穹，目标是打造大模型软硬件一体化最佳解决方案，目前 FlashDecoding++ 已被集成于无问芯穹的大模型盘算引擎「Infini-ACC」中。在「Infini-ACC」的支持下，无问芯穹正在开发一系列大模型软硬件一体化的解决方案，其中包含大模型「无穹天权（Infini-Megrez）」、软硬件一体机等。
据了解，「Infini-Megrez」在处理长文本方面表现非常出色，将可处理的文本长度破纪录地提升到了 256k token，实测处理大约 40 万字的一整本《三体 3：死神永生》也不成问题。这是当前的大模型所能处理的最长文本长度。
此外，「Infini-Megrez」大模型在 CEval (中)、MMLU (英)、CMMLU (中)、AGIEval 等数据集上均取得了第一梯队算法性能，并依托「Infini-ACC」盘算引擎持续进化中。