swactor/examples/pipeline-parallel-inference/SPEC.md

Pipeline-Parallel Inference — Spec (N stages)

This spec covers the N-stage generalisation of the pipeline-parallel example. For the deployment-shape rationale and where this rung sits in the broader buildout, see /docs/ROADMAP.md §3.1. For the test enumeration, see ./TEST_SPEC.md.

Behavioural snapshot of the 2-stage deployment, as it actually runs today, is preserved in ../single-gpu-inference/BEHAVIOR.md for cross-reference and in ./BEHAVIOR.md (to be written once the N-stage path is deployed).

1. Goal

pp-smoke-run --vastai --num-stages N --api-key <key> runs a linear chain of N ≥ 2 GPU nodes that jointly serve one InferenceRequest end-to-end. Acceptance test: N = 5 on vast.ai, llama3.2:1b, prompt "Say hello", produces non-empty text within the 5-minute budget. The localhost (--seed) path supports the same N values via spawning N child processes.

N = 1 is not supported by this example — single-node inference is what examples/single-gpu-inference exists for, and the role-based actor design here is built on the assumption that "first" and "last" are distinct stages. N = 2 is the minimum: first stage handles the embed + first half; last stage handles the second half + head; no middle stages.

2. What changes vs. the 2-stage MVP

Area	2-stage today	N-stage target
pp-gpu-node binary	hard-checks NUM_STAGES == 2 (line 214) and dispatches to run_stage_0 or run_stage_1	accepts any N ≥ 2; dispatches to a single run_stage function parameterised by (stage, num_stages)
Stage actors	Stage0Actor and Stage1Actor are separate types	one StageActor parameterised by role (First / Middle / Last)
Message bridges	three bridge types (Stage0RequestBridge, Stage0NextTokenBridge, Stage1ActivationBridge) bound to specific actor types	same three bridge shapes, but generic over the unified StageMsg
Orchestrator (pp-smoke-run)	spawns exactly two children, hands stage 1 stage 0's PEER_DIRECT	spawns N children in order; each non-first child receives its predecessor's addressing as PEER_DIRECT
Topology helpers	already generic (stage_name, next_stage_name, prev_stage_name)	unchanged
Worker (pp_tinygrad_worker.py)	accepts stage-0 ops (embed_and_forward, decode_step) and stage-N-1 ops (forward_and_sample); already param on STAGE/NUM_STAGES for layer range	adds forward_range op for middle stages; layer-range computation unchanged

The worker contract is the anchor of stability. Anything that doesn't have to change, doesn't.

3. Actor design — a single StageActor

3.1 Role enum

pub enum StageRole {
    /// stage == 0. Owns: prompt entry, tokenizer, embed + first range.
    /// Reacts to: `InferenceRequest` from network; `NextToken` from
    /// the *last* stage (decode-loop feedback).
    First,

    /// 0 < stage < num_stages - 1. Owns: a middle block range, no
    /// tokenizer, no sampler. Reacts to: `StageActivation` from prev.
    Middle,

    /// stage == num_stages - 1. Owns: last block range, output head,
    /// sampler, detokenizer. Reacts to: `StageActivation` from prev.
    /// Emits: `NextToken` to stage 0; `InferenceResponse` to
    /// `reply_to` on terminate.
    Last,
}

Role is set once at construction and never changes. The actor's per-message logic branches on it. The role is computed by the binary from (STAGE, NUM_STAGES) and is a pure function (see TEST_SPEC §3).

3.2 Unified incoming message

pub enum StageMsg {
    Inference(InferenceRequest),     // First only
    Activation(StageActivation),     // Middle and Last
    NextToken(NextToken),            // First only
    Process(ProcessNotification),    // all roles
    SetNeighbors { prev: Option<ActorAddress>,
                   next: Option<ActorAddress>,
                   reply_to: Option<ActorAddress> },
    Reset,
}

Compared to today, this collapses Stage0Msg and Stage1Msg into one type. prev and next are Option because the chain ends are asymmetric. reply_to is Some only for Last. An actor that receives a variant outside its role's set logs and drops it — this is defensive, not load-bearing, because the orchestrator wires bridges matching the role.

3.3 Per-role message behaviour

Role	On Inference	On Activation	On NextToken
First	tokenize → embed_and_forward → send StageActivation { is_prefill: true } to next	drop	done=true → drop; else decode_step → send StageActivation { is_prefill: false, seq_len: 1 } to next
Middle	drop	forward_range → send StageActivation { … } to next, preserving request_id, position, seq_len, is_prefill	drop
Last	drop	forward_and_sample → send NextToken to prev (= the first stage); on terminate, detokenize → send InferenceResponse to reply_to	drop

