Support MXFP6 packing and fused unpack-dequantise kernel

[alex-titterton]

Good afternoon! Following recent developments and increased support for MXFP formats, it would be useful to support efficient packing for MXFP6 to benefit from the decrease in memory consumption and bandwidth requirements vs (MX)FP8.

MXFP6 has shown to perform similarly well compared with MXFP8 in LLM inference tasks, and with sufficient QAT even as well as float32, e.g. in the MXFP paper.

This PR packs the bits representing the FP6 values in a 4+2 fashion as is done in the FP6 LLM paper, and supports both E2M3 and E3M2 variants. Packing is done via standalone Triton kernel, with unpacking and dequantisation performed via a fused kernel for better performance.

Tests have been added in test_custom_cast.py and test_mx_tensor.py to cover accuracy in quantise-pack-unpack-dequantise with various FP6 values (min/max norm, min/max subnorm, -0.0 etc for both E2M3 and E3M2 variants) as well as checking packed tensor dimensions.

Note: due to the 4+2 packing scheme this requires the packing dimension to be a multiple of 4 since the packed dimension will be 3/4 of this. However the typical MX block size is 32 (—> 24 when packed), and HW implementations tend to require dims to be multiples of 16 or 32, so this should not be a problem. The relevant test case dimensions have been changed from 6 to 8 and the MX block sizes from 2 to 4 where applicable in order to accommodate this requirement.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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