llama.cpp 源码完全拆解（二十二）：ggml-cpu 的非 GEMM 算子——注意力、RoPE 与 RMSNorm 如何共用一条线程池节拍

三条路径分别是什么

三者的选择逻辑在源码里是这样的：

1
const bool kv_is_f32_or_f16 = (k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16);
2
const bool use_split_kv_path = !use_ref && (neq1 == 1 && neq3 == 1) &&
3
                               kv_is_f32_or_f16 && (k->type == v->type) &&
4
                               q->type == GGML_TYPE_F32 && nek1 >= 512;

先说清楚为什么 decode 需要一条单独的路径。解码时 $N=1$，按 Q 行切能得到的任务数是 $n_{\text{head}}$（通常 32），对 16 线程的机器来说并行宽度并不是问题。真正的问题是访存：逐行版里每个任务只服务一个 query head，却要把整段 KV 扫一遍；32 个任务并发时，同一份 KV 被 32 个 head 各读一次，单份 KV 的字节数（$L \times D \times 2\text{ B}$，$L=4096$、$D=128$ 时约 1 MB）超出 L1，多个 head 之间只能靠 L2/L3 复用。KV 的逻辑读取字节数被 head 数放大。

split-KV 把循环次序倒过来：每个线程先占住一段 KV，再让所有 query head 复用这一段。只要这段 KV 能留在 L1（16 线程、$L=4096$ 时每段 32 个 key，约 8 KB），一次读入就能服务 32 个 head，KV 的实际访存从「每 head 一遍」降到「整段一遍」。这是它与 GPU 上 FlashDecoding 共享的核心思路——沿序列维切开、再用 partial 合并；站内 FlashAttention 完全拆解（二）：序列维并行讨论过 GPU 版本的做法，这里看 CPU 版本如何实现。

预填充：按 Q 行切，再加一档 4 × nth

非 split-KV 路径的切分代码很简短：

1
const int64_t nr = neq1*neq2*neq3;     // 所有 query 行（含 head 与 stream 维）
2

3
const bool disable_chunking = ggml_is_numa();
4

5
// 4x chunks per thread
6
int nth_scaled = nth * 4;
7
int64_t chunk_size = (nr + nth_scaled - 1) / nth_scaled;
8
int64_t nchunk     = (nr + chunk_size - 1) / chunk_size;
9

10
if (nth == 1 || nchunk < nth || disable_chunking) {
11
    nchunk = nth;
12
}
13
...
14
const int64_t dr = (nr + nchunk - 1) / nchunk;
15

16
int current_chunk = ith;
17
while (current_chunk < nchunk) {
18
    const int64_t ir0 = dr * current_chunk;
19
    const int64_t ir1 = MIN(ir0 + dr, nr);
20
    if (use_tiled) ggml_compute_forward_flash_attn_ext_tiled(params, dst, ir0, ir1);
21
    else           ggml_compute_forward_flash_attn_ext_f16_one_chunk(params, dst, ir0, ir1, 0, nek1, nullptr, 0);
22
    current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1);
23
}

这段代码和 mul_mat 的任务派发几乎同构，值得对照着看：

1. 任务数取线程数的 4 倍。每个任务里只有一个 query 行、要扫一遍全部 KV，工作量差异很大（causal mask 下靠后的行扫得更多），4 倍是为了让动态领取有机会抹平负载差异；
2. NUMA 时退回一维静态切分。ggml_is_numa() 为真时直接把 nchunk 设成 nth，每个线程固定一段行。理由和 mul_mat 里那段 NUMA 兜底一样：任务窃取会打乱「线程只碰自己的内存」的局部性，跨节点访存代价超过负载均衡收益；
3. nchunk < nth 时退化为 nchunk = nth。这时任务数等于线程数，部分线程会拿到空区间（ir0 >= ir1），进函数后循环体一次都不执行，但依然要去 barrier 汇合。

ggml_threadpool_chunk_add 是 mul_mat 里那个原子计数器的复用：同一个线程池字段 current_chunk，同一个 barrier 重置点。也就是说，注意力和矩阵乘共享同一套动态派发机制，只是任务块的含义从「行 × 列」变成了「一组 query 行」。

tiled：把 FlashAttention 的分块算法搬到 CPU

