[electron-rare]

I'm setting up benchmark for vision encoders—ViT models—on Apple Silicon. I had to read through the Metal code to figure out if the fused path was being used; it is not in the documentation.

Without this information, a user has no way of knowing whether their model benefits from the fused kernel, nor how to choose a shape that does benefit from it outside of the usual cases: the fallback is silent and the result is identical—only the speed changes (measured on ViT shapes : the fused kernel runs 15 to 87% faster on the M1 Max). A lightweight ViT with head dimension 48 never benefits from it, without warning.

Motivation: The docstring describes force_fused without mentioning that the fused kernel depends on the shape of the inputs. While setting up an SDPA benchmark for typical ViT shapes (L = 197 and 257, D = 64) on Metal, I wanted to know if the fused path was actually taken, and there’s nothing in the documentation that lets you know this without reading the Metal code or forcing it with force_fused=True. To see what happens in an uncovered case, I tried D = 48: force_fused=True raises the error “no fused kernel is available: the full attention kernel supports head dims {64, 72, 80, 96, 128, 192, 256} …”, whereas D = 64 works. Small head dimensions exist in lightweight ViTs (D = 16 to 48), so this is not just a theoretical issue.

I therefore propose adding a short note to the docstring of mx.fast.scaled_dot_product_attention (python/src/fast.cpp) to state that, on the Metal backend, a fused kernel exists only for certain input shapes (head dimension, query length, mask type), that the unfused path is taken otherwise, and that, if force_fused=True is passed in this case, an error is raised with the reason.

The note refers to the error message rather than listing the dimensions: the list depends on the query length, the mask, and the GPU, and would change with the code. The error is raised in use_fallback (mlx/backend/metal/scaled_dot_product_attention.cpp), which uses the reason provided by has_fused_kernel.

I verified this with mlx 0.32.4 on M1 Max (macOS 27.0.1) and M3 Ultra (macOS 27.2): force_fused=True succeeds for D = 64 (ViT shapes, L = 197 and 257) and raises the above message for D = 48; without force_fused, D = 48 takes the unfused path and yields the same result as matmul + softmax.

Docstring only: one file, +7 lines, no tests added. A separate addition to the sdpa_bench.py benchmark (ViT shapes) may follow in another PR, currently being prepared and reviewed.

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: the patch and the note text were drafted with an AI assistant; I verified the behaviour myself on an M1 Max and an M3 Ultra; this description was written by me in French from AI-prepared notes and machine-translated to English.