dynamo_memory/
external.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! External memory wrapper for memory allocated by external frameworks.
5//!
6//! This module provides `ExternalDeviceMemory` for wrapping pointers to GPU
7//! memory allocated by external frameworks (e.g., vLLM's KV cache). This type
8//! does NOT own the memory - ownership remains with the external framework.
9//!
10//! The primary use case is registering external GPU memory with NIXL for RDMA
11//! transfers without copying.
12
13use crate::nixl::{MemType, NixlCompatible, NixlDescriptor};
14use crate::{MemoryDescriptor, StorageKind};
15use std::any::Any;
16use std::fmt;
17
18/// Wrapper for externally-allocated device (GPU) memory.
19///
20/// This type wraps a raw pointer to GPU memory that is owned by an external
21/// framework (like vLLM). It provides the necessary traits for NIXL registration
22/// without taking ownership of the underlying memory.
23///
24/// # Safety
25///
26/// This type relies on the caller to guarantee that:
27/// - The pointer points to valid GPU memory on the specified device
28/// - The memory remains valid for the lifetime of this wrapper
29/// - The memory size is exactly as specified
30/// - The external framework doesn't free the memory while this wrapper exists
31///
32/// # Example
33///
34/// ```ignore
35/// // vLLM allocates KV cache tensors
36/// let tensor_ptr = tensor.data_ptr();
37/// let tensor_size = tensor.size_bytes();
38/// let device_id = tensor.device.index;
39///
40/// // Wrap without taking ownership
41/// let external = unsafe {
42///     ExternalDeviceMemory::new(tensor_ptr as *const u8, tensor_size, device_id as u64)
43/// };
44///
45/// // Register with NIXL for RDMA
46/// let registered = register_with_nixl(external, &agent, None)?;
47/// ```
48pub struct ExternalDeviceMemory {
49    /// Raw pointer to externally-allocated GPU memory.
50    ptr: *const u8,
51    /// Size of the memory region in bytes.
52    size: usize,
53    /// CUDA device ID where this memory resides.
54    device_id: u64,
55}
56
57// Safety: The external framework (e.g., vLLM) guarantees the memory remains valid
58// for the lifetime of the KV cache. The pointer is only used for NIXL registration
59// and transfer operations which are synchronized by the framework.
60unsafe impl Send for ExternalDeviceMemory {}
61unsafe impl Sync for ExternalDeviceMemory {}
62
63impl ExternalDeviceMemory {
64    /// Create a wrapper for external device memory.
65    ///
66    /// # Safety
67    ///
68    /// Caller must ensure:
69    /// - `ptr` points to valid GPU memory on CUDA device `device_id`
70    /// - The memory remains valid for the lifetime of this wrapper
71    /// - The memory size is exactly `size` bytes
72    /// - The external framework doesn't free the memory while this wrapper exists
73    #[inline]
74    pub unsafe fn new(ptr: *const u8, size: usize, device_id: u64) -> Self {
75        Self {
76            ptr,
77            size,
78            device_id,
79        }
80    }
81
82    /// Get the raw pointer to the external memory.
83    #[inline]
84    pub fn as_ptr(&self) -> *const u8 {
85        self.ptr
86    }
87
88    /// Get the CUDA device ID where this memory resides.
89    #[inline]
90    pub fn device_id(&self) -> u64 {
91        self.device_id
92    }
93}
94
95impl fmt::Debug for ExternalDeviceMemory {
96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97        f.debug_struct("ExternalDeviceMemory")
98            .field("ptr", &format_args!("{:p}", self.ptr))
99            .field("size", &self.size)
100            .field("device_id", &self.device_id)
101            .finish()
102    }
103}
104
105impl MemoryDescriptor for ExternalDeviceMemory {
106    #[inline]
107    fn addr(&self) -> usize {
108        self.ptr as usize
109    }
110
111    #[inline]
112    fn size(&self) -> usize {
113        self.size
114    }
115
116    #[inline]
117    fn storage_kind(&self) -> StorageKind {
118        StorageKind::Device(self.device_id as u32)
119    }
120
121    fn as_any(&self) -> &dyn Any {
122        self
123    }
124
125    fn nixl_descriptor(&self) -> Option<NixlDescriptor> {
126        // External memory doesn't have a pre-existing NIXL descriptor
127        // It will be registered and get one via NixlRegistered wrapper
128        None
129    }
130}
131
132impl NixlCompatible for ExternalDeviceMemory {
133    fn nixl_params(&self) -> (*const u8, usize, MemType, u64) {
134        (self.ptr, self.size, MemType::Vram, self.device_id)
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn test_external_device_memory_traits() {
144        // Create with a dummy pointer (not actually valid GPU memory)
145        let ptr = 0x1000 as *const u8;
146        let size = 1024;
147        let device_id = 0;
148
149        let external = unsafe { ExternalDeviceMemory::new(ptr, size, device_id) };
150
151        // Check MemoryDescriptor
152        assert_eq!(external.addr(), 0x1000);
153        assert_eq!(external.size(), 1024);
154        assert_eq!(external.storage_kind(), StorageKind::Device(0));
155        assert!(external.nixl_descriptor().is_none());
156
157        // Check NixlCompatible
158        let (p, s, mem_type, dev) = external.nixl_params();
159        assert_eq!(p as usize, 0x1000);
160        assert_eq!(s, 1024);
161        assert_eq!(mem_type, MemType::Vram);
162        assert_eq!(dev, 0);
163    }
164}
dynamo_memory/external.rs

dynamo_memory/
external.rs
