diff --git a/Cargo.lock b/Cargo.lock index 20d08baa1..82cd20846 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" @@ -1607,15 +1506,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" @@ -1828,15 +1718,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" @@ -1886,20 +1767,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" @@ -1939,15 +1806,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" @@ -2485,12 +2343,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" @@ -2538,12 +2390,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" @@ -2774,15 +2620,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" @@ -2838,32 +2675,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" @@ -2959,7 +2770,7 @@ dependencies = [ "cc", "cfg-if", "libc", - "memoffset 0.6.5", + "memoffset", ] [[package]] @@ -3071,12 +2882,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" @@ -3268,8 +3073,10 @@ dependencies = [ "log", "openinfer-kv-cache", "pegaflow-core", + "pegaflow-proto", "tokio", - "tokio-stream", + "tonic", + "uuid", "xxhash-rust", ] @@ -3479,16 +3286,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" @@ -3517,17 +3314,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" @@ -3585,16 +3371,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.4" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=29c7fba#29c7fbad851d660ac8c11e9a0a38d5f4b2b4c57b" dependencies = [ "colored", "libc", @@ -3604,8 +3384,8 @@ dependencies = [ [[package]] name = "pegaflow-core" -version = "0.23.3" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=1473c5355d879b4fea23101760cb2a0074642ada#1473c5355d879b4fea23101760cb2a0074642ada" +version = "0.23.4" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=29c7fba#29c7fbad851d660ac8c11e9a0a38d5f4b2b4c57b" dependencies = [ "ahash", "bytesize", @@ -3623,7 +3403,6 @@ dependencies = [ "parking_lot", "pegaflow-common", "pegaflow-proto", - "pegaflow-transfer", "rand 0.10.1", "shared_memory", "smallvec", @@ -3634,8 +3413,8 @@ dependencies = [ [[package]] name = "pegaflow-proto" -version = "0.23.3" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=1473c5355d879b4fea23101760cb2a0074642ada#1473c5355d879b4fea23101760cb2a0074642ada" +version = "0.23.4" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=29c7fba#29c7fbad851d660ac8c11e9a0a38d5f4b2b4c57b" dependencies = [ "prost", "tonic", @@ -3643,35 +3422,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" @@ -3778,7 +3528,7 @@ dependencies = [ "concurrent-queue", "hermit-abi", "pin-project-lite", - "rustix 1.1.4", + "rustix", "windows-sys 0.61.2", ] @@ -3797,19 +3547,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" @@ -3834,32 +3571,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" @@ -4264,19 +3975,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" @@ -4543,19 +4241,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" @@ -4565,7 +4250,7 @@ dependencies = [ "bitflags 2.11.1", "errno", "libc", - "linux-raw-sys 0.12.1", + "linux-raw-sys", "windows-sys 0.61.2", ] @@ -4999,21 +4684,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" @@ -5072,15 +4742,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" @@ -5276,16 +4937,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" @@ -5307,29 +4958,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" @@ -5370,16 +4998,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" @@ -6651,18 +6273,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" @@ -6806,15 +6416,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 19d7ff432..0ddb818e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,6 +159,7 @@ sha2 = "0.11" 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 dae6da79c..f103e8bc9 100644 --- a/docs/index.md +++ b/docs/index.md @@ -117,7 +117,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: OpenInfer registers CUDA IPC allocations and issues save/query/load/flush RPCs; the server owns DRAM/SSD/RDMA. Qwen3 and GLM5.2 cross-process gates on 8×H200 are in progress. | ## 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..06e368f18 --- /dev/null +++ b/docs/subsystems/runtime/external-pegaflow-server.md @@ -0,0 +1,42 @@ +# External PegaFlow server for KV offload + +> **TL;DR:** OpenInfer is a CUDA-IPC/RPC client selected by `--kv-offload-server`; the external PegaFlow process owns the host, SSD, and RDMA tiers. Native registration, terminal Load completion, and Flush are implemented for Qwen3 and GLM5.2, with final-head cross-process replay remaining before deployment. +> +> **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. +- OpenInfer owns model KV allocations and exports them through CUDA IPC. PegaFlow imports those allocations, owns the storage hierarchy, and performs GPU↔host transfers. +- `--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 + +- Each registered layer carries one CUDA IPC allocation handle plus view offset, view size, and block stride. This avoids positionally coupling a second stride array to the layer list. +- The server validates the allocation device and view bounds, opens each allocation once per registration batch, and keeps the mapping alive until unregister. +- Native registration has an exact capability version. Old and new clients fail before importing memory instead of silently assuming a dense layout. +- Native `Load` 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. +- External offload uses exportable CUDA allocations; the default offload-disabled path retains the stream-ordered allocator. + +## Model layouts + +- Qwen3 registers one page-first fused allocation with strided per-layer views. +- GLM5.2 registers rank-local MLA and index-K arenas by name: 78 MLA plus 21 index-K arenas per EP8 rank. +- Both clients preserve PegaFlow query leases and load only after a host-tier hit has produced a valid lease. + +## Validation + +Current PR-head local gates are green: + +- PegaFlow server: 23 unit tests; core: 124 passed with one GPU-only test ignored; CUDA 12/13 checks and Python wheel builds pass. +- OpenInfer KV offload: 9 unit tests; workspace CPU tests, simulated frontend E2E, SM80 CUDA compile, and SM80 CUDA clippy pass. + +Cross-process evidence collected during bring-up established the layout and transfer contract: + +- Qwen page-first save/load restored three blocks into different HBM block IDs byte-for-byte; an untouched block remained zero. +- Qwen3 forced CPU-only and combined GPU/CPU prefix restores without a material head-logprob shift. +- GLM5.2 EP8 registered 99 arenas per rank and restored one 64-token host block with the expected first output tokens. + +Those runs predate the final PR heads. Replay the Qwen and GLM5.2 gates after PegaFlow #407 and OpenInfer #729 are merged; do not treat the historical runs as deployment evidence. diff --git a/docs/subsystems/runtime/pegaflow-offload-integration.md b/docs/subsystems/runtime/pegaflow-offload-integration.md index 4f7c77565..9b52af570 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` 注册 CUDA IPC allocation 并调用 save/query/load/flush RPC，DRAM/hugepage/SSD/RDMA 全归 PegaFlow 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 9468f6733..6b3987f94 100644 --- a/openinfer-glm52/src/lib.rs +++ b/openinfer-glm52/src/lib.rs @@ -43,7 +43,7 @@ use std::{ use anyhow::{Context as _, Result, bail, ensure}; use bytesize::ByteSize; use openinfer_core::engine::{EngineHandle, KvCapacity, LoadSnapshot}; -use openinfer_kv_offload::{HostConfig, KvArena, OffloadEngine, OffloadHost}; +use openinfer_kv_offload::{KvArena, OffloadEngine, OffloadHost}; use remote::Glm52RemoteNode; use runner::{Glm52RankPlacement, Glm52RankWorker, Glm52Worker}; @@ -316,8 +316,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 @@ -325,35 +325,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`). @@ -436,17 +417,6 @@ pub fn launch(model_path: &Path, options: Glm52LaunchOptions) -> Result rank_arenas, - Err(err) => { - for worker in &loaded.workers { - let _ = worker.request_shutdown(); - } - return Err(err); + let rank_arenas = match build_rank_models( + &loaded.workers, + max_model_len, + moe_topo, + dspark_enabled, + kv_offload.is_some(), + ) { + Ok(rank_arenas) => rank_arenas, + Err(err) => { + for worker in &loaded.workers { + let _ = worker.request_shutdown(); } - }; + return Err(err); + } + }; let vllm_compat = kv_offload .as_ref() .and_then(|opts| opts.vllm_compat.clone()); @@ -903,11 +878,14 @@ fn build_rank_models( max_model_len: usize, moe_topo: Glm52MoeTopo, dspark_enabled: bool, + exportable_kv: bool, ) -> Result>> { let build_started = Instant::now(); let responses = workers .iter() - .map(|worker| worker.build_model_async(max_model_len, moe_topo, dspark_enabled)) + .map(|worker| { + worker.build_model_async(max_model_len, moe_topo, dspark_enabled, exportable_kv) + }) .collect::>>()?; let mut rank_arenas = Vec::with_capacity(responses.len()); for (rank, response) in responses.into_iter().enumerate() { @@ -970,20 +948,8 @@ fn build_offload_engines( .all(|arena| arena.bytes_per_block == mla_page_size * mla_bytes_per_token)), "GLM5.2 KV offload ranks disagree on MLA cache layout" ); - let host = OffloadHost::new(HostConfig { - pinned_pool_bytes: opts.pinned_pool_bytes, - use_hugepages: opts.use_hugepages, - runtime_threads: 2, - p2p: opts - .p2p - .as_ref() - .map(|p2p| openinfer_kv_offload::P2pConfig { - metaserver_addr: p2p.metaserver_addr.clone(), - advertise_addr: p2p.advertise_addr.clone(), - rdma_nics: p2p.rdma_nics.clone(), - }), - }) - .map_err(|err| anyhow::anyhow!("GLM5.2 KV offload host: {err}"))?; + let host = OffloadHost::connect(&opts.server_addr, 2) + .map_err(|err| anyhow::anyhow!("GLM5.2 KV offload host: {err}"))?; // vLLM-compat mode joins the *P side's* content domain: the pegaflow // connector derives an 8-hex namespace from vLLM config (and logs it at // startup); reproducing that derivation would mean chasing Python repr @@ -991,7 +957,10 @@ fn build_offload_engines( let namespace = match &opts.vllm_compat { Some(compat) => compat.namespace.clone(), None => format!( - "openinfer-glm52-l{GLM52_LAYERS}-p{}-mla{}-idxk{}", + "openinfer-glm52-{}-l{GLM52_LAYERS}-p{}-mla{}-idxk{}", + opts.namespace + .as_deref() + .context("GLM5.2 native KV offload requires a checkpoint namespace")?, mla_page_size, mla_bytes_per_token, config::GLM52_INDEX_HEAD_DIM + 4, @@ -1029,10 +998,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/model/mod.rs b/openinfer-glm52/src/model/mod.rs index c9c312848..b34fe6261 100644 --- a/openinfer-glm52/src/model/mod.rs +++ b/openinfer-glm52/src/model/mod.rs @@ -21,6 +21,7 @@ use cudarc::driver::{CudaSlice, CudaStream, DevicePtr as _, PinnedHostSlice}; use half::bf16; use openinfer_core::cuda_graph::CudaGraphDumpSummary; use openinfer_core::cuda_graph::CudaGraphState; +use openinfer_kernels::exportable::alloc_ipc_zeros; use openinfer_kernels::ops::{ GLM52_FLASHMLA_SPARSE_BYTES_PER_TOKEN, GLM52_FLASHMLA_SPARSE_PAGE_SIZE, GLM52_FLASHMLA_SPARSE_TOPK, GLM52_GEMV_MMA_SCRATCH_FLOATS_PER_ROW, GLM52_MLA_CACHE_BYTES, @@ -477,6 +478,7 @@ impl Glm52RankModel { moe_topo: crate::Glm52MoeTopo, attn_shard: Option, dspark_enabled: bool, + exportable_kv: bool, ) -> Result { ensure!( moe_topo.uses_tensor_replicated_moe() == attn_shard.is_some(), @@ -523,17 +525,25 @@ impl Glm52RankModel { build::build_decoder_layer(ctx, w, layer, moe_topo, attn_shard) .with_context(|| format!("build GLM5.2 decoder layer {layer}"))?, ); + let mla_len = + contract.num_blocks * GLM52_FLASHMLA_SPARSE_PAGE_SIZE * mla_cache_bytes_per_token; + let mla_cache = if exportable_kv { + alloc_ipc_zeros::(&ctx.stream, mla_len) + } else { + ctx.stream.alloc_zeros::(mla_len) + }?; caches.push(Glm52LayerCaches { - mla_cache: ctx.stream.alloc_zeros::( - contract.num_blocks - * GLM52_FLASHMLA_SPARSE_PAGE_SIZE - * mla_cache_bytes_per_token, - )?, + mla_cache, index_k_cache: glm52_layer_has_full_indexer(layer) .then(|| { - ctx.stream - .alloc_zeros::(index_cache_layout.min_cache_bytes()?) - .map_err(anyhow::Error::from) + let len = index_cache_layout.min_cache_bytes()?; + if exportable_kv { + alloc_ipc_zeros::(&ctx.stream, len).map_err(anyhow::Error::from) + } else { + ctx.stream + .alloc_zeros::(len) + .map_err(anyhow::Error::from) + } }) .transpose()?, }); diff --git a/openinfer-glm52/src/remote.rs b/openinfer-glm52/src/remote.rs index 086dabc2f..9ba982dd3 100644 --- a/openinfer-glm52/src/remote.rs +++ b/openinfer-glm52/src/remote.rs @@ -115,6 +115,7 @@ enum WireRequest { max_model_len: usize, moe_topo: Glm52MoeTopo, dspark_enabled: bool, + exportable_kv: bool, }, SetupComm { unique_id: Vec, @@ -459,6 +460,7 @@ impl Glm52RemoteRankWorker { max_model_len: usize, moe_topo: Glm52MoeTopo, dspark_enabled: bool, + exportable_kv: bool, ) -> Result>>> { let (tx, rx) = bounded(1); self.node.shared.submit( @@ -467,6 +469,7 @@ impl Glm52RemoteRankWorker { max_model_len, moe_topo, dspark_enabled, + exportable_kv, }, PendingResp::BuildModel(tx), )?; @@ -809,10 +812,12 @@ fn host_demux_loop( max_model_len, moe_topo, dspark_enabled, + exportable_kv, } => HostPending::BuildModel(worker.build_model_async( max_model_len, moe_topo, dspark_enabled, + exportable_kv, )?), WireRequest::SetupComm { unique_id, diff --git a/openinfer-glm52/src/runner.rs b/openinfer-glm52/src/runner.rs index c0b3e79ae..20d983cee 100644 --- a/openinfer-glm52/src/runner.rs +++ b/openinfer-glm52/src/runner.rs @@ -98,6 +98,7 @@ enum Glm52RankCommand { max_model_len: usize, moe_topo: crate::Glm52MoeTopo, dspark_enabled: bool, + exportable_kv: bool, resp: Sender>>, }, /// Collective: create the DeepEP context (barriers across ranks). Issued @@ -252,6 +253,7 @@ impl Glm52RankWorker { max_model_len: usize, moe_topo: crate::Glm52MoeTopo, dspark_enabled: bool, + exportable_kv: bool, ) -> Result>>> { let (resp_tx, resp_rx) = bounded(1); self.tx @@ -259,6 +261,7 @@ impl Glm52RankWorker { max_model_len, moe_topo, dspark_enabled, + exportable_kv, resp: resp_tx, }) .map_err(|_| anyhow::anyhow!("GLM5.2 rank worker channel closed"))?; @@ -419,13 +422,14 @@ impl Glm52Worker { max_model_len: usize, moe_topo: crate::Glm52MoeTopo, dspark_enabled: bool, + exportable_kv: bool, ) -> Result>>> { match self { Self::Local(worker) => { - worker.build_model_async(max_model_len, moe_topo, dspark_enabled) + worker.build_model_async(max_model_len, moe_topo, dspark_enabled, exportable_kv) } Self::Remote(worker) => { - worker.build_model_async(max_model_len, moe_topo, dspark_enabled) + worker.build_model_async(max_model_len, moe_topo, dspark_enabled, exportable_kv) } } } @@ -621,6 +625,7 @@ impl Glm52RankThreadState { max_model_len: usize, moe_topo: crate::Glm52MoeTopo, dspark_enabled: bool, + exportable_kv: bool, ) -> Result> { let mut weights = self .loaded @@ -636,6 +641,7 @@ impl Glm52RankThreadState { .uses_tensor_replicated_moe() .then_some(self.placement.rank), dspark_enabled, + exportable_kv, )?); let arenas = model.kv_arenas(&dev_ctx.stream)?; let aux_ctx = self.ctx.auxiliary_device_context("decode aux")?; @@ -923,9 +929,15 @@ fn rank_worker_loop(rx: &Receiver, mut state: Glm52RankThreadS max_model_len, moe_topo, dspark_enabled, + exportable_kv, resp, } => { - let _ = resp.send(state.build_model(max_model_len, moe_topo, dspark_enabled)); + let _ = resp.send(state.build_model( + max_model_len, + moe_topo, + dspark_enabled, + exportable_kv, + )); } Glm52RankCommand::SetupComm { unique_id, diff --git a/openinfer-glm52/src/scheduler/mod.rs b/openinfer-glm52/src/scheduler/mod.rs index ee40bc8c2..6de37daf0 100644 --- a/openinfer-glm52/src/scheduler/mod.rs +++ b/openinfer-glm52/src/scheduler/mod.rs @@ -405,11 +405,11 @@ 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` 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(); + // 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 a40a7fd5f..c30a16f18 100644 --- a/openinfer-glm52/src/scheduler/offload.rs +++ b/openinfer-glm52/src/scheduler/offload.rs @@ -325,7 +325,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 @@ -414,7 +414,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() @@ -578,7 +578,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/exportable.rs b/openinfer-kernels/src/exportable.rs new file mode 100644 index 000000000..65a275a6b --- /dev/null +++ b/openinfer-kernels/src/exportable.rs @@ -0,0 +1,39 @@ +//! CUDA allocations that can be shared through legacy CUDA IPC handles. + +use std::sync::Arc; + +use cudarc::driver::{ + CudaSlice, CudaStream, DeviceRepr, ValidAsZeroBits, + result::DriverError, + sys::{self, CUresult}, +}; + +/// Allocate zeroed, pointer-stable device memory through `cuMemAlloc`. +/// +/// cudarc normally uses `cuMemAllocAsync` on modern GPUs. Those allocations +/// cannot be exported by `cuIpcGetMemHandle`, while persistent KV arenas must +/// be visible to an out-of-process PegaFlow server. `CudaSlice` may release a +/// `cuMemAlloc` allocation through `cuMemFreeAsync`, so the returned value +/// retains the normal stream-ordered drop behavior. +pub fn alloc_ipc_zeros( + stream: &Arc, + len: usize, +) -> Result, DriverError> { + if len == 0 { + return stream.alloc_zeros(0); + } + let bytes = len + .checked_mul(std::mem::size_of::()) + .ok_or(DriverError(CUresult::CUDA_ERROR_INVALID_VALUE))?; + stream.context().bind_to_thread()?; + + let mut ptr = 0; + // SAFETY: `ptr` is valid output storage, the CUDA context is current, and + // the allocation is immediately transferred into a CudaSlice owner. + unsafe { sys::cuMemAlloc_v2(&raw mut ptr, bytes).result()? }; + // SAFETY: `ptr` owns `bytes == len * size_of::()` device bytes on this + // stream's context. `memset_zeros` establishes valid zero bits for T. + let mut slice = unsafe { stream.upgrade_device_ptr(ptr, len) }; + stream.memset_zeros(&mut slice)?; + Ok(slice) +} diff --git a/openinfer-kernels/src/lib.rs b/openinfer-kernels/src/lib.rs index 6ae6e8725..6fdc4e9a0 100644 --- a/openinfer-kernels/src/lib.rs +++ b/openinfer-kernels/src/lib.rs @@ -1,6 +1,7 @@ #![allow(incomplete_features)] #![feature(generic_const_exprs)] +pub mod exportable; pub mod ffi; pub mod forward_pass; pub mod gpu_buffers; diff --git a/openinfer-kv-cache/src/buffer.rs b/openinfer-kv-cache/src/buffer.rs index 381ad7118..0da4d0295 100644 --- a/openinfer-kv-cache/src/buffer.rs +++ b/openinfer-kv-cache/src/buffer.rs @@ -28,12 +28,57 @@ impl KvBuffer { head_dim: usize, page_size: usize, num_blocks: usize, + ) -> anyhow::Result { + Self::allocate( + stream, + num_layers, + num_kv_heads, + head_dim, + page_size, + num_blocks, + false, + ) + } + + /// Allocate through legacy `cuMemAlloc` so the fused arena can be exported + /// to an out-of-process CUDA IPC consumer. + pub fn new_exportable( + stream: &Arc, + num_layers: usize, + num_kv_heads: usize, + head_dim: usize, + page_size: usize, + num_blocks: usize, + ) -> anyhow::Result { + Self::allocate( + stream, + num_layers, + num_kv_heads, + head_dim, + page_size, + num_blocks, + true, + ) + } + + #[allow(clippy::too_many_arguments)] + fn allocate( + stream: &Arc, + num_layers: usize, + num_kv_heads: usize, + head_dim: usize, + page_size: usize, + num_blocks: usize, + exportable: bool, ) -> 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 - .alloc_zeros(total_elements) - .map_err(|e| anyhow::anyhow!("KvBuffer alloc failed: {e}"))?; + let buffer: CudaSlice = if exportable { + openinfer_kernels::exportable::alloc_ipc_zeros(stream, total_elements) + } else { + stream.alloc_zeros(total_elements) + } + .map_err(|e| anyhow::anyhow!("KvBuffer alloc failed: {e}"))?; Ok(Self { inner: Arc::new(Inner { buffer, diff --git a/openinfer-kv-cache/src/manager.rs b/openinfer-kv-cache/src/manager.rs index 8564b72d9..9bc051888 100644 --- a/openinfer-kv-cache/src/manager.rs +++ b/openinfer-kv-cache/src/manager.rs @@ -38,6 +38,26 @@ impl KvCacheManager { Ok(Self { pool, buffer }) } + pub fn new_exportable( + stream: &Arc, + num_layers: usize, + num_kv_heads: usize, + head_dim: usize, + block_size: usize, + num_blocks: usize, + ) -> anyhow::Result { + let buffer = KvBuffer::new_exportable( + stream, + num_layers, + num_kv_heads, + head_dim, + block_size, + num_blocks, + )?; + 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 d6aa7caa1..bd1591a40 100644 --- a/openinfer-kv-offload/Cargo.toml +++ b/openinfer-kv-offload/Cargo.toml @@ -5,26 +5,17 @@ version = "0.1.0" edition = "2024" [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. -pegaflow-core = { git = "https://github.com/novitalabs/pegaflow.git", rev = "1473c5355d879b4fea23101760cb2a0074642ada", default-features = false, features = ["rdma"] } +# The client shares PegaFlow's wire-compatible error/save types but does not +# embed its storage engine or RDMA stack. +pegaflow-core = { git = "https://github.com/novitalabs/pegaflow.git", rev = "29c7fba", default-features = false } +pegaflow-proto = { git = "https://github.com/novitalabs/pegaflow.git", rev = "29c7fba" } openinfer-kv-cache = { workspace = true } cudarc = { workspace = true } half = { workspace = true } log = { workspace = true } tokio = { workspace = true } -tokio-stream = { workspace = true, features = ["net"] } +tonic = { workspace = true } +uuid = { workspace = true } xxhash-rust = { workspace = true } [lints] diff --git a/openinfer-kv-offload/src/engine.rs b/openinfer-kv-offload/src/engine.rs index 4a5ecf7e4..94fa6f375 100644 --- a/openinfer-kv-offload/src/engine.rs +++ b/openinfer-kv-offload/src/engine.rs @@ -1,119 +1,65 @@ -//! [`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. +//! OpenInfer owns the GPU allocations and logical prefix-cache policy. The +//! external server owns every deeper tier and the transfer workers. This +//! module registers CUDA IPC views, then translates scheduler save/query/load +//! operations into PegaFlow RPCs. use std::sync::{Arc, Mutex}; +use std::time::Duration; use cudarc::driver::CudaStream; use openinfer_kv_cache::KvBuffer; -use pegaflow_core::{ - EngineError, LayerSave, P2pTransferService, PegaEngine, PrefetchStatus, QueryLeaseId, - StorageConfig, TransferMode, -}; +use pegaflow_core::{EngineError, LayerSave}; 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, ExternalQuery, ExternalRegistration, 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. + /// 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 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. + /// Content-addressing domain shared by producers and consumers whose KV + /// bytes and block layout are interchangeable. pub namespace: String, - /// CUDA device ordinal whose KV buffer this engine offloads. + /// CUDA device ordinal whose allocations are exported over CUDA IPC. pub device_id: i32, - /// Host pinned-memory pool size in bytes (the CPU KV tier capacity). - pub 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. + /// 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, - /// `Some` joins the cross-instance P2P mesh (see [`P2pConfig`]). - pub p2p: Option, } 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(), } } @@ -122,74 +68,58 @@ 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. pub struct KvArena { pub name: String, pub base_ptr: u64, @@ -198,11 +128,6 @@ 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). struct Registration { layer_names: Vec, data_ptrs: Vec, @@ -215,26 +140,12 @@ struct Registration { } 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. 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(), @@ -247,12 +158,10 @@ impl Registration { 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 { + assert!(!arenas.is_empty(), "KV offload requires at least one arena"); let n = arenas.len(); - let mut reg = Registration { + let mut reg = Self { layer_names: Vec::with_capacity(n), data_ptrs: Vec::with_capacity(n), size_bytes: Vec::with_capacity(n), @@ -263,6 +172,12 @@ impl Registration { block_stride_bytes: Vec::with_capacity(n), }; for arena in arenas { + assert!(arena.num_blocks > 0, "arena {} has no blocks", arena.name); + assert!( + arena.bytes_per_block > 0, + "arena {} has an empty copy unit", + arena.name + ); assert!( arena.bytes_per_block <= arena.block_stride_bytes, "arena {} copy unit {} overruns its block stride {}", @@ -270,12 +185,15 @@ impl Registration { arena.bytes_per_block, arena.block_stride_bytes ); + let strided_prefix = (arena.num_blocks - 1) + .checked_mul(arena.block_stride_bytes) + .expect("arena byte reach overflows usize"); + let size = strided_prefix + .checked_add(arena.bytes_per_block) + .expect("arena byte reach overflows usize"); 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.size_bytes.push(size); 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); @@ -284,217 +202,69 @@ impl Registration { } } -/// 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. - 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, })) } - pub fn p2p_enabled(&self) -> bool { - self.p2p_shutdown.is_some() + 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 set of GPU KV arenas. 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. pub fn new( - config: OffloadConfig, + config: &OffloadConfig, buffer: &KvBuffer, stream: &CudaStream, ) -> Result { let reg = Registration::from_buffer(buffer, stream); - let host = OffloadHost::new(config.host())?; + 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, @@ -502,18 +272,11 @@ impl OffloadEngine { ) } - /// Build the engine over explicit arenas instead of one fused - /// [`KvBuffer`] — for models whose per-layer caches are separate - /// allocations (GLM5.2: MLA latent + index-K per layer). Same contract as - /// [`Self::new`], plus: 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. - pub fn with_arenas(config: OffloadConfig, arenas: &[KvArena]) -> Result { - let host = OffloadHost::new(config.host())?; + pub fn with_arenas(config: &OffloadConfig, arenas: &[KvArena]) -> Result { + let host = OffloadHost::connect(&config.server_addr, config.runtime_threads)?; Self::register( host, - config.instance_id, + &config.instance_id, &config.namespace, config.device_id, Registration::from_arenas(arenas), @@ -521,14 +284,8 @@ impl OffloadEngine { ) } - /// [`Self::with_arenas`] 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`. - /// `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. + /// Register arenas on a connection shared by several ranks in this process. + /// `page_first` must match the producer's host layout for the namespace. pub fn with_arenas_on( host: Arc, instance_id: impl Into, @@ -537,9 +294,10 @@ impl OffloadEngine { arenas: &[KvArena], page_first: bool, ) -> Result { + let instance_id = instance_id.into(); Self::register( host, - instance_id.into(), + &instance_id, namespace, device_id, Registration::from_arenas(arenas), @@ -549,44 +307,37 @@ impl OffloadEngine { 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, - )?; + assert_outside_runtime("register"); + let instance_id = host.instance_id(instance_id); + let session = 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, + data_ptrs: ®.data_ptrs, + 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, + }))?; Ok(Self { + session: Some(session), host, instance_id, device_id, @@ -598,13 +349,11 @@ impl OffloadEngine { }) } - /// 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 { @@ -615,28 +364,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], @@ -648,36 +377,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], @@ -689,25 +405,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, @@ -720,123 +426,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 }) - } - - /// Whether this engine participates in the cross-instance P2P mesh. - pub fn p2p_enabled(&self) -> bool { - self.host.p2p_enabled() + 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) = { @@ -845,24 +501,23 @@ 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(()); @@ -870,11 +525,35 @@ impl OffloadEngine { }); } - /// Drop all resident CPU-tier blocks (test/eviction helper). Saved data in - /// a backing store would survive; the dense v1 path has none, so this - /// empties the CPU tier. - pub fn evict_all(&self) { - self.host.engine.cleanup_memory_cache(); + /// Flush submitted saves, then ask the server to close this instance's + /// imported CUDA mappings. + pub fn shutdown(&mut self) { + if self.session.is_none() { + return; + } + assert_outside_runtime("shutdown"); + self.flush_saves(); + if let Err(err) = self + .host + .runtime + .block_on(self.host.client.unregister(&self.instance_id)) + { + log::error!( + "PegaFlow unregister barrier failed for {}: {err}; aborting before exported CUDA \ + memory can be released", + self.instance_id + ); + std::process::abort(); + } + // The explicit unregister owns mapping cleanup. Aborting the liveness + // stream afterwards only removes its now-empty session entry. + self.session.take(); + } +} + +impl Drop for OffloadEngine { + fn drop(&mut self) { + self.shutdown(); } } @@ -882,15 +561,11 @@ impl OffloadEngine { 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 = [ + let reg = Registration::from_arenas(&[ KvArena { name: "0.mla".into(), base_ptr: 0x1000, @@ -905,8 +580,7 @@ mod tests { 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]); assert_eq!(reg.segments, [1, 1]); @@ -917,9 +591,6 @@ mod tests { 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 { @@ -932,6 +603,15 @@ mod tests { assert_eq!(reg.size_bytes, [3 * 4096 + 512]); } + #[test] + 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); + } + #[test] #[should_panic(expected = "overruns its block stride")] fn arena_copy_unit_must_fit_its_stride() { diff --git a/openinfer-kv-offload/src/external.rs b/openinfer-kv-offload/src/external.rs new file mode 100644 index 000000000..76e693092 --- /dev/null +++ b/openinfer-kv-offload/src/external.rs @@ -0,0 +1,511 @@ +use std::collections::HashMap; +use std::mem::MaybeUninit; + +use cudarc::driver::{CudaContext, sys}; +use pegaflow_core::{EngineError, LayerSave}; +use pegaflow_proto::proto::engine::engine_client::EngineClient; +use pegaflow_proto::proto::engine::{ + CudaIpcTensor, FlushRequest, HealthRequest, LeaseLoad, LoadRequest, QueryRequest, + RegisterContextRequest, ReleaseRequest, ResponseStatus, SaveLayer, SaveRequest, SessionRequest, + TransferMode, UnregisterRequest, query_response, +}; +use tonic::transport::{Channel, Endpoint}; +use tonic::{Request, Streaming}; + +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); + +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], + pub data_ptrs: &'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, +} + +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 }) + } + + pub(super) async fn register( + &self, + registration: ExternalRegistration<'_>, + ) -> Result { + let cuda_ipc_tensors = export_cuda_ipc_tensors( + registration.device_id, + registration.data_ptrs, + 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, + cuda_ipc_tensors, + }; + 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)) => { + let register_error = rpc_error("register_context_batch", &err); + if let Err(cleanup_error) = self.unregister(registration.instance_id).await { + abort_cleanup_failure("failed registration", &cleanup_error); + } + return Err(register_error); + } + Err(_) => abort_ownership_timeout("register_context_batch"), + }; + if let Err(register_error) = require_ok("register_context_batch", response.status) { + if let Err(cleanup_error) = self.unregister(registration.instance_id).await { + abort_cleanup_failure("rejected registration", &cleanup_error); + } + return Err(register_error); + } + Ok(session) + } + + 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( + RPC_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_indeterminate_transfer("save", &err), + Err(_) => abort_ownership_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(), + })) + .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( + RPC_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_indeterminate_transfer("load", &err), + Err(_) => abort_ownership_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(()) + } + + 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) + } + + 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)) => return Err(rpc_error("unregister_context", &err)), + Err(_) => abort_ownership_timeout("unregister_context"), + }; + require_ok("unregister_context", response.status) + } +} + +async fn watch_session(mut stream: Streaming) { + loop { + match stream.message().await { + Ok(Some(_)) => {} + Ok(None) => { + log::error!("external PegaFlow session closed by server"); + break; + } + Err(err) => { + log::error!("external PegaFlow session failed: {err}"); + break; + } + } + } +} + +fn deadline_request(message: T) -> Request { + let mut request = Request::new(message); + request.set_timeout(RPC_DEADLINE); + request +} + +fn abort_indeterminate_transfer(operation: &str, err: &tonic::Status) -> ! { + log::error!( + "external PegaFlow {operation} RPC ended before the server acknowledged DMA completion: \ + {err}; aborting to keep exported CUDA memory from being reused" + ); + std::process::abort(); +} + +fn abort_ownership_timeout(operation: &str) -> ! { + log::error!( + "external PegaFlow {operation} did not acknowledge CUDA ownership within \ + {RPC_DEADLINE:?}; aborting before exported memory can be reused" + ); + std::process::abort(); +} + +fn abort_cleanup_failure(context: &str, err: &EngineError) -> ! { + log::error!( + "external PegaFlow cleanup after {context} failed: {err}; aborting before exported CUDA \ + memory can be released" + ); + std::process::abort(); +} + +fn export_cuda_ipc_tensors( + device_id: i32, + data_ptrs: &[u64], + size_bytes: &[usize], + block_stride_bytes: &[usize], +) -> Result, EngineError> { + if data_ptrs.len() != size_bytes.len() || data_ptrs.len() != block_stride_bytes.len() { + return Err(EngineError::InvalidArgument(format!( + "CUDA IPC export metadata length mismatch: pointers={}, sizes={}, block strides={}", + data_ptrs.len(), + size_bytes.len(), + block_stride_bytes.len() + ))); + } + let device = usize::try_from(device_id).map_err(|_| { + EngineError::InvalidArgument(format!("device_id {device_id} must be non-negative")) + })?; + let context = CudaContext::new(device) + .map_err(|err| EngineError::CudaInit(format!("retain device {device}: {err}")))?; + context + .bind_to_thread() + .map_err(|err| EngineError::CudaInit(format!("bind device {device}: {err}")))?; + + let mut handles = HashMap::>::new(); + data_ptrs + .iter() + .copied() + .zip(size_bytes.iter().copied()) + .zip(block_stride_bytes.iter().copied()) + .map(|((data_ptr, view_size), block_stride)| { + let mut allocation_base = 0; + let mut allocation_size = 0; + // SAFETY: every pointer comes from a live registered KV arena in + // this bound CUDA context; both outputs are valid stack storage. + unsafe { + sys::cuMemGetAddressRange_v2( + &raw mut allocation_base, + &raw mut allocation_size, + data_ptr, + ) + .result() + .map_err(|err| cuda_error("query allocation range", err))?; + } + let offset = data_ptr.checked_sub(allocation_base).ok_or_else(|| { + EngineError::Storage(format!( + "CUDA allocation base {allocation_base:#x} exceeds view pointer {data_ptr:#x}" + )) + })?; + let end = usize::try_from(offset) + .ok() + .and_then(|offset| offset.checked_add(view_size)) + .ok_or_else(|| EngineError::Storage("CUDA IPC view range overflows".to_string()))?; + if end > allocation_size { + return Err(EngineError::InvalidArgument(format!( + "CUDA IPC view {offset}..{end} exceeds allocation size {allocation_size}" + ))); + } + + let handle = if let Some(handle) = handles.get(&allocation_base) { + handle.clone() + } else { + let mut handle = MaybeUninit::::uninit(); + // SAFETY: allocation_base is the live base returned above and + // handle is valid uninitialized output storage. + unsafe { + sys::cuIpcGetMemHandle(handle.as_mut_ptr(), allocation_base) + .result() + .map_err(|err| cuda_error("export CUDA IPC handle", err))?; + } + // SAFETY: CUDA initialized the complete fixed-size handle. + let handle = unsafe { handle.assume_init() }; + // SAFETY: CUipcMemHandle is an opaque byte payload with a fixed + // C layout and remains live for this copy. + let bytes = unsafe { + std::slice::from_raw_parts( + (&raw const handle).cast::(), + std::mem::size_of::(), + ) + } + .to_vec(); + handles.insert(allocation_base, bytes.clone()); + bytes + }; + Ok(CudaIpcTensor { + handle, + 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 cuda_error(operation: &str, err: cudarc::driver::result::DriverError) -> EngineError { + EngineError::Storage(format!("{operation}: {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 346feb207..fdfea93bd 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. +//! deeper tiers (host pinned memory, SSD, RDMA). [`OffloadEngine`] registers +//! the GPU allocation 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,13 +14,13 @@ //! polls its [`LoadHandle`] each scheduler tick. mod engine; +mod external; mod vllm_hash; pub use engine::{ - HostConfig, KvArena, LoadHandle, OffloadConfig, OffloadEngine, OffloadHost, P2pConfig, - QueryHit, QueryOutcome, + KvArena, LoadHandle, OffloadConfig, OffloadEngine, OffloadHost, QueryHit, QueryLeaseId, + QueryOutcome, }; pub use vllm_hash::{VLLM_HASH_BYTES, VllmBlockHasher}; -// Re-exported so callers name pegaflow's engine types through this bridge. -pub use pegaflow_core::{EngineError, PegaEngine, QueryLeaseId}; +pub use pegaflow_core::EngineError; diff --git a/openinfer-kv-offload/tests/cpu_roundtrip.rs b/openinfer-kv-offload/tests/cpu_roundtrip.rs index 576e26992..fa50aa608 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. //! @@ -55,11 +55,15 @@ 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( + let buffer = KvBuffer::new_exportable( &stream, NUM_LAYERS, NUM_KV_HEADS, @@ -93,12 +97,8 @@ 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 config = OffloadConfig::new("roundtrip-test", 0, server_addr); + let engine = OffloadEngine::new(&config, &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(); @@ -121,7 +121,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"); diff --git a/openinfer-qwen3/src/executor.rs b/openinfer-qwen3/src/executor.rs index c6cceddf9..a41648e97 100644 --- a/openinfer-qwen3/src/executor.rs +++ b/openinfer-qwen3/src/executor.rs @@ -935,7 +935,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 @@ -1094,7 +1094,12 @@ impl Qwen3Executor { }), ) } else { - let kv_mgr = KvCacheManager::new( + let allocate = if offload_opts.is_enabled() { + KvCacheManager::new_exportable + } else { + KvCacheManager::new + }; + let kv_mgr = allocate( &model.device_ctx().stream, budget.num_layers, budget.num_kv_heads, @@ -1112,7 +1117,7 @@ impl Qwen3Executor { 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 = build_offload(offload_opts, &kv_mgr, model.device_ctx())?; 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() { @@ -1176,10 +1181,7 @@ impl Qwen3Executor { 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, @@ -1229,7 +1231,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() @@ -1247,7 +1249,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)" ); @@ -1986,7 +1988,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, @@ -2082,21 +2084,17 @@ fn profile_kv_budget_on_worker( fn build_offload( opts: &Qwen3OffloadOptions, kv_mgr: &KvCacheManager, - config: &Config, ctx: &DeviceContext, ) -> Result> { - if !opts.enabled { + 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 @@ -2104,37 +2102,21 @@ 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); + let engine = OffloadEngine::new(&config, kv_mgr.buffer(), &ctx.stream) .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(), - ); + log::info!("KV offload enabled on device {device_id}: server={server_addr}"); Ok(Some(engine)) } @@ -3012,6 +2994,9 @@ mod tests { impl Drop for Qwen3Executor { fn drop(&mut self) { + if let Some(mut offload) = self.offload.take() { + offload.shutdown(); + } self.primary.shutdown(); for worker in &mut self.workers { worker.shutdown(); diff --git a/openinfer-qwen3/src/lib.rs b/openinfer-qwen3/src/lib.rs index b916d990c..8ab1de44a 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 { - pub enabled: bool, - /// Host pinned-memory pool size (the CPU KV-tier capacity), in bytes. - pub 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. - pub 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. pub 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 @@ -151,35 +132,40 @@ pub struct Qwen3VllmCompatOptions { } impl Qwen3OffloadOptions { - /// 8 GiB host tier — a few thousand dense Qwen3-4B blocks. - pub const DEFAULT_PINNED_POOL_BYTES: usize = 8 << 30; - 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); @@ -308,10 +294,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 @@ -161,9 +135,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, @@ -171,7 +145,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, @@ -251,7 +225,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, } @@ -315,12 +289,8 @@ fn consumed_args(model_type: ModelType) -> &'static [&'static str] { "dflash_draft_model_path", "max_model_len", "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", @@ -341,13 +311,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", @@ -441,13 +407,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" @@ -561,17 +538,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 @@ -675,6 +641,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-4b")] #[test] fn qwen35_accepts_tp_size() { diff --git a/openinfer-server/src/main.rs b/openinfer-server/src/main.rs index 8c276837f..b1dd60651 100644 --- a/openinfer-server/src/main.rs +++ b/openinfer-server/src/main.rs @@ -155,25 +155,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, @@ -187,7 +172,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 { @@ -262,8 +239,8 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result