在 x86/ARM 上，预填充的注意力走 ggml_compute_forward_flash_attn_ext_tiled，它是 FlashAttention 算法的 CPU 版本。FlashAttention 的出发点是 GPU 的显存层次：把 $N\times N$ 的注意力矩阵切块，避免中间矩阵落到 HBM 上。CPU 这边的动机略有不同——L1/L2 容量远小于 GPU 的 SRAM，但同样存在「中间矩阵一旦落到内存就变成带宽问题」的困境。

上图展示了分块注意力的核心循环：外层遍历 K/V 块，内层对每个块计算局部 softmax，再用重缩放把结果合并。数学上保证正确的关键是在线 softmax：维护已见元素的当前最大值 $M$ 与指数和 $S$，每遇到更大的值就把已累积的输出乘以 $\exp(M_{\text{old}}-M_{\text{new}})$。

CPU 实现里把这些状态都放在线程私有的工作缓冲里，布局是（tile 尺寸定义在 common.h，GGML_FA_TILE_Q = 64、GGML_FA_TILE_KV = 64）：

1
// Q_q:    Q_TILE_SZ * DK     已转成计算类型的 Q tile
2
// KQ:     Q_TILE_SZ * KV_TILE_SZ   注意力分数
3
// mask:   Q_TILE_SZ * KV_TILE_SZ   掩码
4
// VKQ32:  Q_TILE_SZ * DV     输出累加器
5
// V32:    KV_TILE_SZ * DV    V tile 的 F32 缓冲
6
// K_f32:  KV_TILE_SZ * DK    K tile（GEMM 路径）
7
float * base = (float *) params->wdata + ith*(Q_TILE_SZ*DK + 2*Q_TILE_SZ*KV_TILE_SZ
8
               + Q_TILE_SZ*DV + KV_TILE_SZ*DV + KV_TILE_SZ*DK + CACHE_LINE_SIZE_F32);

这块 scratch 的字节数就是 work_size 估算里 prefill 那一项的来源：

1
size_t prefill = sizeof(float)*(GGML_FA_TILE_Q*DK + 2*GGML_FA_TILE_Q*GGML_FA_TILE_KV +
2
                GGML_FA_TILE_Q*DV + GGML_FA_TILE_KV*DV + GGML_FA_TILE_KV*DK)*n_tasks;
3
cur += MAX(prefill, decode);

注意 work_size 对 prefill 与 decode 取 MAX 而不是相加：同一个节点只可能走其中一条路径，规划器按两者中更大的那个预留。也注意这里乘了 n_tasks——每个线程一份 scratch，这是工作缓冲与 n_tasks 成正比的最典型例子。

tiled 路径里有一个容易被忽略的优化：整块掩码全为 $-\infty$ 时直接跳过。

1
if (mask) {
2
    bool can_skip = true;
3
    for (int tq = 0; tq < tile_rows; tq++) {
4
        for (int tk = 0; tk < kv_tile; tk++) {
5
            mask32[tq * KV_TILE_SZ + tk] = slope * GGML_CPU_FP16_TO_FP32(mp_row[ic + tk]);
6
            if (mask32[tq * KV_TILE_SZ + tk] != -INFINITY) can_skip = false;
7
        }
8
    }
9
    if (can_skip) continue;
10
}

因果掩码下，靠前的 query 行对应的 KV 块里大部分都是「未来」的位置。不做这个检查，每个空块照样要算一遍 $QK^\top$、过一遍 softmax、再累加一遍 $V$——全都是无效工作。跳过之后，一个 tile 的实际计算量与它覆盖的因果范围成正比，而不再是全 KV 长度。这个检查的代价是读一遍 64×64 的掩码，相比被省掉的几百次点积几乎可以忽略。

还有一处值得说的数据摆放：

1
// Pack K tile transposed: K_f32[dk][kv] so KV_TILE is contiguous (SIMD dim)

K 块被转置成 $[D_K, KV]$ 的布局，原因是后面按列算点积时，$\vec{q}\cdot\vec{k}$ 需要 K 在连续维度上。与 mul_mat 里「激活先转成 Q8_K」的思路一致：数据摆放服从内核的访存模式，而不是让内核去迁就原始布局。

split-KV：把 KV 切成段，各算各的，再合并

这是本篇最值得细看的一条路径，因为它是 CPU 上「序列维并行」的完整实现。

当 neq1 == 1 && neq3 == 1（单 token、单 stream）且 nek1 >= 512 时，注意力改按 KV 切：

