studio: skip tilelang on HIP / ROCm torch (Strix Halo crash report)

h34v3nzc0dex tested PR 5434 on Strix Halo (gfx1151, ROCm 7.13,
torch 2.11.0+rocm7.13.0) and hit a hard regression:

  File ".../fla/ops/common/backends/tilelang/__init__.py", line 92,
    in chunk_bwd_dqkwg
  File ".../tilelang/jit/kernel.py", line 137, in __init__
  File ".../tilelang/tileop/gemm/__init__.py", line 143,
    in _select_gemm_instruction
  tvm.error.InternalError: Check failed: (0) is false:
    Unsupported target for gemm:
    hip -keys=hip,gpu -mcpu=gfx1151 ...

`tilelang==0.1.8` ships no HIP GEMM instruction; `_select_gemm_instruction`
raises at lower-time, not import-time. So:
  - pip install succeeds
  - `import tilelang` succeeds
  - `TileLangBackend.is_available()` returns True
  - FLA's dispatcher picks TileLang for `chunk_bwd_dqkwg`
  - training subprocess dies at first GDN backward, no graceful fallback

The PR's existing platform gate (`_tilelang_platform_supported`)
checked only `sys.platform == "linux"` and `platform.machine()`, both
of which look identical on a ROCm box.

Fix has two layers:

1. INSTALL GATE: new `_torch_has_hip()` helper checks
   `torch.version.hip is not None`. `_tilelang_platform_supported`
   now returns False on HIP torch, so the install never fires.

2. RUNTIME GATE: even with the install skipped, a user could have
   tilelang already present (e.g. venv carried over from a CUDA box).
   `_install_fast_path_hooks` now calls
   `os.environ.setdefault("FLA_TILELANG", "0")` when HIP is detected,
   which is the env-var FLA's `TileLangBackend` already honors. Users
   who know they have a HIP-aware tilelang fork can override by
   setting `FLA_TILELANG=1` explicitly.

This costs nothing on CUDA (the gate is a no-op when
`torch.version.hip is None`), and removes the crash for AMD users.
The benchmark numbers in the PR description (1.43x on B200 sm_100)
are not affected.

