40 layers › 2 groups of 6 layers › 5 reusing layers › Compressed sparse attention › Grouped output
Grouped projection
×learnedtorch.einsum("bsgd,grd->bsgr", o, wo_a) · 785–787
Take the heads in 8 groups of 8 and project each group’s 4,096 features to 1,024.
- in
- [B batch, 64 heads, T positions, 512 d_h]
- out
- [B batch, T positions, 8192 groups × rank]
- parameters
- 33,554,432 · × 2 groups × 5 layers = 335,544,320
Stored as one matrix [8,192, 4,096]: a group reads only its own heads (model.py 783–784).
Parameters
33,554,432 parameters owned by this component
| Component | Calculation | Parameters |
|---|---|---|
| weights, no bias | 4,096 × (8 × 1,024) | 33,554,432 |
| Per-instance total | 33,554,432 | |
Across 5 independent layers: parameters.
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.
Codemodel.py 613–698, 765–789
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 xclass 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)Sources2
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
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)