ROCm OCP FP8 Support (#1677)

* document ROCm OCP F8 support

* lint

* Add AMD GPU model and gfx code documentation

Add a comment documenting supported AMD GPU models and their corresponding LLVM gfx codes, including Navi4, MI300X, and MI350.

* lint

* Refactor MI300 float8 dtype detection using utility function

Use is_MI300() utility function to simplify MI300 architecture detection for float8 dtypes

* lint

Author: petrex

torchao/float8/config.py

@@ -11,6 +11,8 @@
 
 import torch
 
+from torchao.utils import is_MI300
+
 logger: logging.Logger = logging.getLogger()
 
 
@@ -52,7 +54,7 @@ class Float8TypeConfig:
     """
     Configuration for selecting the preferred float8 type pair, either e4m3fn/e5m2 or e4m3fnuz/e5m2fnuz.
 
-    Currently, ROCm only supports fnuz variants.
+    Currently, ROCm supports 1. fnuz variants in MI300. 2. OCP F8 variants in MI350/Navi4.
     """
 
     # The preferred e4m3 type.
@@ -62,12 +64,9 @@ class Float8TypeConfig:
     e5m2_dtype = torch.float8_e5m2
 
     def __post_init__(self):
-        if torch.version.hip and torch.cuda.is_available():
-            prop = torch.cuda.get_device_properties(0)
-            MI300_ARCH = ("gfx940", "gfx941", "gfx942")
-            if prop.gcnArchName.split(":")[0] in MI300_ARCH:
-                self.e4m3_dtype = torch.float8_e4m3fnuz
-                self.e5m2_dtype = torch.float8_e5m2fnuz
+        if torch.version.hip and torch.cuda.is_available() and is_MI300():
+            self.e4m3_dtype = torch.float8_e4m3fnuz
+            self.e5m2_dtype = torch.float8_e5m2fnuz
 
 
 # User defined type for using the individual F8 type based on config

torchao/utils.py

@@ -606,6 +606,15 @@ def _torch_version_at_least(min_version):
     return is_fbcode() or version("torch") >= min_version
 
 
+# Supported AMD GPU Models and their LLVM gfx Codes:
+#
+# | AMD GPU Model | LLVM gfx Code          |
+# |---------------|------------------------|
+# | Navi4         | gfx1200, gfx1201       |
+# | MI300X        | gfx940, gfx941, gfx942 |
+# | MI350         | gfx950                 |
+
+
 def is_MI300():
     if torch.cuda.is_available() and torch.version.hip:
         mxArchName = ["gfx940", "gfx941", "gfx942"]
@@ -616,6 +625,22 @@ def is_MI300():
     return False
 
 
+def is_MI350():
+    if torch.cuda.is_available() and torch.version.hip:
+        archName = torch.cuda.get_device_properties(0).gcnArchName
+        if "gfx950" in archName:
+            return True
+    return False
+
+
+def is_Navi4():
+    if torch.cuda.is_available() and torch.version.hip:
+        archName = torch.cuda.get_device_properties(0).gcnArchName
+        if "gfx1200" or "gfx1201" in archName:
+            return True
+    return False
+
+
 def is_sm_at_least_89():
     return (
         torch.cuda.is_available()