1
const int64_t chunk_size = (nek1 + nth - 1) / nth;
2
// Partials buffer layout: [q_head][kv_chunk][M, S, VKQ]
3
const int64_t partial_size  = 2 + DV;
4
float * partials_base = (float *) params->wdata + nth * (DK + 2*DV + CACHE_LINE_SIZE_F32);
5

6
const int64_t ic_start = ith * chunk_size;
7
const int64_t ic_end   = std::min(ic_start + chunk_size, nek1);
8

9
const int64_t partial_stride = nth * partial_size;
10
float * chunk_partials = partials_base + ith * partial_size;
11

12
if (ic_start < nek1) {
13
    for (int64_t q_head = 0; q_head < neq2; q_head++) {
14
        ggml_compute_forward_flash_attn_ext_f16_one_chunk(
15
            params, dst, q_head, q_head + 1, ic_start, ic_end, chunk_partials, partial_stride);
16
    }
17
} else {
18
    for (int64_t q_head = 0; q_head < neq2; q_head++) {
19
        float * q_partials = chunk_partials + q_head * partial_stride;
20
        q_partials[0] = -INFINITY;  // M
21
        q_partials[1] = 0.0f;       // S
22
    }
23
}
24
ggml_barrier(params->threadpool);
25
ggml_flash_attn_ext_reduce_partials(params, dst, nth, chunk_size);

结构是「先分、再 barrier、后合」：每个线程负责一段 KV，为所有 head 各算一遍，得到该段上的局部 $(M, S, VKQ)$ 三元组；全部算完后，全线程在节点内的 barrier 上汇合，再由一部分线程做合并。

几个设计细节：

第一，分的是 KV，不是 query head。 query head 数量通常有 32，本身就是可切的维度，逐行版正是按它切；但按 head 切不能改变「每个 head 读一遍 KV」这个事实。按 KV 切之后，每段 KV 在线程的私有 scratch 里只出现一次，被 32 个 head 轮着用。切分维度变了，访存强度也随之变了——这是 split-KV 与逐行版的本质区别，而不只是并行度。

第二，每个线程还要串行处理所有 query head。 32 个 head 意味着每个线程内部是一个 32 次循环，每次处理一个 head 的一段 KV。head 之间的 scratch 共用，所以 partial_stride = nth * partial_size：同一个线程的不同 head 写到相隔 nth 步的位置，不同线程的同一 head 则是连续的一段。这个布局让合并阶段可以按 (q_head, chunk) 二维索引。

第三，ic_start >= nek1 的线程仍要写合法值。 当 nth * chunk_size 大于 nek1 时末端几个线程负责的区间是空的。它们不能什么都不写，否则合并阶段会读到未初始化的内存；所以显式写入 $M=-\infty$、$S=0$，让合并逻辑通过 if (S_chunk == 0.0f) continue; 跳过。

合并：再一次在线 softmax

ggml_flash_attn_ext_reduce_partials 做的事情，是把各段的局部结果用一个统一的在线 softmax 合成完整结果：

1
for (int64_t q_head = ith; q_head < n_q_heads; q_head += nth) {
2
    float M_final = -INFINITY, S_final = 0.0f;
3
    float * VKQ_final = thread_wdata;
4
    memset(VKQ_final, 0, DV * sizeof(float));
5

6
    for (int64_t chunk_idx = 0; chunk_idx < n_chunks; ++chunk_idx) {
7
        const float * partial = partials_base + (q_head * n_chunks + chunk_idx) * partial_size;
8
        const float   M_chunk = partial[0], S_chunk = partial[1];
9
        const float * VKQ_chunk = partial + 2;
10
        if (S_chunk == 0.0f) continue;
11

12
        const float M_new     = fmaxf(M_final, M_chunk);
13
        const float scale_old = expf(M_final - M_new);
14
        const float scale_new = expf(M_chunk - M_new);
15
        for (int64_t d = 0; d < DV; ++d)
16
            VKQ_final[d] = VKQ_final[d] * scale_old + VKQ_chunk[d] * scale_new;
17
        S_final = S_final * scale_old + S_chunk * scale_new;
18
        M_final = M_new;
19
    }
20
    if (S_final != 0.0f) ggml_vec_scale_f32(DV, VKQ_final, 1.0f / S_final);
21
    memcpy((char *) dst->data + q_head*nb1 + 0*ne1*nb1, VKQ_final, nb1);
22
}

