[Bug] DeepSeek V4.1 TP→EP 布局链路未闭合：使用 ReduceScatter / AllGather 对接标准 EP MoE

[high-cloud]

背景

DeepSeek V4.1 Flash 的 TP4/DP2/EP8 链路尚未闭合：Attention 输出使用 TP all-reduce，MoE ABI 依赖 token_owners / tp_rank 对复制的 token 去重，原始检查提交中的 MoE kernel 仍是占位实现（以下静态复现描述该历史提交，不代表当前 kernel 状态）。本 Issue 的目标是闭合跨层 sequence-parallel（SP）布局：本地 mHC-pre + Norm → TP AllGather → head-TP Attention / output partial → TP ReduceScatter → 本地 mHC / 标准 EP MoE → 下一层本地 mHC-pre + Norm。MoE 只处理各 rank 自己持有的 token，不再承担 Attention TP 副本去重。

本 Issue 是基于源码确认的布局集成缺口，不是已在设备上复现的精度错误。AllReduce + 正确的唯一 owner 方案本身可以正确；问题在于当前链路尚未实现，且接口把 Attention TP 语义带入了 MoE。如果所有 TP 副本直接进入标准 EP dispatch，就会重复派发同一批 token。

责任边界：pypto-lib 模型与 orchestration，没有证据指向编译器或 runtime。

检查环境与代码依据

• pypto-lib：1c594b27cb9936aa1a5000bb31e6a5f74efa8ff0，main，工作区干净；已确认该提交在目标 GitHub 仓库可访问。
• 主机：Linux (aarch64)。目标部署：单机 8 卡 A5，TP4/DP2/EP8。
• 验证方式：Python AST / 源码检查；未编译、未运行 NPU、未进行性能测试。pypto / simpler / PTOAS / pto-isa / CANN 版本及 pin 一致性不作为本次静态检查的结论依据，未检测。

固定提交下的证据：

• TP/EP 配置：DP_SIZE = EP_SIZE // TP_SIZE。
• TP 输入 fixture：x 在 TP 组内复制。
• C1A Full 调用 all-reduce。
• TP all-reduce 实现：每个 rank 得到完整 token 的规约结果。
• MoE 接口与占位实现：包含 token_owners、tp_rank，函数体抛出 NotImplementedError。
• mHC 接口：residual stream 为 [T,4,5120]。

静态复现

在上述提交的仓库根目录执行以下脚本，只依赖 Python 标准库：

import ast
from pathlib import Path

root = Path("models/deepseek_v4_1_flash")
module = ast.parse((root / "moe.py").read_text())
moe = next(n for n in module.body if isinstance(n, ast.FunctionDef) and n.name == "moe")
source = (root / "prefill_c1a_full.py").read_text()
print("prefill calls TP all-reduce:", "        prefill_tp_output_all_reduce(" in source)
print("moe requires token_owners:", any(a.arg == "token_owners" for a in moe.args.args))
print("moe body:", ast.unparse(moe.body[0]))

当前实际输出：

prefill calls TP all-reduce: True
moe requires token_owners: True
moe body: raise NotImplementedError('EP MoE kernel body is assigned independently')

期望修复

1. 从模型入口开始，mHC residual stream 在 TP 组内沿 token 维分片，并跨层保持该布局。每层先在本地完成 Attention mHC pre + norm，再对 [T_local,D] 输入执行 TP AllGather；同一 TP 组的 Attention 输入因此保持相同 token、相同顺序，head / output-group 权重仍按 TP 分片。
2. 在 FP32 output projection partial 上执行沿 token 维的 ReduceScatter(SUM)，替换该模型集成路径上的 all-reduce。不能先 all-reduce 再对相同副本做 SUM reduce-scatter，否则会将值放大 TP 倍。
3. residual / mHC 状态始终保留本地 token 行，与 ReduceScatter 的输出范围一致，不做 SUM；在 token 分片上完成 Attention mHC pre/post、MoE mHC pre/post 和逐行 normalization。兼容原有复制布局的 wrapper 如需切行，应显式处理，不将复制 residual 送入 SUM collective。
4. MoE 使用标准 EP 接口：每卡持有不同的 token；dispatch 发送到专家 owner；combine 将结果返回 token 来源 rank，并按 top-k 权重汇总。MoE 不需要感知 Attention TP，也不需要 TP token 去重。
5. Shared expert 只计算本 rank 的 token，不重复计算全部 TP 副本。
6. MoE mHC post 后不做 residual-stream AllGather，也不对已完整 combine 的 MoE 输出追加 TP AllReduce。直接将 [T_local,4,D] residual 和对应 mHC 状态传给下一层；下一层完成本地 mHC pre + norm 后，才 AllGather [T_local,D] attention 输入。
7. 覆盖 prefill/decode 及支持的 Attention 模式。调整模型集成接口、golden 和文档；独立 Attention 测试如保留完整输出接口，可以显式提供 all-gather wrapper。

