multi_card_cp
FlashKDA 多卡 Context Parallelism (CP) 数学原理详解
本文面向"小白"读者,从最基础的概念讲起,一步步推导多卡 CP 的 forward 和 backward 数学原理,重点讲清线性分解这个核心。 每个公式都配文字解释,最后用可手算的数值例子把整条链走一遍。
阅读顺序建议:§1 直觉 → §2 记号 → §3 forward → §4 backward → §5 跨卡编排 → §6 例子。 如果只想看结论,直接跳 §7 一页速查表。
1. 先建立直觉:CP 到底在切什么?
1.1 一句话版本
FlashKDA 是一个带循环状态的注意力算子。你可以把它想象成一个"读书做笔记"的过程:
- 从左到右一个 token 一个 token 地读;
- 每读一段(一个 chunk = 16 个 token),就更新一次"笔记本" (状态矩阵);
- 笔记本 会一直往后传,后面的输出依赖前面积累的笔记。
问题是:这个"往后传"是串行的。序列长到 64K、256K 时,GPU 上千个核心大部分在干等前一段算完。
Context Parallelism(上下文并行) 就是把这条长序列沿时间维 切成好几段,让不同的段(不同的 GPU / 不同的 SM)同时开工。
- 单卡 CP:在一张 GPU 内部,把序列切成多个 sub-segment,让空闲的 SM 并行跑(见
cp_optimization.md)。 - 多卡 CP:把序列切到多张 GPU上,每张卡(rank)负责一段 token,卡之间只传很小的"边界信息"。
本文讲的是多卡。
1.2 切开之后的麻烦:接缝处的状态
假设有 3 张卡,序列 切成三段:
1 | rank 0 rank 1 rank 2 |
- forward 的麻烦:rank 1 要正确工作,需要知道"读到 时笔记本长什么样"——而这个笔记本是 rank 0 算出来的。状态 从前往后流。
- backward 的麻烦:反向传播时,梯度 从后往前流。rank 1 的右边界梯度,来自 rank 2 的左边界。
CP 的全部技巧,就是如何用最小的通信,把这个"接缝处的状态/梯度"接好。
而 FlashKDA 能把通信压到极小的根本原因,是下面要讲的线性性质:状态对初始值是线性的,所以接缝可以用一个小矩阵一次接好,不用真的把整段重算一遍。
1.3 为什么 KDA 特别适合 CP?
普通 softmax 注意力有因果三角,切开后每段计算量不均(前面的段轻、后面的段重),需要复杂的负载均衡(zigzag 重排)。
KDA 是线性注意力:每个 chunk 的计算量恒定。所以等长切分 = 天然负载均衡,实现上省掉一大块麻烦。
2. 记号与基本设定
先把符号定清楚,后面推导都用这套。
| 符号 | 含义 | 形状 |
|---|---|---|
| 序列总长度(token 数) | 标量 | |
| head 数 | 标量 | |
| key / value 维度,FlashKDA 固定 | 标量 | |
| chunk 大小,固定 | 标量 | |
| rank 数(GPU 数) | 标量 | |
| 状态矩阵(“笔记本”) | ||
| 进入第 个 chunk 前的状态 | ||
| 初始状态 | ||
| query/key/value/gate 输入 | ||
| beta 门控(logit) |
说明:下面所有推导都是**“每个 head 独立”**的。 个 head 之间没有耦合, 所以我们只写一个 head 的公式,实际代码对 维并行即可。同理 batch 维在 varlen 下 。
2.1 一个 chunk 的 forward(来自 autograd.py)
FlashKDA 每个 chunk(16 个 token)做的事,数学上是(全部 fp32 计算):
其中 ,, 是 sigmoid。
别被吓到。对 CP 而言,这一大堆公式里我们只关心一件事:
状态 是怎么从 变成 的?而且这个变换对 是不是线性的?
下一节就把这一坨压缩成一个干净的仿射式子。
3. Forward 的线性分解
3.1 把一个 chunk 写成仿射变换
看 §2.1 最后两行,把含 的部分和不含 的部分分开。
里含 :,所以
代入状态更新 :
于是每个 chunk 就是一个仿射变换:
是 矩阵,叫转移矩阵(transition matrix),只跟 有关,跟 无关。 是 矩阵,是数据项,含 。
小白理解:
是"旧笔记怎么衰减/改写", 是"这一段新读到的内容"。 新笔记 = 旧笔记按 变换一下 + 新内容 。
3.2 一整段的链:
一个 rank 负责一段,里面有
定义两个量:
则整段的输入输出关系是一个干净的仿射式:
这就是 forward 线性分解,是整个 CP 的数学基石。含义:
(转移矩阵):初始状态 穿过这一整段后,被"缩放/旋转"成 。它只跟这一段的 有关,跟 都无关。 (数据驱动终态):假设 ,纯靠这一段的 数据能积累出的终态。
为什么这个分解如此重要? 因为它把"整段递推"这个串行过程,浓缩成了两个可以提前、独立算好的量
。 一旦有了 ,任意给一个 ,我都能一步(一次矩阵乘 + 一次加法)算出 , 不用把整段 个 chunk 重新串一遍。跨卡时,卡之间只需要传 这两个小矩阵。
3.3 怎么算 和 ?(state_only kernel)
state_only kernel 分别从两个不同的初值出发,跑同一条递推,就能"测量"出
- 算
:初始 ,跑 ,终态就是 (因为 )。 - 算
:初始 (单位阵),且去掉数据项 (因为 跟 无关), 跑 ,终态就是 (因为 )。
对应 fwd_state_only.cuh 里的 CalcMt 模板参数:CalcMt=false 从 0 出发算 CalcMt=true 从
3.4 关键简化:强 decay 下 (以真实 warmup kernel 为准)
强 decay 时,
含义:decay 强时,初始状态
但"往回扫几个 chunk 就够"这件事不能拍脑袋定死,得逐段、逐 head、逐维度地实测。
真正决定它的是 csrc/warmup_chunks.cu(get_warmup_chunks_kernel),逻辑如下——
注意它和旧文档 cp_optimization.md 里"mean over D + 0.25 + 中心极限"的说法不一样,
以下这版是当前 kernel 的真实实现。
(1) 衰减在 gate_scale:
于是残差幅度是 gate_scale
把自然底换成了 2 底:
(2) 维度上取 min,不是 mean。 kernel 对 decay 越接近 0
(3) head 上取 max。 各 head 的 g_cumsum 做树形归约取 max(最不负 = 最慢的 head),
只有当最慢的 head 都收敛才算整段收敛。
(4) 阈值 + 逐段自适应 + fallback。 往回扫直到 num_warmup = c+1(该段只需这么多 chunk)。
若扫完整段仍没到阈值(weak decay,fallback=true——
这一段不能用
所以"强 decay 下
4. Backward 的线性分解(forward 的镜像)
这一节是本文重点。我们要证明:backward 也有一个和 forward 完全对称的线性分解,
只是方向反过来(从后往前),转移矩阵用转置
4.1 什么是 ?
训练时有一个 loss
即"进入第
4.2 的递推:从后往前的仿射链
状态
- 本 chunk 的输出
,而 里也含 (通过 )。 - 下一个状态
,进而影响之后所有 chunk。
按链式法则,
第二项
第一项
- 第一块
:来自 ,直接对 求导。 - 第二块:来自
,而 经 又依赖 ,把这条路回传回来(带负号,因为 里是 )。
于是 backward 就是一条从后往前的仿射链:
把它和 forward 的
对照:
- forward:往后走,用
,数据项 。 - backward:往前走,用
,本地源项 。 一模一样的结构,只是方向和转置不同。这就是"镜像"。
这条链对应 bwd_kernel2.cuh 的主循环:for t = t_tiles-1 downto 0(倒序),
Step 8b "dS 向前传播(为上一个 chunk)"就是在算
4.3 一整段的 backward 分解:
现在把这条链在一段
注意第一项的矩阵连乘:
这正是 forward 里那个同一个
这个求和号是怎么来的?( 去哪了)
很多人到这里会困惑:上面展开时明明写的是一个个具体的
答案:
把递推
观察规律:带
| 求和里 |
展开成的具体项 |
|---|---|
所以
和 forward 的
是同一个套路: 也是把 用一个 收起来。backward 完全镜像,只是把 换成本地源 、 换成 ,连乘方向从" 在 右边"变成" 在 左边"(因为链方向反了)。
就得到 backward 线性分解:
含义完全镜像 forward:
:右边界梯度 穿过整段传到左边界,被 变换。 它就是 forward 那个 的转置,免费复用,不用重新算。 :假设 ,纯由本段 数据积累、传到左边界的梯度。 它是 backward 版的"数据驱动项",对应 forward 的 。
4.4 对照表:forward ↔ backward 完美镜像
| 维度 | forward | backward |
|---|---|---|
| 边界量 | 状态 |
状态梯度 |
| 链方向 | 前缀(前 → 后) | 后缀(后 → 前) |
| 单 chunk 递推 | ||
| 转移矩阵 | ||
| 数据项 | ||
| 整段分解 | ||
| 转移分量 | ||
| 数据分量 | ||
| 入口边界 | ||
| 出口 | ||
| 强 decay |
一句话总结:backward 就是把 forward 的图倒过来走一遍, 把
换成 、 换成 、 换成 、“前缀"换成"后缀”、 (左入口)换成 (右入口)。数学结构一模一样。
4.5 强 decay 下 backward 也退化
backward 用的是同一个
含义:右边界梯度
5. 跨卡编排:把线性分解用起来
现在把序列切成
5.1 Forward:跨卡前缀 scan
边界条件:
= rank 的正确 ; = 全局 final_state(在最后一个 rank)。
这是一个前缀 scan(从 rank 0 往后)。
5.2 Backward:跨卡后缀 scan
边界条件:
= rank 的右边界梯度(反向的输入); = 全局 d_initial_state(在 rank 0)。
这是一个后缀 scan(从 rank
5.3 两种落地路线(forward / backward 都适用)
路线 A(精确串行,做对拍基线用) rank 按顺序跑,前一个 rank 的边界算完传给下一个。正确但完全串行, 退化成没有并行收益,只用来验证正确性。
路线 B(并行,推荐,即 fla 的 build_cp_context 思路) 以 backward 为例:
1 | 1. 每个 rank 本地反向,右边界 dS 先置 0 |
forward 的路线 B 完全对称:本地跑 state_only 得
5.4 参数梯度的 all-reduce
反向产出的梯度分两类:
| 梯度 | 性质 | 跨卡处理 |
|---|---|---|
| per-token,按 rank 天然分片 | 不需要 reduce,各卡各自持有自己 token 的梯度 | |
| 全局参数,forward 里被所有 token 用 atomicAdd 累加 | 每卡只累加了自己 token 的贡献 → all-reduce SUM |
为什么
、 要 all-reduce? 因为 (每 head 一个)和 (每 head、每维一个)是所有 token 共享的参数。 完整的梯度 = 所有 token 贡献之和。切卡后每张卡只看到一部分 token, 算出的只是部分和,必须跨卡加起来才是完整梯度。 而 等是每个 token 各自一份,rank 之间不重叠,所以不用加。
5.5 插曲:集合通信原语(all-gather / all-to-all / all-reduce)
上面反复用到 “all-gather”、“all-reduce” 这些词。它们是集合通信(collective communication)原语 —— 多张卡之间约定好的标准数据交换动作。这里用 4 张卡 (rank 0/1/2/3)、每张卡起始各持一块数据的例子讲清楚,重点是 CP 里最常用的 all-gather 和容易混淆的 all-to-all 的区别。
all-gather:人人贡献一块,人人拿到全部(聚合,结果一致)
1 | 起始 结束(all-gather 后) |
- 每张卡贡献 1 块,最后每张卡都拿到全部 4 块(拼接),所有卡结果完全相同。
- CP 里的用途:all-gather 各 rank 的
/ , 让每张卡都凑齐所有 rank 的边界量,好在本地各自冗余跑那条前缀/后缀 scan(§5.3 路线 B)。
all-to-all:每卡按目标把数据切块分发(转置,结果各异)
起始时每张卡持有 4 块,第
1 | 起始 结束(all-to-all 后) |
- rank
的第 块送到 rank ,相当于把这个 数据矩阵做一次转置。 - 每张卡发出去的是不同的块,收到的也来自各卡的不同块,所有卡结果各不相同。
- 典型用途:MoE 的 token dispatch(按专家路由分发)、序列并行
张量并行重排、分布式转置。 FlashKDA 的 CP 边界交换用不到 all-to-all(我们要的是"人人凑齐全量",不是"按收件人分发")。
核心区别一览
| all-gather | all-to-all | |
|---|---|---|
| 每卡发的内容 | 同一份(自己那块) | 按目标切成 |
| 每卡收的结果 | 全量拼接,人人相同 | 各收各的,人人不同 |
| 本质 | 收集 / 广播式聚合 | 转置 / 重分布 |
| 类比 | “把大家的名片收齐,人手一套” | “每人把一叠信按收件人分发出去” |
顺带区分 all-reduce 和 reduce-scatter
- all-reduce = all-gather 再加一步逐元素归约(求和 / 求 max 等),
最后人人拿到的是同一个归约结果(不是拼接)。§5.4 里
、 就是用 all-reduce SUM 把各卡的部分和加成完整梯度。 - reduce-scatter = all-reduce 的前半段:归约后每卡只留结果的一片(不是全量)。
1 | all-gather: 各卡[片] -> 各卡[全量拼接] (结果一致, 无计算) |
一句话记忆:all-gather 收集拼接、all-reduce 收集再求和、all-to-all 是转置分发。 CP 边界交换用 all-gather,参数梯度用 all-reduce SUM,两者都用不到 all-to-all。
5.6 通信量小结
- 每次跨卡只传
个 边界矩阵。 时约 (fp32), 相对 activation 微不足道。 - 强 decay 生产参数下
,连 都不用传,只传 (一个边界量)。 - 主计算(完整 forward / backward)全部本地并行,跨卡只有轻量 warmup + 边界 scan + 两个参数梯度的 all-reduce,开销和单卡 CP 的 ~5–10% 一个量级。
6. 手算例子:把整条链走一遍
为了让链看得见摸得着,我们把状态退化成标量(
| rank |
转移 |
forward 数据项 |
backward 本地源 |
|---|---|---|---|
| 0 | 0.5 | 2.0 | 0.3 |
| 1 | 0.4 | 1.0 | 0.5 |
| 2 | 0.3 | 3.0 | 0.7 |
6.1 Forward 前缀 scan(从 rank 0 往后)
初始状态
每张卡拿到自己的
6.2 Backward 后缀 scan(从 rank 2 往前)
右边界梯度
每张卡实际干的事(路线 B):
1 | 卡2: 本地反向(右边界 = dfinal = 1.0) -> 直接得 dS_2 = 1.0, 本地 dq/dk/... 已精确; 传 dS_2 |
6.3 验证:CP 版和单卡版数值完全一致
如果不切卡,单卡把这 3 段连成一条链一次反向,
为什么是精确相等,而不是近似?
因为
6.4 强 decay 演示:跨卡项被压没
把
forward:
backward:
跨卡回传项全被
6.5 从标量回到矩阵( )要注意什么
上面的例子是标量
都是 矩阵, 不再交换,连乘顺序不能颠倒。- forward 前缀 scan 是
(矩阵乘向量化到 ); backward 后缀 scan 是 ,注意转置。 - 强 decay 下
是接近零矩阵, 也接近零矩阵,退化结论不变。 - 每个 head 独立,实际是
份 的 scan 并行。
7. 一页速查表
Forward(前缀 scan,前 → 后)
Backward(后缀 scan,后 → 前)
单 chunk
跨卡通信
- forward:all-gather
→ 本地前缀 scan → 用 当 并行 forward。 - backward:all-gather
→ 本地后缀 scan → 用 修正并行 backward。 - 参数梯度:
天然分片不 reduce; 需 all-reduce SUM。
强 decay(
跨卡耦合可忽略,每张卡 forward/backward 近似独立并行,开销 ~5–10%。
工程要点(多卡特有)
- conv1d halo:段首需前一 rank 段尾
个 token,一次小通信。 - 切分点对齐到 chunk 边界(
),避免破坏 workspace 的 per-tile 布局。 - varlen:cu_seqlens 要叠一层 rank→段 映射,处理一条序列被切到多卡。
- workspace / all_states 是本地前向产物、本地反向消费,不跨卡(这是通信只剩边界量的前提)。
- 可与单卡 CP 叠加:rank 内再切 sub-segment,两层 scan 可组合(先卡内、边界项再进卡间)。
8. 训练中能用吗?会不会有精度问题?
结论:能用,而且在数学上和单卡训练等价——只要 warmup 的 fallback 打开。
分两层看:
8.1 线性分解本身是精确恒等式,不引入近似
8.2 唯一的近似来自"丢 ",而它被 warmup + fallback 兜底
真正可能引入误差的只有一个动作:强 decay 段把 get_warmup_chunks_kernel 用 min-over-D(最慢维度)+ max-over-head(最慢 head) 的保守估计,
只在 fallback=true 强制算精确
- 训练想要零近似:把 threshold 调得足够负(或直接强制 fallback / 走 §5 的路线 A 精确 scan), 此时 CP 训练与单卡没有任何数学差异,可放心用。
- 训练想要跨卡近零开销:用默认 threshold,误差被
有界压住, 且只落在强 decay( 本就被遗忘)的段上——被丢的项本身量级就 梯度噪声/bf16 舍入,实践中对收敛无影响。
一句话:丢
8.3 精度注意事项(和单卡内核同源,CP 不新增)
CP 不改变内核的数值精度特性,训练需要注意的与单卡一致:
- 状态/中间量 bf16,累加走 fp32:backward 的 workspace、
累加、INV 都在 fp32 下做, 用 fp32atomicAdd累加(跨卡再 all-reduce SUM),避免 bf16 累加掉精度。 - 切分点对齐 chunk 边界(
):错位会破坏 per-tile workspace 布局,属于正确性问题不是精度问题,必须对齐。 必须 all-reduce SUM:这两个是所有 token 共享的参数梯度, 漏掉跨卡归约会让梯度偏小(等价于只用了本卡那段 token),这是 CP 训练最容易踩的坑,和精度近似无关。- fallback 段的开销:weak decay(
接近 0)会大量触发 fallback → 算精确 → 跨卡通信与计算变重。这是"精度换开销"的旋钮,不是精度损失。
附录:符号速记
| 符号 | forward 含义 | backward 对应 |
|---|---|---|
| 状态(笔记本) | ||
| 单 chunk 转移矩阵 | ||
| 单 chunk 数据项(含 |
||
| 整段转移矩阵 | ||
| 整段数据驱动终态( |
||
| 初始状态(左入口) | ||
| 终态(右出口) |
- 标题: multi_card_cp
- 作者: 鱿鱼圈
- 创建于 : 2026-09-12 22:20:00
- 更新于 : 2026-09-12 22:26:40
- 链接: https://yuyanqi.com/2026/09/12/multi_card_cp/
- 版权声明: 本文章采用 CC BY-NC-SA 4.0 进行许可。