Update ARM CPU experimental kernels from AO to leverage pip install

[metascroy]

• torchao experimental CPU kernels are now installed and loaded automatically by pip.
• Switch quantization to use new quantize_ API

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1458

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 67dd729 with merge base e5cf6e5 ():

NEW FAILURES - The following jobs have failed:

• pull / test-torchao-experimental-cpp (macos-14-xlarge) (gh)
  Process completed with exit code 134.
• pull / test-torchao-experimental-et (macos-14-xlarge) (gh)
  AttributeError: '_OpNamespace' 'torchao' object has no attribute '_pack_8bit_act_4bit_weight'

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

[Jack-Khuu]

Thanks for making the pip install work with the subclass APIs!!

[Jack-Khuu]

cc: @manuelcandales

Can have please for MPS? 🥺🥺 (Separate PR)

[Jack-Khuu]

Awaiting pytorch/executorch#7759

[metascroy]

@Jack-Khuu did you update the version AO uses in ET?

[Jack-Khuu]

Yup pytorch/executorch@9836b39 Points to pytorch/ao@11333ba

[nikhil-arm]

Hello @metascroy @Jack-Khuu , what is the plan to get this in mainline? We would like to use KleidiAI kernels from aten via this quantizer path. Let us know if we need to raise a new PR ?

[Jack-Khuu]

Hi @nikhil-arm, we're still planning to land this

Can you share the specific commit hashes y'all need?

[Jack-Khuu]

@nikhil-arm We've bumped the AO pin on main.
Please let me know if you there's any additional support needed to unblock KleidiAI kernels

[Jack-Khuu]

After a suite of rebases, pinbumps, and splitting up tests up we know what we're tackling:

• test-torchao-experimental-cpp (macos-14-xlarge): Tests the AOTI runner and likely failing (also in main) due to not linking to the LibOMP from torch as @malfet mentioned in Bump PT 2025131 and ET pins 20250209 #1493.

• test-torchao-experimental-et (macos-14-xlarge): Tests the ET runner; looks like a install bug where USE_CPP isn't set, but will likely run into the same LibOMP issue as above

[metascroy]

Hello @metascroy @Jack-Khuu , what is the plan to get this in mainline? We would like to use KleidiAI kernels from aten via this quantizer path. Let us know if we need to raise a new PR ?

Sorry about the delay @nikhil-arm.

@Jack-Khuu let's try to get this landed within the next week. Bumping the ao pin in torchchat had various conflicts with the CI, but I think we can dedicate to making this work.

I think it does make sense to first land pytorch/ao#1836 in torchao before bumping because they've already deprecated the old quantizers in quantize_.