从零用CUDA实现高效LLM推理

在不使用任何库的情况下，将单 GPU 推理吞吐量推向极限

本文源代码已发布在 GitHub 上。

Hacker News 上的相关讨论。

本文将介绍如何完全从零开始，使用 C++ 和 CUDA 构建一个不依赖任何库的 LLM 推理引擎。

为什么要这样做？因为在这个过程中，我们可以了解 LLM 推理的完整技术栈——从 CUDA kernel 到模型架构——并切实体会不同优化手段对推理速度的影响。随着推理计算逐渐成为 AI 模型扩展的新维度，以及越来越多的模型被部署到本地边缘设备，这一点正变得愈发重要。而其中一个最重要的使用场景，就是在消费级设备上快速处理单个提示词。

这正是本文关注的重点：构建一个程序，使其能够加载常见开放模型的权重，在一台配备单 CPU 和单 GPU 的服务器上执行单批次推理，并持续迭代提升 token 吞吐量，直至超过 llama.cpp。读者应当对大语言模型、注意力机制和 Transformer 有基本了解。完整源代码已发布在 GitHub 上：yalm（Yet Another Language Model）。

calm——我的实现很大程度上受到了 Arseny Kapoulkine 推理引擎的启发。从某种意义上说，这个项目的起点就是“理解 calm，以及它为何如此之快”。不过，为了便于自己理解，我尽量让代码保持更高的可读性，同时尽可能以科学的方式理解各种优化。这意味着我放弃了 calm 中使用的一些高级技术，例如动态并行。

llama2.c——CPU 后端的部分代码来自 Andrej Karpathy 出色的 Llama 推理 C 语言实现。

接下来，让我们回顾一下 LLM 的工作原理：先从架构开始，然后进入推理机制。这将为优化实现提供一个起点，并帮助我们建立基准测试。

几乎所有主流的开放权重 LLM 都采用相同的架构——也有一些状态空间模型，例如 Mamba，声称相较于 Transformer，它们在处理长序列时更高效、扩展性也更好，但这些模型似乎并未在低功耗机器学习以及音频、视频等非离散数据领域之外取得太大成功。事实上，新发布的基础模型甚至会明确声明自己采用了“标准架构”，而其附加价值主要来自训练过程。这也引发了一些误解，例如可参见 https://blog.eleuther.ai/nyt-yi-34b-response/#how-all-llms-are-similar ——也就是由一系列顺序执行的 Transformer block 组成，只是在 GPT-2 之后出现了一些细微变化和创新：

• 分组查询注意力（以及多查询注意力）
• 基于混合专家的前馈网络
• 使用基于 GLU 的前馈网络，而不是基于 MLP 的前馈网络
• 为前馈网络采用不同的激活函数
• 不同的层归一化方式
• 旋转位置嵌入

因此，加载不同架构的模型，本质上就是定义一个可自定义的 Transformer block 类，然后创建一系列 block，为其配置正确的功能选项，并使用 safetensors 权重进行初始化。本文将只关注一种架构——Mistral v0.2——但如果你感兴趣，也可以阅读 llama.cpp 是如何添加新模型支持的。

1.1 推理概览

从宏观层面来看，推理过程大致如下方的 C++ 伪代码所示：

        
/* PSUEDOCODE */

void generate(Model& model, std::string prompt, int steps) {
  std::vector<int> encoded = tokenizer.encode(prompt);
  InferenceState s(model);
  
  // 1. Prefill step: Forward the model on each prompt token, discarding 
  // the output. This lets the model read the prompt and hydrates the KV 
  // cache.
  for (int token : encoded) {
    model.forward(s, token);
  }
  // 2. Decode step: Forward the model repeatedly, generating 1 token at a time.
  for (int i = 0; i < steps; i++) {
    model.forward(s, encoded.back());
    int next_token = sampler.sample(s.logits);
    encoded.push_back(next_token);
    std::cout << tokenizer.decode_one(next_token) << std::flush;
    if (next_token == tokenizer.EOS) {
      break;
    }
  }
}

我们可以立即看出训练与推理之间的区别。推理——至少是我们关注的本地推理——通常采用单批次。对于提示词补全和文章生成等使用场景，“解码阶段”会占据大部分执行时间，其计算过程涉及在过去的上下文与单个 token（或者说一个查询时间步）之间执行注意力计算。

