Add exhaustive config option to intmm kernel

[jerryzh168]

Summary:
similar to pytorch/pytorch#126220 we added exhaustive option for int8mm and scaled_mm kernels in torchao

Note that there seems to be native int8mm and scaled_mm support in pytorch: https://github.com/pytorch/pytorch/blob/0610b9730e27d066e26396a2d655ba0d98c2012d/torch/_inductor/kernel/mm.py#L305 for int8mm and https://github.com/pytorch/pytorch/blob/0610b9730e27d066e26396a2d655ba0d98c2012d/torch/_inductor/kernel/mm_scaled.py#L575 for scaled mm maybe we should use that at some point. We can do this later as there are slight differences right now.

Also added all options for autoquant.

Test Plan:

cd benchmarks
TORCHAO_AUTOTUNER_ENABLE=1 python intmm.py --file_path intmm_shapes.csv
TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_SEARCH_SPACE=EXHAUSTIVE TORCHAO_AUTOTUNER_ENABLE=1 python intmm.py --file_path intmm_shapes.csv

ran autoquant-all and TORCHAO_AUTOTUNER_ENABLE=1 on sam:

+cuda,vit_h,32,13646,16,21.97087307115583,45.51480484008789,0.5813941740455395,max-autotune,torch.bfloat16,None,False,True,True,32,154,4928,None,None
+cuda,vit_h,32,52727,65,22.11766063236711,45.2127382105047,0.5811011656545951,max-autotune,torch.bfloat16,autoquant-all,False,True,True,32,154,4928,None,None
+cuda,vit_h,32,52671,64,22.292370883752643,44.858395960423856,0.5626328873979952,max-autotune,torch.bfloat16,autoquant-all,False,True,True,32,154,4928,None,None

TORCHAO_AUTOTUNER_ENABLE=1 give a slight boost.
exhaustive mode is taking too long so it probably makes more sense to use it with https://github.com/pytorch/ao/tree/main/torchao/prototype/quantization/mixed_precision

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 29d6032 with merge base 8a805d0 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/integration/test_integration.py::TestSubclass::test_int8_dynamic_quant_subclass_api_5_cuda

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