From 9af23287493decdcb3ef7c061ffce27e0c910500 Mon Sep 17 00:00:00 2001 From: xxhefei Date: Thu, 25 Jun 2026 10:10:50 +0800 Subject: [PATCH 1/2] =?UTF-8?q?=E6=9B=B4=E6=96=B0moe=E6=AD=A5=E9=AA=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../MCTLASS_Fused MoE 算子优化.md | 66 ++++++++++--------- 1 file changed, 35 insertions(+), 31 deletions(-) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md index 5be24f7..7d0df35 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md @@ -91,9 +91,9 @@ * 已获取 Fused MoE 源码。 -## 五、项目实践1-算力平台进行 Benchmark 验证 +## 五、项目实践：从 Benchmark 验证到 XPU-OJ 提交 -**项目目标：**跑通 Fused MoE 算子的 benchmark 脚本，建立性能基线，为后续算子优化提供对比基准。 +**项目目标：**先跑通 Fused MoE 算子的 benchmark 脚本，建立性能基线；再基于同一题目语义完成 XPU-OJ 冒烟提交，确认评测环境能够正确调用 `run_kernel(...)`，为后续正确性修复和性能优化打基础。 ### 步骤 1：检查运行环境 @@ -323,9 +323,9 @@ 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：理解 XPU-OJ 提交要求 -### 6.1 目标 +**目标：**在完成前文的本地验证后，理解 XPU-OJ 的提交接口、评测方式和 Candidate 管理方式。 在完成前文的本地验证后，本节将带你把实现提交到 XPU-OJ，并确认评测环境能够正确调用 `run_kernel(...)`。 @@ -333,9 +333,9 @@ bash scripts/run_fused_moe_i8_tn_benchmark.sh --backend all --warmup 5 --iters 2 本节的冒烟提交只用于验证函数接口、索引逻辑和提交流程；性能优化请在正确性通过后再进行。 -### 6.2 前置准备 +#### 提交前准备 -#### 6.2.1 代码准备 +##### 代码准备 建议在工作目录下保留一个候选版本目录，例如： @@ -351,19 +351,19 @@ 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,7 +384,7 @@ OJ 可以理解为“自动评测机”。 所以，OJ 不是让你提交 benchmark 日志，也不是让你提交本地运行截图，而是让你提交一份符合接口约定的代码。 -#### 6.3.2 什么是 Candidate +##### 什么是 Candidate Candidate 就是一次可复现的候选方案。 @@ -400,7 +400,7 @@ Candidate 就是一次可复现的候选方案。 这样后续多次打榜时，不会忘记哪一版代码对应哪一次提交结果。 -### 6.4 项目实践：Fused MoE GEMM OJ 冒烟提交 +### 步骤 7：实现 Fused MoE GEMM OJ 冒烟代码 本节以当前 XPU-OJ 题目 **1. Fused MoE GEMM** 为例。 @@ -423,7 +423,7 @@ def run_kernel( 评测程序会调用这个函数，并检查你是否把结果正确写入 `out`。 -#### 步骤 1：理解本题要算什么 +#### 题目语义与索引规则 本题计算的是固定规格 `fused_moe_i8_tn`： @@ -462,7 +462,7 @@ expert(r) = expert_ids[r // 128] * 最终结果必须写回传入的 `out`。 -#### 步骤 2：准备 OJ 提交文件 +#### 准备 OJ 提交文件 在终端中创建目录： @@ -481,7 +481,7 @@ touch oj/problem_1_fused_moe/solution001.py 注意：冒烟版的目标只是确认接口、索引和 OJ 提交流程，不追求性能最优。 -#### 步骤 3：让 Agent 生成 OJ 冒烟代码 +#### 生成 OJ 冒烟代码 在 OpenCode 或其他 Agent 中，把题目的接口约定、输入输出说明和下面这段 Prompt 一起发给 Agent： @@ -530,9 +530,11 @@ oj/problem_1_fused_moe/solution001.py 如果 OJ 页面要求选择 `Triton` 或 `Triton Maca` 语言，也可以先让 Agent 生成一个语义正确的冒烟版本。正确性通过后，再让 Agent 把核心计算替换成 Triton kernel 做性能优化。 -#### 步骤 4：提交到 XPU-OJ +### 步骤 8：提交并解读 XPU-OJ 结果 -1. 打开 XPU-OJ： ```texthttps://xpuoj.com/``` +#### 提交到 XPU-OJ + +1. 打开 XPU-OJ： https://xpuoj.com/ 2. 使用组委会发放的账号登录； @@ -556,7 +558,7 @@ oj/problem_1_fused_moe/solution001.py 3. 等待评测结果返回。 -#### 步骤 5：查看 OJ 结果 +#### 查看 OJ 结果 提交后，进入： @@ -589,7 +591,7 @@ oj/problem_1_fused_moe/solution001.py 如果看到用时和内存都是 `0`，很多时候说明代码在正式计时前就失败了，例如函数签名不匹配、导入失败或编译失败。 -#### 步骤 6：理解 OJ 评测流程 +#### 理解 OJ 评测流程 一次 OJ 提交通常会经历下面这些步骤： @@ -622,7 +624,9 @@ torch.allclose(out_target.float(), out_ref.float(), rtol=0.0, atol=1e-2) 也就是说，OJ 允许很小的数值误差，但不是随便差一点都能过。 -#### 步骤 7：保存 Candidate +### 步骤 9：保存 Candidate 并进入优化 + +#### 保存 Candidate 建议每一次能跑的版本都用 Git 保存。 @@ -648,7 +652,7 @@ cp oj/problem_1_fused_moe/solution001.py oj/problem_1_fused_moe/solution002.py 然后让 Agent 基于 `solution002.py` 继续改。 -### 6.5 Agent 使用说明 +#### 使用 Agent 定位问题与优化 本模块中，Agent 主要用来做三件事： @@ -712,7 +716,7 @@ OJ 返回 Runtime Error。 只给出最小修复方案。 ``` -### 6.8 常见问题 +### 常见问题 #### Q1：为什么本地能跑，OJ 上却 Runtime Error？ @@ -776,7 +780,7 @@ Access to torch.Tensor is not allowed 不要只看单次结果。每轮都记录，后面才知道 Agent 的修改到底有没有带来收益。 -### 6.9 从 Benchmark 验证到参赛作品的路径回顾 +### 从 Benchmark 验证到参赛作品的路径回顾 建议按下面顺序推进： @@ -800,13 +804,13 @@ Benchmark 验证代码用来学习，OJ 用来评分，Candidate 用来管理每 ``` -## 七、项目实践2-Kernel Swift 智能算子迁移系统自动调优 +## 六、后续实践：Kernel Swift 智能算子迁移系统自动调优 系统链接：[https://deeplink.org.cn/kernelswift/task](https://deeplink.org.cn/kernelswift/task) **项目目标：**基于 KernelSwift 智能算子迁移系统，对 Fused MoE 算子进行在线自动调优。通过输入算子的 PyTorch 代码，一键生成适配沐曦硬件的高性能实现，高效完成算子优化与全流程追踪。 -### 步骤1：复用算子广场的Fused MoE 算子进行二次优化 +### 6.1 复用算子广场的Fused MoE 算子进行二次优化 **目标：**通过提交算子广场的 fused\_moe 算子代码发起自动优化流程，实现二次优化 @@ -917,7 +921,7 @@ Benchmark 验证代码用来学习，OJ 用来评分，Candidate 用来管理每 [![image2](https://origin.picgo.net/2026/06/23/image268924dc11f138788.png)](https://www.picgo.net/image/image2.4SHscu) -### 步骤2：任务查看与结果管理 +### 6.2 任务查看与结果管理 **目标：**在新建优化任务后可追踪任务进度，获取优化结果 @@ -972,7 +976,7 @@ Benchmark 验证代码用来学习，OJ 用来评分，Candidate 用来管理每 * 排队时间长：可调整提交时间，或联系平台管理员确认资源状态。 -## 八、Agent使用说明 +## 七、Agent使用说明 在本模块中，Agent可以帮助你完成以下任务： @@ -1040,9 +1044,9 @@ Benchmark 验证代码用来学习，OJ 用来评分，Candidate 用来管理每 请帮我在算子广场检索 fused_moe 算子 ``` -## 九、常见问题与注意事项 +## 八、常见问题与注意事项 -### 9.1 算力平台进行 Benchmark 验证 +### 8.1 算力平台进行 Benchmark 验证 1. 环境准备与依赖问题 @@ -1067,7 +1071,7 @@ Benchmark 验证代码用来学习，OJ 用来评分，Candidate 用来管理每 * 性能对比应基于相同随机种子、相同 shape、相同 TopK、相同 batch size​的条件下进行，降低误差。 -### 9.2 Kernel Swift 智能算子迁移系统自动调优项目 +### 8.2 Kernel Swift 智能算子迁移系统自动调优项目 1. 代码规范问题 From 63a5139d88c333166c87c1768ab67bb8460211ba Mon Sep 17 00:00:00 2001 From: xxhefei Date: Fri, 26 Jun 2026 17:17:58 +0800 Subject: [PATCH 2/2] =?UTF-8?q?=E6=9B=B4=E6=96=B0CUDA=E5=86=92=E7=83=9F?= =?UTF-8?q?=E4=BB=A3=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../MCTLASS_Fused MoE 算子优化.md | 523 ++++++++++++++---- 1 file changed, 422 insertions(+), 101 deletions(-) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md index 7d0df35..be3dd01 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/MCTLASS_Fused MoE 算子优化.md @@ -343,10 +343,10 @@ bash scripts/run_fused_moe_i8_tn_benchmark.sh --backend all --warmup 5 --iters 2 mkdir -p oj/problem_1_fused_moe ``` -本教程建议把 OJ 提交代码先保存为： +本教程建议把 OJ 提交源码先保存为： ```text -oj/problem_1_fused_moe/solution001.py +oj/problem_1_fused_moe/solution001.cu ``` 真正提交时，只需要把这个文件里的内容复制到 XPU-OJ 提交框。 @@ -393,48 +393,43 @@ Candidate 就是一次可复现的候选方案。 | 记录项 | 示例 | | --- | --- | | 候选编号 | candidate-001 | -| 代码文件 | `oj/problem_1_fused_moe/solution001.py` | +| 代码文件 | `oj/problem_1_fused_moe/solution001.cu` | | 本地检查结果 | local check passed | | OJ 结果 | WA / RE / AC | | 备注 | 初始冒烟版，只验证接口 | 这样后续多次打榜时，不会忘记哪一版代码对应哪一次提交结果。 -### 步骤 7：实现 Fused MoE GEMM OJ 冒烟代码 +### 步骤 7：实现 Fused MoE GEMM CUDA MACA 冒烟代码 本节以当前 XPU-OJ 题目 **1. Fused MoE GEMM** 为例。 -题目要求你提交一份 Python 代码，并提供固定函数： +题目要求你提交一份 CUDA 源码，并提供固定的 C 符号。评测程序会调用这个符号，并检查你是否把结果正确写入 `out`。 -```python -def run_kernel( - a, - b_col_major, - scale_a, - scale_b, - moe_weights, - token_ids, - expert_ids, - topk, - out, -): - ... +#### 接口约定 + +你必须在提交的 CUDA 源码中提供如下 C 符号，函数名、参数类型、顺序必须完全一致，并使用 `extern "C"` 防止 name mangling： + +```cpp +#include +#include + +extern "C" void run_kernel( + const int8_t* a, + const int8_t* b_col_major, + const float* scale_a, + const float* scale_b, + const float* moe_weights, + const int32_t* token_ids, + const int32_t* expert_ids, + int64_t topk, + __nv_bfloat16* out +); ``` -评测程序会调用这个函数，并检查你是否把结果正确写入 `out`。 - #### 题目语义与索引规则 -本题计算的是固定规格 `fused_moe_i8_tn`： - -```text -N = 128 -K = 128 -EM = num_tokens * topk -EM % 128 == 0 -``` - -核心公式是： +本题计算的是 `fused_moe_i8_tn` 形式的 W8A8 MoE GEMM。核心公式是： ```text out[r, n] = @@ -447,8 +442,8 @@ out[r, n] = 两个索引最容易写错： ```text -token(r) = token_ids[r] // topk -expert(r) = expert_ids[r // 128] +token(r) = token_ids[r] / topk +expert(r) = expert_ids[r / 128] ``` 再说得直白一点： @@ -471,64 +466,263 @@ cd /data/fusedmoe_v2.1 mkdir -p oj/problem_1_fused_moe ``` -新建文件： +新建 CUDA 源码文件： ```bash -touch oj/problem_1_fused_moe/solution001.py +touch oj/problem_1_fused_moe/solution001.cu ``` -下一步不是手写一份固定答案，而是让 Agent 根据题面生成一份“最小冒烟版”。 +注意：冒烟版的目标只是确认 CUDA MACA 提交接口、索引和 OJ 提交流程，不追求性能最优。 -注意：冒烟版的目标只是确认接口、索引和 OJ 提交流程，不追求性能最优。 +#### CUDA MACA 冒烟代码 -#### 生成 OJ 冒烟代码 - -在 OpenCode 或其他 Agent 中，把题目的接口约定、输入输出说明和下面这段 Prompt 一起发给 Agent： +将下面代码保存到： ```text -我正在做 XPU-OJ 的 Fused MoE GEMM 题，需要生成一个最小冒烟提交版本。 - -请根据下面接口写一份完整 Python 代码： - -def run_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 提交框的完整代码。 +oj/problem_1_fused_moe/solution001.cu ``` -Agent 生成后，把代码保存到： +```cpp +#include +#include +#include -```text -oj/problem_1_fused_moe/solution001.py +// xcore1000's CUDA-compatible compiler does not expose NVIDIA's __dp4a. +// This is a correctness-first replacement: each int32 stores four signed +// int8 values in little-endian byte order. +__device__ inline int32_t signed_byte(uint32_t x) { + x &= 0xffu; + return (int32_t)(x ^ 0x80u) - 128; +} + +__device__ inline int32_t dp4a_compat(int32_t a, int32_t b, int32_t acc) { + uint32_t ua = (uint32_t)a; + uint32_t ub = (uint32_t)b; + acc += signed_byte(ua) * signed_byte(ub); + acc += signed_byte(ua >> 8) * signed_byte(ub >> 8); + acc += signed_byte(ua >> 16) * signed_byte(ub >> 16); + acc += signed_byte(ua >> 24) * signed_byte(ub >> 24); + return acc; +} + +__global__ void w8a8_moe_gemm_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b_col_major, + const float* __restrict__ scale_a, + const float* __restrict__ scale_b, + const float* __restrict__ moe_weights, + const int32_t* __restrict__ token_ids, + const int32_t* __restrict__ expert_ids, + int K, int N, int topk, + __nv_bfloat16* __restrict__ out) +{ + int n_base = blockIdx.x * 128; + int m_base = blockIdx.y * 128; + int expert = expert_ids[blockIdx.y]; + + int tid = threadIdx.x; + int warp_id = tid / 32; + int lane_id = tid & 31; + + int warp_y = warp_id / 2; + int warp_x = warp_id & 1; + int my = lane_id / 8; + int mx = lane_id & 7; + + int m_idx[8]; + int n_idx[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + m_idx[i] = warp_y * 32 + my + i * 4; + } +#pragma unroll + for (int j = 0; j < 8; ++j) { + n_idx[j] = warp_x * 64 + mx + j * 8; + } + + __shared__ int32_t smem_A[2][128 * 17]; + __shared__ int32_t smem_B[2][128 * 17]; + + int32_t accum[8][8] = {0}; + +#pragma unroll + for (int step = 0; step < 2; ++step) { + int load_idx = step * 256 + tid; + int row = load_idx / 4; + int col_int4 = load_idx & 3; + + int r = m_base + row; + int token = token_ids[r] / topk; + int64_t a_idx = (int64_t)token * K; + int4 va = ((const int4*)(a + a_idx))[col_int4]; + + int sa = row * 17 + col_int4 * 4; + smem_A[0][sa + 0] = va.x; + smem_A[0][sa + 1] = va.y; + smem_A[0][sa + 2] = va.z; + smem_A[0][sa + 3] = va.w; + + int64_t b_idx = (int64_t)expert * N * K + (int64_t)(n_base + row) * K; + int4 vb = ((const int4*)(b_col_major + b_idx))[col_int4]; + + int sb = row * 17 + col_int4 * 4; + smem_B[0][sb + 0] = vb.x; + smem_B[0][sb + 1] = vb.y; + smem_B[0][sb + 2] = vb.z; + smem_B[0][sb + 3] = vb.w; + } + __syncthreads(); + + for (int k_outer = 0; k_outer < K; k_outer += 64) { + int comp_buf = (k_outer / 64) & 1; + int load_buf = 1 - comp_buf; + int next_k = k_outer + 64; + + if (next_k < K) { +#pragma unroll + for (int step = 0; step < 2; ++step) { + int load_idx = step * 256 + tid; + int row = load_idx / 4; + int col_int4 = load_idx & 3; + + int r = m_base + row; + int token = token_ids[r] / topk; + int64_t a_idx = (int64_t)token * K + next_k; + int4 va = ((const int4*)(a + a_idx))[col_int4]; + + int sa = row * 17 + col_int4 * 4; + smem_A[load_buf][sa + 0] = va.x; + smem_A[load_buf][sa + 1] = va.y; + smem_A[load_buf][sa + 2] = va.z; + smem_A[load_buf][sa + 3] = va.w; + + int64_t b_idx = (int64_t)expert * N * K + + (int64_t)(n_base + row) * K + next_k; + int4 vb = ((const int4*)(b_col_major + b_idx))[col_int4]; + + int sb = row * 17 + col_int4 * 4; + smem_B[load_buf][sb + 0] = vb.x; + smem_B[load_buf][sb + 1] = vb.y; + smem_B[load_buf][sb + 2] = vb.z; + smem_B[load_buf][sb + 3] = vb.w; + } + } + +#pragma unroll + for (int k_step = 0; k_step < 16; ++k_step) { + int32_t reg_A[8]; + int32_t reg_B[8]; + +#pragma unroll + for (int i = 0; i < 8; ++i) { + reg_A[i] = smem_A[comp_buf][m_idx[i] * 17 + k_step]; + } +#pragma unroll + for (int j = 0; j < 8; ++j) { + reg_B[j] = smem_B[comp_buf][n_idx[j] * 17 + k_step]; + } + +#pragma unroll + for (int i = 0; i < 8; ++i) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + accum[i][j] = dp4a_compat(reg_A[i], reg_B[j], accum[i][j]); + } + } + } + __syncthreads(); + } + + float scale_row[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + int r = m_base + m_idx[i]; + int token = token_ids[r] / topk; + scale_row[i] = scale_a[token] * moe_weights[r]; + } + + float scale_col[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + int n = n_base + n_idx[j]; + scale_col[j] = scale_b[(int64_t)expert * N + n]; + } + +#pragma unroll + for (int i = 0; i < 8; ++i) { + int r = m_base + m_idx[i]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + int n = n_base + n_idx[j]; + float v = (float)accum[i][j] * scale_row[i] * scale_col[j]; + out[(int64_t)r * N + n] = __float2bfloat16(v); + } + } +} + +static size_t device_allocation_size(const void* p) { + mcDrvDeviceptr_t base = 0; + size_t size = 0; + (void)wcuMemGetAddressRange(&base, &size, (mcDrvDeviceptr_t)(uintptr_t)p); + return size; +} + +extern "C" void run_kernel( + const int8_t* a, + const int8_t* b_col_major, + const float* scale_a, + const float* scale_b, + const float* moe_weights, + const int32_t* token_ids, + const int32_t* expert_ids, + int64_t topk, + __nv_bfloat16* out) +{ + size_t b_size = device_allocation_size(b_col_major); + size_t out_size = device_allocation_size(out); + + int N = 7168; + int K = 2048; + if (b_size > 5000000000ULL) { + N = 4096; + K = 7168; + } + + int EM = 4096; + if (out_size > 128ULL * 1024ULL * 1024ULL) { + EM = 32768; + } else if (out_size == 0) { + // Last-resort fallback if allocation-size probing is unavailable. + int32_t host_tokens[4096]; + cudaMemcpy(host_tokens, token_ids, sizeof(host_tokens), cudaMemcpyDeviceToHost); + int max_token_id = 0; + for (int i = 0; i < 4096; ++i) { + if (host_tokens[i] > max_token_id) { + max_token_id = host_tokens[i]; + } + } + if (max_token_id >= 4096) { + EM = 32768; + } + } + + dim3 block(256); + dim3 grid(N / 128, EM / 128); + w8a8_moe_gemm_kernel<<>>( + a, b_col_major, scale_a, scale_b, moe_weights, + token_ids, expert_ids, K, N, (int)topk, out); +} ``` 保存前建议人工检查 5 个点： | 检查项 | 应该满足 | | --- | --- | -| 函数名 | 必须是 `run_kernel` | +| C 符号 | 必须是 `extern "C" void 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` | - -如果 OJ 页面要求选择 `Triton` 或 `Triton Maca` 语言，也可以先让 Agent 生成一个语义正确的冒烟版本。正确性通过后，再让 Agent 把核心计算替换成 Triton kernel 做性能优化。 +| token 索引 | 必须使用 `token_ids[r] / topk` | +| expert 索引 | 必须使用 `expert_ids[r / 128]` 或等价的每 128 行一个 expert 逻辑 | +| 输出方式 | 必须原地写入 `out`，不要返回新指针 | ### 步骤 8：提交并解读 XPU-OJ 结果 @@ -544,14 +738,14 @@ oj/problem_1_fused_moe/solution001.py 5. 点击题目进入详情页； -6. 在提交区域选择本题支持的语言，例如： ```textTriton / Triton Maca``` +6. 在提交区域选择本题支持的语言，例如： ```textCUDA / CUDA MACA``` ```plaintext 具体名称以页面下拉框为准。 ``` -1. 将 `solution001.py` 中的代码复制到提交框； +1. 将 `solution001.cu` 中的源码复制到提交框； 2. 点击提交； @@ -573,9 +767,130 @@ oj/problem_1_fused_moe/solution001.py | Accepted / AC | 正确性通过 | 可以继续优化性能 | | Wrong Answer / WA | 输出数值不对 | 检查索引、shape、dtype、缩放系数 | | Runtime Error / RE | 运行时报错 | 点开详情看报错栈 | -| Compile Error / CE | 编译或加载失败 | 检查 import、语法、函数名 | +| Compile Error / CE | 编译或加载失败 | 检查头文件、语法、`extern "C"` 符号和函数签名 | | Time Limit Exceeded / TLE | 超时 | 说明代码太慢，需要做 kernel 优化 | +**测试点结果怎么看：** + +点开提交详情后，优先查看每个测试点是否通过正确性校验，再查看耗时和性能指标。一般按以下顺序判断： + +1. 先看是否通过： + + * 如果出现 `passed`、`pass=true` 或类似通过标记，说明该测试点数值正确； + + * 如果出现 `FAILED`、`Wrong Answer`、`allclose failed`、`Runtime Error`，说明当前实现还没有通过该测试点，不能只看性能时间。 + +2. 再看耗时： + + * `time_ms`：当前提交在该测试点上的运行时间，数值越小越好； + + * `speedup`：相对基线的加速比，数值越大越好； + + * `score_ratio`：当前测试点得分比例，越接近 1 说明越接近该测试点满分。 + +3. 最后看不同测试点之间的差异： + + * 如果某个测试点正确性失败，应先定位该测试点的 shape、topk、dtype、索引规则或容差要求； + + * 如果某个测试点正确但明显更慢，说明该场景下可能存在访存、重复计算或并行度不足问题； + + * 优化时应优先关注失败测试点和耗时占比较高的测试点。 + + +**分数怎么看：** + +XPU-OJ 结果通常需要同时关注正确性和性能分数： + +* `pass=true`：表示该提交通过正确性校验； + +* `pass=false`：表示该提交未通过，通常不会获得有效性能分； + +* `time_ms`：当前提交的运行时间，越小越好； + +* `speedup`：相对基线的加速比，越大越好； + +* `score_ratio`：当前测试点得分比例，越接近 1 说明越接近该测试点满分； + +* `allclose failed`：说明输出与参考结果误差超过阈值，应优先修正确性； + +* `Runtime Error`：说明编译、运行、越界访问或环境限制出错，应先解决运行问题。 + + +评分理解顺序： + +1. 先看 `pass`：不通过时先修正确性； + +2. 再看 `time_ms`：通过后比较耗时； + +3. 再看 `speedup`：判断相对基线是否有提升； + +4. 最后看 `score_ratio`：判断当前优化距离满分还有多远。 + +需要注意：OJ 分数通常不是只由单个测试点决定，而是多个测试点综合计算。因此优化时不要只盯一个最快 case，应优先解决失败测试点和耗时占比较高的测试点。 + +**是否需要优化怎么看：** + +读完 OJ 结果后，可以按下面的顺序判断下一步： + +| 现象 | 说明 | 下一步 | +| --- | --- | --- | +| 提交结果不通过 | 算子语义、索引、shape、dtype、scale、反量化或输出写回存在问题 | 先修正确性 | +| 正确性通过但 `time_ms` 很高 | 初版实现可用，但并行度、访存或计算复用不足 | 进入性能优化 | +| 某个测试点特别慢 | 该场景可能存在重复计算、负载不均或 block 配置不合适 | 针对该测试点单独分析 | +| 多次提交耗时抖动很大 | 可能受 GPU 占用、预热不足或评测波动影响 | 多提交几次或回到本地 benchmark 复查 | + +如果当前版本还没有通过正确性，不建议马上做性能优化。先让冒烟版本跑通，再进入下一轮 Candidate。 + +**怎么优化：** + +在正确性通过后，可以按以下方向逐步优化： + +1. 减少重复计算： + + * 检查同一个 token 或 expert 是否被重复计算； + + * 对 topk2 / topk3 场景，可考虑复用相同 token 的中间结果； + + * 避免每个输出元素都重复加载相同的输入向量。 + +2. 优化访存模式： + + * 尽量让连续线程访问连续内存； + + * 减少非合并访存； + + * 对频繁使用的权重、scale、token id、expert id 做局部缓存； + + * 避免不必要的全局内存读写。 + +3. 提高并行度： + + * 将输出矩阵按行、列或 expert 维度切分； + + * 对小 batch、小 token 场景，适当增加 block 数量； + + * 避免一个 kernel 只有很少 block，导致 GPU 利用率不足。 + +4. 优化计算粒度： + + * 合理设置 `BLOCK_M`、`BLOCK_N`、`BLOCK_K`； + + * 让 tile 大小匹配硬件并行能力和 shared memory 限制； + + * 对不同 topk 或不同 shape 可以采用不同 kernel 配置。 + +5. 对比基线逐步验证： + + * 每次只改一个优化点； + + * 修改后先跑正确性测试； + + * 正确后再跑 benchmark； + + * 记录每次 `time_ms`、`speedup`、`score_ratio` 变化，避免无效优化。 + +具体的 Candidate 保存和下一轮 Agent 优化方式，见步骤 9。 + 如果看到 `0 pts`，通常表示本次提交没有拿到分数。原因可能是： * 样例没过； @@ -589,7 +904,7 @@ oj/problem_1_fused_moe/solution001.py * 输出与参考答案超过容差。 -如果看到用时和内存都是 `0`，很多时候说明代码在正式计时前就失败了，例如函数签名不匹配、导入失败或编译失败。 +如果看到用时和内存都是 `0`，很多时候说明代码在正式计时前就失败了，例如 C 符号不匹配、编译失败或 kernel launch 失败。 #### 理解 OJ 评测流程 @@ -633,7 +948,7 @@ torch.allclose(out_target.float(), out_ref.float(), rtol=0.0, atol=1e-2) ```bash cd /data/fusedmoe_v2.1 git status --short -git add oj/problem_1_fused_moe/solution001.py +git add oj/problem_1_fused_moe/solution001.cu git commit -m "candidate 001 fused moe i8 tn oj smoke" git tag candidate-001-oj-smoke ``` @@ -647,10 +962,10 @@ git log --oneline --decorate -5 如果下一轮要继续优化，可以复制一份新文件： ```bash -cp oj/problem_1_fused_moe/solution001.py oj/problem_1_fused_moe/solution002.py +cp oj/problem_1_fused_moe/solution001.cu oj/problem_1_fused_moe/solution002.cu ``` -然后让 Agent 基于 `solution002.py` 继续改。 +然后让 Agent 基于 `solution002.cu` 继续改。 #### 使用 Agent 定位问题与优化 @@ -676,19 +991,28 @@ cp oj/problem_1_fused_moe/solution001.py oj/problem_1_fused_moe/solution002.py ```text 我正在做 XPU-OJ 的 Fused MoE GEMM 题。 -请只做一件事：根据题目接口写一个最小正确的 run_kernel 冒烟版本。 +请只做一件事：根据题目接口写一个最小正确的 CUDA MACA run_kernel 冒烟版本。 要求： -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] +1. C 符号和函数签名必须完全一致： + extern "C" void run_kernel( + const int8_t* a, + const int8_t* b_col_major, + const float* scale_a, + const float* scale_b, + const float* moe_weights, + const int32_t* token_ids, + const int32_t* expert_ids, + int64_t topk, + __nv_bfloat16* 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 代码 +7. 不要依赖外部文件 +8. 请输出完整可复制提交的 CUDA 源码 ``` 如果 OJ 返回 `Wrong Answer`，可以继续问： @@ -712,7 +1036,7 @@ OJ 返回 Runtime Error。 这是错误日志：[粘贴错误日志] -请先判断是函数签名、import、RestrictedPython 限制、dtype、shape 还是 GPU 算子调用问题。 +请先判断是 extern "C" 符号、头文件、编译选项、dtype、shape、越界访问还是 GPU kernel 调用问题。 只给出最小修复方案。 ``` @@ -720,22 +1044,19 @@ OJ 返回 Runtime Error。 #### Q1：为什么本地能跑，OJ 上却 Runtime Error？ -本地环境和 OJ 沙箱不完全一样。OJ 可能限制某些 Python 写法、外部文件访问或动态编译行为。 +本地环境和 OJ 编译环境不完全一样。OJ 可能对头文件、CUDA/MACA 编译参数、外部文件访问或动态链接行为有限制。 常见例子： -```python -def silu(x: torch.Tensor) -> torch.Tensor: +```cpp +extern "C" void run_kernel(...) { ... +} ``` -这种类型注解可能触发： +常见问题包括：没有使用 `extern "C"` 导致符号名被 C++ name mangling 改写、缺少必要头文件、调用了当前 OJ 编译环境不支持的 CUDA intrinsic，或 kernel launch 参数越界。 -```text -Access to torch.Tensor is not allowed -``` - -处理方式：去掉 `torch.Tensor` 类型注解。 +处理方式：先保证接口符号和参数类型完全一致，再根据错误日志逐项缩小范围。 #### Q2：为什么 OJ 是 Wrong Answer？ @@ -743,18 +1064,18 @@ Access to torch.Tensor is not allowed 1. `token_ids[r]` 是否先除以 `topk`； -2. `expert_ids` 是否按 `r // 128` 取； +2. `expert_ids` 是否按 `r / 128` 取； 3. `b_col_major` 是否按 `[expert, n, k]` 理解； -4. 结果是否写回 `out`，而不是只返回一个新 tensor。 +4. 结果是否写回 `out`，而不是写到临时 buffer 后没有拷回。 #### Q3：为什么冒烟代码很慢？ 冒烟代码的目标是确认接口正确，不是追求性能。 -如果它能过正确性，但耗时很高，这是正常的。下一步才是把核心计算替换成 Triton kernel 或其他更快的 GPU 实现。 +如果它能过正确性，但耗时很高，这是正常的。下一步才是在 CUDA MACA 源码中继续优化访存、并行度、tile 配置和计算复用。 #### Q4：50 分、10 分是什么意思？ @@ -792,7 +1113,7 @@ Access to torch.Tensor is not allowed 4. 如果冒烟代码 WA / RE，先修正确性； -5. 正确性通过后，再让 Agent 生成 Triton kernel 优化版本； +5. 正确性通过后，再让 Agent 在 CUDA MACA 版本上继续优化 kernel； 6. 每一轮提交都保存 candidate、prompt、代码 diff 和 OJ 结果；