模型入口、出口与 token 映射

• Embedding 权重完整复制：每个 TP rank 按统一的连续 token 范围本地切分 token IDs，再查表得到 [T_local,D]；如果各卡已持有相同 IDs，切分本身不需要通信。
• Embedding 按词表 TP 分片：每卡对完整 token batch 产生本地词表贡献，非本卡词表位置为零；对这些 partial 执行 token 维 ReduceScatter(SUM)，直接得到本地完整 embedding。AllReduce 后本地切片可作为正确性基线，但不是目标通信路径。只有 embedding 的有效贡献参与求和，不能对已复制的完整 embedding 再做 SUM。
• 在本地 embedding 上构造 mHC residual streams；首次进入主干时即为 [T_local,4,D]。入口策略由实际 embedding 权重布局决定，不要求为了 SP 复制整张词表。
• 切分对象是本步 scheduled token 行，不是请求的永久 TP 归属；同一请求的 prefill / draft rows 可以跨 TP rank。每一步允许重新确定切分，但本步 token IDs、有效行 mask、residual、mHC 状态以及其他逐 token 输入必须共享一致映射。
• head-TP Attention 使用 AllGather 后的全组 token 顺序及与之对齐的 positions / 请求和 cache 元数据；MoE 只使用本地 IDs 和有效行。不能将仅本地 positions 直接配给已 AllGather 的 hidden。
• token 数不能整除 TP 时补齐；padding 不参与有效路由、shared expert 结果或最终输出。覆盖 token 数小于 TP、空 rank 以及不同 DP 组有效行数不等的情况。
• 主干出口继续在本地完成可逐 token 执行的 hc_head / norm；根据 LM-head、sampling 和 serving ABI 明确最终 gather / redistribution 边界，恢复原始有效 token 顺序。不要为了适配出口接口而在每层收集四路 residual。

设计范围与参考

实现 纯 head TP + SP，保留已有 attention / KV-cache / indexer 的 head-TP 语义。

参考 vllm-ascend 36b7582431326361dd4dae79a78a555e76e715f9：

• V4.1 decoder layer：本地 mHC-pre / norm 后 sp_all_gather，attention 后 sp_reduce_scatter，MoE 使用 already_sequence_parallel=True。
• SP helpers：沿 token 维 gather / reduce-scatter 和 padding。
• MoE prepare/finalize：SP 输入下跳过普通复制布局的切分与出口 AllGather。

这些是布局设计参考，不构成本仓库实现已经通过设备验证或性能提升的证据。

TP4/EP8 多 batch shape 示例

以下是期望的数据流，不是现有实现已通过验证的行为。模型参数：D=5120、64 个 attention heads、head dim 512、4 个 mHC streams、384 个 routed experts、top-k=6、expert intermediate 2304。MoE TP=1，每卡持有 48 个完整 routed experts。

每个 DP 组有 3 条请求，长度为 3、2、3；packed 后 T=8。以下 shape 均为每卡逻辑 shape，不包含量化 scale、物理 padding 或通信 buffer 容量。

分组	物理 ranks	请求	Packed token 顺序
TP 组 0 / DP 组 0	0、1、2、3	A、B、C	A0 A1 A2 B0 B1 C0 C1 C2
TP 组 1 / DP 组 1	4、5、6、7	D、E、F	D0 D1 D2 E0 E1 F0 F1 F2
EP 组	0～7	全部 6 条请求	16 个唯一 token

两个 DP 组处理不同请求。层入口 residual 在同 TP 组内按 token 分片；仅在 Attention mHC pre + norm 后 AllGather，使 Attention 看到相同的完整 token 输入。

