vLLM 分页注意力

vLLM 分页注意力#

  • 目前，vLLM 使用了其自身的多头查询注意力内核实现 (csrc/attention/attention_kernels.cu)。此内核旨在与 vLLM 的分页 KV 缓存兼容，其中键和值缓存存储在单独的块中（请注意，此块概念与 GPU 线程块不同。因此，在后续文档中，我将 vLLM 分页注意力块称为“块”，而将 GPU 线程块称为“线程块”）。

  • 为了实现高性能，此内核依赖于专门设计的内存布局和访问方法，尤其是在线程从全局内存读取数据到共享内存时。本文档旨在逐步提供内核实现的高级解释，以帮助那些希望了解 vLLM 多头查询注意力内核的人。阅读完本文档后，用户可能会对实际实现有更好的理解，并更容易理解实际实现。

  • 请注意，本文档可能未涵盖所有细节，例如如何计算对应数据的正确索引或点乘实现。但是，在阅读本文档并熟悉高级逻辑流程后，您应该更容易阅读实际代码并理解细节。

输入#

  • 内核函数接受一系列参数，供当前线程执行其分配的工作。三个最重要的参数是输入指针 q、k_cache 和 v_cache，它们指向全局内存上需要读取和处理的查询、键和值数据。输出指针 out 指向全局内存，结果应写入其中。这四个指针实际上指的是多维数组，但每个线程仅访问分配给它的数据部分。为了简单起见，我在此省略了所有其他运行时参数。

    template<
typename scalar_t,
int HEAD_SIZE,
int BLOCK_SIZE,
int NUM_THREADS,
int PARTITION_SIZE = 0>
__device__ void paged_attention_kernel(
... // Other side args.
const scalar_t* __restrict__ out,       // [num_seqs, num_heads, max_num_partitions, head_size]
const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
const scalar_t* __restrict__ k_cache,   // [num_blocks, num_kv_heads, head_size/x, block_size, x]
const scalar_t* __restrict__ v_cache,   // [num_blocks, num_kv_heads, head_size, block_size]
... // Other side args.
)

  • 函数签名上方还有一系列模板参数，这些参数在编译时确定。scalar_t 表示查询、键和值数据元素的数据类型，例如 FP16。HEAD_SIZE 表示每个头的元素数量。BLOCK_SIZE 指的是每个块中的令牌数量。NUM_THREADS 表示每个线程块中的线程数。PARTITION_SIZE 表示张量并行 GPU 的数量（为简单起见，我们假设为 0 且禁用张量并行）。

  • 有了这些参数，我们需要执行一系列准备工作。这包括计算当前头索引、块索引和其他必要的变量。但是，就目前而言，我们可以忽略这些准备工作，直接进行实际计算。一旦我们掌握了整个流程，就更容易理解它们了。

概念#

  • 在我们深入研究计算流程之前，我想描述一些后续章节需要的概念。但是，如果您遇到任何令人困惑的术语，您可以跳过本节并在稍后返回。

  • 序列：序列表示客户端请求。例如，q 指向的数据的形状为 [num_seqs, num_heads, head_size]。这表示总共有 num_seqs 个查询序列数据由 q 指向。由于此内核是单查询注意力内核，因此每个序列只有一个查询令牌。因此，num_seqs 等于批处理中处理的令牌总数。

  • 上下文：上下文由序列生成的令牌组成。例如，["What", "is", "your"] 是上下文令牌，输入查询令牌是 "name"。模型可能会生成令牌 "?"。

  • Vec：vec 是同时获取和计算的元素列表。对于查询和键数据，vec 大小 (VEC_SIZE) 的确定使得每个线程组可以一次获取和计算 16 字节的数据。对于值数据，vec 大小 (V_VEC_SIZE) 的确定使得每个线程可以一次获取和计算 16 字节的数据。例如，如果 scalar_t 是 FP16（2 字节）且 THREAD_GROUP_SIZE 为 2，则 VEC_SIZE 将为 4，而 V_VEC_SIZE 将为 8。

  • 线程组：线程组是一小组线程 (THREAD_GROUP_SIZE)，它们一次获取和计算一个查询令牌和一个键令牌。每个线程仅处理部分令牌数据。一个线程组处理的元素总数称为 x。例如，如果线程组包含 2 个线程且头大小为 8，则线程 0 处理索引为 0、2、4、6 的查询和键元素，而线程 1 处理索引为 1、3、5、7 的元素。

  • 块：vLLM 中的键和值缓存数据被拆分为块。每个块在一个头中存储固定数量 (BLOCK_SIZE) 的令牌的数据。每个块可能仅包含整个上下文令牌的一部分。例如，如果块大小为 16，头大小为 128，则对于一个头，一个块可以存储 16 * 128 = 2048 个元素。

  • Warp：Warp 是由 32 个线程 (WARP_SIZE) 组成的小组，这些线程在流式多处理器 (SM) 上同时执行。在此内核中，每个 Warp 一次处理一个查询令牌和一个完整块的键令牌之间的计算（它可能会在多次迭代中处理多个块）。例如，如果一个上下文有 4 个 Warp 和 6 个块，则分配方式类似于 Warp 0 处理第 0 个和第 4 个块，Warp 1 处理第 1 个和第 5 个块，Warp 2 处理第 2 个块，Warp 3 处理第 3 个块。

  • 线程块：线程块是一组线程 (NUM_THREADS)，它们可以访问相同的共享内存。每个线程块包含多个 Warp (NUM_WARPS)，在此内核中，每个线程块处理一个查询令牌和一个完整上下文的键令牌之间的计算。

  • 网格：网格是线程块的集合，并定义了集合的形状。在此内核中，形状为 (num_heads, num_seqs, max_num_partitions)。因此，每个线程块仅处理一个头、一个序列和一个分区的计算。

