From f25d9b0a3dbf57620401f8bdb7c80e377a82032d Mon Sep 17 00:00:00 2001
From: Ryan Rock <ryan.rock@amd.com>
Date: Mon, 23 Mar 2026 20:41:52 +0000
Subject: [PATCH] [Build] Link PYTORCH_ROCM_ARCH specified archs

---
 setup.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/setup.py b/setup.py
index 96745f3..d270768 100644
--- a/setup.py
+++ b/setup.py
@@ -202,6 +202,8 @@
                 "-libverbs",
             ]
         )
+        arch_env = os.environ["PYTORCH_ROCM_ARCH"]
+        extra_link_args.extend([f"--offload-arch={arch}" for arch in arch_env.split(";")])
         if enable_mpi:
             extra_link_args.extend(
                 [