环节	每卡输入 shape	每卡输出 shape	数据含义 / 通信
1. 层入口 mHC residual stream	[2,4,5120]	[2,4,5120]	从 embedding 或上一层继承本地 token 分片
2. Attention mHC pre + norm → TP AllGather	[2,4,5120]	本地 [2,5120] → 全组 [8,5120]	本地合并 streams / norm，只 gather attention 输入
3. Q 投影及 head 切分	[8,5120]	[8,16,512]	每卡计算 16 个 heads
4. Attention	Q [8,16,512] + KV cache	[8,16,512]	每卡自己的 heads；cache 长度由各请求历史决定
5. 分组 output projection	[8,16,512]	[8,5120] partial	每卡负责 2 个 output groups；输出待跨 TP 求和
6. TP ReduceScatter(SUM)	[8,5120] partial	[2,5120]	求和 4 卡贡献，同时沿 token 维切分
7. 保留本地 residual / Attention mHC 系数	residual [2,4,5120]；post [2,4]；mix [2,4,4]	相同本地 shape	与第 6 步 token 行一致；不求和、不通信
8. Attention mHC post	attention [2,5120] + 第 7 步状态	[2,4,5120]	得到本 rank 两个 token 的新 residual stream
9. MoE mHC pre + norm	[2,4,5120]	[2,5120]	生成本地 MoE 输入，保留对应 mHC 状态
10. Gate / top-k	[2,5120]	scores [2,384]；IDs / weights [2,6]	每卡产生 12 条 routed assignments
11. EP dispatch	[2,5120] + routes [2,6]	[Rᵣ,5120]	Rᵣ 为本卡专家实际收到的 assignment 数
12. 本地 routed experts	每个 expert [nₑ,5120]	中间 [nₑ,2304]；输出 [nₑ,5120]	每卡 48 个完整 experts；Σnₑ=Rᵣ
13. EP combine	expert 输出 [Rᵣ,5120]	[2,5120]	返回来源 rank，汇总每个 token 的 6 个加权结果
14. Shared expert 分支	第 9 步输入 [2,5120]	[2,5120]	仅计算本地 token；可与 routed 分支并行
15. Routed + shared	两份 [2,5120]	[2,5120]	完整 MoE 输出
16. MoE mHC post	MoE 输出 + residual [2,4,5120] + 系数	[2,4,5120]	本地完成 stream 更新
17. 跨层传递本地状态	[2,4,5120]	[2,4,5120]	无 TP 通信；下一层回到第 2 步，gather [2,5120]

按连续 token 切分时：

TP rank	DP 组 0	DP 组 1	每卡 MoE 输入
0	rank 0：A0、A1	rank 4：D0、D1	[2,5120]
1	rank 1：A2、B0	rank 5：D2、E0	[2,5120]
2	rank 2：B1、C0	rank 6：E1、F0	[2,5120]
3	rank 3：C1、C2	rank 7：F1、F2	[2,5120]

• token 分片允许跨请求边界，MoE 不要求同一请求的 token 位于同一张卡。
• EP 全局只有 16×6=96 条有效 assignments；ΣRᵣ=96，平均每卡 12 条，但实际接收量由路由决定。若所有 TP 副本重复 dispatch，则变成 384 条。
• AllGather 延迟到下一层 Attention mHC pre + norm 后，收集 [2,5120] 输入，不收集 [2,4,5120] residual stream。相同 dtype 下，此次 gather 的元素数/字节数为 residual-stream gather 的 1/4；实际字节数还取决于各边界 dtype。此比例不是整层通信量或延迟的提升比例，必须测量端到端耗时。

验收标准

• 模型入口根据 embedding 权重布局产生唯一 token 分片；本地 mHC pre + norm 后 AllGather 的 Attention 输入在同 TP 组内相同，DP 组之间可以是不同请求和不同有效 token 数。
• ReduceScatter 结果等价于 TP partial 求和后取对应 token 分片；保持 FP32 累加和既定 dtype/舍入边界。
• residual / mHC 的 token 映射与 ReduceScatter 完全一致，不对复制的 residual 做 SUM。
• MoE ABI 不依赖 Attention TP 的 token_owners / tp_rank；EP combine 正确返回来源 rank。
• 上述案例只有 16 个唯一 token 和 96 条有效 routed assignments；shared expert 每个唯一 token 仅计算一次。
• 跨层 mHC residual / pre-mix 等状态保持 token 分片；本地 mHC pre + norm 后的 AllGather 顺序和数值与完整参考一致，重建后的 residual streams 也与完整参考一致。
• MoE 出口不追加 TP AllReduce 或 residual-stream AllGather；shared expert 权重布局保证本地输出已完整，不能遗漏其所需的归约。
• embedding 到至少两个连续 decoder layers 再到输出边界的 token 映射闭合，覆盖 input IDs、positions、有效行 mask 和最终输出顺序。
• 覆盖 ragged 多请求、不同 DP token 数、token 数不能整除 TP、decode token 数小于 TP，以及空分片；padding 不参与有效路由或 shared expert 输出。
• 覆盖 prefill/decode 的相关 Attention 路径，完成 TP4/EP8 多卡精度验证。
• 对比原复制布局基线，记录 attention-input AllGather、output ReduceScatter、EP dispatch/combine 以及入口/出口布局转换的通信字节数和端到端耗时；分别覆盖小 batch decode、目标 DSpark 验证 batch 和 prefill。
• 更新模型文档、测试和集成 ABI。

关联

• [功能] 打通 DeepSeek-V4.1-Flash 纯文本 TP4/DP2/EP8 推理 #1205：总体 TP4/DP2/EP8 bring-up；本 Issue 将其中允许的 owner / reduce-scatter 两种边界方案明确收敛为后者。
• [Feature] 提取 V4.1 Attention 公共前处理与输出投影，消除重复 Q latent 计算 #1271：Attention 公共前处理与 output projection 提取。该 Issue 原先保留外层 all-reduce；实现本 Issue 时需协调修改模型集成路径上的通信方式，公共 output projection 仍输出 FP32 partial。