MetaX-TileLang工具链

MetaX-TileLang工具链

1.工具链层次

目录

目录
1.1层次概览
1.2测试入口
1.3cuda兼容层
1.4TileLang DSL
1.5TVM和TIR
1.6MACA backend后端
1.7mxcc编译器
1.8JIT运行时编译
1.9Runtime
1.10Driver
1.11GPU
1.12benchmark/smi/mcProfiler
1.13初步的日志结构
1.14总结:在每层该干什么

1.1层次概览

race_tests // PyTorch 输入输出,比赛测试框架,规定  输入 shape 是什么
        ↓
mcPyTorch // torch.cuda 兼容接口
        ↓
TileLang Python DSL //TileLang语言
        ↓
TVM // TIR IR 张量中间表示
        ↓
TileLang MACA backend //接入MACA后端
        ↓
mxcc //沐曦编译器驱动
        ↓
MXMACA Runtime//运行时 API
        ↓
MetaX Driver//设备驱动
        ↓
曦云 C500 GPU//设备执行
        ↓
mcProfiler / mx-smi / benchmark 反向分析 //运行后定位回去,确定问题,之后就可以做出优化

MetaX官方开源算子MoE,MLA,NSA

1.2测试入口

TileLang 本身是一个面向高性能 GPU/CPU kernel 的 Python DSL,底层依赖 TVM 编译基础设施
测试入口:
race_tests/mla/...
race_tests/moe/...
race_tests/nsa/...

在进行比赛时,需要注意

dtype 是什么
layout 是什么
误差容忍度是多少
benchmark 怎么计时
哪些 case 必须通过

1.3cuda兼容层

torch.cuda
很多上层 AI 框架、测试脚本和生态代码历史上都是围绕 PyTorch + CUDA 写的。
所以比赛环境里通常会保留类似:

import torch
x = torch.randn(..., device="cuda")

虽然表面看像 CUDA,实际上底层可能被 mcPyTorch / MACA 兼容层接管。

torch.Tensor
torch.cuda.Event
torch.cuda.synchronize
torch.empty(..., device="cuda")

这些接口可能仍然存在,但它们背后不一定走 NVIDIA CUDA runtime,而是走沐曦的兼容实现。
所以这里的重点是:
PyTorch 负责张量管理、输入输出、reference 对比、benchmark 计时
TileLang 负责自己写的高性能 kernel
MACA 兼容层负责把 PyTorch CUDA 风格调用接到 MetaX 后端
不能简单地把 torch.cuda 当成 NVIDIA CUDA。它在这个环境里更像是一个兼容 API 外壳。
这层的意义主要有三个。
第一,输入输出张量来自 PyTorch。
例如 race_tests 可能给你:

q: torch.Tensor
k: torch.Tensor
v: torch.Tensor
out: torch.Tensor

写的 TileLang kernel 最终要接收这些 tensor 的 device pointer。
第二,benchmark 时间可能由 torch.cuda.Event 或 synchronize 测出来。
所以要搞清楚 benchmark 的计时边界:
是否包含数据预处理?
是否包含 kernel compile?
是否包含 torch reference?
是否只统计 kernel launch 到完成?
是否做了 warmup?
第三,PyTorch 端可能有隐藏开销。
例如 MoE 里如果在 PyTorch 侧做了很多:

sort
index_select
scatter
nonzero
topk
reshape
contiguous

这些操作可能都会成为额外 kernel 或额外 memory movement。比赛可能只允许替换某个核心函数,也可能允许改更大范围。这个必须看 race_tests 的调用边界。

1.4TileLang DSL

TileLang DSL 是主要工作的层,它是用 Python 风格语法描述 tile-based GPU kernel,然后让编译器生成底层代码。官方文档说,TileLang 使用 Pythonic syntax,并基于 TVM 编译基础设施,让开发者在保留底层优化能力的同时提高生产效率。
TileLang kernel 通常会描述这些东西:

block 怎么划分
thread 怎么绑定
shared memory 怎么分配
fragment/register 怎么分配
global memory 怎么读
shared memory 怎么复用
每个 thread / wave 负责多少元素
reduction 怎么做
最终怎么写回 global memory

TileLang 代码会出现这些概念:

@T.prim_func
def kernel(...):
    T.alloc_shared(...)
    T.alloc_fragment(...)
    T.copy(...)
    T.gemm(...)
    T.reduce_sum(...)
    T.thread_binding(...)

TileLang 的编程重点是显式设计 GPU 执行结构
要关注的几个核心概念是:

tile shape
// tile shape
// 指一个 GPU block / CTA 一次负责计算的数据块形状。
// 例如矩阵乘法里,一个 block 可能负责 C 的 128x128 子矩阵,
// 那么 128x128 就可以理解为 tile shape。
// tile shape 决定了每个 block 的计算量、访存量、shared memory 用量、寄存器压力,
// 是算子性能优化里最核心的参数之一。
block shape
// block shape
// 指 GPU grid 中每个 block 的线程组织形状,或者说一个 block 里面有多少线程、如何排列。
// 例如 blockDim = (128, 1, 1),表示一个 block 有 128 个线程。
// block shape 决定了并行粒度,也影响线程如何映射到 tile 的不同元素。
// 在 TileLang 里,你通常要显式考虑:一个 block 算多大的 tile,
// 这个 tile 内部又由多少线程共同完成。
thread layout
// thread layout
// 指 block 内部的线程如何分工、如何映射到数据元素。
// 例如 128 个线程分别负责加载矩阵 A/B 的哪些元素,
// 哪些线程负责计算哪些输出,哪些线程参与 reduction。
// 好的 thread layout 可以让 global memory 访问连续、shared memory 访问不冲突、计算负载均匀;
// 差的 layout 会导致访存不连续、线程空转、性能很差。
memory hierarchy
// memory hierarchy
// 指 GPU 的多级存储结构,以及数据在这些层级之间如何流动。
// 常见层级包括:
// global memory: 显存,容量大但慢;
// shared memory: block 内共享,速度快但容量小;
// register: 每个线程私有,最快但数量有限;
// cache: 硬件自动管理的缓存。
// GPU 优化的核心就是尽量少访问 global memory,
// 多复用 shared memory 和 register 中的数据。
pipeline
// pipeline
// 指把数据搬运和计算重叠起来执行的机制。
// 例如当前正在计算第 k 个 tile,同时提前加载第 k+1 个 tile。
// 这样可以隐藏 global memory 访问延迟。
// 如果没有 pipeline,GPU 可能会出现“等数据”的空泡;
// 有 pipeline 后,load 和 compute 可以交错进行,提高吞吐。
vectorized load/store
// vectorized load/store
// 指一次指令加载或存储多个连续元素。
// 例如不是每个线程一次 load 一个 float,
// 而是一次 load float4、half4、int4 这类向量化数据。
// 好处是减少指令数量,提高内存带宽利用率。
// 前提通常是数据地址连续、对齐,并且 thread layout 合理。
wave-level reduction
// wave-level reduction
// 指在一个 wave / warp 内部做归约操作。
// 例如求和、最大值、最小值、top-k 的部分规约等。
// 在 NVIDIA 里常说 warp,通常是 32 个线程;
// 在沐曦 MXMACA / C500 场景里要特别注意 wave size 可能是 64。
// wave-level reduction 通常比 shared memory reduction 更快,
// 因为它可以利用线程束内部通信指令,减少同步和 shared memory 访问。
bank conflict
// bank conflict
// 指多个线程同时访问 shared memory 时,访问落到了同一个 memory bank,
// 导致本来可以并行完成的访问被串行化。
// shared memory 被划分成多个 bank,理想情况下不同线程访问不同 bank。
// 如果很多线程访问同一个 bank,就会发生 bank conflict,性能下降。
// 常见解决方法包括 padding、改变数据布局、改变 thread layout。
occupancy
// occupancy
// 指一个 GPU SM / CU 上同时驻留的 warps / waves / blocks 的比例。
// occupancy 高,说明有更多线程可以同时挂在硬件上,
// 当一批线程等待内存时,硬件可以切换到另一批线程继续执行。
// 但是 occupancy 不是越高越好。
// 如果为了提高 occupancy 而牺牲 tile size、数据复用、寄存器利用,反而可能变慢。
// 实际优化时要在 occupancy、寄存器数量、shared memory 用量、tile shape 之间平衡。

TileLang 文档中把 DSL 指令分成数据搬运、计算原语、控制辅助、诊断和高级操作几类;例如 T.copy、T.async_copy、T.gemm、T.reduce_sum、warp reducers 等。

1.5TVM:TileLang 代码会先变成TIR中间表示

IR Intermediate Representation
Python经过TVM IRModule / TIR中间层变成中间表示
TIR 可以理解成 TVM 的底层张量程序表示
例如在 TileLang 写:

T.copy(A[...], A_shared)
T.gemm(A_shared, B_shared, C_frag)
T.copy(C_frag, C[...])

编译器内部会逐步把它 lower 成类似:

for blockIdx.x ...
  for threadIdx.x ...
    allocate shared
    load global
    store shared
    sync
    compute
    store global

这一层会做很多 transformation:

