- in
- [B batch, T positions, 4 streams, 5120 d] [B batch, T positions, 2 Engram layers, 24 n-grams × heads]
- out
- [B batch, T positions, 4 streams, 5120 d]
- parameters
- 98,465,597,952
Parameters
98,465,597,952 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Engram table | 98,308,270,592 | |
| Key and value projection | 157,286,400 | |
| Gate per stream | 40,960 | |
| Per-instance total | 98,465,597,952 | |
Derived from the pinned reference implementation and configuration; equal to the element counts of the checkpoint’s own tensors on the text path. The quantisation scales, the vision encoder and the DSpark blocks are not counted.
Explanation
It is applied to the streams before layer 14’s block (model.py 1262–1263). The weights with which that layer’s attention reads the streams were computed before this addition.
The table has 384,016,682 rows of 256 features: each of the 3 n-gram sizes has 8 heads, each with its own range of rows of a prime length (engram.py 102–126).
Codemodel.py 296–1272 · 2 more files
class Engram(nn.Module):
"""Writes an n-gram lookup into the residual stream, gated by how well it matches that stream.
The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per hc copy plus a shared
value. The gate is a normalized dot product of stream against key.
"""
def __init__(self, args: ModelArgs, layer_id: int, layout: EngramLayout):
super().__init__()
self.layer_id = layer_id
self.layer_hash_index = layout.layer_ids.index(layer_id)
self.dim = args.dim
self.hc_mult = args.hc_mult
self.clamp_value = 1e-6
self.embed = ParallelEngramEmbedding(layout.num_embeddings[self.layer_hash_index], layout.head_dim)
n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
self.wkv = Linear(n_hash_cols * layout.head_dim, args.dim * (args.hc_mult + 1))
self.eps = args.norm_eps
self.q_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim))
self.k_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim))
def forward(self, x: torch.Tensor, hash_ids: torch.Tensor, token_mask: torch.Tensor | None = None) -> torch.Tensor:
"""x: [B, L, hc_mult, dim]; hash_ids: [B, L, n_hash_cols]; token_mask: [B, L], False shuts
the gate so those positions pass through untouched."""
kv = self.wkv(self.embed(hash_ids).flatten(-2))
key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1)
key = key.float().unflatten(-1, (self.hc_mult, self.dim))
weight = self.q_weight.float() * self.k_weight.float() # only ever used as a product
h, eps = x.float(), self.eps
# normalized per (token, hc copy) over `dim`, NOT jointly over the copies
rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(key.square().mean(-1) + eps)
dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5
# signed sqrt before the sigmoid, matching the training kernel
gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot))
if token_mask is not None:
gate = gate.masked_fill(~token_mask.unsqueeze(-1), 0)
return (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype)class ParallelEngramEmbedding(nn.Module):
"""The n-gram hash table, sharded over its rows. Stays fp8: rows are dequantized on lookup."""
def __init__(self, num_embeddings: int, dim: int):
super().__init__()
self.num_embeddings = num_embeddings
self.dim = dim
self.part_num_embeddings = (num_embeddings + world_size - 1) // world_size
self.vocab_start_idx = rank * self.part_num_embeddings
self.vocab_end_idx = self.vocab_start_idx + self.part_num_embeddings
self.block_size = fp8_block_size
# the table stays fp8 as stored: rows are dequantized with `scale` on lookup
self.weight = nn.Parameter(torch.empty(self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn))
self.scale = nn.Parameter(torch.empty(self.part_num_embeddings, dim // self.block_size, dtype=scale_dtype))
def forward(self, indices: torch.Tensor) -> torch.Tensor:
mask = (indices < self.vocab_start_idx) | (indices >= self.vocab_end_idx)
local_indices = indices - self.vocab_start_idx
local_indices = local_indices.masked_fill(mask, 0)
values = F.embedding(local_indices, self.weight)
scales = F.embedding(local_indices, self.scale)
values = values.float().unflatten(-1, (-1, self.block_size)) * scales.float().unsqueeze(-1)
values = values.flatten(-2).to(torch.bfloat16)
values = values.masked_fill(mask.unsqueeze(-1), 0)
if world_size > 1:
dist.all_reduce(values)
return values @torch.inference_mode()
def forward(self, input_ids: torch.Tensor, start_pos: int, token_mask: torch.Tensor | None = None) -> torch.Tensor:
"""token_mask: [B, L], False for tokens that take no part in an n-gram (image spans).
Returns the hash ids, shaped [B, L, n_engram_layers, n_hash_cols]."""
batch, seqlen = input_ids.shape
compressed = self.token_map[input_ids]
if token_mask is not None:
compressed = torch.where(token_mask, compressed, self.DEAD)
self.cache[:batch, start_pos : start_pos + seqlen] = compressed
positions = torch.arange(start_pos, start_pos + seqlen, device=input_ids.device).expand(batch, seqlen)
tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool)
for shift in range(self.layout.max_ngram_size):
source = self.cache[:batch].gather(1, (positions - shift).clamp_min(0))
blocked = blocked | (positions < shift) | (source == self.DEAD)
tokens.append(torch.where(blocked, self.pad_id, source))
tokens = torch.stack(tokens, dim=-1) # [B, L, max_ngram_size]
# XOR the multiplied ids together one lookback at a time, so the running value after step i
# is the hash of the (i+1)-gram; each lands in its own prime-sized bucket range
products = tokens.unsqueeze(2) * self.multipliers # [B, L, n_engram_layers, max_ngram_size]
rolling, hashes = products[..., 0], []
for i in range(1, self.layout.max_ngram_size):
rolling = torch.bitwise_xor(rolling, products[..., i])
hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1])
return torch.cat(hashes, dim=-1) + self.offsets @torch.inference_mode()
def forward(
self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None
):
"""input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it
has already filled, so the padding it uses internally never reaches here. token_types /
images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie
inside the first (start_pos 0) chunk."""
image_mask = None if token_types is None else token_types >= 0 # TEXT is -1
# image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask
engram_mask = None if image_mask is None else ~image_mask
engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None
h = self.embed(input_ids)
if images is not None:
assert start_pos == 0, "image spans must be prefilled in a single chunk"
self.merge_image_embeddings(images, h)
# Expand to hc_mult copies for Hyper-Connections
h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
main_hiddens = []
pre_mix = make_identity_pre_mix(h, self.hc_mult)
for i, layer in enumerate(self.layers):
if layer.engram is not None:
h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask)
# the MTP head reads the attention input of its target layers, not their output
if i in self.target_layer_ids:
main_hiddens.append(h.mean(dim=2))
h, pre_mix = layer(h, start_pos, pre_mix, image_mask)
h = layer.hc_pre(h, pre_mix)
logits = self.head(self.norm(h))
output_ids = sample(logits, self.temperature)
main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None
return output_ids, logits, main_hidden{
"vocab_size": 129280,
"dim": 5120,
"moe_inter_dim": 2304,
"n_layers": 40,
"n_mtp_layers": 3,
"dspark_block_size": 5,
"dspark_noise_token_id": 128799,
"dspark_target_layer_ids": [37, 38, 39],
"dspark_markov_rank": 256,
"dspark_n_routed_experts": 128,
"dspark_n_activated_experts": 3,
"n_heads": 64,
"n_routed_experts": 384,
"n_shared_experts": 1,
"n_activated_experts": 6,
"score_func": "sqrtsoftplus",
"route_scale": 1.5,
"swiglu_limit": 10.0,
"q_lora_rank": 1280,
"head_dim": 512,
"rope_head_dim": 64,
"norm_eps": 1e-20,
"o_groups": 8,
"o_lora_rank": 1024,
"window_size": 128,
"kv_source_layers": [2, 8, 14, 20],
"index_source_layers": [2, 8, 14, 20, 24, 28, 32, 36],
"original_seq_len": 65536,
"rope_theta": 10000,
"rope_factor": 16,
"beta_fast": 32,
"beta_slow": 1,
"index_n_heads": 32,
"index_head_dim": 128,
"index_topk": 512,
"candidate_source_layer": 20,
"candidate_topk_blocks": 2048,
"candidate_block_size": 8,
"hc_mult": 4,
"hc_sinkhorn_iters": 20,
"hc_eps": 1e-06,
"engram_layer_ids": [1, 14],
"engram_vocab_size": 16000000,
"engram_num_embeddings": [384006168, 384016682],
"engram_max_ngram_size": 4,
"engram_pad_id": 2,
"engram_compressed_vocab_size": 99092,
"dtype": "fp8",
"expert_dtype": "fp4",
"compress_rope_theta": 160000,
"compress_ratios": [0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
"vision_n_layers": 32,
"vision_dim": 1024,
"vision_n_heads": 16,
"vision_inter_dim": 2816,
"vision_patch_size": 14,
"vision_downsample_ratio": 3,
"vision_max_n_token": 1024,
"vision_min_pixels": 295936,
"vision_max_wh_ratio": null,
"image_token_id": 129264,
"engram_n_heads": 8,
"engram_head_dim": 256,
"vision_rope_theta": 10000
}Sources5
Direct evidence
Engram reads 24 table rows per position, projects them to four keys and one value, gates the value per stream by a normalised dot product of the stream and its key, and adds it to every stream; image positions get none.
engram · deepseek-ai/DeepSeek-V4.1-Flash / inference/model.py
Pinned revision 2cba9e42aa02 ↗ · lines 328–365
Read code excerpt
class Engram(nn.Module):
"""Writes an n-gram lookup into the residual stream, gated by how well it matches that stream.
The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per hc copy plus a shared
value. The gate is a normalized dot product of stream against key.
"""
def __init__(self, args: ModelArgs, layer_id: int, layout: EngramLayout):
super().__init__()
self.layer_id = layer_id
self.layer_hash_index = layout.layer_ids.index(layer_id)
self.dim = args.dim
self.hc_mult = args.hc_mult
self.clamp_value = 1e-6
self.embed = ParallelEngramEmbedding(layout.num_embeddings[self.layer_hash_index], layout.head_dim)
n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
self.wkv = Linear(n_hash_cols * layout.head_dim, args.dim * (args.hc_mult + 1))
self.eps = args.norm_eps
self.q_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim))
self.k_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim))
def forward(self, x: torch.Tensor, hash_ids: torch.Tensor, token_mask: torch.Tensor | None = None) -> torch.Tensor:
"""x: [B, L, hc_mult, dim]; hash_ids: [B, L, n_hash_cols]; token_mask: [B, L], False shuts
the gate so those positions pass through untouched."""
kv = self.wkv(self.embed(hash_ids).flatten(-2))
key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1)
key = key.float().unflatten(-1, (self.hc_mult, self.dim))
weight = self.q_weight.float() * self.k_weight.float() # only ever used as a product
h, eps = x.float(), self.eps
# normalized per (token, hc copy) over `dim`, NOT jointly over the copies
rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(key.square().mean(-1) + eps)
dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5
# signed sqrt before the sigmoid, matching the training kernel
gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot))
if token_mask is not None:
gate = gate.masked_fill(~token_mask.unsqueeze(-1), 0)
return (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype)Direct evidence
The Engram table is an FP8 embedding with one scale per 32 features of a row, dequantised when rows are read.
engram-table · deepseek-ai/DeepSeek-V4.1-Flash / inference/model.py
Pinned revision 2cba9e42aa02 ↗ · lines 296–325
Read code excerpt
class ParallelEngramEmbedding(nn.Module):
"""The n-gram hash table, sharded over its rows. Stays fp8: rows are dequantized on lookup."""
def __init__(self, num_embeddings: int, dim: int):
super().__init__()
self.num_embeddings = num_embeddings
self.dim = dim
self.part_num_embeddings = (num_embeddings + world_size - 1) // world_size
self.vocab_start_idx = rank * self.part_num_embeddings
self.vocab_end_idx = self.vocab_start_idx + self.part_num_embeddings
self.block_size = fp8_block_size
# the table stays fp8 as stored: rows are dequantized with `scale` on lookup
self.weight = nn.Parameter(torch.empty(self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn))
self.scale = nn.Parameter(torch.empty(self.part_num_embeddings, dim // self.block_size, dtype=scale_dtype))
def forward(self, indices: torch.Tensor) -> torch.Tensor:
mask = (indices < self.vocab_start_idx) | (indices >= self.vocab_end_idx)
local_indices = indices - self.vocab_start_idx
local_indices = local_indices.masked_fill(mask, 0)
values = F.embedding(local_indices, self.weight)
scales = F.embedding(local_indices, self.scale)
values = values.float().unflatten(-1, (-1, self.block_size)) * scales.float().unsqueeze(-1)
values = values.flatten(-2).to(torch.bfloat16)
values = values.masked_fill(mask.unsqueeze(-1), 0)
if world_size > 1:
dist.all_reduce(values)
return valuesDirect evidence
Each position hashes its 2-, 3- and 4-gram of normalised token ids, each modulo eight primes, into row indices of each Engram layer’s table; an n-gram does not reach across the start of the sequence or an image position.
engram-hash · deepseek-ai/DeepSeek-V4.1-Flash / inference/engram.py
Pinned revision 2cba9e42aa02 ↗ · lines 159–184
Read code excerpt
@torch.inference_mode()
def forward(self, input_ids: torch.Tensor, start_pos: int, token_mask: torch.Tensor | None = None) -> torch.Tensor:
"""token_mask: [B, L], False for tokens that take no part in an n-gram (image spans).
Returns the hash ids, shaped [B, L, n_engram_layers, n_hash_cols]."""
batch, seqlen = input_ids.shape
compressed = self.token_map[input_ids]
if token_mask is not None:
compressed = torch.where(token_mask, compressed, self.DEAD)
self.cache[:batch, start_pos : start_pos + seqlen] = compressed
positions = torch.arange(start_pos, start_pos + seqlen, device=input_ids.device).expand(batch, seqlen)
tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool)
for shift in range(self.layout.max_ngram_size):
source = self.cache[:batch].gather(1, (positions - shift).clamp_min(0))
blocked = blocked | (positions < shift) | (source == self.DEAD)
tokens.append(torch.where(blocked, self.pad_id, source))
tokens = torch.stack(tokens, dim=-1) # [B, L, max_ngram_size]
# XOR the multiplied ids together one lookback at a time, so the running value after step i
# is the hash of the (i+1)-gram; each lands in its own prime-sized bucket range
products = tokens.unsqueeze(2) * self.multipliers # [B, L, n_engram_layers, max_ngram_size]
rolling, hashes = products[..., 0], []
for i in range(1, self.layout.max_ngram_size):
rolling = torch.bitwise_xor(rolling, products[..., i])
hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1])
return torch.cat(hashes, dim=-1) + self.offsetsDirect evidence
Tokens are embedded and repeated into four streams; each layer, after an Engram step where it has one, reads and writes the streams; the last reading mix collapses them, then the final norm and the head.
transformer-forward · deepseek-ai/DeepSeek-V4.1-Flash / inference/model.py
Pinned revision 2cba9e42aa02 ↗ · lines 1241–1272
Read code excerpt
@torch.inference_mode()
def forward(
self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None
):
"""input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it
has already filled, so the padding it uses internally never reaches here. token_types /
images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie
inside the first (start_pos 0) chunk."""
image_mask = None if token_types is None else token_types >= 0 # TEXT is -1
# image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask
engram_mask = None if image_mask is None else ~image_mask
engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None
h = self.embed(input_ids)
if images is not None:
assert start_pos == 0, "image spans must be prefilled in a single chunk"
self.merge_image_embeddings(images, h)
# Expand to hc_mult copies for Hyper-Connections
h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
main_hiddens = []
pre_mix = make_identity_pre_mix(h, self.hc_mult)
for i, layer in enumerate(self.layers):
if layer.engram is not None:
h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask)
# the MTP head reads the attention input of its target layers, not their output
if i in self.target_layer_ids:
main_hiddens.append(h.mean(dim=2))
h, pre_mix = layer(h, start_pos, pre_mix, image_mask)
h = layer.hc_pre(h, pre_mix)
logits = self.head(self.norm(h))
output_ids = sample(logits, self.temperature)
main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None
return output_ids, logits, main_hiddenDirect evidence
The reference configuration: 40 layers of width 5,120, 64 heads of 512 features with a 64-feature rotary part, a 1,280-feature query latent, a grouped output of 8 groups at rank 1,024, a window of 128, the per-layer compression ratios and the two source lists, the indexer and its candidate stage, hyper-connections of 4 streams, 384 routed experts of which 6 are used with 1 shared, Engram before layers 1 and 14, and three DSpark blocks.
config · deepseek-ai/DeepSeek-V4.1-Flash / inference/config.json
Pinned revision 2cba9e42aa02 ↗ · lines 1–66
Read code excerpt
{
"vocab_size": 129280,
"dim": 5120,
"moe_inter_dim": 2304,
"n_layers": 40,
"n_mtp_layers": 3,
"dspark_block_size": 5,
"dspark_noise_token_id": 128799,
"dspark_target_layer_ids": [37, 38, 39],
"dspark_markov_rank": 256,
"dspark_n_routed_experts": 128,
"dspark_n_activated_experts": 3,
"n_heads": 64,
"n_routed_experts": 384,
"n_shared_experts": 1,
"n_activated_experts": 6,
"score_func": "sqrtsoftplus",
"route_scale": 1.5,
"swiglu_limit": 10.0,
"q_lora_rank": 1280,
"head_dim": 512,
"rope_head_dim": 64,
"norm_eps": 1e-20,
"o_groups": 8,
"o_lora_rank": 1024,
"window_size": 128,
"kv_source_layers": [2, 8, 14, 20],
"index_source_layers": [2, 8, 14, 20, 24, 28, 32, 36],
"original_seq_len": 65536,
"rope_theta": 10000,
"rope_factor": 16,
"beta_fast": 32,
"beta_slow": 1,
"index_n_heads": 32,
"index_head_dim": 128,
"index_topk": 512,
"candidate_source_layer": 20,
"candidate_topk_blocks": 2048,
"candidate_block_size": 8,
"hc_mult": 4,
"hc_sinkhorn_iters": 20,
"hc_eps": 1e-06,
"engram_layer_ids": [1, 14],
"engram_vocab_size": 16000000,
"engram_num_embeddings": [384006168, 384016682],
"engram_max_ngram_size": 4,
"engram_pad_id": 2,
"engram_compressed_vocab_size": 99092,
"dtype": "fp8",
"expert_dtype": "fp4",
"compress_rope_theta": 160000,
"compress_ratios": [0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
"vision_n_layers": 32,
"vision_dim": 1024,
"vision_n_heads": 16,
"vision_inter_dim": 2816,
"vision_patch_size": 14,
"vision_downsample_ratio": 3,
"vision_max_n_token": 1024,
"vision_min_pixels": 295936,
"vision_max_wh_ratio": null,
"image_token_id": 129264,
"engram_n_heads": 8,
"engram_head_dim": 256,
"vision_rope_theta": 10000
}