fix: pass in device_id to KVBM PinnedAllocator instead of hard code to 0 (#6809)

Signed-off-by: Ziqi Fan <ziqif@nvidia.com>

lib/llm/src/block_manager/distributed/worker.rs

@@ -148,7 +148,8 @@ async fn perform_allocation_and_build_handler(
     )?);
     // host
     let host_blocks = if leader_meta.num_host_blocks > 0 {
-        let host_allocator = Arc::new(PinnedAllocator::default());
+        let host_allocator = Arc::new(PinnedAllocator::new(device_id)?);
+
         let host_layout = layout_builder
             .num_blocks(leader_meta.num_host_blocks)
             .build()?

lib/llm/src/block_manager/storage/cuda.rs

@@ -332,10 +332,13 @@ impl Default for PinnedAllocator {
 }
 
 impl PinnedAllocator {
-    /// Create a new pinned allocator
-    pub fn new() -> Result<Self, StorageError> {
+    /// Create a new pinned allocator for the specified device.
+    ///
+    /// The device_id determines which NUMA node pinned memory will be allocated
+    /// on when NUMA-aware allocation is enabled.
+    pub fn new(device_id: usize) -> Result<Self, StorageError> {
         Ok(Self {
-            ctx: Cuda::device_or_create(0)?,
+            ctx: Cuda::device_or_create(device_id)?,
         })
     }
 }