When V4 dropped on April 24, 2026, the community reaction was immediate. Reddit flagged its Codeforces rating of 3,206 as the highest any AI had ever scored. Engineers noticed the 73% reduction in per-token FLOPs compared to V3.2. And everyone noticed the price tag.

DeepSeek V4’s Flash model costs $0.14 per million input tokens. OpenAI’s GPT-5.5 charges $5.00 for the same work. That’s 35x cheaper, and yet V4 hangs with frontier models on most benchmarks.

But the really interesting stuff isn’t in the benchmarks or the press coverage. It’s in the architecture. DeepSeek didn’t just scale up V3. They rebuilt the attention mechanism, redesigned how residual connections work in deep networks, and introduced a hybrid compression scheme that makes million-token contexts tractable.

Multi-Head Latent Attention (MLA): Compressing the Memory Bottleneck

The Problem: KV Cache Is a Memory Hog

In a standard transformer, every time the model generates a token, it needs to look back at all previous tokens. To avoid recomputing everything from scratch, the model stores Key and Value matrices for every past token in what’s called the KV cache. This is what lets autoregressive generation be fast.

The problem? The KV cache grows linearly with sequence length, and it’s per layer, per head. For a model with 16 attention heads each of dimension 128, the cache stores:

Cache per token = 2 × 16 × 128 = 4,096 elements

At a million tokens, that’s billions of floating-point numbers just sitting in GPU memory. This is the bottleneck that makes long-context inference ruinously expensive.

Grouped-Query Attention (GQA), used by models like LLaMA, tackles this by reducing the number of KV heads sharing one set of keys and values across multiple query heads. It helps, but it’s a blunt instrument: you’re throwing away capacity to save memory.

The Idea: Compress, Don’t Discard

Think of it like taking notes in a lecture. GQA is like sharing one person’s notes across four students — it’s cheaper, but everyone sees the same thing. MLA is more like each student having a tiny personal summary card that they can expand back into full notes on demand. The card takes up almost no space, but the information is still there.

That’s exactly what Multi-Head Latent Attention (MLA) does. Instead of caching the full K and V matrices, it caches a single compressed latent vector. When the model needs actual keys and values, it reconstructs them from this latent on the fly.

Cache per token = kv_compress_dim + rope_head_dim

For DeepSeek-V2’s 236B model: 512 elements vs. 12,288 for standard MHA. That’s a 24× reduction.

The Code: RoPE Helpers

Before we build MLA itself, we need a quick helper. Rotary Position Embeddings (RoPE) encode position information into queries and keys via rotation in complex number space. If you’ve seen RoPE before, the key thing to know here is that MLA needs a version that works on arbitrary dimensions, not just the standard head dimension because the positional encoding lives in a separate branch.

def precompute_rope_frequencies(
    dim: int, max_len: int, theta: float = 10000.0, device: str = "cpu"
) -> torch.Tensor:
    """
    Precompute complex-valued RoPE frequency tensor.
    Returns: freqs_complex: (max_len, dim // 2) complex tensor
    """
    # Compute base frequencies for each pair dimension
    # (dim // 2,)
    freq_indices = torch.arange(0, dim, 2, device=device).float()
    freqs = 1.0 / (theta ** (freq_indices / dim))
    # Position indices
    # (max_len,)
    positions = torch.arange(max_len, device=device).float()
    # Outer product: position × frequency
    # (max_len,) × (dim // 2,) → (max_len, dim // 2)
    angles = torch.outer(positions, freqs)
    # Convert to complex exponential: e^{i·angle}
    # (max_len, dim // 2) → (max_len, dim // 2) complex
    return torch.polar(torch.ones_like(angles), angles)

This creates a table of complex exponentials, one per position and frequency pair. Each pair of dimensions in the input gets rotated by an angle that depends on position. The outer product between positions and freqs gives us a (max_len, dim // 2) grid of angles, and torch.polar turns those into complex numbers of unit magnitude.

Now the function that applies these rotations:

def apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
    """Apply RoPE to the last dimension of x.  
    x may have shape (..., rope_dim). We treat the last dim as pairs.
    Args:
        x:     (..., rope_dim)  - any leading dims
        freqs: (seq_len, rope_dim // 2) complex
    Returns:
        (..., rope_dim)
    """
    orig_shape = x.shape
    seq_len = x.shape[-3] if x.dim() >= 3 else x.shape[0]
    # Reshape last dim into pairs → view as complex
    # (..., rope_dim) → (..., rope_dim // 2, 2) → (..., rope_dim // 2) complex
    x_pairs = x.float().reshape(*x.shape[:-1], -1, 2)
    x_complex = torch.view_as_complex(x_pairs)
    # Broadcast freqs to match leading dims
    f = freqs[:seq_len]
    shape = [1] * x_complex.dim()
    if x_complex.dim() >= 3:
        shape[-3] = seq_len
    else:
        shape[0] = seq_len
    shape[-1] = f.shape[-1]
    f = f.view(*shape)
    # Rotate via complex multiplication
    x_rotated = x_complex * f
    # Convert back to real
    # (..., rope_dim // 2) complex → (..., rope_dim // 2, 2) → (..., rope_dim)
    return torch.view_as_real(x_rotated).reshape(orig_shape).type_as(x)

The trick: pair up consecutive dimensions, treat each pair as a complex number, multiply by the frequency (a rotation), then convert back. It’s elegant and cheap.

The Code: Building MLA Piece by Piece

This is the core innovation: The KV compression path

# Down-project hidden state into compressed KV latent
# (batch, seq_len, model_dim) → (batch, seq_len, kv_compress_dim)
self.w_dkv = nn.Linear(model_dim, kv_compress_dim, bias=False)
# Up-project latent to content keys (all heads)
# (batch, seq_len, kv_compress_dim) → (batch, seq_len, num_heads * head_dim)
self.w_uk = nn.Linear(kv_compress_dim, num_heads * head_dim, bias=False)
# Up-project latent to values (all heads)
self.w_uv = nn.Linear(kv_compress_dim, num_heads * head_dim, bias=False)

This is MLA’s compression trick in three lines. w_dkv squeezes the hidden state down to a small latent vector, this is what actually gets cached. w_uk and w_uv reconstruct full keys and values from that latent when needed. It's essentially a bottleneck autoencoder: compress, cache the bottleneck, decompress later.

The decoupled RoPE branch: Here’s a subtle but critical detail:

# Decoupled RoPE key branch (shared across all heads)
# (batch, seq_len, model_dim) → (batch, seq_len, rope_head_dim)
self.w_kr = nn.Linear(model_dim, rope_head_dim, bias=False)

Why is RoPE handled separately? Because RoPE is position-dependent, it encodes “where” a token is in the sequence. The compressed latent c_kv is meant to be content-only (what the token says, not where it is). If you mix position information into the compressed latent, you can't do the low-rank compression trick cleanly. So MLA keeps a tiny separate RoPE branch that goes straight from the hidden state to a small positional key and caches that alongside the compressed latent. What MLA actually caches per token:

content information: c_kv (kv_compress_dim elements)
position information: k_rope (rope_head_dim elements)

The query path mirrors this structure:

# Query compression
self.w_dq = nn.Linear(model_dim, q_compress_dim, bias=False)
self.w_uq = nn.Linear(q_compress_dim, num_heads * head_dim, bias=False)
# Decoupled RoPE query (per head, unlike the shared key RoPE)
self.w_qr = nn.Linear(q_compress_dim, num_heads * rope_head_dim, bias=False)
# Output projection
self.w_o = nn.Linear(num_heads * head_dim, model_dim, bias=False)

Notice that query RoPE is per head (projects to num_heads * rope_head_dim) while key RoPE is shared across heads (projects to just rope_head_dim and gets broadcast). This asymmetry saves memory , the key RoPE is the part that gets cached, so keeping it small matters.

The Forward Pass

Now let’s walk through the forward pass. Here’s the data flow visualized:

h_t  (hidden state at position t)
          │
    ┌─────┼──────────────┐
    │     │              │
  W_DKV  W_DQ          W_KR (decoupled RoPE key)
    │     │              │
  c_kv   c_q         k_rope ← RoPE(·)
  (CACHED)│          (CACHED)
    │     │
 ┌──┴──┐  W_UQ
 W_UK  W_UV │
 │     │   q_c
 k_c   v    │
 │          W_QR → q_rope ← RoPE(·)
 │          │
 k=[k_c;k_rope]    q=[q_c;q_rope]
          │
      Attention(q, k, v)
          │
        W_O → output

Let’s zoom into the forward method’s key sections.

KV compression and reconstruction:

# Compress hidden state into latent
# (batch, seq, model_dim) → (batch, seq, kv_compress_dim)
c_kv = self.w_dkv(x)  
# Reconstruct content keys and values from the compressed latent
k_content = self.w_uk(c_kv).view(batch_num, seq_len, self.num_heads, self.head_dim)
v = self.w_uv(c_kv).view(batch_num, seq_len, self.num_heads, self.head_dim)

Decoupled RoPE for keys:

# Project to small RoPE key space (shared across heads)
# (batch, seq, rope_head_dim)
k_rope = self.w_kr(x)
# (batch, seq, 1, rope_head_dim)                                      
k_rope = k_rope.unsqueeze(2)                                
k_rope = apply_rope(k_rope, rope_freqs)
# broadcast to all heads                     
k_rope = k_rope.expand(-1, -1, self.num_heads, -1)          

The unsqueeze(2) and expand pattern is worth pausing on. We compute RoPE for keys only once (shared across heads), then broadcast. This is why the key RoPE cache is rope_head_dim total, not num_heads * rope_head_dim.

Query compression + decoupled RoPE:

# Compress queries
# (batch, seq, q_compress_dim)
c_q = self.w_dq(x)                    
# Content queries
q_content = self.w_uq(c_q).view(batch_num, seq_len, self.num_heads, self.head_dim)
# RoPE queries (per head - NOT shared like key RoPE)
q_rope = self.w_qr(c_q).view(batch_num, seq_len, self.num_heads, self.rope_head_dim)
q_rope = apply_rope(q_rope, rope_freqs)

Concatenate content + positional parts:

# (batch, seq, heads, head_dim + rope_dim)
q = torch.cat([q_content, q_rope], dim=-1)  
# (batch, seq, heads, head_dim + rope_dim)
k = torch.cat([k_content, k_rope], dim=-1)  

This is where the two streams reunite. The final query and key each have a content part (from the compressed latent) and a positional part (from RoPE). When you compute q @ k^T, both content similarity and positional proximity contribute to the attention score.

Scaled dot-product attention and a critical subtlety about values:

# Transpose to (batch, num_heads, seq_len, attn_dim)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
# V stays at head_dim (NOT attn_dim) - the RoPE part only
# participates in the score, not the value aggregation
v = v.transpose(1, 2)
# (batch, heads, seq, attn_dim) @ (batch, heads, attn_dim, seq) → (batch, heads, seq, seq)
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.attn_dim)
if mask is not None:
    scores = scores.masked_fill(mask == 0, float("-inf"))
