unsloth/studio
danielhanchen 6ce495a42d 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).
2026-05-16 12:54:44 +00:00
..
backend studio: hook transformers' fast-path gates for just-in-time FLA + causal-conv1d install 2026-05-16 12:54:44 +00:00
frontend Stop: drop Ollama API key, clean up code execution UI (#5464) 2026-05-16 02:17:03 +04:00
src-tauri Harden Tauri release flow (#5341) 2026-05-12 20:30:20 -07:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
install_llama_prebuilt.py studio/install: strip top-level dir from repaired symlink target (#5467) 2026-05-15 15:09:50 -07:00
install_python_stack.py studio: skip flash-attn install on Blackwell GPUs (sm_100+) (#5420) 2026-05-14 18:13:50 +04:00
LICENSE.AGPL-3.0 Add AGPL-3.0 license to studio folder 2026-03-09 19:36:25 +00:00
setup.bat Final cleanup 2026-03-12 18:28:04 +00:00
setup.ps1 CI: scope GITHUB_TOKEN permissions, add MLX CI, unblock ~60 skipped tests (#5312) 2026-05-11 03:19:13 -07:00
setup.sh Route CPU-only Linux x86_64 to ggml-org/llama.cpp prebuilts (#5302) 2026-05-05 23:22:22 -07:00
Unsloth_Studio_Colab.ipynb studio: add --local to setup.sh + overlay unsloth-zoo from git main (#5252) 2026-05-02 08:51:56 +04:00