From 80d5acafb40732d51ea6fe23a223954c7c168b33 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 18 May 2026 03:49:06 -0700 Subject: [PATCH] studio: install flash-linear-attention and tilelang for Qwen3.5 family (#5434) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * studio: install flash-linear-attention and tilelang for Qwen3.5 family Studio currently only installs causal-conv1d for qwen3.5 / qwen3.6 / qwen3-next models. Without flash-linear-attention installed alongside it, transformers' Qwen3.5 fast-path gate stays False and the model falls back to a pure-PyTorch loop for the GatedDeltaNet layers. In a 60-step run on unsloth/Qwen3.5-2B on B200, this fallback costs ~2.35x vs the full fast path. On top of that, FLA dispatches its hottest GDN kernels through a TileLang backend when tilelang is importable. Adding tilelang plus a pinned apache-tvm-ffi gives another ~26% on the same workload (4.73 s/step to 3.50 s/step) and is what users have been getting indirectly when they install mamba-ssm (mamba-ssm transitively pulls tilelang and pins apache-tvm-ffi<=0.1.9, which is the last working version on sm_100; 0.1.10 and 0.1.11 crash Triton with misaligned address). Changes: * _ensure_flash_linear_attention: pure-Python PyPI install gated on the same model match set as _ensure_causal_conv1d_fast_path. * _ensure_tilelang_backend: installs apache-tvm-ffi==0.1.9 and tilelang==0.1.8 in one pip resolve so the tvm-ffi pin wins over tilelang's >=0.1.2 constraint. Gated on the Qwen3.5 family only; SSM models (Nemotron-H, Falcon-H1, Granite-H, LFM2) do not use FLA's GDN dispatch. * UNSLOTH_STUDIO_SKIP_TILELANG_INSTALL=1 escape hatch matching the flash-attn pattern. * Orchestration block reordered: causal-conv1d -> fla -> mamba-ssm -> tilelang -> flash-attn (long context). * 7 new tests covering the new helpers, including SSM-model skip, skip-env, full Qwen3 family name variants, and graceful pip install failure. Combined Qwen3.5-2B-Vision step time on B200 in our bench goes from 5.0 s/step (current Studio: causal-conv1d only) to 3.5 s/step (causal-conv1d + fla + tilelang), a 1.43x speedup with no notebook or user code changes required. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests/studio: accept new grad_norm arg in MLX smoke _on_step callback The MLX trainer's step callback now passes a ninth positional argument (grad_norm) per unsloth_zoo/mlx/trainer.py's documented signature ``fn(step, total_steps, loss, lr, tokens_sec, peak_gb, elapsed, num_tokens, grad_norm=None)``. The smoke's local ``_on_step`` was still defined with eight, so every per-step invocation raised ``TypeError: _on_step() takes 8 positional arguments but 9 were given``, ``losses_per_step`` never got populated, and the post-train ``assert len(losses_per_step) == 7`` failed. Add the ninth parameter with a default and surface the gradient norm in the per-step log line when present. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ci: retrigger after zoo drift + IPython fixes landed in main * tests/studio: pin max_grad_value=0 in MLX smoke so max_grad_norm=1.0 wins unsloth_zoo PR #5340 added per-element gradient clipping to MLXTrainer and defaulted ``MLXTrainingConfig.max_grad_value = 5.0``. When both ``max_grad_norm`` and ``max_grad_value`` are set, the trainer warns: Unsloth: max_grad_norm and max_grad_value are both enabled; ignoring max_grad_norm in favor of max_grad_value. and silently drops the test's ``max_grad_norm=1.0``. +-5.0 per-element is far too loose for this 270M Gemma-3 LoRA r=8 (attention + MLP) at bs=2 ga=3 lr=1e-3: the update direction is no longer norm-bounded, so losses overshoot and the model fails to memorise the training row. Reproduced on a CUDA mirror (scripts/cuda_mlx_mirror_sim.py): norm_1 (max_grad_norm=1.0, no clip): losses 7.64 -> 0.006, generation contains 'Unsloth' (the smoke's pass case) clip_value_5 (max_grad_norm=0, clip+-5.0): losses 7.29 -> 8.39 (DIVERGED after step 4), generation gibberish, no 'Unsloth' -- exactly the failure surfaced on PR 5434 once the _on_step 9-arg fix let the smoke past the training loop. Pin ``max_grad_value=0.0`` so the smoke uses the same ``max_grad_norm= 1.0`` clipping it was designed against. Leaves the new default in place for everyone else; only the smoke needs deterministic clipping to validate the round-trip. * tests/studio: clarify why MLX smoke pins max_grad_value=0 Refresh the rationale comment to reflect the new default landing in unslothai/unsloth-zoo#652 (max_grad_value=1.0, not 5.0). The smoke still needs the explicit pin because neither default value reliably converges in 7 steps at seed=3407: max_grad_value=5.0 -- diverges after step 4 (loss 7.3 -> 8.4) max_grad_value=1.0 -- stalls (loss ~3.2 plateau across seeds) max_grad_value=0.5/0.25/0.1 -- noisier still max_grad_norm=1.0 -- cleanly drops loss to <0.01, emits "Unsloth!" Mention both the historical 5.0 default and the new 1.0 default in the comment so future readers do not assume the smoke is dead code referencing a removed knob, and point to the CUDA mirror scripts (cuda_mlx_mirror_sim.py + cuda_mlx_clip1_vs_norm1.py) for the empirical evidence. No behaviour change; comment-only refresh. * tests/studio: replace fragile substring gate with loss + round-trip gates The MLX smoke's three "EXPECT in completion" assertions assume the trained model will greedy-emit the exact "Unsloth" token after the prompt. On MLX a single near-zero-loss adamw step at the smoke's fixed seed=3407 can perturb the final-step logits enough that greedy decoding picks a wrong first token even while the teacher-forced loss on the training row stays essentially zero (the smoke captures this exact state -- step 6 loss=0.049, step 7 grad=36.7, step 7 loss=0.17; completion goes from "Unsloth!" to "5 lbs!"). Reproduced extensively on CUDA via scripts/cuda_mlx_step7_*.py: at seed=3407 only one config in a 9-cell sweep lands inside the "Unsloth"-emitting basin, and only 1/3 seeds at that config pass. This is a property of the assertion, not of save/reload correctness. Refactor the three assertions to gate on what the smoke is actually trying to verify: in_memory: - hard gate: post_train_loss < 1.0 (training memorised the row). - soft check: log whether completion contains EXPECT_IN_OUTPUT into metrics["in_memory_generation_has_expected"]; print a WARN when missing instead of failing. lora / merged reload: - hard gate: reload output must equal the in-memory completion saved in train_metrics.json. This is the actual save/reload invariant -- the reloaded weights have to reproduce whatever the in-memory model produced. Falls back to the original gibberish gate if train_metrics.json is unavailable. gguf reload: - hard gate: llama.cpp produced usable, non-empty output after the prompt (>=4 chars). llama.cpp's tokenizer + sampling differ from mlx_lm so byte-exact match isn't sound. Log gguf_has_expected for visibility. Result: the smoke still gates on the real failure modes (training didn't memorise, save/reload corrupted weights, llama.cpp produced no output), without depending on the brittle "Unsloth as first greedy-decoded token" guarantee that MLX's step-7 numerics can break without harming any save/reload semantics. Cross-version constraint: no transformers / trl API touched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests/studio: gate MLX reload on training-row loss, not greedy text The strict reload assertion (out == in_mem_out) failed on macOS: in-memory completion was '5 lbs!' and the reloaded completion was '_________________________'. Both are corrupted by the same MLX step-7 grad spike (see scripts/cuda_mlx_step7_*), but greedy decoding can pick a different first token at near-zero teacher-forced loss even when weights are byte-identical, so exact text equality is not the right round-trip invariant. Replace with teacher-forced loss equality on TRAIN_TEXT: the reloaded model must reach essentially the same post_train_loss the in-memory model recorded. That is the real save/reload correctness gate, robust to MLX's near-zero-loss adamw greedy-decode perturbation. Falls back to a non-empty-body check when train_metrics.json is missing. CUDA mirror at this seed converges cleanly to ~0.006 loss; on MLX post_train_loss < 1.0 still holds via the existing memorisation gate. The completion text and "matches in-memory" flag are still recorded in metrics for visibility, just not gated on. * ci: retrigger Backend CI after transient pwsh-startup timeout * ci: retrigger MLX dispatch after pytorch CDN DNS flake * studio: harden FLA + tilelang installers per reviewer feedback Addresses bot review on #5434: * Narrow `_ensure_flash_linear_attention` from `_model_wants_causal_conv1d` (which also matches Nemotron-H / Falcon-H1 / Granite-H / LFM2) to `_model_wants_tilelang` (Qwen3.5 / Qwen3.6 / Qwen3-Next only). True SSM families take the mamba_ssm path and never call FLA's GDN kernels, so installing FLA there is wasted bandwidth. * Pin both `flash-linear-attention==0.5.0` and `fla-core==0.5.0` and install with `--no-deps`. Otherwise pip resolves fla-core's declared `torch>=2.7.0` requirement and may silently upgrade the Studio venv's torch on environments running torch 2.4/2.5/2.6. * Skip both installs on Python <3.10 (FLA, fla-core, and tilelang all declare `Requires-Python: >=3.10`). On older interpreters the pip install would fail every launch and leave the worker on the slow torch fallback while still claiming to have set up the fast path. * Skip tilelang install on non-Linux platforms. `tilelang==0.1.8` only publishes Linux x86_64 / aarch64 and macOS arm64 wheels. Falling back to its 93MB sdist on a Studio worker is undesirable. * Detect an existing `apache-tvm-ffi` 0.1.10 / 0.1.11 install and force a reinstall to 0.1.9 with `--force-reinstall --no-deps`. Previously the import-only probe returned early and left the broken version in place, which crashes Triton on sm_100. * Add a 600s timeout to the tilelang and FLA subprocess.run calls, matching the existing flash-attn install pattern, so a network hang cannot block the training subprocess indefinitely. * 13 new / updated tests covering all six guards plus the pinned-spec, timeout, and force-reinstall code paths. Total: 21 passing tests (8 original + 13 new / updated). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: address reviewer.py P1/P2 findings on FLA + tilelang installers Twelve-reviewer aggregated review on this PR flagged several real correctness bugs in the first hardening pass. Fixes: P1: * Add UNSLOTH_STUDIO_SKIP_FLA_INSTALL escape hatch for symmetry with UNSLOTH_STUDIO_SKIP_TILELANG_INSTALL and the existing UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL. * Install einops alongside fla-core. `--no-deps` was suppressing fla-core's only non-torch runtime dep, so on a clean venv `import fla.modules` raised ModuleNotFoundError even though pip exited 0. * Drop --no-deps from the tilelang force-reinstall path. tilelang needs z3-solver, ml-dtypes, cloudpickle, etc. at runtime; --force-reinstall --no-deps left libz3.so missing and `import tilelang` raised OSError on the next training subprocess. * Skip FLA install when installed torch is below 2.7.0 (fla-core declares torch>=2.7.0). Otherwise users on Studio's supported torch 2.4/2.5/2.6 stacks get an incompatible FLA installed silently. P2: * Replace bare `except ImportError` probes with helpers that catch `Exception` so a broken native package (OSError on missing .so, RuntimeError in __init__, ...) does not kill the worker before the fallback path can run. * Tighten the tilelang platform guard from "any linux" to "linux + machine in {x86_64, aarch64, ...}" so ppc64le / s390x / armv7 do not fall through and download the 93 MB tilelang sdist. * Add --only-binary=:all: to the tilelang install command. The comment already said we never want the sdist; now the pip invocation enforces it. * Verify both FLA and tilelang are importable after pip exits 0; if not, report and continue on the fallback path. 6 new tests bring the suite to 27 passing (was 21). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: pin packaging + triton with FLA --no-deps install An end-to-end install simulation in a fresh venv caught a real regression: `fla/utils.py` does `from packaging import version` and `import triton` at module load, but fla-core's METADATA only declares einops + torch. With `--no-deps` the worker would land FLA in any runtime that lacks packaging (e.g. minimal torch builds) and the post-install import probe would fall back to the torch GDN loop silently. Add `packaging` and `triton` to `_FLA_RUNTIME_DEPS` so the install spec list always carries them. Tests updated to assert both are now in the install command. * studio: hook transformers' fast-path gates for just-in-time FLA + causal-conv1d install The substring-based detection in this PR (`_model_wants_tilelang` / `_model_wants_causal_conv1d`) is brittle: it depends on what the user typed for the model name, not on what the architecture actually needs. Users typing custom model paths, future Qwen3.7 / non-Qwen GDN architectures, and any model whose author renamed it would silently fall back to the torch loop. The correct signal is the one transformers itself uses to gate the fast path. `transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py` does at module import time: if is_causal_conv1d_available(): from causal_conv1d import causal_conv1d_fn, causal_conv1d_update if is_flash_linear_attention_available(): from fla.modules import FusedRMSNormGated from fla.ops.gated_delta_rule import ( chunk_gated_delta_rule, fused_recurrent_gated_delta_rule, ) Wrap both gates so the first call (always at modeling import, before any forward pass) installs the matching kernel synchronously and delegates to the original function. Any model whose architecture queries those gates auto-triggers the install; models that never query them (Llama, Gemma, dense Qwen, ...) never pay the cost. Mechanics: - Split `_ensure_flash_linear_attention` and `_ensure_tilelang_backend` into `_unconditional` variants (no substring gate, retains python / torch / platform / skip-env guards) plus thin substring wrappers used by the legacy fallback path. - New `_install_fast_path_hooks(event_queue)` patches both gates on `transformers.utils.import_utils` AND sweeps `sys.modules` so any modeling file that already did `from ... import is_X` sees the wrapper (the local binding survives a module-level reassignment). - Wrappers clear the original's `lru_cache` before delegating, install on False, re-check, and short-circuit on subsequent calls. - Set `UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1` to fall back to the substring path. Verified end-to-end against `transformers.models.qwen3_5_moe`: PRE_STATE fla=False tilelang=False causal_conv1d=False HOOK_INSTALLED Hook fired for is_causal_conv1d_available; installing kernel... Installing prebuilt causal-conv1d wheel... Hook fired for is_flash_linear_attention_available; installing kernel... Installing flash-linear-attention==0.5.0 (with fla-core==0.5.0) for the fast path... Installed flash-linear-attention for the FLA fast path Installing TileLang backend (apache-tvm-ffi==0.1.9, tilelang==0.1.8)... Installed TileLang backend for FLA fast path MODELING_IMPORT_OK FAST_PATH_SYMBOLS {"chunk_gated_delta_rule": true, "fused_recurrent_gated_delta_rule": true, "FusedRMSNormGated": true, "causal_conv1d_fn": true, "causal_conv1d_update": true} POST_STATE fla=True tilelang=True causal_conv1d=True Adds 9 new tests covering: install-on-False, skip-on-True, idempotency, install-failure handling, env-disable, lru_cache clear, sys.modules rebind, missing-transformers fallback, substring fallback. Total test count is now 36 (was 27). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: address reviewer.py n=12 findings on the FLA hook path Eight issues reproduced by parallel reviewers against 6ce495a; all fixed and covered by regression tests. 45 pytest cases pass (was 36); end-to-end Qwen3.5_MoE modeling-import drill still loads all five fast-path symbols. P1 fixes: 1. TileLang loses the Qwen-family guard on the normal FLA hook path (10/12 reviewers, reproduced with allenai/OLMo-Hybrid-1B). The hook unconditionally installed tilelang for any FLA-using model. - Threaded `model_name` through `_install_fast_path_hooks(event_queue, model_name)`. - `_fla_install` now gates tilelang on `_model_wants_tilelang(model_name)` AND a successful FLA install. 2. TileLang repair `--force-reinstall` (without `--no-deps`) could replace `torch==2.12.0+cu130` with `torch==2.12.0`. Split repair into TWO steps: step 1: `--force-reinstall --no-deps apache-tvm-ffi==0.1.9` step 2: regular install of tilelang + apache-tvm-ffi Step 1 surgically downgrades the broken package; step 2 resolves missing transitive deps (z3-solver, ml-dtypes) without --force-reinstall, so it never replaces torch. 3. Hook could return True after the installer's deep import probe failed: when pip exits 0 but `import fla.modules` raises, the old wrapper re-called `original()` (transformers' metadata check) and trusted it. Refactored: - `_ensure_flash_linear_attention_unconditional(...) -> bool` - `_ensure_tilelang_backend_unconditional(...) -> bool` The wrapper now uses the installer's bool directly. 4. SSM models (Nemotron-H, Falcon-H1, Granite-H) use `lazy_load_kernel("causal-conv1d")` and never call `is_causal_conv1d_available()`, so the hook never fires for them. The orchestrator now always runs `_ensure_causal_conv1d_fast_path` outside the hook-mode if/else. P2 fixes: 5. `_rebind_in_already_imported_modules` invoked transformers' lazy module `__getattr__` (hundreds of "Accessing X from .models..." warnings, ~3.4s overhead). Switched to `module.__dict__.get(...)` which only sees real module-level bindings. 6. TileLang installed even when FLA was skipped (Torch <2.7) or failed (timeout, post-install probe failed). Now gated on the installer's bool return. 7. TileLang repair was skipped when FLA was already True but tilelang missing or apache-tvm-ffi on the broken list. Added an optional `post_available_fn` to the wrapper; the FLA hook's `_fla_post_available` runs `_ensure_tilelang_backend_unconditional` when (model wants tilelang) AND (tilelang missing OR tvm-ffi broken). 8. `_flash_linear_attention_importable()` only checks deep import, not version. Added `_flash_linear_attention_current()` that compares against the pinned `flash-linear-attention==0.5.0` / `fla-core==0.5.0`; older versions trigger `--force-reinstall --no-deps` so torch stays untouched. Helpers extracted to keep the surface tight: - `_pip_install_cmd(*args)` builds `uv pip install` or `python -m pip install` depending on uv availability. - `_run_pip(cmd, event_queue, label)` runs a pip command with timeout / failure handling and a status emission. Regression tests added: - test_hook_does_not_install_tilelang_for_non_qwen_fla_model - test_hook_does_install_tilelang_for_qwen35 - test_tilelang_repair_does_not_touch_torch_cuda_stack - test_hook_trusts_installer_bool_not_metadata - test_rebind_does_not_trigger_module_getattr - test_hook_skips_tilelang_when_fla_install_is_skipped - test_hook_runs_tilelang_repair_when_fla_already_true - test_fla_installer_force_reinstalls_when_older_version_present - test_run_training_process_eagerly_installs_causal_conv1d_in_normal_mode Existing tests updated for the new `_install_fast_path_hooks` signature and the two-step tilelang repair flow. End-to-end re-verified against transformers.models.qwen3_5_moe: PRE_STATE fla=False, hook fires for both gates, FLA + tilelang + causal-conv1d install, all 5 fast-path symbols non-None. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: fix double-install of tilelang on the FLA hook install path Backend CI surfaced a test-isolation bug introduced by the post_available_fn mechanism for finding #7. The wrapper ran `post_available_fn` in BOTH paths (install ran AND gate already True), but `_fla_install` already chains tilelang on the install path, so the post-available step then called tilelang install AGAIN. This was masked locally because tilelang was installed in the workspace venv (post_available short-circuited on `_tilelang_importable()` returning True). CI starts with no tilelang, so the second call actually fired and the mock recorded two calls. Fix: only run `post_available_fn` when the install path did NOT run. That preserves the finding #7 semantics (tilelang repair when FLA already True but tilelang missing or tvm-ffi broken) without duplicating the chained install on the gate-was-False path. Also tightened `test_hook_skips_install_when_gate_already_true` to monkeypatch `_tilelang_importable=True` and `_installed_tvm_ffi_version=0.1.9` so it stays a pure "no install at all" test regardless of the venv's actual state. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ci: retrigger Mac Studio GGUF after transient HF DNS resolve flake * 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). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ci: retrigger Windows Studio UI after transient Playwright tab-lookup flake * studio: auto-discover FLA-using model types from installed transformers Drop the hand-maintained `_TILELANG_MODEL_SUBSTRINGS` tuple (qwen3.5 / qwen3_5 / qwen3.6 / qwen3_6 / qwen3-next / qwen3_next) and derive the allowlist by scanning the installed `transformers/models/*/modeling_*.py` for `from fla.` imports. A model "wants tilelang" iff its modeling file imports an FLA op, which is the same signal `is_flash_linear_attention_available()` is the runtime test for. The scan happens once per worker subprocess and is cached for the process lifetime; an empty result (eg transformers not importable) means "no tilelang pre-install" -- the FLA runtime hook still drives the install via the gate when the loaded model actually probes it. Verified against the live installed transformers, the auto-derived set is {qwen3_5, qwen3_5_moe, qwen3_next}, with `_model_wants_tilelang` matching the HF Hub names `unsloth/Qwen3.5-2B`, `Qwen/Qwen3.5-MoE-A3B`, `mlx-community/qwen3-next-80b`, and correctly rejecting Llama, Mistral, Nemotron-H, Falcon-H1, etc. Future GDN models (Qwen3.7, OLMo-Hybrid-FA, ...) are picked up automatically once they ship in transformers; no further worker edits needed. Also trim docstrings / comments through the FLA / tilelang / HIP / hook block: constants get 1-line trailing comments, function docstrings collapse to 1-3 lines, and the fast-path-hooks banner shrinks from a 27-line block to 4 lines. The file drops from 2847 to 2630 lines without losing the load-bearing WHY notes (--no-deps protects torch; `__dict__.get` avoids lazy-module __getattr__; two-step tvm-ffi repair keeps torch off the dep graph; HIP setdefault disables FLA's TileLang dispatch even with tilelang already installed). 7 new tests (50 -> 57 total): discovery returns only FLA-using model_types; discovery cache reuse; missing transformers handled; OSError on a modeling file is non-fatal; `_model_wants_tilelang` matches real HF repo names across separator variants; empty discovery -> always False; normalization across `-`, `.`, `/`, space. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: hermetize the non-allowlist hook test against transformers 5.4.0+ transformers 5.4.0 added `olmo_hybrid` as an FLA-using model_type, so the auto-discovered allowlist now includes it -- and the test's prior choice of `allenai/OLMo-Hybrid-1B` as a "non-Qwen FLA-only" example became an allowlist member. CI on Python 3.11 / 3.13 caught this. Swap to a guaranteed-not-in-allowlist fake model_name AND patch _discover_fla_model_types to a known {qwen3_5, qwen3_5_moe, qwen3_next} set so the test stays valid as upstream transformers adds new FLA-using architectures. Renames the test to reflect the actual semantic under test: "outside-allowlist -> no tilelang". * ci: retrigger Windows Studio API after llama.cpp prebuilt staging WinError 5 flake * tests: move MLX smoke gate changes to dedicated PR #5537 The seven MLX smoke commits in this PR's history (_on_step grad_norm, max_grad_value pin, loss + round-trip gates) are unrelated to the FLA / tilelang work. They now live in #5537 so this PR's diff is limited to the studio worker installer changes. Net effect on tests/studio/run_real_mlx_smoke.py vs main: zero. * studio: friendlier install banners (drop hook / gate-name jargon) User-visible status text now reads: Installing flash-linear-attention== for faster training... Installing TileLang== for faster training... Installing causal-conv1d for faster training... Installing flash-attn for faster training... Removed the transient "Hook fired for is_flash_linear_attention_available; installing kernel..." banner — the install banner that immediately follows already tells the user what is happening, in plain English. The internal logger.info messages (server-side log) still carry the gate names + "Hook fired ..." for debugging. * [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> --- studio/backend/core/training/worker.py | 586 +++++++- .../tests/test_training_worker_flash_attn.py | 1299 ++++++++++++++++- 2 files changed, 1880 insertions(+), 5 deletions(-) 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,