perf: accelerate longhaul prompt loading
This commit is contained in:
Owen Qwen
@@ -0,0 +1,176 @@
|
||||
#include "testing.h"
|
||||
|
||||
#include "../src/llama-longhaul.h"
|
||||
|
||||
#include "ggml-backend.h"
|
||||
|
||||
#include <cstdio>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
struct longhaul_fixture {
|
||||
static constexpr size_t expert_size = 32;
|
||||
|
||||
FILE * file = nullptr;
|
||||
ggml_backend_ptr backend;
|
||||
ggml_context_ptr ctx;
|
||||
ggml_backend_buffer_ptr buffer;
|
||||
ggml_tensor * weights_a = nullptr;
|
||||
ggml_tensor * weights_b = nullptr;
|
||||
ggml_tensor * ids = nullptr;
|
||||
std::unique_ptr<llama_longhaul_cache> cache;
|
||||
|
||||
longhaul_fixture(
|
||||
size_t n_slots,
|
||||
uint32_t n_experts,
|
||||
bool two_sources = false,
|
||||
uint32_t written_experts = 0) {
|
||||
file = tmpfile();
|
||||
GGML_ASSERT(file != nullptr);
|
||||
write_experts(written_experts == 0 ? n_experts : written_experts, two_sources);
|
||||
|
||||
backend.reset(ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr));
|
||||
GGML_ASSERT(backend);
|
||||
|
||||
ggml_init_params params = {
|
||||
/*.mem_size =*/ 16 * 1024,
|
||||
/*.mem_buffer =*/ nullptr,
|
||||
/*.no_alloc =*/ true,
|
||||
};
|
||||
ctx.reset(ggml_init(params));
|
||||
GGML_ASSERT(ctx);
|
||||
|
||||
weights_a = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_I8, expert_size, 1, n_slots);
|
||||
if (two_sources) {
|
||||
weights_b = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_I8, expert_size, 1, n_slots);
|
||||
}
|
||||
ids = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, n_slots);
|
||||
buffer.reset(ggml_backend_alloc_ctx_tensors(ctx.get(), backend.get()));
|
||||
GGML_ASSERT(buffer);
|
||||
|
||||
llama_files files;
|
||||
files.emplace_back(std::make_unique<llama_file>(file));
|
||||
|
||||
std::vector<llama_model_loader::longhaul_source> sources = {
|
||||
{ weights_a, 0, 0, expert_size, 0 },
|
||||
};
|
||||
if (two_sources) {
|
||||
sources.push_back({ weights_b, 0, n_experts * expert_size, expert_size, 0 });
|
||||
}
|
||||
cache = std::make_unique<llama_longhaul_cache>(
|
||||
std::move(files), std::move(sources), n_slots, n_experts, 1);
|
||||
}
|
||||
|
||||
~longhaul_fixture() {
|
||||
cache.reset();
|
||||
if (file) {
|
||||
fclose(file);
|
||||
}
|
||||
}
|
||||
|
||||
void write_experts(uint32_t n_experts, bool two_sources) {
|
||||
GGML_ASSERT(fseek(file, 0, SEEK_END) == 0);
|
||||
for (int source = 0; source < (two_sources ? 2 : 1); ++source) {
|
||||
for (uint32_t expert = 0; expert < n_experts; ++expert) {
|
||||
std::vector<uint8_t> data(expert_size, uint8_t(1 + expert + source * 32));
|
||||
GGML_ASSERT(fwrite(data.data(), 1, data.size(), file) == data.size());
|
||||
}
|
||||
}
|
||||
GGML_ASSERT(fflush(file) == 0);
|
||||
}
|
||||
|
||||
std::vector<int32_t> remap(std::vector<int32_t> values) {
|
||||
GGML_ASSERT(values.size() <= (size_t) ids->ne[0]);
|
||||
ggml_backend_tensor_set(ids, values.data(), 0, values.size() * sizeof(int32_t));
|
||||
ids->ne[0] = values.size();
|
||||
const bool ok = cache->remap(0, ids);
|
||||
if (ok) {
|
||||
ggml_backend_tensor_get(ids, values.data(), 0, values.size() * sizeof(int32_t));
|
||||
}
|
||||
cache->release(0);
|
||||
return ok ? values : std::vector<int32_t>{};
|
||||
}
|
||||
|
||||
uint8_t slot_value(ggml_tensor * tensor, int32_t slot) {
|
||||
uint8_t value = 0;
|
||||
ggml_backend_tensor_get(tensor, &value, size_t(slot) * tensor->nb[2], 1);
|
||||
return value;
|
||||
}
|
||||
};
|
||||
|
||||
static void test_batch_planning(testing & t) {
|
||||
longhaul_fixture fixture(3, 4);
|
||||
|
||||
const auto initial = fixture.remap({0, 1, 2});
|
||||
t.assert_equal(3u, initial.size());
|
||||
t.assert_equal(3u, fixture.cache->misses());
|
||||
|
||||
// Expert 0 is the oldest cache entry. A sequential miss for expert 3 used
|
||||
// to evict it before the later 0 in this same routed batch was observed.
|
||||
const auto remapped = fixture.remap({3, 0, 3});
|
||||
t.assert_equal(3u, remapped.size());
|
||||
t.assert_equal("only expert 3 is loaded", 4u, fixture.cache->misses());
|
||||
t.assert_true("duplicates map to the same slot", remapped[0] == remapped[2]);
|
||||
t.assert_true("different experts map to different slots", remapped[0] != remapped[1]);
|
||||
t.assert_equal(uint8_t(4), fixture.slot_value(fixture.weights_a, remapped[0]));
|
||||
t.assert_equal(uint8_t(1), fixture.slot_value(fixture.weights_a, remapped[1]));
|
||||
}
|
||||
|
||||
static void test_multiple_sources(testing & t) {
|
||||
longhaul_fixture fixture(2, 4, true);
|
||||
|
||||
const auto remapped = fixture.remap({2, 1});
|
||||
t.assert_equal(2u, remapped.size());
|
||||
t.assert_equal(uint64_t(4 * longhaul_fixture::expert_size), fixture.cache->bytes_read_count());
|
||||
t.assert_equal(uint8_t(3), fixture.slot_value(fixture.weights_a, remapped[0]));
|
||||
t.assert_equal(uint8_t(35), fixture.slot_value(fixture.weights_b, remapped[0]));
|
||||
t.assert_equal(uint8_t(2), fixture.slot_value(fixture.weights_a, remapped[1]));
|
||||
t.assert_equal(uint8_t(34), fixture.slot_value(fixture.weights_b, remapped[1]));
|
||||
}
|
||||
|
||||
static void test_invalid_ids(testing & t) {
|
||||
longhaul_fixture fixture(2, 4);
|
||||
|
||||
const auto remapped = fixture.remap({0, 4});
|
||||
t.assert_true(remapped.empty());
|
||||
t.assert_true(fixture.cache->failed());
|
||||
t.assert_equal(0u, fixture.cache->misses());
|
||||
}
|
||||
|
||||
static void test_read_failure_recovery(testing & t) {
|
||||
longhaul_fixture fixture(1, 4, false, 3);
|
||||
|
||||
const auto failed = fixture.remap({3});
|
||||
t.assert_true(failed.empty());
|
||||
t.assert_true(fixture.cache->failed());
|
||||
t.assert_equal(0u, fixture.cache->misses());
|
||||
|
||||
fixture.write_experts(1, false);
|
||||
const auto remapped = fixture.remap({3});
|
||||
t.assert_equal(1u, remapped.size());
|
||||
t.assert_equal(1u, fixture.cache->misses());
|
||||
t.assert_equal(uint8_t(1), fixture.slot_value(fixture.weights_a, remapped[0]));
|
||||
}
|
||||
|
||||
int main(int argc, char ** argv) {
|
||||
testing t;
|
||||
|
||||
const char * verbose = getenv("LLAMA_TEST_VERBOSE");
|
||||
if (verbose) {
|
||||
t.verbose = std::string(verbose) == "1";
|
||||
}
|
||||
if (!t.verbose) {
|
||||
llama_log_set([](ggml_log_level, const char *, void *) {}, nullptr);
|
||||
}
|
||||
|
||||
if (argc > 1) {
|
||||
t.set_filter(argv[1]);
|
||||
}
|
||||
|
||||
t.test("batch_planning", test_batch_planning);
|
||||
t.test("multiple_sources", test_multiple_sources);
|
||||
t.test("invalid_ids", test_invalid_ids);
|
||||
t.test("read_failure", test_read_failure_recovery);
|
||||
|
||||
return t.summary();
|
||||
}
|
||||
Reference in New Issue
Block a user