diff --git a/README.md b/README.md index 7ba4be1b..5078cc6c 100644 --- a/README.md +++ b/README.md @@ -54,9 +54,20 @@ pegaflow-server ```bash vllm serve Qwen/Qwen3-0.6B \ - --kv-transfer-config '{"kv_connector": "PegaKVConnector", "kv_role": "kv_both", "kv_connector_module_path": "pegaflow.connector"}' + --kv-transfer-config '{ + "kv_connector": "PegaKVConnector", + "kv_role": "kv_both", + "kv_connector_module_path": "pegaflow.connector" + }' ``` +To enable direct GPU P2P reads, add +`"kv_connector_extra_config": {"pegaflow.direct_gpu_rdma": true}` to that +JSON. When it is omitted or set to `false`, remote reads use the host-staging +path. Direct mode currently supports dense attention cache group 0 only and +does not fall back to staging if a direct transfer fails. The same setting can +be passed to `examples/run_vllm_with_pega.py` with `--direct-gpu-rdma`. + > For full server options, multi-node setup, and advanced configuration, see [Server Configuration](./docs/server.md). ## Development diff --git a/docs/deployment.md b/docs/deployment.md index ef8cd7d5..0b42262d 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -12,6 +12,13 @@ and use a P/D-aware NIXL router. Configure each server's routable `--addr`, P and D must use compatible model, tokenizer, block size, KV dtype, KV layout, and `PYTHONHASHSEED`. +The PegaKVConnector option `pegaflow.direct_gpu_rdma` is configured inside the +connector's `kv_connector_extra_config`. Set it to `true` only on a +`read_write` PegaFlow connector that will load remote blocks; a `save_only` +connector has no PegaFlow load to accelerate. The option is omitted below, so +this deployment uses the host-staging path. See [P2P](./p2p.md#direct-gpu-performance-notes) +for the direct-load command and the benchmark caveats. + Replace `` and `` with the addresses assigned to the P and D nodes. The example assumes P and D run on separate nodes; to colocate them on one host, keep the distinct NIXL side-channel ports and point both diff --git a/docs/metrics.md b/docs/metrics.md index 3b36e410..f559e5a3 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -193,7 +193,7 @@ The setting remains configurable with `--metric-hll-bucket-bits`. ### Load Metrics (CPU → GPU) - **pegaflow_load_bytes_total** (Counter) - Total bytes loaded from CPU storage to GPU - - Use case: Monitor load throughput + - Use case: Monitor host-staging load throughput - **pegaflow_load_duration_seconds** (Histogram) - Load operation latency distribution @@ -203,6 +203,36 @@ The setting remains configurable with `--metric-hll-bucket-bits`. - Load operation failures (e.g., transfer errors) - Use case: Detect data transfer issues +### Direct GPU RDMA Load Metrics + +These metrics cover the opt-in `pegaflow.direct_gpu_rdma` path. A direct +remote read bypasses requester host staging, so it is not included in +`pegaflow_load_bytes_total` or `pegaflow_load_duration_seconds`. A mixed load +can still emit the ordinary load metrics for its local RAM prefix while the +remote suffix is represented by the direct metrics below. + +- **pegaflow_direct_gpu_load_total** (Counter, `status=success|error`) + - Number of direct GPU load batches reaching a terminal state + - Use case: Confirm that the direct path was exercised and detect failures + +- **pegaflow_direct_gpu_load_duration_seconds** (Histogram) + - End-to-end direct load batch duration + - Includes remote block metadata queries, RDMA completion, GPU visibility + flush, and any concurrent local H2D work + - Use case: Compare direct and host-staging load latency for the same hit set + +- **pegaflow_direct_gpu_mr_registration_failures** (Counter) + - CUDA DMA-BUF or RDMA memory-registration failures + - Use case: Detect capability or allocation-registration problems before a + direct load can start + +The current direct path does not export direct-RDMA bytes or per-stage +histograms. Do not compare `pegaflow_direct_gpu_load_duration_seconds` with +`pegaflow_rdma_fetch_duration` as if they had identical boundaries, and do not +use `pegaflow_rdma_fetch_bytes` to infer direct-GPU bandwidth. Use the direct +load counters together with request-level timings and debug logs for an A/B +comparison. + ### SSD Cache Metrics - **pegaflow_ssd_write_bytes_total** (Counter) - Bytes written to SSD cache - **pegaflow_ssd_write_duration_seconds** (Histogram) - SSD write latency @@ -229,6 +259,7 @@ This metric intentionally records decisions, not completed service outcomes. For backing failure correlation, use: - `pegaflow_rdma_fetch_total{status="error"}` for RDMA fetch failures +- `pegaflow_direct_gpu_load_total{status="error"}` for direct GPU load failures - `pegaflow_ssd_prefetch_failures_total` for SSD prefetch failures The legacy `pegaflow_cache_block_hits_total` and diff --git a/docs/p2p.md b/docs/p2p.md index ab3c6734..87d72130 100644 --- a/docs/p2p.md +++ b/docs/p2p.md @@ -81,20 +81,45 @@ pegaflow-server \ ### 3. Launch inference engine -Same as single-node — PegaFlow server handles P2P transparently. +Same as single-node — PegaFlow server handles host-staged P2P transparently. ```bash vllm serve Qwen/Qwen3-0.6B \ --kv-transfer-config '{"kv_connector": "PegaKVConnector", "kv_role": "kv_both", "kv_connector_module_path": "pegaflow.connector"}' ``` +To select direct GPU reads, pass the connector option through vLLM's +`kv_connector_extra_config`: + +```bash +vllm serve Qwen/Qwen3-0.6B \ + --kv-transfer-config '{ + "kv_connector": "PegaKVConnector", + "kv_role": "kv_both", + "kv_connector_module_path": "pegaflow.connector", + "kv_connector_extra_config": { + "pegaflow.direct_gpu_rdma": true + } + }' +``` + +When enabled, the scheduler carries the remote fetch plan in the query lease +and the worker issues RDMA READs into vLLM's registered GPU KV allocations. +Omitting the option (or setting it to `false`) keeps the existing path of RDMA +READ into host memory followed by host-to-GPU copy. Direct mode currently +supports dense attention cache group 0 only. A failed direct load is surfaced +to vLLM so the request can recompute; it is not automatically retried through +host staging. + ### 4. Verify Use `--log-level debug` to confirm P2P is working. Look for MetaServer registration, RDMA handshake, and RDMA fetch messages in the logs. -## Fallback Behavior +## Failure behavior -P2P is opportunistic. Failures degrade gracefully to single-node operation — no crashes, no significant performance impact in most cases. +Host-staged P2P is opportunistic and can proceed without a remote hit when +discovery or transfer fails. Direct GPU mode has a stricter failure boundary: +the failed load is reported to vLLM and the affected prefix is recomputed. | Scenario | What happens | |---|---| @@ -102,6 +127,62 @@ P2P is opportunistic. Failures degrade gracefully to single-node operation — n | Remote node unreachable | gRPC handshake fails, fetch aborted. Request proceeds without remote blocks. | | RDMA transfer timeout | Connection invalidated, transfer lock force-released. Logged as error. | +For direct GPU mode, the last two cases fail the direct load and do not invoke +the host-staging path. + +## Direct GPU performance notes + +Direct GPU mode removes the requester-side host-to-GPU copy. It does not remove +the other work on the critical path: MetaServer query, the per-segment gRPC +block-metadata query, connection setup on a cold peer, descriptor construction, +RDMA completion waits, transfer-lock release, and the CUDA GPUDirect visibility +flush. The current implementation processes segments in order and flushes the +CUDA context after each segment, so short prefixes or many owner segments can +be latency-bound even when the raw GB300 GPU RDMA bandwidth is higher than the +host-staging path. + +The existing 2P2D replay is not an apples-to-apples direct-load benchmark when +NIXL is configured as the P-to-D transfer connector. NIXL performs the main +P-to-D GPU transfer, while PegaFlow only handles the P-side cache lookup/load +for the requests that hit a remote PegaFlow owner. The decode-side +`save_only` connector records no PegaFlow loads. A small end-to-end difference +in that replay therefore does not measure the direct-vs-host-staging copy in +isolation. + +The GB300 replay illustrates the dilution. In one paired 8,422-request run, +direct mode completed 1,387 direct loads in 16.322 seconds in aggregate +(about 11.8 ms per direct load), while the host-staging control completed 735 +host RDMA fetches in 9.607 seconds (about 13.1 ms per fetch) and had one fetch +error. The direct and host counters cover different cache histories and +different operation boundaries, so these averages are directional evidence, +not a throughput ratio. Overall mean TTFT was 187.527 ms versus 187.460 ms and +throughput was 5.322 versus 5.326 requests/s. The run therefore shows why a +faster direct data path can produce little end-to-end movement: most requests +are misses, and the remaining hit latency includes scheduler work, P-side +prefill, NIXL P-to-D transfer, and decode startup. + +To compare the two PegaFlow paths, use the same model, block count, request +ordering, remote owner, and hit set. Make PegaKVConnector the requester-side +`read_write` connector and run the same requests once with direct mode and +once with host staging. For direct mode, use +`pegaflow_direct_gpu_load_total` and +`pegaflow_direct_gpu_load_duration_seconds`; for host staging, use the +`pegaflow_rdma_fetch_*` metrics. The latter do not include direct GPU loads, +and the current direct path does not export a direct-RDMA byte counter, so do +not combine those counters into one bandwidth number. Split the measurement +with connector timings and debug logs into query, metadata/handshake, RDMA +transfer, visibility flush, and vLLM scheduling time. + +The current direct implementation has four likely latency costs after the +query: one metadata RPC per owner segment, sequential segment processing, +descriptor construction and completion waits, and creating/binding a CUDA +context plus a GPUDirect visibility flush after each segment. Connection reuse +removes the handshake from warm peers, and local RAM H2D and remote GPU RDMA +already run concurrently. The first optimization candidates are therefore +coalescing or parallelizing owner segments and moving the visibility flush to +the end of a load when the CUDA/RDMA contract permits it; measure each change +with the isolated A/B workload above before changing the data path. + ## Tuning ### Hugepages @@ -137,6 +218,9 @@ P2P-related Prometheus metrics (on `:9091/metrics` by default): | `pegaflow_rdma_fetch_bytes` | Counter | Total bytes fetched via RDMA | | `pegaflow_rdma_fetch_plan_segments` | Histogram | Planned segment count per executed RDMA fetch plan | | `pegaflow_rdma_fetch_plan_completed_segments` | Histogram | Completed segment count before a plan stops | +| `pegaflow_direct_gpu_load_total` | Counter | Direct GPU load attempts, labelled by success or error | +| `pegaflow_direct_gpu_load_duration_seconds` | Histogram | End-to-end direct GPU load duration | +| `pegaflow_direct_gpu_mr_registration_failures` | Counter | GPU memory registration failures that prevent direct loads | | `pegaflow_rdma_qps` | Gauge | Active RDMA queue pairs | | `pegaflow_transfer_lock_active` | UpDownCounter | Currently held transfer locks | | `pegaflow_transfer_lock_timeouts_total` | Counter | Transfer lock timeout events | @@ -155,3 +239,6 @@ P2P-related Prometheus metrics (on `:9091/metrics` by default): - Enable hugepages for large pools (`--use-hugepages`). For all P2P issues, `--log-level debug` shows the full handshake and fetch flow. +The direct load duration includes the remote query, RDMA completion, GPU +visibility flush, and any concurrent local H2D load. It is not a raw GPU-RDMA +bandwidth measurement. diff --git a/examples/run_vllm_with_pega.py b/examples/run_vllm_with_pega.py index a951846f..83e00f27 100644 --- a/examples/run_vllm_with_pega.py +++ b/examples/run_vllm_with_pega.py @@ -48,6 +48,14 @@ def parse_args(): default=0.9, help="GPU memory utilization (default: 0.9)", ) + parser.add_argument( + "--direct-gpu-rdma", + action="store_true", + help=( + "Read remote KV blocks directly into GPU memory with GPUDirect RDMA " + "(default: disabled)" + ), + ) parser.add_argument( "--kv-events", action="store_true", @@ -77,6 +85,10 @@ def main(): "kv_role": "kv_both", # Both scheduler and worker roles "kv_connector_module_path": "pegaflow.connector", } + if args.direct_gpu_rdma: + kv_transfer_config["kv_connector_extra_config"] = { + "pegaflow.direct_gpu_rdma": True + } # Build vllm serve command cmd = [ @@ -114,6 +126,7 @@ def main(): print(f"Model: {args.model}") print(f"Endpoint: http://{args.host}:{args.port}") print(f"Tensor Parallel Size: {args.tensor_parallel_size}") + print(f"Direct GPU RDMA: {'enabled' if args.direct_gpu_rdma else 'disabled'}") if args.kv_events: print("Prefix Caching: enabled (required for KV events)") print( diff --git a/pegaflow-core/src/backing/mod.rs b/pegaflow-core/src/backing/mod.rs index b6c2d0fe..159639e6 100644 --- a/pegaflow-core/src/backing/mod.rs +++ b/pegaflow-core/src/backing/mod.rs @@ -27,7 +27,7 @@ use pegaflow_common::NumaNode; #[cfg(feature = "rdma")] pub(crate) use rdma::{RdmaTransport, new_rdma}; #[cfg(feature = "rdma")] -pub(crate) use rdma_fetch::RdmaFetchStore; +pub(crate) use rdma_fetch::{DirectFetchPlan, DirectQueryPlan, GpuReadTarget, RdmaFetchStore}; pub(crate) use ssd::SsdBackingStore; pub(crate) use ssd::new_ssd; diff --git a/pegaflow-core/src/backing/rdma_fetch.rs b/pegaflow-core/src/backing/rdma_fetch.rs index 88e801ed..865706bb 100644 --- a/pegaflow-core/src/backing/rdma_fetch.rs +++ b/pegaflow-core/src/backing/rdma_fetch.rs @@ -24,6 +24,7 @@ use super::transfer_lock_guard::TransferLockGuard; use super::{AllocateFn, PrefetchResult, RdmaTransport}; use crate::block::{BlockKey, RawBlock, SealedBlock, Segment}; use crate::internode::MetaServerClient; +use crate::layout::{BlockCopies, KVCacheLayout}; use crate::metrics::core_metrics; /// Minimum usable transfer timeout. If the server's lock timeout minus the @@ -59,16 +60,134 @@ pub(crate) struct RdmaFetchStore { } #[derive(Debug, Clone, PartialEq, Eq)] -struct FetchPlanSegment { - node: String, - start: usize, - end: usize, +pub(crate) struct FetchPlanSegment { + pub(crate) node: String, + pub(crate) start: usize, + pub(crate) end: usize, } #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct FetchPlan { - segments: Vec, - block_count: usize, + pub(crate) segments: Vec, + pub(crate) block_count: usize, +} + +/// Remote-only plan carried by a query lease until the destination GPU is +/// allocated. It contains no local allocation or transfer lock. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct DirectFetchPlan { + pub(crate) namespace: String, + pub(crate) hashes: Vec>, + pub(crate) fetch_plan: FetchPlan, +} + +impl DirectFetchPlan { + pub(crate) fn block_count(&self) -> usize { + self.fetch_plan.block_count() + } +} + +/// Local sealed prefix and an optional remote suffix, held until GPU allocation. +#[derive(Clone)] +pub struct DirectQueryPlan { + pub(crate) local_blocks: Vec>, + pub(crate) remote: Option, +} + +impl DirectQueryPlan { + pub fn block_count(&self) -> usize { + self.local_blocks.len() + self.remote.as_ref().map_or(0, DirectFetchPlan::block_count) + } +} + +/// Build READ descriptors for one remote slot and one local GPU block. +/// +/// The destination layout owns all pointer/range validation. Remote metadata +/// is checked for the same K/V shape before a WQE is submitted, so a stale or +/// incompatible owner cannot cause a short or shifted GPU write. +#[cfg(test)] +pub(crate) fn build_gpu_read_descs( + layout: &KVCacheLayout, + destination_block: usize, + remote_slot: &pegaflow_proto::proto::engine::TransferSlotInfo, +) -> Result, String> { + let mut descs = Vec::with_capacity(2); + append_gpu_read_descs(layout, destination_block, remote_slot, &mut descs)?; + Ok(descs) +} + +/// Append READ descriptors for one remote slot without allocating a temporary +/// vector. The direct-load path builds one segment-wide descriptor list, so +/// keeping the small K/V expansion in that list avoids a heap allocation for +/// every block and layer. +fn append_gpu_read_descs( + layout: &KVCacheLayout, + destination_block: usize, + remote_slot: &pegaflow_proto::proto::engine::TransferSlotInfo, + descs: &mut Vec, +) -> Result<(), String> { + let copies = layout.block_copies(destination_block)?; + let remote_k = NonNull::new(remote_slot.k_ptr as *mut u8) + .ok_or_else(|| "remote K ptr is null".to_string())?; + let k_size = usize::try_from(remote_slot.k_size) + .map_err(|_| "remote K size exceeds usize".to_string())?; + match copies { + BlockCopies::Contiguous(copy) => { + if remote_slot.v_ptr != 0 || remote_slot.v_size != 0 || k_size != copy.bytes { + return Err(format!( + "remote contiguous slot shape mismatch: remote k={} v_ptr={} v_size={}, local bytes={}", + k_size, remote_slot.v_ptr, remote_slot.v_size, copy.bytes + )); + } + descs.push(TransferDesc { + local_ptr: NonNull::new(copy.addr as *mut u8) + .ok_or_else(|| "local GPU K ptr is null".to_string())?, + remote_ptr: remote_k, + len: copy.bytes, + }); + } + BlockCopies::Split { k, v } => { + let remote_v = NonNull::new(remote_slot.v_ptr as *mut u8) + .ok_or_else(|| "remote V ptr is null for split slot".to_string())?; + let v_size = usize::try_from(remote_slot.v_size) + .map_err(|_| "remote V size exceeds usize".to_string())?; + if k_size != k.bytes || v_size != v.bytes { + return Err(format!( + "remote split slot shape mismatch: remote k/v={k_size}/{v_size}, local k/v={}/{}", + k.bytes, v.bytes + )); + } + descs.push(TransferDesc { + local_ptr: NonNull::new(k.addr as *mut u8) + .ok_or_else(|| "local GPU K ptr is null".to_string())?, + remote_ptr: remote_k, + len: k.bytes, + }); + descs.push(TransferDesc { + local_ptr: NonNull::new(v.addr as *mut u8) + .ok_or_else(|| "local GPU V ptr is null".to_string())?, + remote_ptr: remote_v, + len: v.bytes, + }); + } + } + Ok(()) +} + +fn flush_gpu_visibility(cuda_context: &cudarc::driver::CudaContext) -> Result<(), String> { + cuda_context + .bind_to_thread() + .map_err(|error| format!("direct GPU CUDA visibility fence failed: {error:?}"))?; + let flush = unsafe { + cudarc::driver::sys::cuFlushGPUDirectRDMAWrites( + cudarc::driver::sys::CUflushGPUDirectRDMAWritesTarget_enum::CU_FLUSH_GPU_DIRECT_RDMA_WRITES_TARGET_CURRENT_CTX, + cudarc::driver::sys::CUflushGPUDirectRDMAWritesScope_enum::CU_FLUSH_GPU_DIRECT_RDMA_WRITES_TO_OWNER, + ) + }; + if flush != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + return Err(format!("direct GPU visibility flush failed: {flush:?}")); + } + Ok(()) } impl FetchPlan { @@ -83,6 +202,15 @@ impl FetchPlan { .collect::>() .join(",") } + + pub(crate) fn direct_plan(&self, namespace: &str, hashes: &[Vec]) -> DirectFetchPlan { + let covered = self.block_count.min(hashes.len()); + DirectFetchPlan { + namespace: namespace.to_string(), + hashes: hashes[..covered].to_vec(), + fetch_plan: self.clone(), + } + } } fn validate_fetch_plan( @@ -248,6 +376,19 @@ impl RdmaFetchStore { Some(plan) } + /// Query a remote-only plan without allocating host staging memory or + /// taking a transfer lock. The plan is safe to keep in a query lease until + /// the scheduler has allocated its GPU destination blocks. + pub(crate) async fn query_direct_plan( + &self, + namespace: &str, + hashes: &[Vec], + ) -> Option { + self.query_plan(namespace, hashes) + .await + .map(|plan| plan.direct_plan(namespace, hashes)) + } + pub(crate) async fn fetch_plan( &self, plan: &FetchPlan, @@ -293,6 +434,210 @@ impl RdmaFetchStore { fetched } + /// Execute a remote plan directly into registered CUDA allocations. + /// `targets` contains one entry per local layer and one destination block + /// per planned hash. All segments are fetched in plan order so a stale + /// owner cannot make a later prefix segment appear valid. + pub(crate) async fn fetch_plan_to_gpu( + &self, + plan: &DirectFetchPlan, + req_id: &str, + remote_targets: &[GpuReadTarget], + cuda_context: &cudarc::driver::CudaContext, + transfer_timeout: Duration, + ) -> Result<(), String> { + if remote_targets.is_empty() { + return Err("direct GPU load requires at least one target layer".into()); + } + for target in remote_targets { + if target.destination_block_ids.len() != plan.hashes.len() { + return Err(format!( + "direct GPU layer {} has {} destination blocks for {} hashes", + target.layer_name, + target.destination_block_ids.len(), + plan.hashes.len() + )); + } + } + let mut submitted_gpu_transfer = false; + let transfer_result = async { + let mut offset = 0usize; + for (segment_index, segment) in plan.fetch_plan.segments.iter().enumerate() { + let remote_addr = &segment.node; + let end = segment.end; + let hashes = plan + .hashes + .get(segment.start..end) + .ok_or_else(|| format!("direct plan segment {segment_index} exceeds hash list"))?; + + ensure_connected( + &self.connect_group, + &self.rdma_transport, + &self.grpc_channels, + remote_addr, + &self.advertise_addr, + ) + .await?; + let (client, mut response) = query_remote_blocks( + &self.grpc_channels, + remote_addr, + &plan.namespace, + hashes, + &self.advertise_addr, + ) + .await?; + let lock_guard = TransferLockGuard::new( + client, + std::mem::take(&mut response.transfer_session_id), + remote_addr, + req_id, + ); + if response.blocks.len() != hashes.len() { + return Err(format!( + "direct GPU segment {segment_index} returned {} blocks for {} hashes", + response.blocks.len(), + hashes.len() + )); + } + let (receivers, expected_bytes) = { + let build_result = (|| { + let descriptor_capacity = hashes + .len() + .saturating_mul(remote_targets.len()) + .saturating_mul(2); + let mut descs = Vec::with_capacity(descriptor_capacity); + for (local_index, block) in response.blocks.iter().enumerate() { + if block.block_hash != hashes[local_index] { + return Err(format!( + "direct GPU segment {segment_index} block hash mismatch at offset {}", + segment.start + local_index + )); + } + let destination_index = segment.start + local_index; + for target in remote_targets { + let Some(destination) = target.destination_block_ids[destination_index] + else { + continue; + }; + let slot = block.slots.get(target.remote_slot_id).ok_or_else(|| { + format!( + "remote block has no slot {} for layer {}", + target.remote_slot_id, target.layer_name + ) + })?; + append_gpu_read_descs(&target.layout, destination, slot, &mut descs)?; + } + } + Ok(descs) + })(); + let descs = match build_result { + Ok(descs) => descs, + Err(error) => { + lock_guard.release(); + return Err(error); + } + }; + if descs.is_empty() { + lock_guard.release(); + offset = end; + continue; + } + let expected_bytes: usize = descs.iter().map(|desc| desc.len).sum(); + match self.rdma_transport.engine().batch_transfer_gpu_async( + TransferOp::Read, + remote_addr, + &descs, + ) { + Ok(receivers) => { + submitted_gpu_transfer = true; + (receivers, expected_bytes) + } + Err(error) => { + lock_guard.release(); + return Err(format!("direct GPU RDMA submit failed: {error}")); + } + } + }; + let completions = async { + let mut first_error = None; + let mut completed_bytes = 0usize; + for receiver in receivers { + match receiver.await { + Ok(Ok(bytes)) => completed_bytes += bytes, + Ok(Err(error)) => { + first_error.get_or_insert_with(|| { + format!("direct GPU RDMA completion failed: {error}") + }); + } + Err(_) => { + first_error.get_or_insert_with(|| { + "direct GPU completion channel closed".to_string() + }); + } + } + } + if let Some(error) = first_error { + return Err(error); + } + if completed_bytes != expected_bytes { + return Err(format!( + "direct GPU short completion: {completed_bytes}/{expected_bytes} bytes" + )); + } + Ok(()) + }; + tokio::pin!(completions); + let wait_result = match tokio::time::timeout(transfer_timeout, &mut completions).await { + Ok(result) => result, + Err(_) => { + // Reset the QPs before releasing the destination blocks. A + // dropped receiver alone would leave accepted READs able to + // write into memory that the caller may immediately reuse. + self.rdma_transport + .engine() + .invalidate_connection(remote_addr); + lock_guard.release(); + return Err("direct GPU RDMA transfer timed out".to_string()); + } + }; + if let Err(error) = wait_result { + self.rdma_transport + .engine() + .invalidate_connection(remote_addr); + lock_guard.release(); + return Err(error); + } + lock_guard.release(); + offset = end; + } + if offset != plan.hashes.len() { + return Err(format!( + "direct plan covers {} hashes but request has {}", + offset, + plan.hashes.len() + )); + } + Ok(()) + } + .await; + if !submitted_gpu_transfer { + return transfer_result; + } + + // RDMA READs complete independently of CUDA's device context. Bind + // and flush once after the ordered plan is fully drained instead of + // paying the context/visibility cost for every owner segment. This + // also runs when a later segment fails after an earlier segment + // already wrote to the destination GPU buffer. + match (transfer_result, flush_gpu_visibility(cuda_context)) { + (Ok(()), visibility_result) => visibility_result, + (Err(error), Ok(())) => Err(error), + (Err(error), Err(flush_error)) => Err(format!( + "{error}; GPU visibility flush also failed: {flush_error}" + )), + } + } + /// Fetch `hashes` from `remote_addr`. pub(crate) async fn fetch_blocks( &self, @@ -316,6 +661,13 @@ impl RdmaFetchStore { } } +pub(crate) struct GpuReadTarget { + pub(crate) layer_name: String, + pub(crate) layout: KVCacheLayout, + pub(crate) destination_block_ids: Vec>, + pub(crate) remote_slot_id: usize, +} + /// Execute RDMA fetch against a single remote node. /// /// 1. Ensure RDMA connection (singleflight per remote_addr) @@ -629,7 +981,7 @@ async fn fetch_blocks_via_rdma( }; let wait_start = Instant::now(); - tokio::time::timeout(transfer_timeout, async { + let wait_result = tokio::time::timeout(transfer_timeout, async { for rx in receivers { rx.await .map_err(|_| "RDMA transfer channel closed".to_string())? @@ -637,8 +989,20 @@ async fn fetch_blocks_via_rdma( } Ok::<(), String>(()) }) - .await - .map_err(|_| "RDMA transfer timed out".to_string())??; + .await; + match wait_result { + Ok(Ok(())) => {} + Ok(Err(error)) => { + rdma.engine().invalidate_connection(remote_addr); + return Err(error); + } + Err(_) => { + // Reset the QPs before dropping host allocations. Accepted READs + // must not continue writing into memory that is about to be freed. + rdma.engine().invalidate_connection(remote_addr); + return Err("RDMA transfer timed out".to_string()); + } + } timing.rdma_wait = wait_start.elapsed(); // Build SealedBlocks from allocated memory @@ -914,11 +1278,58 @@ fn transfer_timeout_from_server(lock_timeout_secs: u32) -> Duration { #[cfg(test)] mod tests { use super::*; + use pegaflow_proto::proto::engine::TransferSlotInfo; use std::collections::VecDeque; use std::num::NonZeroU64; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; + #[test] + fn gpu_read_descs_validate_contiguous_shape() { + let layout = KVCacheLayout::new(0x1000, 0x100, 2, 0x20, 0, 1).expect("layout"); + let slot = TransferSlotInfo { + k_ptr: 0x8000, + k_size: 0x20, + v_ptr: 0, + v_size: 0, + numa_node: 0, + }; + let descs = build_gpu_read_descs(&layout, 1, &slot).expect("descriptor"); + assert_eq!(descs.len(), 1); + assert_eq!(descs[0].local_ptr.as_ptr() as u64, 0x1020); + assert_eq!(descs[0].remote_ptr.as_ptr() as u64, 0x8000); + assert_eq!(descs[0].len, 0x20); + } + + #[test] + fn gpu_read_descs_validate_split_shape() { + let layout = KVCacheLayout::new(0x1000, 0x200, 1, 0x20, 0x80, 2).expect("layout"); + let slot = TransferSlotInfo { + k_ptr: 0x8000, + k_size: 0x20, + v_ptr: 0x9000, + v_size: 0x20, + numa_node: 0, + }; + let descs = build_gpu_read_descs(&layout, 0, &slot).expect("descriptors"); + assert_eq!(descs.len(), 2); + assert_eq!(descs[0].local_ptr.as_ptr() as u64, 0x1000); + assert_eq!(descs[1].local_ptr.as_ptr() as u64, 0x1080); + } + + #[test] + fn gpu_read_descs_reject_stale_remote_shape() { + let layout = KVCacheLayout::new(0x1000, 0x100, 1, 0x20, 0, 1).expect("layout"); + let slot = TransferSlotInfo { + k_ptr: 0x8000, + k_size: 0x10, + v_ptr: 0, + v_size: 0, + numa_node: 0, + }; + let error = build_gpu_read_descs(&layout, 0, &slot).expect_err("shape mismatch"); + assert!(error.contains("shape mismatch")); + } fn test_allocate_fn(calls: Arc) -> AllocateFn { let allocator = Arc::new(crate::pinned_pool::PinnedAllocator::new_global( 32 * 1024 * 1024, @@ -1001,6 +1412,31 @@ mod tests { ); } + #[test] + fn direct_plan_preserves_hash_prefix_and_segment_sizes() { + let plan = FetchPlan { + segments: vec![ + FetchPlanSegment { + node: "node-a".into(), + start: 0, + end: 2, + }, + FetchPlanSegment { + node: "node-b".into(), + start: 2, + end: 3, + }, + ], + block_count: 3, + }; + let hashes = vec![vec![1], vec![2], vec![3], vec![4]]; + let direct = plan.direct_plan("ns", &hashes); + assert_eq!(direct.namespace, "ns"); + assert_eq!(direct.hashes, hashes[..3].to_vec()); + assert_eq!(direct.block_count(), 3); + assert_eq!(direct.fetch_plan.segments, plan.segments); + } + #[test] fn rejects_invalid_fetch_plans() { for (segments, expected) in [ diff --git a/pegaflow-core/src/instance.rs b/pegaflow-core/src/instance.rs index 52c6da26..d387d6a1 100644 --- a/pegaflow-core/src/instance.rs +++ b/pegaflow-core/src/instance.rs @@ -295,6 +295,9 @@ pub struct GpuContext { /// Hybrid-cache storage group id by layer name; absent = group 0. layer_groups: HashMap, + #[cfg(feature = "rdma")] + direct_memory: Mutex>, + /// CUDA context handle (kept alive for the lifetime of this context). _cuda_ctx: Arc, @@ -335,6 +338,8 @@ impl GpuContext { preferred_numa: numa_node, kv_caches, layer_groups, + #[cfg(feature = "rdma")] + direct_memory: Mutex::new(None), _cuda_ctx: cuda_ctx, worker_pool, }) @@ -345,6 +350,39 @@ impl GpuContext { self.preferred_numa } + #[cfg(feature = "rdma")] + pub(crate) fn register_direct_memory( + &self, + storage: &crate::storage::StorageEngine, + imported: &[Arc], + ) -> Result<(), EngineError> { + let mut registered = self.direct_memory.lock(); + if registered.is_some() { + return Err(EngineError::InvalidArgument( + "direct GPU memory already registered".into(), + )); + } + if imported.is_empty() { + return Err(EngineError::InvalidArgument( + "direct GPU memory requires owner DMA-BUF exports".into(), + )); + } + self._cuda_ctx + .bind_to_thread() + .map_err(|error| EngineError::CudaInit(error.to_string()))?; + *registered = Some( + storage + .register_device_memory(self.device_id, imported) + .map_err(EngineError::Storage)?, + ); + Ok(()) + } + + #[cfg(feature = "rdma")] + pub(crate) fn direct_memory_registered(&self) -> bool { + self.direct_memory.lock().is_some() + } + /// CUDA device ID represented by this shard. pub(crate) fn device_id(&self) -> i32 { self.device_id @@ -370,6 +408,13 @@ impl GpuContext { &self.worker_pool } + /// Reuse the instance CUDA context for operations that run outside the + /// dedicated GPU worker threads. + #[cfg(feature = "rdma")] + pub(crate) fn cuda_context(&self) -> Arc { + Arc::clone(&self._cuda_ctx) + } + /// Hybrid-cache storage group of a layer; unregistered layers default to /// group 0, preserving single-group behavior. pub(crate) fn group_of_layer(&self, layer_name: &str) -> u32 { diff --git a/pegaflow-core/src/lease.rs b/pegaflow-core/src/lease.rs index 526b8397..ed145992 100644 --- a/pegaflow-core/src/lease.rs +++ b/pegaflow-core/src/lease.rs @@ -6,6 +6,8 @@ use std::time::{Duration, Instant}; use tokio::task::JoinHandle; use uuid::Uuid; +#[cfg(feature = "rdma")] +use crate::backing::DirectQueryPlan; use crate::block::SealedBlock; const DEFAULT_LEASE_TTL: Duration = Duration::from_secs(600); @@ -42,11 +44,17 @@ impl fmt::Debug for QueryLeaseId { struct QueryLease { instance_id: String, - blocks: Vec>, + payload: QueryLeasePayload, remaining_consumers: usize, expires_at: Instant, } +pub(crate) enum QueryLeasePayload { + Cached(Vec>), + #[cfg(feature = "rdma")] + Direct(DirectQueryPlan), +} + pub(crate) struct QueryLeaseManager { inner: Arc, sweeper: Option>, @@ -94,7 +102,26 @@ impl QueryLeaseManager { let token = QueryLeaseId::fresh(); let lease = QueryLease { instance_id: instance_id.to_string(), - blocks, + payload: QueryLeasePayload::Cached(blocks), + remaining_consumers: consumers.max(1), + expires_at: Instant::now() + DEFAULT_LEASE_TTL, + }; + self.inner.insert(token, lease); + token + } + + #[cfg(feature = "rdma")] + pub(crate) fn create_direct( + &self, + instance_id: &str, + plan: DirectQueryPlan, + consumers: usize, + ) -> QueryLeaseId { + self.sweep_expired(); + let token = QueryLeaseId::fresh(); + let lease = QueryLease { + instance_id: instance_id.to_string(), + payload: QueryLeasePayload::Direct(plan), remaining_consumers: consumers.max(1), expires_at: Instant::now() + DEFAULT_LEASE_TTL, }; @@ -106,7 +133,7 @@ impl QueryLeaseManager { &self, instance_id: &str, token: &QueryLeaseId, - ) -> Result>, String> { + ) -> Result { self.sweep_expired(); let mut leases = self .inner @@ -124,13 +151,17 @@ impl QueryLeaseManager { } if lease.remaining_consumers > 1 { lease.remaining_consumers -= 1; - return Ok(lease.blocks.clone()); + return Ok(match &lease.payload { + QueryLeasePayload::Cached(blocks) => QueryLeasePayload::Cached(blocks.clone()), + #[cfg(feature = "rdma")] + QueryLeasePayload::Direct(plan) => QueryLeasePayload::Direct(plan.clone()), + }); } Ok(leases .remove(token) .expect("query lease disappeared during consume") - .blocks) + .payload) } pub(crate) fn release(&self, token: &QueryLeaseId) -> bool { @@ -202,7 +233,7 @@ mod tests { lease_id, QueryLease { instance_id: "inst-a".to_string(), - blocks: Vec::new(), + payload: QueryLeasePayload::Cached(Vec::new()), remaining_consumers: 1, expires_at: Instant::now() + DEFAULT_LEASE_TTL, }, @@ -225,8 +256,12 @@ mod tests { let blocks = vec![Arc::new(SealedBlock::from_slots(Vec::new()))]; let lease_id = manager.create("inst-a", blocks, 2); - assert_eq!(manager.consume("inst-a", &lease_id).unwrap().len(), 1); - assert_eq!(manager.consume("inst-a", &lease_id).unwrap().len(), 1); + assert!( + matches!(manager.consume("inst-a", &lease_id).unwrap(), QueryLeasePayload::Cached(blocks) if blocks.len() == 1) + ); + assert!( + matches!(manager.consume("inst-a", &lease_id).unwrap(), QueryLeasePayload::Cached(blocks) if blocks.len() == 1) + ); let err = manager .consume("inst-a", &lease_id) @@ -234,4 +269,87 @@ mod tests { .expect("lease should be exhausted"); assert!(err.contains("query lease is unknown or expired")); } + + #[cfg(feature = "rdma")] + #[test] + fn direct_lease_round_trips_plan_for_each_consumer() { + for (local_count, remote_count) in [(1, 0), (0, 2), (1, 2)] { + let manager = QueryLeaseManager::default(); + let block = Arc::new(SealedBlock::from_slots(Vec::new())); + let remote = (remote_count > 0).then(|| crate::backing::DirectFetchPlan { + namespace: "ns".into(), + hashes: vec![vec![1], vec![2]], + fetch_plan: crate::backing::rdma_fetch::FetchPlan { + segments: vec![ + crate::backing::rdma_fetch::FetchPlanSegment { + node: "node-a".into(), + start: 0, + end: 1, + }, + crate::backing::rdma_fetch::FetchPlanSegment { + node: "node-b".into(), + start: 1, + end: 2, + }, + ], + block_count: 2, + }, + }); + let plan = DirectQueryPlan { + local_blocks: vec![Arc::clone(&block); local_count], + remote: remote.clone(), + }; + let lease_id = manager.create_direct("inst-a", plan.clone(), 2); + drop(plan); + assert!(manager.consume("inst-b", &lease_id).is_err()); + for _ in 0..2 { + let QueryLeasePayload::Direct(received) = + manager.consume("inst-a", &lease_id).unwrap() + else { + panic!("expected direct query plan"); + }; + assert_eq!(received.block_count(), local_count + remote_count); + assert_eq!(received.remote, remote); + assert_eq!(received.local_blocks.len(), local_count); + for local in &received.local_blocks { + assert!(Arc::ptr_eq(local, &block)); + } + } + assert!(manager.consume("inst-a", &lease_id).is_err()); + assert_eq!(Arc::strong_count(&block), 1); + } + } + + #[cfg(feature = "rdma")] + #[test] + fn direct_lease_release_and_expiry_drop_local_pins() { + let manager = QueryLeaseManager::default(); + let block = Arc::new(SealedBlock::from_slots(Vec::new())); + for expire in [false, true] { + let lease = manager.create_direct( + "inst-a", + DirectQueryPlan { + local_blocks: vec![Arc::clone(&block)], + remote: None, + }, + 2, + ); + assert_eq!(Arc::strong_count(&block), 2); + if expire { + manager + .inner + .leases + .lock() + .unwrap() + .get_mut(&lease) + .unwrap() + .expires_at = Instant::now(); + manager.sweep_expired(); + } else { + assert!(manager.release(&lease)); + } + assert_eq!(Arc::strong_count(&block), 1); + assert!(manager.consume("inst-a", &lease).is_err()); + } + } } diff --git a/pegaflow-core/src/lib.rs b/pegaflow-core/src/lib.rs index 7c15b49b..670706a1 100644 --- a/pegaflow-core/src/lib.rs +++ b/pegaflow-core/src/lib.rs @@ -45,6 +45,8 @@ use layout::KVCacheLayout; pub use lease::QueryLeaseId; pub use pegaflow_common::NumaNode; use pegaflow_common::{NumaTopology, group_hash}; +#[cfg(feature = "rdma")] +pub use pegaflow_transfer::CudaDmaBuf; pub use pinned_pool::PinnedAllocation; pub use seal_offload::SlotMeta; pub use storage::{DEFAULT_RDMA_QPS_PER_PEER, MemoryCacheCleanupStats, StorageConfig}; @@ -61,8 +63,10 @@ use std::{ use log::{debug, info}; use crate::backing::SSD_ALIGNMENT; +#[cfg(feature = "rdma")] +use crate::backing::{DirectFetchPlan, DirectQueryPlan, GpuReadTarget}; use crate::gpu_worker::{HostBlock, LayerTransferData, LoadCompletion, LoadTask, TransferBlock}; -use crate::lease::QueryLeaseManager; +use crate::lease::{QueryLeaseManager, QueryLeasePayload}; use crate::metrics::core_metrics; use crate::storage::StorageEngine; use tokio::sync::oneshot; @@ -580,6 +584,33 @@ impl PegaEngine { Ok(status) } + /// Hold the local RAM prefix and query a direct GPU plan for its missing + /// suffix, without staging remote data or holding remote transfer locks. + #[cfg(feature = "rdma")] + pub async fn query_direct_gpu_plan( + &self, + instance_id: &str, + block_hashes: &[Vec], + ) -> Result, EngineError> { + if !self.has_rdma_transport() { + return Err(EngineError::Storage( + "direct GPU RDMA requires configured RDMA transport".to_string(), + )); + } + let instance = self.get_instance(instance_id)?; + let topology = instance.sealed_topology()?; + if topology.num_groups() != 1 || topology.is_page_first() { + return Err(EngineError::InvalidArgument( + "direct GPU RDMA only supports dense attention group 0".to_string(), + )); + } + let namespace = instance.namespace(); + Ok(self + .storage + .query_direct_plan(namespace, block_hashes) + .await) + } + /// Position-aligned membership query over one hybrid-cache storage group. /// /// Unlike prefix queries, every position reports independently: entry `i` @@ -627,6 +658,24 @@ impl PegaEngine { .create(instance_id, blocks, instance.world_size())) } + /// Lease the local sealed prefix and remote suffix plan until GPU allocation. + #[cfg(feature = "rdma")] + pub fn create_direct_query_lease( + &self, + instance_id: &str, + plan: DirectQueryPlan, + ) -> Result { + let instance = self.get_instance(instance_id)?; + if plan.block_count() == 0 { + return Err(EngineError::InvalidArgument( + "direct query lease requires at least one block".to_string(), + )); + } + Ok(self + .query_leases + .create_direct(instance_id, plan, instance.world_size())) + } + /// Release a query lease. Returns false when the lease is unknown or expired. pub fn release_query_lease(&self, lease: &QueryLeaseId) -> bool { self.query_leases.release(lease) @@ -775,11 +824,9 @@ impl PegaEngine { trace_scope!("load.cache_lookup", _s); let mut block_targets_by_group = vec![Vec::new(); layer_groups.len()]; let mut block_cache = Vec::new(); + #[cfg(feature = "rdma")] + let mut direct_loads = Vec::new(); for (lease, lease_block_ids_by_group) in loads { - let blocks = self - .query_leases - .consume(instance_id, lease) - .map_err(EngineError::Storage)?; if lease_block_ids_by_group.len() != layer_groups.len() { return Err(EngineError::InvalidArgument(format!( "load group count {} does not match layer group count {}", @@ -787,16 +834,71 @@ impl PegaEngine { layer_groups.len() ))); } + let (blocks, block_count) = match self + .query_leases + .consume(instance_id, lease) + .map_err(EngineError::Storage)? + { + QueryLeasePayload::Cached(blocks) => { + let count = blocks.len(); + (blocks, count) + } + #[cfg(feature = "rdma")] + QueryLeasePayload::Direct(plan) => { + if topology.is_page_first() + || topology.num_groups() != 1 + || layer_groups[0].is_empty() + || layer_groups.iter().skip(1).any(|group| !group.is_empty()) + { + return Err(EngineError::InvalidArgument( + "direct GPU load only supports dense attention group 0".to_string(), + )); + } + let count = plan.block_count(); + if lease_block_ids_by_group[0].len() != count { + return Err(EngineError::InvalidArgument(format!( + "query lease block count {count} does not match destination block count {}", + lease_block_ids_by_group[0].len() + ))); + } + if let Some(remote) = plan.remote { + let destinations = &lease_block_ids_by_group[0][plan.local_blocks.len()..]; + let mut targets = Vec::with_capacity(layer_groups[0].len()); + for layer_name in &layer_groups[0] { + let layer_id = topology.layer_id(layer_name)?; + let layout = gpu.get_layout(layer_name).ok_or_else(|| { + EngineError::InvalidArgument(format!( + "layer {layer_name} not registered on device {device_id}" + )) + })?; + for &destination in destinations.iter().flatten() { + layout + .block_copies(destination) + .map_err(EngineError::InvalidArgument)?; + } + targets.push(GpuReadTarget { + layer_name: layer_name.to_string(), + layout, + destination_block_ids: destinations.to_vec(), + remote_slot_id: topology.slot_index(layer_id, tp_rank)?, + }); + } + direct_loads.push((remote, targets)); + } + (plan.local_blocks, count) + } + }; let block_cache_start = block_cache.len(); for (group_index, lease_block_targets) in lease_block_ids_by_group.iter().enumerate() { - if blocks.len() != lease_block_targets.len() { + if block_count != lease_block_targets.len() { return Err(EngineError::InvalidArgument(format!( "query lease block count {} does not match destination block count {} for group {}", - blocks.len(), + block_count, lease_block_targets.len(), group_index ))); } + let lease_block_targets = &lease_block_targets[..blocks.len()]; // A stored block must carry exactly the slot layout of the // storage group this target group loads into. A mismatch // means the namespace is shared by instances with different @@ -882,6 +984,11 @@ impl PegaEngine { } } + #[cfg(feature = "rdma")] + if !direct_loads.is_empty() { + return self.submit_direct_load(gpu, layers, direct_loads, completion); + } + // Complete immediately if no blocks to load if layers.is_empty() { debug!("No blocks to load, completing immediately"); @@ -894,6 +1001,100 @@ impl PegaEngine { .submit_load(LoadTask { layers, completion }) } + #[cfg(feature = "rdma")] + fn submit_direct_load( + &self, + gpu: Arc, + layers: Vec, + direct_loads: Vec<(DirectFetchPlan, Vec)>, + completion: LoadCompletion, + ) -> Result<(), EngineError> { + let metrics = core_metrics(); + if !gpu.direct_memory_registered() { + metrics.direct_gpu_mr_registration_failures.add(1, &[]); + metrics + .direct_gpu_load_total + .add(1, &[opentelemetry::KeyValue::new("status", "error")]); + return Err(EngineError::InvalidArgument( + "direct GPU memory is not registered; owner DMA-BUF import is required".into(), + )); + } + let local_completion = if layers.is_empty() { + None + } else { + let (reply, rx) = oneshot::channel(); + gpu.worker_pool().submit_load(LoadTask { + layers, + completion: LoadCompletion::Channel(reply), + })?; + Some(rx) + }; + let storage = Arc::clone(&self.storage); + let cuda_context = gpu.cuda_context(); + let req_id = format!("direct-load:{}", uuid::Uuid::new_v4()); + let started_at = std::time::Instant::now(); + tokio::spawn(async move { + let _gpu = gpu; + let remote = async { + for (plan, targets) in &direct_loads { + storage + .direct_load(plan, &req_id, targets, &cuda_context) + .await + .map_err(EngineError::Storage)?; + } + Ok(()) + }; + let local = async { + match local_completion { + Some(rx) => rx.await.map_err(|_| { + EngineError::Storage("local GPU load completion channel closed".into()) + })?, + None => Ok(()), + } + }; + // Both sources may already be writing to GPU memory. Even on error, + // drain both before allowing the caller to reuse destination blocks. + let (local_result, remote_result) = tokio::join!(local, remote); + let result = local_result.and(remote_result); + if let Err(error) = &result { + log::error!("direct GPU load failed: {error}"); + } + let metrics = core_metrics(); + metrics.direct_gpu_load_total.add( + 1, + &[opentelemetry::KeyValue::new( + "status", + if result.is_ok() { "success" } else { "error" }, + )], + ); + metrics + .direct_gpu_load_duration_seconds + .record(started_at.elapsed().as_secs_f64(), &[]); + completion.signal(result); + }); + Ok(()) + } + + /// Attach allocation-owner DMA-BUF exports before serving direct loads. + #[cfg(feature = "rdma")] + pub fn register_imported_device_memory( + &self, + instance_id: &str, + device_id: i32, + regions: &[Arc], + ) -> Result<(), EngineError> { + let instance = self.get_instance(instance_id)?; + let gpu = instance + .get_gpu(device_id) + .ok_or_else(|| EngineError::WorkerMissing(instance_id.to_string(), device_id))?; + gpu.register_direct_memory(&self.storage, regions) + .inspect_err(|_| { + core_metrics() + .direct_gpu_mr_registration_failures + .add(1, &[]); + }) + } + /// Wait until all previously submitted save batches have been processed /// by the insert worker. /// diff --git a/pegaflow-core/src/metrics.rs b/pegaflow-core/src/metrics.rs index 4879499b..8e8197e5 100644 --- a/pegaflow-core/src/metrics.rs +++ b/pegaflow-core/src/metrics.rs @@ -113,6 +113,12 @@ pub(crate) struct CoreMetrics { pub rdma_fetch_plan_segments: Histogram, #[cfg(feature = "rdma")] pub rdma_fetch_plan_completed_segments: Histogram, + #[cfg(feature = "rdma")] + pub direct_gpu_load_total: Counter, + #[cfg(feature = "rdma")] + pub direct_gpu_mr_registration_failures: Counter, + #[cfg(feature = "rdma")] + pub direct_gpu_load_duration_seconds: Histogram, } fn init_meter() -> Meter { @@ -506,6 +512,23 @@ pub(crate) fn core_metrics() -> &'static CoreMetrics { ) .with_boundaries(rdma_fetch_plan_segment_boundaries()) .build(), + #[cfg(feature = "rdma")] + direct_gpu_load_total: meter + .u64_counter("pegaflow_direct_gpu_load_total") + .with_description("Direct GPU RDMA load attempts (status=success|error)") + .build(), + #[cfg(feature = "rdma")] + direct_gpu_mr_registration_failures: meter + .u64_counter("pegaflow_direct_gpu_mr_registration_failures") + .with_description("Direct GPU CUDA DMA-BUF/MR registration failures") + .build(), + #[cfg(feature = "rdma")] + direct_gpu_load_duration_seconds: meter + .f64_histogram("pegaflow_direct_gpu_load_duration") + .with_unit("s") + .with_description("Direct GPU RDMA load completion latency") + .with_boundaries(duration_seconds_boundaries()) + .build(), } }) } diff --git a/pegaflow-core/src/storage/mod.rs b/pegaflow-core/src/storage/mod.rs index e17aee96..ee052c5f 100644 --- a/pegaflow-core/src/storage/mod.rs +++ b/pegaflow-core/src/storage/mod.rs @@ -5,6 +5,8 @@ pub(crate) mod transfer_lock; mod write_path; use bytesize::ByteSize; +#[cfg(feature = "rdma")] +use cudarc::driver::CudaContext; use log::{debug, info, warn}; use std::collections::HashSet; use std::num::NonZeroU64; @@ -13,13 +15,19 @@ use std::time::Duration; use crate::backing::{AllocateFn, DEFAULT_MAX_PREFETCH_BLOCKS, SsdBackingStore, SsdCacheConfig}; #[cfg(feature = "rdma")] -use crate::backing::{RdmaFetchStore, RdmaTransport}; +use crate::backing::{ + DirectFetchPlan, DirectQueryPlan, GpuReadTarget, RdmaFetchStore, RdmaTransport, +}; use crate::block::{BlockKey, PrefetchStatus, SealedBlock}; use crate::internode::MetaServerClient; use crate::internode::metaserver_client::MetaServerClientConfig; use crate::metrics::core_metrics; use crate::pinned_pool::{PinnedAllocation, PinnedAllocator}; use pegaflow_common::NumaNode; +#[cfg(feature = "rdma")] +use pegaflow_transfer::CudaDmaBuf; +#[cfg(feature = "rdma")] +use std::ptr::NonNull; use prefetch::PrefetchScheduler; #[cfg(feature = "rdma")] @@ -99,6 +107,8 @@ pub(crate) struct StorageEngine { ssd_store: Option>, #[cfg(feature = "rdma")] rdma_transport: Option>, + #[cfg(feature = "rdma")] + rdma_fetch: Option>, blockwise_alloc: bool, metaserver_client: Option>, transfer_lock: Arc, @@ -219,19 +229,21 @@ impl StorageEngine { .advertise_addr .clone() .unwrap_or_else(|| "127.0.0.1:50055".to_string()); - Some(RdmaFetch::new(Arc::new(RdmaFetchStore::new( + Some(Arc::new(RdmaFetchStore::new( Arc::clone(ms), Arc::clone(rdma), allocate_fn.clone(), advertise, - )))) + ))) }); + #[cfg(feature = "rdma")] + let rdma_fetch_for_prefetch = rdma_fetch.clone().map(RdmaFetch::new); #[cfg(not(feature = "rdma"))] - let rdma_fetch = None; + let rdma_fetch_for_prefetch = None; let prefetch = PrefetchScheduler::new( ssd_store.clone(), - rdma_fetch, + rdma_fetch_for_prefetch, metaserver_client.clone(), max_prefetch_blocks, ); @@ -248,6 +260,8 @@ impl StorageEngine { ssd_store, #[cfg(feature = "rdma")] rdma_transport, + #[cfg(feature = "rdma")] + rdma_fetch, blockwise_alloc, metaserver_client, transfer_lock, @@ -479,6 +493,65 @@ impl StorageEngine { .await } + #[cfg(feature = "rdma")] + pub(crate) async fn query_direct_plan( + &self, + namespace: &str, + hashes: &[Vec], + ) -> Option { + if hashes.is_empty() { + return None; + } + let keys: Vec = hashes + .iter() + .map(|hash| BlockKey::new(namespace.to_string(), hash.clone())) + .collect(); + let (local_count, local_blocks) = self.read_cache.get_prefix_blocks(&keys); + let remote = if local_count < hashes.len() { + match self.rdma_fetch.as_ref() { + Some(fetch) => { + fetch + .query_direct_plan(namespace, &hashes[local_count..]) + .await + } + None => None, + } + } else { + None + }; + let remote_count = remote.as_ref().map_or(0, DirectFetchPlan::block_count); + if local_count == 0 && remote_count == 0 { + return None; + } + Some(DirectQueryPlan { + local_blocks, + remote, + }) + } + + #[cfg(feature = "rdma")] + pub(crate) async fn direct_load( + &self, + plan: &DirectFetchPlan, + req_id: &str, + targets: &[GpuReadTarget], + cuda_context: &CudaContext, + ) -> Result<(), String> { + let fetch = self + .rdma_fetch + .as_ref() + .ok_or_else(|| "direct GPU RDMA is not configured".to_string())?; + fetch + .fetch_plan_to_gpu( + plan, + req_id, + targets, + cuda_context, + self.transfer_lock_timeout(), + ) + .await + } + fn reclaim_until_allocator_can_allocate( &self, required_bytes: u64, @@ -627,6 +700,33 @@ impl StorageEngine { self.rdma_transport.as_ref() } + /// Register owner-exported allocations transactionally on every NIC. + #[cfg(feature = "rdma")] + pub(crate) fn register_device_memory( + &self, + device_id: i32, + regions: &[Arc], + ) -> Result { + let transport = self + .rdma_transport + .as_ref() + .ok_or_else(|| "direct-GPU registration requires RDMA transport".to_string())?; + if device_id < 0 || device_id > u8::MAX as i32 { + return Err(format!("CUDA device id {device_id} is out of range")); + } + let mut allocations = regions.to_vec(); + allocations.sort_unstable_by_key(|region| (region.ptr, region.len)); + allocations.dedup_by_key(|region| (region.ptr, region.len)); + transport + .engine() + .register_dma_buf_memory(&allocations, device_id as u8) + .map_err(|error| error.to_string())?; + Ok(DeviceMemoryRegistration { + transport: Arc::clone(transport), + ptrs: allocations.iter().map(|region| region.ptr).collect(), + }) + } + pub(crate) async fn shutdown_metaserver_client(&self) { if let Some(client) = &self.metaserver_client { client.shutdown().await; @@ -634,6 +734,26 @@ impl StorageEngine { } } +#[cfg(feature = "rdma")] +pub(crate) struct DeviceMemoryRegistration { + transport: Arc, + ptrs: Vec, +} + +#[cfg(feature = "rdma")] +impl Drop for DeviceMemoryRegistration { + fn drop(&mut self) { + let ptrs: Vec<_> = self + .ptrs + .iter() + .filter_map(|&ptr| NonNull::new(ptr as *mut u8)) + .collect(); + if let Err(error) = self.transport.engine().unregister_memory(&ptrs) { + log::error!("Failed to unregister CUDA RDMA memory: {error}"); + } + } +} + #[cfg(test)] impl StorageEngine { /// Insert a block directly into the in-memory cache (test only). diff --git a/pegaflow-proto/proto/engine.proto b/pegaflow-proto/proto/engine.proto index e3bfa8d5..a7531d75 100644 --- a/pegaflow-proto/proto/engine.proto +++ b/pegaflow-proto/proto/engine.proto @@ -131,6 +131,10 @@ message QueryRequest { // prefix-clamped, and the lease pins exactly the hit blocks; local-only // unless wait_for_full_prefix requests the remote-fetching set form. uint32 group_id = 5; + // Reuse the local RAM prefix, then plan RDMA READ of the missing suffix into the + // destination GPU. When false, the existing host-staging query behavior is + // unchanged. Direct mode is currently limited to group 0 dense attention. + bool direct_gpu = 6; } message QueryResponse { diff --git a/pegaflow-server/src/registry.rs b/pegaflow-server/src/registry.rs index d2e33b67..5b4be9e7 100644 --- a/pegaflow-server/src/registry.rs +++ b/pegaflow-server/src/registry.rs @@ -2,6 +2,7 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::PyBytes; use std::collections::{HashMap, HashSet}; +use std::sync::Arc; use tokio::sync::{mpsc, oneshot}; #[derive(Debug, Clone)] @@ -9,6 +10,8 @@ pub struct TensorMetadata { pub data_ptr: u64, pub size_bytes: usize, pub device_id: i32, + #[cfg(feature = "rdma")] + pub dma_buf: Option>, } struct LayerTensor { @@ -16,7 +19,7 @@ struct LayerTensor { dead_code, reason = "holding the Python tensor keeps CUDA IPC memory mapped" )] - tensor: Py, + tensor: Arc>, metadata: TensorMetadata, } @@ -201,7 +204,54 @@ impl CudaTensorRegistry { ) })?; - let tensor_owned = tensor.unbind(); + let tensor_owned = Arc::new(tensor.unbind()); + #[cfg(feature = "rdma")] + let dma_buf = { + use std::os::fd::{FromRawFd, OwnedFd}; + let metadata = if wrapper.hasattr("detach_dma_buf")? { + wrapper + .call_method0("detach_dma_buf")? + .extract::>()? + } else { + None + }; + if let Some((raw_fd, view_offset, allocation_bytes)) = metadata { + if raw_fd < 0 { + return Err(PyValueError::new_err("invalid DMA-BUF descriptor")); + } + // detach_dma_buf transfers a fresh descriptor to this process. + let fd = unsafe { OwnedFd::from_raw_fd(raw_fd) }; + let base = data_ptr + .checked_sub(view_offset) + .ok_or_else(|| PyValueError::new_err("invalid DMA-BUF view offset"))?; + if view_offset + .checked_add(size_bytes as u64) + .is_none_or(|end| end > allocation_bytes as u64) + { + return Err(PyValueError::new_err( + "CUDA tensor exceeds DMA-BUF allocation", + )); + } + cuda.call_method1("set_device", (resolved_device,))?; + Some(Arc::new( + pegaflow_core::CudaDmaBuf::from_import( + base, + allocation_bytes, + fd, + tensor_owned.clone(), + ) + .map_err(|error| PyValueError::new_err(error.to_string()))?, + )) + } else { + None + } + }; + #[cfg(not(feature = "rdma"))] + if wrapper.hasattr("dma_buf")? && !wrapper.getattr("dma_buf")?.is_none() { + return Err(PyValueError::new_err( + "DMA-BUF registration requires an RDMA-enabled server", + )); + } Ok(LayerTensor { tensor: tensor_owned, @@ -209,6 +259,8 @@ impl CudaTensorRegistry { data_ptr, size_bytes, device_id: resolved_device, + #[cfg(feature = "rdma")] + dma_buf, }, }) }) diff --git a/pegaflow-server/src/service.rs b/pegaflow-server/src/service.rs index 1ca3c5a8..c14f47ac 100644 --- a/pegaflow-server/src/service.rs +++ b/pegaflow-server/src/service.rs @@ -337,6 +337,23 @@ impl Engine for GrpcEngineService { return Err(status); } + #[cfg(feature = "rdma")] + if metadatas.iter().any(|metadata| metadata.dma_buf.is_some()) { + let engine = Arc::clone(&self.engine); + let instance_id = req.instance_id.clone(); + let regions: Vec<_> = metadatas.iter().filter_map(|metadata| metadata.dma_buf.clone()).collect(); + let result = tokio::task::spawn_blocking(move || { + engine.register_imported_device_memory(&instance_id, req.device_id, ®ions) + }).await + .map_err(|error| Status::internal(format!("DMA-BUF registration worker failed: {error}"))) + .and_then(|result| result.map_err(Self::map_engine_error)); + if let Err(error) = result { + let _ = self.engine.unregister_instance(&req.instance_id); + self.registry.drop_instance(req.instance_id).await; + return Err(error); + } + } + Ok(Response::new(Self::build_register_context_response())) } .await; @@ -560,7 +577,48 @@ impl Engine for GrpcEngineService { req.block_hashes.len() ); - let outcome = if req.group_id > 0 && req.wait_for_full_prefix { + let outcome = if req.direct_gpu { + if req.group_id != 0 { + return Err(Status::invalid_argument( + "direct_gpu is only supported for dense attention group 0", + )); + } + #[cfg(feature = "rdma")] + { + let plan = self + .engine + .query_direct_gpu_plan(&req.instance_id, &req.block_hashes) + .await + .map_err(Self::map_engine_error)?; + let Some(plan) = plan else { + return Ok(Response::new(QueryResponse { + outcome: Some(query_response::Outcome::Ready(QueryReady { + num_hit_blocks: 0, + lease: Vec::new(), + hit_positions: Vec::new(), + })), + })); + }; + let num_hit_blocks = plan.block_count() as u64; + let lease = self + .engine + .create_direct_query_lease(&req.instance_id, plan) + .map_err(Self::map_engine_error)? + .to_bytes() + .to_vec(); + query_response::Outcome::Ready(QueryReady { + num_hit_blocks, + lease, + hit_positions: Vec::new(), + }) + } + #[cfg(not(feature = "rdma"))] + { + return Err(Status::failed_precondition( + "direct_gpu requires an RDMA-enabled server", + )); + } + } else if req.group_id > 0 && req.wait_for_full_prefix { // All-or-nothing membership fetch: the hash list is an exact // want-set and misses are pulled from SSD / remote peers, so // the query may report Loading before it resolves. A Ready @@ -1086,6 +1144,7 @@ mod tests { req_id: String::new(), wait_for_full_prefix: false, group_id: 0, + direct_gpu: false, }) .expect_err("empty req_id must be rejected before engine lookup"); @@ -1101,6 +1160,7 @@ mod tests { req_id: "request".to_string(), wait_for_full_prefix: false, group_id: 0, + direct_gpu: false, }) .expect("empty block_hashes are a valid zero-hit query"); } diff --git a/pegaflow-server/tests/common/mod.rs b/pegaflow-server/tests/common/mod.rs index f4c669db..f53c10d5 100644 --- a/pegaflow-server/tests/common/mod.rs +++ b/pegaflow-server/tests/common/mod.rs @@ -263,6 +263,7 @@ impl MockVllmRpcHarness { req_id: req_id.to_string(), wait_for_full_prefix: false, group_id: 0, + direct_gpu: false, }; match self.scheduler.query_prefetch(request.clone()).await { Ok(response) => Ok(RpcExchange { diff --git a/pegaflow-transfer/src/cuda_lib/driver.rs b/pegaflow-transfer/src/cuda_lib/driver.rs index 94239232..c8cbefc9 100644 --- a/pegaflow-transfer/src/cuda_lib/driver.rs +++ b/pegaflow-transfer/src/cuda_lib/driver.rs @@ -63,3 +63,14 @@ pub fn cu_get_dma_buf_fd(ptr: NonNull, len: usize) -> Result { )) } } + +pub fn cu_get_address_range(ptr: u64) -> Result<(u64, usize)> { + let mut base = 0; + let mut len = 0; + let ret = unsafe { cuda_sys::cuMemGetAddressRange_v2(&mut base, &mut len, ptr) }; + if ret == cuda_sys::cudaError_enum::CUDA_SUCCESS { + Ok((base, len)) + } else { + Err(CudaDriverError::new(ret as u32, "cuMemGetAddressRange")) + } +} diff --git a/pegaflow-transfer/src/dma_buf.rs b/pegaflow-transfer/src/dma_buf.rs new file mode 100644 index 00000000..15cbb32d --- /dev/null +++ b/pegaflow-transfer/src/dma_buf.rs @@ -0,0 +1,77 @@ +use std::any::Any; +use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd}; +use std::ptr::NonNull; +use std::sync::Arc; + +use crate::cuda_lib::driver::{cu_get_address_range, cu_get_dma_buf_fd}; +use crate::{Result, TransferError}; + +/// An owner-exported DMA-BUF mapped through CUDA IPC in this process. +/// The mapping owner must retain the CUDA IPC tensor until all MRs are gone. +pub struct CudaDmaBuf { + pub ptr: u64, + pub len: usize, + fd: OwnedFd, + _mapping_owner: Arc, +} + +impl CudaDmaBuf { + /// Validate that the FD covers the complete imported CUDA allocation. + pub fn from_import( + ptr: u64, + len: usize, + fd: OwnedFd, + mapping_owner: Arc, + ) -> Result { + let (base, allocation_len) = + cu_get_address_range(ptr).map_err(|error| TransferError::Backend(error.to_string()))?; + if base != ptr || allocation_len != len || len == 0 { + return Err(TransferError::InvalidArgument( + "DMA-BUF range does not match the imported CUDA allocation", + )); + } + Ok(Self { + ptr, + len, + fd, + _mapping_owner: mapping_owner, + }) + } + + pub(crate) fn fd(&self) -> RawFd { + self.fd.as_raw_fd() + } +} + +impl std::fmt::Debug for CudaDmaBuf { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaDmaBuf") + .field("ptr", &self.ptr) + .field("len", &self.len) + .finish_non_exhaustive() + } +} + +/// Export in the allocation-owning process with its CUDA context current. +/// Returns an owned FD, allocation base and full allocation length. +pub fn export_cuda_dma_buf(ptr: u64, len: usize) -> Result<(OwnedFd, u64, usize)> { + let (base, allocation_len) = + cu_get_address_range(ptr).map_err(|error| TransferError::Backend(error.to_string()))?; + if len == 0 + || ptr < base + || ptr.checked_add(len as u64).is_none() + || base.checked_add(allocation_len as u64).is_none() + || ptr + len as u64 > base + allocation_len as u64 + { + return Err(TransferError::InvalidArgument( + "CUDA view exceeds its allocation", + )); + } + let base_ptr = NonNull::new(base as *mut std::ffi::c_void) + .ok_or(TransferError::InvalidArgument("CUDA allocation is null"))?; + let fd = cu_get_dma_buf_fd(base_ptr, allocation_len).map_err(|error| { + TransferError::Backend(format!("CUDA owner DMA-BUF export failed: {error}")) + })?; + // CUDA returned a new process-local descriptor, uniquely owned here. + Ok((unsafe { OwnedFd::from_raw_fd(fd) }, base, allocation_len)) +} diff --git a/pegaflow-transfer/src/engine.rs b/pegaflow-transfer/src/engine.rs index d5b770ad..8bb01886 100644 --- a/pegaflow-transfer/src/engine.rs +++ b/pegaflow-transfer/src/engine.rs @@ -121,6 +121,30 @@ impl TransferEngine { Ok(()) } + /// Register allocation-owner exports received over a local FD channel. + pub fn register_dma_buf_memory( + &self, + regions: &[std::sync::Arc], + device_id: u8, + ) -> Result<()> { + let mut registered = Vec::with_capacity(regions.len()); + for region in regions { + let ptr = NonNull::new(region.ptr as *mut u8) + .ok_or(TransferError::InvalidArgument("CUDA allocation is null"))?; + match self + .backend + .register_dma_buf_memory(region.clone(), device_id) + { + Ok(()) => registered.push(ptr), + Err(error) => { + let _ = self.unregister_memory(®istered); + return Err(error); + } + } + } + Ok(()) + } + pub fn unregister_memory(&self, ptrs: &[NonNull]) -> Result<()> { for &ptr in ptrs { self.backend.unregister_memory(ptr)?; @@ -184,6 +208,22 @@ impl TransferEngine { self.backend.batch_transfer_async(op, remote_addr, descs) } + /// Submit a batch whose local pointers target CUDA device memory. + /// + /// Device pointers cannot be classified with the CPU `move_pages(2)` NUMA + /// query used by the host staging path. Use the same NIC round-robin + /// fallback as an unknown NUMA placement while keeping host transfers + /// NUMA-aware. + pub fn batch_transfer_gpu_async( + &self, + op: TransferOp, + remote_addr: &str, + descs: &[TransferDesc], + ) -> Result>>> { + self.backend + .batch_transfer_gpu_async(op, remote_addr, descs) + } + /// Number of active RC queue pairs across all NICs. pub fn num_qps(&self) -> usize { self.backend.num_qps() diff --git a/pegaflow-transfer/src/lib.rs b/pegaflow-transfer/src/lib.rs index 098cfd5f..0b4fa127 100644 --- a/pegaflow-transfer/src/lib.rs +++ b/pegaflow-transfer/src/lib.rs @@ -1,3 +1,4 @@ +mod dma_buf; mod engine; mod error; mod rc_backend; @@ -8,6 +9,7 @@ mod cuda_sys; mod cudart_sys; pub mod v2; +pub use dma_buf::{CudaDmaBuf, export_cuda_dma_buf}; pub use engine::{ ConnectionStatus, HandshakeMetadata, MemoryRegion, TransferDesc, TransferEngine, TransferOp, }; diff --git a/pegaflow-transfer/src/rc_backend/mod.rs b/pegaflow-transfer/src/rc_backend/mod.rs index 4e88cafd..13192c66 100644 --- a/pegaflow-transfer/src/rc_backend/mod.rs +++ b/pegaflow-transfer/src/rc_backend/mod.rs @@ -34,6 +34,10 @@ impl NicGroup { let idx = self.rr_counter.fetch_add(1, Ordering::Relaxed); self.nic_indices[idx % self.nic_indices.len()] } + + fn next_position(&self) -> usize { + self.rr_counter.fetch_add(1, Ordering::Relaxed) % self.nic_indices.len() + } } struct NumaRoundRobin { @@ -169,6 +173,7 @@ impl RcBackend { base_ptr: raw, len, mrs, + owner: None, })?; debug!( "memory registered: ptr={:#x}, len={}, nics={}", @@ -179,6 +184,80 @@ impl RcBackend { Ok(()) } + pub(crate) fn register_dma_buf_memory( + &self, + region: Arc, + device_id: u8, + ) -> Result<()> { + let ptr = NonNull::new(region.ptr as *mut u8) + .ok_or(TransferError::InvalidArgument("CUDA allocation is null"))?; + let len = region.len; + if len == 0 { + return Err(TransferError::InvalidArgument("len must be non-zero")); + } + + // Replacing an allocation's MR would invalidate exchanged rkeys. + let raw = ptr.as_ptr() as u64; + if self.state.lock().registered.contains_exact(raw, len) { + return Err(TransferError::InvalidArgument( + "CUDA allocation is already registered", + )); + } + + let attrs = crate::cuda_lib::rt::cudaPointerGetAttributes(ptr.cast()).map_err(|error| { + TransferError::Backend(format!( + "CUDA pointer validation failed for device {device_id}: {error}" + )) + })?; + if attrs.type_ != crate::cuda_lib::rt::cudaMemoryTypeDevice { + return Err(TransferError::Backend( + "direct-GPU registration requires a CUDA device pointer".to_string(), + )); + } + if attrs.device != i32::from(device_id) { + return Err(TransferError::Backend(format!( + "CUDA pointer belongs to device {}, requested device {device_id}", + attrs.device + ))); + } + + let mut mrs = Vec::with_capacity(self.nic_count()); + for runtime in &self.runtimes { + let mr = unsafe { + runtime.pd.reg_dmabuf_mr( + 0, + len, + raw, + region.fd(), + AccessFlags::LocalWrite | AccessFlags::RemoteWrite | AccessFlags::RemoteRead, + ) + } + .map_err(|error| { + TransferError::Backend(format!( + "GPUDirect RDMA DMA-BUF registration failed on NIC {}: {error}", + runtime.nic_name + )) + }); + mrs.push(mr?); + } + + let mut state = self.state.lock(); + Arc::make_mut(&mut state.registered).insert(RegisteredMemoryEntry { + base_ptr: raw, + len, + mrs, + owner: Some(region), + })?; + info!( + "CUDA memory registered for direct RDMA: ptr={:#x}, len={}, device={}, nics={}", + raw, + len, + device_id, + self.nic_count() + ); + Ok(()) + } + pub(crate) fn unregister_memory(&self, ptr: NonNull) -> Result<()> { let raw = ptr.as_ptr() as u64; let mut state = self.state.lock(); @@ -389,12 +468,23 @@ impl RcBackend { .map(|c| c.local_nics.clone()) } - /// Remove connection state on transfer failure. The connection owns its - /// sessions, so in-flight work keeps its QPs alive through their Arcs. + /// Remove connection state on transfer failure and reset all of its QPs. + /// Resetting prevents accepted work from continuing to access caller-owned + /// buffers after a timeout or transport error. pub(crate) fn invalidate_connection(&self, remote_addr: &str) { - let mut state = self.state.lock(); - if state.addr_connections.remove(remote_addr).is_some() { + let connection = self.state.lock().addr_connections.remove(remote_addr); + if let Some(connection) = connection { info!("connection invalidated: remote={remote_addr}"); + for nic in connection.nics { + for session in nic.sessions.iter() { + if let Err(error) = session.reset() { + warn!( + "failed to reset invalidated RC session: remote={remote_addr} qpn={} error={error}", + session.local_endpoint.qp_num + ); + } + } + } } } @@ -408,6 +498,29 @@ impl RcBackend { op: TransferOp, remote_addr: &str, descs: &[TransferDesc], + ) -> Result>>> { + self.batch_transfer_async_inner(op, remote_addr, descs, true) + } + + /// Submit a batch whose local pointers are CUDA device addresses. + /// `move_pages(2)` only classifies CPU pages and adds overhead (or returns + /// UNKNOWN) for GPU mappings, so direct GPU transfers use the same + /// round-robin fallback as unknown NUMA placement. + pub(crate) fn batch_transfer_gpu_async( + &self, + op: TransferOp, + remote_addr: &str, + descs: &[TransferDesc], + ) -> Result>>> { + self.batch_transfer_async_inner(op, remote_addr, descs, false) + } + + fn batch_transfer_async_inner( + &self, + op: TransferOp, + remote_addr: &str, + descs: &[TransferDesc], + numa_aware: bool, ) -> Result>>> { if descs.is_empty() { return Ok(Vec::new()); @@ -417,9 +530,26 @@ impl RcBackend { // NUMA-aware NIC assignment: query the NUMA node of each descriptor's // first page and route to a NIC on the same NUMA node. - let mut per_nic: Vec> = (0..nic_count).map(|_| Vec::new()).collect(); - if self.numa_rr.single_numa { - // All NICs on one NUMA node — skip move_pages, plain round-robin. + let per_nic_capacity = if numa_aware { + 0 + } else { + descs.len().div_ceil(nic_count) + }; + let mut per_nic: Vec> = (0..nic_count) + .map(|_| Vec::with_capacity(per_nic_capacity)) + .collect(); + if !numa_aware { + // GPU addresses cannot be classified by move_pages. Reserve one + // atomic increment for the whole batch; per-descriptor fetch_add + // adds contention without improving distribution. + let start = self.numa_rr.fallback.next_position(); + for (index, &desc) in descs.iter().enumerate() { + let nic_idx = self.numa_rr.fallback.nic_indices[(start + index) % nic_count]; + per_nic[nic_idx].push(desc); + } + } else if self.numa_rr.single_numa { + // All host NICs on one NUMA node — skip move_pages, plain + // round-robin as in the original host staging path. for &desc in descs { let nic_idx = self.numa_rr.fallback.next(); per_nic[nic_idx].push(desc); @@ -494,9 +624,10 @@ impl RcBackend { return Err(TransferError::InvalidArgument("len must be non-zero")); } - let local_mr = registered - .find_mr(nic.nic_idx, local_ptr, len) - .ok_or(TransferError::MemoryNotRegistered { ptr: local_ptr })?; + let (local_mr, owner) = + registered + .find_mr_and_owner(nic.nic_idx, local_ptr, len) + .ok_or(TransferError::MemoryNotRegistered { ptr: local_ptr })?; let remote_rkey = nic.remote_memory.find_rkey(remote_ptr, len).ok_or( TransferError::InvalidArgument("remote memory not found in handshake snapshot"), @@ -505,6 +636,7 @@ impl RcBackend { let bucket = nic.rot.wrapping_add(i) % n; buckets[bucket].push(RdmaOp { local_mr, + _owner: owner, local_ptr, remote_ptr, len, @@ -533,7 +665,15 @@ impl RcBackend { // --- Submit outside lock --- let mut receivers = Vec::with_capacity(nic_work.len()); for (session, prepared) in nic_work { - receivers.push(session.transfer_batch_async(prepared, op)?); + match session.transfer_batch_async(prepared, op) { + Ok(receiver) => receivers.push(receiver), + Err(error) => { + // Return every accepted completion even if a later session rejects submission. + let (tx, rx) = oneshot::channel(); + let _ = tx.send(Err(error)); + receivers.push(rx); + } + } } Ok(receivers) } diff --git a/pegaflow-transfer/src/rc_backend/session.rs b/pegaflow-transfer/src/rc_backend/session.rs index 2250d0f1..83accfb8 100644 --- a/pegaflow-transfer/src/rc_backend/session.rs +++ b/pegaflow-transfer/src/rc_backend/session.rs @@ -24,7 +24,10 @@ use super::runtime::RcRuntime; use crate::engine::{RcEndpoint, TransferOp}; use crate::error::{Result, TransferError}; -const MAX_WR_CHAIN_OPS: usize = 4; +// Keep each post_send chain within the QP's read-atomic window. A longer +// chain amortizes ibv_post_send setup while the surrounding loop still caps +// total outstanding WRs at MAX_SEND_WR. +const MAX_WR_CHAIN_OPS: usize = 16; const MAX_SEND_WR: u32 = 128; // One QP per CQ; all WRs are signaled, so CQ depth = SQ depth suffices. // One-sided RDMA only: no recvs are ever posted, so the send CQ doubles as @@ -36,6 +39,7 @@ const PSN_MASK: u32 = 0x00ff_ffff; pub(crate) struct RdmaOp { pub(crate) local_mr: Arc, + pub(crate) _owner: Option>, pub(crate) local_ptr: u64, pub(crate) remote_ptr: u64, pub(crate) len: usize, @@ -202,6 +206,16 @@ impl RcSession { Ok(done_rx) } + /// Move the QP to RESET so outstanding work cannot continue writing to a + /// caller-owned buffer. The next transfer will establish a fresh session. + pub(crate) fn reset(&self) -> Result<()> { + let mut qp = self.qp.lock(); + let mut attr = QueuePairAttribute::new(); + attr.setup_state(QueuePairState::Reset); + qp.modify(&attr) + .map_err(|error| TransferError::Backend(format!("failed to reset RC QP: {error}"))) + } + // The worker holds only a Weak ref: a strong Arc here would keep the // session (and its cmd_tx) alive forever, so recv() could never // disconnect and every invalidated connection would leak its thread + QP. @@ -301,22 +315,41 @@ impl RcSession { let mut next_idx = 0usize; let mut next_wr_id = 1_u64; - let mut inflight: HashMap = HashMap::new(); + // The send window is capped at MAX_SEND_WR, so reserve the complete + // window once instead of growing and rehashing while completions are + // being tracked. + let mut inflight: HashMap = + HashMap::with_capacity(total_ops.min(MAX_SEND_WR as usize)); let mut transferred = 0usize; + let mut first_error = None; - while next_idx < total_ops || !inflight.is_empty() { - while next_idx < total_ops && inflight.len() < MAX_SEND_WR as usize { + // After the first error no new work is posted. Keep polling only to + // drain completions that were accepted before the error, then return. + while (first_error.is_none() && next_idx < total_ops) || !inflight.is_empty() { + while first_error.is_none() + && next_idx < total_ops + && inflight.len() < MAX_SEND_WR as usize + { let available = MAX_SEND_WR as usize - inflight.len(); let remaining = total_ops - next_idx; let chain_len = MAX_WR_CHAIN_OPS.min(available).min(remaining); - let mut qp = session.qp.lock(); - let posted = Self::post_rdma_wr_chain( - &mut qp, - &ops[next_idx..next_idx + chain_len], - next_wr_id, - op, - )?; - drop(qp); + let posted_result = { + let mut qp = session.qp.lock(); + Self::post_rdma_wr_chain( + &mut qp, + &ops[next_idx..next_idx + chain_len], + next_wr_id, + op, + ) + }; + let posted = match posted_result { + Ok(posted) => posted, + Err(error) => { + first_error = Some(error); + let _ = session.reset(); + break; + } + }; if posted == 0 { break; } @@ -335,16 +368,26 @@ impl RcSession { let Some(bytes) = inflight.remove(&wc.wr_id()) else { continue; }; - if wc.status() != WorkCompletionStatus::Success as u32 { - return Err(TransferError::Backend(format!( + if wc.status() != WorkCompletionStatus::Success as u32 + && first_error.is_none() + { + first_error = Some(TransferError::Backend(format!( "send completion failed: local_qpn={}, status={}, opcode={}, vendor_err={}", session.local_endpoint.qp_num, wc.status(), wc.opcode(), wc.vendor_err() ))); + if let Err(error) = session.reset() { + warn!( + "failed to reset RC QP after send completion error: local_qpn={} error={error}", + session.local_endpoint.qp_num + ); + } + } + if wc.status() == WorkCompletionStatus::Success as u32 { + transferred = transferred.saturating_add(bytes); } - transferred = transferred.saturating_add(bytes); } if !did_work { std::hint::spin_loop(); @@ -354,14 +397,18 @@ impl RcSession { std::hint::spin_loop(); } Err(error) => { - return Err(TransferError::Backend(format!( - "poll send CQ failed: local_qpn={}, {error}", - session.local_endpoint.qp_num - ))); + first_error.get_or_insert_with(|| { + TransferError::Backend(format!( + "poll send CQ failed: local_qpn={}, {error}", + session.local_endpoint.qp_num + )) + }); + let _ = session.reset(); + inflight.clear(); } } } - Ok(transferred) + first_error.map_or(Ok(transferred), Err) } } diff --git a/pegaflow-transfer/src/rc_backend/state.rs b/pegaflow-transfer/src/rc_backend/state.rs index 5f2c0c83..ac7459d0 100644 --- a/pegaflow-transfer/src/rc_backend/state.rs +++ b/pegaflow-transfer/src/rc_backend/state.rs @@ -14,6 +14,7 @@ pub(super) struct RegisteredMemoryEntry { pub(super) len: usize, /// One MR per NIC (different PDs → different rkeys). pub(super) mrs: Vec>, + pub(super) owner: Option>, } /// Local registered memory, ordered by base pointer. Insertion rejects @@ -68,24 +69,35 @@ impl LocalMemoryMap { self.entries.remove(&base_ptr) } + pub(super) fn contains_exact(&self, base_ptr: u64, len: usize) -> bool { + self.entries + .get(&base_ptr) + .is_some_and(|entry| entry.len == len) + } + /// Entries in base_ptr order. pub(super) fn iter(&self) -> impl Iterator { self.entries.values() } - /// Find the MR (for `nic_idx`) of the region fully covering `[ptr, ptr+len)`. - pub(super) fn find_mr( + /// Find the MR (for `nic_idx`) and optional owner of the region fully + /// covering `[ptr, ptr+len)`. + /// + /// The transfer hot path needs both values for every descriptor. Resolve + /// the covering B-tree entry once so device batches do not pay for two + /// identical predecessor lookups. + pub(super) fn find_mr_and_owner( &self, nic_idx: usize, ptr: u64, len: usize, - ) -> Option> { + ) -> Option<(Arc, Option>)> { self.find_entry(ptr, len) - .map(|entry| Arc::clone(&entry.mrs[nic_idx])) + .map(|entry| (Arc::clone(&entry.mrs[nic_idx]), entry.owner.clone())) } /// Non-overlap makes the predecessor the only possible covering region. - fn find_entry(&self, ptr: u64, len: usize) -> Option<&RegisteredMemoryEntry> { + pub(super) fn find_entry(&self, ptr: u64, len: usize) -> Option<&RegisteredMemoryEntry> { let end = ptr.checked_add(len as u64)?; let (_, entry) = self.entries.range(..=ptr).next_back()?; (end <= entry.base_ptr + entry.len as u64).then_some(entry) @@ -279,6 +291,7 @@ mod tests { base_ptr, len, mrs: Vec::new(), + owner: None, } } @@ -287,6 +300,8 @@ mod tests { let mut map = LocalMemoryMap::default(); map.insert(entry(0x1000, 0x100)).expect("first region"); map.insert(entry(0x3000, 0x100)).expect("disjoint region"); + assert!(map.contains_exact(0x1000, 0x100)); + assert!(!map.contains_exact(0x1000, 0x80)); // Overlap with predecessor and successor both rejected. assert!(map.insert(entry(0x10ff, 0x10)).is_err()); diff --git a/python/README.md b/python/README.md index c0d1697d..542cd9ff 100644 --- a/python/README.md +++ b/python/README.md @@ -71,6 +71,29 @@ llm = LLM( ) ``` +To enable direct GPU P2P reads, add the connector extra configuration. The +value is read from vLLM's `--kv-transfer-config` in both the scheduler and +worker processes, so no process environment variable is required: + +```json +{ + "kv_connector": "PegaKVConnector", + "kv_role": "kv_both", + "kv_connector_module_path": "pegaflow.connector", + "kv_connector_extra_config": { + "pegaflow.direct_gpu_rdma": true + } +} +``` + +Leave the key out, or set it to `false`, to keep the host-staging path. Direct +GPU mode currently supports dense attention cache group 0 only. A direct load +failure is reported to vLLM for recomputation; it is not silently retried via +host staging. The direct load duration metric covers the full batch, including +remote metadata/completion and any concurrent local H2D work; it is not a raw +GPU-RDMA bandwidth measurement. See [P2P performance notes](../docs/p2p.md#direct-gpu-performance-notes) +for the A/B comparison method. + #### Connector Modes `PegaKVConnector` defaults to `read_write`: it queries PegaFlow for reusable KV diff --git a/python/pegaflow/connector/__init__.py b/python/pegaflow/connector/__init__.py index fbb383db..21413434 100644 --- a/python/pegaflow/connector/__init__.py +++ b/python/pegaflow/connector/__init__.py @@ -132,6 +132,9 @@ def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): "pegaflow.wait_for_full_prefix", False ) ) + direct_gpu_rdma = bool( + vllm_config.kv_transfer_config.get_from_extra_config("pegaflow.direct_gpu_rdma", False) + ) default_endpoint = f"{server_host}:{server_port}" tp_shards = TpShardTopology.from_config( default_endpoint=default_endpoint, @@ -177,6 +180,7 @@ def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): pp_size=pp_size, mode=mode, wait_for_full_prefix=wait_for_full_prefix, + direct_gpu_rdma=direct_gpu_rdma, tp_shards=tp_shards, hash_block_size=hash_block_size, ) @@ -238,7 +242,7 @@ def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): "tp_rank=%s tp_size=%d pp_rank=%d pp_size=%d world_size=%d namespace=%s " "is_mla=%s collapse_mla_tp=%s transfer_backend=%s dcp_world_size=%d " "pcp_world_size=%d dcp_rank=%d tp_shard=%d/%d " - "mode=%s wait_for_full_prefix=%s", + "mode=%s wait_for_full_prefix=%s direct_gpu_rdma=%s", role.name, instance_id, device_id if device_id is not None else "cpu", @@ -258,6 +262,7 @@ def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): tp_shards.shard_count, mode.value, wait_for_full_prefix, + direct_gpu_rdma, ) # ============================== diff --git a/python/pegaflow/connector/common.py b/python/pegaflow/connector/common.py index f4c2f39d..cb29e211 100644 --- a/python/pegaflow/connector/common.py +++ b/python/pegaflow/connector/common.py @@ -143,6 +143,7 @@ class ConnectorContext: pp_size: int = 1 mode: PegaConnectorMode = PegaConnectorMode.READ_WRITE wait_for_full_prefix: bool = False + direct_gpu_rdma: bool = False tp_shards: TpShardTopology | None = None # Token span of one `Request.block_hashes` entry; `None` means one per # scheduler block. diff --git a/python/pegaflow/connector/scheduler.py b/python/pegaflow/connector/scheduler.py index 2179fe6a..03fe0612 100644 --- a/python/pegaflow/connector/scheduler.py +++ b/python/pegaflow/connector/scheduler.py @@ -138,6 +138,9 @@ def __init__( ) self._tp_shard_client = TpShardQueryClient(engine_clients) self._cache_groups = CacheGroupLayout.from_config(kv_cache_config) + self._direct_gpu_rdma = context.direct_gpu_rdma + if self._direct_gpu_rdma and self._cache_groups.has_recurrent_state: + raise ValueError("pegaflow.direct_gpu_rdma only supports dense attention cache group 0") if self._cache_groups.has_recurrent_state and (pd_tail_save or pd_tail_load): raise ValueError("P/D tail-block caching is not supported with HMA") self._gpu_block_pool = None @@ -1170,6 +1173,7 @@ def _count_available_block_prefix( block_hash_list, req_id, self._ctx.wait_for_full_prefix, + self._direct_gpu_rdma, ) if ready is None: if req_id not in self._prefetch_start_times: diff --git a/python/pegaflow/connector/tp_shards.py b/python/pegaflow/connector/tp_shards.py index b032631f..aff6b7c2 100644 --- a/python/pegaflow/connector/tp_shards.py +++ b/python/pegaflow/connector/tp_shards.py @@ -33,16 +33,18 @@ def query( block_hashes: list[bytes], req_id: str, wait_for_full_prefix: bool, + direct_gpu: bool = False, ) -> ShardedQueryReady | None: results: list[QueryReady] = [] try: for client in self._clients: - result = client.query_prefetch( - instance_id, - block_hashes, - req_id=req_id, - wait_for_full_prefix=wait_for_full_prefix, - ) + query_kwargs = { + "req_id": req_id, + "wait_for_full_prefix": wait_for_full_prefix, + } + if direct_gpu: + query_kwargs["direct_gpu"] = True + result = client.query_prefetch(instance_id, block_hashes, **query_kwargs) if isinstance(result, QueryLoading): self.release(tuple(ready.lease for ready in results), req_id) return None @@ -65,12 +67,13 @@ def query( for index, (client, result) in enumerate(zip(self._clients, results, strict=True)): if result.num_hit_blocks == common_blocks: continue - exact = client.query_prefetch( - instance_id, - exact_hashes, - req_id=f"{req_id}:tp-common-{common_blocks}", - wait_for_full_prefix=False, - ) + query_kwargs = { + "req_id": f"{req_id}:tp-common-{common_blocks}", + "wait_for_full_prefix": False, + } + if direct_gpu: + query_kwargs["direct_gpu"] = True + exact = client.query_prefetch(instance_id, exact_hashes, **query_kwargs) if not isinstance(exact, QueryReady): raise RuntimeError( f"TP shard {index} could not lease the common {common_blocks}-block prefix" diff --git a/python/pegaflow/connector/worker.py b/python/pegaflow/connector/worker.py index 73bf7aa5..780a535e 100644 --- a/python/pegaflow/connector/worker.py +++ b/python/pegaflow/connector/worker.py @@ -7,6 +7,7 @@ import threading import time from collections.abc import Iterable, Iterator +from contextlib import nullcontext from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal @@ -22,6 +23,7 @@ logger, parse_env_int, ) +from pegaflow.dma_buf import DmaBufExports from pegaflow.ipc_wrapper import CudaIPCWrapper from pegaflow.pegaflow import PyLoadState @@ -307,6 +309,11 @@ def unregister_context(self) -> None: self._registered_layers.clear() def register_kv_caches(self, kv_caches: dict[str, Any]): + direct = self._ctx.direct_gpu_rdma + with DmaBufExports() if direct else nullcontext() as exports: + self._register_kv_caches(kv_caches, exports) + + def _register_kv_caches(self, kv_caches: dict[str, Any], exports: DmaBufExports | None): """Register exactly the KV caches vLLM built on this device. The engine derives the instance-wide layer-id space once every worker @@ -397,7 +404,11 @@ def register_kv_caches(self, kv_caches: dict[str, Any]): registration_tensor.storage_offset(), ) - wrapper = CudaIPCWrapper(registration_tensor) + wrapper = ( + CudaIPCWrapper(registration_tensor, exports) + if exports + else CudaIPCWrapper(registration_tensor) + ) wrapper_bytes = pickle.dumps(wrapper) registration = _infer_kv_cache_registration( diff --git a/python/pegaflow/dma_buf.py b/python/pegaflow/dma_buf.py new file mode 100644 index 00000000..26a8d483 --- /dev/null +++ b/python/pegaflow/dma_buf.py @@ -0,0 +1,90 @@ +"""Registration-scoped, same-host file descriptor handoff for CUDA DMA-BUF.""" + +import array +import os +import secrets +import socket +import tempfile +import threading +from contextlib import AbstractContextManager + + +class DmaBufExports(AbstractContextManager): + """Keep owner exports alive until the registration RPC has consumed them.""" + + def __init__(self): + self._directory = tempfile.TemporaryDirectory(prefix="pega-dmabuf-") + self.address = os.path.join(self._directory.name, "fd.sock") + self._socket = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET) + self._socket.bind(self.address) + self._socket.listen() + self._socket.settimeout(0.2) + self._fds: dict[bytes, int] = {} + self._lock = threading.Lock() + self._stop = threading.Event() + self._thread = threading.Thread(target=self._serve, name="pega-dmabuf", daemon=True) + self._thread.start() + + def add(self, fd: int) -> tuple[str, bytes]: + """Take ownership of fd and return a serializable one-use capability.""" + token = secrets.token_bytes(32) + with self._lock: + self._fds[token] = fd + return self.address, token + + def _serve(self) -> None: + while not self._stop.is_set(): + try: + connection, _ = self._socket.accept() + except TimeoutError: + continue + with connection: + connection.settimeout(5) + try: + token = connection.recv(33) + with self._lock: + fd = self._fds.pop(token, None) + if fd is None: + connection.send(b"E") + continue + try: + connection.sendmsg( + [b"F"], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, array.array("i", [fd]))] + ) + finally: + os.close(fd) + except OSError: + # The RPC may have been cancelled while this peer connected. + continue + + def __exit__(self, *exc): + self._stop.set() + self._thread.join() + self._socket.close() + with self._lock: + for fd in self._fds.values(): + os.close(fd) + self._fds.clear() + self._directory.cleanup() + + +def receive_dma_buf(address: str, token: bytes) -> int: + """Receive a new process-local descriptor; caller must close it.""" + fds = array.array("i") + try: + with socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET) as connection: + connection.settimeout(10) + connection.connect(address) + connection.sendall(token) + message, ancillary, flags, _ = connection.recvmsg( + 1, socket.CMSG_SPACE(fds.itemsize), socket.MSG_CMSG_CLOEXEC + ) + for level, kind, data in ancillary: + if level == socket.SOL_SOCKET and kind == socket.SCM_RIGHTS: + fds.frombytes(data[: len(data) - len(data) % fds.itemsize]) + if message != b"F" or flags & (socket.MSG_CTRUNC | socket.MSG_TRUNC) or len(fds) != 1: + raise RuntimeError("CUDA DMA-BUF handoff did not return exactly one descriptor") + return fds.pop() + finally: + for fd in fds: + os.close(fd) diff --git a/python/pegaflow/ipc_wrapper.py b/python/pegaflow/ipc_wrapper.py index ce6834ac..9b47ced6 100644 --- a/python/pegaflow/ipc_wrapper.py +++ b/python/pegaflow/ipc_wrapper.py @@ -12,6 +12,8 @@ import torch +from pegaflow.dma_buf import DmaBufExports, receive_dma_buf + class CudaIPCWrapper: """Wrapper for CUDA IPC handle with tensor metadata. @@ -98,7 +100,7 @@ def _get_device_index_from_uuid(device_uuid: str) -> int: ) return device_index - def __init__(self, tensor: torch.Tensor): + def __init__(self, tensor: torch.Tensor, dma_buf_exports: DmaBufExports | None = None): """Create IPC wrapper from a CUDA tensor. Args: @@ -123,6 +125,25 @@ def __init__(self, tensor: torch.Tensor): device_index = tensor.device.index self.device_uuid = CudaIPCWrapper._get_device_uuid(device_index) + self.dma_buf = None + if dma_buf_exports is not None: + from pegaflow.pegaflow import export_cuda_dma_buf + + with torch.cuda.device(tensor.device): + fd, base, allocation_bytes = export_cuda_dma_buf( + storage.data_ptr(), storage.nbytes() + ) + address, token = dma_buf_exports.add(fd) + self.dma_buf = (address, token, tensor.data_ptr() - base, allocation_bytes) + + def detach_dma_buf(self) -> tuple[int, int, int] | None: + """Transfer an owned FD to the local service during registration.""" + metadata = getattr(self, "dma_buf", None) + if metadata is None: + return None + address, token, offset, allocation_bytes = metadata + return receive_dma_buf(address, token), offset, allocation_bytes + def to_tensor(self) -> torch.Tensor: """Reconstruct tensor from IPC handle. diff --git a/python/pegaflow/pegaflow.pyi b/python/pegaflow/pegaflow.pyi index b188fd2e..253626e9 100644 --- a/python/pegaflow/pegaflow.pyi +++ b/python/pegaflow/pegaflow.pyi @@ -224,6 +224,7 @@ class EngineRpcClient: req_id: str, wait_for_full_prefix: bool = False, group_id: int = 0, + direct_gpu: bool = False, ) -> QueryLoading | QueryReady: """Query prefix cache hits with SSD prefetch support. @@ -389,3 +390,10 @@ class PdRdmaEngine: def num_domains(self) -> int: ... def num_groups(self) -> int: ... def aggregated_link_speed(self) -> int: ... + +def export_cuda_dma_buf(ptr: int, len: int) -> tuple[int, int, int]: + """Export in the allocation owner: (owned fd, allocation base, allocation bytes). + + The caller must bind the owning CUDA context and close the returned fd. + """ + ... diff --git a/python/src/lib.rs b/python/src/lib.rs index c97d54b3..dbe04b26 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -476,7 +476,11 @@ impl EngineRpcClient { /// /// Returns: /// QueryLoading while backing fetch is in progress, otherwise QueryReady. - #[pyo3(signature = (instance_id, block_hashes, req_id, wait_for_full_prefix=false, group_id=0))] + #[pyo3(signature = (instance_id, block_hashes, req_id, wait_for_full_prefix=false, group_id=0, direct_gpu=false))] + #[allow( + clippy::too_many_arguments, + reason = "Python API preserves existing query arguments and adds an optional direct GPU flag" + )] fn query_prefetch( &self, py: Python<'_>, @@ -485,6 +489,7 @@ impl EngineRpcClient { req_id: String, wait_for_full_prefix: bool, group_id: u32, + direct_gpu: bool, ) -> PyResult> { let result = py.detach(|| { self.rt_handle.block_on(async { @@ -496,6 +501,7 @@ impl EngineRpcClient { req_id, wait_for_full_prefix, group_id, + direct_gpu, }) .await .map(|resp| resp.into_inner()) @@ -662,6 +668,16 @@ impl PyLoadState { } } +/// Export from the allocation-owning process with its CUDA context current. +#[cfg(feature = "rdma")] +#[pyfunction] +fn export_cuda_dma_buf(ptr: u64, len: usize) -> PyResult<(i32, u64, usize)> { + use std::os::fd::IntoRawFd; + let (fd, base, size) = pegaflow_transfer::export_cuda_dma_buf(ptr, len) + .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + Ok((fd.into_raw_fd(), base, size)) +} + /// A Python module implemented in Rust. #[pymodule] fn pegaflow(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -673,6 +689,8 @@ fn pegaflow(m: &Bound<'_, PyModule>) -> PyResult<()> { pd_rdma::add_classes(m)?; #[cfg(feature = "rdma")] rdma_v1::add_classes(m)?; + #[cfg(feature = "rdma")] + m.add_function(wrap_pyfunction!(export_cuda_dma_buf, m)?)?; // Register custom exceptions for error classification m.add("PegaFlowError", m.py().get_type::())?; m.add("PegaflowInternal", m.py().get_type::())?; diff --git a/python/tests/test_direct_gpu_rdma.py b/python/tests/test_direct_gpu_rdma.py new file mode 100644 index 00000000..a4239879 --- /dev/null +++ b/python/tests/test_direct_gpu_rdma.py @@ -0,0 +1,260 @@ +"""Owner process -> CUDA IPC/FD registration -> gRPC direct load -> GPU bytes. + +Run with PEGAFLOW_IB_DEVICE=mlx5_0 and CUDA_VISIBLE_DEVICES limited to test GPUs: + pytest -m integration tests/test_direct_gpu_rdma.py +Requires release server/metaserver binaries and the matching native extension. +""" + +import hashlib +import os +import pickle +import subprocess +import time +import uuid +from contextlib import ExitStack +from pathlib import Path +from urllib.error import URLError +from urllib.request import urlopen + +import pytest + +from .conftest import find_available_port, wait_for_server_ready + +pytestmark = [pytest.mark.integration, pytest.mark.gpu] + + +def _stop(process): + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + + +def test_ipc_owner_export_direct_rdma_bytes(tmp_path): + import torch + + from pegaflow.dma_buf import DmaBufExports + from pegaflow.ipc_wrapper import CudaIPCWrapper + from pegaflow.pegaflow import EngineRpcClient, PyLoadState + + nic = os.environ.get("PEGAFLOW_IB_DEVICE") + if not nic: + pytest.skip("set PEGAFLOW_IB_DEVICE to explicitly select test NICs") + assert torch.cuda.device_count() >= 2, "select two non-production GPUs" + root = Path(__file__).resolve().parents[2] + meta_port, meta_http, *ports = [find_available_port() for _ in range(8)] + env = { + **os.environ, + "PYTHONPATH": str(root / "python") + os.pathsep + os.environ.get("PYTHONPATH", ""), + } + with ExitStack() as stack: + + def start(name, argv): + log = stack.enter_context((tmp_path / f"{name}.log").open("w")) + process = subprocess.Popen(argv, env=env, stdout=log, stderr=subprocess.STDOUT) + stack.callback(_stop, process) + return process + + meta = start( + "meta", + [ + str(root / "target/release/pegaflow-metaserver"), + "--addr", + f"127.0.0.1:{meta_port}", + "--http-addr", + f"127.0.0.1:{meta_http}", + ], + ) + deadline = time.monotonic() + 30 + while True: + assert meta.poll() is None, (tmp_path / "meta.log").read_text() + try: + with urlopen(f"http://127.0.0.1:{meta_http}/health", timeout=1): + break + except URLError: + assert time.monotonic() < deadline, "MetaServer did not become ready" + time.sleep(0.05) + clients = [] + servers = [] + devices = [0, 1, 0] + for index, device in enumerate(devices): + servers.append( + start( + f"server{index}", + [ + str(root / "target/release/pegaflow-server"), + "--addr", + f"127.0.0.1:{ports[index]}", + "--http-addr", + f"127.0.0.1:{ports[3 + index]}", + "--pool-size", + "64mb", + "--devices", + str(device), + "--nics", + nic, + "--metaserver-addr", + f"http://127.0.0.1:{meta_port}", + "--disable-numa-affinity", + ], + ) + ) + endpoint = f"http://127.0.0.1:{ports[index]}" + assert wait_for_server_ready(endpoint), (tmp_path / f"server{index}.log").read_text() + clients.append(EngineRpcClient(endpoint)) + + namespace = f"ipc-dma-buf-{uuid.uuid4()}" + names = ["layer0", "layer1"] + tensors = [] + for index, (device, client) in enumerate(zip(devices, clients, strict=True)): + # Two split-K/V layers share one allocation, with an unaligned view offset. + allocation = torch.zeros( + 2 * 2 * 8 * 1024 + 8192, dtype=torch.uint8, device=f"cuda:{device}" + ) + layers = [ + allocation[128 + index * 16384 : 128 + (index + 1) * 16384].view(2, 8, 1024) + for index in range(2) + ] + tensors.append(layers) + if device == 0: + for layer_index, layer in enumerate(layers): + for segment in range(2): + for block in range(8): + layer[segment, block].fill_(1 + layer_index * 32 + segment * 8 + block) + torch.cuda.synchronize(device) + with DmaBufExports() as exports: + wrappers = [ + pickle.dumps(CudaIPCWrapper(layer, exports if device else None)) + for layer in layers + ] + ok, message = client.register_context_batch( + f"instance{index}", + namespace, + 0, + 0, + 1, + 1, + device, + names, + wrappers, + [8, 8], + [1024, 1024], + [8192, 8192], + [2, 2], + "direct", + False, + ) + assert ok, message + + hashes = [hashlib.sha256(f"{namespace}-{index}".encode()).digest() for index in range(8)] + # Disjoint owners force the existing ordered multi-node plan to be used. + for owner, indices in [(0, list(range(4))), (2, list(range(4, 8)))]: + ok, message = clients[owner].save( + f"instance{owner}", + 0, + 0, + 0, + [(name, indices, [hashes[i] for i in indices]) for name in names], + ) + assert ok, message + deadline = time.monotonic() + 30 + while True: + ready = clients[1].query_prefetch("instance1", hashes, "owner-ready", direct_gpu=True) + if getattr(ready, "num_hit_blocks", 0) == 8: + clients[1].release(ready.lease) + break + assert time.monotonic() < deadline, "remote owner never became queryable" + time.sleep(0.05) + + source_bytes = [source.cpu() for source in tensors[0]] + + def query(selected_hashes, *, direct=True): + ready = clients[1].query_prefetch( + "instance1", selected_hashes, str(uuid.uuid4()), direct_gpu=direct + ) + assert ready.num_hit_blocks == len(selected_hashes) + return ready.lease + + def load_and_check(loads, source_by_destination, *, success=True): + for layer in tensors[1]: + layer.zero_() + torch.cuda.synchronize(1) + state = PyLoadState() + ok, message = clients[1].load("instance1", 0, 1, state.shm_name(), [names], loads) + assert ok, message + deadline = time.monotonic() + 30 + while not state.is_ready(): + assert time.monotonic() < deadline, "direct RDMA did not complete" + time.sleep(0.01) + assert state.get_state() == (1 if success else -1) + for source, target in zip(source_bytes, tensors[1], strict=True): + expected = torch.zeros_like(target, device="cpu") + for destination, source_index in source_by_destination.items(): + expected[:, destination] = source[:, source_index] + assert torch.equal(target.cpu(), expected), "loaded GPU bytes differ" + + for destinations in ([3], [0, 1, 2], [6, 1, 4], [7, 5, 3, 1, 6, 4, 2, 0]): + load_and_check( + [(query(hashes[: len(destinations)]), [destinations])], + {destination: i for i, destination in enumerate(destinations)}, + ) + + # Seed only the local prefix. Also save hashes with no remote owner. + local_hashes = [ + hashlib.sha256(f"{namespace}-local-{i}".encode()).digest() for i in range(2) + ] + for source, target in zip(source_bytes, tensors[1], strict=True): + target.copy_(source) + torch.cuda.synchronize(1) + for saved_hashes in (hashes[:2], local_hashes): + ok, message = clients[1].save( + "instance1", 0, 0, 1, [(name, [0, 1], saved_hashes) for name in names] + ) + assert ok, message + deadline = time.monotonic() + 30 + while True: + ready = clients[1].query_prefetch( + "instance1", local_hashes, "local-ready", direct_gpu=True + ) + if ready.num_hit_blocks == 2: + clients[1].release(ready.lease) + break + assert time.monotonic() < deadline, "local saves did not become visible" + time.sleep(0.05) + + # Local prefix + two remote owners; None preserves position without a write. + for destinations in ([7, 5, 3, 1, 6, 4, 2, 0], [7, None, 3, 1, None, None, None, None]): + load_and_check( + [(query(hashes), [destinations])], + { + destination: i + for i, destination in enumerate(destinations) + if destination is not None + }, + ) + + # Cached, local-only direct, and remote-only leases work in either batch order. + for reverse in (False, True): + loads = [ + (query(local_hashes, direct=False), [[5, 2]]), + (query(local_hashes), [[6, 0]]), + (query(hashes[4:6]), [[7, 1]]), + ] + load_and_check(loads[::-1] if reverse else loads, {5: 0, 2: 1, 6: 0, 0: 1, 7: 4, 1: 5}) + + failed_lease = query(hashes) + _stop(servers[2]) + # The local copy and first remote segment must settle before the suffix error is signalled. + load_and_check([(failed_lease, [list(range(8))])], {i: i for i in range(4)}, success=False) + + _stop(servers[0]) + _stop(meta) + # No directory or remote owner is available: full RAM hits still work. + load_and_check([(query(local_hashes), [[4, 3]])], {4: 0, 3: 1}) + ready = clients[1].query_prefetch("instance1", hashes, "no-host-admission", direct_gpu=True) + assert ready.num_hit_blocks == 2, "direct-loaded suffix was inserted into the RAM cache" + clients[1].release(ready.lease) + ok, message = clients[1].unregister_context("instance1") + assert ok, message diff --git a/python/tests/test_dma_buf.py b/python/tests/test_dma_buf.py new file mode 100644 index 00000000..8b9f49e6 --- /dev/null +++ b/python/tests/test_dma_buf.py @@ -0,0 +1,44 @@ +"""Real FD handoff across independent processes; no CUDA/native dependency.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from pegaflow.dma_buf import DmaBufExports, receive_dma_buf + +pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="Linux DMA-BUF FD transport") + + +def test_fd_handoff_and_registration_cleanup(tmp_path): + payload = tmp_path / "allocation" + payload.write_bytes(b"owned allocation") + with DmaBufExports() as exports: + fd = os.open(payload, os.O_RDONLY) + address, token = exports.add(fd) + result = subprocess.run( + [ + sys.executable, + "-c", + "import os,sys; from pegaflow.dma_buf import receive_dma_buf; " + "fd=receive_dma_buf(sys.argv[1], bytes.fromhex(sys.argv[2])); " + "assert not os.get_inheritable(fd); " + "print(os.read(fd, 64).decode()); os.close(fd)", + address, + token.hex(), + ], + capture_output=True, + text=True, + timeout=15, + check=True, + ) + assert result.stdout.strip() == "owned allocation" + with pytest.raises(RuntimeError, match="exactly one descriptor"): + receive_dma_buf(address, token) + unclaimed = os.open(payload, os.O_RDONLY) + exports.add(unclaimed) + assert not Path(address).exists() + with pytest.raises(OSError): + os.fstat(unclaimed) diff --git a/python/tests/test_tp_shards.py b/python/tests/test_tp_shards.py index 063d73a0..8d5da086 100644 --- a/python/tests/test_tp_shards.py +++ b/python/tests/test_tp_shards.py @@ -53,12 +53,13 @@ def _context(**kwargs) -> ConnectorContext: return ConnectorContext(**defaults) # type: ignore[arg-type] -def _vllm_config(**parallel_overrides): +def _vllm_config(extra_config=None, **parallel_overrides): extra_config = { "pegaflow.tp_shard_endpoints": [ "http://node-a:50055", "http://node-b:50055", - ] + ], + **(extra_config or {}), } kv_transfer_config = SimpleNamespace( engine_id="instance", @@ -102,6 +103,37 @@ def test_topology_maps_contiguous_global_tp_ranks_to_local_servers(): assert topology.namespace("base", 1) == "base:tp-shard-1-of-2" +@pytest.mark.parametrize("enabled", [None, False, True], ids=["default", "staging", "direct"]) +def test_direct_gpu_config_controls_queries_and_registration(monkeypatch, enabled): + monkeypatch.setattr("pegaflow.connector.get_tensor_model_parallel_rank", lambda: 0) + client = MagicMock() + client.query_prefetch.return_value = QueryReady(1, b"lease") + monkeypatch.setattr("pegaflow.connector.EngineRpcClient", MagicMock(return_value=client)) + monkeypatch.setattr("pegaflow.connector.ServiceStateManager", MagicMock()) + exports = MagicMock() + monkeypatch.setattr("pegaflow.connector.worker.DmaBufExports", exports) + register = MagicMock() + monkeypatch.setattr(WorkerConnector, "_register_kv_caches", register) + extra = {} if enabled is None else {"pegaflow.direct_gpu_rdma": enabled} + scheduler = PegaKVConnector(_vllm_config(extra), KVConnectorRole.SCHEDULER) + worker = PegaKVConnector(_vllm_config(extra), KVConnectorRole.WORKER) + try: + scheduler._scheduler._count_available_block_prefix([b"hash"], "request") + assert len(client.query_prefetch.call_args_list) == 2 + for query in client.query_prefetch.call_args_list: + assert query.kwargs.get("direct_gpu", False) is bool(enabled) + worker.register_kv_caches({}) + if enabled: + exports.assert_called_once_with() + register.assert_called_once_with({}, exports.return_value.__enter__.return_value) + else: + exports.assert_not_called() + register.assert_called_once_with({}, None) + finally: + scheduler.shutdown() + worker.shutdown() + + @pytest.mark.parametrize( ("endpoints", "tp_size", "world_size", "message"), [