From a3a1c3457f75eb3ad8388ef59f41117be7a9d8a5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=87=91=E9=BB=84=E8=89=B2=E8=91=A1=E8=90=84=E7=90=83?=
 =?UTF-8?q?=E5=90=9B=E5=90=9B?= <yueyuan@amd.com>
Date: Sun, 1 Mar 2026 15:59:12 +0800
Subject: [PATCH] =?UTF-8?q?fix(ROCm):=20remove=20fix=5Frocm=5Ftriton=5Fkey?=
 =?UTF-8?q?=5Ferror=20=E2=80=94=20based=20on=20a=20false=20premise=20(#412?=
 =?UTF-8?q?5)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The function (introduced in #3923) assumed that the absence of
`triton.runtime.triton_key` on ROCm means torch.compile will crash.
Investigation shows this is incorrect:

1. `triton.runtime.triton_key` was renamed/removed in the ROCm Triton
   fork — it does not exist at that path.  However,
   `triton.compiler.compiler.triton_key` (the path torch._inductor
   actually imports) EXISTS and works correctly on ROCm.

2. Both call-sites in torch._inductor (codecache.py and
   async_compile.py) already wrap the import in try/except, so even a
   genuinely missing triton_key would be handled gracefully.

3. Comprehensive testing on ROCm 7.1 + Triton 3.4.0 + gfx1100 confirms
   torch.compile works correctly for matmul, cross-entropy, RMSNorm,
   multi-layer transformer forward+backward, and LoRA — all without
   triton.runtime.triton_key.

The original code was also ineffective (environment variables set after
torch import have no effect on torch._dynamo config), so removing it
has zero behavioral change on existing installations.

Supersedes the compile-disable portion of #3923.
---
 unsloth/__init__.py     |  3 ---
 unsloth/import_fixes.py | 33 ---------------------------------
 2 files changed, 36 deletions(-)

diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index 466fbf0628..60c5b3bba6 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -139,7 +139,6 @@ from .import_fixes import (
     fix_vllm_guided_decoding_params,
     fix_vllm_pdl_blackwell,
     fix_triton_compiled_kernel_missing_attrs,
-    fix_rocm_triton_key_error,
     patch_trunc_normal_precision_issue,
     ignore_logger_messages,
     patch_ipykernel_hf_xet,
@@ -161,7 +160,6 @@ check_vllm_torch_sm100_compatibility()
 fix_vllm_guided_decoding_params()
 fix_vllm_pdl_blackwell()
 fix_triton_compiled_kernel_missing_attrs()
-fix_rocm_triton_key_error()
 patch_trunc_normal_precision_issue()
 ignore_logger_messages()
 patch_ipykernel_hf_xet()
@@ -181,7 +179,6 @@ del check_vllm_torch_sm100_compatibility
 del fix_vllm_guided_decoding_params
 del fix_vllm_pdl_blackwell
 del fix_triton_compiled_kernel_missing_attrs
-del fix_rocm_triton_key_error
 del patch_trunc_normal_precision_issue
 del ignore_logger_messages
 del patch_ipykernel_hf_xet
diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index ee82715306..ebd81f9568 100644
--- a/unsloth/import_fixes.py
+++ b/unsloth/import_fixes.py
@@ -941,39 +941,6 @@ def fix_triton_compiled_kernel_missing_attrs():
     )
 
 
-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 patch_trunc_normal_precision_issue():
     """
     Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32.