Skip to content

feat(p2p): load RAM prefixes and remote RDMA suffixes into GPU KV - #454

Draft
GentleCold wants to merge 24 commits into
novitalabs:masterfrom
GentleCold:feat/direct-gpu-rdma-load
Draft

GentleCold wants to merge 24 commits into
novitalabs:masterfrom
GentleCold:feat/direct-gpu-rdma-load

Conversation

@GentleCold

@GentleCold GentleCold commented Sep 11, 2026 •

Copy link
Copy Markdown
Collaborator

Cross-node KV hits normally pass through requester host staging before being copied into GPU KV slots. This draft adds an opt-in RDMA READ path from remote registered host memory directly into the requester GPU allocation after vLLM assigns destination blocks.

Set pegaflow.direct_gpu_rdma=true in the connector's
kv_connector_extra_config to request direct plans. The initial scope is dense
attention storage group 0; page-first and recurrent/HMA layouts are rejected.
Direct-mode errors do not invoke host-staging fallback. The ordinary
host-staging path remains available when the option is omitted or false.

Changes

  • Propagate QueryRequest.direct_gpu through PyO3, Python stubs and TP shard queries. Carry ordered remote plans in the existing query lease, reusing its 600-second expiry and consume/release lifecycle without adding a new TTL mechanism.
  • Add CUDA pointer/device validation, DMA-BUF memory registration on each configured NIC, lazy GPU allocation registration and instance unregister hooks to the v1 pull transport.
  • Resolve contiguous/split K/V GPU destinations during load, issue ordered remote reads, report direct load status/duration and MR failures, and avoid inserting direct destinations into requester host ReadCache or advertising them as owners.
  • Preserve the master ordering of local RAM prefix first and remote direct suffix second; direct errors remain hard load failures and do not fall back to host staging.
  • Extend descriptor/lease tests and the cross-process GPU integration with an independent host-staging control.

Validation

  • Based on upstream/master merge base 1a20960; feature work is isolated on feat/direct-gpu-rdma-load.
  • Local source-only Python gate: 307 passed / 1 skipped / 14 deselected. Ruff format and lint pass for all changed Python files; git diff --check and the example syntax check pass.
  • Remote CUDA 13/RDMA validation on the feature sources: workspace/all-targets Clippy, focused transfer/core/lease tests, native wheel build, same-host cross-process direct GPU byte validation, and the standard host-staging vLLM correctness E2E (5/5) passed.
  • The direct integration covers independent owner/worker processes, DMA-BUF FD handoff, split K/V, nonzero views, contiguous and fragmented blocks, and byte-for-byte GPU verification. It is not a two-physical-node vLLM run.
  • GB300 paired replay (8,422 requests per side, same trace and request order): direct mode completed 1,387 direct loads with 0 errors and 16.322 seconds aggregate direct-load duration; the host-staging control completed 735 host RDMA fetches with 1 fetch error and 9.607 seconds aggregate successful-fetch duration. These metrics have different operation boundaries and hit sets, so they are directional only. Overall mean TTFT was 187.527 ms (direct) versus 187.460 ms (control), and throughput was 5.322 versus 5.326 requests/s; the replay did not establish an end-to-end gain.
  • The replay is same-host 2P2D/TP1 with NIXL handling the main P-to-D transfer. It does not replace the required two-node, multi-TP/multi-NIC validation or isolate direct-vs-host staging on a fixed hit set.
  • Earlier RTX 4090/5090 attempts rejected CUDA DMA-BUF export with driver error 801 (cuMemGetHandleForAddressRange: operation not supported). Those failures are capability limits, not direct-path results.

Before ready for review / merge

  • Add startup capability gating and negotiate direct support so an older server cannot silently ignore the request flag.
  • Ensure GPU allocation/MR ownership spans queued and in-flight transfers, including concurrent registration and instance teardown.
  • Complete the RC post/CQ error drain/reset contract and injected transport-fault coverage; successful completion tests do not cover every in-flight lifetime path.
  • Add direct-RDMA byte and per-stage observability, then compare a fixed hit set against host staging. Current direct duration is an aggregate batch timer and cannot be treated as raw bandwidth.
  • Complete real two-node vLLM query → GPU allocation → RPC load → GPU consumption validation, with split K/V, multilayer/multi-TP cases, multi-NIC coverage, and independent host-staging byte/latency/memory comparison.

This PR remains a draft while these correctness and validation items are open.

Tracking: AIINF-286.

2026-09-24 optimization update

  • Added low-risk RC completion tracking preallocation (min(total_ops, MAX_SEND_WR)) after the ordered READ batch is constructed. This only removes HashMap growth/rehash overhead; it does not change owner ordering, retries, fallback, completion draining, or host staging.
  • Remote CUDA 13/RDMA validation passed: formatting, transfer/core Clippy with -D warnings, 17 transfer tests, 12 backing::rdma_fetch tests, release server/metaserver build, maturin native release build, and feature import-path verification (vLLM 0.29.0).
  • Replayed the same Maiyi B300 2P2D/TP1 Qwen3-0.6B trace prefix (580 selected / 200 eligible, block size 64, concurrency 1, 8 GiB pools, CUDA NIXL buffer, zero retries) on GPUs 0–3. The new opt7 run completed 200/200 requests in 40.610 s at 4.925 req/s; mean TTFT was 202.565 ms. Direct metrics reported 7 correctness-probe successes, 0 errors, and 0 MR registration failures; all three probe cases matched recomputation.
  • The matching master baseline on the same trace/config was 42.740 s at 4.680 req/s with 213.161 ms mean TTFT. Opt7 is +5.24% throughput and -4.97% mean TTFT versus that cold run. This is below the desired 10% target and remains a single cold replay observation; the direct RDMA operation itself accounts for about 0.484 s over 27 loads, so most end-to-end time remains in scheduling and NIXL P/D transfer.
  • GPU 0–3 and all test processes were clean after replay; production kern-serve processes on GPUs 4–7 were left untouched.

@GentleCold GentleCold changed the title feat(p2p): add direct RDMA reads into GPU KV cache feat(p2p): load RAM prefixes and remote RDMA suffixes into GPU KV Sep 17, 2026
Export full CUDA allocations in the owner and hand off live DMA-BUF descriptors during registration. Keep imported tensors with registered and queued RDMA memory, drain accepted completions, and flush writes before reporting success.

Group standardized dense BHNC physical rows into logical cache blocks. Add real cross-process GPU integration and make the same-process E2E load probe exercise a complete dense block.
@GentleCold
GentleCold force-pushed the feat/direct-gpu-rdma-load branch from ea3b70d to 20b1c97 Compare September 21, 2026 14:11
@GentleCold

Copy link
Copy Markdown
Collaborator Author

Reviewed against the latest upstream/master and simplified the PR in 4bc9d7e.

Changes removed:

  • asynchronous direct-GPU RDMA connection warmup during query; handshake now stays on the load path
  • chain-tail-only RDMA signaling and its extra completion accounting; all posted WRs use the existing completion/error-drain contract
  • simplified GPU visibility-fence result handling without changing when the fence runs

Direct-GPU correctness paths remain: DMA-BUF ownership/lifetime, RAM-prefix plus remote-suffix planning, GPU descriptor shape checks, GPU-specific NIC dispatch, QP reset/cleanup, and direct-load metrics.

Validation:

  • Python unit gate: 309 passed, 1 skipped, 14 deselected
  • ruff, rustfmt, and diff checks pass
  • local Rust build is blocked by macOS missing Linux RDMA headers; GitHub cargo/clippy/CUDA checks are pending.

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