一，前言
pytorch原生库加速文本生成模型～
特征：
- 低延迟
- 小于1000行代码
- 仅依赖pytorch的原生库
- int8/int4量化
- Speculative decoding（推测性解码？）
- 张量并行
- 支持Nvidia和AMD的GPU显卡
pytorch团队出品，仅依赖pure pytorch
二，问题来源
系列制作的第二篇，第一篇是Segment Anything over 8x，本篇则是针对文本生成模型。
LLM模型发展很快，功能强大，但不是那么易用，比如需要转化为特定格式或者特殊的依赖项等。可否在原生pytorch上来优化？或者尝试找找优化后的极限性能？
在 Pytorch Developer Conference上指出，在不损失精度的情况下，可以用原生的 pytorch 从头写LLM模型而获得近10x的效率。可优化的点包含：
- Torch.compile->Pytorch的模型编译，一行代码对已有的模型进行加速；
- GPU量化->精度降低，int8/int4；
- Speculative decoding->一种大模型推理加速方法，使用一个小的「draft」模型来预测大的「目标」模型的输出；
- 张量并行->多设备加速；
注意【对比参照数据】：下面所以对比的起点都是(25.5 TOK/S)
测试Llma-7B模型，可以看到的问题是：CPU overhead bound，CPU的开销限制问题，也就是CPU 无法足够快地告诉 GPU 做什么，以便充分利用 GPU。能做些什么来解决或者缓解这个问题？ 第一，我们可以用 C++ 重写实现，甚至可能完全避开框架并编写原始 CUDA，或者可以一次向 GPU 发送更多工作。
通过一次发送大量工作，可以让 GPU 保持忙碌！ 尽管在训练过程中，可能只是通过增加批量大小来完成，但在推理过程中如何做到这一点？
优化点1-通过torch.compile和KV静态缓存降低CPU开销(107.0 TOK/S)
结论：性能提升超过四倍
kv-cache
kv-cache 是一种推理时的时间优化策略，可缓存为先前 token 计算的激活。 然而，当我们生成更多令牌时，kv-cache 的“逻辑长度”就会增长，这是有问题的。首先，每次缓存增长时重新分配（或复制）kv-cache 的成本非常高。 其次，这种动态使得减少开销变得更加困难，因为不再能够利用 cudagraphs 等方法。
所以给出了 “static kv-cach”方案，直接静态分配 kv-cache 的最大大小，然后屏蔽计算注意力时未使用的值。
prefill phase
Transformer 文本生成可被视为一个两阶段过程：1. 处理整个提示的预填充；2. 自回归生成每个标记的解码。一旦 kv-cache 静态化，解码就可以完全静态化，但由于prompt长度可变，预填充阶段仍然需要更多的动态性。 因此，我们实际上需要使用单独的编译策略来编译这两个阶段。
why work?
第一个因素是减少开销。
Torch.compile 能够通过各种优化来减少开销，最有效的优化之一称为 CUDAGraphs。 对用户而言，配置“reduce-overhead”就可以自动应用此功能，无须额外的工作和代码编写。第二个因素是 torch.compile 能生成更快的内核。
在上面的解码 benchmark 中，torch.compile 实际上从头开始生成每个内核，包括矩阵乘法和注意力，而且，这些内核比内置的替代方案（CuBLAS 和 FlashAttention2）更快！
矩阵乘法有: torch.mul(a, b)点乘；torch.mm(a, b)二维的矩阵乘法，超出二维会报错；和torch.matmul(a, b)是高维矩阵乘法。
优化点2-通过 INT8 权重量化缓解内存带宽瓶颈(157.4 TOK/S)
是否可以更快？可以对比下理论峰值，在这个例子中，瓶颈在于将权重从 GPU 全局内存加载到寄存器的成本。 换句话说，每次前向传递都要求我们“接触”GPU 上的每个参数。 那么，理论上我们能够以多快的速度“触及”模型中的每个参数？引入模型带宽利用率参数（MBU），公式如下：
eg:
我们有一个 7B（Llama-7b） 参数模型，每个参数都存储在 fp16 中（每个参数 2 个字节），我们实现了 107 个令牌/秒，Nvidia A100-80GB 理论内存带宽为 2 TB/s。
FP32,4 bytes; FP16, 3 bytes; int8, 1 bytes; int4, 0.5bytes.
即使只是复制内存也很难突破 85%，而本文模型达到了72%，意味着我们非常接近理论极限，并且我们显然在从内存加载权重方面遇到了瓶颈。 如果不从问题源头出发，最多可能只能再争取 10% 的性能。再次回过头来看公式，此时还能优化的就是bytes per param，即参数的内存占用，由此引出int8量化。
注意 仅量化权重，计算本身仍然在 fp16 中完成，使得这种形式的量化易于应用，并且精度几乎没有降低。而且，int8需要与torch.compile融合使用，否则性能可能比fp16（无论是否带torch.compile）的还要低/
结论2：约有50%的性能提升。
优化点3-用推测解码来重构问题
量化之后，下一个问题就凸显出来了，为了生成100个token，需要load模型权重100次。有可以优化的方法吗？乍看貌似没有，因为在自回归生成任务中有严格的序列依赖性，但是本文仍然提出了推测解码来解决该问题？
在 Transformer 模型的推理中，Verity 将由更大的模型发挥作用，我们希望其输出满足我们的任务，称为 Verifier 模型。 同样，Drake 将由一个较小的模型（称为 Draft 模型）来扮演，该模型能够比大模型更快地生成文本。 因此，我们将使用 Draft 模型生成 8 个令牌，然后使用 Verifier 模型并行处理所有 8 个令牌，并丢弃不匹配的令牌。
如上所述，推测解码的一个关键特性是它不会改变输出的质量。 只要使用 Draft 模型生成令牌 + Verifier 模型验证令牌所需的时间少于生成这些令牌所需的时间即可。
尽管推测性解码保证我们在数学上与常规生成相比具有相同的结果，但它在运行时的性能会根据待生成的文本以及 Draft 和 Verifier 模型的对齐程度而变化。 例如，当运行 CodeLlama-34B + CodeLlama-7B 时，我们能够在生成代码时获得 2 倍的 $tokes/s$ 提升。 另一方面，当使用 Llama-7B + TinyLlama-1B 时，我们只能获得大约 1.3 倍的 $tokes/s$ 提升。
在AMD显卡上的实现
如上所述，解码中的每个内核都是由 torch.compile 从头开始生成的，并转换为 OpenAI Triton。 由于 AMD 有 torch.compile 后端（还有 Triton 后端），我们可以简单地完成上述所有优化。在 AMD GPU 上，通过 int8 量化，我们能够使用 MI250x 的 1 个 GCD（即二分之一）实现 102.5 token/s！
优化点4-通过 INT4 量化和 GPTQ (202.1 TOK/S) 进一步减小权重的大小(202.1 TOK/S)
怎么减小int4量化带来的精度损失？
- 第一个是拥有更细粒度的缩放因子。 考虑缩放因子的一种方法是，当我们有量化张量表示时，它处于浮点张量（每个值都有缩放因子）和整数张量（没有值有缩放因子）之间的滑动比例。 例如，对于 int8 量化，我们每行都有一个缩放因子。 然而，如果我们想要更高的精度，我们可以将其更改为“每 32 个元素一个缩放因子”。 我们选择组大小为 32 来最小化准确性下降，这也是社区中的常见选择。
- 另一种是使用比简单地对权重进行舍入更先进的量化策略。 例如，GPTQ 等方法利用示例数据来更准确地校准权重。 在本例中，我们基于 PyTorch 最近发布的 torch.export 在存储库中构建了 GPTQ 实现的原型。此外，我们需要将 int4 反量化与矩阵向量乘法融合的内核。 在这种情况下，torch.compile 不幸地无法从头开始生成这些内核，因此我们在 PyTorch 中利用一些手写的 CUDA 内核。
此外，我们需要内核来将 int4 反量化与矩阵向量乘法融合。 在这种情况下，torch.compile 却无法从头开始生成这些内核，因此我们在 PyTorch 中需要手写一些 CUDA 内核。这些技术需要一些额外的工作，但将它们组合在一起会产生更好的性能！
优化点5-将所有内容组合在一起（244.7 TOK/S）
单GPU->多GPU，是增大了MBU公式的分母，即内存带宽。
选择哪种并行策略呢？请注意，为了减少一个 sample 的延迟，能够同时在更多设备上利用内存带宽，这需要将一个令牌的处理拆分到多个设备上。 换句话说，我们需要使用张量并行性。
幸运的是，PyTorch 还提供了与 torch.compile 组合的张量并行的低级工具，也还在开发用于表达张量并行性的更高级别的 API！
总结
简单性：忽略量化，model.py（244 LOC）+generate.py（371 LOC）+tp.py（151 LOC）得出766 LOC，实现快速推理+推测解码+张量并行。
性能：通过 Llama-7B，我们能够使用编译 + int4 Quant + 推测解码来达到 241 tok/s。 借助 llama-70B，我们还能够引入张量并行性以达到 80 tok/s。 这些都接近或超过 SOTA 性能数据！