这里用的是双状态在线 softmax：与单个线程内部「维护 $M$ 和 $S$」不同，合并时两个状态的贡献要分开缩放——局部 VKQ 存的是「未归一化的加权和」（相当于 $\sum_j e^{s_j - M_k}v_j$），所以合并时既要按 $\exp(M_k - M_{\text{new}})$ 缩放向量，也要按同样比例缩放 $S$。最后统一除以 $S_{\text{final}}$。这是 FlashDecoding 的标准做法，站内 FlashAttention 完全拆解（二）有更完整的推导。

为什么必须先减最大值

合并公式里每一步都先取 M_new = fmaxf(M_final, M_chunk)，再算 expf(M - M_new)。如果跳过这一步，直接累加 $e^{s_j}$，在 $s_j$ 较大时（例如 attention logits 落在 20–30 的范围）$e^{s_j}$ 会迅速超出 FP32 的表示范围（$e^{88}$ 附近就溢出），后面的除法变成 $\infty/\infty$。减去当前最大值之后，所有指数都落在 $(-\infty, 0]$，$e^{x}\le 1$，不会上溢；代价是每遇到新的最大值要把已累积的向量整体缩放一次。

这个「先减最大值」的操作在单线程版里表现为 ggml_vec_scale_f32(DV, VKQ32, ms)，在合并版里表现为对 VKQ_final 的逐元素乘加。两者都是 $O(D_V)$，在 DV=128 时不过一百多次乘加。代价不在算术量，而在于它是串行的：合并阶段每个 head 的这段循环只能由一个线程完成，且必须按 chunk 顺序推进。

还有一个边界情况：所有 KV 都被掩码挡住（例如一个完全空的序列，或在某些受限解码场景）。此时每个 chunk 的 $S=0$，合并循环通过 if (S_chunk == 0.0f) continue; 跳过，最终 S_final 仍为 0，输出通过 S_final != 0.0f 的判断直接保持 memset 的零值，不会出现 0/0。这类边界在正常对话里罕见，但在测试和投机解码的拒绝分支里会出现。

合并阶段本身又做了一次并行：for (q_head = ith; q_head < n_q_heads; q_head += nth) 按 head 做轮转分配。解码时 head 数是 32、线程数可能是 16 或 32，每个线程负责 1–2 个 head，每个 head 要扫过 nth 个分段（例如 16 段，每段 DV=128 个 float）。合并的算术量是 $n_{\text{head}} \times n_{\text{chunk}} \times D_V$ 次乘加，对 32 head、16 段、128 维是约 6.5 万次乘加——相对整层的乘加量很小，但它是单线程宽度的：每个 head 的合并由一个线程独立完成，没有进一步的并行。

为什么 split-KV 的门槛是 512

nek1 >= 512 这个条件不是随手定的。拆开来看，split-KV 相比逐行版增加的开销有三块：

1. partials 缓冲的读写：每个线程要写 $n_{\text{head}} \times (2 + D_V)$ 个 float，合并时要全部读回。以 32 head、$D_V=128$、16 线程为例，partials 是 $16 \times 32 \times 130 \times 4\text{ B} \approx 266\text{ KB}$，已经超出 L2 的舒适范围；
2. 一次节点内的 barrier；
3. 合并阶段的串行等待：合并由部分线程完成，其余线程在合并期间空闲——它们在 ggml_flash_attn_ext_reduce_partials 里按 head 轮转，head 数少于线程数时确实有闲置。

这些开销大致是常数级的，而 split-KV 带来的收益有两部分：并行度从 min(线程数, head 数) 变成线程数（head 数少于线程数时才成为瓶颈），以及 KV 访存被 head 数摊薄。后者只有在「每段 KV 真的能留在缓存里」时才成立。门槛 512 本质上是一个经验交点：$L=512$、16 线程时每段只有 32 个 key，$32\times128\times2\text{ B}=8\text{ KB}$，仍在 L1/L2 的舒适范围内；若门槛降到 128，每段只有 8 个 key，partials 的固定成本占比会迅速上升，而缓存复用的收益反而不明显。源码把阈值写成常量，是承认这个交点需要按典型硬件标定，而不是从公式推导出来的。

还要注意一个不那么明显的结果：split-KV 路径里 dr 意义上的负载均衡几乎不存在。每个线程固定一段长度相等的 KV，causal mask 对单 token query 不起作用（它只与历史 token 相关，全部可见），所以各段工作量确实均匀。这也是这条路径敢用静态切分的原因——它的并行度来自 KV 长度而不是任务数。