Note: the autoregressive feedback edge is last → first, not last → prev. This is the same shape as the 2-stage MVP (where prev is first), and the right shape for arbitrary N — the first stage is the only one that owns the embed step.

3.4 New worker op: forward_range

Middle stages need a worker op that takes a hidden state, runs blocks [lo..hi), and returns the new hidden state. Stage 0's embed_and_forward already runs a block range; stage 1's forward_and_sample runs a range and then samples. The new op is the range-only variant.

→ {"op": "forward_range", "request_id": rid, "hidden_b64": "...",
   "position": p, "seq_len": s}
← {"request_id": rid, "hidden_b64": "...", "seq_len": s}

In pp_tinygrad_worker.py this is the same code path as the prefix of forward_and_sample, factored out and exposed as its own dispatch case. Tinygrad's per-block KV cache continues to be allocated lazily per block, so each middle stage owns the K/V for its own range automatically.

4. Orchestrator changes (pp-smoke-run)

4.1 Seed mode

Today: spawn stage 0, read its PP_GPU_NODE_ADDR, spawn stage 1 with that addressing as PEER_DIRECT. For N stages:

for i in 0..N:
    spawn child i with PEER_DIRECT = (addr of child i-1, if i > 0)
    read child i's PP_GPU_NODE_ADDR

Sequential, not parallel. Each subsequent child needs its predecessor's direct addressing in PEER_DIRECT, so the spawn loop is inherently a chain. Each child only needs its predecessor's addressing (not all earlier addresses): the outbound dial from i → i-1 makes i-1's iroh stack learn i's source-socket reverse-path. The seed/orchestrator is dialed by every child, so it always learns everyone.

4.2 vast.ai mode

Today: find_offer twice (excluding the first offer's id), then create_pipeline_instances with both offer ids. For N stages:

• find_offer N times, each call excluding all previously chosen offer ids.
• create_pipeline_instances takes a slice of N offer ids, sets STAGE, NUM_STAGES, SEED_ADDR, SEED_RELAY on each; in vast.ai mode, no PEER_DIRECT is passed — relay-mediated NAT traversal
  • SWIM gossip handle peer discovery (this is already how vast.ai mode works for 2 stages today, no change).
• wait_for_running once per contract, sequentially (same as today, just N times instead of 2).
• destroy_all_instances is already a slice-aware helper.

4.3 Convergence target

The orchestrator waits for SWIM to report alive >= N peers (today: alive >= 2, i.e. N=2). One-line change in both run_seed and run_vastai.

4.4 Submitting the request

Unchanged: resolve pp-entry, build route, send_to, await_response. The chain is invisible to the orchestrator past the entry node.

5. pp-gpu-node changes

5.1 Boot sequence (generalised)

1. Parse STAGE, NUM_STAGES (any N ≥ 2; remove the == 2 check).
2. Determine role: First if STAGE==0; Last if STAGE==N-1; else Middle.
3. Build iroh driver, join seed.
4. Wait for SWIM convergence (>= 1 alive peer — unchanged).
5. Spawn the StageActor with the role + placeholder addresses.
6. Register the per-index name pp-stage-{STAGE} immediately so
   neighbours can resolve us.
7. Wait for worker ready (unchanged).
8. Resolve neighbours by role:
     First:  resolve pp-stage-1 (next).
             Note: the feedback target (NextToken from last) is *us* —
             the last stage looks up pp-stage-0 itself.
     Middle: resolve pp-stage-{stage+1} (next).
             A middle does not send to prev; nothing to resolve there.
     Last:   resolve pp-stage-0 (the NextToken destination),
             plus pp-orchestrator (for InferenceResponse).
9. Build IrohActorTransport routes for each resolved neighbour.
10. SetNeighbors on the actor with the resolved addresses.
11. Pump 100ms to drain the setup message.
12. Register pp-entry (First) / pp-exit (Last). pp-entry must be
    registered last because the orchestrator may send InferenceRequest
    as soon as it resolves — and First is not ready to forward until
    pp-stage-1 is wired.
13. Enter main pump loop.

The step-8 "who do I resolve" matrix is the only real new logic. It flows from the per-role message table in §3.3.

5.2 Name registry summary

Stage	Registers	Resolves (and routes to)
First (0)	pp-stage-0, pp-entry	pp-stage-1
Middle (i)	pp-stage-i	pp-stage-{i+1}
Last (N-1)	pp-stage-{N-1}, pp-exit	pp-stage-0, pp-orchestrator

This is exactly what topology.rs::register_stage_names already emits — the topology helpers do not change.

6. Failure modes (deltas only)

The 2-stage failure table previously in this spec applies stage-by-stage. N-stage-specific additions:

Failure	Detection	Response
Fewer than N offers available	find_offer returns empty on attempt < N	Destroy any allocated instances, exit 1.
Any one of N instances never reaches running	Per-contract wait_for_running	Destroy all contracts, exit 1.
SWIM never converges to N alive peers	180s timeout in vast.ai mode	Destroy all, exit 1. Fetch logs from each instance.
Any one of N workers never reports ready	Per-node wait_for_worker_ready (600s)	That stage exits 1; SWIM marks it dead; orchestrator times out on convergence and destroys all.
Middle stage drops a StageActivation (mid-decode)	Orchestrator's 600s response timeout	Destroy all, exit 1. Recovery is rung 6 (fault tolerance), out of scope here.

The teardown invariant — "every code path that creates an instance also destroys it" — generalises to "any creation failure rolls back all already-created instances." destroy_all_instances already does this; the loops above just have to call it on every error branch.

7. Test plan

The full enumeration of behavioural tests lives in ./TEST_SPEC.md. Summary of tiers in landing order:

1. Pure-Rust unit — codec, topology helpers, role computation, per-stage layer-range math.
2. Python worker contract — forward_range (new), all existing ops verified against stub mode, layer-range math up to N=8.
3. Stage actor (in-process, stub worker) — StageActor in each of its three roles plus role-mismatch defensive drops.
4. Cluster transport (in-process, real iroh) — N-node clusters (N ∈ {2, 3, 4}) carry every message type intact across every hop.
5. vast.ai client (mocked HTTP) — N-instance create/destroy with distinct STAGE env vars and rollback on partial failure.
6. In-process N-stage integration (stub worker, real iroh) — full chain at N ∈ {2, 3, 4, 5}.
7. Sliced-vs-full equivalence (real tinygrad on CPU, gated) — the load-bearing correctness test: N-stage output token sequence matches single-process Transformer.generate() up to max_tokens. Run at N=2 (carried forward), N=3, and N=4.
8. Localised binary E2E (pp-smoke-run --seed) — the pre-deploy gate: runs the actual pp-gpu-node binaries spawned by the actual pp-smoke-run binary in --seed mode at N ∈ {2, 3, 5}.
9. Manual smoke run on vast.ai at N=5 — gated, rare, the acceptance test for the rung.

A change that breaks any tier 1–4 test must not pass review. Tier 8 must pass before any change is deployed. Tiers 5–7 are CI-gated; real-tinygrad and equivalence are #[ignore] and run on demand or in a slow lane.

8. Migration

Land in this order, each its own PR:

1. Worker op forward_range. Additive in pp_tinygrad_worker.py. No call sites yet. Tested with Python tests in test_worker.py.
2. Unify Stage0Actor + Stage1Actor into StageActor. Keep message-shape behaviour identical for N=2. Delete the old types. Existing t_actor.rs tests adapt to the unified API. This PR does not change the wire or the binary.
3. Add Middle role + actor logic behind a constructor variant. Add N=3 / N=4 integration tests against the stub worker. Binary still rejects N != 2 to avoid premature deploys.
4. Generalise pp-gpu-node (drop the N==2 check; generalise the resolve/route step). Update t_binary.rs to cover N=3 in seed mode.
5. Generalise pp-smoke-run (spawn loop, find-offer loop, convergence target). Add --num-stages flag, default 2.
6. Land tier-8 N=5 localised E2E. This is the pre-deploy gate.
7. Manual vast.ai run at N=5. No code change — just exercise the path and capture the result in BEHAVIOR.md.

Each PR is independently mergeable and ships behind the existing 2-stage acceptance. Tier 8 at N=5 must pass before step 7.

9. Out of scope (this rung)

• Microbatching, multiple in-flight requests, streaming output.
• Dynamic stage assignment / rebalancing. Stages are decided at boot.
• Fault tolerance of any kind beyond "destroy everything on any failure." Rung 6.
• Heterogeneous layer splits — every stage still owns total / N blocks ± 1. The planner (rung 5) takes over splitting.
• N = 1. Use examples/single-gpu-inference.