Skip to content

[RFC]: PD coordination and KV transfer contracts #111

Open
Open
[RFC]: PD coordination and KV transfer contracts#111
Labels
RFCRequests for comments on major architectural changes or design choices

Description

@YouHengfei
YouHengfei
opened on Oct 9, 2026

Progress update — 2026-10-22 (revised)

Baseline: main 4a350344027960a48384626aa646f1765fda6d4f (post #118 unified stage execution, #107 scheduler-owned result application).

Resolved since the original posting:

Newly surfaced, out of scope here: PD teardown logs UCX cuDevicePrimaryCtxGetState errors after Python atexit; a standalone torch/NIXL WRITE reproduces it without vllm-rlt and it disappears with the UCX CUDA IPC mapping cache disabled. This is an environment/dependency lifecycle issue to be handled by a separate PR, not by M9 contract changes.

#119 remains a draft real-checkpoint P→D agreement harness; its BF16 batch/chunk non-invariance finding stands and bounds what "agreement" can claim. Same-host multi-GPU only; no cross-node RDMA qualification. FP8 KV transfer (cf. #129) remains unsupported.

Motivation.

Refs #32.

The pre-#118 PD handoff crossed module boundaries (a dedicated PD worker manipulated scheduler queues, Request fields and runner state; the coordinator faked frontend interfaces; handoff milestones were implicit). Most of that is now resolved structurally by #118. What remains is to agree on and record the contracts of the current implementation, fix demonstrated boundary defects, and establish validation for supported configurations — not to re-extract components that now exist.

The current M9 surface:

  • Coordinator: PDEngine (vllm_rlt/pd/engine.py) — peer routing, admission reservations, cross-worker ordering, output mapping, aggregate profiling.
  • Stage-side connector: PDStageConnector (vllm_rlt/pd/connector.py) — transfer bookkeeping, hidden staging, backpressure, receipt, safe retirement, via the TransferConnector protocol.
  • Transport: NixlConnector (vllm_rlt/pd/transport.py) — registered memory, segment descriptors, NIXL WRITE lifecycle.
  • Shared owners: Scheduler/KVCacheManager (M3/M4), StageEngine/Executor/runners (M2/M5), StageClient (M2 process/control).

Proposed Change.

Current behavior (at the baseline)

D reserves receiving resources before P computes; reservation pins KV but grants no compute slot. P runs shared stage execution, transfers valid KV in bounded chunks plus the final prompt hidden state. P and D may observe different prefix hits: transfer omits D's valid prefix and includes any prefix cached only on P; cached KV alone cannot produce the first sample, so D owns coda, first-token sampling and subsequent RNG.

Commit marks chunk submission, not delivery. After a known commit count and every expected receipt, D's engine publishes imported validity, binds hidden state and queues shared admission. D reports received, which triggers P's ack independently of D's compute admission; later admission produces activated. P retirement additionally requires local staging/compute release and no outstanding NIXL handles.

Cancellation retains staging/KV until local and remote accesses end; a final hidden copy can outlive the model event; async speculative retirement can outlive cancellation. NIXL failure/timeout fails closed: connected workers retain owners while reporting failure, and the coordinator stops/joins P before D. No recoverable quarantine or retry exists — the wording below is corrected accordingly.

Contracts to confirm (replacing the original extraction plan)

  1. Identity: external request IDs are reusable; the wire tid is a unique UUID generation. Local transitions additionally validate Request identity. Generic host control carries independent control/session identities.
  2. Reservation: destination tables, valid prefix and hidden slot stay stable while remote writes can reference them. Reservation is distinct from compute admission and from measured allocator usage; coordinator credits (reserved_blocks, slots) are exactly-once released.
  3. Data validity: transfer only valid [start, end) ranges across actual storage planes/layers (LAST_EXITED per-depth planes; SHARED single plane); preserve unequal prefix hits and the final-hidden requirement.
  4. Completion evidence: distinguish submission, source/staging CUDA completion, NIXL handle completion, destination receipt, import readiness, activation and retirement. Receipt, not activation, acknowledges P. Note: TransferConnector.acknowledge()'s docstring still says "downstream activation" — wording fix needed to match the receipt semantics.
  5. Lifetime: no recycling of staging, KV or execution storage while local tasks or remote writers may still reference it.
  6. Failure: stop new work, retain pending ownership, sender-before-receiver fatal termination; no implied recovery/retry.
  7. Replay/duplicates: generation isolation and exactly-once credit release; specify which repeated messages are tolerated (matching commits, repeated receipts) versus rejected (duplicate reservation, conflicting commit counts, late writes after readiness).
  8. Leases and preemption: transfer leases defer KVCacheManager.free() and exclude victims; imported requests enter normal admission/preemption only after the pinned transition. Preserve growth budgets and prefix references.
  9. Control plane: transfer-scoped messages are distinct from generic control/reply (profiling, stats, lifecycle); no transfer routing of session commands.

Ownership boundaries: use what exists

M9 consumes the existing contracts (can_schedule, prefill_submitted, reserve_transfer, accept_import, suspend_transfer, release_transfer_compute, release_transfer; scheduler admission; KV lease/validity operations) rather than defining replacements. Two open boundary questions remain for review — connector reads of the scheduler registry for retirement detection, and NixlConnector reading cache tensor geometry (pointers/strides). These are ownership/lifetime review items, not demonstrated defects; new read-only views are justified only if a concrete missing capability is shown.

Remaining work sequence

  1. Fix demonstrated defects in small reviewable PRs (done: [Bugfix][M9] Restore model-neutral PD coordinator metadata #140 model metadata).
  2. Align contract wording (acknowledge docstring) and document the wire vocabulary.
  3. Resolve the registry/geometry borrowing questions with evidence.
  4. Adapt Add a real-checkpoint P->D agreement test for the RFC #111 handoff contract #119's real-checkpoint harness to the current baseline and run the validation matrix below; coordinate with @jyizheng.
  5. Handle the UCX teardown issue in a separate PR (evidence above).

Each step stays runnable; no two authoritative stores of scheduler, KV or device state.

Compatibility and non-goals

This RFC does not enable multi-host PD, arbitrary stage graphs, WCPB, FP8 KV transfer, or new model integrations. Supported combinations follow docs/runtime.md: PD speculation is constrained (decode stage, fixed target depth, refill, no preemption; async adds greedy/device-context restrictions); SHARED requires atomic prefill; mixed excludes speculative rounds. Feature owners remain #43, #69, #81, #84; training/rollout (#70) must continue to observe stable PD behavior. Huginn PD is unsupported and now fails fast at coordinator startup (#140); enabling it is #69 roadmap work, not M9.

Alternatives considered

Impact and validation

This RFC has no runtime impact. Evidence recorded so far: #140's CPU regression tests (including a startup-path test that fails on the pre-fix line), full CPU suite 829 passed / 329 skipped, PD GPU suite 37 passed / 1 skipped on 2×H100, and paired GSM8K-100 runs token-identical PD vs colocated for Ouro-1.4B (69/100) and Nanbeige4.2-3B (63/100) under the pinned evaluation environment.

Implementation PRs must still cover:

  • Changed-file Ruff/format and the CPU regression suite; GPU cases on reserved devices only.
  • Commit/receipt reordering, duplicates/conflicts, stale generations, exactly-once credits.
  • Cancellation at reservation, staging, receipt-wait, activation-wait, decode and speculative-retirement boundaries.
  • KV coverage/mapping, prefix asymmetry, leases, growth budgets, safe reuse.
  • Generic control/session routing, partial profiling failure, worker failure, timeout, ordered shutdown.
  • Transferred KV/final hidden, first-token semantics, output/exit trajectories, RNG behavior.
  • Matched unprofiled performance for hot-path changes only; documentation/type-only work must not invent speed claims.

Agreement guarantee. #119's scoping stands: token IDs, exit depths and finish semantics match colocated execution at equal forward geometry; concurrent load checks liveness/shape only. The #140 GSM8K runs are a stronger-than-usual data point (single-request, fixed-depth, greedy: PD token-identical), but they do not generalize to concurrent BF16 batches or speculative paths.

Feedback Period.

Renewed after this baseline update; M1–M5 owners are asked to review the contract list above.

CC List.

@bjf-frz

Module-owner CCs: @carsontung666 @YzXiao101 @BuckMulligan99 @wuli666 @YZJF

Validation collaboration: @jyizheng (#119).

Anything else.

Decisions requested (revised):

  1. Confirm the contract list, especially received → ack independently of compute admission.
  2. Review [Doc][M9] PD and KV transfer walkthrough #142 (the maintained PD walkthrough) as the living documentation referenced by this RFC.
  3. Close the registry/geometry borrowing questions: evidence-based keep-as-is, or define owner views.
  4. Approve the Add a real-checkpoint P->D agreement test for the RFC #111 handoff contract #119 adaptation plan and the validation matrix division of work.
  5. Confirm the UCX teardown investigation goes to a separate PR (dependency/lifecycle), not M9 scope.
  6. Confirm [Bugfix][M9] Restore model-neutral PD coordinator metadata #140 covers the model-metadata regression and Huginn PD stays out of scope.

Activity

  1. YouHengfei
    added
    RFCRequests for comments on major architectural changes or design choices
    on Oct 9, 2026
  2. YouHengfei
    changed the title [-][RFC][M9]: PD coordination and KV transfer contracts[/-] [+][RFC]: PD coordination and KV transfer contracts[/+] on Oct 9, 2026

  3. jyizheng commented on Oct 11, 2026

    @jyizheng
    jyizheng
    on Oct 11, 2026

    Following up from #43, where I claimed the PD-integration slice of scheduling/distributed compatibility (speculative decode on D after the handoff) — that work lands exactly on the contracts this RFC defines, so I'd rather build on them than in parallel.

    Concrete offer: a P→D agreement harness as M9 validation. Identical-prefix runs against a colocated baseline, asserting token-for-token greedy equality plus the first-token/bonus and D-side RNG-ownership semantics described here. It would validate the activation/cancellation milestones on the current non-speculative path first, and then serve as the gate for lifting the "PD speculative decoding remains rejected" restriction. I have single-host multi-GPU H100/B200 (RDMA) matching the current scope.

    If that split works, record me for the validation piece — I'll start from PR #64's walkthrough at 3314c1b.

  4. YouHengfei commented on Oct 12, 2026

    @YouHengfei
    YouHengfei
    on Oct 12, 2026
    ContributorAuthor

    3. Following up from #43, where I claimed the PD-integration slice of scheduling/distributed compatibility (speculative decode on D after the handoff) — that work lands exactly on the contracts this RFC defines, so I'd rather build on them than in parallel.
    Concrete offer: a P→D agreement harness as M9 validation. Identical-prefix runs against a colocated baseline, asserting token-for-token greedy equality plus the first-token/bonus and D-side RNG-ownership semantics described here. It would validate the activation/cancellation milestones on the current non-speculative path first, and then serve as the gate for lifting the "PD speculative decoding remains rejected" restriction. I have single-host multi-GPU H100/B200 (RDMA) matching the current scope.
    If that split works, record me for the validation piece — I'll start from PR #64's walkthrough at 3314c1b.

    Thanks @jyizheng — that split makes sense to me. I’ll continue with the M9 walkthrough and ownership/contracts, and record you as the proposed contributor for the validation harness.

  5. jyizheng commented on Oct 12, 2026

    @jyizheng
    jyizheng
    on Oct 12, 2026

    The agreement harness from the comment above is up as draft #119, validated on 3x H100 (4 passed).

    One finding worth folding into the contract wording: bitwise agreement with colocated execution only holds at equal forward geometry. The colocated engine alone flips a seeded sample with batch composition and a learned-exit decision with prefill chunk size (bf16 tiling ulps — argmax survives, threshold comparisons and multinomial don't). So I'd state the M9 guarantee as "token ids, exit depths and finish semantics match colocated execution for a request at equal prefill geometry", with concurrent-load bitwise equality explicitly out of scope.

  6. YouHengfei
    added 2 commits that reference this issue on Oct 16, 2026

    docs: drop pd_refactoring draft superseded by RFC InferMatrixLab#111

    0afc2f1

    [Doc][M9] PD and KV transfer walkthrough

    db4922c

  7. 0z5a commented on Oct 25, 2026

    @0z5a
    0z5a
    on Oct 25, 2026
    Contributor

    Same-host KV transfer measurement for the current PD path (GH200, NIXL/UCX), in case it is useful for the transport side of this RFC.

    Setup.

    • main 66920da9c1, Ouro-1.4B, 1P on GPU 0 + 1D on GPU 1 of one GH200 node (NVLink between the two GPUs).
    • PDEngine defaults: 64 MiB chunks, 256 MiB in flight, default ExecutionConfig (packed prefill off); attention_backend="triton".
    • ISL 2048 random token ids, OSL 8, ignore_eos. Two workloads, a fresh engine each: seq (4 warm-up, then 16 requests one at a time) and burst (4 warm-up, then 32 requests added at once).
    • Ouro-1.4B full blocks are 1.5 MiB per K/V plane, so a 2048-token request is 1024 descriptors (1.61 GB) in 25 chunks.

    Raw NIXL WRITE, GPU 0 → GPU 1, at the connector's chunk shape (42 × 1.5 MiB): 89.6 GB/s with UCX defaults, 91.1 GB/s with UCX_MODULES=cuda,cma UCX_TLS=self,sm,cuda_copy,cuda_ipc,tcp. Both select cuda_ipc.

    PD cell transfer tail, median (d_received − prefill_compute_completed) TTFT, median mean chunk submit → observed done
    seq 20.9 ms (77 GB/s, 85% of raw) 761 ms 3.0 ms
    burst 2,823 ms 15,449 ms 90 ms

    1. With P idle, the handoff already runs at 85% of the raw link. On one host a different transport backend can only help if it is clearly faster than ~90 GB/s on this link. An earlier same-host NCCL P2P measurement here (different message pattern) also gave ~90 GB/s, which would save ~3 ms per request. We had considered an NCCL backend for same-host PD and are not pursuing it.

      • With the default (non-packed) prefill, a single 2048-token prefill takes ~735 ms, so the transfer is 2.7% of TTFT. With a faster prefill path that share would be larger.

    2. Under load, the handoff runs at ~2.4 GB/s aggregate (~1/38 of the link). In burst the median tail is 135× the idle one, and it sits directly before first_token.

      • The 256 MiB in-flight budget is effectively full: ~3.8 of 4 chunks in flight over the union of the tails.
      • Each chunk is seen complete only ~90 ms after submission, about one P model step; completions are polled from PDStageConnector.progress().
      • The bound budget ÷ completion latency (256 MiB ÷ 90 ms ≈ 3 GB/s) matches the observed rate.

      Candidates, not yet separated:

      • the shared in-flight budget;
      • completion being observed once per model step;
      • copies slowed by P's compute.

      A sweep of max_inflight_bytes (config only) plus a variant that polls completions while the stage waits on the GPU would tell them apart. Would that be in scope for M9? I can run it on the same setup.

    3. UCX teardown (your out-of-scope item). A standalone NIXL WRITE benchmark (no vllm-rlt) shows the same family of exit-time errors on GH200:

      • with the settings above it logs cuDevicePrimaryCtxGetState ... unrecognized error code 4 at exit;
      • with UCX defaults it segfaults at exit, after the data transfer completed and verified.

    Limits.

    • One node, one model, one prompt length.
    • The burst breakdown is observational, not a pre-registered test.
    • Our own NCCL microbenchmark in this run hung on the sender side, so the NCCL comparison rests on the earlier measurement.

  8. 0z5a commented on Oct 25, 2026

    @0z5a
    0z5a
    on Oct 25, 2026
    Contributor

    Follow-up on point 2 above: the load-time bottleneck is already being addressed in #150 (#150, draft).

    #150 fills the in-flight budget on each progress pass instead of sending one chunk per request, and it raises the default budget to 256 MiB chunks / 2 GiB in flight. Its measurements show the commit tail dropping from ~1 s to ~17 ms with 32 concurrent 2048-token requests. The sweep I offered is covered there, so I won't run it separately.

  9. YouHengfei
    added a commit that references this issue on Oct 26, 2026

    [Doc][M9] PD and KV transfer walkthrough

    0203a2e

  10. WEIXILIN0307 commented on Nov 1, 2026

    @WEIXILIN0307
    WEIXILIN0307
    on Nov 1, 2026

    Hi @0z5a and @YouHengfei, thank you for your work on FP8 KV cache and PD disaggregation.

    I’m interested in contributing FP8 KV cache transfer for PD, if this part is still available. I noticed that #129 currently excludes PD because the KV scales are not transferred. I’d like to build on that work by handling the transfer of K/V data together with their scales and ensuring they are ready before the decode side uses them.

    As a starting point, I’m considering same-host 1P1D with Ouro-1.4B. I have access to two GPUs with FP8 support and can test correctness, cancellation and transfer failures, and compare the PD path with the local FP8 path.

    Would this be a useful contribution? If so, I’d be glad to coordinate with you on the scope and interfaces before starting. Thank you!

  11. 0z5a commented on Nov 1, 2026

    @0z5a
    0z5a
    on Nov 1, 2026
    Contributor

    18. Hi @0z5a and @YouHengfei, thank you for your work on FP8 KV cache and PD disaggregation.
    I’m interested in contributing FP8 KV cache transfer for PD, if this part is still available. I noticed that #129 currently excludes PD because the KV scales are not transferred. I’d like to build on that work by handling the transfer of K/V data together with their scales and ensuring they are ready before the decode side uses them.
    As a starting point, I’m considering same-host 1P1D with Ouro-1.4B. I have access to two GPUs with FP8 support and can test correctness, cancellation and transfer failures, and compare the PD path with the local FP8 path.
    Would this be a useful contribution? If so, I’d be glad to coordinate with you on the scope and interfaces before starting. Thank you!

    Which gpus? it takes a huge memory to test PD disaggregation with high concurrency long context test, as you can see I use H100 on #129. And I already have an ongoing local PD experiment about the FlashLoop support issue.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    RFCRequests for comments on major architectural changes or design choices

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions