dynamo_memory/
pinned.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! CUDA pinned host memory storage.
5
6use super::{MemoryDescriptor, Result, StorageError, StorageKind, actions, nixl::NixlDescriptor};
7use cudarc::driver::CudaContext;
8use std::any::Any;
9use std::sync::Arc;
10
11/// Whether to use write-combined pinned allocations.
12///
13/// Probed once at first use: returns `false` if `DYN_KVBM_DISABLE_WRITE_COMBINED`
14/// is set, or if a test allocation reveals the hardware does not support it
15/// (e.g. Grace Hopper / Blackwell with NVLink-C2C). Must be accessed only after
16/// a CUDA context has been bound to the current thread.
17static USE_WRITE_COMBINED: std::sync::LazyLock<bool> = std::sync::LazyLock::new(|| {
18    if crate::env_is_truthy("DYN_KVBM_DISABLE_WRITE_COMBINED") {
19        tracing::debug!("DYN_KVBM_DISABLE_WRITE_COMBINED set; write-combined disabled");
20        return false;
21    }
22    // Probe hardware support with a 1-byte test allocation.
23    // SAFETY: called from an allocation path that has already bound a CUDA context.
24    unsafe {
25        match cudarc::driver::result::malloc_host(
26            1,
27            cudarc::driver::sys::CU_MEMHOSTALLOC_WRITECOMBINED,
28        ) {
29            Ok(ptr) => {
30                let _ = cudarc::driver::result::free_host(ptr);
31                true
32            }
33            Err(_) => {
34                tracing::debug!(
35                    "Write-combined memory not supported on this system; \
36                     will use regular pinned memory"
37                );
38                false
39            }
40        }
41    }
42});
43
44/// Allocates pinned host memory, using write-combined if [`USE_WRITE_COMBINED`]
45/// allows it, otherwise falling back to `CU_MEMHOSTALLOC_DEVICEMAP`.
46///
47/// # Safety
48/// Caller must ensure a valid CUDA context is bound to the current thread.
49unsafe fn malloc_host_prefer_writecombined(size: usize) -> Result<*mut u8> {
50    if *USE_WRITE_COMBINED {
51        // SAFETY: caller guarantees a valid CUDA context is bound to the current thread
52        unsafe {
53            cudarc::driver::result::malloc_host(
54                size,
55                cudarc::driver::sys::CU_MEMHOSTALLOC_WRITECOMBINED,
56            )
57        }
58        .map(|ptr| ptr as *mut u8)
59        .map_err(StorageError::Cuda)
60    } else {
61        // SAFETY: caller guarantees a valid CUDA context is bound to the current thread
62        unsafe {
63            cudarc::driver::result::malloc_host(
64                size,
65                cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP,
66            )
67        }
68        .map(|ptr| ptr as *mut u8)
69        .map_err(StorageError::Cuda)
70    }
71}
72
73/// CUDA pinned host memory allocated via cudaHostAlloc.
74#[derive(Debug)]
75pub struct PinnedStorage {
76    /// Host pointer to the pinned memory.
77    ptr: usize,
78    /// Size of the allocation in bytes.
79    len: usize,
80    /// CUDA context used for allocation and deallocation.
81    ctx: Arc<CudaContext>,
82}
83
84unsafe impl Send for PinnedStorage {}
85unsafe impl Sync for PinnedStorage {}
86
87impl PinnedStorage {
88    /// Allocate new pinned memory of the given size.
89    ///
90    /// This is a convenience method that calls `new_for_device(len, None)`.
91    ///
92    /// # Arguments
93    /// * `len` - Size in bytes to allocate
94    pub fn new(len: usize) -> Result<Self> {
95        Self::new_for_device(len, None)
96    }
97
98    /// Allocate pinned memory, optionally NUMA-aware for a specific GPU.
99    ///
100    /// When `device_id` is `Some`, NUMA-aware allocation is attempted by default:
101    /// a worker thread pinned to the GPU's NUMA node performs the allocation,
102    /// ensuring optimal memory placement via first-touch policy. If the GPU's
103    /// NUMA node cannot be determined, allocation falls back to the direct path.
104    /// Set `DYN_MEMORY_DISABLE_NUMA=1` to skip NUMA optimization entirely.
105    ///
106    /// When `device_id` is `None`, a direct allocation is performed on device 0.
107    ///
108    /// # Arguments
109    /// * `len` - Size in bytes to allocate
110    /// * `device_id` - If Some, use NUMA-aware allocation on the GPU's NUMA node
111    ///
112    /// # Errors
113    /// Returns an error if:
114    /// - `len` is 0
115    /// - CUDA context creation fails
116    /// - Memory allocation fails
117    pub fn new_for_device(len: usize, device_id: Option<u32>) -> Result<Self> {
118        if len == 0 {
119            return Err(StorageError::AllocationFailed(
120                "zero-sized allocations are not supported".into(),
121            ));
122        }
123
124        let gpu_id = device_id.unwrap_or(0);
125        let ctx = crate::device::cuda_context(gpu_id)?;
126
127        // Try NUMA-aware allocation unless explicitly disabled
128        #[cfg(target_os = "linux")]
129        let numa_ptr = if let Some(gpu_id) = device_id {
130            if super::numa::is_numa_enabled() {
131                match super::numa::worker_pool::NumaWorkerPool::global()
132                    .allocate_pinned_for_gpu(len, gpu_id)
133                {
134                    Ok(Some(ptr)) => {
135                        tracing::debug!(
136                            "Using NUMA-aware allocation for {} bytes on GPU {}",
137                            len,
138                            gpu_id
139                        );
140                        Some(ptr as usize)
141                    }
142                    Ok(None) => None, // NUMA node unknown, fall through
143                    Err(e) => return Err(StorageError::AllocationFailed(e)),
144                }
145            } else {
146                None
147            }
148        } else {
149            None
150        };
151
152        #[cfg(not(target_os = "linux"))]
153        let numa_ptr: Option<usize> = None;
154
155        let ptr = if let Some(ptr) = numa_ptr {
156            ptr
157        } else {
158            unsafe {
159                ctx.bind_to_thread().map_err(StorageError::Cuda)?;
160
161                let ptr = malloc_host_prefer_writecombined(len)?;
162
163                assert!(!ptr.is_null(), "Failed to allocate pinned memory");
164                assert!(ptr.is_aligned(), "Pinned memory is not aligned");
165                assert!(len < isize::MAX as usize);
166
167                ptr as usize
168            }
169        };
170
171        Ok(Self { ptr, len, ctx })
172    }
173
174    /// Get a pointer to the underlying memory.
175    ///
176    /// # Safety
177    /// The caller must ensure the pointer is not used after this storage is dropped.
178    pub unsafe fn as_ptr(&self) -> *const u8 {
179        self.ptr as *const u8
180    }
181
182    /// Get a mutable pointer to the underlying memory.
183    ///
184    /// # Safety
185    /// The caller must ensure the pointer is not used after this storage is dropped
186    /// and that there are no other references to this memory.
187    pub unsafe fn as_mut_ptr(&mut self) -> *mut u8 {
188        self.ptr as *mut u8
189    }
190
191    /// Get a reference to the CUDA context used for this allocation.
192    pub fn ctx(&self) -> &Arc<CudaContext> {
193        &self.ctx
194    }
195}
196
197impl Drop for PinnedStorage {
198    fn drop(&mut self) {
199        if let Err(e) = self.ctx.bind_to_thread() {
200            tracing::debug!("failed to bind CUDA context for free: {e}");
201        }
202        unsafe {
203            if let Err(e) = cudarc::driver::result::free_host(self.ptr as _) {
204                tracing::debug!("failed to free pinned memory: {e}");
205            }
206        };
207    }
208}
209
210impl MemoryDescriptor for PinnedStorage {
211    fn addr(&self) -> usize {
212        unsafe { self.as_ptr() as usize }
213    }
214
215    fn size(&self) -> usize {
216        self.len
217    }
218
219    fn storage_kind(&self) -> StorageKind {
220        StorageKind::Pinned
221    }
222
223    fn as_any(&self) -> &dyn Any {
224        self
225    }
226
227    fn nixl_descriptor(&self) -> Option<NixlDescriptor> {
228        None
229    }
230}
231
232// Support for NIXL registration
233impl super::nixl::NixlCompatible for PinnedStorage {
234    fn nixl_params(&self) -> (*const u8, usize, nixl_sys::MemType, u64) {
235        let ptr = unsafe { self.as_ptr() };
236        (ptr, self.len, nixl_sys::MemType::Dram, 0)
237    }
238}
239
240impl actions::Memset for PinnedStorage {
241    fn memset(&mut self, value: u8, offset: usize, size: usize) -> Result<()> {
242        let end = offset
243            .checked_add(size)
244            .ok_or_else(|| StorageError::OperationFailed("memset: offset overflow".into()))?;
245        if end > self.len {
246            return Err(StorageError::OperationFailed(
247                "memset: offset + size > storage size".into(),
248            ));
249        }
250        unsafe {
251            let ptr = (self.ptr as *mut u8).add(offset);
252            std::ptr::write_bytes(ptr, value, size);
253        }
254        Ok(())
255    }
256}
dynamo_memory/pinned.rs

dynamo_memory/
pinned.rs
