Add Float8QuantizedTensor (AQT subclass) and replace to_affine_quantized_floatx with to_affine_quantized_float8 in quantization APIs

[danielvegamyhre]

Context

Currently, AQT has the method from_hp_to_floatx for float8 quantization, and from_hp_to_fpx for low precision floating point data types like fp6 (technically can support fp1-fp7).

from_hp_to_floatx re-uses from_hp_to_intx, which in turn uses these generic quantization primitives.

Overall, in the current state the float8 path is a bit confusing for developers, due to both the naming ("floatx") and the use of generic functions which include a bunch of params which are unrelated to float8 quantization.

Summary of changes

The goal of this PR stack is to refactor this to have a clean separation of concerns, and simpler internal API surfaces for code using in float8 quantization for inference.

Specifically:

• Separate quantization primitives for float8
• Integrate those new quant primitives into AQT
• Integrate new AQT methods into float8 quantization APIs <------------------- (this PR)

Note: I will add float8 static quantization in a separate set of PRs.

[danielvegamyhre]

Stack from ghstack (oldest at bottom):

• -> Add Float8QuantizedTensor (AQT subclass) and replace to_affine_quantized_floatx with to_affine_quantized_float8 in quantization APIs #1599
• integrate new float8 quantization primitives into AQT #1598
• add separate quantization primitives for float8 #1597

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 2f15cc1 with merge base 32d9b0b ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  RuntimeError: Command docker exec -t 84ed1b6238c85976efd407f2dccfb3e367a35c6459cd1613a0be45cb65fb0ba0 /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t 10b1e1a3e3ec936f8a6e94f598ea3637e9c886a117fc6a5bb06ba4438414fdb1 /exec failed with exit code 1

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

[jerryzh168]

thanks, we also want to split out a Float8 (and floatx) specific AQT implementations as well, I talked to @jainapurva before

[danielvegamyhre]

thanks, we also want to split out a Float8 (and floatx) specific AQT implementations as well, I talked to @jainapurva before

Yep that makes sense, when I talked to her earlier she said she is planning to create these AQT subclasses, so I decided to do this part of the refactor.

[jainapurva]

Please remove floatx, float8 should replace floatx.

[danielvegamyhre]

Oh I left it in since it's still in use in other parts of the code base (autoquant, autoquant v2), and I wasn't sure if I should be touching those - is it ok to replace all instances across the whole codebase?

[jainapurva]

thanks, we also want to split out a Float8 (and floatx) specific AQT implementations as well, I talked to @jainapurva before

Yep that makes sense, when I talked to her earlier she said she is planning to create these AQT subclasses, so I decided to do this part of the refactor.

Yes, we want all the instances replaced. Autoquant is using it for Float8. Hence would be better to rename it float8

[danielvegamyhre]

thanks, we also want to split out a Float8 (and floatx) specific AQT implementations as well, I talked to @jainapurva before

Yep that makes sense, when I talked to her earlier she said she is planning to create these AQT subclasses, so I decided to do this part of the refactor.

Yes, we want all the instances replaced. Autoquant is using it for Float8. Hence would be better to rename it float8

Done!

[vkuzo]

I'm not a fan of this, this introduces one more abstraction (Float8QuantizedTensor), while keeping the complexity of AffineQuantizedTensor. I think either staying with AQT or just writing a float8 tensor without using AQT would seem more attractive.

[danielvegamyhre]

Interesting - cc @jainapurva @jerryzh168 thoughts on this?

For context AQT subclassing was part of a BE effort for the week, I'll share the doc with you internally

[jainapurva]

Removing AQT abstraction is easy, but the only reason I felt like keeping it was consistency in all dtypes. Though I do agree that it adds another level of abstraction

[danielvegamyhre]

Discussed offline, closing until internal discussions are finalized.