[Task32] feat: fla mcore chunkwise cp stage 3

[jambow0320]

Task 32 第三阶段：迁移 MCore #5664 的 THD route 预构建，让 Chunkwise 产生正向收益

对应 RFC：redai-infra/Relax#213
第一阶段 PR（FLA 0.4.2 + MCore #3282 backport）：redai-infra/Relax#251
第二阶段 PR（Relax 侧 GDN CP 静态路由接入）：redai-infra/Relax#254

结论：接入后 chunkwise 从"比 headwise 慢 17%"变成静态 CP 下快 45–52%、dynamic CP 下快 23%，在测过的四种 recompute × sequence length × 静态/动态 CP 组合下都是最快的模式。同 harness A/B 显示 chunkwise 训练吞吐 23,371 → 43,397 tok/s（+85.7%），headwise 在噪声范围内不变（+1.3%）。GDN 布局转换引入的 device-host 同步从 65,280 次 / 2.69 s 降到 0，GPU 忙碌率 56.2% → 85.1%。

---

1. 要解决什么问题

headwise 和 chunkwise 的模型算力是相同的——headwise 按 head 切（每卡算全序列 × 1/cp 的头），chunkwise 按时间切（每卡算 1/cp 序列 × 全部头），总 FLOPs 一样，两者都不重复计算；只有 all_gather 每卡跑全序列 × 全部头，重复 CP 倍。trace 直接证实了这点：两种模式的 GEMM kernel 名字、调用次数、耗时全都对得上（例如 nvjet_tss_128x256_64x4_2x1_v_badd_coopA_NTN 是 630.8 ms × 963 对 631.4 ms × 963）。

然而 Stage 2（#254）里朴素实现的 chunkwise，性能既不如 all_gather 也不如 headwise。

既然算力相同，chunkwise 慢就只能是开销。把 GPU kernel 时间拆成 NCCL 与非 NCCL 两部分看（同为 --profile-with-stack 抓取，rank 0）：

模式	span	GPU 忙碌率	kernel 总计	NCCL	非 NCCL 计算
headwise	21.45 s	75.0%	15.66 s	10.27 s	5.40 s
chunkwise	22.73 s	56.2%	12.13 s	5.76 s	6.37 s
all_gather	12.66 s	88.3%	10.56 s	2.65 s	7.92 s

通过 trace 分析，热点集中在 _zigzag_contiguous_thd_swap。每做一次 zigzag↔contiguous 转换，它都要现场推导 all-to-all 的路由：

• 2 × cp_size 次 get_thd_context_parallel_rank_indices，每次内部有 cu[0].item()、cu[-1].item() 和两次 torch.any(...) 判断 → CP=4 时 32 次同步；
• cp_size 次 nonzero() 和 cp_size 次布尔索引 —— 输出形状依赖数据，必然同步；
• 一批 arange / bucketize / argsort / scatter 的小 kernel，规模是 O(T_global)。

而这套路由只依赖 cu_seqlens 和 (cp_size, cp_rank)：一个 micro-batch 内所有 GDN 层、两个转换方向、以及 full recompute 的重放，算出来的结果完全一样。Stage 1 backport 的代码里本来就留着这个 TODO：

# TODO: Let a future CP layout scheduler precompute this routing once per
# microbatch from immutable cu_seqlens and pass it through both THD swaps.
# Do not cache it across microbatches because packed sequence boundaries change.

NVIDIA/Megatron-LM#5664 是一个比较大的 PR 改动，其中包含这个路由缓存的优化。本 PR 只最小化迁移这部分的相关代码。

---

2. 实施方案

从上游 MCore PR 迁移代码，改动只落在 pinned MCore patch 的两个文件，relax/ 目录零改动。

2.1 核心思路（来自 #5664）

两种 layout 都是全局 token 区间的并集——contiguous 每 rank 一段连续区间，zigzag 每 rank 2 × 序列数 段。所以整条 all-to-all 路由可以用区间求交得到，不需要逐 token 的索引张量：把 cu_seqlens 一次性 .tolist() 到 CPU，之后全部用 Python int 做双指针求交，产出 send_rows / recv_rows / input_split_sizes / output_split_sizes。

