Pipeline parallelism for the vLLM runtime by khaiwang · Pull Request #691 · ndif-team/nnsight

khaiwang · 2026-08-20T04:04:37Z

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, and saves working across stage boundaries.

How it works

Every rank runs every intervention block; each rank's forward only visits its own stage's modules. The gap is closed in three moves, all on a PP-aware interleaver (pp_interleaver.py):

Intercept (worker side): a read of a location another stage owns is answered immediately 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 is served in place inside the hook, which is deadlock-free by pipeline order and preserves the read-modify-write pattern.
Publish (producer side): as a rank serves a location it owns, each request's rows of the post-intervention value are cloned into a pull buffer keyed by occurrence tag and request id, and parked peer pulls are dispatched. Pull traffic rides dedicated per-TP-column gloo groups (pp_listener.py), separate from vLLM's own communication.
Serve (driver side): at serve points the runner completes in-flight pulls and resumes the parked workers. A pull is waited for only when its producing sampling round has finished on the local schedule; blocking on a current-round pull would invert the pipeline order into a deadlock.

Ownership is derived from the loaded module tree (pp.py); sub-stub paths under PPMissingLayer resolve through meta-model envoys grafted at load. Saves ship per stage and merge engine-side, refusing silent clobbers on rank divergence.

Supporting fixes that ship with the stack:

Greenlet parks inside torch dispatch corrupted the pybind warning-handler state; a C-level at::ThreadLocalState swap per greenlet switch (JIT-built extension, NNSIGHT_PP_TLS_SWAP=0 opts out) isolates it.
Open-ended tracer.iter[:] is paced by a driver-served step gate instead of spinning the thread, which also covers single-rank engines.
The vLLM prefix cache is disabled by default in both engine modes: a cached prefix presents a truncated prompt slab to hooks, so prompt-position reads and writes were silently wrong.
The within-stage tensor-parallel gather (VLLMFragments) rides the PP interleaver: the base interleaver's serving loop is split out of handle into serve so the publish sees the assembled whole while the model still gets the re-split piece.
Reply-path refusals for mixed-dtype containers, error delivery across the wire onto the worker greenlet, and real-dtype stamping on pull replies.

Validation

CPU PP tree (tests/vllm/pp/, no GPU): 97 passed, 3 skipped, 1 xfailed. Covers two- and three-stage parity against single-rank references, TP=2 x PP=2 topology, ownership derivation, lazy-tensor semantics, merge hardening, error paths, the listener protocol, and the step gate.
Full CPU suite unchanged against the 0.8 base (one pre-existing failure requires transformers >= 5.15.0, an environment constraint).
GPU PP=2: smoke suite and parity suite green; an interp-workload benchmark sweep (logit lens, steering, ablation, activation patching, generation steering, attribution) matches single-rank references at top1=1.00, total variation 0.000 per cell.

🤖 Generated with Claude Code

https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

