从零实现极速LLM推理

图片


作者 | Andrew Chen
翻译|张雪聃、刘乾裕
OneFlow编译

本文旨在从零开始,仅使用C++和CUDA构建一个大语言模型(LLM)推理引擎,且不借助其他外部库。

为何要这样做?通过这种方式,我们能够全面了解LLM推理的整个技术栈——这一点正变得日益重要[1]——从CUDA kernel到模型架构,并且切实体会不同的优化方式如何影响推理速度。其中一个最为重要的应用场景就是,在消费级设备上针对单个输入提示快速运行。

这正是我们将要重点关注的内容:构建一个程序,使其能够加载常见开源模型的权重,并在单CPU+单GPU服务器上对这些模型进行单批次推理,然后迭代式地提升每秒生成的词元吞吐量,直至超过 llama.cpp(https://github.com/ggerganov/llama.cpp)。

阅读本文前请先对LLM、注意力机制以及Transformer有基本的了解。完整的源代码可在 GitHub 上获取:yalm(另一个语言模型,https://github.com/andrewkchan/yalm)。

(本文作者Andrew Chen是技术专家社区South Park Commons软件工程师,毕业于加州大学伯克利分校。本文由OneFlow编译发布,转载请联系授权。原文:

https://andrewkchan.dev/posts/yalm.html)


1

回顾:LLM架构与推理


让我们回顾一下LLM的工作原理,首先从其架构开始,然后再转向推理机制。这将为优化实现提供一个起点,并帮助我们建立基准。

几乎[2]每个主要的开源权重LLM都使用相同的架构[3](序列transformer模块),自GPT-2以来有一些细微的变化/创新:

  • 分组查询注意力(以及多查询注意力)

  • 基于混合专家的前馈网络

  • 基于门控线性单元(GLU)而非多层感知器(MLP)的前馈网络

  • 前馈网络的不同激活函数

  • 不同的层归一化

  • 旋转位置嵌入

因此,从不同架构加载模型本质上是定义一个可定制的transformer block类(class),然后创建一系列配置好的模块(带有适当的附加功能),并使用safetensors(https://huggingface.co/docs/safetensors/en/index)权重对它们进行初始化。本文将仅关注一种架构 ——Mistral v0.2,但如果你感兴趣,可以阅读llama.cpp如何增加对新模型的支持(https://github.com/ggerganov/llama.cpp/blob/master/docs/development/HOWTO-add-model.md)。

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; } }}

我们可以立刻看出训练和推理之间的差异。推理——至少是我们所关注的本地推理——通常是单批次的。对于提示补全以及像生成文章这样的用例,“解码阶段” 占据了执行的大部分时间,并且涉及计算之前的上下文与单个词元(或查询时间步)之间的注意力。

预填充步骤与训练更为相似,因为我们会得到一个完整的序列来进行注意力计算,这里稍后会详细介绍。在聊天机器人中,当向模型传递额外的用户消息时,还有一个 “追加(append)” 步骤,这与预填充类似,但在本文中我不会讨论这个,因为我们的实现将仅支持补全功能。

模型前向传播如下所示:


/* PSUEDOCODE */
// InferenceState is the minimum set of buffers needed to// hold state during the forward pass and exists to avoid// extra allocationsvoid 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缓存来存储每个block过往的键和值。我们将其简化,把它实现为一个简单的环形缓冲区(在文献中称为滑动窗口注意力),这足以支持在达到某个最大上下文长度内的精确注意力计算。一些精确注意力的实现(比如分页注意力,即PagedAttention)使用更复杂的KV缓存来改善诸如内存占用等方面的情况。

1.2 瓶颈与基准


现在我们来讨论一下瓶颈和基准。首先,有一个事实:在现代硬件上,推理受内存带宽的限制。欲了解更多内容,可查看Arseny Kapoulkine撰写的这篇精彩博客文章,其要点如下:

  • 每次我们生成一个词元时,都需要读取整个模型,而对每个权重仅执行少量的浮点运算。

  • 现代CPU和GPU在浮点运算方面速度极快。关键指标是每秒浮点运算次数(FLOPs)与内存带宽的比率(FLOPs / byte)。例如,AMD Ryzen 7950X的这一比率约为40∶1,而英伟达RTX 4090的比率为82∶1。我服务器的AMD EPYC 7702P的这一比率没那么亮眼,但也较为可观,为10∶1。


这就是为什么模型量化(https://huggingface.co/docs/optimum/en/concept_guides/quantization)在提高推理速度方面如此有效的原因。它不仅让硬件能够使用更快的指令(有时候确实如此),还能缩减需要通过带宽瓶颈的输入。

我们可以利用带宽来确立一个理论上的“光速”,即我们能够实现的最大词元吞吐量。在我配备了AMD EPYC 7702P和英伟达RTX 4090的机器上:

  • EPYC 7702P的最大带宽 [4]:204.8GB / 秒

  • RTX 4090的最大带宽 [5]:1008GB / 秒

  • 带有4096(4k)上下文窗口以及32位浮点数(FP32)键值缓存的Mistral - 7B - Instruct - v0.2 模型大小为 29516398592 字节

    • 204.8×10⁹字节 / 秒 ÷29516398592字节 / 词元≈6.9 词元 / 秒(针对EPYC 7702P)

    • 该模型无法装入RTX 4090的24GB显存中,所以我们跳过这部分计算。

  • 带有4096(4k)上下文窗口以及16位浮点数(FP16)键值缓存的Mistral - 7B - Instruct - v0.2模型大小为15020875776字节

    • 204.8×10⁹字节 / 秒 ÷15020875776字节 / 词元≈13.6 词元 / 秒(针对EPYC 7702P)

    • 1008×10⁹字节 / 秒 ÷15020875776字节 / 词元≈67.1 词元 / 秒(针对 RTX 4090)


请注意,实际上我们能接近理论界限的程度会因硬件不同而略有差异。幸运的是,我们有几个流行的推理引擎可供参考,以便设定更贴合实际的目标。在我的机器[6]上,使用16位浮点数(FP16)格式、带有4096(4k)上下文的Mistral - 7B - Instruct - v0.2模型时,我能够达到:

图片


2

CPU上的推理


我们从一个简单的CPU实现开始(代码:https://github.com/andrewkchan/yalm/blob/ec8c8fec911794c788c50dfe5d42ae9e1ef0e905/src/infer.cpp#L240)。这是一个直接的单线程实现,具有4k KV缓存,仅支持FP32权重,并且没有任何显式的SIMD。它实现了每秒0.6个词元的惊人吞吐量。效果如下所示:


2.1 多线程


我们可以进行的第一个优化步骤是开始在线程级别并行化代码。在OpenMP指令的帮助下,我们寻找明显可以并行的部分。我们将优化与llama2.c相同的地方,我会逐一讲解每个优化,展示改进效果。

首先,添加一行代码(https://github.com/andrewkchan/yalm/commit/2a65dcbdff106976aeff1f08a037c6b6ece5b80b)并行矩阵-向量乘法函数,使得每个线程处理输出的一行:


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个词元:


接下来,可以将我们的多头注意力计算并行化(代码:https://github.com/andrewkchan/yalm/compare/2a65dcbdff106976aeff1f08a037c6b6ece5b80b..f39081de410ced856a4ba301234ab3fc6948e4a4),使得每个线程处理一个注意力头。这是一个不那么直观的改进,但它使得短上下文生成的吞吐量达到了每秒4.4个词元,对于长上下文来说可能会更好。


2.2 权重量化与SIMD

下一个潜在的优化机会是使用SIMD。EPYC 7702P CPU支持AVX和AVX2,这使得我们能够一次处理256位的8个压缩的float32值。在我们常见的matmul函数中,我们可以尝试在内部循环中一次加载、相乘和累加8个值,这样每个线程完成其行列点积的速度将提高最多8倍!

不幸的是,通过objdump检查编译后的代码,我们发现matmul实际上已经使用了AVX指令(注意vmovups),在输入足够大的情况下执行向量化的点积。GCC太聪明了:

图片

现在谈谈量化吧。我们不会探讨所有量化格式,因为本文的目标是探索优化的广度,而我们选择的基准测试使用了固定格式。相反,我们仅将权重量化为FP16,这对于将其加载到RTX 4090的显存来说是最低要求。

一个小问题是,许多CPU不支持原生的float16数学运算。但除此之外,我们仍希望尽可能保持计算使用float32,以减少对准确性的影响,并且只要带宽仍然是瓶颈,我们应该能够在不牺牲性能的情况下这样做。

因此,我们利用这一事实:许多CPU仍然支持通过F16C x86扩展(该扩展得到广泛支持超过十年)将float16值转换为float32,以便按需加载float16权重并在计算时转换为float32。为此,我们必须明确地将matmul函数中的加载操作向量化,因为GCC无法处理半精度数组:


// F16C code technically operates on 16-bit unsigned short integerstypedef uint16_t f16_t;
// matmul supporting float16 weights via the F16C extension, which allows// conversion into float32 values before calculations.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++) { // Vectorized dot product of w[i][:] and x[:] where w is a packed float16 array. __m256 sumlo = _mm256_setzero_ps(); __m256 sumhi = _mm256_setzero_ps(); for (int j = 0; j < n; j+=16) { // Extract the next set of 16 float16 weights from `w` and store them // to two separate float32 vectors of width 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); // Extract the next two float32 vectors of width 8 `xveclo`, `xvechi` from `x` __m256 xveclo = _mm256_loadu_ps(&x[j]); __m256 xvechi = _mm256_loadu_ps(&x[j + 8]); // Compute vectorized FMAs: sumlo += wveclo * xveclo, sumhi += wvechi * xvechi sumlo = _mm256_fmadd_ps(wveclo_ps, xveclo, sumlo); sumhi = _mm256_fmadd_ps(wvechi_ps, xvechi, sumhi); } // Horizontally reduce width-8 float32 vectors sumlo, sumhi to a scalar. __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}

最终的实现(https://github.com/andrewkchan/yalm/blob/462b3705a2dc75d2ae22e3b6a4c45d9b92bc20ea/src/infer.cpp)对于短文本来说并没有在困惑度上产生任何变化。但它几乎快了一倍,达到了每秒8.2-8.4个词元:


3

GPU上的推理


在将我们的模型量化为原来的一半大小后,我们可以将其加载到RTX 4090显卡上,并开始进行GPU推理实现。记住我们的规则:仅使用原生C++/CUDA,不使用CUTLASS、cuBLAS、cuDNN或其他库。

如果你以前没有看过CUDA代码,可以参考《An Easy Introduction to CUDA C and C++(https://developer.nvidia.com/blog/easy-introduction-cuda-c-and-c/),这是一个很好的入门教程。简而言之,CUDA允许你在GPU上并行执行C++函数("kernel"),并且这些函数在一组线程的网格上执行,其中:

  • 每个线程接收与其他线程相同的函数参数,但可以创建自己的局部变量,并且被分配一个独立的threadIdx,可以用来确定它负责的任务。

  • 线程被组织成块(block),每个块有一个独立的blockIdx和固定数量的线程(blockDim)。同一块中的线程可以通过共享内存高效地合作,而共享内存的访问速度比网格中的所有线程可能访问的全局内存要快。

  • 在调用kernel时,通过三重尖括号<<< numBlocks, threadsPerBlock >>>来指定块数和每块的线程数,可以是int或dim3类型。


3.1 简单的CUDA移植


首先,我们用270行C++/CUDA代码实现了一个朴素的移植版本(https://github.com/andrewkchan/yalm/blob/ec8c8fec911794c788c50dfe5d42ae9e1ef0e905/src/infer.cpp),试图将我们的CPU操作逐一转化为CUDA kernel,并为向量加法等操作增加了额外的kernel。最终,我们得到了一个相当长的kernel列表,但更重要的是,GPU后端的主机C++代码几乎看起来和CPU后端一样,唯一的区别是函数调用附带了<<< . >>>。例如,以下是我们GPU后端的块前向传递的一部分,其中我们将激活值通过一个前馈网络处理,然后再与残差连接相加:


// mix self.w2(F.silu(self.w1(x)) * self.w3(x))// Note this is a feedforward with a GLU, not a simple 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 residual back into xadd_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 xfor (int i = 0; i < c.dim; ++i) { s.x()[i] += s.xb2()[i];}

之所以这样有效,是因为尽管CUDA kernel是异步执行的,同一个流(或默认流)上的kernel不会并发执行;一个kernel必须在前一个kernel的所有线程完成后才能开始。因此,我们只需让最终的kernel写入主机上映射的设备(device-mapped)内存,并在前向传递结束时执行deviceSync,以确保输出可供采样使用。

3.2 高效矩阵乘法


让我们深入探讨一个有趣的案例——矩阵乘法kernel。正如之前所见,这个函数在CPU上占据了大量的运行时间,并且通过OpenMP对其优化可以带来巨大的性能提升。因此,在继续之前,我们需要确保matmul的优化已经充分完成。

一个基本的矩阵乘法(matmul)实现可能是这样的,每个GPU线程负责计算结果向量的一行(或一个元素),类似于我们之前使用OpenMP并行化的CPU代码:

__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之外,这个kernel还存在内存加载合并(memory load coalescing)问题——但这些细节我们稍后再讨论。简单来说,如果我们直接使用这个kernel作为最终方案,会导致吞吐量下降到2.9 token/s——甚至比我们的CPU后端还慢!

更好的方法是为每行数据分配一个块,这样块内的线程可以高效合作。我们可以利用更多线程来加速单行计算。在这种设置中,每个块包含一个线程束,我们使用线程束步长循环来计算行的总和。最后,通过线程束求和归约,我们可以合并所有线程的结果。

__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);
图片

我们也可以将上述kernel扩展到支持多个线程束的块(最简单的方法是保持每行1个线程束)。我将把这个问题和为什么这样做是更好的处理方式留给读者作为一个练习。:)

这个kernel的线程利用率和内存读取合并得到了显著改善,即使在每个块仅包含1个线程束的情况下,也能够达到51.7 tok/s的吞吐量,这已经是一个相当不错的初步实现结果!然而,我们仍然有很长的路要走才能赶上llama.cpp或calm的性能——还有哪些方面可以改进呢?

3.3 融合和更高效的矩阵乘法


使用nsys对生成过程进行分析显示了一些有趣的结果。首先,我们看到GPU在整个生成时间内几乎一直被使用,尽管每次前向传递结束时都有主机-设备同步。虽然在设备线程中有一些间隙,这表明偶尔会有空闲时间,但这些间隙非常短,仅为微秒级别,这意味着不是CPU受限:

图片

其次,我们可以观察到94.4%的kernel执行时间花在了矩阵乘法上:

图片

这意味着,如果我们去掉所有其他kernel,最多可以将运行时间缩短到当前时间的94.4%,这会使吞吐量从51.7 token/s提高到54.8 token/s,但仍然达不到我们的目标。因此,我们仍然需要优化矩阵乘法。

也就是说,我们可以通过融合kernel的方式来减少其总体数量。特别是,有几个kernel可以与最近的矩阵乘法kernel融合,还有几个矩阵乘法kernel可以融合在一起。例如,我们可以将matmul和add_residuals融合成一个单独的fused_matmul_add_residualskernel,后者直接将结果写入目标位置。此外,Mistral架构在某个地方执行两个矩阵乘法的门控和运算:F.silu(w1(x)) * w3(x),由于所有操作都具有相同的输出维度,并且每个线程只写入输出向量的一个元素,因此,我们可以将所有这些操作融合到一个kernel中。

这使我们能够将两个依赖的操作“融合”在一起。在第一个例子中,add_residuals操作需要等待所有matmul线程完成,如果某些matmul线程执行时间较长,这可能会导致性能瓶颈。通过融合操作,我们可以避免这种问题,并且减少每个线程从全局内存中读取和写入数据的次数。

对于好奇的读者,完整的更改在这里(https://github.com/andrewkchan/yalm/commit/7e3bb32a91117c9892363e8d2536b90cc2b1b2f3)。这使我们的吞吐量提高到54.1 token/s!

现在我们回到优化matmul话题。通过ncu进行的分析显示,matmul写入操作并没有达到最优的合并效果(coalesced)。相关的警告诊断如下:

Section: Memory Workload Analysis --------------------------- ------------ ------------ Metric Name Metric Unit Metric Value --------------------------- ------------ ------------ Memory Throughput Gbyte/second 533.22 Mem Busy % 24.82 Max Bandwidth % 90.26 L1/TEX Hit Rate % 65.94 L2 Compression Success Rate % 0 L2 Compression Ratio 0 L2 Hit Rate % 2.03 Mem Pipes Busy % 28.33 --------------------------- ------------ ------------
Section: Memory Workload Analysis Tables ... ----- -------------------------------------------------------------------------------------------------------------- WRN The memory access pattern for stores from L1TEX to L2 is not optimal. The granularity of an L1TEX request to L2 is a 128 byte cache line. That is 4 consecutive 32-byte sectors per L2 request. However, this kernel only accesses an average of 1.0 sectors out of the possible 4 sectors per cache line. Check the Source Counters section for uncoalesced stores and try to minimize how many cache lines need to be accessed per memory request. ----- -------------------------------------------------------------------------------------------------------------- ...

到这一步,定义什么是CUDA中对全局内存的合并加载或存储是有帮助的。当同一warp(线程束)中的线程对全局内存发出加载或存储请求时,如果这些请求适当对齐并且指向连续的内存区域,它们可以被组合在一起并作为单个事务执行(参见这篇优秀的Stack Overflow帖子中的图示https://stackoverflow.com/questions/5041328/in-cuda-what-is-memory-coalescing-and-how-is-it-achieved/5044424#5044424)。

这很重要,因为GPU全局内存事务只能在特定的粒度级别(32、64或128字节)上进行。因此,如果同一warp中的32个线程正在读取或写入不同的4字节浮点数,如果访问是合并的,将执行单个128字节的事务;而如果访问未合并,则可能执行32个字节的事务,浪费大量带宽。

在我们的matmul kernel中,读取操作是合并的,但写入操作不是。我们每个块有一个warp,每个warp处理一个输出元素,因此我们每个块发出一个最小为32字节的写入操作,但只能更新一个4字节的元素。增加每个块的warp数量并不能解决这个问题,因为合并是在warp级别上进行的,而不是块级别上。

相反,我们可以让每个warp计算自己的结果,然后将所有结果收集到第一个warp的各个线程中,再由第一个warp发出单个合并的写入请求。通过共享内存可以在块范围内快速完成收集。kernel:

__device__ inline float blocktranspose(float v, float def) {  // Performs block-and-warp transpose operation:  //   For a block containing K warps where lane 0 contains val_k,  //   this function returns:  //   - For warp 0, lane K: val_k  //   - For all other warps and lanes: def  int lane = threadIdx.x % warpSize;  int warp = threadIdx.x / warpSize;    // Will hold results of all warps.  // Each lane of the warp accumulates across 1 head element at a time.  // NOTE: Assumes warpSize is 32  __shared__ float sm[32];  if (lane == 0) sm[warp] = v;  __syncthreads();    return lane < blockDim.x / warpSize ? sm[lane] : def;}
template <typename T>__global__void matmul_wide(const T* A, const float* x, int n, int d, float* out) { // A (d,n) @ x (n,) -> out (d,) // PRECOND: Block is 1-D and contains WPB warps. int i = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize; if (i >= d) return; // Warp j computes sum for row at <blockIdx.x*WPB + j> // Lane 0 of each warp will hold result int k = threadIdx.x % warpSize; float rowSum = matmul_row(&A[n * i], x, k, n); // Transpose values so lane k in warp 0 contains row at <blockIdx.x*WPB + k> // For WPB=32, this allows us to coalesce 32 float32 writes into a single 128-byte store rowSum = blocktranspose(rowSum, 1.0); if (threadIdx.x < blockDim.x / warpSize) { int block_start_i = blockIdx.x * blockDim.x / warpSize; out[block_start_i + k] = rowSum; }}

ncu分析显示,在小规模矩阵维度上,性能有近10%的改善(从00微秒改进到550微秒)。在将最终LM分类器和块fused_matmul_add_residual kernel适配为使用这种改进的合并写入后,我们达到了56.1 token/s的吞吐量:


3.4 注意力与长上下文生成


长上下文生成是我们的一个主要改进要点。当生成5k个词元时,模型的平均吞吐量下降至约48 token/s,注意力kernel的运行时间占比从5%增加到超过10%。相比之下,llama.cpp和calm能够维持高50s或更高的吞吐量。问题出在哪里?

在短上下文长度的情况下,注意力操作很快,因为没有太多token需要通信——解码操作主要受其他非上下文敏感操作支配,例如,生成输入的q/k/v向量表示,并将注意力块的输出传递给前馈网络。

然而,长上下文的情况却不同。这种情况受到滑动窗口注意力的限制,但当窗口大小接近隐藏维度时,我预计注意力运行时间将接近前馈网络的运行时间。

更具体地说,让我们考虑两个注意力操作(计算注意力分数和使用它们的混合值)与单个从dim到hidden_dim的矩阵乘法(前馈网络,在现代架构中基于GLU而不是简单的MLP,应该包含3个这样的操作)进行比较。我们有...

  • dim到hidden_dim的矩阵乘法:每行有n_heads个点积,点积大小为head_dim,共有head_dim行

    • 总共有n_heads * head_dim个点积,点积大小为head_dim

  • 序列长度为seq_len的注意力机制:

    • 总共有n_heads * seq_len个点积,点积大小为head_dim

  • 序列长度为seq_len的注意力混合:

    • 总共有n_heads * seq_len个乘加运算(对大小为head_dim的向量进行运算)


在我们的基准测试中,我们使用Mistral-v0.2 7B模型,序列长度seq_len=4096,隐藏维度hidden_dim=14336,因此,我们预计FFN矩阵乘法的运行时间大约是注意力kernel运行时间的3.5倍。

我们可以看到,这与我们的attn(评分kernel)大致相符,它的运行时间约为50μs,而最终FFN矩阵乘法的运行时间约为150μs:

图片

图片

但是,att_mixkernel的运行时间却出乎意料地长,约为150μs,与FFN矩阵乘法的运行时间相同!

图片

使用ncu对att_mixkernel进行分析,发现该kernel的内存带宽利用率非常低,只有8%的内存吞吐量被利用!
att_mix(const float *, const float *, int, int, int, int, int, float *) (32, 128, 1)x(32, 1, 1), Context 1, Stream 7, Device 0, CC 8.6 Section: GPU Speed Of Light Throughput ----------------------- ------------- ------------ Metric Name Metric Unit Metric Value ----------------------- ------------- ------------ DRAM Frequency cycle/nsecond 6.27 SM Frequency cycle/nsecond 1.33 Elapsed Cycles cycle 5765045 Memory Throughput % 8.81 DRAM Throughput % 1.68 Duration msecond 4.35 L1/TEX Cache Throughput % 5.17 L2 Cache Throughput % 8.81 SM Active Cycles cycle 5685841.81 Compute (SM) Throughput % 0.47 ----------------------- ------------- ------------
WRN This kernel exhibits low compute throughput and memory bandwidth utilization relative to the peak performance of this device. Achieved compute throughput and/or memory bandwidth below 60.0% of peak typically indicate latency issues. Look at Scheduler Statistics and Warp State Statistics for potential reasons.

发生了什么?首先,让我们回顾一下这个kernel应该做什么,然后是kernel代码本身。

  • 给定两个以行主序存储的张量:

    • att - 注意力分数,形状为 (n_heads, kv_len)

    • vb - 值向量,形状为 (max_seq_len, n_kv_heads, head_dim)

  • 我们希望输出一个形状为 (n_heads, head_dim)的out张量,其中,out[q] = att[q, :] @ vb[:, q//G, :],其中 G = n_heads//n_kv_heads是分组查询注意力的组大小。

  • 下面,我们假设KV缓存中的时间步数kv_len 等于 max_seq_len:


图片

下面是我们一直用于执行此操作的简单kernel。其出发点是为每个需要计算的元素out[q, i] 启动一个块。块中的线程使用块步长循环分布和混合值vb[:, q//G, i]的范围,然后使用块求和归约(下面,我们假设每个块的大小都是一个线程束的大小,因此我们可以使用高效的线程束归约)来组合它们的结果:

__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;}
__global__void att_mix( const float* vb, // (max_seq_len, n_kv_heads, head_dim) const float* att, // (n_heads, kv_len) int head_dim, int n_heads, int n_kv_heads, int seq_len, int max_seq_len, float* out // (n_heads, head_dim)) {// PRECOND: blocks are 1-D and blockDim.x == WARP_SIZE int h = blockIdx.x; int group_size = n_heads / n_kv_heads; int g = h / group_size; int i = blockIdx.y; int offset = threadIdx.x; int kv_stride = n_kv_heads * head_dim; const float* atth = att + max_seq_len * h; const float* vh = vb + head_dim * g; float* outh = out + head_dim * h; float sum = 0.0;for (int t = offset; t < seq_len; t += WARP_SIZE) { sum += vh[kv_stride * t + i] * atth[t];} sum = warp_reduce_sum(sum);if (offset == 0) outh[i] = sum;}/* usage */dim3 tpb;tpb.x = WARP_SIZE;dim3 blocks;blocks.x = n_heads;blocks.y = head_dim;att_mix<<<blocks, tpb>>>( vb, att, head_dim, n_heads, n_kv_heads, seq_len, max_seq_len, out);

如下图所示,我们发现一个大问题 - 每个块的线程并不是访问连续的内存,因此,它们无法合并加载任何数据!

图片

这是改进合并的第一次尝试:

  • 多个块写入单个输出元素out[h, i]。

  • 将序列划分为每个块处理的连续时间块。这可以使用网格的y维度来完成。然而,我们仍然希望x维度对应于注意力头。

  • 块应该处理多个输出元素。每个块1个线程束。i由线程束ID确定。这允许线程束合并加载。

  • 线程使用atomicAdd将结果累加到out[h, i]。


用代码表示:

__global__void att_mix( const float* vb, // (max_seq_len, n_kv_heads, head_dim) const float* att, // (n_heads, kv_len) int head_dim, int n_heads, int n_kv_heads, int seq_len, int max_seq_len, float* out // (n_heads, head_dim)) {// PRECOND: blocks are 1-D and `out` has been zeroed int h = blockIdx.x; int group_size = n_heads / n_kv_heads; int g = h / group_size; int kv_stride = n_kv_heads * head_dim; const float* atth = att + max_seq_len * h; const float* vh = vb + head_dim * g; float* outh = out + head_dim * h; int t_per_thread = seq_len / gridDim.y; int t_start = blockIdx.y * t_per_thread;for (int i = threadIdx.x; i < head_dim; i += blockDim.x) { float sum = 0.0;for (int t = t_start; t < t_start + t_per_thread; t++) { sum += vh[kv_stride * t + i] * atth[t]; }atomicAdd(&outh[i], sum);}}
int max_t_per_thread = 256;
dim3 tpb;tpb.x = warp_size;dim3 blocks;blocks.x = c.n_heads;blocks.y = (kv_len + max_t_per_thread - 1) / max_t_per_thread;cudaMemset(s.xb2(), 0, c.n_heads * c.head_dim * sizeof(float));att_mix<<<blocks, tpb>>>( vb, s.att(), c.head_dim, c.n_heads, c.n_kv_heads, kv_len, c.max_seq_len, s.xb2());
绘制出来后,我们可以看到每个块的线程在内部循环的每次迭代中加载连续的vb元素:

图片

将这个kernel代入后,长短上下文生成速度都大大提高 - 我们从56 token/s和48 token/s分别提高到63 token/s和57 token/s!分析表明,和我们预期一样,att_mix现在运行时间与attn大致相同。遗憾的是,长上下文生成的主观质量和困惑度也明显下降(对于一个4000词元的文本,困惑度大约增加了5倍)。

事实证明,问题源于全局内存的atomicAdd操作,它会忽略非常小的浮点值(约1e-40),因为这些值被认为是次正常值,而全局原子操作会将其清零(https://stackoverflow.com/questions/20189225/cuda-atomicaddfloat-does-not-add-very-small-values)。根据一个2013年的论坛帖子(https://forums.developer.nvidia.com/t/atomicadd-float-does-not-add-very-small-values/31531/8?u=akc),共享内存的atomicAdd操作不应该有这个问题。因此,我们重写了kernel,以便在块共享内存中累积结果,然后在完成后写入全局内存:

__global__void att_mix( const float* vb, // (max_seq_len, n_kv_heads, head_dim) const float* att, // (n_heads, kv_len) int head_dim, int n_heads, int n_kv_heads, int seq_len, int max_seq_len, float* out // (n_heads, head_dim)) { // PRECOND: blocks are 2-D (warp_size, t_stride) int h = blockIdx.x; int group_size = n_heads / n_kv_heads; int g = h / group_size; int kv_stride = n_kv_heads * head_dim; const float* atth = att + max_seq_len * h; const float* vh = vb + head_dim * g; float* outh = out + head_dim * h; int warp_id = threadIdx.y; int t_stride = blockDim.y; // Capacity 32 since there can be at most 32 warps in a block. __shared__ float shared[32]; for (int i = threadIdx.x; i < head_dim; i += warpSize) { if (warp_id == 0) { shared[threadIdx.x] = 0; } __syncthreads(); float sum = 0.0; for (int t = warp_id; t < seq_len; t += t_stride) { sum += vh[kv_stride * t + i] * atth[t]; } atomicAdd(&shared[threadIdx.x], sum); __syncthreads(); if (warp_id == 0) { outh[i] = shared[threadIdx.x]; shared[threadIdx.x] = 0; } }}
dim3 tpb;tpb.x = warp_size;tpb.y = min(kv_len, max_threads_per_block / warp_size);dim3 blocks;blocks.x = c.n_heads;att_mix<<<blocks, tpb>>>( vb, s.att(), c.head_dim, c.n_heads, c.n_kv_heads, kv_len, c.max_seq_len, s.xb2());
这解决了我们的质量问题,同时保持了短生成的速度在~63.7 tok/s左右:


这终于使我们超越了llama.cpp的短生成吞吐量,达到61.0 token/s(见1.2)!

3.5 KV量化和编译器陷阱


最后这节讨论KV缓存量化,这是一种常见的优化技术,但我们还没有使用过。同时,也会介绍nvccCUDA编译器的一些自动优化技术,当这些优化技术故障时会发生什么,以及如何修复。早些时候,我们看到,当在CPU后端使用FP16权重时,需要手动开启自动向量化。同样,如我们将看到的那样,切换到FP16 KV缓存也需要类似的工作。

首先,介绍一下KV缓存量化:

  • 在我们之前的基准测试设置中(§1.2),llama.cpp和calm实际上使用了FP16 KV缓存条目(因为这是它们的默认设置),我们也假设相同的设置来计算速度。

  • 将KV缓存条目从FP32量化到FP16不会带来与量化权重相同的巨大收益,因为KV缓存占总内存的比例较小。但是,我们仍然期待会有一定的收益,因为每次前向传递的总内存读取量将从15.5 GB减少到15.0 GB。这种效果应该在长上下文和注意力kernel中最为明显,因为KV缓存只在注意力中使用。


仅仅将KV缓存从FP32切换到FP16(仍然使用FP32进行计算,就像权重量化一样)是简单的,只需要将float*替换为half*,并在几个关键地方插入__half2float内在函数:请参见代码更改 (https://github.com/andrewkchan/yalm/commit/19a21912f2bd2e5ff0eb59bf5f3df1756038e3e0)。不过,确保其性能是另外的事情。

然而,简单的切换并没有带来性能提升,反而使得性能下降。长上下文生成尤其受影响,吞吐量从57 tok/s下降到53.6 tok/s,att_mixkernel的运行时间从5.5%(29 μs平均)增加到13.3%(78 μs)。

图片
nsys分析:FP32 KV缓存的长上下文生成

图片nsys分析:FP16 KV缓存的长上下文生成(原始补丁)

在一个测试用例中使用ncu进行分析显示,att_mixkernel的内存吞吐量变得更糟糕,从25.7%的内存速度极限下降到6.8%,并且kernel的运行时间从140 μs增加到309 μs。为什么会这样呢?一个可能的原因是,我们现在在循环的每次迭代中发出比以前更小的内存事务。回顾一下,在att_mix中,每个线程像这样累积一个时间步t上的总和:
float sum = 0.0;for (int t = warp_id; t < seq_len; t += t_stride) { sum += vh[kv_stride * t + i] * atth[t]; }// atomicAdd `sum` when we're done

vh现在是一个半精度数组,所以从中加载的数据需要在使用前转换为 FP32。我们的原始补丁通过简单地插入一个 __half2float 函数调用来实现这一点:
float sum = 0.0;for (int t = warp_id; t < seq_len; t += t_stride) { sum += __half2float(vh[kv_stride * t + i]) * atth[t]; }// atomicAdd `sum` when we're done

这导致了吞吐量降低,因为尽管每个线程束在每个循环迭代中仍然可以合并加载,但它们合并成32个16位加载,或者64字节,这低于线程束可以发出最大合并事务的大小,即128字节。我们可以尝试通过将加载向量转化为half2加载来恢复这种性能,如下所示:

float2 sum01 = make_float2(0.0, 0.0);for (int t = warp_id; t < seq_len; t += t_stride) { float2 v01 = __half22float2(*((half2*)&vh[kv_stride * t + i])); float att_t = atth[t]; // Sadly CUDA does not have float2 SIMD ops sum01.x += v01.x * att_t; sum01.y += v01.y * att_t;}// atomicAdd both `sum01` lanes when we're done

这种尝试(完整补丁https://github.com/andrewkchan/yalm/commit/ce8efcd20a5f1c04e69ff05a334cbc3aa951eac1)使得吞吐量有一定改善,但仍然比浮点32位版本要差(188 μs vs. 140 μs,内存吞吐量 9.8% vs. 25.7%)。

发生了什么?我们的 KV 缓存加载应该与float32位版本的加载大小相同。我们最终写入的数据大小是两倍,但读取操作远远多于写入操作,读取操作也应该是合并的。为了获取更多线索,我们使用-set full 和 -o report选项收集一个更详细的ncu配置文件,这样我们就可以在NCU UI中查看该文件。

通过这种方式,我们可以看到一个非常详细的、逐行的kernel执行情况,包括相关的统计数据和每行对应的PTX和SASS指令。特别是,我们看到在浮点32位kernel中,从vb加载的那一行实际上是执行了8个32位加载:

图片

...而在FP16kernel中对应的那一行只执行了1个32位加载:

图片

深入研究,我们可以看到FP32kernel编译后的SASS指令不仅将循环展开为4次迭代(=8个加载),而且还将加载(LDG.E)和算术(FFMA)指令安排为可加载可被预取(https://developer.nvidia.com/blog/boosting-application-performance-with-gpu-memory-prefetching/):

图片

预取是一个很大的性能提升,特别是在每个迭代都从全局内存加载数据,并且每个迭代的结果与前一个迭代彼此独立。通过展开4个循环迭代并提前重新排列所有加载,编译器可以隐藏后3个迭代的加载延迟。

然而,FP16kernel的SASS代码没有展开任何循环迭代,也没有重新排列加载:

图片

那么,我们如何在FP16kernel中实现预取呢?不幸的是,仅仅在循环中添加#pragma unroll指令是不够的(至少在我的nvcc和机器设置中是这样的)。我们受编译器启发式算法的摆布;后者为我们提供了FP32kernel预取,但现在编译器却不再为FP16kernel提供预取,并且这种情况可能持续,除非我们手动展开预取值(https://developer.nvidia.com/blog/boosting-application-performance-with-gpu-memory-prefetching/)。

下面,我们将循环展开16次迭代(=32个加载),并添加了一些额外的代码来处理16个迭代的余数。我们执行批次预取[9],在每16次迭代的开始处发出一批加载请求,然后一次性检索它们,直到我们的批次耗尽。使用标量而不是数组来持有预取结果允许我们将值存储在寄存器中,而不是较慢的共享内存中:
float2 sum01 = make_float2(0.0, 0.0);constexpr int UNROLL = 16;half2 v01_0; float att_0; half2 v01_1; float att_1; half2 v01_2; float att_2; /* ... SNIP ... */half2 v01_15; float att_15;int t = warp_id;for (int ctr = 0; ctr < seq_len / t_stride; t += t_stride, ctr++) { int ctr_mod = ctr % UNROLL;if (ctr_mod == 0) {// prefetch every UNROLL iterations #define PREFETCH(j) \ v01_##j = *((half2*)&vh[kv_stride * (t + j*t_stride) + i]); \ att_##j = atth[t + j*t_stride];PREFETCH(0)PREFETCH(1)PREFETCH(2)/* ... SNIP ... */PREFETCH(15) #undef PREFETCH}// pull one value out of prefetch batch float2 v01; float att_t; switch (ctr_mod) { #define CASE(j) \ case j: v01 = __half22float2(v01_##j); att_t = att_##j; break;CASE(0)CASE(1)CASE(2)/* ... SNIP ... */CASE(15) #undef CASE}// Sadly CUDA does not have float2 SIMD ops sum01.x += v01.x * att_t; sum01.y += v01.y * att_t;}// Handle any loop remainder that can't be unrolledfor (; t < seq_len; t += t_stride) { float2 v01 = __half22float2(*((half2*)&vh[kv_stride * t + i])); float att_t = atth[t]; sum01.x += v01.x * att_t; sum01.y += v01.y * att_t;}// atomicAdd both `sum01` lanes when we're done


有了这个补丁,我们现在在长上下文生成方面与llama.cpp不相上下,分别达到58.7 tok/s和58.8 tok/s的速度,而在短上下文生成方面,我们则以63.8 tok/s的速度大幅领先于llama.cpp的61.0 tok/s。很不错!

4

下一步计划


通过这篇文章,我们已经基本实现了针对特定用例的最先进性能:仅在一个模型上实现本地、仅生成、单批次和单GPU推理。我们无需使用任何库,可以通过以下几种方式实现这一点:

  • 在CPU上使用多线程和向量化

  • 在GPU后端中使用矩阵乘法线程束归约、合并、kernel融合和更好的注意力机制

  • 对于两个后端,权重和KV缓存量化


然而,我们才刚刚开始触及LLM推理性能的表面。即使在我们的特定用例中,也还有很多可以改进的地方。

提示预填充阶段有很大改进机会,在这个阶段我们已经有了一个类似于应用于训练时的现有序列。一个常见的优化是,在这里使用矩阵-矩阵乘法代替矩阵-向量乘法,以便一次性预填充多个KV缓存条目;这将改善首个词元的生成时间,这个指标对于许多应用程序来说,与原始词元吞吐量一样重要。这种方法也构成了推测性解码生成的基础,该方法利用了预填充K个词元并不是比生成单个词元慢很多的事实。

另一个机会是融合我们的注意力kernel,这些kernel在原始实现中被分割为单独的评分、softmax和混合操作。这与FlashAttention在训练中和FlashDecoding在推理中所做的事情类似,都利用了对底层硬件的巧妙见解。

值得注意的是,我们对于FP16来说已经很快了,但还没有达到最快的速度。要在我们的机器上达到100 token/s的速度,我们需要更加激进地进行量化,使用较小的数据类型,如FP8、INT8或INT4,或者同时量化激活值[10]和KV缓存条目以及权重。或者直接使用较低精度进行计算!

在低层次上,有很多特点kernel的实现需要完成。我们最终得到的许多kernel仍表现出次优性能,其中几个kernel的内存吞吐量低于速度极限的60%。kernel优化是一个巨大的主题,我们讨论的许多优化技术——内存合并、向量化、预取等——需要以复杂的方式组合起来才能真正达到峰值性能

附录


脚注

  1. 尤其是当推理计算成为AI模型扩展的新维度时,模型越来越多地本地部署到边缘设备。

  2. 不幸的是,我无法测试calm CPU,因为我的机器不支持编译所需的扩展。

  3. 线程数量已调整。

  4. 更好的解决方案可能是使用流水线指令进行滚动预取。有关更多信息,请参阅上面链接的NVIDIA博客文章。

  5. 这通常在训练期间完成,其中我们必须保留大量激活,以便在所有层中进行梯度计算。对于单批次推理,它的好处较少,尤其是因为我们可以在层之间重用张量。


致谢

其他人都在看
图片
让超级产品开发者实现“Token自由”