feat(kv-offload): import the server-allocated KV arena over CUDA IPC

Cargo.toml

@@ -171,6 +171,7 @@ tempfile = "3"
 tmq = "0.5.0"
 tokio = { version = "1", features = ["full"] }
 tokio-stream = "0.1"
+tonic = "0.14"
 tokio-util = "0.7"
 toml = "1.1"
 tower = "0.5"

docs/index.md

@@ -119,7 +119,8 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l
 | `subsystems/router/kv-aware-routing.md` | Dynamo KV-aware routing on 8×Qwen3-4B (RTX 5090): cache-affinity routing keeps a multi-turn conversation on its home worker, so follow-up-turn TTFT stays flat ~45ms vs round-robin 160–170ms / random 165–180ms (all-turns p50 3.3–3.8× lower). Router prefix overlap 0.72 under KV, 0 under stateless policies; `kv_hit_rate>0` is the gate that the worker↔router block-hash bridge is actually matching. Includes the per-response `prompt_tokens_details.cached_tokens` signal. |
 | `subsystems/runtime/runtime.md` | Runtime complexity is controlled by a shared `openinfer-core` that owns the generation contract and orchestration; per-model crates implement `ModelForward` so prefill/decode and hybrid attention stay hidden from the caller. State (`&mut`) is separated from weights (`&self`) for future bs > 1. |
 | `subsystems/runtime/kv-cache-design.md` | Dynamo 式 logical/physical 分层 KV cache：BlockManager 管 block 生命周期和 admission，PhysicalBackend trait 管 GPU 内存和布局（FullAttention / MLA）。支持 TP / DP。基于 vLLM/Dynamo/pegaflow 调研。 |
-| `subsystems/runtime/pegaflow-offload-integration.md` | 把 `pegaflow-core` 当进程内 Rust 库做 KV 卸载物理后端（HBM→DRAM/SSD/RDMA），补 kvbm 没写的卸载层。**Qwen3-4B full-attn 首发，端到端已在真实 GPU 跑通并验证**（async SAVE+LOAD 接进 executor/scheduler，纯 CPU-hit 与 GPU+CPU 组合 hit 恢复后 logits 与冷算一致）。pegaflow 经 git rev pin（#331+#333）。默认关，server CLI 已接（#316：`--kv-offload`/`--no-prefix-cache`，plain+LoRA）。linear 排除，sparse 暂缓。 |
+| `subsystems/runtime/pegaflow-offload-integration.md` | Historical design and Qwen3 evidence for the retired embedded PegaFlow backend; current external-server ownership is tracked in `external-pegaflow-server.md`. |
+| `subsystems/runtime/external-pegaflow-server.md` | External-only PegaFlow boundary: the server allocates and owns the fused GPU KV arena (returned as a CUDA IPC handle at registration) plus DRAM/SSD/RDMA; OpenInfer imports the arena and issues save/query/load/flush RPCs. Qwen3 only; GLM5.2 multi-arena is out of the v1 contract. |
 
 ## subsystems / scheduler
 

docs/subsystems/runtime/external-pegaflow-server.md

