大模型加速，超全指南说明！

最新推荐文章于 2025-07-03 10:56:25 发布

原创

最新推荐文章于 2025-07-03 10:56:25 发布 · 1.1k 阅读

21 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #ai #llm #大模型

导读

最近，一位名为 Theia Vogel 的博主整理撰写了一篇长文博客，对加速 LLM 推理的方法进行了全面的总结，对各种方法展开了详细的介绍，值得 LLM 研究人员收藏查阅。

大型语言模型（LLM）以其强大的生成、理解、推理等能力而持续受到高度关注。然而，训练和部署 LLM 非常昂贵，需要大量的计算资源和内存，因此研究人员开发了许多用于加速 LLM 预训练、微调和推理的方法。

以下是博客原文内容。

之前，我使用经典的自回归采样器手动制作了一个 transformer，大致如下：

def generate(prompt: str, tokens_to_generate: int) -> str:  
    tokens = tokenize(prompt)  
    for i in range(tokens_to_generate):  
        next_token = model(tokens)  
        tokens.append(next_token)  
    return detokenize(tokens)

这种推理方法很优雅，是 LLM 工作机制的核心。自回归 LLM 在只有数千个参数的情况下运行得很好，但对于实际模型来说就太慢了。为什么会这样，我们怎样才能让它更快？

本文整理了这个问题的解决方案，从更好的硬件利用率到巧妙的解码技巧。

为什么简单推理这么慢？

使用普通的自回归生成函数进行推理速度缓慢，主要有两个原因：算法原因和硬件原因。

从算法上讲，生成过程必须在每个周期处理越来越多的 token，因为每个周期我们都会将一个新 token 附加到上下文中。这意味着要从 10 个 token prompt 生成 100 个 token，需要在 10 + 11 + 12 + 13 + … + 109 = 5950 个 token 上运行！（初始 prompt 可以并行处理，这就是为什么 prompt token 在推理 API 中通常更便宜的部分原因。）这也意味着模型在生成时会变慢，因为每个连续的 token 生成都有越来越长的前缀：

注意力（至少是普通注意力）也是一种二次算法：所有 token 都关注所有 token，导致 N^2 扩展，使一切变得更糟。

硬件原因是什么呢？很简单：LLM 规模很大。即使像 GPT-2 这样相对较小的模型也有 117M 参数，并且所有数据都必须存储在 RAM 中。RAM 确实很慢，现代处理器（CPU 和 GPU）通过在靠近处理器的地方设置大量高速缓存（cache）来弥补这一点，从而使访问速度更快。其细节根据处理器的类型和型号而有所不同，但关键是 LLM 权重不适合缓存，因此需要花费大量时间等待从 RAM 加载权重。这会产生一些不直观的效果！例如，即使激活张量（tensor）大 10 倍，对 10 个 token 进行操作也不一定比对单个 token 进行操作慢很多，因为主要的时间消耗在于移动模型权重，而不是进行计算。

指标

大模型推理速度「慢」到底是什么意思？谈到 LLM 推理，人们采用的指标有很多：

Time to First Token（TtFT）—— 收到 prompt 和返回第一个 token 之间需要多长时间？
生成延迟 —— 收到 prompt 和返回最终 token 之间需要多长时间？
吞吐量
硬件利用率 —— 我们使用硬件的计算、内存带宽和其他功能的效率如何？

不同的优化对这些指标的影响不同。例如，批处理可以提高吞吐量并更好地利用硬件，但会增加 TtFT 和生成延迟。

硬件

加速推理的一个直接方法就是购买更好的硬件（通常是某种加速器 ——GPU 或 TPU），或者更好地利用您拥有的硬件。

使用加速器可以显著提高速度，但请记住，CPU 和加速器之间存在传输瓶颈。如果模型不适合加速器的内存，则需要在整个前向传播过程中进行交换，这会大大减慢速度。这也是 Apple M1/M2/M3 芯片在推理方面表现出色的原因之一 —— 它们具有统一的 CPU 和 GPU 内存。

关于 CPU 和加速器推理，另一个关键是充分利用硬件，适当优化程序。例如，在 PyTorch 中将注意力写入 F.softmax (q @ k.T/sqrt (k.size (-1)) + mask) @ v，能提供正确的结果，但如果使用 torch.nn.function.scaled_dot_product_attention，会将计算委托给可用的 FlashAttention，这可以更好地利用缓存的手写内核产生 3 倍的加速。

编译器

torch.compile、TinyGrad 和 ONNX 等编译器可以将简单的 Python 代码融合到针对硬件优化的内核中。例如，我可以编写以下函数：

def foo(x):  
  s = torch.sin(x)  
  c = torch.cos(x)  
  return s + c

简单来说，这个函数需要：

1.x.shape () 为 s 分配的内存

2.对 x 进行线性 scan 以计算每个元素的 sin

3.x.shape () 为 c 的另一种内存分配

4.线性 scan x 以计算每个元素的 cos

5.x.shape () 为结果张量分配的内存

6.线性 scan s 和 c，将它们添加到结果中

这些步骤每一个都很慢，并且某些步骤需要跨越 Python 和本机代码之间的界限。如果我使用 torch.compile 编译这个函数会怎样？

>>> compiled_foo = torch.compile(foo, options={"trace.enabled": True, "trace.graph_diagram": True})  
>>> # call with an arbitrary value to trigger JIT  
>>> compiled_foo(torch.tensor(range(10)))  
Writing FX graph to file: .../graph_diagram.svg  
[2023-11-25 17:31:09,833] [6/0] torch._inductor.debug: [WARNING] model__24_inference_60 debug trace: /tmp/...zfa7e2jl.debug  
tensor([ 1.0000,  1.3818,  0.4932, -0.8489, -1.4104, -0.6753,  0.6808,  1.4109,  
         0.8439, -0.4990])

如果进入 debug 跟踪目录并打开其中的 output_code.py 文件，torch 就会为 CPU 生成一个优化的 C++ 内核，将 foo 融合到单个内核中。如果使用 GPU 运行此程序，torch 将为 GPU 生成 CUDA 内核。

#include "/tmp/torchinductor_user/ib/cibrnuq56cxamjj4krp4zpjvsirbmlolpbnmomodzyd46huzhdw7.h"  
extern "C" void kernel(const long* in_ptr0,  
                       float* out_ptr0)  
{  
    {  
        #pragma GCC ivdep  
        for(long i0=static_cast<long>(0L); i0<static_cast<long>(10L); i0+=static_cast<long>(1L))  
        {  
            auto tmp0 = in_ptr0[static_cast<long>(i0)];  
            auto tmp1 = static_cast<float>(tmp0);  
            auto tmp2 = std::sin(tmp1);  
            auto tmp3 = std::cos(tmp1);  
            auto tmp4 = tmp2 + tmp3;  
            out_ptr0[static_cast<long>(i0)] = tmp4;  
        }  
    }  
}

现在，步骤就变成了：
1.x.shape () 为结果张量分配的内存

2.对 x (in_ptr0) 进行线性扫描，计算 sin 和 cos 并将它们相加到结果中

对于大输入来说更简单、更快！

>>> x = torch.rand((10_000, 10_000))  
>>> %timeit foo(x)  
246 ms ± 8.89 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  
>>> %timeit compiled_foo(x)  
91.3 ms ± 14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  
# (for small inputs `compiled_foo` was actually slower--not sure why)

请注意，torch.compile 将上面的代码专门用于传入 ((10,)) 的张量的特定