[Docs] PyPTO 性能优化案例：DeepSeek-V4 MoE 预处理算子融合与同步开销复盘

[Little-oil]

Summary

• Fuse split_pre_post, mix_x, and ffn_norm into one external AIV task.
• Replace the first 48-AIV hard-sync implementation with two indexed 8-AIV
  GM soft barriers, while preserving sync_start=True.
• Reuse the fused normalization outputs in the remaining gate pipeline.
• Move non-critical comb_sinkhorn behind dispatch_gather.
• Add selective dumps, phase-stop debug entries, generated-code audits, and
  synchronization/ABI contract tests.
• Size local routed experts as global / 16 * EP for the updated partition.

阅读导航（ADHD-friendly）

这部分故意写成可以复用的“算子融合操作手册”。不用一次全部读完：

1. 只想知道改了什么：读“0. 一屏结论”。
2. 只想审核正确性：读“2. 依赖怎么改”和“4. 同步怎么选”。
3. 准备融合下一组算子：直接照“8. 可复用逐步清单”执行。
4. 正在定位错误：跳到“6. dump/debug 决策表”。

---

0. 一屏结论

融合前

hc_pre_rms ─────┐
                ├─> split_pre_post ─> mix_x ─> ffn_norm ─> gate/dispatch
hc_pre_linear ──┘          │
                           └─> post

hc_pre_rms + hc_pre_linear ─> comb_sinkhorn ───────────────> hc_post

三个小 AIV 算子在关键路径上串行。任务边界原本提供正确的完成顺序，但也带来
三次 launch 和两个 Scheduler gap。

融合后

hc_pre_rms_tid ───┐
                  ├─explicit deps─> fused_pre_norm（8 AIV）
hc_pre_linear_tid ┘                   │
                                      ├─ split_pre_post
                                      ├─ soft barrier #1
                                      ├─ mix_x
                                      ├─ soft barrier #2
                                      └─ ffn_norm
                                             │
                                             v
                                      gate -> dispatch_gather
                                                   │
                                                   v
                                            comb_sinkhorn -> hc_post

必须同时成立的 7 个不变量

1. Python core_num、C++ grid stride、soft barrier participant 数和
   workspace slot 数描述的是同一组 8 个逻辑 AIV。
2. 即使某个 lane 在当前 phase 没有业务计算，也必须到达两次 barrier。
3. participant ID 必须来自 get_block_idx(args) 的 runtime 逻辑编号，不能使用
   裸 get_block_idx() 的物理核编号。
4. 必须保留 sync_start=True，保证 8 个会忙等的 participant 全部驻留后再启动。
5. 每个 fused grid 必须拥有独立、零初始化的 InOut[64, INT32] workspace。
6. soft barrier 只同步 counter；pre_val_store 和 x_mixed 还需要单独做
   DCCI publish/invalidate + DSB。
7. 三个 generated body 的数学、tiling、valid shape、cast 和业务输出顺序不能改变。

---

1. 如何分析融合前的算子

1.1 先画数据流，不要先看“哪些函数写在一起”

对每个候选 scope 记录：

• core 类型；
• logical work 数；
• 精确的 GM 输入/输出；
• 哪个输出被下一个 scope 读取；
• producer 和 consumer 是否采用相同的切分方式；
• 是否在 latency 关键路径；
• 是否包含通信、atomic 或其他副作用。

本例 decode T=8 的关键表格：

Scope	主要输入	主要输出	decode logical work	角色
hc_pre_rms	x_flat	inv_rms	1	外部 producer
hc_pre_linear	x_flat, hc_fn	mixes_raw	4 split-K	外部 producer
split_pre_post	inv_rms, mixes_raw, base/scale	pre_val_store, post	1	关键路径
mix_x	pre_val_store, x_flat	x_mixed	4	关键路径
ffn_norm	x_mixed, norm_w	gate normalization buffers	0/8	关键路径
comb_sinkhorn	inv_rms, mixes_raw	comb	1	非关键支路

选择 split_pre_post + mix_x + ffn_norm 的原因：

