[Operator Mechanism] Add MoH routing to DSV4 CSA indexer - #1740
[Operator Mechanism] Add MoH routing to DSV4 CSA indexer#1740wangna11BD wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
| 序号 | 位置 | 优先级 | 状态 |
|---|---|---|---|
| 1 | src/paddlefleet/transformer/csa_attention.py:2049 |  |
✅ |
| 2 | src/paddlefleet/transformer/transformer_config.py:2212 | |
✅ |
| 3 | PR #1740 标题 |  |
✅ |
P1（decode 路径不复用 MoH）已接受作者说明：MoH 仅用于 prefill/训练，decode 阶段有意关闭，且已完成对应验证。
P1（MoH head 数约束）已解决：__post_init__ 按 csa_indexer_backend 校验 TileLang/cuDNN 的有效 head 数，并补充边界回归测试。
P2 已解决：标题已更新为 [Operator Mechanism] Add MoH routing to DSV4 CSA indexer。
| weights, _ = self.linear_weights_proj(x) # [b, sq, n_heads] | ||
| weights = weights * (self.index_n_heads**-0.5) | ||
|
|
||
| if self.use_moh: |
There was a problem hiding this comment.
forward_before_topk 中裁剪 q 和 weights；decode 路径 _compute_indexer_compressed_topk_idxs_decode（同一文件约 2515-2543 行）仍直接用全部 index_n_heads，也没有 indexer_moh_bias 路由。因此启用 MoH 后，prefill 与后续增量生成会对同一 token 使用不同的 indexer head 集合，生成结果发生阶段性分叉。请让 decode 复用相同的 MoH 选择/计数逻辑，并补充 prefill→decode 一致性回归测试。
There was a problem hiding this comment.
moh结构在decode阶段收益会劣化，所以仅在prefill/训练使用，已经完成仅在prefill/训练使用moh的验证，训练指标无损，训练速度和推理速度均有提升。
There was a problem hiding this comment.
了解，感谢说明。既然 decode 阶段有意不启用 MoH，且已完成仅在 prefill/训练使用 MoH 的验证，我接受这一设计边界，撤回此前关于 decode 路径必须复用 MoH 的 P1；不再要求 prefill→decode 使用同一组 head。
| "use_moh=True requires num_activated_heads as a positive " | ||
| f"integer, got {self.num_activated_heads!r}." | ||
| ) | ||
| if self.num_activated_heads > self.dsa_index_n_heads: |
There was a problem hiding this comment.
1 <= num_activated_heads <= dsa_index_n_heads，但 MoH 输出随后会传入 fused indexer：TileLang 要求 head 数不超过 64 且为 8 的倍数，cuDNN 更只接受 32/64（csa_indexer_fwd 的输入校验）。例如默认 TileLang 下 num_activated_heads=17 会保留 17 个 head、不会触发 16-padding，首次 forward 即因 head 数不可整除 8 抛错；cuDNN 下 num_activated_heads=8 padding 到 16 也会被拒绝。请按所选 backend 校验或统一 pad 到 kernel 支持的 head 数，并覆盖这些配置组合。
There was a problem hiding this comment.
已在当前 head 核对：post_init 现在按 csa_indexer_backend 校验 MoH 的有效 head 数（TileLang 不超过 64 且为 8 的倍数，cuDNN 为 32/64），并补充了对应边界回归测试。该 P1 已解决，感谢修复。
wangna11BD
changed the title
2 participants
PR Category
Operator Mechanism
PR Types
New features
Description
为DSV4 CSA index增加moh结构
是否引起精度变化
否