[Medium][模型适配] --enable-shared-expert-tp 时标准门控模型的 topk 未随 hidden_states 做 TP 切分，MoE 分发/通信量被高估 tp_size 倍

run_id: repro-moe-topk-mismatch-20260831

Problem

ParallelMoELayer.forward 的 enable_shared_expert_tp 分支（tensor_cast/layers/moe_layer.py，route_after_dp_transform=False 默认路径）中，route() 在 _dp_transform_enter() 之前对全量 token [N, H] 计算路由：

• raw_logits 门控分支（gate_returns_raw_logits=True，如 deepseek_v3）在 route() 内部把 router_logits 切到本 rank 分片，其注释明确该切片用于对齐 _dp_transform_enter() 随后对 hidden_states 的 TP 切片；
• 标准门控分支（gate 返回 (topk_indices, topk_weights) tuple 或裸 logits，即 gate_returns_raw_logits=False 且无 gate_router，如 deepseek_v32）不做任何切片。

随后 _dp_transform_enter() 把 hidden_states 切到 [N/tp, H]，而 topk_indices 仍为 [N, k]。FusedMoETensorCast.forward 以 num_tokens = topk_indices.numel()（= N*k，应为 N*k/tp）驱动全部下游尺寸（init_routing_v2 输出行数、all_to_all 分发量、每专家 GEMM M 维、出口 all_gather/unpermute token 量），全部膨胀 tp_size 倍。末尾截断逻辑使最终输出形状恰好正确——无崩溃、无告警，纯静默的性能高估。

术语：N = 本 DP rank 打包后总 token 数；tp = 全局 tensor parallel size；k（top_k）= 每 token 路由的专家数（DeepSeek-V3.2 为 8）。

Reproduction（纯官方 CLI，两条命令，无任何脚本/探针）

设计说明

本 bug 由官方 CLI 公开参数组合自然触发，且异常直接体现在 CLI 自带的 analytic metrics 表（stdout）中，无需任何辅助工具。触发组与对照组仅差一个 flag，其余（模型、N、tp、ep、层数）完全一致。模型使用仓库自带 config-only 资产 tests/assets/model_config/deepseek_v32，全程离线。

复现步骤

