Mixture of Depths: Dynamic Computation in Transformers

import torch
import torch.nn as nn
import torch.nn.functional as F

class MoDRouting(nn.Module):
    """Routing mechanism for Mixture of Depths."""
    
    def __init__(self, d_model, num_layers, temperature=1.0):
        super().__init__()
        self.d_model = d_model
        self.num_layers = num_layers
        self.temperature = temperature
        
        # Routing network: predicts probability of continuing at each layer
        self.router = nn.Linear(d_model, 1)
        
    def forward(self, x):
        """Compute routing decisions for each token at each layer."""
        batch_size, seq_len, _ = x.shape
        
        # Compute routing scores at each layer position
        # Shape: (batch, seq, num_layers)
        routing_logits = self.router(x).squeeze(-1)  # (batch, seq)
        routing_logits = routing_logits.unsqueeze(-1).expand(-1, -1, self.num_layers)
        
        # Apply temperature and softmax for probability distribution
        routing_probs = F.softmax(routing_logits / self.temperature, dim=-1)
        
        # Sample layer positions based on routing probabilities
        # Each token is assigned to one layer position
        layer_positions = torch.multinomial(routing_probs.view(-1, self.num_layers), 1)
        layer_positions = layer_positions.view(batch_size, seq_len)
        
        return layer_positions, routing_probs


class MoDTransformerBlock(nn.Module):
    """Transformer block with optional early exit capability."""
    
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        
        # Self-attention
        self.qkv = nn.Linear(d_model, d_model * 3)
        self.attention = nn.Linear(d_model, d_model)
        self.attention_dropout = nn.Dropout(dropout)
        
        # Feed-forward
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )
        
        # Layer norms
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
    def forward(self, x, attention_mask=None):
        """Standard transformer block forward pass."""
        # Self-attention with pre-norm
        x_norm = self.norm1(x)
        q, k, v = self.qkv(x_norm).chunk(3, dim=-1)
        
        # Multi-head attention
        attn_output = self._multi_head_attention(q, k, v, attention_mask)
        x = x + self.attention_dropout(attn_output)
        
        # Feed-forward with pre-norm
        x = x + self.ffn(self.norm2(x))
        
        return x
    
    def _multi_head_attention(self, q, k, v, attention_mask=None):
        """Efficient multi-head attention computation."""
        batch_size, seq_len, d_model = q.shape
        q = q.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        
        # Scaled dot-product attention
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        
        if attention_mask is not None:
            scores = scores.masked_fill(attention_mask == 0, -1e9)
        
        attn_weights = F.softmax(scores, dim=-1)
        attn_output = torch.matmul(attn_weights, v)
        
        # Combine heads
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, d_model)
        
        return self.attention(attn_output)


class MoDTransformer(nn.Module):
    """Transformer with Mixture of Depths dynamic computation."""
    
    def __init__(self, vocab_size, d_model=512, n_heads=8, d_ff=2048, 
                 n_layers=12, max_seq_len=4096, dropout=0.1):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList([
            MoDTransformerBlock(d_model, n_heads, d_ff, dropout)
            for _ in range(n_layers)
        ])
        self.router = MoDRouting(d_model, n_layers)
        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size)
        self.max_seq_len = max_seq_len
        
    def forward(self, input_ids, position_ids=None):
        """Forward pass with MoD routing."""
        x = self.embed(input_ids)
        batch_size, seq_len = x.shape[:2]
        
        # Get routing decisions
        layer_positions, routing_probs = self.router(x)
        
        # Process through layers with early exit
        # Tokens exit after their assigned layer
        active_mask = torch.ones(batch_size, seq_len, len(self.layers), 
                                 device=x.device, dtype=torch.bool)
        
        for layer_idx, layer in enumerate(self.layers):
            # Only process tokens that haven't exited yet
            tokens_to_process = active_mask[:, :, layer_idx]
            
            if tokens_to_process.any():
                # Apply layer to active tokens
                x_processed = x[tokens_to_process]
                x_processed = layer(x_processed)
                x[tokens_to_process] = x_processed
            
            # Update active mask for next layer
            if layer_idx < len(self.layers) - 1:
                # Tokens exit if their layer position equals current layer
                exit_condition = (layer_positions == layer_idx)
                active_mask[:, :, layer_idx + 1] = active_mask[:, :, layer_idx] & ~exit_condition
        
        x = self.norm(x)
        return self.head(x), routing_probs

import torch
import torch.nn as nn


