Aiinfra

理解LLM推理中deterministic问题来源 Wiki上对deterministic算法的定义是: “a deterministic algorithm is an algorithm that, given a particular input, will always produce the same output.” 而我们在文中要讨论的，即对于LLM这个context下的deterministic问题，我会先从inference角度（即重复给定一个确定的input，模型的推理为什么无法给定确定的输出）进行问题的理解，再进一步讨论RL工程中的training & inference之间差异，可能会导致RL训练的崩溃问题，并继续讨论业界现在已有的解决方案、与还在working-in-progress的工作。 浮点数的非结合性 thinking machines lab针对batch invariant讨论的文章，详细地解释了在LLM推理中不确定性的来原，即因为精度有限，GPU浮点数运算中的结合性通常不成立： $$(a+b)+c \neq a+(b+c) $$ 这篇arxiv文章，则更深入得说明了这个问题： Floating-point arithmetic in GPUs exhibits non-associativity, meaning (a+b)+c≠a+(b+c) due to finite precision and rounding errors. This property directly impacts the computation of attention scores and logits in the transformer architecture, where parallel operations across multiple threads can yield different results based on execution order. ...

本文从数学原理出发，深入分析FlashAttention的核心思想、算法设计和各版本演进，通过详实的数学推导、直观的流程图表和具体的数值示例，帮助读者真正掌握这一革命性的Attention优化技术。 1. 问题的本质：传统Attention的根本瓶颈 1.1 传统Attention机制的计算模式 传统的Self-Attention机制遵循如下计算流程： $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$让我们用具体数值来理解这个过程的复杂性： 示例场景：考虑一个典型的语言模型场景 序列长度：$n = 2048$（如GPT-2的上下文长度） 特征维度：$d_k = 64$（每个attention head的维度） 输入张量形状：$Q, K, V \in \mathbb{R}^{2048 \times 64}$ 第一步：计算注意力得分矩阵 $$S = \frac{QK^T}{\sqrt{d_k}} \in \mathbb{R}^{2048 \times 2048}$$这一步产生了一个 $2048 \times 2048 = 4,194,304$ 个元素的矩阵，以FP16精度存储需要约8MB内存。 第二步：Softmax归一化 $$P = \text{softmax}(S) \in \mathbb{R}^{2048 \times 2048}$$Softmax计算需要： 计算每行的最大值：$m_i = \max_j S_{i,j}$ 计算指数和：$l_i = \sum_j e^{S_{i,j} - m_i}$ 归一化：$P_{i,j} = \frac{e^{S_{i,j} - m_i}}{l_i}$ 这又需要存储另一个 $2048 \times 2048$ 的矩阵。 ...