diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 048aeeafdc..f47a6bd599 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -52,6 +52,23 @@ _MAMBA_SSM_RELEASE_TAG = "v2.3.1" _MAMBA_SSM_PACKAGE_VERSION = "2.3.1" _FLASH_ATTN_RUNTIME_MIN_SEQ_LEN = 32768 _FLASH_ATTN_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL" +# apache-tvm-ffi 0.1.10/0.1.11 crash Triton with "CUDA: misaligned address" on sm_100. +_TILELANG_PACKAGE_VERSION = "0.1.8" +_APACHE_TVM_FFI_PACKAGE_VERSION = "0.1.9" +_TILELANG_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_TILELANG_INSTALL" +# Pin both so plain pip cannot silently upgrade torch under the worker (fla-core needs torch>=2.7). +_FLA_PACKAGE_VERSION = "0.5.0" +_FLA_CORE_PACKAGE_VERSION = "0.5.0" +_FLA_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FLA_INSTALL" +# `--no-deps` saves torch but loses fla-core's transitive deps; `packaging` is also undeclared upstream. +_FLA_RUNTIME_DEPS = ("einops", "packaging", "triton") +_FLA_MIN_TORCH = (2, 7) +_FLA_MIN_PYTHON = (3, 10) +# tilelang 0.1.8 ships wheels only for these Linux arches and macOS arm64; never fall back to its 93MB sdist. +_TILELANG_SUPPORTED_LINUX_MACHINES = frozenset(("x86_64", "amd64", "aarch64", "arm64")) +_TILELANG_INSTALL_TIMEOUT_S = 600 +_TVM_FFI_BROKEN_VERSIONS = ("0.1.10", "0.1.11") +_FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS" def _model_wants_causal_conv1d(model_name: str) -> bool: @@ -145,7 +162,7 @@ def _install_package_wheel_first( if wheel_url is None: logger.info("No compatible %s wheel candidate", display_name) elif url_exists(wheel_url): - _send_status(event_queue, f"Installing prebuilt {display_name} wheel...") + _send_status(event_queue, f"Installing {display_name} for faster training...") for installer, result in install_wheel( wheel_url, python_executable = sys.executable, @@ -187,7 +204,9 @@ def _install_package_wheel_first( "(this may take several minutes)..." ) else: - pypi_status_message = f"Installing {display_name} from PyPI..." + pypi_status_message = ( + f"Installing {display_name} from PyPI for faster training..." + ) _send_status(event_queue, pypi_status_message) @@ -331,6 +350,168 @@ def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None: ) +def _installed_torch_version_tuple() -> tuple[int, int] | None: + """Return ``(major, minor)`` of the installed torch, else None.""" + try: + from importlib.metadata import version as _pkg_version + + raw = _pkg_version("torch").split("+", 1)[0] + parts = raw.split(".") + return (int(parts[0]), int(parts[1])) + except Exception: + return None + + +def _flash_linear_attention_importable() -> bool: + """Catch any exception (not just ImportError) so a broken native lib doesn't abort the worker.""" + try: + import fla.modules # noqa: F401 + import fla.ops.gated_delta_rule # noqa: F401 + + return True + except Exception as exc: + logger.warning( + "flash-linear-attention is not importable; continuing with install/fallback: %s", + exc, + ) + return False + + +def _flash_linear_attention_current(already_importable: bool | None = None) -> bool: + """True iff FLA imports AND is at the pinned version (older FLA lacks gated_delta_rule kernels).""" + if already_importable is None: + already_importable = _flash_linear_attention_importable() + if not already_importable: + return False + try: + from importlib.metadata import version as _pkg_version + from packaging.version import Version + + fla_v = Version(_pkg_version("flash-linear-attention")) + core_v = Version(_pkg_version("fla-core")) + return fla_v >= Version(_FLA_PACKAGE_VERSION) and core_v >= Version( + _FLA_CORE_PACKAGE_VERSION + ) + except Exception as exc: + logger.warning( + "flash-linear-attention importable but version check failed; treating as stale: %s", + exc, + ) + return False + + +def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool: + """Install pinned FLA + fla-core with --no-deps. Returns True iff importable post-call.""" + if os.getenv(_FLA_SKIP_ENV) == "1": + return False + if sys.version_info < _FLA_MIN_PYTHON: + logger.info( + "Skipping flash-linear-attention install: requires Python >= %d.%d, have %s", + _FLA_MIN_PYTHON[0], + _FLA_MIN_PYTHON[1], + sys.version.split()[0], + ) + return False + torch_ver = _installed_torch_version_tuple() + if torch_ver is not None and torch_ver < _FLA_MIN_TORCH: + _send_status( + event_queue, + ( + f"Skipping flash-linear-attention install: fla-core requires " + f"torch>={_FLA_MIN_TORCH[0]}.{_FLA_MIN_TORCH[1]}, have " + f"{torch_ver[0]}.{torch_ver[1]}" + ), + ) + return False + + # Probe once; reuse result so the --force-reinstall decision and the short-circuit + # share the same call count (stable for tests). + already_importable = _flash_linear_attention_importable() + if already_importable and _flash_linear_attention_current(already_importable = True): + logger.info("flash-linear-attention already importable at the pinned version") + return True + + _send_status( + event_queue, + f"Installing flash-linear-attention=={_FLA_PACKAGE_VERSION} for faster training...", + ) + + # `--no-deps` blocks the silent torch upgrade; we bring the non-torch runtime deps in by hand. + specs = [ + *_FLA_RUNTIME_DEPS, + f"fla-core=={_FLA_CORE_PACKAGE_VERSION}", + f"flash-linear-attention=={_FLA_PACKAGE_VERSION}", + ] + extra_args = ["--no-deps"] + if already_importable: + # Older FLA already imported; pip skips reinstall without this flag. + extra_args.append("--force-reinstall") + + if shutil.which("uv"): + pypi_cmd = [ + "uv", + "pip", + "install", + "--python", + sys.executable, + *extra_args, + *specs, + ] + else: + pypi_cmd = [ + sys.executable, + "-m", + "pip", + "install", + *extra_args, + *specs, + ] + + try: + result = _sp.run( + pypi_cmd, + stdout = _sp.PIPE, + stderr = _sp.STDOUT, + text = True, + timeout = _TILELANG_INSTALL_TIMEOUT_S, + ) + except _sp.TimeoutExpired: + logger.warning("flash-linear-attention install timed out; continuing") + _send_status( + event_queue, "flash-linear-attention install timed out; continuing" + ) + return False + + if result.returncode != 0: + logger.warning( + "flash-linear-attention install failed (continuing on torch fallback):\n%s", + result.stdout, + ) + _send_status( + event_queue, + "flash-linear-attention install failed; continuing without it", + ) + return False + + # pip can exit 0 with a missing transitive runtime dep; verify the import. + if not _flash_linear_attention_importable(): + _send_status( + event_queue, + "flash-linear-attention installed but is not importable; continuing without it", + ) + return False + + logger.info("Installed flash-linear-attention for the FLA fast path") + return True + + +def _ensure_flash_linear_attention(event_queue: Any, model_name: str) -> None: + """Legacy model-name-gated FLA install, used when UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1.""" + if not _model_wants_tilelang(model_name): + return + _ensure_flash_linear_attention_unconditional(event_queue) + + _SSM_MODEL_SUBSTRINGS = ( "nemotron_h", "nemotron-h", @@ -359,6 +540,382 @@ def _ensure_mamba_ssm(event_queue: Any, model_name: str) -> None: ) +# Auto-derived from installed transformers: model_types whose modeling_*.py imports `from fla.*`. +# Cached per process. Empty when transformers can't be inspected -> we skip tilelang pre-install +# (the FLA Triton path still runs via the runtime hook). +_TRANSFORMERS_FLA_MODEL_TYPES_CACHE: frozenset[str] | None = None +_MODEL_NAME_SEP_CHARS = ("-", ".", "/", " ") + + +def _discover_fla_model_types() -> frozenset[str]: + """Model_types in the installed transformers whose modeling file imports `from fla.*`.""" + global _TRANSFORMERS_FLA_MODEL_TYPES_CACHE + if _TRANSFORMERS_FLA_MODEL_TYPES_CACHE is not None: + return _TRANSFORMERS_FLA_MODEL_TYPES_CACHE + found: set[str] = set() + try: + import transformers + + models_root = Path(transformers.__file__).parent / "models" + for modeling in models_root.glob("*/modeling_*.py"): + try: + src = modeling.read_text(encoding = "utf-8", errors = "ignore") + except OSError: + continue + if "from fla." in src: + found.add(modeling.parent.name) + except Exception as exc: + logger.debug("FLA model-type discovery skipped: %s", exc) + _TRANSFORMERS_FLA_MODEL_TYPES_CACHE = frozenset(found) + return _TRANSFORMERS_FLA_MODEL_TYPES_CACHE + + +def _model_wants_tilelang(model_name: str) -> bool: + """True iff model_name normalizes to contain a discovered FLA model_type.""" + types = _discover_fla_model_types() + if not types: + return False + name = model_name.lower() + for sep in _MODEL_NAME_SEP_CHARS: + name = name.replace(sep, "_") + return any(t in name for t in types) + + +def _installed_tvm_ffi_version() -> str | None: + """Installed apache-tvm-ffi version, or None if missing/unimportable.""" + try: + from importlib.metadata import version as _pkg_version + + return _pkg_version("apache-tvm-ffi") + except Exception: + return None + + +def _tilelang_importable() -> bool: + """Catch any exception (not just ImportError) so a broken native lib doesn't abort the worker.""" + try: + import tilelang # noqa: F401 + import tvm_ffi # noqa: F401 + + return True + except Exception as exc: + logger.warning( + "tilelang/tvm_ffi is not importable; continuing with install/fallback: %s", + exc, + ) + return False + + +def _torch_has_hip() -> bool: + """True iff torch is a ROCm build; `torch.version.hip` is the only reliable signal on x86_64 ROCm.""" + 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 a tilelang 0.1.8 wheel will load: Linux x86_64/aarch64, non-HIP torch. + + HIP excluded because tilelang 0.1.8 has no HIP GEMM instruction and crashes mid-backward. + """ + import platform as _platform + + if not sys.platform.startswith("linux"): + return False + 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]: + """`uv pip install` if uv is on PATH, else `python -m pip install`.""" + if shutil.which("uv"): + return ["uv", "pip", "install", "--python", sys.executable, *args] + return [sys.executable, "-m", "pip", "install", *args] + + +def _run_pip(cmd: list[str], event_queue: Any, label: str) -> bool: + """Run a pip install and surface success/failure via status events.""" + try: + result = _sp.run( + cmd, + stdout = _sp.PIPE, + stderr = _sp.STDOUT, + text = True, + timeout = _TILELANG_INSTALL_TIMEOUT_S, + ) + except _sp.TimeoutExpired: + logger.warning("%s install timed out; continuing", label) + _send_status(event_queue, f"{label} install timed out; continuing") + return False + if result.returncode != 0: + logger.warning( + "%s install failed (continuing without it):\n%s", label, result.stdout + ) + _send_status(event_queue, f"{label} install failed; continuing") + return False + return True + + +def _ensure_tilelang_backend_unconditional(event_queue: Any) -> bool: + """Install pinned tilelang + apache-tvm-ffi; two-step repair if a broken tvm-ffi is present. + + Returns True iff both import post-call. Step 1 surgically downgrades a broken tvm-ffi + with --force-reinstall --no-deps so torch / CUDA stay untouched; step 2 is a regular + install for missing transitive deps. Bypass via UNSLOTH_STUDIO_SKIP_TILELANG_INSTALL=1. + """ + if os.getenv(_TILELANG_SKIP_ENV) == "1": + return False + if sys.version_info < _FLA_MIN_PYTHON: + logger.info( + "Skipping tilelang install: requires Python >= %d.%d, have %s", + _FLA_MIN_PYTHON[0], + _FLA_MIN_PYTHON[1], + sys.version.split()[0], + ) + return False + if not _tilelang_platform_supported(): + import platform as _platform + + logger.info( + "Skipping tilelang install: no prebuilt wheel for %s/%s", + sys.platform, + _platform.machine(), + ) + return False + + existing_tvm_ffi = _installed_tvm_ffi_version() + needs_repair = existing_tvm_ffi in _TVM_FFI_BROKEN_VERSIONS + + if not needs_repair and _tilelang_importable(): + logger.info("tilelang + apache-tvm-ffi already installed") + return True + + # Step 1: --no-deps keeps --force-reinstall from touching torch/CUDA via the dep graph. + if needs_repair: + logger.info( + "Forcing apache-tvm-ffi downgrade: %s is on the broken list", + existing_tvm_ffi, + ) + _send_status( + event_queue, + ( + f"Downgrading apache-tvm-ffi {existing_tvm_ffi} -> " + f"{_APACHE_TVM_FFI_PACKAGE_VERSION} (broken-versions list)" + ), + ) + repair_cmd = _pip_install_cmd( + "--only-binary=:all:", + "--force-reinstall", + "--no-deps", + f"apache-tvm-ffi=={_APACHE_TVM_FFI_PACKAGE_VERSION}", + ) + if not _run_pip(repair_cmd, event_queue, "TileLang backend repair"): + return False + + # Step 2: regular install pulls in transitive deps (z3-solver, ml-dtypes) without touching torch. + _send_status( + event_queue, + f"Installing TileLang=={_TILELANG_PACKAGE_VERSION} for faster training...", + ) + install_cmd = _pip_install_cmd( + "--only-binary=:all:", + f"apache-tvm-ffi=={_APACHE_TVM_FFI_PACKAGE_VERSION}", + f"tilelang=={_TILELANG_PACKAGE_VERSION}", + ) + if not _run_pip(install_cmd, event_queue, "TileLang backend"): + return False + + # pip can exit 0 while a native lib (libz3.so) is missing; verify the import. + if not _tilelang_importable(): + _send_status( + event_queue, + "TileLang backend installed but is not importable; continuing on the FLA Triton path", + ) + return False + + logger.info("Installed TileLang backend for FLA fast path") + return True + + +def _ensure_tilelang_backend(event_queue: Any, model_name: str) -> None: + """Legacy substring-gated tilelang installer (opt-out path).""" + if not _model_wants_tilelang(model_name): + return + _ensure_tilelang_backend_unconditional(event_queue) + + +# ── Fast-path hooks ── +# Wrap transformers' is_{flash_linear_attention,causal_conv1d}_available so the first call +# (at modeling import time) drives the install. Any model that queries the gate gets the +# install; models that never query it (Llama, Gemma, dense Qwen) pay nothing. +# UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 falls back to the legacy substring path. + + +def _rebind_in_already_imported_modules( + *, attr_name: str, old_obj: Any, new_obj: Any +) -> int: + """Rebind `attr_name -> new_obj` in every module that already imported `old_obj`. + + `from X import Y` creates a local binding that reassigning X.Y won't reach. + Uses `__dict__.get` (not `getattr`) to skip lazy `__getattr__` aliases. + """ + count = 0 + missing = object() + for mod_name, mod in list(sys.modules.items()): + if mod is None: + continue + module_dict = getattr(mod, "__dict__", None) + if not isinstance(module_dict, dict): + continue + existing = module_dict.get(attr_name, missing) + if existing is old_obj: + try: + setattr(mod, attr_name, new_obj) + count += 1 + except Exception as exc: + logger.debug("Could not rebind %s in %s: %s", attr_name, mod_name, exc) + return count + + +def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: + """Hook transformers' is_*_available gates so the first call drives the install. + + Idempotent. UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 falls back to the substring gate. + """ + if os.getenv(_FAST_PATH_HOOKS_SKIP_ENV) == "1": + logger.info("Fast-path hooks disabled via env; using substring fallback") + return + + # On HIP torch, even already-installed tilelang crashes FLA's TileLang dispatch. + # User can override with FLA_TILELANG=1. + 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 (no HIP GEMM in tilelang 0.1.8)" + ) + + try: + from transformers.utils import import_utils as _iu + except Exception as exc: + logger.warning( + "transformers.utils.import_utils not importable; skipping fast-path hooks: %s", + exc, + ) + return + + def _make_wrapper( + original: Callable[[], bool], + install_fn: Callable[[Any], bool], + gate_name: str, + post_available_fn: Callable[[Any], None] | None = None, + ) -> Callable[[], bool]: + state = {"installed": False} + + def wrapper() -> bool: + if state["installed"]: + return original() + try: + original.cache_clear() # defensive; worker subprocess is fresh + except AttributeError: + pass + ok = original() + ran_install = False + if not ok: + ran_install = True + logger.info("Hook fired for %s; triggering install", gate_name) + try: + ok = bool(install_fn(event_queue)) + except Exception as exc: + logger.warning( + "%s install raised: %s; falling back to torch", gate_name, exc + ) + ok = False + logger.info("%s hook done; available=%s", gate_name, ok) + # post_available_fn handles "gate already True but ancillary kernel broken" (e.g. tilelang + # missing while FLA imports fine); skip when install_fn already chained the follow-up. + if ok and not ran_install and post_available_fn is not None: + try: + post_available_fn(event_queue) + except Exception as exc: + logger.warning( + "%s post-available step raised: %s; continuing", gate_name, exc + ) + state["installed"] = True + return ok + + wrapper.__wrapped__ = original # type: ignore[attr-defined] + wrapper.cache_clear = getattr(original, "cache_clear", lambda: None) # type: ignore[attr-defined] + return wrapper + + def _fla_install(eq: Any) -> bool: + # FLA alone ~2.35x; +tilelang adds ~26%. tilelang is GDN-only (Qwen3.5 family). + if not _ensure_flash_linear_attention_unconditional(eq): + logger.info( + "FLA install did not produce an importable runtime; skipping TileLang" + ) + return False + if _model_wants_tilelang(model_name): + _ensure_tilelang_backend_unconditional(eq) + else: + logger.info( + "Model %r outside TileLang allowlist; FLA Triton path is sufficient", + model_name, + ) + return True + + def _fla_post_available(eq: Any) -> None: + # FLA already imports; repair tilelang if missing or on the broken tvm-ffi list. + if not _model_wants_tilelang(model_name): + return + if ( + _installed_tvm_ffi_version() not in _TVM_FFI_BROKEN_VERSIONS + and _tilelang_importable() + ): + return + _ensure_tilelang_backend_unconditional(eq) + + def _causal_conv1d_install(eq: Any) -> bool: + ok = _install_package_wheel_first( + event_queue = eq, + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = _CAUSAL_CONV1D_PACKAGE_VERSION, + filename_prefix = "causal_conv1d", + release_tag = _CAUSAL_CONV1D_RELEASE_TAG, + release_base_url = ( + "https://github.com/Dao-AILab/causal-conv1d/releases/download" + ), + ) + return bool(ok) + + for gate_name, install_fn, post_fn in ( + ("is_flash_linear_attention_available", _fla_install, _fla_post_available), + ("is_causal_conv1d_available", _causal_conv1d_install, None), + ): + original = getattr(_iu, gate_name, None) + if original is None: + logger.info( + "%s missing on transformers.utils.import_utils; skipping hook", + gate_name, + ) + continue + wrapped = _make_wrapper(original, install_fn, gate_name, post_fn) + setattr(_iu, gate_name, wrapped) + rebound = _rebind_in_already_imported_modules( + attr_name = gate_name, old_obj = original, new_obj = wrapped + ) + logger.info( + "Installed fast-path hook on %s (rebound %d modules)", gate_name, rebound + ) + + def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool: if os.getenv(_FLASH_ATTN_SKIP_ENV) == "1": return False @@ -1199,9 +1756,28 @@ def run_training_process( model_name, ) - # ── 1b. Set up causal-conv1d first, then install mamba-ssm if needed ── + # ── 1b. Install fast-path kernel libraries for the chosen model. + # + # 1) causal-conv1d ALWAYS runs eagerly via the substring path. + # Some SSM modeling files (nemotron_h, falcon_h1, granitemoehybrid) + # use `lazy_load_kernel("causal-conv1d")` directly and never call + # transformers' `is_causal_conv1d_available()`, so the runtime + # hook on that gate would not fire for them. + # 2) FLA + tilelang: primary gate is the runtime hook on transformers' + # `is_flash_linear_attention_available`. Models whose architecture + # queries that gate auto-trigger the install; others never pay. + # `_install_fast_path_hooks` also wraps `is_causal_conv1d_available` + # as a defence in depth for newer modeling files that do use it. + # 3) mamba-ssm + flash-attn keep their existing substring / size gates. + # 4) `UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1` falls back to the + # substring path for FLA / tilelang. try: _ensure_causal_conv1d_fast_path(event_queue, model_name) + if os.getenv(_FAST_PATH_HOOKS_SKIP_ENV) == "1": + _ensure_flash_linear_attention(event_queue, model_name) + _ensure_tilelang_backend(event_queue, model_name) + else: + _install_fast_path_hooks(event_queue, model_name) _ensure_mamba_ssm(event_queue, model_name) _ensure_flash_attn_for_long_context( event_queue, @@ -1213,7 +1789,9 @@ def run_training_process( "type": "error", "error": ( f"Please choose another model to train, since " - f"causal-conv1d / mamba-ssm failed to install " + f"a fast-path kernel library " + f"(causal-conv1d / flash-linear-attention / " + f"mamba-ssm / tilelang) failed to install " f"with error: {exc}" ), "stack": traceback.format_exc(limit = 20), diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 733b726656..94279c28b4 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -70,7 +70,7 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768) - assert statuses == ["Installing prebuilt flash-attn wheel..."] + assert statuses == ["Installing flash-attn for faster training..."] def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): @@ -207,6 +207,1303 @@ def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch): ) +def _force_missing_fla_imports(monkeypatch): + """Make fla.modules / fla.ops.gated_delta_rule imports raise ImportError.""" + real_import = builtins.__import__ + + def fake_import(name, *a, **kw): + if name.startswith("fla.modules") or name.startswith("fla.ops"): + raise ImportError + return real_import(name, *a, **kw) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + +def test_flash_linear_attention_installs_pinned_pair_for_qwen3_5(monkeypatch): + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + _force_missing_fla_imports(monkeypatch) + statuses: list[str] = [] + monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg)) + + worker._ensure_flash_linear_attention( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_called_once() + args = run_mock.call_args[0][0] + assert f"flash-linear-attention=={worker._FLA_PACKAGE_VERSION}" in args + assert f"fla-core=={worker._FLA_CORE_PACKAGE_VERSION}" in args + assert "--no-deps" in args + assert run_mock.call_args.kwargs["timeout"] == worker._TILELANG_INSTALL_TIMEOUT_S + assert any("flash-linear-attention" in s for s in statuses) + + +def test_flash_linear_attention_skips_for_unrelated_models(monkeypatch): + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + + worker._ensure_flash_linear_attention( + event_queue = [], + model_name = "meta-llama/Llama-3.2-1B-Instruct", + ) + + run_mock.assert_not_called() + + +def test_flash_linear_attention_skips_for_ssm_only_models(monkeypatch): + # Nemotron-H / Falcon-H1 / Granite-H / LFM2 take the mamba_ssm path + # and never call FLA's gated_delta_rule kernels. + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + + for name in ( + "tiiuae/Falcon-H1-0.5B-Instruct", + "nvidia/Nemotron-H-8B-Base", + "ibm-granite/granite-4.0-h-tiny", + "LiquidAI/LFM2-1.2B-Instruct", + ): + worker._ensure_flash_linear_attention(event_queue = [], model_name = name) + + run_mock.assert_not_called() + + +def test_flash_linear_attention_matches_full_qwen3_family(monkeypatch): + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + _force_missing_fla_imports(monkeypatch) + monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None) + # Hermetic discovery: pretend installed transformers ships all the Qwen GDN families. + monkeypatch.setattr( + worker, + "_discover_fla_model_types", + lambda: frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_6", "qwen3_next"}), + ) + + for name in ( + "unsloth/Qwen3.5-2B", + "unsloth/Qwen3_5-MoE-A22B", + "unsloth/Qwen3.6-4B", + "unsloth/Qwen3_6-4B", + "unsloth/Qwen3-Next-80B-A3B", + "unsloth/Qwen3_Next-80B-A3B", + ): + worker._ensure_flash_linear_attention(event_queue = [], model_name = name) + + assert run_mock.call_count == 6 + + +def test_flash_linear_attention_skipped_below_python_3_10(monkeypatch): + # sys.version_info is a structseq, not constructible; substitute a + # plain tuple so the `< _FLA_MIN_PYTHON` comparison still works. + monkeypatch.setattr(worker.sys, "version_info", (3, 9, 0, "final", 0)) + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + + worker._ensure_flash_linear_attention( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_not_called() + + +def test_flash_linear_attention_skipped_via_env(monkeypatch): + monkeypatch.setenv(worker._FLA_SKIP_ENV, "1") + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + + worker._ensure_flash_linear_attention( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_not_called() + + +def test_flash_linear_attention_skipped_below_torch_2_7(monkeypatch): + monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False) + monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 5)) + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + statuses: list[str] = [] + monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg)) + + worker._ensure_flash_linear_attention( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_not_called() + assert any("torch>=" in s for s in statuses) + + +def test_flash_linear_attention_install_includes_einops(monkeypatch): + monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9)) + monkeypatch.setattr(worker, "_flash_linear_attention_importable", lambda: False) + 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) + + worker._ensure_flash_linear_attention( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + args = run_mock.call_args[0][0] + assert "--no-deps" in args + # einops is declared by fla-core; packaging and triton are pulled in + # because fla/utils.py imports them at module load but neither is + # declared in fla-core's METADATA (an upstream FLA gap). + assert "einops" in args + assert "packaging" in args + assert "triton" in args + assert f"flash-linear-attention=={worker._FLA_PACKAGE_VERSION}" in args + assert f"fla-core=={worker._FLA_CORE_PACKAGE_VERSION}" in args + + +def test_flash_linear_attention_logs_post_install_import_failure(monkeypatch): + """pip exits 0 but `import fla.modules` still fails (missing transitive).""" + monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9)) + import_calls = {"count": 0} + + def fake_importable(): + import_calls["count"] += 1 + # First call (pre-install probe) -> False so we attempt install. + # Second call (post-install verify) -> still False. + return False + + monkeypatch.setattr(worker, "_flash_linear_attention_importable", fake_importable) + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + statuses: list[str] = [] + monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg)) + + worker._ensure_flash_linear_attention( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + assert import_calls["count"] == 2 + assert any("not importable" in s for s in statuses) + + +def test_tilelang_backend_skipped_on_unsupported_linux_arch(monkeypatch): + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.sys, "platform", "linux") + import platform as _platform + + monkeypatch.setattr(_platform, "machine", lambda: "ppc64le") + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + + worker._ensure_tilelang_backend( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_not_called() + + +def test_tilelang_backend_pins_only_binary(monkeypatch): + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None) + monkeypatch.setattr(worker, "_tilelang_importable", lambda: False) + 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) + # Need to bypass the post-install probe too. + probe_calls = {"count": 0} + + def fake_probe(): + probe_calls["count"] += 1 + # First probe (pre-install): False so install runs. + # Second probe (post-install): True so success branch taken. + return probe_calls["count"] > 1 + + monkeypatch.setattr(worker, "_tilelang_importable", fake_probe) + + worker._ensure_tilelang_backend( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + args = run_mock.call_args[0][0] + assert "--only-binary=:all:" in args + assert "--no-deps" not in args + + +def _force_missing_tilelang_imports(monkeypatch): + real_import = builtins.__import__ + + def fake_import(name, *a, **kw): + if name in ("tilelang", "tvm_ffi"): + raise ImportError + return real_import(name, *a, **kw) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + +def test_tilelang_backend_installs_pinned_pair_for_qwen3_5(monkeypatch): + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None) + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + _force_missing_tilelang_imports(monkeypatch) + statuses: list[str] = [] + monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg)) + + worker._ensure_tilelang_backend( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_called_once() + args = run_mock.call_args[0][0] + assert f"apache-tvm-ffi=={worker._APACHE_TVM_FFI_PACKAGE_VERSION}" in args + assert f"tilelang=={worker._TILELANG_PACKAGE_VERSION}" in args + assert run_mock.call_args.kwargs["timeout"] == worker._TILELANG_INSTALL_TIMEOUT_S + assert any("Installing TileLang" in s for s in statuses) + + +def test_tilelang_backend_reinstalls_when_tvm_ffi_is_broken(monkeypatch): + """Repair path issues TWO pip calls: + + Call 1 (repair): `--force-reinstall --no-deps apache-tvm-ffi==0.1.9` + — surgically downgrades the broken package only. `--no-deps` here + is REQUIRED to prevent --force-reinstall from cascading through + apache-tvm-ffi's dep graph and replacing torch / the CUDA stack. + + Call 2 (install): plain `apache-tvm-ffi==0.1.9 tilelang==0.1.8` + — resolves missing transitive deps (z3-solver, ml-dtypes) without + --force-reinstall, so it never replaces already-correct packages. + """ + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.11") + 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) + + worker._ensure_tilelang_backend( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + assert run_mock.call_count == 2 + repair_args, install_args = (call[0][0] for call in run_mock.call_args_list) + + # Repair: --force-reinstall --no-deps, apache-tvm-ffi ONLY (no tilelang). + assert "--force-reinstall" in repair_args + assert ( + "--no-deps" in repair_args + ), "Repair MUST use --no-deps to avoid replacing torch / CUDA" + assert "--only-binary=:all:" in repair_args + assert f"apache-tvm-ffi=={worker._APACHE_TVM_FFI_PACKAGE_VERSION}" in repair_args + assert all( + "tilelang" not in a for a in repair_args + ), "Repair MUST only touch apache-tvm-ffi" + + # Install: regular dep-resolving install, NO --force-reinstall. + assert "--force-reinstall" not in install_args + assert "--no-deps" not in install_args + assert "--only-binary=:all:" in install_args + assert f"apache-tvm-ffi=={worker._APACHE_TVM_FFI_PACKAGE_VERSION}" in install_args + assert f"tilelang=={worker._TILELANG_PACKAGE_VERSION}" in install_args + + +def test_tilelang_backend_skipped_below_python_3_10(monkeypatch): + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + # sys.version_info is a structseq, not constructible; substitute a + # plain tuple so the `< _FLA_MIN_PYTHON` comparison still works. + monkeypatch.setattr(worker.sys, "version_info", (3, 9, 0, "final", 0)) + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + + worker._ensure_tilelang_backend( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_not_called() + + +def test_tilelang_backend_skipped_on_windows(monkeypatch): + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.sys, "platform", "win32") + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + + worker._ensure_tilelang_backend( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_not_called() + + +def test_tilelang_backend_swallows_install_timeout(monkeypatch): + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None) + _force_missing_tilelang_imports(monkeypatch) + + def raise_timeout(*a, **kw): + raise subprocess.TimeoutExpired(cmd = "pip", timeout = 1) + + monkeypatch.setattr(worker._sp, "run", raise_timeout) + statuses: list[str] = [] + monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg)) + + # Should not raise. + worker._ensure_tilelang_backend( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + assert any("timed out" in s.lower() for s in statuses) + + +def test_tilelang_backend_skipped_for_ssm_models(monkeypatch): + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + + # Nemotron-H / Falcon-H1 / Granite-H take the mamba_ssm path, not FLA's + # gated_delta_rule -> tilelang has no effect on them. + for name in ( + "tiiuae/Falcon-H1-0.5B-Instruct", + "nvidia/Nemotron-H-8B-Base", + "ibm-granite/granite-4.0-h-tiny", + "meta-llama/Llama-3.2-1B-Instruct", + ): + worker._ensure_tilelang_backend(event_queue = [], model_name = name) + + run_mock.assert_not_called() + + +def test_tilelang_backend_skipped_via_env(monkeypatch): + monkeypatch.setenv(worker._TILELANG_SKIP_ENV, "1") + run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = "")) + monkeypatch.setattr(worker._sp, "run", run_mock) + + worker._ensure_tilelang_backend( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_not_called() + + +def test_tilelang_backend_swallows_install_failure(monkeypatch): + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.shutil, "which", lambda name: None) + monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None) + run_mock = mock.Mock(return_value = mock.Mock(returncode = 1, stdout = "boom")) + monkeypatch.setattr(worker._sp, "run", run_mock) + _force_missing_tilelang_imports(monkeypatch) + statuses: list[str] = [] + monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg)) + + # Should not raise even when pip exits non-zero. + worker._ensure_tilelang_backend( + event_queue = [], + model_name = "unsloth/Qwen3.5-2B", + ) + + run_mock.assert_called_once() + assert any("failed" in s.lower() for s in statuses) + + +# ─────────────────────────────────────────────────────────────────── +# Runtime hook on `is_flash_linear_attention_available` / +# `is_causal_conv1d_available`. These are the primary gate in +# normal operation; the substring tests above cover the +# UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 fallback. +# ─────────────────────────────────────────────────────────────────── + + +class _FakeQueue(list): + """List with `.put` so worker._send_status can send into it during tests.""" + + def put(self, item): + self.append(item) + + +def _make_fake_gate(initial_return: bool): + """Build a callable that mimics transformers' lru_cache-decorated gates. + + Tracks call count and exposes a `cache_clear` attribute. The return + value can be flipped to mimic install-then-True behaviour by setting + `.next_return`. + """ + + class Gate: + def __init__(self, initial: bool) -> None: + self.next_return = initial + self.call_count = 0 + self.cache_clear_count = 0 + + def __call__(self) -> bool: + self.call_count += 1 + return self.next_return + + def cache_clear(self) -> None: + self.cache_clear_count += 1 + + return Gate(initial_return) + + +def _patch_iu_gates(monkeypatch, fla_gate, conv_gate): + """Drop fake gates onto transformers.utils.import_utils for the test.""" + from transformers.utils import import_utils as _iu + + monkeypatch.setattr(_iu, "is_flash_linear_attention_available", fla_gate) + monkeypatch.setattr(_iu, "is_causal_conv1d_available", conv_gate) + + +def test_hook_installs_when_gate_returns_false(monkeypatch): + fla_gate = _make_fake_gate(initial_return = False) + conv_gate = _make_fake_gate(initial_return = False) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + def _fla_install_side_effect(eq): + fla_gate.next_return = True + return True + + fla_install = mock.Mock(side_effect = _fla_install_side_effect) + tile_install = mock.Mock(side_effect = lambda eq: None) + + def _conv_install_side_effect(**kw): + conv_gate.next_return = True + return True + + conv_install = mock.Mock(side_effect = _conv_install_side_effect) + + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) + monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + from transformers.utils import import_utils as _iu + + # Both gates are now wrapped. Call them — the hook should drive the install. + assert _iu.is_flash_linear_attention_available() is True + fla_install.assert_called_once() + tile_install.assert_called_once() + assert _iu.is_causal_conv1d_available() is True + conv_install.assert_called_once() + + +def test_hook_skips_install_when_gate_already_true(monkeypatch): + """When both gates are already True AND tilelang is healthy, the hook + must do zero install work. (Tilelang repair on the already-True path + is covered by test_hook_runs_tilelang_repair_when_fla_already_true.) + """ + fla_gate = _make_fake_gate(initial_return = True) + conv_gate = _make_fake_gate(initial_return = True) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + fla_install = mock.Mock() + tile_install = mock.Mock() + conv_install = mock.Mock() + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) + monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install) + # Tilelang healthy so the post_available path is a no-op (otherwise + # it would call tile_install, which is correct behaviour but + # outside the scope of this test). + monkeypatch.setattr(worker, "_tilelang_importable", lambda: True) + monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.9") + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + from transformers.utils import import_utils as _iu + + assert _iu.is_flash_linear_attention_available() is True + assert _iu.is_causal_conv1d_available() is True + fla_install.assert_not_called() + tile_install.assert_not_called() + conv_install.assert_not_called() + + +def test_hook_idempotent_on_repeat_call(monkeypatch): + fla_gate = _make_fake_gate(initial_return = False) + conv_gate = _make_fake_gate(initial_return = False) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + def _fla_install_side_effect(eq): + fla_gate.next_return = True + return True + + fla_install = mock.Mock(side_effect = _fla_install_side_effect) + tile_install = mock.Mock() + + def _conv_install_side_effect(**kw): + conv_gate.next_return = True + return True + + conv_install = mock.Mock(side_effect = _conv_install_side_effect) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) + monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + from transformers.utils import import_utils as _iu + + # First call: hook fires. + _iu.is_flash_linear_attention_available() + # Subsequent calls: must not re-trigger the installer. + _iu.is_flash_linear_attention_available() + _iu.is_flash_linear_attention_available() + assert fla_install.call_count == 1 + assert tile_install.call_count == 1 + + +def test_hook_handles_install_failure_gracefully(monkeypatch): + fla_gate = _make_fake_gate(initial_return = False) + conv_gate = _make_fake_gate(initial_return = True) # bypass to focus on FLA + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + def raising_install(eq): + raise RuntimeError("pip failed to fetch wheel") + + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", raising_install + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", lambda eq: None + ) + monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: None) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + from transformers.utils import import_utils as _iu + + # Must not raise; returns False so transformers falls back to torch loop. + assert _iu.is_flash_linear_attention_available() is False + + +def test_hook_can_be_disabled_via_env(monkeypatch): + fla_gate = _make_fake_gate(initial_return = False) + conv_gate = _make_fake_gate(initial_return = False) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + fla_install = mock.Mock() + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) + monkeypatch.setenv(worker._FAST_PATH_HOOKS_SKIP_ENV, "1") + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + from transformers.utils import import_utils as _iu + + # Hook should NOT have been installed; gates remain the fakes. + assert _iu.is_flash_linear_attention_available is fla_gate + assert _iu.is_causal_conv1d_available is conv_gate + fla_install.assert_not_called() + + +def test_hook_clears_lru_cache_before_first_check(monkeypatch): + fla_gate = _make_fake_gate(initial_return = True) + conv_gate = _make_fake_gate(initial_return = True) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", lambda eq: None + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", lambda eq: None + ) + monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: None) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + from transformers.utils import import_utils as _iu + + _iu.is_flash_linear_attention_available() + # The wrapper called cache_clear at least once before delegating. + assert fla_gate.cache_clear_count >= 1 + + +def test_hook_rewrites_previously_imported_module_bindings(monkeypatch): + """Modeling files bind `is_flash_linear_attention_available` locally + via `from ... import is_X`. Reassigning the attribute on + transformers.utils.import_utils alone does NOT reach those local + bindings. The hook installer sweeps sys.modules and rebinds them. + """ + fla_gate = _make_fake_gate(initial_return = False) + conv_gate = _make_fake_gate(initial_return = True) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + # Create a fake modeling module that did `from ... import is_flash_linear_attention_available`. + fake_mod = sys.modules.setdefault( + "_test_fake_modeling_qwen35", type(sys)("_test_fake_modeling_qwen35") + ) + fake_mod.is_flash_linear_attention_available = fla_gate + + def fake_install(eq): + fla_gate.next_return = True + return True + + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fake_install + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", lambda eq: True + ) + monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + # The fake module's local binding has been rewritten to the wrapper. + assert fake_mod.is_flash_linear_attention_available is not fla_gate + # Calling through the fake module's reference triggers the install. + assert fake_mod.is_flash_linear_attention_available() is True + + del sys.modules["_test_fake_modeling_qwen35"] + + +def test_hook_skips_when_import_utils_unavailable(monkeypatch): + """If transformers.utils.import_utils can't be imported, the hook + installer must log and return cleanly rather than crash the worker.""" + real_import = builtins.__import__ + + def fake_import(name, *a, **kw): + if name == "transformers.utils" or name == "transformers.utils.import_utils": + raise ImportError("transformers missing in worker venv") + return real_import(name, *a, **kw) + + monkeypatch.setattr(builtins, "__import__", fake_import) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + # Should not raise. + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + +def test_substring_fallback_unchanged_when_hook_skipped(monkeypatch): + """Hook disabled -> legacy gate falls back to auto-discovered model types.""" + install_mock = mock.Mock() + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", install_mock + ) + monkeypatch.setattr( + worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_5"}) + ) + monkeypatch.setenv(worker._FAST_PATH_HOOKS_SKIP_ENV, "1") + + worker._ensure_flash_linear_attention( + event_queue = [], model_name = "unsloth/Qwen3.5-2B" + ) + assert install_mock.call_count == 1 + + worker._ensure_flash_linear_attention( + event_queue = [], model_name = "meta-llama/Llama-3.1-8B" + ) + assert install_mock.call_count == 1 + + +# ─────────────────────────────────────────────────────────────────── +# Regression tests for the 10-reviewer findings: +# 1. tilelang Qwen-guard on hook path (non-Qwen FLA models) +# 2. tilelang repair must not replace torch / CUDA stack +# 3. hook must trust installer's bool, not transformers metadata +# 4. causal-conv1d must stay eager for SSM models that bypass the gate +# 5. rebind sweep must not invoke lazy module __getattr__ +# 6. tilelang skipped when FLA was skipped / failed +# 7. tilelang repair runs when FLA is already True +# 8. older FLA detected as stale and reinstalled +# ─────────────────────────────────────────────────────────────────── + + +def test_hook_does_not_install_tilelang_for_model_outside_allowlist(monkeypatch): + """A model whose name is not in the auto-discovered FLA allowlist calls + is_flash_linear_attention_available but should NOT get tilelang.""" + fla_gate = _make_fake_gate(initial_return = False) + conv_gate = _make_fake_gate(initial_return = True) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + def _fla_install(eq): + fla_gate.next_return = True + return True + + fla_install = mock.Mock(side_effect = _fla_install) + tile_install = mock.Mock(return_value = True) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + # Hermetize the auto-discovered set so the test stays valid as new + # transformers releases add FLA-using model_types (eg olmo_hybrid in + # 5.4.0). The semantic under test is "outside-allowlist -> no tilelang". + monkeypatch.setattr( + worker, + "_discover_fla_model_types", + lambda: frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_next"}), + ) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), + model_name = "fake-org/Fictional-FLA-Only-Model-7B", + ) + + from transformers.utils import import_utils as _iu + + assert _iu.is_flash_linear_attention_available() is True + fla_install.assert_called_once() + tile_install.assert_not_called() + + +def test_hook_does_install_tilelang_for_qwen35(monkeypatch): + """Positive control for finding #1: Qwen3.5 still gets tilelang.""" + fla_gate = _make_fake_gate(initial_return = False) + conv_gate = _make_fake_gate(initial_return = True) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + def _fla_install(eq): + fla_gate.next_return = True + return True + + fla_install = mock.Mock(side_effect = _fla_install) + tile_install = mock.Mock(return_value = True) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + from transformers.utils import import_utils as _iu + + _iu.is_flash_linear_attention_available() + fla_install.assert_called_once() + tile_install.assert_called_once() + + +def test_tilelang_repair_does_not_touch_torch_cuda_stack(monkeypatch): + """Finding #2: the broken-tvm-ffi repair must use --no-deps on the + forced step so --force-reinstall does not cascade through + apache-tvm-ffi's dep graph and pull a different torch wheel. + """ + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.10") + 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) + + worker._ensure_tilelang_backend(event_queue = [], model_name = "unsloth/Qwen3.5-2B") + + assert run_mock.call_count == 2 + repair_args = run_mock.call_args_list[0][0][0] + # The forced step MUST be --no-deps so torch / CUDA stack is untouched. + assert "--force-reinstall" in repair_args and "--no-deps" in repair_args + # And it touches ONLY apache-tvm-ffi, not tilelang / torch. + assert all("tilelang" not in a for a in repair_args) + assert all("torch" not in a for a in repair_args) + + +def test_hook_trusts_installer_bool_not_metadata(monkeypatch): + """Finding #3: if pip exits 0 but deep imports fail, the installer + returns False; the hook must propagate False even if the underlying + `original()` gate (which only checks metadata) returns True after + pip succeeds. + + Setup mirrors the real bug: + 1. Pre-install: gate=False (FLA not present) → wrapper triggers install. + 2. Installer's `_flash_linear_attention_importable` post-probe fails, + so the installer returns False. (pip exited 0 but `import fla.modules` + raised because of a missing transitive dep.) + 3. Post-install: gate would return True (metadata check sees fla-core + version) — but the wrapper must IGNORE that and use the installer's + False so transformers takes the torch fallback. + """ + # Gate flips True after install (simulating "metadata sees fla"). + fla_gate = _make_fake_gate(initial_return = False) + conv_gate = _make_fake_gate(initial_return = True) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + # Installer "succeeds" at pip, AND flips the gate to True (metadata + # sees fla post-install), BUT returns False (deep import broken). + def _bad_install(eq): + fla_gate.next_return = True # metadata says yes after pip + return False # but deep import is broken + + fake_fla_install = mock.Mock(side_effect = _bad_install) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fake_fla_install + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", mock.Mock(return_value = True) + ) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + from transformers.utils import import_utils as _iu + + # Hook MUST return False (installer's verdict), not True (metadata lies). + assert _iu.is_flash_linear_attention_available() is False + fake_fla_install.assert_called_once() + + +def test_rebind_does_not_trigger_module_getattr(monkeypatch): + """Finding #5: the rebind sweep must use __dict__, not getattr(), + to avoid invoking transformers' lazy module __getattr__ which spits + out hundreds of "Accessing X from .models..." warnings. + """ + original = object() + replacement = object() + + class _GetattrTripwire(type(sys)): + getattr_called = False + + def __getattr__(self, name): + type(self).getattr_called = True + raise AttributeError(name) + + lazy = _GetattrTripwire("_lazy_test_module") + sys.modules["_lazy_test_module"] = lazy + try: + # No module-level binding to `is_flash_linear_attention_available` + # in __dict__, so the sweep must NOT trip the tripwire. + worker._rebind_in_already_imported_modules( + attr_name = "is_flash_linear_attention_available", + old_obj = original, + new_obj = replacement, + ) + assert ( + not _GetattrTripwire.getattr_called + ), "Rebind sweep invoked __getattr__ — should use __dict__ probe" + finally: + sys.modules.pop("_lazy_test_module", None) + + +def test_hook_skips_tilelang_when_fla_install_is_skipped(monkeypatch): + """Finding #6: env-skipped FLA returns False from + _ensure_flash_linear_attention_unconditional; tilelang must NOT + install in that case. + """ + fla_gate = _make_fake_gate(initial_return = False) + conv_gate = _make_fake_gate(initial_return = True) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + monkeypatch.setenv(worker._FLA_SKIP_ENV, "1") + tile_install = mock.Mock(return_value = True) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + from transformers.utils import import_utils as _iu + + # FLA gate stays False (env-skipped, install never ran). + assert _iu.is_flash_linear_attention_available() is False + tile_install.assert_not_called() + + +def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch): + """Finding #7: when FLA is already importable (gate returns True at + first probe) but tilelang is missing or apache-tvm-ffi is on the + broken list, the post-available action must still run tilelang. + """ + fla_gate = _make_fake_gate(initial_return = True) + conv_gate = _make_fake_gate(initial_return = True) + _patch_iu_gates(monkeypatch, fla_gate, conv_gate) + + fla_install = mock.Mock(return_value = True) + tile_install = mock.Mock(return_value = True) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) + # tilelang missing AND tvm-ffi is on broken list — both trigger repair. + monkeypatch.setattr(worker, "_tilelang_importable", lambda: False) + monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.11") + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) + + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) + + from transformers.utils import import_utils as _iu + + _iu.is_flash_linear_attention_available() + # FLA install was NOT needed; tilelang repair WAS still triggered. + fla_install.assert_not_called() + tile_install.assert_called_once() + + +def test_fla_installer_force_reinstalls_when_older_version_present(monkeypatch): + """Finding #8: when an older `flash-linear-attention` is importable + but below the pin, the installer must force a reinstall (not no-op). + """ + monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False) + monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv") + monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9)) + # Importable but stale (current() reports False even though importable() is True). + monkeypatch.setattr(worker, "_flash_linear_attention_importable", lambda: True) + monkeypatch.setattr(worker, "_flash_linear_attention_current", lambda **kw: False) + 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) + + worker._ensure_flash_linear_attention_unconditional(event_queue = []) + + run_mock.assert_called_once() + args = run_mock.call_args[0][0] + assert ( + "--force-reinstall" in args + ), "Stale FLA must trigger --force-reinstall, otherwise pip is a no-op" + # --no-deps still applies so torch stays untouched. + assert "--no-deps" in args + + +def test_run_training_process_eagerly_installs_causal_conv1d_in_normal_mode(): + """Finding #4: SSM modeling files use `lazy_load_kernel("causal-conv1d")` + and never call `is_causal_conv1d_available()`, so the hook would not + fire for them. The orchestrator must always run the eager + substring installer regardless of hook mode. + + This test reads the worker source rather than running the full + orchestrator (which requires a configured training config). It + asserts the eager install is OUTSIDE the if/else hook branch. + """ + import inspect + + src = inspect.getsource(worker.run_training_process) + # Find the orchestration block. + assert "_ensure_causal_conv1d_fast_path(event_queue, model_name)" in src + assert "_install_fast_path_hooks(event_queue, model_name)" in src + # The eager causal_conv1d call must appear BEFORE the hook-mode if/else, + # not nested inside the `if _FAST_PATH_HOOKS_SKIP_ENV` branch. + eager_pos = src.find("_ensure_causal_conv1d_fast_path(event_queue, model_name)") + skip_check_pos = src.find('os.getenv(_FAST_PATH_HOOKS_SKIP_ENV) == "1"') + assert eager_pos < skip_check_pos, ( + "_ensure_causal_conv1d_fast_path must be called BEFORE the hook-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 + + +# ─────────────────────────────────────────────────────────────────── +# Auto-discovery of FLA model_types from the installed transformers +# ─────────────────────────────────────────────────────────────────── + + +def _make_fake_transformers_tree( + tmp_path, fla_types: list[str], non_fla_types: list[str] +): + """Lay out a tmp dir as `transformers/models/{type}/modeling_{type}.py`.""" + pkg = tmp_path / "transformers" + models = pkg / "models" + models.mkdir(parents = True) + (pkg / "__init__.py").write_text("") + for t in fla_types: + d = models / t + d.mkdir() + (d / f"modeling_{t}.py").write_text( + "from ...utils.import_utils import is_flash_linear_attention_available\n" + "if is_flash_linear_attention_available():\n" + " from fla.modules import FusedRMSNormGated\n" + " from fla.ops.gated_delta_rule import chunk_gated_delta_rule\n" + ) + for t in non_fla_types: + d = models / t + d.mkdir() + (d / f"modeling_{t}.py").write_text("class Foo: pass\n") + return pkg + + +def _reset_fla_cache(monkeypatch): + monkeypatch.setattr(worker, "_TRANSFORMERS_FLA_MODEL_TYPES_CACHE", None) + + +def test_discover_fla_model_types_returns_only_fla_users(tmp_path, monkeypatch): + pkg = _make_fake_transformers_tree( + tmp_path, + fla_types = ["qwen3_5", "qwen3_5_moe", "qwen3_next"], + non_fla_types = ["llama", "gpt2", "mistral"], + ) + fake = mock.MagicMock(__file__ = str(pkg / "__init__.py")) + monkeypatch.setitem(sys.modules, "transformers", fake) + _reset_fla_cache(monkeypatch) + + result = worker._discover_fla_model_types() + assert result == frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_next"}) + assert "llama" not in result + assert "gpt2" not in result + + +def test_discover_fla_model_types_caches_across_calls(tmp_path, monkeypatch): + pkg = _make_fake_transformers_tree( + tmp_path, fla_types = ["qwen3_5"], non_fla_types = [] + ) + fake = mock.MagicMock(__file__ = str(pkg / "__init__.py")) + monkeypatch.setitem(sys.modules, "transformers", fake) + _reset_fla_cache(monkeypatch) + + from pathlib import Path as _Path + + read_calls = [0] + real_read = _Path.read_text + + def counting_read(self, *a, **kw): + read_calls[0] += 1 + return real_read(self, *a, **kw) + + monkeypatch.setattr(_Path, "read_text", counting_read) + + first = worker._discover_fla_model_types() + after_first = read_calls[0] + second = worker._discover_fla_model_types() + + assert first == second + assert read_calls[0] == after_first # cache hit: no extra disk reads + + +def test_discover_fla_model_types_handles_missing_transformers(monkeypatch): + _reset_fla_cache(monkeypatch) + + real_import = builtins.__import__ + + def fake_import(name, globals = None, locals = None, fromlist = (), level = 0): + if name == "transformers": + raise ImportError("transformers not installed") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + result = worker._discover_fla_model_types() + assert result == frozenset() + + +def test_discover_fla_model_types_handles_unreadable_file(tmp_path, monkeypatch): + pkg = _make_fake_transformers_tree( + tmp_path, fla_types = ["qwen3_5"], non_fla_types = [] + ) + fake = mock.MagicMock(__file__ = str(pkg / "__init__.py")) + monkeypatch.setitem(sys.modules, "transformers", fake) + _reset_fla_cache(monkeypatch) + + from pathlib import Path as _Path + + real_read = _Path.read_text + + def boom_read(self, *a, **kw): + if "modeling_qwen3_5.py" in str(self): + raise OSError("permission denied") + return real_read(self, *a, **kw) + + monkeypatch.setattr(_Path, "read_text", boom_read) + result = worker._discover_fla_model_types() + assert result == frozenset() # unreadable file simply doesn't contribute + + +def test_model_wants_tilelang_handles_real_repo_names(monkeypatch): + monkeypatch.setattr( + worker, + "_discover_fla_model_types", + lambda: frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_next"}), + ) + cases = [ + ("unsloth/Qwen3.5-2B", True), + ("Qwen/Qwen3.5-MoE-A3B", True), + ("mlx-community/qwen3-next-80b", True), + ("unsloth/qwen3_5_moe_a3b_lora", True), + ("meta-llama/Llama-3.1-8B", False), + ("nvidia/Nemotron-H-4B", False), + ("mistralai/Mistral-7B-v0.3", False), + ("", False), + ] + for name, expected in cases: + assert worker._model_wants_tilelang(name) is expected, name + + +def test_model_wants_tilelang_empty_when_transformers_has_no_fla(monkeypatch): + monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset()) + assert worker._model_wants_tilelang("unsloth/Qwen3.5-2B") is False + assert worker._model_wants_tilelang("meta-llama/Llama-3.1-8B") is False + + +def test_model_wants_tilelang_normalizes_separators(monkeypatch): + monkeypatch.setattr( + worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_next"}) + ) + for variant in ( + "qwen3-next", + "Qwen3.Next", + "Qwen/Qwen3 Next", + "anyone/qwen3_next", + "qwen3.next-80b", + ): + assert worker._model_wants_tilelang(variant) is True, variant + + # ──────────────────────────────────────────────────────────────────── # HIP source-build gcc-install-dir coverage (h34v3nzc0dex Strix Halo). # Ubuntu 24.04 ships gcc-14's runtime dir without /usr/include/c++/14,