{{ message }}
Repository navigation
Conversation
MPS ignores torch.nn.attention.sdpa_kernel and the torch.backends.cuda SDP toggles, so the execution probe saw flash, mem_efficient, cudnn and math all run and reported sub-quadratic attention on every Mac. Answer MPS from the torch build instead: math only before torch 2.13 or when PYTORCH_MPS_DISABLE_PREFILL_ATTENTION is set to anything but "0", flash otherwise (torch 2.13 added an MPS prefill kernel for the head_dims diffusion transformers use).
Collaborator
Author
|
@codex review |
|
Codex Review: Didn't find any major issues. Nice work! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
danielhanchen
added a commit
that referenced
this pull request
Oct 21, 2026
## Problem Five checks fail on main, so the `Repo tests (CPU, studio)`, `Repo tests (CPU, rest)` and `workflow-trigger lint` shards are red on every open PR (for example #14472, #14482, #14488, #14490 and #14494). All five reproduce on main `4d7ebf5cea`. | Test | Introduced by | Fault | |---|---|---| | `tests/security/test_lint_risky_loaders.py::test_the_tree_matches_the_baseline` | #14437 | Baseline. Two new dynamic-import sites: `import_module` in `studio/backend/core/rag/embeddings.py`, and a third `import_from_string` in `unsloth/models/sentence_transformer.py` (`_require_word_tokenizer`). Both are safe, see below. | | `tests/studio/test_tauri_branding_contract.py::test_desktop_surfaces_do_not_restore_studio_branding` | #14473 | Product. "Studio" became "Unsloth Studio" in the Markdown import errors and the Tauri webview reject message. The desktop contract allows "Unsloth" there, never "Unsloth Studio". "Unsloth" meets both #14473's rule (no bare "Studio") and this one. | | `tests/studio/test_usage_examples_model_source_contract.py::test_auto_switch_section_sits_above_the_usage_examples` | #14455 | Test. The redesign moved the usage examples above the monitor link on purpose, and its own `api-settings-tab.test.ts` pins that. | | `tests/studio/test_usage_examples_model_source_contract.py::test_monitor_can_unload_the_loaded_model` | #14455 | Test. Unloading moved from a standalone button on the page into the model picker's eject (`ApiModelLoadControls`). The `/unload` path itself is unchanged. | | `tests/studio/test_settings_compact_overflow_contract.py::test_api_monitor_entries_and_expanded_text_can_shrink` | #14455 | Test. Entry rows are now rounded cards (no `border-b`), and the prompt/reply box's `rounded-lg` became conditional, so the class literals the test matched are gone. The overflow guards themselves are still in the source. | ## Fix - **Risky loaders:** two baseline entries with reasons. - `embeddings.py` imports only the three `sentence_transformers.models` paths written out in a literal tuple, to wrap their `import_from_string` with the trust gate. - `_require_word_tokenizer` imports only after its caller has refused any `tokenizer_class` outside `sentence_transformers.*`, and only dotted identifiers. - **Branding:** "Unsloth Studio" to "Unsloth" in those seven Markdown import errors and the Rust reject message. `conversation-markdown-import.test.ts` now matches `/Unsloth Markdown/`. - **#14455 contracts:** rewritten to guard what the redesign keeps: - the auto-switch section still follows the monitor link, and the usage examples come before both; - the page passes `unloadActiveModel` to the load controls as `onUnloadActive`, and eject is withheld while `busy` (`modelLoading || unloading`, #11223's rule); - the `/unload` call still reads the internal id from status; - entry rows keep `w-full min-w-0`; - the prompt/reply box keeps `max-h-72 overflow-auto break-words`; - the model id keeps `min-w-0 break-all`. ## Verification - The four Python files: 78 passed (5 failed on main). - `studio/frontend/tests/conversation-markdown-import.test.ts`: 20 passed. - Sabotage, each run against the rewritten tests, each failing exactly one test: - removing `min-w-0` from the entry row, or `break-words` from the text box: the overflow test fails; - dropping `unloading` from `busy`, or passing `onEject` while busy: the unload test fails; - putting the auto-switch section above the monitor link: the order test fails.
danielhanchen
added a commit
that referenced
this pull request
Oct 21, 2026
…utoopt parity checks (#14522) ## Problem #14473 renamed user-facing "Studio" to "Unsloth" in the backend messages, but three tests still expect the old wording. They fail on main and on open PRs (for example #14494 and #14488): | Test | Expects | Product now says | |---|---|---| | `test_sandbox_linux.py::test_an_interpreter_symlinked_out_of_the_workdir_fails_the_call` | `Python that runs Studio` | `Python that runs Unsloth` (`sandbox_linux.py`) | | `test_diffusion_attention.py::test_a_hidden_xformers_is_not_reinstalled_in_the_same_process` | `restart Studio` | `restart Unsloth` (`diffusion_attention.py`) | | `test_diffusion_autoopt_parity.py` (two tests) | `the Studio fast step that replays decode steps from tensors`, 16 times in `fixtures/diffusion_autoopt_parity.json` | `the Unsloth fast step ...` (`diffusion_capture_safe.py`) | The rename was intended, so the tests are what is out of date. ## Fix Test only. Both assertions and the 16 recorded reasons in the parity fixture now use "Unsloth". Those 16 strings were the only "Studio" in the fixture. Nothing else in it changes. ## Verification - `test_sandbox_linux.py` and `test_diffusion_attention.py`: both target tests fail on main and pass with this change. - The one remaining local failure, `test_active_attention_backend_reads_tuple_return`, is an `ImportError` from this host's `huggingface_hub` and fails the same way on main. - `test_diffusion_autoopt_parity.py` (CPU torch 2.10, diffusers 0.41.0.dev0): - with main's fixture: 2 failed, 1 passed; - with this change: 3 passed.
[skip ci]
Collaborator
|
@codex review |
1 similar comment
Collaborator
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.

Problem
Unsloth Studio decides whether diffusion attention is sub-quadratic by running a tiny SDPA call under
torch.nn.attention.sdpa_kernel([backend])for each backend. MPS ignoressdpa_kernel(and thetorch.backends.cudaSDP toggles), so on Apple Silicon every backend "runs" and the probe reported flash / mem-efficient / cuDNN. As a resultsdpa_subquadratic_confirmed(mps)was True andsdpa_math_only(mps)was False on every Mac, although MPS SDPA before torch 2.13 always materialises the full score matrix.User-visible impact on a Mac (torch < 2.13):
diffusion.attention.math_onlywarning never fired._quadratic_attention/_attention_scaling_knownreturned the wrong answer (no current Mac code path acts on them: the image activation guard skips unified memory, learned peaks and calibrated activations are CUDA-only).Fix
MPS is no longer execution-probed. The answer comes from the torch build:
PYTORCH_MPS_DISABLE_PREFILL_ATTENTIONset to anything but"0"(torch's own rule):mathonly.flash+math. torch 2.13 added an MPS flash-style prefill kernel used for head_dim 32/64/72/80/96/128/256 with no mask or a bool / same-dtype mask, which covers the diffusion transformers Studio runs.available_sdpa_kernelsalso stops filtering the MPS answer throughtorch.backends.cuda.*_sdp_enabled(), since MPS dispatch does not honour those toggles. CUDA, ROCm, XPU and CPU paths are unchanged.Tradeoff: the probe is shape-less, so on torch >= 2.13 a model whose head_dim is outside the prefill set (or that passes a float mask of another dtype) is still quadratic but reported fused. On a Mac the only consequence is a missing diagnostic.
Validation
One SDPA call per fresh process, B=1 H=24 N=8192 bf16,
torch.mps.driver_allocated_memory()growth (M3 Max):Z-Image-Turbo transformer forward on MPS, torch 2.10, bf16 (non-Qwen family, head_dim 128):
The probe on the same machine (torch 2.10): before
["flash", "mem_efficient", "cudnn", "math"],_quadratic_attentionFalse; after["math"],_quadratic_attentionTrue.cd studio/backend python -m pytest -q tests/test_diffusion_attention.py -k mps_answersThe new parametrized test covers torch 2.10 / 2.12 / 2.13 / 2.14, the env var values
"0","","false","1", asserts no SDPA call is made on MPS, and runs withflash_sdp_enabled()off. It fails if the MPS branch is removed, the version threshold moves to 2.12, the env check is dropped or narrowed to"1", the CUDA toggle filter is applied to MPS, or 2.13+ never reports fused.test_a_hidden_xformers_is_not_reinstalled_in_the_same_processfails on this machine with and without this change.