From 1e30424eadfc37fb5df9ec331ffa5f6a6f6f6c59 Mon Sep 17 00:00:00 2001
From: Rachel Li <lirachel@mit.edu>
Date: Thu, 22 Jan 2026 18:46:08 -0500
Subject: [PATCH] Guard torch.compile on ROCm when triton_key is missing
 (#3923)

* Guard torch.compile on ROCm when triton_key missing

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update unsloth/import_fixes.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten ROCm Triton import handling

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Rachel Li <rachelliqx07@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 unsloth/__init__.py     |  3 +++
 unsloth/import_fixes.py | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index d3093cf4c0..0b819a546e 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -127,6 +127,7 @@ from .import_fixes import (
     fix_vllm_aimv2_issue,
     fix_vllm_guided_decoding_params,
     fix_vllm_pdl_blackwell,
+    fix_rocm_triton_key_error,
     ignore_logger_messages,
     patch_ipykernel_hf_xet,
     patch_trackio,
@@ -141,6 +142,7 @@ fix_xformers_performance_issue()
 fix_vllm_aimv2_issue()
 fix_vllm_guided_decoding_params()
 fix_vllm_pdl_blackwell()
+fix_rocm_triton_key_error()
 ignore_logger_messages()
 patch_ipykernel_hf_xet()
 patch_trackio()
@@ -154,6 +156,7 @@ del fix_xformers_performance_issue
 del fix_vllm_aimv2_issue
 del fix_vllm_guided_decoding_params
 del fix_vllm_pdl_blackwell
+del fix_rocm_triton_key_error
 del ignore_logger_messages
 del patch_ipykernel_hf_xet
 del patch_trackio
diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index 4f88808c2a..89fd152857 100644
--- a/unsloth/import_fixes.py
+++ b/unsloth/import_fixes.py
@@ -666,6 +666,39 @@ def fix_huggingface_hub():
         )
 
 
+def fix_rocm_triton_key_error():
+    """
+    ROCm + torch.compile can fail if Triton lacks `triton_key`.
+    Disable Inductor/compile only on ROCm when that symbol is missing.
+    """
+    try:
+        import torch
+    except (ImportError, ModuleNotFoundError):
+        return
+
+    if not getattr(torch.version, "hip", None):
+        return
+
+    try:
+        import triton
+    except (ImportError, ModuleNotFoundError):
+        return
+
+    try:
+        from triton.runtime import triton_key  # noqa: F401
+
+        return
+    except ImportError:
+        pass
+
+    os.environ.setdefault("TORCHINDUCTOR_DISABLE", "1")
+    os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
+    logger.info(
+        "Unsloth: ROCm detected and Triton lacks triton_key; "
+        "disabling torch.compile/Inductor to avoid backend crash."
+    )
+
+
 def fix_vllm_pdl_blackwell():
     """
     Fix vLLM PDL (Programmatic Dependent Launch) bug on Blackwell GPUs (SM100).