展开 loop
// 指把 for 循环中的多次迭代显式展开成多条语句。
// 好处:
// 1. 减少循环控制开销,比如 i++、条件判断、跳转;
// 2. 给编译器更多优化机会,比如指令重排、寄存器复用、向量化;
// 3. 对固定小循环特别有用,比如 MMA 内部的小 tile 计算。
// 代价:
// 1. 代码体积变大;
// 2. 可能增加寄存器压力;
// 3. 过度 unroll 可能反而降低 occupancy。
// 在 GPU kernel 中,展开 loop 经常用于 k 维循环、向量 load/store、
// 小矩阵片段计算、reduction 内部循环等。
绑定 thread/block
// bind thread/block
// 指把程序里的循环轴映射到 GPU 的执行层级。
// 常见 GPU 层级包括:
// blockIdx.x / blockIdx.y / blockIdx.z
// threadIdx.x / threadIdx.y / threadIdx.z
// 把普通 for 循环变成 GPU 并行执行结构。
// 绑定得好,线程分工清晰、访存连续、计算负载均匀。
// 绑定得差,可能出现线程空转、访存不连续、block workload 不均衡。
插入 boundary check
// boundary check / predicate guard
// 指在访问数组前插入越界判断,防止线程访问非法地址。
// 在 GPU kernel 中,boundary check 很常见,尤其是:
// 1. 最后一块 tile;
// 2. 非整除 shape;
// 3. dynamic shape;
// 4. ragged tensor;
// 5. MoE 中每个 expert token 数不固定。
// 代价:
// branch / predicate 会带来额外判断开销,
// 也可能导致同一个 wave / warp 内部分线程执行、部分线程不执行。
// 但没有 boundary check 就可能非法访存,导致错误结果甚至 kernel crash。
分析 memory scope
// memory scope analysis
// 指编译器分析每个 buffer 应该放在哪一级 GPU memory 中。
// 常见 scope 包括:
// global memory: 显存,全局可见,容量大但慢;
// shared memory: block 内共享,速度快但容量小;
// local memory: 线程私有,通常可能映射到寄存器或溢出到显存;
// register: 线程私有,最快;
// fragment / accumulator: MMA 或 tensor core 计算中的寄存器片段。
// 输入矩阵 A/B 原始数据一般在 global memory;
// 当前 tile 的 A/B 可以搬到 shared memory;
// 每个线程的累加结果 acc 通常放在 register;
// 最后再写回 global memory。
// 放 global 太多会慢;
// 放 shared 太多会降低 occupancy;
// 放 register 太多会造成寄存器压力,甚至 spill 到 local memory。
做 index simplification
// 指编译器把复杂的索引表达式化简成更简单、更容易生成代码的形式。
// 1. 减少运行时整数计算;
// 2. 消除多余的乘法、除法、取模;
// 3. 帮助编译器识别连续访存;
// 4. 帮助后续 vectorize / coalescing / storage flatten;
// 5. 让生成的底层代码更干净。
// 对 GPU 来说,索引计算不是免费的。
// 复杂的地址计算会增加整数 ALU 开销,也可能增加寄存器使用。
// 所以 index simplification 对 kernel 性能很重要。
做 storage flatten
// 指把高维数组访问转换成一维线性地址访问。
// A_flat[i * stride_j + j]
// B_flat[x * stride_x + y * stride_y + z]
// 编译器最终生成 GPU code 时,不依赖多维数组,
// 要具体的 pointer + offset 地址。
// 1. 明确真实内存布局;
// 2. 计算每次 load/store 的线性地址;
// 3. 配合 index simplification 简化地址表达式;
// 4. 为 vectorized load/store 判断连续性;
// 5. 为 shared memory layout、bank conflict 分析做准备。
做 vectorize / unroll
// vectorize:
// 把多个连续标量操作合并成一个向量操作。
// load float x0
// load float x1
// load float x2
// load float x3
// 可以变成:
// load float4
// 好处是减少指令数,提高内存带宽利用率。
// 但是要求地址连续、对齐,并且数据类型和长度适合向量化。
// unroll:
// 把循环体复制多份,减少循环控制开销,并暴露更多优化机会。
// 例如固定长度的小循环经常被 unroll。
// 在 GPU kernel 里,vectorize 常用于:
// 1. global memory 连续读取;
// 2. shared memory 连续读写;
// 3. hidden dimension 上的连续元素;
// 4. half / bf16 / float 的批量 load/store。
// unroll 常用于:
// 1. reduction 内部循环;
// 2. k tile 循环;
// 3. MMA fragment 内部计算;
// 4. 小维度固定循环。
// vectorize 主要优化访存指令数量;
// unroll 主要优化循环控制和指令调度。
// 两者经常一起出现。
插入同步
// insert synchronization
// 指在需要线程协作的位置插入同步指令。
// 在 GPU block 内,不同线程是并行执行的,需要同步。
// 典型例子:
// 1. 所有线程从 global memory 加载 tile 到 shared memory;
// 2. __syncthreads();
// 3. 所有线程从 shared memory 读取数据并计算。
// 如果没有同步,可能出现 read-after-write hazard。
// 也就是读发生在写完成之前。
// 常见同步位置:
// 1. global -> shared 加载之后;
// 2. shared memory 被复用之前;
// 3. pipeline stage 切换时;
// 4. block-level reduction 中;
// 5. 不同 thread 写入 shared 后,其他 thread 要读取时。
// 多余的同步会让整个 block 等最慢的线程,降低性能。
// 正确优化要做到:
// 必要的同步必须有;
// 不必要的同步尽量删掉。
lower 特殊 intrinsic
// lower special intrinsic
// 指把高层 DSL / IR 中的特殊操作,转换成后端 GPU 能识别的底层指令或 runtime 调用。
// intrinsic 可以理解成“特殊内建操作”。
// 例如高层代码里可能写的是:
// mma_sync(...)
// cp_async(...)
// atomic_add(...)
// warp_shuffle(...)
// exp(...)
// rsqrt(...)
// lowering 后可能变成具体后端指令:
// NVIDIA 上可能变成 mma、ldmatrix、shfl、barrier 等;
// AMD / MetaX / 其他后端会变成对应平台自己的 intrinsic 或 runtime API。
// 很多 GPU 高性能能力不是普通 C 代码能自然表达的,所以需要lower special instrinsic。
// 比如:
// 1. tensor core / matrix core 指令;
// 2. wave / warp shuffle;
// 3. async copy;
// 4. special math function;
// 5. atomic operation;
// 6. barrier / memory fence;
// 7. vectorized memory instruction。
// 如果 intrinsic lower 成功,说明高层算子映射到了硬件特殊能力。
// 如果 lower 失败,可能退化成普通标量计算,性能会差很多。

所以优化不能只看 Python DSL 源码,还要看 lower 后的 TIR / generated code。

1.6TileLang MetaX MACA backend

在这一层,TIR经过MACA后端生成.maca文件,之后.maca将会被mxcc编译器分host和device端编译。
TileLang MACA backend:TileLang/TIR变后端代码
TileLang 本身可以支持不同后端,比如 CUDA、ROCm、Metal、Ascend、MACA 等。
MACA backend 的任务是:
把 TileLang/TIR 中的 GPU 抽象映射到 MXMACA 能接受的代码、intrinsic、memory space 和编译选项
也就是说,它要回答这些问题:

threadIdx / blockIdx 在 MACA 里怎么表达?
shared memory 在 MACA 里叫什么?
barrier 怎么 lower?
T.copy 怎么变成 MACA load/store?
T.gemm 能不能映射到矩阵计算指令?
half / bf16 / fp32 支持情况如何?
vectorized load/store 怎么生成?
atomic 支持情况如何?

所以要注意TileLang语言对MetaX的适配程度,
TileLang 官方主线很多示例可能默认更适配 NVIDIA H100 / A100,甚至有些高级路径和 CUDA/Hopper 特性绑定很深。tilelang-metax 的 issue 页面里也能看到一些 MACA 后端相关问题,例如 FP16 accumulation MMA 编译问题、稀疏 tensorcore 示例缺少可执行 MACA 路径、低精度示例假设 CUDA/Hopper-specific emitter 或 launch 形状等。所以在看 TileLang 示例时要有一个判断:

这是通用 TileLang 写法?
还是 CUDA/Hopper 特化写法?
MACA backend 是否真的支持?
C500 上是否高效?

尤其要检查这些危险点:

warp size 是否写死 32
是否假设 NVIDIA warp-level primitive
是否假设 cp.async / TMA
是否假设 WGMMA
是否使用 CUDA-specific intrinsic
是否使用 unsupported dtype
是否使用 MACA backend 未完善的 lowering path

1.7mxcc:device code 编译器驱动

TileLang MACA backend 生成的.maca文件还需要交给沐曦的编译工具链。
这里的 mxcc / MXMACA LLVM toolchain 可以理解成:MetaX 平台的 GPU 编译器
它的角色类似于 NVIDIA 生态里的nvcc / ptxas / NVVM,负责把更底层的 kernel 源码或 IR 编译成 C500或者其他硬件能执行的 device code。
mxcc 是面向沐曦 GPU 的编译器驱动。官方文档里说:它可以把 MXMACA 源文件编译成同时包含主机部分和设备部分的可执行文件;主机部分在 CPU 上执行,设备部分在 GPU 上执行。最基本的命令是:mxcc a.maca -o a.out。这个命令的意义是使用沐曦公司的编译器驱动 mxcc,将名为 a.maca 的源文件,编译链接成一个名为 a.out 的可执行文件。其中host 部分按普通 C/C++ 编译,device 部分按 GPU 目标编译。
在mxcc编译时,可能做这些事:

前端解析
LLVM IR 生成
目标平台优化
寄存器分配
指令选择
生成设备二进制
链接 runtime 需要的符号

官方文档给出的典型流程是:

1. 把数据从 CPU 内存拷贝到 GPU 内存
2. 调用核函数对 GPU 内存的数据进行处理
3. 将数据从 GPU 内存传送回 CPU 内存

同时文档说明,串行代码在 host CPU 上执行,并行代码在 GPU 上执行;host 代码按标准 C/C++ 编写,device 代码使用 MXMACA C/C++ 编写,mxcc 为 host 和 device 生成可执行代码。
mxcc 的作用类似于 NVIDIA 生态里的 nvcc:它负责处理 MXMACA 源码或后端生成代码,把 host CPU 部分和 device GPU 部分分别编译,最终生成可以在 MetaX 平台上运行的程序或设备侧二进制。

mxcc 编译选项

mxcc 编译选项 是一个总称,指的是传给 mxcc 的所有命令行参数。
这些选项不只控制 GPU 代码生成,还会控制输入文件、输出文件、头文件路径、库路径、宏定义、编译到哪一步、是否生成调试信息、是否链接设备库、是否运行程序等。

可以把 mxcc 选项分成几类:

1. 文件、路径、宏、库选项
2. 编译阶段控制选项
3. 编译器 / 链接器行为选项
4. 给特定阶段透传参数的选项
5. 编译器驱动控制选项
6. GPU codegen 选项
7. 通用工具选项
8. 设备链接器选项

其中对算子优化最重要的是:

编译阶段控制
调试信息 / line info / opt-info
GPU codegen 选项
设备链接 / rdc
寄存器数量控制
fast math / ftz / div / sqrt / fmad
设备库链接

也就是说,mxcc 编译选项 是一个大集合,而 GPU codegen 选项 是其中专门影响 GPU 设备代码生成的一小类。

文件、路径、宏、库选项:控制编译环境

这类选项主要决定编译时使用哪些文件、头文件、库、宏和设备库。
它们通常不是直接优化 kernel 性能的旋钮,但会影响编译是否正确、设备库是否匹配、数学函数和 intrinsic 是否可用。

常见选项包括:

-o / --output-file
-include / --pre-include
-l / --library
-D / --define-macro
-U / --undefine-macro
-I / --include-path
-isystem / --system-include
-L / --library-path
-maca-path
-maca-host-lib-path
-maca-host-lib
-maca-device-lib-path
-maca-device-lib
-maca-host-input
-maca-device-input
-input-is-host
-input-is-device

这些选项在算子优化中主要有三个作用。

