forked from metax-maca/op_optimization
ADD file via upload
This commit is contained in:
parent
fba6cff3e5
commit
b50fab6a97
|
|
@ -0,0 +1,357 @@
|
|||
# Fused MoE GEMM 提交说明
|
||||
|
||||
## 当前结果
|
||||
|
||||
| Status | Score | Case | Time | Memory | Platform | Submit Time |
|
||||
| --- | ---: | --- | ---: | ---: | --- | --- |
|
||||
| Accepted | 50 | muxitest001 | 66 ms | 22.2 G | TileLang Maca C500 / 7.6 K | 06/24 16:35:18 |
|
||||
|
||||
得分说明：
|
||||
|
||||
- 50 分左右基本对应和题目 baseline 的加速比约为 `1:1`。
|
||||
- 当前 `50` 是 baseline 档位附近的 Accepted 结果。
|
||||
- 该结果主要用于确认提交接口、TileLang kernel 调用和输出正确性已经跑通。
|
||||
- 本文档只记录提交模板、改算子位置和当前结果，不展开进一步算子优化。
|
||||
|
||||
## 提交文件
|
||||
|
||||
提交文件为：
|
||||
|
||||
```text
|
||||
race_tests/moe/submission.py
|
||||
```
|
||||
|
||||
评测只要求 Python 文件中暴露 `run_kernel` 函数，函数名、参数顺序必须和题目一致：
|
||||
|
||||
```python
|
||||
def run_kernel(
|
||||
stacked_expert_tokens,
|
||||
gate_w,
|
||||
up_w,
|
||||
down_w,
|
||||
routed_expert_weights,
|
||||
group_sizes,
|
||||
group_offsets,
|
||||
group_padded_offsets,
|
||||
group_idx_for_bx,
|
||||
out,
|
||||
):
|
||||
...
|
||||
```
|
||||
|
||||
注意：`run_kernel` 签名不要加 `torch.Tensor` 类型标注。OJ 沙箱会限制访问 `torch.Tensor`，类型标注可能导致编译阶段失败。
|
||||
|
||||
提交时使用 `submission.py` 的内容即可，不需要提交 benchmark、reference 或测试脚本。
|
||||
|
||||
## 算子来源
|
||||
|
||||
当前提交模板基于 `race_tests/moe/custom_fusedmoe.py` 中的 routed fused MoE kernel 封装而来。
|
||||
|
||||
题目计算语义为：
|
||||
|
||||
```text
|
||||
gate_logits = x @ gate_w[expert]^T
|
||||
up_logits = x @ up_w[expert]^T
|
||||
hidden = silu(gate_logits) * up_logits
|
||||
output = hidden @ down_w[expert]^T
|
||||
output *= routed_expert_weights
|
||||
```
|
||||
|
||||
输入 `stacked_expert_tokens` 已经按 expert 分组，并按 `block_token = 128` 对齐 padding。
|
||||
|
||||
## 模板结构
|
||||
|
||||
`submission.py` 里主要有四部分：
|
||||
|
||||
```text
|
||||
_moe_forward_kernel(...) TileLang MoE fused GEMM kernel
|
||||
_get_kernel(...) 按 shape 缓存编译后的 kernel
|
||||
_get_workspace(...) 缓存 up_logits 中间 workspace
|
||||
run_kernel(...) OJ 调用入口，写入 out
|
||||
```
|
||||
|
||||
`run_kernel` 不做同步，不返回新 tensor，只负责取得 workspace、缓存 kernel 并调用：
|
||||
|
||||
```python
|
||||
up_logits = _get_workspace(stacked_expert_tokens, intermediate)
|
||||
kernel = _get_kernel(...)
|
||||
kernel(..., up_logits, out)
|
||||
```
|
||||
|
||||
## 实现要点
|
||||
|
||||
- `group_idx_for_bx[bx]` 用于把 token block 映射到 expert。
|
||||
- `stacked_expert_tokens`、`up_logits`、`out` 使用 padding 后的 token 下标。
|
||||
- `routed_expert_weights` 使用真实 token 顺序下标，即 `group_offsets[expert] + token_offset`。
|
||||
- padding 行不参与计算和写入。
|
||||
- `up_logits` 是中间 workspace，用于缓存 `silu(gate) * up`，避免 down 阶段重复计算。
|
||||
|
||||
## 改算子的位置
|
||||
|
||||
只改 MoE fused GEMM 时，主要看两个位置：
|
||||
|
||||
```text
|
||||
race_tests/moe/submission.py
|
||||
```
|
||||
|
||||
1. `_moe_forward_kernel(...)`
|
||||
- TileLang kernel 主体。
|
||||
- 第一阶段计算 gate/up 和 `up_logits`。
|
||||
- 第二阶段计算 down projection 并乘 routed weight。
|
||||
|
||||
2. `_get_kernel(...)`
|
||||
- 设置 shape cache key。
|
||||
- 控制不同 `(hidden, intermediate, num_experts, total_padded_tokens, total_valid_tokens, num_blocks_m)` 的 kernel 缓存。
|
||||
|
||||
一般不要改 `run_kernel` 的函数签名；评测器按固定签名调用。
|
||||
|
||||
## 完整算子代码
|
||||
|
||||
下面是当前 `race_tests/moe/submission.py` 的完整提交代码：
|
||||
|
||||
```python
|
||||
import torch
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
|
||||
|
||||
_KERNEL_CACHE = {}
|
||||
_WORKSPACE_CACHE = {}
|
||||
|
||||
|
||||
@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True})
|
||||
def _moe_forward_kernel(
|
||||
hidden,
|
||||
intermediate,
|
||||
num_experts,
|
||||
total_padded_tokens,
|
||||
total_valid_tokens,
|
||||
num_blocks_m,
|
||||
):
|
||||
scale = 1.44269504
|
||||
dtype = T.float16
|
||||
accum_dtype = T.float32
|
||||
block_token = 128
|
||||
block_dhidden = 128
|
||||
block_dexpert = 128
|
||||
threads = 256
|
||||
num_stages = 1
|
||||
|
||||
input_shape = (total_padded_tokens, hidden)
|
||||
intermediate_shape = (total_padded_tokens, intermediate)
|
||||
gate_shape = (num_experts, intermediate, hidden)
|
||||
up_shape = (num_experts, intermediate, hidden)
|
||||
down_shape = (num_experts, hidden, intermediate)
|
||||
|
||||
@T.prim_func
|
||||
def kernel(
|
||||
stacked_expert_tokens: T.Tensor(input_shape, dtype),
|
||||
gate_w: T.Tensor(gate_shape, dtype),
|
||||
up_w: T.Tensor(up_shape, dtype),
|
||||
down_w: T.Tensor(down_shape, dtype),
|
||||
routed_expert_weights: T.Tensor((total_valid_tokens,), T.float32),
|
||||
group_sizes: T.Tensor((num_experts,), T.int32),
|
||||
group_offsets: T.Tensor((num_experts + 1,), T.int32),
|
||||
group_padded_offsets: T.Tensor((num_experts + 1,), T.int32),
|
||||
group_idx_for_bx: T.Tensor((num_blocks_m,), T.int32),
|
||||
up_logits: T.Tensor(intermediate_shape, dtype),
|
||||
out: T.Tensor(input_shape, dtype),
|
||||
):
|
||||
with T.Kernel(num_blocks_m, T.ceildiv(intermediate, block_dexpert), threads=threads) as (bx, by):
|
||||
input_shared = T.alloc_fragment((block_token, block_dhidden), dtype=dtype)
|
||||
gate_shared = T.alloc_shared((block_dexpert, block_dhidden), dtype=dtype)
|
||||
up_shared = T.alloc_shared((block_dexpert, block_dhidden), dtype=dtype)
|
||||
gate_local = T.alloc_fragment((block_token, block_dexpert), dtype=accum_dtype)
|
||||
up_local = T.alloc_fragment((block_token, block_dexpert), dtype=accum_dtype)
|
||||
|
||||
T.use_swizzle(10)
|
||||
|
||||
expert_id = group_idx_for_bx[bx]
|
||||
block_start = bx * block_token
|
||||
group_size = group_sizes[expert_id]
|
||||
padded_start = group_padded_offsets[expert_id]
|
||||
actual_rows = T.max(0, T.min(block_token, group_size - (block_start - padded_start)))
|
||||
|
||||
T.clear(gate_local)
|
||||
T.clear(up_local)
|
||||
|
||||
for k in T.Pipelined(T.ceildiv(hidden, block_dhidden), num_stages=num_stages):
|
||||
T.copy(
|
||||
stacked_expert_tokens[
|
||||
block_start : block_start + block_token,
|
||||
k * block_dhidden : (k + 1) * block_dhidden,
|
||||
],
|
||||
input_shared,
|
||||
)
|
||||
T.copy(
|
||||
gate_w[
|
||||
expert_id,
|
||||
by * block_dexpert : (by + 1) * block_dexpert,
|
||||
k * block_dhidden : (k + 1) * block_dhidden,
|
||||
],
|
||||
gate_shared,
|
||||
)
|
||||
T.gemm(input_shared, gate_shared, gate_local, transpose_B=True)
|
||||
T.copy(
|
||||
up_w[
|
||||
expert_id,
|
||||
by * block_dexpert : (by + 1) * block_dexpert,
|
||||
k * block_dhidden : (k + 1) * block_dhidden,
|
||||
],
|
||||
up_shared,
|
||||
)
|
||||
T.gemm(input_shared, up_shared, up_local, transpose_B=True)
|
||||
|
||||
for i, j in T.Parallel(block_token, block_dexpert):
|
||||
gate_local[i, j] = gate_local[i, j] * (1.0 / (1.0 + T.exp2(-gate_local[i, j] * scale)))
|
||||
up_local[i, j] = up_local[i, j] * gate_local[i, j]
|
||||
|
||||
for i, j in T.Parallel(block_token, block_dexpert):
|
||||
if i < actual_rows:
|
||||
up_logits[block_start + i, by * block_dexpert + j] = up_local[i, j]
|
||||
|
||||
with T.Kernel(num_blocks_m, T.ceildiv(hidden, block_dhidden), threads=threads) as (bx, by):
|
||||
up_shared = T.alloc_fragment((block_token, block_dexpert), dtype=dtype)
|
||||
down_shared = T.alloc_shared((block_dhidden, block_dexpert), dtype=dtype)
|
||||
out_local = T.alloc_fragment((block_token, block_dhidden), dtype=accum_dtype)
|
||||
|
||||
T.use_swizzle(10)
|
||||
|
||||
expert_id = group_idx_for_bx[bx]
|
||||
block_start = bx * block_token
|
||||
group_size = group_sizes[expert_id]
|
||||
raw_start = group_offsets[expert_id]
|
||||
padded_start = group_padded_offsets[expert_id]
|
||||
token_offset = block_start - padded_start
|
||||
actual_rows = T.max(0, T.min(block_token, group_size - token_offset))
|
||||
|
||||
T.clear(out_local)
|
||||
|
||||
for k in T.Pipelined(T.ceildiv(intermediate, block_dexpert), num_stages=num_stages):
|
||||
T.copy(
|
||||
up_logits[
|
||||
block_start : block_start + block_token,
|
||||
k * block_dexpert : (k + 1) * block_dexpert,
|
||||
],
|
||||
up_shared,
|
||||
)
|
||||
T.copy(
|
||||
down_w[
|
||||
expert_id,
|
||||
by * block_dhidden : (by + 1) * block_dhidden,
|
||||
k * block_dexpert : (k + 1) * block_dexpert,
|
||||
],
|
||||
down_shared,
|
||||
)
|
||||
T.gemm(up_shared, down_shared, out_local, transpose_B=True)
|
||||
|
||||
for i, j in T.Parallel(block_token, block_dhidden):
|
||||
if i < actual_rows:
|
||||
out[block_start + i, by * block_dhidden + j] = (
|
||||
out_local[i, j] * routed_expert_weights[raw_start + token_offset + i]
|
||||
)
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
def _get_kernel(
|
||||
hidden,
|
||||
intermediate,
|
||||
num_experts,
|
||||
total_padded_tokens,
|
||||
total_valid_tokens,
|
||||
num_blocks_m,
|
||||
):
|
||||
key = (
|
||||
int(hidden),
|
||||
int(intermediate),
|
||||
int(num_experts),
|
||||
int(total_padded_tokens),
|
||||
int(total_valid_tokens),
|
||||
int(num_blocks_m),
|
||||
)
|
||||
kernel = _KERNEL_CACHE.get(key)
|
||||
if kernel is None:
|
||||
kernel = _moe_forward_kernel(*key)
|
||||
_KERNEL_CACHE[key] = kernel
|
||||
return kernel
|
||||
|
||||
|
||||
def _get_workspace(stacked_expert_tokens, intermediate):
|
||||
key = (
|
||||
int(stacked_expert_tokens.device.index or 0),
|
||||
int(stacked_expert_tokens.shape[0]),
|
||||
int(intermediate),
|
||||
str(stacked_expert_tokens.dtype),
|
||||
)
|
||||
up_logits = _WORKSPACE_CACHE.get(key)
|
||||
if up_logits is None:
|
||||
up_logits = torch.empty(
|
||||
(int(stacked_expert_tokens.shape[0]), int(intermediate)),
|
||||
device=stacked_expert_tokens.device,
|
||||
dtype=stacked_expert_tokens.dtype,
|
||||
)
|
||||
_WORKSPACE_CACHE[key] = up_logits
|
||||
return up_logits
|
||||
|
||||
|
||||
def run_kernel(
|
||||
stacked_expert_tokens,
|
||||
gate_w,
|
||||
up_w,
|
||||
down_w,
|
||||
routed_expert_weights,
|
||||
group_sizes,
|
||||
group_offsets,
|
||||
group_padded_offsets,
|
||||
group_idx_for_bx,
|
||||
out,
|
||||
):
|
||||
hidden = int(stacked_expert_tokens.shape[1])
|
||||
intermediate = int(gate_w.shape[1])
|
||||
num_experts = int(gate_w.shape[0])
|
||||
total_padded_tokens = int(stacked_expert_tokens.shape[0])
|
||||
total_valid_tokens = int(routed_expert_weights.shape[0])
|
||||
num_blocks_m = int(group_idx_for_bx.shape[0])
|
||||
|
||||
up_logits = _get_workspace(stacked_expert_tokens, intermediate)
|
||||
kernel = _get_kernel(
|
||||
hidden,
|
||||
intermediate,
|
||||
num_experts,
|
||||
total_padded_tokens,
|
||||
total_valid_tokens,
|
||||
num_blocks_m,
|
||||
)
|
||||
kernel(
|
||||
stacked_expert_tokens,
|
||||
gate_w,
|
||||
up_w,
|
||||
down_w,
|
||||
routed_expert_weights,
|
||||
group_sizes,
|
||||
group_offsets,
|
||||
group_padded_offsets,
|
||||
group_idx_for_bx,
|
||||
up_logits,
|
||||
out,
|
||||
)
|
||||
```
|
||||
|
||||
## 本地验证
|
||||
|
||||
本地验证时需要使用已经编译好的 TileLang，并把仓库根目录和 TVM Python 路径加入 `PYTHONPATH`。
|
||||
|
||||
在仓库根目录执行：
|
||||
|
||||
```bash
|
||||
export TILELANG_CACHE_DIR=${TILELANG_CACHE_DIR:-/tmp/tilelang-cache}
|
||||
export MACA_PATH=${MACA_PATH:-/opt/maca}
|
||||
export LD_LIBRARY_PATH="$(pwd)/build/lib:${MACA_PATH}/lib:${MACA_PATH}/mxgpu_llvm/lib:${LD_LIBRARY_PATH}"
|
||||
export PATH="${MACA_PATH}/bin:${MACA_PATH}/mxgpu_llvm/bin:${PATH}"
|
||||
export PYTHONPATH="$(pwd):$(pwd)/3rdparty/tvm/python:${PYTHONPATH}"
|
||||
|
||||
python race_tests/moe/test_submission.py
|
||||
```
|
||||
|
||||
不同机器的 Python 环境、MACA 安装路径和动态库路径可能不同，命令中的路径仅作为通用模板。
|
||||
Loading…
Reference in New Issue