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

.github/workflows/ci.yml

@@ -142,6 +142,11 @@ jobs:
           persist-credentials: false
       - name: Run ruff check
         uses: astral-sh/ruff-action@v3
+        with:
+          # Unpinned, this floats to the newest ruff and every new lint rule
+          # breaks CI on untouched files (0.16.0 did exactly that). Bump
+          # deliberately, together with the fixes it demands.
+          version: "0.15.4"
 
   python-unit-tests:
     name: python unit tests

pegaflow-core/src/backing/transfer_lock_guard.rs

@@ -81,11 +81,12 @@ mod tests {
 
     use pegaflow_proto::proto::engine::engine_server::{Engine, EngineServer};
     use pegaflow_proto::proto::engine::{
-        HealthRequest, HealthResponse, LoadRequest, LoadResponse, QueryBlocksForTransferRequest,
-        QueryBlocksForTransferResponse, QueryRequest, QueryResponse, RdmaHandshakeRequest,
-        RdmaHandshakeResponse, RegisterContextRequest, RegisterContextResponse, ReleaseRequest,
-        ReleaseResponse, ReleaseTransferLockResponse, SaveRequest, SaveResponse, SessionEvent,
-        SessionRequest, ShutdownRequest, ShutdownResponse, UnregisterRequest, UnregisterResponse,
+        FlushRequest, FlushResponse, HealthRequest, HealthResponse, LoadRequest, LoadResponse,
+        QueryBlocksForTransferRequest, QueryBlocksForTransferResponse, QueryRequest, QueryResponse,
+        RdmaHandshakeRequest, RdmaHandshakeResponse, RegisterContextRequest,
+        RegisterContextResponse, ReleaseRequest, ReleaseResponse, ReleaseTransferLockResponse,
+        SaveRequest, SaveResponse, SessionEvent, SessionRequest, ShutdownRequest, ShutdownResponse,
+        UnregisterRequest, UnregisterResponse,
     };
     use tokio_stream::wrappers::TcpListenerStream;
     use tonic::transport::Endpoint;
@@ -125,6 +126,12 @@ mod tests {
         ) -> Result<Response<HealthResponse>, Status> {
             Err(Status::unimplemented("stub"))
         }
+        async fn flush(
+            &self,
+            _request: Request<FlushRequest>,
+        ) -> Result<Response<FlushResponse>, Status> {
+            Err(Status::unimplemented("stub"))
+        }
         async fn register_context_batch(
             &self,
             _request: Request<RegisterContextRequest>,

pegaflow-core/src/internode/p2p_service.rs

@@ -20,12 +20,12 @@ use tonic::{Request, Response, Status, async_trait};
 
 use pegaflow_proto::proto::engine::engine_server::{Engine, EngineServer};
 use pegaflow_proto::proto::engine::{
-    HealthRequest, HealthResponse, LoadRequest, LoadResponse, QueryBlocksForTransferRequest,
-    QueryBlocksForTransferResponse, QueryRequest, QueryResponse, RdmaHandshakeRequest,
-    RdmaHandshakeResponse, RegisterContextRequest, RegisterContextResponse, ReleaseRequest,
-    ReleaseResponse, ReleaseTransferLockRequest, ReleaseTransferLockResponse, ResponseStatus,
-    SaveRequest, SaveResponse, SessionEvent, SessionRequest, ShutdownRequest, ShutdownResponse,
-    TransferBlockInfo, TransferSlotInfo, UnregisterRequest, UnregisterResponse,
+    FlushRequest, FlushResponse, HealthRequest, HealthResponse, LoadRequest, LoadResponse,
+    QueryBlocksForTransferRequest, QueryBlocksForTransferResponse, QueryRequest, QueryResponse,
+    RdmaHandshakeRequest, RdmaHandshakeResponse, RegisterContextRequest, RegisterContextResponse,
+    ReleaseRequest, ReleaseResponse, ReleaseTransferLockRequest, ReleaseTransferLockResponse,
+    ResponseStatus, SaveRequest, SaveResponse, SessionEvent, SessionRequest, ShutdownRequest,
+    ShutdownResponse, TransferBlockInfo, TransferSlotInfo, UnregisterRequest, UnregisterResponse,
 };
 
 use crate::{LayerBlock, PegaEngine};
@@ -246,6 +246,13 @@ impl Engine for P2pTransferService {
         Self::not_served("load")
     }
 
+    async fn flush(
+        &self,
+        _request: Request<FlushRequest>,
+    ) -> Result<Response<FlushResponse>, Status> {
+        Self::not_served("flush")
+    }
+
     async fn query_prefetch(
         &self,
         _request: Request<QueryRequest>,

pegaflow-core/src/storage/mod.rs

@@ -5,8 +5,6 @@ pub(crate) mod transfer_lock;
 mod write_path;
 
 use bytesize::ByteSize;
-#[cfg(not(feature = "rdma"))]
-use log::warn;
 use log::{debug, info, warn};
 use std::collections::HashSet;
 use std::num::NonZeroU64;

pegaflow-proto/proto/engine.proto

@@ -46,10 +46,32 @@ message RegisterContextRequest {
   // Set by the connector for MLA; the engine derives per-layer page offsets
   // from the registered layouts. See spec.md.
   bool page_first = 18;
+  // Native (non-Python) registration: per-layer views into one KV arena the
+  // server allocates on `device_id` and shares back as a CUDA IPC handle.
+  // Empty for Python registrations.
+  repeated NativeKvTensor native_kv_tensors = 19;
+  // Total arena size the server allocates for a native registration; every
+  // layer view must fit inside it. Zero for Python registrations.
+  uint64 native_alloc_size = 20;
+}
+
+// One layer's view into the server-allocated KV arena.
+message NativeKvTensor {
+  uint64 offset_bytes = 1;
+  uint64 size_bytes = 2;
+  // Distance between consecutive blocks of this layer inside the arena
+  // (fused allocations interleave layers, so blocks are not contiguous).
+  uint64 block_stride_bytes = 3;
 }
 
 message RegisterContextResponse {
   ResponseStatus status = 1;
+  // CUDA IPC handle (CUipcMemHandle, 64 bytes) of the server-allocated arena.
+  // The client imports it with cuIpcOpenMemHandle and uses the mapping as its
+  // KV cache. Present only for native registrations; the arena is freed when
+  // the instance unregisters, so the client must not touch the mapping after
+  // unregister_context returns.
+  bytes arena_ipc_handle = 2;
 }
 
 message SaveLayer {
@@ -77,6 +99,9 @@ message LoadRequest {
   string load_state_shm = 4;
   repeated string layer_names = 5;
   repeated LeaseLoad loads = 6;
+  // Return only after the H2D copies complete. Native clients must set this:
+  // they have no load_state_shm completion flag to poll.
+  bool wait_for_completion = 7;
 }
 
 message LeaseLoad {
@@ -88,6 +113,12 @@ message LoadResponse {
   ResponseStatus status = 1;
 }
 
+message FlushRequest {}
+
+message FlushResponse {
+  ResponseStatus status = 1;
+}
+
 message QueryRequest {
   string instance_id = 1;
   repeated bytes block_hashes = 2;
@@ -207,6 +238,9 @@ service Engine {
   rpc RegisterContextBatch(RegisterContextRequest) returns (RegisterContextResponse);
   rpc Save(SaveRequest) returns (SaveResponse);
   rpc Load(LoadRequest) returns (LoadResponse);
+  // Durability barrier: returns once every accepted save is cache-visible and
+  // its MetaServer registration has been attempted.
+  rpc Flush(FlushRequest) returns (FlushResponse);
   // Prefetch query: check memory + trigger SSD prefetch or cross-node RDMA fetch for missing blocks
   rpc QueryPrefetch(QueryRequest) returns (QueryResponse);
   rpc Release(ReleaseRequest) returns (ReleaseResponse);

pegaflow-proto/src/lib.rs

@@ -1,3 +1,8 @@
+/// Exact wire capability version. Registration is fenced on this string, so a
+/// client built before the native-arena contract fails before any memory
+/// changes hands.
+pub const VERSION: &str = concat!(env!("CARGO_PKG_VERSION"), "+native-arena-v1");
+
 pub mod proto {
     #[allow(
         clippy::allow_attributes_without_reason,

pegaflow-server/src/http_server.rs

@@ -82,8 +82,13 @@ async fn cleanup_handler(
 ) -> impl IntoResponse {
     match query.id {
         None => {
-            let removed_tensors = state.registry.clear().await;
+            // Engine first, then a save-pipeline barrier, then the registry
+            // drop: the registry drop frees native arenas (cuMemFree), so the
+            // engine must forget its raw pointers and drain accepted saves
+            // before the memory behind them goes away.
             let removed_instances = state.engine.unregister_all_instances();
+            state.engine.flush_saves().await;
+            let removed_tensors = state.registry.clear().await;
 
             if !removed_instances.is_empty() || removed_tensors > 0 {
                 warn!(
@@ -104,8 +109,11 @@ async fn cleanup_handler(
             )
         }
         Some(instance_id) => {
+            // Same ordering as the clear-all branch above.
+            let engine_result = state.engine.unregister_instance(&instance_id);
+            state.engine.flush_saves().await;
             let removed_tensors = state.registry.drop_instance(instance_id.clone()).await;
-            match state.engine.unregister_instance(&instance_id) {
+            match engine_result {
                 Ok(()) => {
                     warn!(
                         "Cleanup instance {}: {} CUDA tensor(s) released",

pegaflow-server/src/lib.rs

@@ -1,6 +1,7 @@
 mod check_cuda_version;
 pub mod http_server;
 pub mod metric;
+mod native_arena;
 pub mod proto;
 pub mod registry;
 pub mod service;
@@ -61,6 +62,11 @@ pub struct Cli {
     #[arg(long, value_delimiter = ',')]
     pub devices: Vec<i32>,
 
+    /// Initialize the torch CUDA registry for Python (vLLM connector) clients.
+    /// `false` runs torch-free and serves native clients only.
+    #[arg(long, default_value_t = true, action = clap::ArgAction::Set)]
+    pub python_registry: bool,
+
     /// Pinned memory pool size (supports units: kb, mb, gb, tb)
     /// Examples: "10gb", "500mb", "1tb"
     #[arg(long, default_value = "30gb", value_parser = parse_memory_size)]
@@ -322,6 +328,30 @@ fn detect_cuda_devices() -> Result<Vec<i32>, std::io::Error> {
     })
 }
 
+/// Torch-free device enumeration for native-only deployments.
+fn detect_native_cuda_devices() -> Result<Vec<i32>, std::io::Error> {
+    let count = cudarc::driver::CudaContext::device_count()
+        .map_err(|e| std::io::Error::other(format!("cudarc device_count: {e}")))?;
+    Ok((0..count).collect())
+}
+
+/// Torch-free CUDA context initialization for native-only deployments.
+fn init_cudarc_cuda(device_ids: &[i32]) -> Result<(), std::io::Error> {
+    if device_ids.is_empty() {
+        return Err(std::io::Error::other("no CUDA devices to initialize"));
+    }
+    for &device_id in device_ids {
+        let ordinal = usize::try_from(device_id)
+            .map_err(|_| std::io::Error::other(format!("device_id {device_id} must be >= 0")))?;
+        let ctx = cudarc::driver::CudaContext::new(ordinal)
+            .map_err(|e| std::io::Error::other(format!("cudarc init device {device_id}: {e}")))?;
+        ctx.bind_to_thread()
+            .map_err(|e| std::io::Error::other(format!("cudarc bind device {device_id}: {e}")))?;
+        info!("Initialized CUDA context for device {device_id} (cudarc, torch-free)");
+    }
+    Ok(())
+}
+
 fn init_python_cuda(device_ids: &[i32]) -> Result<(), std::io::Error> {
     if device_ids.is_empty() {
         return Err(std::io::Error::other("no CUDA devices to initialize"));
@@ -447,7 +477,11 @@ pub fn run() -> Result<(), Box<dyn Error>> {
     // Determine which devices to initialize
     let devices = if cli.devices.is_empty() {
         // Auto-detect all available devices
-        let detected = detect_cuda_devices()?;
+        let detected = if cli.python_registry {
+            detect_cuda_devices()?
+        } else {
+            detect_native_cuda_devices()?
+        };
         info!(
             "Auto-detected {} CUDA device(s): {:?}",
             detected.len(),
@@ -463,17 +497,26 @@ pub fn run() -> Result<(), Box<dyn Error>> {
         return Err("No CUDA devices available".into());
     }
 
-    init_python_cuda(&devices)?;
+    if cli.python_registry {
+        init_python_cuda(&devices)?;
+    } else {
+        init_cudarc_cuda(&devices)?;
+    }
     info!(
         "CUDA runtime initialized for {} device(s): {:?}",
         devices.len(),
         devices
     );
 
-    let registry = CudaTensorRegistry::new().map_err(|err| {
-        let msg = format_py_err(err);
-        std::io::Error::other(format!("failed to initialize torch CUDA context: {msg}"))
-    })?;
+    let registry = if cli.python_registry {
+        CudaTensorRegistry::new().map_err(|err| {
+            let msg = format_py_err(err);
+            std::io::Error::other(format!("failed to initialize torch CUDA context: {msg}"))
+        })?
+    } else {
+        info!("Python/torch registry disabled; serving native clients only");
+        CudaTensorRegistry::empty()
+    };
     // Confine the registry to its own thread: GIL + CUDA work now happens off
     // the async runtime, so a wedged `empty_cache` can't starve tokio workers.
     let registry = RegistryHandle::spawn(registry);