Repository navigation
Conversation
At long context, QwenTarget decomposes prefill into smaller pieces (512 rows from 32K history, 256 rows from 64K, and 128 rows past about 128K history, per rows * (history + rows) <= 2^25) to preserve interactive responsiveness on Apple Silicon displays. Under the baseline heuristic `clamp(32 / tiles, 1, 32)`, pieces of 512, 256 and 128 rows all received only 1 split. For a 256-row piece with 4 KV heads, this dispatched only 128 threadgroups (3.37 threadgroups per core on a 38-core M2 Max), leaving the GPU execution units under-filled. On Apple7/8 (`AttentionTile::Register`), prefillSplits now targets 32 threadgroups per core, capped at the 256 slots of a 2048-row chunk: `min(256, max(maximum, (32 * gpuCores) / kvHeads)) / tiles`. On M1 Max (32 cores) and M2 Max (38 cores) with 4 KV heads, this gives: - 2048-row chunk (256 tiles): 1 split (unchanged) - 1024-row piece (128 tiles): 2 splits (was 1) - 512-row piece (64 tiles): 4 splits (was 1) - 256-row piece (32 tiles): 8 splits (was 1) - 128-row piece (16 tiles): 16 splits (was 1, measured in sweep) On lower core counts (e.g. 16 cores), target splits scale down (2 / 4 / 8), and on higher core counts (76 cores) or 2-KV-head models, the cap binds at 256 slots. The prefill workspace bound never grows on any Mac or model. Numbers: M2 Max (38 cores), Qwen3.8-27B-4bit, bf16 KV (decode-profile): - 64K context (single-session ABBA, 1408 dispatches): prefill_attention_bf16_split_sgf: 126.73 s -> 110.46 s (-12.83%, -16.27 s) prefill fused total: 456.82 s -> 440.69 s (-3.53%, -16.13 s) prefill_attention_reduce: 325.0 ms -> 399.3 ms (+74.4 ms, 0.09% of prefill) A/A run noise: 0.006% Projection control: _a within 0.003%, _r within 0.013%, _g within 0.013% - 128K context (5504 dispatches vs prompt 13 base): prefill_attention_bf16_split_sgf: 530.54 s -> 442.74 s (-16.55%, -87.80 s) prefill fused total: 1192.75 s -> 1104.17 s (-7.43%, -88.58 s) prefill_attention_reduce: 664.9 ms -> 1051.6 ms (+386.7 ms, 0.10% of prefill) Projection control: _a within 0.008%, _r within 0.004%, _g within 0.001% - Isolated step-3 kernel sweep (one 2048-row chunk, 32/core vs baseline S=1): 32K (P=512): 259.38 ms -> 219.58 ms (S=4, -15.3%) 64K (P=256): 538.50 ms -> 433.79 ms (S=8, -19.4%) 96K (P=256): 812.82 ms -> 653.87 ms (S=8, -19.6%) 128K (P=128): 1110.27 ms -> 866.72 ms (S=16, -21.9%) - Correctness: paged-attention-plan runs bf16 register prefill with 3, 10, and 32 splits against CPU reference (split path verified; PASS under shader validation). One attention layer on 64K history (2048 rows in eight 256-row pieces), S=1 vs S=8 production comparison: max |Δ| = 9.77e-4 (2^-10, 1 bf16 ulp), mean |Δ| = 8.19e-7, cosine = 0.99999990.
…ounds Add apple7PrefillAttentionPlans to execution_plans_test.cc to verify that Apple7/8 devices scale prefill splits with GPU cores for 512-, 256-, and 128-row pieces across 32, 38, and 76 cores, with 4 and 2 KV heads, and verify that the prefill attention workspace bound equals the baseline bound.
|
@paperniuk here are the three PRs together (#31, #32, #33) against 1.4.0-m1, measured overnight on my M2 Max in one session, in case it is useful for the next release. It doesn't include the Q2_0 line from abssoul yet. My agent replay (78 coding-agent turns up to 125K context, each sent twice, temperature 0), all five models I use. Prefill speed by context against 1.4.0-m1 (Δ %, + is faster; "paired" is the 120 requests whose prompt and cached prefix are identical in both runs):
Decode speed in the same replay, cycles/s (Δ %; the 32k column is only 5 requests):
So the prefill gain is #33 and grows with context, about the same on every format, and decode only moves on Bonsai (#31). Whole replay, TTFT went down about 8% on every model (MLX 4-bit 1049 -> 966 s), wall time 77-192 s less per model. The texts differ a bit between runs, so wall time is only a rough number. bench-all (up to 14K, decode tok/s at 1 / 2 / 4 lanes, Δ %; the two 1.4.0-m1 passes before and after agree within 0.4%):
Prefill up to 14K doesn't change there, as expected: #33 only acts from about 16K history. #32 shows in its own sweep only (appends of 278-288 rows): 80 ms less TTFT on UD-Q4_K_M and 40 ms on MLX 4-bit, with the Neural Engine idle before the append, as in an agent session. Same token hashes and draft acceptance in all bench-all passes on every model. |
|
Thanks, merged. It conflicts with #39 in prefillSplits, so in the 1.5 branch I kept your core count rule and applied 1.4.1's window cap after it. |
Base:
paperniuk/splashapple7-m1-1.4(branchapple7-m1-prefill-attention-splits: the change, then the test;commit c29a2f5).
This one is for long context. At long history
QwenTargetcuts a 2048-row prefill chunk into pieces so no dispatchkeeps the display waiting, and every piece of 512 rows or less got one history split. On Apple7/8 that leaves the GPU
mostly idle in prefill attention: a 256-row piece is 32 tiles x 4 KV heads = 128 threadgroups, about 3 per core on my
38-core M2 Max. With splits scaled to the core count, prefill on my M2 Max is 3.5% faster at 64K and 7.4% at 128K
(attention 13% and 17% faster), and in my agent replay prefill is 9% faster overall, 13% at 128K. Decode doesn't change.
As before, I used Claude Code to author it and tested it by hand.
What changed
prefillSplits(runtime/ops/PagedAttention.cpp): on the register tile (Apple7/8) it now targets about 32threadgroups per core,
min(256, max(32, 32 * gpuCores / kvHeads)) / tiles, clamped to 1..32 splits. The cap of 256is the slots a full 2048-row chunk already uses, so the prefill attention workspace doesn't grow on any Mac or
model. The Mpp path (Apple9 and later) is unchanged.
ExecutionPlanspasses its planned core count toprefillPlanandprefillWorkspace.prefillAttentionPieceRows), so the display guard is untouched; more splits only makeeach dispatch shorter.
execution_plans_test.ccchecks the split counts and that the workspace bound equals the old one on 32, 38 and76 cores, with 4 and 2 KV heads.
Splits on the M1 Max (32 cores) and M2 Max (38 cores) with 4 KV heads:
A 16-core chip gets half of these; on 76 cores or with 2 KV heads the 256 cap holds.
Results (M2 Max, 38-core GPU, Qwen3.8-27B MLX 4-bit, bf16 KV)
decode-profile, 64K, base / this / this / base in one session (seconds):The two base runs agree within 0.006%, and the projection kernels match within 0.013% in all four runs.
decode-profile, 128K, against an earlier base run (seconds):The projection kernels match the base run within 0.008%, so the two sessions are comparable. The 128K profile never
reaches 128-row pieces (its last chunk is at 129K history); those are only in the sweep below.
attention-sweep, one 2048-row chunk in production pieces (ms):Agent replay (78 coding-agent turns up to 125K context, each sent twice, temperature 0), MLX 4-bit,
1.4.0-m1against1.4.0-m1+ #31 + #32 + this change, on two identical M2 Max machines:On the 120 requests whose prompt and cached prefix are the same in both runs, prefill is 9.0% faster (8.7% on appends
under 1K, 9.1% on larger ones). It is more than in the full-chunk profiles because agent turns are short appends on a
long history, where attention is most of the prefill. #31 doesn't touch MLX and #32 only changes appends of 278–288
rows, so this is almost all from this change. The decode differences are within the two machines' noise.
Testing
paged-attention-plan: runs the bf16 register prefill kernel with 3, 10 and 32 splits against a CPU reference; PASSwithout and with
MTL_SHADER_VALIDATION=1. These split counts come from the planner without core counts, so theycover the kernel's split path, not the new counts.
against 8. Max |Δ| 2^-10 (one bf16 step), mean |Δ| 8.2e-7, cosine 0.9999999: the reduce sums 8 partials instead of
1.
execution-plans: PASS (CPU).q8-chunked-prefillandq8-flash-attention: PASS without and with validation.make check: CPU tests, the Python suite (1320 tests) and ruff pass. In the Metal tests, everything up togguf-moepasses under validation;
gguf-moetimes out under validation onapple7-m1-1.4itself too (the 120 s commandlimit, as perf: column tile for the MMA paths one- and two-row expert tiles #29 reports on an M1 Max), which stops
test-engine-metal, so I ran the Metal tests after itseparately under validation: all pass except one
metal-backendwatchdog check ("a command whose host stepseach ended within the timeout was given up", 0.1 s), which fails the same way on
apple7-m1-1.4; this changedoesn't touch the backend.