Proposal: Qwen3.8-27B at TP=2 on CUDA (two hosts), EXL3 first, with DFlash2 (proposal and gate) #1468

Proposal

Serve the dense Qwen3.8-27B across two CUDA hosts (two GB10s) at TP=2 on the native engine, starting with EXL3 checkpoints (#1465), with DFlash2 drafts. The engine has multi-rank serving for GLM-5.3-Flash and DeepSeek-V4.1, but not for the Qwen family. I'm opening this before the pull request, as CONTRIBUTING asks for an arithmetic change, so we can agree on the gate. It would stack on #1465 (and the EXL3 speed draft after it).

Design (working on our two GB10s)

• Shares. Each rank holds 1/world of every layer's heads and of the MLP width.

  • Column shares: q/k/v, z, gate, up and the head.
  • GDN's fused in_proj_qkv share is three column ranges, one each for its q, k and v heads.
  • Row shares: out_proj, o_proj and down_proj. They write fp32 partials, all-gathered and summed in rank order into bf16.
  • The head's columns are gathered into whole logit rows.
  • Shares are pointer offsets into the whole EXL3 trellis.

• Serving. Rank 0 serves HTTP, drafts, samples and admits. Before each engine step it broadcasts it to the other ranks, which replay it on their own caches:

  • caches made, reset and dropped;
  • forwards, with their keep kind;
  • kept paths;
  • prompt states saved, restored and freed;
  • the drafter head.

  A step longer than the fixed message (a prompt chunk's ids) sends the rest in a second broadcast of its length. The ranks fingerprint their settings, build and checkpoint before the load, and refuse a mismatch by field. They use the existing tp_group, RoCE gathers and tp_watch liveness lines.

• DFlash2. The drafter runs on rank 0. Its head is split by columns on every rank, and each rank's best tokens are merged on rank 0 (logit descending, then id ascending).

• Refused under TP for now: MoE, --kv-cache fp8, KV pages, --slide and --vision.

The arithmetic change

At TP=2 the row-parallel projections round once, after an fp32 sum of the ranks' partials in rank order. A single rank rounds its whole K at once. So TP=2 replies can differ from TP=1. TP=2 is deterministic: the same build and settings give the same bits, and its own window and tree equal its step.

Measured so far (two GB10s, 200 Gb RoCE link; same checkpoint and drafter as #1465)

• Exactness at TP=2.
  • tf-qwen35-cuda ... tp over 2,048 ids: window logits equal step logits byte for byte, and so do tree logits.
  • The column-split drafter head picks the same best tokens as the whole-row head, 1,024/1,024.
  • Served: alone and serial 20/20 equal (drafts on, temperatures 0 and 1.0), 0 failed.
  • Fresh, cached repeat and "draft": false give one reply hash on an 8k prompt.
  • A two-turn 38.7k chat resumes at 36,864 cached tokens.
  • An 80k-token needle is found.
• Fidelity against TP=1, teacher-forced over 2,048 code ids: mean KL 6.3e-4, p99 4.6e-3, max 2.2e-2. Top-1 agreement is 98.88%, or 99.30% without near-ties (gap < 0.05).
• Speed, TP=2 against TP=1 on the same build:
  • bench_openai drafts on: fibonacci 110.0 against 75.3 tok/s, chat 72.3 against 46.4.
  • 4 streams: code 267.7 against 204.6, chat 189.3 against 133.8.
  • Cold prefill: 1,109 tok/s at 4.5k and 1,018 at 69k, against 895 and 776.

Proposed gate

On two GB10s, then any two-GPU host pair you can run:

• TP=2 window and tree equal TP=2 step, byte for byte.
• Drafted equals "draft": false, resumed equals fresh, concurrent equals solo, and cancellation leaves the next request valid. All at TP=2.
• Fidelity against TP=1, teacher-forced on public code and chat ids, with every top-token difference and both margins listed.
• A settings mismatch between ranks is refused before the load; a rank that dies stops the others.
• The new kernels (the rank-order sum, and placing gathered columns) registered with the kernel contract, with an FP64 reference for the rank-order sum.

Questions:

1. Rank placement: the GLM way (TF_GLM_RANK/TF_GLM_WORLD and NCCL_COMM_ID), or --tp/--rank/--master like the HIP backend?
2. Should MLX affine checkpoints share the same share layout in this PR, or follow it?