预填充步骤与训练更加相似，因为此时会给出一个完整的序列，供模型在其上执行注意力计算，稍后我们还会详细讨论这一点。在聊天机器人中，向模型传入额外的用户消息时，还存在一个类似预填充的“追加”步骤。不过，本文不会讨论这一过程，因为我们的实现只支持补全。

模型的前向传递过程如下：

        
/* PSUEDOCODE */

// InferenceState is the minimum set of buffers needed to
// hold state during the forward pass and exists to avoid
// extra allocations
void Model::forward(InferenceState& s, int token) {
  // The embedding table maps token IDs to embedding vectors,
  // which are copied into a buffer of the inference state
  s.x = copy_embedding(token, this->token_embedding_table);
  
  // Models consist of a sequence of transformer blocks which
  // mutate the inference state in order
  for (Block& b : this->blocks) {
    b->block(s);
  }
  
  // Usually there is a layer norm right before the final classifier
  s.x = layernorm(s.x, this->lm_head_prenorm_weights);
  // Typically we end with a linear transform from (dim) -> (vocab_size)
  s.logits = linear(s.x, this->lm_head_classifier_weights); 
}

void Block::block(InferenceState& s) {
  s.x_resid = layernorm(s.x, this->att_prenorm_weights);
  // Multi-head attention typically includes: 
  // 1. RoPE on input (element-wise mutation w/ sines/cosines)
  // 2. QKV matmuls and updating the KV cache
  // 3. Causal self-attention, softmax, and value mixing
  // 4. Projection back into the residual stream
  s.x_resid = multi_head_attn(
    s.x_resid,
    this->wq, 
    this->wk, 
    this->wv, 
    this->key_cache,
    this->value_cache
  );
  s.x += s.x_resid;
  s.x_resid = layernorm(s.x, this->ffn_prenorm_weights);
  // On modern architectures like Llama, this is a GLU feedforward 
  // with 3 linear transforms, not a simple MLP:
  // -> w2(F.silu(w1(x)) * w3(x))
  // Some architectures also split the FFN into a mixture of experts.
  s.x_resid = ffn(s.x_resid, this->w1, this->w2, this->w3);
  s.x += s.x_resid;
}

虽然这里以典型的 C++ 风格进行了更底层的表达，但整体结构应该大致是你所熟悉的。

最值得注意的是，与训练不同，推理可以使用 KV cache 保存每个 block 过去生成的 key 和 value。为了保持简单，我们将其实现为一个普通的环形缓冲区——在相关文献中称为滑动窗口注意力——这足以支持不超过某个最大上下文长度的精确注意力。一些精确注意力实现，例如 PagedAttention，会使用更加复杂的 KV cache，以改善内存占用等方面的表现。

1.2 瓶颈与基准测试

现在可以开始讨论瓶颈和基准测试了。首先明确一个事实：在现代硬件上，推理受限于内存带宽。若想了解更多，可以阅读 Arseny Kapoulkine 的这篇优秀博文，不过其核心要点如下：

每生成一个 token，我们都需要读取整个模型，而每个权重只会参与少数几次浮点运算。

现代 CPU 和 GPU 执行浮点运算的速度极快。这里的关键指标是每秒浮点运算次数与内存带宽之比（FLOPs/byte）。例如，AMD Ryzen 7950X 的比值约为 40:1，而 RTX 4090 的比值为 82:1。我服务器上的 AMD EPYC 7702P 表现没有那么亮眼，但仍有显著的 10:1。

这就是模型量化能够如此有效地提升推理速度的原因。它不仅使硬件可以使用速度更快的指令——在某些情况下确实如此——还可以缩小需要通过带宽瓶颈的输入数据量。

我们可以利用带宽计算出理论上的“光速”，也就是能够实现的最大 token 吞吐量。在我这台配备 AMD EPYC 7702P 和 RTX 4090 的机器上：

• EPYC 7702P 最大带宽（参见官方数据表）：204.8 GB/s
• RTX 4090 最大带宽（参见 Wikipedia）：1008 GB/s

采用 4k 上下文窗口和 FP32 KV-cache 的 Mistral-7B-Instruct-v0.2-FP32 占用 29516398592 字节。

