Skip to content

README_CN.md

Latest commit

 

History

History
182 lines (147 loc) · 9.62 KB

README_CN.md

File metadata and controls

182 lines (147 loc) · 9.62 KB
Raw

中文 | English

MiniTriton

MiniTriton 是一个教学级、但按生产标准做的 tile 编译器：用 Python 内嵌 的 tile DSL 写 kernel，经 MLIR 方言逐级 lower 到 PTX，在 NVIDIA GPU 上 运行；上面再叠一个 eager 张量库（eager 执行和 tile 级编译两种模式， 同一套 API）。所有应用 kernel——flash attention、KDA（Kimi Delta Attention）等——都是用 DSL 自己写的；编译器里只有 tile 词汇原语和 通用 pass，没有面向应用的 intrinsic。设计决策与实测历史见 git 提交 记录和 issue 页。

本项目由 Kimi K3（Moonshot AI）开发： DSL 前端、MLIR 编译器、CUDA kernel、autograd/nn、benchmark、图表与 文档全部由 Moonshot AI 的 K3 模型设计、实现、测量并撰写，维护者 负责工程方向与评审。

声明：本项目是 K3 在编译器设计上的演示（demonstration），不是 Moonshot AI 的官方项目。

稳定复现请用 stable 分支（与 main 同步推进）。下面的每张图都能 用 benchmarks/ 下的脚本重跑出来；数字一律以脚本和它们的 CSV 为准。

概览

CUDA-core roofline (fp32) Tensor-core roofline (tf32/bf16) train_gpt 收敛（默认 fp32 路径）
cuda-core tensor-core convergence

roofline 怎么看：横轴是算术强度（每字节 DRAM 流量对应的浮点运算数）， 纵轴是实测算力（GFLOP/s）。实线是本机实测的性能上限（同机标定），虚线 是厂商标称值。红 = minitriton，蓝 = torch（eager 或 cuBLAS），绿 = torch.compile，紫 = triton。不同实现、同一形状的点会叠在一起；基线系列 画得更大更浅，叠上的点会在红色 minitriton 点周围透出一圈淡色。落在对手 屋顶下面的点我们也照常画出。第三张图是两个栈上同一个小 GPT 的默认 fp32 训练曲线，基本重合。

快速开始

依赖：NVIDIA GPU（sm_89 实测；sm_80+ 理论可行）、系统 CUDA ≥ 12 （ptxas）、Python ≥ 3.10。工具链经 conda-forge 获取：

conda create -p ./.conda-env -c conda-forge python=3.10 mlir-python-bindings mlir
./.conda-env/bin/pip install cuda-python numpy pytest        # torch/triton 仅 benchmark 对照用
PY=./.conda-env/bin/python

$PY examples/vecadd.py                  # DSL → MLIR → PTX → GPU 端到端
$PY -m pytest tests/ -q                 # 1619 passed（数值 oracle 全 numpy/fp64，不 import torch）

也可以作为 git 包安装（MLIR 工具链走 pip 的 mlir-wheels 索引——轮子 ~1.3GB，网络慢的话建议走上一条 conda-forge 路线，同一个 22.1.0 工具链）：

pip install "mlir>=22" -f https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest
pip install cuda-python numpy
pip install "git+https://dev.msh.team/wangzhiyuan/minitriton.git"
python -c "import minitriton; print(minitriton.__version__)"

示例与用法

日常任务一条命令一个（torch/triton 只是对照基线，绝不是运行时依赖）：

# 1. 写 kernel：@tl.kernel DSL → MLIR → PTX（examples/vecadd.py、matmul.py）
$PY examples/vecadd.py

# 2. 融合 eager 代码：@tl.compile + 融合规则引擎（R1-R8/R2' 规则）
$PY -m pytest tests/test_fusion.py -q          # 20+ 个可运行的融合示例

# 3. 端到端训练模型（语料 = 本仓源码，零外部数据）
CUDA_VISIBLE_DEVICES=0 $PY examples/train_gpt.py --steps 200   # ~50M char-GPT

# 4. 混合精度：bf16 训练路径
#    （bf16 参数 + fp32 主权重 + 动态 loss scaling，CUDA graph 步进）
$PY benchmarks/training/train_gpt_precision_bench.py --steps 60 --impls tl_bf16_graph

# 5. 流：独立 kernel 上侧流重叠（opt-in）
MT_STREAM_OVERLAP=1 $PY benchmarks/training/train_gpt_precision_bench.py --steps 60 --impls tl_bf16

# 6. benchmark + CI 闸门（核心套件对 torch eager 与 torch.compile）
CUDA_VISIBLE_DEVICES=0 $PY benchmarks/ops/ce_report.py
CUDA_VISIBLE_DEVICES=0 $PY benchmarks/ci_report.py   # rc=0 即 PASS

# 7. 调度助手：规范化流水骨架与组合行算子参考形，在 minitriton/sched/
#    （直接 import 复用）

分布式（NCCL 数据并行，每步一次梯度 allreduce）：

CUDA_VISIBLE_DEVICES=0,1 $PY examples/train_gpt.py --ddp --steps 120 --loss-csv /tmp/ddp.csv