• 三者都是纯 AIV；
• 由真实数据依赖严格串行；
• 每个 kernel 都很小，launch/Scheduler gap 占比高；
• 不包含分布式通信；
• 都已有经过 golden 验证的 PyPTO/PTOAS generated body。

1.2 找到“重新切分”的数据边

删除 task boundary 后，Scheduler 不再替 phase 提供完成顺序。只有在 consumer
可能读取其他 lane 产生的数据时，才需要 grid barrier。

本例有两条这样的边：

split_pre_post --pre_val_store--> mix_x
mix_x          --x_mixed-------> ffn_norm

第一条边：

• 一个 split logical block 写连续 8 x 8 FP32；
• 同一个 token tile 会被四个 mix logical block 读取；
• 四个 mix block 各负责一个 1024-wide hidden slice；
• decode 时 lane 0 生产，lane 0..3 消费。

第二条边：

• 一个 mix logical block 给 8 行各写一个 1024-wide BF16 slice；
• 四个 mix block共同组成每行完整的 4096 元素；
• ffn 的 lane 0..7 各自读取一整行 [1, 4096]。

所以单个 lane 内的程序顺序不够，必须有两个跨 lane barrier。

1.3 为什么不把更多算子一起融合

没有融合 hc_pre_rms 和 hc_pre_linear：

• 两者本来可以并行；
• linear 有 split-K matmul/atomic 行为；
• 融入后会混合不同执行特性并扩大同步问题；
• 保留两者 TaskId，可以给 fused task 一个清晰的显式起点。

没有融合 comb_sinkhorn：

• gate/dispatch 启动不需要它；
• 20 次 Sinkhorn 迭代会占用 AIV；
• 提前运行会与 fused/gate/dispatch 抢资源；
• 它只需要在 hc_post 之前完成。

通用原则：

融合能消除关键 launch gap 的最小串行区域；不要因为“代码相邻”就吞掉并行
producer 或非关键支路。

---

2. 依赖图如何修改

只写 fused kernel 不够。必须在 Scheduler graph 中重新表达全部 happens-before。

2.1 把外部 producer TaskId 显式返回

hc_pre_moe_producers() 只保留 unfused RMS/linear，并捕获：

with pl.spmd(...) as hc_pre_rms_tid:
    ...

with pl.spmd(...) as hc_pre_linear_tid:
    ...

return hc_pre_rms_tid, hc_pre_linear_tid

fused submit 显式依赖它们：

deps=[hc_pre_rms_tid, hc_pre_linear_tid]

不要在文档或测试中写死数值 TaskId；运行期 TaskId 会随生成结果变化，只检查符号边。

2.2 从 gate 中拆出 ffn_norm

旧 gate() 中的 ffn_norm 产生：

xg_buf
inv_rms_buf
xn_scale_buf
x_norm_scale

融合步骤：

1. 把 standalone ffn_norm 作为 fused 的第三个 generated body。
2. 把以上四个 buffer 变成 fused 的显式业务输出。
3. 剩余 gate 改为 gate_precomputed()，直接消费这些 buffer。
4. 把 fused_pre_norm_tid 作为 norm_tid 传给 gate。
5. x_norm_quant、gate_pre_route、gate 使用 deps=[norm_tid]。

这样 standalone ffn_norm 仍可单测，MoE 路径则复用 fused 结果。

2.3 把 comb_sinkhorn 延后到 dispatch_gather

dispatch() 现在返回 _gather_tid。comb_sinkhorn 改成：

with pl.spmd(
    t_dim // COMB_T_TILE,
    name_hint="comb_sinkhorn",
    deps=[dispatch_gather_tid],
    allow_early_resolve=True,
):
    ob = pl.tile.get_block_idx()

为什么使用 with pl.spmd：

• for-form 主要暴露 block index；
• with ... as tid 表达整个 grid 的 TaskId 和显式 deps；
• body 内通过 pl.tile.get_block_idx() 取得 block index。

延后为什么安全：

hc_pre_rms/linear
  -> fused_pre_norm
  -> gate
  -> dispatch_push/wait/gather
  -> comb_sinkhorn