204.8e9 bytes/s / 29516398592 bytes/tok = ~6.9 tok/s for EPYC 7702P

它无法装入 RTX 4090 的 24GB 显存，因此这里跳过。

采用 4k 上下文窗口和 FP32 KV-cache 的 Mistral-7B-Instruct-v0.2-FP32 占用 29516398592 字节。

204.8e9 bytes/s / 29516398592 bytes/tok = ~6.9 tok/s for EPYC 7702P

它无法装入 RTX 4090 的 24GB 显存，因此这里跳过。

采用 4k 上下文窗口和 FP16 KV-cache 的 Mistral-7B-Instruct-v0.2-FP16 占用 15020875776 字节。

204.8e9 bytes/s / 15020875776 bytes/tok = ~13.6 tok/s for EPYC 7702P

1008e9 bytes/s / 15020875776 bytes/tok = ~67.1 tok/s for RTX 4090

配备4k上下文窗口和FP16 KV缓存的Mistral-7B-Instruct-v0.2-FP16的大小为15020875776字节

204.8e9 bytes/s / 15020875776 bytes/tok = ~13.6 tok/s（EPYC 7702P）

1008e9 bytes/s / 15020875776 bytes/tok = ~67.1 tok/s（RTX 4090）

需要注意的是，我们实际能接近理论上界的程度因硬件而异。幸运的是，我们有几个流行的推理引擎可以参考，以设定更现实的目标。在我的机器上，不幸的是我无法测试calm CPU，因为我的机器不支持编译所需的扩展。使用FP16格式的Mistral-7B-Instruct-v0.2和4k上下文，我能够达到：

我们从CPU的朴素实现开始（代码可在此处获得）。这是一个直接的单线程实现，配备4k KV缓存，仅支持FP32权重，不包含任何显式SIMD。它实现了非常快的0.6 tok/s吞吐量。看起来像这样：

第一个优化步骤是开始在线程级别并行化我们的代码。借助OpenMP pragma，我们去寻找天然可并行的机会。我们将优化与llama2.c相同的部分，我将逐个介绍每一部分以展示改进。

首先，添加单行代码并行化我们广泛使用的矩阵-向量乘法函数，使得每个线程处理输出的一行：

static void matmul(float* xout, float* x, float* w, int n, int d) {
  // W (d,n) @ x (n,) -> xout (d,)
  int i;
#pragma omp parallel for private(i)
  for (i = 0; i < d; i++) {
    float val = 0.0f;
    for (int j = 0; j < n; j++) {
      val += w[i * n + j] * x[j];
    }
    xout[i] = val;
  }
}

这是一个巨大的改进，通过调整以找到正确的线程数，我们达到了4.2 tok/s：

接下来，我们可以并行化多头注意力计算（代码在此处），使得每个线程计算一个注意力头。这不是一个立即显著的改进，但对于短上下文生成达到了4.4 tok/s，对于长上下文可能更好。

2.2 权重量化和SIMD

下一个潜在优化机会是使用SIMD。EPYC 7702P CPU支持AVX和AVX2，它让我们一次处理256位的8个打包float32值的向量。在我们的matmul函数中，我们可以尝试在内循环中一次加载、乘以并累加8个值，这将使得每个线程完成其行-列点积快速最多8倍！

不幸的是，通过objdump检查我们的编译代码显示matmul实际上已经使用AVX指令（注意vmovups）来在输入足够大的情况下执行向量化点积。看起来GCC太聪明了：

1f5:       c4 e3 7d 19 c1 01       vextractf128 xmm1,ymm0,0x1
1fb:       c5 f0 58 c0             vaddps xmm0,xmm1,xmm0
1ff:       c5 f8 12 c8             vmovhlps xmm1,xmm0,xmm0
203:       c5 f0 58 c8             vaddps xmm1,xmm1,xmm0
207:       c5 f0 c6 c1 55          vshufps xmm0,xmm1,xmm1,0x55
20c:       c5 f8 58 c1             vaddps xmm0,xmm0,xmm1

让我们转向量化。我们不会探索全部的量化格式，因为本文的目的是探索优化的广度，我们已经以固定格式选择了基准。相反，我们只是将权重量化为FP16，这是将其加载到RTX 4090 VRAM所需的最低限度。

