Repository navigation
perf: stage the Apple7/8 register attention tile's queries in threadgroup memory - #24
Conversation
|
@sady4850 @paperniuk I measured #24 on my M2 Max (38-core GPU) against apple7-m1-1.4. Outputs match bit for bit in attention-sweep on every case I ran (2K to 64K history, int8 and bf16 KV). attention-sweep, 27B shape, attention kernel only, bf16 KV (A = apple7-m1-1.4, B = pr-24):
int8 KV is similar, 10-20% faster. I only ran the 27B shape, not 35B. Full 27B decode cycle (decode-profile, MLX 4-bit, ABBA), 32K prompt, fused cycle ms:
At 2K it is flat, attention is only ~3% of the cycle there. GPU prefill of 32K tokens goes 199.7 s -> 194.1 s (-2.8%). So it helps exactly the long context agent sessions I care about. q8-flash-attention and q8-chunked-prefill pass with and without MTL_SHADER_VALIDATION=1. BTW your test fix in 03bcee2 is the same one I had on my side, so I will drop mine. Great job, thank you! |
|
@ghisguth thanks for checking on M2 Max, I only have M1 here so this is very useful. good to see it helps on long context, prefill is next thing I want to look at. off-topic question, I am choosing harness for long local agent sessions. from your writeup I see Qwen Code was your pick for some time, and now you file issues in oh-my-pi with 27B on 262K window. did you fully move to omp? what made the switch, compaction/prefix cache behaviour or something else? |
|
@sady4850 you are welcome, it was easy to run and nice to see it helps on M2 Max too. Regarding oh-my-pi and qwen-code, I was back and forth with opencode and qwen-code. I liked qwen-code more, but I found that qwen-code is less reliable, tool calling was failing slightly more often. I tried to move to pi agent, and actually made some progress tuning it for myself. E.g. implemented plugins to save/load context to/from disk in llama.cpp when I switch models. I liked it, and then I tried oh-my-pi, and it really surprised me. It is fast, a lot of things done out of the box. And it has great context management. With Splash + oh-my-pi + Qwen 27B I was able to finish one project, which both AGY and Claude Code were refusing to do. It took one week of running agent almost non stop, but it completed the work. This was real test for omp for me, and it worked really-really well. So I would probably stay with oh-my-pi for now, and as backup will try opencode. |
|
Nice one, thanks! Confirmed on M1 Max (32-core GPU), digests identical in every case.
BF16 KV gains the same: 27B verify 1.18x / 1.21x at 2K / 24K on 1 lane, 1.31x / 1.23x on 4 lanes. End to end on
|
…port) into the 1.4 port Conflicts: - paged_attention.metal, rope.metal: 1.4's own versions, which already hold #24 and the precise RoPE tables. - LinearGguf.cpp, linear_plan_test.mm: #27's one-pass check for the F32 register tile next to 1.4's 16-row tile check; #27's one-pass scratch with LinearPlan::table(). - gguf_projection_test.mm: #27's register one-pass test where the device is not Apple9, before 1.4's prefill. - gguf_projection_benchmark.mm: register_gu rows next to 1.4's staged16 rows; the staged rows take one pass only from a staged policy.
With the sweep fixed (#23) I looked at the
_sgfattention tile, because at agent contexts it is a big part of the M1 cycle: decode-profile on M1 Pro, 27B, B1 at 24K history,verify_attention_q8_split_sgfis 49.7 of 226 ms (22%), and prefill attention 15% of a 24K prefill. Per FLOP it ran ~2x below the linear tiles next to it.Its pipelines report
maxTotalThreadsPerThreadgroup384, the lowest of the hot kernels (the decode linear tile is 832, GDN 448), so register pressure looked like the limit. Each lane keepsuint queries[4][8](32 registers) for the whole page loop, next tooutput[32]. This change moves the queries to threadgroup memory, step-major[b * 8 + t][lane]: every lane reads only the words it wrote, so no barrier and no bank conflicts. 24 KB per threadgroup for G = 6, 32 KB for kv2_g8. The arithmetic is untouched, so output is bit-identical (sameattention-sweepdigests in every case).maxTotalThreadsPerThreadgroupstays 384 (outputstill dominates), but it is faster anyway.attention-sweep, M1 Pro (16-core GPU), 1 lane,
--compare-metallib, medians of 21:In the model it helps more (decode-profile, 27B, 24K prompt, fp32 state):
Agent replay (the public junie_replay from erokhins/splash: warm-up 14.5K + 11 append-only requests growing to 31.7K,
mlx-community/Qwen3.8-27B-4bit,--max-context 64K --max-memory 21G, fresh server each run, ABBA, on AC):All four runs produce the same 1164 tokens with the same acceptance, so token-identical end to end. Decode +6% over the session, ~+10% at 30K; cold 14.5K prompt unchanged (~127 s), short histories don't spend much in attention.
Short context, no regression: your serving methodology with my prompts (256 tokens, temperature 0, reasoning off, ABBA, 2 reps), undisturbed runs before / after: technical 21.1 / 21.3, python_code 27.0 / 27.3, typescript_code 22.0 / 22.2, repeat 45.5 / 45.7 tok/s, same drafted/accepted counts on every prompt.
Tests:
q8_flash_attention_metal_testpasses withMTL_SHADER_VALIDATION=1(it runs the_sgftiles on Apple7/8),attention_sweep_cli.shpasses, build clean. Two things fail onapple7-m1-1.4with the base metallib too, so not from this change:paged-attention-plan(RoPE table precision checks) andq8_chunked_prefill_metal_test.mm:334does not compile (missing field 'block_stride' initializer).Not tried on M1 Max / M2. Next I want to look at
output[32], the rest of the register pressure (splitting D between two simdgroups with a score exchange), if you think it is worth it.