Files
CTI-Inference-Opt/代码/code
OwnerSunshine530 c1d8b91fb2 feat(Phase B): FlexAttention 块对角注意力 + MoE 稠密向量化
- scaled_dot_product 分发:block_mask->FlexAttention(每用户仅自身序列内因果,
  避免对~14000长拼接序列做O(S²)稠密注意力);否则SDPA稠密(回退/对照)。
- CTRModel.build_block_mask 构造块对角因果mask;_use_flex 在SM80+自动启用。
- SMoE 稠密向量化(einsum批量算所有expert后按top-k gather),消除Python循环/同步;
  保留 _smoe_forward_loop 作数值等价对照。CONFIG.vectorize_moe 可切。
- load_model 加可选 torch.compile。
- tests/test_equiv.py:MoE稠密vs循环、Flex vs稠密SDPA 数值等价(无pytest依赖)。
- bench.py 加 --attn/--moe/--compile 便于A800上对比测速。

需 A800(SM80) 实测;CPU/V100 自动回退 SDPA。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-14 23:30:59 +08:00
..