多范式算子开发实战营

模力方舟使用算力券租用沐曦C500的镜像选择：TileLang / 0.1.9 / Python 3.12 / maca 3.3.0.4

镜像选择

开始阶段先完成四个 Assignment，再进入下面的 Llama 算子练习项目。

任务	内容	目录
1	沐曦算力券、`mx-smi`	`assignment/task1`
2	TileLang Add 与 NineToothed Vector Add：Tile、尾块和性能测试	`assignment/task2`
3	TileLang Softmax 与 NineToothed GEMM：归约、数值稳定性和矩阵乘	`assignment/task3`
4	AI Agent 辅助算子开发、验证与优化	`assignment/task4`

cd assignment/task2
python -m pytest -q test_add.py test_ninetoothed_add.py
python benchmark_add.py
python benchmark_ninetoothed_add.py

cd ../task3
python -m pytest -q test_softmax.py test_ninetoothed_gemm.py
python benchmark_softmax.py
python benchmark_ninetoothed_gemm.py

完成四个 Assignment 后，继续 Llama 阶段：算子接入 → 正确性验证 → 性能优化 → 端到端评测。

Llama 算子练习

这是一个小型 Llama 推理项目，用来练习 PyTorch、NineToothed、TileLang 和 MXMACA 算子。

快速开始

安装依赖：

python -m pip install -r requirements.txt

使用 NineToothed 时单独安装可选依赖：

python -m pip install ninetoothed

通过 ModelScope 下载模型：

python -m pip install modelscope
modelscope download --model LLM-Research/Llama-3.2-1B \
  --local_dir models/Llama-3.2-1B

选择后端

--backend 决定使用哪份算子代码：

后端	用途
`torch`	参考实现和 CPU 基线
`tilelang`	TileLang kernel
`maca_cpp`	MXMACA `.maca` kernel
`ninetoothed`	NineToothed kernel

--target 可选 auto、cuda、maca。MACA 版 PyTorch 仍使用 --device cuda。如果后端、target、扩展或某个算子不可用，或者 kernel 运行出错，框架会打印警告并自动使用 PyTorch。推理最终输出的 registered_operators 会列出当前后端已接入的算子；值为 torch_fallback 表示该算子暂时仍调用 PyTorch reference。

NineToothed 测试：

python infer.py --model models/Llama-3.2-1B --prompts "Hello" \
  --max-new-tokens 1 --backend ninetoothed --target maca --device cuda

本仓库只保留 Llama 接入所需的薄包装和一个 RMSNorm kernel。安装和入门见 NineToothed 文档、 Add/Matmul 基础；更多 RMSNorm、RoPE、SDPA、MM 和 SwiGLU 示例见 ninetoothed-examples。

TileLang测试：

先加载预装的 TileLang 开发环境。脚本会自动发现 /app/tilelang-metax：

source ./setup_env.sh
# 自定义源码位置：TILELANG_ROOT=/path/to/tilelang-metax source ./setup_env.sh

python infer.py --model models/Llama-3.2-1B --prompts "Hello" \
  --max-new-tokens 1 --backend tilelang --target maca --device cuda

MXMACA测试：

python operators/maca_cpp/setup.py build_ext --inplace
python infer.py --model models/Llama-3.2-1B --prompts "Hello" \
  --max-new-tokens 1 --backend maca_cpp --target maca --device cuda

学员要改什么

框架已经预置 rms_norm 和 rope 的调用、注册和测试入口。当前 TileLang rms_norm 仅用于演示算子接入流程，并非性能最优实现；学员可以在此基础上继续优化线程布局、访存和归约方式。

优化范围不限于现有示例或文档推荐的某一个算子。学员可以根据自己的能力和目标，自行选择更多适合的模型算子进行实现和优化。RoPE 的接口是：

rope(input, sin_table, cos_table) -> output

学员只需要实现对应的 NineToothed、TileLang 或 MXMACA kernel，还没有实现的槽位会暂时使用 PyTorch reference 并打印警告，所以示例可以直接跑通；这种状态不能用于性能结论。完整说明见 operators/INTEGRATION.md。

RoPE 与 RMSNorm 实现说明

本文记录本次提交对两个算子的实现与优化，覆盖 NineToothed 和 TileLang 两个后端。

本次提交：RoPE 与 RMSNorm

以下章节记录本次提交对两个算子的实现与优化。

变更概览

算子	NineToothed	TileLang
`rope`	新实现（原为 PyTorch reference）	新实现（原为 PyTorch reference）
`rms_norm`	未改动	优化访存与数值精度

优化前 registered_operators 中 rope 在两个后端均为 torch_fallback，不计入后端性能。

RoPE

签名与 operators/torch_ops.py 一致：

rope(input, sin_table, cos_table) -> output

