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>

代码/code/bench.py

@@ -250,6 +250,11 @@ def _parse_args():
                     help="逗号分隔的 keep_fp32_modules，如 linear,rep_encoder.input_norm")
     ap.add_argument("--feasign-none", action="store_true",
                     help="不截断特征（max_feasign_per_slot=None）")
+    ap.add_argument("--attn", choices=["auto", "flex", "sdpa"], default=None,
+                    help="注意力实现：flex=块对角FlexAttention, sdpa=稠密(原), auto=SM80自动")
+    ap.add_argument("--moe", choices=["dense", "loop"], default=None,
+                    help="MoE实现：dense=向量化(新), loop=逐expert循环(原)")
+    ap.add_argument("--compile", action="store_true", help="开启 torch.compile")
     ap.add_argument("--rebuild", action="store_true", help="强制重建过滤缓存")
     return ap.parse_args()
 
@@ -273,5 +278,11 @@ if __name__ == "__main__":
         cfg["merge_threshold"] = a.merge_th
     if a.keep is not None:
         cfg["keep_fp32_modules"] = tuple(x for x in a.keep.split(",") if x)
+    if a.attn is not None:
+        cfg["use_flex_attn"] = {"auto": "auto", "flex": True, "sdpa": False}[a.attn]
+    if a.moe is not None:
+        cfg["vectorize_moe"] = (a.moe == "dense")
+    if a.compile:
+        cfg["compile"] = True
     mf = None if a.feasign_none else {1: 2}
     run_once(cfg, batch_size=a.bs, max_batches=a.smoke, max_feasign_per_slot=mf, rebuild=a.rebuild)