attn_weights = F.softmax(scores, dim=-1)
# (batch, heads, seq, seq) @ (batch, heads, seq, head_dim) → (batch, heads, seq, head_dim)
context = torch.matmul(attn_weights, v)
# Merge heads: (batch, heads, seq, head_dim) → (batch, seq, num_heads * head_dim)
context = context.transpose(1, 2).contiguous().view(
    batch_num, seq_len, self.num_heads * self.head_dim
)
# Final projection
return self.w_o(context)

The RoPE dimensions participate in scoring (which tokens attend to which) but not in value aggregation (what information actually gets passed forward). This makes sense: position tells you where to look, but the content you retrieve doesn’t need position baked into it.

Here’s the Complete MLA Class

Now that we understand every piece, here’s the full assembled version:

class MultiHeadLatentAttention(nn.Module):
    """Multi-head Latent Attention (MLA) from DeepSeek-V2.  
    Key design decisions:
      * Joint KV compression via a single down-projection into c_kv.
        Only c_kv and k_rope need to be cached - dramatically smaller
        than caching full K, V matrices.
      * Decoupled RoPE: position-dependent rotary embeddings are applied
        on a separate, small branch so that the low-rank reconstruction
        of content K remains position-agnostic.
      * Optional query compression further reduces activation memory
        during training.
    """
    def __init__(
        self,
        model_dim: int,
        num_heads: int,
        head_dim: int,
        kv_compress_dim: int,
        q_compress_dim: int,
        rope_head_dim: int,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.kv_compress_dim = kv_compress_dim
        self.rope_head_dim = rope_head_dim
        self.attn_dim = head_dim + rope_head_dim
        # --- KV compression path ---
        self.w_dkv = nn.Linear(model_dim, kv_compress_dim, bias=False)
        self.w_uk = nn.Linear(kv_compress_dim, num_heads * head_dim, bias=False)
        self.w_uv = nn.Linear(kv_compress_dim, num_heads * head_dim, bias=False)
                # --- Decoupled RoPE key branch (shared across all heads) ---
        self.w_kr = nn.Linear(model_dim, rope_head_dim, bias=False)
        # --- Query compression path ---
        self.w_dq = nn.Linear(model_dim, q_compress_dim, bias=False)
        self.w_uq = nn.Linear(q_compress_dim, num_heads * head_dim, bias=False)
        self.w_qr = nn.Linear(q_compress_dim, num_heads * rope_head_dim, bias=False)
                # --- Output projection ---
        self.w_o = nn.Linear(num_heads * head_dim, model_dim, bias=False)
        def forward(self, x, rope_freqs, mask=None):
        batch_num, seq_len, _ = x.shape
        # ── KV compression ──
        c_kv = self.w_dkv(x)
        k_content = self.w_uk(c_kv).view(batch_num, seq_len, self.num_heads, self.head_dim)
        v = self.w_uv(c_kv).view(batch_num, seq_len, self.num_heads, self.head_dim)
                # ── Decoupled RoPE for keys (shared across heads) ──
        k_rope = self.w_kr(x)
        k_rope = k_rope.unsqueeze(2)
        k_rope = apply_rope(k_rope, rope_freqs)
        k_rope = k_rope.expand(-1, -1, self.num_heads, -1)
                # ── Query compression + decoupled RoPE ──
        c_q = self.w_dq(x)
        q_content = self.w_uq(c_q).view(batch_num, seq_len, self.num_heads, self.head_dim)
        q_rope = self.w_qr(c_q).view(
            batch_num, seq_len, self.num_heads, self.rope_head_dim
        )
        q_rope = apply_rope(q_rope, rope_freqs)
        # ── Concatenate content + RoPE parts ──
        q = torch.cat([q_content, q_rope], dim=-1)
        k = torch.cat([k_content, k_rope], dim=-1)
        # ── Scaled dot-product attention ──
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.attn_dim)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        attn_weights = F.softmax(scores, dim=-1)
        context = torch.matmul(attn_weights, v)
        # Merge heads
        context = context.transpose(1, 2).contiguous().view(
            batch_num, seq_len, self.num_heads * self.head_dim
        )
        return self.w_o(context)