一个小问题是许多CPU不支持原生float16数学运算。但除此之外，我们希望尽可能多地保持计算在float32中，以缓解对精度的影响，只要带宽仍然是瓶颈，我们应该能够做到而不需要权衡性能。

所以我们利用许多CPU仍然支持通过F16C x86扩展将float16值转换为float32的事实（在过去十多年来得到了很好的支持），以加载float16权重并在计算前即时将其转换为float32。除了其他事项外，这要求我们显式向量化之前matmul函数中的加载，因为GCC不知道如何处理半精度数组：

// F16C 代码在技术上操作16位无符号短整数
typedef uint16_t f16_t;

// 支持通过F16C扩展的float16权重的matmul，
// 允许在计算前转换为float32值。
static void matmul(float* xout, float* x, f16_t* w, int n, int d) {
#if defined(__AVX2__) && defined(__F16C__)
  // W (d,n) @ x (n,) -> xout (d,)
  assert(n % 16 == 0);
  int i;
#pragma omp parallel for private(i)
  for (i = 0; i < d; i++) {
    // w和x的向量化点积，其中w是打包的float16数组。
    __m256 sumlo = _mm256_setzero_ps();
    __m256 sumhi = _mm256_setzero_ps();
    for (int j = 0; j < n; j+=16) {
      // 从`w`提取下一组16个float16权重并将它们存储
      // 到两个单独的float32向量，宽度为8（`wveclo_ps`，`wvechi_ps`）
      __m256i wvec = _mm256_loadu_si256((__m256i*)&w[i * n + j]);
      __m128i wveclo = _mm256_extractf128_si256(wvec, 0);
      __m128i wvechi = _mm256_extractf128_si256(wvec, 1);
      __m256 wveclo_ps = _mm256_cvtph_ps(wveclo);
      __m256 wvechi_ps = _mm256_cvtph_ps(wvechi);
      // 从`x`提取下一个两个float32向量，宽度为8 `xveclo`，`xvechi`
      __m256 xveclo = _mm256_loadu_ps(&x[j]);
      __m256 xvechi = _mm256_loadu_ps(&x[j + 8]);
      // 计算向量化FMA：sumlo += wveclo * xveclo，sumhi += wvechi * xvechi
      sumlo = _mm256_fmadd_ps(wveclo_ps, xveclo, sumlo);
      sumhi = _mm256_fmadd_ps(wvechi_ps, xvechi, sumhi);
    }
    // 水平约简宽度8的float32向量sumlo，sumhi为标量。
    __m256 sum8 = _mm256_add_ps(sumlo, sumhi);              // sum8[0:8] = sumlo[0:8] + sumhi[0:8]
    __m128 sum4 = _mm_add_ps(                               // sum4[0:4] = sum8[0:4] + sum8[4:8]
      _mm256_extractf128_ps(sum8, 0), 
      _mm256_extractf128_ps(sum8, 1)
    );
    __m128 sum1 = _mm_dp_ps(sum4, _mm_set1_ps(1.0f), 0xf1); // sum1[0] = dot(sum4, [1,1,1,1])
    xout[i] = _mm_cvtss_f32(sum1);
  }
#else
  assert(false && "float16 not supported on this platform");
#endif
}

生成的实现对短文本没有产生任何困惑度的差异。其速度也快了近两倍，达到8.2-8.4 tok/s：

将我们的模型量化为其大小的一半后，我们现在可以将其加载到RTX 4090上并开始GPU推理实现。记住我们的规则：仅原始C++/CUDA，无CUTLASS、cuBLAS、cuDNN或其他库。

如果您之前没有见过CUDA代码，请参阅《CUDA C和C++简介》以获得一个很好的入门。从高层次来说，CUDA允许您在GPU上执行C++函数（"kernel"）平行于线程网格，其中：

每个线程接收与所有其他线程相同的函数参数，但可以创建自己的局部变量，并被分配自己的threadIdx，可以用来确定它负责什么工作。

线程还额外被组织成块，它们有自己的blockIdx和固定数量的线程（blockDim）。同一块中的线程可以通过共享内存有效地合作，共享内存比所有网格中的线程可能访问的全局内存更快。

