老学庵

  Attention 模块是现在几乎所有大模型的核心模块，因此也有很多工作致力于提升注意力计算的性能和效果。其中MHA（Multi-Head Attention）、MQA（Multi-Query Attention）和 GQA（Grouped-Query Attention）这一路线的思路和做法被很多主流模型所采用，因此简单地梳理一些这几个变体的思路和做法，以及会涉及到的KV Cache相关内容。思路比较直白，但也有一些细节和原理值得思考。

  自Transformer架构的开山之作《Attention Is All You Need》发表以来，自注意力机制已成为深度学习模型的基石技术，尤其在自然语言处理领域展现革命性突破。鉴于该机制已渗透至各类先进模型架构，深入理解其运作原理变得至关重要。

  在Transformer架构的生成式模型（如GPT系列）中，推理过程需要逐个生成token。传统方式每次生成都需重新计算所有历史token的注意力信息，导致计算复杂度达到O(n²)。KV Cache技术通过缓存历史token的Key和Value矩阵，将后续生成的计算复杂度降至O(n)，实现推理加速，本质上是一种用空间来换取时间的加速策略。

  在自然语言处理的任务中，位置编码是帮助模型理解序列中每个单词或词片（token）在序列中的位置的一种机制。这是因为像 Transformer 这样的架构本质上是无序的，它们通过注意力机制处理整个序列中的所有元素，但并不能直接感知这些元素在序列中的顺序。因此，我们需要将序列的位置信息编码进模型的输入，编码的方式有绝对位置编码和相对位置编码。

  学习LLM过程中的经典策略记录。

  学习率这个概念在非线性优化中经常出现，在深度学习中模型在反向传播阶段严重依赖于损失函数梯度的链式传播，为了更好的控制参数更新的步长，引入了学习率的概念： \[ w_{new} = w_{old} - \eta \ \Delta J(w) \]

  卡尔曼滤波（Kalman filter）是一种高效率的递归滤波器（自回归滤波器），它能够从一系列的不完全及包含杂讯的测量中，估计动态系统的状态。卡尔曼滤波会根据各测量量在不同时间下的值，考虑各时间下的联合分布，再产生对未知变数的估计，因此会比只以单一测量量为基础的估计方式要准。卡尔曼滤波得名自主要贡献者之一的鲁道夫·卡尔曼，最早用于解决阿波罗计划的轨道预测问题。

  希尔伯特空间是一个完备的内积空间，其标准正交函数系，直观来看就是向量空间中基的延伸。其为基于任意正交系上的多项式表示的傅立叶级数和傅立叶变换提供了一种有效的表述方式，而这也是泛函分析的核心概念之一。下文中我们将通过希尔伯特空间的标准正交函数系推导周期函数和有限区间上函数的傅立叶级数表示，并进一步推出傅里叶积分来表示无穷区间的非周期函数，最后引入复数形式的傅立叶积分，引出傅立叶变换。在这一系列推导中，鉴于篇幅，主动略去了一些比较关键的部分，比如\(f(x)\)可积性及级数收敛性的讨论，有兴趣的读者可以在了解大致原理后，进行细致的理论推导以作补充。为了便于理解希尔伯特空间的概念，引用维基百科中的定义：

  CUDA的共享内存（Shared Memory）是GPU上的一种高速片上存储器，用于同一个线程块（Thread Block）内的线程共享数据。它的访问速度比全局内存（Global Memory）快得多，但访问效率可能受到 Bank Conflicts（内存库冲突）的影响。

什么是 Shared Memory 和 Bank？

• 共享内存：每个CUDA线程块有一个小的共享内存区域（通常几十KB），由该块内的所有线程共享，数据存储在Bank 中。
• Bank：共享内存被分成多个固定大小的存储单元（通常32个Bank），每个Bank可以同时服务一个线程的内存请求。
  • 在NVIDIA GPU中，每个Bank通常存储4字节（32 位）数据。
  • Bank允许一个warp（32个线程）并行访问共享内存，前提是每个线程访问不同的Bank。
  • 如果一个Warp中的两个不同的线程同时访问了Back中同一个元素，那么就会触发广播机制，即可以合并成对Shared Memory的一次访问；

什么是 Bank Conflicts？

• Bank Conflict 发生在多个线程（通常是同一个warp中的线程）试图同时访问同一个Bank中的不同数据。
• 当多个线程访问同一个Bank的不同数据时，GPU 法并行处理这些请求，必须将访问串行化，导致性能下降。
• 理想情况：一个 warp 的 32 个线程访问 32 个不同的 Bank，这样访问是并行的，效率最高。

Bank Conflicts 的简单例子

  假设共享内存有 32 个 Bank，每个 Bank 存储 4 字节，编号为 Bank 0 到 Bank 31。考虑一个 warp（32 个线程）访问共享内存的场景：

• 无 Bank Conflict：
  • 线程 0 访问 Bank 0，线程 1 访问 Bank 1，...，线程 31 访问 Bank 31。
  • 所有线程的请求可以同时处理，效率最高。
  • 示例代码：

    1 2 3

    __shared__ float data[32]; int idx = threadIdx.x; // 线程 ID (0-31) float value = data[idx]; // 每个线程访问不同的 Bank

    • 每个线程访问一个元素（4 字节），均匀分布在 32 个 Bank，无冲突。
• 有 Bank Conflict：
  • 线程 0 和线程 1 都访问 Bank 0（例如，访问 data[0] 和 data[32]，因为它们在同一个 Bank）。
  • GPU 将这两个请求串行处理，访问时间翻倍。
  • 示例代码：

    1 2 3

    __shared__ float data[64]; int idx = threadIdx.x * 32; // 每个线程访问 data[0], data[32], data[64], ... float value = data[idx]; // 多个线程访问同一个 Bank

    • 假设 data[0] 和 data[32] 在 Bank 0，线程 0 和线程 1 的访问会冲突。

Bank Conflicts 的影响

• 性能下降：Bank Conflict 会导致内存访问被串行化。例如：
  • 如果 2 个线程访问同一个 Bank，访问时间加倍（2-way Bank Conflict）。
  • 如果 4 个线程访问同一个 Bank，访问时间增加 4 倍（4-way Bank Conflict）。
• 延迟增加：串行化增加了共享内存访问的延迟，降低了整体性能。

如何避免 Bank Conflicts？

• 均匀分布访问：确保一个 warp 的 32 个线程访问不同的 Bank。例如，线程 i 访问 data[i] 或按步幅 1 访问连续内存。

• 填充（Padding）：在共享内存数组中添加填充，调整数据布局，避免多个线程访问同一 Bank。例如：

  1 2 3

  __shared__ float data[33]; // 填充 1 个元素 int idx = threadIdx.x; float value = data[idx]; // 避免冲突

• 调整步幅：避免线程访问的索引导致 Bank 重叠。例如，步幅为 32（或其倍数）可能导致所有线程访问同一个 Bank。

• 广播访问：如果多个线程读取同一地址（Bank 内的同一数据），CUDA 会通过广播机制避免冲突。

• 对齐访问：确保数据按 4 字节或 8 字节对齐，减少部分冲突。

简单比喻

  想象共享内存像一个有 32 个窗口（Bank）的银行，每个窗口一次只能服务一个客户（线程）。如果 32 个客户去不同的窗口（无冲突），服务很快完成；但如果多个客户挤到同一个窗口（Bank Conflict），就得排队等候，效率降低。

实际例子

假设一个矩阵转置操作：

1
2
3

__shared__ float tile[32][32];
int x = threadIdx.x, y = threadIdx.y;
tile[x][y] = ...; // 可能导致 Bank Conflict

1
2

__shared__ float tile[32][33]; // 填充 1 列
tile[x][y] = ...; // 减少冲突

填充（Padding）如何解决问题？

• 填充的原理：通过在共享内存数组中添加额外的空间（填充），改变数据的内存布局，使线程访问的地址分布到不同的 Bank，避免多个线程同时访问同一 Bank。
• 填充的效果：填充调整了内存对齐，使得相邻线程的访问索引不再映射到相同的 Bank，从而实现并行访问。

简单比喻

  想象共享内存的32个Bank像一个有32个窗口的银行，每个窗口一次只能服务一个客户（线程）。如果客户（线程）都去同一个窗口（Bank），就得排队（串行化，性能下降）。填充就像在窗口之间插入“占位符”，让每个客户去不同的窗口，避免排队。

