Add portable CPU longhaul support

2026-07-30 16:59:40 -05:00
parent a673afad85
commit d5458aa44e
11 changed files with 277 additions and 23 deletions
@@ -259,6 +259,22 @@ set_tests_properties(test-thread-safety PROPERTIES FIXTURES_REQUIRED test-downlo

 llama_build_and_test(test-arg-parser.cpp)
 llama_build_and_test(test-longhaul.cpp)
+llama_test(
+    test-longhaul
+    NAME test-longhaul-cpu-qwen35moe
+    ARGS --cpu-model "${MODEL_DIR}/qwen35moe-moe.gguf" 3145728
+)
+set_tests_properties(test-longhaul-cpu-qwen35moe PROPERTIES
+    FIXTURES_REQUIRED generate-models
+)
+llama_test(
+    test-longhaul
+    NAME test-longhaul-cpu-laguna
+    ARGS --cpu-model "${MODEL_DIR}/laguna-moe.gguf" 2097152
+)
+set_tests_properties(test-longhaul-cpu-laguna PROPERTIES
+    FIXTURES_REQUIRED generate-models
+)
 llama_build_and_test(test-token-cache.cpp)
 target_include_directories(test-token-cache PRIVATE ${PROJECT_SOURCE_DIR}/src)

@@ -9,6 +9,7 @@

 // TODO: replace with #include "llama-ext.h" in the future
 #include "../src/llama-arch.h"
+#include "../src/llama-model.h"
 #include "../src/llama-model-saver.h"

 #include <cinttypes>
@@ -476,7 +477,9 @@ static int save_models(const llm_arch target_arch, const size_t seed, const ggml
            if (!moe && moe_mandatory(arch)) {
                continue;
            }
-            if (!llama_model_saver_supports_arch(arch) || !arch_supported(arch)) {
+            const bool can_save_fixture =
+                llama_model_saver_supports_arch(arch) || arch == LLM_ARCH_LAGUNA;
+            if (!can_save_fixture || !arch_supported(arch)) {
                LOG_INF("%s: %s model (%s) is unsupported, skipping\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense");
                continue;
            }
@@ -484,7 +487,20 @@ static int save_models(const llm_arch target_arch, const size_t seed, const ggml
            auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {});
            const std::string path = dir + "/" + llm_arch_name(arch) + (moe ? "-moe.gguf" : "-dense.gguf");
            LOG_INF("%s: Saving %s model (%s) to %s...\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense", path.c_str());
-            llama_model_save_to_file(model_and_ctx.first.get(), path.c_str());
+            if (llama_model_saver_supports_arch(arch)) {
+                llama_model_save_to_file(model_and_ctx.first.get(), path.c_str());
+            } else {
+                // Laguna is not supported by the production model saver yet.
+                // Preserve the exact synthetic fixture metadata and attach the
+                // initialized model tensors so longhaul can be tested from a
+                // real file without broadening the saver API.
+                gguf_context_ptr fixture(gguf_init_empty());
+                gguf_set_kv(fixture.get(), gguf_ctx.get());
+                for (const auto & item : llama_internal_get_tensor_map(model_and_ctx.first.get())) {
+                    gguf_add_tensor(fixture.get(), item.second);
+                }
+                gguf_write_to_file(fixture.get(), path.c_str(), false);
+            }
        }
    }
    llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);
@@ -1,11 +1,16 @@
 #include "testing.h"

 #include "../src/llama-longhaul.h"
+#include "../src/llama-model.h"

 #include "ggml-backend.h"
+#include "llama-cpp.h"

+#include <cmath>
 #include <cstdio>
 #include <memory>
