perf: faster pipeline decode: TPS 10%-40% and Layer loading 60%+ (token relay + bandwidth-aware placement)

[jgawronek]

Multi-node pipeline decode currently pays three per-token costs that this PR removes or reduces:

1. Token-relay decode — during decode, every rank all_gathers the full hidden state and redundantly runs final norm + lm_head + sampling. With this change only the last pipeline rank samples; the sampled token id is circulated to all ranks with a single-int collective, and the other ranks never compute lm_head at all (it drops out of the lazy graph). Requests with logprobs automatically fall back to the legacy lockstep path so rank 0 keeps real logits.

2. Bandwidth-aware layer allocation — layers were split proportionally to available RAM, which overloads slow nodes in heterogeneous clusters (per-token decode latency is the sum of per-stage times, and stage time is bytes-read / memory-bandwidth). Nodes now report their GPU name (nvidia-smi on Linux, system_profiler on macOS), placement estimates per-chip memory bandwidth from a lookup table, and loads the fastest nodes to their memory caps first. Unrecognised hardware falls back to the existing RAM-proportional split, so behaviour is unchanged for clusters the table doesn't know.

3. Link-speed-aware ring host selection — Linux interfaces were all classified "unknown" (interface typing shelled out to macOS-only networksetup), so ring links between Linux nodes were chosen arbitrarily. Interfaces are now classified via sysfs with their negotiated link speed, and ring neighbours prefer the fastest measured link — e.g. a DGX Spark's 200GbE ConnectX port over its management ethernet. macOS behaviour (Thunderbolt first) is unchanged.

Also fixes a pipeline deadlock: shards of hybrid SSM/attention models (Qwen3-Next, Qwen3.5, Nemotron-H) can restore different KV prefix-cache lengths on different ranks — SSM shards need state snapshots, attention-only shards don't. Ranks then prefill different token counts and hang the ring's send/recv chain on the second request of a conversation. Ranks now gather their hit lengths and discard the hit everywhere when they disagree.

Adds per-token decode communication timing (recv / send / gather) logged at debug level every 64 tokens for future tuning.

Results

Measured on a heterogeneous 4-node cluster (M3 Ultra Mac Studio 256GB + 2× DGX Spark GB10 + RTX 3090) and smaller mixes:

• Generation TPS: +10% to +40% depending on model and cluster composition
• Layer distribution / model load time: over 2× faster
• Multi-turn chat on hybrid SSM models no longer deadlocks on heterogeneous splits

Test plan

• uv run pytest — 448 passed, including new two-process distributed ring tests covering token relay (last rank samples, all ranks converge on its token, non-last ranks skip lm_head)
• New unit tests for throughput-based allocation, bandwidth table, and link-speed ring preference
• uv run basedpyright and uv run ruff check clean
• Soak-tested on the 4-node heterogeneous cluster above (multi-turn chat + exo-bench, Qwen3.5-397B-A17B-8bit and Qwen3-Next-80B)
• Linux/3090, 2x DGX Spark, MacStudio