"""Original reduced MLA subsystem following DeepSeek-V2 equations 9-19. Four heads, content width 2, rotary width 2, value width 2, KV latent width 3, query latent width 4. Cache includes BOTH the KV latent and shared rotary key. No full DeepSeek model, weights, normalization extras, or optimized kernel. """ import torch import torch.nn as nn class PositionPair(nn.Module): def forward(self, x): position = torch.arange(x.shape[1], device=x.device, dtype=x.dtype).reshape( 1, -1, 1 ) a, b = x[..., 0], x[..., 1] return torch.stack( [ a * position.cos() - b * position.sin(), a * position.sin() + b * position.cos(), ], dim=-1, ) # tensorviz: input=1,6,8 class LatentAttention(nn.Module): def __init__(self): super().__init__() self.kv_latent = nn.Linear(8, 3, bias=False) self.query_latent = nn.Linear(8, 4, bias=False) self.query_content = nn.Linear(4, 8, bias=False) self.key_content = nn.Linear(3, 8, bias=False) self.value_content = nn.Linear(3, 8, bias=False) self.query_position = nn.Linear(4, 8, bias=False) self.key_position = nn.Linear(8, 2, bias=False) self.query_rope = PositionPair() self.key_rope = PositionPair() self.softmax = nn.Softmax(dim=-1) self.output_projection = nn.Linear(8, 8, bias=False) def forward(self, x): b, t, _ = x.shape c = self.kv_latent(x) cq = self.query_latent(x) qc = self.query_content(cq).reshape(b, t, 4, 2).transpose(1, 2) kc = self.key_content(c).reshape(b, t, 4, 2).transpose(1, 2) vc = self.value_content(c).reshape(b, t, 4, 2).transpose(1, 2) qr = self.query_rope(self.query_position(cq).reshape(b, t, 4, 2)).transpose( 1, 2 ) kr = self.key_rope(self.key_position(x).reshape(b, t, 1, 2)).transpose(1, 2) scores = (qc @ kc.transpose(-2, -1) + qr @ kr.transpose(-2, -1)) / 2 mask = torch.ones(t, t, dtype=torch.bool, device=x.device).triu(1) weights = self.softmax(scores.masked_fill(mask, float("-inf"))) heads = weights @ vc return self.output_projection(heads.transpose(1, 2).reshape(b, t, 8)) def absorbed_last(model, x): b, t, _ = x.shape # Store only c and kr for the prefix. Query activations are current-token only. c = model.kv_latent(x) kr = model.key_rope(model.key_position(x).reshape(b, t, 1, 2)).transpose(1, 2) cq = model.query_latent(x[:, -1:]) qc = model.query_content(cq).reshape(b, 1, 4, 2).transpose(1, 2) qr = model.query_position(cq).reshape(b, 1, 4, 2) position = torch.tensor(t - 1, dtype=x.dtype, device=x.device) a, z = qr[..., 0], qr[..., 1] qr = torch.stack( [ a * position.cos() - z * position.sin(), a * position.sin() + z * position.cos(), ], -1, ).transpose(1, 2) # Absorb W_UK into the content query instead of reconstructing all keys. compact_query = torch.einsum( "bhqd,hdc->bhqc", qc, model.key_content.weight.reshape(4, 2, 3) ) scores = ( compact_query @ c.unsqueeze(1).transpose(-2, -1) + qr @ kr.transpose(-2, -1) ) / 2 weights = scores.softmax(-1) latent_sum = weights @ c.unsqueeze(1) # Absorb W_UV into each head's block of W_O; never reconstruct prefix values. merged = torch.einsum( "ohd,hdc->hco", model.output_projection.weight.reshape(8, 4, 2), model.value_content.weight.reshape(4, 2, 3), ) output = torch.einsum("bhqc,hco->bqo", latent_sum, merged) return output, c, kr def verify(): torch.manual_seed(6) model = LatentAttention().double().eval() x = torch.randn(1, 6, 8, dtype=torch.float64) with torch.no_grad(): for length in range(1, 7): prefix = x[:, :length] output, c, kr = absorbed_last(model, prefix) torch.testing.assert_close( output, model(prefix)[:, -1:], rtol=1e-11, atol=1e-11 ) assert c.numel() + kr.numel() == length * 5 changed = x.clone() changed[:, 4:] += 5 torch.testing.assert_close(model(x)[:, :4], model(changed)[:, :4]) positions = model.key_position(x).reshape(1, 6, 1, 2) torch.testing.assert_close( model.key_rope(positions).norm(dim=-1), positions.norm(dim=-1) ) full = absorbed_last(model, x)[0] model.key_position.weight.zero_() assert not torch.allclose(full, absorbed_last(model, x)[0]) return [ "Compressed-cache and absorbed-weight inference matches reconstructed attention for every prefix", "Cache stores three latent and two rotary-key values per token; query latents are not cached", "Causal outputs cannot change when only future inputs change", "The decoupled rotary path preserves pair norms and contributes to this example's output", ] def experiment(): torch.manual_seed(6) model = LatentAttention().double().eval() x = torch.randn(1, 6, 8, dtype=torch.float64) cases = [] with torch.no_grad(): for length in [1, 3, 6]: output, c, kr = absorbed_last(model, x[:, :length]) reference = model(x[:, :length])[:, -1:] cases.append( { "id": f"prefix-{length}", "label": f"{length} cached token" + ("s" if length > 1 else ""), "target": "kv_latent", "note": "The cache retains the joint KV latent and a separate rotated key. The content-key up-projection is absorbed into the query calculation; the value up-projection is absorbed into the output projection. The positional path stays separate because token-dependent rotations cannot be freely reordered with learned matrices.", "vectors": [ { "label": "KV latent · final cached token", "values": c[0, -1].tolist(), }, { "label": "Rotary key · final cached token", "values": kr[0, 0, -1].tolist(), }, { "label": "Reconstructed attention output", "values": reference[0, 0].tolist(), }, { "label": "Absorbed-weight output", "values": output[0, 0].tolist(), }, ], "metrics": [ {"label": "Compact cache · scalar values", "value": length * 5}, { "label": "Expanded content K/V + shared rotary key · scalar values", "value": length * 18, }, { "label": "Maximum output difference", "value": (output - reference).abs().max().item(), }, ], } ) return { "kind": "vectors", "title": "Cache a latent plus the positional key.", "description": "Recorded decoding comparisons for 1, 3 and 6 prefix tokens. Expanded storage counts 8 content-key + 8 value + 2 shared rotary-key values per token; the compact cache needs 3 + 2. These teaching dimensions are not DeepSeek-V2's production dimensions.", "controlLabel": "Prefix length", "cases": cases, }