From 38a6201ff01a3878f28f2e5ee4f9a6034d180833 Mon Sep 17 00:00:00 2001 From: xiaguan <751080330@qq.com> Date: Sun, 19 Jul 2026 13:03:21 -0400 Subject: [PATCH 1/5] feat(kv-offload): use external PegaFlow server Signed-off-by: xiaguan <751080330@qq.com> --- Cargo.lock | 447 +--------- Cargo.toml | 1 + docs/index.md | 3 +- .../runtime/external-pegaflow-server.md | 83 ++ .../runtime/pegaflow-offload-integration.md | 4 +- openinfer-glm52/src/lib.rs | 98 +-- openinfer-glm52/src/model/mod.rs | 26 +- openinfer-glm52/src/remote.rs | 5 + openinfer-glm52/src/runner.rs | 18 +- openinfer-glm52/src/scheduler/mod.rs | 10 +- openinfer-glm52/src/scheduler/offload.rs | 6 +- openinfer-kernels/src/exportable.rs | 39 + openinfer-kernels/src/lib.rs | 1 + openinfer-kv-cache/src/buffer.rs | 51 +- openinfer-kv-cache/src/manager.rs | 20 + openinfer-kv-offload/Cargo.toml | 21 +- openinfer-kv-offload/src/engine.rs | 793 ++++++------------ openinfer-kv-offload/src/external.rs | 511 +++++++++++ openinfer-kv-offload/src/lib.rs | 14 +- openinfer-kv-offload/tests/cpu_roundtrip.rs | 20 +- openinfer-qwen3/src/executor.rs | 71 +- openinfer-qwen3/src/lib.rs | 80 +- .../tests/batch_invariance_reject.rs | 3 +- openinfer-qwen3/tests/kv_offload_cpu_hit.rs | 14 +- openinfer-server/src/config.rs | 152 ++-- openinfer-server/src/main.rs | 49 +- 26 files changed, 1229 insertions(+), 1311 deletions(-) create mode 100644 docs/subsystems/runtime/external-pegaflow-server.md create mode 100644 openinfer-kernels/src/exportable.rs create mode 100644 openinfer-kv-offload/src/external.rs diff --git a/Cargo.lock b/Cargo.lock index 20d08baa1..981d39b69 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" @@ -728,7 +651,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -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" @@ -1419,7 +1318,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -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" @@ -2341,7 +2199,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2395,7 +2253,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -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=5b73e5b#5b73e5bb7ee30ca3f6f4765c21b80b4df634d6cd" 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=5b73e5b#5b73e5bb7ee30ca3f6f4765c21b80b4df634d6cd" 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=5b73e5b#5b73e5bb7ee30ca3f6f4765c21b80b4df634d6cd" 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" @@ -4017,7 +3728,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ "heck", - "itertools 0.13.0", + "itertools 0.10.5", "log", "multimap", "petgraph", @@ -4038,7 +3749,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.13.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -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,8 +4250,8 @@ dependencies = [ "bitflags 2.11.1", "errno", "libc", - "linux-raw-sys 0.12.1", - "windows-sys 0.61.2", + "linux-raw-sys", + "windows-sys 0.52.0", ] [[package]] @@ -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", - "windows-sys 0.61.2", + "rustix", + "windows-sys 0.52.0", ] -[[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" @@ -6705,7 +6315,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -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..2c793b662 --- /dev/null +++ b/docs/subsystems/runtime/external-pegaflow-server.md @@ -0,0 +1,83 @@ +# External PegaFlow server for KV offload + +> **TL;DR:** KV offload now has one ownership boundary: OpenInfer is a CUDA-IPC/RPC client selected by `--kv-offload-server`, while the external PegaFlow server owns storage and acknowledges transfer/drain barriers before exported GPU memory can be reused; Qwen3 and GLM5.2 cross-process restore gates pass on 8×H200. +> +> **Last touched:** 2026-07 + +## Preparation + +- **Read**: + - `docs/index.md` — routes this cross-model change to the runtime subsystem. + - `docs/subsystems/runtime/pegaflow-offload-integration.md` — records the current in-process ownership model, transfer ordering, leases, and Qwen3 CPU-hit gate. + - `docs/models/qwen3/prefix-cache.md` — establishes full-block matching, the one-token recompute boundary, and cache-lifetime invariants. + - `docs/models/qwen3/serving-performance.md` — retains the current host-tier offload baseline and pure-L2 measurement shape. + - `docs/models/glm52/pd-m2-execution.md` — establishes GLM5.2's 99-arena page-first layout, PegaFlow version boundary, and strict restore semantics. + - `openinfer-kv-offload/src/{lib.rs,engine.rs}` — `OffloadEngine` currently owns an in-process `PegaEngine`; the scheduler-facing save/query/load API is already shared by the model integrations. + - `openinfer-qwen3/src/{lib.rs,executor.rs}` — Qwen3 constructs one offload engine around its fused page-first `KvBuffer` and already has asynchronous prefetch admission. + - `openinfer-glm52/src/lib.rs` — GLM5.2 shares one `OffloadHost` across eight rank instances and registers 78 MLA plus 21 index-K arenas per rank. +- **Relevant history**: + - PegaFlow PRs #331/#333 made strided page-first registration and in-process completion usable by OpenInfer. + - OpenInfer #316 proved Qwen3 host-tier save/load correctness with an embedded pool. + - OpenInfer #600/#657 proved GLM5.2 host offload and cross-node target-KV restore with an embedded pool. + - The development cluster already runs a PegaFlow DaemonSet on selected nodes; the first remote step is read-only inventory, and the production daemon must not be modified for bring-up. +- **Plan**: + 1. On the reserved 8×H200 development node, inspect the PegaFlow DaemonSet endpoint, version, server flags, CUDA IPC registration contract, available model paths, and existing OpenInfer/PegaFlow worktrees without changing the running daemon. + 2. Replace the embedded `PegaEngine` with an external PegaFlow gRPC client while preserving the scheduler-visible query/lease/load/save semantics. + 3. Reuse the server's CUDA IPC registration when it can represent Qwen3's strided fused buffer and GLM5.2's named multi-arena/page-first layout. If it cannot, extend PegaFlow's protocol and server with the smallest layout-general registration primitive, verify its own integration tests, and prepare an upstream PegaFlow PR. + 4. Expose one explicit OpenInfer server configuration for the external endpoint; fail at startup on incompatible combinations, version/layout mismatch, missing server health, or registration failure. + 5. Gate Qwen3 with an external-server register → save → HBM eviction → load path, byte/logit parity, and a repeated-prompt serving check that proves a host-tier hit rather than a GPU-prefix hit. + 6. Gate GLM5.2 with all eight ranks registered against one external server, restore of both MLA and index-K arenas, token parity after forced host-tier restore, and a repeated-prompt serving check with no local GPU-prefix hit. + 7. Run formatting, focused unit/integration tests, and cross-process correctness/latency measurements before review. Record deployment commands without internal hostnames or private paths, then submit the required OpenInfer and PegaFlow PRs. +- **Risks / open questions**: + - Cross-process registration requires CUDA IPC handles rather than raw device pointers; handle lifetime, device identity, and the server's embedded Python/CUDA context must be verified on the real daemon. + - The existing gRPC schema may assume vLLM tensor metadata and may not expose GLM5.2's named arena ordering. A protocol change is acceptable, but silent layout coercion is not. + - Qwen3 uses one engine per device while GLM5.2 uses one host shared across eight rank registrations. The external client must preserve that ownership distinction without introducing a generic abstraction above the transfer layer. + - The daemon's pinned pool is shared infrastructure. Bring-up needs isolated namespaces and must not evict, stop, or reconfigure unrelated cache contents. + - Correctness evidence must force HBM misses; a warm request that hits OpenInfer's local prefix cache does not prove external offload. + +## Execution Log + +### Step 1: development-cluster inventory + +- Connected to the reserved 8×H200 development node and inspected the cluster without modifying workloads. +- The cluster runs a host-networked PegaFlow DaemonSet with host IPC, all GPUs initialized, and a MetaServer connection. Its deployed protocol predates the native Rust registration added here, so compatibility must be checked rather than assumed. +- The reserved node was occupied by unrelated inference processes, which were left untouched. Verification moved to an isolated 8×H200 development pod with local Qwen3 and GLM5.2 checkpoints. +- The production daemon already owns CUDA contexts and host hugepages. External mode must use that pool instead of creating another embedded pool. + +### Step 2: external-only ownership boundary + +- Removed the embedded `PegaEngine`, local pinned-pool construction, local P2P service, and their CLI knobs. `--kv-offload-server` is now the sole enablement switch; server resource configuration stays outside OpenInfer. +- Added a native PegaFlow registration payload carrying CUDA IPC handles, per-view offsets/sizes, and optional block strides. This represents both Qwen3's page-interleaved fused allocation and GLM5.2's MLA/index-K arenas without Python tensor wrappers. +- Added a PegaFlow `Flush` RPC for the existing save/registration visibility barrier. Native loads now await the server's GPU worker through the Load RPC; the Python connector retains its shared-memory completion path. +- KV allocations use legacy-exportable CUDA allocation only when external offload is enabled. The default path retains the stream-ordered allocator. A direct GPU probe confirmed that `cudaIpcGetMemHandle` rejects `cudaMallocAsync` allocations and accepts `cudaMalloc` allocations. +- Local gates passed: PegaFlow server unit tests (26), OpenInfer KV-offload unit tests (8), Qwen3 unit tests (75), default-feature server compile/tests, and a GLM5.2 feature compile using NCCL 2.30.7. The local GPU is not SM90, so real GLM5.2 kernels and behavior still require the 8×H200 target. + +### Step 3: 8×H200 cross-process gates + +- Built the PegaFlow server without RDMA and built the GLM5.2 OpenInfer feature on the target SM90a system. The build exercised TileLang 0.1.12 code generation, FlashMLA, DeepGEMM, NCCL 2.30.7, and the final server link. +- The raw Qwen page-first gate saved three blocks through a separate PegaFlow process, loaded them into different HBM block ids, and compared every layer/segment byte-for-byte. The negative-control block stayed zero. The test body took 0.44 s; a hot build plus test took 1.20 s. +- The Qwen3-4B executor gate forced both a CPU-only prefix restore and a combined GPU+CPU restore. Warm/cold head-logprob mean deltas were 0.0105 and 0.0242 nat; maxima were 0.0652 and 0.0632. The test body took 5.61 s. +- GLM5.2 EP8 registered eight rank instances with 99 arenas each (78 MLA plus 21 index-K) through native CUDA IPC. A 66-token cold request saved one logical block across all 99 arenas: 3,452,160 bytes, with the PegaFlow D2H worker reporting 0.50 ms. +- Nine distinct 3,900-token requests pinned to the same data-parallel rank displaced the original prefix from its 520-block HBM LRU. Replaying the original request restored one block from PegaFlow, returned `cached_tokens=64`, reproduced the same first two output tokens, and logged `GLM5.2 host-tier restore: 1 blocks committed`. PegaFlow performed 99 H2D copies (3,452,160 bytes) in 0.62 ms. End-to-end request time was 0.256 s cold and 0.044 s after host-tier restore; this is a correctness-gate observation, not a serving benchmark. +- The isolated OpenInfer and PegaFlow processes were stopped after the gate; all eight GPUs returned to zero allocated memory. The development pod and source mounts remain available. + +### Step 4: failure and multi-tenant review + +- Scoped PegaFlow's asynchronous deeper-tier prefetch state by `(instance_id, req_id)`. A regression test keeps two equal request IDs in separate active entries and proves that polling one instance cannot consume the other's result. +- Replaced process-ID instance suffixes with UUIDv4 and made native content namespaces explicit. Native Qwen3 and GLM5.2 launches now require `--kv-offload-namespace`, a stable checkpoint/deployment identity; vLLM compatibility continues to use the connector namespace. +- Deadlines remain on health, query, release, and flush. Ownership RPCs do not carry a cancelable gRPC deadline because an accepted CUDA DMA cannot be canceled with its handler; a 30-second local watchdog or an indeterminate transport failure is fail-stop. Native load completion is returned by the Load RPC instead of a shared-memory poll. +- Added an owned session guard and explicit client shutdown. Shutdown waits for accepted saves, sends `UnregisterContext`, and releases the session only after the server acknowledges both GPU queue drain and CUDA mapping close. Failure to obtain that ownership acknowledgement is fail-stop. +- PegaFlow unregister keeps a closing instance addressable to concurrent cleanup calls, rejects new transfers, places barriers behind every accepted load/save on every registered device, and removes the instance only after the shared drain completes. The worker threads retain join handles for deterministic teardown. +- Registration and cleanup are serialized per instance and by session generation. A canceled registry reply rolls back a newly published CUDA mapping; session cleanup still closes registry state when engine registration never completed. +- Kept legacy-exportable allocation behind the offload option so an offload-disabled launch uses the original allocator. +- Failure-path local gates pass: PegaFlow server 28/28, PegaFlow core 126/126 with one GPU-only test ignored, and OpenInfer KV-offload 9/9. +- Re-ran the cross-process gates after the lifecycle fixes. The raw Qwen round trip remained byte-exact at 0.43 s versus 0.44 s before the fixes. The hot Qwen3 executor gate took 3.32 s; CPU-only and combined GPU/CPU restores kept mean head-logprob deltas at 0.0105 and 0.0272 nat. +- Re-ran GLM5.2 with a 520-block rank-local HBM pool. After rank-pinned displacement, the original 66-token request restored one host block, returned `cached_tokens=64`, reproduced the same first two tokens, and completed in 0.0438 s versus 0.257 s cold. A graceful stop drained all eight rank instances, closed 99 CUDA mappings per rank, emptied the server instance list, and returned all eight GPUs to zero allocated memory after the PegaFlow process exited. + +## Debrief + +The embedded and external modes represented two owners for the same host-tier state and doubled every lifecycle decision. Removing the embedded mode left a stable boundary: model crates describe GPU arenas, the client exports allocation-relative CUDA IPC views and preserves lease semantics, and PegaFlow owns storage resources and process-level cleanup. + +The runtime gates caught environment assumptions that are worth keeping outside the product interface: a server build needs a full protobuf toolchain, GLM5.2 code generation needs the validated TileLang version, and a NIC-less EP8 development node must disable DeepEP GIN. None of these became OpenInfer offload flags. + +Independent failure-path review is READY. Merge the PegaFlow protocol/server PR before the dependent OpenInfer client PR, then deploy matching revisions together. 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..07c41e886 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 = "5b73e5b", default-features = false } +pegaflow-proto = { git = "https://github.com/novitalabs/pegaflow.git", rev = "5b73e5b" } 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..3b5cc2e61 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,36 @@ 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(); + /// Stop this client instance and wait for the server to drain its GPU + /// queues and close every imported CUDA mapping. + 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 cleanup. Once its response arrives the + // server has drained workers and closed the imported mappings; aborting + // the liveness stream only removes its now-empty session entry. + self.session.take(); + } +} + +impl Drop for OffloadEngine { + fn drop(&mut self) { + self.shutdown(); } } @@ -882,15 +562,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 +581,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 +592,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 +604,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..6986e9376 --- /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, + )?; + 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, + block_stride_bytes: registration + .block_stride_bytes + .iter() + .map(|&value| as_u64(value, "block_stride_bytes")) + .collect::>()?, + }; + 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], +) -> Result, EngineError> { + if data_ptrs.len() != size_bytes.len() { + return Err(EngineError::InvalidArgument(format!( + "CUDA IPC export metadata length mismatch: pointers={}, sizes={}", + data_ptrs.len(), + size_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()) + .map(|(data_ptr, view_size)| { + 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")?, + }) + }) + .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 Date: Mon, 20 Jul 2026 03:54:14 -0400 Subject: [PATCH 2/5] fix(kv-offload): align PegaFlow wire schema Signed-off-by: xiaguan <751080330@qq.com> --- Cargo.lock | 6 +++--- .../runtime/external-pegaflow-server.md | 11 +++++++++-- openinfer-kv-offload/Cargo.toml | 4 ++-- openinfer-kv-offload/src/external.rs | 18 +++++++++--------- 4 files changed, 23 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 981d39b69..31984a97e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3374,7 +3374,7 @@ dependencies = [ [[package]] name = "pegaflow-common" version = "0.23.4" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=5b73e5b#5b73e5bb7ee30ca3f6f4765c21b80b4df634d6cd" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=933f548#933f548119b6d75772d9d92f8166379827d862e9" dependencies = [ "colored", "libc", @@ -3385,7 +3385,7 @@ dependencies = [ [[package]] name = "pegaflow-core" version = "0.23.4" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=5b73e5b#5b73e5bb7ee30ca3f6f4765c21b80b4df634d6cd" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=933f548#933f548119b6d75772d9d92f8166379827d862e9" dependencies = [ "ahash", "bytesize", @@ -3414,7 +3414,7 @@ dependencies = [ [[package]] name = "pegaflow-proto" version = "0.23.4" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=5b73e5b#5b73e5bb7ee30ca3f6f4765c21b80b4df634d6cd" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=933f548#933f548119b6d75772d9d92f8166379827d862e9" dependencies = [ "prost", "tonic", diff --git a/docs/subsystems/runtime/external-pegaflow-server.md b/docs/subsystems/runtime/external-pegaflow-server.md index 2c793b662..2799180fc 100644 --- a/docs/subsystems/runtime/external-pegaflow-server.md +++ b/docs/subsystems/runtime/external-pegaflow-server.md @@ -1,6 +1,6 @@ # External PegaFlow server for KV offload -> **TL;DR:** KV offload now has one ownership boundary: OpenInfer is a CUDA-IPC/RPC client selected by `--kv-offload-server`, while the external PegaFlow server owns storage and acknowledges transfer/drain barriers before exported GPU memory can be reused; Qwen3 and GLM5.2 cross-process restore gates pass on 8×H200. +> **TL;DR:** KV offload now has one ownership boundary: OpenInfer is a CUDA-IPC/RPC client selected by `--kv-offload-server`, while the external PegaFlow server owns storage and acknowledges transfer/drain barriers before exported GPU memory can be reused. The rewritten PR stack passes local compile/test gates; its recorded Qwen3 and GLM5.2 cross-process results predate the split and must be replayed on the final heads. > > **Last touched:** 2026-07 @@ -74,10 +74,17 @@ - Re-ran the cross-process gates after the lifecycle fixes. The raw Qwen round trip remained byte-exact at 0.43 s versus 0.44 s before the fixes. The hot Qwen3 executor gate took 3.32 s; CPU-only and combined GPU/CPU restores kept mean head-logprob deltas at 0.0105 and 0.0272 nat. - Re-ran GLM5.2 with a 520-block rank-local HBM pool. After rank-pinned displacement, the original 66-token request restored one host block, returned `cached_tokens=64`, reproduced the same first two tokens, and completed in 0.0438 s versus 0.257 s cold. A graceful stop drained all eight rank instances, closed 99 CUDA mappings per rank, emptied the server instance list, and returned all eight GPUs to zero allocated memory after the PegaFlow process exited. +### Step 5: upstream PR scope reduction + +- Review found that the PegaFlow PR combined three concerns: native CUDA IPC registration, drain-before-unmap ownership, and an unrelated prefetch-key isolation fix. +- The rewritten feature PR keeps native registration plus synchronous native load and flush, depends on a separate GPU-worker drain and unregister-ordering PR, and removes the prefetch-key change from this delivery chain. +- Native registration carries `block_stride_bytes` inside each CUDA IPC view instead of adding a second positionally coupled request array. The existing Python-wrapper request shape remains unchanged. +- The registration handshake has a distinct native-CUDA-IPC wire capability, so mixed old/new clients and servers fail before importing a CUDA mapping instead of silently falling back to a dense layout. + ## Debrief The embedded and external modes represented two owners for the same host-tier state and doubled every lifecycle decision. Removing the embedded mode left a stable boundary: model crates describe GPU arenas, the client exports allocation-relative CUDA IPC views and preserves lease semantics, and PegaFlow owns storage resources and process-level cleanup. The runtime gates caught environment assumptions that are worth keeping outside the product interface: a server build needs a full protobuf toolchain, GLM5.2 code generation needs the validated TileLang version, and a NIC-less EP8 development node must disable DeepEP GIN. None of these became OpenInfer offload flags. -Independent failure-path review is READY. Merge the PegaFlow protocol/server PR before the dependent OpenInfer client PR, then deploy matching revisions together. +Independent failure-path review is READY. Merge the PegaFlow drain PR, then the protocol/server PR, then the dependent OpenInfer client PR. Replay the Qwen3 and GLM5.2 cross-process gates on those final heads before deployment. diff --git a/openinfer-kv-offload/Cargo.toml b/openinfer-kv-offload/Cargo.toml index 07c41e886..49b0834ab 100644 --- a/openinfer-kv-offload/Cargo.toml +++ b/openinfer-kv-offload/Cargo.toml @@ -7,8 +7,8 @@ edition = "2024" [dependencies] # 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 = "5b73e5b", default-features = false } -pegaflow-proto = { git = "https://github.com/novitalabs/pegaflow.git", rev = "5b73e5b" } +pegaflow-core = { git = "https://github.com/novitalabs/pegaflow.git", rev = "933f548", default-features = false } +pegaflow-proto = { git = "https://github.com/novitalabs/pegaflow.git", rev = "933f548" } openinfer-kv-cache = { workspace = true } cudarc = { workspace = true } half = { workspace = true } diff --git a/openinfer-kv-offload/src/external.rs b/openinfer-kv-offload/src/external.rs index 6986e9376..76e693092 100644 --- a/openinfer-kv-offload/src/external.rs +++ b/openinfer-kv-offload/src/external.rs @@ -91,6 +91,7 @@ impl ExternalClient { 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")?; @@ -148,11 +149,6 @@ impl ExternalClient { transfer_mode: TransferMode::Direct as i32, page_first: registration.page_first, cuda_ipc_tensors, - block_stride_bytes: registration - .block_stride_bytes - .iter() - .map(|&value| as_u64(value, "block_stride_bytes")) - .collect::>()?, }; let mut client = self.client.clone(); let response = match tokio::time::timeout( @@ -394,12 +390,14 @@ 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() { + 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={}", + "CUDA IPC export metadata length mismatch: pointers={}, sizes={}, block strides={}", data_ptrs.len(), - size_bytes.len() + size_bytes.len(), + block_stride_bytes.len() ))); } let device = usize::try_from(device_id).map_err(|_| { @@ -416,7 +414,8 @@ fn export_cuda_ipc_tensors( .iter() .copied() .zip(size_bytes.iter().copied()) - .map(|(data_ptr, view_size)| { + .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 @@ -474,6 +473,7 @@ fn export_cuda_ipc_tensors( handle, offset_bytes: offset, size_bytes: as_u64(view_size, "size_bytes")?, + block_stride_bytes: as_u64(block_stride, "block_stride_bytes")?, }) }) .collect() From 9602cb239da308b217d765fc135815b86e02d5d0 Mon Sep 17 00:00:00 2001 From: xiaguan <751080330@qq.com> Date: Mon, 20 Jul 2026 04:25:34 -0400 Subject: [PATCH 3/5] docs(kv-offload): reflect combined PegaFlow PR Signed-off-by: xiaguan <751080330@qq.com> --- docs/subsystems/runtime/external-pegaflow-server.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/subsystems/runtime/external-pegaflow-server.md b/docs/subsystems/runtime/external-pegaflow-server.md index 2799180fc..639cae083 100644 --- a/docs/subsystems/runtime/external-pegaflow-server.md +++ b/docs/subsystems/runtime/external-pegaflow-server.md @@ -1,6 +1,6 @@ # External PegaFlow server for KV offload -> **TL;DR:** KV offload now has one ownership boundary: OpenInfer is a CUDA-IPC/RPC client selected by `--kv-offload-server`, while the external PegaFlow server owns storage and acknowledges transfer/drain barriers before exported GPU memory can be reused. The rewritten PR stack passes local compile/test gates; its recorded Qwen3 and GLM5.2 cross-process results predate the split and must be replayed on the final heads. +> **TL;DR:** KV offload now has one ownership boundary: OpenInfer is a CUDA-IPC/RPC client selected by `--kv-offload-server`, while the external PegaFlow server owns storage and acknowledges transfer/drain barriers before exported GPU memory can be reused. The rewritten PRs pass local compile/test gates; the recorded Qwen3 and GLM5.2 cross-process results predate the final scope and must be replayed on the final heads. > > **Last touched:** 2026-07 @@ -74,10 +74,10 @@ - Re-ran the cross-process gates after the lifecycle fixes. The raw Qwen round trip remained byte-exact at 0.43 s versus 0.44 s before the fixes. The hot Qwen3 executor gate took 3.32 s; CPU-only and combined GPU/CPU restores kept mean head-logprob deltas at 0.0105 and 0.0272 nat. - Re-ran GLM5.2 with a 520-block rank-local HBM pool. After rank-pinned displacement, the original 66-token request restored one host block, returned `cached_tokens=64`, reproduced the same first two tokens, and completed in 0.0438 s versus 0.257 s cold. A graceful stop drained all eight rank instances, closed 99 CUDA mappings per rank, emptied the server instance list, and returned all eight GPUs to zero allocated memory after the PegaFlow process exited. -### Step 5: upstream PR scope reduction +### Step 5: upstream PR scope - Review found that the PegaFlow PR combined three concerns: native CUDA IPC registration, drain-before-unmap ownership, and an unrelated prefetch-key isolation fix. -- The rewritten feature PR keeps native registration plus synchronous native load and flush, depends on a separate GPU-worker drain and unregister-ordering PR, and removes the prefetch-key change from this delivery chain. +- The rewritten PegaFlow PR keeps native registration, synchronous native load and flush, and the required GPU-worker drain/unregister ordering together; the unrelated prefetch-key change is removed. - Native registration carries `block_stride_bytes` inside each CUDA IPC view instead of adding a second positionally coupled request array. The existing Python-wrapper request shape remains unchanged. - The registration handshake has a distinct native-CUDA-IPC wire capability, so mixed old/new clients and servers fail before importing a CUDA mapping instead of silently falling back to a dense layout. @@ -87,4 +87,4 @@ The embedded and external modes represented two owners for the same host-tier st The runtime gates caught environment assumptions that are worth keeping outside the product interface: a server build needs a full protobuf toolchain, GLM5.2 code generation needs the validated TileLang version, and a NIC-less EP8 development node must disable DeepEP GIN. None of these became OpenInfer offload flags. -Independent failure-path review is READY. Merge the PegaFlow drain PR, then the protocol/server PR, then the dependent OpenInfer client PR. Replay the Qwen3 and GLM5.2 cross-process gates on those final heads before deployment. +Independent failure-path review is READY. Merge the combined PegaFlow protocol/lifecycle PR before the dependent OpenInfer client PR. Replay the Qwen3 and GLM5.2 cross-process gates on those final heads before deployment. From e30010ac4c7fbe1b16311fc5e9a3c8ca6f327a97 Mon Sep 17 00:00:00 2001 From: xiaguan <751080330@qq.com> Date: Mon, 20 Jul 2026 05:40:22 -0400 Subject: [PATCH 4/5] chore(kv-offload): pin minimal pegaflow protocol Signed-off-by: xiaguan <751080330@qq.com> --- Cargo.lock | 24 +++++++++---------- .../runtime/external-pegaflow-server.md | 16 +++++-------- openinfer-kv-offload/Cargo.toml | 4 ++-- openinfer-kv-offload/src/engine.rs | 9 ++++--- 4 files changed, 24 insertions(+), 29 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 31984a97e..82cd20846 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -651,7 +651,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -1318,7 +1318,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2199,7 +2199,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2253,7 +2253,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3374,7 +3374,7 @@ dependencies = [ [[package]] name = "pegaflow-common" version = "0.23.4" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=933f548#933f548119b6d75772d9d92f8166379827d862e9" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=29c7fba#29c7fbad851d660ac8c11e9a0a38d5f4b2b4c57b" dependencies = [ "colored", "libc", @@ -3385,7 +3385,7 @@ dependencies = [ [[package]] name = "pegaflow-core" version = "0.23.4" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=933f548#933f548119b6d75772d9d92f8166379827d862e9" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=29c7fba#29c7fbad851d660ac8c11e9a0a38d5f4b2b4c57b" dependencies = [ "ahash", "bytesize", @@ -3414,7 +3414,7 @@ dependencies = [ [[package]] name = "pegaflow-proto" version = "0.23.4" -source = "git+https://github.com/novitalabs/pegaflow.git?rev=933f548#933f548119b6d75772d9d92f8166379827d862e9" +source = "git+https://github.com/novitalabs/pegaflow.git?rev=29c7fba#29c7fbad851d660ac8c11e9a0a38d5f4b2b4c57b" dependencies = [ "prost", "tonic", @@ -3728,7 +3728,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ "heck", - "itertools 0.10.5", + "itertools 0.13.0", "log", "multimap", "petgraph", @@ -3749,7 +3749,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.13.0", "proc-macro2", "quote", "syn 2.0.117", @@ -4251,7 +4251,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4999,7 +4999,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -6315,7 +6315,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/docs/subsystems/runtime/external-pegaflow-server.md b/docs/subsystems/runtime/external-pegaflow-server.md index 639cae083..51a1594f0 100644 --- a/docs/subsystems/runtime/external-pegaflow-server.md +++ b/docs/subsystems/runtime/external-pegaflow-server.md @@ -1,6 +1,6 @@ # External PegaFlow server for KV offload -> **TL;DR:** KV offload now has one ownership boundary: OpenInfer is a CUDA-IPC/RPC client selected by `--kv-offload-server`, while the external PegaFlow server owns storage and acknowledges transfer/drain barriers before exported GPU memory can be reused. The rewritten PRs pass local compile/test gates; the recorded Qwen3 and GLM5.2 cross-process results predate the final scope and must be replayed on the final heads. +> **TL;DR:** KV offload now has one ownership boundary: OpenInfer is a CUDA-IPC/RPC client selected by `--kv-offload-server`, while the external PegaFlow server owns storage. The minimal protocol/server PR passes local compile/test gates; the recorded Qwen3 and GLM5.2 cross-process results predate the final scope and must be replayed on the final heads. > > **Last touched:** 2026-07 @@ -63,21 +63,17 @@ ### Step 4: failure and multi-tenant review -- Scoped PegaFlow's asynchronous deeper-tier prefetch state by `(instance_id, req_id)`. A regression test keeps two equal request IDs in separate active entries and proves that polling one instance cannot consume the other's result. - Replaced process-ID instance suffixes with UUIDv4 and made native content namespaces explicit. Native Qwen3 and GLM5.2 launches now require `--kv-offload-namespace`, a stable checkpoint/deployment identity; vLLM compatibility continues to use the connector namespace. - Deadlines remain on health, query, release, and flush. Ownership RPCs do not carry a cancelable gRPC deadline because an accepted CUDA DMA cannot be canceled with its handler; a 30-second local watchdog or an indeterminate transport failure is fail-stop. Native load completion is returned by the Load RPC instead of a shared-memory poll. -- Added an owned session guard and explicit client shutdown. Shutdown waits for accepted saves, sends `UnregisterContext`, and releases the session only after the server acknowledges both GPU queue drain and CUDA mapping close. Failure to obtain that ownership acknowledgement is fail-stop. -- PegaFlow unregister keeps a closing instance addressable to concurrent cleanup calls, rejects new transfers, places barriers behind every accepted load/save on every registered device, and removes the instance only after the shared drain completes. The worker threads retain join handles for deterministic teardown. -- Registration and cleanup are serialized per instance and by session generation. A canceled registry reply rolls back a newly published CUDA mapping; session cleanup still closes registry state when engine registration never completed. +- Added an owned session guard and explicit client shutdown. Shutdown flushes accepted saves before `UnregisterContext`; native Load returns terminal completion, so callers retain responsibility for completing load handles before engine teardown. - Kept legacy-exportable allocation behind the offload option so an offload-disabled launch uses the original allocator. -- Failure-path local gates pass: PegaFlow server 28/28, PegaFlow core 126/126 with one GPU-only test ignored, and OpenInfer KV-offload 9/9. -- Re-ran the cross-process gates after the lifecycle fixes. The raw Qwen round trip remained byte-exact at 0.43 s versus 0.44 s before the fixes. The hot Qwen3 executor gate took 3.32 s; CPU-only and combined GPU/CPU restores kept mean head-logprob deltas at 0.0105 and 0.0272 nat. -- Re-ran GLM5.2 with a 520-block rank-local HBM pool. After rank-pinned displacement, the original 66-token request restored one host block, returned `cached_tokens=64`, reproduced the same first two tokens, and completed in 0.0438 s versus 0.257 s cold. A graceful stop drained all eight rank instances, closed 99 CUDA mappings per rank, emptied the server instance list, and returned all eight GPUs to zero allocated memory after the PegaFlow process exited. +- Final-scope local gates pass: PegaFlow server 23/23, PegaFlow core 124/124 with one GPU-only test ignored, and OpenInfer KV-offload 9/9. +- A discarded lifecycle prototype also passed the raw Qwen, Qwen3 executor, and GLM5.2 cross-process gates. Those results are historical only; the final minimal heads still require replay. ### Step 5: upstream PR scope - Review found that the PegaFlow PR combined three concerns: native CUDA IPC registration, drain-before-unmap ownership, and an unrelated prefetch-key isolation fix. -- The rewritten PegaFlow PR keeps native registration, synchronous native load and flush, and the required GPU-worker drain/unregister ordering together; the unrelated prefetch-key change is removed. +- The rewritten PegaFlow PR contains only native registration plus synchronous native load and flush. GPU-worker lifecycle and the unrelated prefetch-key change are outside this PR. - Native registration carries `block_stride_bytes` inside each CUDA IPC view instead of adding a second positionally coupled request array. The existing Python-wrapper request shape remains unchanged. - The registration handshake has a distinct native-CUDA-IPC wire capability, so mixed old/new clients and servers fail before importing a CUDA mapping instead of silently falling back to a dense layout. @@ -87,4 +83,4 @@ The embedded and external modes represented two owners for the same host-tier st The runtime gates caught environment assumptions that are worth keeping outside the product interface: a server build needs a full protobuf toolchain, GLM5.2 code generation needs the validated TileLang version, and a NIC-less EP8 development node must disable DeepEP GIN. None of these became OpenInfer offload flags. -Independent failure-path review is READY. Merge the combined PegaFlow protocol/lifecycle PR before the dependent OpenInfer client PR. Replay the Qwen3 and GLM5.2 cross-process gates on those final heads before deployment. +Independent failure-path review is READY. Merge the PegaFlow protocol/server PR before the dependent OpenInfer client PR. Replay the Qwen3 and GLM5.2 cross-process gates on those final heads before deployment. diff --git a/openinfer-kv-offload/Cargo.toml b/openinfer-kv-offload/Cargo.toml index 49b0834ab..bd1591a40 100644 --- a/openinfer-kv-offload/Cargo.toml +++ b/openinfer-kv-offload/Cargo.toml @@ -7,8 +7,8 @@ edition = "2024" [dependencies] # 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 = "933f548", default-features = false } -pegaflow-proto = { git = "https://github.com/novitalabs/pegaflow.git", rev = "933f548" } +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 } diff --git a/openinfer-kv-offload/src/engine.rs b/openinfer-kv-offload/src/engine.rs index 3b5cc2e61..94fa6f375 100644 --- a/openinfer-kv-offload/src/engine.rs +++ b/openinfer-kv-offload/src/engine.rs @@ -525,8 +525,8 @@ impl OffloadEngine { }); } - /// Stop this client instance and wait for the server to drain its GPU - /// queues and close every imported CUDA mapping. + /// 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; @@ -545,9 +545,8 @@ impl OffloadEngine { ); std::process::abort(); } - // The explicit unregister owns cleanup. Once its response arrives the - // server has drained workers and closed the imported mappings; aborting - // the liveness stream only removes its now-empty session entry. + // The explicit unregister owns mapping cleanup. Aborting the liveness + // stream afterwards only removes its now-empty session entry. self.session.take(); } } From 0e3d01cc007b65716c52dc4e74b4f6bfc5c43db3 Mon Sep 17 00:00:00 2001 From: xiaguan <751080330@qq.com> Date: Mon, 20 Jul 2026 06:19:14 -0400 Subject: [PATCH 5/5] docs(kv-offload): remove discarded design history Signed-off-by: xiaguan <751080330@qq.com> --- .../runtime/external-pegaflow-server.md | 96 +++++-------------- 1 file changed, 26 insertions(+), 70 deletions(-) diff --git a/docs/subsystems/runtime/external-pegaflow-server.md b/docs/subsystems/runtime/external-pegaflow-server.md index 51a1594f0..06e368f18 100644 --- a/docs/subsystems/runtime/external-pegaflow-server.md +++ b/docs/subsystems/runtime/external-pegaflow-server.md @@ -1,86 +1,42 @@ # External PegaFlow server for KV offload -> **TL;DR:** KV offload now has one ownership boundary: OpenInfer is a CUDA-IPC/RPC client selected by `--kv-offload-server`, while the external PegaFlow server owns storage. The minimal protocol/server PR passes local compile/test gates; the recorded Qwen3 and GLM5.2 cross-process results predate the final scope and must be replayed on the final heads. +> **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 -## Preparation +## Ownership boundary -- **Read**: - - `docs/index.md` — routes this cross-model change to the runtime subsystem. - - `docs/subsystems/runtime/pegaflow-offload-integration.md` — records the current in-process ownership model, transfer ordering, leases, and Qwen3 CPU-hit gate. - - `docs/models/qwen3/prefix-cache.md` — establishes full-block matching, the one-token recompute boundary, and cache-lifetime invariants. - - `docs/models/qwen3/serving-performance.md` — retains the current host-tier offload baseline and pure-L2 measurement shape. - - `docs/models/glm52/pd-m2-execution.md` — establishes GLM5.2's 99-arena page-first layout, PegaFlow version boundary, and strict restore semantics. - - `openinfer-kv-offload/src/{lib.rs,engine.rs}` — `OffloadEngine` currently owns an in-process `PegaEngine`; the scheduler-facing save/query/load API is already shared by the model integrations. - - `openinfer-qwen3/src/{lib.rs,executor.rs}` — Qwen3 constructs one offload engine around its fused page-first `KvBuffer` and already has asynchronous prefetch admission. - - `openinfer-glm52/src/lib.rs` — GLM5.2 shares one `OffloadHost` across eight rank instances and registers 78 MLA plus 21 index-K arenas per rank. -- **Relevant history**: - - PegaFlow PRs #331/#333 made strided page-first registration and in-process completion usable by OpenInfer. - - OpenInfer #316 proved Qwen3 host-tier save/load correctness with an embedded pool. - - OpenInfer #600/#657 proved GLM5.2 host offload and cross-node target-KV restore with an embedded pool. - - The development cluster already runs a PegaFlow DaemonSet on selected nodes; the first remote step is read-only inventory, and the production daemon must not be modified for bring-up. -- **Plan**: - 1. On the reserved 8×H200 development node, inspect the PegaFlow DaemonSet endpoint, version, server flags, CUDA IPC registration contract, available model paths, and existing OpenInfer/PegaFlow worktrees without changing the running daemon. - 2. Replace the embedded `PegaEngine` with an external PegaFlow gRPC client while preserving the scheduler-visible query/lease/load/save semantics. - 3. Reuse the server's CUDA IPC registration when it can represent Qwen3's strided fused buffer and GLM5.2's named multi-arena/page-first layout. If it cannot, extend PegaFlow's protocol and server with the smallest layout-general registration primitive, verify its own integration tests, and prepare an upstream PegaFlow PR. - 4. Expose one explicit OpenInfer server configuration for the external endpoint; fail at startup on incompatible combinations, version/layout mismatch, missing server health, or registration failure. - 5. Gate Qwen3 with an external-server register → save → HBM eviction → load path, byte/logit parity, and a repeated-prompt serving check that proves a host-tier hit rather than a GPU-prefix hit. - 6. Gate GLM5.2 with all eight ranks registered against one external server, restore of both MLA and index-K arenas, token parity after forced host-tier restore, and a repeated-prompt serving check with no local GPU-prefix hit. - 7. Run formatting, focused unit/integration tests, and cross-process correctness/latency measurements before review. Record deployment commands without internal hostnames or private paths, then submit the required OpenInfer and PegaFlow PRs. -- **Risks / open questions**: - - Cross-process registration requires CUDA IPC handles rather than raw device pointers; handle lifetime, device identity, and the server's embedded Python/CUDA context must be verified on the real daemon. - - The existing gRPC schema may assume vLLM tensor metadata and may not expose GLM5.2's named arena ordering. A protocol change is acceptable, but silent layout coercion is not. - - Qwen3 uses one engine per device while GLM5.2 uses one host shared across eight rank registrations. The external client must preserve that ownership distinction without introducing a generic abstraction above the transfer layer. - - The daemon's pinned pool is shared infrastructure. Bring-up needs isolated namespaces and must not evict, stop, or reconfigure unrelated cache contents. - - Correctness evidence must force HBM misses; a warm request that hits OpenInfer's local prefix cache does not prove external offload. +- 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. -## Execution Log +## Wire contract -### Step 1: development-cluster inventory +- 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. -- Connected to the reserved 8×H200 development node and inspected the cluster without modifying workloads. -- The cluster runs a host-networked PegaFlow DaemonSet with host IPC, all GPUs initialized, and a MetaServer connection. Its deployed protocol predates the native Rust registration added here, so compatibility must be checked rather than assumed. -- The reserved node was occupied by unrelated inference processes, which were left untouched. Verification moved to an isolated 8×H200 development pod with local Qwen3 and GLM5.2 checkpoints. -- The production daemon already owns CUDA contexts and host hugepages. External mode must use that pool instead of creating another embedded pool. +## Model layouts -### Step 2: external-only ownership boundary +- 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. -- Removed the embedded `PegaEngine`, local pinned-pool construction, local P2P service, and their CLI knobs. `--kv-offload-server` is now the sole enablement switch; server resource configuration stays outside OpenInfer. -- Added a native PegaFlow registration payload carrying CUDA IPC handles, per-view offsets/sizes, and optional block strides. This represents both Qwen3's page-interleaved fused allocation and GLM5.2's MLA/index-K arenas without Python tensor wrappers. -- Added a PegaFlow `Flush` RPC for the existing save/registration visibility barrier. Native loads now await the server's GPU worker through the Load RPC; the Python connector retains its shared-memory completion path. -- KV allocations use legacy-exportable CUDA allocation only when external offload is enabled. The default path retains the stream-ordered allocator. A direct GPU probe confirmed that `cudaIpcGetMemHandle` rejects `cudaMallocAsync` allocations and accepts `cudaMalloc` allocations. -- Local gates passed: PegaFlow server unit tests (26), OpenInfer KV-offload unit tests (8), Qwen3 unit tests (75), default-feature server compile/tests, and a GLM5.2 feature compile using NCCL 2.30.7. The local GPU is not SM90, so real GLM5.2 kernels and behavior still require the 8×H200 target. +## Validation -### Step 3: 8×H200 cross-process gates +Current PR-head local gates are green: -- Built the PegaFlow server without RDMA and built the GLM5.2 OpenInfer feature on the target SM90a system. The build exercised TileLang 0.1.12 code generation, FlashMLA, DeepGEMM, NCCL 2.30.7, and the final server link. -- The raw Qwen page-first gate saved three blocks through a separate PegaFlow process, loaded them into different HBM block ids, and compared every layer/segment byte-for-byte. The negative-control block stayed zero. The test body took 0.44 s; a hot build plus test took 1.20 s. -- The Qwen3-4B executor gate forced both a CPU-only prefix restore and a combined GPU+CPU restore. Warm/cold head-logprob mean deltas were 0.0105 and 0.0242 nat; maxima were 0.0652 and 0.0632. The test body took 5.61 s. -- GLM5.2 EP8 registered eight rank instances with 99 arenas each (78 MLA plus 21 index-K) through native CUDA IPC. A 66-token cold request saved one logical block across all 99 arenas: 3,452,160 bytes, with the PegaFlow D2H worker reporting 0.50 ms. -- Nine distinct 3,900-token requests pinned to the same data-parallel rank displaced the original prefix from its 520-block HBM LRU. Replaying the original request restored one block from PegaFlow, returned `cached_tokens=64`, reproduced the same first two output tokens, and logged `GLM5.2 host-tier restore: 1 blocks committed`. PegaFlow performed 99 H2D copies (3,452,160 bytes) in 0.62 ms. End-to-end request time was 0.256 s cold and 0.044 s after host-tier restore; this is a correctness-gate observation, not a serving benchmark. -- The isolated OpenInfer and PegaFlow processes were stopped after the gate; all eight GPUs returned to zero allocated memory. The development pod and source mounts remain available. +- 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. -### Step 4: failure and multi-tenant review +Cross-process evidence collected during bring-up established the layout and transfer contract: -- Replaced process-ID instance suffixes with UUIDv4 and made native content namespaces explicit. Native Qwen3 and GLM5.2 launches now require `--kv-offload-namespace`, a stable checkpoint/deployment identity; vLLM compatibility continues to use the connector namespace. -- Deadlines remain on health, query, release, and flush. Ownership RPCs do not carry a cancelable gRPC deadline because an accepted CUDA DMA cannot be canceled with its handler; a 30-second local watchdog or an indeterminate transport failure is fail-stop. Native load completion is returned by the Load RPC instead of a shared-memory poll. -- Added an owned session guard and explicit client shutdown. Shutdown flushes accepted saves before `UnregisterContext`; native Load returns terminal completion, so callers retain responsibility for completing load handles before engine teardown. -- Kept legacy-exportable allocation behind the offload option so an offload-disabled launch uses the original allocator. -- Final-scope local gates pass: PegaFlow server 23/23, PegaFlow core 124/124 with one GPU-only test ignored, and OpenInfer KV-offload 9/9. -- A discarded lifecycle prototype also passed the raw Qwen, Qwen3 executor, and GLM5.2 cross-process gates. Those results are historical only; the final minimal heads still require replay. +- 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. -### Step 5: upstream PR scope - -- Review found that the PegaFlow PR combined three concerns: native CUDA IPC registration, drain-before-unmap ownership, and an unrelated prefetch-key isolation fix. -- The rewritten PegaFlow PR contains only native registration plus synchronous native load and flush. GPU-worker lifecycle and the unrelated prefetch-key change are outside this PR. -- Native registration carries `block_stride_bytes` inside each CUDA IPC view instead of adding a second positionally coupled request array. The existing Python-wrapper request shape remains unchanged. -- The registration handshake has a distinct native-CUDA-IPC wire capability, so mixed old/new clients and servers fail before importing a CUDA mapping instead of silently falling back to a dense layout. - -## Debrief - -The embedded and external modes represented two owners for the same host-tier state and doubled every lifecycle decision. Removing the embedded mode left a stable boundary: model crates describe GPU arenas, the client exports allocation-relative CUDA IPC views and preserves lease semantics, and PegaFlow owns storage resources and process-level cleanup. - -The runtime gates caught environment assumptions that are worth keeping outside the product interface: a server build needs a full protobuf toolchain, GLM5.2 code generation needs the validated TileLang version, and a NIC-less EP8 development node must disable DeepEP GIN. None of these became OpenInfer offload flags. - -Independent failure-path review is READY. Merge the PegaFlow protocol/server PR before the dependent OpenInfer client PR. Replay the Qwen3 and GLM5.2 cross-process gates on those final heads before deployment. +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.