DeepSeek-V3 & R1 Architecture: Why Multi-Head Latent Attention (MLA) Broke Silicon Valley

import torch
import torch.nn as nn

class MultiHeadLatentAttention(nn.Module):
    def __init__(self, d_model=4096, num_heads=64, d_latent_kv=512, d_head=128):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_latent_kv = d_latent_kv
        self.d_head = d_head
        
        # Down-projection to compressed latent vector
        self.W_down_kv = nn.Linear(d_model, d_latent_kv, bias=False)
        
        # Up-projections for keys and values
        self.W_up_k = nn.Linear(d_latent_kv, num_heads * d_head, bias=False)
        self.W_up_v = nn.Linear(d_latent_kv, num_heads * d_head, bias=False)
        self.W_q = nn.Linear(d_model, num_heads * d_head, bias=False)
        self.W_out = nn.Linear(num_heads * d_head, d_model, bias=False)

    def forward(self, x, kv_cache=None):
        B, S, D = x.shape
        # Compress token representation to low-rank latent
        latent_kv = self.W_down_kv(x) # [B, S, d_latent_kv]
        
        # Keys and Values decompressed on-the-fly
        k = self.W_up_k(latent_kv).view(B, S, self.num_heads, self.d_head)
        v = self.W_up_v(latent_kv).view(B, S, self.num_heads, self.d_head)
        q = self.W_q(x).view(B, S, self.num_heads, self.d_head)
        
        # Standard scaled dot-product attention
        scores = torch.einsum('bshd,bthd->bsht', q, k) / (self.d_head ** 0.5)
        attn = torch.softmax(scores, dim=-1)
        out = torch.einsum('bsht,bthd->bshd', attn, v)
        return self.W_out(out.reshape(B, S, -1))