# 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