Module engram_v41

Module engram_v41 

Source
Expand description

2026-09-25: DeepSeek-V4.1 Flash engram on the GPU: n-gram hash, row gather, wkv projection and gate.

Owner: model-arch (DeepSeek-V4.1 engram). Invariants:

  • rows_from_q2k, rows_from_bf16 and apply refuse more rows or tokens than the max_tokens workspace before any copy or launch.
  • hc is in 1..=4 (EngramV41::new).

Engram runs before each layer in 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:

  1. Hash (EngramHasher): integer-exact, on the CPU, from the tables in the GGUF metadata.
  2. Rows: the tables are Q2_K, one ENGRAM_ROW_BYTES block per 256-wide row, read by id through EngramRowReader (metrale_model_weights) and dequantised on the device by dequant_q2_k_to_bf16 into the [tokens * cols, head_dim] bf16 input (EngramV41::rows_from_q2k).
  3. Projection: [tokens, cols * head_dim] x wkv^T -> [tokens, dim * (hc + 1)], a GEMM for several tokens and a GEMV for one.
  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].

engram_v41_tests.rs checks the hash exactly against the CPU reference in crate::deepseek_v41_ref::engram, and (on a GPU) the projection and gate against the golden captures and the device dequant against the CPU decoder.

Structs§

EngramHashTables
2026-09-25: The hash tables, in the shapes the CPU reference’s EngramTables uses.
EngramHasher
2026-09-25: The n-gram hash for one sequence. It keeps the compressed id of every position, so decode steps see the prefill’s look-back.
EngramLayerWeights
2026-09-25: One engram layer’s resident weights on the device.
EngramV41
2026-09-25: The GPU side: kernels, per-layer weights, and workspaces for up to max_tokens positions per call.

Constants§

ENGRAM_ROW_BYTES
2026-09-25: Bytes of one Q2_K super-block, which is one engram row.