[Performance] cummax/cummin 的 CUDA kernel 硬编码 dim3(16,32)，长 innermost 维扫描比 torch.cummax 慢 19 倍、比 paddle.cumsum 慢 87 倍

[DanielSun11]

问题描述

CUDA 上的 paddle.cummax / paddle.cummin 在被扫描维度很长、而独立行数很少时性能急剧下降。最差（也是很常见）的情形是一维 tensor，或者扫描维就是最内维、外层大小为 1 的 tensor：此时 kernel 在一块 148 SM 的卡上只跑出 约 0.69 GB/s 的有效带宽，具体表现为：

• 同样 65536 个 int64 元素，比 paddle.cumsum 慢 87 倍（2.285 ms vs 0.026 ms）
• 同一张卡、同一 shape、同一 dtype 下，比 torch.cummax 慢 19 倍（2.285 ms vs 0.120 ms）
• 耗时严格与扫描长度成正比 O(N)，而 cumsum 和 torch.cummax 基本是平的（并行度正常）

根因看起来是 ScanWithIndicesKernel 里 innermost 分支硬编码的 dim3 threads(16, 32) —— 它完全不随 tensor shape 变化。PyTorch 的 kernel 结构几乎一样，但会根据 (num_rows, row_size) 推导 block 形状，因此快 19 倍。详细分析和修复建议见下。

这不是 micro-benchmark 里的边角问题：我们在长上下文 LLM 预训练里踩到了它 —— 对一个 [seqlen]（seqlen = 65536）的 document 边界 tensor 调 paddle.cummax，单次 2.3 ms、单个训练 iteration 累计 17 ms，并且直接体现为流水线气泡。

环境

• paddle：3.4.0.post20260808（当前 develop 的 kernel 源码也是同样的实现）
• GPU：NVIDIA B30Z，SM 10.3，148 SMs，Driver API 13.0，Runtime API 13.2
• OS：Linux 6.8.0，Python 3.12
• 对比基线：torch 2.12.0+cu132，同一台机器、同一张卡

复现代码

repro_cummax.py（只依赖 paddle，可直接跑）

import paddle

paddle.set_device("gpu:0")


def gpu_ms(fn, *a, warmup=10, iters=20, trials=10):
    for _ in range(warmup):
        fn(*a)
    paddle.device.synchronize()
    best = float("inf")
    for _ in range(trials):
        s = paddle.device.cuda.Event(enable_timing=True)
        e = paddle.device.cuda.Event(enable_timing=True)
        paddle.device.synchronize()
        s.record()
        for _ in range(iters):
            fn(*a)
        e.record()
        paddle.device.synchronize()
        best = min(best, s.elapsed_time(e) / iters)
    return best


cummax = lambda x, ax: paddle.cummax(x, axis=ax)
cumsum = lambda x, ax: paddle.cumsum(x, axis=ax)

print("A) 最内维扫描，只有 1 行：[1, N], axis=-1, int64")
print(f"{'N':>9} {'cummax':>11} {'cumsum':>11} {'ratio':>7} {'cummax GB/s':>12}")
for n in (16384, 32768, 65536, 131072, 262144):
    x = paddle.arange(n, dtype="int64").reshape([1, n])
    tm, ts = gpu_ms(cummax, x, -1), gpu_ms(cumsum, x, -1)
    print(f"{n:>9} {tm:>9.3f} ms {ts:>9.3f} ms {tm / ts:>6.0f}x "
          f"{n * 8 * 3 / (tm * 1e-3) / 1e9:>11.2f}")

print("\nB) 同样 65536 个 int64 元素，只改行数，axis=-1")
for outer in (1, 2, 4, 8, 16, 32, 64, 256, 1024):
    inner = 65536 // outer
    x = paddle.arange(65536, dtype="int64").reshape([outer, inner])
    tm, ts = gpu_ms(cummax, x, -1), gpu_ms(cumsum, x, -1)
    print(f"{f'[{outer},{inner}]':>15} {tm:>9.3f} ms {ts:>9.3f} ms {tm / ts:>6.0f}x")

print("\nC) 非最内维，65536 个 int64 元素，axis=0")
for outer, inner in ((65536, 1), (32768, 2), (8192, 8), (1024, 64), (128, 512)):
    x = paddle.arange(outer * inner, dtype="int64").reshape([outer, inner])
    print(f"{f'[{outer},{inner}]':>15} num_irows={inner:>5} "
          f"{gpu_ms(cummax, x, 0):>9.3f} ms {gpu_ms(cumsum, x, 0):>9.3f} ms")

测量数据

以下均为 CUDA event 计时、10 组 x 20 次取 min，未特别说明的均为 int64。

A）单行，扫描长度扫描 —— [1, N]，axis=-1

N	paddle.cummax	paddle.cumsum	倍数	cummax 有效带宽
16384	0.574 ms	0.024 ms	24x	0.69 GB/s
32768	1.145 ms	0.025 ms	45x	0.69 GB/s
65536	2.285 ms	0.026 ms	87x	0.69 GB/s
131072	4.564 ms	0.027 ms	167x	0.69 GB/s
262144	9.128 ms	0.027 ms	339x	0.69 GB/s

