From fba6cff3e5d6b075842ec4bf3743a99372c1b933 Mon Sep 17 00:00:00 2001 From: Beckylu <648245013@qq.com> Date: Wed, 24 Jun 2026 17:52:49 +0800 Subject: [PATCH] =?UTF-8?q?Revert=20"Merge=20pull=20request=20'=E6=9B=B4?= =?UTF-8?q?=E6=96=B0moe=E6=95=99=E7=A8=8B=E5=92=8C=E7=BB=93=E6=9E=84'=20(#?= =?UTF-8?q?51)=20from=20xxhefei/op=5Foptimization:update-moark-agent-doc?= =?UTF-8?q?=20into=20master"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 2ea21d458d26bc06624766e1f75fb0c2c5da7033, reversing changes made to 30de9e7d28967dfbca56a8b16bef450dc542829b. --- .../MCTLASS_Fused MoE 算子优化.md | 545 ++++++++---------- 1 file changed, 229 insertions(+), 316 deletions(-) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md index fddafc6..5be24f7 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md @@ -1,7 +1,5 @@ # Fused MoE 算子入门:从 Benchmark 验证到 XPU-OJ 接口提交 -# Fused MoE 算子入门:从 Benchmark 验证到 XPU-OJ 接口提交 - ## 一、教程定位 本教程是赛题二 **Fused MoE** 任务的“benchmark 性能基线与 XPU-OJ 提交衔接”模块,主要帮助学员跑通目标算子的 benchmark 脚本,理解原库 API、输入输出结构、性能指标和性能基线结果,并进一步读懂 XPU-OJ 题目包中的接口约定、测试数据、参考输出和精度要求。 @@ -31,7 +29,7 @@ 3. 学习如何基于 Trition 与 MXMACA C++ 编写 Fused MOE 算子。 -4. 完成数值正确性测试,即验证 reference 计算、pybind 计算、Triton 计算这三种方式的结果是否在规定容差内一致。 +4. 完成数值正确性测试,即验证 reference 计算、pybind 计算、Triton 计算这三种方式计算结果是否数值完全一致。 * reference:基于 PyTorch 架构在 CPU 上运行的**数值基准**实现。 @@ -39,7 +37,7 @@ * Triton:基于 Python 编写的高效 GPU Kernel,可利用 Agent 自动调优,**开发效率高、易于迁移**。 - * 要求 pybind 和 Triton 结果均在测试脚本规定的容差内与 reference 一致,鼓励参赛选手持续调优 Triton ,使其性能逼近甚至超越 pybind 性能。 + * 要求 pybind 和 Triton 结果均与 reference 一致,鼓励参赛选手持续调优 Triton ,使其性能逼近甚至超越 pybind 性能。 5. 区分 benchmark 性能基线、OJ 参考实现和选手提交代码。 @@ -93,15 +91,15 @@ * 已获取 Fused MoE 源码。 -## 五、项目实践:从 Benchmark 验证到 XPU-OJ 提交 +## 五、项目实践1-算力平台进行 Benchmark 验证 -\*\*项目目标:\*\*依次完成 Fused MoE 算子的本地 benchmark 验证、XPU-OJ 冒烟提交、结果解读和候选版本管理,为后续性能优化建立可复现的基线。 +**项目目标:**跑通 Fused MoE 算子的 benchmark 脚本,建立性能基线,为后续算子优化提供对比基准。 ### 步骤 1:检查运行环境 -\*\*目标:\*\*确认当前环境满足本模块运行要求,包括编译器、MXMACA 工具链及 Python 依赖库。 +**目标:**确认当前环境满足本模块运行要求,包括编译器、MXMACA 工具链及 Python 依赖库。 -\*\*操作:\*\*检查 Python、编译工具、MXMACA 编译器及关键 Python 包(numpy、torch、triton)是否存在。 +**操作:**检查 Python、编译工具、MXMACA 编译器及关键 Python 包(numpy、torch、triton)是否存在。 **命令示例:** @@ -115,9 +113,7 @@ which mxcc # 确认 MACA 编译器存在 python - << 'EOF' import sys deps = ["numpy", "torch", "triton"] - -missing = [ ] - +missing = [] for d in deps: try: __import__(d) @@ -129,7 +125,6 @@ if missing: else: print("[OK] numpy, torch, triton are installed.") EOF - ``` **预期结果:** @@ -140,7 +135,7 @@ EOF * /opt/maca/mxgpu\_llvm/bin/mxcc -* [OK] numpy, torch, triton are installed. +* \[OK\] numpy, torch, triton are installed. **常见问题:** @@ -153,60 +148,50 @@ EOF ### 步骤 2:进入项目目录 -\*\*目标:\*\*进入本模块所需的源码目录 [benchmark](https://www.gitlink.org.cn/metax-maca/op_optimization/tree/master/%E5%9F%BA%E4%BA%8EAI%20Agent%E5%BC%80%E5%8F%91%E8%8C%83%E5%BC%8F%E7%9A%84%E5%9B%BD%E4%BA%A7GPU%E5%A4%A7%E6%A8%A1%E5%9E%8B%E6%8E%A8%E7%90%86%E7%AE%97%E5%AD%90%E5%BA%93%E4%BC%98%E5%8C%96%2Foperator_task_package%2Ffused_moe_task_package%2Fbenchmark)。 +**目标:**进入本模块所需的源码目录。 -1. 克隆代码仓库。 - - ```bash - git clone https://gitlink.org.cn/metax-maca/op\_optimization.git - ``` - -2. 准备 [fused\_moe\_task\_package](https://www.gitlink.org.cn/metax-maca/op_optimization/tree/master/%E5%9F%BA%E4%BA%8EAI%20Agent%E5%BC%80%E5%8F%91%E8%8C%83%E5%BC%8F%E7%9A%84%E5%9B%BD%E4%BA%A7GPU%E5%A4%A7%E6%A8%A1%E5%9E%8B%E6%8E%A8%E7%90%86%E7%AE%97%E5%AD%90%E5%BA%93%E4%BC%98%E5%8C%96%2Foperator_task_package%2Ffused_moe_task_package)。从仓库根目录开始,在 `基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package` 下找到该任务包,并复制到工作目录。 - - ```bash - mkdir -p data/operator_task_package - cp -r "基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/fused_moe_task_package" data/operator_task_package/ - ``` - -3. 切换到 `benchmark` 目录。 - +**操作:**切换到指定项目路径。 -```bash -cd data/operator_task_package/fused_moe_task_package/benchmark -ls +**命令示例:** + +```apl +#克隆代码仓库 +git clone https://gitlink.org.cn/metax-maca/op_optimization.git +#切换到fused moe目录下benchmark项目 +cd op_optimization/基于AI\ Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/fused_moe_task_package/benchmark ``` + ### 步骤 3:pybind 编译 -\*\*目标:\*\*将用 C++ 编写的 fused\_moe 算子编译为 Python 可调用的 pybind 模块。 +**目标:**将用 C++ 编写的 fused\_moe 算子编译为 Python 可调用的 pybind 模块。 -\*\*操作:\*\*运行 `fused_moe/scripts/build_fused_moe_i8_tn_pybind.sh` 脚本 +**操作:**运行 `fused_moe/scripts/build_fused_moe_i8_tn_pybind.sh` 脚本 **命令示例:** ```apl bash scripts/build_fused_moe_i8_tn_pybind.sh - ``` 切换 Python 环境命令示例: ```apl PYTHON_BIN=/path/to/python bash scripts/build_fused_moe_i8_tn_pybind.sh - ``` **预期结果:** 编译成功无报错,终端显示: -* \[SUCCESS\] /root/Project/fused\_moe/standalone/fused\_moe\_i8\_tn/build/fused\_moe\_i8\_tn\_ pybind.so + * \[SUCCESS\] /root/Project/fused\_moe/standalone/fused\_moe\_i8\_tn/build/fused\_moe\_i8\_tn\_ pybind.so + + 且成功生成 `fused_moe/standalone/fused_moe_i8_tn/build/fused_moe_i8_tn_pybind.cpython-310-x86_64-linux-gnu.so` 文件 - -且成功生成 `fused_moe/standalone/fused_moe_i8_tn/build/fused_moe_i8_tn_pybind.cpython-310-x86_64-linux-gnu.so` 文件 **常见问题:** + | 报错 | 原因 | 解决办法 | | --- | --- | --- | | `Python.h: No such file or directory` | Python 头文件路径未找到 | 确认 `PYTHON_BIN` 路径正确,脚本自动探测 `sysconfig.get_path('include')` | @@ -214,15 +199,15 @@ PYTHON_BIN=/path/to/python bash scripts/build_fused_moe_i8_tn_pybind.sh | `recompile with -fPIC` | 编译未开启位置无关代码 | 确保 `mxcc`/ `g++`编译参数中有 `-fPIC` | | `permission denied` | 无脚本执行权限 | `chmod +x scripts/*.sh` | | `undefined reference to Py_...` | Python 版本不匹配 | 确认编译脚本中`PYTHON_BIN`路径与当前运行的 Python 环境完全一致 | - + ### 步骤 4:正确性测试 - -\*\*目标:\*\*验证 reference 计算、pybind 计算、Triton 计算这三种方式的结果是否在测试容差内一致。 - -\*\*操作:\*\*运行 `fused_moe/scripts/run_fused_moe_i8_tn_pybind_test.sh` 脚本 - + +**目标:**验证 reference 计算、pybind 计算、Triton 计算这三种方式计算结果的数值是否一致。 + +**操作:**运行 `fused_moe/scripts/run_fused_moe_i8_tn_pybind_test.sh` 脚本 + **命令示例:** - + ```apl bash scripts/run_fused_moe_i8_tn_pybind_test.sh --backend all # 运行全部计算方式 @@ -233,43 +218,44 @@ bash scripts/run_fused_moe_i8_tn_pybind_test.sh --backend pybind bash scripts/run_fused_moe_i8_tn_pybind_test.sh --backend triton # 只测 reference: bash scripts/run_fused_moe_i8_tn_pybind_test.sh --backend reference - ``` - + **预期结果:** - + 编译成功无报错,输出示例如下: - + > pybind:fused\_moe\_i8\_tn\_topk1 passed: rows=256, cols=128, sample C\[0\]=0.69531, C\[last\]=-0.44531 - + > pybind:fused\_moe\_i8\_tn\_topk2 passed: rows=512, cols=128, sample C\[0\]=-0.57813, C\[last\]=-0.49805 - + > pybind:fused\_moe\_i8\_tn\_topk3 passed: rows=384, cols=128, sample C\[0\]=-1.08594, C\[last\]=-0.33594 - + > reference:fused\_moe\_i8\_tn\_topk1 passed: rows=256, cols=128, sample C\[0\]=0.6934, C\[last\]=-0.4451 - + > reference:fused\_moe\_i8\_tn\_topk2 passed: rows=512, cols=128, sample C\[0\]=-0.5768, C\[last\]=-0.4975 - + > reference:fused\_moe\_i8\_tn\_topk3 passed: rows=384, cols=128, sample C\[0\]=-1.0875, C\[last\]=-0.3362 - + > triton:fused\_moe\_i8\_tn\_topk1 passed: rows=256, cols=128, sample C\[0\]=0.69337, C\[last\]=-0.44513 - + > triton:fused\_moe\_i8\_tn\_topk2 passed: rows=512, cols=128, sample C\[0\]=-0.57678, C\[last\]=-0.49749 - + > triton:fused\_moe\_i8\_tn\_topk3 passed: rows=384, cols=128, sample C\[0\]=-1.08748, C\[last\]=-0.33618 - + **结果解释:** + + * “pybind/reference/Triton”:三种计算方式; + + * “fused\_moe\_i8\_tn\_topk1/2/3 passed”:测试算子通过数值校验,数值误差在允许范围内且无明显异常,否则会报错 FAILED; + + * ”rows=... , cols=...“:输出 Tensor 的形状; + + * ”sample C\[0\]=... , C\[last\]=...“:首尾采样值,用于辅助定位数值偏差,不作为精度判定依据。 + -* “pybind/reference/Triton”:三种计算方式; - -* “fused\_moe\_i8\_tn\_topk1/2/3 passed”:测试算子通过数值校验,数值误差在允许范围内且无明显异常,否则会报错 FAILED; - -* ”rows=... , cols=...“:输出 Tensor 的形状; - -* ”sample C\[0\]=... , C\[last\]=...“:首尾采样值,用于辅助定位数值偏差,不作为精度判定依据。 - **常见问题:** + | 报错 | 原因 | 解决办法 | | --- | --- | --- | @@ -278,57 +264,56 @@ bash scripts/run_fused_moe_i8_tn_pybind_test.sh --backend reference | `FAILED: shape mismatch` | 输出张量形状不一致 | 检查 Token Permute / Unpermute 逻辑;确认 expert 维度对齐 | | `FAILED: NaN or Inf detected` | 溢出或未初始化内存 | 检查 INT8 乘加是否溢出;确认 GEMM 输出是否反量化 | | 终端长时间无输出 | Kernel 死锁或 Launch 失败 | 减小测试 shape;检查是否触发 MACA 硬件限制 | - -### 步骤 5:性能测试 - -\*\*目标:\*\*输出 benchmark 结果对比表 - -\*\*操作:\*\*运行 `fused_moe/scripts/run_fused_moe_i8_tn_benchmark.sh` 脚本 - + +### 步骤5:性能测试 + +**目标:**输出 benchmark 结果对比表 + +**操作:**运行 `fused_moe/scripts/run_fused_moe_i8_tn_benchmark.sh` 脚本 + **命令示例:** - + ```apl bash scripts/run_fused_moe_i8_tn_benchmark.sh --backend all --warmup 5 --iters 20 # --backend:选择计算方式 # --warmup:设置预热次数 # --iters:设置迭代次数 - ``` - + **预期结果:** - + 编译成功无报错,输出示例如下: - + > pybind:fused\_moe\_i8\_tn\_topk1 benchmark: avg\_ms=0.308978, TOPS=0.027149, warmup=5, iters=20 - + > pybind:fused\_moe\_i8\_tn\_topk2 benchmark: avg\_ms=0.304500, TOPS=0.055098, warmup=5, iters=20 - + > pybind:fused\_moe\_i8\_tn\_topk3 benchmark: avg\_ms=0.297775, TOPS=0.042256, warmup=5, iters=20 - + > reference:fused\_moe\_i8\_tn\_topk1 benchmark: avg\_ms=1685.43, TOPS=0.000005, warmup=5, iters=20 - + > reference:fused\_moe\_i8\_tn\_topk2 benchmark: avg\_ms=3384.52, TOPS=0.000005, warmup=5, iters=20 - + > reference:fused\_moe\_i8\_tn\_topk3 benchmark: avg\_ms=2532.14, TOPS=0.000005, warmup=5, iters=20 - + > triton:fused\_moe\_i8\_tn\_topk1 benchmark: avg\_ms=19.013421, TOPS=0.000441, warmup=5, iters=20 - + > triton:fused\_moe\_i8\_tn\_topk2 benchmark: avg\_ms=16.745914, TOPS=0.001002, warmup=5, iters=20 - + > triton:fused\_moe\_i8\_tn\_topk3 benchmark: avg\_ms=19.630328, TOPS=0.000641, warmup=5, iters=20 - + **结果解释:** - + * “pybind/reference/Triton”:三种计算方式; - + * “fused\_moe\_i8\_tn\_topk1/2/3”:分别对应选择前 1 / 2 / 3 个专家场景下的 MoE 算子; - + * “avg\_ms”:平均算子执行耗时(毫秒),这里不计算预热时间,只计算正式迭代的时间; - + * “TOPS”:Tera Operations Per Second,本次 MoE 算子的总运算量 / 实际耗时; - + * “warmup=5, iters=20”:预热轮数和正式迭代数。 - + **常见错误:** @@ -338,7 +323,9 @@ bash scripts/run_fused_moe_i8_tn_benchmark.sh --backend all --warmup 5 --iters 2 | 终端长时间无输出 | Kernel 死锁或 MACA 驱动异常 | 减小测试 shape;重启容器或设备 | | avg\_ms 异常抖动(±50%) | 其他进程占用 GPU | 关闭其他占用显存的进程,单机单任务运行 | -### 步骤 6:理解 XPU-OJ 提交要求 +## 六、XPU-OJ 冒烟提交 + +### 6.1 目标 在完成前文的本地验证后,本节将带你把实现提交到 XPU-OJ,并确认评测环境能够正确调用 `run_kernel(...)`。 @@ -346,40 +333,37 @@ bash scripts/run_fused_moe_i8_tn_benchmark.sh --backend all --warmup 5 --iters 2 本节的冒烟提交只用于验证函数接口、索引逻辑和提交流程;性能优化请在正确性通过后再进行。 -#### 提交前准备 +### 6.2 前置准备 -**代码准备** +#### 6.2.1 代码准备 建议在工作目录下保留一个候选版本目录,例如: ```bash mkdir -p oj/problem_1_fused_moe - ``` 本教程建议把 OJ 提交代码先保存为: ```text oj/problem_1_fused_moe/solution001.py - ``` 真正提交时,只需要把这个文件里的内容复制到 XPU-OJ 提交框。 -**账号准备** +#### 6.2.2 账号准备 XPU-OJ 账号由组委会统一发放。登录入口: ```text https://xpuoj.com/ - ``` 如果登录后看不到比赛或题目,请联系助教或赛事运营确认账号是否已经加入对应比赛或用户组。 -#### OJ 运行机制与精度要求 +### 6.3 知识预备 -**OJ 如何评测提交代码** +#### 6.3.1 什么是 OJ OJ 可以理解为“自动评测机”。 @@ -400,29 +384,23 @@ OJ 可以理解为“自动评测机”。 所以,OJ 不是让你提交 benchmark 日志,也不是让你提交本地运行截图,而是让你提交一份符合接口约定的代码。 -**精度要求:int8 计算、float32 累加与 bf16 写回** +#### 6.3.2 什么是 Candidate -与 Benchmark 的本地对拍类似,OJ 会把你写入 `out` 的结果与参考结果比较;本题的精度要求以题包为准: +Candidate 就是一次可复现的候选方案。 -```python -torch.allclose(out_target.float(), out_ref.float(), rtol=0.0, atol=1e-2) +建议每一轮都记录: -``` - -`rtol=0.0` 表示不使用相对误差,所有元素都必须满足绝对误差不超过 `1e-2`。因此,结果不是比较 int8 累加值本身,而是比较完成缩放并写回 `bfloat16` 后的最终输出。 - -| 计算阶段 | 数据类型与要求 | +| 记录项 | 示例 | | --- | --- | -| 输入矩阵 `a`、`b_col_major` | 均为 `int8`;按 `[token, k]` 和 `[expert, n, k]` 读取。 | -| 矩阵乘加 | 每个输出元素累加 128 个 `int8 × int8` 乘积;可使用 `int32` 或 `float32` 累加,不能用 `int8` 累加。 | -| 反量化与路由加权 | 累加结果依次乘 `scale_a`、`scale_b`、`moe_weights`;三者均为 `float32`,建议在 `float32` 中完成。 | -| 输出 | 将最终结果转换为 `bfloat16`,并原地写入 `out`。 | +| 候选编号 | candidate-001 | +| 代码文件 | `oj/problem_1_fused_moe/solution001.py` | +| 本地检查结果 | local check passed | +| OJ 结果 | WA / RE / AC | +| 备注 | 初始冒烟版,只验证接口 | -本节提供的 TileLang 冒烟代码采用 `float32` 累加:先将 `a` 和 `b_col_major` 的元素转换为 `float32`,完成 K 维归约和缩放后再转换为 `bfloat16`。对于本题 `K=128`,`int8 × int8` 的乘积及其整数累加都可被 `float32` 精确表示;主要误差来自缩放计算以及最终写回 `bfloat16` 时的舍入。 +这样后续多次打榜时,不会忘记哪一版代码对应哪一次提交结果。 -常见错误包括:在乘加过程中提前转换为 `bfloat16`、将累加器声明为 `int8`,或在乘完 `scale_a` 后漏乘 `scale_b`、`moe_weights`。这些错误都会使最终 `out` 偏离参考结果,即使矩阵索引本身正确也会导致 `Wrong Answer`。 - -### 步骤 7:实现 Fused MoE GEMM OJ 冒烟代码 +### 6.4 项目实践:Fused MoE GEMM OJ 冒烟提交 本节以当前 XPU-OJ 题目 **1. Fused MoE GEMM** 为例。 @@ -441,12 +419,11 @@ def run_kernel( out, ): ... - ``` 评测程序会调用这个函数,并检查你是否把结果正确写入 `out`。 -#### 算子语义 +#### 步骤 1:理解本题要算什么 本题计算的是固定规格 `fused_moe_i8_tn`: @@ -455,7 +432,6 @@ N = 128 K = 128 EM = num_tokens * topk EM % 128 == 0 - ``` 核心公式是: @@ -466,7 +442,6 @@ out[r, n] = * scale_a[token(r)] * scale_b[expert(r), n] * moe_weights[r] - ``` 两个索引最容易写错: @@ -474,7 +449,6 @@ out[r, n] = ```text token(r) = token_ids[r] // topk expert(r) = expert_ids[r // 128] - ``` 再说得直白一点: @@ -488,168 +462,87 @@ expert(r) = expert_ids[r // 128] * 最终结果必须写回传入的 `out`。 -#### 准备并保存提交文件 +#### 步骤 2:准备 OJ 提交文件 在终端中创建目录: ```bash cd /data/fusedmoe_v2.1 mkdir -p oj/problem_1_fused_moe - ``` 新建文件: ```bash touch oj/problem_1_fused_moe/solution001.py - ``` -下一步直接使用下面的 TileLang 冒烟实现。它用于确认接口、索引和 OJ 提交流程,不追求性能最优。 +下一步不是手写一份固定答案,而是让 Agent 根据题面生成一份“最小冒烟版”。 -#### TileLang OJ 冒烟代码 +注意:冒烟版的目标只是确认接口、索引和 OJ 提交流程,不追求性能最优。 -将以下代码保存为 `oj/problem_1_fused_moe/solution001.py`,再复制到 XPU-OJ 提交框: +#### 步骤 3:让 Agent 生成 OJ 冒烟代码 -```python -import tilelang -import tilelang.language as T -from tilelang import jit +在 OpenCode 或其他 Agent 中,把题目的接口约定、输入输出说明和下面这段 Prompt 一起发给 Agent: -real_kernel = None -real_key = None +```text +我正在做 XPU-OJ 的 Fused MoE GEMM 题,需要生成一个最小冒烟提交版本。 +请根据下面接口写一份完整 Python 代码: -@jit -def fused_moe_i8_tn_kernel(num_tokens, EM, num_experts): - """ - Compile a TileLang fused-MoE kernel for the given static shapes. +def run_kernel(a, b_col_major, scale_a, scale_b, moe_weights, token_ids, expert_ids, topk, out): + ... - Grid : num_tiles = EM // 128 blocks - Block : N = 128 threads (thread t owns output column t) - - For every tile (block): - - All 128 rows share one expert: expert = expert_ids[tile_id] - - Thread t accumulates: - acc[t] = sum_k A[token(m), k] * B[expert, t, k] - then writes: - Out[row, t] = cast(acc[t] * scale_a[token] * scale_b[expert,t] - * moe_weights[row], bfloat16) - - topk is a runtime T.int32 scalar so the same compiled kernel works for - every topk value (1, 2, 3, 4, …) without recompilation. - """ - N = 128 - K = 128 - BM = 128 - num_tiles = EM // BM - - @T.prim_func - def kernel( - A: T.Tensor((num_tokens, K), T.int8), # [token, k] - B: T.Tensor((num_experts, N, K), T.int8), # [expert, n, k] col-major - SA: T.Tensor((num_tokens,), T.float32), # per-token scale_a - SB: T.Tensor((num_experts, N), T.float32), # per-expert per-col scale_b - MW: T.Tensor((EM,), T.float32), # moe_weights - TI: T.Tensor((EM,), T.int32), # token_ids - EI: T.Tensor((num_tiles,), T.int32), # expert_ids (one per tile) - topk: T.int32, # runtime scalar - Out: T.Tensor((EM, N), T.bfloat16), - ): - with T.Kernel(num_tiles, threads=N) as (tile_id,): - # 1 float32 register accumulator per thread (= per output column). - acc = T.alloc_fragment((N,), T.float32) - - # All rows in this tile share the same expert. - expert = EI[tile_id] - - # B[expert, :, :] is identical for all BM=128 rows in this tile. - # Cache each thread's own row (B[expert, t, :]) into shared memory - # once per tile instead of re-reading it from global memory on - # every one of the 128 row iterations below. - b_cache = T.alloc_shared((N, K), T.int8) - for n in T.Parallel(N): - for k in range(K): - b_cache[n, k] = B[expert, n, k] - - # Serial over BM=128 output rows. - for m in range(BM): - row = tile_id * BM + m - tok = TI[row] // topk # token_ids[r] // topk - - T.fill(acc, 0.0) - - # K-reduction: a_k is a block-wide broadcast (all threads - # read the same A[tok,k]), each thread t accumulates b_cache[t,k]. - for k in range(K): - a_k = T.cast(A[tok, k], T.float32) - for n in T.Parallel(N): - acc[n] = acc[n] + a_k * T.cast(b_cache[n, k], T.float32) - - # Dequantize and write bfloat16 output. - for n in T.Parallel(N): - Out[row, n] = T.cast( - acc[n] * SA[tok] * SB[expert, n] * MW[row], - T.bfloat16, - ) - - return kernel - - -def run_kernel( - a, # Tensor[int8] (num_tokens, 128) - b_col_major, # Tensor[int8] (num_experts, 128, 128) layout [e, n, k] - scale_a, # Tensor[float32] (num_tokens,) - scale_b, # Tensor[float32] (num_experts, 128) - moe_weights, # Tensor[float32] (EM,) - token_ids, # Tensor[int32] (EM,) - expert_ids, # Tensor[int32] (EM // 128,) - topk, # int (runtime; passed through to the compiled kernel) - out, # Tensor[bfloat16] (EM, 128) — written in-place -): - global real_kernel, real_key - - num_tokens = a.shape[0] - EM = out.shape[0] - num_experts = b_col_major.shape[0] - key = (num_tokens, EM, num_experts) - - if real_kernel is None or real_key != key: - real_kernel = fused_moe_i8_tn_kernel(num_tokens, EM, num_experts) - real_key = key - - # topk is passed as a runtime T.int32 scalar to the compiled kernel. - real_kernel(a, b_col_major, scale_a, scale_b, moe_weights, - token_ids, expert_ids, topk, out) +题目语义: +1. N = 128,K = 128; +2. EM = num_tokens * topk,且 EM 是 128 的倍数; +3. token(r) = token_ids[r] // topk; +4. expert(r) = expert_ids[r // 128]; +5. b_col_major 的布局是 [expert, n, k]; +6. 结果必须原地写入 out; +7. out 的 dtype 是 bfloat16; +8. 正确性优先,不需要优化性能。 +代码要求: +1. 函数名和参数顺序必须完全一致; +2. 不要添加 torch.Tensor 类型注解; +3. 不要依赖外部文件; +4. 不要打印调试信息; +5. 不要返回新 tensor,只写入 out; +6. 请输出一份可以直接复制到 XPU-OJ 提交框的完整代码。 ``` -提交前确认以下内容: +Agent 生成后,把代码保存到: + +```text +oj/problem_1_fused_moe/solution001.py +``` + +保存前建议人工检查 5 个点: | 检查项 | 应该满足 | | --- | --- | | 函数名 | 必须是 `run_kernel` | | 参数顺序 | 必须是 `a, b_col_major, scale_a, scale_b, moe_weights, token_ids, expert_ids, topk, out` | -| token 索引 | 使用 `token_ids[r] // topk` | -| expert 索引 | 使用 `expert_ids[r // 128]` | -| 输出方式 | 原地写入 `out` | -| 提交语言 | 选择 `TileLang` | +| token 索引 | 必须使用 `token_ids[r] // topk` | +| expert 索引 | 必须使用 `expert_ids[r // 128]` | +| 输出方式 | 必须原地写入 `out` | -### 步骤 8:提交并解读 XPU-OJ 结果 +如果 OJ 页面要求选择 `Triton` 或 `Triton Maca` 语言,也可以先让 Agent 生成一个语义正确的冒烟版本。正确性通过后,再让 Agent 把核心计算替换成 Triton kernel 做性能优化。 -#### 提交到 XPU-OJ +#### 步骤 4:提交到 XPU-OJ -1. 打开 XPU-OJ: \`\`\`texthttps://xpuoj.com/\`\`\` +1. 打开 XPU-OJ: ```texthttps://xpuoj.com/``` 2. 使用组委会发放的账号登录; -3. 进入比赛页面; [![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/44e87199-1017-43a1-9370-80195e3d2452.png)](https://www.picgo.net/image/image6.4ScJM4) +3. 进入比赛页面; [![image6](https://origin.picgo.net/2026/06/23/image6047f2ac4bd2a0f08.png)](https://www.picgo.net/image/image6.4ScJM4) -4. 找到题目: `text1. Fused MoE GEMM` +4. 找到题目: ```text1. Fused MoE GEMM``` 5. 点击题目进入详情页; -6. 在提交区域选择 `TileLang` 语言; +6. 在提交区域选择本题支持的语言,例如: ```textTriton / Triton Maca``` ```plaintext @@ -663,13 +556,12 @@ def run_kernel( 3. 等待评测结果返回。 -#### 查看 OJ 结果 +#### 步骤 5:查看 OJ 结果 提交后,进入: ```text 我的提交 - ``` 常见状态含义如下: @@ -682,8 +574,6 @@ def run_kernel( | Compile Error / CE | 编译或加载失败 | 检查 import、语法、函数名 | | Time Limit Exceeded / TLE | 超时 | 说明代码太慢,需要做 kernel 优化 | -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/77c57195-2ba6-4d85-bc4e-38af37d82e77.png) - 如果看到 `0 pts`,通常表示本次提交没有拿到分数。原因可能是: * 样例没过; @@ -699,7 +589,7 @@ def run_kernel( 如果看到用时和内存都是 `0`,很多时候说明代码在正式计时前就失败了,例如函数签名不匹配、导入失败或编译失败。 -#### 理解 OJ 评测流程 +#### 步骤 6:理解 OJ 评测流程 一次 OJ 提交通常会经历下面这些步骤: @@ -724,11 +614,15 @@ def run_kernel( 10. 在排行榜或提交记录中更新结果。 -本题会比较写回 `out` 的 `bfloat16` 结果;完整的 int8 计算与精度口径见本步骤前的“OJ 运行机制与精度要求”。通过正确性后才会统计运行耗时。 +本题的正确性校验口径是: -### 步骤 9:保存 Candidate 并进入优化 +```python +torch.allclose(out_target.float(), out_ref.float(), rtol=0.0, atol=1e-2) +``` -Candidate 是一份可复现的候选方案:它应当对应一份固定代码、一次 OJ 提交及其结果。建议每轮至少记录候选编号、代码文件、本地检查结果、OJ 状态和备注,避免后续无法追溯某次分数对应的实现。 +也就是说,OJ 允许很小的数值误差,但不是随便差一点都能过。 + +#### 步骤 7:保存 Candidate 建议每一次能跑的版本都用 Git 保存。 @@ -738,66 +632,87 @@ git status --short git add oj/problem_1_fused_moe/solution001.py git commit -m "candidate 001 fused moe i8 tn oj smoke" git tag candidate-001-oj-smoke - ``` 查看最近候选版本: ```bash git log --oneline --decorate -5 - ``` 如果下一轮要继续优化,可以复制一份新文件: ```bash cp oj/problem_1_fused_moe/solution001.py oj/problem_1_fused_moe/solution002.py - ``` 然后让 Agent 基于 `solution002.py` 继续改。 -#### 使用 Agent 定位问题与优化 +### 6.5 Agent 使用说明 本模块中,Agent 主要用来做三件事: -1. 解读现有 TileLang 冒烟代码与题目接口的对应关系; +1. 读题目接口; -2. 根据 OJ 结果定位正确性或运行时问题; +2. 生成最小可提交代码; -3. 在正确性通过后分析性能瓶颈并提出优化方案。 +3. 根据 OJ 报错定位问题。 建议不要一开始就让 Agent “直接写最快版本”。更稳的流程是: ```text -第一步:提交本文提供的冒烟版本,确认接口和链路。 -第二步:根据 OJ 结果修复正确性或运行时问题。 +第一步:先写一个能过正确性的最小版本。 +第二步:提交 OJ,看是否 AC。 第三步:AC 后再优化性能。 - ``` -排查问题时,可以将提交结果和当前代码一起交给 Agent: +可以使用下面的 Prompt: ```text 我正在做 XPU-OJ 的 Fused MoE GEMM 题。 -当前使用的是 TileLang 冒烟版本,OJ 返回:[粘贴 AC / WA / RE / CE 和错误日志]。 +请只做一件事:根据题目接口写一个最小正确的 run_kernel 冒烟版本。 -请基于现有代码定位问题,不要重写整份代码。重点检查: -1. 函数签名和参数顺序是否完全一致: +要求: +1. 函数签名必须完全一致: run_kernel(a, b_col_major, scale_a, scale_b, moe_weights, token_ids, expert_ids, topk, out) 2. token(r) = token_ids[r] // topk 3. expert(r) = expert_ids[r // 128] 4. b_col_major 的布局是 [expert, n, k] 5. 结果必须原地写入 out -6. TileLang 编译、导入或 dtype 是否与题目环境兼容 - -请给出最小修改建议,并说明修改原因。 - +6. 不要做性能优化 +7. 不要添加 torch.Tensor 类型注解 +8. 不要依赖外部文件 +9. 请输出完整可复制提交的 Python 代码 ``` -### 常见问题 +如果 OJ 返回 `Wrong Answer`,可以继续问: + +```text +OJ 返回 Wrong Answer。 + +请不要重写整份代码,先根据下面四点检查可能原因: +1. token_ids 是否正确除以 topk; +2. expert_ids 是否按每 128 行一个 expert 使用; +3. b_col_major 是否按 [expert, n, k] 读取; +4. 是否把结果写入 out,且 dtype 与 out 保持一致。 + +请给出最小修改建议。 +``` + +如果 OJ 返回 `Runtime Error`,可以问: + +```text +OJ 返回 Runtime Error。 + +这是错误日志:[粘贴错误日志] + +请先判断是函数签名、import、RestrictedPython 限制、dtype、shape 还是 GPU 算子调用问题。 +只给出最小修复方案。 +``` + +### 6.8 常见问题 #### Q1:为什么本地能跑,OJ 上却 Runtime Error? @@ -808,14 +723,12 @@ cp oj/problem_1_fused_moe/solution001.py oj/problem_1_fused_moe/solution002.py ```python def silu(x: torch.Tensor) -> torch.Tensor: ... - ``` 这种类型注解可能触发: ```text Access to torch.Tensor is not allowed - ``` 处理方式:去掉 `torch.Tensor` 类型注解。 @@ -837,7 +750,7 @@ Access to torch.Tensor is not allowed 冒烟代码的目标是确认接口正确,不是追求性能。 -如果它能过正确性,但耗时很高,这是正常的。下一步才是在保持上述 int8 计算与 bf16 输出精度口径不变的前提下,优化 TileLang kernel 或替换为更快的 GPU 实现。 +如果它能过正确性,但耗时很高,这是正常的。下一步才是把核心计算替换成 Triton kernel 或其他更快的 GPU 实现。 #### Q4:50 分、10 分是什么意思? @@ -854,9 +767,16 @@ Access to torch.Tensor is not allowed 先看自己的提交是否通过正确性,再看耗时和分数。 -不要只看单次结果。按步骤 9 保存每轮的 Candidate、代码 diff、Prompt 和 OJ 结果,后面才能判断 Agent 的修改是否带来收益。 +建议记录: -### 从 Benchmark 验证到参赛作品的路径回顾 +| Candidate | OJ 状态 | 用时 | 分数 | 备注 | +| --- | --- | --- | --- | --- | +| candidate-001 | AC / WA / RE | 以页面为准 | 以页面为准 | 冒烟版 | +| candidate-002 | AC / WA / RE | 以页面为准 | 以页面为准 | 第一轮优化 | + +不要只看单次结果。每轮都记录,后面才知道 Agent 的修改到底有没有带来收益。 + +### 6.9 从 Benchmark 验证到参赛作品的路径回顾 建议按下面顺序推进: @@ -868,7 +788,7 @@ Access to torch.Tensor is not allowed 4. 如果冒烟代码 WA / RE,先修正确性; -5. 正确性通过后,再让 Agent 针对现有 TileLang kernel 生成优化版本; +5. 正确性通过后,再让 Agent 生成 Triton kernel 优化版本; 6. 每一轮提交都保存 candidate、prompt、代码 diff 和 OJ 结果; @@ -877,18 +797,18 @@ Access to torch.Tensor is not allowed ```text Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每一轮结果。 - ``` -## 六、后续实践:Kernel Swift 智能算子迁移系统自动调优 + +## 七、项目实践2-Kernel Swift 智能算子迁移系统自动调优 系统链接:[https://deeplink.org.cn/kernelswift/task](https://deeplink.org.cn/kernelswift/task) -\*\*项目目标:\*\*基于 KernelSwift 智能算子迁移系统,对 Fused MoE 算子进行在线自动调优。通过输入算子的 PyTorch 代码,一键生成适配沐曦硬件的高性能实现,高效完成算子优化与全流程追踪。 +**项目目标:**基于 KernelSwift 智能算子迁移系统,对 Fused MoE 算子进行在线自动调优。通过输入算子的 PyTorch 代码,一键生成适配沐曦硬件的高性能实现,高效完成算子优化与全流程追踪。 -### 6.1 复用算子广场的Fused MoE 算子进行二次优化 +### 步骤1:复用算子广场的Fused MoE 算子进行二次优化 -\*\*目标:\*\*通过提交算子广场的 fused\_moe 算子代码发起自动优化流程,实现二次优化 +**目标:**通过提交算子广场的 fused\_moe 算子代码发起自动优化流程,实现二次优化 **操作:** @@ -971,10 +891,7 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 return [hidden_states, w1, w2, topk_weights, topk_idx, top_k, True] def get_init_inputs(): - - return [ ] - - + return [] ``` 2. 进入新建任务页:点击左侧导航栏【新建任务】 ,进入算子提交页面。 @@ -994,15 +911,15 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 5. 提交优化任务:点击右下角 \[优化\] 按钮,系统将提交任务并进入 \[生成中\] 状态 -[![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/7c947aaf-ab14-4556-aebc-03d183f3746a.png)](https://www.picgo.net/image/image1.4SHrb4) + [![image1](https://origin.picgo.net/2026/06/23/image1d46e08e5a17fd767.png)](https://www.picgo.net/image/image1.4SHrb4) 完成上述步骤将看到如下界面: -[![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/27b79d35-3c53-40b1-8db4-73b8897f7744.png)](https://www.picgo.net/image/image2.4SHscu) + [![image2](https://origin.picgo.net/2026/06/23/image268924dc11f138788.png)](https://www.picgo.net/image/image2.4SHscu) -### 6.2 任务查看与结果管理 +### 步骤2:任务查看与结果管理 -\*\*目标:\*\*在新建优化任务后可追踪任务进度,获取优化结果 +**目标:**在新建优化任务后可追踪任务进度,获取优化结果 **操作:** @@ -1015,7 +932,7 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 * 操作按钮:查看详情、删除任务 - [![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/0c30baf1-e231-4922-84ec-bfde914433c4.png)](https://www.picgo.net/image/image3.4SHDeY) + [![image3](https://origin.picgo.net/2026/06/23/image33091601c9a68bd18.png)](https://www.picgo.net/image/image3.4SHDeY) 2. 追踪任务进度:当前任务状态为【运行中】时,点击任务列表中的【查看详情】按钮,追踪任务进度: @@ -1031,11 +948,12 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 4. 性能调优:按设定的演化轮次迭代优化算子性能 + * 顶部:任务名称、创建/更新时间、适配硬件、当前轮次进度 - - [![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/915f7bb1-43a3-45f7-94fc-acc061efb364.png)](https://www.picgo.net/image/image4.4SHVpp) - + + [![image4](https://origin.picgo.net/2026/06/23/image403daf417d165a79f.png)](https://www.picgo.net/image/image4.4SHVpp) + 3. 获取优化结果:当前任务状态为【已完成】时,可在详情页查看优化结果: * 优化后算子代码支持一键复制 @@ -1045,7 +963,7 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 * 可点击【Diff 对比】查看优化前后代码差异,理解性能提升逻辑 - [![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/d0b4267f-135e-4158-a022-7ffddcd78018.png)](https://www.picgo.net/image/image5.4ScbBr) + [![image5](https://origin.picgo.net/2026/06/23/image58f2b2ea36dad2ef0.png)](https://www.picgo.net/image/image5.4ScbBr) 4. 任务异常处理 @@ -1054,7 +972,7 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 * 排队时间长:可调整提交时间,或联系平台管理员确认资源状态。 -## 七、Agent使用说明 +## 八、Agent使用说明 在本模块中,Agent可以帮助你完成以下任务: @@ -1071,14 +989,12 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 请帮我确认: 1. 当前环境是否满足编译与运行要求? 2. 是否有潜在的不兼容风险(如 Python 与 libpython 版本)? - ``` 2. **运行测试** ```plaintext 请帮我运行 scripts/run_fused_moe_i8_tn_pybind_test.sh 脚本 - ``` 3. **分析结果** @@ -1093,7 +1009,6 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 1. 为什么 pybind 比 Triton 快这么多? 2. TOPS 指标是否可信? 3. 当前结果是否已经具备提交价值? - ``` 4. **报错检查** @@ -1110,33 +1025,30 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 1. 错误原因是什么? 2. 如何用 find 命令定位 libpython3.10.so? 3. 如何在 build_fused_moe_i8_tn_pybind.sh 中正确指定路径? - ``` 5. **代码理解** ```plaintext 请帮我梳理释 benchmark_fused_moe_i8_tn.py 代码整体框架 - ``` 6. **KernelSwift 系统搜索算子** -```plaintext -请帮我在算子广场检索 fused_moe 算子 + ```plaintext + 请帮我在算子广场检索 fused_moe 算子 + ``` -``` +## 九、常见问题与注意事项 -## 八、常见问题与注意事项 - -### 8.1 算力平台进行 Benchmark 验证 +### 9.1 算力平台进行 Benchmark 验证 1. 环境准备与依赖问题 * 确保算力平台已正确安装 Python 和 C++、MACA 编译器及相关运行时库,避免因环境缺失导致编译失败; - * 镜像环境使用 Conda Python 作为默认运行环境,避免系统 Python 与 Conda Python 混用,防止 `Python.h`或 `libpython`路径错误。 + * 镜像环境使用 Conda Python​ 作为默认运行环境,避免系统 Python 与 Conda Python 混用,防止 `Python.h`或 `libpython`路径错误。 2. pybind 编译与链接 @@ -1148,14 +1060,14 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 3. 性能测试建议 - * benchmark 应在关闭其他占用 GPU 的任务后执行,避免干扰性能数据; + * benchmark 应在关闭其他占用 GPU 的任务​后执行,避免干扰性能数据; * 多次运行取平均值,避免单次抖动影响结果; - * 性能对比应基于相同随机种子、相同 shape、相同 TopK、相同 batch size的条件下进行,降低误差。 + * 性能对比应基于相同随机种子、相同 shape、相同 TopK、相同 batch size​的条件下进行,降低误差。 -### 8.2 Kernel Swift 智能算子迁移系统自动调优项目 +### 9.2 Kernel Swift 智能算子迁移系统自动调优项目 1. 代码规范问题 @@ -1177,4 +1089,5 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 * 提交任务前确认目标硬件支持的算子类型; - * 优化失败时,可尝试更换适配硬件,或调整算子实现逻辑。 \ No newline at end of file + * 优化失败时,可尝试更换适配硬件,或调整算子实现逻辑。 +