segment_anything: fuse attention with mx.fast.scaled_dot_product_attention

While building a torch-mlx-based SAM port (unrelated project, not proposing it here), I read through this repo's own native segment_anything implementation for reference and noticed both of its attention blocks are hand-rolled matmul+softmax+matmul rather than mx.fast.scaled_dot_product_attention:

• segment_anything/image_encoder.py's Attention.__call__ (~line 224): computes attn = (q * scale) @ k.transpose(...), optionally adds add_decomposed_rel_pos's relative-position bias, then mx.softmax(attn, axis=-1), then attn @ v.
• segment_anything/transformer.py's Attention.__call__ (~line 213, used by TwoWayTransformer in the mask decoder): the same manual pattern, without the rel-pos bias.

mx.fast.scaled_dot_product_attention's mask argument accepts an arbitrary additive array broadcast-compatible with [B, N, T_q, T_kv] (confirmed via its docstring), not just a boolean/causal mask — so the image encoder's add_decomposed_rel_pos output can be passed straight through as mask= instead of being added to a manually materialized attention matrix. The mask decoder's transformer attention has no such bias and would fuse even more directly (mask=None).

The win would be avoiding materializing the full H*W x H*W attention matrix (4096x4096 per head at the ViT backbone's fixed 64x64 patch grid for sam-vit-base) and running a separate mx.softmax pass, in favor of a single fused kernel call — the same class of fusion (mx.fast.scaled_dot_product_attention / mx.fast.layer_norm) that gave real, measured wins in a few other MLX ports I've been working on recently (depth-anything-mlx: ~1.35-1.5x from fp16 combined with this fusion; a similar suggestion for Blaizzy/mlx-vlm's Video-Depth-Anything port, #2180).

I haven't benchmarked this specific change against your repo's own implementation (I don't currently have it converted/running end-to-end locally), so I can't give you a measured number the way I'd want to before actually proposing a diff — flagging it as something worth trying rather than a verified fix. Happy to attempt the actual PR if that'd be useful, or happy to close this if it's already been considered and ruled out for some reason (e.g. numerical-stability concerns with the additive rel-pos mask at fp32 vs fp16).

Generated with Claude Code