完全线性：约 34.9 ns / 元素，与 N 无关。按"读 8 字节 + 写 8 字节 value + 写 8 字节 index"计算，有效带宽被钉死在 0.69 GB/s —— 而这张卡的 HBM 带宽在 1 TB/s 以上。

B）同样 65536 个元素、同样的数据，只改行数 —— axis=-1

shape	paddle.cummax	paddle.cumsum	倍数
[1, 65536]	2.285 ms	0.026 ms	88x
[2, 32768]	1.147 ms	0.046 ms	25x
[4, 16384]	0.585 ms	0.025 ms	23x
[8, 8192]	0.312 ms	0.015 ms	21x
[16, 4096]	0.181 ms	0.009 ms	19x
[32, 2048]	0.118 ms	0.009 ms	13x
[64, 1024]	0.061 ms	0.009 ms	7x
[256, 256]	0.018 ms	0.009 ms	2x
[1024, 64]	0.008 ms	0.009 ms	1x

耗时基本等于 2.285 ms / num_rows，直到 num_rows 达到几百才饱和。这张表里每一行的总计算量完全相同，只是 shape 不同，却有 285 倍的差距。

C）paddle.cummax vs torch.cummax（同卡、同 dtype、同 shape）

shape	axis	paddle.cummax	torch.cummax	paddle / torch
[1, 16384]	-1	0.574 ms	0.033 ms	17x
[1, 32768]	-1	1.145 ms	0.062 ms	18x
[1, 65536]	-1	2.285 ms	0.120 ms	19x
[1, 131072]	-1	4.564 ms	0.236 ms	19x
[1, 262144]	-1	9.128 ms	0.470 ms	19x
[4, 16384]	-1	0.585 ms	0.033 ms	18x
[16, 4096]	-1	0.181 ms	0.017 ms	11x
[64, 1024]	-1	0.061 ms	0.015 ms	4x
[1024, 64]	-1	0.008 ms	0.008 ms	1x
[65536, 1]	0	4.284 ms	4.184 ms	1.0x（两边都慢）

D）一维 tensor，axis=0 —— 实际代码踩到的就是这个

N	cummax	cummin	cumsum
16384	0.574 ms	0.581 ms	0.024 ms
32768	1.145 ms	1.159 ms	0.025 ms
65536	2.285 ms	2.314 ms	0.026 ms
131072	4.564 ms	4.623 ms	0.027 ms

cummin 表现完全一致（走同一个 ScanWithIndicesKernel）。

E）非最内维是另一个独立问题 —— axis=0，65536 个元素

shape	num_irows	paddle.cummax	paddle.cumsum
[65536, 1]	1	4.284 ms	0.025 ms
[32768, 2]	2	3.190 ms	0.058 ms
[8192, 8]	8	1.342 ms	0.023 ms
[1024, 64]	64	0.180 ms	0.014 ms
[128, 512]	512	0.027 ms	0.014 ms

F）dtype —— [1, 65536]，axis=-1

dtype	cummax
int32	2.255 ms
int64	2.285 ms
float32	1.982 ms
float64	3.619 ms

（float16 / bfloat16 会抛 RuntimeError，属于不支持，另说。）

原因分析

1. innermost 分支硬编码 block 形状，因此完全忽略 row_size。

cum_maxmin_kernel.cu#L270-L283：

if (axis == out_dims.size() - 1) {
  int64_t row_size = x.dims()[ndim - 1];
  int64_t num_rows = x.numel() / row_size;

  dim3 threads(16, 32);                       // <-- 固定值，与 row_size 无关
  dim3 grid(std::min(maxGridDim,
      ceil(float(num_rows) / float(threads.y))));   // <-- 只取决于 num_rows

  KernelScanInnerWithIndices<T1, T2, 16, 32><<<grid, threads, ...>>>(...);
}

当 num_rows == 1 时：

• grid = ceil(1 / 32) = 1 —— 148 个 SM 的卡上只起 1 个 block
• num_threads_x = 16，所以 #L117-L118 的分块循环每次只前进 2 * 16 = 32 个元素 —— 65536 / 32 = 2048 次串行迭代，每次都要做一遍 Blelloch up-sweep + down-sweep、约 12 个 __syncthreads() —— 也就是单个 block 内串行执行约 2.4 万次 block 级 barrier
• block 里 512 个线程中，只有 threadIdx.y == 0 的 16 个线程拿到了合法的 row，其余 496 个被 if (row < num_rows) 屏蔽掉，但依然要参与每一次 __syncthreads()

也就是说，实际只有约 0.02% 的算力在干活，这与实测的 0.69 GB/s 和严格 O(N) 的伸缩完全吻合。

2. PyTorch 的 kernel 结构几乎相同，但 block 形状是根据 tensor shape 推导出来的。

ScanUtils.cuh#L18-L40（v2.12.0）：