+#include <stdexcept>
+#include <string>
 #include <vector>

 struct longhaul_fixture {
@@ -128,6 +133,32 @@ static void test_multiple_sources(testing & t) {
    t.assert_equal(uint8_t(34), fixture.slot_value(fixture.weights_b, remapped[1]));
 }

+static void test_concurrent_source_reads(testing & t) {
+    longhaul_fixture fixture(2, 4, true);
+
+    for (int iteration = 0; iteration < 64; ++iteration) {
+        const int32_t first = iteration % 2 == 0 ? 0 : 2;
+        const int32_t second = first + 1;
+        const auto remapped = fixture.remap({first, second});
+
+        if (!t.assert_equal(2u, remapped.size())) {
+            return;
+        }
+        t.assert_equal(
+            uint8_t(1 + first),
+            fixture.slot_value(fixture.weights_a, remapped[0]));
+        t.assert_equal(
+            uint8_t(33 + first),
+            fixture.slot_value(fixture.weights_b, remapped[0]));
+        t.assert_equal(
+            uint8_t(1 + second),
+            fixture.slot_value(fixture.weights_a, remapped[1]));
+        t.assert_equal(
+            uint8_t(33 + second),
+            fixture.slot_value(fixture.weights_b, remapped[1]));
+    }
+}
+
 static void test_invalid_ids(testing & t) {
    longhaul_fixture fixture(2, 4);

@@ -152,8 +183,110 @@ static void test_read_failure_recovery(testing & t) {
    t.assert_equal(uint8_t(1), fixture.slot_value(fixture.weights_a, remapped[0]));
 }

+struct cpu_decode_result {
+    std::vector<float> logits;
+    uint32_t capacity = 0;
+    uint64_t misses = 0;
+    uint64_t bytes_read = 0;
+    bool expert_buffer_is_host = false;
+};
+
+static cpu_decode_result decode_cpu_model(
+        const std::string & path,
+        llama_load_mode load_mode,
+        uint64_t cache_bytes) {
+    llama_model_params model_params = llama_model_default_params();
+    model_params.n_gpu_layers = 0;
+    model_params.load_mode = load_mode;
+    model_params.longhaul_cache_bytes = cache_bytes;
+    model_params.use_extra_bufts = true;
+
+    llama_model_ptr model(llama_model_load_from_file(path.c_str(), model_params));
+    if (!model) {
+        throw std::runtime_error("failed to load CPU model: " + path);
+    }
+
+    llama_context_params context_params = llama_context_default_params();
+    context_params.n_ctx = 32;
+    context_params.n_batch = 8;
+    context_params.n_ubatch = 8;
+    context_params.n_threads = 2;
+    context_params.n_threads_batch = 2;
+
+    llama_context_ptr context(llama_init_from_model(model.get(), context_params));
+    if (!context) {
+        throw std::runtime_error("failed to create CPU context: " + path);
+    }
+
+    std::vector<llama_token> tokens = {1, 2, 3, 4, 5, 6, 7, 8};
+    llama_batch batch = llama_batch_get_one(tokens.data(), tokens.size());
+    if (llama_decode(context.get(), batch) != 0) {
+        throw std::runtime_error("failed to decode CPU model: " + path);
+    }
+
+    const int32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model.get()));
+    const float * logits = llama_get_logits_ith(context.get(), -1);
+    if (!logits) {
+        throw std::runtime_error("CPU model produced no logits: " + path);
+    }
+
+    cpu_decode_result result;
+    result.logits.assign(logits, logits + n_vocab);
+
+    if (load_mode == LLAMA_LOAD_MODE_LONGHAUL) {
+        llama_longhaul_cache * cache = model->longhaul_cache();
+        if (!cache) {
+            throw std::runtime_error("longhaul cache was not created: " + path);
+        }
+        result.capacity = cache->capacity();
+        result.misses = cache->misses();
+        result.bytes_read = cache->bytes_read_count();
+
+        for (const auto & item : llama_internal_get_tensor_map(model.get())) {
+            const std::string & name = item.first;
+            if (name.find("ffn_") != std::string::npos &&
+                name.find("_exps.weight") != std::string::npos &&
+                item.second->ne[2] == result.capacity) {
+                result.expert_buffer_is_host = ggml_backend_buffer_is_host(item.second->buffer);
+                break;
+            }
+        }
+    }
+
+    return result;
+}
+
+static void test_cpu_model(
+        testing & t,
+        const std::string & path,
+        uint64_t cache_bytes) {
+    const cpu_decode_result regular =
+        decode_cpu_model(path, LLAMA_LOAD_MODE_NONE, 0);
+    const cpu_decode_result longhaul =
+        decode_cpu_model(path, LLAMA_LOAD_MODE_LONGHAUL, cache_bytes);
+
+    t.assert_equal(1u, longhaul.capacity);
+    t.assert_true("longhaul loaded at least one expert", longhaul.misses > 0);
+    t.assert_true("longhaul read expert bytes", longhaul.bytes_read > 0);
+    t.assert_true("streamed experts use a directly writable host buffer", longhaul.expert_buffer_is_host);
+    t.assert_equal(regular.logits.size(), longhaul.logits.size());
+
+    double squared_error = 0.0;
+    double squared_reference = 0.0;
+    for (size_t i = 0; i < regular.logits.size(); ++i) {
+        const double delta = double(regular.logits[i]) - double(longhaul.logits[i]);
+        squared_error += delta * delta;
+        squared_reference += double(regular.logits[i]) * double(regular.logits[i]);
+    }
+    const double nmse = squared_reference > 0.0 ? squared_error / squared_reference : squared_error;
+    t.assert_true(
+        "longhaul CPU logits NMSE = " + std::to_string(nmse),
+        std::isfinite(nmse) && nmse <= 1e-6);
+}
+
 int main(int argc, char ** argv) {
    testing t;
+    llama_backend_init();

    const char * verbose = getenv("LLAMA_TEST_VERBOSE");
    if (verbose) {
@@ -163,14 +296,28 @@ int main(int argc, char ** argv) {
        llama_log_set([](ggml_log_level, const char *, void *) {}, nullptr);
    }

+    if (argc == 4 && std::string(argv[1]) == "--cpu-model") {
+        const std::string path = argv[2];
+        const uint64_t cache_bytes = std::stoull(argv[3]);
+        t.test("cpu_model", [&](testing & current) {
+            test_cpu_model(current, path, cache_bytes);
+        });
+        const int result = t.summary();
+        llama_backend_free();
+        return result;
+    }
+
    if (argc > 1) {
        t.set_filter(argv[1]);
    }

    t.test("batch_planning",  test_batch_planning);
    t.test("multiple_sources", test_multiple_sources);
+    t.test("concurrent_reads", test_concurrent_source_reads);
    t.test("invalid_ids",     test_invalid_ids);
    t.test("read_failure",    test_read_failure_recovery);

-    return t.summary();
+    const int result = t.summary();
+    llama_backend_free();
+    return result;
 }