9. 具体例子：矩阵转置中的 Bank Conflicts

假设一个 CUDA 核函数用于矩阵转置，共享内存定义为一个 32 × 32 的二维数组：

1
2
3
4

__shared__ float tile[32][32];
int x = threadIdx.x, y = threadIdx.y;
tile[x][y] = ...; // 写入共享内存
float value = tile[x][y]; // 读取共享内存

• 添加填充：
  • 修改数组为 32 × 33（添加一列填充）：

    1	__shared__ float tile[32][33];

  • 现在，tile[x][y] 的地址为 x + y * 33。
  • 对于 x = 0 的行，线程 y = 0 到 y = 31 访问 tile[0][0], tile[0][1], ..., tile[0][31]，地址为 0, 33, 66, ..., 990。
  • Bank 计算：0 % 32 = 0, 33 % 32 = 1, 66 % 32 = 2, ..., 990 % 32 = 30，这些地址落在 Bank 0, 1, 2, ..., 30，避免了冲突。
• 为什么有效：
  • 填充使每列的起始地址偏移（y * 33 代替 y * 32），打破了地址对 32 取模的规律，确保线程访问的地址分布到不同的 Bank。
  • 步幅（stride）从 32 变为 33（非 32 的倍数），避免了多个线程的请求集中在同一 Bank。

填充的通用规则

• 填充大小：通常在数组维度上加 1（例如 32 × 33 而不是 32 × 32），因为这足以打乱 Bank 分配，避免冲突。
• 适用场景：填充特别适用于二维数组或矩阵操作（如转置、卷积），其中线程按行或列访问可能导致冲突。
• 代价：填充会增加共享内存使用量（例如 32 × 33 比 32 × 32 多 32 个 float），需确保不超过共享内存限制（通常 48KB 或 96KB 每线程块）。

性能影响

• 无冲突：一个 warp 的 32 个线程并行访问 32 个 Bank，单周期完成。
• 有冲突：例如 32-way Bank Conflict 需要 32 个周期，性能下降 32 倍。
• 填充效果：通过填充将冲突降为 1-way（无冲突），显著提升性能。

简单例子

假设线程块为 32 × 1，访问共享内存数组：

1
2

__shared__ float data[32];
float value = data[threadIdx.x]; // 无冲突，每个线程访问不同 Bank

1
2

__shared__ float data[32];
float value = data[threadIdx.x * 32]; // 冲突，所有线程访问 Bank 0

1
2

__shared__ float data[33];
float value = data[threadIdx.x]; // 无冲突，地址分布在不同 Bank

总结

• Bank Conflicts：当多个线程同时访问共享内存的同一 Bank 时，访问被串行化，降低性能。
• 解决办法：确保线程访问不同 Bank，使用填充或调整索引。
• 影响：增加内存访问延迟，通常使性能下降 2-4 倍（视冲突程度）。
• 验证：使用 NVIDIA Nsight Compute 检查共享内存访问模式。

  NVIDIA GPU 上的内存结构从慢到快分为 Global Memory、L2 缓存、L1TEX 缓存 / Shared Memory 和寄存器，从 Volta/Turning 开始其中 L1TEX 缓存和 Shared Memory 在物理上放在了同一块芯片上，拥有相似的延时和带宽 1，因此如何掌握 Shared Memory 对性能而言就变得尤为重要。可惜的是，NVIDIA 官方的 CUDA 编程手册中介绍 Shared Memory 的教程其实只介绍了在每个线程访问一个 4 字节（32 位宽）的元素时 Bank Conflict 和广播机制，但对使用常用的向量化访存指令如 LDS.64 或 LDS.128 这种一次能访问 8 个字节（64 位宽）或 16 个字节（128 位宽）元素的情况却鲜有资料讨论。这篇文章大概就是想结合网络上的一些讨论以及通过 Microbenchmark 对这些细节来一探究竟。需要注意的是这篇文章的结论仅在 Turing 架构的 GPU 上验证过，其他架构的 GPU 可能会产生变化（欢迎评论区交流🤪）。