相比于之前的朴素实现（逐 token 计算索引），复杂度从 O(cp_size × T_global)（GPU）降到 O(cp_size × 序列条数)（CPU）。

与此同时把这四个对象存进 route 缓存，之后 forward 时只需取出直接做 all-to-all，不必每层重算一遍 CPU/GPU 上的索引：

send_buf = _pack_thd_cp_route_send_buffer(x, route.local_source_length, route.send_rows)
recv_buf = all_to_all(cp_group, send_buf, route.output_split_sizes, route.input_split_sizes)
out      = _scatter_thd_cp_route_recv_buffer(recv_buf, route.recv_rows, out_shape)

2.2 原封不动迁移过来的函数（6 个）

符号	作用
_cp_layout_nvtx_range	nvtx range 的 contextmanager
_compact_thd_cu_seqlens_to_list	cu_seqlens 一次性 .tolist() + 去掉重复边界（padding 空槽）
_append_range	rows.extend(range(start, start + length))
_row_list_is_identity	判断置换是否恒等
_pack_thd_cp_route_send_buffer	恒等则零拷贝，否则 index_select
_scatter_thd_cp_route_recv_buffer	恒等则零拷贝，否则 index_copy_

2.3 只改了命名和报错文案，核心逻辑不修改（4 个）

符号	具体差异
_validate_thd_route_partitioning	报错文案加了 "/contiguous" 一个词
_build_thd_layout_segments	参数 cp_partition_mode → layout；加了 docstring
_intersect_thd_layout_segments	加了 1 行 docstring
_thd_cp_partition_route_attr_name	同样的参数重命名

为什么改名：#5664 在它自己的重构里把 layout 全局重命名成了 cp_partition_mode。跟着改会污染 Stage 1 已经暴露出去的接口（get_thd_context_parallel_rank_indices(..., layout=...) 以及已合入的单测），所以保留 Stage 1 的命名。

2.4 迁移过来但做了一定改动（3 个）

build_thd_cp_partition_route

求交主体（区间构建 → 双指针求交 → 展开行号 → 两个完整性 assert）与上游逐字节相同，nvtx range 也保留。唯一实质差异是返回类型：

• 上游把结果序列化成一个扁平 torch.Tensor，每次转换再 decode_thd_cp_partition_route() 解开；
• 我们直接返回解码后的 ThdCpPartitionRoute。

原因：上游那层编码是为了让 route 能当 CUDA graph 的捕获输入，而 decode 里有 3 处 .cpu()，有一定性能影响，而且 Megatron 本身对 THD 布局转换是拒绝 full-iteration CUDA graph 的。

兼容性：ThdCpPartitionRoute 的前 6 个字段就是上游 decode_thd_cp_partition_route() 的返回元组，连"恒等用 None 表示"这个约定都一样（上游 _encode 把恒等的 payload 存成空列表，decode 出来就是 None）。这条偏离已写进函数 docstring，并注明：若将来 GDN 需要在 graph capture 下运行，这里必须重新考虑。

get_thd_cp_partition_route / prebuild_thd_cp_partition_routes

接口与上游一致，函数体不同。上游假定 route 由数据侧（get_batch）主动 prebuild，所以查找是一次裸 getattr，build-on-miss 只是兼容后路并会 warnings.warn(FutureWarning)。我们这边 build-on-miss 是预期路径（原因见 2.5），warn 只会变成噪声；取而代之的是每次复用前先过一道有效性校验。

2.5 Relax 所必须要的特有逻辑（2 个）

当前不采用上游"在数据准备侧提前算好 route"的做法，改成了懒构建 + 缓存。

上游 #5664 的用法是：数据侧在 get_batch 里调 prebuild_thd_cp_partition_routes(packed_seq_params)，把两个方向的 route 一次性写到 PackedSeqParams 上；模型里所有消费点只管 getattr 取用。

当前 PR 考虑到两个问题暂时没有使用这套机制：