把字节账算出来

用一组具体数字比较逐行版与 split-KV 的访存。取 32 个 query head、8 个 KV head（每 4 个 query head 共享 1 个 KV head）、head 维 128、KV 长度 4096、KV 类型 F16、单层：

• 每个 KV head 的 K 是 $4096 \times 128 \times 2\text{ B} = 1\text{ MB}$，V 同样 1 MB；8 个 KV head 合计 16 MB/层；
• 逐行版：每个 query head 处理一行，读取它对应的 KV head 的 K 与 V。32 个 head 合计 $32 \times 2 \times 1\text{ MB} = 64\text{ MB}$ 的逻辑读取。如果 L2/L3 能把这些数据留住，物理 DRAM 流量可能接近 16 MB；但在 32 层连续执行、每层 16 MB 的工作集下，跨层复用几乎不可能，层内能否复用取决于 L3 容量；
• split-KV 版：每个线程拿到一段 KV（16 线程、4096 长度时每段 256 个 key），在段内轮询 32 个 query head。同一段 K/V 在 L1 里被读入一次、复用 4 次（每个 KV head 对应 4 个 query head），理论上每层的唯一 KV 读取量就是 16 MB；
• 额外的 partials 流量：$32 \text{ head} \times 16 \text{ chunk} \times (2 + 128) \times 4\text{ B} \approx 266\text{ KB/层}$，约占 16 MB 的 1.7%。写 partials 时是每个线程写自己的 32 个 head 结果，读 partials 时按 head 轮转——两次都是顺序访问，对带宽友好。

这组数字说明了一件反直觉的事：split-KV 的主要收益不是「并行度变高」，而是「同样的 KV 只读一遍」。对 32 head 的模型，逐行版与 split 版的并行宽度都可以是 16，差别在于每个线程的 KV 工作集是「整段 1 MB」还是「一段 8–64 KB」。后者能进 L1/L2，前者不能。门槛 512 也就可以从缓存角度理解：只有当每段足够小、整段 KV 又能被所有 head 复用时，改写循环次序才有意义。

当然，实际硬件上逐行版并非一定读 64 MB：多线程并发时 L3 可能替它兜住一部分重复读取。这也是为什么 split-KV 的收益在不同 CPU 上差异很大——L3 越大的服务器，逐行版被缓存救得越多，split 的相对优势越小；L3 较小的桌面/边缘设备，split 的收益更明显。源码把选择条件写成「KV 长度 ≥ 512」再加类型限制，是一个对多种硬件都成立的保守门槛，而不是针对某一台机器调出的最优值。

上图是 FlashAttention-2 论文里关于线程块划分的示意。CPU 版的 split-KV 和 GPU 版的 FlashDecoding 在形式上是一致的：都沿着序列维切开，都靠 partial $(m, l, \text{acc})$ 三元组合并。差别在于 CPU 版的「线程块」只有几十个，每个线程的 scratch 是普通的工作缓冲而非共享内存；GPU 版能同时跑成百上千个 block，因此更频繁地走这条路径。

与 RoPE 的位置关系

大多数模型在写入缓存之前完成 K 的旋转：建图时为 K 建一个 ROPE 节点，旋转后的张量再交给 cpy_k。llama.cpp 新版本里更进一步：当 K/V 是量化类型且 head 维是 64 的倍数时，会在缓存上启用一层注意力旋转（Hadamard rotation），参数在缓存构造时决定：

1
bool attn_rot_k = !attn_rot_disable && n_embd_head_k_all > 0 &&
2
                  ggml_is_quantized(type_k) && hparams.n_embd_head_k() % 64 == 0;

它的动机与量化 KV 有关：先旋转再量化可以让量化的误差在 head 维上更均匀，和站内 vLLM 源码完全拆解（十九）：TurboQuant KV Cache 量化讨论的 Hadamard 旋转是同一类手段。这一段代码只决定是否启用以及旋转维度 n_rot_k（从 64 起倍增到不超过 head 维的一半）。旋转本身不是 set_rows 内部完成的：建图时会在 K/V 投影之后、写缓存之前插入一个 llama_mul_mat_hadamard 节点，cpy_k 拿到的已经是旋转后的张量。这多出来的一对 MUL_MAT 节点（K 与 V 各一次）也计入前面那张节点表之外的额外开销，量化 KV 的收益要扣掉这部分。本篇不展开旋转的数学，后文讲融合时会再次遇到它。

