User-defined triton kernels fail with import error

[Chillee]

User-defined triton kernels fail with import error

[oulgen]

Do you have both pytorch-triton and regular triton installed?

[oulgen]

@Chillee and I looked at this, I am able to repro this. Although not exactly sure what the problem is. The problem is coming from https://github.com/openai/triton/blob/751ba68f364779f18b7135a759083da922586348/third_party/nvidia/backend/driver.py#L371
in the shape of nvi.CudaUtils but CudaUtils is defined in the same file. So it is not clear to me where this nvi is coming from. Triton might be doing some on the fly code generation here with exec and friends.

@jeromeku this is another instance of the problem you were seeing but now without the matmuls.

[oulgen]

I have a super minimal repro

import torch
import triton

@triton.jit
def kernel(X):
    return

def forward(x):
    kernel[(1, )](x, num_ctas=1)
    return x

torch.compile(forward)(torch.randn(1))

presence of num_ctas in this repro is causing the No module named 'nvi', not sure why yet.

[oulgen]

@Chillee in your repro, you use

num_warps=2, num_stages=1

so perhaps that's also triggering it.