1. 不大规模重构成 #5664 的 block-level layout 调度，布局转换仍留在 GDN.forward 里（#3282 的位置），route 的消费点在模型内部而不是数据侧；
2. 数据侧 prebuild 在 Relax 这边不可靠——Bridge/VLM 路径的 preprocess_packed_seqs 会在 embedding 之后重打包出新的 PackedSeqParams，GDN 拿到的未必是数据侧建的那个，提前 prebuild 有可能挂在一个随后被替换掉的对象上。

所以改成：第一个 GDN 层用到时才构建，缓存到本次 forward 实际收到的那个对象上，同 micro-batch 的后续层与 recompute 重放直接命中。

上游那套单写多读天然不会陈旧，我们这套懒构建 + 缓存复用则必须自己证明缓存没过期（下一个 micro-batch 换边界、dynamic CP 换 cp_size/cp_rank）。校验要拿当前请求和 route 的来源逐项比对，route 因此得自带来源信息——这就引出两个 Relax 特有的符号：

符号	说明
ThdCpPartitionRoute（NamedTuple）	前 6 个字段 = 上游 decode_thd_cp_partition_route() 的返回值，参与计算；后 5 个（cu_seqlens / cp_size / cp_rank / source_layout / target_layout）是 Relax 加的来源凭据，不参与任何计算
_thd_cp_partition_route_is_reusable	复用前逐项核对上述凭据。其中 cu_seqlens 用对象身份（is）而不是数值比较

prebuild_thd_cp_partition_routes() 也按上游接口导出了，但本 PR 未调用；等 Bridge/VLM 的最终 PackedSeqParams 对象归属理清后，可以直接切回上游那套数据侧 prebuild 的用法。

2.6 老代码的改动面

context_parallel_layout.py 里 Stage 1 已有的函数：

| 函数 | 状态 |
| _zigzag_contiguous_thd_swap | 96 → 53 行。头（cp_size==1 短路、movedim、contiguous）、尾（movedim 回去、contiguous）和 all_to_all 的调用形式全部没动，只把中间 60 行的路由推导换成"取 route + 三步执行" |

2.8 一个不属于 #5664 的附带优化

route 优化之后，profile 里第一名同步点变成 _resolve_cu_seqlens（943 ms / 2,880 次）。它每层都要做 cu_seqlens[-1].item()、(seq_lengths % cp_size != 0).any() 和 torch.equal(cu_q, cu_kv)。

新增 _resolve_thd_cu_seqlens()，用与 route 相同的缓存键（四个源张量的对象身份 + seq_len_global + cp_size）把这套校验收敛到每 micro-batch 一次。

影响面：对 headwise 也生效（A/B 实测 +1.3%，噪声量级）；对 all_gather 不生效（Relax fallback 整个替换了 MCore 的 forward）

2.8 改动清单

文件	改动
docker/patch/megatron/20260805-85bced0ae.patch	context_parallel_layout.py 307 → 691 行；gated_delta_net.py 的 forward + 新增 _resolve_thd_cu_seqlens
tests/backends/megatron/test_gdn_chunkwise_cp_route.py	新增，10 个测试函数（参数化展开 33 项）

---

3. Trace 对比

同一 recipe（8×H200 / Qwen3.5-9B / TP2 / 静态 CP4 / SP / full recompute / max-tokens-per-gpu 8192），seed 1234，用 Relax 自带的 --use-pytorch-profiler --profile-target train_overall --profile-with-stack 抓 step 5，取 rank 0。

优化前（Stage 2）

优化后（本 PR）

3.1 同步热点

调用点	优化前	优化后
context_parallel_layout.get_thd_context_parallel_rank_indices	2,417.6 ms / 61,440 次	0
gated_delta_net._resolve_cu_seqlens	943.5 ms / 2,880 次	1.5 ms / 60 次
裸 Tensor.nonzero	267.4 ms / 3,840 次	0
全部 sync 类算子	73,771 次	871 次

_resolve_cu_seqlens 剩下的 60 次 = 10 个 micro-batch × 2（q/kv）× 3 次检查，正好是"每 micro-batch 一次"。

3.2 整步指标

指标	优化前	优化后	变化
profiled step 墙钟	22.73 s	13.36 s	−41.2%
GPU 忙碌率	56.2%	85.1%	+28.9 pt
kernel 总耗时	12.13 s	10.94 s	−9.8%
其中计算 kernel（非 NCCL）	6.37 s	4.88 s	−23.4%
gpu_memset	47.0 ms	5.2 ms	−89%
NCCL	5.76 s	6.06 s	+5%

计算 kernel 少掉的 1.49 s 是 route 构建原先在 GPU 上跑的那批索引小算子——GEMM kernel 的名字、调用次数、耗时前后完全不变，模型算力一点没动。优化后 chunkwise 的计算 kernel（4.88 s）从"比 headwise 多 0.97 s"变成"比它少 0.52 s"；这 0.52 s 同样不是算力差异，而是 headwise 自己的开销：它按 packed 序列逐条做 a2a 再 torch.cat 拼回，光 CatArrayBatchedCopy 就是 262 ms / 15,984 次，chunkwise 优化后的前 8 大 kernel 则全是 GEMM。

NCCL 略升是因为 CPU 不再拖后腿，collective 发得更密，单次 kernel 里等待 peer 的时间变长。

---

4. 端到端性能

4.1 口径

8×H200，Qwen3.5-9B，OpenMathReasoning-mini SFT，TP2 + SP，seed 1234，35 step，取 step 10–34 共 25 step。为让每一步都是训练步，关掉了 eval。每组 25 步的 token 数完全相同（7,270,905），吞吐可直接比。

个别 step 会出现 2–4 倍的偶发 stall（TFLOPs 同步塌陷，说明是同样的工作被拉长的环境噪声）。下表统一剔除 > 1.6 × 中位数的步并给出剔除数量——最终表中所有配置的剔除数均为 0。S3 整组复跑过一次以排除噪声，两轮吻合在 2.3% 以内（chunkwise 48,087 / 46,998，headwise 31,692 / 32,306，all_gather 23,043 / 23,186），表中用的是干净的复跑。

4.2 S1 — Stage 2 的配置（full recompute，max-tokens-per-gpu 8192）

模式	train 时间	train 吞吐	MFU	相对 headwise	mean loss
chunkwise	6.70 s	43,397 tok/s	0.304	1.524	0.35117
all_gather	9.89 s	29,416 tok/s	0.205	1.033	0.35116
headwise	10.21 s	28,483 tok/s	0.199	1.000	0.35115

4.3 S2 — 关掉 recompute（max-tokens-per-gpu 8192）

模式	train 时间	train 吞吐	MFU	相对 headwise	mean loss
chunkwise	5.36 s	54,254 tok/s	0.381	1.499	0.35114
headwise	8.04 s	36,185 tok/s	0.254	1.000	0.35119
all_gather	启动即失败	—	—	—	—

all_gather 在这个场景跑不了：它在每张卡上重放完整序列的 scan，GDN 激活按全局上下文长度增长，Relax 的 _assert_gdn_full_recompute() 因此强制要求 full recompute：

GatedDeltaNet context-parallel (cp>1) requires whole-layer activation recompute:
pass `--recompute-granularity full --recompute-method uniform --recompute-num-layers 1`.
Got recompute_granularity=None.

chunkwise 和 headwise 全程保持 1/cp 分片，所以两者在这个配置下都能跑；差别在速度——chunkwise 快 50%，而且是全部实验里绝对吞吐最高的一组（54,254 tok/s，MFU 0.381）。

4.4 S3 — 拉长 pack（full recompute，max-tokens-per-gpu 24576，即每 micro-batch 全局 98,304 token）

模式	train 时间	train 吞吐	MFU	相对 headwise	mean loss
chunkwise	6.19 s	46,998 tok/s	0.329	1.455	0.35112
headwise	9.00 s	32,306 tok/s	0.225	1.000	0.35119
all_gather	12.54 s	23,186 tok/s	0.162	0.718	0.35113

序列拉长后 all_gather 从"略快于 headwise"掉到"慢 28%"——它重复 CP 倍的 scan，代价随上下文长度线性增长。

4.5 S4 — dynamic CP（--dynamic-context-parallel，max CP=4，full recompute，8192）

