4. 并行模式：归约、扫描与模板计算

前一章讲“怎么把数据搬对”，这一章讲“怎么把算法摆成 GPU 的形状”。核心思想只有一句：把串行的累加改成树形，把跨线程的依赖换成每一层同时完成。代价是总工作量变大、结果不再逐位可复现。

4.1 归约：从 O(n) 步到 O(log n) 步

串行求和需要 \(n-1\) 次加法、\(n-1\) 步。树形求和的第一层让 \(n/2\) 对元素同时相加，第二层 \(n/4\) 对，依此类推，总共仍是 \(n-1\) 次加法（工作量不变），但关键路径长度从 \(n\) 降到 \(\log_2 n\)。GPU 有几千个并行单元，正好把这个深度摊开。

一个 block 内可以完全用寄存器 + shuffle 完成，不碰 shared：

reduce.cu

__inline__ __device__ float warpReduce(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, off);
    return v;                                  // lane 0 拿到整个 warp 的和
}

__global__ void blockSum(const float* x, float* out, int n) {
    __shared__ float partial[32];              // 每 block 最多 32 个 warp
    int tid = threadIdx.x;

    float v = 0.0f;
    for (int i = blockIdx.x * blockDim.x + tid; i < n; i += blockDim.x * gridDim.x)
        v += x[i];                             // 先做串行部分，尽量少用 shuffle

    v = warpReduce(v);
    if ((tid & 31) == 0) partial[tid >> 5] = v;
    __syncthreads();

    if (tid < 32) {
        int nwarps = blockDim.x >> 5;          // 这里假设 blockDim 是 32 的倍数
        v = (tid < nwarps) ? partial[tid] : 0.0f;
        v = warpReduce(v);
        if (tid == 0) atomicAdd(out, v);       // 每个 block 一次原子
    }
}

三个设计点：

• 先串行、后并行。每个线程在 grid-stride 循环里已经把多次加法做掉了，shuffle 只处理“每个线程一个值”的最后一层。反过来（每个元素都 shuffle）会浪费指令。
• 两级归约。warp 内用 shuffle，warp 之间用 shared。这是 GPU 归约的标准骨架。
• 跨 block 汇总。atomicAdd 最简单，但顺序不定、结果不可复现。要确定性就写两阶段：kernel 1 把每个 block 的部分和写进数组，kernel 2 再归约那几千个数。

固定归约树：四步。把上面说的“两阶段”写成可以照着抄的样子：

1. kernel 1 写部分和，不用原子。每个 block 把结果写到 partial[blockIdx.x]；block 数变了，partial 的长度跟着变，但每个元素怎么算不变。
2. kernel 2 用单 block 归约 partial。块数一般几千，一个 block 足够；shuffle 树和本节开头那段一样。
3. 把树的上界写死。for (off = 16; off > 0; off >>= 1) 的起点由 warp 宽度决定，不能随 blockDim 或数据量变。否则同一份数据在 8 核和 80 核上会得到不同结果——这就是“换了个机器，数值就不一样”的来源。
4. 补齐，而不是跳过。变长输入（partial 不满一轮）在尾部填 0（求和）或 -INFINITY（求 max），让每个位置都参与同一棵树。跳过空槽会让树形随数据长度改变。

• atomicAdd 汇总：不可复现（顺序由调度决定）；不占额外显存；块数任意。
• 两阶段固定树：可复现，但前提是块数与树形固定；多花“块数 × 4 B”；需要第二次 launch。
• 单 block 归约：可复现；零额外显存；但一个 block 最多 1024 线程，扩不动。

注意第二条那个前提，它比“别用原子”更重要：可复现的意思是“结果只由数据决定”，不是“两次碰巧一样”。如果你的 block 数取决于 batch 大小或序列长度，那么换个 batch 组合，树的形状就变了，数值照样会漂——这正是生产框架要专门开一个“batch 不变”模式的原因。

生产里怎么做。某开源推理框架为此提供三组开关：kernel 库的确定性算法模式、注意力内核的 batch-invariant 模式，以及关掉集合通信里的多重归约。代价是吞吐下降，README 里直接写着 slower。验收方式也值得抄：64 条序列批量跑一遍，再逐条单独跑，比较 (token, logprob, top-k ids, top-k logprobs) 四元组——只比 token 是看不出问题的，argmax 对 1e-7 的漂移不敏感。出处：DeepSeek DeepInfer main@774ce24，Apache-2.0（worker/runtime_config.py、attention_backend.py）。

