# Pipeline-Parallel Inference Spec: Two Nodes, Say Hello

## 1. Wanted Behavior

`pp-smoke-run --vastai --api-key <key>` runs on the local machine. It:

1. Rents two GPUs on vast.ai (cheapest available matching the constraint).
2. Starts a local swactor node with iroh transport.
3. Both rented instances boot the `gpu-node` image. One is assigned `STAGE=0`, the other `STAGE=1`, via env var at instance-create time. Both also receive `NUM_STAGES=2` and `SEED_ADDR=<local node id>`.
4. Both nodes join the local cluster via iroh/QUIC.
5. Stage 0 loads layers `[0..k)` of `llama3.2:1b`; stage 1 loads layers `[k..N)` plus the final norm and output head. `k = num_blocks / 2`.
6. The orchestrator resolves the SWIM name `"pp-entry"` (registered by stage 0) and sends `InferenceRequest { prompt: "Say hello", max_tokens: 64 }` to it.
7. Stage 0 tokenizes the prompt, runs embedding + its layer range, sends a `StageActivation` message to stage 1.
8. Stage 1 runs its layer range + output norm + LM head, samples a token, and sends it back to stage 0 as the next decode input. Loop until EOS or `max_tokens`.
9. Stage 1 sends the final detokenized text to the orchestrator as `InferenceResponse`.
10. Orchestrator asserts response is non-empty, prints it, destroys both instances, exits 0.

If anything fails, both instances are destroyed (best-effort) and the command exits 1.

Total wall-clock budget: 5 minutes.

```
local (pp-smoke-run)        vast.ai stage-0           vast.ai stage-1
       |                          |                         |
       |  rent x2, set STAGE env  |                         |
       | -----------------------> |                         |
       | --------------------------------------------> .... |
       |                          | [load blocks 0..k)      |
       |                          |                         | [load blocks k..N, output head]
       |  SWIM join (iroh/QUIC)   |                         |
       | <----------------------> | <---------------------> |
       |                          |                         |
       |  InferenceRequest        |                         |
       |  "Say hello"             |                         |
       | -----------------------> |                         |
       |                          |  StageActivation        |
       |                          |  (hidden state, pos)    |
       |                          | ----------------------> |
       |                          |                         | [sample token]
       |                          |  next-token feedback    |
       |                          | <---------------------- |
       |                          | [decode loop ...]       |
       |                          |                         |
       |       InferenceResponse "Hello! ..."               |
       | <------------------------------------------------- |
       |                          |                         |
       |  DELETE both instances   |                         |
       |  exit 0                  |                         |
```

---

## 2. Scope and Non-Goals

**In scope.** Exactly two nodes. One static layer split, decided by the orchestrator at boot. One `InferenceRequest` per run. tinygrad worker with manually sliced `Transformer.blk`. Reuses the single-GPU example's vast.ai client, iroh transport, and swactor process bridge unchanged.

**Out of scope.** N > 2 nodes; dynamic stage assignment or rebalancing; microbatching; multiple in-flight requests; streaming output; fault tolerance of any kind; per-token latency or throughput goals; tensor parallelism; non-tinygrad workers; quantization choices beyond what the single-GPU example already uses.

---

## 3. Architecture

### Topology

A linear chain of two stages. Stage 0 is the chain entry point and registers itself under the SWIM name `"pp-entry"`. Stage 1 is the chain exit and registers under `"pp-exit"`. The orchestrator resolves `"pp-entry"` to submit the request and `"pp-exit"` is the address stage 1 uses for replies. Stage 0 resolves `"pp-stage-1"` to know where to send activations; stage 1 resolves `"pp-stage-0"` to send sampled tokens back for the next decode step.

Each stage knows its own `STAGE` index and `NUM_STAGES` from env. The "next stage" name (`pp-stage-{i+1}`) and "previous stage" name (`pp-stage-{i-1}`) are computed locally — no central topology config.

### Message types

```
InferenceRequest {
  reply_to:  ActorAddress,
  prompt:    String,
  max_tokens: u32,
}

StageActivation {
  request_id: u64,
  position:   u32,        // start_pos for KV cache
  hidden:     Vec<u8>,    // bf16 tensor bytes, shape [seq_len, hidden_dim]
  seq_len:    u32,
  is_prefill: bool,
}

NextToken {
  request_id: u64,
  token_id:   u32,
  position:   u32,        // position to use for the next forward
  done:       bool,       // EOS or max_tokens reached
}

InferenceResponse {
  text: String,
}
```

`StageActivation` flows stage 0 → stage 1. `NextToken` flows stage 1 → stage 0 (closes the autoregressive loop). `InferenceResponse` flows stage 1 → orchestrator once decode terminates.

### Per-request flow

1. Orchestrator → stage 0: `InferenceRequest { reply_to, prompt, max_tokens }`.
2. Stage 0 tokenizes, runs `token_embd` + blocks `[0..k)` on the full prompt (prefill), produces hidden state for the final position.
3. Stage 0 → stage 1: `StageActivation { request_id, position: 0, hidden, seq_len: prompt_len, is_prefill: true }`.
4. Stage 1 runs blocks `[k..N)` + `output_norm` + `output`, samples next token.
5. If EOS or `max_tokens` reached: stage 1 → orchestrator: `InferenceResponse { text }`. Done.
6. Otherwise: stage 1 → stage 0: `NextToken { request_id, token_id, position: prompt_len, done: false }`. Stage 1 also accumulates `token_id` into its local output buffer for detokenization at the end.
7. Stage 0 runs `token_embd` + blocks `[0..k)` on the single new token at `position`, sends `StageActivation { ..., seq_len: 1, is_prefill: false }`. Back to step 4.

KV cache is held on the actor that owns its layers — stage 0 caches K/V for blocks `[0..k)`, stage 1 for `[k..N)`. This falls out of tinygrad's per-block cache without any extra plumbing.

### Wire format

Hidden-state tensors are sent as raw little-endian bf16 bytes plus the shape. No protobuf, no msgpack — the existing swactor codec carries `Vec<u8>` plus typed scalars and that's enough.

For `llama3.2:1b` (`hidden=2048`, bf16): one decode-step activation is `2048 * 2 = 4096` bytes per hop. Prefill is `prompt_len * 4096` bytes. Negligible.

---

## 4. Worker

The worker is `pp_tinygrad_worker.py`, a fork of `tinygrad_worker.py` from the single-GPU example. Same stdin/stdout JSON protocol with the swactor process bridge, different operations.

### Loading

On startup, the worker reads `STAGE`, `NUM_STAGES`, and `MODEL` from env. It calls `Transformer.from_gguf(...)` to load the full model. This is wasteful — both nodes hold all weights in VRAM — but it's the simplest path and `llama3.2:1b` fits trivially. Pruning unused blocks is a later optimization.

The worker computes its layer range:

```
total = len(model.blk)
k = total // NUM_STAGES
my_range = (STAGE * k, (STAGE + 1) * k if STAGE < NUM_STAGES - 1 else total)
```

### Per-stage operations

**Stage 0** handles two operations:

- `embed_and_forward(token_ids, position)`: run `model.token_embd(tokens)`, then iterate `model.blk[0:k]` calling `block(x, start_pos=position)`. Return hidden state bytes.
- `decode_step(token_id, position)`: same as above with a single-token input.

**Stage 1** handles one operation:

- `forward_and_sample(hidden_bytes, position, seq_len)`: rebuild the tensor, iterate `model.blk[k:total]` calling `block(x, start_pos=position)`, apply `model.output_norm`, apply `model.output`, take logits at the last position, sample (argmax for MVP), return `token_id`.

Sampling is `argmax` only. No temperature, no top-p — matches what `Transformer.forward()` does at line 178 of `tinygrad/apps/llm.py`.

### What we bypass in tinygrad

- `Transformer.forward()` is not used — it runs the entire model. We call `model.blk[i:j]` directly.
- `Transformer.forward_jit` is not used — it wraps the full forward. Each stage's own forward is uncompiled. Slower, fine for MVP.
- `Transformer.generate()` is not used — the autoregressive loop is split across stages, not a single Python generator.

### What we get for free

- Per-block KV cache (`self.cache_kv` in `TransformerBlock._attention`) is allocated lazily on first call and persists across decode steps. Stage 0's cache holds K/V for blocks 0..k; stage 1's holds k..N. Correct by construction.
- The tokenizer (`SimpleTokenizer.from_gguf_kv`) is loaded on stage 0 only.
- Detokenization happens on stage 1, which accumulates sampled token ids and decodes once on EOS/max_tokens via `tok.decode([...])`. The tokenizer is also loaded on stage 1 for this — cheap.

### Worker protocol additions

In addition to the single-GPU example's `{"status": "ready", "pid": <pid>}`, the worker accepts and replies to stage-specific request types over the same stdin/stdout JSON channel:

```
Stage 0:
→ {"op": "embed_and_forward", "request_id": 1, "tokens": [...], "position": 0}
← {"request_id": 1, "hidden_b64": "...", "seq_len": <int>}

Stage 1:
→ {"op": "forward_and_sample", "request_id": 1, "hidden_b64": "...", "position": 0, "seq_len": 42}
← {"request_id": 1, "token_id": 1234}
```

Hidden-state tensors are base64-encoded on the stdin/stdout boundary (process bridge is text JSON). On the actor-to-actor wire they are raw bytes — no base64 — via the swactor codec.

---

## 5. What's Reused from single-gpu-inference

| Component | Reused as-is | Notes |
|---|---|---|
| `vastai.rs` REST client | yes | Just called twice instead of once. Image name parameterized. |
| `iroh_transport.rs` | yes | Same `RelayMode::Default` over WAN. |
| swactor process bridge | yes | Same stdin/stdout JSON pattern; different ops inside. |
| `Dockerfile` | mostly | Adds `pp_tinygrad_worker.py`. `STAGE`/`NUM_STAGES` env are runtime, not baked. |
| `gpu-node` binary | forked | New `pp-gpu-node` reads `STAGE`/`NUM_STAGES`, registers stage-specific SWIM names, instantiates the stage actor. |
| `smoke-run` binary | forked | New `pp-smoke-run` rents two instances, waits for both, sends one request, tears both down. |
| Message types | extended | New `StageActivation`, `NextToken` in addition to the existing `InferenceRequest`/`InferenceResponse`. |
| Tinygrad worker | forked | New `pp_tinygrad_worker.py` exposing per-stage ops instead of full inference. |

---

## 6. Error Modes

| Failure | Detection | Response |
|---|---|---|
| Fewer than 2 GPU offers available | vast.ai search returns < 2 | Print error, exit 1. No cleanup needed. |
| Either instance creation rejected | API non-success on either call | Destroy any instance already created, exit 1. |
| Either instance never reaches `running` | 60 polls exhausted | Destroy both, exit 1. |
| Cluster never converges to 3 members | SWIM timeout (2 min) | Destroy both, exit 1. |
| `pp-entry` or `pp-exit` never resolves | Timeout (1 min) | Destroy both, exit 1. |
| Either worker never reports ready | Process timeout | Destroy both, exit 1. |
| Stage 1 reports `done` but no text accumulated | Assertion on response | Destroy both, exit 1. |
| Either node drops out mid-decode | Iroh connection error / SWIM marks dead | Destroy both, exit 1. No retry. |
| Destroy fails for either instance | Catch around cleanup | Print warning, exit 1. |

Any path that allocates an instance also destroys it (best-effort).

---

## 7. Environment Requirements

- Rust toolchain.
- Python 3.10+ with tinygrad 0.12.0 + numpy (in `.venv/`).
- Docker.
- `VAST_API_KEY` (vast.ai deployment only).
- Network access to iroh relay servers.

Same as single-gpu-inference. No new dependencies.

---

## 8. What This Proves

1. swactor nodes on rented GPUs can form a chain that carries activations stage-to-stage over WAN.
2. Model weights can be split across machines and the resulting forward pass still produces coherent output.
3. KV cache locality (each stage holds K/V for its own layers) is preserved across decode steps without explicit cache management on the orchestrator.
4. The autoregressive feedback edge (stage 1 → stage 0) works as a swactor message and closes the loop without involving the orchestrator on every token.
5. The full lifecycle (rent x2 → cluster → distribute layers → infer → teardown) is automated.

This is the foundation for the next rung: N > 2 stages, where the message types and per-stage worker contract here generalize directly (each intermediate stage just forwards activations to `pp-stage-{i+1}`).