CP 逐 micro-batch 在 {1,2,4} 中选择，本轮实际用到 CP1 × 8、CP2 × 41、CP4 × 7 个 micro-batch。

模式	train 时间	train 吞吐	MFU	相对 headwise	mean loss
chunkwise	6.23 s	46,671 tok/s	0.326	1.228	0.35113
headwise	7.65 s	38,004 tok/s	0.267	1.000	0.35114

优势收窄符合预期：CP=1 的 micro-batch 根本不做布局转换，CP=2 时 headwise 的 all-to-all 也便宜得多。这一组同时验证了 route 缓存在 cp_size/cp_rank 逐 micro-batch 变化下的正确失效。

4.6 A/B：收益确实来自本次改动

同一 harness、同一 spec，把 MCore 的两个文件换回 Stage 1 版本（无 route 预构建）再跑一遍：

配置	模式	Stage 1	本 PR	变化
S1（8192）	chunkwise	23,371 tok/s	43,397 tok/s	+85.7%
S1（8192）	headwise	28,124 tok/s	28,483 tok/s	+1.3%
S3（24576）	chunkwise	34,947 tok/s	46,998 tok/s	+34.5%
S3（24576）	headwise	32,090 tok/s	32,306 tok/s	+0.7%

两点交叉验证：

1. Stage 1 在 S1 下 chunkwise/headwise = 0.831，与 Stage 2 §12.3 报的 19,802/23,081 = 0.858 一致，说明这套更短的 harness 复现了原结论；
2. headwise 前后在 ±1.3% 内，收益是 chunkwise 专属的。

---

5. 正确性

• CPU 单测：新增 test_gdn_chunkwise_cp_route.py，10 个测试函数、参数化展开 33 项，分三组：
  • 逐 token 等价（1 个函数，24 项）：对 cp_size ∈ {1,2,4,8} × 两个方向 × 三种 packed 边界（单序列、不等长多序列、含重复边界即空 padding 槽），在本地模拟全体 rank 的 all-to-all，要求 route 产出的每一行都与 Stage 1 未改动的 get_thd_context_parallel_rank_indices 分区完全相同，同时断言相邻 rank 的 split sizes 互相对得上。
  • fail-fast 对拍（3 个）：长度不能被 2×cp 整除、cu_seqlens 不从 0 开始、非单调、非法方向，新旧实现必须同样报 ValueError。
  • 缓存有效性（6 个）：这是 Relax 特有、上游没有对应实现的部分，也是最需要覆盖的——同对象命中、两个方向分开缓存、换 cu_seqlens 对象重建、dynamic CP 换 cp_size/cp_rank 重建、prebuild 填充两个方向、非 THD / CP=1 时 prebuild 为 no-op。
• Stage 1 + Stage 2 回归：test_gdn_chunkwise_cp_layout.py（51）+ test_gdn_cp_mode_stage2.py（19）全过。三个文件合计 103 项全过。
• GPU / NCCL：test_gdn_chunkwise_cp_gpu.py 10 项全过（20 分 36 秒），含真实 CP 组上 zigzag→contiguous→zigzag 的 token 级往返、CP=2 vs CP=1 的 fp32/bf16 前反向对齐、state_dict / sharded_state_dict 一致性。
• 训练 loss：上述 14 组 run 的 mean train/loss 全部落在 0.35111–0.35119，跨模式、跨 recompute、跨 sequence length、跨静态/动态 CP 一致。

---

6. 已知边界

1. #5664 尚未合入上游（open，目标分支 dev，对应 main 的是 #6233）。本 PR 的定位是"提前迁移其核心优化并验证收益"，正确性由"与 Stage 1 已合入实现逐 token 对拍"保证。
2. route 走懒构建而非上游推荐的数据侧 prebuild，prebuild_thd_cp_partition_routes() 已导出，待对象流理清后可直接切换。
3. CUDA graph：本 PR 不支持在 full-iteration graph capture 下运行 THD 布局转换（Megatron 本身也拒绝这个组合）。若将来需要，build_thd_cp_partition_route 的返回类型要换回上游的编码张量形式。
4. #5664 合入后的清理路径：ThdCpPartitionRoute、_thd_cp_partition_route_is_reusable 和 GDN forward 里的粘合代码可整体删除，改用上游的 block-level 调度器 + 数据侧 prebuild。

