From 74c5167c7f8107404364c9161158652a1eadb1d6 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:15:59 -0700 Subject: [PATCH 1/5] engram : add bounded disk table and hash core Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/CMakeLists.txt | 2 + src/llama-bounded-file.cpp | 237 ++++++++++++++++++++++ src/llama-bounded-file.h | 44 ++++ src/llama-engram.cpp | 305 ++++++++++++++++++++++++++++ src/llama-engram.h | 73 +++++++ src/llama-ple-disk.cpp | 118 +++-------- tests/CMakeLists.txt | 1 + tests/test-engram.cpp | 398 +++++++++++++++++++++++++++++++++++++ 8 files changed, 1085 insertions(+), 93 deletions(-) create mode 100644 src/llama-bounded-file.cpp create mode 100644 src/llama-bounded-file.h create mode 100644 src/llama-engram.cpp create mode 100644 src/llama-engram.h create mode 100644 tests/test-engram.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6094ca17339a..94b133686abe 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -35,6 +35,8 @@ add_library(llama 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..6175255dbc47 --- /dev/null +++ b/src/llama-bounded-file.cpp @@ -0,0 +1,237 @@ +#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 ssize_t n = pread(fd, (uint8_t *) dst + done, len - done, (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 (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..5f3d9847ad62 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 @@ -12,23 +13,16 @@ #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 +34,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 @@ -63,6 +57,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 +72,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 +89,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 +109,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 = file->make_buffer(rs); uint64_t seen = 0; for (;;) { { @@ -213,13 +144,12 @@ struct llama_ple_disk::impl { } } } - 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); @@ -295,7 +225,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 ec74cf2d3492..f4125364431e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -198,6 +198,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-deepseek41-schema.cpp) llama_build_and_test(test-deepseek41-runtime.cpp) + llama_build_and_test(test-engram.cpp) llama_build_and_test(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..f5ad1a7971a6 --- /dev/null +++ b/tests/test-engram.cpp @@ -0,0 +1,398 @@ +#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"); +} + +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"); + } + } + 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(); + test_ple_disk_reader(); +#endif + std::puts("Engram layout, hash and bounded disk rows: PASS"); + return 0; +} From 8a59f344389368b591a11fa717bca3a7cdf329bf Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:27:03 -0700 Subject: [PATCH 2/5] engram : fix bounded read failure handling Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-bounded-file.cpp | 9 +++++++- src/llama-ple-disk.cpp | 31 ++++++++++++++++++++----- tests/test-engram.cpp | 46 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 79 insertions(+), 7 deletions(-) diff --git a/src/llama-bounded-file.cpp b/src/llama-bounded-file.cpp index 6175255dbc47..3ae6d5c31a18 100644 --- a/src/llama-bounded-file.cpp +++ b/src/llama-bounded-file.cpp @@ -134,7 +134,8 @@ struct llama_bounded_file::impl { void pread_full(void * dst, size_t len, uint64_t offset, size_t need) const { size_t done = 0; while (done < len) { - const ssize_t n = pread(fd, (uint8_t *) dst + done, len - done, (off_t) (offset + done)); + 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; @@ -146,6 +147,12 @@ struct llama_bounded_file::impl { 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", diff --git a/src/llama-ple-disk.cpp b/src/llama-ple-disk.cpp index 5f3d9847ad62..45f38a82c603 100644 --- a/src/llama-ple-disk.cpp +++ b/src/llama-ple-disk.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -47,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; @@ -119,7 +121,7 @@ struct llama_ple_disk::impl { } void worker() { - llama_bounded_file::buffer bounce = file->make_buffer(rs); + llama_bounded_file::buffer bounce; uint64_t seen = 0; for (;;) { { @@ -130,15 +132,26 @@ 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); + } + 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); } - 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(); } @@ -166,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) { diff --git a/tests/test-engram.cpp b/tests/test-engram.cpp index f5ad1a7971a6..e741f231262b 100644 --- a/tests/test-engram.cpp +++ b/tests/test-engram.cpp @@ -1,3 +1,4 @@ +#include "../src/llama-bounded-file.h" #include "../src/llama-engram.h" #include "../src/llama-ple-disk.h" @@ -351,6 +352,30 @@ static void test_disk_rows() { 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); @@ -381,6 +406,24 @@ static void test_ple_disk_reader() { "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 @@ -391,6 +434,9 @@ int main() { 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"); From 4a214505044d8e0a8d136a8f1737bab727be6814 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 09:04:45 -0700 Subject: [PATCH 3/5] deepseek41 : add disk-backed Engram runtime Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/CMakeLists.txt | 1 + src/llama-dsv41-engram.cpp | 356 +++++++++++++++++++++++++ src/llama-dsv41-engram.h | 111 ++++++++ src/llama-dsv41.cpp | 31 ++- src/llama-dsv41.h | 6 + src/llama-model-loader.cpp | 1 + src/llama-model.h | 4 + src/models/deepseek41.cpp | 55 +++- src/models/models.h | 4 + tests/CMakeLists.txt | 1 + tests/test-deepseek41-engram.cpp | 414 ++++++++++++++++++++++++++++++ tests/test-deepseek41-runtime.cpp | 1 - 12 files changed, 971 insertions(+), 14 deletions(-) create mode 100644 src/llama-dsv41-engram.cpp create mode 100644 src/llama-dsv41-engram.h create mode 100644 tests/test-deepseek41-engram.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 94b133686abe..06864dfb679f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -18,6 +18,7 @@ add_library(llama llama-context.cpp llama-cparams.cpp llama-dsv41.cpp + llama-dsv41-engram.cpp llama-grammar.cpp llama-graph.cpp llama-hparams.cpp diff --git a/src/llama-dsv41-engram.cpp b/src/llama-dsv41-engram.cpp new file mode 100644 index 000000000000..578177b7a2d7 --- /dev/null +++ b/src/llama-dsv41-engram.cpp @@ -0,0 +1,356 @@ +#include "llama-dsv41-engram.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include + +static llama_dsv41_engram_sequence_state dsv41_engram_initial_state() { + llama_dsv41_engram_sequence_state state; + state.history.reset(); + return state; +} + +void llama_dsv41_validate_engram_extent(const llama_dsv41_engram_extent & extent) { + if (extent.fname.empty() || extent.rows == 0 || extent.columns != LLAMA_ENGRAM_ROW_BYTES || + extent.row_count != extent.rows || extent.type != GGML_TYPE_I8) { + throw std::invalid_argument("DeepSeek V4.1 Engram tensor must be I8 [264, rows]"); + } + const uint64_t bytes = (uint64_t) extent.rows*LLAMA_ENGRAM_ROW_BYTES; + if (extent.offset > (uint64_t) INT64_MAX || bytes > (uint64_t) INT64_MAX - extent.offset) { + throw std::invalid_argument("DeepSeek V4.1 Engram tensor extent overflows"); + } +} + +struct llama_dsv41_engram_transaction::impl { + const llama_dsv41_engram_runtime * owner = nullptr; + uint64_t generation = 0; + bool active = false; + size_t count = 0; + std::vector ids; + std::array, LLAMA_ENGRAM_LAYERS> decoded; + std::vector mask; + std::vector mask_f32; + std::map next; +}; + +llama_dsv41_engram_transaction::llama_dsv41_engram_transaction() : pimpl(std::make_unique()) {} +llama_dsv41_engram_transaction::~llama_dsv41_engram_transaction() = default; +llama_dsv41_engram_transaction::llama_dsv41_engram_transaction(llama_dsv41_engram_transaction && other) noexcept = default; +llama_dsv41_engram_transaction & llama_dsv41_engram_transaction::operator=(llama_dsv41_engram_transaction && other) noexcept = default; + +size_t llama_dsv41_engram_transaction::token_count() const { + return pimpl->count; +} + +const uint32_t * llama_dsv41_engram_transaction::row_ids(uint32_t layer) const { + if (layer >= LLAMA_ENGRAM_LAYERS || pimpl->count == 0) { + return nullptr; + } + return pimpl->ids.data() + layer*LLAMA_ENGRAM_COLS; +} + +const float * llama_dsv41_engram_transaction::rows(uint32_t layer) const { + if (layer >= LLAMA_ENGRAM_LAYERS || pimpl->decoded[layer].empty()) { + return nullptr; + } + return pimpl->decoded[layer].data(); +} + +const uint8_t * llama_dsv41_engram_transaction::text_mask() const { + return pimpl->mask.empty() ? nullptr : pimpl->mask.data(); +} + +void llama_dsv41_engram_transaction::upload_layer( + uint32_t layer, + size_t token_offset, + size_t token_count, + ggml_tensor * rows_input, + ggml_tensor * text_mask_input) const { + if (!pimpl->active) { + throw std::invalid_argument("DeepSeek V4.1 Engram transaction is not active"); + } + if (layer >= LLAMA_ENGRAM_LAYERS || token_offset > pimpl->count || + token_count > pimpl->count - token_offset) { + throw std::invalid_argument("DeepSeek V4.1 Engram upload range is invalid"); + } + const int64_t row_width = LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM; + if (rows_input == nullptr || text_mask_input == nullptr || + rows_input->type != GGML_TYPE_F32 || rows_input->ne[0] != row_width || + ggml_nelements(rows_input) != row_width*(int64_t) token_count || + text_mask_input->type != GGML_TYPE_F32 || + ggml_nelements(text_mask_input) != (int64_t) token_count) { + throw std::invalid_argument("DeepSeek V4.1 Engram input tensor shape mismatch"); + } + + const size_t row_offset = token_offset*row_width; + ggml_backend_tensor_set( + rows_input, + pimpl->decoded[layer].data() + row_offset, + 0, + token_count*row_width*sizeof(float)); + ggml_backend_tensor_set( + text_mask_input, + pimpl->mask_f32.data() + token_offset, + 0, + token_count*sizeof(float)); +} + +struct llama_dsv41_engram_runtime::impl { + llama_engram_hasher hasher; + std::array, LLAMA_ENGRAM_LAYERS> tables; + std::map sequences; + size_t max_tokens; + uint64_t generation = 0; + + impl( + llama_engram_layout layout, + const std::array & extents, + size_t max_tokens) : + hasher(std::move(layout)), + max_tokens(max_tokens) { + if (max_tokens == 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram token bound must be non-zero"); + } + for (size_t i = 0; i < LLAMA_ENGRAM_LAYERS; ++i) { + llama_dsv41_validate_engram_extent(extents[i]); + if (extents[i].rows != hasher.layout().rows[i]) { + throw std::invalid_argument("DeepSeek V4.1 Engram tensor rows do not match metadata"); + } + tables[i] = std::make_unique( + extents[i].fname, extents[i].offset, extents[i].rows); + } + } +}; + +llama_dsv41_engram_runtime::llama_dsv41_engram_runtime( + llama_engram_layout layout, + const std::array & extents, + size_t max_tokens) : + pimpl(std::make_unique(std::move(layout), extents, max_tokens)) {} + +llama_dsv41_engram_runtime::~llama_dsv41_engram_runtime() = default; + +llama_dsv41_engram_transaction llama_dsv41_engram_runtime::prepare( + const std::vector & tokens) { + if (tokens.size() > pimpl->max_tokens) { + throw std::invalid_argument("DeepSeek V4.1 Engram batch exceeds the configured bound"); + } + + llama_dsv41_engram_transaction result; + result.pimpl->owner = this; + result.pimpl->generation = pimpl->generation; + result.pimpl->count = tokens.size(); + result.pimpl->ids.resize(tokens.size()*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS); + result.pimpl->mask.resize(tokens.size()); + result.pimpl->mask_f32.resize(tokens.size()); + result.pimpl->next = pimpl->sequences; + + for (size_t i = 0; i < tokens.size(); ++i) { + const llama_dsv41_engram_token & token = tokens[i]; + if (token.seq_ids.empty()) { + throw std::invalid_argument("DeepSeek V4.1 Engram token has no sequence"); + } + + uint32_t expected[LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS] = {}; + bool have_expected = false; + std::set seen; + for (llama_seq_id seq_id : token.seq_ids) { + if (seq_id < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + if (!seen.insert(seq_id).second) { + throw std::invalid_argument("DeepSeek V4.1 Engram token repeats a sequence ID"); + } + auto inserted = result.pimpl->next.emplace(seq_id, dsv41_engram_initial_state()); + llama_dsv41_engram_sequence_state & state = inserted.first->second; + if (token.pos != state.pos + 1) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence position is not contiguous"); + } + + uint32_t current[LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS] = {}; + pimpl->hasher.hash(state.history, &token.token, &token.text, 1, current); + state.pos = token.pos; + if (have_expected && std::memcmp(expected, current, sizeof(expected)) != 0) { + throw std::invalid_argument("DeepSeek V4.1 coupled sequence histories differ"); + } + std::memcpy(expected, current, sizeof(expected)); + have_expected = true; + } + + std::memcpy( + result.pimpl->ids.data() + i*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + expected, + sizeof(expected)); + result.pimpl->mask[i] = token.text != 0; + result.pimpl->mask_f32[i] = token.text != 0 ? 1.0f : 0.0f; + } + + const size_t row_values = tokens.size()*LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + result.pimpl->decoded[layer].resize(row_values); + pimpl->tables[layer]->read_batch( + result.pimpl->ids.data() + layer*LLAMA_ENGRAM_COLS, + tokens.size(), + LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + result.pimpl->decoded[layer].data()); + } + result.pimpl->active = true; + return result; +} + +void llama_dsv41_engram_runtime::commit(llama_dsv41_engram_transaction & transaction) { + if (!transaction.pimpl->active || transaction.pimpl->owner != this) { + throw std::invalid_argument("DeepSeek V4.1 Engram transaction is not active"); + } + if (transaction.pimpl->generation != pimpl->generation) { + throw std::runtime_error("DeepSeek V4.1 Engram transaction is stale"); + } + auto next = transaction.pimpl->next; + pimpl->sequences.swap(next); + ++pimpl->generation; + transaction.pimpl->active = false; +} + +void llama_dsv41_engram_runtime::rollback(llama_dsv41_engram_transaction & transaction) { + if (transaction.pimpl->owner != this) { + throw std::invalid_argument("DeepSeek V4.1 Engram transaction belongs to another runtime"); + } + transaction.pimpl->active = false; +} + +void llama_dsv41_engram_runtime::seq_reset(llama_seq_id seq_id) { + if (seq_id < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + pimpl->sequences[seq_id] = dsv41_engram_initial_state(); + ++pimpl->generation; +} + +void llama_dsv41_engram_runtime::seq_copy(llama_seq_id seq_id_src, llama_seq_id seq_id_dst) { + if (seq_id_src < 0 || seq_id_dst < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + const auto it = pimpl->sequences.find(seq_id_src); + pimpl->sequences[seq_id_dst] = it == pimpl->sequences.end() ? + dsv41_engram_initial_state() : it->second; + ++pimpl->generation; +} + +void llama_dsv41_engram_runtime::seq_remove(llama_seq_id seq_id) { + if (seq_id < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + pimpl->sequences.erase(seq_id); + ++pimpl->generation; +} + +llama_dsv41_engram_snapshot llama_dsv41_engram_runtime::checkpoint() const { + return { pimpl->sequences }; +} + +void llama_dsv41_engram_runtime::restore(const llama_dsv41_engram_snapshot & snapshot) { + auto restored = snapshot.sequences; + for (auto & item : restored) { + if (item.first < 0 || item.second.pos < -1) { + throw std::invalid_argument("DeepSeek V4.1 Engram snapshot is invalid"); + } + pimpl->hasher.hash(item.second.history, nullptr, nullptr, 0, nullptr); + } + pimpl->sequences.swap(restored); + ++pimpl->generation; +} + +llama_dsv41_engram_sequence_state llama_dsv41_engram_runtime::sequence(llama_seq_id seq_id) const { + const auto it = pimpl->sequences.find(seq_id); + return it == pimpl->sequences.end() ? dsv41_engram_initial_state() : it->second; +} + +size_t llama_dsv41_engram_runtime::max_tokens() const { + return pimpl->max_tokens; +} + +static ggml_tensor * dsv41_bf16_f32(ggml_context * ctx, ggml_tensor * tensor) { + return ggml_cast(ctx, ggml_cast(ctx, tensor, GGML_TYPE_BF16), GGML_TYPE_F32); +} + +ggml_tensor * llama_dsv41_build_engram_add( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * projected, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_mask, + float rms_eps) { + if (ctx == nullptr || residual == nullptr || projected == nullptr || q_norm == nullptr || k_norm == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 Engram graph input is null"); + } + + const int64_t width = residual->ne[0]; + const int64_t streams = residual->ne[1]; + const int64_t tokens = residual->ne[2]; + if (width <= 0 || streams != 4 || tokens <= 0 || + projected->ne[0] != 5*width || projected->ne[1] != tokens || + q_norm->ne[0] != width || q_norm->ne[1] != streams || + k_norm->ne[0] != width || k_norm->ne[1] != streams || + (text_mask != nullptr && (text_mask->ne[0] != 1 || text_mask->ne[1] != tokens))) { + throw std::invalid_argument("DeepSeek V4.1 Engram graph shape mismatch"); + } + + projected = dsv41_bf16_f32(ctx, projected); + ggml_tensor * value = ggml_view_2d( + ctx, projected, width, tokens, projected->nb[1], 4*projected->nb[0]*width); + ggml_tensor * result = nullptr; + for (int64_t stream = 0; stream < streams; ++stream) { + ggml_tensor * hidden = ggml_view_2d( + ctx, residual, width, tokens, residual->nb[2], stream*residual->nb[1]); + ggml_tensor * key = ggml_view_2d( + ctx, projected, width, tokens, projected->nb[1], stream*projected->nb[0]*width); + ggml_tensor * qw = ggml_view_1d(ctx, q_norm, width, stream*q_norm->nb[1]); + ggml_tensor * kw = ggml_view_1d(ctx, k_norm, width, stream*k_norm->nb[1]); + + ggml_tensor * hidden_norm = ggml_rms_norm(ctx, hidden, rms_eps); + ggml_tensor * key_norm = ggml_rms_norm(ctx, key, rms_eps); + ggml_tensor * dot = ggml_mul(ctx, hidden_norm, qw); + dot = ggml_mul(ctx, dot, kw); + dot = ggml_mul(ctx, dot, key_norm); + dot = ggml_scale(ctx, ggml_sum_rows(ctx, dot), 1.0f/std::sqrt((float) width)); + + ggml_tensor * magnitude = ggml_sqrt(ctx, ggml_clamp(ctx, ggml_abs(ctx, dot), 1.0e-6f, INFINITY)); + ggml_tensor * gate = ggml_sigmoid(ctx, ggml_mul(ctx, ggml_sgn(ctx, dot), magnitude)); + if (text_mask != nullptr) { + gate = ggml_mul(ctx, gate, text_mask); + } + + ggml_tensor * updated = ggml_add(ctx, hidden, ggml_mul(ctx, value, gate)); + updated = dsv41_bf16_f32(ctx, updated); + updated = ggml_reshape_3d(ctx, updated, width, 1, tokens); + result = result == nullptr ? updated : ggml_concat(ctx, result, updated, 1); + } + return result; +} + +ggml_tensor * llama_dsv41_build_engram( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * rows, + ggml_tensor * engram_kv, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_mask, + float rms_eps) { + if (ctx == nullptr || residual == nullptr || rows == nullptr || engram_kv == nullptr || + rows->ne[0] != LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM || + engram_kv->ne[0] != rows->ne[0] || + engram_kv->ne[1] != 5*residual->ne[0]) { + throw std::invalid_argument("DeepSeek V4.1 Engram projection shape mismatch"); + } + ggml_tensor * projected = ggml_mul_mat(ctx, engram_kv, rows); + return llama_dsv41_build_engram_add( + ctx, residual, projected, q_norm, k_norm, text_mask, rms_eps); +} diff --git a/src/llama-dsv41-engram.h b/src/llama-dsv41-engram.h new file mode 100644 index 000000000000..db0ae04dc853 --- /dev/null +++ b/src/llama-dsv41-engram.h @@ -0,0 +1,111 @@ +#pragma once + +#include "llama-engram.h" +#include "llama.h" + +#include +#include +#include +#include +#include +#include +#include + +struct ggml_context; +struct ggml_tensor; + +struct llama_dsv41_engram_extent { + std::string fname; + uint64_t offset = 0; + uint32_t rows = 0; + int64_t columns = 0; + int64_t row_count = 0; + int32_t type = 0; +}; + +void llama_dsv41_validate_engram_extent(const llama_dsv41_engram_extent & extent); + +struct llama_dsv41_engram_token { + int32_t token = -1; + llama_pos pos = -1; + std::vector seq_ids; + uint8_t text = 1; +}; + +struct llama_dsv41_engram_sequence_state { + llama_engram_history history = {}; + llama_pos pos = -1; +}; + +struct llama_dsv41_engram_snapshot { + std::map sequences; +}; + +struct llama_dsv41_engram_transaction { + llama_dsv41_engram_transaction(); + ~llama_dsv41_engram_transaction(); + llama_dsv41_engram_transaction(llama_dsv41_engram_transaction && other) noexcept; + llama_dsv41_engram_transaction & operator=(llama_dsv41_engram_transaction && other) noexcept; + + llama_dsv41_engram_transaction(const llama_dsv41_engram_transaction &) = delete; + llama_dsv41_engram_transaction & operator=(const llama_dsv41_engram_transaction &) = delete; + + size_t token_count() const; + const uint32_t * row_ids(uint32_t layer) const; + const float * rows(uint32_t layer) const; + const uint8_t * text_mask() const; + void upload_layer( + uint32_t layer, + size_t token_offset, + size_t token_count, + ggml_tensor * rows_input, + ggml_tensor * text_mask_input) const; + + struct impl; + std::unique_ptr pimpl; +}; + +class llama_dsv41_engram_runtime { +public: + llama_dsv41_engram_runtime( + llama_engram_layout layout, + const std::array & extents, + size_t max_tokens); + ~llama_dsv41_engram_runtime(); + + llama_dsv41_engram_transaction prepare(const std::vector & tokens); + void commit(llama_dsv41_engram_transaction & transaction); + void rollback(llama_dsv41_engram_transaction & transaction); + + void seq_reset(llama_seq_id seq_id); + void seq_copy(llama_seq_id seq_id_src, llama_seq_id seq_id_dst); + void seq_remove(llama_seq_id seq_id); + llama_dsv41_engram_snapshot checkpoint() const; + void restore(const llama_dsv41_engram_snapshot & snapshot); + llama_dsv41_engram_sequence_state sequence(llama_seq_id seq_id) const; + + size_t max_tokens() const; + +private: + struct impl; + std::unique_ptr pimpl; +}; + +ggml_tensor * llama_dsv41_build_engram_add( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * projected, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_mask, + float rms_eps); + +ggml_tensor * llama_dsv41_build_engram( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * rows, + ggml_tensor * engram_kv, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_mask, + float rms_eps); diff --git a/src/llama-dsv41.cpp b/src/llama-dsv41.cpp index 769f7c9096b8..5de36da9da8f 100644 --- a/src/llama-dsv41.cpp +++ b/src/llama-dsv41.cpp @@ -75,10 +75,39 @@ void llama_dsv41_validate_config(const llama_dsv41_config & config) { dsv41_require(config.engram_token_map_size == LLAMA_DSV41_N_VOCAB, "engram.token_map must contain 129280 entries"); dsv41_require(config.engram_primes_size == LLAMA_DSV41_ENGRAM_PRIMES_COUNT, "engram.primes must contain 48 entries"); dsv41_require(config.engram_multipliers_size == LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT, "engram.multipliers must contain 8 entries"); + if (!config.engram_token_map.empty() || !config.engram_primes.empty() || !config.engram_multipliers.empty()) { + llama_dsv41_make_engram_layout(config); + } +} + +llama_engram_layout llama_dsv41_make_engram_layout(const llama_dsv41_config & config) { + dsv41_require(config.engram_token_map.size() == LLAMA_DSV41_N_VOCAB, "engram.token_map data must contain 129280 entries"); + dsv41_require(config.engram_primes.size() == LLAMA_DSV41_ENGRAM_PRIMES_COUNT, "engram.primes data must contain 48 entries"); + dsv41_require(config.engram_multipliers.size() == LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT, "engram.multipliers data must contain 8 entries"); + + llama_engram_layout layout; + layout.encoding = config.engram_encoding; + std::copy(config.engram_layers.begin(), config.engram_layers.end(), layout.layer_ids.begin()); + layout.token_map = config.engram_token_map; + layout.compressed_vocab_size = config.engram_compressed_vocab_size; + layout.pad_id = config.engram_pad_id; + std::copy(config.engram_rows.begin(), config.engram_rows.end(), layout.rows.begin()); + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + std::copy_n( + config.engram_primes.begin() + layer*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS, + layout.primes[layer].begin()); + std::copy_n( + config.engram_multipliers.begin() + layer*LLAMA_ENGRAM_NGRAM, + LLAMA_ENGRAM_NGRAM, + layout.multipliers[layer].begin()); + } + llama_engram_hasher validate(layout); + return layout; } const char * llama_dsv41_runtime_dependency_error() { - return "DeepSeek V4.1 execution requires disk-backed Engram and routed-expert streaming support"; + return "DeepSeek V4.1 execution requires routed-expert streaming support"; } static int32_t dsv41_source_layer(const int32_t * sources, size_t n, uint32_t il) { diff --git a/src/llama-dsv41.h b/src/llama-dsv41.h index 47e6df22c76a..19d73a2f04dd 100644 --- a/src/llama-dsv41.h +++ b/src/llama-dsv41.h @@ -1,7 +1,9 @@ #pragma once +#include "llama-engram.h" #include "llama.h" +#include #include #include #include @@ -92,9 +94,13 @@ struct llama_dsv41_config { uint32_t engram_token_map_size; uint32_t engram_primes_size; uint32_t engram_multipliers_size; + std::vector engram_token_map; + std::vector engram_primes; + std::vector engram_multipliers; }; void llama_dsv41_validate_config(const llama_dsv41_config & config); +llama_engram_layout llama_dsv41_make_engram_layout(const llama_dsv41_config & config); const char * llama_dsv41_runtime_dependency_error(); struct llama_dsv41_compression_plan { diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 2dfcd6eb9074..0c609c2486b0 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -429,6 +429,7 @@ namespace GGUFMeta { template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); diff --git a/src/llama-model.h b/src/llama-model.h index 4c4a30e018bc..f277ed8e14ee 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -557,6 +557,10 @@ struct llama_layer { struct ggml_tensor * indexer_attn_k = nullptr; struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias + struct ggml_tensor * engram_kv = nullptr; + struct ggml_tensor * engram_q_norm = nullptr; + struct ggml_tensor * engram_k_norm = nullptr; + // MSA struct ggml_tensor * index_q_proj = nullptr; struct ggml_tensor * index_k_proj = nullptr; diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 03ac0a6ab375..4312290507c0 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -1,8 +1,10 @@ #include "llama-dsv41.h" +#include "llama-dsv41-engram.h" #include "llama-hparams.h" #include "models.h" #include +#include #include #include #include @@ -12,6 +14,11 @@ static float dsv41_rope_attn_factor(float freq_scale) { return 1.0f/(1.0f + 0.1f*logf(1.0f/freq_scale)); } +struct llama_model_deepseek41::engram_model { + llama_engram_layout layout; + std::array extents; +}; + void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { llama_dsv41_config config = {}; std::string raw_config; @@ -65,9 +72,14 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { ml.get_arr_n(LLM_KV_DSV41_ENGRAM_TOKEN_MAP, config.engram_token_map_size); ml.get_arr_n(LLM_KV_DSV41_ENGRAM_PRIMES, config.engram_primes_size); ml.get_arr_n(LLM_KV_DSV41_ENGRAM_MULTIPLIERS, config.engram_multipliers_size); + ml.get_arr(LLM_KV_DSV41_ENGRAM_TOKEN_MAP, config.engram_token_map); + ml.get_arr(LLM_KV_DSV41_ENGRAM_PRIMES, config.engram_primes); + ml.get_arr(LLM_KV_DSV41_ENGRAM_MULTIPLIERS, config.engram_multipliers); config.n_ff_dense = LLAMA_DSV41_N_FF_DENSE; llama_dsv41_validate_config(config); + engram = std::make_shared(); + engram->layout = llama_dsv41_make_engram_layout(config); if (raw_config.empty()) { throw std::runtime_error("DeepSeek V4.1 metadata: config must not be empty"); @@ -199,23 +211,42 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_exp*n_expert_shared }, 0); if (hparams.dsv41_engram_layers.test(il)) { - const llm_tensor engram_tensors[] = { - LLM_TENSOR_ENGRAM_EMBD, - LLM_TENSOR_ENGRAM_Q_NORM, - LLM_TENSOR_ENGRAM_K_NORM, - LLM_TENSOR_ENGRAM_KV, - }; - for (llm_tensor tensor : engram_tensors) { - const std::string name = tn(tensor, "weight", il).str(); - if (ml.get_weight(name.c_str()) == nullptr) { - throw std::runtime_error("DeepSeek V4.1 is missing required Engram tensor " + name); - } + const size_t index = il == (int32_t) engram->layout.layer_ids[0] ? 0 : 1; + const std::string table_name = tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il).str(); + const auto * table = ml.get_weight(table_name.c_str()); + if (table == nullptr) { + throw std::runtime_error("DeepSeek V4.1 is missing required Engram tensor " + table_name); } + llama_dsv41_engram_extent & extent = engram->extents[index]; + extent.fname = ml.fnames.at(table->idx); + extent.offset = table->offs; + extent.rows = engram->layout.rows[index]; + extent.columns = table->tensor->ne[0]; + extent.row_count = table->tensor->ne[1]; + extent.type = table->tensor->type; + llama_dsv41_validate_engram_extent(extent); + + create_tensor( + tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il), + { LLAMA_ENGRAM_ROW_BYTES, (int64_t) extent.rows }, + TENSOR_SKIP); + layer.engram_q_norm = create_tensor( + tn(LLM_TENSOR_ENGRAM_Q_NORM, "weight", il), + { n_embd, hc_mult }, + 0); + layer.engram_k_norm = create_tensor( + tn(LLM_TENSOR_ENGRAM_K_NORM, "weight", il), + { n_embd, hc_mult }, + 0); + layer.engram_kv = create_tensor( + tn(LLM_TENSOR_ENGRAM_KV, "weight", il), + { LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, (hc_mult + 1)*n_embd }, + 0); } } throw std::runtime_error( - std::string("DeepSeek V4.1 tensor metadata is valid, but tensor payloads cannot be mapped: ") + + std::string("DeepSeek V4.1 Engram metadata and disk extents are valid, but execution is blocked: ") + llama_dsv41_runtime_dependency_error()); } diff --git a/src/models/models.h b/src/models/models.h index 231936009d49..695ff2ccfbb8 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1315,6 +1315,10 @@ struct llama_model_deepseek4 : public llama_model_base { struct llama_model_deepseek41 : public llama_model_deepseek4 { llama_model_deepseek41(const struct llama_model_params & params) : llama_model_deepseek4(params) {} + + struct engram_model; + std::shared_ptr engram; + void load_arch_hparams(llama_model_loader & ml) override; [[noreturn]] void load_arch_tensors(llama_model_loader & ml) override; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index f4125364431e..38fe6e4e2930 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-deepseek41-engram.cpp) llama_build_and_test(test-deepseek41-runtime.cpp) llama_build_and_test(test-engram.cpp) llama_build_and_test(test-llama-archs.cpp) diff --git a/tests/test-deepseek41-engram.cpp b/tests/test-deepseek41-engram.cpp new file mode 100644 index 000000000000..0eaf340a744c --- /dev/null +++ b/tests/test-deepseek41-engram.cpp @@ -0,0 +1,414 @@ +#include "../src/llama-dsv41-engram.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.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_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(32); + for (size_t i = 0; i < layout.token_map.size(); ++i) { + layout.token_map[i] = (uint32_t) i; + } + layout.compressed_vocab_size = 32; + layout.pad_id = 2; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t i = 0; i < LLAMA_ENGRAM_NGRAM; ++i) { + layout.multipliers[layer][i] = 101 + 8*layer + 2*i; + } + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + layout.primes[layer][col] = 2; + layout.rows[layer] += 2; + } + } + return layout; +} + +static void fill_row(uint8_t row[LLAMA_ENGRAM_ROW_BYTES], uint32_t id) { + const uint8_t code = (uint8_t) (8 + id%100); + std::fill(row, row + LLAMA_ENGRAM_DIM, code); + std::fill(row + LLAMA_ENGRAM_DIM, row + LLAMA_ENGRAM_ROW_BYTES, 127); +} + +#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 DeepSeek V4.1 Engram test data"); +} + +struct test_file { + std::string path; + int fd = -1; + std::array extents; + + test_file(const llama_engram_layout & layout) { + char name[] = "/tmp/llama-dsv41-engram-XXXXXX"; + fd = mkstemp(name); + check(fd >= 0, "failed to create DeepSeek V4.1 Engram test file"); + path = name; + + uint64_t offset = 4096; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + extents[layer] = { + path, + offset, + layout.rows[layer], + LLAMA_ENGRAM_ROW_BYTES, + layout.rows[layer], + GGML_TYPE_I8, + }; + for (uint32_t row_id = 0; row_id < layout.rows[layer]; ++row_id) { + uint8_t row[LLAMA_ENGRAM_ROW_BYTES]; + fill_row(row, row_id + 7*layer); + write_full(fd, row, sizeof(row), offset + (uint64_t) row_id*sizeof(row)); + } + offset += (uint64_t) layout.rows[layer]*LLAMA_ENGRAM_ROW_BYTES + 4096; + } + } + + ~test_file() { + if (fd >= 0) { + close(fd); + } + if (!path.empty()) { + unlink(path.c_str()); + } + } +}; + +static void test_extent_validation() { + llama_dsv41_engram_extent extent = { + "/tmp/model.gguf", 4096, 48, LLAMA_ENGRAM_ROW_BYTES, 48, GGML_TYPE_I8, + }; + llama_dsv41_validate_engram_extent(extent); + extent.type = GGML_TYPE_F32; + expect_invalid([&] { llama_dsv41_validate_engram_extent(extent); }, "non-I8 Engram extent was accepted"); + extent.type = GGML_TYPE_I8; + extent.columns = LLAMA_ENGRAM_ROW_BYTES - 1; + expect_invalid([&] { llama_dsv41_validate_engram_extent(extent); }, "short Engram row was accepted"); +} + +static void test_transactions_and_sequences() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime runtime(layout, file.extents, 8); + + std::vector tokens = { + { 3, 0, { 7, 9 }, 1 }, + { 5, 1, { 7, 9 }, 0 }, + { 11, 0, { 12 }, 1 }, + }; + llama_dsv41_engram_transaction transaction = runtime.prepare(tokens); + check(transaction.token_count() == tokens.size(), "Engram transaction token count mismatch"); + check(transaction.row_ids(1) == transaction.row_ids(0) + LLAMA_ENGRAM_COLS, + "Engram layer row selection is not the second 24-ID half"); + llama_engram_hasher hasher(layout); + llama_engram_history expected_history_7; + llama_engram_history expected_history_12; + expected_history_7.reset(); + expected_history_12.reset(); + uint32_t expected[3*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS]; + const int32_t coupled_tokens[] = { 3, 5 }; + const uint8_t coupled_mask[] = { 1, 0 }; + hasher.hash(expected_history_7, coupled_tokens, coupled_mask, 2, expected); + hasher.hash( + expected_history_12, + &tokens[2].token, + &tokens[2].text, + 1, + expected + 2*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS); + for (size_t token = 0; token < tokens.size(); ++token) { + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + check(std::memcmp( + transaction.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + expected + (token*LLAMA_ENGRAM_LAYERS + layer)*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(uint32_t)) == 0, + "DeepSeek V4.1 Engram IDs or 48-ID token stride differ from the core hasher"); + } + } + check(transaction.text_mask()[0] == 1 && transaction.text_mask()[1] == 0, + "Engram text mask changed"); + check(runtime.sequence(7).pos == -1, "Engram prepare committed sequence state early"); + + runtime.commit(transaction); + check(runtime.sequence(7).pos == 1 && runtime.sequence(9).pos == 1, + "Engram commit did not advance coupled sequences"); + check(runtime.sequence(7).history.tail[0] == LLAMA_ENGRAM_DEAD, + "masked Engram token did not break sequence history"); + + runtime.seq_copy(7, 13); + check(runtime.sequence(13).history.tail == runtime.sequence(7).history.tail, + "Engram sequence copy changed history"); + const llama_dsv41_engram_snapshot snapshot = runtime.checkpoint(); + runtime.seq_reset(7); + check(runtime.sequence(7).pos == -1, "Engram sequence reset did not clear position"); + runtime.restore(snapshot); + check(runtime.sequence(7).pos == 1, "Engram checkpoint restore lost position"); + runtime.seq_remove(13); + check(runtime.sequence(13).pos == -1, "Engram sequence remove retained state"); + + llama_dsv41_engram_transaction stale = runtime.prepare({ { 7, 2, { 7 }, 1 } }); + runtime.seq_copy(7, 14); + expect_runtime([&] { runtime.commit(stale); }, "stale Engram transaction was committed"); + runtime.rollback(stale); + + llama_dsv41_engram_snapshot invalid = runtime.checkpoint(); + invalid.sequences[7].history.tail[0] = (int32_t) layout.compressed_vocab_size; + expect_invalid([&] { runtime.restore(invalid); }, "invalid Engram snapshot was restored"); + expect_invalid( + [&] { runtime.prepare({ { 8, 2, { 7, 7 }, 1 } }); }, + "duplicate Engram sequence ID was accepted"); +} + +static void test_scheduler_upload() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime runtime(layout, file.extents, 8); + llama_dsv41_engram_transaction transaction = runtime.prepare({ + { 3, 0, { 0 }, 1 }, + { 5, 1, { 0 }, 0 }, + { 7, 2, { 0 }, 1 }, + }); + + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create DeepSeek V4.1 Engram upload context"); + ggml_tensor * rows = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, transaction.token_count()); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, transaction.token_count()); + + ggml_backend_t backend = ggml_backend_cpu_init(); + check(backend != nullptr, "failed to create DeepSeek V4.1 Engram upload backend"); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + check(buffer != nullptr, "failed to allocate DeepSeek V4.1 Engram upload tensors"); + + std::vector actual_rows(ggml_nelements(rows)); + std::vector actual_mask(ggml_nelements(mask)); + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + transaction.upload_layer(layer, 0, transaction.token_count(), rows, mask); + ggml_backend_tensor_get(rows, actual_rows.data(), 0, ggml_nbytes(rows)); + ggml_backend_tensor_get(mask, actual_mask.data(), 0, ggml_nbytes(mask)); + check(std::memcmp( + actual_rows.data(), + transaction.rows(layer), + ggml_nbytes(rows)) == 0, + "scheduler-backed Engram row upload changed the bounded pack"); + check(actual_mask == std::vector({ 1.0f, 0.0f, 1.0f }), + "scheduler-backed Engram text mask upload changed"); + } + + runtime.commit(transaction); + expect_invalid( + [&] { transaction.upload_layer(0, 0, 3, rows, mask); }, + "committed Engram transaction uploaded stale input"); + + ggml_backend_buffer_free(buffer); + ggml_backend_free(backend); + ggml_free(ctx); +} + +static void test_chunked_prefill() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime whole_runtime(layout, file.extents, 8); + llama_dsv41_engram_runtime chunked_runtime(layout, file.extents, 8); + const std::vector tokens = { + { 3, 0, { 0 }, 1 }, + { 5, 1, { 0 }, 1 }, + { 7, 2, { 0 }, 0 }, + { 9, 3, { 0 }, 1 }, + }; + + llama_dsv41_engram_transaction whole = whole_runtime.prepare(tokens); + llama_dsv41_engram_transaction first = chunked_runtime.prepare({ tokens[0], tokens[1] }); + for (size_t token = 0; token < 2; ++token) { + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + check(std::memcmp( + whole.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + first.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(uint32_t)) == 0, + "first Engram prefill chunk differs from whole-chunk hashing"); + } + } + chunked_runtime.commit(first); + + llama_dsv41_engram_transaction second = chunked_runtime.prepare({ tokens[2], tokens[3] }); + for (size_t token = 0; token < 2; ++token) { + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + check(std::memcmp( + whole.row_ids(layer) + (token + 2)*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + second.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(uint32_t)) == 0, + "later Engram prefill chunk lost committed history"); + } + } +} + +static void test_transactional_read_failure() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime runtime(layout, file.extents, 8); + + llama_dsv41_engram_transaction first = runtime.prepare({ { 1, 0, { 0 }, 1 } }); + runtime.commit(first); + const llama_dsv41_engram_sequence_state before = runtime.sequence(0); + + const uint64_t truncated = file.extents[1].offset + LLAMA_ENGRAM_ROW_BYTES; + check(ftruncate(file.fd, (off_t) truncated) == 0, "failed to truncate Engram transaction test file"); + expect_runtime( + [&] { runtime.prepare({ { 2, 1, { 0 }, 1 } }); }, + "Engram read failure was not surfaced"); + const llama_dsv41_engram_sequence_state after = runtime.sequence(0); + check(after.pos == before.pos && after.history.tail == before.history.tail, + "failed Engram transaction advanced sequence state"); +} +#endif + +static float bf16(float value) { + return ggml_bf16_to_fp32(ggml_fp32_to_bf16(value)); +} + +static void test_graph_gate() { + constexpr int64_t width = 8; + constexpr int64_t streams = 4; + constexpr int64_t tokens = 2; + + ggml_init_params params = { + /*.mem_size =*/ 8*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create DeepSeek V4.1 Engram graph context"); + + ggml_tensor * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, width, streams, tokens); + ggml_tensor * rows = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, tokens); + ggml_tensor * engram_kv = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, 5*width); + ggml_tensor * q_norm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, streams); + ggml_tensor * k_norm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, streams); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, tokens); + + float * residual_data = static_cast(residual->data); + float * rows_data = static_cast(rows->data); + float * engram_kv_data = static_cast(engram_kv->data); + float * q_data = static_cast(q_norm->data); + float * k_data = static_cast(k_norm->data); + for (int64_t i = 0; i < width*streams*tokens; ++i) { + residual_data[i] = bf16(0.125f + (float) (i%13)/16.0f); + } + std::fill(rows_data, rows_data + ggml_nelements(rows), 0.0f); + std::fill(engram_kv_data, engram_kv_data + ggml_nelements(engram_kv), 0.0f); + std::vector projected_data(5*width*tokens); + for (int64_t token = 0; token < tokens; ++token) { + rows_data[token*rows->ne[0]] = 1.0f; + for (int64_t i = 0; i < 5*width; ++i) { + projected_data[token*5*width + i] = 0.0625f + (float) ((token*5*width + i)%11)/32.0f; + engram_kv_data[i*engram_kv->ne[0] + token] = projected_data[token*5*width + i]; + } + } + for (int64_t i = 0; i < width*streams; ++i) { + q_data[i] = 0.5f + (float) (i%5)/8.0f; + k_data[i] = 0.75f - (float) (i%3)/16.0f; + } + static_cast(mask->data)[0] = 1.0f; + static_cast(mask->data)[1] = 0.0f; + + const std::vector original(residual_data, residual_data + width*streams*tokens); + ggml_tensor * output = llama_dsv41_build_engram( + ctx, residual, rows, engram_kv, q_norm, k_norm, mask, 1.0e-20f); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + check(ggml_graph_compute_with_ctx(ctx, graph, 1) == GGML_STATUS_SUCCESS, + "DeepSeek V4.1 Engram graph execution failed"); + + const float * actual = static_cast(output->data); + for (int64_t stream = 0; stream < streams; ++stream) { + double hidden_sq = 0.0; + double key_sq = 0.0; + double dot = 0.0; + for (int64_t i = 0; i < width; ++i) { + const float hidden = original[stream*width + i]; + const float key = bf16(projected_data[stream*width + i]); + hidden_sq += hidden*hidden; + key_sq += key*key; + dot += hidden*q_data[stream*width + i]*k_data[stream*width + i]*key; + } + dot /= std::sqrt(hidden_sq/width + 1.0e-20); + dot /= std::sqrt(key_sq/width + 1.0e-20); + dot /= std::sqrt((double) width); + const double gate = 1.0/(1.0 + std::exp(-std::copysign(std::sqrt(std::max(std::abs(dot), 1.0e-6)), dot))); + for (int64_t i = 0; i < width; ++i) { + const float value = bf16(projected_data[4*width + i]); + const float expected = bf16(original[stream*width + i] + (float) gate*value); + check(std::abs(actual[stream*width + i] - expected) <= std::max(1.0e-6f, std::abs(expected)/128.0f), + "DeepSeek V4.1 Engram gate differs from scalar reference"); + const size_t masked = width*streams + stream*width + i; + check(std::memcmp(actual + masked, original.data() + masked, sizeof(float)) == 0, + "masked DeepSeek V4.1 Engram row changed"); + } + } + + ggml_free(ctx); +} + +int main() { +#if !defined(_WIN32) + test_extent_validation(); + test_transactions_and_sequences(); + test_scheduler_upload(); + test_chunked_prefill(); + test_transactional_read_failure(); +#endif + test_graph_gate(); + std::puts("DeepSeek V4.1 Engram runtime and graph: PASS"); + return 0; +} diff --git a/tests/test-deepseek41-runtime.cpp b/tests/test-deepseek41-runtime.cpp index 3b996208d472..2d5945b407fb 100644 --- a/tests/test-deepseek41-runtime.cpp +++ b/tests/test-deepseek41-runtime.cpp @@ -111,7 +111,6 @@ static void test_hparams() { expect_throw([&]() { llama_dsv41_validate_config(config); }, "truncated Engram prime table was accepted"); const std::string dependency_error = llama_dsv41_runtime_dependency_error(); - check(dependency_error.find("disk-backed Engram") != std::string::npos, "dependency error omits Engram"); check(dependency_error.find("routed-expert streaming") != std::string::npos, "dependency error omits expert streaming"); } From 78175bfc23c0038558100b589c07a4862e8aee32 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 09:26:53 -0700 Subject: [PATCH 4/5] deepseek41 : preserve exact Engram gate semantics Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-dsv41-engram.cpp | 74 +++++++++++++++++++++----------- src/llama-dsv41-engram.h | 11 +++-- tests/test-deepseek41-engram.cpp | 70 ++++++++++++++++++++++++------ 3 files changed, 116 insertions(+), 39 deletions(-) diff --git a/src/llama-dsv41-engram.cpp b/src/llama-dsv41-engram.cpp index 578177b7a2d7..eac31309eba8 100644 --- a/src/llama-dsv41-engram.cpp +++ b/src/llama-dsv41-engram.cpp @@ -35,7 +35,6 @@ struct llama_dsv41_engram_transaction::impl { std::vector ids; std::array, LLAMA_ENGRAM_LAYERS> decoded; std::vector mask; - std::vector mask_f32; std::map next; }; @@ -71,7 +70,7 @@ void llama_dsv41_engram_transaction::upload_layer( size_t token_offset, size_t token_count, ggml_tensor * rows_input, - ggml_tensor * text_mask_input) const { + ggml_tensor * text_select_input) const { if (!pimpl->active) { throw std::invalid_argument("DeepSeek V4.1 Engram transaction is not active"); } @@ -80,11 +79,11 @@ void llama_dsv41_engram_transaction::upload_layer( throw std::invalid_argument("DeepSeek V4.1 Engram upload range is invalid"); } const int64_t row_width = LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM; - if (rows_input == nullptr || text_mask_input == nullptr || + if (rows_input == nullptr || text_select_input == nullptr || rows_input->type != GGML_TYPE_F32 || rows_input->ne[0] != row_width || ggml_nelements(rows_input) != row_width*(int64_t) token_count || - text_mask_input->type != GGML_TYPE_F32 || - ggml_nelements(text_mask_input) != (int64_t) token_count) { + text_select_input->type != GGML_TYPE_I32 || + ggml_nelements(text_select_input) != (int64_t) token_count) { throw std::invalid_argument("DeepSeek V4.1 Engram input tensor shape mismatch"); } @@ -94,11 +93,15 @@ void llama_dsv41_engram_transaction::upload_layer( pimpl->decoded[layer].data() + row_offset, 0, token_count*row_width*sizeof(float)); + std::vector select(token_count); + for (size_t i = 0; i < token_count; ++i) { + select[i] = pimpl->mask[token_offset + i] != 0 ? (int32_t) (token_count + i) : (int32_t) i; + } ggml_backend_tensor_set( - text_mask_input, - pimpl->mask_f32.data() + token_offset, + text_select_input, + select.data(), 0, - token_count*sizeof(float)); + token_count*sizeof(int32_t)); } struct llama_dsv41_engram_runtime::impl { @@ -114,8 +117,8 @@ struct llama_dsv41_engram_runtime::impl { size_t max_tokens) : hasher(std::move(layout)), max_tokens(max_tokens) { - if (max_tokens == 0) { - throw std::invalid_argument("DeepSeek V4.1 Engram token bound must be non-zero"); + if (max_tokens == 0 || max_tokens > (size_t) INT32_MAX/2) { + throw std::invalid_argument("DeepSeek V4.1 Engram token bound is invalid"); } for (size_t i = 0; i < LLAMA_ENGRAM_LAYERS; ++i) { llama_dsv41_validate_engram_extent(extents[i]); @@ -148,7 +151,6 @@ llama_dsv41_engram_transaction llama_dsv41_engram_runtime::prepare( result.pimpl->count = tokens.size(); result.pimpl->ids.resize(tokens.size()*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS); result.pimpl->mask.resize(tokens.size()); - result.pimpl->mask_f32.resize(tokens.size()); result.pimpl->next = pimpl->sequences; for (size_t i = 0; i < tokens.size(); ++i) { @@ -188,7 +190,6 @@ llama_dsv41_engram_transaction llama_dsv41_engram_runtime::prepare( expected, sizeof(expected)); result.pimpl->mask[i] = token.text != 0; - result.pimpl->mask_f32[i] = token.text != 0 ? 1.0f : 0.0f; } const size_t row_values = tokens.size()*LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM; @@ -279,13 +280,40 @@ static ggml_tensor * dsv41_bf16_f32(ggml_context * ctx, ggml_tensor * tensor) { return ggml_cast(ctx, ggml_cast(ctx, tensor, GGML_TYPE_BF16), GGML_TYPE_F32); } +static void dsv41_engram_gate_f32( + ggml_tensor * dst, + const ggml_tensor * src, + int ith, + int nth, + void *) { + GGML_ASSERT(dst->type == GGML_TYPE_F32 && src->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src)); + const float * input = static_cast(src->data); + float * output = static_cast(dst->data); + const int64_t count = ggml_nelements(src); + for (int64_t i = ith; i < count; i += nth) { + const float signed_root = std::copysign(std::sqrt(std::max(std::abs(input[i]), 1.0e-6f)), input[i]); + output[i] = 1.0f/(1.0f + std::exp(-signed_root)); + } +} + +ggml_tensor * llama_dsv41_build_engram_gate( + ggml_context * ctx, + ggml_tensor * dot) { + if (ctx == nullptr || dot == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 Engram gate input is null"); + } + // This small CPU fallback preserves copysign for signed zero on every scheduler backend. + return ggml_map_custom1(ctx, dot, dsv41_engram_gate_f32, GGML_N_TASKS_MAX, nullptr); +} + ggml_tensor * llama_dsv41_build_engram_add( ggml_context * ctx, ggml_tensor * residual, ggml_tensor * projected, ggml_tensor * q_norm, ggml_tensor * k_norm, - ggml_tensor * text_mask, + ggml_tensor * text_select, float rms_eps) { if (ctx == nullptr || residual == nullptr || projected == nullptr || q_norm == nullptr || k_norm == nullptr) { throw std::invalid_argument("DeepSeek V4.1 Engram graph input is null"); @@ -294,11 +322,12 @@ ggml_tensor * llama_dsv41_build_engram_add( const int64_t width = residual->ne[0]; const int64_t streams = residual->ne[1]; const int64_t tokens = residual->ne[2]; - if (width <= 0 || streams != 4 || tokens <= 0 || + if (width <= 0 || streams != 4 || tokens <= 0 || tokens > INT32_MAX/2 || projected->ne[0] != 5*width || projected->ne[1] != tokens || q_norm->ne[0] != width || q_norm->ne[1] != streams || k_norm->ne[0] != width || k_norm->ne[1] != streams || - (text_mask != nullptr && (text_mask->ne[0] != 1 || text_mask->ne[1] != tokens))) { + (text_select != nullptr && + (text_select->type != GGML_TYPE_I32 || ggml_nelements(text_select) != tokens))) { throw std::invalid_argument("DeepSeek V4.1 Engram graph shape mismatch"); } @@ -321,14 +350,11 @@ ggml_tensor * llama_dsv41_build_engram_add( dot = ggml_mul(ctx, dot, key_norm); dot = ggml_scale(ctx, ggml_sum_rows(ctx, dot), 1.0f/std::sqrt((float) width)); - ggml_tensor * magnitude = ggml_sqrt(ctx, ggml_clamp(ctx, ggml_abs(ctx, dot), 1.0e-6f, INFINITY)); - ggml_tensor * gate = ggml_sigmoid(ctx, ggml_mul(ctx, ggml_sgn(ctx, dot), magnitude)); - if (text_mask != nullptr) { - gate = ggml_mul(ctx, gate, text_mask); + ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot); + ggml_tensor * updated = dsv41_bf16_f32(ctx, ggml_add(ctx, hidden, ggml_mul(ctx, value, gate))); + if (text_select != nullptr) { + updated = ggml_get_rows(ctx, ggml_concat(ctx, hidden, updated, 1), text_select); } - - ggml_tensor * updated = ggml_add(ctx, hidden, ggml_mul(ctx, value, gate)); - updated = dsv41_bf16_f32(ctx, updated); updated = ggml_reshape_3d(ctx, updated, width, 1, tokens); result = result == nullptr ? updated : ggml_concat(ctx, result, updated, 1); } @@ -342,7 +368,7 @@ ggml_tensor * llama_dsv41_build_engram( ggml_tensor * engram_kv, ggml_tensor * q_norm, ggml_tensor * k_norm, - ggml_tensor * text_mask, + ggml_tensor * text_select, float rms_eps) { if (ctx == nullptr || residual == nullptr || rows == nullptr || engram_kv == nullptr || rows->ne[0] != LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM || @@ -352,5 +378,5 @@ ggml_tensor * llama_dsv41_build_engram( } ggml_tensor * projected = ggml_mul_mat(ctx, engram_kv, rows); return llama_dsv41_build_engram_add( - ctx, residual, projected, q_norm, k_norm, text_mask, rms_eps); + ctx, residual, projected, q_norm, k_norm, text_select, rms_eps); } diff --git a/src/llama-dsv41-engram.h b/src/llama-dsv41-engram.h index db0ae04dc853..0a92a039b3e0 100644 --- a/src/llama-dsv41-engram.h +++ b/src/llama-dsv41-engram.h @@ -59,7 +59,7 @@ struct llama_dsv41_engram_transaction { size_t token_offset, size_t token_count, ggml_tensor * rows_input, - ggml_tensor * text_mask_input) const; + ggml_tensor * text_select_input) const; struct impl; std::unique_ptr pimpl; @@ -91,15 +91,20 @@ class llama_dsv41_engram_runtime { std::unique_ptr pimpl; }; +// text_select row i keeps the original residual; row tokens+i selects the BF16-updated residual. ggml_tensor * llama_dsv41_build_engram_add( ggml_context * ctx, ggml_tensor * residual, ggml_tensor * projected, ggml_tensor * q_norm, ggml_tensor * k_norm, - ggml_tensor * text_mask, + ggml_tensor * text_select, float rms_eps); +ggml_tensor * llama_dsv41_build_engram_gate( + ggml_context * ctx, + ggml_tensor * dot); + ggml_tensor * llama_dsv41_build_engram( ggml_context * ctx, ggml_tensor * residual, @@ -107,5 +112,5 @@ ggml_tensor * llama_dsv41_build_engram( ggml_tensor * engram_kv, ggml_tensor * q_norm, ggml_tensor * k_norm, - ggml_tensor * text_mask, + ggml_tensor * text_select, float rms_eps); diff --git a/tests/test-deepseek41-engram.cpp b/tests/test-deepseek41-engram.cpp index 0eaf340a744c..1bc4e9bee700 100644 --- a/tests/test-deepseek41-engram.cpp +++ b/tests/test-deepseek41-engram.cpp @@ -222,7 +222,7 @@ static void test_scheduler_upload() { check(ctx != nullptr, "failed to create DeepSeek V4.1 Engram upload context"); ggml_tensor * rows = ggml_new_tensor_2d( ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, transaction.token_count()); - ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, transaction.token_count()); + ggml_tensor * select = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, transaction.token_count()); ggml_backend_t backend = ggml_backend_cpu_init(); check(backend != nullptr, "failed to create DeepSeek V4.1 Engram upload backend"); @@ -230,23 +230,23 @@ static void test_scheduler_upload() { check(buffer != nullptr, "failed to allocate DeepSeek V4.1 Engram upload tensors"); std::vector actual_rows(ggml_nelements(rows)); - std::vector actual_mask(ggml_nelements(mask)); + std::vector actual_select(ggml_nelements(select)); for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { - transaction.upload_layer(layer, 0, transaction.token_count(), rows, mask); + transaction.upload_layer(layer, 0, transaction.token_count(), rows, select); ggml_backend_tensor_get(rows, actual_rows.data(), 0, ggml_nbytes(rows)); - ggml_backend_tensor_get(mask, actual_mask.data(), 0, ggml_nbytes(mask)); + ggml_backend_tensor_get(select, actual_select.data(), 0, ggml_nbytes(select)); check(std::memcmp( actual_rows.data(), transaction.rows(layer), ggml_nbytes(rows)) == 0, "scheduler-backed Engram row upload changed the bounded pack"); - check(actual_mask == std::vector({ 1.0f, 0.0f, 1.0f }), - "scheduler-backed Engram text mask upload changed"); + check(actual_select == std::vector({ 3, 1, 5 }), + "scheduler-backed Engram text selection upload changed"); } runtime.commit(transaction); expect_invalid( - [&] { transaction.upload_layer(0, 0, 3, rows, mask); }, + [&] { transaction.upload_layer(0, 0, 3, rows, select); }, "committed Engram transaction uploaded stale input"); ggml_backend_buffer_free(buffer); @@ -335,7 +335,7 @@ static void test_graph_gate() { ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, 5*width); ggml_tensor * q_norm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, streams); ggml_tensor * k_norm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, streams); - ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, tokens); + ggml_tensor * select = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens); float * residual_data = static_cast(residual->data); float * rows_data = static_cast(rows->data); @@ -343,8 +343,12 @@ static void test_graph_gate() { float * q_data = static_cast(q_norm->data); float * k_data = static_cast(k_norm->data); for (int64_t i = 0; i < width*streams*tokens; ++i) { - residual_data[i] = bf16(0.125f + (float) (i%13)/16.0f); + residual_data[i] = i < width*streams ? + bf16(0.125f + (float) (i%13)/16.0f) : + 0.1234567f + (float) (i%13)/17.0f; } + residual_data[width*streams] = -0.0f; + residual_data[width*streams + 1] = 0.0f; std::fill(rows_data, rows_data + ggml_nelements(rows), 0.0f); std::fill(engram_kv_data, engram_kv_data + ggml_nelements(engram_kv), 0.0f); std::vector projected_data(5*width*tokens); @@ -359,12 +363,12 @@ static void test_graph_gate() { q_data[i] = 0.5f + (float) (i%5)/8.0f; k_data[i] = 0.75f - (float) (i%3)/16.0f; } - static_cast(mask->data)[0] = 1.0f; - static_cast(mask->data)[1] = 0.0f; + static_cast(select->data)[0] = tokens; + static_cast(select->data)[1] = 1; const std::vector original(residual_data, residual_data + width*streams*tokens); ggml_tensor * output = llama_dsv41_build_engram( - ctx, residual, rows, engram_kv, q_norm, k_norm, mask, 1.0e-20f); + ctx, residual, rows, engram_kv, q_norm, k_norm, select, 1.0e-20f); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, output); check(ggml_graph_compute_with_ctx(ctx, graph, 1) == GGML_STATUS_SUCCESS, @@ -400,6 +404,47 @@ static void test_graph_gate() { ggml_free(ctx); } +static void test_signed_zero_gate() { + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create DeepSeek V4.1 signed-zero graph context"); + + ggml_tensor * dot = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2); + ggml_set_input(dot); + ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot); + ggml_set_output(gate); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, gate); + + ggml_backend_t backend = ggml_backend_cpu_init(); + check(backend != nullptr, "failed to create DeepSeek V4.1 signed-zero backend"); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 16, false, true); + check(sched != nullptr, "failed to create DeepSeek V4.1 signed-zero scheduler"); + check(ggml_backend_sched_alloc_graph(sched, graph), "failed to allocate DeepSeek V4.1 signed-zero graph"); + + const float input[] = { 0.0f, -0.0f }; + ggml_backend_tensor_set(dot, input, 0, sizeof(input)); + check(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS, + "DeepSeek V4.1 signed-zero scheduler execution failed"); + const float positive = 1.0f/(1.0f + std::exp(-0.001f)); + const float negative = 1.0f/(1.0f + std::exp(0.001f)); + float actual[2]; + ggml_backend_tensor_get(gate, actual, 0, sizeof(actual)); + check(std::abs(actual[0] - positive) < 1.0e-7f && actual[0] > 0.5f, + "DeepSeek V4.1 positive-zero gate lost copysign semantics"); + check(std::abs(actual[1] - negative) < 1.0e-7f && actual[1] < 0.5f, + "DeepSeek V4.1 negative-zero gate lost copysign semantics"); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + int main() { #if !defined(_WIN32) test_extent_validation(); @@ -409,6 +454,7 @@ int main() { test_transactional_read_failure(); #endif test_graph_gate(); + test_signed_zero_gate(); std::puts("DeepSeek V4.1 Engram runtime and graph: PASS"); return 0; } From 4ac4d4a8433dc5862256874db1fd4257a6a90cf9 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 09:42:48 -0700 Subject: [PATCH 5/5] deepseek41 : pin Engram gate to local CPU Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-dsv41-engram.cpp | 34 ++++++--- src/llama-dsv41-engram.h | 13 +++- tests/test-deepseek41-engram.cpp | 119 ++++++++++++++++++++++++------- 3 files changed, 130 insertions(+), 36 deletions(-) diff --git a/src/llama-dsv41-engram.cpp b/src/llama-dsv41-engram.cpp index eac31309eba8..b2273a5236f2 100644 --- a/src/llama-dsv41-engram.cpp +++ b/src/llama-dsv41-engram.cpp @@ -2,6 +2,7 @@ #include "ggml.h" #include "ggml-backend.h" +#include "ggml-cpu.h" #include #include @@ -299,12 +300,25 @@ static void dsv41_engram_gate_f32( ggml_tensor * llama_dsv41_build_engram_gate( ggml_context * ctx, - ggml_tensor * dot) { - if (ctx == nullptr || dot == nullptr) { + ggml_tensor * dot, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || dot == nullptr || sched == nullptr || backend_cpu == nullptr) { throw std::invalid_argument("DeepSeek V4.1 Engram gate input is null"); } - // This small CPU fallback preserves copysign for signed zero on every scheduler backend. - return ggml_map_custom1(ctx, dot, dsv41_engram_gate_f32, GGML_N_TASKS_MAX, nullptr); + if (!ggml_backend_is_cpu(backend_cpu)) { + throw std::invalid_argument("DeepSeek V4.1 Engram gate backend is not local CPU"); + } + bool found = false; + for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { + found = found || ggml_backend_sched_get_backend(sched, i) == backend_cpu; + } + if (!found) { + throw std::invalid_argument("DeepSeek V4.1 Engram gate CPU backend is not in the scheduler"); + } + ggml_tensor * gate = ggml_map_custom1(ctx, dot, dsv41_engram_gate_f32, GGML_N_TASKS_MAX, nullptr); + ggml_backend_sched_set_tensor_backend(sched, gate, backend_cpu); + return gate; } ggml_tensor * llama_dsv41_build_engram_add( @@ -314,7 +328,9 @@ ggml_tensor * llama_dsv41_build_engram_add( ggml_tensor * q_norm, ggml_tensor * k_norm, ggml_tensor * text_select, - float rms_eps) { + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { if (ctx == nullptr || residual == nullptr || projected == nullptr || q_norm == nullptr || k_norm == nullptr) { throw std::invalid_argument("DeepSeek V4.1 Engram graph input is null"); } @@ -350,7 +366,7 @@ ggml_tensor * llama_dsv41_build_engram_add( dot = ggml_mul(ctx, dot, key_norm); dot = ggml_scale(ctx, ggml_sum_rows(ctx, dot), 1.0f/std::sqrt((float) width)); - ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot); + ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot, sched, backend_cpu); ggml_tensor * updated = dsv41_bf16_f32(ctx, ggml_add(ctx, hidden, ggml_mul(ctx, value, gate))); if (text_select != nullptr) { updated = ggml_get_rows(ctx, ggml_concat(ctx, hidden, updated, 1), text_select); @@ -369,7 +385,9 @@ ggml_tensor * llama_dsv41_build_engram( ggml_tensor * q_norm, ggml_tensor * k_norm, ggml_tensor * text_select, - float rms_eps) { + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { if (ctx == nullptr || residual == nullptr || rows == nullptr || engram_kv == nullptr || rows->ne[0] != LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM || engram_kv->ne[0] != rows->ne[0] || @@ -378,5 +396,5 @@ ggml_tensor * llama_dsv41_build_engram( } ggml_tensor * projected = ggml_mul_mat(ctx, engram_kv, rows); return llama_dsv41_build_engram_add( - ctx, residual, projected, q_norm, k_norm, text_select, rms_eps); + ctx, residual, projected, q_norm, k_norm, text_select, rms_eps, sched, backend_cpu); } diff --git a/src/llama-dsv41-engram.h b/src/llama-dsv41-engram.h index 0a92a039b3e0..2edca59c12d0 100644 --- a/src/llama-dsv41-engram.h +++ b/src/llama-dsv41-engram.h @@ -1,5 +1,6 @@ #pragma once +#include "ggml-backend.h" #include "llama-engram.h" #include "llama.h" @@ -99,11 +100,15 @@ ggml_tensor * llama_dsv41_build_engram_add( ggml_tensor * q_norm, ggml_tensor * k_norm, ggml_tensor * text_select, - float rms_eps); + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); ggml_tensor * llama_dsv41_build_engram_gate( ggml_context * ctx, - ggml_tensor * dot); + ggml_tensor * dot, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); ggml_tensor * llama_dsv41_build_engram( ggml_context * ctx, @@ -113,4 +118,6 @@ ggml_tensor * llama_dsv41_build_engram( ggml_tensor * q_norm, ggml_tensor * k_norm, ggml_tensor * text_select, - float rms_eps); + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); diff --git a/tests/test-deepseek41-engram.cpp b/tests/test-deepseek41-engram.cpp index 1bc4e9bee700..ee73c48628f2 100644 --- a/tests/test-deepseek41-engram.cpp +++ b/tests/test-deepseek41-engram.cpp @@ -323,7 +323,7 @@ static void test_graph_gate() { ggml_init_params params = { /*.mem_size =*/ 8*1024*1024, /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ false, + /*.no_alloc =*/ true, }; ggml_context * ctx = ggml_init(params); check(ctx != nullptr, "failed to create DeepSeek V4.1 Engram graph context"); @@ -337,11 +337,18 @@ static void test_graph_gate() { ggml_tensor * k_norm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, streams); ggml_tensor * select = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens); - float * residual_data = static_cast(residual->data); - float * rows_data = static_cast(rows->data); - float * engram_kv_data = static_cast(engram_kv->data); - float * q_data = static_cast(q_norm->data); - float * k_data = static_cast(k_norm->data); + ggml_set_input(residual); + ggml_set_input(rows); + ggml_set_input(engram_kv); + ggml_set_input(q_norm); + ggml_set_input(k_norm); + ggml_set_input(select); + + std::vector residual_data(ggml_nelements(residual)); + std::vector rows_data(ggml_nelements(rows)); + std::vector engram_kv_data(ggml_nelements(engram_kv)); + std::vector q_data(ggml_nelements(q_norm)); + std::vector k_data(ggml_nelements(k_norm)); for (int64_t i = 0; i < width*streams*tokens; ++i) { residual_data[i] = i < width*streams ? bf16(0.125f + (float) (i%13)/16.0f) : @@ -349,8 +356,6 @@ static void test_graph_gate() { } residual_data[width*streams] = -0.0f; residual_data[width*streams + 1] = 0.0f; - std::fill(rows_data, rows_data + ggml_nelements(rows), 0.0f); - std::fill(engram_kv_data, engram_kv_data + ggml_nelements(engram_kv), 0.0f); std::vector projected_data(5*width*tokens); for (int64_t token = 0; token < tokens; ++token) { rows_data[token*rows->ne[0]] = 1.0f; @@ -363,18 +368,31 @@ static void test_graph_gate() { q_data[i] = 0.5f + (float) (i%5)/8.0f; k_data[i] = 0.75f - (float) (i%3)/16.0f; } - static_cast(select->data)[0] = tokens; - static_cast(select->data)[1] = 1; + const int32_t select_data[] = { (int32_t) tokens, 1 }; - const std::vector original(residual_data, residual_data + width*streams*tokens); + const std::vector original = residual_data; + ggml_backend_t backend = ggml_backend_cpu_init(); + check(backend != nullptr, "failed to create DeepSeek V4.1 Engram graph backend"); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 512, false, true); + check(sched != nullptr, "failed to create DeepSeek V4.1 Engram graph scheduler"); ggml_tensor * output = llama_dsv41_build_engram( - ctx, residual, rows, engram_kv, q_norm, k_norm, select, 1.0e-20f); + ctx, residual, rows, engram_kv, q_norm, k_norm, select, 1.0e-20f, sched, backend); + ggml_set_output(output); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, output); - check(ggml_graph_compute_with_ctx(ctx, graph, 1) == GGML_STATUS_SUCCESS, + check(ggml_backend_sched_alloc_graph(sched, graph), "failed to allocate DeepSeek V4.1 Engram graph"); + ggml_backend_tensor_set(residual, residual_data.data(), 0, ggml_nbytes(residual)); + ggml_backend_tensor_set(rows, rows_data.data(), 0, ggml_nbytes(rows)); + ggml_backend_tensor_set(engram_kv, engram_kv_data.data(), 0, ggml_nbytes(engram_kv)); + ggml_backend_tensor_set(q_norm, q_data.data(), 0, ggml_nbytes(q_norm)); + ggml_backend_tensor_set(k_norm, k_data.data(), 0, ggml_nbytes(k_norm)); + ggml_backend_tensor_set(select, select_data, 0, sizeof(select_data)); + check(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS, "DeepSeek V4.1 Engram graph execution failed"); - const float * actual = static_cast(output->data); + std::vector actual(ggml_nelements(output)); + ggml_backend_tensor_get(output, actual.data(), 0, ggml_nbytes(output)); for (int64_t stream = 0; stream < streams; ++stream) { double hidden_sq = 0.0; double key_sq = 0.0; @@ -396,15 +414,19 @@ static void test_graph_gate() { check(std::abs(actual[stream*width + i] - expected) <= std::max(1.0e-6f, std::abs(expected)/128.0f), "DeepSeek V4.1 Engram gate differs from scalar reference"); const size_t masked = width*streams + stream*width + i; - check(std::memcmp(actual + masked, original.data() + masked, sizeof(float)) == 0, + check(std::memcmp(actual.data() + masked, original.data() + masked, sizeof(float)) == 0, "masked DeepSeek V4.1 Engram row changed"); } } + ggml_backend_sched_free(sched); + ggml_backend_free(backend); ggml_free(ctx); } static void test_signed_zero_gate() { + constexpr int64_t matrix_size = 32; + ggml_init_params params = { /*.mem_size =*/ 1024*1024, /*.mem_buffer =*/ nullptr, @@ -414,34 +436,81 @@ static void test_signed_zero_gate() { check(ctx != nullptr, "failed to create DeepSeek V4.1 signed-zero graph context"); ggml_tensor * dot = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2); + ggml_tensor * matrix_a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, matrix_size, matrix_size); + ggml_tensor * matrix_b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, matrix_size, matrix_size); ggml_set_input(dot); - ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot); - ggml_set_output(gate); - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, gate); + ggml_set_input(matrix_a); + ggml_set_input(matrix_b); + ggml_tensor * accelerated = ggml_mul_mat(ctx, matrix_a, matrix_b); + + ggml_backend_load_all(); + ggml_backend_t backend_accel = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t device = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(device) == GGML_BACKEND_DEVICE_TYPE_ACCEL && + ggml_backend_dev_supports_op(device, accelerated)) { + backend_accel = ggml_backend_dev_init(device, nullptr); + if (backend_accel != nullptr) { + break; + } + } + } - ggml_backend_t backend = ggml_backend_cpu_init(); - check(backend != nullptr, "failed to create DeepSeek V4.1 signed-zero backend"); - ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); - ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 16, false, true); + ggml_backend_t backend_cpu = ggml_backend_cpu_init(); + check(backend_cpu != nullptr, "failed to create DeepSeek V4.1 signed-zero CPU backend"); + ggml_backend_t backends[] = { + backend_accel != nullptr ? backend_accel : backend_cpu, + backend_cpu, + }; + ggml_backend_buffer_type_t bufts[] = { + ggml_backend_get_default_buffer_type(backends[0]), + ggml_backend_cpu_buffer_type(), + }; + const int n_backends = backend_accel != nullptr ? 2 : 1; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, bufts, n_backends, 32, false, true); check(sched != nullptr, "failed to create DeepSeek V4.1 signed-zero scheduler"); + if (backend_accel != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, accelerated, backend_accel); + expect_invalid( + [&] { llama_dsv41_build_engram_gate(ctx, dot, sched, backend_accel); }, + "DeepSeek V4.1 gate accepted a non-CPU backend"); + } + + ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot, sched, backend_cpu); + ggml_tensor * output = ggml_add(ctx, gate, ggml_repeat(ctx, ggml_sum(ctx, accelerated), gate)); + ggml_set_output(output); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + check(ggml_backend_sched_alloc_graph(sched, graph), "failed to allocate DeepSeek V4.1 signed-zero graph"); + if (backend_accel != nullptr) { + check(ggml_backend_sched_get_tensor_backend(sched, accelerated) == backend_accel, + "DeepSeek V4.1 gate moved surrounding work off its accelerator"); + } + check(ggml_backend_sched_get_tensor_backend(sched, gate) == backend_cpu, + "DeepSeek V4.1 gate was not assigned to the local CPU backend"); const float input[] = { 0.0f, -0.0f }; + std::vector zeros(matrix_size*matrix_size); ggml_backend_tensor_set(dot, input, 0, sizeof(input)); + ggml_backend_tensor_set(matrix_a, zeros.data(), 0, ggml_nbytes(matrix_a)); + ggml_backend_tensor_set(matrix_b, zeros.data(), 0, ggml_nbytes(matrix_b)); check(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS, "DeepSeek V4.1 signed-zero scheduler execution failed"); const float positive = 1.0f/(1.0f + std::exp(-0.001f)); const float negative = 1.0f/(1.0f + std::exp(0.001f)); float actual[2]; - ggml_backend_tensor_get(gate, actual, 0, sizeof(actual)); + ggml_backend_tensor_get(output, actual, 0, sizeof(actual)); check(std::abs(actual[0] - positive) < 1.0e-7f && actual[0] > 0.5f, "DeepSeek V4.1 positive-zero gate lost copysign semantics"); check(std::abs(actual[1] - negative) < 1.0e-7f && actual[1] < 0.5f, "DeepSeek V4.1 negative-zero gate lost copysign semantics"); ggml_backend_sched_free(sched); - ggml_backend_free(backend); + if (backend_accel != nullptr) { + ggml_backend_free(backend_accel); + } + ggml_backend_free(backend_cpu); ggml_free(ctx); }