ROCm mx-fp8 Gemm

[petrex]

TLDR: This pull request introduces support for AMD MI355x GPUs with HIPBLASLT kernels in the MX formats prototype. Note that this feature requires ROCm 6.5+ and gfx950

alongside several updates to improve compatibility and functionality for these GPUs. Key changes include updates to configuration options, validation logic, and GEMM kernel handling to integrate HIPBLASLT support.

AMD MI355x GPU Support:

• torchao/prototype/mx_formats/config.py:

  • Added HIPBLASLT as a new MXGemmKernelChoice and included it in the MXLinearRecipeName for configuration presets. [1] [2]
  • Updated _validate_gemm_kernel_choice to include validation logic for HIPBLASLT, ensuring proper block size, data type, and ROCm availability.

• torchao/prototype/mx_formats/mx_ops.py:

  • Extended mx_mm to support HIPBLASLT for scaled matrix multiplication and real GEMM operations. [1] [2]
  • Adjusted error messaging for unsupported kernel choices in FP4 operations.

Documentation Updates:

• torchao/prototype/mx_formats/README.md:
  • Updated the README to reflect AMD MI355x GPU support, including instructions for using HIPBLASLT kernels and ongoing optimization efforts for AMD hardware. [1] [2] [3]

Minor Code Refinements:

• torchao/prototype/mx_formats/mx_ops.py:
  • Improved readability in mx_view_op by reformatting conditions for FP6 element packing.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 129a6d6 with merge base 801af03 ():
💚 Looks good so far! There are no failures yet. 💚

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