From 06345585fcd348b55169e2030ea7446fef45f4c9 Mon Sep 17 00:00:00 2001
From: Abhishek Enaguthi <enaguthiabhishek@gmail.com>
Date: Wed, 8 Jul 2026 18:38:38 -0700
Subject: [PATCH] feat: add AMD Instinct peak FLOPS for MFU on ROCm

get_peak_flops() fell back to A100 for MI250/MI300, skewing MFU on
ROCm hosts. Add published peak BF16/FP16 matrix FLOPS for MI100,
MI210, MI250, and MI300 with unit tests.

Complements #255 (ROCm support).
---
 src/zeroband/utils/__init__.py | 16 +++++++++++++++-
 tests/test_peak_flops.py       | 25 +++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_peak_flops.py

diff --git a/src/zeroband/utils/__init__.py b/src/zeroband/utils/__init__.py
index fafa9c7b..b8a93f1f 100644
--- a/src/zeroband/utils/__init__.py
+++ b/src/zeroband/utils/__init__.py
@@ -30,7 +30,7 @@ def get_sharding_strategy(sharding_strategy: str) -> ShardingStrategy:
 ### code above inspired and copied from https://github.com/pytorch/torchtitan/blob/4b3f2e41a084bf79a8540068ed525539d1244edd/torchtitan/utils.py#L119
 
 
-# hardcoded BF16 type peak flops for NVIDIA A100 and H100 GPU
+# hardcoded BF16 type peak flops for NVIDIA A100/H100 and AMD Instinct MI GPUs
 def get_peak_flops(device_name: str) -> int:
     if "A100" in device_name:
         # data from https://www.nvidia.com/en-us/data-center/a100/
@@ -44,6 +44,20 @@ def get_peak_flops(device_name: str) -> int:
             return 756e12
         else:  # for H100 SXM and other variants
             return 989e12
+    elif "MI300" in device_name:
+        # AMD Instinct MI300X/MI300A dense matrix peak (BF16/FP16)
+        # https://www.amd.com/en/products/accelerators/instinct/mi300/mi300x.html
+        return 1300e12
+    elif "MI250" in device_name:
+        # AMD Instinct MI250X dense matrix peak (BF16/FP16)
+        # https://www.amd.com/en/products/accelerators/instinct/mi250.html
+        return 383e12
+    elif "MI210" in device_name:
+        # AMD Instinct MI210 dense matrix peak (BF16/FP16)
+        return 181e12
+    elif "MI100" in device_name:
+        # AMD Instinct MI100 dense matrix peak (FP16)
+        return 184.6e12
     else:  # for other GPU types, assume A100
         return 312e12
 
diff --git a/tests/test_peak_flops.py b/tests/test_peak_flops.py
new file mode 100644
index 00000000..a33cd6d4
--- /dev/null
+++ b/tests/test_peak_flops.py
@@ -0,0 +1,25 @@
+import pytest
+
+from zeroband.utils import get_peak_flops
+
+
+@pytest.mark.parametrize(
+    ("device_name", "expected"),
+    [
+        ("NVIDIA A100-SXM4-80GB", 312e12),
+        ("NVIDIA H100 80GB HBM3", 989e12),
+        ("NVIDIA H100 PCIe", 756e12),
+        ("NVIDIA H100 NVL", 835e12),
+        ("AMD Instinct MI300X", 1300e12),
+        ("AMD Instinct MI300A", 1300e12),
+        ("AMD Instinct MI250X", 383e12),
+        ("AMD Instinct MI210", 181e12),
+        ("AMD Instinct MI100", 184.6e12),
+    ],
+)
+def test_get_peak_flops_known_devices(device_name: str, expected: float):
+    assert get_peak_flops(device_name) == expected
+
+
+def test_get_peak_flops_unknown_defaults_to_a100():
+    assert get_peak_flops("Some Unknown GPU") == 312e12