4.2 扫描：每个位置都要前缀和

给定 \(x_0,\dots,x_{n-1}\)，扫描要输出

串行是 \(n\) 步。Hillis-Steele 的并行版本在每一轮让每个位置加上它前面 \(2^{d}\) 个位置的当前值：

scan.cu

// 单 block，要求 blockDim.x == n，动态 shared 传入 n 个 float
__global__ void scan(float* x, int n) {
    extern __shared__ float s[];
    int t = threadIdx.x;
    s[t] = (t < n) ? x[t] : 0.0f;
    __syncthreads();

    for (int off = 1; off < n; off <<= 1) {
        float add = (t >= off) ? s[t - off] : 0.0f;
        __syncthreads();          // 先让所有人读完上一轮
        s[t] += add;
        __syncthreads();          // 再让所有人写完，供下一轮读
    }

    if (t < n) x[t] = s[t];
}

启动：scan<<<1, n, n * sizeof(float)>>>(d_x, n)。复杂度是 \(O(n\log n)\) 次加法——比串行多做了 \(\log n\) 倍的工作，但深度只有 \(\log n\)。GPU 的并行度高，这个交换通常是划算的。Blelloch 的上行/下行两趟可以把工作量降回 \(O(n)\)，代价是常数更大、实现更绕；生产代码直接用 CUB 的 DeviceScan。

扫描是很多算法的地基：stream compaction（把满足条件的元素挤到前面）、稀疏矩阵组装、排序的基数排序阶段，都要它。

4.3 模板计算：halo 与 shared

一维卷积、差分、平滑都属于 stencil：每个输出依赖输入的一个邻域。把邻域搬进 shared，一次加载、多次复用：

blur.cu

constexpr int R = 4;                       // 半径，窗口 2R+1

__global__ void blur1d(const float* in, float* out, int n) {
    extern __shared__ float s[];           // 大小 = blockDim.x + 2R
    int base = blockIdx.x * blockDim.x;

    for (int j = threadIdx.x; j < blockDim.x + 2 * R; j += blockDim.x) {
        int g = base - R + j;              // 带 halo 的全局下标
        s[j] = (g >= 0 && g < n) ? in[g] : 0.0f;
    }
    __syncthreads();

    int g = base + threadIdx.x;
    if (g < n) {
        float acc = 0.0f;
        #pragma unroll
        for (int k = -R; k <= R; ++k) acc += s[threadIdx.x + R + k];
        out[g] = acc / (2 * R + 1);
    }
}

边界处理是这类 kernel 的主要工作量：halo 超出数组时补 0、复制边缘或跳过。二维 stencil 的复用率更高——每个 block 加载 \((T+2R)^2\) 个元素，服务 \(T^2\) 个输出，算术强度随 \(T\) 线性增长。

4.4 直方图：原子冲突与私有化

最直接的直方图是每个元素对全局数组做一次 atomicAdd。256 个 bin、几百万个元素时，热点地址会被反复争用。标准解法是“每个 block 一份 shared 副本，最后合并”：

hist.cu

__global__ void histogram(const unsigned char* data, int n, unsigned int* hist) {
    __shared__ unsigned int local[256];
    for (int i = threadIdx.x; i < 256; i += blockDim.x) local[i] = 0;
    __syncthreads();

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
        atomicAdd(&local[data[i]], 1u);        // 打在 shared 上，不抢全局
    __syncthreads();

    for (int i = threadIdx.x; i < 256; i += blockDim.x)
        if (local[i]) atomicAdd(&hist[i], local[i]);
}

shared 上的原子仍然会在“多个 lane 命中同一个 bin”时串行化，但比全局原子便宜一个数量级，而且冲突被限制在一个 block 内。如果数据分布非常集中（比如 90% 都是 0），可以再复制几份 local[K][256]，按 threadIdx.x % K 选副本，最后合并，用空间换冲突。

注意 hist 必须在 host 侧清零（cudaMemset），kernel 只做累加。