An interleaver on a pipeline-parallel rank owns only its stage's modules, so a worker must not park on a location no local hook will ever fire. Mediator.event now consults interleaver.intercept(mediator, event, location, rest) before parking: the base interleaver never intercepts (returns None); a distributed override can serve the event in place — a remote-owned read answered with a lazy handle, a remote-owned write absorbed. The 1-tuple return distinguishes 'serve None' from 'park normally'. pp.py ports the derived ownership map from the 0.7 PP branch: per-module owning stage from the load-time meta exchange, longest-owned-ancestor resolution, memoized lookups, unknown-path-is-local default. The readiness-gate timeouts do not port — the gate existed to pace preemptible worker threads against the forward thread, and workers here are parked greenlets whenever the forward runs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Producer side ports from the 0.7 PP branch nearly unchanged: fixed-size self-identifying requests on a shared tag, self-describing replies (true dtype stamped on the wire) on per-pull tags, parked serving of not-yet-produced values, a bounded reply pool, error replies for unserializable or abandoned pulls, scoped buffer clears, and the rank-ordered drain barrier. The listener thread is legal under the greenlet engine: it never touches a greenlet. Consumer side is redesigned from a blocking worker-thread recv into a split-phase Pull: begin_pull sends the request and arms a waiter-pool thread that blocks in the reply recvs and assembles the value on CPU; complete() collects it at a serve point (device placement on the collecting thread) and re-raises any carried error there. Issuance happens when a forced lazy parks, so the transfer overlaps the remainder of the forward; probed on this torch build, gloo irecv handles cannot replace the waiter (is_completed() never flips outside wait(), and wait(timeout) expiry closes the whole peer pair). tests/vllm/pp/test_listener_protocol.py round-trips the protocol across two spawned gloo ranks on CPU: buffered serve, parked-then-dispatched serve, tuple and int32 round-trips, producer-side error reply, abandoned-pull error reply on scoped clear, and the drain barrier. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…sh/serve) LazyRemoteTensor ports the 0.7 proxy surface (torch-function dispatch, operator/comparison forcing, iteration/len materialization, deferred-index children, tuple-output guidance, write absorption) with one structural change: materialization parks the worker on an encoded pull location instead of calling a blocking pull closure on a worker thread. Deferred children hold a parent reference rather than a closure, and .save() is no longer overridden: the mounted save marks the proxy's id like any value, so the saved name ships as a NOT_ON_THIS_RANK sentinel for the engine-side merge. PPInterleaver implements the core intercept seam for one PP rank: a read of a remote-owned location is answered immediately with a lazy (no traffic; an unconsumed read costs nothing), a write to one is absorbed, and a forced lazy's park issues its cross-stage pull at that exact moment so the transfer overlaps the remainder of the forward. The producer side publishes each request's rows of the post-intervention value under (provider.i{visit}, req_id) from handle(), dispatching parked peer pulls. serve_pulls() drains at a serve point: complete each worker's pull, switch the worker back in, repeat while resumed workers force further lazies; a failed pull is thrown into the owning worker (recovering workers keep their new park) and recorded, tearing down only when not deferring. Occurrence tags mirror the owning rank's visit counts: pinned reads use the tracer.iter step, relaxed reads use the interleaver's forward count. A module visited several times in one forward is mis-tagged on the non-owning side, the same limitation as the 0.7 branch. tests/vllm/pp/test_pp_interleaver_crossstage.py drives the whole chain over two spawned gloo ranks with no vLLM: a replicated block reads both stages' modules, forces the remote one on each side, and must finish with the same value on both ranks. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

load_model builds the PP stack before the envoy tree when the PP world is >1: the load-time meta exchange (allgather of each rank's real-module names + dtypes over the PP cpu group) derives ownership, structural last-rank claims cover the build-everywhere/fire-on-last modules (logits, samples, logits_processor, with Llama's missing non-last-rank logits_processor stubbed for persistent-id symmetry), a dedicated per-TP-column gloo pull group carries the listener traffic, and the envoy tree is constructed over the resulting PPInterleaver via Envoy's interleaver= kwarg. Request ingest stamps each deserialized worker with its request id, scoping cross-stage pulls and publishes per request. execute_model gains the two serve points with the pipeline-order discipline: a BLOCKING serve before the step (stragglers from the previous step — the whole pipeline has finished it) and a NON-blocking serve after the forward (a downstream stage's value is produced only after this method returns, so waiting there would deadlock the pipeline); serve_pulls grows the block= parameter accordingly. collect finalizes on every rank: blocking serve, drain barrier, then a scoped buffer clear for the finished requests only. Known gaps for follow-up commits: per-stage save shipping + engine-side merge (collect still ships from PP rank 0 only), and the meta-envoy graft under PPMissingLayer stubs (sub-stub paths like layers.5.attn do not yet resolve at deserialization on non-owning ranks). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

collect.py ports the merge machinery from the 0.7 PP branch: strip_lazy turns saved lazies into NOT_ON_THIS_RANK sentinels (a name purely owned by another stage is skipped; the owner ships it), and merge_saved assembles the per-stage partial save trees — dicts union by key (disjoint per-stage cache keys combine), lists union position-wise with overshoot-tail dropping and a warning for one-sided real entries, equal-length tuples rebuild NamedTuple-safely, and real leaves win over sentinels, with PPRankDivergenceWarning on genuinely different reals (in-trace RNG, per-rank environment values). merge_collected adapts it to 0.8's payload shape ({engine_id: {saves, error}} per rank): the runner now ships from every PP stage's TP-rank-0 instead of PP rank 0 only, and all four consumers (sync engine step, async attach/free, serve handler) merge the per-rank payloads instead of taking the first non-None one, which under PP silently dropped every stage but one. CPU unit tests cover sentinel preference, positional list union, disjoint dict union, the divergence tripwire (silent on equal reals), ownership stripping, and cross-payload assembly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

