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
@@ -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;
 }