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.

❌ 1 New Failure, 1 Unrelated Failure

As of commit 9b7b602 with merge base a37f4e4 ():

NEW FAILURE - The following job has failed:

• Code Analysis with Ruff / build (3.9) (gh)

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests on ROCm / test-nightly (ROCM Nightly, linux.rocm.gpu.mi300.2, --pre torch --index-url https://download.pyto... / linux-job (gh) (trunk failure)
  test/dtypes/test_affine_quantized.py::TestAffineQuantized::test_print_quantized_module

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