Module engram_v41

Module engram_v41 

Source
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:

  1. 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 to engram.py.
  2. 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::EngramRowReader and dequantised on the device by the Q2_K arm of dequant_gguf_bf16 straight into the [tokens * cols, head_dim] bf16 input (EngramV41::rows_from_q2k).
  3. Projection: the dense bf16 GEMM, [tokens, cols * head_dim] x wkv^T -> [tokens, dim * (hc + 1)].
  4. 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§

EngramHashTables
The hash tables, shaped as the reference indexes them.
EngramHasher
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.
EngramLayerWeights
One engram layer’s resident weights on the device.
EngramV41
The GPU side: kernels, per-layer weights, and workspaces for up to max_tokens positions per call.

Constants§

ENGRAM_ROW_BYTES
Bytes of one Q2_K super-block = one engram row.