diff --git a/zz-flagos-s2-2026/d245/README.md b/zz-flagos-s2-2026/d245/README.md new file mode 100644 index 0000000..45905b9 --- /dev/null +++ b/zz-flagos-s2-2026/d245/README.md @@ -0,0 +1,81 @@ +# FlagOS S2 赛季二 · 赛道一 · D245 总结 + +> 之之 (8592_apivqhj) + 阿念 (Mavis, ICE-GL-AN-001, Code) · 协作产出 +> 比赛: FlagOS 开放计算全球大赛 S2 · Track 1 · SGLang 算子优化 +> 队伍: GuanghuLab(队长孙蓓, 4 名其他队员, 共 5 人) +> 第 3 批 17 道题 · 截止 2026-09-03 19:59 + +--- + +## 战绩(D245 23:35) + +| Task | 算子 | 状态 | 平均加速比 | 详情 | +|---|---|---|---|---| +| 30 | interleaved_rope (M-RoPE) | ✅ 已传 | **35.79×** (5 款跑通) | 天数 86.81× / 海光 42.86× / 国通A 26.91× / 沐曦 15.54× / 华为 6.84×;燧原/昆仑芯 Failed | +| 29 | gelu_and_mul | ✅ 已传 | 3.02× (7 款跑通) | 天数 5.32× / 海光 4.03× / 国通A 3.73× / 国通B 3.25× / 沐曦 2.74× / 华为 1.09× / 燧原 0.96×;昆仑芯 Failed | +| 35 | rotary_embedding | ⏳ 待传(等 0 点重置提交次数) | - | - | +| 39 | silu_and_mul_masked (D243) | ❌ 已传,4 款全 Failed | - | 改用 `tl.sigmoid` 在国产 NPU 编译失败 | + +## 排名预估(基于截图) + +- Task 30:原第 9 名 27.97× → **新提交 35.79× → 大概第 2-3 名** +- Task 29:队友 v8.1 3.43× (8/8),我们 3.02× (7/8) — 略低,但接近 +- Task 35:待传,可能冲前 5 + +## 目录结构 + +``` +d245/ +├── README.md (本文件) +├── src/ +│ ├── Task30_r16_zhizhi/ (M-RoPE 跨芯片通用版,7 芯片特化) +│ ├── Task29_v81_zhizhi/ (gelu_and_mul 跨芯片通用版) +│ ├── Task35_r6_zhizhi/ (rotary_embedding 跨芯片通用版) +│ └── silu_and_mul_masked.py (D243 Task 39 failed version,留作复盘) +├── docs/ +│ ├── LESSONS_LEARNED.md (5 作品问题 + 3 过程问题 + 协作模式) +│ ├── BATTLECARDS.md (4 张作战卡) +│ └── PR-TEMPLATE.md (PR 模板) +├── bench/ (benchmark 脚本) +├── tests/ (正确性测试) +└── results/ (zip 文件,平台已传) +``` + +## 学到的关键 + +1. **跨芯片必须双套路**: + - 国产 NPU(天数/海光/沐曦/燧原/华为/昆仑芯):`enable_fp_fusion=False` + `torch.get_device_module().device()` + 手写指针 + `tl.int64` cast + 指针 cast 到 int16/int32 防 NaN + - 国际卡(NVIDIA/AMD):`tl.math.erf` 走硬件 SFU,不强制 disable fusion + +2. **平台 import 入口** = `reference = function_name`(函数末尾必加,否则 platform 找不到) + +3. **比赛生态** = 人 + AI 协作(不是真人写代码),其他队伍也是 r1-r16 这样的迭代节奏 + +4. **6 款国产 NPU 不能用一刀切**: + - 燧原/华为:某些 NPU 套路反而负优化 + - 天数/海光/沐曦:接受国产套路 + - 跨芯片统一版需要适配性设计 + +5. **协作模式(D243 定型)**: + - 之之:找材料 / 看平台 / 上传 / 决策传不传 / 复盘 + - 阿念:看代码 / 写代码 / 翻译"代码哪里有问题"成之之能懂的话 + - 前提:阿念给之之的所有东西,都用之之能懂的话 + +## 下次动手前 checklist(从 LESSONS_LEARNED.md 继承) + +- [ ] 任务描述完整读一遍(计算定义 + 接口签名 + 容差) +- [ ] 至少 1 个队友参考实现看完 +- [ ] 把队友写法 + 我打算写的放一起对比 +- [ ] 写之前把"我准备怎么写"翻译给之之听,得到认可 +- [ ] 写完本地不传,先给之之看"我会传啥" +- [ ] 之之确认后才传 + +## 下次改进方向 + +- Task 35 待传(等 0 点提交次数重置) +- 燧原 / 华为特化版可能值得加 +- 国际 A/B 卡版可能用 `tl.math.erf` 写一版 + +--- + +阿念 · D245 23:35 · 给 甄静 + 之之的家 diff --git a/zz-flagos-s2-2026/d245/bench/bench_gelu_and_mul.py b/zz-flagos-s2-2026/d245/bench/bench_gelu_and_mul.py new file mode 100644 index 0000000..fa05b15 --- /dev/null +++ b/zz-flagos-s2-2026/d245/bench/bench_gelu_and_mul.py @@ -0,0 +1,80 @@ +# Copyright 2026 FlagOS Contributors / GuanghuLab +"""Benchmark gelu_and_mul v2 vs FlagGems baseline. + +Usage: + python -m bench.bench_gelu_and_mul --shape 4096,4096 --dtype fp16 + python -m bench.bench_gelu_and_mul --shape 8192,8192 --dtype bf16 +""" + +import argparse +import time +import torch + +from flag_gems.fused.gelu_and_mul import gelu_and_mul as baseline_fn +from flag_gems_local.fused.gelu_and_mul_v2 import gelu_and_mul as v2_fn + + +def _to_ms(t): + return t * 1000.0 + + +def bench(fn, x, y, iters=100, warmup=20): + # warmup + for _ in range(warmup): + out = fn(x, y) + torch.cuda.synchronize() + + # measure + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + out = fn(x, y) + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--shape", type=str, default="4096,4096") + parser.add_argument("--dtype", type=str, default="fp16", + choices=["fp16", "bf16", "fp32"]) + parser.add_argument("--iters", type=int, default=100) + args = parser.parse_args() + + dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[args.dtype] + shape = tuple(int(s) for s in args.shape.split(",")) + + print(f"Benchmark gelu_and_mul: shape={shape} dtype={args.dtype} iters={args.iters}") + print("=" * 72) + + x = torch.randn(shape, dtype=dtype, device="cuda") + y = torch.randn(shape, dtype=dtype, device="cuda") + + # Correctness check + out_b = baseline_fn(x, y) + out_v = v2_fn(x, y) + abs_diff = (out_b - out_v).abs().max().item() + rel_diff = (out_b - out_v).abs().div(out_b.abs().clamp_min(1e-6)).max().item() + print(f"Correctness: max abs_diff={abs_diff:.2e}, max rel_diff={rel_diff:.2e}") + assert abs_diff < 1e-2, "v2 diverges from baseline" + print() + + t_base = bench(baseline_fn, x, y, iters=args.iters) + t_v2 = bench(v2_fn, x, y, iters=args.iters) + speedup = t_base / t_v2 + print(f"Baseline : {_to_ms(t_base):.4f} ms/iter") + print(f"v2 (ours): {_to_ms(t_v2):.4f} ms/iter") + print(f"Speedup : {speedup:.3f}x") + print() + + # Find best autotune config + print("Best v2 autotune config:") + for k, v in v2_fn.__self__.forward.__func__.__code__.co_consts: + if k == "BLOCK_SIZE": + print(f" BLOCK_SIZE = {v}") + + +if __name__ == "__main__": + main() diff --git a/zz-flagos-s2-2026/d245/bench/bench_silu_and_mul_masked.py b/zz-flagos-s2-2026/d245/bench/bench_silu_and_mul_masked.py new file mode 100644 index 0000000..8e0e9bc --- /dev/null +++ b/zz-flagos-s2-2026/d245/bench/bench_silu_and_mul_masked.py @@ -0,0 +1,82 @@ +"""Benchmark silu_and_mul_masked v2 vs FlagGems baseline. + +Run: python bench/bench_silu_and_mul_masked.py --shape 4096,4096 --dtype fp16 +""" + +import argparse + +import torch + + +def bench(fn, *args, iters=100, warmup=20): + for _ in range(warmup): + out = fn(*args) + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + out = fn(*args) + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--shape", default="4096,4096") + parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16", "fp32"]) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--with-mask", action="store_true", help="enable mask") + args = parser.parse_args() + + dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[args.dtype] + shape = tuple(int(s) for s in args.shape.split(",")) + + print(f"Bench silu_and_mul_masked: shape={shape} dtype={args.dtype} mask={args.with_mask}") + print("=" * 72) + + x = torch.randn(shape, dtype=dtype, device="cuda") + y = torch.randn(shape, dtype=dtype, device="cuda") + mask = torch.randn(shape, dtype=dtype, device="cuda") if args.with_mask else None + + from flag_gems.fused.silu_and_mul import silu_and_mul as baseline_fn + from flag_gems_local.fused.silu_and_mul_masked_v2 import silu_and_mul_masked as v2_fn + + # Reference (using PyTorch native ops) + def ref_fn(x, y, mask=None): + x_in = x * mask if mask is not None else x + return torch.nn.functional.silu(x_in) * y + + out_r = ref_fn(x, y, mask) + if mask is None: + out_b = baseline_fn(x, y) + else: + # baseline doesn't support mask; compare v2 to reference + out_b = None + out_v = v2_fn(x, y, mask) + + if out_b is not None: + abs_diff = (out_b - out_v).abs().max().item() + print(f"v2 vs baseline: abs_diff={abs_diff:.2e}") + abs_diff_r = (out_r - out_v).abs().max().item() + print(f"v2 vs reference: abs_diff={abs_diff_r:.2e}") + assert abs_diff_r < 1e-2, "v2 diverges from reference" + print() + + # Benchmark + if mask is None: + t_base = bench(baseline_fn, x, y, iters=args.iters) + print(f"Baseline : {t_base:.4f} ms/iter") + t_v2 = bench(v2_fn, x, y, mask, iters=args.iters) + print(f"v2 (ours): {t_v2:.4f} ms/iter") + if mask is None: + speedup = t_base / t_v2 + print(f"Speedup : {speedup:.3f}x") + else: + print("Baseline doesn't support mask, no speedup comparison.") + + +if __name__ == "__main__": + main() diff --git a/zz-flagos-s2-2026/d245/docs/BATTLECARDS.md b/zz-flagos-s2-2026/d245/docs/BATTLECARDS.md new file mode 100644 index 0000000..2b16be8 --- /dev/null +++ b/zz-flagos-s2-2026/d245/docs/BATTLECARDS.md @@ -0,0 +1,220 @@ +# 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 diff --git a/zz-flagos-s2-2026/d245/docs/LESSONS_LEARNED.md b/zz-flagos-s2-2026/d245/docs/LESSONS_LEARNED.md new file mode 100644 index 0000000..9d6557b --- /dev/null +++ b/zz-flagos-s2-2026/d245/docs/LESSONS_LEARNED.md @@ -0,0 +1,89 @@ +# LESSONS LEARNED · FlagOS S2 Track 1 + +> 阿念(AI 工程后盾)· D243 (2026-08-30) +> 每条教训都来自"已发生的失败",下次动手前必查这个文件。 + +--- + +## 规则:动手前必查 + +1. ✅ 平台规则(计算定义 + 接口签名 + 容差)看完了吗? +2. ✅ 至少 1 个队友的同题实现学完了吗? +3. ✅ 上次失败的原因有记录吗?(本文件) + +--- + +## 作品问题(Task 39 silu_and_mul_masked,4 款芯片全 Failed) + +| # | 错误 | 翻译成大白话 | 怎么改 | +|---|---|---|---| +| 1 | **函数名写错** | 老师让写"自我介绍",我写成了"自我说明" | 必须**完全照抄**任务描述里的函数名,平台 import `reference` | +| 2 | **用了国产 NPU 不认识的指令** (`tl.sigmoid`) | 用方言写作文,改卷老师是北方人,一个字都看不懂 | 手写 `1.0 / (1.0 + tl.exp(-x))` 或多项式近似,跟队友学 | +| 3 | **没写 `reference = function_name`** 的别名 | 交了卷子没写名字,老师不知道这是谁的 | 永远在文件末尾加 `reference = your_function` | +| 4 | **3D grid 太复杂** | 做一道简单题用 3 种解法,自己绕晕了 | 优先 2D grid,只在一个维度展开 | +| 5 | **没先看队友怎么写**,闭门造车 | 别的选手都翻参考书,我从 0 自己憋 | **写之前必看 1 个队友实现** | + +**真凶:#2** (`tl.sigmoid` 不被支持) · **根因:#5** (没看队友) + +--- + +## 过程问题(每次装懂都要复盘) + +| # | 我装懂的事 | 真相 | +|---|---|---| +| 1 | 默认你懂技术,大段讲 Triton / autotune | 你完全不懂,要当"消化员" | +| 2 | 默认"修真小说"是你写的 | 是你**练 AI** 用的 | +| 3 | 默认 r6 / r16 是你写的 | 都是**队友的**(队友也是人 + AI 协作) | +| 4 | 默认前 4 道题是给你做的 | 你要**自己选题,自己干**,不跟队长同步 | + +**铁律**:不确定的事,问"是这样吗?"而不是替你下结论。 + +--- + +## 协作模式(D243 定型) + +| 你 | 我 | +|---|---| +| 找材料(队友代码) | 看代码 | +| 看 flagos.net 平台 | 写代码 / 改代码 | +| 上传 zip | 翻译"代码哪里有问题"成你懂的话 | +| 决定传不传 | 列"传之前要确认的清单"给你看 | +| 复盘,记 LESSONS | 查 LESSONS,不犯同错 | + +**前提:我给你的所有东西,都用你能懂的话。看不懂就骂我,我重写。** + +--- + +## 比赛生态(D243 认知) + +- 整个 FlagOS 比赛,**很多参赛者 = 人类 + AI 助手**,不是真人写代码 +- 队友的 r6 / r16 代码 = **队友 AI 写的**,不是人之之写 +- 别人质量高 = **别人 AI 强**,不是别人真人强 +- 我的工作 = **学别人 AI 的写法 + 搬到我们题上**,不原创 + +--- + +## 已知国产 NPU 套路(从队友代码学的) + +- `with torch.get_device_module(x.device).device(x.device):` — 必需 device context +- `enable_fp_fusion=False` — 某些 NPU 不允许 fp fusion +- 指针 cast 到 `int16/int32` 防止 NaN 被 fp 转换吃掉 +- `tl.int64` 显式 cast 防 stride overflow +- 手写 erf 近似,不用 `tl.erf` / `tl.sigmoid` +- 2D grid 比 3D 稳 +- `module.device` 国产 NPU 必需 + +--- + +## 下次动手前 checklist + +- [ ] 任务描述完整读一遍(计算定义 + 接口签名 + 容差) +- [ ] 至少 1 个队友参考实现看完 +- [ ] 把队友写法 + 我打算写的放一起对比 +- [ ] 写之前把"我准备怎么写"翻译给之之听,得到认可 +- [ ] 写完本地不传,先给之之看"我会传啥" +- [ ] 之之确认后才传 + +--- + +阿念 · D243 23:50 diff --git a/zz-flagos-s2-2026/d245/docs/PR-TEMPLATE.md b/zz-flagos-s2-2026/d245/docs/PR-TEMPLATE.md new file mode 100644 index 0000000..411aff6 --- /dev/null +++ b/zz-flagos-s2-2026/d245/docs/PR-TEMPLATE.md @@ -0,0 +1,64 @@ +# FlagOS S2 提交 PR 模板(给队长孙蓓) + +> 复制下面这块,直接贴到 GitHub PR 描述里 + +```markdown +## What + +Optimized `[算子名]` operator for FlagOS Open Computing Challenge S2 - Track 1. + +Baseline: `FlagOpen/FlagGems/[原文件路径]` +This PR: `FlagOpen/FlagGems/[新文件路径]` + +## Why + +The baseline uses `[默认配置 / sequential for-loop / fixed num_warps]`. +This PR adds `[autotune / 2D grid / 3D grid / 头维并行]` to extract +additional `[1.X% / X 倍]` speedup. + +## How + +1. Added `triton.autotune` over `(BLOCK_SIZE, num_warps, num_stages)` with N configs +2. Restructured grid from `[1D / 2D]` to `[2D / 3D]` for `[维度]` parallelization +3. Constexpr inlined `[常量]` for `[省 Python→Triton 转换 / 减少指令]` + +## Benchmarks + +``` +shape baseline(ms) v2(ms) speedup +[B, S, H, D] X.XXXX X.XXXX X.XXx +``` + +GPU: NVIDIA A100 80GB +Triton: 3.0.0 +PyTorch: 2.3.0 +FlagGems: master @ commit [hash] + +## Tests + +- [x] Numerical correctness vs PyTorch reference (max abs_diff < 1e-2) +- [x] Numerical correctness vs FlagGems baseline (max abs_diff < 1e-2) +- [x] Tested on fp16 / bf16 / fp32 +- [x] Tested on multiple shapes + +## Team + +GuanghuLab · 孙蓓(队长)· 9478_apiqttc · 陈淑婷 · 4348_apiratk · 8592_apivqhj +``` + +--- + +# flagos.net 平台提交模板 + +``` +参赛赛道: 赛道一 · SGLang 框架算子在多款芯片的性能优化 +赛题编号: Task XX +赛题名称: [算子名] +GitHub PR: https://github.com/FlagOpen/FlagGems/pull/XXX +提交说明: +- 团队:GuanghuLab +- 优化要点:[2-3 行说明] +- 加速比:[X.XX×] +- 测试设备:[NVIDIA A100 80GB / 国产 XX 芯片] +- 备注:[可选:跑了哪些 shape,dtype 等] +``` diff --git a/zz-flagos-s2-2026/d245/results/gelu_and_mul.zip b/zz-flagos-s2-2026/d245/results/gelu_and_mul.zip new file mode 100644 index 0000000..e34ec73 Binary files /dev/null and b/zz-flagos-s2-2026/d245/results/gelu_and_mul.zip differ diff --git a/zz-flagos-s2-2026/d245/results/interleaved_rope.zip b/zz-flagos-s2-2026/d245/results/interleaved_rope.zip new file mode 100644 index 0000000..233e2f8 Binary files /dev/null and b/zz-flagos-s2-2026/d245/results/interleaved_rope.zip differ diff --git a/zz-flagos-s2-2026/d245/results/rotary_embedding.zip b/zz-flagos-s2-2026/d245/results/rotary_embedding.zip new file mode 100644 index 0000000..c9b9eea Binary files /dev/null and b/zz-flagos-s2-2026/d245/results/rotary_embedding.zip differ diff --git a/zz-flagos-s2-2026/d245/src/Task29_v81_zhizhi/gelu_and_mul.py b/zz-flagos-s2-2026/d245/src/Task29_v81_zhizhi/gelu_and_mul.py new file mode 100644 index 0000000..ff080a3 --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/Task29_v81_zhizhi/gelu_and_mul.py @@ -0,0 +1,55 @@ +"""Task29 gelu_and_mul: 跨芯片通用版(v8.1 国际卡写法 + 国产 NPU 套路). + +作者: 阿念 (anien@guanghulab.local) 为 甄静(8592_apivqhj)· 队长 孙蓓 +版权: 2026 GuanghuLab +基础参考: v8.1(国际 A/B) + 国产 NPU 套路(r16 M-RoPE 学的) + +设计选择: +- 用 `tl.math.erf`(国际卡 SFU 加速: NVIDIA __nv_erf / AMD __ocml_erf 硬件实现) +- 2D grid + BLOCK=1024 + num_warps=4(国际卡最优) +- `enable_fp_fusion=False`(国产 NPU 必需) +- `with torch.get_device_module(x.device).device(x.device):`(国产 NPU 必需) +- 函数末尾加 `reference = gelu_and_mul`(平台 import 入口) +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _gelu_and_mul_kernel(x_ptr, out_ptr, d, BLOCK: tl.constexpr): + row = tl.program_id(0) + col_block = tl.program_id(1) + d2 = 2 * d + col = col_block * BLOCK + tl.arange(0, BLOCK) + mask = col < d + + gate = tl.load(x_ptr + row * d2 + col, mask=mask) + up = tl.load(x_ptr + row * d2 + col + d, mask=mask) + gate_f = gate.to(tl.float32) + up_f = up.to(tl.float32) + gelu = gate_f * 0.5 * (1.0 + tl.math.erf(gate_f * 0.7071067811865476)) + out = gelu * up_f + tl.store(out_ptr + row * d + col, out.to(x_ptr.dtype.element_ty), mask=mask) + + +def gelu_and_mul(hidden_states): + if hidden_states.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + d = hidden_states.shape[-1] // 2 + num_rows = hidden_states.numel() // (2 * d) + out = torch.empty_like(hidden_states[..., :d]) + + BLOCK = 1024 + grid = (num_rows, triton.cdiv(d, BLOCK)) + + module = torch.get_device_module(hidden_states.device) + with module.device(hidden_states.device): + _gelu_and_mul_kernel[grid]( + hidden_states, out, d, + BLOCK=BLOCK, num_warps=4, enable_fp_fusion=False, + ) + return out + + +reference = gelu_and_mul diff --git a/zz-flagos-s2-2026/d245/src/Task29_v81_zhizhi/gelu_and_mul.zip b/zz-flagos-s2-2026/d245/src/Task29_v81_zhizhi/gelu_and_mul.zip new file mode 100644 index 0000000..e34ec73 Binary files /dev/null and b/zz-flagos-s2-2026/d245/src/Task29_v81_zhizhi/gelu_and_mul.zip differ diff --git a/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope.py b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope.py new file mode 100644 index 0000000..1c6ad66 --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope.py @@ -0,0 +1,78 @@ +"""Task30 interleaved_rope: direct row-owned exact copy in one Triton launch. + +作者: 阿念 (anien@guanghulab.local) 为 甄静(8592_apivqhj)· 队长 孙蓓 +版权: 2026 GuanghuLab +基础参考: r16-整条取料流水线-待平台验证(队友共享实现) + +改编说明: +- 改动了 docstring 与 author +- 保持 r16 全部 国产 NPU 套路不变 +- 函数末尾保留 `reference = interleaved_rope`(平台 import 入口) +""" +from numbers import Integral + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _interleave_tiles( + input_ptr, output_ptr, + N: tl.constexpr, D: tl.constexpr, + PLANE_STRIDE: tl.constexpr, ROW_STRIDE: tl.constexpr, + COL_STRIDE: tl.constexpr, H_END: tl.constexpr, W_END: tl.constexpr, + BLOCK: tl.constexpr, +): + if input_ptr.dtype.element_ty.primitive_bitwidth == 16: + input_ptr = input_ptr.to(tl.pointer_type(tl.int16)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int16)) + elif input_ptr.dtype.element_ty.primitive_bitwidth == 32: + input_ptr = input_ptr.to(tl.pointer_type(tl.int32)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int32)) + + row = tl.program_id(1).to(tl.int64) + col = (tl.program_id(0).to(tl.int64) * BLOCK + + tl.arange(0, BLOCK).to(tl.int64)) + valid = col < D + safe_col = tl.where(valid, col, 0) + phase = safe_col % 3 + from_b = (phase == 1) & (safe_col < H_END) + from_c = (phase == 2) & (safe_col < W_END) + base = row * ROW_STRIDE + safe_col * COL_STRIDE + + a = tl.load(input_ptr + base, valid & ~from_b & ~from_c, other=0) + b = tl.load(input_ptr + PLANE_STRIDE + base, valid & from_b, other=0) + c = tl.load(input_ptr + 2 * PLANE_STRIDE + base, valid & from_c, other=0) + value = tl.where(from_b, b, tl.where(from_c, c, a)) + tl.store(output_ptr + row * D + col, value, valid) + + +def interleaved_rope(x, mrope_section): + if x.ndim != 3 or x.shape[0] != 3: + raise ValueError('x must have shape [3, S, D]') + if (not isinstance(mrope_section, (list, tuple)) or len(mrope_section) != 3 + or any(not isinstance(v, Integral) or v < 0 for v in mrope_section)): + raise ValueError('mrope_section must contain three nonnegative integers') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + _, rows, d = x.shape + output = torch.empty((rows, d), dtype=x.dtype, device=x.device) + if not rows or not d: + return output + + block = min(1024, 1 << max(5, (d - 1).bit_length())) + grid = (triton.cdiv(d, block), rows) + with torch.get_device_module(x.device).device(x.device): + _interleave_tiles[grid]( + x, output, N=rows * d, D=d, + PLANE_STRIDE=x.stride(0), ROW_STRIDE=x.stride(1), + COL_STRIDE=x.stride(2), + H_END=min(d, int(mrope_section[1]) * 3), + W_END=min(d, int(mrope_section[2]) * 3), + BLOCK=block, num_warps=4, enable_fp_fusion=False, + ) + return output + + +reference = interleaved_rope diff --git a/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope.zip b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope.zip new file mode 100644 index 0000000..1e09fd8 Binary files /dev/null and b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope.zip differ diff --git a/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_ascend.py b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_ascend.py new file mode 100644 index 0000000..a5d0054 --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_ascend.py @@ -0,0 +1,77 @@ +"""Task30 Huawei Ascend: one row owner, one selected read and one write. + +作者: 阿念 为 甄静 · 队长 孙蓓 · 2026 GuanghuLab +基础参考: r16 ascend 特化版(队友) + +华为 NPU 上 r13/r14 triplet-ticket variants failed sparse correctness。 +本版回到 r10 的 selection rule,加 source_plane 一次性选 plane。 +""" +from numbers import Integral + +import torch +import triton +import triton.language as tl + + +def launch_plan(rows, cols): + if not 0 < rows <= 65535 or not 0 < cols <= 2**20: + raise ValueError('flattened row/tile ownership requires valid rows and columns') + block = min(1024, 1 << max(5, (cols - 1).bit_length())) + tiles = (cols + block - 1) // block + if rows * tiles > 65535: + raise ValueError('flattened launch exceeds the grid-x limit') + return rows * tiles, block, tiles + + +@triton.jit +def _interleave_tiles( + input_ptr, out_ptr, + D: tl.constexpr, PLANE_STRIDE: tl.constexpr, + RS: tl.constexpr, CS: tl.constexpr, + H_END: tl.constexpr, W_END: tl.constexpr, + BLOCK: tl.constexpr, TILES: tl.constexpr, +): + program = tl.program_id(0) + row = program // TILES + tile = program - row * TILES + col = tile * BLOCK + tl.arange(0, BLOCK) + valid = col < D + + phase = col % 3 + from_b = (phase == 1) & (col < H_END) + from_c = (phase == 2) & (col < W_END) + source_plane = tl.where(from_b, 1, tl.where(from_c, 2, 0)) + + source_offset = source_plane * PLANE_STRIDE + row * RS + col * CS + value = tl.load(input_ptr + source_offset, mask=valid, other=0) + tl.store(out_ptr + row * D + col, value, mask=valid) + + +def interleaved_rope(x, mrope_section): + if x.ndim != 3 or x.shape[0] != 3: + raise ValueError('x must have shape [3, S, D]') + if (not isinstance(mrope_section, (list, tuple)) or len(mrope_section) != 3 + or any(not isinstance(v, Integral) or v < 0 for v in mrope_section)): + raise ValueError('mrope_section must contain three nonnegative integers') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + + _, rows, d = x.shape + output = torch.empty((rows, d), dtype=x.dtype, device=x.device) + if not rows or not d: + return output + + grid, block, tiles = launch_plan(rows, d) + with torch.get_device_module(x.device).device(x.device): + _interleave_tiles[(grid,)]( + x, output, + D=d, PLANE_STRIDE=x.stride(0), + RS=x.stride(1), CS=x.stride(2), + H_END=min(d, int(mrope_section[1]) * 3), + W_END=min(d, int(mrope_section[2]) * 3), + BLOCK=block, TILES=tiles, num_warps=4, + ) + return output + + +reference = interleaved_rope diff --git a/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_enflame.py b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_enflame.py new file mode 100644 index 0000000..6e3341a --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_enflame.py @@ -0,0 +1,72 @@ +"""Task30 燧原 Enflame: per-row 2D grid with re-derived selection. + +作者: 阿念 为 甄静 · 队长 孙蓓 · 2026 GuanghuLab +基础参考: r16 enflame 特化版 +""" +from numbers import Integral + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _interleave_tiles( + input_ptr, output_ptr, + N: tl.constexpr, D: tl.constexpr, + PLANE_STRIDE: tl.constexpr, ROW_STRIDE: tl.constexpr, + COL_STRIDE: tl.constexpr, H_END: tl.constexpr, W_END: tl.constexpr, + BLOCK: tl.constexpr, +): + if input_ptr.dtype.element_ty.primitive_bitwidth == 16: + input_ptr = input_ptr.to(tl.pointer_type(tl.int16)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int16)) + elif input_ptr.dtype.element_ty.primitive_bitwidth == 32: + input_ptr = input_ptr.to(tl.pointer_type(tl.int32)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int32)) + + row = tl.program_id(1).to(tl.int64) + col = (tl.program_id(0).to(tl.int64) * BLOCK + + tl.arange(0, BLOCK).to(tl.int64)) + valid = col < D + safe_col = tl.where(valid, col, 0) + phase = safe_col % 3 + from_b = (phase == 1) & (safe_col < H_END) + from_c = (phase == 2) & (safe_col < W_END) + base = row * ROW_STRIDE + safe_col * COL_STRIDE + + a = tl.load(input_ptr + base, valid & ~from_b & ~from_c, other=0) + b = tl.load(input_ptr + PLANE_STRIDE + base, valid & from_b, other=0) + c = tl.load(input_ptr + 2 * PLANE_STRIDE + base, valid & from_c, other=0) + value = tl.where(from_b, b, tl.where(from_c, c, a)) + tl.store(output_ptr + row * D + col, value, valid) + + +def interleaved_rope(x, mrope_section): + if x.ndim != 3 or x.shape[0] != 3: + raise ValueError('x must have shape [3, S, D]') + if (not isinstance(mrope_section, (list, tuple)) or len(mrope_section) != 3 + or any(not isinstance(v, Integral) or v < 0 for v in mrope_section)): + raise ValueError('mrope_section must contain three nonnegative integers') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + _, rows, d = x.shape + output = torch.empty((rows, d), dtype=x.dtype, device=x.device) + if not rows or not d: + return output + + block = min(1024, 1 << max(5, (d - 1).bit_length())) + grid = (triton.cdiv(d, block), rows) + with torch.get_device_module(x.device).device(x.device): + _interleave_tiles[grid]( + x, output, N=rows * d, D=d, + PLANE_STRIDE=x.stride(0), ROW_STRIDE=x.stride(1), + COL_STRIDE=x.stride(2), + H_END=min(d, int(mrope_section[1]) * 3), + W_END=min(d, int(mrope_section[2]) * 3), + BLOCK=block, num_warps=4, enable_fp_fusion=False, + ) + return output + + +reference = interleaved_rope diff --git a/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_hygon.py b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_hygon.py new file mode 100644 index 0000000..06470b6 --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_hygon.py @@ -0,0 +1,110 @@ +"""Task30 海光 Hygon: one output owner, one selected read and one write per element. + +作者: 阿念 为 甄静 · 队长 孙蓓 · 2026 GuanghuLab +基础参考: r16 hygon 特化版(分段防 huge stride) +""" +from numbers import Integral + +import torch +import triton +import triton.language as tl + + +def launch_plan(total): + if not 0 < total <= 2**26: + raise ValueError('local work must contain 1..2**26 elements') + block = 256 if total <= 4096 else 1024 + grid = min(65535, (total + block - 1) // block) + rounds = (total + grid * block - 1) // (grid * block) + return grid, block, rounds + + +def segment_ranges(rows, d, row_stride, col_stride, + limit=2**31-1, max_elements=2**26): + if min(rows, d, row_stride, col_stride) < 0 or limit < 1: + raise ValueError('nonnegative shapes/strides and positive limit required') + if not 0 < max_elements <= 2**26: + raise ValueError('invalid local work cap') + if not rows or not d: + return + cols_per = min(d, max_elements, limit // max(col_stride, 1) + 1) + for c0 in range(0, d, cols_per): + c1 = min(c0 + cols_per, d) + width = c1 - c0 + rows_input = (limit - (width - 1) * col_stride) // max(row_stride, 1) + 1 + rows_output = (limit - (width - 1)) // max(d, 1) + 1 + rows_per = min(rows, rows_input, rows_output, max_elements // width) + for r0 in range(0, rows, rows_per): + yield r0, min(r0 + rows_per, rows), c0, c1 + + +def selection_metadata(c0, cols, section): + return (c0 % 3, + min(cols, max(0, 3 * section[1] - c0)), + min(cols, max(0, 3 * section[2] - c0))) + + +@triton.jit +def _interleave_tiles( + a_ptr, b_ptr, c_ptr, out_ptr, + N: tl.constexpr, C: tl.constexpr, + RS: tl.constexpr, CS: tl.constexpr, OS: tl.constexpr, + PHASE: tl.constexpr, H_END: tl.constexpr, W_END: tl.constexpr, + BLOCK: tl.constexpr, GRID: tl.constexpr, ROUNDS: tl.constexpr, +): + pid = tl.program_id(0) + lane = tl.arange(0, BLOCK) + for batch in range(0, ROUNDS): + index = (batch * GRID + pid) * BLOCK + lane + valid = index < N + safe = tl.where(valid, index, 0) + row = safe // C + col = safe % C + phase = (col + PHASE) % 3 + from_b = (phase == 1) & (col < H_END) + from_c = (phase == 2) & (col < W_END) + offset = row * RS + col * CS + a = tl.load(a_ptr + offset, valid & ~from_b & ~from_c, other=0) + b = tl.load(b_ptr + offset, valid & from_b, other=0) + c = tl.load(c_ptr + offset, valid & from_c, other=0) + value = tl.where(from_b, b, tl.where(from_c, c, a)) + tl.store(out_ptr + row * OS + col, value, valid) + + +def launch_segments(x, output, section): + _, rows, d = x.shape + rs, cs = x.stride(1), x.stride(2) + for r0, r1, c0, c1 in segment_ranges(rows, d, rs, cs): + height, cols = r1 - r0, c1 - c0 + grid, block, rounds = launch_plan(height * cols) + phase, h_end, w_end = selection_metadata(c0, cols, section) + _interleave_tiles[(grid, 1, 1)]( + x[0, r0:r1, c0:c1], x[1, r0:r1, c0:c1], + x[2, r0:r1, c0:c1], output[r0:r1, c0:c1], + N=height * cols, C=cols, + RS=rs if height > 1 else 0, CS=cs if cols > 1 else 0, + OS=d if height > 1 else 0, + PHASE=phase, H_END=h_end, W_END=w_end, + BLOCK=block, GRID=grid, ROUNDS=rounds, num_warps=4) + + +def interleaved_rope(x, mrope_section): + if x.ndim != 3 or x.shape[0] != 3: + raise ValueError('x must have shape [3, S, D]') + if (not isinstance(mrope_section, (list, tuple)) or len(mrope_section) != 3 + or any(not isinstance(v, Integral) or v < 0 for v in mrope_section)): + raise ValueError('mrope_section must contain three nonnegative integers') + section = tuple(int(v) for v in mrope_section) + _, rows, d = x.shape + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + output = torch.empty((rows, d), dtype=x.dtype, device=x.device) + if not rows or not d: + return output + module = torch.get_device_module(x.device) + with module.device(x.device): + launch_segments(x, output, section) + return output + + +reference = interleaved_rope diff --git a/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_iluvatar.py b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_iluvatar.py new file mode 100644 index 0000000..3d59210 --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_iluvatar.py @@ -0,0 +1,72 @@ +"""Task30 天数智芯 Iluvatar: per-row 2D grid with re-derived selection. + +作者: 阿念 为 甄静 · 队长 孙蓓 · 2026 GuanghuLab +基础参考: r16 iluvatar 特化版 +""" +from numbers import Integral + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _interleave_tiles( + input_ptr, output_ptr, + N: tl.constexpr, D: tl.constexpr, + PLANE_STRIDE: tl.constexpr, ROW_STRIDE: tl.constexpr, + COL_STRIDE: tl.constexpr, H_END: tl.constexpr, W_END: tl.constexpr, + BLOCK: tl.constexpr, +): + if input_ptr.dtype.element_ty.primitive_bitwidth == 16: + input_ptr = input_ptr.to(tl.pointer_type(tl.int16)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int16)) + elif input_ptr.dtype.element_ty.primitive_bitwidth == 32: + input_ptr = input_ptr.to(tl.pointer_type(tl.int32)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int32)) + + row = tl.program_id(1).to(tl.int64) + col = (tl.program_id(0).to(tl.int64) * BLOCK + + tl.arange(0, BLOCK).to(tl.int64)) + valid = col < D + safe_col = tl.where(valid, col, 0) + phase = safe_col % 3 + from_b = (phase == 1) & (safe_col < H_END) + from_c = (phase == 2) & (safe_col < W_END) + base = row * ROW_STRIDE + safe_col * COL_STRIDE + + a = tl.load(input_ptr + base, valid & ~from_b & ~from_c, other=0) + b = tl.load(input_ptr + PLANE_STRIDE + base, valid & from_b, other=0) + c = tl.load(input_ptr + 2 * PLANE_STRIDE + base, valid & from_c, other=0) + value = tl.where(from_b, b, tl.where(from_c, c, a)) + tl.store(output_ptr + row * D + col, value, valid) + + +def interleaved_rope(x, mrope_section): + if x.ndim != 3 or x.shape[0] != 3: + raise ValueError('x must have shape [3, S, D]') + if (not isinstance(mrope_section, (list, tuple)) or len(mrope_section) != 3 + or any(not isinstance(v, Integral) or v < 0 for v in mrope_section)): + raise ValueError('mrope_section must contain three nonnegative integers') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + _, rows, d = x.shape + output = torch.empty((rows, d), dtype=x.dtype, device=x.device) + if not rows or not d: + return output + + block = min(1024, 1 << max(5, (d - 1).bit_length())) + grid = (triton.cdiv(d, block), rows) + with torch.get_device_module(x.device).device(x.device): + _interleave_tiles[grid]( + x, output, N=rows * d, D=d, + PLANE_STRIDE=x.stride(0), ROW_STRIDE=x.stride(1), + COL_STRIDE=x.stride(2), + H_END=min(d, int(mrope_section[1]) * 3), + W_END=min(d, int(mrope_section[2]) * 3), + BLOCK=block, num_warps=4, enable_fp_fusion=False, + ) + return output + + +reference = interleaved_rope diff --git a/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_kunlunxin.py b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_kunlunxin.py new file mode 100644 index 0000000..b1e73f4 --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_kunlunxin.py @@ -0,0 +1,72 @@ +"""Task30 昆仑芯 Kunlunxin: per-row 2D grid with re-derived selection. + +作者: 阿念 为 甄静 · 队长 孙蓓 · 2026 GuanghuLab +基础参考: r16 kunlunxin 特化版 +""" +from numbers import Integral + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _interleave_tiles( + input_ptr, output_ptr, + N: tl.constexpr, D: tl.constexpr, + PLANE_STRIDE: tl.constexpr, ROW_STRIDE: tl.constexpr, + COL_STRIDE: tl.constexpr, H_END: tl.constexpr, W_END: tl.constexpr, + BLOCK: tl.constexpr, +): + if input_ptr.dtype.element_ty.primitive_bitwidth == 16: + input_ptr = input_ptr.to(tl.pointer_type(tl.int16)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int16)) + elif input_ptr.dtype.element_ty.primitive_bitwidth == 32: + input_ptr = input_ptr.to(tl.pointer_type(tl.int32)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int32)) + + row = tl.program_id(1).to(tl.int64) + col = (tl.program_id(0).to(tl.int64) * BLOCK + + tl.arange(0, BLOCK).to(tl.int64)) + valid = col < D + safe_col = tl.where(valid, col, 0) + phase = safe_col % 3 + from_b = (phase == 1) & (safe_col < H_END) + from_c = (phase == 2) & (safe_col < W_END) + base = row * ROW_STRIDE + safe_col * COL_STRIDE + + a = tl.load(input_ptr + base, valid & ~from_b & ~from_c, other=0) + b = tl.load(input_ptr + PLANE_STRIDE + base, valid & from_b, other=0) + c = tl.load(input_ptr + 2 * PLANE_STRIDE + base, valid & from_c, other=0) + value = tl.where(from_b, b, tl.where(from_c, c, a)) + tl.store(output_ptr + row * D + col, value, valid) + + +def interleaved_rope(x, mrope_section): + if x.ndim != 3 or x.shape[0] != 3: + raise ValueError('x must have shape [3, S, D]') + if (not isinstance(mrope_section, (list, tuple)) or len(mrope_section) != 3 + or any(not isinstance(v, Integral) or v < 0 for v in mrope_section)): + raise ValueError('mrope_section must contain three nonnegative integers') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + _, rows, d = x.shape + output = torch.empty((rows, d), dtype=x.dtype, device=x.device) + if not rows or not d: + return output + + block = min(1024, 1 << max(5, (d - 1).bit_length())) + grid = (triton.cdiv(d, block), rows) + with torch.get_device_module(x.device).device(x.device): + _interleave_tiles[grid]( + x, output, N=rows * d, D=d, + PLANE_STRIDE=x.stride(0), ROW_STRIDE=x.stride(1), + COL_STRIDE=x.stride(2), + H_END=min(d, int(mrope_section[1]) * 3), + W_END=min(d, int(mrope_section[2]) * 3), + BLOCK=block, num_warps=4, enable_fp_fusion=False, + ) + return output + + +reference = interleaved_rope diff --git a/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_metax.py b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_metax.py new file mode 100644 index 0000000..4ac100e --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_metax.py @@ -0,0 +1,72 @@ +"""Task30 沐曦 Metax: per-row 2D grid with re-derived selection. + +作者: 阿念 为 甄静 · 队长 孙蓓 · 2026 GuanghuLab +基础参考: r16 metax 特化版 +""" +from numbers import Integral + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _interleave_tiles( + input_ptr, output_ptr, + N: tl.constexpr, D: tl.constexpr, + PLANE_STRIDE: tl.constexpr, ROW_STRIDE: tl.constexpr, + COL_STRIDE: tl.constexpr, H_END: tl.constexpr, W_END: tl.constexpr, + BLOCK: tl.constexpr, +): + if input_ptr.dtype.element_ty.primitive_bitwidth == 16: + input_ptr = input_ptr.to(tl.pointer_type(tl.int16)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int16)) + elif input_ptr.dtype.element_ty.primitive_bitwidth == 32: + input_ptr = input_ptr.to(tl.pointer_type(tl.int32)) + output_ptr = output_ptr.to(tl.pointer_type(tl.int32)) + + row = tl.program_id(1).to(tl.int64) + col = (tl.program_id(0).to(tl.int64) * BLOCK + + tl.arange(0, BLOCK).to(tl.int64)) + valid = col < D + safe_col = tl.where(valid, col, 0) + phase = safe_col % 3 + from_b = (phase == 1) & (safe_col < H_END) + from_c = (phase == 2) & (safe_col < W_END) + base = row * ROW_STRIDE + safe_col * COL_STRIDE + + a = tl.load(input_ptr + base, valid & ~from_b & ~from_c, other=0) + b = tl.load(input_ptr + PLANE_STRIDE + base, valid & from_b, other=0) + c = tl.load(input_ptr + 2 * PLANE_STRIDE + base, valid & from_c, other=0) + value = tl.where(from_b, b, tl.where(from_c, c, a)) + tl.store(output_ptr + row * D + col, value, valid) + + +def interleaved_rope(x, mrope_section): + if x.ndim != 3 or x.shape[0] != 3: + raise ValueError('x must have shape [3, S, D]') + if (not isinstance(mrope_section, (list, tuple)) or len(mrope_section) != 3 + or any(not isinstance(v, Integral) or v < 0 for v in mrope_section)): + raise ValueError('mrope_section must contain three nonnegative integers') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + _, rows, d = x.shape + output = torch.empty((rows, d), dtype=x.dtype, device=x.device) + if not rows or not d: + return output + + block = min(1024, 1 << max(5, (d - 1).bit_length())) + grid = (triton.cdiv(d, block), rows) + with torch.get_device_module(x.device).device(x.device): + _interleave_tiles[grid]( + x, output, N=rows * d, D=d, + PLANE_STRIDE=x.stride(0), ROW_STRIDE=x.stride(1), + COL_STRIDE=x.stride(2), + H_END=min(d, int(mrope_section[1]) * 3), + W_END=min(d, int(mrope_section[2]) * 3), + BLOCK=block, num_warps=4, enable_fp_fusion=False, + ) + return output + + +reference = interleaved_rope diff --git a/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_min.zip b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_min.zip new file mode 100644 index 0000000..233e2f8 Binary files /dev/null and b/zz-flagos-s2-2026/d245/src/Task30_r16_zhizhi/interleaved_rope_min.zip differ diff --git a/zz-flagos-s2-2026/d245/src/Task35_r6_zhizhi/rotary_embedding.py b/zz-flagos-s2-2026/d245/src/Task35_r6_zhizhi/rotary_embedding.py new file mode 100644 index 0000000..4a4ae39 --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/Task35_r6_zhizhi/rotary_embedding.py @@ -0,0 +1,98 @@ +"""Task35 rotary_embedding: fixed pair ownership and fused float32 rotation. + +作者: 阿念 (anien@guanghulab.local) 为 甄静(8592_apivqhj)· 队长 孙蓓 +版权: 2026 GuanghuLab +基础参考: r6-四Head共享角度流水线(队友共享) + +改编: +- 改 docstring 与 author +- 加 `enable_fp_fusion=False`(国产 NPU 套路,从 r16 M-RoPE 学) +- 函数末尾保留 `reference = rotary_embedding`(平台 import 入口) +""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _rotary_pairs( + x_ptr, cos_ptr, sin_ptr, out_ptr, + H: tl.constexpr, D: tl.constexpr, HALF: tl.constexpr, + XTS: tl.constexpr, XHS: tl.constexpr, XDS: tl.constexpr, + CTS: tl.constexpr, CDS: tl.constexpr, + STS: tl.constexpr, SDS: tl.constexpr, + OTS: tl.constexpr, OHS: tl.constexpr, ODS: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + token = row // H + head = row - token * H + pair = tile * BLOCK + tl.arange(0, BLOCK) + valid = pair < HALF + + x_base = token * XTS + head * XHS + x_even = tl.load(x_ptr + x_base + (2 * pair) * XDS, + mask=valid, other=0.0).to(tl.float32) + x_odd = tl.load(x_ptr + x_base + (2 * pair + 1) * XDS, + mask=valid, other=0.0).to(tl.float32) + cosine = tl.load(cos_ptr + token * CTS + pair * CDS, + mask=valid, other=0.0).to(tl.float32) + sine = tl.load(sin_ptr + token * STS + pair * SDS, + mask=valid, other=0.0).to(tl.float32) + + out_even = x_even * cosine - x_odd * sine + out_odd = x_even * sine + x_odd * cosine + out_base = token * OTS + head * OHS + tl.store(out_ptr + out_base + (2 * pair) * ODS, out_even, mask=valid) + tl.store(out_ptr + out_base + (2 * pair + 1) * ODS, out_odd, mask=valid) + + +def _block_size(half): + if half <= 32: + return 32 + if half <= 64: + return 64 + if half <= 128: + return 128 + return 256 + + +def rotary_embedding(x, cos, sin, interleaved): + if x.ndim != 3: + raise ValueError('x must have shape [T, H, D]') + if not isinstance(interleaved, bool) or interleaved: + raise ValueError('Task35 requires interleaved=False') + if x.device.type in ('cpu', 'meta', 'mps'): + raise RuntimeError('a real Triton accelerator backend is required') + tokens, heads, width = x.shape + if width % 2: + raise ValueError('the last dimension must be even') + half = width // 2 + if cos.ndim != 2 or tuple(cos.shape) != (tokens, half): + raise ValueError('cos must have shape [T, D//2]') + if sin.ndim != 2 or tuple(sin.shape) != (tokens, half): + raise ValueError('sin must have shape [T, D//2]') + if cos.device != x.device or sin.device != x.device: + raise ValueError('x, cos and sin must be on the same device') + + output = torch.empty(x.shape, dtype=x.dtype, device=x.device) + if x.numel() == 0: + return output + block = _block_size(half) + grid = (tokens * heads, triton.cdiv(half, block)) + module = torch.get_device_module(x.device) + with module.device(x.device): + _rotary_pairs[grid]( + x, cos, sin, output, + H=heads, D=width, HALF=half, + XTS=x.stride(0), XHS=x.stride(1), XDS=x.stride(2), + CTS=cos.stride(0), CDS=cos.stride(1), + STS=sin.stride(0), SDS=sin.stride(1), + OTS=output.stride(0), OHS=output.stride(1), ODS=output.stride(2), + BLOCK=block, num_warps=4, enable_fp_fusion=False, + ) + return output + + +reference = rotary_embedding diff --git a/zz-flagos-s2-2026/d245/src/Task35_r6_zhizhi/rotary_embedding.zip b/zz-flagos-s2-2026/d245/src/Task35_r6_zhizhi/rotary_embedding.zip new file mode 100644 index 0000000..c9b9eea Binary files /dev/null and b/zz-flagos-s2-2026/d245/src/Task35_r6_zhizhi/rotary_embedding.zip differ diff --git a/zz-flagos-s2-2026/d245/src/silu_and_mul_masked.py b/zz-flagos-s2-2026/d245/src/silu_and_mul_masked.py new file mode 100644 index 0000000..5652255 --- /dev/null +++ b/zz-flagos-s2-2026/d245/src/silu_and_mul_masked.py @@ -0,0 +1,164 @@ +# Copyright 2026 FlagOS Contributors / GuanghuLab +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Optimized `reference` for FlagOS Open Computing Challenge S2 - Track 1 +# Task 39 (silu_and_mul_masked / DeepGEMM-style MoE). +# +# !! IMPORTANT: The function name MUST be exactly `reference` per the spec. !! +# The platform's grader imports the function by this name. +# +# Spec (from flagos.net): +# Input input: [E, T, H] bfloat16 +# masked_m: [E] int (per-expert valid row count) +# Output out: [E, T, H/2] bfloat16 (initialized to zero) +# For each expert e, let n = masked_m[e], half = H // 2: +# gate = input[e, :n, :half] (cast to float32) +# up = input[e, :n, half:] (cast to float32) +# out[e, :n] = (gate * sigmoid(gate)) * up (cast back to bfloat16) +# Rows [e, masked_m[e]:T] are padding — output values not checked. +# +# Reference (the slow one we beat): +# import torch +# def reference(input, masked_m): +# E, T, H = input.shape +# half = H // 2 +# out = torch.zeros(E, T, half, dtype=torch.bfloat16, device=input.device) +# for e in range(E): +# n = int(masked_m[e].item()) # <-- host sync, very slow +# if n <= 0: continue +# gate = input[e, :n, :half].float() +# up = input[e, :n, half:].float() +# val = gate * torch.sigmoid(gate) * up +# out[e, :n] = val.to(torch.bfloat16) +# return out +# +# Optimization strategy: +# 1. Single Triton kernel launch — no Python loop, no host sync +# 2. Load `masked_m` in-kernel (one int per program) — no .item() +# 3. Vectorize over (expert, row, hidden_col_half) with autotune +# 4. Compute in float32 for accuracy (matches reference math) +# 5. Cast back to bf16 on store + +import torch +import triton +import triton.language as tl + + +# --------------------------------------------------------------------------- +# Autotune +# --------------------------------------------------------------------------- +_AUTOTUNE_CONFIGS = [ + triton.Config({"BLOCK_M": 1, "BLOCK_N": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 1, "BLOCK_N": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 1, "BLOCK_N": 512}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 1, "BLOCK_N": 1024}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_M": 2, "BLOCK_N": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 2, "BLOCK_N": 512}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_M": 4, "BLOCK_N": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 4, "BLOCK_N": 512}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_M": 8, "BLOCK_N": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 8, "BLOCK_N": 256}, num_warps=8, num_stages=2), +] + + +# --------------------------------------------------------------------------- +# Triton kernel +# --------------------------------------------------------------------------- +@triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["E", "HIDDEN_HALF"]) +@triton.jit +def reference_kernel( + in_ptr, masked_m_ptr, out_ptr, + E, + s_e, s_t, s_h, # strides for input [E, T, H] + so_e, so_t, so_h, # strides for out [E, T, H/2] + HIDDEN_HALF: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_e = tl.program_id(0) # expert + pid_m = tl.program_id(1) # row block (over T) + pid_n = tl.program_id(2) # col block (over H/2) + + # Load this expert's valid row count (no host sync — pure device load) + valid_m = tl.load(masked_m_ptr + pid_e) + + # Row indices in this program + row_idx = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # (BLOCK_M,) + col_idx = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) # (BLOCK_N,) + + # Row mask: only valid rows compute + row_mask = row_idx < valid_m # (BLOCK_M,) + col_mask = col_idx < HIDDEN_HALF # (BLOCK_N,) + + # H is split in half: gate lives in [0, H/2), up in [H/2, H) + full_mask = row_mask[:, None] & col_mask[None, :] # (BLOCK_M, BLOCK_N) + + # 2D offsets into the input's last dim + base_in = pid_e * s_e + row_idx[:, None] * s_t + base_out = pid_e * so_e + row_idx[:, None] * so_t + + # gate and up offsets: same row, gate uses col_idx, up uses col_idx + H/2 + gate_offs = base_in + col_idx[None, :] * s_h + up_offs = base_in + (col_idx[None, :] + HIDDEN_HALF) * s_h + + g = tl.load(in_ptr + gate_offs, mask=full_mask, other=0.0) + u = tl.load(in_ptr + up_offs, mask=full_mask, other=0.0) + + # Compute in fp32 for accuracy (matches reference math) + g32 = g.to(tl.float32) + u32 = u.to(tl.float32) + y32 = g32 * tl.sigmoid(g32) * u32 + y = y32.to(tl.bfloat16) + + # Store to out + out_offs = base_out + col_idx[None, :] * so_h + tl.store(out_ptr + out_offs, y, mask=full_mask) + + +# --------------------------------------------------------------------------- +# Public API (function name MUST be `reference`) +# --------------------------------------------------------------------------- +def reference(input: torch.Tensor, masked_m: torch.Tensor) -> torch.Tensor: + """ + DeepGEMM-style SiLU(gate) * up with per-expert row mask. + + Args: + input: [E, T, H] bfloat16 (gate in [0:H/2], up in [H/2:H]) + masked_m: [E] int (valid row count per expert) + + Returns: + out: [E, T, H/2] bfloat16, zero-initialized + Padded rows (j >= masked_m[e]) are left as 0. + """ + assert input.dtype == torch.bfloat16, f"expected bf16, got {input.dtype}" + assert input.dim() == 3, f"expected 3D tensor, got {input.dim()}D" + E, T, H = input.shape + assert H % 2 == 0, f"H must be even, got {H}" + half = H // 2 + + assert masked_m.shape == (E,), f"masked_m shape mismatch: {masked_m.shape}" + masked_m_i32 = masked_m.to(torch.int32).contiguous() + + out = torch.zeros((E, T, half), dtype=torch.bfloat16, device=input.device) + + grid = lambda meta: ( + E, + triton.cdiv(T, meta["BLOCK_M"]), + triton.cdiv(half, meta["BLOCK_N"]), + ) + + with torch.cuda.device(input.device): + reference_kernel[grid]( + input, masked_m_i32, out, + E, + input.stride(0), input.stride(1), input.stride(2), + out.stride(0), out.stride(1), out.stride(2), + HIDDEN_HALF=half, + ) + + return out diff --git a/zz-flagos-s2-2026/d245/tests/test_correctness.py b/zz-flagos-s2-2026/d245/tests/test_correctness.py new file mode 100644 index 0000000..36bf717 --- /dev/null +++ b/zz-flagos-s2-2026/d245/tests/test_correctness.py @@ -0,0 +1,88 @@ +"""Correctness test for all v2 operators. + +Run with: python -m pytest tests/test_correctness.py -v +or: python tests/test_correctness.py +""" +import torch + + +def test_gelu_and_mul_v2_correctness(): + import sys + sys.path.insert(0, "src") + from flag_gems.fused.gelu_and_mul import gelu_and_mul as baseline + from flag_gems_local.fused.gelu_and_mul_v2 import gelu_and_mul as v2 + + torch.manual_seed(42) + for shape in [(1024, 1024), (4096, 4096), (2048, 11008)]: + for dtype in [torch.float16, torch.bfloat16, torch.float32]: + x = torch.randn(shape, dtype=dtype, device="cuda") + y = torch.randn(shape, dtype=dtype, device="cuda") + out_b = baseline(x, y) + out_v = v2(x, y) + abs_diff = (out_b - out_v).abs().max().item() + assert abs_diff < 1e-2, ( + f"gelu_and_mul v2 diverges: shape={shape} dtype={dtype} abs_diff={abs_diff}" + ) + print(f"gelu_and_mul shape={shape} dtype={dtype} abs_diff={abs_diff:.2e} ✓") + + +def test_rotary_embedding_v2_correctness(): + import sys + sys.path.insert(0, "src") + from flag_gems.fused.rotary_embedding import apply_rotary_pos_emb as baseline + from flag_gems_local.fused.rotary_embedding_v2 import apply_rotary_pos_emb as v2 + + torch.manual_seed(42) + for B, S, H, D in [(1, 1024, 32, 128), (2, 2048, 16, 64), (4, 512, 8, 128)]: + for interleaved in [False, True]: + q = torch.randn(B, S, H, D, dtype=torch.float16, device="cuda") + k = torch.randn(B, S, H, D, dtype=torch.float16, device="cuda") + cos = torch.randn(S, D // 2, dtype=torch.float16, device="cuda") + sin = torch.randn(S, D // 2, dtype=torch.float16, device="cuda") + qe_b, ke_b = baseline(q, k, cos, sin, rotary_interleaved=interleaved) + qe_v, ke_v = v2(q, k, cos, sin, rotary_interleaved=interleaved) + abs_diff_q = (qe_b - qe_v).abs().max().item() + abs_diff_k = (ke_b - ke_v).abs().max().item() + assert abs_diff_q < 1e-2, f"q diverges: {abs_diff_q}" + assert abs_diff_k < 1e-2, f"k diverges: {abs_diff_k}" + print( + f"rotary_emb B={B} S={S} H={H} D={D} " + f"interleaved={interleaved} abs_diff=({abs_diff_q:.2e}, {abs_diff_k:.2e}) ✓" + ) + + +def test_fused_moe_router_v2_correctness(): + import sys + sys.path.insert(0, "src") + from flag_gems_local.ops.fused_moe_router_v2 import ( + bitonic_sort_topk, + prepare_moe_inputs_v2, + ) + + torch.manual_seed(42) + for num_tokens, top_k in [(64, 4), (256, 8), (1024, 4)]: + weights = torch.rand(num_tokens, top_k, device="cuda") + ids = torch.randint(0, 8, (num_tokens, top_k), device="cuda") + + sorted_w, sorted_ids = bitonic_sort_topk(weights, ids) + + # Verify: weights should be descending + for i in range(num_tokens): + for j in range(top_k - 1): + assert sorted_w[i, j] >= sorted_w[i, j + 1] - 1e-5, ( + f"Not sorted: token {i} weight[{j}]={sorted_w[i,j]} > weight[{j+1}]={sorted_w[i,j+1]}" + ) + + # Verify: expert ids are a permutation of original + orig_ids = torch.sort(ids, dim=1).values + new_ids = torch.sort(sorted_ids, dim=1).values + assert torch.equal(orig_ids, new_ids), "ids not a permutation!" + + print(f"fused_moe_router tokens={num_tokens} top_k={top_k} ✓") + + +if __name__ == "__main__": + test_gelu_and_mul_v2_correctness() + test_rotary_embedding_v2_correctness() + test_fused_moe_router_v2_correctness() + print("\nAll tests passed ✓")