Antkillerfarm Hacking V8.0

DL acceleration » 并行 & 框架 & 优化（八）——LLM Inference

2024-09-26 :: 5757 Words

快速Transformer

FlashAttention-4（续）

为了H100要做个FlashAttention-3，为了B200要做个FlashAttention-4，为了Rubin要做个FlashAttention-5。只要黄皮衣发卡不停，FlashAttention生意就能一直做下去。

每一代高性能算子都有每一代的要求：

  • Flash attention 1：tensor core和vector/SIMT core要能读写同一个SRAM。
  • Flash attention 3/DeepGEMM：tensor core和vector/SIMT core要能完全异步。
  • Deepseek Sparse Attention：vector/SIMT core需要高效地支持topK，较为高效地支持hadamard transform。

https://www.zhihu.com/question/1964791844773822881

NPU为什么很难支持FlashAttention?

过去两年，Mamba、RWKV、各种线性注意力架构之所以受到关注，一个核心论点是：Attention的二次复杂度是个根本缺陷，序列越长代价越高，所以需要用线性复杂度的替代方案。

FA4现在证明了虽然Attention的理论复杂度确实是\(O(N^2)\)，但当实现层面的常数因子被压缩到极致之后，这个二次方的实际开销比你想象的小得多。

Mamba和RWKV并不是没有价值。但它们的生存空间，正在被FA4这样的工作一寸一寸地压缩。它们需要找到一个”即使Attention再怎么优化也搞不定”的场景来证明自己——目前来看，那个场景还在地平线之外了。

FlashInfer

FlashInfer是由NVIDIA与CMU、UC Berkeley等机构联合开源的GPU专用LLM推理内核库。将FlashAttention、PageAttention、稀疏 Attention、采样、通信等全部打包，并针对LLM推理服务做了 JIT编译、Paged KV-Cache、变长批量调度等工程化增强。

https://zhuanlan.zhihu.com/p/681506469

用FlashInfer加速大语言模型推理中的自注意力操作

Dual Chunk FlashAttention

超长文本（论文、代码库、百万token级对话）在原始预训练窗口外直接推理会严重掉精度。

纯FlashAttention虽然节省显存，但序列长度N增加后仍然O(N²)地吃显存，单卡80 GB也很快就OOM。

DCFA把长序列切成≤预训练长度的chunk，先算chunk内注意力（intra-chunk），再算chunk间注意力（inter-chunk），把显存复杂度压到O(chunk_size²)，理论上可以无限外推长度。

Differential FlashAttention

传统Transformer在长上下文时会把大量注意力权重放到“噪声token”上，产生幻觉和上下文丢失。

Differential Transformer：

\[\text{DiffAttn}(X) = \underbrace{\vphantom{\lambda}\operatorname{softmax}\!\left(\frac{Q_{1}K_{1}^{\!\top}}{\sqrt{d_{k}}}\right)V}_{\text{主注意力}} - \lambda\,\underbrace{\vphantom{\lambda}\operatorname{softmax}\!\left(\frac{Q_{2}K_{2}^{\!\top}}{\sqrt{d_{k}}}\right)V}_{\text{噪声注意力}}\] \[\lambda = \exp(\lambda_{q1}\!\cdot\!\lambda_{k1}) - \exp(\lambda_{q2}\!\cdot\!\lambda_{k2}) + \lambda_{\text{init}}\]

相当于让第二个分支专门学“噪声模式”，然后显式减掉。

但这样就带来两倍KV Cache + 两次FlashAttention计算的开销。于是作者直接写了一个“一次kernel launch里跑两路”的专用CUDA kernel，起名叫Differential FlashAttention。

ThunderKittens

ThunderKittens是一个由HazyResearch团队开发的轻量级、高性能深度学习框架。对长度<1K的序列，ThunderKittens的速度比FlashAttention-2快2-5倍。

代码：

https://github.com/HazyResearch/ThunderKittens

Stream-K

我们将GPU上最经典的切块方式称为Data Parallel，即每一个计算单元都计算输出一个完整的子矩阵，所有计算单元计算的数据互不相干，从而达到并行计算的效果。

如上左图所示，我们将矩阵乘法拆分成9个子矩阵，每一个子矩阵交由一个CTA完成，这样会有9个CTA。假设我们的GPU上只有4个SM，4个SM每一wave都会计算出各自的一个子矩阵。但在最后一wave中我们只剩下一个子矩阵没有计算完成，这样就会导致我们的吞吐量只有75%=9/12。

当然，我们可以将切块变小，比如说切分成18块（如上右图所示），这样能将我们的吞吐量提升到90%=18/20。

但是这样会削弱cache作为暂存器的作用，因为每次从gmem读到smem的数据量都变小了。

我们将这种矩阵的形状不能被计算单元整除的情况称之为Tile Quantization，导致计算单元利用率低的情况称之为Quantization Inefficiency（不是量化的quantization）。

我们也可以通过在累加维度并行来减小分配给每个CTA的任务分配粒度，这样既能够减小Quantization Inefficiency，也能够防止因为切块变小带来的smem数据复用问题。这就是所谓的Split-K（如上左图所示）。FlashDecoding就用到了Split-K的思想。

但这也并不能完美解决Quantization Inefficiency，甚至还会带来一个额外的成本，即需要先将其中一部分数据写到gmem中后再做一次reduce，我们将这个操作称之为fixup。

