Skip to content

Enabling fused attention for head dim 128#1899

Merged
jagrit06 merged 7 commits into
mainfrom
attn128
Feb 26, 2025
Merged

Enabling fused attention for head dim 128#1899
jagrit06 merged 7 commits into
mainfrom
attn128

Conversation

@jagrit06

@jagrit06 jagrit06 commented Feb 24, 2025

Copy link
Copy Markdown
Member

Proposed changes

  • Share KV smem
  • Enable head dim 128 fused attention

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@jagrit06

jagrit06 commented Feb 24, 2025 •
edited
Loading

Copy link
Copy Markdown
Member Author

Adding some numbers here for M2 Ultra:

Before, 128 routing disabled:

  B,   qsl,   ksl, hdim, n_qh, n_kvh, tp,   dtype, t_unfs, t_fuse, diff%
  1,    32,    32,   64,   32,    32,  0, float16,  0.026,  0.009, +189.68%
  1,    64,    64,   64,   32,    32,  0, float16,  0.017,  0.010, +64.14%
  1,   128,   128,   64,   32,    32,  0, float16,  0.020,  0.013, +59.84%
  1,   256,   256,   64,   32,    32,  0, float16,  0.028,  0.018, +54.00%
  1,   512,   512,   64,   32,    32,  0, float16,  0.064,  0.049, +29.63%
  1,  1024,  1024,   64,   32,    32,  0, float16,  0.174,  0.151, +15.30%
  1,  2048,  2048,   64,   32,    32,  0, float16,  0.620,  0.526, +17.70%
  1,  4096,  4096,   64,   32,    32,  0, float16,  2.302,  1.986, +15.95%
  1,  1024,  1024,   80,   32,    32,  0, float16,  0.242,  0.184, +31.04%
  1,  2048,  2048,   80,   32,    32,  0, float16,  0.863,  0.641, +34.67%
  1,  4096,  4096,   80,   32,    32,  0, float16,  3.280,  2.432, +34.85%
  1,  1024,  1024,  128,   32,    32,  0, float16,  0.270,  0.270, -0.21%
  1,  2048,  2048,  128,   32,    32,  0, float16,  0.963,  0.965, -0.14%
  1,  4096,  4096,  128,   32,    32,  0, float16,  3.685,  3.685, -0.00%
  1,    32,    32,   64,   32,    32,  1, float16,  0.016,  0.010, +64.55%
  1,    64,    64,   64,   32,    32,  1, float16,  0.017,  0.010, +63.29%
  1,   128,   128,   64,   32,    32,  1, float16,  0.019,  0.013, +54.24%
  1,   256,   256,   64,   32,    32,  1, float16,  0.029,  0.020, +49.04%
  1,   512,   512,   64,   32,    32,  1, float16,  0.065,  0.050, +29.80%
  1,  1024,  1024,   64,   32,    32,  1, float16,  0.174,  0.151, +15.31%
  1,  2048,  2048,   64,   32,    32,  1, float16,  0.620,  0.526, +17.71%
  1,  4096,  4096,   64,   32,    32,  1, float16,  2.302,  1.986, +15.92%
  1,  1024,  1024,   80,   32,    32,  1, float16,  0.243,  0.185, +31.02%
  1,  2048,  2048,   80,   32,    32,  1, float16,  0.864,  0.643, +34.32%
  1,  4096,  4096,   80,   32,    32,  1, float16,  3.292,  2.447, +34.50%
  1,  1024,  1024,  128,   32,    32,  1, float16,  0.270,  0.271, -0.26%
  1,  2048,  2048,  128,   32,    32,  1, float16,  0.964,  0.965, -0.13%
  1,  4096,  4096,  128,   32,    32,  1, float16,  3.685,  3.685, -0.01%

Before, 128 routing enabled:

