Pipeline parallelism for the vLLM runtime

[khaiwang]

Adds pipeline parallelism to the vLLM runtime: a trace written against the client's meta tree runs unchanged on an engine sharded across PP stages (including TP within each stage), with reads, writes, saves, and caches working across stage boundaries. Four commits on the current 0.8.

How it works

Every rank runs every intervention block; each rank's forward only visits its own stage's modules. Three seams on the base interleaver close the gap, each a no-op on a single rank:

• Intercept (Mediator.event): the PP interleaver sees a worker's event before it parks. A read of a location another stage owns is answered with a LazyRemoteTensor; forcing it issues a cross-stage pull at the moment the worker parks, so the transfer overlaps the rest of the forward. A write to a remote-owned location is absorbed, since the owning rank runs the same line locally. A force of an upstream-owned value whose round this rank has opened is served in place inside the handoff; a force for a later round parks like a downstream pull.
• Publish (Interleaver.publish, called after the workers parked on a visit are served and before a fragment is re-split): the PP interleaver copies the value to host on the forward thread and buffers it, per served request, under the occurrence those workers waited for. A location enters the buffer only when a worker read it. Pull traffic uses dedicated per-TP-column gloo groups (pp_listener.py).
• Step gate: an open-ended tracer.iter[:] whose body never parks (a body of remote reads builds lazies and returns) parks on a gate the driver serves once per generation step; the vLLM runner serves it at its own step boundary.

Serve points along a step: before the forward, pulls whose round has run complete; after it, the round count advances and landed pulls complete without waiting; the sampling path completes produced-round stragglers before the logits and samples offers. At collect, the finished requests' workers are served under the same produced-round gate, the drain barrier holds every rank until the peers' pulls are in, and the finished requests' buffer entries are cleared.

Ownership is derived from a load-time exchange of which modules are real on which rank (pp.py); modules an architecture builds only on some ranks are stubbed from the full meta tree the worker builds before its groups exist, and the meta tree's children are grafted under each PPMissingLayer envoy so sub-stub paths resolve. Each stage's TP rank 0 ships the slots it owns; the engine's merge unions them slot-wise (collect.py), including per-stage tracer.cache() recordings.

Also included: a C-level per-greenlet swap of torch's thread-local dispatcher state (pp_tls_swap.py), so a worker parked inside a torch call leaves the forward's state intact.

Tests

tests/vllm/pp runs as one command in three tiers: unit (one process), harness (two gloo ranks over the real interleaver and listener, no engine), and engine (a shared PP=2 engine for the behavior tests, per-configuration subprocess engines for PP=1 vs PP=2 parity, PP=3, and TP=2 x PP=2 topology). pytest tests/vllm/pp -m "not gpu" runs the CPU tiers in about 40 seconds (101 tests); -m gpu runs the engines (29 tests, about 18 minutes).

Validation

• CPU tiers: 101 passed, 1 xfailed.
• GPU tier: 29 passed, including every tracer.iter shape against the pull machinery with elapsed-time bounds, the request-count gauge at zero on every rank, and a rank-gated architecture (granite) under PP.
• Upstream vLLM suites (tracing, requests, taps) and the non-vllm CPU suite match a clean origin/0.8 checkout in this environment exactly.

🤖 Generated with Claude Code

https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR