71 lines
2.9 KiB
C
71 lines
2.9 KiB
C
#ifndef DS4_ENGRAM_H
|
|
#define DS4_ENGRAM_H
|
|
|
|
#include <stdbool.h>
|
|
#include <stddef.h>
|
|
#include <stdint.h>
|
|
|
|
enum {
|
|
DS4_ENGRAM_LAYERS = 2,
|
|
DS4_ENGRAM_NGRAM = 4,
|
|
DS4_ENGRAM_HEADS = 8,
|
|
DS4_ENGRAM_COLS = 24,
|
|
DS4_ENGRAM_DIM = 256,
|
|
DS4_ENGRAM_ROW_BYTES = 264,
|
|
DS4_ENGRAM_DEAD = -1
|
|
};
|
|
|
|
typedef struct {
|
|
const uint32_t *token_map;
|
|
uint32_t vocab_size, compressed_vocab_size, pad_id;
|
|
uint32_t rows[DS4_ENGRAM_LAYERS];
|
|
uint64_t multipliers[DS4_ENGRAM_LAYERS][DS4_ENGRAM_NGRAM];
|
|
uint32_t primes[DS4_ENGRAM_LAYERS][DS4_ENGRAM_COLS];
|
|
} ds4_engram_layout;
|
|
|
|
/* Newest first; copying this value also snapshots the complete hash state. */
|
|
typedef struct {
|
|
int32_t tail[DS4_ENGRAM_NGRAM - 1];
|
|
} ds4_engram_history;
|
|
|
|
bool ds4_engram_layout_valid(const ds4_engram_layout *layout);
|
|
void ds4_engram_history_reset(ds4_engram_history *history);
|
|
/* Validate the layout once at model load. Mask zero breaks all n-grams that
|
|
* cross that position. Output is [token][layer][column], including masked rows;
|
|
* the graph must suppress Engram at masked positions, as in the reference. */
|
|
bool ds4_engram_hash(const ds4_engram_layout *layout,
|
|
ds4_engram_history *history,
|
|
const int *tokens, const uint8_t *mask, size_t count,
|
|
uint32_t *rows);
|
|
|
|
typedef struct {
|
|
int fd;
|
|
uint64_t offset;
|
|
uint32_t rows;
|
|
} ds4_engram_table;
|
|
|
|
/* A separate uncached file descriptor, never an mmap or Metal model view.
|
|
* Each GGUF I8 row is 256 E4M3 bytes followed by 8 original E8M0 scales. */
|
|
bool ds4_engram_table_open(ds4_engram_table *table, const char *path,
|
|
uint64_t offset, uint32_t rows);
|
|
void ds4_engram_table_close(ds4_engram_table *table);
|
|
/* Output uses F32 storage for the reference's BF16-rounded values. No whole
|
|
* table allocation; caller owns count * DIM floats. Failure invalidates output. */
|
|
bool ds4_engram_read(const ds4_engram_table *table, const uint32_t *rows,
|
|
size_t count, float *out);
|
|
/* Read COLS rows per token, restoring token order after deduplicated disk reads.
|
|
* Input stride is in row IDs; output is packed [token][COLS][DIM]. Temporary
|
|
* storage is bounded to 384 KiB, independent of the table and prefix size.
|
|
* On macOS, large batches use bounded concurrent pread readers. */
|
|
bool ds4_engram_read_batch(const ds4_engram_table *table, const uint32_t *rows,
|
|
size_t tokens, size_t stride, float *out);
|
|
|
|
/* The same read on a worker: `rows` and `out` stay the caller's until finish,
|
|
* which joins the worker and returns its result. */
|
|
typedef struct ds4_engram_prefetch ds4_engram_prefetch;
|
|
ds4_engram_prefetch *ds4_engram_read_batch_start(const ds4_engram_table *table, const uint32_t *rows,
|
|
size_t tokens, size_t stride, float *out);
|
|
bool ds4_engram_read_batch_finish(ds4_engram_prefetch *prefetch);
|
|
|
|
#endif
|