RFC: a layer-range (pipeline) layout across mixed nodes: one hidden row per token over Ethernet, GB10 and Mac in one model #2005

Sections 1-5 cite main at 933b969fa and Starlight (#1922) at 39ecd68b3; section 6 cites main at 071cefe08.

Conventions:

• cluster/… is zig/src/cluster/… and sl/… is starlight/….
• Code claims are cited by file:line. Everything marked (guess) is an estimate, not measured.
• The cluster evidence is from a read-only audit of the cluster runtime: CLUSTER-AUDIT.md.

1. Problem

• Tensor parallel needs fast links and the same bits on every rank.

  • Every layer all-reduces partial sums, so one decode step makes dozens of round trips.
  • The cluster refuses Metal and CUDA in one tensor-parallel group: "Metal and CUDA kernels give different bits" (cluster/config.zig:288-299).
  • Starlight's tensor_parallel is CUDA (NCCL) or ROCm only (sl/internal/control/wire/ranks.go:13-34; sl/docs/fleet.md:105-111).

• Measured: mixed tensor parallel is slower than either machine alone (ThinkOffApp, in this thread). Qwen3.8-27B MLX 4-bit on an M5 Max and a GB10 over a direct ConnectX-7 cable, 128 in / 64 out, median of 5, decode tok/s at 1 / 4 / 16 streams, no drafter:

  placement	1	4	16
  M5 Max alone	24.8	94.4	279
  GB10 alone	14.4	51.8	147
  mixed TP, FFN split only	17.0	—	113.8
  mixed TP, attention, GDN and FFN split (2 host reductions a layer)	15.7	43.0	84.8

  These are sequential comparisons, not a controlled A/B, and the per-layer synchronization time was not isolated. Mixed TP's value today is fitting a model that one machine cannot hold.

• So a GB10 and a Mac, or two nodes on 1/10 GbE, can't share a model today at a useful speed.

• A layer split instead sends one hidden row per token per stage boundary. That fits Ethernet, and each stage runs its own backend.

2. The layout

A stage is one node (or one tensor-parallel group of identical nodes) holding layers a..b of one model:

• rank 0 serves: it holds the lane host, the embeddings and the first layer range, and also the final norm, the head, the sampler and any drafter;
• the other ranks hold their layer ranges only, and each stage keeps KV / recurrent state for its own layers;
• the rows go round a ring: rank 0, 1, ..., the last, and back to rank 0 for the head. That costs one more hidden hop (20 KiB a row for Qwen 27B) than a chain, and keeps sampling, penalties, stop strings, the prompt cache's host side and MTP on one node, unchanged.

What crosses the wire each step:

• the hidden rows after layer b, in the dtype the engine itself keeps between layers. A narrower wire type would change bits (see 5);
• the token ids, for families whose layers read them (Flash-Next's PLE n-gram embeddings, (guess) that each stage needs them);
• the round's row map (slots, positions), as the cluster's round commands carry today (cluster/round.zig:17-89).

Sizes per stage boundary. Qwen 27B is hidden_size 5120 and 64 layers (zig/src/families/qwen27/gdn_test.zig:9). Between layers its CUDA engine keeps two bf16 rows, the hidden state and the residual (e.h and e.r), so its own dtype is already bf16 and an exact boundary carries 2 × 5,120 bf16 values. Flash-Next carries 4 residual streams of 2560 (zig/src/families/flashnext/cuda_prompt_wt.zig:10-15, wide = streams * dims).

Latency budget per boundary. I assume about 115 MB/s of payload at 1 GbE and about 1.15 GB/s at 10 GbE, with a LAN one-way latency of about 0.1 ms (guess).

• Decode: on the ring a token costs one hidden hop per stage, back to rank 0 included: two hops for two stages, about 0.56 ms at 1 GbE on the wire. A plain Qwen 27B (MLX 4-bit) decode step is about 64 ms on a GB10 (measured, section 2a), so the wire is about 1 %. GPU synchronization, staging and waking the peer come on top: the prototype measured 0.2-1.6 ms in all on localhost, and step 1 measures them apart from the wire.
• Prompts: the chunks pipeline. Chunk k+1 computes on stage 1 while chunk k crosses the link and computes on stage 2. So the link only costs time when it is slower than a stage.
  • Qwen 27B (measured): a cold 32k prompt takes 25.0 s on one GB10, about 3.1 s a 4096-row chunk, so about 1.55 s a chunk on each of two stages. The 0.73 s a chunk at 1 GbE hides behind that.
  • The last stage need not send a whole chunk back to rank 0 when only the last row's logits are wanted: one row suffices. MTP absorbs the last layer's rows for every prompt row, and full prompt logprobs need every row, so those keep the full return.
  • At 1 GbE a slower link than a stage is likely for Flash-Next, whose MoE compute per row is small (guess).
  • 10 GbE or more is the comfortable target for long prompts.
• Speed against fit:
  • Decode at one stream is not faster than one node: a token still passes every weight in order, and decode is bound by reading them. Overlapped micro-batches keep both nodes busy, but in the memory-bound regime that buys room for more streams (twice the KV), not faster streams.
  • Prompts can be faster at one stream. With chunks overlapped, a prompt of n chunks on two stages takes about (n + 1) stage-times plus one link transfer: 9 x 1.55 + 0.73 = 14.7 s for a 32k Qwen 27B prompt against 25.0 s, so about 1.7x to the first token, before scheduling overhead. That is the ideal; short prompts (one or two chunks) can be slower than one node, since there is nothing to overlap the fill and the transfer with.
  • What a pipeline always buys is fit: a model larger than one node, or a GB10 and a Mac serving one model.

2a. A prototype: Qwen 27B in two stages on one GB10

To check the boundary before proposing it, I split Qwen3.8-27B (MLX 4-bit) into layers 0-32 and 32-64 in two processes on one GB10, with the boundary over localhost TCP (no Nagle). It is a test-tool mode, not a server, and both processes still load every weight.

• Exact: the last stage's logits equal one process's logits at all 128 greedy steps of a 2,000-token prompt, byte for byte, in both runs.
• Decode: 65.2 and 65.3 ms a step split, against 63.6 and 65.1 ms whole: up to 1.6 ms for two device copies and the hop, on a step that runs nothing in parallel.
• Prompt: one 2,000-row chunk takes 1,632-1,637 ms split against 1,536-1,543 ms whole. The 40 MiB of rows cost about 95 ms, all of it visible, because a single chunk leaves nothing to overlap.
• One trap: the CUDA engine's init sets its fields one by one, so a new field's default is never applied. A stage field left unset gave NaN logits on the whole path.

The next steps for the prototype are loading only each stage's layers, chunked prompts with overlap, and a second machine.

3. Planning

What exists:

• The cluster config already has a pipeline degree, models.<id>.parallel.pipeline (cluster/config.zig:224, 285-289), and its mixed-backend refusal already suggests it: "put each backend in its own pipeline stage" (cluster/config.zig:298).
• plan.zig already cuts uneven, contiguous layer ranges. Each stage's share is weighted by its ranks' gpu_limit (cluster/plan.zig:153-159), through splitLayers (cluster/plan.zig:177-196).
• Speed today only splits work within a stage (cluster/plan.zig:141-144, 162-171).
• Every stage has the same width (cluster/plan.zig:123), and stages take nodes in id order (cluster/plan.zig:156, 161).
• budget.zig checks each rank against its own gpu_limit (cluster/budget.zig:121-123, 137-148).
• cost.zig prices a stage at its slowest rank (cluster/cost.zig:137-149), with FLOPs per backend (cluster/cost.zig:92-95).

Proposed:

1. Place the ends explicitly. Today placeOf puts .head and .final_norm on the last stage (cluster/plan.zig:202). The ring needs them on rank 0, with the embeddings. A pipeline layout names where the head goes, and rank 0's budget counts the embeddings, final norm, head, any drafter and the boundary's staging buffers.
2. Search the boundaries, capped by memory. Layer bytes are a poor measure of time: Qwen's DeltaNet and full-attention layers cost differently, and the gap moves with context and batch. With a measured time per layer for each node and a few stages over 64 layers, an exhaustive search over contiguous boundaries is cheap. It minimises decode time (the sum of stage times plus the hops), or prompt time (the sum plus (chunks - 1) times the slowest stage or link), within each stage's memory. It is no longer a change to splitLayers' weights alone.
3. Allow width 1 per stage on mixed backends. Keep the tensor-parallel same-backend rule (cluster/config.zig:298).

What a node reports (beyond the cluster's Inventory: backend, memory, gpu_limit, bandwidth, cores, unified, cluster/node.zig:46-65; and Starlight's Facts: platform, backend, memory, GPUs, rank_host, sl/internal/control/wire/wire.go:56-68):

• for each family it may serve, the measured ms per layer at decode (1 row, and B rows) and per 4096-row prompt chunk. One calibration run per family and engine version (guess: about 30 s each);
• usable memory for weights and KV after the OS floor;
• for each link: the peer, the measured bandwidth and the round-trip time. A probe over the same TCP port the stages use.

4. Churn

Today:

• Membership tracks alive / suspect / dead / left (cluster/membership.zig:12, 304-315), but nothing reads a view change (cluster/status.zig:83, cluster/bringup.zig:33-44).
• A lost peer ends the round with PeerTimeout (cluster/exchange.zig:39, 107; cluster/barrier.zig:51).
• Starlight drains and stops any replica it no longer lists (sl/docs/nodes.md:23-25) and gives a node 180 s to drain on shutdown (sl/docs/nodes.md:69).

What a client sees: a request that fails before its first byte gets a retryable 503. A stream that has already sent tokens cannot become a 503: it ends with the protocol's error event where the API has one, and otherwise with a closed connection and no final finish_reason, so a client can tell it is incomplete.

Options when a node leaves mid-request:

• (a) Re-plan, then rebuild the lost stage's KV by re-prefilling from the request's tokens. The server has the prompt and the tokens it generated.
• (b) Fail the in-flight requests on that replica with a retryable error, re-plan, and serve new requests on the new placement.

Recommendation: (b) for the first version.

• A re-plan can move layers to the other backend, so a resumed reply would change bits halfway (see 5).
• (a) also needs KV for every layer rebuilt in order, which is a full prefill anyway.
• (b) reuses Starlight's drain-and-restart path, and clients already retry 503s (sl/docs/nodes.md:49).
• (a) can follow as an option for same-backend re-plans, where bits are unchanged.

Joins: a new node changes nothing until the next planned re-placement. That happens at a drain boundary, never mid-request.

5. Exactness

Today the cluster promises the same bits as one node:

• through its fixed reduction tree (cluster/canon.zig:1; cluster/exchange.zig:179);
• tested at 1, 2 and 4 nodes on the toy model (cluster/backend_test.zig:111, 160).

For a pipeline:

• Same backend, same kernels, hidden rows sent in the engine's own dtype: the same bits as one node running the same kernels with the same row batches. A stage boundary moves data but does no arithmetic.
  • The condition is that each stage sees the same batch composition as one node would. Kernels that are not batch-invariant would otherwise differ (guess: true for the current engines' row-batched paths; to be checked per family).
• Mixed CUDA and Metal: no promise of equality with either single node, because each stage's bits come from its own backend.
  • What we can promise: the same placement, engine versions and inputs give the same bits (deterministic per placement).
  • Per-stage tolerance does not compose: an early stage's small error can grow through later layers and the recurrent state. So the check is end to end, against a whole-model reference with the same weights and inputs: the KL divergence and the largest logit difference, teacher-forced top-1 agreement, greedy split positions over long contexts, the DeltaNet states at chosen positions, and repeated runs of one placement. Determinism per placement also assumes the same execution shapes (chunk sizes and row batches).
  • The served reply should say which placement produced it (proposal).
• Not promised: bits across a re-placement, which is why section 4 recommends (b).

6. Fit with the current design

The CUDA engines already have a coordination layer: the tensor-parallel fabric (zig/src/fabric/tp). The pipeline should join it, not the cluster runtime (zig/src/cluster). The cluster runtime is MCDMA-only, builds only on macOS, has no CUDA shard, and its rounds never read membership (CLUSTER-AUDIT.md, sections 3 and 4). That is main today: #1957 (#1594 step 1) adds ConnectX links to fabric/mcdma, and #1594's next step brings the cluster runtime up on a Linux CUDA host, which makes the cluster runtime a real second answer to question 1.

6.1 Communication

What exists in fabric/tp:

• TCP rank bootstrap, a star on rank 0, and a settings agreement (bootstrap.zig).
• Rank 0's round plan to the followers over its own TCP connection (planlink.zig).
• Fail-fast: a failed or vanished rank ends all ranks (fate.zig).
• bootstrap.zig and sock.zig import only std, and fate.zig already branches on the OS (fate.zig:65). planlink.zig is Linux-specific: linux.poll, CPU pinning with sched_setaffinity, /proc reads and MSG_MORE (planlink.zig:7-13, 213-217, 254-287). The build wires tp only beside a CUDA module (zig/build/tp.zig, zig/build/cuda_engines.zig:13), so a Metal build needs these four split into a control module without cuda, and planlink ported to std.posix.poll with pinning Linux-only.
• The Collective is CUDA-typed (collective.zig:3-6: cuda.abi.DevicePtr, cuda.abi.Stream). A Metal follower cannot implement it, so the activation hop does not go through it.

Proposed: a stage link beside the control plane.

• Control: bootstrap, the plan link and fail-fast as they are, with rank 0 the hub.
• Data: direct TCP connections between ring neighbours (rank i to i + 1, the last to rank 0), never relayed through rank 0. Framing carries the round's sequence number and the slot rows' count, and every read and write has a deadline.
• The backend's side is its own: the producing stage's work is complete on its device before the bytes leave (a stream synchronize on CUDA, a completed command buffer on Metal), and the received bytes are in a buffer the next kernel may read before it is launched. Both machines have unified memory, so the boundary buffer can be host-visible on both and no copy is needed past the socket.
• v1 is blocking: send, then receive, in turn. v2 needs a nonblocking form, with a completion per buffer and two buffers a boundary so a chunk's rows are not overwritten while in flight.
• On two CUDA nodes, NCCL send / recv can later replace the TCP data path and sit inside a captured graph.

Trust: fabric/tp's listener binds every IPv4 address (sock.zig:34-43), and nothing authenticates a peer there. The stage link (#2012) refuses a wildcard listen address, checks the checkpoint's and layout's fingerprint in its hello, and proves a shared key both ways before any plan or rows frame; it does not encrypt. v1 is for a trusted network only. Encrypted links for untrusted networks are a separate requirement.

6.2 Node coordination

• Launch and roles work as for tensor_parallel today. Starlight starts each rank with its environment (rank, master, port) and its layer range, and rank 0 serves.
• Ranges come from cluster/plan.zig's splitLayers used as a library: by memory today, by measured speed later (section 3).
• Churn is fail-fast as it stands: one rank lost ends the replica, and Starlight restarts it. That is option (b) of section 4, so v1 needs no membership work.
• Per-row metadata (positions, slots, segments, draw rules) rides the plan link, as in tensor parallel.
• New work on the followers, the largest part of this layer:
  • each stage holds KV and recurrent state for its own layers a slot;
  • keep (the rollback to an accepted path after a verify) reaches every stage through the plan;
  • the prompt cache saves and restores each stage's part under one key;
  • rank 0's admission charges by the tightest stage. The followers report their pools at bootstrap, which fits the byte budget planned in Proposal: one KV admission ledger for unpaged streams: a request reserves its prompt and reply, not a whole window per --parallel stream (#781 follow-up) #1943.
• State consistency:
  • Every plan carries a round sequence number, and every slot a state generation that a keep, a release or a restore advances. A stage refuses rows whose generation is not its own, which turns a missed rollback into an error instead of a wrong answer.
  • v1 runs one round at a time, so the plan link's order is the only order. From v2 on, a slot's buffers are reused only after every stage has completed the round that last wrote them.
  • A prompt-cache entry counts as reusable only once every stage has acknowledged its part, so a hit never restores different positions on different nodes. This matters most for the DeltaNet state, which a token count cannot repair.

6.3 Scheduling

Today one lane host runs one round at a time on one engine thread:

• step lands any in-flight token before it plans (core/lanes/engine.zig:128-131);
• verify reads its tokens back before step returns (families/qwen3_5/cuda_native.zig:416-474, cuda_engine.zig:544-546);
• the engine has one CUDA stream, one set of activation buffers, one staged-table buffer and one graph map (cuda_engine.zig:82-156).

v3 gains room for streams and compute-bound batches, not per-stream speed (section 2), so it comes last.

Any number of stages: nothing here is specific to two. The ring runs rank 0, 1, ..., N - 1 and back to rank 0; the plan link is a star of up to 64 ranks (planlink.max_ranks); splitLayers cuts N ranges. With N stages a v1 token adds N hops (about 2 % at four stages on 1 GbE for Qwen 27B), each node is busy 1/N of the time, v2 approaches N times faster for long prompts, and v3 needs at least N groups.

Drafters:

• MTP taps only the last layer and reuses the head (families/qwen3_5/cuda_mtp.zig:121, 157, 354), so it stays on rank 0.
• DFlash2 taps the drafter's target_layer_ids (for example 5, 19, 33, 47 and 61) on whichever stage holds them (cuda_forward.zig:116-118). Raw taps would ride back with the row, about 50 KiB more a row, relayed through every later stage. Since the drafter's input projection is linear (W·[t1..t5] = Σ Wᵢ·tᵢ), each stage can instead apply its own slice when its tap layer finishes and send one partial sum a row (section 6.4); rank 0 adds them in a fixed order. That is deterministic per placement, not the single-node bits. v1 serves MTP or no drafter.

Graphs: decode is captured from embed to head today (cuda_engine.zig:511-526). Each stage needs its own capture, and on TCP the hop stays outside the captured region.

6.4 Boundary format v0 (proposal)

One format for every backend, so a CUDA stage and a Metal stage can be built separately against it. Everything is little-endian. Framing is #2012's (fabric/stage/frame.zig); what v0 adds is marked (new).

Frame header, 32 bytes:

Boot:

1. Hello (fabric: the pipeline stage link, a TCP ring that carries each round's plan and hidden rows, with a shared key (Linux and macOS, #2005 step 1) #2012): rank, world and a 32-byte fingerprint of the checkpoint and the whole layout (every rank's layer range), then the shared-key proof both ways. A neighbour with another fingerprint, world or key is refused.
2. Description (new), each stage to rank 0 over the plan link: engine version, backend (CUDA, Metal), KV dtype; layer range a..b; the boundary format (hidden width, row dtype, residual planes); execution shapes (prompt chunk, most rows a round, --parallel, context); the drafter tap layers it owns; capacity (KV rows and recurrent states it can hold). Rank 0 refuses any mismatch by its name, before a round runs.
3. Ready (new): each stage reports loaded and warm: executable for the shapes it declared. That is all warm promises: graph capture is a CUDA detail, and a Metal stage need not offer it. Rank 0 opens its server only after every stage has.

Plan, rank 0 to every stage, once a round (new payload): for each stream in the round, in a fixed order:

• stream id, and an op: open, continue, close, cancel, or resume from a kept prompt state;
• first position and row count, with the round's chunk cuts;
• the previous verify's accepted count, so every stage rolls its KV and recurrent state back to the same position;
• the slot's state generation (section 6.1), which a keep, a release or a restore advances: a stage refuses a stream whose generation is not its own, by name.

Token ids are left out: no later stage reads them for Qwen, and a family that needs them would add them here. A plan with no streams is the keep-alive: rank 0 sends one once 2 s pass without a plan, idle or in the middle of a round, and a stage that hears no frame for three intervals declares the ring dead (fail-fast, section 4). A keep-alive never changes a stream's state, and it takes a frame round number like any plan. So a long forward on one stage, a slow chunk on a Mac for example, never reads as a dead ring elsewhere. A round's own deadline is separate: it bounds each send and receive, not the forward.

Rows, stage i to stage i + 1: the boundary's hidden rows in the engine's own dtype, bit for bit, as raw little-endian values. Plane-major, in the family's fixed plane order (Qwen 3.5: plane 0 the hidden state h, plane 1 the residual r, both bf16, as its CUDA engine keeps them and ThinkOffApp's prototype sends them), rows in plan order, hidden values a row. Width, row dtype and plane count are in the description and checked at boot; the boundary site is the sending stage's layer range end; the family and checkpoint are bound by the hello's fingerprint. After the rows, when a drafter is on, the tap partial sums: one draft_width row per row in fp32. The first stage starts them at zero; each stage adds the contributions of its own tap layers, in ascending layer id, and forwards the sum. A split sum rounds differently from one GEMM on one node, so each placement with taps needs its own gate; v0 carries the format only, and no engine implements taps before step 10. The last stage sends its rows back to rank 0 for the head and the sampler.

v0's decisions (step 1b's code and fixtures follow them):

• Payload prefix: every description, ready and plan payload starts with 8 bytes: magic 0x50534654, whose little-endian bytes read "TFSP", a u16 protocol version (1), the message (1 describe, 2 ready, 3 plan) and a reserved 0. The protocol version travels here, so fabric: the pipeline stage link, a TCP ring that carries each round's plan and hidden rows, with a shared key (Linux and macOS, #2005 step 1) #2012's frame header stays at version 1. Rows payloads have no prefix.
• The description also carries the rank and world, and a drafter width (0 without taps), so a tap-sum size mismatch is refused at boot too.
• Layer ranges are half-open: first..end, end excluded. Rank 0 checks the layout: contiguous from 0 to the model's layer count, with no gap, overlap or overrun.
• Version: every stage must report the same engine version string as rank 0's placement, the release both sides were built from. The backend is a separate field, so a CUDA and a Metal stage of one release pass. Allowing different releases per stage is left for later.
• Capacity floor at boot: KV rows ≥ the context and recurrent states ≥ --parallel. Per-request admission against the live budgets comes later, in step 7.
• Plan: streams in ascending id order. open starts at position 0 with 0 accepted; resume names a kept state's position with 0 < resume ≤ first; close and cancel carry no rows. A chunk cut is a row offset across the whole round, strictly ascending and inside the round.
• Keep-alive: a stage counts any frame as heard (a plan, a keep-alive or rows), and rank 0's keep-alive is due whenever 2 s pass without a plan, mid-round too.
• Stream entry, 28 bytes: id, op, 3 reserved, first, rows, accepted, resume point, generation (u32 each, op u8). An open carries generation 0.
• Golden bytes: fabric stage: the protocol layer, boot description, round plan, rows layout and keep-alive, with golden bytes for every message (#2005 step 1b, stacked on #2012) #2060's zig/src/fabric/stage/fixtures/ pins every width, enum value, length and reserved byte, good and bad, with the error each bad file must raise; a host encoder and decoder can be built against them before any engine.

Where sampling runs. v0 keeps the ring: tokens, the lane host, the sampler and any drafter stay on rank 0, so the API's behaviour (stop strings, logprobs, grammars) has one home. ThinkOffApp's prototype is a chain instead: the last stage runs the head and sampler and returns a 4-byte token id, which saves one hidden row a round on the way back. That could be a boot option (sampler: first | last) if it pays; see question 8.

7. Steps

The test hardware is one GB10 and one Mac. That sets the order:

• two CUDA nodes over NCCL can't be tested, and NCCL refuses two ranks on one GPU, so the TCP stage link comes first;
• the only two-machine test is GB10 + Mac, so a Metal follower comes before any speed claim;
• two processes on one GB10 share its GPU, so they show exactness and memory, not overlap gains.

• No day estimates yet. The Metal side of the stage link and control plane, range-only loading and the follower state protocol are each larger than a first estimate showed, so the steps carry estimates once a spike of each exists. Steps 1, 1b and 2 come first, on the GB10 alone; 1b's fixtures let a Metal stage be built in parallel against the same bytes.
• Not testable here: two CUDA nodes over NCCL / RoCE. That path would need a second CUDA node, or your pair.

8. Open questions

1. Where should the layout live? I'd propose the fabric/tp rank model, with engine flags placed by Starlight and the cluster planner as a library (section 6), rather than the cluster runtime's MCDMA rounds. Is that the direction you want, or should the cluster runtime gain CUDA and TCP instead?
2. Hidden rows on the wire: the engine's own dtype (exact, larger), or bf16 (half the bytes, different bits)? For Qwen 27B these are the same, since its own rows are bf16 (section 2); the question is for Flash-Next.
3. First family: Qwen 27B (simpler, dense GDN), or Flash-Next (the one we want to serve)?
4. Is "deterministic per placement, within single-node tolerance" acceptable as the promise for CUDA + Metal pipelines?
5. Micro-batching (v3): it changes the lanes core every family shares (section 6.3). Is a later change there acceptable, or should pipelines stay at v1 and v2?
6. A second CUDA node: we have one GB10 and one Mac. Could the NCCL path (two CUDA nodes) be checked on your pair once steps 1 and 2 land?
7. Trust: is "trusted network only, with fingerprints and a shared key checked at bootstrap" enough for v1?
8. Ring or chain: should the last stage be allowed to sample and return token ids (a chain, as in ThinkOffApp's prototype), or does sampling stay on rank 0 (section 6.4)?