pytorch 生成Triton代码之后流程详解

最新推荐文章于 2025-08-24 21:33:16 发布

MindShare AGI

最新推荐文章于 2025-08-24 21:33:16 发布

阅读量499

点赞数 20

CC 4.0 BY-SA版权

分类专栏： pytorch 大模型推理引擎文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/m0_49133355/article/details/150268410

Triton 代码是在运行过程中进行即时编译（Just-in-Time, JIT）的。 这也是 torch.compile 强大和灵活的关键所在。

下面我们来详细分解这个过程，看看它到底“何时”以及“如何”发生。

JIT 编译的完整流程（以第一次运行为例）

我们从用户代码开始：

# 1. 准备模型和输入
model = MyModel().cuda()
input_tensor = torch.randn(16, 3, 224, 224).cuda()

# 2. 应用 torch.compile
# 注意：这一步本身几乎不做什么事情！它只是用一个“代理”或“钩子”包装了你的模型。
# 编译是“懒惰的”（lazy），只在需要时才发生。
model_compiled = torch.compile(m