feat(distributed): multi-host capability exchange + distributed speculative decoding (rebase of #105 onto main)

[FluffyAIcode]

What

Revives the v0.5-M1 multi-host distributed-inference milestone (originally #105, closed un-merged, 107 commits behind), rebased onto current main with conflicts resolved and tests green. Adds:

• ADR 0009 — mlx.distributed spec-decode + capability exchange.
• inference_engine/distributed/ — capability registry, gossip exchange, mlx.distributed ring probe, model-free NGramProposer, RemoteProposer + servicer, pure accept_block, DistributedSpeculativeDecoder, placement planner.
• proto proto/kakeya/v1/distributed.proto (+ regenerated stubs): CapabilityService + ProposerService.
• gRPC integration — create_grpc_server optional capability_registry/proposers/default_proposer_model_id planes (off by default); fleet flags; scripts/demo_distributed_spec_decode.py.
• Unit tests (tests/inference_engine/distributed/*) + an integration test.

Conflicts resolved (union of both sides)

ci.yaml (test+cover bridge/+distributed/), docs/adr/README.md (ADR 0009 in order), README.md (roadmap + v0.5-M1 rows), grpc_app.py (union signature), start_grpc_runtime_server.py (union kwargs).

Follow-up CI fixes (this revision)

• proto stub drift (red→green): regenerated distributed_pb2_grpc.py with the pinned grpcio-tools==1.81.1 (v0.5-M1 milestone: agent capability exchange + distributed spec decode on multi-host fleets (ADR 0009) #105 shipped 1.81.0).
• transformers 4.x/5.x compat: apply_chat_template returns a dict on 5.x (the Mac engine's runtime) → verifier.prefill raised 'str' object cannot be interpreted as an integer. Adopted the proven tokenize=True, return_dict=False convention (kv_cache_proposer.proposer.encode_chat) in the demo + the distributed integration-test fixture. Found by the real on-device runs.
• integration.yaml install (pre-existing bug): pip install -e . always failed (repo is not a pip package; runs via PYTHONPATH) — switched to pip install -r requirements.txt.
• Added on-device tooling: scripts/run_distributed_demo.sh + scripts/bench_distributed_spec_decode.py + bridge presets mlx-distributed-spec-decode-demo / -bench.

Testing — CI

• ✅ unit tests + 100% coverage (3.12), proto stub drift, proto lint, TypeScript SDK, docker build, package import smoke — all green.
• ✅ pytest -m integration on Mac M4 — the 2 distributed integration tests pass. The other 13 failures are all the documented attention_type transformers-5.x known issue (requirements.txt:14-25, legacy Qwen3 dLM, deferred to "K2.B Qwen backport," "NOT on the K3 critical path") — pre-existing, unrelated to distributed inference, only visible because the install fix let the long-broken job run (69 pass / 13 known-issue fail).

Testing — real multi-host (byte-identical to local greedy)

2-node gossip → placement (colocated=False) → remote n-gram drafts + LOCAL greedy verify → assert byte-identical. Qwen/Qwen3-0.6B, 48 tok, acceptance 0.127 (23/181) — identical across all three: GPU H200 (2-proc localhost), Mac mini (2-proc, via bridge, transformers 5.x), cross-host (verifier@cloud-vm ↔ proposer@H200 over SSH -L tunnel).

Performance comparison (scripts/bench_distributed_spec_decode.py)

verifier Qwen3-0.6B CPU/bf16 sink=4/window=64, block=4, RTT n=300:

Env	gRPC RTT p50	RTT p99	baseline	distributed	bounded-KV
Mac mini (localhost)	0.219 ms	0.264 ms	31.13 tok/s	15.76 tok/s	7.80 MB const
GPU H200 (localhost)	0.323 ms	0.539 ms	15.29 tok/s	8.47 tok/s	7.80 MB const
cross-host VM→H200 (SSH -L)	51.6 ms	140.1 ms	15.56 tok/s	7.41 tok/s	7.80 MB const

• Bounded memory (headline): at a 1072-token context the verifier's resident K/V stays at the 68-slot sink+window bound = 7.80 MB, vs 122.9 MB full-attention → 15.8× smaller, O(1) in context length (identical on all three).
• RTT: localhost ProposeBlock is sub-ms (0.22–0.32 ms); cross-host over the tunnel is ~52 ms p50 (~160× higher) — the network cost of remote drafts.
• Throughput (honest): with the model-free n-gram at low acceptance (0.097 on this prompt), distributed throughput is below local baseline everywhere — remote-draft RTT + a full block forward with mostly-rejected drafts (no CPU batching win). Throughput gains require higher acceptance / a GPU-batched verifier; this PR delivers the correctness-containment + bounded-memory guarantees, with throughput left to v0.5-GA (dLM/DFlash-over-ring, ADR 0009).

distributed_perf_comparison.md

Draft. The mlx.distributed data-plane / DFlash-over-ring is left as v0.5-GA hardening (per ADR 0009).

_{To show artifacts inline, enable in settings.}