dispatch_gather 是 producer 的传递后继。gather 完成时，inv_rms/mixes_raw
必然已经完成，所以 comb 此时读取它们是安全的。

2.4 为什么需要 manual_dep=True

显式 deps 会与自动依赖合并。如果 inv_rms/mixes_raw 仍由 OverlapMap 自动建边，
comb 的实际 fanin 会重新变成：

dispatch_gather + hc_pre_rms + hc_pre_linear

这会把想迁移的旧依赖加回来。生产路径因此对这两个跨显式边界的 buffer 使用
manual_dep=True：

• fused submit 用两个 producer TaskId 恢复正确性边；
• comb 只保留 dispatch_gather_tid；
• codegen audit 检查最终生成结果，而不只相信 Python 源码。

manual_dep=True 不是通用性能开关。只有能够给出完整的传递顺序证明时才可使用。

2.5 sync_start 和 allow_early_resolve 是两件事

sync_start=True：

• 是 fused grid 内 barrier 的正确性要求；
• 要求 8 blocks 作为一个 cohort 全部集结后启动；
• 防止部分 lane 忙等并占满资源，导致剩余 lane 永远无法调度。

allow_early_resolve=True：

• 是 producer 对 Scheduler 的性能提示；
• 允许后继在 producer 完成前预装/解析，但不允许提前读未完成数据；
• RMS/linear 保留它，fused 才可能在两个直接 producer 发布后提前准备；
• fused 保留它，gate/x_norm 后继才可能提前准备；
• 它不能替代 sync_start。

依赖修改后的检查项：

• fused 是否显式依赖所有外部 producer？
• 每个 fused 输出的 consumer 是否依赖 fused TaskId 或对应 Tensor RAW？
• 被屏蔽的自动依赖是否有明确的传递替代？
• comb_sinkhorn 是否只有 dispatch_gather 一条显式 dep？
• 生成 orchestration 是否与源码设计一致？

---

3. 如何构造 fused external kernel

3.1 复用 external CCE bridge 的结构

整体结构参考 external-kernel PR（例如
#765）：

Python @pl.jit.extern 声明
        ↓
entry.cpp ABI 入口
        ↓
独立 C++ kernel body
        ↓
model 中 pl.spmd_submit

这里复用的是 bridge/ABI 分层，不是 #765 的 attention 算法，也不照搬它的同步模板。

3.2 generated body 是数学实现的 source of truth

三个 standalone PyPTO scope 分别编译成：

split_pre_post_generated.hpp
mix_x_generated.hpp
ffn_norm_generated.hpp

做法：

1. 单独编译原 scope 并取得 PTOAS 生成的 C++ body。
2. 每个 body 放进独立 namespace/header。
3. 把原 entry 调整成可调用函数，显式接收 GM pointer/scalar。
4. 显式传入原始 logical_block 和原 phase 的 logical work 数。
5. fused_body.hpp 只负责 ABI、循环、同步和调用顺序。

generated instruction sequence 尽量保持原样。不要手工“简化”：

• cast/rounding；
• valid shape/padding；
• pipeline；
• TASSIGN 的 UB 地址；
• generated tail barrier。

standalone scope 改变后，应重新生成对应 header，并重新跑 UB high-water/contract。

generated body 末尾的 PIPE_ALL 只 drain 本核流水，不是跨核 barrier，也不能保证
其他核看到业务 GM 的新值。

3.3 Python/C++ ABI 必须一一对应

PyPTO 会把全部 Tensor 参数放在 scalar 参数之前。production 使用 13 个 Tensor：

Index	参数	Direction
0	x_mixed	Out
1	x_flat	Input
2	inv_rms	Input
3	mixes_raw	Input
4	hc_base	Input
5	norm_w	Input
6	pre_val_store	Out
7	post	Out
8	xg_buf	Out
9	ffn_inv_rms_buf	Out
10	xn_scale_buf	Out
11	x_norm_scale	Out
12	sync_workspace	InOut

随后是：

scale0, scale1, num_tokens

debug ABI 最后再追加 stop_after。

C++ 侧锁定：

kSyncWorkspace = 12
kTensorArgCount = 13
kScale0 = kTensorArgCount

