metal: add optional polled release fence for the TP gate (~23% decode gain with --tensor-parallel)

[ryan5rdx]

This PR ports the fast cpu/gpu sync primitive in ml-explore/mlx#1773 to ds4. With this, and setting an already existing DS4 flag to increase the GPU duty cycle, TP decode performance roughly doubled in my testing.

I'm leaving this opt-in because it's still off by default in MLX/exo so there seems to be some hesitancy given it's largely undocumented, but works fine in my testing. I'm not a metal eng and all credit goes to prior art in ml-explore/mlx#1773, I've just had an LLM help port it and track down why the GPU was idle for so long during TP.

Speedup

All testing on 2x M3 Ultra. no MTP, 1M context:

coord:

 DS4_TP_NO_KEEPALIVE=1 DS4_METAL_FAST_SYNC=1  ./ds4 -m ~/Downloads/DeepSeek-V4-Flash-MXFP4Experts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-mxfp4-0731.gguf --role coordinator --tensor-parallel --transport rdma  --ctx 1048576 --listen 0.0.0.0 1234  --rdma-device rdma_en7

worker:

DS4_METAL_FAST_SYNC=1 DS4_TP_NO_KEEPALIVE=1 ./ds4 -m ~/Downloads/DeepSeek-V4-Flash-MXFP4Experts-F16HC-F16Compressor-F16Indexer-Q8Attn-Q8Shared-Q8Out-chat-v2-mxfp4-0731.gguf --role worker  --coordinator 192.168.0.14 1234 --transport rdma --ctx 1048576 --tensor-parallel

Enabled the fast fence alone(via DS4_METAL_FAST_SYNC=1) results in a decode gain of ~23%
tg2048@0 16.56t/s -> 20.44t/s

However - by setting one other (existing)flag that drastically increase GPU util in TP we can achieve a net ~2x gain, because speedup from the fast fence is magnified:

With DS4_TP_NO_KEEPALIVE=1 and DS4_METAL_FAST_SYNC=1
tg2048@0 goes from ~16.56t/s(baseline TP) to 33.12t/s <--------------------------------------!!

DS4_TP_NO_KEEPALIVE spins a dummy kernel on the GPU for power gating reasons, but in testing it drastically reduced GPU util because while there is some pause logic to stop it from spinning while real work is happening, we don't(can't?) interrupt the spin when work arrives, so we can end up with a lot of "spin" time per token. Disabling it completely resolves that. With it off - total system draw increases by ~10W during decode, imo the strongest signal that this is problematic.

With both DS4_TP_NO_KEEPALIVE=1 and DS4_METAL_FAST_SYNC=1 - system draw goes from ~60W to 88W during decode. (Closing the gap we see vs prefill at ~115W, though I don't expect we'll reach that during decode because of the nature of the workload).

With this - TP eclipses layer parallel distributed decode tg/s in my 2 node setup(even after enabling RDMA for layer parallel in #715)

cc @antirez - I'm making this all opt-in and not changing the default DS4_TP_NO_KEEPALIVE behaviour because looks like some extensive work has gone into adding it, so don't want to break an edge case I haven't run into.

[ryan5rdx]

decode/prefill appears to scale alright:
200 t/s prefill and ~24.6 t/s decode @~100k context.

ds4-bench/server/agent don't allow passing --tensor-parallel yet so this is just me testing in the repl