查询#

  • 本节将介绍查询数据如何在内存中存储以及每个线程如何获取。如上所述，每个线程组获取一个查询令牌数据，而每个线程本身仅处理一个查询令牌数据的一部分。在每个 Warp 中，每个线程组都将获取相同的查询令牌数据，但会将其与不同的键令牌数据相乘。

    const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
    query

    #

  • 每个线程都定义了自己的 q_ptr，它指向全局内存上分配的查询令牌数据。例如，如果 VEC_SIZE 为 4 且 HEAD_SIZE 为 128，则 q_ptr 指向的数据总共包含 128 个元素，分为 128 / 4 = 32 个 vec。

    q_vecs

    #

    __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];

  • 接下来，我们需要将 q_ptr 指向的全局内存数据读取到共享内存中作为 q_vecs。重要的是要注意，每个 vec 都分配给不同的行。例如，如果 THREAD_GROUP_SIZE 为 2，则线程 0 将处理第 0 行 vec，而线程 1 将处理第 1 行 vec。通过以这种方式读取查询数据，相邻线程（如线程 0 和线程 1）可以读取相邻内存，从而实现内存合并以提高性能。

键#

  • 与“查询”部分类似，本节介绍键的内存布局和分配。虽然每个线程组在一个内核运行中仅处理一个查询令牌，但它可能会在多次迭代中处理多个键令牌。同时，每个 Warp 将在多次迭代中处理多个键令牌块，确保在内核运行后整个线程组处理所有上下文令牌。在此上下文中，“处理”是指执行查询数据和键数据之间的点乘。

    const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride
                    + kv_head_idx * kv_head_stride
                    + physical_block_offset * x;

  • 与 q_ptr 不同，每个线程中的 k_ptr 将在不同的迭代中指向不同的键令牌。如上所示，k_ptr 指向键令牌数据，该数据基于分配的块、分配的头和分配的令牌处的 k_cache。

    key

    #

  • 上图说明了键数据的内存布局。它假设 BLOCK_SIZE 为 16，HEAD_SIZE 为 128，x 为 8，THREAD_GROUP_SIZE 为 2，并且总共有 4 个 Warp。每个矩形代表一个头中一个键令牌的所有元素，这些元素将由一个线程组处理。左半部分显示了 Warp 0 的键令牌数据的总共 16 个块，而右半部分代表了其他 Warp 或迭代的剩余键令牌数据。在每个矩形内部，总共有 32 个 vec（一个令牌 128 个元素），这些 vec 将由 2 个线程（一个线程组）分别处理。

    k_vecs

    #

    K_vec k_vecs[NUM_VECS_PER_THREAD]

  • 接下来，我们需要从 k_ptr 读取键令牌数据，并将它们存储在寄存器内存中作为 k_vecs。我们对 k_vecs 使用寄存器内存，因为它只会被一个线程访问一次，而 q_vecs 将被多个线程多次访问。每个 k_vecs 将包含多个向量，以供后续计算使用。每个 vec 将在每次内部迭代中设置。vec 的分配允许 Warp 中的相邻线程一起读取相邻内存，这再次促进了内存合并。例如，线程 0 将读取 vec 0，而线程 1 将读取 vec 1。在下一个内部循环中，线程 0 将读取 vec 2，而线程 1 将读取 vec 3，依此类推。

  • 您可能仍然对整体流程感到有些困惑。不用担心，请继续阅读下一节“QK”。它将以更清晰和更高级的方式说明查询和键的计算流程。

QK#

  • 如下面的伪代码所示，在整个 for 循环块之前，我们获取一个令牌的查询数据并将其存储在 q_vecs 中。然后，在外部 for 循环中，我们迭代指向不同令牌的不同 k_ptrs，并在内部 for 循环中准备 k_vecs。最后，我们执行 q_vecs 和每个 k_vecs 之间的点乘。

    q_vecs = ...
for ... {
   k_ptr = ...
   for ... {
      k_vecs[i] = ...
   }
   ...
   float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(q_vecs[thread_group_offset], k_vecs);
}

  • 如前所述，对于每个线程，它一次只获取部分查询和键令牌数据。但是，在 Qk_dot<>::dot 中会发生跨线程组缩减。因此，此处返回的 qk 不仅仅是部分查询和键令牌点乘之间的结果，而实际上是整个查询和键令牌数据之间的完整结果。

  • 例如，如果 HEAD_SIZE 的值为 128 且 THREAD_GROUP_SIZE 为 2，则每个线程的 k_vecs 将总共包含 64 个元素。但是，返回的 qk 实际上是 128 个查询元素和 128 个键元素之间点乘的结果。如果您想了解有关点乘和缩减的更多细节，您可以参考 Qk_dot<>::dot 的实现。但是，为了简单起见，我不会在本文档中介绍它。

Softmax#

  • 接下来，我们需要计算所有 qk 的归一化 softmax，如上所示，其中每个 $x$ 代表一个 qk。为此，我们必须获得所有 qk 的缩减值 qk_max($m(x)$) 和 exp_sum($\ell(x)$)。缩减应在整个线程块中执行，涵盖查询令牌与所有上下文键令牌之间的结果。

    \begin{gather*} m(x):=\max _i \quad x_i \\ \quad f(x):=\left[\begin{array}{lll}e^{x_1-m(x)} & \ldots & e^{x_B-m(x)}\end{array}\right]\\ \quad \ell(x):=\sum_i f(x)_i \\ \quad \operatorname{softmax}(x):=\frac{f(x)}{\ell(x)} \end{gather*}

qk_max 和 logits#

  • 就在我们获得 qk 结果之后，我们可以使用 qk 设置临时的 logits 结果（最后，logits 应存储归一化的 softmax 结果）。我们还可以比较和收集当前线程组计算的所有 qk 的 qk_max。

    if (thread_group_offset == 0) {
   const bool mask = token_idx >= context_len;
   logits[token_idx - start_token_idx] = mask ? 0.f : qk;
   qk_max = mask ? qk_max : fmaxf(qk_max, qk);
}

  • 请注意，此处的 logits 位于共享内存上，因此每个线程组都会为其自身分配的上下文令牌设置字段。总体而言，logits 的大小应为上下文令牌的数量。

    for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
}

if (lane == 0) {
   red_smem[warp_idx] = qk_max;
}

  • 然后，我们需要获取每个 Warp 的缩减 qk_max。主要思想是使 Warp 中的线程相互通信，并获得最终的最大 qk。

    for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
}
qk_max = VLLM_SHFL_SYNC(qk_max, 0);

  • 最后，我们可以通过比较此线程块中所有 Warp 的 qk_max 来获得来自整个线程块的缩减 qk_max。然后，我们需要将最终结果广播到每个线程。

exp_sum#

  • 与 qk_max 类似，我们也需要从整个线程块中获取缩减的总和值。

    for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
    float val = __expf(logits[i] - qk_max);
    logits[i] = val;
    exp_sum += val;
}
...
exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);

  • 首先，将每个线程组的所有 exp 值求和，同时，将 logits 的每个条目从 qk 转换为 exp(qk - qk_max)。请注意，此处的 qk_max 已经是整个线程块中的最大 qk。然后，我们可以像 qk_max 一样对整个线程块进行 exp_sum 的缩减。

    const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);