Verifying It Works: Cache Size Comparison

# ── Test MLA and compare cache sizes ──
model_dim = 2048
num_heads = 16
head_dim = 128
kv_compress_dim = 512
q_compress_dim = 1536
rope_head_dim = 64
max_len = 128
mla = MultiHeadLatentAttention(
    model_dim=model_dim,
    num_heads=num_heads,
    head_dim=head_dim,
    kv_compress_dim=kv_compress_dim,
    q_compress_dim=q_compress_dim,
    rope_head_dim=rope_head_dim,
)
rope_freqs = precompute_rope_frequencies(rope_head_dim, max_len)
x = torch.randn(2, 64, model_dim)
causal_mask = torch.tril(torch.ones(64, 64)).unsqueeze(0).unsqueeze(0)
out = mla(x, rope_freqs, mask=causal_mask)
print(f"MLA output shape: {out.shape}")  # (2, 64, 2048) ✓
# ── KV cache size comparison ──
mha_cache_per_token = 2 * num_heads * head_dim          # = 4,096
gqa_kv_heads = 4
gqa_cache_per_token = 2 * gqa_kv_heads * head_dim       # = 1,024
mla_cache_per_token = kv_compress_dim + rope_head_dim    # = 576
print(f"KV cache elements per token:")
print(f"  Standard MHA : {mha_cache_per_token:>6}")      # 4,096
print(f"  GQA (4 heads): {gqa_cache_per_token:>6}")      # 1,024
print(f"  MLA          : {mla_cache_per_token:>6}")       # 576
print(f"  MLA reduction vs MHA: {mha_cache_per_token / mla_cache_per_token:.1f}×")  # 7.1×

That’s the entire trick: MLA compresses the KV cache into a bottleneck that’s 7× smaller than standard MHA and nearly 2× smaller than GQA — while outperforming GQA on quality benchmarks. DeepSeek’s V2 ablations showed that MLA even slightly exceeds full MHA performance.

Mixture of Expert (MoE): When Not Every Neuron Needs to Fire

The Problem: Dense Models Are Wasteful

Standard transformers are dense: every parameter participates in every forward pass. If you want a smarter model, you add more parameters, and all of them activate for every token. That’s like hiring 100 specialists but making all of them attend every single meeting.

The Idea: Route Tokens to Specialists

Think of it like a university’s tutoring center. Instead of one enormous tutor who knows everything (dense model), you have a front desk that reads your question and sends you to the 2–3 tutors who are actually experts in that topic. The other tutors don’t even know you came in. That’s Mixture of Experts (MoE).

DeepSeek’s take on MoE has three twists that separate it from earlier designs like Switch Transformer:

1. Fine-grained experts: Instead of 8–16 large experts, DeepSeek uses 64+ small experts. Finer granularity = better specialization.
2. Shared experts: Some knowledge (basic syntax, common phrases) is needed by every token. A few experts are “always on” to capture this, freeing the routed experts to truly specialize.
3. Auxiliary-loss-free load balancing: Traditional MoE adds a loss term to push tokens toward underused experts. DeepSeek V3 replaced this with a learnable bias, simpler and doesn’t degrade model quality.

The Code: A Single Expert

Each expert is a small SwiGLU feed-forward network:

class ExpertFFN(nn.Module):
    """Single fine-grained expert: a SwiGLU feed-forward network.  
    Compared to standard Transformer FFN (4 × model_dim intermediate),
    each fine-grained expert uses a much smaller intermediate dimension
    (expert_dim), since many experts are activated in parallel.
    """
    def __init__(self, model_dim: int, expert_dim: int):
        super().__init__()
        self.w_gate = nn.Linear(model_dim, expert_dim, bias=False)
        self.w_up = nn.Linear(model_dim, expert_dim, bias=False)
        self.w_down = nn.Linear(expert_dim, model_dim, bias=False)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU: silu(gate) * up → down
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

Why three linear layers instead of two? SwiGLU uses a gated activation: silu(gate) * up. The gate projection controls how much of the up projection passes through. This consistently outperforms the classic two-layer ReLU FFN. The w_down then projects back to model dimension.

Note that expert_dim is much smaller than what you'd use in a dense FFN (which typically uses 4 × model_dim). Each individual expert is narrow, the aggregate capacity comes from having many of them.

The Code: The Router

This is where the magic of “auxiliary-loss-free balancing” lives:

