ivanfioravanti
bdae7640ef
Prefill reads 2 x 24 Engram rows per token, spread across 189 GiB. That random pattern is the worst case for external SSDs: a USB drive served about 6K rows/s over the whole span, uncached or buffered alike, so a 2K prompt stalled about 12 s before layers 1 and 14. Q2 single-Mac prefill ran at 205-290 t/s while its GPU work alone supported ~420. When the resident model, graph buffers, a 64 GiB reserve and the tables fit in physical memory (Q2 on a 512 GB Mac; not Q4), a low-priority background reader streams the tables once into the kernel page cache at the drive's sequential rate (~1 GB/s over USB, about 3.5 minutes). The uncached row readers then hit memory. No table is allocated or wired, the pages stay reclaimable, and the warmer stops at memory pressure. The same bytes are read either way, so logits do not change. Q2 from a USB SSD, cold prompts: 2K 198 -> 420 t/s, 8K 214 -> 518 t/s. DS4_DISABLE_V41_ENGRAM_WARM=1 skips the warmer; DS4_V41_ENGRAM_WARM_WAIT=1 blocks until it finishes, for benchmarks.
91 lines
4.1 KiB
C
91 lines
4.1 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);
|
|
|
|
/* One reusable reader for a token's two tables. Tables/IDs are copied at
|
|
* start; descriptors and output must remain valid until wait(reader, 1) or
|
|
* free, even if wait(reader, 0) reports an error. No table data is cached.
|
|
* One caller owns the reader. Creation failure permits synchronous fallback. */
|
|
typedef struct ds4_engram_reader ds4_engram_reader;
|
|
ds4_engram_reader *ds4_engram_reader_create(void);
|
|
bool ds4_engram_reader_start(ds4_engram_reader *reader,
|
|
const ds4_engram_table tables[2], const uint32_t ids[2][DS4_ENGRAM_COLS], float *out);
|
|
bool ds4_engram_reader_wait(ds4_engram_reader *reader, unsigned table);
|
|
void ds4_engram_reader_free(ds4_engram_reader *reader);
|
|
|
|
/* Background sequential read of [offset, offset + bytes) into the kernel page
|
|
* cache, so the uncached row readers above hit memory. No table allocation;
|
|
* stops early at memory pressure. Wait reports whether it covered the span. */
|
|
typedef struct ds4_engram_warm ds4_engram_warm;
|
|
uint64_t ds4_engram_physical_memory(void);
|
|
ds4_engram_warm *ds4_engram_warm_start(const char *path, uint64_t offset, uint64_t bytes);
|
|
bool ds4_engram_warm_wait(ds4_engram_warm *warm, double *seconds);
|
|
void ds4_engram_warm_stop(ds4_engram_warm *warm);
|
|
|
|
#endif
|