A stub has no children, so sub-stub paths (model.layers.5.attn on a rank that doesn't own layer 5) neither resolved at request deserialization nor were reachable in a block. The worker builds a full meta-device copy of the architecture before the real distributed groups exist (bootstrap PP=1/TP=1 env, dummy loader, torn down after), hands it to the runner, and the runner walks the envoy tree grafting each stub's meta children via Envoy._wrap_envoy (which handles recursive construction and shadowed names). Grafted envoys wrap meta modules that never run; reads on them resolve by ownership to lazies like the stub itself. The graft runs before the persistent-object map is built, so grafted paths resolve for deserialization. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ite green A worker that forces an upstream value mid-block can still be parked on its pull when the end-of-forward (non-blocking) serve runs, if the transfer hasn't landed. Its next park is typically this step's logits — offered exactly once, in sample_tokens — so under max_tokens=1 the worker missed them forever and surfaced a spurious out-of-order error at collect. sample_tokens and _sample now complete in-flight pulls (blocking) before offering logits/samples, which is deadlock-free there: sampling runs on the last stage, every other stage is upstream, so the wait is transfer-only. tests/vllm/pp/test_pp_smoke.py passes end to end on a 2-GPU Qwen2.5-0.5B engine (vllm 0.19.1): cross-stage reads from both stages with the logits merge (greedy decode intact), a cross-stage write on a stage-0 layer changing the logits, and multi-token generation collecting a stage-1 layer per step through tracer.iter. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Each scenario runs identical intervention code at PP=1 and PP=2 in separate subprocesses and compares logits argmax plus cosine similarity: final logits, early/late hidden states, a stage-local write, three-token generation, and two concurrent invokes with different-length prompts (per-request narrowing, request-keyed pulls, and the save merge keeping both invokes distinct). One strict xfail pins a semantics regression against the 0.7 thread-based PP: a write that follows a forced cross-stage read loses its swap window on the owning rank. The worker parks on the pull inside the target layer's hook, pulls are only served at serve points outside the forward, and by the time the worker resumes the forward has passed the write site, so the swap raises OutOfOrderError. The thread implementation supported this because the hook blocked until the worker's pull completed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

…indow A worker that forced a cross-stage value used to park unconditionally, resuming only at a serve point outside the forward. For a value owned by a downstream stage that is the only correct option: the value does not exist until this rank's forward returns, and blocking would deadlock the pipeline. But for an upstream-owned value the pipeline order guarantees the owning stage already produced and published it before this stage's forward started, so the wait is transfer only. Parking there surrendered the swap window: a block that read a local layer's output, forced an upstream value, and wrote that same layer's output resumed after the forward had passed the write site and raised OutOfOrderError. The 0.7 thread-based PP supported this pattern because the hook blocked until the pull completed. The intercept now distinguishes by source rank: upstream forces block on the worker right where they are and return the value without parking; downstream forces keep the issue-at-park path. The parity suite's cross-stage read-modify-write test drops its expected-failure mark and passes. The scenario itself needed two corrections the suite surfaced: the early-layer read must clone (holding the live activation across the intervening layers lets vLLM's buffer reuse mutate it; the pull path clones at publish, so the two PP sizes were grafting different values, logits cosine 0.94), and the graft must scale to the residual element's norm (on Qwen2 the layer tuple is (hidden, residual) with the stream riding the residual, so a perturbation at hidden's own scale moves the logits by less than the parity threshold). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

Wire level: a value that overflows the reply's shape header comes back as an error reply the consumer raises, never a hang or a desynced recv. A strict xfail pins a gap: a mixed-dtype tuple is neither error-replied nor faithfully delivered — torch.cat promotes instead of raising (the same-dtype validation _serve_reply's comment claims does not exist), the shape header carries only the first tensor's dtype, and the consumer rebuilds every element in it, so an int64 element arrives as float32. Interleaver level, against a live two-rank listener pair: a failed downstream pull is thrown into the worker at the line that forced the value, where user code catches it, pulls again, and the drain loop serves the retry; an uncaught failure with defer_exceptions set unwinds only its own mediator while a concurrent healthy one finishes (one request dies, the engine survives); with defer_exceptions off the failure propagates out of serve_pulls after being recorded; and the upstream in-place serve raises at the force line directly, catchable in the block, with no park involved. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

