ggml-vulkan: pre-stage host weights when import is unavailable

VK_EXT_external_memory_host import fails on RADV for file-backed mmap pages,
so register_host_buffer no-op'd and every MoE-expert upload paid a slow
single-threaded pageable memcpy into the staging buffer each eval (~3 GB/s
effective on the 35B, ~25% of PCIe bandwidth).

When the import fails, fall back to a one-time copy of the region into a
host-visible Vulkan buffer registered in device->pinned_memory under the
mmap's address range. The existing pinned fast path then DMAs straight from
that buffer each eval at full PCIe bandwidth - no per-eval CPU memcpy and no
blocking sync. Costs one-time host-visible memory equal to the CPU-resident
weight region.

Assisted-by: opencode

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

@@ -3173,6 +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) {
+            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++) {
@@ -7856,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
-        std::vector<vk::BufferCopy> slices(1);
-        if (width == spitch && width == dpitch) {
-            // Only do single write if stride is equal
-            slices[0].srcOffset = buf_offset;
-            slices[0].dstOffset = offset;
-            slices[0].size = width * height;
-        } 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;
+        // extent of the read in pinned source memory; guard against tensors that
+        // straddle a pinned-chunk boundary (they fall back to staging below)
+        size_t src_extent = (width == spitch) ? (size_t) width * height
+                                              : (height > 0 ? (height - 1) * spitch + width : 0);
+        if (buf_offset + src_extent <= buf->size) {
+            // Memory is pinned, use as staging buffer
+            std::vector<vk::BufferCopy> slices(1);
+            if (width == spitch && width == dpitch) {
+                // Only do single write if stride is equal
+                slices[0].srcOffset = buf_offset;
+                slices[0].dstOffset = offset;
+                slices[0].size = width * height;
+            } 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);
-        subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
-        return true;
+            ggml_vk_sync_buffers(nullptr, subctx);
+            subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
+            return true;
+        }
+        // straddles a chunk boundary: fall through to staging
     }
     VK_LOG_DEBUG("STAGING");
 
@@ -12206,9 +12214,108 @@ static void ggml_vk_opt_step_sgd(ggml_backend_vk_context * ctx, vk_context& subc
     ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, src1, src2, nullptr, dst, GGML_OP_OPT_STEP_SGD, { (uint32_t)n, 0, 0.0f, 0.0f, 0.0f, 0.0f });
 }
 