(base) jdigani@mlx-m2-macstudio-4 mlx % python tmp/bench_sdpa.py
  B,   qsl,   ksl, hdim, n_qh, n_kvh, tp,   dtype, t_unfs, t_fuse, diff%
  1,    32,    32,   64,   32,    32,  0, float16,  0.017,  0.009, +81.53%
  1,    64,    64,   64,   32,    32,  0, float16,  0.017,  0.010, +66.47%
  1,   128,   128,   64,   32,    32,  0, float16,  0.021,  0.013, +55.93%
  1,   256,   256,   64,   32,    32,  0, float16,  0.029,  0.019, +52.56%
  1,   512,   512,   64,   32,    32,  0, float16,  0.064,  0.049, +30.18%
  1,  1024,  1024,   64,   32,    32,  0, float16,  0.174,  0.151, +15.44%
  1,  2048,  2048,   64,   32,    32,  0, float16,  0.620,  0.526, +17.86%
  1,  4096,  4096,   64,   32,    32,  0, float16,  2.303,  1.985, +15.99%
  1,  1024,  1024,   80,   32,    32,  0, float16,  0.242,  0.184, +31.32%
  1,  2048,  2048,   80,   32,    32,  0, float16,  0.862,  0.640, +34.74%
  1,  4096,  4096,   80,   32,    32,  0, float16,  3.281,  2.432, +34.88%
  1,  1024,  1024,  128,   32,    32,  0, float16,  0.269,  0.304, -11.26%
  1,  2048,  2048,  128,   32,    32,  0, float16,  0.964,  1.072, -10.07%
  1,  4096,  4096,  128,   32,    32,  0, float16,  3.685,  4.111, -10.36%
  1,    32,    32,   64,   32,    32,  1, float16,  0.016,  0.010, +65.75%
  1,    64,    64,   64,   32,    32,  1, float16,  0.017,  0.010, +61.03%
  1,   128,   128,   64,   32,    32,  1, float16,  0.020,  0.013, +54.54%
  1,   256,   256,   64,   32,    32,  1, float16,  0.029,  0.019, +49.72%
  1,   512,   512,   64,   32,    32,  1, float16,  0.064,  0.050, +29.03%
  1,  1024,  1024,   64,   32,    32,  1, float16,  0.174,  0.151, +15.43%
  1,  2048,  2048,   64,   32,    32,  1, float16,  0.620,  0.526, +17.82%
  1,  4096,  4096,   64,   32,    32,  1, float16,  2.302,  1.986, +15.94%
  1,  1024,  1024,   80,   32,    32,  1, float16,  0.243,  0.185, +31.00%
  1,  2048,  2048,   80,   32,    32,  1, float16,  0.865,  0.643, +34.51%
  1,  4096,  4096,   80,   32,    32,  1, float16,  3.292,  2.447, +34.54%
  1,  1024,  1024,  128,   32,    32,  1, float16,  0.269,  0.304, -11.36%
  1,  2048,  2048,  128,   32,    32,  1, float16,  0.964,  1.072, -10.11%
  1,  4096,  4096,  128,   32,    32,  1, float16,  3.685,  4.111, -10.35%

After, with 128 routing enabled:
  B,   qsl,   ksl, hdim, n_qh, n_kvh, tp,   dtype, t_unfs, t_fuse, diff%
  1,    32,    32,   64,   32,    32,  0, float16,  0.016,  0.009, +82.04%
  1,    64,    64,   64,   32,    32,  0, float16,  0.017,  0.009, +90.65%
  1,   128,   128,   64,   32,    32,  0, float16,  0.021,  0.012, +67.30%
  1,   256,   256,   64,   32,    32,  0, float16,  0.030,  0.018, +69.45%
  1,   512,   512,   64,   32,    32,  0, float16,  0.064,  0.045, +42.37%
  1,  1024,  1024,   64,   32,    32,  0, float16,  0.174,  0.139, +25.74%
  1,  2048,  2048,   64,   32,    32,  0, float16,  0.620,  0.475, +30.35%
  1,  4096,  4096,   64,   32,    32,  0, float16,  2.304,  1.797, +28.27%
  1,  1024,  1024,   80,   32,    32,  0, float16,  0.241,  0.175, +37.84%
  1,  2048,  2048,   80,   32,    32,  0, float16,  0.862,  0.608, +41.76%
  1,  4096,  4096,   80,   32,    32,  0, float16,  3.282,  2.312, +41.96%
  1,  1024,  1024,  128,   32,    32,  0, float16,  0.269,  0.268, +0.68%
  1,  2048,  2048,  128,   32,    32,  0, float16,  0.964,  0.947, +1.74%
  1,  4096,  4096,  128,   32,    32,  0, float16,  3.687,  3.640, +1.29%
  1,    32,    32,   64,   32,    32,  1, float16,  0.017,  0.009, +87.16%
  1,    64,    64,   64,   32,    32,  1, float16,  0.017,  0.010, +71.19%
  1,   128,   128,   64,   32,    32,  1, float16,  0.019,  0.012, +61.74%
  1,   256,   256,   64,   32,    32,  1, float16,  0.030,  0.019, +57.69%
  1,   512,   512,   64,   32,    32,  1, float16,  0.064,  0.046, +40.44%
  1,  1024,  1024,   64,   32,    32,  1, float16,  0.175,  0.139, +25.75%
  1,  2048,  2048,   64,   32,    32,  1, float16,  0.620,  0.476, +30.23%
  1,  4096,  4096,   64,   32,    32,  1, float16,  2.304,  1.797, +28.21%
  1,  1024,  1024,   80,   32,    32,  1, float16,  0.242,  0.177, +37.26%
  1,  2048,  2048,   80,   32,    32,  1, float16,  0.865,  0.613, +41.05%
  1,  4096,  4096,   80,   32,    32,  1, float16,  3.293,  2.333, +41.16%
  1,  1024,  1024,  128,   32,    32,  1, float16,  0.270,  0.268, +0.78%
  1,  2048,  2048,  128,   32,    32,  1, float16,  0.964,  0.947, +1.77%
  1,  4096,  4096,  128,   32,    32,  1, float16,  3.687,  3.639, +1.30%

