[shipiyouniao]

Extends the CPU offloading path from TP=1/2/4, PP=1 to validated single-node TP/PP layouts on vLLM 0.29–0.30. Pipeline stages may have different KV payloads: all ranks now share one chunk count, with native pinned CPU buffers sized for each rank.

Stacked on #598 (feat/offload-copy-pins). Its copy pins, retirement fence and ordered compute-release barrier remain the foundation.

Implementation and scope

• Before native KV initialization, derive each rank's logical KV payload and broadcast a capacity plan. Chunk count is floor(cpu_bytes / (blocks_per_chunk * sum(rank_payload_bytes))); reserved VMM storage size is not payload.
• A small native CPUOffloadingSpec subclass selects rank-private buffers and adapts the 0.29 num_blocks / 0.30 num_chunks API. Native workers still perform transfers. Validate canonical tensors against the plan and preserve empty pipeline-local group indices.
• Validate worker counts and copy-completion quorum against TP * PP. Queued execution requires the existing all-rank compute barrier, including PP with async scheduling disabled. CLI float-valued byte budgets retain native truncation behavior.
• Scope: positive TP/PP sizes allowed by the model/upstream, one node, mp (uni for one worker only), Model Runner V2, sole OffloadingConnector, DP/PCP/DCP=1. Reject canonical_layout, ranks without KV tensors, and strided state-only partitions unsupported by native registration. Existing backend and page-geometry constraints still apply.

Before / after architecture

flowchart TB
  subgraph Before["Before: TP 1/2/4, PP 1"]
    BG["Topology guard"] --> BN["Native CPU spec and workers"]
    BN --> BC["Native buffers; shared slots where eligible"]
    BG --> BR["PP greater than 1 rejected"]
  end
  subgraph After["After: validated single-node TP x PP"]
    D["Per-rank logical KV descriptors"] --> P["Common chunk budget"]
    P --> S["Scheduler: common chunk IDs and count"]
    P --> RPC["Collective RPC before native KV initialization"]
    RPC --> R0["Rank 0 native worker"]
    RPC --> RN["Rank N native worker"]
    R0 <--> C0["Rank 0 private pinned CPU buffer"]
    RN <--> CN["Rank N private pinned CPU buffer"]
    R0 --> Q["Copy-completion quorum: TP x PP"]
    RN --> Q
    Q --> F["Existing retirement and compute-release fence"]
    F --> U["Unmap when all holds are clear"]
  end

Loading

Before: restricted path

sequenceDiagram
  participant E as Engine
  participant W as TP workers
  participant M as KV manager
  Note over E,W: TP=1/2/4, PP=1, PP rejected before initialization
  E->>W: Native KV and CPU offloader initialization
  E->>M: Pin transfer pages before dispatch
  E->>W: Dispatch native load/store jobs
  W-->>E: Copy-complete reports
  E->>M: Unpin after configured worker quorum
  opt Queued execution
    E->>W: Existing ordered compute barrier
    W-->>E: All workers finish earlier CUDA work
  end
  E->>M: Drain eligible retired pages
  Note over E,M: Pinning and release fence inherited from #598

Loading

After: rank-private capacity and all-rank release

sequenceDiagram
  participant E as Engine
  participant W as All TP x PP workers
  participant S as Native CPU specs
  participant M as KV manager
  E->>E: Derive rank payloads and common chunk count
  E->>W: Collective RPC installs capacity plan
  W-->>E: Verify every expected rank acknowledged
  E->>W: Native KV initialization
  W->>S: Validate local tensors and create private pinned buffers
  E->>S: Scheduler uses the same chunk count
  E->>M: Pin transfer pages before dispatch
  E->>W: Dispatch native load/store jobs
  W-->>E: Copy-complete counts
  E->>M: Unpin only after TP x PP completions
  opt Queued execution, including synchronous PP
    E->>W: Existing ordered all-rank compute barrier
    W-->>E: All ranks finish earlier CUDA work
  end
  E->>M: Drain fenced retirements with no remaining holds
  Note over E,M: Copy completion and compute safety both precede unmap

Loading

Validation

• Exact commit: CPU aggregate suite 1659 passed, 15 skipped; test classification, GPU-free C++ tests, applicable pre-commit hooks and full-package/test mypy for Python 3.10–3.13 passed. Dedicated private Codex review found no actionable regressions.
• Focused regressions: 324 passed, 1 skipped (the skip requires an installed vLLM runtime).
• Pinned vLLM 0.29/0.30 native spec/factory API harness: 4 cases passed, with dependencies stubbed.
• CPU tensor layout/copy harness using pinned upstream grouping, PP projection and registration code: 80 configurations, 760 stages, including 136 state-only stages passed. Covers TP=1/2/4/8 and PP=1/2/3/4/5/7/8/16/17/32 for a dense hybrid attention/state layout with equal padded pages, align mode and BLNHC. This is source-contract and real CPU tensor coverage.