add separate quantization primitives for float8

[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 float8 quantization for inference.

Specifically:

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

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

[danielvegamyhre]

Stack from ghstack (oldest at bottom):

• 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/1597

• 📄 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 2838d50 with merge base 860da26 ():
💚 Looks good so far! There are no failures yet. 💚

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

[jainapurva]

@drisspg Does these look good?

[danielvegamyhre]

Alternatively I could just make these functions wrappers around the generic primitives which pass in the appropriate params for float8, as that may be more maintainable - although it does add a step of indirection and hides how the scale is actually computed. Any thoughts?

For example:

def choose_qparams_float8(input: torch.Tensor, float8_dtype: torch.dtype):
    scale, _ = choose_qparams_affine(
            input,
            MappingType.SYMMETRIC,
            input.shape,     # only tensorwise scaling is supported at the moment
            float8_dtype,
            eps=float8_eps,  # use same EPS as float8 training
            scale_dtype=torch.float32,
            quant_min=torch.finfo(float8_dtype).min,
            quant_max=torch.finfo(float8_dtype).max,
        )
    return scale