40 layers › 6 layers › 5 reusing layers › Compressed sparse attention › Sparse attention
Join rows
×no parameterstorch.cat([kv, compress_kv], dim=1) · 777
Put the compressed keys and values after the window’s.
- in
- [B batch, 128 window slots, 512 d_h] [B batch, 500 compressed positions, 512 d_h]
- out
- [B batch, 628 rows, 512 d_h]
Codemodel.py 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 xSources1
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 x