40 layers › 2 groups of 6 layers › 5 reusing layers › Compressed sparse attention › Grouped output
Output projection
×learnedself.wo_b(o.flatten(2)) · 788
Map the 8,192 group features to the model width.
- in
- [B batch, T positions, 8192 groups × rank]
- out
- [B batch, T positions, 5120 d]
- parameters
- 41,943,040 · × 2 groups × 5 layers = 419,430,400
Parameters
41,943,040 parameters per instance.
| Component | Calculation | Parameters |
|---|---|---|
| Weight matrix | 41,943,040 | |
| Per-instance total | 41,943,040 | |
This projection has no bias parameter.
The reference module holds the weight output by input, [out, in].
Across 5 independent layers: parameters.
Derived from the configured matrix dimensions; the model total matches the checkpoint’s tensors.
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)