…ilently promoting The reply's data message is one flat tensor and the header carries one dtype, so every element of a container value must share it. The flattening cat cannot enforce that: torch.cat silently promotes a mixed input, and the consumer rebuilds every element in the first tensor's dtype. The reachable case is a remote layer's .inputs, which bundles int64 positions with bf16 hidden states; positions above 256 are not exactly representable in bf16, so they arrived corrupted with no error anywhere. An explicit same-dtype check before the cat turns the mismatch into an error reply naming the location and telling the user to read the elements separately. Supporting mixed containers instead would take a per-tensor dtype code in the header plus a single uint8 byte-blob data message; the comment records that path. The mixed-dtype wire test drops its expected-failure mark and asserts the error reply. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

Pins the proxy contract in a single process, no wire. Creation, indexing (deferred children), absorbed writes, and saves pull nothing, and a proxy leaked past its trace raises instead of hanging. Any real read materializes: arithmetic operators and their reflected forms, torch functions (with the proxy anywhere in the arguments), method and metadata access, elementwise comparisons (the identity fallback would silently diverge between ranks), iteration and unpacking and len (the sequence-protocol fallback never raises IndexError, so tuple(lazy) would spin forever on the non-owning rank). Materialization happens exactly once per proxy tree: repeated reads and sibling children reuse one pull, verified by counting Mediator.value calls, and the pull location round-trips through the codec including the park's doubled occurrence tag. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

Structure preservation: strip_lazy and merge_saved keep namedtuple types, recurse into dicts, union nested lists inside dicts, prefer real over sentinel at shared keys, and copy through dict.__getitem__ so a dict subclass with overridden lookups cannot corrupt the merge. Divergence tripwire: low-order float noise and matching NaNs merge silently; genuine divergence warns with the slot label (including nested list and dict positions) and the magnitude; shape, dtype, tensor-vs-scalar, and list-vs-tuple clashes warn and keep the later rank's value; a user type whose comparison raises merges silently instead of false-positiving. Length tolerance: one-sided real entries warn as a stalled worker and the complete side is kept, an empty list yields to a populated one, and a trailing container of nothing but sentinels drops as overshoot. merge_collected: three stages' owned slots all arrive, disjoint save names combine, the first non-None error wins. A strict xfail pins a defect the three-stage case surfaced: the pairwise fold truncates the trailing sentinel slot after the first two payloads merge, so the third stage's real entry at that position counts as one-sided and trips the stalled-worker warning spuriously. Values merge correctly; the drop is only valid after the last payload folds in, so the warning is wrong at three or more stages. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

Ownership comes from where modules are real: derive_owners reduces the allgathered per-stage real-module lists with the exactly-one-stage rule (a path real on several stages is dropped and later resolves as local). The reduction moves out of the runner's allgather method into pp.py unchanged, so it is callable without a distributed group; the runner now calls it. The tests run the derivation and PPModuleMap resolution against module trees named decoder_blocks / word_embeddings / output_projection, so any naming-convention dependence would fail: envoy-path and raw-name lookups, eproperty suffix stripping, submodule inheritance from the nearest owned ancestor, the unknown-path and empty-map local defaults, a custom root path, and re-installing owners replacing memoized results. is_pp_missing detects the stub by class name alone; resolve_meta strips exactly one root component and never more. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

