[Float8] add non-decomposed version of quantize/dequantize ops for fp8

[shiyang-weng]

Fix #2896

What we want to do is to enable FP8 quantization in PyTorch. Similar to INT8 quantization, this requires inserting quantize and dequantize operations into the computational graph. In order to reuse pattern matching logic of int8, we need register FP8 quant and dequant.

To address this, we attempted to register quant in #2379, but the PR was reverted in #2672 because it caused performance regression on H100 GPUs. And there is no need to register q/dq on CUDA.

Based on the above reasons, I register quant specifically for CPU.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 85614a4 with merge base 18dbe87 ():
💚 Looks good so far! There are no failures yet. 💚

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

[jerryzh168]

seems OK to me, wondering if @vkuzo has additional thoughts, not sure if there is a better alternative here to support preserving ops for cpu

[shiyang-weng]

@vkuzo Could you help review this PR?

[jerryzh168]

@vkuzo Could you help review this PR?

is this urgent? Vasiliy is not available recently and will be back next week

[shiyang-weng]

is this urgent? Vasiliy is not available recently and will be back next week

Thanks for letting me know. Not urgent. We can wait for him back next week

[shiyang-weng]

@vkuzo Could you help review this PR?

[vkuzo]

we should keep cuda and cpu logic consistent, device is supposed to be orthogonal to quantization workflows

I'd recommend a flag named around something like AOTI or pt2e (cc @jerryzh168 for the right name) to control whether you want to decompose the quant/dequant or not decompose them.

[Xia-Weiwen]

Hi @vkuzo Thanks for your suggestion. May I know what kind of flag you were talking about? A global flag, or an argument passed to quantization APIs?

[Xia-Weiwen]

Hi @vkuzo Thanks for your suggestion. May I know what kind of flag you were talking about? A global flag, or an argument passed to quantization APIs?

Hi @vkuzo @jerryzh168 Could you share a little more about the design? Thanks.

[jerryzh168]

I think the decision of decompose or not decompose should be static? if we want consistent behavior for the same op across cuda and cpu, it might be better to have separate ops I feel

[Xia-Weiwen]

Hi @jerryzh168 Would you think it better to have a non-decomposed and a decomposed version of the op than a CPU and a CUDA version? We did a similar thing here: https://github.com/pytorch/pytorch/blob/df4ebddbe0fa2306fb8acd09b20265046d968c10/torch/ao/quantization/fx/_decomposed.py#L1206
also @vkuzo

[jerryzh168]

yeah just a different op seems to be the only alternative here

[shiyang-weng]

Created quantize_affine_float8_non_decomposed and dequantize_affine_float8_non_decomposed separately for non-decomposed

[jerryzh168]

LGTM, cc @vkuzo can you take a look again

[jerryzh168]

looks good to me