diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md index 5be24f7..fddafc6 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md @@ -1,5 +1,7 @@ # Fused MoE 算子入门:从 Benchmark 验证到 XPU-OJ 接口提交 +# Fused MoE 算子入门:从 Benchmark 验证到 XPU-OJ 接口提交 + ## 一、教程定位 本教程是赛题二 **Fused MoE** 任务的“benchmark 性能基线与 XPU-OJ 提交衔接”模块,主要帮助学员跑通目标算子的 benchmark 脚本,理解原库 API、输入输出结构、性能指标和性能基线结果,并进一步读懂 XPU-OJ 题目包中的接口约定、测试数据、参考输出和精度要求。 @@ -29,7 +31,7 @@ 3. 学习如何基于 Trition 与 MXMACA C++ 编写 Fused MOE 算子。 -4. 完成数值正确性测试,即验证 reference 计算、pybind 计算、Triton 计算这三种方式计算结果是否数值完全一致。 +4. 完成数值正确性测试,即验证 reference 计算、pybind 计算、Triton 计算这三种方式的结果是否在规定容差内一致。 * reference:基于 PyTorch 架构在 CPU 上运行的**数值基准**实现。 @@ -37,7 +39,7 @@ * Triton:基于 Python 编写的高效 GPU Kernel,可利用 Agent 自动调优,**开发效率高、易于迁移**。 - * 要求 pybind 和 Triton 结果均与 reference 一致,鼓励参赛选手持续调优 Triton ,使其性能逼近甚至超越 pybind 性能。 + * 要求 pybind 和 Triton 结果均在测试脚本规定的容差内与 reference 一致,鼓励参赛选手持续调优 Triton ,使其性能逼近甚至超越 pybind 性能。 5. 区分 benchmark 性能基线、OJ 参考实现和选手提交代码。 @@ -91,15 +93,15 @@ * 已获取 Fused MoE 源码。 -## 五、项目实践1-算力平台进行 Benchmark 验证 +## 五、项目实践:从 Benchmark 验证到 XPU-OJ 提交 -**项目目标:**跑通 Fused MoE 算子的 benchmark 脚本,建立性能基线,为后续算子优化提供对比基准。 +\*\*项目目标:\*\*依次完成 Fused MoE 算子的本地 benchmark 验证、XPU-OJ 冒烟提交、结果解读和候选版本管理,为后续性能优化建立可复现的基线。 ### 步骤 1:检查运行环境 -**目标:**确认当前环境满足本模块运行要求,包括编译器、MXMACA 工具链及 Python 依赖库。 +\*\*目标:\*\*确认当前环境满足本模块运行要求,包括编译器、MXMACA 工具链及 Python 依赖库。 -**操作:**检查 Python、编译工具、MXMACA 编译器及关键 Python 包(numpy、torch、triton)是否存在。 +\*\*操作:\*\*检查 Python、编译工具、MXMACA 编译器及关键 Python 包(numpy、torch、triton)是否存在。 **命令示例:** @@ -113,7 +115,9 @@ which mxcc # 确认 MACA 编译器存在 python - << 'EOF' import sys deps = ["numpy", "torch", "triton"] -missing = [] + +missing = [ ] + for d in deps: try: __import__(d) @@ -125,6 +129,7 @@ if missing: else: print("[OK] numpy, torch, triton are installed.") EOF + ``` **预期结果:** @@ -135,7 +140,7 @@ EOF * /opt/maca/mxgpu\_llvm/bin/mxcc -* \[OK\] numpy, torch, triton are installed. +* [OK] numpy, torch, triton are installed. **常见问题:** @@ -148,50 +153,60 @@ 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)。 -**操作:**切换到指定项目路径。 - -**命令示例:** - -```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 -``` +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 +``` ### 步骤 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 - - 且成功生成 `fused_moe/standalone/fused_moe_i8_tn/build/fused_moe_i8_tn_pybind.cpython-310-x86_64-linux-gnu.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` 文件 + **常见问题:** - | 报错 | 原因 | 解决办法 | | --- | --- | --- | | `Python.h: No such file or directory` | Python 头文件路径未找到 | 确认 `PYTHON_BIN` 路径正确,脚本自动探测 `sysconfig.get_path('include')` | @@ -199,15 +214,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 # 运行全部计算方式 @@ -218,44 +233,43 @@ 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: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\]=...“:首尾采样值,用于辅助定位数值偏差,不作为精度判定依据。 + **常见问题:** - | 报错 | 原因 | 解决办法 | | --- | --- | --- | @@ -264,56 +278,57 @@ 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”:预热轮数和正式迭代数。 - + **常见错误:** @@ -323,9 +338,7 @@ bash scripts/run_fused_moe_i8_tn_benchmark.sh --backend all --warmup 5 --iters 2 | 终端长时间无输出 | Kernel 死锁或 MACA 驱动异常 | 减小测试 shape;重启容器或设备 | | avg\_ms 异常抖动(±50%) | 其他进程占用 GPU | 关闭其他占用显存的进程,单机单任务运行 | -## 六、XPU-OJ 冒烟提交 - -### 6.1 目标 +### 步骤 6:理解 XPU-OJ 提交要求 在完成前文的本地验证后,本节将带你把实现提交到 XPU-OJ,并确认评测环境能够正确调用 `run_kernel(...)`。 @@ -333,37 +346,40 @@ 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/ + ``` 如果登录后看不到比赛或题目,请联系助教或赛事运营确认账号是否已经加入对应比赛或用户组。 -### 6.3 知识预备 +#### OJ 运行机制与精度要求 -#### 6.3.1 什么是 OJ +**OJ 如何评测提交代码** OJ 可以理解为“自动评测机”。 @@ -384,23 +400,29 @@ OJ 可以理解为“自动评测机”。 所以,OJ 不是让你提交 benchmark 日志,也不是让你提交本地运行截图,而是让你提交一份符合接口约定的代码。 -#### 6.3.2 什么是 Candidate +**精度要求:int8 计算、float32 累加与 bf16 写回** -Candidate 就是一次可复现的候选方案。 +与 Benchmark 的本地对拍类似,OJ 会把你写入 `out` 的结果与参考结果比较;本题的精度要求以题包为准: -建议每一轮都记录: +```python +torch.allclose(out_target.float(), out_ref.float(), rtol=0.0, atol=1e-2) -| 记录项 | 示例 | +``` + +`rtol=0.0` 表示不使用相对误差,所有元素都必须满足绝对误差不超过 `1e-2`。因此,结果不是比较 int8 累加值本身,而是比较完成缩放并写回 `bfloat16` 后的最终输出。 + +| 计算阶段 | 数据类型与要求 | | --- | --- | -| 候选编号 | candidate-001 | -| 代码文件 | `oj/problem_1_fused_moe/solution001.py` | -| 本地检查结果 | local check passed | -| OJ 结果 | WA / RE / AC | -| 备注 | 初始冒烟版,只验证接口 | +| 输入矩阵 `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`。 | -这样后续多次打榜时,不会忘记哪一版代码对应哪一次提交结果。 +本节提供的 TileLang 冒烟代码采用 `float32` 累加:先将 `a` 和 `b_col_major` 的元素转换为 `float32`,完成 K 维归约和缩放后再转换为 `bfloat16`。对于本题 `K=128`,`int8 × int8` 的乘积及其整数累加都可被 `float32` 精确表示;主要误差来自缩放计算以及最终写回 `bfloat16` 时的舍入。 -### 6.4 项目实践:Fused MoE GEMM OJ 冒烟提交 +常见错误包括:在乘加过程中提前转换为 `bfloat16`、将累加器声明为 `int8`,或在乘完 `scale_a` 后漏乘 `scale_b`、`moe_weights`。这些错误都会使最终 `out` 偏离参考结果,即使矩阵索引本身正确也会导致 `Wrong Answer`。 + +### 步骤 7:实现 Fused MoE GEMM OJ 冒烟代码 本节以当前 XPU-OJ 题目 **1. Fused MoE GEMM** 为例。 @@ -419,11 +441,12 @@ def run_kernel( out, ): ... + ``` 评测程序会调用这个函数,并检查你是否把结果正确写入 `out`。 -#### 步骤 1:理解本题要算什么 +#### 算子语义 本题计算的是固定规格 `fused_moe_i8_tn`: @@ -432,6 +455,7 @@ N = 128 K = 128 EM = num_tokens * topk EM % 128 == 0 + ``` 核心公式是: @@ -442,6 +466,7 @@ out[r, n] = * scale_a[token(r)] * scale_b[expert(r), n] * moe_weights[r] + ``` 两个索引最容易写错: @@ -449,6 +474,7 @@ out[r, n] = ```text token(r) = token_ids[r] // topk expert(r) = expert_ids[r // 128] + ``` 再说得直白一点: @@ -462,87 +488,168 @@ 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 + ``` -下一步不是手写一份固定答案,而是让 Agent 根据题面生成一份“最小冒烟版”。 +下一步直接使用下面的 TileLang 冒烟实现。它用于确认接口、索引和 OJ 提交流程,不追求性能最优。 -注意:冒烟版的目标只是确认接口、索引和 OJ 提交流程,不追求性能最优。 +#### TileLang OJ 冒烟代码 -#### 步骤 3:让 Agent 生成 OJ 冒烟代码 +将以下代码保存为 `oj/problem_1_fused_moe/solution001.py`,再复制到 XPU-OJ 提交框: -在 OpenCode 或其他 Agent 中,把题目的接口约定、输入输出说明和下面这段 Prompt 一起发给 Agent: +```python +import tilelang +import tilelang.language as T +from tilelang import jit -```text -我正在做 XPU-OJ 的 Fused MoE GEMM 题,需要生成一个最小冒烟提交版本。 +real_kernel = None +real_key = None -请根据下面接口写一份完整 Python 代码: -def run_kernel(a, b_col_major, scale_a, scale_b, moe_weights, token_ids, expert_ids, topk, out): - ... +@jit +def fused_moe_i8_tn_kernel(num_tokens, EM, num_experts): + """ + Compile a TileLang fused-MoE kernel for the given static shapes. -题目语义: -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. 正确性优先,不需要优化性能。 + 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. 函数名和参数顺序必须完全一致; -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` | +| token 索引 | 使用 `token_ids[r] // topk` | +| expert 索引 | 使用 `expert_ids[r // 128]` | +| 输出方式 | 原地写入 `out` | +| 提交语言 | 选择 `TileLang` | -如果 OJ 页面要求选择 `Triton` 或 `Triton Maca` 语言,也可以先让 Agent 生成一个语义正确的冒烟版本。正确性通过后,再让 Agent 把核心计算替换成 Triton kernel 做性能优化。 +### 步骤 8:提交并解读 XPU-OJ 结果 -#### 步骤 4:提交到 XPU-OJ +#### 提交到 XPU-OJ -1. 打开 XPU-OJ: ```texthttps://xpuoj.com/``` +1. 打开 XPU-OJ: \`\`\`texthttps://xpuoj.com/\`\`\` 2. 使用组委会发放的账号登录; -3. 进入比赛页面; [![image6](https://origin.picgo.net/2026/06/23/image6047f2ac4bd2a0f08.png)](https://www.picgo.net/image/image6.4ScJM4) +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) -4. 找到题目: ```text1. Fused MoE GEMM``` +4. 找到题目: `text1. Fused MoE GEMM` 5. 点击题目进入详情页; -6. 在提交区域选择本题支持的语言,例如: ```textTriton / Triton Maca``` +6. 在提交区域选择 `TileLang` 语言; ```plaintext @@ -556,12 +663,13 @@ oj/problem_1_fused_moe/solution001.py 3. 等待评测结果返回。 -#### 步骤 5:查看 OJ 结果 +#### 查看 OJ 结果 提交后,进入: ```text 我的提交 + ``` 常见状态含义如下: @@ -574,6 +682,8 @@ oj/problem_1_fused_moe/solution001.py | 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`,通常表示本次提交没有拿到分数。原因可能是: * 样例没过; @@ -589,7 +699,7 @@ oj/problem_1_fused_moe/solution001.py 如果看到用时和内存都是 `0`,很多时候说明代码在正式计时前就失败了,例如函数签名不匹配、导入失败或编译失败。 -#### 步骤 6:理解 OJ 评测流程 +#### 理解 OJ 评测流程 一次 OJ 提交通常会经历下面这些步骤: @@ -614,15 +724,11 @@ oj/problem_1_fused_moe/solution001.py 10. 在排行榜或提交记录中更新结果。 -本题的正确性校验口径是: +本题会比较写回 `out` 的 `bfloat16` 结果;完整的 int8 计算与精度口径见本步骤前的“OJ 运行机制与精度要求”。通过正确性后才会统计运行耗时。 -```python -torch.allclose(out_target.float(), out_ref.float(), rtol=0.0, atol=1e-2) -``` +### 步骤 9:保存 Candidate 并进入优化 -也就是说,OJ 允许很小的数值误差,但不是随便差一点都能过。 - -#### 步骤 7:保存 Candidate +Candidate 是一份可复现的候选方案:它应当对应一份固定代码、一次 OJ 提交及其结果。建议每轮至少记录候选编号、代码文件、本地检查结果、OJ 状态和备注,避免后续无法追溯某次分数对应的实现。 建议每一次能跑的版本都用 Git 保存。 @@ -632,87 +738,66 @@ 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` 继续改。 -### 6.5 Agent 使用说明 +#### 使用 Agent 定位问题与优化 本模块中,Agent 主要用来做三件事: -1. 读题目接口; +1. 解读现有 TileLang 冒烟代码与题目接口的对应关系; -2. 生成最小可提交代码; +2. 根据 OJ 结果定位正确性或运行时问题; -3. 根据 OJ 报错定位问题。 +3. 在正确性通过后分析性能瓶颈并提出优化方案。 建议不要一开始就让 Agent “直接写最快版本”。更稳的流程是: ```text -第一步:先写一个能过正确性的最小版本。 -第二步:提交 OJ,看是否 AC。 +第一步:提交本文提供的冒烟版本,确认接口和链路。 +第二步:根据 OJ 结果修复正确性或运行时问题。 第三步:AC 后再优化性能。 + ``` -可以使用下面的 Prompt: +排查问题时,可以将提交结果和当前代码一起交给 Agent: ```text 我正在做 XPU-OJ 的 Fused MoE GEMM 题。 -请只做一件事:根据题目接口写一个最小正确的 run_kernel 冒烟版本。 +当前使用的是 TileLang 冒烟版本,OJ 返回:[粘贴 AC / WA / RE / CE 和错误日志]。 -要求: -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. 不要做性能优化 -7. 不要添加 torch.Tensor 类型注解 -8. 不要依赖外部文件 -9. 请输出完整可复制提交的 Python 代码 +6. TileLang 编译、导入或 dtype 是否与题目环境兼容 + +请给出最小修改建议,并说明修改原因。 + ``` -如果 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? @@ -723,12 +808,14 @@ OJ 返回 Runtime Error。 ```python def silu(x: torch.Tensor) -> torch.Tensor: ... + ``` 这种类型注解可能触发: ```text Access to torch.Tensor is not allowed + ``` 处理方式:去掉 `torch.Tensor` 类型注解。 @@ -750,7 +837,7 @@ Access to torch.Tensor is not allowed 冒烟代码的目标是确认接口正确,不是追求性能。 -如果它能过正确性,但耗时很高,这是正常的。下一步才是把核心计算替换成 Triton kernel 或其他更快的 GPU 实现。 +如果它能过正确性,但耗时很高,这是正常的。下一步才是在保持上述 int8 计算与 bf16 输出精度口径不变的前提下,优化 TileLang kernel 或替换为更快的 GPU 实现。 #### Q4:50 分、10 分是什么意思? @@ -767,16 +854,9 @@ Access to torch.Tensor is not allowed 先看自己的提交是否通过正确性,再看耗时和分数。 -建议记录: +不要只看单次结果。按步骤 9 保存每轮的 Candidate、代码 diff、Prompt 和 OJ 结果,后面才能判断 Agent 的修改是否带来收益。 -| Candidate | OJ 状态 | 用时 | 分数 | 备注 | -| --- | --- | --- | --- | --- | -| candidate-001 | AC / WA / RE | 以页面为准 | 以页面为准 | 冒烟版 | -| candidate-002 | AC / WA / RE | 以页面为准 | 以页面为准 | 第一轮优化 | - -不要只看单次结果。每轮都记录,后面才知道 Agent 的修改到底有没有带来收益。 - -### 6.9 从 Benchmark 验证到参赛作品的路径回顾 +### 从 Benchmark 验证到参赛作品的路径回顾 建议按下面顺序推进: @@ -788,7 +868,7 @@ Access to torch.Tensor is not allowed 4. 如果冒烟代码 WA / RE,先修正确性; -5. 正确性通过后,再让 Agent 生成 Triton kernel 优化版本; +5. 正确性通过后,再让 Agent 针对现有 TileLang kernel 生成优化版本; 6. 每一轮提交都保存 candidate、prompt、代码 diff 和 OJ 结果; @@ -797,18 +877,18 @@ Access to torch.Tensor is not allowed ```text Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每一轮结果。 + ``` - -## 七、项目实践2-Kernel Swift 智能算子迁移系统自动调优 +## 六、后续实践:Kernel Swift 智能算子迁移系统自动调优 系统链接:[https://deeplink.org.cn/kernelswift/task](https://deeplink.org.cn/kernelswift/task) -**项目目标:**基于 KernelSwift 智能算子迁移系统,对 Fused MoE 算子进行在线自动调优。通过输入算子的 PyTorch 代码,一键生成适配沐曦硬件的高性能实现,高效完成算子优化与全流程追踪。 +\*\*项目目标:\*\*基于 KernelSwift 智能算子迁移系统,对 Fused MoE 算子进行在线自动调优。通过输入算子的 PyTorch 代码,一键生成适配沐曦硬件的高性能实现,高效完成算子优化与全流程追踪。 -### 步骤1:复用算子广场的Fused MoE 算子进行二次优化 +### 6.1 复用算子广场的Fused MoE 算子进行二次优化 -**目标:**通过提交算子广场的 fused\_moe 算子代码发起自动优化流程,实现二次优化 +\*\*目标:\*\*通过提交算子广场的 fused\_moe 算子代码发起自动优化流程,实现二次优化 **操作:** @@ -891,7 +971,10 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 return [hidden_states, w1, w2, topk_weights, topk_idx, top_k, True] def get_init_inputs(): - return [] + + return [ ] + + ``` 2. 进入新建任务页:点击左侧导航栏【新建任务】 ,进入算子提交页面。 @@ -911,15 +994,15 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 5. 提交优化任务:点击右下角 \[优化\] 按钮,系统将提交任务并进入 \[生成中\] 状态 - [![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/7c947aaf-ab14-4556-aebc-03d183f3746a.png)](https://www.picgo.net/image/image1.4SHrb4) 完成上述步骤将看到如下界面: - [![image2](https://origin.picgo.net/2026/06/23/image268924dc11f138788.png)](https://www.picgo.net/image/image2.4SHscu) +[![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/27b79d35-3c53-40b1-8db4-73b8897f7744.png)](https://www.picgo.net/image/image2.4SHscu) -### 步骤2:任务查看与结果管理 +### 6.2 任务查看与结果管理 -**目标:**在新建优化任务后可追踪任务进度,获取优化结果 +\*\*目标:\*\*在新建优化任务后可追踪任务进度,获取优化结果 **操作:** @@ -932,7 +1015,7 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 * 操作按钮:查看详情、删除任务 - [![image3](https://origin.picgo.net/2026/06/23/image33091601c9a68bd18.png)](https://www.picgo.net/image/image3.4SHDeY) + [![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/0c30baf1-e231-4922-84ec-bfde914433c4.png)](https://www.picgo.net/image/image3.4SHDeY) 2. 追踪任务进度:当前任务状态为【运行中】时,点击任务列表中的【查看详情】按钮,追踪任务进度: @@ -948,12 +1031,11 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 4. 性能调优:按设定的演化轮次迭代优化算子性能 - * 顶部:任务名称、创建/更新时间、适配硬件、当前轮次进度 - - [![image4](https://origin.picgo.net/2026/06/23/image403daf417d165a79f.png)](https://www.picgo.net/image/image4.4SHVpp) - + + [![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/915f7bb1-43a3-45f7-94fc-acc061efb364.png)](https://www.picgo.net/image/image4.4SHVpp) + 3. 获取优化结果:当前任务状态为【已完成】时,可在详情页查看优化结果: * 优化后算子代码支持一键复制 @@ -963,7 +1045,7 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 * 可点击【Diff 对比】查看优化前后代码差异,理解性能提升逻辑 - [![image5](https://origin.picgo.net/2026/06/23/image58f2b2ea36dad2ef0.png)](https://www.picgo.net/image/image5.4ScbBr) + [![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4j6OJ5j0d0Ne6q3p/img/d0b4267f-135e-4158-a022-7ffddcd78018.png)](https://www.picgo.net/image/image5.4ScbBr) 4. 任务异常处理 @@ -972,7 +1054,7 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 * 排队时间长:可调整提交时间,或联系平台管理员确认资源状态。 -## 八、Agent使用说明 +## 七、Agent使用说明 在本模块中,Agent可以帮助你完成以下任务: @@ -989,12 +1071,14 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 请帮我确认: 1. 当前环境是否满足编译与运行要求? 2. 是否有潜在的不兼容风险(如 Python 与 libpython 版本)? + ``` 2. **运行测试** ```plaintext 请帮我运行 scripts/run_fused_moe_i8_tn_pybind_test.sh 脚本 + ``` 3. **分析结果** @@ -1009,6 +1093,7 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 1. 为什么 pybind 比 Triton 快这么多? 2. TOPS 指标是否可信? 3. 当前结果是否已经具备提交价值? + ``` 4. **报错检查** @@ -1025,30 +1110,33 @@ 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 算子 -## 九、常见问题与注意事项 +``` -### 9.1 算力平台进行 Benchmark 验证 +## 八、常见问题与注意事项 + +### 8.1 算力平台进行 Benchmark 验证 1. 环境准备与依赖问题 * 确保算力平台已正确安装 Python 和 C++、MACA 编译器及相关运行时库,避免因环境缺失导致编译失败; - * 镜像环境使用 Conda Python​ 作为默认运行环境,避免系统 Python 与 Conda Python 混用,防止 `Python.h`或 `libpython`路径错误。 + * 镜像环境使用 Conda Python 作为默认运行环境,避免系统 Python 与 Conda Python 混用,防止 `Python.h`或 `libpython`路径错误。 2. pybind 编译与链接 @@ -1060,14 +1148,14 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 3. 性能测试建议 - * benchmark 应在关闭其他占用 GPU 的任务​后执行,避免干扰性能数据; + * benchmark 应在关闭其他占用 GPU 的任务后执行,避免干扰性能数据; * 多次运行取平均值,避免单次抖动影响结果; - * 性能对比应基于相同随机种子、相同 shape、相同 TopK、相同 batch size​的条件下进行,降低误差。 + * 性能对比应基于相同随机种子、相同 shape、相同 TopK、相同 batch size的条件下进行,降低误差。 -### 9.2 Kernel Swift 智能算子迁移系统自动调优项目 +### 8.2 Kernel Swift 智能算子迁移系统自动调优项目 1. 代码规范问题 @@ -1089,5 +1177,4 @@ Benchmark 验证代码用来学习,OJ 用来评分,Candidate 用来管理每 * 提交任务前确认目标硬件支持的算子类型; - * 优化失败时,可尝试更换适配硬件,或调整算子实现逻辑。 - + * 优化失败时,可尝试更换适配硬件,或调整算子实现逻辑。 \ No newline at end of file