Overview

PyTorch 2.3, released on June 14, 2024, improves torch.export for model export and integrates custom Triton kernels into torch.compile.

Main Features

Improved torch.export

torch.export produces a clean IR graph for deployment, without Python runtime dependency.

python 
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

model = MyModel()
example = (torch.randn(1, 10),)
exported = torch.export.export(model, example)
print(exported.module()(torch.randn(1, 10)))

Triton kernels

Custom Triton kernels can now be integrated into the torch.compile pipeline, enabling optimization of specific GPU operations.

python 
import torch

# torch.compile integrates Triton kernels
@torch.compile
def my_function(x, y):
    return torch.matmul(x, y) + torch.relu(x.sum(dim=-1, keepdim=True))

a = torch.randn(32, 64, device='cuda')
b = torch.randn(64, 128, device='cuda')
result = my_function(a, b)
print(result.shape)

Sources