第一,确认使用的是正确的 SDK 和设备库。
如果机器上同时存在多个 MXMACA 版本,路径选错可能导致编译能过,但性能异常,或者某些 dtype、intrinsic、math 函数行为和预期不一致。

第二,通过宏定义切换实验版本。
例如:

mxcc kernel.maca -D TILE_M=128 -D TILE_N=128 -D USE_PIPELINE=1

这样同一个 kernel 可以通过宏切换不同 tile size、pipeline stage 或 fast path。

第三,处理设备侧自定义库。
如果后续把 device helper、数学函数、intrinsic wrapper 拆成单独设备库,就需要理解 -maca-device-lib-path-maca-device-lib

编译阶段控制:控制生成到哪一步

编译阶段控制选项决定 mxcc 是直接生成最终可执行文件,还是只生成某种中间产物。

常见选项包括:

-c
-fatbin
-device-obj
-device-bin
-E
-M
-MM
-MD
-MMD
-run
-lib
-dlink-obj
-dlink-asm

作用如下:

-E:
    只做预处理。

-c:
    编译成目标文件,不直接链接。

-fatbin:
    生成仅设备侧 fat binary,丢弃 host 代码。

-device-obj:
    生成仅设备侧目标文件。

-device-bin:
    生成仅设备侧二进制。

-run:
    编译、链接并运行。

对 TileLang 调优来说,这类选项的意义是:
当 kernel 出错或者性能异常时,可以把编译流程停在某个阶段,检查中间产物,判断问题到底来自 TileLang DSL、TIR lowering、MACA backend codegen,还是 mxcc 后端编译。

-lineinfo:让 profiler 能定位到源码行

-lineinfo 的作用是为设备代码生成行号信息。

mxcc kernel.maca -lineinfo -o a.out

这对 profiler 很重要。
如果没有 line info,profiler 可能只能告诉你某个 kernel 很慢,或者某类指令耗时高。
有了 line info,profiler 才更容易把热点映射回源码行,例如:

热点在某一行 global load
热点在某一行 mma loop
热点在某一行 store

对于 TileLang kernel,理想链路是:

TileLang DSL 行
    ↓
生成的 MACA kernel 行
    ↓
设备 line info
    ↓
profiler 热点

如果拿不到 TileLang 源码行映射,至少也要能映射到后端生成的 .maca 代码行。这样才能把 profiler 结果反向定位到具体的 load、store、MMA、同步或 epilogue 逻辑。

-opt-info:看编译器到底做了哪些优化

-opt-info 用来输出优化报告。

mxcc kernel.maca -opt-info=<kind> -o a.out

用于确认“编译器实际做了什么优化”。

优化报告可以帮助回答:

哪些循环被优化了?
哪些函数成功 inline?
哪些优化没有发生?
有没有 register / resource 相关提示?
有没有 vectorize、unroll、fusion、消除等信息?

具体 <kind> 支持哪些值,需要以本机 mxcc --help 或当前版本文档为准。

对算子优化来说,-opt-info 的价值在于:
如果 TileLang 里写了 pipeline、vectorized load、unroll 或某种内存布局优化,但性能没有提升,就可以通过优化报告和 lowered code 判断这些优化是否真的生效。

-O:主机优化级别

-O 用来指定优化级别,例如:

mxcc kernel.maca -O3 -o a.out

需要注意的是,这里更偏向 host 代码优化级别。
对于 GPU kernel 性能来说,-O3 不是唯一核心。真正直接影响 device kernel 的,通常是后面的 GPU codegen 选项,例如寄存器限制、fast math、ftz、prec-div、prec-sqrt、fmad 等。
所以确定哪种优化更快的合理的做法是:

host 编译优化级别固定
device codegen 选项单独 sweep
最终用 benchmark + profiler 判断

-default-stream:默认 stream 行为会影响 benchmark

-default-stream 用来指定 MXMACA 命令默认发送到哪个 stream,例如:

mxcc kernel.maca -default-stream legacy
mxcc kernel.maca -default-stream null
mxcc kernel.maca -default-stream per-thread

这个选项不直接改变 kernel body,但会影响 runtime 调度和 benchmark 结果。

影响:

kernel 是否被串行化
host-device copy 是否能和 kernel overlap
多 stream benchmark 是否真的并发
profiler timeline 是否符合预期

有时 kernel 本体不慢,但由于默认 stream 语义导致前后操作被同步或串行化,最终端到端 latency 仍然不好。因此在正式性能报告中,stream 行为也应该记录清楚。

GPU codegen 选项

GPU codegen 的意思是 GPU code generation,也就是“GPU 设备代码生成”。
在 mxcc 选项体系里,GPU codegen 选项专门控制 device kernel 代码如何生成。

它们影响:

寄存器使用
occupancy
spilling
数学函数近似
FMA 融合
denormal 处理
设备代码链接方式
矩阵计算 / fragment 相关资源使用

常见 GPU codegen 选项包括:

-fgpu-rdc
-maxrregcount
-maxmregcount
-maxsregcount
-use-fast-math
-ftz
-prec-div
-prec-sqrt
-fmad

作用如下:

-fgpu-rdc:
    启用 GPU 可重定位设备代码。
    适合 device code 跨编译单元链接,例如多个 device helper 或设备库协作。
    但它可能影响编译、链接和优化方式,最终是否打开要根据工程结构和性能测试决定。

-maxrregcount:
    限制每个线程使用的通用寄存器数量。
    寄存器少一些可能提升 occupancy,但限制过紧会导致 spill,把寄存器变量溢出到更慢的本地/全局内存。

-maxmregcount:
    限制每线程 mreg 数量。
    可能影响矩阵计算、MMA、fragment 相关资源使用。

-maxsregcount:
    限制每线程 sreg 数量。
    可能影响标量控制、地址计算、索引逻辑等资源使用。

-use-fast-math:
    启用快速数学近似优化。
    可能提升 exp、div、sqrt 等数学路径性能,但可能降低严格数值精度。

-ftz:
    flush-to-zero。
    将 denormal / subnormal 极小浮点数直接当作 0 处理。
    可能提升性能,但会改变极小数行为。

-prec-div:
    使用精确除法。
    如果关闭精确除法,可能使用近似除法以提升性能,但会影响数值精度。

-prec-sqrt:
    使用精确平方根。
    如果关闭精确 sqrt,可能使用近似 sqrt 或 rsqrt,提高速度但影响精度。

-fmad:
    允许把 a*b+c 融合成 FMA/FMAD。
    通常有利于性能,但可能改变浮点舍入结果。

gpu codegen选项的调优应用

如果 profiler 显示 occupancy 低,且寄存器使用过高:
    可以尝试减小 tile / fragment,或者 sweep -maxrregcount。

如果限制寄存器后性能更差:
    可能发生 spill,说明寄存器限制过紧。

如果 attention / softmax 中数学函数较多:
    可以测试 -use-fast-math、-ftz、-prec-div、-prec-sqrt,但必须重新做 correctness。

如果普通乘加密集:
    可以确认 -fmad 是否启用,以及是否改变误差边界。

如果 device helper 被拆到多个编译单元:
    需要关注 -fgpu-rdc 和设备链接行为。

mxcc编译选项总结图

mxcc 编译选项
├── 文件 / 路径 / 宏 / 库
├── 编译阶段控制
├── 调试信息
├── 优化报告
├── stream / runtime 行为
├── 设备链接
└── GPU codegen 选项
    ├── 寄存器限制
    ├── fast math
    ├── ftz
    ├── div / sqrt 精度
    ├── fmad
    └── rdc

与 TileLang 调优的关系

一个可能的调优链如下

1. 修改 TileLang DSL
   设计 tile size、block/thread 映射、fragment、shared memory、pipeline、layout。

2. 查看 lowered TIR
   确认 shared memory、thread binding、MMA intrinsic、pipeline 是否真的出现。

3. 生成 MACA 后端代码
   如果能保存临时 .maca 文件,就检查 kernel 结构是否符合预期。

4. 使用 mxcc 编译
   开启 -lineinfo,必要时开启 -opt-info,记录编译选项和设备库版本。

5. 使用 profiler 分析
   查看 kernel latency、occupancy、register、global memory、shared memory、compute utilization、scheduling stall。

6. 做 GPU codegen sweep
   测试 -maxrregcount、-maxsregcount、-use-fast-math、-ftz、-prec-div、-prec-sqrt、-fmad、-fgpu-rdc 等选项。

7. 选择 correctness 和 performance 都稳定的配置
   最终把编译选项写入 benchmark 记录和性能报告。

最重要的因果链是:

TileLang tile / thread / fragment / pipeline 设计
    ↓
每个 thread 使用多少寄存器
    ↓
每个 block 能驻留多少 wave / block
    ↓
occupancy
    ↓
是否能隐藏访存和执行延迟
    ↓
最终性能

推荐的编译配置分层

实际工程中,可以把 mxcc 配置分成三类。

debug 配置

用于查错

-g
-device-debug
-lineinfo

目标是方便定位非法访存、错误行号、kernel crash 和设备侧问题。

profile 配置

用于性能定位:

-lineinfo
-opt-info

目标是让 profiler 能把热点映射到源码或生成代码,同时观察编译器优化情况。

release 配置

用于最终性能测试:

不开 device-debug
固定最终 fast math / reg 限制 / fmad 等选项
记录完整编译参数

目标是测真实性能,并保证结果可复现。

最终性能报告中建议记录:

mxcc version
MXMACA SDK version
driver version
TileLang commit
是否开启 lineinfo
是否开启 device-debug
是否开启 fast math
是否设置 maxrregcount / maxsregcount
是否开启 fmad
是否开启 ftz
是否使用精确 div / sqrt
是否使用 rdc
default-stream 设置

这样在看 benchmark 结果时,才能知道性能差异到底来自 TileLang kernel 设计,还是来自编译选项、设备库版本、stream 行为或 profiler 配置。

在这一层会遇到的问题通常是编译器/后端问题:

某个 intrinsic 不支持
某个 dtype 不支持
寄存器爆了
shared memory 超限制
编译器优化失败
生成非法代码
编译时间过长
kernel launch 参数不合法

调试时要区分:

TileLang DSL 写法错误
TIR lowering 错误
MACA backend 代码生成错误
mxcc 编译错误
runtime launch 错误
kernel 结果错误
性能不达标

这几个问题的定位方法完全不同。

1.8JIT即时编译,动态加载

