update
MetaX C500 GPU 上的 CUDA / TileLang 高性能计算 Kernel 仓库。 聚焦 DeepSeek 风格 MoE (Mixture of Experts)、FlashMLA 和 Native Sparse Attention (NSA) 等核心算子的 TileLang kernel 实现与性能优化。 配套 mcProfiler 性能分析工具链与 mxmaca-performance-tuning-guide 性能优化指南。
本仓库的核心目标是在 MetaX(沐曦)C500 GPU 上,使用 TileLang 与 CUDA 编写并优化大语言模型推理中的关键计算 Kernel,包括:
推荐阅读顺序(对 C500 高性能 Kernel 开发者):
GPU └─ 8 × DPC └─ 每个 DPC 包含 13 × AP ├─ Wave 调度器(4 组,对应 4 个 PEU) ├─ MMA 矩阵计算单元 ├─ MTE 数据搬运单元 ├─ STE 标量执行单元 └─ 存储层次: ├─ 向量寄存器堆:512 KiB/AP(2048 × 32-bit/lane) ├─ 标量寄存器堆:800 × 32-bit/AP ├─ WSM (Shared Memory):64 KiB/AP ├─ VL1 (Vector L1 Cache):32 KiB/AP(默认关闭) ├─ L2 Cache:8 MB(全芯片共享) └─ 全局内存 (DRAM):64 GB,峰值带宽 1843 GB/s
C500 的存储子系统采用三层架构,类似认知科学的三层记忆模型:
工作记忆区(寄存器堆) ↕ 全双工 短期记忆区(WSM / VL1 / L2 Cache) ↕ 半双工 长期记忆区(DRAM 全局内存)
关键特性:
每周期最多发射 5 条指令(来自不同 warp):
/data/metaxRace/ ├── README.md ← 本文档 ├── init_tilelang.sh ← TileLang 安装/注册脚本 ├── skills/ ← 技能目录 │ ├── install_skills.sh │ └── mcprof-report-skill/ ← mcProfiler 报告分析技能 │ ├── SKILL.md │ ├── metax-kernel-guide.md ← C500 Kernel 开发指南 │ └── tests/conformance/ ← 一致性测试 ├── doc/ ← 文档目录 │ ├── record.md ← 术语对照与架构速查 │ ├── metax_c500_arch.md ← C500 架构说明 │ ├── image.png │ ├── GPU扫盲/ │ │ └── gemm.md ← GEMM Tile 切分基础 │ ├── problem_description/ ← OJ 题目描述 │ │ ├── oj_requirement.md ← 评测系统规范(CUDA/Triton/TileLang) │ │ └── moe.md ← MoE 题目描述 │ ├── mxmaca-performance-tuning-guide/ ← 性能优化指南 │ │ ├── README.md │ │ ├── guide/ ← 10 章优化文档 │ │ │ ├── ch1.初探异构编程.vectoradd.md │ │ │ ├── ch2.曦云C500芯片架构.md │ │ │ ├── ch3.Kernel编程入门.reduction.md │ │ │ ├── ch4.Kernel性能建模.sgemm.md │ │ │ ├── ch5.Kernel性能优化技巧.md │ │ │ ├── ch6.Kernel性能分析工具.md │ │ │ ├── ch7.Host代码性能优化.md │ │ │ ├── ch8.张量编程.hgemm.md │ │ │ ├── ch9.C500_HW_limitation.md │ │ │ └── ch10.常用Compiler参数和Driver环境变量.md │ │ ├── case/ ← 文档对应示例代码 │ │ │ ├── 1.vector_add/ ← 向量加法 │ │ │ ├── 2.reduction/ ← 归约 │ │ │ ├── 3.sgemv/ ← 矩阵-向量乘 │ │ │ └── 4.sgemm_tt/ ← 矩阵乘法 (18 个版本) │ │ ├── microbenchmark/ ← 微架构参数测试 │ │ └── requirements.txt │ ├── kernel-optimize/ ← Kernel 优化笔记 │ │ └── moe/ │ │ ├── moe.md ← MoE Kernel 全面优化记录 │ │ ├── analyze.md ← 流水线与卡顿分析 │ │ ├── plan.md ← 优化计划 │ │ ├── case.md ← 案例研究 │ │ ├── fused.md ← 融合分析 │ │ ├── kernel.md ← Kernel 细节 │ │ ├── analyze-data-reuse.md ← 数据复用分析 │ │ ├── optimize.md ← 优化分析整理版 │ │ └── optimize/ │ │ ├── HANDOFF.md │ │ ├── bsm_instruction_num.md │ │ ├── plan.md │ │ ├── state.md │ │ └── tl_thread.md │ └── skill/ │ └── review.md ← 代码评审指南 ├── example/ ← 参考实现 │ ├── README.md ← 学习路线指南 │ ├── FlashMLA/ ← FlashMLA 参考实现 │ │ ├── README.md │ │ ├── setup.py │ │ ├── flash_mla/ │ │ ├── csrc/ ← CUDA 源码 │ │ └── tests/ │ ├── mcoplib/ ← 沐曦算子库源码 │ │ ├── README.md │ │ ├── kernel/ ← CUDA kernel │ │ ├── op/ ← 算子实现 (CUDA) │ │ └── benchmark/ ← 性能测试 │ │ ├── config/ ← 130 个算子 benchmark 配置 │ │ └── runners/ ← 131 个 benchmark runner │ ├── kernels/ ← 各种 kernel 参考 │ │ ├── FlashMLA/ ← FlashMLA 参考 │ │ ├── TileKernels-Metax/ ← TileLang kernel 集 │ │ ├── TileOPs-Metax/ ← TileLang ops │ │ ├── mcoplib/ ← mcoplib 符号链接 │ │ └── vLLM-metax/ ← vLLM 适配代码 │ └── mcoplib/ ← 沐曦官方算子库(完整版) │ ├── CMakeLists.txt │ ├── setup.py / pyproject.toml │ ├── include/ ← 头文件 │ ├── kernel/ ← CUDA kernel (14 个文件) │ ├── op/ ← 算子入口 (31 个 CUDA 文件) │ ├── benchmark/ ← 性能测试框架 │ └── unit_test/ ├── race_tests/ ← 竞赛测试用例 │ ├── moe/ ← MoE Fused Kernel │ │ ├── README.md ← 使用说明 │ │ ├── doc.md ← 算法拆解 │ │ ├── custom_fusedmoe.py ← TileLang kernel(主入口) │ │ ├── ref_fusedmoe.py ← PyTorch 参考实现 │ │ ├── fusedmoe_benchmark.py ← host 侧编排 │ │ ├── moe_test_configs.json ← 测试配置 │ │ ├── moe_prof_driver.py ← Profiling 被测进程 │ │ ├── profile_mcprof.sh ← 一键 profiling 脚本 │ │ ├── mcprof_summary.py ← Profiler 报告摘要解析 │ │ ├── run.sh / run_min.sh ← 运行脚本 │ │ └── submission*/ ← 各版本提交 │ ├── mla/ ← Multi-head Latent Attention │ │ ├── test_tilelang_mla.py ← TileLang MLA Kernel │ │ └── test_cases_mla_batch_ctx.json │ └── nsa/ ← Native Sparse Attention │ ├── test_tilelang_nsa_fwd.py ← TileLang NSA Kernel │ ├── test_cases_nsa_fwd.json │ └── reference.py ← PyTorch 参考实现
OJ 平台支持三种语言提交:CUDA(C/C++ 源码)、Triton、TileLang。
单测试点执行步骤:
S(Tk)=1001+(1s−1)Tk−ThTb−ThS(T_k)=\dfrac{100}{1+\left(\dfrac{1}{s}-1\right)\dfrac{T_k-T_h}{T_b-T_h}}S(Tk)=1+(s1−1)Tb−ThTk−Th100
run_kernel
run_kernel(*args)
torch
tilelang
tilelang.language
tilelang.intrinsics
math
@tilelang.jit
@T.prim_func
T.Tensor((M, K), T.float16)
T.Tensor[(M, K), T.float16]
os
sys
subprocess
socket
pickle
exec
eval
compile
TorchProxyError
实现 DeepSeek 风格 MoE 中的 pre-routed fused expert GEMM。
计算链条:
for each expert e: gate_logits = stacked_expert_tokens @ gate_w[e]^T [d_hidden → d_expert] up_logits = stacked_expert_tokens @ up_w[e]^T [d_hidden → d_expert] hidden = SiLU(gate_logits) * up_logits [d_expert] output = hidden @ down_w[e]^T [d_expert → d_hidden] output *= routed_expert_weights [逐行缩放]
重要:routed_expert_weights 为 torch.float32(非 float16),需在 TileLang kernel 中单独声明 dtype。
routed_expert_weights
torch.float32
设计上采用 2 次 launch(相对朴素 5 次 launch 的融合):
Stage 1:gate/up 双 GEMM + SiLU 融合
grid: (M, ceil(d_expert/128)) K 维 = d_hidden, T.Pipelined 循环 每轮: 读 input[128,128] + gate_w[128,128] + up_w[128,128] 两次 T.gemm, fp32 累加 Epilogue: SiLU(gate) * up → 写回 up_logits
Stage 2:down GEMM + 权重缩放融合
grid: (M, ceil(d_hidden/128)) K 维 = d_expert, T.Pipelined 循环 每轮: 读 up_logits[128,128] + down_w[128,128] Epilogue: output *= routed_expert_weights[m]
MoEGate
x
gate_logits
cfg0 瓶颈:卡在 Stage 2(K=8192 的 64 轮长依赖链)
cfg1 瓶颈:卡在 Stage 1(K=7168 的 56 轮 + smem 锁死 1 block/AP)
关键发现:K 链长的一侧(16 轮 → 30K per-wave;56-64 轮 → 104-118K per-wave)构成”结构性必要等待”——所有并行度杠杆(split-K / 小 tile / t1024 / 双缓冲)均实测无效或负收益。
不是一次 launch 的终极融合,而是两次连续 launch 的”两段式融合”。
“gate/up 产出 → down 消费”之间必须经全局内存——因为一个 128-token tile 的 up_logits 是 128×2048×2B=512KB,远超 C500 的 64 KiB WSM。这是所有主流实现都采用的结构。
up_logits
128×2048×2B=512KB
文件:race_tests/mla/test_tilelang_mla.py
race_tests/mla/test_tilelang_mla.py
TileLang 实现的 Flash Attention 风格 MLA kernel,支持 split 与 no-split 两种模式。
核心特性:
Q_pe
K_pe
num_split > 1
T.use_swizzle(10)
参考实现:example/FlashMLA/(含 CUDA 源码与 PyTorch 接口)
example/FlashMLA/
文件:race_tests/nsa/test_tilelang_nsa_fwd.py、race_tests/nsa/reference.py
race_tests/nsa/test_tilelang_nsa_fwd.py
race_tests/nsa/reference.py
基于 block-sparse 的原生稀疏注意力机制。
S
groups = HQ // H
naive_nsa
当前限制:
__restrict__
分析任意 Kernel 时按以下三步:
典型思考框架:
墙钟时间 ≈ Σ(每 AP 串行 block 数 × 每 block K 循环轮数 × 每轮串行下界) - 可被并行掩盖的延迟
mcProfiler 3.8.1.4 是 MetaX C500 配套的性能分析工具,基于 Client-Server 架构。
/usr/local/bin/mcProfiler
/opt/mcProfiler-ubuntu18.04/profiler_server
127.0.0.1:50123
/opt/mcProfiler-ubuntu18.04/gui-profiler-0.1.0.AppImage
/opt/mcProfiler-ubuntu18.04/config/*.pcd
cd /data/metaxRace/race_tests/moe ./profile_mcprof.sh # 使用默认配置 MOE_PROF_CFG=1 ./profile_mcprof.sh # 使用 cfg1 PER_KERNEL=1 ./profile_mcprof.sh 0 # 逐 kernel detail
HEAD: Summary — Total Instructions/Cycles, AP busy Duty, RoofLine CE Statistics — WORKGROUPS, WAVES, Average Wave life cycles ISU Statistics — stall cycles layout, DPC comparison Memory Statistics — Global/Private/Shared 指令数, L2C/VL1 Hit Rate Occupancy — Achieved/Dispatched waves GPU Throughput — AP/MMA/MTE/STE/VLS/L2C Duty ratio Compute workload — IPC, instruction throughput/efficiency
config/
/opt/mcProfiler-ubuntu18.04
perf done
report.txt.json
--headnames
--single-pass
race_tests/moe/custom_fusedmoe.py
race_tests/moe/ref_fusedmoe.py
race_tests/moe/fusedmoe_benchmark.py
race_tests/moe/moe_test_configs.json
race_tests/moe/moe_prof_driver.py
race_tests/moe/profile_mcprof.sh
race_tests/moe/mcprof_summary.py
example/FlashMLA/csrc/
example/mcoplib/kernel/
example/mcoplib/op/
example/mcoplib/benchmark/
example/kernels/
skills/mcprof-report-skill/SKILL.md
skills/mcprof-report-skill/metax-kernel-guide.md
skills/mcprof-report-skill/tests/conformance/
doc/mxmaca-performance-tuning-guide/guide/
doc/mxmaca-performance-tuning-guide/case/
doc/mxmaca-performance-tuning-guide/microbenchmark/
run_cfg0_20260810_025910
run_cfg0_20260810_032333
run_cfg0_20260810_033517
run_cfg0_20260810_041838
run_cfg0_20260810_050659
run_cfg1_20260810_055801
# 启动服务端 /opt/mcProfiler-ubuntu18.04/profiler_server & # 手动 perf_exec cd /opt/mcProfiler-ubuntu18.04 ./mcProfiler perf_exec \ --cmdline '/opt/conda/bin/python /data/metaxRace/race_tests/moe/moe_prof_driver.py --iters 30 --warmup 3' \ --casename 'moe_fused_gemm' \ [--per-kernel] [--counts 50] [--single-pass] # 查看指标 ./mcProfiler show_metrics ./mcProfiler version ./mcProfiler help
bash /data/metaxRace/init_tilelang.sh
安装脚本支持三级回退:
/opt/tilelang-metax-v0.1.10
.pth
pip install -e .
git clone + install
沐曦比赛的算子开发代码
版权所有:中国计算机学会技术支持:开源发展技术委员会 京ICP备13000930号-9 京公网安备 11010802047560号
metaxRace — MetaX C500 GPU 高性能 Kernel 优化项目
目录
1. 项目概览
1.1 目标
本仓库的核心目标是在 MetaX(沐曦)C500 GPU 上,使用 TileLang 与 CUDA 编写并优化大语言模型推理中的关键计算 Kernel,包括:
1.2 技术栈
1.3 学习路线
2. 硬件平台:MetaX C500 架构
2.1 MetaX ↔ NVIDIA ↔ AMD 术语对照
2.2 硬件层次结构
2.3 C500 关键参数速记
2.4 存储层次与数据局部性
C500 的存储子系统采用三层架构,类似认知科学的三层记忆模型:
关键特性:
2.5 指令发射单元
每周期最多发射 5 条指令(来自不同 warp):
3. 项目结构
4. OJ 提交评测系统
4.1 评测流程
OJ 平台支持三种语言提交:CUDA(C/C++ 源码)、Triton、TileLang。
单测试点执行步骤:
4.2 评分公式
S(Tk)=1+(s1−1)Tb−ThTk−Th100
4.3 TileLang 提交规范
run_kernel(无装饰器,按run_kernel(*args)调用)torch、tilelang(含tilelang.language、tilelang.intrinsics)、math@tilelang.jit包装 +@T.prim_func内层T.Tensor((M, K), T.float16)或T.Tensor[(M, K), T.float16]os、sys、subprocess、socket、pickle、exec/eval/compileTorchProxyError5. MoE Kernel:核心攻关项目
5.1 题目描述
实现 DeepSeek 风格 MoE 中的 pre-routed fused expert GEMM。
计算链条:
5.2 测试用例
重要:
routed_expert_weights为torch.float32(非 float16),需在 TileLang kernel 中单独声明 dtype。5.3 两阶段 Kernel 结构
设计上采用 2 次 launch(相对朴素 5 次 launch 的融合):
Stage 1:gate/up 双 GEMM + SiLU 融合
Stage 2:down GEMM + 权重缩放融合
5.4 当前最优参数
5.5 性能现状(cfg0)
5.6 职责划分:Host vs Kernel
MoEGate)5.7 融合的四个层面
x分块x只读一次;少一次 launchgate_logits中间张量5.8 瓶颈诊断(mcProfiler 证据)
cfg0 瓶颈:卡在 Stage 2(K=8192 的 64 轮长依赖链)
cfg1 瓶颈:卡在 Stage 1(K=7168 的 56 轮 + smem 锁死 1 block/AP)
关键发现:K 链长的一侧(16 轮 → 30K per-wave;56-64 轮 → 104-118K per-wave)构成”结构性必要等待”——所有并行度杠杆(split-K / 小 tile / t1024 / 双缓冲)均实测无效或负收益。
5.9 关于”融合”的澄清
“gate/up 产出 → down 消费”之间必须经全局内存——因为一个 128-token tile 的
up_logits是128×2048×2B=512KB,远超 C500 的 64 KiB WSM。这是所有主流实现都采用的结构。6. MLA / NSA 算子
6.1 FlashMLA(Multi-head Latent Attention)
文件:
race_tests/mla/test_tilelang_mla.pyTileLang 实现的 Flash Attention 风格 MLA kernel,支持 split 与 no-split 两种模式。
核心特性:
Q_pe/K_pe)num_split > 1)与 combine kernelT.use_swizzle(10)改善 L2 命中参考实现:
example/FlashMLA/(含 CUDA 源码与 PyTorch 接口)6.2 NSA(Native Sparse Attention)
文件:
race_tests/nsa/test_tilelang_nsa_fwd.py、race_tests/nsa/reference.py基于 block-sparse 的原生稀疏注意力机制。
核心特性:
S个 key block 计算注意力groups = HQ // H)naive_nsa参考实现对比验证当前限制:
7. 性能优化指南
7.1 优化策略速查表
__restrict__7.2 优化方法论(数据路径图方法)
分析任意 Kernel 时按以下三步:
典型思考框架:
7.3 C500 关键优化约束
7.4 优化案例:MoE Kernel 的经验教训
8. mcProfiler 性能分析工具
8.1 概述
mcProfiler 3.8.1.4 是 MetaX C500 配套的性能分析工具,基于 Client-Server 架构。
/usr/local/bin/mcProfiler/opt/mcProfiler-ubuntu18.04/profiler_server127.0.0.1:50123/opt/mcProfiler-ubuntu18.04/gui-profiler-0.1.0.AppImage/opt/mcProfiler-ubuntu18.04/config/*.pcd8.2 一键 Profiling 流程
8.3 关键口径修正(重要)
8.4 常用指标
8.5 已知问题
config/按 CWD 解析:必须在/opt/mcProfiler-ubuntu18.04下运行 mcProfilerperf done后约 5-6 分钟才生成report.txt.json--headnames900s+ timeout:不做无人值守采集--single-pass**:不稳定(出现过 108050% duty 异常)9. 工程文件清单
9.1 竞赛测试
race_tests/moe/custom_fusedmoe.pyrace_tests/moe/ref_fusedmoe.pyrace_tests/moe/fusedmoe_benchmark.pyrace_tests/moe/moe_test_configs.jsonrace_tests/moe/moe_prof_driver.pyrace_tests/moe/profile_mcprof.shrace_tests/moe/mcprof_summary.pyrace_tests/mla/test_tilelang_mla.pyrace_tests/nsa/test_tilelang_nsa_fwd.pyrace_tests/nsa/reference.py9.2 参考实现与算子库
example/FlashMLA/csrc/example/mcoplib/kernel/example/mcoplib/op/example/mcoplib/benchmark/example/kernels/9.3 性能分析工具与脚本
skills/mcprof-report-skill/SKILL.mdskills/mcprof-report-skill/metax-kernel-guide.mdskills/mcprof-report-skill/tests/conformance/doc/mxmaca-performance-tuning-guide/guide/doc/mxmaca-performance-tuning-guide/case/doc/mxmaca-performance-tuning-guide/microbenchmark/10. 开发计划与状态
10.1 McProfiler 报告分析技能状态
10.2 仍需验证的项目
--headnames900s timeout10.3 后续顺序
附录 A:MoE 优化 A/B 报告索引
run_cfg0_20260810_025910run_cfg0_20260810_032333run_cfg0_20260810_033517run_cfg0_20260810_041838run_cfg0_20260810_050659run_cfg1_20260810_055801附录 B:mcProfiler 一键命令参考
附录 C:TileLang 安装
安装脚本支持三级回退:
/opt/tilelang-metax-v0.1.10已编译 →.pth注册(秒级)pip install -e .(触发编译)git clone + install