for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
   logits[i] *= inv_sum;
}

  • 最后，使用缩减的 qk_max 和 exp_sum，我们可以获得最终的归一化 softmax 结果作为 logits。此 logits 变量将在后续步骤中用于与值数据进行点乘。现在，它应存储所有分配的上下文令牌的 qk 的归一化 softmax 结果。

值#

value

#

logits_vec

#

v_vec

#

  • 现在我们需要检索值数据并执行与 logits 的点乘。与查询和键不同，值数据没有线程组概念。如图所示，与键令牌内存布局不同，同一列的元素对应于相同的值令牌。对于一个值数据块，有 HEAD_SIZE 行和 BLOCK_SIZE 列，这些行列被拆分为多个 v_vecs。

  • 每个线程始终一次从相同 V_VEC_SIZE 的令牌中获取 V_VEC_SIZE 个元素。因此，单个线程通过多次内部迭代从不同行和相同列中检索多个 v_vec。对于每个 v_vec，它需要与相应的 logits_vec 进行点乘，logits_vec 也是来自 logits 的 V_VEC_SIZE 个元素。总而言之，通过多次内部迭代，每个 Warp 将处理一个值令牌块。通过多次外部迭代，整个上下文值令牌将被处理

    float accs[NUM_ROWS_PER_THREAD];
for ... { // Iteration over different blocks.
    logits_vec = ...
    for ... { // Iteration over different rows.
        v_vec = ...
        ...
        accs[i] += dot(logits_vec, v_vec);
    }
}

  • 如上面的伪代码所示，在外部循环中，类似于 k_ptr，logits_vec 迭代不同的块并从 logits 读取 V_VEC_SIZE 个元素。在内部循环中，每个线程从与 v_vec 相同的令牌中读取 V_VEC_SIZE 个元素，并执行点乘。重要的是要注意，在每次内部迭代中，线程都会为相同的令牌获取不同的头位置元素。然后将点乘结果累加到 accs 中。因此，accs 的每个条目都映射到分配给当前线程的头位置。

  • 例如，如果 BLOCK_SIZE 为 16 且 V_VEC_SIZE 为 8，则每个线程一次获取 8 个令牌的 8 个值元素。每个元素来自同一头位置的不同令牌。如果 HEAD_SIZE 为 128 且 WARP_SIZE 为 32，则对于每个内部循环，一个 Warp 需要获取 WARP_SIZE * V_VEC_SIZE = 256 个元素。这意味着 Warp 需要总共 128 * 16 / 256 = 8 次内部迭代来处理整个值令牌块。并且每个线程中的每个 accs 都包含 8 个元素，这些元素是在 8 个不同头位置累积的。对于线程 0，accs 变量将具有 8 个元素，这些元素是从所有分配的 8 个令牌累积的值头的第 0、32、...、224 个元素。

LV#

  • 现在，我们需要在每个 Warp 内对 accs 执行缩减。此过程允许每个线程累积一个块中所有令牌的分配头位置的 accs。

    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
   float acc = accs[i];
   for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
      acc += VLLM_SHFL_XOR_SYNC(acc, mask);
   }
   accs[i] = acc;
}

  • 接下来，我们跨所有 Warp 对 accs 执行缩减，从而允许每个线程拥有所有上下文令牌的分配头位置的 accs 累积。请注意，每个线程中的每个 accs 仅存储所有上下文令牌的整个头的部分元素的累积。但是，总而言之，输出的所有结果都已计算出来，但只是存储在不同的线程寄存器内存中。

    float* out_smem = reinterpret_cast<float*>(shared_mem);
for (int i = NUM_WARPS; i > 1; i /= 2) {
    // Upper warps write to shared memory.
    ...
        float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];
        for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
                ...
        dst[row_idx] = accs[i];
    }

    // Lower warps update the output.
        const float* src = &out_smem[warp_idx * HEAD_SIZE];
    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
                ...
        accs[i] += src[row_idx];
    }

        // Write out the accs.
}

输出#

  • 现在，我们可以将所有计算出的结果从本地寄存器内存写入到最终输出全局内存。

    scalar_t* out_ptr = out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE
                + head_idx * max_num_partitions * HEAD_SIZE
                + partition_idx * HEAD_SIZE;

  • 首先，我们需要定义 out_ptr 变量，该变量指向分配的序列和分配的头的起始地址。

    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
    from_float(*(out_ptr + row_idx), accs[i]);
}
}

  • 最后，我们需要迭代不同的分配头位置，并根据 out_ptr 写出相应的累积结果。