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).
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.
* studio: skip flash-attn install on Blackwell GPUs (sm_100+)
Dao-AILab does not publish prebuilt flash-attn wheels for sm_100, sm_120,
or sm_121, and the older-arch wheels fail to load on Blackwell. Add a
shared has_blackwell_gpu() helper and gate both the install-time
(install_python_stack._ensure_flash_attn) and runtime
(worker._ensure_flash_attn_for_long_context) paths on it. Detection uses
nvidia-smi --query-gpu=compute_cap, which works on Linux and Windows.
* test: stub has_blackwell_gpu in pre-existing runtime flash-attn tests
prefers_prebuilt_wheel and falls_back_to_pypi exercise the install
paths that the Blackwell guard now short-circuits. Make them explicit
about non-Blackwell so they pass on real Blackwell hosts.
* studio: cache has_blackwell_gpu, skip Blackwell warning under NO_TORCH
- Wrap has_blackwell_gpu in functools.lru_cache so repeated calls in a
single process avoid redundant nvidia-smi spawns. Tests clear the
cache via setup_method/teardown_method.
- In _ensure_flash_attn, run the NO_TORCH short-circuit before the
Blackwell check so GGUF-only users (who never install torch anyway)
do not see a Blackwell warning. Blackwell check still runs above the
IS_WINDOWS / IS_MACOS gates so Blackwell-on-Windows users still see
the explicit reason rather than a silent OS skip.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* test: add has_blackwell_gpu to mlx worker test wheel_utils stub
test_mlx_training_worker_config loads worker.py against a hand-rolled
utils.wheel_utils stub. Adding has_blackwell_gpu to the stub symbol
list so worker's import line resolves.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* qwen3.6 unsloth studio support
* Add qwen3.6 causal-conv1d detection
* Update model_mappings.py
moved qwen3.6-27B to thinking train on completion template
* [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] Install flash attn at setup time for linux
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* cleanup changes
Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Test cases
* wheel_utils: narrow url_exists exceptions and log at debug level
---------
Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
