RPC: improve loading

ggml/src/ggml-backend-meta.cpp

@@ -865,7 +865,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
         ggml_backend_meta_split_state split_state;
         switch (tensor->op) {
             case GGML_OP_NONE: {
-                split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+                if (tensor->view_src != nullptr) {
+                    // full-tensor view created with ggml_view_tensor, transparent for the split state
+                    split_state = ggml_backend_meta_get_split_state(stc, tensor->view_src, assume_sync);
+                } else {
+                    split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+                }
             } break;
             case GGML_OP_DUP: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
@@ -2289,6 +2294,14 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend,
                 cgraph_ij->uid = ggml_graph_next_uid();
             }
         }
+
+        // Aux graph contents are rewritten on every compute but are identical across calls while the subgraphs are reused,
+        // so they can get stable uids on rebuild. Only safe without a comm backend, where the fallback usage is deterministic.
+        if (backend_ctx->comm_ctx == nullptr) {
+            for (ggml_cgraph * cgraph_aux : backend_ctx->cgraphs_aux) {
+                cgraph_aux->uid = ggml_graph_next_uid();
+            }
+        }
     }
 
     size_t iga = 0; // i graph aux

ggml/src/ggml-rpc/transport.cpp

@@ -69,6 +69,11 @@ struct rdma_conn {
     struct ibv_mr * rx_mr  = nullptr;
     int             rx_head = 0;
 
+    // received slot that is only partly consumed, so recv sizes need not match send sizes
+    int             rx_cur  = -1;
+    size_t          rx_off  = 0;
+    size_t          rx_len  = 0;
+
     uint32_t        max_inline = 0;
 
     uint8_t * rx_slot(int i) const {
@@ -458,14 +463,24 @@ bool socket_t::impl::rdma_recv(void * data, size_t size) {
     uint8_t * dst = (uint8_t *)data;
     size_t rem = size;
     while (rem > 0) {
-        struct ibv_wc wc;
-        if (!rdma_poll(c->rcq, &wc)) return false;
+        if (c->rx_cur < 0) {
+            struct ibv_wc wc;
+            if (!rdma_poll(c->rcq, &wc)) return false;
 
-        int slot = (int)wc.wr_id;
-        size_t got = wc.byte_len;
-        memcpy(dst, c->rx_slot(slot), got);
+            c->rx_cur = (int)wc.wr_id;
+            c->rx_off = 0;
+            c->rx_len = wc.byte_len;
+        }
 
-        if (!c->post_rx(slot)) return false;
+        size_t got = std::min(rem, c->rx_len - c->rx_off);
+        memcpy(dst, c->rx_slot(c->rx_cur) + c->rx_off, got);
+        c->rx_off += got;
+
+        // repost the slot only once it is fully consumed
+        if (c->rx_off == c->rx_len) {
+            if (!c->post_rx(c->rx_cur)) return false;
+            c->rx_cur = -1;
+        }
 
         dst += got;
         rem -= got;

src/llama-model-loader.cpp

@@ -1622,6 +1622,8 @@ bool llama_model_loader::load_all_data(
         });
     }
 
+    // staging buffer for non-host buffers without async uploads, reused because a fresh allocation per tensor page faults on every load
+    std::vector<no_init<uint8_t>> read_buf;
     for (struct ggml_tensor * cur : tensors) {
         const auto * weight = get_weight(ggml_get_name(cur));
         if (weight == nullptr) {
@@ -1735,12 +1737,25 @@ bool llama_model_loader::load_all_data(
                         buffer_idx %= n_buffers;
                     }
                 } else {
-                    // scoped to one tensor so only one staging buffer is alive at a time
-                    std::vector<no_init<uint8_t>> read_buf(n_size);
-                    file->seek(weight->offs, SEEK_SET);
-                    file->read_raw(read_buf.data(), n_size);
-                    ggml_backend_tensor_set(cur, read_buf.data(), 0, n_size);
-                    if (check_tensors && !ggml_validate_row_data(cur->type, read_buf.data(), n_size)) {
+                    // with direct IO read the aligned range around the tensor and use the data inside it, so no bounce buffer is needed
+                    const size_t align          = file->read_alignment();
+                    const size_t aligned_offset = weight->offs & ~(align - 1);
+                    const size_t pad            = weight->offs - aligned_offset;
+                    const size_t read_size      = (pad + n_size + align - 1) & ~(align - 1);
+                    const size_t buf_size       = read_size + align - 1;
+
+                    // tensors come biggest-first: shrink the buffer when it gets too large, so peak memory stays close to one tensor
+                    if (read_buf.size() < buf_size || read_buf.size() > 2*buf_size) {
+                        read_buf = {};
+                        read_buf.resize(buf_size);
+                    }
+                    uint8_t * read_dst = (uint8_t *) GGML_PAD((uintptr_t) read_buf.data(), align);
+                    file->seek(aligned_offset, SEEK_SET);
+                    file->read_raw_unsafe(read_dst, read_size);
+
+                    const uint8_t * data = read_dst + pad;
+                    ggml_backend_tensor_set(cur, data, 0, n_size);
+                    if (check_tensors && !ggml_validate_row_data(cur->type, data, n_size)) {
                         throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));
                     }
                 }