# Step 1: 触发组（bug 分支）
uv run python -m cli.inference.text_generate `
    tests/assets/model_config/deepseek_v32 `
    --num-queries 4 --query-length 128 --device TEST_DEVICE `
    --num-devices 2 --tp-size 2 --ep-size 2 `
    --enable-shared-expert-tp `
    --num-hidden-layers-override 4 --no-repetition

# Step 2: 对照组（同命令，仅去掉 --enable-shared-expert-tp）
uv run python -m cli.inference.text_generate `
    tests/assets/model_config/deepseek_v32 `
    --num-queries 4 --query-length 128 --device TEST_DEVICE `
    --num-devices 2 --tp-size 2 --ep-size 2 `
    --num-hidden-layers-override 4 --no-repetition

Step 1 预期：MoE 分发链路（all_to_all / unpermute_tokens / init_routing_v2）时延与对照组一致（该 flag 只影响共享专家的 TP 化，不改变 routed expert 的分发量，见下方数值推理）。
Step 1 实际：三项全部约 2× 于对照组（all_to_all 2.999ms vs 1.501ms），总时延 +12.7%。
Step 2 预期/实际：正常完成，作为健康基线。
判定：必现（确定性；两组均 exit=0、无告警——静默）。

Test Results（stock CLI stdout，commit 7a07630）

从两组 metrics 表摘取关键行（完整表见各自 stdout）：

MoE 分发链路（时延 ∝ 路由 token 数）：

与路由 token 数无关的算子（权重/激活带宽型，应不变）：

汇总：

数值推理（正确值 2048、实际 4096 的推导）

全部输入来自 CLI 参数与模型 config，不需要观测任何内部张量：

1. 账本：N = 4 queries × 128 tokens = 512（stdout "Number of Queries per DP rank: 4"）；k = num_experts_per_tok = 8；--num-devices 2 --tp-size 2 → world=2，dp = 2/(tp×pp) = 1。
2. TP 语义：tp=2 时每 rank 只负责 N/tp = 256 个 token 的计算（生产代码 _dp_transform_enter() 即做此切片，出口 all_gather 重组）。因此每 rank 应分发的 (token, expert) 对 = 256 × 8 = 2048。
3. 对照组标定：对照组 all_to_all = 1.501ms ↔ 2048 对（解析模型时延随消息量线性，两组 ep/dtype/hidden_size 完全一致）。对照组的 attention/量化/共享权重算子均为 1.00×，是健康基线。
4. 触发组实测：all_to_all = 2.999ms = 1.998× 对照组 → 消息量翻倍 → 4096 对 = N×k——恰等于 tp=1（不做 TP 切分）时的量。即：开了 TP，MoE 分发量却和不切分一样。unpermute（1.965×）与 init_routing_v2（1.836×）两个独立算子给出同源印证。
5. 选择性指纹：只有时延随路由 token 数变化的三个算子集体 ≈2×，与 token 数无关的算子全部 1.00×。若该 flag 合法地改变了计算语义，不会呈现这种精确选择性的翻倍。
6. 守恒：dp=1 时两个 rank 持有完全相同的 512 个 token，TP 语义是各负责一半。触发组每 rank 按全量 512 个 token 分发路由 → 同一份路由工作全局执行了两遍。
7. 合法代价核算（推断）：该 flag 的合法语义是共享专家 TP 化（CLI help："dense-MLP TP for shared_experts with delayed down_proj reduction"），其合理代价是共享专家输出引入的 1 次 all_reduce——实测 all_reduce 恰好 7→8 次、per-call 均值两组完全相同（377.091μs），合法增量 ≈ 0.38ms。而触发组总增量 2.52ms 中，仅 all_to_all + unpermute + init_routing 三项就多出 1.58ms（合法代价的 4.2 倍）。

Expected

--enable-shared-expert-tp 开启时，标准门控模型每 rank 的 MoE 分发量应为 (N/tp)×k 对（本例 2048）：触发组的 all_to_all / unpermute_tokens / init_routing_v2 应与对照组一致（约 1.5ms / 58μs / 33μs），总时延回到约 19.8ms 量级。

Actual

每 rank 分发量为 N×k 对（本例 4096，= tp_size × 正确值），MoE 分发/通信量被高估 tp_size 倍：总时延 +12.7%，吞吐被低估 11.3%（tp=2 实测；tp 更大时线性放大）。输出形状靠截断修正，无任何告警。

Root Cause

tensor_cast/layers/moe_layer.py 三处叠加：

# ParallelMoELayer.forward — enable_shared_expert_tp 分支，route_after_dp_transform=False（默认）
topk_indices, topk_weights = self._inner.route(
    hidden_states, ...                       # hidden_states: [N, H] 全量（切片前）
)                                             # topk_indices: [N, k] ← 标准门控分支不做 TP 切片
hidden_states, num_tokens = self._dp_transform_enter(hidden_states)
#                                              hidden_states: [N/tp, H]（global_tp_group.slice）

# FusedMoETensorCast.forward
num_tokens = topk_indices.numel()             # = N*k，应为 N*k/tp（仅由 topk 尺寸决定）

MoELayer.route 内部仅 raw_logits 分支（deepseek_v3 等模型）内置 pad + tensor_split(tp_size)[tp_rank] 切片，注释明确该切片用于对齐 _dp_transform_enter() 的 hidden_states 切片；标准门控分支（gate 返回 tuple / 裸 logits）无对应逻辑。

同仓正确对照（修复时可参考）：

• route_after_dp_transform=True（kimi_k3/kimi_k25 设置）：在 _dp_transform_enter() 之后路由，天然一致；
• enable_shared_expert_tp=False 路径：MoELayer.forward 内部先 enter 后 route，一致，且有回归测试断言（test_parallel_moe_ep_route_before_tp_slice_smoke：fused_moe 入参 ((3,16),(3,2)) 即 hidden 与 topk 的 token 数必须一致）。

测试盲区：tests/regression/tensor_cast/test_parallel_moe.py 的两个 shared_expert_tp 测试均使用 _make_fake_gate(shard_by_tp=True)（假门控自行模拟 raw_logits 式 TP 切片），标准门控模型在该分支上零覆盖。

Fix 建议

在 route_after_dp_transform=False 路径中，对 route() 返回的 topk_indices / topk_weights 执行与 raw_logits 分支相同的 pad + tensor_split(tp_size)[tp_rank] 切片（或统一在 route() 内置切片，使两个门控分支行为一致）。验收标准（即上述两条 CLI 命令）：触发组的 all_to_all / unpermute_tokens / init_routing_v2 与对照组一致，总时延回到对照组量级；并补一条标准门控（shard_by_tp=False）的 shared_expert_tp 回归测试。

Environment

• OS: Windows（PowerShell）
• Python: 3.11.15（uv 管理）/ torch 2.10.0+cpu / transformers 5.13.1
• 模型: tests/assets/model_config/deepseek_v32（仓库自带 config-only 资产，离线）
• 仓库 commit: 7a07630156dc164827c3fcb56e763e34887a492d（2026-08-29 15:29:13 +0800）

Impact / Severity

• Medium：性能仿真数字错误（静默高估），非功能崩溃。
• 受影响组合：gate_returns_raw_logits=False 且无 gate_router 的 MoE 模型（deepseek_v32、qwen3_moe、ling、ernie4_5_moe、minimax_m2 等）在 --ep-size >1 --tp-size >1 --enable-shared-expert-tp 下。DeepSeek-V3.2 是该特性的旗舰用户（共享专家 + EP 部署）。
• 不受影响：deepseek_v3（raw_logits 分支内置切片）、kimi_k3/kimi_k25（route_after_dp=True）、不开该 flag 的所有模型。
• 规避方式：不使用 --enable-shared-expert-tp（但该 flag 正是 shared-expert TP 部署形态的建模入口，规避即丢失目标场景）。

相似 Issue 检索

已检索现有 issue 关键字（shared_expert_tp / topk / MoE / DP transform），未发现重复。

SIG 路由

tensor_cast/layers/ → 模型适配 SIG（chair ChenHuiwen，reviewers stormchasingg / jhon-117 / weixin_43113933，approver lutean）。