loader: validate GPU split sidecars before applying them

CMakeLists.txt

@@ -678,6 +678,7 @@ endif()
 add_library(llama
             llama.cpp
             llama.h
+            llama-gpu-split.h
             )
 
 target_include_directories(llama PUBLIC .)

llama-gpu-split.h

@@ -0,0 +1,78 @@
+#pragma once
+
+#include "ggml.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+// Internal validation shared by the sidecar loader and the row-copy path.
+namespace llama_gpu_split {
+
+inline bool is_i32_vector(const ggml_tensor * tensor) {
+    return tensor && tensor->type == GGML_TYPE_I32 && tensor->ne[0] >= 0 &&
+        tensor->ne[1] == 1 && tensor->ne[2] == 1 && tensor->ne[3] == 1 && ggml_is_contiguous(tensor);
+}
+
+inline bool validate_layout(const ggml_tensor * index, const ggml_tensor * bucket,
+                            int64_t rows, std::string & error) {
+    if (rows <= 0 || !is_i32_vector(index) || !is_i32_vector(bucket) || index->ne[0] != rows || bucket->ne[0] > rows) {
+        error = "GPU split tensors must be I32 vectors matching the FFN row count";
+        return false;
+    }
+    return true;
+}
+
+inline bool validate_bucket(const ggml_tensor * bucket, int64_t rows, std::string & error) {
+    if (rows <= 0 || !is_i32_vector(bucket) || bucket->ne[0] > rows || (bucket->ne[0] > 0 && !bucket->data)) {
+        error = "invalid GPU bucket shape, type, or data";
+        return false;
+    }
+    const int32_t * data = static_cast<const int32_t *>(bucket->data);
+    for (int64_t i = 0; i < bucket->ne[0]; ++i) {
+        if (data[i] < 0 || int64_t(data[i]) >= rows) {
+            error = "gpu_bucket[" + std::to_string(i) + "]=" + std::to_string(data[i]) +
+                " is outside [0, " + std::to_string(rows) + ")";
+            return false;
+        }
+    }
+    return true;
+}
+
+inline bool validate(const ggml_tensor * index, const ggml_tensor * bucket,
+                     int64_t rows, int64_t & selected, std::string & error) {
+    if (!validate_layout(index, bucket, rows, error) || !validate_bucket(bucket, rows, error)) {
+        return false;
+    }
+    if (!index->data) {
+        error = "missing GPU index data";
+        return false;
+    }
+    const int32_t * mask = static_cast<const int32_t *>(index->data);
+    int64_t count = 0;
+    for (int64_t row = 0; row < rows; ++row) {
+        if (mask[row] != 0 && mask[row] != 1) {
+            error = "GPU index mask must contain only zero or one";
+            return false;
+        }
+        count += mask[row];
+    }
+    if (count != bucket->ne[0]) {
+        error = "GPU bucket length does not match the number of selected neurons";
+        return false;
+    }
+    std::vector<bool> seen(static_cast<size_t>(rows), false);
+    const int32_t * data = static_cast<const int32_t *>(bucket->data);
+    for (int64_t i = 0; i < count; ++i) {
+        const int32_t row = data[i]; // bounds were checked above
+        if (seen[row] || mask[row] != 1) {
+            error = "GPU bucket contains a duplicate or an unselected row";
+            return false;
+        }
+        seen[row] = true;
+    }
+    selected = count;
+    return true;
+}
+
+} // namespace llama_gpu_split

llama.cpp

@@ -6,6 +6,7 @@
 #include "ggml.h"
 
 #include "ggml-alloc.h"
+#include "llama-gpu-split.h"
 
 #ifdef GGML_USE_CUBLAS
 #  include "ggml-cuda.h"
@@ -2754,30 +2755,6 @@ static void llm_load_print_meta(llama_model_loader & ml, llama_model & model) {
 }
 
 
-static int64_t sum_gpu_index(struct ggml_tensor * gpu_index) {
-    ggml_context * ctx_aux = ggml_init({
-        /* mem_size */ 1 << 10,
-    });
-
-    GGML_ASSERT(ctx_aux);
-
-    ggml_cgraph * gf = ggml_new_graph_custom(ctx_aux, 1, false);
-    ggml_tensor * sum = ggml_sum(ctx_aux, gpu_index);
-
-    ggml_set_name(sum, "gpu_index_sum");
-    ggml_build_forward_expand(gf, sum);
-
-    // TODO: +1 worker for GPU under hybrid inference but no use
-    // ggml_graph_compute_helper(work_buffer, gf, 2);
-    ggml_graph_compute_with_ctx(ctx_aux, gf, 2);
-
-    int32_t sum_val = ggml_get_i32_1d(sum, 0);
-
-    ggml_free(ctx_aux);
-
-    return sum_val;
-}
-
 struct llama_gpu_split_loader {
     int n_tensors = 0;
     size_t n_bytes = 0; // tensor data bytes
@@ -2827,27 +2804,60 @@ struct llama_gpu_split_loader {
         LLAMA_LOG_INFO("%s: applying gpu_idx adapter from '%s' - please wait ...\n", __func__, fname.c_str());
         const int64_t t_start_mlp_us = ggml_time_us();
 
+        std::vector<ggml_tensor *> gpu_indices(n_layers);
+        std::vector<ggml_tensor *> gpu_buckets(n_layers);
+        std::vector<int64_t> selected(n_layers);
         for (int il = 0; il < n_layers; il++) {
-            llama_layer &model_layer = model->layers[il];
-            ggml_tensor * gpu_idx = idx_loader->get_tensor_meta(il*2);
-            ggml_tensor * gpu_bucket = idx_loader->get_tensor_meta(il*2+1);
-            if (gpu_idx == nullptr || gpu_bucket == nullptr) {
-                LLAMA_LOG_ERROR("%s: error: failed to load gpu index or bucket\n", __func__);
+            const llama_layer & layer = model->layers[il];
+            ggml_tensor * gpu_idx = ggml_get_tensor(idx_loader->ctx_meta, format("blk.%d.gpu_idx", il).c_str());
+            ggml_tensor * gpu_bucket = ggml_get_tensor(idx_loader->ctx_meta, format("blk.%d.gpu_bucket", il).c_str());
+            std::string error;
+            if (!layer.ffn_up || !layer.ffn_down_t ||
+                layer.ffn_down_t->ne[1] != layer.ffn_up->ne[1] ||
+                (layer.ffn_gate && layer.ffn_gate->ne[1] != layer.ffn_up->ne[1]) ||
+                !llama_gpu_split::validate_layout(gpu_idx, gpu_bucket, layer.ffn_up->ne[1], error)) {
+                LLAMA_LOG_ERROR("%s: error: invalid GPU split for layer %d: %s\n", __func__, il,
+                    error.empty() ? "FFN tensors have incompatible row counts" : error.c_str());
                 return 1;
             }
-            model_layer.gpu_idx = idx_loader->create_tensor_for(ctx_meta, gpu_idx, GGML_BACKEND_CPU);
-            model_layer.gpu_bucket = idx_loader->create_tensor_for(ctx_meta, gpu_bucket, GGML_BACKEND_CPU);
+            // Validate sidecar data ranges before mapping and dereferencing any values.
+            // Subtraction avoids overflow when a file-controlled offset is very large.
+            for (const ggml_tensor * tensor : { gpu_idx, gpu_bucket }) {
+                const int id = gguf_find_tensor(idx_loader->ctx_gguf, ggml_get_name(tensor));
+                const uint64_t base = gguf_get_data_offset(idx_loader->ctx_gguf);
+                const uint64_t offset = gguf_get_tensor_offset(idx_loader->ctx_gguf, id);
+                const uint64_t size = idx_loader->file.size;
+                if (base > size || offset > size - base || ggml_nbytes(tensor) > size - base - offset) {
+                    LLAMA_LOG_ERROR("%s: error: GPU split tensor '%s' extends beyond the sidecar file\n",
+                        __func__, ggml_get_name(tensor));
+                    return 1;
+                }
+            }
+            gpu_indices[il] = idx_loader->create_tensor_for(ctx_meta, gpu_idx, GGML_BACKEND_CPU);
+            gpu_buckets[il] = idx_loader->create_tensor_for(ctx_meta, gpu_bucket, GGML_BACKEND_CPU);
         }
         llama_progress_callback cb = [](float progress, void *ctx) {
             LLAMA_LOG_INFO(".");
         };
         idx_loader->load_all_data(ctx_meta, cb, nullptr, nullptr);
 
+        // Reject the entire cache before publishing any layer state or uploading data.
+        // The caller may regenerate a rejected cache or fall back to an empty split.
+        for (int il = 0; il < n_layers; il++) {
+            std::string error;
+            if (!llama_gpu_split::validate(gpu_indices[il], gpu_buckets[il],
+                    model->layers[il].ffn_up->ne[1], selected[il], error)) {
+                LLAMA_LOG_ERROR("%s: error: invalid GPU split for layer %d: %s\n", __func__, il, error.c_str());
+                return 1;
+            }
+        }
         for (int il = 0; il < n_layers; il++) {
             llama_layer &model_layer = model->layers[il];
-            ggml_tensor * gpu_idx = model_layer.gpu_idx;
-            ggml_tensor * gpu_bucket = model_layer.gpu_bucket;
-            int64_t gpu_neurons = sum_gpu_index(gpu_idx);
+            ggml_tensor * gpu_idx = gpu_indices[il];
+            ggml_tensor * gpu_bucket = gpu_buckets[il];
+            model_layer.gpu_idx = gpu_idx;
+            model_layer.gpu_bucket = gpu_bucket;
+            int64_t gpu_neurons = selected[il];
             model_layer.gpu_offload_ratio = (double)gpu_neurons / gpu_idx->ne[0];
             if (gpu_neurons == 0 || gpu_neurons == gpu_idx->ne[0]) {
                 // no hybrid inference for this layer, unset gpu_bucket
@@ -2897,16 +2907,23 @@ struct llama_augmentation_model_loader {
         // allocate and copy selected weights to gpu
     ggml_tensor * create_striped_mat_to_gpu(struct ggml_tensor *src, struct ggml_tensor * gpu_bucket) {
 #ifdef GGML_USE_CUBLAS
+        if (!src || !src->data || src->ne[1] <= 0) {
+            throw std::runtime_error(format("%s: invalid source tensor", __func__));
+        }
         if (gpu_bucket == NULL) {
             // offload the whole tensor to gpu
             ggml_set_backend(src, GGML_BACKEND_GPU);
             ggml_cuda_transform_tensor(src->data, src);
             return src;
         }
 
+        std::string error;
+        if (!llama_gpu_split::validate_bucket(gpu_bucket, src->ne[1], error) || gpu_bucket->ne[0] == 0) {
+            throw std::runtime_error(format("%s: invalid GPU bucket: %s", __func__,
+                error.empty() ? "cannot slice an empty bucket" : error.c_str()));
+        }
         int64_t row_len = src->ne[0];
         int64_t gpu_rows = gpu_bucket->ne[0];
-        GGML_ASSERT(0 < gpu_rows && gpu_rows <= src->ne[1]);
 
         ggml_set_no_alloc(aux_ctx, true);
         ggml_tensor * gpu_dst = ggml_new_tensor_2d(aux_ctx, src->type, row_len, gpu_rows);

tests/CMakeLists.txt

@@ -45,6 +45,11 @@ llama_build_and_test_executable(test-grad0.cpp) # SLOW
 
 llama_build_and_test_executable(test-rope.cpp)
 
+# Host-side sidecar validation does not require a GPU or model weights.
+add_executable(test-gpu-split test-gpu-split.cpp)
+target_link_libraries(test-gpu-split PRIVATE ggml)
+add_test(NAME test-gpu-split COMMAND test-gpu-split)
+
 # dummy executable - not installed
 get_filename_component(TEST_TARGET test-c.c NAME_WE)
 add_executable(${TEST_TARGET} test-c.c)

tests/test-gpu-split.cpp

@@ -0,0 +1,108 @@
+#include "llama-gpu-split.h"
+
+#include <cstdio>
+#include <cstring>
+
+struct fixture {
+    ggml_context * ctx;
+    ggml_tensor * index;
+    ggml_tensor * bucket;
+
+    fixture() : ctx(ggml_init({ 1024 * 1024, nullptr, false })) {
+        index = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 4);
+        bucket = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 2);
+        ggml_set_zero(index);
+        ggml_set_i32_1d(index, 0, 1);
+        ggml_set_i32_1d(index, 3, 1);
+        ggml_set_i32_1d(bucket, 0, 0);
+        ggml_set_i32_1d(bucket, 1, 3);
+    }
+
+    ~fixture() { ggml_free(ctx); }
+};
+
+int main() {
+    const char * cases[] = {
+        "valid", "unsorted", "empty", "full", "negative", "upper-bound", "int-max",
+        "duplicate", "unselected", "count-mismatch", "non-binary", "index-type", "bucket-type",
+        "index-shape", "bucket-shape", "index-size", "oversized-bucket", "missing-index-data",
+        "missing-bucket-data", "null-index", "null-bucket", "zero-source-rows", "strided-bucket"
+    };
+    int failures = 0;
+    for (const char * name : cases) {
+        fixture f;
+        ggml_tensor * index = f.index;
+        ggml_tensor * bucket = f.bucket;
+        int64_t rows = 4;
+        int64_t expected_count = 2;
+        const bool expected_valid = std::strcmp(name, "valid") == 0 || std::strcmp(name, "unsorted") == 0 ||
+            std::strcmp(name, "empty") == 0 || std::strcmp(name, "full") == 0;
+        if (std::strcmp(name, "unsorted") == 0) {
+            ggml_set_i32_1d(f.bucket, 0, 3);
+            ggml_set_i32_1d(f.bucket, 1, 0);
+        } else if (std::strcmp(name, "empty") == 0) {
+            ggml_set_zero(f.index);
+            bucket = ggml_new_tensor_1d(f.ctx, GGML_TYPE_I32, 0);
+            bucket->data = nullptr; // no dereference is needed for an empty split
+            expected_count = 0;
+        } else if (std::strcmp(name, "full") == 0) {
+            bucket = ggml_new_tensor_1d(f.ctx, GGML_TYPE_I32, 4);
+            for (int i = 0; i < 4; ++i) {
+                ggml_set_i32_1d(f.index, i, 1);
+                ggml_set_i32_1d(bucket, i, i);
+            }
+            expected_count = 4;
+        } else if (std::strcmp(name, "negative") == 0) {
+            ggml_set_i32_1d(f.bucket, 1, -1);
+        } else if (std::strcmp(name, "upper-bound") == 0) {
+            ggml_set_i32_1d(f.bucket, 1, 4);
+        } else if (std::strcmp(name, "int-max") == 0) {
+            ggml_set_i32_1d(f.bucket, 1, INT32_MAX);
+        } else if (std::strcmp(name, "duplicate") == 0) {
+            ggml_set_i32_1d(f.bucket, 1, 0);
+        } else if (std::strcmp(name, "unselected") == 0) {
+            ggml_set_i32_1d(f.bucket, 1, 1);
+        } else if (std::strcmp(name, "count-mismatch") == 0) {
+            ggml_set_i32_1d(f.index, 1, 1);
+        } else if (std::strcmp(name, "non-binary") == 0) {
+            ggml_set_i32_1d(f.index, 0, 2);
+        } else if (std::strcmp(name, "index-type") == 0) {
+            index->type = GGML_TYPE_F32;
+        } else if (std::strcmp(name, "bucket-type") == 0) {
+            bucket->type = GGML_TYPE_F32;
+        } else if (std::strcmp(name, "index-shape") == 0) {
+            index = ggml_new_tensor_2d(f.ctx, GGML_TYPE_I32, 4, 2);
+        } else if (std::strcmp(name, "bucket-shape") == 0) {
+            bucket = ggml_new_tensor_2d(f.ctx, GGML_TYPE_I32, 2, 2);
+        } else if (std::strcmp(name, "index-size") == 0) {
+            rows = 3;
+        } else if (std::strcmp(name, "oversized-bucket") == 0) {
+            bucket = ggml_new_tensor_1d(f.ctx, GGML_TYPE_I32, 5);
+        } else if (std::strcmp(name, "missing-index-data") == 0) {
+            index->data = nullptr;
+        } else if (std::strcmp(name, "missing-bucket-data") == 0) {
+            bucket->data = nullptr;
+        } else if (std::strcmp(name, "null-index") == 0) {
+            index = nullptr;
+        } else if (std::strcmp(name, "null-bucket") == 0) {
+            bucket = nullptr;
+        } else if (std::strcmp(name, "zero-source-rows") == 0) {
+            rows = 0;
+        } else if (std::strcmp(name, "strided-bucket") == 0) {
+            bucket->nb[0] *= 2;
+        }
+
+        int64_t selected = -1;
+        std::string error;
+        const bool valid = llama_gpu_split::validate(index, bucket, rows, selected, error);
+        if (valid != expected_valid || (valid && selected != expected_count) ||
+            (!valid && (error.empty() || selected != -1))) {
+            std::fprintf(stderr, "FAIL %s: valid=%d selected=%lld error=%s\n",
+                name, int(valid), static_cast<long long>(selected), error.c_str());
+            ++failures;
+        } else {
+            std::printf("PASS %s\n", name);
+        }
+    }
+    return failures == 0 ? 0 : 1;
+}