Remove int_scaled_mm's dependency on triton for cpu

[Xia-Weiwen]

The int_scaled_mm op in Torchao is designed for CUDA at the beginning. However, this op is also needed for CPU. This PR adds a path for CPU in intmm.int_scaled_mm. It is not registered as and implementation for torchao.int_scaled_mm because we want to use Inductor for further optimization which cannot recognize the torchao.int_scaled_mm op.

This change requires this PR pytorch/pytorch#136942. Otherwise, there might be numerical issues. So, it works with pytorch nightly since 20241026.

Test is covered by test/kernel/test_autotuner.py and test/prototype/test_smoothquant.py.

[Xia-Weiwen]

Hi @cpuhrsch Could you please review and see if the changes are reasonable to you? Thanks.

[Xia-Weiwen]

Hi @cpuhrsch Could you please suggest how to deal with the issue (CPU impl availability depends on triton and AUTOTUNER_ENABLE)? Thanks!

[cpuhrsch]

Hey @Xia-Weiwen - Thank you for the PR! Sorry for the delay in review. Also, please note the CI hasn't run green.

Another way to resolve this could be to move

@torch.library.impl(lib, "int_scaled_matmul", "CPU")
def int_scaled_matmul_cpu(a, b, scales1):
    c = torch._int_mm(a, b)
    return c.to(scales1.dtype) * scales1

into torchao/kernel/intmm.py which shouldn't have a dependency on triton. Just be sure to also define lib = torch.library.Library("torchao", "FRAGMENT")

[Xia-Weiwen]

@cpuhrsch Thanks! I will give it a try. A question is what AUTOTUNER_ENABLE is and whether CPU impl should depend on it or not.

[cpuhrsch]

@Xia-Weiwen - it's used for a Triton autotuner that allows us to cycle over a very large number of configs for a given fixed input shape. See https://github.com/pytorch-labs/ao/tree/main/torchao/kernel#autotuner-and-custom-triton-kernels

[Xia-Weiwen]

Thank you @cpuhrsch. Looks like CPU impl does not need this.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/128

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 97dfea8 with merge base cbd90e3 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Xia-Weiwen]

Hi @cpuhrsch This PR requires latest torch nightly (after 20241026) to pass CI. May I know when torch nightly will be updated in the CI? Thanks.

[cpuhrsch]

@Xia-Weiwen - can you try merging the latest version of main? You might be built on top of a commit that pinned the nightly version. I see dev20241022 here: https://github.com/pytorch/ao/actions/runs/11550979799/job/32147052309?pr=128

[Xia-Weiwen]

@Xia-Weiwen - can you try merging the latest version of main? You might be built on top of a commit that pinned the nightly version. I see dev20241022 here: https://github.com/pytorch/ao/actions/runs/11550979799/job/32147052309?pr=128

Thanks

[Xia-Weiwen]

Hi @cpuhrsch CI is green. Could you please review? Thanks.