From 0203a2e323ea5f53e959c4c5d1bf5b3d89f67258 Mon Sep 17 00:00:00 2001 From: youhengfei <474029121@qq.com> Date: Fri, 23 Oct 2026 11:52:54 +0800 Subject: [PATCH] [Doc][M9] PD and KV transfer walkthrough Add a maintained architecture document for PD disaggregation: component ownership, the reserve/prefill/commit/received/ack/activation/release handoff, prefix and layout handling, reservation credits, cancellation and failure ordering, lease and preemption interaction, profiling control routing, and supported scope. Index it in docs/README.md. Follows the maintained kv_cache_walkthrough.md style: current-state mechanics without baseline pinning or review framing. Refs #32, #111. Signed-off-by: youhengfei <474029121@qq.com> --- docs/README.md | 1 + docs/pd_walkthrough.md | 226 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 227 insertions(+) create mode 100644 docs/pd_walkthrough.md diff --git a/docs/README.md b/docs/README.md index c3507a8a..6dec2b06 100644 --- a/docs/README.md +++ b/docs/README.md @@ -20,6 +20,7 @@ an environment to receiving the first generated response. The | [FlashAttention](https://github.com/InferMatrixLab/vllm-rlt/pull/30) | FA2/FA3/FA4 installation, hardware selection, and constraints | | [Cache and scheduling features](https://github.com/InferMatrixLab/vllm-rlt/pull/31) | Prefix reuse, incremental KV, priorities, and preemption | | [Prefill/decode disaggregation](https://github.com/InferMatrixLab/vllm-rlt/pull/31) | Single-host GPU worker pools and NIXL transfer | +| [PD walkthrough](pd_walkthrough.md) | Prefill/decode disaggregation: ownership, KV/hidden handoff, cancellation and scope | | [HTTP serving](serving.md) | Completions API, streaming, and service lifecycle | | [Profiling](profiling.md) | Core profiler options, Python/HTTP controls, PD collection, and per-rank archives | | [Ouro-Thinking recipe](recipes/ouro_thinking.md) | Shared 1.4B/2.6B setup, JSON/SSE client, budgets and validation limits | diff --git a/docs/pd_walkthrough.md b/docs/pd_walkthrough.md new file mode 100644 index 00000000..e93e19a3 --- /dev/null +++ b/docs/pd_walkthrough.md @@ -0,0 +1,226 @@ +# PD disaggregation and KV transfer + +Prefill/decode (PD) disaggregation runs the prompt phase and the generation +phase in separate stage processes on the same host, each with a complete model +replica and its own KV pool. A CPU coordinator routes requests across them; +NIXL WRITEs move KV and the final prompt hidden state directly between +registered worker memory. Control channels carry only metadata and outputs. + +Stage is the deployment and resource boundary; loop is the local scheduling +unit. This page covers the cross-stage handoff; local execution follows +[unified execution](unified_execution.md) and the +[loop scheduler](scheduler.md). + +## Ownership + +| Component | Owns or performs | +| --- | --- | +| `PDEngine` (coordinator) | External request ID → transfer generation mapping, peer assignment, admission reservations, cross-worker ordering, output mapping, aggregate profiling | +| `StageClient` | Worker process, control channel, readiness, generic control replies; connected-worker shutdown is ordered by the coordinator | +| `StageEngine` | Local request lifecycle, import binding, compute release and asynchronous retirement; coordinates scheduler and runner operations | +| `PDStageConnector` / `RequestTransfer` | Transfer slot, destination mapping, chunk sequences, commit/receipt flags, cancellation state, safe retirement | +| `NixlConnector` | Registered tensors, segment descriptors, NIXL handles, notifications, in-flight bytes | +| Scheduler | Request registry, phase queues, admission, local selection; never invoked directly by transport code | +| KVCacheManager | Allocation, block mappings, prefix validity, references, transfer leases | +| Executor / runners | Execution, persistent buffers, stream/event dependencies, result retirement | + +`RequestTransfer.request` borrows the stage-local `Request`; it is not a second +copy of loop/token progress. Connector-owned hidden staging is separate from +runner storage: the prefill stage copies the final prompt hidden state into a +staging slot, and the decode engine binds its received slot for coda. + +Code: [coordinator](../vllm_rlt/pd/engine.py), +[connector](../vllm_rlt/pd/connector.py), +[transport](../vllm_rlt/pd/transport.py), +[connector protocol](../vllm_rlt/engine/connector.py), +[stage engine](../vllm_rlt/engine/stage_engine.py). + +## Configuration and startup + +Deployment entry points create and validate `ModelConfig` before spawning +workers. It retains the concrete model configuration and delegates vocabulary, +loop and context limits, including derived properties. `PDEngine` uses those +limits for request validation and passes the configuration through each +`StageClient` to its worker. Spawned processes receive copies, not shared mutable +configuration. + +Checkpoint workers use the supplied concrete configuration to load weights +without reparsing JSON. Config-only inputs construct models; module inputs retain +the supplied weights. Models and `StageEngine` use the concrete configuration. +Unsupported PD model families are rejected explicitly before workers start. + +Workers report runtime results in `ready`: actual KV capacity and storage depths, +admission limits, registered transport information and compatibility fingerprints. +The coordinator does not reconstruct model configuration from this message. It +still waits for every worker, checks model/weight/KV compatibility, connects peers +and waits for connection acknowledgment before accepting work. Entry configuration +alone does not establish runtime compatibility. + +Code: [model configuration](../vllm_rlt/models/config.py), +[stage client](../vllm_rlt/engine/llm_engine.py), +[worker startup](../vllm_rlt/engine/worker.py). + +## Handoff + +```text +Coordinator P stage D stage + | reserve ------------------------------------------------>| + |<----------------------- reserved(tables, slot, prefix) ---| + | prefill ---------------->| + | | shared admission / prefill + | | KV chunks + final hidden ---->| + |<-- prefill_complete -----| | + |<-- commit(count) --------| | + | commit(count) ------------------------------------------>| + | | all receipts + import binding + |<--------------------------------------------- received --| + | ack -------------------->| | + | | shared compute admission + |<-------------------------------------------- activated --| + |<-- released -------------| | + |<--------------------------------------- output/released --| +``` + +The decode stage reserves first: it allocates its full prompt-plus-output +mapping, pins a transfer lease, and replies with block tables, a hidden receive +slot and its own cached prefix length. Only then does the coordinator dispatch +prefill to the prefill stage. + +Completion evidence stays distinct at every step: + +1. **Submission.** The engine applies logical prompt progress and supplies a + completion event. Submission is not device completion. +2. **Staging readiness.** At an export boundary, `prefill_submitted()` enqueues + a stream wait on the producer event; the final chunk copies the prompt hidden + state and records a new event covering that copy. This is a device dependency, + not a blocking host wait. +3. **Bounded WRITE submission.** The connector posts at most one bounded chunk + per transfer per tick with round-robin ordering; `can_schedule()` applies + backpressure without touching scheduler queues. +4. **Commit.** All queued writes have been submitted. The advertised count is + not proof of NIXL completion. +5. **Receipt and import readiness.** D holds a known commit count and every + sequence in `[0, count)`, then `accept_import()` publishes KV validity, + binds hidden state and queues shared admission. D reports `received`, + which triggers P's `ack` independently of D's compute admission: a receiver + waiting for capacity must not hold sender KV. +6. **Activation.** The scheduler admits the imported request alongside local + work, unpins its lease and queues CODA. The connector observes the phase + change and reports `activated`. D owns first-token sampling and RNG. +7. **P retirement.** Requires the acknowledgment, committed submission, + released compute state and no outstanding NIXL handles for the generation. +8. **D retirement.** Local execution retires; the staging slot returns only + when safe. + +Receipt notifications may precede the control-plane commit. P release and D +activation are unordered relative to each other. + +## Prefill scheduling and export boundaries + +The scheduler chooses execution boundaries through `ExecutionBatch.loop_count`. +Ordinary prompt ranges can traverse their remaining depths in one submission; +staged scheduling exposes intermediate depth boundaries. These are local execution +choices, not additional cross-process scheduling commands. + +With `SchedulerConfig(wavefront_chunked_prefill=True)`, WCPB can batch ready ranges +from several chunks and depths. The scheduler applies grouped progress for each +request before the engine notifies the connector. The callback uses +`PrefillProgress.transfer_start` when provided, so its export origin is not inferred +from an arbitrary range in the batch. + +`num_prefilled_tokens` tracks the final-depth frontier, while +`prefill_reclaimable_tokens` advances at completed chunk boundaries. Intermediate +prelude/depth submissions do not by themselves queue a transfer. When +`progress.chunk_complete` is set, the connector queues the newly exportable KV +range under the producer event; `progress.prompt_complete` additionally stages +the final hidden state. The connector keeps at most one bounded WRITE submission +per transfer per tick, independently of how many local ranges shared a batch. + +Imported prompts initialize the same prefill frontiers before entering D's shared +admission path. WCPB requires LAST_EXITED KV; it does not change the receipt, +acknowledgment or retirement protocol. + +See [split and wavefront prefill](split-prefill-wavefront.md) for the hidden-state +ring, dependency frontiers and numerical limitations. Different BF16 batch/chunk +geometries are not guaranteed to produce identical tokens; transfer correctness +and cross-geometry numerical agreement are separate claims. + +## Ranges, prefixes and layouts + +P and D may hold different valid prefixes. D advertises its hit; P transfers +any additional prefix cached only on P plus the recomputed suffix, and +`queue_chunk()` clips the transfer start to D's frontier. Final hidden state is +always included: cached KV alone cannot produce the first sample. + +`plane_block_tables()` exposes one table per storage plane — per loop depth for +LAST_EXITED, a single plane for SHARED. `kv_segments()` maps a valid +`[start, end)` range to byte segments: full blocks cover all layers, partial +blocks emit per-layer ranges so unused tails are never advertised as +initialized. `partition_segments()` bounds bytes and descriptor count per +WRITE without an intermediate KV copy. + +Peer startup verifies model/weight/configuration fingerprints and KV geometry; +mismatches abort connection. + +## Reservations and statistics + +The coordinator charges both peers before dispatching D's reservation: +`depths * ceil(tokens / block_size)` blocks per peer (D additionally counts +`max_tokens - 1` output positions) plus one transfer slot each. These are +coordinator reservations, not measured worker usage; prefix hits and +incremental allocation do not reduce them. Credits return on each peer's first +`released` event and cannot be released twice. + +`PDEngine.stats()` reports `reserved_blocks`; physical usage belongs to +worker-side cache statistics. Frontend queries go through `request_ids()`, +`has_request()` and `stats()`. + +## Cancellation, failure and leases + +| Situation | Ordering | +| --- | --- | +| Waiting at coordinator | Drop the generation; neither peer owns resources. | +| D reservation outstanding | Cancel both. A late `reserved` gets `cancel_safe`: no prefill was dispatched. | +| P executing or transferring | Stop future scheduling; retain staging/KV until the final staging event and submitted NIXL handles finish. | +| D not fully received | Retain receive resources until P's release proves remote writers stopped, then `cancel_safe`. | +| D ready or active | Completed receipt proves remote writes ended; retire local work. Async speculative retirement may still delay slot release. | +| External ID reuse | Wire generations are unique UUIDs; local transitions also check `Request` identity, so retired messages cannot target a reused ID. | +| Transport error or timeout | Fail closed: no pool reuse or deregistration. Connected workers retain owners while reporting the failure; the coordinator terminates and joins P before D. | + +Transfer leases defer `KVCacheManager.free()` and exclude allocations from +preemption. Imports keep their leases while awaiting compute admission; +admission unpins them, after which normal supported preemption applies. + +Logical completion is not physical retirement: an asynchronous speculative +request may be `FINISHED` while submitted work still retains its storage. The +connector keeps its staging slot until the engine retires that ownership; +frontend `has_request()` alone is not a storage-lifetime check. + +## Profiling and control + +Generic host operations — profiling sessions, statistics, lifecycle — use the +shared `control`/`reply` channel with control and session identities, never +transfer routing. The coordinator fans profiling control out to all workers, +aggregates per-rank manifests and rolls back partial starts. + +## Scope + +The launcher supports same-host P/D pools with full model replicas; this does not +qualify cross-host transport. Model-family admission is checked before spawning, +while backend and execution compatibility are checked during worker initialization. + +Decode-side speculation requires LAST_EXITED, full target depth, fixed-depth +`ouro` exits and refill scheduling. Asynchronous speculation additionally requires +CUDA, greedy sampling, a backend supporting device-owned context lengths and no +preemption. Synchronous speculation does not impose that same preemption +rejection; transfer leases still protect receiving/exporting allocations. +Mixed scheduling excludes speculative rounds. + +SHARED completes all depths for each prompt position before advancing; staged +mixed/WCPB paths require LAST_EXITED. The prefill-only stage has its own prompt +scheduling policy, so its mode does not imply decode interleaving. + +See [runtime](runtime.md) and +[async speculation](async-speculative.md) for execution details and validation +limits. Feature availability is not evidence that every checkpoint, batching +geometry or feature combination has been validated.