input 形状为 [batch, sequence, heads, head_dim]，两个 table 为 [sequence, head_dim // 2]。

实现要点

两个后端都把 head 切成两半，在 float32 寄存器里完成旋转后一次写回，因此 reference 里的 torch.cat 不再需要额外的 kernel 启动和数据往返。累积全程使用 float32，只在写回时舍入一次。

位置索引放在 Python wrapper 完成：把 [sequence, head_dim // 2] 展开为每行对应一个 (batch, sequence, head)。若放进 kernel 内部，需要 % sequence 形式的索引表达式，NineToothed 的 autotune 求解器无法为其求出上下界。

NineToothed 的三个约束

实现时踩到的框架行为，记录下来供后续算子参考：

application 的签名里不能出现 constexpr symbol。 NineToothed 从该签名生成 kernel 函数体，constexpr 会漏成未定义名。
constexpr symbol 只能在 arrangement 中引用。 在 application 中引用不会被重写。
tl.split / tl.join 的 tile 形状与表格对不齐。 这里的 ninetoothed/triton 组合下，split 后的形状无法与 sin/cos 广播对齐。最终把 head 的两半作为独立张量传入，split 上移到 wrapper，代价只是两个 strided view。

RMSNorm 优化（TileLang）

原实现分两遍扫描：先把每个 chunk 拷到 shared memory 累加平方和，再在归一化阶段从 global memory 重新读取同一批元素。对一个纯带宽受限的算子，这等于把读流量翻倍。

优化后整行只入 shared memory 一次并复用，kernel 每行只做一次全局读和一次全局写。同时把累加与缩放保持在 float32，写回时才舍入一次——原实现是在乘 weight 之前就舍入，白白损失 bfloat16 精度。

staging 数量限制为每线程 8 个元素。 把 2048 列全部放进 fragment 需要 5 KiB/thread，而 C500 驱动将私有内存上限设为 4 KiB/thread，行数较大时会以 mcErrorMemoryValueTooLarge 启动失败。这是实现过程中实际触发过的错误。

正确性验证

python -m pytest -q tests/

三个测试文件分别对比 operators.torch_ops 的 reference，覆盖真实 Llama-3.2-1B 前向过程中出现的形状，并断言算子以原生方式注册：

tests/test_rope_ninetoothed.py
tests/test_rope_tilelang.py
tests/test_rms_norm_tilelang.py

在 MetaX C500 上结果为 56 passed、3 skipped，三个后端输出的 token IDs 逐个一致。

注意：这些测试每个都会重新设置 backend。backends.configure_backend 设置的是进程级全局状态，而算子 wrapper 会读取它——不这样做的话，先运行的测试文件会导致后运行的文件报 “operator called without an active backend”。

性能测量

bench_ops.py 用批量 CUDA event 计时测单个算子：

python bench_ops.py --backend tilelang --repeats 20

单次 launch 加一次 host 同步的开销（约 0.1 ms）与这些算子本身的耗时同量级，直接测一次 launch 测到的是测量框架而不是 kernel，所以脚本每次采样连续发射多次再取平均。

参考量级（bf16，rows=4096、columns=2048）：

后端	耗时	有效带宽
ninetoothed	0.0707 ms	474 GB/s
tilelang	0.0885 ms	379 GB/s
torch	0.1451 ms	231 GB/s

关于端到端数字

infer.py 的端到端吞吐在这台共享机器上不稳定：同一份 torch 基线连跑三次分别为 101.90、83.57、74.60 tok/s，波动约 ±16%，来源是同机其他负载。仓库 README 也提示该配置「结果波动较大，不作为正式性能结论」。

若需要引用具体数字，请在独占环境重跑。

参考数据见 benchmarks/results/，由仓库自带的 compare_results.py 生成，其中 output_tokens_match 为 true。

未处理

MXMACA C++（.maca）后端未实现 rope，仍走 PyTorch reference。
attention 中的 repeat_interleave 和 SDPA 未做融合，仍是 PyTorch 路径。

测试和性能

下面提供轻量性能对比，统一生成 16 个 token，使用 1 次 warmup、3 次测量，并保持模型、 prompt、seed、精度和设备完全一致。该配置用于快速反馈，结果波动较大，不作为正式性能结论。下面命令假设在 MACA 机器上运行，模型目录是 models/Llama-3.2-1B。

先测 PyTorch 基线：

python infer.py \
  --model models/Llama-3.2-1B \
  --prompts "Hello" \
  --max-new-tokens 16 \
  --backend torch --target maca --device cuda \
  --num-warmup-iterations 1 \
  --num-profiling-iterations 3 \
  --seed 0 \
  --output-json benchmarks/results/torch_maca.json

再测 TileLang：

python infer.py \
  --model models/Llama-3.2-1B \
  --prompts "Hello" \
  --max-new-tokens 16 \
  --backend tilelang --target maca --device cuda \
  --num-warmup-iterations 1 \
  --num-profiling-iterations 3 \
  --seed 0 \
  --output-json benchmarks/results/tilelang_maca.json

NineToothed 使用同一组参数，运行以下命令并保存结果：

python -m pip install ninetoothed
python infer.py \
  --model models/Llama-3.2-1B \
  --prompts "Hello" \
  --max-new-tokens 16 \
  --backend ninetoothed --target maca --device cuda \
  --num-warmup-iterations 1 \
  --num-profiling-iterations 3 \
  --seed 0 \
  --output-json benchmarks/results/ninetoothed_maca.json

MXMACA 原生算子需要先构建扩展：

python operators/maca_cpp/setup.py build_ext --inplace
python infer.py \
  --model models/Llama-3.2-1B \
  --prompts "Hello" \
  --max-new-tokens 16 \
  --backend maca_cpp --target maca --device cuda \
  --num-warmup-iterations 1 \
  --num-profiling-iterations 3 \
  --seed 0 \
  --output-json benchmarks/results/maca_cpp_maca.json

比较 TileLang 和 PyTorch：

python benchmarks/compare_results.py \
  benchmarks/results/torch_maca.json \
  benchmarks/results/tilelang_maca.json \
  --output-json benchmarks/results/torch_vs_tilelang_maca.json

比较 NineToothed 和 PyTorch：

python benchmarks/compare_results.py \
  benchmarks/results/torch_maca.json \
  benchmarks/results/ninetoothed_maca.json \
  --output-json benchmarks/results/torch_vs_ninetoothed_maca.json

比较 MXMACA 原生实现和 PyTorch：

python benchmarks/compare_results.py \
  benchmarks/results/torch_maca.json \
  benchmarks/results/maca_cpp_maca.json \
  --output-json benchmarks/results/torch_vs_maca_cpp_maca.json

比较器会校验测试条件和生成的 token IDs，并输出吞吐 speedup、性能变化比例及实际替换的算子。 TileLang 调用包含 JIT 编译，不要直接拿第一次运行的时间评价性能。如果 RoPE 还在使用临时 PyTorch reference，也不能把该结果当作完整后端性能；应先实现对应 kernel。

.
├── assignment/                    # 四个入门实战
│   ├── README.md                  # Assignment 总览与通用要求
│   ├── task1/                     # 沐曦 GPU 与 MXMACA 环境
│   │   └── README.md
│   ├── task2/                     # TileLang Add 与 NineToothed Vector Add
│   │   ├── README.md
│   │   ├── solution.py            # TileLang kernel 作业入口
│   │   ├── ninetoothed_add.py     # NineToothed kernel 作业入口
│   │   ├── test_add.py            # TileLang 正确性测试
│   │   ├── test_ninetoothed_add.py # NineToothed 正确性测试
│   │   ├── benchmark_add.py       # TileLang 综合性能测试
│   │   └── benchmark_ninetoothed_add.py # NineToothed 性能测试
│   ├── task3/                     # TileLang Softmax 与 NineToothed GEMM
│   │   ├── README.md
│   │   ├── solution.py            # TileLang kernel 作业入口
│   │   ├── ninetoothed_gemm.py    # NineToothed kernel 作业入口
│   │   ├── test_softmax.py        # TileLang 正确性测试
│   │   ├── test_ninetoothed_gemm.py # NineToothed 正确性测试
│   │   ├── benchmark_softmax.py   # TileLang 综合性能测试
│   │   └── benchmark_ninetoothed_gemm.py # NineToothed 性能测试
│   └── task4/                     # AI Agent 辅助开发
│       ├── README.md
│       ├── prompts.md             # Prompt 记录模板
│       └── reflection.md          # 实践总结模板
│
├── infer.py                       # Llama 推理入口
├── llama.py                       # 模型与算子调用点
├── backends.py                    # 后端与 target 配置
├── setup_env.sh                   # TileLang/MXMACA 环境变量配置
├── requirements.txt               # Python 依赖
│
├── operators/                     # 算子实现
│   ├── registry.py                # 算子注册与分发
│   ├── torch_ops.py               # PyTorch 参考实现
│   ├── ninetoothed_ops.py         # NineToothed 薄包装
│   ├── ninetoothed_kernels/       # NineToothed 示例 kernel
│   ├── tilelang_ops.py            # TileLang 实现槽位
│   ├── maca_cpp/                  # MXMACA 原生扩展
│   └── INTEGRATION.md             # 算子接入教程
│
├── benchmarks/                    # 性能评测
│   ├── compare_results.py         # 结果比较工具
│   └── results/                   # 性能结果 JSON
│
└── tests/                         # 单元测试与后端集成测试