- in
- [B batch, 1 positions, 5120 d]
- out
- [B batch, 1 positions, 5120 d]
- parameters
- 126,617,408
- mode
- decode · one new token after P = 1,000 positions
- state
- this call writes slot P mod 128, then reads the ring: 128 of 128 slots hold a position · fixed size: 65,536 values at any P
Parameters
126,617,408 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Query compression | 6,554,880 across its components | 6,554,880 |
| Query heads | 41,943,040 across its components | 41,943,040 |
| Window keys and values | 2,621,952 across its components | 2,621,952 |
| Window attention | 64 across its components | 64 |
| Grouped output | 75,497,472 across its components | 75,497,472 |
| Per-instance total | 126,617,408 | |
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
One 512-feature vector per position is key and value for all 64 heads (model.py 643; kernel.py 365, 380). Layers with no compressed positions use the rotary base 10,000 without YaRN (model.py 685–687).
Codemodel.py 613–698, 765–789 · config.json 1–66
class Attention(nn.Module):
"""Latent attention over two KV sources at once, concatenated into one `sparse_attn` call: a
sliding window of raw KV, plus -- when compress_ratio > 0 -- `index_topk` compressed positions
reaching further back. Q and the output projection are both low-rank, the latter grouped.
compress_ratio > 0 does not mean the layer compresses its own KV: only kv_source_layers do,
the rest read that same cache.
"""
def __init__(self, layer_id: int, args: ModelArgs):
super().__init__()
self.layer_id = layer_id
self.dim = args.dim
self.n_heads = args.n_heads
self.n_local_heads = args.n_heads // world_size
self.q_lora_rank = args.q_lora_rank
self.o_lora_rank = args.o_lora_rank
self.head_dim = args.head_dim
self.rope_head_dim = args.rope_head_dim
self.nope_head_dim = args.head_dim - args.rope_head_dim
self.n_groups = args.o_groups
self.n_local_groups = self.n_groups // world_size
self.window_size = args.window_size
self.compress_ratio = args.compress_ratios[layer_id]
self.eps = args.norm_eps
self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32))
self.wq_a = Linear(self.dim, self.q_lora_rank)
self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim)
self.wkv = Linear(self.dim, self.head_dim)
self.kv_norm = RMSNorm(self.head_dim, self.eps)
self.wo_a = ColumnParallelLinear(
self.n_heads * self.head_dim // self.n_groups,
self.n_groups * args.o_lora_rank,
dtype=torch.bfloat16,
)
self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim)
self.softmax_scale = self.head_dim**-0.5
is_backbone = layer_id < args.n_layers
self.is_kv_source = is_backbone and layer_id in args.kv_source_layers
self.is_index_source = is_backbone and layer_id in args.index_source_layers
self.compressor: Compressor | None = None
self.indexer: Indexer | None = None
if self.is_kv_source:
self.compressor = Compressor(args, layer_id)
if self.is_index_source:
self.indexer = Indexer(args, layer_id)
self.window_kv_cache: torch.Tensor
self.register_buffer(
"window_kv_cache",
torch.zeros(args.max_batch_size, args.window_size, self.head_dim),
persistent=False,
)
if self.is_kv_source:
self.compress_kv_cache: torch.Tensor
self.register_buffer(
"compress_kv_cache",
torch.zeros(
args.max_batch_size,
args.max_seq_len // self.compress_ratio,
self.head_dim,
),
persistent=False,
)
if self.compress_ratio:
original_seq_len, rope_theta = (
args.original_seq_len,
args.compress_rope_theta,
)
else:
# disable YaRN and use base rope_theta in pure sliding-window attention
original_seq_len, rope_theta = 0, args.rope_theta
freqs_cis = precompute_freqs_cis(
self.rope_head_dim,
args.max_seq_len,
original_seq_len,
rope_theta,
args.rope_factor,
args.beta_fast,
args.beta_slow,
)
self.freqs_cis: torch.Tensor
self.register_buffer("freqs_cis", freqs_cis, persistent=False) def forward(self, x: torch.Tensor, start_pos: int):
bsz, seqlen, _ = x.size()
freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
rd = self.rope_head_dim
qr = self.q_norm(self.wq_a(x))
q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim))
apply_rotary_emb(q[..., -rd:], freqs_cis)
kv, topk_idxs = self._window_kv(x, freqs_cis, start_pos)
if self.compress_ratio:
compress_kv, compress_idxs = self._compress_kv(x, qr, start_pos, kv.size(1))
kv = torch.cat([kv, compress_kv], dim=1)
topk_idxs = torch.cat([topk_idxs, compress_idxs], dim=-1)
o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
apply_rotary_emb(o[..., -rd:], freqs_cis, True)
# wo_a is block-diagonal over groups (each projects only its own heads), hence einsum not
# Linear. convert.py dequantizes it to bf16; an fp8 grouped GEMM would halve the memory.
o = o.view(bsz, seqlen, self.n_local_groups, -1)
wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
x = self.wo_b(o.flatten(2))
return x{
"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
}Sources3
Direct evidence
Attention owns a low-rank query, one key–value projection, a learned sink per head, a grouped low-rank output, a ring buffer for the window, and, by the two source lists, a compressor and an indexer; layers that attend to compressed positions use the larger rotary base with YaRN.
attention-init · deepseek-ai/DeepSeek-V4.1-Flash / inference/model.py
Pinned revision 2cba9e42aa02 ↗ · lines 613–698
Read code excerpt
class Attention(nn.Module):
"""Latent attention over two KV sources at once, concatenated into one `sparse_attn` call: a
sliding window of raw KV, plus -- when compress_ratio > 0 -- `index_topk` compressed positions
reaching further back. Q and the output projection are both low-rank, the latter grouped.
compress_ratio > 0 does not mean the layer compresses its own KV: only kv_source_layers do,
the rest read that same cache.
"""
def __init__(self, layer_id: int, args: ModelArgs):
super().__init__()
self.layer_id = layer_id
self.dim = args.dim
self.n_heads = args.n_heads
self.n_local_heads = args.n_heads // world_size
self.q_lora_rank = args.q_lora_rank
self.o_lora_rank = args.o_lora_rank
self.head_dim = args.head_dim
self.rope_head_dim = args.rope_head_dim
self.nope_head_dim = args.head_dim - args.rope_head_dim
self.n_groups = args.o_groups
self.n_local_groups = self.n_groups // world_size
self.window_size = args.window_size
self.compress_ratio = args.compress_ratios[layer_id]
self.eps = args.norm_eps
self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32))
self.wq_a = Linear(self.dim, self.q_lora_rank)
self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim)
self.wkv = Linear(self.dim, self.head_dim)
self.kv_norm = RMSNorm(self.head_dim, self.eps)
self.wo_a = ColumnParallelLinear(
self.n_heads * self.head_dim // self.n_groups,
self.n_groups * args.o_lora_rank,
dtype=torch.bfloat16,
)
self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim)
self.softmax_scale = self.head_dim**-0.5
is_backbone = layer_id < args.n_layers
self.is_kv_source = is_backbone and layer_id in args.kv_source_layers
self.is_index_source = is_backbone and layer_id in args.index_source_layers
self.compressor: Compressor | None = None
self.indexer: Indexer | None = None
if self.is_kv_source:
self.compressor = Compressor(args, layer_id)
if self.is_index_source:
self.indexer = Indexer(args, layer_id)
self.window_kv_cache: torch.Tensor
self.register_buffer(
"window_kv_cache",
torch.zeros(args.max_batch_size, args.window_size, self.head_dim),
persistent=False,
)
if self.is_kv_source:
self.compress_kv_cache: torch.Tensor
self.register_buffer(
"compress_kv_cache",
torch.zeros(
args.max_batch_size,
args.max_seq_len // self.compress_ratio,
self.head_dim,
),
persistent=False,
)
if self.compress_ratio:
original_seq_len, rope_theta = (
args.original_seq_len,
args.compress_rope_theta,
)
else:
# disable YaRN and use base rope_theta in pure sliding-window attention
original_seq_len, rope_theta = 0, args.rope_theta
freqs_cis = precompute_freqs_cis(
self.rope_head_dim,
args.max_seq_len,
original_seq_len,
rope_theta,
args.rope_factor,
args.beta_fast,
args.beta_slow,
)
self.freqs_cis: torch.Tensor
self.register_buffer("freqs_cis", freqs_cis, persistent=False)Direct evidence
Queries through the latent with their rotary part rotated; window and compressed keys and values joined; sparse attention over the listed positions; the output rotated back, projected per group of heads and then to the model width.
attention-forward · deepseek-ai/DeepSeek-V4.1-Flash / inference/model.py
Pinned revision 2cba9e42aa02 ↗ · lines 765–789
Read code excerpt
def forward(self, x: torch.Tensor, start_pos: int):
bsz, seqlen, _ = x.size()
freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
rd = self.rope_head_dim
qr = self.q_norm(self.wq_a(x))
q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim))
apply_rotary_emb(q[..., -rd:], freqs_cis)
kv, topk_idxs = self._window_kv(x, freqs_cis, start_pos)
if self.compress_ratio:
compress_kv, compress_idxs = self._compress_kv(x, qr, start_pos, kv.size(1))
kv = torch.cat([kv, compress_kv], dim=1)
topk_idxs = torch.cat([topk_idxs, compress_idxs], dim=-1)
o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
apply_rotary_emb(o[..., -rd:], freqs_cis, True)
# wo_a is block-diagonal over groups (each projects only its own heads), hence einsum not
# Linear. convert.py dequantizes it to bf16; an fp8 grouped GEMM would halve the memory.
o = o.view(bsz, seqlen, self.n_local_groups, -1)
wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
x = self.wo_b(o.flatten(2))
return xDirect 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
}