@@ -0,0 +1,35 @@
+# External PegaFlow server for KV offload
+
+> **TL;DR:** OpenInfer is an RPC client selected by `--kv-offload-server`; the external PegaFlow process **allocates and owns the GPU KV arena** and the host/SSD/RDMA tiers. OpenInfer registers its KV layout, imports the arena over CUDA IPC, and runs attention kernels on the imported mapping. Implemented for Qwen3; GLM5.2's multi-arena layout is not supported by the v1 contract.
+>
+> **Last touched:** 2026-07
+
+## Ownership boundary
+
+- External mode is the only PegaFlow integration. OpenInfer no longer embeds a `PegaEngine` or constructs a second pinned-memory pool.
+- **PegaFlow owns the fused GPU KV arena.** Registration sends the layout (per-layer offsets, sizes, block strides, total size); the server allocates the arena on the requested device, zeroes it, registers the raw layer pointers into its engine, and returns a 64-byte CUDA IPC handle in the registration response. OpenInfer imports the handle and builds its `KvBuffer` as a non-owning view.
+- Why this direction: GPUDirect RDMA registration (`ibv_reg_mr`/dma-buf) only works on memory the registering process owns — an IPC-imported pointer can never back it. Server-side allocation makes the arena NIC-registerable in the server without any fd side-channel; the IPC handle is plain bytes and rides the existing gRPC response.
+- `--kv-offload-server` enables the client. `--kv-offload-namespace` identifies the checkpoint/deployment content domain; vLLM compatibility uses the connector namespace.
+- Server capacity, SSD, RDMA, routing, and topology configuration stay outside OpenInfer.
+
+## Wire contract
+
+- One fused arena per `(instance, device)`. `RegisterContextRequest.native_kv_tensors` carries per-layer `{offset_bytes, size_bytes, block_stride_bytes}` views into it, and `native_alloc_size` the total size; `RegisterContextResponse.arena_ipc_handle` returns the CUDA IPC handle.
+- Native registration has an exact capability version (`+native-arena-v1`). Old and new clients fail before any memory changes hands.
+- Native `Load` sets `wait_for_completion` and returns only after the GPU transfer completes. The Python connector retains its shared-memory completion path.
+- `Flush` waits for previously submitted saves to become cache-visible and for queued MetaServer registrations to be attempted.
+
+## Lifetime and failure policy
+
+- The arena lives exactly as long as the registration: unregister (or the server's session/HTTP cleanup, if the client dies) frees it. Teardown order on the client is workers → close IPC mapping → unregister.
+- The client treats the server as load-bearing for its own GPU memory: a broken session stream, a failed save/load transport, or an unregister failure exits the process. There is no reconnect.
+
+## Model layouts
+
+- Qwen3 registers one page-first fused arena with strided per-layer views.
+- GLM5.2's rank-local MLA and index-K arenas (78 + 21 per EP8 rank) need multiple allocations per instance, which the v1 contract does not cover; `OffloadEngine::with_arenas_on` fails before touching the server.
+
+## Validation
+
+- PegaFlow `native_arena_rpc_e2e` (real GPU): register → child process imports the handle and writes a pattern → save → wipe → `wait_for_completion` load → bit-exact restore → unregister frees the arena.
+- OpenInfer `cpu_roundtrip` and `kv_offload_cpu_hit` run against a live server via `OPENINFER_PEGAFLOW_SERVER` (same host; CUDA IPC is host-local).

docs/subsystems/runtime/pegaflow-offload-integration.md

@@ -1,8 +1,8 @@
 # pegaflow KV 卸载接入 Spec
 
-> **TL;DR**: 把 `pegaflow-core` 当**进程内 Rust 库**做 KV 卸载的物理后端（HBM→DRAM/SSD/RDMA），补上 kvbm 留着没写的卸载层。connector 大脑（决定 load/save 哪些 block）用 kvbm logical/physical 分层思想自建，pegaflow 退为语义无关的 raw block transfer 后端。**路线已调整为 Qwen3-4B full-attn 首发**（原计划 Kimi 首发）：page-first 单 buffer 经 pegaflow `block_stride_bytes`（PR #331）适配。**端到端已在真实 GPU 上跑通并验证**：async SAVE + async LOAD 接进 `Qwen3Executor` + scheduler，`tests/kv_offload_cpu_hit.rs` 覆盖纯 CPU-hit 与 GPU+CPU 组合 hit，恢复后 logits 与冷算一致；连接层 `OffloadEngine` + `tests/cpu_roundtrip.rs` 字节级一致。默认关（builder flag opt-in）；**server CLI 已接**（#316：`--kv-offload` / `--kv-offload-host-gib` / `--no-prefix-cache`，plain 与 `--enable-lora` 两条启动路径都透传）。纯-L2 基准实测 Qwen3-4B mean TTFT 195→40ms（−79%，evict-before-probe → `gpu_hit=0`，全前缀从 host tier 恢复）。**Qwen3.5 linear/SSM state 明确排除**；**DeepSeek sparse 暂缓**。
+> **TL;DR**: 本文保留最初把 `pegaflow-core` 嵌入 OpenInfer 的设计与 Qwen3 验证记录；当前实现已改为 external-only：OpenInfer 通过 `--kv-offload-server` 注册 KV 布局，PegaFlow server 分配并持有 GPU KV arena（注册响应返回 CUDA IPC handle，OpenInfer import 后当 KV buffer 用），DRAM/hugepage/SSD/RDMA 也全归 server。当前边界和跨进程验证见 [external-pegaflow-server.md](external-pegaflow-server.md)。
 >
-> Last touched: 2026-06
+> Last touched: 2026-07
 
 ## 0. 实现状态（2026-06）
 

openinfer-glm52/src/lib.rs

@@ -51,7 +51,6 @@ pub use config::probe_config_json;
 use openinfer_core::engine::EngineHandle;
 use openinfer_core::engine::KvCapacity;
 use openinfer_core::engine::LoadSnapshot;
-use openinfer_kv_offload::HostConfig;
 use openinfer_kv_offload::KvArena;
 use openinfer_kv_offload::OffloadEngine;
 use openinfer_kv_offload::OffloadHost;
@@ -343,44 +342,25 @@ mod topology_tests {
     }
 }
 
