ggml-vulkan: chunk host-memory imports + bound-check pinned reads

The mmap region kept in system RAM for -ncmoe can be tens of GB, larger than
device->max_buffer_size (~4 GiB), so importing it as a single Vulkan buffer
failed and pinning silently no-op'd. Import in page-aligned chunks up to
max_buffer_size, registering each in device->pinned_memory. Add a bound check
in the pinned copy path so a tensor that straddles a chunk boundary falls
back to staging instead of reading out of bounds.

Assisted-by: opencode

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -3173,7 +3173,7 @@ static vk_buffer ggml_vk_create_buffer(vk_device& device, size_t size, const std
             import_info.setPNext(&mem_flags_info);
             buf->device_memory = device->device.allocateMemory({ size, memory_type_idx, &import_info });
         } catch (const vk::SystemError& e) {
-            fprintf(stderr, "PROBE import allocateMemory threw: %s\n", e.what());
+            GGML_LOG_WARN("ggml_vulkan: host pointer memory import failed (%s)\n", e.what());
         }
     } else {
         for (auto it = req_flags_list.begin(); it != req_flags_list.end(); it++) {
@@ -7857,25 +7857,32 @@ static bool ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, siz
     ggml_vk_host_get(dst->device, src, buf, buf_offset);
 
     if (buf != nullptr) {
-        // Memory is pinned, use as staging buffer
+        // extent of the read in pinned source memory; guard against tensors that
-        std::vector<vk::BufferCopy> slices(1);
+        // straddle a pinned-chunk boundary (they fall back to staging below)
-        if (width == spitch && width == dpitch) {
+        size_t src_extent = (width == spitch) ? (size_t) width * height
-            // Only do single write if stride is equal
+                                              : (height > 0 ? (height - 1) * spitch + width : 0);
-            slices[0].srcOffset = buf_offset;
+        if (buf_offset + src_extent <= buf->size) {
-            slices[0].dstOffset = offset;
+            // Memory is pinned, use as staging buffer
-            slices[0].size = width * height;
+            std::vector<vk::BufferCopy> slices(1);
-        } else {
+            if (width == spitch && width == dpitch) {
-            slices.resize(height);
+                // Only do single write if stride is equal
-            for (size_t i = 0; i < height; i++) {
+                slices[0].srcOffset = buf_offset;
-                slices[i].srcOffset = buf_offset + i * spitch;
+                slices[0].dstOffset = offset;
-                slices[i].dstOffset = offset + i * dpitch;
+                slices[0].size = width * height;
-                slices[i].size = width;
+            } else {
+                slices.resize(height);
+                for (size_t i = 0; i < height; i++) {
+                    slices[i].srcOffset = buf_offset + i * spitch;
+                    slices[i].dstOffset = offset + i * dpitch;
+                    slices[i].size = width;
+                }
             }
-        }
 
-        ggml_vk_sync_buffers(nullptr, subctx);
+            ggml_vk_sync_buffers(nullptr, subctx);
-        subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
+            subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
-        return true;
+            return true;
+        }
+        // straddles a chunk boundary: fall through to staging
     }
     VK_LOG_DEBUG("STAGING");
 
@@ -17813,17 +17820,14 @@ static void ggml_backend_vk_device_event_synchronize(ggml_backend_dev_t dev, ggm
 
 static vk_buffer ggml_vk_buffer_from_host_ptr(vk_device & device, void * ptr, size_t size) {
     if (!device->external_memory_host) {
-        fprintf(stderr, "PROBE from_host_ptr: ext_mem_host disabled\n");
         return {};
     }
 
     uintptr_t uptr = reinterpret_cast<uintptr_t>(ptr);
     if (uptr & (device->min_imported_host_pointer_alignment - 1)) {
-        fprintf(stderr, "PROBE from_host_ptr: ptr %p not aligned to %zu\n", ptr, device->min_imported_host_pointer_alignment);
         return {};
     }
     if (size & (device->min_imported_host_pointer_alignment - 1)) {
-        fprintf(stderr, "PROBE from_host_ptr: size %zu not aligned to %zu\n", size, device->min_imported_host_pointer_alignment);
         return {};
     }
 
@@ -17833,10 +17837,9 @@ static vk_buffer ggml_vk_buffer_from_host_ptr(vk_device & device, void * ptr, si
     try {
         buf = ggml_vk_create_buffer(device, size, { property_flags }, ptr);
     } catch (vk::SystemError& e) {
-        fprintf(stderr, "PROBE from_host_ptr: create_buffer threw (%s)\n", e.what());
+        GGML_LOG_WARN("ggml_vulkan: Failed ggml_vk_create_buffer (%s)\n", e.what());
     }
 
-    fprintf(stderr, "PROBE from_host_ptr: ptr=%p size=%zu -> buf=%p buf->buffer=%p\n", ptr, size, buf.get(), buf ? (void*)buf->buffer.operator VkBuffer() : nullptr);
     return buf;
 }
 
@@ -17927,9 +17930,13 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
 // instead of bouncing through the staging buffer + host memcpy. Mirrors the
 // GGML_CUDA_REGISTER_HOST path; populates device->pinned_memory, which the
 // existing pinned fast path in ggml_vk_buffer_write_2d_async looks up.
+//
+// A single Vulkan buffer cannot cover a whole multi-GB mmap (it is capped at
+// device->max_buffer_size), so the region is imported in page-aligned chunks.
+// ggml_vk_host_get resolves a tensor pointer to the chunk that contains it,
+// and ggml_vk_buffer_write_2d_async falls back to staging for any tensor that
+// straddles a chunk boundary, so correctness is preserved.
 static bool ggml_backend_vk_register_host_buffer(void * buffer, size_t size) {
-    fprintf(stderr, "PROBE register_host CALLED buffer=%p size=%zu cuda_env=%d vk_env=%d\n", buffer, size,
-            getenv("GGML_CUDA_REGISTER_HOST") != nullptr, getenv("GGML_VK_REGISTER_HOST") != nullptr);
     if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr && getenv("GGML_VK_REGISTER_HOST") == nullptr) {
         return false;
     }
@@ -17940,20 +17947,36 @@ static bool ggml_backend_vk_register_host_buffer(void * buffer, size_t size) {
     bool success = false;
     for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
         vk_device& device = vk_instance.devices[i];
-        if (!device || !device->external_memory_host) {
+        if (!device || !device->external_memory_host || device->max_buffer_size == 0) {
             continue;
         }
 
-        vk_buffer buf = ggml_vk_buffer_from_host_ptr(device, buffer, size);
+        const size_t align = device->min_imported_host_pointer_alignment;
-        if (!buf || !buf->buffer) {
+        size_t chunk = device->max_buffer_size & ~(align - 1);
+        if (chunk == 0) {
             continue;
         }
 
-        {
+        uint8_t * p = static_cast<uint8_t *>(buffer);
-            std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
+        size_t remaining = size;
-            device->pinned_memory.emplace_back(buffer, size, buf);
+        bool dev_success = false;
+        while (remaining > 0) {
+            size_t cur = std::min(remaining, chunk);
+            vk_buffer buf = ggml_vk_buffer_from_host_ptr(device, p, cur);
+            if (!buf || !buf->buffer) {
+                break;
+            }
+            {
+                std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
+                device->pinned_memory.emplace_back(p, cur, buf);
+            }
+            dev_success = true;
+            p += cur;
+            remaining -= cur;
+        }
+        if (dev_success) {
+            success = true;
         }
-        success = true;
     }
     return success;
 }