[xiaoliang0601]

这个正确性验证是如何做的？有跑 >150 个 step 比对 loss、grad norm、reward、mismatch 吗？

[xiaoliang0601]

另外，cu_seqlens、_resolve_thd_cu_seqlens() 作为缓存对象，只检查是不是同一个 Tensor 对象，是不是不够？还应该检查这个 Tensor 对象 有没有被修改过？或者你能判断这个值没有被修改的可能吗？

[jambow0320]

@xiaoliang0601 感谢review，以下是两个问题的回答~

这个正确性验证是如何做的？有跑 >150 个 step 比对 loss、grad norm、reward、mismatch 吗？

端到端逐 step 精确对拍我试过，但是这条路径 run-to-run 本身就不确定（FLA 的 gated delta rule 反向用 atomics，加上 NCCL 规约顺序不固定），同一份配置、同一个 seed 跑两遍，loss 就能差 1.5e-4 ~ 5.3e-4，grad_norm 相对差能到 6% ~ 58%；(主要是grad_norm抖动大，单纯比loss和有效token的话，本 PR chunkwise 对比 Stage 0 老镜像老代码headwise 220步中: loss 逐步绝对差 max 2.08e-4，mean 6.29e-5，中位 5.81e-5，逐步有效token数完全相同)

我主要比较的是：step 0 的 loss 逐 bit 相等。我比了四组：同配置跑两遍、chunkwise 改动前后、chunkwise vs headwise、chunkwise vs all_gather；他们step 0 的 loss 全部相同，grad_norm 相对差 0.01%。

二是单测，现在的单测会涵盖cpu端的route计算（确保我们迁移后的新的route计算和老版本一致），以及gpu端的一次实际的cp chunkwise forward正确性结果检查。

另外，cu_seqlens、_resolve_thd_cu_seqlens() 作为缓存对象，只检查是不是同一个 Tensor 对象，是不是不够？还应该检查这个 Tensor 对象 有没有被修改过？或者你能判断这个值没有被修改的可能吗？

现在这条路径上是不会被静默修改的；Relax 和 MCore 里所有对 cu_seqlens 的原地写全是 torch.zeros 出来立刻 cu[1:] = cumsum(...) 填一次，发生在张量进 PackedSeqParams 之前。但确实可能出现，比如 MCore 有这种写法，mamba_metadata.py:228 的 self._cu_seqlens_buffer[0] = 0 就是常驻 buffer 原地写（为了 CUDA graph 复用缓冲区），这类模式哪天进到 GDN 路径，就会静默拿上一批的 route 去重排这一批 token而且不报错。

刚刚提交了新的commit，新加了一些守卫是：身份之外再比一下 autograd 的版本计数器 tensor._version，避免这个被静默改写不重新构建route。_resolve_thd_cu_seqlens() 同样处理，键改成四个源张量的 (id, _version) 指纹。

这样实现不知道可以不，或者你有什么建议吗？

[xiaoliang0601]

grad_norm 相对差能到 6% ~ 58%

绝对值差多少？如果绝对值太小的话，grad_norm 的相对值意义并不大。

我希望你能贴一个 reward、loss 和 grad norm 的 A/A 和 A/B 对比的曲线（可以用绝对值，也可以取 log），这样直观地判断训练是否有问题。

[xiaoliang0601]

如果 A/B 能落到 A/A 的误差范围内的话，我觉得是没问题的，不需要 bitwise 对齐。

[jambow0320]

grad_norm 相对差能到 6% ~ 58%

绝对值差多少？如果绝对值太小的话，grad_norm 的相对值意义并不大。

我希望你能贴一个 reward、loss 和 grad norm 的 A/A 和 A/B 对比的曲线（可以用绝对值，也可以取 log），这样直观地判断训练是否有问题。

ok，我拿之前存的日志信息画了个图，我认为基本上都是在允许误差范围内的；A是当前pr的chunkwise，B是stage0就是当前Relax main分支headwise跑出来的

[xiaoliang0601]

好，我没什么问题了。