# 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 安装路径和动态库路径可能不同，命令中的路径仅作为通用模板。