workspace 是内部同步状态，不是业务输出，所以返回 tuple 保持七项。随意改成八项会
移动 alias/解包顺序，给所有下游制造 ABI 风险。

3.4 保留每个 phase 原来的 logical mapping

wrapper 计算：

split_work = tokens / 8;
mix_work   = split_work * 4;
ffn_work   = clamp(round_up(num_tokens, 16), tokens);

三个 phase 均使用：

for (int32_t logical_block = lane;
     logical_block < phase_work;
     logical_block += kAivLanes) {
  generated_body(..., logical_block, phase_work);
}

decode T=8 的 work 是 (1, 4, 8)。三个 phase 串行，所以参与核数是：

max(1, 4, 8) = 8

不是 1 + 4 + 8 = 13。

prefill T=128 可能是 (16, 64, 128)。8 lane 通过 grid-stride 多轮执行，因此
正确性仍成立；但性能最优 participant 数需要单独调优。

3.5 Python submit

((seven_business_outputs), fused_pre_norm_tid) = pl.spmd_submit(
    self.fused_pre_norm_cce,
    ...,
    sync_workspace,
    scale0,
    scale1,
    num_tokens_index,
    core_num=FUSED_AIV_CORES,  # 8
    sync_start=True,
    deps=[hc_pre_rms_tid, hc_pre_linear_tid],
    allow_early_resolve=True,
)

if False 中的 direct extern call 仅用于 materialize
self.fused_pre_norm_cce specialization，最终会被删除。

---

4. 如何选择同步

4.1 通用决策顺序

每删除一条 task edge，都按顺序问：

1. consumer 是否会读到其他 lane 产生的数据？
   • 否：不需要 grid barrier。
   • 是：继续。
2. 原 task boundary 是否提供了 producer-complete 顺序？
   • 是：融合后必须在 kernel 内补回。
3. 串行区域实际需要多少 participant？
   • 取各 phase 最大并行 work，不取总和。
4. busy-wait participant 能否全部同时驻留？
   • 需要 sync_start=True。
5. barrier 是否同时保证业务数据 cache 可见？
   • 本例不保证，必须单独 publish/acquire。

4.2 第一版：48-AIV 硬同步

第一版 correctness baseline 使用：

AscendC::SyncAll<>();

并启动 48 个 AIV。它证明了：

• 三个 generated body 可以在同一 entry 中运行；
• 两个 task boundary 可由两个 kernel barrier 替代；
• 所有 lane 采用一致控制流。

但它不是最终性能方案：

• decode 最多只需要 8 个 AIV；
• hard sync 的同步集合是全部 48 个物理 AIV；
• fused 启动前 Scheduler drain 要收集 48 个资源；
• 上游 hc_pre_rms 和后继 gate 也需要 AIV；
• 全占用降低 early-dispatch/overlap；
• 把 launch 改成 8 却保留全 AIV hard barrier 会挂死。

本 kernel 是 AIV-only，不能直接照抄 mixed-core kernel 的
SyncAll<false>() 模板；同步类型必须从实际 core topology 推导。

4.3 最终版：indexed 8-AIV soft barrier

pto::SYNCALL_SOFT_AIV_BARRIER(
    gm_workspace,
    ub_workspace,
    kAivLanes,        // 8
    participant_idx  // logical 0..7
);

为什么不用当前 public 三参数 soft API：

• simpler 可能把逻辑 block 0..7 动态调度到任意物理 AIV；
• 裸 get_block_idx() 返回物理编号，可能稀疏或不从 0 开始；
• 用物理编号索引 [0,8) workspace 会越界或永远等不到某些 slot；
• get_block_idx(args) 返回 dense runtime logical ID。

所以统一使用：

const int32_t lane = static_cast<int32_t>(get_block_idx(args));

同一个 lane 同时用于：

• 三个 grid-stride loop；
• 两次 indexed soft barrier；
• producer/consumer ownership range。

长期应给 PTO-ISA 增加公共四参数 indexed overload；本 PR 使用 simpler 当前 pin
中已有的 internal helper。

4.4 为什么绝不能删除 sync_start=True

