From b50fab6a9743a7d125284daddb4a16b8b0c2e160 Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Thu, 25 Jun 2026 10:31:23 +0800 Subject: [PATCH] ADD file via upload --- .../赛题一moe教程 | 357 ++++++++++++++++++ 1 file changed, 357 insertions(+) create mode 100644 基于国产软件栈大模型推理前沿算子优化/赛题一moe教程 diff --git a/基于国产软件栈大模型推理前沿算子优化/赛题一moe教程 b/基于国产软件栈大模型推理前沿算子优化/赛题一moe教程 new file mode 100644 index 0000000..4f9b444 --- /dev/null +++ b/基于国产软件栈大模型推理前沿算子优化/赛题一moe教程 @@ -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 安装路径和动态库路径可能不同，命令中的路径仅作为通用模板。