PP=3 on Qwen2.5-0.5B (stages hold layers 0-7, 8-15, 16-23): reads at layers 2, 12, and 20 land one value on each stage and all match the single-GPU reference, so the payload merge holds across three ranks; the cross-stage write pulls from a non-adjacent stage (the last stage reads the first stage's layer, past the middle one) and the grafted logits match the reference. TP=2 x PP=2 against a TP=2 single-stage reference, so the comparison isolates what PP adds under sharding: cross-stage reads and the grafted write both match, which also covers saves shipping from each stage's TP-rank-0 and pull traffic riding the per-TP-column groups. The parity worker gains a --tp argument and a three-stage read scenario. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

…s finished The step-start serve drained every parked pull, and the sampling-path serves (which vLLM invokes on every rank, contrary to the comment that assumed the last rank only) did the same. A worker chaining per-step forces under tracer.iter re-parks, from inside the serve, on the NEXT round's value; that value is produced by forwards the serve was holding up, so the serve waited against its own rank's progress until the pull deadline. Every multi-token trace that forces a cross-stage value per step deadlocked this way (found by the interp-serve-bench generation-steering cell; the earlier multi-token test appended unforced lazies, which ship as sentinels and pull nothing, so the pattern had no coverage). The pipeline schedule decides locally which pulls may be waited for: the engine schedules a request's round k only after round k-1 sampled, and sampling runs on the last stage after every stage finished round k-1, so when a rank opens round k all stages have completed rounds 0..k-1 for the request. The runner now counts completed rounds per request, and per-step serves complete a not-yet-arrived pull only when its occurrence is below that count (occurrence tags and round counts share the sampling-round clock for pinned iter reads). Later-round pulls stay parked; their transfers are already in flight and the round's own serve boundary completes them. A worker several rounds behind still catches up in one serve, since each chained pull it re-parks on is a past round until it reaches the current one. Collect keeps the unconditional drain: no rounds remain there, and what is left resolves by reply or by the peers' finalize clears error-replying overshoot pulls. The parity suite gains the forced per-step read scenario that deadlocked. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

Materializing an unforced LazyRemoteTensor inside __torch_function__ parks the worker greenlet while torch's dispatcher frames are live on its stack. Every torch binding installs thread-local state on entry (the pybind warning handler is a pointer into this greenlet's C stack) that stays installed across the switch; the forward that then runs on the thread segfaults in C++ the next time that state is consulted. Observed as a silent rank-0 worker death, SIGSEGV in c10::warn under ProcessGroupNCCL::send in the core dump, minimal repro torch.ones_like(lazy) in any PP trace. Parks from plain Python frames (the proxy's own operators, methods, properties) carry no dispatcher state and are safe, which the bisection confirmed case by case. The guard raises before parking, with the workaround in the message: force the value with a method or with it leading the expression, then call the torch function. An operator with a plain tensor on the LEFT routes through the same torch machinery, so the cross-stage harness block now leads with the lazy. A deferred child of an already-pulled parent passes the guard, since indexing a cached value never parks. Raising is state-safe where parking is not: the exception unwinds the dispatcher's RAII guards in order on the worker's own stack. Also in-tree, env-gated off (NNSIGHT_PP_TLS_SWAP=1): a prototype that swaps the Python-reachable torch thread-local bundle (grad mode, dispatcher include/exclude key sets) at every greenlet switch via greenlet.settrace, standing the guard down while active. The canary run answers the coverage question: the crash persists under the swap, so the poisoned state is outside the Python-reachable bundle, consistent with the warning handler, and the full per-greenlet isolation fix requires a C binding to at::ThreadLocalState capture/replace. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

NNSIGHT_PP_TLS_SWAP=1 swaps torch's thread-local state at every greenlet switch and throw via greenlet.settrace, giving each worker and the forward their own effective copy — the isolation 0.7's real threads provided. The bundle is at::ThreadLocalState plus the c10 warning handler, captured by a small extension JIT-built against the installed torch on first use (torch.utils.cpp_extension, no link configuration in-tree). The Python-reachable bundle (grad mode, dispatcher key sets) was prototyped first and failed the canary: the crashing word is the warning handler, a pointer every torch binding installs on entry that aims into its caller's C stack — memory greenlet slices away on a switch — and it has no Python surface. With the C bundle the canary passes: torch.ones_like on an unforced cross-stage value per generation step, the exact construction that segfaulted the rank-0 worker in c10::warn under ProcessGroupNCCL::send, runs the full steered generation with correct values, as does the original crashing write repro. While the swap is active the __torch_function__ materialization guard stands down (already wired: the guard consults the loaded state, not the env var, so a failed build keeps the guard up). Measured cost: 4.6us capture + 2.9us restore per switch, tens of switches per trace, sub-percent at trace level. Off by default until soak coverage widens; the guard remains the default protection. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

The swap is the fix, so it runs whenever a PP engine does: the runner installs it on the forward thread at every engine, NNSIGHT_PP_TLS_SWAP=0 disables it (leaving parks inside torch calls unsafe), and a failed extension build raises at engine start instead of running unprotected. The __torch_function__ materialization guard is deleted along with its tests: it existed only to convert the crash into an error while no fix existed, and with the swap always on, forcing a cross-stage value inside torch's dispatcher (torch functions, and operators with a plain tensor on the left) is simply correct again, restoring 0.7 behavior. The cross-stage harness block returns to its natural operand order, which now exercises dispatch-path materialization. Validated under the default: the canary (per-step torch.ones_like on an unforced cross-stage value) runs the full steered generation with correct values, GPU smoke passes 3/3, CPU tree 78 passed and 1 expected failure. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

NNSIGHT_PP_DEBUG_STACKS=1 registers faulthandler on SIGUSR1 at runner import, so kill -USR1 a wedged worker prints every thread's stack to its log. Needed on hosts where ptrace is restricted (py-spy and gdb attach are refused). First use diagnosed the generation-patching PP hang: rank 0's main thread wedged inside the garbage collector during a lazy-tensor indexing allocation, rank 1 waiting in vLLM's intermediate-tensor receive. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

An open-ended iter loop has no termination condition of its own: the model stopping ends it through a dangling park, which assumes every step's body parks at least once. A body that never parks spins the thread forever inside Iterations.__iter__. Under pipeline parallelism ordinary bodies become park-free on non-owning ranks, because reads of remote-owned locations return lazies immediately; the interp-serve-bench generation-patching cell, whose per-step body is a logits read appended to a list, deadlocked both ranks this way (rank 0's main thread spinning at iterator.py:126, rank 1 waiting in vLLM's intermediate-tensor receive; diagnosed via /proc thread states, a gc.disable run that eliminated the collector hypothesis, and SIGUSR1 stack dumps). On a single-rank engine the same spin exists for any park-free body, so the fix is in the iterator's contract. The mediator counts its parks; after an open-ended step whose body did not park, the iterator parks on STEP_GATE, which the runner serves once per execute_model, so the loop advances at the model's pace. Generation ending leaves the gate read dangling, and the standard dangling-worker unwind ends the loop, phrased as "generation ended before the loop's next step". The PP interleaver does not publish the gate. Bounded loops and parking bodies are untouched, and the whole previously green set stays green (full CPU suite, PP tree, smoke, multi-token parity, generation-steering spec). Also env-gated NNSIGHT_PP_DEBUG_NOGC=1: run workers with the cyclic collector off, the diagnostic that separated the spin from the collector. The generation-patching spec has two further defects behind this one, under investigation: a prefill injection binding to a wrong 3-token slab, and a later engine-core timeout wedge. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

…visibility Cached prefix blocks skip the forward for the tokens they cover, so on a repeated prompt the hooks see only the uncached tail: a 19-token prompt after one repetition presents a 3-token slab (19 minus one 16-token block). Prompt-position reads and writes then bind to the wrong rows with no error anywhere; last-position reads stay correct, which is why single-forward cells kept scoring clean while the generation-patching cell's prefill injection hit its own length guard on every repetition after the first. The 0.7 branch disabled the cache for exactly this reason and the port dropped it; the benchmark cell's docstring even records the old contract. Off by default in both engine modes now; an explicit caller setting still wins. Also re-pin the step gate to its step before parking: a body that created a lazy relaxed the iteration pin without parking, and a relaxed gate park is tagged from a count the driver's serve loop has not advanced yet, so it lands on the same location and the serve loop re-serves the worker forever (the second generation-patching wedge, caught by SIGUSR1 stack dump inside Mediator.handle). Pinned, each gate park is a new location and one serve releases exactly one step; a new test covers the pin-relaxing park-free body via an intercepting interleaver. With both changes the generation-patching spec scores EQUIVALENT on both cells at PP=2 (top1=1.00 tv=0.000), completing the benchmark's default PP pass. Regression green: smoke, logits and both multi-token parity tests, single-rank tracing suite, PP CPU tree. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018XVbQoV7jMDQftovuoZToD

The base interleaver's mediator-and-cache loop lives in a serve method; handle keeps the fragments bracket around it (gather, serve, re-split). serve is the point where the post-intervention whole value exists, which a distributed interleaver needs to see. PPInterleaver takes a fragments argument, forwards it to the base, and publishes from serve: the pull buffer holds the assembled whole a local worker saw, and handle hands the model the re-split piece. A two-rank test drives a fake Fragments (whole = piece * 2, no collective) through the real pull machinery and asserts all three sides: the local worker reads the whole, the peer's pull receives the whole, the model's forward continues from the re-split piece. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

… workers on every rank collect_nnsight's blocking pull serve covers exactly the finished requests' workers (serve_pulls takes an `only` request-id filter); a concurrent request's parked pulls stay parked and are resumed by its own step serves. An exception a worker raises while collect's serve resumes it is captured onto the request's deferred error, so it reaches the client alongside the saves. Finalization (finish_dangling, releasing finished workers, dropping errored entries) runs on every rank; only the payload construction and return are gated to the shipping rank (each stage's TP rank 0 under PP, the single rank otherwise), so a long-lived engine's non-shipping ranks release their mediators, saved tensors, and worker greenlets as requests retire. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

