feat(kv-offload): export KV via VMM dma-buf fd instead of CUDA IPC - #740
feat(kv-offload): export KV via VMM dma-buf fd instead of CUDA IPC#740xiaguan wants to merge 6 commits into
Conversation
Signed-off-by: xiaguan <751080330@qq.com>
Signed-off-by: xiaguan <751080330@qq.com>
Signed-off-by: xiaguan <751080330@qq.com>
Signed-off-by: xiaguan <751080330@qq.com>
Signed-off-by: xiaguan <751080330@qq.com>
Switch the external-PegaFlow KV offload client from CUDA IPC handles to VMM (cuMemCreate) allocations exported as POSIX file descriptors. A CUDA IPC imported pointer cannot be registered into the NIC, so it can never back GPUDirect RDMA; a VMM dma-buf fd can. host-tier save/load is unchanged in behavior — this swaps the sharing primitive so RDMA becomes possible later. - exportable: VmmExportableBuffer allocates the fused KV arena via cuMemCreate (POSIX-fd handle), maps it, and exports an fd. The gpuDirectRDMACapable flag is set only where the device attribute reports support (consumer parts reject it in cuMemCreate). - buffer: KvBuffer holds the VMM allocation and a ManuallyDrop CudaSlice view for the attention kernels; the view never frees the VMM pointer (teardown goes through unmap/free/release). - external: send the export fd to PegaFlow over its Unix-domain fd side-channel (SCM_RIGHTS) before register; the gRPC request carries only per-layer strided metadata plus the aligned allocation size. Requires the matching PegaFlow native-VMM server (pinned rev). Verified end to end on Qwen3-4B: an evicted prefix restores from the CPU tier bit-exact. Co-Authored-By: Claude <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e56f8a71d8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self.flush_saves(); | ||
| if let Err(err) = self | ||
| .host | ||
| .runtime | ||
| .block_on(self.host.client.unregister(&self.instance_id)) |
There was a problem hiding this comment.
When a save or server flush takes longer than the 5s FLUSH_DEADLINE, flush_saves_then reports completion but drops the JoinHandles rather than aborting the save tasks, so those RPCs can continue using the exported CUDA allocation. shutdown() then unregisters the context here and the subsequent drop can release the VMM buffer while the server may still be DMAing from it (the save RPC timeout is 30s), which defeats the new ownership-safety checks. The shutdown path should either wait for/abort on all outstanding saves instead of using the bounded flush, or not unregister/release the mapping until they are definitively complete.
Useful? React with 👍 / 👎.
| // See `with_arenas`: multi-arena VMM export is future work. | ||
| None, | ||
| None, |
There was a problem hiding this comment.
This multi-arena path is still reachable from GLM5.2 build_offload_engines() for every --kv-offload-server launch, but ExternalClient::register now treats a missing export_fd/alloc_size as a hard InvalidArgument because VMM registration has no fallback sharing mechanism. As a result, GLM5.2 KV offload always fails during rank registration despite the CLI/docs enabling it; either make these arenas VMM-exportable and pass the fd metadata, or reject GLM offload before startup until that migration is implemented.
Useful? React with 👍 / 👎.
1 participant
Summary
Switches the external-PegaFlow KV offload client from CUDA IPC handles to VMM allocations exported as POSIX file descriptors.
Why: a CUDA-IPC-imported pointer cannot be registered into the NIC (
ibv_reg_mron acuIpcOpenMemHandlepointer fails, and no dma-buf fd can be derived), so it can never back GPUDirect RDMA. A VMM allocation (cuMemCreate+ POSIX-fd export) can be handed toibv_reg_dmabuf_mr. This PR swaps the sharing primitive so RDMA becomes possible; host-tier save/load behavior is unchanged.A POSIX fd cannot travel inside a gRPC message, so the export fd is sent to PegaFlow over its host-local Unix-domain fd side-channel (
SCM_RIGHTS) before registration; the gRPC request carries only per-layer strided metadata plus the aligned allocation size.Changes
exportable.rs:VmmExportableBufferallocates the fused KV arena viacuMemCreate(POSIX-fd handle), reserves/maps it, and exports an fd.gpuDirectRDMACapableis set only where the device attribute reports support — consumer parts reject it incuMemCreate, data-center parts (H100/H200/B200) accept it.buffer.rs:KvBufferholds the VMM allocation plus aManuallyDrop<CudaSlice>view for the attention kernels; the view never frees the VMM pointer (teardown goes through unmap/free/release, notcuMemFree).external.rs: send the export fd over PegaFlow's fd side-channel (SCM_RIGHTS) beforeregister; build per-layerNativeKvTensormetadata and pass the granularity-alignednative_alloc_size.engine.rs: thread the export fd + alloc size from the buffer through registration. The multi-arena path (GLM5.2) is left onNone— VMM export for it is future work; this PR covers the single fused Qwen3 buffer.Dependency
Requires the matching PegaFlow native-VMM server: novitalabs/pegaflow#414. The dep is pinned to that commit (
6e0baa7); repin to the merged rev before this lands.Testing
kv_offload_cpu_hit::live_gpu_and_cpu_prefix_hits, head-logprob delta 0.0000).--python-registry false) mode.🤖 Generated with Claude Code