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.

❌ 10 New Failures

As of commit 4ee1b96 with merge base 2766a95 ():

NEW FAILURES - The following jobs have failed:

• pull / test-gpu-aoti-bfloat16 (cuda, stories15M) / linux-job (gh)
  RuntimeError: Command docker exec -t 14af4e47aa0b71f95b564b3c6208b7c8088259f1543a164fd19d4eeb38a0ed1b /exec failed with exit code 1
• pull / test-gpu-aoti-float16 (cuda, stories15M) / linux-job (gh)
  RuntimeError: Command docker exec -t 18799958a831eaab7384987a6665e64d2639f089c2224b9f0ddb1528bdd9818d /exec failed with exit code 1
• pull / test-gpu-aoti-float32 (cuda, stories15M) / linux-job (gh)
  RuntimeError: Command docker exec -t 37537919830b4df74f15e11a92f9584ef376358dc7845c54d818892f4cab84d9 /exec failed with exit code 1
• pull / test-gpu-compile (cuda, stories15M) / linux-job (gh)
  RuntimeError: Command docker exec -t a3e1466bc90dfcb3308c3dac084b3e47ffe5c57219adf3d563ab11009cd9c666 /exec failed with exit code 1
• pull / test-gpu-eval-sanity-check (cuda, stories15M) / linux-job (gh)
  RuntimeError: Command docker exec -t 908eda4f8983d58c7052923caded0ba63fecbe35ad908612530a0dcb5b7a19d7 /exec failed with exit code 1
• pull / test-torchao-experimental-cpp (macos-14-xlarge) (gh)
  AttributeError: '_OpNamespace' 'torchao' object has no attribute '_pack_8bit_act_4bit_weight'
• pull / test-torchao-experimental-mps (macos-m1-stable) (gh)
  Process completed with exit code 134.
• Run parallel prefill / test-cuda / linux-job (gh)
  RuntimeError: Command docker exec -t 4a2fc24ab338c14b6a25a0568bbf7749fb7417c267722a90dd3dde718e2d8381 /exec failed with exit code 1
• Run parallel prefill / test-sdpa-backends-export / linux-job (gh)
  RuntimeError: Command docker exec -t 11be5904b0a7fbd249b6d822b2bf2774e03348266cc9fcf0eb2bedb2204d3759 /exec failed with exit code 1
• Run the aoti runner with CUDA using stories / test-runner-aot-cuda / linux-job (gh)
  RuntimeError: Command docker exec -t f78b73896a8fb2bd312a22fd3a288e0e9d1505d9a9c2da692265bc9adf58de8a /exec failed with exit code 1

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]

Is there a specific case we are anticipating here?

[metascroy]

If using this on cuda machine of x86 machine

[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?