From 0dc3be24b38af2b1dc7919f790a2b1ff6ee0541a Mon Sep 17 00:00:00 2001 From: dandan20031031 <2425926169@qq.com> Date: Tue, 27 Oct 2026 16:37:27 +0800 Subject: [PATCH] loader: validate GPU split sidecars before applying them --- CMakeLists.txt | 1 + llama-gpu-split.h | 78 ++++++++++++++++++++++++++++ llama.cpp | 87 ++++++++++++++++++------------- tests/CMakeLists.txt | 5 ++ tests/test-gpu-split.cpp | 108 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 244 insertions(+), 35 deletions(-) create mode 100644 llama-gpu-split.h create mode 100644 tests/test-gpu-split.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 87e9c814e..09a9a5367 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -678,6 +678,7 @@ endif() add_library(llama llama.cpp llama.h + llama-gpu-split.h ) target_include_directories(llama PUBLIC .) diff --git a/llama-gpu-split.h b/llama-gpu-split.h new file mode 100644 index 000000000..5ba50134c --- /dev/null +++ b/llama-gpu-split.h @@ -0,0 +1,78 @@ +#pragma once + +#include "ggml.h" + +#include +#include +#include + +// 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(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(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 seen(static_cast(rows), false); + const int32_t * data = static_cast(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 diff --git a/llama.cpp b/llama.cpp index ac52908ac..f1ba5064e 100644 --- a/llama.cpp +++ b/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 gpu_indices(n_layers); + std::vector gpu_buckets(n_layers); + std::vector 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,6 +2907,9 @@ 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); @@ -2904,9 +2917,13 @@ struct llama_augmentation_model_loader { 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); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c8b4bc254..0afe1b1b6 100644 --- a/tests/CMakeLists.txt +++ b/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) diff --git a/tests/test-gpu-split.cpp b/tests/test-gpu-split.cpp new file mode 100644 index 000000000..bf37cb894 --- /dev/null +++ b/tests/test-gpu-split.cpp @@ -0,0 +1,108 @@ +#include "llama-gpu-split.h" + +#include +#include + +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(selected), error.c_str()); + ++failures; + } else { + std::printf("PASS %s\n", name); + } + } + return failures == 0 ? 0 : 1; +}