forked from metax-maca/op_optimization
更新CUDA冒烟代码
This commit is contained in:
parent
9af2328749
commit
63a5139d88
|
|
@ -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 <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
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 <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
```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<<<grid, block>>>(
|
||||
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 结果；
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue