metal: pack element-wise kernels into full-width threadgroups

[czoli1976]

What

The flat, one-thread-per-element Metal kernels dispatched n threadgroups of a single thread (dispatch_thread_groups(grid = n, group = 1)). On Apple GPUs each threadgroup owns its own SIMD-group, so a 1-thread threadgroup leaves 31 of every 32 lanes idle.

This adds utils::dispatch_threads_1d, which uses dispatch_threads (non-uniform threadgroups — already used by apply_rope) to pack the same n threads into threadgroups of up to the pipeline maximum. The kernels index by thread_position_in_grid, which is unchanged, so this is a dispatch-side-only change — no kernel edits.

Converted (all index by thread_position_in_grid): element_wise, cast, copy_unicast, silu, gelu_approx, leaky_relu, the bin_op 1-row path, and the iff/select iff_generic kernel.

Left as-is: the threadgroup_position_in_grid kernels (gather, diag_gather, broadcast bin_op) — they structurally encode one threadgroup per element and need separate reworks; the broadcast bin_op already uses a real threadgroup via build_metal_grid_and_groups_for_el_wise_op.

Honest scope

The "31/32 SIMD lanes idle" framing oversells it: these ops are memory-bandwidth-bound, and on large tensors the GPU already has enough 1-thread threadgroups in flight to saturate bandwidth, so the change is neutral there. The win shows up on small/medium dispatches — which is exactly the per-token LLM-decode regime (element-wise/activation/cast over hidden_size × 1 token).

Benchmark

Added metal/examples/threadgroup_bench.rs (M-series GPU, µs/call, 200 iters; baseline = 1-thread threadgroups):

n_elements	silu base→fix	speedup	cast (f32→f16) base→fix	speedup
16 384	6.08 → 3.32	1.8×	6.85 → 2.96	2.3×
65 536	24.91 → 5.00	5.0×	13.71 → 6.15	2.2×
262 144	31.64 → 13.12	2.4×	54.60 → 10.67	5.1×
1 048 576	143.1 → 68.7	2.1×	69.1 → 48.6	1.4×
4 194 304	385 → 391	1.0×	279 → 280	1.0×
16 777 216	1541 → 1556	1.0×	1149 → 1160	1.0×

1.4–5× on small/medium element-wise dispatches; neutral on large memory-bound tensors. No regression.

Testing

All 63 tract-metal GPU tests pass (run on Apple Silicon). The change is +31/−33 in the library (the helper replaces 8 duplicated grid/group blocks); the example adds the benchmark.

Files

• metal/src/kernels/utils.rs — dispatch_threads_1d helper
• metal/src/kernels/{element_wise,bin_ops}.rs, kernels/array/{cast,copy}.rs, kernels/nn/{silu,gelu_approximate,leaky_relu}.rs — use it
• metal/examples/threadgroup_bench.rs — benchmark

🤖 Generated with Claude Code

[kali]

⚠️⚠️⚠️ Just rebased! ⚠️⚠️⚠️

[github-actions]

✅ Bench vs main — no regressions

Reference: main nightly, latest 2026-07-01 (0d old) · PR 3900cab4d · ran on apple-m1-max, i9-11900kb_rtx-4060, jetson-orin-nx · 727 metrics compared

Speed — evaltime · prefill · decode

no inference-speed regressions

lower is better except prefill/decode (tok/s) · adaptive thresholds (max(floor, k×noise) vs the series' own history) · single-shot vs nightly reference · full report → run

[kali]

@czoli1976 did we not expect to observe speedups on metal ?

[czoli1976]

will get back to this this afternoon or tomorrow, a bit swamped on other internal projects Il giorno gio 2 lug 2026 alle ore 08:24 Mathieu Poumeyrol < ***@***.***> ha scritto:

…

*kali* left a comment (sonos/tract#2365) <#2365 (comment)> @czoli1976 <https://github.com/czoli1976> did we not expect to observe speedups on metal ? — Reply to this email directly, view it on GitHub <#2365?email_source=notifications&email_token=APL2Z6X3DNAE463A4TNSJKL5CYE3HA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBWGMZDEOJZGU42M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-4863229959>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/APL2Z6RBQVHE22LH4Y4NFS35CYE3HAVCNFSNUABEKJSXA33TNF2G64TZHM4TSNJWGEZDCMZ3JFZXG5LFHM2DMNJVGEZDSMZYG2QXMAQ> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/APL2Z6T7QU4AISPB5ZCSCZD5CYE3HA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBWGMZDEOJZGU42M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/APL2Z6RI2RGCBF74G7U66N35CYE3HA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBWGMZDEOJZGU42M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

-- Best Regards Ckristian Zoli Email: ***@***.***

[kali]

@czoli1976 moving this one to draft, because it's all green and i fear i will merge it.