…n-place upstream serves A pinned tracer.iter step tags every read of its body, so all of a step's cross-stage reads name the same round. An unpinned read tags with the request's completed-round count, which advances in step with the owning rank's per-request visit counter for once-per-step modules; a run outside the engine counts its own forwards. An in-place upstream serve blocks the forward thread, so it is limited to rounds this rank has opened, where the value exists and the wait is transfer only. A pull for a later round (a loop body running ahead of the model) parks like a downstream pull and is resumed by the serve point once its round has run. The .i{n} parse shared by the serve gate moves into _occurrence. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

A saved tracer.cache() ships one CacheView per stage, each holding the modules that fired locally. merge_saved merges the views' entry dicts: disjoint module paths union, a path both stages recorded merges entry-wise (Entry fields recurse, so tensors compare under the standard tolerance). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

The producer-side publish is gated per mediator on the same string match handle() serves by: a location enters the pull buffer only when some worker's park names this visit of it. Every rank runs the same block, so a peer's pull always has a matching local park on the owning rank, and locations nothing reads cost neither the clone nor the buffer entry. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

…erves already seen An open-ended tracer.iter[:] whose step body never parks waits on the step gate, and the gate is served once per generation step. The base interleaver serves it as the root module's forward completes, which is the step boundary on every backend; a driver that owns the boundary itself (the vLLM runner, whose forward may replay as a captured graph with no hooks) constructs the interleaver with step_gate_at_root=False and keeps serving the gate once per engine step. A gate park pins to the count of gate serves the worker has already seen, and at least one past the loop's previous gate park, so serves that pass while the worker is parked elsewhere (a save before the loop) are skipped over and each serve releases exactly one step. A worker left parked on the gate when the run ends is the loop's exit: both dangling checks unwind it without a warning, keeping the reached steps' values. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