GM soft barrier 会忙等并持续占用 AIV。如果 grid 只启动一部分：

1. 先启动的 lane 到达 barrier；
2. 它们占住 AIV 等待未启动 lane；
3. 剩余 lane 因无可用资源无法 dispatch；
4. 形成死锁/饥饿。

sync_start=True 负责先集结完整的 8-lane cohort，再允许 kernel 运行。

soft sync 的作用是把同步集合从 48 缩到 8；它不替代 Scheduler 的同时驻留保证。

4.5 GM workspace 和 UB scratch

PTO helper 的 slot ABI 是每个 participant 8 个 INT32：

8 participants * 8 INT32 = 64 INT32 = 256 B

每个 fused grid 创建：

sync_workspace = pl.create_tensor(
    [FUSED_SOFT_SYNC_WORDS],
    dtype=pl.INT32,
    init_value=0,
)

必须满足：

• 每个可能并发的 grid 独占一份；
• 初始值为 0；
• 声明为 InOut；
• 不能用 module-global workspace；
• 不能用 pl.no_dep/manual_dep 隐藏它。

generation：

初始：    每个 slot head = 0
barrier1：每个 slot head = 1
barrier2：每个 slot head = 2

共享 workspace 会让不同 grid 互相“代打卡”。

UB scratch 固定在 176 KiB：

generated body 当前 UB high-water：< 65 KiB
soft-sync scratch：[176 KiB, 176 KiB + 256 B)
PTO temp reserved：[184 KiB, 192 KiB)

static/contract 检查：

• 64B 对齐；
• scratch end 不进入 PTO temp；
• 扫描 generated header 的 TASSIGN，确保所有 tile end 低于 176 KiB。

注意两个不同概念：

• PTO sync slot stride 当前固定为 8 INT32 = 32B；
• A2/A3 业务数据 DCCI 按 64B cache line。

不能为了“对齐 64B”只在 wrapper 中擅自把 slot stride 改为 16，否则会和 helper
寻址/Python ABI 不一致。这个 32B slot/64B cache-line 差异需要靠 A2/A3 并发压力
测试持续覆盖。

4.6 barrier 到达不等于业务 GM 可见

soft barrier 自带的 DCCI 只维护同步 counter，不会刷新 pre_val_store/x_mixed。

正确顺序：

generated producer TSTORE 完成
  -> pipe_barrier(PIPE_ALL)
  -> producer 对自己拥有的业务 cache line 做 CACHELINE_OUT
  -> dsb(DSB_DDR)
  -> indexed soft barrier
  -> consumer invalidate 自己将要读取的业务 cache line
  -> dsb(DSB_DDR)
  -> generated consumer TLOAD

范围：

• pre_val_store：一个 split block 写 8*8*4=256B，即 4 条 64B line；
• x_mixed：一个 mix block 在 8 个不连续 row 中分别写 2048B slice；
• ffn consumer invalidate 一整行 4096*2=8192B。

x_mixed 的 8 个 slice 不是连续区域，必须逐 row 处理；不能把它误当成一段连续
16KiB，否则会触碰其他 lane 拥有的 slice，增加 false sharing。

当前 DCCI 是 correctness-first，可能比较贵。必须先通过真机 golden/压力测试，再
基于泳道证据缩小或合并范围。

---

5. InOut、specialization 与 dump

5.1 workspace 为什么是 InOut

• Input 会隐藏 native write；
• Out 会隐藏旧 generation read；
• InOut 才会生成 add_inout(sync_workspace)。

workspace 必须放在全部 Tensor 之后、scalar 之前；TaskArgs 不允许 scalar 后再放
Tensor。

5.2 if False 为什么使用另一份 workspace

PyPTO 会在 constant-false elimination 前检查 direct extern call。一个变量传入
InOut 后，原 pre-call SSA value 就被消费。

如果 specialization direct call 和真实 spmd_submit 共用同一 workspace，会触发：

InOutUseDiscipline

因此分支中单独创建：

_sync_workspace_specialize = pl.create_tensor(
    [FUSED_SOFT_SYNC_WORDS],
    dtype=pl.INT32,
    init_value=0,
)

