`triton` package that ships with torch is not compatible with `torch.ops.matmul` about pytorch

[oulgen]

"No module named 'nvi'" means that there's a version mismatch. Are you able to run your triton code in eager mode (without torch.compile)?

[jeromeku]

yes, in the minimal repro I provided, the code first runs the kernel without compile first. Only errors out when compiling.

[oulgen]

@jeromeku looking at your relevant libraries, i see you have both triton and pytorch-triton installed, could you pip uninstall triton so that there's only one triton?

[jeromeku]

I tried to do that, but uninstalling triton breaks my installation. When I try to import triton.ops, I get:

  "No module named 'triton.ops'"

[oulgen]

That's because ops is likely not defined in the __init__.py. If you add from triton import ops and then change triton.ops.matmul to ops.matmul, it makes progress, but then it runs into another issue in triton's autotuner.

import triton
import torch

from triton import ops

def test_triton_matmul_compile():
    dtype = torch.float16
    A = torch.randn(4096, 4096, dtype=dtype, device="cuda")
    B = torch.randn(4096, 4096, dtype=dtype, device="cuda")
    acc_dtype = None
    input_precision = "ieee"
    fp8_fast_accum = False
    output_dtype = None

    out = ops.matmul(A, B, acc_dtype, input_precision, fp8_fast_accum, output_dtype)

    #triton_compiled = torch.compile(ops.matmul)
    #compiled_out = triton_compiled(A, B, acc_dtype, input_precision, fp8_fast_accum, output_dtype)

if __name__ == "__main__":
    test_triton_matmul_compile()

notice that there's no torch compile anymore, still runs into the same autotuner issue. I suspect the triton package that ships with torch is not compatible with torch.ops.matmul.

I'll update the issue, for proper triage.

[arthursunbao]

I have the same issue and will update more details later