…-gated modules by meta diff The PP meta model builds from a deep copy of the engine's vllm_config, narrowed to one rank and the meta device, so every loading option the user gave (trust_remote_code, revision, hf_overrides, quantization) shapes the meta tree exactly as it shaped the real model. Rank-gated modules resolve architecture-agnostically: any module the meta tree has and this rank's tree does not gets a PPMissingLayer stub (stub_rank_gated_modules), so a request serialized against the client's full meta tree deserializes on every rank whatever the architecture gates behind its rank checks. Modules under a stage's stub stay with the meta-envoy graft. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

…he pull gates Both consumers of the per-request round count (the in-place upstream serve's guard and the relaxed-read occurrence tag) treat a request the runner has not yet counted a round for as being in round zero, which is what its first forward is. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

… produced rounds collect_nnsight names the finished requests' workers when serving their remaining pulls: a later step's scheduling rebuilds the interleaver's per-step mediator list, so the finalize serve reaches them through the request registry (serve_pulls takes a `mediators` argument). The serve keeps the produced-round gate: a pull for a produced round waits on its transfer only, and a pull past generation end stays parked and is unwound by finish_dangling as the loop's exit. A released worker's in-flight pull records are dropped with it (discard_pulls); the owner's clear_buffer error-replies their wire requests. finish_dangling phrases a pull park by the provider it was pulling. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

khaiwang and others added 30 commits August 19, 2026 20:26

style: plain phrasing in the fix series' comments and docstrings

bbb5704

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015fELrvmHWpBLcv7akbQ1oR

khaiwang force-pushed the pp-on-08 branch from 29c2ed9 to bbb5704 Compare August 26, 2026 18:26

Navigation Menu

Uh oh!

Pipeline parallelism for the vLLM runtime - #691

Pipeline parallelism for the vLLM runtime#691
khaiwang wants to merge 32 commits into
0.8from
pp-on-08

khaiwang commented Aug 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

khaiwang commented Aug 20, 2026

How it works

Validation

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant