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
While building a torch-mlx-based SAM port (unrelated project, not proposing it here), I read through this repo's own native
segment_anythingimplementation for reference and noticed both of its attention blocks are hand-rolled matmul+softmax+matmul rather thanmx.fast.scaled_dot_product_attention:segment_anything/image_encoder.py'sAttention.__call__(~line 224): computesattn = (q * scale) @ k.transpose(...), optionally addsadd_decomposed_rel_pos's relative-position bias, thenmx.softmax(attn, axis=-1), thenattn @ v.segment_anything/transformer.py'sAttention.__call__(~line 213, used byTwoWayTransformerin the mask decoder): the same manual pattern, without the rel-pos bias.mx.fast.scaled_dot_product_attention'smaskargument 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'sadd_decomposed_rel_posoutput can be passed straight through asmask=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*Wattention matrix (4096x4096 per head at the ViT backbone's fixed 64x64 patch grid forsam-vit-base) and running a separatemx.softmaxpass, 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 forBlaizzy/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