ao icon indicating copy to clipboard operation
ao copied to clipboard

[CPU] Enable DA8W4 on CPU

Open Xia-Weiwen opened this issue 8 months ago • 5 comments

Summary This PR enables DA8W4 on CPU.

  • It adds a new layout Int8DynamicActInt4WeightCPULayout and its implementation
  • It adds two custom ops: da8w4_linear_prepack_cpu for weight packing and da8w4_linear_cpu for DA8W4 GEMM.
  • It adds C++ kernels for the two new custom ops

The ops and kernels won't be available unless torchao is built from source with USE_CPP_KERNELS=1 on Linux only.

Test plan

pytest test/quantization/test_quant_api.py -k test_8da4w_cpu

Xia-Weiwen avatar Apr 25 '25 10:04 Xia-Weiwen

:link: Helpful Links

:test_tube: See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2128

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

:white_check_mark: No Failures

As of commit e3731f720f2dd7da50f6ef37bbdbb53895fa5b6b with merge base 4ebc9c042565e16af249b7cec8ebb2dc9fa0274f (image): :green_heart: Looks good so far! There are no failures yet. :green_heart:

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

pytorch-bot[bot] avatar Apr 25 '25 10:04 pytorch-bot[bot]

@leslie-fang-intel This PR is updated to use a new layout. Please review again. Thanks.

Xia-Weiwen avatar May 14 '25 10:05 Xia-Weiwen

Hi @jerryzh168 Could you please review this PR? Thanks.

Xia-Weiwen avatar May 16 '25 09:05 Xia-Weiwen

Hi @jerryzh168 Could you please review this PR? Thanks.

Xia-Weiwen avatar May 19 '25 02:05 Xia-Weiwen

Hi @jerryzh168 Could you please review this PR? Thanks.

Xia-Weiwen avatar May 20 '25 14:05 Xia-Weiwen

Hi @leslie-fang-intel Please review this PR again. I have also added the kernel code in this PR. It showed reasonable performance in internal benchmarks. Thanks.

Xia-Weiwen avatar Jun 04 '25 06:06 Xia-Weiwen

Please also describe how we choose different implementations based on the CPU Info.

I have added more details in the description. Thanks.

Xia-Weiwen avatar Jun 04 '25 15:06 Xia-Weiwen

Hi @jerryzh168 Could you please review this PR? Thanks. It's changed a lot since your last review.

Xia-Weiwen avatar Jun 06 '25 01:06 Xia-Weiwen

Hi @jerryzh168 Could you please review this PR? Thanks.

Xia-Weiwen avatar Jun 11 '25 03:06 Xia-Weiwen