@jagrit06 jagrit06 requested review from angeloskath and awni February 24, 2025 23:44
@awni

awni commented Feb 25, 2025

Copy link
Copy Markdown
Member

Very nice! If I understand correctly it's slightly faster than the baseline for head_dim=128 as opposed to a bit slower.

awni
awni reviewed Feb 25, 2025
View reviewed changes
Comment thread mlx/backend/metal/kernels/steel/attn/mma.h Outdated
// }
// }
// }
// }

@awni awni Feb 25, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason to keep that around?

awni
awni approved these changes Feb 25, 2025
View reviewed changes

@awni awni left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me!! 🚀

@jagrit06

jagrit06 commented Feb 25, 2025

Copy link
Copy Markdown
Member Author

Very nice! If I understand correctly it's slightly faster than the baseline for head_dim=128 as opposed to a bit slower.

Yes! Its a few percent faster, and saves memory over the unfused baseline

@angeloskath

angeloskath commented Feb 25, 2025

Copy link
Copy Markdown
Member

I tested it on the flux example and the memory savings are super nice 🚀

This is flux-schnell generating 4 images of size 1024x512. No speed regression noticed.

Before

No quantization
-----------------
Peak memory used for the text:       9.239GB
Peak memory used for the generation: 24.693GB
Peak memory used for the decoding:   7.607GB
Peak memory used overall:            24.693GB

Quantization
-------------
Peak memory used for the text:       3.703GB
Peak memory used for the generation: 8.902GB
Peak memory used for the decoding:   7.607GB
Peak memory used overall:            8.902GB

After

No quantization
----------------
Peak memory used for the text:       9.239GB
Peak memory used for the generation: 24.008GB
Peak memory used for the decoding:   7.607GB
Peak memory used overall:            24.008GB

Quantization
-------------
Peak memory used for the text:       3.703GB
Peak memory used for the generation: 8.094GB
Peak memory used for the decoding:   7.607GB
Peak memory used overall:            8.094GB

@jagrit06 there is a compilation warning for the constexpr in the kernel. If there is a way to suppress that it might be a good idea.

@awni

awni commented Feb 25, 2025

Copy link
Copy Markdown
Member

Yes let's get rid of that warning. I am fine with allowing c++17 stuff in metal. You pass -Wno-c++17-extensions to the Metal compiler to suppress the warnings about it.

@jagrit06 jagrit06 merged commit 89d3270 into main Feb 26, 2025
@jagrit06 jagrit06 deleted the attn128 branch February 26, 2025 18:02
faisalmemon pushed a commit to faisalmemon/mlx that referenced this pull request Oct 30, 2025
@jagrit06 @faisalmemon
Enabling fused attention for head dim 128 (ml-explore#1899)
a45cc43 
* Share KV smem

* Fix bfloat error

* Unroll O = S @ V loop

* Perf upgrade

* Remove commented out function

* Add -Wno-c++17-extensions flag to metal flags

* Add -Wno-c++17-extensions flag to metal extension flags
giaki3003 added a commit to giaki3003/mlx that referenced this pull request Jun 29, 2026
@giaki3003 @claude
fix(steel attn): apply the BD==128 simdgroup barriers to all BD>=128 …
1175b4d 
…(fixes hd256 garbage)

The fused steel attention kernel guards its V-tile load + MMA accumulation with
two simdgroup_barrier(mem_none) calls, but only 'if constexpr (BD == 128)'. Those
barriers were added by upstream ml-explore#1899 ('Enabling fused attention for head dim
128') as a CORRECTNESS fix for the large-head-dim reuse hazard: a single Vtile
fragment is reloaded and MMA'd every (iq,id,ik) iteration, so without the
barriers the next load races the prior MMA's read.

head_dim=256 was instantiated but never ROUTED to the fused path (old use_fallback
forced it unfused), so the barriers were never extended to it. PR ml-explore#3660's routing
(key_sequence_length > 16384 -> fused for hd 192/256) finally exercises BD=256,
and the missing barriers race -> wrong attention -> multilingual word-salad on
long-context prefill.

Fix: gate on 'BD >= 128' so 128/192/256 all get the barriers (64/80 keep their
working barrier-free path). Adding a simdgroup execution barrier can only remove a
race, never introduce one.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@awni awni awni approved these changes
@angeloskath angeloskath Awaiting requested review from angeloskath

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@jagrit06 @awni @angeloskath