JIT 是 Just-In-Time compilation,即运行时编译调度机制。意思是:代码不是完全提前编译好,而是在程序运行过程中,根据当前的参数、shape、target 后端,临时生成或获取可执行设备代码。
JIT跨越了TileLang DSL,TVM lower,MACA backend,mxcc这些阶段。当 @tilelang.jit kernel 被调用时,如果发现某个(kernel+shape+dtype+tile参数+target backend+编译选项)组合还没有编译过,或者未命中cache,就触发JIT编译,调用TileLang DSL,TVM lower,MACA backend,mxcc去实例化这个kernel组合,生成GPU可使用的device code,然后被runtime API动态加载成compiled module。
在 MoE 仓库中,JIT 入口是 custom_fusedmoe.py 里的 @tilelang.jit。它修饰的是 moe_forward_tilelang_routed(),这个函数根据 hidden size、expert size、expert 数量、routed token 总数和 block 参数生成具体的 TileLang kernel。

RoutedMoEKernel.__init__() 中调用这个 JIT 函数,得到 self.impl;
之后 RoutedMoEKernel.__call__() 调用 self.impl(...),
才是真正把 packed token、expert weights、group metadata 和输出 buffer 传给设备 kernel 执行。

动态加载在仓库里没有显式写 mcModuleLoad,因为 TileLang 后端封装了这层。但从 MXMACA 官方运行时模型看,运行时编译和动态加载会经历生成 bitcode/binary、module load、获取 kernel function、module launch kernel 的过程。TileLang 生成的后端设备代码最终也需要经过类似的 runtime 加载和 launch 才能在 C500 上执行。
当前 benchmark 会在 main 开头删除 ~/.tilelang/cache,所以第一次运行可能触发 TileLang JIT 编译;但性能分支有 warmup 10 次,正式 event 计时在 warmup 之后开始,所以首次冷启动 JIT 不进入最终平均 latency。需要注意的是,正式 repeat 中每次仍然重新构造 RoutedMoEKernel 和 MoE,所以 cache lookup、Python wrapper 构造、buffer 分配等仍然会计入当前脚本定义的端到端耗时。
从 MXMACA runtime 角度看,JIT 和动态加载可以理解为:运行时生成设备代码,加载 module,获取 kernel function,然后通过 runtime launch API 送入 GPU 执行。官方文档中 MCRTC 会通过 mcrtcGetBitCode 生成 bitcode,再通过 mcModuleLoad 加载,并在 mcModuleLaunchKernel 时执行设备代码。TileLang 把这部分封装起来了,所以仓库里看不到显式的 mcModuleLoad,但底层必须完成类似的加载和 launch 流程。

1.9Runtime API:运行时资源管理、load kernel module 和 launch kernel

MetaX通用运行时API
在JIT组织mxcc编译好了 device code 之后,还需要 runtime 来启动它。
runtime API ≈ load + launch + memory + stream + sync
runtime 层主要体现为 PyTorch CUDA/Tensor API 和 TileLang JIT kernel 调用。虽然代码里没有直接出现 mcMalloc、mcMemcpy、mcLaunchKernel 这类底层接口,但 torch.empty、torch.zeros、tensor.clone、torch.tensor(..., device="cuda")、torch.cuda.Event、torch.cuda.synchronize、self.routed_kernel(...) 等调用背后都对应设备内存分配、数据搬运、event 计时、stream 调度、kernel launch 和同步。
runtime 层主要负责这些事情:

1. 选择设备
2. 初始化设备上下文
3. 管理设备内存
4. 管理 host 和 device 之间的数据搬运
5. 创建和管理 stream
6. 创建和记录 event
7. 启动 kernel
8. 等待 kernel 完成
9. 管理 JIT 编译后的 module / kernel function
10. 返回错误状态或触发异常

runtime 层会影响:

kernel launch overhead
stream 同步行为
event 计时准确性
显存分配开销
错误码定位

与优化算子的关系:
如果一个算子很小,比如 MoE 某些 routing 阶段,kernel launch overhead 可能占比很高。此时优化方向可能是:

融合多个 kernel
减少 PyTorch 侧操作
减少 launch 次数
把 preprocessing 合进 TileLang kernel

那么Runtime API如何进行内存管理?
runtime API负责的事项之一是设备端内存分配(device memory allocation),指的是在 GPU设备内存上分配空间。
device memory allocation 可以分成两种:global memory 和 kernel 内部临时存储
在底层 MXMACA runtime 里,device global memory分配的典型形式是:

mcSetDevice(device_id);
mcMalloc(&ptr, sizeBytes);
...
mcFree(ptr);

MXMACA 文档说明,设备内存分配前可以通过 mcSetDevice(deviceId) 选择设备,如果不显式指定则默认设备为 0;设备内存通过 mcMalloc(void **ptr, size_t sizeBytes) 分配,释放时使用 mcFree。文档还说明 mcMalloc 会在设备上分配指定字节数的线性内存,并返回设备内存指针。
MoE,MLA,NSA算子仓库
在 MoE 仓库写:

torch.empty(..., device="cuda")
torch.zeros(..., device="cuda")
torch.randn(..., device="cuda")
tensor.clone()
torch.stack(...)
torch.tensor(..., device=self.device)

这些 PyTorch/TileLang 操作背后会触发设备内存分配、设备内 copy、host/device copy 或 allocator 缓存复用。
host 端进行操作,申请 device 端 global memory 主要有以下几类

第一类分配:generate_input() 里的输入和权重
第二类分配:clone_data(data) 每次 repeat 都会分配
第三类分配:custom_kernel() 每次构造 RoutedMoEKernel 和 MoE
第四类分配:权重转置、contiguous、stack
第五类分配:forward 中的 metadata tensor
第六类分配:scatter_reduce 前的大 index tensor

上面讲的是 host 侧创建的 GPU tensor,也就是全局设备内存。
custom_fusedmoe.py 里还有另一类“内存分配”:kernel 内部的片上存储/寄存器/共享内存资源声明

类型 代码形式 生命周期 是否由 host 分配大 tensor
global device memory torch.empty, torch.zeros, torch.stack, tensor.clone API 调用期间或 tensor 生命周期内
kernel shared memory T.alloc_shared 一个 kernel block 执行期间 否,属于 kernel 资源
fragment/register/local T.alloc_fragment 一个 kernel / block / 线程计算期间 否,通常由编译器映射到寄存器或本地存储
intermediate global buffer up_logits_routed Step 1 和 Step 2 之间

Runtime API如何进行kernel launch?
kernel launch指的是 host 侧把一个已经编译好的设备函数提交给 GPU 执行。
在底层 MXMACA 里,kernel 启动可以用类似:kernel_name<<<grid, block>>>(args)表示
MXMACA 文档说明,<<<grid, block>>> 中第一个值是线程网格维度,即启动多少线程块;第二个值是线程块维度,即每个线程块包含多少线程。经典三尖号启动语法在编译时会被替换为 runtime 库提供的 mcLaunchKernel API。
MXMACA Runtime API 参考中,mcLaunchKernel 的参数包括:

mcLaunchKernel(
    const void *function_address,
    dim3 numBlocks,
    dim3 dimBlocks,
    void **args,
    size_t sharedMemBytes,
    mcStream_t stream
)

它负责启动设备函数;其中 numBlocks 是 block 数量,dimBlocks 是 block 尺寸,args 是 kernel 参数,sharedMemBytes 是动态共享内存大小,stream 是调度该 kernel 的流。
在仓库的MoE算子中,custom_fusedmoe.py 中的 TileLang kernel 有两个主要阶段。
Step 1:gate/up projection
这是 MoE routed experts 的第一段 MLP 计算,也就是对每个被路由到 expert 的 token 做:

gate projection + up projection + SiLU 激活 + 逐元素相乘

数学公式是

gate = x @ W_gate^T
up   = x @ W_up^T
hidden = SiLU(gate) * up

计算得到的 hidden 会保存为up_logits_routed
代码是:

with T.Kernel(
    M,
    T.ceildiv(dexpert, block_dexpert),
    threads=threads
) as (bx, by):
    ...

默认参数是:

block_token   = 128
block_dhidden = 128
block_dexpert = 128
threads       = 256
num_stages    = 1

所以 Step 1 的逻辑 launch 配置可以理解成:

grid_x = M
grid_y = ceildiv(d_expert, block_dexpert)
block_threads = 256

其中

M = ceil(group_sum / block_token) + group_count
group_sum = B * S * top_k,也就是 routed token 总数。

Step 1 每个 block 负责:
一块 routed token 行:block_token = 128
一块 expert intermediate 维度:block_dexpert = 128
它做:

input @ routed_expert_gate^T
input @ routed_expert_up^T
SiLU(gate)
up * SiLU(gate)

结果写入 up_logits
代码里就是先 T.copy inputgate/up weight,再 T.gemm,然后做 SiLUelementwise multiply,最后写 up_logits
Step 2:down projection
第二段是:

with T.Kernel(
    M,
    T.ceildiv(dhidden, block_dhidden),
    threads=threads
) as (bx, by):
    ...

所以 Step 2 的逻辑 launch 配置是:

grid_x = M
grid_y = ceildiv(d_hidden, block_dhidden)
block_threads = 256

它每个 block 负责:
一块 routed token 行:block_token = 128
一块 hidden 输出维度:block_dhidden = 128
它做:

up_logits @ routed_expert_down^T
乘 routed_expert_weights
写入 expert_output_routed

代码里 Step 2 从 up_logitsrouted_expert_down 读数据,用 T.gemm 累加到 output_local,最后写:

output[...] = output_local[...] * routed_expert_weights[...]
一次 kernel launch 需要传哪些参数
RoutedMoEKernel.__call__() 的参数是:
input
routed_expert_gate
routed_expert_up
routed_expert_down
routed_expert_weights
group_sizes
group_offsets
group_padded_offsets
group_idx_for_bx
up_logits
output

对应到底层 runtime,launch 时传入的是这些 tensor 的底层设备指针和 shape/stride 等必要元信息。核心可以理解为:

input:
    stacked_expert_tokens 的 device pointer

routed_expert_gate/up/down:
    stacked expert weights 的 device pointer

routed_expert_weights:
    top-k router score 的 device pointer

group_sizes / group_offsets / group_padded_offsets / group_idx_for_bx:
    routing metadata 的 device pointer

up_logits:
    中间输出 buffer 的 device pointer

output:
    expert_output_routed 的 device pointer

device memory allocation 是 launch 之前的准备;
kernel launch 是把已分配的 device memory 交给设备 kernel 执行;
kernel 内部再用 shared/fragment/register 做片上临时计算。

1.10MetaX Driver:驱动层

Runtime 再往下就是 Driver。Driver 负责更底层的事情:

设备上下文管理
命令提交
显存页管理
kernel 调度
硬件队列
错误恢复
和操作系统交互

一般不会直接操作 driver,但 profiler、mx-smi、runtime 都会依赖 driver。
driver/runtime 层可能出现的问题:

device not found
permission denied
kernel launch failed
illegal memory access
device-side assert
显存占用异常
GPU utilization 异常

