Expand description
DeepSeek-V4.1 Flash engram on the GPU.
Engram sits before layers 1 and 14 (engram_layer_ids). Per token it hashes
the n-gram window of compressed token ids into n_hash_cols row ids of a
per-layer table, gathers those rows, projects their concatenation through
wkv into one key per hyper-connection stream plus one shared value, and
adds gate * value to each stream on the mHC highway. Four stages, each
held to the CPU reference crate::layers::deepseek_v41_ref::engram:
- Hash (
EngramHasher): integer-exact on the CPU; 48 ids per token is not a GPU problem. The tables (token_map, multipliers, primes, offsets) come from the GGUF metadata, proven identical toengram.py. - Rows: the checkpoint’s tables are Q2_K, one 84-byte block per
256-wide row, read by id through
spark_runtime::weights::expert_stream::EngramRowReaderand dequantised on the device by the Q2_K arm ofdequant_gguf_bf16straight into the[tokens * cols, head_dim]bf16 input (EngramV41::rows_from_q2k). - Projection: the dense bf16 GEMM,
[tokens, cols * head_dim]xwkv^T->[tokens, dim * (hc + 1)]. - Gate (
engram_v41_gate,kernels/gb10/deepseek-v4-flash/nvfp4/engram_v41.cu): normalised dot of stream against key, signed-sqrt sigmoid,h += gate * value, in place on the FP32 highway[T, hc, H]the mHC kernels keep.
Oracle: engram_v41_tests.rs (the hash exact against the reference; the
GPU path against the golden’s engram_out captures; real-table rows
bit-identical to the CPU decoder).
Structs§
- Engram
Hash Tables - The hash tables, shaped as the reference indexes them.
- Engram
Hasher - The n-gram hash for one sequence: keeps the compressed ids of every position so decode steps see the prefill’s look-back, exactly as the reference’s cache does.
- Engram
Layer Weights - One engram layer’s resident weights on the device.
- Engram
V41 - The GPU side: kernels, per-layer weights, and workspaces for up to
max_tokenspositions per call.
Constants§
- ENGRAM_
ROW_ BYTES - Bytes of one Q2_K super-block = one engram row.