diff --git a/Cargo.lock b/Cargo.lock index 1eff1890a..a07c0c7de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -151,7 +151,7 @@ dependencies = [ "futures-lite", "parking", "polling", - "rustix 1.1.4", + "rustix", "slab", "windows-sys 0.61.2", ] @@ -235,15 +235,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "atomic-polyfill" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cf2bce30dfe09ef0bfaef228b9d414faaf7e563035494d7fe092dba54b300f4" -dependencies = [ - "critical-section", -] - [[package]] name = "atomic-waker" version = "1.1.2" @@ -363,29 +354,6 @@ dependencies = [ "virtue", ] -[[package]] -name = "bindgen" -version = "0.66.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2b84e06fc203107bfbad243f4aba2af864eb7db3b1cf46ea0a023b0b433d2a7" -dependencies = [ - "bitflags 2.11.1", - "cexpr", - "clang-sys", - "lazy_static", - "lazycell", - "log", - "peeking_take_while", - "prettyplease", - "proc-macro2", - "quote", - "regex", - "rustc-hash 1.1.0", - "shlex", - "syn 2.0.117", - "which", -] - [[package]] name = "bit-set" version = "0.5.3" @@ -428,16 +396,6 @@ version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" -[[package]] -name = "bitmask-enum" -version = "2.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6cbbb8f56245b5a479b30a62cdc86d26e2f35c2b9f594bc4671654b03851380" -dependencies = [ - "quote", - "syn 2.0.117", -] - [[package]] name = "blake3" version = "1.8.5" @@ -496,12 +454,6 @@ version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" -[[package]] -name = "bytecount" -version = "0.6.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "175812e0be2bccb6abe50bb8d566126198344f707e304f45c648fd8f2cc0365e" - [[package]] name = "bytemuck" version = "1.25.0" @@ -576,15 +528,6 @@ dependencies = [ "shlex", ] -[[package]] -name = "cexpr" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" -dependencies = [ - "nom", -] - [[package]] name = "cfg-if" version = "1.0.4" @@ -641,17 +584,6 @@ dependencies = [ "half", ] -[[package]] -name = "clang-sys" -version = "1.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" -dependencies = [ - "glob", - "libc", - "libloading 0.8.9", -] - [[package]] name = "clap" version = "4.6.1" @@ -701,15 +633,6 @@ dependencies = [ "cc", ] -[[package]] -name = "cobs" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fa961b519f0b462e3a3b4a34b64d119eeaca1d59af726fe450bbba07a9fc0a1" -dependencies = [ - "thiserror 2.0.18", -] - [[package]] name = "color_quant" version = "1.1.0" @@ -914,12 +837,6 @@ dependencies = [ "itertools 0.13.0", ] -[[package]] -name = "critical-section" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b" - [[package]] name = "crossbeam-channel" version = "0.5.15" @@ -973,7 +890,7 @@ dependencies = [ "crossterm_winapi", "document-features", "parking_lot", - "rustix 1.1.4", + "rustix", "winapi", ] @@ -1244,12 +1161,6 @@ dependencies = [ "litrs", ] -[[package]] -name = "downcast" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1435fa1053d8b2fbbe9be7e97eca7f33d37b28409959813daefc1446a14247f1" - [[package]] name = "dtoa" version = "1.0.11" @@ -1313,18 +1224,6 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" -[[package]] -name = "embedded-io" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef1a6892d9eef45c8fa6b9e0086428a2cca8491aca8f787c534a3d6d0bcb3ced" - -[[package]] -name = "embedded-io" -version = "0.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "edd0f118536f44f5ccd48bcb8b111bdc3de888b58c74639dfb034a357d0f206d" - [[package]] name = "encode_unicode" version = "1.0.0" @@ -1620,15 +1519,6 @@ dependencies = [ "percent-encoding", ] -[[package]] -name = "fragile" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8878864ba14bb86e818a412bfd6f18f9eabd4ec0f008a28e8f7eb61db532fcf9" -dependencies = [ - "futures-core", -] - [[package]] name = "futures" version = "0.3.32" @@ -1841,15 +1731,6 @@ dependencies = [ "zerocopy", ] -[[package]] -name = "hash32" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0c35f58762feb77d74ebe43bdbc3210f09be9fe6742234d573bacc26ed92b67" -dependencies = [ - "byteorder", -] - [[package]] name = "hashbrown" version = "0.12.3" @@ -1899,20 +1780,6 @@ dependencies = [ "hashbrown 0.16.1", ] -[[package]] -name = "heapless" -version = "0.7.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdc6457c0eb62c71aac4bc17216026d8410337c4126773b9c5daba343f17964f" -dependencies = [ - "atomic-polyfill", - "hash32", - "rustc_version", - "serde", - "spin", - "stable_deref_trait", -] - [[package]] name = "heck" version = "0.5.0" @@ -1952,15 +1819,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "home" -version = "0.5.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" -dependencies = [ - "windows-sys 0.61.2", -] - [[package]] name = "hound" version = "3.5.1" @@ -2498,12 +2356,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lazycell" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" - [[package]] name = "leb128fmt" version = "0.1.0" @@ -2551,12 +2403,6 @@ dependencies = [ "libc", ] -[[package]] -name = "linux-raw-sys" -version = "0.4.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" - [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -2787,15 +2633,6 @@ dependencies = [ "autocfg", ] -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - [[package]] name = "mime" version = "0.3.17" @@ -2851,32 +2688,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "mockall" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39a6bfcc6c8c7eed5ee98b9c3e33adc726054389233e201c95dab2d41a3839d2" -dependencies = [ - "cfg-if", - "downcast", - "fragile", - "mockall_derive", - "predicates", - "predicates-tree", -] - -[[package]] -name = "mockall_derive" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25ca3004c2efe9011bd4e461bd8256445052b9615405b4f7ea43fc8ca5c20898" -dependencies = [ - "cfg-if", - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "monostate" version = "0.1.18" @@ -2972,7 +2783,7 @@ dependencies = [ "cc", "cfg-if", "libc", - "memoffset 0.6.5", + "memoffset", ] [[package]] @@ -3084,12 +2895,6 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" -[[package]] -name = "oneshot" -version = "0.1.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "269bca4c2591a28585d6bf10d9ed0332b7d76900a1b02bec41bdc3a2cdcda107" - [[package]] name = "onig" version = "6.5.3" @@ -3287,10 +3092,13 @@ dependencies = [ "cudarc", "half", "log", + "openinfer-kernels", "openinfer-kv-cache", "pegaflow-core", + "pegaflow-proto", "tokio", - "tokio-stream", + "tonic", + "uuid", "xxhash-rust", ] @@ -3563,16 +3371,6 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" -[[package]] -name = "os_socketaddr" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa149f3ec09cadeb41968390fcd71e6ea038a9a373d1d4eb070f4d9484d4b244" -dependencies = [ - "libc", - "winapi", -] - [[package]] name = "oss-harmony" version = "0.0.11" @@ -3601,17 +3399,6 @@ dependencies = [ "winapi", ] -[[package]] -name = "papergrid" -version = "0.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b915f831b85d984193fdc3d3611505871dc139b2534530fa01c1a6a6707b6723" -dependencies = [ - "bytecount", - "fnv", - "unicode-width", -] - [[package]] name = "parking" version = "2.2.1" @@ -3669,16 +3456,10 @@ dependencies = [ "pkg-config", ] -[[package]] -name = "peeking_take_while" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" - [[package]] name = "pegaflow-common" -version = "0.23.3" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=1473c5355d879b4fea23101760cb2a0074642ada#1473c5355d879b4fea23101760cb2a0074642ada" +version = "0.23.5" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=6043c290000a1cc5a384cff59bda6ecaaf868027#6043c290000a1cc5a384cff59bda6ecaaf868027" dependencies = [ "colored", "libc", @@ -3688,8 +3469,8 @@ dependencies = [ [[package]] name = "pegaflow-core" -version = "0.23.3" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=1473c5355d879b4fea23101760cb2a0074642ada#1473c5355d879b4fea23101760cb2a0074642ada" +version = "0.23.5" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=6043c290000a1cc5a384cff59bda6ecaaf868027#6043c290000a1cc5a384cff59bda6ecaaf868027" dependencies = [ "ahash", "bytesize", @@ -3707,7 +3488,6 @@ dependencies = [ "parking_lot", "pegaflow-common", "pegaflow-proto", - "pegaflow-transfer", "rand 0.10.1", "shared_memory", "smallvec", @@ -3718,8 +3498,8 @@ dependencies = [ [[package]] name = "pegaflow-proto" -version = "0.23.3" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=1473c5355d879b4fea23101760cb2a0074642ada#1473c5355d879b4fea23101760cb2a0074642ada" +version = "0.23.5" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=6043c290000a1cc5a384cff59bda6ecaaf868027#6043c290000a1cc5a384cff59bda6ecaaf868027" dependencies = [ "prost", "tonic", @@ -3727,35 +3507,6 @@ dependencies = [ "tonic-prost-build", ] -[[package]] -name = "pegaflow-transfer" -version = "0.23.3" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=1473c5355d879b4fea23101760cb2a0074642ada#1473c5355d879b4fea23101760cb2a0074642ada" -dependencies = [ - "anyhow", - "bincode", - "bytes", - "clap", - "crossbeam-channel", - "cudarc", - "dashmap", - "libc", - "log", - "mea", - "mockall", - "oneshot", - "parking_lot", - "pegaflow-common", - "postcard", - "rdma-mummy-sys", - "serde", - "sideway", - "smallvec", - "syscalls", - "thiserror 2.0.18", - "tokio", -] - [[package]] name = "percent-encoding" version = "2.3.2" @@ -3862,7 +3613,7 @@ dependencies = [ "concurrent-queue", "hermit-abi", "pin-project-lite", - "rustix 1.1.4", + "rustix", "windows-sys 0.61.2", ] @@ -3887,19 +3638,6 @@ dependencies = [ "portable-atomic", ] -[[package]] -name = "postcard" -version = "1.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6764c3b5dd454e283a30e6dfe78e9b31096d9e32036b5d1eaac7a6119ccb9a24" -dependencies = [ - "cobs", - "embedded-io 0.4.0", - "embedded-io 0.6.1", - "heapless", - "serde", -] - [[package]] name = "potential_utf" version = "0.1.5" @@ -3924,32 +3662,6 @@ dependencies = [ "zerocopy", ] -[[package]] -name = "predicates" -version = "3.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ada8f2932f28a27ee7b70dd6c1c39ea0675c55a36879ab92f3a715eaa1e63cfe" -dependencies = [ - "anstyle", - "predicates-core", -] - -[[package]] -name = "predicates-core" -version = "1.0.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cad38746f3166b4031b1a0d39ad9f954dd291e7854fcc0eed52ee41a0b50d144" - -[[package]] -name = "predicates-tree" -version = "1.0.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0de1b847b39c8131db0467e9df1ff60e6d0562ab8e9a16e568ad0fdb372e2f2" -dependencies = [ - "predicates-core", - "termtree", -] - [[package]] name = "prettyplease" version = "0.2.37" @@ -4354,19 +4066,6 @@ dependencies = [ "crossbeam-utils", ] -[[package]] -name = "rdma-mummy-sys" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56f8ead4d7835190d761b9365d9e53906a226d7699007b35bdca49d7edb89a15" -dependencies = [ - "bindgen", - "cmake", - "libc", - "memoffset 0.9.1", - "paste", -] - [[package]] name = "realfft" version = "3.5.0" @@ -4634,19 +4333,6 @@ dependencies = [ "transpose", ] -[[package]] -name = "rustix" -version = "0.38.44" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" -dependencies = [ - "bitflags 2.11.1", - "errno", - "libc", - "linux-raw-sys 0.4.15", - "windows-sys 0.59.0", -] - [[package]] name = "rustix" version = "1.1.4" @@ -4656,7 +4342,7 @@ dependencies = [ "bitflags 2.11.1", "errno", "libc", - "linux-raw-sys 0.12.1", + "linux-raw-sys", "windows-sys 0.61.2", ] @@ -5090,21 +4776,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" -[[package]] -name = "sideway" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7669500361477013bc37bcf21f6c9cd6ffcf4839b6fbffb3cb9df11e30a78c3e" -dependencies = [ - "bitmask-enum", - "libc", - "os_socketaddr", - "rdma-mummy-sys", - "serde", - "tabled", - "thiserror 2.0.18", -] - [[package]] name = "signal-hook-registry" version = "1.4.8" @@ -5163,15 +4834,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "spin" -version = "0.9.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" -dependencies = [ - "lock_api", -] - [[package]] name = "spm_precompiled" version = "0.1.4" @@ -5367,16 +5029,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "syscalls" -version = "0.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90db46b5b4962319605d435986c775ea45a0ad2561c09e1d5372b89afeb49cf4" -dependencies = [ - "serde", - "serde_repr", -] - [[package]] name = "system-configuration" version = "0.7.0" @@ -5398,29 +5050,6 @@ dependencies = [ "libc", ] -[[package]] -name = "tabled" -version = "0.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "121d8171ee5687a4978d1b244f7d99c43e7385a272185a2f1e1fa4dc0979d444" -dependencies = [ - "papergrid", - "tabled_derive", -] - -[[package]] -name = "tabled_derive" -version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52d9946811baad81710ec921809e2af67ad77719418673b2a3794932d57b7538" -dependencies = [ - "heck", - "proc-macro-error2", - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "task-local" version = "0.1.1" @@ -5461,16 +5090,10 @@ dependencies = [ "fastrand", "getrandom 0.4.2", "once_cell", - "rustix 1.1.4", + "rustix", "windows-sys 0.61.2", ] -[[package]] -name = "termtree" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" - [[package]] name = "thiserror" version = "1.0.69" @@ -6742,18 +6365,6 @@ version = "0.1.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a28ac98ddc8b9274cb41bb4d9d4d5c425b6020c50c46f25559911905610b4a88" -[[package]] -name = "which" -version = "4.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" -dependencies = [ - "either", - "home", - "once_cell", - "rustix 0.38.44", -] - [[package]] name = "win-sys" version = "0.3.1" @@ -6897,15 +6508,6 @@ dependencies = [ "windows-targets", ] -[[package]] -name = "windows-sys" -version = "0.59.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" -dependencies = [ - "windows-targets", -] - [[package]] name = "windows-sys" version = "0.61.2" diff --git a/Cargo.toml b/Cargo.toml index 05c388cfb..cd382fc35 100644 --- a/Cargo.toml +++ b/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" diff --git a/docs/index.md b/docs/index.md index d4c38c8ea..99ac4c1e5 100644 --- a/docs/index.md +++ b/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 diff --git a/docs/subsystems/runtime/external-pegaflow-server.md b/docs/subsystems/runtime/external-pegaflow-server.md new file mode 100644 index 000000000..91d6d3007 --- /dev/null +++ b/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). diff --git a/docs/subsystems/runtime/pegaflow-offload-integration.md b/docs/subsystems/runtime/pegaflow-offload-integration.md index 5db431c6c..9ecbee249 100644 --- a/docs/subsystems/runtime/pegaflow-offload-integration.md +++ b/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） diff --git a/openinfer-glm52/src/lib.rs b/openinfer-glm52/src/lib.rs index 300b89f55..f1bfa0d19 100644 --- a/openinfer-glm52/src/lib.rs +++ b/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,8 +342,8 @@ 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 @@ -352,35 +351,16 @@ mod topology_tests { /// 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, + /// Out-of-process PegaFlow gRPC endpoint. + pub server_addr: String, + /// Stable deployment/checkpoint identity for native OpenInfer peers. + pub namespace: Option, /// `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, } -/// 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, -} - /// 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 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) diff --git a/openinfer-glm52/src/scheduler/mod.rs b/openinfer-glm52/src/scheduler/mod.rs index 1847a04a8..dfe6feb40 100644 --- a/openinfer-glm52/src/scheduler/mod.rs +++ b/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); } diff --git a/openinfer-glm52/src/scheduler/offload.rs b/openinfer-glm52/src/scheduler/offload.rs index 1171d6528..0d5c41db3 100644 --- a/openinfer-glm52/src/scheduler/offload.rs +++ b/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); diff --git a/openinfer-kernels/src/imported.rs b/openinfer-kernels/src/imported.rs new file mode 100644 index 000000000..b8c8c52f9 --- /dev/null +++ b/openinfer-kernels/src/imported.rs @@ -0,0 +1,114 @@ +//! GPU memory imported from an out-of-process PegaFlow server. +//! +//! PegaFlow allocates the fused KV arena in its own process (which is what +//! lets it register the memory into the NIC for GPUDirect RDMA — owner-side +//! `ibv_reg_mr`/dma-buf registration works, an IPC-imported pointer never +//! does) and returns a CUDA IPC handle in the registration response. This side +//! only runs compute kernels on the imported mapping, which CUDA IPC fully +//! supports. + +use std::sync::Arc; + +use cudarc::driver::CudaStream; +use cudarc::driver::result::DriverError; +use cudarc::driver::sys::CUresult; +use cudarc::driver::sys::{self}; + +/// Byte length of a serialized `CUipcMemHandle`. +pub const IPC_HANDLE_BYTES: usize = 64; + +/// A device mapping of a PegaFlow-owned allocation, imported via CUDA IPC. +/// +/// The server owns the memory and frees it on unregister; this mapping must be +/// closed (drop) before the engine unregisters. The pointer is stable for the +/// mapping's lifetime. +pub struct ImportedKvArena { + stream: Arc, + ptr: sys::CUdeviceptr, + size_bytes: usize, +} + +impl ImportedKvArena { + /// Import `ipc_handle` (the 64-byte `CUipcMemHandle` from PegaFlow's + /// registration response) on `stream`'s device. `size_bytes` is the arena + /// size this side needs; the actual allocation is queried from the driver + /// and must cover it — trusting the client's own number would let a + /// server-side under-allocation turn into silent out-of-bounds writes. + pub fn open( + stream: &Arc, + ipc_handle: &[u8], + size_bytes: usize, + ) -> Result { + if ipc_handle.len() != IPC_HANDLE_BYTES || size_bytes == 0 { + return Err(DriverError(CUresult::CUDA_ERROR_INVALID_VALUE)); + } + let ctx = stream.context(); + ctx.bind_to_thread()?; + + let mut handle = sys::CUipcMemHandle { reserved: [0; 64] }; + for (dst, src) in handle.reserved.iter_mut().zip(ipc_handle) { + *dst = *src as i8; + } + let mut ptr: sys::CUdeviceptr = 0; + // SAFETY: the handle references a live server-owned allocation; the + // server keeps it alive until this instance unregisters. + unsafe { + sys::cuIpcOpenMemHandle_v2( + &raw mut ptr, + handle, + sys::CUipcMem_flags_enum::CU_IPC_MEM_LAZY_ENABLE_PEER_ACCESS as u32, + ) + .result()?; + } + + let mut range_base: sys::CUdeviceptr = 0; + let mut range_size: usize = 0; + // SAFETY: ptr is a live imported mapping; the driver reports the + // allocation it belongs to. + let range = unsafe { + sys::cuMemGetAddressRange_v2(&raw mut range_base, &raw mut range_size, ptr).result() + }; + let actual = match range { + Ok(()) => range_size - (ptr - range_base) as usize, + Err(err) => { + // SAFETY: close the mapping we just opened before bailing. + unsafe { sys::cuIpcCloseMemHandle(ptr).result().ok() }; + return Err(err); + } + }; + if actual < size_bytes { + // SAFETY: as above. + unsafe { sys::cuIpcCloseMemHandle(ptr).result().ok() }; + return Err(DriverError(CUresult::CUDA_ERROR_INVALID_VALUE)); + } + + Ok(Self { + stream: Arc::clone(stream), + ptr, + size_bytes, + }) + } + + /// Base device address of the imported arena. Stable for the mapping's + /// lifetime. + pub fn device_ptr(&self) -> u64 { + self.ptr + } + + /// Arena size in bytes as registered with the server. + pub fn size_bytes(&self) -> usize { + self.size_bytes + } +} + +impl Drop for ImportedKvArena { + fn drop(&mut self) { + self.stream + .context() + .bind_to_thread() + .expect("bind CUDA context before closing imported KV arena"); + // SAFETY: `ptr` came from cuIpcOpenMemHandle in `open`. + unsafe { sys::cuIpcCloseMemHandle(self.ptr).result() } + .expect("close imported KV arena mapping"); + } +} diff --git a/openinfer-kernels/src/lib.rs b/openinfer-kernels/src/lib.rs index 6ae6e8725..202ea6084 100644 --- a/openinfer-kernels/src/lib.rs +++ b/openinfer-kernels/src/lib.rs @@ -4,6 +4,7 @@ pub mod ffi; pub mod forward_pass; pub mod gpu_buffers; +pub mod imported; pub mod ops; pub mod paged_kv; pub mod tensor; diff --git a/openinfer-kv-cache/src/buffer.rs b/openinfer-kv-cache/src/buffer.rs index a055b003d..94451435d 100644 --- a/openinfer-kv-cache/src/buffer.rs +++ b/openinfer-kv-cache/src/buffer.rs @@ -1,14 +1,52 @@ +use std::mem::ManuallyDrop; use std::sync::Arc; use cudarc::driver::CudaSlice; use cudarc::driver::CudaStream; use cudarc::driver::DevicePtr; use half::bf16; +use openinfer_kernels::imported::ImportedKvArena; use crate::KvLayout; +/// Backing memory for the fused KV arena. +/// +/// The default path owns a normal cudarc `CudaSlice`. The offload path is a +/// non-owning view over a PegaFlow-allocated arena the executor imported over +/// CUDA IPC; the `CudaSlice` must not free that pointer, so it is wrapped in +/// `ManuallyDrop`. Whoever owns the `ImportedKvArena` controls the mapping +/// lifetime: it must outlive every kernel that touches this buffer and be +/// closed before the server frees the allocation. +enum Backing { + Owned(CudaSlice), + ImportedView(ManuallyDrop>), +} + +impl Backing { + fn view(&self) -> &CudaSlice { + match self { + Backing::Owned(slice) => slice, + Backing::ImportedView(view) => view, + } + } +} + +impl Drop for Backing { + fn drop(&mut self) { + if let Backing::ImportedView(view) = self { + // Drop the view WITHOUT freeing the imported pointer: leak() + // reclaims the raw ptr and runs the CudaSlice field teardown + // (events/stream) without any cuMemFree. The mapping itself is + // closed by the `ImportedKvArena` owner. + // SAFETY: `view` is not used again; ManuallyDrop::take moves it out. + let slice = unsafe { ManuallyDrop::take(view) }; + let _raw_ptr = slice.leak(); + } + } +} + struct Inner { - buffer: CudaSlice, + backing: Backing, layout: KvLayout, num_blocks: usize, } @@ -33,16 +71,60 @@ impl KvBuffer { ) -> anyhow::Result { let layout = KvLayout::new(num_layers, num_kv_heads, head_dim, page_size); let total_elements = num_blocks * layout.page_stride; - let buffer: CudaSlice = stream + let slice: CudaSlice = stream .alloc_zeros(total_elements) .map_err(|e| anyhow::anyhow!("KvBuffer alloc failed: {e}"))?; - Ok(Self { + Ok(Self::from_backing( + Backing::Owned(slice), + layout, + num_blocks, + )) + } + + /// Build the KV buffer as a view over a PegaFlow-allocated arena the + /// caller imported via CUDA IPC. The caller keeps the `ImportedKvArena` + /// alive for as long as any kernel touches this buffer and closes it + /// before the server frees the allocation. The server zeroed the arena at + /// allocation, so a KV hit never reads stale device memory. + pub fn new_imported( + stream: &Arc, + num_layers: usize, + num_kv_heads: usize, + head_dim: usize, + page_size: usize, + num_blocks: usize, + arena: &ImportedKvArena, + ) -> anyhow::Result { + let layout = KvLayout::new(num_layers, num_kv_heads, head_dim, page_size); + let total_elements = num_blocks * layout.page_stride; + let bytes = total_elements + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| anyhow::anyhow!("KvBuffer size overflows usize"))?; + anyhow::ensure!( + arena.size_bytes() >= bytes, + "imported arena is {} bytes but the KV layout needs {bytes}", + arena.size_bytes() + ); + // Wrap the imported pointer as a CudaSlice view for the attention + // kernels. + // SAFETY: the caller keeps the mapping valid for `total_elements` bf16; + // the view is ManuallyDrop so it never frees the pointer. + let view = unsafe { stream.upgrade_device_ptr::(arena.device_ptr(), total_elements) }; + Ok(Self::from_backing( + Backing::ImportedView(ManuallyDrop::new(view)), + layout, + num_blocks, + )) + } + + fn from_backing(backing: Backing, layout: KvLayout, num_blocks: usize) -> Self { + Self { inner: Arc::new(Inner { - buffer, + backing, layout, num_blocks, }), - }) + } } pub fn layout(&self) -> &KvLayout { @@ -50,18 +132,15 @@ impl KvBuffer { } pub fn buffer(&self) -> &CudaSlice { - &self.inner.buffer + self.inner.backing.view() } /// Base device address of the fused KV buffer. /// - /// Stable for the buffer's lifetime — cudarc allocations don't move — so - /// the KV-offload connector registers this once with pegaflow and the - /// page-first [`KvLayout`] strides reach every (layer, block, K/V) segment - /// from it. The returned address outlives the transient stream-ordering - /// guard precisely because the `Arc` keeps the slice alive. + /// Stable for the buffer's lifetime, so the page-first [`KvLayout`] strides + /// reach every (layer, block, K/V) segment from it. pub fn device_ptr(&self, stream: &CudaStream) -> u64 { - let (ptr, _guard) = self.inner.buffer.device_ptr(stream); + let (ptr, _guard) = self.inner.backing.view().device_ptr(stream); ptr } diff --git a/openinfer-kv-cache/src/manager.rs b/openinfer-kv-cache/src/manager.rs index 8564b72d9..a44805b76 100644 --- a/openinfer-kv-cache/src/manager.rs +++ b/openinfer-kv-cache/src/manager.rs @@ -38,6 +38,30 @@ impl KvCacheManager { Ok(Self { pool, buffer }) } + /// Like [`new`](Self::new) but over a PegaFlow-allocated arena imported + /// via CUDA IPC (see [`KvBuffer::new_imported`]). + pub fn new_imported( + stream: &Arc, + num_layers: usize, + num_kv_heads: usize, + head_dim: usize, + block_size: usize, + num_blocks: usize, + arena: &openinfer_kernels::imported::ImportedKvArena, + ) -> anyhow::Result { + let buffer = KvBuffer::new_imported( + stream, + num_layers, + num_kv_heads, + head_dim, + block_size, + num_blocks, + arena, + )?; + let pool = BlockPool::new(block_size, num_blocks)?; + Ok(Self { pool, buffer }) + } + /// Like [`new`](Self::new) but the pool emits KV block events; returns the /// receiver to drain. See [`BlockPool::with_events`]. pub fn new_with_events( diff --git a/openinfer-kv-offload/Cargo.toml b/openinfer-kv-offload/Cargo.toml index 0997cd005..758eba31b 100644 --- a/openinfer-kv-offload/Cargo.toml +++ b/openinfer-kv-offload/Cargo.toml @@ -5,29 +5,22 @@ name = "openinfer-kv-offload" version = "0.1.0" [dependencies] -# Embedded in-process: pegaflow-core is the host/SSD/RDMA KV offload tier. -# Pinned to the upstream master rev that landed PR #395 (P-side partial-tail -# save + first-token router forwarding), including PR #381's metaserver flush -# barrier and embeddable P2pTransferService. The supported producer versions -# and default layouts are pinned in the P/D deployment docs; they register the -# same block-local bytes OpenInfer restores, so no cross-layout conversion is -# required. -# per the workflow in docs/subsystems/runtime/pegaflow-offload-integration.md -# §5.2. default-features=false drops its `cuda-12` (which would pull a -# clashing cudarc/cuda-12080 selector); the workspace cudarc already provides -# cuda-12090 + nvrtc + runtime, which covers pegaflow-core and -# pegaflow-transfer. `rdma` enables pegaflow-transfer: cross-node block fetch -# (RDMA READ) + MetaServer discovery — the P/D disaggregation data plane. +# The client shares PegaFlow's wire-compatible error/save types but does not +# embed its storage engine or RDMA stack. Pinned to the native-arena server: +# registration returns a CUDA IPC handle for the server-allocated KV arena. +pegaflow-core = { git = "https://github.com/novitalabs/pegaflow.git", rev = "6043c290000a1cc5a384cff59bda6ecaaf868027", default-features = false } +pegaflow-proto = { git = "https://github.com/novitalabs/pegaflow.git", rev = "6043c290000a1cc5a384cff59bda6ecaaf868027" } cudarc = { workspace = true } half = { workspace = true } log = { workspace = true } openinfer-kv-cache = { workspace = true } -pegaflow-core = { git = "https://github.com/novitalabs/pegaflow.git", rev = "1473c5355d879b4fea23101760cb2a0074642ada", default-features = false, features = [ - "rdma", -] } tokio = { workspace = true } -tokio-stream = { workspace = true, features = ["net"] } +tonic = { workspace = true } +uuid = { workspace = true } xxhash-rust = { workspace = true } +[dev-dependencies] +openinfer-kernels = { workspace = true } + [lints] workspace = true diff --git a/openinfer-kv-offload/src/engine.rs b/openinfer-kv-offload/src/engine.rs index fe5b367e2..2ef79bb0f 100644 --- a/openinfer-kv-offload/src/engine.rs +++ b/openinfer-kv-offload/src/engine.rs @@ -1,124 +1,70 @@ -//! [`OffloadEngine`]: the in-process connector that moves KV blocks between -//! openinfer's GPU paged cache and pegaflow's host/SSD tiers. +//! Client-side bridge between OpenInfer GPU KV cache and a PegaFlow server. //! -//! It owns a [`PegaEngine`] plus a small tokio runtime to drive pegaflow's -//! async save/query, and translates openinfer's page-first [`KvLayout`] into -//! pegaflow's per-layer strided registration. Block content hashes are opaque -//! `Vec` here — the caller (scheduler) derives them from kvbm sequence -//! hashes, so this layer never depends on the logical-cache hashing scheme. +//! The server owns the fused KV arena (it allocates it at registration and +//! returns a CUDA IPC handle) plus every deeper tier and the transfer workers. +//! This module registers the KV layout, hands the arena handle back to the +//! caller for import, and translates scheduler save/query/load operations into +//! PegaFlow RPCs. use std::sync::Arc; use std::sync::Mutex; +use std::time::Duration; -use cudarc::driver::CudaStream; -use openinfer_kv_cache::KvBuffer; +use openinfer_kv_cache::KvLayout; use pegaflow_core::EngineError; use pegaflow_core::LayerSave; -use pegaflow_core::P2pTransferService; -use pegaflow_core::PegaEngine; -use pegaflow_core::PrefetchStatus; -use pegaflow_core::QueryLeaseId; -use pegaflow_core::StorageConfig; -use pegaflow_core::TransferMode; use tokio::runtime::Runtime; use tokio::sync::oneshot; use tokio::task::JoinHandle; -/// Single-GPU, single-rank topology. The dense Qwen3 path runs one offload -/// engine per executor rank, each owning one GPU's KV buffer. +use crate::external::ExternalClient; +use crate::external::ExternalQuery; +use crate::external::ExternalRegistration; +use crate::external::ExternalSession; + const TP_RANK: usize = 0; const PP_RANK: usize = 0; const TP_SIZE: usize = 1; const WORLD_SIZE: usize = 1; - -/// bf16 KV cache: every layout stride is counted in elements, bytes are ×2. const ELEM_SIZE: usize = std::mem::size_of::(); +const FLUSH_DEADLINE: Duration = Duration::from_secs(5); -/// Upper bound on the [`OffloadEngine::flush_saves_then`] barrier. Generous -/// for the normal case (D2H drain + a few local RPCs complete in -/// milliseconds); the cap only bites when the MetaServer connection stalls -/// mid-RPC, where the alternative is withholding finished requests' responses -/// for the TCP keepalive window. -const FLUSH_DEADLINE: std::time::Duration = std::time::Duration::from_secs(5); - -/// Guard the `block_on` entry points: tokio panics with an opaque message if -/// you block on a runtime from within any runtime. These methods are meant for -/// the synchronous scheduler thread — fail loud and specific if that's violated. fn assert_outside_runtime(op: &str) { debug_assert!( tokio::runtime::Handle::try_current().is_err(), "OffloadEngine::{op} drives the offload runtime with block_on and must be \ - called from a synchronous thread, never from within a tokio runtime" + called from a synchronous thread" ); } -/// Cross-instance P2P sharing over pegaflow's MetaServer + RDMA data plane. -/// -/// With this set, the engine (a) registers saved block hashes with the -/// MetaServer, (b) serves peer RDMA fetches on `listen_addr`, and (c) on a -/// local host-tier miss, discovers and pulls the missing prefix from whichever -/// peer owns it (one-sided RDMA READ into the local pinned pool, then a normal -/// H2D load). This is the P/D disaggregation data plane: a decode node finds -/// the prefill node's KV by content hash — no handle protocol. -#[derive(Clone, Debug)] -pub struct P2pConfig { - /// MetaServer gRPC address, e.g. `http://10.0.0.100:50056`. - pub metaserver_addr: String, - /// This engine's routable `IP:port` (a literal socket address — it doubles - /// as the embedded transfer service's bind address, so hostnames are - /// rejected at startup). Peers dial it for RDMA handshakes and block - /// queries, and the MetaServer records it as the block owner. Must not be - /// 0.0.0.0/127.0.0.1 for cross-node use. - pub advertise_addr: String, - /// RDMA NIC device names to register the pinned pool on (e.g. `mlx5_0`). - pub rdma_nics: Vec, -} - -/// Tuning knobs for a new [`OffloadEngine`]. +/// Connection settings for one rank-level offload engine. pub struct OffloadConfig { - /// Stable identifier shared across this engine's lifetime so prefix blocks - /// saved by one request are query-visible to the next. - instance_id: String, - /// Content-addressing domain shared with P2P peers: two engines see each - /// other's blocks iff their namespaces match. Callers derive it from - /// whatever makes KV layouts interchange-safe (model, dtype, block - /// geometry). Single-node offload can use any constant. - namespace: String, - /// CUDA device ordinal whose KV buffer this engine offloads. - device_id: i32, - /// Host pinned-memory pool size in bytes (the CPU KV tier capacity). - pinned_pool_bytes: usize, - /// Back the pinned pool with hugepages (see [`HostConfig::use_hugepages`]). - pub use_hugepages: bool, - /// Worker threads for the embedded runtime that drives pegaflow's async - /// save/query. Two is plenty: save is fire-and-forget, query is a brief - /// memory-cache lookup. - runtime_threads: usize, - /// `Some` joins the cross-instance P2P mesh (see [`P2pConfig`]). - p2p: Option, + /// Base identifier for this rank. A process-local suffix is added so a + /// restarted client cannot collide with a session still being cleaned up. + pub instance_id: String, + /// Content-addressing domain shared by producers and consumers whose KV + /// bytes and block layout are interchangeable. + pub namespace: String, + /// CUDA device ordinal the server allocates the KV arena on. + pub device_id: i32, + /// PegaFlow gRPC endpoint, for example `http://127.0.0.1:50055`. + pub server_addr: String, + /// Worker threads for the client runtime. + pub runtime_threads: usize, } impl OffloadConfig { - pub fn new(instance_id: impl Into, device_id: i32, pinned_pool_bytes: usize) -> Self { + pub fn new( + instance_id: impl Into, + device_id: i32, + server_addr: impl Into, + ) -> Self { Self { instance_id: instance_id.into(), namespace: "openinfer".to_string(), device_id, - pinned_pool_bytes, - use_hugepages: false, + server_addr: server_addr.into(), runtime_threads: 2, - p2p: None, - } - } - - /// The host-tier half of this config (private-host constructors split it - /// off before consuming the instance fields). - fn host(&self) -> HostConfig { - HostConfig { - pinned_pool_bytes: self.pinned_pool_bytes, - use_hugepages: self.use_hugepages, - runtime_threads: self.runtime_threads, - p2p: self.p2p.clone(), } } @@ -127,74 +73,59 @@ impl OffloadConfig { self.namespace = namespace.into(); self } - - #[must_use] - pub fn with_p2p(mut self, p2p: P2pConfig) -> Self { - self.p2p = Some(p2p); - self - } } -/// A query hit: how many prefix blocks pegaflow can return from its CPU tier, -/// and the lease that owns those blocks until [`OffloadEngine::load`] consumes -/// it. `num_blocks == 0` means a full miss and `lease` is `None`. +/// Opaque ownership token returned by a PegaFlow prefix query. +pub struct QueryLeaseId(Vec); + pub struct QueryHit { pub lease: Option, pub num_blocks: usize, } -/// Outcome of [`OffloadEngine::query`]. pub enum QueryOutcome { - /// Terminal: `hit.num_blocks` prefix blocks are host-resident and leased. Ready(QueryHit), - /// pegaflow kicked off an async fetch of the missing prefix from a remote - /// peer (P2P) or SSD. Not terminal: re-`query` with the same `req_id` next - /// tick to poll; the fetch resolves to `Ready` (with the pulled blocks) or - /// falls back to a plain local hit count. Only occurs with a deeper tier - /// configured — never in the host-memory-only setup. + /// PegaFlow is fetching the prefix from a deeper tier. Re-query with the + /// same request id on a later scheduler tick. Loading, } -/// In-flight handle for a CPU→GPU load submitted to pegaflow's worker. -/// -/// The load runs on pegaflow's GPU worker thread; this resolves when the DMA -/// completes. [`Self::poll`] keeps scheduler admission non-blocking; [`Self::wait`] -/// blocks for tests and non-pipelined callers. +/// In-flight server-side CPU-to-GPU load. pub struct LoadHandle { - rx: oneshot::Receiver>, + submission: Option>>, } impl LoadHandle { - /// Non-blocking check for a scheduler tick. `None` while still loading. pub fn poll(&mut self) -> Option> { - match self.rx.try_recv() { - Ok(result) => Some(result), - Err(oneshot::error::TryRecvError::Empty) => None, - Err(oneshot::error::TryRecvError::Closed) => Some(Err(EngineError::Storage( - "load worker dropped reply".into(), - ))), + if let Some(rx) = &mut self.submission { + match rx.try_recv() { + Ok(Ok(())) => self.submission = None, + Ok(Err(err)) => return Some(Err(err)), + Err(oneshot::error::TryRecvError::Empty) => return None, + Err(oneshot::error::TryRecvError::Closed) => { + return Some(Err(EngineError::Storage( + "external load submission dropped reply".into(), + ))); + } + } } + Some(Ok(())) } - /// Block the current thread until the load settles. - pub fn wait(self) -> Result<(), EngineError> { - self.rx - .blocking_recv() - .unwrap_or_else(|_| Err(EngineError::Storage("load worker dropped reply".into()))) + pub fn wait(mut self) -> Result<(), EngineError> { + if let Some(rx) = self.submission.take() { + rx.blocking_recv().unwrap_or_else(|_| { + Err(EngineError::Storage( + "external load submission dropped reply".into(), + )) + })?; + } + Ok(()) } } -/// One strided GPU arena to register as one pegaflow "layer": `num_blocks` -/// copy units of `bytes_per_block`, sitting `block_stride_bytes` apart from -/// `base_ptr`. A fused buffer (qwen3) contributes one arena per model layer; -/// a model with sidecar caches (GLM5.2: MLA latent + index-K per layer, two -/// separate allocations sharing pool block ids) contributes several arenas -/// per model layer — pegaflow moves whatever arenas are registered under one -/// block id together, which is what keeps sidecars in lockstep with their -/// main cache. -/// -/// `name` keys the arena for the whole engine lifetime (save/load fan across -/// every registered name); it must be unique within the engine. +/// One strided GPU arena registered as one PegaFlow layer. Only meaningful for +/// the (unsupported) multi-arena path; see [`OffloadEngine::with_arenas_on`]. pub struct KvArena { pub name: String, pub base_ptr: u64, @@ -203,300 +134,120 @@ pub struct KvArena { pub block_stride_bytes: usize, } -/// Per-layer registration geometry fed to pegaflow's one batched call. -/// -/// Only `data_ptrs` and `size_bytes` differ per layer; the rest are the same -/// scalar broadcast across all layers (kept as vectors only to feed pegaflow's -/// one batched registration call). +/// Per-layer views into the single server-allocated arena, as offsets. struct Registration { layer_names: Vec, - data_ptrs: Vec, + offset_bytes: Vec, size_bytes: Vec, num_blocks: Vec, bytes_per_block: Vec, kv_stride_bytes: Vec, segments: Vec, block_stride_bytes: Vec, + /// Total arena size the server must allocate. + alloc_size: usize, } impl Registration { - /// Map the fused page-first buffer to pegaflow's per-layer view. - /// - /// Each model layer registers as one pegaflow "layer". Within a page the - /// layout is K then V back-to-back (`layer_stride = 2·kv_block_len`), so a - /// layer's K and V are *contiguous* — one single segment of `layer_stride` - /// bytes copies both, and pegaflow's K/V-split path (which needs the two - /// segments set apart, `kv_stride > bytes_per_block`) does not apply here. - /// What is *not* contiguous is consecutive blocks of one layer: the fused - /// buffer interleaves all layers within a page, so they sit `page_stride` - /// apart. That gap (stride ≠ copy size) is exactly what `block_stride_bytes` - /// decouples. - fn from_buffer(buffer: &KvBuffer, stream: &CudaStream) -> Self { - let layout = buffer.layout(); - let num_blocks = buffer.num_blocks(); - let base_ptr = buffer.device_ptr(stream); - - // One block's copy unit for a layer = its whole [K|V] span in a page. + /// Carve one layer view per model layer out of the page-first fused arena: + /// block `b` of layer `l` lives at `b * page_stride + l * layer_stride` + /// (in elements), so a layer's blocks are `page_stride` apart. + fn from_layout(layout: &KvLayout, num_blocks: usize) -> Self { + assert!(num_blocks > 0, "KV offload requires at least one block"); + assert!(layout.num_layers > 0, "KV offload requires layers"); let layer_bytes = layout.layer_stride * ELEM_SIZE; let page_stride_bytes = layout.page_stride * ELEM_SIZE; - - let arenas: Vec = (0..layout.num_layers) - .map(|layer| KvArena { - name: layer.to_string(), - base_ptr: base_ptr + (layer * layer_bytes) as u64, - num_blocks, - bytes_per_block: layer_bytes, - block_stride_bytes: page_stride_bytes, - }) - .collect(); - Self::from_arenas(&arenas) - } - - /// One pegaflow layer per arena, single-segment (an arena is one copy - /// unit per block by definition; K/V split segments only exist for the - /// symmetric-pair layouts vLLM registers). - fn from_arenas(arenas: &[KvArena]) -> Self { - let n = arenas.len(); - let mut reg = Registration { - layer_names: Vec::with_capacity(n), - data_ptrs: Vec::with_capacity(n), - size_bytes: Vec::with_capacity(n), - num_blocks: Vec::with_capacity(n), - bytes_per_block: Vec::with_capacity(n), + assert!( + layer_bytes > 0 && layer_bytes <= page_stride_bytes, + "layer copy unit {layer_bytes} overruns its page stride {page_stride_bytes}" + ); + let alloc_size = num_blocks + .checked_mul(page_stride_bytes) + .expect("KV arena size overflows usize"); + let reach = (num_blocks - 1) * page_stride_bytes + layer_bytes; + let n = layout.num_layers; + Self { + layer_names: (0..n).map(|layer| layer.to_string()).collect(), + offset_bytes: (0..n).map(|layer| (layer * layer_bytes) as u64).collect(), + size_bytes: vec![reach; n], + num_blocks: vec![num_blocks; n], + bytes_per_block: vec![layer_bytes; n], kv_stride_bytes: vec![0; n], segments: vec![1; n], - block_stride_bytes: Vec::with_capacity(n), - }; - for arena in arenas { - assert!( - arena.bytes_per_block <= arena.block_stride_bytes, - "arena {} copy unit {} overruns its block stride {}", - arena.name, - arena.bytes_per_block, - arena.block_stride_bytes - ); - reg.layer_names.push(arena.name.clone()); - reg.data_ptrs.push(arena.base_ptr); - // The arena's region must cover the strided reach of its last - // block (pegaflow validates copies against this bound). - reg.size_bytes - .push((arena.num_blocks - 1) * arena.block_stride_bytes + arena.bytes_per_block); - reg.num_blocks.push(arena.num_blocks); - reg.bytes_per_block.push(arena.bytes_per_block); - reg.block_stride_bytes.push(arena.block_stride_bytes); + block_stride_bytes: vec![page_stride_bytes; n], + alloc_size, } - reg } } -/// Host-tier knobs for a shared [`OffloadHost`]. -pub struct HostConfig { - /// Host pinned-memory pool size in bytes (the CPU KV tier capacity). - pub pinned_pool_bytes: usize, - /// Back the pinned pool with hugepages (pegaflow supports it natively). - /// Verify the box actually holds a reservation (`HugePages_Total`) — - /// some cluster platforms re-claim it across reboots. - pub use_hugepages: bool, - /// Worker threads for the runtime that drives pegaflow's async save/query. - pub runtime_threads: usize, - /// `Some` joins the cross-instance P2P mesh (see [`P2pConfig`]). - pub p2p: Option, -} - -/// The shared side of the offload: one [`PegaEngine`] (one host pool), the -/// tokio runtime that drives it, and the optional P2P serving lifecycle. -/// -/// One host serves any number of rank-level [`OffloadEngine`]s. That is the -/// DP-rank sharing model: each rank registers its own GPU arenas as its own -/// pegaflow *instance*, but blocks land in the one host tier keyed by -/// `(namespace, hash)` — with a shared namespace, any rank restores what any -/// rank saved. Callers share a namespace only when their KV is -/// interchangeable across instances: for replicated-weight DP ranks that -/// holds to the same tolerance as reusing a rank's own prefix cache (the -/// bytes may differ by FP reduction order across batch shapes, exactly like -/// two local recomputations of the same prefix would). -/// -/// Dropping the last handle drops the [`Runtime`], which abandons any -/// in-flight fire-and-forget saves (acceptable — the host tier is a cache) -/// and stops the P2P serving tasks; peers degrade to their own local -/// prefill. In-flight [`OffloadEngine::flush_saves_then`] barriers are -/// cancelled too, dropping their `then` callbacks unrun. +/// Shared connection and runtime for rank-level engines in one process. pub struct OffloadHost { - engine: Arc, + client: ExternalClient, runtime: Runtime, - /// `Some` when P2P is on: resolves the P2P serving tasks (gRPC transfer - /// service + transfer-lock GC) on drop. - #[allow(dead_code)] - p2p_shutdown: Option>, + instance_suffix: String, } impl OffloadHost { - pub fn new(config: HostConfig) -> Result, EngineError> { + pub fn connect(server_addr: &str, runtime_threads: usize) -> Result, EngineError> { + if server_addr.is_empty() { + return Err(EngineError::InvalidArgument( + "PegaFlow server address must not be empty".into(), + )); + } let runtime = tokio::runtime::Builder::new_multi_thread() - .worker_threads(config.runtime_threads.max(1)) + .worker_threads(runtime_threads.max(1)) .enable_all() .build() - .map_err(|e| EngineError::Storage(format!("offload runtime build: {e}")))?; - - let mut storage_config = StorageConfig::default(); - if let Some(p2p) = &config.p2p { - if p2p.rdma_nics.is_empty() { - return Err(EngineError::InvalidArgument( - "P2P requires at least one RDMA NIC".into(), - )); - } - storage_config.rdma_nic_names = Some(p2p.rdma_nics.clone()); - storage_config.metaserver_addr = Some(p2p.metaserver_addr.clone()); - storage_config.advertise_addr = Some(p2p.advertise_addr.clone()); - } - // pegaflow's MetaServerClient spawns its background registration loop - // with tokio::spawn, so the engine must be built inside our runtime. - let engine = { - let _guard = runtime.enter(); - Arc::new(PegaEngine::new_with_config( - config.pinned_pool_bytes, - config.use_hugepages, - storage_config, - )?) - }; - - // P2P serving side: peers discovered us via the MetaServer and dial - // `advertise_addr` for the RDMA handshake + block queries. Same - // lifecycle as the engine — shut down (via the oneshot) on drop. - let p2p_shutdown = match config.p2p { - Some(p2p) => { - let listen: std::net::SocketAddr = p2p.advertise_addr.parse().map_err(|e| { - EngineError::InvalidArgument(format!( - "P2P advertise_addr {:?} is not a socket address: {e}", - p2p.advertise_addr - )) - })?; - let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); - let serve_engine = Arc::clone(&engine); - let (ready_tx, ready_rx) = std::sync::mpsc::channel::>(); - runtime.spawn(async move { - // Bind eagerly so startup fails loud on a taken port - // instead of P2P silently never serving. - let bound = tokio::net::TcpListener::bind(listen).await; - let listener = match bound { - Ok(l) => { - let _ = ready_tx.send(Ok(())); - l - } - Err(e) => { - let _ = ready_tx.send(Err(format!("bind {listen}: {e}"))); - return; - } - }; - let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener); - if let Err(e) = - P2pTransferService::serve_with_incoming(serve_engine, incoming, async { - let _ = shutdown_rx.await; - }) - .await - { - log::error!("P2P transfer service exited: {e}"); - } - }); - ready_rx - .recv() - .map_err(|_| EngineError::Storage("P2P serve task died at startup".into()))? - .map_err(EngineError::Storage)?; - - // Background GC, mirroring pegaflow-server's task. Two sweeps: - // expired transfer locks (a crashed peer must not pin our - // blocks past the lock timeout) and stale prefetch state — an - // abandoned remote fetch (request dropped mid-RemoteFetch, or - // the executor's re-query deadline fired) leaves an orphaned - // entry whose completed task pins its fetched blocks in the - // pinned pool until this sweep drops it. - let gc_engine = Arc::clone(&engine); - runtime.spawn(async move { - const STALE_MAX_AGE: std::time::Duration = std::time::Duration::from_mins(5); - let mut tick = tokio::time::interval(std::time::Duration::from_mins(1)); - tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - loop { - tick.tick().await; - let expired = gc_engine.gc_expired_transfer_locks(); - if expired > 0 { - log::warn!("P2P GC released {expired} expired transfer locks"); - } - let (stale, failed) = gc_engine - .gc_stale_inflight(STALE_MAX_AGE, STALE_MAX_AGE) - .await; - if stale > 0 || failed > 0 { - log::info!( - "P2P GC dropped {stale} stale prefetch entries, \ - {failed} failed-remote markers" - ); - } - } - }); - log::info!( - "KV offload P2P enabled: serving on {listen}, metaserver={}", - p2p.metaserver_addr - ); - Some(shutdown_tx) - } - None => None, - }; - + .map_err(|err| EngineError::Storage(format!("offload runtime build: {err}")))?; + let client = runtime.block_on(ExternalClient::connect(server_addr))?; + let instance_suffix = new_instance_suffix(); Ok(Arc::new(Self { - engine, + client, runtime, - p2p_shutdown, + instance_suffix, })) } + + fn instance_id(&self, base: &str) -> String { + format!("{base}-{}", self.instance_suffix) + } } -/// In-process bridge from one rank's GPU KV cache to pegaflow's offload -/// tiers, over a shared or private [`OffloadHost`]. -/// -/// Save is best-effort fire-and-forget (a lost save only forfeits a future -/// hit, never inference correctness); saves that must survive a handoff -/// (eviction) use the synchronous [`Self::save_blocking`]. Runtime and P2P -/// lifetime live on the host — see [`OffloadHost`] for drop semantics. +fn new_instance_suffix() -> String { + uuid::Uuid::new_v4().simple().to_string() +} + +/// Rank-level client for one server-allocated KV arena. pub struct OffloadEngine { + session: Option, host: Arc, instance_id: String, device_id: i32, - /// Owned per-layer names; load borrows these as `&[&str]`. layer_names: Vec, - /// In-flight fire-and-forget save tasks plus the completion signal of the - /// latest flush barrier. One lock so a barrier's "drain handles + chain - /// behind the previous barrier" is atomic — two racing barriers can never - /// each take half the coverage (see [`Self::flush_saves_then`]). write_barrier: Mutex, } -/// Save handles and barrier chain behind [`OffloadEngine::write_barrier`]. struct WriteBarrierState { - /// In-flight fire-and-forget save tasks; finished handles are pruned on - /// each [`OffloadEngine::save`]. pending_saves: Vec>, - /// Completion signal of the latest spawned flush barrier (fires on - /// success and deadline alike). `None` before the first barrier. prev_flush_done: Option>, } impl OffloadEngine { - /// Build a private host and register `buffer` as the GPU side of the - /// offload. - /// - /// `stream` must be the stream that owns `buffer` (used only to read its - /// base device address). pegaflow attaches the device's primary CUDA - /// context for its own worker transfers — the same context openinfer runs - /// on — so the registered pointers are valid across both. + /// Register the KV layout with the server, which allocates the fused arena + /// on `config.device_id` and returns its CUDA IPC handle. The caller + /// imports the handle (`ImportedKvArena::open`) and builds the `KvBuffer` + /// over it; the mapping must be closed before [`Self::shutdown`] lets the + /// server free the arena. pub fn new( - config: OffloadConfig, - buffer: &KvBuffer, - stream: &CudaStream, - ) -> Result { - let reg = Registration::from_buffer(buffer, stream); - let host = OffloadHost::new(config.host())?; + config: &OffloadConfig, + layout: &KvLayout, + num_blocks: usize, + ) -> Result<(Self, Vec), EngineError> { + let reg = Registration::from_layout(layout, num_blocks); + let host = OffloadHost::connect(&config.server_addr, config.runtime_threads)?; Self::register( host, - config.instance_id, + &config.instance_id, &config.namespace, config.device_id, reg, @@ -504,76 +255,58 @@ impl OffloadEngine { ) } - /// Build the engine over explicit arenas (instead of one fused - /// [`KvBuffer`]) onto an existing shared host: this rank becomes - /// one more pegaflow instance over the host's single pool. Ranks that - /// should see each other's blocks pass the same `namespace`. Every - /// arena's device allocation must stay live and pointer-stable for the - /// engine's lifetime (the registration bakes raw device addresses), and - /// all arenas must be indexed by the same pool block ids. - /// `page_first` must match how the namespace's writer stores blocks: the - /// vLLM connector stores MLA-model blocks page-first (all layers of a - /// block concatenated into one host page, offsets by lexicographic layer - /// name), so joining a vLLM MLA namespace requires `true` — with layer - /// names and per-layer block bytes identical to the writer's. + /// Multi-arena registration (GLM5.2's rank-local MLA/index-K arenas) is not + /// supported by the native-arena server contract, which allocates exactly + /// one fused arena per instance. Fails before touching the server. pub fn with_arenas_on( - host: Arc, + _host: Arc, instance_id: impl Into, - namespace: &str, - device_id: i32, - arenas: &[KvArena], - page_first: bool, + _namespace: &str, + _device_id: i32, + _arenas: &[KvArena], + _page_first: bool, ) -> Result { - Self::register( - host, - instance_id.into(), - namespace, - device_id, - Registration::from_arenas(arenas), - page_first, - ) + Err(EngineError::InvalidArgument(format!( + "multi-arena KV offload for instance {} is not supported by the native-arena \ + PegaFlow contract (one server-allocated arena per instance)", + instance_id.into() + ))) } fn register( host: Arc, - instance_id: String, + instance_id: &str, namespace: &str, device_id: i32, reg: Registration, page_first: bool, - ) -> Result { - host.engine.register_context_layer_batch_strided( - &instance_id, - namespace, - device_id, - TP_RANK, - PP_RANK, - TP_SIZE, - WORLD_SIZE, - ®.layer_names, - ®.data_ptrs, - ®.size_bytes, - ®.num_blocks, - ®.bytes_per_block, - ®.kv_stride_bytes, - ®.segments, - Some(reg.block_stride_bytes.as_slice()), - // Direct (cuMemcpyAsync on the DMA engines). The Kernel backend - // was A/B'd for the fragmented bulk-restore batches (#704) and - // measured WORSE for co-resident decode: its grid-strided copy - // kernels compete for SMs with decode kernels, stretching the - // stall (two ~110ms waves vs Direct's one). - TransferMode::Direct, - // Layer-first (false): one pegaflow layer per model layer, the - // page-interleaved gap expressed via `block_stride_bytes` — the - // native openinfer layout. Page-first (true) instead stores each - // block as one host page holding every layer at its - // name-sorted offset; used only to join a namespace whose writer - // (the vLLM connector on MLA models) stores blocks that way. - page_first, - )?; - - Ok(Self { + ) -> Result<(Self, Vec), EngineError> { + assert_outside_runtime("register"); + let instance_id = host.instance_id(instance_id); + let (session, arena_ipc_handle) = + host.runtime + .block_on(host.client.register(ExternalRegistration { + instance_id: &instance_id, + namespace, + device_id, + tp_rank: TP_RANK, + pp_rank: PP_RANK, + tp_size: TP_SIZE, + world_size: WORLD_SIZE, + layer_names: ®.layer_names, + offset_bytes: ®.offset_bytes, + size_bytes: ®.size_bytes, + num_blocks: ®.num_blocks, + bytes_per_block: ®.bytes_per_block, + kv_stride_bytes: ®.kv_stride_bytes, + block_stride_bytes: ®.block_stride_bytes, + segments: ®.segments, + page_first, + alloc_size: reg.alloc_size, + }))?; + + let engine = Self { + session: Some(session), host, instance_id, device_id, @@ -582,16 +315,15 @@ impl OffloadEngine { pending_saves: Vec::new(), prev_flush_done: None, }), - }) + }; + Ok((engine, arena_ipc_handle)) } - /// Fan one (block_id, hash) list across every layer — the device data - /// differs per layer, the ids and hashes don't. fn build_saves(&self, block_ids: &[i32], block_hashes: &[Vec]) -> Vec { - // pegaflow indexes GPU blocks by `usize`; openinfer carries them as - // `i32` (its kvbm/CUDA convention). Convert once at this boundary — - // block ids are slot indices, always non-negative. - let block_ids: Vec = block_ids.iter().map(|&id| id as usize).collect(); + let block_ids: Vec = block_ids + .iter() + .map(|&id| usize::try_from(id).expect("KV block id must be non-negative")) + .collect(); self.layer_names .iter() .map(|name| LayerSave { @@ -602,28 +334,8 @@ impl OffloadEngine { .collect() } - /// Save the named GPU blocks to the host tier — fire-and-forget. - /// - /// Best-effort by contract: the GPU→CPU copy runs on pegaflow's worker and - /// any failure (pinned pool full, copy error) is logged, never surfaced. - /// `block_hashes[i]` is the content hash of `block_ids[i]`; all layers share - /// the same (block_id, hash) pairing — only the device data differs. - /// - /// ORDERING CONTRACT: pegaflow's D2H runs on *its own* stream, with no - /// dependency on openinfer's compute stream. The caller must therefore only - /// save blocks whose KV writes are already complete — i.e. call this after - /// the producing forward step has synchronized (block-seal time, which is - /// post-step-sync in the executor). Saving a block whose attention write is - /// still in flight reads torn data. This connector cannot enforce the - /// invariant (it does not own the compute stream); the wiring must uphold it. - /// - /// REUSE CONTRACT: the copy reads the GPU block asynchronously *after* this - /// returns, so the block must stay stable until the copy lands. `keep_alive` - /// is an opaque payload (e.g. the source blocks' allocator guards) held for - /// the lifetime of the spawned save and dropped only once it finishes — so - /// the caller's blocks cannot be evicted and overwritten under the in-flight - /// D2H (which would snapshot the wrong KV and persist it under the old hash). - /// Pass `()` only when the blocks are owned elsewhere for the whole save. + /// Save GPU blocks without blocking the scheduler. `keep_alive` pins the + /// source blocks until the server confirms the copy submission completed. pub fn save( &self, block_ids: &[i32], @@ -635,36 +347,23 @@ impl OffloadEngine { return; } let saves = self.build_saves(block_ids, block_hashes); - let engine = Arc::clone(&self.host.engine); + let client = self.host.client.clone(); let instance_id = self.instance_id.clone(); let device_id = self.device_id; let handle = self.host.runtime.spawn(async move { - if let Err(e) = engine - .batch_save_kv_blocks_from_ipc(&instance_id, TP_RANK, PP_RANK, device_id, saves) + if let Err(err) = client + .save(&instance_id, TP_RANK, PP_RANK, device_id, saves) .await { - log::warn!("pegaflow save failed (best-effort): {e}"); + log::warn!("PegaFlow save failed (best-effort): {err}"); } - // Release the source-block pins only now the D2H has landed; before - // this point the blocks must not be reused (see REUSE CONTRACT). drop(keep_alive); }); - // Track for the flush barrier; prune the ones that already settled so - // the list stays bounded by the genuinely in-flight saves. let mut barrier = self.write_barrier.lock().expect("write_barrier poisoned"); - barrier.pending_saves.retain(|h| !h.is_finished()); + barrier.pending_saves.retain(|handle| !handle.is_finished()); barrier.pending_saves.push(handle); } - /// Save the named GPU blocks and block until the GPU→CPU copy has captured - /// the data into the host tier (the insert may still be in flight; pair with - /// [`Self::flush_saves`] for cache visibility). - /// - /// The synchronous contract is what makes this safe at eviction handoff: the - /// GPU block can be reused the moment this returns. Errors surface, unlike - /// the fire-and-forget [`Self::save`]. The same compute-stream ORDERING - /// CONTRACT as [`Self::save`] applies: blocking waits on pegaflow's D2H, not - /// on openinfer's compute stream, so the writes must already be complete. pub fn save_blocking( &self, block_ids: &[i32], @@ -676,25 +375,15 @@ impl OffloadEngine { } assert_outside_runtime("save_blocking"); let saves = self.build_saves(block_ids, block_hashes); - self.host - .runtime - .block_on(self.host.engine.batch_save_kv_blocks_from_ipc( - &self.instance_id, - TP_RANK, - PP_RANK, - self.device_id, - saves, - )) + self.host.runtime.block_on(self.host.client.save( + &self.instance_id, + TP_RANK, + PP_RANK, + self.device_id, + saves, + )) } - /// Look up how long a prefix of `block_hashes` is resident in the CPU tier. - /// - /// Returns [`QueryOutcome::Ready`] with the hit-block count and a lease - /// owning those blocks (pass the lease to [`Self::load`] to copy them to - /// GPU), or [`QueryOutcome::Loading`] when pegaflow is fetching the missing - /// prefix from a remote peer / SSD in the background — re-`query` with the - /// same `req_id` to poll. `req_id` must be non-empty and unique enough to - /// scope an in-flight prefetch (the request id works). pub fn query( &self, req_id: &str, @@ -707,118 +396,73 @@ impl OffloadEngine { })); } assert_outside_runtime("query"); - let status = - self.host - .runtime - .block_on(self.host.engine.count_prefix_hit_blocks_with_prefetch( - &self.instance_id, - req_id, - block_hashes, - ))?; - - match status { - PrefetchStatus::Loading => Ok(QueryOutcome::Loading), - PrefetchStatus::Ready { blocks, .. } => { - if blocks.is_empty() { - return Ok(QueryOutcome::Ready(QueryHit { - lease: None, - num_blocks: 0, - })); - } - let num_blocks = blocks.len(); - let lease = self - .host - .engine - .create_query_lease(&self.instance_id, blocks)?; - Ok(QueryOutcome::Ready(QueryHit { - lease: Some(lease), - num_blocks, - })) - } + match self.host.runtime.block_on(self.host.client.query( + &self.instance_id, + req_id, + block_hashes, + ))? { + ExternalQuery::Loading => Ok(QueryOutcome::Loading), + ExternalQuery::Ready { num_blocks, lease } => Ok(QueryOutcome::Ready(QueryHit { + lease: (num_blocks > 0).then_some(QueryLeaseId(lease)), + num_blocks, + })), } } - /// Copy the leased CPU blocks into the GPU blocks named by `dst_block_ids`, - /// across every registered layer. Returns a non-blocking [`LoadHandle`]. - /// - /// `dst_block_ids.len()` must equal the lease's block count (the - /// `num_blocks` from [`Self::query`]); pegaflow maps the i-th leased block - /// onto `dst_block_ids[i]` for each layer. pub fn load( &self, - lease: QueryLeaseId, + lease: &QueryLeaseId, dst_block_ids: Vec, ) -> Result { - let layer_refs: Vec<&str> = self.layer_names.iter().map(String::as_str).collect(); - // pegaflow indexes GPU blocks by `usize` (see `build_saves`). - let dst_block_ids: Vec = dst_block_ids.into_iter().map(|id| id as usize).collect(); - let loads = [(lease, dst_block_ids)]; - let rx = self.host.engine.batch_load_kv_blocks_multi_layer_inproc( - &self.instance_id, - TP_RANK, - self.device_id, - &layer_refs, - &loads, - )?; - Ok(LoadHandle { rx }) + let dst_block_ids: Vec = dst_block_ids + .into_iter() + .map(|id| usize::try_from(id).expect("KV block id must be non-negative")) + .collect(); + let instance_id = self.instance_id.clone(); + let layer_names = self.layer_names.clone(); + let client = self.host.client.clone(); + let device_id = self.device_id; + let lease = lease.0.clone(); + let (tx, rx) = oneshot::channel(); + self.host.runtime.spawn(async move { + let result = client + .load( + &instance_id, + TP_RANK, + device_id, + &layer_names, + lease, + dst_block_ids, + ) + .await; + let _ = tx.send(result); + }); + Ok(LoadHandle { + submission: Some(rx), + }) } - /// Release a query lease without loading it. - /// - /// [`Self::query`] pins its hit blocks behind a lease until [`Self::load`] - /// consumes it. When the caller decides not to load (e.g. no GPU - /// destination blocks are free), it must release the lease here — a dropped - /// [`QueryLeaseId`] is an inert token, so without this the pinned host - /// blocks would sit unevictable until the lease's TTL expires. A no-op if - /// the lease was already consumed by a `load`. pub fn release_query_lease(&self, lease: QueryLeaseId) { - self.host.engine.release_query_lease(&lease); + let client = self.host.client.clone(); + self.host.runtime.spawn(async move { + if let Err(err) = client.release(lease.0).await { + log::warn!("PegaFlow lease release failed: {err}"); + } + }); } - /// Blocking form of [`Self::flush_saves_then`], for tests and eviction - /// handoff on synchronous threads. Bounded by the same [`FLUSH_DEADLINE`] - /// chain. pub fn flush_saves(&self) { assert_outside_runtime("flush_saves"); let (tx, rx) = oneshot::channel(); self.flush_saves_then(move || { let _ = tx.send(()); }); - // block_on (not a bare channel wait) keeps the call-from-a-runtime - // misuse a loud tokio panic instead of a silently deadlocked worker. let _ = self.host.runtime.block_on(rx); } - /// Barrier the save pipeline, then call `then` — without blocking the - /// caller. Once `then` runs, a following [`Self::query`] (local or from a - /// P2P peer) observes every block saved before this call: this is the P/D - /// KV-ready signal, where the prefill node withholds a request's - /// `Finished` event until its KV is peer-visible. - /// - /// The barrier first awaits every in-flight fire-and-forget [`Self::save`] - /// (their D2H copy + write-pipeline submit), then drains the write - /// pipeline, then waits for the queued MetaServer registrations to be - /// delivered (or dropped after a failed attempt — registration stays - /// best-effort, the barrier only bounds *when* delivery is attempted, - /// never *whether* it succeeds; a peer that misses a registration - /// degrades to recompute). Without P2P the last step is a no-op. - /// - /// Barriers chain: each first awaits the previous barrier's completion, - /// so — as long as no barrier in the chain hit its deadline — it - /// transitively covers every save submitted before its own call, - /// including handles an earlier barrier drained whose D2H had not yet - /// submitted into the write pipeline (e.g. a chunked prefill's early - /// chunks flushed by another request's finish). Without the chain, the - /// pipeline drain cannot see such saves and the barrier would falsely - /// report them visible. A predecessor that timed out may leave its - /// drained handles permanently uncovered; that is the same accepted - /// degradation as the deadline itself — peers recompute. - /// - /// Each barrier is capped at [`FLUSH_DEADLINE`] (the wait on the - /// predecessor counts against it, and the predecessor is itself capped, - /// so delays never accumulate): a stalled MetaServer connection degrades - /// to "registrations still in flight" — semantically the same as a - /// dropped registration — and `then` still runs. + /// Wait for all saves submitted before this call and the server visibility + /// barrier, then invoke `then`. The deadline prevents a failed server from + /// withholding a finished request indefinitely. pub fn flush_saves_then(&self, then: impl FnOnce() + Send + 'static) { let (done_tx, done_rx) = oneshot::channel(); let (handles, prev_done) = { @@ -827,95 +471,102 @@ impl OffloadEngine { let prev_done = barrier.prev_flush_done.replace(done_rx); (handles, prev_done) }; - let engine = Arc::clone(&self.host.engine); + let client = self.host.client.clone(); self.host.runtime.spawn(async move { let flushed = tokio::time::timeout(FLUSH_DEADLINE, async { if let Some(prev) = prev_done { - // A cancelled predecessor (runtime teardown) resolves as - // an error immediately — don't let it stall the chain. let _ = prev.await; } for handle in handles { let _ = handle.await; } - engine.flush_saves_and_registrations().await; + if let Err(err) = client.flush().await { + log::warn!("PegaFlow flush failed: {err}"); + } }) .await; if flushed.is_err() { log::warn!( - "KV offload flush timed out after {FLUSH_DEADLINE:?}; \ - saves/registrations still in flight (peers may recompute)" + "KV offload flush timed out after {FLUSH_DEADLINE:?}; peers may recompute" ); } let _ = done_tx.send(()); then(); }); } + + /// Flush submitted saves, then unregister — the server frees the arena. + /// The caller must have closed its imported mapping (dropped the + /// `KvBuffer`/`ImportedKvArena`) first: CUDA IPC does not defend an open + /// importer mapping against the owner freeing the allocation. + pub fn shutdown(&mut self) { + if self.session.is_none() { + return; + } + assert_outside_runtime("shutdown"); + self.flush_saves(); + // Retire the dead-server watchdog BEFORE unregister: the server may + // close the session stream while (or right after) it handles the + // unregister, and the watchdog aborts the process on stream loss — + // it cannot tell a clean teardown from a crash. From here on, + // unregister's own error handling is the only liveness check needed. + self.session.take(); + if let Err(err) = self + .host + .runtime + .block_on(self.host.client.unregister(&self.instance_id)) + { + // The arena's fate is unknown; the server may free it while this + // process still believes the KV bytes exist. Nothing here can + // continue safely. + log::error!( + "PegaFlow unregister failed for {}: {err}; exiting", + self.instance_id + ); + std::process::abort(); + } + } +} + +impl Drop for OffloadEngine { + fn drop(&mut self) { + self.shutdown(); + } } #[cfg(test)] mod tests { use super::*; - /// The GLM5.2 shape: two arenas per model layer (MLA latent + index-K - /// sidecar) with different copy units, sharing pool block ids. Pins the - /// per-arena mapping and the exact strided-reach size bound pegaflow - /// validates copies against. #[test] - fn arena_registration_geometry() { - const MLA: usize = 656 * 64; - const IDXK: usize = 132 * 64; - let arenas = [ - KvArena { - name: "0.mla".into(), - base_ptr: 0x1000, - num_blocks: 10, - bytes_per_block: MLA, - block_stride_bytes: MLA, - }, - KvArena { - name: "0.idxk".into(), - base_ptr: 0x9000, - num_blocks: 10, - bytes_per_block: IDXK, - block_stride_bytes: IDXK, - }, - ]; - let reg = Registration::from_arenas(&arenas); - assert_eq!(reg.layer_names, ["0.mla", "0.idxk"]); - assert_eq!(reg.data_ptrs, [0x1000, 0x9000]); + fn layout_registration_geometry() { + // 2 layers, page-first: layer stride 512 elems, page stride 1024 elems. + let layout = KvLayout::new(2, 8, 16, 2); + let reg = Registration::from_layout(&layout, 10); + let layer_bytes = layout.layer_stride * ELEM_SIZE; + let page_bytes = layout.page_stride * ELEM_SIZE; + assert_eq!(reg.layer_names, ["0", "1"]); + assert_eq!(reg.offset_bytes, [0, layer_bytes as u64]); assert_eq!(reg.segments, [1, 1]); assert_eq!(reg.kv_stride_bytes, [0, 0]); assert_eq!(reg.num_blocks, [10, 10]); - assert_eq!(reg.bytes_per_block, [MLA, IDXK]); - assert_eq!(reg.block_stride_bytes, [MLA, IDXK]); - assert_eq!(reg.size_bytes, [10 * MLA, 10 * IDXK]); - } - - /// A page-interleaved arena (the qwen3 fused layout expressed as arenas): - /// stride exceeds the copy unit, and the size bound is the reach of the - /// last block, not `num_blocks * stride`. - #[test] - fn interleaved_arena_size_is_last_block_reach() { - let reg = Registration::from_arenas(&[KvArena { - name: "3".into(), - base_ptr: 0x100, - num_blocks: 4, - bytes_per_block: 512, - block_stride_bytes: 4096, - }]); - assert_eq!(reg.size_bytes, [3 * 4096 + 512]); + assert_eq!(reg.bytes_per_block, [layer_bytes, layer_bytes]); + assert_eq!(reg.block_stride_bytes, [page_bytes, page_bytes]); + assert_eq!(reg.alloc_size, 10 * page_bytes); + // A layer's reach is the last block's copy unit, not the full arena. + assert_eq!(reg.size_bytes, [9 * page_bytes + layer_bytes; 2]); + // Every view fits the arena. + for (offset, size) in reg.offset_bytes.iter().zip(®.size_bytes) { + assert!(*offset as usize + size <= reg.alloc_size); + } } #[test] - #[should_panic(expected = "overruns its block stride")] - fn arena_copy_unit_must_fit_its_stride() { - let _ = Registration::from_arenas(&[KvArena { - name: "bad".into(), - base_ptr: 0, - num_blocks: 1, - bytes_per_block: 4096, - block_stride_bytes: 512, - }]); + fn instance_suffix_is_process_independent_uuid() { + let a = new_instance_suffix(); + let b = new_instance_suffix(); + assert_eq!(a.len(), 32); + assert!(a.chars().all(|ch| ch.is_ascii_hexdigit())); + assert_ne!(a, b); } } diff --git a/openinfer-kv-offload/src/external.rs b/openinfer-kv-offload/src/external.rs new file mode 100644 index 000000000..321fedcf9 --- /dev/null +++ b/openinfer-kv-offload/src/external.rs @@ -0,0 +1,486 @@ +use pegaflow_core::EngineError; +use pegaflow_core::LayerSave; +use pegaflow_proto::proto::engine::FlushRequest; +use pegaflow_proto::proto::engine::HealthRequest; +use pegaflow_proto::proto::engine::LeaseLoad; +use pegaflow_proto::proto::engine::LoadRequest; +use pegaflow_proto::proto::engine::NativeKvTensor; +use pegaflow_proto::proto::engine::QueryRequest; +use pegaflow_proto::proto::engine::RegisterContextRequest; +use pegaflow_proto::proto::engine::ReleaseRequest; +use pegaflow_proto::proto::engine::ResponseStatus; +use pegaflow_proto::proto::engine::SaveLayer; +use pegaflow_proto::proto::engine::SaveRequest; +use pegaflow_proto::proto::engine::SessionRequest; +use pegaflow_proto::proto::engine::TransferMode; +use pegaflow_proto::proto::engine::UnregisterRequest; +use pegaflow_proto::proto::engine::engine_client::EngineClient; +use pegaflow_proto::proto::engine::query_response; +use tonic::Request; +use tonic::Streaming; +use tonic::transport::Channel; +use tonic::transport::Endpoint; + +const MAX_GRPC_MESSAGE_SIZE: usize = 64 * 1024 * 1024; +const CONNECT_DEADLINE: std::time::Duration = std::time::Duration::from_secs(5); +const RPC_DEADLINE: std::time::Duration = std::time::Duration::from_secs(30); +/// Save/load move data through the server's GPU/host pipelines; under load they +/// legitimately take much longer than control RPCs, and hitting this deadline +/// is fatal (see `abort_server_timeout`). Keep it far above worst-case D2H/H2D. +const DATA_DEADLINE: std::time::Duration = std::time::Duration::from_secs(120); +/// Byte length of a serialized `CUipcMemHandle`. +const IPC_HANDLE_BYTES: usize = 64; + +pub(super) struct ExternalRegistration<'a> { + pub instance_id: &'a str, + pub namespace: &'a str, + pub device_id: i32, + pub tp_rank: usize, + pub pp_rank: usize, + pub tp_size: usize, + pub world_size: usize, + pub layer_names: &'a [String], + /// Per-layer view offsets into the server-allocated arena. + pub offset_bytes: &'a [u64], + pub size_bytes: &'a [usize], + pub num_blocks: &'a [usize], + pub bytes_per_block: &'a [usize], + pub kv_stride_bytes: &'a [usize], + pub block_stride_bytes: &'a [usize], + pub segments: &'a [usize], + pub page_first: bool, + /// Total arena size the server allocates; every layer view fits inside. + pub alloc_size: usize, +} + +pub(super) enum ExternalQuery { + Loading, + Ready { num_blocks: usize, lease: Vec }, +} + +#[derive(Clone)] +pub(super) struct ExternalClient { + client: EngineClient, +} + +pub(super) struct ExternalSession { + task: tokio::task::JoinHandle<()>, +} + +impl Drop for ExternalSession { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl ExternalClient { + pub(super) async fn connect(server_addr: &str) -> Result { + let endpoint = Endpoint::from_shared(server_addr.to_string()).map_err(|err| { + EngineError::InvalidArgument(format!( + "invalid external PegaFlow server address {server_addr:?}: {err}" + )) + })?; + let channel = endpoint + .connect_timeout(CONNECT_DEADLINE) + .connect() + .await + .map_err(|err| { + EngineError::Storage(format!( + "connect external PegaFlow server {server_addr}: {err}" + )) + })?; + let mut client = EngineClient::new(channel) + .max_decoding_message_size(MAX_GRPC_MESSAGE_SIZE) + .max_encoding_message_size(MAX_GRPC_MESSAGE_SIZE); + let response = client + .health(deadline_request(HealthRequest {})) + .await + .map_err(|err| rpc_error("health", &err))? + .into_inner(); + require_ok("health", response.status)?; + Ok(Self { client }) + } + + /// Register the KV layout. The server allocates the arena and returns its + /// CUDA IPC handle. + pub(super) async fn register( + &self, + registration: ExternalRegistration<'_>, + ) -> Result<(ExternalSession, Vec), EngineError> { + let native_kv_tensors = build_native_kv_tensors( + registration.offset_bytes, + registration.size_bytes, + registration.block_stride_bytes, + )?; + let tp_rank = as_u32(registration.tp_rank, "tp_rank")?; + let pp_rank = as_u32(registration.pp_rank, "pp_rank")?; + let tp_size = as_u32(registration.tp_size, "tp_size")?; + let world_size = as_u32(registration.world_size, "world_size")?; + + let mut session_client = self.client.clone(); + let session = tokio::time::timeout( + RPC_DEADLINE, + session_client.session(SessionRequest { + instance_id: registration.instance_id.to_string(), + namespace: registration.namespace.to_string(), + tp_size, + world_size, + }), + ) + .await + .map_err(|_| EngineError::Storage("external PegaFlow session RPC timed out".into()))? + .map_err(|err| rpc_error("session", &err))?; + let session = ExternalSession { + task: tokio::spawn(watch_session(session.into_inner())), + }; + + let request = RegisterContextRequest { + instance_id: registration.instance_id.to_string(), + namespace: registration.namespace.to_string(), + tp_rank, + tp_size, + world_size, + device_id: registration.device_id, + layer_names: registration.layer_names.to_vec(), + wrapper_bytes: Vec::new(), + num_blocks: registration + .num_blocks + .iter() + .map(|&value| as_u64(value, "num_blocks")) + .collect::>()?, + bytes_per_block: registration + .bytes_per_block + .iter() + .map(|&value| as_u64(value, "bytes_per_block")) + .collect::>()?, + kv_stride_bytes: registration + .kv_stride_bytes + .iter() + .map(|&value| as_u64(value, "kv_stride_bytes")) + .collect::>()?, + segments: registration + .segments + .iter() + .map(|&value| as_u32(value, "segments")) + .collect::>()?, + pp_rank, + client_version: pegaflow_proto::VERSION.to_string(), + transfer_mode: TransferMode::Direct as i32, + page_first: registration.page_first, + native_kv_tensors, + native_alloc_size: as_u64(registration.alloc_size, "alloc_size")?, + }; + let mut client = self.client.clone(); + let response = match tokio::time::timeout( + RPC_DEADLINE, + client.register_context_batch(Request::new(request)), + ) + .await + { + Ok(Ok(response)) => response.into_inner(), + Ok(Err(err)) => return Err(rpc_error("register_context_batch", &err)), + Err(_) => { + return Err(EngineError::Storage( + "external PegaFlow register_context_batch timed out".into(), + )); + } + }; + require_ok("register_context_batch", response.status)?; + if response.arena_ipc_handle.len() != IPC_HANDLE_BYTES { + return Err(EngineError::Storage(format!( + "register_context_batch returned a {}-byte arena handle, expected \ + {IPC_HANDLE_BYTES} (is the server running the native-arena build?)", + response.arena_ipc_handle.len() + ))); + } + Ok((session, response.arena_ipc_handle)) + } + + pub(super) async fn save( + &self, + instance_id: &str, + tp_rank: usize, + pp_rank: usize, + device_id: i32, + saves: Vec, + ) -> Result<(), EngineError> { + let saves = saves + .into_iter() + .map(|save| { + Ok(SaveLayer { + layer_name: save.layer_name, + block_ids: save + .block_ids + .into_iter() + .map(|id| as_u32(id, "block_id")) + .collect::>()?, + block_hashes: save.block_hashes, + }) + }) + .collect::>()?; + let mut client = self.client.clone(); + let response = match tokio::time::timeout( + DATA_DEADLINE, + client.save(Request::new(SaveRequest { + instance_id: instance_id.to_string(), + tp_rank: as_u32(tp_rank, "tp_rank")?, + device_id, + saves, + pp_rank: as_u32(pp_rank, "pp_rank")?, + })), + ) + .await + { + Ok(Ok(response)) => response.into_inner(), + Ok(Err(err)) => abort_lost_server("save", &err), + Err(_) => abort_server_timeout("save"), + }; + require_ok("save", response.status) + } + + pub(super) async fn query( + &self, + instance_id: &str, + req_id: &str, + block_hashes: &[Vec], + ) -> Result { + let mut client = self.client.clone(); + let response = client + .query_prefetch(deadline_request(QueryRequest { + instance_id: instance_id.to_string(), + block_hashes: block_hashes.to_vec(), + req_id: req_id.to_string(), + wait_for_full_prefix: false, + })) + .await + .map_err(|err| rpc_error("query_prefetch", &err))? + .into_inner(); + match response.outcome { + Some(query_response::Outcome::Loading(_)) => Ok(ExternalQuery::Loading), + Some(query_response::Outcome::Ready(ready)) => { + let num_blocks = usize::try_from(ready.num_hit_blocks).map_err(|_| { + EngineError::Storage(format!( + "query hit count {} does not fit usize", + ready.num_hit_blocks + )) + })?; + if num_blocks > block_hashes.len() { + return Err(EngineError::Storage(format!( + "query returned {num_blocks} blocks for {} requested hashes", + block_hashes.len() + ))); + } + if (num_blocks == 0) != ready.lease.is_empty() { + return Err(EngineError::Storage(format!( + "query returned inconsistent hit/lease: blocks={num_blocks}, lease_bytes={}", + ready.lease.len() + ))); + } + Ok(ExternalQuery::Ready { + num_blocks, + lease: ready.lease, + }) + } + None => Err(EngineError::Storage( + "query_prefetch response omitted outcome".to_string(), + )), + } + } + + pub(super) async fn load( + &self, + instance_id: &str, + tp_rank: usize, + device_id: i32, + layer_names: &[String], + lease: Vec, + block_ids: Vec, + ) -> Result<(), EngineError> { + let mut client = self.client.clone(); + let response = match tokio::time::timeout( + DATA_DEADLINE, + client.load(Request::new(LoadRequest { + instance_id: instance_id.to_string(), + tp_rank: as_u32(tp_rank, "tp_rank")?, + device_id, + load_state_shm: String::new(), + layer_names: layer_names.to_vec(), + loads: vec![LeaseLoad { + lease, + block_ids: block_ids + .into_iter() + .map(|id| as_u32(id, "block_id")) + .collect::>()?, + }], + wait_for_completion: true, + })), + ) + .await + { + Ok(Ok(response)) => response.into_inner(), + Ok(Err(err)) => abort_lost_server("load", &err), + Err(_) => abort_server_timeout("load"), + }; + require_ok("load", response.status) + } + + pub(super) async fn release(&self, lease: Vec) -> Result<(), EngineError> { + let mut client = self.client.clone(); + client + .release(deadline_request(ReleaseRequest { lease })) + .await + .map_err(|err| rpc_error("release", &err))?; + Ok(()) + } + + /// Server-wide durability barrier: `Flush` has no instance scope, so on a + /// shared server this also waits out other instances' save tails. + pub(super) async fn flush(&self) -> Result<(), EngineError> { + let mut client = self.client.clone(); + let response = client + .flush(deadline_request(FlushRequest {})) + .await + .map_err(|err| rpc_error("flush", &err))? + .into_inner(); + require_ok("flush", response.status) + } + + /// Ensure the instance is gone server-side. "Not found" counts as + /// success: dropping the liveness stream just before this call may have + /// already triggered the server's session cleanup for the same instance. + pub(super) async fn unregister(&self, instance_id: &str) -> Result<(), EngineError> { + let mut client = self.client.clone(); + let response = match tokio::time::timeout( + RPC_DEADLINE, + client.unregister_context(Request::new(UnregisterRequest { + instance_id: instance_id.to_string(), + })), + ) + .await + { + Ok(Ok(response)) => response.into_inner(), + Ok(Err(err)) + if matches!( + err.code(), + tonic::Code::FailedPrecondition | tonic::Code::NotFound + ) => + { + log::debug!( + "unregister_context: instance {instance_id} already cleaned up ({err})" + ); + return Ok(()); + } + Ok(Err(err)) => return Err(rpc_error("unregister_context", &err)), + Err(_) => abort_server_timeout("unregister_context"), + }; + require_ok("unregister_context", response.status) + } +} + +/// The liveness stream doubles as a dead-server detector. The server owns the +/// KV arena this process has mapped; if the server goes away the mapping is +/// backed by freed memory, so the only safe reaction is to exit. There is no +/// reconnect. +async fn watch_session(mut stream: Streaming) { + loop { + match stream.message().await { + Ok(Some(_)) => {} + Ok(None) => { + log::error!( + "external PegaFlow session closed by server; the imported KV arena is gone, \ + exiting" + ); + std::process::abort(); + } + Err(err) => { + log::error!( + "external PegaFlow session failed: {err}; the imported KV arena is gone, \ + exiting" + ); + std::process::abort(); + } + } + } +} + +fn deadline_request(message: T) -> Request { + let mut request = Request::new(message); + request.set_timeout(RPC_DEADLINE); + request +} + +fn abort_lost_server(operation: &str, err: &tonic::Status) -> ! { + log::error!( + "external PegaFlow {operation} RPC failed mid-transfer: {err}; the server owns the KV \ + arena this process has mapped, exiting" + ); + std::process::abort(); +} + +fn abort_server_timeout(operation: &str) -> ! { + log::error!( + "external PegaFlow {operation} did not respond in time; the server owns this process's \ + KV arena, so its state is unknowable — exiting" + ); + std::process::abort(); +} + +/// Build per-layer strided views for the register RPC from explicit offsets +/// into the server-allocated arena. +fn build_native_kv_tensors( + offset_bytes: &[u64], + size_bytes: &[usize], + block_stride_bytes: &[usize], +) -> Result, EngineError> { + if offset_bytes.is_empty() + || offset_bytes.len() != size_bytes.len() + || offset_bytes.len() != block_stride_bytes.len() + { + return Err(EngineError::InvalidArgument(format!( + "native layer metadata length mismatch: offsets={}, sizes={}, block strides={}", + offset_bytes.len(), + size_bytes.len(), + block_stride_bytes.len() + ))); + } + offset_bytes + .iter() + .copied() + .zip(size_bytes.iter().copied()) + .zip(block_stride_bytes.iter().copied()) + .map(|((offset, view_size), block_stride)| { + Ok(NativeKvTensor { + offset_bytes: offset, + size_bytes: as_u64(view_size, "size_bytes")?, + block_stride_bytes: as_u64(block_stride, "block_stride_bytes")?, + }) + }) + .collect() +} + +fn require_ok(operation: &str, status: Option) -> Result<(), EngineError> { + let status = status + .ok_or_else(|| EngineError::Storage(format!("{operation} response omitted status")))?; + if status.ok { + Ok(()) + } else { + Err(EngineError::Storage(format!( + "{operation} failed: {}", + status.message + ))) + } +} + +fn rpc_error(operation: &str, err: &tonic::Status) -> EngineError { + EngineError::Storage(format!("external PegaFlow {operation} RPC: {err}")) +} + +fn as_u32(value: usize, field: &str) -> Result { + u32::try_from(value) + .map_err(|_| EngineError::InvalidArgument(format!("{field}={value} does not fit u32"))) +} + +fn as_u64(value: usize, field: &str) -> Result { + u64::try_from(value) + .map_err(|_| EngineError::InvalidArgument(format!("{field}={value} does not fit u64"))) +} diff --git a/openinfer-kv-offload/src/lib.rs b/openinfer-kv-offload/src/lib.rs index f916f8730..02d7ff4aa 100644 --- a/openinfer-kv-offload/src/lib.rs +++ b/openinfer-kv-offload/src/lib.rs @@ -1,9 +1,9 @@ -//! In-process KV cache offload bridge between openinfer and pegaflow. +//! KV cache offload client between OpenInfer and an external PegaFlow server. //! -//! openinfer owns the GPU paged-KV (`openinfer-kv-cache::KvBuffer`, page-first -//! layout) and the logical prefix cache (kvbm `BlockPool`). pegaflow owns the -//! deeper tiers (host pinned memory, SSD, RDMA). [`OffloadEngine`] is the -//! connector "brain" that moves blocks between them and decides when. +//! openinfer owns the logical prefix cache (kvbm `BlockPool`); pegaflow owns +//! the GPU KV arena it allocated for us plus the deeper tiers (host pinned +//! memory, SSD, RDMA). [`OffloadEngine`] registers the KV layout, imports the +//! server-allocated arena over CUDA IPC, and issues save/query/load RPCs. //! //! Dense-attention v1 (Qwen3-4B): the GPU prefix hit stays native to kvbm's //! `BlockPool`; this connector covers the CPU tier and stacks a CPU-hit prefix @@ -14,18 +14,17 @@ //! polls its [`LoadHandle`] each scheduler tick. mod engine; +mod external; mod vllm_hash; -pub use engine::HostConfig; pub use engine::KvArena; pub use engine::LoadHandle; pub use engine::OffloadConfig; pub use engine::OffloadEngine; pub use engine::OffloadHost; -pub use engine::P2pConfig; pub use engine::QueryHit; +pub use engine::QueryLeaseId; pub use engine::QueryOutcome; -// Re-exported so callers name pegaflow's engine types through this bridge. -pub use pegaflow_core::{EngineError, PegaEngine, QueryLeaseId}; +pub use pegaflow_core::EngineError; pub use vllm_hash::VLLM_HASH_BYTES; pub use vllm_hash::VllmBlockHasher; diff --git a/openinfer-kv-offload/tests/cpu_roundtrip.rs b/openinfer-kv-offload/tests/cpu_roundtrip.rs index f1ebce4f0..4062635d6 100644 --- a/openinfer-kv-offload/tests/cpu_roundtrip.rs +++ b/openinfer-kv-offload/tests/cpu_roundtrip.rs @@ -4,8 +4,8 @@ //! to pegaflow's host tier, evicts the GPU-side data implicitly by loading into //! a *different* set of blocks, and checks the bytes match. This exercises the //! whole connector — strided per-layer registration (`block_stride` ≠ copy -//! size), the K/V split, the async save, the prefix query, and the in-process -//! oneshot load — on actual device memory. If the layout math were wrong the +//! size), the K/V split, the async save, the prefix query, and the cross-process +//! load state — on actual device memory. If the layout math were wrong the //! loaded bytes would land in the wrong layer/segment/block and the compare //! would fail. //! @@ -14,7 +14,9 @@ use cudarc::driver::CudaContext; use cudarc::driver::result; use half::bf16; +use openinfer_kernels::imported::ImportedKvArena; use openinfer_kv_cache::KvBuffer; +use openinfer_kv_cache::KvLayout; use openinfer_kv_offload::OffloadConfig; use openinfer_kv_offload::OffloadEngine; use openinfer_kv_offload::QueryOutcome; @@ -58,19 +60,32 @@ fn block_hash(logical: usize) -> Vec { #[test] fn gpu_cpu_gpu_roundtrip_preserves_kv_bytes() { + let Ok(server_addr) = std::env::var("OPENINFER_PEGAFLOW_SERVER") else { + eprintln!("skipping cpu_roundtrip: set OPENINFER_PEGAFLOW_SERVER to run it"); + return; + }; let ctx = CudaContext::new(0).expect("cuda device 0"); ctx.bind_to_thread().expect("bind ctx to test thread"); let stream = ctx.default_stream(); - let buffer = KvBuffer::new( + // ── Register first: the server allocates the arena and returns its + // CUDA IPC handle; the buffer is a view over the imported mapping. ── + let layout = KvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HEAD_DIM, PAGE_SIZE); + let config = OffloadConfig::new("roundtrip-test", 0, server_addr); + let (mut engine, handle) = + OffloadEngine::new(&config, &layout, NUM_BLOCKS).expect("build OffloadEngine"); + let arena_bytes = NUM_BLOCKS * PAGE_STRIDE * std::mem::size_of::(); + let arena = ImportedKvArena::open(&stream, &handle, arena_bytes).expect("import arena"); + let buffer = KvBuffer::new_imported( &stream, NUM_LAYERS, NUM_KV_HEADS, HEAD_DIM, PAGE_SIZE, NUM_BLOCKS, + &arena, ) - .expect("alloc KvBuffer"); + .expect("wrap imported arena"); // Sanity: our test-local geometry constants match the buffer's layout. assert_eq!(buffer.layout().page_stride, PAGE_STRIDE); assert_eq!(buffer.layout().kv_block_len, SEGMENT_LEN); @@ -95,14 +110,6 @@ fn gpu_cpu_gpu_roundtrip_preserves_kv_bytes() { } stream.synchronize().expect("sync after fill"); - // ── Build the offload engine (registers the fused buffer) ── - let engine = OffloadEngine::new( - OffloadConfig::new("roundtrip-test", 0, 64 * 1024 * 1024), - &buffer, - &stream, - ) - .expect("build OffloadEngine"); - let hashes: Vec> = (0..src_blocks.len()).map(block_hash).collect(); let src_ids: Vec = src_blocks.iter().map(|&b| b as i32).collect(); @@ -124,7 +131,7 @@ fn gpu_cpu_gpu_roundtrip_preserves_kv_bytes() { // ── Load CPU→GPU into a *different* set of blocks ── let dst_ids: Vec = dst_blocks.iter().map(|&b| b as i32).collect(); engine - .load(lease, dst_ids) + .load(&lease, dst_ids) .expect("submit load") .wait() .expect("load completes"); @@ -158,4 +165,10 @@ fn gpu_cpu_gpu_roundtrip_preserves_kv_bytes() { zero.iter().all(|v| v.to_bits() == 0), "an unloaded block must remain zeroed — load must not scribble outside its destinations" ); + + // ── Teardown order mirrors the executor: close the mapping, then let + // unregister free the server-side arena. ── + drop(buffer); + drop(arena); + engine.shutdown(); } diff --git a/openinfer-qwen3/src/executor.rs b/openinfer-qwen3/src/executor.rs index 71dc97c82..c55e99a41 100644 --- a/openinfer-qwen3/src/executor.rs +++ b/openinfer-qwen3/src/executor.rs @@ -22,6 +22,7 @@ use openinfer_core::tensor::DeviceContext; use openinfer_core::tensor::HiddenStates; use openinfer_core::weight_loader::WeightPrefetch; use openinfer_core::weight_loader::load_shard_info; +use openinfer_kernels::imported::ImportedKvArena; use openinfer_kv_cache::KvBlockGuard; use openinfer_kv_cache::KvBuffer; use openinfer_kv_cache::KvCacheEvent; @@ -997,6 +998,10 @@ pub struct Qwen3Executor { /// pegaflow KV-offload bridge; `None` unless offload is opted in on the /// single-GPU path. Drives both the SAVE hook and the async LOAD prefetch. offload: Option, + /// Mapping of the server-owned KV arena backing `kv_mgr`. Closed in + /// [`Drop`] after the workers stop and before `offload` unregisters (the + /// unregister lets the server free the allocation). + kv_arena: Option, /// Per-request count of sealed blocks already saved to the host tier, so /// each step only saves blocks that newly sealed. Initialized to the /// GPU-hit prefix (already resident) on first save. @@ -1014,7 +1019,7 @@ pub struct Qwen3Executor { l1_retention_disabled: bool, /// P/D prefill role: withhold each step's `Finished` events until offload /// saves + MetaServer registrations are peer-visible, so the HTTP response - /// doubles as the KV-ready signal (see `Qwen3P2pOptions::flush_on_finish`). + /// doubles as the KV-ready signal. flush_offload_on_finish: bool, /// P/D decode role with a vLLM prefill peer: offload query keys derive /// with vLLM's hash scheme, a zero hit waits out the producer's @@ -1144,6 +1149,20 @@ impl Qwen3Executor { dflash_kv_bytes_per_token, memory_options, )?; + // Pure config validation goes before registration: past this point a + // failure means unwinding a server-side arena allocation. + if let Some(c) = offload_opts.vllm_compat.as_ref() { + ensure!( + c.miss_wait < REMOTE_FETCH_DEADLINE, + "kv-pd miss wait ({:?}) must stay below the {:?} remote-fetch \ + deadline, which would otherwise silently cap it", + c.miss_wait, + REMOTE_FETCH_DEADLINE, + ); + } + // With offload on, register with PegaFlow before any KV allocation: + // the server owns the arena and hands back its CUDA IPC handle. + let mut offload_setup = None; let (kv_mgr, kv_events) = if enable_kv_events { let (kv_mgr, rx) = KvCacheManager::new_with_events( &model.device_ctx().stream, @@ -1161,15 +1180,10 @@ impl Qwen3Executor { }), ) } else { - let kv_mgr = KvCacheManager::new( - &model.device_ctx().stream, - budget.num_layers, - budget.num_kv_heads, - budget.head_dim, - budget.block_size, - budget.num_blocks, - )?; - (kv_mgr, None) + ( + Self::build_kv_manager(&model, &budget, offload_opts, &mut offload_setup)?, + None, + ) }; let device_ordinal = model.device_ctx().device_ordinal; let metadata = Qwen3ExecutorMetadata { @@ -1178,21 +1192,15 @@ impl Qwen3Executor { config: model.config().clone(), }; let kv_buffer = kv_mgr.buffer().clone(); - // Build the offload engine while the model's stream is still in hand - // (it moves into the RankWorker below). Registers the fused KV buffer. - let offload = build_offload(offload_opts, &kv_mgr, model.config(), model.device_ctx())?; + let (offload, kv_arena) = match offload_setup { + Some((arena, engine)) => (Some(engine), Some(arena)), + None => (None, None), + }; let total_blocks = kv_mgr.pool().total_blocks(); let padding_block_id = kv_mgr.pool().padding_block_id(); let vllm_compat = match offload_opts.vllm_compat.as_ref() { None => None, Some(c) => { - ensure!( - c.miss_wait < REMOTE_FETCH_DEADLINE, - "kv-pd miss wait ({:?}) must stay below the {:?} remote-fetch \ - deadline, which would otherwise silently cap it", - c.miss_wait, - REMOTE_FETCH_DEADLINE, - ); let hasher = openinfer_kv_offload::VllmBlockHasher::new( &c.python_hash_seed, budget.block_size, @@ -1239,15 +1247,13 @@ impl Qwen3Executor { prefix_cache_enabled: true, lora_options: Qwen3LoraOptions::default(), offload, + kv_arena, saved_cursor: HashMap::new(), prefetch: HashMap::new(), l1_retention_disabled: false, // Derived here, not via a post-construction setter, so every // launch path (plain and LoRA alike) honors the P/D contract. - flush_offload_on_finish: offload_opts - .p2p - .as_ref() - .is_some_and(|p2p| p2p.flush_on_finish), + flush_offload_on_finish: offload_opts.flush_on_finish, vllm_compat, overlap: None, async_prefill: None, @@ -1298,7 +1304,7 @@ impl Qwen3Executor { "Qwen3 executor requires at least one device" ); anyhow::ensure!( - !offload_options.enabled || device_ordinals.len() == 1, + !offload_options.is_enabled() || device_ordinals.len() == 1, "KV offload is only supported on the single-GPU path (tensor parallel \ shards KV per rank); got {} devices", device_ordinals.len() @@ -1316,7 +1322,7 @@ impl Qwen3Executor { // hash, which the cursor + lineage→seq map do not model — so the two are // mutually exclusive by construction rather than silently mis-announced. anyhow::ensure!( - !enable_kv_events || !offload_options.enabled, + !enable_kv_events || !offload_options.is_enabled(), "KV block events and KV offload are mutually exclusive (the event cursor \ assumes GPU-resident block reuse)" ); @@ -1568,6 +1574,7 @@ impl Qwen3Executor { lora_options, // Offload is single-GPU only (asserted above); never built here. offload: None, + kv_arena: None, saved_cursor: HashMap::new(), prefetch: HashMap::new(), l1_retention_disabled: false, @@ -2061,7 +2068,7 @@ impl Qwen3Executor { return Err(()); }; let page_ids = reservation.page_ids(); - match offload.load(lease, page_ids) { + match offload.load(&lease, page_ids) { Ok(handle) => { self.prefetch.insert( id, @@ -2154,24 +2161,65 @@ fn profile_kv_budget_on_worker( /// is disabled. Registers the fused KV buffer with pegaflow against the model's /// device/stream — must be called while that stream is still owned by the model /// (before it moves into the `RankWorker`). +impl Qwen3Executor { + /// KV manager for the no-events path. With offload enabled this registers + /// with PegaFlow first (the server allocates the arena), imports the + /// returned CUDA IPC handle, and builds the manager over that mapping. + fn build_kv_manager( + model: &Qwen3Model, + budget: &crate::weights::KvBudget, + offload_opts: &Qwen3OffloadOptions, + // Arena before engine: if construction bails after this is populated, + // the tuple's field-order drop closes the IPC mapping before the + // engine's Drop unregisters and the server frees the allocation. + offload_setup: &mut Option<(ImportedKvArena, OffloadEngine)>, + ) -> Result { + if !offload_opts.is_enabled() { + return KvCacheManager::new( + &model.device_ctx().stream, + budget.num_layers, + budget.num_kv_heads, + budget.head_dim, + budget.block_size, + budget.num_blocks, + ); + } + let layout = openinfer_kv_cache::KvLayout::new( + budget.num_layers, + budget.num_kv_heads, + budget.head_dim, + budget.block_size, + ); + let setup = build_offload(offload_opts, &layout, budget.num_blocks, model.device_ctx())? + .expect("offload enabled implies a server address"); + let kv_mgr = KvCacheManager::new_imported( + &model.device_ctx().stream, + budget.num_layers, + budget.num_kv_heads, + budget.head_dim, + budget.block_size, + budget.num_blocks, + &setup.0, + )?; + *offload_setup = Some(setup); + Ok(kv_mgr) + } +} + fn build_offload( opts: &Qwen3OffloadOptions, - kv_mgr: &KvCacheManager, - config: &Config, + layout: &openinfer_kv_cache::KvLayout, + num_blocks: usize, ctx: &DeviceContext, -) -> Result> { - if !opts.enabled { +) -> Result> { + let Some(server_addr) = &opts.server_addr else { return Ok(None); - } + }; let device_id = ctx.device_ordinal as i32; - let layout = kv_mgr.buffer().layout(); // Content-addressing domain: two engines may cross-hit only when the same - // token prefix produces interchangeable KV bytes. That needs the *model* - // to match, not just the KV geometry — Qwen3-4B and 8B share - // layers/heads/head_dim and a tokenizer, so geometry alone would let a - // mixed mesh silently feed one model the other's KV. hidden_size + - // intermediate_size + vocab_size discriminate the model line's sizes; - // the layout fields pin the block geometry the transfer relies on. + // token prefix produces interchangeable KV bytes. The deployment identity + // distinguishes checkpoints with identical shapes; layout facts prevent a + // reused identity from crossing incompatible cache geometry. // vLLM-compat mode joins the *P side's* content domain instead: the // pegaflow connector derives an 8-hex namespace from vLLM config (and logs // it at startup); reproducing that derivation would mean chasing Python @@ -2179,38 +2227,27 @@ fn build_offload( let namespace = match &opts.vllm_compat { Some(compat) => compat.namespace.clone(), None => format!( - "openinfer-qwen3-hs{}-is{}-v{}-l{}h{}d{}p{}", - config.hidden_size, - config.intermediate_size, - config.vocab_size, + "openinfer-qwen3-{}-l{}h{}d{}p{}", + opts.namespace + .as_deref() + .context("Qwen3 native KV offload requires a checkpoint namespace")?, layout.num_layers, layout.num_kv_heads, layout.head_dim, layout.page_size ), }; - let mut config = OffloadConfig::new( - format!("qwen3-dev{device_id}"), - device_id, - opts.pinned_pool_bytes, - ) - .with_namespace(namespace); - config.use_hugepages = opts.use_hugepages; - if let Some(p2p) = &opts.p2p { - config = config.with_p2p(openinfer_kv_offload::P2pConfig { - metaserver_addr: p2p.metaserver_addr.clone(), - advertise_addr: p2p.advertise_addr.clone(), - rdma_nics: p2p.rdma_nics.clone(), - }); - } - let engine = OffloadEngine::new(config, kv_mgr.buffer(), &ctx.stream) + let config = OffloadConfig::new(format!("qwen3-dev{device_id}"), device_id, server_addr) + .with_namespace(namespace); + // Register first: the server allocates the fused arena and returns its + // CUDA IPC handle; the KV buffer is then built over the imported mapping. + let (engine, arena_ipc_handle) = OffloadEngine::new(&config, layout, num_blocks) .map_err(|e| anyhow::anyhow!("KV offload engine init failed: {e}"))?; - log::info!( - "KV offload enabled on device {device_id} ({} MiB host tier, p2p={})", - opts.pinned_pool_bytes >> 20, - opts.p2p.is_some(), - ); - Ok(Some(engine)) + let arena_bytes = num_blocks * layout.page_stride * std::mem::size_of::(); + let arena = ImportedKvArena::open(&ctx.stream, &arena_ipc_handle, arena_bytes) + .map_err(|e| anyhow::anyhow!("import PegaFlow KV arena: {e}"))?; + log::info!("KV offload enabled on device {device_id}: server={server_addr}"); + Ok(Some((arena, engine))) } fn ensure_lora_capacity( @@ -3125,10 +3162,27 @@ mod tests { impl Drop for Qwen3Executor { fn drop(&mut self) { + // Ordering: stop every kernel producer first, wait out in-flight + // host→GPU loads (their DMA targets the arena), then close our + // mapping of the server-owned arena, and only then unregister — + // which lets the server free the allocation. self.primary.shutdown(); for worker in &mut self.workers { worker.shutdown(); } + if self.offload.is_some() { + for (_, state) in self.prefetch.drain() { + if let PrefetchPhase::Loading { handle, .. } = state.phase { + if let Err(err) = handle.wait() { + log::warn!("KV offload load abandoned during shutdown: {err}"); + } + } + } + } + self.kv_arena.take(); + if let Some(mut offload) = self.offload.take() { + offload.shutdown(); + } } } diff --git a/openinfer-qwen3/src/lib.rs b/openinfer-qwen3/src/lib.rs index 3d5691dcd..30d5917ca 100644 --- a/openinfer-qwen3/src/lib.rs +++ b/openinfer-qwen3/src/lib.rs @@ -82,7 +82,7 @@ impl Default for Qwen3LoraOptions { /// stream plumbing in [`green_ctx`]. pub use green_ctx::DecodeOverlap; -/// KV-offload (pegaflow) opt-in for the single-GPU Qwen3 path. +/// External PegaFlow KV-offload opt-in for the single-GPU Qwen3 path. /// /// Disabled by default — the existing GPU-only prefix cache is unchanged. /// When enabled, the executor saves sealed KV blocks to pegaflow's host tier @@ -91,40 +91,21 @@ pub use green_ctx::DecodeOverlap; /// single-GPU topology is supported (tensor parallel shards KV per rank). #[derive(Clone, Debug, Eq, PartialEq)] pub struct Qwen3OffloadOptions { - enabled: bool, - /// Host pinned-memory pool size (the CPU KV-tier capacity), in bytes. - pinned_pool_bytes: usize, - /// Back the pool with 2 MiB hugepages (the box must hold a reservation). - pub use_hugepages: bool, - /// `Some` joins the cross-instance P2P mesh: block hashes register with a - /// MetaServer, peers pull missing prefixes over RDMA, and this engine - /// serves theirs. The P/D disaggregation data plane. - p2p: Option, + /// Out-of-process PegaFlow gRPC endpoint. `None` disables offload. + pub server_addr: Option, + /// Stable deployment/checkpoint identity shared by compatible native + /// producers and consumers. Required unless vLLM compatibility supplies + /// the connector namespace. + pub namespace: Option, + /// Barrier saves before emitting `Finished`. Prefill deployments use this + /// to make the HTTP response a KV-ready signal. + pub flush_on_finish: bool, /// `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. vllm_compat: Option, } -/// Cross-instance P2P KV sharing (see `openinfer_kv_offload::P2pConfig`). -#[derive(Clone, Debug, Eq, PartialEq)] -pub struct Qwen3P2pOptions { - /// MetaServer gRPC address, e.g. `http://127.0.0.1:50056`. - pub metaserver_addr: String, - /// This engine's routable `IP:port` (literal socket address; also the - /// P2P gRPC listen address, so hostnames are rejected at startup). Peers - /// dial it for RDMA handshakes and block queries. - pub advertise_addr: String, - /// RDMA NIC device names (e.g. `mlx5_0`). - pub rdma_nics: Vec, - /// Barrier a request's KV saves (host tier + MetaServer registration) - /// before its `Finished` event is emitted. The prefill role in a P/D - /// deployment turns this on so its HTTP response *is* the KV-ready signal; - /// costs one write-pipeline + registration drain per finishing step, so - /// leave it off on decode/serving instances. - pub flush_on_finish: bool, -} - /// Decode-node settings for a P/D deployment whose prefill node is vLLM with /// the pegaflow connector. vLLM registers KV under its own prefix-cache block /// hashes (`xxh3_128` over canonical-CBOR chained tuples — see @@ -153,30 +134,38 @@ pub struct Qwen3VllmCompatOptions { impl Qwen3OffloadOptions { pub fn disabled() -> Self { Self { - enabled: false, - pinned_pool_bytes: 0, - use_hugepages: false, - p2p: None, + server_addr: None, + namespace: None, + flush_on_finish: false, vllm_compat: None, } } - pub fn enabled(pinned_pool_bytes: usize) -> Self { + pub fn external(server_addr: impl Into) -> Self { Self { - enabled: true, - pinned_pool_bytes, - use_hugepages: false, - p2p: None, + server_addr: Some(server_addr.into()), + namespace: None, + flush_on_finish: false, vllm_compat: None, } } #[must_use] - pub fn with_p2p(mut self, p2p: Qwen3P2pOptions) -> Self { - self.p2p = Some(p2p); + pub fn with_flush_on_finish(mut self, flush_on_finish: bool) -> Self { + self.flush_on_finish = flush_on_finish; self } + #[must_use] + pub fn with_namespace(mut self, namespace: impl Into) -> Self { + self.namespace = Some(namespace.into()); + self + } + + pub fn is_enabled(&self) -> bool { + self.server_addr.is_some() + } + #[must_use] pub fn with_vllm_compat(mut self, compat: Qwen3VllmCompatOptions) -> Self { self.vllm_compat = Some(compat); @@ -289,10 +278,10 @@ pub fn launch(model_path: &Path, options: Qwen3LaunchOptions) -> Result /proc/sys/vm/nr_hugepages` as root) — allocation fails at - /// startup otherwise. - #[arg(long, default_value_t = false, requires = "kv_offload")] - pub kv_offload_hugepages: bool, - - /// Join the cross-instance KV P2P mesh: pegaflow MetaServer gRPC address - /// (e.g. `http://127.0.0.1:50056`). Saved block hashes register there and - /// missing prefixes are pulled from peer instances over RDMA — the P/D - /// disaggregation data plane. Requires --kv-offload, --kv-p2p-advertise-addr - /// and --kv-p2p-nics. - #[arg(long, requires_all = ["kv_offload", "kv_p2p_advertise_addr", "kv_p2p_nics"])] - pub kv_p2p_metaserver_addr: Option, - - /// This instance's routable IP:port for KV P2P — a literal socket address - /// (it is also the embedded transfer-service bind address, so hostnames - /// are rejected at startup). Peers dial it for RDMA handshakes and block - /// queries. Must be reachable by every peer; not 0.0.0.0. - #[arg(long, requires = "kv_p2p_metaserver_addr")] - pub kv_p2p_advertise_addr: Option, - - /// RDMA NIC device names for KV P2P (e.g. `mlx5_0`), comma-separated. - #[arg(long, value_delimiter = ',', requires = "kv_p2p_metaserver_addr")] - pub kv_p2p_nics: Vec, - - /// P/D prefill role: barrier each request's KV saves (host tier + - /// MetaServer registration) before its final token event, so this + /// Connect to an external PegaFlow gRPC server and enable KV offload. The + /// server owns host memory, hugepages, SSD, RDMA, and peer discovery. + #[arg(long)] + pub kv_offload_server: Option, + + /// Stable checkpoint/deployment identity for native OpenInfer KV sharing. + /// Every compatible producer and consumer must use the same value. + #[arg( + long, + requires = "kv_offload_server", + conflicts_with = "kv_pd_vllm_seed" + )] + pub kv_offload_namespace: Option, + + /// P/D prefill role: barrier each request's KV saves before its final + /// token event, so this /// instance's HTTP response doubles as the KV-ready signal a router can /// act on. Leave off on decode instances. - #[arg(long, default_value_t = false, requires = "kv_p2p_metaserver_addr")] - pub kv_p2p_flush_on_finish: bool, + #[arg(long, default_value_t = false, requires = "kv_offload_server")] + pub kv_offload_flush_on_finish: bool, /// P/D decode role with a vLLM prefill peer: the shared PYTHONHASHSEED /// value set on every vLLM prefill process. Switches offload query keys to /// vLLM's prefix-cache hash scheme (requires the P side to run /// --prefix-caching-hash-algo xxhash_cbor) and makes a cold request wait /// out the producer's registration tail instead of prefilling locally. - /// Requires --kv-pd-vllm-namespace and the P2P mesh flags. - #[arg(long, value_parser = parse_pythonhashseed, requires_all = ["kv_p2p_metaserver_addr", "kv_pd_vllm_namespace"])] + /// Requires --kv-offload-server and --kv-pd-vllm-namespace. + #[arg(long, value_parser = parse_pythonhashseed, requires_all = ["kv_offload_server", "kv_pd_vllm_namespace"])] pub kv_pd_vllm_seed: Option, /// The vLLM prefill peer's pegaflow-connector namespace (an 8-hex digest @@ -164,9 +138,9 @@ pub(crate) struct Args { #[arg(long, default_value_t = false, requires = "kv_pd_vllm_seed")] pub kv_pd_allow_local_prefill: bool, - /// vLLM-style no-prefix-cache. Without --kv-offload it disables prefix + /// vLLM-style no-prefix-cache. Without --kv-offload-server it disables prefix /// matching outright (every prefill recomputes the full prompt). With - /// --kv-offload it is the pure-L2 mode: no cross-request HBM reuse, so every + /// --kv-offload-server it is the pure-L2 mode: no cross-request HBM reuse, so every /// prefix is restored from the host tier — for measuring the L2 TTFT win. #[arg(long, default_value_t = false)] pub no_prefix_cache: bool, @@ -174,7 +148,7 @@ pub(crate) struct Args { /// Speculative drafter model path: Qwen3 DFlash/DSpark decoding, or the /// GLM5.2 DSpark drafter (greedy AND sampled requests speculate; /// per-request accept stats logged). For Qwen3: single-GPU greedy only; - /// incompatible with --enable-lora and --kv-offload, and forces the + /// incompatible with --enable-lora and --kv-offload-server, and forces the /// prefix cache off (it needs clean target hidden states). #[arg(long = "dflash-draft-model-path")] pub dflash_draft_model_path: Option, @@ -274,7 +248,7 @@ pub(crate) struct Args { /// Enable single-GPU Qwen3 batch-invariant serving by pinning the numeric paths and cutting /// each prompt's prefill chunks on its own grid. Off by default. Requires `--no-prefix-cache`; - /// incompatible with `--kv-offload`, which keeps prefix matching on regardless. + /// incompatible with `--kv-offload-server`, which keeps prefix matching on regardless. #[arg(long, default_value_t = false)] pub batch_invariant: bool, } @@ -360,12 +334,8 @@ fn consumed_args(model_type: ModelType) -> &'static [&'static str] { "glm52_prefill_only", "glm52_prefill_chunk_size", "no_prefix_cache", - "kv_offload", - "kv_offload_host_gib", - "kv_offload_hugepages", - "kv_p2p_metaserver_addr", - "kv_p2p_advertise_addr", - "kv_p2p_nics", + "kv_offload_server", + "kv_offload_namespace", "kv_pd_vllm_seed", "kv_pd_vllm_namespace", "kv_pd_miss_wait_ms", @@ -387,13 +357,9 @@ fn consumed_args(model_type: ModelType) -> &'static [&'static str] { "max_lora_rank", "device_ordinal", "tp_size", - "kv_offload", - "kv_offload_host_gib", - "kv_offload_hugepages", - "kv_p2p_metaserver_addr", - "kv_p2p_advertise_addr", - "kv_p2p_nics", - "kv_p2p_flush_on_finish", + "kv_offload_server", + "kv_offload_namespace", + "kv_offload_flush_on_finish", "kv_pd_vllm_seed", "kv_pd_vllm_namespace", "kv_pd_miss_wait_ms", @@ -488,13 +454,24 @@ impl Args { "--batch-invariant is not compatible with --decode-overlap; the stream override would force the pinned GEMM to bail at runtime" ); } - if self.batch_invariant && self.kv_offload { + if self.batch_invariant && self.kv_offload_server.is_some() { bail!( - "--batch-invariant is not supported with --kv-offload: offload keeps prefix matching \ + "--batch-invariant is not supported with --kv-offload-server: offload keeps prefix matching \ on (--no-prefix-cache only disables HBM retention there), and a host-tier prefix hit \ shifts a prompt's chunk boundaries off the request-local grid" ); } + if self.kv_offload_server.is_some() + && self.kv_pd_vllm_seed.is_none() + && self + .kv_offload_namespace + .as_deref() + .is_none_or(str::is_empty) + { + bail!( + "--kv-offload-namespace is required with --kv-offload-server unless vLLM compatibility supplies --kv-pd-vllm-namespace" + ); + } if self.batch_invariant && self.dflash_draft_model_path.is_some() { bail!( "--batch-invariant is not supported with DFlash speculative decoding; enable one at a time" @@ -658,17 +635,6 @@ pub(crate) fn parse_lora_modules_arg(value: &str) -> Result } } -fn parse_offload_gib(value: &str) -> Result { - let gib = value - .parse::() - .map_err(|error| format!("invalid --kv-offload-host-gib: {error}"))?; - if gib.is_finite() && gib > 0.0 { - Ok(gib) - } else { - Err("--kv-offload-host-gib must be a positive, finite number of GiB".to_owned()) - } -} - #[cfg(feature = "qwen3")] pub(crate) fn parse_max_lora_rank_arg(value: &str) -> Result { let rank = value @@ -772,6 +738,40 @@ mod tests { } } + #[cfg(feature = "qwen3")] + #[test] + fn qwen3_parses_external_kv_offload_endpoint() { + let (args, provided) = parse_with_provided(&[ + "openinfer", + "--kv-offload-server", + "http://127.0.0.1:50055", + "--kv-offload-namespace", + "qwen3-checkpoint-v1", + ]); + assert_eq!( + args.kv_offload_server.as_deref(), + Some("http://127.0.0.1:50055") + ); + assert_eq!( + args.kv_offload_namespace.as_deref(), + Some("qwen3-checkpoint-v1") + ); + args.validate(ModelType::Qwen3, &provided) + .expect("Qwen3 should accept an external PegaFlow endpoint"); + } + + #[cfg(feature = "qwen3")] + #[test] + fn qwen3_native_kv_offload_requires_checkpoint_namespace() { + let (args, provided) = + parse_with_provided(&["openinfer", "--kv-offload-server", "http://127.0.0.1:50055"]); + let error = args + .validate(ModelType::Qwen3, &provided) + .expect_err("native KV sharing without checkpoint identity must be rejected") + .to_string(); + assert!(error.contains("--kv-offload-namespace")); + } + #[cfg(feature = "qwen35")] #[test] fn qwen35_accepts_tp_size() { diff --git a/openinfer-server/src/main.rs b/openinfer-server/src/main.rs index 627c5a35c..b08ef9cfa 100644 --- a/openinfer-server/src/main.rs +++ b/openinfer-server/src/main.rs @@ -187,25 +187,10 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result { - Some(openinfer_glm52::Glm52P2pOptions { - metaserver_addr, - advertise_addr, - rdma_nics: args.kv_p2p_nics.clone(), - }) - } - _ => None, - }, + kv_offload: args.kv_offload_server.clone().map(|server_addr| { + openinfer_glm52::Glm52KvOffloadOptions { + server_addr, + namespace: args.kv_offload_namespace.clone(), vllm_compat: args.kv_pd_vllm_seed.clone().map(|seed| { openinfer_glm52::Glm52VllmCompatOptions { python_hash_seed: seed, @@ -219,7 +204,8 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result anyhow::Result { - let offload = if args.kv_offload { - let bytes = (args.kv_offload_host_gib * f64::from(1u32 << 30)) as usize; - let mut offload = Qwen3OffloadOptions::enabled(bytes); - offload.use_hugepages = args.kv_offload_hugepages; - if let (Some(metaserver_addr), Some(advertise_addr)) = ( - args.kv_p2p_metaserver_addr.clone(), - args.kv_p2p_advertise_addr.clone(), - ) { - offload = offload.with_p2p(openinfer_qwen3::Qwen3P2pOptions { - metaserver_addr, - advertise_addr, - rdma_nics: args.kv_p2p_nics.clone(), - flush_on_finish: args.kv_p2p_flush_on_finish, - }); + let offload = if let Some(server_addr) = &args.kv_offload_server { + let mut offload = Qwen3OffloadOptions::external(server_addr) + .with_flush_on_finish(args.kv_offload_flush_on_finish); + if let Some(namespace) = args.kv_offload_namespace.clone() { + offload = offload.with_namespace(namespace); } if let Some(seed) = args.kv_pd_vllm_seed.clone() { offload = offload.with_vllm_compat(openinfer_qwen3::Qwen3VllmCompatOptions { @@ -295,8 +272,8 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result