在比赛优化里更多通过上层工具间接观察 driver/hardware 状态。

1.11GPU:最终执行层

这一层是硬件本身。
所有上层设计最后都要落到 C500 的硬件约束上:

wave size = 64
每个 CU / SM-like 单元的资源限制
寄存器数量
shared/local memory 容量
L2 / cache 行为
global memory bandwidth
矩阵计算单元能力
支持的 dtype
访存合并规则
bank conflict 规则
occupancy 限制

它会影响:

block size 最好是 64 的倍数
lane id 应按 0~63 设计
reduction 要覆盖 64 lanes
一个 wave 的连续访存跨度要重新计算
分支发散的单位是 64 lanes
MoE 小 expert 可能浪费大量 lane
MLA head_dim=64/128 可能天然适配
NSA 稀疏 pattern 要避免 wave 内严重发散

在 GPU 上优化,不能只问:
这个算法复杂度是多少?
还要问:

每个 wave 干什么?
每个 block 有几个 wave?
每个 wave 是否满载?
每个 wave 是否连续访存?
每个 wave 是否分支一致?
寄存器够不够?
shared memory 是否冲突?
occupancy 是否够?

1.12 mcProfiler / mx-smi / benchmark:结果分析工具:从端到端耗时反向定位瓶颈

benchmark可以看主要结果:

运行时长多少
相对 reference 快多少
不同 shape 下表现如何

这里提到了reference
常用的baseline包括
沐曦软件栈提供的 mcDNN、mcBLAS、mcSolverIT、mcRAND、mcFFT、mcThrust、mcCUB 等库。
实践中建议:

标准 GEMM:先对比 mcBLAS
标准 Conv / Norm / Activation:先对比 mcDNN / 框架 kernel
通信:先对比 MCCL
自定义 fused kernel:TileLang / MXMACA native
需要极致硬件特化:MXMACA native + profiler + 必要时补 TileLang 后端

1.12.1用 mcBLAS / mcDNN 建立 MoE 性能参照系

MoE 的核心计算本质上仍然是 MLP,只不过普通 dense MLP 变成了按 expert 分组后的 grouped MLP。
对于每个被路由到 expert 的 token,expert 内部主要包含三次矩阵乘:

gate projection:
    X @ W_gate

up projection:
    X @ W_up

down projection:
    hidden @ W_down

其中 gate 和 up projection 先分别计算:

G = X @ W_gate
U = X @ W_up

然后经过激活和逐元素乘法:

Z = SiLU(G) * U

最后再做 down projection:

Y = Z @ W_down

所以从计算主体看,MoE expert MLP 的主要 FLOPs 来自三类 GEMM:

gate GEMM
up GEMM
down GEMM

这也是为什么在分析 MoE kernel 性能时,不能只看端到端 latency,还需要给这三类 GEMM 找一个合理的性能参照系。

mcBLAS 作为参照的作用

mcBLAS 是 MXMACA 软件栈中的基础线性代数库,可以理解为沐曦平台上的标准 GEMM 实现。
它不一定是我们最终提交的实现方式,但它非常适合作为性能边界和 sanity check。

对 MoE 来说,mcBLAS 可以提供三类参考价值。

第一,作为纯 GEMM 层面的性能上界参考。
如果一个相同规模的标准 GEMM 用 mcBLAS 能跑得很快,而 TileLang 版本明显慢很多,那么说明 TileLang kernel 的 tile shape、thread layout、MMA lowering、shared memory 复用或寄存器使用可能存在问题。

第二,作为 grouped / batched expert GEMM 的 baseline。
MoE 的 expert 计算可以粗略拆成:

for each expert:
    mcBLAS GEMM for gate
    mcBLAS GEMM for up
    activation + multiply
    mcBLAS GEMM for down

这种做法不一定最优,因为它可能产生很多小 GEMM 和多次 kernel launch,但它可以帮助我们理解:如果不写 fused TileLang kernel,而是用成熟 GEMM 库逐 expert 执行,性能大概处在哪个水平。

第三,用来判断 TileLang T.gemm 的 tile 参数是否合理。
如果 TileLang 的 grouped GEMM 远低于 mcBLAS 在类似矩阵规模下的效率,就要回头检查:

block_token 是否合适
block_dhidden / block_dexpert 是否合适
threads 是否适合 C500 wave size
shared memory 是否复用充分
register 是否过多导致 occupancy 下降
MMA 是否真的成功 lower
global load/store 是否连续
mcDNN 作为参照的作用

mcDNN 适合作为标准 DNN primitive 的参考。
对于 MoE,它可以提供两类间接参照。

第一,参考标准 MLP / activation / DNN primitive 的性能水平。
MoE expert 本质上是 gated MLP,所以 Linear + SiLU + Linear 这类路径可以帮助我们判断标准框架实现大概能达到什么性能。

第二,和 mcPyTorch / PyTorch backend 下的 nn.Linear + SiLU 组合做间接对比。
当前 MoE 仓库中,router、topk、scatter 等部分仍然使用 PyTorch op,而 routed expert MLP 使用 TileLang 实现。因此 mcDNN / mcPyTorch 的标准执行路径可以作为“框架默认实现”的参考,而 TileLang kernel 则是更定制化、更融合的实现。

mcBLAS GEMM 时间并非 MoE 端到端 latency

最容易误解的一点是:
mcBLAS 的单个 GEMM 时间,不能直接和当前 MoE benchmark 的端到端 latency 比。

原因是当前 MoE benchmark 的计时边界不是单个 GEMM,而是:

custom_kernel(clone_data(data))

这个调用内部包含很多内容:

clone_data
RoutedMoEKernel / MoE 对象构造
权重 stack / contiguous
router GEMM
softmax
topk
argsort
bincount
token packing
group metadata 构造
TileLang routed grouped GEMM
scatter_reduce
cuda synchronize

也就是说,MoE 的端到端 latency 不只是 expert MLP 的 GEMM 时间,还包含 routing、数据重排、metadata 构造、scatter reduce、同步、临时 buffer 分配等开销。

因此更合理的对比方式应该分成三层:

第一层:pure GEMM baseline
    用 mcBLAS 测 gate/up/down 单个 GEMM 或 batched GEMM 的性能边界。

第二层:expert MLP baseline
    用 mcBLAS 或 mcDNN 组合出 gate + up + SiLU + multiply + down,
    作为 expert MLP 的参考实现。

第三层:MoE end-to-end baseline
    保持官方 benchmark 边界,统计 router、topk、packing、TileLang kernel、
    scatter_reduce 和同步之后的整体耗时。

只有做到第三层之后才能和比赛端到端 latency 直接比较。
第一层和第二层主要用于分析 TileLang kernel 是否接近硬件和库的计算性能边界。

优化闭环

mcBLAS / mcDNN 的作用总结:

mcBLAS / mcDNN 是性能参照系。
它们帮助我们判断 TileLang grouped GEMM 是否高效,以及当前瓶颈到底在 kernel 内部,
还是在 router、packing、scatter_reduce、同步和 host 侧开销上。

如果 routed_kernel_only 很慢,而且明显低于 mcBLAS 类似 GEMM 的效率,说明要重点优化 TileLang kernel 内部。
如果 routed_kernel_only 已经接近合理水平,但 official_e2e 仍然很慢,说明主要瓶颈可能在 PyTorch 前后处理、token packing、metadata 构造、scatter_reduce 或同步上。

优化工具

mx-smi 是系统级观察工具,可以看:

GPU 利用率
显存占用
功耗
温度
进程
设备状态

Profiler 是性能诊断工具,作用是回答“慢在哪里、为什么慢、下一步应该改什么”。
在 MetaX-TileLang 工具链中,benchmark 只能告诉我们端到端 latency、平均耗时和是否通过测试;而 profiler 需要进一步把耗时拆开,定位瓶颈属于 host 侧调度、数据搬运、kernel launch、global memory 访问、计算利用率、occupancy、同步等待,还是负载不均衡。

因此 profiler 阶段应该采用两层分析方法:

第一层是 Python / host 侧分段计时,用来定位 API 端到端耗时由哪些阶段构成。
第二层是 GPU kernel 级 mcProfiler,用来分析具体 device kernel 的访存、计算、调度和资源使用情况。

1.12.2 Python / host 侧分段计时 profiling

对于 MoE 算子,首先应该做 Python 侧分段计时,因为 MoE 的端到端路径中包含大量不属于 TileLang kernel 本体的开销。
官方 MoE 流程中,router、softmax、topk、argsort、bincount、expert grouping、token packing、group metadata 构造、TileLang routed kernel、scatter_reduce 和 synchronize 分布在不同阶段。如果只看总耗时,很难判断到底是 TileLang grouped GEMM 慢,还是前后的 routing / packing / scatter 慢。

因此 MoE 至少应该拆成以下几个 profiling scope:

official_e2e:
    保持官方 benchmark 逻辑。
    计时范围包含 clone_data、对象构造、router、topk、argsort、bincount、
    token packing、metadata 构造、TileLang routed kernel、scatter_reduce 和同步。
    这个结果最接近比赛 API 端到端耗时。

cached_e2e:
    复用 RoutedMoEKernel / MoE 对象,尽量避免每轮重复构造 kernel 和权重 buffer。
    用来观察对象构造、JIT cache lookup、权重 stack、buffer 分配等 host 侧开销。

routed_kernel_only:
    预先准备好 stacked_expert_tokens、group_sizes、group_offsets、
    group_padded_offsets、group_idx_for_bx、routed_expert_weights 等输入。
    只计时 TileLang routed expert MLP kernel。
    用来观察自定义 TileLang kernel 本体的性能。

通过这三种 scope,可以把 MoE 的总耗时拆成:

official_e2e latency
    = clone / object construction / routing / packing / metadata
    + TileLang routed kernel
    + scatter_reduce
    + synchronize

cached_e2e latency
    = routing / packing / metadata
    + TileLang routed kernel
    + scatter_reduce
    + synchronize

routed_kernel_only latency
    = TileLang grouped expert MLP kernel

如果 official_e2e 明显慢,但 routed_kernel_only 不慢,说明主要瓶颈不在 TileLang kernel 内部,而在 Python / PyTorch / runtime 调度、token packing、scatter_reduce 或同步上。
如果 routed_kernel_only 本身就慢,则需要进入 mcProfiler,继续分析 kernel 内部的访存、计算和资源瓶颈。

1.12.3 GPU kernel 级 mcProfiler

mcProfiler 是 MXMACA 软件栈中的 GPU 性能分析工具,适合分析具体 kernel 在 C500 上的执行情况。
它可以从 SOL / RoofLine 全局视角,以及 Memory、Computing、Scheduling 等维度观察 kernel 是否充分利用硬件资源。

