Add support for fbgemm fp8 kernels

[jerryzh168]

Summary:
fp8 per row quantized weight with fp8 dynamic per row quantization only for now

	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
fp8dq-per-row - 1	95.95	0.0525	9.01 GB	7.51 GB
fp8dq-per-row - 128	94.29	0.0655	19.90 GB	7.51 GB
fbgemm-fp8 - 1 (no compile)	37.02	0.0486	16.76 GB	7.51 GB
fbgemm-fp8 - 128 (no compile)	11.7	0.0768	21.92 GB	7.51 GB

Test Plan:
python test/dtypes/test_fbgemm_fp8.py

in torchao/_models/llama folder:

export CHECKPOINT_PATH=../../../checkpoints # path to checkpoints folder 
export MODEL_REPO=meta-llama/Meta-Llama-3.1-8B-Instruct

python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --quantization fbgemm-fp8 --batch_size 1

compile doesn't work yet

batch size 1

Average overall tokens/sec: 37.02
Average decode tokens/sec: 37.4586 s
Average TTFT: 0.0486 s
Average tokens/sec: 37.02
Average Bandwidth: 278.05 GB/s
Peak Memory Usage: 16.76 GB
Model Size: 7.51 GB

batch size 128
Average overall tokens/sec: 11.70
Average decode tokens/sec: 11.7677 s
Average TTFT: 0.0768 s
Average tokens/sec: 11.70
Average tokens/sec including batches 1498.09
Average Bandwidth: 87.91 GB/s
Peak Memory Usage: 21.92 GB
Model Size: 7.51 GB

float8dq-row batch size 1, w/ compile
Average overall tokens/sec: 95.95
Average decode tokens/sec: 99.0108 s
Average TTFT: 0.0525 s
Average tokens/sec: 95.95
Average Bandwidth: 720.68 GB/s
Peak Memory Usage: 9.01 GB
Model Size: 7.51 GB

float8dq-row batch size 128, w/ compile
Average overall tokens/sec: 94.29
Average decode tokens/sec: 97.3500 s
Average TTFT: 0.0655 s
Average tokens/sec: 94.29
Average tokens/sec including batches 12069.68
Average Bandwidth: 708.26 GB/s
Peak Memory Usage: 19.90 GB
Model Size: 7.51 GB

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit ff66c46 with merge base d963a88 ():
💚 Looks good so far! There are no failures yet. 💚

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