"""
DeepSeek-V4 Model Implementation
=================================
Open reproduction of DeepSeek-V4 architecture as a HuggingFace PreTrainedModel.

This is a pure PyTorch implementation that runs on standard GPUs (no custom CUDA kernels).
For production inference, the official repo uses tilelang kernels for FP4/FP8 GEMM and sparse attention.

Architecture summary:
- Multi-head Latent Attention (MLA) with compressed sparse attention
- DeepSeekMoE with hash routing + sqrtsoftplus scoring
- Manifold-Constrained Hyper-Connections (mHC) 
- Multi-Token Prediction (MTP)

Reference: DeepSeek-V4 Technical Report (2026)
           deepseek-ai/DeepSeek-V4-Pro inference/model.py
"""

import math
from dataclasses import dataclass
from typing import Optional, Tuple, List, Union
from functools import lru_cache

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from transformers import PreTrainedModel
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

from configuration_deepseek_v4 import DeepSeekV4Config


# =============================================================================
# Core Building Blocks
# =============================================================================

class DeepSeekV4RMSNorm(nn.Module):
    """RMSNorm with float32 computation for numerical stability."""
    
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps
    
    def forward(self, x: Tensor) -> Tensor:
        dtype = x.dtype
        x = x.float()
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return (self.weight * x).to(dtype)


# =============================================================================
# Rotary Position Embeddings (RoPE) with YaRN scaling
# =============================================================================