mcProfiler 主要用来回答以下问题:

1. kernel 是 memory-bound 还是 compute-bound?
2. global memory / HBM 读写是否过高?
3. load/store 是否连续,是否存在严重非合并访存?
4. register 使用是否过高,是否发生 spill?
5. shared memory / WSM 使用是否过高?
6. occupancy 是否被 register、shared memory 或 block size 限制?
7. Scheduling stall 主要来自哪里?
8. kernel launch 数量是否过多,是否存在大量小 kernel?
9. 不同 shape 下瓶颈是否变化?
10. 当前 TileLang kernel 是否真的走到了预期的 MMA / pipeline / vectorized load 路径?

在使用 mcProfiler 时,应该优先记录以下信息:

case_id
op_name
shape
dtype
layout
kernel name
latency
kernel launch count
global memory read/write
shared memory usage
register usage
occupancy
compute utilization
memory utilization
scheduling stall
是否发生 spill
是否存在异常同步
profiler 结论
对应优化动作

1.12.4 MoE 的 profiling 重点

MoE 的 profiling 需要同时关注 host 侧和 device 侧。

host 侧重点是:

router / softmax / topk 是否占用明显时间
argsort 和 bincount 是否产生额外 kernel 或同步
bincount().cpu().numpy() 是否造成 D2H copy 和强同步
token packing 是否由 Python loop 产生大量小操作
group metadata 是否每轮重新创建
scatter_reduce 是否占用大量时间
是否每轮都 clone_data
是否每轮都重新构造 RoutedMoEKernel / MoE

device 侧重点是:

TileLang routed kernel 的 Step 1 和 Step 2 是否耗时均衡
gate/up GEMM 是否充分利用矩阵计算单元
down GEMM 是否成为主要瓶颈
up_logits global intermediate 是否产生大量 global memory 写读
block_token / block_dhidden / block_dexpert 是否适合当前 shape
expert token 分布不均是否导致 block 内大量空行
register / shared memory 是否限制 occupancy

MoE 中尤其要关注 up_logits 这个 global intermediate。当前实现中,Step 1 会把 gate/up/SiLU/multiply 后的中间结果写入 global memory,Step 2 再从 global memory 读回来做 down projection。对于大 hidden、大 expert_dim、大 routed token 数的 case,这个中间 buffer 可能成为显著的 global memory traffic 来源。

如果 mcProfiler 显示 memory utilization 高而 compute utilization 低,则说明优化方向应该是减少 global memory 读写、改善 coalescing、减少中间结果落 global memory,或者尝试更强的 kernel fusion / 分块融合。
如果 compute utilization 低但 memory utilization 也不高,则需要检查 expert 负载不均、block 内有效行数不足、launch overhead 或调度 stall。
如果 register 使用过高导致 occupancy 低,则需要调整 tile size、fragment 大小、num_stages 或使用 mxcc 的寄存器限制选项做 sweep。

1.12.5 MLA 的 profiling 重点

MLA 的主要瓶颈通常来自长 KV context 扫描。
profiling 时要重点观察:

KV / K_pe 读取是否成为 memory bandwidth 瓶颈
QK score 计算是否充分利用矩阵计算单元
online softmax 的 max/sum reduction 是否占用明显时间
softmax 中 exp/div 等数学操作是否成为热点
PV 累加是否产生过高 register pressure
block_N / block_H / num_split 是否适合长上下文
长 kv_ctx 下是否需要 split KV 并行化

如果长上下文下单个 block 扫描完整 KV,可能出现 block 数不足、并行度不够、单 block 工作过重的问题。此时需要考虑增大并行度,例如启用 split KV、调整 BLOCK_N、调整 head block 映射,或者降低单个 block 的寄存器和 shared memory 压力。

1.12.6 NSA 的 profiling 重点

NSA 的主要瓶颈通常来自稀疏索引和不规则访存。
profiling 时要重点观察:

BlockIndices 间接访存是否导致 global load 不连续
selected blocks 循环是否造成负载不均
不同 block 的工作量是否差异过大
causal mask 是否造成明显分支发散
稀疏 gather 是否拖慢 memory pipeline
index 计算是否消耗过多 register

如果 NSA 的 memory stall 很高,通常不是简单增加计算 tile 就能解决,而要优先优化 sparse block layout、索引局部性、selected block 并行策略和访存合并。

1.12.7 从 profiler 结论反推优化动作

Profiler 阶段最后必须落到优化动作,而不是只展示指标。

可以按照下面的规则整理:

现象:global memory bandwidth 高,compute utilization 低
结论:memory-bound
优化:减少中间 tensor、改善 coalescing、复用 shared memory、融合 epilogue

现象:compute utilization 高,memory utilization 不高
结论:compute-bound
优化:检查 MMA lower、调整 tile size、提高矩阵计算单元利用率

现象:kernel launch 数量多,每个 kernel 很短
结论:latency-bound / launch overhead 高
优化:kernel fusion、减少 PyTorch 侧小 op、使用 cached_e2e 或 graph capture 思路

现象:occupancy 低,register 使用高
结论:register pressure 限制并发
优化:减小 tile / fragment、降低 num_stages、尝试 maxrregcount sweep

现象:shared memory 使用高或存在 bank conflict
结论:shared memory 成为限制
优化:调整 shared layout、padding、swizzle、tile shape

现象:某些 expert 对应 block 很慢或有效行很少
结论:MoE expert load imbalance
优化:优化 routing 分组、改 block_token、减少 padding 浪费

现象:MoE official_e2e 慢,但 routed_kernel_only 快
结论:瓶颈在 host/PyTorch 前后处理
优化:减少 clone、缓存对象和 buffer、device 侧构造 metadata、优化 packing 和 scatter_reduce

最终报告中建议把 profiler 结果整理成表格:

scope case latency 主要瓶颈 证据 优化动作
official_e2e MoE case 1 xx ms host + scatter scatter_reduce / packing 占比高 优化 packing,减少 index repeat
cached_e2e MoE case 1 xx ms routing / scatter kernel-only 占比较低 复用 buffer,减少同步
routed_kernel_only MoE case 1 xx ms global memory up_logits 读写高 减少 intermediate global traffic
kernel MLA 64K xx ms KV scan memory-bound memory utilization 高 split KV,调 BLOCK_N
kernel NSA 64K xx ms sparse gather load 不连续 / stall 高 优化 block index layout

Profiler 阶段的核心结论应该写成:

benchmark 给出性能结果,profiler 解释性能原因。
Python/host 分段计时用于拆解 API 端到端耗时;
mcProfiler 用于分析 GPU kernel 内部瓶颈;
最终根据 profiler 证据决定优化方向,包括 tile 参数调整、kernel fusion、
减少 global memory 中间结果、优化 token packing、降低同步和减少 launch overhead。

1.13日志结构

官方仓库目前没有给出具体的日志结构,下述日志是被认为需要保存的。

raw log:
    完整终端输出,用于复现问题和提交原始证据。

JSON / JSONL:
    结构化结果,适合脚本聚合、生成报告。

CSV:
    人能快速打开看,适合比赛提交和横向比较。

1.14总结

层级 应当注意的问题 典型操作
race_tests 比赛测什么 读 shape、dtype、reference、benchmark
PyTorch 输入输出和计时边界 看是否有额外 torch op
mcPyTorch CUDA 兼容是否透明 确认 tensor/device/event 行为
TileLang DSL 算子怎么写 设计 tile、thread、shared、reduction
TVM/TIR 编译后结构是否符合预期 dump IR,看 loop/thread/memory
MACA backend 是否适配 GPU 查 unsupported path,避免 CUDA 假设
mxcc 能否正确编译 看编译报错、寄存器、资源限制
Runtime kernel 如何启动 减少 launch,同步和 event 要明确
Driver 设备状态和错误 看 runtime error、设备状态
GPU 性能极限 按 wave=64、带宽、矩阵单元优化
profiler 慢在哪里 用数据反推优化方向

2. 评测工程

第一章讲的是 MetaX-TileLang 工具链如何把 TileLang kernel 编译并运行到曦云 C500 GPU 上。第二章要解决的是另一个问题:如何证明这个 kernel 是可评测、可复现、正确且有性能收益的工程实现

在比赛场景里,单纯写出一个能运行的 kernel 还不够。最终提交通常需要同时满足:

correctness:输出必须和 golden / reference 在误差阈值内一致
performance:端到端耗时必须有竞争力
report:需要有功能测试报告、性能测试报告和算法设计报告
reproducibility:别人能够按照 README 在相同环境中复现结果

因此,评测工程的核心不是“跑一次 benchmark”,而是建立完整闭环:

C500 在线容器
    ↓
环境检查:driver / SDK / mxcc / TileLang / Python deps
    ↓
TileLang kernel 实现
    ↓
TileLang 编译与后端 lowering
    ↓
MXMACA 编译 / JIT / 动态加载
    ↓
Runtime API 调度 kernel
    ↓
correctness:golden 对比 + 日志保存
    ↓
benchmark:warmup + repeat + sync + 端到端计时
    ↓
profiler:host 分段计时 + mcProfiler kernel 分析
    ↓
性能报告:case 表格 + 平均耗时 + 失败样例 + profiler 摘要
    ↓
打榜提交:源码 + README + 功能报告 + 性能报告 + 算法设计报告

2.1 race_tests:比赛评测入口

race_tests 是评测工程的入口。它规定了每个算子的输入输出、shape、dtype、layout、reference、correctness 规则和 benchmark 计时方式。

当前仓库中主要包含三个方向:

race_tests/
  moe/
    README.md
    custom_fusedmoe.py
    fusedmoe_benchmark.py
    moe_test_configs.json
    ref_fusedmoe.py
    run.sh

  mla/
    test_cases_mla_batch_ctx.json
    test_tilelang_mla.py

  nsa/
    reference.py
    test_cases_nsa_fwd.json
    test_tilelang_nsa_fwd.py

对每个算子,首先要读清楚下面这些内容:

输入 tensor 的 shape 是什么
输出 tensor 的 shape 是什么
dtype 是什么
layout 是什么
golden / reference 是什么
误差容忍度是多少
benchmark 怎么计时
warmup / repeat 次数是多少
哪些 case 必须通过
哪些文件是正式提交接口
哪些文件只是测试辅助

以 MoE 为例,目录中最关键的文件是:

custom_fusedmoe.py       # 要优化和提交的 TileLang MoE kernel 接口
fusedmoe_benchmark.py    # 功能测试 + 性能测试驱动
ref_fusedmoe.py          # PyTorch golden reference
moe_test_configs.json    # 功能 / 性能 case 配置
run.sh                   # 启动脚本
README.md                # 使用说明

其中 custom_fusedmoe.py 是正式提交接口。评测脚本会按照固定方式调用:

