Add basic fp8/bfp8 support, future implementations of 8-bit floating …

[Geramy]

…point functions next.

Proposed changes

The proposed changes are to include fp8 and bfp8 support.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works

[zcbenz]

Rather than float8/bfloat8, I think it is more common using names like float8_e4m3/float8_e5m2, and there are also a few other float8 types (e.g. e8m0) that we should consider in advance.

JAX's float8 types: https://github.com/jax-ml/ml_dtypes
CUTLASS: https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/float8.h

I think we should have a consensus on the type names with the team first.

[Geramy]

Rather than float8/bfloat8, I think it is more common using names like float8_e4m3/float8_e5m2, and there are also a few other float8 types (e.g. e8m0) that we should consider in advance.

JAX's float8 types: https://github.com/jax-ml/ml_dtypes CUTLASS: https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/float8.h

I think we should have a consensus on the type names with the team first.

Sounds good, I actually think that is a lot better since we have to think about compatibility between all possible backends. It could turn into a big problem later I would assume. Let me know what you guys decide on I 100% agree on changing the names upfront.