def precompute_freqs_cis(
    dim: int,
    seq_len: int,
    theta: float = 10000.0,
    rope_scaling: dict = None,
) -> Tensor:
    """Precompute complex-valued rotary embeddings with optional YaRN scaling."""
    
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
    
    if rope_scaling is not None and rope_scaling.get("type") == "yarn":
        factor = rope_scaling.get("factor", 1.0)
        original_max_pos = rope_scaling.get("original_max_position_embeddings", seq_len)
        beta_fast = rope_scaling.get("beta_fast", 32)
        beta_slow = rope_scaling.get("beta_slow", 1)
        
        def find_correction_dim(num_rotations, dim, base, max_seq_len):
            return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base))
        
        def find_correction_range(low_rot, high_rot, dim, base, max_seq_len):
            low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))
            high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))
            return max(low, 0), min(high, dim - 1)
        
        low, high = find_correction_range(beta_fast, beta_slow, dim, theta, original_max_pos)
        
        if low == high:
            high += 0.001
        linear_func = (torch.arange(dim // 2, dtype=torch.float32) - low) / (high - low)
        smooth = torch.clamp(linear_func, 0, 1)
        
        freqs = freqs / factor * (1 - smooth) + freqs * smooth
    
    t = torch.arange(seq_len, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis


def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
    """Apply rotary embeddings to the last dimension of x."""
    xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))
    
    # Reshape freqs_cis for broadcasting
    shape = [1] * (xc.ndim - 2) + list(xc.shape[-2:])
    if freqs_cis.ndim == 2:
        # [seq_len, dim//2] -> broadcast over batch and heads
        freqs_cis = freqs_cis.view(1, freqs_cis.size(0), 1, freqs_cis.size(1))
        if xc.ndim == 3:
            freqs_cis = freqs_cis.squeeze(2)  # no head dim
    
    # Slice freqs_cis to match sequence length
    seq_len = xc.size(-2) if xc.ndim >= 3 else xc.size(0)
    if freqs_cis.size(-2) > seq_len:
        freqs_cis = freqs_cis[..., :seq_len, :]
    
    result = torch.view_as_real(xc * freqs_cis).flatten(-2)
    return result.to(x.dtype)


# =============================================================================
# KV Compressor (Compressed Sparse Attention)
# =============================================================================

class KVCompressor(nn.Module):
    """Compresses KV cache via learned gated pooling.
    
    For compress_ratio=4: uses overlapping windows for smoother boundaries.
    For compress_ratio=128: simple non-overlapping pooling (HCA).
    
    Each group of `compress_ratio` consecutive tokens is compressed into a single
    KV vector using attention-weighted pooling with learnable positional biases (ape).
    """
    
    def __init__(self, config: DeepSeekV4Config, compress_ratio: int, head_dim: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.head_dim = head_dim
        self.rope_head_dim = config.qk_rope_head_dim
        self.compress_ratio = compress_ratio
        self.overlap = (compress_ratio == 4)  # only overlap for CSA, not HCA
        coff = 1 + int(self.overlap)
        
        # Learnable absolute position embeddings within each compression window
        self.ape = nn.Parameter(torch.empty(compress_ratio, coff * head_dim))
        
        # Projection + gating for compression
        self.wkv = nn.Linear(self.hidden_size, coff * head_dim, bias=False)
        self.wgate = nn.Linear(self.hidden_size, coff * head_dim, bias=False)
        self.norm = DeepSeekV4RMSNorm(head_dim, config.rms_norm_eps)
    
    def _overlap_transform(self, tensor: Tensor, fill_value: float = 0) -> Tensor:
        """For CSA (ratio=4): create overlapping windows for smoother compression."""
        b, s, _, _ = tensor.size()
        ratio, d = self.compress_ratio, self.head_dim
        new_tensor = tensor.new_full((b, s, 2 * ratio, d), fill_value)
        new_tensor[:, :, ratio:] = tensor[:, :, :, d:]
        new_tensor[:, 1:, :ratio] = tensor[:, :-1, :, :d]
        return new_tensor
    
    def forward(self, x: Tensor, freqs_cis: Tensor) -> Optional[Tensor]:
        """Compress a sequence of hidden states into pooled KV vectors.
        
        Args:
            x: [batch, seq_len, hidden_size]
            freqs_cis: precomputed rotary embeddings
        
        Returns:
            Compressed KV: [batch, seq_len // compress_ratio, head_dim] or None if seq too short
        """
        bsz, seqlen, _ = x.size()
        ratio = self.compress_ratio
        rd = self.rope_head_dim
        dtype = x.dtype
        
        if seqlen < ratio:
            return None
        
        x_f = x.float()
        kv = self.wkv(x_f)
        score = self.wgate(x_f)
        
        remainder = seqlen % ratio
        cutoff = seqlen - remainder
        
        if remainder > 0:
            kv = kv[:, :cutoff]
            score = score[:, :cutoff]
        
        # Reshape into compression groups
        kv = kv.unflatten(1, (-1, ratio))          # [b, n_groups, ratio, coff*d]
        score = score.unflatten(1, (-1, ratio))     # [b, n_groups, ratio, coff*d]
        score = score + self.ape                     # add positional bias
        
        if self.overlap:
            kv = self._overlap_transform(kv, 0)
            score = self._overlap_transform(score, float("-inf"))
        
        # Gated pooling: softmax over positions within each group, weighted sum
        kv = (kv * score.softmax(dim=2)).sum(dim=2)  # [b, n_groups, head_dim]
        kv = self.norm(kv.to(dtype))
        
        # Apply RoPE to rope dimensions of compressed KV
        if freqs_cis is not None:
            rope_slice = kv[..., -rd:]
            compressed_freqs = freqs_cis[:cutoff:ratio]
            if compressed_freqs.size(0) > rope_slice.size(1):
                compressed_freqs = compressed_freqs[:rope_slice.size(1)]
            kv = torch.cat([
                kv[..., :-rd],
                apply_rotary_emb(rope_slice, compressed_freqs)
            ], dim=-1)
        
        return kv


# =============================================================================
# Indexer (Learned Sparse Attention Selection)
# =============================================================================

class AttentionIndexer(nn.Module):
    """Selects top-k compressed KV positions for sparse attention via learned scoring.
    
    Has its own KV compressor to build compressed representations for scoring.
    Uses dot-product attention scores between queries and compressed keys,
    weighted by learned per-head importance weights.
    """
    
    def __init__(self, config: DeepSeekV4Config, compress_ratio: int = 4):
        super().__init__()
        self.n_heads = config.index_n_heads
        self.head_dim = config.index_head_dim
        self.rope_head_dim = config.qk_rope_head_dim
        self.index_topk = config.index_topk
        self.q_lora_rank = config.q_lora_rank
        self.compress_ratio = compress_ratio
        self.softmax_scale = self.head_dim ** -0.5
        
        # Query projection from Q lora bottleneck
        self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)
        # Per-head importance weights
        self.weights_proj = nn.Linear(config.hidden_size, self.n_heads, bias=False)
        
        # Dedicated compressor for indexing
        self.compressor = KVCompressor(config, compress_ratio, self.head_dim)
    
    def forward(
        self,
        x: Tensor,
        qr: Tensor,
        freqs_cis: Tensor,
        seq_len: int,
    ) -> Tensor:
        """Compute top-k indices of compressed KV positions to attend to.
        
        Args:
            x: [batch, seq_len, hidden_size] — input for building compressed KV
            qr: [batch, seq_len, q_lora_rank] — Q after first low-rank projection
            freqs_cis: rotary embeddings
            seq_len: sequence length
            
        Returns:
            topk_idxs: [batch, seq_len, topk] — indices into compressed KV cache
        """
        bsz = x.size(0)
        ratio = self.compress_ratio
        rd = self.rope_head_dim
        
        # Build query
        q = self.wq_b(qr)  # [b, s, n_heads * head_dim]
        q = q.view(bsz, seq_len, self.n_heads, self.head_dim)
        
        # Apply RoPE to query rope dims
        q_rope = apply_rotary_emb(q[..., -rd:], freqs_cis)
        q = torch.cat([q[..., :-rd], q_rope], dim=-1)
        
        # Build compressed KV for scoring
        compressed_kv = self.compressor(x, freqs_cis)  # [b, n_compressed, head_dim]
        if compressed_kv is None:
            return None
        
        n_compressed = compressed_kv.size(1)
        
        # Compute importance weights per token
        weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)
        
        # Score: q dot compressed_kv, weighted by importance
        # [b, s, n_heads, d] x [b, n_compressed, d] -> [b, s, n_heads, n_compressed]
        index_score = torch.einsum("bshd,btd->bsht", q, compressed_kv)
        index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)  # [b, s, n_compressed]
        
        # Causal masking: can't attend to future compressed positions
        positions = torch.arange(seq_len, device=x.device)
        compressed_positions = torch.arange(n_compressed, device=x.device)
        mask = compressed_positions.unsqueeze(0) >= (positions.unsqueeze(1) + 1) // ratio
        index_score = index_score + torch.where(mask.unsqueeze(0), float("-inf"), 0.0)
        
        # Select top-k
        k = min(self.index_topk, n_compressed)
        topk_idxs = index_score.topk(k, dim=-1)[1]  # [b, s, k]
        
        # Mask out invalid indices
        topk_idxs = torch.where(
            topk_idxs >= (positions.unsqueeze(1) + 1).unsqueeze(0) // ratio,
            -1,
            topk_idxs
        )
        
        return topk_idxs


