Support distributed prefill with local decode

[lobanov]

It should be possible to run generation locally on one device even if prefill is distributed. For example, it may be practical in a setup that uses DGX Spark to offload some of even all the prefill and then a Mac for generation. DGX Spark has a very fast GPU in theory, and Mac has very fast memory, so offloading prefill to the former while keeping all the generation on the latter could be the way to speed up large-context generation. This is also practical on homogeneous hardware, such as two Macs, because local decode is faster.

Conceptually it's similar to Kimi's Mooncake architecture, but at much smaller scale.

Goal: allow ds4 to pipeline prefill across two nodes while having the node owning the output layer to decode fully locally.

[lobanov]

Some back of the napkin calculation for feasibility:

Pipelining of prefills is already an existing mechanism, the only additional overhead is shipping KV cache for earlier (or all layers) to the machine doing the inference. It makes sense to do that per chunk alongside the layer activation data being pipelined anyway. If full 1M tokens context is taking 26GB (from README), then assuming context chunking by 4096 tokens, full context is 244 chunks with ~106MB per activation chunk. Practical bandwidth of 2.5GbE connection is around 280Mbytes/s with 1-2ms RTT. Activation chunk itself is around 64MB (4096 tokens × hidden size 4096 × 4 bytes FP32). Sending 170MB (chunk + KV cache) instead of 64MB over that link adds only extra 0.4-0.5s latency. Since inference is strictly feed-forward, and assuming efficient pipelining implementation, that 0.5s latency should be added only once per prefill.

This makes it feasible.

On Thunderbolt 5 vs 2.5Gbps, I'm yet to test this, but I suspect practical difference is negligible here. It's reported (from README) that setting --dist-activation-bits to a smaller value doesn't provide significant improvement despite cutting the payload size in half (if set to 16), from which I conclude that distributed inference at this scale is not bandwidth-bound. Will report back after practical testing once some hardware arrives later this week.

[lobanov]

I did some basic testing to prove feasibility over 5GbE direct link.

Setup:
Code main@ba00a8a
Model DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf

1. DGX Spark as worker node:
   ./ds4 --role worker --layers 22:output --coordinator 10.77.0.1 1234
2. MacBook M5 Max 128GB as coordinator:
   ./ds4 --role coordinator --layers 0:21 --listen 10.77.0.1 1234 --prompt-file README.md --nothink -n 100

Output:
ds4: prefill: 598.73 t/s, generation: 14.98 t/s

When switching Spark to --layers 22:42: ds4: prefill: 592.16 t/s, generation: 15.81 t/s

This so far establishes:

1. Heterogenous distributed inference works across CUDA and Metal
2. Spreading prefill across two devices is faster even on 5GbE with 1-2ms ping, it doesn't need Thunderbolt
3. Generation is slow as expected

If this idea works out, I should be able to see ~600 t/s prefills (due to pipelining) and 30+ t/s generation (due to local generation).

[lobanov]

See #343 draft PR for early implementation.

Current status: usable for evaluation, but still largely research-focused work. Generation with local decode can work with ds4 and ds4-server' frontends. There is also work-in-progress implementation of ds4-evalfrontend for quality controls. Code compiles formetalandcuda-spark` targets.

Current focus of work: evaluation of quality of local decode after distributed prefill. Official logits gate passes, but I see divergent generation trajectories. I'm focusing on longer evaluations using modified ds4-eval. After that I will look at options to speed up the distributed handoff, as I sense lots of improvement opportunities.

How to run:

1. Connect two nodes with high-speed link. I use USB-C-to-5GbE adapter on Mac side and a short Cat8 cable plugged into DGX Spark, but I have not seen traffic more than 200MB/s, so 2.5GbE with Cat6 cable should be adequate, 1GbE will be too tight unless you switch to --dist-activation-bits 16 (not tested).
2. Run controller on the node that will own early layers prefill, e.g. ./ds4 --role coordinator --layers 0:21 --listen 0.0.0.0 1234 --prompt-file speed-bench/promessi_sposi_1k.txt --ctx 200000 -n 1000 where speed-bench/promessi_sposi_1k.txt is the first 1000 lines of speed-bench/promessi_sposi.txt created by head -n 1000 speed-bench/promessi_sposi.txt > speed-bench/promessi_sposi_1k.txt, it's around 18k tokens. I use DGX Spark for this role. Only the owned layers (0-21 in this case) will be loaded into this instance.
3. Run controller on the node that will own later layers prefill and local generation, e.g. ./ds4 --role worker --layers 22:output --coordinator <coordinator_ip> 1234 --local-decode --ctx 200000. I use MacBook M5 Max 128GB for that. Note that the worker with --local-decode must have enough memory to fit the entire model plus context.

On this reference setup I'm currently getting ds4: prefill: 602.78 t/s, generation: 30.10 t/s. Prefill speed is consistent with running fully distributed prefill acrss DGX Spark and M5 Max, and with fully local generation on M5 Max. Best of both worlds!

The implementation hands off KV cache for early layers to the node owning local decode, and catches of KV cache of early layers once the decode produces tokens, so repeated generation can resume correctly.

[lobanov]

PR #401 contains implementation of this feature.