c.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=-1)

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

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

# 0.037ms v.s. 0.025ms
print(ms)