# =============================================================================
# MLA with Compressed Sparse Attention
# =============================================================================

class DeepSeekV4Attention(nn.Module):
    """Multi-head Latent Attention with Compressed Sparse Attention.
    
    Key design choices in V4:
    1. Unified KV: Single head_dim=512 (instead of V3's separate nope/rope/v dims)
    2. Low-rank Q: q_lora_rank=1536 bottleneck with RMSNorm
    3. Grouped low-rank O: wo split into o_groups, each compressed to o_lora_rank
    4. Sliding window (128 tokens) + compressed global context
    5. Learnable attention sink per head
    6. QK normalization via RMSNorm on Q
    """
    
    def __init__(self, config: DeepSeekV4Config, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = config.head_dim
        self.rope_head_dim = config.qk_rope_head_dim
        self.nope_head_dim = config.head_dim - config.qk_rope_head_dim
        self.q_lora_rank = config.q_lora_rank
        self.o_groups = config.o_groups
        self.o_lora_rank = config.o_lora_rank
        self.window_size = config.window_size
        self.eps = config.rms_norm_eps
        
        # Compression ratio for this layer
        if layer_idx < len(config.compress_ratios):
            self.compress_ratio = config.compress_ratios[layer_idx]
        else:
            self.compress_ratio = 0
        
        self.softmax_scale = self.head_dim ** -0.5
        
        # Learnable attention sink (per head)
        self.attn_sink = nn.Parameter(torch.zeros(self.num_heads))
        
        # Q path: low-rank compression
        self.wq_a = nn.Linear(self.hidden_size, self.q_lora_rank, bias=False)
        self.q_norm = DeepSeekV4RMSNorm(self.q_lora_rank, self.eps)
        self.wq_b = nn.Linear(self.q_lora_rank, self.num_heads * self.head_dim, bias=False)
        
        # KV path: direct projection to unified head_dim
        self.wkv = nn.Linear(self.hidden_size, self.head_dim, bias=False)
        self.kv_norm = DeepSeekV4RMSNorm(self.head_dim, self.eps)
        
        # Output path: grouped low-rank projection
        # wo_a: [num_heads * head_dim // o_groups, o_groups * o_lora_rank]
        self.wo_a = nn.Linear(
            self.num_heads * self.head_dim // self.o_groups,
            self.o_groups * self.o_lora_rank,
            bias=False,
        )
        self.wo_b = nn.Linear(self.o_groups * self.o_lora_rank, self.hidden_size, bias=False)
        
        # Compressed attention components
        if self.compress_ratio > 0:
            self.compressor = KVCompressor(config, self.compress_ratio, self.head_dim)
            if self.compress_ratio == 4:  # CSA layers get an indexer
                self.indexer = AttentionIndexer(config, self.compress_ratio)
            else:
                self.indexer = None
    
    def forward(
        self,
        hidden_states: Tensor,
        attention_mask: Optional[Tensor] = None,
        position_ids: Optional[Tensor] = None,
        past_key_value: Optional[Tuple[Tensor]] = None,
        freqs_cis: Optional[Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        bsz, seq_len, _ = hidden_states.size()
        rd = self.rope_head_dim
        
        # ---- Q projection ----
        qr = self.q_norm(self.wq_a(hidden_states))  # [b, s, q_lora_rank]
        q = self.wq_b(qr)  # [b, s, num_heads * head_dim]
        q = q.view(bsz, seq_len, self.num_heads, self.head_dim)
        
        # QK normalization (V4-style: RMSNorm on Q)
        q = q * torch.rsqrt(q.float().pow(2).mean(-1, keepdim=True) + self.eps)
        q = q.to(hidden_states.dtype)
        
        # Apply RoPE to Q rope dims
        if freqs_cis is not None:
            q_nope = q[..., :-rd]
            q_rope = apply_rotary_emb(q[..., -rd:], freqs_cis)
            q = torch.cat([q_nope, q_rope], dim=-1)
        
        # ---- KV projection ----
        kv = self.kv_norm(self.wkv(hidden_states))  # [b, s, head_dim]
        
        # Apply RoPE to KV rope dims
        if freqs_cis is not None:
            kv_nope = kv[..., :-rd]
            kv_rope = apply_rotary_emb(kv[..., -rd:], freqs_cis)
            kv = torch.cat([kv_nope, kv_rope], dim=-1)
        
        # ---- Build full KV (window + compressed) ----
        # For training: always use full sequence with causal mask + window constraint
        # kv is [b, s, head_dim] — shared across all heads (MQA at KV level)
        
        # Standard attention with sliding window
        # Expand kv for multi-head: [b, s, 1, head_dim] -> broadcast to all heads
        k = kv.unsqueeze(2)  # [b, s, 1, head_dim]
        v = kv.unsqueeze(2)  # [b, s, 1, head_dim]
        
        # Compute attention scores
        # q: [b, s, num_heads, head_dim], k: [b, s, 1, head_dim]
        attn_weights = torch.einsum("bshd,bthd->bsht", q, k.expand(-1, -1, self.num_heads, -1))
        attn_weights = attn_weights * self.softmax_scale
        
        # Sliding window causal mask
        positions_q = torch.arange(seq_len, device=hidden_states.device)
        positions_k = torch.arange(seq_len, device=hidden_states.device)
        
        # Causal: q can attend to k where k_pos <= q_pos
        causal_mask = positions_k.unsqueeze(0) > positions_q.unsqueeze(1)
        
        # Window: q can only attend to k within window_size
        window_mask = (positions_q.unsqueeze(1) - positions_k.unsqueeze(0)) >= self.window_size
        
        combined_mask = causal_mask | window_mask  # [s, s]
        attn_weights = attn_weights.masked_fill(combined_mask.unsqueeze(0).unsqueeze(2), float("-inf"))
        
        # Add compressed KV attention if this layer has compression
        if self.compress_ratio > 0 and hasattr(self, 'compressor'):
            compress_freqs = freqs_cis  # use same freqs for compression
            compressed_kv = self.compressor(hidden_states, compress_freqs)
            
            if compressed_kv is not None:
                n_compressed = compressed_kv.size(1)
                ck = compressed_kv.unsqueeze(2)  # [b, n_c, 1, head_dim]
                
                # Attention to compressed positions
                compressed_attn = torch.einsum(
                    "bshd,bthd->bsht",
                    q,
                    ck.expand(-1, -1, self.num_heads, -1)
                )
                compressed_attn = compressed_attn * self.softmax_scale
                
                # Causal mask for compressed positions
                compressed_positions = torch.arange(n_compressed, device=hidden_states.device)
                cmask = compressed_positions.unsqueeze(0) >= (positions_q.unsqueeze(1) + 1) // self.compress_ratio
                compressed_attn = compressed_attn.masked_fill(cmask.unsqueeze(0).unsqueeze(2), float("-inf"))
                
                # Concatenate window attention and compressed attention
                attn_weights = torch.cat([attn_weights, compressed_attn], dim=-1)
                v_all = torch.cat([
                    v.expand(-1, -1, self.num_heads, -1),
                    ck.expand(-1, -1, self.num_heads, -1)
                ], dim=1)
            else:
                v_all = v.expand(-1, -1, self.num_heads, -1)
        else:
            v_all = v.expand(-1, -1, self.num_heads, -1)
        
        # Add attention sink
        attn_sink = self.attn_sink.view(1, 1, -1, 1)  # [1, 1, num_heads, 1]
        attn_weights = torch.cat([attn_weights, attn_sink.expand(bsz, seq_len, -1, 1)], dim=-1)
        sink_v = torch.zeros(bsz, 1, self.num_heads, self.head_dim, device=v_all.device, dtype=v_all.dtype)
        v_all = torch.cat([v_all, sink_v], dim=1)
        
        # Softmax + weighted sum
        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(hidden_states.dtype)
        
        # attn_output: [b, s, num_heads, head_dim]
        attn_output = torch.einsum("bsht,bthd->bshd", attn_weights, v_all)
        
        # ---- De-rotate RoPE on output (V4 applies inverse RoPE) ----
        if freqs_cis is not None:
            o_nope = attn_output[..., :-rd]
            # Inverse RoPE: conjugate of freqs_cis
            inv_freqs = freqs_cis.conj() if freqs_cis.is_complex() else freqs_cis
            o_rope = apply_rotary_emb(attn_output[..., -rd:], inv_freqs)
            attn_output = torch.cat([o_nope, o_rope], dim=-1)
        
        # ---- Output projection (grouped low-rank) ----
        # Split heads into groups, compress each group
        # attn_output: [b, s, num_heads, head_dim] -> [b, s, o_groups, heads_per_group * head_dim]
        attn_output = attn_output.view(bsz, seq_len, self.o_groups, -1)
        
        # wo_a weight shape: [o_groups * o_lora_rank, heads_per_group * head_dim]
        # Reshape for group-wise matmul
        wo_a_weight = self.wo_a.weight.view(self.o_groups, self.o_lora_rank, -1)
        
        # Group-wise low-rank compression
        # [b, s, o_groups, d] x [o_groups, o_lora_rank, d] -> [b, s, o_groups, o_lora_rank]
        o_compressed = torch.einsum("bsgd,grd->bsgr", attn_output, wo_a_weight)
        
        # Final projection
        output = self.wo_b(o_compressed.flatten(2))  # [b, s, hidden_size]
        
        return output, None


# =============================================================================
# MoE (Mixture of Experts) with Hash Routing
# =============================================================================

class DeepSeekV4Gate(nn.Module):
    """MoE routing gate supporting hash-based and score-based routing.
    
    V4 innovations:
    - Hash routing for first `num_hash_layers` layers (deterministic, no load balancing needed)
    - sqrtsoftplus scoring function: sqrt(softplus(x)) for smoother gradients
    - Bias-based load balancing for score-based routing
    """
    
    def __init__(self, config: DeepSeekV4Config, layer_idx: int):
        super().__init__()
        self.topk = config.num_experts_per_tok
        self.n_routed_experts = config.n_routed_experts
        self.scoring_func = config.scoring_func
        self.route_scale = config.routed_scaling_factor
        self.is_hash = layer_idx < config.num_hash_layers
        
        self.weight = nn.Parameter(torch.empty(config.n_routed_experts, config.hidden_size))
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        
        if self.is_hash:
            # Hash routing: deterministic expert assignment per token ID
            self.tid2eid = nn.Parameter(
                torch.randint(0, config.n_routed_experts, (config.vocab_size, self.topk)),
                requires_grad=False,
            )
            self.register_parameter("bias", None)
        else:
            # Score-based routing with learnable bias for load balancing
            self.bias = nn.Parameter(torch.zeros(config.n_routed_experts))
    
    def forward(
        self, x: Tensor, input_ids: Optional[Tensor] = None
    ) -> Tuple[Tensor, Tensor]:
        """
        Args:
            x: [batch * seq_len, hidden_size]
            input_ids: [batch * seq_len] (needed for hash routing)
            
        Returns:
            weights: [batch * seq_len, topk] — routing weights
            indices: [batch * seq_len, topk] — expert indices
        """
        scores = F.linear(x.float(), self.weight.float())
        
        if self.scoring_func == "softmax":
            scores = scores.softmax(dim=-1)
        elif self.scoring_func == "sigmoid":
            scores = scores.sigmoid()
        elif self.scoring_func == "sqrtsoftplus":
            scores = F.softplus(scores).sqrt()
        else:
            raise ValueError(f"Unknown scoring function: {self.scoring_func}")
        
        original_scores = scores
        
        # Add bias for expert selection (doesn't affect routing weights)
        if self.bias is not None:
            scores = scores + self.bias
        
        if self.is_hash and input_ids is not None:
            indices = self.tid2eid[input_ids]
        else:
            indices = scores.topk(self.topk, dim=-1)[1]
        
        weights = original_scores.gather(1, indices)
        
        # Normalize weights
        if self.scoring_func != "softmax":
            weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-8)
        
        weights = weights * self.route_scale
        
        return weights, indices


class DeepSeekV4Expert(nn.Module):
    """Single MoE expert: SwiGLU FFN with optional activation clamping."""
    
    def __init__(self, config: DeepSeekV4Config):
        super().__init__()
        self.w1 = nn.Linear(config.hidden_size, config.moe_intermediate_size, bias=False)
        self.w2 = nn.Linear(config.moe_intermediate_size, config.hidden_size, bias=False)
        self.w3 = nn.Linear(config.hidden_size, config.moe_intermediate_size, bias=False)
        self.swiglu_limit = config.swiglu_limit
    
    def forward(self, x: Tensor) -> Tensor:
        gate = self.w1(x).float()
        up = self.w3(x).float()
        
        if self.swiglu_limit > 0:
            up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit)
            gate = torch.clamp(gate, max=self.swiglu_limit)
        
        x = F.silu(gate) * up
        return self.w2(x.to(self.w2.weight.dtype))


class DeepSeekV4MoE(nn.Module):
    """Mixture-of-Experts layer with routed + shared experts."""
    
    def __init__(self, config: DeepSeekV4Config, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.n_routed_experts = config.n_routed_experts
        self.n_activated_experts = config.num_experts_per_tok
        
        self.gate = DeepSeekV4Gate(config, layer_idx)
        self.experts = nn.ModuleList([
            DeepSeekV4Expert(config) for _ in range(config.n_routed_experts)
        ])
        
        # Shared expert(s) — always active
        assert config.n_shared_experts == 1, "V4 uses exactly 1 shared expert"
        self.shared_expert = DeepSeekV4Expert(config)
    
    def forward(self, x: Tensor, input_ids: Optional[Tensor] = None) -> Tensor:
        """
        Args:
            x: [batch, seq_len, hidden_size]
            input_ids: [batch, seq_len] for hash routing
        """
        shape = x.size()
        x_flat = x.view(-1, self.hidden_size)
        
        flat_ids = input_ids.flatten() if input_ids is not None else None
        weights, indices = self.gate(x_flat, flat_ids)
        
        # Compute routed expert outputs
        y = torch.zeros_like(x_flat, dtype=torch.float32)
        
        for i in range(self.n_routed_experts):
            # Find which tokens route to this expert
            idx, top = torch.where(indices == i)
            if idx.numel() == 0:
                continue
            expert_input = x_flat[idx]
            expert_output = self.experts[i](expert_input)
            y[idx] += expert_output.float() * weights[idx, top, None]
        
        # Add shared expert output (always computed for all tokens)
        shared_output = self.shared_expert(x_flat)
        y = y + shared_output.float()
        
        return y.to(x.dtype).view(shape)


# =============================================================================
# Hyper-Connections (mHC)
# =============================================================================

class HyperConnection(nn.Module):
    """Manifold-Constrained Hyper-Connections.
    
    Replaces standard residual connections with a learned mixing scheme that
    maintains `hc_mult` copies of the hidden state. Uses Sinkhorn normalization
    for doubly-stochastic combination matrices.
    
    Pre-connection: Reduces hc_mult copies → 1 via learned weighted sum
    Post-connection: Expands 1 → hc_mult copies via learned weights + combination
    """
    
    def __init__(self, config: DeepSeekV4Config):
        super().__init__()
        self.hc_mult = config.hc_mult
        self.hidden_size = config.hidden_size
        self.eps = config.hc_eps
        self.norm_eps = config.rms_norm_eps
        self.sinkhorn_iters = config.hc_sinkhorn_iters
        
        hc_dim = config.hc_mult * config.hidden_size
        mix_hc = (2 + config.hc_mult) * config.hc_mult
        
        # Learnable parameters for mixing
        self.hc_fn = nn.Parameter(torch.empty(mix_hc, hc_dim))
        self.hc_base = nn.Parameter(torch.empty(mix_hc))
        self.hc_scale = nn.Parameter(torch.ones(3))
        
        nn.init.normal_(self.hc_fn, std=0.02)
        nn.init.zeros_(self.hc_base)
    
    def _sinkhorn(self, comb: Tensor) -> Tensor:
        """Apply Sinkhorn normalization for doubly-stochastic matrix."""
        hc = self.hc_mult
        eps = self.eps
        
        # Initial: softmax over rows + normalize columns
        comb = comb.softmax(dim=-1) + eps
        comb = comb / (comb.sum(dim=-2, keepdim=True) + eps)
        
        for _ in range(self.sinkhorn_iters - 1):
            comb = comb / (comb.sum(dim=-1, keepdim=True) + eps)
            comb = comb / (comb.sum(dim=-2, keepdim=True) + eps)
        
        return comb
    
    def pre(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
        """Pre-connection: reduce hc_mult copies to 1.
        
        Args:
            x: [batch, seq_len, hc_mult, hidden_size]
            
        Returns:
            y: [batch, seq_len, hidden_size] — reduced
            post_weights: [batch, seq_len, hc_mult] — for post-connection
            comb: [batch, seq_len, hc_mult, hc_mult] — combination matrix
        """
        bsz, seq_len, hc, d = x.size()
        dtype = x.dtype
        
        x_flat = x.flatten(2).float()  # [b, s, hc*d]
        
        # RMSNorm-like scaling
        rsqrt = torch.rsqrt(x_flat.pow(2).mean(-1, keepdim=True) + self.norm_eps)
        
        # Compute mixing coefficients
        mixes = F.linear(x_flat, self.hc_fn.float()) * rsqrt  # [b, s, mix_hc]
        
        # Split into pre, post, combination
        pre_raw = mixes[..., :hc]
        post_raw = mixes[..., hc:2*hc]
        comb_raw = mixes[..., 2*hc:].view(bsz, seq_len, hc, hc)
        
        # Pre-weights: sigmoid gating
        pre_weights = torch.sigmoid(pre_raw * self.hc_scale[0] + self.hc_base[:hc]) + self.eps
        
        # Post-weights: scaled sigmoid
        post_weights = 2 * torch.sigmoid(
            post_raw * self.hc_scale[1] + self.hc_base[hc:2*hc]
        )
        
        # Combination matrix: Sinkhorn-normalized
        comb = comb_raw * self.hc_scale[2] + self.hc_base[2*hc:].view(hc, hc)
        comb = self._sinkhorn(comb)
        
        # Weighted sum: reduce hc_mult copies to 1
        y = (pre_weights.unsqueeze(-1) * x).sum(dim=2)  # [b, s, d]
        
        return y.to(dtype), post_weights, comb
    
    def post(self, x: Tensor, residual: Tensor, post_weights: Tensor, comb: Tensor) -> Tensor:
        """Post-connection: expand 1 copy to hc_mult.
        
        Args:
            x: [batch, seq_len, hidden_size] — sublayer output
            residual: [batch, seq_len, hc_mult, hidden_size] — pre-sublayer state
            post_weights: [batch, seq_len, hc_mult]
            comb: [batch, seq_len, hc_mult, hc_mult]
            
        Returns:
            y: [batch, seq_len, hc_mult, hidden_size]
        """
        # x contributes via post_weights, residual mixes via combination matrix
        y = (post_weights.unsqueeze(-1) * x.unsqueeze(2) +
             torch.einsum("bsij,bsjd->bsid", comb, residual))
        return y.to(x.dtype)


class HyperConnectionHead(nn.Module):
    """Head-level hyper-connection: reduces hc_mult copies to 1 for the final output.
    
    Uses simpler sigmoid gating (no Sinkhorn) since this is the final reduction.
    """
    
    def __init__(self, config: DeepSeekV4Config):
        super().__init__()
        self.hc_mult = config.hc_mult
        self.norm_eps = config.rms_norm_eps
        self.hc_eps = config.hc_eps
        
        hc_dim = config.hc_mult * config.hidden_size
        
        self.hc_fn = nn.Parameter(torch.empty(config.hc_mult, hc_dim))
        self.hc_base = nn.Parameter(torch.empty(config.hc_mult))
        self.hc_scale = nn.Parameter(torch.ones(1))
        
        nn.init.normal_(self.hc_fn, std=0.02)
        nn.init.zeros_(self.hc_base)
    
    def forward(self, x: Tensor) -> Tensor:
        """
        Args:
            x: [batch, seq_len, hc_mult, hidden_size]
        Returns:
            y: [batch, seq_len, hidden_size]
        """
        shape = x.size()
        dtype = x.dtype
        
        x_flat = x.flatten(2).float()  # [b, s, hc*d]
        rsqrt = torch.rsqrt(x_flat.pow(2).mean(-1, keepdim=True) + self.norm_eps)
        
        mixes = F.linear(x_flat, self.hc_fn.float()) * rsqrt  # [b, s, hc]
        pre = torch.sigmoid(mixes * self.hc_scale + self.hc_base) + self.hc_eps
        
        y = (pre.unsqueeze(-1) * x).sum(dim=2)  # [b, s, d]
        return y.to(dtype)


# =============================================================================
# Transformer Block
# =============================================================================

class DeepSeekV4Block(nn.Module):
    """Transformer block with Hyper-Connections.
    
    Flow: hc_pre -> attn_norm -> attention -> hc_post -> hc_pre -> ffn_norm -> moe -> hc_post
    """
    
    def __init__(self, config: DeepSeekV4Config, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        
        self.attn = DeepSeekV4Attention(config, layer_idx)
        self.ffn = DeepSeekV4MoE(config, layer_idx)
        
        self.attn_norm = DeepSeekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.ffn_norm = DeepSeekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
        
        # Hyper-connections for attention and FFN sublayers
        self.hc_attn = HyperConnection(config)
        self.hc_ffn = HyperConnection(config)
    
    def forward(
        self,
        hidden_states: Tensor,
        attention_mask: Optional[Tensor] = None,
        position_ids: Optional[Tensor] = None,
        past_key_value: Optional[Tuple[Tensor]] = None,
        freqs_cis: Optional[Tensor] = None,
        input_ids: Optional[Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[Tensor, ...]:
        """
        Args:
            hidden_states: [batch, seq_len, hc_mult, hidden_size]
        """
        # ---- Attention sublayer ----
        residual = hidden_states
        x, post_weights, comb = self.hc_attn.pre(hidden_states)
        x = self.attn_norm(x)
        x, attn_weights = self.attn(
            x,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            freqs_cis=freqs_cis,
            output_attentions=output_attentions,
        )
        hidden_states = self.hc_attn.post(x, residual, post_weights, comb)
        
        # ---- FFN/MoE sublayer ----
        residual = hidden_states
        x, post_weights, comb = self.hc_ffn.pre(hidden_states)
        x = self.ffn_norm(x)
        x = self.ffn(x, input_ids)
        hidden_states = self.hc_ffn.post(x, residual, post_weights, comb)
        
        outputs = (hidden_states,)
        if output_attentions:
            outputs += (attn_weights,)
        
        return outputs


# =============================================================================
# MTP (Multi-Token Prediction) Block
# =============================================================================

class MTPBlock(nn.Module):
    """Multi-Token Prediction block.
    
    Projects both the previous layer's hidden state and the next token's embedding,
    combines them, processes through a full transformer block, and predicts the
    next-next token. Shares embedding and head weights with the main model.
    """
    
    def __init__(self, config: DeepSeekV4Config, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.hc_mult = config.hc_mult
        
        # Projections for combining embedding and hidden state
        self.e_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.h_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.enorm = DeepSeekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.hnorm = DeepSeekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
        
        # Full transformer block for processing
        self.block = DeepSeekV4Block(config, layer_idx)
        
        # Head-level HC for final prediction
        self.hc_head = HyperConnectionHead(config)
        self.norm = DeepSeekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
    
    def forward(
        self,
        hidden_states: Tensor,
        embed_tokens: nn.Module,
        lm_head: nn.Module,
        input_ids: Tensor,
        freqs_cis: Optional[Tensor] = None,
    ) -> Tensor:
        """
        Args:
            hidden_states: [batch, seq_len, hc_mult, hidden_size] from last transformer block
            embed_tokens: shared embedding layer
            lm_head: shared output head
            input_ids: [batch, seq_len] — token ids for embedding lookup
            
        Returns:
            logits: [batch, seq_len, vocab_size]
        """
        bsz, seq_len, hc, d = hidden_states.size()
        
        # Get embeddings of current tokens (shifted by 1 for next-token prediction)
        e = embed_tokens(input_ids)  # [b, s, d]
        e = self.enorm(e)
        
        # Reduce HC copies of hidden state
        h = self.hnorm(hidden_states)
        
        # Combine: embedding + hidden state projection
        # e_proj output is broadcast to hc copies
        x = self.e_proj(e).unsqueeze(2) + self.h_proj(h)  # [b, s, hc, d]
        
        # Process through transformer block
        x = self.block(
            x, freqs_cis=freqs_cis, input_ids=input_ids
        )[0]
        
        # Reduce to single hidden state and project to vocab
        x = self.hc_head(x)  # [b, s, d]
        x = self.norm(x)
        logits = lm_head(x)  # [b, s, vocab_size]
        
        return logits


# =============================================================================
# Full Model
# =============================================================================

class DeepSeekV4PreTrainedModel(PreTrainedModel):
    config_class = DeepSeekV4Config
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["DeepSeekV4Block"]
    
    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)


class DeepSeekV4Model(DeepSeekV4PreTrainedModel):
    """DeepSeek-V4 base model (no LM head)."""
    
    def __init__(self, config: DeepSeekV4Config):
        super().__init__(config)
        self.config = config
        
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        
        self.layers = nn.ModuleList([
            DeepSeekV4Block(config, i) for i in range(config.num_hidden_layers)
        ])
        
        self.norm = DeepSeekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
        
        # Head-level HC for reducing copies
        self.hc_head = HyperConnectionHead(config)
        
        # Precompute RoPE frequencies
        self._precompute_freqs(config)
        
        self.gradient_checkpointing = False
        self.post_init()
    
    def _precompute_freqs(self, config):
        """Precompute rotary embedding frequencies."""
        max_seq_len = min(config.max_position_embeddings, 8192)  # limit for memory
        freqs_cis = precompute_freqs_cis(
            config.qk_rope_head_dim,
            max_seq_len,
            config.rope_theta,
            config.rope_scaling,
        )
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
    
    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        position_ids: Optional[Tensor] = None,
        past_key_values: Optional[List[Tuple[Tensor]]] = None,
        inputs_embeds: Optional[Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        
        bsz, seq_len, dim = inputs_embeds.size()
        
        # Expand to hc_mult copies for Hyper-Connections
        hidden_states = inputs_embeds.unsqueeze(2).repeat(1, 1, self.config.hc_mult, 1)
        # [batch, seq_len, hc_mult, hidden_size]
        
        # Get RoPE frequencies
        if self.freqs_cis.device != hidden_states.device:
            self.freqs_cis = self.freqs_cis.to(hidden_states.device)
        
        if seq_len <= self.freqs_cis.size(0):
            freqs_cis = self.freqs_cis[:seq_len]
        else:
            # Recompute for longer sequences
            freqs_cis = precompute_freqs_cis(
                self.config.qk_rope_head_dim,
                seq_len,
                self.config.rope_theta,
                self.config.rope_scaling,
            ).to(hidden_states.device)
        
        all_hidden_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None
        
        for i, layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            
            if self.gradient_checkpointing and self.training:
                layer_outputs = torch.utils.checkpoint.checkpoint(
                    layer,
                    hidden_states,
                    attention_mask,
                    position_ids,
                    None,  # past_key_value
                    freqs_cis,
                    input_ids,
                    output_attentions,
                    use_reentrant=False,
                )
            else:
                layer_outputs = layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    freqs_cis=freqs_cis,
                    input_ids=input_ids,
                    output_attentions=output_attentions,
                )
            
            hidden_states = layer_outputs[0]
        
        # Reduce HC copies for output
        # hidden_states: [b, s, hc_mult, d] -> [b, s, d]
        reduced = self.hc_head(hidden_states)
        reduced = self.norm(reduced)
        
        return BaseModelOutputWithPast(
            last_hidden_state=reduced,
            past_key_values=None,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )


class DeepSeekV4ForCausalLM(DeepSeekV4PreTrainedModel):
    """DeepSeek-V4 for causal language modeling with optional MTP."""
    
    def __init__(self, config: DeepSeekV4Config):
        super().__init__(config)
        self.model = DeepSeekV4Model(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        # Multi-Token Prediction heads
        if config.num_nextn_predict_layers > 0:
            self.mtp = nn.ModuleList([
                MTPBlock(config, config.num_hidden_layers + i)
                for i in range(config.num_nextn_predict_layers)
            ])
        else:
            self.mtp = None
        
        # Head-level HC (shared between main and MTP)
        self.hc_head = HyperConnectionHead(config)
        
        self.post_init()
    
    def get_input_embeddings(self):
        return self.model.embed_tokens
    
    def set_input_embeddings(self, value):
        self.model.embed_tokens = value
    
    def get_output_embeddings(self):
        return self.lm_head
    
    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings
    
    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        position_ids: Optional[Tensor] = None,
        past_key_values: Optional[List[Tuple[Tensor]]] = None,
        inputs_embeds: Optional[Tensor] = None,
        labels: Optional[Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        
        # Get base model outputs
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=True,  # need for MTP
            return_dict=True,
        )
        
        hidden_states = outputs.last_hidden_state  # [b, s, d]
        logits = self.lm_head(hidden_states)
        
        loss = None
        if labels is not None:
            # Main language modeling loss
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )
            
            # MTP loss (predict token at position t+2, t+3, etc.)
            # This requires the full HC hidden states from before reduction
            # For simplicity in this open reproduction, MTP is optional during training
        
        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output
        
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


# =============================================================================
# Model Registration
# =============================================================================

# Auto-register with HF
DeepSeekV4Config.register_for_auto_class()
DeepSeekV4ForCausalLM.register_for_auto_class("AutoModelForCausalLM")