-/// Host-tier KV offload knobs. One `PegaEngine` (one pinned pool) backs all
-/// 8 DP ranks under a single namespace: the MLA latent has no TP sharding
+/// External PegaFlow KV offload for all 8 DP ranks under one namespace. The
+/// MLA latent has no TP sharding
 /// and the non-expert weights are replicated, so any rank's KV for a token
 /// prefix is as good as any other's — the same tolerance as reusing a
 /// rank's own prefix cache (FP reduction order may differ across the batch
 /// shapes that computed it, never the semantics). Any rank restores what
 /// any rank saved.
 #[derive(Clone, Debug)]
 pub struct Glm52KvOffloadOptions {
-    /// Host pinned-memory pool size in bytes, shared by all ranks.
-    pub pinned_pool_bytes: usize,
-    /// Back the pool with hugepages (the box must hold a reservation —
-    /// check `HugePages_Total`).
-    pub use_hugepages: bool,
-    /// `Some` joins the cross-instance P2P mesh: saved block hashes register
-    /// with the MetaServer and missing prefixes are pulled from peer
-    /// instances over RDMA — the P/D disaggregation data plane.
-    pub p2p: Option<Glm52P2pOptions>,
+    /// Out-of-process PegaFlow gRPC endpoint.
+    pub server_addr: String,
+    /// Stable deployment/checkpoint identity for native OpenInfer peers.
+    pub namespace: Option<String>,
     /// `Some` when the P/D prefill peer is vLLM (pegaflow connector): offload
     /// query keys switch from kvbm lineage hashes to vLLM's prefix-cache hash
     /// scheme so this decode node can find the blocks vLLM registered.
-    /// Requires `p2p` (the peer's KV lives in its pegaflow-server's pool,
-    /// even on the same host).
     pub vllm_compat: Option<Glm52VllmCompatOptions>,
 }
 
-/// Cross-instance P2P KV sharing (see `openinfer_kv_offload::P2pConfig`).
-#[derive(Clone, Debug)]
-pub struct Glm52P2pOptions {
-    /// MetaServer gRPC address, e.g. `http://10.0.0.100:50056`.
-    pub metaserver_addr: String,
-    /// This engine's routable `IP:port` (doubles as the embedded transfer
-    /// service's bind address). Must be reachable by every peer.
-    pub advertise_addr: String,
-    /// RDMA NIC device names to register the pinned pool on.
-    pub rdma_nics: Vec<String>,
-}
-
 /// Decode-node settings for a P/D deployment whose prefill node is vLLM with
 /// the pegaflow connector (see `openinfer_kv_offload::VllmBlockHasher` and
 /// `docs/models/glm52/pd-vllm-prefill.md`).
@@ -497,16 +477,14 @@ pub fn launch(model_path: &Path, options: Glm52LaunchOptions) -> Result<EngineHa
         "GLM5.2 --kv-offload requires the EP8 topology (tp8 replicates KV on all ranks; \
          a host-tier restore would land on one)"
     );
-    // The vLLM prefill peer's KV lives in its pegaflow-server's pool (a
-    // separate process even on the same host); without the P2P mesh the
-    // compat keys would query an empty local tier and every request would
-    // wait out the full miss window.
+    // The native-arena server allocates exactly one fused arena per instance;
+    // GLM5.2's rank-local MLA/index-K arenas need multi-arena registration,
+    // which the v1 contract does not cover. Refuse here, before the multi-GPU
+    // weight load — `with_arenas_on` would only fail after it.
     ensure!(
-        kv_offload
-            .as_ref()
-            .is_none_or(|kv| kv.vllm_compat.is_none() || kv.p2p.is_some()),
-        "GLM5.2 --kv-pd-vllm-seed requires the KV P2P mesh (--kv-p2p-metaserver-addr, \
-         --kv-p2p-advertise-addr, --kv-p2p-nics)"
+        kv_offload.is_none(),
+        "GLM5.2 KV offload is not supported by the native-arena PegaFlow contract \
+         (one server-allocated arena per instance); drop --kv-offload-server"
     );
     // The miss window must sit inside the in-flight-fetch ceiling, or the
     // registration phase could never hand over to the fetch phase.
@@ -1135,28 +1113,19 @@ fn build_offload_engines(
             .all(|arena| arena.bytes_per_block == mla_page_size * mla_bytes_per_token)),
         "GLM5.2 KV offload ranks disagree on MLA cache layout"
     );
