unsloth/studio/backend
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
..
assets Add Qwen3.6 inference defaults for Studio (#5065) 2026-04-16 11:42:42 -07:00
auth studio: security and hardening pass (auth rate-limit, sandbox, path containment, schema validation, headers) (#5375) 2026-05-13 06:12:18 -07:00
core studio: hook transformers' fast-path gates for just-in-time FLA + causal-conv1d install 2026-05-16 12:54:44 +00:00
loggers Studio: stop truncating long log lines as suspected base64 (#5335) 2026-05-08 13:07:18 +04:00
models studio/chat: OpenAI container picker delete reliability (#5466) 2026-05-16 01:53:13 +04:00
plugins fix(gh_client): fail fast on 401/403 auth errors instead of retrying forever (#5325) (#5329) 2026-05-08 21:57:41 +04:00
requirements studio: API external provider support for chat (OpenAI, Mistral, Gemini, Cohere, Anthropic, OpenRouter, DeepSeek, custom providers) (#4706) 2026-05-14 16:13:59 +04:00
routes studio/chat: OpenAI container picker delete reliability (#5466) 2026-05-16 01:53:13 +04:00
state unsloth run: add --enable-tools/--disable-tools server-side tool policy (#5277) 2026-05-05 12:45:15 +04:00
storage studio: API external provider support for chat (OpenAI, Mistral, Gemini, Cohere, Anthropic, OpenRouter, DeepSeek, custom providers) (#4706) 2026-05-14 16:13:59 +04:00
tests studio: hook transformers' fast-path gates for just-in-time FLA + causal-conv1d install 2026-05-16 12:54:44 +00:00
utils studio/mmproj: skip unwanted GGUF values via seek instead of read (#5431) 2026-05-14 21:57:04 -07:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
_platform_compat.py Fix Studio crash on Anaconda/conda-forge Python (#4484) 2026-03-22 05:36:55 -07:00
colab.py Fix/studio colab button message: Add fallback message for Colab Studio button when proxy URL fails (#4866) 2026-04-05 21:57:45 -07:00
main.py studio: API external provider support for chat (OpenAI, Mistral, Gemini, Cohere, Anthropic, OpenRouter, DeepSeek, custom providers) (#4706) 2026-05-14 16:13:59 +04:00
run.py studio: security and hardening pass (auth rate-limit, sandbox, path containment, schema validation, headers) (#5375) 2026-05-13 06:12:18 -07:00
startup_banner.py studio: unify Windows installer/setup logging style, verbosity controls, and startup messaging (#4651) 2026-03-30 00:53:23 -07:00