为此，NV又提出了Stream-K，如上右图所示。简单的说法，就是按照SM的数量划分wave。每个wave可以执行分数个tile的运算，而不是整数个tile的运算。

原理就是在tile的基础上，在K上划分sub tile。显然这里的逻辑会变的更加复杂，什么时候要进行累加，什么时候要写到gmem跟其他的SM进行reduce是我们需要进行额外判断的。

参考

https://mp.weixin.qq.com/s/1R_plHqxTLE-Fw3TjYnlJQ

GPU BERT上线性能不合格，看看微信AI的PPoPP论文

https://mp.weixin.qq.com/s/OgTQ3O_6lvOG07U-tjpTDA

如何让Transformer在GPU上跑得更快？快手：需要GPU底层优化

https://zhuanlan.zhihu.com/p/638468472

从FlashAttention到PagedAttention, 如何进一步优化Attention性能

https://blog.csdn.net/v_JULY_v/article/details/144218958

一文通透vLLM与其核心技术PagedAttention：减少KV Cache碎片、提高GPU显存利用率(推理加速利器)

LLM Inference

A Survey on Efficient Inference for Large Language Models

https://zhuanlan.zhihu.com/p/653352979

LLM七种推理服务框架总结

https://zhuanlan.zhihu.com/p/671347964

大模型(LLM)推理框架汇总

https://zhuanlan.zhihu.com/p/642412124

LLM的推理优化技术纵览

https://github.com/DefTruth/Awesome-LLM-Inference

Awesome LLM Inference

https://www.zhihu.com/tardis/zm/art/647813179

大模型文本生成——解码策略（Top-k & Top-p & Temperature）

https://b23.tv/OfdfBnz

如何设置大模型推理参数，top_k，top_p, temperature, num_beams

https://blog.csdn.net/HUSTHY/article/details/125990877

Contrastive Search Decoding——一种对比搜索解码文本生成算法

https://zhuanlan.zhihu.com/p/656707466

DoLa：层对比解码改善LLM的真实性

Speculative Decoding：先用一个轻量级“提议器”快速生成K个候选token，再用主模型并行打分，能少跑很多步主模型。

《Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads》提出：

不额外训练完整草稿模型，而是在原模型最后一层隐藏状态上挂3-5个轻量级前馈头（Medusa Heads），每个头负责预测第2,3,…,K+1位置的token。这些头与原模型共享KV缓存，因此几乎不占额外显存。

EagleProposer复用主模型的Embedding与LMHead，只额外训练一个轻量的Auto-regression Head，接受率显著高于Medusa/普通小草稿模型。

https://zhuanlan.zhihu.com/p/651359908

大模型推理妙招—投机采样（Speculative Decoding）

LLM Inference的性能评估主要有以下几个方面：

  • Time To First Token (TTFT)
  • Time Per Output Token (TPOT)：解码阶段总耗时 / 生成token数
  • Latency：模型为用户生成完整响应所需的总时间。latency = (TTFT) + (TPOT) * (the number of tokens to be generated)
  • Throughput：一个推理服务器每秒可以为所有用户和请求生成的输出token数量。
  • TBT（Time Between Tokens）：第i+1个token时间戳 − 第i个token时间戳。流式/对话场景（逐token吐字）更关心TBT。

https://www.databricks.com/blog/llm-inference-performance-engineering-best-practices

LLM Inference Performance Engineering: Best Practices

PD分离

一次用户请求，实际上既包含prefill，也包含decode。一个是计算密集型，一个是访存密集型。

prefill（用户输入）和decode（模型输出）的token量在不同场景下也是不一样的。如果是简单对话场景，通常模型的decode输出会更多一些，而如果是超长上下文场景，用户先上传一本几十万字的书再进行问答，这本书的prefill会直接起飞。在Agent场景下，大量预设的prompt也会占据非常多的prefill，不过prompt的prefill有不少机会可以提前算好KV而无需每个用户请求单独重复计算。

当整个推理系统服务几千万用户时，一个batch的几十个用户请求只是开胃菜。每个用户会不间断地和大模型进行交互，发出大量请求，但这些请求的间隔时间短则几秒，长则几分钟几小时。考虑人机交互的频率，一个用户请求结束后，对应的KV-cache继续常驻在高速内存中实际意义不大。

这个从今年年中开始，各家都陆续上了PD分离方案（如MoonCake）。

Prefilling阶段是计算密集型，少量Query就可以打满GPU，大量Query反而会增加首token延迟；Decoding阶段是访存密集型，必须依赖大量Query提高GPU计算利用率。因此可以通过多台机器处理Prefilling、单台机器处理Decoding的PD分离方案，实现综合效率最大化、首token延迟（TTFT）最低、DecodeSpeed（TPS）最高。

Rubin CPX，是专门为推理阶段prefill准备的低成本卡。prefill阶段所需带宽较小，配置“瘦带宽”而“肥容量”的GDDR7，显然更为合适。与此同时，英伟达与Groq达成了200亿美元战略授权协议。“Groq 风味”的硅——英伟达正在将其整合进自身推理路线图——将承担高速decode引擎的角色。

https://www.zhihu.com/question/1949040585416636414

如何看待英伟达最新发布的Rubin CPX及相应的上下文理解/生成分离设计?

Antkillerfarm
antkillerfarm@sohu.com

github.com/antkillerfarm

Fork me on GitHub