feat(server): native KV clients backed by server-owned CUDA arenas

[xiaguan]

Summary

Adds a native (non-Python) KV registration path with inverted ownership: the server allocates the fused GPU KV arena and shares it back to the client as a CUDA IPC handle inside the registration response. This is a redesign of #414/#416 that keeps their goal (RDMA-capable KV memory) while deleting the fd side-channel and the ownership machinery those PRs needed.

Why invert the ownership

The end goal is registering KV memory into the NIC (GPUDirect RDMA). That registration only works on memory the registering process owns — ibv_reg_mr/dma-buf on an IPC-imported pointer fails, which is why #414 had the client allocate via VMM and export a POSIX fd. But a POSIX fd cannot travel in a gRPC message, which forced:

• a Unix-socket SCM_RIGHTS side-channel (fd_channel.rs, ~520 lines with its hardening) with an out-of-band (instance_id, device_id) keying protocol, reordering waits, TTLs, and stale-socket recovery;
• fix(server): make native VMM registration ownership explicit #416's ownership contract (Arc<VmmRegistration>, operation gates, two-phase teardown, spawn shields) to keep the imported mapping alive against client crashes and RPC cancellation.

Allocating in the server dissolves both problems at once:

• a legacy CUDA IPC handle is 64 plain bytes — it rides the existing RegisterContextResponse. Registration is one round trip; there is no side-channel, no fd, no ordering invariant.
• ownership is trivially correct: the server owns the arena for the registration's lifetime and frees it on unregister/session-cleanup/HTTP-cleanup. The client only runs compute kernels on its imported mapping (fully supported by CUDA IPC). Owner-side NIC registration in the follow-up RDMA work is the standard path (what Mooncake/NCCL do).
• failure policy is the simple one: the client treats a lost server as fatal and exits; a dead client leaks nothing on the server. No reconnect, no drain gates.

Verified empirically on a real GPU before building this: cross-process cudaMalloc + cuIpcGetMemHandle → import → kernel read/write in both directions, bit-exact.

Changes

• proto: NativeKvTensor{offset,size,block_stride} views + native_alloc_size on RegisterContextRequest; arena_ipc_handle on the response; LoadRequest.wait_for_completion; Flush RPC; version tag +native-arena-v1.
• native_arena.rs (new, isolated): cuMemAlloc + zero + cuIpcGetMemHandle; layer views are base + offset arithmetic; freed on registry drop, on the registry actor thread.
• registry: native contexts live in a separate map beside the Python bookkeeping — the torch/GIL code paths are untouched; the shared actor thread runs allocation and teardown.
• service: one native branch in register_context_batch (both paths converge on the pre-existing register_context_layer_batch_strided); synchronous native load via the pre-existing _inproc engine path; teardown order flipped to engine-unregister-before-registry-drop so raw pointers never dangle over a freed arena.
• --python-registry=false: torch-free startup (cudarc device init), serving native clients only. Default true keeps the vLLM connector path identical.
• engine (pegaflow-core): zero data-path changes — _strided registration and _inproc load already existed on master.

Diff is +937/−79 vs #414+#416's +2712/−751, with near-zero churn on the Python path.

Scope (native v1)

One GPU, one fused arena, multiple layer views per registration (tp_size=1, world_size=1 enforced). Multi-arena (GLM5.2) is out of scope and fails validation early. Same-host only — CUDA IPC is host-local, same constraint as the fd design.

Testing

• native_arena_rpc_e2e on a real GPU: register → a child process imports the handle (CUDA forbids importing in the exporting process) and writes a pattern → save → child wipes → wait_for_completion load → child verifies bit-exact restore → unregister frees the arena.
• Full release suite: server 23 unit + mock-vLLM 8 + cleanup-hang repro + native e2e, all green; cd python && uv run pytest — 202 passed.
• Consumer PR (openinfer): byte-level cpu_roundtrip and Qwen3-4B kv_offload_cpu_hit (HBM-evicted prefix restored from the CPU tier) pass against this server in torch-free mode on an RTX 5070 Ti.

Review hardening (second commit)

An adversarial review pass surfaced and fixed: HTTP-cleanup freeing arenas before the engine forgot its pointers (the one teardown path with the old ordering); a save-pipeline flush barrier on all three teardown paths (saves still queued inside a GPU worker remain unfenced until the engine-level drain-before-unregister lands — the comments state this window explicitly); a post-publish registry confirmation so a cleanup racing the two-step registration can no longer leave the engine pointing at a freed arena; cuCtxSynchronize after the arena zero-fill (host-async memset + no cross-process stream ordering).

Not in this PR

• NIC registration (ibv_reg_mr/dma-buf) of the arena — the allocation is owner-side and RDMA-ready; wiring RDMA is a follow-up.
• Multi-arena registration (GLM5.2 layout).

Replaces #414 and #416.

[GentleCold]

The arena being allocated once at client/server startup and remaining fixed during normal inference looks reasonable; I do not think allocation frequency or fragmentation should block this PR.

The remaining blocker is teardown ordering. unregister_instance() removes the instance from the map, but it does not drain load/save tasks that have already been handed to GpuWorkerPool. flush_saves() only drains the storage write pipeline. The subsequent registry.drop_instance() drops NativeArena and calls cuMemFree, while an in-flight H2D/D2H task can still hold a KVCacheLayout containing pointers into that arena. Session-close cleanup, HTTP cleanup, or UnregisterContext racing a native load/save can therefore access freed device memory. The comment in cleanup_instance already acknowledges this residual window.

Please add an engine-level drain-before-unregister operation: stop accepting new work for the instance, wait for both GPU worker queues and their active transfers to complete, remove the engine instance, and only then drop the registry arena.

One related contract to confirm/document: before a native client sends Save, it must synchronize the CUDA stream that wrote the imported arena. The server stream synchronization cannot order work from the client process. The current e2e test uses synchronous cuMemcpyHtoD_v2, so it does not exercise an asynchronous kernel-write followed by Save.

[xiaguan]

@GentleCold re teardown ordering:

Agreed on allocation: fixed-size arena at registration is fine; that is not a blocker.

On teardown ordering — we are intentionally not adding engine-level drain-before-unregister for this PR.

Client contract: before UnregisterContext (or equivalent cleanup the client initiates), the client must ensure every Save / Load it started has finished. For native clients that means:

• every Save RPC has returned (the server already awaits GPU D2H before replying);
• every Load with wait_for_completion=true has returned;
• do not race unregister with in-flight transfer RPCs.

Under that contract, GPU workers for the instance have no outstanding work when the engine drops layouts and the registry frees the arena, so the residual window does not apply to a correct client.

We are not building a shared-ref shutdown / generation / drain barrier in the engine for this. A previous attempt at full GPU drain-on-unregister was large and messy relative to the real failure mode, which is almost always a client lifecycle bug (or forced session/HTTP cleanup while traffic is still in flight). Session-disconnect and admin cleanup can still race inflight work; that is accepted for now the same way the Python IPC path never fenced GPU workers before unmap — native just makes a buggy client more obviously fatal (cuMemFree vs unmap).

Stream sync before Save: yes — the client must synchronize the CUDA stream that wrote the imported arena before calling Save. Server-side stream sync cannot order work issued in the client process. Will call that out in the native client contract / docs (the e2e uses synchronous cuMemcpyHtoD_v2 and does not cover async kernel → Save).

So: no drain-before-unregister in this PR; client guarantees idle-before-unregister + pre-Save stream sync. Happy to revisit a minimal disconnect-only join later if forced cleanup races become a real operational issue, but not as a merge blocker here.