b.py

import torch
from triton.testing import do_bench
import functools
from torch._inductor import config
from torch._dynamo.decorators import mark_dynamic
import os

@torch.compile
def f(x):
    return x.sum(dim=0)

N = 50
C = functools.partial(torch.randn, device="cuda")
x_large = C(4096, N)

f(x_large)
ms = do_bench(lambda: f(x_large))

# 0.01 v.s. 0.009
print(ms)