The other halves of the PR are confirmed working on gfx1151 by the
same report:
  - `flash-linear-attention 0.5.0` runs at production scale
    (B=1 T=8192 H=16 K=128 V=128 and others) with no patches.
  - `causal-conv1d` runs at the shapes the fast-path gate cares
    about. (A separate Ubuntu 24.04 `--gcc-install-dir` build
    workaround is needed for the source-build path; that mirrors
    bbf004c's llama.cpp fix and is out of scope here.)

Tests added:
  - test_tilelang_platform_unsupported_on_hip_torch
  - test_tilelang_install_skipped_on_hip_torch
  - test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip
  - test_install_fast_path_hooks_respects_user_fla_tilelang_override
  - test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda

Total 50 passing (was 45).

studio/backend/core/training/worker.py

@@ -604,18 +604,49 @@ def _tilelang_importable() -> bool:
         return False
 
 
+def _torch_has_hip() -> bool:
+    """True iff the installed torch is a HIP / ROCm build.
+
+    We check `torch.version.hip` (non-None on ROCm wheels). This is the
+    reliable signal even on x86_64 Linux Strix Halo / MI300, where
+    `sys.platform` and `platform.machine()` look identical to a CUDA box.
+
+    Importing torch here is acceptable in the worker subprocess context:
+    the next step after kernel installers is the model load, which
+    imports torch anyway. We swallow import errors so a missing torch
+    (extremely unusual at this point) is treated as "not HIP" and the
+    rest of the gate stack handles it.
+    """
+    try:
+        import torch as _torch
+        return getattr(_torch.version, "hip", None) is not None
+    except Exception:
+        return False
+
+
 def _tilelang_platform_supported() -> bool:
-    """True iff the current platform has a tilelang 0.1.8 wheel.
+    """True iff the current platform has a usable tilelang 0.1.8 backend.
 
     tilelang publishes manylinux x86_64/aarch64 and macOS arm64 wheels
     plus a 93MB sdist; we never want the sdist on a Studio worker, so
     we restrict to Linux x86_64/aarch64 explicitly.
+
+    Excludes HIP / ROCm torch builds: tilelang 0.1.8 has no HIP GEMM
+    instruction, so `_select_gemm_instruction` raises `Unsupported
+    target for gemm: hip` mid-compile during Qwen3.5 GDN backward.
+    Reported by h34v3nzc0dex on Strix Halo (gfx1151, ROCm 7.13). The
+    pip wheel installs fine and imports cleanly, but FLA's TileLang
+    dispatcher then crashes at first training step. See PR 5434.
     """
     import platform as _platform
 
     if not sys.platform.startswith("linux"):
         return False
-    return _platform.machine().lower() in _TILELANG_SUPPORTED_LINUX_MACHINES
+    if _platform.machine().lower() not in _TILELANG_SUPPORTED_LINUX_MACHINES:
+        return False
+    if _torch_has_hip():
+        return False
+    return True
 
 
 def _pip_install_cmd(*args: str) -> list[str]:
@@ -854,6 +885,23 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
         logger.info("Fast-path hooks disabled via env; using substring fallback")
         return
 
+    # Defensive: on HIP/ROCm torch builds, FLA's TileLang backend (when
+    # tilelang is installed for any reason — e.g. a stale CUDA env that
+    # was reused for ROCm) crashes mid-backward with
+    # "Unsupported target for gemm: hip" inside
+    # `tilelang.tileop.gemm._select_gemm_instruction`. The install gate
+    # in `_ensure_tilelang_backend_unconditional` prevents NEW installs
+    # on HIP; this env-var setdefault disables FLA's TileLang dispatch
+    # for already-installed tilelang too. Users can override by setting
+    # FLA_TILELANG=1 explicitly. Reported by h34v3nzc0dex on Strix Halo.
+    if _torch_has_hip() and os.environ.get("FLA_TILELANG") is None:
+        os.environ["FLA_TILELANG"] = "0"
+        logger.info(
+            "HIP/ROCm torch detected; setting FLA_TILELANG=0 to keep "
+            "FLA on the safe Triton path (tilelang 0.1.8 has no HIP "
+            "GEMM backend)"
+        )
+
     try:
         from transformers.utils import import_utils as _iu
     except Exception as exc:

studio/backend/tests/test_training_worker_flash_attn.py

@@ -1221,3 +1221,91 @@ def test_run_training_process_eagerly_installs_causal_conv1d_in_normal_mode():
         "branch, so SSM models that bypass is_causal_conv1d_available() still "
         "get the eager install"
     )
+
+
+# ───────────────────────────────────────────────────────────────────
+# HIP / ROCm regression coverage (h34v3nzc0dex Strix Halo report).
+# tilelang 0.1.8 has no HIP GEMM backend; FLA's TileLang dispatch
+# crashes mid-backward on AMD with "Unsupported target for gemm: hip".
+# The fix: skip the install on HIP-built torch AND setdefault
+# FLA_TILELANG=0 so already-installed tilelang doesn't get used either.
+# ───────────────────────────────────────────────────────────────────
+
+
+def test_tilelang_platform_unsupported_on_hip_torch(monkeypatch):
+    """Strix Halo / MI300 with ROCm torch: linux + x86_64 looks
+    identical to a CUDA box at the OS level, so the platform check
+    must consult torch.version.hip explicitly.
+    """
+    monkeypatch.setattr(worker, "_torch_has_hip", lambda: True)
+    assert worker._tilelang_platform_supported() is False
+
+
+def test_tilelang_install_skipped_on_hip_torch(monkeypatch):
+    """End-to-end: the unconditional installer must not call pip on HIP torch."""
+    monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising=False)
+    monkeypatch.setattr(worker, "_torch_has_hip", lambda: True)
+    run_mock = mock.Mock(return_value=mock.Mock(returncode=0, stdout=""))
+    monkeypatch.setattr(worker._sp, "run", run_mock)
+    monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
+
+    result = worker._ensure_tilelang_backend_unconditional(event_queue=[])
+
+    assert result is False
+    run_mock.assert_not_called()
+
+
+def test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip(monkeypatch):
+    """When HIP torch is detected, hook installer must set
+    FLA_TILELANG=0 (via setdefault — respects user override) so any
+    PRE-EXISTING tilelang install isn't used by FLA's dispatcher.
+    """
+    import os as _os
+    monkeypatch.delenv("FLA_TILELANG", raising=False)
+    monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising=False)
+    monkeypatch.setattr(worker, "_torch_has_hip", lambda: True)
+    monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True)
+    monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True)
+    monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True)
+
+    worker._install_fast_path_hooks(
+        event_queue=_FakeQueue(), model_name="unsloth/Qwen3.5-2B"
+    )
+
+    assert _os.environ.get("FLA_TILELANG") == "0"
+
+
+def test_install_fast_path_hooks_respects_user_fla_tilelang_override(monkeypatch):
+    """If the user explicitly set FLA_TILELANG (even on HIP), don't
+    overwrite — they may know they have a HIP-aware tilelang fork.
+    """
+    import os as _os
+    monkeypatch.setenv("FLA_TILELANG", "1")
+    monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising=False)
+    monkeypatch.setattr(worker, "_torch_has_hip", lambda: True)
+    monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True)
+    monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True)
+    monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True)
+
+    worker._install_fast_path_hooks(
+        event_queue=_FakeQueue(), model_name="unsloth/Qwen3.5-2B"
+    )
+
+    assert _os.environ["FLA_TILELANG"] == "1"
+
+
+def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch):
+    """CUDA path must NOT set FLA_TILELANG (tilelang is wanted there)."""
+    import os as _os
+    monkeypatch.delenv("FLA_TILELANG", raising=False)
+    monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising=False)
+    monkeypatch.setattr(worker, "_torch_has_hip", lambda: False)
+    monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True)
+    monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True)
+    monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True)
+
+    worker._install_fast_path_hooks(
+        event_queue=_FakeQueue(), model_name="unsloth/Qwen3.5-2B"
+    )
+
+    assert _os.environ.get("FLA_TILELANG") is None