旋转的配对方式：NORMAL 与 NEOX 差在哪

RoPE 的数学形式是「把 head 维的每两个分量看成一对，按位置旋转一个角度」。但哪两个分量算一对，不同模型并不一样。ggml 用 mode 参数区分：

1
switch (mode) {
2
    case GGML_ROPE_TYPE_NORMAL:
3
        rotate_pairs<T>(n_dims, 1, cache, src + n_offs, dst_data + n_offs, 1);
4
        break;
5
    case GGML_ROPE_TYPE_NEOX:
6
    case GGML_ROPE_TYPE_MROPE:
7
    case GGML_ROPE_TYPE_IMROPE:
8
        rotate_pairs<T>(n_dims, n_dims/2, cache, src + n_offs, dst_data + n_offs);
9
        break;
10
    case GGML_ROPE_TYPE_VISION:
11
        rotate_pairs<T>(ne0, n_dims, cache, src, dst_data);
12
        break;
13
}

rotate_pairs 的参数是「参与旋转的维数、配对步长、cos/sin 缓存、源指针、目标指针」。两种主要模式的配对规则：

• NORMAL：相邻分量配对，$(x_{2i}, x_{2i+1})$ 使用第 $i$ 组频率。这是原始 RoPE 论文的排布，也是 GPT-NeoX 之前多数模型的写法；
• NEOX：前后两半配对，$(x_i, x_{i+n_{\text{dims}}/2})$ 使用第 $i$ 组频率。GPT-NeoX、LLaMA 系列采用这种排布，好处是与 Q/K 投影权重的列顺序更自然地对齐，也让向量化的旋转实现更容易；
• VISION：视觉编码器的二维 RoPE，配对步长等于 n_dims，每个位置有独立的频率组。

n_offs 是旋转的起始偏移，允许只旋转 head 维的一段（例如某些模型只旋转前 64 维，剩下的是 NoPE）。源码里旋转完成后还有一段「把未旋转的通道从源拷贝到目标」的循环：

1
if (!is_vision) {
2
    for (int64_t i0 = 0; i0 < ne0; i0 += 2) {
3
        if (i0 == n_offs) { i0 += n_dims - 2; continue; }
4
        dst_data[0] = src[0]; dst_data[1] = src[1];
5
    }
6
}

这段循环的存在说明：ROPE 节点的输出并不只是「旋转过的部分」，而是整个 head 维。未旋转的通道必须原样搬运，否则下游注意力会读到未初始化数据。这也意味着即使 n_dims 只有 64、head 维是 128，RoPE 依然要读写完整的 128 维——一半的带宽花在拷贝上。某些架构因此把 NoPE 部分单独切出来，避免这次无谓的拷贝；这是模型层面的优化，不是 ROPE 算子本身的。

理解配对规则对读后面的代码有直接帮助：rotate_pairs 的第二个参数决定了 src 与 dst 的索引步长，也决定了每行的工作量是 $n_{\text{dims}}/2$ 次旋转（NORMAL）还是 $n_{\text{dims}}/2$ 次跨半区旋转（NEOX）。两者计算量相同，但访存模式不同：NORMAL 是顺序读相邻两元素，NEOX 是一前一后两个相距 n_dims/2 的元素。对 CPU 缓存来说，NORMAL 更友好；这也是部分模型在转换到 ggml 时需要重排权重的原因之一。

行切分的实际并行度

RoPE 的行数是 $n_{\text{head}} \times N$。decode 时 $N=1$，行数就是 head 数（32）。如果线程数超过 32，dr = 1，前 32 个线程各处理一行，后面的线程拿到 ir0 >= ir1 的空区间。注意力里 head 是 32，RoPE 这里也是 32，但 RoPE 没有像注意力那样按 KV 切——因为 RoPE 的计算量与 KV 长度无关，只有 head 维的 $D/2$ 组旋转对，每行的工作量固定且很小。

这一点决定了 RoPE 在 decode 阶段是一个「并行度有限的小节点」：即使线程很多，它也只能用满 min(nth, 32) 个线程，其余线程在 barrier 处等待。单次 RoPE 的时间可能只有几微秒，但它连同 barrier 一起，在 32 层里重复出现——barrier 的数量不随 RoPE 变快而减少。

判定条件来自 ggml_can_fuse

ggml_can_fuse 定义在 ggml-impl.h，它检查的是三个条件：

