Add support for fbgemm int4 mm kernel

[jerryzh168]

Summary:
we also plan to expose some other kernels like fp8xint4 and bf16xfp8, fp8xfp8 to compare with existing torchao kernels

Test Plan:
test/dtypes/test_fbgemm_int4_tensor.py

H100, with compile:

TODO: update

	overall tokens/sec	TTFT	Peak Memory	Model Size
baseline - 1	131.65	0.0220	16.24 GB	15.01 GB
baseline - 128	76.38	0.0544	26.92 GB	15.01 GB
int4wo - 1	207.69	0.0288	6.41 GB	3.99 GB
int4wo - 128	12.85	0.4223	16.01 GB	3.99 GB
fbgemm-int4 - 1 (no compile)	40.00	0.0508	29.03 GB	4.22 GB
fbgemm-int4 - 128 (no compile)	11.46	0.0846	28.96 GB	4.22 GB

export CHECKPOINT_PATH=../../../checkpoints # path to checkpoints folder
export MODEL_REPO=meta-llama/Meta-Llama-3.1-8B-Instruct
# default batch size 1
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --compile --write_result benchmark_results.txt
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --compile --quantization int4wo-128 --write_result benchmark_results.txt
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --quantization fbgemm-int4-128 --write_result benchmark_results.txt

python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --compile --write_result benchmark_results.txt --batch_size 128
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --compile --quantization int4wo-128 --write_result benchmark_results.txt --batch_size 128
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --quantization fbgemm-int4-128 --write_result benchmark_results.txt --batch_size 128

Note: fbgemm-int4-128 does not work with compile yet since the fbgemm op does not have meta device implementation.

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[samanamp]

Thank you! community really needs this.