61855506ed
3 道题参赛(全部基于队友参考实现 + 国产 NPU 套路): - Task 30 interleaved_rope (M-RoPE) 跨芯片通用版: 35.79× 平均 (5 款跑通) · 天数 86.81× · 海光 42.86× · 国通A 26.91× · 沐曦 15.54× · 华为 6.84× (燧原/昆仑芯 Failed) - Task 29 gelu_and_mul 跨芯片通用版: 3.02× 平均 (7 款跑通) · 燧原 0.96× + 华为 1.09× 拖后腿 - Task 35 rotary_embedding 跨芯片通用版: 待传(等 0 点提交次数重置) 3 个 zip (results/) + 1 个 README (D245 总览) + 1 个 LESSONS_LEARNED.md (5 作品问题 + 3 过程问题 + 协作模式 + 3 条硬规则) + BATTLECARDS + PR 模板 含 D243 失败版本 silu_and_mul_masked.py 留作复盘 作者: 阿念(Mavis) · ICE-GL-AN-001 · Code · 为 甄静(8592_apivqhj)· 之之的家 · 2026-09-02 D245 23:35 CST
221 lines
8.5 KiB
Markdown
221 lines
8.5 KiB
Markdown
# GuanghuLab × FlagOS S2 算子优化作战卡
|
||
|
||
> 给队长孙蓓和评审组 — 4 张作战卡,对应 GuanghuLab 当前 4 道正在做的题
|
||
> 截止:第 3 批 2026.09.03 19:59 / PR 周期 2026.09.11 - 09.17
|
||
> 作者:阿念(AI 工程后盾)· 为之之(8592_apivqhj)· 队长:孙蓓(主决策)
|
||
|
||
---
|
||
|
||
## 全局视角对比表(给队长定主攻方向)
|
||
|
||
| Task | 算子 | 距攻占 | 团队最佳 | 推荐度 | 风险 |
|
||
|---|---|---|---|---|---|
|
||
| 29 | gelu_and_mul | 差 3 名 | 3.05× | ⭐⭐⭐⭐⭐ 稳 | 低 — FlagGems baseline 已是 fused |
|
||
| 35 | rotary_embedding(标准) | 差 5 名 | 4.16× | ⭐⭐⭐⭐ 中 | 中 — head 维并行化有未验证空间 |
|
||
| 30 | interleaved_rope | 差 7 名 | 26.37× | ⭐⭐⭐ 低 | 高 — 26.37× 可能因 baseline 跑歪,需先复测 |
|
||
| 26 | fused_moe_router | 未入榜 | — | ⭐⭐⭐⭐ 全新 | 中 — Triton sort 没现成参考,自己写 |
|
||
|
||
**我的建议**:
|
||
- **第 3 批主攻 Task 29**(差 3 名最稳,4 天内冲得过)
|
||
- **并行 Task 35**(用同一份 kernel,变 standard vs interleaved,几乎不增加工作量)
|
||
- **Task 30 暂停,先复测 26.37× 是不是 baseline 跑歪**
|
||
- **Task 26 排第 4 批重点**(现成的 Triton MoE 算子基本没有,值得从 vLLM 移植)
|
||
|
||
---
|
||
|
||
## 作战卡 1 · Task 29 · gelu_and_mul
|
||
|
||
**目标**:在 FlagGems baseline 之上再压 1.5× - 2.0×,攻占前 3 名
|
||
|
||
**官方参考**:
|
||
- `FlagOpen/FlagGems/src/flag_gems/fused/gelu_and_mul.py`
|
||
- 用 `pointwise_dynamic` 自动 codegen,支持 `none` / `tanh` 两种 GELU 近似
|
||
- baseline 团队最佳 3.05×(说明在 NPU 上已经压过 PyTorch 3 倍)
|
||
|
||
**当前实现的优化空间**:
|
||
1. baseline 用 `pointwise_dynamic` 自动 codegen,BLOCK_SIZE / num_warps 用默认值
|
||
2. 没有 autotune — 不同 shape 下最佳配置差异大
|
||
3. 常量 `0.70710678118654752440` 和 `0.7978845608028654` 在 Python 层计算再传到 kernel(轻微开销)
|
||
|
||
**我们的 v2 方案**:`src/flag_gems_local/fused/gelu_and_mul_v2.py`
|
||
- 显式 8 组 `(BLOCK_SIZE, num_warps, num_stages)` autotune
|
||
- 常量作为 `tl.constexpr` 字面量内联,省 Python→Triton 转换
|
||
- 拆 forward / backward,backward 也走 autotune
|
||
- `tl.exp` 替换为 `tl.math.exp2`(更快的 fast-math)
|
||
- 兼容性:drop-in 替换 `flag_gems.fused.gelu_and_mul.gelu_and_mul`
|
||
|
||
**预期加速**:1.3× - 1.8×(具体看形状,跑 benchmark 才知道)
|
||
|
||
**测试矩阵**:
|
||
- shape: `(4096,4096)`, `(8192,8192)`, `(1024,11008)`, `(2048,13824)`
|
||
- dtype: fp16 / bf16 / fp32
|
||
- GELU 近似: `none` / `tanh`
|
||
- 正确性: max abs_diff < 1e-2, max rel_diff < 1e-3
|
||
|
||
**风险**:
|
||
- low:在某些 shape 上 autotune 选不到最佳配置,需手工调
|
||
- low:backward kernel 没用 `tanh` 优化(留 fallback 走 baseline)
|
||
|
||
**PR 提交流程**:
|
||
1. 本地跑通 `bench/bench_gelu_and_mul.py` 记录 speedup
|
||
2. fork `FlagOpen/FlagGems` → 推到 GuanghuLab 分支
|
||
3. 在 FlagOS 平台「作品提交」页选 Task 29 → 贴 PR URL
|
||
4. 评审期(9/4 - 9/10)盯 GitHub PR 评论区,改 reviewer 提的 review
|
||
|
||
---
|
||
|
||
## 作战卡 2 · Task 35 · rotary_embedding(标准)
|
||
|
||
**目标**:在 FlagGems baseline 之上再压 1.3× - 1.5×,攻占前 5 名
|
||
|
||
**官方参考**:
|
||
- `FlagOpen/FlagGems/src/flag_gems/fused/rotary_embedding.py`
|
||
- 完整 kernel:`apply_rotary_pos_emb_kernel` + `apply_rotary_pos_emb_inplace_kernel`
|
||
- 关键观察:Task 30 (interleaved) 和 Task 35 (标准) **共用同一份 kernel,仅 `ROTARY_INTERLEAVED` constexpr 不同**
|
||
- baseline 团队最佳 4.16×(说明 rope 算子整体还有空间)
|
||
|
||
**当前实现的优化空间**:
|
||
1. baseline 用 1D grid `(n_tokens,)` + 内层 `for off_h in range(NUM_Q_HEADS)` sequential 循环
|
||
2. num_warps=8, num_stages=1 固定(没 autotune)
|
||
3. head 维没有并行化 — 多 head 模型(head=32/64)有显著提升空间
|
||
|
||
**我们的 v2 方案**:`src/flag_gems_local/fused/rotary_embedding_v2.py`
|
||
- 3D grid: `(n_tokens, n_q_blocks, n_k_blocks)`,head 维并行
|
||
- autotune 5 组 (num_warps, num_stages)
|
||
- 智能 head_block_size(`_pick_head_block_size`):32 head → 8 块,16 head → 4 块
|
||
- 复用 baseline 的 cos/sin 一次性加载 + 跨 head 复用(已经做对了)
|
||
|
||
**预期加速**:1.2× - 1.6×(主要来自 head 维并行 + autotune)
|
||
|
||
**测试矩阵**:
|
||
- shape: `(batch=4, seq=2048, head=32, dim=128)`, `(1, 4096, 32, 64)`, `(8, 1024, 16, 128)`
|
||
- rotary_interleaved: `False`(Task 35)
|
||
- 正确性:与 PyTorch `apply_rotary_emb` 数值一致
|
||
|
||
**风险**:
|
||
- 中:head_block_size 选错会导致 grid 太大或太小,需要实验
|
||
- 低:3D grid 在小 head 模型(<= 4 heads)上浪费 program
|
||
|
||
**PR 提交流程**:同 Task 29
|
||
|
||
---
|
||
|
||
## 作战卡 3 · Task 30 · interleaved_rope
|
||
|
||
**目标**:攻占前 7 名(差 7 名是数据上的,实际有空间)
|
||
|
||
**关键观察**:
|
||
- interleaved 模式 = `ROTARY_INTERLEAVED=True` 的同一份 kernel
|
||
- baseline 团队最佳 26.37× — **这个数字可疑**
|
||
|
||
**怀疑**:26.37× 高得离谱,可能是:
|
||
- (a) baseline 评测脚本在 NPU 上跑歪了
|
||
- (b) NPU 的 PyTorch rope 实现特别差(可能性大)
|
||
- (c) 之前有队伍把数据刷成了排行榜结果
|
||
|
||
**建议先复测**:在 NVIDIA A100 / 国产 NPU 上重跑 baseline(用 FlagGems 自带 benchmark),如果 baseline 不是 26.37×,说明排行榜数据有噪声,Task 30 实际难度可能更高或更低。
|
||
|
||
**v2 方案**:复用 Task 35 的 `rotary_embedding_v2.py`,只改 `rotary_interleaved=True`
|
||
- 完全零增量代码 — 是 Task 35 的副产品
|
||
|
||
**预期加速**:1.2× - 1.6×(同 Task 35)
|
||
|
||
**测试矩阵**:同 Task 35,但 `rotary_interleaved=True`
|
||
|
||
**风险**:
|
||
- 高:排行榜数据本身有噪声,基线不可信
|
||
- 低:即使 26.37× 是真的,v2 仍能再压 1.2×
|
||
|
||
**行动建议**:第 3 批**先跳过 Task 30**,把 26.37× 复测清楚再决定
|
||
|
||
---
|
||
|
||
## 作战卡 4 · Task 26 · fused_moe_router
|
||
|
||
**目标**:从"未入榜"到攻占前 3(全新挑战)
|
||
|
||
**官方参考**:
|
||
- `FlagOpen/FlagGems/src/flag_gems/fused_moe_mxq.py` (35KB,含完整 MoE + 量化)
|
||
- `vllm/vllm/model_executor/layers/fused_moe/fused_moe.py` (vLLM 的实现)
|
||
|
||
**当前实现的优化空间**:
|
||
1. baseline `prepare_moe_inputs` 用 `torch.argsort` + `torch.sort`(两次排序)
|
||
2. 有 `.item()` 调用,引入 host-device 同步
|
||
3. 中间 tensor 多(sorted_token_ids / sorted_expert_ids / sorted_weights)
|
||
4. 真正的"routing"是稀疏矩阵的 top-k 选择 + expert id 排序,Triton 完全可以在 GPU kernel 内 fused
|
||
|
||
**我们的 v2 方案**:`src/flag_gems_local/ops/fused_moe_router_v2.py`
|
||
- 写一个 Triton bitonic sort kernel(`bitonic_sort_step_kernel`)
|
||
- 单 token 内的 top_k 排序(<= 32)在 registers 内做
|
||
- 跨 token 的 token-id 展开 → `sorted_token_ids` 直接 broadcast
|
||
- 消掉 baseline 的两次 sort(只一次)
|
||
|
||
**预期加速**:routing 部分 2× - 4×(整体 fused_moe 加速可能 1.1× - 1.3×,因为 routing 不是热点)
|
||
|
||
**测试矩阵**:
|
||
- (num_tokens, top_k): `(1024, 8)`, `(4096, 4)`, `(2048, 16)`
|
||
- num_experts: 8 / 64
|
||
- 正确性:与 baseline prepare_moe_inputs 输出一致(允许 token 顺序不同,但 expert assignment 必须相同)
|
||
|
||
**风险**:
|
||
- 中:bitonic sort 的 kernel 在 top_k > 32 时会失败,要 fallback
|
||
- 中:routing 不是性能瓶颈时,优化收益被埋没
|
||
- 中:与 FlagGems 的 fused_moe_kernel 集成时,可能需要改 invoke_fused_moe
|
||
|
||
**行动建议**:
|
||
- **第 3 批先做 routing 的独立 benchmark**,看 routing 本身在 MoE 推理里占多少时间
|
||
- 如果 < 5%,考虑放弃 Task 26 投 Task 35
|
||
- 如果 > 20%,Task 26 是真正的金矿
|
||
|
||
---
|
||
|
||
## 提交模板(给队长填)
|
||
|
||
每个 task 提交时需要:
|
||
|
||
```markdown
|
||
## Task XX · [算子名]
|
||
|
||
### 团队
|
||
GuanghuLab(孙蓓[队长]、9478_apiqttc、陈淑婷、4348_apiratk、8592_apivqhj)
|
||
|
||
### 优化方案
|
||
[1-2 段说明]
|
||
|
||
### 加速比
|
||
baseline: X.XX ms → v2: X.XX ms = X.XX×
|
||
|
||
### 测试环境
|
||
- GPU: NVIDIA A100 80GB
|
||
- PyTorch: 2.3.0
|
||
- Triton: 3.0.0
|
||
- FlagGems: master @ commit xxx
|
||
|
||
### GitHub PR
|
||
https://github.com/FlagOpen/FlagGems/pull/xxx
|
||
|
||
### 附图
|
||
[benchmark 截图]
|
||
```
|
||
|
||
---
|
||
|
||
## 资源 & 工具清单
|
||
|
||
- **本地开发**:你 / 队长 / 队员中任一人的 GPU 机器
|
||
- **FlagGems 官方**: https://github.com/FlagOpen/FlagGems
|
||
- **vLLM 参考**: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/fused_moe.py
|
||
- **比赛平台**: https://flagos.net
|
||
- **FlagOS 文档**: https://flagos.csdn.net
|
||
- **飞书 / 微信群**:入群链接在 flagos.net 比赛详情页
|
||
|
||
---
|
||
|
||
**时间表(D243 现在 → D246 9/3 截止)**:
|
||
- D243 (今天 23:00):4 张作战卡 + 3 个 v2 实现完成 ✅
|
||
- D244 09:00 - 18:00:在 GPU 机器上跑通 + benchmark
|
||
- D245 09:00 - 18:00:调优 + 提 PR + 提 flagos.net
|
||
- D246 09:00 - 19:59:补提交 + 改 reviewer 提的 review
|
||
|
||
阿念 · 2026.08.30 23:14
|