// 让 x:y 线程比例接近 row_size:num_rows 的比例，block 总线程数保持在 512 左右
integer diff = log_num_threads_x - log_num_threads_y;
log_num_threads_x = ((integer)9 + diff) / (integer)2;          // 9 == log2(512)
log_num_threads_x = std::min(std::max((integer)4, log_num_threads_x), (integer)9);

对 num_rows = 1, row_size = 65536：log_x = 16、log_y = 0、diff = 16，于是 log_num_threads_x = min(max(4, 12), 9) = 9 —— num_threads_x = 512、num_threads_y = 1。分块循环因此每次前进 2 * 512 = 1024 个元素，64 次迭代而不是 2048 次（少 32 倍），而且 512 个线程全部有效。实测差距 19 倍，与"barrier 轮数少 32 倍、但每轮开销略高"是自洽的。

Paddle 目前做不到这一点，是因为 shared memory 是静态、模板固定的数组（#L100-L101）：

__shared__ T1 vbuf[num_threads_y][2 * num_threads_x];
__shared__ T2 ibuf[num_threads_y][2 * num_threads_x];

而 PyTorch 传的是 2 * num_threads * (sizeof(scalar_t) + sizeof(int64_t)) 的动态 shared memory，这正是它能自由调整 num_threads_x 的前提。

3. 非最内维分支在 num_irows 很小时退化成单线程。

#L299-L305：

dim3 threads(std::min(512, static_cast<int>(num_irows)));
dim3 grid(std::min(maxGridDim, num_orows), ...);

对 [65536, 1]、axis=0：num_irows = 1 → threads = 1，num_orows = 1 → grid = (1, 1)，即一个线程串行走 65536 步（KernelScanOuterWithIndices 的循环）。实测 4.28 ms。PyTorch 这里的 launch 配置一模一样，也一样慢（4.18 ms），所以这条算是上游共有的短板，而不是 Paddle 的退化 —— 但仍然值得修，因为 [N, 1]、或者 unsqueeze 之后按 axis=0 扫描这类 shape 很常见。

修复建议

按改动量从小到大：

1. innermost 分支改成自适应 block 形状 —— 移植 PyTorch 的 get_log_num_threads_x_inner_scan(num_rows, row_size)，并把 vbuf / ibuf 改成动态 shared memory。仅此一项预计就能拿回那 19 倍。算法不变，改动小、风险低。

2. 把退化 shape 路由到真正的并行 scan。 当 num_rows * ceil(row_size / (2 * num_threads_x)) 不足以填满 GPU 时，走分块两遍 scan：把长行切成 tile，先在 tile 内 cummax，再对每个 tile 的最大值做 cummax，最后把 carry 折回去。纯 Python 版本就已经比现在的 kernel 快 29 倍且逐位一致：

   part  = paddle.cummax(x.reshape([512, 128]), axis=-1)[0]
   carry = paddle.cummax(part[:, -1], axis=0)[0]
   carry = paddle.concat([paddle.full([1], LOWEST, dtype=x.dtype), carry[:-1]])
   out   = paddle.maximum(part, carry.unsqueeze(1)).reshape([-1])

   x.shape == [65536] 时 paddle.cummax(x, 0)：2.285 ms → 0.073 ms（bit-equal）。如果用户态两行 reshape 就能快 29 倍，说明 kernel 侧还有很大空间。

3. 考虑用 CUB 实现 value 部分。 paddle.cumsum 已经在用 cub::BlockScan / cub::DeviceScan（cum_kernel.cu），在这些 shape 上快 87 倍；用 cub::DeviceScan::InclusiveScan 配一个 (value, index) 的自定义 reduce op（argmax 语义、并列时取更早的下标）可以直接覆盖 cummax / cummin。

4. 修掉非最内维分支 num_irows == 1 的情况（比如把扫描轴 transpose 到最内维，或者对扫描轴分块后用两遍折叠），并重新审视 dim3 threads(std::min(512, num_irows))。

5. 补一个对 shape 敏感的性能回归用例 —— 现在这个行为在任何"[batch, seq] 且 batch 较大"的 benchmark 里都看不出来。

实际影响

在一个 ERNIE 长上下文预训练作业里（seqlen 65536），每个 attention 层都会调一次 paddle.cummax，用来把 packed document mask 前向填充成每个位置所属文档的起始位置：

doc_start_per_pos = paddle.cummax(is_boundary.cast("int64") * positions, axis=0).values

is_boundary 是 [65536] 的一维 tensor，所以每次调用都精确命中最差情形。从一次训练 iteration 的 nsys 采集看：KernelScanInnerWithIndices 出现 7 次、累计 17.2 ms，且在关键路径上 —— 紧跟其后的 paddle.nonzero 需要 D2H 同步，host 会一直阻塞等它，于是直接转化成流水线气泡。

我们的临时规避是把这次调用改写成 cumsum + gather（2.79 ms → 0.30 ms，9.3 倍，逐位一致）。但只要是在长序列上构造边界表 / 段表，就很容易踩到同一个坑，而 cummax 恰恰是这种场景下最自然的 API。

如果需要，我可以在这台机器上验证你们的修复补丁。