Skip to content

feat: add native physical KV reservation APIs - #487

Open
shipiyouniao wants to merge 1 commit into
ovg-project:mainfrom
shipiyouniao:feature/physical-kv-reservations
Open

shipiyouniao wants to merge 1 commit into
ovg-project:mainfrom
shipiyouniao:feature/physical-kv-reservations

Conversation

@shipiyouniao

@shipiyouniao shipiyouniao commented Sep 14, 2026 •

Copy link
Copy Markdown
Collaborator

Problem

Two KVCached instances on one GPU can both observe enough free memory for their next growth operation, then allocate against the same observation. For example, with 12 GiB free, each can independently decide to grow by 8 GiB. The CUDA driver will reject physical over-allocation, but repeated partial transactions and retries are not a useful admission policy for serving.

This PR separates physical page reservation from virtual-address mapping. The process-shared turn covers a fresh capacity check and acquisition of physical handles. Mapping, CUDA synchronization and IPC remain outside that turn. It is the native foundation for #455, not an independently enabled distributed protocol.

Root Cause and Lock Boundary

The earlier FTensor map path acquires a physical page and maps it in the same operation:

auto page = make_unique_page(dev_, page_id, page_size_);
if (!page->map(vaddr)) {
  throw std::runtime_error("physical page map returned false");
}

A free-memory observation made before that operation is not a reservation. Putting a shared lock around the entire map/broadcast path would serialize more than necessary: after physical handles have been acquired, the GPU's free memory already reflects their cost. Another instance can safely make a fresh admission decision without waiting for those handles to be mapped.

The following is an illustrative race, not a claim that every rejected allocation follows this exact timing.

Before: Observation Without Reservation

sequenceDiagram
    participant A as Instance A
    participant G as GPU physical memory
    participant B as Instance B
    A->>G: Read free memory
    G-->>A: 12 GiB
    B->>G: Read free memory
    G-->>B: 12 GiB
    A->>G: Allocate 8 GiB and map
    B->>G: Attempt another 8 GiB
    G-->>B: Physical allocation cannot satisfy both requests
Loading

After: Reserve Under a Short Turn

sequenceDiagram
    participant A as Instance A
    participant T as Per-GPU shared admission turn
    participant G as GPU physical memory
    participant B as Instance B
    A->>T: Wait for turn and lock physical GPU
    A->>G: Fresh headroom check, then acquire handles
    A->>T: Release shared turn
    par Mapping outside the shared turn
        A->>G: Map already-reserved handles
    and Next admission
        B->>T: Acquire next turn
        B->>G: Re-read capacity including A's reservation
        G-->>B: Reserve if sufficient, otherwise reject before mapping
        B->>T: Release shared turn
    end
Loading

Implementation

The capacity check selects the target device and compares the complete new physical cost against the remaining headroom-adjusted capacity:

const size_t usable_bytes = free_bytes > headroom ? free_bytes - headroom : 0;
if (required_bytes > usable_bytes) {
  // Record shortfall and reject physical growth.
}

This excerpt abbreviates the error/statistics branch; the linked source contains the full calculation. required_bytes sums the actual FTensor targets, not just the nominal size of one logical page. Existing mapped targets being adopted are not charged again.

Area Change
FTensor reservation and mapping Acquire a physical handle separately from attaching it. The caller retains ownership if attachment fails.
Native reservation scope Hold the shared turn only for admission and handle acquisition. Deduplicate offsets before reserving.
Transaction APIs Prepare, commit, abort and inspect reservations. Checked release retains cleanup ownership if handle release fails; an incomplete reservation cannot commit.
FIFO admission Identify the physical GPU by PCI Bus ID and use process-shared locking with abandoned-ticket cleanup.
Python bindings Release the GIL around blocking native operations and return phase statistics.

The allocator rejects unmap of offsets with pending reservations while allowing disjoint offsets to be released. Mapping failures retain their error classification; this layer alone does not decide a multi-worker outcome.

Dependencies and Intermediate State

Review only this PR's change.

Rebased onto main 65bbe6606f1e4f35a6400da5c8ba85548792b28e after #410 merged. #464 is also merged. This branch now contains one reservation commit on main, with the same PR-specific changes. The reservation target list now preserves #357 single-buffer MLA mapping and #529/#540 per-pool physical page sizes; current lifecycle, shutdown and ordered-release behavior is retained.

Merge order: #487 -> #455, with #488 / #489 optional. This PR adds native APIs without switching legacy one-phase callers. #455 connects distributed IPC callers to the new protocol; direct local legacy callers remain unchanged. #488 and #489 are independent follow-ups.

The updated integration and tests remove the stale-base blocker, but single-GPU tests do not replace current-head multi-GPU serving acceptance.

Deployment and Limits

  • All participating processes must see the same underlying lock files, with compatible filesystem locking and permissions. KVCACHED_PHYSICAL_GROWTH_LOCK_DIR pointing to the same string inside isolated containers is not sufficient; use a genuinely shared mount.
  • FIFO orders admission attempts. It does not divide memory equally, reclaim another instance's live pages, or guarantee bounded service time when a holder stalls.
  • The guard coordinates participating KV allocations only. Unrelated Torch/CUDA allocations do not take it and can still consume headroom.
  • Rebuild the native extension with the matching Python code. Local allocator mutexes still protect native state; narrowing the cross-process turn is not a claim that every operation is lock-free.

A failed commit may attach only part of a logical offset. Attached pages remain owned by their FTensor and are adoptable by a matching retry; abort releases only unattached handles. Immediate cancellation cannot necessarily reclaim a partially populated offset without completing it first. Under persistent capacity pressure, retained physical usage can therefore last until worker teardown. Restoration or physical-handle release failures raise a terminal state-consistency error; handles retained by an FTensor may require teardown rather than another transaction abort. The dedicated review identified this limitation. It is documented rather than addressed with an unsafe online unmap path.

Validation

Exact head: 0e1ec0531dfb1a114e1415ea8f8549cabdb1ac17.

Check Result
Pre-commit, classification, C++ CPU tests Passed
CPU tests, Python 3.10–3.13 1,234 passed, 12 skipped per version
MyPy, Python 3.10–3.13 Passed
Rebuilt stable-ABI extension on T4 Passed
Native VMM/failure/retry, lifecycle and geometry tests on T4 121 passed; 3 manager API cases passed
vLLM 0.29 / 0.30 integration contracts 118 passed per version
Deferred release/reuse/clear, both layouts; native import isolation Passed; 23 import-isolation tests passed
Dedicated Codex review Completed; existing partial-map cancellation and terminal handle-release limits remain documented above

The 12 CPU skips require an installed SGLang radix-tree implementation; those cases are not claimed as validated here.

Added native coverage combines 2/4 MiB per-pool pages, single/two-buffer storage and both layouts, including reservation/abort/commit/adoption and CUDA readback. Fault tests include duplicate offsets, competing reservations, abort/release failure with retained cleanup ownership, partial attachment and adoption, and typed failure propagation. The test build injects real driver-operation failures; these are not just mocked Python callbacks.

Historical serving investigations in #467 motivate this boundary but are not an exact-head benchmark of this split. No new TP4/MPS serving-chaos or throughput claim is made here.

@shipiyouniao
shipiyouniao force-pushed the feature/physical-kv-reservations branch from 53925d6 to 0e1ec05 Compare October 8, 2026 20:25
@shipiyouniao
shipiyouniao marked this pull request as ready for review October 8, 2026 20:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant