Studio: stop the SDPA probe reporting fused attention on MPS by Lyxot · Pull Request #14488 · unslothai/unsloth · GitHub
Skip to content

Studio: stop the SDPA probe reporting fused attention on MPS - #14488

Open
Lyxot wants to merge 2 commits into
unslothai:mainfrom
Lyxot:fix/mps-sdpa-probe
Open

Lyxot wants to merge 2 commits into
unslothai:mainfrom
Lyxot:fix/mps-sdpa-probe

Conversation

@Lyxot

@Lyxot Lyxot commented Oct 21, 2026

Copy link
Copy Markdown
Collaborator

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 ignores sdpa_kernel (and the torch.backends.cuda SDP toggles), so on Apple Silicon every backend "runs" and the probe reported flash / mem-efficient / cuDNN. As a result sdpa_subquadratic_confirmed(mps) was True and sdpa_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):

  • A video generation that runs out of memory got the generic "device ran out of memory" message instead of the quadratic-attention diagnosis that names the real cause (resolution x frames, squared).
  • The load-time diffusion.attention.math_only warning never fired.
  • _quadratic_attention / _attention_scaling_known returned 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:

  • torch < 2.13, or PYTORCH_MPS_DISABLE_PREFILL_ATTENTION set to anything but "0" (torch's own rule): math only.
  • torch >= 2.13: 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_kernels also stops filtering the MPS answer through torch.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):

torch 2.10.0  any head_dim / dtype / mask, FLASH-only included   +15368 MiB  (5.0x the 3072 MiB score matrix)
torch 2.13.0  D=64/128, fp16/bf16/fp32, no mask or bool mask     ~0 MiB
torch 2.14.1  same                                               ~0 MiB
torch 2.13/14 D=48, D=160, fp32 mask with bf16 q                 +15368 MiB
torch 2.14.1  enable_flash_sdp(False)                            ~0 MiB   (toggle ignored)
torch 2.14.1  PYTORCH_MPS_DISABLE_PREFILL_ATTENTION= "" / false / 1  +15368 MiB;  "0"  ~0 MiB

Z-Image-Turbo transformer forward on MPS, torch 2.10, bf16 (non-Qwen family, head_dim 128):

1024x1024   4224 tokens  +10956 MiB
1536x1536   9344 tokens  +51356 MiB   -> 4.69x memory for 2.21x tokens (quadratic predicts 4.89x)

The probe on the same machine (torch 2.10): before ["flash", "mem_efficient", "cudnn", "math"], _quadratic_attention False; after ["math"], _quadratic_attention True.

cd studio/backend
python -m pytest -q tests/test_diffusion_attention.py -k mps_answers

The 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 with flash_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_process fails on this machine with and without this change.

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).
@Lyxot
Lyxot requested a review from danielhanchen as a code owner October 21, 2026 10:47
@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Oct 21, 2026 •

Copy link
Copy Markdown

@Lyxot

Lyxot commented Oct 21, 2026

Copy link
Copy Markdown
Collaborator Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Nice work!

Reviewed commit: 7e194480ac

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

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.
@NilayYadav NilayYadav self-assigned this Oct 21, 2026
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.
@NilayYadav

Copy link
Copy Markdown
Collaborator

@codex review

1 similar comment
@NilayYadav

Copy link
Copy Markdown
Collaborator

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Oct 21, 2026
@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Oct 21, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants