[CPU] Fix registration of int4wo linear implementation on CPU

[Xia-Weiwen]

Summary
Int4wo on CPU does not run into expected mm op (torch.ops.aten._weight_int4pack_mm_for_cpu). It seems to be a regression after some refactoring of related code. This PR fixes it by registering a linear impl for the Int4CPULayout, which calls torch.ops.aten._weight_int4pack_mm_for_cpu for computation. The new impl is enabled for torch>=2.6. The new impl does not require dtype to be bfloat16. It supports fp32, fp16, bf16 for both weight and activation.

Test plan

python test/quantization/test_quant_api.py -k test_int4wo_cpu

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 6783a22 with merge base de5c6e1 ():
💚 Looks good so far! There are no failures yet. 💚

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

[leslie-fang-intel]

Since the model and inputs have already been converted to target data type, why we still need to enable autocast here?

[leslie-fang-intel]

why we change _weight_int4pack_mm_for_cpu to _weight_int4pack_mm? I remember _weight_int4pack_mm didn't register for CPU.

[sanchitintel]

In this PR, @Xia-Weiwen moved the CPU implementation to torchao/dtypes/uintx/int4_cpu_layout.py, where support for more activation dtypes is also being added.

This code now exclusively pertains to CUDA.

[sanchitintel]

Int4wo on CPU does not run into expected mm op (torch.ops.aten._weight_int4pack_mm_for_cpu)

The FX IR pattern you shared with me after running a toy model did have a call to torch.ops.aten._weight_int4pack_mm_for_cpu, but the whole pattern corresponding to it was very weird (even had a aten.mm call at the end), so something is indeed broken.
If we had run an LLM with torchchat with int4 WoQ & simply searched for torch.ops.aten._weight_int4pack_mm_for_cpu in the FX graph of the model, we may probably not have discovered this issue.

Is it possible to add a UT for a small model that uses torch.compile, and also somehow checks if the pattern is as expected?
Perhaps, the corresponding UT in test/integration/test_integration.py could be modified.

Thanks!