-    let host = OffloadHost::new(HostConfig {
-        pinned_pool_bytes: opts.pinned_pool_bytes,
-        use_hugepages: opts.use_hugepages,
-        runtime_threads: 2,
-        p2p: opts
-            .p2p
-            .as_ref()
-            .map(|p2p| openinfer_kv_offload::P2pConfig {
-                metaserver_addr: p2p.metaserver_addr.clone(),
-                advertise_addr: p2p.advertise_addr.clone(),
-                rdma_nics: p2p.rdma_nics.clone(),
-            }),
-    })
-    .map_err(|err| anyhow::anyhow!("GLM5.2 KV offload host: {err}"))?;
+    let host = OffloadHost::connect(&opts.server_addr, 2)
+        .map_err(|err| anyhow::anyhow!("GLM5.2 KV offload host: {err}"))?;
     // vLLM-compat mode joins the *P side's* content domain: the pegaflow
     // connector derives an 8-hex namespace from vLLM config (and logs it at
     // startup); reproducing that derivation would mean chasing Python repr
     // of vLLM internals, so the operator passes it through explicitly.
     let namespace = match &opts.vllm_compat {
         Some(compat) => compat.namespace.clone(),
         None => format!(
-            "openinfer-glm52-l{GLM52_LAYERS}-p{}-mla{}-idxk{}",
+            "openinfer-glm52-{}-l{GLM52_LAYERS}-p{}-mla{}-idxk{}",
+            opts.namespace
+                .as_deref()
+                .context("GLM5.2 native KV offload requires a checkpoint namespace")?,
             mla_page_size,
             mla_bytes_per_token,
             config::GLM52_INDEX_HEAD_DIM + 4,
@@ -1194,10 +1163,9 @@ fn build_offload_engines(
             .filter(|&layer| config::glm52_layer_has_full_indexer(layer))
             .count();
     log::info!(
-        "GLM5.2 KV offload up: {} pinned host pool (hugepages: {}), namespace {namespace}, \
+        "GLM5.2 KV offload up: server={}, namespace {namespace}, \
          {} rank instances x {arenas_per_rank} arenas",
-        ByteSize(opts.pinned_pool_bytes as u64),
-        opts.use_hugepages,
+        opts.server_addr,
         engines.len(),
     );
     Ok(engines)

openinfer-glm52/src/scheduler/mod.rs

@@ -490,12 +490,12 @@ pub(crate) fn run_dp8_coordinator(
 
     // Drain in-flight release saves and drop the offload engines BEFORE the
     // workers drop the models: the registered arenas' device memory must
-    // outlive every D2H copy (the `with_arenas_on` contract), and pegaflow's
-    // save worker cannot cancel a copy already handed to it. `flush_saves`
-    // is deadline-bounded, so a stuck host tier cannot hang teardown.
-    if let Some(offload) = offload {
-        for rank in &offload {
-            rank.engine.flush_saves();
+    // outlive every D2H copy (the `with_arenas` contract), and pegaflow's
+    // save worker cannot cancel a copy already handed to it. Shutdown waits
+    // for the server's unregister acknowledgement before the arenas are freed.
+    if let Some(mut offload) = offload {
+        for rank in &mut offload {
+            rank.engine.shutdown();
         }
         drop(offload);
     }

openinfer-glm52/src/scheduler/offload.rs

@@ -334,7 +334,7 @@ pub(super) fn admit_vllm_pd(
                     // feeding the breaker forever.
                     state.consecutive_miss_windows = 0;
                     if let Some(reservation) = pool.reserve_loaded_blocks(hit.num_blocks) {
-                        match offload.engine.load(lease, reservation.page_ids()) {
+                        match offload.engine.load(&lease, reservation.page_ids()) {
                             Ok(handle) => {
                                 // After the H2D lands, rewrite the pages'
                                 // RoPE dims from the peer's interleaved
@@ -423,7 +423,7 @@ pub(super) fn admit_vllm_pd(
                             // so tail_len tokens open exactly one fresh page).
                             let pages = kv.step_page_indices(tail_len);
                             let tail_page = *pages.last().expect("tail step has a page");
-                            match offload.engine.load(lease, vec![tail_page]) {
+                            match offload.engine.load(&lease, vec![tail_page]) {
                                 Ok(handle) => {
                                     let landed = handle
                                         .wait()
@@ -587,7 +587,7 @@ pub(super) fn restore_host_prefix(
     };
     let page_ids = reservation.page_ids();
     let restored = reservation.len();
-    match engine.load(lease, page_ids) {
+    match engine.load(&lease, page_ids) {
         Ok(handle) => match handle.wait() {
             Ok(()) => {
                 pool.commit_loaded_blocks(&mut probe, reservation);