RoutedMoEKernel(...)
kernel(...)

因此优化时可以重写内部实现,但不应该随意修改外部接口。ref_fusedmoe.py 只作为 PyTorch reference,不参与性能优化。fusedmoe_benchmark.py 是测试入口,负责生成输入、调用 reference、调用 custom kernel、做 correctness 对比和 performance 计时。

2.2 评测标准:correctness + performance + report + reproducibility

比赛评测通常分成四个维度:

1. correctness
   输出是否和 golden/reference 在误差阈值内一致。

2. performance
   是否在官方计时规则下有更低 latency 或更高 throughput。

3. report
   是否能清楚说明测试 case、误差、性能、优化思路和 profiler 证据。

4. reproducibility
   是否能在 C500 环境中按照 README 复现结果。

初赛核心是 Fused MoE GEMM。基本门槛是功能 test 集通过率达到要求,并且相对参考样例有性能提升。性能排名使用算法测试 API 的端到端耗时,通常按照 warmup 10 次、repeat 100 次取平均耗时。

决赛可以选择 MLA 或 NSA。MLA 的技术分通常由 correctness 和 performance 组成,性能同样按照 warmup 10 次、repeat 100 次取平均。NSA 则重点考察 64K sequence length 输入下的 correctness 和 performance,通常使用 warmup 10 次、repeat 50 次取平均。

所以评测工程的第一原则是:必须严格对齐官方计时边界和测试规则。自己开发阶段可以额外做 kernel-only、cached_e2e、profiler 分析,但最终提交报告必须说明哪些结果是官方端到端计时,哪些结果是辅助分析计时。

2.3 环境检查:先确认评测环境可信

进入 C500 在线容器或赛事环境后,首先检查设备和软件栈:

which mxcc
mxcc --version

python -c "import tilelang; print(tilelang.__version__)"
python -c "import torch; print(torch.__version__)"

建议把这些信息写进日志和报告:

GPU 型号
driver version
MXMACA SDK version
mxcc version
TileLang version / commit
PyTorch / mcPyTorch version
Python version
当前 git commit
是否清理 TileLang cache
是否开启 fast math
是否开启 lineinfo / debug
编译选项

原因是性能结果和环境强相关。不同 driver、SDK、TileLang commit、mxcc codegen 选项、fast math 配置,都可能影响 correctness 和 performance。

2.4 correctness:golden 对比与功能测试报告

correctness 的目标是证明自定义 TileLang kernel 的输出和 reference 一致。

以 MoE 为例,functional 分支的核心逻辑是:

data = generate_input(**config)

ref_output = ref_kernel(clone_data(data)).to(torch.float32)
tilelang_output = custom_kernel(clone_data(data)).to(torch.float32)

torch.testing.assert_close(
    ref_output,
    tilelang_output,
    atol=1e-2,
    rtol=1e-2
)

这段逻辑可以拆成 6 步:

1. 根据 config 生成输入和权重
2. clone 一份 data 给 PyTorch reference
3. ref_kernel(...) 生成 golden 输出
4. clone 一份 data 给 TileLang custom kernel
5. custom_kernel(...) 生成 TileLang 输出
6. assert_close 比较 golden 和 TileLang 输出

这里的 golden 不是 CPU 手写结果,而是 ref_fusedmoe.py 里的 PyTorch reference 输出。
ref_fusedmoe.py 定义了 ExpertTorchMoEGateTorchMoETorchref_kernel。其中 expert MLP 的数学逻辑是:

gate = SiLU(x @ W_gate)
up   = x @ W_up
mid  = gate * up
out  = mid @ W_down

router 的逻辑是:

logits = W_g(x)
scores = softmax(logits)
topk_scores, topk_indices = topk(scores)

也就是:

router logits → softmax → top-k expert id 和 top-k score

功能测试不应该只打印 pass / fail。报告里至少应该记录:

case_id
test_type
status
config
seed
input shape
output shape
input dtype
compare dtype
golden function
custom function
atol
rtol
max_abs_err
max_rel_err
num_elements
num_failed
element_pass_rate
error_type
error_msg
timestamp
git_commit
tilelang version
torch version
device name

其中最关键的是:

config
golden
atol / rtol
max_abs_err / max_rel_err
status

最终功能测试报告应该回答:

这个 case 的参数是什么?
golden 是什么?
误差阈值是多少?
最大绝对误差是多少?
最大相对误差是多少?
是否通过?
失败时失败在哪里?
失败日志是否足够复现?

2.5 benchmark:warmup、repeat 与端到端计时

benchmark 的目标是得到稳定、可信、可复现的性能数据。

开发阶段可以使用不同计时范围,但比赛性能排名通常看算法测试 API 的端到端耗时。典型形式是:

start_timer()
output = your_api(input)
sync()
end_timer()

只要发生在 your_api(input) 里面的事情,都可能算进端到端耗时。它可能包括:

host 侧参数检查
Python wrapper 调用
Tensor reshape / view / contiguous
临时 buffer 分配
H2D / D2H / D2D 数据搬运
sort / topk / gather / scatter 等前后处理
TileLang kernel launch
一个或多个 device kernel 执行
stream synchronize
output 后处理

所以端到端优化不是只优化 kernel body。MoE 尤其明显,因为完整 API 里可能包含:

router
topk
argsort
bincount
token packing
group metadata 构造
grouped GEMM
scatter_reduce
synchronize

如果 benchmark 计时边界包住这些步骤,那么它们全部都会影响最终成绩。

2.5.1 为什么要 warmup?

warmup 的作用是排除第一次运行中不稳定的一次性开销。第一次运行可能包含:

TileLang JIT 编译
MXMACA 后端 lowering / codegen
mxcc 编译
kernel module 加载
runtime 初始化
CUDA/MXMACA context 初始化
memory allocator 第一次分配
cache 冷启动
Python wrapper 第一次触发某些路径

如果不做 warmup,第一次运行的冷启动开销会污染性能结果。比赛更关心稳定态推理性能,所以通常先 warmup,再正式 repeat 计时。

2.5.2 为什么要 repeat?

repeat 的作用是降低随机波动。单次运行可能受到 runtime 调度、stream 前后任务、host Python 调用、allocator 状态、cache 状态和系统负载影响。

最终报告中除了 mean latency,最好也保存:

min latency
max latency
median latency
p50
p90
p99
std

排名可能看平均耗时,但调优时要看分布。如果 p99 很差,说明存在不稳定因素,例如同步、内存分配、cache miss 或系统负载干扰。

2.6 MoE benchmark 计时边界分析

MoE 的性能分支通常是:

for i in range(warm_up):
    _ = custom_kernel(clone_data(data))

start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)

start_event.record()
for i in range(iteration):
    _ = custom_kernel(clone_data(data))
end_event.record()

torch.cuda.synchronize()
elapsed_ms = start_event.elapsed_time(end_event)
elapsed_ms = elapsed_ms / iteration

因此 benchmark 打印的耗时是平均每次:

custom_kernel(clone_data(data))

的时间。

这不是纯 TileLang kernel 时间,而是 MoE 开发样例中的 Python 端到端时间。它大致包含:

clone input / weights
RoutedMoEKernel / MoE object construction
weight transpose / contiguous / to / stack
intermediate buffer allocation
router GEMM
softmax
topk
argsort
bincount + CPU sync
token packing
group metadata construction
explicit cuda synchronize
TileLang routed grouped GEMM: gate / up / SiLU / down
scatter_reduce
final cuda synchronize

因此可以写成:

MoE measured latency
= average over 100 repeats of custom_kernel(clone_data(data))
= clone_data
  + object construction
  + weight preparation
  + buffer allocation
  + router / topk / sort
  + packing / metadata
  + TileLang routed kernel
  + scatter_reduce
  + synchronize

这点很重要。因为如果端到端耗时很差,不一定说明 TileLang GEMM kernel 很差,也可能是 clone_data、Python 分组、bincount().cpu().numpy()scatter_reduce 或同步导致的。

2.7 host 侧开销:识别和量化

host 侧开销不是 TileLang kernel 本身的计算,但会出现在 Python / PyTorch / runtime / 数据准备 / kernel launch 周围,并且可能进入官方端到端计时。

MoE 中常见 host 侧开销包括:

1. 临时 buffer 分配
2. tensor clone / copy
3. CPU ↔ GPU 数据搬运
4. GPU 内部 D2D 数据搬运
5. 显式或隐式同步
6. Python / PyTorch dispatch
7. kernel launch 开销
8. JIT 编译 / cache / 对象构造开销

可以整理成表格:

类别 官方代码中的例子 是否属于 TileLang kernel 对性能的影响 报告中如何处理
临时 buffer torch.empty, torch.zeros, torch.stack 显存分配、allocator、初始化 单独列 buffer 大小
clone clone_data(data) D2D copy、显存分配 比较 official vs no_clone
D2H .cpu().numpy() 强同步、GPU → CPU copy 标记为 host sync
H2D torch.tensor(..., device="cuda") 小 tensor copy、dispatch 尝试 device 侧 metadata
Python loop expert 分组 loop 多次小 op 调度 单独测 grouping_ms
PyTorch dispatch softmax, topk, argsort, scatter_reduce 多 kernel launch E2E 与 kernel-only 分开
显式同步 torch.cuda.synchronize() 阻塞异步执行 外层统一同步
TileLang launch self.routed_kernel(...) 真实 kernel launch kernel-only 计时
global intermediate up_logits_routed 是 / 算法设计 global 写读很大 算法报告重点分析

MoE 官方样例里典型 host 开销包括:

1. clone_data(data)
   每轮 clone 输入和权重,产生 D2D copy 和显存分配。

2. 每轮 custom_kernel 都构造 RoutedMoEKernel 和 MoE
   产生 Python 对象构造、buffer 分配、权重 stack、可能的 JIT/cache 开销。

3. counts = flat_expert_indices.bincount().cpu().numpy()
   产生 D2H copy 和强同步。

4. Python for-loop 做 expert 分组
   产生多次 gather、slice assignment 和 Python dispatch。

5. torch.tensor(..., device="cuda")
   每轮创建 group metadata,产生小 tensor H2D copy。

6. forward 内部 torch.cuda.synchronize()
   阻塞 GPU 异步执行。

7. torch.scatter_reduce
   属于 PyTorch op,不是 TileLang grouped GEMM。

8. up_logits_routed
   是巨大的 device global intermediate,Step 1 写、Step 2 读。

2.8 profiling:从总耗时反向定位瓶颈

benchmark 告诉我们“快还是慢”,profiler 负责解释“为什么快或慢”。

Profiler 阶段建议分成两层:

第一层:Python / host 侧分段计时
第二层:GPU kernel 级 mcProfiler

