diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index b3d17330ca54..3fa23c44da93 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -33,6 +33,8 @@ set(LLAMA_CORE_SOURCES llama-memory-hybrid-idx.cpp llama-memory-recurrent.cpp llama-mmap.cpp + llama-bounded-file.cpp + llama-engram.cpp llama-ple-disk.cpp llama-model-loader.cpp llama-model-saver.cpp diff --git a/src/llama-bounded-file.cpp b/src/llama-bounded-file.cpp new file mode 100644 index 000000000000..3ae6d5c31a18 --- /dev/null +++ b/src/llama-bounded-file.cpp @@ -0,0 +1,244 @@ +#include "llama-bounded-file.h" + +#include "llama-impl.h" + +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#endif + +struct llama_bounded_file::buffer::impl { + void * ptr = nullptr; + size_t bytes = 0; + + ~impl() { + free(ptr); + } +}; + +llama_bounded_file::buffer::buffer() : pimpl(std::make_unique()) {} +llama_bounded_file::buffer::~buffer() = default; +llama_bounded_file::buffer::buffer(buffer && other) noexcept = default; +llama_bounded_file::buffer & llama_bounded_file::buffer::operator=(buffer && other) noexcept = default; + +bool llama_bounded_file::buffer::empty() const { + return pimpl->ptr == nullptr; +} + +struct llama_bounded_file::impl { + std::string fname; + int fd = -1; + uint64_t file_size = 0; + size_t block = 4096; + bool direct = false; + + impl(const std::string & fname, const params & p) : fname(fname) { +#if defined(_WIN32) + GGML_UNUSED(p); + throw std::runtime_error("llama_bounded_file: not supported on Windows"); +#else + int direct_error = 0; + if (p.direct_io) { +#if defined(O_DIRECT) + fd = open(fname.c_str(), O_RDONLY | O_DIRECT | O_CLOEXEC); + direct = fd >= 0; + direct_error = direct ? 0 : errno; +#else + fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); + direct_error = fd >= 0 ? ENOTSUP : errno; +#if defined(F_NOCACHE) + if (fd >= 0) { + if (fcntl(fd, F_NOCACHE, 1) == 0) { + direct = true; +#if defined(F_RDAHEAD) + if (fcntl(fd, F_RDAHEAD, 0) != 0) { + direct = false; + direct_error = errno; + } +#endif + } else { + direct_error = errno; + } + } +#endif +#endif + if (!direct && p.direct_io_required) { + const int saved = direct_error ? direct_error : ENOTSUP; + if (fd >= 0) { + close(fd); + fd = -1; + } + throw std::runtime_error(format("llama_bounded_file: uncached open of %s failed: %s", + fname.c_str(), strerror(saved))); + } + if (!direct) { + LLAMA_LOG_WARN("%s: uncached open of %s failed (%s); falling back to buffered reads\n", + __func__, fname.c_str(), strerror(direct_error ? direct_error : ENOTSUP)); + } + if (!direct && fd >= 0) { + close(fd); + fd = -1; + } + } + if (fd < 0) { + fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); + if (fd < 0) { + throw std::runtime_error(format("llama_bounded_file: failed to open %s: %s", + fname.c_str(), strerror(errno))); + } + } + + struct stat st = {}; + if (fstat(fd, &st) != 0) { + const int saved = errno; + close(fd); + fd = -1; + throw std::runtime_error(format("llama_bounded_file: fstat of %s failed: %s", + fname.c_str(), strerror(saved))); + } + if (!S_ISREG(st.st_mode) || st.st_size < 0) { + close(fd); + fd = -1; + throw std::runtime_error(format("llama_bounded_file: %s is not a regular file", fname.c_str())); + } + file_size = (uint64_t) st.st_size; +#endif + } + + ~impl() { +#if !defined(_WIN32) + if (fd >= 0) { + close(fd); + } +#endif + } + + size_t read_size(uint64_t offset, size_t size) const { + if (!direct || size == 0) { + return size; + } + const size_t prefix = (size_t) (offset & (block - 1)); + if (size > SIZE_MAX - prefix || size + prefix > SIZE_MAX - (block - 1)) { + throw std::overflow_error("llama_bounded_file: aligned read size overflow"); + } + return (size + prefix + block - 1) & ~(block - 1); + } + +#if !defined(_WIN32) + void pread_full(void * dst, size_t len, uint64_t offset, size_t need) const { + size_t done = 0; + while (done < len) { + const size_t remaining = len - done; + const ssize_t n = pread(fd, (uint8_t *) dst + done, remaining, (off_t) (offset + done)); + if (n < 0) { + if (errno == EINTR) { + continue; + } + throw std::runtime_error(format("llama_bounded_file: pread of %s failed at %llu: %s", + fname.c_str(), (unsigned long long) (offset + done), strerror(errno))); + } + if (n == 0) { + break; + } + done += (size_t) n; + if (direct && done >= need) { + return; + } + if (direct && (size_t) n < remaining) { + break; + } + } + if (done < need) { + throw std::runtime_error(format("llama_bounded_file: short read in %s: %zu of %zu bytes at %llu", + fname.c_str(), done, need, (unsigned long long) offset)); + } + } +#endif +}; + +llama_bounded_file::llama_bounded_file(const std::string & fname, const params & p) + : pimpl(std::make_unique(fname, p)) {} + +llama_bounded_file::~llama_bounded_file() = default; + +llama_bounded_file::buffer llama_bounded_file::make_buffer(size_t size) const { + buffer result; + if (!pimpl->direct || size == 0) { + return result; + } +#if defined(_WIN32) + GGML_UNUSED(size); + throw std::runtime_error("llama_bounded_file: not supported on Windows"); +#else + const size_t bytes = pimpl->read_size(pimpl->block - 1, size); + void * ptr = nullptr; + const int err = posix_memalign(&ptr, pimpl->block, bytes); + if (err != 0) { + throw std::runtime_error(format("llama_bounded_file: posix_memalign failed: %s", strerror(err))); + } + result.pimpl->ptr = ptr; + result.pimpl->bytes = bytes; + return result; +#endif +} + +void llama_bounded_file::read(uint64_t offset, void * dst, size_t size, buffer & scratch) const { +#if defined(_WIN32) + GGML_UNUSED(offset); + GGML_UNUSED(dst); + GGML_UNUSED(size); + GGML_UNUSED(scratch); + throw std::runtime_error("llama_bounded_file: not supported on Windows"); +#else + if (size == 0) { + return; + } + if (dst == nullptr || offset > pimpl->file_size || size > pimpl->file_size - offset || + offset > (uint64_t) INT64_MAX || size > (uint64_t) INT64_MAX - offset) { + throw std::invalid_argument(format("llama_bounded_file: invalid read of %zu bytes at %llu in %s", + size, (unsigned long long) offset, pimpl->fname.c_str())); + } + if (!pimpl->direct) { + pimpl->pread_full(dst, size, offset, size); + return; + } + + const uint64_t aligned_offset = offset & ~(uint64_t) (pimpl->block - 1); + const size_t prefix = (size_t) (offset - aligned_offset); + const size_t aligned_size = pimpl->read_size(offset, size); + if (aligned_size > (uint64_t) INT64_MAX - aligned_offset) { + throw std::overflow_error("llama_bounded_file: aligned read offset overflow"); + } + if (scratch.pimpl->ptr == nullptr || scratch.pimpl->bytes < aligned_size) { + throw std::invalid_argument("llama_bounded_file: aligned scratch buffer is too small"); + } + pimpl->pread_full(scratch.pimpl->ptr, aligned_size, aligned_offset, prefix + size); + memcpy(dst, (uint8_t *) scratch.pimpl->ptr + prefix, size); +#endif +} + +uint64_t llama_bounded_file::size() const { + return pimpl->file_size; +} + +size_t llama_bounded_file::alignment() const { + return pimpl->block; +} + +size_t llama_bounded_file::read_size(uint64_t offset, size_t size) const { + return pimpl->read_size(offset, size); +} + +bool llama_bounded_file::direct_io() const { + return pimpl->direct; +} + +const std::string & llama_bounded_file::name() const { + return pimpl->fname; +} diff --git a/src/llama-bounded-file.h b/src/llama-bounded-file.h new file mode 100644 index 000000000000..3af763580a36 --- /dev/null +++ b/src/llama-bounded-file.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include +#include + +struct llama_bounded_file { + struct params { + bool direct_io = true; + bool direct_io_required = false; // fail instead of using the page cache + }; + + struct buffer { + buffer(); + ~buffer(); + buffer(buffer && other) noexcept; + buffer & operator=(buffer && other) noexcept; + + buffer(const buffer &) = delete; + buffer & operator=(const buffer &) = delete; + + bool empty() const; + + struct impl; + std::unique_ptr pimpl; + }; + + llama_bounded_file(const std::string & fname, const params & p); + ~llama_bounded_file(); + + buffer make_buffer(size_t size) const; + // Read only the requested extent. Direct I/O may read its containing aligned blocks into scratch. + void read(uint64_t offset, void * dst, size_t size, buffer & scratch) const; + + uint64_t size() const; + size_t alignment() const; + size_t read_size(uint64_t offset, size_t size) const; + bool direct_io() const; + const std::string & name() const; + + struct impl; + std::unique_ptr pimpl; +}; diff --git a/src/llama-engram.cpp b/src/llama-engram.cpp new file mode 100644 index 000000000000..336aae39f2d7 --- /dev/null +++ b/src/llama-engram.cpp @@ -0,0 +1,305 @@ +#include "llama-engram.h" + +#include "llama-bounded-file.h" +#include "llama-impl.h" + +#include +#include +#include +#include +#include +#include + +static void llama_engram_validate_layout(const llama_engram_layout & layout) { + if (layout.encoding != "e4m3_e8m0_32_row264") { + throw std::invalid_argument("llama_engram: unsupported row encoding"); + } + if (layout.layer_ids[0] == layout.layer_ids[1]) { + throw std::invalid_argument("llama_engram: layer IDs must be distinct"); + } + if (layout.token_map.empty() || layout.token_map.size() > (size_t) INT32_MAX) { + throw std::invalid_argument("llama_engram: invalid token map size"); + } + if (layout.compressed_vocab_size == 0 || layout.compressed_vocab_size > (uint32_t) INT32_MAX || + layout.pad_id >= layout.compressed_vocab_size) { + throw std::invalid_argument("llama_engram: invalid compressed vocabulary"); + } + for (uint32_t token : layout.token_map) { + if (token >= layout.compressed_vocab_size) { + throw std::invalid_argument("llama_engram: token map entry is out of range"); + } + } + + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t i = 0; i < LLAMA_ENGRAM_NGRAM; ++i) { + const uint64_t multiplier = layout.multipliers[layer][i]; + if ((multiplier & 1) == 0 || + multiplier > (uint64_t) INT64_MAX / layout.compressed_vocab_size) { + throw std::invalid_argument("llama_engram: invalid hash multiplier"); + } + } + + uint64_t row_count = 0; + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + const uint32_t prime = layout.primes[layer][col]; + if (prime < 2) { + throw std::invalid_argument("llama_engram: invalid hash prime"); + } + row_count += prime; + } + if (row_count > UINT32_MAX || row_count != layout.rows[layer]) { + throw std::invalid_argument("llama_engram: row extent does not match hash buckets"); + } + } +} + +void llama_engram_history::reset() { + tail.fill(LLAMA_ENGRAM_DEAD); +} + +struct llama_engram_hasher::impl { + llama_engram_layout layout; + + explicit impl(llama_engram_layout layout) : layout(std::move(layout)) { + llama_engram_validate_layout(this->layout); + } + + void validate_history(const llama_engram_history & history) const { + for (int32_t token : history.tail) { + if (token < LLAMA_ENGRAM_DEAD || + (token >= 0 && (uint32_t) token >= layout.compressed_vocab_size)) { + throw std::invalid_argument("llama_engram: invalid token history"); + } + } + } + + void hash( + llama_engram_history & history, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * rows) const { + if (count > SIZE_MAX / (LLAMA_ENGRAM_LAYERS * LLAMA_ENGRAM_COLS * sizeof(*rows)) || + (count != 0 && (tokens == nullptr || rows == nullptr))) { + throw std::invalid_argument("llama_engram: invalid hash buffers"); + } + validate_history(history); + for (size_t i = 0; i < count; ++i) { + if (tokens[i] < 0 || (size_t) tokens[i] >= layout.token_map.size()) { + throw std::invalid_argument("llama_engram: token ID is out of range"); + } + } + + llama_engram_history next = history; + uint32_t * output = rows; + for (size_t pos = 0; pos < count; ++pos) { + const int32_t current = mask != nullptr && mask[pos] == 0 ? + LLAMA_ENGRAM_DEAD : (int32_t) layout.token_map[tokens[pos]]; + uint32_t ids[LLAMA_ENGRAM_NGRAM]; + bool blocked = false; + for (size_t depth = 0; depth < LLAMA_ENGRAM_NGRAM; ++depth) { + const int32_t id = depth == 0 ? current : next.tail[depth - 1]; + blocked = blocked || id == LLAMA_ENGRAM_DEAD; + ids[depth] = blocked ? layout.pad_id : (uint32_t) id; + } + + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + uint64_t hash = (uint64_t) ids[0] * layout.multipliers[layer][0]; + uint64_t offset = 0; + for (size_t depth = 1; depth < LLAMA_ENGRAM_NGRAM; ++depth) { + hash ^= (uint64_t) ids[depth] * layout.multipliers[layer][depth]; + for (size_t head = 0; head < LLAMA_ENGRAM_HEADS; ++head) { + const size_t col = (depth - 1) * LLAMA_ENGRAM_HEADS + head; + const uint32_t prime = layout.primes[layer][col]; + *output++ = (uint32_t) (hash % prime + offset); + offset += prime; + } + } + } + + for (size_t i = next.tail.size() - 1; i > 0; --i) { + next.tail[i] = next.tail[i - 1]; + } + next.tail[0] = current; + } + history = next; + } +}; + +llama_engram_hasher::llama_engram_hasher(llama_engram_layout layout) + : pimpl(std::make_unique(std::move(layout))) {} + +llama_engram_hasher::~llama_engram_hasher() = default; + +void llama_engram_hasher::hash( + llama_engram_history & history, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * rows) const { + pimpl->hash(history, tokens, mask, count, rows); +} + +const llama_engram_layout & llama_engram_hasher::layout() const { + return pimpl->layout; +} + +static float llama_engram_e4m3(uint8_t code) { + const int exponent = (code >> 3) & 15; + const int mantissa = code & 7; + const float value = exponent != 0 ? + std::ldexp((float) (8 + mantissa), exponent - 10) : + std::ldexp((float) mantissa, -9); + return (code & 128) != 0 ? -value : value; +} + +void llama_engram_decode_row(const uint8_t row[LLAMA_ENGRAM_ROW_BYTES], float output[LLAMA_ENGRAM_DIM]) { + if (row == nullptr || output == nullptr) { + throw std::invalid_argument("llama_engram: invalid row decode buffers"); + } + + float decoded[LLAMA_ENGRAM_DIM]; + for (size_t i = 0; i < LLAMA_ENGRAM_DIM; ++i) { + const uint8_t code = row[i]; + const uint8_t scale = row[LLAMA_ENGRAM_DIM + i / 32]; + if ((code & 127) == 127) { + throw std::domain_error("llama_engram: E4M3 NaN encoding"); + } + if (scale == 255) { + throw std::domain_error("llama_engram: E8M0 scale 255"); + } + + float value = std::ldexp(llama_engram_e4m3(code), (int) scale - 127); + uint32_t bits; + memcpy(&bits, &value, sizeof(bits)); + bits = (bits + 0x7fffu + ((bits >> 16) & 1u)) & 0xffff0000u; + memcpy(&value, &bits, sizeof(value)); + if (!std::isfinite(value)) { + throw std::domain_error("llama_engram: decoded value is not finite"); + } + decoded[i] = value; + } + memcpy(output, decoded, sizeof(decoded)); +} + +struct llama_engram_table::impl { + struct request { + uint32_t row; + uint32_t output; + }; + + static constexpr size_t BATCH_TOKENS = 2048; + + llama_bounded_file file; + llama_bounded_file::buffer scratch; + uint64_t offset; + uint32_t rows; + std::mutex mutex; + std::vector requests; + + impl(const std::string & fname, uint64_t offset, uint32_t rows) + : file(fname, { true, true }), + scratch(file.make_buffer(LLAMA_ENGRAM_ROW_BYTES)), + offset(offset), + rows(rows) { + const uint64_t bytes = (uint64_t) rows * LLAMA_ENGRAM_ROW_BYTES; + if (rows == 0 || offset > (uint64_t) INT64_MAX || bytes > (uint64_t) INT64_MAX - offset || + offset > file.size() || bytes > file.size() - offset) { + throw std::invalid_argument("llama_engram: invalid table extent"); + } + requests.reserve(BATCH_TOKENS * LLAMA_ENGRAM_COLS); + } + + void validate_rows(const uint32_t * row_ids, size_t count) const { + if (count > SIZE_MAX / (LLAMA_ENGRAM_DIM * sizeof(float)) || + (count != 0 && row_ids == nullptr)) { + throw std::invalid_argument("llama_engram: invalid row list"); + } + for (size_t i = 0; i < count; ++i) { + if (row_ids[i] >= rows) { + throw std::invalid_argument("llama_engram: row ID is out of range"); + } + } + } + + void read_one(uint32_t row, float * output) { + uint8_t raw[LLAMA_ENGRAM_ROW_BYTES]; + file.read(offset + (uint64_t) row * LLAMA_ENGRAM_ROW_BYTES, raw, sizeof(raw), scratch); + llama_engram_decode_row(raw, output); + } + + void read(const uint32_t * row_ids, size_t count, float * output) { + validate_rows(row_ids, count); + if (count != 0 && output == nullptr) { + throw std::invalid_argument("llama_engram: invalid row output"); + } + + std::lock_guard lock(mutex); + for (size_t i = 0; i < count; ++i) { + read_one(row_ids[i], output + i * LLAMA_ENGRAM_DIM); + } + } + + void read_batch(const uint32_t * row_ids, size_t tokens, size_t stride, float * output) { + if (tokens > SIZE_MAX / (LLAMA_ENGRAM_COLS * LLAMA_ENGRAM_DIM * sizeof(float)) || + (tokens != 0 && (row_ids == nullptr || output == nullptr || stride < LLAMA_ENGRAM_COLS)) || + (tokens != 0 && tokens - 1 > (SIZE_MAX / sizeof(*row_ids) - LLAMA_ENGRAM_COLS) / stride)) { + throw std::invalid_argument("llama_engram: invalid batch buffers"); + } + for (size_t token = 0; token < tokens; ++token) { + validate_rows(row_ids + token * stride, LLAMA_ENGRAM_COLS); + } + if (tokens == 0) { + return; + } + + std::lock_guard lock(mutex); + for (size_t start = 0; start < tokens; start += BATCH_TOKENS) { + const size_t count_tokens = std::min(tokens - start, BATCH_TOKENS); + const size_t count_rows = count_tokens * LLAMA_ENGRAM_COLS; + requests.resize(count_rows); + for (size_t i = 0; i < count_rows; ++i) { + requests[i].row = row_ids[(start + i / LLAMA_ENGRAM_COLS) * stride + i % LLAMA_ENGRAM_COLS]; + requests[i].output = (uint32_t) i; + } + std::sort(requests.begin(), requests.end(), [](const request & a, const request & b) { + return a.row < b.row; + }); + + float * chunk = output + start * LLAMA_ENGRAM_COLS * LLAMA_ENGRAM_DIM; + const float * previous = nullptr; + for (size_t i = 0; i < count_rows; ++i) { + float * dst = chunk + (size_t) requests[i].output * LLAMA_ENGRAM_DIM; + if (i != 0 && requests[i].row == requests[i - 1].row) { + memcpy(dst, previous, LLAMA_ENGRAM_DIM * sizeof(*dst)); + } else { + read_one(requests[i].row, dst); + previous = dst; + } + } + } + } +}; + +llama_engram_table::llama_engram_table(const std::string & fname, uint64_t offset, uint32_t rows) + : pimpl(std::make_unique(fname, offset, rows)) {} + +llama_engram_table::~llama_engram_table() = default; + +void llama_engram_table::read(const uint32_t * rows, size_t count, float * output) { + pimpl->read(rows, count, output); +} + +void llama_engram_table::read_batch(const uint32_t * rows, size_t tokens, size_t stride, float * output) { + pimpl->read_batch(rows, tokens, stride, output); +} + +uint32_t llama_engram_table::n_rows() const { + return pimpl->rows; +} + +std::string llama_engram_table::describe() const { + return format("%s @ %llu: %u rows x %u bytes, uncached aligned reads", + pimpl->file.name().c_str(), (unsigned long long) pimpl->offset, + pimpl->rows, LLAMA_ENGRAM_ROW_BYTES); +} diff --git a/src/llama-engram.h b/src/llama-engram.h new file mode 100644 index 000000000000..6c1c0f78b34c --- /dev/null +++ b/src/llama-engram.h @@ -0,0 +1,73 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +enum { + LLAMA_ENGRAM_LAYERS = 2, + LLAMA_ENGRAM_NGRAM = 4, + LLAMA_ENGRAM_HEADS = 8, + LLAMA_ENGRAM_COLS = 24, + LLAMA_ENGRAM_DIM = 256, + LLAMA_ENGRAM_ROW_BYTES = 264, + LLAMA_ENGRAM_DEAD = -1, +}; + +struct llama_engram_layout { + std::string encoding; + std::array layer_ids = {}; + std::vector token_map; + uint32_t compressed_vocab_size = 0; + uint32_t pad_id = 0; + std::array rows = {}; + std::array, LLAMA_ENGRAM_LAYERS> multipliers = {}; + std::array, LLAMA_ENGRAM_LAYERS> primes = {}; +}; + +struct llama_engram_history { + // Newest compressed token first. LLAMA_ENGRAM_DEAD breaks all n-grams that cross it. + std::array tail = {}; + + void reset(); +}; + +struct llama_engram_hasher { + explicit llama_engram_hasher(llama_engram_layout layout); + ~llama_engram_hasher(); + + // Output is [token][layer][column]. Invalid input does not change history or output. + void hash( + llama_engram_history & history, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * rows) const; + + const llama_engram_layout & layout() const; + + struct impl; + std::unique_ptr pimpl; +}; + +// Decode 256 E4M3 values and eight E8M0 scales to BF16-rounded values in F32 storage. +void llama_engram_decode_row(const uint8_t row[LLAMA_ENGRAM_ROW_BYTES], float output[LLAMA_ENGRAM_DIM]); + +struct llama_engram_table { + // The table is never mapped or cached. Construction fails unless uncached reads are available. + llama_engram_table(const std::string & fname, uint64_t offset, uint32_t rows); + ~llama_engram_table(); + + void read(const uint32_t * rows, size_t count, float * output); + // Read 24 rows per token. Input uses row-ID stride; output is packed [token][column][dimension]. + void read_batch(const uint32_t * rows, size_t tokens, size_t stride, float * output); + + uint32_t n_rows() const; + std::string describe() const; + + struct impl; + std::unique_ptr pimpl; +}; diff --git a/src/llama-ple-disk.cpp b/src/llama-ple-disk.cpp index 08263746b863..45f38a82c603 100644 --- a/src/llama-ple-disk.cpp +++ b/src/llama-ple-disk.cpp @@ -1,5 +1,6 @@ #include "llama-ple-disk.h" +#include "llama-bounded-file.h" #include "llama-impl.h" #include @@ -7,28 +8,22 @@ #include #include #include +#include #include #include #include #include -#if !defined(_WIN32) -#include -#include -#include -#endif - struct llama_ple_disk::impl { std::string fname; - int fd = -1; bool direct = false; size_t offs = 0; ggml_type type = GGML_TYPE_COUNT; int64_t ne0 = 0; int64_t nrows = 0; size_t rs = 0; // bytes per row - size_t block = 4096; // O_DIRECT alignment for offset, length and buffer ggml_to_float_t to_float = nullptr; + std::unique_ptr file; // direct-mapped cache of raw rows: slot = row & mask, tag = row size_t n_slots = 0; @@ -40,7 +35,7 @@ struct llama_ple_disk::impl { std::vector uniq; std::vector raw; // [uniq.size(), rs] std::vector> misses; // (row, index into uniq) - uint8_t * bounce0 = nullptr; // main-thread bounce buffer + llama_bounded_file::buffer bounce0; std::mutex mtx; // one gather at a time; a model is shared by every context built on it @@ -53,6 +48,7 @@ struct llama_ple_disk::impl { uint64_t gen = 0; size_t pending = 0; std::atomic next{0}; + std::exception_ptr worker_error; bool stop = false; uint64_t st_calls = 0, st_rows = 0, st_uniq = 0, st_hits = 0, st_reads = 0, st_bytes = 0; @@ -63,6 +59,9 @@ struct llama_ple_disk::impl { #if defined(_WIN32) throw std::runtime_error("llama_ple_disk: not supported on Windows"); #else + if (ne0 <= 0 || nrows <= 0) { + throw std::runtime_error("llama_ple_disk: invalid table shape"); + } if (ggml_is_quantized(type) && ne0 % ggml_blck_size(type) != 0) { throw std::runtime_error(format("llama_ple_disk: row of %lld %s elements is not a whole number of blocks", (long long) ne0, ggml_type_name(type))); @@ -75,35 +74,16 @@ struct llama_ple_disk::impl { } } - direct = p.direct_io; - if (direct) { - // only Linux has O_DIRECT; Darwin turns the page cache off per descriptor - // with F_NOCACHE after the open -#if defined(O_DIRECT) - fd = open(fname.c_str(), O_RDONLY | O_DIRECT | O_CLOEXEC); -#else - fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); -#if defined(F_NOCACHE) - if (fd >= 0 && fcntl(fd, F_NOCACHE, 1) < 0) { - LLAMA_LOG_WARN("%s: F_NOCACHE on %s failed (%s); reads go through the page cache\n", - __func__, fname.c_str(), strerror(errno)); - direct = false; - } -#else - direct = false; -#endif -#endif - if (fd < 0) { - LLAMA_LOG_WARN("%s: direct open of %s failed (%s); falling back to buffered reads\n", - __func__, fname.c_str(), strerror(errno)); - direct = false; - } + llama_bounded_file::params fp; + fp.direct_io = p.direct_io; + file = std::make_unique(fname, fp); + direct = file->direct_io(); + if (rs == 0 || (uint64_t) nrows > UINT64_MAX / rs) { + throw std::runtime_error("llama_ple_disk: table extent overflow"); } - if (fd < 0) { - fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); - if (fd < 0) { - throw std::runtime_error(format("llama_ple_disk: failed to open %s: %s", fname.c_str(), strerror(errno))); - } + const uint64_t bytes = (uint64_t) nrows * rs; + if (offs > file->size() || bytes > file->size() - offs) { + throw std::runtime_error("llama_ple_disk: table extent is outside the file"); } n_threads = std::max(1, p.n_threads); @@ -111,7 +91,7 @@ struct llama_ple_disk::impl { if (p.cache_bytes >= rs) { size_t n = p.cache_bytes / rs; n_slots = 1; - while (n_slots * 2 <= n) { + while (n_slots <= n / 2) { n_slots *= 2; } mask = n_slots - 1; @@ -131,64 +111,17 @@ struct llama_ple_disk::impl { for (auto & w : workers) { w.join(); } - free(bounce0); - if (fd >= 0) { - close(fd); - } #endif } #if !defined(_WIN32) - size_t bounce_size() const { - return ((rs + block - 1) / block) * block + 2 * block; - } - - uint8_t * alloc_bounce() const { - void * ptr = nullptr; - if (posix_memalign(&ptr, block, bounce_size()) != 0) { - throw std::runtime_error("llama_ple_disk: posix_memalign failed"); - } - return (uint8_t *) ptr; - } - - // read `len` bytes at `off` into `dst`; a short read is only tolerated past `need` - void pread_full(uint8_t * dst, size_t len, off_t off, size_t need) const { - size_t got = 0; - while (got < len) { - const ssize_t r = pread(fd, dst + got, len - got, off + (off_t) got); - if (r < 0) { - if (errno == EINTR) { - continue; - } - GGML_ABORT("llama_ple_disk: pread(%s, %zu @ %lld) failed: %s", - fname.c_str(), len, (long long) off, strerror(errno)); - } - if (r == 0) { - break; // EOF - } - got += (size_t) r; - } - if (got < need) { - GGML_ABORT("llama_ple_disk: short read in %s: %zu of %zu bytes at %lld", - fname.c_str(), got, need, (long long) off); - } - } - - void read_row(int64_t row, uint8_t * dst, uint8_t * bounce) const { - const off_t off = (off_t) offs + (off_t) row * (off_t) rs; - if (!direct) { - pread_full(dst, rs, off, rs); - return; - } - const off_t a0 = off & ~(off_t) (block - 1); - const size_t need = (size_t) (off - a0) + rs; - const size_t len = ((need + block - 1) / block) * block; - pread_full(bounce, len, a0, need); - memcpy(dst, bounce + (off - a0), rs); + void read_row(int64_t row, uint8_t * dst, llama_bounded_file::buffer & bounce) const { + const uint64_t off = (uint64_t) offs + (uint64_t) row * rs; + file->read(off, dst, rs, bounce); } void worker() { - uint8_t * bounce = direct ? alloc_bounce() : nullptr; + llama_bounded_file::buffer bounce; uint64_t seen = 0; for (;;) { { @@ -199,27 +132,37 @@ struct llama_ple_disk::impl { } seen = gen; } - for (;;) { - const size_t i = next.fetch_add(1); - if (i >= misses.size()) { - break; + std::exception_ptr error; + try { + if (direct && bounce.empty()) { + bounce = file->make_buffer(rs); } - read_row(misses[i].first, raw.data() + (size_t) misses[i].second * rs, bounce); + for (;;) { + const size_t i = next.fetch_add(1); + if (i >= misses.size()) { + break; + } + read_row(misses[i].first, raw.data() + (size_t) misses[i].second * rs, bounce); + } + } catch (...) { + error = std::current_exception(); } { std::lock_guard lk(pm); + if (error && !worker_error) { + worker_error = error; + } if (--pending == 0) { cv_done.notify_one(); } } } - free(bounce); } void run_misses() { if (n_threads <= 1 || misses.size() <= 2) { - if (direct && bounce0 == nullptr) { - bounce0 = alloc_bounce(); + if (direct && bounce0.empty()) { + bounce0 = file->make_buffer(rs); } for (const auto & m : misses) { read_row(m.first, raw.data() + (size_t) m.second * rs, bounce0); @@ -236,11 +179,17 @@ struct llama_ple_disk::impl { std::lock_guard lk(pm); next = 0; pending = workers.size(); + worker_error = nullptr; ++gen; } cv_work.notify_all(); std::unique_lock lk(pm); cv_done.wait(lk, [&] { return pending == 0; }); + const std::exception_ptr error = worker_error; + lk.unlock(); + if (error) { + std::rethrow_exception(error); + } } void gather(const int32_t * idx, size_t n, float * dst) { @@ -295,7 +244,9 @@ struct llama_ple_disk::impl { st_uniq += uniq.size(); st_hits += uniq.size() - misses.size(); st_reads += misses.size(); - st_bytes += misses.size() * (direct ? block : rs); + for (const auto & miss : misses) { + st_bytes += file->read_size((uint64_t) offs + (uint64_t) miss.first * rs, rs); + } st_ms += std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); } #endif diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5d73a0984bc7..345ac26c6a6e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -197,6 +197,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # llama_build_and_test(test-double-float.cpp) # SLOW llama_build_and_test(test-deepseek41-schema.cpp) + llama_build_and_test(test-engram.cpp) llama_build(test-llama-archs.cpp) set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") diff --git a/tests/test-engram.cpp b/tests/test-engram.cpp new file mode 100644 index 000000000000..e741f231262b --- /dev/null +++ b/tests/test-engram.cpp @@ -0,0 +1,444 @@ +#include "../src/llama-bounded-file.h" +#include "../src/llama-engram.h" +#include "../src/llama-ple-disk.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#endif + +static void check(bool condition, const char * message) { + if (!condition) { + std::fprintf(stderr, "%s\n", message); + std::exit(1); + } +} + +static void expect_invalid(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::invalid_argument &) { + return; + } + check(false, message); +} + +static void expect_domain(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::domain_error &) { + return; + } + check(false, message); +} + +static void expect_runtime(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::runtime_error &) { + return; + } + check(false, message); +} + +static llama_engram_layout make_layout() { + llama_engram_layout layout; + layout.encoding = "e4m3_e8m0_32_row264"; + layout.layer_ids = { 1, 14 }; + layout.token_map.resize(256); + for (size_t i = 0; i < layout.token_map.size(); ++i) { + layout.token_map[i] = (uint32_t) i / 2; + } + layout.compressed_vocab_size = 128; + layout.pad_id = 1; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t depth = 0; depth < LLAMA_ENGRAM_NGRAM; ++depth) { + layout.multipliers[layer][depth] = 35184372088831ull - 2 * (depth + 4 * layer); + } + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + layout.primes[layer][col] = 16000057; + layout.rows[layer] += layout.primes[layer][col]; + } + } + return layout; +} + +static void reference_hash( + const llama_engram_layout & layout, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * output) { + for (size_t pos = 0; pos < count; ++pos) { + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + uint64_t offset = 0; + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + uint64_t hash = 0; + bool blocked = false; + for (size_t shift = 0; shift < col / LLAMA_ENGRAM_HEADS + 2; ++shift) { + const bool before_start = shift > pos; + const size_t source = before_start ? 0 : pos - shift; + blocked = blocked || before_start || (mask != nullptr && mask[source] == 0); + const uint32_t id = blocked ? layout.pad_id : layout.token_map[tokens[source]]; + hash ^= (uint64_t) id * layout.multipliers[layer][shift]; + } + *output++ = (uint32_t) (hash % layout.primes[layer][col] + offset); + offset += layout.primes[layer][col]; + } + } + } +} + +static void test_layout_validation() { + llama_engram_layout layout = make_layout(); + llama_engram_hasher valid(layout); + check(valid.layout().rows == layout.rows, "valid Engram layout changed"); + + llama_engram_layout bad = layout; + bad.encoding = "e4m3"; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad encoding was accepted"); + bad = layout; + bad.layer_ids[1] = bad.layer_ids[0]; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "duplicate layer IDs were accepted"); + bad = layout; + bad.token_map.clear(); + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "empty token map was accepted"); + bad = layout; + bad.compressed_vocab_size = 0; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "empty compressed vocabulary was accepted"); + bad = layout; + bad.pad_id = bad.compressed_vocab_size; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad pad ID was accepted"); + bad = layout; + bad.token_map[10] = bad.compressed_vocab_size; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad token map entry was accepted"); + bad = layout; + bad.multipliers[0][0]--; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "even multiplier was accepted"); + bad = layout; + bad.multipliers[0][0] = UINT64_MAX; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "overflowing multiplier was accepted"); + bad = layout; + bad.primes[0][0] = 1; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad prime was accepted"); + bad = layout; + bad.rows[0]--; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad row extent was accepted"); + bad = layout; + bad.primes[0].fill(UINT32_MAX); + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "overflowing row extent was accepted"); +} + +static void test_hash() { + const llama_engram_layout layout = make_layout(); + const llama_engram_hasher hasher(layout); + constexpr size_t count = 513; + constexpr size_t width = LLAMA_ENGRAM_LAYERS * LLAMA_ENGRAM_COLS; + std::vector tokens(count); + std::vector mask(count); + std::vector expected(count * width); + std::vector actual(count * width); + for (size_t i = 0; i < count; ++i) { + tokens[i] = (int32_t) ((i * 97 + i / 3) % layout.token_map.size()); + mask[i] = i % 17 != 0 && (i < 125 || i > 131); + } + + for (int masked = 0; masked < 2; ++masked) { + const uint8_t * current_mask = masked != 0 ? mask.data() : nullptr; + reference_hash(layout, tokens.data(), current_mask, count, expected.data()); + llama_engram_history expected_history; + expected_history.reset(); + hasher.hash(expected_history, tokens.data(), current_mask, count, actual.data()); + check(actual == expected, "single-chunk Engram hash differs from full-history reference"); + for (size_t chunk = 1; chunk <= count; ++chunk) { + llama_engram_history history; + history.reset(); + for (size_t pos = 0; pos < count; pos += chunk) { + const size_t size = std::min(chunk, count - pos); + hasher.hash(history, tokens.data() + pos, + current_mask != nullptr ? current_mask + pos : nullptr, + size, actual.data() + pos * width); + } + check(actual == expected, "rolling Engram hash differs from full-history reference"); + check(history.tail == expected_history.tail, "chunking changed final Engram history"); + } + } + + llama_engram_history history; + history.reset(); + const llama_engram_history before = history; + uint32_t output[2 * width]; + std::fill(output, output + 2 * width, UINT32_MAX); + const int32_t invalid_high[] = { 0, (int32_t) layout.token_map.size() }; + expect_invalid([&] { hasher.hash(history, invalid_high, nullptr, 2, output); }, + "high invalid token was accepted"); + check(history.tail == before.tail, "invalid token mutated Engram history"); + check(output[0] == UINT32_MAX, "invalid token changed Engram output"); + const int32_t invalid_low[] = { 0, -1 }; + expect_invalid([&] { hasher.hash(history, invalid_low, nullptr, 2, output); }, + "negative token was accepted"); + check(history.tail == before.tail, "negative token mutated Engram history"); + hasher.hash(history, nullptr, nullptr, 0, nullptr); + expect_invalid([&] { hasher.hash(history, invalid_high, nullptr, SIZE_MAX, output); }, + "overflowing hash count was accepted"); + history.tail[0] = (int32_t) layout.compressed_vocab_size; + expect_invalid([&] { hasher.hash(history, invalid_high, nullptr, 1, output); }, + "invalid history was accepted"); +} + +static float decode_reference(uint8_t code, uint8_t scale) { + const int exponent = (code >> 3) & 15; + double value = exponent != 0 ? + (1.0 + (code & 7) / 8.0) * std::pow(2.0, exponent - 7) : + (code & 7) / 512.0; + if ((code & 128) != 0) { + value = -value; + } + float result = (float) (value * std::pow(2.0, (int) scale - 127)); + uint32_t bits; + memcpy(&bits, &result, sizeof(bits)); + bits = (bits + 0x7fffu + ((bits >> 16) & 1u)) & 0xffff0000u; + memcpy(&result, &bits, sizeof(result)); + return result; +} + +static void test_decode() { + uint8_t row[LLAMA_ENGRAM_ROW_BYTES]; + float output[LLAMA_ENGRAM_DIM]; + for (uint32_t code = 0; code < 256; ++code) { + memset(row, (int) code, LLAMA_ENGRAM_DIM); + for (uint32_t scale = 0; scale < 256; ++scale) { + memset(row + LLAMA_ENGRAM_DIM, (int) scale, LLAMA_ENGRAM_ROW_BYTES - LLAMA_ENGRAM_DIM); + const float expected = decode_reference((uint8_t) code, (uint8_t) scale); + const bool valid = (code & 127) != 127 && scale != 255 && std::isfinite(expected); + if (valid) { + llama_engram_decode_row(row, output); + for (float value : output) { + check(memcmp(&value, &expected, sizeof(value)) == 0, "Engram decode differs from BF16 reference"); + } + } else { + output[0] = 123456.0f; + expect_domain([&] { llama_engram_decode_row(row, output); }, "invalid Engram value was accepted"); + check(output[0] == 123456.0f, "invalid Engram value changed output"); + } + } + } + + memset(row, 0, sizeof(row)); + memset(row + LLAMA_ENGRAM_DIM, 127, LLAMA_ENGRAM_ROW_BYTES - LLAMA_ENGRAM_DIM); + row[0] = 128; + llama_engram_decode_row(row, output); + check(output[0] == 0.0f && std::signbit(output[0]), "negative zero was not preserved"); + check(output[1] == 0.0f && !std::signbit(output[1]), "positive zero was not preserved"); +} + +#if !defined(_WIN32) +static void write_full(int fd, const void * data, size_t size, uint64_t offset) { + const ssize_t written = pwrite(fd, data, size, (off_t) offset); + check(written == (ssize_t) size, "failed to write Engram test row"); +} + +static void test_disk_rows() { + char path[] = "/tmp/llama-engram-XXXXXX"; + const int fd = mkstemp(path); + check(fd >= 0, "failed to create Engram test file"); + + const uint64_t offset = (1ull << 33) + 32; + uint8_t raw[3][LLAMA_ENGRAM_ROW_BYTES]; + for (size_t row = 0; row < 3; ++row) { + for (size_t i = 0; i < LLAMA_ENGRAM_DIM; ++i) { + raw[row][i] = (uint8_t) i; + } + raw[row][127] = 0; + raw[row][255] = 128; + for (size_t i = 0; i < LLAMA_ENGRAM_ROW_BYTES - LLAMA_ENGRAM_DIM; ++i) { + raw[row][LLAMA_ENGRAM_DIM + i] = (uint8_t) (126 + row); + } + } + write_full(fd, raw, sizeof(raw), offset); + + llama_engram_table table(path, offset, 3); + check(table.n_rows() == 3, "Engram table row count changed"); + const uint32_t rows[] = { 2, 0, 2, 1 }; + float output[4 * LLAMA_ENGRAM_DIM]; + table.read(rows, 4, output); + for (size_t row = 0; row < 4; ++row) { + float expected[LLAMA_ENGRAM_DIM]; + llama_engram_decode_row(raw[rows[row]], expected); + check(memcmp(output + row * LLAMA_ENGRAM_DIM, expected, sizeof(expected)) == 0, + "Engram row order was not preserved"); + } + + constexpr size_t tokens = 2051; + constexpr size_t stride = LLAMA_ENGRAM_LAYERS * LLAMA_ENGRAM_COLS; + constexpr size_t width = LLAMA_ENGRAM_COLS * LLAMA_ENGRAM_DIM; + std::vector batch_rows(tokens * stride, UINT32_MAX); + std::vector batch(tokens * width + 1); + for (size_t token = 0; token < tokens; ++token) { + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + batch_rows[token * stride + col] = (uint32_t) ((token * 7 + col * 11) % 3); + } + } + const size_t sizes[] = { 1, 2, 31, 65, 257, 2047, 2048, 2049, tokens }; + float decoded[3][LLAMA_ENGRAM_DIM]; + for (size_t row = 0; row < 3; ++row) { + llama_engram_decode_row(raw[row], decoded[row]); + } + for (size_t count : sizes) { + batch[count * width] = 123456.0f; + table.read_batch(batch_rows.data(), count, stride, batch.data()); + for (size_t token = 0; token < count; ++token) { + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + const uint32_t row = batch_rows[token * stride + col]; + check(memcmp(batch.data() + token * width + col * LLAMA_ENGRAM_DIM, + decoded[row], sizeof(decoded[row])) == 0, + "batched Engram row differs from direct decode"); + } + } + check(batch[count * width] == 123456.0f, "batched Engram read overflowed output"); + } + + table.read(nullptr, 0, nullptr); + table.read_batch(nullptr, 0, 0, nullptr); + expect_invalid([&] { table.read_batch(batch_rows.data(), 1, LLAMA_ENGRAM_COLS - 1, batch.data()); }, + "short Engram batch stride was accepted"); + expect_invalid([&] { table.read_batch(batch_rows.data(), 2, SIZE_MAX, batch.data()); }, + "overflowing Engram batch stride was accepted"); + expect_invalid([&] { table.read_batch(batch_rows.data(), SIZE_MAX, stride, batch.data()); }, + "overflowing Engram batch count was accepted"); + + const uint32_t bad_row = 3; + output[0] = 123456.0f; + expect_invalid([&] { table.read(&bad_row, 1, output); }, "invalid Engram row was accepted"); + check(output[0] == 123456.0f, "invalid Engram row changed output"); + batch_rows[(tokens - 1) * stride] = bad_row; + batch[0] = 123456.0f; + expect_invalid([&] { table.read_batch(batch_rows.data(), tokens, stride, batch.data()); }, + "invalid batched Engram row was accepted"); + check(batch[0] == 123456.0f, "invalid batched Engram row changed output"); + batch_rows[(tokens - 1) * stride] = 0; + + uint8_t invalid = 127; + write_full(fd, &invalid, 1, offset); + const uint32_t first_row = 0; + expect_domain([&] { table.read(&first_row, 1, output); }, "E4M3 NaN row was accepted"); + write_full(fd, raw, sizeof(raw), offset); + invalid = 255; + write_full(fd, &invalid, 1, offset + LLAMA_ENGRAM_DIM); + expect_domain([&] { table.read(&first_row, 1, output); }, "E8M0 scale 255 row was accepted"); + write_full(fd, raw, sizeof(raw), offset); + + check(ftruncate(fd, (off_t) (offset + 260)) == 0, "failed to truncate Engram test file"); + expect_runtime([&] { table.read(&first_row, 1, output); }, "short Engram read was accepted"); + close(fd); + + expect_invalid([&] { llama_engram_table invalid_table(path, offset, 1); }, + "truncated Engram extent was accepted"); + expect_invalid([&] { llama_engram_table invalid_table(path, UINT64_MAX - 1, 3); }, + "overflowing Engram extent was accepted"); + expect_invalid([&] { llama_engram_table invalid_table(path, offset, 0); }, + "empty Engram table was accepted"); + check(unlink(path) == 0, "failed to remove Engram test file"); +} + +#if defined(__linux__) +static void test_direct_tail_read() { + char path[] = "/tmp/llama-direct-tail-XXXXXX"; + const int fd = mkstemp(path); + check(fd >= 0, "failed to create direct tail test file"); + + constexpr uint64_t offset = 4096; + const uint8_t expected[] = { 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2 }; + write_full(fd, expected, sizeof(expected), offset); + close(fd); + + llama_bounded_file::params params; + params.direct_io = true; + params.direct_io_required = true; + llama_bounded_file file(path, params); + check(file.direct_io(), "direct tail test did not use O_DIRECT"); + llama_bounded_file::buffer scratch = file.make_buffer(sizeof(expected)); + uint8_t actual[sizeof(expected)] = {}; + file.read(offset, actual, sizeof(actual), scratch); + check(memcmp(actual, expected, sizeof(actual)) == 0, "valid direct tail read failed"); + check(unlink(path) == 0, "failed to remove direct tail test file"); +} +#endif + +static void test_ple_disk_reader() { + char path[] = "/tmp/llama-ple-reader-XXXXXX"; + const int fd = mkstemp(path); + check(fd >= 0, "failed to create PLE test file"); + + constexpr uint64_t offset = 32; + constexpr size_t columns = 4; + constexpr size_t rows = 3; + const float table[rows][columns] = { + { 1.0f, 2.0f, 3.0f, 4.0f }, + { -1.0f, -2.0f, -3.0f, -4.0f }, + { 0.5f, 0.25f, 0.125f, 0.0625f }, + }; + write_full(fd, table, sizeof(table), offset); + close(fd); + + for (int direct = 0; direct < 2; ++direct) { + llama_ple_disk::params params; + params.n_threads = 2; + params.cache_bytes = sizeof(table); + params.direct_io = direct != 0; + llama_ple_disk disk(path, offset, GGML_TYPE_F32, columns, rows, params); + const int32_t ids[] = { 2, 0, 2, 1 }; + float output[4][columns]; + disk.gather(ids, 4, output[0]); + for (size_t i = 0; i < 4; ++i) { + check(memcmp(output[i], table[ids[i]], sizeof(output[i])) == 0, + "shared bounded reader changed PLE row output"); + } + } + + llama_ple_disk::params params; + params.n_threads = 4; + params.cache_bytes = 0; + params.direct_io = false; + llama_ple_disk disk(path, offset, GGML_TYPE_F32, columns, rows, params); + check(truncate(path, (off_t) (offset + sizeof(float))) == 0, "failed to truncate PLE test file"); + const int32_t ids[] = { 0, 1, 2 }; + float output[3][columns]; + expect_runtime([&] { disk.gather(ids, 3, output[0]); }, "threaded PLE read failure did not propagate"); + + const int repair_fd = open(path, O_WRONLY); + check(repair_fd >= 0, "failed to reopen PLE test file"); + write_full(repair_fd, table, sizeof(table), offset); + close(repair_fd); + disk.gather(ids, 3, output[0]); + check(memcmp(output, table, sizeof(table)) == 0, "PLE reader did not recover after worker failure"); + + check(unlink(path) == 0, "failed to remove PLE test file"); +} +#endif + +int main() { + test_layout_validation(); + test_hash(); + test_decode(); +#if !defined(_WIN32) + test_disk_rows(); +#if defined(__linux__) + test_direct_tail_read(); +#endif + test_ple_disk_reader(); +#endif + std::puts("Engram layout, hash and bounded disk rows: PASS"); + return 0; +}