class HardwareAwareMoDLayer(nn.Module):
    """
    MoD layer with hardware-optimized depth bucketing.
    Groups tokens into fixed-depth buckets to enable efficient batched GEMM.
    """

    DEPTH_BUCKETS = [0.25, 0.5, 0.75, 1.0]  # 25%, 50%, 75%, 100% of layers

    def __init__(self, layers: nn.ModuleList, router: nn.Module):
        super().__init__()
        self.layers = layers
        self.router = router
        self.num_layers = len(layers)

    def _assign_bucket(self, depth_fraction: float) -> int:
        """Map a continuous depth to the nearest discrete bucket."""
        bucket_idx = min(
            range(len(self.DEPTH_BUCKETS)),
            key=lambda i: abs(self.DEPTH_BUCKETS[i] - depth_fraction),
        )
        return int(self.DEPTH_BUCKETS[bucket_idx] * self.num_layers)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        batch_size, seq_len, d_model = x.shape
        _, routing_probs = self.router(x)

        # Assign each token to a depth bucket
        per_token_depth = routing_probs.mean(dim=0)  # (seq_len,)
        token_depths = torch.tensor(
            [self._assign_bucket(d.item()) for d in per_token_depth],
            device=x.device,
        )

        output = x.clone()
        # Process each bucket as a contiguous batch (avoids scatter/gather overhead)
        for bucket_depth in set(token_depths.tolist()):
            mask = token_depths == bucket_depth
            if not mask.any():
                continue
            bucket_tokens = x[:, mask, :]  # (batch, bucket_size, d_model)
            for layer_idx in range(int(bucket_depth)):
                bucket_tokens = self.layers[layer_idx](bucket_tokens)
            output[:, mask, :] = bucket_tokens

        return output, routing_probs

"""
Minimal MoD transformer for experimentation.
Suitable for training on a single GPU with sequences up to 2K tokens.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class MinimalMoD(nn.Module):
    def __init__(
        self,
        vocab_size: int = 32000,
        d_model: int = 256,
        n_heads: int = 4,
        n_layers: int = 8,
        target_depth: float = 0.5,
    ):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=n_heads,
                dim_feedforward=d_model * 4,
                batch_first=True,
                norm_first=True,
            )
            for _ in range(n_layers)
        ])
        # Single router: predicts depth fraction [0, 1] per token
        self.router = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )
        self.head = nn.Linear(d_model, vocab_size)
        self.target_depth = target_depth
        self.n_layers = n_layers

    def forward(self, input_ids: torch.Tensor):
        x = self.embed(input_ids)
        depth_fractions = self.router(x).squeeze(-1)  # (batch, seq)

        # Convert fraction to number of layers (1 to n_layers)
        layer_counts = (depth_fractions * self.n_layers).long().clamp(1, self.n_layers)

        output = torch.zeros_like(x)
        for layer_idx, layer in enumerate(self.layers):
            # Only process tokens that need this layer
            active = layer_counts > layer_idx  # (batch, seq)
            if active.any():
                # Gather active tokens, process, scatter back
                active_x = x[active]
                # Transformer layer expects (batch, seq, d) — treat each token as a batch of 1
                active_x = layer(active_x.unsqueeze(0)).squeeze(0)
                x = x.clone()
                x[active] = active_x
            # Tokens that exit here write their output
            exiting = layer_counts == (layer_idx + 1)
            output[exiting] = x[exiting]

        logits = self.head(output)

        # Auxiliary loss: encourage target average depth
        depth_loss = F.mse_loss(depth_fractions.mean(), torch.tensor(self.target_depth))
        return logits, depth_loss

import time
import torch
from contextlib import contextmanager

@contextmanager
def timer(name: str):
    start = time.perf_counter()
    yield
    elapsed = (time.perf_counter() - start) * 1000
    print(f"{name}: {elapsed:.2f}ms")


def profile_mod_routing(model, input_ids, num_runs: int = 50):
    """Profile routing overhead vs transformer compute."""
    device = input_ids.device
    warmup = 5

    # Warmup
    for _ in range(warmup):
        with torch.no_grad():
            _ = model(input_ids)

    # Profile full forward pass
    forward_times = []
    routing_times = []

    for _ in range(num_runs):
        with torch.no_grad():
            # Full pass
            t0 = time.perf_counter()
            logits, routing_probs = model(input_ids)
            torch.cuda.synchronize()
            forward_times.append((time.perf_counter() - t0) * 1000)

            # Routing only
            t1 = time.perf_counter()
            _ = model.router(model.embed(input_ids))
            torch.cuda.synchronize()
            routing_times.append((time.perf_counter() - t1) * 1000)

    avg_forward = sum(forward_times) / num_runs
    avg_routing = sum(routing_times) / num_runs
    routing_pct = (avg_routing / avg_forward) * 100

    print(f"Average forward pass: {avg_forward:.2f}ms")
    print(f"Average routing overhead: {avg_routing:.2f}ms ({routing_pct:.1f}%)")

    # Compute effective compute utilization
    batch_size, seq_len = input_ids.shape
    with torch.no_grad():
        _, routing_probs = model(input_ids)
    avg_depth = routing_probs.mean().item()
    print(f"Average token depth: {avg_depth:.2%} of max layers")
    print(f"Compute savings: {(1 - avg_depth) * 100:.1f}%")

import torch
import torch.nn as nn
from dataclasses import dataclass
from typing import Optional

@dataclass
class MoDInferenceConfig:
    max_batch_size: int = 32
    max_seq_len: int = 4096
    target_avg_depth: float = 0.5  # Target 50% average depth
    min_depth_fraction: float = 0.25  # At least 25% of layers


class MoDProductionInference:
    """Production inference engine for MoD models with batching support."""

    def __init__(self, model: nn.Module, config: MoDInferenceConfig):
        self.model = model
        self.config = config
        self.model.eval()

    @torch.inference_mode()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 256,
        temperature: float = 1.0,
    ) -> torch.Tensor:
        """Generate tokens with MoD routing."""
        batch_size, prompt_len = input_ids.shape
        generated = input_ids.clone()

        # KV cache per layer (only populated for tokens that reached that layer)
        kv_cache = [None] * len(self.model.layers)

        for step in range(max_new_tokens):
            logits, routing_probs = self.model(generated)
            next_token_logits = logits[:, -1, :] / temperature
            next_tokens = torch.argmax(next_token_logits, dim=-1, keepdim=True)
            generated = torch.cat([generated, next_tokens], dim=-1)

            # Early stopping on EOS
            if (next_tokens == self.model.eos_token_id).all():
                break

        return generated[:, prompt_len:]

    def benchmark_throughput(
        self, batch_sizes: list[int], seq_len: int = 512
    ) -> dict:
        """Measure throughput at different batch sizes."""
        results = {}
        device = next(self.model.parameters()).device

        for bs in batch_sizes:
            input_ids = torch.randint(0, 32000, (bs, seq_len), device=device)

            # Warmup
            for _ in range(3):
                self.generate(input_ids, max_new_tokens=10)

            # Measure
            torch.cuda.synchronize()
            start = torch.cuda.Event(enable_timing=True)
            end = torch.cuda.Event(enable_timing=True)

            num_tokens = 100
            start.record()
            self.generate(input_ids, max_new_tokens=num_tokens)
            end.record()
            torch.cuda.synchronize()

            elapsed_ms = start.elapsed_time(end)
            tokens_per_sec = (bs * num_tokens) / (elapsed_ms / 1000)
            results[bs] = {
                "tokens_per_sec": round(tokens_per_sec, 1),
                "latency_ms_per_token": round(elapsed_ms / num_tokens, 2),
            }

        return results

from collections import defaultdict


class DepthAwareBatcher:
    """Batch requests by predicted depth for efficient GPU utilization."""

    def __init__(self, model, depth_buckets: list[float] = [0.25, 0.5, 0.75, 1.0]):
        self.model = model
        self.depth_buckets = depth_buckets

    def predict_depths(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Predict routing depths for a batch without full forward pass."""
        with torch.no_grad():
            embeddings = self.model.embed(input_ids)
            _, routing_probs = self.model.router(embeddings)
            # Average expected depth per sequence
            expected_depths = routing_probs.mean(dim=[1, 2])  # (batch,)
        return expected_depths

    def bucket_requests(
        self, requests: list[torch.Tensor]
    ) -> dict[float, list[torch.Tensor]]:
        """Group requests into depth buckets for efficient batching."""
        buckets = defaultdict(list)

        for req in requests:
            depth = self.predict_depths(req.unsqueeze(0)).item()
            # Assign to nearest bucket
            bucket = min(self.depth_buckets, key=lambda b: abs(b - depth))
            buckets[bucket].append(req)

        return dict(buckets)

