翻译(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。

Image 1: GPU vs CPU Processing Flow
Image 2: GPU vs CPU Processing Flow

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。

Image 3: GPU Schedule Manager
Image 4: GPU Schedule Manager

理解这些调度细节对于优化运行时性能至关重要。例如,选择合适的 block 大小可以通过平衡并行性与资源利用来显著提升性能。


Quiz

线程 block 大小在 CUDA 内核性能中起到了什么作用?


线程块(Thread Blocks)和网格(Grids)

我们已经看到线程如何被组织成 warp(例如,每组 32 个线程),继而被组合进 block,再由 block 被调度到 Streaming Multiprocessor (SM) 上。现在让我们拉远视角看整个计算布局。

  • 线程块是一组可以通过共享内存和同步来协同工作的线程,并且它们在同一个 SM 上执行。可把 block 看成是一个团队,它们共同处理数据的一部分。
  • 网格由所有启动的 block 组成,代表了所有团队共同解决整个问题。在图示中,可以看到网格如何划分成多个 block,每个 block 又细分为多个 warp,每个 warp 包含多个线程。

Image 5: CUDA Thread Blocks Organization
Image 6: CUDA Thread Blocks Organization

使用 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);
}

这个示例展示了几个关键概念:

  1. 内核使用 blockIdx.x、blockDim.x 和 threadIdx.x 为每个线程计算出唯一的全局索引
  2. 由于启动的线程数可能会舍入到下一个 block 大小,所以需要检查数组边界
  3. 主机代码使用向上取整计算所需的 block 数目
  4. 每个线程正好处理一个元素,演示了线程与数据元素之间一一对应的关系

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++ 中,代码通常按以下顺序操作:

  1. 在 GPU 上分配内存(使用 cudaMalloc)
  2. 将数据从主机(CPU)拷贝到设备(GPU)(使用带有 Host-to-Device 标志的 cudaMemcpy)
  3. 启动内核对那些数据进行计算(数据会保持在 GPU 上,直到所有内核执行完毕)
  4. 将结果从设备拷贝回主机(使用带有 Device-to-Host 标志的 cudaMemcpy)
  5. 释放 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?


以上就是完整的中文翻译。

posted @ 2025-02-21 17:11  ffl  阅读(266)  评论(0)    收藏  举报