fix(ROCm): Comprehensive RDNA GPU support - fix Gemma3 NaN & add is_rdna() (#4109)

* fix(ROCm): comprehensive RDNA GPU support - fix Gemma3 NaN & add is_rdna()

- Add is_rdna() detection for RDNA3/3.5/RDNA4 consumer GPUs (gfx11xx, gfx1151, gfx12xx)
- Disable torch.compile for Gemma3 on HIP to fix NaN loss (fixes #3385, #4029)
- Export is_cdna/is_rdna from kernels for downstream use
- Import is_rdna into cross_entropy_loss for future RDNA-specific tuning

Tested on AMD Radeon PRO W7900 (gfx1100) with ROCm 7.1:
  ✓ Gemma3-1B: loss 3.37→3.25 (no NaN)
  ✓ Llama-3.2-1B: loss 2.44→2.37 (no NaN)
  ✓ Qwen2.5-1.5B: loss 1.89→1.85 (no NaN)
  ✓ RMS LayerNorm Triton kernel: bf16/fp16 PASSED
  ✓ Cross Entropy Loss Triton kernel: 32K/256K vocab PASSED

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

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

* Address review: scope compile disable to RDNA only, use partial mode, remove unused import

Changes based on Daniel's review:
1. (HIGH) Replace DEVICE_TYPE=='hip' with is_rdna() to avoid disabling
   torch.compile on CDNA GPUs (MI250X/MI300X/MI350) where it works fine
2. (MEDIUM) Use 'partial' instead of '1' for UNSLOTH_COMPILE_DISABLE to
   only disable model forward compilation while keeping loss compilation,
   matching the existing Sesame pattern
3. (LOW) Remove unused is_rdna import from cross_entropy_loss.py (F401)

* Remove redundant is_cdna/is_rdna exports from kernels/__init__.py

These functions are imported directly from .utils where needed
(e.g. cross_entropy_loss.py, loader.py). No external code imports
them from the unsloth.kernels namespace.

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

unsloth/kernels/utils.py

@@ -86,6 +86,15 @@ def is_cdna():
     )
 
 
+@functools.lru_cache(1)
+def is_rdna():
+    """Detect RDNA consumer/workstation GPUs (RDNA3, RDNA3.5, RDNA4)."""
+    if not is_hip():
+        return False
+    arch = triton.runtime.driver.active.get_current_target().arch
+    return arch.startswith("gfx1") and not is_cdna()
+
+
 def calculate_settings(
     n: int,
 ) -> (

unsloth/models/loader.py

@@ -1131,6 +1131,14 @@ class FastModel(FastBaseModel):
             # Set norms to float32 since anyways they get upcasted to float32
             # common in both gemma-3 and gemma-3n
             os.environ["UNSLOTH_HIGH_PRECISION_LAYERNORM"] = "1"
+            # ROCm/HIP: Gemma3 compiled forward produces NaN on RDNA GPUs
+            # (gfx1100, gfx1101, gfx1102, gfx1150, gfx1151, etc.).
+            # Disable torch.compile for model forward; loss compilation is fine.
+            # See https://github.com/unslothai/unsloth/issues/3385
+            from unsloth.kernels.utils import is_rdna
+
+            if is_rdna():
+                os.environ["UNSLOTH_COMPILE_DISABLE"] = "partial"
         # Cohere
         elif "cohere2" in model_types_all and transformers_version < Version(
             "4.50.0.dev0"