翻译(by o3-mini):CUDA 编程入门:面向 Python 开发者
原文:https://www.pyspur.dev/blog/introduction_cuda_programming
下面是文章的中文翻译:
CUDA 编程入门:面向 Python 开发者
发布日期:2025-02-03
注意:目标 URL 返回错误 500:内部服务器错误
GPU 与 CPU 的基本概念
GPU 是专为同时处理大量运算而设计的大规模并行处理器,内含数千个核心。用软件开发的类比来说,可以将 CPU 看作擅长复杂、顺序任务的单线程应用——一次只能处理有限的操作;而 GPU 则类似于拥有数千个线程的多线程应用,每个线程同时处理任务中的一小部分。CPU 适用于处理精细、逐步的过程,而 GPU 强项在于利用并行性快速处理海量数据,例如深度学习中常用的矩阵运算。
在 GPU 编程中,我们会显式地利用这种并行计算能力。例如,一块现代消费级 GPU(如 NVIDIA RTX 4090)包含 16,384 个 CUDA 核心(更小、专用的计算单元),而一台高端 CPU 只有 16 到 24 个通用核心。虽然单个 GPU 核心处理速度不及 CPU 核心,但大量核心协同工作能使 GPU 同时执行海量计算,非常适合深度学习中需要的矩阵操作。
NVIDIA 的 CUDA(Compute Unified Device Architecture)既是一个平台,也是 C++ 的扩展,它允许我们编写在 GPU 上运行的程序。CUDA 提供了编程模型和 API,使开发者能够直接在 GPU 上编写代码,将可并行化的工作从 CPU 卸载到 GPU,从而实现显著的性能提升。其基本思路是将问题拆分成许多可以同时解决的小块(就像给每个 GPU 核心分配一个小任务)。
如果你曾经从事机器学习工作,或许已经使用过 PyTorch、JAX 或 TensorFlow。这些框架屏蔽了许多 GPU 编程的复杂性:你只需分配张量,然后调用诸如 tensor.cuda()
或 tensor.to(device)
的操作,而在底层,框架自动处理 CUDA 的执行。它们依靠高度优化的 CUDA 库(如用于矩阵乘法的 cuBLAS 及用于深度学习基本操作的 cuDNN),同时保持用户友好的 Python 风格接口。
例如,在 PyTorch 中实现向量相加就非常简单:
import torch
# 在 GPU 上创建两个大向量
a = torch.rand(1000000, device='cuda')
b = torch.rand(1000000, device='cuda')
# 对应元素相加
c = a + b
在这个例子中,a + b
看似只是一个单一的数学运算,但实际上,在底层 PyTorch 会启动一个 GPU 内核(即在 GPU 上运行的函数),该内核并行地对 a
和 b
的对应元素进行相加。PyTorch 抽象了显式的 GPU 内存管理和线程调度,允许开发者专注于高层次的操作。
这种抽象机制非常强大,但了解更底层的 CUDA 工作原理有助于针对特定任务优化性能。接下来,我们将探讨 CUDA 内核(kernel)的工作原理以及它们如何将计算映射到 GPU 的硬件上。
Quiz
如果将 CPU 比作数据库、GPU 比作 MapReduce 工作者,那么哪种操作理论上在 GPU 上的表现会比在 CPU 上差?
CUDA 内核与线程模型
在 CUDA 中,内核(kernel)是你编写并运行在 GPU 上的函数。当你启动一个内核时,并不是调用一个普通函数,而是同时生成数百甚至数千个并行线程,这些线程在不同数据上同时执行同一个函数。这通常被称为单指令多线程模型(SIMT)——“一条程序,多份数据”。
举个例子,假设我们要将两个数组 A 和 B 中的数字逐元素相加得到输出数组 C。使用 CUDA C++ 我们可以编写如下内核函数:
__global__ void vecAddKernel(float *A, float *B, float *C, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
C[idx] = A[idx] + B[idx];
}
}
这里先不去关心 blockIdx
和 blockDim
,本质上这段计算为每个线程在数组中生成了一个唯一的索引 idx
,并让该线程处理 A 和 B 中相应的元素相加。__global__
限定符表示这是一个运行在设备(GPU)上的 CUDA 内核,它可以从主机(CPU)代码中调用。if (idx < n)
的判断是为了防止线程数目略大于数组长度时,多出来的线程什么也不做。
当我们从主 CPU 程序(即“主机代码”)中启动一个内核时,需要指定线程数目。内核本身是在 GPU 上运行,但需要 CPU 代码来配置和启动内核。例如:
int N = 1000000;
int threadsPerBlock = 256;
int numberOfBlocks = (N + threadsPerBlock - 1) / threadsPerBlock;
vecAddKernel<<< numberOfBlocks, threadsPerBlock >>>(d_A, d_B, d_C, N); // 启动配置
这段 CPU 代码配置并启动了 GPU 内核,指定总共需要 numberOfBlocks * threadsPerBlock
个线程,每个 block 中有 threadsPerBlock
个线程,总共生成 numberOfBlocks
个 block。
CUDA 将线程组织为 warps(通常为 32 个线程一组并一起执行),而 warps 又被组合进 block 中。每个 block 都在一个 Streaming Multiprocessor (SM) 上运行,而 SM 拥有有限的资源,如寄存器和共享内存。block 的大小会影响这些资源如何分配以及同时能运行多少个 warp(即所谓的 occupancy)。
当 warp 中的线程遇到 if 语句时,如果部分线程走一条分支而其他线程走另一条,则执行会序列化。硬件使用掩码位来追踪哪些线程应执行哪条分支,虽然能保证正确性,但可能会影响性能。下面是一个展示 warp 分歧的例子:
__global__ void divergentKernel(float *data, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
// 这个条件会导致同一 warp 内部的线程走不同路径,从而产生分歧
if (data[idx] > 0.5f) {
data[idx] *= 2.0f; // 部分线程执行此操作
} else {
data[idx] += 1.0f; // 其他线程执行此操作
}
}
}
// 考虑 SM 资源的启动配置
int maxThreadsPerSM = 1024; // 举例的资源限制
int registersPerThread = 32;
int sharedMemoryPerBlock = 1024; // 单位:字节
// 选择合适的 block 大小以在满足硬件限制的前提下最大化 occupancy
int threadsPerBlock = 256; // 必须为 warp 大小(32)的倍数
int numberOfBlocks = (N + threadsPerBlock - 1) / threadsPerBlock;
divergentKernel<<<numberOfBlocks, threadsPerBlock>>>(d_data, N);
GPU 的调度器负责将 block 分配到可用的 SM 上。如果 block 数目超过 SM 数量,多余的 block 将进入队列,待资源空闲后被调度执行。这种调度依赖于每个 SM 可用的共享内存和寄存器的大小,从而限制了同时能运行多少个 block。
理解这些调度细节对于优化运行时性能至关重要。例如,选择合适的 block 大小可以通过平衡并行性与资源利用来显著提升性能。
Quiz
线程 block 大小在 CUDA 内核性能中起到了什么作用?
线程块(Thread Blocks)和网格(Grids)
我们已经看到线程如何被组织成 warp(例如,每组 32 个线程),继而被组合进 block,再由 block 被调度到 Streaming Multiprocessor (SM) 上。现在让我们拉远视角看整个计算布局。
- 线程块是一组可以通过共享内存和同步来协同工作的线程,并且它们在同一个 SM 上执行。可把 block 看成是一个团队,它们共同处理数据的一部分。
- 网格由所有启动的 block 组成,代表了所有团队共同解决整个问题。在图示中,可以看到网格如何划分成多个 block,每个 block 又细分为多个 warp,每个 warp 包含多个线程。
使用 block 的原因之一是出于硬件的实际限制:单个 block 中可启动的线程数有限(当前 GPU 通常最多 1024 个线程)。如果问题需要比这更多的线程,就必须将其拆分到多个 block 中。例如,假设你需要对两个包含 2048 个元素的数组进行相加,可以使用 2 个 block,每个 block 各包含 1024 个线程 —— block 0 处理索引 0–1023,而 block 1 处理索引 1024–2047。一般来说,如果你有 N 个元素且单个 block 最多可以包含 B 个线程,则你需要启动 ceil(N/B) 个 block 来确保所有元素都能被处理。在之前的向量加法示例中,我们用 256 个线程的 block 启动了 3907 个 block 来处理 1,000,000 个元素,而这 3907 个 block 组成的集合就是网格。
另一个使用 block 的原因是它提供了扩展性和调度灵活性。例如,假设你的 GPU 有 20 个 SM,每个 SM 可以同时运行几个 block(取决于资源的可用性)。如果你启动了 100 个 block,而同时只有 20 个 block 可以运行(每个 SM 一个),GPU 就会并行启动 20 个 block,并且一旦有某个 block 执行完毕,就会在空闲的 SM 上调度下一个 block。从程序员的角度看,所有 100 个 block 共同完成了计算(结果就好像 100 × blockSize 个线程同时运行一样),而 GPU 负责将这些 block 分配到其硬件资源上。这意味着即使问题启动的线程数目超过了 GPU 实际能同时执行的数量,运行时也会通过时间分片来依次执行这些 block。另外,block 也为在多个 GPU 之间分配任务或者限制并行处理的工作量(从而解决诸如共享内存、寄存器等资源约束)提供了一种自然的方式。
同一 block 内的线程具有特殊能力,使它们在协作时更加强大。它们可以共享一块快速的片上内存,并通过同步屏障进行协调,从而高效地完成需要线程间通信的任务。但这种协作仅限于同一 block 内的线程 —— 不同 block 间必须使用较慢的全局内存进行通信,且不能直接同步。这一设计选择使得 block 能够灵活地在 SM 之间调度,同时让 CUDA 程序能适配具有不同 SM 数量的 GPU 架构。
在选择每个 block 的线程数时,开发者通常会使用 2 的幂次(如 128、256 或 512),以便与 warp 大小(32 个线程)以及硬件特性对齐。这一选择会影响资源利用率、内存访问模式以及 GPU 通过线程调度隐藏内存延迟的能力。
下面是一个简单的示例,展示了 block 与网格如何协同处理大数组:
__global__ void processLargeArray(float* input, float* output, int N) {
// 根据 block 和线程索引计算全局线程索引
int idx = blockIdx.x * blockDim.x + threadIdx.x;
// 确保不会越界访问
if (idx < N) {
// 每个线程处理一个元素
// 本例中,仅将每个元素乘以 2
output[idx] = input[idx] * 2.0f;
}
}
// 主机代码用于启动内核
void launchProcessing(float* d_input, float* d_output, int N) {
// 选择每个 block 的线程数(2 的幂次,小于等于 1024)
const int threadsPerBlock = 256;
// 计算处理 N 个元素需要多少 block(向上取整)
int numBlocks = (N + threadsPerBlock - 1) / threadsPerBlock;
// 以计算得到的网格维度启动内核
processLargeArray<<<numBlocks, threadsPerBlock>>>(d_input, d_output, N);
}
这个示例展示了几个关键概念:
- 内核使用
blockIdx.x
、blockDim.x
和threadIdx.x
为每个线程计算出唯一的全局索引 - 由于启动的线程数可能会舍入到下一个 block 大小,所以需要检查数组边界
- 主机代码使用向上取整计算所需的 block 数目
- 每个线程正好处理一个元素,演示了线程与数据元素之间一一对应的关系
Quiz
为什么不能仅启动一个包含 10,000 个线程的 block 来处理 10,000 个任务?
CUDA 中的内存管理
CUDA 编程不仅涉及在 GPU 上运行代码,还需要管理 CPU(主机)与 GPU(设备)之间的数据传输。CPU 和 GPU 拥有各自独立的内存空间——CPU 使用计算机的 RAM,而 GPU 使用板载显存(VRAM)。你不能在 CPU 上直接访问 GPU 内存,反之亦然——必须通过 PCIe(或 NVLink)总线显式地拷贝数据。这一点与 PyTorch 的使用场景有很大不同,在 PyTorch 中只需一行代码便可将张量移动到 GPU 上,而通常这也是你唯一需要关注的部分。
在原生 CUDA C/C++ 中,代码通常按以下顺序操作:
- 在 GPU 上分配内存(使用
cudaMalloc
) - 将数据从主机(CPU)拷贝到设备(GPU)(使用带有 Host-to-Device 标志的
cudaMemcpy
) - 启动内核对那些数据进行计算(数据会保持在 GPU 上,直到所有内核执行完毕)
- 将结果从设备拷贝回主机(使用带有 Device-to-Host 标志的
cudaMemcpy
) - 释放 GPU 内存(使用
cudaFree
)
例如,若要使用之前的 vecAddKernel
内核,可以按如下方式编写代码:
int N = 1000000;
size_t bytes = N * sizeof(float);
// 分配主机内存并初始化
float *h_A = (float*)malloc(bytes);
float *h_B = (float*)malloc(bytes);
float *h_C = (float*)malloc(bytes);
// ... 初始化 h_A 和 h_B 数据 ...
// 分配设备内存
float *d_A, *d_B, *d_C;
cudaMalloc(&d_A, bytes);
cudaMalloc(&d_B, bytes);
cudaMalloc(&d_C, bytes);
// 将输入数组从主机拷贝到设备
cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice);
cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice);
// 启动内核(例如,继续采用每个 block 256 个线程)
int threads = 256;
int blocks = (N + threads - 1) / threads;
vecAddKernel<<<blocks, threads>>>(d_A, d_B, d_C, N);
// 将结果数组拷贝回主机
cudaMemcpy(h_C, d_C, bytes, cudaMemcpyDeviceToHost);
// 释放设备内存
cudaFree(d_A);
cudaFree(d_B);
cudaFree(d_C);
Quiz
在处理 GPU 数据时,CUDA 内存操作的正确顺序是什么?
共享内存与同步
与 Python 的自动内存管理不同,CUDA 提供了一个快速的片上内存——共享内存,所有同一 block 内的线程都可以访问。可以将共享内存看作是团队会议室中的白板,各成员(线程)可以快速协作完成任务——这类似于在 Python 多进程中使用共享数据结构,但性能上有所不同。
CUDA 中的同步通过 __syncthreads()
(或者在 Numba 中使用 cuda.syncthreads()
)实现,它充当一个屏障,确保同一 block 内所有线程都达到某个同步点后再继续执行,从而防止出现类似于 Python 多线程中因竞争条件引起的错误。
下面是一个使用共享内存和同步的示例:
__global__ void incrementElements(float *data, int n) {
__shared__ float tile[256]; // 声明共享内存数组
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int tid = threadIdx.x;
if (idx < n) {
// 将全局内存中的元素加载到共享内存
tile[tid] = data[idx];
__syncthreads(); // 确保所有加载完成
// 每个线程对共享内存中的元素加 1
tile[tid] += 1.0f;
__syncthreads(); // 确保所有线程完成更新
// 将计算结果写回全局内存
data[idx] = tile[tid];
}
}
Quiz
在 CUDA 编程中,使用共享内存的主要目的是什么?
面向大语言模型(LLMs)的融合 CUDA 内核
现代大语言模型(LLM)的工作负载推动了自定义 CUDA 内核的发展,目的是将多个操作融合在一起,从而降低内存开销并提高效率。其中一个著名的例子是 FlashAttention,它通过减少内存读写实现了 Transformer 自注意力机制的优化。
传统的 Transformer 自注意力计算公式为:
Attention(Q, K, V) = softmax(QK^T / √dₖ) ⋅ V
其中 Q、K 和 V 分别表示查询(query)、键(key)和值(value)矩阵,而 dₖ 表示键向量的维度。
FlashAttention 通过在 GPU 的共享内存中对计算进行分块处理,从而减少从全局内存中读写数据的次数。这在处理长序列时可以显著提高效率。
通过了解 CUDA 内核底层的工作原理,开发者可以编写定制的融合内核,从而超越高层 PyTorch 实现的性能表现。这在深度学习中尤为重要,因为内存带宽往往是性能瓶颈。
FlashAttention:PyTorch 与 CUDA 实现对比
PyTorch 实现
下面展示的是在 PyTorch 中如何使用分块(tiled)计算方式实现 FlashAttention,以降低内存需求:
import torch, math
def flash_attention_pytorch(Q, K, V, block_size=16):
"""
使用内存高效的分块操作计算注意力分数。
参数:
Q: 查询矩阵 [N_out x d]
K: 键矩阵 [N_inp x d]
V: 值矩阵 [N_inp x d]
block_size: 分块大小,用于分块计算
返回:
O: 输出矩阵 [N_out x d]
"""
N_out, d = Q.shape
N_inp = K.shape[0]
# 初始化输出张量
O = torch.zeros(N_out, d, device=Q.device)
L = torch.zeros(N_out, 1, device=Q.device)
# 计算所需块数
T_c = (N_inp + block_size - 1) // block_size # 向上取整
T_r = (N_out + block_size - 1) // block_size
scale_factor = 1 / math.sqrt(d)
# 对 Q 和 O 分成 T_r 块;对 K 和 V 分成 T_c 块进行处理
for i in range(T_r):
# 获取当前块的查询
q_start = i * block_size
q_end = min((i + 1) * block_size, N_out)
Q_block = Q[q_start:q_end]
# 初始化块内累加器
O_block = torch.zeros(q_end - q_start, d, device=Q.device)
L_block = torch.zeros(q_end - q_start, 1, device=Q.device)
m_block = torch.full((q_end - q_start, 1), float('-inf'), device=Q.device)
last_m = m_block.clone()
# 对 K、V 分块处理
for j in range(T_c):
k_start = j * block_size
k_end = min((j + 1) * block_size, N_inp)
K_block = K[k_start:k_end]
V_block = V[k_start:k_end]
# 计算当前块的注意力分数
S_block = scale_factor * (Q_block @ K_block.T) # [B_r x B_c]
# 为数值稳定性更新运行最大值
m_block = torch.maximum(m_block, S_block.max(dim=-1, keepdim=True).values)
# 使用数值稳定性计算注意力权重
P_block = torch.exp(S_block - m_block) # [B_r x B_c]
# 利用更新的最大值计算缩放因子,更新累加器
scaling_factor = torch.exp(last_m - m_block)
L_block = scaling_factor * L_block + P_block.sum(dim=-1, keepdim=True)
O_block = scaling_factor * O_block + P_block @ V_block
last_m = m_block.clone()
# 保存当前块的结果
O[q_start:q_end] = O_block / L_block # 使用累加和进行归一化
L[q_start:q_end] = L_block
return O
PyTorch 实现总结:
- 将矩阵分块处理以节省内存
- 使用最大值技巧(max-trick)保证 softmax 的数值稳定性
- 维护运行最大值(
last_m
)以缩放之前的计算 - 逐步累加结果以处理大序列
- 对每个块单独归一化以提高精度
CUDA 实现
下面是等效的 CUDA 内核,通过融合操作实现最高效的计算:
constexpr int BLOCK_SIZE = 16;
constexpr int HIDDEN_DIM = 128;
extern "C" __global__ void flash_attention_cuda(
float* output, // [N_out x d]
float* output_lse, // [N_out]
const float* query, // [N_out x d]
const float* key, // [N_inp x d]
const float* value, // [N_inp x d]
const float scale,
const int N_out,
const int N_inp) {
// 用于块内计算的共享内存
__shared__ float q_block[BLOCK_SIZE][HIDDEN_DIM];
__shared__ float k_block[BLOCK_SIZE][HIDDEN_DIM];
__shared__ float v_block[BLOCK_SIZE][HIDDEN_DIM];
// 线程索引
const int tx = threadIdx.x;
const int ty = threadIdx.y;
const int row = blockIdx.x * BLOCK_SIZE + tx;
// 存放在寄存器中的局部累加器
float m_i = -INFINITY;
float l_i = 0.0f;
float o_i[HIDDEN_DIM] = {0.0f};
// 以块为单位处理输入数据
for (int tile = 0; tile < (N_inp + BLOCK_SIZE - 1) / BLOCK_SIZE; ++tile) {
// 加载查询块(在每个外层循环中只做一次)
if (tile == 0 && row < N_out) {
for (int d = 0; d < HIDDEN_DIM; d += blockDim.y) {
int d_idx = d + ty;
if (d_idx < HIDDEN_DIM) {
q_block[tx][d_idx] = query[row * HIDDEN_DIM + d_idx];
}
}
}
__syncthreads();
// 加载键和值块
if (tile * BLOCK_SIZE + ty < N_inp && row < N_out) {
for (int d = 0; d < HIDDEN_DIM; d += blockDim.y) {
int d_idx = d + ty;
if (d_idx < HIDDEN_DIM) {
k_block[tx][d_idx] = key[(tile * BLOCK_SIZE + tx) * HIDDEN_DIM + d_idx];
v_block[tx][d_idx] = value[(tile * BLOCK_SIZE + tx) * HIDDEN_DIM + d_idx];
}
}
}
__syncthreads();
// 计算注意力分数并更新累加器
if (row < N_out) {
float m_prev = m_i;
// 计算分数并找到数值稳定所需的最大值
float max_score = -INFINITY;
float scores[BLOCK_SIZE];
#pragma unroll
for (int j = 0; j < BLOCK_SIZE && tile * BLOCK_SIZE + j < N_inp; ++j) {
float score = 0.0f;
#pragma unroll
for (int d = 0; d < HIDDEN_DIM; ++d) {
score += q_block[tx][d] * k_block[j][d];
}
scores[j] = scale * score;
max_score = max(max_score, scores[j]);
}
// 更新运行最大值并用之前的结果进行缩放
m_i = max(m_i, max_score);
float scale_factor = exp(m_prev - m_i);
l_i *= scale_factor;
#pragma unroll
for (int d = 0; d < HIDDEN_DIM; ++d) {
o_i[d] *= scale_factor;
}
// 计算注意力并更新输出
#pragma unroll
for (int j = 0; j < BLOCK_SIZE && tile * BLOCK_SIZE + j < N_inp; ++j) {
float p_ij = exp(scores[j] - m_i);
l_i += p_ij;
#pragma unroll
for (int d = 0; d < HIDDEN_DIM; ++d) {
o_i[d] += p_ij * v_block[j][d];
}
}
}
__syncthreads();
}
// 写入最终输出
if (row < N_out) {
float inv_l = 1.0f / l_i;
for (int d = 0; d < HIDDEN_DIM; ++d) {
output[row * HIDDEN_DIM + d] = o_i[d] * inv_l;
}
output_lse[row] = l_i;
}
}
CUDA 实现总结(主要优势):
- 更快的内存访问:使用
__shared__
内存存储q_block
、k_block
和v_block
,相比于 PyTorch 中对K_block = K[k_start:k_end]
的重复切片 - 寄存器使用:将累加器(
o_i
、l_i
、m_i
)存储在快速寄存器中,而非在 PyTorch 中申请张量 - 并行处理:每个线程通过
row = blockIdx.x * BLOCK_SIZE + tx
计算处理一行数据,而非 PyTorch 中的顺序for i in range(T_r)
- 内存合并访问:利用
d += blockDim.y
结构化加载内存,相比于 PyTorch 中的矩阵切片操作 - 循环优化:使用
#pragma unroll
指令减少循环开销,并通过优化内存访问和寄存器使用提高性能,而 PyTorch 的循环解释执行速度较慢
Quiz
在 PyTorch 实现中,为什么要同时跟踪 m_block 和 last_m,而不能仅使用 m_block?
以上就是完整的中文翻译。
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· DeepSeek 开源周回顾「GitHub 热点速览」
· 物流快递公司核心技术能力-地址解析分单基础技术分享
· .NET 10首个预览版发布:重大改进与新特性概览!
· AI与.NET技术实操系列(二):开始使用ML.NET
· 单线程的Redis速度为什么快?