它随 if False 一起消除。codegen audit 会拒绝最终 orchestration 中残留的
specialization-only buffer。

5.3 正确的 dump 位置

pl.dump_tag 是 forward-sticky 标记。只在真实 submit 前执行一次：

pl.dump_tag(sync_workspace)

runtime 会保存：

• BEFORE_DISPATCH；
• AFTER_COMPLETION。

submit 后不能再次引用旧的 workspace SSA value，否则违反 InOut use discipline。
七个业务输出在 submit 后单独 dump_tag。

---

6. dump/debug 定位方法

6.1 phase-stop debug entry

--debug-stop	执行范围
split	只执行 split，停在 barrier1 前
barrier1	split + barrier1
mix	split + barrier1 + mix，停在 barrier2 前
barrier2	执行到 barrier2 后
full	执行完整三个 phase

示例：

python models/deepseek/v4-flash/fused_pre_norm_cce.py \
  -p a2a3 -d 0 \
  --mode decode \
  --num-tokens 8 \
  --debug-stop barrier1 \
  --dump-args 1

6.2 workspace dump 期望值

8 个 slot 的 head 位于：

0, 8, 16, 24, 32, 40, 48, 56

期望：

before dispatch：全部为 0
after barrier1：8 个 slot head 为 1，其余 padding 为 0
after barrier2/full：8 个 slot head 为 2，其余 padding 为 0

6.3 症状 -> 第一检查点

症状	第一检查点
barrier1 前后挂住	core_num==8、kAivLanes==8、sync_start=True、logical ID
某个 slot 永远不到 1	lane 提前 return、ID 错、workspace 共享/过小
barrier1 counter 正常但 mix_x 错	pre_val_store publish/invalidate
mix stop 正确但 full 错	x_mixed 可见性或 ffn logical mapping
输出整体错位/alias	Python/C++ Tensor ABI 和七项返回解包
只在重复/并发时失败	workspace 跨 grid 共享或 cache 可见性遗漏
compile 报 InOutUseDiscipline	specialization 共用 workspace 或 submit 后读旧 InOut
comb 太早运行	manual_dep 缺失或生成 deps 又含 rms/linear
偶发无 hang 但结果错	不要依赖 helper poll timeout；participant/驻留必须本身正确

---

7. 分层验证

Level 1：contract/static

python -m pytest -q tests/contract/test_deepseek_v4_moe_fusion.py

contract 锁定：

• 13-Tensor Python/C++ ABI 和七项业务返回；
• 8 blocks、sync_start=true；
• 恰好两次 indexed soft barrier；
• 禁止 AscendC::SyncAll、隐式 public hard sync、FFTS sync、裸物理
  get_block_idx()；
• phase 顺序与 grid-stride；
• DCCI publish/acquire 顺序；
• UB scratch 边界；
• 每次 submit 独立、全零的 InOut workspace；
• dump placement；
• delayed comb dependency；
• 生成 orchestration 的 alias/参数顺序。

Level 2：compile-only

分别编译：

• production entry；
• debug entry 的所有 stop specialization；
• decode/prefill；
• 完整 MoE orchestration。

检查生成 C++：

block_num = 8
require_sync_start = true
workspace = [64] INT32, initial value = 0
add_inout(sync_workspace)
fused dump 包含 sync_workspace
comb_sinkhorn 只依赖 dispatch_gather

Level 3：单 kernel 真机 golden

至少覆盖：

decode： num_tokens = 0, 1, 8
prefill：num_tokens = 0, 1, 128

0/1 很重要：业务 loop 可以为空/很小，但 8 个 lane 仍必须完整通过 barrier。

python models/deepseek/v4-flash/fused_pre_norm_cce.py \
  -p a2a3 -d 0 \
  --mode decode \
  --num-tokens 8 \
  --dump-args 1 \
  --enable-l2-swimlane 4

Level 4：debug stop + counter

按顺序跑：

split -> barrier1 -> mix -> barrier2 -> full

同时检查业务输出和 workspace generation。

Level 5：完整分布式 MoE

至少两 rank、balanced routing，检查：

• 最终 x_next golden；
• codegen audit；
• 多轮无随机 stale read；
• fused 只有 8 个 AIV block；
• 剩余 AIV 能运行其他任务。

Level 6：并发/泳道/性能

需要覆盖：

• fresh workspace 多轮 0 -> 1 -> 2；
• 多个独立 fused grid；
• 与长 AIV task overlap；
• 不同 rank 落到不同物理 AIV 集合；
• 48-hard vs 8-soft 的完整 step 对比。

Perfetto/泳道重点看：

• rms -> fused、linear -> fused；
• fused -> x_norm_quant/gate；
• dispatch_gather -> comb，且 comb fanin 不应再出现 rms/linear；
• fused 是否正好 8 条 AIV worker；
• drain 规模是否从 48 降到 8；
• producer finish -> fused start gap；
• fused finish -> gate start gap；
• comb 是否移出 producer -> dispatch 关键窗口。

不要跨不同 profiling level/输入比较绝对时间。CPU sim 也不能证明动态物理核映射
和真实 DCCI 行为。

---

8. 可复用逐步清单

Phase A：证明融合边界

• 画出当前 task DAG。
• 为每个候选算子列 input/output/core/work-count 表。
• 标记真正关键路径。
• 标记所有跨 lane 重切分的数据边。
• 排除可并行 producer、通信和非关键支路。
• 保存 standalone golden/dump。

Phase B：准备 Scheduler 依赖

• 捕获全部外部 producer TaskId。
• fused submit 显式依赖所有 producer。
• 把下游拆成“接受 precomputed output”的 consumer。
• 非关键支路延后前，写出完整传递顺序证明。
• 只在确实要替换自动边时使用 manual_dep=True。
• 对生成 orchestration 增加依赖 audit。

Phase C：构建 external kernel

• 单独编译并保存每个 standalone generated body。
• 每个 body 使用独立 namespace/header。
• 在 Python/C++ 定义完全一致的 Tensor-before-scalar ABI。
• 保留业务输出顺序。
• 保留每个 phase 的 logical work 数和 mapping。
• 用 grid-stride 解耦 resource cap 与大 shape work count。
• 增加 production entry 和 phase-stop debug entry。

Phase D：选择同步

• 对每条被删除的 task edge 判断是否需要跨 lane barrier。
• participant 数取串行 phase 的最大并发 work，不取总和。
• 动态调度下使用 runtime logical ID。
• busy-wait barrier 保留 sync_start=True。
• 每个 grid 创建独立全零 workspace。
• native 写 workspace 必须建模成 InOut。
• 为 UB scratch 选不重叠区间并加静态检查。
• 单独处理业务 GM publish/invalidate。

Phase E：让问题可定位

• extern 前后 dump 所有业务输入/输出。
• submit 前 dump sync workspace。
• 每个 barrier 前后都能 debug-stop。
• 明确定义 counter 期望值。
• 对 partial-valid 输出使用 poison-aware compare。

Phase F：逐级验收

• contract/static。
• production/debug compile-only。
• decode 边界。
• prefill/grid-stride。
• debug-stop dump。
• 完整分布式 golden。
• 多轮/并发 A2/A3 压力。
• 最后才优化 DCCI、participant 数或 UB 复用。

---

9. 常见错误

1. 把串行 phase 的核数相加：(1,4,8) 需要 8，不是 13。
2. launch 减为 8，却保留全 AIV hard barrier：可能直接挂死。
3. 用裸 get_block_idx()：物理 ID 不保证是 dense 0..N-1。
4. 删除 sync_start=True：partial grid 可占住全部资源等待未调度 lane。
5. 让空闲 lane 提前 return：它仍欠两次 barrier arrival。
6. 认为 soft barrier 自动 flush 业务 Tensor：counter 与业务数据可见性是两件事。
7. 多个 grid 共用 workspace：generation 会互相提前放行。
8. workspace 声明为 Input：runtime 看不到 native write。
9. specialization/真实 submit 共用 InOut：会被 PyPTO 线性使用检查拒绝。
10. 只改变源码顺序，不检查生成 deps：Scheduler graph 可能完全不是预期。
11. 手改 generated 数学指令：应该从 standalone scope 重新生成。
12. 只测 num_tokens=max：空/tail work 才容易暴露 barrier bug。
13. 只看 fused kernel 自身时长：还要看 drain、后继启动、完整 step 和整体 AIV
    利用率。