class TopKRouter(nn.Module):
    """Token-level top-K router with auxiliary-loss-free load balancing.   
    Design decisions:
      * A bias vector is added to routing logits for expert SELECTION only.
      * The softmax weights used to COMBINE expert outputs are computed
        from the ORIGINAL (unbiased) logits. This prevents the bias from
        distorting the learned representations.
      * The bias is updated via a simple heuristic (not gradient descent):
        increase bias for underloaded experts, decrease for overloaded.
    """
    def __init__(self, model_dim: int, num_experts: int, top_k: int):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        # Routing projection: token → expert scores
        self.gate = nn.Linear(model_dim, num_experts, bias=False)
        # Auxiliary-loss-free bias (NOT trained by gradient; updated heuristically)
        self.expert_bias = nn.Parameter(
            torch.zeros(num_experts), requires_grad=False
        )
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Route tokens to top-K experts.
        Args:
            x: (num_tokens, model_dim)
        Returns:
            weights: (num_tokens, top_k) - normalised expert weights
            indices: (num_tokens, top_k) - selected expert indices
        """
        # Compute raw routing logits
        logits = self.gate(x)                          # (num_tokens, num_experts)
        # Biased logits for SELECTION (bias helps balance load)
        routing_logits = logits + self.expert_bias     # (num_tokens, num_experts)
        # Select top-K experts per token
        _, indices = torch.topk(routing_logits, self.top_k, dim=-1)
        # Compute combination weights from ORIGINAL (unbiased) logits
        selected_logits = logits.gather(-1, indices)   # (num_tokens, top_k)
        weights = F.softmax(selected_logits, dim=-1)   # (num_tokens, top_k)
        return weights, indices

The router is literally just a linear layer with a softmax. That’s it. A token goes in, expert scores come out. But focus on expert_bias, this is the clever part.

The bias affects which experts get selected (topk uses biased logits), but the combination weights are computed from the original, unbiased logits. Why separate selection from weighting? If the bias influenced the softmax weights, it would distort the model's learned representations. You'd be telling the model "pretend you like this expert more than you actually do." By only affecting the selection threshold, the bias nudges load balancing without polluting the learned expert affinities. It's a clean separation of concerns.

The Code: The Full MoE Layer

Now let’s assemble the complete MoE layer with shared and routed experts:

class DeepSeekMoE(nn.Module):
    """DeepSeek Mixture-of-Experts layer.  
    Three key components:
      1. Shared experts - always activated for common knowledge.
      2. Routed (fine-grained) experts - dynamically selected per token.
      3. Top-K router with auxiliary-loss-free bias balancing.
    """
    def __init__(
        self,
        model_dim: int,
        num_shared_experts: int,
        num_routed_experts: int,
        num_active_experts: int,
        expert_dim: int,
    ):
        super().__init__()
        self.num_shared = num_shared_experts
        self.num_routed = num_routed_experts
        self.num_active = num_active_experts
        # Shared experts: always contribute to every token
        self.shared_experts = nn.ModuleList(
            [ExpertFFN(model_dim, expert_dim) for _ in range(num_shared_experts)]
        )
        # Routed (fine-grained) experts: dynamically selected
        self.routed_experts = nn.ModuleList(
            [ExpertFFN(model_dim, expert_dim) for _ in range(num_routed_experts)]
        )
        # Token-level router
        self.router = TopKRouter(model_dim, num_routed_experts, num_active_experts)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_num, seq_len, model_dim = x.shape
        # ── Shared experts (always activated) ──
        # Sum outputs from all shared experts
        shared_output = sum(expert(x) for expert in self.shared_experts)
        # ── Routed experts ──
        # Flatten batch and sequence dims for routing
        num_tokens = batch_num * seq_len
        flat_x = x.view(num_tokens, model_dim)
        # Get top-K expert assignments
        weights, indices = self.router(flat_x)
        # Accumulate weighted expert outputs
        routed_output = torch.zeros_like(flat_x)
        # Iterate over each active expert slot
        for k in range(self.num_active):
            # Expert indices for slot k: (num_tokens,)
            expert_ids = indices[:, k]
            # Corresponding weights: (num_tokens, 1)
            expert_weights = weights[:, k].unsqueeze(-1)
            # Process tokens assigned to each expert
            for e_idx in range(self.num_routed):
                # Boolean mask of tokens routed to expert e_idx in slot k
                token_mask = expert_ids == e_idx
                if not token_mask.any():
                    continue
                # Gather tokens for this expert, run FFN, weighted addition
                expert_input = flat_x[token_mask]
                expert_out = self.routed_experts[e_idx](expert_input)
                routed_output[token_mask] += expert_weights[token_mask] * expert_out
                # Reshape back and combine shared + routed
        routed_output = routed_output.view(batch_num, seq_len, model_dim)
        return shared_output + routed_output

The flattening step is important: routing happens per token, not per sequence. Every token independently chooses its experts. The nested loop is the naive implementation (fine for learning — production uses batched kernels). For each of the top_k expert slots, it finds which tokens chose each expert, runs them through that expert, and accumulates the weighted output.

Verifying It Works: Efficiency Numbers

moe = DeepSeekMoE(
    model_dim=2048,
    num_shared_experts=2,
    num_routed_experts=16,
    num_active_experts=4,
    expert_dim=1024,
)
x = torch.randn(2, 64, 2048)
out = moe(x)
print(f"MoE output shape: {out.shape}")  # (2, 64, 2048) ✓
total_params = sum(p.numel() for p in moe.parameters())
print(f"Total MoE parameters: {total_params / 1e6:.1f}M")
# Per forward pass, only shared + top-K routed experts are activated
shared_params = sum(p.numel() for e in moe.shared_experts for p in e.parameters())
per_expert_params = sum(p.numel() for p in moe.routed_experts[0].parameters())
active_params = shared_params + 4 * per_expert_params
print(f"Activated parameters per token: {active_params / 1e6:.1f}M / {total_params / 1e6:.1f}M")

The model “knows” ~113M parameters worth of FFN knowledge, but each token only uses ~38M of it. Scale that up: V4-Pro has 1.6T total, 49B active. Same principle.

Multi-Token Prediction (MTP): Training on the Future

The Problem: One Token at a Time Is Slow

Standard language models predict one next token per position. That’s fine, but it wastes information — why not also predict the token two steps ahead? And three?

The Idea: Study Groups, Not Solo Study

Think about how studying works. If you only quiz yourself on the very next flashcard, you learn the immediate pattern. But if you also try to predict two and three cards ahead, you build deeper understanding of the overall structure. MTP works the same way.

During training, each position predicts not just the next token (standard), but also tokens 2, 3, … k+1 steps into the future. This gives the model a richer gradient signal. During inference, these extra heads enable speculative decoding , the model drafts multiple future tokens in parallel, then verifies them.

The Code: A Single MTP Head

Each MTP head sits at a specific “depth” — depth 1 predicts t+2, depth 2 predicts t+3, etc.

class MTPHead(nn.Module):
    """
    Single Multi-Token Prediction head.
    Each MTP head receives the previous head's hidden state (or the main
    Transformer output for head 0), concatenated with the embedding of the
    predicted token from the previous depth, then processes it through a
    small Transformer layer to predict a token further into the future.
    """
    def __init__(self, model_dim: int, vocab_size: int, num_heads: int = 4):
        super().__init__()
        # Project concatenated [hidden_state; prev_embedding] back to model_dim
        self.proj = nn.Linear(2 * model_dim, model_dim, bias=False)
        # Lightweight self-attention layer for this prediction depth
        self.norm = nn.RMSNorm(model_dim)
        self.attn = nn.MultiheadAttention(
            model_dim, num_heads, batch_first=True
        )
        # Vocabulary projection
        self.head = nn.Linear(model_dim, vocab_size, bias=False)
    def forward(self, hidden, prev_embed, causal_mask):
        """
        Args:
            hidden:      (batch, seq_len, model_dim)
            prev_embed:  (batch, seq_len, model_dim)
            causal_mask: (seq_len, seq_len)
        Returns:
            logits:      (batch, seq_len, vocab_size)
            new_hidden:  (batch, seq_len, model_dim)
        """
        # Concatenate hidden state with previous depth's embedding
        combined = torch.cat([hidden, prev_embed], dim=-1)
        # Project back to model_dim
        h = self.proj(combined)
        # Self-attention with causal mask
        h_norm = self.norm(h)
        attn_out, _ = self.attn(
            h_norm, h_norm, h_norm, attn_mask=causal_mask, is_causal=True
        )
        new_hidden = h + attn_out
        # Predict token at this depth
        logits = self.head(self.norm(new_hidden))
        return logits, new_hidden

Why concatenate with the previous embedding? This creates a sequential chain: head 0 predicts token t+1, head 1 takes head 0’s prediction as context to predict t+2, and so on. Each depth “knows” what the earlier depths predicted.

The Full MTP Module

class MultiTokenPredictor(nn.Module):
    """Multi-Token Prediction module with D extra prediction depths."""
    def __init__(self, model_dim, vocab_size, num_depths=2, num_heads=4):
        super().__init__()
        self.num_depths = num_depths
        # Shared embedding (same as main model)
        self.embed = nn.Embedding(vocab_size, model_dim)
        # Main LM head (depth 0): predict next token
        self.main_head = nn.Linear(model_dim, vocab_size, bias=False)
        # Extra MTP heads (depths 1..D): predict tokens further ahead
        self.mtp_heads = nn.ModuleList(
            [MTPHead(model_dim, vocab_size, num_heads) for _ in range(num_depths)]
        )
    def forward(self, hidden, target_ids):
        """Compute combined MTP loss.
        Args:
            hidden:     (batch, seq_len, model_dim) - Transformer output
            target_ids: (batch, seq_len) - ground-truth token IDs
        Returns:
            loss: scalar - weighted sum of per-depth cross-entropy losses
        """
        batch_num, seq_len, model_dim = hidden.shape
        # Causal mask for internal attention in MTP heads
        causal_mask = torch.triu(
            torch.full((seq_len, seq_len), float("-inf"), device=hidden.device),
            diagonal=1,
        )
        # ── Depth 0: standard next-token prediction ──
        logits_0 = self.main_head(hidden)
        loss_0 = F.cross_entropy(
            logits_0[:, :-1].reshape(-1, logits_0.size(-1)),
            target_ids[:, 1:].reshape(-1),
            ignore_index=0,
        )
        total_loss = loss_0
        # ── Depths 1..D: predict further future tokens ──
        current_hidden = hidden
        prev_embed = self.embed(target_ids)
        for d, mtp_head in enumerate(self.mtp_heads):
            depth = d + 1
            # MTP head produces logits and updated hidden state
            logits_d, current_hidden = mtp_head(
                current_hidden, prev_embed, causal_mask
            )
            # Target for depth d: token at position t + d + 1
            shift = depth + 1
            if shift >= seq_len:
                break
            loss_d = F.cross_entropy(
                logits_d[:, :-shift].reshape(-1, logits_d.size(-1)),
                target_ids[:, shift:].reshape(-1),
                ignore_index=0,
            )
            # Weight deeper predictions less (simple linear decay)
            weight = 1.0 / (depth + 1)
            total_loss = total_loss + weight * loss_d
            # Prepare prev_embed for next depth using ground-truth token
            if depth < self.num_depths:
                prev_embed = self.embed(target_ids.roll(-depth, dims=1))
        return total_loss

The 1.0 / (depth + 1) weighting is a design choice: predicting two tokens ahead is harder and noisier than predicting one, so it contributes less to the loss. The total loss is a weighted sum across all depths.

Manifold-Constrained Hyper-Connections (mHC): The Residual Connection, Reinvented

The Problem: One Stream Isn’t Enough for 100+ Layers

This is where DeepSeek V4 makes its most mathematically interesting contribution.

Standard residual connections are beautifully simple:

The hidden state gets one path forward, and each layer adds to it. But there’s a bottleneck: everything has to share that one stream. At 100+ layers deep, this single-stream residual becomes a limiting factor. Some earlier DeepSeek experiments tried Hyper-Connections — multiple parallel streams that exchange information at each layer. The idea was good, but unconstrained mixing led to signal explosion exceeding 3,000× in a 27B model, causing training to diverge catastrophically.

The Idea: Traffic Signals for Information Flow

Think about a multi-lane highway. Without rules, cars constantly merge into each other’s lanes and some lanes get exponentially more traffic while others empty out — chaos. A doubly stochastic mixing matrix is like a traffic system where the total flow into every lane equals 1 and the total flow out of every lane equals 1. Traffic gets redistributed but never amplified. No lane overflows, no lane starves.

That’s what manifold-constrained hyper-connections (mHC) do. The mixing matrix lives on the Birkhoff Polytope, the set of doubly stochastic matrices (non-negative entries, rows and columns each sum to 1). This mathematically guarantees that signal magnitude is preserved across layers, no matter how deep.

The Code: Sinkhorn-Knopp Algorithm

How do you force a learnable matrix to be doubly stochastic? With the Sinkhorn-Knopp algorithm: alternately normalize rows and columns until convergence.

def sinkhorn_knopp(
    log_weights: torch.Tensor, num_iters: int = 5, eps: float = 1e-6
) -> torch.Tensor:
    """Project a raw weight matrix onto the Birkhoff Polytope via Sinkhorn-Knopp. 
    We operate in log-space for numerical stability: start from log_weights,
    exponentiate to get a non-negative matrix, then alternately normalise
    rows and columns.
    Args:
        log_weights: (num_streams, num_streams) - unconstrained learnable params
        num_iters:   number of alternating normalisation sweeps
        eps:         small constant to prevent division by zero
    Returns:
        doubly_stochastic: (num_streams, num_streams) - rows & cols each sum to 1
    """
    # Exponentiate to ensure non-negativity
    matrix = torch.exp(log_weights)
    for _ in range(num_iters):
        # Row normalisation: each row sums to 1
        matrix = matrix / (matrix.sum(dim=-1, keepdim=True) + eps)
        # Column normalisation: each column sums to 1
        matrix = matrix / (matrix.sum(dim=-2, keepdim=True) + eps)
    return matrix

That’s the whole algorithm. Starting from any positive matrix, alternately normalizing rows and columns converges to a doubly stochastic matrix. Working in log-space (torch.exp(log_weights)) means the learnable parameters are unconstrained reals — the constraint is enforced by the projection, not the parameterization.

Wild, right? Five lines of core logic to tame signal explosion in hundred-layer networks.

The Code: The Full mHC Layer

Each mHC layer wraps a sub-layer (attention or MoE) and manages three mixing matrices:

class ManifoldHyperConnection(nn.Module):
    """Manifold-Constrained Hyper-Connection (mHC) wrapper.
    Wraps any sub-layer (attention, FFN, etc.) with multi-stream residual
    connections whose mixing matrices are constrained to preserve signal
    magnitude across deep networks.
    Design decisions:
      * H_res uses Sinkhorn-Knopp to stay on the Birkhoff Polytope,
        guaranteeing that residual streams neither explode nor vanish
        regardless of network depth.
      * H_pre and H_post use softplus (not exp) for non-negativity to
        give gentler gradients than raw exponentiation.
      * stream_index selects which stream is fed to the sub-layer,
        avoiding the cost of running F on all streams.
    """
    def __init__(
        self,
        num_streams: int,
        model_dim: int,
        stream_index: int = 0,
        sinkhorn_iters: int = 5,
    ):
        super().__init__()
        self.num_streams = num_streams
        self.model_dim = model_dim
        self.stream_index = stream_index
        self.sinkhorn_iters = sinkhorn_iters
        # Learnable raw weights - will be projected to Birkhoff Polytope
        self.h_res_raw = nn.Parameter(torch.zeros(num_streams, num_streams))
        # Pre-layer mixing (non-negative via softplus)
        self.h_pre_raw = nn.Parameter(torch.zeros(num_streams, num_streams))
        # Post-layer mixing (non-negative via softplus)
        self.h_post_raw = nn.Parameter(torch.zeros(num_streams, num_streams))
        self._init_weights()
    def _init_weights(self):
        """Initialise mixing matrices close to identity.
        H_res starts as uniform (which Sinkhorn maps to identity-like).
        H_pre and H_post start as identity so the initial behaviour
        closely matches a standard residual connection.
        """
        nn.init.zeros_(self.h_res_raw)
        nn.init.zeros_(self.h_pre_raw)
        with torch.no_grad():
            self.h_pre_raw.fill_diagonal_(1.0)
            self.h_post_raw.fill_diagonal_(1.0)
    def forward(self, streams, sublayer_fn, **sublayer_kwargs):
        """Apply mHC-wrapped sub-layer to multi-stream input.
        Args:
            streams:       (batch, seq_len, num_streams, model_dim)
            sublayer_fn:   callable(x, **kwargs) → x
            sublayer_kwargs: extra args forwarded to sublayer_fn
        Returns:
            streams_out:   (batch, seq_len, num_streams, model_dim)
        """
        batch_num, seq_len, n_s, d = streams.shape
        # === Compute constrained mixing matrices ===
        h_res = sinkhorn_knopp(self.h_res_raw, self.sinkhorn_iters)
        h_pre = F.softplus(self.h_pre_raw)
        h_post = F.softplus(self.h_post_raw)
        # ── Residual path: mix streams via doubly stochastic matrix ──
        flat = streams.view(-1, n_s, d)
        # h_res: (n_s, n_s) @ (n_s, d) per token → (batch*seq, n_s, d)
        residual = torch.einsum("ij, bje -> bie", h_res, flat)
        # ── Sub-layer path: select stream, run F, scatter back ──
        # Pre-mix: combine streams before sub-layer
        pre_mixed = torch.einsum("ij, bje -> bie", h_pre, flat)
        # Select the designated stream for the sub-layer
        sublayer_input = pre_mixed[:, self.stream_index, :].view(
            batch_num, seq_len, d
        )
        # Run the wrapped sub-layer (attention, MoE, etc.)
        sublayer_output = sublayer_fn(sublayer_input, **sublayer_kwargs)
        # Expand sub-layer output to all streams via h_post^T
        f_out = sublayer_output.view(-1, 1, d)
        # h_post^T[:, stream_index] gives the column weights for scattering
        post_weights = h_post[:, self.stream_index].unsqueeze(0).unsqueeze(-1)
        # Broadcast: (batch*seq, num_streams, model_dim)
        sublayer_contribution = f_out * post_weights
        # === Combine residual + sub-layer ===
        out = residual + sublayer_contribution
        return out.view(batch_num, seq_len, n_s, d)

Three matrices, three roles:

• H_res (doubly stochastic via Sinkhorn): mixes the residual streams. Because it’s doubly stochastic, total signal magnitude is preserved (the mathematical guarantee against explosion).
• H_pre (non-negative via softplus): mixes streams before the sub-layer. Determines what information the sub-layer sees.
• H_post (non-negative via softplus): scatters the sub-layer output back into multiple streams.

Why are H_pre and H_post only required to be non-negative (not doubly stochastic)? Because the sub-layer itself can change magnitudes, it’s the residual path that needs the hard constraint.

The einsum("ij, bje -> bie", h_res, flat) applies the doubly stochastic matrix to each token's streams, redistributing information across streams without amplifying it. The sub-layer only runs on one stream (stream_index=0 by default) to keep cost reasonable, then its output gets scattered back to all streams via h_post.

The layer update, in equation form:

streams_out = Sinkhorn(H_res) @ streams_in + softplus(H_post)ᵀ @ F(softplus(H_pre) @ streams_in)

Verifying It Works: Doubly Stochastic Property

model_dim_mhc = 512
num_streams = 4
mhc = ManifoldHyperConnection(
    num_streams=num_streams,
    model_dim=model_dim_mhc,
)
# Initialise streams: replicate input across all streams
x_base = torch.randn(2, 32, model_dim_mhc)
streams_in = x_base.unsqueeze(2).expand(-1, -1, num_streams, -1).clone()
# Dummy sub-layer: simple linear transform
dummy_layer = nn.Linear(model_dim_mhc, model_dim_mhc, bias=False)
streams_out = mhc(streams_in, sublayer_fn=lambda x: dummy_layer(x))
print(f"mHC input:  {streams_in.shape}")   # (2, 32, 4, 512)
print(f"mHC output: {streams_out.shape}")  # (2, 32, 4, 512)
# Verify doubly stochastic property of H_res
h_res = sinkhorn_knopp(mhc.h_res_raw, num_iters=10)
print(f"\nH_res row sums:    {h_res.sum(dim=-1).detach().cpu().numpy().round(4)}")
print(f"H_res col sums:    {h_res.sum(dim=-2).detach().cpu().numpy().round(4)}")
print(f"All entries >= 0:  {(h_res >= 0).all().item()}")
print("✓ H_res is doubly stochastic - signal magnitude preserved across layers")

Hybrid Attention: CSA + HCA — Taming Million-Token Contexts

The Problem: Quadratic Attention at a Million Tokens

Standard attention computes scores between every pair of tokens. That’s O(n²), at a million tokens, it's a trillion operations per layer. The KV cache alone would consume hundreds of gigabytes. This is the wall that prevents most models from truly using long contexts.

The Idea: Two Zoom Levels

DeepSeek V4’s solution is to combine two complementary attention mechanisms:

Mechanism Compression Selection Best for CSA (Compressed Sparse) 4× Top-k sparse Fine-grained, query-specific retrieval HCA (Heavily Compressed) 128× Dense (all) Broad global context, very cheap

Think of it like Google Maps. Sometimes you need the street-level view to find a specific building (CSA: mild compression, precise selection). Other times you just need the satellite view to understand the overall geography (HCA — heavy compression, global view). V4 interleaves both across its layers. Both mechanisms also include a sliding window for the most recent tokens because your immediate local context is almost always relevant.

The Code: Softmax-Gated Compressor

Both CSA and HCA share a compression module. Instead of simple average pooling, it uses a learned gate:

class SoftmaxGatedCompressor(nn.Module):
    """Compress KV sequences by grouping m tokens into one via learned gating.
    This is the shared compressor used by both CSA (m=4) and HCA (m=128).
    A learned positional bias inside each group controls how much each
    position contributes to the compressed representation.
    Design decision: softmax gating over the group is preferred over
    simple average pooling because different positions carry different
    information density. For example, the first token of a sentence
    often carries more semantic weight than mid-sentence tokens.
    """
    def __init__(self, group_size: int, num_heads: int):
        super().__init__()
        self.group_size = group_size
        # Learned positional bias for gating within each group
        # One bias vector per attention head for expressivity
        self.positional_bias = nn.Parameter(torch.zeros(num_heads, group_size))
    def forward(self, kv: torch.Tensor) -> torch.Tensor:
        """Compress a KV tensor along the sequence dimension.
        Args:
            kv: (batch, num_heads, seq_len, head_dim)
        Returns:
            compressed: (batch, num_heads, seq_len // group_size, head_dim)
        """
        batch_num, num_heads, seq_len, head_dim = kv.shape
        m = self.group_size
        # Truncate sequence to be divisible by group_size
        usable_len = (seq_len // m) * m
        kv = kv[:, :, :usable_len, :]
        num_groups = usable_len // m
        # Reshape into groups
        # (batch, heads, num_groups, group_size, head_dim)
        kv_grouped = kv.view(batch_num, num_heads, num_groups, m, head_dim)
        # Compute softmax gate weights from learned positional bias
        # (num_heads, group_size) → (1, num_heads, 1, group_size, 1)
        gate = F.softmax(self.positional_bias, dim=-1)
        gate = gate.view(1, num_heads, 1, m, 1)
        # Weighted sum within each group
        compressed = (kv_grouped * gate).sum(dim=3)
        return compressed

For CSA (group_size=4): 1,024 tokens → 256 compressed entries. For HCA (group_size=128): 1,024 tokens → 8 compressed entries.

The Code: Compressed Sparse Attention (CSA)

CSA does mild compression followed by query-dependent sparse selection:

class CompressedSparseAttention(nn.Module):
    """Compressed Sparse Attention (CSA) from DeepSeek-V4.   
    Two-stage process:
      1. Compress KV by group_size (4×) via learned softmax-gated pooling.
      2. For each query, score all compressed blocks (lightning indexer)
         and attend only to the top-k blocks.
      3. A sliding window branch attends to the most recent w uncompressed
         tokens for strong local context.
    """
    def __init__(
        self,
        num_heads: int,
        head_dim: int,
        group_size: int = 4,
        top_k: int = 64,
        window_size: int = 64,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.group_size = group_size
        self.top_k = top_k
        self.window_size = window_size
        self.compressor = SoftmaxGatedCompressor(group_size, num_heads)
    def forward(self, q, k, v):
        """Compute compressed sparse attention.
        Args:
            q, k, v: each (batch, num_heads, seq_len, head_dim)
        Returns:
            output: (batch, num_heads, seq_len, head_dim)
        """
        batch_num, num_heads, seq_len, head_dim = q.shape
        scale = 1.0 / math.sqrt(head_dim)
        # ── Compress K and V ──
        k_compressed = self.compressor(k)      # (batch, heads, seq//4, dim)
        v_compressed = self.compressor(v)
        num_blocks = k_compressed.shape[2]
        # Clamp top_k to available blocks
        effective_k = min(self.top_k, num_blocks)
        # ── Lightning indexer: score compressed blocks per query ──
        # (batch, heads, seq_len, num_blocks)
        block_scores = torch.matmul(q, k_compressed.transpose(-2, -1)) * scale
        # Select top-k compressed blocks per query position
        # (batch, heads, seq_len, top_k)
        topk_scores, topk_indices = torch.topk(
            block_scores, effective_k, dim=-1
        )
        # Gather the selected compressed KV entries
        idx_expanded = topk_indices.unsqueeze(-1).expand(
            -1, -1, -1, -1, head_dim
        )
        # (batch, heads, seq_len, top_k, head_dim)
        k_selected = k_compressed.unsqueeze(2).expand(
            -1, -1, seq_len, -1, -1
        ).gather(3, idx_expanded)
        v_selected = v_compressed.unsqueeze(2).expand(
            -1, -1, seq_len, -1, -1
        ).gather(3, idx_expanded)
        # ── Sparse attention over selected blocks ──
        # (batch, heads, seq, 1, dim) @ (batch, heads, seq, dim, top_k)
        sparse_scores = torch.matmul(
            q.unsqueeze(3), k_selected.transpose(-2, -1)
        ).squeeze(3) * scale
        sparse_weights = F.softmax(sparse_scores, dim=-1)
        # (batch, heads, seq, top_k) @ (batch, heads, seq, top_k, dim)
        sparse_context = torch.matmul(
            sparse_weights.unsqueeze(3), v_selected
        ).squeeze(3)
        # ── Sliding window for local context ──
        w = min(self.window_size, seq_len)
        k_recent = k[:, :, -w:, :]
        v_recent = v[:, :, -w:, :]
        q_recent = q[:, :, -w:, :]
        # Causal mask for local window
        local_mask = torch.tril(torch.ones(w, w, device=q.device))
        local_scores = torch.matmul(q_recent, k_recent.transpose(-2, -1)) * scale
        local_scores = local_scores.masked_fill(local_mask == 0, float("-inf"))
        local_weights = F.softmax(local_scores, dim=-1)
        local_context = torch.matmul(local_weights, v_recent)
        # ── Combine sparse global + local window ──
        output = sparse_context.clone()
        output[:, :, -w:, :] = (
            0.5 * sparse_context[:, :, -w:, :] + 0.5 * local_context
        )
        return output

This is the query-dependent part: different queries attend to different compressed blocks. A query about “climate” might select blocks containing climate-related tokens, while a query about “cooking” selects entirely different blocks. The gather operation collects only the top-k blocks per query.

The Code: Heavily Compressed Attention (HCA)

HCA takes the opposite trade-off: extreme compression, but no sparsity needed.

class HeavilyCompressedAttention(nn.Module):
    """Heavily Compressed Attention (HCA) from DeepSeek-V4.
    Compresses KV by 128× and runs full dense attention over the
    compressed sequence. Because the compressed sequence is extremely
    short, dense attention is cheap even at million-token contexts.
    Design decisions:
      * 128× compression is far more aggressive than CSA's 4×, but HCA
        compensates with dense (not sparse) attention, so no information
        is lost within the compressed representation.
      * CSA layers handle fine-grained, query-specific retrieval;
        HCA layers handle broad global context cheaply. The interleaving
        of both gives the model both capabilities.
    """
    def __init__(
        self,
        num_heads: int,
        head_dim: int,
        group_size: int = 128,
        window_size: int = 64,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.group_size = group_size
        self.window_size = window_size
        self.compressor = SoftmaxGatedCompressor(group_size, num_heads)
    def forward(self, q, k, v):
        """Compute heavily compressed dense attention.
        Args:
            q, k, v: each (batch, num_heads, seq_len, head_dim)
        Returns:
            output: (batch, num_heads, seq_len, head_dim)
        """
        batch_num, num_heads, seq_len, head_dim = q.shape
        scale = 1.0 / math.sqrt(head_dim)
        # ── Heavily compress K and V ──
        # 1M tokens → 128× compression → ~7,800 entries
        k_compressed = self.compressor(k)
        v_compressed = self.compressor(v)
        # ── Dense attention over ALL compressed entries ──
        # No top-k selection: the compressed sequence is short enough
        scores = torch.matmul(q, k_compressed.transpose(-2, -1)) * scale
        weights = F.softmax(scores, dim=-1)
        global_context = torch.matmul(weights, v_compressed)
        # ── Sliding window for local context ──
        w = min(self.window_size, seq_len)
        k_recent = k[:, :, -w:, :]
        v_recent = v[:, :, -w:, :]
        q_recent = q[:, :, -w:, :]
        local_mask = torch.tril(torch.ones(w, w, device=q.device))
        local_scores = torch.matmul(q_recent, k_recent.transpose(-2, -1)) * scale
        local_scores = local_scores.masked_fill(local_mask == 0, float("-inf"))
        local_weights = F.softmax(local_scores, dim=-1)
        local_context = torch.matmul(local_weights, v_recent)
        # ── Combine global compressed + local window ──
        output = global_context.clone()
        output[:, :, -w:, :] = (
            0.5 * global_context[:, :, -w:, :] + 0.5 * local_context
        )
        return output

Why no sparse selection? With 128× compression, a million tokens become ~7,800 entries. Dense attention over 7,800 entries is trivially cheap. And by attending to all compressed entries, HCA avoids the approximation error inherent in top-k selection.

The Complete V4 Block: Everything Snaps Together

Now let’s see how all these pieces assemble into a single DeepSeek V4 decoder block.

streams (batch, seq, num_streams, model_dim)
      │
  ┌───┴──────────────────────────┐
  │   mHC wrapper                │
  │   ├── Pre-mix streams        │
  │   ├── RMSNorm → MLA          │
  │   └── Post-mix → residual    │
  └───┬──────────────────────────┘
      │
  ┌───┴──────────────────────────┐
  │   mHC wrapper                │
  │   ├── Pre-mix streams        │
  │   ├── RMSNorm → DeepSeekMoE  │
  │   └── Post-mix → residual    │
  └───┬──────────────────────────┘
      │
  streams (batch, seq, num_streams, model_dim)

The V4 block has the same two sub-layers as a standard transformer (attention + FFN), but with three key differences: attention uses MLA paired with CSA or HCA, FFN uses DeepSeekMoE, and both are wrapped in mHC.

Here’s the full block with nothing abbreviated:

class DeepSeekV4Block(nn.Module):
    """Single DeepSeek-V4 decoder block.   
    Combines all V4 architectural innovations:
      * mHC replaces standard residual connections with multi-stream
        doubly-stochastic mixing for stable deep networks.
      * Hybrid attention alternates between CSA (fine-grained sparse)
        and HCA (broad dense) across layers.
      * DeepSeekMoE for the feed-forward component.
    Design decisions:
      * attn_type selects CSA or HCA per layer. In V4-Pro, layers 0-1
        use HCA, layers 2-60 alternate. This lets different layers
        specialise in global vs. local attention patterns.
      * mHC wraps both attention and MoE sub-layers independently,
        each with its own set of mixing matrices.
      * stream_index=0 means the "primary" stream is always the one
        processed by sub-layers.
    """
    def __init__(self, config: dict, attn_type: str = "csa"):
        super().__init__()
        model_dim = config["model_dim"]
        num_streams = config.get("num_streams", 4)
        num_heads = config["num_heads"]
        head_dim = config["head_dim"]
        # Pre-norms
        self.attn_norm = nn.RMSNorm(model_dim)
        self.moe_norm = nn.RMSNorm(model_dim)
        # MLA for Q/K/V projection (reused from V2)
        self.mla = MultiHeadLatentAttention(
            model_dim=model_dim,
            num_heads=num_heads,
            head_dim=head_dim,
            kv_compress_dim=config["kv_compress_dim"],
            q_compress_dim=config["q_compress_dim"],
            rope_head_dim=config["rope_head_dim"],
        )
        # Hybrid attention: CSA or HCA depending on layer position
        if attn_type == "csa":
            self.hybrid_attn = CompressedSparseAttention(
                num_heads=num_heads,
                head_dim=head_dim + config["rope_head_dim"],
                group_size=config.get("csa_group_size", 4),
                top_k=config.get("csa_top_k", 64),
                window_size=config.get("window_size", 64),
            )
        else:
            self.hybrid_attn = HeavilyCompressedAttention(
                num_heads=num_heads,
                head_dim=head_dim + config["rope_head_dim"],
                group_size=config.get("hca_group_size", 128),
                window_size=config.get("window_size", 64),
            )
        self.attn_type = attn_type
        # DeepSeekMoE
        self.moe = DeepSeekMoE(
            model_dim=model_dim,
            num_shared_experts=config["num_shared_experts"],
            num_routed_experts=config["num_routed_experts"],
            num_active_experts=config["num_active_experts"],
            expert_dim=config["expert_dim"],
        )
        # mHC wrappers for attention and MoE
        self.mhc_attn = ManifoldHyperConnection(
            num_streams=num_streams, model_dim=model_dim
        )
        self.mhc_moe = ManifoldHyperConnection(
            num_streams=num_streams, model_dim=model_dim
        )
    def _attn_sublayer(self, x, rope_freqs, mask):
        """Attention sub-layer: norm → MLA → output."""
        return self.mla(self.attn_norm(x), rope_freqs, mask)
    def _moe_sublayer(self, x):
        """MoE sub-layer: norm → MoE → output."""
        return self.moe(self.moe_norm(x))
    def forward(self, streams, rope_freqs, mask=None):
        """Forward pass through V4 block.
        Args:
            streams:    (batch, seq_len, num_streams, model_dim)
            rope_freqs: (max_len, rope_head_dim // 2) complex
            mask:       (batch, 1, seq_len, seq_len) or None
        Returns:
            streams:    (batch, seq_len, num_streams, model_dim)
        """
        # mHC-wrapped attention sub-layer
        streams = self.mhc_attn(
            streams,
            sublayer_fn=self._attn_sublayer,
            rope_freqs=rope_freqs,
            mask=mask,
        )
        # mHC-wrapped MoE sub-layer
        streams = self.mhc_moe(
            streams,
            sublayer_fn=self._moe_sublayer,
        )
        return streams

The _attn_sublayer and _moe_sublayer methods are tiny wrappers that add the norm. The mHC wrapper handles all the multi-stream mixing — the sub-layers themselves are unaware they're operating inside a multi-stream residual. They just see a regular (batch, seq, model_dim) input.

Running the Full V4 Block: CSA + HCA Interleaved

v4_config = dict(
    model_dim=512,
    num_heads=8,
    head_dim=64,
    kv_compress_dim=256,
    q_compress_dim=384,
    rope_head_dim=32,
    num_shared_experts=2,
    num_routed_experts=8,
    num_active_experts=2,
    expert_dim=256,
    num_streams=4,
    csa_group_size=4,
    csa_top_k=32,
    hca_group_size=128,
    window_size=32,
)
# Create one CSA layer and one HCA layer (as in the interleaved pattern)
block_csa = DeepSeekV4Block(v4_config, attn_type="csa")
block_hca = DeepSeekV4Block(v4_config, attn_type="hca")
# Initialise multi-stream input
x_init = torch.randn(2, 64, v4_config["model_dim"])
streams = x_init.unsqueeze(2).expand(-1, -1, v4_config["num_streams"], -1).clone()
rope_freqs_v4 = precompute_rope_frequencies(
    v4_config["rope_head_dim"], 256
)
causal_mask_v4 = torch.tril(
    torch.ones(64, 64)
).unsqueeze(0).unsqueeze(0)
# Forward through CSA block → HCA block (as in real V4)
streams_after_csa = block_csa(streams, rope_freqs_v4, causal_mask_v4)
print(f"After CSA block: {streams_after_csa.shape}")   # (2, 64, 4, 512) ✓
streams_after_hca = block_hca(streams_after_csa, rope_freqs_v4, causal_mask_v4)
print(f"After HCA block: {streams_after_hca.shape}")   # (2, 64, 4, 512) ✓
# Extract primary stream as final output
output = streams_after_hca[:, :, 0, :]
print(f"Final output (stream 0): {output.shape}")      # (2, 64, 512) ✓
# ── Parameter comparison ──
v2_params = sum(p.numel() for p in block_csa.mla.parameters()) + sum(
    p.numel() for p in block_csa.moe.parameters()
)
mhc_params = sum(p.numel() for p in block_csa.mhc_attn.parameters()) + sum(
    p.numel() for p in block_csa.mhc_moe.parameters()
)
hybrid_params = sum(p.numel() for p in block_csa.hybrid_attn.parameters())
print(f"\nParameter breakdown (CSA block):")
print(f"  MLA + MoE (V2/V3):       {v2_params / 1e6:.2f}M")
print(f"  mHC overhead:             {mhc_params / 1e3:.1f}K  (negligible)")
print(f"  Hybrid attention (CSA):   {hybrid_params / 1e3:.1f}K  (compressor only)")
print(f"  ✓ V4 innovations add <0.1% parameter overhead")

The mHC and hybrid attention components are almost free in terms of parameters. The efficiency gains come from how computation is organized (sparse experts, compressed KV cache, constrained residuals), not from adding more trainable parameters.

What This Unlocks

Let’s step back and see what DeepSeek V4’s architecture actually achieves:

• MLA shrinks the KV cache by 7–24× while matching or exceeding full MHA quality.
• MoE lets the model “know” 1.6 trillion parameters but only “use” 49 billion per token.
• MTP trains the model to anticipate multiple future tokens, giving richer gradients and enabling speculative decoding.
• mHC prevents signal explosion in deep networks through a mathematically elegant constraint from convex geometry.
• CSA + HCA compress million-token contexts to manageable sizes via two complementary zoom levels.

The combined result: V4-Pro uses 27% of V3.2’s FLOPs and 10% of its KV cache at 1M tokens, while maintaining frontier-level quality. And the pricing: $0.14 per million tokens for Flash, reflects these genuine architectural efficiencies.

The most exciting implication isn’t any single technique. It’s the composition. MoE, MLA, hybrid attention, and mHC are each independently published ideas, but DeepSeek showed that they stack cleanly. As context windows keep growing toward 10M+ tokens, and as models push past 200 layers, these techniques aren’t optional optimizations, they’re the foundation of what comes next.

Whether V4 reshapes the industry like R1 did remains to be seen. But architecturally, it’s the most interesting open-weight model released in 2026 so far. And now you know exactly how it works.

References & Resources

• Full code notebook: GitHub — ChanCheeKean/DecodeAI/…/C04 — DeepSeek.ipynb
• DeepSeek-V4 Technical Report: “DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence” (April 2026) — available on the HuggingFace model card
• DeepSeek-V2 Paper: arXiv 2405.04434 — where MLA and DeepSeekMoE were introduced
• Manifold-Constrained Hyper-Connections: arXiv 2512.24880 (DeepSeek, Jan 2026)
• Multi-Head Latent Attention explainer: Sebastian Raschka’s MLA walkthrough

A message from our Founder

Hey, Sunil here. I wanted to take a moment to thank you for reading until the end and for being a part of this community. Did you know that our team run these publications as a volunteer effort to over 3.5m monthly readers? We don’t receive any funding, we do this to support the community.

If you want to show some love, please take a moment to follow me on LinkedIn, TikTok, Instagram. You can also subscribe to our weekly newsletter. And before you go, don’t forget to clap and follow the writer️!