+// Fast path for concat along dim 0 where one source is stored "transposed"
+// (nb[1] == type_size, dim1 innermost) and the other source + dst are
+// contiguous along dim 0. The generic concat shader reads the transposed
+// source with a catastrophic uncoalesced stride; here we instead copy the
+// contiguous source with copy.comp and transpose the other source into the
+// matching dst sub-region with the tiled copy_transpose shader (shared-memory
+// transpose, coalesced read+write). Mirrors the precedent in
+// ggml_vk_cpy_to_contiguous: direct dispatch with custom push constants.
+static void ggml_vk_concat_transpose_fastpath(ggml_backend_vk_context * ctx, vk_context& subctx,
+                                              const ggml_tensor * ctg, const ggml_tensor * trp,
+                                              ggml_tensor * dst, uint32_t off_ctg, uint32_t off_trp) {
+    const uint32_t ts = ggml_type_size(dst->type);
+
+    vk_pipeline pipeline_cpy = (ts == 4) ? ctx->device->pipeline_cpy_f32_f32
+                                          : ctx->device->pipeline_cpy_f16_f16;
+    vk_pipeline pipeline_trp = (ts == 4) ? ctx->device->pipeline_cpy_transpose_32
+                                          : ctx->device->pipeline_cpy_transpose_16;
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline_cpy, 1);
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline_trp, 1);
+
+    vk_subbuffer ctg_buf = ggml_vk_tensor_subbuffer(ctx, ctg, true);
+    vk_subbuffer trp_buf = ggml_vk_tensor_subbuffer(ctx, trp, true);
+    vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
+
+    const uint32_t a_misalign_ctg = get_misalign_bytes(ctx, ctg) / ts;
+    const uint32_t a_misalign_trp = get_misalign_bytes(ctx, trp) / ts;
+    const uint32_t d_misalign     = get_misalign_bytes(ctx, dst) / ts;
+
+    // Dispatch A: contiguous copy of `ctg` into dst[off_ctg : off_ctg + ctg->ne[0], :]
+    if (ctg->ne[0] > 0) {
+        const uint32_t ne_ctg = (uint32_t) ggml_nelements(ctg);
+        vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(ctg, dst, ne_ctg);
+        pc.ne10 = (uint32_t) ctg->ne[0];                 // only ctg's columns in dst
+        pc.misalign_offsets = (a_misalign_ctg << 16) | (d_misalign + off_ctg);
+        init_pushconst_fastdiv(pc);
+        std::array<uint32_t, 3> el = ne_ctg > 262144 ? std::array<uint32_t,3>{512, 512, CEIL_DIV(ne_ctg, 262144)}
+                                      : ne_ctg > 512  ? std::array<uint32_t,3>{512, CEIL_DIV(ne_ctg, 512), 1}
+                                                      : std::array<uint32_t,3>{ne_ctg, 1, 1};
+        ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_cpy, { ctg_buf, dst_buf }, pc, el);
+    }
+
+    // Dispatch B: tiled transpose of `trp` into dst[off_trp : off_trp + trp->ne[0], :]
+    if (trp->ne[0] > 0) {
+        vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(trp, dst, ggml_nelements(trp));
+        pc.ne10 = (uint32_t) trp->ne[0];                 // dst bound = trp columns (NOT dst->ne[0])
+        pc.misalign_offsets = (a_misalign_trp << 16) | (d_misalign + off_trp);
+        init_pushconst_fastdiv(pc);
+        std::array<uint32_t, 3> el = {
+            (uint32_t) CEIL_DIV(trp->ne[0], 32),
+            (uint32_t) CEIL_DIV(trp->ne[1], 32),
+            (uint32_t) (trp->ne[2] * trp->ne[3]),
+        };
+        el[0] = std::min(el[0], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
+        el[1] = std::min(el[1], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[1]);
+        el[2] = std::min(el[2], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
+        ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_trp, { trp_buf, dst_buf }, pc, el);
+    }
+
+    ggml_vk_sync_buffers(ctx, subctx);
+}
+
 static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     int * op_params = (int *)dst->op_params;
 
+    // Fast path: concat along dim 0 with one source "transposed" (nb[1]==type_size,
+    // dim1 innermost) and the other source + dst contiguous along dim 0. The generic
+    // shader reads the transposed source uncoalesced (~5.6ms on RX 580 vs ~130us for
+    // a coalesced copy); here we instead use a tiled shared-memory transpose.
+    auto src_transposed_2d = [&](const ggml_tensor * s) {
+        const uint32_t ts = ggml_type_size(s->type);
+        return s->nb[1] == ts                                              // dim1 innermost
+            && s->nb[0] == s->ne[1] * ts                                   // consistent 2D transpose
+            && s->ne[2] == 1 && s->ne[3] == 1;
+    };
+    const uint32_t dst_ts = ggml_type_size(dst->type);
+    const bool dim0_ok   = (op_params[0] == 0) && (dst->nb[0] == dst_ts);
+    const bool types_ok  = (dst_ts == 4 || dst_ts == 2)
+                        && (src0->type == src1->type && src0->type == dst->type)
+                        && (ggml_blck_size(dst->type) == 1);
+    const bool shapes_ok = (src0->ne[1] == src1->ne[1] && src1->ne[1] == dst->ne[1])
+                        && (src0->ne[2] == src1->ne[2] && src0->ne[2] == dst->ne[2])
+                        && (src0->ne[3] == src1->ne[3] && src0->ne[3] == dst->ne[3])
+                        && (src0->ne[0] + src1->ne[0] == dst->ne[0]);
+    // doffset is 16-bit for unary shaders; guard against overflow
+    const bool off_fits  = ((get_misalign_bytes(ctx, dst)/dst_ts + dst->ne[0])) < 0xFFFFu;
+
+    const bool s1_trans = dim0_ok && types_ok && shapes_ok && off_fits
+                        && src_transposed_2d(src1) && ggml_is_contiguous(src0);
+    const bool s0_trans = dim0_ok && types_ok && shapes_ok && off_fits
+                        && src_transposed_2d(src0) && ggml_is_contiguous(src1);
+
+    if (s1_trans || s0_trans) {
+        const ggml_tensor * ctg = s1_trans ? src0 : src1;     // contiguous source
+        const ggml_tensor * trp = s1_trans ? src1 : src0;     // transposed source
+        // dst offsets (in elements) where each source region begins
+        const uint32_t off_ctg = s1_trans ? 0u                       : (uint32_t) src1->ne[0];
+        const uint32_t off_trp = s1_trans ? (uint32_t) src0->ne[0]   : 0u;
+        ggml_vk_concat_transpose_fastpath(ctx, subctx, ctg, trp, dst, off_ctg, off_trp);
+        return;
+    }
+
     const uint32_t src0_type_size = ggml_type_size(src0->type);
     const uint32_t src1_type_size = ggml_type_size(src1->type);
     const uint32_t dst_type_size = ggml_type_size(dst->type);
@@ -17917,11 +18024,112 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
     return devices[device];
 }
 
+// Import an mmap-backed host region as a Vulkan pinned buffer via
+// VK_EXT_external_memory_host so H2D uploads DMA straight from system RAM
+// 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) {
+    if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr && getenv("GGML_VK_REGISTER_HOST") == nullptr) {
+        return false;
+    }
+    if (size == 0) {
+        return false;
+    }
+
+    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 || device->max_buffer_size == 0) {
+            continue;
+        }
+
+        const size_t align = device->min_imported_host_pointer_alignment;
+        size_t chunk = device->max_buffer_size & ~(align - 1);
+        if (chunk == 0) {
+            continue;
+        }
+
+        uint8_t * p = static_cast<uint8_t *>(buffer);
+        size_t remaining = size;
+        bool dev_success = false;
+        bool import_ok = true;   // flips to false once VK_EXT_external_memory_host import fails
+        while (remaining > 0) {
+            size_t cur = std::min(remaining, chunk);
+            vk_buffer buf;
+            if (import_ok) {
+                buf = ggml_vk_buffer_from_host_ptr(device, p, cur);
+            }
+            if (!buf || !buf->buffer) {
+                // Fallback for drivers that can't import file-backed mmap pages (e.g. RADV):
+                // make a one-time copy of the region into a host-visible Vulkan buffer. The GPU
+                // then DMAs straight from it every eval at full PCIe bandwidth, instead of paying
+                // a slow single-threaded pageable memcpy into the staging buffer each time.
+                import_ok = false;
+                buf = ggml_vk_create_buffer_check(device, cur,
+                    vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached,
+                    vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent);
+                if (buf && buf->buffer && buf->ptr) {
+                    memcpy(buf->ptr, p, cur);
+                } else {
+                    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;
+        }
+    }
+    return success;
+}
+
+static void ggml_backend_vk_unregister_host_buffer(void * buffer) {
+    for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
+        vk_device& device = vk_instance.devices[i];
+        if (!device) {
+            continue;
+        }
+        std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
+        for (auto it = device->pinned_memory.begin(); it != device->pinned_memory.end(); ++it) {
+            if (std::get<0>(*it) == buffer) {
+                vk_buffer buf = std::get<2>(*it);
+                device->pinned_memory.erase(it);
+                ggml_vk_destroy_buffer(buf);
+                break;
+            }
+        }
+    }
+}
+
+static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
+    GGML_UNUSED(reg);
+    if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
+        return (void *) ggml_backend_vk_register_host_buffer;
+    }
+    if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) {
+        return (void *) ggml_backend_vk_unregister_host_buffer;
+    }
+    return nullptr;
+}
+
 static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = {
     /* .get_name         = */ ggml_backend_vk_reg_get_name,
     /* .get_device_count = */ ggml_backend_vk_reg_get_device_count,
     /* .get_device       = */ ggml_backend_vk_reg_get_device,
-    /* .get_proc_address = */ NULL,
+    /* .get_proc_address = */ ggml_backend_vk_reg_get_proc_address,
 };
 
 ggml_backend_reg_t ggml_backend_vk_reg() {