4.5 work / span：为什么 GPU 愿意“多干活”

并行算法的两个量：work \(W\) 是总操作数，span \(S\) 是关键路径长度。用 \(p\) 个执行单元时，下界是

CPU 上 \(p\) 小（几十），所以优化 \(W/p\) 更重要；GPU 上 \(p\) 大（几万），\(W/p\) 很快降到 1 以下，span 变成瓶颈。这解释了为什么 Hillis-Steele 扫描（\(W=O(n\log n)\)）在 GPU 上仍然常用：它把 span 压到 \(O(\log n)\)，而多出来的工作量被几千个 lane 摊掉了。

4.6 两阶段剪枝：不排序也能 top-k

每生成一个 token，都要在词表上做一次 top-k + top-p + 采样。以 \(V = 129{,}280\) 为例：一行的 logits 是 253 KB（bf16），512 行就是 126 MB。而你真正需要的只是前 8 个 token 和一条阈值线。

先排除两个错误答案。

• 整行排序。按 4.5 的表，排序的 work 是 \(O(V\log V)\)、span 是 \(O(\log^2 V)\)，还要把整行写回。只要前 8 个，却让 13 万个元素都参加了比赛。
• 每线程一个 bin 的直方图。词表的概率质量集中在少数 token 上，热点 bin 的原子冲突会把收益吃掉（4.4）。

两阶段。先粗筛，再精算：

129280 个 logits
   │  粗筛：每 16 个取最大值
   ▼
8080 个块分数（16 KB）
   │  选 top-8 块
   ▼
128 个候选 token
   │  精确 top-k / CDF / 采样
   ▼
8 个 token + 一条阈值线

第一步把数组缩小到 1/16，第二步只留下 128 个候选：剪枝率约 1000 倍（\(129{,}280 / 128\)）。第三步是原本就有的采样逻辑，只是现在只跑在候选上。

topk.cu

constexpr int B = 16;   // 每块 16 个 token

// 阶段 1：每块取最大值；末块给 +inf
// lg = 一行 logits，sc = 块分数
__global__ void block_max(const float* lg,
                          float* sc, int V) {
    int tid = threadIdx.x;
    int b = blockIdx.x*blockDim.x + tid;
    int lo = b*B, hi = min(lo + B, V);
    if (lo >= V) return;

    float m = -INFINITY;
    for (int i = lo; i < hi; ++i)
        m = fmaxf(m, lg[i]);
    // 末块给 +inf，保证入选
    sc[b] = (hi == V) ? INFINITY : m;
}

// 阶段 2：选 top-k 个块号（专用 topk kernel），
//     只对这 k*B 个 token 做精确 top-k / CDF
void refine(const float* lg, const int* cand,
            int k, float* p, int* id);

三个只有落到生产里才会踩到的细节：

1. 块大小取 8 或 16。块越小，候选越准、块内精算越省；但块分数越多，索引本身的字节数线性上升。\(B = 16\) 时索引是 16 KB/行，正好是 logits 的 \(1/16\)；\(B\) 再小，省下的精算可能抵不过索引的写入。
2. 最后一个有效块给 \(+\infty\)。长度可变时末块只有一半的有效元素，它几乎拿不到最高分。给末块一个哨兵值保证它入选 —— 否则“正确的 token 不在候选里”，采样会静默出错，而不是报错。
3. 行按 1024 B 对齐。下游的选择 kernel 要求行对齐，否则向量化读退化成按元素读。上游多写几列 padding，比让下游慢一个数量级划算。

省的是什么。logits 本来就要读一遍，字节数省不掉，所以这不是访存优化。省的是排序与重复扫描：全排序 \(W = O(V\log V)\)；两阶段 \(W \approx V + K + kB\)，其中 \(K\) 是块数、\(k\) 是保留的块数。后者的 span 也只有 \(\log(V/B) + \log K + \log(kB)\)，每一层都是一个块内的归约（4.1），而不是一条 \(V\) 长的串行链。

练习。设 \(V = 129{,}280\)、\(B = 16\)：算出块分数占 logits 的百分比、\(k = 8\) 时精算要扫多少个 token；再把 \(B\) 改成 64，索引字节和精算量各自往哪个方向变？

4.7 容易写错的地方