Module engram

Module engram 

Source
Expand description

DeepSeek-V4.1 engram, CPU reference, checked sub-op by sub-op against the golden.

Four stages, each independently testable against its own capture:

  1. Hash (NgramHashState): compressed token map, n-gram window with pad and a cache that carries look-back across the prefill/decode split, per-lookback multiply, XOR-rolled, then % prime per (n-gram size, head) plus that slot’s offset. Integer-exact.
  2. Table rows (embed_rows): the fp8 e4m3 table, dequantised per 32-block by an e8m0 scale, gathered by hash id. e4m3 values are exactly representable in bf16, so with the fixture’s unit scales this stage is exact too.
  3. Projection (wkv): one bf16 GEMM, [n_hash_cols * head_dim] -> [dim * (hc_mult + 1)], split into one key per hyper-connection copy plus one shared value.
  4. Gate (gate_and_add): normalised dot product of stream against key, signed-sqrt sigmoid, h + gate * value, back to bf16.

The tables (token_map, multipliers, primes, offsets) are consumed as GIVEN, never recomputed: in production they come from the GGUF’s deepseek41.engram.* metadata, which gen_ds41_engram_hash_check.py proved identical to the reference’s own computation.

Structs§

EngramTables
The hash tables and geometry, exactly as the loader will read them from the GGUF.
NgramHashState
engram.py::NgramHashState. The cache holds compressed ids for every position seen, so a decode step’s look-back reaches into the prefill.

Functions§

e4m3_to_f32
embed_rows
ParallelEngramEmbedding.forward with unit scales: gather rows by id. Output [n_ids, head_dim], bf16-exact because e4m3 ⊂ bf16.
f32_to_e4m3_rne
Encode with round-to-nearest-even, as torch’s .to(float8_e4m3fn) does. Saturates above 448 (callers clamp first; the fixture never gets there).
f32_to_e8m0_rne
Torch’s f32 -> e8m0 cast, as a power of two: nearest exponent, ties (exactly 1.5 * 2^e) to the even exponent. Pinned against the fixture’s shipped scales by e8m0_rounding_rule_matches_torch.
gate_and_add
Engram.forward after the projection: x [tokens, hc, dim] (bf16 values), kv [tokens, dim*(hc+1)], q/k weights [hc, dim]. Returns bf16-rounded [tokens, hc, dim].
linear_bf16
F.linear on bf16: f32 accumulate, bf16 output. x: [rows, in], w: [out, in].
regen_bf16_matrix
A bf16-stored [out, in] weight regenerated by the fixture rule scale = in^-0.5.
regen_bf16_qk
A bf16-stored [hc, dim] engram q/k weight regenerated by the fixture rule (offset 1, ±0.05).
regen_table
The fixture’s fp8 table for layer_id, regenerated and dequantised with the SHIPPED per-block e8m0 scales (ParallelEngramEmbedding.forward: values.float() * scales.float() per block, then bf16). [rows * head_dim] f32; exact, since e4m3 times a power of two is bf16-exact.