单卡 vs 双卡 DDP loss 曲线——逐位一致（120 步 diff 0.0000；benchmarks/training/plot_ddp_overlay.py）

图形/物理（minitriton.viz，窗口或无头 PNG 帧）：

$PY examples/rigid_demo.py --steps 240 --frames /tmp/rigid_frames   # 刚体粒子
$PY examples/sph_demo.py   --steps 240 --frames /tmp/sph_frames     # SPH 流体
$PY examples/mpm_demo.py   --stats 50                               # MPM 连续介质（控制台统计）

SPH 溃坝，5 万粒子：按密度场着色的水体（深青底、青色波面、压密处泛 白），480 步 @ 48 fps；examples/render_gallery.py 可重生成。 静态接触表 · PDF（打印用）。

限制（Limitations）

  • 硬件：仅在 NVIDIA L20（sm_89）验证；sm_80+ 理论可行但未实测；无 AMD/ROCm；CPU 后端是测试用的 numpy oracle，不是可用运行时。
  • 数据类型：fp32 与 bf16 是一等路径，tf32 为 opt-in；无 fp16、 fp8、量化路径。默认路径与 torch 不逐位一致（归约顺序不同、 ex2.approx）。
  • 算子覆盖：算子集即 benchmark 与示例所及——不是 torch 替代品 （无 conv/pool/循环类；多数算子要求连续输入；新的 shape/stride 组合会触发重新编译 + autotune）。
  • 性能上限如实登记：部分 kernel 族低于 torch/cuBLAS（见 roofline 图；证据链在 git 记录和 issue 里）。冷启动编译进程内串行（测试闸门 用 pytest-xdist 在进程级并行）。
  • 分布式：仅 NCCL 数据并行，实测 2 卡单机。sparse/ 与 distributed/ 是最年轻的包。
  • 生态与成熟度：API 是 torch-风格而非 torch-兼容；无 HuggingFace/ONNX 桥。教学级但测量诚实——每个数字都有可重跑脚本； 已知缺口记录在 git 提交记录和 issue 里。
  • 输入校验刻意从简：embedding 与 cross_entropy 的索引不做边界 检查（越界索引是未定义行为，反向散写可能写出界）；张量对张量的 pow 在负底数时返回 NaN（整数标量指数如 x ** 3 正常）。
  • 没有中段优化器：DSL 直接 lower 到手工排布的 MLIR（软件流水、 smem swizzle、ldmatrix/mma 编排和约 25 处 inline PTX 热点都按 kernel 族手写在 compiler/builder.py 里），其上只有三个通用 W-IR 小 pass。本项目是教学级手工排布 tile DSL，不是 optimizing compiler—— 新增一类 kernel 要手写几百行调度，而不是新写一个 pass。

仓库布局

├── minitriton/          # 包本体
│   ├── frontend/        # @tl.kernel AST 前端（types/code_generator/semantics）
│   ├── compiler/        # builder（tile→MLIR） layout（布局代数）
│   │                    #   wir/passes（通用 pass） lowering（→PTX→cubin）
│   ├── fusion.py        # 图级融合规则引擎（注册表 + 仲裁）
│   ├── compile.py       # @tl.compile（trace + 执行器）
│   ├── sched/           # 调度编排骨架 + 行算子参考形
│   ├── runtime/         # cuda_driver / buffer / cache / autotune / allocator / streams
│   ├── device/
│   │   ├── cuda/        # ops.py（路由壳） elementwise.py（生成器 + 表达式词表）
│   │   │   └── kernels/ # 唯一 kernel 库（matmul/attention/norm/ce/misc/kda，纯 DSL）
│   │   └── cpu/         # numpy oracle
│   ├── ops/             # eager op 共享语义（广播/提升/分发）
│   ├── nn/  autograd/  distributed/  sparse/  precision.py（精度开关）
├── examples/            # vecadd → matmul → train_gpt（~50M char-GPT，语料自包含）
├── benchmarks/
│   ├── ci_report.py     # CI 性能闸门（geomean 阈值 + 非零退出码）
│   ├── roofline/        # plot_style（统一风格）+ 双 roofline + solve_tril
│   ├── matmul/  attention/  kda/  ops/  training/  probes/
├── tests/               # pytest（数值 oracle 全 numpy/fp64，不 import torch）
├── docs/                # logo（architecture_map.py 是开发期 IR 地图）
└── build/               # 编译中间产物（gitignore，本地调试用）

工程纪律

  • 一切性能对照必带图 + CSV + 基线实现（torch eager / torch.compile / triton，同口径）；图统一走 benchmarks/roofline/plot_style.py 风格族（(dark, light) × (png, svg, pdf)；light 版 PDF 供打印/LaTeX）
  • 负结果与未达标项记录在案，不遮掩（git 提交记录 + issue 页）
  • DSL 优先：应用 kernel 用 Python DSL 书写、经通用管线 lower；编译器 intrinsic 仅限 tile 词汇集 + 逐案论证的表达能力例外 （AGENTS.md §1 三分法）
  • 每个 kernel 必须有 numpy/fp64 oracle 测试； benchmarks/training/grad_check.py（对 torch autograd，atol 1e-4） 不许回退