1. 节点序列的算子类型与期望的 ops 数组一致；
2. 除最后一个外，每个节点只有一个使用者（ggml_node_has_n_uses(cgraph, idx, 1)），且不是 view、不是用户显式请求的输出；
3. 后一个节点的 src[0] 或 src[1] 是前一个节点，且形状相同。

为什么这三个条件是必须的？

• 条件 2 的「只有一个使用者」：如果 RMSNorm 的结果还被别的节点使用（例如被残差 ADD 用），那么融合执行时就不会写出中间结果，那个节点就拿不到数据。所以只有在「中间结果没人关心」时才能省掉它；
• 条件 3 的「形状相同」：融合实现里一个线程处理一行，同时读写两个张量的同一行，形状不同就没法共用同一个行索引。

RMSNorm + MUL 这个模式在 build_norm 里是普遍存在的：

1
if (mw) {
2
    cur = ggml_mul(ctx0, cur, mw);
3
}

只要模型带 norm 权重（几乎全部 RMSNorm 模型都带），每层就有两次「RMS_NORM 紧跟 MUL」。融合之后，原来需要「RMSNorm 节点 + barrier + MUL 节点 + barrier」变成「融合节点 + barrier」。

融合实现真的省了什么

看 ggml_compute_forward_rms_norm_f32 的模板参数分支：

1
if constexpr (FUSE_OP == GGML_RMS_NORM_FUSE_OP_MUL) {
2
    src1 = (dst_fused->src[0] == dst_rms_norm) ? dst_fused->src[1] : dst_fused->src[0];
3
    dst  = dst_fused;
4
}
5
...
6
if constexpr (FUSE_OP == GGML_RMS_NORM_FUSE_OP_MUL) {
7
    const float * w = src1 行指针;
8
    for (int64_t i00 = 0; i00 < ne00; i00++) y[i00] = x[i00] * scale * w[i00];
9
} else {
10
    memcpy(y, x, ne00 * sizeof(float));
11
    ggml_vec_scale_f32(ne00, y, scale);
12
}

原来的两步是：

1. RMSNorm 把 $x$ 的缩放结果写到 y；
2. 下一个节点 MUL 再从 y 读、从权重读，写到最终输出。

融合后是：

1. 一次读 $x$、一次读权重 $w$、一次写 $y = x \cdot \text{scale} \cdot w$。

省下的是：中间张量的一次写 + 一次读（两个 $n$ 字节的访存），以及一次节点 barrier + 一次节点调度。对每层两次 norm、32 层共 64 次融合，这就是省掉 64 次 barrier、以及 128 次中间张量的读写。在隐藏维 4096、F32 的情况下，每次中间张量是 16 KB；64 层 × 2 × 32 KB（一读一写）= 4 MB 的总流量，相对权重读取可以忽略，但 barrier 的节省是可观的。

但也要看到它还没有融合的部分：RMSNorm 与它前面的 MUL_MAT 之间没有融合，与残差 ADD 也没有融合。源码用一条注释明确了目前的定位：

1
// TODO: move fused-op detection into ggml_graph_plan so fusion decisions are made once at planning time

也就是说，融合判定在每次执行时对每个节点做一遍——虽然每层只有一次判断，但仍然是热路径上的检查。更重要的是，融合的实现目前只有一个组合（RMS_NORM + MUL）。其他后端有更丰富的融合表，例如 CANN 后端在它的调度器里识别 {ADD, RMS_NORM}、{MUL_MAT, ADD} 等组合；CPU 后端没有对应的通用表。这一点从后端自己的代码里能看出来：

1
// ggml-cann: if (ggml_cann_can_fuse(cgraph, i, { GGML_OP_ADD, GGML_OP_RMS_NORM }))
2
// ggml-et:   if (ggml_et_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }))
3
//            if (ggml_et_can_fuse(cgraph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD }))

CPU 后端只在 ggml_cpu_try_fuse_ops 里写死了一个组合。这是一个明确的可扩展点，也是 CPU 推理继续优化的方向之一。

关闭融合的入口

两个开关都能让融合失效：ggml_cpu_disable_fusion（在 ggml_cpu_init() 里读一次环境变量后只读）和 cplan->use_ref。后者是与调试相关的——当用户要求走参考实现时，融合路径不参与，保证输出与参考路径一致。这也间接说明，融合实现是「等价但不同代码路径」的，测试时需要两条路径都覆盖。