TVM是50万行C++，PyTorch堆着Dynamo、Inductor和Triton。现在有人用约5000行纯Python写了编译器，直接输出原生CUDA。
- 纯Python写编译器：约5000行代码，告别C++巨兽，易读易改
- 支持主流小模型：能把TinyLlama和Qwen2.5-7B直接转成CUDA内核
- 极致算力优化：6层IR降级，自动融合算子、消除中间显存占用
为什么我们要关注底层编译器
接活做AIGC，算力成本是大头。想省显存提速，就得看底层。但现在的编译器栈太折磨了。PyTorch那一套Dynamo、Inductor、Triton层层套娃，还有XLA、MLIR、Halide、Mojo，根本没教程讲清楚高层设计，直接把你扔进源码里。出了问题没法调，只能干瞪眼加钱买卡。
这哥们从零写了约5000行纯Python的参考编译器。它把模型通过6个IR（中间表示）一步步降级成原生CUDA。目标不是干翻Triton，而是做一个容易跟进、随便改的编译器。对我们来说，这意味着能真正看懂并控制模型是怎么跑在显卡上的，从哪省显存，从哪挤速度。
6层IR降级到底干了什么
拿torch.relu(torch.matmul(x + bias, w))举例，看它怎么把代码一步步压进显卡。
Torch IR：抓取FX图，和PyTorch算子一一对应，原汁原味。
Tensor IR：把算子拆成Elementwise、Reduction、IndexMap。统一算子面，未来ONNX、JAX前端接入不用动下游。那个(16, 64, 16)的中间变量看着吓人，但下一阶段直接融合掉，永远不会真正生成。
Loop IR：算子融合发生地。广播乘法、归约、输出布局、relu收尾，全压进一个循环嵌套，中间缓冲区全部干掉。
Tile IR：开始感知GPU。循环轴分配给线程/块，Stage把共享输入提入共享内存，2x2寄存器分块让一个线程算4个输出。K轴切成2次32宽的外部迭代。三大优化：buffers=2@a2双缓冲让加载和计算重叠；async用cp.async.ca.shared.global不阻塞warp；pad=(0,1,0)加1个元素填充消除bank冲突。
Kernel IR：变成硬件原语。THREAD/BLOCK变threadIdx/blockIdx，异步Stage变Smem加cp.async填充和commit/wait围栏，同步Stage变跨步填充循环。框架无关，也能降级到Metal或HIP。
CUDA：一对一转换，读x、bias、w各一次，写relu一次，全在一个内核完成，零中间变量。
自己动手改编译器
项目名叫deplodock，代码量极小，非常适合想搞懂编译器原理的创作者自己拉下来改。想给TinyLlama或Qwen2.5-7B写专属优化，或者研究怎么把显存占用压到极致，这是最好的切入点。
详细文章：
https://medium.com/data-science-collective/a-principled-ml-compiler-stack-in-5-000-lines-of-python-17f2db9549d4
代码仓库：
https://github.com/cloudrift-ai/deplodock
留言聊聊
你平时本地跑图，最头疼哪个环节的显存爆炸？
来源：Reddit MachineLearning｜原文：A Hackable ML Compiler Stack in 5,000 Lines of Python