2.8.1 Python / host 侧分段计时

MoE 要先做 Python 侧分段计时,因为它的 host 开销很多。
官方 MoE 代码里 router/topk、expert grouping、TileLang routed kernel、scatter_reduce 明确分布在不同阶段,适合做分段。

可以设置三种 scope:

official_e2e:
    保持官方 fusedmoe_benchmark.py 逻辑。
    包含 clone_data、对象构造、CPU grouping、同步、scatter_reduce。
    最接近比赛 API 端到端耗时。

cached_e2e:
    不在循环内构造 RoutedMoEKernel / MoE。
    不 clone_data。
    仍包含 routing、grouping、scatter_reduce。
    用来评估对象构造、JIT cache lookup、buffer 分配等开销。

kernel_only:
    预先准备 stacked_expert_tokens、group_sizes、offsets、weights。
    只测 self.routed_kernel。
    用来评估 TileLang routed expert MLP kernel 本体性能。

通过这三种 scope,可以判断:

official_e2e 慢,但 kernel_only 不慢:
    瓶颈主要在 host/PyTorch 前后处理,例如 clone、packing、scatter、同步。

kernel_only 本身慢:
    瓶颈在 TileLang kernel 内部,需要看 TIR、MACA backend、mxcc 编译和 mcProfiler。

cached_e2e 比 official_e2e 明显快:
    说明对象构造、buffer 分配、clone 或 JIT/cache 查询占比明显。

2.8.2 GPU kernel mcProfiler

mcProfiler 是 MXMACA 软件栈中的可视化 GPU 性能指标分析工具,可以从 SOL / RoofLine 全局视角,以及 Memory、Computing、Scheduling 等角度分析 kernel。

mcProfiler 适合回答:

这个 kernel 是 memory-bound 还是 compute-bound?
HBM / global memory 读写是否过高?
load/store 是否连续?
register / spill 是否限制 occupancy?
shared / WSM 是否过高?
Scheduling stall 主要来自哪里?
kernel launch 数量是否过多?
MLA 的长 KV scan 是否被 memory bandwidth 限制?
NSA 的 BlockIndices 间接访存是否拖慢?
MoE 的 up_logits global intermediate 是否是主要瓶颈?

MoE 的 profiler 重点:

TileLang Step 1 和 Step 2 是否耗时均衡
gate/up GEMM 是否充分利用矩阵计算单元
down GEMM 是否成为主要瓶颈
up_logits global intermediate 是否造成大量 global memory 读写
block_token / block_dhidden / block_dexpert 是否适合当前 shape
expert token 分布不均是否导致 block 内浪费
register / shared memory 是否限制 occupancy

MLA 的 profiler 重点:

长 KV context scan 是否 memory-bound
QK score 是否充分利用矩阵计算单元
online softmax 的 max/sum reduction 是否成为瓶颈
PV 累加是否造成 register pressure
block_N / block_H / num_split 是否合适
长上下文下是否需要 split KV

NSA 的 profiler 重点:

BlockIndices 间接访存是否导致不连续 load
selected blocks 循环是否造成负载不均
causal mask 是否造成分支发散
稀疏 gather 是否拖慢 memory pipeline
index 计算是否消耗过多 register

最终 profiler 结果不要只贴图,而要转化成优化动作:

现象 结论 优化方向
global memory bandwidth 高,compute utilization 低 memory-bound 减少中间 tensor、改善 coalescing、复用 shared memory
compute utilization 高 compute-bound 检查 MMA lower、调 tile size、提高矩阵单元利用率
kernel launch 多,每个 kernel 很短 latency-bound kernel fusion、减少 PyTorch 小 op、缓存对象和 buffer
occupancy 低,register 高 register pressure 减小 tile / fragment,调整 num_stages,sweep maxrregcount
shared memory 高或 bank conflict shared memory 限制 调 shared layout、padding、swizzle
某些 expert 有效行很少 MoE load imbalance 调 block_token、优化路由分组、减少 padding 浪费
official_e2e 慢但 kernel_only 快 host/PyTorch 开销 优化 packing、metadata、scatter_reduce、同步

2.9 性能参照系:mcBLAS / mcDNN baseline

MoE 的核心计算本质上是 expert MLP。对于每个 routed token,主要有三次 GEMM:

gate projection:
    X @ W_gate

up projection:
    X @ W_up

down projection:
    hidden @ W_down

因此,mcBLAS 可以作为三类性能参照:

1. 纯 GEMM 层面的性能边界
2. grouped / batched expert GEMM 分解实现的 baseline
3. 判断 TileLang T.gemm tile 参数是否合理的参照

一种粗略 baseline 是:

for each expert:
    mcBLAS GEMM for gate
    mcBLAS GEMM for up
    activation + multiply
    mcBLAS GEMM for down

mcDNN 可以作为两类间接参考:

1. 标准 MLP / activation / DNN primitive 的性能参考
2. 和 mcPyTorch / PyTorch backend 下的 nn.Linear + SiLU 组合做间接对比

需要注意的是,当前 MoE 仓库本身没有显式调用 mcBLAS / mcDNN,而是用 PyTorch op 做 router/topk/scatter,用 TileLang 实现 routed expert grouped GEMM。

因此不能直接用单个 mcBLAS GEMM 的时间和 MoE 端到端 latency 比较。
原因是当前 MoE 计时边界是:

custom_kernel(clone_data(data))

它包含 router、topk、sort、packing、metadata、TileLang kernel、scatter_reduce 和同步。单个 mcBLAS GEMM 只能作为 expert MLP 计算部分的性能边界,而不能代表完整 MoE API 的端到端耗时。

更合理的对比方式是:

pure GEMM baseline:
    mcBLAS 测 gate / up / down 单个 GEMM 或 batched GEMM 的性能边界。

expert MLP baseline:
    mcBLAS / mcDNN 组合 gate + up + SiLU + multiply + down。

MoE end-to-end baseline:
    保持官方 benchmark 边界,统计 router、topk、packing、TileLang kernel、
    scatter_reduce 和同步后的整体耗时。

如果 kernel_only 明显低于 mcBLAS 类似 GEMM 的效率,说明 TileLang kernel 内部还需要优化。
如果 kernel_only 已经接近合理水平,但 official_e2e 仍然很慢,说明主要瓶颈可能在 PyTorch 前后处理、token packing、metadata 构造、scatter_reduce 或同步上。

2.10 结果校验:不只是 assert pass

结果校验需要确认:

1. shape 是否正确
2. dtype 是否符合预期
3. device 是否正确
4. 输出中是否有 NaN / Inf
5. 输出是否和 golden/reference 在误差阈值内一致
6. 多次运行是否 deterministic / consistent
7. 边界 case 是否通过
8. 官方测试 case 是否全部覆盖
9. 如果大 case 跳过 golden,是否有替代校验策略
10. 每个失败 case 是否保存足够日志用于复现

对于每个失败 case,至少保存:

config
seed
input/output shape
dtype
atol/rtol
max_abs_err
max_rel_err
num_failed
failing index
error message
raw log
环境信息
git commit

这部分最终进入 correctness_report.md

2.11 报告生成与目录结构

评测工程最好不要只依赖终端输出,而应该统一生成日志、CSV、JSON 和 Markdown 报告。

推荐目录结构:

reports/
  performance_report.md
  correctness_report.md
  algorithm_design_report.md

logs/
  raw/
    moe.log
    mla.log
    nsa.log
  json/
    moe_perf.json
    mla_perf.json
    nsa_perf.json
  csv/
    moe_results.csv
    mla_results.csv
    nsa_results.csv
  profiler/
    moe/
    mla/
    nsa/

报告生成链路可以固定为:

官方 race_tests 原始脚本
    ↓
每个算子生成原始 CSV / JSON / log
    ↓
统一 report collector 汇总
    ↓
生成 performance_report.md
    ↓
生成 correctness_report.md
    ↓
把 host overhead audit 写进 algorithm_design_report.md

其中:

raw/*.log:
    保存完整终端输出,用于复现问题和提交原始证据。

json/*.json:
    保存结构化结果,适合脚本聚合。

csv/*.csv:
    保存表格结果,适合横向比较和画图。

profiler/:
    保存 mcProfiler 或 TileLang profiler 的截图、导出文件和文字摘要。

2.12 提交材料:源码、README、功能报告、性能报告、算法设计报告

最终提交材料可以分成五类。

2.12.1 源码

源码要保证可运行、可评测、接口稳定。

以 MoE 为例,需要说明:

custom_fusedmoe.py 是正式提交 kernel。
RoutedMoEKernel 的接口固定,只优化内部实现。
ref_fusedmoe.py 只作为 PyTorch reference,不参与性能优化。
fusedmoe_benchmark.py 是测试入口。
moe_test_configs.json 是 case 配置。
run.sh 只是执行 python fusedmoe_benchmark.py。

2.12.2 README

README 要说明:

环境要求
依赖安装
如何进入 C500 容器
如何运行 correctness
如何运行 benchmark
如何保存日志
输出文件在哪里
已知限制

2.12.3 功能测试报告

功能测试报告要说明:

每个 case 的参数
golden/reference 是什么
误差阈值
max_abs_err
max_rel_err
是否通过
失败 case 和错误信息

2.12.4 性能测试报告

性能测试报告要说明:

op_name
shape
dtype
layout
device
warmup 次数
repeat 次数
计时范围
mean / median / p50 / p90 / p99
TFLOPS / GB/s / token/s
profiler 摘要
编译选项
环境版本

2.12.5 算法设计报告

算法设计报告要说明:

tile / block / thread 如何设计
shared memory 如何使用
register / fragment 如何使用
MMA 是否成功 lower
pipeline 是否生效
layout 是否适配访存
host 侧开销有哪些
当前瓶颈是什么
下一步优化方向是什么

2.13 本章小结

评测工程的核心是建立从 correctness 到 performance 再到 profiler 的闭环。

correctness 证明结果对;
benchmark 证明速度快;
profiler 解释为什么慢;
report 证明结果可信;
README 保证别人能复现。

对于 MetaX-TileLang kernel 优化,尤其不能只看单个 kernel body。
MoE 这类算子的端到端耗时会包含 router、topk、argsort、bincount、token packing、metadata 构造、TileLang grouped GEMM、scatter_reduce、同步和 buffer 分配。
因此评测工程必须同时关注:

官方端到端计时
kernel-only 计时
host overhead
GPU kernel profiler
correctness 日志
性能报告
提交接口稳定性

只有把这些链路全部打通,后续的 tile size、thread layout、shared memory、register、MMA、pipeline、layout 优化才有可靠依据。

posted @ 2026-06-01 19:04  White_Swan  阅读(39)  评论(0)    收藏  举报