块的数量和每块的线程数在使用三重尖括号<<<numBlocks, threadsPerBlock>>>调用kernel时指定，可以作为int或dim3给出。

3.1 到CUDA的朴素移植

第一个朴素的270行C++/CUDA实现试图将我们的CPU操作1-1翻译成kernel，加上用于向量加法等操作的额外kernel。所以我们最终得到相当长的kernel列表，但更重要的是，我们GPU后端中的主机C++代码看起来几乎像我们的CPU后端，但在函数调用上附加了<<<...>>>。例如，这是我们块前向传递的一半，其中我们在添加回残差之前将激活通过前馈网络传递：

// mix self.w2(F.silu(self.w1(x)) * self.w3(x))
// 注意这是一个具有GLU的前馈网络，而不是简单的MLP。
matmul<<<c.hidden_dim, WARP_SIZE>>>(
  w1(), s.xb(), c.dim, c.hidden_dim, s.hb()
);
matmul<<<c.hidden_dim, WARP_SIZE>>>(
  w3(), s.xb(), c.dim, c.hidden_dim, s.hb2()
);
glu_gelu<<<
  (c.hidden_dim + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK, 
  MAX_THREADS_PER_BLOCK,
>>>(
  s.hb(), s.hb2(), s.hb()
);
matmul<<<c.dim, WARP_SIZE>>>(
  w2(), s.hb(), c.hidden_dim, c.dim, s.xb2()
);
// ffn 残差加回到x
add_residuals<<<
  (c.dim + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK,
  MAX_THREADS_PER_BLOCK
>>>(
  s.x(), s.xb2(), c.dim, s.x()
);

和相应的CPU代码：

        
// mix self.w2(F.silu(self.w1(x)) * self.w3(x))
// Note this is a feedforward with a GLU, not a simple MLP.
matmul(s.hb(), s.xb(), w1<T>(), c.dim, c.hidden_dim);
matmul(s.hb2(), s.xb(), w3<T>(), c.dim, c.hidden_dim);
for (int i = 0; i < c.hidden_dim; ++i) {
  s.hb()[i] = gelu(s.hb()[i]) * s.hb2()[i];
}
matmul(s.xb2(), s.hb(), w2<T>(), c.hidden_dim, c.dim);
// residual connection back into x
for (int i = 0; i < c.dim; ++i) {
  s.x()[i] += s.xb2()[i];
}

这是可行的，因为即使 CUDA 内核是异步执行的，同一个流（或默认流）中的内核也绝不会并发执行；只有当前一个内核的所有线程都执行完毕后，下一个内核才会开始。因此，我们只需要让最后一个内核写入主机上的设备映射内存，并在前向传播结束时执行一次 deviceSync，这样输出就可用于采样了。

这里一个值得深入研究的地方是我们的 matmul 内核。前面我们已经看到，这个函数占据了 CPU 运行时间的很大一部分，而通过 OpenMP 优化它取得了巨大的性能提升。因此，在继续之前，我们要确保 matmul 已得到充分优化。

matmul 的朴素实现可能类似于经过 OpenMP 并行化的 CPU 代码：让每个 GPU 线程负责计算结果向量中的一行（一个元素）：

        
__global__
void matmul(const float* A, const float* x, int n, int d, float* out) {
  // A (d,n) @ x (n,) -> out (d,)
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= d) return;
  float sum = 0.0;
  for (int j = 0; j < n; j++) {
    sum += A[n * i + j] * x[j];
  }
  out[i] = sum;
}

/* usage */
int MAX_THREADS_PER_BLOCK = 1024;
matmul<<<
  (d + MAX_THREADS_PER_BLOCK - 1)/MAX_THREADS_PER_BLOCK, 
  MAX_THREADS_PER_BLOCK
>>>(A, x, n, d, out);

这种方法的一大问题是，它无法充分利用 CUDA 核心。Mistral-7B 的 Transformer 输入/输出维度为 4096，因此，如果我们正在计算（例如）输出之前的最后一次 matmul，就会启动 4096 个线程。但 RTX 4090 可以同时运行 16384 个线程！许多核心将处于空闲状态，我们也无法达到满额 FLOPs/s。

除了 FLOPs，这个内核还存在内存加载合并方面的问题——稍后会详细讨论。这里只需知道，如果用这个内核替换我们最终采用的内核，吞吐量将只有 2.9 tok/s，甚至比 CPU 后端还慢！

更好的思路是每一行使用一个块。块内的线程可以高效协作，因此我们能够利用更多线程，加速单行结果的计算。在这种配置中，每个块恰好包含一个 warp（warp 是块内更小的一组线程，是硬件上的基本执行单元，并拥有自己的协作原语），然后使用 warp 步长循环对整行求和。最后，可以执行 warp 求和归约，将所有线程的结果合并起来：

        
__device__ 
inline float warp_reduce_sum(float val) {
  for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
    val += __shfl_down_sync(0xffffffff, val, offset);

return val;
}

__device__
inline float matmul_row(const float* row, const float* x, int offset, int dim) {
  float sum = 0.0;
  for (int j = offset; j < dim; j += WARP_SIZE) {
    float v = row[j] * x[j];
    sum += v;
  }
  return warp_reduce_sum(sum);
}

__global__
void matmul(const float* A, const float* x, int n, int d, float* out) {
  // A (d,n) @ x (n,) -> out (d,)
  // PRECOND: Blocks are 1-D and same size as warp.
  int i = blockIdx.x;
  if (i >= d) return;
  int offset = threadIdx.x;
  float rowSum = matmul_row(&A[n * i], x, offset, n);
  if (threadIdx.x == 0) {
    out[i] = rowSum;
  }
}

/* usage */
int BLOCK_SIZE = WARP_SIZE;
matmul<<<d, BLOCK_SIZE>>>(A, x, n, d, out);

我们也可以调整上述内核，使其支持包含多个 warp 的块（最简单的做法仍是每行使用一个 warp）。至于该如何实现，以及为什么这样做可能是个好主意，就留给读者作为练习吧。:)