def detect_routing_collapse(routing_probs: torch.Tensor, threshold: float = 0.1) -> bool:
    """
    Detect if routing has collapsed.
    routing_probs: (batch, seq_len, num_layers)
    Returns True if collapse detected.
    """
    avg_depth = routing_probs.mean().item()
    depth_std = routing_probs.std().item()

    if avg_depth > 0.95:
        print(f"WARNING: Routing collapsed to full depth (avg={avg_depth:.3f})")
        return True
    if avg_depth < 0.05:
        print(f"WARNING: Routing collapsed to early exit (avg={avg_depth:.3f})")
        return True
    if depth_std < 0.02:
        print(f"WARNING: Low routing variance (std={depth_std:.4f}) — uniform routing")
        return True

    return False

def routing_entropy_loss(routing_probs: torch.Tensor) -> torch.Tensor:
    """Auxiliary loss to prevent routing collapse."""
    # Encourage high entropy (diverse depth assignments)
    entropy = -(routing_probs * torch.log(routing_probs + 1e-10)).sum(-1)
    # Maximize entropy = minimize negative entropy
    return -entropy.mean()

# Set CUDA graphs for fixed-depth buckets to avoid kernel launch overhead
# In production, pre-compile separate CUDA graphs for each depth bucket
export TORCH_COMPILE_DEPTH_BUCKETS="0.25,0.5,0.75,1.0"