14. 依赖 soft helper timeout 兜底：正确性必须来自一致 participant、fresh
    workspace 和 sync_start，不能依赖超时后行为。

主要文件

• models/deepseek/v4-flash/hc_pre.py
  • standalone source scopes；
  • RMS/linear producer TaskId；
  • delayed comb_sinkhorn。
• models/deepseek/v4-flash/gate.py
  • standalone ffn_norm；
  • gate_precomputed consumer。
• models/deepseek/v4-flash/fused_pre_norm_cce.py
  • Python extern ABI；
  • golden/debug/dump harness。
• models/deepseek/v4-flash/kernels/fused_pre_norm_cce/
  • production/debug entry；
  • 三个 generated body；
  • fused wrapper、soft barrier、GM visibility。
• models/deepseek/v4-flash/moe.py
  • production submit；
  • 依赖重写；
  • delayed comb；
  • generated-code audit。
• tests/contract/test_deepseek_v4_moe_fusion.py
  • 全部结构正确性不变量。

Verification performed

• tests/contract/test_deepseek_v4_moe_fusion.py: 14 passed.
• Production/debug/decode/prefill/full-MoE compile and link passed.
• A2/A3 decode golden passed for num_tokens=0/1/8.
• A2/A3 prefill golden passed for num_tokens=0/1/128.
• Debug barrier1/barrier2 outputs and workspace generations passed.
• Two-rank balanced-routing MoE passed final x_next golden and ABI audit.
• Full workspace dump observed slot-head generations 0 -> 1 -> 2.

Follow-ups

• 给 PTO-ISA 增加公共 indexed soft-barrier overload，替代 internal helper。
• 在完整 A2/A3 压力测试后优化 DCCI 重复范围。
• decode/prefill 分别调 participant cap；8 是 decode 的资源/正确性选择，prefill
  可能需要更大 cap。

[coderabbitai]

📝 Walkthrough

Walkthrough

The MoE pipeline now computes normalization intermediates during preprocessing, passes them into a dedicated gating entry point, and produces quantized routing inputs for dispatch. Expert sizing derived from --ep is also updated.

Changes

Deferred MoE normalization and gating

Layer / File(s)	Summary
MoE preprocessing and normalization models/deepseek/v4-flash/hc_pre.py	Adds MoE/FFN sizing constants and introduces hc_pre_moe_norm, which computes mixed outputs, combination data, and normalization buffers.
Deferred gate routing models/deepseek/v4-flash/gate.py	Adds gate_from_norm and refactors gate to produce normalization intermediates before delegating quantization and routing.
MoE pipeline integration models/deepseek/v4-flash/moe.py	Updates expert sizing and wires the new preprocessing and gating functions into dispatch input production.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant moe
  participant hc_pre_moe_norm
  participant gate_from_norm
  participant dispatch
  moe->>hc_pre_moe_norm: preprocess tokens and produce norm buffers
  moe->>gate_from_norm: route using norm buffers
  gate_from_norm->>dispatch: provide x_norm_i8, x_norm_scale, indices, weights

Loading

Possibly related PRs

• hw-native-sys/pypto-lib#784: Implements a related deferred-RMSNorm gating refactor using shared normalization intermediates.
• hw-native-sys/pypto-lib#790: Reworks related RMSNorm and INT8 quantization intermediates in the gating path.

Poem

I hop through norm buffers, tidy and bright,
Then guide tiny tokens to routers just right.
With scales in my basket and weights in a row,
The dispatch path follows wherever they go.
Squeak! Deferred gating makes pipelines glow.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title describes the MoE preprocessing fusion and sync-overhead review, which matches the core change set.
Description check	✅ Passed	The description is detailed and clearly aligned with the fused preprocessing, gating reuse, and expert sizing changes.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}