这个内核的线程利用率和内存读取合并效果都好得多。将其配置为每个块一个 warp 后，即便函数与内核几乎保持一一对应，我们仍然获得了 51.7 tok/s 的吞吐量——对于初步实现来说，这个性能相当不错！不过，要追上 llama.cpp 或 calm，我们仍有一段路要走——还能改进哪些地方呢？

3.3 融合与更出色的 matmul

使用 nsys 对一次生成过程进行性能分析，可以发现几个有趣的现象。首先，尽管每次前向传播结束时都会进行主机与设备同步，但在几乎整个生成期间，GPU 都处于使用状态。设备线程中确实存在一些空隙，说明它偶尔会处于空闲状态，但这些空隙只有微秒量级，这意味着我们始终没有受到 CPU 的限制：

其次，我们可以看到，94.4% 的内核执行都是矩阵乘法：

这意味着，即使去掉其他所有内核，理想情况下运行时间也只能缩短到当前的 94.4%，吞吐量将从 51.7 tok/s 提升到 54.8 tok/s，仍未达到我们的目标。因此，我们依然需要优化矩阵乘法。

尽管如此，我们可以通过将其他内核融合在一起来近似实现消除它们。具体而言，有几个内核可以融合到距离最近的 matmul 中，也有几个 matmul 可以相互融合。例如，可以将 matmul 和 add_residuals 融合成一个 fused_matmul_add_residuals，直接将结果累加到目标位置。Mistral 架构还会在某个位置对两次 matmul 的结果执行门控乘法：F.silu(w1(x)) * w3(x)——我们可以在一个内核中完成所有这些操作，因为每项操作都是由一个线程向输出向量的一个元素写入数据，并且它们的输出维度都相同。

这样可以让我们将两个相互依赖的操作“混合”在一起。以第一个例子来说，只有等到 matmul 的所有线程都完成之后，add_residuals 才能开始执行；如果 matmul 中的个别线程耗时更长，这就会成为问题。此外，我们还可以避免每个线程对全局内存进行一次额外的读取和写入。

感兴趣的读者可以在这里查看完整改动。它让我们的吞吐量提升到了 54.1 tok/s！

现在，我们可以回过头继续优化 matmul。通过 ncu 进行性能分析后可以发现，matmul 的写入并未以最优方式进行合并。相关的警告诊断信息是 s